backport: cpu: aarch64: fix potential overflow in ACL strides

Co-authored-by: Siddhartha Menon <[email protected]>
Signed-off-by: Fadi Arafeh <[email protected]>

Author: fadara01

src/cpu/aarch64/acl_convolution_utils.cpp

@@ -278,8 +278,8 @@ status_t acl_init_conf(acl_conv_conf_t &acp, memory_desc_t &src_md,
             return status::unimplemented;
     }
 
-    acl_utils::reorder_to_weight_format(acp.wei_tensor_info, weights_md,
-            expected_weight_format, I_dim, O_dim, {W_dim, H_dim}, {});
+    CHECK(acl_utils::reorder_to_weight_format(acp.wei_tensor_info, weights_md,
+            expected_weight_format, I_dim, O_dim, {W_dim, H_dim}, {}));
 
     return status::success;
 }

src/cpu/aarch64/acl_inner_product.cpp

@@ -264,8 +264,8 @@ status_t acl_inner_product_fwd_t::pd_t::init_conf_ip(
     }
 
     const memory_desc_t weights_md_received = weights_md_;
-    acl_utils::reorder_to_weight_format(aip_.wei_tensor_info, weights_md_,
-            expected_weight_format, inner_dim, o_dim, remaining_dims, {});
+    CHECK(acl_utils::reorder_to_weight_format(aip_.wei_tensor_info, weights_md_,
+            expected_weight_format, inner_dim, o_dim, remaining_dims, {}));
 
     ACL_CHECK_SUPPORT((weights_format_kind_received == format_kind::blocked)
                     && !(dnnl_memory_desc_equal(

src/cpu/aarch64/acl_utils.cpp

@@ -15,6 +15,7 @@
 *******************************************************************************/
 
 #include "cpu/aarch64/acl_utils.hpp"
+#include <limits>
 
 namespace dnnl {
 namespace impl {
@@ -26,6 +27,17 @@ namespace acl_utils {
 using namespace dnnl::impl::alg_kind;
 using namespace data_type;
 
+status_t safe_set_strides(arm_compute::Strides &strides, size_t dim, size_t val,
+        bool inc_dim = true) {
+    // ACL stride value is uint32, check for overflow
+    if (val > std::numeric_limits<uint32_t>::max()) {
+        return status::unimplemented;
+    }
+
+    strides.set(dim, val, inc_dim);
+    return status::success;
+}
+
 arm_compute::DataType get_acl_data_t(
         const dnnl_data_type_t dt, const bool is_quantized) {
     switch (dt) {
@@ -154,8 +166,8 @@ status_t tensor_info(
     for (int i = md.ndims() - 1; i >= 0; --i) {
         // ACL strides are in bytes, oneDNN strides are in numbers of elements,
         // multiply by data type size to convert
-        strides_in_bytes.set(
-                acl_stride_i, blocking_desc.strides[i] * md.data_type_size());
+        CHECK(safe_set_strides(strides_in_bytes, acl_stride_i,
+                blocking_desc.strides[i] * md.data_type_size()));
         ++acl_stride_i;
     }
 
@@ -190,10 +202,12 @@ status_t insert_singleton_dimension(arm_compute::TensorInfo &ti, size_t dim_i) {
     arm_compute::Strides strides;
     for (size_t old_i = 0, new_i = 0; old_i < ti.num_dimensions(); ++old_i) {
         if (old_i == dim_i) {
-            strides.set(new_i, ti.strides_in_bytes()[old_i], false);
+            CHECK(safe_set_strides(
+                    strides, new_i, ti.strides_in_bytes()[old_i], false));
             ++new_i;
         }
-        strides.set(new_i, ti.strides_in_bytes()[old_i], false);
+        CHECK(safe_set_strides(
+                strides, new_i, ti.strides_in_bytes()[old_i], false));
         ++new_i;
     }
 
@@ -274,9 +288,9 @@ int reorder_dimensions_by_stride(std::vector<memory_desc_t *> permuted_mds,
     return reordered_dims;
 }
 
-void reorder_to_weight_format(arm_compute::TensorInfo &info, memory_desc_t &md,
-        arm_compute::WeightFormat wf, dim_t I_dim, dim_t O_dim,
-        const std::vector<dim_t> &spatial_dims,
+status_t reorder_to_weight_format(arm_compute::TensorInfo &info,
+        memory_desc_t &md, arm_compute::WeightFormat wf, dim_t I_dim,
+        dim_t O_dim, const std::vector<dim_t> &spatial_dims,
         const std::vector<dim_t> &batch_dims) {
 
     md.format_kind = format_kind::blocked;
@@ -336,12 +350,14 @@ void reorder_to_weight_format(arm_compute::TensorInfo &info, memory_desc_t &md,
     //   the interleaving). Note that we use the innermost_batch_stride
     //   because all the batched dimensions are collapsed (as required by ACL).
     arm_compute::Strides new_strides_in_bytes = info.strides_in_bytes();
-    new_strides_in_bytes.set(1, ldb * info.element_size());
-    new_strides_in_bytes.set(2, innermost_batch_stride * info.element_size());
+    CHECK(safe_set_strides(new_strides_in_bytes, 1, ldb * info.element_size()));
+    CHECK(safe_set_strides(new_strides_in_bytes, 2,
+            innermost_batch_stride * info.element_size()));
 
     info.init(info.tensor_shape(), info.num_channels(), info.data_type(),
             new_strides_in_bytes, info.offset_first_element_in_bytes(),
             memory_desc_wrapper(md).size());
+    return status::success;
 }
 
 } // namespace acl_utils

src/cpu/aarch64/acl_utils.hpp

@@ -92,9 +92,9 @@ int reorder_dimensions_by_stride(std::vector<memory_desc_t *> permuted_mds,
 //            matmul, ordered from innermost to outermost. ACL calls these
 //            the multi_stride_b. These will become the outermost (least dense)
 //            dimensions and will be collapsed.
-void reorder_to_weight_format(arm_compute::TensorInfo &info, memory_desc_t &md,
-        arm_compute::WeightFormat wf, dim_t I_dim, dim_t O_dim,
-        const std::vector<dim_t> &spatial_dims,
+status_t reorder_to_weight_format(arm_compute::TensorInfo &info,
+        memory_desc_t &md, arm_compute::WeightFormat wf, dim_t I_dim,
+        dim_t O_dim, const std::vector<dim_t> &spatial_dims,
         const std::vector<dim_t> &batch_dims = {});
 
 // Logs a custom 'info' line describing an unsupported case

src/cpu/aarch64/matmul/acl_matmul_utils.cpp

@@ -215,8 +215,8 @@ status_t init_conf_matmul(acl_matmul_conf_t &amp, memory_desc_t &src_md,
         for (dim_t i = K_dim - 1; i >= 0; --i)
             batch_dims.push_back(i);
 
-        acl_utils::reorder_to_weight_format(amp.wei_tensor_info, wei_md,
-                expected_weight_format, K_dim, N_dim, {}, batch_dims);
+        CHECK(acl_utils::reorder_to_weight_format(amp.wei_tensor_info, wei_md,
+                expected_weight_format, K_dim, N_dim, {}, batch_dims));
     }
 
     return status::success;