[FORK][FEATURE] Added custom vesrion of JIT DW FP32/BF16 Convolution with 5D input support

Author: dmitry-gorokhov

include/oneapi/dnnl/dnnl.hpp

@@ -1454,6 +1454,7 @@ struct memory : public handle<dnnl_memory_t> {
         BAcde16b16a = dnnl_BAcde16b16a,
         BAcde16a16b = dnnl_BAcde16a16b,
         aBdec32b = dnnl_aBdec32b,
+        Abcdef8a = dnnl_Abcdef8a,
         Abcdef16a = dnnl_Abcdef16a,
         Abcdef32a = dnnl_Abcdef32a,
         Acdb32a = dnnl_Acdb32a,
@@ -1701,6 +1702,7 @@ struct memory : public handle<dnnl_memory_t> {
         IOdhw16i16o = dnnl_IOdhw16i16o,
         gIOhw16i16o = dnnl_gIOhw16i16o,
         gOhwi32o = dnnl_gOhwi32o,
+        Goidhw8g = dnnl_Goidhw8g,
         Goidhw16g = dnnl_Goidhw16g,
         IOw16o16i = dnnl_IOw16o16i,
         OIw16i16o = dnnl_OIw16i16o,

include/oneapi/dnnl/dnnl_types.h

@@ -398,6 +398,7 @@ typedef enum {
     dnnl_aCBdef16c16b,
     dnnl_aBdefc4b,
     dnnl_aBdefc8b,
+    dnnl_Abcdef8a,
     dnnl_Abcdef16a,
     dnnl_Abcdef32a,
     dnnl_aBedc16b,
@@ -1620,6 +1621,7 @@ typedef enum {
     dnnl_gIOdhw8o16i2o = dnnl_aCBdef8b16c2b,
     dnnl_gOIdhw8o8i = dnnl_aBCdef8b8c,
     dnnl_gOIdhw8o4i = dnnl_aBCdef8b4c,
+    dnnl_Goidhw8g = dnnl_Abcdef8a,
     dnnl_Goidhw16g = dnnl_Abcdef16a,
     dnnl_Goidhw32g = dnnl_Abcdef32a,
     dnnl_gOIdhw2i4o2i = dnnl_aBCdef2c4b2c,

src/common/c_types_map.hpp

@@ -677,6 +677,7 @@ const format_tag_t ABcd40a32b = dnnl_ABcd40a32b;
 const format_tag_t ABcde40a32b = dnnl_ABcde40a32b;
 const format_tag_t BAcde16b16a = dnnl_BAcde16b16a;
 const format_tag_t aBdec32b = dnnl_aBdec32b;
+const format_tag_t Abcdef8a = dnnl_Abcdef8a;
 const format_tag_t Abcdef16a = dnnl_Abcdef16a;
 const format_tag_t Abcdef32a = dnnl_Abcdef32a;
 const format_tag_t Acdb32a = dnnl_Acdb32a;
@@ -1174,6 +1175,7 @@ const format_tag_t IOhw16i16o = dnnl_IOhw16i16o;
 const format_tag_t Ohwi32o = dnnl_Ohwi32o;
 const format_tag_t gIOhw16i16o = dnnl_gIOhw16i16o;
 const format_tag_t gOhwi32o = dnnl_gOhwi32o;
+const format_tag_t Goidhw8g = dnnl_Goidhw8g;
 const format_tag_t Goidhw16g = dnnl_Goidhw16g;
 const format_tag_t IOw16o16i = dnnl_IOw16o16i;
 const format_tag_t IOw16i16o = dnnl_IOw16i16o;

src/common/dnnl_debug_autogenerated.cpp

@@ -302,6 +302,7 @@ const char *dnnl_fmt_tag2str(dnnl_format_tag_t v) {
     if (v == dnnl_aCBdef16c16b) return "aCBdef16c16b";
     if (v == dnnl_aBdefc4b) return "aBdefc4b";
     if (v == dnnl_aBdefc8b) return "aBdefc8b";
+    if (v == dnnl_Abcdef8a) return "Abcdef8a";
     if (v == dnnl_Abcdef16a) return "Abcdef16a";
     if (v == dnnl_Abcdef32a) return "Abcdef32a";
     if (v == dnnl_aBedc16b) return "aBedc16b";
@@ -1397,6 +1398,7 @@ const char *dnnl_fmt_tag2str(dnnl_format_tag_t v) {
     if (v == dnnl_gIOdhw8o16i2o) return "gIOdhw8o16i2o";
     if (v == dnnl_gOIdhw8o8i) return "gOIdhw8o8i";
     if (v == dnnl_gOIdhw8o4i) return "gOIdhw8o4i";
+    if (v == dnnl_Goidhw8g) return "Goidhw8g";
     if (v == dnnl_Goidhw16g) return "Goidhw16g";
     if (v == dnnl_Goidhw32g) return "Goidhw32g";
     if (v == dnnl_gOIdhw2i4o2i) return "gOIdhw2i4o2i";

src/common/memory.cpp

@@ -164,7 +164,8 @@ status_t dnnl_memory_create(memory_t **memory, const memory_desc_t *md,
             : memory_flags_t::use_runtime_ptr;
     void *handle_ptr = (handle == DNNL_MEMORY_ALLOCATE) ? nullptr : handle;
     auto _memory = new memory_t(engine, md, flags, handle_ptr);
-    if (_memory == nullptr) return out_of_memory;
+    if (_memory == nullptr)
+        return out_of_memory;
     if (_memory->memory_storage() == nullptr) {
         delete _memory;
         return out_of_memory;

src/common/memory_desc_wrapper.cpp

@@ -586,6 +586,7 @@ status_t memory_desc_wrapper::compute_blocking(
         C(aBdec32b, {0, 1, 3, 4, 2}, {32}, {1});
         C(aCBdef16c16b, {0, 2, 1, 3, 4, 5}, {16, 16}, {2, 1});
         C(aCBdef16b16c, {0, 2, 1, 3, 4, 5}, {16, 16}, {1, 2});
+        C(Abcdef8a, {0, 1, 2, 3, 4, 5}, {8}, {0});
         C(Abcdef16a, {0, 1, 2, 3, 4, 5}, {16}, {0});
         C(Abcdef32a, {0, 1, 2, 3, 4, 5}, {32}, {0});
         C(aCBd16c16b, {0, 2, 1, 3}, {16, 16}, {2, 1});

src/common/tag_traits.hpp

@@ -795,6 +795,7 @@ DECL_TRAITS(aBCde2b8c8b2c, _BC, _2b8c8b2c, 5);
 DECL_TRAITS(aBdec32b, _B, _32b, 5);
 DECL_TRAITS(aCBdef16c16b, _BC, _16c16b, 6);
 DECL_TRAITS(aCBdef16b16c, _BC, _16b16c, 6);
+DECL_TRAITS(Abcdef8a, _A, _8a, 6);
 DECL_TRAITS(Abcdef16a, _A, _16a, 6);
 DECL_TRAITS(aCBd16c16b, _BC, _16c16b, 4);
 DECL_TRAITS(aCBde16c16b, _BC, _16c16b, 4);

src/cpu/cpu_convolution_list.cpp

@@ -52,6 +52,7 @@
 #include "cpu/x64/jit_sse41_1x1_convolution.hpp"
 #include "cpu/x64/jit_sse41_convolution.hpp"
 #include "cpu/x64/jit_uni_dw_convolution.hpp"
+#include "cpu/x64/jit_uni_fork_dw_convolution.hpp"
 #include "cpu/x64/jit_uni_x8s8s32x_1x1_convolution.hpp"
 #include "cpu/x64/jit_uni_x8s8s32x_convolution.hpp"
 using namespace dnnl::impl::cpu::x64;
@@ -126,13 +127,16 @@ const std::map<pk_dt_impl_key_t, std::vector<impl_list_item_t>> &impl_list_map()
             CPU_INSTANCE_AVX512(brgemm_1x1_convolution_fwd_t<avx512_core>)
             CPU_INSTANCE_AVX512(brgemm_convolution_fwd_t<avx512_core>)
             CPU_INSTANCE_AVX512(jit_avx512_common_dw_convolution_fwd_t)
+            CPU_INSTANCE_AVX512(jit_avx512_common_fork_dw_convolution_fwd_t)
             CPU_INSTANCE_AVX512(jit_avx512_common_1x1_convolution_fwd_f32_t)
             CPU_INSTANCE_AVX512(jit_avx512_common_convolution_fwd_t<f32>)
             CPU_INSTANCE_AVX2(jit_avx2_dw_convolution_fwd_t)
             CPU_INSTANCE_AVX2(brgemm_1x1_convolution_fwd_t<avx2>)
             CPU_INSTANCE_AVX2(brgemm_convolution_fwd_t<avx2>)
+            CPU_INSTANCE_AVX2(jit_avx2_fork_dw_convolution_fwd_t)
             CPU_INSTANCE_AVX2(jit_avx2_1x1_convolution_fwd_t)
             CPU_INSTANCE_SSE41(jit_sse41_dw_convolution_fwd_t)
+            CPU_INSTANCE_SSE41(jit_sse41_fork_dw_convolution_fwd_t)
             CPU_INSTANCE_SSE41(jit_sse41_1x1_convolution_fwd_t)
             CPU_INSTANCE_AVX2(jit_avx2_convolution_fwd_t)
             CPU_INSTANCE_SSE41(jit_sse41_convolution_fwd_t)
@@ -166,6 +170,7 @@ const std::map<pk_dt_impl_key_t, std::vector<impl_list_item_t>> &impl_list_map()
             CPU_INSTANCE_AVX512(brgemm_1x1_convolution_fwd_t<avx512_core_bf16>)
             CPU_INSTANCE_AVX512(brgemm_convolution_fwd_t<avx512_core_bf16>)
             CPU_INSTANCE_AVX512(jit_uni_dw_convolution_fwd_t<avx512_core, bf16, f32>)
+            CPU_INSTANCE_AVX512(jit_uni_fork_dw_convolution_fwd_t<avx512_core, bf16, f32>)
             CPU_INSTANCE_AVX512(jit_avx512_core_bf16_1x1_convolution_fwd_t<f32>)
             CPU_INSTANCE_AVX512(jit_avx512_core_bf16_convolution_fwd_t)
             CPU_INSTANCE_AVX512(gemm_bf16_convolution_fwd_t<f32>)
@@ -184,6 +189,7 @@ const std::map<pk_dt_impl_key_t, std::vector<impl_list_item_t>> &impl_list_map()
             CPU_INSTANCE_AVX512(brgemm_1x1_convolution_fwd_t<avx512_core_bf16>)
             CPU_INSTANCE_AVX512(brgemm_convolution_fwd_t<avx512_core_bf16>)
             CPU_INSTANCE_AVX512(jit_uni_dw_convolution_fwd_t<avx512_core, bf16, bf16>)
+            CPU_INSTANCE_AVX512(jit_uni_fork_dw_convolution_fwd_t<avx512_core, bf16, bf16>)
             CPU_INSTANCE_AVX512(jit_avx512_core_bf16_1x1_convolution_fwd_t<bf16>)
             CPU_INSTANCE_AVX512(jit_avx512_core_bf16_convolution_fwd_t)
             CPU_INSTANCE_AVX512(gemm_bf16_convolution_fwd_t<bf16>)
@@ -246,13 +252,15 @@ const std::map<pk_dt_impl_key_t, std::vector<impl_list_item_t>> &impl_list_map()
             CPU_INSTANCE_AVX512(brgemm_convolution_bwd_t<avx512_core>)
             CPU_INSTANCE_AVX512(brgemm_convolution_bwd_strided_t<avx512_core>)
             CPU_INSTANCE_AVX512(jit_avx512_common_dw_convolution_bwd_data_t)
+            CPU_INSTANCE_AVX512(jit_avx512_common_fork_dw_convolution_bwd_data_t)
             CPU_INSTANCE_AVX512(jit_avx512_common_1x1_convolution_bwd_data_f32_t)
             CPU_INSTANCE_AVX512(jit_avx512_common_convolution_bwd_data_t<f32>)
             CPU_INSTANCE_AVX2(brgemm_convolution_bwd_t<avx2>)
             CPU_INSTANCE_AVX2(brgemm_convolution_bwd_strided_t<avx2>)
             CPU_INSTANCE_AVX2(jit_avx2_dw_convolution_bwd_data_t)
+            CPU_INSTANCE_AVX2(jit_avx2_fork_dw_convolution_bwd_data_t)
             CPU_INSTANCE_AVX2(jit_avx2_1x1_convolution_bwd_data_t)
-            CPU_INSTANCE_SSE41(jit_sse41_dw_convolution_bwd_data_t)
+            CPU_INSTANCE_SSE41(jit_sse41_fork_dw_convolution_bwd_data_t)
             CPU_INSTANCE_AVX2(jit_avx2_convolution_bwd_data_t)
             CPU_INSTANCE_AARCH64(jit_uni_dw_convolution_bwd_data_t<sve_512,data_type::f32>)
             CPU_INSTANCE_AARCH64(jit_sve_512_1x1_convolution_bwd_data_f32_t)
@@ -271,6 +279,7 @@ const std::map<pk_dt_impl_key_t, std::vector<impl_list_item_t>> &impl_list_map()
             CPU_INSTANCE_AVX512(brgemm_convolution_bwd_t<avx512_core_bf16>)
             CPU_INSTANCE_AVX512(brgemm_convolution_bwd_strided_t<avx512_core_bf16>)
             CPU_INSTANCE_AVX512(jit_uni_dw_convolution_bwd_data_t<avx512_core, bf16, f32>)
+            CPU_INSTANCE_AVX512(jit_uni_fork_dw_convolution_bwd_data_t<avx512_core, bf16, f32>)
             CPU_INSTANCE_AVX512(jit_avx512_core_bf16_1x1_convolution_bwd_data_t<f32>)
             CPU_INSTANCE_AVX512(jit_avx512_core_bf16_convolution_bwd_data_t)
             CPU_INSTANCE_AVX512(gemm_bf16_convolution_bwd_data_t<f32>)
@@ -287,6 +296,7 @@ const std::map<pk_dt_impl_key_t, std::vector<impl_list_item_t>> &impl_list_map()
             CPU_INSTANCE_AVX512(brgemm_convolution_bwd_t<avx512_core_bf16>)
             CPU_INSTANCE_AVX512(brgemm_convolution_bwd_strided_t<avx512_core_bf16>)
             CPU_INSTANCE_AVX512(jit_uni_dw_convolution_bwd_data_t<avx512_core, bf16, bf16>)
+            CPU_INSTANCE_AVX512(jit_uni_fork_dw_convolution_bwd_data_t<avx512_core, bf16, bf16>)
             CPU_INSTANCE_AVX512(jit_avx512_core_bf16_1x1_convolution_bwd_data_t<bf16>)
             CPU_INSTANCE_AVX512(jit_avx512_core_bf16_convolution_bwd_data_t)
             CPU_INSTANCE_AVX512(gemm_bf16_convolution_bwd_data_t<bf16>)
@@ -783,4 +793,4 @@ const impl_list_item_t *get_convolution_impl_list(
 
 } // namespace cpu
 } // namespace impl
-} // namespace dnnl
+} // namespace dnnl