From 97f1c0e602049ee43537ad6c8a08f476eb5c722b Mon Sep 17 00:00:00 2001
From: TatWai Chong <tatwai.chong@arm.com>
Date: Mon, 5 Feb 2024 11:56:46 -0800
Subject: Fix comparing unsigned int and int in slice evaluation

also check if tensors of start and size are allocated.

Change-Id: I6a72d11ebcb8d0725fe267058dfd792102459427
Signed-off-by: TatWai Chong <tatwai.chong@arm.com>
---
 reference_model/src/ops/data_layout.cc | 22 +++++++++++++++-------
 reference_model/src/ops/data_layout.h  |  6 +++---
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/reference_model/src/ops/data_layout.cc b/reference_model/src/ops/data_layout.cc
index ddf0713..3bff047 100644
--- a/reference_model/src/ops/data_layout.cc
+++ b/reference_model/src/ops/data_layout.cc
@@ -453,11 +453,11 @@ int OpSlice<Rank, Dtype>::checkTensorAttributes()
     }
 
     in    = dynamic_cast<TosaReference::TensorTemplate<TIn>*>(inputs[0]);
-    start = dynamic_cast<TosaReference::TensorTemplate<TInShape>*>(inputs[1]);
-    size  = dynamic_cast<TosaReference::TensorTemplate<TInShape>*>(inputs[2]);
+    start = dynamic_cast<TosaReference::TensorTemplate<TSlicing>*>(inputs[1]);
+    size  = dynamic_cast<TosaReference::TensorTemplate<TSlicing>*>(inputs[2]);
     out   = dynamic_cast<TosaReference::TensorTemplate<TOut>*>(outputs[0]);
 
-    ASSERT_MEM(in && out);
+    ASSERT_MEM(in && out && start && size);
 
     return 0;
 }
@@ -465,13 +465,21 @@ int OpSlice<Rank, Dtype>::checkTensorAttributes()
 template <int Rank, TOSA_REF_TYPE Dtype>
 int OpSlice<Rank, Dtype>::eval()
 {
-    ERROR_IF(start->getElementCount() != in->getRank(), "OpSlice: start array length needs to be rank(input)");
-    ERROR_IF(size->getElementCount() != in->getRank(), "OpSlice: size array length needs to be rank(input)");
+    TSlicing start_tensor = start->getTensor();
+    TSlicing size_tensor  = size->getTensor();
+
+    // According to https://eigen.tuxfamily.org/dox/unsupported/eigen_tensors.html
+    // The type of size() is <Tensor-Type>::Index, but can always handily use it like an int.
+    // However, apply explicit cast to int32_t is preferred.
+    ERROR_IF(static_cast<int32_t>(start_tensor.size()) != in->getRank(),
+             "OpSlice: start array length needs to be rank(input)");
+    ERROR_IF(static_cast<int32_t>(size_tensor.size()) != in->getRank(),
+             "OpSlice: size array length needs to be rank(input)");
 
     for (int32_t i = 0; i < in->getRank(); i++)
     {
-        int32_t b = start->getTensor()(i);
-        int32_t s = size->getTensor()(i);
+        int32_t b = start_tensor(i);
+        int32_t s = size_tensor(i);
         ERROR_IF(b < 0 || b >= in->getShape()[i], "OpSlice: start out of boundary");
         ERROR_IF((b + s) < 0 || (b + s) > in->getShape()[i], "OpSlice: (start+size) out of boundary");
         ERROR_IF(s <= 0, "OpSlice: output must be positive");
diff --git a/reference_model/src/ops/data_layout.h b/reference_model/src/ops/data_layout.h
index 6ab5ebd..dee2ae0 100644
--- a/reference_model/src/ops/data_layout.h
+++ b/reference_model/src/ops/data_layout.h
@@ -149,14 +149,14 @@ public:
     using InEigenShapeType = typename GetEigenType<TOSA_REF_TYPE_SHAPE>::type;
     using OutEigenType     = typename GetEigenType<Dtype>::type;
     using TIn              = Eigen::Tensor<InEigenType, Rank>;
-    using TInShape         = Eigen::Tensor<InEigenShapeType, 1>;
+    using TSlicing         = Eigen::Tensor<InEigenShapeType, 1>;
     using TOut             = Eigen::Tensor<OutEigenType, Rank>;
 
 protected:
     Eigen::array<Eigen::Index, Rank> begin_array;
     Eigen::array<Eigen::Index, Rank> size_array;
-    TosaReference::TensorTemplate<TInShape>* start;
-    TosaReference::TensorTemplate<TInShape>* size;
+    TosaReference::TensorTemplate<TSlicing>* start;
+    TosaReference::TensorTemplate<TSlicing>* size;
     TosaReference::TensorTemplate<TIn>* in;
     TosaReference::TensorTemplate<TOut>* out;
 };
-- 
cgit v1.2.1