tests/python/tirx-base/test_tir_buffer.py - tvm - Git at Google

 # 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.
 # ruff: noqa: E741, F401, F841

 import numpy as np
 import pytest

 import tvm
 import tvm.testing
 from tvm.script import tirx as T
 from tvm.tirx import BufferAccessKind


 def test_buffer():
     m = tvm.tirx.Var("m", "int32")
     n = tvm.tirx.Var("n", "int32")
     l = tvm.tirx.Var("l", "int32")
     Ab = tvm.tirx.decl_buffer((m, n), "float32")
     Bb = tvm.tirx.decl_buffer((n, l), "float32")

     assert type(Ab) is tvm.ir.Var
     assert tvm.tirx.is_buffer_var(Ab)
     assert isinstance(Ab.ty, tvm.tirx.BufferType)
     assert Ab.ty.dtype == tvm.ir.PrimType("float32")
     assert tuple(Ab.ty.shape) == (m, n)
     assert not tvm.tirx.is_buffer_var(m)


 def test_buffer_compatibility_alias_and_global_var_properties():
     scalar = tvm.ir.Var("scalar", tvm.ir.PrimType("int32"))
     buffer = tvm.tirx.decl_buffer((8,), "float32")

     assert tvm.tirx.Buffer is tvm.ir.Var
     assert isinstance(scalar, tvm.tirx.Buffer)
     assert not tvm.tirx.is_buffer_var(scalar)
     assert tvm.tirx.is_buffer_var(buffer)

     assert tuple(buffer.shape) == (8,)
     assert buffer.dtype == tvm.DataType("float32")
     assert buffer.data.args[0].same_as(buffer)

     for name in ("shape", "dtype", "data"):
         assert not hasattr(scalar, name)
         with pytest.raises(AttributeError, match="only available on a Var with BufferType"):
             getattr(scalar, name)


 def test_buffer_data_is_typed_projection():
     buffer = tvm.tirx.decl_buffer((8,), "bool", scope="shared")

     assert buffer.ty.dtype == tvm.ir.PrimType("bool")
     assert tvm.tirx.buffer_data_pointer_type(buffer) == tvm.ir.PointerType(
         tvm.ir.PrimType("bool"), "shared"
     )
     assert buffer.data.op.name == "tirx.buffer_data"
     assert buffer.data.args[0].same_as(buffer)
     assert buffer.data.ty == tvm.tirx.buffer_data_pointer_type(buffer)


 def test_buffer_pointer_type_derived_from_dtype_and_scope():
     data = tvm.ir.Var("storage", tvm.ir.PointerType(tvm.ir.PrimType("uint8"), "local"))
     buffer = tvm.tirx.decl_buffer((8,), "float16", data=data)

     assert buffer.ty.dtype == tvm.ir.PrimType("float16")
     assert buffer.ty.storage_scope == "local"
     assert buffer.data.ty == tvm.ir.PointerType(tvm.ir.PrimType("float16"), "local")


 def test_decl_buffer_requires_physical_data_binding():
     buffer = tvm.tirx.decl_buffer((8,), "float32")
     data = tvm.tirx.Var("data", buffer.data.ty)

     with pytest.raises(TypeError, match="requires a physical data binding"):
         tvm.tirx.DeclBuffer(buffer)

     decl = tvm.tirx.DeclBuffer(buffer, data=data)
     assert decl.buffer.same_as(buffer)
     assert decl.data.same_as(data)


 def test_buffer_access_ptr():
     m = tvm.tirx.Var("m", "int32")
     n = tvm.tirx.Var("n", "int32")
     Ab = tvm.tirx.decl_buffer((m, n), "float32", strides=[n + 1, 1])
     aptr = Ab.access_ptr("rw")
     assert isinstance(aptr.ty, tvm.ir.PointerType)
     assert aptr.ty.element_type == tvm.ir.PrimType("void")
     tvm.ir.assert_structural_equal(aptr.args[3], Ab.ty.strides[0] * m)
     assert aptr.args[0].ty == Ab.ty.dtype
     assert aptr.args[4].value == BufferAccessKind.READ | BufferAccessKind.WRITE
     typed_ptr = Ab.access_ptr("r", ptr_type="uint8")
     assert typed_ptr.ty == tvm.ir.PointerType(tvm.ir.PrimType("uint8"))
     shared = tvm.tirx.decl_buffer((m, n), "float32", scope="shared")
     assert shared.access_ptr("r").ty == tvm.ir.PointerType(tvm.ir.PrimType("void"), "shared")
     assert shared.access_ptr("r", ptr_type="uint8").ty == tvm.ir.PointerType(
         tvm.ir.PrimType("uint8"), "shared"
     )
     aptr = Ab.access_ptr("w")
     assert aptr.args[4].value == BufferAccessKind.WRITE


 def test_buffer_access_ptr_offset():
     m = tvm.tirx.Var("m", "int32")
     n = tvm.tirx.Var("n", "int32")
     Ab = tvm.tirx.decl_buffer((m, n), "float32")
     aptr = Ab.access_ptr("rw", offset=100)
     tvm.testing.assert_prim_expr_equal(aptr.args[2], 100)
     assert aptr.args[4].value == BufferAccessKind.READ | BufferAccessKind.WRITE
     v = tvm.tirx.Var("int32", "int32")
     aptr = Ab.access_ptr("rw", offset=100 + 100 + v)
     tvm.testing.assert_prim_expr_equal(aptr.args[2], 200 + v)
     assert aptr.args[4].value == BufferAccessKind.READ | BufferAccessKind.WRITE
     aptr = Ab.access_ptr("rw", offset=tvm.tirx.call_extern("int32", "test_call", 100 + 100 + v))
     tvm.testing.assert_prim_expr_equal(
         aptr.args[2], tvm.tirx.call_extern("int32", "test_call", 200 + v)
     )
     assert aptr.args[4].value == BufferAccessKind.READ | BufferAccessKind.WRITE


 def test_buffer_access_ptr_extent():
     m = tvm.tirx.Var("m", "int32")
     n = tvm.tirx.Var("n", "int32")
     Ab = tvm.tirx.decl_buffer((m, n), "float32")
     aptr = Ab.access_ptr("rw")
     tvm.ir.assert_structural_equal(aptr.args[3], m * n)
     aptr = Ab.access_ptr("rw", offset=100)
     tvm.ir.assert_structural_equal(aptr.args[3], m * n - 100)
     Ab = tvm.tirx.decl_buffer((m, n), "float32", strides=[n + 1, 1])
     aptr = Ab.access_ptr("rw", offset=100)
     tvm.ir.assert_structural_equal(aptr.args[3], Ab.ty.strides[0] * m - 100)

     # Test extent from input params
     aptr = Ab.access_ptr("rw", extent=200)
     tvm.ir.assert_structural_equal(aptr.args[3], T.int32(200))
     aptr = Ab.access_ptr("rw", offset=100, extent=100)
     tvm.ir.assert_structural_equal(aptr.args[3], T.int32(100))


 def test_buffer_vload():
     m = tvm.tirx.Var("m", "int32")
     n = tvm.tirx.Var("n", "int32")
     Ab = tvm.tirx.decl_buffer((m, n), "float32", elem_offset=100)
     load = Ab.vload([2, 3])
     tvm.ir.assert_structural_equal(load.indices, [T.int32(2), T.int32(3)])


 def test_buffer_offset_of():
     m = tvm.tirx.Var("m", "int32")
     n = tvm.tirx.Var("n", "int32")
     Ab = tvm.tirx.decl_buffer((m, n), "float32", elem_offset=100)
     offset = Ab.offset_of([2, 3])
     tvm.ir.assert_structural_equal(offset, [n * 2 + 103])


 def test_buffer_index_merge_mult_mod():
     m = tvm.tirx.Var("m", "int32")
     n = tvm.tirx.Var("n", "int32")
     s = tvm.tirx.Var("s", "int32")
     k0 = tvm.tirx.Var("k0", "int32")
     k1 = tvm.tirx.Var("k1", "int32")
     A = tvm.tirx.decl_buffer((m, n), "float32")
     A_stride = tvm.tirx.decl_buffer((m, n), "float32", strides=(s, 1))

     def assert_simplified_equal(index_simplified, index_direct):
         (
             tvm.ir.assert_structural_equal(index_simplified, index_direct),
             f"index_simplified={index_simplified}, index_direct={index_direct}",
         )

     idxd = tvm.tirx.indexdiv
     idxm = tvm.tirx.indexmod

     # Test Case1
     index_simplified = A_stride.offset_of(
         (idxd(idxm(k0, k1), s), idxm(idxm(k0, k1), s) + idxd(k0, k1) * k1)
     )
     index_direct = A_stride.offset_of((0, k0))
     assert_simplified_equal(index_simplified, index_direct)

     # Test Case2
     index_simplified = A.offset_of(
         (idxd(idxm(k0, idxd(k1, s)), n), idxm(idxm(k0, idxd(k1, s)), n) + idxm(k0, k1))
     )
     index_direct = A.offset_of((0, idxm(k0, idxd(k1, s)) + idxm(k0, k1)))
     assert_simplified_equal(index_simplified, index_direct)
     # Test Case3
     index_simplified = A.offset_of(
         (
             idxd((idxd(k0, idxd(k1, s)) * idxd(k1, s)), n) + idxd(idxm(k0, idxd(k1, s)), n),
             idxm((idxd(k0, idxd(k1, s)) * idxd(k1, s)), n) + idxm(idxm(k0, idxd(k1, s)), n),
         )
     )
     index_direct = A.offset_of((0, k0))
     assert_simplified_equal(index_simplified, index_direct)
     # Test Case4 (not able to simplify)
     index_simplified = A.offset_of(
         (idxd(idxm(k0, idxd(k1, s)), n), idxm(idxm(k0, idxd(k1, n)), n) + idxm(k0, k1))
     )
     index_direct = A.offset_of(
         (0, idxd(idxm(k0, idxd(k1, s)), n) * n + (idxm(idxm(k0, idxd(k1, n)), n) + idxm(k0, k1)))
     )
     assert_simplified_equal(index_simplified, index_direct)

     # Test Case5
     B = tvm.tirx.decl_buffer((1, 14, 14, 1024))
     i = tvm.tirx.Var("i", "int32")
     j = tvm.tirx.Var("j", "int32")
     k = tvm.tirx.Var("k", "int32")

     index_simplified1 = B.offset_of(
         (
             idxd(idxd(idxd((i * 50176 + j * 28672 + k), 1024), 14), 14),
             idxm(idxd(idxd((i * 50176 + j * 28672 + k), 1024), 14), 14),
             idxm(idxd((i * 50176 + j * 28672 + k), 1024), 14),
             idxm((i * 50176 + j * 28672 + k), 1024),
         )
     )
     index_simplified2 = B.offset_of(
         (
             idxd(idxd(i * 49 + j * 28 + idxd(k, 1024), 14), 14),
             idxm(idxd(i * 49 + j * 28 + idxd(k, 1024), 14), 14),
             idxm(i * 7 + idxd(k, 1024), 14),
             idxm(k, 1024),
         )
     )
     index_direct = B.offset_of((0, 0, 0, (i * 50176 + j * 28672 + k)))
     assert_simplified_equal(index_simplified1, index_direct)
     assert_simplified_equal(index_simplified2, index_direct)


 def test_buffer_flatten():
     """A buffer should flatten to a 1-d shape"""
     buf = tvm.tirx.decl_buffer([16, 32])
     flat = buf.get_flattened_buffer()
     # A metadata-changing rewrite creates a fresh typed Var.  The physical
     # pointer is always derived from that Var instead of being stored as a
     # second buffer identity.
     assert not buf.same_as(flat)
     assert flat.data.args[0].same_as(flat)
     assert flat.data.op.name == "tirx.buffer_data"
     tvm.ir.assert_structural_equal(flat.ty.shape, [T.int32(16 * 32)])


 def test_buffer_flatten_preserves_identity():
     """Flattening a 1-d buffer should return the original"""
     buf = tvm.tirx.decl_buffer([16])
     flat = buf.get_flattened_buffer()
     assert buf.same_as(flat)


 if __name__ == "__main__":
     tvm.testing.main()
	# 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.
	# ruff: noqa: E741, F401, F841

	import numpy as np
	import pytest

	import tvm
	import tvm.testing
	from tvm.script import tirx as T
	from tvm.tirx import BufferAccessKind


	def test_buffer():
	m = tvm.tirx.Var("m", "int32")
	n = tvm.tirx.Var("n", "int32")
	l = tvm.tirx.Var("l", "int32")
	Ab = tvm.tirx.decl_buffer((m, n), "float32")
	Bb = tvm.tirx.decl_buffer((n, l), "float32")

	assert type(Ab) is tvm.ir.Var
	assert tvm.tirx.is_buffer_var(Ab)
	assert isinstance(Ab.ty, tvm.tirx.BufferType)
	assert Ab.ty.dtype == tvm.ir.PrimType("float32")
	assert tuple(Ab.ty.shape) == (m, n)
	assert not tvm.tirx.is_buffer_var(m)


	def test_buffer_compatibility_alias_and_global_var_properties():
	scalar = tvm.ir.Var("scalar", tvm.ir.PrimType("int32"))
	buffer = tvm.tirx.decl_buffer((8,), "float32")

	assert tvm.tirx.Buffer is tvm.ir.Var
	assert isinstance(scalar, tvm.tirx.Buffer)
	assert not tvm.tirx.is_buffer_var(scalar)
	assert tvm.tirx.is_buffer_var(buffer)

	assert tuple(buffer.shape) == (8,)
	assert buffer.dtype == tvm.DataType("float32")
	assert buffer.data.args[0].same_as(buffer)

	for name in ("shape", "dtype", "data"):
	assert not hasattr(scalar, name)
	with pytest.raises(AttributeError, match="only available on a Var with BufferType"):
	getattr(scalar, name)


	def test_buffer_data_is_typed_projection():
	buffer = tvm.tirx.decl_buffer((8,), "bool", scope="shared")

	assert buffer.ty.dtype == tvm.ir.PrimType("bool")
	assert tvm.tirx.buffer_data_pointer_type(buffer) == tvm.ir.PointerType(
	tvm.ir.PrimType("bool"), "shared"
	)
	assert buffer.data.op.name == "tirx.buffer_data"
	assert buffer.data.args[0].same_as(buffer)
	assert buffer.data.ty == tvm.tirx.buffer_data_pointer_type(buffer)


	def test_buffer_pointer_type_derived_from_dtype_and_scope():
	data = tvm.ir.Var("storage", tvm.ir.PointerType(tvm.ir.PrimType("uint8"), "local"))
	buffer = tvm.tirx.decl_buffer((8,), "float16", data=data)

	assert buffer.ty.dtype == tvm.ir.PrimType("float16")
	assert buffer.ty.storage_scope == "local"
	assert buffer.data.ty == tvm.ir.PointerType(tvm.ir.PrimType("float16"), "local")


	def test_decl_buffer_requires_physical_data_binding():
	buffer = tvm.tirx.decl_buffer((8,), "float32")
	data = tvm.tirx.Var("data", buffer.data.ty)

	with pytest.raises(TypeError, match="requires a physical data binding"):
	tvm.tirx.DeclBuffer(buffer)

	decl = tvm.tirx.DeclBuffer(buffer, data=data)
	assert decl.buffer.same_as(buffer)
	assert decl.data.same_as(data)


	def test_buffer_access_ptr():
	m = tvm.tirx.Var("m", "int32")
	n = tvm.tirx.Var("n", "int32")
	Ab = tvm.tirx.decl_buffer((m, n), "float32", strides=[n + 1, 1])
	aptr = Ab.access_ptr("rw")
	assert isinstance(aptr.ty, tvm.ir.PointerType)
	assert aptr.ty.element_type == tvm.ir.PrimType("void")
	tvm.ir.assert_structural_equal(aptr.args[3], Ab.ty.strides[0] * m)
	assert aptr.args[0].ty == Ab.ty.dtype
	assert aptr.args[4].value == BufferAccessKind.READ \| BufferAccessKind.WRITE
	typed_ptr = Ab.access_ptr("r", ptr_type="uint8")
	assert typed_ptr.ty == tvm.ir.PointerType(tvm.ir.PrimType("uint8"))
	shared = tvm.tirx.decl_buffer((m, n), "float32", scope="shared")
	assert shared.access_ptr("r").ty == tvm.ir.PointerType(tvm.ir.PrimType("void"), "shared")
	assert shared.access_ptr("r", ptr_type="uint8").ty == tvm.ir.PointerType(
	tvm.ir.PrimType("uint8"), "shared"
	)
	aptr = Ab.access_ptr("w")
	assert aptr.args[4].value == BufferAccessKind.WRITE


	def test_buffer_access_ptr_offset():
	m = tvm.tirx.Var("m", "int32")
	n = tvm.tirx.Var("n", "int32")
	Ab = tvm.tirx.decl_buffer((m, n), "float32")
	aptr = Ab.access_ptr("rw", offset=100)
	tvm.testing.assert_prim_expr_equal(aptr.args[2], 100)
	assert aptr.args[4].value == BufferAccessKind.READ \| BufferAccessKind.WRITE
	v = tvm.tirx.Var("int32", "int32")
	aptr = Ab.access_ptr("rw", offset=100 + 100 + v)
	tvm.testing.assert_prim_expr_equal(aptr.args[2], 200 + v)
	assert aptr.args[4].value == BufferAccessKind.READ \| BufferAccessKind.WRITE
	aptr = Ab.access_ptr("rw", offset=tvm.tirx.call_extern("int32", "test_call", 100 + 100 + v))
	tvm.testing.assert_prim_expr_equal(
	aptr.args[2], tvm.tirx.call_extern("int32", "test_call", 200 + v)
	)
	assert aptr.args[4].value == BufferAccessKind.READ \| BufferAccessKind.WRITE


	def test_buffer_access_ptr_extent():
	m = tvm.tirx.Var("m", "int32")
	n = tvm.tirx.Var("n", "int32")
	Ab = tvm.tirx.decl_buffer((m, n), "float32")
	aptr = Ab.access_ptr("rw")
	tvm.ir.assert_structural_equal(aptr.args[3], m * n)
	aptr = Ab.access_ptr("rw", offset=100)
	tvm.ir.assert_structural_equal(aptr.args[3], m * n - 100)
	Ab = tvm.tirx.decl_buffer((m, n), "float32", strides=[n + 1, 1])
	aptr = Ab.access_ptr("rw", offset=100)
	tvm.ir.assert_structural_equal(aptr.args[3], Ab.ty.strides[0] * m - 100)

	# Test extent from input params
	aptr = Ab.access_ptr("rw", extent=200)
	tvm.ir.assert_structural_equal(aptr.args[3], T.int32(200))
	aptr = Ab.access_ptr("rw", offset=100, extent=100)
	tvm.ir.assert_structural_equal(aptr.args[3], T.int32(100))


	def test_buffer_vload():
	m = tvm.tirx.Var("m", "int32")
	n = tvm.tirx.Var("n", "int32")
	Ab = tvm.tirx.decl_buffer((m, n), "float32", elem_offset=100)
	load = Ab.vload([2, 3])
	tvm.ir.assert_structural_equal(load.indices, [T.int32(2), T.int32(3)])


	def test_buffer_offset_of():
	m = tvm.tirx.Var("m", "int32")
	n = tvm.tirx.Var("n", "int32")
	Ab = tvm.tirx.decl_buffer((m, n), "float32", elem_offset=100)
	offset = Ab.offset_of([2, 3])
	tvm.ir.assert_structural_equal(offset, [n * 2 + 103])


	def test_buffer_index_merge_mult_mod():
	m = tvm.tirx.Var("m", "int32")
	n = tvm.tirx.Var("n", "int32")
	s = tvm.tirx.Var("s", "int32")
	k0 = tvm.tirx.Var("k0", "int32")
	k1 = tvm.tirx.Var("k1", "int32")
	A = tvm.tirx.decl_buffer((m, n), "float32")
	A_stride = tvm.tirx.decl_buffer((m, n), "float32", strides=(s, 1))

	def assert_simplified_equal(index_simplified, index_direct):
	(
	tvm.ir.assert_structural_equal(index_simplified, index_direct),
	f"index_simplified={index_simplified}, index_direct={index_direct}",
	)

	idxd = tvm.tirx.indexdiv
	idxm = tvm.tirx.indexmod

	# Test Case1
	index_simplified = A_stride.offset_of(
	(idxd(idxm(k0, k1), s), idxm(idxm(k0, k1), s) + idxd(k0, k1) * k1)
	)
	index_direct = A_stride.offset_of((0, k0))
	assert_simplified_equal(index_simplified, index_direct)

	# Test Case2
	index_simplified = A.offset_of(
	(idxd(idxm(k0, idxd(k1, s)), n), idxm(idxm(k0, idxd(k1, s)), n) + idxm(k0, k1))
	)
	index_direct = A.offset_of((0, idxm(k0, idxd(k1, s)) + idxm(k0, k1)))
	assert_simplified_equal(index_simplified, index_direct)
	# Test Case3
	index_simplified = A.offset_of(
	(
	idxd((idxd(k0, idxd(k1, s)) * idxd(k1, s)), n) + idxd(idxm(k0, idxd(k1, s)), n),
	idxm((idxd(k0, idxd(k1, s)) * idxd(k1, s)), n) + idxm(idxm(k0, idxd(k1, s)), n),
	)
	)
	index_direct = A.offset_of((0, k0))
	assert_simplified_equal(index_simplified, index_direct)
	# Test Case4 (not able to simplify)
	index_simplified = A.offset_of(
	(idxd(idxm(k0, idxd(k1, s)), n), idxm(idxm(k0, idxd(k1, n)), n) + idxm(k0, k1))
	)
	index_direct = A.offset_of(
	(0, idxd(idxm(k0, idxd(k1, s)), n) * n + (idxm(idxm(k0, idxd(k1, n)), n) + idxm(k0, k1)))
	)
	assert_simplified_equal(index_simplified, index_direct)

	# Test Case5
	B = tvm.tirx.decl_buffer((1, 14, 14, 1024))
	i = tvm.tirx.Var("i", "int32")
	j = tvm.tirx.Var("j", "int32")
	k = tvm.tirx.Var("k", "int32")

	index_simplified1 = B.offset_of(
	(
	idxd(idxd(idxd((i * 50176 + j * 28672 + k), 1024), 14), 14),
	idxm(idxd(idxd((i * 50176 + j * 28672 + k), 1024), 14), 14),
	idxm(idxd((i * 50176 + j * 28672 + k), 1024), 14),
	idxm((i * 50176 + j * 28672 + k), 1024),
	)
	)
	index_simplified2 = B.offset_of(
	(
	idxd(idxd(i * 49 + j * 28 + idxd(k, 1024), 14), 14),
	idxm(idxd(i * 49 + j * 28 + idxd(k, 1024), 14), 14),
	idxm(i * 7 + idxd(k, 1024), 14),
	idxm(k, 1024),
	)
	)
	index_direct = B.offset_of((0, 0, 0, (i * 50176 + j * 28672 + k)))
	assert_simplified_equal(index_simplified1, index_direct)
	assert_simplified_equal(index_simplified2, index_direct)


	def test_buffer_flatten():
	"""A buffer should flatten to a 1-d shape"""
	buf = tvm.tirx.decl_buffer([16, 32])
	flat = buf.get_flattened_buffer()
	# A metadata-changing rewrite creates a fresh typed Var. The physical
	# pointer is always derived from that Var instead of being stored as a
	# second buffer identity.
	assert not buf.same_as(flat)
	assert flat.data.args[0].same_as(flat)
	assert flat.data.op.name == "tirx.buffer_data"
	tvm.ir.assert_structural_equal(flat.ty.shape, [T.int32(16 * 32)])


	def test_buffer_flatten_preserves_identity():
	"""Flattening a 1-d buffer should return the original"""
	buf = tvm.tirx.decl_buffer([16])
	flat = buf.get_flattened_buffer()
	assert buf.same_as(flat)


	if __name__ == "__main__":
	tvm.testing.main()