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#ifndef TVM_META_SCHEDULE_SPACE_GENERATOR_H_
20#define TVM_META_SCHEDULE_SPACE_GENERATOR_H_
21
22#include <tvm/ir/module.h>
23#include <tvm/meta_schedule/mutator.h>
24#include <tvm/meta_schedule/postproc.h>
25#include <tvm/meta_schedule/schedule_rule.h>
26#include <tvm/node/reflection.h>
27#include <tvm/runtime/container/array.h>
28#include <tvm/runtime/object.h>
29#include <tvm/runtime/packed_func.h>
30#include <tvm/target/target.h>
31#include <tvm/tir/schedule/schedule.h>
32
33namespace tvm {
34namespace meta_schedule {
35
36// Forward declaration
37class TuneContext;
38class SpaceGenerator;
39
40/*!
41 * \brief The abstract class for design space generation.
42 * \note The relationship between SpaceGenerator and other classes are as follows:
43 ┌──────────────────────────────────────────────────────────────┐
44 ┌──┴───────────────────────────────────────────────────────────┐ │
45┌──┴────────────────── Tune Context ───────────────────────────┐ │ │
46│ ┌─────────────────────┐ │ │ │
47│ │ │ Generate │ │ │
48│ │ Space Generator ├──────────────┐ │ │ │
49│ │ │ │ │ │ │
50│ └─────────────────────┘ ▼ │ │ │
51│ Design Space │ │ │
52│ ┌─────────────────────┐ │ │ │ │
53│ Generate │ │ Pretuning │ │ │ │
54│ ┌───────────┤ Search Strategy │◄─────────────┘ │ │ │
55│ │ │ │ │ ├──┘
56│ │ └─────────────────────┘ ├──┘
57└────┼─────────────────────────────────────────────────────────┘
58 │
59 │
60┌────┼──────────────── Managed By Task Scheduler ─────────────────────┐
61│ │ ┌───────────┐ │
62│ │ Send to │ │ Send to │
63│ ▼ ┌─────────────►│ Builder ├──────────┐ │
64│ Measure Candidate │ Builder │ │ Runner │ │
65│ │ │ └───────────┘ │ │
66│ │ ┌────────────┴────────┐ │ │
67│ │ │ │ ┌───────────┐ │ │
68│ └────►│ Task Scheduler │ │ │ │ │
69│ │ │ │ Runner │◄─────────┘ │
70│ └─────────────────────┘ │ │ │
71│ ▲ └─────┬─────┘ │
72│ │ │ │
73│ └─── Runner Future ◄────┘ │
74└─────────────────────────────────────────────────────────────────────┘
75*/
76class SpaceGeneratorNode : public runtime::Object {
77 public:
78 /*! \brief The schedule rules. */
79 Optional<Array<ScheduleRule>> sch_rules;
80 /*! \brief The postprocessors. */
81 Optional<Array<Postproc>> postprocs;
82 /*! \brief The probability of using certain mutator. */
83 Optional<Map<Mutator, FloatImm>> mutator_probs;
84
85 void VisitAttrs(tvm::AttrVisitor* v) {
86 v->Visit("sch_rules", &sch_rules);
87 v->Visit("postprocs", &postprocs);
88 v->Visit("mutator_probs", &mutator_probs);
89 }
90
91 /*! \brief Default destructor */
92 virtual ~SpaceGeneratorNode() = default;
93
94 /*!
95 * \brief Initialize the design space generator with tuning context.
96 * \param context The tuning context for initialization.
97 * \note This method is supposed to be called only once before every other method.
98 */
99 virtual void InitializeWithTuneContext(const TuneContext& context);
100
101 /*!
102 * \brief Generate design spaces given a module.
103 * \param mod The module used for design space generation.
104 * \return The generated design spaces, i.e., schedules.
105 */
106 virtual Array<tir::Schedule> GenerateDesignSpace(const IRModule& mod) = 0;
107
108 /*!
109 * \brief Clone the space generator.
110 * \return The cloned space generator.
111 */
112 virtual SpaceGenerator Clone() const = 0;
113
114 static constexpr const char* _type_key = "meta_schedule.SpaceGenerator";
115 TVM_DECLARE_BASE_OBJECT_INFO(SpaceGeneratorNode, Object);
116};
117
118/*!
119 * \brief Managed reference to SpaceGeneratorNode.
120 * \sa SpaceGeneratorNode
121 */
122class SpaceGenerator : public runtime::ObjectRef {
123 public:
124 /*!
125 * \brief The function type of `InitializeWithTuneContext` method.
126 * \param context The tuning context for initialization.
127 */
128 using FInitializeWithTuneContext = runtime::TypedPackedFunc<void(const TuneContext&)>;
129 /*!
130 * \brief The function type of `GenerateDesignSpace` method.
131 * \param mod The module used for design space generation.
132 * \return The generated design spaces, i.e., schedules.
133 */
134 using FGenerateDesignSpace = runtime::TypedPackedFunc<Array<tir::Schedule>(const IRModule&)>;
135 /*!
136 * \brief The function type of `Clone` method.
137 * \return The cloned space generator.
138 */
139 using FClone = runtime::TypedPackedFunc<SpaceGenerator()>;
140
141 protected:
142 SpaceGenerator() = default;
143
144 public:
145 /*!
146 * \brief Create a design space generator with customized methods on the python-side.
147 * \param sch_rules The schedule rules.
148 * \param postprocs The postprocessors.
149 * \param mutator_probs The probability of using certain mutator.
150 * \param f_initialize_with_tune_context The packed function of `InitializeWithTuneContext`.
151 * \param f_generate_design_space The packed function of `GenerateDesignSpace`.
152 * \param f_clone The packed function of `Clone`.
153 * \return The design space generator created.
154 */
155 TVM_DLL static SpaceGenerator PySpaceGenerator(
156 Optional<Array<ScheduleRule>> sch_rules, Optional<Array<Postproc>> postprocs,
157 Optional<Map<Mutator, FloatImm>> mutator_probs,
158 FInitializeWithTuneContext f_initialize_with_tune_context,
159 FGenerateDesignSpace f_generate_design_space, FClone f_clone);
160 /*!
161 * \brief Create a design space generator with customized schedule function.
162 * \param schedule_fn The schedule function, which can have the following signatures:
163 * 1) void(Schedule)
164 * 2) Schedule(Schedule)
165 * 3) Array<Schedule>(Schedule)
166 * \param sch_rules The schedule rules.
167 * \param postprocs The postprocessors.
168 * \param mutator_probs The probability of using certain mutator.
169 */
170 TVM_DLL static SpaceGenerator ScheduleFn(PackedFunc schedule_fn,
171 Optional<Array<ScheduleRule>> sch_rules,
172 Optional<Array<Postproc>> postprocs,
173 Optional<Map<Mutator, FloatImm>> mutator_probs);
174 /*!
175 * \brief Create a design space generator that is union of multiple design space generators.
176 * \param space_generators An array of design space generators to be unioned.
177 * \param sch_rules The schedule rules.
178 * \param postprocs The postprocessors.
179 * \param mutator_probs The probability of using certain mutator.
180 * \return The design space generator created.
181 */
182 TVM_DLL static SpaceGenerator SpaceGeneratorUnion(Array<SpaceGenerator, void> space_generators,
183 Optional<Array<ScheduleRule>> sch_rules,
184 Optional<Array<Postproc>> postprocs,
185 Optional<Map<Mutator, FloatImm>> mutator_probs);
186 /*!
187 * \brief Create a design space generator that generates design spaces by applying schedule
188 * rules to blocks in post-DFS order.
189 * \param f_block_filter The filter function to filter blocks to be applied with schedule rules.
190 * \param sch_rules The schedule rules.
191 * \param postprocs The postprocessors.
192 * \param mutator_probs The probability of using certain mutator.
193 * \return The design space generator created.
194 */
195 TVM_DLL static SpaceGenerator PostOrderApply(runtime::PackedFunc f_block_filter,
196 Optional<Array<ScheduleRule>> sch_rules,
197 Optional<Array<Postproc>> postprocs,
198 Optional<Map<Mutator, FloatImm>> mutator_probs);
199 TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(SpaceGenerator, ObjectRef, SpaceGeneratorNode);
200};
201
202/*! \brief The design space generator with customized methods on the python-side. */
203class PySpaceGeneratorNode : public SpaceGeneratorNode {
204 public:
205 using FInitializeWithTuneContext = SpaceGenerator::FInitializeWithTuneContext;
206 using FGenerateDesignSpace = SpaceGenerator::FGenerateDesignSpace;
207 using FClone = SpaceGenerator::FClone;
208 /*! \brief The packed function to the `InitializeWithTuneContext` function. */
209 FInitializeWithTuneContext f_initialize_with_tune_context;
210 /*! \brief The packed function to the `GenerateDesignSpace` function. */
211 FGenerateDesignSpace f_generate_design_space;
212 /*! \brief The packed function to the `Clone` function. */
213 FClone f_clone;
214
215 void VisitAttrs(tvm::AttrVisitor* v) {
216 SpaceGeneratorNode::VisitAttrs(v);
217 // `f_initialize_with_tune_context` is not visited
218 // `f_generate_design_space` is not visited
219 // `f_clone` is not visited
220 }
221
222 void InitializeWithTuneContext(const TuneContext& context) final;
223 Array<tir::Schedule> GenerateDesignSpace(const IRModule& mod) final;
224 SpaceGenerator Clone() const final;
225
226 static constexpr const char* _type_key = "meta_schedule.PySpaceGenerator";
227 TVM_DECLARE_FINAL_OBJECT_INFO(PySpaceGeneratorNode, SpaceGeneratorNode);
228};
229
230} // namespace meta_schedule
231} // namespace tvm
232
233#endif // TVM_META_SCHEDULE_SPACE_GENERATOR_H_
234