ATen Compliance: masked_fill (Dtype, Shape & Broadcast)

Reviewed By: kirklandsign

Differential Revision: D48276291

fbshipit-source-id: 656f39a8303e5c5fb231f8f421f5881b4d2721e0

Author: manuelcandales

kernels/portable/cpu/op_masked_fill.cpp

@@ -17,87 +17,40 @@ namespace native {
 using Tensor = exec_aten::Tensor;
 using ScalarType = exec_aten::ScalarType;
 using Scalar = exec_aten::Scalar;
-namespace {
 
-/**
- * Fills the `out` with values of `self` or `value` based on mask.
- *
- * Assumes that the tensors are contiguous, are the same shape,
- * input and output have the same time and mask is tensor of bools.
- * CTYPE should be the C type (like `float` or `int`) that matches
- * the dtype of the tensors.
- */
-template <class CTYPE>
-void masked_fill_kernel(
-    const Tensor& self,
-    const Tensor& mask,
-    const Scalar& value,
-    Tensor& out) {
-  ET_DCHECK(self.numel() == mask.numel() && self.numel() == out.numel());
-  CTYPE value_v = 0;
-  bool ok = utils::extract_scalar(value, &value_v);
-  ET_CHECK_MSG(ok, "Invalid fill value: wrong type or out of range");
-  const size_t n = self.numel();
-  const auto data_self = self.const_data_ptr<CTYPE>();
-  const auto data_mask = mask.const_data_ptr<bool>();
-  auto data_out = out.mutable_data_ptr<CTYPE>();
-  for (size_t i = 0; i < n; ++i) {
-    data_out[i] = data_mask[i] ? value_v : data_self[i];
-  }
-}
-
-} // namespace
-
-/**
- * Copies `self` to `out` masking some elemnts with `value`.
- *
- * Asserts that `mask` tensor can be broadcasted to `self`, self and out should
- * have same dtype and size, and mask should be boolean tensor.
- *
- * masked_fill_Scalar_out(Tensor self, Tensor other, *, Scalar alpha=1.0,
- * Tensor(a!) out) -> Tensor(a!)
- */
 Tensor& masked_fill_scalar_out(
     RuntimeContext& ctx,
-    const Tensor& self,
+    const Tensor& in,
     const Tensor& mask,
     const Scalar& value,
     Tensor& out) {
-  ET_CHECK_MSG(
-      tensor_is_broadcastable_to(mask, self),
-      "masked_fill_scalar_out operateor can not broadcast mask to self");
+  (void)ctx;
 
-  // The mask needs to be broadcasted iff its size differnet from the target one
-  // (self.size())
-  bool broadcasted = !self.sizes().equals(mask.sizes());
-  const Tensor& broadcast_mask =
-      broadcasted ? torch::executor::broadcast_tensor(mask, self) : mask;
+  ET_KERNEL_CHECK(ctx, tensors_have_same_dtype(in, out), InvalidArgument, out);
 
-  torch::executor::Error err = resize_tensor(out, self.sizes());
-  ET_CHECK_MSG(
-      err == torch::executor::Error::Ok,
-      "Failed to resize out Tensor in masked_fill_scalar_out");
+  ScalarType in_type = in.scalar_type();
+  ScalarType mask_type = mask.scalar_type();
+  ScalarType val_type = utils::get_scalar_dtype(value);
 
-  ET_CHECK_SAME_SHAPE_AND_DTYPE2(self, out);
-  ET_CHECK_SAME_SHAPE2(self, broadcast_mask);
-  ET_CHECK_MSG(
-      broadcast_mask.scalar_type() == ScalarType::Bool, "Unexpected mask type");
+  ET_KERNEL_CHECK(ctx, mask_type == ScalarType::Bool, InvalidArgument, out);
 
-#define MASKED_FILL(ctype, dtype)                                \
-  case ScalarType::dtype:                                        \
-    masked_fill_kernel<ctype>(self, broadcast_mask, value, out); \
-    break;
+  resize_to_broadcast_target_size(in, mask, out);
 
-  switch (self.scalar_type()) {
-    ET_FORALL_REAL_TYPES_AND(Bool, MASKED_FILL)
-    default:
-      ET_CHECK_MSG(false, "Unhandled dtype %hhd", self.scalar_type());
-  }
+  ET_SWITCH_REAL_TYPES_AND(Bool, in_type, ctx, __func__, CTYPE, [&]() {
+    ET_SWITCH_REAL_TYPES_AND(Bool, val_type, ctx, __func__, CTYPE_VAL, [&]() {
+      CTYPE_VAL value_v;
+      ET_EXTRACT_SCALAR(value, value_v);
+      CTYPE val = static_cast<CTYPE>(value_v);
 
-#undef MASKED_FILL
-  if (broadcasted) {
-    free_broadcast_tensor(broadcast_mask);
-  }
+      apply_binary_elementwise_fn<CTYPE, bool, CTYPE>(
+          [val](const CTYPE val_in, const bool val_mask) {
+            return val_mask ? val : val_in;
+          },
+          in,
+          mask,
+          out);
+    });
+  });
 
   return out;
 }