1//
2// MIT license
3// Copyright (C) 2024 Intel Corporation
4// SPDX-License-Identifier: MIT
5//
6
7//
8// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
9// See https://llvm.org/LICENSE.txt for license information.
10// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
11//
12
13#include "concat.hpp"
14
15static inline size_t elem_size(ggml_type t) {
16 return ggml_type_size(t) / ggml_blck_size(t);
17}
18
19template <typename T>
20static void concat_T_dim0(const T *x, const T *y, T *dst,
21 const int ne0, const int ne00,
22 const sycl::nd_item<3> &item_ct1) {
23 int nidx = item_ct1.get_local_id(2) +
24 item_ct1.get_group(2) * item_ct1.get_local_range(2);
25 if (nidx >= ne0) {
26 return;
27 }
28 // operation
29 int offset_dst = nidx + item_ct1.get_group(1) * ne0 +
30 item_ct1.get_group(0) * ne0 * item_ct1.get_group_range(1);
31 if (nidx < ne00) { // src0
32 int offset_src = nidx + item_ct1.get_group(1) * ne00 +
33 item_ct1.get_group(0) * ne00 * item_ct1.get_group_range(1);
34 dst[offset_dst] = x[offset_src];
35 } else {
36 int offset_src =
37 nidx - ne00 + item_ct1.get_group(1) * (ne0 - ne00) +
38 item_ct1.get_group(0) * (ne0 - ne00) * item_ct1.get_group_range(1);
39 dst[offset_dst] = y[offset_src];
40 }
41}
42
43template <typename T>
44static void concat_T_dim1(const T *x, const T *y, T *dst,
45 const int ne0, const int ne01,
46 const sycl::nd_item<3> &item_ct1) {
47 int nidx = item_ct1.get_local_id(2) +
48 item_ct1.get_group(2) * item_ct1.get_local_range(2);
49 if (nidx >= ne0) {
50 return;
51 }
52 // operation
53 int offset_dst = nidx + item_ct1.get_group(1) * ne0 +
54 item_ct1.get_group(0) * ne0 * item_ct1.get_group_range(1);
55 if (item_ct1.get_group(1) < (size_t) ne01) { // src0
56 int offset_src =
57 nidx + item_ct1.get_group(1) * ne0 + item_ct1.get_group(0) * ne0 * ne01;
58 dst[offset_dst] = x[offset_src];
59 } else {
60 int offset_src =
61 nidx + (item_ct1.get_group(1) - ne01) * ne0 +
62 item_ct1.get_group(0) * ne0 * (item_ct1.get_group_range(1) - ne01);
63 dst[offset_dst] = y[offset_src];
64 }
65}
66
67template <typename T>
68static void concat_T_dim2(const T *x, const T *y, T *dst,
69 const int ne0, const int ne02,
70 const sycl::nd_item<3> &item_ct1) {
71 int nidx = item_ct1.get_local_id(2) +
72 item_ct1.get_group(2) * item_ct1.get_local_range(2);
73 if (nidx >= ne0) {
74 return;
75 }
76 // operation
77 int offset_dst = nidx + item_ct1.get_group(1) * ne0 +
78 item_ct1.get_group(0) * ne0 * item_ct1.get_group_range(1);
79 if (item_ct1.get_group(0) < (size_t) ne02) { // src0
80 int offset_src = nidx + item_ct1.get_group(1) * ne0 +
81 item_ct1.get_group(0) * ne0 * item_ct1.get_group_range(1);
82 dst[offset_dst] = x[offset_src];
83 } else {
84 int offset_src =
85 nidx + item_ct1.get_group(1) * ne0 +
86 (item_ct1.get_group(0) - ne02) * ne0 * item_ct1.get_group_range(1);
87 dst[offset_dst] = y[offset_src];
88 }
89}
90
91template <typename T>
92static void concat_T_sycl(const T *x, const T *y, T *dst,
93 int ne00, int ne01, int ne02, int ne0, int ne1,
94 int ne2, int dim, queue_ptr stream) {
95 int num_blocks = (ne0 + SYCL_CONCAT_BLOCK_SIZE - 1) / SYCL_CONCAT_BLOCK_SIZE;
96 sycl::range<3> gridDim(ne2, ne1, num_blocks);
97 switch (dim) {
98 case 0:
99 stream->parallel_for(sycl::nd_range<3>(gridDim * sycl::range<3>(1, 1, SYCL_CONCAT_BLOCK_SIZE),
100 sycl::range<3>(1, 1, SYCL_CONCAT_BLOCK_SIZE)),
101 [=](sycl::nd_item<3> item_ct1) { concat_T_dim0<T>(x, y, dst, ne0, ne00, item_ct1); });
102 break;
103 case 1:
104 stream->parallel_for(sycl::nd_range<3>(gridDim * sycl::range<3>(1, 1, SYCL_CONCAT_BLOCK_SIZE),
105 sycl::range<3>(1, 1, SYCL_CONCAT_BLOCK_SIZE)),
106 [=](sycl::nd_item<3> item_ct1) { concat_T_dim1<T>(x, y, dst, ne0, ne01, item_ct1); });
107 break;
108 // dim >=2 will be dispatched to the default path
109 default:
110 stream->parallel_for(sycl::nd_range<3>(gridDim * sycl::range<3>(1, 1, SYCL_CONCAT_BLOCK_SIZE),
111 sycl::range<3>(1, 1, SYCL_CONCAT_BLOCK_SIZE)),
112 [=](sycl::nd_item<3> item_ct1) { concat_T_dim2<T>(x, y, dst, ne0, ne02, item_ct1); });
113 break;
114 }
115}
116
117// non-contiguous kernel (slow)
118template<typename T>
119static void concat_T_sycl_non_cont(
120 queue_ptr stream, const char *src0, const char *src1, char *dst,
121 int64_t ne00, int64_t ne01, int64_t ne02, int64_t ne03, uint64_t nb00,
122 uint64_t nb01, uint64_t nb02, uint64_t nb03, int64_t /*ne10*/,
123 int64_t /*ne11*/, int64_t /*ne12*/, int64_t /*ne13*/, uint64_t nb10,
124 uint64_t nb11, uint64_t nb12, uint64_t nb13, int64_t ne0, int64_t ne1,
125 int64_t ne2, int64_t ne3, uint64_t nb0, uint64_t nb1, uint64_t nb2,
126 uint64_t nb3, int32_t dim) {
127 sycl::range<3> gridDim(ne3, ne2, ne1);
128 stream->parallel_for(sycl::nd_range<3>(gridDim, sycl::range<3>(1, 1, 1)), [=](sycl::nd_item<3> item_ct1) {
129 int64_t i3 = item_ct1.get_group(0);
130 int64_t i2 = item_ct1.get_group(1);
131 int64_t i1 = item_ct1.get_group(2);
132
133 int64_t o[4] = { 0, 0, 0, 0 };
134 o[dim] = dim == 0 ? ne00 : (dim == 1 ? ne01 : (dim == 2 ? ne02 : ne03));
135
136 const T * x;
137
138 for (int i0 = item_ct1.get_local_id(2); i0 < ne0; i0 += item_ct1.get_local_range(2)) {
139 if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) {
140 x = (const T *) (src0 + (i3) *nb03 + (i2) *nb02 + (i1) *nb01 + (i0) *nb00);
141 } else {
142 x = (const T *) (src1 + (i3 - o[3]) * nb13 + (i2 - o[2]) * nb12 + (i1 - o[1]) * nb11 +
143 (i0 - o[0]) * nb10);
144 }
145
146 T *y = (T *)(dst + i3 * nb3 + i2 * nb2 + i1 * nb1 + i0 * nb0);
147
148 *y = *x;
149 }
150 });
151}
152
153template <typename T>
154void concat_impl_sycl(ggml_backend_sycl_context & ctx, ggml_tensor *dst) {
155 scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
156 const ggml_tensor * src0 = dst->src[0];
157 const ggml_tensor * src1 = dst->src[1];
158 queue_ptr stream = ctx.stream();
159
160 const int32_t dim = ((int32_t *) dst->op_params)[0];
161
162 if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1)) {
163 const T * src0_d = (const T *) src0->data;
164 const T * src1_d = (const T *) src1->data;
165 T * dst_d = (T *) dst->data;
166 size_t type_size = elem_size(dst->type);
167 if (dim != 3) {
168 for (int i3 = 0; i3 < dst->ne[3]; i3++) {
169 concat_T_sycl<T>(src0_d + i3 * (src0->nb[3] / type_size), src1_d + i3 * (src1->nb[3] / type_size),
170 dst_d + i3 * (dst->nb[3] / type_size), src0->ne[0], src0->ne[1], src0->ne[2], dst->ne[0],
171 dst->ne[1], dst->ne[2], dim, stream);
172 }
173 } else {
174 const size_t size0 = ggml_nbytes(src0);
175 const size_t size1 = ggml_nbytes(src1);
176
177 SYCL_CHECK(CHECK_TRY_ERROR(stream->memcpy(dst_d, src0_d, size0).wait()));
178 SYCL_CHECK(CHECK_TRY_ERROR(stream->memcpy(dst_d + size0 / type_size, src1_d, size1).wait()));
179 }
180 } else {
181 concat_T_sycl_non_cont<T>(stream, (const char *) src0->data, (const char *) src1->data, (char *) dst->data,
182 src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], src0->nb[0], src0->nb[1],
183 src0->nb[2], src0->nb[3], src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3],
184 src1->nb[0], src1->nb[1], src1->nb[2], src1->nb[3], dst->ne[0], dst->ne[1], dst->ne[2],
185 dst->ne[3], dst->nb[0], dst->nb[1], dst->nb[2], dst->nb[3], dim);
186 }
187}
188
189void ggml_sycl_op_concat(ggml_backend_sycl_context & ctx, ggml_tensor *dst) {
190
191 switch (dst->type) {
192 case GGML_TYPE_F32:
193 concat_impl_sycl<float>(ctx, dst);
194 break;
195 case GGML_TYPE_I32:
196 concat_impl_sycl<int32_t>(ctx, dst);
197 break;
198 default:
199 GGML_ASSERT(false && "ggml_sycl_op_concat: unsupported type");
200 break;
201 }
202}