1 #pragma once
2
3 #include <ATen/Tensor.h>
4 #include <ATen/core/Scalar.h>
5 #include <ATen/TensorUtils.h>
6 #include <ATen/native/ReductionType.h>
7 #include <ATen/native/cpu/SpmmReduceKernel.h>
8
9 namespace at::native::sparse::impl {
10
11 // Returns true if all entries of self are zero
12 // TODO: This has potential to be a generic helper
_is_sparse_and_zero(const Tensor & self)13 inline bool _is_sparse_and_zero(const Tensor& self) {
14 if (self.layout() == kSparse || self.layout() == kSparseCsr ||
15 self.layout() == kSparseCsc || self.layout() == kSparseBsr ||
16 self.layout() == kSparseBsc) {
17 if (self._nnz() == 0) {
18 return true;
19 }
20 }
21 return false;
22 }
23
_check_is_cpu(const Tensor & self,c10::string_view name)24 inline void _check_is_cpu(const Tensor& self, c10::string_view name) {
25 TORCH_CHECK(
26 self.is_cpu(),
27 "Expected all tensors to be on the same device. addmm expected '",
28 name,
29 "' to be CPU tensor, but got ",
30 self.device(),
31 " tensor");
32 }
33
_check_is_cuda(const Tensor & self,c10::string_view name)34 inline void _check_is_cuda(const Tensor& self, c10::string_view name) {
35 TORCH_CHECK(
36 self.is_cuda(),
37 "Expected all tensors to be on the same device. addmm expected '",
38 name,
39 "' to be CUDA tensor, but got ",
40 self.device(),
41 " tensor");
42 }
43
_check_dim(const Tensor & self,int64_t target_dim,c10::string_view name)44 inline void _check_dim(const Tensor& self, int64_t target_dim, c10::string_view name) {
45 if (target_dim == 2) {
46 TORCH_CHECK(
47 self.dim() == target_dim,
48 name, " must be a matrix, ",
49 "got ", self.dim(), "-D tensor");
50 }
51 TORCH_CHECK(
52 self.dim() == target_dim,
53 "Expected ",
54 name,
55 " to be of dimension ",
56 target_dim,
57 " but got ",
58 self.dim(),
59 " instead.");
60 }
61
62 template <bool train>
check_sparse_mm_reduce_impl_inputs(const Tensor & self,const Tensor & grad_out,const Tensor & other)63 inline void check_sparse_mm_reduce_impl_inputs(
64 const Tensor& self,
65 const Tensor& grad_out,
66 const Tensor& other) {
67 TORCH_INTERNAL_ASSERT(self.is_sparse_csr());
68
69 const auto input_scalar_type = self.values().scalar_type();
70 CheckedFrom c = train ? "sparse_mm_reduce_backward" : "sparse_mm_reduce";
71 if (train) {
72 checkLayout(c, grad_out, kStrided);
73 checkScalarType(c, {grad_out, "grad_out", 1}, input_scalar_type);
74 check_dim_size(grad_out, 2, 0, self.size(0));
75 check_dim_size(grad_out, 2, 1, other.size(1));
76 }
77
78 int pos = train ? 2 : 1;
79 checkLayout(c, other, kStrided);
80 checkScalarType(c, {other, "other", pos}, input_scalar_type);
81 check_dim_size(other, 2, 0, self.size(1));
82 }
83
84 } // at::native::sparse::impl
85