Forbid having TensorImpl with zero number of elements and null data.

Differential Revision: D61810277

Pull Request resolved: #4909

Author: shoumikhin

runtime/core/exec_aten/testing_util/targets.bzl

@@ -23,6 +23,7 @@ def define_common_targets():
                 # list.
                 "//executorch/runtime/core/exec_aten/util/test/...",
                 "//executorch/runtime/core/exec_aten/testing_util/test/...",
+                "//executorch/runtime/core/portable_type/test/...",
                 "//executorch/kernels/prim_ops/test/...",
                 "//executorch/kernels/portable/test/...",
                 "//executorch/kernels/portable/cpu/util/test/...",

runtime/core/portable_type/tensor_impl.cpp

@@ -25,8 +25,16 @@ namespace {
  * Compute the number of elements based on the sizes of a tensor.
  */
 ssize_t compute_numel(const TensorImpl::SizesType* sizes, ssize_t dim) {
+  ET_CHECK_MSG(
+      dim == 0 || sizes != nullptr,
+      "Sizes must be provided for non-scalar tensors");
   ssize_t numel = 1; // Zero-dimensional tensors (scalars) have numel == 1.
   for (ssize_t i = 0; i < dim; ++i) {
+    ET_CHECK_MSG(
+        sizes[i] >= 0,
+        "Size must be non-negative, got %d at dimension %zd",
+        sizes[i],
+        i);
     numel *= sizes[i];
   }
   return numel;
@@ -52,6 +60,7 @@ TensorImpl::TensorImpl(
       shape_dynamism_(dynamism) {
   ET_CHECK_MSG(
       isValid(type_), "Invalid type %" PRId8, static_cast<int8_t>(type_));
+  ET_CHECK_MSG(dim_ >= 0, "Dimension must be non-negative, got %zd", dim_);
 }
 
 size_t TensorImpl::nbytes() const {

runtime/core/portable_type/test/tensor_impl_test.cpp

@@ -8,12 +8,13 @@
 
 #include <executorch/runtime/core/portable_type/tensor_impl.h>
 
-#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
-#include <executorch/runtime/platform/runtime.h>
-
 #include <gtest/gtest.h>
 #include <random>
 
+#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
+#include <executorch/runtime/platform/runtime.h>
+#include <executorch/test/utils/DeathTest.h>
+
 using namespace ::testing;
 
 namespace torch {
@@ -29,7 +30,7 @@ class TensorImplTest : public ::testing::Test {
   void SetUp() override {
     // Since these tests cause ET_LOG to be called, the PAL must be initialized
     // first.
-    torch::executor::runtime_init();
+    runtime_init();
   }
 };
 
@@ -370,5 +371,81 @@ TEST_F(TensorImplTest, TestWriteRead) {
   EXPECT_EQ(y[0], 22.0);
 }
 
+TEST_F(TensorImplTest, TestInvalidScalarType) {
+  SizesType sizes[2] = {3, 2};
+  ET_EXPECT_DEATH(TensorImpl t(static_cast<ScalarType>(-1), 2, sizes), "");
+}
+
+TEST_F(TensorImplTest, TestNegativeDimension) {
+  SizesType sizes[2] = {3, 2};
+  ET_EXPECT_DEATH(TensorImpl t(ScalarType::Float, -1, sizes), "");
+}
+
+TEST_F(TensorImplTest, TestNullSizesNonZeroDim) {
+  ET_EXPECT_DEATH(TensorImpl t(ScalarType::Float, 2, nullptr), "");
+}
+
+TEST_F(TensorImplTest, TestNonNegativeSizes) {
+  SizesType sizes[2] = {3, -2};
+  ET_EXPECT_DEATH(TensorImpl t(ScalarType::Float, 2, sizes), "");
+}
+
+TEST_F(TensorImplTest, TestEmptyTensor) {
+  SizesType sizes[2] = {0, 0};
+  TensorImpl t(ScalarType::Float, 2, sizes);
+  EXPECT_EQ(t.numel(), 0);
+  EXPECT_EQ(t.data(), nullptr);
+}
+
+TEST_F(TensorImplTest, TestTensorWithNoElementsButAllocatedMemory) {
+  SizesType sizes[2] = {0, 0};
+  float data[1] = {1.0};
+  TensorImpl t(ScalarType::Float, 2, sizes, data);
+  EXPECT_EQ(t.numel(), 0);
+  EXPECT_EQ(t.data(), data);
+}
+
+TEST_F(TensorImplTest, TestTensorWithShapeButNoMemory) {
+  SizesType sizes[2] = {3, 2};
+  TensorImpl t(ScalarType::Float, 2, sizes);
+  EXPECT_GT(t.numel(), 0);
+  EXPECT_EQ(t.data(), nullptr);
+}
+
+TEST_F(TensorImplTest, TestNormalTensor) {
+  SizesType sizes[2] = {3, 2};
+  float data[6] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0};
+  TensorImpl t(ScalarType::Float, 2, sizes, data);
+  EXPECT_GT(t.numel(), 0);
+  EXPECT_EQ(t.data(), data);
+}
+
+TEST_F(TensorImplTest, TestResizingTensorToZeroAndBack) {
+  SizesType sizes[2] = {3, 2};
+  TensorImpl t(
+      ScalarType::Float,
+      2,
+      sizes,
+      nullptr,
+      nullptr,
+      nullptr,
+      TensorShapeDynamism::DYNAMIC_BOUND);
+
+  float data[6] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0};
+  t.set_data(data);
+  EXPECT_GT(t.numel(), 0);
+  EXPECT_EQ(t.data(), data);
+
+  SizesType zero_sizes[2] = {0, 0};
+  t.set_sizes_contiguous({zero_sizes, 2});
+  EXPECT_EQ(t.numel(), 0);
+  EXPECT_EQ(t.data(), data);
+
+  SizesType new_sizes[2] = {3, 2};
+  t.set_sizes_contiguous({new_sizes, 2});
+  EXPECT_GT(t.numel(), 0);
+  EXPECT_EQ(t.data(), data);
+}
+
 } // namespace executor
 } // namespace torch