[REFACTOR] Cleanup legacy TE-based passes

include/tvm/tir/function.h

@@ -89,8 +89,7 @@ class PrimFuncNode : public BaseFuncNode {
    *  normal statements, making buffer_map as first class citizen of PrimFunc
    *  will make program analysis much easier.
    *
-   *  Prior to buffer flattening, which is performed either in
-   *  StorageFlatten for TE-based schedules or in FlattenBuffer for
+   *  Prior to buffer flattening, which is performed FlattenBuffer for
    *  TIR-based schedules, these buffer objects are used directly in
    *  the body of the function.  After buffer flattening, these buffer
    *  objects remain unflattened for use in argument validation, but

include/tvm/tir/transform.h

@@ -58,55 +58,6 @@ TVM_DLL Pass CreatePrimFuncPass(
     const runtime::TypedPackedFunc<PrimFunc(PrimFunc, IRModule, PassContext)>& pass_func,
     int opt_level, String name, tvm::Array<String> required, bool traceable = false);
 
-/*!
- * \brief Inject prefetch instructions into stmt.
- *
- * \return The pass.
- */
-TVM_DLL Pass InjectPrefetch();
-
-// TODO(tvm-team): consolidate configs to the PassContext
-/*!
- * \brief Flatten the multi-dimensional read/write
- *  to single dimensional Load/Store
- *
- * \param cache_line_size The size of CPU cache line.
- * \param create_bound_attribute Whether to create bound attributes.
- *
- * \return The Pass
- */
-TVM_DLL Pass StorageFlatten(int cache_line_size, bool create_bound_attribute = false);
-
-/*!
- * \brief Inject copy intrinsics with optional pad.
- *
- * \param pragma_key The pragma key for hint of copy.
- * \param fintrin The function with signature
- *
- *   Stmt fintrin(Buffer src,
- *                Buffer dst,
- *                Array<Expr> pad_before,
- *                Array<Expr> pad_after,
- *                Expr pad_value)
- * \return The pass.
- */
-TVM_DLL Pass InjectCopyIntrin(String pragma_key, runtime::PackedFunc fintrin);
-
-/*!
- * \brief Detect and insert sync points to co-processor.
- *
- * \return The pass.
- */
-TVM_DLL Pass CoProcSync();
-
-/*!
- * \brief Lift common attrs with attr_key to outer scope.
- *
- * \param attr_key The attribute key to be checked.
- * \return The pass.
- */
-TVM_DLL Pass LiftAttrScope(String attr_key);
-
 /*!
  * \brief partition loops in the stmt.
  *
@@ -573,15 +524,6 @@ TVM_DLL Pass LowerOpaqueBlock();
  */
 TVM_DLL Pass FlattenBuffer();
 
-/*
- * \brief Flatten the multi-dimensional read/write
- *  to two dimensional texture Load/Store and realize
- *  texture buffer allocations.
- *
- * \return The Pass
- */
-TVM_DLL Pass TextureFlatten();
-
 /*
  * \brief Lower VTCM allocations
  *

python/tvm/tir/pipeline.py

@@ -31,11 +31,6 @@ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.I
         pass_ctx = tvm.transform.PassContext.current()
         config = pass_ctx.config
         passes = [
-            tir.transform.InjectPrefetch(),
-            tir.transform.TextureFlatten(),
-            tir.transform.StorageFlatten(
-                64, bool(config.get("tir.instrument_bound_checkers", False))
-            ),
             tir.transform.LowerCrossThreadReduction(),
             tir.transform.LowerInitBlock(),
             tir.transform.PlanAndUpdateBufferAllocationLocation(),

python/tvm/tir/transform/transform.py

@@ -48,112 +48,6 @@ def _transform(func, mod, ctx):
     return _fpass.prim_func_pass(_transform, opt_level=0, name="Apply")  # type: ignore
 
 
-def InjectPrefetch():
-    """Inject prefetch instructions into stmt.
-
-    Returns
-    -------
-    fpass : tvm.transform.Pass
-        The result pass
-    """
-    return _ffi_api.InjectPrefetch()  # type: ignore
-
-
-def ApplyLayoutTransforms():
-    """Reshape buffers that appear in the "layout_transform_map"
-    fucntion attribute.
-
-    Returns
-    -------
-    fpass : tvm.transform.Pass
-        The result pass
-
-    """
-    return _ffi_api.ApplyLayoutTransforms()  # type: ignore
-
-
-def StorageFlatten(cache_line_size, create_bound_attribute: bool = False):
-    """Flatten the multi-dimensional read/write to 1D.
-
-
-    Parameters
-    ----------
-    cache_line_size: int
-        The size of CPU cache line.
-
-    create_bound_attribute:
-        Whether to create bound attributes.
-
-
-    Returns
-    -------
-    fpass : tvm.transform.Pass
-        The result pass
-    """
-    return _ffi_api.StorageFlatten(cache_line_size, create_bound_attribute)  # type: ignore
-
-
-def TextureFlatten():
-    """Flatten the multi-dimensional read/write to 2D.
-
-
-    Parameters
-    ----------
-
-    Returns
-    -------
-    fpass : tvm.transform.Pass
-        The result pass
-    """
-    return _ffi_api.TextureFlatten()  # type: ignore
-
-
-def InjectCopyIntrin(pragma_key: str, fintrin):
-    """Inject virtual thread loops.
-
-    Parameters
-    ----------
-    pragma_key : str
-        The pragma key for hint of copy.
-
-    fintrin : function
-        The function with signature copyintrin(src, dst, pad_before, pad_after, pad_value)
-
-    Returns
-    -------
-    fpass : tvm.transform.Pass
-        The result pass
-    """
-    return _ffi_api.InjectCopyIntrin(pragma_key, fintrin)  # type: ignore
-
-
-def CoProcSync():
-    """Detect and insert sync points to co-processor.
-
-    Returns
-    -------
-    fpass : tvm.transform.Pass
-        The result pass
-    """
-    return _ffi_api.CoProcSync()  # type: ignore
-
-
-def LiftAttrScope(attr_key: str):
-    """Lift common attrs with attr_key to outer scope.
-
-    Parameters
-    ----------
-    attr_key : str
-        The attribute key to be checked.
-
-    Returns
-    -------
-    fpass : tvm.transform.Pass
-        The result pass
-    """
-    return _ffi_api.LiftAttrScope(attr_key)  # type: ignore
-
-
 def LoopPartition():
     """Inject virtual thread loops.
 
@@ -682,7 +576,7 @@ def NarrowDataType(target_bits: int):
 
     Note
     ----
-    Run this pass after StorageFlatten.
+    Run this pass after FlattenBuffer.
     """
     return _ffi_api.NarrowDataType(target_bits)  # type: ignore
 

src/meta_schedule/postproc/verify_gpu_code.cc

@@ -153,10 +153,6 @@ class VerifyGPUCodeNode : public PostprocNode {
         try {
           auto pass_list = Array<tvm::transform::Pass>();
           // Phase 1
-          // First three passes are not needed in TIR schedule.
-          // pass_list.push_back(tir::transform::InjectPrefetch());
-          // pass_list.push_back(tir::transform::TextureFlatten());
-          // pass_list.push_back(tir::transform::StorageFlatten(64, instrument_bound_checkers));
           pass_list.push_back(tir::transform::LowerCrossThreadReduction());
           pass_list.push_back(tir::transform::LowerInitBlock());
           pass_list.push_back(tir::transform::PlanAndUpdateBufferAllocationLocation());

src/target/stackvm/codegen_stackvm.cc

@@ -141,8 +141,7 @@ int CodeGenStackVM::GetVarID(const VarNode* v) const {
 
 void CodeGenStackVM::VisitExpr_(const BufferLoadNode* op) {
   ICHECK_EQ(op->indices.size(), 1) << "StackVM expects flat 1-d buffers.  "
-                                   << "Has StorageFlatten (TE-based schedules) or "
-                                   << "FlattenBuffer (TIR-based schedules) been run?";
+                                   << "Has FlattenBuffer  been run?";
   auto index = op->indices[0];
 
   this->Push(op->buffer->data);
@@ -160,8 +159,7 @@ void CodeGenStackVM::VisitExpr_(const BufferLoadNode* op) {
 
 void CodeGenStackVM::VisitStmt_(const BufferStoreNode* op) {
   ICHECK_EQ(op->indices.size(), 1) << "StackVM expects flat 1-d buffers.  "
-                                   << "Has StorageFlatten (TE-based schedules) or "
-                                   << "FlattenBuffer (TIR-based schedules) been run?";
+                                   << "Has FlattenBuffer been run?";
   auto index = op->indices[0];
 
   this->Push(op->buffer->data);

src/tir/transforms/bound_checker.cc

@@ -40,7 +40,7 @@ namespace tvm {
 namespace tir {
 
 // TODO(Lunderberg): Move this pass to be before
-// StorageFlatten/FlattenBuffer.  That will simplify this pass,
+// FlattenBuffer.  That will simplify this pass,
 // because it can check directly against the buffer limits.
 class BoundCollector : public StmtVisitor {
  public: