Add scatter/segment bf16 support (#316)

csrc/cpu/scatter_cpu.cpp

@@ -57,7 +57,7 @@ scatter_cpu(torch::Tensor src, torch::Tensor index, int64_t dim,
   auto N = out.size(dim);
 
   auto index_info = getTensorInfo<int64_t>(index);
-  AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, src.scalar_type(), "_", [&] {
+  AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "scatter_cpu", [&] {
     auto src_data = src.data_ptr<scalar_t>();
     auto out_data = out.data_ptr<scalar_t>();
 

csrc/cpu/segment_coo_cpu.cpp

@@ -69,7 +69,7 @@ segment_coo_cpu(torch::Tensor src, torch::Tensor index,
   auto index_info = getTensorInfo<int64_t>(index);
   auto stride = index_info.strides[index_info.dims - 1];
   std::vector<int64_t> args(K);
-  AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, src.scalar_type(), "_", [&] {
+  AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "segment_coo_cpu", [&] {
     auto src_data = src.data_ptr<scalar_t>();
     auto out_data = out.data_ptr<scalar_t>();
     scalar_t *count_data = nullptr;
@@ -178,7 +178,7 @@ torch::Tensor gather_coo_cpu(torch::Tensor src, torch::Tensor index,
 
   auto index_info = getTensorInfo<int64_t>(index);
   auto stride = index_info.strides[index_info.dims - 1];
-  AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, src.scalar_type(), "_", [&] {
+  AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "gather_coo_cpu", [&] {
     auto src_data = src.data_ptr<scalar_t>();
     auto out_data = out.data_ptr<scalar_t>();
 

csrc/cpu/segment_csr_cpu.cpp

@@ -57,7 +57,7 @@ segment_csr_cpu(torch::Tensor src, torch::Tensor indptr,
   auto indptr_info = getTensorInfo<int64_t>(indptr);
   auto stride = indptr_info.strides[indptr_info.dims - 1];
   std::vector<int64_t> args(K);
-  AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, src.scalar_type(), "_", [&] {
+  AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "segment_csr_cpu", [&] {
     auto src_data = src.data_ptr<scalar_t>();
     auto out_data = out.data_ptr<scalar_t>();
 
@@ -135,7 +135,7 @@ torch::Tensor gather_csr_cpu(torch::Tensor src, torch::Tensor indptr,
 
   auto indptr_info = getTensorInfo<int64_t>(indptr);
   auto stride = indptr_info.strides[indptr_info.dims - 1];
-  AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, src.scalar_type(), "_", [&] {
+  AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "gather_csr_cpu", [&] {
     auto src_data = src.data_ptr<scalar_t>();
     auto out_data = out.data_ptr<scalar_t>();
 

test/utils.py

@@ -2,7 +2,8 @@ import torch
 
 reductions = ['sum', 'add', 'mean', 'min', 'max']
 
-dtypes = [torch.half, torch.float, torch.double, torch.int, torch.long]
+dtypes = [torch.half, torch.bfloat16, torch.float, torch.double,
+          torch.int, torch.long]
 grad_dtypes = [torch.float, torch.double]
 
 devices = [torch.device('cpu')]