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 <cmath>
10
11 #include <executorch/kernels/portable/cpu/util/functional_util.h>
12 #include <executorch/runtime/kernel/kernel_includes.h>
13 #include <executorch/runtime/platform/assert.h>
14
15 namespace torch {
16 namespace executor {
17 namespace native {
18
19 using exec_aten::Tensor;
20
21 namespace {
22
23 // Rounds a floating point value to the closest integer. Values with a
24 // fractional part of exactly 0.5 are rounded to the closest even integer. Uses
25 // the implementation from torch/src/jit/runtime/register_ops_utils.h.
26 template <typename CTYPE>
round_to_even(CTYPE a)27 inline CTYPE round_to_even(CTYPE a) {
28 return a - std::floor(a) == 0.5 ? (std::round(a * 0.5) * 2.0) : std::round(a);
29 }
30
31 } // namespace
32
round_out(KernelRuntimeContext & ctx,const Tensor & in,Tensor & out)33 Tensor& round_out(KernelRuntimeContext& ctx, const Tensor& in, Tensor& out) {
34 (void)ctx;
35
36 // Resize for dynamic shape
37 ET_KERNEL_CHECK_MSG(
38 ctx,
39 resize_tensor(out, in.sizes()) == Error::Ok,
40 InvalidArgument,
41 out,
42 "Failed to resize output tensor.");
43
44 ET_KERNEL_CHECK(
45 ctx, tensors_have_same_shape_and_dtype(in, out), InvalidArgument, out);
46 ET_KERNEL_CHECK(ctx, tensor_is_real_type(out), InvalidArgument, out);
47
48 ET_KERNEL_CHECK(
49 ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
50
51 auto in_scalar_type = in.scalar_type();
52
53 ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, "round.out", CTYPE, [&] {
54 apply_unary_map_fn(
55 [in_scalar_type](const CTYPE val_in) {
56 if (isIntegralType(in_scalar_type, /*includeBool=*/false)) {
57 return val_in;
58 } else {
59 return static_cast<CTYPE>(round_to_even<CTYPE>(val_in));
60 }
61 },
62 in.const_data_ptr<CTYPE>(),
63 out.mutable_data_ptr<CTYPE>(),
64 in.numel());
65 });
66
67 return out;
68 }
69
70 } // namespace native
71 } // namespace executor
72 } // namespace torch
73