103 lines
3.9 KiB
C
103 lines
3.9 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/**
|
|
* @file fm_kernel_matvec_lane_split.h
|
|
* @brief Shared grid arithmetic for the K-quant matvecs: how many of a SIMD
|
|
* group's 32 lanes one output row gets, and therefore how many rows share a
|
|
* threadgroup. Host-side only; each family emits its own MSL reduction.
|
|
*
|
|
* @status Used by the q3_k, q4_k, q5_k and q6_k matvec families. q8_0 does not
|
|
* need it — its lane unit is a single weight, so every row already fills the
|
|
* SIMD group at any cols % 32 == 0.
|
|
* @issues None
|
|
* @todo None
|
|
*/
|
|
|
|
#ifndef FOUNDRY_METAL_KERNEL_MATVEC_LANE_SPLIT_H
|
|
#define FOUNDRY_METAL_KERNEL_MATVEC_LANE_SPLIT_H
|
|
|
|
/* ================================================================
|
|
* Why this exists
|
|
*
|
|
* Every K-quant matvec reduces its per-lane partials with SIMD-group
|
|
* primitives, so one threadgroup is exactly one SIMD group of 32 lanes. Give
|
|
* that group a single output row and the row's work units — super-blocks for
|
|
* q3_k/q4_k, 32-weight sub-blocks for q5_k/q6_k — have to fill all 32 lanes on
|
|
* their own. The MoE expert shapes cannot: a 512-column q3_k row is 2
|
|
* super-blocks, so 30 of 32 lanes idle and the kernel runs at ~5 GB/s against a
|
|
* ~90 GB/s roof (metal/tests/RESULTS_matvec_bandwidth.md).
|
|
*
|
|
* The fix that does NOT touch the decode path is to give the row only as many
|
|
* lanes as it has units and put several rows in the same SIMD group. Lane
|
|
* occupancy goes to 100% while each lane still decodes a whole unit exactly
|
|
* once, so the super-block header is not re-decoded — the measured trap that
|
|
* pins the 32-weight-unit kernels at 44-52 GB/s.
|
|
*
|
|
* Lanes per row is a POWER OF TWO so the cross-lane reduction is a
|
|
* simd_shuffle_xor butterfly that never crosses a row boundary: with a
|
|
* row-group based at a multiple of `lanes`, `lane ^ offset` for every
|
|
* offset < lanes stays inside that row's lanes.
|
|
* ================================================================ */
|
|
|
|
/* One threadgroup is one SIMD group; the reductions require it. */
|
|
#define FM_MATVEC_SIMD_LANES 32
|
|
|
|
/**
|
|
* Lanes to give one output row: the largest power of two that is <= the row's
|
|
* work-unit count and <= the SIMD width. A row with more units than lanes
|
|
* keeps the whole group and strides, exactly as before this split existed.
|
|
*
|
|
* @param units_per_row Lane-sized work units in one row (must be >= 1)
|
|
* @return Lanes per row, a power of two in [1, 32]
|
|
*
|
|
* Thread-safe: yes (pure)
|
|
* Reentrant: yes
|
|
* Signal-safe: yes
|
|
*/
|
|
static inline int fm_matvec_lanes_per_row(int units_per_row)
|
|
{
|
|
int lanes = 1;
|
|
|
|
while (lanes < FM_MATVEC_SIMD_LANES && lanes * 2 <= units_per_row)
|
|
lanes *= 2;
|
|
return lanes;
|
|
}
|
|
|
|
/**
|
|
* Output rows one threadgroup covers, given the lanes each row gets.
|
|
*
|
|
* @param lanes_per_row From fm_matvec_lanes_per_row
|
|
* @return Rows per threadgroup, a power of two in [1, 32]
|
|
*
|
|
* Thread-safe: yes (pure)
|
|
* Reentrant: yes
|
|
* Signal-safe: yes
|
|
*/
|
|
static inline int fm_matvec_rows_per_threadgroup(int lanes_per_row)
|
|
{
|
|
return FM_MATVEC_SIMD_LANES / lanes_per_row;
|
|
}
|
|
|
|
/**
|
|
* Threadgroups to dispatch for `rows` output rows. Rounds up, so the last
|
|
* group may carry row indices past the end — the kernels must accumulate zero
|
|
* and store nothing for those, while still entering the reduction so every
|
|
* lane of the SIMD group is active at the shuffles.
|
|
*
|
|
* @param rows Output rows (must be >= 1)
|
|
* @param lanes_per_row From fm_matvec_lanes_per_row
|
|
* @return Threadgroup count for the 1-D grid
|
|
*
|
|
* Thread-safe: yes (pure)
|
|
* Reentrant: yes
|
|
* Signal-safe: yes
|
|
*/
|
|
static inline int fm_matvec_threadgroup_count(int rows, int lanes_per_row)
|
|
{
|
|
int rows_per_group = fm_matvec_rows_per_threadgroup(lanes_per_row);
|
|
|
|
return (rows + rows_per_group - 1) / rows_per_group;
|
|
}
|
|
|
|
#endif /* FOUNDRY_METAL_KERNEL_MATVEC_LANE_SPLIT_H */
|