// cpu version voidsoftmax(float *input, float *output, int size){ float max = -INFINITY; int i; // 找到这组数据的最大值 for (i = 0; i < size; i++) { if(input[i] > max) max = input[i]; } float sum = 0.0f; for (i = 0; i < size; i++) { // 减去最大值防止溢出 output[i] = expf(input[i]-max); sum += output[i]; } for (i = 0; i < size; i++) { output[i] /= sum; } }
// gpu version __global__ voidsoftmaxKernel(float *input, float *output, int cols, int rows){ int row = blockDim.x * blockIdx.x + threadIdx.x; if (row >= rows) return; float max = -INFINITY; int i; for (i = 0; i < cols; i++) { float tmp = input[row * cols + i]; if(tmp > max) max = tmp; } float sum = 0.f; for (i = 0; i < cols; i++) { float tmp = input[row * cols + i]; float tmp_exp = expf(tmp - max); output[row * cols + i] = tmp_exp; sum += tmp_exp; } for (i = 0; i < cols; i++) { output[row * cols + i] /= sum; } }
// cpu version voidSgemm(int M, int N, int K, float alpha, float beta, float *A, float *B, float *C){ for (int i = 0; i < M; i++) { for (int j = 0; j < N; j++) { float sum = 0.0f; for (int k = 0; k < K; k++) { sum += A[i * K + k] * B[k * N + j]; } C[i * N + j] = alpha * sum + beta * C[i * N + j]; } } }
// gpu version __global__ voidSgemmKernel(int M, int N, int K, float alpha, float beta, float *A, float *B, float *C){ int col = blockDim.x * blockIdx.x + threadIdx.x; int row = blockDim.y * blockIdx.y + threadIdx.y; if (row >= M || col >= N) return; float sum = 0.f; for (int k = 0; k < K; k++) { sum += A[row * K + k] * B[k * N + col]; } C[row * N + col] = alpha * sum + beta * C[row * N + col]; }
// share Memory template<constint BLOCK_SIZE> __global__ voidSgemmKernel(int M, int N, int K, float alpha, float beta, float *A, float *B, float *C){ int BM = BLOCK_SIZE; int BN = BLOCK_SIZE; int BK = BLOCK_SIZE; int bx = blockIdx.x; int by = blockIdx.y; int tx = threadIdx.x % BN; int ty = threadIdx.y / BN; // 偏移到当前线程的索引 A = &A[by * BM * K]; B = &B[bx * BN]; C = &C[by * BM * K + bx * BN]; __shared__ float As[BM * BK]; __shared__ float Bs[BK * BN]; float sum =0.f; for (int i = 0; i < K; i+= BK) { As[ty * BK + tx] = A[ty * K + tx]; Bs[ty * BN + tx] = B[ty * N + tx]; __syncthreads(); //偏移到下一步 A += BK; B += BK * N; for(int k = 0; k < BK; k++) { sum += As[ty * BK + k] * Bs[k * BN + tx]; } __syncthreads(); } C[ty * N + tx] = alpha * sum + C[ty * N + tx]; }
// 0维和3维互换 voidtranspose(float *dst, float *src, int M, int N, int K, int W){ // 原形状 M*N*K*W int ostride0 = N * K * W; int ostride1 = K * W; int ostride2 = W; // 转置后 W*N*K*M int nstride0 = M * N * K; int nstride1 = K * M; int nstride2 = M; for (int m = 0; m < M; m++) { for (int n = 0; n < N; n++) { for (int k = 0; k < K; k++) { for (int w = 0; w < W; w++) { int oldidx = m * ostride0 + n * ostride1 + k * ostride2 + w; int newidx = w * nstride0 + n * nstride1 + k * nstride2 + m; dst[newidx] = src[oldidx]; } } } } }
__global__ voidtranspose(float *dst, float *src, int M, int N, int K, int W){ int m = blockDim.x * blockIdx.x + threadIdx.x; int n = blockDim.y * blockIdx.y + threadIdx.y; // 原形状 M*N*K*W int ostride0 = N * K * W; int ostride1 = K * W; int ostride2 = W; // 转置后 W*N*K*M int nstride0 = M * N * K; int nstride1 = K * M; int nstride2 = M; for (int k = 0; k < K; k++) { for (int w = 0; w < W; w++) { int oldidx = m * ostride0 + n * ostride1 + k * ostride2 + w; int newidx = w * nstride0 + n * nstride1 + k * nstride2 + m; dst[newidx] = src[oldidx]; } } }
// cpu version voidconv2d(constfloat *input, constfloat *kernel, float *output, int iH, int iW, int kH, int kW){ int oH = iH - kH + 1; int oW = iW - kW + 1; for (int h = 0; h < oH; h++) { for (int w = 0; w < oW; w++) { float sum = 0.0f; for (int ky = 0; ky < kH; ky++) { for (int kx = 0; kx < kW; kx ++) { sum += input[(h + ky) * iW + (w + kx)] * kernel[ky * kW + kx]; } } output[oW * h + w] = sum; } } }
// gpu version __global__ voidconv2dKernel(constfloat *input, constfloat *kernel, float *output, int iH, int iW, int kH, int kW, int oH, int oW){ int h = blockDim.y * blockIdx.y + threadIdx.y; int w = blockDim.x * blockIdx.x + threadIdx.x; if(h < oH && w < oW) { float sum = 0.f; for (int ky = 0; ky < kH; ky ++) { for (int kx = 0; kx < kW; kx++) { sum += input[(h + ky) * iW + (kx + w)] * kernel[ky * kW + kx]; } } output[oW * h + w] = sum; } }