121 lines
4.6 KiB
Plaintext
121 lines
4.6 KiB
Plaintext
#include "tri.cuh"
|
|
#include "convert.cuh"
|
|
|
|
template<typename T, bool prefix_keep, int add_to_split>
|
|
static __global__ void tri_kernel(
|
|
const T * src, T * dst,
|
|
const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03,
|
|
const int64_t nb00, const int64_t nb01, const int64_t nb02, const int64_t nb03,
|
|
const int64_t nb0, const int64_t nb1, const int64_t nb2, const int64_t nb3) {
|
|
const int64_t i3 = blockIdx.z;
|
|
const int64_t i2 = blockIdx.y;
|
|
const int64_t i1 = blockIdx.x;
|
|
const int64_t split_point = i1 + add_to_split;
|
|
|
|
(void) nb00;
|
|
(void) nb0;
|
|
|
|
if (i3 >= ne03 || i2 >= ne02 || i1 >= ne01) {
|
|
return;
|
|
}
|
|
|
|
const T * src_row = src + i1 * nb01 + i2 * nb02 + i3 * nb03;
|
|
T * dst_row = dst + i1 * nb1 + i2 * nb2 + i3 * nb3;
|
|
|
|
if constexpr (prefix_keep) {
|
|
for (int64_t i0 = threadIdx.x; i0 < split_point; i0 += blockDim.x) {
|
|
dst_row[i0] = src_row[i0];
|
|
}
|
|
for (int64_t i0 = threadIdx.x + split_point; i0 < ne00; i0 += blockDim.x) {
|
|
dst_row[i0] = ggml_cuda_cast<T, float>(0.0f);
|
|
}
|
|
} else {
|
|
for (int64_t i0 = threadIdx.x; i0 < split_point; i0 += blockDim.x) {
|
|
dst_row[i0] = ggml_cuda_cast<T, float>(0.0f);
|
|
}
|
|
for (int64_t i0 = threadIdx.x + split_point; i0 < ne00; i0 += blockDim.x) {
|
|
dst_row[i0] = src_row[i0];
|
|
}
|
|
}
|
|
}
|
|
|
|
template<typename T>
|
|
static void tri_cuda(
|
|
const T * src, T * dst,
|
|
const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03,
|
|
const int64_t nb00, const int64_t nb01, const int64_t nb02, const int64_t nb03,
|
|
const int64_t nb0, const int64_t nb1, const int64_t nb2, const int64_t nb3,
|
|
const ggml_tri_type ttype,
|
|
cudaStream_t stream) {
|
|
dim3 block_dims(CUDA_TRI_BLOCK_SIZE, 1, 1);
|
|
dim3 grid_dims(ne01, ne02, ne03);
|
|
const size_t type_size = sizeof(T);
|
|
|
|
const int add_to_split = (ttype == GGML_TRI_TYPE_LOWER_DIAG || ttype == GGML_TRI_TYPE_UPPER) ? 1 : 0;
|
|
const bool prefix_keep = (ttype == GGML_TRI_TYPE_LOWER || ttype == GGML_TRI_TYPE_LOWER_DIAG);
|
|
|
|
if (prefix_keep) {
|
|
if (add_to_split == 0) {
|
|
tri_kernel<T, true, 0><<<grid_dims, block_dims, 0, stream>>>(
|
|
src, dst,
|
|
ne00, ne01, ne02, ne03,
|
|
nb00 / type_size, nb01 / type_size, nb02 / type_size, nb03 / type_size,
|
|
nb0 / type_size, nb1 / type_size, nb2 / type_size, nb3 / type_size
|
|
);
|
|
} else {
|
|
tri_kernel<T, true, 1><<<grid_dims, block_dims, 0, stream>>>(
|
|
src, dst,
|
|
ne00, ne01, ne02, ne03,
|
|
nb00 / type_size, nb01 / type_size, nb02 / type_size, nb03 / type_size,
|
|
nb0 / type_size, nb1 / type_size, nb2 / type_size, nb3 / type_size
|
|
);
|
|
}
|
|
} else {
|
|
if (add_to_split == 0) {
|
|
tri_kernel<T, false, 0><<<grid_dims, block_dims, 0, stream>>>(
|
|
src, dst,
|
|
ne00, ne01, ne02, ne03,
|
|
nb00 / type_size, nb01 / type_size, nb02 / type_size, nb03 / type_size,
|
|
nb0 / type_size, nb1 / type_size, nb2 / type_size, nb3 / type_size
|
|
);
|
|
} else {
|
|
tri_kernel<T, false, 1><<<grid_dims, block_dims, 0, stream>>>(
|
|
src, dst,
|
|
ne00, ne01, ne02, ne03,
|
|
nb00 / type_size, nb01 / type_size, nb02 / type_size, nb03 / type_size,
|
|
nb0 / type_size, nb1 / type_size, nb2 / type_size, nb3 / type_size
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
void ggml_cuda_op_tri(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
|
|
const ggml_tensor * src0 = dst->src[0];
|
|
const ggml_tri_type ttype = static_cast<ggml_tri_type>(((const int32_t *) dst->op_params)[0]);
|
|
|
|
GGML_ASSERT(src0->type == dst->type);
|
|
|
|
switch (src0->type) {
|
|
case GGML_TYPE_F32:
|
|
tri_cuda(
|
|
(const float *) src0->data, (float *) dst->data,
|
|
src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3],
|
|
src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3],
|
|
dst->nb[0], dst->nb[1], dst->nb[2], dst->nb[3],
|
|
ttype, ctx.stream()
|
|
);
|
|
break;
|
|
case GGML_TYPE_F16:
|
|
tri_cuda(
|
|
(const half *) src0->data, (half *) dst->data,
|
|
src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3],
|
|
src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3],
|
|
dst->nb[0], dst->nb[1], dst->nb[2], dst->nb[3],
|
|
ttype, ctx.stream()
|
|
);
|
|
break;
|
|
default:
|
|
GGML_ABORT("fatal error");
|
|
}
|
|
}
|