Skip to content

Commit 87ec806

Browse files
tqchentqchen
committed
[REFACTOR] Cleanup legacy TE-based passes
This PR cleans up legacy TE-based passes. - StorageFlatten/TextureFlatten/InjectPrefetch were used in TE-lowering. - CoProcSync/LiftAttrScope were used in VTA.
1 parent cc03780 commit 87ec806

26 files changed

+21
-3976
lines changed

‎apps/README.md‎

Lines changed: 0 additions & 28 deletions
This file was deleted.

‎include/tvm/tir/function.h‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -89,8 +89,7 @@ class PrimFuncNode : public BaseFuncNode {
8989
* normal statements, making buffer_map as first class citizen of PrimFunc
9090
* will make program analysis much easier.
9191
*
92-
* Prior to buffer flattening, which is performed either in
93-
* StorageFlatten for TE-based schedules or in FlattenBuffer for
92+
* Prior to buffer flattening, which is performed FlattenBuffer for
9493
* TIR-based schedules, these buffer objects are used directly in
9594
* the body of the function. After buffer flattening, these buffer
9695
* objects remain unflattened for use in argument validation, but

‎include/tvm/tir/transform.h‎

Lines changed: 0 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -58,55 +58,6 @@ TVM_DLL Pass CreatePrimFuncPass(
5858
const runtime::TypedPackedFunc<PrimFunc(PrimFunc, IRModule, PassContext)>& pass_func,
5959
int opt_level, String name, tvm::Array<String> required, bool traceable = false);
6060

61-
/*!
62-
* \brief Inject prefetch instructions into stmt.
63-
*
64-
* \return The pass.
65-
*/
66-
TVM_DLL Pass InjectPrefetch();
67-
68-
// TODO(tvm-team): consolidate configs to the PassContext
69-
/*!
70-
* \brief Flatten the multi-dimensional read/write
71-
* to single dimensional Load/Store
72-
*
73-
* \param cache_line_size The size of CPU cache line.
74-
* \param create_bound_attribute Whether to create bound attributes.
75-
*
76-
* \return The Pass
77-
*/
78-
TVM_DLL Pass StorageFlatten(int cache_line_size, bool create_bound_attribute = false);
79-
80-
/*!
81-
* \brief Inject copy intrinsics with optional pad.
82-
*
83-
* \param pragma_key The pragma key for hint of copy.
84-
* \param fintrin The function with signature
85-
*
86-
* Stmt fintrin(Buffer src,
87-
* Buffer dst,
88-
* Array<Expr> pad_before,
89-
* Array<Expr> pad_after,
90-
* Expr pad_value)
91-
* \return The pass.
92-
*/
93-
TVM_DLL Pass InjectCopyIntrin(String pragma_key, runtime::PackedFunc fintrin);
94-
95-
/*!
96-
* \brief Detect and insert sync points to co-processor.
97-
*
98-
* \return The pass.
99-
*/
100-
TVM_DLL Pass CoProcSync();
101-
102-
/*!
103-
* \brief Lift common attrs with attr_key to outer scope.
104-
*
105-
* \param attr_key The attribute key to be checked.
106-
* \return The pass.
107-
*/
108-
TVM_DLL Pass LiftAttrScope(String attr_key);
109-
11061
/*!
11162
* \brief partition loops in the stmt.
11263
*
@@ -573,15 +524,6 @@ TVM_DLL Pass LowerOpaqueBlock();
573524
*/
574525
TVM_DLL Pass FlattenBuffer();
575526

576-
/*
577-
* \brief Flatten the multi-dimensional read/write
578-
* to two dimensional texture Load/Store and realize
579-
* texture buffer allocations.
580-
*
581-
* \return The Pass
582-
*/
583-
TVM_DLL Pass TextureFlatten();
584-
585527
/*
586528
* \brief Lower VTCM allocations
587529
*

‎python/tvm/tir/pipeline.py‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -31,11 +31,6 @@ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.I
3131
pass_ctx = tvm.transform.PassContext.current()
3232
config = pass_ctx.config
3333
passes = [
34-
tir.transform.InjectPrefetch(),
35-
tir.transform.TextureFlatten(),
36-
tir.transform.StorageFlatten(
37-
64, bool(config.get("tir.instrument_bound_checkers", False))
38-
),
3934
tir.transform.LowerCrossThreadReduction(),
4035
tir.transform.LowerInitBlock(),
4136
tir.transform.PlanAndUpdateBufferAllocationLocation(),

‎python/tvm/tir/transform/transform.py‎

Lines changed: 1 addition & 107 deletions
Original file line numberDiff line numberDiff line change
@@ -48,112 +48,6 @@ def _transform(func, mod, ctx):
4848
return _fpass.prim_func_pass(_transform, opt_level=0, name="Apply") # type: ignore
4949

5050

51-
def InjectPrefetch():
52-
"""Inject prefetch instructions into stmt.
53-
54-
Returns
55-
-------
56-
fpass : tvm.transform.Pass
57-
The result pass
58-
"""
59-
return _ffi_api.InjectPrefetch() # type: ignore
60-
61-
62-
def ApplyLayoutTransforms():
63-
"""Reshape buffers that appear in the "layout_transform_map"
64-
fucntion attribute.
65-
66-
Returns
67-
-------
68-
fpass : tvm.transform.Pass
69-
The result pass
70-
71-
"""
72-
return _ffi_api.ApplyLayoutTransforms() # type: ignore
73-
74-
75-
def StorageFlatten(cache_line_size, create_bound_attribute: bool = False):
76-
"""Flatten the multi-dimensional read/write to 1D.
77-
78-
79-
Parameters
80-
----------
81-
cache_line_size: int
82-
The size of CPU cache line.
83-
84-
create_bound_attribute:
85-
Whether to create bound attributes.
86-
87-
88-
Returns
89-
-------
90-
fpass : tvm.transform.Pass
91-
The result pass
92-
"""
93-
return _ffi_api.StorageFlatten(cache_line_size, create_bound_attribute) # type: ignore
94-
95-
96-
def TextureFlatten():
97-
"""Flatten the multi-dimensional read/write to 2D.
98-
99-
100-
Parameters
101-
----------
102-
103-
Returns
104-
-------
105-
fpass : tvm.transform.Pass
106-
The result pass
107-
"""
108-
return _ffi_api.TextureFlatten() # type: ignore
109-
110-
111-
def InjectCopyIntrin(pragma_key: str, fintrin):
112-
"""Inject virtual thread loops.
113-
114-
Parameters
115-
----------
116-
pragma_key : str
117-
The pragma key for hint of copy.
118-
119-
fintrin : function
120-
The function with signature copyintrin(src, dst, pad_before, pad_after, pad_value)
121-
122-
Returns
123-
-------
124-
fpass : tvm.transform.Pass
125-
The result pass
126-
"""
127-
return _ffi_api.InjectCopyIntrin(pragma_key, fintrin) # type: ignore
128-
129-
130-
def CoProcSync():
131-
"""Detect and insert sync points to co-processor.
132-
133-
Returns
134-
-------
135-
fpass : tvm.transform.Pass
136-
The result pass
137-
"""
138-
return _ffi_api.CoProcSync() # type: ignore
139-
140-
141-
def LiftAttrScope(attr_key: str):
142-
"""Lift common attrs with attr_key to outer scope.
143-
144-
Parameters
145-
----------
146-
attr_key : str
147-
The attribute key to be checked.
148-
149-
Returns
150-
-------
151-
fpass : tvm.transform.Pass
152-
The result pass
153-
"""
154-
return _ffi_api.LiftAttrScope(attr_key) # type: ignore
155-
156-
15751
def LoopPartition():
15852
"""Inject virtual thread loops.
15953
@@ -682,7 +576,7 @@ def NarrowDataType(target_bits: int):
682576
683577
Note
684578
----
685-
Run this pass after StorageFlatten.
579+
Run this pass after FlattenBuffer.
686580
"""
687581
return _ffi_api.NarrowDataType(target_bits) # type: ignore
688582

‎src/README.md‎

Lines changed: 0 additions & 37 deletions
This file was deleted.

‎src/meta_schedule/postproc/verify_gpu_code.cc‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -153,10 +153,6 @@ class VerifyGPUCodeNode : public PostprocNode {
153153
try {
154154
auto pass_list = Array<tvm::transform::Pass>();
155155
// Phase 1
156-
// First three passes are not needed in TIR schedule.
157-
// pass_list.push_back(tir::transform::InjectPrefetch());
158-
// pass_list.push_back(tir::transform::TextureFlatten());
159-
// pass_list.push_back(tir::transform::StorageFlatten(64, instrument_bound_checkers));
160156
pass_list.push_back(tir::transform::LowerCrossThreadReduction());
161157
pass_list.push_back(tir::transform::LowerInitBlock());
162158
pass_list.push_back(tir::transform::PlanAndUpdateBufferAllocationLocation());

‎src/target/stackvm/codegen_stackvm.cc‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -141,8 +141,7 @@ int CodeGenStackVM::GetVarID(const VarNode* v) const {
141141

142142
void CodeGenStackVM::VisitExpr_(const BufferLoadNode* op) {
143143
ICHECK_EQ(op->indices.size(), 1) << "StackVM expects flat 1-d buffers. "
144-
<< "Has StorageFlatten (TE-based schedules) or "
145-
<< "FlattenBuffer (TIR-based schedules) been run?";
144+
<< "Has FlattenBuffer been run?";
146145
auto index = op->indices[0];
147146

148147
this->Push(op->buffer->data);
@@ -160,8 +159,7 @@ void CodeGenStackVM::VisitExpr_(const BufferLoadNode* op) {
160159

161160
void CodeGenStackVM::VisitStmt_(const BufferStoreNode* op) {
162161
ICHECK_EQ(op->indices.size(), 1) << "StackVM expects flat 1-d buffers. "
163-
<< "Has StorageFlatten (TE-based schedules) or "
164-
<< "FlattenBuffer (TIR-based schedules) been run?";
162+
<< "Has FlattenBuffer been run?";
165163
auto index = op->indices[0];
166164

167165
this->Push(op->buffer->data);

‎src/tir/transforms/bound_checker.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ namespace tvm {
4040
namespace tir {
4141

4242
// TODO(Lunderberg): Move this pass to be before
43-
// StorageFlatten/FlattenBuffer. That will simplify this pass,
43+
// FlattenBuffer. That will simplify this pass,
4444
// because it can check directly against the buffer limits.
4545
class BoundCollector : public StmtVisitor {
4646
public:

0 commit comments

Comments
 (0)