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 bool _type_mutable = true;
67  TVM_FFI_DECLARE_OBJECT_INFO("meta_schedule.Postproc", PostprocNode, Object);
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<ffi::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 ffi::Array<Postproc, void> DefaultLLVM();
168  TVM_DLL static ffi::Array<Postproc, void> DefaultCPUTensorization();
170  TVM_DLL static ffi::Array<Postproc, void> DefaultRISCV();
172  TVM_DLL static ffi::Array<Postproc, void> DefaultCUDA();
174  TVM_DLL static ffi::Array<Postproc, void> DefaultCUDATensorCore();
176  TVM_DLL static ffi::Array<Postproc, void> DefaultHexagon();
177 
179 };
180 
182 class PyPostprocNode : public PostprocNode {
183  public:
196 
197  static void RegisterReflection() {
198  // `f_initialize_with_tune_context` is not registered
199  // `f_apply` is not registered
200  // `f_clone` is not registered
201  // `f_as_string` is not registered
202  }
203 
204  void InitializeWithTuneContext(const TuneContext& context) final;
205  bool Apply(const tir::Schedule& sch) final;
206  Postproc Clone() const final;
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.
static constexpr const bool _type_mutable
Definition: postproc.h:66
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.
static void RegisterReflection()
Definition: postproc.h:42
TVM_FFI_DECLARE_OBJECT_INFO("meta_schedule.Postproc", PostprocNode, Object)
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.
static Postproc RewriteUnboundBlock(int max_threadblocks)
Create a postprocessor that adds thread binding to unbound blocks.
static ffi::Array< Postproc, void > DefaultHexagon()
Create default postprocessors for Hexagon.
ffi::TypedFunction< ffi::String()> FAsString
Get the postprocessor function as string with name.
Definition: postproc.h:96
static ffi::Array< Postproc, void > DefaultCUDATensorCore()
Create default postprocessors for CUDA with TensorCore.
ffi::TypedFunction< void(const TuneContext &)> FInitializeWithTuneContext
The function type of InitializeWithTuneContext method.
Definition: postproc.h:80
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 ffi::Array< Postproc, void > DefaultRISCV()
Create default postprocessors for RISCV.
static ffi::Array< Postproc, void > DefaultLLVM()
Create default postprocessors for LLVM.
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.
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Postproc, ObjectRef, PostprocNode)
static ffi::Array< Postproc, void > DefaultCPUTensorization()
Create default postprocessors for x86 (AVX512 and VNNI)
static ffi::Array< Postproc, void > DefaultCUDA()
Create default postprocessors for CUDA.
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 Postproc DisallowDynamicLoop()
Create a postprocessor that checks if all loops are static.
The postprocessor with customized methods on the python-side.
Definition: postproc.h:182
Postproc Clone() const final
Clone the postprocessor.
FAsString f_as_string
The packed function to the AsString function.
Definition: postproc.h:195
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:191
bool Apply(const tir::Schedule &sch) final
Apply a postprocessor to the given schedule.
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("meta_schedule.PyPostproc", PyPostprocNode, PostprocNode)
static void RegisterReflection()
Definition: postproc.h:197
FInitializeWithTuneContext f_initialize_with_tune_context
The packed function to the InitializeWithTuneContext function.
Definition: postproc.h:189
Postproc::FInitializeWithTuneContext FInitializeWithTuneContext
Definition: postproc.h:184
FClone f_clone
The packed function to the Clone function.
Definition: postproc.h:193
Postproc::FClone FClone
Definition: postproc.h:186
Postproc::FAsString FAsString
Definition: postproc.h:187
Postproc::FApply FApply
Definition: postproc.h:185
Managed reference to TuneContextNode.
Definition: tune_context.h:98
Managed reference to ScheduleNode.
Definition: schedule.h:885
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
A managed object in the TVM runtime.