cpu: x64: jit_direct_reorder: add identical extra buffers support

If extra data is identical for both buffers, direct copy can handle
that. Those buffers are never big, no need to parallelize the copy.

Author: dzarukin

src/cpu/reorder/cpu_reorder_comp_s8_s8.cpp

@@ -1,5 +1,5 @@
 /*******************************************************************************
-* Copyright 2020-2024 Intel Corporation
+* Copyright 2020-2025 Intel Corporation
 * Copyright 2023 FUJITSU LIMITED
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
@@ -28,6 +28,7 @@ const impl_list_map_t &comp_s8_s8_impl_list_map() {
         // s8 -> s8
         {{s8, s8, 2}, {
             DNNL_X64_ONLY(CPU_REORDER_INSTANCE(x64::brgemm_matmul_copy_reorder_t))
+            DNNL_X64_ONLY(CPU_REORDER_INSTANCE(x64::jit_uni_reorder_direct_copy_t))
             DNNL_X64_ONLY(CPU_REORDER_INSTANCE(x64::jit_uni_reorder_t))
             DNNL_AARCH64_ONLY(CPU_REORDER_INSTANCE(aarch64::jit_uni_reorder_t))
             DNNL_NON_X64_ONLY(REG_SR(s8, oi, s8, OI4i16o4i, fmt_order::keep, spec::conv_req_comp))
@@ -50,6 +51,7 @@ const impl_list_map_t &comp_s8_s8_impl_list_map() {
         // s8 -> s8
         {{s8, s8, 3}, {
             DNNL_X64_ONLY(CPU_REORDER_INSTANCE(x64::brgemm_matmul_copy_reorder_t))
+            DNNL_X64_ONLY(CPU_REORDER_INSTANCE(x64::jit_uni_reorder_direct_copy_t))
             DNNL_X64_ONLY(CPU_REORDER_INSTANCE(x64::jit_uni_reorder_t))
             DNNL_AARCH64_ONLY(CPU_REORDER_INSTANCE(aarch64::jit_uni_reorder_t))
             DNNL_NON_X64_ONLY(REG_SR(s8, any, s8, wio, fmt_order::keep, spec::conv_req_comp))
@@ -88,6 +90,7 @@ const impl_list_map_t &comp_s8_s8_impl_list_map() {
             nullptr,
         }},
         {{s8, s8, 4}, {
+            DNNL_X64_ONLY(CPU_REORDER_INSTANCE(x64::jit_uni_reorder_direct_copy_t))
             DNNL_X64_ONLY(CPU_REORDER_INSTANCE(x64::jit_uni_reorder_t))
             DNNL_AARCH64_ONLY(CPU_REORDER_INSTANCE(aarch64::jit_uni_reorder_t))
             DNNL_NON_X64_ONLY(REG_SR(s8, any, s8, hwio, fmt_order::keep, spec::conv_req_comp))
@@ -137,6 +140,7 @@ const impl_list_map_t &comp_s8_s8_impl_list_map() {
             nullptr,
         }},
         {{s8, s8, 5}, {
+            DNNL_X64_ONLY(CPU_REORDER_INSTANCE(x64::jit_uni_reorder_direct_copy_t))
             DNNL_X64_ONLY(CPU_REORDER_INSTANCE(x64::jit_uni_reorder_t))
             DNNL_AARCH64_ONLY(CPU_REORDER_INSTANCE(aarch64::jit_uni_reorder_t))
             DNNL_NON_X64_ONLY(REG_SR(s8, any, s8, hwigo, fmt_order::keep, spec::conv_req_comp))
@@ -183,6 +187,7 @@ const impl_list_map_t &comp_s8_s8_impl_list_map() {
             nullptr,
         }},
         {{s8, s8, 6}, {
+            DNNL_X64_ONLY(CPU_REORDER_INSTANCE(x64::jit_uni_reorder_direct_copy_t))
             DNNL_X64_ONLY(CPU_REORDER_INSTANCE(x64::jit_uni_reorder_t))
             DNNL_AARCH64_ONLY(CPU_REORDER_INSTANCE(aarch64::jit_uni_reorder_t))
             DNNL_NON_X64_ONLY(REG_SR(s8, any, s8, dhwigo, fmt_order::keep, spec::conv_req_comp))

src/cpu/x64/jit_uni_reorder_direct_copy.cpp

@@ -309,9 +309,13 @@ status_t jit_uni_reorder_direct_copy_t::pd_t::init(
     VDISPATCH_REORDER(src_d.similar_to(dst_d, true, false, 0),
             VERBOSE_TENSOR_FORMAT_MISMATCH, "src", "dst");
 
-    VDISPATCH_REORDER(
-            utils::everyone_is(0UL, src_d.extra().flags, dst_d.extra().flags),
-            VERBOSE_UNSUPPORTED_MD_FLAG);
+    VDISPATCH_REORDER(src_d.extra().flags == dst_d.extra().flags,
+            VERBOSE_UNSUPPORTED_MD_FLAG, "src or dst");
+
+    VDISPATCH_REORDER(IMPLICATION(src_d.extra().flags > 0UL,
+                              src_d.additional_buffer_size()
+                                      == dst_d.additional_buffer_size()),
+            VERBOSE_UNSUPPORTED_MD_FLAG, "src or dst");
 
     VDISPATCH_REORDER(attr()->has_default_values(), VERBOSE_UNSUPPORTED_ATTR);
 
@@ -374,6 +378,17 @@ status_t jit_uni_reorder_direct_copy_t::execute(const exec_ctx_t &ctx) const {
                 out + (start + dst_d.offset0()) * dst_dt_size, end - start);
     });
 
+    if (src_d.is_additional_buffer()) {
+        // Verified in pd_t::init();
+        assert(src_d.extra().flags == dst_d.extra().flags);
+
+        const auto additional_size = src_d.additional_buffer_size();
+        const auto data_size = src_d.size(/* index = */ 0,
+                /* include_additional_size = */ false);
+        std::memcpy(out + data_size * dst_dt_size, in + data_size * src_dt_size,
+                additional_size);
+    }
+
     return status::success;
 }