Skip to content

Commit bc012c8

Browse files
manuelcandalesfacebook-github-bot
manuelcandales
authored and
facebook-github-bot
committed
Add op: masked_scatter (#6167)
Summary: Pull Request resolved: #6167 Differential Revision: D64243532
1 parent d094b09 commit bc012c8

File tree

6 files changed

+237
-0
lines changed

6 files changed

+237
-0
lines changed

‎kernels/aten/functions.yaml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -241,6 +241,8 @@
241241

242242
- op: masked_fill.Scalar_out
243243

244+
- op: masked_scatter.out
245+
244246
- op: max_pool2d_with_indices.out
245247

246248
- op: max.dim_max

‎kernels/portable/cpu/op_masked_scatter.cpp

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#include <executorch/kernels/portable/cpu/util/broadcast_util.h>
10+
#include <executorch/runtime/kernel/kernel_includes.h>
11+
12+
namespace torch {
13+
namespace executor {
14+
namespace native {
15+
16+
Tensor& masked_scatter_out(
17+
KernelRuntimeContext& ctx,
18+
const Tensor& in,
19+
const Tensor& mask,
20+
const Tensor& src,
21+
Tensor& out) {
22+
ScalarType in_type = in.scalar_type();
23+
24+
ET_KERNEL_CHECK(
25+
ctx,
26+
executorch::runtime::tensor_is_realhbbf16_type(in),
27+
InvalidArgument,
28+
out);
29+
30+
ET_KERNEL_CHECK(
31+
ctx, mask.scalar_type() == ScalarType::Bool, InvalidArgument, out);
32+
ET_KERNEL_CHECK(ctx, src.scalar_type() == in_type, InvalidArgument, out);
33+
ET_KERNEL_CHECK(ctx, out.scalar_type() == in_type, InvalidArgument, out);
34+
35+
ET_KERNEL_CHECK(
36+
ctx, tensors_have_same_dim_order(in, mask, out), InvalidArgument, out);
37+
38+
ET_KERNEL_CHECK(
39+
ctx,
40+
resize_to_broadcast_target_size(in, mask, out) == Error::Ok,
41+
InvalidArgument,
42+
out);
43+
44+
constexpr auto op_name = "masked_scatter.out";
45+
46+
int64_t idx = 0;
47+
int64_t src_numel = src.numel();
48+
bool src_numel_check = true;
49+
50+
ET_SWITCH_REALHBBF16_TYPES(in_type, ctx, op_name, CTYPE, [&]() {
51+
const CTYPE* const src_data = src.const_data_ptr<CTYPE>();
52+
apply_binary_elementwise_fn<CTYPE, bool, CTYPE>(
53+
[src_data, &idx, &src_numel, &src_numel_check](
54+
const CTYPE val_in, const bool val_mask) {
55+
if (val_mask && idx >= src_numel) {
56+
src_numel_check = false;
57+
return val_in;
58+
}
59+
return val_mask ? src_data[idx++] : val_in;
60+
},
61+
in,
62+
mask,
63+
out);
64+
});
65+
66+
ET_KERNEL_CHECK_MSG(
67+
ctx,
68+
src_numel_check,
69+
InvalidArgument,
70+
out,
71+
"masked_scatter: src doesn't have enough elements");
72+
73+
return out;
74+
}
75+
76+
} // namespace native
77+
} // namespace executor
78+
} // namespace torch

‎kernels/portable/functions.yaml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -542,6 +542,11 @@
542542
- arg_meta: null
543543
kernel_name: torch::executor::masked_fill_scalar_out
544544

545+
- op: masked_scatter.out
546+
kernels:
547+
- arg_meta: null
548+
kernel_name: torch::executor::masked_scatter_out
549+
545550
- op: max.dim_max
546551
kernels:
547552
- arg_meta: null

‎kernels/test/op_masked_scatter_test.cpp

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#include <executorch/kernels/test/FunctionHeaderWrapper.h> // Declares the operator
10+
#include <executorch/kernels/test/TestUtil.h>
11+
#include <executorch/kernels/test/supported_features.h>
12+
#include <executorch/runtime/core/exec_aten/exec_aten.h>
13+
#include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
14+
#include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
15+
16+
#include <gtest/gtest.h>
17+
18+
using namespace ::testing;
19+
using exec_aten::ScalarType;
20+
using exec_aten::Tensor;
21+
using torch::executor::testing::SupportedFeatures;
22+
using torch::executor::testing::TensorFactory;
23+
24+
class OpMaskedScatterOutTest : public OperatorTest {
25+
protected:
26+
Tensor& op_masked_scatter_out(
27+
const Tensor& in,
28+
const Tensor& mask,
29+
const Tensor& src,
30+
Tensor& out) {
31+
return torch::executor::aten::masked_scatter_outf(
32+
context_, in, mask, src, out);
33+
}
34+
};
35+
36+
TEST_F(OpMaskedScatterOutTest, SmokeTest) {
37+
TensorFactory<ScalarType::Int> tf;
38+
TensorFactory<ScalarType::Bool> tfBool;
39+
40+
Tensor in = tf.make({2, 3}, {1, 2, 3, 4, 5, 6});
41+
Tensor mask = tfBool.make({2, 3}, {true, false, false, true, false, true});
42+
Tensor src = tf.make({3}, {10, 20, 30});
43+
44+
Tensor out = tf.zeros({2, 3});
45+
46+
op_masked_scatter_out(in, mask, src, out);
47+
EXPECT_TENSOR_EQ(out, tf.make({2, 3}, {10, 2, 3, 20, 5, 30}));
48+
}
49+
50+
TEST_F(OpMaskedScatterOutTest, BroadcastInput) {
51+
TensorFactory<ScalarType::Int> tf;
52+
TensorFactory<ScalarType::Bool> tfBool;
53+
54+
Tensor in = tf.make({3}, {1, 2, 3});
55+
Tensor mask = tfBool.make({2, 3}, {true, false, false, true, false, true});
56+
Tensor src = tf.make({3}, {10, 20, 30});
57+
58+
Tensor out = tf.zeros({2, 3});
59+
60+
op_masked_scatter_out(in, mask, src, out);
61+
EXPECT_TENSOR_EQ(out, tf.make({2, 3}, {10, 2, 3, 20, 2, 30}));
62+
}
63+
64+
TEST_F(OpMaskedScatterOutTest, BroadcastMask) {
65+
TensorFactory<ScalarType::Int> tf;
66+
TensorFactory<ScalarType::Bool> tfBool;
67+
68+
Tensor in = tf.make({2, 3}, {1, 2, 3, 4, 5, 6});
69+
Tensor mask = tfBool.make({3}, {false, true, false});
70+
Tensor src = tf.make({2}, {10, 20});
71+
72+
Tensor out = tf.zeros({2, 3});
73+
74+
op_masked_scatter_out(in, mask, src, out);
75+
EXPECT_TENSOR_EQ(out, tf.make({2, 3}, {1, 10, 3, 4, 20, 6}));
76+
}
77+
78+
TEST_F(OpMaskedScatterOutTest, SrcWithMoreElements) {
79+
TensorFactory<ScalarType::Int> tf;
80+
TensorFactory<ScalarType::Bool> tfBool;
81+
82+
Tensor in = tf.make({2, 3}, {1, 2, 3, 4, 5, 6});
83+
Tensor mask = tfBool.make({2, 3}, {true, false, false, true, false, true});
84+
Tensor src = tf.make({4}, {10, 20, 30, 40});
85+
86+
Tensor out = tf.zeros({2, 3});
87+
88+
op_masked_scatter_out(in, mask, src, out);
89+
EXPECT_TENSOR_EQ(out, tf.make({2, 3}, {10, 2, 3, 20, 5, 30}));
90+
}
91+
92+
TEST_F(OpMaskedScatterOutTest, SrcWithLessElementsFails) {
93+
TensorFactory<ScalarType::Int> tf;
94+
TensorFactory<ScalarType::Bool> tfBool;
95+
96+
Tensor in = tf.make({2, 3}, {1, 2, 3, 4, 5, 6});
97+
Tensor mask = tfBool.make({2, 3}, {true, false, false, true, false, true});
98+
Tensor src = tf.make({2}, {10, 20});
99+
100+
Tensor out = tf.zeros({2, 3});
101+
102+
ET_EXPECT_KERNEL_FAILURE(context_, op_masked_scatter_out(in, mask, src, out));
103+
}
104+
105+
TEST_F(OpMaskedScatterOutTest, EmptyMask) {
106+
TensorFactory<ScalarType::Int> tf;
107+
TensorFactory<ScalarType::Bool> tfBool;
108+
109+
Tensor in = tf.make({2, 1}, {100, 200});
110+
Tensor mask = tfBool.make({2, 0}, {});
111+
Tensor src = tf.make({4}, {10, 20, 30, 40});
112+
113+
Tensor out = tf.zeros({2, 0});
114+
115+
op_masked_scatter_out(in, mask, src, out);
116+
EXPECT_TENSOR_EQ(out, tf.make({2, 0}, {}));
117+
}
118+
119+
TEST_F(OpMaskedScatterOutTest, EmptySrc) {
120+
TensorFactory<ScalarType::Int> tf;
121+
TensorFactory<ScalarType::Bool> tfBool;
122+
123+
Tensor in = tf.make({2, 1}, {100, 200});
124+
Tensor mask = tfBool.make({2, 1}, {false, false});
125+
Tensor src = tf.make({0}, {});
126+
127+
Tensor out = tf.zeros({2, 1});
128+
129+
op_masked_scatter_out(in, mask, src, out);
130+
EXPECT_TENSOR_EQ(out, tf.make({2, 1}, {100, 200}));
131+
}
132+
133+
TEST_F(OpMaskedScatterOutTest, EmptyMaskAndSrc) {
134+
TensorFactory<ScalarType::Int> tf;
135+
TensorFactory<ScalarType::Bool> tfBool;
136+
137+
Tensor in = tf.make({2, 1}, {100, 200});
138+
Tensor mask = tfBool.make({0}, {});
139+
Tensor src = tf.make({0}, {});
140+
141+
Tensor out = tf.zeros({2, 0});
142+
143+
op_masked_scatter_out(in, mask, src, out);
144+
EXPECT_TENSOR_EQ(out, tf.make({2, 0}, {}));
145+
}

‎kernels/test/targets.bzl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -254,6 +254,7 @@ def define_common_targets():
254254
_common_op_test("op_logit_test", ["aten", "portable"])
255255
_common_op_test("op_lt_test", ["aten", "portable"])
256256
_common_op_test("op_masked_fill_test", ["aten", "portable"])
257+
_common_op_test("op_masked_scatter_test", ["aten", "portable"])
257258
_common_op_test("op_max_test", ["aten", "portable"])
258259
_common_op_test("op_max_pool2d_with_indices_test", ["aten", "portable"])
259260
_common_op_test("op_maximum_test", ["aten", "portable"])

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -776,6 +776,12 @@ ATEN_OPS = (
776776
":scalar_utils",
777777
],
778778
),
779+
op_target(
780+
name = "op_masked_scatter",
781+
deps = [
782+
"//executorch/kernels/portable/cpu/util:broadcast_util",
783+
],
784+
),
779785
op_target(
780786
name = "op_max",
781787
deps = [

0 commit comments

Comments
 (0)