tvm
space_generator.h
Go to the documentation of this file.
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/ffi/container/array.h>
23 #include <tvm/ffi/function.h>
24 #include <tvm/ffi/reflection/registry.h>
25 #include <tvm/ir/module.h>
29 #include <tvm/runtime/object.h>
30 #include <tvm/target/target.h>
32 
33 namespace tvm {
34 namespace meta_schedule {
35 
36 // Forward declaration
37 class TuneContext;
38 class SpaceGenerator;
39 
76 class SpaceGeneratorNode : public runtime::Object {
77  public:
79  Optional<Array<ScheduleRule>> sch_rules;
81  Optional<Array<Postproc>> postprocs;
83  Optional<Map<Mutator, FloatImm>> mutator_probs;
84 
85  static void RegisterReflection() {
86  namespace refl = tvm::ffi::reflection;
87  refl::ObjectDef<SpaceGeneratorNode>()
88  .def_ro("sch_rules", &SpaceGeneratorNode::sch_rules)
89  .def_ro("postprocs", &SpaceGeneratorNode::postprocs)
90  .def_ro("mutator_probs", &SpaceGeneratorNode::mutator_probs);
91  }
92 
94  virtual ~SpaceGeneratorNode() = default;
95 
101  virtual void InitializeWithTuneContext(const TuneContext& context);
102 
108  virtual Array<tir::Schedule> GenerateDesignSpace(const IRModule& mod) = 0;
109 
114  virtual SpaceGenerator Clone() const = 0;
115 
116  static constexpr const char* _type_key = "meta_schedule.SpaceGenerator";
118 };
119 
124 class SpaceGenerator : public runtime::ObjectRef {
125  public:
130  using FInitializeWithTuneContext = ffi::TypedFunction<void(const TuneContext&)>;
136  using FGenerateDesignSpace = ffi::TypedFunction<Array<tir::Schedule>(const IRModule&)>;
141  using FClone = ffi::TypedFunction<SpaceGenerator()>;
142 
143  protected:
144  SpaceGenerator() = default;
145 
146  public:
158  Optional<Array<ScheduleRule>> sch_rules, Optional<Array<Postproc>> postprocs,
159  Optional<Map<Mutator, FloatImm>> mutator_probs,
160  FInitializeWithTuneContext f_initialize_with_tune_context,
161  FGenerateDesignSpace f_generate_design_space, FClone f_clone);
172  TVM_DLL static SpaceGenerator ScheduleFn(ffi::Function schedule_fn,
173  Optional<Array<ScheduleRule>> sch_rules,
174  Optional<Array<Postproc>> postprocs,
175  Optional<Map<Mutator, FloatImm>> mutator_probs);
184  TVM_DLL static SpaceGenerator SpaceGeneratorUnion(Array<SpaceGenerator, void> space_generators,
185  Optional<Array<ScheduleRule>> sch_rules,
186  Optional<Array<Postproc>> postprocs,
187  Optional<Map<Mutator, FloatImm>> mutator_probs);
197  TVM_DLL static SpaceGenerator PostOrderApply(ffi::Function f_block_filter,
198  Optional<Array<ScheduleRule>> sch_rules,
199  Optional<Array<Postproc>> postprocs,
200  Optional<Map<Mutator, FloatImm>> mutator_probs);
202 };
203 
206  public:
216 
217  static void RegisterReflection() {
218  // `f_initialize_with_tune_context` is not registered
219  // `f_generate_design_space` is not registered
220  // `f_clone` is not registered
221  }
222 
223  void InitializeWithTuneContext(const TuneContext& context) final;
224  Array<tir::Schedule> GenerateDesignSpace(const IRModule& mod) final;
225  SpaceGenerator Clone() const final;
226 
227  static constexpr const char* _type_key = "meta_schedule.PySpaceGenerator";
229 };
230 
231 } // namespace meta_schedule
232 } // namespace tvm
233 
234 #endif // TVM_META_SCHEDULE_SPACE_GENERATOR_H_
Managed reference class to IRModuleNode.
Definition: module.h:257
The design space generator with customized methods on the python-side.
Definition: space_generator.h:205
void InitializeWithTuneContext(const TuneContext &context) final
Initialize the design space generator with tuning context.
static constexpr const char * _type_key
Definition: space_generator.h:227
FClone f_clone
The packed function to the Clone function.
Definition: space_generator.h:215
Array< tir::Schedule > GenerateDesignSpace(const IRModule &mod) final
Generate design spaces given a module.
SpaceGenerator::FClone FClone
Definition: space_generator.h:209
SpaceGenerator::FInitializeWithTuneContext FInitializeWithTuneContext
Definition: space_generator.h:207
FGenerateDesignSpace f_generate_design_space
The packed function to the GenerateDesignSpace function.
Definition: space_generator.h:213
static void RegisterReflection()
Definition: space_generator.h:217
TVM_DECLARE_FINAL_OBJECT_INFO(PySpaceGeneratorNode, SpaceGeneratorNode)
SpaceGenerator Clone() const final
Clone the space generator.
SpaceGenerator::FGenerateDesignSpace FGenerateDesignSpace
Definition: space_generator.h:208
FInitializeWithTuneContext f_initialize_with_tune_context
The packed function to the InitializeWithTuneContext function.
Definition: space_generator.h:211
The abstract class for design space generation.
Definition: space_generator.h:76
Optional< Array< Postproc > > postprocs
The postprocessors.
Definition: space_generator.h:81
static constexpr const char * _type_key
Definition: space_generator.h:116
Optional< Array< ScheduleRule > > sch_rules
The schedule rules.
Definition: space_generator.h:79
virtual ~SpaceGeneratorNode()=default
Default destructor.
static void RegisterReflection()
Definition: space_generator.h:85
TVM_DECLARE_BASE_OBJECT_INFO(SpaceGeneratorNode, Object)
virtual void InitializeWithTuneContext(const TuneContext &context)
Initialize the design space generator with tuning context.
Optional< Map< Mutator, FloatImm > > mutator_probs
The probability of using certain mutator.
Definition: space_generator.h:83
virtual Array< tir::Schedule > GenerateDesignSpace(const IRModule &mod)=0
Generate design spaces given a module.
virtual SpaceGenerator Clone() const =0
Clone the space generator.
Managed reference to SpaceGeneratorNode.
Definition: space_generator.h:124
ffi::TypedFunction< void(const TuneContext &)> FInitializeWithTuneContext
The function type of InitializeWithTuneContext method.
Definition: space_generator.h:130
static SpaceGenerator SpaceGeneratorUnion(Array< SpaceGenerator, void > space_generators, Optional< Array< ScheduleRule >> sch_rules, Optional< Array< Postproc >> postprocs, Optional< Map< Mutator, FloatImm >> mutator_probs)
Create a design space generator that is union of multiple design space generators.
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(SpaceGenerator, ObjectRef, SpaceGeneratorNode)
static SpaceGenerator PySpaceGenerator(Optional< Array< ScheduleRule >> sch_rules, Optional< Array< Postproc >> postprocs, Optional< Map< Mutator, FloatImm >> mutator_probs, FInitializeWithTuneContext f_initialize_with_tune_context, FGenerateDesignSpace f_generate_design_space, FClone f_clone)
Create a design space generator with customized methods on the python-side.
ffi::TypedFunction< SpaceGenerator()> FClone
The function type of Clone method.
Definition: space_generator.h:141
ffi::TypedFunction< Array< tir::Schedule >(const IRModule &)> FGenerateDesignSpace
The function type of GenerateDesignSpace method.
Definition: space_generator.h:136
static SpaceGenerator ScheduleFn(ffi::Function schedule_fn, Optional< Array< ScheduleRule >> sch_rules, Optional< Array< Postproc >> postprocs, Optional< Map< Mutator, FloatImm >> mutator_probs)
Create a design space generator with customized schedule function.
static SpaceGenerator PostOrderApply(ffi::Function f_block_filter, Optional< Array< ScheduleRule >> sch_rules, Optional< Array< Postproc >> postprocs, Optional< Map< Mutator, FloatImm >> mutator_probs)
Create a design space generator that generates design spaces by applying schedule rules to blocks in ...
Managed reference to TuneContextNode.
Definition: tune_context.h:98
IRModule that holds the functions and type definitions.
Definition: repr_printer.h:91
tvm::relax::Function Function
Definition: transform.h:42
tvm::PrimExpr mod(const tvm::PrimExpr &a, const tvm::PrimExpr &b)
Definition: broadcast.h:306
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
A managed object in the TVM runtime.
Compilation target object.