char-fork/llama.cpp

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

Intel AI Get-to Market Customer Success and Solutionssycl: add FILL, CUMSUM, DIAG, SOLVE_TRI, SSM_SCAN, GATED_DELTA_NET (#22149)ad0922465

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
4.9 KiB148 linesraw
1#include "cumsum.hpp"
2#include "common.hpp"
3
4#include <algorithm>
5
6#define SYCL_CUMSUM_BLOCK_SIZE 256
7
8static __dpct_inline__ float warp_prefix_inclusive_sum_f32(float x, const sycl::nd_item<3> & item) {
9    return sycl::inclusive_scan_over_group(item.get_sub_group(), x, sycl::plus<float>());
10}
11
12static void cumsum_f32_kernel(
13        const float * __restrict__ src, float * __restrict__ dst,
14        const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03,
15        const int64_t s01, const int64_t s02, const int64_t s03,
16        const int64_t  d1, const int64_t  d2, const int64_t  d3,
17        const sycl::nd_item<3> & item, float * smem) {
18
19    const int tid = item.get_local_id(2);
20    const int block_size = item.get_local_range(2);
21    const int lane = tid % WARP_SIZE;
22    const int warp = tid / WARP_SIZE;
23    const int warps_per_block = block_size / WARP_SIZE;
24
25    float * s_vals      = smem;
26    float * s_warp_sums = smem + block_size;
27    float * s_carry     = smem + block_size + warps_per_block;
28
29    if (tid == 0) {
30        s_carry[0] = 0.0f;
31    }
32    item.barrier(sycl::access::fence_space::local_space);
33
34    const int64_t i3 = item.get_group(0);
35    const int64_t i2 = item.get_group(1);
36    const int64_t i1 = item.get_group(2);
37    if (i3 >= ne03 || i2 >= ne02 || i1 >= ne01) {
38        return;
39    }
40
41    const float * src_row = src + i1 * s01 + i2 * s02 + i3 * s03;
42    float       * dst_row = dst + i1 * d1  + i2 * d2  + i3 * d3;
43
44    constexpr int num_unroll = 4;
45    float temp[num_unroll];
46
47    for (int64_t i = 0; i < ne00; i += num_unroll * block_size) {
48        int64_t idx = i + tid * num_unroll;
49
50        temp[0] = (idx < ne00 ? src_row[idx] : 0.0f);
51#pragma unroll
52        for (int j = 1; j < num_unroll; j++) {
53            temp[j] = temp[j - 1];
54            if (idx + j < ne00) {
55                temp[j] += src_row[idx + j];
56            }
57        }
58
59        float val = (idx < ne00) ? temp[num_unroll - 1] : 0.0f;
60
61        val = warp_prefix_inclusive_sum_f32(val, item);
62        s_vals[tid] = val;
63
64        if (lane == WARP_SIZE - 1) {
65            s_warp_sums[warp] = val;
66        }
67        item.barrier(sycl::access::fence_space::local_space);
68
69        if (warp == 0) {
70            float w = (tid < warps_per_block) ? s_warp_sums[tid] : 0.0f;
71            float inc = warp_prefix_inclusive_sum_f32(w, item);
72            if (tid < warps_per_block) {
73                s_warp_sums[tid] = inc - w;
74            }
75            if (tid == warps_per_block - 1) {
76                s_carry[1] = inc;
77            }
78        }
79        item.barrier(sycl::access::fence_space::local_space);
80
81        float carry = s_carry[0];
82        float final_offset = s_vals[tid] + s_warp_sums[warp] + carry - temp[num_unroll - 1];
83
84#pragma unroll
85        for (int j = 0; j < num_unroll; j++) {
86            if (idx + j < ne00) {
87                dst_row[idx + j] = temp[j] + final_offset;
88            }
89        }
90
91        item.barrier(sycl::access::fence_space::local_space);
92
93        if (tid == 0) {
94            s_carry[0] += s_carry[1];
95        }
96    }
97}
98
99inline void ggml_sycl_op_cumsum(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
100    const ggml_tensor * src0 = dst->src[0];
101
102    GGML_ASSERT(src0->type == GGML_TYPE_F32);
103    GGML_ASSERT(dst->type == GGML_TYPE_F32);
104
105    dpct::queue_ptr stream = ctx.stream();
106    SYCL_CHECK(ggml_sycl_set_device(ctx.device));
107
108    const float * src_d = static_cast<const float *>(src0->data);
109    float       * dst_d = static_cast<float *>(dst->data);
110
111    const int64_t ne00 = src0->ne[0];
112    const int64_t ne01 = src0->ne[1];
113    const int64_t ne02 = src0->ne[2];
114    const int64_t ne03 = src0->ne[3];
115
116    const size_t ts = sizeof(float);
117    const int64_t s01 = src0->nb[1] / ts;
118    const int64_t s02 = src0->nb[2] / ts;
119    const int64_t s03 = src0->nb[3] / ts;
120    const int64_t d1  = dst->nb[1] / ts;
121    const int64_t d2  = dst->nb[2] / ts;
122    const int64_t d3  = dst->nb[3] / ts;
123
124    const int num_warps = (ne00 + WARP_SIZE - 1) / WARP_SIZE;
125    int block_size = num_warps * WARP_SIZE;
126    block_size = std::min(block_size, SYCL_CUMSUM_BLOCK_SIZE);
127    const int warps_per_block = block_size / WARP_SIZE;
128    const int smem_size = block_size + warps_per_block + 2;
129
130    const sycl::range<3> grid(ne03, ne02, ne01);
131    const sycl::range<3> block(1, 1, block_size);
132
133    stream->submit([&](sycl::handler & cgh) {
134        sycl::local_accessor<float, 1> smem_acc(sycl::range<1>(smem_size), cgh);
135        cgh.parallel_for(
136            sycl::nd_range<3>(grid * block, block),
137            [=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
138                cumsum_f32_kernel(src_d, dst_d, ne00, ne01, ne02, ne03,
139                                  s01, s02, s03, d1, d2, d3,
140                                  item, get_pointer(smem_acc));
141            });
142    });
143}
144
145void ggml_sycl_cumsum(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
146    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/1);
147    ggml_sycl_op_cumsum(ctx, dst);
148}