Validate inputs to _nested_view_from_buffer to prevent overflows (pytorch#147356)

Pull Request resolved: pytorch#147356 Approved by: https://github.com/albanD, https://github.com/jbschlosser ghstack dependencies: pytorch#147352, pytorch#147354

Author: mikaylagawarecki

aten/src/ATen/native/nested/NestedTensorMath.cpp

@@ -12,6 +12,7 @@
 #include <ATen/WrapDimUtilsMulti.h>
 #include <ATen/core/Tensor.h>
 #include <ATen/native/layer_norm.h>
+#include <ATen/native/Resize.h>
 #include <ATen/native/nested/NestedTensorUtils.h>
 
 #include <tuple>
@@ -879,19 +880,28 @@ Tensor _nested_view_from_buffer(
  "Can only a create Nested Tensor from a normal tensor buffer");
  TORCH_INTERNAL_ASSERT(buffer.dim() == 1, "The input buffer must be flat");
  TORCH_INTERNAL_ASSERT(nested_sizes.dim() == 2, "Expected the nested size tensor to be two dimensional.");
- uint64_t num_elements_nested_size = at::prod(nested_sizes, 1).sum().item<int64_t>();
- uint64_t buffer_storage_size = buffer.storage().nbytes()/buffer.dtype().itemsize();
- TORCH_INTERNAL_ASSERT(
- buffer_storage_size == num_elements_nested_size,
- "The number of elements in the buffer must equal the nested tensor size but buffer size: ",
- buffer_storage_size,
- " and nested tensor size: ",
- num_elements_nested_size,
- ".");
-
  TORCH_INTERNAL_ASSERT(nested_strides.dim() == 2, "Expected the nested stride tensor to be two dimensional.");
  TORCH_INTERNAL_ASSERT(nested_sizes.size(0) == nested_strides.size(0), "Expected the first dimension of nested size and nested stride tensor to be equal.");
  TORCH_INTERNAL_ASSERT(nested_strides.size(0) == storage_offsets.size(0), "Expected the first dimension of nested stride tensor to equal the length of offsets.");
+
+
+ std::vector<at::Tensor> all_sizes = nested_sizes.unbind();
+ std::vector<at::Tensor> all_strides = nested_strides.unbind();
+ std::vector<at::Tensor> all_offsets = storage_offsets.unbind();
+ auto size_dim = nested_sizes.size(1);
+
+ for (const auto i : c10::irange(nested_sizes.size(0))) {
+ const int64_t* sizemat_ptr = all_sizes[i].const_data_ptr<int64_t>();
+ const int64_t* stridemat_ptr = all_strides[i].const_data_ptr<int64_t>();
+ const int64_t* offset_ptr = all_offsets[i].const_data_ptr<int64_t>();
+ checkInBoundsForStorage(
+ IntArrayRef(sizemat_ptr, sizemat_ptr + size_dim),
+ IntArrayRef(stridemat_ptr, stridemat_ptr + size_dim),
+ *offset_ptr,
+ buffer.dtype(),
+ buffer.storage());
+ }
+
  return at::detail::make_tensor<NestedTensorImpl>(
  c10::TensorImpl::VIEW,
  buffer,

test/test_nestedtensor.py

@@ -857,6 +857,21 @@ def test_cat(self):
  ):
  torch.cat([x, y], dim=-1)
 
+ def test_nested_view_from_buffer_overflow_errors(self):
+ buffer = torch.tensor([1])
+ sizes = torch.tensor([[2**63 - 1], [2**63 - 1], [3]], dtype=torch.int64)
+ strides = torch.tensor(
+ [[0x41414141], [0x41414141], [0x41414141]], dtype=torch.int64
+ )
+ offsets = torch.tensor(
+ [[0x41414141], [0x41414141], [0x41414141]], dtype=torch.int64
+ )
+ with self.assertRaisesRegex(
+ RuntimeError,
+ r"Storage size calculation overflowed with sizes=\[9223372036854775807\] and strides=\[1094795585\]",
+ ):
+ nt = torch._nested_view_from_buffer(buffer, sizes, strides, offsets)
+
 
 @markDynamoStrictTest
 class TestNestedTensorDeviceType(NestedTensorTestCase):