1#include "repeat_back.hpp"
 2
 3#include "common.hpp"
 4
 5#define GGML_ASSERT_TENSOR_FITS_INT(t) \
 6    GGML_ASSERT((t)->ne[0] < INT_MAX && (t)->ne[1] < INT_MAX && (t)->ne[2] < INT_MAX && (t)->ne[3] < INT_MAX)
 7
 8void ggml_sycl_op_repeat_back(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
 9    GGML_ASSERT(dst->src[0]->type == GGML_TYPE_F32);
10    GGML_ASSERT(dst->type == GGML_TYPE_F32);
11
12    const float * src0_dd = (const float *) dst->src[0]->data;
13    float *       dst_dd  = (float *) dst->data;
14
15    GGML_ASSERT_TENSOR_FITS_INT(dst);
16    GGML_ASSERT_TENSOR_FITS_INT(dst->src[0]);
17
18    const int ne0 = dst->ne[0], ne1 = dst->ne[1], ne2 = dst->ne[2], ne3 = dst->ne[3];
19    const int ne00 = dst->src[0]->ne[0], ne01 = dst->src[0]->ne[1], ne02 = dst->src[0]->ne[2],
20              ne03 = dst->src[0]->ne[3];
21
22    const int nr0 = ne00 / ne0;
23    const int nr1 = ne01 / ne1;
24    const int nr2 = ne02 / ne2;
25    const int nr3 = ne03 / ne3;
26
27    const int nb0 = dst->src[0]->nb[0];
28    const int nb1 = dst->src[0]->nb[1];
29    const int nb2 = dst->src[0]->nb[2];
30    const int nb3 = dst->src[0]->nb[3];
31
32    const char * base = (const char *) src0_dd;
33
34    const size_t  total      = (size_t) ne0 * ne1 * ne2 * ne3;
35    constexpr int BLOCK_SIZE = 256;
36    const int     num_blocks = (total + BLOCK_SIZE - 1) / BLOCK_SIZE;
37
38    const float inv_ne0      = 1.0f / ne0;
39    const float inv_ne_01    = 1.0f / (ne0 * ne1);
40    const float inv_ne_012   = 1.0f / (ne0 * ne1 * ne2);
41    const int   repeat_count = nr0 * nr1 * nr2 * nr3;
42
43    queue_ptr stream = ctx.stream();
44
45    stream->parallel_for(
46        sycl::nd_range<1>(sycl::range<1>(num_blocks * BLOCK_SIZE), sycl::range<1>(BLOCK_SIZE)),
47        [=](sycl::nd_item<1> item_ct1) {
48            const size_t i = item_ct1.get_global_linear_id();
49            if (i >= total) {
50                return;
51            }
52
53            const int i3 = (int) (i * inv_ne_012);
54            const int i2 = (int) (i * inv_ne_01) - i3 * ne2;
55            const int i1 = (int) (i * inv_ne0) - (int) (i * inv_ne_01) * ne1;
56            const int i0 = i - (int) (i * inv_ne0) * ne0;
57
58            int   j0 = 0, j1 = 0, j2 = 0, j3 = 0;
59            float acc = 0.0f;
60
61            for (int j = 0; j < repeat_count; ++j) {
62                const float * ptr = (const float *) (base + (i0 + j0 * ne0) * nb0 + (i1 + j1 * ne1) * nb1 +
63                    (i2 + j2 * ne2) * nb2 + (i3 + j3 * ne3) * nb3);
64                acc += *ptr;
65
66                int carry = (++j0 >= nr0);
67                j0 -= carry * nr0;
68                carry = (carry && (++j1 >= nr1));
69                j1 -= carry * nr1;
70                carry = (carry && (++j2 >= nr2));
71                j2 -= carry * nr2;
72                j3 += carry;
73            }
74            dst_dd[i] = acc;
75        });
76}