refactor: split nvidia operator kernels

Co-authored-by: wawahejun <hejunlbbc@gmail.com>

ops/copy/nvidia/copy_cuda.cu

@@ -4,7 +4,7 @@
 #include "operator_runtime/elementwise.h"
 #include "operator_runtime/tensor_checks.h"
 #include "operator_runtime/cuda_helpers.h"
-#include "ops/common/nvidia/elementwise.cuh"
+#include "ops/copy/nvidia/kernel.cuh"
 
 #include <cuda_fp16.h>
 
@@ -25,8 +25,8 @@ template <typename T>
 oprt_status_t launch_copy(const CopyDescriptor *desc, void *dst, const void *src, oprt_stream_t stream) {
     constexpr int threads = 256;
     int blocks = oprt::blocks_for(desc->elements, threads);
-    oprt::nvidia::unary_contiguous_kernel<T><<<blocks, threads, 0, oprt::as_cuda_stream(stream)>>>(
-        static_cast<T *>(dst), static_cast<const T *>(src), desc->elements, oprt::nvidia::CopyOp{});
+    oprt::copy::nvidia::copy_contiguous_kernel<T><<<blocks, threads, 0, oprt::as_cuda_stream(stream)>>>(
+        static_cast<T *>(dst), static_cast<const T *>(src), desc->elements);
     OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError());
     return OPRT_SUCCESS;
 }

ops/copy/nvidia/kernel.cuh

@@ -0,0 +1,16 @@
+#pragma once
+
+#include <stdint.h>
+
+namespace oprt::copy::nvidia {
+
+template <typename T>
+__global__ void copy_contiguous_kernel(T *dst, const T *src, int64_t n) {
+    int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+    int64_t stride = int64_t(blockDim.x) * gridDim.x;
+    for (int64_t i = idx; i < n; i += stride) {
+        dst[i] = src[i];
+    }
+}
+
+} // namespace oprt::copy::nvidia

ops/copy/operator.yaml

@@ -10,6 +10,7 @@ backends:
       - ops/copy/nvidia/copy_cuda.cu
     headers:
       - ops/copy/nvidia/copy_cuda.h
+      - ops/copy/nvidia/kernel.cuh
     symbols:
       create: oprt_create_copy_descriptor_nvidia
       workspace: oprt_get_copy_workspace_size_nvidia
@@ -43,4 +44,3 @@ performance_model:
   bound: memory
   bytes: "2 * numel * elem_bytes"
   flops: "0"
-

ops/reduce_sum/nvidia/kernel.cuh

@@ -0,0 +1,29 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <stdint.h>
+
+namespace oprt::reduce_sum::nvidia {
+
+__global__ void reduce_sum_rowwise_kernel(float *out, const float *in, int64_t rows, int64_t cols) {
+    extern __shared__ float smem[];
+    int row = blockIdx.x;
+    float sum = 0.0f;
+    for (int64_t col = threadIdx.x; col < cols; col += blockDim.x) {
+        sum += in[int64_t(row) * cols + col];
+    }
+    smem[threadIdx.x] = sum;
+    __syncthreads();
+
+    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+        if (threadIdx.x < stride) {
+            smem[threadIdx.x] += smem[threadIdx.x + stride];
+        }
+        __syncthreads();
+    }
+    if (threadIdx.x == 0) {
+        out[row] = smem[0];
+    }
+}
+
+} // namespace oprt::reduce_sum::nvidia

ops/reduce_sum/nvidia/reduce_sum_cuda.cu

@@ -3,9 +3,9 @@
 #include "operator_runtime/descriptor.h"
 #include "operator_runtime/tensor_checks.h"
 #include "operator_runtime/cuda_helpers.h"
+#include "ops/reduce_sum/nvidia/kernel.cuh"
 
 #include <cuda_runtime.h>
-#include <float.h>
 
 namespace {
 
@@ -21,27 +21,6 @@ struct ReduceSumDescriptor final : oprt_operator_descriptor {
     }
 };
 
-__global__ void reduce_sum_rowwise_kernel(float *out, const float *in, int64_t rows, int64_t cols) {
-    extern __shared__ float smem[];
-    int row = blockIdx.x;
-    float sum = 0.0f;
-    for (int64_t col = threadIdx.x; col < cols; col += blockDim.x) {
-        sum += in[int64_t(row) * cols + col];
-    }
-    smem[threadIdx.x] = sum;
-    __syncthreads();
-
-    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
-        if (threadIdx.x < stride) {
-            smem[threadIdx.x] += smem[threadIdx.x + stride];
-        }
-        __syncthreads();
-    }
-    if (threadIdx.x == 0) {
-        out[row] = smem[0];
-    }
-}
-
 bool is_rowwise_case(const oprt_tensor_view_t &out, const oprt_tensor_view_t &in, int64_t axis) {
     return in.dtype == OPRT_DTYPE_F32 &&
            out.dtype == OPRT_DTYPE_F32 &&
@@ -112,7 +91,7 @@ extern "C" OPRT_EXPORT oprt_status_t oprt_execute_reduce_sum_nvidia(
     }
     auto *typed = static_cast<const ReduceSumDescriptor *>(desc);
     constexpr int threads = 256;
-    reduce_sum_rowwise_kernel<<<typed->rows, threads, threads * sizeof(float), oprt::as_cuda_stream(stream)>>>(
+    oprt::reduce_sum::nvidia::reduce_sum_rowwise_kernel<<<typed->rows, threads, threads * sizeof(float), oprt::as_cuda_stream(stream)>>>(
         static_cast<float *>(out), static_cast<const float *>(in), typed->rows, typed->cols);
     OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError());
     return OPRT_SUCCESS;
@@ -123,4 +102,3 @@ extern "C" OPRT_EXPORT oprt_status_t oprt_destroy_reduce_sum_descriptor_nvidia(
     delete desc;
     return OPRT_SUCCESS;
 }
-

ops/reduce_sum/operator.yaml

@@ -10,6 +10,7 @@ backends:
       - ops/reduce_sum/nvidia/reduce_sum_cuda.cu
     headers:
       - ops/reduce_sum/nvidia/reduce_sum_cuda.h
+      - ops/reduce_sum/nvidia/kernel.cuh
     symbols:
       create: oprt_create_reduce_sum_descriptor_nvidia
       workspace: oprt_get_reduce_sum_workspace_size_nvidia
@@ -42,4 +43,3 @@ performance_model:
   bound: memory
   bytes: "(numel + rows) * elem_bytes"
   flops: "numel"
-

ops/softmax/nvidia/kernel.cuh

@@ -0,0 +1,51 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <float.h>
+#include <math.h>
+#include <stdint.h>
+
+namespace oprt::softmax::nvidia {
+
+__global__ void softmax_rowwise_kernel(float *out, const float *in, int64_t rows, int64_t cols) {
+    extern __shared__ float smem[];
+    int row = blockIdx.x;
+    float local_max = -FLT_MAX;
+    for (int64_t col = threadIdx.x; col < cols; col += blockDim.x) {
+        float value = in[int64_t(row) * cols + col];
+        local_max = fmaxf(local_max, value);
+    }
+    smem[threadIdx.x] = local_max;
+    __syncthreads();
+
+    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+        if (threadIdx.x < stride) {
+            smem[threadIdx.x] = fmaxf(smem[threadIdx.x], smem[threadIdx.x + stride]);
+        }
+        __syncthreads();
+    }
+    float row_max = smem[0];
+
+    float local_sum = 0.0f;
+    for (int64_t col = threadIdx.x; col < cols; col += blockDim.x) {
+        float value = expf(in[int64_t(row) * cols + col] - row_max);
+        out[int64_t(row) * cols + col] = value;
+        local_sum += value;
+    }
+    smem[threadIdx.x] = local_sum;
+    __syncthreads();
+
+    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+        if (threadIdx.x < stride) {
+            smem[threadIdx.x] += smem[threadIdx.x + stride];
+        }
+        __syncthreads();
+    }
+    float row_sum = smem[0];
+
+    for (int64_t col = threadIdx.x; col < cols; col += blockDim.x) {
+        out[int64_t(row) * cols + col] /= row_sum;
+    }
+}
+
+} // namespace oprt::softmax::nvidia

ops/softmax/nvidia/softmax_cuda.cu

@@ -3,10 +3,9 @@
 #include "operator_runtime/descriptor.h"
 #include "operator_runtime/tensor_checks.h"
 #include "operator_runtime/cuda_helpers.h"
+#include "ops/softmax/nvidia/kernel.cuh"
 
 #include <cuda_runtime.h>
-#include <float.h>
-#include <math.h>
 
 namespace {
 
@@ -22,47 +21,6 @@ struct SoftmaxDescriptor final : oprt_operator_descriptor {
     }
 };
 
-__global__ void softmax_rowwise_kernel(float *out, const float *in, int64_t rows, int64_t cols) {
-    extern __shared__ float smem[];
-    int row = blockIdx.x;
-    float local_max = -FLT_MAX;
-    for (int64_t col = threadIdx.x; col < cols; col += blockDim.x) {
-        float value = in[int64_t(row) * cols + col];
-        local_max = fmaxf(local_max, value);
-    }
-    smem[threadIdx.x] = local_max;
-    __syncthreads();
-
-    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
-        if (threadIdx.x < stride) {
-            smem[threadIdx.x] = fmaxf(smem[threadIdx.x], smem[threadIdx.x + stride]);
-        }
-        __syncthreads();
-    }
-    float row_max = smem[0];
-
-    float local_sum = 0.0f;
-    for (int64_t col = threadIdx.x; col < cols; col += blockDim.x) {
-        float value = expf(in[int64_t(row) * cols + col] - row_max);
-        out[int64_t(row) * cols + col] = value;
-        local_sum += value;
-    }
-    smem[threadIdx.x] = local_sum;
-    __syncthreads();
-
-    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
-        if (threadIdx.x < stride) {
-            smem[threadIdx.x] += smem[threadIdx.x + stride];
-        }
-        __syncthreads();
-    }
-    float row_sum = smem[0];
-
-    for (int64_t col = threadIdx.x; col < cols; col += blockDim.x) {
-        out[int64_t(row) * cols + col] /= row_sum;
-    }
-}
-
 bool is_rowwise_case(const oprt_tensor_view_t &out, const oprt_tensor_view_t &in, int64_t axis) {
     return in.dtype == OPRT_DTYPE_F32 &&
            out.dtype == OPRT_DTYPE_F32 &&
@@ -133,7 +91,7 @@ extern "C" OPRT_EXPORT oprt_status_t oprt_execute_softmax_nvidia(
     }
     auto *typed = static_cast<const SoftmaxDescriptor *>(desc);
     constexpr int threads = 256;
-    softmax_rowwise_kernel<<<typed->rows, threads, threads * sizeof(float), oprt::as_cuda_stream(stream)>>>(
+    oprt::softmax::nvidia::softmax_rowwise_kernel<<<typed->rows, threads, threads * sizeof(float), oprt::as_cuda_stream(stream)>>>(
         static_cast<float *>(out), static_cast<const float *>(in), typed->rows, typed->cols);
     OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError());
     return OPRT_SUCCESS;
@@ -144,4 +102,3 @@ extern "C" OPRT_EXPORT oprt_status_t oprt_destroy_softmax_descriptor_nvidia(
     delete desc;
     return OPRT_SUCCESS;
 }
-

ops/softmax/operator.yaml

@@ -10,6 +10,7 @@ backends:
       - ops/softmax/nvidia/softmax_cuda.cu
     headers:
       - ops/softmax/nvidia/softmax_cuda.h
+      - ops/softmax/nvidia/kernel.cuh
     symbols:
       create: oprt_create_softmax_descriptor_nvidia
       workspace: oprt_get_softmax_workspace_size_nvidia
@@ -42,4 +43,3 @@ performance_model:
   bound: mixed
   bytes: "5 * numel * elem_bytes"
   flops: "4 * numel"
-

ops/vector_add/nvidia/kernel.cuh

@@ -0,0 +1,27 @@
+#pragma once
+
+#include <cuda_fp16.h>
+#include <stdint.h>
+
+namespace oprt::vector_add::nvidia {
+
+template <typename T>
+__device__ T add_values(T a, T b) {
+    return a + b;
+}
+
+template <>
+__device__ inline half add_values<half>(half a, half b) {
+    return __hadd(a, b);
+}
+
+template <typename T>
+__global__ void vector_add_contiguous_kernel(T *out, const T *a, const T *b, int64_t n) {
+    int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+    int64_t stride = int64_t(blockDim.x) * gridDim.x;
+    for (int64_t i = idx; i < n; i += stride) {
+        out[i] = add_values(a[i], b[i]);
+    }
+}
+
+} // namespace oprt::vector_add::nvidia

ops/vector_add/nvidia/vector_add_cuda.cu

@@ -4,7 +4,7 @@
 #include "operator_runtime/elementwise.h"
 #include "operator_runtime/tensor_checks.h"
 #include "operator_runtime/cuda_helpers.h"
-#include "ops/common/nvidia/elementwise.cuh"
+#include "ops/vector_add/nvidia/kernel.cuh"
 
 #include <cuda_fp16.h>
 
@@ -25,8 +25,8 @@ template <typename T>
 oprt_status_t launch_vector_add(const VectorAddDescriptor *desc, void *out, const void *a, const void *b, oprt_stream_t stream) {
     constexpr int threads = 256;
     int blocks = oprt::blocks_for(desc->elements, threads);
-    oprt::nvidia::binary_contiguous_kernel<T><<<blocks, threads, 0, oprt::as_cuda_stream(stream)>>>(
-        static_cast<T *>(out), static_cast<const T *>(a), static_cast<const T *>(b), desc->elements, oprt::nvidia::AddOp{});
+    oprt::vector_add::nvidia::vector_add_contiguous_kernel<T><<<blocks, threads, 0, oprt::as_cuda_stream(stream)>>>(
+        static_cast<T *>(out), static_cast<const T *>(a), static_cast<const T *>(b), desc->elements);
     OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError());
     return OPRT_SUCCESS;
 }

ops/vector_add/operator.yaml

@@ -10,6 +10,7 @@ backends:
       - ops/vector_add/nvidia/vector_add_cuda.cu
     headers:
       - ops/vector_add/nvidia/vector_add_cuda.h
+      - ops/vector_add/nvidia/kernel.cuh
     symbols:
       create: oprt_create_vector_add_descriptor_nvidia
       workspace: oprt_get_vector_add_workspace_size_nvidia
@@ -43,4 +44,3 @@ performance_model:
   bound: memory
   bytes: "3 * numel * elem_bytes"
   flops: "numel"
-