1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19#include "pareto.h"
20
21#include <tvm/runtime/container/array.h>
22#include <tvm/runtime/container/map.h>
23#include <tvm/runtime/object.h>
24#include <tvm/runtime/registry.h>
25
26#include <array>
27#include <utility>
28#include <vector>
29
30#include "common.h"
31#include "plan.h"
32#include "proposal.h"
33#include "tensor_config.h"
34
35namespace tvm {
36namespace contrib {
37namespace ethosu {
38namespace cascader {
39
40template <int N>
41std::vector<bool> GetParetoFrontier(const std::vector<std::array<float, N>>& costs) {
42 std::vector<bool> is_optimal(costs.size(), true);
43 for (size_t i = 0; i < costs.size(); i++) {
44 if (is_optimal[i]) {
45 for (size_t j = 0; j < costs.size(); j++) {
46 if (is_optimal[j]) {
47 bool optimal = false;
48 for (size_t k = 0; k < N; k++) {
49 if (costs[i][k] > costs[j][k]) {
50 optimal = true;
51 break;
52 }
53 }
54 is_optimal[j] = optimal;
55 }
56 }
57 is_optimal[i] = true;
58 }
59 }
60 return is_optimal;
61}
62
63template <class T>
64std::vector<T> ThinVector(const std::vector<T>& vec, size_t max_size) {
65 if (max_size < 1) {
66 return std::vector<T>();
67 }
68 if (vec.size() <= max_size || vec.size() == 0) {
69 return vec;
70 }
71 if (max_size == 1) {
72 return std::vector<T>{vec[0]};
73 }
74 std::vector<T> thin_vec;
75 float step = static_cast<float>(vec.size()) / static_cast<float>(max_size - 1);
76 for (float i = 0; i < vec.size() - 1; i += step) {
77 thin_vec.push_back(vec[static_cast<int>(i)]);
78 }
79 thin_vec.push_back(vec.back());
80 return thin_vec;
81}
82
83std::vector<Plan> ParetoCullPlans(std::vector<Plan> plans, size_t max_plans,
84 bool disable_pareto_metric) {
85 if (plans.size() <= max_plans) {
86 return plans;
87 }
88 if (disable_pareto_metric) {
89 // Sample from all plans
90 return ThinVector(plans, max_plans);
91 }
92
93 std::sort(plans.begin(), plans.end(), [](const Plan& a, const Plan& b) -> bool {
94 if (a->GetMemoryUsage() == b->GetMemoryUsage()) {
95 return a->GetCycles() < b->GetCycles();
96 }
97 return a->GetMemoryUsage() < b->GetMemoryUsage();
98 });
99 std::vector<std::array<float, 2>> costs;
100 for (const auto& plan : plans) {
101 std::array<float, 2> cost = {static_cast<float>(plan->GetMemoryUsage()),
102 static_cast<float>(plan->GetCycles())};
103 costs.emplace_back(cost);
104 }
105 std::vector<bool> is_optimal = GetParetoFrontier<2>(costs);
106 std::vector<Plan> optimal_plans;
107 size_t i = 0;
108 for (bool optimal : is_optimal) {
109 if (optimal) {
110 optimal_plans.push_back(plans[i]);
111 }
112 i++;
113 }
114 if (optimal_plans.size() <= max_plans) {
115 return optimal_plans;
116 }
117 return ThinVector(optimal_plans, max_plans);
118}
119
120std::vector<Proposal> ParetoCullProposals(std::vector<Proposal> proposals, size_t max_proposals,
121 bool disable_pareto_metric) {
122 if (disable_pareto_metric) {
123 // Sample from all Proposals
124 return ThinVector(proposals, max_proposals);
125 }
126
127 std::sort(proposals.begin(), proposals.end(), [](const Proposal& a, const Proposal& b) -> bool {
128 if (a->GetMemoryUsage() == b->GetMemoryUsage()) {
129 return a->GetCycles() < b->GetCycles();
130 }
131 return a->GetMemoryUsage() < b->GetMemoryUsage();
132 });
133 std::vector<std::array<float, 2>> costs;
134 for (const auto& proposal : proposals) {
135 std::array<float, 2> cost = {static_cast<float>(proposal->GetMemoryUsage()),
136 static_cast<float>(proposal->GetCycles())};
137 costs.emplace_back(cost);
138 }
139 std::vector<bool> is_optimal = GetParetoFrontier<2>(costs);
140 std::vector<Proposal> optimal_proposals;
141 size_t i = 0;
142 for (bool optimal : is_optimal) {
143 if (optimal) {
144 optimal_proposals.push_back(proposals[i]);
145 }
146 i++;
147 }
148 if (optimal_proposals.size() <= max_proposals) {
149 return optimal_proposals;
150 }
151 return ThinVector(optimal_proposals, max_proposals);
152}
153
154TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.GetParetoFrontier")
155 .set_body_typed([](Array<Array<FloatImm>> tcosts) {
156 std::vector<std::array<float, 2>> costs;
157 for (const auto& tcost : tcosts) {
158 ICHECK_EQ(tcost.size(), 2);
159 std::array<float, 2> point = {static_cast<float>(tcost[0]->value),
160 static_cast<float>(tcost[1]->value)};
161 costs.push_back(point);
162 }
163 Array<Bool> is_optimal;
164 for (bool opt : GetParetoFrontier<2>(costs)) {
165 is_optimal.push_back(Bool(opt));
166 }
167 return is_optimal;
168 });
169
170TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.ThinVector")
171 .set_body_typed([](Array<ObjectRef> vec, int max_size) {
172 std::vector<ObjectRef> vvec(vec.begin(), vec.end());
173 return Array<ObjectRef>(ThinVector<ObjectRef>(vvec, max_size));
174 });
175
176TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.ParetoCullPlans")
177 .set_body_typed([](Array<Plan> plans, int max_size, bool disable_pareto_metric) {
178 std::vector<Plan> vplans(plans.begin(), plans.end());
179 return Array<Plan>(ParetoCullPlans(vplans, max_size, disable_pareto_metric));
180 });
181
182} // namespace cascader
183} // namespace ethosu
184} // namespace contrib
185} // namespace tvm
186