x64: conv: avoid overflows and add limit for huge spatial sizes

tczeszun · densamoilov · commit a68b8ae9b40a · 2025-02-07T17:43:56.000-08:00
diff --git a/src/cpu/x64/jit_avx2_conv_kernel_f32.cpp b/src/cpu/x64/jit_avx2_conv_kernel_f32.cpp
@@ -1,5 +1,5 @@
 /*******************************************************************************
-* Copyright 2016-2024 Intel Corporation
+* Copyright 2016-2025 Intel Corporation
 * Copyright 2018 YANDEX LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
@@ -1459,6 +1459,26 @@ status_t jit_avx2_conv_bwd_weights_kernel_f32::init_conf(jit_conv_conf_t &jcp,
     jcp.nb_oc = div_up(jcp.oc, jcp.oc_block);
     jcp.nb_ic_blocking = jcp.nb_oc_blocking = 1;
 
+    jcp.typesize_in = types::data_type_size(src_d.data_type());
+    jcp.typesize_out = types::data_type_size(diff_dst_d.data_type());
+
+    const bool is_src_layout_blocked = jcp.src_tag == dat_tag_nCx8c;
+    const bool is_dst_layout_blocked = jcp.dst_tag == dat_tag_nCx8c;
+
+    dim_t src_size = static_cast<dim_t>(jcp.mb)
+            * (is_src_layout_blocked ? rnd_up(jcp.ic, jcp.ic_block) : jcp.ic)
+            * jcp.id * jcp.ih * jcp.iw * jcp.typesize_in;
+
+    VDISPATCH_CONV_IC(src_size <= INT_MAX, VERBOSE_UNSUPPORTED_FEATURE,
+            "src size > INT_MAX is not supported");
+
+    dim_t diff_dst_size = static_cast<dim_t>(jcp.mb)
+            * (is_dst_layout_blocked ? rnd_up(jcp.oc, jcp.oc_block) : jcp.oc)
+            * jcp.id * jcp.ih * jcp.iw * jcp.typesize_in;
+
+    VDISPATCH_CONV_IC(diff_dst_size <= INT_MAX, VERBOSE_UNSUPPORTED_FEATURE,
+            "diff_dst size > INT_MAX is not supported");
+
     return status::success;
 }
 
diff --git a/src/cpu/x64/jit_avx512_common_conv_kernel.cpp b/src/cpu/x64/jit_avx512_common_conv_kernel.cpp
@@ -4118,14 +4118,24 @@ status_t jit_avx512_common_conv_bwd_weights_kernel_f32::init_conf(
     jcp.typesize_in = typesize;
     jcp.typesize_out = typesize;
 
+    dim_t src_size = static_cast<dim_t>(jcp.mb)
+            * (is_data_layout_nxc ? jcp.ic : rnd_up(jcp.ic, jcp.ic_block))
+            * jcp.id * jcp.ih * jcp.iw * jcp.typesize_in;
+
+    VDISPATCH_CONV_IC(src_size <= INT_MAX, VERBOSE_UNSUPPORTED_FEATURE,
+            "src size > INT_MAX is not supported");
+
+    dim_t diff_dst_size = static_cast<dim_t>(jcp.mb)
+            * (is_data_layout_nxc ? jcp.oc : rnd_up(jcp.oc, jcp.oc_block))
+            * jcp.id * jcp.ih * jcp.iw * jcp.typesize_in;
+
+    VDISPATCH_CONV_IC(diff_dst_size <= INT_MAX, VERBOSE_UNSUPPORTED_FEATURE,
+            "diff_dst size > INT_MAX is not supported");
+
     bool use_nxc_harness = false;
     if (is_data_layout_nxc) {
         dim_t kernel_size = static_cast<dim_t>(jcp.ic) * jcp.oc * jcp.kd
                 * jcp.kh * jcp.kw * jcp.typesize_out;
-        dim_t src_size = static_cast<dim_t>(jcp.mb) * jcp.ic * jcp.id * jcp.ih
-                * jcp.iw * jcp.typesize_in;
-        dim_t diff_dst_size = static_cast<dim_t>(jcp.mb) * jcp.oc * jcp.id
-                * jcp.ih * jcp.iw * jcp.typesize_in;
         dim_t data_size = src_size + diff_dst_size;
 
         // The advantage of the nxc kernel is cache traversal, this comes at a
diff --git a/src/cpu/x64/jit_avx512_core_bf16_conv_kernel.cpp b/src/cpu/x64/jit_avx512_core_bf16_conv_kernel.cpp
@@ -4449,8 +4449,15 @@ status_t jit_avx512_core_bf16_conv_bwd_weights_kernel_f32::init_conf(
             && jcp.oc <= diff_weights_d.padded_dims()[with_groups + 0];
     if (!args_ok) return status::unimplemented;
 
-    int inp_row_size = jcp.ic_block * jcp.tr_iw * jcp.typesize_in;
-    int out_row_size = jcp.oc_block * jcp.tr_ow * jcp.typesize_in;
+    const auto inp_row_size
+            = static_cast<size_t>(jcp.ic_block) * jcp.tr_iw * jcp.typesize_in;
+    VDISPATCH_CONV_IC(inp_row_size <= INT_MAX, VERBOSE_UNSUPPORTED_FEATURE,
+            "inp_row_size > INT_MAX is not supported");
+    const auto out_row_size
+            = static_cast<size_t>(jcp.oc_block) * jcp.tr_ow * jcp.typesize_in;
+    VDISPATCH_CONV_IC(out_row_size <= INT_MAX, VERBOSE_UNSUPPORTED_FEATURE,
+            "out_row_size > INT_MAX is not supported");
+
     int full_spat_min_h_block_size
             = nstl::max(1, nstl::max(jcp.b_pad, jcp.t_pad));
     int full_spat_working_set_size
diff --git a/src/cpu/x64/jit_brgemm_conv_bwd_w.cpp b/src/cpu/x64/jit_brgemm_conv_bwd_w.cpp
@@ -139,9 +139,20 @@ status_t brgemm_convolution_bwd_weights_t::pd_t::init(engine_t *engine) {
                 brgattr.max_top_vpad = 0;
                 brgattr.max_bottom_vpad = 0;
 
-                brgattr.LDA2 = jcp_.tr_iw * jcp_.ih_block * jcp_.id;
-                brgattr.LDB2
-                        = jcp_.tr_ow * jcp_.oc_block * jcp_.oh_block * jcp_.od;
+                const auto lda2_size = static_cast<size_t>(jcp_.tr_iw)
+                        * jcp_.ih_block * jcp_.id;
+                VDISPATCH_CONV_IC(lda2_size <= INT_MAX,
+                        VERBOSE_UNSUPPORTED_FEATURE,
+                        "lda2_size > INT_MAX is not supported");
+                brgattr.LDA2 = lda2_size;
+
+                const auto ldb2_size = static_cast<size_t>(jcp_.tr_ow)
+                        * jcp_.oc_block * jcp_.oh_block * jcp_.od;
+                VDISPATCH_CONV_IC(ldb2_size <= INT_MAX,
+                        VERBOSE_UNSUPPORTED_FEATURE,
+                        "ldb2_size > INT_MAX is not supported");
+                brgattr.LDB2 = ldb2_size;
+
                 brgattr.LDC2_M = jcp_.oc_block * jcp_.kd * jcp_.kh * jcp_.kw;
                 brgattr.LDC2_N = jcp_.nb_ic * jcp_.ic_block * jcp_.oc_block
                         * jcp_.kd * jcp_.kh * jcp_.kw;
diff --git a/src/cpu/x64/jit_brgemm_conv_utils.cpp b/src/cpu/x64/jit_brgemm_conv_utils.cpp
@@ -3306,9 +3306,10 @@ status_t init_conf_bwd_w(jit_brgemm_conv_conf_t &jcp,
     jcp.tr_diff_dst_buf_count = jcp.global_transpose
             ? jcp.nthr_mb * jcp.nb_oc * jcp.ngroups
             : jcp.nthr;
-    jcp.tr_src_block_size = jcp.tr_iw * jcp.ic_block * jcp.ih_block * jcp.id;
-    jcp.tr_diff_dst_block_size
-            = jcp.tr_ow * jcp.oc_block * jcp.oh_block * jcp.od;
+    jcp.tr_src_block_size = static_cast<size_t>(jcp.tr_iw) * jcp.ic_block
+            * jcp.ih_block * jcp.id;
+    jcp.tr_diff_dst_block_size = static_cast<size_t>(jcp.tr_ow) * jcp.oc_block
+            * jcp.oh_block * jcp.od;
 
     jcp.tr_src_buf_size = jcp.tr_src_block_size
             * (jcp.global_transpose ? 1 : jcp.nb_ic_blocking);
@@ -3368,7 +3369,7 @@ status_t init_scratchpad_bwd_w(memory_tracking::registrar_t &scratchpad,
     // (jcp.tr_diff_dst_buf_size + jcp.tr_iw * jcp.oc_block)
     const auto tr_diff_dst_size
             = jcp.tr_diff_dst_buf_count * jcp.tr_diff_dst_buf_size
-            + jcp.tr_iw * jcp.oc_block;
+            + static_cast<size_t>(jcp.tr_iw) * jcp.oc_block;
 
     const size_t min_align = 64;
     scratchpad.book(
