ggml-metal: FWHT kernel for metal backend (#25924)

This commit is contained in:
Georgi Gerganov
2026-07-28 19:05:42 +03:00
committed by YiChen Lv
parent b2d79a7e7f
commit db0a8ce844
+63
View File
@@ -373,3 +373,66 @@ template [[host_name("kernel_snake_f16")]] kernel void kernel_snake<half>(const
#if defined(GGML_METAL_HAS_BF16)
template [[host_name("kernel_snake_bf16")]] kernel void kernel_snake<bfloat>(constant ggml_metal_kargs_snake &, device const bfloat *, device const float *, device const float *, device bfloat *, uint, uint, uint);
#endif
template<int N>
kernel void kernel_fwht_f32(
constant ggml_metal_kargs_fwht & args,
device const float * src,
device float * dst,
uint3 tgpig[[threadgroup_position_in_grid]],
ushort sgitg[[simdgroup_index_in_threadgroup]],
ushort tiisg[[thread_index_in_simdgroup]],
ushort3 ntg[[threads_per_threadgroup]]) {
constexpr int NW = N_SIMDWIDTH;
constexpr int NE = N / NW;
const float scale = 1.0f / sqrt((float) N);
const int sg_per_tg = ntg.x / NW;
const int64_t r = tgpig.x * sg_per_tg + sgitg;
if (r >= args.nrows) {
return;
}
src += r * N;
dst += r * N;
const int lane = tiisg;
float reg[NE];
for (int i = 0; i < NE; i++) {
reg[i] = src[i*NW + lane]*scale;
}
for (int i = 1; i < NW; i *= 2) {
for (int j = 0; j < NE; j++) {
const float val = reg[j];
const float val2 = simd_shuffle_xor(val, i);
reg[j] = (lane & i) == 0 ? val2 + val : val2 - val;
}
}
for (int i = NW; i < N; i *= 2) {
const int step = i / NW;
for (int j = 0; j < NE; j += (2 * step)) {
for (int k = 0; k < step; k++) {
const float x = reg[j + k ];
const float y = reg[j + k + step];
reg[j + k] = x + y;
reg[j + k + step] = x - y;
}
}
}
for (int i = 0; i < NE; i++) {
dst[i*NW + lane] = reg[i];
}
}
typedef decltype(kernel_fwht_f32<64>) kernel_fwht_t;
template [[host_name("kernel_fwht_f32_64")]] kernel kernel_fwht_t kernel_fwht_f32<64>;
template [[host_name("kernel_fwht_f32_128")]] kernel kernel_fwht_t kernel_fwht_f32<128>;
template [[host_name("kernel_fwht_f32_256")]] kernel kernel_fwht_t kernel_fwht_f32<256>;
template [[host_name("kernel_fwht_f32_512")]] kernel kernel_fwht_t kernel_fwht_f32<512>;