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/scalar_utils.h>
12 #include <executorch/kernels/portable/cpu/util/functional_util.h>
13 #include <executorch/kernels/portable/cpu/util/math_util.h>
14 #include <executorch/runtime/kernel/kernel_includes.h>
15
16 namespace torch {
17 namespace executor {
18 namespace native {
19
20 using Tensor = exec_aten::Tensor;
21 using ScalarType = exec_aten::ScalarType;
22
hardtanh_out(KernelRuntimeContext & ctx,const Tensor & in,const Scalar & min,const Scalar & max,Tensor & out)23 Tensor& hardtanh_out(
24 KernelRuntimeContext& ctx,
25 const Tensor& in,
26 const Scalar& min,
27 const Scalar& max,
28 Tensor& out) {
29 (void)ctx;
30
31 // Resize for dynamic shape
32 ET_KERNEL_CHECK_MSG(
33 ctx,
34 resize_tensor(out, in.sizes()) == Error::Ok,
35 InvalidArgument,
36 out,
37 "Failed to resize output tensor.");
38
39 ET_KERNEL_CHECK(
40 ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
41
42 ScalarType in_type = in.scalar_type();
43 ScalarType min_type = utils::get_scalar_dtype(min);
44 ScalarType max_type = utils::get_scalar_dtype(max);
45 ScalarType out_type = out.scalar_type();
46
47 ET_KERNEL_CHECK(ctx, in_type == out_type, InvalidArgument, out);
48
49 ET_SWITCH_REAL_TYPES(in_type, ctx, "hardtanh.out", CTYPE, [&]() {
50 CTYPE min_casted;
51 ET_SWITCH_SCALAR_OBJ_TYPES(min_type, ctx, "hardtanh.out", CTYPE_MIN, [&]() {
52 CTYPE_MIN min_val;
53 utils::extract_scalar(min, &min_val);
54 min_casted = static_cast<CTYPE>(min_val);
55 });
56
57 CTYPE max_casted;
58 ET_SWITCH_SCALAR_OBJ_TYPES(max_type, ctx, "hardtanh.out", CTYPE_MAX, [&]() {
59 CTYPE_MAX max_val;
60 utils::extract_scalar(max, &max_val);
61 max_casted = static_cast<CTYPE>(max_val);
62 });
63
64 apply_unary_map_fn(
65 [min_casted, max_casted](const CTYPE val_in) {
66 return utils::min_override(
67 utils::max_override(val_in, min_casted), max_casted);
68 },
69 in.const_data_ptr<CTYPE>(),
70 out.mutable_data_ptr<CTYPE>(),
71 in.numel());
72 });
73
74 return out;
75 }
76
77 } // namespace native
78 } // namespace executor
79 } // namespace torch
80