COMPMID-3887: NEGEMMConvolutionLayer hangs up on num_threads>18

When store the Tensor shape information, the size was clamped to 32bit, which cause the memory allocated is not big enough.

Signed-off-by: Sheri Zhang <sheri.zhang@arm.com>
Change-Id: I9f7dfcd5595a143b0ed4f6973e20bcd9d776b673
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/4331
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Georgios Pinitas <georgios.pinitas@arm.com>

3 files changed, 23 insertions, 9 deletions

diff --git a/arm_compute/core/TensorShape.h b/arm_compute/core/TensorShape.h
index 218774360e..b455a07767 100644
--- a/arm_compute/core/TensorShape.h
+++ b/arm_compute/core/TensorShape.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2019 Arm Limited.
+ * Copyright (c) 2016-2020 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -36,7 +36,7 @@
 namespace arm_compute
 {
 /** Shape of a tensor */
-class TensorShape : public Dimensions<uint32_t>
+class TensorShape : public Dimensions<size_t>
 {
 public:
     /** Constructor to initialize the tensor shape.
diff --git a/arm_compute/core/Window.inl b/arm_compute/core/Window.inl
index 14a432a0c0..6100d09a1c 100644
--- a/arm_compute/core/Window.inl
+++ b/arm_compute/core/Window.inl
@@ -197,15 +197,15 @@ inline Window Window::split_window(size_t dimension, size_t id, size_t total) co
     {
         if(d == dimension)
         {
-            int start        = _dims[d].start();
-            int end          = _dims[d].end();
-            const int step   = _dims[d].step();
+            int       start = _dims[d].start();
+            int       end   = _dims[d].end();
+            const int step  = _dims[d].step();
 
             const int num_it = num_iterations(d);
             const int rem    = num_it % total;
-            int work         = num_it / total;
+            int       work   = num_it / total;
 
-            int it_start     = work * id;
+            int it_start = work * id;
 
             if(int(id) < rem)
             {
@@ -277,7 +277,7 @@ inline void Window::use_tensor_dimensions(const TensorShape &shape, size_t first
 {
     for(unsigned int n = first_dimension; n < shape.num_dimensions(); ++n)
     {
-        set(n, Window::Dimension(0, std::max(shape[n], static_cast<uint32_t>(1))));
+        set(n, Window::Dimension(0, std::max(shape[n], static_cast<size_t>(1))));
     }
 }
 
diff --git a/utils/GraphUtils.cpp b/utils/GraphUtils.cpp
index 84f04161a8..e543cabea9 100644
--- a/utils/GraphUtils.cpp
+++ b/utils/GraphUtils.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2019 Arm Limited.
+ * Copyright (c) 2017-2020 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -272,9 +272,16 @@ bool ImageAccessor::access_tensor(ITensor &tensor)
         {
             std::tie(permuted_shape, perm) = compute_permutation_parameters(tensor.info()->tensor_shape(), tensor.info()->data_layout());
         }
+
+#ifdef __arm__
         ARM_COMPUTE_EXIT_ON_MSG_VAR(image_loader->width() != permuted_shape.x() || image_loader->height() != permuted_shape.y(),
                                     "Failed to load image file: dimensions [%d,%d] not correct, expected [%" PRIu32 ",%" PRIu32 "].",
                                     image_loader->width(), image_loader->height(), permuted_shape.x(), permuted_shape.y());
+#else  // __arm__
+        ARM_COMPUTE_EXIT_ON_MSG_VAR(image_loader->width() != permuted_shape.x() || image_loader->height() != permuted_shape.y(),
+                                    "Failed to load image file: dimensions [%d,%d] not correct, expected [%" PRIu64 ",%" PRIu64 "].",
+                                    image_loader->width(), image_loader->height(), permuted_shape.x(), permuted_shape.y());
+#endif // __arm__
 
         // Fill the tensor with the PPM content (BGR)
         image_loader->fill_planar_tensor(tensor, _bgr);
@@ -348,9 +355,16 @@ bool ValidationInputAccessor::access_tensor(arm_compute::ITensor &tensor)
             std::tie(permuted_shape, perm) = compute_permutation_parameters(tensor.info()->tensor_shape(),
                                                                             tensor.info()->data_layout());
         }
+
+#ifdef __arm__
         ARM_COMPUTE_EXIT_ON_MSG_VAR(jpeg.width() != permuted_shape.x() || jpeg.height() != permuted_shape.y(),
                                     "Failed to load image file: dimensions [%d,%d] not correct, expected [%" PRIu32 ",%" PRIu32 "].",
                                     jpeg.width(), jpeg.height(), permuted_shape.x(), permuted_shape.y());
+#else  // __arm__
+        ARM_COMPUTE_EXIT_ON_MSG_VAR(jpeg.width() != permuted_shape.x() || jpeg.height() != permuted_shape.y(),
+                                    "Failed to load image file: dimensions [%d,%d] not correct, expected [%" PRIu64 ",%" PRIu64 "].",
+                                    jpeg.width(), jpeg.height(), permuted_shape.x(), permuted_shape.y());
+#endif // __arm__
 
         // Fill the tensor with the JPEG content (BGR)
         jpeg.fill_planar_tensor(tensor, _bgr);