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/optimized/vec/functional.h>
12 #include <executorch/kernels/optimized/vec/vec.h>
13 #include <executorch/runtime/kernel/kernel_includes.h>
14
15 namespace torch {
16 namespace executor {
17 namespace native {
18
19 namespace {
20
21 /**
22 * Fast path of natural exponential function. When no casting is required, CPU
23 * vector intrinsics can be used.
24 */
25 template <
26 typename CTYPE_IN,
27 typename CTYPE_OUT,
28 typename std::enable_if<
29 std::is_same_v<CTYPE_IN, CTYPE_OUT> &&
30 !std::is_same_v<CTYPE_IN, exec_aten::Half> &&
31 !std::is_same_v<CTYPE_OUT, exec_aten::BFloat16>,
32 int>::type = 0>
exp_data(const CTYPE_IN * in_data,const size_t numel,CTYPE_OUT * out_data)33 void exp_data(
34 const CTYPE_IN* in_data,
35 const size_t numel,
36 CTYPE_OUT* out_data) {
37 using Vec = executorch::vec::Vectorized<CTYPE_IN>;
38 executorch::vec::map<CTYPE_IN>(
39 [](Vec x) { return x.exp(); }, out_data, in_data, numel);
40 }
41
42 /**
43 * Slow path of natural exponential function.
44 */
45 template <
46 typename CTYPE_IN,
47 typename CTYPE_OUT,
48 typename std::enable_if<
49 !std::is_same_v<CTYPE_IN, CTYPE_OUT> ||
50 std::is_same_v<CTYPE_IN, exec_aten::Half> ||
51 std::is_same_v<CTYPE_IN, exec_aten::BFloat16> ||
52 std::is_same_v<CTYPE_OUT, exec_aten::Half> ||
53 std::is_same_v<CTYPE_OUT, exec_aten::BFloat16>,
54 int>::type = 0>
exp_data(const CTYPE_IN * in_data,const size_t numel,CTYPE_OUT * out_data)55 void exp_data(
56 const CTYPE_IN* in_data,
57 const size_t numel,
58 CTYPE_OUT* out_data) {
59 for (size_t i = 0; i < numel; i++) {
60 CTYPE_OUT xi = static_cast<CTYPE_OUT>(in_data[i]);
61 out_data[i] = std::exp(xi);
62 }
63 }
64
65 } // namespace
66
opt_exp_out(KernelRuntimeContext & ctx,const Tensor & in,Tensor & out)67 Tensor& opt_exp_out(KernelRuntimeContext& ctx, const Tensor& in, Tensor& out) {
68 (void)ctx;
69
70 // Resize for dynamic shape
71 auto error = resize_tensor(out, in.sizes());
72 ET_KERNEL_CHECK_MSG(
73 ctx,
74 error == Error::Ok,
75 InvalidArgument,
76 out,
77 "Failed to resize output tensor.");
78
79 ET_KERNEL_CHECK(ctx, tensor_is_floating_type(out), InvalidArgument, out);
80
81 ET_SWITCH_REALHBBF16_TYPES(in.scalar_type(), ctx, "exp.out", CTYPE_IN, [&] {
82 ET_SWITCH_FLOATHBF16_TYPES(
83 out.scalar_type(), ctx, "exp.out", CTYPE_OUT, [&] {
84 exp_data<CTYPE_IN, CTYPE_OUT>(
85 in.const_data_ptr<CTYPE_IN>(),
86 in.numel(),
87 out.mutable_data_ptr<CTYPE_OUT>());
88 });
89 });
90
91 return out;
92 }
93
94 } // namespace native
95 } // namespace executor
96 } // namespace torch
97