view in certain noncontig case (#4062)

ssnl · soumith · commit e0d5d1b7c92f · 2017-12-18T02:08:17.000-05:00
diff --git a/aten/src/TH/generic/THTensor.c b/aten/src/TH/generic/THTensor.c
@@ -226,14 +226,61 @@ THTensor *THTensor_(newUnfold)(THTensor *tensor, int dimension_, int64_t size_,
   return self;
 }
 
+// Also sets new_stride if viewable.
+//
+// On a high level,
+// 1. separate tensor->size into chunks of dimensions, where the dimensions are
+//    ``contiguous'' in each chunk, i.e., stride[i] = size[i+1] * stride[i+1]
+// 2. view_size must be able to be separated into same number of chunks, where
+//    each chunk pair has matching ``numel'', i.e., number of subspaces.
+static int THTensor_(isViewable)(THTensor *tensor, THLongStorage *view_size, THLongStorage *new_stride) {
+  // dim indices
+  int64_t tensor_d = tensor->nDimension - 1;
+  if (tensor_d < 0) {
+    return 1;
+  }
+  int64_t view_d = view_size->size - 1;
+  // stride for each subspace in the chunk
+  int64_t chunk_base_stride = tensor->stride[tensor_d];
+  // numel in current chunk
+  int64_t tensor_numel = 1;
+  int64_t view_numel = 1;
+  for (; tensor_d >= 0; tensor_d--) {
+    tensor_numel *= tensor->size[tensor_d];
+    // if end of tensor size chunk, check view
+    if ((tensor_d == 0) ||
+        (tensor->size[tensor_d - 1] != 1 && tensor->stride[tensor_d - 1] != tensor_numel * chunk_base_stride)) {
+      while ((view_numel < tensor_numel || view_size->data[view_d] == 1) && view_d >= 0) {
+        new_stride->data[view_d] = view_numel * chunk_base_stride;
+        view_numel *= view_size->data[view_d];
+        view_d--;
+      }
+      if (view_numel != tensor_numel) {
+        return 0;
+      }
+      if (tensor_d > 0) {
+        chunk_base_stride = tensor->stride[tensor_d - 1];
+        tensor_numel = 1;
+        view_numel = 1;
+      }
+    }
+  }
+  // check that we iterated through all view size
+  return view_d == -1;
+}
+
 THTensor *THTensor_(newView)(THTensor *tensor, THLongStorage *size)
 {
-  THArgCheck(THTensor_(isContiguous)(tensor), 1, "input is not contiguous");
   ptrdiff_t numel = THTensor_(nElement)(tensor);
   THTensor *self = THTensor_(new)();
   THLongStorage *inferred_size = THLongStorage_newInferSize(size, numel);
-  THTensor_(setStorage)(self, tensor->storage, tensor->storageOffset, inferred_size, NULL);
+  THLongStorage *new_stride = THLongStorage_newWithSize(size->size);
+  THArgCheck(THTensor_(isViewable)(tensor, inferred_size, new_stride), 2, "view size is "
+    "not compatible with input tensor's size and stride (at least one dimension spans "
+    "across two contiguous subspaces). Call .contiguous() before .view().");
+  THTensor_(setStorage)(self, tensor->storage, tensor->storageOffset, inferred_size, new_stride);
   THLongStorage_free(inferred_size);
+  THLongStorage_free(new_stride);
   return self;
 }
 
diff --git a/aten/src/THC/generic/THCTensor.c b/aten/src/THC/generic/THCTensor.c
@@ -222,14 +222,61 @@ THCTensor *THCTensor_(newUnfold)(THCState *state, THCTensor *tensor, int dimensi
   return self;
 }
 
+// Also sets new_stride if viewable.
+//
+// On a high level,
+// 1. separate tensor->size into chunks of dimensions, where the dimensions are
+//    ``contiguous'' in each chunk, i.e., stride[i] = size[i+1] * stride[i+1]
+// 2. view_size must be able to be separated into same number of chunks, where
+//    each chunk pair has matching ``numel'', i.e., number of subspaces.
+static int THCTensor_(isViewable)(THCState *state, THCTensor *tensor, THLongStorage *view_size, THLongStorage *new_stride) {
+  // dim indices
+  int64_t tensor_d = tensor->nDimension - 1;
+  if (tensor_d < 0) {
+    return 1;
+  }
+  int64_t view_d = view_size->size - 1;
+  // stride for each subspace in the chunk
+  int64_t chunk_base_stride = tensor->stride[tensor_d];
+  // numel in current chunk
+  int64_t tensor_numel = 1;
+  int64_t view_numel = 1;
+  for (; tensor_d >= 0; tensor_d--) {
+    tensor_numel *= tensor->size[tensor_d];
+    // if end of tensor size chunk, check view
+    if ((tensor_d == 0) ||
+        (tensor->size[tensor_d - 1] != 1 && tensor->stride[tensor_d - 1] != tensor_numel * chunk_base_stride)) {
+      while ((view_numel < tensor_numel || view_size->data[view_d] == 1) && view_d >= 0) {
+        new_stride->data[view_d] = view_numel * chunk_base_stride;
+        view_numel *= view_size->data[view_d];
+        view_d--;
+      }
+      if (view_numel != tensor_numel) {
+        return 0;
+      }
+      if (tensor_d > 0) {
+        chunk_base_stride = tensor->stride[tensor_d - 1];
+        tensor_numel = 1;
+        view_numel = 1;
+      }
+    }
+  }
+  // check that we iterated through all view size
+  return view_d == -1;
+}
+
 THCTensor *THCTensor_(newView)(THCState *state, THCTensor *tensor, THLongStorage *size)
 {
-  THArgCheck(THCTensor_(isContiguous)(state, tensor), 2, "input is not contiguous");
   ptrdiff_t numel = THCTensor_(nElement)(state, tensor);
   THCTensor *self = THCTensor_(new)(state);
   THLongStorage *inferred_size = THLongStorage_newInferSize(size, numel);
-  THCTensor_(setStorage)(state, self, tensor->storage, tensor->storageOffset, inferred_size, NULL);
+  THLongStorage *new_stride = THLongStorage_newWithSize(size->size);
+  THArgCheck(THCTensor_(isViewable)(state, tensor, inferred_size, new_stride), 2, "View size is "
+    "not compatible with input tensor's size and stride (at least one dimension spans "
+    "across two contiguous subspaces). Call .contiguous() before .view().");
+  THCTensor_(setStorage)(state, self, tensor->storage, tensor->storageOffset, inferred_size, new_stride);
   THLongStorage_free(inferred_size);
+  THLongStorage_free(new_stride);
   return self;
 }
 
diff --git a/test/test_cuda.py b/test/test_cuda.py
@@ -302,7 +302,7 @@ def tmp(t):
     ('triu', medium_2d, lambda t: [-2], 'negative'),
     ('unsqueeze', new_t(2, 3, 4), lambda t: [2],),
     ('unsqueeze', new_t(2, 3, 4), lambda t: [-2], 'neg_dim'),
-    ('view', small_3d, lambda t: [100, 10],),
+    ('view', small_3d, lambda t: [100, 10], 'contiguous'),
     ('view_as', small_3d, lambda t: [t(100, 10)],),
     ('zero', small_3d, lambda t: [],),
     ('zeros', small_3d, lambda t: [1, 2, 3, 4],),
@@ -989,6 +989,9 @@ def _select_broadcastable_dims(dims_full=None):
     def test_det(self):
         TestTorch._test_det(self, lambda t: t.cuda())
 
+    def test_view(self):
+        TestTorch._test_view(self, lambda t: t.cuda())
+
     def test_broadcast(self):
         TestTorch._test_broadcast(self, lambda t: t.cuda())
 
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -3828,10 +3828,11 @@ def test_var_stability(self):
         self.assertEqual(tensor.var(0)[0], 0.03125)
         self.assertEqual(tensor.var(), 0.03125)
 
-    def test_view(self):
-        tensor = torch.rand(15)
-        template = torch.rand(3, 5)
-        empty = torch.Tensor()
+    @staticmethod
+    def _test_view(self, cast):
+        tensor = cast(torch.rand(15))
+        template = cast(torch.rand(3, 5))
+        empty = cast(torch.Tensor())
         target = template.size()
         self.assertEqual(tensor.view_as(template).size(), target)
         self.assertEqual(tensor.view(3, 5).size(), target)
@@ -3848,6 +3849,52 @@ def test_view(self):
         self.assertRaises(RuntimeError, lambda: tensor.view(15, 0))
         self.assertRaises(RuntimeError, lambda: tensor.view(7, -1))
         self.assertRaises(RuntimeError, lambda: tensor.view(15, -1, -1))
+        # test view when tensor is not contiguous in every dimension, but only
+        # contiguous dimensions are touched.
+        tensor = cast(torch.rand(4, 2, 5, 1, 6, 2, 9, 3)).transpose(-1, 2).transpose(-2, 3)
+        # size:                      [   4,    2,    3,    9,    6,    2,    1,    5]
+        # stride:                    [3840, 1620,    1,    3,   54,   27,  324,  324]
+        # contiguous dim chunks:     [__________, ____, ____, __________, ____, ____]
+        # merging 1 to chunk after:  [__________, ____, ____, __________, __________]
+        contig_tensor = tensor.clone()
+        # [4, 2] => [8, 1]
+        # [3] => [3]
+        # [9] => [3, 3]
+        # [6, 2] => [4, 1, 3]
+        # [1, 5] => [5]
+        view_size = [8, 1, 3, 3, 3, 4, 1, 3, 5]
+        self.assertEqual(tensor.view(*view_size), contig_tensor.view(*view_size))
+        # [4, 2] => [2, 4]
+        # [3] => [3]
+        # [9] => [1, 9]
+        # [6, 2] => [2, 2, 3]
+        # [1, 5] => [5, 1]
+        view_size = [2, 4, 3, 1, 9, 2, 2, 3, 5, 1]
+        self.assertEqual(tensor.view(*view_size), contig_tensor.view(*view_size))
+        # adding size 1 dims
+        view_size = [1, 1, 2, 1, 4, 3, 1, 1, 9, 1, 2, 1, 2, 3, 1, 5, 1, 1]
+        self.assertEqual(tensor.view(*view_size), contig_tensor.view(*view_size))
+
+        # invalid views
+        self.assertRaises(RuntimeError, lambda: tensor.view(-1))
+        # crossing [4, 2], [3]
+        self.assertRaises(RuntimeError, lambda: tensor.view(24, 9, 6, 2, 1, 5))
+        # crossing [6, 2], [1, 5]
+        self.assertRaises(RuntimeError, lambda: tensor.view(8, 3, 9, 6, 10))
+        # crossing [9], [6, 2]
+        self.assertRaises(RuntimeError, lambda: tensor.view(8, 3, 54, 2, 1, 5))
+
+        # view with stride 0 dims
+        tensor = cast(torch.Tensor(1, 1)).expand(3, 4)  # all dims are contiguous
+        contig_tensor = tensor.clone()
+        self.assertEqual(tensor.view(-1), contig_tensor.view(-1))
+        self.assertEqual(tensor.view(1, -1, 1), contig_tensor.view(1, -1, 1))
+        self.assertEqual(tensor.view(-1, 1), contig_tensor.view(-1, 1))
+        self.assertEqual(tensor.view(6, 2, 1), contig_tensor.view(6, 2, 1))
+        self.assertEqual(tensor.view(1, 6, 2, 1), contig_tensor.view(1, 6, 2, 1))
+
+    def test_view(self):
+        TestTorch._test_view(self, lambda x: x)
 
     def test_expand(self):
         tensor = torch.rand(1, 8, 1)
diff --git a/torch/_tensor_docs.py b/torch/_tensor_docs.py
@@ -1789,8 +1789,18 @@ def callable(a, b) -> number
 different size.
 
 The returned tensor shares the same data and must have the same number
-of elements, but may have a different size. A tensor must be
-:func:`contiguous` to be viewed.
+of elements, but may have a different size. For a tensor to be viewed, the new
+view size must be compatible with its original size and stride, i.e., each new
+view dimension must either be a subspace of an original dimension, or only span
+across original dimensions :math:`d, d+1, \dots, d+k` that satisfy the following
+contiguity-like condition that :math:`\forall i = 0, \dots, k-1`,
+
+.. math::
+
+  stride[i] = stride[i+1] \times size[i+1]
+
+Otherwise, :func:`contiguous` needs to be called before the tensor can be
+viewed.
 
 Args:
     args (torch.Size or int...): the desired size
