tvm
postproc.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 
20 #ifndef TVM_META_SCHEDULE_POSTPROC_H_
21 #define TVM_META_SCHEDULE_POSTPROC_H_
22 
23 #include <tvm/ffi/function.h>
24 #include <tvm/ffi/reflection/registry.h>
25 #include <tvm/runtime/object.h>
27 
28 namespace tvm {
29 namespace meta_schedule {
30 
31 class TuneContext;
32 class Postproc;
33 
37 class PostprocNode : public runtime::Object {
38  public:
40  virtual ~PostprocNode() = default;
41 
42  static void RegisterReflection() {
43  // No fields to register
44  }
45 
51  virtual void InitializeWithTuneContext(const TuneContext& context) = 0;
52 
58  virtual bool Apply(const tir::Schedule& sch) = 0;
59 
64  virtual Postproc Clone() const = 0;
65 
66  static constexpr const char* _type_key = "meta_schedule.Postproc";
68 };
69 
74 class Postproc : public runtime::ObjectRef {
75  public:
80  using FInitializeWithTuneContext = ffi::TypedFunction<void(const TuneContext&)>;
86  using FApply = ffi::TypedFunction<bool(const tir::Schedule&)>;
91  using FClone = ffi::TypedFunction<Postproc()>;
96  using FAsString = ffi::TypedFunction<String()>;
105  TVM_DLL static Postproc PyPostproc(FInitializeWithTuneContext f_initialize_with_tune_context, //
106  FApply f_apply, //
107  FClone f_clone, //
108  FAsString f_as_string);
113  TVM_DLL static Postproc DisallowDynamicLoop();
135  TVM_DLL static Postproc RewriteReductionBlock();
141  TVM_DLL static Postproc RewriteUnboundBlock(int max_threadblocks);
148  TVM_DLL static Postproc RewriteTensorize(bool vectorize_init_loop = false);
153  TVM_DLL static Postproc VerifyGPUCode();
158  TVM_DLL static Postproc VerifyVTCMLimit();
164  TVM_DLL static Postproc RewriteLayout();
166  TVM_DLL static Array<Postproc, void> DefaultLLVM();
168  TVM_DLL static Array<Postproc, void> DefaultCPUTensorization();
170  TVM_DLL static Array<Postproc, void> DefaultCUDA();
172  TVM_DLL static Array<Postproc, void> DefaultCUDATensorCore();
174  TVM_DLL static Array<Postproc, void> DefaultHexagon();
175 
177 };
178 
180 class PyPostprocNode : public PostprocNode {
181  public:
194 
195  static void RegisterReflection() {
196  // `f_initialize_with_tune_context` is not registered
197  // `f_apply` is not registered
198  // `f_clone` is not registered
199  // `f_as_string` is not registered
200  }
201 
202  void InitializeWithTuneContext(const TuneContext& context) final;
203  bool Apply(const tir::Schedule& sch) final;
204  Postproc Clone() const final;
205 
206  static constexpr const char* _type_key = "meta_schedule.PyPostproc";
208 };
209 
210 } // namespace meta_schedule
211 } // namespace tvm
212 
213 #endif // TVM_META_SCHEDULE_POSTPROC_H_
Rules to apply a postprocessor to a schedule.
Definition: postproc.h:37
virtual void InitializeWithTuneContext(const TuneContext &context)=0
Initialize the design space generator with tuning context.
virtual Postproc Clone() const =0
Clone the postprocessor.
virtual ~PostprocNode()=default
Virtual destructor.
virtual bool Apply(const tir::Schedule &sch)=0
Apply a postprocessor to the given schedule.
TVM_DECLARE_BASE_OBJECT_INFO(PostprocNode, Object)
static void RegisterReflection()
Definition: postproc.h:42
static constexpr const char * _type_key
Definition: postproc.h:66
Managed reference to PostprocNode.
Definition: postproc.h:74
static Postproc DisallowAsyncStridedMemCopy()
Create a postprocessor that checks if all async mem copies are not strided.
static Postproc RewriteReductionBlock()
Create a postprocessor that rewrites reduction block by moving the init block out.
ffi::TypedFunction< String()> FAsString
Get the postprocessor function as string with name.
Definition: postproc.h:96
static Postproc RewriteUnboundBlock(int max_threadblocks)
Create a postprocessor that adds thread binding to unbound blocks.
ffi::TypedFunction< void(const TuneContext &)> FInitializeWithTuneContext
The function type of InitializeWithTuneContext method.
Definition: postproc.h:80
TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(Postproc, ObjectRef, PostprocNode)
static Array< Postproc, void > DefaultCUDATensorCore()
Create default postprocessors for CUDA with TensorCore.
static Array< Postproc, void > DefaultCPUTensorization()
Create default postprocessors for x86 (AVX512 and VNNI)
static Array< Postproc, void > DefaultLLVM()
Create default postprocessors for LLVM.
ffi::TypedFunction< bool(const tir::Schedule &)> FApply
Apply a postprocessor to the given schedule.
Definition: postproc.h:86
static Postproc VerifyGPUCode()
Creates a postprocessor that verifies if the GPU code is correct.
static Array< Postproc, void > DefaultCUDA()
Create default postprocessors for CUDA.
static Postproc RewriteLayout()
Creates a postprocessor that rewrites the layout of input tensor.
static Postproc RewriteCooperativeFetch()
Create a postprocessor that rewrites the cooperative fetch annotation to actual vectorized cooperativ...
static Postproc RewriteTensorize(bool vectorize_init_loop=false)
Create a postprocessor that applies tensorization to annotated blocks.
static Postproc VerifyVTCMLimit()
Verifies that the VTCM usage of a given schedule is within the provided limit.
static Postproc RewriteParallelVectorizeUnroll()
Creates a postprocessor that applies parallelization, vectorization and auto unrolling according to t...
ffi::TypedFunction< Postproc()> FClone
Clone the postprocessor.
Definition: postproc.h:91
static Postproc PyPostproc(FInitializeWithTuneContext f_initialize_with_tune_context, FApply f_apply, FClone f_clone, FAsString f_as_string)
Create a postprocessor with customized methods on the python-side.
static Array< Postproc, void > DefaultHexagon()
Create default postprocessors for Hexagon.
static Postproc DisallowDynamicLoop()
Create a postprocessor that checks if all loops are static.
The postprocessor with customized methods on the python-side.
Definition: postproc.h:180
Postproc Clone() const final
Clone the postprocessor.
FAsString f_as_string
The packed function to the AsString function.
Definition: postproc.h:193
void InitializeWithTuneContext(const TuneContext &context) final
Initialize the design space generator with tuning context.
FApply f_apply
The packed function to the Apply function.
Definition: postproc.h:189
bool Apply(const tir::Schedule &sch) final
Apply a postprocessor to the given schedule.
static constexpr const char * _type_key
Definition: postproc.h:206
static void RegisterReflection()
Definition: postproc.h:195
FInitializeWithTuneContext f_initialize_with_tune_context
The packed function to the InitializeWithTuneContext function.
Definition: postproc.h:187
TVM_DECLARE_FINAL_OBJECT_INFO(PyPostprocNode, PostprocNode)
Postproc::FInitializeWithTuneContext FInitializeWithTuneContext
Definition: postproc.h:182
FClone f_clone
The packed function to the Clone function.
Definition: postproc.h:191
Postproc::FClone FClone
Definition: postproc.h:184
Postproc::FAsString FAsString
Definition: postproc.h:185
Postproc::FApply FApply
Definition: postproc.h:183
Managed reference to TuneContextNode.
Definition: tune_context.h:98
Managed reference to ScheduleNode.
Definition: schedule.h:880
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
A managed object in the TVM runtime.