[TIR][Schedule] Transform layout

include/tvm/tir/index_map.h

@@ -106,18 +106,34 @@ class IndexMapNode : public Object {
  */
  Array<PrimExpr> MapShape(const Array<PrimExpr>& shape) const;
 
+ /*!
+ * \brief Convert to string representation in Python.
+ * \return The stringified lambda expression in Python.
+ */
+ String ToPythonString() const;
+
  void VisitAttrs(AttrVisitor* v) {
  v->Visit("initial_indices", &initial_indices);
  v->Visit("final_indices", &final_indices);
  }
 
+ static constexpr const char* _type_key = "tir.IndexMap";
+
  TVM_DECLARE_FINAL_OBJECT_INFO(IndexMapNode, Object);
 };
 
 class IndexMap : public ObjectRef {
  public:
  IndexMap(Array<Var> initial_indices, Array<PrimExpr> final_indices);
 
+ /*!
+ * \brief Create an index map from a packed function
+ * \param ndim The number of dimensions
+ * \param func The function to be applied
+ * \return The created index map
+ */
+ static IndexMap FromFunc(int ndim, runtime::TypedPackedFunc<Array<PrimExpr>(Array<Var>)> func);
+
  /*! \brief Generate the inverse mapping.
  *
  * The range of the input indices is required in order to ensure

include/tvm/tir/schedule/schedule.h

@@ -20,6 +20,7 @@
 #define TVM_TIR_SCHEDULE_SCHEDULE_H_
 
 #include <tvm/support/random_engine.h>
+#include <tvm/tir/index_map.h>
 #include <tvm/tir/schedule/state.h>
 #include <tvm/tir/schedule/trace.h>
 
@@ -36,6 +37,14 @@ enum class ScheduleErrorRenderLevel : int32_t {
  kNone = 2,
 };
 
+/*! \brief Type of buffer index */
+enum class BufferIndexType : int32_t {
+ /*! \brief Index of a read buffer */
+ kRead = 0,
+ /*! \brief Index of a written buffer */
+ kWrite = 1,
+};
+
 /**************** Random variable: BlockRV ****************/
 
 /*! \brief A random variable that evaluates to a TensorIR block */
@@ -521,6 +530,21 @@ class ScheduleNode : public runtime::Object {
  */
  virtual void Unannotate(const BlockRV& block_rv, const String& ann_key) = 0;
 
+ /******** Schedule: Layout transformation ********/
+ /*!
+ * \brief Apply a transformation represented by IndexMap to buffer
+ * \details The indices and the access region to the target buffer is transformed by the given
+ * index_map. The index_map is used to infer the new shape of the buffer. Buffer must be either
+ * a function parameter, or allocated in a block (it cannot be a buffer subregion created via
+ * 'match_buffer').
+ * \param block_rv The block that accesses the target buffer.
+ * \param buffer_index The index of the buffer in block's read or write region.
+ * \param buffer_index_type The type of the buffer index, kRead or kWrite.
+ * \param index_map The transformation to apply.
+ */
+ virtual void TransformLayout(const BlockRV& block_rv, int buffer_index,
+ BufferIndexType buffer_index_type, const IndexMap& index_map) = 0;
+
  /******** Schedule: Misc ********/
  /*! \brief A no-op that marks the start of postprocessing phase of scheduling */
  virtual void EnterPostproc() = 0;

python/tvm/tir/__init__.py

@@ -57,7 +57,7 @@
 from .op import comm_reducer, min, max, sum
 from .op import q_multiply_shift
 
-from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError
+from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError, BufferType
 
 from . import schedule
 from . import ir_builder

python/tvm/tir/function.py

@@ -16,7 +16,8 @@
 # under the License.
 """Function data types."""
 
-from typing import Mapping, Union
+from typing import Callable, List, Mapping, Optional, Union
+import inspect
 
 import tvm._ffi
 import tvm.runtime
@@ -239,3 +240,58 @@ def get(name: str):
  The TensorIntrin with the specified name.
  """
  return _ffi_api.TensorIntrinGet(name) # pylint: type: ignore
+
+
+@tvm._ffi.register_object("tir.IndexMap")
+class IndexMap(Object):
+ """A mapping from multi-dimensional indices to another set of multi-dimensional indices
+
+ Parameters
+ ----------
+ initial_indices : List[Var]
+ Variables representing the indices prior to remapping.
+ final_indices : List[PrimExpr]
+ Expressions defining the indices after remapping.
+ """
+
+ initial_indices: List[Var]
+ final_indices: List[PrimExpr]
+
+ def __init__(self, initial_indices, final_indices):
+ self.__init_handle_by_constructor__(_ffi_api.IndexMap, initial_indices, final_indices)
+
+ @staticmethod
+ def from_func(mapping_function: Callable, ndim: Optional[int] = None):
+ """Create an index map from a function
+
+ Parameters
+ ----------
+ mapping_function : Callable
+ The function to map from source indices to target indices
+ """
+ params = inspect.signature(mapping_function).parameters
+ default_index_dtype = "int32"
+ args = []
+ var_arg_name = None
+ for name, param in params.items():
+ if param.kind in [
+ inspect.Parameter.POSITIONAL_ONLY,
+ inspect.Parameter.POSITIONAL_OR_KEYWORD,
+ ]:
+ args.append(tvm.tir.Var(name, default_index_dtype))
+ elif param.kind == inspect.Parameter.VAR_POSITIONAL:
+ var_arg_name = name
+ else:
+ raise ValueError("transform_layout mapping may not have *args or **kwargs")
+
+ # Now that all the named arguments have been collected,
+ # everything that remains should go to the *args, if
+ # specified.
+ if var_arg_name is not None:
+ assert ndim is not None, "ndim must be specified when *args is used"
+ num_var_args = ndim - len(args)
+ for i in range(num_var_args):
+ args.append(tvm.tir.Var(f"{var_arg_name}_{i}", default_index_dtype))
+
+ final_indices = mapping_function(*args)
+ return IndexMap(args, final_indices)

python/tvm/tir/schedule/__init__.py

@@ -19,6 +19,6 @@
 
 from .block_scope import BlockScope, Dependency, DepKind, StmtSRef
 from .instruction import Instruction, InstructionKind
-from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError
+from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError, BufferType
 from .state import ScheduleDebugMask, ScheduleState
 from .trace import Trace

python/tvm/tir/schedule/schedule.py

@@ -15,13 +15,15 @@
 # specific language governing permissions and limitations
 # under the License.
 """The TensorIR schedule class"""
-from typing import Dict, List, Optional, Union
+import enum
+from typing import Callable, Dict, List, Optional, Union
 
 from tvm._ffi import register_object as _register_object
 from tvm.error import TVMError, register_error
 from tvm.ir import IRModule, PrimExpr
 from tvm.runtime import Object, String
 from tvm.tir import Block, FloatImm, For, IntImm, PrimFunc
+from ..function import IndexMap
 
 from . import _ffi_api
 from .state import ScheduleState, StmtSRef, _parse_debug_mask, _parse_mod
@@ -71,6 +73,13 @@ def __init__(self) -> None:
 }
 
 
+class BufferType(enum.IntEnum):
+ """Type of buffer in access regions of a block"""
+
+ READ = 0
+ WRITE = 1
+
+
 def _parse_error_render_level(error_render_level: str) -> int:
  if error_render_level not in _ERROR_RENDER_LEVEL:
  raise ValueError(
@@ -2111,6 +2120,82 @@ def after_unannotate(a: T.handle, b: T.handle) -> None:
  self, block_or_loop, ann_key
  )
 
+ ########## Schedule: Layout transformation ##########
+
+ @type_checked
+ def transform_layout(
+ self,
+ block: BlockRV,
+ buffer_index: int,
+ buffer_type: BufferType,
+ index_map: Union[IndexMap, Callable],
+ ) -> None:
+ """Apply a transformation represented by IndexMap to buffer
+ Parameters
+ ----------
+ block_rv : BlockRV
+ The block that accesses the target buffer
+ buffer_index: int
+ The index of the buffer in block's read or write region
+ buffer_type : BufferType
+ Type of the buffer, READ or WRITE.
+ index_map : Union[IndexMap, Callable]
+ The transformation to apply
+
+ Examples
+ --------
+ Before transform_layout, in TensorIR, the IR is:
+
+ .. code-block:: python
+
+ @T.prim_func
+ def before_transform_layout(a: T.handle, c: T.handle) -> None:
+ A = T.match_buffer(a, (128, 128), "float32")
+ B = T.alloc_buffer((128, 128), "float32")
+ C = T.match_buffer(c, (128, 128), "float32")
+ for i, j in T.grid(128, 128):
+ with T.block("B"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ B[vi, vj] = A[vi, vj] * 2.0
+ for i, j in T.grid(128, 128):
+ with T.block("C"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ C[vi, vj] = B[vi, vj] + 1.0
+
+ Create the schedule and do transform_layout:
+
+ .. code-block:: python
+
+ sch = tir.Schedule(before_storage_align)
+ sch.transform_layout(sch.get_block("B"), buffer_index=0, BufferType.WRITE,
+ index_map=lambda m, n: (m // 16, n // 16, m % 16, n % 16))
+ print(sch.mod["main"].script())
+
+ After applying transform_layout, the IR becomes:
+
+ .. code-block:: python
+
+ @T.prim_func
+ def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> None:
+ A = T.match_buffer(a, (128, 128), "float32")
+ B = T.alloc_buffer((8, 8, 16, 16), "float32")
+ C = T.match_buffer(c, (128, 128), "float32")
+ for i, j in T.grid(128, 128):
+ with T.block("B"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ B[vi // 16, vj // 16, vi % 16, vj % 16] = A[vi, vj] * 2.0
+ for i, j in T.grid(128, 128):
+ with T.block("C"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ C[vi, vj] = B[vi // 16, vj // 16, vi % 16, vj % 16] + 1.0
+
+ """
+ if callable(index_map):
+ index_map = IndexMap.from_func(index_map)
+ _ffi_api.ScheduleTransformLayout( # type: ignore # pylint: disable=no-member
+ self, block, buffer_index, buffer_type, index_map
+ )
+
  ########## Schedule: Misc ##########
 
  @type_checked

src/tir/ir/index_map.cc

@@ -27,6 +27,7 @@
 #include <tvm/arith/int_set.h>
 #include <tvm/arith/iter_affine_map.h>
 #include <tvm/tir/op.h>
+#include <tvm/tir/stmt_functor.h>
 
 #include <sstream>
 
@@ -40,6 +41,15 @@ IndexMap::IndexMap(Array<Var> initial_indices, Array<PrimExpr> final_indices) {
  data_ = std::move(n);
 }
 
+IndexMap IndexMap::FromFunc(int ndim, runtime::TypedPackedFunc<Array<PrimExpr>(Array<Var>)> func) {
+ Array<Var> initial_indices;
+ initial_indices.reserve(ndim);
+ for (int i = 0; i < ndim; ++i) {
+ initial_indices.push_back(Var("i" + std::to_string(i), DataType::Int(32)));
+ }
+ return IndexMap(initial_indices, func(initial_indices));
+}
+
 IndexMap IndexMap::Inverse(Array<Range> initial_ranges) const {
  // Dummy variables to represent the inverse's inputs.
  Array<Var> output_vars;
@@ -142,13 +152,52 @@ Array<PrimExpr> IndexMapNode::MapShape(const Array<PrimExpr>& shape) const {
  return output;
 }
 
+String IndexMapNode::ToPythonString() const {
+ std::unordered_set<std::string> used_names;
+ Map<Var, PrimExpr> var_remap;
+ for (const Var& initial_index : initial_indices) {
+ if (used_names.count(initial_index->name_hint)) {
+ std::string new_name = initial_index->name_hint + std::to_string(used_names.size());
+ used_names.insert(new_name);
+ var_remap.Set(initial_index, Var(new_name));
+ } else {
+ used_names.insert(initial_index->name_hint);
+ }
+ }
+ std::ostringstream oss;
+ oss << "lambda ";
+ for (size_t i = 0; i < initial_indices.size(); ++i) {
+ if (i != 0) {
+ oss << ", ";
+ }
+ auto it = var_remap.find(initial_indices[i]);
+ if (it != var_remap.end()) {
+ oss << (*it).second;
+ } else {
+ oss << initial_indices[i];
+ }
+ }
+ oss << ": (";
+ for (size_t i = 0; i < final_indices.size(); ++i) {
+ oss << Substitute(final_indices[i], var_remap);
+ oss << ", ";
+ }
+ oss << ")";
+ return String(oss.str());
+}
+
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
  .set_dispatch<IndexMapNode>([](const ObjectRef& node, ReprPrinter* p) {
  auto* op = static_cast<const IndexMapNode*>(node.get());
- p->stream << "index_map(" << op->initial_indices << ", " << op->final_indices << ")";
+ p->stream << "index_map(" << op->ToPythonString() << ")";
  });
 
 TVM_REGISTER_NODE_TYPE(IndexMapNode);
 
+TVM_REGISTER_GLOBAL("tir.IndexMap")
+ .set_body_typed([](Array<Var> initial_indices, Array<PrimExpr> final_indices) {
+ return IndexMap(initial_indices, final_indices);
+ });
+
 } // namespace tir
 } // namespace tvm

src/tir/schedule/analysis.h

@@ -401,6 +401,16 @@ struct ProducerConsumerSplit {
  */
 Buffer GetNthAccessBuffer(const ScheduleState& self, const Block& block, int n, bool is_write);
 
+/*!
+ * \brief Find the defining site of the buffer in the given block and its ancestors
+ * \param block_sref The block sref
+ * \param buffer The buffer
+ * \return The defining site of the buffer and whether the buffer is allocated (otherwise the
+ * buffer is from match_buffer).
+ */
+std::pair<Optional<StmtSRef>, bool> GetBufferDefiningSite(const StmtSRef& block_sref,
+ const Buffer& buffer);
+
 /******** Reduction Block Related ********/
 
 /*!

include/tvm/tir/index_map.h

@@ -106,18 +106,34 @@ class IndexMapNode : public Object {
  */
  Array<PrimExpr> MapShape(const Array<PrimExpr>& shape) const;
 
+ /*!
+ * \brief Convert to string representation in Python.
+ * \return The stringified lambda expression in Python.
+ */
+ String ToPythonString() const;
+
  void VisitAttrs(AttrVisitor* v) {
  v->Visit("initial_indices", &initial_indices);
  v->Visit("final_indices", &final_indices);
  }
 
+ static constexpr const char* _type_key = "tir.IndexMap";
+
  TVM_DECLARE_FINAL_OBJECT_INFO(IndexMapNode, Object);
 };
 
 class IndexMap : public ObjectRef {
  public:
  IndexMap(Array<Var> initial_indices, Array<PrimExpr> final_indices);
 
+ /*!
+ * \brief Create an index map from a packed function
+ * \param ndim The number of dimensions
+ * \param func The function to be applied
+ * \return The created index map
+ */
+ static IndexMap FromFunc(int ndim, runtime::TypedPackedFunc<Array<PrimExpr>(Array<Var>)> func);
+
  /*! \brief Generate the inverse mapping.
  *
  * The range of the input indices is required in order to ensure

include/tvm/tir/schedule/schedule.h

@@ -20,6 +20,7 @@
 #define TVM_TIR_SCHEDULE_SCHEDULE_H_
 
 #include <tvm/support/random_engine.h>
+#include <tvm/tir/index_map.h>
 #include <tvm/tir/schedule/state.h>
 #include <tvm/tir/schedule/trace.h>
 
@@ -36,6 +37,14 @@ enum class ScheduleErrorRenderLevel : int32_t {
  kNone = 2,
 };
 
+/*! \brief Type of buffer index */
+enum class BufferIndexType : int32_t {
+ /*! \brief Index of a read buffer */
+ kRead = 0,
+ /*! \brief Index of a written buffer */
+ kWrite = 1,
+};
+
 /**************** Random variable: BlockRV ****************/
 
 /*! \brief A random variable that evaluates to a TensorIR block */
@@ -521,6 +530,21 @@ class ScheduleNode : public runtime::Object {
  */
  virtual void Unannotate(const BlockRV& block_rv, const String& ann_key) = 0;
 
+ /******** Schedule: Layout transformation ********/
+ /*!
+ * \brief Apply a transformation represented by IndexMap to buffer
+ * \details The indices and the access region to the target buffer is transformed by the given
+ * index_map. The index_map is used to infer the new shape of the buffer. Buffer must be either
+ * a function parameter, or allocated in a block (it cannot be a buffer subregion created via
+ * 'match_buffer').
+ * \param block_rv The block that accesses the target buffer.
+ * \param buffer_index The index of the buffer in block's read or write region.
+ * \param buffer_index_type The type of the buffer index, kRead or kWrite.
+ * \param index_map The transformation to apply.
+ */
+ virtual void TransformLayout(const BlockRV& block_rv, int buffer_index,
+ BufferIndexType buffer_index_type, const IndexMap& index_map) = 0;
+
  /******** Schedule: Misc ********/
  /*! \brief A no-op that marks the start of postprocessing phase of scheduling */
  virtual void EnterPostproc() = 0;

python/tvm/tir/__init__.py

@@ -57,7 +57,7 @@
 from .op import comm_reducer, min, max, sum
 from .op import q_multiply_shift
 
-from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError
+from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError, BufferType
 
 from . import schedule
 from . import ir_builder

python/tvm/tir/function.py

@@ -16,7 +16,8 @@
 # under the License.
 """Function data types."""
 
-from typing import Mapping, Union
+from typing import Callable, List, Mapping, Optional, Union
+import inspect
 
 import tvm._ffi
 import tvm.runtime
@@ -239,3 +240,58 @@ def get(name: str):
  The TensorIntrin with the specified name.
  """
  return _ffi_api.TensorIntrinGet(name) # pylint: type: ignore
+
+
+@tvm._ffi.register_object("tir.IndexMap")
+class IndexMap(Object):
+ """A mapping from multi-dimensional indices to another set of multi-dimensional indices
+
+ Parameters
+ ----------
+ initial_indices : List[Var]
+ Variables representing the indices prior to remapping.
+ final_indices : List[PrimExpr]
+ Expressions defining the indices after remapping.
+ """
+
+ initial_indices: List[Var]
+ final_indices: List[PrimExpr]
+
+ def __init__(self, initial_indices, final_indices):
+ self.__init_handle_by_constructor__(_ffi_api.IndexMap, initial_indices, final_indices)
+
+ @staticmethod
+ def from_func(mapping_function: Callable, ndim: Optional[int] = None):
+ """Create an index map from a function
+
+ Parameters
+ ----------
+ mapping_function : Callable
+ The function to map from source indices to target indices
+ """
+ params = inspect.signature(mapping_function).parameters
+ default_index_dtype = "int32"
+ args = []
+ var_arg_name = None
+ for name, param in params.items():
+ if param.kind in [
+ inspect.Parameter.POSITIONAL_ONLY,
+ inspect.Parameter.POSITIONAL_OR_KEYWORD,
+ ]:
+ args.append(tvm.tir.Var(name, default_index_dtype))
+ elif param.kind == inspect.Parameter.VAR_POSITIONAL:
+ var_arg_name = name
+ else:
+ raise ValueError("transform_layout mapping may not have *args or **kwargs")
+
+ # Now that all the named arguments have been collected,
+ # everything that remains should go to the *args, if
+ # specified.
+ if var_arg_name is not None:
+ assert ndim is not None, "ndim must be specified when *args is used"
+ num_var_args = ndim - len(args)
+ for i in range(num_var_args):
+ args.append(tvm.tir.Var(f"{var_arg_name}_{i}", default_index_dtype))
+
+ final_indices = mapping_function(*args)
+ return IndexMap(args, final_indices)

python/tvm/tir/schedule/__init__.py

@@ -19,6 +19,6 @@
 
 from .block_scope import BlockScope, Dependency, DepKind, StmtSRef
 from .instruction import Instruction, InstructionKind
-from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError
+from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError, BufferType
 from .state import ScheduleDebugMask, ScheduleState
 from .trace import Trace

python/tvm/tir/schedule/schedule.py

@@ -15,13 +15,15 @@
 # specific language governing permissions and limitations
 # under the License.
 """The TensorIR schedule class"""
-from typing import Dict, List, Optional, Union
+import enum
+from typing import Callable, Dict, List, Optional, Union
 
 from tvm._ffi import register_object as _register_object
 from tvm.error import TVMError, register_error
 from tvm.ir import IRModule, PrimExpr
 from tvm.runtime import Object, String
 from tvm.tir import Block, FloatImm, For, IntImm, PrimFunc
+from ..function import IndexMap
 
 from . import _ffi_api
 from .state import ScheduleState, StmtSRef, _parse_debug_mask, _parse_mod
@@ -71,6 +73,13 @@ def __init__(self) -> None:
 }
 
 
+class BufferType(enum.IntEnum):
+ """Type of buffer in access regions of a block"""
+
+ READ = 0
+ WRITE = 1
+
+
 def _parse_error_render_level(error_render_level: str) -> int:
  if error_render_level not in _ERROR_RENDER_LEVEL:
  raise ValueError(
@@ -2111,6 +2120,82 @@ def after_unannotate(a: T.handle, b: T.handle) -> None:
  self, block_or_loop, ann_key
  )
 
+ ########## Schedule: Layout transformation ##########
+
+ @type_checked
+ def transform_layout(
+ self,
+ block: BlockRV,
+ buffer_index: int,
+ buffer_type: BufferType,
+ index_map: Union[IndexMap, Callable],
+ ) -> None:
+ """Apply a transformation represented by IndexMap to buffer
+ Parameters
+ ----------
+ block_rv : BlockRV
+ The block that accesses the target buffer
+ buffer_index: int
+ The index of the buffer in block's read or write region
+ buffer_type : BufferType
+ Type of the buffer, READ or WRITE.
+ index_map : Union[IndexMap, Callable]
+ The transformation to apply
+
+ Examples
+ --------
+ Before transform_layout, in TensorIR, the IR is:
+
+ .. code-block:: python
+
+ @T.prim_func
+ def before_transform_layout(a: T.handle, c: T.handle) -> None:
+ A = T.match_buffer(a, (128, 128), "float32")
+ B = T.alloc_buffer((128, 128), "float32")
+ C = T.match_buffer(c, (128, 128), "float32")
+ for i, j in T.grid(128, 128):
+ with T.block("B"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ B[vi, vj] = A[vi, vj] * 2.0
+ for i, j in T.grid(128, 128):
+ with T.block("C"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ C[vi, vj] = B[vi, vj] + 1.0
+
+ Create the schedule and do transform_layout:
+
+ .. code-block:: python
+
+ sch = tir.Schedule(before_storage_align)
+ sch.transform_layout(sch.get_block("B"), buffer_index=0, BufferType.WRITE,
+ index_map=lambda m, n: (m // 16, n // 16, m % 16, n % 16))
+ print(sch.mod["main"].script())
+
+ After applying transform_layout, the IR becomes:
+
+ .. code-block:: python
+
+ @T.prim_func
+ def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> None:
+ A = T.match_buffer(a, (128, 128), "float32")
+ B = T.alloc_buffer((8, 8, 16, 16), "float32")
+ C = T.match_buffer(c, (128, 128), "float32")
+ for i, j in T.grid(128, 128):
+ with T.block("B"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ B[vi // 16, vj // 16, vi % 16, vj % 16] = A[vi, vj] * 2.0
+ for i, j in T.grid(128, 128):
+ with T.block("C"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ C[vi, vj] = B[vi // 16, vj // 16, vi % 16, vj % 16] + 1.0
+
+ """
+ if callable(index_map):
+ index_map = IndexMap.from_func(index_map)
+ _ffi_api.ScheduleTransformLayout( # type: ignore # pylint: disable=no-member
+ self, block, buffer_index, buffer_type, index_map
+ )
+
  ########## Schedule: Misc ##########
 
  @type_checked

src/tir/ir/index_map.cc

@@ -27,6 +27,7 @@
 #include <tvm/arith/int_set.h>
 #include <tvm/arith/iter_affine_map.h>
 #include <tvm/tir/op.h>
+#include <tvm/tir/stmt_functor.h>
 
 #include <sstream>
 
@@ -40,6 +41,15 @@ IndexMap::IndexMap(Array<Var> initial_indices, Array<PrimExpr> final_indices) {
  data_ = std::move(n);
 }
 
+IndexMap IndexMap::FromFunc(int ndim, runtime::TypedPackedFunc<Array<PrimExpr>(Array<Var>)> func) {
+ Array<Var> initial_indices;
+ initial_indices.reserve(ndim);
+ for (int i = 0; i < ndim; ++i) {
+ initial_indices.push_back(Var("i" + std::to_string(i), DataType::Int(32)));
+ }
+ return IndexMap(initial_indices, func(initial_indices));
+}
+
 IndexMap IndexMap::Inverse(Array<Range> initial_ranges) const {
  // Dummy variables to represent the inverse's inputs.
  Array<Var> output_vars;
@@ -142,13 +152,52 @@ Array<PrimExpr> IndexMapNode::MapShape(const Array<PrimExpr>& shape) const {
  return output;
 }
 
+String IndexMapNode::ToPythonString() const {
+ std::unordered_set<std::string> used_names;
+ Map<Var, PrimExpr> var_remap;
+ for (const Var& initial_index : initial_indices) {
+ if (used_names.count(initial_index->name_hint)) {
+ std::string new_name = initial_index->name_hint + std::to_string(used_names.size());
+ used_names.insert(new_name);
+ var_remap.Set(initial_index, Var(new_name));
+ } else {
+ used_names.insert(initial_index->name_hint);
+ }
+ }
+ std::ostringstream oss;
+ oss << "lambda ";
+ for (size_t i = 0; i < initial_indices.size(); ++i) {
+ if (i != 0) {
+ oss << ", ";
+ }
+ auto it = var_remap.find(initial_indices[i]);
+ if (it != var_remap.end()) {
+ oss << (*it).second;
+ } else {
+ oss << initial_indices[i];
+ }
+ }
+ oss << ": (";
+ for (size_t i = 0; i < final_indices.size(); ++i) {
+ oss << Substitute(final_indices[i], var_remap);
+ oss << ", ";
+ }
+ oss << ")";
+ return String(oss.str());
+}
+
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
  .set_dispatch<IndexMapNode>([](const ObjectRef& node, ReprPrinter* p) {
  auto* op = static_cast<const IndexMapNode*>(node.get());
- p->stream << "index_map(" << op->initial_indices << ", " << op->final_indices << ")";
+ p->stream << "index_map(" << op->ToPythonString() << ")";
  });
 
 TVM_REGISTER_NODE_TYPE(IndexMapNode);
 
+TVM_REGISTER_GLOBAL("tir.IndexMap")
+ .set_body_typed([](Array<Var> initial_indices, Array<PrimExpr> final_indices) {
+ return IndexMap(initial_indices, final_indices);
+ });
+
 } // namespace tir
 } // namespace tvm

src/tir/schedule/analysis.h

@@ -401,6 +401,16 @@ struct ProducerConsumerSplit {
  */
 Buffer GetNthAccessBuffer(const ScheduleState& self, const Block& block, int n, bool is_write);
 
+/*!
+ * \brief Find the defining site of the buffer in the given block and its ancestors
+ * \param block_sref The block sref
+ * \param buffer The buffer
+ * \return The defining site of the buffer and whether the buffer is allocated (otherwise the
+ * buffer is from match_buffer).
+ */
+std::pair<Optional<StmtSRef>, bool> GetBufferDefiningSite(const StmtSRef& block_sref,
+ const Buffer& buffer);
+
 /******** Reduction Block Related ********/
 
 /*!

include/tvm/tir/index_map.h

@@ -106,18 +106,34 @@ class IndexMapNode : public Object {
  */
  Array<PrimExpr> MapShape(const Array<PrimExpr>& shape) const;
 
+ /*!
+ * \brief Convert to string representation in Python.
+ * \return The stringified lambda expression in Python.
+ */
+ String ToPythonString() const;
+
  void VisitAttrs(AttrVisitor* v) {
  v->Visit("initial_indices", &initial_indices);
  v->Visit("final_indices", &final_indices);
  }
 
+ static constexpr const char* _type_key = "tir.IndexMap";
+
  TVM_DECLARE_FINAL_OBJECT_INFO(IndexMapNode, Object);
 };
 
 class IndexMap : public ObjectRef {
  public:
  IndexMap(Array<Var> initial_indices, Array<PrimExpr> final_indices);
 
+ /*!
+ * \brief Create an index map from a packed function
+ * \param ndim The number of dimensions
+ * \param func The function to be applied
+ * \return The created index map
+ */
+ static IndexMap FromFunc(int ndim, runtime::TypedPackedFunc<Array<PrimExpr>(Array<Var>)> func);
+
  /*! \brief Generate the inverse mapping.
  *
  * The range of the input indices is required in order to ensure

include/tvm/tir/schedule/schedule.h

@@ -20,6 +20,7 @@
 #define TVM_TIR_SCHEDULE_SCHEDULE_H_
 
 #include <tvm/support/random_engine.h>
+#include <tvm/tir/index_map.h>
 #include <tvm/tir/schedule/state.h>
 #include <tvm/tir/schedule/trace.h>
 
@@ -36,6 +37,14 @@ enum class ScheduleErrorRenderLevel : int32_t {
  kNone = 2,
 };
 
+/*! \brief Type of buffer index */
+enum class BufferIndexType : int32_t {
+ /*! \brief Index of a read buffer */
+ kRead = 0,
+ /*! \brief Index of a written buffer */
+ kWrite = 1,
+};
+
 /**************** Random variable: BlockRV ****************/
 
 /*! \brief A random variable that evaluates to a TensorIR block */
@@ -521,6 +530,21 @@ class ScheduleNode : public runtime::Object {
  */
  virtual void Unannotate(const BlockRV& block_rv, const String& ann_key) = 0;
 
+ /******** Schedule: Layout transformation ********/
+ /*!
+ * \brief Apply a transformation represented by IndexMap to buffer
+ * \details The indices and the access region to the target buffer is transformed by the given
+ * index_map. The index_map is used to infer the new shape of the buffer. Buffer must be either
+ * a function parameter, or allocated in a block (it cannot be a buffer subregion created via
+ * 'match_buffer').
+ * \param block_rv The block that accesses the target buffer.
+ * \param buffer_index The index of the buffer in block's read or write region.
+ * \param buffer_index_type The type of the buffer index, kRead or kWrite.
+ * \param index_map The transformation to apply.
+ */
+ virtual void TransformLayout(const BlockRV& block_rv, int buffer_index,
+ BufferIndexType buffer_index_type, const IndexMap& index_map) = 0;
+
  /******** Schedule: Misc ********/
  /*! \brief A no-op that marks the start of postprocessing phase of scheduling */
  virtual void EnterPostproc() = 0;

python/tvm/tir/__init__.py

@@ -57,7 +57,7 @@
 from .op import comm_reducer, min, max, sum
 from .op import q_multiply_shift
 
-from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError
+from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError, BufferType
 
 from . import schedule
 from . import ir_builder

python/tvm/tir/function.py

@@ -16,7 +16,8 @@
 # under the License.
 """Function data types."""
 
-from typing import Mapping, Union
+from typing import Callable, List, Mapping, Optional, Union
+import inspect
 
 import tvm._ffi
 import tvm.runtime
@@ -239,3 +240,58 @@ def get(name: str):
  The TensorIntrin with the specified name.
  """
  return _ffi_api.TensorIntrinGet(name) # pylint: type: ignore
+
+
+@tvm._ffi.register_object("tir.IndexMap")
+class IndexMap(Object):
+ """A mapping from multi-dimensional indices to another set of multi-dimensional indices
+
+ Parameters
+ ----------
+ initial_indices : List[Var]
+ Variables representing the indices prior to remapping.
+ final_indices : List[PrimExpr]
+ Expressions defining the indices after remapping.
+ """
+
+ initial_indices: List[Var]
+ final_indices: List[PrimExpr]
+
+ def __init__(self, initial_indices, final_indices):
+ self.__init_handle_by_constructor__(_ffi_api.IndexMap, initial_indices, final_indices)
+
+ @staticmethod
+ def from_func(mapping_function: Callable, ndim: Optional[int] = None):
+ """Create an index map from a function
+
+ Parameters
+ ----------
+ mapping_function : Callable
+ The function to map from source indices to target indices
+ """
+ params = inspect.signature(mapping_function).parameters
+ default_index_dtype = "int32"
+ args = []
+ var_arg_name = None
+ for name, param in params.items():
+ if param.kind in [
+ inspect.Parameter.POSITIONAL_ONLY,
+ inspect.Parameter.POSITIONAL_OR_KEYWORD,
+ ]:
+ args.append(tvm.tir.Var(name, default_index_dtype))
+ elif param.kind == inspect.Parameter.VAR_POSITIONAL:
+ var_arg_name = name
+ else:
+ raise ValueError("transform_layout mapping may not have *args or **kwargs")
+
+ # Now that all the named arguments have been collected,
+ # everything that remains should go to the *args, if
+ # specified.
+ if var_arg_name is not None:
+ assert ndim is not None, "ndim must be specified when *args is used"
+ num_var_args = ndim - len(args)
+ for i in range(num_var_args):
+ args.append(tvm.tir.Var(f"{var_arg_name}_{i}", default_index_dtype))
+
+ final_indices = mapping_function(*args)
+ return IndexMap(args, final_indices)

python/tvm/tir/schedule/__init__.py

@@ -19,6 +19,6 @@
 
 from .block_scope import BlockScope, Dependency, DepKind, StmtSRef
 from .instruction import Instruction, InstructionKind
-from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError
+from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError, BufferType
 from .state import ScheduleDebugMask, ScheduleState
 from .trace import Trace

python/tvm/tir/schedule/schedule.py

@@ -15,13 +15,15 @@
 # specific language governing permissions and limitations
 # under the License.
 """The TensorIR schedule class"""
-from typing import Dict, List, Optional, Union
+import enum
+from typing import Callable, Dict, List, Optional, Union
 
 from tvm._ffi import register_object as _register_object
 from tvm.error import TVMError, register_error
 from tvm.ir import IRModule, PrimExpr
 from tvm.runtime import Object, String
 from tvm.tir import Block, FloatImm, For, IntImm, PrimFunc
+from ..function import IndexMap
 
 from . import _ffi_api
 from .state import ScheduleState, StmtSRef, _parse_debug_mask, _parse_mod
@@ -71,6 +73,13 @@ def __init__(self) -> None:
 }
 
 
+class BufferType(enum.IntEnum):
+ """Type of buffer in access regions of a block"""
+
+ READ = 0
+ WRITE = 1
+
+
 def _parse_error_render_level(error_render_level: str) -> int:
  if error_render_level not in _ERROR_RENDER_LEVEL:
  raise ValueError(
@@ -2111,6 +2120,82 @@ def after_unannotate(a: T.handle, b: T.handle) -> None:
  self, block_or_loop, ann_key
  )
 
+ ########## Schedule: Layout transformation ##########
+
+ @type_checked
+ def transform_layout(
+ self,
+ block: BlockRV,
+ buffer_index: int,
+ buffer_type: BufferType,
+ index_map: Union[IndexMap, Callable],
+ ) -> None:
+ """Apply a transformation represented by IndexMap to buffer
+ Parameters
+ ----------
+ block_rv : BlockRV
+ The block that accesses the target buffer
+ buffer_index: int
+ The index of the buffer in block's read or write region
+ buffer_type : BufferType
+ Type of the buffer, READ or WRITE.
+ index_map : Union[IndexMap, Callable]
+ The transformation to apply
+
+ Examples
+ --------
+ Before transform_layout, in TensorIR, the IR is:
+
+ .. code-block:: python
+
+ @T.prim_func
+ def before_transform_layout(a: T.handle, c: T.handle) -> None:
+ A = T.match_buffer(a, (128, 128), "float32")
+ B = T.alloc_buffer((128, 128), "float32")
+ C = T.match_buffer(c, (128, 128), "float32")
+ for i, j in T.grid(128, 128):
+ with T.block("B"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ B[vi, vj] = A[vi, vj] * 2.0
+ for i, j in T.grid(128, 128):
+ with T.block("C"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ C[vi, vj] = B[vi, vj] + 1.0
+
+ Create the schedule and do transform_layout:
+
+ .. code-block:: python
+
+ sch = tir.Schedule(before_storage_align)
+ sch.transform_layout(sch.get_block("B"), buffer_index=0, BufferType.WRITE,
+ index_map=lambda m, n: (m // 16, n // 16, m % 16, n % 16))
+ print(sch.mod["main"].script())
+
+ After applying transform_layout, the IR becomes:
+
+ .. code-block:: python
+
+ @T.prim_func
+ def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> None:
+ A = T.match_buffer(a, (128, 128), "float32")
+ B = T.alloc_buffer((8, 8, 16, 16), "float32")
+ C = T.match_buffer(c, (128, 128), "float32")
+ for i, j in T.grid(128, 128):
+ with T.block("B"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ B[vi // 16, vj // 16, vi % 16, vj % 16] = A[vi, vj] * 2.0
+ for i, j in T.grid(128, 128):
+ with T.block("C"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ C[vi, vj] = B[vi // 16, vj // 16, vi % 16, vj % 16] + 1.0
+
+ """
+ if callable(index_map):
+ index_map = IndexMap.from_func(index_map)
+ _ffi_api.ScheduleTransformLayout( # type: ignore # pylint: disable=no-member
+ self, block, buffer_index, buffer_type, index_map
+ )
+
  ########## Schedule: Misc ##########
 
  @type_checked

src/tir/ir/index_map.cc

@@ -27,6 +27,7 @@
 #include <tvm/arith/int_set.h>
 #include <tvm/arith/iter_affine_map.h>
 #include <tvm/tir/op.h>
+#include <tvm/tir/stmt_functor.h>
 
 #include <sstream>
 
@@ -40,6 +41,15 @@ IndexMap::IndexMap(Array<Var> initial_indices, Array<PrimExpr> final_indices) {
  data_ = std::move(n);
 }
 
+IndexMap IndexMap::FromFunc(int ndim, runtime::TypedPackedFunc<Array<PrimExpr>(Array<Var>)> func) {
+ Array<Var> initial_indices;
+ initial_indices.reserve(ndim);
+ for (int i = 0; i < ndim; ++i) {
+ initial_indices.push_back(Var("i" + std::to_string(i), DataType::Int(32)));
+ }
+ return IndexMap(initial_indices, func(initial_indices));
+}
+
 IndexMap IndexMap::Inverse(Array<Range> initial_ranges) const {
  // Dummy variables to represent the inverse's inputs.
  Array<Var> output_vars;
@@ -142,13 +152,52 @@ Array<PrimExpr> IndexMapNode::MapShape(const Array<PrimExpr>& shape) const {
  return output;
 }
 
+String IndexMapNode::ToPythonString() const {
+ std::unordered_set<std::string> used_names;
+ Map<Var, PrimExpr> var_remap;
+ for (const Var& initial_index : initial_indices) {
+ if (used_names.count(initial_index->name_hint)) {
+ std::string new_name = initial_index->name_hint + std::to_string(used_names.size());
+ used_names.insert(new_name);
+ var_remap.Set(initial_index, Var(new_name));
+ } else {
+ used_names.insert(initial_index->name_hint);
+ }
+ }
+ std::ostringstream oss;
+ oss << "lambda ";
+ for (size_t i = 0; i < initial_indices.size(); ++i) {
+ if (i != 0) {
+ oss << ", ";
+ }
+ auto it = var_remap.find(initial_indices[i]);
+ if (it != var_remap.end()) {
+ oss << (*it).second;
+ } else {
+ oss << initial_indices[i];
+ }
+ }
+ oss << ": (";
+ for (size_t i = 0; i < final_indices.size(); ++i) {
+ oss << Substitute(final_indices[i], var_remap);
+ oss << ", ";
+ }
+ oss << ")";
+ return String(oss.str());
+}
+
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
  .set_dispatch<IndexMapNode>([](const ObjectRef& node, ReprPrinter* p) {
  auto* op = static_cast<const IndexMapNode*>(node.get());
- p->stream << "index_map(" << op->initial_indices << ", " << op->final_indices << ")";
+ p->stream << "index_map(" << op->ToPythonString() << ")";
  });
 
 TVM_REGISTER_NODE_TYPE(IndexMapNode);
 
+TVM_REGISTER_GLOBAL("tir.IndexMap")
+ .set_body_typed([](Array<Var> initial_indices, Array<PrimExpr> final_indices) {
+ return IndexMap(initial_indices, final_indices);
+ });
+
 } // namespace tir
 } // namespace tvm

src/tir/schedule/analysis.h

@@ -401,6 +401,16 @@ struct ProducerConsumerSplit {
  */
 Buffer GetNthAccessBuffer(const ScheduleState& self, const Block& block, int n, bool is_write);
 
+/*!
+ * \brief Find the defining site of the buffer in the given block and its ancestors
+ * \param block_sref The block sref
+ * \param buffer The buffer
+ * \return The defining site of the buffer and whether the buffer is allocated (otherwise the
+ * buffer is from match_buffer).
+ */
+std::pair<Optional<StmtSRef>, bool> GetBufferDefiningSite(const StmtSRef& block_sref,
+ const Buffer& buffer);
+
 /******** Reduction Block Related ********/
 
 /*!

include/tvm/tir/index_map.h

@@ -106,18 +106,34 @@ class IndexMapNode : public Object {
  */
  Array<PrimExpr> MapShape(const Array<PrimExpr>& shape) const;
 
+ /*!
+ * \brief Convert to string representation in Python.
+ * \return The stringified lambda expression in Python.
+ */
+ String ToPythonString() const;
+
  void VisitAttrs(AttrVisitor* v) {
  v->Visit("initial_indices", &initial_indices);
  v->Visit("final_indices", &final_indices);
  }
 
+ static constexpr const char* _type_key = "tir.IndexMap";
+
  TVM_DECLARE_FINAL_OBJECT_INFO(IndexMapNode, Object);
 };
 
 class IndexMap : public ObjectRef {
  public:
  IndexMap(Array<Var> initial_indices, Array<PrimExpr> final_indices);
 
+ /*!
+ * \brief Create an index map from a packed function
+ * \param ndim The number of dimensions
+ * \param func The function to be applied
+ * \return The created index map
+ */
+ static IndexMap FromFunc(int ndim, runtime::TypedPackedFunc<Array<PrimExpr>(Array<Var>)> func);
+
  /*! \brief Generate the inverse mapping.
  *
  * The range of the input indices is required in order to ensure

include/tvm/tir/schedule/schedule.h

@@ -20,6 +20,7 @@
 #define TVM_TIR_SCHEDULE_SCHEDULE_H_
 
 #include <tvm/support/random_engine.h>
+#include <tvm/tir/index_map.h>
 #include <tvm/tir/schedule/state.h>
 #include <tvm/tir/schedule/trace.h>
 
@@ -36,6 +37,14 @@ enum class ScheduleErrorRenderLevel : int32_t {
  kNone = 2,
 };
 
+/*! \brief Type of buffer index */
+enum class BufferIndexType : int32_t {
+ /*! \brief Index of a read buffer */
+ kRead = 0,
+ /*! \brief Index of a written buffer */
+ kWrite = 1,
+};
+
 /**************** Random variable: BlockRV ****************/
 
 /*! \brief A random variable that evaluates to a TensorIR block */
@@ -521,6 +530,21 @@ class ScheduleNode : public runtime::Object {
  */
  virtual void Unannotate(const BlockRV& block_rv, const String& ann_key) = 0;
 
+ /******** Schedule: Layout transformation ********/
+ /*!
+ * \brief Apply a transformation represented by IndexMap to buffer
+ * \details The indices and the access region to the target buffer is transformed by the given
+ * index_map. The index_map is used to infer the new shape of the buffer. Buffer must be either
+ * a function parameter, or allocated in a block (it cannot be a buffer subregion created via
+ * 'match_buffer').
+ * \param block_rv The block that accesses the target buffer.
+ * \param buffer_index The index of the buffer in block's read or write region.
+ * \param buffer_index_type The type of the buffer index, kRead or kWrite.
+ * \param index_map The transformation to apply.
+ */
+ virtual void TransformLayout(const BlockRV& block_rv, int buffer_index,
+ BufferIndexType buffer_index_type, const IndexMap& index_map) = 0;
+
  /******** Schedule: Misc ********/
  /*! \brief A no-op that marks the start of postprocessing phase of scheduling */
  virtual void EnterPostproc() = 0;

python/tvm/tir/__init__.py

@@ -57,7 +57,7 @@
 from .op import comm_reducer, min, max, sum
 from .op import q_multiply_shift
 
-from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError
+from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError, BufferType
 
 from . import schedule
 from . import ir_builder

python/tvm/tir/function.py

@@ -16,7 +16,8 @@
 # under the License.
 """Function data types."""
 
-from typing import Mapping, Union
+from typing import Callable, List, Mapping, Optional, Union
+import inspect
 
 import tvm._ffi
 import tvm.runtime
@@ -239,3 +240,58 @@ def get(name: str):
  The TensorIntrin with the specified name.
  """
  return _ffi_api.TensorIntrinGet(name) # pylint: type: ignore
+
+
+@tvm._ffi.register_object("tir.IndexMap")
+class IndexMap(Object):
+ """A mapping from multi-dimensional indices to another set of multi-dimensional indices
+
+ Parameters
+ ----------
+ initial_indices : List[Var]
+ Variables representing the indices prior to remapping.
+ final_indices : List[PrimExpr]
+ Expressions defining the indices after remapping.
+ """
+
+ initial_indices: List[Var]
+ final_indices: List[PrimExpr]
+
+ def __init__(self, initial_indices, final_indices):
+ self.__init_handle_by_constructor__(_ffi_api.IndexMap, initial_indices, final_indices)
+
+ @staticmethod
+ def from_func(mapping_function: Callable, ndim: Optional[int] = None):
+ """Create an index map from a function
+
+ Parameters
+ ----------
+ mapping_function : Callable
+ The function to map from source indices to target indices
+ """
+ params = inspect.signature(mapping_function).parameters
+ default_index_dtype = "int32"
+ args = []
+ var_arg_name = None
+ for name, param in params.items():
+ if param.kind in [
+ inspect.Parameter.POSITIONAL_ONLY,
+ inspect.Parameter.POSITIONAL_OR_KEYWORD,
+ ]:
+ args.append(tvm.tir.Var(name, default_index_dtype))
+ elif param.kind == inspect.Parameter.VAR_POSITIONAL:
+ var_arg_name = name
+ else:
+ raise ValueError("transform_layout mapping may not have *args or **kwargs")
+
+ # Now that all the named arguments have been collected,
+ # everything that remains should go to the *args, if
+ # specified.
+ if var_arg_name is not None:
+ assert ndim is not None, "ndim must be specified when *args is used"
+ num_var_args = ndim - len(args)
+ for i in range(num_var_args):
+ args.append(tvm.tir.Var(f"{var_arg_name}_{i}", default_index_dtype))
+
+ final_indices = mapping_function(*args)
+ return IndexMap(args, final_indices)

python/tvm/tir/schedule/__init__.py

@@ -19,6 +19,6 @@
 
 from .block_scope import BlockScope, Dependency, DepKind, StmtSRef
 from .instruction import Instruction, InstructionKind
-from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError
+from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError, BufferType
 from .state import ScheduleDebugMask, ScheduleState
 from .trace import Trace

python/tvm/tir/schedule/schedule.py

@@ -15,13 +15,15 @@
 # specific language governing permissions and limitations
 # under the License.
 """The TensorIR schedule class"""
-from typing import Dict, List, Optional, Union
+import enum
+from typing import Callable, Dict, List, Optional, Union
 
 from tvm._ffi import register_object as _register_object
 from tvm.error import TVMError, register_error
 from tvm.ir import IRModule, PrimExpr
 from tvm.runtime import Object, String
 from tvm.tir import Block, FloatImm, For, IntImm, PrimFunc
+from ..function import IndexMap
 
 from . import _ffi_api
 from .state import ScheduleState, StmtSRef, _parse_debug_mask, _parse_mod
@@ -71,6 +73,13 @@ def __init__(self) -> None:
 }
 
 
+class BufferType(enum.IntEnum):
+ """Type of buffer in access regions of a block"""
+
+ READ = 0
+ WRITE = 1
+
+
 def _parse_error_render_level(error_render_level: str) -> int:
  if error_render_level not in _ERROR_RENDER_LEVEL:
  raise ValueError(
@@ -2111,6 +2120,82 @@ def after_unannotate(a: T.handle, b: T.handle) -> None:
  self, block_or_loop, ann_key
  )
 
+ ########## Schedule: Layout transformation ##########
+
+ @type_checked
+ def transform_layout(
+ self,
+ block: BlockRV,
+ buffer_index: int,
+ buffer_type: BufferType,
+ index_map: Union[IndexMap, Callable],
+ ) -> None:
+ """Apply a transformation represented by IndexMap to buffer
+ Parameters
+ ----------
+ block_rv : BlockRV
+ The block that accesses the target buffer
+ buffer_index: int
+ The index of the buffer in block's read or write region
+ buffer_type : BufferType
+ Type of the buffer, READ or WRITE.
+ index_map : Union[IndexMap, Callable]
+ The transformation to apply
+
+ Examples
+ --------
+ Before transform_layout, in TensorIR, the IR is:
+
+ .. code-block:: python
+
+ @T.prim_func
+ def before_transform_layout(a: T.handle, c: T.handle) -> None:
+ A = T.match_buffer(a, (128, 128), "float32")
+ B = T.alloc_buffer((128, 128), "float32")
+ C = T.match_buffer(c, (128, 128), "float32")
+ for i, j in T.grid(128, 128):
+ with T.block("B"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ B[vi, vj] = A[vi, vj] * 2.0
+ for i, j in T.grid(128, 128):
+ with T.block("C"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ C[vi, vj] = B[vi, vj] + 1.0
+
+ Create the schedule and do transform_layout:
+
+ .. code-block:: python
+
+ sch = tir.Schedule(before_storage_align)
+ sch.transform_layout(sch.get_block("B"), buffer_index=0, BufferType.WRITE,
+ index_map=lambda m, n: (m // 16, n // 16, m % 16, n % 16))
+ print(sch.mod["main"].script())
+
+ After applying transform_layout, the IR becomes:
+
+ .. code-block:: python
+
+ @T.prim_func
+ def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> None:
+ A = T.match_buffer(a, (128, 128), "float32")
+ B = T.alloc_buffer((8, 8, 16, 16), "float32")
+ C = T.match_buffer(c, (128, 128), "float32")
+ for i, j in T.grid(128, 128):
+ with T.block("B"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ B[vi // 16, vj // 16, vi % 16, vj % 16] = A[vi, vj] * 2.0
+ for i, j in T.grid(128, 128):
+ with T.block("C"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ C[vi, vj] = B[vi // 16, vj // 16, vi % 16, vj % 16] + 1.0
+
+ """
+ if callable(index_map):
+ index_map = IndexMap.from_func(index_map)
+ _ffi_api.ScheduleTransformLayout( # type: ignore # pylint: disable=no-member
+ self, block, buffer_index, buffer_type, index_map
+ )
+
  ########## Schedule: Misc ##########
 
  @type_checked

src/tir/ir/index_map.cc

@@ -27,6 +27,7 @@
 #include <tvm/arith/int_set.h>
 #include <tvm/arith/iter_affine_map.h>
 #include <tvm/tir/op.h>
+#include <tvm/tir/stmt_functor.h>
 
 #include <sstream>
 
@@ -40,6 +41,15 @@ IndexMap::IndexMap(Array<Var> initial_indices, Array<PrimExpr> final_indices) {
  data_ = std::move(n);
 }
 
+IndexMap IndexMap::FromFunc(int ndim, runtime::TypedPackedFunc<Array<PrimExpr>(Array<Var>)> func) {
+ Array<Var> initial_indices;
+ initial_indices.reserve(ndim);
+ for (int i = 0; i < ndim; ++i) {
+ initial_indices.push_back(Var("i" + std::to_string(i), DataType::Int(32)));
+ }
+ return IndexMap(initial_indices, func(initial_indices));
+}
+
 IndexMap IndexMap::Inverse(Array<Range> initial_ranges) const {
  // Dummy variables to represent the inverse's inputs.
  Array<Var> output_vars;
@@ -142,13 +152,52 @@ Array<PrimExpr> IndexMapNode::MapShape(const Array<PrimExpr>& shape) const {
  return output;
 }
 
+String IndexMapNode::ToPythonString() const {
+ std::unordered_set<std::string> used_names;
+ Map<Var, PrimExpr> var_remap;
+ for (const Var& initial_index : initial_indices) {
+ if (used_names.count(initial_index->name_hint)) {
+ std::string new_name = initial_index->name_hint + std::to_string(used_names.size());
+ used_names.insert(new_name);
+ var_remap.Set(initial_index, Var(new_name));
+ } else {
+ used_names.insert(initial_index->name_hint);
+ }
+ }
+ std::ostringstream oss;
+ oss << "lambda ";
+ for (size_t i = 0; i < initial_indices.size(); ++i) {
+ if (i != 0) {
+ oss << ", ";
+ }
+ auto it = var_remap.find(initial_indices[i]);
+ if (it != var_remap.end()) {
+ oss << (*it).second;
+ } else {
+ oss << initial_indices[i];
+ }
+ }
+ oss << ": (";
+ for (size_t i = 0; i < final_indices.size(); ++i) {
+ oss << Substitute(final_indices[i], var_remap);
+ oss << ", ";
+ }
+ oss << ")";
+ return String(oss.str());
+}
+
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
  .set_dispatch<IndexMapNode>([](const ObjectRef& node, ReprPrinter* p) {
  auto* op = static_cast<const IndexMapNode*>(node.get());
- p->stream << "index_map(" << op->initial_indices << ", " << op->final_indices << ")";
+ p->stream << "index_map(" << op->ToPythonString() << ")";
  });
 
 TVM_REGISTER_NODE_TYPE(IndexMapNode);
 
+TVM_REGISTER_GLOBAL("tir.IndexMap")
+ .set_body_typed([](Array<Var> initial_indices, Array<PrimExpr> final_indices) {
+ return IndexMap(initial_indices, final_indices);
+ });
+
 } // namespace tir
 } // namespace tvm

src/tir/schedule/analysis.h

@@ -401,6 +401,16 @@ struct ProducerConsumerSplit {
  */
 Buffer GetNthAccessBuffer(const ScheduleState& self, const Block& block, int n, bool is_write);
 
+/*!
+ * \brief Find the defining site of the buffer in the given block and its ancestors
+ * \param block_sref The block sref
+ * \param buffer The buffer
+ * \return The defining site of the buffer and whether the buffer is allocated (otherwise the
+ * buffer is from match_buffer).
+ */
+std::pair<Optional<StmtSRef>, bool> GetBufferDefiningSite(const StmtSRef& block_sref,
+ const Buffer& buffer);
+
 /******** Reduction Block Related ********/
 
 /*!

include/tvm/tir/index_map.h

@@ -106,18 +106,34 @@ class IndexMapNode : public Object {
  */
  Array<PrimExpr> MapShape(const Array<PrimExpr>& shape) const;
 
+ /*!
+ * \brief Convert to string representation in Python.
+ * \return The stringified lambda expression in Python.
+ */
+ String ToPythonString() const;
+
  void VisitAttrs(AttrVisitor* v) {
  v->Visit("initial_indices", &initial_indices);
  v->Visit("final_indices", &final_indices);
  }
 
+ static constexpr const char* _type_key = "tir.IndexMap";
+
  TVM_DECLARE_FINAL_OBJECT_INFO(IndexMapNode, Object);
 };
 
 class IndexMap : public ObjectRef {
  public:
  IndexMap(Array<Var> initial_indices, Array<PrimExpr> final_indices);
 
+ /*!
+ * \brief Create an index map from a packed function
+ * \param ndim The number of dimensions
+ * \param func The function to be applied
+ * \return The created index map
+ */
+ static IndexMap FromFunc(int ndim, runtime::TypedPackedFunc<Array<PrimExpr>(Array<Var>)> func);
+
  /*! \brief Generate the inverse mapping.
  *
  * The range of the input indices is required in order to ensure

include/tvm/tir/schedule/schedule.h

@@ -20,6 +20,7 @@
 #define TVM_TIR_SCHEDULE_SCHEDULE_H_
 
 #include <tvm/support/random_engine.h>
+#include <tvm/tir/index_map.h>
 #include <tvm/tir/schedule/state.h>
 #include <tvm/tir/schedule/trace.h>
 
@@ -36,6 +37,14 @@ enum class ScheduleErrorRenderLevel : int32_t {
  kNone = 2,
 };
 
+/*! \brief Type of buffer index */
+enum class BufferIndexType : int32_t {
+ /*! \brief Index of a read buffer */
+ kRead = 0,
+ /*! \brief Index of a written buffer */
+ kWrite = 1,
+};
+
 /**************** Random variable: BlockRV ****************/
 
 /*! \brief A random variable that evaluates to a TensorIR block */
@@ -521,6 +530,21 @@ class ScheduleNode : public runtime::Object {
  */
  virtual void Unannotate(const BlockRV& block_rv, const String& ann_key) = 0;
 
+ /******** Schedule: Layout transformation ********/
+ /*!
+ * \brief Apply a transformation represented by IndexMap to buffer
+ * \details The indices and the access region to the target buffer is transformed by the given
+ * index_map. The index_map is used to infer the new shape of the buffer. Buffer must be either
+ * a function parameter, or allocated in a block (it cannot be a buffer subregion created via
+ * 'match_buffer').
+ * \param block_rv The block that accesses the target buffer.
+ * \param buffer_index The index of the buffer in block's read or write region.
+ * \param buffer_index_type The type of the buffer index, kRead or kWrite.
+ * \param index_map The transformation to apply.
+ */
+ virtual void TransformLayout(const BlockRV& block_rv, int buffer_index,
+ BufferIndexType buffer_index_type, const IndexMap& index_map) = 0;
+
  /******** Schedule: Misc ********/
  /*! \brief A no-op that marks the start of postprocessing phase of scheduling */
  virtual void EnterPostproc() = 0;

python/tvm/tir/__init__.py

@@ -57,7 +57,7 @@
 from .op import comm_reducer, min, max, sum
 from .op import q_multiply_shift
 
-from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError
+from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError, BufferType
 
 from . import schedule
 from . import ir_builder

python/tvm/tir/function.py

@@ -16,7 +16,8 @@
 # under the License.
 """Function data types."""
 
-from typing import Mapping, Union
+from typing import Callable, List, Mapping, Optional, Union
+import inspect
 
 import tvm._ffi
 import tvm.runtime
@@ -239,3 +240,58 @@ def get(name: str):
  The TensorIntrin with the specified name.
  """
  return _ffi_api.TensorIntrinGet(name) # pylint: type: ignore
+
+
+@tvm._ffi.register_object("tir.IndexMap")
+class IndexMap(Object):
+ """A mapping from multi-dimensional indices to another set of multi-dimensional indices
+
+ Parameters
+ ----------
+ initial_indices : List[Var]
+ Variables representing the indices prior to remapping.
+ final_indices : List[PrimExpr]
+ Expressions defining the indices after remapping.
+ """
+
+ initial_indices: List[Var]
+ final_indices: List[PrimExpr]
+
+ def __init__(self, initial_indices, final_indices):
+ self.__init_handle_by_constructor__(_ffi_api.IndexMap, initial_indices, final_indices)
+
+ @staticmethod
+ def from_func(mapping_function: Callable, ndim: Optional[int] = None):
+ """Create an index map from a function
+
+ Parameters
+ ----------
+ mapping_function : Callable
+ The function to map from source indices to target indices
+ """
+ params = inspect.signature(mapping_function).parameters
+ default_index_dtype = "int32"
+ args = []
+ var_arg_name = None
+ for name, param in params.items():
+ if param.kind in [
+ inspect.Parameter.POSITIONAL_ONLY,
+ inspect.Parameter.POSITIONAL_OR_KEYWORD,
+ ]:
+ args.append(tvm.tir.Var(name, default_index_dtype))
+ elif param.kind == inspect.Parameter.VAR_POSITIONAL:
+ var_arg_name = name
+ else:
+ raise ValueError("transform_layout mapping may not have *args or **kwargs")
+
+ # Now that all the named arguments have been collected,
+ # everything that remains should go to the *args, if
+ # specified.
+ if var_arg_name is not None:
+ assert ndim is not None, "ndim must be specified when *args is used"
+ num_var_args = ndim - len(args)
+ for i in range(num_var_args):
+ args.append(tvm.tir.Var(f"{var_arg_name}_{i}", default_index_dtype))
+
+ final_indices = mapping_function(*args)
+ return IndexMap(args, final_indices)

python/tvm/tir/schedule/__init__.py

@@ -19,6 +19,6 @@
 
 from .block_scope import BlockScope, Dependency, DepKind, StmtSRef
 from .instruction import Instruction, InstructionKind
-from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError
+from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError, BufferType
 from .state import ScheduleDebugMask, ScheduleState
 from .trace import Trace

python/tvm/tir/schedule/schedule.py

@@ -15,13 +15,15 @@
 # specific language governing permissions and limitations
 # under the License.
 """The TensorIR schedule class"""
-from typing import Dict, List, Optional, Union
+import enum
+from typing import Callable, Dict, List, Optional, Union
 
 from tvm._ffi import register_object as _register_object
 from tvm.error import TVMError, register_error
 from tvm.ir import IRModule, PrimExpr
 from tvm.runtime import Object, String
 from tvm.tir import Block, FloatImm, For, IntImm, PrimFunc
+from ..function import IndexMap
 
 from . import _ffi_api
 from .state import ScheduleState, StmtSRef, _parse_debug_mask, _parse_mod
@@ -71,6 +73,13 @@ def __init__(self) -> None:
 }
 
 
+class BufferType(enum.IntEnum):
+ """Type of buffer in access regions of a block"""
+
+ READ = 0
+ WRITE = 1
+
+
 def _parse_error_render_level(error_render_level: str) -> int:
  if error_render_level not in _ERROR_RENDER_LEVEL:
  raise ValueError(
@@ -2111,6 +2120,82 @@ def after_unannotate(a: T.handle, b: T.handle) -> None:
  self, block_or_loop, ann_key
  )
 
+ ########## Schedule: Layout transformation ##########
+
+ @type_checked
+ def transform_layout(
+ self,
+ block: BlockRV,
+ buffer_index: int,
+ buffer_type: BufferType,
+ index_map: Union[IndexMap, Callable],
+ ) -> None:
+ """Apply a transformation represented by IndexMap to buffer
+ Parameters
+ ----------
+ block_rv : BlockRV
+ The block that accesses the target buffer
+ buffer_index: int
+ The index of the buffer in block's read or write region
+ buffer_type : BufferType
+ Type of the buffer, READ or WRITE.
+ index_map : Union[IndexMap, Callable]
+ The transformation to apply
+
+ Examples
+ --------
+ Before transform_layout, in TensorIR, the IR is:
+
+ .. code-block:: python
+
+ @T.prim_func
+ def before_transform_layout(a: T.handle, c: T.handle) -> None:
+ A = T.match_buffer(a, (128, 128), "float32")
+ B = T.alloc_buffer((128, 128), "float32")
+ C = T.match_buffer(c, (128, 128), "float32")
+ for i, j in T.grid(128, 128):
+ with T.block("B"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ B[vi, vj] = A[vi, vj] * 2.0
+ for i, j in T.grid(128, 128):
+ with T.block("C"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ C[vi, vj] = B[vi, vj] + 1.0
+
+ Create the schedule and do transform_layout:
+
+ .. code-block:: python
+
+ sch = tir.Schedule(before_storage_align)
+ sch.transform_layout(sch.get_block("B"), buffer_index=0, BufferType.WRITE,
+ index_map=lambda m, n: (m // 16, n // 16, m % 16, n % 16))
+ print(sch.mod["main"].script())
+
+ After applying transform_layout, the IR becomes:
+
+ .. code-block:: python
+
+ @T.prim_func
+ def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> None:
+ A = T.match_buffer(a, (128, 128), "float32")
+ B = T.alloc_buffer((8, 8, 16, 16), "float32")
+ C = T.match_buffer(c, (128, 128), "float32")
+ for i, j in T.grid(128, 128):
+ with T.block("B"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ B[vi // 16, vj // 16, vi % 16, vj % 16] = A[vi, vj] * 2.0
+ for i, j in T.grid(128, 128):
+ with T.block("C"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ C[vi, vj] = B[vi // 16, vj // 16, vi % 16, vj % 16] + 1.0
+
+ """
+ if callable(index_map):
+ index_map = IndexMap.from_func(index_map)
+ _ffi_api.ScheduleTransformLayout( # type: ignore # pylint: disable=no-member
+ self, block, buffer_index, buffer_type, index_map
+ )
+
  ########## Schedule: Misc ##########
 
  @type_checked

src/tir/ir/index_map.cc

@@ -27,6 +27,7 @@
 #include <tvm/arith/int_set.h>
 #include <tvm/arith/iter_affine_map.h>
 #include <tvm/tir/op.h>
+#include <tvm/tir/stmt_functor.h>
 
 #include <sstream>
 
@@ -40,6 +41,15 @@ IndexMap::IndexMap(Array<Var> initial_indices, Array<PrimExpr> final_indices) {
  data_ = std::move(n);
 }
 
+IndexMap IndexMap::FromFunc(int ndim, runtime::TypedPackedFunc<Array<PrimExpr>(Array<Var>)> func) {
+ Array<Var> initial_indices;
+ initial_indices.reserve(ndim);
+ for (int i = 0; i < ndim; ++i) {
+ initial_indices.push_back(Var("i" + std::to_string(i), DataType::Int(32)));
+ }
+ return IndexMap(initial_indices, func(initial_indices));
+}
+
 IndexMap IndexMap::Inverse(Array<Range> initial_ranges) const {
  // Dummy variables to represent the inverse's inputs.
  Array<Var> output_vars;
@@ -142,13 +152,52 @@ Array<PrimExpr> IndexMapNode::MapShape(const Array<PrimExpr>& shape) const {
  return output;
 }
 
+String IndexMapNode::ToPythonString() const {
+ std::unordered_set<std::string> used_names;
+ Map<Var, PrimExpr> var_remap;
+ for (const Var& initial_index : initial_indices) {
+ if (used_names.count(initial_index->name_hint)) {
+ std::string new_name = initial_index->name_hint + std::to_string(used_names.size());
+ used_names.insert(new_name);
+ var_remap.Set(initial_index, Var(new_name));
+ } else {
+ used_names.insert(initial_index->name_hint);
+ }
+ }
+ std::ostringstream oss;
+ oss << "lambda ";
+ for (size_t i = 0; i < initial_indices.size(); ++i) {
+ if (i != 0) {
+ oss << ", ";
+ }
+ auto it = var_remap.find(initial_indices[i]);
+ if (it != var_remap.end()) {
+ oss << (*it).second;
+ } else {
+ oss << initial_indices[i];
+ }
+ }
+ oss << ": (";
+ for (size_t i = 0; i < final_indices.size(); ++i) {
+ oss << Substitute(final_indices[i], var_remap);
+ oss << ", ";
+ }
+ oss << ")";
+ return String(oss.str());
+}
+
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
  .set_dispatch<IndexMapNode>([](const ObjectRef& node, ReprPrinter* p) {
  auto* op = static_cast<const IndexMapNode*>(node.get());
- p->stream << "index_map(" << op->initial_indices << ", " << op->final_indices << ")";
+ p->stream << "index_map(" << op->ToPythonString() << ")";
  });
 
 TVM_REGISTER_NODE_TYPE(IndexMapNode);
 
+TVM_REGISTER_GLOBAL("tir.IndexMap")
+ .set_body_typed([](Array<Var> initial_indices, Array<PrimExpr> final_indices) {
+ return IndexMap(initial_indices, final_indices);
+ });
+
 } // namespace tir
 } // namespace tvm

src/tir/schedule/analysis.h

@@ -401,6 +401,16 @@ struct ProducerConsumerSplit {
  */
 Buffer GetNthAccessBuffer(const ScheduleState& self, const Block& block, int n, bool is_write);
 
+/*!
+ * \brief Find the defining site of the buffer in the given block and its ancestors
+ * \param block_sref The block sref
+ * \param buffer The buffer
+ * \return The defining site of the buffer and whether the buffer is allocated (otherwise the
+ * buffer is from match_buffer).
+ */
+std::pair<Optional<StmtSRef>, bool> GetBufferDefiningSite(const StmtSRef& block_sref,
+ const Buffer& buffer);
+
 /******** Reduction Block Related ********/
 
 /*!

include/tvm/tir/index_map.h

@@ -106,18 +106,34 @@ class IndexMapNode : public Object {
  */
  Array<PrimExpr> MapShape(const Array<PrimExpr>& shape) const;
 
+ /*!
+ * \brief Convert to string representation in Python.
+ * \return The stringified lambda expression in Python.
+ */
+ String ToPythonString() const;
+
  void VisitAttrs(AttrVisitor* v) {
  v->Visit("initial_indices", &initial_indices);
  v->Visit("final_indices", &final_indices);
  }
 
+ static constexpr const char* _type_key = "tir.IndexMap";
+
  TVM_DECLARE_FINAL_OBJECT_INFO(IndexMapNode, Object);
 };
 
 class IndexMap : public ObjectRef {
  public:
  IndexMap(Array<Var> initial_indices, Array<PrimExpr> final_indices);
 
+ /*!
+ * \brief Create an index map from a packed function
+ * \param ndim The number of dimensions
+ * \param func The function to be applied
+ * \return The created index map
+ */
+ static IndexMap FromFunc(int ndim, runtime::TypedPackedFunc<Array<PrimExpr>(Array<Var>)> func);
+
  /*! \brief Generate the inverse mapping.
  *
  * The range of the input indices is required in order to ensure

include/tvm/tir/schedule/schedule.h

@@ -20,6 +20,7 @@
 #define TVM_TIR_SCHEDULE_SCHEDULE_H_
 
 #include <tvm/support/random_engine.h>
+#include <tvm/tir/index_map.h>
 #include <tvm/tir/schedule/state.h>
 #include <tvm/tir/schedule/trace.h>
 
@@ -36,6 +37,14 @@ enum class ScheduleErrorRenderLevel : int32_t {
  kNone = 2,
 };
 
+/*! \brief Type of buffer index */
+enum class BufferIndexType : int32_t {
+ /*! \brief Index of a read buffer */
+ kRead = 0,
+ /*! \brief Index of a written buffer */
+ kWrite = 1,
+};
+
 /**************** Random variable: BlockRV ****************/
 
 /*! \brief A random variable that evaluates to a TensorIR block */
@@ -521,6 +530,21 @@ class ScheduleNode : public runtime::Object {
  */
  virtual void Unannotate(const BlockRV& block_rv, const String& ann_key) = 0;
 
+ /******** Schedule: Layout transformation ********/
+ /*!
+ * \brief Apply a transformation represented by IndexMap to buffer
+ * \details The indices and the access region to the target buffer is transformed by the given
+ * index_map. The index_map is used to infer the new shape of the buffer. Buffer must be either
+ * a function parameter, or allocated in a block (it cannot be a buffer subregion created via
+ * 'match_buffer').
+ * \param block_rv The block that accesses the target buffer.
+ * \param buffer_index The index of the buffer in block's read or write region.
+ * \param buffer_index_type The type of the buffer index, kRead or kWrite.
+ * \param index_map The transformation to apply.
+ */
+ virtual void TransformLayout(const BlockRV& block_rv, int buffer_index,
+ BufferIndexType buffer_index_type, const IndexMap& index_map) = 0;
+
  /******** Schedule: Misc ********/
  /*! \brief A no-op that marks the start of postprocessing phase of scheduling */
  virtual void EnterPostproc() = 0;

python/tvm/tir/__init__.py

@@ -57,7 +57,7 @@
 from .op import comm_reducer, min, max, sum
 from .op import q_multiply_shift
 
-from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError
+from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError, BufferType
 
 from . import schedule
 from . import ir_builder

python/tvm/tir/function.py

@@ -16,7 +16,8 @@
 # under the License.
 """Function data types."""
 
-from typing import Mapping, Union
+from typing import Callable, List, Mapping, Optional, Union
+import inspect
 
 import tvm._ffi
 import tvm.runtime
@@ -239,3 +240,58 @@ def get(name: str):
  The TensorIntrin with the specified name.
  """
  return _ffi_api.TensorIntrinGet(name) # pylint: type: ignore
+
+
+@tvm._ffi.register_object("tir.IndexMap")
+class IndexMap(Object):
+ """A mapping from multi-dimensional indices to another set of multi-dimensional indices
+
+ Parameters
+ ----------
+ initial_indices : List[Var]
+ Variables representing the indices prior to remapping.
+ final_indices : List[PrimExpr]
+ Expressions defining the indices after remapping.
+ """
+
+ initial_indices: List[Var]
+ final_indices: List[PrimExpr]
+
+ def __init__(self, initial_indices, final_indices):
+ self.__init_handle_by_constructor__(_ffi_api.IndexMap, initial_indices, final_indices)
+
+ @staticmethod
+ def from_func(mapping_function: Callable, ndim: Optional[int] = None):
+ """Create an index map from a function
+
+ Parameters
+ ----------
+ mapping_function : Callable
+ The function to map from source indices to target indices
+ """
+ params = inspect.signature(mapping_function).parameters
+ default_index_dtype = "int32"
+ args = []
+ var_arg_name = None
+ for name, param in params.items():
+ if param.kind in [
+ inspect.Parameter.POSITIONAL_ONLY,
+ inspect.Parameter.POSITIONAL_OR_KEYWORD,
+ ]:
+ args.append(tvm.tir.Var(name, default_index_dtype))
+ elif param.kind == inspect.Parameter.VAR_POSITIONAL:
+ var_arg_name = name
+ else:
+ raise ValueError("transform_layout mapping may not have *args or **kwargs")
+
+ # Now that all the named arguments have been collected,
+ # everything that remains should go to the *args, if
+ # specified.
+ if var_arg_name is not None:
+ assert ndim is not None, "ndim must be specified when *args is used"
+ num_var_args = ndim - len(args)
+ for i in range(num_var_args):
+ args.append(tvm.tir.Var(f"{var_arg_name}_{i}", default_index_dtype))
+
+ final_indices = mapping_function(*args)
+ return IndexMap(args, final_indices)

python/tvm/tir/schedule/__init__.py

@@ -19,6 +19,6 @@
 
 from .block_scope import BlockScope, Dependency, DepKind, StmtSRef
 from .instruction import Instruction, InstructionKind
-from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError
+from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError, BufferType
 from .state import ScheduleDebugMask, ScheduleState
 from .trace import Trace

python/tvm/tir/schedule/schedule.py

@@ -15,13 +15,15 @@
 # specific language governing permissions and limitations
 # under the License.
 """The TensorIR schedule class"""
-from typing import Dict, List, Optional, Union
+import enum
+from typing import Callable, Dict, List, Optional, Union
 
 from tvm._ffi import register_object as _register_object
 from tvm.error import TVMError, register_error
 from tvm.ir import IRModule, PrimExpr
 from tvm.runtime import Object, String
 from tvm.tir import Block, FloatImm, For, IntImm, PrimFunc
+from ..function import IndexMap
 
 from . import _ffi_api
 from .state import ScheduleState, StmtSRef, _parse_debug_mask, _parse_mod
@@ -71,6 +73,13 @@ def __init__(self) -> None:
 }
 
 
+class BufferType(enum.IntEnum):
+ """Type of buffer in access regions of a block"""
+
+ READ = 0
+ WRITE = 1
+
+
 def _parse_error_render_level(error_render_level: str) -> int:
  if error_render_level not in _ERROR_RENDER_LEVEL:
  raise ValueError(
@@ -2111,6 +2120,82 @@ def after_unannotate(a: T.handle, b: T.handle) -> None:
  self, block_or_loop, ann_key
  )
 
+ ########## Schedule: Layout transformation ##########
+
+ @type_checked
+ def transform_layout(
+ self,
+ block: BlockRV,
+ buffer_index: int,
+ buffer_type: BufferType,
+ index_map: Union[IndexMap, Callable],
+ ) -> None:
+ """Apply a transformation represented by IndexMap to buffer
+ Parameters
+ ----------
+ block_rv : BlockRV
+ The block that accesses the target buffer
+ buffer_index: int
+ The index of the buffer in block's read or write region
+ buffer_type : BufferType
+ Type of the buffer, READ or WRITE.
+ index_map : Union[IndexMap, Callable]
+ The transformation to apply
+
+ Examples
+ --------
+ Before transform_layout, in TensorIR, the IR is:
+
+ .. code-block:: python
+
+ @T.prim_func
+ def before_transform_layout(a: T.handle, c: T.handle) -> None:
+ A = T.match_buffer(a, (128, 128), "float32")
+ B = T.alloc_buffer((128, 128), "float32")
+ C = T.match_buffer(c, (128, 128), "float32")
+ for i, j in T.grid(128, 128):
+ with T.block("B"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ B[vi, vj] = A[vi, vj] * 2.0
+ for i, j in T.grid(128, 128):
+ with T.block("C"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ C[vi, vj] = B[vi, vj] + 1.0
+
+ Create the schedule and do transform_layout:
+
+ .. code-block:: python
+
+ sch = tir.Schedule(before_storage_align)
+ sch.transform_layout(sch.get_block("B"), buffer_index=0, BufferType.WRITE,
+ index_map=lambda m, n: (m // 16, n // 16, m % 16, n % 16))
+ print(sch.mod["main"].script())
+
+ After applying transform_layout, the IR becomes:
+
+ .. code-block:: python
+
+ @T.prim_func
+ def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> None:
+ A = T.match_buffer(a, (128, 128), "float32")
+ B = T.alloc_buffer((8, 8, 16, 16), "float32")
+ C = T.match_buffer(c, (128, 128), "float32")
+ for i, j in T.grid(128, 128):
+ with T.block("B"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ B[vi // 16, vj // 16, vi % 16, vj % 16] = A[vi, vj] * 2.0
+ for i, j in T.grid(128, 128):
+ with T.block("C"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ C[vi, vj] = B[vi // 16, vj // 16, vi % 16, vj % 16] + 1.0
+
+ """
+ if callable(index_map):
+ index_map = IndexMap.from_func(index_map)
+ _ffi_api.ScheduleTransformLayout( # type: ignore # pylint: disable=no-member
+ self, block, buffer_index, buffer_type, index_map
+ )
+
  ########## Schedule: Misc ##########
 
  @type_checked

src/tir/ir/index_map.cc

@@ -27,6 +27,7 @@
 #include <tvm/arith/int_set.h>
 #include <tvm/arith/iter_affine_map.h>
 #include <tvm/tir/op.h>
+#include <tvm/tir/stmt_functor.h>
 
 #include <sstream>
 
@@ -40,6 +41,15 @@ IndexMap::IndexMap(Array<Var> initial_indices, Array<PrimExpr> final_indices) {
  data_ = std::move(n);
 }
 
+IndexMap IndexMap::FromFunc(int ndim, runtime::TypedPackedFunc<Array<PrimExpr>(Array<Var>)> func) {
+ Array<Var> initial_indices;
+ initial_indices.reserve(ndim);
+ for (int i = 0; i < ndim; ++i) {
+ initial_indices.push_back(Var("i" + std::to_string(i), DataType::Int(32)));
+ }
+ return IndexMap(initial_indices, func(initial_indices));
+}
+
 IndexMap IndexMap::Inverse(Array<Range> initial_ranges) const {
  // Dummy variables to represent the inverse's inputs.
  Array<Var> output_vars;
@@ -142,13 +152,52 @@ Array<PrimExpr> IndexMapNode::MapShape(const Array<PrimExpr>& shape) const {
  return output;
 }
 
+String IndexMapNode::ToPythonString() const {
+ std::unordered_set<std::string> used_names;
+ Map<Var, PrimExpr> var_remap;
+ for (const Var& initial_index : initial_indices) {
+ if (used_names.count(initial_index->name_hint)) {
+ std::string new_name = initial_index->name_hint + std::to_string(used_names.size());
+ used_names.insert(new_name);
+ var_remap.Set(initial_index, Var(new_name));
+ } else {
+ used_names.insert(initial_index->name_hint);
+ }
+ }
+ std::ostringstream oss;
+ oss << "lambda ";
+ for (size_t i = 0; i < initial_indices.size(); ++i) {
+ if (i != 0) {
+ oss << ", ";
+ }
+ auto it = var_remap.find(initial_indices[i]);
+ if (it != var_remap.end()) {
+ oss << (*it).second;
+ } else {
+ oss << initial_indices[i];
+ }
+ }
+ oss << ": (";
+ for (size_t i = 0; i < final_indices.size(); ++i) {
+ oss << Substitute(final_indices[i], var_remap);
+ oss << ", ";
+ }
+ oss << ")";
+ return String(oss.str());
+}
+
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
  .set_dispatch<IndexMapNode>([](const ObjectRef& node, ReprPrinter* p) {
  auto* op = static_cast<const IndexMapNode*>(node.get());
- p->stream << "index_map(" << op->initial_indices << ", " << op->final_indices << ")";
+ p->stream << "index_map(" << op->ToPythonString() << ")";
  });
 
 TVM_REGISTER_NODE_TYPE(IndexMapNode);
 
+TVM_REGISTER_GLOBAL("tir.IndexMap")
+ .set_body_typed([](Array<Var> initial_indices, Array<PrimExpr> final_indices) {
+ return IndexMap(initial_indices, final_indices);
+ });
+
 } // namespace tir
 } // namespace tvm

src/tir/schedule/analysis.h

@@ -401,6 +401,16 @@ struct ProducerConsumerSplit {
  */
 Buffer GetNthAccessBuffer(const ScheduleState& self, const Block& block, int n, bool is_write);
 
+/*!
+ * \brief Find the defining site of the buffer in the given block and its ancestors
+ * \param block_sref The block sref
+ * \param buffer The buffer
+ * \return The defining site of the buffer and whether the buffer is allocated (otherwise the
+ * buffer is from match_buffer).
+ */
+std::pair<Optional<StmtSRef>, bool> GetBufferDefiningSite(const StmtSRef& block_sref,
+ const Buffer& buffer);
+
 /******** Reduction Block Related ********/
 
 /*!

include/tvm/tir/index_map.h

@@ -106,18 +106,34 @@ class IndexMapNode : public Object {
  */
  Array<PrimExpr> MapShape(const Array<PrimExpr>& shape) const;
 
+ /*!
+ * \brief Convert to string representation in Python.
+ * \return The stringified lambda expression in Python.
+ */
+ String ToPythonString() const;
+
  void VisitAttrs(AttrVisitor* v) {
  v->Visit("initial_indices", &initial_indices);
  v->Visit("final_indices", &final_indices);
  }
 
+ static constexpr const char* _type_key = "tir.IndexMap";
+
  TVM_DECLARE_FINAL_OBJECT_INFO(IndexMapNode, Object);
 };
 
 class IndexMap : public ObjectRef {
  public:
  IndexMap(Array<Var> initial_indices, Array<PrimExpr> final_indices);
 
+ /*!
+ * \brief Create an index map from a packed function
+ * \param ndim The number of dimensions
+ * \param func The function to be applied
+ * \return The created index map
+ */
+ static IndexMap FromFunc(int ndim, runtime::TypedPackedFunc<Array<PrimExpr>(Array<Var>)> func);
+
  /*! \brief Generate the inverse mapping.
  *
  * The range of the input indices is required in order to ensure

include/tvm/tir/schedule/schedule.h

@@ -20,6 +20,7 @@
 #define TVM_TIR_SCHEDULE_SCHEDULE_H_
 
 #include <tvm/support/random_engine.h>
+#include <tvm/tir/index_map.h>
 #include <tvm/tir/schedule/state.h>
 #include <tvm/tir/schedule/trace.h>
 
@@ -36,6 +37,14 @@ enum class ScheduleErrorRenderLevel : int32_t {
  kNone = 2,
 };
 
+/*! \brief Type of buffer index */
+enum class BufferIndexType : int32_t {
+ /*! \brief Index of a read buffer */
+ kRead = 0,
+ /*! \brief Index of a written buffer */
+ kWrite = 1,
+};
+
 /**************** Random variable: BlockRV ****************/
 
 /*! \brief A random variable that evaluates to a TensorIR block */
@@ -521,6 +530,21 @@ class ScheduleNode : public runtime::Object {
  */
  virtual void Unannotate(const BlockRV& block_rv, const String& ann_key) = 0;
 
+ /******** Schedule: Layout transformation ********/
+ /*!
+ * \brief Apply a transformation represented by IndexMap to buffer
+ * \details The indices and the access region to the target buffer is transformed by the given
+ * index_map. The index_map is used to infer the new shape of the buffer. Buffer must be either
+ * a function parameter, or allocated in a block (it cannot be a buffer subregion created via
+ * 'match_buffer').
+ * \param block_rv The block that accesses the target buffer.
+ * \param buffer_index The index of the buffer in block's read or write region.
+ * \param buffer_index_type The type of the buffer index, kRead or kWrite.
+ * \param index_map The transformation to apply.
+ */
+ virtual void TransformLayout(const BlockRV& block_rv, int buffer_index,
+ BufferIndexType buffer_index_type, const IndexMap& index_map) = 0;
+
  /******** Schedule: Misc ********/
  /*! \brief A no-op that marks the start of postprocessing phase of scheduling */
  virtual void EnterPostproc() = 0;

python/tvm/tir/__init__.py

@@ -57,7 +57,7 @@
 from .op import comm_reducer, min, max, sum
 from .op import q_multiply_shift
 
-from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError
+from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError, BufferType
 
 from . import schedule
 from . import ir_builder

python/tvm/tir/function.py

@@ -16,7 +16,8 @@
 # under the License.
 """Function data types."""
 
-from typing import Mapping, Union
+from typing import Callable, List, Mapping, Optional, Union
+import inspect
 
 import tvm._ffi
 import tvm.runtime
@@ -239,3 +240,58 @@ def get(name: str):
  The TensorIntrin with the specified name.
  """
  return _ffi_api.TensorIntrinGet(name) # pylint: type: ignore
+
+
+@tvm._ffi.register_object("tir.IndexMap")
+class IndexMap(Object):
+ """A mapping from multi-dimensional indices to another set of multi-dimensional indices
+
+ Parameters
+ ----------
+ initial_indices : List[Var]
+ Variables representing the indices prior to remapping.
+ final_indices : List[PrimExpr]
+ Expressions defining the indices after remapping.
+ """
+
+ initial_indices: List[Var]
+ final_indices: List[PrimExpr]
+
+ def __init__(self, initial_indices, final_indices):
+ self.__init_handle_by_constructor__(_ffi_api.IndexMap, initial_indices, final_indices)
+
+ @staticmethod
+ def from_func(mapping_function: Callable, ndim: Optional[int] = None):
+ """Create an index map from a function
+
+ Parameters
+ ----------
+ mapping_function : Callable
+ The function to map from source indices to target indices
+ """
+ params = inspect.signature(mapping_function).parameters
+ default_index_dtype = "int32"
+ args = []
+ var_arg_name = None
+ for name, param in params.items():
+ if param.kind in [
+ inspect.Parameter.POSITIONAL_ONLY,
+ inspect.Parameter.POSITIONAL_OR_KEYWORD,
+ ]:
+ args.append(tvm.tir.Var(name, default_index_dtype))
+ elif param.kind == inspect.Parameter.VAR_POSITIONAL:
+ var_arg_name = name
+ else:
+ raise ValueError("transform_layout mapping may not have *args or **kwargs")
+
+ # Now that all the named arguments have been collected,
+ # everything that remains should go to the *args, if
+ # specified.
+ if var_arg_name is not None:
+ assert ndim is not None, "ndim must be specified when *args is used"
+ num_var_args = ndim - len(args)
+ for i in range(num_var_args):
+ args.append(tvm.tir.Var(f"{var_arg_name}_{i}", default_index_dtype))
+
+ final_indices = mapping_function(*args)
+ return IndexMap(args, final_indices)

python/tvm/tir/schedule/__init__.py

@@ -19,6 +19,6 @@
 
 from .block_scope import BlockScope, Dependency, DepKind, StmtSRef
 from .instruction import Instruction, InstructionKind
-from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError
+from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError, BufferType
 from .state import ScheduleDebugMask, ScheduleState
 from .trace import Trace

python/tvm/tir/schedule/schedule.py

@@ -15,13 +15,15 @@
 # specific language governing permissions and limitations
 # under the License.
 """The TensorIR schedule class"""
-from typing import Dict, List, Optional, Union
+import enum
+from typing import Callable, Dict, List, Optional, Union
 
 from tvm._ffi import register_object as _register_object
 from tvm.error import TVMError, register_error
 from tvm.ir import IRModule, PrimExpr
 from tvm.runtime import Object, String
 from tvm.tir import Block, FloatImm, For, IntImm, PrimFunc
+from ..function import IndexMap
 
 from . import _ffi_api
 from .state import ScheduleState, StmtSRef, _parse_debug_mask, _parse_mod
@@ -71,6 +73,13 @@ def __init__(self) -> None:
 }
 
 
+class BufferType(enum.IntEnum):
+ """Type of buffer in access regions of a block"""
+
+ READ = 0
+ WRITE = 1
+
+
 def _parse_error_render_level(error_render_level: str) -> int:
  if error_render_level not in _ERROR_RENDER_LEVEL:
  raise ValueError(
@@ -2111,6 +2120,82 @@ def after_unannotate(a: T.handle, b: T.handle) -> None:
  self, block_or_loop, ann_key
  )
 
+ ########## Schedule: Layout transformation ##########
+
+ @type_checked
+ def transform_layout(
+ self,
+ block: BlockRV,
+ buffer_index: int,
+ buffer_type: BufferType,
+ index_map: Union[IndexMap, Callable],
+ ) -> None:
+ """Apply a transformation represented by IndexMap to buffer
+ Parameters
+ ----------
+ block_rv : BlockRV
+ The block that accesses the target buffer
+ buffer_index: int
+ The index of the buffer in block's read or write region
+ buffer_type : BufferType
+ Type of the buffer, READ or WRITE.
+ index_map : Union[IndexMap, Callable]
+ The transformation to apply
+
+ Examples
+ --------
+ Before transform_layout, in TensorIR, the IR is:
+
+ .. code-block:: python
+
+ @T.prim_func
+ def before_transform_layout(a: T.handle, c: T.handle) -> None:
+ A = T.match_buffer(a, (128, 128), "float32")
+ B = T.alloc_buffer((128, 128), "float32")
+ C = T.match_buffer(c, (128, 128), "float32")
+ for i, j in T.grid(128, 128):
+ with T.block("B"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ B[vi, vj] = A[vi, vj] * 2.0
+ for i, j in T.grid(128, 128):
+ with T.block("C"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ C[vi, vj] = B[vi, vj] + 1.0
+
+ Create the schedule and do transform_layout:
+
+ .. code-block:: python
+
+ sch = tir.Schedule(before_storage_align)
+ sch.transform_layout(sch.get_block("B"), buffer_index=0, BufferType.WRITE,
+ index_map=lambda m, n: (m // 16, n // 16, m % 16, n % 16))
+ print(sch.mod["main"].script())
+
+ After applying transform_layout, the IR becomes:
+
+ .. code-block:: python
+
+ @T.prim_func
+ def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> None:
+ A = T.match_buffer(a, (128, 128), "float32")
+ B = T.alloc_buffer((8, 8, 16, 16), "float32")
+ C = T.match_buffer(c, (128, 128), "float32")
+ for i, j in T.grid(128, 128):
+ with T.block("B"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ B[vi // 16, vj // 16, vi % 16, vj % 16] = A[vi, vj] * 2.0
+ for i, j in T.grid(128, 128):
+ with T.block("C"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ C[vi, vj] = B[vi // 16, vj // 16, vi % 16, vj % 16] + 1.0
+
+ """
+ if callable(index_map):
+ index_map = IndexMap.from_func(index_map)
+ _ffi_api.ScheduleTransformLayout( # type: ignore # pylint: disable=no-member
+ self, block, buffer_index, buffer_type, index_map
+ )
+
  ########## Schedule: Misc ##########
 
  @type_checked

src/tir/ir/index_map.cc

@@ -27,6 +27,7 @@
 #include <tvm/arith/int_set.h>
 #include <tvm/arith/iter_affine_map.h>
 #include <tvm/tir/op.h>
+#include <tvm/tir/stmt_functor.h>
 
 #include <sstream>
 
@@ -40,6 +41,15 @@ IndexMap::IndexMap(Array<Var> initial_indices, Array<PrimExpr> final_indices) {
  data_ = std::move(n);
 }
 
+IndexMap IndexMap::FromFunc(int ndim, runtime::TypedPackedFunc<Array<PrimExpr>(Array<Var>)> func) {
+ Array<Var> initial_indices;
+ initial_indices.reserve(ndim);
+ for (int i = 0; i < ndim; ++i) {
+ initial_indices.push_back(Var("i" + std::to_string(i), DataType::Int(32)));
+ }
+ return IndexMap(initial_indices, func(initial_indices));
+}
+
 IndexMap IndexMap::Inverse(Array<Range> initial_ranges) const {
  // Dummy variables to represent the inverse's inputs.
  Array<Var> output_vars;
@@ -142,13 +152,52 @@ Array<PrimExpr> IndexMapNode::MapShape(const Array<PrimExpr>& shape) const {
  return output;
 }
 
+String IndexMapNode::ToPythonString() const {
+ std::unordered_set<std::string> used_names;
+ Map<Var, PrimExpr> var_remap;
+ for (const Var& initial_index : initial_indices) {
+ if (used_names.count(initial_index->name_hint)) {
+ std::string new_name = initial_index->name_hint + std::to_string(used_names.size());
+ used_names.insert(new_name);
+ var_remap.Set(initial_index, Var(new_name));
+ } else {
+ used_names.insert(initial_index->name_hint);
+ }
+ }
+ std::ostringstream oss;
+ oss << "lambda ";
+ for (size_t i = 0; i < initial_indices.size(); ++i) {
+ if (i != 0) {
+ oss << ", ";
+ }
+ auto it = var_remap.find(initial_indices[i]);
+ if (it != var_remap.end()) {
+ oss << (*it).second;
+ } else {
+ oss << initial_indices[i];
+ }
+ }
+ oss << ": (";
+ for (size_t i = 0; i < final_indices.size(); ++i) {
+ oss << Substitute(final_indices[i], var_remap);
+ oss << ", ";
+ }
+ oss << ")";
+ return String(oss.str());
+}
+
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
  .set_dispatch<IndexMapNode>([](const ObjectRef& node, ReprPrinter* p) {
  auto* op = static_cast<const IndexMapNode*>(node.get());
- p->stream << "index_map(" << op->initial_indices << ", " << op->final_indices << ")";
+ p->stream << "index_map(" << op->ToPythonString() << ")";
  });
 
 TVM_REGISTER_NODE_TYPE(IndexMapNode);
 
+TVM_REGISTER_GLOBAL("tir.IndexMap")
+ .set_body_typed([](Array<Var> initial_indices, Array<PrimExpr> final_indices) {
+ return IndexMap(initial_indices, final_indices);
+ });
+
 } // namespace tir
 } // namespace tvm

src/tir/schedule/analysis.h

@@ -401,6 +401,16 @@ struct ProducerConsumerSplit {
  */
 Buffer GetNthAccessBuffer(const ScheduleState& self, const Block& block, int n, bool is_write);
 
+/*!
+ * \brief Find the defining site of the buffer in the given block and its ancestors
+ * \param block_sref The block sref
+ * \param buffer The buffer
+ * \return The defining site of the buffer and whether the buffer is allocated (otherwise the
+ * buffer is from match_buffer).
+ */
+std::pair<Optional<StmtSRef>, bool> GetBufferDefiningSite(const StmtSRef& block_sref,
+ const Buffer& buffer);
+
 /******** Reduction Block Related ********/
 
 /*!

include/tvm/tir/index_map.h

@@ -106,18 +106,34 @@ class IndexMapNode : public Object {
  */
  Array<PrimExpr> MapShape(const Array<PrimExpr>& shape) const;
 
+ /*!
+ * \brief Convert to string representation in Python.
+ * \return The stringified lambda expression in Python.
+ */
+ String ToPythonString() const;
+
  void VisitAttrs(AttrVisitor* v) {
  v->Visit("initial_indices", &initial_indices);
  v->Visit("final_indices", &final_indices);
  }
 
+ static constexpr const char* _type_key = "tir.IndexMap";
+
  TVM_DECLARE_FINAL_OBJECT_INFO(IndexMapNode, Object);
 };
 
 class IndexMap : public ObjectRef {
  public:
  IndexMap(Array<Var> initial_indices, Array<PrimExpr> final_indices);
 
+ /*!
+ * \brief Create an index map from a packed function
+ * \param ndim The number of dimensions
+ * \param func The function to be applied
+ * \return The created index map
+ */
+ static IndexMap FromFunc(int ndim, runtime::TypedPackedFunc<Array<PrimExpr>(Array<Var>)> func);
+
  /*! \brief Generate the inverse mapping.
  *
  * The range of the input indices is required in order to ensure

include/tvm/tir/schedule/schedule.h

@@ -20,6 +20,7 @@
 #define TVM_TIR_SCHEDULE_SCHEDULE_H_
 
 #include <tvm/support/random_engine.h>
+#include <tvm/tir/index_map.h>
 #include <tvm/tir/schedule/state.h>
 #include <tvm/tir/schedule/trace.h>
 
@@ -36,6 +37,14 @@ enum class ScheduleErrorRenderLevel : int32_t {
  kNone = 2,
 };
 
+/*! \brief Type of buffer index */
+enum class BufferIndexType : int32_t {
+ /*! \brief Index of a read buffer */
+ kRead = 0,
+ /*! \brief Index of a written buffer */
+ kWrite = 1,
+};
+
 /**************** Random variable: BlockRV ****************/
 
 /*! \brief A random variable that evaluates to a TensorIR block */
@@ -521,6 +530,21 @@ class ScheduleNode : public runtime::Object {
  */
  virtual void Unannotate(const BlockRV& block_rv, const String& ann_key) = 0;
 
+ /******** Schedule: Layout transformation ********/
+ /*!
+ * \brief Apply a transformation represented by IndexMap to buffer
+ * \details The indices and the access region to the target buffer is transformed by the given
+ * index_map. The index_map is used to infer the new shape of the buffer. Buffer must be either
+ * a function parameter, or allocated in a block (it cannot be a buffer subregion created via
+ * 'match_buffer').
+ * \param block_rv The block that accesses the target buffer.
+ * \param buffer_index The index of the buffer in block's read or write region.
+ * \param buffer_index_type The type of the buffer index, kRead or kWrite.
+ * \param index_map The transformation to apply.
+ */
+ virtual void TransformLayout(const BlockRV& block_rv, int buffer_index,
+ BufferIndexType buffer_index_type, const IndexMap& index_map) = 0;
+
  /******** Schedule: Misc ********/
  /*! \brief A no-op that marks the start of postprocessing phase of scheduling */
  virtual void EnterPostproc() = 0;

python/tvm/tir/__init__.py

@@ -57,7 +57,7 @@
 from .op import comm_reducer, min, max, sum
 from .op import q_multiply_shift
 
-from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError
+from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError, BufferType
 
 from . import schedule
 from . import ir_builder

python/tvm/tir/function.py

@@ -16,7 +16,8 @@
 # under the License.
 """Function data types."""
 
-from typing import Mapping, Union
+from typing import Callable, List, Mapping, Optional, Union
+import inspect
 
 import tvm._ffi
 import tvm.runtime
@@ -239,3 +240,58 @@ def get(name: str):
  The TensorIntrin with the specified name.
  """
  return _ffi_api.TensorIntrinGet(name) # pylint: type: ignore
+
+
+@tvm._ffi.register_object("tir.IndexMap")
+class IndexMap(Object):
+ """A mapping from multi-dimensional indices to another set of multi-dimensional indices
+
+ Parameters
+ ----------
+ initial_indices : List[Var]
+ Variables representing the indices prior to remapping.
+ final_indices : List[PrimExpr]
+ Expressions defining the indices after remapping.
+ """
+
+ initial_indices: List[Var]
+ final_indices: List[PrimExpr]
+
+ def __init__(self, initial_indices, final_indices):
+ self.__init_handle_by_constructor__(_ffi_api.IndexMap, initial_indices, final_indices)
+
+ @staticmethod
+ def from_func(mapping_function: Callable, ndim: Optional[int] = None):
+ """Create an index map from a function
+
+ Parameters
+ ----------
+ mapping_function : Callable
+ The function to map from source indices to target indices
+ """
+ params = inspect.signature(mapping_function).parameters
+ default_index_dtype = "int32"
+ args = []
+ var_arg_name = None
+ for name, param in params.items():
+ if param.kind in [
+ inspect.Parameter.POSITIONAL_ONLY,
+ inspect.Parameter.POSITIONAL_OR_KEYWORD,
+ ]:
+ args.append(tvm.tir.Var(name, default_index_dtype))
+ elif param.kind == inspect.Parameter.VAR_POSITIONAL:
+ var_arg_name = name
+ else:
+ raise ValueError("transform_layout mapping may not have *args or **kwargs")
+
+ # Now that all the named arguments have been collected,
+ # everything that remains should go to the *args, if
+ # specified.
+ if var_arg_name is not None:
+ assert ndim is not None, "ndim must be specified when *args is used"
+ num_var_args = ndim - len(args)
+ for i in range(num_var_args):
+ args.append(tvm.tir.Var(f"{var_arg_name}_{i}", default_index_dtype))
+
+ final_indices = mapping_function(*args)
+ return IndexMap(args, final_indices)

python/tvm/tir/schedule/__init__.py

@@ -19,6 +19,6 @@
 
 from .block_scope import BlockScope, Dependency, DepKind, StmtSRef
 from .instruction import Instruction, InstructionKind
-from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError
+from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError, BufferType
 from .state import ScheduleDebugMask, ScheduleState
 from .trace import Trace

python/tvm/tir/schedule/schedule.py

@@ -15,13 +15,15 @@
 # specific language governing permissions and limitations
 # under the License.
 """The TensorIR schedule class"""
-from typing import Dict, List, Optional, Union
+import enum
+from typing import Callable, Dict, List, Optional, Union
 
 from tvm._ffi import register_object as _register_object
 from tvm.error import TVMError, register_error
 from tvm.ir import IRModule, PrimExpr
 from tvm.runtime import Object, String
 from tvm.tir import Block, FloatImm, For, IntImm, PrimFunc
+from ..function import IndexMap
 
 from . import _ffi_api
 from .state import ScheduleState, StmtSRef, _parse_debug_mask, _parse_mod
@@ -71,6 +73,13 @@ def __init__(self) -> None:
 }
 
 
+class BufferType(enum.IntEnum):
+ """Type of buffer in access regions of a block"""
+
+ READ = 0
+ WRITE = 1
+
+
 def _parse_error_render_level(error_render_level: str) -> int:
  if error_render_level not in _ERROR_RENDER_LEVEL:
  raise ValueError(
@@ -2111,6 +2120,82 @@ def after_unannotate(a: T.handle, b: T.handle) -> None:
  self, block_or_loop, ann_key
  )
 
+ ########## Schedule: Layout transformation ##########
+
+ @type_checked
+ def transform_layout(
+ self,
+ block: BlockRV,
+ buffer_index: int,
+ buffer_type: BufferType,
+ index_map: Union[IndexMap, Callable],
+ ) -> None:
+ """Apply a transformation represented by IndexMap to buffer
+ Parameters
+ ----------
+ block_rv : BlockRV
+ The block that accesses the target buffer
+ buffer_index: int
+ The index of the buffer in block's read or write region
+ buffer_type : BufferType
+ Type of the buffer, READ or WRITE.
+ index_map : Union[IndexMap, Callable]
+ The transformation to apply
+
+ Examples
+ --------
+ Before transform_layout, in TensorIR, the IR is:
+
+ .. code-block:: python
+
+ @T.prim_func
+ def before_transform_layout(a: T.handle, c: T.handle) -> None:
+ A = T.match_buffer(a, (128, 128), "float32")
+ B = T.alloc_buffer((128, 128), "float32")
+ C = T.match_buffer(c, (128, 128), "float32")
+ for i, j in T.grid(128, 128):
+ with T.block("B"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ B[vi, vj] = A[vi, vj] * 2.0
+ for i, j in T.grid(128, 128):
+ with T.block("C"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ C[vi, vj] = B[vi, vj] + 1.0
+
+ Create the schedule and do transform_layout:
+
+ .. code-block:: python
+
+ sch = tir.Schedule(before_storage_align)
+ sch.transform_layout(sch.get_block("B"), buffer_index=0, BufferType.WRITE,
+ index_map=lambda m, n: (m // 16, n // 16, m % 16, n % 16))
+ print(sch.mod["main"].script())
+
+ After applying transform_layout, the IR becomes:
+
+ .. code-block:: python
+
+ @T.prim_func
+ def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> None:
+ A = T.match_buffer(a, (128, 128), "float32")
+ B = T.alloc_buffer((8, 8, 16, 16), "float32")
+ C = T.match_buffer(c, (128, 128), "float32")
+ for i, j in T.grid(128, 128):
+ with T.block("B"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ B[vi // 16, vj // 16, vi % 16, vj % 16] = A[vi, vj] * 2.0
+ for i, j in T.grid(128, 128):
+ with T.block("C"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ C[vi, vj] = B[vi // 16, vj // 16, vi % 16, vj % 16] + 1.0
+
+ """
+ if callable(index_map):
+ index_map = IndexMap.from_func(index_map)
+ _ffi_api.ScheduleTransformLayout( # type: ignore # pylint: disable=no-member
+ self, block, buffer_index, buffer_type, index_map
+ )
+
  ########## Schedule: Misc ##########
 
  @type_checked

src/tir/ir/index_map.cc

@@ -27,6 +27,7 @@
 #include <tvm/arith/int_set.h>
 #include <tvm/arith/iter_affine_map.h>
 #include <tvm/tir/op.h>
+#include <tvm/tir/stmt_functor.h>
 
 #include <sstream>
 
@@ -40,6 +41,15 @@ IndexMap::IndexMap(Array<Var> initial_indices, Array<PrimExpr> final_indices) {
  data_ = std::move(n);
 }
 
+IndexMap IndexMap::FromFunc(int ndim, runtime::TypedPackedFunc<Array<PrimExpr>(Array<Var>)> func) {
+ Array<Var> initial_indices;
+ initial_indices.reserve(ndim);
+ for (int i = 0; i < ndim; ++i) {
+ initial_indices.push_back(Var("i" + std::to_string(i), DataType::Int(32)));
+ }
+ return IndexMap(initial_indices, func(initial_indices));
+}
+
 IndexMap IndexMap::Inverse(Array<Range> initial_ranges) const {
  // Dummy variables to represent the inverse's inputs.
  Array<Var> output_vars;
@@ -142,13 +152,52 @@ Array<PrimExpr> IndexMapNode::MapShape(const Array<PrimExpr>& shape) const {
  return output;
 }
 
+String IndexMapNode::ToPythonString() const {
+ std::unordered_set<std::string> used_names;
+ Map<Var, PrimExpr> var_remap;
+ for (const Var& initial_index : initial_indices) {
+ if (used_names.count(initial_index->name_hint)) {
+ std::string new_name = initial_index->name_hint + std::to_string(used_names.size());
+ used_names.insert(new_name);
+ var_remap.Set(initial_index, Var(new_name));
+ } else {
+ used_names.insert(initial_index->name_hint);
+ }
+ }
+ std::ostringstream oss;
+ oss << "lambda ";
+ for (size_t i = 0; i < initial_indices.size(); ++i) {
+ if (i != 0) {
+ oss << ", ";
+ }
+ auto it = var_remap.find(initial_indices[i]);
+ if (it != var_remap.end()) {
+ oss << (*it).second;
+ } else {
+ oss << initial_indices[i];
+ }
+ }
+ oss << ": (";
+ for (size_t i = 0; i < final_indices.size(); ++i) {
+ oss << Substitute(final_indices[i], var_remap);
+ oss << ", ";
+ }
+ oss << ")";
+ return String(oss.str());
+}
+
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
  .set_dispatch<IndexMapNode>([](const ObjectRef& node, ReprPrinter* p) {
  auto* op = static_cast<const IndexMapNode*>(node.get());
- p->stream << "index_map(" << op->initial_indices << ", " << op->final_indices << ")";
+ p->stream << "index_map(" << op->ToPythonString() << ")";
  });
 
 TVM_REGISTER_NODE_TYPE(IndexMapNode);
 
+TVM_REGISTER_GLOBAL("tir.IndexMap")
+ .set_body_typed([](Array<Var> initial_indices, Array<PrimExpr> final_indices) {
+ return IndexMap(initial_indices, final_indices);
+ });
+
 } // namespace tir
 } // namespace tvm

src/tir/schedule/analysis.h

@@ -401,6 +401,16 @@ struct ProducerConsumerSplit {
  */
 Buffer GetNthAccessBuffer(const ScheduleState& self, const Block& block, int n, bool is_write);
 
+/*!
+ * \brief Find the defining site of the buffer in the given block and its ancestors
+ * \param block_sref The block sref
+ * \param buffer The buffer
+ * \return The defining site of the buffer and whether the buffer is allocated (otherwise the
+ * buffer is from match_buffer).
+ */
+std::pair<Optional<StmtSRef>, bool> GetBufferDefiningSite(const StmtSRef& block_sref,
+ const Buffer& buffer);
+
 /******** Reduction Block Related ********/
 
 /*!