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 →
6.0 KiB172 linesraw
1#include "solve_tri.hpp"
2#include "common.hpp"
3#include <oneapi/mkl/blas.hpp>
4
5template <int n_template, int k_template>
6static void solve_tri_f32_fast(const float * __restrict__ A,
7                               const float * __restrict__ B,
8                               float * __restrict__ X,
9                               const int64_t ne02, [[maybe_unused]] const int64_t ne03,
10                               const int64_t nb02, const int64_t nb03,
11                               const int64_t nb12, const int64_t nb13,
12                               const int64_t nb2,  const int64_t nb3,
13                               const int n_arg, const int k_arg,
14                               const sycl::nd_item<2> & item, float * sA) {
15
16    const int n = n_template == 0 ? n_arg : n_template;
17    const int k = k_template == 0 ? k_arg : k_template;
18
19    const int batch_idx = item.get_group(1);
20    const int lane      = item.get_local_id(1) % WARP_SIZE;
21    const int col_idx   = item.get_local_id(0);
22
23    if (col_idx >= k) {
24        return;
25    }
26
27    const int64_t i03 = batch_idx / ne02;
28    const int64_t i02 = batch_idx % ne02;
29
30    const float * A_batch = (const float *) ((const char *) A + i02 * nb02 + i03 * nb03);
31    const float * B_batch = (const float *) ((const char *) B + i02 * nb12 + i03 * nb13);
32    float *       X_batch = (float *)       ((char *)       X + i02 * nb2  + i03 * nb3);
33
34    const int offset = item.get_local_id(1) + item.get_local_id(0) * item.get_local_range(1);
35
36#pragma unroll
37    for (int i = 0; i < n * n; i += k * WARP_SIZE) {
38        const int i0 = i + offset;
39        if (i0 < n * n) {
40            sA[i0] = A_batch[i0];
41        }
42    }
43
44    item.barrier(sycl::access::fence_space::local_space);
45
46    float x_low  = (lane < n) ? B_batch[lane * k + col_idx] : 0.0f;
47    float x_high = (WARP_SIZE + lane < n) ? B_batch[(WARP_SIZE + lane) * k + col_idx] : 0.0f;
48
49    const int half      = WARP_SIZE;
50    const int nrows_low = (n < half) ? n : half;
51
52#pragma unroll
53    for (int row = 0; row < nrows_low; ++row) {
54        float sum = 0.0f;
55        if (lane < row) {
56            sum += sA[row * n + lane] * x_low;
57        }
58        sum = warp_reduce_sum<WARP_SIZE>(sum);
59        if (lane == row) {
60            x_low = (x_low - sum) / sA[row * n + row];
61        }
62    }
63
64#pragma unroll
65    for (int row = half; row < n; ++row) {
66        float     sum = sA[row * n + lane] * x_low;
67        const int j   = half + lane;
68        if (j < row) {
69            sum += sA[row * n + j] * x_high;
70        }
71        sum = warp_reduce_sum<WARP_SIZE>(sum);
72        if (lane == row - half) {
73            x_high = (x_high - sum) / sA[row * n + row];
74        }
75    }
76
77#pragma unroll
78    for (int rr = 0; rr < 2; ++rr) {
79        const int row = rr * WARP_SIZE + lane;
80        if (row < n) {
81            const float val            = (row < half) ? x_low : x_high;
82            X_batch[row * k + col_idx] = val;
83        }
84    }
85}
86
87static void solve_tri_f32_mkl(dpct::queue_ptr stream,
88                               const float * A, float * X,
89                               int n, int k,
90                               int64_t ne02, [[maybe_unused]] int64_t ne03,
91                               int64_t nb02, [[maybe_unused]] int64_t nb03,
92                               int64_t nb2,  [[maybe_unused]] int64_t nb3) {
93    const float alpha = 1.0f;
94    const int64_t total_batches = ne02 * ne03;
95    if (total_batches == 0) {
96        return;
97    }
98
99    const int64_t stride_a = nb02 / sizeof(float);
100    const int64_t stride_x = nb2 / sizeof(float);
101
102    oneapi::mkl::blas::trsm_batch(
103        *stream,
104        oneapi::mkl::side::right,
105        oneapi::mkl::uplo::upper,
106        oneapi::mkl::transpose::nontrans,
107        oneapi::mkl::diag::nonunit,
108        k, n, alpha,
109        A, n, stride_a,
110        X, k, stride_x,
111        total_batches);
112}
113
114inline void ggml_sycl_op_solve_tri(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
115    const ggml_tensor * src0 = dst->src[0];
116    const ggml_tensor * src1 = dst->src[1];
117
118    GGML_ASSERT(ggml_is_contiguous(src0));
119    GGML_ASSERT(ggml_is_contiguous(src1));
120    GGML_ASSERT(src0->type == GGML_TYPE_F32);
121
122    dpct::queue_ptr stream = ctx.stream();
123    SYCL_CHECK(ggml_sycl_set_device(ctx.device));
124
125    const int n    = src0->ne[0];
126    const int k    = src1->ne[0];
127    const int64_t ne02 = src0->ne[2];
128    const int64_t ne03 = src0->ne[3];
129
130    GGML_ASSERT(n <= SYCL_SOLVE_TRI_MAX_N && k <= SYCL_SOLVE_TRI_MAX_K);
131
132    const float * A_d = static_cast<const float *>(src0->data);
133    const float * B_d = static_cast<const float *>(src1->data);
134    float * X_d       = static_cast<float *>(dst->data);
135
136    if (X_d != B_d) {
137        const int64_t total_elements = (int64_t)n * k * ne02 * ne03;
138        stream->memcpy(X_d, B_d, total_elements * sizeof(float));
139    }
140
141    const int64_t nb02 = src0->nb[2];
142    const int64_t nb03 = src0->nb[3];
143    const int64_t nb12 = src1->nb[2];
144    const int64_t nb13 = src1->nb[3];
145    const int64_t nb2  = dst->nb[2];
146    const int64_t nb3  = dst->nb[3];
147
148    const int64_t total_batches = ne02 * ne03;
149
150    if (n <= 2 * WARP_SIZE && k <= 32) {
151        const int smem_size = 2 * WARP_SIZE * 2 * WARP_SIZE;
152        const sycl::range<2> grid(1, total_batches);
153        const sycl::range<2> block(k, WARP_SIZE);
154        stream->submit([&](sycl::handler & cgh) {
155            sycl::local_accessor<float, 1> smem_acc(sycl::range<1>(smem_size), cgh);
156            cgh.parallel_for(
157                sycl::nd_range<2>(grid * block, block),
158                [=](sycl::nd_item<2> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
159                    solve_tri_f32_fast<0, 0>(A_d, B_d, X_d, ne02, ne03,
160                                              nb02, nb03, nb12, nb13, nb2, nb3,
161                                              n, k, item, get_pointer(smem_acc));
162                });
163        });
164    } else {
165        solve_tri_f32_mkl(stream, A_d, X_d, n, k, ne02, ne03, nb02, nb03, nb2, nb3);
166    }
167}
168
169void ggml_sycl_solve_tri(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
170    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
171    ggml_sycl_op_solve_tri(ctx, dst);
172}