cpu: jit: avx512: check code size against upper bound

src/cpu/jit_avx512_common_conv_kernel.cpp

@@ -1322,6 +1322,7 @@ status_t jit_avx512_common_conv_fwd_kernel::init_conf(
     if (jcp.l_pad > 0 && r_pad > 0)
         n_oi--;
 
+    // Heuristic to optimize code size on KNX
     bool large_code_size = jcp.ur_w != jcp.ow && jcp.l_pad > 0 && r_pad > 0
             && ((jcp.l_pad <= 0 && n_oi > 0) || (jcp.l_pad > 0 && n_oi > 1));
     if (large_code_size) {
@@ -1597,6 +1598,18 @@ status_t jit_avx512_common_conv_fwd_kernel::init_conf(
         }
     }
 
+    // A rough check on code size
+    // TODO: come up with a tighter bound
+    {
+        const int max_code_size = 256 * 1024; // default size of jit generator
+        int mult = 1 + (jcp.l_pad > 0) + (r_pad > 0);
+        const float max_instruction_size = 15;
+        float ur_fac
+                = (float)jcp.kw * jcp.ic_block * jcp.nb_oc_blocking * jcp.ur_w;
+        float code_size = mult * ur_fac * max_instruction_size;
+        if (code_size > max_code_size) return status::unimplemented;
+    }
+
     return status::success;
 }
 
@@ -2368,6 +2381,7 @@ status_t jit_avx512_common_conv_bwd_data_kernel_f32::init_conf(
 
     jcp.loop_order = loop_gnc;
 
+    // Heuristic to optimize code size on KNX
     bool large_code_size = (jcp.ur_w != jcp.ow)
          && ((l_overflow <= 0 && n_oi > 0) ||(l_overflow > 0 && n_oi > 1))
          && (r_overflow1 > 0) && (l_overflow > 0);
@@ -2469,6 +2483,18 @@ status_t jit_avx512_common_conv_bwd_data_kernel_f32::init_conf(
         && jcp.oc <= weights_d.padded_dims()[with_groups + 0];
     if (!args_ok) return status::unimplemented;
 
+    // A rough check on code size
+    // TODO: come up with a tighter bound
+    {
+        const int max_code_size = 256 * 1024; // default size of jit generator
+        int mult = 1 + (l_overflow > 0) + (r_overflow1 > 0);
+        const float max_instruction_size = 15;
+        float ur_fac
+                = (float)jcp.kw * jcp.oc_block * jcp.nb_ic_blocking * jcp.ur_w;
+        float code_size = mult * ur_fac * max_instruction_size;
+        if (code_size > max_code_size) return status::unimplemented;
+    }
+
     return status::success;
 }
 

tests/benchdnn/inputs/test_conv_regression_general

@@ -84,3 +84,8 @@ mb1_g2ic22oc32_ih8oh6kh3sh1dh0ph0_iw8ow6kw3sw1dw0pw0
 #MFDNN-1707
 --reset --mb=1 --dir=FWD_D --cfg=f32
 mb1_g2ic6oc16_ih5oh5kh1sh1dh0ph0_iw5ow5kw1sw1dw0pw0
+
+# Code too big error
+--reset --mb=1 --dir=FWD_D,BWD_D --cfg=f32
+mb1_ic16oc16_iw80ow80kw31sw1dw0pw15
+