x64: conv: relax large size check for 3d f32/int8 shapes

Author: tczeszun

src/cpu/x64/jit_brgemm_conv_utils.cpp

@@ -1679,9 +1679,10 @@ status_t init_jcp(jit_brgemm_conv_conf_t &jcp, cpu_isa_t isa,
     const memory_desc_wrapper dst_d(&dst_md);
     const memory_desc_wrapper bias_d(&bias_md);
 
-    // Big int (> INT_MAX) values are unsupported and jcp fields may overflow
+    // Per-tensor spatial product, weights nelems, and per-dim
+    // strides/paddings/dilations must fit in int.
     // TODO: change data type of jcp fields to size_t
-    VDISPATCH_CONV_IC(!has_large_size(cd, src_d, weights_d, dst_d),
+    VDISPATCH_CONV_IC(!has_large_size_relaxed(cd, src_d, weights_d, dst_d),
             VERBOSE_BAD_PARAM, "large size is not supported");
 
     const bool with_groups = weights_d.ndims() == src_d.ndims() + 1;
@@ -2658,7 +2659,7 @@ status_t init_1x1_conf(jit_brgemm_conv_conf_t &jcp, cpu_isa_t isa,
 
     const size_t rtus_buffer_size = jcp.is_reduced_rtus
             ? jcp.rtus_padded_ic_size * jcp.os_block
-            : jcp.LDA * jcp.os;
+            : static_cast<size_t>(jcp.LDA) * jcp.os;
     jcp.inp_buffer_size
             = jcp.is_rtus ? rnd_up(rtus_buffer_size, align_size) : 0;
 

src/cpu/x64/jit_primitive_conf.hpp

@@ -305,6 +305,36 @@ inline bool has_large_size(const convolution_desc_t &cd,
     return false;
 }
 
+// Relaxed version of has_large_size(). The only difference from
+// has_large_size() is that this variant checks the per-tensor spatial product
+// instead of the count of all consecutive elements excluding the minibatch
+// (i.e. spatial * channels). Weights nelems and per-dim strides/paddings/dilations
+// are checked the same way. Use this when element offsets into memory are computed
+// as dim_t and only per-dim / spatial values need to fit in int.
+inline bool has_large_size_relaxed(const convolution_desc_t &cd,
+        const memory_desc_wrapper &src_d, const memory_desc_wrapper &weights_d,
+        const memory_desc_wrapper &dst_d) {
+    auto is_large = [](const dim_t val) { return val > INT_MAX; };
+
+    const int ndims = src_d.ndims();
+    dim_t src_spatial = 1, dst_spatial = 1;
+    for (int d = 2; d < ndims; d++) {
+        src_spatial *= src_d.dims()[d];
+        dst_spatial *= dst_d.dims()[d];
+    }
+    if (is_large(src_spatial) || is_large(dst_spatial)) return true;
+
+    if (is_large(weights_d.nelems())) return true;
+
+    for (int d = 3; d <= ndims; d++) {
+        if (utils::one_of(true, is_large(cd.strides[ndims - d]),
+                    is_large(cd.padding[0][ndims - d]),
+                    is_large(cd.dilates[ndims - d])))
+            return true;
+    }
+    return false;
+}
+
 struct jit_conv_args_t {
     const void *src = nullptr; /* hack, non-const for backward_data */
     const void *dst = nullptr; /* hack, non-const for forward */