char-fork/llama.cpp

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

Ruben Ortlamllama: add default load-mode auto, which avoids mmap on iGPUs (#26081)153d324bc

char/aliceai-t5-moe
30.9 KiB838 linesraw
1#include "ggml-zendnn.h"
2
3#include "ggml-backend-impl.h"
4#include "ggml-impl.h"
5
6#define GGML_COMMON_DECL_CPP
7#include "ggml-common.h"
8
9#include "zendnnl.hpp"
10
11#include <cstring>
12
13
14struct ggml_backend_zendnn_context {
15    int n_threads = GGML_DEFAULT_N_THREADS;
16    std::unique_ptr<char[]> work_data;
17    size_t work_size = 0;
18};
19
20template<typename T>
21zendnnl::common::data_type_t ggml_to_zendnn_type() {
22    if constexpr (std::is_same_v<T, float>) {
23        return zendnnl::common::data_type_t::f32;
24    } else if constexpr (std::is_same_v<T, ggml_bf16_t>) {
25        return zendnnl::common::data_type_t::bf16;
26    } else if constexpr (std::is_same_v<T, block_q8_0>) {
27        return zendnnl::common::data_type_t::s8;
28    } else {
29        return zendnnl::common::data_type_t::none;
30    }
31}
32
33/**
34 * Builds the matmul_params shared by ggml_zendnn_matmul() and ggml_zendnn_group_matmul():
35 * dtype selection plus, for Q8_0 weights, dynamic-quant setup. Callers still need to set
36 * quant_params.src_scale.dims themselves, since that depends on the batch size(s) in use.
37 */
38template <typename TA, typename TB, typename TC>
39static zendnnl::lowoha::matmul::matmul_params ggml_zendnn_make_matmul_params(ggml_backend_zendnn_context * ctx) {
40    zendnnl::lowoha::matmul::matmul_params params;
41    params.dtypes.src = ggml_to_zendnn_type<TB>();
42    params.dtypes.wei = ggml_to_zendnn_type<TA>();
43    params.dtypes.dst = ggml_to_zendnn_type<TC>();
44    params.num_threads = ctx->n_threads;
45
46    if constexpr (std::is_same_v<TA, block_q8_0>) {
47        params.dtypes.compute = zendnnl::common::data_type_t::s8;
48        params.dynamic_quant = true;
49        params.quant_params.src_scale.buff = nullptr;
50        params.quant_params.src_scale.dt   = zendnnl::common::data_type_t::bf16;
51        params.packing.pack_format_b = 1;
52    }
53    return params;
54}
55
56/**
57 * ZenDNN matmul: computes C = B * A.
58 *
59 * - A: weights, shape (k, m), column-major (each column is a weight vector for one output).
60 * - B: input, shape (n, k), row-major (each row is an input sample).
61 * - C: output, shape (n, m), row-major.
62 *
63 * Dimensions:
64 *   m = output features (columns of C, columns of A)
65 *   n = batch size      (rows of C, rows of B)
66 *   k = inner dimension (columns of B, rows of A)
67 */
68template <typename TA, typename TB, typename TC>
69static bool ggml_zendnn_matmul(ggml_backend_zendnn_context * ctx, int64_t m, int64_t n, int64_t k,
70                               const TA * A, int64_t lda, const TB * B, int64_t ldb, TC * C,
71                               int64_t ldc) {
72
73    zendnnl::lowoha::matmul::matmul_params params = ggml_zendnn_make_matmul_params<TA, TB, TC>(ctx);
74
75    zendnnl::lowoha::matmul::matmul_batch_params_t batch_params;
76
77    if constexpr (std::is_same_v<TA, block_q8_0>) {
78        params.quant_params.src_scale.dims = {n, k / QK8_0};
79    }
80
81    zendnnl::error_handling::status_t status = zendnnl::lowoha::matmul::matmul_direct(
82        'r', false, true,   // row-major, don't transpose B, transpose A (because it's column-major)
83        n,                  // M: rows of B and C
84        m,                  // N: cols of A^T and C
85        k,                  // K: cols of B, rows of A
86        1.0f,               // alpha
87        B, ldb,             // src: B[n,k]
88        A, lda,             // weight: A[k,m] column-major (transposed)
89        nullptr,            // bias
90        0.0f,               // beta
91        C, ldc,             // output C[n,m]
92        true,               // is_weights_const
93        batch_params,       // batch_params
94        params              // params
95    );
96
97    if (status != zendnnl::error_handling::status_t::success) {
98        GGML_LOG_ERROR("%s, ZenDNN matmul failed: status=%d\n", __func__, static_cast<int>(status));
99        return false;
100    }
101    return true;
102}
103
104static bool ggml_zendnn_gemm(ggml_backend_zendnn_context * ctx, int64_t m, int64_t n, int64_t k,
105                              const void * A, int64_t lda, const void * B, int64_t ldb, void * C,
106                              int64_t ldc, int Atype, int Btype, int Ctype) {
107
108    assert(m >= 0);
109    assert(n >= 0);
110    assert(k >= 0);
111    assert(lda >= k);
112    assert(ldb >= k);
113    assert(ldc >= m);
114
115    // categorize types
116    switch (Atype) {
117        case GGML_TYPE_F32:
118            if (Btype != GGML_TYPE_F32 || Ctype != GGML_TYPE_F32)
119                return false;
120            return ggml_zendnn_matmul<float, float, float>(
121                ctx, m, n, k,
122                (const float *)A, lda,
123                (const float *)B, ldb,
124                (float *)C, ldc);
125        case GGML_TYPE_BF16:
126            if (Btype != GGML_TYPE_BF16)
127                return false;
128            if (Ctype == GGML_TYPE_BF16)
129                return ggml_zendnn_matmul<ggml_bf16_t, ggml_bf16_t, ggml_bf16_t>(
130                    ctx, m, n, k,
131                    (const ggml_bf16_t *)A, lda,
132                    (const ggml_bf16_t *)B, ldb,
133                    (ggml_bf16_t *)C, ldc);
134            if (Ctype == GGML_TYPE_F32)
135                return ggml_zendnn_matmul<ggml_bf16_t, ggml_bf16_t, float>(
136                    ctx, m, n, k,
137                    (const ggml_bf16_t *)A, lda,
138                    (const ggml_bf16_t *)B, ldb,
139                    (float *)C, ldc);
140            return false;
141        case GGML_TYPE_Q8_0:
142            if (Btype != GGML_TYPE_F32 || Ctype != GGML_TYPE_F32)
143                return false;
144            return ggml_zendnn_matmul<block_q8_0, float, float>(
145                ctx, m, n, k,
146                (const block_q8_0 *)A, lda,
147                (const float *)B, ldb,
148                (float *)C, ldc);
149        default:
150            return false; // unsupported type
151    }
152}
153
154static void ggml_zendnn_compute_forward_mul_mat(
155    ggml_backend_zendnn_context * ctx,
156    ggml_tensor * dst) {
157
158    const ggml_tensor * src0 = dst->src[0];  // weights
159    const ggml_tensor * src1 = dst->src[1];  // inputs
160
161    GGML_TENSOR_BINARY_OP_LOCALS
162
163    ggml_type         const vec_dot_type = src0->type;
164    ggml_from_float_t const from_float = ggml_get_type_traits(vec_dot_type)->from_float_ref;
165
166    GGML_ASSERT(ne0 == ne01);
167    GGML_ASSERT(ne1 == ne11);
168    GGML_ASSERT(ne2 == ne12);
169    GGML_ASSERT(ne3 == ne13);
170
171    // we don't support permuted src0 or src1
172    GGML_ASSERT(nb00 == ggml_type_size(src0->type));
173    GGML_ASSERT(nb10 == ggml_type_size(src1->type));
174
175    // dst cannot be transposed or permuted
176    GGML_ASSERT(nb0 == sizeof(float));
177    GGML_ASSERT(nb0 <= nb1);
178    GGML_ASSERT(nb1 <= nb2);
179    GGML_ASSERT(nb2 <= nb3);
180
181    // broadcast factors
182    const int64_t r2 = ne12/ne02;
183    const int64_t r3 = ne13/ne03;
184
185    void * work_data = ctx->work_data.get();
186
187    // ZenDNN requires FP32 for dynamic quantization, so conversion is skipped
188    if (src1->type != vec_dot_type && src0->type != GGML_TYPE_Q8_0) {
189        const size_t nbw1 = ggml_row_size(vec_dot_type, ne10);
190        const size_t nbw2 = nbw1 * ne11;
191        const size_t nbw3 = nbw2 * ne12;
192        const size_t desired_wsize = ne13 * nbw3;
193        if (ctx->work_size < desired_wsize) {
194            ctx->work_data.reset(new char[desired_wsize]);
195            ctx->work_size = desired_wsize;
196        }
197        work_data = ctx->work_data.get();
198
199        // #pragma omp parallel for num_threads(ctx->n_threads)
200        #pragma omp parallel for collapse(3) num_threads(ctx->n_threads) schedule(static)
201        for (int64_t i13 = 0; i13 < ne13; ++i13) {
202            for (int64_t i12 = 0; i12 < ne12; ++i12) {
203                for (int64_t i11 = 0; i11 < ne11; ++i11) {
204                    const float * src1_f32 = (float *)((char *)src1->data + i11*nb11 + i12*nb12 + i13*nb13);
205                    void * src1_conv = (char *)work_data + i11*nbw1 + i12*nbw2 + i13*nbw3;
206                    from_float(src1_f32, src1_conv, ne10);
207                }
208            }
209        }
210    }
211
212    for (int64_t i13 = 0; i13 < ne13; i13++) {
213        for (int64_t i12 = 0; i12 < ne12; i12++) {
214            const void* wdata = (src1->type == vec_dot_type || src0->type == GGML_TYPE_Q8_0) ? src1->data : work_data;
215            const size_t row_size = ggml_row_size(vec_dot_type, ne10);
216            if (!ggml_zendnn_gemm(ctx,
217                                  ne01,     // m
218                                  ne11,     // n
219                                  ne10,     // k
220                                  static_cast<const char *>(src0->data) + (i12/r2)*nb02 + (i13/r3)*nb03,
221                                  ne00,     // lda
222                                  static_cast<const char *>(wdata) + (i12*ne11 + i13*ne12*ne11)*row_size,
223                                  ne10,     // ldb
224                                  static_cast<char *>(dst->data) + i12*nb2 + i13*nb3,
225                                  ne01,     // ldc
226                                  src0->type,
227                                  src0->type == GGML_TYPE_Q8_0 ? GGML_TYPE_F32 : vec_dot_type,
228                                  dst->type))
229                GGML_ABORT("%s: ZenDNN gemm failed\n", __func__);
230        }
231    }
232}
233
234struct mmid_row_mapping {
235    int32_t i1;
236    int32_t i2;
237};
238
239/**
240 * ZenDNN batched matmul: computes C[i] = B[i] * A[i] for every active expert i via a single
241 * group_matmul_direct() call. Batched analogue of ggml_zendnn_matmul() - see its docs for the
242 * per-expert A/B/C shape convention. m and k are shared by every expert; n (batch size) varies
243 * per expert, hence the vector.
244 */
245template <typename TA, typename TB, typename TC>
246static bool ggml_zendnn_group_matmul(ggml_backend_zendnn_context * ctx, int64_t m, int64_t k,
247                                     const std::vector<int64_t> & n,
248                                     const std::vector<const void *> & A, int64_t lda,
249                                     const std::vector<const void *> & B, int64_t ldb,
250                                     const std::vector<void *> & C, int64_t ldc) {
251
252    const int n_experts = n.size();
253
254    zendnnl::lowoha::matmul::matmul_params base_params = ggml_zendnn_make_matmul_params<TA, TB, TC>(ctx);
255
256    std::vector<char>         layout(n_experts, 'r');
257    std::vector<bool>         trans_a(n_experts, false);
258    std::vector<bool>         trans_b(n_experts, true);
259    std::vector<int>          batch_m(n_experts);
260    std::vector<int>          batch_n(n_experts, m);
261    std::vector<int>          batch_k(n_experts, k);
262    std::vector<float>        alpha(n_experts, 1.0f);
263    std::vector<float>        beta(n_experts, 0.0f);
264    std::vector<const void *> bias(n_experts, nullptr);
265    std::vector<int>          lda_v(n_experts, lda);
266    std::vector<int>          ldb_v(n_experts, ldb);
267    std::vector<int>          ldc_v(n_experts, ldc);
268    std::vector<bool>         is_wei_const(n_experts, true);
269    std::vector<zendnnl::lowoha::matmul::matmul_params> params(n_experts, base_params);
270
271    for (int i = 0; i < n_experts; i++) {
272        batch_m[i] = n[i];
273
274        // src_scale.dims depends on this expert's row count, unlike the rest of base_params
275        if constexpr (std::is_same_v<TA, block_q8_0>) {
276            params[i].quant_params.src_scale.dims = {n[i], k / QK8_0};
277        }
278    }
279
280    zendnnl::error_handling::status_t status = zendnnl::lowoha::matmul::group_matmul_direct(
281        layout, trans_a, trans_b, batch_m, batch_n, batch_k, alpha,
282        B, ldb_v, A, lda_v, bias, beta,
283        C, ldc_v, is_wei_const, params);
284
285    if (status != zendnnl::error_handling::status_t::success) {
286        GGML_LOG_ERROR("%s, ZenDNN group matmul failed: status=%d\n", __func__, static_cast<int>(status));
287        return false;
288    }
289    return true;
290}
291
292static bool ggml_zendnn_group_gemm(ggml_backend_zendnn_context * ctx, int64_t m, int64_t k,
293                                    const std::vector<int64_t> & n,
294                                    const std::vector<const void *> & A, int64_t lda,
295                                    const std::vector<const void *> & B, int64_t ldb,
296                                    const std::vector<void *> & C, int64_t ldc,
297                                    int Atype, int Btype, int Ctype) {
298
299    assert(m >= 0);
300    for (size_t i = 0; i < n.size(); i++) {
301        assert(n[i] >= 0);
302    }
303    assert(k >= 0);
304    assert(lda >= k);
305    assert(ldb >= k);
306    assert(ldc >= m);
307
308    // categorize types
309    switch (Atype) {
310        case GGML_TYPE_F32:
311            if (Btype != GGML_TYPE_F32 || Ctype != GGML_TYPE_F32)
312                return false;
313            return ggml_zendnn_group_matmul<float, float, float>(ctx, m, k, n, A, lda, B, ldb, C, ldc);
314        case GGML_TYPE_BF16:
315            if (Btype != GGML_TYPE_BF16)
316                return false;
317            if (Ctype == GGML_TYPE_BF16)
318                return ggml_zendnn_group_matmul<ggml_bf16_t, ggml_bf16_t, ggml_bf16_t>(
319                    ctx, m, k, n, A, lda, B, ldb, C, ldc);
320            if (Ctype == GGML_TYPE_F32)
321                return ggml_zendnn_group_matmul<ggml_bf16_t, ggml_bf16_t, float>(ctx, m, k, n, A, lda, B, ldb, C, ldc);
322            return false;
323        case GGML_TYPE_Q8_0:
324            if (Btype != GGML_TYPE_F32 || Ctype != GGML_TYPE_F32)
325                return false;
326            return ggml_zendnn_group_matmul<block_q8_0, float, float>(ctx, m, k, n, A, lda, B, ldb, C, ldc);
327        default:
328            return false; // unsupported type
329    }
330}
331
332static void ggml_zendnn_compute_forward_mul_mat_id(
333    ggml_backend_zendnn_context * ctx,
334    ggml_tensor * dst) {
335
336    const ggml_tensor * src0 = dst->src[0];  // expert weights
337    const ggml_tensor * src1 = dst->src[1];  // inputs
338    const ggml_tensor * ids  = dst->src[2];  // expert ids
339
340    GGML_TENSOR_BINARY_OP_LOCALS
341
342    // exit for no tokens to process
343    if (ne2 == 0 || ne11 == 0) {
344        return;
345    }
346
347    ggml_type         const vec_dot_type = src0->type;
348    ggml_from_float_t const from_float = ggml_get_type_traits(vec_dot_type)->from_float_ref;
349
350    // we don't support permuted src0 or src1
351    GGML_ASSERT(nb00 == ggml_type_size(src0->type));
352    GGML_ASSERT(nb10 == ggml_type_size(src1->type));
353
354    // dst cannot be transposed or permuted
355    GGML_ASSERT(nb0 == sizeof(float));
356    GGML_ASSERT(nb0 <= nb1);
357    GGML_ASSERT(nb1 <= nb2);
358    GGML_ASSERT(nb2 <= nb3);
359
360    GGML_ASSERT(ne03 == 1);
361    GGML_ASSERT(ne13 == 1);
362    GGML_ASSERT(ne3  == 1);
363
364    // row groups
365    const int n_ids = ids->ne[0]; // n_expert_used
366    const int n_as  = ne02;       // n_experts
367
368    std::vector<int64_t> matrix_row_counts(n_as, 0);
369    std::vector<std::vector<mmid_row_mapping>> matrix_rows(n_as);
370
371    int64_t total_rows = 0;
372    int     n_active_experts = 0;
373    // group rows by expert (preprocessing step)
374    for (int64_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
375        for (int id = 0; id < n_ids; ++id) {
376            const int32_t i02 = *(const int32_t *)((const char *)ids->data + iid1*ids->nb[1] + id*ids->nb[0]);
377
378            GGML_ASSERT(i02 >= 0 && i02 < n_as);
379
380            if (matrix_row_counts[i02] == 0) {
381                n_active_experts++;
382            }
383            matrix_rows[i02].push_back({id, iid1});
384            matrix_row_counts[i02]++;
385            total_rows++;
386        }
387    }
388
389    if (total_rows == 0) {
390        return; // no rows to process
391    }
392
393    const size_t row_size = ggml_row_size(vec_dot_type, ne10);
394
395    // For Q8_0, src1 is always F32; the gather buffer must hold F32 rows (ne10*4 bytes),
396    // not Q8_0-encoded rows (row_size ≈ ne10/32*34 bytes) — they differ by ~4x.
397    const size_t f32_row_size = (size_t)ne10 * sizeof(float);
398    const size_t gather_row_size = (src0->type == GGML_TYPE_Q8_0) ? f32_row_size : row_size;
399
400    if (src1->type != vec_dot_type && src0->type != GGML_TYPE_Q8_0) {
401        GGML_ASSERT(src1->type == GGML_TYPE_F32);
402    }
403
404    // size for MoE gather/scatter buffers
405    const size_t wdata_cur_size = total_rows * gather_row_size;
406    const size_t dst_cur_size = total_rows * ggml_row_size(dst->type, ne01);
407
408    // allocate single buffer for all needs
409    const size_t total_size = wdata_cur_size + dst_cur_size;
410    if (ctx->work_size < total_size) {
411        ctx->work_data.reset(new char[total_size]);
412        ctx->work_size = total_size;
413    }
414
415    // partition the buffer
416    char * wdata_cur = ctx->work_data.get();
417    char * dst_cur = wdata_cur + wdata_cur_size;
418
419    // per-expert data collected during gather, handed to ggml_zendnn_group_gemm() as one batch
420    std::vector<int64_t>      expert_row_count(n_active_experts);
421    std::vector<const void *> batch_src(n_active_experts);
422    std::vector<const void *> batch_wei(n_active_experts);
423    std::vector<void *>       batch_dst(n_active_experts);
424
425    // precompute per-expert buffer offsets and batch indices for the parallel loop below
426    std::vector<int64_t> expert_wdata_off(n_as, 0);
427    std::vector<int64_t> expert_dst_off(n_as, 0);
428    std::vector<int>     expert_batch_idx(n_as, -1);
429    {
430        int64_t w_off = 0;
431        int64_t d_off = 0;
432        int     batch_idx = 0;
433        for (int64_t cur_a = 0; cur_a < n_as; ++cur_a) {
434            if (matrix_row_counts[cur_a] == 0) {
435                continue;
436            }
437            expert_wdata_off[cur_a] = w_off;
438            expert_dst_off[cur_a]   = d_off;
439            expert_batch_idx[cur_a] = batch_idx;
440            w_off += matrix_row_counts[cur_a] * gather_row_size;
441            d_off += matrix_row_counts[cur_a] * ggml_row_size(dst->type, ne01);
442            batch_idx++;
443        }
444    }
445
446    // gather + inline-convert input rows into each expert's batch slot
447    #pragma omp parallel for num_threads(ctx->n_threads) schedule(static)
448    for (int64_t cur_a = 0; cur_a < n_as; ++cur_a) {
449        const int64_t cne1 = matrix_row_counts[cur_a];
450
451        if (cne1 == 0) {
452            continue;
453        }
454
455        const int64_t w_off     = expert_wdata_off[cur_a];
456        const int64_t d_off     = expert_dst_off[cur_a];
457        const int     batch_idx = expert_batch_idx[cur_a];
458
459        for (int64_t ir1 = 0; ir1 < cne1; ++ir1) {
460            const mmid_row_mapping & row_mapping = matrix_rows[cur_a][ir1];
461            const int64_t id  = row_mapping.i1;
462            const int64_t i11 = id % ne11;
463            const int64_t i12 = row_mapping.i2;
464
465            const char * src_row = (const char *) src1->data + i11*nb11 + i12*nb12;
466            void *       dst_row = wdata_cur + w_off + ir1 * gather_row_size;
467
468            if (src1->type != vec_dot_type && src0->type != GGML_TYPE_Q8_0) {
469                from_float((const float *) src_row, dst_row, ne10);
470            } else {
471                // no conversion: src1 already matches vec_dot_type, or src0 is Q8_0, whose
472                // ZenDNN dynamic quantization requires the row to stay in F32
473                std::memcpy(dst_row, src_row, gather_row_size);
474            }
475        }
476
477        expert_row_count[batch_idx] = cne1;
478        batch_src[batch_idx]        = wdata_cur + w_off;
479        batch_wei[batch_idx]        = (const char *) src0->data + cur_a * nb02;
480        batch_dst[batch_idx]        = dst_cur + d_off;
481    }
482
483    if (!ggml_zendnn_group_gemm(ctx,
484                                ne01,             // m
485                                ne10,             // k
486                                expert_row_count, // n (per expert)
487                                batch_wei, ne00,  // A: weights, lda
488                                batch_src, ne10,  // B: input, ldb
489                                batch_dst, ne01,  // C: output, ldc
490                                src0->type,
491                                src0->type == GGML_TYPE_Q8_0 ? GGML_TYPE_F32 : vec_dot_type,
492                                dst->type))
493        GGML_ABORT("%s: ZenDNN group gemm failed\n", __func__);
494
495    // scatter output rows to destination
496    #pragma omp parallel for num_threads(ctx->n_threads) schedule(static)
497    for (int64_t cur_a = 0; cur_a < n_as; ++cur_a) {
498        const int64_t cne1 = matrix_row_counts[cur_a];
499
500        if (cne1 == 0) {
501            continue;
502        }
503
504        const int64_t d_off = expert_dst_off[cur_a];
505
506        for (int64_t ir1 = 0; ir1 < cne1; ++ir1) {
507            const mmid_row_mapping & row_mapping = matrix_rows[cur_a][ir1];
508            const int64_t id = row_mapping.i1;
509            const int64_t i1 = id;
510            const int64_t i2 = row_mapping.i2;
511
512            std::memcpy(
513                (char *) dst->data + i1*nb1 + i2*nb2,
514                dst_cur + d_off + ir1 * ggml_row_size(dst->type, ne01),
515                ggml_row_size(dst->type, ne01)
516            );
517        }
518    }
519}
520
521// backend interface
522
523static const char * ggml_backend_zendnn_get_name(ggml_backend_t backend) {
524    return "ZenDNN";
525
526    GGML_UNUSED(backend);
527}
528
529static void ggml_backend_zendnn_free(ggml_backend_t backend) {
530    ggml_backend_zendnn_context * ctx = (ggml_backend_zendnn_context *)backend->context;
531    delete ctx;
532    delete backend;
533}
534
535static ggml_status ggml_backend_zendnn_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) {
536    ggml_backend_zendnn_context * ctx = (ggml_backend_zendnn_context *)backend->context;
537
538    for (int i = 0; i < cgraph->n_nodes; i++) {
539        struct ggml_tensor * node = cgraph->nodes[i];
540
541        if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) {
542            continue;
543        }
544
545        switch (node->op) {
546            case GGML_OP_MUL_MAT:
547                ggml_zendnn_compute_forward_mul_mat(ctx, node);
548                break;
549            case GGML_OP_MUL_MAT_ID:
550                ggml_zendnn_compute_forward_mul_mat_id(ctx, node);
551                break;
552            case GGML_OP_NONE:
553            case GGML_OP_RESHAPE:
554            case GGML_OP_VIEW:
555            case GGML_OP_PERMUTE:
556            case GGML_OP_TRANSPOSE:
557                break;
558
559            default:
560                GGML_ABORT("%s: unsupported op %s\n", __func__, ggml_op_desc(node));
561        }
562    }
563
564    return GGML_STATUS_SUCCESS;
565
566    GGML_UNUSED(backend);
567}
568
569static struct ggml_backend_i ggml_backend_zendnn_i = {
570    /* .get_name                = */ ggml_backend_zendnn_get_name,
571    /* .free                    = */ ggml_backend_zendnn_free,
572    /* .set_tensor_async        = */ NULL,
573    /* .get_tensor_async        = */ NULL,
574    /* .set_tensor_2d_async     = */ NULL,
575    /* .get_tensor_2d_async     = */ NULL,
576    /* .cpy_tensor_async        = */ NULL,
577    /* .synchronize             = */ NULL,
578    /* .graph_plan_create       = */ NULL,
579    /* .graph_plan_free         = */ NULL,
580    /* .graph_plan_update       = */ NULL,
581    /* .graph_plan_compute      = */ NULL,
582    /* .graph_compute           = */ ggml_backend_zendnn_graph_compute,
583    /* .event_record            = */ NULL,
584    /* .event_wait              = */ NULL,
585    /* .graph_optimize          = */ NULL,
586};
587
588static ggml_guid_t ggml_backend_zendnn_guid(void) {
589    static const char * guid_str = "AMD-ZENDNN-ACCEL";
590    return reinterpret_cast<ggml_guid_t>(const_cast<char*>(guid_str));
591}
592
593ggml_backend_t ggml_backend_zendnn_init(void) {
594    ggml_backend_zendnn_context * ctx = new ggml_backend_zendnn_context;
595
596    ggml_backend_t backend = new ggml_backend {
597        /* .guid    = */ ggml_backend_zendnn_guid(),
598        /* .iface   = */ ggml_backend_zendnn_i,
599        /* .device  = */ ggml_backend_reg_dev_get(ggml_backend_zendnn_reg(), 0),
600        /* .context = */ ctx,
601    };
602
603    return backend;
604}
605
606bool ggml_backend_is_zendnn(ggml_backend_t backend) {
607    return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_zendnn_guid());
608}
609
610void ggml_backend_zendnn_set_n_threads(ggml_backend_t backend_zendnn, int n_threads) {
611    GGML_ASSERT(ggml_backend_is_zendnn(backend_zendnn));
612
613    ggml_backend_zendnn_context * ctx = (ggml_backend_zendnn_context *)backend_zendnn->context;
614    ctx->n_threads = n_threads;
615}
616
617// device interface
618static const char * ggml_backend_zendnn_device_get_name(ggml_backend_dev_t dev) {
619    return "ZenDNN";
620
621    GGML_UNUSED(dev);
622}
623/**
624 * ZenDNN is AMD's performance library providing optimized primitives and implementations
625 * for deep learning workloads on AMD CPUs. It targets improved performance for common
626 * neural network operations on AMD architectures. For more information, see:
627 * https://www.amd.com/en/developer/zendnn.html
628 */
629static const char * ggml_backend_zendnn_device_get_description(ggml_backend_dev_t dev) {
630    return "ZenDNN: AMD optimized primitives backend for GGML (optimized for AMD CPUs)";
631
632    GGML_UNUSED(dev);
633}
634
635static void ggml_backend_zendnn_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) {
636    *free  = 0;
637    *total = 0;
638
639    GGML_UNUSED(dev);
640}
641
642static enum ggml_backend_dev_type ggml_backend_zendnn_device_get_type(ggml_backend_dev_t dev) {
643    return GGML_BACKEND_DEVICE_TYPE_ACCEL;
644
645    GGML_UNUSED(dev);
646}
647
648static void ggml_backend_zendnn_device_get_props(ggml_backend_dev_t dev, struct ggml_backend_dev_props * props) {
649    props->name        = ggml_backend_zendnn_device_get_name(dev);
650    props->description = ggml_backend_zendnn_device_get_description(dev);
651    props->type        = ggml_backend_zendnn_device_get_type(dev);
652    ggml_backend_zendnn_device_get_memory(dev, &props->memory_free, &props->memory_total);
653    props->caps = {
654        /* .async                = */ false,
655        /* .host_buffer          = */ false,
656        /* .buffer_from_host_ptr = */ true,
657        /* .events               = */ false,
658        /* .mmap_support         = */ true,
659    };
660}
661
662static ggml_backend_t ggml_backend_zendnn_device_init_backend(ggml_backend_dev_t dev, const char * params) {
663    ggml_backend_t backend = ggml_backend_zendnn_init();
664    if (backend == NULL) {
665        GGML_LOG_ERROR("%s: error: failed to initialize ZenDNN backend\n", __func__);
666        return NULL;
667    }
668
669    return backend;
670
671    GGML_UNUSED(dev);
672    GGML_UNUSED(params);
673}
674
675static ggml_backend_buffer_type_t ggml_backend_zendnn_device_get_buffer_type(ggml_backend_dev_t dev) {
676    return ggml_backend_cpu_buffer_type();
677
678    GGML_UNUSED(dev);
679}
680
681static ggml_backend_buffer_t ggml_backend_zendnn_device_buffer_from_host_ptr(ggml_backend_dev_t dev, void * ptr, size_t size, size_t max_tensor_size) {
682    return ggml_backend_cpu_buffer_from_ptr(ptr, size);
683
684    GGML_UNUSED(dev);
685    GGML_UNUSED(max_tensor_size);
686}
687
688static bool ggml_zendnn_adaptive_fallback_enabled() {
689    static const bool enabled = std::getenv("GGML_ZENDNN_ADAPTIVE_FALLBACK") == nullptr ||
690                                std::atoi(std::getenv("GGML_ZENDNN_ADAPTIVE_FALLBACK")) != 0;
691    return enabled;
692}
693
694static bool ggml_backend_zendnn_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) {
695    switch (op->op) {
696        case GGML_OP_NONE:
697        case GGML_OP_RESHAPE:
698        case GGML_OP_VIEW:
699        case GGML_OP_PERMUTE:
700        case GGML_OP_TRANSPOSE:
701            return true;
702
703        case GGML_OP_MUL_MAT:
704        case GGML_OP_MUL_MAT_ID:
705        {
706            const ggml_tensor * weights = op->src[0];
707            const ggml_tensor * inputs = op->src[1];
708
709            const int64_t ne10 = inputs->ne[0];
710            const int64_t ne0 = op->ne[0];
711            const int64_t ne1 = op->ne[1];
712            const int64_t min_batch = 1;
713
714            if(!ggml_is_contiguous(weights) || !ggml_is_contiguous(inputs)) {
715                return false;
716            }
717
718            if (ggml_zendnn_adaptive_fallback_enabled()) {
719                const int64_t K = inputs->ne[0];
720                const int64_t N = (inputs->ne[1]*inputs->ne[2]*inputs->ne[3]);
721                const int64_t M = weights->ne[1];
722                if(K <= 256 || N <= 128 || M <= 96) {
723                    return false;
724                }
725
726                // MUL_MAT_ID's gather+matmul+scatter approach favors a moderate expert count
727                if (op->op == GGML_OP_MUL_MAT_ID) {
728                    const int64_t n_experts = weights->ne[2];
729                    const int64_t max_experts = 32;
730                    if (n_experts > max_experts) {
731                        return false;
732                    }
733
734                    // fall back once the average rows per expert (N / n_experts) is too thin
735                    // to amortize each per-expert GEMM's overhead
736                    if (N / n_experts <= 32) {
737                        return false;
738                    }
739                }
740            }
741            else if (ne0 < min_batch || ne1 < min_batch || ne10 < min_batch) {
742                return false;
743            }
744
745            switch (weights->type) {
746                case GGML_TYPE_F32:
747                case GGML_TYPE_BF16:
748                case GGML_TYPE_Q8_0:
749                    return true;
750                default:
751                    return false;
752            }
753        } break;
754
755        default:
756            return false;
757    }
758
759    GGML_UNUSED(dev);
760}
761
762static bool ggml_backend_zendnn_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) {
763    return ggml_backend_buft_is_host(buft);
764
765    GGML_UNUSED(dev);
766}
767
768static const struct ggml_backend_device_i ggml_backend_zendnn_device_i = {
769    /* .get_name               = */ ggml_backend_zendnn_device_get_name,
770    /* .get_description        = */ ggml_backend_zendnn_device_get_description,
771    /* .get_memory             = */ ggml_backend_zendnn_device_get_memory,
772    /* .get_type               = */ ggml_backend_zendnn_device_get_type,
773    /* .get_props              = */ ggml_backend_zendnn_device_get_props,
774    /* .init_backend           = */ ggml_backend_zendnn_device_init_backend,
775    /* .get_buffer_type        = */ ggml_backend_zendnn_device_get_buffer_type,
776    /* .get_host_buffer_type   = */ NULL,
777    /* .buffer_from_host_ptr   = */ ggml_backend_zendnn_device_buffer_from_host_ptr,
778    /* .supports_op            = */ ggml_backend_zendnn_device_supports_op,
779    /* .supports_buft          = */ ggml_backend_zendnn_device_supports_buft,
780    /* .offload_op             = */ NULL,
781    /* .event_new              = */ NULL,
782    /* .event_free             = */ NULL,
783    /* .event_synchronize      = */ NULL,
784};
785
786// backend reg interface
787static const char * ggml_backend_zendnn_reg_get_name(ggml_backend_reg_t reg) {
788    return "ZenDNN";
789
790    GGML_UNUSED(reg);
791}
792
793static size_t ggml_backend_zendnn_reg_get_device_count(ggml_backend_reg_t reg) {
794    return 1;
795
796    GGML_UNUSED(reg);
797}
798
799static ggml_backend_dev_t ggml_backend_zendnn_reg_get_device(ggml_backend_reg_t reg, size_t index) {
800    GGML_ASSERT(index == 0);
801
802    static ggml_backend_device ggml_backend_zendnn_device = {
803        /* .iface   = */ ggml_backend_zendnn_device_i,
804        /* .reg     = */ reg,
805        /* .context = */ nullptr,
806    };
807
808    return &ggml_backend_zendnn_device;
809}
810
811static void * ggml_backend_zendnn_get_proc_address(ggml_backend_reg_t reg, const char * name) {
812    if (std::strcmp(name, "ggml_backend_set_n_threads") == 0) {
813        return (void *) ggml_backend_zendnn_set_n_threads;
814    }
815    return NULL;
816
817    GGML_UNUSED(reg);
818    GGML_UNUSED(name);
819}
820
821static const struct ggml_backend_reg_i ggml_backend_zendnn_reg_i = {
822    /* .get_name         = */ ggml_backend_zendnn_reg_get_name,
823    /* .get_device_count = */ ggml_backend_zendnn_reg_get_device_count,
824    /* .get_device       = */ ggml_backend_zendnn_reg_get_device,
825    /* .get_proc_address = */ ggml_backend_zendnn_get_proc_address,
826};
827
828ggml_backend_reg_t ggml_backend_zendnn_reg(void) {
829    static struct ggml_backend_reg ggml_backend_zendnn_reg = {
830        /* .api_version = */ GGML_BACKEND_API_VERSION,
831        /* .iface       = */ ggml_backend_zendnn_reg_i,
832        /* .context     = */ NULL,
833    };
834
835    return &ggml_backend_zendnn_reg;
836}
837
838GGML_BACKEND_DL_IMPL(ggml_backend_zendnn_reg)