cpu: x64: brgemm convolution blocking: update oh_block selection

oneapi-src · Feb 9, 2023 · c9cb51d · c9cb51d
1 parent 4e2c573
commit c9cb51d
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/src/cpu/x64/jit_brgemm_conv_utils.cpp b/src/cpu/x64/jit_brgemm_conv_utils.cpp
@@ -572,7 +572,9 @@ void brg_blocking_t::select_ic_block() {
                 ic_block = simd_w;
         }
     } else {
-        const auto est_ur = nstl::min(sp_block, estimate_ur(oc_block));
+        const auto est_ur = sp_block > 0
+                ? nstl::min(sp_block, estimate_ur(oc_block))
+                : estimate_ur(oc_block);
         const auto inp_ur = is_os_blocking ? est_ur : inp_w(est_ur, kw_block);
 
         if (kw_block > 1) {
@@ -1050,13 +1052,16 @@ void brg_blocking_t::iterate_ker_block(brg_blocking_t &best_brgb, int kd_block_,
         kw_block_pad = kw;
     }
 
+    sp_block = -1;
+    select_ic_block();
+
     if (exec_type == exec_vpad) {
         od_block = 1;
         oh_block = 1;
     } else if (exec_type == exec_trans) {
         const auto w_block_size
-                = 2 * src_dsz * ic * iwp + dst_dsz * ow * oc_block;
-        const auto other_size = wei_dsz * kd * kh * kw * ic * oc_block
+                = 2 * src_dsz * ic_block * iwp + dst_dsz * ow * oc_block;
+        const auto other_size = wei_dsz * kd * kh * kw * ic_block * oc_block
                 + acc_dsz * 2 * amx_h * oc_block;
         const auto L2_available = nstl::min(static_cast<size_t>(div_up(L2, 2)),
                 other_size > L2 ? 0 : L2 - other_size);