# 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. import pytest import tvm import tvm.testing from tvm import s_tir from tvm.ir import Range from tvm.script import tirx as T @T.prim_func(s_tir=True) def func() -> None: A = T.sblock_alloc_buffer((128, 128), "float32") B = T.sblock_alloc_buffer((128, 128), "float32") C = T.sblock_alloc_buffer((128, 128), "float32") D = T.sblock_alloc_buffer((128, 128), "float32") with T.sblock(): # Need add read/write region manually to avoid triggering block access region detector T.reads([B[0, 0], C[0:16, 0:16], A[4:12, 4:12]]) T.writes([A[0:12, 0:12]]) for i, j in T.grid(8, 8): A[i, j] = B[0, 0] + C[0, 0] for i, j in T.grid(2, 2): with T.sblock(): vi, vj = T.axis.remap("SS", [i, j]) T.reads([A[vi * 4 + 4 : vi * 4 + 8, vj * 4 + 4 : vj * 4 + 8], C[12:16, 12:16]]) T.writes([A[vi * 4 + 4 : vi * 4 + 8, vj * 4 + 4 : vj * 4 + 8]]) for i, j in T.grid(4, 4): A[vi * 4 + 4 + i, vj * 4 + 4 + j] += C[i + 12, j + 12] T.evaluate(D.data) @T.prim_func(s_tir=True) def match_buffer_func() -> None: with T.sblock("root"): A = T.sblock_alloc_buffer((128, 128), "float32") B = T.sblock_alloc_buffer((128, 128), "float32") T.reads([]) T.writes([]) # Need add read/write region manually to avoid triggering block access region detector for i, j in T.grid(8, 8): with T.sblock("block"): vi, vj = T.axis.remap("SS", [i, j]) T.reads(B[vi * 16 + 2 : vi * 16 + 12, vj * 16 + 2 : vj * 16 + 16]) T.writes(A[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16]) AA = T.match_buffer(A[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16], (16, 16)) B0 = T.match_buffer(B[vi * 16 + 2 : vi * 16 + 6, vj * 16 + 2 : vj * 16 + 6], (4, 4)) B1 = T.match_buffer( B[vi * 16 + 8 : vi * 16 + 12, vj * 16 + 8 : vj * 16 + 16], (4, 8) ) for ii, jj in T.grid(16, 16): with T.sblock("AAA"): vii, vjj = T.axis.remap("SS", [ii, jj]) T.reads([]) T.writes(AA[vii, vjj]) AAA = T.match_buffer(AA[vii, vjj], ()) AAA[()] = 1.0 T.evaluate(B0.data) T.evaluate(B1.data) @T.prim_func(s_tir=True) def opaque_block_func() -> None: with T.sblock("root"): A = T.sblock_alloc_buffer((16, 16), "float32") B = T.sblock_alloc_buffer((16, 16), "float32") T.reads([]) T.writes([]) # Need add read/write region manually to avoid triggering block access region detector for i in range(0, 16): with T.sblock(): T.reads(A[i, 0:16]) T.writes([B[i, 0:16]]) for j in range(0, 16): with T.sblock(): T.reads(A[i, j]) T.writes(B[i, j]) B[i, j] = A[i, j] + 1.0 @T.prim_func(s_tir=True) def opaque_access_func() -> None: A = T.sblock_alloc_buffer([1024]) B = T.sblock_alloc_buffer([1024]) for i in T.serial(0, 8): with T.sblock(): v = T.axis.S(8, i) T.reads([A[v * 128 : v * 128 + 128]]) T.writes([B[v * 128 : v * 128 + 128]]) T.evaluate( T.call_extern("test", B.data, v * 128, 128, A.data, v * 128, 128, dtype="float32") ) @T.prim_func(s_tir=True) def opaque_access_with_tvm_access_ptr_func() -> None: A = T.sblock_alloc_buffer([1024]) B = T.sblock_alloc_buffer([1024]) C = T.sblock_alloc_buffer([1024]) with T.sblock("opaque"): T.reads(A[0:1024], C[0:1024]) T.writes(B[0:1024], C[0:1024]) T.evaluate(A.access_ptr("r")) T.evaluate(B.access_ptr("w")) T.evaluate(C.access_ptr("rw")) @T.prim_func(s_tir=True) def access_in_if_then_else_func() -> None: A = T.sblock_alloc_buffer([8]) B = T.sblock_alloc_buffer([8]) with T.sblock(): T.reads([A[0:5]]) T.writes([B[0:8]]) for i in T.serial(0, 8): B[i] = T.if_then_else(i < 5, A[i], 0.0, dtype="float32") @T.prim_func(s_tir=True) def access_in_branch_func() -> None: A = T.sblock_alloc_buffer([8]) B = T.sblock_alloc_buffer([8]) with T.sblock(): T.reads([A[0:7]]) T.writes([B[0:8]]) for i in T.serial(0, 8): if i < 5: B[i] = A[i] + 1.0 else: B[i] = A[i - 1] @T.prim_func(s_tir=True) def gemm() -> None: A = T.sblock_alloc_buffer([16, 16], "float32") B = T.sblock_alloc_buffer([16, 16], "float32") C = T.sblock_alloc_buffer([16, 16], "float32") for i, j, k, ii, jj in T.grid(4, 4, 16, 4, 4): with T.sblock("update"): vi = T.axis.S(16, i * 4 + ii) vj = T.axis.S(16, j * 4 + jj) vk = T.axis.R(16, k) T.reads(A[vi, vk], B[vj, vk]) T.writes(C[vi, vj]) with T.init(): C[vi, vj] = 0 C[vi, vj] += A[vi, vk] * B[vj, vk] @T.prim_func(s_tir=True) def decomposed_gemm() -> None: A = T.sblock_alloc_buffer([16, 16], "float32") B = T.sblock_alloc_buffer([16, 16], "float32") C = T.sblock_alloc_buffer([16, 16], "float32") for i, j in T.grid(4, 4): for ii, jj in T.grid(4, 4): with T.sblock("init"): vi = T.axis.S(16, i * 4 + ii) vj = T.axis.S(16, j * 4 + jj) T.reads([]) T.writes(C[vi, vj]) C[vi, vj] = 0 for k, ii, jj in T.grid(16, 4, 4): with T.sblock("update"): vi = T.axis.S(16, i * 4 + ii) vj = T.axis.S(16, j * 4 + jj) vk = T.axis.R(16, k) T.reads(C[vi, vj], A[vi, vk], B[vj, vk]) T.writes(C[vi, vj]) C[vi, vj] += A[vi, vk] * B[vj, vk] @T.prim_func(s_tir=True) def access_of_padding_pattern() -> None: X = T.sblock_alloc_buffer([28, 28]) X_pad = T.sblock_alloc_buffer([32, 32]) Y = T.sblock_alloc_buffer([28, 28]) for i, j in T.grid(32, 32): with T.sblock("padding"): vi, vj = T.axis.remap("SS", [i, j]) T.reads([X[vi - 2, vj - 2]]) T.writes([X_pad[vi, vj]]) X_pad[vi, vj] = T.if_then_else( 2 <= vi and vi < 30 and 2 <= vj and vj < 30, X[vi - 2, vj - 2], 0.0, dtype="float32" ) with T.sblock("padding_reverse"): vi, vj = T.axis.remap("SS", [i, j]) T.reads([X_pad[vi, vj]]) T.writes([Y[vi - 2, vj - 2]]) if 2 <= vi and vi < 30 and 2 <= vj and vj < 30: Y[vi - 2, vj - 2] = X_pad[vi, vj] def test_block_access_region_detector(): block = func.body.block.body.block alloc_buffers = func.body.block.alloc_buffers buffer_var_map = {buf.data: buf for buf in alloc_buffers} ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(block.reads, ret[0]) tvm.ir.assert_structural_equal(block.writes, ret[1]) D = alloc_buffers[-1] tvm.ir.assert_structural_equal( [tvm.tirx.BufferRegion(D, [Range(0, 128), Range(0, 128)])], ret[2] ) def test_opaque_block(): alloc_buffers = opaque_block_func.body.block.alloc_buffers buffer_var_map = {buf.data: buf for buf in alloc_buffers} block0 = opaque_block_func.body.block.body.body.block ret = s_tir.analysis.get_sblock_access_region(block0, buffer_var_map) tvm.ir.assert_structural_equal(block0.reads, ret[0]) tvm.ir.assert_structural_equal(block0.writes, ret[1]) block1 = block0.body.body.block ret = s_tir.analysis.get_sblock_access_region(block1, buffer_var_map) tvm.ir.assert_structural_equal(block1.reads, ret[0]) tvm.ir.assert_structural_equal(block1.writes, ret[1]) def test_opaque_access(): block = opaque_access_func.body.block.body.body.block alloc_buffers = opaque_access_func.body.block.alloc_buffers buffer_var_map = {buf.data: buf for buf in alloc_buffers} ret0 = s_tir.analysis.get_sblock_read_write_region(block, buffer_var_map) ret1 = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) with pytest.raises(ValueError): tvm.ir.assert_structural_equal(ret0[0], ret1[0]) with pytest.raises(ValueError): tvm.ir.assert_structural_equal(ret0[1], ret1[1]) def test_opaque_access_with_tvm_access_ptr(): block = opaque_access_with_tvm_access_ptr_func.body.block.body.block alloc_buffers = opaque_access_with_tvm_access_ptr_func.body.block.alloc_buffers buffer_var_map = {buf.data: buf for buf in alloc_buffers} ret0 = s_tir.analysis.get_sblock_read_write_region(block, buffer_var_map) ret1 = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(block.reads, ret0[0]) tvm.ir.assert_structural_equal(block.writes, ret0[1]) with pytest.raises(ValueError): tvm.ir.assert_structural_equal(ret0[0], ret1[0]) with pytest.raises(ValueError): tvm.ir.assert_structural_equal(ret0[1], ret1[1]) def test_match_buffer(): root_block = match_buffer_func.body.block block = root_block.body.body.body.block block_inner = block.body[0].body.body.block alloc_buffers = match_buffer_func.body.block.alloc_buffers buffer_var_map = {buf.data: buf for buf in alloc_buffers} # Check block ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(block.writes, ret[1]) # B is opaque access tvm.ir.assert_structural_equal(block.reads, ret[2]) # Check inner block AAA without updating buffer_var_map ret = s_tir.analysis.get_sblock_access_region(block_inner, buffer_var_map) # Since AA is not in the buffer_var_map, region of AA will not be collected. tvm.ir.assert_structural_equal([], ret[1]) # Check inner block AAA for match_buffer in block.match_buffers: target_buffer = match_buffer.buffer buffer_var_map[target_buffer.data] = target_buffer ret = s_tir.analysis.get_sblock_access_region(block_inner, buffer_var_map) tvm.ir.assert_structural_equal(block_inner.reads, ret[0]) tvm.ir.assert_structural_equal(block_inner.writes, ret[1]) def test_access_in_if_then_else_func(): block = access_in_if_then_else_func.body.block.body.block alloc_buffers = access_in_if_then_else_func.body.block.alloc_buffers buffer_var_map = {buf.data: buf for buf in alloc_buffers} ret0 = s_tir.analysis.get_sblock_read_write_region(block, buffer_var_map) ret1 = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(ret0[0], ret1[0]) tvm.ir.assert_structural_equal(ret0[1], ret1[1]) def test_access_in_branch_func(): block = access_in_branch_func.body.block.body.block alloc_buffers = access_in_branch_func.body.block.alloc_buffers buffer_var_map = {buf.data: buf for buf in alloc_buffers} ret0 = s_tir.analysis.get_sblock_read_write_region(block, buffer_var_map) ret1 = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(ret0[0], ret1[0]) tvm.ir.assert_structural_equal(ret0[1], ret1[1]) def test_access_of_padding_pattern(): s = tvm.s_tir.schedule.Schedule(access_of_padding_pattern) alloc_buffers = s.get_sref(s.get_sblock("root")).stmt.alloc_buffers buffer_var_map = {buf.data: buf for buf in alloc_buffers} def do_compare_buffer_region(region, expect): assert region.buffer == expect.buffer analyzer = tvm.arith.Analyzer() for observed_range, expected_range in zip(region.region, expect.region): analyzer.can_prove_equal(observed_range.min, expected_range.min) analyzer.can_prove_equal(observed_range.extent, expected_range.extent) def do_check_block(block_name): block = s.get_sref(s.get_sblock(block_name)).stmt expect_reads = block.reads expect_writes = block.writes ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) for i, read in enumerate(ret[0]): do_compare_buffer_region(read, expect_reads[i]) for i, write in enumerate(ret[1]): do_compare_buffer_region(write, expect_writes[i]) do_check_block("padding") do_check_block("padding_reverse") def test_access_of_reduction(): block = gemm.body.block.body.body.body.body.body.body.block alloc_buffers = gemm.body.block.alloc_buffers buffer_var_map = {buf.data: buf for buf in alloc_buffers} ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(block.reads, ret[0]) tvm.ir.assert_structural_equal(block.writes, ret[1]) def test_access_of_decompose_reduction(): init = decomposed_gemm.body.block.body.body.body[0].body.body.block update = decomposed_gemm.body.block.body.body.body[1].body.body.body.block alloc_buffers = decomposed_gemm.body.block.alloc_buffers buffer_var_map = {buf.data: buf for buf in alloc_buffers} for block in [init, update]: ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(block.reads, ret[0]) tvm.ir.assert_structural_equal(block.writes, ret[1]) def test_buffer_access_with_let_binding(): @T.prim_func(s_tir=True) def func( storage: T.Buffer((16, 16, 16), "float32"), seq_slot_ids: T.Buffer((16,), "int32"), history_slot_ids: T.Buffer((16,), "int32"), output: T.Buffer((16, 16), "float32"), ): for i, s in T.grid(16, 16): with T.sblock("copy"): vi, vs = T.axis.remap("SS", [i, s]) T.reads( seq_slot_ids[vi], history_slot_ids[vi], storage[seq_slot_ids[vi], history_slot_ids[vi], vs], ) T.writes(output[vi, vs]) seq_id: T.let[T.int32] = seq_slot_ids[vi] history_id: T.let[T.int32] = history_slot_ids[vi] output[vi, vs] = storage[seq_id, history_id, vs] block = func.body.block.body.body.body.block buffer_var_map = {buf.data: buf for buf in func.buffer_map.values()} ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(block.reads, ret[0]) tvm.ir.assert_structural_equal(block.writes, ret[1]) def test_buffer_access_with_nested_let_binding(): @T.prim_func(s_tir=True) def func( A: T.Buffer((16, 16), "float32"), B: T.Buffer((16, 16), "float32"), C: T.Buffer((16, 16), "float32"), ): for i, s in T.grid(16, 16): with T.sblock("copy"): vi, vs = T.axis.remap("SS", [i, s]) T.reads(A[vi, vs], B[vi, vs]) T.writes(C[vi, vs]) vi1: T.let[T.int32] = vi vi2: T.let[T.int32] = vi1 vs1: T.let[T.int32] = vs vs2: T.let[T.int32] = vs1 vs3: T.let[T.int32] = vs2 C[vi, vs1] = A[vi1, vs2] + B[vi2, vs3] block = func.body.block.body.body.body.block buffer_var_map = {buf.data: buf for buf in func.buffer_map.values()} ret = s_tir.analysis.get_sblock_access_region(block, buffer_var_map) tvm.ir.assert_structural_equal(block.reads, ret[0]) tvm.ir.assert_structural_equal(block.writes, ret[1]) if __name__ == "__main__": tvm.testing.main()