Update to check tensor size is valid

Validate tensor arguments

Signed-off-by: Eric Kunze <eric.kunze@arm.com>
Change-Id: Iac398f2c3e600944c6f65c3d8433bd17095f820b

1 files changed, 9 insertions, 0 deletions

diff --git a/reference_model/src/subgraph_traverser.cc b/reference_model/src/subgraph_traverser.cc
index 15d82e6..ae216d8 100644
--- a/reference_model/src/subgraph_traverser.cc
+++ b/reference_model/src/subgraph_traverser.cc
@@ -14,6 +14,7 @@
 //    limitations under the License.
 
 #include "subgraph_traverser.h"
+#include "tosa_model_types.h"
 
 #ifndef SUBGRAPH_ERROR_IF
 #define SUBGRAPH_ERROR_IF(COND, fmt, ...)                                                                              \
@@ -323,6 +324,7 @@ int SubgraphTraverser::allocateTensor()
         auto got = used_tensor_name_set.find(ts->GetName());
         if (got != used_tensor_name_set.end())
         {
+            uint32_t elements = 1;
             for (auto& dim : ts->GetShape())
             {
                 if (dim <= 0)
@@ -331,6 +333,13 @@ int SubgraphTraverser::allocateTensor()
                     this->setGraphStatus(GraphStatus::TOSA_UNPREDICTABLE);
                     return 1;
                 }
+                if (dim > static_cast<int32_t>(TOSA_MAX_TENSOR_SIZE / elements))
+                {
+                    // Size greather than maximum defined in spec
+                    DEBUG_INFO(GT, "Tensor %s size is greater than allowed maximum", ts->GetName().c_str());
+                    this->setGraphStatus(GraphStatus::TOSA_UNPREDICTABLE);
+                    return 1;
+                }
             }
         }