#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