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  ffi::Optional<ffi::Array<ScheduleRule>> sch_rules;
81  ffi::Optional<ffi::Array<Postproc>> postprocs;
83  ffi::Optional<ffi::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 ffi::Array<tir::Schedule> GenerateDesignSpace(const IRModule& mod) = 0;
109 
114  virtual SpaceGenerator Clone() const = 0;
115 
116  static constexpr const bool _type_mutable = true;
117  TVM_FFI_DECLARE_OBJECT_INFO("meta_schedule.SpaceGenerator", SpaceGeneratorNode, Object);
118 };
119 
124 class SpaceGenerator : public runtime::ObjectRef {
125  public:
130  explicit SpaceGenerator(ObjectPtr<SpaceGeneratorNode> data) : ObjectRef(data) {
131  TVM_FFI_ICHECK(data != nullptr);
132  }
137  using FInitializeWithTuneContext = ffi::TypedFunction<void(const TuneContext&)>;
143  using FGenerateDesignSpace = ffi::TypedFunction<ffi::Array<tir::Schedule>(const IRModule&)>;
148  using FClone = ffi::TypedFunction<SpaceGenerator()>;
149 
150  protected:
151  SpaceGenerator() = default;
152 
153  public:
165  ffi::Optional<ffi::Array<ScheduleRule>> sch_rules,
166  ffi::Optional<ffi::Array<Postproc>> postprocs,
167  ffi::Optional<ffi::Map<Mutator, FloatImm>> mutator_probs,
168  FInitializeWithTuneContext f_initialize_with_tune_context,
169  FGenerateDesignSpace f_generate_design_space, FClone f_clone);
180  TVM_DLL static SpaceGenerator ScheduleFn(
181  ffi::Function schedule_fn, ffi::Optional<ffi::Array<ScheduleRule>> sch_rules,
182  ffi::Optional<ffi::Array<Postproc>> postprocs,
183  ffi::Optional<ffi::Map<Mutator, FloatImm>> mutator_probs);
193  ffi::Array<SpaceGenerator, void> space_generators,
194  ffi::Optional<ffi::Array<ScheduleRule>> sch_rules,
195  ffi::Optional<ffi::Array<Postproc>> postprocs,
196  ffi::Optional<ffi::Map<Mutator, FloatImm>> mutator_probs);
207  ffi::Function f_block_filter, ffi::Optional<ffi::Array<ScheduleRule>> sch_rules,
208  ffi::Optional<ffi::Array<Postproc>> postprocs,
209  ffi::Optional<ffi::Map<Mutator, FloatImm>> mutator_probs);
211 };
212 
215  public:
225 
226  static void RegisterReflection() {
227  // `f_initialize_with_tune_context` is not registered
228  // `f_generate_design_space` is not registered
229  // `f_clone` is not registered
230  }
231 
232  void InitializeWithTuneContext(const TuneContext& context) final;
233  ffi::Array<tir::Schedule> GenerateDesignSpace(const IRModule& mod) final;
234  SpaceGenerator Clone() const final;
235  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("meta_schedule.PySpaceGenerator", PySpaceGeneratorNode,
237 };
238 
239 } // namespace meta_schedule
240 } // namespace tvm
241 
242 #endif // TVM_META_SCHEDULE_SPACE_GENERATOR_H_
Managed reference class to IRModuleNode.
Definition: module.h:256
The design space generator with customized methods on the python-side.
Definition: space_generator.h:214
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("meta_schedule.PySpaceGenerator", PySpaceGeneratorNode, SpaceGeneratorNode)
void InitializeWithTuneContext(const TuneContext &context) final
Initialize the design space generator with tuning context.
FClone f_clone
The packed function to the Clone function.
Definition: space_generator.h:224
SpaceGenerator::FClone FClone
Definition: space_generator.h:218
SpaceGenerator::FInitializeWithTuneContext FInitializeWithTuneContext
Definition: space_generator.h:216
FGenerateDesignSpace f_generate_design_space
The packed function to the GenerateDesignSpace function.
Definition: space_generator.h:222
static void RegisterReflection()
Definition: space_generator.h:226
SpaceGenerator Clone() const final
Clone the space generator.
SpaceGenerator::FGenerateDesignSpace FGenerateDesignSpace
Definition: space_generator.h:217
ffi::Array< tir::Schedule > GenerateDesignSpace(const IRModule &mod) final
Generate design spaces given a module.
FInitializeWithTuneContext f_initialize_with_tune_context
The packed function to the InitializeWithTuneContext function.
Definition: space_generator.h:220
The abstract class for design space generation.
Definition: space_generator.h:76
TVM_FFI_DECLARE_OBJECT_INFO("meta_schedule.SpaceGenerator", SpaceGeneratorNode, Object)
virtual ~SpaceGeneratorNode()=default
Default destructor.
static void RegisterReflection()
Definition: space_generator.h:85
ffi::Optional< ffi::Array< ScheduleRule > > sch_rules
The schedule rules.
Definition: space_generator.h:79
virtual void InitializeWithTuneContext(const TuneContext &context)
Initialize the design space generator with tuning context.
ffi::Optional< ffi::Array< Postproc > > postprocs
The postprocessors.
Definition: space_generator.h:81
static constexpr const bool _type_mutable
Definition: space_generator.h:116
ffi::Optional< ffi::Map< Mutator, FloatImm > > mutator_probs
The probability of using certain mutator.
Definition: space_generator.h:83
virtual ffi::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
SpaceGenerator(ObjectPtr< SpaceGeneratorNode > data)
Constructor from ObjectPtr<SpaceGeneratorNode>.
Definition: space_generator.h:130
ffi::TypedFunction< void(const TuneContext &)> FInitializeWithTuneContext
The function type of InitializeWithTuneContext method.
Definition: space_generator.h:137
static SpaceGenerator SpaceGeneratorUnion(ffi::Array< SpaceGenerator, void > space_generators, ffi::Optional< ffi::Array< ScheduleRule >> sch_rules, ffi::Optional< ffi::Array< Postproc >> postprocs, ffi::Optional< ffi::Map< Mutator, FloatImm >> mutator_probs)
Create a design space generator that is union of multiple design space generators.
ffi::TypedFunction< SpaceGenerator()> FClone
The function type of Clone method.
Definition: space_generator.h:148
static SpaceGenerator ScheduleFn(ffi::Function schedule_fn, ffi::Optional< ffi::Array< ScheduleRule >> sch_rules, ffi::Optional< ffi::Array< Postproc >> postprocs, ffi::Optional< ffi::Map< Mutator, FloatImm >> mutator_probs)
Create a design space generator with customized schedule function.
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SpaceGenerator, ObjectRef, SpaceGeneratorNode)
ffi::TypedFunction< ffi::Array< tir::Schedule >(const IRModule &)> FGenerateDesignSpace
The function type of GenerateDesignSpace method.
Definition: space_generator.h:143
static SpaceGenerator PostOrderApply(ffi::Function f_block_filter, ffi::Optional< ffi::Array< ScheduleRule >> sch_rules, ffi::Optional< ffi::Array< Postproc >> postprocs, ffi::Optional< ffi::Map< Mutator, FloatImm >> mutator_probs)
Create a design space generator that generates design spaces by applying schedule rules to blocks in ...
static SpaceGenerator PySpaceGenerator(ffi::Optional< ffi::Array< ScheduleRule >> sch_rules, ffi::Optional< ffi::Array< Postproc >> postprocs, ffi::Optional< ffi::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.
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:308
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
A managed object in the TVM runtime.
Compilation target object.