Use 1-d kernel

Author: am17an

ggml/src/ggml-cuda/set-rows.cu

@@ -17,33 +17,34 @@ __device__ __forceinline__ void set_rows_1<float, float>(const float * src_f, fl
     *dst_f = *src_f;
 }
 
-//TODO: consolidate kernels from cpy.cu, get_rows etc to make this function generic
 template<typename src_t, typename dst_t>
 static __global__ void k_set_rows(
         const src_t * __restrict__ src0, const int64_t * __restrict__ src1, dst_t * __restrict__ dst,
         const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03,
         const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t ne13,
         const size_t nb01, const size_t nb02, const size_t nb03,
         const size_t nb10, const size_t nb11, const size_t nb12,
-        const size_t nb1, const size_t nb2, const size_t nb3,
-        const size_t src_type_size, const size_t dst_type_size) {
+        const size_t nb1, const size_t nb2, const size_t nb3) {
 
-    const int i03 = blockIdx.z / ne02;
-    const int i02 = blockIdx.z % ne02;
-    const int i01 = blockDim.x * blockIdx.x +  threadIdx.x;
-    const int i00 = blockIdx.y;
+    const int64_t i = blockDim.x * blockIdx.x + threadIdx.x;
+    const int64_t ne_total = ne00 * ne01 * ne02 * ne03;
 
-    if (i01 >= ne01) {
+    if (i >= ne_total) {
         return;
     }
 
-    const int i12 = i03 % ne12;
-    const int i11 = i02 % ne11;
-    const int i10 = i01;
+    const int64_t i03 = i / (ne00 * ne01 * ne02);
+    const int64_t i02 = (i - i03 * ne00 * ne01 * ne02) / (ne00 * ne01);
+    const int64_t i01 = (i - i03 * ne00 * ne01 * ne02 - i02 * ne00 * ne01) / ne00;
+    const int64_t i00 = i - i03 * ne00 * ne01 * ne02 - i02 * ne00 * ne01 - i01 * ne00;
+
+    const int64_t i12 = i03 % ne12;
+    const int64_t i11 = i02 % ne11;
+    const int64_t i10 = i01;
 
     const int64_t dst_row = *(src1 + i10*nb10 + i11*nb11 + i12*nb12);
 
-    const src_t * src0_row = (const src_t *)src0 + i01*nb01 + i02*nb02 + i03*nb03;
+    const src_t * src0_row = src0 + i01*nb01 + i02*nb02 + i03*nb03;
     dst_t * dst_row_ptr    = dst + dst_row*nb1 + i02*nb2 + i03*nb3;
 
     const src_t* src_elem = src0_row + i00;
@@ -59,38 +60,32 @@ static void set_rows_cuda(
         const size_t nb01, const size_t nb02, const size_t nb03,
         const size_t nb10, const size_t nb11, const size_t nb12,
         const size_t nb1, const size_t nb2, const size_t nb3,
-        const size_t src_type_size, const size_t dst_type_size,
         cudaStream_t stream) {
 
+    const int64_t ne_total = ne00 * ne01 * ne02 * ne03;
+    const int num_blocks = (ne_total + CUDA_SET_ROWS_BLOCK_SIZE - 1) / CUDA_SET_ROWS_BLOCK_SIZE;
     const dim3 block_size(CUDA_SET_ROWS_BLOCK_SIZE);
-    const dim3 grid_size(
-        (ne01 + CUDA_SET_ROWS_BLOCK_SIZE - 1)/CUDA_SET_ROWS_BLOCK_SIZE,
-        ne00,
-        ne03*ne02
-    );
-
-    const int s1 = nb01 / sizeof(src_t);
-    const int s2 = nb02 / sizeof(src_t);
-    const int s3 = nb03 / sizeof(src_t);
+    const dim3 grid_size(num_blocks);
 
-    const int s10 = nb10 / sizeof(int64_t);
-    const int s11 = nb11 / sizeof(int64_t);
-    const int s12 = nb12 / sizeof(int64_t);
 
-    const int s_dst = nb1 / sizeof(dst_t);
-    const int s_dst2 = nb2 / sizeof(dst_t);
-    const int s_dst3 = nb3 / sizeof(dst_t);
+    const int64_t s01 = nb01/sizeof(src_t);
+    const int64_t s02 = nb02/sizeof(src_t);
+    const int64_t s03 = nb03/sizeof(src_t);
+    const int64_t s10 = nb10/sizeof(int64_t);
+    const int64_t s11 = nb11/sizeof(int64_t);
+    const int64_t s12 = nb12/sizeof(int64_t);
+    const int64_t s1  = nb1/sizeof(dst_t);
+    const int64_t s2  = nb2/sizeof(dst_t);
+    const int64_t s3  = nb3/sizeof(dst_t);
 
-
-    if(ne01 > 0 && ne00 > 0) {
+    if (ne_total > 0) {
         k_set_rows<<<grid_size, block_size, 0, stream>>>(
             src0_d, src1_d, dst_d,
             ne00, ne01, ne02, ne03,
             ne10, ne11, ne12, ne13,
-            s1, s2, s3,
+            s01, s02, s03,
             s10, s11, s12,
-            s_dst, s_dst2, s_dst3,
-            src_type_size, dst_type_size);
+            s1, s2, s3);
     }
 }
 
@@ -109,6 +104,8 @@ void ggml_cuda_op_set_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
 
     cudaStream_t stream = ctx.stream();
 
+
+
     if (dst->type == GGML_TYPE_F32) {
         set_rows_cuda(
             src0_d, src1_d, (float*)dst->data,
@@ -117,7 +114,6 @@ void ggml_cuda_op_set_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
             nb01, nb02, nb03,
             nb10, nb11, nb12,
             nb1, nb2, nb3,
-            sizeof(float), sizeof(float),
             stream
         );
     } else if (dst->type == GGML_TYPE_F16) {
@@ -128,7 +124,6 @@ void ggml_cuda_op_set_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
             nb01, nb02, nb03,
             nb10, nb11, nb12,
             nb1, nb2, nb3,
-            sizeof(float), sizeof(half),
             stream
         );
     } else {

ggml/src/ggml-cuda/set-rows.cuh

@@ -2,6 +2,6 @@
 
 #include "common.cuh"
 
-#define CUDA_SET_ROWS_BLOCK_SIZE 64
+#define CUDA_SET_ROWS_BLOCK_SIZE 256
 
 void ggml_cuda_op_set_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst);