[REFACTOR][TIR] Remove buffer type and axis separators

include/tvm/relax/attrs/manipulate.h

@@ -64,19 +64,6 @@ struct LayoutTransformAttrs : public AttrsNode {
   // pad_value is chosen to be of PrimExpr type, as it represents constant TIR POD expression. This
   // needs to be revisited in case PrimExpr is evolved to represent symbolic expression in future.
   ffi::Optional<PrimExpr> pad_value;
-  /*!
-   * axis_separators between input axes when generating flattened output axes. For buffers
-   * representing flat 1-d memory (e.g. any buffer in RAM), this should be an empty array.
-   * For buffers representing non-flat memory, each entry in axis_separators should be the
-   * first input axis that is part of a new flattened axis.
-   */
-  ffi::Optional<ffi::Array<IntImm>> axis_separators;
-  /*!
-   * axis_separators for input buffers.
-   * Needed to identify if the input buffer to layout_transform
-   * contains axis separator.
-   */
-  ffi::Optional<ffi::Array<IntImm>> input_axis_separators;
 
   static void RegisterReflection() {
     namespace refl = tvm::ffi::reflection;
@@ -86,11 +73,7 @@ struct LayoutTransformAttrs : public AttrsNode {
         .def_ro(
             "pad_value", &LayoutTransformAttrs::pad_value,
             "The specific value to be used to pad if the layout transform would result in implicit "
-            "padding. If not specified, the compiler is free to choose any value.")
-        .def_ro("axis_separators", &LayoutTransformAttrs::axis_separators,
-                "The separators between input axes when generating flat output axes")
-        .def_ro("input_axis_separators", &LayoutTransformAttrs::input_axis_separators,
-                "The separators between axes to regenerate output");
+            "padding. If not specified, the compiler is free to choose any value.");
   }
   TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LayoutTransformAttrs", LayoutTransformAttrs,
                                     AttrsNode);

include/tvm/relax/transform.h

@@ -595,16 +595,11 @@ TVM_DLL Pass DecomposeOpsForTraining(ffi::Optional<ffi::String> func_name);
  * \param op_impl_map Map from kOperatorName attr (e.g., relax.conv2d) to replacement PrimFunc
  * \param op_buffer_transforms Map from kOperatorName attr to layout transformations on each of the
  * PrimFunc i/o buffers.
- * \param axis_separators Map from kOperatorName attr to axis_separators of each buffer_transforms
- * \param input_axis_separators Map from kOperatorName attr to axis_separator for input buffer
  * \return The Pass.
  */
-TVM_DLL Pass AlterOpImpl(
-    const ffi::Map<ffi::String, tirx::PrimFunc>& op_impl_map,
-    const ffi::Map<ffi::String, ffi::Array<tirx::IndexMap>>& op_buffer_transforms,
-    const ffi::Map<ffi::String, ffi::Optional<ffi::Array<ffi::Array<IntImm>>>>& axis_separators,
-    const ffi::Map<ffi::String, ffi::Optional<ffi::Array<ffi::Array<IntImm>>>>&
-        input_axis_separators);
+TVM_DLL Pass
+AlterOpImpl(const ffi::Map<ffi::String, tirx::PrimFunc>& op_impl_map,
+            const ffi::Map<ffi::String, ffi::Array<tirx::IndexMap>>& op_buffer_transforms);
 
 /*!
  * \brief Layout conversion pass.

include/tvm/s_tir/schedule/schedule.h

@@ -796,18 +796,6 @@ class ScheduleNode : public ffi::Object {
    */
   virtual void TransformBlockLayout(const SBlockRV& block_rv, const IndexMap& index_map) = 0;
 
-  /*!
-   * \brief Set the axis separator of a buffer, where the buffer is specified by a block and a read
-   * or write index
-   * \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 axis_separators The axis separator of the buffer
-   */
-  virtual void SetAxisSeparator(const SBlockRV& block_rv, int buffer_index,
-                                BufferIndexType buffer_index_type,
-                                const ffi::Array<IntImm>& axis_separators) = 0;
-
   /******** Schedule: Padding ********/
   /*!
    * \brief Decompose a padding block into a block filling const pad values and a block

include/tvm/s_tir/stmt.h

@@ -226,15 +226,6 @@ constexpr const char* warp_execution = "warp_execution";
  */
 constexpr const char* layout_transforms = "layout_transforms";
 
-/*!
- * \brief Marks the physical axis separators
- *
- * Only applies to a DataProducer, as it should be made part of the
- * Buffer definition in a PrimFunc.  See `BufferNode::axis_separators`
- * for more details.
- */
-constexpr const char* axis_separators = "axis_separators";
-
 /*!
  * \brief Mark that the kernel is hand threaded and doesn't need syncs inserted
  */

include/tvm/tirx/buffer.h

@@ -60,13 +60,6 @@ inline DLDataType DefaultIndexType() {
 // forward declare Stmt
 class Stmt;
 
-/*! \brief buffer type */
-enum BufferType : int {
-  kDefault = 1,
-  // Maps buffer[i][j][k] -> buffer[i][0][k] if dimension i's shape equals 1.
-  kAutoBroadcast = 2,
-};
-
 /*! \brief Node to represent a buffer */
 class BufferNode : public ffi::Object {
  public:
@@ -85,15 +78,6 @@ class BufferNode : public ffi::Object {
    * generators.
    */
   ffi::Array<PrimExpr> shape;
-  /*!
-   * \brief Separators between input axes when generating flattened output axes
-   *
-   * For buffers representing flat 1-d memory (e.g. any buffer in
-   * RAM), this should be an empty array.  For buffers representing
-   * non-flat memory, each entry in axis_separators should be the
-   * first input axis that is part of a new flattened axis.
-   */
-  ffi::Array<IntImm> axis_separators;
   /*!
    * \brief The strides of each dimension
    *  This can be an empty array, indicating array is contiguous
@@ -111,8 +95,6 @@ class BufferNode : public ffi::Object {
    *  elem_offset is guaranteed to be multiple of offset_factor.
    */
   int offset_factor;
-  /*! \brief buffer type */
-  BufferType buffer_type;
   /*!
    * \brief Span that points to the original source code.
    *        Reserved debug information.
@@ -141,15 +123,12 @@ class BufferNode : public ffi::Object {
         .def_ro("shape", &BufferNode::shape, refl::AttachFieldFlag::SEqHashDefRecursive())
         // TODO(tqchen): use SEqHashDefNonRecursive after the next pypi tvm-ffi release
         .def_ro("strides", &BufferNode::strides, refl::AttachFieldFlag::SEqHashDefRecursive())
-        .def_ro("axis_separators", &BufferNode::axis_separators,
-                refl::AttachFieldFlag::SEqHashDefRecursive())
         // TODO(tqchen): use SEqHashDefNonRecursive after the next pypi tvm-ffi release
         .def_ro("elem_offset", &BufferNode::elem_offset,
                 refl::AttachFieldFlag::SEqHashDefRecursive())
         .def_ro("name", &BufferNode::name, refl::AttachFieldFlag::SEqHashIgnore())
         .def_ro("data_alignment", &BufferNode::data_alignment)
         .def_ro("offset_factor", &BufferNode::offset_factor)
-        .def_ro("buffer_type", &BufferNode::buffer_type)
         .def_ro("span", &BufferNode::span, refl::AttachFieldFlag::SEqHashIgnore())
         .def_ro("layout", &BufferNode::layout)
         .def_ro("allocated_addr", &BufferNode::allocated_addr);
@@ -190,7 +169,6 @@ class Buffer : public ffi::ObjectRef {
   // A default value will be picked.
   TVM_DLL Buffer(Var data, PrimType dtype, ffi::Array<PrimExpr> shape, ffi::Array<PrimExpr> strides,
                  PrimExpr elem_offset, ffi::String name, int data_alignment, int offset_factor,
-                 BufferType buffer_type, ffi::Array<IntImm> axis_separators = {},
                  Span span = Span(), ffi::Optional<Layout> layout = std::nullopt,
                  ffi::Array<PrimExpr> allocated_addr = {});
 
@@ -299,14 +277,12 @@ class Buffer : public ffi::ObjectRef {
  * \param dtype The content data type.
  * \param name The name of the buffer
  * \param storage_scope The storage scope associated with this buffer
- * \param axis_separators Divisions defining the groups of axes that will be flattened together.
  * \param span The location of this object in the source code.
  * \return The created buffer.
  * \sa Buffer for complete constructor.
  */
 TVM_DLL Buffer decl_buffer(ffi::Array<PrimExpr> shape, PrimType dtype = PrimType::Float(32),
                            ffi::String name = "buffer", ffi::String storage_scope = "",
-                           ffi::Optional<ffi::Array<IntImm>> axis_separators = std::nullopt,
                            Span span = Span());
 
 /*!
@@ -362,13 +338,11 @@ class DataProducer : public PrimExprConvertible {
  *                      multiple of offset_factor
                         User can specify data_alignment and offset_factor to be 0
  *                      A default value will be picked.
- * \param compact If the statement has already bound to a compact buffer.
  * \param memory_scope memory scope of the buffer
  */
 TVM_DLL tirx::Buffer BufferWithOffsetAlignment(ffi::Array<PrimExpr> shape, PrimType dtype,
                                                std::string name, int data_alignment,
-                                               int offset_factor, bool compact,
-                                               std::string memory_scope = "");
+                                               int offset_factor, std::string memory_scope = "");
 }  // namespace tirx
 }  // namespace tvm
 #endif  // TVM_TIR_BUFFER_H_

include/tvm/tirx/script/builder/ir.h

@@ -53,16 +53,12 @@ using tvm::tirx::Var;
  * \param storage_scope The optional storage scope of buffer data pointer.
  * \param align The alignment requirement of data pointer in bytes.
  * \param offset_factor The factor of elem_offset field.
- * \param buffer_type The buffer type.
- * \param axis_separators The separators between input axes when generating flattened output axes.
  * \return The declared buffer.
  */
 Buffer BufferDecl(ffi::Array<PrimExpr> shape, PrimType dtype, ffi::String buffer_name,
                   ffi::Optional<Var> data, ffi::Optional<ffi::Array<PrimExpr>> strides,
                   ffi::Optional<PrimExpr> elem_offset, ffi::String storage_scope, int align,
-                  int offset_factor, ffi::String buffer_type,
-                  ffi::Optional<ffi::Array<IntImm>> axis_separators,
-                  ffi::Optional<Layout> layout = std::nullopt,
+                  int offset_factor, ffi::Optional<Layout> layout = std::nullopt,
                   ffi::Array<PrimExpr> allocated_addr = {});
 
 /*!
@@ -117,16 +113,12 @@ Type FuncRet(Type ret_type);
  * \param storage_scope The optional storage scope of buffer data pointer.
  * \param align The alignment requirement of data pointer in bytes.
  * \param offset_factor The factor of elem_offset field.
- * \param buffer_type The buffer type.
- * \param axis_separators The separators between input axes when generating flattened output axes.
  * \return The matched buffer.
  */
 Buffer MatchBuffer(ffi::ObjectRef param, ffi::Array<PrimExpr> shape,
                    PrimType dtype = PrimType::Float(32), ffi::Optional<Var> data = std::nullopt,
                    ffi::Array<PrimExpr> strides = {}, PrimExpr elem_offset = PrimExpr(),
                    ffi::String storage_scope = "global", int align = -1, int offset_factor = 0,
-                   ffi::String buffer_type = "default",
-                   ffi::Optional<ffi::Array<IntImm>> axis_separators = std::nullopt,
                    ffi::Optional<Layout> layout = std::nullopt);
 
 /*!
@@ -189,8 +181,6 @@ void BlockAttrs(ffi::Map<ffi::String, ffi::Any> attrs);
  * \param storage_scope The optional storage scope of buffer data pointer.
  * \param align The alignment requirement of data pointer in bytes.
  * \param offset_factor The factor of elem_offset field.
- * \param buffer_type The buffer type.
- * \param axis_separators The separators between input axes when generating flattened output axes.
  * \param layout The layout of the buffer.
  * \param allocated_addr The allocated address of the buffer. Might be multi-dimensional.
  * \return The allocated buffer or the AllocBufferFrame if the function is called under
@@ -200,9 +190,8 @@ ffi::Variant<Buffer, AllocBufferFrame> SBlockAllocBuffer(
     ffi::Array<PrimExpr> shape, PrimType dtype = PrimType::Float(32),
     ffi::Optional<Var> data = std::nullopt, ffi::Array<PrimExpr> strides = {},
     PrimExpr elem_offset = PrimExpr(), ffi::String storage_scope = "", int align = -1,
-    int offset_factor = 0, ffi::String buffer_type = "default",
-    ffi::Optional<ffi::Array<IntImm>> axis_separators = std::nullopt,
-    ffi::Optional<Layout> layout = std::nullopt, ffi::Array<PrimExpr> allocated_addr = {});
+    int offset_factor = 0, ffi::Optional<Layout> layout = std::nullopt,
+    ffi::Array<PrimExpr> allocated_addr = {});
 
 namespace axis {
 
@@ -413,16 +402,13 @@ ElseFrame Else();
  * \param storage_scope The optional storage scope of buffer data pointer.
  * \param align The alignment requirement of data pointer in bytes.
  * \param offset_factor The factor of elem_offset field.
- * \param buffer_type The buffer type.
- * \param axis_separators The separators between input axes when generating flattened output axes.
  * \param layout The layout of the buffer.
  * \return The declaration frame.
  */
 DeclBufferFrame DeclBuffer(ffi::Array<PrimExpr> shape, PrimType dtype, ffi::String buffer_name,
                            ffi::Optional<Var> data, ffi::Optional<ffi::Array<PrimExpr>> strides,
                            ffi::Optional<PrimExpr> elem_offset, ffi::String storage_scope,
-                           int align, int offset_factor, ffi::String buffer_type,
-                           ffi::Optional<ffi::Array<IntImm>> axis_separators,
+                           int align, int offset_factor,
                            ffi::Optional<Layout> layout = std::nullopt,
                            ffi::Optional<PrimExpr> allocated_addr = std::nullopt);
 

python/tvm/backend/cuda/lang/alloc_pool.py

@@ -448,8 +448,6 @@ def alloc(
         strides=None,
         scope="shared.dyn",
         align=0,
-        buffer_type="",
-        axis_separators=None,
         layout="default",
     ):
         ir = _get_ir()
@@ -463,8 +461,6 @@ def alloc(
             byte_offset=self.offset,
             scope=scope,
             align=align,
-            buffer_type=buffer_type,
-            axis_separators=axis_separators,
             layout=layout,
         )
         # Advance in bits then round up to bytes so sub-byte dtypes (e.g.

python/tvm/contrib/hexagon/tools.py

@@ -19,7 +19,6 @@
 """Tools/compilers/linkers for Hexagon"""
 
 import io
-import itertools
 import os
 import pathlib
 import re
@@ -401,13 +400,8 @@ def export_module(module, out_dir, binary_name="test_binary.so"):
     return binary_path
 
 
-def allocate_hexagon_array(
-    dev, tensor_shape=None, dtype=None, data=None, axis_separators=None, mem_scope=None
-):
-    """
-    Allocate a hexagon array which could be a 2D array
-    on physical memory defined by axis_separators
-    """
+def allocate_hexagon_array(dev, tensor_shape=None, dtype=None, data=None, mem_scope=None):
+    """Allocate a Hexagon array backed by flat physical memory."""
     if tensor_shape is None:
         assert data is not None, "Must provide either tensor shape or numpy data array"
         tensor_shape = data.shape
@@ -422,13 +416,7 @@ def allocate_hexagon_array(
     elif data is not None:
         assert dtype == data.dtype, "Mismatch between provided dtype and numpy data array dtype"
 
-    if axis_separators is None:
-        axis_separators = []
-
-    boundaries = [0, *axis_separators, len(tensor_shape)]
-    physical_shape = [
-        numpy.prod(tensor_shape[dim_i:dim_f]) for dim_i, dim_f in itertools.pairwise(boundaries)
-    ]
+    physical_shape = [numpy.prod(tensor_shape)]
 
     arr = tvm.runtime.empty(physical_shape, dtype=dtype, device=dev, mem_scope=mem_scope)
 

python/tvm/relax/op/manipulate.py

@@ -116,8 +116,6 @@ def layout_transform(
     x: Expr,
     index_map: Callable | IndexMap,
     pad_value: int | float | Expr | None = None,
-    axis_separators: int | str | None = None,  # str for IndexMap.AXIS_SEPARATOR
-    input_axis_separators: int | str | None = None,  # str for IndexMap.AXIS_SEPARATOR
 ):
     """Modifies the layout of a tensor.
 
@@ -133,9 +131,6 @@ def layout_transform(
         The value used for padding if the transformation results in implicit padding.
         If not specified, any value can be used.
 
-    axis_separators : Optional[int | IndexMap.AXIS_SEPARATOR]
-        The axis_separators for index_map to create non flat buffers.
-
     Returns
     -------
     result : relax.Expr
@@ -160,15 +155,7 @@ def layout_transform(
             pad_value = FloatImm(x_dtype.dtype, float(pad_value))
         pad_value = prim_value(pad_value)
 
-    if axis_separators is None:
-        axis_separators = []
-
-    if input_axis_separators is None:
-        input_axis_separators = []
-
-    return _ffi_api.layout_transform(
-        x, index_map, pad_value, axis_separators, input_axis_separators
-    )
+    return _ffi_api.layout_transform(x, index_map, pad_value)
 
 
 def permute_dims(x: Expr, axes: list[int] | None = None) -> Expr:

python/tvm/relax/transform/legalize_ops/manipulate.py

@@ -19,7 +19,7 @@
 """Default legalization function for manipulate operators."""
 
 import tvm
-from tvm import DataTypeCode, relax, s_tir, te, tirx, topi
+from tvm import DataTypeCode, relax, te, tirx, topi
 from tvm.ir import Call
 from tvm.relax.op.base import call_tir
 from tvm.relax.type import TensorType
@@ -329,9 +329,6 @@ def te_layout_transform(data, name):
             name=name,
         )
 
-    def set_axis_sep(axis_sep: list, sch: s_tir.schedule, buffer_type: str):
-        sch.set_axis_separator(primfunc_name, (buffer_type, 0), axis_separators=axis_sep)
-
     index_map: tvm.tirx.IndexMap = call.attrs.index_map
     pad_value = call.attrs.pad_value
     if pad_value is not None:
@@ -342,24 +339,14 @@ def set_axis_sep(axis_sep: list, sch: s_tir.schedule, buffer_type: str):
         else:
             pad_value = 0.0
 
-    axis_separators: tvm.tirx.IndexMap.AXIS_SEPARATOR = call.attrs.axis_separators
-    input_axis_separators: tvm.tirx.IndexMap.AXIS_SEPARATOR = call.attrs.input_axis_separators
-
-    # Convert to list from array
-    axis_separators = [int(sep) for sep in axis_separators]
     primfunc_name = "te_layout_transform"
     _, padding_predicate = index_map.non_surjective_inverse(call.args[0].ty.shape)
     if not isinstance(padding_predicate, tvm.tirx.expr.IntImm):
         primfunc_name += "_with_pad"
-    if len(axis_separators) != 0:
-        primfunc_name += "_axis_separator"
     tir_func, call_args, _ = gen_call_tir_inputs(te_layout_transform, call.args[0], primfunc_name)
-    # Create TIR schedule to apply layout changes with axis separators
+    # Create a TIR schedule to apply the layout change.
     sch = tvm.s_tir.Schedule(tir_func)
     sch.transform_layout(primfunc_name, ("write", 0), index_map, pad_value)
-    set_axis_sep(axis_separators, sch, "write")
-    if input_axis_separators is not None:
-        set_axis_sep(input_axis_separators, sch, "read")
     gvar = bb.add_func(sch.mod["main"], primfunc_name)
     output_shape = index_map.map_shape(list(call_args[0].ty.shape))
     output_dtype = call_args[0].ty.dtype