xref: /aosp_15_r20/external/executorch/kernels/portable/cpu/pattern/pattern.h (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
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 /* Naming Convention
10 --------------------
11 Each pattern name should be all lowercase words separated by underscores.
12 
13 It should start with unary/binary/ternary depending on the number of input
14 tensors.
15 
16 It should then specify the type of function that it is. E.g: ufunc, reduce, etc
17 Ufunc is a concept used in NumPy and PyTorch (short for universal function)
18 See NumPy's definition here: https://numpy.org/doc/stable/reference/ufuncs.html
19 
20 It should then specify the sets of dtypes accepted as inputs and output.
21 We currently support the following sets of dtypes:
22  - bool   {Bool}
23  - int    {Byte, Char, Short, Int, Long}
24  - intb   {Byte, Char, Short, Int, Long, Bool}
25  - float  {Float, Double}
26  - floath {Half, Float, Double}
27  - floatb {Float, Double, Bool}
28  - real   {Byte, Char, Short, Int, Long, Float, Double}
29  - realb  {Byte, Char, Short, Int, Long, Float, Double, Bool}
30  - realh  {Byte, Char, Short, Int, Long, Half, Float, Double}
31  - realhb {Byte, Char, Short, Int, Long, Half, Float, Double, Bool}
32 
33 Input types are separated from output types by the "to" word.
34 Input types are separated by underscores. Output types as well, in the cases
35 when there are multiple outputs.
36 In the case of ops that only allow the same dtype (for example abs/neg where
37 the input and output dtype must be the same, only that dtype is specified.
38 E.g, there is a difference between:
39  - unary_ufunc_real_to_real
40      (output and input types must be real, but might be different)
41  - unary_ufunc_real
42      (output must be same type than input, which must be real)
43 
44 A pattern named using this convention will be quite generic. If the pattern in
45 question is a bit more specific, then add a descriptive sufix. */
46 
47 #pragma once
48 
49 #include <executorch/runtime/kernel/kernel_includes.h>
50 
51 namespace torch {
52 namespace executor {
53 namespace native {
54 namespace internal {
55 
56 /**
57  * Implements an op pattern for ops that take a single input tensor of any
58  * realh dtye, no additional arguments, and outputs a tensor of the same size
59  * and dtype. The function fn specifies the math operation which is applied to
60  * the input tensor element-wise.
61  */
62 Tensor& unary_ufunc_realh(
63     double (*fn)(double),
64     KernelRuntimeContext& ctx,
65     const Tensor& in,
66     Tensor& out);
67 
68 /**
69  * Implements an op pattern for ops that take a single input tensor of any
70  * realhb dtye (real, half and boolean), no additional arguments, and outputs a
71  * boolean tensor of the same size. The function fn specifies the math
72  * operation which is applied to the input tensor element-wise.
73  */
74 Tensor& unary_ufunc_realhb_to_bool(
75     bool (*fn)(double),
76     KernelRuntimeContext& ctx,
77     const Tensor& in,
78     Tensor& out);
79 
80 /**
81  * Implements an op pattern for ops that take a single input tensor of any
82  * realhbbf16 dtype (real/half/bool/bfloat16), no additional arguments, and
83  * outputs a floating point tensor of the same size. The function fn specifies
84  * the math operation which is applied to the input tensor element-wise.
85  */
86 Tensor& unary_ufunc_realhbbf16_to_floathbf16(
87     double (*fn)(double),
88     KernelRuntimeContext& ctx,
89     const Tensor& in,
90     Tensor& out);
91 
92 } // namespace internal
93 } // namespace native
94 } // namespace executor
95 } // namespace torch
96