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
34namespace tvm {
35namespace contrib {
36namespace ethosu {
37namespace cascader {
38
39class CascaderGraph;
40class MemoryRegion;
41class Part;
42class Tensor;
43class StripeConfig;
44class Plan;
45class CascaderOptions;
46
47using 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 */
58std::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 */
73std::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 */
103std::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