Skip to content

Commit e359903

Browse files
committed
x64: brgemm conv: heuristic for big convolution on avx2
1 parent 0503b68 commit e359903

File tree

1 file changed

+17
-6
lines changed

1 file changed

+17
-6
lines changed

src/cpu/x64/jit_brgemm_conv_utils.cpp

+17-6
Original file line numberDiff line numberDiff line change
@@ -1940,12 +1940,17 @@ status_t init_conf(jit_brgemm_conv_conf_t &jcp, cpu_isa_t isa,
19401940
using namespace data_type;
19411941
// ======================= blocking =================================
19421942

1943-
auto bcast_amount
1944-
= static_cast<size_t>(jcp.id) * jcp.ih * jcp.iw * jcp.src_dsz;
1943+
auto bcast_amount = static_cast<size_t>(jcp.id) * jcp.ih * jcp.iw
1944+
* jcp.src_dsz * jcp.ic;
19451945
auto wei_amount = static_cast<size_t>(jcp.oc) * jcp.kd * jcp.kh * jcp.kw
1946-
* jcp.wei_dsz;
1947-
1948-
jcp.loop_order = (bcast_amount < wei_amount) ? loop_ngcdhw : loop_ndhwgc;
1946+
* jcp.wei_dsz * jcp.ic;
1947+
1948+
jcp.loop_order
1949+
= (one_of(isa, avx2, avx2_vnni, avx2_vnni_2) && jcp.mb > jcp.nthr
1950+
&& bcast_amount > brg_blocking_t::L2
1951+
&& wei_amount > brg_blocking_t::L2)
1952+
? loop_gcndhw
1953+
: ((bcast_amount < wei_amount) ? loop_ngcdhw : loop_ndhwgc);
19491954
jcp.brgemm_kernel_loop_order
19501955
= brgemm_kernel_loop_order_t::brgemm_lo_default;
19511956

@@ -1963,6 +1968,13 @@ status_t init_conf(jit_brgemm_conv_conf_t &jcp, cpu_isa_t isa,
19631968
const bool small_amx_job = est_amx_job < 64 || jcp.oc < 256;
19641969
auto start_ocb
19651970
= (is_amx(isa) && jcp.is_os_blocking && small_amx_job) ? 2 : 4;
1971+
if (one_of(isa, avx2, avx2_vnni, avx2_vnni_2)
1972+
&& jcp.loop_order == loop_gcndhw)
1973+
start_ocb = 2;
1974+
if (one_of(isa, avx2, avx2_vnni, avx2_vnni_2)
1975+
&& jcp.oh * jcp.ow >= 150 * 150)
1976+
start_ocb = 2;
1977+
19661978
start_ocb = nstl::min(div_up(jcp.oc, jcp.acc_simd_w), start_ocb);
19671979

19681980
auto finish_ocb = 1;
@@ -2215,7 +2227,6 @@ status_t init_conf(jit_brgemm_conv_conf_t &jcp, cpu_isa_t isa,
22152227
#endif
22162228

22172229
// ============ end blocking ===========================================
2218-
22192230
jcp.brg_type
22202231
= (jcp.use_uker && one_of(jcp.exec_type, exec_base, exec_trans))
22212232
? brgemm_static_offs

0 commit comments

Comments
 (0)