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/reduce_util.h>
10 #include <executorch/runtime/kernel/kernel_includes.h>
11 #include <executorch/runtime/platform/assert.h>
12
13 namespace torch {
14 namespace executor {
15 namespace native {
16 namespace {
17
check_flip_args(const Tensor & in,IntArrayRef dims,const Tensor & out)18 bool check_flip_args(const Tensor& in, IntArrayRef dims, const Tensor& out) {
19 ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
20 return check_dim_list_is_valid(in, dims);
21 }
22
unflip_flat_ix(size_t ix,const Tensor & in,ArrayRef<bool> flip_dim)23 size_t unflip_flat_ix(size_t ix, const Tensor& in, ArrayRef<bool> flip_dim) {
24 size_t ix_coord[kTensorDimensionLimit];
25 indexToCoordinate(in, ix, ix_coord);
26
27 size_t unflip_coord[kTensorDimensionLimit];
28 for (size_t d = 0; d < in.dim(); d++) {
29 if (flip_dim[d]) {
30 unflip_coord[d] = in.size(d) - ix_coord[d] - 1;
31 } else {
32 unflip_coord[d] = ix_coord[d];
33 }
34 }
35
36 return coordinateToIndex(in, unflip_coord);
37 }
38
39 } // namespace
40
flip_out(KernelRuntimeContext & ctx,const Tensor & in,IntArrayRef dims,Tensor & out)41 Tensor& flip_out(
42 KernelRuntimeContext& ctx,
43 const Tensor& in,
44 IntArrayRef dims,
45 Tensor& out) {
46 (void)ctx;
47
48 ET_KERNEL_CHECK(
49 ctx, resize_tensor(out, in.sizes()) == Error::Ok, InvalidArgument, out);
50
51 ET_KERNEL_CHECK(
52 ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
53
54 ET_KERNEL_CHECK(ctx, check_flip_args(in, dims, out), InvalidArgument, out);
55
56 bool flip_dim_data[kTensorDimensionLimit];
57 for (size_t i = 0; i < in.dim(); i++) {
58 flip_dim_data[i] = false;
59 }
60 for (size_t i = 0; i < dims.size(); i++) {
61 const auto d = dims[i] < 0 ? dims[i] + nonzero_dim(in) : dims[i];
62 flip_dim_data[d] = true;
63 }
64 size_t flip_dim_length = static_cast<size_t>(in.dim()); // NOLINT
65 ArrayRef<bool> flip_dim(flip_dim_data, flip_dim_length);
66
67 constexpr auto name = "flip.out";
68
69 ET_SWITCH_REALHB_TYPES(in.scalar_type(), ctx, name, CTYPE, [&] {
70 const CTYPE* in_data = in.const_data_ptr<CTYPE>();
71 CTYPE* out_data = out.mutable_data_ptr<CTYPE>();
72
73 for (size_t ix = 0; ix < out.numel(); ++ix) {
74 out_data[ix] = in_data[unflip_flat_ix(ix, in, flip_dim)];
75 }
76 });
77
78 return out;
79 }
80
81 } // namespace native
82 } // namespace executor
83 } // namespace torch
84