Hugging Face's logo

HelloSun
/
Experiment_1

GGUF
conversational
Model card Files
xet
 Community
Experiment_1 / onednn /onednn.patch
HelloSun's picture
HelloSun
Integrate oneDNN/OpenVINO CPU backend study; merge docs; point build script at Experiment_1
da11688 verified
Raw History Blame Contribute Delete
24.4 kB
diff --git a/ggml/CMakeLists.txt b/ggml/CMakeLists.txt
index 35252a8..8a99468 100644
--- a/ggml/CMakeLists.txt
+++ b/ggml/CMakeLists.txt
@@ -246,6 +246,7 @@ set (GGML_METAL_STD "" CACHE STRING "ggml: metal standard version (-std
set (GGML_METAL_TARGET_OS "macos" CACHE STRING
"ggml: metal -mtargetos OS name (macos, ios, xros, tvos)")
option(GGML_OPENMP "ggml: use OpenMP" ON)
+option(GGML_CPU_ONEDNN "ggml: use oneDNN for quantized matmul in the CPU backend" OFF)
option(GGML_OPENMP_FETCH "ggml: fetch LLVM OpenMP" OFF)
option(GGML_RPC "ggml: use RPC" OFF)
option(GGML_SYCL "ggml: use SYCL" OFF)
diff --git a/ggml/src/ggml-cpu/CMakeLists.txt b/ggml/src/ggml-cpu/CMakeLists.txt
index d347757..3b0109f 100644
--- a/ggml/src/ggml-cpu/CMakeLists.txt
+++ b/ggml/src/ggml-cpu/CMakeLists.txt
@@ -56,11 +56,25 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
ggml-cpu/vec.cpp
ggml-cpu/ops.h
ggml-cpu/ops.cpp
+ ggml-cpu/onednn.h
+ ggml-cpu/onednn.cpp
)
target_compile_features(${GGML_CPU_NAME} PRIVATE c_std_11 cxx_std_17)
target_include_directories(${GGML_CPU_NAME} PRIVATE . ggml-cpu)
+ if (GGML_CPU_ONEDNN)
+ find_library(DNNL_LIBRARY dnnl REQUIRED)
+ find_path(DNNL_INCLUDE_DIR dnnl.hpp PATH_SUFFIXES oneapi)
+ if (NOT DNNL_INCLUDE_DIR)
+ message(FATAL_ERROR "oneDNN headers not found (dnnl.hpp)")
+ endif()
+ message(STATUS "oneDNN: ${DNNL_LIBRARY} (${DNNL_INCLUDE_DIR})")
+ target_compile_definitions(${GGML_CPU_NAME} PRIVATE GGML_USE_ONEDNN)
+ target_include_directories(${GGML_CPU_NAME} PRIVATE ${DNNL_INCLUDE_DIR})
+ target_link_libraries(${GGML_CPU_NAME} PRIVATE ${DNNL_LIBRARY})
+ endif()
+
if (APPLE AND GGML_ACCELERATE)
find_library(ACCELERATE_FRAMEWORK Accelerate)
if (ACCELERATE_FRAMEWORK)
diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c
index ab82e3d..e14b23c 100644
--- a/ggml/src/ggml-cpu/ggml-cpu.c
+++ b/ggml/src/ggml-cpu/ggml-cpu.c
@@ -15,6 +15,7 @@
#include "ggml.h"
#include "common.h"
#include "tiled/tiled.h"
+#include "onednn.h"
#if defined(_MSC_VER) || defined(__MINGW32__)
#include <malloc.h> // using malloc.h with MSC/MINGW
@@ -1260,11 +1261,30 @@ void ggml_compute_forward_mul_mat(
const struct ggml_tensor * src1 = dst->src[1];
const int32_t hint = ggml_get_op_params_i32(dst, 1);
+#ifdef GGML_USE_ONEDNN
+ if (getenv("GGML_ONEDNN_DEBUG")) {
+ fprintf(stderr, "[onednn] CALL src0=%s ne=[%lld,%lld,%lld,%lld] src1=%s ne=[%lld,%lld,%lld,%lld]\n",
+ ggml_type_name(src0->type), (long long)src0->ne[0], (long long)src0->ne[1], (long long)src0->ne[2], (long long)src0->ne[3],
+ ggml_type_name(src1->type), (long long)src1->ne[0], (long long)src1->ne[1], (long long)src1->ne[2], (long long)src1->ne[3]);
+ }
+#endif
if (hint == GGML_HINT_SRC0_IS_HADAMARD && !params->use_ref) {
ggml_compute_forward_fwht(params, dst);
return;
}
+#ifdef GGML_USE_ONEDNN
+ // Quantized weight x activation projections are executed by oneDNN. Only one
+ // thread enters oneDNN (which has its own OpenMP runtime); the others wait.
+ if (!params->use_ref && ggml_cpu_onednn_supported(dst)) {
+ if (params->ith == 0) {
+ ggml_cpu_onednn_mul_mat(dst);
+ }
+ ggml_barrier(params->threadpool);
+ return;
+ }
+#endif
+
// If tiled is supported, it will execute the full op here and we return
if (ggml_compute_forward_mul_mat_tiled(params, dst)) {
return;
diff --git a/ggml/src/ggml-cpu/onednn.cpp b/ggml/src/ggml-cpu/onednn.cpp
new file mode 100644
index 0000000..dc97b7f
--- /dev/null
+++ b/ggml/src/ggml-cpu/onednn.cpp
@@ -0,0 +1,508 @@
+#include "onednn.h"
+
+#ifdef GGML_USE_ONEDNN
+
+#include "ggml-impl.h"
+
+#include <dnnl.hpp>
+
+#include <immintrin.h>
+
+#include <algorithm>
+#include <type_traits>
+#include <chrono>
+#include <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+namespace {
+
+using dt = dnnl::memory::data_type;
+using tag = dnnl::memory::format_tag;
+
+// weight quantization used by the oneDNN path
+enum class wq {
+ bf16,
+ int8,
+ int4,
+};
+
+wq get_mode() {
+ static const wq mode = []() {
+ const char * s = std::getenv("GGML_ONEDNN_WEIGHT_TYPE");
+ if (s != nullptr && std::strcmp(s, "int8") == 0) return wq::int8;
+ if (s != nullptr && std::strcmp(s, "int4") == 0) return wq::int4;
+ return wq::bf16;
+ }();
+ return mode;
+}
+
+const char * mode_name(wq mode) {
+ switch (mode) {
+ case wq::bf16: return "bf16";
+ case wq::int8: return "int8";
+ case wq::int4: return "int4";
+ }
+ return "?";
+}
+
+inline bool dbg_on() {
+ static const bool on = std::getenv("GGML_ONEDNN_DEBUG") != nullptr;
+ return on;
+}
+
+inline uint16_t f32_to_bf16(float f) {
+ uint32_t u;
+ std::memcpy(&u, &f, sizeof(u));
+ if ((u & 0x7f800000u) == 0x7f800000u && (u & 0x007fffffu) != 0) {
+ return (uint16_t) ((u >> 16) | 0x0040u);
+ }
+ u += 0x7fffu + ((u >> 16) & 1u);
+ return (uint16_t) (u >> 16);
+}
+
+struct weight_cache {
+ int64_t k = 0;
+ int64_t n = 0;
+ ggml_type type = GGML_TYPE_COUNT;
+ wq mode = wq::bf16;
+
+ std::vector<uint16_t> bf16; // k*n
+ std::vector<int8_t> i8; // k*n, layout matches tag::ba
+ std::vector<int8_t> i4; // k*n/2 packed s4, layout matches tag::ba
+ std::vector<float> scale; // n
+};
+
+std::mutex g_mutex;
+std::unordered_map<const void *, std::unique_ptr<weight_cache>> g_weights;
+
+// convert one weight row (k elements) to float
+void row_to_float(const struct ggml_tensor * w, int64_t j, float * out) {
+ const struct ggml_type_traits * traits = ggml_get_type_traits(w->type);
+ GGML_ASSERT(traits->to_float != nullptr);
+ const char * src = (const char *) w->data + (size_t) j * w->nb[1];
+ traits->to_float(src, out, w->ne[0]);
+}
+
+const weight_cache * get_weight(const struct ggml_tensor * w) {
+ std::lock_guard<std::mutex> lock(g_mutex);
+
+ const int64_t k = w->ne[0];
+ const int64_t n = w->ne[1];
+ const wq mode = get_mode();
+
+ auto it = g_weights.find(w->data);
+ if (it != g_weights.end() &&
+ it->second->k == k && it->second->n == n &&
+ it->second->type == w->type && it->second->mode == mode) {
+ return it->second.get();
+ }
+
+ auto entry = std::make_unique<weight_cache>();
+ entry->k = k;
+ entry->n = n;
+ entry->type = w->type;
+ entry->mode = mode;
+
+ const auto t_start = std::chrono::steady_clock::now();
+
+ std::vector<float> row((size_t) k);
+
+ if (mode == wq::bf16) {
+ entry->bf16.resize((size_t) k * (size_t) n);
+ for (int64_t j = 0; j < n; ++j) {
+ row_to_float(w, j, row.data());
+ uint16_t * dst = entry->bf16.data() + (size_t) j * (size_t) k;
+ for (int64_t i = 0; i < k; ++i) dst[i] = f32_to_bf16(row[(size_t) i]);
+ }
+ } else if (mode == wq::int8) {
+ entry->i8.resize((size_t) k * (size_t) n);
+ entry->scale.resize((size_t) n);
+ for (int64_t j = 0; j < n; ++j) {
+ row_to_float(w, j, row.data());
+ float amax = 0.0f;
+ for (int64_t i = 0; i < k; ++i) amax = std::max(amax, std::fabs(row[(size_t) i]));
+ const float s = amax > 0.0f ? amax / 127.0f : 1e-8f;
+ entry->scale[(size_t) j] = s;
+ int8_t * dst = entry->i8.data() + (size_t) j * (size_t) k;
+ for (int64_t i = 0; i < k; ++i) {
+ int v = (int) std::lround(row[(size_t) i] / s);
+ dst[i] = (int8_t) std::max(-127, std::min(127, v));
+ }
+ }
+ } else {
+ GGML_ASSERT(k % 2 == 0);
+ entry->i4.resize((size_t) k * (size_t) n / 2);
+ entry->scale.resize((size_t) n);
+ for (int64_t j = 0; j < n; ++j) {
+ row_to_float(w, j, row.data());
+ float amax = 0.0f;
+ for (int64_t i = 0; i < k; ++i) amax = std::max(amax, std::fabs(row[(size_t) i]));
+ const float s = amax > 0.0f ? amax / 7.0f : 1e-8f;
+ entry->scale[(size_t) j] = s;
+ int8_t * dst = entry->i4.data() + (size_t) j * (size_t) k / 2;
+ for (int64_t i = 0; i < k; i += 2) {
+ int v0 = (int) std::lround(row[(size_t) i + 0] / s);
+ int v1 = (int) std::lround(row[(size_t) i + 1] / s);
+ v0 = std::max(-8, std::min(7, v0));
+ v1 = std::max(-8, std::min(7, v1));
+ dst[i / 2] = (int8_t) ((v0 & 0x0f) | ((v1 & 0x0f) << 4));
+ }
+ }
+ }
+
+ const weight_cache * result = entry.get();
+ g_weights[w->data] = std::move(entry);
+
+ if (dbg_on()) {
+ const double ms = std::chrono::duration<double, std::milli>(
+ std::chrono::steady_clock::now() - t_start).count();
+ std::fprintf(stderr, "[onednn] to_%s %s %lldx%lld in %.1f ms (cache %zu)\n",
+ mode_name(mode), ggml_type_name(w->type),
+ (long long) k, (long long) n, ms, g_weights.size());
+ }
+ return result;
+}
+
+struct primitive_key {
+ int64_t m;
+ int64_t k;
+ int64_t n;
+ wq mode;
+
+ bool operator==(const primitive_key & o) const {
+ return m == o.m && k == o.k && n == o.n && mode == o.mode;
+ }
+};
+
+struct primitive_hash {
+ size_t operator()(const primitive_key & key) const {
+ size_t h = std::hash<int64_t>{}(key.m);
+ h ^= std::hash<int64_t>{}(key.k) + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2);
+ h ^= std::hash<int64_t>{}(key.n) + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2);
+ h ^= std::hash<int>{}((int) key.mode) + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2);
+ return h;
+ }
+};
+
+dnnl::engine & engine() {
+ static dnnl::engine eng(dnnl::engine::kind::cpu, 0);
+ return eng;
+}
+
+dnnl::stream & stream() {
+ static dnnl::stream stm(engine());
+ return stm;
+}
+
+const dnnl::matmul & get_primitive(int64_t m, int64_t k, int64_t n, wq mode) {
+ static std::mutex mutex;
+ static std::unordered_map<primitive_key, dnnl::matmul, primitive_hash> cache;
+
+ std::lock_guard<std::mutex> lock(mutex);
+
+ const primitive_key key{m, k, n, mode};
+ auto it = cache.find(key);
+ if (it != cache.end()) {
+ return it->second;
+ }
+
+ dnnl::primitive_attr attr;
+ dnnl::memory::desc src_md;
+ dnnl::memory::desc wei_md;
+ const auto dst_md = dnnl::memory::desc({m, n}, dt::f32, tag::ab);
+
+ switch (mode) {
+ case wq::bf16:
+ src_md = dnnl::memory::desc({m, k}, dt::bf16, tag::ab);
+ wei_md = dnnl::memory::desc({k, n}, dt::bf16, tag::ba);
+ break;
+ case wq::int8:
+ src_md = dnnl::memory::desc({m, k}, dt::u8, tag::ab);
+ wei_md = dnnl::memory::desc({k, n}, dt::s8, tag::ba);
+ attr.set_scales(DNNL_ARG_SRC, 0, {});
+ attr.set_scales(DNNL_ARG_WEIGHTS, 1 << 1, {});
+ attr.set_zero_points(DNNL_ARG_SRC, 0, {});
+ break;
+ case wq::int4:
+ src_md = dnnl::memory::desc({m, k}, dt::u8, tag::ab);
+ wei_md = dnnl::memory::desc({k, n}, dt::s4, tag::ba);
+ attr.set_scales(DNNL_ARG_SRC, 0, {});
+ attr.set_scales(DNNL_ARG_WEIGHTS, 1 << 1, {});
+ attr.set_zero_points(DNNL_ARG_SRC, 0, {});
+ break;
+ }
+
+ auto pd = dnnl::matmul::primitive_desc(engine(), src_md, wei_md, dst_md, attr);
+
+ const auto t_start = std::chrono::steady_clock::now();
+ auto res = cache.emplace(key, dnnl::matmul(pd));
+ if (dbg_on()) {
+ const double ms = std::chrono::duration<double, std::milli>(
+ std::chrono::steady_clock::now() - t_start).count();
+ std::fprintf(stderr, "[onednn] primitive %s M=%lld K=%lld N=%lld impl=%s in %.1f ms (cache %zu)\n",
+ mode_name(mode), (long long) m, (long long) k, (long long) n,
+ pd.impl_info_str(), ms, cache.size());
+ }
+ return res.first->second;
+}
+
+template <typename T>
+inline float cvt_f32(T v);
+
+template <> inline float cvt_f32<float> (float v) { return v; }
+template <> inline float cvt_f32<ggml_fp16_t> (ggml_fp16_t v) { return GGML_COMPUTE_FP16_TO_FP32(v); }
+template <> inline float cvt_f32<ggml_bf16_t> (ggml_bf16_t v) { return GGML_BF16_TO_FP32(v); }
+
+#if defined(__AVX512F__) && defined(__AVX512BW__)
+inline void minmax_f32_avx512(const float * row, int64_t k, __m512 & vmin, __m512 & vmax) {
+ vmin = _mm512_set1_ps(0.0f);
+ vmax = _mm512_set1_ps(0.0f);
+ int64_t j = 0;
+ for (; j + 16 <= k; j += 16) {
+ const __m512 v = _mm512_loadu_ps(row + j);
+ vmin = _mm512_min_ps(vmin, v);
+ vmax = _mm512_max_ps(vmax, v);
+ }
+ if (j < k) {
+ const __mmask16 mask = (__mmask16) ((1u << (k - j)) - 1u);
+ const __m512 v = _mm512_maskz_loadu_ps(mask, row + j);
+ vmin = _mm512_min_ps(vmin, v);
+ vmax = _mm512_max_ps(vmax, v);
+ }
+}
+
+inline void quant_f32_avx512(const float * row, int64_t k, uint8_t * dst, float inv, int32_t zp) {
+ const __m512 vinv = _mm512_set1_ps(inv);
+ const __m512i vzp = _mm512_set1_epi32(zp);
+ int64_t j = 0;
+ for (; j + 16 <= k; j += 16) {
+ const __m512 v = _mm512_mul_ps(_mm512_loadu_ps(row + j), vinv);
+ const __m512i q = _mm512_add_epi32(_mm512_cvtps_epi32(v), vzp);
+ _mm_storeu_si128((__m128i *) (dst + j), _mm512_cvtusepi32_epi8(q));
+ }
+ for (; j < k; ++j) {
+ const int q = (int) std::lround(row[j] * inv) + zp;
+ dst[j] = (uint8_t) std::max(0, std::min(255, q));
+ }
+}
+#endif
+
+template <typename T>
+void quantize_u8_typed(const char * base, size_t row_bytes, int64_t m, int64_t k,
+ std::vector<uint8_t> & out, float & scale, int32_t & zp) {
+ const T * p0 = (const T *) base;
+ const int64_t stride = (int64_t) (row_bytes / sizeof(T));
+
+ float vmin = 0.0f;
+ float vmax = 0.0f;
+#if defined(__AVX512F__) && defined(__AVX512BW__)
+ if (std::is_same<T, float>::value) {
+ __m512 vminv, vmaxv;
+ minmax_f32_avx512((const float *) p0, k, vminv, vmaxv);
+ for (int64_t i = 1; i < m; ++i) {
+ __m512 a, b;
+ minmax_f32_avx512((const float *) p0 + (size_t) i * (size_t) stride, k, a, b);
+ vminv = _mm512_min_ps(vminv, a);
+ vmaxv = _mm512_max_ps(vmaxv, b);
+ }
+ vmin = _mm512_reduce_min_ps(vminv);
+ vmax = _mm512_reduce_max_ps(vmaxv);
+ } else
+#endif
+ {
+ for (int64_t i = 0; i < m; ++i) {
+ const T * row = p0 + (size_t) i * (size_t) stride;
+ for (int64_t j = 0; j < k; ++j) {
+ const float v = cvt_f32<T>(row[j]);
+ vmin = std::min(vmin, v);
+ vmax = std::max(vmax, v);
+ }
+ }
+ }
+
+ scale = vmax > vmin ? (vmax - vmin) / 255.0f : 1e-8f;
+ const float inv = 1.0f / scale;
+ zp = (int32_t) std::lround(-vmin * inv);
+ zp = std::max(0, std::min(255, zp));
+
+ for (int64_t i = 0; i < m; ++i) {
+ const T * row = p0 + (size_t) i * (size_t) stride;
+ uint8_t * dst = out.data() + (size_t) i * (size_t) k;
+#if defined(__AVX512F__) && defined(__AVX512BW__)
+ if (std::is_same<T, float>::value) {
+ quant_f32_avx512((const float *) row, k, dst, inv, zp);
+ continue;
+ }
+#endif
+ for (int64_t j = 0; j < k; ++j) {
+ const int q = (int) std::lround(cvt_f32<T>(row[j]) * inv) + zp;
+ dst[j] = (uint8_t) std::max(0, std::min(255, q));
+ }
+ }
+}
+
+// quantize activations to u8, per-tensor asymmetric
+void quantize_u8(const struct ggml_tensor * src1, int64_t m, int64_t k,
+ std::vector<uint8_t> & out, float & scale, int32_t & zp) {
+ out.resize((size_t) m * (size_t) k);
+
+ const char * base = (const char *) src1->data;
+ const size_t row_bytes = (size_t) src1->nb[1];
+
+ switch (src1->type) {
+ case GGML_TYPE_F32: quantize_u8_typed<float> (base, row_bytes, m, k, out, scale, zp); break;
+ case GGML_TYPE_F16: quantize_u8_typed<ggml_fp16_t> (base, row_bytes, m, k, out, scale, zp); break;
+ case GGML_TYPE_BF16: quantize_u8_typed<ggml_bf16_t> (base, row_bytes, m, k, out, scale, zp); break;
+ default: GGML_ABORT("unsupported activation type for quantization");
+ }
+}
+
+} // namespace
+
+static void dbg_unsupported(const char * why, const struct ggml_tensor * dst) {
+ if (!dbg_on()) return;
+ static int n = 0;
+ if (n++ >= 400) return;
+ const struct ggml_tensor * a = dst->src[0];
+ const struct ggml_tensor * b = dst->src[1];
+ std::fprintf(stderr, "[onednn] skip %s: dst_type=%s src0=%s ne=[%lld,%lld,%lld,%lld] src1=%s ne=[%lld,%lld,%lld,%lld] cont0=%d cont1=%d\n",
+ why,
+ ggml_type_name(dst->type),
+ a ? ggml_type_name(a->type) : "-", a?(long long)a->ne[0]:0, a?(long long)a->ne[1]:0, a?(long long)a->ne[2]:0, a?(long long)a->ne[3]:0,
+ b ? ggml_type_name(b->type) : "-", b?(long long)b->ne[0]:0, b?(long long)b->ne[1]:0, b?(long long)b->ne[2]:0, b?(long long)b->ne[3]:0,
+ a?ggml_is_contiguous(a):0, b?ggml_is_contiguous(b):0);
+}
+
+bool ggml_cpu_onednn_supported(const struct ggml_tensor * dst) {
+ const struct ggml_tensor * src0 = dst->src[0];
+ const struct ggml_tensor * src1 = dst->src[1];
+
+ if (src0 == nullptr || src1 == nullptr) { dbg_unsupported("null", dst); return false; }
+ if (dst->type != GGML_TYPE_F32) { dbg_unsupported("dst_f32", dst); return false; }
+ if (!ggml_is_quantized(src0->type)) { dbg_unsupported("not_quant", dst); return false; }
+ if (src1->type != GGML_TYPE_F32 && src1->type != GGML_TYPE_F16 && src1->type != GGML_TYPE_BF16) { dbg_unsupported("src1_type", dst); return false; }
+
+ // only plain 2D weight x 2D activation (projections and lm_head)
+ if (src0->ne[2] != 1 || src0->ne[3] != 1) { dbg_unsupported("src0_ndim", dst); return false; }
+ if (src1->ne[2] != 1 || src1->ne[3] != 1) { dbg_unsupported("src1_ndim", dst); return false; }
+ if (dst->ne[2] != 1 || dst->ne[3] != 1) { dbg_unsupported("dst_ndim", dst); return false; }
+
+ if (!ggml_is_contiguous(src0)) { dbg_unsupported("cont0", dst); return false; }
+ if (!ggml_is_contiguous(src1)) { dbg_unsupported("cont1", dst); return false; }
+
+ if (dst->nb[0] != sizeof(float)) { dbg_unsupported("dst_nb0", dst); return false; }
+ if (dst->nb[1] != (int64_t) dst->ne[0] * (int64_t) sizeof(float)) { dbg_unsupported("dst_nb1", dst); return false; }
+
+ if (src0->ne[0] <= 0 || src0->ne[1] <= 0) { dbg_unsupported("dim", dst); return false; }
+ if (src0->ne[0] % ggml_blck_size(src0->type) != 0) { dbg_unsupported("block", dst); return false; }
+
+ if (get_mode() == wq::int4 && (src0->ne[0] % 2) != 0) { dbg_unsupported("int4_k_odd", dst); return false; }
+
+ return true;
+}
+
+void ggml_cpu_onednn_mul_mat(const struct ggml_tensor * dst) {
+ const struct ggml_tensor * src0 = dst->src[0];
+ const struct ggml_tensor * src1 = dst->src[1];
+
+ const int64_t k = src0->ne[0];
+ const int64_t n = src0->ne[1];
+ const int64_t m = src1->ne[1];
+ const wq mode = get_mode();
+
+ const weight_cache * w = get_weight(src0);
+ const dnnl::matmul & prim = get_primitive(m, k, n, mode);
+
+ const auto dst_md = dnnl::memory::desc({m, n}, dt::f32, tag::ab);
+ dnnl::memory dst_mem(dst_md, engine(), dst->data);
+
+ if (mode == wq::bf16) {
+ std::vector<uint16_t> act((size_t) m * (size_t) k);
+ const char * src1_data = (const char *) src1->data;
+ const size_t src1_row = (size_t) k * ggml_type_size(src1->type);
+ for (int64_t i = 0; i < m; ++i) {
+ const char * row = src1_data + (size_t) i * src1_row;
+ uint16_t * out = act.data() + (size_t) i * (size_t) k;
+ if (src1->type == GGML_TYPE_F32) {
+ const float * f = (const float *) row;
+ for (int64_t j = 0; j < k; ++j) out[j] = f32_to_bf16(f[j]);
+ } else if (src1->type == GGML_TYPE_F16) {
+ const ggml_fp16_t * h = (const ggml_fp16_t *) row;
+ for (int64_t j = 0; j < k; ++j) out[j] = f32_to_bf16(GGML_COMPUTE_FP16_TO_FP32(h[j]));
+ } else {
+ std::memcpy(out, row, (size_t) k * sizeof(uint16_t));
+ }
+ }
+ auto src_md = dnnl::memory::desc({m, k}, dt::bf16, tag::ab);
+ auto wei_md = dnnl::memory::desc({k, n}, dt::bf16, tag::ba);
+ dnnl::memory src_mem(src_md, engine(), act.data());
+ dnnl::memory wei_mem(wei_md, engine(), const_cast<uint16_t *>(w->bf16.data()));
+ prim.execute(stream(), {
+ {DNNL_ARG_SRC, src_mem}, {DNNL_ARG_WEIGHTS, wei_mem}, {DNNL_ARG_DST, dst_mem},
+ });
+ stream().wait();
+ return;
+ }
+
+ const auto t_q0 = std::chrono::steady_clock::now();
+ std::vector<uint8_t> act;
+ float scale_a = 1.0f;
+ int32_t zp_a = 0;
+ quantize_u8(src1, m, k, act, scale_a, zp_a);
+ const auto t_q1 = std::chrono::steady_clock::now();
+
+ auto src_md = dnnl::memory::desc({m, k}, dt::u8, tag::ab);
+ dnnl::memory src_mem(src_md, engine(), act.data());
+
+ dnnl::memory wei_mem;
+ if (mode == wq::int8) {
+ auto wei_md = dnnl::memory::desc({k, n}, dt::s8, tag::ba);
+ wei_mem = dnnl::memory(wei_md, engine(), const_cast<int8_t *>(w->i8.data()));
+ } else {
+ auto wei_md = dnnl::memory::desc({k, n}, dt::s4, tag::ba);
+ wei_mem = dnnl::memory(wei_md, engine(), const_cast<int8_t *>(w->i4.data()));
+ }
+
+ dnnl::memory scale_a_mem(dnnl::memory::desc({1}, dt::f32, tag::x), engine(), &scale_a);
+ dnnl::memory scale_w_mem(dnnl::memory::desc({n}, dt::f32, tag::x), engine(), const_cast<float *>(w->scale.data()));
+ dnnl::memory zp_a_mem(dnnl::memory::desc({1}, dt::s32, tag::x), engine(), &zp_a);
+
+ prim.execute(stream(), {
+ {DNNL_ARG_SRC, src_mem},
+ {DNNL_ARG_WEIGHTS, wei_mem},
+ {DNNL_ARG_DST, dst_mem},
+ {DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC, scale_a_mem},
+ {DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS, scale_w_mem},
+ {DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_SRC, zp_a_mem},
+ });
+ stream().wait();
+
+ if (dbg_on()) {
+ const auto t_e = std::chrono::steady_clock::now();
+ const double qms = std::chrono::duration<double, std::milli>(t_q1 - t_q0).count();
+ const double ems = std::chrono::duration<double, std::milli>(t_e - t_q1).count();
+ std::fprintf(stderr, "[onednn] exec %s M=%lld K=%lld N=%lld quant=%.2f exec=%.2f\n",
+ mode_name(mode), (long long) m, (long long) k, (long long) n, qms, ems);
+ }
+}
+
+#else
+
+bool ggml_cpu_onednn_supported(const struct ggml_tensor * dst) {
+ (void) dst;
+ return false;
+}
+
+void ggml_cpu_onednn_mul_mat(const struct ggml_tensor * dst) {
+ (void) dst;
+}
+
+#endif
diff --git a/ggml/src/ggml-cpu/onednn.h b/ggml/src/ggml-cpu/onednn.h
new file mode 100644
index 0000000..0f36d5d
--- /dev/null
+++ b/ggml/src/ggml-cpu/onednn.h
@@ -0,0 +1,18 @@
+#pragma once
+
+#include "ggml.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Returns true when the given MUL_MAT node can be executed by the oneDNN path.
+// Must be cheap and deterministic (called by every worker thread).
+bool ggml_cpu_onednn_supported(const struct ggml_tensor * dst);
+
+// Executes the MUL_MAT node with oneDNN. Must be called by a single thread.
+void ggml_cpu_onednn_mul_mat(const struct ggml_tensor * dst);
+
+#ifdef __cplusplus
+}
+#endif