Skip to content

Commit 572132f

Browse files
weiyangfbfacebook-github-bot
weiyangfb
authored and
facebook-github-bot
committed
copy_(Sparse, Sparse) for sparse tensor (#9005)
Summary: - fix #8330 - add `torch.copy_(Sparse, Sparse)` with autograd support Pull Request resolved: #9005 Differential Revision: D8987885 Pulled By: weiyangfb fbshipit-source-id: b317a41da22ee1eae2835622a0ed28a6771a3a06
1 parent 93ecf4d commit 572132f

File tree

9 files changed

+90
-26
lines changed

9 files changed

+90
-26
lines changed

‎aten/src/ATen/native/native_functions.yaml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2068,8 +2068,8 @@
20682068
SparseCPU: hspmm_sparse_cpu
20692069
SparseCUDA: hspmm_sparse_cuda
20702070

2071-
# This "raw copy" doesn't handle conversions NOR does it handle non-blocking.
2072-
- func: raw_copy_sparse_(Tensor self, Tensor src) -> Tensor
2071+
- func: copy_sparse_to_sparse_(Tensor self, Tensor src, bool non_blocking=false) -> Tensor
2072+
variants: function
20732073
dispatch:
20742074
SparseCPU: copy_sparse_
20752075
SparseCUDA: copy_sparse_

‎aten/src/ATen/native/sparse/SparseTensor.cpp‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -204,7 +204,7 @@ SparseTensor new_with_tensor_and_size_sparse(const LongTensor& indices, const Te
204204

205205
SparseTensor clone_sparse(const SparseTensor& self) {
206206
SparseTensor other = new_with_dims_and_size_sparse(self.type(), self._sparseDims(), self._denseDims(), self.sizes());
207-
_copy_into_sparse(other, _get_sparse_impl(self)->indices(), _get_sparse_impl(self)->values());
207+
_copy_into_sparse(other, _get_sparse_impl(self)->indices(), _get_sparse_impl(self)->values(), true);
208208
_get_sparse_impl(other)->set_coalesced(self.is_coalesced());
209209
return other;
210210
}
@@ -243,11 +243,11 @@ Tensor sparse_to_dense(const SparseTensor& self) {
243243
return dst.add_(self);
244244
}
245245

246-
SparseTensor& copy_sparse_(SparseTensor& self, const SparseTensor& src) {
246+
SparseTensor& copy_sparse_(SparseTensor& self, const SparseTensor& src, bool non_blocking) {
247247
if (isSameTensor(self, src)) return self;
248248
_get_sparse_impl(self)->resize_(src._sparseDims(), src._denseDims(), src.sizes());
249249
// NB: This seems to copy the underlying full indices/values buffer
250-
_copy_into_sparse(self, _get_sparse_impl(src)->indices(), _get_sparse_impl(src)->values());
250+
_copy_into_sparse(self, _get_sparse_impl(src)->indices(), _get_sparse_impl(src)->values(), non_blocking);
251251
_get_sparse_impl(self)->set_coalesced(src.is_coalesced());
252252
return self;
253253
}

‎aten/src/ATen/native/sparse/SparseTensorMath.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,7 @@ SparseTensor& log1p_out_sparse(SparseTensor& r, const SparseTensor& t) {
9898
r.is_coalesced(), "log1p: in-place on uncoalesced tensors is not supported yet!");
9999
}
100100
else {
101-
r = raw_copy_sparse_(r, t.coalesce());
101+
copy_sparse_to_sparse_(r, t.coalesce());
102102
}
103103
r._values().log1p_();
104104
return r;
@@ -192,7 +192,7 @@ SparseTensor& add_out_sparse_cpu(SparseTensor& r, const SparseTensor& t, const S
192192
AT_CHECK(t.sizes().equals(src.sizes()), "add: expected sizes of 'self' and 'other' to match, but ", t.sizes(), " != ", src.sizes());
193193

194194
if (src._nnz() == 0) {
195-
return raw_copy_sparse_(r, t);
195+
return copy_sparse_to_sparse_(r, t);
196196
}
197197
if (t._nnz() == 0) {
198198
return mul_out_sparse_scalar(r, src, value);

‎aten/src/ATen/native/sparse/SparseUtils.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,8 @@ inline void _alias_into_sparse(const SparseTensor& self, const LongTensor& indic
5050

5151
// Take indices and values and makes a (data) copy of them to put into the sparse
5252
// indices/values. This used to be called THSTensor_(_set)
53-
inline void _copy_into_sparse(const SparseTensor& self, const LongTensor& indices, const Tensor& values) {
54-
_alias_into_sparse(self, indices.clone(), values.clone());
53+
inline void _copy_into_sparse(const SparseTensor& self, const LongTensor& indices, const Tensor& values, bool non_blocking) {
54+
_alias_into_sparse(self, self._indices().type().copy(indices, non_blocking), self._values().type().copy(values, non_blocking));
5555
}
5656

5757
// Does NOT make copies of indices/values

‎aten/src/ATen/native/sparse/cuda/SparseCUDATensorMath.cu‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -363,7 +363,7 @@ SparseTensor& add_out_sparse_cuda(SparseTensor& r_, const SparseTensor& t, const
363363
AT_CHECK(t.sizes().equals(src.sizes()), "add: expected 'self' and 'other' to have same size, but ", t.sizes(), " != ", src.sizes());
364364

365365
if (src._nnz() == 0) {
366-
return raw_copy_sparse_(r_, t);
366+
return copy_sparse_to_sparse_(r_, t);
367367
}
368368
if (t._nnz() == 0) {
369369
return mul_out_sparse_scalar(r_, src, value);

‎aten/src/ATen/templates/TypeDefault.cpp‎

Lines changed: 7 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,8 @@ namespace at {
1818

1919
Tensor & TypeDefault::copy_(Tensor & self, const Tensor & src, bool non_blocking) const {
2020
Tensor b_src;
21-
std::tie(b_src) = expand_inplace(self, src, "copy");
21+
if (is_sparse()) b_src = src;
22+
else std::tie(b_src) = expand_inplace(self, src, "copy");
2223
return s_copy_(self, b_src, non_blocking);
2324
}
2425

@@ -28,19 +29,11 @@ Tensor TypeDefault::copy(const Tensor & src, bool non_blocking, optional<Device>
2829
device_guard.set_index(to_device.value().index());
2930
}
3031
AT_CHECK(src.defined(), "attempt to copy an undefined tensor");
31-
if (is_sparse()) {
32-
auto indices = src._indices();
33-
auto values = src._values();
34-
auto & this_dense = toBackend(is_cuda() ? Backend::CUDA : Backend::CPU);
35-
auto & this_dense_idx = this_dense.toScalarType(ScalarType::Long);
36-
auto indices_copy = this_dense_idx.copy(indices, non_blocking);
37-
auto values_copy = this_dense.copy(values, non_blocking);
38-
return _sparse_coo_tensor_unsafe(indices_copy, values_copy, src.sizes());
39-
} else {
40-
Tensor r = this->tensor(src.sizes());
41-
r.copy_(src, non_blocking);
42-
return r;
43-
}
32+
Tensor r;
33+
if (is_sparse()) r = this->native_tensor();
34+
else r = this->tensor(src.sizes());
35+
r.copy_(src, non_blocking);
36+
return r;
4437
}
4538

4639
void TypeDefault::backward(Tensor & self, at::optional<Tensor> gradient, bool keep_graph, bool create_graph) const {

‎test/test_sparse.py‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -496,6 +496,74 @@ def test_shape(sparse_dims, nnz, with_size):
496496
test_shape(3, 10, [100, 100, 100, 5, 5, 5, 0])
497497
test_shape(3, 0, [0, 0, 100, 5, 5, 5, 0])
498498

499+
def test_Sparse_to_Sparse_copy_(self):
500+
# This is for testing torch.copy_(SparseTensor, SparseTensor)
501+
sparse_dims = 3
502+
nnz = 10
503+
sizes = [2, 3, 4, 5] # hybrid sparse
504+
x1, _, _ = self._gen_sparse(sparse_dims, nnz, sizes)
505+
x2, _, _ = self._gen_sparse(sparse_dims, nnz + 10, sizes)
506+
507+
# test copy
508+
x2_dense = x2.to_dense()
509+
x1.copy_(x2)
510+
self.assertEqual(x2_dense, x1.to_dense())
511+
512+
# test type conversion (when x1.copy_(x2), x1.dtype should stay the same)
513+
x1 = x1.to(torch.float32)
514+
x2 = x2.to(torch.float64)
515+
x1_dtype = x1.dtype
516+
x1.copy_(x2)
517+
self.assertEqual(x1_dtype, x1.dtype)
518+
519+
# test no broadcast
520+
self.assertRaises(RuntimeError, lambda: x1.copy_(x2.narrow_copy(0, 0, 1)))
521+
522+
# test raise error on copy_() between dense and sparse Tensors
523+
self.assertRaises(RuntimeError, lambda: x1.copy_(torch.randn(5, 5)))
524+
525+
# test autograd
526+
x1, _, _ = self._gen_sparse(sparse_dims, nnz, sizes)
527+
x2, _, _ = self._gen_sparse(sparse_dims, nnz + 10, sizes)
528+
x2.requires_grad_(True)
529+
x1.copy_(x2)
530+
y = x1 * 2
531+
x2_clone = x2.clone()
532+
y.backward(x2_clone)
533+
expected_grad = x2_clone * 2
534+
self.assertEqual(expected_grad.to_dense(), x2.grad.to_dense())
535+
self.assertEqual(None, x1.grad)
536+
537+
@unittest.skipIf(torch.cuda.device_count() < 2, "no multi-GPU")
538+
def test_Sparse_to_Sparse_copy_multi_gpu(self):
539+
# This is for testing torch.copy_(SparseTensor, SparseTensor) across GPU devices
540+
sparse_dims = 3
541+
nnz = 10
542+
sizes = [2, 3, 4, 5] # hybrid sparse
543+
x1, _, _ = self._gen_sparse(sparse_dims, nnz, sizes)
544+
x2, _, _ = self._gen_sparse(sparse_dims, nnz + 10, sizes)
545+
x1 = x1.to('cuda:0')
546+
547+
def test_cross_device(x1, x2):
548+
x1_device = x1.device
549+
x1.copy_(x2)
550+
self.assertEqual(x2.to('cuda:0').to_dense(), x1.to_dense())
551+
self.assertEqual(x1_device, x1.device)
552+
553+
test_cross_device(x1, x2.to('cuda:1')) # test across gpu devices
554+
test_cross_device(x1, x2.to('cpu')) # test between cpu and gpu
555+
556+
# test autograd
557+
x2 = x2.to('cuda:1')
558+
x2.requires_grad_(True)
559+
x1.copy_(x2)
560+
y = x1 * 2
561+
x2_clone = x2.clone().to('cuda:0')
562+
y.backward(x2_clone)
563+
expected_grad = x2_clone * 2
564+
self.assertEqual(expected_grad.to_dense(), x2.grad.to('cuda:0').to_dense())
565+
self.assertEqual(None, x1.grad)
566+
499567
@cuda_only
500568
def test_cuda_empty(self):
501569
def test_tensor(x):

‎tools/autograd/gen_python_functions.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,8 @@
2828
'_cumsum.*', '_cumprod.*', '_sum.*', '_prod.*', '_th_.*',
2929
'arange.*', 'range.*', '_gesv.*', '_getri.*', 'slice', 'randint(_out)?',
3030
'_local_scalar', '_local_scalar_dense',
31-
'max_pool1d', 'max_pool2d', 'max_pool3d', 'linear', 'to'
31+
'max_pool1d', 'max_pool2d', 'max_pool3d', 'linear', 'to',
32+
'copy_sparse_to_sparse_'
3233
]
3334

3435
# These function signatures are not exposed to Python. Note that this signature

‎tools/autograd/templates/VariableType.cpp‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -416,7 +416,9 @@ Tensor & VariableType::s_copy_(Tensor & self, const Tensor & src, bool non_block
416416
grad_fn->src_device = src.get_device();
417417
}
418418
}
419-
baseType->s_copy_(self_, src_, non_blocking);
419+
if (self.is_sparse() && src.is_sparse()) baseType->copy_sparse_to_sparse_(self_, src_, non_blocking);
420+
else if (!self.is_sparse() && !src.is_sparse()) baseType->s_copy_(self_, src_, non_blocking);
421+
else AT_ERROR("copy_() between dense and sparse Tensors is not implemented! Found self type = ", self.type(), " and src type = ", src.type());
420422
increment_version(self);
421423
rebase_history(as_variable_ref( self ), std::move(grad_fn));
422424
if(torch::jit::tracer::isTracing()) {

0 commit comments

Comments
 (0)