Make coalesce() out of place

Author: adamlerer

test/common.py

@@ -150,8 +150,8 @@ def assertTensorsEqual(a, b):
  self.assertLessEqual(max_err, prec, message)
  self.assertEqual(x.is_sparse, y.is_sparse, message)
  if x.is_sparse:
- x = x.clone().coalesce_()
- y = y.clone().coalesce_()
+ x = x.coalesce()
+ y = y.coalesce()
  assertTensorsEqual(x.indices(), y.indices())
  assertTensorsEqual(x.values(), y.values())
  else:

test/test_sparse.py

@@ -156,7 +156,7 @@ def _test_contig(self, is_cuda):
  [31, 92, 65, 50, 34, 62, 22, 56, 74, 89],
  ])
  exp_v = ValueTensor([2, 1, 6, 4, 10, 3, 5, 9, 8, 7])
- x.coalesce_()
+ x = x.coalesce()
  self.assertEqual(exp_i, x.indices())
  self.assertEqual(exp_v, x.values())
 
@@ -174,7 +174,7 @@ def _test_contig(self, is_cuda):
  ])
  exp_v = ValueTensor([2, 1, 3, 4])
 
- x.coalesce_()
+ x = x.coalesce()
  self.assertEqual(exp_i, x.indices())
  self.assertEqual(exp_v, x.values())
 
@@ -193,7 +193,7 @@ def _test_contig(self, is_cuda):
  ])
  exp_v = ValueTensor([6, 4])
 
- x.coalesce_()
+ x = x.coalesce()
  self.assertEqual(exp_i, x.indices())
  self.assertEqual(exp_v, x.values())
 
@@ -224,7 +224,7 @@ def _test_contig_hybrid(self, is_cuda):
  [2, 3], [1, 2], [6, 7], [4, 5], [10, 11],
  [3, 4], [5, 6], [9, 10], [8, 9], [7, 8],
  ])
- x.coalesce_()
+ x = x.coalesce()
  self.assertEqual(exp_i, x.indices())
  self.assertEqual(exp_v, x.values())
 
@@ -242,7 +242,7 @@ def _test_contig_hybrid(self, is_cuda):
  ])
  exp_v = ValueTensor([[2, 2, 2], [1, 1, 1], [3, 3, 3], [4, 4, 4]])
 
- x.coalesce_()
+ x = x.coalesce()
  self.assertEqual(exp_i, x.indices())
  self.assertEqual(exp_v, x.values())
 
@@ -261,7 +261,7 @@ def _test_contig_hybrid(self, is_cuda):
  ])
  exp_v = ValueTensor([[6, 4, 5], [4, 3, 4]])
 
- x.coalesce_()
+ x = x.coalesce()
  self.assertEqual(exp_i, x.indices())
  self.assertEqual(exp_v, x.values())
 
@@ -490,6 +490,16 @@ def _test_basic_ops_shape(self, is_cuda, shape_i, shape_v=None):
  expected = torch.zeros(x1.size())
  self.assertEqual(y.to_dense(), expected)
 
+ self.assertFalse(x1.is_coalesced())
+ y = x1.coalesce()
+ z = x1.coalesce()
+ self.assertFalse(x1.is_coalesced())
+ self.assertTrue(y.is_coalesced())
+ self.assertEqual(x1, y)
+ # check that coalesce is out of place
+ y.values().add_(1)
+ self.assertEqual(z.values() + 1, y.values())
+
  def _test_basic_ops(self, is_cuda):
  self._test_basic_ops_shape(is_cuda, [5, 6])
  self._test_basic_ops_shape(is_cuda, [10, 10, 10])

torch/csrc/generic/methods/SparseTensor.cwrap

@@ -90,9 +90,9 @@ PyObject * THSPTensor_(size)(PyObject *self, PyObject *args, PyObject *kwargs)
 
 [[
  name: coalesce
- python_name: coalesce_
+ cname: newCoalesce
  sparse: yes
- return: argument 0
+ return: THSTensor*
  arguments:
  - THSTensor* self
 ]]

torch/lib/THCS/generic/THCSTensor.cu

@@ -35,10 +35,16 @@ THCTensor *THCSTensor_(toDense)(THCState *state, THCSTensor *self) {
  return dst;
 }
 
-void THCSTensor_(coalesce)(THCState *state, THCSTensor *self) {
- if (self->coalesced) return;
- int nnz = self->nnz;
- if (nnz < 2) return;
+THCSTensor *THCSTensor_(newCoalesce)(THCState *state, THCSTensor *self) {
+ ptrdiff_t nnz = self->nnz;
+ if (nnz < 2) {
+ self->coalesced = 1;
+ }
+ if (self->coalesced) {
+ THCSTensor_(retain)(state, self);
+ return self;
+ }
+
 #if CUDA_VERSION >= 7000
  THCThrustAllocator thrustAlloc(state);
 #define THRUST_EXEC(fn, ...) fn(thrust::cuda::par(thrustAlloc).on(THCState_getCurrentStream(state)), ##__VA_ARGS__)
@@ -93,10 +99,10 @@ void THCSTensor_(coalesce)(THCState *state, THCSTensor *self) {
  long newNnz = newEnd.first - indicesIter;
 
  THCIndexTensor_(resize2d)(state, indices1D, 1, newNnz);
- THLongStorage *newValuesSize = THCTensor_(newSizeOf)(state, values);
- newValuesSize->data[0] = newNnz;
- THCTensor *newValues = THCTensor_(newWithSize)(state, newValuesSize, NULL);
- THLongStorage_free(newValuesSize);
+ THCTensor *newValues = THCTensor_(new)(state);
+ THCTensor_(resizeNd)(state, newValues, values->nDimension, values->size, NULL);
+ newValues->size[0] = newNnz;
+
 
  dim3 grid(THCCeilDiv(newNnz, (long) 4), THCCeilDiv(stride, (long) 128));
  dim3 block(32, 4);
@@ -152,16 +158,16 @@ void THCSTensor_(coalesce)(THCState *state, THCSTensor *self) {
  THCIndexTensor_(free)(state, indicesSlice);
  }
  ////////////////////////////////////////////////////////////
- self->nnz = newNnz;
- THCIndexTensor_(free)(state, self->indices);
- self->indices = newIndices;
+ THLongStorage *size = THCSTensor_(newSizeOf)(state, self);
+ THCSTensor *dst = THCSTensor_(newWithTensorAndSize)(state, newIndices, newValues, size);
+ THLongStorage_free(size);
 
+ THCIndexTensor_(free)(state, indices);
  THCTensor_(free)(state, values);
- THCTensor_(free)(state, self->values);
- self->values = newValues;
 
- self->coalesced = 1;
+ dst->coalesced = 1;
  THCudaCheck(cudaGetLastError());
+ return dst;
 #undef THRUST_EXEC
 }
 

torch/lib/THCS/generic/THCSTensor.h

@@ -58,7 +58,7 @@ TH_API void THCSTensor_(copy)(THCState *state, THCSTensor *self, THCSTensor *src
 
 TH_API void THCSTensor_(transpose)(THCState *state, THCSTensor *self, int dimension1_, int dimension2_);
 TH_API int THCSTensor_(isCoalesced)(THCState *state, const THCSTensor *self);
-TH_API void THCSTensor_(coalesce)(THCState *state, THCSTensor *self);
+TH_API THCSTensor *THCSTensor_(newCoalesce)(THCState *state, THCSTensor *self);
 
 TH_API void THCTensor_(sparseMask)(THCState *state, THCSTensor *r_, THCTensor *t, THCSTensor *mask);
 

torch/lib/THCS/generic/THCSTensorMath.cu

@@ -40,22 +40,22 @@ void THCTensor_(spaddcdiv)(THCState *state, THCTensor *r_, THCTensor *t, real va
  THError("WARNING: Sparse Cuda Tensor op spaddcdiv is not implemented");
 }
 
-void THCSTensor_(spaddmm)(THCState *state, THCTensor *r_, real beta, THCTensor *t, real alpha, THCSTensor *sparse, THCTensor *dense) {
+void THCSTensor_(spaddmm)(THCState *state, THCTensor *r_, real beta, THCTensor *t, real alpha, THCSTensor *sparse_, THCTensor *dense) {
 #if defined(THCS_REAL_IS_FLOAT) || defined(THCS_REAL_IS_DOUBLE)
- THCAssertSameGPU(THCSTensor_(checkGPU)(state, 1, 4, sparse, r_, t, dense));
+ THCAssertSameGPU(THCSTensor_(checkGPU)(state, 1, 4, sparse_, r_, t, dense));
  THCudaIntTensor *csr;
  THCIndexTensor *indices;
  THCTensor *values, *r__, *dense_;
 
- THArgCheck(sparse->nDimensionI == 2, 2,
- "matrices expected, got %dD tensor", sparse->nDimensionI);
- THArgCheck(sparse->nDimensionV == 0, 2,
- "scalar values expected, got %dD values", sparse->nDimensionV);
+ THArgCheck(sparse_->nDimensionI == 2, 2,
+ "matrices expected, got %dD tensor", sparse_->nDimensionI);
+ THArgCheck(sparse_->nDimensionV == 0, 2,
+ "scalar values expected, got %dD values", sparse_->nDimensionV);
  THArgCheck(dense->nDimension == 2, 2,
  "matrices expected, got %dD tensor", dense->nDimension);
 
- long m = THCSTensor_(size)(state, sparse, 0);
- long k = THCSTensor_(size)(state, sparse, 1);
+ long m = THCSTensor_(size)(state, sparse_, 0);
+ long k = THCSTensor_(size)(state, sparse_, 1);
  long n = THCTensor_(size)(state, dense, 1);
 
  THCTensor_(resize2d)(state, r_, m, n);
@@ -67,7 +67,7 @@ void THCSTensor_(spaddmm)(THCState *state, THCTensor *r_, real beta, THCTensor *
  THArgCheck(THCTensor_(size)(state, dense, 0) == k, 3,
  "Expected dim 0 size %d, got %d", k, THCTensor_(size)(state, dense, 0));
 
- THCSTensor_(coalesce)(state, sparse);
+ THCSTensor *sparse = THCSTensor_(newCoalesce)(state, sparse_);
 
  long nnz = THCSTensor_(nnz)(state, sparse);
  indices = THCSTensor_(newIndices)(state, sparse);
@@ -146,6 +146,7 @@ void THCSTensor_(spaddmm)(THCState *state, THCTensor *r_, real beta, THCTensor *
  THCIndexTensor_(free)(state, rowIndices);
  THCIndexTensor_(free)(state, colIndices);
  THCTensor_(free)(state, values);
+ THCSTensor_(free)(state, sparse);
 #else
  THError("unimplemented data type");
 #endif
@@ -156,40 +157,42 @@ void THCSTensor_(sspaddmm)(THCState *state, THCSTensor *r_, real beta, THCSTenso
  // TODO Write some kernels
 }
 
-void THCSTensor_(hspmm)(THCState *state, THCSTensor *r_, real alpha, THCSTensor *sparse, THCTensor *dense) {
+void THCSTensor_(hspmm)(THCState *state, THCSTensor *r_, real alpha, THCSTensor *sparse_, THCTensor *dense) {
 #if CUDA_VERSION >= 7000
  THCThrustAllocator thrustAlloc(state);
 #define THRUST_EXEC(fn, ...) fn(thrust::cuda::par(thrustAlloc).on(THCState_getCurrentStream(state)), ##__VA_ARGS__)
 #else
 #define THRUST_EXEC(fn, ...) fn(##__VA_ARGS__)
 #endif
 
- THCAssertSameGPU(THCSTensor_(checkGPU)(state, 2, 3, r_, sparse, dense));
+ THCAssertSameGPU(THCSTensor_(checkGPU)(state, 2, 3, r_, sparse_, dense));
 
- THArgCheck(sparse->nDimensionI == 2, 3,
- "matrices expected, got %dD tensor", sparse->nDimensionI);
- THArgCheck(sparse->nDimensionV == 0, 3,
- "scalar values expected, got %dD values", sparse->nDimensionV);
+ THArgCheck(sparse_->nDimensionI == 2, 3,
+ "matrices expected, got %dD tensor", sparse_->nDimensionI);
+ THArgCheck(sparse_->nDimensionV == 0, 3,
+ "scalar values expected, got %dD values", sparse_->nDimensionV);
  THArgCheck(dense->nDimension == 2, 4,
  "matrices expected, got %dD tensor", dense->nDimension);
 
- long m = THCSTensor_(size)(state, sparse, 0);
- long k = THCSTensor_(size)(state, sparse, 1);
+ long m = THCSTensor_(size)(state, sparse_, 0);
+ long k = THCSTensor_(size)(state, sparse_, 1);
  long n = THCTensor_(size)(state, dense, 1);
 
  THArgCheck(THCTensor_(size)(state, dense, 0) == k, 4,
  "Expected dim 0 size %d, got %d", k, THCTensor_(size)(state, dense, 0));
  long size[2] = {m, n};
  THCSTensor_(rawResize)(state, r_, 1, 1, size);
 
- THCSTensor_(coalesce)(state, sparse);
+ THCSTensor *sparse = THCSTensor_(newCoalesce)(state, sparse_);
 
  long nnz = THCSTensor_(nnz)(state, sparse);
  THCIndexTensor *indices = THCIndexTensor_(newWithSize2d)(state, 1, nnz);
  // create values in column-major format to avoid copying in spaddmm
  THCTensor *values = THCTensor_(newWithSize2d)(state, n, nnz);
  THCTensor_(transpose)(state, values, NULL, 0, 1);
 
+ // why does sparse need to be cloned? If this is really necessary maybe we
+ // need to fuse this with newCoalesce
  THCSTensor *newSparse = THCSTensor_(newClone)(state, sparse);
  THCIndexTensor *spIndices = THCSTensor_(newIndices)(state, newSparse);
  THCIndexTensor *dstIndices = THCIndexTensor_(newSelect)(state, spIndices, 0, 0);
@@ -206,6 +209,7 @@ void THCSTensor_(hspmm)(THCState *state, THCSTensor *r_, real alpha, THCSTensor
  THCSTensor_(free)(state, newSparse);
  THCIndexTensor_(free)(state, spIndices);
  THCIndexTensor_(free)(state, dstIndices);
+ THCSTensor_(free)(state, sparse);
 
 #undef THRUST_EXEC
 }
@@ -348,8 +352,6 @@ void THCSTensor_(cadd)(THCState *state, THCSTensor *r_, THCSTensor *t, real valu
  if(!THCSTensor_(isSameSizeAs)(state, t, src)) {
  THError("cadd operands have incompatible sizes or dimension types");
  }
- THCSTensor_(coalesce)(state, t);
- THCSTensor_(coalesce)(state, src);
 
  if (src->nnz == 0) {
  THCSTensor_(copy)(state, r_, t);
@@ -399,13 +401,13 @@ void THCSTensor_(csub)(THCState *state, THCSTensor *r_, THCSTensor *t, real valu
  THCSTensor_(cadd)(state, r_, t, ScalarNegate<real>::to(value), src);
 }
 
-void THCSTensor_(cmul)(THCState *state, THCSTensor *r_, THCSTensor *t, THCSTensor *src) {
- THCAssertSameGPU(THCSTensor_(checkGPU)(state, 3, 3, r_, t, src));
- if(!THCSTensor_(isSameSizeAs)(state, t, src)) {
+void THCSTensor_(cmul)(THCState *state, THCSTensor *r_, THCSTensor *t_, THCSTensor *src_) {
+ THCAssertSameGPU(THCSTensor_(checkGPU)(state, 3, 3, r_, t_, src_));
+ if(!THCSTensor_(isSameSizeAs)(state, t_, src_)) {
  THError("cmul operands have incompatible sizes or dimension types");
  }
- THCSTensor_(coalesce)(state, t);
- THCSTensor_(coalesce)(state, src);
+ THCSTensor *t = THCSTensor_(newCoalesce)(state, t_);
+ THCSTensor *src = THCSTensor_(newCoalesce)(state, src_);
 
  if (t->nnz == 0 || src->nnz == 0) {
  THCSTensor_(zero)(state, r_);
@@ -453,6 +455,8 @@ void THCSTensor_(cmul)(THCState *state, THCSTensor *r_, THCSTensor *t, THCSTenso
  THCTensor_(free)(state, t_values_);
  THCIndexTensor_(free)(state, s_indices_);
  THCTensor_(free)(state, s_values_);
+ THCSTensor_(free)(state, t);
+ THCSTensor_(free)(state, src);
 }
 
 #if defined(THCS_REAL_IS_FLOAT) || defined(THCS_REAL_IS_DOUBLE)