[FFI] Construct NDArray.strides by default (apache#18272)

This PR updates NDArray.strides to construct strides by default

Author: cyx-6

include/tvm/ffi/container/ndarray.h

@@ -151,6 +151,7 @@ class NDArrayObj : public Object, public DLTensor {
  protected:
   // backs up the shape of the NDArray
   Optional<Shape> shape_data_;
+  Optional<Shape> stride_data_;
 
   static void DLManagedTensorDeleter(DLManagedTensor* tensor) {
     NDArrayObj* obj = static_cast<NDArrayObj*>(tensor->manager_ctx);
@@ -184,9 +185,11 @@ class NDArrayObjFromNDAlloc : public NDArrayObj {
     this->ndim = static_cast<int>(shape.size());
     this->dtype = dtype;
     this->shape = const_cast<int64_t*>(shape.data());
-    this->strides = nullptr;
+    Shape strides = Shape(details::MakeStridesFromShape(this->ndim, this->shape));
+    this->strides = const_cast<int64_t*>(strides.data());
     this->byte_offset = 0;
     this->shape_data_ = std::move(shape);
+    this->stride_data_ = std::move(strides);
     alloc_.AllocData(static_cast<DLTensor*>(this), std::forward<ExtraArgs>(extra_args)...);
   }
 
@@ -202,9 +205,10 @@ class NDArrayObjFromDLPack : public NDArrayObj {
  public:
   explicit NDArrayObjFromDLPack(TDLPackManagedTensor* tensor) : tensor_(tensor) {
     *static_cast<DLTensor*>(this) = tensor_->dl_tensor;
-    // set strides to nullptr if the tensor is contiguous.
-    if (IsContiguous(tensor->dl_tensor)) {
-      this->strides = nullptr;
+    if (tensor_->dl_tensor.strides == nullptr) {
+      Shape strides = Shape(details::MakeStridesFromShape(ndim, shape));
+      this->strides = const_cast<int64_t*>(strides.data());
+      this->stride_data_ = std::move(strides);
     }
   }
 

include/tvm/ffi/container/shape.h

@@ -91,6 +91,17 @@ TVM_FFI_INLINE ObjectPtr<ShapeObj> MakeInplaceShape(IterType begin, IterType end
   return p;
 }
 
+TVM_FFI_INLINE ObjectPtr<ShapeObj> MakeStridesFromShape(int64_t ndim, int64_t* shape) {
+  int64_t* strides_data;
+  ObjectPtr<ShapeObj> strides = details::MakeEmptyShape(ndim, &strides_data);
+  int64_t stride = 1;
+  for (int i = ndim - 1; i >= 0; --i) {
+    strides_data[i] = stride;
+    stride *= shape[i];
+  }
+  return strides;
+}
+
 }  // namespace details
 
 /*!

tests/cpp/test_ndarray.cc

@@ -69,7 +69,9 @@ TEST(NDArray, DLPack) {
   EXPECT_EQ(dlpack->dl_tensor.device.device_type, kDLCPU);
   EXPECT_EQ(dlpack->dl_tensor.device.device_id, 0);
   EXPECT_EQ(dlpack->dl_tensor.byte_offset, 0);
-  EXPECT_EQ(dlpack->dl_tensor.strides, nullptr);
+  EXPECT_EQ(dlpack->dl_tensor.strides[0], 6);
+  EXPECT_EQ(dlpack->dl_tensor.strides[1], 3);
+  EXPECT_EQ(dlpack->dl_tensor.strides[2], 1);
   EXPECT_EQ(nd.use_count(), 2);
   {
     NDArray nd2 = NDArray::FromDLPack(dlpack);
@@ -96,7 +98,7 @@ TEST(NDArray, DLPackVersioned) {
   EXPECT_EQ(dlpack->dl_tensor.device.device_type, kDLCPU);
   EXPECT_EQ(dlpack->dl_tensor.device.device_id, 0);
   EXPECT_EQ(dlpack->dl_tensor.byte_offset, 0);
-  EXPECT_EQ(dlpack->dl_tensor.strides, nullptr);
+  EXPECT_EQ(dlpack->dl_tensor.strides[0], 1);
 
   EXPECT_EQ(nd.use_count(), 2);
   {