metal : int -> short

This commit is contained in:
Georgi Gerganov 2024-11-08 15:38:25 +02:00
parent 80b5b51bd8
commit fc1a76af11
No known key found for this signature in database
GPG Key ID: 449E073F9DC10735

View File

@ -6313,7 +6313,7 @@ kernel void kernel_mul_mm(device const uchar * src0,
simdgroup_T8x8 ma[4];
simdgroup_half8x8 mb[2];
simdgroup_half8x8 mc[8];
for (int i = 0; i < 8; i++){
for (short i = 0; i < 8; i++){
mc[i] = make_filled_simdgroup_matrix<half, 8>(0.h);
}
@ -6339,7 +6339,7 @@ kernel void kernel_mul_mm(device const uchar * src0,
threadgroup_barrier(mem_flags::mem_threadgroup);
#pragma unroll(16)
for (int i = 0; i < 16; i++) {
for (short i = 0; i < 16; i++) {
*(sa + SG_MAT_SIZE * ((tiitg / THREAD_PER_ROW / 8) \
+ (tiitg % THREAD_PER_ROW) * 16 + (i / 8) * 8) \
+ (tiitg / THREAD_PER_ROW) % 8 + (i & 7) * 8) = temp_a[i/4][i%4];
@ -6358,14 +6358,14 @@ kernel void kernel_mul_mm(device const uchar * src0,
threadgroup half * lsmb = (sb + THREAD_MAT_N * SG_MAT_SIZE * (sgitg / 2));
#pragma unroll(4)
for (int ik = 0; ik < BLOCK_SIZE_K / 8; ik++) {
for (short ik = 0; ik < BLOCK_SIZE_K / 8; ik++) {
#pragma unroll(4)
for (int i = 0; i < 4; i++) {
for (short i = 0; i < 4; i++) {
simdgroup_load(ma[i],lsma + SG_MAT_SIZE * i);
}
simdgroup_barrier(mem_flags::mem_none);
#pragma unroll(2)
for (int i = 0; i < 2; i++) {
for (short i = 0; i < 2; i++) {
simdgroup_load(mb[i],lsmb + SG_MAT_SIZE * i);
}
@ -6373,7 +6373,7 @@ kernel void kernel_mul_mm(device const uchar * src0,
lsmb += BLOCK_SIZE_N / SG_MAT_ROW * SG_MAT_SIZE;
#pragma unroll(8)
for (int i = 0; i < 8; i++){
for (short i = 0; i < 8; i++){
simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
}
}
@ -6382,7 +6382,7 @@ kernel void kernel_mul_mm(device const uchar * src0,
if ((r0 + 1) * BLOCK_SIZE_M <= ne0 && (r1 + 1) * BLOCK_SIZE_N <= ne1) {
device float * C = dst + (BLOCK_SIZE_M * r0 + 32 * (sgitg & 1)) \
+ (BLOCK_SIZE_N * r1 + 16 * (sgitg >> 1)) * ne0 + im*ne1*ne0;
for (int i = 0; i < 8; i++) {
for (short i = 0; i < 8; i++) {
// cast to f32
simdgroup_float8x8 mc_f32(1.0f);
simdgroup_multiply(mc_f32, mc[i], mc_f32);
@ -6394,7 +6394,7 @@ kernel void kernel_mul_mm(device const uchar * src0,
threadgroup_barrier(mem_flags::mem_threadgroup);
threadgroup float * temp_str = ((threadgroup float *)shared_memory) \
+ 32 * (sgitg&1) + (16 * (sgitg>>1)) * BLOCK_SIZE_M;
for (int i = 0; i < 8; i++) {
for (short i = 0; i < 8; i++) {
simdgroup_float8x8 mc_f32(1.0f);
simdgroup_multiply(mc_f32, mc[i], mc_f32);
simdgroup_store(mc_f32, temp_str + 8 * (i%4) + 8 * BLOCK_SIZE_M * (i/4), BLOCK_SIZE_M);