chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,415 @@
|
||||
# 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()
|
||||
Reference in New Issue
Block a user