#define TORCH_ASSERT_ONLY_METHOD_OPERATORS #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #ifndef AT_PER_OPERATOR_HEADERS #include #include #include #else #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #endif #include namespace at { namespace meta { TORCH_META_FUNC(_convert_indices_from_coo_to_csr) (const Tensor& self, const int64_t size, const bool out_int32) { TORCH_CHECK(self.dim() <= 1, "Input is supposed to be a vector, but got ", self.dim(), " dimensional tensor."); ScalarType scalar_type = out_int32 ? ScalarType::Int : ScalarType::Long; c10::TensorOptions options = TensorOptions().device(self.options().device()).dtype(scalar_type); set_output_raw_strided(0, size + 1, {}, options); } TORCH_META_FUNC(_convert_indices_from_csr_to_coo) (const Tensor& crow_indices, const Tensor& col_indices, const bool out_int32, const bool transpose) { TORCH_CHECK( crow_indices.dim() == col_indices.dim(), "crow_indices and col_indices are supposed to have" " the same dimensionality, but got ", crow_indices.dim(), " and ", crow_indices.dim(), " dimensional tensors, respectively."); ScalarType scalar_type = out_int32 ? ScalarType::Int : ScalarType::Long; c10::TensorOptions options = crow_indices.options().dtype(scalar_type); set_output_raw_strided(0, {col_indices.dim() + 1, col_indices.numel()}, {}, options, {}); } } // namespace meta namespace { template Tensor& unary_op_out(F op_out, const Tensor& self, Tensor& result) { TORCH_INTERNAL_ASSERT(self.is_sparse_csr()); TORCH_INTERNAL_ASSERT(result.is_sparse_csr()); if (!result.is_same(self)) { // For the case of (0x0) result tensor, manually resize `result` tensor // to the size of `self` tensor if (result.numel() == 0) { at::native::resize_as_sparse_compressed_(result, self); } // copy_sparse_compressed_ internally checks the sizes of result and self tensors // Hence no external size check required at::native::copy_sparse_compressed_(result, self); } auto self_values = self.values(); auto result_values = result.values(); op_out(self_values, result_values); return result; } template Tensor& unary_op_inplace(Tensor& self, const F& op_inplace, Args&&... args) { AT_DISPATCH_ALL_SPARSE_COMPRESSED_LAYOUTS(self.layout(), "unary_op_inplace", [](){}); auto self_values = self.values(); (self_values.*op_inplace)(std::forward(args)...); return self; } } // end anonymous namespace namespace native { using namespace at::sparse_csr; // certain utility functions are usable from sparse COO. using namespace at::sparse; Tensor& mul_out_sparse_csr(const Tensor& t_, const Tensor& src_, Tensor& r) { // // TODO: Use a specialized CSR kernel for performance if needed if (t_.is_sparse_csr() && src_.layout() == kStrided) { return mul_out_sparse_csr(t_, src_.sparse_mask(t_), r); } if (t_.layout() == kStrided && src_.is_sparse_csr()) { return mul_out_sparse_csr(t_.sparse_mask(src_), src_, r); } TORCH_CHECK(r.is_sparse_csr(), "Expected result Tensor to be of format CSR"); Tensor t = t_.to_sparse(); Tensor src = src_.to_sparse(); Tensor tmp_result = t.mul(src); auto r_sparse_csr = tmp_result.to_sparse_csr(); r.resize_as_sparse_(r_sparse_csr); r.copy_(r_sparse_csr); return r; } template Tensor intersection_binary_op_with_wrapped_scalar(const Tensor& sparse, const Tensor& scalar, const op_t& op) { // NOTE: intersection_binary_op_with_wrapped_scalar assumes scalar.numel() == 1. const auto result_values = op(sparse.values(), scalar.squeeze()).to(at::result_type(sparse, scalar)); const auto result_sizes = infer_size(sparse.sizes(), scalar.sizes()); auto [compressed_indices, plain_indices] = getCompressedPlainIndices(sparse); return at::_sparse_compressed_tensor_unsafe( compressed_indices.clone(), plain_indices.clone(), result_values, result_sizes, sparse.options().dtype(result_values.scalar_type())); } template Tensor& intersection_binary_op_with_wrapped_scalar_(Tensor& sparse, const Tensor& scalar, const string& op_name, const op_t& op) { // NOTE: intersection_binary_op_with_wrapped_scalar_ assumes scalar.numel() == 1. const auto broadcasted_shape = infer_size(sparse.sizes(), scalar.sizes()); if (sparse.sizes() != broadcasted_shape) { TORCH_CHECK(false, op_name, "(): output with shape ", sparse.sizes(), " does not match ", "the broadcast shape ", broadcasted_shape); } auto values = sparse.values(); // Safe to use squeeze here, we already know that scalar safely broadcasts. op(values, scalar.squeeze()); return sparse; } Tensor mul_sparse_csr(const Tensor& self, const Tensor& other) { // Check if either of the arguments is a wrapped Scalar if (self.layout() == kStrided && self.dim() == 0) { return intersection_binary_op_with_wrapped_scalar(other, self, [](const Tensor& a, const Tensor& b) -> Tensor { return a.mul(b); }); } if (other.layout() == kStrided && other.dim() == 0) { return intersection_binary_op_with_wrapped_scalar(self, other, [](const Tensor& a, const Tensor& b) -> Tensor { return a.mul(b); }); } if (self.is_sparse_csr() && other.layout() == kStrided) { return mul_sparse_csr(self, other.sparse_mask(self)); } if (self.layout() == kStrided && other.is_sparse_csr()) { return mul_sparse_csr(self.sparse_mask(other), other); } auto commonDtype = at::result_type(self, other); auto result_options = self.options().dtype(commonDtype); // CSR is 2d! Tensor result = at::empty({0, 0}, result_options); return at::mul_out(result, self, other); // redispatch! } Tensor& mul_sparse_csr_(Tensor& self, const Tensor& other) { if (other.layout() == kStrided && other.dim() == 0) { return intersection_binary_op_with_wrapped_scalar_(self, other, "mul_", [](Tensor& a, const Tensor& b) -> Tensor& { return a.mul_(b); }); } return at::mul_out(self, self, other); // redispatch! } namespace { template inline Tensor get_result_tensor_for_unary_op(F op, const Tensor& input) { auto values = input.values(); // To handle type promotion for inputs to unary ops, // we first get the result from the underlined op, and use the result // to create a sparse compressed tensor, which is used as the input to the out= // variant auto result_values = op(values); auto compressed_indices = AT_DISPATCH_ROW_SPARSE_COMPRESSED_LAYOUTS(input.layout(), "get_result_tensor_for_unary_op", [&]{ return input.crow_indices(); }, [&]{ return input.ccol_indices(); }); auto plain_indices = AT_DISPATCH_ROW_SPARSE_COMPRESSED_LAYOUTS(input.layout(), "get_result_tensor_for_unary_op", [&]{ return input.col_indices(); }, [&]{ return input.row_indices(); }); auto result = at::_sparse_compressed_tensor_unsafe( compressed_indices.clone(), plain_indices.clone(), result_values, input.sizes(), input.options().dtype(result_values.scalar_type())); return result; } } // namespace Tensor& normal_sparse_csr_( Tensor& self, double mean, double std, std::optional gen) { return unary_op_inplace(self, &Tensor::normal_, mean, std, gen); } Tensor& fill_sparse_csr_(Tensor& self, const Scalar& value) { return unary_op_inplace(self, &TensorBase::fill_, value); } Tensor sparse_mask_sparse_compressed( const Tensor& self, const Tensor& mask) { TORCH_CHECK(at::sparse_csr::is_sparse_compressed(mask), "sparse_mask_sparse_compressed expects mask to have sparse compressed layout, got ", mask.layout()); TORCH_CHECK( mask.sizes().equals(self.sizes()), "sparse_mask(): operands have incompatible sizes; self has size ", self.sizes(), " but mask has size ", mask.sizes()); if (self.is_same(mask)) { return self; } if (!mask.numel() || !mask._nnz()) { return mask.clone().to(self.device(), self.scalar_type()); } if (self.layout() == kStrided) { auto [compressed_indices, plain_indices] = at::sparse_csr::getCompressedPlainIndices(mask); auto mask_values = mask.values(); auto dense_mask = at::_sparse_compressed_tensor_unsafe( compressed_indices, plain_indices, at::ones({1}, self.options().dtype(kBool)).expand_as(mask_values), self.sizes(), self.options().dtype(kBool).layout(mask.layout())).to_dense(); return AT_DISPATCH_PLAIN_SPARSE_COMPRESSED_LAYOUTS( mask.layout(), "sparse_mask_sparse_compressed", [&] { return at::native::dense_to_sparse_with_mask(self, dense_mask, mask.layout(), {}, mask.dense_dim()); }, [&] { auto blocksize = at::sparse_csr::getBlockSize(mask); return at::native::dense_to_sparse_with_mask(self, dense_mask, mask.layout(), blocksize, mask.dense_dim()); }); } else if (self.layout() == mask.layout()) { // TODO: keeping this for BC but the method used here may lead to // incorrect indices. return self.mul(at::ones_like(mask)).to(self.scalar_type()); } else { // TODO: keeping this for BC but the method used here cannot // support batch dimensions because sparse COO tensors are batch // dimension ignorant. return AT_DISPATCH_PLAIN_SPARSE_COMPRESSED_LAYOUTS( mask.layout(), "sparse_mask_sparse_compressed", [&] { return self.sparse_mask(mask.to_sparse()).to_sparse(mask.layout()); }, [&] { auto blocksize = at::sparse_csr::getBlockSize(mask); return self.sparse_mask(mask.to_sparse()).to_sparse(mask.layout(), blocksize); }); } } Tensor mul_scalar_sparse_csr(const Tensor& self, const Scalar& other) { auto result_values = self.values().mul(other); return at::native::_sparse_csr_tensor_unsafe( self.crow_indices().clone(), self.col_indices().clone(), result_values, self.sizes(), result_values.scalar_type(), self.layout(), result_values.device()); } Tensor& zero_sparse_csr_(Tensor& self) { /* csr.zero_() resets nnz to 0. If the original sparsity pattern needs to be preserved, use `csr.values().zero_()` instead. The above behavior also implies that torch.zeros_like(csr) returns a new tensor with nnz == 0. If one needs a zeros_like semantics where the result has the same sparsity pattern as input, then use `result = csr.clone(); result.values.zero_();` */ AT_DISPATCH_ALL_SPARSE_COMPRESSED_LAYOUTS(self.layout(), "zero_sparse_csr_", [](){}); get_sparse_csr_impl(self)->resize_and_clear_(self.sparse_dim(), self.dense_dim(), self.sizes()); return self; } /* Implementation of Unary Ufuncs, those supported for Sparse CSR Layout * Only simple funcs, with 0->0 correspondence are currently supported. */ #define CREATE_UNARY_UFUNC_OUT(op_name) \ Tensor& op_name##_sparse_csr_out(const Tensor& self, Tensor& result) { \ return unary_op_out(&at::op_name##_outf, self, result); \ } #define CREATE_UNARY_UFUNC_FUNCTIONAL(op_name) \ Tensor op_name##_sparse_csr(const Tensor& self) { \ return get_result_tensor_for_unary_op(&at::op_name, self); \ } #define CREATE_UNARY_UFUNC_INPLACE(op_name) \ Tensor& op_name##_sparse_csr_(Tensor& self) { \ return unary_op_inplace(self, &Tensor::op_name##_); \ } #define CREATE_UNARY_UFUNC(op_name) \ CREATE_UNARY_UFUNC_OUT(op_name); \ CREATE_UNARY_UFUNC_FUNCTIONAL(op_name); \ CREATE_UNARY_UFUNC_INPLACE(op_name); #define CREATE_UNARY_UFUNC_NO_INPLACE(op_name) \ CREATE_UNARY_UFUNC_OUT(op_name); \ CREATE_UNARY_UFUNC_FUNCTIONAL(op_name); // Exhaustive list of the unary ufuncs supported by sparse compressed CREATE_UNARY_UFUNC(abs); CREATE_UNARY_UFUNC(asin); CREATE_UNARY_UFUNC(asinh); CREATE_UNARY_UFUNC(atan); CREATE_UNARY_UFUNC(atanh); CREATE_UNARY_UFUNC(ceil); CREATE_UNARY_UFUNC(deg2rad); CREATE_UNARY_UFUNC(erf); CREATE_UNARY_UFUNC(erfinv); CREATE_UNARY_UFUNC(expm1); CREATE_UNARY_UFUNC(floor); CREATE_UNARY_UFUNC(frac); CREATE_UNARY_UFUNC(log1p); CREATE_UNARY_UFUNC(neg); CREATE_UNARY_UFUNC(rad2deg); CREATE_UNARY_UFUNC(sign); CREATE_UNARY_UFUNC(sin); CREATE_UNARY_UFUNC(sinh); CREATE_UNARY_UFUNC(sgn); CREATE_UNARY_UFUNC(sqrt); CREATE_UNARY_UFUNC(tan); CREATE_UNARY_UFUNC(tanh); CREATE_UNARY_UFUNC(trunc); CREATE_UNARY_UFUNC(conj_physical); C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wunused-function") static CREATE_UNARY_UFUNC(relu); C10_DIAGNOSTIC_POP() // With addition of `round.decimals` overload, using CREATE_UNARY_UFUNC leads // to unresolved overload. Tensor& round_sparse_csr_out(const Tensor& self, Tensor& result) { return unary_op_out(&at::_ops::round_out::call, self, result); } Tensor round_sparse_csr(const Tensor& self) { return get_result_tensor_for_unary_op(&at::_ops::round::call, self); } Tensor& round_sparse_csr_(Tensor& self) { TORCH_INTERNAL_ASSERT(self.is_sparse_csr()); self.values().round_(); return self; } Tensor threshold_backward_sparse_compressed( const Tensor& grad_output, const Tensor& self, const Scalar& threshold) { return get_result_tensor_for_unary_op( [&](const Tensor& t) { return at::threshold_backward(t, self.values(), threshold); }, grad_output); } Tensor& threshold_backward_sparse_compressed_out( const Tensor& grad_output, const Tensor& self, const Scalar& threshold, Tensor& grad_input) { return unary_op_out( [&](const Tensor& t, Tensor& out) { return at::threshold_backward_outf(t, self.values(), threshold, out); }, grad_output, grad_input); } // angle, isneginf, isposinf and signbit currently don't have an inplace variant CREATE_UNARY_UFUNC_NO_INPLACE(angle); CREATE_UNARY_UFUNC_NO_INPLACE(isneginf); CREATE_UNARY_UFUNC_NO_INPLACE(isposinf); CREATE_UNARY_UFUNC_NO_INPLACE(signbit); // isnan and isinf don't have an out variant CREATE_UNARY_UFUNC_FUNCTIONAL(isnan); CREATE_UNARY_UFUNC_FUNCTIONAL(isinf); template void addmm_out_sparse_csr_native_cpu( const Tensor& sparse, const Tensor& dense, const Tensor& r, Scalar alpha, Scalar beta) { auto dim_i = sparse.size(0); auto dim_k = dense.size(1); auto csr = sparse.crow_indices(); auto col_indices = sparse.col_indices(); auto values = sparse.values(); scalar_t cast_alpha = alpha.to(); r.mul_(beta); AT_DISPATCH_INDEX_TYPES( col_indices.scalar_type(), "csr_mm_crow_indices", [&]() { auto csr_accessor = csr.accessor(); auto col_indices_accessor = col_indices.accessor(); auto values_accessor = values.accessor(); scalar_t* dense_ptr = dense.data_ptr(); scalar_t* r_ptr = r.data_ptr(); int64_t dense_stride0 = dense.stride(0); int64_t dense_stride1 = dense.stride(1); int64_t r_stride0 = r.stride(0); int64_t r_stride1 = r.stride(1); at::parallel_for( 0, dim_i, internal::GRAIN_SIZE, [&](int64_t irow_start, int64_t irow_end) { for (index_t h = irow_start; h < irow_end; ++h) { index_t i_start = csr_accessor[h]; index_t i_end = csr_accessor[h + 1]; for (index_t i = i_start; i < i_end; i++) { scalar_t val = values_accessor[i]; index_t col = col_indices_accessor[i]; at::native::cpublas::axpy( dim_k, cast_alpha * val, dense_ptr + col * dense_stride0, dense_stride1, r_ptr + h * r_stride0, r_stride1); } } }); }); } // Functions for matrix multiplication. // result = beta * self + alpha (mat1 @ mat2) Tensor& addmm_out_sparse_compressed_cpu( const Tensor& self, const Tensor& mat1, const Tensor& mat2, const Scalar& beta, const Scalar& alpha, Tensor& result) { // All the checks are from addmm_out_cuda_impl (ATen/native/cuda/Blas.cpp) and // TORCH_META_FUNC(addmm) (ATen/native/LinearAlgebra.cpp) // TODO: remove code duplication and unify code sparse::impl::_check_dim(mat1, 2, "mat1"); sparse::impl::_check_dim(mat2, 2, "mat2"); TORCH_CHECK( mat1.size(1) == mat2.size(0), "mat1 and mat2 shapes cannot be multiplied (", mat1.size(0), "x", mat1.size(1), " and ", mat2.sizes()[0], "x", mat2.sizes()[1], ")"); c10::MaybeOwned self_; // Don't expand self if this is an in-place operation if (&result == &self) { self_ = c10::MaybeOwned::borrowed(self); } else { self_ = expand_size(self, {mat1.size(0), mat2.size(1)}, "addmm"); } TORCH_CHECK(((self_->dim() == 2) && (self_->size(0) == mat1.size(0)) && (self_->size(1) == mat2.size(1))), "The input tensor must be a matrix with size ", mat1.size(0), "x", mat2.size(1), ", but got a ", self_->dim(), "-D tensor with size ", self_->size(0), "x", self_->size(1)); if (&result != &self) { if (result.layout() == kStrided) { at::native::resize_output(result, self_->sizes()); } else { result.resize_as_sparse_(*self_); } result.copy_(*self_); } if (result.numel() == 0) { // If result gets resized and is sparse compressed, // it's compressed_indices tensor will contain junk values // so the whole tensor is not a valid compressed tensor. // To combat that, result needs to get zeroed out. if (at::sparse_csr::is_sparse_compressed(result)) { result.zero_(); } return result; } if (sparse::impl::_is_sparse_and_zero(mat1) || sparse::impl::_is_sparse_and_zero(mat2)) { // According to docs, when beta==0 values in self should be ignored. // nans and infs should not propagate if (beta.toComplexDouble() == 0.) { result.zero_(); } else { result.mul_(beta); } return result; } #if !AT_USE_MKL_SPARSE() // The custom impl addmm_out_sparse_csr_native_cpu only supports CSR @ // strided -> strided if (mat1.layout() == kStrided) { if (mat2.layout() == kSparseCsr) { if (result.layout() == kStrided) { AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES( result.scalar_type(), "addmm_sparse_dense", [&] { addmm_out_sparse_csr_native_cpu( mat2.transpose(-2, -1).to_sparse_csr(), mat1.transpose(-2, -1), result.transpose(-2, -1), alpha, beta); }); return result; } } if (mat2.layout() == kSparseCsc) { if (result.layout() == kStrided) { AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES( result.scalar_type(), "addmm_sparse_dense", [&] { addmm_out_sparse_csr_native_cpu( mat2.transpose(-2, -1), mat1.transpose(-2, -1), result.transpose(-2, -1), alpha, beta); }); return result; } } } else if (mat1.layout() == kSparseCsr) { if (mat2.layout() == kStrided) { if (result.layout() == kStrided) { AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES( result.scalar_type(), "addmm_sparse_dense", [&] { addmm_out_sparse_csr_native_cpu( mat1, mat2, result, alpha, beta); }); return result; } } } else if (mat1.layout() == kSparseCsc) { if (mat2.layout() == kStrided) { if (result.layout() == kStrided) { AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES( result.scalar_type(), "addmm_sparse_dense", [&] { addmm_out_sparse_csr_native_cpu( mat1.to_sparse_csr(), mat2, result, alpha, beta); }); return result; } } } TORCH_CHECK( false, "addmm: computation on CPU is not implemented for ", result.layout(), " + ", mat1.layout(), " @ ", mat2.layout(), " without MKL. PyTorch built with MKL has better support for addmm with sparse CPU tensors."); #else sparse::impl::mkl::addmm_out_sparse_csr(mat1, mat2, beta, alpha, result); #endif return result; } Tensor addmm_sparse_compressed_dense( const Tensor& self, const SparseCsrTensor& sparse, const Tensor& dense, const Scalar& beta, const Scalar& alpha) { Tensor r = at::empty({0, 0}, self.options()); at::addmm_out(r, self, sparse, dense, beta, alpha); return r; } Tensor& _sparse_csr_mm_out( const Tensor& mat1, const Tensor& mat2, Tensor& result) { auto zero = at::zeros_like(result); return at::addmm_out(result, zero, mat1, mat2, 0.0, 1.0); } Tensor _sparse_csr_mm(const Tensor& mat1, const Tensor& mat2) { if (mat1.is_sparse_csr() && mat2.is_sparse_csr()) { // Return sparse return at::addmm( at::zeros({mat1.size(0), mat2.size(1)}, mat2.options()), mat1, mat2, 0.0, 1.0); } if ((mat1.layout() == kSparseCsc || mat1.layout() == kSparseCsr) && (mat2.layout() == kSparseCsc || mat2.layout() == kSparseCsr)) { // TODO: Expensive conversion to CSR. Should add native support for CSC. // Covers CSC @ CSR // Covers CSR @ CSC // Covers CSC @ CSC return _sparse_csr_mm(mat1.to_sparse_csr(), mat2.to_sparse_csr()); } if (mat1.layout() == kSparseCsc && mat2.layout() == c10::kStrided) { // TODO: This is a costly conversion. We should have // native support for CSC. return _sparse_csr_mm(mat1.to_sparse_csr(), mat2); } // Default to taking options from mat1 auto result_options = mat1.options(); if (mat2.layout() == kStrided) { // if either arg is strided we return strided, so update the options if // mat2 is strided. result_options = result_options.layout(kStrided); } return at::addmm( at::zeros({mat1.size(0), mat2.size(1)}, result_options), mat1, mat2, 0.0, 1.0); } // Functions for element-wise addition. Tensor add_sparse_csr( const Tensor& self, const Tensor& other, const Scalar& alpha) { auto commonDtype = at::result_type(self, other); alpha_check(commonDtype, alpha); Tensor result; if (self.layout() != kStrided && other.layout() == kStrided) { // add(sparse, dense) -> dense result = at::empty_like( other, other.options() .dtype(commonDtype) .memory_format(at::MemoryFormat::Contiguous)); } else { // add(dense, sparse) -> dense AND add(sparse, sparse) -> sparse result = at::empty_like( self, self.options() .dtype(commonDtype) .memory_format(at::MemoryFormat::Contiguous)); } return at::add_out(result, self, other, alpha); // redispatch! } Tensor& add_sparse_csr_( Tensor& self, const Tensor& other, const Scalar& alpha) { return at::add_out(self, self, other, alpha); // redispatch! } static void add_out_dense_sparse_compressed_cpu( const Tensor& out, const Tensor& dense, const SparseCsrTensor& src, const Scalar& alpha) { TORCH_INTERNAL_ASSERT(dense.layout() == kStrided); TORCH_INTERNAL_ASSERT( src.layout() == kSparseCsr || src.layout() == kSparseCsc); TORCH_INTERNAL_ASSERT(dense.device() == kCPU || dense.device() == kMeta); TORCH_CHECK( out.is_contiguous(), "out argument must be contiguous, but got: ", out.suggest_memory_format()); TORCH_CHECK( out.device() == dense.device(), "add: expected 'out' to match dense tensor, but got tensor on device: ", out.device()); TORCH_CHECK( src.device() == dense.device(), "add: expected 'src' to match dense tensor, but got tensor on device: ", src.device()); TORCH_CHECK( dense.sizes().equals(src.sizes()), "add: expected 'self' and 'other' to have same size, but self has size ", dense.sizes(), " while other has size ", src.sizes(), " (FYI: op2-sparse addition does not currently support broadcasting)"); auto commonDtype = promoteTypes(dense.scalar_type(), src.scalar_type()); TORCH_CHECK( canCast(commonDtype, out.scalar_type()), "Can't convert result type ", commonDtype, " to output ", out.scalar_type(), " in add operation"); auto src_values = src.values(); resize_output(out, dense.sizes()); Tensor resultBuffer = out; if (out.scalar_type() != commonDtype) { resultBuffer = dense.to(commonDtype); } else if (!is_same_tensor(out, dense)) { resultBuffer.copy_(dense); } if (src._nnz() == 0) { return; } TORCH_INTERNAL_ASSERT(dense.device() == kCPU); auto valuesBuffer = src_values.to(commonDtype).reshape({-1, src_values.size(-1)}); resultBuffer = resultBuffer.view({-1, out.size(-2), out.size(-1)}); Tensor src_compressed_indices; Tensor src_plain_indices; std::tie(src_compressed_indices, src_plain_indices) = at::sparse_csr::getCompressedPlainIndices(src); src_compressed_indices = src_compressed_indices.reshape({-1, src_compressed_indices.size(-1)}); src_plain_indices = src_plain_indices.reshape({-1, src_plain_indices.size(-1)}); auto src_layout = src.layout(); AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND4( kComplexHalf, kHalf, kBool, kBFloat16, commonDtype, "add_out_op2_sparse_csr", [&valuesBuffer, &resultBuffer, &alpha, &src_compressed_indices, &src_plain_indices, &src_layout]() { AT_DISPATCH_INDEX_TYPES( src_compressed_indices.scalar_type(), "csr_add_out_crow_indices", [&valuesBuffer, &resultBuffer, &alpha, &src_compressed_indices, &src_plain_indices, &src_layout]() { auto batch_count = resultBuffer.dim() > 2 ? resultBuffer.size(-3) : 1; auto values_accessor = valuesBuffer.accessor(); scalar_t* out_ptr = resultBuffer.data_ptr(); scalar_t cast_value = alpha.to(); auto compressed_indices_accessor = src_compressed_indices.accessor(); auto plain_indices_accessor = src_plain_indices.accessor(); auto out_strides = resultBuffer.strides(); auto const out_stride_batch = out_strides[0]; auto const out_stride_compressed = AT_DISPATCH_ROW_SPARSE_COMPRESSED_LAYOUTS( src_layout, "add_out_dense_sparse_compressed_cpu", [&out_strides] { return out_strides[1]; }, [&out_strides] { return out_strides[2]; }); auto const out_stride_plain = AT_DISPATCH_ROW_SPARSE_COMPRESSED_LAYOUTS( src_layout, "add_out_dense_sparse_compressed_cpu", [&out_strides] { return out_strides[2]; }, [&out_strides] { return out_strides[1]; }); for (const auto batch_idx : c10::irange(batch_count)) { for (const auto i_compressed : c10::irange(src_compressed_indices.size(-1) - 1)) { index_t start_index = compressed_indices_accessor[batch_idx][i_compressed]; index_t end_index = compressed_indices_accessor[batch_idx][i_compressed + 1]; for (const auto i : c10::irange(start_index, end_index)) { auto i_plain = plain_indices_accessor[batch_idx][i]; auto index = batch_idx * out_stride_batch + i_compressed * out_stride_compressed + i_plain * out_stride_plain; out_ptr[index] += cast_value * values_accessor[batch_idx][i]; } } } }); }); if (out.scalar_type() != commonDtype) { out.copy_(resultBuffer); } } Tensor& add_out_sparse_compressed_cpu( const Tensor& self, const SparseCsrTensor& other, const Scalar& alpha, SparseCsrTensor& out) { if (self.layout() == kStrided) { add_out_dense_sparse_compressed_cpu(out, self, other, alpha); } else if (other.layout() == kStrided) { add_out_dense_sparse_compressed_cpu(out, other, self, alpha); } else { TORCH_CHECK( self.sizes().equals(other.sizes()), "torch.add: Expected input tensors to have the same shape, but got tensor `self` with shape ", self.sizes(), " and tensor `other` with shape ", other.sizes()); if (only_sparse_compressed_add_trivial_cases(self, other, alpha, out)) { return out; } at::native::resize_as_sparse_compressed_(out, self); sparse::impl::cpu::add_out_sparse_csr(self, other, alpha, out); } return out; } /* Reductions on sparse CSR tensors using masked semantics. - A CSR tensor is a 2D tensor that is specified by a 3-tuple (crow_indices, col_indices, values). - To support a reduction operator on a CSR tensor, define: template struct Reduction...Op { inline scalar_t operator()(const scalar_t& a, const scalar_t& b) const { return a ... b; } inline scalar_t identity() const { return ...; } }; Tensor _sparse_csr_..._cpu(const Tensor& input, IntArrayRef dims_to_sum, bool keepdim, std::optional dtype) { ... result = reduce_sparse_csr_cpu_template(input_, dims_to_sum, keepdim, Reduction...Op()); ... return result; } and add the following - func: _sparse_csr_op.dim_dtype(Tensor self, int[1] dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor dispatch: SparseCsrCUDA: _sparse_csr_..._cpu to native_functions.yaml Use ReductionAddOp and _sparse_csr_sum implementation as an example. - Since a CSR tensor dimensionality is always 2, only reductions with keepdim=True can be supported. */ namespace { template Tensor reduce_sparse_csr_dim0_cpu_template(const Tensor& sparse, ReductionOp rop) { /* Consider the following sparse tensor: 1 * * * * * * * 2 * * * 3 * * * * * * * 4 * 5 * * that has CSR representation crow_indices = [0, 1, 2, 3, 3, 5] col_indices = [0, 3, 2, 0, 2] values = [1, 2, 3, 4, 5] Reduction with dim=0 results: rop(1,4) * rop(3,5) 2 * that has CSR representation new_crow_indices = [0, 3] new_col_indices = [0, 2, 3] new_values = [rop(1, 4], rop(3, 5), 2] In general, the CSR representation data can be computed as follows: new_col_indices, col_map = col_indices.unique(sorted=True, return_inverse=True) nnz = new_col_indices.numel() new_crow_indices = [0, nnz] new_values.resize(nnz); new_values.fill_(identity) for i in range(col_indices.numel()): new_values[col_map[i]] = rop(new_values[col_map[i], values[i]) */ Tensor col_indices = sparse.col_indices(); Tensor values = sparse.values(); auto numel = values.numel(); /* Calling at::_unique constitutes the main bottleneck of this function. However, it is still about 5x faster than using the invariant: csr.sum(dim=0) == csr.transpose(0, 1).sum(dim=1) */ auto [new_col_indices, columns_map] = at::_unique(col_indices, true, true); auto nnz = new_col_indices.numel(); Tensor new_crow_indices = at::empty({2}, col_indices.options()); new_crow_indices[0] = 0; new_crow_indices[1] = nnz; // Set `is_cuda` = `true` in acc_type in CPU backend. Because the accumulate type // of float should be float in current scenario. In CUDA, float is the accumulate type // of float, while in CPU, double is the accumulate type of float. using acc_t = at::acc_type; auto acc_buffer = at::sparse_csr::create_acc_buffer( values.options(), values.scalar_type(), nnz); Tensor new_values = std::get<0>(acc_buffer); Tensor new_values_acc = std::get<1>(acc_buffer); new_values_acc.fill_(rop.identity()); int64_t* columns_map_ptr = columns_map.data_ptr(); scalar_t* values_ptr = values.data_ptr(); acc_t* new_values_acc_ptr = new_values_acc.data_ptr(); // There is no point in parallelizing the following for-loop // because about 99.3% of the computation time is spent in the // at::_unique call above. for (const auto i : c10::irange(numel)) { int64_t col = columns_map_ptr[i]; scalar_t val = values_ptr[i]; new_values_acc_ptr[col] = rop(new_values_acc_ptr[col], static_cast(val)); } copy_from_acc_buffer(new_values, new_values_acc); return at::native::_sparse_csr_tensor_unsafe(new_crow_indices, new_col_indices, new_values, {1, sparse.size(1)}, new_values.scalar_type(), sparse.layout(), new_values.device()); } template Tensor reduce_sparse_csr_dim1_cpu_template(const Tensor& sparse, ReductionOp rop) { /* Consider the following sparse tensor: 1 * * * * * * * 2 * * * 3 * * * * * * * 4 * 5 * * that has CSR representation crow_indices = [0, 1, 2, 3, 3, 5] col_indices = [0, 3, 2, 0, 2] values = [1, 2, 3, 4, 5] Reduction with dim=1 results: 1 2 3 * rop(4, 5) that has CSR representation new_crow_indices = [0, 1, 2, 3, 3, 4] new_col_indices = [0, 0, 0, 0] new_values = [1, 2, 3, rop(4, 5)] In general, the result CSR data can be computed as follows: new_crow_indices = [0] for i in range(1, nrows+1): new_crow_indices[i] = new_crow_indices[i-1] + (crow_indices[i] == crow_indices[i-1]) nnz = new_crow_indices[-1] new_col_indices = zeros(nnz) new_values.resize(nnz) j = -1 for i in range(1, nrows+1): if crow_indices[i] == crow_indices[i-1]: continue j += 1 new_values[j] = rop(values[crow_indices[i] : crow_indices[i-1]]) */ Tensor crow_indices = sparse.crow_indices(); auto ioptions = crow_indices.options(); Tensor values = sparse.values(); auto nrows = sparse.size(0); Tensor new_crow_indices = at::empty({crow_indices.numel()}, ioptions); Tensor new_col_indices = at::empty({}, ioptions); Tensor row_map = at::empty({nrows}, ioptions); // Set `is_cuda` = `true` in acc_type in CPU backend. Because the accumulate type // of float should be float in current scenario. In CUDA, float is the accumulate type // of float, while in CPU, double is the accumulate type of float. using acc_t = at::acc_type; auto acc_buffer = at::sparse_csr::create_acc_buffer( values.options(), values.scalar_type()); Tensor new_values = std::get<0>(acc_buffer); Tensor new_values_acc = std::get<1>(acc_buffer); AT_DISPATCH_INDEX_TYPES(crow_indices.scalar_type(), "reduce_sparse_csr_dim1_cpu_indices", [&]() { index_t* crow_indices_ptr = crow_indices.data_ptr(); index_t* new_crow_indices_ptr = new_crow_indices.data_ptr(); index_t* row_map_ptr = row_map.data_ptr(); int64_t nnz = 0; new_crow_indices_ptr[0] = 0; for(int64_t i=0; i(); acc_t* new_values_acc_ptr = new_values_acc.data_ptr(); at::parallel_for( 0, nrows, internal::GRAIN_SIZE, [&](int64_t irow_start, int64_t irow_end) { index_t i_end = crow_indices_ptr[irow_start]; for (index_t h = irow_start; h < irow_end; ++h) { index_t i_start = i_end; i_end = crow_indices_ptr[h+1]; if (i_start != i_end) { acc_t res = static_cast(values_ptr[i_start]); for (index_t i = i_start + 1; i < i_end; i++) { res = rop(res, static_cast(values_ptr[i])); } new_values_acc_ptr[row_map_ptr[h]] = res; } } }); }); copy_from_acc_buffer(new_values, new_values_acc); return at::native::_sparse_csr_tensor_unsafe(new_crow_indices, new_col_indices, new_values, {sparse.size(0), 1}, new_values.scalar_type(), sparse.layout(), new_values.device()); } template Tensor reduce_sparse_csr_dim01_cpu_template(const Tensor& sparse, ReductionOp rop) { auto ioptions = sparse.col_indices().options(); Tensor values = sparse.values(); auto numel = values.numel(); auto nnz = std::min(1, numel); /* TODO: we can likely do about 3x better than parallel_reduce: In [2]: t=torch.randn(5000, 5000).to_sparse_csr() In [3]: %timeit torch._sparse_csr_sum(t, dim=(0, 1), keepdim=True) 3.39 ms ± 898 ns per loop (mean ± std. dev. of 7 runs, 100 loops each) In [4]: %timeit torch.sum(t.values()) 1.07 ms ± 291 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each) */ // Set `is_cuda` = `true` in acc_type in CPU backend. Because the accumulate type // of float should be float in current scenario. In CUDA, float is the accumulate type // of float, while in CPU, double is the accumulate type of float. using acc_t = at::acc_type; scalar_t* values_ptr = values.data_ptr(); acc_t value = at::parallel_reduce( 0, numel, internal::GRAIN_SIZE, rop.identity(), [&](int64_t i_start, int64_t i_end, scalar_t identity) { acc_t res = acc_t(identity); for (int64_t i=i_start; i{0, nnz}, ioptions); Tensor new_values; auto result_dtype = at::isIntegralType(values.scalar_type(), /*includeBool=*/true) ? ScalarType::Long : values.scalar_type(); if (numel > 0) { new_values = at::empty({1}, values.options().dtype(result_dtype)); new_values.fill_(value); } else { new_values = at::empty({}, values.options().dtype(result_dtype)); } return at::native::_sparse_csr_tensor_unsafe(new_crow_indices, new_col_indices, new_values, {1, std::min(1, sparse.size(1))}, new_values.scalar_type(), sparse.layout(), new_values.device()); } template Tensor reduce_sparse_csr_cpu_template(const Tensor& sparse, std::vector dims, ReductionOp rop) { if (dims.size() == 1) { if (dims[0] == 0) { return reduce_sparse_csr_dim0_cpu_template(sparse, rop); } else { TORCH_INTERNAL_ASSERT(dims[0] == 1); return reduce_sparse_csr_dim1_cpu_template(sparse, rop); } } else if (dims.size() == 2) { TORCH_INTERNAL_ASSERT(((dims[0] == 0 && dims[1] == 1) || (dims[0] == 1 && dims[1] == 0))); return reduce_sparse_csr_dim01_cpu_template(sparse, rop); } TORCH_INTERNAL_ASSERT(dims.empty()); // effective after gh-29137 has been resolved return sparse.clone(); } template Tensor reduce_sparse_csr_cpu_template(const Tensor& sparse, IntArrayRef dims_to_sum, bool keepdim, ReductionOp rop) { TORCH_INTERNAL_ASSERT(sparse.is_sparse_csr()); TORCH_CHECK(keepdim, "reduction operations on CSR tensors with keepdim=False is unsupported"); TORCH_INTERNAL_ASSERT(sparse.device() == kCPU); const int64_t input_dim = sparse.dim(); TORCH_INTERNAL_ASSERT(input_dim == 2); auto dims = dims_to_sum.vec(); maybe_wrap_dims(dims, input_dim); if (dims.empty()) { // after gh-29137 is resolved, delete this if-block dims.emplace_back(0); dims.emplace_back(1); } return reduce_sparse_csr_cpu_template(sparse, dims, rop); } template struct ReductionAddOp { inline scalar_t operator()(const scalar_t& a, const scalar_t& b) const { return a + b; } inline scalar_t identity() const { return 0; } }; template struct ReductionMulOp { inline scalar_t operator()(const scalar_t& a, const scalar_t& b) const { return a * b; } inline scalar_t identity() const { return 1; } }; } // namespace Tensor _sparse_csr_sum_cpu(const Tensor& input, IntArrayRef dims_to_sum, bool keepdim, std::optional dtype) { ScalarType dtype_ = dtype.value_or(input.scalar_type()); Tensor input_ = at::sparse_csr::to_type(input, dtype_); Tensor result; AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2( kHalf, kBFloat16, input_.scalar_type(), "_sparse_csr_sum_cpu", [&] { // Set `is_cuda` = `true` in acc_type in CPU backend. Because the accumulate type // of float should be float in current scenario. In CUDA, float is the accumulate type // of float, while in CPU, double is the accumulate type of float. using acc_t = at::acc_type; result = reduce_sparse_csr_cpu_template( input_, dims_to_sum, keepdim, ReductionAddOp()); }); return result; } Tensor _sparse_csr_prod_cpu(const Tensor& input, IntArrayRef dims_to_reduce, bool keepdim, std::optional dtype) { ScalarType dtype_ = dtype.value_or(input.scalar_type()); Tensor input_ = input.to(dtype_); Tensor result; AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2( kHalf, kBFloat16, input_.scalar_type(), "_sparse_csr_prod_cpu", [&] { result = reduce_sparse_csr_cpu_template(input_, dims_to_reduce, keepdim, ReductionMulOp()); }); return result; } std::tuple _sparse_mm_reduce_impl_sparse_csr_cpu( const Tensor& self, const Tensor& other, const c10::string_view reduce) { auto layout = self.layout(); TORCH_CHECK(layout == kSparseCsr, "sparse_mm_reduce: expect self to be SparseCsr, got ", layout); TORCH_CHECK(self.dense_dim() == 0, "sparse_mm_reduce: expected non-hybrid self tensor."); TORCH_CHECK(self.dim() == 2, "sparse_mm_reduce: expected self to be a 2-D tensor, got ", self.dim(), "-D tensor."); sparse::impl::check_sparse_mm_reduce_impl_inputs( self, Tensor(), other); auto op = get_reduction_enum(reduce); TORCH_CHECK(op != ReductionType::PROD, "sparse_mm_reduce: reduce type of prod has not been enabled.") auto crow = self.crow_indices(); auto col = self.col_indices(); auto val = self.values(); // init output to be all zeros, for `rows` that has no nonzero elements, // the corresponding rows in the output will be zero. auto out = at::zeros({self.size(0), other.size(1)}, other.options()); auto arg_out = at::empty({0}, col.options()); int64_t nnz = self._nnz(); if (nnz == 0) { return std::make_tuple(out, arg_out); } // only need to calculate the out args // for reduce type "amax" and "amin" for training bool need_arg_out = at::GradMode::is_enabled() && (self.requires_grad() || other.requires_grad()) && (op == ReductionType::MAX || op == ReductionType::MIN); if (!need_arg_out) { spmm_reduce_stub(kCPU, out, crow, col, val, other, op); } else { // allocate memory and init with invalid index arg_out.resize_(out.sizes()); arg_out.fill_(nnz); spmm_reduce_arg_stub(kCPU, out, arg_out, crow, col, val, other, op); } return std::make_tuple(std::move(out), std::move(arg_out)); } std::tuple _sparse_mm_reduce_impl_backward_sparse_csr_cpu( const Tensor& self, const Tensor& grad_out, const Tensor& other, const c10::string_view reduce, const Tensor& arg_out, std::array output_mask) { auto layout = self.layout(); TORCH_CHECK(layout == kSparseCsr, "sparse_mm_reduce: expect self to be SparseCsr, got ", layout); sparse::impl::check_sparse_mm_reduce_impl_inputs( self, grad_out, other); auto op = get_reduction_enum(reduce); auto crow = self.crow_indices(); auto col = self.col_indices(); auto val = self.values(); // `row`: row indices of COO format // `ccol`: ccol indices of CSC format (with permute) // `permute`: permute pattern from CSR to CSC // // TODO: optimize the following section, // currently `argsort` is sequential. Tensor row, ccol, permute; { bool out_int32 = crow.scalar_type() == ScalarType::Int; Tensor coo_indices = at::_convert_indices_from_csr_to_coo( crow, col, out_int32, /*transpose*/false); row = coo_indices.select(0, 0); // calculate the global index for CSC // and get the conversion permute pattern Tensor index = col.mul(self.size(0)).add_(row); permute = index.argsort(); ccol = at::_convert_indices_from_coo_to_csr( /*column indices*/col.index_select(0, permute), /*column count*/self.size(1), out_int32); } Tensor grad_self, grad_other; if (output_mask[0]) { // grad_input has the same indices and nnz with input grad_self = at::empty_like(self); grad_self.values().zero_(); if (op == ReductionType::MAX || op == ReductionType::MIN) { spmm_reduce_backward_input_arg_stub(kCPU, grad_self, grad_out, col, other, arg_out, op); } else { spmm_reduce_backward_input_stub(kCPU, grad_self, grad_out, crow, col, other, row, op); } } if (output_mask[1]) { grad_other = at::zeros(other.sizes(), other.options()); if (op == ReductionType::MAX || op == ReductionType::MIN) { spmm_reduce_backward_other_arg_stub(kCPU, grad_other, grad_out, col, val, arg_out, op); } else { spmm_reduce_backward_other_stub(kCPU, grad_other, grad_out, crow, val, row, ccol, permute, op); } } return std::make_tuple(std::move(grad_self), std::move(grad_other)); } DEFINE_DISPATCH(spmm_reduce_stub); DEFINE_DISPATCH(spmm_reduce_arg_stub); DEFINE_DISPATCH(spmm_reduce_backward_input_stub); DEFINE_DISPATCH(spmm_reduce_backward_input_arg_stub); DEFINE_DISPATCH(spmm_reduce_backward_other_stub); DEFINE_DISPATCH(spmm_reduce_backward_other_arg_stub); } // namespace native } // namespace at