From 775e48abb21db533284e0880c0342f293901bd41 Mon Sep 17 00:00:00 2001 From: bssrdf Date: Mon, 17 Nov 2025 10:02:28 -0500 Subject: [PATCH] remove some repeated index computation; various code/comments clean up --- ggml/src/ggml-cuda/conv2d-implicit.cu | 187 ++++---------------- ggml/src/ggml-cuda/conv2d-implicit.cuh | 225 +++++-------------------- 2 files changed, 72 insertions(+), 340 deletions(-) diff --git a/ggml/src/ggml-cuda/conv2d-implicit.cu b/ggml/src/ggml-cuda/conv2d-implicit.cu index ce6c2b69d8..1ffce7a9d7 100644 --- a/ggml/src/ggml-cuda/conv2d-implicit.cu +++ b/ggml/src/ggml-cuda/conv2d-implicit.cu @@ -786,14 +786,6 @@ static __global__ void conv2d_implicit_kernel(const half * __restrict__ input, constexpr unsigned int MMA_M = 16; constexpr unsigned int MMA_N = 8; - // const unsigned int K = param.c; - // const uint inChannelOffset = param.c * param.w; - // const uint weightKOffset = param.c * param.r * param.s; - - // const unsigned int PQ = param.Ow * param.Oh; - // const unsigned int KPQ = param.k * PQ; - // const unsigned int NKPQ = param.n * KPQ; - // loop bounds, constexpr where possible allows for loop unrolling #if __CUDA_ARCH__ >= GGML_CUDA_CC_RUBIN constexpr unsigned int mma_tiles_per_warp_k = 2; @@ -817,6 +809,7 @@ static __global__ void conv2d_implicit_kernel(const half * __restrict__ input, unsigned int masks_a[A_K_STRID][2]; int64_t element_offset_a[A_K_STRID]; + int64_t element_offset_b; // calculate block/warp indices const unsigned int block_m = blockIdx.y; @@ -833,7 +826,7 @@ static __global__ void conv2d_implicit_kernel(const half * __restrict__ input, half* B_block_smem = &shmem[BM * BK]; constexpr int BUFFER_SIZE = BM * BK + BK * BN; -#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#ifdef CP_ASYNC_AVAILABLE half* SA1 = A_block_smem; half* SB1 = B_block_smem; half* SA2 = &shmem[BUFFER_SIZE]; @@ -841,6 +834,7 @@ static __global__ void conv2d_implicit_kernel(const half * __restrict__ input, #else float4 A_gmem_cache_reg[4]; float4 B_gmem_cache_reg[4]; + int offset_direction = 1; #endif // declare register storage // ptx instructions expect uint32_t registers, where each uint32_t is 2 halfs packed together @@ -883,21 +877,6 @@ static __global__ void conv2d_implicit_kernel(const half * __restrict__ input, prepareIteratorA(thread_row, masks_a, element_offset_a, param); - // for(int kk =0; kk < A_K_STRID; kk++){ - // if(element_offset_a[kk] >= 327680) - // printf("%d, %d, %d, %d, %d, %lld \n", - // threadIdx.x, threadIdx.y, blockIdx.x, blockIdx.y, blockIdx.z, - // element_offset_a[kk]); - // } - - // if(threadIdx.x == 64 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0){ - // printf("A["); - // for(int kk =0; kk < A_K_STRID; kk++) - // printf("%f,", element_offset_a[kk]); - // printf("]\n"); - // } - - // prefetch the first block tile of A,B into shared memory const half* A_block_gmem = input; @@ -905,17 +884,19 @@ static __global__ void conv2d_implicit_kernel(const half * __restrict__ input, unsigned int curC = tileMemcpySwizzleA(A_block_gmem, A_block_smem, 0, 0, masks_a, element_offset_a, thread_row, thread_col, start_k, end_k, param); - tileMemcpySwizzleB(B_block_gmem, B_block_smem, 0, 0, start_k, end_k, thread_row, thread_col, param); -#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + element_offset_b = curC; + tileMemcpySwizzleB(B_block_gmem, B_block_smem, 0, 0, curC, element_offset_b, start_k, end_k, thread_row, thread_col, param); + +#ifdef CP_ASYNC_AVAILABLE asm volatile("cp.async.commit_group;\n" ::); #endif - int offset_direction = 1; + unsigned int block_k = 0; unsigned int block_krs = 1; - // for (unsigned int block_k = 1; block_k <= num_block_tiles_k; block_k++){ int s = 0; int r = 0; -#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + +#ifdef CP_ASYNC_AVAILABLE while (block_krs < num_block_tiles_krs) { asm volatile("cp.async.wait_group %0;\n" ::"n"(0)); @@ -944,44 +925,26 @@ static __global__ void conv2d_implicit_kernel(const half * __restrict__ input, ++block_k; } - // if(threadIdx.x == 64 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0){ - // printf("B %d,%d,%d [", s, r, block_k); - // for(int kk =0; kk < A_K_STRID; kk++){ - // if(element_offset_a[kk] >= 327680) - // printf("%d, %d, %d, %d, %d, %lld, %d, %d, %d %d, %lld\n", - // threadIdx.x, threadIdx.y, blockIdx.x, blockIdx.y, blockIdx.z, - // element_offset_a[kk], r, s, block_k, next_idx, param.inc_next[next_idx]); - // } - // threadIdx.x == 64 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0){ - // printf("%f,", element_offset_a[kk]); - // printf("]\n"); - // if(block_k == num_block_tiles_k) - // break; - - // if(thread_idx == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0){ - // printf(" s = %d, r = %d, block_k = %d, next_idx = %d , %d, %d, %d \n", s, r, block_k, next_idx, - // block_krs, num_block_tiles_k, num_block_tiles_krs); - // } - - // if (block_k != num_block_tiles_k){ - if (block_krs != num_block_tiles_krs){ -#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + if (block_krs != num_block_tiles_krs) { +#ifdef CP_ASYNC_AVAILABLE curC = tileMemcpyAsyncLoadA(A_block_gmem, SA2, r, s, masks_a, element_offset_a, thread_row, thread_col, block_k * BK, start_k, end_k, curC, param); - tileMemcpyAsyncLoadB(B_block_gmem, SB2, r, s, block_k * BK, + element_offset_b = (r*param.s+s)*param.c + curC; + tileMemcpyAsyncLoadB(B_block_gmem, SB2, r, s, curC, element_offset_b, block_k * BK, start_k, end_k, thread_row, thread_col, param); asm volatile("cp.async.commit_group;\n" ::); #else curC = tileMemcpyLoadA(A_block_gmem, A_gmem_cache_reg, r, s, masks_a, element_offset_a, thread_row, thread_col, block_k * BK, start_k, end_k, curC, param); - tileMemcpyLoadB(B_block_gmem, B_gmem_cache_reg, r, s, block_k * BK, + element_offset_b = (r*param.s+s)*param.c + curC; + tileMemcpyLoadB(B_block_gmem, B_gmem_cache_reg, r, s, curC, element_offset_b, block_k * BK, start_k, end_k, thread_row, thread_col, param); #endif } -#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#ifdef CP_ASYNC_AVAILABLE half* A_warp_tile = SA1 + A_warp_tile_offset; half* B_warp_tile = SB1 + B_warp_tile_offset; #else @@ -994,11 +957,11 @@ static __global__ void conv2d_implicit_kernel(const half * __restrict__ input, // outer product between mma tiles #pragma unroll - for (unsigned int mma_k = 0; mma_k < mma_tiles_per_warp_k; mma_k++){ + for (unsigned int mma_k = 0; mma_k < mma_tiles_per_warp_k; mma_k++) { #pragma unroll - for (unsigned int mma_n = 0; mma_n < mma_tiles_per_warp_n; mma_n++){ + for (unsigned int mma_n = 0; mma_n < mma_tiles_per_warp_n; mma_n++) { #pragma unroll - for (unsigned int mma_m = 0; mma_m < mma_tiles_per_warp_m; mma_m++){ + for (unsigned int mma_m = 0; mma_m < mma_tiles_per_warp_m; mma_m++) { #if __CUDA_ARCH__ >= GGML_CUDA_CC_RUBIN asm volatile ( "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " @@ -1026,49 +989,11 @@ static __global__ void conv2d_implicit_kernel(const half * __restrict__ input, #endif } } - - // if(threadIdx.x >= 8 && threadIdx.x < 12 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0){ - // printf("A %d, %d, %d: %f, %f \n", block_krs, mma_k, threadIdx.x, - // __half2float(A_register_[1][mma_k][0]), - // __half2float(A_register_[1][mma_k][1])); - // } - // if(threadIdx.x < 4 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0){ - // printf("B %d, %d, %d: %f, %f\n", block_krs, mma_k, threadIdx.x, - // __half2float(B_register_[mma_k][1][0]), - // __half2float(B_register_[mma_k][1][1])); - // } - // if(threadIdx.x == 8 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0){ - // printf("C %d, %d, %d: %f, %f, %f, %f\n", block_krs, mma_k, threadIdx.x, - // __half2float(acc_register_[1][1][0]), - // __half2float(acc_register_[1][1][1]), - // __half2float(acc_register_[1][1][2]), - // __half2float(acc_register_[1][1][3])); - // } - - // if(threadIdx.x < 4 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0){ - // printf("A %d, %d, (%d, %d) %d: %f, %f \n", block_krs, mma_k, r, s, threadIdx.x, - // __half2float(A_register_[0][mma_k][0]), - // __half2float(A_register_[0][mma_k][1])); - // } - // if(threadIdx.x < 4 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0){ - // printf("B %d, %d, (%d, %d) %d: %f, %f\n", block_krs, mma_k, r, s, threadIdx.x, - // __half2float(B_register_[mma_k][0][0]), - // __half2float(B_register_[mma_k][0][1])); - // } - // if(threadIdx.x == 0 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0){ - // printf("C %d, %d, (%d, %d) %d: %f, %f, %f, %f\n", block_krs, mma_k, r, s, threadIdx.x, - // __half2float(acc_register_[0][0][0]), - // __half2float(acc_register_[0][0][1]), - // __half2float(acc_register_[0][0][2]), - // __half2float(acc_register_[0][0][3])); - // } - } - // if (block_k != num_block_tiles_k) if (block_krs != num_block_tiles_krs) { -#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#ifdef CP_ASYNC_AVAILABLE half *tmp = SA1; SA1 = SA2; SA2 = tmp; tmp = SB1; SB1 = SB2; SB2 = tmp; #else @@ -1085,7 +1010,7 @@ static __global__ void conv2d_implicit_kernel(const half * __restrict__ input, } -#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#ifdef CP_ASYNC_AVAILABLE asm volatile("cp.async.wait_group %0;\n" ::"n"(0)); __syncthreads(); half* A_warp_tile = SA1 + A_warp_tile_offset; @@ -1094,11 +1019,11 @@ static __global__ void conv2d_implicit_kernel(const half * __restrict__ input, ldmatrix_b(B_warp_tile, B_register_); // outer product between mma tiles #pragma unroll - for (unsigned int mma_k = 0; mma_k < mma_tiles_per_warp_k; mma_k++){ + for (unsigned int mma_k = 0; mma_k < mma_tiles_per_warp_k; mma_k++) { #pragma unroll - for (unsigned int mma_n = 0; mma_n < mma_tiles_per_warp_n; mma_n++){ + for (unsigned int mma_n = 0; mma_n < mma_tiles_per_warp_n; mma_n++) { #pragma unroll - for (unsigned int mma_m = 0; mma_m < mma_tiles_per_warp_m; mma_m++){ + for (unsigned int mma_m = 0; mma_m < mma_tiles_per_warp_m; mma_m++) { #if __CUDA_ARCH__ >= GGML_CUDA_CC_RUBIN asm volatile ( "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " @@ -1126,42 +1051,10 @@ static __global__ void conv2d_implicit_kernel(const half * __restrict__ input, #endif } } - // if(threadIdx.x < 4 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0){ - // printf("A %d, %d, (%d, %d) %d: %f, %f \n", block_krs, mma_k, r, s, threadIdx.x, - // __half2float(A_register_[0][mma_k][0]), - // __half2float(A_register_[0][mma_k][1])); - // } - // if(threadIdx.x < 4 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0){ - // printf("B %d, %d, (%d, %d) %d: %f, %f\n", block_krs, mma_k, r, s, threadIdx.x, - // __half2float(B_register_[mma_k][0][0]), - // __half2float(B_register_[mma_k][0][1])); - // } - // if(threadIdx.x == 0 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0){ - // printf("C %d, %d, (%d, %d) %d: %f, %f, %f, %f\n", block_krs, mma_k, r, s, threadIdx.x, - // __half2float(acc_register_[0][0][0]), - // __half2float(acc_register_[0][0][1]), - // __half2float(acc_register_[0][0][2]), - // __half2float(acc_register_[0][0][3])); - // } } #endif - // if(threadIdx.x == 8 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0){ - // printf(" %u, %f, %f, %f, %f\n", blockIdx.z, - // __half2float(acc_register_[1][1][0]), - // __half2float(acc_register_[1][1][1]), - // __half2float(acc_register_[1][1][2]), - // __half2float(acc_register_[1][1][3])); - // } - // if(threadIdx.x == 0 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0){ - // printf(" %u, %f, %f, %f, %f\n", blockIdx.z, - // __half2float(acc_register_[0][1][0]), - // __half2float(acc_register_[0][1][1]), - // __half2float(acc_register_[0][1][2]), - // __half2float(acc_register_[0][1][3])); - // } - // reuse smem half *smemoutput = shmem; const uint lane_id = threadIdx.x % WARPSIZE; @@ -1174,16 +1067,13 @@ static __global__ void conv2d_implicit_kernel(const half * __restrict__ input, const uint n_idx = block_m * BM + warp_m * WM + lane_id; #pragma unroll - for (int i = 0; i < 2; ++i) - { + for (int i = 0; i < 2; ++i) { const unsigned int i_offset = i * mma_tiles_per_warp_n/2; __syncthreads(); #pragma unroll - for (unsigned int mma_m = 0; mma_m < mma_tiles_per_warp_m; mma_m++) - { + for (unsigned int mma_m = 0; mma_m < mma_tiles_per_warp_m; mma_m++) { const unsigned int mma_m_offset = output_sts_addr + mma_m * MMA_M * BN / 2; - for (unsigned int mma_n = i_offset; mma_n < (i+1)*mma_tiles_per_warp_n/2; mma_n++) - { + for (unsigned int mma_n = i_offset; mma_n < (i+1)*mma_tiles_per_warp_n/2; mma_n++) { uint32_t (®_)[2] = reinterpret_cast(acc_register_[mma_m][mma_n]); uint idx = mma_m_offset + (mma_n - i_offset) * MMA_N; idx = idx ^ ((idx & 0b110000000000) >> 9); @@ -1199,13 +1089,13 @@ static __global__ void conv2d_implicit_kernel(const half * __restrict__ input, const unsigned int m_i_wn = m_idx + i * WN / 2; #pragma unroll - for (int subk = 0; subk < WN / 4; ++subk){ + for (int subk = 0; subk < WN / 4; ++subk) { const uint row = m_i_wn + subk*2; uint idx = output_lds_addr + subk*2; idx = idx ^ ((idx & 0b110000000000) >> 9); idx = idx ^ ((idx & 0b1110000000) >> 4); #pragma unroll - for (int j = 0; j < 4; ++j){ + for (int j = 0; j < 4; ++j) { const uint gemm_i = n_idx + j*32; const int n = fastdiv(gemm_i, param.OHOW_fastdiv); const int col = fastmodulo(gemm_i, param.OHOW_fastdiv); @@ -1213,14 +1103,10 @@ static __global__ void conv2d_implicit_kernel(const half * __restrict__ input, half (&res_)[2] = reinterpret_cast(dst_ptr); if (n < param.n && row < param.k && col < param.PQ) { const uint outOffset = ((ksplit > 0) ? z * param.NKPQ : 0) + n * param.KPQ + row * param.PQ + col; - // if(row == 8 && col == 18) - // printf("A %u, %u, %f \n", outOffset, z, ggml_cuda_cast(res_[0])); output[outOffset] = ggml_cuda_cast(res_[0]); } if (n < param.n && row+1 < param.k && col < param.PQ) { const uint outOffset = ((ksplit > 0) ? z * param.NKPQ : 0) + n * param.KPQ + (row+1) * param.PQ + col; - // if(row+1 == 8 && col == 17) - // printf("B %u, %u, %f \n", outOffset, z, ggml_cuda_cast(res_[0])); output[outOffset] = ggml_cuda_cast(res_[1]); } } @@ -1532,13 +1418,7 @@ void ggml_cuda_op_conv2d_implicit(ggml_backend_cuda_context & ctx, ggml_tensor * const uint PD_Y = p[3]; // padding_y const uint DL_X = p[4]; // dilation_x const uint DL_Y = p[5]; // dilation_y - // const int LT = p[6]; // layout - // GGML_ASSERT(LT == 0 || LT == 1); - - // same number of input channels - // GGML_ASSERT(LT == 0 ? input->ne[0] == kernel->ne[0] : input->ne[2] == kernel->ne[2]); - // No cwhn GGML_ASSERT(p[6] == false); const uint IW = input->ne[0]; // input_w @@ -1554,13 +1434,6 @@ void ggml_cuda_op_conv2d_implicit(ggml_backend_cuda_context & ctx, ggml_tensor * int64_t pp[3] = {0}; - // const unsigned int K = param.c; -// const uint inChannelOffset = param.c * param.w; -// const uint weightKOffset = param.c * param.r * param.s; -// const unsigned int PQ = param.Ow * param.Oh; -// const unsigned int KPQ = param.k * PQ; -// const unsigned int NKPQ = param.n * KPQ; - param_t params = { B, IC, IH, IW, OC, KH, KW, ST_Y, ST_X, PD_Y, PD_X, DL_Y, DL_X, OH, OW, init_fastdiv_values(KW*IC), diff --git a/ggml/src/ggml-cuda/conv2d-implicit.cuh b/ggml/src/ggml-cuda/conv2d-implicit.cuh index ee56c80b7f..2cf03f268f 100644 --- a/ggml/src/ggml-cuda/conv2d-implicit.cuh +++ b/ggml/src/ggml-cuda/conv2d-implicit.cuh @@ -1,6 +1,11 @@ #pragma once #include "common.cuh" +constexpr unsigned int SWIZZLE_MASK_1 = 0b10000; +constexpr unsigned int SWIZZLE_BITS_1 = 4; +constexpr unsigned int SWIZZLE_MASK_2 = 0b1100; +constexpr unsigned int SWIZZLE_BITS_2 = 2; + typedef struct{ unsigned int n; //batch size unsigned int c; //number if channels @@ -24,7 +29,6 @@ typedef struct{ uint3 S_fastdiv; uint3 OHOW_fastdiv; int64_t inc_next[3]; - // unsigned int K; unsigned int inChannelOffset; unsigned int weightKOffset; unsigned int PQ; @@ -37,7 +41,6 @@ typedef struct{ /// Clears the predicates template -// __host__ __device__ void clear_mask(unsigned int masks_[][2], bool clear = true) { __device__ void clear_mask(unsigned int masks_[][2], bool clear = true) { #pragma unroll @@ -48,8 +51,7 @@ __device__ void clear_mask(unsigned int masks_[][2], bool clear = true) { } template -// __host__ __device__ void add_byte_offset(int64_t element_offset[], const int64_t offset){ -__device__ void add_byte_offset(int64_t element_offset[], const int64_t offset){ +__device__ void add_byte_offset(int64_t element_offset[], const int64_t offset) { #pragma unroll for (int s = 0; s < K_STRID; ++s) { element_offset[s] += offset; @@ -63,21 +65,14 @@ template(ptr); - - // int offset_npq = threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; const unsigned int gemm_i = blockIdx.y * TILE_ROWS + thread_row; offset_n[s] = fastdiv(gemm_i, param.OHOW_fastdiv); unsigned int npq_res = fastmodulo(gemm_i, param.OHOW_fastdiv); @@ -86,19 +81,8 @@ __device__ void prepareIteratorA(unsigned int thread_row, const int h = offset_p[s] * (int)param.u - (int) param.p; const int w = offset_q[s] * (int)param.v - (int) param.q; - // if(threadIdx.x < 32 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0) - // printf("%d, %d : %d, %d, %d, %d offset (%d, %d, %d), kele %llu Kcont %d\n ", thread_idx, s, - // // printf("[%s - %d] %d, %d : %d, %d, %d, %d\n ", __FUNCTION__, __LINE__, thread_idx, s, - // threadblock_offset.row(), thread_coord.strided(), ThreadMap::Delta::kStrided, - // offset_npq, offset_n[s], offset_p[s], offset_q[s], AccessType::kElements, - // ThreadMap::Iterations::kContiguous); - element_offset[s] = offset_n[s] * (int64_t)param.CHW + h * (int64_t)(param.inChannelOffset) + w * (int64_t)param.c; - // if(element_offset[s] >= 327680) - // printf("(%d, %d, %d, %d, %d), %d, %lld, %d, %d, %d, %d, %d, %u, %u, %u \n", - // threadIdx.x, threadIdx.y, blockIdx.x, blockIdx.y, blockIdx.z, - // s, element_offset[s], offset_n[s], offset_p[s], offset_q[s], h, w, chw, param.c * param.w, param.c); thread_row += ROW_STEP; } @@ -126,8 +110,7 @@ __device__ void prepareIteratorA(unsigned int thread_row, template __device__ void cp_async_zfill(void *ptr, void const *global_ptr, bool pred_guard = true) { -#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE - +#ifdef CP_ASYNC_AVAILABLE unsigned int smem_ptr; int src_in_bytes = pred_guard ? preload : 0; @@ -154,19 +137,16 @@ __device__ __forceinline__ void tileMemcpySwizzleB( half* __restrict__ dst, const unsigned int curR, const unsigned int curS, + const unsigned int curC, + const int64_t ki, const unsigned int start_k, const unsigned int end_k, unsigned int thread_row, const unsigned int thread_col, - // const unsigned int src_stride, param_t param -){ +) { #if __CUDA_ARCH__ >= GGML_CUDA_TURING - constexpr unsigned int SWIZZLE_MASK_1 = 0b10000; - constexpr unsigned int SWIZZLE_BITS_1 = 4; - constexpr unsigned int SWIZZLE_MASK_2 = 0b1100; - constexpr unsigned int SWIZZLE_BITS_2 = 2; constexpr unsigned int TILE_COLS = 32; float4* dst_float4 = reinterpret_cast(dst); @@ -174,39 +154,27 @@ __device__ __forceinline__ void tileMemcpySwizzleB( // # of threads is multiple of # of columns in the tile constexpr unsigned int TILE_COLS_VECTORIZED = TILE_COLS / 8; static_assert(NUM_THREADS % TILE_COLS_VECTORIZED == 0); - // flatten out 2d grid of threads into in order of increasing threadIdx.x - // const unsigned int thread_idx = threadIdx.y * blockDim.x + threadIdx.x; // assign each thread a row/column in the tile, calculate how many iterations we need // to cover the whole tile constexpr unsigned int ROW_STEP = NUM_THREADS / TILE_COLS_VECTORIZED; constexpr unsigned int NUM_ITERS = TILE_ROWS / ROW_STEP; - // unsigned int thread_row = thread_idx / TILE_COLS_VECTORIZED; - // const unsigned int thread_col = thread_idx % TILE_COLS_VECTORIZED; - - // const unsigned int ki = (curR*param.s+curS)*param.c + start_k+thread_col*8; - // const unsigned int curR = fastdiv(ki, param.SC_fastdiv); // channel offset - // const unsigned int curS = fastdiv(fastmodulo(ki, param.SC_fastdiv), param.C_fastdiv); // kernel r offset - // const unsigned int curC = fastmodulo(fastmodulo(ki, param.SC_fastdiv), param.C_fastdiv); // - const unsigned int curC = start_k+thread_col*8; - const unsigned int ki = (curR*param.s+curS)*param.c + curC; #pragma unroll - for (unsigned int i = 0; i < NUM_ITERS; i++){ + for (unsigned int i = 0; i < NUM_ITERS; i++) { // apply swizzle to the dst index const unsigned int src_index = thread_row * param.weightKOffset + ki; unsigned int dst_index = thread_row * TILE_COLS_VECTORIZED + thread_col; dst_index = dst_index ^ ((dst_index & SWIZZLE_MASK_1) >> SWIZZLE_BITS_1); dst_index = dst_index ^ ((dst_index & SWIZZLE_MASK_2) >> SWIZZLE_BITS_2); -#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE - +#ifdef CP_ASYNC_AVAILABLE cp_async_zfill((void *)(&dst_float4[dst_index]), (void const *)(&src[src_index]), thread_row + blockIdx.x * TILE_ROWS < param.k && curC < end_k); #else - if (thread_row + blockIdx.x * TILE_ROWS < param.k && curC < end_k){ + if (thread_row + blockIdx.x * TILE_ROWS < param.k && curC < end_k) { dst_float4[dst_index] = reinterpret_cast(&src[src_index])[0]; - }else{ // read 4 halves + } else { // read 4 halves dst_float4[dst_index] = make_float4(0.f, 0.f, 0.f, 0.f); } #endif @@ -217,6 +185,7 @@ __device__ __forceinline__ void tileMemcpySwizzleB( GGML_UNUSED(dst); GGML_UNUSED(curR); GGML_UNUSED(curS); + GGML_UNUSED(ki); GGML_UNUSED(start_k); GGML_UNUSED(end_k); GGML_UNUSED(thread_row); @@ -242,14 +211,9 @@ __device__ __forceinline__ unsigned int tileMemcpySwizzleA( const unsigned int start_k, const unsigned int end_k, param_t param -) -{ +) { #if __CUDA_ARCH__ >= GGML_CUDA_TURING - constexpr unsigned int SWIZZLE_MASK_1 = 0b10000; - constexpr unsigned int SWIZZLE_BITS_1 = 4; - constexpr unsigned int SWIZZLE_MASK_2 = 0b1100; - constexpr unsigned int SWIZZLE_BITS_2 = 2; constexpr unsigned int TILE_COLS = 32; float4* dst_float4 = reinterpret_cast(dst); @@ -257,42 +221,26 @@ __device__ __forceinline__ unsigned int tileMemcpySwizzleA( // # of threads is multiple of # of columns in the tile constexpr unsigned int TILE_COLS_VECTORIZED = TILE_COLS / 8; static_assert(NUM_THREADS % TILE_COLS_VECTORIZED == 0); - // flatten out 2d grid of threads into in order of increasing threadIdx.x - // const unsigned int thread_idx = threadIdx.y * blockDim.x + threadIdx.x; // assign each thread a row/column in the tile, calculate how many iterations we need // to cover the whole tile constexpr unsigned int ROW_STEP = NUM_THREADS / TILE_COLS_VECTORIZED; constexpr unsigned int NUM_ITERS = TILE_ROWS / ROW_STEP; - // unsigned int thread_row = thread_idx / TILE_COLS_VECTORIZED; - // const unsigned int thread_col = thread_idx % TILE_COLS_VECTORIZED; - // const unsigned int ki = start_k+thread_col*8; - // const unsigned int chw = param.c * param.h * param.w; - // const unsigned int curR = fastdiv(ki, param.SC_fastdiv); // channel offset - // const unsigned int curS = fastdiv(fastmodulo(ki, param.SC_fastdiv), param.C_fastdiv); // kernel r offset - // const unsigned int curC = fastmodulo(fastmodulo(ki, param.SC_fastdiv), param.C_fastdiv); // kernel r offset const unsigned int curC = start_k+thread_col*8; clear_mask(masks, curC >= end_k); #pragma unroll - for (unsigned int i = 0; i < NUM_ITERS; i++){ + for (unsigned int i = 0; i < NUM_ITERS; i++) { bool valid = (masks[i][0] & (1u << curR)) && (masks[i][1] & (1u << curS)); // apply swizzle to the dst index unsigned int dst_index = thread_row * TILE_COLS_VECTORIZED + thread_col; dst_index = dst_index ^ ((dst_index & SWIZZLE_MASK_1) >> SWIZZLE_BITS_1); dst_index = dst_index ^ ((dst_index & SWIZZLE_MASK_2) >> SWIZZLE_BITS_2); - // if(threadIdx.x == 3 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 1){ - // printf(" %u, %u, %u, %u, %lld, %d\n", i, curR, curS, curC, element_offset[i], valid?1:0); - // } - // if (valid && curC < end_k){ -#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#ifdef CP_ASYNC_AVAILABLE cp_async_zfill((void *)(&dst_float4[dst_index]), (void const *)(&src[element_offset[i]+curC]), valid); #else - if (valid){ - // if(element_offset[i] >= 327680 || element_offset[i] < 0) - // printf("%d, %d, %d, %d, %d, %d, %d, %d, %d \n", threadIdx.x, threadIdx.y, blockIdx.x, blockIdx.y, - // i, element_offset[i], curR, curS, curC); + if (valid) { dst_float4[dst_index] = reinterpret_cast(&src[element_offset[i]+curC])[0]; } else { dst_float4[dst_index] = make_float4(0.f, 0.f, 0.f, 0.f); @@ -300,29 +248,6 @@ __device__ __forceinline__ unsigned int tileMemcpySwizzleA( #endif thread_row += ROW_STEP; } - // #pragma unroll - // for (unsigned int i = 0; i < NUM_ITERS; i++){ - // unsigned int gemm_i = blockIdx.y * TILE_ROWS + thread_row; - // unsigned int n = fastdiv(gemm_i, param.OHOW_fastdiv); - // unsigned int npq_res = fastmodulo(gemm_i, param.OHOW_fastdiv); - // int posh_ori = fastdiv(npq_res, param.OW_fastdiv) * param.u - param.p; - // int posw_ori = fastmodulo(npq_res, param.OW_fastdiv) * param.v - param.q; - // // unsigned int inOffset = n * param.c * param.h * param.w; - // int curH = posh_ori + curR * param.d_h; // input h - // int curW = posw_ori + curS * param.d_w; // input w - // // apply swizzle to the dst index - // unsigned int dst_index = thread_row * TILE_COLS_VECTORIZED + thread_col; - // dst_index = dst_index ^ ((dst_index & SWIZZLE_MASK_1) >> SWIZZLE_BITS_1); - // dst_index = dst_index ^ ((dst_index & SWIZZLE_MASK_2) >> SWIZZLE_BITS_2); - // if (curH >= 0 && curW >= 0 && curW < param.w && curH < param.h && - // curR < param.r && curS < param.s && curC < param.c && n < param.n && ki < end_k){ - // const unsigned int inOffsetTmp = curH * inChannelOffset + curW * param.c + curC; - // dst_float4[dst_index] = reinterpret_cast(&src[n * chw + inOffsetTmp])[0]; - // } else{ - // dst_float4[dst_index] = make_float4(0.f, 0.f, 0.f, 0.f); - // } - // thread_row += ROW_STEP; - // } return curC; #else GGML_UNUSED(src); @@ -357,42 +282,29 @@ __device__ __forceinline__ unsigned int tileMemcpyLoadA( const unsigned int start_k, const unsigned int end_k, unsigned int oldC, - // const unsigned int inChannelOffset, param_t param -){ +) { #if __CUDA_ARCH__ >= GGML_CUDA_TURING // # of threads is multiple of # of columns in the tile constexpr unsigned int TILE_COLS_VECTORIZED = TILE_COLS / 8; static_assert(NUM_THREADS % TILE_COLS_VECTORIZED == 0); - // flatten out 2d grid of threads into in order of increasing threadIdx.x // assign each thread a row/column in the tile, calculate how many iterations we need // to cover the whole tile constexpr unsigned int ROW_STEP = NUM_THREADS / TILE_COLS_VECTORIZED; constexpr unsigned int NUM_ITERS = TILE_ROWS / ROW_STEP; - // unsigned int thread_row = thread_idx / TILE_COLS_VECTORIZED; - // const unsigned int thread_col = thread_idx % TILE_COLS_VECTORIZED; // compile time check that we provided the right amount of registers for storage static_assert(ELEMENTS_PER_THREAD == NUM_ITERS); - // const unsigned int ki = start_k+block_k+thread_col*8; - // const unsigned int chw = param.c * param.h * param.w; - - // const unsigned int curR = fastdiv(ki, param.SC_fastdiv); // channel offset - // const unsigned int curS = fastdiv(fastmodulo(ki, param.SC_fastdiv), param.C_fastdiv); // kernel r offset - // const unsigned int curC = fastmodulo(fastmodulo(ki, param.SC_fastdiv), param.C_fastdiv); // kernel r offset const unsigned int curC = start_k+block_k+thread_col*8; if (curC > oldC) clear_mask(masks, curC >= end_k); #pragma unroll - for (unsigned int i = 0; i < NUM_ITERS; i++){ + for (unsigned int i = 0; i < NUM_ITERS; i++) { bool valid = (masks[i][0] & (1u << curR)) && (masks[i][1] & (1u << curS)); - // if(threadIdx.x == 3 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 1){ - // printf(" %u, %u, %u, %u, %u, %lld, %d\n", i, curR, curS, oldC, curC, element_offset[i], valid?1:0); - // } if (valid) { dst_reg[i] = reinterpret_cast(&src[element_offset[i]+curC])[0]; } else{ @@ -435,50 +347,32 @@ __device__ __forceinline__ unsigned int tileMemcpyAsyncLoadA( const unsigned int start_k, const unsigned int end_k, unsigned int oldC, - // const unsigned int inChannelOffset, param_t param -){ -#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE - // # of threads is multiple of # of columns in the tile - constexpr unsigned int SWIZZLE_MASK_1 = 0b10000; - constexpr unsigned int SWIZZLE_BITS_1 = 4; - constexpr unsigned int SWIZZLE_MASK_2 = 0b1100; - constexpr unsigned int SWIZZLE_BITS_2 = 2; +) { +#ifdef CP_ASYNC_AVAILABLE constexpr unsigned int TILE_COLS_VECTORIZED = TILE_COLS / 8; static_assert(NUM_THREADS % TILE_COLS_VECTORIZED == 0); float4* dst_float4 = reinterpret_cast(dst); - // flatten out 2d grid of threads into in order of increasing threadIdx.x // assign each thread a row/column in the tile, calculate how many iterations we need // to cover the whole tile constexpr unsigned int ROW_STEP = NUM_THREADS / TILE_COLS_VECTORIZED; constexpr unsigned int NUM_ITERS = TILE_ROWS / ROW_STEP; constexpr unsigned int ITER_STEPS = ROW_STEP * TILE_COLS_VECTORIZED; - // unsigned int thread_row = thread_idx / TILE_COLS_VECTORIZED; - // const unsigned int thread_col = thread_idx % TILE_COLS_VECTORIZED; // compile time check that we provided the right amount of registers for storage static_assert(ELEMENTS_PER_THREAD == NUM_ITERS); - // const unsigned int ki = start_k+block_k+thread_col*8; - // const unsigned int chw = param.c * param.h * param.w; - - // const unsigned int curR = fastdiv(ki, param.SC_fastdiv); // channel offset - // const unsigned int curS = fastdiv(fastmodulo(ki, param.SC_fastdiv), param.C_fastdiv); // kernel r offset - // const unsigned int curC = fastmodulo(fastmodulo(ki, param.SC_fastdiv), param.C_fastdiv); // kernel r offset const unsigned int curC = start_k+block_k+thread_col*8; if (curC > oldC) clear_mask(masks, curC >= end_k); unsigned int iter_idx = thread_row * TILE_COLS_VECTORIZED + thread_col; #pragma unroll - for (unsigned int i = 0; i < NUM_ITERS; i++){ - bool valid = (masks[i][0] & (1u << curR)) && (masks[i][1] & (1u << curS)); - // if(threadIdx.x == 0 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 1){ - // printf(" %u, %u, %u, %u, %u, %lld, %d\n", i, curR, curS, oldC, curC, element_offset[i], valid?1:0); - // } + for (unsigned int i = 0; i < NUM_ITERS; i++) { + bool valid = (masks[i][0] & (1u << curR)) && (masks[i][1] & (1u << curS)); unsigned int dst_index = iter_idx; dst_index = dst_index ^ ((dst_index & SWIZZLE_MASK_1) >> SWIZZLE_BITS_1); dst_index = dst_index ^ ((dst_index & SWIZZLE_MASK_2) >> SWIZZLE_BITS_2); @@ -515,57 +409,40 @@ __device__ __forceinline__ void tileMemcpyLoadB( float4 (&dst_reg)[ELEMENTS_PER_THREAD], const unsigned int curR, const unsigned int curS, + const unsigned int curC, + const int64_t ki, const unsigned int block_k, const unsigned int start_k, const unsigned int end_k, unsigned int thread_row, const unsigned int thread_col, - // const unsigned int src_stride, param_t param -){ +) { #if __CUDA_ARCH__ >= GGML_CUDA_TURING - - constexpr unsigned int SWIZZLE_MASK_1 = 0b10000; - constexpr unsigned int SWIZZLE_BITS_1 = 4; - constexpr unsigned int SWIZZLE_MASK_2 = 0b1100; - constexpr unsigned int SWIZZLE_BITS_2 = 2; - // # of threads is multiple of # of columns in the tile constexpr unsigned int TILE_COLS_VECTORIZED = TILE_COLS / 8; static_assert(NUM_THREADS % TILE_COLS_VECTORIZED == 0); - // flatten out 2d grid of threads into in order of increasing threadIdx.x - // const unsigned int thread_idx = threadIdx.y * blockDim.x + threadIdx.x; // assign each thread a row/column in the tile, calculate how many iterations we need // to cover the whole tile constexpr unsigned int ROW_STEP = NUM_THREADS / TILE_COLS_VECTORIZED; constexpr unsigned int NUM_ITERS = TILE_ROWS / ROW_STEP; - // unsigned int thread_row = thread_idx / TILE_COLS_VECTORIZED; - // const unsigned int thread_col = thread_idx % TILE_COLS_VECTORIZED; // compile time check that we provided the right amount of registers for storage static_assert(ELEMENTS_PER_THREAD == NUM_ITERS); - // const unsigned int curR = fastdiv(ki, param.SC_fastdiv); // channel offset - // const unsigned int curS = fastdiv(fastmodulo(ki, param.SC_fastdiv), param.C_fastdiv); // kernel r offset - // const unsigned int curC = fastmodulo(fastmodulo(ki, param.SC_fastdiv), param.C_fastdiv); // - const unsigned int curC = start_k+block_k+thread_col*8; - const unsigned int ki = (curR*param.s+curS)*param.c + curC; - unsigned int iter_idx = thread_row * param.weightKOffset + ki; unsigned int krow_idx = thread_row + blockIdx.x * TILE_ROWS; const int ITER_STEPS = ROW_STEP * param.weightKOffset; #pragma unroll - for (unsigned int i = 0; i < NUM_ITERS; i++){ - // const unsigned int src_index = thread_row * param.weightKOffset + ki; + for (unsigned int i = 0; i < NUM_ITERS; i++) { const unsigned int src_index = iter_idx; - // if (thread_row + blockIdx.x * TILE_ROWS < param.k && curC < end_k){ - if (krow_idx < param.k && curC < end_k){ + if (krow_idx < param.k && curC < end_k) { dst_reg[i] = reinterpret_cast(&src[src_index])[0]; - }else{ // read 4 halves + } else { // read 4 halves dst_reg[i] = make_float4(0.f, 0.f, 0.f, 0.f); } krow_idx += ROW_STEP; @@ -577,6 +454,7 @@ __device__ __forceinline__ void tileMemcpyLoadB( GGML_UNUSED(block_k); GGML_UNUSED(curR); GGML_UNUSED(curS); + GGML_UNUSED(ki); GGML_UNUSED(start_k); GGML_UNUSED(end_k); GGML_UNUSED(thread_row); @@ -595,27 +473,22 @@ __device__ __forceinline__ void tileMemcpyAsyncLoadB( half *dst, const unsigned int curR, const unsigned int curS, + const unsigned int curC, + const int64_t ki, const unsigned int block_k, const unsigned int start_k, const unsigned int end_k, unsigned int thread_row, const unsigned int thread_col, param_t param -){ +) { -#if __CUDA_ARCH__ >= GGML_CUDA_AMPERE - - constexpr unsigned int SWIZZLE_MASK_1 = 0b10000; - constexpr unsigned int SWIZZLE_BITS_1 = 4; - constexpr unsigned int SWIZZLE_MASK_2 = 0b1100; - constexpr unsigned int SWIZZLE_BITS_2 = 2; +#ifdef CP_ASYNC_AVAILABLE // # of threads is multiple of # of columns in the tile constexpr unsigned int TILE_COLS_VECTORIZED = TILE_COLS / 8; static_assert(NUM_THREADS % TILE_COLS_VECTORIZED == 0); - // flatten out 2d grid of threads into in order of increasing threadIdx.x - // const unsigned int thread_idx = threadIdx.y * blockDim.x + threadIdx.x; float4* dst_float4 = reinterpret_cast(dst); // assign each thread a row/column in the tile, calculate how many iterations we need @@ -627,17 +500,13 @@ __device__ __forceinline__ void tileMemcpyAsyncLoadB( // compile time check that we provided the right amount of registers for storage static_assert(ELEMENTS_PER_THREAD == NUM_ITERS); - const unsigned int curC = start_k+block_k+thread_col*8; - const unsigned int ki = (curR*param.s+curS)*param.c + curC; - unsigned int iter_src_idx = thread_row * param.weightKOffset + ki; unsigned int iter_dst_idx = thread_row * TILE_COLS_VECTORIZED + thread_col; unsigned int krow_idx = thread_row + blockIdx.x * TILE_ROWS; const int ITER_SRC_STEPS = ROW_STEP * param.weightKOffset; #pragma unroll - for (unsigned int i = 0; i < NUM_ITERS; i++){ - // const unsigned int src_index = thread_row * param.weightKOffset + ki; + for (unsigned int i = 0; i < NUM_ITERS; i++) { const unsigned int src_index = iter_src_idx; unsigned int dst_index = iter_dst_idx; dst_index = dst_index ^ ((dst_index & SWIZZLE_MASK_1) >> SWIZZLE_BITS_1); @@ -655,6 +524,7 @@ __device__ __forceinline__ void tileMemcpyAsyncLoadB( GGML_UNUSED(block_k); GGML_UNUSED(curR); GGML_UNUSED(curS); + GGML_UNUSED(ki); GGML_UNUSED(start_k); GGML_UNUSED(end_k); GGML_UNUSED(thread_row); @@ -676,14 +546,10 @@ __device__ __forceinline__ void tileMemcpySwizzleStore( half* __restrict__ dst, unsigned int thread_row, const unsigned int thread_col -) -{ +) { #if __CUDA_ARCH__ >= GGML_CUDA_TURING - constexpr unsigned int SWIZZLE_MASK_1 = 0b10000; - constexpr unsigned int SWIZZLE_BITS_1 = 4; - constexpr unsigned int SWIZZLE_MASK_2 = 0b1100; - constexpr unsigned int SWIZZLE_BITS_2 = 2; + constexpr unsigned int TILE_COLS = 32; // reinterpret input/output as float4 @@ -693,26 +559,19 @@ __device__ __forceinline__ void tileMemcpySwizzleStore( constexpr unsigned int TILE_COLS_VECTORIZED = TILE_COLS / 8; static_assert(NUM_THREADS % TILE_COLS_VECTORIZED == 0); - // flatten out 2d grid of threads into in order of increasing threadIdx.x - // const unsigned int thread_idx = threadIdx.y * blockDim.x + threadIdx.x; - // assign each thread a row/column in the tile, calculate how many iterations we need // to cover the whole tile constexpr unsigned int ROW_STEP = NUM_THREADS / TILE_COLS_VECTORIZED; constexpr unsigned int NUM_ITERS = TILE_ROWS / ROW_STEP; constexpr unsigned int ITER_STEPS = ROW_STEP * TILE_COLS_VECTORIZED; - // unsigned int thread_row = thread_idx / TILE_COLS_VECTORIZED; - // const unsigned int thread_col = thread_idx % TILE_COLS_VECTORIZED; // compile time check that we provided the right amount of registers for storage static_assert(ELEMENTS_PER_THREAD == NUM_ITERS); unsigned int iter_idx = thread_row * TILE_COLS_VECTORIZED + thread_col; #pragma unroll - for (unsigned int i = 0; i < NUM_ITERS; i++) - { + for (unsigned int i = 0; i < NUM_ITERS; i++) { // apply swizzle to the dst index - // unsigned int dst_index = thread_row * TILE_COLS_VECTORIZED + thread_col; unsigned int dst_index = iter_idx; dst_index = dst_index ^ ((dst_index & SWIZZLE_MASK_1) >> SWIZZLE_BITS_1); dst_index = dst_index ^ ((dst_index & SWIZZLE_MASK_2) >> SWIZZLE_BITS_2);