torch.half support for spmm (CPU)

csrc/cpu/reducer.h

@@ -72,7 +72,7 @@ template <typename scalar_t, ReductionType REDUCE> struct Reducer {
     if (REDUCE == SUM || REDUCE == MUL || REDUCE == DIV)
       *address = val;
     else if (REDUCE == MEAN)
-      *address = val / (count > 0 ? count : (scalar_t)1);
+      *address = val / (scalar_t)(count > 0 ? count : 1);
     else if (REDUCE == MIN || REDUCE == MAX) {
       if (count > 0) {
         *address = val;

csrc/cpu/spmm_cpu.cpp

@@ -44,7 +44,8 @@ spmm_cpu(torch::Tensor rowptr, torch::Tensor col,
   auto K = mat.size(-1);
   auto B = mat.numel() / (N * K);
 
-  AT_DISPATCH_ALL_TYPES(mat.scalar_type(), "spmm", [&] {
+  AT_DISPATCH_ALL_TYPES_AND(
+      at::ScalarType::Half, mat.scalar_type(), "spmm", [&] {
         scalar_t *value_data = nullptr;
         auto mat_data = mat.data_ptr<scalar_t>();
         auto out_data = out.data_ptr<scalar_t>();
@@ -55,8 +56,10 @@ spmm_cpu(torch::Tensor rowptr, torch::Tensor col,
               value_data = optional_value.value().data_ptr<scalar_t>();
             }
 
-        int64_t grain_size = at::internal::GRAIN_SIZE / (K * (col.numel() / M));
-        at::parallel_for(0, B * M, grain_size, [&](int64_t begin, int64_t end) {
+            int64_t grain_size =
+                at::internal::GRAIN_SIZE / (K * (col.numel() / M));
+            at::parallel_for(
+                0, B * M, grain_size, [&](int64_t begin, int64_t end) {
                   scalar_t val;
                   std::vector<scalar_t> vals(K);
                   int64_t row_start, row_end, b, m, c;
@@ -78,18 +81,20 @@ spmm_cpu(torch::Tensor rowptr, torch::Tensor col,
                       for (auto k = 0; k < K; k++) {
                         if (HAS_VALUE)
                           Reducer<scalar_t, REDUCE>::update(
-                      &vals[k], val * mat_data[offset + c * K + k], &args[k],
-                      e);
+                              &vals[k], val * mat_data[offset + c * K + k],
+                              &args[k], e);
                         else
                           Reducer<scalar_t, REDUCE>::update(
-                      &vals[k], mat_data[offset + c * K + k], &args[k], e);
+                              &vals[k], mat_data[offset + c * K + k], &args[k],
+                              e);
                       }
                     }
                     offset = b * M * K + m * K;
                     for (auto k = 0; k < K; k++)
-              Reducer<scalar_t, REDUCE>::write(out_data + offset + k, vals[k],
-                                               arg_out_data + offset + k,
-                                               args[k], row_end - row_start);
+                      Reducer<scalar_t, REDUCE>::write(
+                          out_data + offset + k, vals[k],
+                          arg_out_data + offset + k, args[k],
+                          row_end - row_start);
                   }
                 });
           });
@@ -122,7 +127,8 @@ torch::Tensor spmm_value_bw_cpu(torch::Tensor row, torch::Tensor rowptr,
   auto row_data = row.data_ptr<int64_t>();
   auto rowptr_data = rowptr.data_ptr<int64_t>();
   auto col_data = col.data_ptr<int64_t>();
-  AT_DISPATCH_ALL_TYPES(mat.scalar_type(), "spmm_value_bw", [&] {
+  AT_DISPATCH_ALL_TYPES_AND(
+      at::ScalarType::Half, mat.scalar_type(), "spmm_value_bw", [&] {
         auto mat_data = mat.data_ptr<scalar_t>();
         auto grad_data = grad.data_ptr<scalar_t>();
         auto out_data = out.data_ptr<scalar_t>();
@@ -138,7 +144,8 @@ torch::Tensor spmm_value_bw_cpu(torch::Tensor row, torch::Tensor rowptr,
                        grad_data[b * M * K + row * K + k];
               }
               if (REDUCE == MEAN) {
-            int row_start = rowptr_data[row], row_end = rowptr_data[row + 1];
+                int row_start = rowptr_data[row],
+                    row_end = rowptr_data[row + 1];
                 val /= (scalar_t)std::max(row_end - row_start, 1);
               }
               out_data[e] += val;

test/test_matmul.py

@@ -45,6 +45,20 @@ def test_spmm(dtype, device, reduce):
     assert torch.allclose(expected_grad_other, other.grad, atol=1e-6)
 
 
+def test_spmm_half_precision():
+    src_dense = torch.randn((10, 8), dtype=torch.half, device='cpu')
+    src_dense[2:4, :] = 0  # Remove multiple rows.
+    src_dense[:, 2:4] = 0  # Remove multiple columns.
+    src = SparseTensor.from_dense(src_dense)
+
+    other = torch.randn((2, 8, 2), dtype=torch.half, device='cpu')
+
+    expected = src_dense @ other
+    out = src @ other
+
+    assert torch.allclose(expected, out, atol=1e-6)
+
+
 @pytest.mark.parametrize('dtype,device', product(grad_dtypes, devices))
 def test_spspmm(dtype, device):
     src = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=dtype,