1#include "common.cuh"
2#include "count-equal.cuh"
3
4#include <cstdint>
5
6template <typename T>
7static __global__ void count_equal(const T * __restrict__ x, const T * __restrict__ y, int64_t * __restrict__ dst, const int64_t dk, const int64_t k) {
8 const int64_t i0 = (int64_t) blockIdx.x*dk;
9 const int64_t i1 = min(a: i0 + dk, b: k);
10
11 int nequal = 0;
12
13 for (int64_t i = i0 + threadIdx.x; i < i1; i += WARP_SIZE) {
14 const T xi = x[i];
15 const T yi = y[i];
16 nequal += xi == yi;
17 }
18
19 nequal = warp_reduce_sum(x: nequal);
20
21 if (threadIdx.x != 0) {
22 return;
23 }
24
25 atomicAdd(address: (int *) dst, val: nequal);
26}
27
28void ggml_cuda_count_equal(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
29 const ggml_tensor * src0 = dst->src[0];
30 const ggml_tensor * src1 = dst->src[1];
31
32 GGML_ASSERT(src0->type == src1->type);
33 GGML_ASSERT( dst->type == GGML_TYPE_I64);
34
35 GGML_ASSERT(ggml_are_same_shape(src0, src1));
36 GGML_ASSERT(ggml_is_contiguous(src0));
37 GGML_ASSERT(ggml_is_contiguous(src1));
38 GGML_ASSERT(ggml_is_contiguous(dst));
39
40 int64_t * dst_d = (int64_t *) dst->data;
41
42 cudaStream_t stream = ctx.stream();
43 const int nsm = ggml_cuda_info().devices[ggml_cuda_get_device()].nsm;
44
45 const int64_t ne = ggml_nelements(src0);
46 GGML_ASSERT(ne < (1 << 30) && "atomicAdd implementation only supports int");
47 const int64_t dne = GGML_PAD((ne + 4*nsm - 1) / (4*nsm), CUDA_COUNT_EQUAL_CHUNK_SIZE);
48
49 CUDA_CHECK(cudaMemsetAsync(dst_d, 0, ggml_nbytes(dst), stream));
50
51 const dim3 blocks_dim(WARP_SIZE, 1, 1);
52 const dim3 blocks_num(std::min(a: (int64_t)4*nsm, b: (ne + CUDA_COUNT_EQUAL_CHUNK_SIZE - 1)/CUDA_COUNT_EQUAL_CHUNK_SIZE), 1, 1);
53
54 switch (src0->type) {
55 case GGML_TYPE_I32: {
56 const int * src0_d = (const int *) src0->data;
57 const int * src1_d = (const int *) src1->data;
58 count_equal<<<gridDim: blocks_num, blockDim: blocks_dim, sharedMem: 0, stream>>>(x: src0_d, y: src1_d, dst: dst_d, dk: dne, k: ne);
59 } break;
60 default:
61 GGML_ASSERT(false);
62 break;
63 }
64}
65