add narrow() support for sparse tensors re: #8853 (#11342)

realdoug · facebook-github-bot · commit c2f8f5076c58 · 2018-09-26T12:24:54.000-07:00
Summary: Couple questions: 1) I used the log1p implementation in #8969 as a guide especially for testing. I'm not sure what the ```skipIfROCM``` annotation is for, so unsure if i need it for my test. 2) I implemented the branching logic in the narrow function itself; is this the right place to do so? I noticed that there a number of places where sparse-specific logic is handled with just an if statement in this file. Or should I implement a separate dispatch in native_functions.yml as in the log1p? And of course, happy to make any any other updates/changes that I may have missed as well. This is my first PR to the project. Pull Request resolved: #11342 Differential Revision: D9978430 Pulled By: weiyangfb fbshipit-source-id: e73dc20302ab58925afb19e609e31f4a38c634ad
diff --git a/aten/src/ATen/core/Tensor.h b/aten/src/ATen/core/Tensor.h
@@ -533,6 +533,7 @@ class CAFFE2_API Tensor {
   Tensor mv(const Tensor & vec) const;
   Tensor mvlgamma(int64_t p) const;
   Tensor & mvlgamma_(int64_t p);
+  Tensor narrow_copy(int64_t dim, int64_t start, int64_t length) const;
   Tensor narrow(int64_t dim, int64_t start, int64_t length) const;
   Tensor permute(IntList dims) const;
   Tensor pin_memory() const;
diff --git a/aten/src/ATen/core/TensorMethods.h b/aten/src/ATen/core/TensorMethods.h
@@ -902,6 +902,9 @@ inline Tensor Tensor::mvlgamma(int64_t p) const {
 inline Tensor & Tensor::mvlgamma_(int64_t p) {
     return type().mvlgamma_(*this, p);
 }
+inline Tensor Tensor::narrow_copy(int64_t dim, int64_t start, int64_t length) const {
+    return type().narrow_copy(*this, dim, start, length);
+}
 inline Tensor Tensor::narrow(int64_t dim, int64_t start, int64_t length) const {
     return type().narrow(*this, dim, start, length);
 }
diff --git a/aten/src/ATen/core/Type.h b/aten/src/ATen/core/Type.h
@@ -492,6 +492,7 @@ struct CAFFE2_API Type {
   virtual Tensor mv(const Tensor & self, const Tensor & vec) const = 0;
   virtual Tensor mvlgamma(const Tensor & self, int64_t p) const = 0;
   virtual Tensor & mvlgamma_(Tensor & self, int64_t p) const = 0;
+  virtual Tensor narrow_copy(const Tensor & self, int64_t dim, int64_t start, int64_t length) const = 0;
   virtual Tensor narrow(const Tensor & self, int64_t dim, int64_t start, int64_t length) const = 0;
   virtual Tensor permute(const Tensor & self, IntList dims) const = 0;
   virtual Tensor pin_memory(const Tensor & self) const = 0;
diff --git a/aten/src/ATen/native/TensorShape.cpp b/aten/src/ATen/native/TensorShape.cpp
@@ -148,6 +148,45 @@ Tensor &as_strided_(Tensor& self, IntList size, IntList stride) {
   return at::as_strided_(self, size, stride, self.storage_offset());
 }
 
+Tensor narrow_copy_sparse(const Tensor& self, int64_t dim, int64_t start, int64_t length){
+  int64_t allDim = self.dim();
+  int64_t end = start+length;
+  AT_CHECK(allDim > 0, "narrow() cannot be applied to a 0-dim tensor.");
+  AT_CHECK(dim >= 0 && dim < allDim, 
+    "Dimension ", dim, " out of range. Expecting 0 <= dim < ", allDim, ".");
+  AT_CHECK(start >= 0 && length >= 0 && end <= self.size(dim),
+    "Invalid range to narrow. range(start, start+length) must be a subset of range(0, ", self.size(dim), ").")
+  LongTensor indices = self._indices();
+  int64_t sparseDims = self._sparseDims();
+  
+  std::vector<int64_t> newSizes = self.sizes().vec();
+  newSizes[dim]=length;
+  
+  Tensor newValues;
+  LongTensor newIndices;
+  if(dim < sparseDims){
+    Tensor mask = (indices[dim] >= start).__and__((indices[dim] < end));
+    newIndices = indices.masked_select(mask).view({sparseDims, -1});
+    newIndices[dim].add_(-start);
+    Tensor nzIndices = mask.nonzero().view(-1);
+    newValues = self._values().index_select(0, nzIndices);
+  }else{
+    /* This means we are narrowing on a dense dim, which is in effect just a
+        regular narrow on _values() */
+    newIndices = indices;
+    int64_t ddim = dim - sparseDims + 1;
+    newValues = self._values().narrow_copy(ddim, start, length);
+  }
+
+  SparseTensor newTensor = at::sparse_coo_tensor(newIndices, newValues, newSizes, self.type().options());
+  _get_sparse_impl(newTensor)->set_coalesced(self.is_coalesced());
+  return newTensor;
+}
+
+Tensor narrow_copy_dense(const Tensor& self, int64_t dim, int64_t start, int64_t length){
+    return self.narrow(dim, start, length).clone();
+}
+
 Tensor narrow(const Tensor& self, int64_t dim, int64_t start, int64_t length) {
   AT_CHECK(self.dim() > 0, "narrow() cannot be applied to a 0-dim tensor.");
   auto cur_size = self.size(dim);
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -1174,6 +1174,14 @@
 - func: mvlgamma_(Tensor self, int64_t p) -> Tensor
   variants: method
 
+- func: narrow_copy(Tensor self, int64_t dim, int64_t start, int64_t length) -> Tensor
+  variants: method
+  dispatch:
+    CPU: narrow_copy_dense
+    CUDA: narrow_copy_dense
+    SparseCPU: narrow_copy_sparse
+    SparseCUDA: narrow_copy_sparse
+
 - func: narrow(Tensor self, int64_t dim, int64_t start, int64_t length) -> Tensor
   variants: function, method
 
diff --git a/docs/source/sparse.rst b/docs/source/sparse.rst
@@ -110,6 +110,7 @@ An empty sparse tensor can be constructed by specifying its size:
     .. method:: mm
     .. method:: mul
     .. method:: mul_
+    .. method:: narrow_copy
     .. method:: resizeAs_
     .. method:: size
     .. method:: spadd
diff --git a/test/test_sparse.py b/test/test_sparse.py
@@ -1023,6 +1023,33 @@ def test_shape(i_shapes, v_shapes, nnzs):
         test_shape([0, 3, 4], [3, 4, 5, 6], [0])
         test_shape([2, 3, 4], [0, 4, 5, 6], [9, 12])
 
+    def _test_narrow(self, input, narrow_args):
+        expected = input.to_dense().narrow(*narrow_args)
+        self.assertEqual(expected, input.narrow_copy(*narrow_args).to_dense())
+
+    def _all_narrow_combs(self, shape):
+        for dim, dim_sz in enumerate(shape):
+            for start in range(dim_sz):
+                for length in range(dim_sz - start):
+                    yield [dim, start, length]
+
+    def test_narrow(self):
+        shape = [3, 3, 4, 2]
+        input, _, _ = self._gen_sparse(4, 19, shape)
+        for narrow_args in self._all_narrow_combs(shape):
+            self._test_narrow(input, narrow_args)
+
+        self.assertRaises(RuntimeError, lambda: input.narrow_copy(-1, 0, 3))  # dim < 0
+        self.assertRaises(RuntimeError, lambda: input.narrow_copy(10, 0, 3))  # dim > input.dim()
+        self.assertRaises(RuntimeError, lambda: input.narrow_copy(0, shape[0] + 1, 3))  # start > size of dim
+        self.assertRaises(RuntimeError, lambda: input.narrow_copy(0, 2, shape[0]))  # start+length > size of dim
+
+        with_dense, _, _ = self._gen_sparse(2, 7, shape)
+        for narrow_args in self._all_narrow_combs(shape):
+            self._test_narrow(with_dense, narrow_args)
+
+        self.assertRaises(RuntimeError, lambda: with_dense.narrow_copy(10, 0, 3))  # dim > sparseDim + denseDim
+
     def _test_log1p_tensor(self, input, dense_tensor):
         expected_output = torch.tensor(dense_tensor).log1p_()
         self.assertEqual(expected_output, input.log1p().to_dense())
diff --git a/torch/_tensor_docs.py b/torch/_tensor_docs.py
@@ -1453,6 +1453,17 @@ def callable(a, b) -> number
             [ 8,  9]])
 """)
 
+add_docstr_all('narrow_copy',
+               r"""
+narrow_copy(dimension, start, length) -> Tensor
+
+Same as :meth:`Tensor.narrow` except returning a copy rather
+than shared storage.  This is primarily for sparse tensors, which
+do not have a shared-storage narrow method.  Calling ```narrow_copy``
+with ```dimemsion > self._sparseDims()``` will return a copy with the
+relevant dense dimension narrowed, and ```self.shape``` updated accordingly.
+""")
+
 add_docstr_all('ndimension',
                r"""
 ndimension() -> int

Original file line number	Diff line number	Diff line change
`@@ -902,6 +902,9 @@ inline Tensor Tensor::mvlgamma(int64_t p) const {`
`902`	`902`	`inline Tensor & Tensor::mvlgamma_(int64_t p) {`
`903`	`903`	`return type().mvlgamma_(*this, p);`
`904`	`904`	`}`
	`905`	`+inline Tensor Tensor::narrow_copy(int64_t dim, int64_t start, int64_t length) const {`
	`906`	`+ return type().narrow_copy(*this, dim, start, length);`
	`907`	`+}`
`905`	`908`	`inline Tensor Tensor::narrow(int64_t dim, int64_t start, int64_t length) const {`
`906`	`909`	`return type().narrow(*this, dim, start, length);`
`907`	`910`	`}`