happyz
/
llama.cpp
mirror of https://github.com/ggerganov/llama.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

remove some repeated index computation; various code/comments clean up

This commit is contained in:
bssrdf 2025-11-17 10:02:28 -05:00
parent f54cd74ed0
commit 775e48abb2
2 changed files with 72 additions and 340 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

187
ggml/src/ggml-cuda/conv2d-implicit.cu
View File

@ -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<BM, BK, A_K_STRID, ROW_STEP>(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<BM, NUM_THREADS>(A_block_gmem, A_block_smem, 0, 0, masks_a, element_offset_a,
thread_row, thread_col, start_k, end_k, param);
tileMemcpySwizzleB<BN, NUM_THREADS>(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<BN, NUM_THREADS>(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<BM, BK, NUM_THREADS, 4>(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<BN, BK, NUM_THREADS, 4>(B_block_gmem, SB2, r, s, block_k * BK,
element_offset_b = (r*param.s+s)*param.c + curC;
tileMemcpyAsyncLoadB<BN, BK, NUM_THREADS, 4>(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<BM, BK, NUM_THREADS, 4>(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<BN, BK, NUM_THREADS, 4>(B_block_gmem, B_gmem_cache_reg, r, s, block_k * BK,
element_offset_b = (r*param.s+s)*param.c + curC;
tileMemcpyLoadB<BN, BK, NUM_THREADS, 4>(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<mma_tiles_per_warp_k, mma_tiles_per_warp_n, BK>(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 (&reg_)[2] = reinterpret_cast<uint32_t(&)[2]>(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<half(&)[2]>(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<float>(res_[0]));
output[outOffset] = ggml_cuda_cast<T>(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<float>(res_[0]));
output[outOffset] = ggml_cuda_cast<T>(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),

225
ggml/src/ggml-cuda/conv2d-implicit.cuh
View File

@ -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<const unsigned int K_STRID>
// __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<const unsigned int K_STRID>
// __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<const unsigned int TILE_ROWS,
__device__ void prepareIteratorA(unsigned int thread_row,
unsigned int masks[][2],
int64_t element_offset[],
const param_t param){
const param_t param) {
int offset_n[A_K_STRID];
int offset_p[A_K_STRID];
int offset_q[A_K_STRID];
// constexpr unsigned int TILE_COLS_VECTORIZED = TILE_COLS / 8;
// unsigned int thread_row = thread_idx / TILE_COLS_VECTORIZED;
// const unsigned int chw = param.c * param.h * param.w;
#pragma unroll
for (int s = 0; s < A_K_STRID; ++s) {
// pointer_[s] = reinterpret_cast<char const *>(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 <int preload=16>
__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<float4*>(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<const float4 *>(&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<float4*>(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<NUM_ITERS>(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<const float4 *>(&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<const float4 *>(&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<NUM_ITERS>(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<const float4 *>(&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<float4*>(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<NUM_ITERS>(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<const float4 *>(&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<float4*>(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);