diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 1f785796014bd..0a132c2a8d205 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -44,6 +44,7 @@ #include "ggml-cuda/upscale.cuh" #include "ggml-cuda/wkv.cuh" #include "ggml-cuda/gla.cuh" +#include "ggml-cuda/set.cuh" #include "ggml-cuda/set-rows.cuh" #include "ggml.h" @@ -2233,6 +2234,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_GET_ROWS_BACK: ggml_cuda_op_get_rows_back(ctx, dst); break; + case GGML_OP_SET: + ggml_cuda_op_set(ctx, dst); + break; case GGML_OP_SET_ROWS: ggml_cuda_op_set_rows(ctx, dst); break; @@ -3275,6 +3279,10 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g { return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->ne[2] == 1 && op->ne[3] == 1; } break; + case GGML_OP_SET: + { + return op->type == GGML_TYPE_F32; + } break; case GGML_OP_SET_ROWS: { return (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_BF16 || diff --git a/ggml/src/ggml-cuda/set.cu b/ggml/src/ggml-cuda/set.cu new file mode 100644 index 0000000000000..fdd343e26808a --- /dev/null +++ b/ggml/src/ggml-cuda/set.cu @@ -0,0 +1,82 @@ +#include "ggml-cuda/common.cuh" +#include "set.cuh" + +static __global__ void set_f32_cuda(const float * __restrict__ src1, + float * __restrict__ dst, + const size_t ne10, + const size_t ne11, + const size_t ne12, + const size_t ne13, + const size_t nb10, + const size_t nb11, + const size_t nb12, + const size_t nb13, + const size_t nb0, + const size_t nb1, + const size_t nb2, + const size_t nb3, + const size_t offset + +) { + const size_t total = ne10 * ne11 * ne12 * ne13; + const size_t gid = blockIdx.x * blockDim.x + threadIdx.x; + if (gid >= total) { + return; + } + + size_t tmp = gid; + + const size_t i0 = tmp % ne10; + tmp /= ne10; + const size_t i1 = tmp % ne11; + tmp /= ne11; + const size_t i2 = tmp % ne12; + tmp /= ne12; + const size_t i3 = tmp; + + size_t dst_offset = offset + i0 * nb0 + i1 * nb1 + i2 * nb2 + i3 * nb3; + size_t src1_offset = i0 * nb10 + i1 * nb11 + i2 * nb12 + i3 * nb13; + + *((float *) ((char *) dst + dst_offset)) = *((const float *) ((const char *) src1 + src1_offset)); +} + +void ggml_cuda_op_set(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + // nb0 is implicitly element_size because src0 and dst are contiguous + const int32_t nb1 = dst->op_params[0]; + const int32_t nb2 = dst->op_params[1]; + const int32_t nb3 = dst->op_params[2]; + const int32_t offset = dst->op_params[3]; + const bool inplace = dst->op_params[4]; + + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + + // TODO: support more dtypes. + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + GGML_TENSOR_BINARY_OP_LOCALS01; + const int nb0 = ggml_element_size(dst); + + const float * src0_d = (const float *) src0->data; + const float * src1_d = (const float *) src1->data; + float * dst_d = (float *) dst->data; + + cudaStream_t stream = ctx.stream(); + + if (!inplace) { + // copy whole src0 -> dst. + CUDA_CHECK(cudaMemcpyAsync(dst_d, src0_d, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream)); + } + + // set: src1 -> dst + + const size_t total = ne10 * ne11 * ne12 * ne13; + const size_t num_blocks = (total + CUDA_SET_BLOCK_SIZE - 1) / CUDA_SET_BLOCK_SIZE; + + set_f32_cuda<<>>( + src1_d, dst_d, ne10, ne11, ne12, ne13, nb10, nb11, nb12, nb13, nb0, nb1, nb2, nb3, offset); +} diff --git a/ggml/src/ggml-cuda/set.cuh b/ggml/src/ggml-cuda/set.cuh new file mode 100644 index 0000000000000..dd09529f3e42b --- /dev/null +++ b/ggml/src/ggml-cuda/set.cuh @@ -0,0 +1,7 @@ +#pragma once + +#include "common.cuh" + +#define CUDA_SET_BLOCK_SIZE 256 + +void ggml_cuda_op_set(ggml_backend_cuda_context & ctx, ggml_tensor * dst);