char-fork/llama.cpp

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

Łukasz Ślusarczyksycl : fix oneDNN scratchpad breaking the pool free order (#28704)661643e43

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
3.4 KiB95 linesraw
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#ifndef GGML_SYCL_GEMM_HPP
14#define GGML_SYCL_GEMM_HPP
15
16#include "ggml-sycl.h"
17
18#if GGML_SYCL_DNNL
19
20#include "dnnl.hpp"
21#include "dnnl_sycl.hpp"
22
23class DnnlGemmWrapper {
24public:
25    using dt = dnnl::memory::data_type;
26    using tag = dnnl::memory::format_tag;
27
28    template<typename T>
29    static constexpr dt to_dt() {
30        if constexpr (std::is_same_v<T, float>) return dt::f32;
31        else if constexpr (std::is_same_v<T, sycl::half>) return dt::f16;
32#ifdef GGML_SYCL_HAS_BF16
33        else if constexpr (std::is_same_v<T, sycl::ext::oneapi::bfloat16>) return dt::bf16;
34#endif
35        else static_assert(0);
36    }
37
38    static void gemm(ggml_backend_sycl_context & ctx, int m, int n, int k,
39        const void * a, dt at, dnnl_dim_t stra0, dnnl_dim_t stra1, dnnl_dim_t stra2,
40        const void * b, dt bt, dnnl_dim_t strb0, dnnl_dim_t strb1, dnnl_dim_t strb2,
41        void * c, dt ct, const queue_ptr & q, dnnl_dim_t batches_a, dnnl_dim_t batches_b) {
42
43        auto stream = ctx.stream_dnnl(q);
44        auto eng = ctx.engine_dnnl(q);
45
46        dnnl::memory::dims a_dims = {batches_a, m, k };
47        dnnl::memory::dims a_strides = {stra2, stra1, stra0};
48        const auto a_in_md = dnnl::memory::desc(a_dims, at, a_strides);
49
50        dnnl::memory::dims b_dims = {batches_b, k, n };
51        dnnl::memory::dims b_strides = {strb2, strb0, strb1};
52        const auto b_in_md = dnnl::memory::desc(b_dims, bt, b_strides);
53
54        dnnl::memory::dims c_dims = { std::max(batches_a, batches_b), m, n};
55        dnnl::memory::dims c_strides = {m*n, 1,  m };
56        const auto c_md    = dnnl::memory::desc(c_dims, ct, c_strides);
57        dnnl::primitive_attr primitive_attr;
58        primitive_attr.set_scratchpad_mode(dnnl::scratchpad_mode::user);
59
60#ifdef GGML_SYCL_F16
61        primitive_attr.set_fpmath_mode(dnnl::fpmath_mode::f16);
62#endif
63
64        auto a_mem = dnnl::memory(a_in_md, eng, const_cast<void*>(a));
65        auto b_mem = dnnl::memory(b_in_md, eng, const_cast<void*>(b));
66        auto matmul_pd = dnnl::matmul::primitive_desc(eng, a_in_md, b_in_md, c_md, primitive_attr);
67        auto c_mem = dnnl::memory(matmul_pd.dst_desc(), eng, c);
68
69        const auto scratchpad_md = matmul_pd.scratchpad_desc();
70        ggml_sycl_pool_alloc<uint8_t> scratchpad(ctx.pool());
71        void * scratchpad_ptr = scratchpad_md.get_size() > 0 ? scratchpad.alloc(scratchpad_md.get_size()) : nullptr;
72        auto scratchpad_mem = dnnl::memory(scratchpad_md, eng, scratchpad_ptr);
73
74        auto matmul_prim = dnnl::matmul(matmul_pd);
75
76        std::unordered_map<int, dnnl::memory> matmul_args;
77        matmul_args.insert({ DNNL_ARG_SRC, a_mem });
78        matmul_args.insert({ DNNL_ARG_WEIGHTS, b_mem });
79
80        matmul_args.insert({ DNNL_ARG_DST, c_mem });
81        matmul_args.insert({ DNNL_ARG_SCRATCHPAD, scratchpad_mem });
82
83        matmul_prim.execute(stream, matmul_args);
84    }
85
86    static void row_gemm(ggml_backend_sycl_context & ctx, int m, int n, int k,
87        const void * a, dt at, const void * b, dt bt, void * c, dt ct, const queue_ptr & q) {
88
89        gemm(ctx, m, n, k, a, at, 1, k, k * m, b, bt, 1, k, n * k, c, ct, q, 1, 1);
90    }
91};
92
93#endif
94
95#endif // GGML_SYCL_GEMM_HPP