Skip to content

Commit 9cf458f

Browse files
ThireusThireus
committed
cuda: make im2col faster
Port from ggml-org/llama.cpp#15025
1 parent cd0d7f0 commit 9cf458f

File tree

1 file changed

+45
-36
lines changed

1 file changed

+45
-36
lines changed

‎ggml/src/ggml-cuda/im2col.cu‎

Lines changed: 45 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,65 +1,75 @@
11
#include "im2col.cuh"
22

3+
#define MIN(a, b) (a) < (b) ? (a) : (b)
4+
5+
#define MAX_GRIDDIM_Z 65535
6+
37
template <typename T>
48
static __global__ void im2col_kernel(
5-
const float * x, T * dst, int64_t batch_offset,
6-
int64_t offset_delta, int64_t IC, int64_t IW, int64_t IH, int64_t OH, int64_t OW, int64_t KW, int64_t KH, int64_t pelements, int64_t CHW,
9+
const float * x, T * dst,
10+
int64_t IC, int64_t IW, int64_t IH, int64_t OH, int64_t OW, int64_t KW, int64_t KH,
11+
int64_t IC_IH_IW, int64_t IH_IW, int64_t N_OH, int64_t KH_KW, int64_t IC_KH_KW,
712
int s0, int s1, int p0, int p1, int d0, int d1) {
813
const int64_t i = threadIdx.x + blockIdx.x * blockDim.x;
9-
if (i >= pelements) {
14+
if (i >= IC_KH_KW) {
1015
return;
1116
}
1217

13-
const int64_t ksize = OW * (KH > 1 ? KW : 1);
14-
const int64_t kx = i / ksize;
15-
const int64_t kd = kx * ksize;
16-
const int64_t ky = (i - kd) / OW;
17-
const int64_t ix = i % OW;
18+
const int64_t iic = i / (KH_KW);
19+
const int64_t rem = i - iic * KH_KW;
20+
const int64_t ikh = rem / KW;
21+
const int64_t ikw = rem - ikh * KW;
1822

19-
const int64_t oh = blockIdx.y;
20-
const int64_t batch = blockIdx.z / IC;
21-
const int64_t ic = blockIdx.z % IC;
23+
const int64_t iow = blockIdx.y;
24+
for (int64_t iz = blockIdx.z; iz < N_OH; iz+=MAX_GRIDDIM_Z) {
25+
const int64_t in = iz / OH;
26+
const int64_t ioh = iz - in * OH;
2227

23-
const int64_t iiw = ix * s0 + kx * d0 - p0;
24-
const int64_t iih = oh * s1 + ky * d1 - p1;
28+
const int64_t iiw = iow * s0 + ikw * d0 - p0;
29+
const int64_t iih = ioh * s1 + ikh * d1 - p1;
2530

26-
const int64_t offset_dst =
27-
((batch * OH + oh) * OW + ix) * CHW +
28-
(ic * (KW * KH) + ky * KW + kx);
31+
const int64_t offset_dst =
32+
((in * OH + ioh) * OW + iow) * IC_KH_KW + iic * KH_KW + ikh * KW + ikw;
2933

30-
if (iih < 0 || iih >= IH || iiw < 0 || iiw >= IW) {
31-
dst[offset_dst] = 0.0f;
32-
} else {
33-
const int64_t offset_src = ic * offset_delta + batch * batch_offset;
34-
dst[offset_dst] = x[offset_src + iih * IW + iiw];
34+
if (iih < 0 || iih >= IH || iiw < 0 || iiw >= IW) {
35+
dst[offset_dst] = 0.0f;
36+
} else {
37+
const int64_t offset_src = iic * IC_IH_IW + in * IH_IW;
38+
dst[offset_dst] = x[offset_src + iih * IW + iiw];
39+
}
3540
}
3641
}
3742

43+
// im2col: [N, IC, IH, IW] => [N, OH, OW, IC*KH*KW]
3844
template <typename T>
3945
static void im2col_cuda(const float * x, T* dst,
4046
int64_t IW, int64_t IH, int64_t OW, int64_t OH, int64_t KW, int64_t KH, int64_t IC,
41-
int64_t batch, int64_t batch_offset, int64_t offset_delta,
47+
int64_t N, int64_t IC_IH_IW, int64_t IH_IW,
4248
int s0,int s1,int p0,int p1,int d0,int d1, cudaStream_t stream) {
43-
const int parallel_elements = OW * KW * KH;
44-
const int num_blocks = (parallel_elements + CUDA_IM2COL_BLOCK_SIZE - 1) / CUDA_IM2COL_BLOCK_SIZE;
45-
dim3 block_nums(num_blocks, OH, batch * IC);
46-
im2col_kernel<<<block_nums, CUDA_IM2COL_BLOCK_SIZE, 0, stream>>>(x, dst, batch_offset, offset_delta, IC, IW, IH, OH, OW, KW, KH, parallel_elements, (IC * KH * KW), s0, s1, p0, p1, d0, d1);
49+
const int64_t IC_KH_KW = IC * KH * KW;
50+
const int64_t num_blocks = (IC_KH_KW + CUDA_IM2COL_BLOCK_SIZE - 1) / CUDA_IM2COL_BLOCK_SIZE;
51+
const int64_t N_OH = N * OH;
52+
const int64_t KH_KW = KW*KH;
53+
dim3 block_nums(num_blocks, OW, MIN(N_OH, MAX_GRIDDIM_Z));
54+
im2col_kernel<<<block_nums, MIN(IC_KH_KW, CUDA_IM2COL_BLOCK_SIZE) , 0, stream>>>(x, dst, IC, IW, IH, OH, OW, KW, KH,
55+
IC_IH_IW, IH_IW, N_OH, KH_KW, IC_KH_KW,
56+
s0, s1, p0, p1, d0, d1);
4757
}
4858

4959
static void im2col_cuda_f16(const float * x, half * dst,
5060
int64_t IW, int64_t IH, int64_t OW, int64_t OH, int64_t KW, int64_t KH, int64_t IC,
51-
int64_t batch, int64_t batch_offset, int64_t offset_delta,
61+
int64_t N, int64_t IC_IH_IW, int64_t IH_IW,
5262
int s0,int s1,int p0,int p1,int d0,int d1, cudaStream_t stream) {
5363

54-
im2col_cuda<half>(x, dst, IW, IH, OW, OH, KW, KH, IC, batch, batch_offset, offset_delta, s0, s1, p0, p1, d0, d1, stream);
64+
im2col_cuda<half>(x, dst, IW, IH, OW, OH, KW, KH, IC, N, IC_IH_IW, IH_IW, s0, s1, p0, p1, d0, d1, stream);
5565
}
5666

5767
static void im2col_cuda_f32(const float * x, float * dst,
5868
int64_t IW, int64_t IH, int64_t OW, int64_t OH, int64_t KW, int64_t KH, int64_t IC,
59-
int64_t batch, int64_t batch_offset, int64_t offset_delta,
69+
int64_t N, int64_t IC_IH_IW, int64_t IH_IW,
6070
int s0,int s1,int p0,int p1,int d0,int d1, cudaStream_t stream) {
6171

62-
im2col_cuda<float>(x, dst, IW, IH, OW, OH, KW, KH, IC, batch, batch_offset, offset_delta, s0, s1, p0, p1, d0, d1, stream);
72+
im2col_cuda<float>(x, dst, IW, IH, OW, OH, KW, KH, IC, N, IC_IH_IW, IH_IW, s0, s1, p0, p1, d0, d1, stream);
6373
}
6474

6575
void ggml_cuda_op_im2col(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
@@ -69,7 +79,6 @@ void ggml_cuda_op_im2col(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
6979
float * dst_d = (float *)dst->data;
7080
cudaStream_t stream = ctx.stream();
7181

72-
GGML_ASSERT(src0->type == GGML_TYPE_F16);
7382
GGML_ASSERT(src1->type == GGML_TYPE_F32);
7483
GGML_ASSERT( dst->type == GGML_TYPE_F16 || dst->type == GGML_TYPE_F32);
7584

@@ -92,13 +101,13 @@ void ggml_cuda_op_im2col(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
92101
const int64_t OH = is_2D ? dst->ne[2] : 1;
93102
const int64_t OW = dst->ne[1];
94103

95-
const size_t delta_offset = src1->nb[is_2D ? 2 : 1] / 4; // nb is byte offset, src is type float32
96-
const int64_t batch = src1->ne[3];
97-
const size_t batch_offset = src1->nb[3] / 4; // nb is byte offset, src is type float32
104+
const int64_t IC_IH_IW = src1->nb[is_2D ? 2 : 1] / 4; // nb is byte offset, src is type float32
105+
const int64_t N = src1->ne[is_2D ? 3 : 2];
106+
const int64_t IH_IW = src1->nb[is_2D ? 3 : 2] / 4; // nb is byte offset, src is type float32
98107

99108
if(dst->type == GGML_TYPE_F16) {
100-
im2col_cuda_f16(src1_d, (half *) dst_d, IW, IH, OW, OH, KW, KH, IC, batch, batch_offset, delta_offset, s0, s1, p0, p1, d0, d1, stream);
109+
im2col_cuda_f16(src1_d, (half *) dst_d, IW, IH, OW, OH, KW, KH, IC, N, IC_IH_IW, IH_IW, s0, s1, p0, p1, d0, d1, stream);
101110
} else {
102-
im2col_cuda_f32(src1_d, (float *) dst_d, IW, IH, OW, OH, KW, KH, IC, batch, batch_offset, delta_offset, s0, s1, p0, p1, d0, d1, stream);
111+
im2col_cuda_f32(src1_d, (float *) dst_d, IW, IH, OW, OH, KW, KH, IC, N, IC_IH_IW, IH_IW, s0, s1, p0, p1, d0, d1, stream);
103112
}
104113
}

0 commit comments

Comments
 (0)