1 //*****************************************************************************
2 // Copyright 2017-2020 Intel Corporation
4 // Licensed under the Apache License, Version 2.0 (the "License");
5 // you may not use this file except in compliance with the License.
6 // You may obtain a copy of the License at
8 // http://www.apache.org/licenses/LICENSE-2.0
10 // Unless required by applicable law or agreed to in writing, software
11 // distributed under the License is distributed on an "AS IS" BASIS,
12 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 // See the License for the specific language governing permissions and
14 // limitations under the License.
15 //*****************************************************************************
23 #include "ngraph/coordinate_transform.hpp"
24 #include "ngraph/op/topk.hpp"
26 NGRAPH_SUPPRESS_DEPRECATED_START
34 // Had to split out these two functions. They used to be lambda expressions but
35 // MSVC had difficulty compiling. This way is more explicit.
36 template <typename T, typename U>
37 inline bool compare_max(const std::tuple<T, U>& a, const std::tuple<T, U>& b)
39 // this is intentional to be able to compare floats directly
40 // without using relative or absolute tolerance
42 #pragma GCC diagnostic push
43 #pragma GCC diagnostic ignored "-Wfloat-equal"
45 if (std::get<0>(a) == std::get<0>(b))
47 return std::get<1>(a) < std::get<1>(b);
50 #pragma GCC diagnostic pop
55 template <typename T, typename U>
56 inline bool compare_min(const std::tuple<T, U>& a, const std::tuple<T, U>& b)
61 template <typename T, typename U>
62 inline bool sort_indices_descending(const std::tuple<T, U>& a,
63 const std::tuple<T, U>& b)
65 return std::get<1>(a) < std::get<1>(b);
68 template <typename T, typename U>
69 inline bool sort_indices_ascending(const std::tuple<T, U>& a, const std::tuple<T, U>& b)
71 return std::get<1>(a) > std::get<1>(b);
74 template <typename T, typename U>
75 void topk(const T* arg,
78 const Shape& in_shape,
79 const Shape& out_shape,
83 op::TopK::SortType sort = op::TopK::SortType::NONE)
86 // reorder source axis visit order and make "axis" inner most
87 size_t ndim = static_cast<size_t>(in_shape.size());
88 Coordinate start_corner(ndim, 0);
89 Coordinate end_corner(in_shape);
91 Strides strides(ndim, 1);
92 AxisVector axis_order(ndim);
93 iota(axis_order.begin(), axis_order.end(), 0);
94 axis_order.erase(axis_order.begin() + axis);
95 axis_order.push_back(axis);
96 // Create CoordinateTransforms that visits only the first element along "axis"
97 CoordinateTransform input_transform(
98 in_shape, start_corner, end_corner, strides, axis_order);
99 CoordinateTransform output_transform(
100 out_shape, start_corner, end_corner, strides, axis_order);
101 // Create temp vector for sorting.
102 vector<tuple<T, U>> workspace(in_shape[axis]);
103 vector<size_t> in_strides = ngraph::row_major_strides(in_shape);
104 vector<size_t> out_strides = ngraph::row_major_strides(out_shape);
105 auto in_axis_stride = in_strides[axis];
106 auto out_axis_stride = out_strides[axis];
107 for (const Coordinate& coord : input_transform)
109 auto arg_index = input_transform.index(coord);
110 auto out_index = output_transform.index(coord);
111 // Fill the temp vector
113 for (tuple<T, U>& entry : workspace)
115 get<0>(entry) = arg[arg_index];
117 arg_index += in_axis_stride;
120 // Sort the temp vector
123 nth_element(workspace.begin(),
124 workspace.begin() + k,
130 nth_element(workspace.begin(),
131 workspace.begin() + k,
135 // Write temp vector to output
140 case op::TopK::SortType::NONE: break;
141 case op::TopK::SortType::SORT_INDICES:
142 std::sort(workspace.begin(),
143 workspace.begin() + k,
144 sort_indices_descending<T, U>);
146 case op::TopK::SortType::SORT_VALUES:
147 std::sort(workspace.begin(), workspace.begin() + k, compare_max<T, U>);
155 case op::TopK::SortType::NONE: break;
156 case op::TopK::SortType::SORT_INDICES:
157 std::sort(workspace.begin(),
158 workspace.begin() + k,
159 sort_indices_ascending<T, U>);
161 case op::TopK::SortType::SORT_VALUES:
162 std::sort(workspace.begin(), workspace.begin() + k, compare_min<T, U>);
166 for (size_t j = 0; j < k; j++)
168 tuple<T, U> entry = workspace[j];
169 out_values[out_index] = get<0>(entry);
170 out_indices[out_index] = get<1>(entry);
171 out_index += out_axis_stride;
179 NGRAPH_SUPPRESS_DEPRECATED_END