Skip to content

Commit bb15126

Browse files
committed
Revert "ggml: add ops for WAN video model (cuda && cpu) (ggml-org#15669)" CUDA
1 parent 8035c93 commit bb15126

6 files changed

Lines changed: 67 additions & 225 deletions

File tree

ggml/src/ggml-cuda/getrows.cu

Lines changed: 35 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,6 @@
22
#include "dequantize.cuh"
33
#include "convert.cuh"
44

5-
#define MAX_GRIDDIM_Y 65535
6-
75
template<int qk, int qr, dequantize_kernel_t dequantize_kernel, typename dst_t>
86
static __global__ void k_get_rows(
97
const void * __restrict__ src0, const int32_t * __restrict__ src1, dst_t * __restrict__ dst,
@@ -13,29 +11,32 @@ static __global__ void k_get_rows(
1311
/*const size_t nb00,*/ const size_t nb01, const size_t nb02, const size_t nb03,
1412
const size_t s10, const size_t s11, const size_t s12/*, const size_t s13*/) {
1513

16-
for (int64_t i00 = 2*(blockIdx.y*blockDim.x + threadIdx.x); i00 < ne00; i00 += gridDim.y*blockDim.x) {
17-
// The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher.
18-
const int i10 = blockIdx.x;
19-
const int i11 = blockIdx.z / ne12;
20-
const int i12 = blockIdx.z % ne12;
14+
// The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher.
15+
const int i00 = (blockIdx.y * blockDim.x + threadIdx.x)*2;
16+
const int i10 = blockIdx.x;
17+
const int i11 = blockIdx.z / ne12;
18+
const int i12 = blockIdx.z % ne12;
2119

22-
const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
20+
if (i00 >= ne00) {
21+
return;
22+
}
2323

24-
dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3;
25-
const void * src0_row = (const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03;
24+
const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
2625

27-
const int ib = i00/qk; // block index
28-
const int iqs = (i00%qk)/qr; // quant index
29-
const int iybs = i00 - i00%qk; // dst block start index
30-
const int y_offset = qr == 1 ? 1 : qk/2;
26+
dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3;
27+
const void * src0_row = (const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03;
3128

32-
// dequantize
33-
float2 v;
34-
dequantize_kernel(src0_row, ib, iqs, v);
29+
const int ib = i00/qk; // block index
30+
const int iqs = (i00%qk)/qr; // quant index
31+
const int iybs = i00 - i00%qk; // dst block start index
32+
const int y_offset = qr == 1 ? 1 : qk/2;
3533

36-
dst_row[iybs + iqs + 0] = ggml_cuda_cast<dst_t>(v.x);
37-
dst_row[iybs + iqs + y_offset] = ggml_cuda_cast<dst_t>(v.y);
38-
}
34+
// dequantize
35+
float2 v;
36+
dequantize_kernel(src0_row, ib, iqs, v);
37+
38+
dst_row[iybs + iqs + 0] = ggml_cuda_cast<dst_t>(v.x);
39+
dst_row[iybs + iqs + y_offset] = ggml_cuda_cast<dst_t>(v.y);
3940
}
4041

4142
template<typename src0_t, typename dst_t>
@@ -47,23 +48,22 @@ static __global__ void k_get_rows_float(
4748
/*const size_t nb00,*/ const size_t nb01, const size_t nb02, const size_t nb03,
4849
const size_t s10, const size_t s11, const size_t s12/*, const size_t s13*/) {
4950

50-
for (int64_t i00 = blockIdx.y*blockDim.x + threadIdx.x; i00 < ne00; i00 += gridDim.y*blockDim.x) {
51-
// The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher.
52-
const int i10 = blockIdx.x;
53-
const int i11 = blockIdx.z / ne12;
54-
const int i12 = blockIdx.z % ne12;
51+
// The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher.
52+
const int i00 = blockIdx.y * blockDim.x + threadIdx.x;
53+
const int i10 = blockIdx.x;
54+
const int i11 = blockIdx.z / ne12;
55+
const int i12 = blockIdx.z % ne12;
5556

56-
if (i00 >= ne00) {
57-
return;
58-
}
57+
if (i00 >= ne00) {
58+
return;
59+
}
5960

60-
const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
61+
const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
6162

62-
dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3;
63-
const src0_t * src0_row = (const src0_t *)((const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03);
63+
dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3;
64+
const src0_t * src0_row = (const src0_t *)((const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03);
6465

65-
dst_row[i00] = ggml_cuda_cast<dst_t>(src0_row[i00]);
66-
}
66+
dst_row[i00] = ggml_cuda_cast<dst_t>(src0_row[i00]);
6767
}
6868

6969
template<typename grad_t, typename dst_t>
@@ -98,7 +98,7 @@ static void get_rows_cuda_q(
9898
cudaStream_t stream) {
9999
const dim3 block_dims(CUDA_GET_ROWS_BLOCK_SIZE, 1, 1);
100100
const int block_num_y = (ne00 + 2*CUDA_GET_ROWS_BLOCK_SIZE - 1) / (2*CUDA_GET_ROWS_BLOCK_SIZE);
101-
const dim3 block_nums(ne10, MIN(block_num_y, MAX_GRIDDIM_Y), ne11*ne12);
101+
const dim3 block_nums(ne10, block_num_y, ne11*ne12);
102102

103103
// strides in elements
104104
// const size_t s0 = nb0 / sizeof(dst_t);
@@ -131,7 +131,7 @@ static void get_rows_cuda_float(
131131
cudaStream_t stream) {
132132
const dim3 block_dims(CUDA_GET_ROWS_BLOCK_SIZE, 1, 1);
133133
const int block_num_y = (ne00 + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE;
134-
const dim3 block_nums(ne10, MIN(block_num_y, MAX_GRIDDIM_Y), ne11*ne12);
134+
const dim3 block_nums(ne10, block_num_y, ne11*ne12);
135135

136136
// strides in elements
137137
// const size_t s0 = nb0 / sizeof(dst_t);

ggml/src/ggml-cuda/ggml-cuda.cu

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2465,9 +2465,6 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
24652465
case GGML_OP_IM2COL:
24662466
ggml_cuda_op_im2col(ctx, dst);
24672467
break;
2468-
case GGML_OP_IM2COL_3D:
2469-
ggml_cuda_op_im2col_3d(ctx, dst);
2470-
break;
24712468
case GGML_OP_CONV_2D:
24722469
ggml_cuda_op_conv2d(ctx, dst);
24732470
break;
@@ -3575,7 +3572,6 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
35753572
return op->src[0]->nb[0] == ggml_type_size(op->src[0]->type) && ggml_is_contiguous_2(op->src[0]);
35763573
}
35773574
case GGML_OP_IM2COL:
3578-
case GGML_OP_IM2COL_3D:
35793575
case GGML_OP_CONV_2D:
35803576
case GGML_OP_CONV_2D_DW:
35813577
case GGML_OP_CONV_TRANSPOSE_2D:

ggml/src/ggml-cuda/im2col.cu

Lines changed: 0 additions & 129 deletions
Original file line numberDiff line numberDiff line change
@@ -112,132 +112,3 @@ void ggml_cuda_op_im2col(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
112112
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);
113113
}
114114
}
115-
116-
// [N*IC, ID, IH, IW] => [N*OD, OH, OW, IC * KD * KH * KW]
117-
template <typename T>
118-
static __global__ void im2col_3d_kernel(
119-
const float * src, T * dst,
120-
int64_t N, int64_t IC, int64_t ID, int64_t IH, int64_t IW, int64_t OC,
121-
int64_t KD, int64_t KH, int64_t KW, int64_t OD, int64_t OH, int64_t OW,
122-
int64_t OH_OW, int64_t KD_KH_KW, int64_t ID_IH_IW, int64_t KH_KW, int64_t IH_IW, int64_t IC_ID_IH_IW,
123-
int64_t IC_KD_KH_KW, int64_t OW_KD_KH_KW, int64_t OD_OH_OW_IC_KD_KH_KW, int64_t OH_OW_IC_KD_KH_KW,
124-
int64_t OW_IC_KD_KH_KW, int64_t N_OD_OH, int64_t OD_OH,
125-
int s0, int s1, int s2, int p0, int p1, int p2, int d0, int d1, int d2) {
126-
const int64_t i = threadIdx.x + blockIdx.x * blockDim.x;
127-
if (i >= IC_KD_KH_KW) {
128-
return;
129-
}
130-
131-
const int64_t iic = i / KD_KH_KW;
132-
const int64_t ikd = (i - iic * KD_KH_KW) / KH_KW;
133-
const int64_t ikh = (i - iic * KD_KH_KW - ikd * KH_KW) / KW;
134-
const int64_t ikw = i % KW;
135-
136-
const int64_t iow = blockIdx.y;
137-
for (int64_t iz = blockIdx.z; iz < N_OD_OH; iz+=MAX_GRIDDIM_Z) {
138-
const int64_t in = iz / OD_OH;
139-
const int64_t iod = (iz - in*OD_OH) / OH;
140-
const int64_t ioh = iz % OH;
141-
142-
const int64_t iiw = iow * s0 + ikw * d0 - p0;
143-
const int64_t iih = ioh * s1 + ikh * d1 - p1;
144-
const int64_t iid = iod * s2 + ikd * d2 - p2;
145-
146-
const int64_t offset_dst = in*OD_OH_OW_IC_KD_KH_KW + iod*OH_OW_IC_KD_KH_KW + ioh*OW_IC_KD_KH_KW + iow*IC_KD_KH_KW + iic*KD_KH_KW + ikd * KH_KW + ikh*KW + ikw;
147-
148-
if (iih < 0 || iih >= IH || iiw < 0 || iiw >= IW || iid < 0 || iid >= ID) {
149-
dst[offset_dst] = 0.0f;
150-
} else {
151-
const int64_t offset_src = in*IC_ID_IH_IW + iic*ID_IH_IW + iid*IH_IW + iih*IW + iiw;
152-
dst[offset_dst] = src[offset_src];
153-
}
154-
}
155-
}
156-
157-
// [N*IC, ID, IH, IW] => [N*OD, OH, OW, IC * KD * KH * KW]
158-
template <typename T>
159-
static void im2col_3d_cuda(const float * src, T* dst,
160-
int64_t N, int64_t IC, int64_t ID, int64_t IH, int64_t IW, int64_t OC,
161-
int64_t KD, int64_t KH, int64_t KW, int64_t OD, int64_t OH, int64_t OW,
162-
int s0, int s1, int s2, int p0, int p1, int p2, int d0, int d1, int d2, cudaStream_t stream) {
163-
const int64_t OH_OW = OH*OW;
164-
const int64_t KD_KH_KW = KD*KH*KW;
165-
const int64_t ID_IH_IW = ID*IH*IW;
166-
const int64_t KH_KW = KH*KW;
167-
const int64_t IH_IW = IH*IW;
168-
const int64_t IC_KD_KH_KW = IC*KD*KH*KW;
169-
const int64_t OW_KD_KH_KW = OW*KD*KH*KW;
170-
const int64_t N_OD_OH = N*OD*OH;
171-
const int64_t OD_OH = OD*OH;
172-
const int64_t IC_ID_IH_IW = IC*ID*IH*IW;
173-
const int64_t OD_OH_OW_IC_KD_KH_KW = OD*OH*OW*IC*KD*KH*KW;
174-
const int64_t OH_OW_IC_KD_KH_KW = OH*OW*IC*KD*KH*KW;
175-
const int64_t OW_IC_KD_KH_KW = OW*IC*KD*KH*KW;
176-
const int64_t num_blocks = (IC_KD_KH_KW + CUDA_IM2COL_BLOCK_SIZE - 1) / CUDA_IM2COL_BLOCK_SIZE;
177-
dim3 block_nums(num_blocks, OW, MIN(N_OD_OH, MAX_GRIDDIM_Z));
178-
im2col_3d_kernel<<<block_nums, MIN(IC_KD_KH_KW, CUDA_IM2COL_BLOCK_SIZE) , 0, stream>>>(src, dst, N, IC, ID, IH, IW, OC, KD, KH, KW, OD, OH, OW,
179-
OH_OW, KD_KH_KW, ID_IH_IW, KH_KW, IH_IW, IC_ID_IH_IW,
180-
IC_KD_KH_KW, OW_KD_KH_KW, OD_OH_OW_IC_KD_KH_KW,
181-
OH_OW_IC_KD_KH_KW, OW_IC_KD_KH_KW, N_OD_OH, OD_OH,
182-
s0, s1, s2, p0, p1, p2, d0, d1, d2);
183-
}
184-
185-
static void im2col_3d_cuda_f16(const float * src, half * dst,
186-
int64_t N, int64_t IC, int64_t ID, int64_t IH, int64_t IW, int64_t OC,
187-
int64_t KD, int64_t KH, int64_t KW, int64_t OD, int64_t OH, int64_t OW,
188-
int s0, int s1, int s2, int p0, int p1, int p2, int d0, int d1, int d2, cudaStream_t stream) {
189-
190-
im2col_3d_cuda<half>(src, dst, N, IC, ID, IH, IW, OC, KD, KH, KW, OD, OH, OW, s0, s1, s2, p0, p1, p2, d0, d1, d2, stream);
191-
}
192-
193-
static void im2col_3d_cuda_f32(const float * src, float * dst,
194-
int64_t N, int64_t IC, int64_t ID, int64_t IH, int64_t IW, int64_t OC,
195-
int64_t KD, int64_t KH, int64_t KW, int64_t OD, int64_t OH, int64_t OW,
196-
int s0, int s1, int s2, int p0, int p1, int p2, int d0, int d1, int d2, cudaStream_t stream) {
197-
198-
im2col_3d_cuda<float>(src, dst, N, IC, ID, IH, IW, OC, KD, KH, KW, OD, OH, OW, s0, s1, s2, p0, p1, p2, d0, d1, d2, stream);
199-
}
200-
201-
void ggml_cuda_op_im2col_3d(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
202-
const ggml_tensor * src0 = dst->src[0];
203-
const ggml_tensor * src1 = dst->src[1];
204-
const float * src1_d = (const float *)src1->data;
205-
float * dst_d = (float *)dst->data;
206-
cudaStream_t stream = ctx.stream();
207-
208-
GGML_ASSERT(src1->type == GGML_TYPE_F32);
209-
GGML_ASSERT( dst->type == GGML_TYPE_F16 || dst->type == GGML_TYPE_F32);
210-
211-
GGML_TENSOR_BINARY_OP_LOCALS
212-
213-
const int32_t s0 = ((const int32_t *)(dst->op_params))[0];
214-
const int32_t s1 = ((const int32_t *)(dst->op_params))[1];
215-
const int32_t s2 = ((const int32_t *)(dst->op_params))[2];
216-
const int32_t p0 = ((const int32_t *)(dst->op_params))[3];
217-
const int32_t p1 = ((const int32_t *)(dst->op_params))[4];
218-
const int32_t p2 = ((const int32_t *)(dst->op_params))[5];
219-
const int32_t d0 = ((const int32_t *)(dst->op_params))[6];
220-
const int32_t d1 = ((const int32_t *)(dst->op_params))[7];
221-
const int32_t d2 = ((const int32_t *)(dst->op_params))[8];
222-
const int32_t IC = ((const int32_t *)(dst->op_params))[9];
223-
224-
const int64_t N = ne13 / IC;
225-
const int64_t ID = ne12;
226-
const int64_t IH = ne11;
227-
const int64_t IW = ne10;
228-
229-
const int64_t OC = ne03 / IC;
230-
const int64_t KD = ne02;
231-
const int64_t KH = ne01;
232-
const int64_t KW = ne00;
233-
234-
const int64_t OD = ne3 / N;
235-
const int64_t OH = ne2;
236-
const int64_t OW = ne1;
237-
238-
if(dst->type == GGML_TYPE_F16) {
239-
im2col_3d_cuda_f16(src1_d, (half *) dst_d, N, IC, ID, IH, IW, OC, KD, KH, KW, OD, OH, OW, s0, s1, s2, p0, p1, p2, d0, d1, d2, stream);
240-
} else {
241-
im2col_3d_cuda_f32(src1_d, (float *) dst_d, N, IC, ID, IH, IW, OC, KD, KH, KW, OD, OH, OW, s0, s1, s2, p0, p1, p2, d0, d1, d2, stream);
242-
}
243-
}

ggml/src/ggml-cuda/im2col.cuh

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,4 +3,3 @@
33
#define CUDA_IM2COL_BLOCK_SIZE 256
44

55
void ggml_cuda_op_im2col(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
6-
void ggml_cuda_op_im2col_3d(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

ggml/src/ggml-cuda/pad.cu

Lines changed: 23 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -1,50 +1,36 @@
11
#include "pad.cuh"
22

3-
static __global__ void pad_f32(const float * src, float * dst,
4-
const int lp0, const int rp0, const int lp1, const int rp1,
5-
const int lp2, const int rp2, const int lp3, const int rp3,
6-
const int ne0, const int ne1, const int ne2, const int ne3) {
7-
// blockIdx.z: i3*ne2+i2
8-
// blockIdx.y: i1
9-
// blockIDx.x: i0 / CUDA_PAD_BLOCK_SIZE
10-
// gridDim.y: ne1
11-
int i0 = threadIdx.x + blockIdx.x * blockDim.x;
12-
int i1 = blockIdx.y;
13-
int i2 = blockIdx.z % ne2;
14-
int i3 = blockIdx.z / ne2;
15-
if (i0 >= ne0 || i1 >= ne1 || i2 >= ne2 || i3 >= ne3) {
3+
static __global__ void pad_f32(const float * x, float * dst, const int ne0, const int ne00, const int ne01, const int ne02, const int ne03) {
4+
// blockIdx.z: idx of ne2*ne3, aka ne02*ne03
5+
// blockIdx.y: idx of ne1
6+
// blockIDx.x: idx of ne0 / BLOCK_SIZE
7+
int nidx = threadIdx.x + blockIdx.x * blockDim.x;
8+
if (nidx >= ne0) {
169
return;
1710
}
1811

1912
// operation
20-
const int64_t dst_idx = i3*(ne0*ne1*ne2) + i2*(ne0*ne1) + i1*ne0 + i0;
21-
if ((i0 >= lp0 && i0 < ne0 - rp0) &&
22-
(i1 >= lp1 && i1 < ne1 - rp1) &&
23-
(i2 >= lp2 && i2 < ne2 - rp2) &&
24-
(i3 >= lp3 && i3 < ne3 - rp3)) {
25-
const int64_t i00 = i0 - lp0;
26-
const int64_t i01 = i1 - lp1;
27-
const int64_t i02 = i2 - lp2;
28-
const int64_t i03 = i3 - lp3;
29-
const int64_t ne02 = ne2 - lp2 - rp2;
30-
const int64_t ne01 = ne1 - lp1 - rp1;
31-
const int64_t ne00 = ne0 - lp0 - rp0;
32-
33-
const int64_t src_idx = i03*(ne00*ne01*ne02) + i02*(ne00*ne01) + i01*ne00 + i00;
34-
35-
dst[dst_idx] = src[src_idx];
13+
int offset_dst =
14+
nidx +
15+
blockIdx.y * ne0 +
16+
blockIdx.z * ne0 * gridDim.y;
17+
if (nidx < ne00 && blockIdx.y < (unsigned)ne01 && blockIdx.z < (unsigned)(ne02*ne03)) {
18+
int offset_src =
19+
nidx +
20+
blockIdx.y * ne00 +
21+
blockIdx.z * ne00 * ne01;
22+
dst[offset_dst] = x[offset_src];
3623
} else {
37-
dst[dst_idx] = 0.0f;
24+
dst[offset_dst] = 0.0f;
3825
}
3926
}
4027

41-
static void pad_f32_cuda(const float * src, float * dst,
42-
const int lp0, const int rp0, const int lp1, const int rp1,
43-
const int lp2, const int rp2, const int lp3, const int rp3,
28+
static void pad_f32_cuda(const float * x, float * dst,
29+
const int ne00, const int ne01, const int ne02, const int ne03,
4430
const int ne0, const int ne1, const int ne2, const int ne3, cudaStream_t stream) {
4531
int num_blocks = (ne0 + CUDA_PAD_BLOCK_SIZE - 1) / CUDA_PAD_BLOCK_SIZE;
4632
dim3 gridDim(num_blocks, ne1, ne2*ne3);
47-
pad_f32<<<gridDim, CUDA_PAD_BLOCK_SIZE, 0, stream>>>(src, dst, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3, ne0, ne1, ne2, ne3);
33+
pad_f32<<<gridDim, CUDA_PAD_BLOCK_SIZE, 0, stream>>>(x, dst, ne0, ne00, ne01, ne02, ne03);
4834
}
4935

5036
void ggml_cuda_op_pad(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
@@ -55,18 +41,9 @@ void ggml_cuda_op_pad(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
5541

5642
GGML_ASSERT(src0->type == GGML_TYPE_F32);
5743
GGML_ASSERT(dst->type == GGML_TYPE_F32);
58-
GGML_ASSERT(ggml_is_contiguous(src0));
59-
60-
const int32_t lp0 = ((const int32_t*)(dst->op_params))[0];
61-
const int32_t rp0 = ((const int32_t*)(dst->op_params))[1];
62-
const int32_t lp1 = ((const int32_t*)(dst->op_params))[2];
63-
const int32_t rp1 = ((const int32_t*)(dst->op_params))[3];
64-
const int32_t lp2 = ((const int32_t*)(dst->op_params))[4];
65-
const int32_t rp2 = ((const int32_t*)(dst->op_params))[5];
66-
const int32_t lp3 = ((const int32_t*)(dst->op_params))[6];
67-
const int32_t rp3 = ((const int32_t*)(dst->op_params))[7];
44+
GGML_ASSERT(src0->ne[3] == 1 && dst->ne[3] == 1); // just 3D tensors
6845

6946
pad_f32_cuda(src0_d, dst_d,
70-
lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3,
71-
dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], stream);
47+
src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3],
48+
dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], stream);
7249
}

ggml/src/ggml-cuda/scale.cu

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,18 @@
11
#include "scale.cuh"
22

3-
#define MAX_GRIDDIM_X 0x7FFFFFFF
3+
static __global__ void scale_f32(const float * x, float * dst, const float scale, const float bias, const int k) {
4+
const int i = blockDim.x*blockIdx.x + threadIdx.x;
45

5-
static __global__ void scale_f32(const float * x, float * dst, const float scale, const float bias, const int64_t nelements) {
6-
int64_t tid = (int64_t)blockIdx.x * (int64_t)blockDim.x + (int64_t)threadIdx.x;
7-
int64_t stride = (int64_t)blockDim.x * (int64_t)gridDim.x;
8-
9-
for (int64_t i = tid; i < nelements; i += stride) {
10-
dst[i] = scale * x[i] + bias;
6+
if (i >= k) {
7+
return;
118
}
9+
10+
dst[i] = scale * x[i] + bias;
1211
}
1312

14-
static void scale_f32_cuda(const float * x, float * dst, const float scale, const float bias, const int64_t nelements, cudaStream_t stream) {
15-
const int64_t num_blocks = (nelements + CUDA_SCALE_BLOCK_SIZE - 1) / CUDA_SCALE_BLOCK_SIZE;
16-
scale_f32<<<MIN(MAX_GRIDDIM_X, num_blocks), CUDA_SCALE_BLOCK_SIZE, 0, stream>>>(x, dst, scale, bias, nelements);
13+
static void scale_f32_cuda(const float * x, float * dst, const float scale, const float bias, const int k, cudaStream_t stream) {
14+
const int num_blocks = (k + CUDA_SCALE_BLOCK_SIZE - 1) / CUDA_SCALE_BLOCK_SIZE;
15+
scale_f32<<<num_blocks, CUDA_SCALE_BLOCK_SIZE, 0, stream>>>(x, dst, scale, bias, k);
1716
}
1817

1918
void ggml_cuda_op_scale(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {

0 commit comments

Comments
 (0)