xref: /aosp_15_r20/external/ComputeLibrary/tests/datasets/GatherDataset.h (revision c217d954acce2dbc11938adb493fc0abd69584f3)
1 /*
2  * Copyright (c) 2018-2019, 2022 Arm Limited.
3  *
4  * SPDX-License-Identifier: MIT
5  *
6  * Permission is hereby granted, free of charge, to any person obtaining a copy
7  * of this software and associated documentation files (the "Software"), to
8  * deal in the Software without restriction, including without limitation the
9  * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
10  * sell copies of the Software, and to permit persons to whom the Software is
11  * furnished to do so, subject to the following conditions:
12  *
13  * The above copyright notice and this permission notice shall be included in all
14  * copies or substantial portions of the Software.
15  *
16  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17  * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18  * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19  * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20  * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
21  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
22  * SOFTWARE.
23  */
24 
25 #ifndef ARM_COMPUTE_TEST_GATHER_DATASET
26 #define ARM_COMPUTE_TEST_GATHER_DATASET
27 
28 #include "utils/TypePrinter.h"
29 
30 #include "arm_compute/core/Types.h"
31 
32 namespace arm_compute
33 {
34 namespace test
35 {
36 namespace datasets
37 {
38 class GatherDataset
39 {
40 public:
41     using type = std::tuple<TensorShape, TensorShape, int>;
42 
43     struct iterator
44     {
iteratoriterator45         iterator(std::vector<TensorShape>::const_iterator input_shapes_it,
46                  std::vector<TensorShape>::const_iterator starts_values_it,
47                  std::vector<int>::const_iterator         axis_it)
48             : _input_shapes_it{ std::move(input_shapes_it) },
49               _indices_shapes_it{ std::move(starts_values_it) },
50               _axis_it{ std::move(axis_it) }
51         {
52         }
53 
descriptioniterator54         std::string description() const
55         {
56             std::stringstream description;
57             description << "InputShape=" << *_input_shapes_it << ":";
58             description << "IndicesShape=" << *_indices_shapes_it << ":";
59             description << "Axis=" << *_axis_it << ":";
60             return description.str();
61         }
62 
63         GatherDataset::type operator*() const
64         {
65             return std::make_tuple(*_input_shapes_it, *_indices_shapes_it, *_axis_it);
66         }
67 
68         iterator &operator++()
69         {
70             ++_input_shapes_it;
71             ++_indices_shapes_it;
72             ++_axis_it;
73             return *this;
74         }
75 
76     private:
77         std::vector<TensorShape>::const_iterator _input_shapes_it;
78         std::vector<TensorShape>::const_iterator _indices_shapes_it;
79         std::vector<int>::const_iterator         _axis_it;
80     };
81 
begin()82     iterator begin() const
83     {
84         return iterator(_input_shapes.begin(), _indices_shapes.begin(), _axis.begin());
85     }
86 
size()87     int size() const
88     {
89         return std::min(_input_shapes.size(), std::min(_indices_shapes.size(), _axis.size()));
90     }
91 
add_config(TensorShape input_shape,TensorShape indices_shape,int axis)92     void add_config(TensorShape input_shape, TensorShape indices_shape, int axis)
93     {
94         _input_shapes.emplace_back(std::move(input_shape));
95         _indices_shapes.emplace_back(std::move(indices_shape));
96         _axis.emplace_back(std::move(axis));
97     }
98 
99 protected:
100     GatherDataset()                 = default;
101     GatherDataset(GatherDataset &&) = default;
102 
103 private:
104     std::vector<TensorShape> _input_shapes{};
105     std::vector<TensorShape> _indices_shapes{};
106     std::vector<int>         _axis{};
107 };
108 
109 class SmallGatherMultiDimIndicesDataset final : public GatherDataset
110 {
111 public:
SmallGatherMultiDimIndicesDataset()112     SmallGatherMultiDimIndicesDataset()
113     {
114         add_config(TensorShape(2U, 6U), TensorShape(4U, 9U), 1);
115         add_config(TensorShape(15U, 15U), TensorShape(3U, 2U, 2U), 1);
116         add_config(TensorShape(15U, 15U), TensorShape(2U, 11U), 1);
117         add_config(TensorShape(5U, 3U, 4U), TensorShape(2U, 7U), 1);
118         add_config(TensorShape(1U, 5U, 3U), TensorShape(1U, 7U, 3U), 1);
119     }
120 };
121 
122 class SmallGatherDataset final : public GatherDataset
123 {
124 public:
SmallGatherDataset()125     SmallGatherDataset()
126     {
127         // 2D input
128         add_config(TensorShape(15U, 15U), TensorShape(5U), 0);
129         add_config(TensorShape(15U, 15U), TensorShape(5U), 1);
130         add_config(TensorShape(5U, 5U), TensorShape(80U), -1);
131 
132         // 3D input
133         add_config(TensorShape(5U, 5U, 5U), TensorShape(19U), 0);
134         add_config(TensorShape(5U, 4U, 6U), TensorShape(30U), 1);
135         add_config(TensorShape(3U, 5U, 7U), TensorShape(20U), 2);
136         add_config(TensorShape(5U, 4U, 6U), TensorShape(30U), -1);
137         add_config(TensorShape(3U, 5U, 7U), TensorShape(20U), -2);
138 
139         // 4D input
140         add_config(TensorShape(4U, 3U, 4U, 5U), TensorShape(4U), 0);
141         add_config(TensorShape(4U, 3U, 5U, 5U), TensorShape(5U), 1);
142         add_config(TensorShape(4U, 3U, 2U, 5U), TensorShape(6U), 2);
143         add_config(TensorShape(3U, 4U, 4U, 6U), TensorShape(7U), 3);
144         add_config(TensorShape(4U, 3U, 5U, 5U), TensorShape(5U), -1);
145         add_config(TensorShape(4U, 3U, 2U, 5U), TensorShape(6U), -2);
146         add_config(TensorShape(3U, 4U, 4U, 6U), TensorShape(7U), -3);
147     }
148 };
149 
150 class LargeGatherDataset final : public GatherDataset
151 {
152 public:
LargeGatherDataset()153     LargeGatherDataset()
154     {
155         // 2D input
156         add_config(TensorShape(150U, 150U), TensorShape(50U), 0);
157         add_config(TensorShape(150U, 150U), TensorShape(50U), 1);
158         add_config(TensorShape(150U, 150U), TensorShape(50U), -1);
159 
160         // 3D input
161         add_config(TensorShape(50U, 40U, 60U), TensorShape(33U), 0);
162         add_config(TensorShape(40U, 50U, 60U), TensorShape(24U), 1);
163         add_config(TensorShape(70U, 80U, 100U), TensorShape(50U), 2);
164         add_config(TensorShape(40U, 50U, 60U), TensorShape(24U), -1);
165         add_config(TensorShape(70U, 80U, 100U), TensorShape(50U), -2);
166 
167         // 4D input
168         add_config(TensorShape(30U, 40U, 20U, 20U), TensorShape(33U), 0);
169         add_config(TensorShape(23U, 10U, 60U, 20U), TensorShape(24U), 1);
170         add_config(TensorShape(14U, 20U, 10U, 31U), TensorShape(30U), 2);
171         add_config(TensorShape(34U, 10U, 40U, 20U), TensorShape(50U), 3);
172         add_config(TensorShape(23U, 10U, 60U, 20U), TensorShape(24U), -1);
173         add_config(TensorShape(14U, 20U, 10U, 31U), TensorShape(30U), -2);
174         add_config(TensorShape(34U, 10U, 40U, 20U), TensorShape(50U), -3);
175     }
176 };
177 
178 } // namespace datasets
179 } // namespace test
180 } // namespace arm_compute
181 
182 #endif /* ARM_COMPUTE_TEST_GATHER_DATASET */
183