Add strides to ManagedTensor

Differential Revision: D61509079

Pull Request resolved: #4786

Author: jackzhxng

extension/runner_util/managed_tensor.h

@@ -48,13 +48,23 @@ class ManagedTensor {
 #ifdef USE_ATEN_LIB
     tensor_ = torch::from_blob(data, sizes, dtype);
 #else
+    // Calculate strides.
+    strides_ = std::vector<StridesType>(sizes_.size());
+    if (sizes_.size() > 0) {
+      strides_.back() = 1;
+      for (size_t i = strides_.size() - 1; i > 0; --i) {
+        strides_[i - 1] = strides_[i] * sizes_[i];
+      }
+    }
+
+    // Allocate TensorImpl.
     tensor_impl_ = std::make_unique<TensorImpl>(
         dtype,
         sizes_.size(),
         sizes_.data(),
         data,
-        nullptr,
-        nullptr,
+        /*dim_order=*/nullptr,
+        strides_.data(),
         TensorShapeDynamism::DYNAMIC_BOUND);
 #endif
   }
@@ -80,6 +90,7 @@ class ManagedTensor {
  private:
   std::unique_ptr<TensorImpl> tensor_impl_;
   std::vector<SizesType> sizes_;
+  std::vector<StridesType> strides_;
 #ifdef USE_ATEN_LIB
   Tensor tensor_;
 #endif

extension/runner_util/test/managed_tensor_test.cpp

@@ -25,15 +25,17 @@ class ManagedTensorTest : public ::testing::Test {
   void SetUp() override {
     torch::executor::runtime_init();
 
-    data_ = {1, 2, 3, 4, 5, 6};
-    sizes_ = {2, 3};
+    data_ = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
+    sizes_ = {2, 3, 4};
+    expected_strides_ = {12, 4, 1};
     managed_tensor_ =
         std::make_unique<ManagedTensor>(data_.data(), sizes_, ScalarType::Long);
   }
 
  protected:
   std::vector<int64_t> data_;
   std::vector<SizesType> sizes_;
+  std::vector<int> expected_strides_;
   std::unique_ptr<ManagedTensor> managed_tensor_;
 };
 
@@ -43,24 +45,27 @@ TEST_F(ManagedTensorTest, Smoke) {
   EXPECT_EQ(tensor.sizes(), ArrayRef<SizesType>(sizes_.data(), sizes_.size()));
   EXPECT_EQ(tensor.scalar_type(), ScalarType::Long);
   EXPECT_EQ(tensor.const_data_ptr(), data_.data());
+  for (size_t i = 0; i < expected_strides_.size(); ++i) {
+    EXPECT_EQ(tensor.strides()[i], expected_strides_[i]);
+  }
 }
 
 TEST_F(ManagedTensorTest, ResizeWithUpdatedRank) {
   // gtest death test doesn't work on iOS:
   // https://github.com/google/googletest/issues/2834
 #if !GTEST_OS_IOS
   EXPECT_EXIT(
-      managed_tensor_->resize(std::vector<SizesType>{2, 3, 4}),
+      managed_tensor_->resize(std::vector<SizesType>{2, 3, 4, 5}),
       ::testing::KilledBySignal(SIGABRT),
       "");
 #endif
 }
 
 TEST_F(ManagedTensorTest, ResizeShrink) {
-  managed_tensor_->resize(std::vector<SizesType>{2, 2});
+  managed_tensor_->resize(std::vector<SizesType>{2, 2, 2});
   const auto tensor = managed_tensor_->get_aliasing_tensor();
 
-  std::vector<SizesType> expected_sizes = {2, 2};
+  std::vector<SizesType> expected_sizes = {2, 2, 2};
   EXPECT_EQ(
       tensor.sizes(),
       ArrayRef<SizesType>(expected_sizes.data(), expected_sizes.size()));
@@ -69,10 +74,10 @@ TEST_F(ManagedTensorTest, ResizeShrink) {
 }
 
 TEST_F(ManagedTensorTest, Resize) {
-  managed_tensor_->resize(std::vector<SizesType>{3, 2});
+  managed_tensor_->resize(std::vector<SizesType>{4, 3, 2});
   const auto tensor = managed_tensor_->get_aliasing_tensor();
 
-  std::vector<SizesType> expected_sizes = {3, 2};
+  std::vector<SizesType> expected_sizes = {4, 3, 2};
   EXPECT_EQ(
       tensor.sizes(),
       ArrayRef<SizesType>(expected_sizes.data(), expected_sizes.size()));