Add op: masked_scatter

kernels/aten/functions.yaml

@@ -241,6 +241,8 @@
 
 - op: masked_fill.Scalar_out
 
+- op: masked_scatter.out
+
 - op: max_pool2d_with_indices.out
 
 - op: max.dim_max

kernels/portable/cpu/op_masked_scatter.cpp

@@ -0,0 +1,78 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/kernels/portable/cpu/util/broadcast_util.h>
+#include <executorch/runtime/kernel/kernel_includes.h>
+
+namespace torch {
+namespace executor {
+namespace native {
+
+Tensor& masked_scatter_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& in,
+    const Tensor& mask,
+    const Tensor& src,
+    Tensor& out) {
+  ScalarType in_type = in.scalar_type();
+
+  ET_KERNEL_CHECK(
+      ctx,
+      executorch::runtime::tensor_is_realhbbf16_type(in),
+      InvalidArgument,
+      out);
+
+  ET_KERNEL_CHECK(
+      ctx, mask.scalar_type() == ScalarType::Bool, InvalidArgument, out);
+  ET_KERNEL_CHECK(ctx, src.scalar_type() == in_type, InvalidArgument, out);
+  ET_KERNEL_CHECK(ctx, out.scalar_type() == in_type, InvalidArgument, out);
+
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(in, mask, out), InvalidArgument, out);
+
+  ET_KERNEL_CHECK(
+      ctx,
+      resize_to_broadcast_target_size(in, mask, out) == Error::Ok,
+      InvalidArgument,
+      out);
+
+  constexpr auto op_name = "masked_scatter.out";
+
+  int64_t idx = 0;
+  int64_t src_numel = src.numel();
+  bool src_numel_check = true;
+
+  ET_SWITCH_REALHBBF16_TYPES(in_type, ctx, op_name, CTYPE, [&]() {
+    const CTYPE* const src_data = src.const_data_ptr<CTYPE>();
+    apply_binary_elementwise_fn<CTYPE, bool, CTYPE>(
+        [src_data, &idx, &src_numel, &src_numel_check](
+            const CTYPE val_in, const bool val_mask) {
+          if (val_mask && idx >= src_numel) {
+            src_numel_check = false;
+            return val_in;
+          }
+          return val_mask ? src_data[idx++] : val_in;
+        },
+        in,
+        mask,
+        out);
+  });
+
+  ET_KERNEL_CHECK_MSG(
+      ctx,
+      src_numel_check,
+      InvalidArgument,
+      out,
+      "masked_scatter: src doesn't have enough elements");
+
+  return out;
+}
+
+} // namespace native
+} // namespace executor
+} // namespace torch

kernels/portable/functions.yaml

@@ -542,6 +542,11 @@
     - arg_meta: null
       kernel_name: torch::executor::masked_fill_scalar_out
 
+- op: masked_scatter.out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::masked_scatter_out
+
 - op: max.dim_max
   kernels:
     - arg_meta: null

kernels/test/op_masked_scatter_test.cpp

@@ -0,0 +1,145 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/kernels/test/FunctionHeaderWrapper.h> // Declares the operator
+#include <executorch/kernels/test/TestUtil.h>
+#include <executorch/kernels/test/supported_features.h>
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
+#include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
+
+#include <gtest/gtest.h>
+
+using namespace ::testing;
+using exec_aten::ScalarType;
+using exec_aten::Tensor;
+using torch::executor::testing::SupportedFeatures;
+using torch::executor::testing::TensorFactory;
+
+class OpMaskedScatterOutTest : public OperatorTest {
+ protected:
+  Tensor& op_masked_scatter_out(
+      const Tensor& in,
+      const Tensor& mask,
+      const Tensor& src,
+      Tensor& out) {
+    return torch::executor::aten::masked_scatter_outf(
+        context_, in, mask, src, out);
+  }
+};
+
+TEST_F(OpMaskedScatterOutTest, SmokeTest) {
+  TensorFactory<ScalarType::Int> tf;
+  TensorFactory<ScalarType::Bool> tfBool;
+
+  Tensor in = tf.make({2, 3}, {1, 2, 3, 4, 5, 6});
+  Tensor mask = tfBool.make({2, 3}, {true, false, false, true, false, true});
+  Tensor src = tf.make({3}, {10, 20, 30});
+
+  Tensor out = tf.zeros({2, 3});
+
+  op_masked_scatter_out(in, mask, src, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 3}, {10, 2, 3, 20, 5, 30}));
+}
+
+TEST_F(OpMaskedScatterOutTest, BroadcastInput) {
+  TensorFactory<ScalarType::Int> tf;
+  TensorFactory<ScalarType::Bool> tfBool;
+
+  Tensor in = tf.make({3}, {1, 2, 3});
+  Tensor mask = tfBool.make({2, 3}, {true, false, false, true, false, true});
+  Tensor src = tf.make({3}, {10, 20, 30});
+
+  Tensor out = tf.zeros({2, 3});
+
+  op_masked_scatter_out(in, mask, src, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 3}, {10, 2, 3, 20, 2, 30}));
+}
+
+TEST_F(OpMaskedScatterOutTest, BroadcastMask) {
+  TensorFactory<ScalarType::Int> tf;
+  TensorFactory<ScalarType::Bool> tfBool;
+
+  Tensor in = tf.make({2, 3}, {1, 2, 3, 4, 5, 6});
+  Tensor mask = tfBool.make({3}, {false, true, false});
+  Tensor src = tf.make({2}, {10, 20});
+
+  Tensor out = tf.zeros({2, 3});
+
+  op_masked_scatter_out(in, mask, src, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 3}, {1, 10, 3, 4, 20, 6}));
+}
+
+TEST_F(OpMaskedScatterOutTest, SrcWithMoreElements) {
+  TensorFactory<ScalarType::Int> tf;
+  TensorFactory<ScalarType::Bool> tfBool;
+
+  Tensor in = tf.make({2, 3}, {1, 2, 3, 4, 5, 6});
+  Tensor mask = tfBool.make({2, 3}, {true, false, false, true, false, true});
+  Tensor src = tf.make({4}, {10, 20, 30, 40});
+
+  Tensor out = tf.zeros({2, 3});
+
+  op_masked_scatter_out(in, mask, src, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 3}, {10, 2, 3, 20, 5, 30}));
+}
+
+TEST_F(OpMaskedScatterOutTest, SrcWithLessElementsFails) {
+  TensorFactory<ScalarType::Int> tf;
+  TensorFactory<ScalarType::Bool> tfBool;
+
+  Tensor in = tf.make({2, 3}, {1, 2, 3, 4, 5, 6});
+  Tensor mask = tfBool.make({2, 3}, {true, false, false, true, false, true});
+  Tensor src = tf.make({2}, {10, 20});
+
+  Tensor out = tf.zeros({2, 3});
+
+  ET_EXPECT_KERNEL_FAILURE(context_, op_masked_scatter_out(in, mask, src, out));
+}
+
+TEST_F(OpMaskedScatterOutTest, EmptyMask) {
+  TensorFactory<ScalarType::Int> tf;
+  TensorFactory<ScalarType::Bool> tfBool;
+
+  Tensor in = tf.make({2, 1}, {100, 200});
+  Tensor mask = tfBool.make({2, 0}, {});
+  Tensor src = tf.make({4}, {10, 20, 30, 40});
+
+  Tensor out = tf.zeros({2, 0});
+
+  op_masked_scatter_out(in, mask, src, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 0}, {}));
+}
+
+TEST_F(OpMaskedScatterOutTest, EmptySrc) {
+  TensorFactory<ScalarType::Int> tf;
+  TensorFactory<ScalarType::Bool> tfBool;
+
+  Tensor in = tf.make({2, 1}, {100, 200});
+  Tensor mask = tfBool.make({2, 1}, {false, false});
+  Tensor src = tf.make({0}, {});
+
+  Tensor out = tf.zeros({2, 1});
+
+  op_masked_scatter_out(in, mask, src, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 1}, {100, 200}));
+}
+
+TEST_F(OpMaskedScatterOutTest, EmptyMaskAndSrc) {
+  TensorFactory<ScalarType::Int> tf;
+  TensorFactory<ScalarType::Bool> tfBool;
+
+  Tensor in = tf.make({2, 1}, {100, 200});
+  Tensor mask = tfBool.make({0}, {});
+  Tensor src = tf.make({0}, {});
+
+  Tensor out = tf.zeros({2, 0});
+
+  op_masked_scatter_out(in, mask, src, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 0}, {}));
+}

kernels/test/targets.bzl

@@ -254,6 +254,7 @@ def define_common_targets():
     _common_op_test("op_logit_test", ["aten", "portable"])
     _common_op_test("op_lt_test", ["aten", "portable"])
     _common_op_test("op_masked_fill_test", ["aten", "portable"])
+    _common_op_test("op_masked_scatter_test", ["aten", "portable"])
     _common_op_test("op_max_test", ["aten", "portable"])
     _common_op_test("op_max_pool2d_with_indices_test", ["aten", "portable"])
     _common_op_test("op_maximum_test", ["aten", "portable"])

shim/xplat/executorch/kernels/portable/op_registration_util.bzl

@@ -776,6 +776,12 @@ ATEN_OPS = (
             ":scalar_utils",
         ],
     ),
+    op_target(
+        name = "op_masked_scatter",
+        deps = [
+            "//executorch/kernels/portable/cpu/util:broadcast_util",
+        ],
+    ),
     op_target(
         name = "op_max",
         deps = [