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
masked_scatter_out(KernelRuntimeContext & ctx,const Tensor & in,const Tensor & mask,const Tensor & src,Tensor & out)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
79