diff --git a/src/cpu/x64/jit_brgemm_conv_bwd_utils.cpp b/src/cpu/x64/jit_brgemm_conv_bwd_utils.cpp index 2bb39f12dc6..f3e9291e0ee 100644 --- a/src/cpu/x64/jit_brgemm_conv_bwd_utils.cpp +++ b/src/cpu/x64/jit_brgemm_conv_bwd_utils.cpp @@ -1601,9 +1601,6 @@ status_t init_jcp(jit_brgemm_conv_conf_t &jcp, cpu_isa_t isa, VDISPATCH_CONV_IC(IMPLICATION(is_f32, one_of(isa, avx512_core, avx2)), "unsupported isa for current datatype combination"); - VDISPATCH_CONV_IC(post_ops_ok(jcp, attr, diff_src_d, is_deconv), - VERBOSE_UNSUPPORTED_POSTOP); - jcp.amx_h = 16; jcp.amx_w = 64 / jcp.src_dsz; @@ -1655,6 +1652,11 @@ status_t init_jcp(jit_brgemm_conv_conf_t &jcp, cpu_isa_t isa, format_tag_t src_tag = pick(jcp.ndims - 3, nwc, nhwc, ndhwc); CHECK(init_tag(jcp.src_tag, diff_dst_md, diff_dst_d, src_tag)); + CHECK(init_tag(jcp.dst_tag, diff_src_md, diff_src_d, src_tag)); + CHECK(attr.set_default_formats(&diff_src_md)); + + VDISPATCH_CONV_IC(post_ops_ok(jcp, attr, diff_src_d, is_deconv), + VERBOSE_UNSUPPORTED_POSTOP); return status::success; } @@ -2010,7 +2012,6 @@ status_t init_conf(jit_brgemm_conv_conf_t &jcp, cpu_isa_t isa, = div_up(rnd_up(jcp.gemm_batch_size * sc_size, P4K), sc_size); CHECK(pick_tags(jcp, diff_dst_md, weights_md, diff_src_md, bias_md)); - CHECK(attr.set_default_formats(&diff_src_md)); jcp.buffer_size = jcp.LDC * (jcp.M > 0 ? jcp.M : jcp.M_tail); diff --git a/src/cpu/x64/jit_brgemm_conv_utils.cpp b/src/cpu/x64/jit_brgemm_conv_utils.cpp index 71449db3ca1..f961288f153 100644 --- a/src/cpu/x64/jit_brgemm_conv_utils.cpp +++ b/src/cpu/x64/jit_brgemm_conv_utils.cpp @@ -1773,9 +1773,6 @@ status_t init_jcp(jit_brgemm_conv_conf_t &jcp, cpu_isa_t isa, IMPLICATION(is_f32, one_of(isa, avx512_core, avx2) || jcp.is_bf32), "unsupported isa for current datatype combination"); - VDISPATCH_CONV_IC( - post_ops_ok(jcp, attr, dst_d), VERBOSE_UNSUPPORTED_POSTOP); - jcp.amx_h = 16; jcp.amx_w = 64 / (jcp.is_bf32 ? types::data_type_size(bf16) : jcp.src_dsz); @@ -1817,9 +1814,16 @@ status_t init_jcp(jit_brgemm_conv_conf_t &jcp, cpu_isa_t isa, if (!jcp.wei_plain && jcp.prop_kind != prop_kind::backward_weights) { // fast check data layout before spending time for blocking selection format_tag_t src_tag = pick(jcp.ndims - 3, nwc, nhwc, ndhwc); - CHECK(init_tag( - jcp.src_tag, src_md, src_d, src_tag, is_any_eligible(jcp))); + const bool any_eligible = is_any_eligible(jcp); + CHECK(init_tag(jcp.src_tag, src_md, src_d, src_tag, any_eligible)); + CHECK(init_tag(jcp.dst_tag, dst_md, dst_d, src_tag, any_eligible)); } + + CHECK(attr.set_default_formats(&dst_md)); + + VDISPATCH_CONV_IC( + post_ops_ok(jcp, attr, dst_d), VERBOSE_UNSUPPORTED_POSTOP); + if (jcp.with_bias) { if (bias_d.format_kind() == format_kind::any) CHECK(memory_desc_init_by_tag(bias_md, x)); @@ -2217,7 +2221,6 @@ status_t init_conf(jit_brgemm_conv_conf_t &jcp, bool use_inversion, if (!jcp.wei_plain) CHECK(pick_tags(jcp, src_md, weights_md, dst_md, bias_md)); - CHECK(attr.set_default_formats(&dst_md)); jcp.buffer_size = jcp.LDC * jcp.M;