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 #pragma once 10 11 #include <executorch/runtime/core/exec_aten/exec_aten.h> 12 #include <executorch/runtime/core/exec_aten/util/tensor_util.h> 13 14 namespace torch { 15 namespace executor { 16 17 // 18 // Reduction 19 // 20 21 /** 22 * For `size` elements of `data_in`, accumulates the modified values into a 23 * value using `reduce_fun`, and returns the accumulated value. The `stride` can 24 * also be defined; by default it is set to 1. 25 */ 26 template <typename CTYPE, typename ReduceOp> 27 inline CTYPE apply_unary_reduce_fn( 28 const ReduceOp& reduce_fun, 29 const CTYPE* const data_in, 30 const int64_t size, 31 const int64_t stride = 1) { 32 CTYPE acc_val = data_in[0]; 33 for (size_t i = 1; i < size; i++) { 34 acc_val = reduce_fun(data_in[i * stride], acc_val); 35 } 36 return acc_val; 37 } 38 39 // 40 // Mapping 41 // 42 43 /** 44 * Applies `map_fun` to `size` elements of `data_in`, writing results to 45 * `data_out`. The `stride` can also be defined; by default it is set to 1. 46 */ 47 template <typename CTYPE_IN, typename CTYPE_OUT, typename MapOp> 48 inline void apply_unary_map_fn( 49 const MapOp& map_fun, 50 const CTYPE_IN* const data_in, 51 CTYPE_OUT* const data_out, 52 const int64_t size, 53 const int64_t stride = 1) { 54 for (size_t i = 0; i < size; i++) { 55 data_out[i * stride] = map_fun(data_in[i * stride]); 56 } 57 } 58 59 // 60 // Mapping + Reduction 61 // 62 63 /** 64 * Applies `map_fun` to `size` elements of `data_in`, accumulates the modified 65 * values into a value using `reduce_fun`, and returns the accumulated value. 66 * The `stride` can also be defined; by default it is set to 1. 67 */ 68 template < 69 typename CTYPE_IN, 70 typename CTYPE_OUT, 71 typename MapOp, 72 typename ReduceOp> 73 inline CTYPE_OUT apply_unary_map_reduce_fn( 74 const MapOp& map_fun, 75 const ReduceOp& reduce_fun, 76 const CTYPE_IN* const data_in, 77 const int64_t size, 78 const int64_t stride = 1) { 79 CTYPE_OUT acc_val = map_fun(data_in[0]); 80 for (size_t i = 1; i < size; ++i) { 81 acc_val = reduce_fun(map_fun(data_in[i * stride]), acc_val); 82 } 83 return acc_val; 84 } 85 86 } // namespace executor 87 } // namespace torch 88