Add extra checks for the provided dim order and strides. (#5390)

Add extra checks for the provided dim order and strides. (#5377)

Summary:
Pull Request resolved: #5377

Also make sure the scalar tensors stay static.

Reviewed By: kirklandsign

Differential Revision: D62681004

fbshipit-source-id: 57e035bcc5d9a4ffb2571144d7b3ed72ff84b193
(cherry picked from commit 0a501eb)

Co-authored-by: Anthony Shoumikhin <[email protected]>

Author: pytorchbot

extension/tensor/tensor_impl_ptr.cpp

@@ -62,6 +62,13 @@ TensorImplPtr make_tensor_impl_ptr(
     exec_aten::TensorShapeDynamism dynamism,
     std::function<void(void*)> deleter) {
   const auto dim = sizes.size();
+  ET_CHECK_MSG(
+      dim_order.empty() || dim_order.size() == dim,
+      "dim_order size must match sizes or be empty.");
+  ET_CHECK_MSG(
+      strides.empty() || strides.size() == dim,
+      "strides size must match sizes or be empty.");
+
   if (dim_order.empty()) {
     dim_order.resize(dim);
     std::iota(dim_order.begin(), dim_order.end(), 0);
@@ -89,7 +96,7 @@ TensorImplPtr make_tensor_impl_ptr(
       data,
       dim_order.data(),
       strides.data(),
-      dynamism);
+      dim > 0 ? dynamism : exec_aten::TensorShapeDynamism::STATIC);
   return TensorImplPtr(
       tensor_impl.release(),
       TensorImplPtrDeleter{

extension/tensor/tensor_impl_ptr.h

@@ -87,7 +87,7 @@ TensorImplPtr make_tensor_impl_ptr(
  * @return A TensorImplPtr that manages the newly created TensorImpl.
  */
 template <typename T = float>
-TensorImplPtr make_tensor_impl_ptr(
+inline TensorImplPtr make_tensor_impl_ptr(
     std::vector<exec_aten::SizesType> sizes,
     std::vector<T> data,
     std::vector<exec_aten::DimOrderType> dim_order = {},
@@ -123,23 +123,13 @@ TensorImplPtr make_tensor_impl_ptr(
  * @return A TensorImplPtr that manages the newly created TensorImpl.
  */
 template <typename T = float>
-TensorImplPtr make_tensor_impl_ptr(
+inline TensorImplPtr make_tensor_impl_ptr(
     std::vector<T> data,
     exec_aten::TensorShapeDynamism dynamism =
         exec_aten::TensorShapeDynamism::DYNAMIC_BOUND) {
-  constexpr exec_aten::ScalarType scalar_type =
-      runtime::CppTypeToScalarType<T>::value;
   std::vector<exec_aten::SizesType> sizes{exec_aten::SizesType(data.size())};
-  const auto raw_data_ptr = data.data();
-  auto data_ptr = std::make_shared<std::vector<T>>(std::move(data));
   return make_tensor_impl_ptr(
-      scalar_type,
-      std::move(sizes),
-      raw_data_ptr,
-      {0},
-      {1},
-      dynamism,
-      [data_ptr = std::move(data_ptr)](void*) {});
+      std::move(sizes), std::move(data), {0}, {1}, dynamism);
 }
 
 /**

extension/tensor/tensor_ptr.h

@@ -235,7 +235,19 @@ TensorPtr make_tensor_ptr(
     std::initializer_list<T> data,
     exec_aten::TensorShapeDynamism dynamism =
         exec_aten::TensorShapeDynamism::DYNAMIC_BOUND) {
-  return make_tensor_ptr(std::vector<T>(data), dynamism);
+  return make_tensor_ptr(make_tensor_impl_ptr(std::vector<T>(data), dynamism));
+}
+
+/**
+ * Creates a TensorPtr that manages a Tensor with a single scalar value.
+ *
+ * @tparam T The C++ type of the scalar value.
+ * @param value The scalar value to be used for the Tensor.
+ * @return A TensorPtr that manages the newly created TensorImpl.
+ */
+template <typename T>
+TensorPtr make_tensor_ptr(T value) {
+  return make_tensor_ptr(make_tensor_impl_ptr({}, std::vector<T>{value}));
 }
 
 /**

extension/tensor/tensor_ptr_maker.h

@@ -421,10 +421,8 @@ inline TensorPtr full(
  */
 inline TensorPtr scalar_tensor(
     exec_aten::Scalar value,
-    exec_aten::ScalarType type = exec_aten::ScalarType::Float,
-    exec_aten::TensorShapeDynamism dynamism =
-        exec_aten::TensorShapeDynamism::DYNAMIC_BOUND) {
-  return full({}, value, type, dynamism);
+    exec_aten::ScalarType type = exec_aten::ScalarType::Float) {
+  return full({}, value, type);
 }
 
 /**

extension/tensor/test/tensor_impl_ptr_test.cpp

@@ -377,3 +377,19 @@ TEST_F(TensorImplPtrTest, TensorImplUint8BufferTooLarge) {
   EXPECT_EQ(tensor_impl->strides()[0], 2);
   EXPECT_EQ(tensor_impl->strides()[1], 1);
 }
+
+TEST_F(TensorImplPtrTest, StridesAndDimOrderMustMatchSizes) {
+  float data[12] = {0};
+  ET_EXPECT_DEATH(
+      {
+        auto _ = make_tensor_impl_ptr(
+            exec_aten::ScalarType::Float, {3, 4}, data, {}, {1});
+      },
+      "");
+  ET_EXPECT_DEATH(
+      {
+        auto _ = make_tensor_impl_ptr(
+            exec_aten::ScalarType::Float, {3, 4}, data, {0}, {4, 1});
+      },
+      "");
+}

extension/tensor/test/tensor_ptr_test.cpp

@@ -44,6 +44,40 @@ TEST_F(TensorPtrTest, ScalarTensorOwningData) {
   EXPECT_EQ(tensor->const_data_ptr<float>()[0], 3.14f);
 }
 
+TEST_F(TensorPtrTest, ScalarTensorSingleValueCreation) {
+  auto tensor_float = make_tensor_ptr(3.14f);
+  EXPECT_EQ(tensor_float->dim(), 0);
+  EXPECT_EQ(tensor_float->numel(), 1);
+  EXPECT_EQ(tensor_float->sizes().size(), 0);
+  EXPECT_EQ(tensor_float->strides().size(), 0);
+  EXPECT_EQ(tensor_float->const_data_ptr<float>()[0], 3.14f);
+  EXPECT_EQ(tensor_float->scalar_type(), exec_aten::ScalarType::Float);
+
+  auto tensor_int32 = make_tensor_ptr(42);
+  EXPECT_EQ(tensor_int32->dim(), 0);
+  EXPECT_EQ(tensor_int32->numel(), 1);
+  EXPECT_EQ(tensor_int32->sizes().size(), 0);
+  EXPECT_EQ(tensor_int32->strides().size(), 0);
+  EXPECT_EQ(tensor_int32->const_data_ptr<int32_t>()[0], 42);
+  EXPECT_EQ(tensor_int32->scalar_type(), exec_aten::ScalarType::Int);
+
+  auto tensor_double = make_tensor_ptr(2.718);
+  EXPECT_EQ(tensor_double->dim(), 0);
+  EXPECT_EQ(tensor_double->numel(), 1);
+  EXPECT_EQ(tensor_double->sizes().size(), 0);
+  EXPECT_EQ(tensor_double->strides().size(), 0);
+  EXPECT_EQ(tensor_double->const_data_ptr<double>()[0], 2.718);
+  EXPECT_EQ(tensor_double->scalar_type(), exec_aten::ScalarType::Double);
+
+  auto tensor_int64 = make_tensor_ptr(static_cast<int64_t>(10000000000));
+  EXPECT_EQ(tensor_int64->dim(), 0);
+  EXPECT_EQ(tensor_int64->numel(), 1);
+  EXPECT_EQ(tensor_int64->sizes().size(), 0);
+  EXPECT_EQ(tensor_int64->strides().size(), 0);
+  EXPECT_EQ(tensor_int64->const_data_ptr<int64_t>()[0], 10000000000);
+  EXPECT_EQ(tensor_int64->scalar_type(), exec_aten::ScalarType::Long);
+}
+
 TEST_F(TensorPtrTest, CreateTensorWithStridesAndDimOrder) {
   float data[20] = {2};
   auto tensor = make_tensor_ptr(
@@ -299,12 +333,7 @@ TEST_F(TensorPtrTest, TensorSharingImplModifiesSharedDataVector) {
 TEST_F(TensorPtrTest, TensorSharingImplResizingAffectsBothVector) {
   std::vector<float> data = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
 
-  auto tensor1 = make_tensor_ptr(
-      {3, 4},
-      std::move(data),
-      {},
-      {},
-      exec_aten::TensorShapeDynamism::DYNAMIC_UNBOUND);
+  auto tensor1 = make_tensor_ptr({3, 4}, std::move(data));
   auto tensor2 = make_tensor_ptr(tensor1);
 
   EXPECT_EQ(resize_tensor_ptr(tensor1, {2, 6}), Error::Ok);