Arrcttacsrks commited on
Commit
b648d1d
·
verified ·
1 Parent(s): fe13945

Upload llama.cpp/ggml/src/ggml-cuda/im2col.cu with huggingface_hub

Browse files
llama.cpp/ggml/src/ggml-cuda/im2col.cu ADDED
@@ -0,0 +1,103 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #include "im2col.cuh"
2
+
3
+ template <typename T>
4
+ 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,
7
+ int s0, int s1, int p0, int p1, int d0, int d1) {
8
+ const int64_t i = threadIdx.x + blockIdx.x * blockDim.x;
9
+ if (i >= pelements) {
10
+ return;
11
+ }
12
+
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
+
19
+ const int64_t oh = blockIdx.y;
20
+ const int64_t batch = blockIdx.z / IC;
21
+ const int64_t ic = blockIdx.z % IC;
22
+
23
+ const int64_t iiw = ix * s0 + kx * d0 - p0;
24
+ const int64_t iih = oh * s1 + ky * d1 - p1;
25
+
26
+ const int64_t offset_dst =
27
+ ((batch * OH + oh) * OW + ix) * CHW +
28
+ (ic * (KW * KH) + ky * KW + kx);
29
+
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];
35
+ }
36
+ }
37
+
38
+ template <typename T>
39
+ static void im2col_cuda(const float * x, T* dst,
40
+ 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,
42
+ 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);
47
+ }
48
+
49
+ static void im2col_cuda_f16(const float * x, half * dst,
50
+ 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,
52
+ int s0,int s1,int p0,int p1,int d0,int d1, cudaStream_t stream) {
53
+
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);
55
+ }
56
+
57
+ static void im2col_cuda_f32(const float * x, float * dst,
58
+ 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,
60
+ int s0,int s1,int p0,int p1,int d0,int d1, cudaStream_t stream) {
61
+
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);
63
+ }
64
+
65
+ void ggml_cuda_op_im2col(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
66
+ const ggml_tensor * src0 = dst->src[0];
67
+ const ggml_tensor * src1 = dst->src[1];
68
+ const float * src1_d = (const float *)src1->data;
69
+ float * dst_d = (float *)dst->data;
70
+ cudaStream_t stream = ctx.stream();
71
+
72
+ GGML_ASSERT(src1->type == GGML_TYPE_F32);
73
+ GGML_ASSERT( dst->type == GGML_TYPE_F16 || dst->type == GGML_TYPE_F32);
74
+
75
+ const int32_t s0 = ((const int32_t*)(dst->op_params))[0];
76
+ const int32_t s1 = ((const int32_t*)(dst->op_params))[1];
77
+ const int32_t p0 = ((const int32_t*)(dst->op_params))[2];
78
+ const int32_t p1 = ((const int32_t*)(dst->op_params))[3];
79
+ const int32_t d0 = ((const int32_t*)(dst->op_params))[4];
80
+ const int32_t d1 = ((const int32_t*)(dst->op_params))[5];
81
+
82
+ const bool is_2D = ((const int32_t*)(dst->op_params))[6] == 1;
83
+
84
+ const int64_t IC = src1->ne[is_2D ? 2 : 1];
85
+ const int64_t IH = is_2D ? src1->ne[1] : 1;
86
+ const int64_t IW = src1->ne[0];
87
+
88
+ const int64_t KH = is_2D ? src0->ne[1] : 1;
89
+ const int64_t KW = src0->ne[0];
90
+
91
+ const int64_t OH = is_2D ? dst->ne[2] : 1;
92
+ const int64_t OW = dst->ne[1];
93
+
94
+ const size_t delta_offset = src1->nb[is_2D ? 2 : 1] / 4; // nb is byte offset, src is type float32
95
+ const int64_t batch = src1->ne[is_2D ? 3 : 2];
96
+ const size_t batch_offset = src1->nb[is_2D ? 3 : 2] / 4; // nb is byte offset, src is type float32
97
+
98
+ if(dst->type == GGML_TYPE_F16) {
99
+ 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);
100
+ } else {
101
+ 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);
102
+ }
103
+ }