ITADN

Weird performance when using shared memory in GEMV

#12ClosedFdyCN 创建于 2023-12-11
F
FdyCNcommented
I try to optimize GEMV using shared memory to speed up I\O,theoretically speaking,GEMV with sram will have better bandwidth. BUT here comes a weird performance result. **Device: M2 Ultra 128GB** **kernel cost from: GPUEndTime and GPUStartTime** 1. Fistly, i build a xcode metal project for original GEMV(your codes) and sram GEMV(my codes),I found sram GEMV is 30% faster than original GEMV; ``` // transA = false, transB = true. // and this optimization on performance is as my wish gemv [1,2048] @ [4096,2048] **0.098ms(original) --> 0.068ms(sram)** gemv [1,2048] @ [11001,2048] **0.271ms(original) --> 0.195ms(sram)** ``` 2. Secondly, I add my sram GEMV kernel into your project(because i want to combine them into one metal lib),and call then in my other C++\OC project,then comes the strange thing: ``` // original much fast than xcode project test, and even faster than sram GEMV gemv [1,2048] @ [4096,2048] **0.040ms(original) vs. 0.047ms(sram)** gemv [1,2048] @ [11001,2048] **0.175ms(original) vs. 0.173ms(sram)** ``` 3. my kernel code is : warpPerBlock: 4, GridSize: {UP_ROUND(K, warpPerBlock), 1, 1},GroupSize: {32 * warpPerBlock, 1, 1}, ``` // ONLY support M = 1, tranA = false, transB = true now. template <typename T, int Align> void _gemv_sram_impl(device T *A [[buffer(0)]], device T *B [[buffer(1)]], device T *C [[buffer(2)]], device void *D [[buffer(3), function_constant(use_activation)]], threadgroup T *threadgroup_block [[threadgroup(0)]], constant ulong4 *matrix_offsets [[buffer(10), function_constant(batched)]], constant uint *activation_type [[buffer(13), function_constant(fused_activation)]], uint3 gid [[threadgroup_position_in_grid]], ushort warp_num [[dispatch_simdgroups_per_threadgroup]], ushort sidx [[simdgroup_index_in_threadgroup]], ushort lane_id [[thread_index_in_simdgroup]]) { if (gid.x * warp_num + sidx >= N || gid.y >= M) return; if (batched) { // TODO: Re-compute every inner loop iteration for FP64 accumulate. ulong3 offsets = matrix_offsets[gid.z].xyz; A = (device T*)((device uchar*)A + offsets[0]); B = (device T*)((device uchar*)B + offsets[1]); C = (device T*)((device uchar*)C + offsets[2]); } B += gid.x * warp_num * K; C += gid.y * N + gid.x * warp_num + sidx; T acc_sum = 0; device vec<T, Align> * Aalign = (device vec<T, Align> *)A; device vec<T, Align> * Balign = (device vec<T, Align> *)B; // move data into smem threadgroup vec<T, Align> * smem = (threadgroup vec<T, Align> *)threadgroup_block; for (uint k = sidx * 32 + lane_id; k < K / Align; k += 32 * warp_num) { smem[k] = Aalign[k]; } threadgroup_barrier(mem_flags::mem_threadgroup); for (uint k = lane_id; k < K / Align; k += 32) { device vec<T, Align> * BalignSIMD = Balign + K / Align * sidx; for (uint i = 0; i < Align; ++i) { acc_sum += smem[k][i] * BalignSIMD[k][i]; } } T all_sum = simd_sum(acc_sum); if (lane_id == 0) { device T* BWarp = B + sidx * K; for (uint k = Align * (K / Align); k < K; ++k) { all_sum += A[k] * BWarp[k]; } if (use_bias) { // not supported now... } if (fused_activation) { // not supported now... } *C = all_sum; } } ``` Question: 1. Why is different between xcode testbed and metallib call?sram GEMV basiclly the same in xcode and metallib call, but orignal GEMV is much better in metallib call 2. Is there any compiling optimization i missed in sram GEMV? 3. According to my code and the situation i described, could you give me some advise about the potential cause of the performance gap? Thank you for your help!
关闭于 2024-08-27 8 条评论