[sparse] use _sparse_coo_tensor_unsafe in coalesce for speedup

Aapo Kyrola · Aapo Kyrola · commit c79fb30f57f3 · 2019-05-31T13:50:09.000-07:00
diff --git a/aten/src/ATen/native/sparse/cuda/SparseCUDATensor.cu b/aten/src/ATen/native/sparse/cuda/SparseCUDATensor.cu
@@ -143,8 +143,10 @@ SparseTensor coalesce_sparse_cuda(const SparseTensor& self) {
     }
   }
   ////////////////////////////////////////////////////////////
-
-  SparseTensor dst = ::at::native::sparse_coo_tensor(newIndices, newValues, self.sizes())._coalesced_(true);
+  // We can use unsafe sparse tensor constructor because the indices do not
+  // need to be revalidated as we do not add or change indices, just remove
+  // duplicates.
+  SparseTensor dst = ::at::native::_sparse_coo_tensor_unsafe(newIndices, newValues, self.sizes())._coalesced_(true);
 
   THCudaCheck(cudaGetLastError());
   return dst;

Original file line number	Diff line number	Diff line change
`@@ -143,8 +143,10 @@ SparseTensor coalesce_sparse_cuda(const SparseTensor& self) {`
`143`	`143`	`}`
`144`	`144`	`}`
`145`	`145`	`////////////////////////////////////////////////////////////`
`146`		`-`
`147`		`- SparseTensor dst = ::at::native::sparse_coo_tensor(newIndices, newValues, self.sizes())._coalesced_(true);`
	`146`	`+ // We can use unsafe sparse tensor constructor because the indices do not`
	`147`	`+ // need to be revalidated as we do not add or change indices, just remove`
	`148`	`+ // duplicates.`
	`149`	`+ SparseTensor dst = ::at::native::_sparse_coo_tensor_unsafe(newIndices, newValues, self.sizes())._coalesced_(true);`
`148`	`150`
`149`	`151`	`THCudaCheck(cudaGetLastError());`
`150`	`152`	`return dst;`