!13038 Separate tensor.resize_ & storage.resize_

Merge pull request !13038 from 姜怡文/master_resize

Author: 姜怡文

test/npu/test_combine_tensors.py

@@ -81,6 +81,12 @@ def test_storage_resize(self, dtype, device="npu"):
             self.assertEqual(tensor_new.storage(), tensor.storage())
             idx += torch_npu.get_storage_size(tensor)
 
+    @Dtypes(torch.half, torch.float, torch.bfloat16)
+    def test_untyped_storage_resize(self, dtype, device="npu"):
+        a = torch.randn((2, 2, 2, 2, 2), device=device, dtype=dtype)
+        a.untyped_storage().resize_(0)
+        a.untyped_storage().resize_(128)
+
     @Dtypes(torch.half, torch.float, torch.bfloat16)
     def test_combine_tensors(self, dtype, device="npu"):
         x = torch.zeros((2, 2, 2, 2), device=device, dtype=dtype)

torch_npu/csrc/aten/common/ResizeNpu.h

@@ -14,10 +14,13 @@
 namespace at_npu {
 namespace native {
 
+// Refresh storage_desc to ND if set force_refresh = true,
+// mainly used in storage.resize_
 static void storage_resize_npu(
     torch_npu::NPUStorageImpl& storage,
     ptrdiff_t size,
-    c10::IntArrayRef new_size)
+    c10::IntArrayRef new_size,
+    bool force_refresh = false)
 {
     if (!storage.resizable()) {
         TORCH_CHECK(false, "Trying to resize storage that is not resizable", OPS_ERROR(ErrCode::NOT_SUPPORT));
@@ -50,25 +53,29 @@ static void storage_resize_npu(
     };
     // It is necessary to properly refresh the storage according to sizes and strides,
     // not just new sizes.
-    int64_t new_data_numel = c10::multiply_integers(resize_shape);
-    int64_t new_shape_numel = c10::multiply_integers(new_size);
-    const c10::IntArrayRef &refresh_size = new_data_numel > new_shape_numel ? resize_shape : new_size;
-
-    // 计算连续场景下size对应的stride值
-    int64_t dim_ = static_cast<int64_t>(refresh_size.size());
-    c10::SmallVector<int64_t, 5> new_stride(dim_);
-    if (dim_ > 0) {
-        int64_t last_idx = dim_ - 1;
-        new_stride[last_idx] = 1;
-        for (auto i = last_idx - 1; i >= 0; --i) {
-            new_stride[i] = new_stride[i + 1] * std::max<int64_t>(refresh_size[i + 1], 1);
+    if (force_refresh) {
+        int64_t new_data_numel = c10::multiply_integers(resize_shape);
+        int64_t new_shape_numel = c10::multiply_integers(new_size);
+        const c10::IntArrayRef &refresh_size = new_data_numel > new_shape_numel ? resize_shape : new_size;
+
+        // 计算连续场景下size对应的stride值
+        int64_t dim_ = static_cast<int64_t>(refresh_size.size());
+        c10::SmallVector<int64_t, 5> new_stride(dim_);
+        if (dim_ > 0) {
+            int64_t last_idx = dim_ - 1;
+            new_stride[last_idx] = 1;
+            for (auto i = last_idx - 1; i >= 0; --i) {
+                new_stride[i] = new_stride[i + 1] * std::max<int64_t>(refresh_size[i + 1], 1);
+            }
         }
-    }
 
-    storage_desc.base_sizes_ = refresh_size;
-    storage_desc.base_strides_ = new_stride;
-    storage_desc.npu_format_ = ACL_FORMAT_ND;
-    storage_desc.storage_sizes_ = storage_desc.base_sizes_;
+        storage_desc.base_sizes_ = refresh_size;
+        storage_desc.base_strides_ = new_stride;
+        storage_desc.npu_format_ = ACL_FORMAT_ND;
+        storage_desc.storage_sizes_ = storage_desc.base_sizes_;
+    } else {
+        StorageDescHelper::UpdateDesc(storage_desc, resize_shape, new_size);
+    }
 
     if (old_data != nullptr) {
         ptrdiff_t copy_size = old_size;

torch_npu/csrc/core/npu/NPUHooksInterface.cpp

@@ -38,7 +38,7 @@ void NPUHooksInterface::resizePrivateUse1Bytes(const c10::Storage &storage, size
 
     auto itemsize = storage_impl->npu_desc_.data_type_.itemsize();
     std::vector<int64_t> new_size = {static_cast<int64_t>(new_bytes) / (ptrdiff_t)itemsize};
-    at_npu::native::storage_resize_npu(*storage_impl, new_bytes, new_size);
+    at_npu::native::storage_resize_npu(*storage_impl, new_bytes, new_size, true);
 }
 
 at::PrivateUse1HooksInterface* get_npu_hooks()