From e120b11ca4e46bf2ccd8f1ec91733446671860b6 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 16 Jul 2026 13:51:23 +0000 Subject: [PATCH] [REFACTOR][TIR] Remove buffer type and axis separators --- include/tvm/relax/attrs/manipulate.h | 19 +- include/tvm/relax/transform.h | 11 +- include/tvm/s_tir/schedule/schedule.h | 12 - include/tvm/s_tir/stmt.h | 9 - include/tvm/tirx/buffer.h | 28 +-- include/tvm/tirx/script/builder/ir.h | 22 +- python/tvm/backend/cuda/lang/alloc_pool.py | 4 - python/tvm/contrib/hexagon/tools.py | 18 +- python/tvm/relax/op/manipulate.py | 15 +- .../transform/legalize_ops/manipulate.py | 17 +- python/tvm/relax/transform/transform.py | 16 +- python/tvm/s_tir/schedule/schedule.py | 110 +-------- python/tvm/te/__init__.py | 2 +- python/tvm/te/operation.py | 3 - python/tvm/tirx/buffer.py | 15 -- python/tvm/tirx/function.py | 78 +------ python/tvm/tirx/script/builder/ir.py | 51 ----- python/tvm/tirx/script/builder/tirx.py | 3 - python/tvm/tirx/script/parser/entry.py | 4 - python/tvm/tirx/script/parser/parser.py | 4 - src/backend/hexagon/runtime/ops/conv2d.h | 6 +- .../trn/transform/lower_trainium_layout.cc | 2 - src/relax/op/tensor/manipulate.cc | 6 +- src/relax/op/tensor/manipulate.h | 7 +- src/relax/transform/alter_op_impl.cc | 102 ++------- src/relax/transform/convert_layout.cc | 4 - src/relax/transform/fuse_tir.cc | 3 +- src/s_tir/schedule/concrete_schedule.cc | 10 - src/s_tir/schedule/concrete_schedule.h | 3 - src/s_tir/schedule/primitive.h | 12 - src/s_tir/schedule/primitive/cache_index.cc | 4 +- .../primitive/layout_transformation.cc | 120 ---------- src/s_tir/schedule/schedule.cc | 9 +- src/s_tir/schedule/traced_schedule.cc | 15 -- src/s_tir/schedule/traced_schedule.h | 3 - .../transform/lower_cross_thread_reduction.cc | 3 +- .../transform/memhammer_tensorcore_rewrite.cc | 18 +- .../merge_shared_memory_allocations.cc | 2 +- .../transform/transform_mma_buffer_layout.cc | 9 +- src/tirx/ir/buffer.cc | 159 ++----------- src/tirx/ir/stmt.cc | 5 - src/tirx/script/builder/ir.cc | 38 ++-- src/tirx/script/printer/buffer.cc | 15 +- src/tirx/script/printer/function.cc | 3 +- src/tirx/transform/ir_utils.h | 7 +- src/tirx/transform/lower_intrin.cc | 3 +- src/tirx/transform/lower_tirx_cleanup.cc | 2 - src/tirx/transform/lower_warp_memory.cc | 2 +- src/tirx/transform/storage_rewrite.cc | 20 +- src/tirx/transform/tvm_ffi_binder.cc | 19 -- src/tirx/transform/tvm_ffi_binder.h | 12 - .../transform/unsupported_dtype_legalize.cc | 16 +- tests/cpp/ir_functor_test.cc | 7 +- .../relax/test_optimize_layout_transform.py | 6 - .../relax/test_transform_alter_op_impl.py | 157 ------------- .../relax/test_transform_convert_layout.py | 48 ---- tests/python/relax/test_transform_fuse_tir.py | 127 ----------- .../test_transform_legalize_ops_manipulate.py | 40 ---- .../test_tir_schedule_set_axis_separator.py | 208 ------------------ .../test_tir_schedule_transform_layout.py | 118 ---------- tests/python/tirx-base/test_tir_buffer.py | 13 -- tests/python/tirx-base/test_tir_intrin.py | 4 - .../test_tir_transform_flatten_buffer.py | 68 ------ .../test_tir_transform_storage_rewrite.py | 92 -------- .../tvmscript/test_tvmscript_roundtrip.py | 92 -------- 65 files changed, 125 insertions(+), 1935 deletions(-) delete mode 100644 tests/python/s_tir/schedule/test_tir_schedule_set_axis_separator.py diff --git a/include/tvm/relax/attrs/manipulate.h b/include/tvm/relax/attrs/manipulate.h index f916e8244ad7..6a3061b112f5 100644 --- a/include/tvm/relax/attrs/manipulate.h +++ b/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 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> axis_separators; - /*! - * axis_separators for input buffers. - * Needed to identify if the input buffer to layout_transform - * contains axis separator. - */ - ffi::Optional> 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); diff --git a/include/tvm/relax/transform.h b/include/tvm/relax/transform.h index 94382cf60210..df66358455e0 100644 --- a/include/tvm/relax/transform.h +++ b/include/tvm/relax/transform.h @@ -595,16 +595,11 @@ TVM_DLL Pass DecomposeOpsForTraining(ffi::Optional 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& op_impl_map, - const ffi::Map>& op_buffer_transforms, - const ffi::Map>>>& axis_separators, - const ffi::Map>>>& - input_axis_separators); +TVM_DLL Pass +AlterOpImpl(const ffi::Map& op_impl_map, + const ffi::Map>& op_buffer_transforms); /*! * \brief Layout conversion pass. diff --git a/include/tvm/s_tir/schedule/schedule.h b/include/tvm/s_tir/schedule/schedule.h index 55d6ab72ef58..dd3b91877777 100644 --- a/include/tvm/s_tir/schedule/schedule.h +++ b/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& axis_separators) = 0; - /******** Schedule: Padding ********/ /*! * \brief Decompose a padding block into a block filling const pad values and a block diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h index 72aae3642db6..6eaee11ca031 100644 --- a/include/tvm/s_tir/stmt.h +++ b/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 */ diff --git a/include/tvm/tirx/buffer.h b/include/tvm/tirx/buffer.h index 8d31e23cc1b5..e2a3e167d8db 100644 --- a/include/tvm/tirx/buffer.h +++ b/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 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 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 shape, ffi::Array strides, PrimExpr elem_offset, ffi::String name, int data_alignment, int offset_factor, - BufferType buffer_type, ffi::Array axis_separators = {}, Span span = Span(), ffi::Optional layout = std::nullopt, ffi::Array 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 shape, PrimType dtype = PrimType::Float(32), ffi::String name = "buffer", ffi::String storage_scope = "", - ffi::Optional> 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 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_ diff --git a/include/tvm/tirx/script/builder/ir.h b/include/tvm/tirx/script/builder/ir.h index 9997e3ec524b..ff005566417e 100644 --- a/include/tvm/tirx/script/builder/ir.h +++ b/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 shape, PrimType dtype, ffi::String buffer_name, ffi::Optional data, ffi::Optional> strides, ffi::Optional elem_offset, ffi::String storage_scope, int align, - int offset_factor, ffi::String buffer_type, - ffi::Optional> axis_separators, - ffi::Optional layout = std::nullopt, + int offset_factor, ffi::Optional layout = std::nullopt, ffi::Array 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 shape, PrimType dtype = PrimType::Float(32), ffi::Optional data = std::nullopt, ffi::Array strides = {}, PrimExpr elem_offset = PrimExpr(), ffi::String storage_scope = "global", int align = -1, int offset_factor = 0, - ffi::String buffer_type = "default", - ffi::Optional> axis_separators = std::nullopt, ffi::Optional layout = std::nullopt); /*! @@ -189,8 +181,6 @@ void BlockAttrs(ffi::Map 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 SBlockAllocBuffer( ffi::Array shape, PrimType dtype = PrimType::Float(32), ffi::Optional data = std::nullopt, ffi::Array strides = {}, PrimExpr elem_offset = PrimExpr(), ffi::String storage_scope = "", int align = -1, - int offset_factor = 0, ffi::String buffer_type = "default", - ffi::Optional> axis_separators = std::nullopt, - ffi::Optional layout = std::nullopt, ffi::Array allocated_addr = {}); + int offset_factor = 0, ffi::Optional layout = std::nullopt, + ffi::Array 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 shape, PrimType dtype, ffi::String buffer_name, ffi::Optional data, ffi::Optional> strides, ffi::Optional elem_offset, ffi::String storage_scope, - int align, int offset_factor, ffi::String buffer_type, - ffi::Optional> axis_separators, + int align, int offset_factor, ffi::Optional layout = std::nullopt, ffi::Optional allocated_addr = std::nullopt); diff --git a/python/tvm/backend/cuda/lang/alloc_pool.py b/python/tvm/backend/cuda/lang/alloc_pool.py index 725264a35b16..c635999b865f 100644 --- a/python/tvm/backend/cuda/lang/alloc_pool.py +++ b/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. diff --git a/python/tvm/contrib/hexagon/tools.py b/python/tvm/contrib/hexagon/tools.py index 85c456014b64..1241c704206c 100644 --- a/python/tvm/contrib/hexagon/tools.py +++ b/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) diff --git a/python/tvm/relax/op/manipulate.py b/python/tvm/relax/op/manipulate.py index 057ffa4e9f15..9c7ad20197aa 100644 --- a/python/tvm/relax/op/manipulate.py +++ b/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: diff --git a/python/tvm/relax/transform/legalize_ops/manipulate.py b/python/tvm/relax/transform/legalize_ops/manipulate.py index ac72f032f951..4301641b05c7 100644 --- a/python/tvm/relax/transform/legalize_ops/manipulate.py +++ b/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 diff --git a/python/tvm/relax/transform/transform.py b/python/tvm/relax/transform/transform.py index 4538e7e6cdb4..ff4601290af7 100644 --- a/python/tvm/relax/transform/transform.py +++ b/python/tvm/relax/transform/transform.py @@ -1304,8 +1304,6 @@ def DecomposeOpsForTraining(func_name: str | None = None) -> tvm.ir.transform.Pa def AlterOpImpl( op_impl_map: dict[str, PrimFunc], op_buffer_transforms: dict[str, list[IndexMap | Callable]], - op_buffer_axis_separators: dict[str, list[str | Callable]], # str=IndexMap.AXIS_SEPARATOR - op_buffer_input_axis_separators: dict[str, list[str | Callable]], # str=IndexMap.AXIS_SEPARATOR ): """Replace all PrimFunc's which have matching 'operator_name' attribute, with replacement PrimFunc that could possibly have different layouts on i/o buffers. The layout @@ -1319,11 +1317,6 @@ def AlterOpImpl( op_kind to PrimFunc map op_buffer_transforms: Dict[str, List[Union[IndexMap, Callable]] op_kind to layout transformation map for each of the buffers - op_buffer_axis_separators: Dict[str, List[Union[IndexMap.AXIS_SEPARATOR, Callable]]] - op_kind to axis_separator for each index_map - op_buffer_input_axis_separators: Dict[str, List[Union[IndexMap.AXIS_SEPARATOR, Callable]]] - op_kind to axis_separator for input index_map - Returns ------- ret: tvm.ir.transform.Pass @@ -1333,18 +1326,13 @@ def AlterOpImpl( for transform in transform_list: # Extract the index_map if isinstance(transform, Callable): - transform = IndexMap.from_func_with_separators(transform)[0] + transform = IndexMap.from_func(transform) elif isinstance(transform, Array | tuple) and isinstance(transform[0], IndexMap): transform = transform[0] l.append(transform) op_buffer_transforms[operator_name] = l - return _ffi_api.AlterOpImpl( - op_impl_map, - op_buffer_transforms, - op_buffer_axis_separators, - op_buffer_input_axis_separators, - ) # type: ignore + return _ffi_api.AlterOpImpl(op_impl_map, op_buffer_transforms) # type: ignore def ConvertLayout( diff --git a/python/tvm/s_tir/schedule/schedule.py b/python/tvm/s_tir/schedule/schedule.py index 43df7c597ad3..ecad5f1cb578 100644 --- a/python/tvm/s_tir/schedule/schedule.py +++ b/python/tvm/s_tir/schedule/schedule.py @@ -3349,11 +3349,6 @@ def transform_layout( The transformation to apply. - If `index_map` is a callable, and the returned list - contains IndexMap.AXIS_SEPARATOR, the SetAxisSeparators - primitive will be called in addition to the - TransformLayout primitive. - pad_value: Optional[int | float | Expr | IndexMap | Callable] The value to be used for any padding introduced by the @@ -3444,13 +3439,11 @@ def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> ndim = len(buffer_obj.shape) if callable(index_map): - index_map, axis_separators = IndexMap.from_func_with_separators( + index_map = IndexMap.from_func( index_map, ndim=ndim, index_dtype=_get_sblock_default_dtype(self.get(block)), ) - else: - axis_separators = [] if pad_value is None: pass @@ -3489,10 +3482,6 @@ def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> pad_value, assume_injective_transform, ) - if axis_separators: - _ffi_api.ScheduleSetAxisSeparator( # type: ignore # pylint: disable=no-member - self, block, buffer_index, buffer_index_type_enum, axis_separators - ) @type_checked def transform_block_layout(self, block: SBlockRV | str, index_map: IndexMap | Callable) -> None: @@ -3555,103 +3544,6 @@ def after_transform_block_layout( self, block, index_map ) - def set_axis_separator( - self, - block: SBlockRV | str, - buffer: tuple[str, int] | str | Buffer, - axis_separators: list[int] | None, - ) -> None: - """Set the axis separator of a buffer, where the buffer is specified by a block and a read - or write index. - - Parameters - ---------- - block : SBlockRV | str - - The block that accesses the target buffer. If a string, - this must uniquely identify a block. - - buffer: Union[Tuple[str,int], Buffer, str] - - The buffer to be transformed, or a specification of how to - identify the buffer to be transformed. - - If `buffer` if a tuple of ``(str,int)``, the first item - should be either "read" or "write", and the second item is - an index into the block's read or write regions. - - If `buffer` is a string, it is the name of the buffer, - which must exist within the reads/writes of the block. In - addition, the reads/writes of the block may not contain - more than one buffer with this name. - - If `buffer` is a Buffer object, it must exist within the - reads/writes of the block. - - axis_separators : Optional[List[int]] - - The axis separators. - - Examples - -------- - - Before set_axis_separator, in TensorIR, the IR is: - - .. code-block:: python - - @T.prim_func(s_tir=True) - def before_set_axis_separator( - A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32") - ) -> None: - B = T.sblock_alloc_buffer((128, 128), dtype="float32") - - for i, j in T.grid(128, 128): - with T.sblock("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.sblock("C"): - vi, vj = T.axis.remap("SS", [i, j]) - C[vi, vj] = B[vi, vj] + 1.0 - - Create the schedule and do set_axis_separator: - - .. code-block:: python - - sch = tvm.s_tir.Schedule(before_set_axis_separator) - sch.set_axis_separators(sch.get_sblock("B"), buffer=("write", 0), - axis_separators=[1]) - print(sch.mod["main"].script()) - - After applying set_axis_separator, the IR becomes: - - .. code-block:: python - - @T.prim_func(s_tir=True) - def after_set_axis_separators( - A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32") - ) -> None: - B = T.sblock_alloc_buffer([128, 128], dtype="float32", axis_separators=[1]) - - for i, j in T.grid(128, 128): - with T.sblock("B"): - vi, vj = T.axis.remap("SS", [i, j]) - B[vi, vj] = A[vi, vj] * T.float32(2) - for i, j in T.grid(128, 128): - with T.sblock("C"): - vi, vj = T.axis.remap("SS", [i, j]) - C[vi, vj] = B[vi, vj] + T.float32(1) - """ - axis_separators = axis_separators or [] - - block = self._normalize_block_arg(block) - buffer_index_type, buffer_index, _ = self._normalize_buffer_arg(block, buffer) - - buffer_index_type_enum = 0 if buffer_index_type == "read" else 1 - _ffi_api.ScheduleSetAxisSeparator( # type: ignore # pylint: disable=no-member - self, block, buffer_index, buffer_index_type_enum, axis_separators - ) - ########## Schedule: Padding decomposition ######### @type_checked def decompose_padding(self, block: SBlockRV | str, loop: LoopRV) -> SBlockRV: diff --git a/python/tvm/te/__init__.py b/python/tvm/te/__init__.py index 4773b8bcb99b..94b16ec5f635 100644 --- a/python/tvm/te/__init__.py +++ b/python/tvm/te/__init__.py @@ -30,7 +30,7 @@ from .tensor import TensorSlice, Tensor from .tag import tag_scope from .operation import placeholder, compute, scan, extern, var, const -from .operation import thread_axis, reduce_axis, AXIS_SEPARATOR +from .operation import thread_axis, reduce_axis from .operation import create_prim_func from .operation import extern_primfunc diff --git a/python/tvm/te/operation.py b/python/tvm/te/operation.py index b0f639e64305..73c8d7f26301 100644 --- a/python/tvm/te/operation.py +++ b/python/tvm/te/operation.py @@ -590,6 +590,3 @@ def tir_matmul(a: T.handle, b: T.handle, c: T.handle) -> None: if not isinstance(ops, list | tuple | Array): ops = [ops] return _ffi_api.CreatePrimFunc(ops, index_dtype_override) - - -AXIS_SEPARATOR = tvm.tirx.IndexMap.AXIS_SEPARATOR diff --git a/python/tvm/tirx/buffer.py b/python/tvm/tirx/buffer.py index 1d6a380adb6d..b22688a09d06 100644 --- a/python/tvm/tirx/buffer.py +++ b/python/tvm/tirx/buffer.py @@ -350,8 +350,6 @@ def _infer_shape(shape): self.scope(), self.data_alignment, self.offset_factor, - "", - self.axis_separators, layout, ) else: @@ -383,8 +381,6 @@ def _infer_shape(shape): self.scope(), self.data_alignment, self.offset_factor, - "", - self.axis_separators, self.layout if layout is None else layout, ) @@ -425,8 +421,6 @@ def local(self, *shape, layout=None) -> "Buffer": self.scope(), self.data_alignment, self.offset_factor, - "", - self.axis_separators, self.layout.storage() if layout is None else layout, ) @@ -463,8 +457,6 @@ def permute(self, *dims) -> "Buffer": self.scope(), self.data_alignment, self.offset_factor, - "", - self.axis_separators, new_layout, ) @@ -529,8 +521,6 @@ def decl_buffer( scope="", data_alignment=-1, offset_factor=0, - buffer_type="", - axis_separators=None, span=None, layout="default", ): @@ -542,9 +532,6 @@ def decl_buffer( dtype = "float32" if dtype is None else dtype strides = () if strides is None else strides - if axis_separators is None: - axis_separators = [] - if layout == "default": layout = TileLayout(S[tuple(shape)]) if shape else None @@ -565,8 +552,6 @@ def decl_buffer( name, data_alignment, offset_factor, - buffer_type, - axis_separators, span, layout, ) diff --git a/python/tvm/tirx/function.py b/python/tvm/tirx/function.py index ab7b6435a242..80b8fcaab0a6 100644 --- a/python/tvm/tirx/function.py +++ b/python/tvm/tirx/function.py @@ -250,11 +250,6 @@ class IndexMap(Object): initial_indices: list[Var] final_indices: list[Expr] - # Sentinel value used to indicate which groups of pre-flattening axes - # should be used to post-flattening axes axes. See - # Stage.transform_layout for more details. - AXIS_SEPARATOR = "axis_separator" - def __init__(self, initial_indices, final_indices, inverse_index_map): if isinstance(inverse_index_map, Callable): inverse_index_map = IndexMap.from_func(inverse_index_map) @@ -297,74 +292,14 @@ def from_func( It is the user's responsibility to ensure the correctness of the pre-defined inverse index map. - Returns - ------- - index_map: IndexMap - - Returns an IndexMap representing the `mapping_function`. - - """ - index_map, axis_separators = IndexMap.from_func_with_separators( - mapping_function, - ndim, - inverse_index_map, - index_dtype=index_dtype, - ) - assert not axis_separators, ( - "The mapping_function provided to IndexMap.from_func " - "may not return IndexMap.AXIS_SEPARATOR. " - "If required, please use IndexMap.from_func_with_separators instead." - ) - return index_map - - @staticmethod - def from_func_with_separators( - mapping_function: Callable, - ndim: int | None = None, - inverse_index_map: Callable | Optional["IndexMap"] = None, - *, - index_dtype: str = "int64", - ): - """Create an index map from a function - - Parameters - ---------- - mapping_function : Callable - - The function to map from source indices to target indices. - The function should accept tirx.Var parameters and return - either a `tirx.Expr` or a list. Each element of the - returned list should be either a `tirx.Expr` or the - object `IndexMap.AXIS_SEPARATOR`. Returning a - `tirx.Expr` is equivalent to returning a list of length - 1 containing that `tirx.Expr`. - - ndim: Optional[int] - - The dimensionality of the buffer to which this - transformation should be applied. If mapping_function uses - variadic argument `*args`, ndim must be specified. If - mapping_function does not use variadic arguments, ndim is - optional. - - inverse_index_map : Union[Callable, Optional[IndexMap]] - The optional pre-defined inverse index map. - When this is defined, IndexMap::Inverse will return the pre-defined inverse index map. - Otherwise, the inverse index map will be computed on the fly. - It is the user's responsibility to ensure the correctness of the pre-defined inverse - index map. - index_dtype : str The default index dtype to use for input iters in the mapping function. Returns ------- - ret: Tuple[IndexMap, List[int]] + index_map: IndexMap - Returns a tuple whose first element is an IndexMap - representing the `mapping_function`, and whose second index - is a list of indices at which `IndexMap.AXIS_SEPARATOR` - occurred. + Returns an IndexMap representing the `mapping_function`. """ params = inspect.signature(mapping_function).parameters @@ -403,7 +338,6 @@ def from_func_with_separators( initial_indices = args + list(kwargs.values()) final_indices = [] - axis_separators = [] if tvm.ir.is_prim_expr(mapping): is_iterable = False @@ -418,18 +352,16 @@ def from_func_with_separators( for val in mapping: if tvm.ir.is_prim_expr(val): final_indices.append(val) - elif val is IndexMap.AXIS_SEPARATOR: - axis_separators.append(len(final_indices)) else: raise TypeError( - "Expected mapping function to return list of " - "either tvm.ir.Expr or IndexMap.AXIS_SEPARATOR. " + "Expected mapping function to return tvm.ir.Expr " + "or a list of tvm.ir.Expr. " f"Instead received {val} of type {type(val)}." ) else: final_indices.append(mapping) - return IndexMap(initial_indices, final_indices, inverse_index_map), axis_separators + return IndexMap(initial_indices, final_indices, inverse_index_map) def is_equivalent_to(self, other_map: "IndexMap", analyzer=None) -> bool: """Return if the index maps are equivalent. diff --git a/python/tvm/tirx/script/builder/ir.py b/python/tvm/tirx/script/builder/ir.py index ccd091ae6309..560a3b152b0b 100644 --- a/python/tvm/tirx/script/builder/ir.py +++ b/python/tvm/tirx/script/builder/ir.py @@ -286,8 +286,6 @@ def buffer( scope: str = "global", align: int = 0, offset_factor: int = 0, - buffer_type: str = "", - axis_separators: list[int] | None = None, layout: str | Layout | None = "default", allocated_addr: int | tuple[int, ...] | None = None, buffer_name: str = "", @@ -320,12 +318,6 @@ def buffer( offset_factor : int The factor of elem_offset field. - buffer_type : str - The buffer type. - - axis_separators : List[int] - The separators between input axes when generating flattened output axes. - buffer_name : str The name of the buffer. @@ -353,8 +345,6 @@ def buffer( scope, align, offset_factor, - buffer_type, - axis_separators, _get_layout(layout, shape, scope), allocated_addr, ) @@ -467,8 +457,6 @@ def match_buffer( scope: str = "global", align: int = -1, offset_factor: int = 0, - buffer_type: str = "default", - axis_separators: list[int] | None = None, layout: str | Layout | None = "default", ) -> Buffer: """The buffer match function. @@ -518,12 +506,6 @@ def match_buffer( offset_factor : int The factor of elem_offset field. - buffer_type : str - The buffer type. - - axis_separators : List[int] - The separators between input axes when generating flattened output axes. - layout: Optional[Union[str, Layout]] The layout of the buffer. @@ -554,8 +536,6 @@ def match_buffer( scope, align, offset_factor, - buffer_type, - axis_separators, _get_layout(layout, shape, scope), ) @@ -810,8 +790,6 @@ def alloc_buffer( scope: str = "global", align: int = -1, offset_factor: int = 0, - buffer_type: str = "default", - axis_separators: list[int] | None = None, layout: str | Layout | None = "default", allocated_addr: int | tuple[int, ...] | None = None, annotations: dict[str, Any] | None = None, @@ -845,10 +823,6 @@ def alloc_buffer( Alignment requirement in bytes. offset_factor : int Offset factor. - buffer_type : str - Buffer type. - axis_separators : Optional[List[int]] - Optional axis separators. layout : Optional[Union[str, Layout]] Optional layout. allocated_addr : Optional[Union[int, Tuple[int, ...]]] @@ -872,8 +846,6 @@ def alloc_buffer( scope=scope, align=align, offset_factor=offset_factor, - buffer_type=buffer_type, - axis_separators=axis_separators, layout=layout, allocated_addr=allocated_addr, buffer_name="", @@ -930,8 +902,6 @@ def sblock_alloc_buffer( scope: str = "global", align: int = -1, offset_factor: int = 0, - buffer_type: str = "default", - axis_separators: list[int] | None = None, layout: str | Layout | None = "default", allocated_addr: int | tuple[int, ...] | None = None, ) -> Buffer: @@ -955,11 +925,6 @@ def sblock_alloc_buffer( The alignment requirement of data pointer in bytes. offset_factor : int The factor of elem_offset field. - buffer_type : str - The buffer type. - axis_separators : List[int] - The separators between input axes when generating flattened output axes. - layout: Optional[Union[str, Layout]] The layout of the buffer. @@ -979,8 +944,6 @@ def sblock_alloc_buffer( strides = [Var(s, "int32") if isinstance(s, str) else s for s in strides] else: strides = [] - if axis_separators is None: - axis_separators = [] if allocated_addr is None: allocated_addr = [] if not isinstance(allocated_addr, list | tuple): @@ -994,8 +957,6 @@ def sblock_alloc_buffer( scope, align, offset_factor, - buffer_type, - axis_separators, _get_layout(layout, shape, scope), allocated_addr, ) @@ -1766,8 +1727,6 @@ def decl_buffer( scope="global", align=0, offset_factor=0, - buffer_type="", - axis_separators=None, layout="default", allocated_addr=None, ) -> Buffer: @@ -1805,12 +1764,6 @@ def decl_buffer( offset_factor : int The factor of elem_offset field. - buffer_type : str - The buffer type. - - axis_separators : List[int] - The separators between input axes when generating flattened output axes. - layout : Layout The layout of the buffer. @@ -1835,8 +1788,6 @@ def decl_buffer( scope, align, offset_factor, - buffer_type, - axis_separators, _get_layout(layout, shape, scope), allocated_addr, ) @@ -2057,8 +2008,6 @@ def decl_scalar(dtype, data, scope, elem_offset=None, byte_offset=None) -> Buffe strides=None, align=-1, offset_factor=0, - buffer_type="default", - axis_separators=None, layout=TileLayout(S[1]), ) assert isinstance(buf, Buffer) diff --git a/python/tvm/tirx/script/builder/tirx.py b/python/tvm/tirx/script/builder/tirx.py index 960875bef430..6f2442ed82ef 100644 --- a/python/tvm/tirx/script/builder/tirx.py +++ b/python/tvm/tirx/script/builder/tirx.py @@ -1561,7 +1561,6 @@ def reshape(buffer: Buffer, shape: list[Expr]): + " are not compatible" ) - assert buffer.buffer_type == 1 return decl_buffer( shape, buffer.dtype, @@ -1572,8 +1571,6 @@ def reshape(buffer: Buffer, shape: list[Expr]): buffer.scope(), buffer.data_alignment, buffer.offset_factor, - "", - buffer.axis_separators, buffer.layout, ) diff --git a/python/tvm/tirx/script/parser/entry.py b/python/tvm/tirx/script/parser/entry.py index 383d2621449e..9778a41794fd 100644 --- a/python/tvm/tirx/script/parser/entry.py +++ b/python/tvm/tirx/script/parser/entry.py @@ -447,8 +447,6 @@ def __call__( scope="global", align=0, offset_factor=0, - buffer_type="", - axis_separators=None, layout="default", ) -> Buffer: return buffer( @@ -461,8 +459,6 @@ def __call__( scope=scope, align=align, offset_factor=offset_factor, - buffer_type=buffer_type, - axis_separators=axis_separators, layout=layout, ) diff --git a/python/tvm/tirx/script/parser/parser.py b/python/tvm/tirx/script/parser/parser.py index 076569a9f3b9..ddbfac29ab18 100644 --- a/python/tvm/tirx/script/parser/parser.py +++ b/python/tvm/tirx/script/parser/parser.py @@ -64,8 +64,6 @@ def slice_buffer_from_region(br: BufferRegion) -> Buffer: buf.scope(), buf.data_alignment, buf.offset_factor, - "", - buf.axis_separators, sliced_layout, ) # Fallback: compute elem_offset for default/no layout @@ -89,8 +87,6 @@ def slice_buffer_from_region(br: BufferRegion) -> Buffer: buf.scope(), buf.data_alignment, buf.offset_factor, - "", - buf.axis_separators, buf.layout, ) diff --git a/src/backend/hexagon/runtime/ops/conv2d.h b/src/backend/hexagon/runtime/ops/conv2d.h index 604eb2353e10..253d59a7dea6 100644 --- a/src/backend/hexagon/runtime/ops/conv2d.h +++ b/src/backend/hexagon/runtime/ops/conv2d.h @@ -149,12 +149,10 @@ inline uintptr_t hwio_at(const DLTensor& f, int y, int x, int i, int o) { * transform layout): * * For uint8_t type - * lambda n, h, w, c: n, h//8, w//8, c//32, AXIS_SEPARATOR, h%8, w%8, c%32 + * lambda n, h, w, c: n, h//8, w//8, c//32, h%8, w%8, c%32 * * For uint16_t type - * lambda n, h, w, c: n, h//8, w//4, c//32, AXIS_SEPARATOR, h%8, (w%4)//2, c%32, w%2 - * - * where AXIS_SEPARATOR represents split up in the physical layout + * lambda n, h, w, c: n, h//8, w//4, c//32, h%8, (w%4)//2, c%32, w%2 * * @param out Pre-allocated output memory pointer * @param inp_flat Flat input data pointer diff --git a/src/backend/trn/transform/lower_trainium_layout.cc b/src/backend/trn/transform/lower_trainium_layout.cc index 3799ec43e76d..5354ff5b34f7 100644 --- a/src/backend/trn/transform/lower_trainium_layout.cc +++ b/src/backend/trn/transform/lower_trainium_layout.cc @@ -142,7 +142,6 @@ class TrainiumLayoutApplier : public arith::IRMutatorWithAnalyzer { writer = flattened.CopyOnWrite(); writer->shape = new_shape; writer->strides = {}; - writer->axis_separators = {}; } else if (is_alloc) { if (auto tile_layout = buf->layout.as(); tile_layout && tile_layout->HasThreadAxis()) { @@ -167,7 +166,6 @@ class TrainiumLayoutApplier : public arith::IRMutatorWithAnalyzer { writer = flattened.CopyOnWrite(); writer->shape = {ana->Simplify(mem_span)}; writer->strides = {}; - writer->axis_separators = {}; } else { flattened = buf.GetFlattenedBuffer(); writer = flattened.CopyOnWrite(); diff --git a/src/relax/op/tensor/manipulate.cc b/src/relax/op/tensor/manipulate.cc index 6421769905b0..7d709afde8db 100644 --- a/src/relax/op/tensor/manipulate.cc +++ b/src/relax/op/tensor/manipulate.cc @@ -699,14 +699,10 @@ TVM_REGISTER_OP("relax.index_tensor") /* relax.layout_transform */ -Expr layout_transform(Expr x, tirx::IndexMap index_map, ffi::Optional pad_value, - ffi::Optional> axis_separators, - ffi::Optional> input_axis_separators) { +Expr layout_transform(Expr x, tirx::IndexMap index_map, ffi::Optional pad_value) { ffi::ObjectPtr attrs = ffi::make_object(); attrs->index_map = std::move(index_map); attrs->pad_value = std::move(pad_value); - attrs->axis_separators = std::move(axis_separators); - attrs->input_axis_separators = std::move(input_axis_separators); static const Op& op = Op::Get("relax.layout_transform"); return Call(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {}); diff --git a/src/relax/op/tensor/manipulate.h b/src/relax/op/tensor/manipulate.h index a818998545c2..2066c7b3b4d4 100644 --- a/src/relax/op/tensor/manipulate.h +++ b/src/relax/op/tensor/manipulate.h @@ -67,14 +67,9 @@ Expr flatten(Expr x); * \param index_map The transformation to apply. * \param pad_value The value used for padding if the transformation results in implicit padding. If * not specified, any value can be used. - * \param axis_separators Array of values to differentiate between input axes - * when generating flattened output axes. - * \param input axis_separators Array of values for input buffer. * \return The transformed result. */ -Expr layout_transform(Expr x, tirx::IndexMap index_map, ffi::Optional pad_value, - ffi::Optional> axis_separators, - ffi::Optional> input_axis_separators = std::nullopt); +Expr layout_transform(Expr x, tirx::IndexMap index_map, ffi::Optional pad_value); /*! * \brief Permutes the dimensions of an array. diff --git a/src/relax/transform/alter_op_impl.cc b/src/relax/transform/alter_op_impl.cc index 96c1d6c232da..19b65a5532eb 100644 --- a/src/relax/transform/alter_op_impl.cc +++ b/src/relax/transform/alter_op_impl.cc @@ -81,18 +81,12 @@ bool IsTransformBijective(const Expr& expr, const IndexMap& transform) { */ class AlterOpImplMutator : public ExprMutator { public: - AlterOpImplMutator( - const IRModule& mod, const ffi::Map& op_impl_map, - const ffi::Map>& op_buffer_transforms_, - const ffi::Map>>>& axis_separators_, - const ffi::Map>>>& - input_axis_separators_) + AlterOpImplMutator(const IRModule& mod, const ffi::Map& op_impl_map, + const ffi::Map>& op_buffer_transforms_) : ExprMutator(mod), mod_(mod), op_impl_map_(op_impl_map), - op_buffer_transforms__(op_buffer_transforms_), - op_buffer_axis_separators__(axis_separators_), - op_buffer_input_axis_separators__(input_axis_separators_) {} + op_buffer_transforms__(op_buffer_transforms_) {} IRModule Run() { for (const auto& gv : mod_->GetGlobalVars()) { @@ -132,13 +126,7 @@ class AlterOpImplMutator : public ExprMutator { const auto& replacement_func = op_impl_map_[op_kind]; ffi::Array buffer_transforms; - ffi::Optional>> axis_separators; - ffi::Optional>> input_axis_separators; if (op_buffer_transforms__.count(op_kind)) buffer_transforms = op_buffer_transforms__[op_kind]; - if (op_buffer_axis_separators__.count(op_kind)) - axis_separators = op_buffer_axis_separators__[op_kind]; - if (op_buffer_input_axis_separators__.count(op_kind)) - input_axis_separators = op_buffer_input_axis_separators__[op_kind]; TVM_FFI_ICHECK(buffer_transforms.empty() || buffer_transforms.size() == replacement_func->params.size()) @@ -150,8 +138,7 @@ class AlterOpImplMutator : public ExprMutator { GlobalVar replacement_gv = GetOrCreateGlobalVarForFunc(replacement_func, op_kind); auto call_tir_inputs_tuple = ffi::GetRef(call->args[1].as()); - Tuple updated_inputs = UpdateInputs(call_tir_inputs_tuple, buffer_transforms, axis_separators, - input_axis_separators); + Tuple updated_inputs = UpdateInputs(call_tir_inputs_tuple, buffer_transforms); TVM_FFI_ICHECK_EQ(call->ty_args.size(), 1) << "call_tir ty_args.size() is expected to be 1"; Type updated_ret_ty = UpdateOutputType(call->ty_args[0], buffer_transforms); @@ -160,8 +147,7 @@ class AlterOpImplMutator : public ExprMutator { call->attrs, {updated_ret_ty})); // Now transform each of the outputs to previous layout. - return TransformOutputs(updated_call, buffer_transforms, call->ty_args[0], axis_separators, - input_axis_separators); + return TransformOutputs(updated_call, buffer_transforms, call->ty_args[0]); } ffi::Array GetTensorTypePerOutput(const Type& output_ty) { @@ -187,9 +173,7 @@ class AlterOpImplMutator : public ExprMutator { return false; } - Expr TransformLayout(const Expr& expr, const IndexMap& index_map, - const ffi::Array& axis_separators, - const ffi::Array& input_axis_separators) { + Expr TransformLayout(const Expr& expr, const IndexMap& index_map) { if (IsScalarConstant(expr) || index_map.get() == nullptr) { return expr; } @@ -198,8 +182,6 @@ class AlterOpImplMutator : public ExprMutator { // identical. The scope of vars used in index map initial indices is local to the op. Not doing // so would confuse the structural equality check. attrs->index_map = DeepCopyIndexMap(index_map); - attrs->axis_separators = std::move(axis_separators); - attrs->input_axis_separators = std::move(input_axis_separators); return Call(Type::Missing(), layout_transform_op_, {expr}, Attrs{std::move(attrs)}, {}); } @@ -246,9 +228,7 @@ class AlterOpImplMutator : public ExprMutator { } Expr TransformLayoutInverse(const Expr& expr, const IndexMap& index_map, - const TensorType& old_tensor_ty, - const ffi::Array& axis_separator, - const ffi::Array& input_axis_separator) { + const TensorType& old_tensor_ty) { if (IsScalarConstant(expr) || index_map.get() == nullptr) { return expr; } @@ -259,10 +239,9 @@ class AlterOpImplMutator : public ExprMutator { index_map.NonSurjectiveInverse(initial_ranges, analyzer); if (tirx::is_zero(padding_predicate)) { - return TransformLayout(expr, inverse_index_map, axis_separator, input_axis_separator); + return TransformLayout(expr, inverse_index_map); } else { - auto padded_expr = builder_->Normalize( - TransformLayout(expr, inverse_index_map, axis_separator, input_axis_separator)); + auto padded_expr = builder_->Normalize(TransformLayout(expr, inverse_index_map)); const auto& tensor_ty = padded_expr->ty.as_or_throw(); GlobalVar gv_remove_pad = GetOrCreateRemovePadOp(old_shape, tensor_ty->dtype.value()->dtype); @@ -294,27 +273,14 @@ class AlterOpImplMutator : public ExprMutator { /*! * \brief Updates call inputs with layout transformed inputs */ - Tuple UpdateInputs(const Tuple& inputs, const ffi::Array& transforms, - const ffi::Optional>>& axis_separators, - const ffi::Optional>>& input_axis_separators) { + Tuple UpdateInputs(const Tuple& inputs, const ffi::Array& transforms) { if (transforms.empty()) return inputs; ffi::Array updated_inputs; int index = 0; for (const auto& input : inputs->fields) { - ffi::Array axis_separator; - ffi::Array input_axis_separator; - if (axis_separators.has_value()) { - ffi::Array> axis_separators_value = axis_separators.value(); - axis_separator = axis_separators_value[index]; - } - if (input_axis_separators.has_value()) { - ffi::Array> input_axis_separators_value = input_axis_separators.value(); - input_axis_separator = input_axis_separators_value[index]; - } auto transform = transforms[index++]; - updated_inputs.push_back( - TransformLayout(input, transform, axis_separator, input_axis_separator)); + updated_inputs.push_back(TransformLayout(input, transform)); } return Tuple(updated_inputs); } @@ -359,15 +325,12 @@ class AlterOpImplMutator : public ExprMutator { return TensorType(ShapeExpr(new_shape), tensor_ty->dtype); } - Expr TransformOutputs( - const Expr& expr, const ffi::Array& buffer_transforms, const Type& old_ty, - const ffi::Optional>>& axis_separators, - const ffi::Optional>>& input_axis_separators) { + Expr TransformOutputs(const Expr& expr, const ffi::Array& buffer_transforms, + const Type& old_ty) { if (buffer_transforms.empty()) return expr; ffi::Array old_output_ty = GetTensorTypePerOutput(old_ty); - ffi::Array axis_sep, input_axis_sep; size_t num_outputs = old_output_ty.size(); if (num_outputs == 0) return expr; @@ -375,15 +338,7 @@ class AlterOpImplMutator : public ExprMutator { // If there is a single output, return the transformed output. if (num_outputs == 1) { IndexMap output_map = buffer_transforms[first_output_index]; - if (axis_separators.has_value()) { - ffi::Array> axis_separators_value = axis_separators.value(); - axis_sep = axis_separators_value[first_output_index]; - } - if (input_axis_separators.has_value()) { - ffi::Array> input_axis_separators_value = input_axis_separators.value(); - input_axis_sep = input_axis_separators_value[first_output_index]; - } - return TransformLayoutInverse(expr, output_map, old_output_ty[0], axis_sep, input_axis_sep); + return TransformLayoutInverse(expr, output_map, old_output_ty[0]); } // In case of more than one output, we would have to get each item of the output tuple, @@ -391,17 +346,8 @@ class AlterOpImplMutator : public ExprMutator { ffi::Array transformed_outputs; for (size_t i = 0; i + first_output_index < buffer_transforms.size(); ++i) { const auto& output_map = buffer_transforms[i + first_output_index]; - if (axis_separators.has_value()) { - ffi::Array> axis_separators_value = axis_separators.value(); - axis_sep = axis_separators_value[i + first_output_index]; - } - if (input_axis_separators.has_value()) { - ffi::Array> input_axis_separators_value = input_axis_separators.value(); - input_axis_sep = input_axis_separators_value[i + first_output_index]; - } auto output = builder_->Normalize(TupleGetItem(expr, static_cast(i))); - transformed_outputs.push_back( - TransformLayoutInverse(output, output_map, old_output_ty[i], axis_sep, input_axis_sep)); + transformed_outputs.push_back(TransformLayoutInverse(output, output_map, old_output_ty[i])); } return Tuple(transformed_outputs); } @@ -417,12 +363,6 @@ class AlterOpImplMutator : public ExprMutator { const ffi::Map& op_impl_map_; /*! \brief Map from kOperatorName attribute to the layout transforms on i/o buffers */ const ffi::Map>& op_buffer_transforms__; - /*! \brief Map from kOperatorName attribute to the axis separatos on i/o buffers */ - const ffi::Map>>>& - op_buffer_axis_separators__; - /*! \brief Map from kOperatorName attribute to the input axis separatos */ - const ffi::Map>>>& - op_buffer_input_axis_separators__; const Op& call_tir_op_ = Op::Get("relax.call_tir"); const Op& layout_transform_op_ = Op::Get("relax.layout_transform"); @@ -430,16 +370,10 @@ class AlterOpImplMutator : public ExprMutator { namespace transform { -Pass AlterOpImpl( - const ffi::Map& op_impl_map, - const ffi::Map>& op_buffer_transforms_, - const ffi::Map>>>& axis_separators_, - const ffi::Map>>>& - input_axis_separators_) { +Pass AlterOpImpl(const ffi::Map& op_impl_map, + const ffi::Map>& op_buffer_transforms_) { auto pass_func = [=](IRModule mod, PassContext pc) { - return AlterOpImplMutator(mod, op_impl_map, op_buffer_transforms_, axis_separators_, - input_axis_separators_) - .Run(); + return AlterOpImplMutator(mod, op_impl_map, op_buffer_transforms_).Run(); }; return CreateModulePass(/*pass_function=*/pass_func, // /*opt_level=*/0, // diff --git a/src/relax/transform/convert_layout.cc b/src/relax/transform/convert_layout.cc index 28c30b708d57..f33f70e8ac06 100644 --- a/src/relax/transform/convert_layout.cc +++ b/src/relax/transform/convert_layout.cc @@ -131,11 +131,7 @@ class LayoutConvertMutator : public ExprMutator { auto index_map = LayoutIndexMap(from.LeafValue()->layout.ndim(), from.LeafValue()->layout, to.LeafValue()->layout); ffi::ObjectPtr attrs = ffi::make_object(); - ffi::Array axis_separator; - ffi::Array input_axis_separator; attrs->index_map = ffi::FromJSONGraph(ffi::ToJSONGraph(index_map)).as_or_throw(); - attrs->axis_separators = std::move(axis_separator); - attrs->input_axis_separators = std::move(input_axis_separator); const Op& layout_transform_op_ = Op::Get("relax.layout_transform"); auto ret_expr = Call(Type::Missing(), layout_transform_op_, {expr}, Attrs{std::move(attrs)}, {}); diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc index 3a653ea7534a..5db564c9bf8a 100644 --- a/src/relax/transform/fuse_tir.cc +++ b/src/relax/transform/fuse_tir.cc @@ -941,8 +941,7 @@ class FusedTIRConstructor : public ExprVisitor { tirx::Buffer buffer; if (tir_buffer_param.has_value()) { buffer = tirx::decl_buffer(shape_expr->values, dtype, name_hint, - tir_buffer_param.value().scope(), - tir_buffer_param.value()->axis_separators); + tir_buffer_param.value().scope()); } else { buffer = tirx::decl_buffer(shape_expr->values, dtype, name_hint); } diff --git a/src/s_tir/schedule/concrete_schedule.cc b/src/s_tir/schedule/concrete_schedule.cc index 0d2fa9c2614a..2443d7966e7a 100644 --- a/src/s_tir/schedule/concrete_schedule.cc +++ b/src/s_tir/schedule/concrete_schedule.cc @@ -1049,16 +1049,6 @@ void ConcreteScheduleNode::TransformBlockLayout(const SBlockRV& block_rv, TVM_TIR_SCHEDULE_END("transform_block_layout", this->error_render_level_); } -void ConcreteScheduleNode::SetAxisSeparator(const SBlockRV& block_rv, int buffer_index, - BufferIndexType buffer_index_type, - const ffi::Array& axis_separators) { - TVM_TIR_SCHEDULE_BEGIN(); - s_tir::SetAxisSeparator(state_, this->GetSRef(block_rv), buffer_index, buffer_index_type, - axis_separators); - TVM_TIR_SCHEDULE_END("set-axis-separator", this->error_render_level_); - this->state_->DebugVerify(); -} - /******** Schedule: Padding ********/ SBlockRV ConcreteScheduleNode::DecomposePadding(const SBlockRV& block_rv, const LoopRV& loop_rv) { diff --git a/src/s_tir/schedule/concrete_schedule.h b/src/s_tir/schedule/concrete_schedule.h index 1372aa9dbe8f..ab78f8af923d 100644 --- a/src/s_tir/schedule/concrete_schedule.h +++ b/src/s_tir/schedule/concrete_schedule.h @@ -181,9 +181,6 @@ class ConcreteScheduleNode : public ScheduleNode { const ffi::Optional& pad_value, bool assume_injective_transform = false) override; void TransformBlockLayout(const SBlockRV& block_rv, const IndexMap& index_map) override; - void SetAxisSeparator(const SBlockRV& block_rv, int buffer_index, - BufferIndexType buffer_index_type, - const ffi::Array& axis_separators) override; /******** Schedule: Padding decomposition ********/ SBlockRV DecomposePadding(const SBlockRV& block_rv, const LoopRV& loop_rv) override; /******** Schedule: Buffer transformation ********/ diff --git a/src/s_tir/schedule/primitive.h b/src/s_tir/schedule/primitive.h index aa48a5fb9c3d..0c2c02af86e3 100644 --- a/src/s_tir/schedule/primitive.h +++ b/src/s_tir/schedule/primitive.h @@ -587,18 +587,6 @@ TVM_DLL void SetScope(ScheduleState self, const StmtSRef& block_sref, int buffer */ TVM_DLL void UnsafeSetDType(ScheduleState self, const StmtSRef& block_sref, int buffer_index, const ffi::String& dtype); -/*! - * \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 - */ -TVM_DLL void SetAxisSeparator(ScheduleState self, const StmtSRef& block_sref, int buffer_index, - BufferIndexType buffer_index_type, - const ffi::Array& axis_separators); - /******** Schedule: Blockize & Tensorize ********/ /*! diff --git a/src/s_tir/schedule/primitive/cache_index.cc b/src/s_tir/schedule/primitive/cache_index.cc index e0c075544f33..d3775e0ca330 100644 --- a/src/s_tir/schedule/primitive/cache_index.cc +++ b/src/s_tir/schedule/primitive/cache_index.cc @@ -269,8 +269,8 @@ ffi::Array MakeIndexCacheStage(IndexInfo* info, const ffi::String& stora buffer_shape.push_back( arith::EvalSet(info->var_binding.at(it), arith::AsIntSet(info->range_map)).max() + 1); } - info->cache_buffer.push_back(Buffer(index_buffer_var, data_ty, buffer_shape, {1}, {0}, - index_buffer_var->name, 0, 0, kDefault)); + info->cache_buffer.push_back( + Buffer(index_buffer_var, data_ty, buffer_shape, {1}, {0}, index_buffer_var->name, 0, 0)); // Create loop vars and block vars' binding_value std::vector loop_vars; diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc b/src/s_tir/schedule/primitive/layout_transformation.cc index 9ebdd599a2d0..97c0f00a759e 100644 --- a/src/s_tir/schedule/primitive/layout_transformation.cc +++ b/src/s_tir/schedule/primitive/layout_transformation.cc @@ -1496,89 +1496,6 @@ void TransformBlockLayout(ScheduleState self, const StmtSRef& block_sref, self->Replace(scope_sref, body, {{block, new_block}}); } -class BufferAxisSeparatorMutator : private ReplaceBufferMutator { - public: - static SBlock Mutate(const SBlock& scope_block, const Buffer& old_buffer, Buffer new_buffer, - ffi::Map* block_sref_reuse) { - BufferAxisSeparatorMutator mutator(old_buffer, std::move(new_buffer), block_sref_reuse); - return mutator.VisitStmt(scope_block).as_or_throw(); - } - - private: - BufferAxisSeparatorMutator(const Buffer& old_buffer, Buffer new_buffer, - ffi::Map* block_sref_reuse) - : ReplaceBufferMutator(old_buffer, new_buffer, block_sref_reuse) {} - - MatchBufferRegion VisitMatchBufferRegion(const MatchBufferRegion& match_buffer) final { - auto it = buffer_var_map_.find(match_buffer->source->buffer->data.get()); - if (it != buffer_var_map_.end()) { - const Buffer& new_source_buffer = it->second; - Buffer new_target_buffer = match_buffer->buffer; - - if (new_target_buffer->shape.size() == new_source_buffer->shape.size()) { - new_target_buffer.CopyOnWrite()->axis_separators = new_source_buffer->axis_separators; - } else { - new_target_buffer.CopyOnWrite()->axis_separators = - ffi::Array(new_source_buffer->axis_separators.size(), IntImm::Int32(0)); - LOG(WARNING) << "Buffer view " << new_target_buffer - << " has different dimensionality than backing buffer " << new_source_buffer - << ". The `axis_separators` for " << new_target_buffer << "." - << "`axis_separators` for the view might be incorrect."; - } - buffer_var_map_[new_target_buffer->data.get()] = new_target_buffer; - return MatchBufferRegion(new_target_buffer, - BufferRegion(new_source_buffer, match_buffer->source->region)); - } - return match_buffer; - } -}; - -void SetAxisSeparator(ScheduleState self, const StmtSRef& block_sref, int buffer_index, - BufferIndexType buffer_index_type, - const ffi::Array& axis_separators) { - const SBlockNode* block_ptr = TVM_SREF_TO_SBLOCK(block_sref); - Buffer old_buffer = - GetNthAccessBuffer(self, ffi::GetRef(block_ptr), buffer_index, buffer_index_type); - auto [defining_site_sref, is_alloc] = GetBufferDefiningSite(block_sref, old_buffer); - if (defining_site_sref.has_value() && !is_alloc) { - throw BufferIsSubregionError(self->mod, old_buffer); - } - - StmtSRef scope_sref = defining_site_sref.has_value() - ? defining_site_sref.value() - : GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false); - const SBlockNode* scope_block = TVM_SREF_TO_SBLOCK(scope_sref); - - // Step 1: Check and update axis_separators of the buffer. - Buffer new_buffer = old_buffer; - new_buffer.CopyOnWrite()->axis_separators = axis_separators; - - ffi::Map block_sref_reuse; - - // Step 2: Rewrite alloc_buffer of the block or buffer_map of the PrimFunc. - SBlock new_scope_block = BufferAxisSeparatorMutator::Mutate( - ffi::GetRef(scope_block), old_buffer, new_buffer, &block_sref_reuse); - if (!defining_site_sref.has_value()) { - // mutate buffer_map of the PrimFunc - GlobalVar g_var; - GetRootPrimFunc(self->mod, scope_block, &g_var); - IRModuleNode* new_mod = self->mod.CopyOnWrite(); - ffi::MapObj* new_map = new_mod->functions.CopyOnWrite(); - PrimFunc ref_new_func = std::move(new_map->at(g_var)).as_or_throw(); - PrimFuncNode* new_func = ref_new_func.CopyOnWrite(); - ffi::MapObj* new_buffer_map = new_func->buffer_map.CopyOnWrite(); - for (auto it = new_buffer_map->begin(); it != new_buffer_map->end(); ++it) { - if ((*it).second.same_as(old_buffer)) { - (*it).second = new_buffer; - } - } - new_map->at(g_var) = std::move(ref_new_func); - } - - // Step 4: Replace the scope block with the new block - self->Replace(scope_sref, new_scope_block, block_sref_reuse); -} - /******** InstructionKind Registration ********/ struct TransformLayoutTraits : public UnpackedInstTraits { @@ -1696,45 +1613,8 @@ struct TransformBlockLayoutTraits : public UnpackedInstTraits { - static constexpr const char* kName = "SetAxisSeparator"; - static constexpr bool kIsPure = false; - - private: - static constexpr size_t kNumInputs = 1; - static constexpr size_t kNumAttrs = 3; - static constexpr size_t kNumDecisions = 0; - - static void UnpackedApplyToSchedule(Schedule sch, SBlockRV block_rv, IntImm buffer_index, - IntImm buffer_index_type, - ffi::Array axis_separators) { - return sch->SetAxisSeparator(block_rv, buffer_index->value, - static_cast(buffer_index_type->value), - axis_separators); - } - - static ffi::String UnpackedAsPython(ffi::Array outputs, ffi::String block_rv, - IntImm buffer_index, IntImm buffer_index_type, - ffi::Array axis_separators) { - PythonAPICall py("set_axis_separator"); - py.Input("block", block_rv); - - std::ostringstream os; - os << "(\"" << BufferIndexType2Str(static_cast(buffer_index_type->value)) - << "\", " << buffer_index << ")"; - py.Input("buffer", os.str()); - - py.Input("axis_separators", axis_separators); - return py.Str(); - } - - template - friend struct ::tvm::s_tir::UnpackedInstTraits; -}; - TVM_REGISTER_INST_KIND_TRAITS(TransformLayoutTraits); TVM_REGISTER_INST_KIND_TRAITS(TransformBlockLayoutTraits); -TVM_REGISTER_INST_KIND_TRAITS(SetAxisSeparatorTraits); } // namespace s_tir } // namespace tvm diff --git a/src/s_tir/schedule/schedule.cc b/src/s_tir/schedule/schedule.cc index a4713a09fcb2..e8e666672f08 100644 --- a/src/s_tir/schedule/schedule.cc +++ b/src/s_tir/schedule/schedule.cc @@ -323,14 +323,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { index_map, pad_value, assume_injective_transform); }) .def_method("s_tir.schedule.ScheduleTransformBlockLayout", - &ScheduleNode::TransformBlockLayout) - .def("s_tir.schedule.ScheduleSetAxisSeparator", - [](Schedule self, const SBlockRV& block_rv, int buffer_index, int buffer_index_type, - const ffi::Array& axis_separators) { - return self->SetAxisSeparator(block_rv, buffer_index, - static_cast(buffer_index_type), - axis_separators); - }); + &ScheduleNode::TransformBlockLayout); } /******** (FFI) Padding decomposition ********/ diff --git a/src/s_tir/schedule/traced_schedule.cc b/src/s_tir/schedule/traced_schedule.cc index 4845f88c6b20..646269abfb7e 100644 --- a/src/s_tir/schedule/traced_schedule.cc +++ b/src/s_tir/schedule/traced_schedule.cc @@ -737,21 +737,6 @@ void TracedScheduleNode::TransformBlockLayout(const SBlockRV& block_rv, const In /*outputs=*/{})); } -void TracedScheduleNode::SetAxisSeparator(const SBlockRV& block_rv, int buffer_index, - BufferIndexType buffer_index_type, - const ffi::Array& axis_separators) { - ConcreteScheduleNode::SetAxisSeparator(block_rv, buffer_index, buffer_index_type, - axis_separators); - static const InstructionKind& kind = InstructionKind::Get("SetAxisSeparator"); - trace_->Append(/*inst=*/Instruction( - /*kind=*/kind, - /*inputs=*/{block_rv}, - /*attrs=*/ - {IntImm::Int32(buffer_index), IntImm::Int32(static_cast(buffer_index_type)), - axis_separators}, - /*outputs=*/{})); -} - /******** Schedule: Padding ********/ SBlockRV TracedScheduleNode::DecomposePadding(const SBlockRV& block_rv, const LoopRV& loop_rv) { SBlockRV new_block = ConcreteScheduleNode::DecomposePadding(block_rv, loop_rv); diff --git a/src/s_tir/schedule/traced_schedule.h b/src/s_tir/schedule/traced_schedule.h index 89daa8f4f6ae..cd8cb8b6cdeb 100644 --- a/src/s_tir/schedule/traced_schedule.h +++ b/src/s_tir/schedule/traced_schedule.h @@ -136,9 +136,6 @@ class TracedScheduleNode : public ConcreteScheduleNode { const ffi::Optional& pad_value, bool assume_injective_transform) override; void TransformBlockLayout(const SBlockRV& block_rv, const IndexMap& index_map) override; - void SetAxisSeparator(const SBlockRV& block_rv, int buffer_index, - BufferIndexType buffer_index_type, - const ffi::Array& axis_separators) final; /******** Schedule: Padding ********/ SBlockRV DecomposePadding(const SBlockRV& block_rv, const LoopRV& loop_rv) final; void PadEinsum(const SBlockRV& block_rv, const ffi::Array& padding) final; diff --git a/src/s_tir/transform/lower_cross_thread_reduction.cc b/src/s_tir/transform/lower_cross_thread_reduction.cc index 23011943b4ec..994f461c2007 100644 --- a/src/s_tir/transform/lower_cross_thread_reduction.cc +++ b/src/s_tir/transform/lower_cross_thread_reduction.cc @@ -154,8 +154,7 @@ ffi::Array MakeScratchpads(const ffi::Array& reduction_buffers, /*elem_offset=*/PrimExpr{nullptr}, /*name=*/name, /*data_alignment=*/0, - /*offset_factor=*/0, - /*buffer_type=*/kDefault)); + /*offset_factor=*/0)); } return new_buffers; } diff --git a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc index 04578aedfdc7..07bddb3f1295 100644 --- a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc +++ b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc @@ -142,8 +142,7 @@ Stmt RewriteWmmaLoad(Stmt stmt) { /*elem_offset=*/PrimVar("src_elem_offset", int32_ty), /*name=*/"src", /*data_alignment=*/64, - /*offset_factor=*/16, - /*buffer_type=*/kDefault); + /*offset_factor=*/16); Buffer new_tgt_buffer( /*data=*/Var("tgt", PointerType(dtype_ty, tgt_buffer.scope())), /*dtype=*/dtype, @@ -152,8 +151,7 @@ Stmt RewriteWmmaLoad(Stmt stmt) { /*elem_offset=*/PrimVar("tgt_elem_offset", int32_ty), /*name=*/"tgt", /*data_alignment=*/64, - /*offset_factor=*/16, - /*buffer_type=*/kDefault); + /*offset_factor=*/16); ffi::Array read_region = RelaxIndices(buf_load->indices, src_buffer->shape, var_dom); ffi::Array write_region = RelaxIndices(buf_store->indices, tgt_buffer->shape, var_dom); static const Op& tvm_load_matrix_sync_op = Op::Get("tirx.tvm_load_matrix_sync"); @@ -253,8 +251,7 @@ Stmt RewriteWmmaStore(Stmt stmt) { /*elem_offset=*/PrimVar("src_elem_offset", int32_ty), /*name=*/"src", /*data_alignment=*/64, - /*offset_factor=*/16, - /*buffer_type=*/kDefault); + /*offset_factor=*/16); Buffer new_tgt_buffer(/*data=*/Var("tgt", PointerType(dtype_ty, tgt_buffer.scope())), /*dtype=*/dtype, /*shape=*/{IntImm::Int32(16), IntImm::Int32(16)}, @@ -262,8 +259,7 @@ Stmt RewriteWmmaStore(Stmt stmt) { /*elem_offset=*/PrimVar("tgt_elem_offset", int32_ty), /*name=*/"tgt", /*data_alignment=*/64, - /*offset_factor=*/16, - /*buffer_type=*/kDefault); + /*offset_factor=*/16); ffi::Array read_region = RelaxIndices(buf_load->indices, src_buffer->shape, var_dom); ffi::Array write_region = RelaxIndices(buf_store->indices, tgt_buffer->shape, var_dom); @@ -481,8 +477,7 @@ Stmt RewriteMmaStore(Stmt stmt) { /*elem_offset=*/PrimVar("src_elem_offset", int32_ty), /*name=*/"src", /*data_alignment=*/64, - /*offset_factor=*/8, - /*buffer_type=*/kDefault); + /*offset_factor=*/8); Buffer new_tgt_buffer(/*data=*/Var("tgt", PointerType(dtype_ty, tgt_buffer.scope())), /*dtype=*/dtype, /*shape=*/{IntImm::Int32(8), IntImm::Int32(8)}, @@ -490,8 +485,7 @@ Stmt RewriteMmaStore(Stmt stmt) { /*elem_offset=*/PrimVar("tgt_elem_offset", int32_ty), /*name=*/"tgt", /*data_alignment=*/64, - /*offset_factor=*/8, - /*buffer_type=*/kDefault); + /*offset_factor=*/8); // Step 3.2. Generate new r/w region ffi::Array read_region = RelaxIndices(buf_load->indices, src_buffer->shape, var_dom); diff --git a/src/s_tir/transform/merge_shared_memory_allocations.cc b/src/s_tir/transform/merge_shared_memory_allocations.cc index 952468f6322c..6c6956d22b39 100644 --- a/src/s_tir/transform/merge_shared_memory_allocations.cc +++ b/src/s_tir/transform/merge_shared_memory_allocations.cc @@ -394,7 +394,7 @@ class SharedMemoryRewriter : public StmtExprMutator { // 7. Wrap with the merged-buffer AllocBuffer. Buffer merged_buf(scope.merged_buf_var, PrimType::UInt(8), {scope.merged_alloc_size}, {}, - PrimExpr(), scope.merged_buf_var->name, 0, 0, BufferType::kDefault); + PrimExpr(), scope.merged_buf_var->name, 0, 0); ffi::Map annotations; if (scope.has_volatile_alloc) { annotations.Set(tirx::attr::kVolatile, true); diff --git a/src/s_tir/transform/transform_mma_buffer_layout.cc b/src/s_tir/transform/transform_mma_buffer_layout.cc index 002e4a082749..ee60179657ea 100644 --- a/src/s_tir/transform/transform_mma_buffer_layout.cc +++ b/src/s_tir/transform/transform_mma_buffer_layout.cc @@ -70,8 +70,7 @@ class MmaBufferLayoutTransformer : public StmtExprMutator { new_shape.insert(new_shape.end(), {IntImm::Int32(dim0->value / 16), IntImm::Int32(dim1->value / 8), 2, 2}); - Buffer new_buffer = decl_buffer(std::move(new_shape), buffer->dtype, buffer->name, "local", - buffer->axis_separators); + Buffer new_buffer = decl_buffer(std::move(new_shape), buffer->dtype, buffer->name, "local"); this->buffer_map_.insert({buffer, new_buffer}); this->buffer_var_map_.insert({buffer->data, new_buffer->data}); return new_buffer; @@ -92,8 +91,7 @@ class MmaBufferLayoutTransformer : public StmtExprMutator { new_shape.insert(new_shape.end(), {IntImm::Int32(dim0->value / 32), IntImm::Int32(dim1->value / 8), 4, 2}); - Buffer new_buffer = decl_buffer(std::move(new_shape), buffer->dtype, buffer->name, "local", - buffer->axis_separators); + Buffer new_buffer = decl_buffer(std::move(new_shape), buffer->dtype, buffer->name, "local"); this->buffer_map_.insert({buffer, new_buffer}); this->buffer_var_map_.insert({buffer->data, new_buffer->data}); return new_buffer; @@ -114,8 +112,7 @@ class MmaBufferLayoutTransformer : public StmtExprMutator { new_shape.insert(new_shape.end(), {IntImm::Int32(dim0->value / 8), IntImm::Int32(dim1->value / 32), 1, 8}); - Buffer new_buffer = decl_buffer(std::move(new_shape), buffer->dtype, buffer->name, "local", - buffer->axis_separators); + Buffer new_buffer = decl_buffer(std::move(new_shape), buffer->dtype, buffer->name, "local"); this->buffer_map_.insert({buffer, new_buffer}); this->buffer_var_map_.insert({buffer->data, new_buffer->data}); return new_buffer; diff --git a/src/tirx/ir/buffer.cc b/src/tirx/ir/buffer.cc index 3d52c94d8647..6deea08ac617 100644 --- a/src/tirx/ir/buffer.cc +++ b/src/tirx/ir/buffer.cc @@ -52,12 +52,10 @@ ffi::Array SimplifyArray(arith::AnalyzerObj* ana, ffi::Array } Buffer decl_buffer(ffi::Array shape, PrimType dtype, ffi::String name, - ffi::String storage_scope, ffi::Optional> axis_separators, - Span span) { + ffi::String storage_scope, Span span) { PrimType storage_type = (dtype == PrimType::Bool() ? PrimType::Int(8) : dtype); return Buffer(Var(name, PointerType(storage_type, storage_scope), span), dtype, shape, - ffi::Array(), PrimExpr(), name, 0, 0, kDefault, - axis_separators.value_or(ffi::Array()), span, std::nullopt, + ffi::Array(), PrimExpr(), name, 0, 0, span, std::nullopt, ffi::Array()); } @@ -272,33 +270,10 @@ ffi::Array BufferNode::ElemOffset(ffi::Array input_indices, << "the index's dimensionality must match the dimensionality of the index given."; } - // TODO(Lunderberg): Better handling for cases where there is more - // than one output index. Currently, this only allows elem_offset - // to be non-zero for flat memory allocations. - ffi::Array elem_offsets = {}; - if (elem_offset.defined() && !is_zero(elem_offset) && !inner) { - elem_offsets = {elem_offset}; - } - - if (elem_offsets.size()) { - TVM_FFI_ICHECK_EQ(elem_offsets.size(), axis_separators.size() + 1) - << "If element offsets are defined, " - << "there must be one element offset for each output index."; - } - - ffi::Array output_indices(axis_separators.size() + 1, 0); - - size_t current_output_axis = 0; - + PrimExpr output_index = 0; arith::Analyzer ana; for (size_t i = 0; i < input_indices.size(); i++) { - if ((current_output_axis < axis_separators.size()) && - (i == size_t(axis_separators[current_output_axis]->value))) { - current_output_axis++; - } - - PrimExpr output_index = output_indices[current_output_axis]; if (strides.size()) { output_index = output_index + input_indices[i] * strides[i]; } else { @@ -308,17 +283,13 @@ ffi::Array BufferNode::ElemOffset(ffi::Array input_indices, if (i > 0) { output_index = MergeMulMod(ana.get(), output_index); } - - output_indices.Set(current_output_axis, output_index); } - if (elem_offsets.size()) { - for (size_t i = 0; i < output_indices.size(); i++) { - output_indices.Set(i, output_indices[i] + elem_offsets[i]); - } + if (elem_offset.defined() && !is_zero(elem_offset) && !inner) { + output_index = output_index + elem_offset; } - return SimplifyArray(ana.get(), output_indices); + return SimplifyArray(ana.get(), {output_index}); } inline ffi::Array BufferOffset(const BufferNode* n, ffi::Array index, @@ -342,78 +313,29 @@ inline ffi::Array BufferOffset(const BufferNode* n, ffi::Array& axis_separators, size_t buffer_dim) { - // These checks ensure that all output axes contain at least one - // input axis. - for (size_t i = 0; (i + 1) < axis_separators.size(); i++) { - auto sep = axis_separators[i]->value; - auto next_sep = axis_separators[i + 1]->value; - TVM_FFI_CHECK_LE(sep, next_sep, ValueError) - << "ValueError: " - << "Axis separators must be in increasing order, " - << "but axis_separators[" << i << "] = " << sep - << " is greater than or equal to axis_separators[" << (i + 1) << "] = " << next_sep << "."; - } - if (axis_separators.size()) { - auto first_sep = axis_separators[0]->value; - TVM_FFI_CHECK_GE(first_sep, 0, ValueError) << "ValueError: " - << "All axis separators must be non-negative. " - << "However, the axis_separators[0] = " << first_sep; - auto last_sep = axis_separators[axis_separators.size() - 1]->value; - TVM_FFI_CHECK_LE(last_sep, buffer_dim, ValueError) - << "ValueError: " - << "All axis separators must be within the range " - << "0 <= sep <= buffer_dim. " - << "However, the last axis_separators[" << (axis_separators.size() - 1) - << "] = " << last_sep << " is greater than the buffer's dimensionality of " << buffer_dim; - } -} - Buffer Buffer::GetFlattenedBuffer() const { auto self = operator->(); - ValidateAxisSeparators(self->axis_separators, self->shape.size()); - - ffi::Array output_shape; + ffi::Array output_shape{1}; if (self->strides.size()) { - // If strides are defined, then the extent of each flattened - // buffer is the stride*size for the first input axis used for - // each output axis. + // If strides are defined, the flattened extent is the span of the + // outermost input axis. TVM_FFI_ICHECK_EQ(self->shape.size(), self->strides.size()); - output_shape.push_back(self->strides[0] * self->shape[0]); - for (const auto& sep : self->axis_separators) { - output_shape.push_back(self->strides[sep->value] * self->shape[sep->value]); - } - + output_shape.Set(0, self->strides[0] * self->shape[0]); } else { - // Otherwise, the extent of each flattened buffer is the product - // of the extents of each input axis used to generate that output - // axis. This also "flattens" rank-0 tensors to a rank-1 buffer - // of shape [1]. - output_shape = ffi::Array(self->axis_separators.size() + 1, 1); - size_t current_output_index = 0; + // Otherwise, the flattened extent is the product of the input extents. + // This also flattens rank-0 tensors to a rank-1 buffer of shape [1]. for (size_t i = 0; i < self->shape.size(); i++) { - if ((current_output_index < self->axis_separators.size()) && - (i == size_t(self->axis_separators[current_output_index]->value))) { - current_output_index += 1; - } - output_shape.Set(current_output_index, output_shape[current_output_index] * self->shape[i]); + output_shape.Set(0, output_shape[0] * self->shape[i]); } } - // The axis_separators for the output buffer. - ffi::Array output_axis_separators; - for (size_t i = 0; i < self->axis_separators.size(); i++) { - output_axis_separators.push_back(IntImm(self->axis_separators[i].ty(), i + 1)); - } - if (output_shape.size() == self->shape.size() && self->strides.empty()) { return *this; } else { Buffer output = *this; auto writer = output.CopyOnWrite(); writer->shape = output_shape; - writer->axis_separators = output_axis_separators; writer->strides = {}; // Keep `layout` in sync with `shape`. The old layout describes the // pre-flatten N-D shape (e.g. `S[(16,16):(16,1)]`); after collapsing @@ -531,18 +453,7 @@ Buffer Buffer::MakeSlice(ffi::Array begins, ffi::Array exten } } Buffer slice(n->data, n->dtype, extents, strides, elem_offset[0], n->name + "_slice", - n->data_alignment, 0, n->buffer_type); - - // Buffer must be constructed with a singular element offset which means there is no - // support for n-dimensional buffers where n > 1. Insert sentinel value for - // TVMFFIABIBuilder::BindBuffer to state that any usage of element offset is invalid - // in this case. This allows for construction of a Buffer with multiple element offsets - // but disallows any usage of those element offsets. See PR #10816 for discussion on - // supporting multiple element offsets in TIR Buffer. - // TODO(Lunderberg): Remove if/when TIR supports multiple element offsets in TIR Buffer - if (elem_offset.size() != 1) { - slice.CopyOnWrite()->elem_offset = PrimExpr(); - } + n->data_alignment, 0); return slice; } @@ -587,8 +498,7 @@ Expr Buffer::access_ptr(int access_mask, PointerType ptr_type, int content_lanes Buffer::Buffer(Var data, PrimType dtype, ffi::Array shape, ffi::Array strides, PrimExpr elem_offset, ffi::String name, int data_alignment, int offset_factor, - BufferType buffer_type, ffi::Array axis_separators, Span span, - ffi::Optional layout, ffi::Array allocated_addr) { + Span span, ffi::Optional layout, ffi::Array allocated_addr) { DLDataType storage_dtype = dtype->dtype; // specially handle bool if (storage_dtype == DLDataType{kDLBool, 8, 1}) { @@ -610,15 +520,12 @@ Buffer::Buffer(Var data, PrimType dtype, ffi::Array shape, ffi::Array< TVM_FFI_ICHECK(data->ty.as()->element_type.as()) << "Variable " << data->name << " does not point to a primitive."; - ValidateAxisSeparators(axis_separators, shape.size()); - auto n = ffi::make_object(); n->data = std::move(data); n->dtype = dtype; n->shape = std::move(shape); n->strides = std::move(strides); - n->axis_separators = std::move(axis_separators); n->name = std::move(name); if (!elem_offset.defined()) { elem_offset = IntImm(PrimType(n->DefaultIndexType()), 0); @@ -632,12 +539,6 @@ Buffer::Buffer(Var data, PrimType dtype, ffi::Array shape, ffi::Array< n->elem_offset = std::move(elem_offset); n->data_alignment = data_alignment; n->offset_factor = offset_factor; - n->buffer_type = buffer_type; - if (n->buffer_type == kAutoBroadcast && n->shape.size() > 0 && n->strides.empty()) { - for (size_t i = 0; i < n->shape.size(); ++i) { - n->strides.push_back(PrimVar("stride", n->shape[i].ty())); - } - } n->span = std::move(span); // `layout=nullopt` is a meaningful sentinel: it tells the printer that the // user opted out of layout sugar (e.g., the `local_scalar` shorthand keys @@ -649,20 +550,10 @@ Buffer::Buffer(Var data, PrimType dtype, ffi::Array shape, ffi::Array< } tirx::Buffer BufferWithOffsetAlignment(ffi::Array shape, PrimType dtype, std::string name, - int data_alignment, int offset_factor, bool compact, + int data_alignment, int offset_factor, std::string memory_scope) { PrimType storage_type = (dtype == PrimType::Bool() ? PrimType::Int(8) : dtype); auto data = tirx::Var(name, PointerType(storage_type, memory_scope)); - bool has_any = false; - if (!compact) { - for (const auto& it : shape) { - if (it.as()) { - has_any = true; - break; - } - } - } - tirx::BufferType buffer_type = has_any ? tirx::kAutoBroadcast : tirx::kDefault; PrimExpr elem_offset; if (offset_factor != 0) { @@ -672,7 +563,7 @@ tirx::Buffer BufferWithOffsetAlignment(ffi::Array shape, PrimType dtyp } return tirx::Buffer(data, dtype, shape, ffi::Array(), elem_offset, name, data_alignment, - offset_factor, buffer_type); + offset_factor); } Buffer Buffer::with_allocated_addr(ffi::Array allocated_addr) const { @@ -704,10 +595,9 @@ PrimExpr Buffer::OffsetOf_p(const Array& indices) const { bool Buffer::IsScalar(bool alloc_or_decl) const { // TODO(@bohan): logical scope is not considered return (*this)->shape.size() == 1 && is_one((*this)->shape[0]) && (*this)->strides.size() == 0 && - (*this)->axis_separators.size() == 0 && (!alloc_or_decl || tirx::is_zero((*this)->elem_offset)) && (*this)->data_alignment == 64 && - (*this)->offset_factor == 1 && (*this)->buffer_type == tirx::BufferType::kDefault && - (*this)->allocated_addr.size() == 0 && (*this)->layout.has_value() && + (*this)->offset_factor == 1 && (*this)->allocated_addr.size() == 0 && + (*this)->layout.has_value() && ffi::StructuralEqual()((*this)->layout.value(), TileLayoutNode::DefaultLayout({1})); } @@ -716,9 +606,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef() .def_packed("tirx.Buffer", [](ffi::PackedArgs args, ffi::Any* ret) { - TVM_FFI_ICHECK_EQ(args.size(), 12); - auto buffer_type = args[8].cast(); - BufferType type = (buffer_type == "auto_broadcast") ? kAutoBroadcast : kDefault; + TVM_FFI_ICHECK_EQ(args.size(), 10); auto data = args[0].cast(); auto dtype = args[1].cast(); auto shape = args[2].cast>(); @@ -727,11 +615,10 @@ TVM_FFI_STATIC_INIT_BLOCK() { auto name = args[5].cast(); auto data_alignment = args[6].cast(); auto offset_factor = args[7].cast(); - auto axis_separators = args[9].cast>(); - auto span = args[10].cast(); - auto layout = args[11].cast(); + auto span = args[8].cast(); + auto layout = args[9].cast(); *ret = Buffer(data, dtype, shape, strides, elem_offset, name, data_alignment, - offset_factor, type, axis_separators, span, layout); + offset_factor, span, layout); }) .def_method("tirx.BufferAccessPtr", static_castregion.size() >= buffer->shape.size()) << "Dimension of source ffi::Array expected to be larger or equal than target buffer " diff --git a/src/tirx/script/builder/ir.cc b/src/tirx/script/builder/ir.cc index e9dc5c14da99..c1f37c7f7b0d 100644 --- a/src/tirx/script/builder/ir.cc +++ b/src/tirx/script/builder/ir.cc @@ -44,12 +44,8 @@ using tvm::tirx::Layout; Buffer BufferDecl(ffi::Array shape, PrimType dtype, ffi::String buffer_name, ffi::Optional data, ffi::Optional> strides, ffi::Optional elem_offset, ffi::String storage_scope, int align, - int offset_factor, ffi::String buffer_type, - ffi::Optional> axis_separators, ffi::Optional layout, + int offset_factor, ffi::Optional layout, ffi::Array allocated_addr) { - TVM_FFI_CHECK(buffer_type == "auto" || buffer_type == "default" || buffer_type.empty(), - ValueError) - << "ValueError: `buffer_type` must be `auto` or `default` or empty"; if (!allocated_addr.empty()) { TVM_FFI_ICHECK(!data.has_value() && !elem_offset.has_value() && !offset_factor) << "ValueError: `allocated_addr` can only be used with `data`, `elem_offset`, and " @@ -70,9 +66,8 @@ Buffer BufferDecl(ffi::Array shape, PrimType dtype, ffi::String buffer elem_offset = tvm::tirx::PrimVar("elem_offset", shape_dtype); } return Buffer(buffer_data, dtype, shape, strides.value_or(ffi::Array()), - elem_offset.value_or(PrimExpr()), buffer_name, align, offset_factor, - (buffer_type == "auto" ? tvm::tirx::kAutoBroadcast : tvm::tirx::kDefault), - axis_separators.value_or(ffi::Array()), Span(), layout, allocated_addr); + elem_offset.value_or(PrimExpr()), buffer_name, align, offset_factor, Span(), layout, + allocated_addr); } PrimFuncFrame PrimFunc(bool is_private, bool s_tir, bool persistent) { @@ -153,10 +148,9 @@ tvm::Type FuncRet(tvm::Type ret_type) { Buffer MatchBuffer(ffi::ObjectRef param, ffi::Array shape, PrimType dtype, ffi::Optional data, ffi::Array strides, PrimExpr elem_offset, ffi::String storage_scope, int align, int offset_factor, - ffi::String buffer_type_str, ffi::Optional> axis_separators, ffi::Optional layout) { Buffer buffer = BufferDecl(shape, dtype, "", data, strides, elem_offset, storage_scope, align, - offset_factor, buffer_type_str, axis_separators, layout, {}); + offset_factor, layout, {}); if (auto var = param.as()) { PrimFuncFrame frame = FindPrimFuncFrame("T.match_buffer"); Var v = var.value(); @@ -375,9 +369,7 @@ void BlockAttrs(ffi::Map attrs) { ffi::Variant SBlockAllocBuffer( ffi::Array shape, PrimType dtype, ffi::Optional data, ffi::Array strides, PrimExpr elem_offset, ffi::String storage_scope, int align, - int offset_factor, ffi::String buffer_type_str, - ffi::Optional> axis_separators, ffi::Optional layout, - ffi::Array allocated_addr) { + int offset_factor, ffi::Optional layout, ffi::Array allocated_addr) { std::string scope = static_cast(storage_scope); if (scope.empty()) { scope = "global"; @@ -389,9 +381,8 @@ ffi::Variant SBlockAllocBuffer( } ffi::Optional opt_elem_offset = elem_offset.defined() ? ffi::Optional(elem_offset) : std::nullopt; - Buffer buffer = - BufferDecl(shape, dtype, "", std::nullopt, strides, opt_elem_offset, storage_scope, align, - offset_factor, buffer_type_str, axis_separators, layout, allocated_addr); + Buffer buffer = BufferDecl(shape, dtype, "", std::nullopt, strides, opt_elem_offset, + storage_scope, align, offset_factor, layout, allocated_addr); IRBuilder builder = IRBuilder::Current(); auto opt_func_frame = builder->FindFrame(); if (opt_func_frame.has_value()) { @@ -815,9 +806,8 @@ void BufferStore(Buffer buffer, PrimExpr value, ffi::Array indices, DeclBufferFrame DeclBuffer(ffi::Array shape, PrimType dtype, ffi::String buffer_name, ffi::Optional data, ffi::Optional> strides, ffi::Optional elem_offset, ffi::String storage_scope, - int align, int offset_factor, ffi::String buffer_type, - ffi::Optional> axis_separators, - ffi::Optional layout, ffi::Optional allocated_addr) { + int align, int offset_factor, ffi::Optional layout, + ffi::Optional allocated_addr) { std::string scope = static_cast(storage_scope); if (scope.empty()) { scope = "global"; @@ -846,9 +836,8 @@ DeclBufferFrame DeclBuffer(ffi::Array shape, PrimType dtype, ffi::Stri } ffi::ObjectPtr n = ffi::make_object(); - n->buffer = - BufferDecl(shape, dtype, buffer_name, data, strides, elem_offset, storage_scope, align, - offset_factor, buffer_type, axis_separators, layout, allocated_addr_arr); + n->buffer = BufferDecl(shape, dtype, buffer_name, data, strides, elem_offset, storage_scope, + align, offset_factor, layout, allocated_addr_arr); // For tmem, even without `data`, we should not emit an Allocate node. n->allocated = (scope == "tmem") || data.has_value(); return DeclBufferFrame(n); @@ -857,7 +846,7 @@ DeclBufferFrame DeclBuffer(ffi::Array shape, PrimType dtype, ffi::Stri Buffer AllocBuffer(ffi::Array shape, PrimType dtype, ffi::String storage_scope, ffi::Optional> annotations) { Buffer buffer = BufferDecl(shape, dtype, "", std::nullopt, std::nullopt, std::nullopt, - storage_scope, 0, 0, "", std::nullopt); + storage_scope, 0, 0, std::nullopt, {}); AddToParent( tvm::tirx::AllocBuffer(buffer, annotations.value_or(ffi::Map()))); return buffer; @@ -921,8 +910,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .def("script.ir_builder.tirx.Buffer", static_cast, PrimType, ffi::String, ffi::Optional, ffi::Optional>, ffi::Optional, - ffi::String, int, int, ffi::String, - ffi::Optional>, ffi::Optional, + ffi::String, int, int, ffi::Optional, ffi::Array)>(BufferDecl)) .def("script.ir_builder.tirx.PrimFunc", PrimFunc) .def("script.ir_builder.tirx.Arg", diff --git a/src/tirx/script/printer/buffer.cc b/src/tirx/script/printer/buffer.cc index 3a43f6fdd4c7..732de462a285 100644 --- a/src/tirx/script/printer/buffer.cc +++ b/src/tirx/script/printer/buffer.cc @@ -175,16 +175,7 @@ ffi::Map BufferAttrs(tirx::Buffer buffer, const AccessPath kwargs.Set("offset_factor", LiteralDoc::Int(buffer->offset_factor, buffer_p->Attr("offset_factor"))); } - // Step 9. Handle `buffer.buffer_type` - if (buffer->buffer_type != tirx::BufferType::kDefault) { - kwargs.Set("buffer_type", LiteralDoc::Str("auto", buffer_p->Attr("buffer_type"))); - } - // Step 10. Handle `buffer.axis_separator` - if (!buffer->axis_separators.empty()) { - kwargs.Set("axis_separators", - d->AsDoc(buffer->axis_separators, buffer_p->Attr("axis_separators"))); - } - // Step 12. Handle `buffer.layout`. Track the enclosing PrimFunc's `s_tir` + // Step 9. Handle `buffer.layout`. Track the enclosing PrimFunc's `s_tir` // attr — in `s_tir=True` mode the parser fills `layout=None` by default, // in `s_tir=False` (tirx) mode it fills `DefaultLayout(shape)`. Mirror // that here so the implicit default is omitted and the non-default value @@ -209,7 +200,7 @@ ffi::Map BufferAttrs(tirx::Buffer buffer, const AccessPath } else if (!enclosing_s_tir) { kwargs.Set("layout", LiteralDoc::None(buffer_p->Attr("layout"))); } - // Step 13. Handle `buffer.allocated_addr` + // Step 10. Handle `buffer.allocated_addr` if (!buffer->allocated_addr.empty()) { if (buffer->allocated_addr.size() == 1) { // Unwrap single-element array: DeclBuffer expects Optional, not Array. @@ -262,7 +253,7 @@ ExprDoc BufferCall(const ExprDoc& prefix, const ffi::Map& } } for (ffi::String s : {"data", "strides", "elem_offset", "scope", "align", "offset_factor", - "buffer_type", "axis_separators", "layout", "allocated_addr"}) { + "layout", "allocated_addr"}) { if (ffi::Optional doc = attrs.Get(s)) { kwargs_keys.push_back(s); kwargs_values.push_back(doc.value()); diff --git a/src/tirx/script/printer/function.cc b/src/tirx/script/printer/function.cc index ce66059ed5fd..b3a553c17905 100644 --- a/src/tirx/script/printer/function.cc +++ b/src/tirx/script/printer/function.cc @@ -62,8 +62,7 @@ bool IsSimpleBuffer(const tirx::Buffer& buf, bool s_tir) { return false; } return buf.scope() == "global" && buf->data_alignment == runtime::kAllocAlignment && - buf->offset_factor == 1 && buf->buffer_type == tirx::BufferType::kDefault && - !buf->axis_separators.size(); + buf->offset_factor == 1; } int CountVarOccurrence(const tirx::PrimFunc& f, const tirx::Var& v) { diff --git a/src/tirx/transform/ir_utils.h b/src/tirx/transform/ir_utils.h index 1313fe8918cf..716f26a5f1e5 100644 --- a/src/tirx/transform/ir_utils.h +++ b/src/tirx/transform/ir_utils.h @@ -114,8 +114,7 @@ inline PrimExpr TVMStructGet(PrimType type, Var handle, int index, inline Call AddressOffset(Var handle, PrimType dtype, int offset) { PrimExpr offset_expr = IntImm::Int32(offset * dtype.lanes()); ffi::Array shape = {offset_expr + 1}; - Buffer dummy_buf(handle, dtype, shape, {}, 0, handle->name, 0, 0, kDefault, {}, Span(), - std::nullopt); + Buffer dummy_buf(handle, dtype, shape, {}, 0, handle->name, 0, 0, Span(), std::nullopt); BufferLoad buf_load(dummy_buf, {offset_expr}); return Call(handle->ty, builtin::address_of(), {buf_load}); @@ -135,8 +134,8 @@ inline Call AddressOffset(Var handle, PrimType dtype, PrimExpr offset) { } ffi::Array shape = {offset + 1}; - Buffer dummy_buf(handle, dtype.WithLanes(1), shape, {}, 0, handle->name, 0, 0, kDefault, {}, - Span(), std::nullopt); + Buffer dummy_buf(handle, dtype.WithLanes(1), shape, {}, 0, handle->name, 0, 0, Span(), + std::nullopt); BufferLoad buf_load(dummy_buf, {offset}); return Call(handle->ty, builtin::address_of(), {buf_load}); diff --git a/src/tirx/transform/lower_intrin.cc b/src/tirx/transform/lower_intrin.cc index f08bf49ddd25..4c93b16440e8 100644 --- a/src/tirx/transform/lower_intrin.cc +++ b/src/tirx/transform/lower_intrin.cc @@ -76,8 +76,7 @@ static Expr LowerAccessPtr(const CallNode* call) { offset = offset * IntImm(offset_ty, dtype.lanes()); offset = Ramp(offset, IntImm(offset_ty, 1), dtype.lanes()); } - Buffer dummy_buf(buffer_var, dtype.WithLanes(1), {offset + 1}, {}, 0, buffer_var->name, 0, 0, - kDefault); + Buffer dummy_buf(buffer_var, dtype.WithLanes(1), {offset + 1}, {}, 0, buffer_var->name, 0, 0); BufferLoad buf_load(dummy_buf, {offset}); return Call(call->ty, builtin::address_of(), {buf_load}); } diff --git a/src/tirx/transform/lower_tirx_cleanup.cc b/src/tirx/transform/lower_tirx_cleanup.cc index b903dda87a63..8a977a729d72 100644 --- a/src/tirx/transform/lower_tirx_cleanup.cc +++ b/src/tirx/transform/lower_tirx_cleanup.cc @@ -134,7 +134,6 @@ class LayoutApplier : public arith::IRMutatorWithAnalyzer { writer = flattened.CopyOnWrite(); writer->shape = new_shape; writer->strides = {}; - writer->axis_separators = {}; } else if (is_alloc) { if (auto tile_layout = buf->layout.as(); tile_layout && tile_layout->HasThreadAxis()) { @@ -160,7 +159,6 @@ class LayoutApplier : public arith::IRMutatorWithAnalyzer { writer = flattened.CopyOnWrite(); writer->shape = {ana->Simplify(mem_span)}; writer->strides = {}; - writer->axis_separators = {}; } else { flattened = buf.GetFlattenedBuffer(); writer = flattened.CopyOnWrite(); diff --git a/src/tirx/transform/lower_warp_memory.cc b/src/tirx/transform/lower_warp_memory.cc index 62eb5f479863..4981dee72558 100644 --- a/src/tirx/transform/lower_warp_memory.cc +++ b/src/tirx/transform/lower_warp_memory.cc @@ -279,7 +279,7 @@ class WarpAccessRewriter : protected StmtExprMutator { alloc_size = warp_group_ * factor; Buffer new_buf(op->buffer->data, op->buffer->dtype, {IntImm::Int32(alloc_size / width_)}, {}, - PrimExpr(), op->buffer->data->name, 0, 0, BufferType::kDefault); + PrimExpr(), op->buffer->data->name, 0, 0); Stmt rewritten_body = this->VisitStmt(body); return SeqStmt::Flatten(AllocBuffer(new_buf, op->annotations), rewritten_body); } diff --git a/src/tirx/transform/storage_rewrite.cc b/src/tirx/transform/storage_rewrite.cc index 650021d1256d..947f787c1344 100644 --- a/src/tirx/transform/storage_rewrite.cc +++ b/src/tirx/transform/storage_rewrite.cc @@ -132,11 +132,11 @@ class LinearAccessPatternFinder final : public StmtExprVisitor { TVM_FFI_ICHECK_LT(it->second.level, scope_.size()); scope_[it->second.level].touched.push_back(buffer_var); - TVM_FFI_ICHECK_EQ(op->buffer->axis_separators.size() + 1, it->second.num_physical_dimensions) + TVM_FFI_ICHECK_EQ(1, it->second.num_physical_dimensions) << "Buffer " << op->buffer->name << " is allocated with " << it->second.num_physical_dimensions << " physical dimensions, but is accessed as having " - << op->buffer->axis_separators.size() + 1 << " physical dimensions" << std::endl; + << "1 physical dimension" << std::endl; } StmtEntry e = scope_.back(); scope_.pop_back(); @@ -159,11 +159,11 @@ class LinearAccessPatternFinder final : public StmtExprVisitor { << "Load memory in places other than store."; scope_[it->second.level].touched.push_back(buffer_var); - TVM_FFI_ICHECK_EQ(op->buffer->axis_separators.size() + 1, it->second.num_physical_dimensions) + TVM_FFI_ICHECK_EQ(1, it->second.num_physical_dimensions) << "Buffer " << op->buffer->name << " is allocated with " << it->second.num_physical_dimensions << " physical dimensions, but is accessed as having " - << op->buffer->axis_separators.size() + 1 << " physical dimensions" << std::endl; + << "1 physical dimension" << std::endl; } } @@ -477,10 +477,9 @@ class StoragePlanRewriter : public StmtExprMutator { return it->second; } - Buffer remapped = - Buffer(new_backing_array, buf->dtype, buf->shape, buf->strides, buf->elem_offset, - new_backing_array->name, buf->data_alignment, buf->offset_factor, buf->buffer_type, - buf->axis_separators, buf->span, buf->layout, buf->allocated_addr); + Buffer remapped = Buffer(new_backing_array, buf->dtype, buf->shape, buf->strides, + buf->elem_offset, new_backing_array->name, buf->data_alignment, + buf->offset_factor, buf->span, buf->layout, buf->allocated_addr); buffer_remap_[key] = remapped; return remapped; } @@ -786,7 +785,7 @@ class StoragePlanRewriter : public StmtExprMutator { } combo_size = analyzer_->Simplify(combo_size); Buffer buf(e->alloc_var, alloc_type, {combo_size}, {}, PrimExpr(), e->alloc_var->name, 0, - 0, BufferType::kDefault); + 0); ffi::Map annotations; if (e->is_volatile) { annotations.Set(attr::kVolatile, true); @@ -823,8 +822,7 @@ class StoragePlanRewriter : public StmtExprMutator { uint64_t type_bits = e->elem_type.bits() * e->elem_type.lanes(); PrimExpr alloc_size = MakeConst(e->allocs[0]->buffer->shape[0].ty(), (total_bits + type_bits - 1) / type_bits); - Buffer buf(e->alloc_var, e->elem_type, {alloc_size}, {}, PrimExpr(), e->alloc_var->name, 0, 0, - BufferType::kDefault); + Buffer buf(e->alloc_var, e->elem_type, {alloc_size}, {}, PrimExpr(), e->alloc_var->name, 0, 0); bool any_volatile = e->is_volatile; for (StorageEntry* child : e->merged_children) { if (child->is_volatile) any_volatile = true; diff --git a/src/tirx/transform/tvm_ffi_binder.cc b/src/tirx/transform/tvm_ffi_binder.cc index 8c487c842fd1..a2d44f6a4e9f 100644 --- a/src/tirx/transform/tvm_ffi_binder.cc +++ b/src/tirx/transform/tvm_ffi_binder.cc @@ -663,23 +663,6 @@ void TVMFFIABIBuilder::BindCompactStrides(const Buffer& buffer, const Var& strid } } -void TVMFFIABIBuilder::BindAutoBroadcastStrides(const Buffer& buffer, const Var& strides_ptr, - const PrimExpr& v_strides_is_null, - const ffi::reflection::AccessPath& param_path) { - PrimType stype(buffer->DefaultIndexType()); - PrimExpr stride = IntImm(stype, 1); - for (size_t i = buffer->shape.size(); i != 0; --i) { - size_t k = i - 1; - PrimExpr value = cast(buffer->shape[k].ty(), LoadInt64ArrayElem(strides_ptr, k)); - value = tvm::if_then_else(v_strides_is_null, stride, value); - value = tvm::if_then_else(buffer->shape[k] == 1, 0, value); - ffi::reflection::AccessPath strides_k_path = - param_path->Attr(ffi::String("strides"))->ArrayItem(k); - BindScalar(buffer->strides[k], value, strides_k_path, true); - stride = analyzer_->Simplify(stride * buffer->shape[k]); - } -} - void TVMFFIABIBuilder::BindRegularStrides(const Buffer& buffer, const Var& strides_ptr, const Var& shape_ptr, const PrimExpr& v_strides_is_null, const ffi::reflection::AccessPath& param_path) { @@ -761,8 +744,6 @@ void TVMFFIABIBuilder::DecodeParamDLTensor(const Buffer& buffer, const PrimExpr& Call(PrimType::Bool(), builtin::isnullptr(), {strides_ptr}).as_or_throw(); if (buffer->strides.size() == 0) { BindCompactStrides(buffer, strides_ptr, v_strides_is_null, param_path); - } else if (buffer->buffer_type == kAutoBroadcast) { - BindAutoBroadcastStrides(buffer, strides_ptr, v_strides_is_null, param_path); } else { BindRegularStrides(buffer, strides_ptr, shape_ptr, v_strides_is_null, param_path); } diff --git a/src/tirx/transform/tvm_ffi_binder.h b/src/tirx/transform/tvm_ffi_binder.h index 7e6effe28d76..a90a295f492d 100644 --- a/src/tirx/transform/tvm_ffi_binder.h +++ b/src/tirx/transform/tvm_ffi_binder.h @@ -325,18 +325,6 @@ class TVMFFIABIBuilder { const PrimExpr& v_strides_is_null, const ffi::reflection::AccessPath& param_path); - /*! - * \brief Bind strides for auto-broadcast buffers: stride=0 for shape==1 dims. - * - * \param buffer The expected buffer definition. - * \param strides_ptr The strides pointer variable. - * \param v_strides_is_null Expression checking if strides pointer is NULL. - * \param param_path ffi::reflection::AccessPath for the tensor parameter. - */ - void BindAutoBroadcastStrides(const Buffer& buffer, const Var& strides_ptr, - const PrimExpr& v_strides_is_null, - const ffi::reflection::AccessPath& param_path); - /*! * \brief Bind strides with C-contiguous fallback when strides pointer is NULL. * diff --git a/src/tirx/transform/unsupported_dtype_legalize.cc b/src/tirx/transform/unsupported_dtype_legalize.cc index 8da543e8f982..912bf189f0e9 100644 --- a/src/tirx/transform/unsupported_dtype_legalize.cc +++ b/src/tirx/transform/unsupported_dtype_legalize.cc @@ -145,8 +145,7 @@ class ComputeLegalizePlanner : public StmtExprVisitor { Buffer new_buffer(var_it->second, promote_dtype_.WithLanes(buf->dtype.lanes()), buf->shape, buf->strides, buf->elem_offset, buf->name, buf->data_alignment, - buf->offset_factor, buf->buffer_type, buf->axis_separators, buf->span, - buf->layout, buf->allocated_addr); + buf->offset_factor, buf->span, buf->layout, buf->allocated_addr); (*buffer_remap_)[buf] = new_buffer; } @@ -586,8 +585,8 @@ class StorageLegalizer : public StmtExprMutator { Var new_data = Var(buf->data->name, PointerType(new_dtype, storage_scope)); var_remap_[buf->data] = new_data; buf = Buffer(new_data, new_dtype, buf->shape, buf->strides, buf->elem_offset, buf->name, - buf->data_alignment, buf->offset_factor, buf->buffer_type, buf->axis_separators, - buf->span, buf->layout, buf->allocated_addr); + buf->data_alignment, buf->offset_factor, buf->span, buf->layout, + buf->allocated_addr); buffer_remap_[op->buffer] = buf; } if (buf.same_as(op->buffer)) { @@ -606,9 +605,8 @@ class StorageLegalizer : public StmtExprMutator { // force remap here if (MatchType(buf->dtype)) { buf = Buffer(buf->data, GetStorageUIntDType(buf->dtype), buf->shape, buf->strides, - buf->elem_offset, buf->name, buf->data_alignment, buf->offset_factor, - buf->buffer_type, buf->axis_separators, buf->span, buf->layout, - buf->allocated_addr); + buf->elem_offset, buf->name, buf->data_alignment, buf->offset_factor, buf->span, + buf->layout, buf->allocated_addr); buffer_remap_[op->buffer] = buf; } if (buf.same_as(op->buffer)) { @@ -767,8 +765,8 @@ class StorageLegalizer : public StmtExprMutator { if (var_it != var_remap_.end()) { PrimType dtype = MatchType(buf->dtype) ? GetStorageUIntDType(buf->dtype) : buf->dtype; new_buf = Buffer(var_it->second, dtype, buf->shape, buf->strides, buf->elem_offset, buf->name, - buf->data_alignment, buf->offset_factor, buf->buffer_type, - buf->axis_separators, buf->span, buf->layout, buf->allocated_addr); + buf->data_alignment, buf->offset_factor, buf->span, buf->layout, + buf->allocated_addr); } else { TVM_FFI_ICHECK(!MatchType(buf->dtype)) << "Cannot find var remap for " << buf; } diff --git a/tests/cpp/ir_functor_test.cc b/tests/cpp/ir_functor_test.cc index 857a7a790f67..190a81d66574 100644 --- a/tests/cpp/ir_functor_test.cc +++ b/tests/cpp/ir_functor_test.cc @@ -154,7 +154,7 @@ TEST(IRF, StmtVisitor) { Stmt eval_body = Evaluate(z); PrimType dtype = PrimType::Float(32); tirx::Var data_var("b", PointerType(dtype)); - Buffer buf(data_var, dtype, {z, z}, {}, PrimExpr(), "b", 0, 0, BufferType::kDefault); + Buffer buf(data_var, dtype, {z, z}, {}, PrimExpr(), "b", 0, 0); // AllocBuffer is flat (no body). Return as SeqStmt with eval. return SeqStmt({AllocBuffer(buf), eval_body}); }; @@ -208,7 +208,7 @@ TEST(IRF, StmtMutator) { auto z = x + 1; PrimType dtype = PrimType::Float(32); tirx::Var data_var("b", PointerType(dtype)); - Buffer buf(data_var, dtype, {1, z}, {}, PrimExpr(), "b", 0, 0, BufferType::kDefault); + Buffer buf(data_var, dtype, {1, z}, {}, PrimExpr(), "b", 0, 0); return AllocBuffer(buf); }; @@ -342,8 +342,7 @@ TEST(IRF, Substitute) { /*elem_offset=*/PrimExpr(), /*name=*/"buf", /*data_alignment=*/1, - /*offset_factor=*/1, - /*buffer_type=*/BufferType::kDefault}; + /*offset_factor=*/1}; }; { diff --git a/tests/python/relax/test_optimize_layout_transform.py b/tests/python/relax/test_optimize_layout_transform.py index 0ef21c2daab2..78fa4555d23a 100644 --- a/tests/python/relax/test_optimize_layout_transform.py +++ b/tests/python/relax/test_optimize_layout_transform.py @@ -306,7 +306,6 @@ def main(x: R.Tensor((14,), dtype="float32")) -> R.Tensor((14,), dtype="float32" x, index_map=T.index_map(lambda i: (i % 16,)), pad_value=None, - axis_separators=[], ) lv1 = R.call_tir( Before.relax_relu_replacement, @@ -317,7 +316,6 @@ def main(x: R.Tensor((14,), dtype="float32")) -> R.Tensor((14,), dtype="float32" lv1, index_map=T.index_map(lambda axis0: (axis0,)), pad_value=None, - axis_separators=[], ) lv_1 = R.call_tir( Before.remove_pad, (lv2,), out_ty=R.Tensor((14,), dtype="float32") @@ -326,7 +324,6 @@ def main(x: R.Tensor((14,), dtype="float32")) -> R.Tensor((14,), dtype="float32" lv_1, index_map=T.index_map(lambda i: (i % 16,)), pad_value=None, - axis_separators=[], ) lv4 = R.call_tir( Before.relax_relu_replacement, @@ -337,7 +334,6 @@ def main(x: R.Tensor((14,), dtype="float32")) -> R.Tensor((14,), dtype="float32" lv4, index_map=T.index_map(lambda axis0: (axis0,)), pad_value=None, - axis_separators=[], ) lv_2 = R.call_tir( Before.remove_pad, (lv5,), out_ty=R.Tensor((14,), dtype="float32") @@ -383,7 +379,6 @@ def main(x: R.Tensor((14,), dtype="float32")) -> R.Tensor((14,), dtype="float32" x, index_map=T.index_map(lambda i: (i % 16,)), pad_value=None, - axis_separators=[], ) lv1 = R.call_tir( Expected.relax_relu_replacement, @@ -399,7 +394,6 @@ def main(x: R.Tensor((14,), dtype="float32")) -> R.Tensor((14,), dtype="float32" lv4, index_map=T.index_map(lambda axis0: (axis0,)), pad_value=None, - axis_separators=[], ) gv = R.call_tir( Expected.remove_pad, (lv5,), out_ty=R.Tensor((14,), dtype="float32") diff --git a/tests/python/relax/test_transform_alter_op_impl.py b/tests/python/relax/test_transform_alter_op_impl.py index cd433d0a30c3..0d927568529b 100644 --- a/tests/python/relax/test_transform_alter_op_impl.py +++ b/tests/python/relax/test_transform_alter_op_impl.py @@ -22,7 +22,6 @@ from tvm.script import ir as I from tvm.script import relax as R from tvm.script import tirx as T -from tvm.tirx import IndexMap kOperatorName = "operator_name" @@ -33,14 +32,10 @@ def _check( operator_name, replacement_primfunc, layout_changes, - axis_separator=None, - input_axis_separator=None, ): after = relax.transform.AlterOpImpl( {operator_name: replacement_primfunc}, {operator_name: layout_changes}, - {operator_name: axis_separator}, - {operator_name: input_axis_separator}, )(before) after = relax.transform.DeadCodeElimination()(after) tvm.ir.assert_structural_equal(after, expected) @@ -239,79 +234,6 @@ def some_op_2d(arg0: T.Buffer((4, 4), "float32"), arg1: T.Buffer((4, 4), "float3 ) -def test_multiple_outputs_with_axis_sep(): - # fmt: off - @I.ir_module(s_tir=True) - class Before: - @T.prim_func(private=True, s_tir=True) - def some_op(arg0: T.Buffer((16,), "float32"), arg1: T.Buffer((16,), "float32"), output0: T.Buffer((16,), "float32"), output1: T.Buffer((16,), "float32")): - T.func_attr({"operator_name": "relax.some_op"}) - for ax0 in range(16): - with T.sblock("T_add"): - v_ax0 = T.axis.spatial(16, ax0) - T.reads(arg0[v_ax0], arg1[v_ax0]) - T.writes(output0[v_ax0], output1[v_ax0]) - output0[v_ax0] = arg0[v_ax0] + arg1[v_ax0] - output1[v_ax0] = arg0[v_ax0] - arg1[v_ax0] - - @R.function - def main(x: R.Tensor((16,), dtype="float32"), y: R.Tensor((16,), dtype="float32")) -> R.Tuple(R.Tensor((16,), dtype="float32"), R.Tensor((16,), dtype="float32")): - with R.dataflow(): - gv = R.call_tir(Before.some_op, (x, y), out_ty=[R.Tensor((16,), dtype="float32"), R.Tensor((16,), dtype="float32")]) - R.output(gv) - return gv - - @I.ir_module(s_tir=True) - class Expected: - @T.prim_func(private=True, s_tir=True) - def relax_some_op_replacement(arg0: T.Buffer((4, 4), "float32"), arg1: T.Buffer((4, 4), "float32"), output0: T.Buffer((4, 4), "float32"), output1: T.Buffer((4, 4), "float32")): - T.func_attr({"operator_name": "relax.some_op"}) - for ax0, ax1 in T.grid(4, 4): - with T.sblock("T_add"): - v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) - T.reads(arg0[v_ax0, v_ax1], arg1[v_ax0, v_ax1]) - T.writes(output0[v_ax0, v_ax1], output1[v_ax0, v_ax1]) - output0[v_ax0, v_ax1] = arg0[v_ax0, v_ax1] + arg1[v_ax0, v_ax1] - output1[v_ax0, v_ax1] = arg0[v_ax0, v_ax1] - arg1[v_ax0, v_ax1] - - @R.function - def main(x: R.Tensor((16,), dtype="float32"), y: R.Tensor((16,), dtype="float32")) -> R.Tuple(R.Tensor((16,), dtype="float32"), R.Tensor((16,), dtype="float32")): - with R.dataflow(): - lv: R.Tensor((4, 4), dtype="float32") = R.layout_transform(x, index_map=lambda i: (i // 4, i % 4), pad_value=None, axis_separators=[1]) - lv1: R.Tensor((4, 4), dtype="float32") = R.layout_transform(y, index_map=lambda i: (i // 4, i % 4), pad_value=None, axis_separators=[1]) - lv2 = R.call_tir(Expected.relax_some_op_replacement, (lv, lv1), out_ty=[R.Tensor((4, 4), dtype="float32"), R.Tensor((4, 4), dtype="float32")]) - lv3: R.Tensor((4, 4), dtype="float32") = lv2[0] - lv4: R.Tensor((16,), dtype="float32") = R.layout_transform(lv3, index_map=lambda axis0, axis1: (axis0 * 4 + axis1,), pad_value=None, axis_separators=[1]) - lv5: R.Tensor((4, 4), dtype="float32") = lv2[1] - lv6: R.Tensor((16,), dtype="float32") = R.layout_transform(lv5, index_map=lambda axis0, axis1: (axis0 * 4 + axis1,), pad_value=None, axis_separators=[1]) - gv: R.Tuple(R.Tensor((16,), dtype="float32"), R.Tensor((16,), dtype="float32")) = (lv4, lv6) - R.output(gv) - return gv - - @T.prim_func(private=True, s_tir=True) - def some_op_2d(arg0: T.Buffer((4, 4), "float32"), arg1: T.Buffer((4, 4), "float32"), output0: T.Buffer((4, 4), "float32"), output1: T.Buffer((4, 4), "float32")): - for ax0, ax1 in T.grid(4, 4): - with T.sblock("T_add"): - v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) - T.reads(arg0[v_ax0, v_ax1], arg1[v_ax0, v_ax1]) - T.writes(output0[v_ax0, v_ax1], output1[v_ax0, v_ax1]) - output0[v_ax0, v_ax1] = arg0[v_ax0, v_ax1] + arg1[v_ax0, v_ax1] - output1[v_ax0, v_ax1] = arg0[v_ax0, v_ax1] - arg1[v_ax0, v_ax1] - # fmt: on - - index_map, axis_sep = IndexMap.from_func_with_separators( - lambda i: (i // 4, IndexMap.AXIS_SEPARATOR, i % 4) - ) - _check( - Before, - Expected, - operator_name="relax.some_op", - replacement_primfunc=some_op_2d, - layout_changes=[index_map, index_map, index_map, index_map], - axis_separator=[axis_sep, axis_sep, axis_sep, axis_sep], - ) - - def test_supported_implicit_padding(): @I.ir_module(s_tir=True) class Before: @@ -342,7 +264,6 @@ def foo(x: R.Tensor((14,), dtype="float32")) -> R.Tensor((14,), dtype="float32") x, index_map=T.index_map(lambda i: (i % 16,)), pad_value=None, - axis_separators=[], ) lv1 = R.call_tir( Expected.relax_relu_replacement, @@ -353,7 +274,6 @@ def foo(x: R.Tensor((14,), dtype="float32")) -> R.Tensor((14,), dtype="float32") lv1, index_map=T.index_map(lambda axis0: (axis0,)), pad_value=None, - axis_separators=[], ) lv_1 = R.call_tir( Expected.remove_pad, (lv2,), out_ty=R.Tensor((14,), dtype="float32") @@ -543,7 +463,6 @@ def main(x: R.Tensor((850, 2048), dtype="float16")) -> R.Tensor( x, index_map=T.index_map(lambda i, j: (i, j // 1024, j % 1024)), pad_value=None, - axis_separators=[], ) lv_1 = R.call_tir( cls.relax_reshape_replacement, @@ -579,81 +498,5 @@ def reshape_new( ) -def test_input_axis_separator(): - # fmt: off - @I.ir_module(s_tir=True) - class Before: - @T.prim_func(private=True, s_tir=True) - def some_op(arg0: T.Buffer((16,), "float32"), arg1: T.Buffer((16,), "float32"), output0: T.Buffer((16,), "float32"), output1: T.Buffer((16,), "float32")): - T.func_attr({"operator_name": "relax.some_op"}) - for ax0 in range(16): - with T.sblock("T_add"): - v_ax0 = T.axis.spatial(16, ax0) - T.reads(arg0[v_ax0], arg1[v_ax0]) - T.writes(output0[v_ax0], output1[v_ax0]) - output0[v_ax0] = arg0[v_ax0] + arg1[v_ax0] - output1[v_ax0] = arg0[v_ax0] - arg1[v_ax0] - - @R.function - def main(x: R.Tensor((16,), dtype="float32"), y: R.Tensor((16,), dtype="float32")) -> R.Tuple(R.Tensor((16,), dtype="float32"), R.Tensor((16,), dtype="float32")): - with R.dataflow(): - gv = R.call_tir(Before.some_op, (x, y), out_ty=[R.Tensor((16,), dtype="float32"), R.Tensor((16,), dtype="float32")]) - R.output(gv) - return gv - - @I.ir_module(s_tir=True) - class Expected: - @T.prim_func(private=True, s_tir=True) - def relax_some_op_replacement(arg0: T.Buffer((4, 4), "float32"), arg1: T.Buffer((4, 4), "float32"), output0: T.Buffer((4, 4), "float32"), output1: T.Buffer((4, 4), "float32")): - T.func_attr({"operator_name": "relax.some_op"}) - for ax0, ax1 in T.grid(4, 4): - with T.sblock("T_add"): - v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) - output0[v_ax0, v_ax1] = arg0[v_ax0, v_ax1] + arg1[v_ax0, v_ax1] - output1[v_ax0, v_ax1] = arg0[v_ax0, v_ax1] - arg1[v_ax0, v_ax1] - - @R.function - def main(x: R.Tensor((16,), dtype="float32"), y: R.Tensor((16,), dtype="float32")) -> R.Tuple(R.Tensor((16,), dtype="float32"), R.Tensor((16,), dtype="float32")): - with R.dataflow(): - lv: R.Tensor((4, 4), dtype="float32") = R.layout_transform(x, index_map=lambda i: (i // 4, i % 4), pad_value=None, axis_separators=[1]) - lv1: R.Tensor((4, 4), dtype="float32") = R.layout_transform(y, index_map=lambda i: (i // 4, i % 4), pad_value=None, axis_separators=[1]) - lv2 = R.call_tir(Expected.relax_some_op_replacement, (lv, lv1), out_ty=[R.Tensor((4, 4), dtype="float32"), R.Tensor((4, 4), dtype="float32")]) - lv3: R.Tensor((4, 4), dtype="float32") = lv2[0] - lv4: R.Tensor((16,), dtype="float32") = R.layout_transform(lv3, index_map=lambda axis0, axis1: (axis0 * 4 + axis1,), pad_value=None, axis_separators=[], input_axis_separators=[1]) - lv5: R.Tensor((4, 4), dtype="float32") = lv2[1] - lv6: R.Tensor((16,), dtype="float32") = R.layout_transform(lv5, index_map=lambda axis0, axis1: (axis0 * 4 + axis1,), pad_value=None, axis_separators=[], input_axis_separators=[1]) - gv: R.Tuple(R.Tensor((16,), dtype="float32"), R.Tensor((16,), dtype="float32")) = (lv4, lv6) - R.output(gv) - return gv - - @T.prim_func(private=True, s_tir=True) - def some_op_2d(arg0: T.Buffer((4, 4), "float32"), arg1: T.Buffer((4, 4), "float32"), output0: T.Buffer((4, 4), "float32"), output1: T.Buffer((4, 4), "float32")): - for ax0, ax1 in T.grid(4, 4): - with T.sblock("T_add"): - v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) - output0[v_ax0, v_ax1] = arg0[v_ax0, v_ax1] + arg1[v_ax0, v_ax1] - output1[v_ax0, v_ax1] = arg0[v_ax0, v_ax1] - arg1[v_ax0, v_ax1] - # fmt: on - - index_map_axis_sep = IndexMap.from_func_with_separators( - lambda i: (i // 4, IndexMap.AXIS_SEPARATOR, i % 4) - ) - - _check( - Before, - Expected, - operator_name="relax.some_op", - replacement_primfunc=some_op_2d, - layout_changes=[ - index_map_axis_sep, - index_map_axis_sep, - index_map_axis_sep, - index_map_axis_sep, - ], - axis_separator=[index_map_axis_sep[1], index_map_axis_sep[1], [], []], - input_axis_separator=[[], [], index_map_axis_sep[1], index_map_axis_sep[1]], - ) - - if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/relax/test_transform_convert_layout.py b/tests/python/relax/test_transform_convert_layout.py index fed702b7ba29..68ee8064419d 100644 --- a/tests/python/relax/test_transform_convert_layout.py +++ b/tests/python/relax/test_transform_convert_layout.py @@ -1676,8 +1676,6 @@ def main( lambda i0, i1, i2, i3: (i0, i1 // 4, i2, i3, i1 % 4), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) lv1: R.Tensor((1, 16, 3, 3, 4), dtype="float32") = R.layout_transform( w, @@ -1685,8 +1683,6 @@ def main( lambda i0, i1, i2, i3: (i0 // 4, i1, i2, i3, i0 % 4), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) lv2: R.Tensor((2, 1, 26, 26, 4), dtype="float32") = R.nn.conv2d( lv, @@ -1706,8 +1702,6 @@ def main( lambda i0, i1, i2, i3, i4: (i0, i1 * 4 + i4, i2, i3), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) R.output(gv) return gv @@ -1728,8 +1722,6 @@ def main( lambda i0, i1, i2, i3: (i0, i2, i3, i1 // 4, i1 % 4), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) lv1: R.Tensor((1, 3, 3, 16, 4), dtype="float32") = R.layout_transform( w, @@ -1737,8 +1729,6 @@ def main( lambda i0, i1, i2, i3: (i0 // 4, i2, i3, i1, i0 % 4), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) lv2: R.Tensor((2, 26, 26, 1, 4), dtype="float32") = R.nn.conv2d( lv, @@ -1758,8 +1748,6 @@ def main( lambda i0, i1, i2, i3, i4: (i0, i3 * 4 + i4, i1, i2), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) R.output(gv) return gv @@ -1799,8 +1787,6 @@ def main( lambda i0, i1, i2, i3: (i0, i3 // 4, i1, i2, i3 % 4), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) lv1: R.Tensor((1, 16, 3, 3, 4), dtype="float32") = R.layout_transform( w, @@ -1808,8 +1794,6 @@ def main( lambda i0, i1, i2, i3: (i0 // 4, i3, i1, i2, i0 % 4), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) lv2: R.Tensor((2, 1, 26, 26, 4), dtype="float32") = R.nn.conv2d( lv, @@ -1829,8 +1813,6 @@ def main( lambda i0, i1, i2, i3, i4: (i0, i2, i3, i1 * 4 + i4), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) R.output(gv) return gv @@ -1851,8 +1833,6 @@ def main( lambda i0, i1, i2, i3: (i0, i1, i2, i3 // 4, i3 % 4), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) lv1: R.Tensor((1, 3, 3, 16, 4), dtype="float32") = R.layout_transform( w, @@ -1860,8 +1840,6 @@ def main( lambda i0, i1, i2, i3: (i0 // 4, i1, i2, i3, i0 % 4), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) lv2: R.Tensor((2, 26, 26, 1, 4), dtype="float32") = R.nn.conv2d( lv, @@ -1881,8 +1859,6 @@ def main( lambda i0, i1, i2, i3, i4: (i0, i1, i2, i3 * 4 + i4), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) R.output(gv) return gv @@ -1979,8 +1955,6 @@ def main( lambda i0, i1, i2, i3: (i0, i1 // 4, i2, i3, i1 % 4), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) lv1_1: R.Tensor((1, 16, Hw, Ww, 4), dtype="float32") = R.layout_transform( lv1, @@ -1988,8 +1962,6 @@ def main( lambda i0, i1, i2, i3: (i0 // 4, i1, i2, i3, i0 % 4), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) lv2: R.Tensor((N, 1, H + 1 - Hw, W + 1 - Ww, 4), dtype="float32") = R.nn.conv2d( lv, @@ -2009,8 +1981,6 @@ def main( lambda i0, i1, i2, i3, i4: (i0, i1 * 4 + i4, i2, i3), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) R.output(gv) return gv @@ -2073,8 +2043,6 @@ def main( lambda i0, i1, i2, i3: (i0, i2, i3, i1 // 4, i1 % 4), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) lv1_1: R.Tensor((1, Hw, Ww, 16, 4), dtype="float32") = R.layout_transform( lv1, @@ -2082,8 +2050,6 @@ def main( lambda i0, i1, i2, i3: (i0 // 4, i2, i3, i1, i0 % 4), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) lv2: R.Tensor((N, H + 1 - Hw, W + 1 - Ww, 1, 4), dtype="float32") = R.nn.conv2d( lv, @@ -2108,8 +2074,6 @@ def main( index_dtype="int32", ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) ) lv3: R.Tensor(dtype="float32", ndim=5) = R.add(lv2, lv2_1) @@ -2119,8 +2083,6 @@ def main( lambda i0, i1, i2, i3, i4: (i0, i3 * 4 + i4, i1, i2), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) R.output(gv) return gv @@ -2437,8 +2399,6 @@ def main( lambda i0, i1, i2, i3: (i0 // 4, i2, i3, i1, i0 % 4), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) gv: R.Tensor((2, 26, 26, 1, 4), dtype="float32") = R.nn.conv2d( lv, @@ -2455,8 +2415,6 @@ def main( lambda i0, i1, i2, i3, i4: (i0, i3 * 4 + i4, i1, i2), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) R.output(gv2) return gv2 @@ -3422,8 +3380,6 @@ def main( lambda i0, i1, i2, i3: (i0, i2, i3, i1 // 4, i1 % 4), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) lv1: R.Tensor((1, 3, 3, 16, 4), dtype="float32") = R.layout_transform( w, @@ -3431,8 +3387,6 @@ def main( lambda i0, i1, i2, i3: (i0 // 4, i2, i3, i1, i0 % 4), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) gv: R.Tensor((1, 26, 26, 1, 4), dtype="float32") = R.nn.conv2d( lv, @@ -3452,8 +3406,6 @@ def main( lambda i0, i1, i2, i3, i4: (i0, i3 * 4 + i4, i1, i2), index_dtype="int32" ), pad_value=None, - axis_separators=[], - input_axis_separators=[], ) gv2: R.Tensor((4, 26, 26), dtype="float32") = R.squeeze(lv2, axis=[0]) R.output(gv2) diff --git a/tests/python/relax/test_transform_fuse_tir.py b/tests/python/relax/test_transform_fuse_tir.py index 006621e4a482..5108116abca9 100644 --- a/tests/python/relax/test_transform_fuse_tir.py +++ b/tests/python/relax/test_transform_fuse_tir.py @@ -15,7 +15,6 @@ # specific language governing permissions and limitations # under the License. -import pytest import tvm import tvm.testing @@ -2315,132 +2314,6 @@ def take( _check(Before, Before) -def test_fuse_with_axis_separators(): - @I.ir_module(s_tir=True) - class Before: - @T.prim_func(private=True, s_tir=True) - def add(a: T.handle, b: T.handle, c: T.handle): - A = T.match_buffer(a, [T.int64(16), T.int64(32)], "float32", axis_separators=[1]) - B = T.match_buffer(b, [T.int64(16), T.int64(32)], "float32", axis_separators=[1]) - C = T.match_buffer(c, [T.int64(16), T.int64(32)], "float32", axis_separators=[1]) - - for iters in T.grid(T.int64(16), T.int64(32)): - with T.sblock("compute"): - i, j = T.axis.remap("SS", iters) - C[i, j] = A[i, j] + B[i, j] - - @R.function(private=True) - def fused_function( - x: R.Tensor([T.int64(16), T.int64(32)], "float32"), - y: R.Tensor([T.int64(16), T.int64(32)], "float32"), - z: R.Tensor([T.int64(16), T.int64(32)], "float32"), - ) -> R.Tensor([T.int64(16), T.int64(32)], dtype="float32"): - R.func_attr({"Primitive": True}) - cls = Before - with R.dataflow(): - w = R.call_tir( - cls.add, [x, y], out_ty=R.Tensor([T.int64(16), T.int64(32)], "float32") - ) - out = R.call_tir( - cls.add, [w, z], out_ty=R.Tensor([T.int64(16), T.int64(32)], "float32") - ) - R.output(out) - return out - - @R.function - def main( - x: R.Tensor([T.int64(16), T.int64(32)], "float32"), - y: R.Tensor([T.int64(16), T.int64(32)], "float32"), - z: R.Tensor([T.int64(16), T.int64(32)], "float32"), - ) -> R.Tensor([T.int64(16), T.int64(32)], dtype="float32"): - cls = Before - with R.dataflow(): - gv = cls.fused_function(x, y, z) - R.output(gv) - return gv - - @I.ir_module(s_tir=True) - class Expected: - @T.prim_func(private=True, s_tir=True) - def fused_function(x: T.handle, y: T.handle, z: T.handle, c: T.handle): - T.func_attr({"tirx.noalias": True}) - X = T.match_buffer(x, [T.int64(16), T.int64(32)], "float32", axis_separators=[1]) - Y = T.match_buffer(y, [T.int64(16), T.int64(32)], "float32", axis_separators=[1]) - Z = T.match_buffer(z, [T.int64(16), T.int64(32)], "float32", axis_separators=[1]) - C = T.match_buffer(c, [T.int64(16), T.int64(32)], "float32", axis_separators=[1]) - Temp = T.sblock_alloc_buffer(X.shape, "float32", axis_separators=[1]) - for iters in T.grid(*X.shape): - with T.sblock("compute_Y"): - i, j = T.axis.remap("SS", iters) - Temp[i, j] = X[i, j] + Y[i, j] - - for iters in T.grid(*X.shape): - with T.sblock("compute_Z"): - i, j = T.axis.remap("SS", iters) - C[i, j] = Temp[i, j] + Z[i, j] - - @R.function - def main( - x: R.Tensor([T.int64(16), T.int64(32)], "float32"), - y: R.Tensor([T.int64(16), T.int64(32)], "float32"), - z: R.Tensor([T.int64(16), T.int64(32)], "float32"), - ) -> R.Tensor([T.int64(16), T.int64(32)], dtype="float32"): - cls = Expected - with R.dataflow(): - gv = R.call_tir( - cls.fused_function, - [x, y, z], - out_ty=R.Tensor([T.int64(16), T.int64(32)], "float32"), - ) - R.output(gv) - return gv - - _check(Before, Expected) - - -def test_fuse_with_axis_separators_inconsistent_buffer_mapping(): - @I.ir_module(s_tir=True) - class Before: - @T.prim_func(private=True, s_tir=True) - def mul(a: T.handle, b: T.handle, c: T.handle): - A = T.match_buffer(a, [T.int64(16), T.int64(32)], "float32", axis_separators=[1]) - B = T.match_buffer(b, [T.int64(16), T.int64(32)], "float32", axis_separators=[]) - C = T.match_buffer(c, [T.int64(16), T.int64(32)], "float32", axis_separators=[1]) - - for iters in T.grid(T.int64(16), T.int64(32)): - with T.sblock("compute"): - i, j = T.axis.remap("SS", iters) - C[i, j] = A[i, j] * B[i, j] - - @R.function(private=True) - def fused_function( - x: R.Tensor([T.int64(16), T.int64(32)], "float32"), - ) -> R.Tensor([T.int64(16), T.int64(32)], dtype="float32"): - R.func_attr({"Primitive": True}) - cls = Before - with R.dataflow(): - out = R.call_tir( - cls.mul, [x, x], out_ty=R.Tensor([T.int64(16), T.int64(32)], "float32") - ) - R.output(out) - return out - - @R.function - def main( - x: R.Tensor([T.int64(16), T.int64(32)], "float32"), - ) -> R.Tensor([T.int64(16), T.int64(32)], dtype="float32"): - cls = Before - with R.dataflow(): - gv = cls.fused_function(x) - R.output(gv) - return gv - - with pytest.raises( - RuntimeError, match=r"Inconsistent buffers.*and.*mapped to the same relax var:.*" - ): - relax.transform.FuseTIR()(Before) - - def test_block_name_numeric_suffix_deduplication(): @I.ir_module(s_tir=True) class Before: diff --git a/tests/python/relax/test_transform_legalize_ops_manipulate.py b/tests/python/relax/test_transform_legalize_ops_manipulate.py index 250e953f646a..ce70948157d6 100644 --- a/tests/python/relax/test_transform_legalize_ops_manipulate.py +++ b/tests/python/relax/test_transform_legalize_ops_manipulate.py @@ -1752,46 +1752,6 @@ def main(x: R.Tensor(("a", "b", "c"), dtype="float32")) -> R.Tensor(("a", "c", " tvm.ir.assert_structural_equal(mod, Expected) -def test_layout_transform_with_pad_axis_sep(): - transformation = lambda a, b, c: (a, c, b // 3, b % 3) - pad_value = 2 - axis_separator = [3] - - # fmt: off - @I.ir_module(s_tir=True) - class LayoutTransform: - @R.function - def main(x: R.Tensor((10, 20, 30), "float32")): - gv = R.layout_transform( - x, index_map=transformation, pad_value=pad_value, axis_separators=axis_separator, - ) - return gv - - @I.ir_module(s_tir=True) - class Expected: - @T.prim_func(private=True, s_tir=True) - def te_layout_transform_with_pad_axis_separator(A: T.Buffer((T.int64(10), T.int64(20), T.int64(30)), "float32"), var_te_layout_transform_with_pad_axis_separator: T.handle): - T.func_attr({"tirx.noalias": True}) - te_layout_transform_with_pad_axis_separator_1 = T.match_buffer(var_te_layout_transform_with_pad_axis_separator, (T.int64(10), T.int64(30), T.int64(7), T.int64(3)), axis_separators=[3]) - # with T.sblock("root"): - for axis0, axis1, axis2, axis3 in T.grid(T.int64(10), T.int64(30), T.int64(7), T.int64(3)): - with T.sblock("te_layout_transform_with_pad_axis_separator"): - v_axis0, v_axis1, v_axis2, v_axis3 = T.axis.remap("SSSS", [axis0, axis1, axis2, axis3]) - T.reads(A[v_axis0, v_axis2 * T.int64(3) + v_axis3, v_axis1]) - T.writes(te_layout_transform_with_pad_axis_separator_1[v_axis0, v_axis1, v_axis2, v_axis3]) - te_layout_transform_with_pad_axis_separator_1[v_axis0, v_axis1, v_axis2, v_axis3] = T.if_then_else(v_axis2 == T.int64(6) and v_axis3 == T.int64(2), T.float32(2), A[v_axis0, v_axis2 * T.int64(3) + v_axis3, v_axis1]) - - @R.function - def main(x: R.Tensor((10, 20, 30), dtype="float32")) -> R.Tensor((10, 30, 7, 3), dtype="float32"): - cls = Expected - gv = R.call_tir(cls.te_layout_transform_with_pad_axis_separator, (x,), out_ty=R.Tensor((10, 30, 7, 3), dtype="float32")) - return gv - # fmt: on - - mod = LegalizeOps()(LayoutTransform) - tvm.ir.assert_structural_equal(mod, Expected) - - def test_func_ty_of_legalized_layout_transform(): """PrimFunc shape information must be correct diff --git a/tests/python/s_tir/schedule/test_tir_schedule_set_axis_separator.py b/tests/python/s_tir/schedule/test_tir_schedule_set_axis_separator.py deleted file mode 100644 index 175d41a168c9..000000000000 --- a/tests/python/s_tir/schedule/test_tir_schedule_set_axis_separator.py +++ /dev/null @@ -1,208 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -# pylint: disable=missing-function-docstring,missing-module-docstring -# ruff: noqa: E501, F401 -import pytest - -import tvm -import tvm.testing -from tvm import tirx -from tvm.s_tir.schedule.testing import ( - assert_structural_equal_ignore_global_symbol, - verify_trace_roundtrip, -) -from tvm.script import ir as I -from tvm.script import tirx as T -from tvm.tirx import IndexMap - -# fmt: off -# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg - -@T.prim_func(s_tir=True) -def element_wise(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")) -> None: - B = T.sblock_alloc_buffer((128, 128), dtype="float32") - - for i, j in T.grid(128, 128): - with T.sblock("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.sblock("C"): - vi, vj = T.axis.remap("SS", [i, j]) - C[vi, vj] = B[vi, vj] + 1.0 - - -@T.prim_func(s_tir=True) -def element_wise_set_axis_separator(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")) -> None: - B = T.sblock_alloc_buffer([128, 128], dtype="float32", axis_separators=[1]) - - for i, j in T.grid(128, 128): - with T.sblock("B"): - vi, vj = T.axis.remap("SS", [i, j]) - B[vi, vj] = A[vi, vj] * T.float32(2) - for i, j in T.grid(128, 128): - with T.sblock("C"): - vi, vj = T.axis.remap("SS", [i, j]) - C[vi, vj] = B[vi, vj] + T.float32(1) - - -@T.prim_func(s_tir=True) -def element_wise_set_axis_separator_input_buffer(A: T.Buffer(shape=(128, 128), dtype="float32", axis_separators=(1,)), C: T.Buffer((128, 128), "float32")) -> None: - B = T.sblock_alloc_buffer([128, 128], dtype="float32") - - for i, j in T.grid(128, 128): - with T.sblock("B"): - vi, vj = T.axis.remap("SS", [i, j]) - B[vi, vj] = A[vi, vj] * T.float32(2) - for i, j in T.grid(128, 128): - with T.sblock("C"): - vi, vj = T.axis.remap("SS", [i, j]) - C[vi, vj] = B[vi, vj] + T.float32(1) - - -@T.prim_func(s_tir=True) -def element_wise_subregion_match(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")) -> None: - B = T.sblock_alloc_buffer((128, 128), dtype="float32") - - for i, j in T.grid(128, 128): - with T.sblock("B"): - vi, vj = T.axis.remap("SS", [i, j]) - B_subregion0 = T.match_buffer(B[vi, vj], [], offset_factor=1) - B_subregion0[()] = A[vi, vj] * 2.0 - for i, j in T.grid(128, 128): - with T.sblock("C"): - vi, vj = T.axis.remap("SS", [i, j]) - B_subregion1 = T.match_buffer(B[vi, vj], [], offset_factor=1) - C[vi, vj] = B_subregion1[()] + 1.0 - - -@T.prim_func(s_tir=True) -def element_wise_subregion_match_set_axis_separator(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")) -> None: - B = T.sblock_alloc_buffer([128, 128], dtype="float32", axis_separators=[1]) - - for i, j in T.grid(128, 128): - with T.sblock("B"): - vi, vj = T.axis.remap("SS", [i, j]) - B_subregion0 = T.match_buffer(B[vi, vj], [], dtype="float32", offset_factor=1, axis_separators=[0]) - B_subregion0[()] = A[vi, vj] * T.float32(2) - for i, j in T.grid(128, 128): - with T.sblock("C"): - vi, vj = T.axis.remap("SS", [i, j]) - B_subregion1 = T.match_buffer(B[vi, vj], [], dtype="float32", offset_factor=1, axis_separators=[0]) - C[vi, vj] = B_subregion1[()] + T.float32(1) - - -# pylint: enable=no-member,invalid-name,unused-variable,unexpected-keyword-arg - -argument_style = tvm.testing.parameter('set_axis_separators', - 'transform_layout_named', - 'transform_layout_buffer_object', - ) - - -def test_set_axis_separator(argument_style): - func = element_wise - s = tvm.s_tir.Schedule(func, debug_mask='all') - - if argument_style=='set_axis_separators': - s.set_axis_separator(s.get_sblock("B"), ("write",0), [1]) - elif argument_style=='transform_layout_named': - s.transform_layout(block='B', buffer='B', index_map=lambda i,j: [i,IndexMap.AXIS_SEPARATOR,j]) - elif argument_style =='transform_layout_buffer_object': - B = s.get(s.get_sblock('B')).writes[0].buffer - s.transform_layout(block='B', buffer=B, index_map=lambda i,j: [i,IndexMap.AXIS_SEPARATOR,j]) - else: - raise ValueError(f'Unexpected argument_style: {argument_style}') - - assert_structural_equal_ignore_global_symbol(element_wise_set_axis_separator, s.mod["main"]) - verify_trace_roundtrip(sch=s, mod=func) - - -def test_set_scope_fail_on_index_out_of_bound(): - func = element_wise - s = tvm.s_tir.Schedule(func, debug_mask='all') - with pytest.raises(AssertionError): - s.set_axis_separator(s.get_sblock("B"), ("write",1),[1]) - with pytest.raises(AssertionError): - s.set_axis_separator(s.get_sblock("B"), ("read",-1),[1]) - - -def test_set_axis_separator_input_buffer(argument_style): - func = element_wise - s = tvm.s_tir.Schedule(func, debug_mask='all') - - if argument_style=='set_axis_separators': - s.set_axis_separator(s.get_sblock("B"), ("read",0), [1]) - elif argument_style=='transform_layout_named': - s.transform_layout(block='B', buffer='A', index_map=lambda i,j: [i,IndexMap.AXIS_SEPARATOR,j]) - elif argument_style =='transform_layout_buffer_object': - A = s.get(s.get_sblock('B')).reads[0].buffer - s.transform_layout(block='B', buffer=A, index_map=lambda i,j: [i,IndexMap.AXIS_SEPARATOR,j]) - else: - raise ValueError(f'Unexpected argument_style: {argument_style}') - - - assert_structural_equal_ignore_global_symbol(element_wise_set_axis_separator_input_buffer, s.mod["main"]) - verify_trace_roundtrip(sch=s, mod=func) - - -def test_set_axis_separator_subregion(argument_style): - func = element_wise_subregion_match - s = tvm.s_tir.Schedule(func, debug_mask='all') - - if argument_style=='set_axis_separators': - s.set_axis_separator(s.get_sblock("B"), ("write",0), [1]) - elif argument_style=='transform_layout_named': - s.transform_layout(block='B', buffer='B', index_map=lambda i,j: [i,IndexMap.AXIS_SEPARATOR,j]) - elif argument_style =='transform_layout_buffer_object': - B = s.get(s.get_sblock('B')).writes[0].buffer - s.transform_layout(block='B', buffer=B, index_map=lambda i,j: [i,IndexMap.AXIS_SEPARATOR,j]) - else: - raise ValueError(f'Unexpected argument_style: {argument_style}') - - assert_structural_equal_ignore_global_symbol(element_wise_subregion_match_set_axis_separator, s.mod["main"]) - verify_trace_roundtrip(sch=s, mod=func) - -def test_indexed_lookup(): - @I.ir_module(s_tir=True) - class Before: - @T.prim_func(s_tir=True) - def main(): - A = T.sblock_alloc_buffer([4,4], dtype="int32") - B = T.sblock_alloc_buffer([1,1], dtype="int32") - for j in T.serial(4): - with T.sblock('block'): - A[B[0,0],j] = 0 - - @I.ir_module(s_tir=True) - class Expected: - @T.prim_func(s_tir=True) - def main(): - A = T.sblock_alloc_buffer([4,4], dtype="int32") - B = T.sblock_alloc_buffer([1,1], dtype="int32", axis_separators=[1]) - for j in T.serial(4): - with T.sblock('block'): - A[B[0,0],j] = 0 - - sch = tvm.s_tir.Schedule(Before) - sch.set_axis_separator('block', 'B', [1]) - After = sch.mod - tvm.ir.assert_structural_equal(After, Expected) - - -if __name__ == "__main__": - tvm.testing.main() diff --git a/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py b/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py index 0143e2f41945..1948814a506d 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py @@ -177,58 +177,6 @@ def two_elementwise_unit_dim(A: T.Buffer((1, 128), "float32"), C: T.Buffer((1, 1 vi, vj = T.axis.remap("SS", [i, j]) C[vi, vj] = B[vi, vj] + 1.0 -def test_transform_layout_with_cache_write_and_axis_separators(): - """ - transform_layout with axis_separator on a buffer from cache_write should work as expected - """ - - @I.ir_module(s_tir=True) - class Before: - @T.prim_func(s_tir=True) - def main( - p0: T.Buffer((T.int64(33), T.int64(128)), "float32"), - p1: T.Buffer((T.int64(33), T.int64(128)), "float32"), - T_add: T.Buffer((T.int64(33), T.int64(128)), "float32"), - ): - T.func_attr({"global_symbol": "main", "tirx.noalias": True}) - # with T.sblock("root"): - for ax0, ax1 in T.grid(T.int64(33), T.int64(128)): - with T.sblock("T_add"): - v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) - T.reads(p0[v_ax0, v_ax1], p1[v_ax0, v_ax1]) - T.writes(T_add[v_ax0, v_ax1]) - T_add[v_ax0, v_ax1] = p0[v_ax0, v_ax1] + p1[v_ax0, v_ax1] - - @I.ir_module(s_tir=True) - class Expected: - @T.prim_func(s_tir=True) - def main(p0: T.Buffer((T.int64(33), T.int64(128)), "float32"), p1: T.Buffer((T.int64(33), T.int64(128)), "float32"), T_add: T.Buffer((T.int64(33), T.int64(128)), "float32")): - T.func_attr({"global_symbol": "main", "tirx.noalias": True}) - # with T.sblock("root"): - T_add_global = T.sblock_alloc_buffer((T.int64(2), T.int64(128), T.int64(32)), axis_separators=[2]) - for axis0, axis1, axis2 in T.grid(T.int64(2), T.int64(128), T.int64(32)): - with T.sblock("T_add"): - v_axis0, v_axis1, v_axis2 = T.axis.remap("SSS", [axis0, axis1, axis2]) - T.reads(p0[v_axis0 * T.int64(32) + v_axis2, v_axis1], p1[v_axis0 * T.int64(32) + v_axis2, v_axis1]) - T.writes(T_add_global[v_axis0, v_axis1, v_axis2]) - T_add_global[v_axis0, v_axis1, v_axis2] = T.if_then_else(v_axis0 == T.int64(1) and T.int64(1) <= v_axis2, T.float32(0), p0[v_axis0 * T.int64(32) + v_axis2, v_axis1] + p1[v_axis0 * T.int64(32) + v_axis2, v_axis1]) - for ax0, ax1 in T.grid(T.int64(33), T.int64(128)): - with T.sblock("T_add_global"): - v0, v1 = T.axis.remap("SS", [ax0, ax1]) - T.reads(T_add_global[v0 // T.int64(32), v1, v0 % T.int64(32)]) - T.writes(T_add[v0, v1]) - T_add[v0, v1] = T_add_global[v0 // T.int64(32), v1, v0 % T.int64(32)] - - def transform_fn(x, y): - return [x // 32, y, tvm.te.AXIS_SEPARATOR, x % 32] - - sch = tvm.s_tir.Schedule(Before, debug_mask="all") - block_rv = sch.get_sblock("T_add") - sch.cache_write(block_rv, 0, "global") - sch.transform_layout(block_rv, ("write", 0), transform_fn, pad_value=0.0) - After = sch.mod - tvm.ir.assert_structural_equal(After, Expected) - # pylint: enable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks # fmt: on @@ -1129,72 +1077,6 @@ def main(A: T.Buffer(16, "int32"), n: T.int32): tvm.ir.assert_structural_equal(After, Expected) -def test_transform_with_axis_separators(): - """Axis separators may be specified in a transform""" - - @I.ir_module(s_tir=True) - class Before: - @T.prim_func(s_tir=True) - def main(a: T.handle): - A = T.match_buffer(a, [14], "int32") - for i in T.serial(14): - with T.sblock("block"): - vi = T.axis.remap("S", [i]) - A[vi] = 42 - - @I.ir_module(s_tir=True) - class Expected: - @T.prim_func(s_tir=True) - def main(a: T.handle): - A = T.match_buffer(a, [4, 4], "int32", axis_separators=[1]) - for i, j in T.grid(4, 4): - with T.sblock("block"): - vi, vj = T.axis.remap("SS", [i, j]) - A[vi, vj] = T.if_then_else(vi == 3 and 2 <= vj, 0, 42, dtype="int32") - - sch = tvm.s_tir.Schedule(Before) - sch.transform_layout( - "block", - "A", - lambda i: [i // 4, tvm.tirx.IndexMap.AXIS_SEPARATOR, i % 4], - pad_value=0, - ) - After = sch.mod - tvm.ir.assert_structural_equal(After, Expected) - - -def test_transform_with_axis_separators_opaque_block(): - """Axis separators may be specified in a transform of opaque block""" - - @I.ir_module(s_tir=True) - class Before: - @T.prim_func(s_tir=True) - def main(a: T.handle): - A = T.match_buffer(a, [14], "int32") - for i in T.serial(14): - with T.sblock("block"): - A[i] = 42 - - @I.ir_module(s_tir=True) - class Expected: - @T.prim_func(s_tir=True) - def main(a: T.handle): - A = T.match_buffer(a, [4, 4], "int32", axis_separators=[1]) - for i, j in T.grid(4, 4): - with T.sblock("block"): - A[i, j] = T.if_then_else(i == 3 and 2 <= j, 0, 42, dtype="int32") - - sch = tvm.s_tir.Schedule(Before) - sch.transform_layout( - "block", - "A", - lambda i: [i // 4, tvm.tirx.IndexMap.AXIS_SEPARATOR, i % 4], - pad_value=0, - ) - After = sch.mod - tvm.ir.assert_structural_equal(After, Expected) - - def test_index_map_dtype_legalize(): """Test dtype legalization of the index map indices.""" diff --git a/tests/python/tirx-base/test_tir_buffer.py b/tests/python/tirx-base/test_tir_buffer.py index 560742a4ad01..42e24ce127a5 100644 --- a/tests/python/tirx-base/test_tir_buffer.py +++ b/tests/python/tirx-base/test_tir_buffer.py @@ -202,18 +202,5 @@ def test_buffer_flatten_preserves_identity(): assert buf.same_as(flat) -def test_buffer_flatten_uses_axis_separators(): - """Flattening to N-d physical buffers uses the axis separators""" - buf = tvm.tirx.decl_buffer([4, 16, 32], axis_separators=[2]) - flat = buf.get_flattened_buffer() - tvm.ir.assert_structural_equal(flat.axis_separators, [T.int32(1)]) - tvm.ir.assert_structural_equal(flat.shape, [T.int32(4 * 16), T.int32(32)]) - - -def test_invalid_axis_separators_raises_exception(): - with pytest.raises(ValueError): - tvm.tirx.decl_buffer([1], axis_separators=[1, 2]) - - if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/tirx-base/test_tir_intrin.py b/tests/python/tirx-base/test_tir_intrin.py index 4765bd2f9404..6b31133b1307 100644 --- a/tests/python/tirx-base/test_tir_intrin.py +++ b/tests/python/tirx-base/test_tir_intrin.py @@ -357,7 +357,6 @@ def test_tir_fma(A: T.handle, B: T.handle, C: T.handle, d: T.handle) -> None: elem_offset=0, align=64, offset_factor=1, - buffer_type="auto", ) B_1 = T.match_buffer( B, @@ -366,7 +365,6 @@ def test_tir_fma(A: T.handle, B: T.handle, C: T.handle, d: T.handle) -> None: elem_offset=0, align=64, offset_factor=1, - buffer_type="auto", ) C_1 = T.match_buffer( C, @@ -375,7 +373,6 @@ def test_tir_fma(A: T.handle, B: T.handle, C: T.handle, d: T.handle) -> None: elem_offset=0, align=64, offset_factor=1, - buffer_type="auto", ) d_1 = T.match_buffer( d, @@ -384,7 +381,6 @@ def test_tir_fma(A: T.handle, B: T.handle, C: T.handle, d: T.handle) -> None: elem_offset=0, align=64, offset_factor=1, - buffer_type="auto", ) # body for i in T.serial(0, n): diff --git a/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py b/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py index a43715665532..ea2f9dbb3f18 100644 --- a/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py +++ b/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py @@ -368,73 +368,5 @@ def main(): tvm.ir.assert_structural_equal(After, Expected) -def test_no_change_to_2d_physical_buffer(): - """Flattening preserves axis separators.""" - - @I.ir_module(s_tir=True) - class Before: - @T.prim_func(s_tir=True) - def main(): - A = T.sblock_alloc_buffer([32, 32], axis_separators=[1]) - for i, j in T.grid(32, 32): - T.evaluate(A[i, j]) - - Expected = Before - - After = _transform()(Before) - tvm.ir.assert_structural_equal(After, Expected) - - -def test_flatten_alloc_buffer_with_axis_separators(): - """Flattening preserves axis separators""" - - @I.ir_module(s_tir=True) - class Before: - @T.prim_func(s_tir=True) - def main(): - A = T.sblock_alloc_buffer([2, 3, 5, 7, 11, 13], axis_separators=[3]) - for i0, i1, i2, i3, i4, i5 in T.grid(2, 3, 5, 7, 11, 13): - T.evaluate(A[i0, i1, i2, i3, i4, i5]) - - @I.ir_module(s_tir=True) - class Expected: - @T.prim_func(s_tir=True) - def main(): - A = T.sblock_alloc_buffer([30, 1001], axis_separators=[1]) - for i0, i1, i2, i3, i4, i5 in T.grid(2, 3, 5, 7, 11, 13): - T.evaluate(A[i0 * 15 + i1 * 5 + i2, i3 * 143 + i4 * 13 + i5]) - - After = _transform()(Before) - tvm.ir.assert_structural_equal(After, Expected) - - -def test_flatten_decl_buffer_with_axis_separators(): - """Flattening preserves axis separators - - Like test_flatten_alloc_buffer_with_axis_separators, but the allocations - is done using Allocate/DeclBuffer, rather than through - BlockNode::alloc_buffers. - """ - - @I.ir_module(s_tir=True) - class Before: - @T.prim_func(s_tir=True) - def main(): - A = T.decl_buffer([2, 3, 5, 7, 11, 13], axis_separators=[3]) - for i0, i1, i2, i3, i4, i5 in T.grid(2, 3, 5, 7, 11, 13): - T.evaluate(A[i0, i1, i2, i3, i4, i5]) - - @I.ir_module(s_tir=True) - class Expected: - @T.prim_func(s_tir=True) - def main(): - A = T.decl_buffer([30, 1001], axis_separators=[1]) - for i0, i1, i2, i3, i4, i5 in T.grid(2, 3, 5, 7, 11, 13): - T.evaluate(A[i0 * 15 + i1 * 5 + i2, i3 * 143 + i4 * 13 + i5]) - - After = _transform()(Before) - tvm.ir.assert_structural_equal(After, Expected) - - if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py b/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py index 28c5819df289..2e519c229e8a 100644 --- a/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py +++ b/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py @@ -407,98 +407,6 @@ def main() -> None: tvm.ir.assert_structural_equal(After, Expected) -def test_rewrite_in_place_use_of_non_flat_buffer(): - """A non-flat buffer may be re-used for in-place operations""" - - @I.ir_module - class Before: - @T.prim_func(s_tir=True) - def main(A: T.Buffer((16, 16), "float32"), D: T.Buffer((16, 16), "float32")): - B = T.decl_buffer( - [16, 16], - dtype="float32", - axis_separators=[1], - ) - C = T.decl_buffer( - [16, 16], - dtype="float32", - axis_separators=[1], - ) - - for i, j in T.grid(16, 16): - B[i, j] = A[i, j] - - for i, j in T.grid(16, 16): - C[i, j] = 2.0 * B[i, j] - - for i, j in T.grid(16, 16): - D[i, j] = C[i, j] - - @I.ir_module - class Expected: - @T.prim_func(s_tir=True) - def main(A: T.Buffer((16, 16), "float32"), D: T.Buffer((16, 16), "float32")): - B = T.decl_buffer([16, 16], dtype="float32", axis_separators=[1]) - C = T.decl_buffer( - [16, 16], - dtype="float32", - axis_separators=[1], - data=B.data, - ) - - for i, j in T.grid(16, 16): - B[i, j] = A[i, j] - - for i, j in T.grid(16, 16): - C[i, j] = 2.0 * B[i, j] - - for i, j in T.grid(16, 16): - D[i, j] = C[i, j] - - After = tvm.tirx.transform.StorageRewrite()(Before) - tvm.ir.assert_structural_equal(After, Expected) - - -def test_no_rewrite_of_shared_non_flat_buffer(): - """In general, sharing of non-flat buffer isn't supported - - The current packing algorithms in StorageRewrite assume a flat - memory space, and do not support packing of N-d buffers. For - buffers with axis separators, normal buffer sharing should be - disabled. - - Like test_rewrite_in_place_use_of_non_flat_buffer, except that B and C do - not have matching shapes. - """ - - @T.prim_func(s_tir=True) - def Before(A: T.Buffer((16, 16), "float32"), D: T.Buffer((16, 16), "float32")): - B = T.decl_buffer( - [16, 16], - dtype="float32", - axis_separators=[1], - ) - C = T.decl_buffer( - [20, 20], - dtype="float32", - axis_separators=[1], - ) - - for i, j in T.grid(16, 16): - B[i, j] = A[i, j] - - for i, j in T.grid(16, 16): - C[i, j] = 2.0 * B[i, j] - - for i, j in T.grid(16, 16): - D[i, j] = C[i, j] - - Expected = Before - - After = tvm.tirx.transform.StorageRewrite()(tvm.IRModule.from_expr(Before)) - tvm.ir.assert_structural_equal(After["Before"], Expected) - - def test_rewrite_decl_buffer(): """A DeclBuffer node may appear in StorageRewrite's input""" diff --git a/tests/python/tvmscript/test_tvmscript_roundtrip.py b/tests/python/tvmscript/test_tvmscript_roundtrip.py index 71aedaecbeb3..94d4169e382f 100644 --- a/tests/python/tvmscript/test_tvmscript_roundtrip.py +++ b/tests/python/tvmscript/test_tvmscript_roundtrip.py @@ -2416,25 +2416,6 @@ def func_with_ptr_type_annotations(x: T.handle("int32"), y: T.handle("int32", "s return func_with_ptr_type_annotations -def buffer_axis_separator(): - @T.prim_func(s_tir=True) - def element_wise(a: T.handle, c: T.handle) -> None: - A = T.match_buffer(a, (128, 128), "float32", axis_separators=[1]) - C = T.match_buffer(c, (128, 128), "float32") - B = T.sblock_alloc_buffer((128, 128), "float32", axis_separators=[1]) - - for i, j in T.grid(128, 128): - with T.sblock("B"): - vi, vj = T.axis.remap("SS", [i, j]) - B[vi, vj] = A[vi, vj] * T.float32(2) - for i, j in T.grid(128, 128): - with T.sblock("C"): - vi, vj = T.axis.remap("SS", [i, j]) - C[vi, vj] = B[vi, vj] + T.float32(1) - - return element_wise - - def buffer_ramp_access_as_slice_index(): @T.prim_func(s_tir=True) def buffer_ramp_access(a: T.handle, b: T.handle, c: T.handle) -> None: @@ -2769,75 +2750,6 @@ def func(): return func -def string_stride(): - @T.prim_func(s_tir=True) - def main(a: T.handle, b: T.handle): - T.func_attr({"global_symbol": "main", "tirx.noalias": True}) - n = T.int32() - A = T.match_buffer(a, (n,), strides=("A_s0",), buffer_type="auto") - B = T.match_buffer(b, (n,), strides=("B_s0",), buffer_type="auto") - blockIdx_x = T.launch_thread("blockIdx.x", (n + 63) // 64) - threadIdx_x = T.launch_thread("threadIdx.x", 64) - if T.likely(blockIdx_x * 64 + threadIdx_x < n): - B2 = T.decl_buffer((B.strides[0] * n,), data=B.data) - A2 = T.decl_buffer((A.strides[0] * n,), data=A.data) - B2[(blockIdx_x * 64 + threadIdx_x) * B.strides[0]] = A2[ - (blockIdx_x * 64 + threadIdx_x) * A.strides[0] - ] * T.float32(2) - - return main - - -def string_stride_int64(): - @T.prim_func(s_tir=True) - def main(a: T.handle, b: T.handle): - T.func_attr({"global_symbol": "main", "tirx.noalias": True}) - n = T.int64() - A_s0 = T.int64() - B_s0 = T.int64() - A = T.match_buffer(a, (n,), strides=(A_s0,), buffer_type="auto") - B = T.match_buffer(b, (n,), strides=(B_s0,), buffer_type="auto") - for i in range(n): - B[i] = A[i] - - return main - - -def merge_shape_var_def(): - # uninitialized vars - @T.prim_func(check_well_formed=False, s_tir=True) - def main(A: T.handle, B: T.handle): - # fmt: off - T.func_attr({"global_symbol": "main", "tirx.noalias": True}) - m, n = T.int32(), T.int32() - A_1 = T.match_buffer(A, (m, n), strides=("A_1_s0", "A_1_s1"), buffer_type="auto") - B_1 = T.match_buffer(B, (m, n), strides=("B_1_s0", "B_1_s1"), buffer_type="auto") - for i_outer, j_outer, i_inner in T.grid((m + 9) // 10, (n + 4) // 5, 10): - if T.likely(i_outer * 10 + i_inner < m): - for j_inner in range(5): - if T.likely(j_outer * 5 + j_inner < n): - cse_v2: T.let[T.int32] = j_outer * 5 + j_inner - cse_v1: T.let[T.int32] = i_outer * 10 + i_inner - B_2 = T.decl_buffer( - (B_1.strides[0] * m,), - data=B_1.data, - strides=("B_2_s0",), - buffer_type="auto", - ) - A_2 = T.decl_buffer( - (A_1.strides[0] * m,), - data=A_1.data, - strides=("A_2_s0",), - buffer_type="auto", - ) - B_2[cse_v1 * B_1.strides[0] + cse_v2 * B_1.strides[1]] = A_2[ - cse_v1 * A_1.strides[0] + cse_v2 * A_1.strides[1] - ] - # fmt: on - - return main - - def if_then_else_var(): @T.prim_func(s_tir=True) def main(n: T.int32): @@ -3335,7 +3247,6 @@ def func(value: R.Prim("float16")): int64_support, string_annotation_escaping, pointer_type, - buffer_axis_separator, buffer_ramp_access_as_slice_index, ramp_int64, scalable_vectors, @@ -3361,9 +3272,6 @@ def func(value: R.Prim("float16")): multi_env_threads, intrinsic_pow, bind_var, - string_stride, - string_stride_int64, - merge_shape_var_def, if_then_else_var, tvm_shfl_builtins, make_packed_api_result,