// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

#pragma once
#include "cpu/x64/jit_generator.hpp"
#include "emitters/plugin/x64/jit_load_store_emitters.hpp"
#include <common/nstl.hpp>
#include <type_traits>
#include <functional>
#include <vector>
#include <array>
#include <tuple>

namespace ov {
namespace intel_cpu {

struct jit_kernel;

namespace internal {

template<size_t S>
struct reg_traits_by_size;
template<typename T>
struct reg_traits;
template<typename T, size_t N>
struct reg_traits<T[N]>;
template<dnnl::impl::cpu::x64::cpu_isa_t isa>
struct isa_traits;

template<>
struct reg_traits_by_size<1> {
    using type = Xbyak::Reg8;
    constexpr static size_t size = 1;           // in bytes
    constexpr static dnnl::impl::cpu::x64::cpu_isa_t isa
                        = dnnl::impl::cpu::x64::cpu_isa_t::isa_undef;
};

template<>
struct reg_traits_by_size<2> {
    using type = Xbyak::Reg16;
    constexpr static size_t size = 2;           // in bytes
    constexpr static dnnl::impl::cpu::x64::cpu_isa_t isa
                        = dnnl::impl::cpu::x64::cpu_isa_t::isa_undef;
};

template<>
struct reg_traits_by_size<4> {
    using type = Xbyak::Reg32;
    constexpr static size_t size = 4;           // in bytes
    constexpr static dnnl::impl::cpu::x64::cpu_isa_t isa
                        = dnnl::impl::cpu::x64::cpu_isa_t::isa_undef;
};

template<>
struct reg_traits_by_size<8> {
    using type = Xbyak::Reg64;
    constexpr static size_t size = 8;           // in bytes
    constexpr static dnnl::impl::cpu::x64::cpu_isa_t isa
                        = dnnl::impl::cpu::x64::cpu_isa_t::isa_undef;
};

template<>
struct reg_traits_by_size<16> {
    using type = Xbyak::Xmm;
    constexpr static size_t size = 16;          // in bytes
    constexpr static dnnl::impl::cpu::x64::cpu_isa_t isa
                        = dnnl::impl::cpu::x64::cpu_isa_t::sse41;
};

template<>
struct reg_traits_by_size<32> {
    using type = Xbyak::Ymm;
    constexpr static size_t size = 32;          // in bytes
    constexpr static dnnl::impl::cpu::x64::cpu_isa_t isa
                        = dnnl::impl::cpu::x64::cpu_isa_t::avx2;
};

template<>
struct reg_traits_by_size<64> {
    using type = Xbyak::Zmm;
    constexpr static size_t size = 64;          // in bytes
    constexpr static dnnl::impl::cpu::x64::cpu_isa_t isa
                        = dnnl::impl::cpu::x64::cpu_isa_t::avx512_core;
};

template<typename T>
struct reg_traits : public reg_traits_by_size<sizeof(T)> {};

template<size_t N>
struct vec_min_size {
    constexpr static size_t size = N <= 16 ? 16 :
                                   N <= 32 ? 32 :
                                   64;
};

template<typename T, size_t N>
struct reg_traits<T[N]> : public reg_traits_by_size<vec_min_size<sizeof(T[N])>::size> {};

template<>
struct reg_traits<float> {
    using type = Xbyak::Fpu;
    constexpr static size_t size = 10;          // in bytes
    constexpr static dnnl::impl::cpu::x64::cpu_isa_t isa
                        = dnnl::impl::cpu::x64::cpu_isa_t::isa_undef;
};
template<>
struct reg_traits<double> : public reg_traits<float> {};

template<>
struct isa_traits<dnnl::impl::cpu::x64::cpu_isa_t::sse41> {
    struct reg {
        using type = Xbyak::Xmm;
        constexpr static size_t size = 4 * 4;   // in bytes
        constexpr static size_t length = 4;     // in dwords
    };
};

template<>
struct isa_traits<dnnl::impl::cpu::x64::cpu_isa_t::avx2> {
    struct reg {
        using type = Xbyak::Ymm;
        constexpr static size_t size = 8 * 4;   // in bytes
        constexpr static size_t length = 8;     // in dwords
    };
};

template<>
struct isa_traits<dnnl::impl::cpu::x64::cpu_isa_t::avx512_core> {
    struct reg {
        using type = Xbyak::Zmm;
        constexpr static size_t size = 16 * 4;  // in bytes
        constexpr static size_t length = 16;    // in dwords
    };
};

template<typename T, typename Tag>
class variable;
template<typename T>
class if_expression;
template<typename T>
class then_expression;
template<typename Reg>
using shared_reg = std::shared_ptr<Reg>;

template<typename Reg>
shared_reg<Reg> make_shared(Reg & reg, jit_kernel & kernel);

template<typename T>
class boolean_expression {
public:
    using reg_type = const typename reg_traits<T>::type;

    enum class type {
        eq,     // ==
        neq,    // !=
        ls,     // <
        gt,     // >
        le,     // <=
        ge      // >=
    };

    boolean_expression(jit_kernel & kernel, type t, const shared_reg<reg_type> & lhs, const shared_reg<reg_type> & rhs);
    boolean_expression(jit_kernel & kernel, type t, const shared_reg<reg_type> & lhs, T rhs);

private:
    void cmp(const Xbyak::Label & exit) const;

    jit_kernel & _kernel;
    type _type;
    shared_reg<reg_type> _lhs;
    shared_reg<reg_type> _rhs;
    T _rvalue;

    friend class if_expression<T>;
    friend class then_expression<T>;
};

template<typename T>
class then_expression {
public:
    then_expression(if_expression<T> & expr);

    template<typename F>
    void _else(F && fn);

private:
    if_expression<T> & _if_expr;
};

template<typename T>
class if_expression {
public:
    if_expression(const boolean_expression<T> & expr)
        : _expr(expr) {}

    ~if_expression() {
        try {
            if (!_is_exit_valid)
                _expr._kernel.assignL(_exit, _else);
        } catch(...) {}
    }

    template<typename F>
    then_expression<T> _then(F && fn) {
        using namespace Xbyak;

        _expr.cmp(_else);
        fn();
        _expr._kernel.jmp(_exit, Xbyak::CodeGenerator::T_NEAR);
        _expr._kernel.L(_else);

        return then_expression<T>(*this);
    }

private:
    const boolean_expression<T> & _expr;
    Xbyak::Label _exit;
    Xbyak::Label _else;
    bool _is_exit_valid = false;

    friend class then_expression<T>;
};

typedef struct register_tag {} register_tag;
typedef struct memory_tag {} memory_tag;

template<typename T, typename Tag>
class variable_base;

template<typename T>
class variable_base<T, register_tag> {
public:
    using reg_type = const typename reg_traits<T>::type;

    variable_base & operator = (const variable_base &) = delete;

    variable_base(const variable_base &);
    variable_base(variable_base &&);

    reg_type & reg() const {
        return  *_reg;
    }

    const shared_reg<reg_type> & shreg() const {
        return _reg;
    }

    operator reg_type &() const {
        return reg();
    }

    operator Xbyak::RegExp () const {
        return reg();
    }

protected:
    variable_base(jit_kernel & krnl, const shared_reg<reg_type> & reg);
    ~variable_base() = default;

    jit_kernel & _kernel;
    shared_reg<reg_type> _reg;
};

template<typename T>
class variable_base<T, memory_tag> {
public:
    using reg_type = const typename reg_traits<T*>::type;

    variable_base & operator = (const variable_base &) = delete;

    variable_base(const variable_base &);
    variable_base(variable_base &&);

    reg_type & reg() const {
        return  *_addr;
    }

protected:
    variable_base(jit_kernel & krnl, const shared_reg<reg_type> & addr);
    ~variable_base() = default;

    jit_kernel & _kernel;
    shared_reg<const reg_type> _addr;
};

template<typename T>
class variable<T, register_tag> : public variable_base<typename std::enable_if<!std::is_floating_point<T>::value, T>::type, register_tag> {
public:
    using type = T;
    using base = variable_base<type, register_tag>;
    using reg_type = const typename base::reg_type;
    using arithmetic_type = typename std::conditional<std::is_pointer<T>::value, size_t, T>::type;

    variable(variable &&) = default;
    variable(jit_kernel & krnl);
    variable(jit_kernel & krnl, const shared_reg<reg_type> & reg);

    typename std::conditional<std::is_pointer<T>::value
                    && !std::is_pointer<typename std::remove_pointer<T>::type>::value,
                    variable<typename std::remove_pointer<T>::type, memory_tag>, void>::type
    operator *() const {
        return variable<typename std::remove_pointer<T>::type, memory_tag>(base::_kernel, base::shreg());
    }

    const variable & operator = (reg_type & rhs) const {
        base::_kernel.mov(base::reg(), rhs);
        return *this;
    }
    template<typename U>
    const variable & operator = (U *rhs) const {
        // interpret pointers as size_t
        base::_kernel.mov(base::reg(), reinterpret_cast<size_t>(rhs));
        return *this;
    }
    const variable & operator = (arithmetic_type rhs) const {
        base::_kernel.mov(base::reg(), static_cast<size_t>(rhs));
        return *this;
    }
    const variable & operator += (reg_type & rhs) const {
        base::_kernel.add(base::reg(), rhs);
        return *this;
    }
    variable operator + (reg_type & rhs) const {
        variable res(base::_kernel);
        res = base::reg();
        res += rhs;
        return res;
    }
    const variable & operator += (arithmetic_type rhs) const {
        base::_kernel.add(base::reg(), rhs);
        return *this;
    }
    variable operator + (arithmetic_type rhs) const {
        variable res(base::_kernel);
        res = base::reg();
        res += rhs;
        return res;
    }
    const variable & operator -= (reg_type & rhs) const {
        base::_kernel.sub(base::reg(), rhs);
        return *this;
    }
    variable operator - (reg_type & rhs) const {
        variable res(base::_kernel);
        res = base::reg();
        res -= rhs;
        return res;
    }
    const variable & operator -= (arithmetic_type rhs) const {
        base::_kernel.sub(base::reg(), rhs);
        return *this;
    }
    variable operator - (arithmetic_type rhs) const {
        variable res(base::_kernel);
        res = base::reg();
        res -= rhs;
        return res;
    }
    const variable & operator *= (reg_type & rhs) const {
        base::_kernel.imul(base::reg(), rhs);
        return *this;
    }
    variable operator * (reg_type & rhs) const {
        variable res(base::_kernel);
        res = base::reg();
        res *= rhs;
        return res;
    }
    const variable & operator *= (arithmetic_type rhs) const {
        base::_kernel.imul(base::reg(), base::reg(), static_cast<int>(rhs));
        return *this;
    }
    variable operator * (arithmetic_type rhs) const {
        variable res(base::_kernel);
        res = base::reg();
        res *= rhs;
        return res;
    }
    const variable & operator &= (reg_type & rhs) const {
        base::_kernel.and_(base::reg(), rhs);
        return *this;
    }
    variable operator & (reg_type & rhs) const {
        variable res(base::_kernel);
        res = base::reg();
        res &= rhs;
        return res;
    }
    const variable & operator &= (T rhs) const {
        base::_kernel.and_(base::reg(), rhs);
        return *this;
    }
    variable operator & (T rhs) const {
        variable res(base::_kernel);
        res = base::reg();
        res &= rhs;
        return res;
    }
    const variable & operator |= (reg_type & rhs) const {
        base::_kernel.or_(base::reg(), rhs);
        return *this;
    }
    variable operator | (reg_type & rhs) const {
        variable res(base::_kernel);
        res = base::reg();
        res |= rhs;
        return res;
    }
    const variable & operator |= (T rhs) const {
        base::_kernel.or_(base::reg(), rhs);
        return *this;
    }
    variable operator | (T rhs) const {
        variable res(base::_kernel);
        res = base::reg();
        res |= rhs;
        return res;
    }
    const variable & operator >>= (size_t rhs) const {
        base::_kernel.shr(base::reg(), rhs);
        return *this;
    }
    variable operator >> (size_t rhs) const {
        variable res(base::_kernel);
        res = base::reg();
        res >>= rhs;
        return res;
    }
    const variable & operator <<= (size_t rhs) const {
        base::_kernel.shl(base::reg(), rhs);
        return *this;
    }
    variable operator << (size_t rhs) const {
        variable res(base::_kernel);
        res = base::reg();
        res <<= rhs;
        return res;
    }

    boolean_expression<T> operator == (const variable & rhs) const {
        return boolean_expression<T>(base::_kernel, boolean_expression<T>::type::eq, base::shreg(), rhs.shreg());
    }

    boolean_expression<T> operator == (T rhs) const {
        return boolean_expression<T>(base::_kernel, boolean_expression<T>::type::eq, base::shreg(), rhs);
    }

    boolean_expression<T> operator != (const variable & rhs) const {
        return boolean_expression<T>(base::_kernel, boolean_expression<T>::type::neq, base::shreg(), rhs.shreg());
    }

    boolean_expression<T> operator != (T rhs) const {
        return boolean_expression<T>(base::_kernel, boolean_expression<T>::type::neq, base::shreg(), rhs);
    }

    boolean_expression<T> operator < (const variable & rhs) const {
        return boolean_expression<T>(base::_kernel, boolean_expression<T>::type::ls, base::shreg(), rhs.shreg());
    }

    boolean_expression<T> operator < (T rhs) const {
        return boolean_expression<T>(base::_kernel, boolean_expression<T>::type::ls, base::shreg(), rhs);
    }

    boolean_expression<T> operator > (const variable & rhs) const {
        return boolean_expression<T>(base::_kernel, boolean_expression<T>::type::gt, base::shreg(), rhs.shreg());
    }

    boolean_expression<T> operator > (T rhs) const {
        return boolean_expression<T>(base::_kernel, boolean_expression<T>::type::gt, base::shreg(), rhs);
    }

    boolean_expression<T> operator <= (const variable & rhs) const {
        return boolean_expression<T>(base::_kernel, boolean_expression<T>::type::le, base::shreg(), rhs.shreg());
    }

    boolean_expression<T> operator <= (T rhs) const {
        return boolean_expression<T>(base::_kernel, boolean_expression<T>::type::le, base::shreg(), rhs);
    }

    boolean_expression<T> operator >= (const variable & rhs) const {
        return boolean_expression<T>(base::_kernel, boolean_expression<T>::type::ge, base::shreg(), rhs.shreg());
    }

    boolean_expression<T> operator >= (T rhs) const {
        return boolean_expression<T>(base::_kernel, boolean_expression<T>::type::ge, base::shreg(), rhs);
    }

    // TODO: add necessary operations
};

template<typename T>
class variable<T, memory_tag> : public variable_base<T*, memory_tag> {
public:
    using type = T;
    using base = variable_base<type*, memory_tag>;
    using reg_type = const typename base::reg_type;

    variable(variable &&) = default;
    variable(jit_kernel & krnl, const shared_reg<reg_type> & reg);

    const variable & operator = (const variable<T, register_tag> & rhs) const;
};

template<typename T, size_t N>
class variable<T[N], register_tag> : public variable_base<T[N], register_tag> {
public:
    using type = T[N];
    using base = variable_base<type, register_tag>;
    using reg_type = const typename base::reg_type;
    constexpr static size_t length = N;

    variable(variable &&) = default;
    variable(jit_kernel & krnl);
    variable(jit_kernel & krnl, const shared_reg<reg_type> & reg);

    const variable & operator = (reg_type & rhs) const {
        base::_kernel.uni_vmovups(base::reg(), rhs);
        return *this;
    }

    const variable & operator = (const type & rhs) const {
        const type & cref = base::_kernel.constant(rhs);
        variable<const type*, register_tag> creg(base::_kernel);
        creg = &cref;
        base::_kernel.uni_vmovdqu(base::reg(), base::_kernel.ptr[creg]);
        return *this;
    }

    const variable & blend(reg_type & rhs, uint16_t mask) const {
        base::_kernel.uni_vblendps(base::reg(), rhs, mask);
        return *this;
    }

    const variable & permute(const std::array<uint8_t, N> & order) const {
        base::_kernel.uni_vpermps(base::reg(), order.data(), base::reg());
        return *this;
    }

    const variable & permute(const uint8_t * order) const {
        base::_kernel.uni_vpermps(base::reg(), order, base::reg());
        return *this;
    }

    // TODO: implement vector arithmetic
};

class stack_frame {
    stack_frame(const stack_frame &) = delete;
    stack_frame & operator = (const stack_frame &) = delete;

public:
    stack_frame(jit_kernel & kernel, size_t size, uint32_t alignment = 1);
    stack_frame(stack_frame && rhs);
    ~stack_frame();
    const Xbyak::Reg64 & pointer() const;
    void clear() const;

private:
    jit_kernel & _kernel;
    size_t _size;
    uint32_t _alignment;
};

template<typename T>
ov::element::Type type2precision();

dnnl::impl::cpu::x64::cpu_isa_t get_current_isa();

class consts_table {
    consts_table(const consts_table &) = delete;
    consts_table & operator = (const consts_table &) = delete;

public:
    consts_table() = default;
    const void * store(const void *data, size_t size);

private:
    static constexpr const size_t chunk_size = 512;
    using chunk = std::array<uint8_t, chunk_size>;
    std::list<chunk> _chunks;
    size_t _size {};
};

}   // namespace internal

struct jit_kernel : public dnnl::impl::cpu::x64::jit_generator {
    using reg_indices = std::vector<int>;
    template<typename T>
    using reg_traits = internal::reg_traits<T>;
    template<size_t S>
    using reg_traits_by_size = internal::reg_traits_by_size<S>;
    template<dnnl::impl::cpu::x64::cpu_isa_t isa>
    using isa_traits = internal::isa_traits<isa>;
    using stack_frame = internal::stack_frame;
    using register_tag = internal::register_tag;
    using memory_tag = internal::memory_tag;
    template<typename T, typename Tag = register_tag>
    using variable = internal::variable<T, Tag>;
    template<typename T>
    using if_expression = internal::if_expression<T>;
    template<typename T>
    using boolean_expression = internal::boolean_expression<T>;

    template<typename T, typename U>
    Xbyak::Address argPtr(U T::*member) const {
        auto memPtr = &(reinterpret_cast<const T*>(0)->*member);
        const size_t offs =  reinterpret_cast<const char*>(memPtr) - reinterpret_cast<const char*>(0);
        return address_frame(sizeof(U))[param1 + offs];
    }

    template<typename T, typename U>
    variable<U> arg(U T::*member) {
        using traits = internal::reg_traits<U>;
        using reg_type = typename traits::type;
        const auto & res = reserve<reg_type>();
        if (sizeof(T) < traits::size)
            movzx(res, argPtr(member));
        else
            mov(res, argPtr(member));
        return { *this, internal::make_shared(res, *this) };
    }

    template<typename CastU, typename T, typename U>
    variable<CastU> arg(U T::*member) {
        using traits = internal::reg_traits<U>;
        using reg_type = typename traits::type;
        const auto & res = reserve<reg_type>();
        if (sizeof(T) < traits::size)
            movzx(res, argPtr(member));
        else
            mov(res, argPtr(member));
        return { *this, internal::make_shared(res, *this) };
    }

    jit_kernel(const char *name);

    template<typename RegType>
    const RegType & reserve();

    template<typename RegType>
    void free(const RegType & reg);

    template<typename T>
    void copy(const Xbyak::Reg64& dst,
              const Xbyak::Reg64& src,
              const Xbyak::Reg64& size);
    template<typename T>
    void copy(const Xbyak::Address& dst,
              const Xbyak::Reg64& src,
              const Xbyak::Reg64& size);

    template<typename DstT, size_t N, typename SrcT>
    void load(const variable<DstT[N]> & dst, const variable<SrcT> & src, size_t length = N);
    template<typename DstT, size_t N, typename SrcT>
    void load(const variable<DstT[N]> & dst, const variable<SrcT> & src, const variable<size_t> & length);
    template<typename DstT, typename SrcT, size_t N>
    void store(const variable<DstT> & dst, const variable<SrcT[N]> & src, size_t length = N);
    template<typename DstT, typename SrcT, size_t N>
    void store(const variable<DstT> & dst, const variable<SrcT[N]> & src, const variable<size_t> & length);

    template<typename B, typename E, typename S = size_t>
    void foreach(const B & begin,
                 const E & end,
                 std::function<void(const variable<size_t>&)> && fn,
                 const S & step = 1);

    template<typename T>
    variable<T> var();
    template<typename T>
    variable<T> var(const T & val);

    template<typename T>
    const T & constant(const T & c);
    template<typename T>
    const T * constant(const T * c, size_t size);

    stack_frame stack(size_t size, uint32_t alignment = 1);

    template<typename T>
    if_expression<T> _if(const boolean_expression<T> & expr) const;

    void uni_vpermps(const Xbyak::Xmm& x1, const uint8_t mask[4], const Xbyak::Operand& op);
    void uni_vpermps(const Xbyak::Ymm& y1, const uint8_t mask[8], const Xbyak::Operand& op);
    void uni_vpermps(const Xbyak::Zmm& z1, const uint8_t mask[16], const Xbyak::Operand& op);
    void uni_vblendps(const Xbyak::Xmm& x1, const Xbyak::Xmm& x2, uint16_t mask);
    void uni_vblendps(const Xbyak::Ymm& y1, const Xbyak::Ymm& y2, uint16_t mask);
    void uni_vblendps(const Xbyak::Zmm& z1, const Xbyak::Zmm& z2, uint16_t mask);

    void postamble();

    const Xbyak::AddressFrame & address_frame(size_t size) const;
    const reg_indices & free_x64regs() const;
    const reg_indices & free_rmmregs() const;

private:
    reg_indices _free_x64regs;
    reg_indices _free_rmmregs;
    internal::consts_table _consts;
    std::unordered_map<size_t, std::unique_ptr<jit_emitter>> _emitters;
};

template <>
const Xbyak::Reg64& jit_kernel::reserve<Xbyak::Reg64>();

template<typename T>
void jit_kernel::copy(const Xbyak::Reg64& dst,
                      const Xbyak::Reg64& src,
                      const Xbyak::Reg64& size) {
    const auto & addr_frame = address_frame(sizeof(T));
    auto p = reserve<typename reg_traits_by_size<sizeof(T)>::type>();
    foreach(0, size, [&](const Xbyak::Reg64& idx) {
        mov(p, addr_frame[src + idx * sizeof(T)]);
        mov(addr_frame[dst + idx * sizeof(T)], p);
    });
    free(p);
}

template<typename T>
void jit_kernel::copy(const Xbyak::Address& dst,
                      const Xbyak::Reg64& src,
                      const Xbyak::Reg64& size) {
    const auto & addr_frame = address_frame(sizeof(T));
    auto p = reserve<typename reg_traits_by_size<sizeof(T)>::type>();
    auto d = reserve<Xbyak::Reg64>();
    lea(d, dst);
    foreach(0, size, [&](const Xbyak::Reg64& idx) {
        mov(p, addr_frame[src + idx * sizeof(T)]);
        mov(addr_frame[d + idx * sizeof(T)], p);
    });
    free(d);
    free(p);
}

template<typename DstT, size_t N, typename SrcT>
void jit_kernel::load(const variable<DstT[N]> & dst, const variable<SrcT> & src, size_t length) {
    static_assert(std::is_same<typename variable<SrcT>::reg_type, const Xbyak::Reg64>::value,
        "Source register must be Reg64");

    using src_type = typename std::remove_cv<
                        typename std::remove_pointer<SrcT>::type>::type;
    using dst_type = typename std::remove_cv<
                        typename std::remove_pointer<DstT>::type>::type;

    const std::vector<size_t> pool_vec_idxs(_free_rmmregs.begin(), _free_rmmregs.end());
    const std::vector<size_t> pool_gpr_idxs(_free_x64regs.begin(), _free_x64regs.end());

    const auto src_prc = internal::type2precision<src_type>();
    const auto dst_prc = internal::type2precision<dst_type>();

    const auto key = load_emitter_params(src_prc, dst_prc, length).hash();
    if (!_emitters[key]) {
        _emitters[key].reset(new jit_load_emitter(this, internal::get_current_isa(), src_prc, dst_prc, length));
    }
    _emitters[key]->emit_code(
        { static_cast<size_t>(static_cast<const Xbyak::Operand&>(src).getIdx()) },
        { static_cast<size_t>(static_cast<const Xbyak::Operand&>(dst).getIdx()) },
        pool_vec_idxs,
        pool_gpr_idxs);
}

template<typename DstT, size_t N, typename SrcT>
void jit_kernel::load(const variable<DstT[N]> & dst, const variable<SrcT> & src, const variable<size_t> & length) {
    using src_type = typename std::remove_cv<
                        typename std::remove_pointer<SrcT>::type>::type;

    auto s = stack(N * sizeof(src_type));
    s.clear();

    auto tmp = var<SrcT>();
    tmp = s.pointer();

    copy<src_type>(tmp, src, length);

    load(dst, tmp);
}

template<typename DstT, typename SrcT, size_t N>
void jit_kernel::store(const variable<DstT> & dst, const variable<SrcT[N]> & src, size_t length) {
    static_assert(std::is_same<typename variable<DstT>::reg_type, const Xbyak::Reg64>::value,
        "Destination register must be Reg64");

    using src_type = typename std::remove_cv<
                        typename std::remove_pointer<SrcT>::type>::type;
    using dst_type = typename std::remove_cv<
                        typename std::remove_pointer<DstT>::type>::type;

    const std::vector<size_t> pool_vec_idxs(_free_rmmregs.begin(), _free_rmmregs.end());
    const std::vector<size_t> pool_gpr_idxs(_free_x64regs.begin(), _free_x64regs.end());

    const auto src_prc = internal::type2precision<src_type>();
    const auto dst_prc = internal::type2precision<dst_type>();

    const auto key = store_emitter_params(src_prc, dst_prc, length).hash();
    if (!_emitters[key]) {
        _emitters[key].reset(new jit_store_emitter(this, internal::get_current_isa(), src_prc, dst_prc, length));
    }
    _emitters[key]->emit_code(
        { static_cast<size_t>(static_cast<const Xbyak::Operand&>(src).getIdx()) },
        { static_cast<size_t>(static_cast<const Xbyak::Operand&>(dst).getIdx()) },
        pool_vec_idxs,
        pool_gpr_idxs);
}

template<typename DstT, typename SrcT, size_t N>
void jit_kernel::store(const variable<DstT> & dst, const variable<SrcT[N]> & src, const variable<size_t> & length) {
    using dst_type = typename std::remove_cv<
                        typename std::remove_pointer<DstT>::type>::type;

    auto s = stack(N * sizeof(dst_type));

    auto tmp = var<DstT>();
    tmp = s.pointer();

    store(tmp, src);

    copy<dst_type>(dst, tmp, length);
}

template<typename B, typename E, typename S>
void jit_kernel::foreach(const B & begin,
                         const E & end,
                         std::function<void(const variable<size_t>&)> && fn,
                         const S & step) {
    using namespace Xbyak;

    Label loop, exit;

    auto idx = var<size_t>();

    idx = begin;

    L(loop);
    cmp(idx, end);
    jge(exit, T_NEAR);

    fn(idx);

    add(idx, step);
    jmp(loop, T_NEAR);
    L(exit);
}

template<typename T>
jit_kernel::variable<T> jit_kernel::var() {
    using reg_type = typename reg_traits<T>::type;
    const auto & reg = reserve<reg_type>();
    return variable<T>(*this, internal::make_shared(reg, *this));
}

template<typename T>
jit_kernel::variable<T> jit_kernel::var(const T & val) {
    using reg_type = typename reg_traits<T>::type;
    const auto & reg = reserve<reg_type>();
    variable<T> res(*this, internal::make_shared(reg, *this));
    res = val;
    return res;
}

template<typename T>
const T & jit_kernel::constant(const T & c) {
    auto res = _consts.store(&c, sizeof c);
    return *reinterpret_cast<const T*>(res);
}

template<typename T>
const T * jit_kernel::constant(const T * c, size_t size) {
    auto res = _consts.store(c, size * sizeof(T));
    return reinterpret_cast<const T*>(res);
}

template<typename T>
jit_kernel::if_expression<T> jit_kernel::_if(const boolean_expression<T> & expr) const {
    return if_expression<T>(expr);
}

namespace internal {

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// shared_reg

template<typename Reg>
shared_reg<Reg> make_shared(Reg & reg, jit_kernel & kernel) {
    std::shared_ptr<Reg> ptr(&reg, [&kernel](Reg *preg) {
        try {
            kernel.free(*preg);
        } catch(...) {}
    });
    return ptr;
}

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// boolean_expression

template<typename T>
boolean_expression<T>::boolean_expression(jit_kernel & kernel, type t, const shared_reg<reg_type> & lhs, const shared_reg<reg_type> & rhs)
    : _kernel(kernel)
    , _type(t)
    , _lhs(lhs)
    , _rhs(rhs)
    , _rvalue {} {
}

template<typename T>
boolean_expression<T>::boolean_expression(jit_kernel & kernel, type t, const shared_reg<reg_type> & lhs, T rhs)
    : _kernel(kernel)
    , _type(t)
    , _lhs(lhs)
    , _rvalue(rhs) {
}

template<typename T>
void boolean_expression<T>::cmp(const Xbyak::Label & exit) const {
    if (_rhs)
        _kernel.cmp(*_lhs, *_rhs);
    else
        _kernel.cmp(*_lhs, _rvalue);

    switch (_type) {
        case type::eq: {
            _kernel.jne(exit, Xbyak::CodeGenerator::T_NEAR);
            break;
        }
        case type::neq: {
            _kernel.je(exit, Xbyak::CodeGenerator::T_NEAR);
            break;
        }
        case type::ls: {
            _kernel.jge(exit, Xbyak::CodeGenerator::T_NEAR);
            break;
        }
        case type::gt: {
            _kernel.jle(exit, Xbyak::CodeGenerator::T_NEAR);
            break;
        }
        case type::le: {
            _kernel.jg(exit, Xbyak::CodeGenerator::T_NEAR);
            break;
        }
        case type::ge: {
            _kernel.jl(exit, Xbyak::CodeGenerator::T_NEAR);
            break;
        }
    }
}

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// then_expression

template<typename T>
then_expression<T>::then_expression(if_expression<T> & expr)
    : _if_expr(expr) {}

template<typename T>
template<typename F>
void then_expression<T>::_else(F && fn) {
    fn();
    _if_expr._expr._kernel.L(_if_expr._exit);
    _if_expr._is_exit_valid = true;
}

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// variable

template<typename T>
variable_base<T, register_tag>::variable_base(jit_kernel & krnl, const shared_reg<reg_type> & reg)
    : _kernel(krnl)
    , _reg(reg) {
}

template<typename T>
variable_base<T, register_tag>::variable_base(const variable_base & rhs)
    : _kernel(rhs._kernel)
    , _reg(rhs._reg) {
}

template<typename T>
variable_base<T, register_tag>::variable_base(variable_base && rhs)
    : _kernel(rhs._kernel)
    , _reg(std::move(rhs._reg)) {
}

template<typename T>
variable_base<T, memory_tag>::variable_base(jit_kernel & krnl, const shared_reg<reg_type> & addr)
    : _kernel(krnl)
    , _addr(addr) {
}

template<typename T>
variable_base<T, memory_tag>::variable_base(const variable_base & rhs)
    : _kernel(rhs._kernel)
    , _addr(rhs._addr) {
}

template<typename T>
variable_base<T, memory_tag>::variable_base(variable_base && rhs)
    : _kernel(rhs._kernel)
    , _addr(std::move(rhs._addr)) {
}

template<typename T>
variable<T, register_tag>::variable(jit_kernel & krnl)
    : base(krnl, make_shared(krnl.reserve<typename reg_traits<T>::type>(), krnl)) {
}

template<typename T>
variable<T, register_tag>::variable(jit_kernel & krnl, const shared_reg<reg_type> & reg)
    : base(krnl, reg) {
}

template<typename T>
variable<T, memory_tag>::variable(jit_kernel & krnl, const shared_reg<reg_type> & reg)
    : base(krnl, reg) {
}

template<typename T>
const variable<T, memory_tag> & variable<T, memory_tag>::operator = (const variable<T, register_tag> & rhs) const {
    const auto & addr_frame = base::_kernel.address_frame(sizeof(T));
    base::_kernel.mov(addr_frame[base::reg()], rhs);
    return *this;
}

template<typename T, size_t N>
variable<T[N], register_tag>::variable(jit_kernel & krnl)
    : base(krnl, make_shared(krnl.reserve<typename reg_traits<T[N]>::type>(), krnl)) {
}

template<typename T, size_t N>
variable<T[N], register_tag>::variable(jit_kernel & krnl, const shared_reg<reg_type> & reg)
    : base(krnl, reg) {
}

}   // namespace internal

}   // namespace intel_cpu
}   // namespace ov
