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 | |
20 | /*! |
21 | * \file src/contrib/ethosu/cascader/plan_generator.h |
22 | * \brief Algorithm to generate possible Plans in the NPU cascader |
23 | */ |
24 | #ifndef TVM_CONTRIB_ETHOSU_CASCADER_PLAN_GENERATOR_H_ |
25 | #define TVM_CONTRIB_ETHOSU_CASCADER_PLAN_GENERATOR_H_ |
26 | |
27 | #include <tvm/node/reflection.h> |
28 | #include <tvm/runtime/object.h> |
29 | |
30 | #include <unordered_map> |
31 | #include <unordered_set> |
32 | #include <vector> |
33 | |
34 | namespace tvm { |
35 | namespace contrib { |
36 | namespace ethosu { |
37 | namespace cascader { |
38 | |
39 | class CascaderGraph; |
40 | class MemoryRegion; |
41 | class Part; |
42 | class Tensor; |
43 | class StripeConfig; |
44 | class Plan; |
45 | class CascaderOptions; |
46 | |
47 | using HomeMap = |
48 | std::unordered_map<Tensor, std::vector<MemoryRegion>, ObjectPtrHash, ObjectPtrEqual>; |
49 | |
50 | /*! |
51 | * \brief Generate possible output StripeConfigs that could be applied to a Part's output. |
52 | * \param part The Part to generate StripeConfigs for. |
53 | * \param stripe_factors How many striping factors to try per axis. |
54 | * \param enable_striping Whether striping is enabled |
55 | * \param multi_dimensional Whether to stripe in more than one dimension. |
56 | * \return The generated StripeConfigs for the Part's output. |
57 | */ |
58 | std::vector<StripeConfig> GenerateOutputStripeConfigs(const Part& part, int stripe_factors, |
59 | bool enable_striping, bool multi_dimensional); |
60 | |
61 | /*! |
62 | * \brief Generate single-Part Plans for a Part for a given list of output StripeConfigs. |
63 | * \param part The Part to generate Plans for. |
64 | * \param output_stripe_configs The output StripeConfigs to generate Plans with. |
65 | * \param home_map The Tensor homing map defining valid memory homes for Tensors. |
66 | * \param options The configuration options with which to run the generator. |
67 | * \return The generated Plans covering the Part. |
68 | * \note For each of the output StripeConfigs provided, this algorithm will produce a number |
69 | * of Plans corresponding to different choices of Tensor homing/copying, buffer modes |
70 | * and INTERIOR/BOUNDARY states. For each of these variants, the Part's performance will |
71 | * be queried and the memory usage will be calculated. |
72 | */ |
73 | std::vector<Plan> GenerateSinglePlans(const Part& part, |
74 | const std::vector<StripeConfig>& output_stripe_configs, |
75 | const HomeMap& home_map, const CascaderOptions& options); |
76 | |
77 | /*! |
78 | * \brief Generate pareto optimal Plans for a Graph. |
79 | * \param graph The Graph to generate Plans for. |
80 | * \param home_map The Tensor homing map defining valid memory homes for Tensors. |
81 | * \param options The configuration options with which to run the generator. |
82 | * \return A map between Part groups and a list of pareto optimal Plans which cover that group. |
83 | * \note This algorithm does the following: |
84 | * |
85 | * Iterate Part-by-Part in a reversed topological ordering (starting at the output Parts and |
86 | * working towards the input Parts). |
87 | * |
88 | * For each Part: |
89 | * 1. Determine the possible StripeConfigs we might want to use to stripe the Part using |
90 | * GenerateOutputStripeConfigs. |
91 | * 2. Additionally, collect all the StripeConfigs of open Plans that could connect to this |
92 | * Part (i.e. the Plan has an open TensorConfig for the Part's output Tensor). |
93 | * 3. Use these two lists of StripeConfigs to produce single Part Plans with GenerateSinglePlans. |
94 | * 4. For the generated Plans that have an open output TensorConfig, try and merge these into |
95 | * existing Plans which share an open input TensorConfig. |
96 | * 5. All Plans are then indexed by both the Part group they cover and their open TensorConfigs. |
97 | * 6. Plans which cover the same Part group and share the same open TensorConfigs are culled |
98 | * using ParetoCullPlans. |
99 | * |
100 | * Once every Part has been visited, return the Plans with no open TensorConfigs indexed by Part |
101 | * group. |
102 | */ |
103 | std::unordered_map<std::vector<Part>, std::vector<Plan>> GenerateGraphPlans( |
104 | const CascaderGraph& graph, const HomeMap& home_map, const CascaderOptions& options); |
105 | |
106 | } // namespace cascader |
107 | } // namespace ethosu |
108 | } // namespace contrib |
109 | } // namespace tvm |
110 | |
111 | #endif // TVM_CONTRIB_ETHOSU_CASCADER_PLAN_GENERATOR_H_ |
112 | |