Add preamble for atomicAdd
@@ -153,10 +153,26 @@ class _compressed_sparse_matrix(sparse_data._data_matrix):
start = self.indptr[major]
end = self.indptr[major + 1]
answer = cupy.zeros((), self.dtype)
+ preamble = '''
+#if __CUDA_ARCH__ < 600
+__device__ double atomicAdd(double* address, double val) {
+ unsigned long long int* address_as_ull =
+ (unsigned long long int*)address;
+ unsigned long long int old = *address_as_ull, assumed;
+ do {
+ assumed = old;
+ old = atomicCAS(address_as_ull, assumed,
+ __double_as_longlong(val +
+ __longlong_as_double(assumed)));
+ } while (assumed != old);
+ return __longlong_as_double(old);
+}
+#endif
+'''
kern = cupy.ElementwiseKernel(
'T d, S ind, int32 minor', 'raw T answer',
'if (ind == minor) atomicAdd(&answer[0], d);',
- 'compress_getitem')
+ 'compress_getitem', preamble=preamble)
kern(self.data[start:end], self.indices[start:end], minor, answer)
return answer[()]
