Revert "port masked_select from TH to ATen and optimize perf on CPU (pytorch#33269)" (pytorch#41829)

gchanan · web-flow · commit 4081fdd3df04 · 2020-07-22T09:52:30.000-07:00
This reverts commit fe66bdb. This also makes a sense to THTensorEvenMoreMath because sumall was removed, see THTensor_wrap.
diff --git a/aten/src/ATen/Declarations.cwrap b/aten/src/ATen/Declarations.cwrap
@@ -56,6 +56,38 @@
     - THBoolTensor* mask
     - THTensor* source
 ]]
+[[
+  name: _th_masked_select
+  cname: maskedSelect
+  cpu_bool: True
+  cpu_bfloat16: True
+  variants:
+    - function
+  backends:
+    - CPU
+  return: argument 0
+  arguments:
+    - arg: THTensor* result
+      output: True
+    - THTensor* self
+    - THByteTensor* mask
+]]
+[[
+  name: _th_masked_select_bool
+  cname: maskedSelectBool
+  cpu_bool: True
+  cpu_bfloat16: True
+  variants:
+    - function
+  backends:
+    - CPU
+  return: argument 0
+  arguments:
+    - arg: THTensor* result
+      output: True
+    - THTensor* self
+    - THBoolTensor* mask
+]]
 [[
   name: _th_nonzero
   cname: nonzero
diff --git a/aten/src/ATen/native/LegacyDefinitions.cpp b/aten/src/ATen/native/LegacyDefinitions.cpp
@@ -22,6 +22,32 @@ Tensor & masked_scatter__cpu(Tensor& self, const Tensor & mask, const Tensor & s
   }
 }
 
+Tensor masked_select_cpu(const Tensor & self, const Tensor & mask) {
+  namedinference::compute_broadcast_outnames(self, mask);
+
+  Tensor b_self, b_mask;
+  std::tie(b_self, b_mask) = expand_outplace(self, mask, "masked_select");
+  if (b_mask.dtype() == at::ScalarType::Byte) {
+    TORCH_WARN("masked_select received a mask with dtype torch.uint8, this behavior is now deprecated," \
+            "please use a mask with dtype torch.bool instead.");
+    return legacy::cpu::_th_masked_select(b_self, b_mask);
+  } else {
+    return legacy::cpu::_th_masked_select_bool(b_self, b_mask);
+  }
+}
+
+Tensor & masked_select_out_cpu(Tensor & result, const Tensor & self, const Tensor & mask) {
+  namedinference::compute_broadcast_outnames(self, mask);
+
+  Tensor b_self, b_mask;
+  std::tie(b_self, b_mask) = expand_outplace(self, mask, "masked_select_out");
+  if (b_mask.dtype() == at::ScalarType::Bool) {
+    return legacy::cpu::_th_masked_select_bool_out(result, b_self, b_mask);
+  } else {
+    return legacy::cpu::_th_masked_select_out(result, b_self, b_mask);
+  }
+}
+
 Tensor argsort(const Tensor & self, int64_t dim, bool descending) {
   return std::get<1>(at::sort(self, dim, descending));
 }
diff --git a/aten/src/ATen/native/TensorAdvancedIndexing.cpp b/aten/src/ATen/native/TensorAdvancedIndexing.cpp
@@ -71,8 +71,6 @@ DEFINE_DISPATCH(index_put_stub);
 DEFINE_DISPATCH(index_put_accum_stub);
 DEFINE_DISPATCH(masked_fill_stub);
 REGISTER_NO_CPU_DISPATCH(index_put_accum_stub, index_put_accum_fn);
-DEFINE_DISPATCH(masked_select_serial_stub);
-DEFINE_DISPATCH(masked_select_stub);
 
 DEFINE_DISPATCH(gather_stub);
 DEFINE_DISPATCH(scatter_stub);
@@ -629,82 +627,6 @@ Tensor masked_fill(const Tensor & self, const Tensor & mask, const Tensor & sour
   return result;
 }
 
-static Tensor & masked_select_out_impl_cpu(Tensor & result, const Tensor & self, const Tensor & mask) {
-  NoNamesGuard guard;
-
-  TORCH_CHECK(mask.scalar_type() == ScalarType::Byte || mask.scalar_type() == ScalarType::Bool,
-              "masked_select: expected BoolTensor or ByteTensor for mask");
-  TORCH_CHECK(self.scalar_type() == result.scalar_type(),
-              "masked_select(): self and result must have the same scalar type");
-
-  if (mask.dtype() == at::ScalarType::Byte) {
-    TORCH_WARN("masked_select received a mask with dtype torch.uint8, this behavior is now deprecated," \
-            "please use a mask with dtype torch.bool instead.");
-  }
-
-  Tensor _mask, _self;
-  std::tie(_mask, _self) = expand_outplace(mask, self);
-
-  auto shape = _self.sizes();
-  int64_t numel = _mask.sum().item().toLong();
-  result.resize_({numel});
-  if (numel == 0) {
-    return result;
-  }
-
-  // Create strided view of result before feeding into TensorIterator
-  auto strides = DimVector(shape.size(), 0);
-  auto result_strided = result.as_strided(shape, strides);
-
-  // serial kernel
-  bool use_serial_kernel = self.numel() < at::internal::GRAIN_SIZE || at::get_num_threads() == 1;
-  if (use_serial_kernel) {
-    auto iter = TensorIteratorConfig()
-      .check_all_same_dtype(false)
-      .resize_outputs(false)
-      .add_output(result_strided)
-      .add_input(_self)
-      .add_input(_mask)
-      .build();
-
-    masked_select_serial_stub(iter.device_type(), iter);
-    return result;
-  }
-
-  // Use a prefix sum to record the output locations of the masked elements,
-  // so as to parallel with TensorIterator.
-  auto mask_long = at::empty(shape, self.options().dtype(at::kLong)).copy_(_mask);
-  auto mask_prefix_sum = at::empty(shape, self.options().dtype(at::kLong));
-  auto mask_long_data = mask_long.data_ptr<int64_t>();
-  auto mask_prefix_sum_data = mask_prefix_sum.data_ptr<int64_t>();
-  // TODO: Here can only use std::partial_sum for C++14,
-  // use std::exclusive_scan when PyTorch upgrades to C++17, which have better peformance.
-  // std::exclusive_scan(mask_long_data, mask_long_data + mask_long.numel(), mask_prefix_sum_data, 0);
-  std::partial_sum(mask_long_data, mask_long_data + mask_long.numel(), mask_prefix_sum_data);
-
-  auto iter = TensorIteratorConfig()
-    .check_all_same_dtype(false)
-    .resize_outputs(false)
-    .add_output(result_strided)
-    .add_input(_self)
-    .add_input(_mask)
-    .add_input(mask_prefix_sum)
-    .build();
-
-  masked_select_stub(iter.device_type(), iter);
-  return result;
-}
-
-Tensor & masked_select_out_cpu(Tensor & result, const Tensor & self, const Tensor & mask) {
-  namedinference::compute_broadcast_outnames(self, mask);
-  return masked_select_out_impl_cpu(result, self, mask);
-}
-
-Tensor masked_select_cpu(const Tensor & self, const Tensor & mask) {
-  Tensor result = at::empty({0}, self.options());
-  return masked_select_out_cpu(result, self, mask);
-}
-
 Tensor _gather_sparse_backward(const Tensor& self, int64_t dim, const Tensor& index, const Tensor& grad){
 // special case scalar input and/or index
     if (self.ndimension() == 0) return at::_sparse_coo_tensor_unsafe(at::empty({0,grad.numel()}, index.options()), grad, self.sizes());
diff --git a/aten/src/ATen/native/TensorAdvancedIndexing.h b/aten/src/ATen/native/TensorAdvancedIndexing.h
@@ -15,7 +15,6 @@ using index_fn = void(*)(TensorIterator &, IntArrayRef indexed_sizes, IntArrayRe
 using index_put_fn = void(*)(TensorIterator &, IntArrayRef indexed_sizes, IntArrayRef indexed_strides, bool accumulate);
 using index_put_accum_fn = void(*)(Tensor &, TensorList , const Tensor &, bool unsafe);
 using masked_fill_fn = void(*)(TensorIterator &, Scalar scalar);
-using masked_select_fn = void(*)(TensorIterator &);
 
 using gather_fn = void (*)(Tensor & result, const Tensor & self, int64_t dim, const Tensor & index);
 using scatter_fn = void(*)(Tensor& self, int64_t dim, const Tensor& index, const Tensor& src);
@@ -26,8 +25,6 @@ DECLARE_DISPATCH(index_fn, index_stub);
 DECLARE_DISPATCH(index_put_fn, index_put_stub);
 DECLARE_DISPATCH(index_put_accum_fn, index_put_accum_stub);
 DECLARE_DISPATCH(masked_fill_fn, masked_fill_stub);
-DECLARE_DISPATCH(masked_select_fn, masked_select_serial_stub);
-DECLARE_DISPATCH(masked_select_fn, masked_select_stub);
 
 DECLARE_DISPATCH(gather_fn, gather_stub);
 DECLARE_DISPATCH(scatter_fn, scatter_stub);
diff --git a/aten/src/ATen/native/cpu/IndexKernel.cpp b/aten/src/ATen/native/cpu/IndexKernel.cpp
@@ -163,90 +163,11 @@ void masked_fill_kernel(TensorIterator& iter, Scalar value) {
     });
 }
 
-template <typename scalar_t, typename mask_t, typename func_t>
-void cpu_masked_select_serial_kernel(TensorIterator& iter, const func_t& f) {
-  auto is_mask_bool = std::is_same<mask_t, bool>::value;
-  int64_t offset = 0;
-  auto loop = [&](char** data, const int64_t* strides, int64_t n) {
-    char* dst = data[0];
-    char* src = data[1];
-    char* mask = data[2];
-    for (int64_t i = 0; i < n; i++) {
-      mask_t mask_value = *(mask_t*)(mask + strides[2] * i);
-      if (!is_mask_bool) {
-        TORCH_CHECK(mask_value == 0 || mask_value == 1, "Mask tensor can take 0 and 1 values only");
-      }
-      if (mask_value) {
-        int64_t offset_bytes = offset * sizeof(scalar_t);
-        f(dst, src + strides[1] * i, offset_bytes);
-        offset++;
-      }
-    }
-  };
-  iter.serial_for_each(loop, {0, iter.numel()});
-}
-
-void masked_select_serial_kernel(TensorIterator& iter) {
-  AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(at::ScalarType::Bool, at::ScalarType::BFloat16,
-    iter.dtype(), "masked_select", [&] {
-      auto mask_dtype = iter.input_dtype(1);
-      if (mask_dtype == at::ScalarType::Bool) {
-        cpu_masked_select_serial_kernel<scalar_t, bool>(iter, [](char* dst, char* src, int64_t offset) {
-          *(scalar_t*)(dst + offset) = *(scalar_t*)src;
-        });
-      } else {
-        cpu_masked_select_serial_kernel<scalar_t, unsigned char>(iter, [](char* dst, char* src, int64_t offset) {
-          *(scalar_t*)(dst + offset) = *(scalar_t*)src;
-        });
-      }
-    });
-}
-
-template <typename scalar_t, typename mask_t, typename func_t>
-void cpu_masked_select_kernel(TensorIterator& iter, const func_t& f) {
-  auto is_mask_bool = std::is_same<mask_t, bool>::value;
-  auto loop = [&](char** data, const int64_t* strides, int64_t n) {
-    char* dst = data[0];
-    char* src = data[1];
-    char* mask = data[2];
-    char* mask_prefix_sum = data[3];
-    for (int64_t i = 0; i < n; i++) {
-      mask_t mask_value = *(mask_t*)(mask + strides[2] * i);
-      if (!is_mask_bool) {
-        TORCH_CHECK(mask_value == 0 || mask_value == 1, "Mask tensor can take 0 and 1 values only");
-      }
-      if (mask_value) {
-        int64_t offset = *(int64_t*)(mask_prefix_sum + strides[3] * i);
-        int64_t offset_bytes = (offset - 1) * sizeof(scalar_t);
-        f(dst, src + strides[1] * i, offset_bytes);
-      }
-    }
-  };
-  iter.for_each(loop);
-}
-
-void masked_select_kernel(TensorIterator& iter) {
-  AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(at::ScalarType::Bool, at::ScalarType::BFloat16,
-    iter.dtype(), "masked_select", [&] {
-      auto mask_dtype = iter.input_dtype(1);
-      if (mask_dtype == at::ScalarType::Bool) {
-        cpu_masked_select_kernel<scalar_t, bool>(iter, [](char* dst, char* src, int64_t offset) {
-          *(scalar_t*)(dst + offset) = *(scalar_t*)src;
-        });
-      } else {
-        cpu_masked_select_kernel<scalar_t, unsigned char>(iter, [](char* dst, char* src, int64_t offset) {
-          *(scalar_t*)(dst + offset) = *(scalar_t*)src;
-        });
-      }
-    });
-}
-
 } // anonymous namespace
 
+
 REGISTER_DISPATCH(index_stub, &index_kernel);
 REGISTER_DISPATCH(index_put_stub, &index_put_kernel);
 REGISTER_DISPATCH(masked_fill_stub, &masked_fill_kernel);
-REGISTER_DISPATCH(masked_select_serial_stub, &masked_select_serial_kernel);
-REGISTER_DISPATCH(masked_select_stub, &masked_select_kernel);
 
 }} // namespace at::native
diff --git a/aten/src/TH/generic/THTensorEvenMoreMath.cpp b/aten/src/TH/generic/THTensorEvenMoreMath.cpp
@@ -99,6 +99,50 @@ accreal THTensor_(dot)(THTensor *tensor, THTensor *src)
 
 #if !defined(TH_REAL_IS_HALF) /* non half part */
 
+void THTensor_(maskedSelect)(THTensor *tensor, THTensor *src, THByteTensor *mask)
+{
+  at::NoNamesGuard guard;
+  ptrdiff_t numel = THTensor_wrap(mask).sum().item<int64_t>();
+  scalar_t *tensor_data;
+
+#ifdef DEBUG
+  THAssert(numel <= LONG_MAX);
+#endif
+  THTensor_(resize1d)(tensor,numel);
+  tensor_data = tensor->data<scalar_t>();
+  TH_TENSOR_APPLY2(scalar_t, src, unsigned char, mask,
+                   if (*mask_data > 1)
+                   {
+                     THFree(mask_counter);
+                     THFree(src_counter);
+                     THError("Mask tensor can take 0 and 1 values only");
+                   }
+                   else if (*mask_data == 1)
+                   {
+                     *tensor_data = *src_data;
+                     tensor_data++;
+                   });
+}
+
+void THTensor_(maskedSelectBool)(THTensor *tensor, THTensor *src, THBoolTensor *mask)
+{
+  at::NoNamesGuard guard;
+  ptrdiff_t numel = THTensor_wrap(mask).sum().item<int64_t>();
+  scalar_t *tensor_data;
+
+#ifdef DEBUG
+  THAssert(numel <= LONG_MAX);
+#endif
+  THTensor_(resize1d)(tensor,numel);
+  tensor_data = tensor->data<scalar_t>();
+  TH_TENSOR_APPLY2(scalar_t, src, bool, mask,
+                   if (*mask_data)
+                   {
+                     *tensor_data = *src_data;
+                     tensor_data++;
+                   });
+}
+
 void THTensor_(maskedCopy)(THTensor *tensor, THByteTensor *mask, THTensor* src )
 {
   THTensor *srct = THTensor_(newContiguous)(src);
diff --git a/aten/src/TH/generic/THTensorMath.h b/aten/src/TH/generic/THTensorMath.h
@@ -9,6 +9,8 @@ TH_API int THTensor_(equal)(THTensor *ta, THTensor *tb);
 
 #if !defined(TH_REAL_IS_HALF)
 
+TH_API void THTensor_(maskedSelect)(THTensor *tensor, THTensor* src, THByteTensor *mask);
+TH_API void THTensor_(maskedSelectBool)(THTensor *tensor, THTensor* src, THBoolTensor *mask);
 TH_API void THTensor_(maskedCopy)(THTensor *tensor, THByteTensor *mask, THTensor* src);
 TH_API void THTensor_(maskedCopyBool)(THTensor *tensor, THBoolTensor *mask, THTensor* src);
 
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -12644,7 +12644,7 @@ def test_masked_select(self, device, dtype):
             src = torch.tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype=dtype, device=device)
             mask = torch.rand(num_src, device=device).clamp(0, 1).mul(2).floor().to(maskType)
 
-            if dtype == torch.half and torch.device(device).type == 'cpu':
+            if (dtype.is_complex or dtype == torch.half) and torch.device(device).type == 'cpu':
                 self.assertRaises(RuntimeError, lambda: src.masked_select(mask))
                 continue
 
