diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma.glsl new file mode 100644 index 00000000000..96c1915766e --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma.glsl @@ -0,0 +1,98 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * 8da4w (dynamically quantized int8 activations, 4-bit group-symmetric + * weights) int8 cooperative-matrix linear for subgroup-16 devices with an + * 8x16x32 int8 MMA, used by the architecture-family path + * (impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.cpp). Reads the 4h4w + * int8 activations of quantize_and_pack_4h4w_with_group_sums. Texel-wise B + * staging, multi-block A staging and an optional workgroup order + * (ORDER_NB_MAX). Template header only; the body is + * linear_dq8ca_q4gsw_wmma_body.glslh. + * + * Preconditions (not checked in the shader), enforced by the selection in + * CoopmatFamily.cpp: M % WG_TILE_M == 0, N % WG_TILE_N == 0, + * K % WG_TILE_K == 0, group_size % WG_TILE_K == 0, no bias, batch 1. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +// When the A staging thread map exactly covers the workgroup +// (A_ACTIVE_THREADS == WG_SIZE) the `a_active` guard is always true, but the +// driver compiler may not fold it because gl_LocalInvocationID.x's bound comes +// from a spec constant. Set A_MAP_FULL only for tiles where the equality holds. +$if A_MAP_FULL: + #define A_ALWAYS_ACTIVE + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings, in the argument order of the dispatch node: +// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3), +// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), +// weight_scales(8), bias(9). +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +// t_input is unread here -- the activations arrive already quantized in +// t_packed_int8_input -- but stays declared so the binding layout matches the +// dispatch site. +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +$if A_MULTI_BLOCK: + #define A_MULTI_BLOCK +const uint A_BLOCKS = ${A_BLOCKS}; + +// Workgroup order: 0 = dispatch order; > 0 = blocks of at most ORDER_NB_MAX +// N tiles, all M tiles of a block before the next block. +#define ORDER_NB_MAX ${ORDER_NB_MAX} + +#include "linear_dq8ca_q4gsw_wmma_body.glslh" diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma.yaml new file mode 100644 index 00000000000..0143c0bfd32 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma.yaml @@ -0,0 +1,51 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Variants of the architecture-family 8da4w int8 cooperative-matrix linear +# (4h4w int8 activations). Each one is referenced by a row in +# impl/gemm/linear/qw/CoopmatFamily.cpp, which also carries its tile geometry; +# the two must stay in sync. +# +# Name: txkgs +# m8 (MMA_M = 8), then i (dispatch order) or n (blocked order). + +linear_dq8ca_q4gsw_wmma: + parameter_names_with_default_values: + PRECISION: highp + IO_STORAGE: buffer + WEIGHT_STORAGE: texture2d + A_MAP_FULL: false + A_MULTI_BLOCK: false + A_BLOCKS: 1 + MMA_M: 8 + MMA_N: 16 + MMA_K: 32 + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 16 + ORDER_NB_MAX: 0 + shader_variants: + # Intel Xe2: subgroup 16, int8 MMA 8x16x32, dispatch order. + - NAME: linear_dq8ca_q4gsw_wmma_t128x128k64g84s16m8i_texture3d_texture2d_half + IO_STORAGE: texture3d + A_MAP_FULL: true + A_MULTI_BLOCK: true + - NAME: linear_dq8ca_q4gsw_wmma_t128x128k64g84s16m8i_buffer_texture2d_half + A_MAP_FULL: true + A_MULTI_BLOCK: true + # Same tile, blocked workgroup order (weight-heavy projections). + - NAME: linear_dq8ca_q4gsw_wmma_t128x128k64g84s16m8n13_texture3d_texture2d_half + IO_STORAGE: texture3d + A_MAP_FULL: true + A_MULTI_BLOCK: true + ORDER_NB_MAX: 13 + - NAME: linear_dq8ca_q4gsw_wmma_t128x128k64g84s16m8n13_buffer_texture2d_half + A_MAP_FULL: true + A_MULTI_BLOCK: true + ORDER_NB_MAX: 13 diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_body.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_body.glslh new file mode 100644 index 00000000000..9d3902b96c3 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_body.glslh @@ -0,0 +1,559 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Body of linear_dq8ca_q4gsw_wmma; included by its template header +// (linear_dq8ca_q4gsw_wmma.glsl), which defines bindings, spec constants and +// tile constants. Must not contain template syntax. +// +// Per quantization group the int8 x int8 products accumulate in int32 +// coopmats; each group adds float(acc) * weight_scale into an fp32 total. The +// activation zero point and scale are applied once at the end: +// out = input_scale * (total - input_zp * SUM_g weight_scale_g * weight_sum_g) + +#ifndef LINEAR_DQ8CA_Q4GSW_WMMA_BODY_GLSLH +#define LINEAR_DQ8CA_Q4GSW_WMMA_BODY_GLSLH + +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; +const uint B_USEFUL_U32 = MMA_K / 4u; +// B rows are not padded (stride 4 words), and each thread writes contiguous +// shared-memory words (see the staging index setup below). +const uint B_STRIDE_U32 = B_USEFUL_U32; +const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; +const uint NUM_K_SLABS = WG_TILE_K / MMA_K; + +const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; +const uint A_STRIDE_U32 = MMA_K / 4u; + +// One ping-pong slice covers all K-slabs of one chunk. +const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; +const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; + +// Double-buffered MMA operand staging. +shared uint Ash_int8[2u * ASH_SLICE_U32]; +shared uint Bsh_int8[2u * BSH_SLICE_U32]; + +// Per-WG-tile-row activation params (loaded ONCE at WG start; constant +// across groups). +shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast +shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast + +// Per-(group, output-channel) weight params, ping-ponged by group parity. +// (For per-channel INT8 only slice 0 is ever used.) +shared float wsc_sh[2u * WG_TILE_N]; +// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is +// accumulated once in the prologue. +shared float wcorr_sh[WG_TILE_N]; + + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue: SG_GRID_Y bands of MMA_M rows, +// each WG_TILE_N wide, row-major. A full WG_TILE_M x WG_TILE_N buffer would +// cost SG_GRID_Y/MMAS_PER_SG_M x more shared memory. float16_t-typed because +// coopMatStore needs it. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; +#endif + +// Running fp32 accumulator (across all groups). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Per-group int32 MMA accumulator. +coopmat + accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; + + +// Byte-parallel int4 -> int8 widening. +// +// The four nibbles this shader needs from one packed uint are ALREADY one per +// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four +// can be widened at once instead of with a per-nibble +// shift/mask/bias-subtract/mask chain. +// +// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit +// two's-complement pattern of v-8, because -8 == +8 (mod 16): +// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 +// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 +// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. +// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, +// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. +// +// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes +// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is +// logical rather than arithmetic. +uint widen_nibbles(const uint w, const uint parity) { + const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); + const uint p = nib ^ 0x08080808u; + const uint sgn = p & 0x08080808u; + return p | (sgn * 0x1Eu); +} + +// All four components of a packed-weight texel at once: `shift` is 0 or 4 +// (nibble parity). Same arithmetic as widen_nibbles, per component. +uvec4 widen_nibbles4(const uvec4 w, const uint shift) { + const uvec4 nib = (w >> shift) & uvec4(0x0F0F0F0Fu); + const uvec4 p = nib ^ uvec4(0x08080808u); + const uvec4 sgn = p & uvec4(0x08080808u); + return p | (sgn * 0x1Eu); +} + +void main() { + // The tile this workgroup computes is a remap of its launch index (the grid + // is (N tiles, M tiles), launched x-fastest). ORDER_NB_MAX == 0: dispatch + // order. ORDER_NB_MAX > 0: blocks of at most ORDER_NB_MAX N tiles, evened + // out; all M tiles of a block run before the next block (N-fastest inside a + // block; the last block may be narrower). A bijection of the tiles. + const uint ord_nx = gl_NumWorkGroups.x; + const uint ord_lin = gl_WorkGroupID.y * ord_nx + gl_WorkGroupID.x; +#if ORDER_NB_MAX > 0 + const uint ord_ny = gl_NumWorkGroups.y; + const uint ord_nblk = (ord_nx + uint(ORDER_NB_MAX) - 1u) / uint(ORDER_NB_MAX); + const uint ord_nb = (ord_nx + ord_nblk - 1u) / ord_nblk; + const uint ord_full = ord_nb * ord_ny; + const uint ord_b = ord_lin / ord_full; + const uint ord_n0 = ord_b * ord_nb; + const uint ord_w = min(ord_nb, ord_nx - ord_n0); + const uint ord_r = ord_lin - ord_b * ord_full; + const uvec2 tileID = uvec2(ord_n0 + ord_r % ord_w, ord_r / ord_w); +#else + const uvec2 tileID = uvec2(ord_lin % ord_nx, ord_lin / ord_nx); +#endif + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint N = uint(output_sizes.x); + const uint N4 = (N + 3u) / 4u; + const uint nblocks_x_A = (K + 3u) >> 2u; + + const uint num_groups = uint(num_groups_arg); + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; + const uint num_chunks = num_groups * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + accum_int32[i][j] = coopmat(0); + } + } + + // --- A staging thread map: one (m4, k4) ivec4 block per active thread --- + const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; + const uint A_ACTIVE_THREADS = (WG_TILE_M >> 2u) * K_BLOCKS_PER_CHUNK; + const uint a_m_block = gl_LocalInvocationID.x / K_BLOCKS_PER_CHUNK; + const uint a_k_block = gl_LocalInvocationID.x % K_BLOCKS_PER_CHUNK; +#ifdef A_ALWAYS_ACTIVE + // A_ACTIVE_THREADS == WG_SIZE for this variant's tile, so every thread stages + // A and the guard is unconditionally true. + const bool a_active = true; +#else + const bool a_active = gl_LocalInvocationID.x < A_ACTIVE_THREADS; +#endif + + // --- B staging thread map: (block, col) slots; each slot extracts one + // ColumnMajor shared uint (4 K-contiguous sign-extended int8) --- + // Texel-wise: a slot is (texel, nibble parity) and yields 4 shared words + // (the texel's 4 components = 4 output columns x 4 K). + const uint N8_PER_TILE = WG_TILE_N >> 3u; + const uint B_TOTAL_SLOTS = 2u * K_BLOCKS_PER_CHUNK * N8_PER_TILE; + // Round up; B_SLOT_ACTIVE masks the slots past the last one. + const uint B_SLOTS_PER_THREAD = (B_TOTAL_SLOTS + WG_SIZE - 1u) / WG_SIZE; +#define B_SLOT_ACTIVE(si) (B_TOTAL_SLOTS % WG_SIZE == 0u || gl_LocalInvocationID.x + (si) * WG_SIZE < B_TOTAL_SLOTS) + + // Loop-invariant staging index math, hoisted: every term below depends only + // on gl_LocalInvocationID.x, si, tile_n_start and tile_m_start, not on + // `chunk` or `group_i`. Only the k4 texel row varies with the chunk, by a + // constant stride. + uint b_lds_off[B_SLOTS_PER_THREAD]; // shared store offset within a slice + uint b_comp[B_SLOTS_PER_THREAD]; // which ivec4 component feeds this slot + uint b_par[B_SLOTS_PER_THREAD]; // nibble parity for this slot + uint b_n8blk[B_SLOTS_PER_THREAD]; // global texel column (N/8 blocks) + uint b_k4off[B_SLOTS_PER_THREAD]; // k4 offset of this slot within a chunk + // Coalesced writes: start from the contiguous shared-memory index this + // thread writes and derive which global element it needs, rather than + // deriving the shared address from the natural fetch grouping (which gives + // B_STRIDE_U32-strided writes). + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint s = gl_LocalInvocationID.x + si * WG_SIZE; + const uint par = s & 1u; + const uint tex = s >> 1u; + const uint k4_in_chunk = tex % K_BLOCKS_PER_CHUNK; + const uint n8_in_tile = tex / K_BLOCKS_PER_CHUNK; + const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); + // Shared word of component 0; component c is c * B_STRIDE_U32 further. + b_lds_off[si] = slab_idx * B_SLAB_U32 + + (n8_in_tile * 8u + par * 4u) * B_STRIDE_U32 + k4_in_slab; + b_comp[si] = 0u; // unused: all four components are staged + b_par[si] = par * 4u; // nibble shift + b_n8blk[si] = (tile_n_start >> 3u) + n8_in_tile; + b_k4off[si] = k4_in_chunk; + } + // A staging: same argument. base_row/slab/k_uint are all invariant. +#ifdef A_MULTI_BLOCK + // A_BLOCKS int8 blocks per thread (blocks tid + i * WG_SIZE). + uint a_lds_off[A_BLOCKS]; // slice-relative shared offset of the block's row 0 + uint a_glb[A_BLOCKS]; // global int8 block index at chunk 0 + [[unroll]] for (uint ai = 0; ai < A_BLOCKS; ++ai) { + const uint blk = gl_LocalInvocationID.x + ai * WG_SIZE; + const uint mb = blk / K_BLOCKS_PER_CHUNK; + const uint kb = blk % K_BLOCKS_PER_CHUNK; + a_lds_off[ai] = (kb / (MMA_K >> 2u)) * A_SLAB_U32 + + (mb * 4u) * A_STRIDE_U32 + + (kb % (MMA_K >> 2u)); + a_glb[ai] = ((tile_m_start >> 2u) + mb) * nblocks_x_A + kb; + } + ivec4 temp_A[A_BLOCKS]; +#else + const uint a_lds_off0 = + (a_k_block / (MMA_K >> 2u)) * A_SLAB_U32 + + (a_m_block * 4u) * A_STRIDE_U32 + + (a_k_block % (MMA_K >> 2u)); + const uint a_glb_row = ((tile_m_start >> 2u) + a_m_block) * nblocks_x_A; + + // Prefetch temp registers. + ivec4 temp_A; +#endif + ivec4 temp_B[B_SLOTS_PER_THREAD]; + float temp_wsc; + + // ========================================================= + // PROLOGUE + // ========================================================= + if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { + const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; + const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); + const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); + const uint base = gl_LocalInvocationID.x * 4u; + ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; + ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; + izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; + izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; + } + // Group 0 weight scales -> slice 0, and the hoisted weight-side correction + // SUM_g wsc[g][n]*wsum[g][n] accumulated across all groups. The loop is + // prologue-only and replaces per-group wsum work inside the loop body. + if (gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; + wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); + + float corr = 0.0; + for (uint g = 0; g < num_groups; ++g) { + f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; + corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); + } + wcorr_sh[gl_LocalInvocationID.x] = corr; + } + memoryBarrierShared(); + barrier(); + + // izp/ifs are needed only after the group loop, so they are loaded there + // (not kept live across it). + + // Prefetch chunk 0 into temp registers, then store to slice 0 (no barrier + // here -- the main loop's first iteration barriers before reading slice 0). + if (a_active) { +#ifdef A_MULTI_BLOCK + [[unroll]] for (uint ai = 0; ai < A_BLOCKS; ++ai) { + temp_A[ai] = t_packed_int8_input[a_glb[ai]]; + } +#else + const uint m4_global = (tile_m_start >> 2u) + a_m_block; + temp_A = t_packed_int8_input[m4_global * nblocks_x_A + a_k_block]; +#endif + } + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) if (B_SLOT_ACTIVE(si)) { + temp_B[si] = texelFetch(t_packed_weight, ivec2(b_k4off[si], b_n8blk[si]), 0); + } + { + // store chunk 0 -> slice 0 +#ifndef A_ALWAYS_ACTIVE + if (a_active) +#endif + { +#ifdef A_MULTI_BLOCK + [[unroll]] for (uint ai = 0; ai < A_BLOCKS; ++ai) { + [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { + Ash_int8[a_lds_off[ai] + m4i * A_STRIDE_U32] = uint(temp_A[ai][m4i]); + } + } +#else + [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { + Ash_int8[a_lds_off0 + m4i * A_STRIDE_U32] = uint(temp_A[m4i]); + } +#endif + } + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) if (B_SLOT_ACTIVE(si)) { + const uvec4 bw = widen_nibbles4(uvec4(temp_B[si]), b_par[si]); + [[unroll]] for (uint c = 0; c < 4u; ++c) { + Bsh_int8[b_lds_off[si] + c * B_STRIDE_U32] = bw[c]; + } + } + } + + // ========================================================= + // MAIN LOOP: nested groups x chunks (kept nested; flattening it with a + // conditional coopmat epilog has crashed a driver compiler at large + // spec-resolved trip counts). One barrier per chunk. Chunk + // iteration (global index `chunk`): + // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk + // of group g, wsum/wsc slice (g%2) is too. + // 2. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 + // starts a new group, also its wsum/wsc element. Skipped + // entirely on the final chunk. + // 3. int8 MMA — on slice (chunk%2) into accum_int32. + // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; + // on a group boundary, wsum/wsc -> slice ((g+1)%2). + // The group epilog runs unconditionally at the tail of each group. + // ========================================================= + uint chunk = 0; + for (uint group_i = 0; group_i < num_groups; ++group_i) { + for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { + const bool has_next = chunk + 1u < num_chunks; + const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); + const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; + const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; + const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; + const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; + + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + + // --- 2. prefetch chunk+1 -> temp --- + if (has_next) { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; +#ifndef A_ALWAYS_ACTIVE + if (a_active) +#endif + { +#ifdef A_MULTI_BLOCK + [[unroll]] for (uint ai = 0; ai < A_BLOCKS; ++ai) { + temp_A[ai] = t_packed_int8_input[a_glb[ai] + (chunkK_nxt >> 2u)]; + } +#else + temp_A = t_packed_int8_input[a_glb_row + (chunkK_nxt >> 2u) + a_k_block]; +#endif + } + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) if (B_SLOT_ACTIVE(si)) { + const uint k4_blk = (chunkK_nxt >> 2u) + b_k4off[si]; + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8blk[si]), 0); + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + } + } + + // --- 3. int8 MMA on the cur slice --- + [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { + const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; + const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash_int8, + slab_a_base_u32 + row_a * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh_int8, + slab_b_base_u32 + col_b * B_STRIDE_U32, + B_STRIDE_U32, + gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); + } + } + } + + // --- 4. store temp (chunk+1) -> nxt slice --- + if (has_next) { +#ifndef A_ALWAYS_ACTIVE + if (a_active) +#endif + { +#ifdef A_MULTI_BLOCK + [[unroll]] for (uint ai = 0; ai < A_BLOCKS; ++ai) { + [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { + Ash_int8[nxt_a + a_lds_off[ai] + m4i * A_STRIDE_U32] = uint(temp_A[ai][m4i]); + } + } +#else + [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { + Ash_int8[nxt_a + a_lds_off0 + m4i * A_STRIDE_U32] = uint(temp_A[m4i]); + } +#endif + } + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) if (B_SLOT_ACTIVE(si)) { + const uvec4 bw = widen_nibbles4(uvec4(temp_B[si]), b_par[si]); + [[unroll]] for (uint c = 0; c < 4u; ++c) { + Bsh_int8[nxt_b + b_lds_off[si] + c * B_STRIDE_U32] = bw[c]; + } + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; + wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + } + } + } // chunks + + // --- Group epilog: scale-only accumulate, reset accum --- + // Just result += float(acc) * wsc. The zero-point subtract and the ifs + // multiply are hoisted out of the group loop (applied once below). + { + const uint wbase = (group_i % 2u) * WG_TILE_N; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + coopmat wsc_bcast; + coopMatLoad( + wsc_bcast, wsc_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] += + coopmat( + accum_int32[i][j]) * wsc_bcast; + accum_int32[i][j] = coopmat(0); + } + } + } + } // groups + + // --- Hoisted correction, applied ONCE: --------------------------------- + // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) + // izp/ifs are loaded here rather than before the group loop so they are not + // live across it. + { + coopmat + izpf_bcast[MMAS_PER_SG_M]; + coopmat + ifs_bcast[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopmat izp_i; + coopMatLoad( + izp_i, izp_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + izpf_bcast[i] = + coopmat(izp_i); + coopMatLoad( + ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + } + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat wcorr_bcast; + coopMatLoad( + wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); + } + } + } + + // --- Store result tile --- + // N for the store address math MUST come from the spec constant, not the + // sizes UBO (see out_N_arg above). +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global + // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the + // writing subgroup's warpInTile.y, so the global row reproduces the buffer + // path's gi exactly. + // + // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled + // despite [[unroll]], so result[i][j] is dynamically indexed. Coopmat arrays + // are opaque per-lane storage and dynamic indexing is a construct shader + // compilers have broken before -- check this first if the texture variants + // miscompile on a new device. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh_out against the previous iteration's readers. Inert on i == 0 + // but must stay unconditional to remain workgroup-uniform. + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, Csh_out, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh_out[base]), + float(Csh_out[base + 1u]), + float(Csh_out[base + 2u]), + float(Csh_out[base + 3u]))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} + +#endif // LINEAR_DQ8CA_Q4GSW_WMMA_BODY_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor.glsl new file mode 100644 index 00000000000..35980ece92d --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor.glsl @@ -0,0 +1,101 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * 8da4w int8 cooperative-matrix linear that reads row-major int8 activations + * (utils::kPackedInt8_4W, written by quantize_and_pack_4w_with_group_sums), + * used by the architecture-family path + * (impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.cpp). The 4h4w layout + * is not row-major, so no coopMatLoad can address it; this kernel needs the + * row-major packer. Template header only; the body is + * linear_dq8ca_q4gsw_wmma_rowmajor_body.glslh, which: + * - stages A as raw 16-byte vectors (uvec4 global load -> uvec4 shared + * store) instead of coopMatLoad(global) -> coopMatStore(shared); needs + * WG_TILE_M * WG_TILE_K / 16 to be a multiple of the workgroup size; + * - lets one weight texel fetch feed both nibble parities of a component + * (output columns c and c + 4); + * - stages the texture3d drain band in the A shared-memory tile, which is + * dead after the last MMA. + * + * Preconditions (not checked in the shader), enforced by the selection in + * CoopmatFamily.cpp: N % WG_TILE_N == 0, K % WG_TILE_K == 0, + * group_size % WG_TILE_K == 0, no bias, batch 1. M is rounded up by the grid. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +// A is bound as a scalar int8_t array so that a coopMatLoad from it has a +// matching component type. +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings, in the argument order of the dispatch node: +// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3), +// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), +// weight_scales(8), bias(9). +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +// t_input is unread here -- the activations arrive already quantized in +// t_packed_int8_input -- but stays declared so the binding layout matches the +// dispatch site. +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +// Row-major packed activations (row stride = K int8). +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} +// uvec4 view of binding 2 (t_packed_int8_input) for 16-byte A staging. +// Spelled layout(std430, set...) on purpose: gen_vulkan_spv.py collects the +// descriptor list from every line matching ^layout\(set, so the usual spelling +// would add a descriptor. +layout(std430, set = 0, binding = 2) buffer restrict readonly t_packed_int8_input_v4Buffer { + uvec4 t_packed_int8_input_v4[]; +}; +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; + +#include "linear_dq8ca_q4gsw_wmma_rowmajor_body.glslh" diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor.yaml new file mode 100644 index 00000000000..6351a1f65f0 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor.yaml @@ -0,0 +1,30 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Variants of the row-major-activation 8da4w int8 cooperative-matrix linear. +# Each one is referenced by a row in impl/gemm/linear/qw/CoopmatFamily.cpp, +# which also carries its tile geometry; the two must stay in sync. +# +# Name: txkgs +# mk32 (MMA_K = 32). + +linear_dq8ca_q4gsw_wmma_rowmajor: + parameter_names_with_default_values: + PRECISION: highp + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 32 + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + shader_variants: + # NVIDIA: subgroup 32, int8 MMA 16x16x32. + - NAME: linear_dq8ca_q4gsw_wmma_rowmajor_t128x128k64g24s32mk32_texture3d_texture2d_half diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor_body.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor_body.glslh new file mode 100644 index 00000000000..61658d0ede5 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor_body.glslh @@ -0,0 +1,500 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Body of linear_dq8ca_q4gsw_wmma_rowmajor; included by its template header +// (linear_dq8ca_q4gsw_wmma_rowmajor.glsl), which defines bindings, spec +// constants and tile constants. Must not contain template syntax. +// +// Same arithmetic as linear_dq8ca_q4gsw_wmma_body.glslh (int32 accumulation +// per quantization group, fp32 total, zero-point correction applied once at +// the end); differs in the A layout (row-major int8, staged as raw 16-byte +// vectors) and in B staging (one thread per whole weight texel). +// +// The last M tile may be partial: its rows at or past M read row M - 1 of the +// activations and their params (each output row depends only on its own +// activation row) and are not stored, so no buffer is read past M rows. + +#ifndef LINEAR_DQ8CA_Q4GSW_WMMA_ROWMAJOR_BODY_GLSLH +#define LINEAR_DQ8CA_Q4GSW_WMMA_ROWMAJOR_BODY_GLSLH + +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; +const uint B_USEFUL_U32 = MMA_K / 4u; +// B rows are not padded (stride 4 words). +const uint B_STRIDE_U32 = B_USEFUL_U32; +const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; +const uint NUM_K_SLABS = WG_TILE_K / MMA_K; + +const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; +const uint A_STRIDE_U32 = MMA_K / 4u; + +// One ping-pong slice covers all K-slabs of one chunk. +const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; +const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; + +// MMA operand staging in two ping-pong slices. +const uint SLICES = 2u; +shared uvec4 Ash_int8[SLICES * ASH_SLICE_U32 / 4u]; +const uint ASH_UNIT = 4u; // u32 per Ash_int8 element +shared uint Bsh_int8[SLICES * BSH_SLICE_U32]; + +// Per-WG-tile-row activation params (loaded ONCE at WG start; constant +// across groups). +shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast +shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast + +// Per-(group, output-channel) weight params, ping-ponged by group parity. +// (For per-channel INT8 only slice 0 is ever used.) +shared float wsc_sh[2u * WG_TILE_N]; +// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is +// accumulated once in the prologue. +shared float wcorr_sh[WG_TILE_N]; + + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue: SG_GRID_Y bands of MMA_M rows, +// each WG_TILE_N wide, row-major, staged in Ash_int8 (dead after the last +// MMA; every drain iteration starts with a barrier). +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +// Compile-time check: the drain band (in u32) must fit in Ash_int8. +const int kCshFitsInAsh[(CSH_ROWS * WG_TILE_N / 2u <= SLICES * ASH_SLICE_U32) ? 1 : -1] = int[](0); +#endif + +// Running fp32 accumulator (across all groups). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Per-group int32 MMA accumulator. +coopmat + accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; + + +// Byte-parallel int4 -> int8 widening. +// +// The four nibbles this shader needs from one packed uint are ALREADY one per +// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four +// can be widened at once instead of with a per-nibble +// shift/mask/bias-subtract/mask chain. +// +// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit +// two's-complement pattern of v-8, because -8 == +8 (mod 16): +// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 +// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 +// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. +// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, +// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. +// +// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes +// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is +// logical rather than arithmetic. +uint widen_nibbles(const uint w, const uint parity) { + const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); + const uint p = nib ^ 0x08080808u; + const uint sgn = p & 0x08080808u; + return p | (sgn * 0x1Eu); +} + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint N = uint(output_sizes.x); + const uint N4 = (N + 3u) / 4u; + const uint nblocks_x_A = (K + 3u) >> 2u; + // A row stride in int8 elements (the binding is row-major int8), derived + // from nblocks_x_A so it matches the packer's `m_row * K4 + k4` addressing + // exactly (K % 4 == 0 makes it equal to K). + const uint a_row_stride_i8 = nblocks_x_A * 4u; + + const uint num_groups = uint(num_groups_arg); + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; + const uint num_chunks = num_groups * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + const uint M_rows = uint(input_sizes.y); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + accum_int32[i][j] = coopmat(0); + } + } + + const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; + + // --- A tile counts (MMA_M x MMA_K tiles per chunk) --- + const uint A_TILES_M = WG_TILE_M / MMA_M; + const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS + const uint NUM_A_TILES = A_TILES_M * A_TILES_K; + const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; + + // --- B staging thread map: (block, col) slots; each slot extracts one + // ColumnMajor shared uint (4 K-contiguous sign-extended int8) --- + const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; + // A slot is a whole texel (k4, n8) producing eight shared uints. Threads + // beyond the texel count of a chunk have no slot (b_on). + const uint B_TEXELS = B_TOTAL_SLOTS / 8u; + const uint B_SLOTS_PER_THREAD = (B_TEXELS + WG_SIZE - 1u) / WG_SIZE; + const uint N8_PER_TILE = WG_TILE_N >> 3u; + + // Loop-invariant B staging index math, hoisted out of the chunk loop. + uint b_lds_off[B_SLOTS_PER_THREAD]; // shared store offset within a slice + uint b_comp[B_SLOTS_PER_THREAD]; // which ivec4 component feeds this slot + uint b_par[B_SLOTS_PER_THREAD]; // nibble parity for this slot + uint b_n8blk[B_SLOTS_PER_THREAD]; // global texel column (N/8 blocks) + uint b_k4off[B_SLOTS_PER_THREAD]; // k4 offset of this slot within a chunk + bool b_on[B_SLOTS_PER_THREAD]; // this thread owns a texel in this slot + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint q = gl_LocalInvocationID.x + si * WG_SIZE; + const uint K4_PER_SLAB = MMA_K >> 2u; + const uint k4_in_slab = q % K4_PER_SLAB; + // Lanes run over k4 within a slab fastest, then slab, then the texel + // column block. + const uint slab_idx = (q / K4_PER_SLAB) % NUM_K_SLABS; + const uint n8_in_tile = (q / (K4_PER_SLAB * NUM_K_SLABS)) % N8_PER_TILE; + const uint k4_in_chunk = slab_idx * K4_PER_SLAB + k4_in_slab; + const uint n_col = n8_in_tile * 8u; // component c is column n_col + c, parity 1 is + 4 + b_lds_off[si] = slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab; + b_comp[si] = 0u; + b_on[si] = q < B_TEXELS; + b_par[si] = 0u; + b_n8blk[si] = (tile_n_start >> 3u) + n8_in_tile; + b_k4off[si] = k4_in_chunk; + } + + // A staging thread map: 16-byte part fastest, then row, then K slab, so 8 + // lanes write 128 contiguous shared bytes and a subgroup reads whole + // 32-byte global sectors. + const uint A_PARTS = MMA_K / 16u; + const uint A_VECS = WG_TILE_M * WG_TILE_K / 16u; + const uint A_VECS_PER_THREAD = A_VECS / WG_SIZE; + const int kAVecsDivide[(A_VECS % WG_SIZE == 0u && A_VECS >= WG_SIZE) ? 1 : -1] = int[](0); + uint a_lds[A_VECS_PER_THREAD]; // Ash_int8 u32 offset within a slice + uint a_gv4[A_VECS_PER_THREAD]; // uvec4 index of this thread's chunk-0 data + [[unroll]] for (uint v = 0; v < A_VECS_PER_THREAD; ++v) { + const uint q = gl_LocalInvocationID.x + v * WG_SIZE; + const uint part = q % A_PARTS; + const uint row = (q / A_PARTS) % WG_TILE_M; + const uint slab = q / (A_PARTS * WG_TILE_M); + a_lds[v] = slab * A_SLAB_U32 + row * A_STRIDE_U32 + part * 4u; + const uint m_row = min(tile_m_start + row, M_rows - 1u); + a_gv4[v] = (m_row * a_row_stride_i8 + slab * MMA_K + part * 16u) / 16u; + } + uvec4 temp_Av[A_VECS_PER_THREAD]; + ivec4 temp_B[B_SLOTS_PER_THREAD]; + float temp_wsc; + + // ========================================================= + // PROLOGUE + // ========================================================= + if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { + const uint m4 = min( + (tile_m_start >> 2u) + gl_LocalInvocationID.x, (M_rows + 3u) / 4u - 1u); + const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); + const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); + const uint base = gl_LocalInvocationID.x * 4u; + ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; + ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; + izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; + izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; + } + // Group 0 weight scales -> slice 0, and the hoisted weight-side correction + // SUM_g wsc[g][n]*wsum[g][n] accumulated across all groups. + if (gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; + wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); + + float corr = 0.0; + for (uint g = 0; g < num_groups; ++g) { + f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; + corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); + } + wcorr_sh[gl_LocalInvocationID.x] = corr; + } + memoryBarrierShared(); + barrier(); + + // izp/ifs are needed only after the group loop, so they are loaded there. + + // Prefetch chunk 0 into temp registers, then store to slice 0 (no barrier + // here -- the main loop's first iteration barriers before reading slice 0). + [[unroll]] for (uint v = 0; v < A_VECS_PER_THREAD; ++v) { + temp_Av[v] = t_packed_int8_input_v4[a_gv4[v]]; + } + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + if (b_on[si]) { + temp_B[si] = texelFetch(t_packed_weight, ivec2(b_k4off[si], b_n8blk[si]), 0); + } + } + { + // store chunk 0 -> slice 0 + [[unroll]] for (uint v = 0; v < A_VECS_PER_THREAD; ++v) { + Ash_int8[a_lds[v] / 4u] = temp_Av[v]; + } + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + if (b_on[si]) { + Bsh_int8[b_lds_off[si] + 0u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][0]), 0u); + Bsh_int8[b_lds_off[si] + 4u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][0]), 1u); + Bsh_int8[b_lds_off[si] + 1u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][1]), 0u); + Bsh_int8[b_lds_off[si] + 5u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][1]), 1u); + Bsh_int8[b_lds_off[si] + 2u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][2]), 0u); + Bsh_int8[b_lds_off[si] + 6u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][2]), 1u); + Bsh_int8[b_lds_off[si] + 3u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][3]), 0u); + Bsh_int8[b_lds_off[si] + 7u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][3]), 1u); + } + } + } + + // ========================================================= + // MAIN LOOP: nested groups x chunks (kept nested; flattening it with a + // conditional coopmat epilog has crashed a driver compiler at large + // spec-resolved trip counts). One barrier per chunk. Chunk + // iteration (global index `chunk`): + // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk + // of group g, wsc slice (g%2) is too. + // 2. prefetch — chunk+1 (A tiles, B blocks) into temp; when chunk+1 + // starts a new group, also its wsc element. Skipped + // entirely on the final chunk. + // 3. int8 MMA — on slice (chunk%2) into accum_int32. + // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; + // on a group boundary, wsc -> slice ((g+1)%2). + // The group epilog runs unconditionally at the tail of each group. + // ========================================================= + uint chunk = 0; + for (uint group_i = 0; group_i < num_groups; ++group_i) { + for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { + const bool has_next = chunk + 1u < num_chunks; + const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); + const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; + const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; + const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; + const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; + + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + + // --- 2. prefetch chunk+1 -> temp --- + if (has_next) { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + [[unroll]] for (uint v = 0; v < A_VECS_PER_THREAD; ++v) { + temp_Av[v] = t_packed_int8_input_v4[a_gv4[v] + chunkK_nxt / 16u]; + } + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint k4_blk = (chunkK_nxt >> 2u) + b_k4off[si]; + if (b_on[si]) { + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8blk[si]), 0); + } + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + } + } + + // --- 3. int8 MMA on the cur slice --- + [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { + const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; + const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash_int8, + (slab_a_base_u32 + row_a * A_STRIDE_U32) / ASH_UNIT, + A_STRIDE_U32 / ASH_UNIT, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh_int8, + slab_b_base_u32 + col_b * B_STRIDE_U32, + B_STRIDE_U32, + gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); + } + } + } + + // --- 4. store temp (chunk+1) -> nxt slice --- + if (has_next) { + [[unroll]] for (uint v = 0; v < A_VECS_PER_THREAD; ++v) { + Ash_int8[(nxt_a + a_lds[v]) / 4u] = temp_Av[v]; + } + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + if (b_on[si]) { + Bsh_int8[nxt_b + b_lds_off[si] + 0u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][0]), 0u); + Bsh_int8[nxt_b + b_lds_off[si] + 4u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][0]), 1u); + Bsh_int8[nxt_b + b_lds_off[si] + 1u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][1]), 0u); + Bsh_int8[nxt_b + b_lds_off[si] + 5u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][1]), 1u); + Bsh_int8[nxt_b + b_lds_off[si] + 2u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][2]), 0u); + Bsh_int8[nxt_b + b_lds_off[si] + 6u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][2]), 1u); + Bsh_int8[nxt_b + b_lds_off[si] + 3u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][3]), 0u); + Bsh_int8[nxt_b + b_lds_off[si] + 7u * B_STRIDE_U32] = widen_nibbles(uint(temp_B[si][3]), 1u); + } + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; + wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + } + } + } // chunks + + // --- Group epilog: scale-only accumulate, reset accum --- + // Just result += float(acc) * wsc. The zero-point subtract and the ifs + // multiply are hoisted out of the group loop (applied once below). + { + const uint wbase = (group_i % 2u) * WG_TILE_N; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + coopmat wsc_bcast; + coopMatLoad( + wsc_bcast, wsc_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] += + coopmat( + accum_int32[i][j]) * wsc_bcast; + accum_int32[i][j] = coopmat(0); + } + } + } + } // groups + + // --- Hoisted correction, applied ONCE: -------------------------------- + // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) + // izp/ifs are loaded here rather than before the group loop so they are not + // live across it. + { + coopmat + izpf_bcast[MMAS_PER_SG_M]; + coopmat + ifs_bcast[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopmat izp_i; + coopMatLoad( + izp_i, izp_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + izpf_bcast[i] = + coopmat(izp_i); + coopMatLoad( + ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + } + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat wcorr_bcast; + coopMatLoad( + wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); + } + } + } + + // --- Store result tile --- + // N for the store address math MUST come from the spec constant, not the + // sizes UBO (see out_N_arg above). +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into the drain band at once, so the SG_GRID_Y bands it holds are disjoint global + // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the + // writing subgroup's warpInTile.y, so the global row reproduces the buffer + // path's gi exactly. + // + // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled + // despite [[unroll]], so result[i][j] is dynamically indexed. Coopmat arrays + // are opaque per-lane storage and dynamic indexing is a construct shader + // compilers have broken before -- check this first if the texture variants + // miscompile on a new device. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards the drain band against the previous iteration's readers. Inert on i == 0 + // but must stay unconditional to remain workgroup-uniform. + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, Ash_int8, + (warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j)) / (2u * ASH_UNIT), + WG_TILE_N / (2u * ASH_UNIT), + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + // The last M tile may be partial: rows at or past M are computed but + // not stored. + if (m >= uint(output_sizes.y)) { + continue; + } + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4(unpackHalf2x16((base % 8u) != 0u ? Ash_int8[base / 8u].z : Ash_int8[base / 8u].x), + unpackHalf2x16((base % 8u) != 0u ? Ash_int8[base / 8u].w : Ash_int8[base / 8u].y))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} + +#endif // LINEAR_DQ8CA_Q4GSW_WMMA_ROWMAJOR_BODY_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/quantize_and_pack_4w_with_group_sums.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/quantize_and_pack_4w_with_group_sums.glsl new file mode 100644 index 00000000000..25a48549cb3 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/quantize_and_pack_4w_with_group_sums.glsl @@ -0,0 +1,157 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Row-major (utils::kPackedInt8_4W) counterpart of +// quantize_and_pack_4h4w_with_group_sums.glsl. +// +// Identical quantization and per-group sum reduction; only the output layout +// differs. The 4h4w version writes an ivec4 per (m4, k4) block whose component +// selects one of 4 rows -- as a scalar array that is +// index = m4*(4*K4) + k4*4 + r, which is not affine in the row index and so +// cannot be addressed by a cooperative-matrix load. This version writes each +// of the 4 quantized rows to its own row-major slot, +// t_packed_int8_input[m * K4 + k4], giving plain row-major int8 with 4 +// K-contiguous values per int32 and a row stride of K4. +// +// Consumed only by linear_dq8ca_q4gsw_wmma_rowmajor. Buffer output only. + +#version 450 core + +${define_required_extensions(INPUT_STORAGE, DTYPE)} +${define_required_extensions("texture3d", "int8")} + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, INPUT_STORAGE)} +#define T ${texel_load_component_type(DTYPE, INPUT_STORAGE)} + +$if OUTPUT_STORAGE == "buffer": + #define OUTPUT_BUFFER +$if INPUT_STORAGE == "buffer": + #define INPUT_BUFFER + +#extension GL_EXT_integer_dot_product : require + +#define NUM_GROUPS_PER_WG ${NUM_GROUPS_PER_WG} +#define NUM_WORKERS_PER_GROUP ${NUM_WORKERS_PER_GROUP} + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_packed_int8_input", "int", OUTPUT_STORAGE, is_scalar_array=True)} +${layout_declare_tensor(B, "w", "t_int8_input_sums", "int", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_input", DTYPE, INPUT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", DTYPE, "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8" if ZP_DTYPE_MODE == "zpint8" else DTYPE, "texture3d")} + +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} + +shared ivec4 shared_sums[NUM_GROUPS_PER_WG][NUM_WORKERS_PER_GROUP]; + +#define TILE_M4 1 +#define TILE_K4 1 + +#define TILE_M 4 + +// The shared helper's write_block() assumes an ivec4-typed output resource; +// this shader's output is a scalar int array, so opt out and write our own. +#define SKIP_BLOCK_WRITE_HELPERS +#include "gemm/tile_utils/linear_int8_input_block.glslh" + +// Scatters the 4 rows of a quantized block into row-major slots. Unlike the +// 4h4w layout, kPackedInt8_4W has NO row padding (outer_dim_align == 1), so +// the tail rows of a workgroup whose m4 block runs past M must be dropped -- +// hence the explicit bound check that write_block() does not need. +void write_block_rowmajor( + const Int8InputBlock block, + const int k4, + const int m, + const int K4, + const int M) { + for (int row = 0; row < 4; ++row) { + const int m_row = m + row; + if (m_row < M) { + t_packed_int8_input[m_row * K4 + k4] = block.data[row]; + } + } +} +#include "gemm/tile_utils/linear_int8_input_scales_zps_load.glslh" +#include "gemm/tile_utils/linear_fp_input_tile_load.glslh" + +void main() { + const int group_idx = int(gl_GlobalInvocationID.x); + const int m4 = int(gl_GlobalInvocationID.y); + + const int worker_id = int(gl_LocalInvocationID.z); + const int group_offset = int(gl_LocalInvocationID.x); + + const int K = input_sizes.x; + const int M = input_sizes.y; + + // K4 and M4 represent the number of blocks in each dimension. + const int K4 = div_up_4(K); + const int M4 = div_up_4(M); + + const int num_groups = K4 / K4_per_group; + + if (group_idx >= num_groups || m4 >= M4) { + return; + } + + const int start_k4 = group_idx * K4_per_group + worker_id; + const int end_k4 = (group_idx + 1) * K4_per_group; + + Int8InputScales input_scales; + Int8InputZeroPoints input_zps; + load_int8_input_scales_and_zps(input_scales, input_zps, m4); + + // row of the input tensor to start loading from + const int m = mul_4(m4); + + FPInputTile in_tile; + Int8InputBlock packed_block; + + ivec4 local_sum = ivec4(0, 0, 0, 0); + const int packed_ones = 0x01010101; + + for (int k4 = start_k4; k4 < end_k4; k4 += NUM_WORKERS_PER_GROUP) { + load_input_tile_no_checks(in_tile, k4, m, K4, M); + quantize_and_pack(packed_block, in_tile, input_scales, input_zps); + + // Sum the quantized values in the block + [[unroll]] for (int m = 0; m < TILE_M; m++) { + local_sum[m] += dotPacked4x8AccSatEXT( + packed_block.data[m], packed_ones, local_sum[m]); + } + write_block_rowmajor(packed_block, k4, m, K4, M); + } + + shared_sums[group_offset][worker_id] = local_sum; + + memoryBarrierShared(); + barrier(); + + // Tree reduction to compute the overall result + for (int i = NUM_WORKERS_PER_GROUP / 2; i > 0; i >>= 1) { + if (worker_id < i) { + shared_sums[group_offset][worker_id] = + shared_sums[group_offset][worker_id] + + shared_sums[group_offset][worker_id + i]; + } + memoryBarrierShared(); + barrier(); + } + + if (worker_id == 0) { + t_int8_input_sums[group_idx * M4 + m4] = shared_sums[group_offset][0]; + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/quantize_and_pack_4w_with_group_sums.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/quantize_and_pack_4w_with_group_sums.yaml new file mode 100644 index 00000000000..dfa6f7313a5 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/quantize_and_pack_4w_with_group_sums.yaml @@ -0,0 +1,32 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Row-major (utils::kPackedInt8_4W) counterpart of +# quantize_and_pack_4h4w_with_group_sums: same quantization, per-group sum +# reduction and workgroup shapes (_o2w32 / _o4w16); only the output layout +# differs. Consumed by linear_dq8ca_q4gsw_wmma_rowmajor, whose cooperative +# matrix A staging needs row-major int8 activations. Buffer output only. + +quantize_and_pack_4w_with_group_sums: + parameter_names_with_default_values: + DTYPE: float + OUTPUT_STORAGE: buffer + INPUT_STORAGE: texture3d + NUM_GROUPS_PER_WG: 2 + NUM_WORKERS_PER_GROUP: 32 + ZP_DTYPE_MODE: zpint8 + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + ZP_DTYPE_MODE: + - VALUE: zpint8 + - VALUE: zpinherit + shader_variants: + - NAME: quantize_and_pack_4w_with_group_sums_o2w32_buffer_texture3d + - NAME: quantize_and_pack_4w_with_group_sums_o4w16_buffer_texture3d + NUM_GROUPS_PER_WG: 4 + NUM_WORKERS_PER_GROUP: 16 diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma.glsl new file mode 100644 index 00000000000..84c1cb73fd6 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma.glsl @@ -0,0 +1,97 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * fp16 q4gsw (4-bit group-symmetric weight) cooperative-matrix linear with + * per-variant tile geometry, used by the architecture-family path + * (impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp). + * + * This file is the template header: bindings, spec constants and the tile + * geometry. The kernel body is linear_q4gsw_wmma_body.glslh. + * + * Loop structure: double-buffered, prefetch-peeled. The prologue prefetches + * chunk 0; each iteration does barrier -> prefetch(next) -> MMA(cur) -> + * store(next). + * + * Feature defines (yaml parameters, all default off): + * B_COLMAJOR B staged N-major in shared memory, ColumnMajor B coopMatLoad + * CSH_BAND texture3d drain one SG_GRID_Y band at a time + * FRAG_LAYOUT fragment-contiguous shared memory, no padding + * IMG_A storage-image loads for A (texture3d IO only) + * MMA_M = 8 for devices whose fp16 cooperative matrix is 8x16x16 only + * + * Preconditions (not checked in the shader), enforced by the selection in + * CoopmatFamily.cpp: M % WG_TILE_M == 0, N % WG_TILE_N == 0, + * K % WG_TILE_K == 0, group_size % WG_TILE_K == 0, no bias, batch 1. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +$if IMG_A and IO_STORAGE == "texture3d": + ${layout_declare_image(B, "r", "t_input", "half")} +$else: + ${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +#define ACC_T float16_t + +$if B_COLMAJOR: + #define B_COLMAJOR + +$if CSH_BAND: + #define CSH_BAND + +$if FRAG_LAYOUT: + #define FRAG_LAYOUT + +$if IMG_A and IO_STORAGE == "texture3d": + #define IMG_A + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; + +#include "linear_q4gsw_wmma_body.glslh" diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma.yaml new file mode 100644 index 00000000000..75ca4ee943c --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma.yaml @@ -0,0 +1,60 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Variants of the architecture-family q4gsw cooperative-matrix linear. Each +# one is referenced by a row in impl/gemm/linear/qw/CoopmatFamily.cpp, which +# also carries its tile geometry; the two must stay in sync. +# +# Name: txkgs +# followed by m8 (MMA_M = 8), f (FRAG_LAYOUT), i (IMG_A), b (CSH_BAND) or +# bt (B_COLMAJOR). + +linear_q4gsw_wmma: + parameter_names_with_default_values: + PRECISION: highp + IO_STORAGE: buffer + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + FRAG_LAYOUT: false + IMG_A: false + CSH_BAND: false + B_COLMAJOR: false + shader_variants: + # Intel Xe2: subgroup 16, MMA 8x16x16 (the only fp16 shape Xe2 exposes), + # fragment-contiguous shared memory and image loads for A. + - NAME: linear_q4gsw_wmma_t128x128k16g82s16m8flib_texture3d_texture2d_half + IO_STORAGE: texture3d + MMA_M: 8 + SG_GRID_X: 8 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 16 + FRAG_LAYOUT: true + IMG_A: true + CSH_BAND: true + - NAME: linear_q4gsw_wmma_t128x128k16g44s16m8fli_buffer_texture2d_half + MMA_M: 8 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 16 + FRAG_LAYOUT: true + IMG_A: true + # NVIDIA: subgroup 32, MMA 16x16x16 with fp16 accumulation (K <= 8192). + - NAME: linear_q4gsw_wmma_t256x128k16g42s32bt_texture3d_texture2d_half + IO_STORAGE: texture3d + WG_TILE_M: 256 + SG_GRID_X: 4 + SG_GRID_Y: 2 + B_COLMAJOR: true + - NAME: linear_q4gsw_wmma_t128x128k16g22s32_texture3d_texture2d_half + IO_STORAGE: texture3d diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_body.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_body.glslh new file mode 100644 index 00000000000..2879f226cd0 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_body.glslh @@ -0,0 +1,592 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Body of linear_q4gsw_wmma. Included by its template header +// (linear_q4gsw_wmma.glsl), which defines the bindings, the spec constants, +// ACC_T, the optional IO_TEXTURE / FRAG_LAYOUT / IMG_A / B_COLMAJOR / CSH_BAND +// macros and the tile constants MMA_{M,N,K}, WG_TILE_{M,N,K}, SG_GRID_{X,Y}, +// SUBGROUP_SIZE. Must not contain template syntax: .glslh files are not +// expanded per variant. + +#ifndef LINEAR_Q4GSW_WMMA_BODY_GLSLH +#define LINEAR_Q4GSW_WMMA_BODY_GLSLH + +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +// fp16: 8 elements per uvec4 (128-bit) +const uint FP16_PER_VEC4 = 8; +#ifdef FRAG_LAYOUT +// Fragment-contiguous shared memory: A as [k16 block][m][16 fp16] and B as +// [n16 block][k][16 fp16], so every MMA_M x 16 A fragment and 16 x MMA_N B +// fragment is one contiguous run with a 32-byte row stride. The row-padded +// layout below scatters a fragment across strided 32-byte rows, which some +// compilers lower to many narrow shared-memory messages; where shared memory +// has no bank conflicts the padding buys nothing. Requires +// MMA_K == MMA_N == 16. +const uint FRAG_ROW_VEC4 = 16u / FP16_PER_VEC4; +const uint A_STRIDE_VEC4 = FRAG_ROW_VEC4; +const uint B_STRIDE_VEC4 = FRAG_ROW_VEC4; +const uint ASH_SLICE = WG_TILE_M * (WG_TILE_K / FP16_PER_VEC4); +const uint BSH_SLICE = WG_TILE_K * (WG_TILE_N / FP16_PER_VEC4); +// (row, uvec4 column) -> slice-relative index. +#define A_SH_IDX(m, c) ((((c) / FRAG_ROW_VEC4) * WG_TILE_M + (m)) * FRAG_ROW_VEC4 + ((c) % FRAG_ROW_VEC4)) +#define B_SH_IDX(k, c) ((((c) / FRAG_ROW_VEC4) * WG_TILE_K + (k)) * FRAG_ROW_VEC4 + ((c) % FRAG_ROW_VEC4)) +#else +const uint A_STRIDE_VEC4 = (WG_TILE_K + FP16_PER_VEC4) / FP16_PER_VEC4; +const uint B_STRIDE_VEC4 = (WG_TILE_N + FP16_PER_VEC4) / FP16_PER_VEC4; + +// One ping-pong slice of each shared-memory buffer (in uvec4 units). +const uint ASH_SLICE = WG_TILE_M * A_STRIDE_VEC4; +const uint BSH_SLICE = WG_TILE_K * B_STRIDE_VEC4; +// No outer parentheses: every use is `base + X_SH_IDX(...)` or a subscript, +// so `base + row * STRIDE + col` keeps its natural association. +#define A_SH_IDX(m, c) (m) * A_STRIDE_VEC4 + (c) +#define B_SH_IDX(k, c) (k) * B_STRIDE_VEC4 + (c) +#endif + +// Double-buffered shared memory. +shared uvec4 Ash[2 * ASH_SLICE]; +#ifdef B_COLMAJOR +// B_COLMAJOR: on devices whose MMA B operand is one N column of MMA_K +// consecutive K values per lane, coopMatLoad from the row-major Bsh compiles +// to per-element gathers. Stage B N-major (BshT[n][k], K contiguous, padded +// like Ash) and load it ColumnMajor instead. +// Units: uvec2 = 4 fp16. +#if defined(FRAG_LAYOUT) +#error "B_COLMAJOR is not combined with FRAG_LAYOUT" +#endif +const uint BT_STRIDE_H = WG_TILE_K + 8u; +const uint BT_STRIDE_V2 = BT_STRIDE_H / 4u; +const uint BT_SLICE = WG_TILE_N * BT_STRIDE_V2; +shared uvec2 BshT[2 * BT_SLICE]; +#else +shared uvec4 Bsh[2 * BSH_SLICE]; +#endif +// Shared-memory stores of one staged uvec4. +#define ASH_STORE(idx, v) Ash[idx] = v +#define BSH_STORE(idx, v) Bsh[idx] = v + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue: SG_GRID_Y bands of MMA_M rows, +// each WG_TILE_N wide, row-major. A full WG_TILE_M x WG_TILE_N buffer would +// cost SG_GRID_Y/MMAS_PER_SG_M x more shared memory. Not aliased +// with Ash/Bsh: GLSL has no shared unions and coopMatStore needs a +// float16_t-typed array. +#if defined(CSH_BAND) +// CSH_BAND: drain one SG_GRID_Y band at a time, so Csh is MMA_M rows x +// WG_TILE_N. +const uint CSH_ROWS = MMA_M; +#else +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +#endif +shared float16_t Csh[CSH_ROWS * WG_TILE_N]; +#endif + +// Staging thread maps: each thread covers one uvec4 (8 fp16) per pass. +const uint INVS_PER_ROW_A = WG_TILE_K / FP16_PER_VEC4; +const uint A_ROWS_PER_PASS = WG_SIZE / INVS_PER_ROW_A; +const uint A_PASSES = WG_TILE_M / A_ROWS_PER_PASS; +const uint INVS_PER_ROW_B = WG_TILE_N / FP16_PER_VEC4; +const uint B_ROWS_PER_PASS = WG_SIZE / INVS_PER_ROW_B; +const uint B_PASSES = WG_TILE_K / B_ROWS_PER_PASS; +#ifdef B_COLMAJOR +// One item = half of one packed-weight texel (4 K x 8 N nibbles): 4 K x 4 N, +// half 0 = N 0..3 (even nibbles), half 1 = N 4..7 (odd nibbles). K varies +// fastest across lanes: each (n8, half) group of BT_KQ lanes writes one +// contiguous run per column, and the groups land on disjoint shared-memory +// banks. +const uint BT_KQ = WG_TILE_K / 4u; +const uint BT_N8 = WG_TILE_N / 8u; +const uint BT_ITEMS = BT_KQ * BT_N8 * 2u; +const uint BT_PASSES = (BT_ITEMS + WG_SIZE - 1u) / WG_SIZE; +#endif + +// Accumulator coopmats (MMAS_PER_SG_M x MMAS_PER_SG_N per subgroup). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Dequant one packed INT4 block column-pair into 8 scaled fp16 weights +// (one Bsh uvec4). col_lo/col_hi select the K row within the block. +#define WFETCH(coord) texelFetch(t_packed_weight, coord, 0) + +uvec4 dequant_block( + const ivec4 wb, + const uint col_lo, + const uint col_hi, + const f16vec4 s0, + const f16vec4 s1) { + const f16vec4 v0 = f16vec4(((wb >> int(4u * col_lo)) & 0xF) - 8) * s0; + const f16vec4 v1 = f16vec4(((wb >> int(4u * col_hi)) & 0xF) - 8) * s1; + return uvec4( + packFloat2x16(v0.xy), packFloat2x16(v0.zw), + packFloat2x16(v1.xy), packFloat2x16(v1.zw)); +} + +#ifdef B_COLMAJOR +// Store one item: 4 N columns (x..w of wb, nibble parity hf) x 4 K rows +// (nibbles 2*kk + hf), each column's 4 K values as one uvec2. +void store_bt_item( + const uint base, + const ivec4 wb, + const uint n_local, + const uint k_local, + const uint hf, + const f16vec4 sc) { + [[unroll]] for (uint j = 0; j < 4u; ++j) { + const int w = wb[j]; + const f16vec4 v = f16vec4( + ((w >> int(4u * (0u + hf))) & 0xF) - 8, + ((w >> int(4u * (2u + hf))) & 0xF) - 8, + ((w >> int(4u * (4u + hf))) & 0xF) - 8, + ((w >> int(4u * (6u + hf))) & 0xF) - 8) * sc[j]; + BshT[base + (n_local + j) * BT_STRIDE_V2 + k_local / 4u] = + uvec2(packFloat2x16(v.xy), packFloat2x16(v.zw)); + } +} +#define BT_ACTIVE(p) (BT_ITEMS % WG_SIZE == 0u || (p) * WG_SIZE + gl_LocalInvocationID.x < BT_ITEMS) +#endif + +// Fetch 8 consecutive fp16 activations of `row` starting at half-vec4 index +// `k_hv4`, packed into one uvec4. Both spellings address the same bytes: a +// width-packed texture3d holds elements [4x, 4x+3] of row m at texel (x, m, 0), +// so buffer index (row * K4 + x) and texel coord (x, row, 0) are one address in +// two notations. +uvec4 load_a_vec4(const uint row_in, const uint k_hv4, const uint K4) { +#ifdef IO_TEXTURE + // Rows of a partial last M tile read row M - 1; their results are not + // stored (each output row depends only on its own A row). + const uint row = min(row_in, uint(input_sizes.y) - 1u); + // Narrow to f16vec4 at the fetch: a half sampler is typed to return vec4 + // (fp32); narrowing immediately lets the compiler fold the conversion into + // a 16-bit image load. Lossless -- the source is rgba16f. +#ifdef IMG_A + // Storage-image load instead of a sampled fetch. + const f16vec4 v0 = f16vec4(imageLoad(t_input, ivec3(k_hv4, row, 0))); + const f16vec4 v1 = f16vec4(imageLoad(t_input, ivec3(k_hv4 + 1u, row, 0))); +#else + const f16vec4 v0 = f16vec4(texelFetch(t_input, ivec3(k_hv4, row, 0), 0)); + const f16vec4 v1 = f16vec4(texelFetch(t_input, ivec3(k_hv4 + 1u, row, 0), 0)); +#endif + return uvec4( + packFloat2x16(v0.xy), packFloat2x16(v0.zw), + packFloat2x16(v1.xy), packFloat2x16(v1.zw)); +#else + // Buffer IO requires M to be a multiple of WG_TILE_M (CoopmatFamily.cpp). + const uint row = row_in; + const f16vec4 v0 = t_input[row * K4 + k_hv4]; + const f16vec4 v1 = t_input[row * K4 + k_hv4 + 1u]; + return uvec4( + packFloat2x16(v0.xy), packFloat2x16(v0.zw), + packFloat2x16(v1.xy), packFloat2x16(v1.zw)); +#endif +} + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint K4 = (K + 3u) / 4u; + const uint N4 = (uint(output_sizes.x) + 3u) / 4u; + + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; + const uint num_chunks = uint(num_groups_arg) * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + } + } + + const uint a_col = gl_LocalInvocationID.x % INVS_PER_ROW_A; + const uint a_row_offset = gl_LocalInvocationID.x / INVS_PER_ROW_A; + const uint b_col = gl_LocalInvocationID.x % INVS_PER_ROW_B; + const uint b_row_offset = gl_LocalInvocationID.x / INVS_PER_ROW_B; + + const uint n8_blk = (tile_n_start + b_col * 8u) >> 3u; + const uint col_lo = 2u * (b_row_offset & 3u); + const uint col_hi = col_lo + 1u; + + const uint sc_n4 = (tile_n_start + b_col * 8u) >> 2u; + uint cached_group = 0xFFFFFFFFu; + f16vec4 sc0; + f16vec4 sc1; + + uvec4 temp_A[A_PASSES]; +#ifdef B_COLMAJOR + ivec4 temp_B[BT_PASSES]; + f16vec4 sc_bt[BT_PASSES]; + uint bt_kq[BT_PASSES], bt_n8[BT_PASSES], bt_hf[BT_PASSES]; + [[unroll]] for (uint p = 0; p < BT_PASSES; ++p) { + const uint item = p * WG_SIZE + gl_LocalInvocationID.x; + bt_kq[p] = item % BT_KQ; + bt_hf[p] = (item / BT_KQ) & 1u; + bt_n8[p] = (item / BT_KQ) >> 1u; + } +#else + ivec4 temp_B[B_PASSES]; +#endif + + // ========================================================= + // PROLOGUE: prefetch chunk 0 into temp registers, then store to + // slice 0. No barrier here -- the main loop's first iteration barriers + // before reading slice 0. + // ========================================================= + { + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + const uint row = tile_m_start + p * A_ROWS_PER_PASS + a_row_offset; + const uint k_hv4 = (a_col * FP16_PER_VEC4) / 4u; + temp_A[p] = load_a_vec4(row, k_hv4, K4); + } + cached_group = 0u; +#ifdef B_COLMAJOR + [[unroll]] for (uint p = 0; p < BT_PASSES; ++p) { + if (BT_ACTIVE(p)) { + const uint n8g = (tile_n_start >> 3u) + bt_n8[p]; + sc_bt[p] = t_weight_scales[n8g * 2u + bt_hf[p]]; + temp_B[p] = WFETCH(ivec2(bt_kq[p], n8g)); + } + } +#else + sc0 = t_weight_scales[sc_n4]; + sc1 = t_weight_scales[sc_n4 + 1u]; + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + const uint k_row = p * B_ROWS_PER_PASS + b_row_offset; + ivec4 wblock; + wblock = WFETCH(ivec2(k_row >> 2u, n8_blk)); + temp_B[p] = wblock; + } +#endif + } + { + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + ASH_STORE(A_SH_IDX(p * A_ROWS_PER_PASS + a_row_offset, a_col), temp_A[p]); + } +#ifdef B_COLMAJOR + [[unroll]] for (uint p = 0; p < BT_PASSES; ++p) { + if (BT_ACTIVE(p)) { + store_bt_item(0u, temp_B[p], bt_n8[p] * 8u + 4u * bt_hf[p], bt_kq[p] * 4u, bt_hf[p], sc_bt[p]); + } + } +#else + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + BSH_STORE(B_SH_IDX(p * B_ROWS_PER_PASS + b_row_offset, b_col), + dequant_block(temp_B[p], col_lo, col_hi, sc0, sc1)); + } +#endif + } + + // ========================================================= + // MAIN LOOP: one barrier per iteration, unconditional body, loop + // peeled (bound excludes the last chunk). Iteration `chunk` does: + // 1. barrier — slice (chunk%2) fully written + // 2. prefetch — chunk+1 from global into temp (in flight during math) + // 3. MMA math — on slice (chunk%2) + // 4. store — temp (chunk+1, dequantized) into slice ((chunk+1)%2) + // ========================================================= + uint chunk; + for (chunk = 0; chunk + 1u < num_chunks; ++chunk) { + const uint cur_base_A = (chunk % 2u) * ASH_SLICE; +#ifdef B_COLMAJOR + const uint cur_base_B = (chunk % 2u) * BT_SLICE; + const uint nxt_base_B = ((chunk + 1u) % 2u) * BT_SLICE; +#else + const uint cur_base_B = (chunk % 2u) * BSH_SLICE; +#endif + const uint nxt_base_A = ((chunk + 1u) % 2u) * ASH_SLICE; +#if !defined(B_COLMAJOR) + const uint nxt_base_B = ((chunk + 1u) % 2u) * BSH_SLICE; +#endif + + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + + // --- prefetch chunk+1 -> temp --- + { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + const uint row = tile_m_start + p * A_ROWS_PER_PASS + a_row_offset; + const uint k_hv4 = (chunkK_nxt + a_col * FP16_PER_VEC4) / 4u; + temp_A[p] = load_a_vec4(row, k_hv4, K4); + } +#ifdef B_COLMAJOR + const uint group_nxt = (chunk + 1u) / CHUNKS_PER_GROUP; + const bool new_group = group_nxt != cached_group; + cached_group = group_nxt; + [[unroll]] for (uint p = 0; p < BT_PASSES; ++p) { + if (BT_ACTIVE(p)) { + const uint n8g = (tile_n_start >> 3u) + bt_n8[p]; + temp_B[p] = WFETCH(ivec2(chunkK_nxt / 4u + bt_kq[p], n8g)); + if (new_group) { + sc_bt[p] = t_weight_scales[group_nxt * N4 + n8g * 2u + bt_hf[p]]; + } + } + } +#else + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + const uint k_row = chunkK_nxt + p * B_ROWS_PER_PASS + b_row_offset; + temp_B[p] = WFETCH(ivec2(k_row >> 2u, n8_blk)); + } + const uint group_nxt = (chunk + 1u) / CHUNKS_PER_GROUP; + if (group_nxt != cached_group) { + cached_group = group_nxt; + sc0 = t_weight_scales[group_nxt * N4 + sc_n4]; + sc1 = t_weight_scales[group_nxt * N4 + sc_n4 + 1u]; + } +#endif + } + + // --- MMA math on the cur slice --- + [[unroll]] for (uint k = 0; k < WG_TILE_K / MMA_K; ++k) { + const uint k_start = MMA_K * k; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash, + cur_base_A + A_SH_IDX(row_a, k_start / FP16_PER_VEC4), + A_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef B_COLMAJOR + const uint n_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, BshT, + cur_base_B + n_b * BT_STRIDE_V2 + k_start / 4u, + BT_STRIDE_V2, + gl_CooperativeMatrixLayoutColumnMajor); +#else + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j) / FP16_PER_VEC4; + coopMatLoad( + matB, Bsh, + cur_base_B + B_SH_IDX(k_start, col_b), + B_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); +#endif // B_COLMAJOR + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = coopMatMulAdd(matA[i], matB, result[i][j]); + } + } + } + + + // --- store temp (chunk+1) -> nxt slice, dequantizing B --- + { + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + ASH_STORE(nxt_base_A + A_SH_IDX(p * A_ROWS_PER_PASS + a_row_offset, a_col), + temp_A[p]); + } +#ifdef B_COLMAJOR + [[unroll]] for (uint p = 0; p < BT_PASSES; ++p) { + if (BT_ACTIVE(p)) { + store_bt_item(nxt_base_B, temp_B[p], bt_n8[p] * 8u + 4u * bt_hf[p], bt_kq[p] * 4u, bt_hf[p], sc_bt[p]); + } + } +#else + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + BSH_STORE(nxt_base_B + B_SH_IDX(p * B_ROWS_PER_PASS + b_row_offset, b_col), + dequant_block(temp_B[p], col_lo, col_hi, sc0, sc1)); + } +#endif + } + } + + // --- epilogue: barrier, then MMA on the last chunk (loop peeled) --- + { + const uint cur_base_A = (chunk % 2u) * ASH_SLICE; +#ifdef B_COLMAJOR + const uint cur_base_B = (chunk % 2u) * BT_SLICE; +#else + const uint cur_base_B = (chunk % 2u) * BSH_SLICE; +#endif + + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + + [[unroll]] for (uint k = 0; k < WG_TILE_K / MMA_K; ++k) { + const uint k_start = MMA_K * k; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash, + cur_base_A + A_SH_IDX(row_a, k_start / FP16_PER_VEC4), + A_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef B_COLMAJOR + const uint n_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, BshT, + cur_base_B + n_b * BT_STRIDE_V2 + k_start / 4u, + BT_STRIDE_V2, + gl_CooperativeMatrixLayoutColumnMajor); +#else + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j) / FP16_PER_VEC4; + coopMatLoad( + matB, Bsh, + cur_base_B + B_SH_IDX(k_start, col_b), + B_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); +#endif // B_COLMAJOR + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = coopMatMulAdd(matA[i], matB, result[i][j]); + } + } + } + } + + + // --- Store result tile --- +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh at once, so the SG_GRID_Y bands it holds are disjoint global row + // ranges; the whole workgroup then imageStores them. The band-to-global-row + // map below reproduces the buffer path's gi exactly: + // buffer: gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i) + // = tile_m_start + warpInTile.y * SG_TILE_M + i * MMA_M + // and lr / MMA_M is the writing subgroup's warpInTile.y. + // + // PORTABILITY NOTE: [[unroll]] is only a hint and glslc does not honor it + // here -- the barrier() in the loop body keeps the loop rolled, so + // result[i][j] is dynamically indexed. Coopmat arrays are opaque per-lane + // storage and dynamic indexing is a construct shader compilers have broken + // before, so check this first if the texture variants miscompile on a new + // device. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; +#if defined(CSH_BAND) + // Band b = the subgroups with warpInTile.y == b; per (i, b): barrier, that + // band's subgroups store accumulator row-block i into Csh, barrier, the + // whole workgroup imageStores the MMA_M rows. Same global row map as below. + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + for (uint b = 0; b < SG_GRID_Y; ++b) { + memoryBarrierShared(); + barrier(); + if (warpInTile.y == b) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + coopMatStore( + coopmat(result[i][j]), Csh, + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + } + memoryBarrierShared(); + barrier(); + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = tile_m_start + b * SG_TILE_M + i * MMA_M + lr; + const uint base = lr * WG_TILE_N + lc4 * 4u; + // The last M tile may be partial: rows at or past M are computed but + // not stored. + if (m >= uint(output_sizes.y)) { + continue; + } + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh[base]), + float(Csh[base + 1u]), + float(Csh[base + 2u]), + float(Csh[base + 3u]))); + } + } + } +#else + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh against the previous iteration's readers. Inert on i == 0 but + // must stay unconditional to remain workgroup-uniform. + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + coopMatStore( + coopmat(result[i][j]), Csh, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + // The last M tile may be partial: rows at or past M are computed but + // not stored. + if (m >= uint(output_sizes.y)) { + continue; + } + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh[base]), + float(Csh[base + 1u]), + float(Csh[base + 2u]), + float(Csh[base + 3u]))); + } + } +#endif // CSH_BAND +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} + +#endif // LINEAR_Q4GSW_WMMA_BODY_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx.glsl new file mode 100644 index 00000000000..ecc8fde5bd3 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx.glsl @@ -0,0 +1,83 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * fp16 q4gsw cooperative-matrix linear with texel-wise weight staging and an + * optional fp32 accumulator, used by the architecture-family path + * (impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp) for long-K shapes where + * fp16 accumulation loses accuracy. Template header only; the body is + * linear_q4gsw_wmma_bx_body.glslh. + * + * Feature defines (yaml parameters, all default off): + * ACC_FP32 fp32 accumulator + * CSH_IN_ASH texture3d drain staged in the A shared-memory tile, which is + * dead after the last MMA + * + * Preconditions as for linear_q4gsw_wmma.glsl. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +$if ACC_FP32: + #define ACC_FP32 + +#ifdef ACC_FP32 +#define ACC_T float +#else +#define ACC_T float16_t +#endif + +$if CSH_IN_ASH: + #define CSH_IN_ASH + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; + +#include "linear_q4gsw_wmma_bx_body.glslh" diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx.yaml new file mode 100644 index 00000000000..7951b48ff16 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx.yaml @@ -0,0 +1,36 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Variants of the texel-wise-staging q4gsw cooperative-matrix linear. Each one +# is referenced by a row in impl/gemm/linear/qw/CoopmatFamily.cpp, which also +# carries its tile geometry; the two must stay in sync. Name as in +# linear_q4gsw_wmma.yaml, plus f32 (ACC_FP32) and c (CSH_IN_ASH). + +linear_q4gsw_wmma_bx: + parameter_names_with_default_values: + PRECISION: highp + IO_STORAGE: buffer + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + ACC_FP32: false + CSH_IN_ASH: false + shader_variants: + # NVIDIA, K > 8192: fp32 accumulation keeps the long-K error in bounds. + - NAME: linear_q4gsw_wmma_bx_t128x128k32g42s32f32c_texture3d_texture2d_half + IO_STORAGE: texture3d + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + ACC_FP32: true + CSH_IN_ASH: true diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx_body.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx_body.glslh new file mode 100644 index 00000000000..adaa9790beb --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_wmma_bx_body.glslh @@ -0,0 +1,426 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Body of linear_q4gsw_wmma_bx: linear_q4gsw_wmma_body.glslh with texel-wise +// B staging (one packed-weight texel feeds every B pass of a thread). +// Included by its template header (linear_q4gsw_wmma_bx.glsl), which defines +// the bindings, the spec constants, ACC_T, the optional IO_TEXTURE / +// CSH_IN_ASH macros and the tile constants MMA_{M,N,K}, WG_TILE_{M,N,K}, +// SG_GRID_{X,Y}, SUBGROUP_SIZE. Must not contain template syntax. + +#ifndef LINEAR_Q4GSW_WMMA_BX_BODY_GLSLH +#define LINEAR_Q4GSW_WMMA_BX_BODY_GLSLH + +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +// fp16: 8 elements per uvec4 (128-bit) +const uint FP16_PER_VEC4 = 8; +const uint A_STRIDE_VEC4 = (WG_TILE_K + FP16_PER_VEC4) / FP16_PER_VEC4; +const uint B_STRIDE_VEC4 = (WG_TILE_N + FP16_PER_VEC4) / FP16_PER_VEC4; + +// One ping-pong slice of each shared-memory buffer (in uvec4 units). +const uint ASH_SLICE = WG_TILE_M * A_STRIDE_VEC4; +const uint BSH_SLICE = WG_TILE_K * B_STRIDE_VEC4; +// No outer parentheses: every use is `base + X_SH_IDX(...)` or a subscript, +// so `base + row * STRIDE + col` keeps its natural association. +#define A_SH_IDX(m, c) (m) * A_STRIDE_VEC4 + (c) +#define B_SH_IDX(k, c) (k) * B_STRIDE_VEC4 + (c) + +// Double-buffered shared memory. +shared uvec4 Ash[2 * ASH_SLICE]; +shared uvec4 Bsh[2 * BSH_SLICE]; +// Shared-memory stores of one staged uvec4. +#define ASH_STORE(idx, v) Ash[idx] = v +#define BSH_STORE(idx, v) Bsh[idx] = v + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue: SG_GRID_Y bands of MMA_M rows, +// each WG_TILE_N wide, row-major. A full WG_TILE_M x WG_TILE_N buffer would +// cost SG_GRID_Y/MMAS_PER_SG_M x more shared memory. Not aliased +// with Ash/Bsh: GLSL has no shared unions and coopMatStore needs a +// float16_t-typed array. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +// CSH_IN_ASH stages the drain band in Ash instead: Ash is dead after the +// last MMA and every drain iteration starts with a barrier. +#ifndef CSH_IN_ASH +shared float16_t Csh[CSH_ROWS * WG_TILE_N]; +#endif +#endif + +// Staging thread maps: each thread covers one uvec4 (8 fp16) per pass. +const uint INVS_PER_ROW_A = WG_TILE_K / FP16_PER_VEC4; +const uint A_ROWS_PER_PASS = WG_SIZE / INVS_PER_ROW_A; +const uint A_PASSES = WG_TILE_M / A_ROWS_PER_PASS; +const uint INVS_PER_ROW_B = WG_TILE_N / FP16_PER_VEC4; +const uint B_ROWS_PER_PASS = WG_SIZE / INVS_PER_ROW_B; +const uint B_PASSES = WG_TILE_K / B_ROWS_PER_PASS; +// Texel-wise B staging: one packed-weight texel (4 K rows x 8 N) feeds all +// B_PASSES rows of a thread; B_RP threads share a texel when B_PASSES < 4. +// Row of pass p within the chunk, and its first nibble index in the texel. +const uint B_RP = 4u / B_PASSES; +#define B_TROW(p) ((b_row_offset / B_RP) * 4u + (b_row_offset % B_RP) * B_PASSES + (p)) +#define B_TCOL(p) (2u * ((b_row_offset % B_RP) * B_PASSES + (p))) + +// Accumulator coopmats (MMAS_PER_SG_M x MMAS_PER_SG_N per subgroup). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Dequant one packed INT4 block column-pair into 8 scaled fp16 weights +// (one Bsh uvec4). col_lo/col_hi select the K row within the block. +#define WFETCH(coord) texelFetch(t_packed_weight, coord, 0) + +uvec4 dequant_block( + const ivec4 wb, + const uint col_lo, + const uint col_hi, + const f16vec4 s0, + const f16vec4 s1) { + const f16vec4 v0 = f16vec4(((wb >> int(4u * col_lo)) & 0xF) - 8) * s0; + const f16vec4 v1 = f16vec4(((wb >> int(4u * col_hi)) & 0xF) - 8) * s1; + return uvec4( + packFloat2x16(v0.xy), packFloat2x16(v0.zw), + packFloat2x16(v1.xy), packFloat2x16(v1.zw)); +} + + +// Fetch 8 consecutive fp16 activations of `row` starting at half-vec4 index +// `k_hv4`, packed into one uvec4. Both spellings address the same bytes: a +// width-packed texture3d holds elements [4x, 4x+3] of row m at texel (x, m, 0), +// so buffer index (row * K4 + x) and texel coord (x, row, 0) are one address in +// two notations. +uvec4 load_a_vec4(const uint row_in, const uint k_hv4, const uint K4) { +#ifdef IO_TEXTURE + // Rows of a partial last M tile read row M - 1; their results are not + // stored (each output row depends only on its own A row). + const uint row = min(row_in, uint(input_sizes.y) - 1u); + // Narrow to f16vec4 at the fetch: a half sampler is typed to return vec4 + // (fp32); narrowing immediately lets the compiler fold the conversion into + // a 16-bit image load. Lossless -- the source is rgba16f. + const f16vec4 v0 = f16vec4(texelFetch(t_input, ivec3(k_hv4, row, 0), 0)); + const f16vec4 v1 = f16vec4(texelFetch(t_input, ivec3(k_hv4 + 1u, row, 0), 0)); + return uvec4( + packFloat2x16(v0.xy), packFloat2x16(v0.zw), + packFloat2x16(v1.xy), packFloat2x16(v1.zw)); +#else + // Buffer IO requires M to be a multiple of WG_TILE_M (CoopmatFamily.cpp). + const uint row = row_in; + const f16vec4 v0 = t_input[row * K4 + k_hv4]; + const f16vec4 v1 = t_input[row * K4 + k_hv4 + 1u]; + return uvec4( + packFloat2x16(v0.xy), packFloat2x16(v0.zw), + packFloat2x16(v1.xy), packFloat2x16(v1.zw)); +#endif +} + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint K4 = (K + 3u) / 4u; + const uint N4 = (uint(output_sizes.x) + 3u) / 4u; + + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; + const uint num_chunks = uint(num_groups_arg) * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + } + } + + const uint a_col = gl_LocalInvocationID.x % INVS_PER_ROW_A; + const uint a_row_offset = gl_LocalInvocationID.x / INVS_PER_ROW_A; + const uint b_col = gl_LocalInvocationID.x % INVS_PER_ROW_B; + const uint b_row_offset = gl_LocalInvocationID.x / INVS_PER_ROW_B; + + const uint n8_blk = (tile_n_start + b_col * 8u) >> 3u; + const uint col_lo = 2u * (b_row_offset & 3u); + const uint col_hi = col_lo + 1u; + + const uint sc_n4 = (tile_n_start + b_col * 8u) >> 2u; + uint cached_group = 0xFFFFFFFFu; + f16vec4 sc0; + f16vec4 sc1; + + uvec4 temp_A[A_PASSES]; + ivec4 temp_B[B_PASSES]; + + // ========================================================= + // PROLOGUE: prefetch chunk 0 into temp registers, then store to + // slice 0. No barrier here -- the main loop's first iteration barriers + // before reading slice 0. + // ========================================================= + { + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + const uint row = tile_m_start + p * A_ROWS_PER_PASS + a_row_offset; + const uint k_hv4 = (a_col * FP16_PER_VEC4) / 4u; + temp_A[p] = load_a_vec4(row, k_hv4, K4); + } + cached_group = 0u; + sc0 = t_weight_scales[sc_n4]; + sc1 = t_weight_scales[sc_n4 + 1u]; + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + if (p > 0u) { + temp_B[p] = temp_B[0]; + continue; + } + const uint k_row = B_TROW(0u); + ivec4 wblock; + wblock = WFETCH(ivec2(k_row >> 2u, n8_blk)); + temp_B[p] = wblock; + } + } + { + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + ASH_STORE(A_SH_IDX(p * A_ROWS_PER_PASS + a_row_offset, a_col), temp_A[p]); + } + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + BSH_STORE(B_SH_IDX(B_TROW(p), b_col), + dequant_block(temp_B[p], B_TCOL(p), B_TCOL(p) + 1u, sc0, sc1)); + } + } + + // ========================================================= + // MAIN LOOP: one barrier per iteration, unconditional body, loop + // peeled (bound excludes the last chunk). Iteration `chunk` does: + // 1. barrier — slice (chunk%2) fully written + // 2. prefetch — chunk+1 from global into temp (in flight during math) + // 3. MMA math — on slice (chunk%2) + // 4. store — temp (chunk+1, dequantized) into slice ((chunk+1)%2) + // ========================================================= + uint chunk; + for (chunk = 0; chunk + 1u < num_chunks; ++chunk) { + const uint cur_base_A = (chunk % 2u) * ASH_SLICE; + const uint cur_base_B = (chunk % 2u) * BSH_SLICE; + const uint nxt_base_A = ((chunk + 1u) % 2u) * ASH_SLICE; + const uint nxt_base_B = ((chunk + 1u) % 2u) * BSH_SLICE; + + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + + // --- prefetch chunk+1 -> temp --- + { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + const uint row = tile_m_start + p * A_ROWS_PER_PASS + a_row_offset; + const uint k_hv4 = (chunkK_nxt + a_col * FP16_PER_VEC4) / 4u; + temp_A[p] = load_a_vec4(row, k_hv4, K4); + } + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + if (p > 0u) { + temp_B[p] = temp_B[0]; + continue; + } + const uint k_row = chunkK_nxt + B_TROW(0u); + temp_B[p] = WFETCH(ivec2(k_row >> 2u, n8_blk)); + } + const uint group_nxt = (chunk + 1u) / CHUNKS_PER_GROUP; + if (group_nxt != cached_group) { + cached_group = group_nxt; + sc0 = t_weight_scales[group_nxt * N4 + sc_n4]; + sc1 = t_weight_scales[group_nxt * N4 + sc_n4 + 1u]; + } + } + + // --- MMA math on the cur slice --- + [[unroll]] for (uint k = 0; k < WG_TILE_K / MMA_K; ++k) { + const uint k_start = MMA_K * k; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash, + cur_base_A + A_SH_IDX(row_a, k_start / FP16_PER_VEC4), + A_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j) / FP16_PER_VEC4; + coopMatLoad( + matB, Bsh, + cur_base_B + B_SH_IDX(k_start, col_b), + B_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = coopMatMulAdd(matA[i], matB, result[i][j]); + } + } + } + + + // --- store temp (chunk+1) -> nxt slice, dequantizing B --- + { + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + ASH_STORE(nxt_base_A + A_SH_IDX(p * A_ROWS_PER_PASS + a_row_offset, a_col), + temp_A[p]); + } + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + BSH_STORE(nxt_base_B + B_SH_IDX(B_TROW(p), b_col), + dequant_block(temp_B[p], B_TCOL(p), B_TCOL(p) + 1u, sc0, sc1)); + } + } + } + + // --- epilogue: barrier, then MMA on the last chunk (loop peeled) --- + { + const uint cur_base_A = (chunk % 2u) * ASH_SLICE; + const uint cur_base_B = (chunk % 2u) * BSH_SLICE; + + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + + [[unroll]] for (uint k = 0; k < WG_TILE_K / MMA_K; ++k) { + const uint k_start = MMA_K * k; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash, + cur_base_A + A_SH_IDX(row_a, k_start / FP16_PER_VEC4), + A_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j) / FP16_PER_VEC4; + coopMatLoad( + matB, Bsh, + cur_base_B + B_SH_IDX(k_start, col_b), + B_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = coopMatMulAdd(matA[i], matB, result[i][j]); + } + } + } + } + + + // --- Store result tile --- +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh at once, so the SG_GRID_Y bands it holds are disjoint global row + // ranges; the whole workgroup then imageStores them. The band-to-global-row + // map below reproduces the buffer path's gi exactly: + // buffer: gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i) + // = tile_m_start + warpInTile.y * SG_TILE_M + i * MMA_M + // and lr / MMA_M is the writing subgroup's warpInTile.y. + // + // PORTABILITY NOTE: [[unroll]] is only a hint and glslc does not honor it + // here -- the barrier() in the loop body keeps the loop rolled, so + // result[i][j] is dynamically indexed. Coopmat arrays are opaque per-lane + // storage and dynamic indexing is a construct shader compilers have broken + // before, so check this first if the texture variants miscompile on a new + // device. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh against the previous iteration's readers. Inert on i == 0 but + // must stay unconditional to remain workgroup-uniform. + // barrier() alone does not order shared-memory stores against a later + // coopMatLoad on at least one driver (a stale MMA_M-row band of A in a + // small fraction of runs), hence the explicit memoryBarrierShared(). + memoryBarrierShared(); + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef CSH_IN_ASH + coopMatStore( + coopmat(result[i][j]), Ash, + (warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j)) / FP16_PER_VEC4, + WG_TILE_N / FP16_PER_VEC4, + gl_CooperativeMatrixLayoutRowMajor); +#else + coopMatStore( + coopmat(result[i][j]), Csh, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); +#endif + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + // The last M tile may be partial: rows at or past M are computed but + // not stored. + if (m >= uint(output_sizes.y)) { + continue; + } +#ifdef CSH_IN_ASH + // base is a multiple of 4 halves: the texel is one half of a uvec4. + const uvec4 q = Ash[base / FP16_PER_VEC4]; + const bool hi = (base % FP16_PER_VEC4) != 0u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4(unpackHalf2x16(hi ? q.z : q.x), unpackHalf2x16(hi ? q.w : q.y))); +#else + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh[base]), + float(Csh[base + 1u]), + float(Csh[base + 2u]), + float(Csh[base + 3u]))); +#endif + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} + +#endif // LINEAR_Q4GSW_WMMA_BX_BODY_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_block.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_block.glslh index 8f19418cd19..2e0f12f95bb 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_block.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_block.glslh @@ -69,6 +69,12 @@ void quantize_and_pack( } } +// The 4h4w block writer assumes t_packed_int8_input is an ivec4-typed +// (block-packed) resource and would not compile against a scalar int array. +// A shader that reuses the quantize helpers above but writes another layout +// (quantize_and_pack_4w_with_group_sums.glsl) defines SKIP_BLOCK_WRITE_HELPERS +// before including this file. +#ifndef SKIP_BLOCK_WRITE_HELPERS void write_block( const Int8InputBlock block, const int block_x, @@ -80,5 +86,6 @@ void write_block( imageStore(t_packed_int8_input, ivec3(block_x, block_y, 0), block.data); #endif // OUTPUT_BUFFER } +#endif // SKIP_BLOCK_WRITE_HELPERS #endif // LINEAR_INT8_INPUT_BLOCK_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.h b/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.h index 96e9cc7c1d3..2adf41469da 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.h +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.h @@ -34,6 +34,14 @@ void add_quantize_and_pack_4h4w_node( const ValueRef packed_int_input, const ValueRef group_size); +// Global work grid of the *_with_group_sums packers (_o2w32 / _o4w16 +// variants); empty for GEMV inputs. +GlobalWorkGrid pick_quantize_and_pack_4h4w_with_group_sums_gwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args); + void add_quantize_and_pack_4h4w_with_group_sums_node( ComputeGraph& graph, const QuantizationConfig& input_quant_config, diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp new file mode 100644 index 00000000000..06a3efea361 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.cpp @@ -0,0 +1,277 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +namespace vkcompute { + +namespace { + +bool has_shape( + const CoopmatDeviceKey& device, + const uint32_t m, + const uint32_t n, + const uint32_t k, + const CoopmatComponentType ab, + const CoopmatComponentType result) { + for (const CoopmatShape& s : device.coopmat_shapes) { + if (s.m == m && s.n == n && s.k == k && s.ab == ab && s.result == result) { + return true; + } + } + return false; +} + +bool subgroup_size_available( + const CoopmatDeviceKey& device, + const uint32_t size) { + return device.subgroup_size == size || + (device.subgroup_size_control && size >= device.min_subgroup_size && + size <= device.max_subgroup_size); +} + +// +// Intel Xe2 rows. MMA 8x16x16 / 8x16x32 (the only fp16 / int8 cooperative +// matrix shapes Xe2 exposes), subgroup 16. +// + +// Weight-heavy projections run their workgroups in blocks of N tiles. The +// threshold lies between the largest shape where the blocked order is slower +// (K * N = 16.8e6) and the smallest where it is faster (21.2e6) on both the +// Arc B580 and the Arc Pro B70. +bool xe2_weight_heavy(const WmmaShape& s) { + return s.K * s.N >= 20000000; +} + +const std::vector& xe2_rows() { + static const std::vector rows = { + // 4w, texture3d: banded drain, 8 x 2 subgroups. + {WmmaOp::kQ4gsw, + "linear_q4gsw_wmma_t128x128k16g82s16m8flib", + {128, 128, 16, 8, 2, 16, 8, false}, + kWmmaTex3dTex2d, + nullptr, + /*rowmajor_a=*/false, + /*image_a=*/true}, + // 4w, buffer IO: 4 x 4 subgroups. + {WmmaOp::kQ4gsw, + "linear_q4gsw_wmma_t128x128k16g44s16m8fli", + {128, 128, 16, 4, 4, 16, 8, false}, + kWmmaBufTex2d, + nullptr}, + // 8da4w: texel-wise B staging and a full A thread map; blocked + // workgroup order on weight-heavy projections, dispatch order elsewhere. + {WmmaOp::kDq8caQ4gsw, + "linear_dq8ca_q4gsw_wmma_t128x128k64g84s16m8n13", + {128, 128, 64, 8, 4, 16, 8, false}, + kWmmaTex3dTex2d | kWmmaBufTex2d, + xe2_weight_heavy}, + {WmmaOp::kDq8caQ4gsw, + "linear_dq8ca_q4gsw_wmma_t128x128k64g84s16m8i", + {128, 128, 64, 8, 4, 16, 8, false}, + kWmmaTex3dTex2d | kWmmaBufTex2d, + nullptr}, + }; + return rows; +} + +// +// NVIDIA rows. MMA 16x16x16 (fp16) / 16x16x32 (int8), subgroup 32, +// texture3d IO only. +// + +// fp16 accumulation loses accuracy beyond K = 8192: fp32 accumulation there. +bool nv_large_k(const WmmaShape& s) { + return s.K > 8192; +} +bool nv_256(const WmmaShape& s) { + return s.K <= 8192 && (s.M != 256 || s.N % 2048 == 0); +} +bool nv_default(const WmmaShape& s) { + return s.K <= 8192; +} + +const std::vector& nvidia_rows() { + static const std::vector rows = { + // 4w, K > 8192: fp32 accumulation, texel-wise B staging, drain in the + // A tile. + {WmmaOp::kQ4gsw, + "linear_q4gsw_wmma_bx_t128x128k32g42s32f32c", + {128, 128, 32, 4, 2, 32, 16, true}, + kWmmaTex3dTex2d, + nv_large_k}, + // 4w, K <= 8192: N-major B in shared memory. + {WmmaOp::kQ4gsw, + "linear_q4gsw_wmma_t256x128k16g42s32bt", + {256, 128, 16, 4, 2, 32, 16, false}, + kWmmaTex3dTex2d, + nv_256}, + // 4w, K <= 8192, the remaining shapes (M not a multiple of 256, and + // M = 256 at N % 2048 != 0). + {WmmaOp::kQ4gsw, + "linear_q4gsw_wmma_t128x128k16g22s32", + {128, 128, 16, 2, 2, 32, 16, false}, + kWmmaTex3dTex2d, + nv_default}, + // 8da4w: row-major int8 activations, 2 x 4 subgroups. + {WmmaOp::kDq8caQ4gsw, + "linear_dq8ca_q4gsw_wmma_rowmajor_t128x128k64g24s32mk32", + {128, 128, 64, 2, 4, 32, 16, true}, + kWmmaTex3dTex2d, + nullptr, + /*rowmajor_a=*/true}, + }; + return rows; +} + +} // namespace + +CoopmatFamily coopmat_family(const CoopmatDeviceKey& device) { + if (!device.coopmat || !device.int8_coopmat) { + return CoopmatFamily::kNone; + } + using T = CoopmatComponentType; + if (device.vendor_id == kVendorIdIntel && + device.driver_id == kDriverIdIntelMesa && + has_shape(device, 8, 16, 16, T::kFp16, T::kFp16) && + has_shape(device, 8, 16, 32, T::kSint8, T::kSint32) && + subgroup_size_available(device, 16)) { + return CoopmatFamily::kXe2; + } + if (device.vendor_id == kVendorIdNvidia && + device.driver_id == kDriverIdNvidiaProprietary && + has_shape(device, 16, 16, 16, T::kFp16, T::kFp16) && + has_shape(device, 16, 16, 16, T::kFp16, T::kFp32) && + has_shape(device, 16, 16, 32, T::kSint8, T::kSint32) && + subgroup_size_available(device, 32)) { + return CoopmatFamily::kNvidia; + } + return CoopmatFamily::kNone; +} + +const char* coopmat_family_name(const CoopmatFamily family) { + switch (family) { + case CoopmatFamily::kXe2: + return "xe2"; + case CoopmatFamily::kNvidia: + return "nvidia"; + case CoopmatFamily::kNone: + break; + } + return "none"; +} + +const std::vector& coopmat_family_rows(const CoopmatFamily family) { + static const std::vector kNoRows; + switch (family) { + case CoopmatFamily::kXe2: + return xe2_rows(); + case CoopmatFamily::kNvidia: + return nvidia_rows(); + case CoopmatFamily::kNone: + break; + } + return kNoRows; +} + +bool wmma_row_fits( + const CoopmatDeviceKey& device, + const WmmaShape& shape, + const WmmaRow& row) { + const WmmaTile& tile = row.tile; + if (shape.op != row.op) { + return false; + } + // The kernels are built without bias and read no leading dim. + if (shape.has_bias || shape.gemv || shape.batch != 1 || !shape.half) { + return false; + } + if (!device.coopmat) { + return false; + } + if (shape.op == WmmaOp::kDq8caQ4gsw) { + if (!device.int8_coopmat) { + return false; + } + const WmmaInt8Layout want = + row.rowmajor_a ? WmmaInt8Layout::kRowMajor : WmmaInt8Layout::k4H4W; + if (shape.int8_layout != WmmaInt8Layout::kAny && + shape.int8_layout != want) { + return false; + } + } + // The pipeline requires the kernel's subgroup size. + if (!subgroup_size_available(device, tile.subgroup_size)) { + return false; + } + // One IO_STORAGE parameter covers input and output. + uint8_t combo = 0; + if (shape.output == WmmaStorage::kBuffer && + shape.input == WmmaStorage::kBuffer && + shape.weight == WmmaStorage::kTexture2D) { + combo = kWmmaBufTex2d; + } else if ( + shape.output == WmmaStorage::kTexture3D && + shape.input == WmmaStorage::kTexture3D && shape.io_width_packed && + shape.weight == WmmaStorage::kTexture2D) { + combo = kWmmaTex3dTex2d; + // The texture epilogue stages SG_GRID_Y * MMA_M rows x WG_TILE_N fp16 on + // top of the A/B tiles unless it reuses the A tile. A tile over the + // shared-memory limit must not reach pipeline creation. + if (!tile.csh_in_ash) { + const uint64_t csh_bytes = + uint64_t(tile.sg_grid_y) * tile.mma_m * tile.n * 2u; + if (csh_bytes >= device.max_shared_bytes) { + return false; + } + } + } + if ((combo & row.storages) == 0) { + return false; + } + WmmaShape sel = shape; + if (shape.op == WmmaOp::kQ4gsw && combo == kWmmaTex3dTex2d && + shape.M >= kWmmaPartialTileM) { + sel.M = (shape.M + kWmmaPartialTileM - 1) / kWmmaPartialTileM * + kWmmaPartialTileM; + } + if (!sel.ignore_alignment && + (sel.M % tile.m != 0 || sel.N % tile.n != 0 || sel.K % tile.k != 0 || + sel.group_size % tile.k != 0)) { + return false; + } + return row.shape_ok == nullptr || row.shape_ok(sel); +} + +std::optional select_wmma_kernel( + const CoopmatDeviceKey& device, + const WmmaShape& shape) { + for (const WmmaRow& row : coopmat_family_rows(coopmat_family(device))) { + if (wmma_row_fits(device, shape, row)) { + return WmmaChoice{row.kernel_base, row.tile, row.rowmajor_a, row.image_a}; + } + } + return std::nullopt; +} + +std::optional wmma_tile_for_kernel(const std::string& kernel_name) { + for (const CoopmatFamily family : + {CoopmatFamily::kXe2, CoopmatFamily::kNvidia}) { + for (const WmmaRow& row : coopmat_family_rows(family)) { + const std::string base(row.kernel_base); + if (kernel_name.size() > base.size() && + kernel_name.compare(0, base.size(), base) == 0 && + kernel_name[base.size()] == '_') { + return row.tile; + } + } + } + return std::nullopt; +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h new file mode 100644 index 00000000000..ed810af0149 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/CoopmatFamily.h @@ -0,0 +1,187 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// Cooperative-matrix architecture families for the 4-bit quantized linear ops +// (et_vk.linear_q4gsw / et_vk.q4gsw_linear and et_vk.linear_dq8ca_q4gsw). +// +// A device belongs to a family by its vendor ID, its driver ID and a property +// signature (cooperative-matrix shapes, available subgroup sizes), never by +// its name. Each family has a small set of kernel rows; a device that matches +// no family, or a shape no row of its family accepts, keeps the existing +// kernels unchanged. +// +// Pure C++ (no Vulkan or ComputeGraph types) so it can be unit-tested on a +// host without a GPU (test/custom_ops/coopmat_family_test.cpp). The Vulkan +// side builds a CoopmatDeviceKey from the adapter (WmmaLinear.h). + +#include +#include +#include +#include + +namespace vkcompute { + +// Component types of a cooperative-matrix configuration (VkComponentTypeKHR, +// reduced to the types the family keys test). +enum class CoopmatComponentType { kOther, kFp16, kFp32, kSint8, kSint32 }; + +// One subgroup-scope VkCooperativeMatrixPropertiesKHR entry with A == B and +// C == Result types. +struct CoopmatShape { + uint32_t m; + uint32_t n; + uint32_t k; + CoopmatComponentType ab; // A and B + CoopmatComponentType result; // C and Result +}; + +// VkPhysicalDeviceProperties::vendorID and VkDriverId values the family keys +// use. +constexpr uint32_t kVendorIdIntel = 0x8086; +constexpr uint32_t kVendorIdNvidia = 0x10de; +// VK_DRIVER_ID_NVIDIA_PROPRIETARY +constexpr uint32_t kDriverIdNvidiaProprietary = 4; +// VK_DRIVER_ID_INTEL_OPEN_SOURCE_MESA +constexpr uint32_t kDriverIdIntelMesa = 6; + +struct CoopmatDeviceKey { + uint32_t vendor_id = 0; + uint32_t driver_id = 0; // VkDriverId, 0 when not reported + uint32_t subgroup_size = 0; // default subgroup size + uint32_t min_subgroup_size = 0; + uint32_t max_subgroup_size = 0; + bool subgroup_size_control = false; // a required subgroup size can be set + bool coopmat = false; // VK_KHR_cooperative_matrix is usable + bool int8_coopmat = false; // an int8 cooperative-matrix shape is enumerated + uint32_t max_shared_bytes = 0; // maxComputeSharedMemorySize + std::vector coopmat_shapes; // subgroup scope only +}; + +enum class CoopmatFamily { + kNone, + // Intel Xe2 on the Mesa driver: fp16 8x16x16 (fp16 result), int8 8x16x32, + // subgroup size 16 available. + kXe2, + // NVIDIA proprietary driver: fp16 16x16x16 (fp16 and fp32 result), int8 + // 16x16x32, subgroup size 32. Measured on Ampere (Jetson Orin) and Ada; the + // key does not separate them from other NVIDIA generations with the same + // shapes. + kNvidia, +}; + +// The only keying point. kNone keeps the existing kernels. +CoopmatFamily coopmat_family(const CoopmatDeviceKey& device); + +const char* coopmat_family_name(CoopmatFamily family); + +// +// Kernel selection +// + +enum class WmmaOp { + kQ4gsw, // 4w: fp16 activations, 4-bit group-symmetric weights + kDq8caQ4gsw, // 8da4w: dynamic int8 activations, 4-bit weights +}; + +enum class WmmaStorage { kBuffer, kTexture2D, kTexture3D, kOther }; + +// Layout of the int8 activations of kDq8caQ4gsw. The 4h4w kernels read +// utils::kPackedInt8_4H4W blocks; the row-major kernels read +// utils::kPackedInt8_4W. +enum class WmmaInt8Layout { kAny, k4H4W, kRowMajor }; + +struct WmmaShape { + WmmaOp op = WmmaOp::kQ4gsw; + int64_t M = 0; + int64_t N = 0; + int64_t K = 0; + int64_t group_size = 0; + int64_t batch = 1; // product of the leading dims + bool gemv = false; + bool has_bias = false; + bool half = false; + WmmaStorage input = WmmaStorage::kOther; + WmmaStorage output = WmmaStorage::kOther; + WmmaStorage weight = WmmaStorage::kOther; + bool io_width_packed = false; + WmmaInt8Layout int8_layout = WmmaInt8Layout::kAny; // kDq8caQ4gsw only + // Skip the M/N/K tile-alignment check (row-major 8da4w kernels, whose + // activation layout is fixed at build time; the grid rounds M up). + bool ignore_alignment = false; +}; + +// The kQ4gsw texture-IO kernels accept a partial last M tile. Their rows take +// an M of at least this many rows that is not a multiple of their tile, and are +// selected as for M rounded up to a multiple of it (the smallest WG_TILE_M of +// the family rows), so an unaligned M gets the kernel of the next aligned M. +// Smaller M, buffer IO and kDq8caQ4gsw keep the alignment requirement. +constexpr int64_t kWmmaPartialTileM = 128; + +struct WmmaTile { + uint32_t m; + uint32_t n; + uint32_t k; + uint32_t sg_grid_x; + uint32_t sg_grid_y; + uint32_t subgroup_size; + uint32_t mma_m; + bool csh_in_ash; // texture3d drain staged in the A tile (no extra Csh) + uint32_t wg_size() const { + return sg_grid_x * sg_grid_y * subgroup_size; + } +}; + +// Storage combinations (IO_STORAGE x WEIGHT_STORAGE) a kernel is built for; +// the yaml contains ___half for each set bit. +enum WmmaStorageSet : uint8_t { + kWmmaTex3dTex2d = 1u << 0, // texture3d input/output, texture2d weight + kWmmaBufTex2d = 1u << 1, // buffer input/output, texture2d weight +}; + +struct WmmaRow { + WmmaOp op; + const char* kernel_base; // shader name without ___ + WmmaTile tile; + uint8_t storages; // WmmaStorageSet bits + // Extra shape requirement (beyond tile alignment), or null. Rows are tried + // in order, so a later row can serve the shapes an earlier one rejects. + bool (*shape_ok)(const WmmaShape&); + bool rowmajor_a = false; // kDq8caQ4gsw: reads row-major int8 activations + // kQ4gsw: the texture3d variant reads A through a storage image, so the + // input must be bound in VK_IMAGE_LAYOUT_GENERAL. + bool image_a = false; +}; + +struct WmmaChoice { + std::string kernel_base; + WmmaTile tile; + bool rowmajor_a = false; + bool image_a = false; +}; + +// The rows of a family, in selection order (empty for kNone). +const std::vector& coopmat_family_rows(CoopmatFamily family); + +// The device, storage, alignment and shape constraints of `row`. +bool wmma_row_fits( + const CoopmatDeviceKey& device, + const WmmaShape& shape, + const WmmaRow& row); + +// The first row of the device's family that fits `shape`, or std::nullopt +// (no family, or no row applies): the op keeps its existing kernels. +std::optional select_wmma_kernel( + const CoopmatDeviceKey& device, + const WmmaShape& shape); + +// Tile of a kernel named _ for any family row. +std::optional wmma_tile_for_kernel(const std::string& kernel_name); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/WmmaLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/WmmaLinear.cpp new file mode 100644 index 00000000000..02d94544a0e --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/WmmaLinear.cpp @@ -0,0 +1,118 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include + +namespace vkcompute { + +namespace { + +#ifdef VK_KHR_cooperative_matrix +CoopmatComponentType to_component_type(const VkComponentTypeKHR t) { + switch (t) { + case VK_COMPONENT_TYPE_FLOAT16_KHR: + return CoopmatComponentType::kFp16; + case VK_COMPONENT_TYPE_FLOAT32_KHR: + return CoopmatComponentType::kFp32; + case VK_COMPONENT_TYPE_SINT8_KHR: + return CoopmatComponentType::kSint8; + case VK_COMPONENT_TYPE_SINT32_KHR: + return CoopmatComponentType::kSint32; + default: + return CoopmatComponentType::kOther; + } +} +#endif /* VK_KHR_cooperative_matrix */ + +} // namespace + +CoopmatDeviceKey coopmat_device_key(ComputeGraph* graph) { + vkapi::Adapter* adapter = graph->context()->adapter_ptr(); + CoopmatDeviceKey key; + key.vendor_id = adapter->vendor_id(); + key.driver_id = adapter->driver_id(); + key.subgroup_size = adapter->subgroup_size(); + key.min_subgroup_size = adapter->min_subgroup_size(); + key.max_subgroup_size = adapter->max_subgroup_size(); + key.subgroup_size_control = adapter->supports_subgroup_size_control(); + key.coopmat = adapter->supports_cooperative_matrix(); + key.int8_coopmat = adapter->supports_int8_cooperative_matrix(); + key.max_shared_bytes = adapter->max_compute_shared_memory_size(); +#ifdef VK_KHR_cooperative_matrix + for (const VkCooperativeMatrixPropertiesKHR& p : + adapter->cooperative_matrix_properties()) { + if (p.scope != VK_SCOPE_SUBGROUP_KHR || p.AType != p.BType || + p.CType != p.ResultType) { + continue; + } + key.coopmat_shapes.push_back( + {p.MSize, + p.NSize, + p.KSize, + to_component_type(p.AType), + to_component_type(p.ResultType)}); + } +#endif /* VK_KHR_cooperative_matrix */ + return key; +} + +WmmaStorage to_wmma_storage(const utils::StorageType storage) { + switch (storage) { + case utils::kBuffer: + return WmmaStorage::kBuffer; + case utils::kTexture2D: + return WmmaStorage::kTexture2D; + case utils::kTexture3D: + return WmmaStorage::kTexture3D; + default: + return WmmaStorage::kOther; + } +} + +utils::StorageType predicted_q4_weight_storage( + ComputeGraph& graph, + const ValueRef weight_data) { + const std::vector sizes = graph.sizes_of(weight_data); + const int64_t K = utils::val_at(-1, sizes) * 2; + const int64_t N = utils::val_at(-2, sizes); + // Mirrors the 4-bit packed extents in prepack_quantized_linear_weight(). + const int64_t height = utils::div_up(N, int64_t(8)); + const int64_t width = utils::div_up(K, int64_t(4)) * 4; + const uint32_t max_extent = + graph.context()->adapter_ptr()->max_texture2d_dim(); + return (width > int64_t(max_extent) * 4 || height > int64_t(max_extent)) + ? utils::kBuffer + : utils::kTexture2D; +} + +GlobalWorkGrid wmma_linear_gwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + const std::optional tile = wmma_tile_for_kernel(shader.kernel_name); + if (!tile.has_value()) { + return quantized_linear_gwg(graph, shader, args, resize_args); + } + const std::vector out_sizes = graph->sizes_of(args.at(0).refs.at(0)); + const uint32_t N = + utils::safe_downcast(utils::val_at(-1, out_sizes)); + const uint32_t M = + utils::safe_downcast(utils::val_at(-2, out_sizes)); + // The x extent is multiplied by the workgroup size because the framework + // divides it by the local size. + const uint32_t wg_size = tile->wg_size(); + return GlobalWorkGrid( + {utils::div_up(N, tile->n) * wg_size, utils::div_up(M, tile->m), 1u}, + kTiledWorkGrid, + LocalWorkGroup(wg_size, 1u, 1u)); +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/WmmaLinear.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/WmmaLinear.h new file mode 100644 index 00000000000..ea55cb17cc3 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/WmmaLinear.h @@ -0,0 +1,39 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// ComputeGraph side of the cooperative-matrix family kernels +// (CoopmatFamily.h), shared by the q4gsw and dq8ca_q4gsw node builders. + +#include +#include + +namespace vkcompute { + +// The family key inputs of the graph's device. +CoopmatDeviceKey coopmat_device_key(ComputeGraph* graph); + +WmmaStorage to_wmma_storage(utils::StorageType storage); + +// The storage prepack_quantized_linear_weight() picks for a 4-bit weight: +// texture2d unless the packed extents exceed the device limit. +utils::StorageType predicted_q4_weight_storage( + ComputeGraph& graph, + const ValueRef weight_data); + +// One workgroup per WG_TILE_M x WG_TILE_N output tile for a family kernel; +// quantized_linear_gwg() for any other kernel name. args.at(0).refs.at(0) is +// the output. +GlobalWorkGrid wmma_linear_gwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinear.cpp index dd13ec2ebb5..a1cfbe7c0ef 100644 --- a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinear.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include namespace vkcompute { @@ -153,6 +154,12 @@ void add_linear_dqa_qw_node( void linear_dq8ca_q4gsw( ComputeGraph& graph, const std::vector& args) { + // Devices of a cooperative-matrix family (CoopmatFamily.h) take the + // family's kernels where one of its rows applies. + if (try_add_dq8ca_q4gsw_linear_wmma_node(graph, args)) { + return; + } + int32_t idx = 0; const ValueRef fp_input = args.at(idx++); const ValueRef input_scale = args.at(idx++); diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.cpp new file mode 100644 index 00000000000..6800c6ceb2a --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.cpp @@ -0,0 +1,303 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace vkcompute { + +namespace { + +// Dispatch node args: +// {{output}, {fp_input, packed_int_input, int_input_sums, input_scale, +// input_zp, packed_weight, weight_sums, weight_scales, bias}} +// resize_args: {group_size, weight_data, bias_data} +WmmaShape dq8ca_runtime_shape( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef output = args.at(0).refs.at(0); + const ValueRef fp_input = args.at(1).refs.at(0); + const ValueRef packed_int_input = args.at(1).refs.at(1); + const ValueRef packed_weight = args.at(1).refs.at(5); + + const std::vector out_sizes = graph->sizes_of(output); + WmmaShape s; + s.op = WmmaOp::kDq8caQ4gsw; + s.N = utils::val_at(-1, out_sizes); + s.M = utils::val_at(-2, out_sizes); + s.K = utils::val_at(-1, graph->sizes_of(fp_input)); + s.group_size = graph->extract_scalar(resize_args.at(0)); + for (int64_t d = 0; d < graph->dim_of(output) - 2; d++) { + s.batch *= utils::val_at(d, out_sizes); + } + s.gemv = is_gemv(graph, fp_input); + s.has_bias = !graph->val_is_none(resize_args.at(2)); + s.half = graph->dtype_of(output) == vkapi::kHalf; + s.input = to_wmma_storage(graph->storage_type_of(fp_input)); + s.output = to_wmma_storage(graph->storage_type_of(output)); + s.weight = to_wmma_storage(graph->storage_type_of(packed_weight)); + s.io_width_packed = graph->packed_dim_of(output) == WHCN::kWidthDim && + graph->packed_dim_of(fp_input) == WHCN::kWidthDim; + s.int8_layout = graph->estimate_memory_layout_of(packed_int_input) == + utils::kPackedInt8_4W + ? WmmaInt8Layout::kRowMajor + : WmmaInt8Layout::k4H4W; + return s; +} + +// args: {input, input_scale, input_zp, weight, weight_sums, weight_scales, +// group_size, bias, output} +WmmaShape dq8ca_build_time_shape( + ComputeGraph& graph, + const std::vector& args) { + const ValueRef fp_input = args.at(0); + const ValueRef output = args.at(8); + const std::vector out_sizes = graph.sizes_of(output); + WmmaShape s; + s.op = WmmaOp::kDq8caQ4gsw; + s.N = utils::val_at(-1, out_sizes); + s.M = utils::val_at(-2, out_sizes); + s.K = utils::val_at(-1, graph.sizes_of(fp_input)); + s.group_size = graph.extract_scalar(args.at(6)); + for (int64_t d = 0; d < graph.dim_of(output) - 2; d++) { + s.batch *= utils::val_at(d, out_sizes); + } + s.gemv = is_gemv(&graph, fp_input); + s.has_bias = !graph.val_is_none(args.at(7)); + s.half = graph.dtype_of(output) == vkapi::kHalf; + s.input = to_wmma_storage(graph.storage_type_of(fp_input)); + s.output = to_wmma_storage(graph.storage_type_of(output)); + s.weight = to_wmma_storage(predicted_q4_weight_storage(graph, args.at(3))); + s.io_width_packed = graph.packed_dim_of(output) == WHCN::kWidthDim && + graph.packed_dim_of(fp_input) == WHCN::kWidthDim; + s.int8_layout = WmmaInt8Layout::kAny; + return s; +} + +vkapi::ShaderInfo pick_dq8ca_q4gsw_linear_wmma_shader( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef output = args.at(0).refs.at(0); + const ValueRef fp_input = args.at(1).refs.at(0); + const ValueRef input_zp = args.at(1).refs.at(4); + const ValueRef packed_weight = args.at(1).refs.at(5); + + const CoopmatDeviceKey device = coopmat_device_key(graph); + WmmaShape shape = dq8ca_runtime_shape(graph, args, resize_args); + std::optional choice = select_wmma_kernel(device, shape); + if (!choice.has_value() && shape.int8_layout == WmmaInt8Layout::kRowMajor && + !shape.gemv) { + // The row-major activation layout was chosen at build time and, of the + // M > 1 kernels, only the row-major family kernels read it, so keep one + // for unaligned M: the grid rounds M up. Decode (M == 1) takes the + // linear_dq8ca_q4gsw_coop kernel below, which reads the fp activations + // (the quantize/pack node dispatches nothing at M == 1). + shape.ignore_alignment = true; + choice = select_wmma_kernel(device, shape); + VK_CHECK_COND( + choice.has_value(), "no row-major dq8ca_q4gsw kernel for this shape"); + } + + std::string kernel_name; + if (choice.has_value()) { + kernel_name = choice->kernel_base; + add_storage_type_suffix(kernel_name, graph->storage_type_of(output)); + add_storage_type_suffix(kernel_name, graph->storage_type_of(packed_weight)); + add_dtype_suffix(kernel_name, graph->dtype_of(output)); + } else { + // The kernels for the 4h4w layout (as in pick_linear_dqa_qw_shader). + kernel_name = "linear_dq8ca_q4gsw"; + kernel_name += is_gemv(graph, fp_input) ? "_coop" : "_tiled"; + add_storage_type_suffix(kernel_name, graph->storage_type_of(output)); + add_storage_type_suffix(kernel_name, graph->storage_type_of(packed_weight)); + add_dtype_suffix(kernel_name, graph->dtype_of(output)); + add_zp_dtype_mode_suffix(kernel_name, graph->dtype_of(input_zp)); + } + return VK_KERNEL_FROM_STR(kernel_name); +} + +// Row-major (utils::kPackedInt8_4W) counterpart of +// add_quantize_and_pack_4h4w_with_group_sums_node: same quantization, group +// sums and workgroup shape. +void add_quantize_and_pack_4w_with_group_sums_node( + ComputeGraph& graph, + const ValueRef fp_input, + const ValueRef int_input_sums, + const ValueRef packed_input_scales, + const ValueRef packed_input_zps, + const ValueRef packed_int_input, + const ValueRef group_size) { + const int32_t group_size_val = graph.extract_scalar(group_size); + const int32_t blocks_per_group = utils::div_up(group_size_val, int32_t(4)); + + std::string shader_name = "quantize_and_pack_4w_with_group_sums"; + shader_name += group_size_val >= 128 ? "_o2w32" : "_o4w16"; + add_storage_type_suffix(shader_name, graph.storage_type_of(packed_int_input)); + add_storage_type_suffix(shader_name, graph.storage_type_of(fp_input)); + add_dtype_suffix(shader_name, graph.dtype_of(fp_input)); + add_zp_dtype_mode_suffix(shader_name, graph.dtype_of(packed_input_zps)); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(shader_name), + pick_quantize_and_pack_4h4w_with_group_sums_gwg, + pick_required_lwg, + {{{packed_int_input, int_input_sums}, vkapi::kWrite}, + {{fp_input, packed_input_scales, packed_input_zps}, vkapi::kRead}}, + {graph.sizes_ubo(fp_input)}, + {}, + {blocks_per_group}, + {group_size})); +} + +} // namespace + +bool try_add_dq8ca_q4gsw_linear_wmma_node( + ComputeGraph& graph, + const std::vector& args) { + const ValueRef fp_input = args.at(0); + const ValueRef input_scale = args.at(1); + const ValueRef input_zp = args.at(2); + const ValueRef weight_data = args.at(3); + const ValueRef weight_sums_data = args.at(4); + const ValueRef weight_scales_data = args.at(5); + const ValueRef group_size = args.at(6); + const ValueRef bias_data = args.at(7); + const ValueRef output = args.at(8); + + if (graph.dtype_of(fp_input) != vkapi::kHalf) { + return false; + } + // The decode and fallback kernels of this path use int8 dot products. + if (!graph.can_use_int8_dot_product()) { + return false; + } + const CoopmatDeviceKey device = coopmat_device_key(&graph); + const WmmaShape shape = dq8ca_build_time_shape(graph, args); + // Take the op over only if a row applies to its build-time shape; otherwise + // it stays entirely on the existing path. The kernel chosen here also fixes + // the int8 activation layout. + const std::optional build_choice = + select_wmma_kernel(device, shape); + if (!build_choice.has_value()) { + return false; + } + const bool rowmajor = build_choice->rowmajor_a; + + const int64_t group_size_val = graph.extract_scalar(group_size); + const QuantizationConfig input_quant_config(8, kPerChannel, {}, false, true); + const QuantizationConfig weight_quant_config(4, kPerGroup, {group_size_val}); + + const ValueRef packed_weight = + prepack_quantized_linear_weight(graph, weight_quant_config, weight_data); + const ValueRef packed_weight_scales = prepack_standard( + graph, weight_scales_data, utils::kBuffer, utils::kWidthPacked); + const ValueRef packed_weight_sums = prepack_standard( + graph, weight_sums_data, utils::kBuffer, utils::kWidthPacked); + + TmpTensor dummy_bias( + &graph, {}, graph.dtype_of(output), utils::kBuffer, utils::kWidthPacked); + ValueRef packed_bias = dummy_bias.vref; + uint32_t apply_bias = 0; + if (graph.val_is_not_none(bias_data)) { + packed_bias = + prepack_standard(graph, bias_data, utils::kBuffer, utils::kWidthPacked); + apply_bias = 1; + } + + ValueRef packed_input_scale = input_scale; + ValueRef packed_input_zp = input_zp; + if (graph.val_is_tref(input_scale)) { + VK_CHECK_COND(graph.val_is_tref(input_zp)); + packed_input_scale = prepack_standard( + graph, input_scale, utils::kTexture3D, utils::kWidthPacked); + packed_input_zp = prepack_standard( + graph, input_zp, utils::kTexture3D, utils::kWidthPacked); + } + + TmpTensor packed_int_input( + &graph, + graph.sizes_of(fp_input), + vkapi::kInt8x4, + utils::kBuffer, + rowmajor ? utils::kPackedInt8_4W : utils::kPackedInt8_4H4W); + + // Per-group int8 input sums, indexed as ivec4[group_idx * M4 + m4]. + const int64_t num_groups = graph.size_at(-2, weight_scales_data); + const int64_t M4 = utils::div_up(shape.M, int64_t(4)); + TmpTensor int_input_sums( + &graph, + {num_groups * M4 * 4}, + vkapi::kInt, + utils::kBuffer, + utils::kWidthPacked); + + if (rowmajor) { + add_quantize_and_pack_4w_with_group_sums_node( + graph, + fp_input, + int_input_sums, + packed_input_scale, + packed_input_zp, + packed_int_input, + group_size); + } else { + add_quantize_and_pack_4h4w_with_group_sums_node( + graph, + input_quant_config, + fp_input, + int_input_sums, + packed_input_scale, + packed_input_zp, + packed_int_input, + group_size); + } + + const int32_t K4_per_group = + utils::div_up(static_cast(group_size_val), int32_t(4)); + const int32_t num_groups_arg = graph.size_at(-1, fp_input) / + static_cast(group_size_val); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + pick_dq8ca_q4gsw_linear_wmma_shader, + wmma_linear_gwg, + quantized_linear_lwg, + {{output, vkapi::kWrite}, + {{fp_input, + packed_int_input.vref, + int_input_sums.vref, + packed_input_scale, + packed_input_zp, + packed_weight, + packed_weight_sums, + packed_weight_scales, + packed_bias}, + vkapi::kRead}}, + {graph.sizes_ubo(output), graph.sizes_ubo(fp_input)}, + {}, + // The spec constants of add_linear_dqa_qw_node. + {apply_bias, + K4_per_group, + num_groups_arg, + graph.size_at(-1, output)}, + {group_size, weight_data, bias_data}, + resize_linear_qw_node)); + return true; +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.h new file mode 100644 index 00000000000..ce94047574e --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinearWmma.h @@ -0,0 +1,26 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace vkcompute { + +// Cooperative-matrix family path of et_vk.linear_dq8ca_q4gsw +// (CoopmatFamily.h). If the device belongs to a family and one of its rows +// applies to the op's build-time shape, adds the whole op (activation +// quantize/pack and the int8 cooperative-matrix linear) and returns true; +// otherwise adds nothing and returns false, and the caller builds its usual +// nodes. `args`: {input, input_scale, input_zp, weight, weight_sums, +// weight_scales, group_size, bias, output}. +bool try_add_dq8ca_q4gsw_linear_wmma_node( + ComputeGraph& graph, + const std::vector& args); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinear.cpp index 8e7d16092c7..aebc0acf5b2 100644 --- a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinear.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include namespace vkcompute { @@ -634,6 +635,12 @@ void add_q4gsw_linear_tin_w_4x8_node( } void q4gsw_linear(ComputeGraph& graph, const std::vector& args) { + // Devices of a cooperative-matrix family (CoopmatFamily.h) take the + // family's kernels where one of its rows applies. + if (try_add_q4gsw_linear_wmma_node(graph, args)) { + return; + } + int32_t idx = 0; const ValueRef fp_input = args.at(idx++); const ValueRef weight_data = args.at(idx++); diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp new file mode 100644 index 00000000000..c6ac5b46424 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.cpp @@ -0,0 +1,184 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace vkcompute { + +namespace { + +// Dispatch node args: +// {{output}, {fp_input}, {packed_weight, packed_weight_scales, packed_bias}} +// resize_args: {group_size, weight_data, bias_data} +WmmaShape q4gsw_runtime_shape( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef output = args.at(0).refs.at(0); + const ValueRef fp_input = args.at(1).refs.at(0); + const ValueRef packed_weight = args.at(2).refs.at(0); + + const std::vector out_sizes = graph->sizes_of(output); + WmmaShape s; + s.op = WmmaOp::kQ4gsw; + s.N = utils::val_at(-1, out_sizes); + s.M = utils::val_at(-2, out_sizes); + s.K = utils::val_at(-1, graph->sizes_of(fp_input)); + s.group_size = graph->extract_scalar(resize_args.at(0)); + for (int64_t d = 0; d < graph->dim_of(output) - 2; d++) { + s.batch *= utils::val_at(d, out_sizes); + } + s.gemv = is_gemv(graph, fp_input); + s.has_bias = !graph->val_is_none(resize_args.at(2)); + s.half = graph->dtype_of(output) == vkapi::kHalf; + s.input = to_wmma_storage(graph->storage_type_of(fp_input)); + s.output = to_wmma_storage(graph->storage_type_of(output)); + s.weight = to_wmma_storage(graph->storage_type_of(packed_weight)); + s.io_width_packed = graph->packed_dim_of(output) == WHCN::kWidthDim && + graph->packed_dim_of(fp_input) == WHCN::kWidthDim; + return s; +} + +// args: {input, weight, weight_scales, group_size, bias, output} +WmmaShape q4gsw_build_time_shape( + ComputeGraph& graph, + const std::vector& args) { + const ValueRef fp_input = args.at(0); + const ValueRef output = args.at(5); + const std::vector out_sizes = graph.sizes_of(output); + WmmaShape s; + s.op = WmmaOp::kQ4gsw; + s.N = utils::val_at(-1, out_sizes); + s.M = utils::val_at(-2, out_sizes); + s.K = utils::val_at(-1, graph.sizes_of(fp_input)); + s.group_size = graph.extract_scalar(args.at(3)); + for (int64_t d = 0; d < graph.dim_of(output) - 2; d++) { + s.batch *= utils::val_at(d, out_sizes); + } + s.gemv = is_gemv(&graph, fp_input); + s.has_bias = !graph.val_is_none(args.at(4)); + s.half = graph.dtype_of(output) == vkapi::kHalf; + s.input = to_wmma_storage(graph.storage_type_of(fp_input)); + s.output = to_wmma_storage(graph.storage_type_of(output)); + s.weight = to_wmma_storage(predicted_q4_weight_storage(graph, args.at(1))); + s.io_width_packed = graph.packed_dim_of(output) == WHCN::kWidthDim && + graph.packed_dim_of(fp_input) == WHCN::kWidthDim; + return s; +} + +vkapi::ShaderInfo pick_q4gsw_linear_wmma_shader( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef output = args.at(0).refs.at(0); + const ValueRef fp_input = args.at(1).refs.at(0); + const ValueRef packed_weight = args.at(2).refs.at(0); + + const std::optional choice = select_wmma_kernel( + coopmat_device_key(graph), q4gsw_runtime_shape(graph, args, resize_args)); + + std::string kernel_name; + if (choice.has_value()) { + kernel_name = choice->kernel_base; + } else { + // Shapes no family row covers after a resize (decode, M below + // kWmmaPartialTileM, buffer IO at unaligned M): the kernels that read + // this packed weight layout. + kernel_name = + is_gemv(graph, fp_input) ? "linear_q4gsw_coop" : "linear_q4gsw_tiled"; + } + add_storage_type_suffix(kernel_name, graph->storage_type_of(output)); + add_storage_type_suffix(kernel_name, graph->storage_type_of(packed_weight)); + add_dtype_suffix(kernel_name, graph->dtype_of(output)); + return VK_KERNEL_FROM_STR(kernel_name); +} + +} // namespace + +bool try_add_q4gsw_linear_wmma_node( + ComputeGraph& graph, + const std::vector& args) { + const ValueRef fp_input = args.at(0); + const ValueRef weight_data = args.at(1); + const ValueRef weight_scales_data = args.at(2); + const ValueRef group_size = args.at(3); + const ValueRef bias_data = args.at(4); + const ValueRef output = args.at(5); + + // The family kernels are fp16 only. + if (graph.dtype_of(fp_input) != vkapi::kHalf) { + return false; + } + // Take the op over only if a row applies to its build-time shape; otherwise + // it stays entirely on the existing path, including for later resizes. + const std::optional build_choice = select_wmma_kernel( + coopmat_device_key(&graph), q4gsw_build_time_shape(graph, args)); + if (!build_choice.has_value()) { + return false; + } + // A storage-image read of the input needs VK_IMAGE_LAYOUT_GENERAL + // (VUID-vkCmdDispatch-imageLayout-00344), which read-write access selects. + // The fallback kernels sample the input, which GENERAL also allows. + const vkapi::MemoryAccessFlags input_access = build_choice->image_a && + graph.storage_type_of(fp_input) == utils::kTexture3D + ? vkapi::kReadWrite + : vkapi::kRead; + + const int64_t group_size_val = graph.extract_scalar(group_size); + const QuantizationConfig weight_quant_config(4, kPerGroup, {group_size_val}); + + const ValueRef packed_weight = + prepack_quantized_linear_weight(graph, weight_quant_config, weight_data); + const ValueRef packed_weight_scales = prepack_standard( + graph, weight_scales_data, utils::kBuffer, utils::kWidthPacked); + + TmpTensor dummy_bias( + &graph, {}, graph.dtype_of(output), utils::kBuffer, utils::kWidthPacked); + ValueRef packed_bias = dummy_bias.vref; + uint32_t apply_bias = 0; + if (graph.val_is_not_none(bias_data)) { + packed_bias = + prepack_standard(graph, bias_data, utils::kBuffer, utils::kWidthPacked); + apply_bias = 1; + } + + const int32_t K4_per_group = + utils::div_up(static_cast(group_size_val), int32_t(4)); + const int32_t num_groups = graph.size_at(-1, fp_input) / + static_cast(group_size_val); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + pick_q4gsw_linear_wmma_shader, + wmma_linear_gwg, + quantized_linear_lwg, + {{output, vkapi::kWrite}, + {fp_input, input_access}, + {{packed_weight, packed_weight_scales, packed_bias}, vkapi::kRead}}, + {graph.sizes_ubo(output), graph.sizes_ubo(fp_input)}, + {}, + // The spec constants of the linear_q4gsw_{tiled,coop} kernels. The + // output width N is a spec constant for the coopMatStore address math. + {apply_bias, + K4_per_group, + num_groups, + graph.size_at(-1, output)}, + {group_size, weight_data, bias_data}, + resize_linear_qw_node)); + return true; +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.h new file mode 100644 index 00000000000..3d452c6bbed --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinearWmma.h @@ -0,0 +1,25 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace vkcompute { + +// Cooperative-matrix family path of et_vk.linear_q4gsw / et_vk.q4gsw_linear +// (CoopmatFamily.h). If the device belongs to a family and one of its rows +// applies to the op's build-time shape, adds the op's nodes and returns true; +// otherwise adds nothing and returns false, and the caller builds its usual +// nodes. `args` is the op's argument list: +// {input, weight, weight_scales, group_size, bias, output}. +bool try_add_q4gsw_linear_wmma_node( + ComputeGraph& graph, + const std::vector& args); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/vk_api/Adapter.cpp b/backends/vulkan/runtime/vk_api/Adapter.cpp index eea61a8fe3f..501d80a7130 100644 --- a/backends/vulkan/runtime/vk_api/Adapter.cpp +++ b/backends/vulkan/runtime/vk_api/Adapter.cpp @@ -10,6 +10,7 @@ #include +#include #include #include @@ -134,6 +135,13 @@ VkDevice create_logical_device( #endif /* VK_EXT_subgroup_size_control */ }; +#ifdef VK_KHR_vulkan_memory_model + if (physical_device.vulkan_memory_model_features.vulkanMemoryModel) { + requested_device_extensions.push_back( + VK_KHR_VULKAN_MEMORY_MODEL_EXTENSION_NAME); + } +#endif /* VK_KHR_vulkan_memory_model */ + std::vector enabled_device_extensions; find_requested_device_extensions( physical_device.handle, @@ -170,6 +178,25 @@ VkDevice create_logical_device( void* extension_list_top = nullptr; +#ifdef VK_KHR_vulkan_memory_model + VkPhysicalDeviceVulkanMemoryModelFeaturesKHR vulkan_memory_model_features{ + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_MEMORY_MODEL_FEATURES_KHR}; + for (const char* extension : enabled_device_extensions) { + if (std::strcmp(extension, VK_KHR_VULKAN_MEMORY_MODEL_EXTENSION_NAME) == + 0) { + vulkan_memory_model_features.vulkanMemoryModel = VK_TRUE; + // Existing shaders use Device memory scope + // (VUID-RuntimeSpirv-vulkanMemoryModel-06265). + vulkan_memory_model_features.vulkanMemoryModelDeviceScope = + physical_device.vulkan_memory_model_features + .vulkanMemoryModelDeviceScope; + vulkan_memory_model_features.pNext = extension_list_top; + extension_list_top = &vulkan_memory_model_features; + break; + } + } +#endif /* VK_KHR_vulkan_memory_model */ + #ifdef VK_KHR_16bit_storage VkPhysicalDevice16BitStorageFeatures shader_16bit_storage{ physical_device.shader_16bit_storage}; @@ -300,7 +327,12 @@ Adapter::Adapter( shader_layout_cache_(device_.handle), shader_cache_(device_.handle), pipeline_layout_cache_(device_.handle), - compute_pipeline_cache_(device_.handle, cache_data_path), + compute_pipeline_cache_( + device_.handle, + cache_data_path, + // Full subgroups: what create_logical_device enabled + supports_subgroup_size_control() && + supports_compute_full_subgroups()), sampler_cache_(device_.handle), vma_(instance_, physical_device_.handle, device_.handle), linear_tiling_3d_enabled_{test_linear_tiling_3d_image_support( diff --git a/backends/vulkan/runtime/vk_api/Adapter.h b/backends/vulkan/runtime/vk_api/Adapter.h index 0c2784ae473..933b16ac70c 100644 --- a/backends/vulkan/runtime/vk_api/Adapter.h +++ b/backends/vulkan/runtime/vk_api/Adapter.h @@ -146,6 +146,16 @@ class Adapter final { return physical_device_.device_type; } + // VkPhysicalDeviceProperties::vendorID. + inline uint32_t vendor_id() const { + return physical_device_.properties.vendorID; + } + + // VkDriverId of the driver, or 0 when the device does not report it. + inline uint32_t driver_id() const { + return physical_device_.driver_id; + } + // Driver-reported physical device type. Use device_type() above for // vendor-specific quirks; this for capability-tier decisions. inline bool is_integrated_gpu() const { @@ -326,6 +336,20 @@ class Adapter final { #endif /* VK_KHR_cooperative_matrix */ } +#ifdef VK_KHR_cooperative_matrix + // Every cooperative-matrix configuration the device enumerates (empty when + // cooperative matrices are unsupported). + inline const std::vector& + cooperative_matrix_properties() const { +#ifdef ETVK_FORCE_NO_EXTENSIONS + static const std::vector kNone; + return kNone; +#else + return physical_device_.cooperative_matrix_properties; +#endif + } +#endif /* VK_KHR_cooperative_matrix */ + inline bool supports_int16_shader_types() { #ifdef ETVK_FORCE_NO_EXTENSIONS return false; @@ -491,6 +515,10 @@ class Adapter final { return physical_device_.properties.limits.maxComputeWorkGroupInvocations; } + inline uint32_t max_compute_shared_memory_size() const { + return physical_device_.properties.limits.maxComputeSharedMemorySize; + } + inline uint32_t recommended_lwg_nthreads() const { return 64u; } diff --git a/backends/vulkan/runtime/vk_api/Device.cpp b/backends/vulkan/runtime/vk_api/Device.cpp index 56902f05cd8..169d9b58c70 100644 --- a/backends/vulkan/runtime/vk_api/Device.cpp +++ b/backends/vulkan/runtime/vk_api/Device.cpp @@ -47,6 +47,10 @@ PhysicalDevice::PhysicalDevice( handle(physical_device_handle), properties{}, memory_properties{}, +#ifdef VK_KHR_vulkan_memory_model + vulkan_memory_model_features{ + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_MEMORY_MODEL_FEATURES_KHR}, +#endif /* VK_KHR_vulkan_memory_model */ #ifdef VK_KHR_16bit_storage shader_16bit_storage{ VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_16BIT_STORAGE_FEATURES, @@ -74,6 +78,7 @@ PhysicalDevice::PhysicalDevice( cooperative_matrix_features{ VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR}, supports_int8_coopmat{false}, + cooperative_matrix_properties{}, #endif /* VK_KHR_cooperative_matrix */ #ifdef VK_NV_cooperative_matrix2 cooperative_matrix2_features{ @@ -107,7 +112,8 @@ PhysicalDevice::PhysicalDevice( supports_subgroup_size_control(false), supports_compute_full_subgroups(false), device_name{}, - device_type{DeviceType::UNKNOWN} { + device_type{DeviceType::UNKNOWN}, + driver_id(0) { // Extract physical device properties vkGetPhysicalDeviceProperties(handle, &properties); @@ -262,6 +268,19 @@ void PhysicalDevice::query_extensions_vk_1_1() { // Create linked list to query availability of extensions void* extension_list_top = nullptr; +#ifdef VK_KHR_vulkan_memory_model + // Runtime requests Vulkan 1.1, so use the advertised extension even on 1.2+. + std::vector memory_model_extensions; + find_requested_device_extensions( + handle, + memory_model_extensions, + {VK_KHR_VULKAN_MEMORY_MODEL_EXTENSION_NAME}); + if (!memory_model_extensions.empty()) { + vulkan_memory_model_features.pNext = extension_list_top; + extension_list_top = &vulkan_memory_model_features; + } +#endif /* VK_KHR_vulkan_memory_model */ + #ifdef VK_KHR_16bit_storage shader_16bit_storage.pNext = extension_list_top; extension_list_top = &shader_16bit_storage; @@ -340,13 +359,15 @@ void PhysicalDevice::query_extensions_vk_1_1() { p.sType = VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_KHR; p.pNext = nullptr; } - vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR_fn( - handle, &count, props.data()); - for (const auto& p : props) { - if (p.AType == VK_COMPONENT_TYPE_SINT8_KHR) { - supports_int8_coopmat = true; - break; + if (vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR_fn( + handle, &count, props.data()) == VK_SUCCESS) { + props.resize(count); + for (const auto& p : props) { + if (p.AType == VK_COMPONENT_TYPE_SINT8_KHR) { + supports_int8_coopmat = true; + } } + cooperative_matrix_properties = std::move(props); } } } @@ -374,10 +395,34 @@ void PhysicalDevice::query_extensions_vk_1_1() { properties_list_top = &subgroup_size_control_properties; #endif /* VK_EXT_subgroup_size_control */ + // VkPhysicalDeviceDriverProperties is core in Vulkan 1.2; on a 1.1 device it + // needs VK_KHR_driver_properties. + VkPhysicalDeviceDriverProperties driver_properties{ + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DRIVER_PROPERTIES}; + bool query_driver_properties = properties.apiVersion >= VK_API_VERSION_1_2; +#ifdef VK_KHR_driver_properties + if (!query_driver_properties) { + std::vector driver_properties_extensions; + find_requested_device_extensions( + handle, + driver_properties_extensions, + {VK_KHR_DRIVER_PROPERTIES_EXTENSION_NAME}); + query_driver_properties = !driver_properties_extensions.empty(); + } +#endif /* VK_KHR_driver_properties */ + if (query_driver_properties) { + driver_properties.pNext = properties_list_top; + properties_list_top = &driver_properties; + } + properties2.pNext = properties_list_top; vkGetPhysicalDeviceProperties2(handle, &properties2); + if (query_driver_properties) { + driver_id = static_cast(driver_properties.driverID); + } + subgroup_size = subgroup_properties.subgroupSize; supported_subgroup_ops = subgroup_properties.supportedOperations; supported_subgroup_stages = subgroup_properties.supportedStages; diff --git a/backends/vulkan/runtime/vk_api/Device.h b/backends/vulkan/runtime/vk_api/Device.h index 587d3965c3e..f530c913065 100644 --- a/backends/vulkan/runtime/vk_api/Device.h +++ b/backends/vulkan/runtime/vk_api/Device.h @@ -36,6 +36,9 @@ struct PhysicalDevice final { VkPhysicalDeviceMemoryProperties memory_properties; // Additional features available from extensions +#ifdef VK_KHR_vulkan_memory_model + VkPhysicalDeviceVulkanMemoryModelFeaturesKHR vulkan_memory_model_features; +#endif /* VK_KHR_vulkan_memory_model */ #ifdef VK_KHR_16bit_storage VkPhysicalDevice16BitStorageFeatures shader_16bit_storage; #endif /* VK_KHR_16bit_storage */ @@ -57,6 +60,8 @@ struct PhysicalDevice final { // True when VK_COMPONENT_TYPE_SINT8_KHR appears in the enumerated coopmat // property list — required for coopmat shaders (e.g. dq8ca_q4gsw). bool supports_int8_coopmat; + // Every entry of vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR. + std::vector cooperative_matrix_properties; #endif /* VK_KHR_cooperative_matrix */ #ifdef VK_NV_cooperative_matrix2 @@ -103,6 +108,9 @@ struct PhysicalDevice final { // Device identity std::string device_name; DeviceType device_type; + // VkDriverId from VkPhysicalDeviceDriverProperties (Vulkan 1.2 or + // VK_KHR_driver_properties); 0 when the device does not report it. + uint32_t driver_id; explicit PhysicalDevice(VkInstance instance, VkPhysicalDevice); diff --git a/backends/vulkan/runtime/vk_api/Pipeline.cpp b/backends/vulkan/runtime/vk_api/Pipeline.cpp index 2b42e25c92e..bf469a49c15 100644 --- a/backends/vulkan/runtime/vk_api/Pipeline.cpp +++ b/backends/vulkan/runtime/vk_api/Pipeline.cpp @@ -266,6 +266,28 @@ void swap(PipelineLayout& lhs, PipelineLayout& rhs) noexcept { rhs.handle_ = tmp_handle; } +namespace { + +// Cooperative-matrix shaders below SPIR-V 1.6 need full subgroups +// (VUID-RuntimeSpirv-OpTypeCooperativeMatrixKHR-10770). The flag is only valid +// when the local size in X, specialization constant 0, is a multiple of the +// required subgroup size (VUID-VkPipelineShaderStageCreateInfo-pNext-02757). +VkPipelineShaderStageCreateFlags full_subgroups_flag( + const ComputePipeline::Descriptor& descriptor, + const bool full_subgroups) { +#ifdef VK_EXT_subgroup_size_control + const SpecVarList& spec_vars = descriptor.specialization_constants; + const uint32_t subgroup_size = descriptor.required_subgroup_size; + if (full_subgroups && subgroup_size > 0u && spec_vars.size() > 0u && + spec_vars.at(0).value.as_uint32 % subgroup_size == 0u) { + return VK_PIPELINE_SHADER_STAGE_CREATE_REQUIRE_FULL_SUBGROUPS_BIT_EXT; + } +#endif /* VK_EXT_subgroup_size_control */ + return 0u; +} + +} // namespace + // // ComputePipeline // @@ -276,7 +298,8 @@ ComputePipeline::ComputePipeline(VkDevice device, VkPipeline handle) ComputePipeline::ComputePipeline( VkDevice device, const ComputePipeline::Descriptor& descriptor, - VkPipelineCache pipeline_cache) + VkPipelineCache pipeline_cache, + const bool full_subgroups) : device_(device), handle_{VK_NULL_HANDLE} { map_entries_ = descriptor.specialization_constants.generate_map_entries(); @@ -302,7 +325,7 @@ ComputePipeline::ComputePipeline( const VkPipelineShaderStageCreateInfo shader_stage_create_info{ VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO, // sType shader_stage_pnext, // pNext - 0u, // flags + full_subgroups_flag(descriptor, full_subgroups), // flags VK_SHADER_STAGE_COMPUTE_BIT, // stage descriptor.shader_module, // module "main", // pName @@ -417,9 +440,11 @@ void PipelineLayoutCache::purge() { ComputePipelineCache::ComputePipelineCache( VkDevice device, - const std::string& cache_data_path) + const std::string& cache_data_path, + const bool full_subgroups) : cache_mutex_{}, device_(device), + full_subgroups_(full_subgroups), pipeline_cache_{VK_NULL_HANDLE}, cache_{}, cache_data_path_(cache_data_path) { @@ -443,6 +468,7 @@ ComputePipelineCache::ComputePipelineCache( ComputePipelineCache&& other) noexcept : cache_mutex_{}, device_(other.device_), + full_subgroups_(other.full_subgroups_), pipeline_cache_(other.pipeline_cache_), cache_(std::move(other.cache_)) { std::lock_guard lock(other.cache_mutex_); @@ -537,7 +563,7 @@ void ComputePipelineCache::create_pipelines( shader_stage_create_infos.push_back(VkPipelineShaderStageCreateInfo{ VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO, // sType shader_stage_pnext, // pNext - 0u, // flags + full_subgroups_flag(key, full_subgroups_), // flags VK_SHADER_STAGE_COMPUTE_BIT, // stage key.shader_module, // module "main", // pName @@ -579,7 +605,8 @@ VkPipeline ComputePipelineCache::retrieve( it = cache_ .insert( {key, - ComputePipelineCache::Value(device_, key, pipeline_cache_)}) + ComputePipelineCache::Value( + device_, key, pipeline_cache_, full_subgroups_)}) .first; } return it->second.handle(); diff --git a/backends/vulkan/runtime/vk_api/Pipeline.h b/backends/vulkan/runtime/vk_api/Pipeline.h index 5e286ad77d7..fa37387a28c 100644 --- a/backends/vulkan/runtime/vk_api/Pipeline.h +++ b/backends/vulkan/runtime/vk_api/Pipeline.h @@ -172,7 +172,8 @@ class ComputePipeline final { explicit ComputePipeline( VkDevice device, const Descriptor& descriptor, - VkPipelineCache pipeline_cache); + VkPipelineCache pipeline_cache, + const bool full_subgroups = false); ComputePipeline(const ComputePipeline&) = delete; ComputePipeline& operator=(const ComputePipeline&) = delete; @@ -244,7 +245,8 @@ class ComputePipelineCache final { public: explicit ComputePipelineCache( VkDevice device, - const std::string& cache_data_path); + const std::string& cache_data_path, + const bool full_subgroups = false); ComputePipelineCache(const ComputePipelineCache&) = delete; ComputePipelineCache& operator=(const ComputePipelineCache&) = delete; @@ -306,6 +308,9 @@ class ComputePipelineCache final { std::mutex cache_mutex_; VkDevice device_; + // computeFullSubgroups of the device: pipelines with a required subgroup + // size are created with REQUIRE_FULL_SUBGROUPS. + bool full_subgroups_; VkPipelineCache pipeline_cache_; std::unordered_map cache_; const std::string cache_data_path_; diff --git a/backends/vulkan/test/custom_ops/CMakeLists.txt b/backends/vulkan/test/custom_ops/CMakeLists.txt new file mode 100644 index 00000000000..fdf2b485010 --- /dev/null +++ b/backends/vulkan/test/custom_ops/CMakeLists.txt @@ -0,0 +1,50 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Host-only tests that need neither a GPU nor the Vulkan backend library. +# +# ~~~ +# cmake backends/vulkan/test/custom_ops -Bcmake-out/backends/vulkan/test/custom_ops +# cmake --build cmake-out/backends/vulkan/test/custom_ops +# ctest --test-dir cmake-out/backends/vulkan/test/custom_ops +# ~~~ + +cmake_minimum_required(VERSION 3.19) +project(vulkan_custom_ops_tests CXX) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +if(NOT EXECUTORCH_ROOT) + set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../..) +endif() + +if(NOT TARGET GTest::gtest_main) + set(CMAKE_POLICY_VERSION_MINIMUM 3.5) + add_subdirectory( + ${EXECUTORCH_ROOT}/third-party/googletest + ${CMAKE_CURRENT_BINARY_DIR}/googletest EXCLUDE_FROM_ALL + ) +endif() + +enable_testing() + +set(VULKAN_OPS_PATH ${EXECUTORCH_ROOT}/backends/vulkan/runtime/graph/ops) + +add_executable( + vulkan_coopmat_family_test + ${CMAKE_CURRENT_SOURCE_DIR}/coopmat_family_test.cpp + ${VULKAN_OPS_PATH}/impl/gemm/linear/qw/CoopmatFamily.cpp +) +target_include_directories( + vulkan_coopmat_family_test PRIVATE ${EXECUTORCH_ROOT}/.. +) +target_compile_definitions( + vulkan_coopmat_family_test + PRIVATE COOPMAT_FAMILY_TEST_GLSL_DIR="${VULKAN_OPS_PATH}/glsl" +) +target_link_libraries(vulkan_coopmat_family_test PRIVATE GTest::gtest_main) +add_test(vulkan_coopmat_family_test vulkan_coopmat_family_test) diff --git a/backends/vulkan/test/custom_ops/coopmat_family_test.cpp b/backends/vulkan/test/custom_ops/coopmat_family_test.cpp new file mode 100644 index 00000000000..a3f05b3924b --- /dev/null +++ b/backends/vulkan/test/custom_ops/coopmat_family_test.cpp @@ -0,0 +1,541 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +// Host-only tests of the cooperative-matrix family key and the family kernel +// selection (CoopmatFamily.h). No GPU needed. + +#include + +#include + +#include +#include +#include +#include + +namespace vkcompute { +namespace { + +using T = CoopmatComponentType; + +// VkDriverId values used below. +constexpr uint32_t kDriverMesaRadv = 3; +constexpr uint32_t kDriverIntelWindows = 5; +constexpr uint32_t kDriverMesaLlvmpipe = 13; +constexpr uint32_t kDriverMesaNvk = 24; + +CoopmatDeviceKey +base_key(const uint32_t sg, const uint32_t sg_min, const uint32_t sg_max) { + CoopmatDeviceKey d; + d.subgroup_size = sg; + d.min_subgroup_size = sg_min; + d.max_subgroup_size = sg_max; + d.subgroup_size_control = true; + d.coopmat = true; + d.int8_coopmat = true; + d.max_shared_bytes = 65536; + return d; +} + +// Keys as Xe2 (Arc B580 / Arc Pro B70) and NVIDIA (Ampere / Ada) devices +// report them. +CoopmatDeviceKey xe2_key(const uint32_t driver = kDriverIdIntelMesa) { + CoopmatDeviceKey d = base_key(32, 8, 32); + d.vendor_id = kVendorIdIntel; + d.driver_id = driver; + d.coopmat_shapes = { + {8, 16, 16, T::kFp16, T::kFp16}, + {8, 16, 16, T::kFp16, T::kFp32}, + {8, 16, 32, T::kSint8, T::kSint32}}; + return d; +} + +CoopmatDeviceKey nvidia_key( + const uint32_t driver = kDriverIdNvidiaProprietary) { + CoopmatDeviceKey d = base_key(32, 32, 32); + d.vendor_id = kVendorIdNvidia; + d.driver_id = driver; + d.coopmat_shapes = { + {16, 16, 16, T::kFp16, T::kFp16}, + {16, 8, 16, T::kFp16, T::kFp16}, + {16, 16, 16, T::kFp16, T::kFp32}, + {16, 16, 32, T::kSint8, T::kSint32}}; + return d; +} + +// Mesa llvmpipe (lavapipe): 8x8x8 shapes only, fixed subgroup 8. +CoopmatDeviceKey llvmpipe_key() { + CoopmatDeviceKey d = base_key(8, 8, 8); + d.vendor_id = 0x10005; + d.driver_id = kDriverMesaLlvmpipe; + d.max_shared_bytes = 32768; + d.coopmat_shapes = { + {8, 8, 8, T::kFp16, T::kFp16}, + {8, 8, 8, T::kFp16, T::kFp32}, + {8, 8, 8, T::kOther, T::kOther}, // u8 x u8 -> u32 + {8, 8, 8, T::kSint8, T::kSint32}}; + return d; +} + +// AMD RDNA3 under RADV: no family yet. +CoopmatDeviceKey radv_key() { + CoopmatDeviceKey d = base_key(64, 32, 64); + d.vendor_id = 0x1002; + d.driver_id = kDriverMesaRadv; + d.coopmat_shapes = { + {16, 16, 16, T::kFp16, T::kFp16}, + {16, 16, 16, T::kFp16, T::kFp32}, + {16, 16, 16, T::kSint8, T::kSint32}}; + return d; +} + +WmmaShape prefill( + const WmmaOp op, + const int64_t M, + const int64_t K, + const int64_t N, + const WmmaStorage io) { + WmmaShape s; + s.op = op; + s.M = M; + s.K = K; + s.N = N; + s.group_size = 128; + s.half = true; + s.input = io; + s.output = io; + s.weight = WmmaStorage::kTexture2D; + s.io_width_packed = true; + return s; +} + +std::string selected(const CoopmatDeviceKey& d, const WmmaShape& s) { + const std::optional c = select_wmma_kernel(d, s); + return c.has_value() ? c->kernel_base : ""; +} + +// (K, N) of every 4-bit linear in Llama 3.2 1B / 3B and Llama 3.1 8B. +const int64_t kShapes[][2] = { + {2048, 2048}, + {2048, 512}, + {2048, 8192}, + {8192, 2048}, + {3072, 3072}, + {3072, 1024}, + {3072, 8192}, + {8192, 3072}, + {4096, 4096}, + {4096, 1024}, + {4096, 14336}, + {14336, 4096}, +}; + +const std::string kQ = "linear_q4gsw_wmma_"; +const std::string kD = "linear_dq8ca_q4gsw_wmma_"; + +using Rule = std::function; + +// The intended kernel per family, restated independently of the tables. +std::string xe2_rule(const WmmaShape& s) { + if (s.op == WmmaOp::kQ4gsw) { + return kQ + + (s.output == WmmaStorage::kTexture3D ? "t128x128k16g82s16m8flib" + : "t128x128k16g44s16m8fli"); + } + return kD + + (s.K * s.N >= 20000000 ? "t128x128k64g84s16m8n13" + : "t128x128k64g84s16m8i"); +} + +std::string nvidia_rule(const WmmaShape& s) { + if (s.output != WmmaStorage::kTexture3D) { + return ""; + } + if (s.op == WmmaOp::kQ4gsw) { + return kQ + + (s.K > 8192 ? "bx_t128x128k32g42s32f32c" : "t256x128k16g42s32bt"); + } + return kD + "rowmajor_t128x128k64g24s32mk32"; +} + +std::string none_rule(const WmmaShape&) { + return ""; +} + +void check_llama_shapes( + const char* label, + const CoopmatDeviceKey& d, + const Rule& rule) { + for (const WmmaOp op : {WmmaOp::kQ4gsw, WmmaOp::kDq8caQ4gsw}) { + for (const auto& kn : kShapes) { + for (const WmmaStorage io : + {WmmaStorage::kTexture3D, WmmaStorage::kBuffer}) { + const WmmaShape s = prefill(op, 2048, kn[0], kn[1], io); + EXPECT_EQ(selected(d, s), rule(s)) + << label << " op=" << static_cast(op) << " K=" << kn[0] + << " N=" << kn[1] + << " io=" << (io == WmmaStorage::kBuffer ? "buffer" : "texture3d"); + } + } + } +} + +TEST(CoopmatFamilyTest, FamilyKey) { + EXPECT_EQ(coopmat_family(xe2_key()), CoopmatFamily::kXe2); + EXPECT_EQ(coopmat_family(nvidia_key()), CoopmatFamily::kNvidia); + + // Same silicon under another driver: no family. + EXPECT_EQ(coopmat_family(xe2_key(kDriverIntelWindows)), CoopmatFamily::kNone); + EXPECT_EQ(coopmat_family(nvidia_key(kDriverMesaNvk)), CoopmatFamily::kNone); + + // The int8 shape is part of the key. + CoopmatDeviceKey no_int8 = nvidia_key(); + no_int8.int8_coopmat = false; + EXPECT_EQ(coopmat_family(no_int8), CoopmatFamily::kNone); + CoopmatDeviceKey no_int8_shape = xe2_key(); + no_int8_shape.coopmat_shapes.pop_back(); + EXPECT_EQ(coopmat_family(no_int8_shape), CoopmatFamily::kNone); + + // Xe2 needs subgroup size 16 to be available. + CoopmatDeviceKey xe2_sg8 = xe2_key(); + xe2_sg8.subgroup_size = 8; + xe2_sg8.min_subgroup_size = 8; + xe2_sg8.max_subgroup_size = 8; + EXPECT_EQ(coopmat_family(xe2_sg8), CoopmatFamily::kNone); + CoopmatDeviceKey xe2_no_control = xe2_key(); + xe2_no_control.subgroup_size_control = false; + EXPECT_EQ(coopmat_family(xe2_no_control), CoopmatFamily::kNone); + + // No family for vendors without validated kernels. + EXPECT_EQ(coopmat_family(radv_key()), CoopmatFamily::kNone); + + // No cooperative matrix at all. + CoopmatDeviceKey no_coopmat = xe2_key(); + no_coopmat.coopmat = false; + EXPECT_EQ(coopmat_family(no_coopmat), CoopmatFamily::kNone); +} + +TEST(CoopmatFamilyTest, LlvmpipeHasNoFamily) { + const CoopmatDeviceKey d = llvmpipe_key(); + EXPECT_EQ(coopmat_family(d), CoopmatFamily::kNone); + EXPECT_TRUE(coopmat_family_rows(coopmat_family(d)).empty()); + check_llama_shapes("llvmpipe", d, none_rule); + // Shapes aligned to every family tile. + for (const WmmaOp op : {WmmaOp::kQ4gsw, WmmaOp::kDq8caQ4gsw}) { + for (const WmmaStorage io : + {WmmaStorage::kTexture3D, WmmaStorage::kBuffer}) { + EXPECT_FALSE( + select_wmma_kernel(d, prefill(op, 256, 256, 256, io)).has_value()); + } + } +} + +TEST(CoopmatFamilyTest, UnmatchedKeysSelectNothing) { + check_llama_shapes("radv", radv_key(), none_rule); + check_llama_shapes("xe2-windows", xe2_key(kDriverIntelWindows), none_rule); + check_llama_shapes("nvk", nvidia_key(kDriverMesaNvk), none_rule); +} + +TEST(CoopmatFamilyTest, Xe2Rows) { + check_llama_shapes("xe2", xe2_key(), xe2_rule); + // The weight-heavy rule: K * N >= 20e6. + const CoopmatDeviceKey d = xe2_key(); + EXPECT_EQ( + selected( + d, + prefill(WmmaOp::kDq8caQ4gsw, 256, 2048, 9856, WmmaStorage::kBuffer)), + kD + "t128x128k64g84s16m8n13"); // 20.2e6 + EXPECT_EQ( + selected( + d, + prefill(WmmaOp::kDq8caQ4gsw, 256, 2048, 9728, WmmaStorage::kBuffer)), + kD + "t128x128k64g84s16m8i"); // 19.9e6 + // Any aligned shape, not only the Llama projections. + EXPECT_EQ( + selected(d, prefill(WmmaOp::kQ4gsw, 128, 256, 384, WmmaStorage::kBuffer)), + kQ + "t128x128k16g44s16m8fli"); +} + +TEST(CoopmatFamilyTest, NvidiaRows) { + check_llama_shapes("nvidia", nvidia_key(), nvidia_rule); + const CoopmatDeviceKey d = nvidia_key(); + struct Edge { + int64_t M, K, N; + const char* want; // nullptr: existing kernels + }; + const Edge edges[] = { + // M not a multiple of 256: the 128-row tiles. + {384, 2048, 2048, "t128x128k16g22s32"}, + // M = 256 at N % 2048 != 0: not the 256-row tile. + {256, 3072, 1024, "t128x128k16g22s32"}, + {256, 2048, 8192, "t256x128k16g42s32bt"}, + // K > 8192 with M not a multiple of 256: fp32 accumulation. + {384, 14336, 4096, "bx_t128x128k32g42s32f32c"}, + // Any aligned shape, not only the Llama projections (Qwen2.5 0.5B, + // Gemma 2 2B), and M beyond 2048. + {512, 896, 4864, "t256x128k16g42s32bt"}, + {512, 9216, 2304, "bx_t128x128k32g42s32f32c"}, + {4096, 2048, 2048, "t256x128k16g42s32bt"}, + // N not a multiple of the 128-column tile. + {2048, 2048, 4160, nullptr}, + }; + for (const Edge& e : edges) { + const WmmaShape s = + prefill(WmmaOp::kQ4gsw, e.M, e.K, e.N, WmmaStorage::kTexture3D); + EXPECT_EQ(selected(d, s), e.want == nullptr ? "" : kQ + e.want) + << "M=" << e.M << " K=" << e.K << " N=" << e.N; + } + // Group sizes other than 128, down to the K tile of the row. + for (const WmmaOp op : {WmmaOp::kQ4gsw, WmmaOp::kDq8caQ4gsw}) { + WmmaShape g64 = prefill(op, 2048, 2048, 2048, WmmaStorage::kTexture3D); + g64.group_size = 64; + EXPECT_EQ(selected(d, g64), nvidia_rule(g64)); + } + WmmaShape g32 = + prefill(WmmaOp::kDq8caQ4gsw, 2048, 2048, 2048, WmmaStorage::kTexture3D); + g32.group_size = 32; + EXPECT_EQ(selected(d, g32), ""); +} + +TEST(CoopmatFamilyTest, ShapesThatKeepTheExistingKernels) { + for (const CoopmatDeviceKey& d : {xe2_key(), nvidia_key()}) { + for (const WmmaOp op : {WmmaOp::kQ4gsw, WmmaOp::kDq8caQ4gsw}) { + const WmmaShape s = + prefill(op, 2048, 4096, 4096, WmmaStorage::kTexture3D); + ASSERT_FALSE(selected(d, s).empty()); + + // Below kWmmaPartialTileM no row takes an unaligned M. + WmmaShape unaligned = s; + unaligned.M = 120; + WmmaShape gemv = s; + gemv.M = 1; + gemv.gemv = true; + WmmaShape bias = s; + bias.has_bias = true; + WmmaShape fp32 = s; + fp32.half = false; + WmmaShape mixed = s; + mixed.input = WmmaStorage::kBuffer; + WmmaShape batched = s; + batched.batch = 2; + WmmaShape buffer_weight = s; + buffer_weight.weight = WmmaStorage::kBuffer; + WmmaShape channels_packed = s; + channels_packed.io_width_packed = false; + for (const WmmaShape* x : + {&unaligned, + &gemv, + &bias, + &fp32, + &mixed, + &batched, + &buffer_weight, + &channels_packed}) { + EXPECT_FALSE(select_wmma_kernel(d, *x).has_value()) + << coopmat_family_name(coopmat_family(d)) + << " op=" << static_cast(op); + } + } + } +} + +TEST(CoopmatFamilyTest, PartialMTile) { + struct Case { + CoopmatDeviceKey device; + WmmaOp op; + int64_t M, K, N; + WmmaStorage io; + std::string want; // empty: existing kernels + }; + const Case cases[] = { + // 4w texture IO: selected as for M rounded up to 128 rows. + {xe2_key(), + WmmaOp::kQ4gsw, + 511, + 2048, + 2048, + WmmaStorage::kTexture3D, + kQ + "t128x128k16g82s16m8flib"}, + {xe2_key(), + WmmaOp::kQ4gsw, + 129, + 2048, + 2048, + WmmaStorage::kTexture3D, + kQ + "t128x128k16g82s16m8flib"}, + {nvidia_key(), + WmmaOp::kQ4gsw, + 511, + 2048, + 2048, + WmmaStorage::kTexture3D, + kQ + "t256x128k16g42s32bt"}, + {nvidia_key(), + WmmaOp::kQ4gsw, + 300, + 3072, + 1024, + WmmaStorage::kTexture3D, + kQ + "t128x128k16g22s32"}, + {nvidia_key(), + WmmaOp::kQ4gsw, + 129, + 14336, + 4096, + WmmaStorage::kTexture3D, + kQ + "bx_t128x128k32g42s32f32c"}, + {nvidia_key(), + WmmaOp::kQ4gsw, + 2047, + 2048, + 2048, + WmmaStorage::kTexture3D, + kQ + "t256x128k16g42s32bt"}, + {nvidia_key(), + WmmaOp::kQ4gsw, + 2049, + 2048, + 2048, + WmmaStorage::kTexture3D, + kQ + "t128x128k16g22s32"}, + // Below kWmmaPartialTileM. + {xe2_key(), WmmaOp::kQ4gsw, 127, 2048, 2048, WmmaStorage::kTexture3D, ""}, + {nvidia_key(), + WmmaOp::kQ4gsw, + 127, + 2048, + 2048, + WmmaStorage::kTexture3D, + ""}, + // Buffer IO and 8da4w keep the alignment requirement. + {xe2_key(), WmmaOp::kQ4gsw, 511, 2048, 2048, WmmaStorage::kBuffer, ""}, + {xe2_key(), + WmmaOp::kDq8caQ4gsw, + 511, + 2048, + 2048, + WmmaStorage::kTexture3D, + ""}, + {nvidia_key(), + WmmaOp::kDq8caQ4gsw, + 511, + 2048, + 2048, + WmmaStorage::kTexture3D, + ""}, + }; + for (const Case& c : cases) { + EXPECT_EQ(selected(c.device, prefill(c.op, c.M, c.K, c.N, c.io)), c.want) + << coopmat_family_name(coopmat_family(c.device)) + << " op=" << static_cast(c.op) << " M=" << c.M << " K=" << c.K + << " N=" << c.N << " io=" << static_cast(c.io); + } +} + +TEST(CoopmatFamilyTest, Int8Layout) { + const CoopmatDeviceKey d = nvidia_key(); + WmmaShape s = + prefill(WmmaOp::kDq8caQ4gsw, 2048, 4096, 4096, WmmaStorage::kTexture3D); + const std::optional aligned = select_wmma_kernel(d, s); + ASSERT_TRUE(aligned.has_value()); + EXPECT_TRUE(aligned->rowmajor_a); + + // A row-major op keeps its kernel for unaligned M when alignment is waived. + s.int8_layout = WmmaInt8Layout::kRowMajor; + s.M = 1304; + s.ignore_alignment = true; + const std::optional forced = select_wmma_kernel(d, s); + ASSERT_TRUE(forced.has_value()); + EXPECT_EQ(forced->kernel_base, aligned->kernel_base); + + // 4h4w activations never get a row-major kernel, and vice versa. + WmmaShape h = s; + h.M = 2048; + h.ignore_alignment = false; + h.int8_layout = WmmaInt8Layout::k4H4W; + EXPECT_FALSE(select_wmma_kernel(d, h).has_value()); + WmmaShape x = + prefill(WmmaOp::kDq8caQ4gsw, 2048, 4096, 4096, WmmaStorage::kTexture3D); + x.int8_layout = WmmaInt8Layout::kRowMajor; + EXPECT_FALSE(select_wmma_kernel(xe2_key(), x).has_value()); +} + +TEST(CoopmatFamilyTest, SharedMemoryLimit) { + // The texture3d drain of the Xe2 4w tile needs 2 * 8 * 128 * 2 = 4 KiB of + // shared memory on top of the A/B tiles. + CoopmatDeviceKey d = xe2_key(); + d.max_shared_bytes = 4096; + EXPECT_EQ( + selected( + d, + prefill(WmmaOp::kQ4gsw, 2048, 4096, 4096, WmmaStorage::kTexture3D)), + ""); + EXPECT_EQ( + selected( + d, prefill(WmmaOp::kQ4gsw, 2048, 4096, 4096, WmmaStorage::kBuffer)), + kQ + "t128x128k16g44s16m8fli"); +} + +TEST(CoopmatFamilyTest, TileForKernel) { + for (const CoopmatFamily f : {CoopmatFamily::kXe2, CoopmatFamily::kNvidia}) { + for (const WmmaRow& row : coopmat_family_rows(f)) { + const std::optional tile = wmma_tile_for_kernel( + std::string(row.kernel_base) + "_texture3d_texture2d_half"); + ASSERT_TRUE(tile.has_value()) << row.kernel_base; + EXPECT_EQ(tile->m, row.tile.m) << row.kernel_base; + EXPECT_EQ(tile->n, row.tile.n) << row.kernel_base; + EXPECT_EQ(tile->wg_size(), row.tile.wg_size()) << row.kernel_base; + } + } + // The existing kernels are not family kernels. + EXPECT_FALSE( + wmma_tile_for_kernel("linear_q4gsw_coopmat_buffer_texture2d_half") + .has_value()); + EXPECT_FALSE( + wmma_tile_for_kernel("linear_q4gsw_tiled_texture3d_texture2d_half") + .has_value()); + EXPECT_FALSE(wmma_tile_for_kernel("linear_q4gsw_wmma_t128x128k16g82s16m8fli") + .has_value()); +} + +#ifdef COOPMAT_FAMILY_TEST_GLSL_DIR +std::string read_file(const std::string& path) { + std::ifstream in(path); + std::stringstream ss; + ss << in.rdbuf(); + return ss.str(); +} + +// Every kernel name a row can produce must be a variant in its yaml. +TEST(CoopmatFamilyTest, EveryRowHasShaderVariants) { + const std::string dir = + std::string(COOPMAT_FAMILY_TEST_GLSL_DIR) + "/gemm/linear/qw/"; + const std::string yamls = read_file(dir + "q4gsw/linear_q4gsw_wmma.yaml") + + read_file(dir + "q4gsw/linear_q4gsw_wmma_bx.yaml") + + read_file(dir + "dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma.yaml") + + read_file(dir + "dq8ca_q4gsw/linear_dq8ca_q4gsw_wmma_rowmajor.yaml"); + if (yamls.empty()) { + GTEST_SKIP() << "shader sources not found under " << dir; + } + const std::pair combos[] = { + {kWmmaTex3dTex2d, "_texture3d_texture2d_half"}, + {kWmmaBufTex2d, "_buffer_texture2d_half"}}; + for (const CoopmatFamily f : {CoopmatFamily::kXe2, CoopmatFamily::kNvidia}) { + for (const WmmaRow& row : coopmat_family_rows(f)) { + for (const auto& c : combos) { + if (row.storages & c.first) { + const std::string name = std::string(row.kernel_base) + c.second; + EXPECT_NE(yamls.find("NAME: " + name + "\n"), std::string::npos) + << "no shader variant " << name; + } + } + } + } +} +#endif // COOPMAT_FAMILY_TEST_GLSL_DIR + +} // namespace +} // namespace vkcompute diff --git a/backends/vulkan/test/custom_ops/targets.bzl b/backends/vulkan/test/custom_ops/targets.bzl index 2355e6a4f2a..77ff22a20c1 100644 --- a/backends/vulkan/test/custom_ops/targets.bzl +++ b/backends/vulkan/test/custom_ops/targets.bzl @@ -41,3 +41,13 @@ def define_common_targets(is_fbcode = False): "//executorch/backends/vulkan:vulkan_graph_runtime", ], ) + + runtime.cxx_test( + name = "coopmat_family_test", + srcs = ["coopmat_family_test.cpp"], + platforms = get_platforms(), + deps = [ + "//third-party/googletest:gtest_main", + "//executorch/backends/vulkan:vulkan_graph_runtime", + ], + ) diff --git a/backends/vulkan/test/ops/benchmark/cases/q4gsw_linear/test_sets/wmma_sets.cpp b/backends/vulkan/test/ops/benchmark/cases/q4gsw_linear/test_sets/wmma_sets.cpp new file mode 100644 index 00000000000..b24f111fbf2 --- /dev/null +++ b/backends/vulkan/test/ops/benchmark/cases/q4gsw_linear/test_sets/wmma_sets.cpp @@ -0,0 +1,191 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +// Correctness cases for the cooperative-matrix family kernels of linear_q4gsw +// and linear_dq8ca_q4gsw (impl/gemm/linear/qw/CoopmatFamily.cpp). The shapes +// are aligned to every family tile (M a multiple of 256, N of 128, K of 64, +// group size 64 or 128) and include the shapes that select each row (K * N +// >= 20e6 for the blocked Xe2 8da4w order, K > 8192 for the NVIDIA fp32 +// accumulation): Llama projections and some of Qwen2.5 0.5B and Gemma 2 2B. +// linear_q4gsw also runs at M that is not a multiple of the tile: a partial +// last M tile on texture IO (kWmmaPartialTileM), the existing kernels below it +// and on buffer IO. On a device without a family the same cases exercise the +// existing kernels. + +#include + +#include +#include +#include +#include +#include +#include + +namespace executorch { +namespace vulkan { +namespace prototyping { +namespace q4gsw_linear { + +namespace { + +std::vector to_float(const ValueSpec& s) { + if (s.dtype == vkapi::kFloat) { + return s.get_float_data(); + } + const auto& h = s.get_half_data(); + std::vector o(h.size()); + for (size_t i = 0; i < h.size(); ++i) { + o[i] = half_to_float(h[i]); + } + return o; +} + +// Zero-mean data with exactly representable activations: multiples of 1/16 +// in [-1, 1) (not periodic in K) and the random int4 weights and scales of +// make_case. Partial sums stay small, so the fp16-accumulating kernels (family +// and existing) stay within the tolerance; outputs are of order 1 to 5, so a +// structural error (a missing or misplaced tile, a wrong scale group) does +// not. +// For dq8ca the per-row activation scale is 1/16 so the dynamic int8 +// quantization is exact; the zero points vary per row to exercise the +// zero-point correction (quantized values stay within int8). +void set_family_data(TestCase& t, const CoopmatLinearConfig& cfg) { + auto& hin = t.inputs()[0].get_half_data(); + for (size_t i = 0; i < hin.size(); ++i) { + hin[i] = float_to_half(float(int((i * 7 + i / 61) % 32) - 16) / 16.0f); + } + if (is_dq8ca(cfg.op_name)) { + auto& hs = t.inputs()[1].get_half_data(); + std::fill(hs.begin(), hs.end(), float_to_half(0.0625f)); + auto& zp = t.inputs()[2].get_int8_data(); + for (size_t i = 0; i < zp.size(); ++i) { + zp[i] = int8_t(int(i % 17) - 8); + } + } + // fp16 accumulation error grows with K: at K = 14336 it reaches about 0.2 + // for the fp16 family kernels and 0.7 for the existing fp16 q4gsw kernel + // (llvmpipe), against outputs of order 10. + t.set_abs_tolerance(cfg.K > 8192 ? 1.0f : 0.25f); + t.set_rel_tolerance(0.02f); +} + +// Reference for both ops with 4-bit weights and no bias: +// q4gsw: out[m][n] = SUM_k a[m][k] * w[n][k] * s[k / group][n] +// dq8ca: a is first quantized per row (scale, zero point) and dequantized. +// Weights are dequantized once and the output columns are split across +// threads, so the larger shapes stay fast. +void family_reference(TestCase& tc) { + const bool dq = tc.operator_name().find("dq8ca") != std::string::npos; + const ValueSpec& in = tc.inputs()[0]; + ValueSpec& out = tc.outputs()[0]; + const auto in_sizes = in.get_tensor_sizes(); + const int64_t M = in_sizes[0]; + const int64_t K = in_sizes[1]; + const int64_t N = out.get_tensor_sizes()[1]; + const ValueSpec& w = tc.inputs()[dq ? 3 : 1]; + const std::vector scales = to_float(tc.inputs()[dq ? 5 : 2]); + const int64_t group = tc.inputs()[dq ? 6 : 3].get_int_value(); + const std::vector& w4 = w.get_uint8_data(); // [N, K/2] nibbles + + std::vector a = to_float(in); + std::vector row_scale(M, 1.0f); + if (dq) { + const std::vector in_scale = to_float(tc.inputs()[1]); + const std::vector& in_zp = tc.inputs()[2].get_int8_data(); + for (int64_t m = 0; m < M; ++m) { + const float s = in_scale[m]; + const int zp = int(in_zp[m]); + for (int64_t k = 0; k < K; ++k) { + float q = std::round(a[m * K + k] / s) + float(zp); + q = std::min(std::max(q, -128.0f), 127.0f); + a[m * K + k] = q - float(zp); + } + row_scale[m] = s; + } + } + + auto& ref = out.get_ref_float_data(); + ref.assign(M * N, 0.0f); + const int64_t num_threads = std::max( + 1, std::min(N, std::thread::hardware_concurrency())); + auto work = [&](const int64_t t) { + std::vector wn(K); + for (int64_t n = t; n < N; n += num_threads) { + for (int64_t k = 0; k < K; ++k) { + const uint8_t byte = w4[n * (K / 2) + k / 2]; + const int nib = (k & 1) ? ((byte >> 4) & 0xF) : (byte & 0xF); + wn[k] = float(nib - 8) * scales[(k / group) * N + n]; + } + for (int64_t m = 0; m < M; ++m) { + const float* am = a.data() + m * K; + float acc = 0.0f; + for (int64_t k = 0; k < K; ++k) { + acc += am[k] * wn[k]; + } + ref[m * N + n] = acc * row_scale[m]; + } + } + }; + std::vector threads; + for (int64_t t = 1; t < num_threads; ++t) { + threads.emplace_back(work, t); + } + work(0); + for (auto& th : threads) { + th.join(); + } +} + +} // namespace + +REGISTER_TEST_CASE_SET("q4gsw_linear", "wmma_correctness") { + // {M, K, N, group_size} + const std::vector shapes = { + {256, 256, 384, 128, ""}, // small, multi-workgroup + {2048, 2048, 512, 128, ""}, // long M + {256, 3072, 1024, 128, ""}, // M = 256 with N % 2048 != 0 + {256, 2048, 8192, 128, ""}, // M = 256 with N % 2048 == 0 + {256, 4096, 8192, 128, ""}, // K * N >= 20e6 + {256, 14336, 4096, 128, ""}, // K > 8192 + {256, 896, 4864, 128, ""}, // Qwen2.5 0.5B MLP + {256, 9216, 2304, 128, ""}, // Gemma 2 2B down projection, K > 8192 + {256, 2048, 2048, 64, ""}, // group size 64 + }; + // linear_q4gsw at M not a multiple of the tile. + const std::vector unaligned_shapes = { + {511, 2048, 2048, 128, ""}, // 256-row tile (NVIDIA), 255-row last tile + {300, 3072, 1024, 128, ""}, // 128-row tile, 44-row last tile + {129, 14336, 4096, 128, ""}, // K > 8192, 1-row last tile + {127, 2048, 512, 128, ""}, // below kWmmaPartialTileM + }; + std::vector cases; + auto add = [&](const CoopmatLinearConfig& cfg) { + for (const auto st : {utils::kTexture3D, utils::kBuffer}) { + TestCase t = make_case(cfg, st); + set_family_data(t, cfg); + cases.push_back(std::move(t)); + } + }; + for (const std::string op : {"linear_q4gsw", "linear_dq8ca_q4gsw"}) { + for (const auto& shape : shapes) { + add({shape.M, shape.K, shape.N, shape.group_size, op}); + } + } + for (const auto& shape : unaligned_shapes) { + add({shape.M, shape.K, shape.N, shape.group_size, "linear_q4gsw"}); + } + return { + std::move(cases), + family_reference, + flop_calc, + /*warmup_runs=*/1, + /*benchmark_runs=*/3}; +} + +} // namespace q4gsw_linear +} // namespace prototyping +} // namespace vulkan +} // namespace executorch