diff --git a/ops/copy/nvidia/copy_cuda.cu b/ops/copy/nvidia/copy_cuda.cu index 4d42b5e..f73b48a 100644 --- a/ops/copy/nvidia/copy_cuda.cu +++ b/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 @@ -25,8 +25,8 @@ template 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<<>>( - static_cast(dst), static_cast(src), desc->elements, oprt::nvidia::CopyOp{}); + oprt::copy::nvidia::copy_contiguous_kernel<<>>( + static_cast(dst), static_cast(src), desc->elements); OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError()); return OPRT_SUCCESS; } diff --git a/ops/copy/nvidia/kernel.cuh b/ops/copy/nvidia/kernel.cuh new file mode 100644 index 0000000..83c0e8b --- /dev/null +++ b/ops/copy/nvidia/kernel.cuh @@ -0,0 +1,16 @@ +#pragma once + +#include + +namespace oprt::copy::nvidia { + +template +__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 diff --git a/ops/copy/operator.yaml b/ops/copy/operator.yaml index 4405960..c190fee 100644 --- a/ops/copy/operator.yaml +++ b/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" - diff --git a/ops/reduce_sum/nvidia/kernel.cuh b/ops/reduce_sum/nvidia/kernel.cuh new file mode 100644 index 0000000..124e0de --- /dev/null +++ b/ops/reduce_sum/nvidia/kernel.cuh @@ -0,0 +1,29 @@ +#pragma once + +#include +#include + +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 diff --git a/ops/reduce_sum/nvidia/reduce_sum_cuda.cu b/ops/reduce_sum/nvidia/reduce_sum_cuda.cu index f150391..21428e5 100644 --- a/ops/reduce_sum/nvidia/reduce_sum_cuda.cu +++ b/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 -#include 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(desc); constexpr int threads = 256; - reduce_sum_rowwise_kernel<<rows, threads, threads * sizeof(float), oprt::as_cuda_stream(stream)>>>( + oprt::reduce_sum::nvidia::reduce_sum_rowwise_kernel<<rows, threads, threads * sizeof(float), oprt::as_cuda_stream(stream)>>>( static_cast(out), static_cast(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; } - diff --git a/ops/reduce_sum/operator.yaml b/ops/reduce_sum/operator.yaml index 31d2c52..3e2b5e8 100644 --- a/ops/reduce_sum/operator.yaml +++ b/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" - diff --git a/ops/softmax/nvidia/kernel.cuh b/ops/softmax/nvidia/kernel.cuh new file mode 100644 index 0000000..9306d5d --- /dev/null +++ b/ops/softmax/nvidia/kernel.cuh @@ -0,0 +1,51 @@ +#pragma once + +#include +#include +#include +#include + +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 diff --git a/ops/softmax/nvidia/softmax_cuda.cu b/ops/softmax/nvidia/softmax_cuda.cu index 00479af..bb98930 100644 --- a/ops/softmax/nvidia/softmax_cuda.cu +++ b/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 -#include -#include 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(desc); constexpr int threads = 256; - softmax_rowwise_kernel<<rows, threads, threads * sizeof(float), oprt::as_cuda_stream(stream)>>>( + oprt::softmax::nvidia::softmax_rowwise_kernel<<rows, threads, threads * sizeof(float), oprt::as_cuda_stream(stream)>>>( static_cast(out), static_cast(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; } - diff --git a/ops/softmax/operator.yaml b/ops/softmax/operator.yaml index a89082a..11258b0 100644 --- a/ops/softmax/operator.yaml +++ b/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" - diff --git a/ops/vector_add/nvidia/kernel.cuh b/ops/vector_add/nvidia/kernel.cuh new file mode 100644 index 0000000..9422bdb --- /dev/null +++ b/ops/vector_add/nvidia/kernel.cuh @@ -0,0 +1,27 @@ +#pragma once + +#include +#include + +namespace oprt::vector_add::nvidia { + +template +__device__ T add_values(T a, T b) { + return a + b; +} + +template <> +__device__ inline half add_values(half a, half b) { + return __hadd(a, b); +} + +template +__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 diff --git a/ops/vector_add/nvidia/vector_add_cuda.cu b/ops/vector_add/nvidia/vector_add_cuda.cu index 94cc52c..1217e61 100644 --- a/ops/vector_add/nvidia/vector_add_cuda.cu +++ b/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 @@ -25,8 +25,8 @@ template 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<<>>( - static_cast(out), static_cast(a), static_cast(b), desc->elements, oprt::nvidia::AddOp{}); + oprt::vector_add::nvidia::vector_add_contiguous_kernel<<>>( + static_cast(out), static_cast(a), static_cast(b), desc->elements); OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError()); return OPRT_SUCCESS; } diff --git a/ops/vector_add/operator.yaml b/ops/vector_add/operator.yaml index 69e7ed7..5bd12fa 100644 --- a/ops/vector_add/operator.yaml +++ b/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" -