char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Xuan-Son Nguyenggml: support ggml_rope_set_offset on opencl, sycl, wgpu, hexagon (#27345)749f688fc

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
27.9 KiB651 linesraw
1#include "rope.hpp"
2#include "convert.hpp"
3#include "ggml-sycl/common.hpp"
4#include "ggml.h"
5
6struct rope_corr_dims {
7    float v[2];
8};
9
10struct mrope_sections {
11    int v[4];
12};
13
14static float rope_yarn_ramp(const float low, const float high, const int i0) {
15    const float y = (i0 / 2 - low) / sycl::max(0.001f, high - low);
16    return 1.0f - sycl::min(1.0f, sycl::max(0.0f, y));
17}
18
19template <bool forward>
20static void rope_yarn(const float theta_extrap, const float freq_scale,
21                      const rope_corr_dims corr_dims, const int64_t i0,
22                      const float ext_factor, float mscale, float &cos_theta,
23                      float &sin_theta) {
24    float theta_interp = freq_scale * theta_extrap;
25    float theta = theta_interp;
26    if (ext_factor != 0.0f) {
27        float ramp_mix =
28            rope_yarn_ramp(corr_dims.v[0], corr_dims.v[1], i0) * ext_factor;
29        theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix;
30
31        mscale *= 1.0f + 0.1f * sycl::log(1.0f / freq_scale);
32    }
33    cos_theta = sycl::cos(theta) * mscale;
34    sin_theta = sycl::sin(theta) * mscale;
35    if (!forward) {
36        sin_theta *= -1.0f;
37    }
38}
39
40template <bool forward, bool has_ff, typename T, typename D>
41static void rope_norm(const T *x, D *dst, const int ne00, const int ne01,
42                      const int ne02, const int s01, const int s02,
43                      const int s03, const int s1, const int s2, const int s3,
44                      const int n_dims, const int n_offs, const int32_t *pos,
45                      const float freq_scale, const float ext_factor,
46                      const float attn_factor, const rope_corr_dims corr_dims,
47                      const float theta_scale, const float *freq_factors,
48                      const int64_t *row_indices, const int set_rows_stride) {
49    auto item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>();
50    const int i0 = 2 * (item_ct1.get_local_range(1) * item_ct1.get_group(1) +
51                        item_ct1.get_local_id(1));
52
53    if (i0 >= ne00) {
54        return;
55    }
56
57    const int row_dst = item_ct1.get_local_range(2) * item_ct1.get_group(2) +
58                        item_ct1.get_local_id(2);
59
60    const uint32_t i3 = row_dst / (ne01 * ne02);
61    const uint32_t i2 = (row_dst - i3 * ne01 * ne02) / ne01;
62    const uint32_t i1 = row_dst - i3 * ne01 * ne02 - i2 * ne01;
63
64    int idst = i0 + i1 * s1 + i2 * s2 + i3 * s3;
65    const int ix = i0 + i1 * s01 + i2 * s02 + i3 * s03;
66
67    if (set_rows_stride != 0) {
68        idst = i1 * s1 + i0;
69        idst += row_indices[i2] * set_rows_stride;
70    }
71
72    const auto &store_coaelsced = [&](float x0, float x1) {
73        if constexpr (std::is_same_v<float, D>) {
74            sycl::float2 v = sycl::float2(x0, x1);
75            ggml_sycl_memcpy_1<8>(dst + idst, &v);
76        } else if constexpr (std::is_same_v<sycl::half, D>) {
77            sycl::half2 v = sycl::half2(x0, x1);
78            ggml_sycl_memcpy_1<4>(dst + idst, &v);
79        }
80    };
81    if (i0 < n_offs || i0 >= n_offs + n_dims) {
82        store_coaelsced(x[ix + 0], x[ix + 1]);
83        return;
84    }
85
86    const int iw = i0 - n_offs; // relative idx
87
88    const float theta_base = pos[i2] * dpct::pow(theta_scale, iw / 2.0f);
89
90    const float freq_factor = has_ff ? freq_factors[iw / 2] : 1.0f;
91
92    float cos_theta;
93    float sin_theta;
94
95    rope_yarn<forward>(theta_base / freq_factor, freq_scale, corr_dims, iw,
96                       ext_factor, attn_factor, cos_theta, sin_theta);
97
98    const float x0 = x[ix + 0];
99    const float x1 = x[ix + 1];
100
101    store_coaelsced(x0 * cos_theta - x1 * sin_theta,
102                    x0 * sin_theta + x1 * cos_theta);
103}
104
105template <bool forward, bool has_ff, typename T, typename D>
106static void rope_neox(const T *x, D *dst, const int ne00, const int ne01,
107                      const int ne02, const int s01, const int s02,
108                      const int s03, const int s1, const int s2, const int s3,
109                      const int n_dims, const int n_offs, const int32_t *pos,
110                      const float freq_scale, const float ext_factor,
111                      const float attn_factor, const rope_corr_dims corr_dims,
112                      const float theta_scale, const float *freq_factors,
113                      const int64_t *row_indices, const int set_rows_stride) {
114    auto item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>();
115    const int i0 = 2 * (item_ct1.get_local_range(1) * item_ct1.get_group(1) +
116                        item_ct1.get_local_id(1));
117
118    if (i0 >= ne00) {
119        return;
120    }
121
122    const int row_dst = item_ct1.get_local_range(2) * item_ct1.get_group(2) +
123                        item_ct1.get_local_id(2);
124
125    const uint32_t i3 = row_dst / (ne01 * ne02);
126    const uint32_t i2 = (row_dst - i3 * ne01 * ne02) / ne01;
127    const uint32_t i1 = row_dst - i3 * ne01 * ne02 - i2 * ne01;
128
129    int idst = i0 / 2 + i1 * s1 + i2 * s2 + i3 * s3;
130    const int ix = i0 / 2 + i1 * s01 + i2 * s02 + i3 * s03;
131
132    if (set_rows_stride != 0) {
133        idst = i1 * s1 + i0 / 2;
134        idst += row_indices[i2] * set_rows_stride;
135    }
136
137    if (i0 < n_offs || i0 >= n_offs + n_dims) {
138        dst[idst + i0 / 2 + 0] = ggml_sycl_cast<D>(x[ix + i0 / 2 + 0]);
139        dst[idst + i0 / 2 + 1] = ggml_sycl_cast<D>(x[ix + i0 / 2 + 1]);
140
141        return;
142    }
143
144    const int iw = i0 - n_offs; // relative idx
145
146    const float theta_base = pos[i2] * dpct::pow(theta_scale, iw / 2.0f);
147
148    const float freq_factor = has_ff ? freq_factors[iw / 2] : 1.0f;
149
150    float cos_theta;
151    float sin_theta;
152
153    rope_yarn<forward>(theta_base / freq_factor, freq_scale, corr_dims, iw,
154                       ext_factor, attn_factor, cos_theta, sin_theta);
155
156    // idst/ix point at channel i0/2; the first channel of the rotated pair is n_offs + iw/2 = i0/2 + n_offs/2
157    const float x0 = x[ix + n_offs / 2 + 0];
158    const float x1 = x[ix + n_offs / 2 + n_dims / 2];
159
160    dst[idst + n_offs / 2 + 0] = ggml_sycl_cast<D>(x0 * cos_theta - x1 * sin_theta);
161    dst[idst + n_offs / 2 + n_dims / 2] = ggml_sycl_cast<D>(x0 * sin_theta + x1 * cos_theta);
162}
163
164template <bool forward, bool has_ff, typename T>
165static void rope_multi(const T *x, T *dst, const int ne00, const int ne01,
166                       const int ne02, const int s01, const int s02,
167                       const int s03, const int s1, const int s2, const int s3,
168                       const int n_dims, const int n_offs, const int32_t *pos,
169                       const float freq_scale, const float ext_factor,
170                       const float attn_factor, const rope_corr_dims corr_dims,
171                       const float theta_scale, const float *freq_factors,
172                       const mrope_sections sections, const bool is_imrope) {
173    auto item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>();
174    const int i0 = 2 * (item_ct1.get_local_range(1) * item_ct1.get_group(1) +
175                        item_ct1.get_local_id(1));
176
177    if (i0 >= ne00) {
178        return;
179    }
180
181    const int row_dst = item_ct1.get_local_range(2) * item_ct1.get_group(2) +
182                        item_ct1.get_local_id(2);
183
184    const uint32_t i3 = row_dst / (ne01 * ne02);
185    const uint32_t i2 = (row_dst - i3 * ne01 * ne02) / ne01;
186    const uint32_t i1 = row_dst - i3 * ne01 * ne02 - i2 * ne01;
187
188    int idst = i0 / 2 + i1 * s1 + i2 * s2 + i3 * s3;
189    const int ix = i0 / 2 + i1 * s01 + i2 * s02 + i3 * s03;
190
191    if (i0 < n_offs || i0 >= n_offs + n_dims) {
192        dst[idst + i0 / 2 + 0] = x[ix + i0 / 2 + 0];
193        dst[idst + i0 / 2 + 1] = x[ix + i0 / 2 + 1];
194
195        return;
196    }
197
198    const int iw = i0 - n_offs; // relative idx
199
200    const int sect_dims =
201        sections.v[0] + sections.v[1] + sections.v[2] + sections.v[3];
202    const int sec_w = sections.v[1] + sections.v[0];
203    const int sector = (iw / 2) % sect_dims;
204
205    float theta_base = 0.0;
206    if (is_imrope) {
207        if (sector % 3 == 1 && sector < 3 * sections.v[1]) { // h
208            theta_base = pos[i2 + ne02 * 1] * dpct::pow(theta_scale, iw / 2.0f);
209        } else if (sector % 3 == 2 && sector < 3 * sections.v[2]) { // w
210            theta_base = pos[i2 + ne02 * 2] * dpct::pow(theta_scale, iw / 2.0f);
211        } else if (sector % 3 == 0 && sector < 3 * sections.v[0]) { // t
212            theta_base = pos[i2] * dpct::pow(theta_scale, iw / 2.0f);
213        } else {
214            theta_base = pos[i2 + ne02 * 3] * dpct::pow(theta_scale, iw / 2.0f);
215        }
216    } else {
217        if (sector < sections.v[0]) {
218            theta_base = pos[i2] * dpct::pow(theta_scale, iw / 2.0f);
219        } else if (sector >= sections.v[0] && sector < sec_w) {
220            theta_base = pos[i2 + ne02 * 1] * dpct::pow(theta_scale, iw / 2.0f);
221        } else if (sector >= sec_w && sector < sec_w + sections.v[2]) {
222            theta_base = pos[i2 + ne02 * 2] * dpct::pow(theta_scale, iw / 2.0f);
223        } else if (sector >= sec_w + sections.v[2]) {
224            theta_base = pos[i2 + ne02 * 3] * dpct::pow(theta_scale, iw / 2.0f);
225        }
226    }
227
228    const float freq_factor = has_ff ? freq_factors[iw / 2] : 1.0f;
229
230    float cos_theta;
231    float sin_theta;
232
233    rope_yarn<forward>(theta_base / freq_factor, freq_scale, corr_dims, iw,
234                       ext_factor, attn_factor, cos_theta, sin_theta);
235
236    // idst/ix point at channel i0/2; the first channel of the rotated pair is n_offs + iw/2 = i0/2 + n_offs/2
237    const float x0 = x[ix + n_offs / 2 + 0];
238    const float x1 = x[ix + n_offs / 2 + n_dims / 2];
239
240    dst[idst + n_offs / 2 + 0] = x0 * cos_theta - x1 * sin_theta;
241    dst[idst + n_offs / 2 + n_dims / 2] = x0 * sin_theta + x1 * cos_theta;
242}
243
244template <bool forward, bool has_ff, typename T>
245static void rope_vision(const T *x, T *dst, const int ne00, const int ne01,
246                        const int ne02, const int s01, const int s02,
247                        const int s03, const int s1, const int s2, const int s3,
248                        const int n_dims, const int32_t *pos,
249                        const float freq_scale, const float ext_factor,
250                        const float attn_factor, const rope_corr_dims corr_dims,
251                        const float theta_scale, const float *freq_factors,
252                        const mrope_sections sections) {
253    auto item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>();
254    const int i0 = 2 * (item_ct1.get_local_range(1) * item_ct1.get_group(1) +
255                        item_ct1.get_local_id(1));
256
257    if (i0 >= ne00) {
258        return;
259    }
260
261    const int row_dst = item_ct1.get_local_range(2) * item_ct1.get_group(2) +
262                        item_ct1.get_local_id(2);
263
264    const uint32_t i3 = row_dst / (ne01 * ne02);
265    const uint32_t i2 = (row_dst - i3 * ne01 * ne02) / ne01;
266    const uint32_t i1 = row_dst - i3 * ne01 * ne02 - i2 * ne01;
267
268    int idst = i0 / 2 + i1 * s1 + i2 * s2 + i3 * s3;
269    const int ix = i0 / 2 + i1 * s01 + i2 * s02 + i3 * s03;
270
271    const int sect_dims = sections.v[0] + sections.v[1];
272    const int sec_w = sections.v[1] + sections.v[0];
273    const int sector = (i0 / 2) % sect_dims;
274
275    float theta_base = 0.0;
276    if (sector < sections.v[0]) {
277        const int p = sector;
278        theta_base = pos[i2] * dpct::pow(theta_scale, p);
279    } else if (sector >= sections.v[0] && sector < sec_w) {
280        const int p = sector - sections.v[0];
281        theta_base = pos[i2 + ne02] * dpct::pow(theta_scale, p);
282    }
283
284    const float freq_factor = has_ff ? freq_factors[i0 / 2] : 1.0f;
285
286    float cos_theta;
287    float sin_theta;
288
289    rope_yarn<forward>(theta_base / freq_factor, freq_scale, corr_dims, i0,
290                       ext_factor, attn_factor, cos_theta, sin_theta);
291
292    const float x0 = x[ix + 0];
293    const float x1 = x[ix + n_dims];
294
295    dst[idst + 0] = x0 * cos_theta - x1 * sin_theta;
296    dst[idst + n_dims] = x0 * sin_theta + x1 * cos_theta;
297}
298
299template <bool forward, typename T, typename D>
300static void
301rope_norm_sycl(const T *x, D *dst, const int ne00, const int ne01,
302               const int ne02, const int s01, const int s02, const int s03,
303               const int s1, const int s2, const int s3, const int n_dims,
304               const int n_offs, const int nr, const int32_t *pos, const float freq_scale,
305               const float freq_base, const float ext_factor,
306               const float attn_factor, const rope_corr_dims corr_dims,
307               const float *freq_factors, const int64_t *row_indices,
308               const int set_rows_stride, dpct::queue_ptr stream) {
309    GGML_ASSERT(ne00 % 2 == 0);
310    const dpct::dim3 block_dims(1, SYCL_ROPE_BLOCK_SIZE, 1);
311    const int n_blocks_x =
312        (ne00 + 2 * SYCL_ROPE_BLOCK_SIZE - 1) / (2 * SYCL_ROPE_BLOCK_SIZE);
313    const dpct::dim3 block_nums(nr, n_blocks_x, 1);
314
315    const float theta_scale = powf(freq_base, -2.0f / n_dims);
316
317    if (freq_factors == nullptr) {
318        stream->parallel_for(
319            sycl::nd_range<3>(block_nums * block_dims, block_dims),
320            [=](sycl::nd_item<3> item_ct1) {
321                GGML_UNUSED(item_ct1);
322                rope_norm<forward, false>(
323                    x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims,
324                    n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims,
325                    theta_scale, freq_factors, row_indices, set_rows_stride);
326            });
327    } else {
328        stream->parallel_for(
329            sycl::nd_range<3>(block_nums * block_dims, block_dims),
330            [=](sycl::nd_item<3> item_ct1) {
331                GGML_UNUSED(item_ct1);
332                rope_norm<forward, true>(
333                    x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims,
334                    n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims,
335                    theta_scale, freq_factors, row_indices, set_rows_stride);
336            });
337    }
338}
339
340template <bool forward, typename T, typename D>
341static void
342rope_neox_sycl(const T *x, D *dst, const int ne00, const int ne01,
343               const int ne02, const int s01, const int s02, const int s03,
344               const int s1, const int s2, const int s3, const int n_dims,
345               const int n_offs, const int nr, const int32_t *pos, const float freq_scale,
346               const float freq_base, const float ext_factor,
347               const float attn_factor, const rope_corr_dims corr_dims,
348               const float *freq_factors, const int64_t *row_indices,
349               const int set_rows_stride, dpct::queue_ptr stream) {
350    GGML_ASSERT(ne00 % 2 == 0);
351    const dpct::dim3 block_dims(1, SYCL_ROPE_BLOCK_SIZE, 1);
352    const int n_blocks_x =
353        (ne00 + 2 * SYCL_ROPE_BLOCK_SIZE - 1) / (2 * SYCL_ROPE_BLOCK_SIZE);
354    const dpct::dim3 block_nums(nr, n_blocks_x, 1);
355
356    const float theta_scale = powf(freq_base, -2.0f / n_dims);
357
358    if (freq_factors == nullptr) {
359        stream->parallel_for(
360            sycl::nd_range<3>(block_nums * block_dims, block_dims),
361            [=](sycl::nd_item<3> item_ct1) {
362                GGML_UNUSED(item_ct1);
363                rope_neox<forward, false>(
364                    x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims,
365                    n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims,
366                    theta_scale, freq_factors, row_indices, set_rows_stride);
367            });
368    } else {
369        stream->parallel_for(
370            sycl::nd_range<3>(block_nums * block_dims, block_dims),
371            [=](sycl::nd_item<3> item_ct1) {
372                GGML_UNUSED(item_ct1);
373                rope_neox<forward, true>(
374                    x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims,
375                    n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims,
376                    theta_scale, freq_factors, row_indices, set_rows_stride);
377            });
378    }
379}
380
381template <bool forward, typename T>
382static void
383rope_multi_sycl(const T *x, T *dst, const int ne00, const int ne01,
384                const int ne02, const int s01, const int s02, const int s03,
385                const int s1, const int s2, const int s3, const int n_dims,
386                const int n_offs, const int nr, const int32_t *pos, const float freq_scale,
387                const float freq_base, const float ext_factor,
388                const float attn_factor, const rope_corr_dims corr_dims,
389                const float *freq_factors, const mrope_sections sections,
390                const bool is_imrope, dpct::queue_ptr stream) {
391    GGML_ASSERT(ne00 % 2 == 0);
392    const dpct::dim3 block_dims(1, SYCL_ROPE_BLOCK_SIZE, 1);
393    const int n_blocks_x =
394        (ne00 + 2 * SYCL_ROPE_BLOCK_SIZE - 1) / (2 * SYCL_ROPE_BLOCK_SIZE);
395    const dpct::dim3 block_nums(nr, n_blocks_x, 1);
396
397    const float theta_scale = powf(freq_base, -2.0f / n_dims);
398
399    if (freq_factors == nullptr) {
400        stream->parallel_for(
401            sycl::nd_range<3>(block_nums * block_dims, block_dims),
402            [=](sycl::nd_item<3> item_ct1) {
403                GGML_UNUSED(item_ct1);
404                rope_multi<forward, false, T>(
405                    x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims,
406                    n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims,
407                    theta_scale, freq_factors, sections, is_imrope);
408            });
409    } else {
410        stream->parallel_for(
411            sycl::nd_range<3>(block_nums * block_dims, block_dims),
412            [=](sycl::nd_item<3> item_ct1) {
413                GGML_UNUSED(item_ct1);
414                rope_multi<forward, true, T>(
415                    x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims,
416                    n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims,
417                    theta_scale, freq_factors, sections, is_imrope);
418            });
419    }
420}
421
422template <bool forward, typename T>
423static void
424rope_vision_sycl(const T *x, T *dst, const int ne00, const int ne01,
425                 const int ne02, const int s01, const int s02, const int s03,
426                 const int s1, const int s2, const int s3, const int n_dims,
427                 const int nr, const int32_t *pos, const float freq_scale,
428                 const float freq_base, const float ext_factor,
429                 const float attn_factor, const rope_corr_dims corr_dims,
430                 const float *freq_factors, const mrope_sections sections,
431                 dpct::queue_ptr stream) {
432    GGML_ASSERT(ne00 % 2 == 0);
433    const dpct::dim3 block_dims(1, SYCL_ROPE_BLOCK_SIZE, 1);
434    const int n_blocks_x =
435        (ne00 + 2 * SYCL_ROPE_BLOCK_SIZE - 1) / (2 * SYCL_ROPE_BLOCK_SIZE);
436    const dpct::dim3 block_nums(nr, n_blocks_x, 1);
437
438    const float theta_scale = powf(freq_base, -2.0f / n_dims);
439
440    if (freq_factors == nullptr) {
441        stream->parallel_for(
442            sycl::nd_range<3>(block_nums * block_dims, block_dims),
443            [=](sycl::nd_item<3> item_ct1) {
444                GGML_UNUSED(item_ct1);
445                rope_vision<forward, false, T>(
446                    x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims,
447                    pos, freq_scale, ext_factor, attn_factor, corr_dims,
448                    theta_scale, freq_factors, sections);
449            });
450    } else {
451        stream->parallel_for(
452            sycl::nd_range<3>(block_nums * block_dims, block_dims),
453            [=](sycl::nd_item<3> item_ct1) {
454                GGML_UNUSED(item_ct1);
455                rope_vision<forward, true, T>(
456                    x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims,
457                    pos, freq_scale, ext_factor, attn_factor, corr_dims,
458                    theta_scale, freq_factors, sections);
459            });
460    }
461}
462
463template <bool forward>
464void ggml_sycl_op_rope_impl(ggml_backend_sycl_context &ctx, ggml_tensor *dst,
465                            const ggml_tensor *set_rows = nullptr) {
466    const ggml_tensor *src0 = dst->src[0];
467    const ggml_tensor *src1 = dst->src[1];
468    const ggml_tensor *src2 = dst->src[2];
469
470    const float *src0_d = (const float *)src0->data;
471    const float *src1_d = (const float *)src1->data;
472
473    void *dst_d = dst->data;
474    const int64_t *row_indices = nullptr;
475    ggml_type dst_type = dst->type;
476    int set_rows_stride = 0;
477
478    if (set_rows != nullptr) {
479        GGML_ASSERT(forward);
480        dst_d = set_rows->data;
481        row_indices = (const int64_t *)set_rows->src[1]->data;
482        dst_type = set_rows->type;
483        set_rows_stride = set_rows->nb[1] / ggml_type_size(set_rows->type);
484    }
485    dpct::queue_ptr stream = ctx.stream();
486
487    GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);
488    GGML_ASSERT(dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_F16);
489    GGML_ASSERT(src0->type == dst->type ||
490                (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16));
491
492    const int64_t ne00 = src0->ne[0]; // head dims
493    const int64_t ne01 = src0->ne[1]; // num heads
494    const int64_t ne02 = src0->ne[2]; // num heads
495    const int64_t nr = ggml_nrows(src0);
496
497    const size_t s01 = src0->nb[1] / ggml_type_size(src0->type);
498    const size_t s02 = src0->nb[2] / ggml_type_size(src0->type);
499    const size_t s03 = src0->nb[3] / ggml_type_size(src0->type);
500
501    const size_t s1 = dst->nb[1] / ggml_type_size(dst->type);
502    const size_t s2 = dst->nb[2] / ggml_type_size(dst->type);
503    const size_t s3 = dst->nb[3] / ggml_type_size(dst->type);
504
505    const int n_dims = ((int32_t *)dst->op_params)[1];
506    const int mode = ((int32_t *)dst->op_params)[2];
507    const int n_ctx_orig = ((int32_t *)dst->op_params)[4];
508    const int n_offs = ((int32_t *)dst->op_params)[15];
509    mrope_sections sections;
510
511    float freq_base;
512    float freq_scale;
513    float ext_factor;
514    float attn_factor;
515    float beta_fast;
516    float beta_slow;
517
518    memcpy(&freq_base, (int32_t *)dst->op_params + 5, sizeof(float));
519    memcpy(&freq_scale, (int32_t *)dst->op_params + 6, sizeof(float));
520    memcpy(&ext_factor, (int32_t *)dst->op_params + 7, sizeof(float));
521    memcpy(&attn_factor, (int32_t *)dst->op_params + 8, sizeof(float));
522    memcpy(&beta_fast, (int32_t *)dst->op_params + 9, sizeof(float));
523    memcpy(&beta_slow, (int32_t *)dst->op_params + 10, sizeof(float));
524    memcpy(&sections.v, (int32_t *)dst->op_params + 11, sizeof(int) * 4);
525
526    const bool is_neox = mode & GGML_ROPE_TYPE_NEOX;
527    const bool is_mrope = mode & GGML_ROPE_TYPE_MROPE;
528    const bool is_imrope = mode == GGML_ROPE_TYPE_IMROPE;
529    const bool is_vision = mode == GGML_ROPE_TYPE_VISION;
530
531    if (is_mrope) {
532        GGML_ASSERT(sections.v[0] > 0 || sections.v[1] > 0 ||
533                    sections.v[2] > 0);
534    }
535
536    if (is_vision) {
537        GGML_ASSERT(n_dims == ne00 / 2);
538        GGML_ASSERT(n_offs == 0); // offset not supported for vision, as the rotated pairs span the whole row
539    }
540
541    const int32_t *pos = (const int32_t *)src1_d;
542
543    const float *freq_factors = nullptr;
544    if (src2 != nullptr) {
545        freq_factors = (const float *)src2->data;
546    }
547
548    rope_corr_dims corr_dims;
549    ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast,
550                             beta_slow, corr_dims.v);
551
552    // compute
553    if (is_neox) {
554        GGML_SYCL_DEBUG("%s: neox path\n", __func__);
555        if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) {
556            rope_neox_sycl<forward, float, float>(
557                (const float *)src0_d, (float *)dst_d, ne00, ne01, ne02, s01,
558                s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
559                ext_factor, attn_factor, corr_dims, freq_factors, row_indices,
560                set_rows_stride, stream);
561        } else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) {
562            rope_neox_sycl<forward, float, sycl::half>(
563                (const float *)src0_d, (sycl::half *)dst_d, ne00, ne01, ne02,
564                s01, s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale,
565                freq_base, ext_factor, attn_factor, corr_dims, freq_factors,
566                row_indices, set_rows_stride, stream);
567        } else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) {
568            rope_neox_sycl<forward, sycl::half, sycl::half>(
569                (const sycl::half *)src0_d, (sycl::half *)dst_d, ne00, ne01,
570                ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale,
571                freq_base, ext_factor, attn_factor, corr_dims, freq_factors,
572                row_indices, set_rows_stride, stream);
573        } else {
574            GGML_ABORT("Fatal error: Tensor type unsupported!");
575        }
576    } else if (is_mrope && !is_vision) {
577        GGML_SYCL_DEBUG("%s: mrope path\n", __func__);
578        if (src0->type == GGML_TYPE_F32) {
579            rope_multi_sycl<forward>((const float *)src0_d, (float *)dst_d,
580                                     ne00, ne01, ne02, s01, s02, s03, s1, s2,
581                                     s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
582                                     ext_factor, attn_factor, corr_dims,
583                                     freq_factors, sections, is_imrope, stream);
584        } else if (src0->type == GGML_TYPE_F16) {
585            rope_multi_sycl<forward>(
586                (const sycl::half *)src0_d, (sycl::half *)dst_d, ne00, ne01,
587                ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale,
588                freq_base, ext_factor, attn_factor, corr_dims, freq_factors,
589                sections, is_imrope, stream);
590        } else {
591            GGML_ABORT("Fatal error: Tensor type unsupported!");
592        }
593    } else if (is_vision) {
594        GGML_SYCL_DEBUG("%s: vision path\n", __func__);
595        if (src0->type == GGML_TYPE_F32) {
596            rope_vision_sycl<forward>(
597                (const float *)src0_d, (float *)dst_d, ne00, ne01, ne02, s01,
598                s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base,
599                ext_factor, attn_factor, corr_dims, freq_factors, sections,
600                stream);
601        } else if (src0->type == GGML_TYPE_F16) {
602            rope_vision_sycl<forward>(
603                (const sycl::half *)src0_d, (sycl::half *)dst_d, ne00, ne01,
604                ne02, s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale,
605                freq_base, ext_factor, attn_factor, corr_dims, freq_factors,
606                sections, stream);
607        } else {
608            GGML_ABORT("Fatal error: Tensor type unsupported!");
609        }
610    } else {
611        GGML_SYCL_DEBUG("%s: norm path\n", __func__);
612        if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) {
613            rope_norm_sycl<forward, float, float>(
614                (const float *)src0_d, (float *)dst_d, ne00, ne01, ne02, s01,
615                s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
616                ext_factor, attn_factor, corr_dims, freq_factors, row_indices,
617                set_rows_stride, stream);
618        } else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) {
619            rope_norm_sycl<forward, float, sycl::half>(
620                (const float *)src0_d, (sycl::half *)dst_d, ne00, ne01, ne02,
621                s01, s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale,
622                freq_base, ext_factor, attn_factor, corr_dims, freq_factors,
623                row_indices, set_rows_stride, stream);
624        } else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) {
625            rope_norm_sycl<forward, sycl::half, sycl::half>(
626                (const sycl::half *)src0_d, (sycl::half *)dst_d, ne00, ne01,
627                ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale,
628                freq_base, ext_factor, attn_factor, corr_dims, freq_factors,
629                row_indices, set_rows_stride, stream);
630        } else {
631            GGML_ABORT("Fatal error: Tensor type unsupported!");
632        }
633    }
634}
635
636void ggml_sycl_rope(ggml_backend_sycl_context &ctx, ggml_tensor *dst) {
637    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/3);
638
639    ggml_sycl_op_rope_impl<true>(ctx, dst);
640}
641
642void ggml_sycl_rope_back(ggml_backend_sycl_context &ctx, ggml_tensor *dst) {
643    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/3);
644    ggml_sycl_op_rope_impl<false>(ctx, dst);
645}
646
647void ggml_sycl_rope_fused(ggml_backend_sycl_context &ctx, ggml_tensor *rope,
648                          ggml_tensor *set_rows) {
649    scope_op_debug_print scope_dbg_print(__func__, rope, /*num_src=*/3);
650    ggml_sycl_op_rope_impl<true>(ctx, rope, set_rows);
651}