ggml-webgpu: add mulmat with overlapping src0/src1 (e.g., for minimax-01) (#27321)

This commit is contained in:
Masashi Yoshimura
2026-08-19 16:29:33 +09:00
committed by GitHub
parent 0adcc3bb57
commit 5112b9738b
10 changed files with 189 additions and 92 deletions
@@ -1,3 +1,7 @@
#ifndef SRC0
#define SRC0 src0
#endif
#ifdef BYTE_HELPERS
fn get_byte(value: u32, index: u32) -> u32 {
return (value >> (index * 8)) & 0xFF;
@@ -46,7 +50,7 @@ fn load_f16_as_f32_at_src(byte_offset: u32) -> f32 {
#ifdef DECLARE_BYTE_LOADERS_SRC0
fn load_u16_at_src0(byte_offset: u32) -> u32 {
let word = src0[byte_offset / 4u];
let word = SRC0[byte_offset / 4u];
let shift = (byte_offset & 0x2u) * 8u;
return (word >> shift) & 0xFFFFu;
}
@@ -55,14 +59,14 @@ fn load_u16_at_src0(byte_offset: u32) -> u32 {
// Caller extracts the 16-bit half it needs via & 0xFFFFu or >> 16u.
// this is used in k-quants for better performance
fn load_u32_at_src0_aligned(byte_offset: u32) -> u32 {
return src0[(byte_offset & ~3u) / 4u];
return SRC0[(byte_offset & ~3u) / 4u];
}
fn load_u32_at_src0(byte_offset: u32) -> u32 {
let word_idx = byte_offset / 4u;
let shift = (byte_offset & 0x3u) * 8u;
let lo = src0[word_idx];
let hi = src0[word_idx + 1u];
let lo = SRC0[word_idx];
let hi = SRC0[word_idx + 1u];
let shifted = (lo >> shift) | (hi << (32u - shift));
return select(shifted, lo, shift == 0u);
}
@@ -73,7 +77,7 @@ fn load_f16_at_src0(byte_offset: u32) -> f16 {
}
fn load_f16_as_f32_at_src0(byte_offset: u32) -> f32 {
let word = src0[byte_offset / 4u];
let word = SRC0[byte_offset / 4u];
let shift = (byte_offset & 0x2u) * 8u;
let d_bits = (word >> shift) & 0xFFFFu;
return unpack2x16float(d_bits)[0];
@@ -1,3 +1,10 @@
#ifndef SRC0
#define SRC0 src0
#endif
#ifndef SRC1
#define SRC1 src1
#endif
#ifdef VEC
#define VEC_SIZE 4
#define SHMEM_TYPE vec4<f16>
@@ -39,7 +46,7 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3
let src0_idx = batch_offset + global_m * params.stride_01 + global_k;
let src0_val = select( // taking a slight performance hit to avoid oob
SRC0_TYPE(0.0),
src0[src0_idx/VEC_SIZE],
SRC0[src0_idx/VEC_SIZE],
global_m < params.m && global_k < params.k);
store_shmem(SHMEM_TYPE(src0_val), elem_idx);
}
@@ -57,7 +64,7 @@ fn init_shmem_src1(thread_id: u32, batch_offset: u32, offset_n: u32, k_outer: u3
let src1_idx = batch_offset + global_n * params.stride_11 + global_k;
let src1_val = select(
SRC1_TYPE(0.0),
src1[src1_idx/VEC_SIZE],
SRC1[src1_idx/VEC_SIZE],
global_n < params.n && global_k < params.k);
store_shmem(SHMEM_TYPE(src1_val), TILE_SRC0_SHMEM + elem_idx);
}
@@ -1,8 +1,12 @@
enable f16;
#define DECLARE_BYTE_LOADERS_SRC0
#include "common_decls.tmpl"
#ifdef SRC_OVERLAP
#define SRC0 merged_src
#define SRC1 merged_src
#endif
#include "common_decls.tmpl"
#include "mul_mat_decls.tmpl"
#ifdef VEC
@@ -36,11 +40,17 @@ struct MulMatParams {
broadcast3: u32
};
#ifdef SRC_OVERLAP
@group(0) @binding(0) var<storage, read_write> merged_src: array<SRC0_TYPE>;
#define DST_BINDING 1
#else
@group(0) @binding(0) var<storage, read_write> src0: array<SRC0_TYPE>; // M rows, K columns
@group(0) @binding(1) var<storage, read_write> src1: array<SRC1_TYPE>; // K rows, N columns (transposed)
@group(0) @binding(2) var<storage, read_write> dst: array<DST_TYPE>; // M rows, N columns (transposed)
#define DST_BINDING 2
#endif
@group(0) @binding(3) var<uniform> params: MulMatParams;
@group(0) @binding(DST_BINDING) var<storage, read_write> dst: array<DST_TYPE>; // M rows, N columns (transposed)
@group(0) @binding(DST_BINDING + 1) var<uniform> params: MulMatParams;
fn get_local_n(thread_id: u32) -> u32 {
return thread_id / WORKGROUP_SIZE_M;
@@ -4,6 +4,10 @@ enable subgroups;
enable chromium_experimental_subgroup_matrix;
#define DECLARE_BYTE_LOADERS_SRC0
#ifdef SRC_OVERLAP
#define SRC0 merged_src
#define SRC1 merged_src
#endif
#include "common_decls.tmpl"
#include "mul_mat_decls.tmpl"
@@ -48,11 +52,17 @@ struct MulMatParams {
};
// SRC0_TYPE and SRC1_TYPE are defined in mul_mat_decls, which is included
#ifdef SRC_OVERLAP
@group(0) @binding(0) var<storage, read_write> merged_src: array<SRC0_TYPE>;
#define DST_BINDING 1
#else
@group(0) @binding(0) var<storage, read_write> src0: array<SRC0_TYPE>; // M rows, K columns
@group(0) @binding(1) var<storage, read_write> src1: array<SRC1_TYPE>; // K rows, N columns (transposed)
@group(0) @binding(2) var<storage, read_write> dst: array<DST_TYPE>; // M rows, N columns (transposed)
#define DST_BINDING 2
#endif
@group(0) @binding(3) var<uniform> params: MulMatParams;
@group(0) @binding(DST_BINDING) var<storage, read_write> dst: array<DST_TYPE>; // M rows, N columns (transposed)
@group(0) @binding(DST_BINDING + 1) var<uniform> params: MulMatParams;
const WG_M_SG_TILE_SIZE = SUBGROUP_M * SUBGROUP_MATRIX_M * SUBGROUP_MATRIX_M_SIZE;
const WG_N_SG_TILE_SIZE = SUBGROUP_N * SUBGROUP_MATRIX_N * SUBGROUP_MATRIX_N_SIZE;
@@ -7,6 +7,11 @@ enable f16;
requires packed_4x8_integer_dot_product;
#endif
#ifdef SRC_OVERLAP
#define SRC0 merged_src
#define SRC1 merged_src
#endif
#define DECLARE_BYTE_LOADERS_SRC0
#include "common_decls.tmpl"
@@ -35,17 +40,22 @@ struct MulMatParams {
broadcast3: u32
};
#if defined(MMVQ)
@group(0) @binding(0) var<storage, read_write> src0: array<SRC0_TYPE>;
#ifdef MMVQ
@group(0) @binding(1) var<storage, read_write> src1q: array<q8_1>;
#define DST_BINDING 2
#elif defined(SRC_OVERLAP)
@group(0) @binding(0) var<storage, read_write> merged_src: array<SRC0_TYPE>;
#define DST_BINDING 1
#else
@group(0) @binding(0) var<storage, read_write> src0: array<SRC0_TYPE>;
@group(0) @binding(1) var<storage, read_write> src1: array<SRC1_TYPE>;
#define DST_BINDING 2
#endif
@group(0) @binding(2) var<storage, read_write> dst: array<f32>;
@group(0) @binding(DST_BINDING) var<storage, read_write> dst: array<f32>;
// "mul_mat_vec_acc.tmpl" requires params.k, params.m, params.stride_01
@group(0) @binding(3) var<uniform> params: MulMatParams;
@group(0) @binding(DST_BINDING + 1) var<uniform> params: MulMatParams;
// Flattened as [row][thread] to keep each row's reduction contiguous in memory.
var<workgroup> partial_sums: array<f32, OUTPUTS_PER_WG * WG_SIZE>;
@@ -1,3 +1,10 @@
#ifndef SRC0
#define SRC0 src0
#endif
#ifndef SRC1
#define SRC1 src1
#endif
#ifdef U32_DEQUANT_HELPERS
#define SRC0_TYPE u32
@@ -43,13 +50,13 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
for (var k = thread_id; k < k_vec; k += WG_SIZE) {
var x_vals: array<SRC1_TYPE, NUM_COLS>;
for (var col = 0u;col < NUM_COLS;col += 1) {
x_vals[col] = src1[src1_idx_base_vec + col * (params.stride_11 / VEC_SIZE) + k];
x_vals[col] = SRC1[src1_idx_base_vec + col * (params.stride_11 / VEC_SIZE) + k];
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
let output_row = row_base + row;
if (output_row < params.m) {
let src0_idx = (src0_batch_offset + output_row * params.stride_01) / VEC_SIZE + k;
let w = src0[src0_idx];
let w = SRC0[src0_idx];
for (var col = 0u;col < NUM_COLS;col += 1) {
acc[col][row] += inner_dot(w, x_vals[col]);
}
@@ -76,7 +83,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < ELEMS_PER_THREAD; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -116,8 +123,8 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i + 4] = f32(src1[x_base + col * params.stride_11 + i + 16]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
x_block[col][i + 4] = f32(SRC1[x_base + col * params.stride_11 + i + 16]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -160,8 +167,8 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i + 4] = f32(src1[x_base + col * params.stride_11 + i + 16]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
x_block[col][i + 4] = f32(SRC1[x_base + col * params.stride_11 + i + 16]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -205,8 +212,8 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i + 4] = f32(src1[x_base + col * params.stride_11 + i + 16]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
x_block[col][i + 4] = f32(SRC1[x_base + col * params.stride_11 + i + 16]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -253,8 +260,8 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i + 4] = f32(src1[x_base + col * params.stride_11 + i + 16]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
x_block[col][i + 4] = f32(SRC1[x_base + col * params.stride_11 + i + 16]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -302,7 +309,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < ELEMS_PER_THREAD; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -347,7 +354,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < ELEMS_PER_THREAD; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -409,10 +416,10 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, 16>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < 4u; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i + 4u] = f32(src1[x_base + col * params.stride_11 + 32u + i]);
x_block[col][i + 8u] = f32(src1[x_base + col * params.stride_11 + 64u + i]);
x_block[col][i + 12u] = f32(src1[x_base + col * params.stride_11 + 96u + i]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
x_block[col][i + 4u] = f32(SRC1[x_base + col * params.stride_11 + 32u + i]);
x_block[col][i + 8u] = f32(SRC1[x_base + col * params.stride_11 + 64u + i]);
x_block[col][i + 12u] = f32(SRC1[x_base + col * params.stride_11 + 96u + i]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -518,8 +525,8 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, 16>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < 8u; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i + 8u] = f32(src1[x_base + col * params.stride_11 + 32u + i]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
x_block[col][i + 8u] = f32(SRC1[x_base + col * params.stride_11 + 32u + i]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -610,10 +617,10 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
for (var col = 0u; col < NUM_COLS;col += 1) {
let col_base = x_base + col * params.stride_11;
for (var i = 0u; i < 4u; i++) {
x_block[col][i] = f32(src1[col_base + i]);
x_block[col][i + 4u] = f32(src1[col_base + 32u + i]);
x_block[col][i + 8u] = f32(src1[col_base + 128u + i]);
x_block[col][i + 12u] = f32(src1[col_base + 160u + i]);
x_block[col][i] = f32(SRC1[col_base + i]);
x_block[col][i + 4u] = f32(SRC1[col_base + 32u + i]);
x_block[col][i + 8u] = f32(SRC1[col_base + 128u + i]);
x_block[col][i + 12u] = f32(SRC1[col_base + 160u + i]);
}
}
@@ -713,10 +720,10 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
for (var col = 0u; col < NUM_COLS;col += 1) {
let col_base = x_base + col * params.stride_11;
for (var i = 0u; i < 4u; i++) {
x_block[col][i] = f32(src1[col_base + i]);
x_block[col][i + 4u] = f32(src1[col_base + 32u + i]);
x_block[col][i + 8u] = f32(src1[col_base + 128u + i]);
x_block[col][i + 12u] = f32(src1[col_base + 160u + i]);
x_block[col][i] = f32(SRC1[col_base + i]);
x_block[col][i + 4u] = f32(SRC1[col_base + 32u + i]);
x_block[col][i + 8u] = f32(SRC1[col_base + 128u + i]);
x_block[col][i + 12u] = f32(SRC1[col_base + 160u + i]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -823,10 +830,10 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
for (var col = 0u; col < NUM_COLS;col += 1) {
let col_base = x_base + col * params.stride_11;
for (var l = 0u; l < 4u; l++) {
x_block[col][l] = f32(src1[col_base + l]);
x_block[col][l + 4u] = f32(src1[col_base + 32u + l]);
x_block[col][l + 8u] = f32(src1[col_base + 64u + l]);
x_block[col][l + 12u] = f32(src1[col_base + 96u + l]);
x_block[col][l] = f32(SRC1[col_base + l]);
x_block[col][l + 4u] = f32(SRC1[col_base + 32u + l]);
x_block[col][l + 8u] = f32(SRC1[col_base + 64u + l]);
x_block[col][l + 12u] = f32(SRC1[col_base + 96u + l]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -899,7 +906,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, 16>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < 16u; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -960,7 +967,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, 16>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < 16u; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -1039,7 +1046,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, 16>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < 16u; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -1101,7 +1108,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, 16>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < 16u; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -1168,7 +1175,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, 16>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < 16u; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -1234,7 +1241,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, 16>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < 16u; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -1302,7 +1309,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, 16>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < 16u; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -1367,8 +1374,8 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < ELEMS_PER_THREAD / 2u; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i + 4u] = f32(src1[x_base + col * params.stride_11 + i + 16u]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
x_block[col][i + 4u] = f32(SRC1[x_base + col * params.stride_11 + i + 16u]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -1418,7 +1425,7 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, 16>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < 16u; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -1476,8 +1483,8 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i + 4] = f32(src1[x_base + col * params.stride_11 + i + 16]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
x_block[col][i + 4] = f32(SRC1[x_base + col * params.stride_11 + i + 16]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
@@ -1521,8 +1528,8 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
for (var col = 0u; col < NUM_COLS;col += 1) {
for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]);
x_block[col][i + 8] = f32(src1[x_base + col * params.stride_11 + i + 8]);
x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
x_block[col][i + 8] = f32(SRC1[x_base + col * params.stride_11 + i + 8]);
}
}
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {