xref: /aosp_15_r20/external/executorch/kernels/portable/cpu/util/slice_util.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 #pragma once
10 
11 #include <executorch/runtime/kernel/kernel_includes.h>
12 
13 namespace torch {
14 namespace executor {
15 
16 bool check_narrow_copy_args(
17     const Tensor& in,
18     int64_t dim,
19     int64_t start,
20     int64_t length,
21     Tensor& out);
22 
23 void get_narrow_copy_out_target_size(
24     const Tensor& in,
25     int64_t dim,
26     int64_t length,
27     exec_aten::SizesType* out_sizes,
28     size_t* out_ndim);
29 
30 bool check_slice_copy_args(
31     const Tensor& in,
32     int64_t dim,
33     int64_t step,
34     Tensor& out);
35 
36 void get_slice_copy_out_target_size(
37     const Tensor& in,
38     int64_t dim,
39     int64_t num_values,
40     exec_aten::SizesType* out_sizes,
41     size_t* out_ndim);
42 
43 bool check_slice_scatter_args(
44     const Tensor& input,
45     const Tensor& src,
46     int64_t dim,
47     int64_t num_values,
48     int64_t step,
49     Tensor output);
50 
51 int64_t adjust_slice_indices(
52     int64_t dim_length,
53     int64_t* start,
54     int64_t* end,
55     int64_t step);
56 
57 void compute_slice(
58     const Tensor& in,
59     int64_t dim,
60     int64_t start,
61     int64_t length,
62     int64_t step,
63     Tensor& out);
64 
65 } // namespace executor
66 } // namespace torch
67