[CPU] NV12toRGB and NV12toBGR operations for CPU plugin (#8628)

src/plugins/intel_cpu/src/cpu_types.cpp

@@ -103,6 +103,8 @@ const InferenceEngine::details::caseless_unordered_map<std::string, Type> type_t
         { "ReadValue", MemoryInput},  // for construction from name ctor, arbitrary name is used
         { "Assign", MemoryOutput },  // for construction from layer ctor
         { "Convert", Convert },
+        { "NV12toRGB", ColorConvert },
+        { "NV12toBGR", ColorConvert },
         { "MVN", MVN},
         { "NormalizeL2", NormalizeL2},
         { "ScatterUpdate", ScatterUpdate},
@@ -275,6 +277,8 @@ std::string NameFromType(const Type type) {
             return "TensorIterator";
         case Convert:
             return "Convert";
+        case ColorConvert:
+            return "ColorConvert";
         case NormalizeL2:
             return "NormalizeL2";
         case ScatterUpdate:
@@ -470,6 +474,8 @@ std::string algToString(const Algorithm alg) {
     CASE(MathTan);
     CASE(TensorIteratorCommon);
     CASE(TensorIteratorLoop);
+    CASE(ColorConvertNV12toRGB);
+    CASE(ColorConvertNV12toBGR);
 #undef CASE
     return "Undefined";
 }

src/plugins/intel_cpu/src/cpu_types.h

@@ -55,6 +55,7 @@ enum Type {
     DeformableConvolution,
     TensorIterator,
     Convert,
+    ColorConvert,
     MVN,
     NormalizeL2,
     ScatterUpdate,
@@ -220,10 +221,12 @@ enum Algorithm {
     MathSoftPlus,
     MathSoftsign,
     MathTan,
-
     // TensorIterator
     TensorIteratorCommon,
-    TensorIteratorLoop
+    TensorIteratorLoop,
+    // Color conversions
+    ColorConvertNV12toRGB,
+    ColorConvertNV12toBGR,
 };
 
 extern const InferenceEngine::details::caseless_unordered_map<std::string, Type> type_to_name_tbl;

src/plugins/intel_cpu/src/emitters/jit_load_store_emitters.cpp

@@ -763,7 +763,9 @@ template <typename Vmm>
                 if (is_signed) {
                     h->vpmovsdb(addr(0), vmm);
                 } else {
-                    h->vpmaxsd(vmm, vmm, Vmm(aux_vec_idxs[0]));
+                    Vmm zero(aux_vec_idxs[0]);
+                    h->uni_vpxor(zero, zero, zero);
+                    h->vpmaxsd(vmm, vmm, zero);
                     h->vpmovusdb(addr(0), vmm);
                 }
             } else {
@@ -774,7 +776,9 @@ template <typename Vmm>
                 if (is_signed) {
                     h->vpmovsdb(addr(0), vmm | k_mask);
                 } else {
-                    h->vpmaxsd(vmm, vmm, Vmm(aux_vec_idxs[0]));
+                    Vmm zero(aux_vec_idxs[0]);
+                    h->uni_vpxor(zero, zero, zero);
+                    h->vpmaxsd(vmm, vmm, zero);
                     h->vpmovusdb(addr(0), vmm | k_mask);
                 }
             }
@@ -845,7 +849,9 @@ template <typename Vmm>
                     if (is_signed) {
                         h->vpmovsdw(ptr[reg + offset], vmm);  // singed int32 saturate to signed int16.
                     } else {
-                        h->vmaxsd(vmm, Vmm(aux_vec_idxs[0]), vmm);        // if singed bit is 1, set value as 0.
+                        Vmm zero(aux_vec_idxs[0]);
+                        h->uni_vpxor(zero, zero, zero);
+                        h->vmaxsd(vmm, zero, vmm);        // if singed bit is 1, set value as 0.
                         h->vpmovusdw(ptr[reg + offset], vmm); // unsinged int32 saturate to unsigned int16.
                     }
                 } else {
@@ -856,7 +862,9 @@ template <typename Vmm>
                     if (is_signed) {
                         h->vpmovsdw(ptr[reg + offset], vmm | k_mask);
                     } else {
-                        h->vmaxsd(vmm, Vmm(aux_vec_idxs[0]), vmm);
+                        Vmm zero(aux_vec_idxs[0]);
+                        h->uni_vpxor(zero, zero, zero);
+                        h->vmaxsd(vmm, zero, vmm);
                         h->vpmovusdw(ptr[reg + offset], vmm | k_mask);
                     }
                 }

src/plugins/intel_cpu/src/mkldnn_node.h

@@ -127,6 +127,9 @@ private:
 
 class MKLDNNNode {
 public:
+    MKLDNNNode(const MKLDNNNode &) = delete;
+    MKLDNNNode & operator = (const MKLDNNNode &) = delete;
+
     using AttrPtr = std::shared_ptr<mkldnn::primitive_attr>;
 
 public:
@@ -444,7 +447,7 @@ public:
         return execIndex;
     }
 
-    std::string getTypeStr() const {
+    const std::string & getTypeStr() const {
         return typeStr;
     }
 

src/plugins/intel_cpu/src/mkldnn_nodes_factory.cpp

@@ -81,6 +81,7 @@
 #include "nodes/mkldnn_if_node.h"
 #include "nodes/mkldnn_ctc_greedy_decoder_node.h"
 #include "nodes/mkldnn_non_zero.h"
+#include "nodes/mkldnn_color_convert_node.h"
 #include "nodes/subgraph.h"
 
 #define MKLDNN_NODE(__prim, __type) \
@@ -172,4 +173,5 @@ MKLDNNPlugin::MKLDNNNode::NodesFactory::NodesFactory()
     MKLDNN_NODE(MKLDNNGRNNode, GRN);
     MKLDNN_NODE(MKLDNNNonZeroNode, NonZero);
     MKLDNN_NODE(MKLDNNSnippetNode, Subgraph);
+    MKLDNN_NODE(MKLDNNColorConvertNode, ColorConvert);
 }

src/plugins/intel_cpu/src/nodes/mkldnn_color_convert_node.cpp

@@ -0,0 +1,659 @@
+// Copyright (C) 2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "mkldnn_color_convert_node.h"
+#include <memory_desc/dnnl_blocked_memory_desc.h>
+#include <openvino/op/nv12_to_bgr.hpp>
+#include <openvino/op/nv12_to_rgb.hpp>
+#include <openvino/core/type.hpp>
+#include <ie/ie_parallel.hpp>
+#include <utils/jit_kernel.hpp>
+
+using namespace InferenceEngine;
+using namespace mkldnn::impl::utils;
+using namespace mkldnn::impl::cpu::x64;
+using namespace Xbyak;
+
+namespace MKLDNNPlugin {
+
+namespace {
+
+std::tuple<Algorithm, std::string> getAlgorithmFor(const std::shared_ptr<const ngraph::Node>& op) {
+    if (ov::is_type<ov::op::v8::NV12toRGB>(op))
+        return std::make_tuple(Algorithm::ColorConvertNV12toRGB, std::string());
+    if (ov::is_type<ov::op::v8::NV12toBGR>(op))
+        return std::make_tuple(Algorithm::ColorConvertNV12toBGR, std::string());
+    return std::make_tuple(Algorithm::Default, "Only v8::NV12toRGB or v8::NV12toBGR operation is supported");
+}
+
+namespace nv12 {
+
+MKLDNNColorConvertNode::Converter::PrimitiveDescs supportedPrimitiveDescs(MKLDNNNode *node) {
+    const LayoutType layout = LayoutType::ncsp; // 0,1,2,3
+
+    const Precision precision = node->getOriginalInputPrecisionAtPort(0) == Precision::U8
+                                    ? Precision::U8
+                                    : Precision::FP32;
+
+    MKLDNNColorConvertNode::Converter::PrimitiveDescs descs;
+
+    descs.emplace_back(std::vector<PortConfigurator> { node->getOriginalInputsNumber(), { layout, precision } },
+                        std::vector<PortConfigurator> { { layout, precision } },
+                        mayiuse(cpu_isa_t::sse41)
+                            ? impl_desc_type::jit_uni
+                            : impl_desc_type::ref,
+                        true);
+
+    return std::move(descs);
+}
+
+class Converter : public MKLDNNColorConvertNode::Converter {
+    using Base = MKLDNNColorConvertNode::Converter;
+
+public:
+    Converter(MKLDNNNode *node);
+
+protected:
+    Shapes shapeInfer() const override;
+    bool singlePlane() const;
+
+    template<typename T>
+    void convert(const T* y,
+                 const T* uv,
+                 T* dst,
+                 size_t batch_size,
+                 size_t height,
+                 size_t width,
+                 size_t stride_y,
+                 size_t stride_uv);
+};
+
+Converter::Converter(MKLDNNNode *node)
+    : Base(node, node->getAlgorithm() == Algorithm::ColorConvertNV12toRGB
+                        ? ColorFormat { { 0, 1, 2 } }
+                        : ColorFormat { { 2, 1, 0 } }) {
+    if (node->getOriginalInputsNumber() != (singlePlane() ? 1: 2))
+        IE_THROW() <<"NV12Converter node has incorrect number of inputs";
+    if (!node->getOriginalOutputsNumber())
+        IE_THROW() <<"NV12Converter node has incorrect number of outputs";
+}
+
+MKLDNNColorConvertNode::Converter::Shapes
+Converter::shapeInfer() const {
+    const auto & dims = inputDims(0);
+    if (dims.size() != 4)
+        IE_THROW() <<"NV12Converter node has incorrect input dimensions";
+    return singlePlane()
+                ? Shapes { { dims[N_DIM], dims[H_DIM] * 2 / 3, dims[W_DIM], 3 } }
+                : Shapes { { dims[N_DIM], dims[H_DIM], dims[W_DIM], 3 } };
+}
+
+bool Converter::singlePlane() const {
+    return _node->getOriginalInputsNumber() == 1;
+}
+
+template<typename T>
+void Converter::convert(const T* y,
+                        const T* uv,
+                        T* dst,
+                        size_t batch_size,
+                        size_t height,
+                        size_t width,
+                        size_t stride_y,
+                        size_t stride_uv) {
+    InferenceEngine::parallel_for2d(batch_size, height, [&](int batch, int h) {
+        T* out = dst + batch * width * height * 3;
+        auto y_ptr = y + batch * stride_y;
+        auto uv_ptr = uv + batch * stride_uv;
+
+        for (int w = 0; w < width; w++) {
+            auto y_index = h * width + w;
+            auto y_val = static_cast<float>(y_ptr[y_index]);
+            auto uv_index = (h / 2) * width + (w / 2) * 2;
+            auto u_val = static_cast<float>(uv_ptr[uv_index]);
+            auto v_val = static_cast<float>(uv_ptr[uv_index + 1]);
+            auto c = y_val - 16.f;
+            auto d = u_val - 128.f;
+            auto e = v_val - 128.f;
+            auto clip = [](float a) -> T {
+                if (std::is_integral<T>()) {
+                    return static_cast<T>(std::min(std::max(std::round(a), 0.f), 255.f));
+                } else {
+                    return static_cast<T>(std::min(std::max(a, 0.f), 255.f));
+                }
+            };
+
+            auto r = clip(1.164f * c + 1.596f * e);
+            auto g = clip(1.164f * c - 0.391f * d - 0.813f * e);
+            auto b = clip(1.164f * c + 2.018f * d);
+
+            out[y_index * 3 + _colorFormat[0]] = r;
+            out[y_index * 3 + _colorFormat[1]] = g;
+            out[y_index * 3 + _colorFormat[2]] = b;
+        }
+    });
+}
+
+template<typename T, impl_desc_type I>
+class SinglePlaneConvert;
+template<typename T, impl_desc_type I>
+class TwoPlaneConvert;
+
+template<typename T>
+class SinglePlaneConvert<T, impl_desc_type::ref> : public Converter {
+public:
+    using Converter::Converter;
+
+    void execute(mkldnn::stream strm) override {
+        const auto & dims = inputDims(0);
+
+        const size_t batch_size = dims[N_DIM];
+        const size_t height = dims[H_DIM] * 2 / 3;
+        const size_t width = dims[W_DIM];
+
+        const T* y = static_cast<const T*>(input(0));
+        const T* uv = y + width * height;
+        T* dst = static_cast<T*>(output(0));
+
+        convert<T>(y, uv, dst,
+                   batch_size,
+                   height,
+                   width,
+                   height * width * 3 / 2,
+                   height * width * 3 / 2);
+    }
+};
+
+template<typename T>
+class TwoPlaneConvert<T, impl_desc_type::ref> : public Converter {
+public:
+    using Converter::Converter;
+
+    void execute(mkldnn::stream strm) override {
+        const auto & dims = inputDims(0);
+
+        const T* y = static_cast<const T*>(input(0));
+        const T* uv = static_cast<const T*>(input(1));
+        T* dst = static_cast<T*>(output(0));
+
+        const size_t batch_size = dims[N_DIM];
+        const size_t height = dims[H_DIM];
+        const size_t width = dims[W_DIM];
+
+        convert<T>(y, uv, dst,
+                   batch_size,
+                   height,
+                   width,
+                   height * width,
+                   height * width / 2);
+    }
+};
+
+struct jit_uni_converter : public jit_kernel {
+    DECLARE_CPU_JIT_AUX_FUNCTIONS(jit_uni_converter)
+
+    struct Params {
+        const void * y;
+        const void * uv;
+        void * dst;
+        size_t width;
+        uint8_t colorFormat;    // RGB: 0, BGR: !=0
+    };
+
+    typedef void (*function_t)(const Params *);
+
+    void operator()(const Params & args) const {
+        _fn(&args);
+    }
+
+    template<typename T>
+    static const jit_uni_converter & get();
+
+protected:
+    jit_uni_converter() = default;
+
+    function_t _fn;
+};
+
+template<typename T, cpu_isa_t isa>
+class jit_uni_converter_impl : public jit_uni_converter {
+    void generate() override;
+};
+
+template<typename T, cpu_isa_t isa>
+void jit_uni_converter_impl<T, isa>::generate() {
+    using reg_type = typename isa_traits<isa>::reg::type;
+    using var_type = variable<float[isa_traits<isa>::reg::length]>;
+    constexpr auto reg_capacity = isa_traits<isa>::reg::length;
+
+    preamble();
+
+    // Get arguments addresses
+    auto y = arg<const T *>(&Params::y);
+    auto uv = arg<const T *>(&Params::uv);
+    auto dst = arg<T *>(&Params::dst);
+    auto width = arg(&Params::width);
+    auto colorFormat = arg(&Params::colorFormat);
+
+    // Reserve registars
+    auto consts = reserve<Reg64>();
+    auto tmp = var<float[reg_capacity]>();
+    auto y_val = var<float[reg_capacity]>();
+    auto u_val = var<float[reg_capacity]>();
+    auto v_val = var<float[reg_capacity]>();
+    auto r = var<float[reg_capacity]>();
+    auto g = var<float[reg_capacity]>();
+    auto b = var<float[reg_capacity]>();
+
+    // Aliases
+    const auto & c = y_val;
+    const auto & d = u_val;
+    const auto & e = v_val;
+    const auto & uv_val = tmp;
+
+    const uint8_t even_mask = 0xA0;         // 0b10100000
+    const uint8_t odd_mask = 0xF5;          // 0b11110101
+    static const float data[8] = { 16.f, 128.f, 1.164f, 1.596f, 0.391f, 2.018f, 0.813f, 255.f };
+
+    mov(consts, (size_t)data);
+
+    auto clip = [this](const reg_type & op, const reg_type & a, const reg_type & b) {
+        if (std::is_integral<T>())
+            uni_vroundps(op, op, 0);
+        uni_vmaxps(op, op, a);
+        uni_vminps(op, op, b);
+    };
+
+    auto blend = [&, this](const reg_type & r, const reg_type & g, const reg_type & b) {
+        /*
+            Input:
+            r0,r1,r2,r3,r4,r5,r6,r7
+            g0,g1,g2,g3,g4,g5,g6,g7
+            b0,b1,b2,b3,b4,b5,b6,b7
+
+            Permutation:
+            r0,r3,r6,r1,r4,r7,r2,r5
+            g5,g0,g3,g6,g1,g4,g7,g2
+            b2,b5,b0,b3,b6,b1,b4,b7
+
+            Blend
+            r0,g0,xx,r1,g1,xx,r2,g2     blend 1+2 by mask 10210210
+            r0,g0,b0,r1,g1,b1,r2,g2     blend +3  by mask 00100100
+
+            xx,r3,g3,xx,r4,g4,xx,r5     blend 1+2 by mask 02102102
+            b2,r3,g3,b3,r4,g4,b4,r5     blend +3  by mask 01001001
+
+            g5,xx,r6,g6,xx,r7,g7,xx     blend 1+2 by mask 21021021
+            g5,b5,r6,g6,b6,r7,g7,b7     blend +3  by mask 10010010
+
+            Result
+            c = r0,g0,b0,r1,g1,b1,r2,g2
+            d = b2,r3,g3,b3,r4,g4,b4,r5
+            e = g5,b5,r6,g6,b6,r7,g7,b7
+        */
+
+        auto genPermutationMask = [](int offset) -> std::array<int, isa_traits<isa>::reg::length> {
+            std::array<int, isa_traits<isa>::reg::length> mask {};
+
+            if (!!(isa & cpu_isa_t::avx)) {
+                for (int i = 0; i < mask.size(); ++i)
+                    mask[(i * 3 + offset) % mask.size()] = i;
+            } else {
+                int & m0 = mask.front();
+                for (int i = 0; i < 4; ++i)
+                    m0 |= i << ((i * 3 + offset) % 4) * 2;
+            }
+
+            return std::move(mask);
+        };
+
+        static const auto permutationMask4r = genPermutationMask(0);
+        static const auto permutationMask4g = genPermutationMask(1);
+        static const auto permutationMask4b = genPermutationMask(2);
+
+        uni_vpermps(r, permutationMask4r.data(), r);
+        uni_vpermps(g, permutationMask4g.data(), g);
+        uni_vpermps(b, permutationMask4b.data(), b);
+
+        auto blendWithMask = [&](int offset, const var_type & result) {
+            static const uint32_t blendMasks[2] = {
+                0x92492492,
+                0x24924924
+            };
+            const uint16_t mask0 = static_cast<const uint16_t>(blendMasks[0] >> ((offset * reg_capacity) % 3));
+            const uint16_t mask1 = static_cast<const uint16_t>(blendMasks[1] >> ((offset * reg_capacity) % 3));
+
+            result = r;
+            result.blend(g, mask0);
+            result.blend(b, mask1);
+        };
+
+        blendWithMask(0, c);
+        blendWithMask(1, d);
+        blendWithMask(2, e);
+    };  // blend
+
+    auto colorConvert = [&](const reg_type & y_val, const reg_type & uv_val) {
+        uni_vshufps(u_val, uv_val, uv_val, even_mask);              // u_val = tmp[0,0,2,2,4,4,6,6]
+        uni_vshufps(v_val, uv_val, uv_val, odd_mask);               // v_val = tmp[1,1,3,3,5,5,7,7]
+
+        uni_vbroadcastss(tmp, ptr[consts + 0 * sizeof(float)]);     // tmp = [16.0f,16.0f,...]
+        uni_vsubps(c, y_val, tmp);                                  // c = y_val - tmp
+        uni_vbroadcastss(tmp, ptr[consts + 1 * sizeof(float)]);     // tmp = [128.f,128.f,...]
+        uni_vsubps(d, u_val, tmp);                                  // d = u_val - tmp
+        uni_vsubps(e, v_val, tmp);                                  // e = v_val - tmp
+
+        uni_vbroadcastss(tmp, ptr[consts + 2 * sizeof(float)]);     // tmp = [1.164f,1.164f,...]
+        uni_vmulps(c, c, tmp);                                      // c = c * tmp
+
+        uni_vbroadcastss(r, ptr[consts + 3 * sizeof(float)]);       // r = [1.596f,1.596f,...]
+        uni_vmulps(r, r, e);                                        // r = r * e
+        uni_vaddps(r, r, c);                                        // r = r + c
+
+        uni_vbroadcastss(g, ptr[consts + 4 * sizeof(float)]);       // g = [0.391f,0.391f,...]
+        uni_vmulps(g, g, d);                                        // g = g * d
+        uni_vsubps(g, c, g);                                        // g = c - g
+        uni_vbroadcastss(tmp, ptr[consts + 6 * sizeof(float)]);     // tmp = [0.813f,0.813f,...]
+        uni_vmulps(tmp, tmp, e);                                    // tmp = tmp * e
+        uni_vsubps(g, g, tmp);                                      // g = g - tmp
+
+        uni_vbroadcastss(b, ptr[consts + 5 * sizeof(float)]);       // b = [2.018f,2.018f,...]
+        uni_vmulps(b, b, d);                                        // b = b * d
+        uni_vaddps(b, b, c);                                        // b = b + c
+
+        // clip
+        uni_vxorps(c, c, c);
+        uni_vbroadcastss(d, ptr[consts + 7 * sizeof(float)]);
+
+        clip(r, c, d);
+        clip(g, c, d);
+        clip(b, c, d);
+
+        _if(colorFormat == 0)
+        ._then([&]{ blend(r, g, b); })
+        ._else([&]{ blend(b, g, r); });
+    };
+
+    const size_t reg_capacity_log = static_cast<size_t>(std::logb(reg_capacity));
+    const size_t step = reg_capacity * sizeof(T);
+
+    width >>= reg_capacity_log;
+
+    foreach(0, width, [&](const Reg64 & idx) {
+        load(y_val, y);
+        load(uv_val, uv);
+
+        colorConvert(y_val, uv_val);
+
+        store(dst, c);  dst += step;
+        store(dst, d);  dst += step;
+        store(dst, e);  dst += step;
+
+        y += step;
+        uv += step;
+    });
+
+    mov(width, argPtr(&Params::width));
+    width &= reg_capacity - 1;
+
+    _if(width != 0)
+    ._then([&] {
+        auto s = stack(3 * step);
+        s.clear();
+
+        copy<T>(s.pointer(), y, width);
+        copy<T>(ptr[s.pointer() + step], uv, width);
+
+        y = s.pointer();
+        lea(uv, ptr[s.pointer() + step]);
+
+        load(y_val, y);
+        load(uv_val, uv);
+
+        colorConvert(y_val, uv_val);
+
+        store(y, c);    y += step;
+        store(y, d);    y += step;
+        store(y, e);
+
+        lea(width, ptr[width + width * 2]);
+        copy<T>(ptr[dst], s.pointer(), width);
+    });
+
+    postamble();
+}
+
+template<typename T>
+const jit_uni_converter & jit_uni_converter::get() {
+    auto createKernel = []() {
+        std::unique_ptr<jit_uni_converter> kernel;
+
+        if (mayiuse(cpu_isa_t::avx512_common)) {
+            kernel.reset(new jit_uni_converter_impl<T, cpu_isa_t::avx512_common>);
+        } else if (mayiuse(cpu_isa_t::avx2)) {
+            kernel.reset(new jit_uni_converter_impl<T, cpu_isa_t::avx2>);
+        } else if (mayiuse(cpu_isa_t::sse41)) {
+            kernel.reset(new jit_uni_converter_impl<T, cpu_isa_t::sse41>);
+        } else {
+            IE_THROW() << "Can't create jit color converter kernel";
+        }
+
+        if (kernel->create_kernel() != status::success)
+            IE_THROW() << "Can't generate jit color converter kernel";
+        kernel->_fn = (function_t)kernel->jit_ker();
+
+        return std::move(kernel);
+    };
+
+    static auto kernel = createKernel();
+
+    return *kernel;
+}
+
+template<typename T>
+class SinglePlaneConvert<T, impl_desc_type::jit_uni> : public Converter {
+public:
+    using Converter::Converter;
+
+    void execute(mkldnn::stream strm) override {
+        const auto & kernel = jit_uni_converter::get<T>();
+        const auto & dims = inputDims(0);
+
+        const size_t batch_size = dims[N_DIM];
+        const size_t height = dims[H_DIM] * 2 / 3;
+        const size_t width = dims[W_DIM];
+
+        const T* y = static_cast<const T*>(input(0));
+        const T* uv = y + width * height;
+        T* dst = static_cast<T*>(output(0));
+
+        const size_t stride_y = height * width * 3 / 2;
+        const size_t stride_uv = height * width * 3 / 2;
+
+        InferenceEngine::parallel_for2d(batch_size, height, [&](int batch, int h) {
+            typename jit_uni_converter::Params args;
+            args.y = y + batch * stride_y + h * width;
+            args.uv = uv + batch * stride_uv + (h / 2) * width;
+            args.dst = dst + (batch * width * height + h * width) * 3;
+            args.width = width;
+            args.colorFormat = _colorFormat[0]; // The first byte is enough to determine the RGB or BGR format.
+            kernel(args);
+        });
+    }
+};
+
+template<typename T>
+class TwoPlaneConvert<T, impl_desc_type::jit_uni> : public Converter {
+public:
+    using Converter::Converter;
+
+    void execute(mkldnn::stream strm) override {
+        const auto & kernel = jit_uni_converter::get<T>();
+        const auto & dims = inputDims(0);
+
+        const size_t batch_size = dims[N_DIM];
+        const size_t height = dims[H_DIM];
+        const size_t width = dims[W_DIM];
+
+        const T* y = static_cast<const T*>(input(0));
+        const T* uv = static_cast<const T*>(input(1));
+        T* dst = static_cast<T*>(output(0));
+
+        const size_t stride_y = height * width;
+        const size_t stride_uv = height * width / 2;
+
+        InferenceEngine::parallel_for2d(batch_size, height, [&](int batch, int h) {
+            typename jit_uni_converter::Params args;
+            args.y = y + batch * stride_y + h * width;
+            args.uv = uv + batch * stride_uv + (h / 2) * width;
+            args.dst = dst + (batch * width * height + h * width) * 3;
+            args.width = width;
+            args.colorFormat = _colorFormat[0]; // The first byte is enough to determine the RGB or BGR format.
+            kernel(args);
+        });
+    }
+};
+
+}   // namespace nv12
+}   // namespace
+
+MKLDNNColorConvertNode::Converter::Converter(MKLDNNNode *node, const ColorFormat & colorFormat)
+    : _node(node)
+    , _colorFormat(colorFormat) {
+}
+
+InferenceEngine::Precision MKLDNNColorConvertNode::Converter::inputPrecision(size_t idx) const {
+    return _node->getParentEdgeAt(idx)->getMemory().GetDescWithType<BlockedMemoryDesc>()->getPrecision();
+}
+
+InferenceEngine::Precision MKLDNNColorConvertNode::Converter::outputPrecision(size_t idx) const {
+    return _node->getChildEdgeAt(idx)->getMemory().GetDescWithType<BlockedMemoryDesc>()->getPrecision();
+}
+
+const void * MKLDNNColorConvertNode::Converter::input(size_t idx) const {
+    return _node->getParentEdgeAt(idx)->getMemoryPtr()->GetPtr();
+}
+
+void * MKLDNNColorConvertNode::Converter::output(size_t idx) const {
+    return _node->getChildEdgeAt(idx)->getMemoryPtr()->GetPtr();
+}
+
+const VectorDims & MKLDNNColorConvertNode::Converter::inputDims(size_t idx) const {
+    return _node->getParentEdgesAtPort(idx)[0]->getMemory().getStaticDims();
+}
+
+bool MKLDNNColorConvertNode::isSupportedOperation(const std::shared_ptr<const ngraph::Node>& op, std::string& errorMessage) noexcept {
+    Algorithm alg;
+    std::tie(alg, errorMessage) = getAlgorithmFor(op);
+    return alg != Algorithm::Default;
+}
+
+MKLDNNColorConvertNode::MKLDNNColorConvertNode(const std::shared_ptr<ngraph::Node>& op,
+                                               const mkldnn::engine& eng,
+                                               MKLDNNWeightsSharing::Ptr &cache)
+    : MKLDNNNode(op, eng, cache) {
+    std::string errorMessage;
+    std::tie(algorithm, errorMessage) = getAlgorithmFor(op);
+    if (algorithm == Algorithm::Default)
+        IE_THROW(NotImplemented) << errorMessage;
+}
+
+void MKLDNNColorConvertNode::getSupportedDescriptors() {}
+
+void MKLDNNColorConvertNode::initSupportedPrimitiveDescriptors() {
+    if (supportedPrimitiveDescriptors.empty()) {
+        switch (algorithm) {
+            case Algorithm::ColorConvertNV12toRGB:
+            case Algorithm::ColorConvertNV12toBGR: {
+                for (const auto &desc : nv12::supportedPrimitiveDescs(this)) {
+                    const auto & inPortConfigs = std::get<0>(desc);
+                    const auto & outPortConfigs = std::get<1>(desc);
+                    const auto implType = std::get<2>(desc);
+                    const auto dynBatchSupport = std::get<3>(desc);
+                    addSupportedPrimDesc(inPortConfigs, outPortConfigs, implType, dynBatchSupport);
+                }
+                initSupportedNV12Impls();
+                break;
+            default:
+                break;
+            }
+        }
+    }
+}
+
+void MKLDNNColorConvertNode::initSupportedNV12Impls() {
+    #define SUPPORTED_IMPL(Impl, type, desc_type)                           \
+        [](MKLDNNNode *node) {                                              \
+            return new nv12::Impl<type, impl_desc_type::desc_type>(node);   \
+        };
+
+    // ref
+    {
+        auto &impls = _supportedImpls[impl_desc_type::ref][algorithm];
+        impls[Precision::U8][true] = SUPPORTED_IMPL(SinglePlaneConvert, uint8_t, ref);
+        impls[Precision::U8][false] = SUPPORTED_IMPL(TwoPlaneConvert, uint8_t, ref);
+        impls[Precision::FP32][true] = SUPPORTED_IMPL(SinglePlaneConvert, float, ref);
+        impls[Precision::FP32][false] = SUPPORTED_IMPL(TwoPlaneConvert, float, ref);
+    }
+
+    // jit_uni
+    {
+        auto &impls = _supportedImpls[impl_desc_type::jit_uni][algorithm];
+        impls[Precision::U8][true] = SUPPORTED_IMPL(SinglePlaneConvert, uint8_t, jit_uni);
+        impls[Precision::U8][false] = SUPPORTED_IMPL(TwoPlaneConvert, uint8_t, jit_uni);
+        impls[Precision::FP32][true] = SUPPORTED_IMPL(SinglePlaneConvert, float, jit_uni);
+        impls[Precision::FP32][false] = SUPPORTED_IMPL(TwoPlaneConvert, float, jit_uni);
+    }
+
+    #undef SUPPORTED_IMPL
+}
+
+void MKLDNNColorConvertNode::createPrimitive() {
+    const NodeDesc *desc = getSelectedPrimitiveDescriptor();
+    if (!desc)
+        IE_THROW() << getTypeStr() + " node with name '" + getName() + "' "
+                   << "no optimal primitive descriptor selected";
+
+    if (!_impl) {
+        const auto & cfg = desc->getConfig();
+        const auto precision = cfg.inConfs[0].desc->getPrecision();
+        const bool isSinglePlane = cfg.inConfs.size() == 1;
+
+        _impl = std::unique_ptr<Converter>(_supportedImpls
+                                            .at(desc->getImplementationType())
+                                            .at(algorithm)
+                                            .at(precision)
+                                            .at(isSinglePlane)(this));
+    }
+}
+
+void MKLDNNColorConvertNode::execute(mkldnn::stream strm) {
+    if (!_impl)
+        IE_THROW() << getTypeStr() + " node with name '" + getName() + "' "
+                   << "has no any implemented converter";
+    _impl->execute(strm);
+}
+
+bool MKLDNNColorConvertNode::created() const {
+    return getType() == ColorConvert;
+}
+
+std::vector<VectorDims> MKLDNNColorConvertNode::shapeInfer() const {
+    if (!_impl)
+        IE_THROW() << getTypeStr() + " node with name '" + getName() + "' "
+                   << "has no any implemented converter";
+    return _impl->shapeInfer();
+}
+
+bool MKLDNNColorConvertNode::needPrepareParams() const {
+    return false;
+}
+
+void MKLDNNColorConvertNode::executeDynamicImpl(mkldnn::stream strm) {
+    execute(strm);
+}
+
+REG_MKLDNN_PRIM_FOR(MKLDNNColorConvertNode, ColorConvert);
+
+}  // namespace MKLDNNPlugin

src/plugins/intel_cpu/src/nodes/mkldnn_color_convert_node.h

@@ -0,0 +1,75 @@
+// Copyright (C) 2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <mkldnn_node.h>
+#include <utils/multidim_map.hpp>
+#include <functional>
+#include <tuple>
+#include <array>
+
+namespace MKLDNNPlugin {
+
+class MKLDNNColorConvertNode : public MKLDNNNode {
+public:
+    MKLDNNColorConvertNode(const std::shared_ptr<ngraph::Node>& op,
+                           const mkldnn::engine& eng,
+                           MKLDNNWeightsSharing::Ptr &cache);
+    class Converter;
+
+public:
+    void getSupportedDescriptors() override;
+    void initSupportedPrimitiveDescriptors() override;
+    void createPrimitive() override;
+    void execute(mkldnn::stream strm) override;
+    bool created() const override;
+    std::vector<VectorDims> shapeInfer() const override;
+    bool needPrepareParams() const override;
+    void executeDynamicImpl(mkldnn::stream strm) override;
+
+    static bool isSupportedOperation(const std::shared_ptr<const ngraph::Node>& op, std::string& errorMessage) noexcept;
+
+private:
+    void initSupportedNV12Impls();
+
+private:
+    using ConverterBuilder = std::function<Converter*(MKLDNNNode *)>;
+    using SupportedImpls = multidim_map<impl_desc_type, Algorithm, InferenceEngine::Precision::ePrecision, bool, ConverterBuilder>;
+
+    std::unique_ptr<Converter> _impl;
+    SupportedImpls _supportedImpls;
+};
+
+class MKLDNNColorConvertNode::Converter {
+public:
+    using PrimitiveDescs = std::vector<std::tuple<std::vector<PortConfigurator>,
+                                                    std::vector<PortConfigurator>,
+                                                    impl_desc_type,
+                                                    bool>>;
+    using Shapes = std::vector<VectorDims>;
+
+    static constexpr size_t N_DIM = 0;
+    static constexpr size_t H_DIM = 1;
+    static constexpr size_t W_DIM = 2;
+    static constexpr size_t C_DIM = 3;
+
+    using ColorFormat = std::array<uint8_t, 3>;
+
+    Converter(MKLDNNNode *node, const ColorFormat & colorFormat);
+    virtual ~Converter() = default;
+    InferenceEngine::Precision inputPrecision(size_t idx) const;
+    InferenceEngine::Precision outputPrecision(size_t idx) const;
+    const void * input(size_t idx) const;
+    void * output(size_t idx) const;
+    const VectorDims & inputDims(size_t idx) const;
+    virtual Shapes shapeInfer() const = 0;
+    virtual void execute(mkldnn::stream strm) = 0;
+
+protected:
+    MKLDNNNode *_node;
+    ColorFormat _colorFormat;   // RGB: {0,1,2}, BGR: {2,1,0}
+};
+
+}  // namespace MKLDNNPlugin

src/plugins/intel_cpu/src/utils/blob_dump.cpp

@@ -18,8 +18,8 @@ using namespace InferenceEngine;
 namespace MKLDNNPlugin {
 
 // IEB file format routine
-static unsigned char IEB_MAGIC[4] = {'I', 'E', 'B', '0'};
-static unsigned char NO_SCALES = 0xFF;
+static const unsigned char IEB_MAGIC[4] = {'I', 'E', 'B', '0'};
+static const unsigned char NO_SCALES = 0xFF;
 
 struct IEB_HEADER {
     unsigned char magic[4];

src/plugins/intel_cpu/src/utils/jit_kernel.cpp

@@ -0,0 +1,334 @@
+// Copyright (C) 2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "jit_kernel.hpp"
+#include <stdexcept>
+
+using namespace dnnl::impl::cpu::x64;
+using namespace Xbyak;
+
+namespace MKLDNNPlugin {
+
+namespace {
+
+template<typename RegType>
+using registers = std::array<std::reference_wrapper<const RegType>, 16>;
+
+template<typename RegType>
+const RegType & reserveReg(jit_kernel::reg_indices & freeRegs, const registers<RegType> & regs) {
+    if (freeRegs.empty())
+        throw std::runtime_error("No free registers");
+    const auto idx = freeRegs.back();
+    freeRegs.pop_back();
+    return regs[idx];
+}
+
+template<typename RegType>
+void freeReg(jit_kernel::reg_indices & freeRegs, const registers<RegType> & regs, const RegType & reg) {
+    const auto idx = reg.getIdx();
+    // Debug:
+    // auto it = std::find(freeRegs.begin(), freeRegs.end(), idx);
+    // if (it != freeRegs.end())
+    //     throw std::runtime_error("Some register was freed twice");
+    freeRegs.emplace_back(idx);
+    if (freeRegs.size() > regs.size())
+        throw std::runtime_error("Some register was freed twice");
+}
+
+const registers<Reg64> & x64regs() {
+    using namespace Xbyak::util;
+    static const registers<Reg64> _x64regs {{
+        rax, rcx, rdx, rbx,
+        rsp, rbp, rsi, rdi,
+        r8,  r9,  r10, r11,
+        r12, r13, r14, r15,
+    }};
+    return _x64regs;
+}
+
+const registers<Reg32> & x32regs() {
+    using namespace Xbyak::util;
+    static const registers<Reg32> _x32regs {{
+        eax,  ecx,  edx,  ebx,
+        esp,  ebp,  esi,  edi,
+        r8d,  r9d,  r10d, r11d,
+        r12d, r13d, r14d, r15d,
+    }};
+    return _x32regs;
+}
+
+const registers<Reg16> & x16regs() {
+    using namespace Xbyak::util;
+    static const registers<Reg16> _x16regs {{
+        ax,   cx,   dx,   bx,
+        sp,   bp,   si,   di,
+        r8w,  r9w,  r10w, r11w,
+        r12w, r13w, r14w, r15w,
+    }};
+    return _x16regs;
+}
+
+const registers<Reg8> & x8regs() {
+    using namespace Xbyak::util;
+    static const registers<Reg8> _x8regs {{
+        al,   cl,   dl,   bl,
+        spl,  bpl,  sil,  dil,
+        r8b,  r9b,  r10b, r11b,
+        r12b, r13b, r14b, r15b,
+    }};
+    return _x8regs;
+}
+
+const registers<Xmm> & xmmregs() {
+    static const registers<Xmm> _xmmregs {{
+        util::xmm0,  util::xmm1,  util::xmm2,  util::xmm3,
+        util::xmm4,  util::xmm5,  util::xmm6,  util::xmm7,
+        util::xmm8,  util::xmm9,  util::xmm10, util::xmm11,
+        util::xmm12, util::xmm13, util::xmm14, util::xmm15,
+    }};
+    return _xmmregs;
+}
+
+const registers<Ymm> & ymmregs() {
+    static const registers<Ymm> _ymmregs {{
+        util::ymm0,  util::ymm1,  util::ymm2,  util::ymm3,
+        util::ymm4,  util::ymm5,  util::ymm6,  util::ymm7,
+        util::ymm8,  util::ymm9,  util::ymm10, util::ymm11,
+        util::ymm12, util::ymm13, util::ymm14, util::ymm15,
+    }};
+    return _ymmregs;
+}
+
+const registers<Zmm> & zmmregs() {
+    static const registers<Zmm> _zmmregs {{
+        util::zmm0,  util::zmm1,  util::zmm2,  util::zmm3,
+        util::zmm4,  util::zmm5,  util::zmm6,  util::zmm7,
+        util::zmm8,  util::zmm9,  util::zmm10, util::zmm11,
+        util::zmm12, util::zmm13, util::zmm14, util::zmm15,
+    }};
+    return _zmmregs;
+}
+
+}   // namespace
+
+namespace internal {
+
+template<>
+InferenceEngine::Precision type2precision<float>() {
+    return InferenceEngine::Precision::FP32;
+}
+
+template<>
+InferenceEngine::Precision type2precision<uint8_t>() {
+    return InferenceEngine::Precision::U8;
+}
+
+cpu_isa_t get_current_isa() {
+    if (mayiuse(cpu_isa_t::avx512_common))
+        return cpu_isa_t::avx512_common;
+    if (mayiuse(cpu_isa_t::avx2))
+        return cpu_isa_t::avx2;
+    return cpu_isa_t::sse41;
+}
+
+stack_frame::stack_frame(MKLDNNPlugin::jit_kernel & kernel, size_t size)
+    : _kernel(kernel)
+    , _size(size) {
+    if (_size)
+        _kernel.sub(_kernel.rsp, _size);
+}
+
+stack_frame::stack_frame(stack_frame && rhs)
+    : _kernel(rhs._kernel)
+    , _size(rhs._size) {
+    rhs._size = 0;
+}
+
+stack_frame::~stack_frame() {
+    if (_size)
+        _kernel.add(_kernel.rsp, _size);
+}
+
+const Xbyak::Reg64 & stack_frame::pointer() const {
+    return _kernel.rsp;
+}
+
+void stack_frame::clear() const {
+    const size_t end = _size & ~(size_t)7u;
+
+    _kernel.foreach(0, end, [&](const Reg64 & idx) {
+        _kernel.mov(_kernel.qword[pointer() + idx], 0);
+    }, sizeof(size_t));
+
+    if (end < _size) {
+        _kernel.foreach(end, _size, [&](const Reg64 & idx) {
+            _kernel.mov(_kernel.byte[pointer() + idx], 0);
+        });
+    }
+}
+
+}   // namespace internal
+
+jit_kernel::jit_kernel()
+    : _load_emitter(this, internal::get_current_isa())
+    , _store_emitter(this, internal::get_current_isa()) {
+    _free_rmmregs.reserve(16);
+    _free_rmmregs.reserve(16);
+
+    auto isRegReserved = [this](int idx) {
+        return idx == param1.getIdx()           // function argument
+                || idx == Operand::Code::RSP    // stack pointer
+                || idx == Operand::Code::RBP;   // frame pointer
+    };
+
+    for (int reg = Operand::Code::RAX; reg <= Operand::Code::R15; ++reg) {
+        if (!isRegReserved(reg))
+            _free_x64regs.emplace_back(reg);
+        _free_rmmregs.emplace_back(reg);
+    }
+}
+
+template<>
+const Reg64 & jit_kernel::reserve<Reg64>() {
+    return reserveReg(_free_x64regs, x64regs());
+}
+
+template<>
+const Reg32 & jit_kernel::reserve<Reg32>() {
+    return reserveReg(_free_x64regs, x32regs());
+}
+
+template<>
+const Reg16 & jit_kernel::reserve<Reg16>() {
+    return reserveReg(_free_x64regs, x16regs());
+}
+
+template<>
+const Reg8 & jit_kernel::reserve<Reg8>() {
+    return reserveReg(_free_x64regs, x8regs());
+}
+
+template<>
+void jit_kernel::free<Reg64>(const Reg64 & reg) {
+    freeReg(_free_x64regs, x64regs(), reg);
+}
+
+template<>
+void jit_kernel::free<Reg32>(const Reg32 & reg) {
+    freeReg(_free_x64regs, x32regs(), reg);
+}
+
+template<>
+void jit_kernel::free<Reg16>(const Reg16 & reg) {
+    freeReg(_free_x64regs, x16regs(), reg);
+}
+
+template<>
+void jit_kernel::free<Reg8>(const Reg8 & reg) {
+    freeReg(_free_x64regs, x8regs(), reg);
+}
+
+template<>
+const Xmm & jit_kernel::reserve<Xmm>() {
+    return reserveReg(_free_rmmregs, xmmregs());
+}
+
+template<>
+void jit_kernel::free<Xmm>(const Xmm & reg) {
+    freeReg(_free_rmmregs, xmmregs(), reg);
+}
+
+template<>
+const Ymm & jit_kernel::reserve<Ymm>() {
+    return reserveReg(_free_rmmregs, ymmregs());
+}
+
+template<>
+void jit_kernel::free<Ymm>(const Ymm & reg) {
+    freeReg(_free_rmmregs, ymmregs(), reg);
+}
+
+template<>
+const Zmm & jit_kernel::reserve<Zmm>() {
+    return reserveReg(_free_rmmregs, zmmregs());
+}
+
+template<>
+void jit_kernel::free<Zmm>(const Zmm & reg) {
+    freeReg(_free_rmmregs, zmmregs(), reg);
+}
+
+void jit_kernel::postamble() {
+    jit_generator::postamble();
+    if (_is_load_emitter_used)
+        _load_emitter.emit_data();
+    if (_is_store_emitter_used)
+        _store_emitter.emit_data();
+}
+
+const AddressFrame & jit_kernel::address_frame(size_t size) const {
+        switch (size) {
+            case 1: return byte;
+            case 2: return word;
+            case 4: return dword;
+            case 8: return qword;
+            case 16: return xword;
+            case 32: return yword;
+            case 64: return zword;
+            default:
+                break;
+        }
+        return ptr;
+}
+
+jit_kernel::stack_frame jit_kernel::stack(size_t size) {
+    return stack_frame(*this, size);
+}
+
+void jit_kernel::uni_vpermps(const Xmm& x1, const int *mask, const Operand& op) {
+    uint8_t imm8 = static_cast<uint8_t>(*mask);
+    mov(x1, op);
+    shufps(x1, op, imm8);
+}
+
+void jit_kernel::uni_vpermps(const Ymm& y1, const int *mask, const Operand& op) {
+    auto mreg = reserve<Ymm>();
+    auto mptr = reserve<Reg64>();
+
+    mov(mptr, (size_t)mask);
+    uni_vmovdqu(mreg, ptr[mptr]);
+    vpermps(y1, mreg, op);
+
+    free(mreg);
+    free(mptr);
+}
+
+void jit_kernel::uni_vpermps(const Zmm& z1, const int *mask, const Operand& op) {
+    auto mreg = reserve<Zmm>();
+    auto mptr = reserve<Reg64>();
+
+    mov(mptr, (size_t)mask);
+    uni_vmovdqu(mreg, ptr[mptr]);
+    vpermps(z1, mreg, op);
+
+    free(mreg);
+    free(mptr);
+}
+
+void jit_kernel::uni_vblendps(const Xbyak::Xmm& x1, const Xbyak::Xmm& x2, uint16_t mask) {
+    blendps(x1, x2, mask);
+}
+
+void jit_kernel::uni_vblendps(const Xbyak::Ymm& y1, const Xbyak::Ymm& y2, uint16_t mask) {
+    vblendps(y1, y1, y2, static_cast<uint8_t>(mask));
+}
+
+void jit_kernel::uni_vblendps(const Xbyak::Zmm& z1, const Xbyak::Zmm& z2, uint16_t mask) {
+    auto reg = var<uint32_t>();
+    mov(reg, mask);
+    kmovw(k1, reg);
+    vblendmps(z1 | k1, z1, z2);
+}
+
+}   // namespace MKLDNNPlugin

src/plugins/intel_cpu/src/utils/jit_kernel.hpp

@@ -0,0 +1,697 @@
+// Copyright (C) 2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+#include <cpu/x64/jit_generator.hpp>
+#include <emitters/jit_load_store_emitters.hpp>
+#include <ie/ie_precision.hpp>
+#include <common/nstl.hpp>
+#include <type_traits>
+#include <functional>
+#include <vector>
+#include <array>
+#include <tuple>
+
+namespace MKLDNNPlugin {
+
+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_any;
+};
+
+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_any;
+};
+
+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_any;
+};
+
+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_any;
+};
+
+template<typename T>
+struct reg_traits : public reg_traits_by_size<sizeof(T)> {};
+
+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_any;
+};
+template<>
+struct reg_traits<double> : public reg_traits<float> {};
+
+template<typename T>
+struct reg_traits<T[1]> : public reg_traits<T> {};
+
+template<typename T>
+struct reg_traits<T[4]> {
+    using type = Xbyak::Xmm;
+    constexpr static size_t size = 4 * 4;       // in bytes
+    constexpr static dnnl::impl::cpu::x64::cpu_isa_t isa
+                        = dnnl::impl::cpu::x64::cpu_isa_t::sse41;
+};
+
+template<typename T>
+struct reg_traits<T[8]> {
+    using type = Xbyak::Ymm;
+    constexpr static size_t size = 8 * 4;       // in bytes
+    constexpr static dnnl::impl::cpu::x64::cpu_isa_t isa
+                        = dnnl::impl::cpu::x64::cpu_isa_t::avx2;
+};
+
+template<typename T>
+struct reg_traits<T[16]> {
+    using type = Xbyak::Zmm;
+    constexpr static size_t size = 16 * 4;      // in bytes
+    constexpr static dnnl::impl::cpu::x64::cpu_isa_t isa
+                        = dnnl::impl::cpu::x64::cpu_isa_t::avx512_common;
+};
+
+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_common> {
+    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>
+class variable;
+template<typename T>
+class if_expression;
+template<typename T>
+class then_expression;
+
+template<typename T>
+class boolean_expression {
+private:
+    using reg_type = typename reg_traits<T>::type;
+
+    enum class type {
+        eq, neq
+    };
+
+    boolean_expression(jit_kernel & kernel, type t, const reg_type & lhs, const reg_type & rhs);
+    boolean_expression(jit_kernel & kernel, type t, const reg_type & lhs, T rhs);
+    void cmp(const Xbyak::Label & exit) const;
+
+    jit_kernel & _kernel;
+    type _type;
+    const reg_type & _lhs;
+
+    bool _is_ref;
+
+    union datum {
+        datum(const reg_type & r)
+            : reg(&r) {}
+        datum(T v)
+            : value(v) {}
+        const reg_type * reg;
+        T value;
+    } _rhs;
+
+    friend class variable<T>;
+    friend class if_expression<T>;
+    friend class then_expression<T>;
+};
+
+template<typename T>
+struct then_expression {
+    then_expression(if_expression<T> & expr);
+
+    template<typename F>
+    void _else(F && fn);
+
+private:
+    if_expression<T> & _if_expr;
+};
+
+template<typename T>
+struct if_expression {
+    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>;
+};
+
+template<typename T>
+class variable_base {
+public:
+    using reg_type = typename reg_traits<T>::type;
+
+    variable_base(const variable_base &) = delete;
+    variable_base & operator = (const variable_base &) = delete;
+    variable_base(variable_base &&);
+
+    operator const reg_type &() const {
+        return _reg;
+    }
+
+    operator Xbyak::RegExp () const {
+        return _reg;
+    }
+
+    jit_kernel & kernel;
+
+protected:
+    variable_base(jit_kernel & krnl, const reg_type & reg);
+    ~variable_base();
+
+    bool _manage_lifetime = true;
+    const reg_type & _reg;
+};
+
+template<typename T>
+class variable : public variable_base<typename std::enable_if<!std::is_floating_point<T>::value, T>::type> {
+public:
+    using type = T;
+    using base = variable_base<type>;
+    using reg_type = typename base::reg_type;
+
+    variable(variable &&) = default;
+    variable(jit_kernel & krnl);
+    variable(jit_kernel & krnl, const reg_type & reg);
+
+    const variable & operator = (const reg_type & rhs) const {
+        base::kernel.mov(base::_reg, rhs);
+        return *this;
+    }
+    const variable & operator = (T rhs) const {
+        base::kernel.mov(base::_reg, rhs);
+        return *this;
+    }
+    const variable & operator += (const reg_type & rhs) const {
+        base::kernel.add(base::_reg, rhs);
+        return *this;
+    }
+    const variable & operator += (typename std::conditional<std::is_pointer<T>::value, size_t, T>::type rhs) const {
+        base::kernel.add(base::_reg, rhs);
+        return *this;
+    }
+    const variable & operator -= (const reg_type & rhs) const {
+        base::kernel.sub(base::_reg, rhs);
+        return *this;
+    }
+    const variable & operator -= (typename std::conditional<std::is_pointer<T>::value, size_t, T>::type rhs) const {
+        base::kernel.sub(base::_reg, rhs);
+        return *this;
+    }
+    const variable & operator &= (const reg_type & rhs) const {
+        base::kernel.and_(base::_reg, rhs);
+        return *this;
+    }
+    const variable & operator &= (T rhs) const {
+        base::kernel.and_(base::_reg, rhs);
+        return *this;
+    }
+    const variable & operator |= (const reg_type & rhs) const {
+        base::kernel.or_(base::_reg, rhs);
+        return *this;
+    }
+    const variable & operator |= (T rhs) const {
+        base::kernel.or_(base::_reg, rhs);
+        return *this;
+    }
+    const variable & operator >>= (size_t rhs) const {
+        base::kernel.shr(base::_reg, rhs);
+        return *this;
+    }
+    const variable & operator <<= (size_t rhs) const {
+        base::kernel.shl(base::_reg, rhs);
+        return *this;
+    }
+
+    boolean_expression<T> operator == (const reg_type & rhs) const {
+        return boolean_expression<T>(base::kernel, boolean_expression<T>::type::eq, base::_reg, rhs);
+    }
+
+    boolean_expression<T> operator == (T rhs) const {
+        return boolean_expression<T>(base::kernel, boolean_expression<T>::type::eq, base::_reg, rhs);
+    }
+
+    boolean_expression<T> operator != (const reg_type & rhs) const {
+        return boolean_expression<T>(base::kernel, boolean_expression<T>::type::neq, base::_reg, rhs);
+    }
+
+    boolean_expression<T> operator != (T rhs) const {
+        return boolean_expression<T>(base::kernel, boolean_expression<T>::type::neq, base::_reg, rhs);
+    }
+
+    // TODO: add necessary operations
+};
+
+template<typename T, size_t N>
+class variable<T[N]> : public variable_base<T[N]> {
+public:
+    using type = T[N];
+    using base = variable_base<type>;
+    using reg_type = typename base::reg_type;
+    constexpr static size_t length = N;
+
+    variable(variable &&) = default;
+    variable(jit_kernel & krnl);
+    variable(jit_kernel & krnl, const reg_type & reg);
+
+    const variable & operator = (const reg_type & rhs) const {
+        base::kernel.uni_vmovups(base::_reg, rhs);
+        return *this;
+    }
+
+    const variable & blend(const reg_type & rhs, uint16_t mask) const {
+        base::kernel.uni_vblendps(base::_reg, rhs, mask);
+        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);
+    stack_frame(stack_frame && rhs);
+    ~stack_frame();
+    const Xbyak::Reg64 & pointer() const;
+    void clear() const;
+
+private:
+    jit_kernel & _kernel;
+    size_t _size;
+};
+
+template<typename T>
+InferenceEngine::Precision type2precision();
+
+dnnl::impl::cpu::x64::cpu_isa_t get_current_isa();
+
+}   // 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;
+    template<typename T>
+    using variable = internal::variable<T>;
+    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, res };
+    }
+
+    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, res };
+    }
+
+    jit_kernel();
+
+    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, typename SrcT>
+    void load(const variable<DstT> & dst, const variable<SrcT> & src);
+    template<typename DstT, typename SrcT>
+    void store(const variable<DstT> & dst, const variable<SrcT> & src);
+
+    template<typename B, typename E, typename S = size_t>
+    void foreach(const B & begin,
+                 const E & end,
+                 std::function<void(const Xbyak::Reg64&)> && fn,
+                 const S & step = 1);
+
+    template<typename T>
+    variable<T> var();
+
+    stack_frame stack(size_t size);
+
+    template<typename T>
+    if_expression<T> _if(const boolean_expression<T> & expr);
+
+    void uni_vpermps(const Xbyak::Xmm& x1, const int *mask, const Xbyak::Operand& op);
+    void uni_vpermps(const Xbyak::Ymm& y1, const int *mask, const Xbyak::Operand& op);
+    void uni_vpermps(const Xbyak::Zmm& z1, const int *mask, 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();
+
+private:
+    const Xbyak::AddressFrame & address_frame(size_t size) const;
+
+    reg_indices _free_x64regs;
+    reg_indices _free_rmmregs;
+    bool _is_load_emitter_used = false;
+    bool _is_store_emitter_used = false;
+    jit_load_emitter _load_emitter;
+    jit_store_emitter _store_emitter;
+};
+
+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, typename SrcT>
+void jit_kernel::load(const variable<DstT> & dst, const variable<SrcT> & src) {
+    static_assert(std::is_same<typename variable<SrcT>::reg_type, Xbyak::Reg64>::value,
+        "Source register must be Reg64");
+
+    using src_type = typename std::remove_cv<
+                        typename std::remove_pointer<
+                            typename std::decay<SrcT>::type>::type>::type;
+    using dst_type = typename std::remove_cv<
+                        typename std::remove_pointer<
+                            typename std::decay<DstT>::type>::type>::type;
+    constexpr size_t length = variable<DstT>::length;
+
+    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());
+
+    _load_emitter.emit_code(
+        { static_cast<size_t>(static_cast<const Xbyak::Operand&>(src).getIdx()) },
+        { static_cast<size_t>(static_cast<const Xbyak::Operand&>(dst).getIdx()) },
+        std::make_shared<load_emitter_context>(
+            internal::type2precision<src_type>(),
+            internal::type2precision<dst_type>(),
+            static_cast<int>(length)),
+        pool_vec_idxs,
+        pool_gpr_idxs);
+
+    _is_load_emitter_used = true;
+}
+
+template<typename DstT, typename SrcT>
+void jit_kernel::store(const variable<DstT> & dst, const variable<SrcT> & src) {
+    static_assert(std::is_same<typename variable<DstT>::reg_type, Xbyak::Reg64>::value,
+        "Destibnation register must be Reg64");
+
+    using src_type = typename std::remove_cv<
+                        typename std::remove_pointer<
+                            typename std::decay<SrcT>::type>::type>::type;
+    using dst_type = typename std::remove_cv<
+                        typename std::remove_pointer<
+                            typename std::decay<DstT>::type>::type>::type;
+    constexpr size_t length = variable<SrcT>::length;
+
+    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());
+
+    _store_emitter.emit_code(
+        { static_cast<size_t>(static_cast<const Xbyak::Operand&>(src).getIdx()) },
+        { static_cast<size_t>(static_cast<const Xbyak::Operand&>(dst).getIdx()) },
+        std::make_shared<store_emitter_context>(
+            internal::type2precision<src_type>(),
+            internal::type2precision<dst_type>(),
+            static_cast<int>(length)),
+        pool_vec_idxs,
+        pool_gpr_idxs);
+
+    _is_store_emitter_used = true;
+}
+
+template<typename B, typename E, typename S>
+void jit_kernel::foreach(const B & begin,
+                         const E & end,
+                         std::function<void(const Xbyak::Reg64&)> && fn,
+                         const S & step) {
+    using namespace Xbyak;
+
+    Label loop, exit;
+
+    auto idx = reserve<Reg64>();
+
+    mov(idx, begin);
+
+    L(loop);
+    cmp(idx, end);
+    jge(exit, T_NEAR);
+
+    fn(idx);
+
+    add(idx, step);
+    jmp(loop, T_NEAR);
+    L(exit);
+
+    free<Reg64>(idx);
+}
+
+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, reg);
+}
+
+template<typename T>
+jit_kernel::if_expression<T> jit_kernel::_if(const boolean_expression<T> & expr) {
+    return if_expression<T>(expr);
+}
+
+namespace internal {
+
+// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+// boolean_expression
+
+template<typename T>
+boolean_expression<T>::boolean_expression(jit_kernel & kernel, type t, const reg_type & lhs, const reg_type & rhs)
+    : _kernel(kernel)
+    , _type(t)
+    , _lhs(lhs)
+    , _is_ref(true)
+    , _rhs(rhs) {
+}
+
+template<typename T>
+boolean_expression<T>::boolean_expression(jit_kernel & kernel, type t, const reg_type & lhs, T rhs)
+    : _kernel(kernel)
+    , _type(t)
+    , _lhs(lhs)
+    , _is_ref(false)
+    , _rhs(rhs) {
+}
+
+template<typename T>
+void boolean_expression<T>::cmp(const Xbyak::Label & exit) const {
+    if (_is_ref)
+        _kernel.cmp(_lhs, *_rhs.reg);
+    else
+        _kernel.cmp(_lhs, _rhs.value);
+
+    switch (_type) {
+        case type::eq: {
+            _kernel.jne(exit, Xbyak::CodeGenerator::T_NEAR);
+            break;
+        }
+        case type::neq: {
+            _kernel.je(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>::variable_base(jit_kernel & krnl, const reg_type & reg)
+    : kernel(krnl)
+    , _reg(reg) {
+}
+
+template<typename T>
+variable_base<T>::variable_base(variable_base && rhs)
+    : kernel(rhs.kernel)
+    , _reg(rhs._reg) {
+    rhs._manage_lifetime = false;
+}
+
+template<typename T>
+variable_base<T>::~variable_base() {
+    if (_manage_lifetime)
+        kernel.free(_reg);
+}
+
+template<typename T>
+variable<T>::variable(jit_kernel & krnl)
+    : base(krnl, krnl.reserve<typename reg_traits<T>::type>()) {
+}
+
+template<typename T>
+variable<T>::variable(jit_kernel & krnl, const reg_type & reg)
+    : base(krnl, reg) {
+}
+
+template<typename T, size_t N>
+variable<T[N]>::variable(jit_kernel & krnl)
+    : base(krnl, krnl.reserve<typename reg_traits<T[N]>::type>()) {
+}
+
+template<typename T, size_t N>
+variable<T[N]>::variable(jit_kernel & krnl, const reg_type & reg)
+    : base(krnl, reg) {
+}
+
+}   // namespace internal
+
+}   // namespace MKLDNNPlugin

src/plugins/intel_cpu/src/utils/multidim_map.hpp

@@ -0,0 +1,64 @@
+// Copyright (C) 2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+#include <unordered_map>
+#include <type_traits>
+#include <functional>
+
+namespace MKLDNNPlugin {
+
+namespace internal {
+
+template <typename K>
+struct enum_hash {
+    std::size_t operator()(K t) const {
+        return static_cast<std::size_t>(t);
+    }
+};
+
+template<typename K>
+using hash_t = typename std::conditional<std::is_enum<K>::value, enum_hash<K>, std::hash<K>>::type;
+
+}   // namespace internal
+
+template<typename K, typename ...Ts>
+struct multidim_map {
+    using key_type = K;
+    using mapped_type = multidim_map<Ts...>;
+    using hash_type = internal::hash_t<K>;
+
+public:
+    mapped_type & operator[](const key_type & key) {
+        return _map[key];
+    }
+
+    const mapped_type & at(const key_type & key) const {
+        return _map.at(key);
+    }
+
+private:
+    std::unordered_map<key_type, mapped_type, hash_type> _map;
+};
+
+template<typename K, typename T>
+struct multidim_map<K, T> {
+    using key_type = K;
+    using mapped_type = T;
+    using hash_type = internal::hash_t<K>;
+
+public:
+    mapped_type & operator[](const key_type & key) {
+        return _map[key];
+    }
+
+    const mapped_type & at(const key_type & key) const {
+        return _map.at(key);
+    }
+
+private:
+    std::unordered_map<key_type, mapped_type, hash_type> _map;
+};
+
+}   // namespace MKLDNNPlugin