mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-19 17:24:57 +02:00
ggml-metal: FWHT kernel for metal backend (#25924)
This commit is contained in:
committed by
YiChen Lv
parent
b2d79a7e7f
commit
db0a8ce844
@@ -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>;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user