diff --git a/ggml/src/ggml-metal.m b/ggml/src/ggml-metal.m index 04ec5117f..57a0981b8 100644 --- a/ggml/src/ggml-metal.m +++ b/ggml/src/ggml-metal.m @@ -1973,7 +1973,7 @@ static void ggml_metal_encode_node( [encoder setBytes:&ne1 length:sizeof(ne1) atIndex:14]; [encoder setBytes:&r2 length:sizeof(r2) atIndex:15]; [encoder setBytes:&r3 length:sizeof(r3) atIndex:16]; - [encoder setThreadgroupMemoryLength:8192 atIndex:0]; + [encoder setThreadgroupMemoryLength:(4096 + 2048) atIndex:0]; [encoder dispatchThreadgroups:MTLSizeMake( (ne11 + 31)/32, (ne01 + 63)/64, ne12*ne13) threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; } else { int nth0 = 32; diff --git a/ggml/src/ggml-metal.metal b/ggml/src/ggml-metal.metal index 413661c8a..e0eb37c5c 100644 --- a/ggml/src/ggml-metal.metal +++ b/ggml/src/ggml-metal.metal @@ -6310,8 +6310,8 @@ kernel void kernel_mul_mm(device const uchar * src0, uint tiitg[[thread_index_in_threadgroup]], uint sgitg[[simdgroup_index_in_threadgroup]]) { - threadgroup T * sa = (threadgroup T *)(shared_memory); - threadgroup float * sb = (threadgroup float *)(shared_memory + 4096); + threadgroup T * sa = (threadgroup T *)(shared_memory); + threadgroup half * sb = (threadgroup half *)(shared_memory + 4096); const uint r0 = tgpig.y; const uint r1 = tgpig.x; @@ -6325,12 +6325,12 @@ kernel void kernel_mul_mm(device const uchar * src0, short thread_row = ((short)tiitg/THREAD_PER_ROW) < n_rows ? ((short)tiitg/THREAD_PER_ROW) : n_rows - 1; short thread_col = ((short)tiitg/THREAD_PER_COL) < n_cols ? ((short)tiitg/THREAD_PER_COL) : n_cols - 1; - simdgroup_T8x8 ma[4]; - simdgroup_float8x8 mb[2]; - simdgroup_float8x8 mc[8]; + simdgroup_T8x8 ma[4]; + simdgroup_half8x8 mb[2]; + simdgroup_half8x8 mc[8]; for (short i = 0; i < 8; i++){ - mc[i] = make_filled_simdgroup_matrix(0.f); + mc[i] = make_filled_simdgroup_matrix(0.h); } short il = (tiitg % THREAD_PER_ROW); @@ -6361,17 +6361,17 @@ kernel void kernel_mul_mm(device const uchar * src0, + (tiitg/THREAD_PER_ROW)%8 + (i&7)*8) = temp_a[i/4][i%4]; } - *(threadgroup float2x4 *)(sb + (tiitg % THREAD_PER_COL)*8*32 + 8*(tiitg/THREAD_PER_COL)) = *((device float2x4 *) y); + *(threadgroup half2x4 *)(sb + (tiitg%THREAD_PER_COL)*8*32 + 8*(tiitg/THREAD_PER_COL)) = (half2x4)(*((device float2x4 *)y)); il = (il + 2 < nl) ? il + 2 : il % 2; - x = (il < 2) ? x + (2+nl-1)/nl : x; + x = (il < 2) ? x + (2 + nl - 1)/nl : x; y += BLOCK_SIZE_K; threadgroup_barrier(mem_flags::mem_threadgroup); // load matrices from threadgroup memory and conduct outer products - threadgroup T * lsma = (sa + THREAD_MAT_M*SG_MAT_SIZE*(sgitg%2)); - threadgroup float * lsmb = (sb + THREAD_MAT_N*SG_MAT_SIZE*(sgitg/2)); + threadgroup T * lsma = (sa + THREAD_MAT_M*SG_MAT_SIZE*(sgitg%2)); + threadgroup half * lsmb = (sb + THREAD_MAT_N*SG_MAT_SIZE*(sgitg/2)); #pragma unroll(4) for (short ik = 0; ik < BLOCK_SIZE_K / 8; ik++) { @@ -6399,15 +6399,22 @@ kernel void kernel_mul_mm(device const uchar * src0, device float * C = dst + (BLOCK_SIZE_M * r0 + 32 * (sgitg & 1)) \ + (BLOCK_SIZE_N * r1 + 16 * (sgitg >> 1)) * ne0 + im*ne1*ne0; for (short i = 0; i < 8; i++) { - simdgroup_store(mc[i], C + 8 * (i%4) + 8 * ne0 * (i/4), ne0); + // cast to f32 + simdgroup_float8x8 mc_f32(1.0f); + simdgroup_multiply(mc_f32, mc[i], mc_f32); + simdgroup_store(mc_f32, C + 8 * (i%4) + 8 * ne0 * (i/4), ne0); + //simdgroup_store(mc[i], C + 8 * (i%4) + 8 * ne0 * (i/4), ne0); } } else { // block is smaller than 64x32, we should avoid writing data outside of the matrix threadgroup_barrier(mem_flags::mem_threadgroup); threadgroup float * temp_str = ((threadgroup float *) shared_memory) \ - + 32 * (sgitg&1) + (16 * (sgitg>>1))*BLOCK_SIZE_M; + + 32*(sgitg&1) + (16*(sgitg >> 1))*BLOCK_SIZE_M; for (short i = 0; i < 8; i++) { - simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*BLOCK_SIZE_M*(i/4), BLOCK_SIZE_M); + 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); + //simdgroup_store(mc[i], temp_str + 8 * (i%4) + 8 * BLOCK_SIZE_M * (i/4), BLOCK_SIZE_M); } threadgroup_barrier(mem_flags::mem_threadgroup);