chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,347 @@
|
||||
# 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 tvm.testing
|
||||
from tvm.ir import Range
|
||||
from tvm.script import tirx as T
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
B = T.match_buffer(b, [128, 128])
|
||||
C = T.match_buffer(c, [128, 128])
|
||||
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
with T.sblock("update"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = T.float32(0)
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def matmul_original(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
B = T.match_buffer(b, [128, 128])
|
||||
C = T.match_buffer(c, [128, 128])
|
||||
|
||||
for i, j in T.grid(32, 32):
|
||||
with T.sblock("init"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
for ii, jj in T.grid(4, 4):
|
||||
C[vi * 4 + ii, vj * 4 + jj] = T.float32(0)
|
||||
|
||||
for k in range(0, 32):
|
||||
with T.sblock("update"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
for ii, jj, kk in T.grid(4, 4, 4):
|
||||
C[vi * 4 + ii, vj * 4 + jj] = (
|
||||
C[vi * 4 + ii, vj * 4 + jj]
|
||||
+ A[vi * 4 + ii, vk * 4 + kk] * B[vj * 4 + jj, vk * 4 + kk]
|
||||
)
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def elementwise_with_root(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
B = T.match_buffer(b, [128, 128])
|
||||
C = T.match_buffer(c, [128, 128])
|
||||
|
||||
with T.sblock():
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
B[vi, vj] = A[vi, vj] + T.float32(1)
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
C[vi, vj] = B[vi, vj] + T.float32(1)
|
||||
|
||||
|
||||
def func_with_opaque_block(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
B = T.match_buffer(b, [128, 128])
|
||||
C = T.match_buffer(c, [128, 128])
|
||||
|
||||
with T.sblock():
|
||||
with T.sblock():
|
||||
B[0, 0] = A[0, 0] + T.float32(1)
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
C[vi, vj] = B[vi, vj] + T.float32(1)
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func_with_part_access_region(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
B = T.match_buffer(b, [128, 128])
|
||||
C = T.match_buffer(c, [128, 128])
|
||||
|
||||
with T.sblock():
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
T.reads(A[vi, vj])
|
||||
B[vi, vj] = A[vi, vj] + T.float32(1)
|
||||
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
T.writes(C[vi, vj])
|
||||
C[vi, vj] = B[vi, vj] + T.float32(1)
|
||||
|
||||
|
||||
def test_complete_matmul():
|
||||
func = matmul
|
||||
A, B, C = [func.buffer_map[x] for x in func.params]
|
||||
|
||||
block = func.body.block.body.body.body.body.block
|
||||
assert isinstance(block, tvm.tirx.SBlock)
|
||||
vi, vj, vk = [x.var for x in block.iter_vars]
|
||||
access_A = tvm.tirx.BufferRegion(
|
||||
A, [Range.from_min_extent(vi, 1), Range.from_min_extent(vk, 1)]
|
||||
)
|
||||
access_B = tvm.tirx.BufferRegion(
|
||||
B, [Range.from_min_extent(vj, 1), Range.from_min_extent(vk, 1)]
|
||||
)
|
||||
access_C = tvm.tirx.BufferRegion(
|
||||
C, [Range.from_min_extent(vi, 1), Range.from_min_extent(vj, 1)]
|
||||
)
|
||||
tvm.ir.assert_structural_equal(block.reads, [access_A, access_B])
|
||||
tvm.ir.assert_structural_equal(block.writes, [access_C])
|
||||
|
||||
|
||||
def test_complete_matmul_original():
|
||||
func = matmul_original
|
||||
A, B, C = [func.buffer_map[x] for x in func.params]
|
||||
|
||||
block1 = func.body.block.body.body.body[0].block
|
||||
assert isinstance(block1, tvm.tirx.SBlock)
|
||||
vi, vj = [x.var for x in block1.iter_vars]
|
||||
access_C = tvm.tirx.BufferRegion(
|
||||
C, [Range.from_min_extent(vi * 4, 4), Range.from_min_extent(vj * 4, 4)]
|
||||
)
|
||||
tvm.ir.assert_structural_equal(block1.reads, [])
|
||||
tvm.ir.assert_structural_equal(block1.writes, [access_C])
|
||||
|
||||
block2 = func.body.block.body.body.body[1].body.block
|
||||
assert isinstance(block2, tvm.tirx.SBlock)
|
||||
vi, vj, vk = [x.var for x in block2.iter_vars]
|
||||
access_A = tvm.tirx.BufferRegion(
|
||||
A, [Range.from_min_extent(vi * 4, 4), Range.from_min_extent(vk * 4, 4)]
|
||||
)
|
||||
access_B = tvm.tirx.BufferRegion(
|
||||
B, [Range.from_min_extent(vj * 4, 4), Range.from_min_extent(vk * 4, 4)]
|
||||
)
|
||||
access_C = tvm.tirx.BufferRegion(
|
||||
C, [Range.from_min_extent(vi * 4, 4), Range.from_min_extent(vj * 4, 4)]
|
||||
)
|
||||
tvm.ir.assert_structural_equal(block2.reads, [access_C, access_A, access_B])
|
||||
tvm.ir.assert_structural_equal(block2.writes, [access_C])
|
||||
|
||||
|
||||
def _check_elementwise(func):
|
||||
A, B, C = [func.buffer_map[x] for x in func.params]
|
||||
|
||||
root_block = func.body.block
|
||||
assert len(root_block.reads) == 0
|
||||
assert len(root_block.writes) == 0
|
||||
|
||||
block1 = func.body.block.body[0].body.body.block
|
||||
assert isinstance(block1, tvm.tirx.SBlock)
|
||||
vi, vj = [x.var for x in block1.iter_vars]
|
||||
|
||||
tvm.ir.assert_structural_equal(
|
||||
block1.reads,
|
||||
[tvm.tirx.BufferRegion(A, [Range.from_min_extent(vi, 1), Range.from_min_extent(vj, 1)])],
|
||||
)
|
||||
tvm.ir.assert_structural_equal(
|
||||
block1.writes,
|
||||
[tvm.tirx.BufferRegion(B, [Range.from_min_extent(vi, 1), Range.from_min_extent(vj, 1)])],
|
||||
)
|
||||
|
||||
block2 = func.body.block.body[1].body.body.block
|
||||
assert isinstance(block2, tvm.tirx.SBlock)
|
||||
vi, vj = [x.var for x in block2.iter_vars]
|
||||
tvm.ir.assert_structural_equal(
|
||||
block2.reads,
|
||||
[tvm.tirx.BufferRegion(B, [Range.from_min_extent(vi, 1), Range.from_min_extent(vj, 1)])],
|
||||
)
|
||||
tvm.ir.assert_structural_equal(
|
||||
block2.writes,
|
||||
[tvm.tirx.BufferRegion(C, [Range.from_min_extent(vi, 1), Range.from_min_extent(vj, 1)])],
|
||||
)
|
||||
|
||||
|
||||
def test_complete_with_root():
|
||||
_check_elementwise(elementwise_with_root)
|
||||
|
||||
|
||||
def test_complete_part_region():
|
||||
_check_elementwise(func_with_part_access_region)
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func_with_bufferslice_indices(data: T.handle, index: T.handle) -> None:
|
||||
data_buf = T.match_buffer(data, (16, 16), "float32")
|
||||
index_buf = T.match_buffer(index, (1,), "int32")
|
||||
out_buf = T.sblock_alloc_buffer((16, 16), "float32")
|
||||
|
||||
for i, j in T.grid(16, 16):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
out_buf[vi, vj] = data_buf[vi, index_buf[0]]
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def expected_bufferslice_indices(data: T.handle, index: T.handle) -> None:
|
||||
index_buf = T.match_buffer(index, [1], dtype="int32", elem_offset=0, align=64, offset_factor=1)
|
||||
data_buf = T.match_buffer(data, [16, 16], elem_offset=0, align=64, offset_factor=1)
|
||||
with T.sblock("root"):
|
||||
T.reads([])
|
||||
T.writes([])
|
||||
out_buf = T.sblock_alloc_buffer([16, 16], elem_offset=0, align=64, offset_factor=1)
|
||||
for i0, i1 in T.grid(16, 16):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i0, i1])
|
||||
T.reads([data_buf[vi, index_buf[0]], index_buf[0]])
|
||||
T.writes([out_buf[vi, vj]])
|
||||
out_buf[vi, vj] = data_buf[vi, index_buf[0]]
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func_with_recursive_bufferslice_indices(data: T.handle, index: T.handle) -> None:
|
||||
data_buf = T.match_buffer(data, (16, 16), "float32")
|
||||
index_buf = T.match_buffer(index, (1,), "int32")
|
||||
out_buf = T.sblock_alloc_buffer((16, 16), "float32")
|
||||
|
||||
for i, j in T.grid(16, 16):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
out_buf[vi, vj] = data_buf[index_buf[index_buf[0]], index_buf[0]]
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def expected_recursive_bufferslice_indices(data: T.handle, index: T.handle) -> None:
|
||||
index_buf = T.match_buffer(index, [1], dtype="int32", elem_offset=0, align=64, offset_factor=1)
|
||||
data_buf = T.match_buffer(data, [16, 16], elem_offset=0, align=64, offset_factor=1)
|
||||
with T.sblock("root"):
|
||||
T.reads([])
|
||||
T.writes([])
|
||||
out_buf = T.sblock_alloc_buffer([16, 16], elem_offset=0, align=64, offset_factor=1)
|
||||
for i0, i1 in T.grid(16, 16):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(
|
||||
[
|
||||
data_buf[index_buf[index_buf[0]], index_buf[0]],
|
||||
index_buf[T.min(index_buf[0], 0) : T.max(index_buf[0], 0) + 1],
|
||||
]
|
||||
)
|
||||
T.writes([out_buf[vi, vj]])
|
||||
out_buf[vi, vj] = data_buf[index_buf[index_buf[0]], index_buf[0]]
|
||||
|
||||
|
||||
def test_complete_buffer_indices():
|
||||
new_func = tvm.script.from_source(func_with_bufferslice_indices.script()).with_attr(
|
||||
"global_symbol", "main"
|
||||
)
|
||||
tvm.ir.assert_structural_equal(
|
||||
new_func, expected_bufferslice_indices.with_attr("global_symbol", "main")
|
||||
)
|
||||
new_func = tvm.script.from_source(func_with_recursive_bufferslice_indices.script()).with_attr(
|
||||
"global_symbol", "main"
|
||||
)
|
||||
tvm.ir.assert_structural_equal(
|
||||
new_func, expected_recursive_bufferslice_indices.with_attr("global_symbol", "main")
|
||||
)
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def match_buffer_func(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, (16, 16))
|
||||
for i in range(0, 16):
|
||||
with T.sblock():
|
||||
A0 = T.match_buffer(A[i, 0:16], (16))
|
||||
with T.sblock():
|
||||
for j in range(0, 16):
|
||||
with T.sblock():
|
||||
A1 = T.match_buffer(A0[j], ())
|
||||
A1[()] = 1.0
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def expected_match_buffer_func(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, (16, 16))
|
||||
for i in range(0, 16):
|
||||
with T.sblock():
|
||||
T.reads([])
|
||||
T.writes(A[i, 0:16])
|
||||
A0 = T.match_buffer(A[i, 0:16], (16))
|
||||
with T.sblock():
|
||||
T.reads([])
|
||||
T.writes(A0[0:16])
|
||||
for j in range(0, 16):
|
||||
with T.sblock():
|
||||
T.reads([])
|
||||
T.writes(A0[j])
|
||||
A1 = T.match_buffer(A0[j], ())
|
||||
A1[()] = 1.0
|
||||
|
||||
|
||||
def test_complete_match_buffer():
|
||||
tvm.ir.assert_structural_equal(
|
||||
match_buffer_func.with_attr("global_symbol", "main"),
|
||||
expected_match_buffer_func.with_attr("global_symbol", "main"),
|
||||
)
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def alloc_buffer_func(a: T.handle, b: T.handle) -> None:
|
||||
A = T.match_buffer(a, [2, 2], dtype="float32")
|
||||
B = T.match_buffer(b, [2, 2], dtype="float32")
|
||||
C = T.sblock_alloc_buffer([2, 2], dtype="float32")
|
||||
A[(0, 0)] = T.float32(2)
|
||||
C[(0, 0)] = A[(0, 0)] + B[(0, 0)]
|
||||
B[(0, 0)] = C[(0, 0)]
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def expect_alloc_buffer_func(a: T.handle, b: T.handle) -> None:
|
||||
A = T.match_buffer(a, [2, 2], dtype="float32", elem_offset=0, align=64, offset_factor=1)
|
||||
B = T.match_buffer(b, [2, 2], dtype="float32", elem_offset=0, align=64, offset_factor=1)
|
||||
with T.sblock("root"):
|
||||
T.reads([])
|
||||
T.writes([])
|
||||
C = T.sblock_alloc_buffer([2, 2], dtype="float32", elem_offset=0, align=64, offset_factor=1)
|
||||
A[(0, 0)] = T.float32(2)
|
||||
C[(0, 0)] = A[(0, 0)] + B[(0, 0)]
|
||||
B[(0, 0)] = C[(0, 0)]
|
||||
|
||||
|
||||
def test_complete_alloc_buffer():
|
||||
rt_func = tvm.script.from_source(alloc_buffer_func.script()).with_attr("global_symbol", "main")
|
||||
tvm.ir.assert_structural_equal(
|
||||
rt_func, expect_alloc_buffer_func.with_attr("global_symbol", "main")
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,666 @@
|
||||
# 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, F821, F841, RUF005
|
||||
import inspect
|
||||
import re
|
||||
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm import tirx
|
||||
from tvm.script import from_source
|
||||
from tvm.script import tirx as T
|
||||
|
||||
|
||||
def check_error(func, rel_lineno):
|
||||
"""check if TIR script throws error"""
|
||||
check_error_re = re.compile(r"^.*# check_error: (.+)$")
|
||||
source_code = inspect.getsource(func)
|
||||
indent = len(re.match(r"^\s*", source_code).group(0))
|
||||
source_code = "@T.prim_func(s_tir=True)\n" + "\n".join(
|
||||
line[indent:] for line in source_code.splitlines()
|
||||
)
|
||||
# Parse errors now raise DiagnosticError with formatted source location.
|
||||
with pytest.raises(tvm.error.DiagnosticError) as execinfo:
|
||||
from_source(source_code)
|
||||
err_str = str(execinfo.value)
|
||||
if rel_lineno is None:
|
||||
return
|
||||
# The error message contains " --> <source>:<lineno>:<col>" formatted by Diagnostics.error().
|
||||
# Accept either rel_lineno or rel_lineno+1 to match old tolerance.
|
||||
assert f":{rel_lineno}:" in err_str or f":{rel_lineno + 1}:" in err_str, (
|
||||
f"Expected error message to contain line {rel_lineno}, got:\n{err_str}"
|
||||
)
|
||||
error_line = source_code.split("\n")[rel_lineno]
|
||||
m = check_error_re.match(error_line)
|
||||
if m:
|
||||
expected_error_text = m.group(1)
|
||||
assert expected_error_text in err_str, (
|
||||
f'check_error expects "{expected_error_text}" in error: {err_str}'
|
||||
)
|
||||
|
||||
|
||||
def test_buffer_bind():
|
||||
def buffer_bind_missing_args(a: T.handle) -> None:
|
||||
A = T.match_buffer((16, 16), "float32") # error
|
||||
|
||||
check_error(buffer_bind_missing_args, 2)
|
||||
|
||||
|
||||
def test_undefined_buffer():
|
||||
def undefined_buffer(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, (16, 16), "float32")
|
||||
|
||||
for i in T.serial(16):
|
||||
for j in T.serial(0, 16):
|
||||
C[i, j] = 0.0 # error
|
||||
|
||||
check_error(undefined_buffer, 6)
|
||||
|
||||
|
||||
def test_unsupported_function_call():
|
||||
def unsupported_function_call(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, (16, 16), "float32")
|
||||
|
||||
for i in T.const_range(16): # error
|
||||
for j in T.serial(0, 16):
|
||||
A[i, j] = 0.0
|
||||
|
||||
check_error(unsupported_function_call, 4)
|
||||
|
||||
|
||||
def test_missing_type_annotation():
|
||||
def missing_type_annotation(a) -> None: # error
|
||||
T.evaluate(0.0)
|
||||
|
||||
check_error(missing_type_annotation, 1)
|
||||
|
||||
|
||||
def test_invalid_for_function():
|
||||
def invalid_for_function(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, (16, 16), "float32")
|
||||
for i in T.evaluate(0.0): # error
|
||||
for j in T.serial(0, 16):
|
||||
A[i, j] = 0.0
|
||||
|
||||
check_error(invalid_for_function, 4)
|
||||
|
||||
|
||||
def test_invalid_block_function():
|
||||
def invalid_block_function(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, (16, 16), "float32")
|
||||
|
||||
with T.evaluate(0.0): # error
|
||||
T.evaluate(1.0)
|
||||
|
||||
check_error(invalid_block_function, 4)
|
||||
|
||||
|
||||
def test_return_not_allowed():
|
||||
def return_not_allowed(a: T.handle) -> None:
|
||||
return T.evaluate(0) # error
|
||||
|
||||
check_error(return_not_allowed, 2)
|
||||
|
||||
|
||||
def test_no_body():
|
||||
def no_body(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, (16, 16), "float32")
|
||||
T.realize(A, "") # error
|
||||
|
||||
check_error(no_body, 3)
|
||||
|
||||
|
||||
def test_inconsistent_binding():
|
||||
def inconsistent_binding_value() -> None:
|
||||
for i, j in T.grid(16, 16):
|
||||
vi, vj = T.axis.remap("SS", [i]) # error
|
||||
T.evaluate(1.0)
|
||||
|
||||
def inconsistent_binding_type() -> None:
|
||||
for i, j in T.grid(16, 16):
|
||||
vi, vj = T.axis.remap("S", [i, j]) # error
|
||||
T.evaluate(1.0)
|
||||
|
||||
check_error(inconsistent_binding_value, 3)
|
||||
check_error(inconsistent_binding_type, 3)
|
||||
|
||||
|
||||
def test_error_remap_args():
|
||||
def error_remap_type() -> None:
|
||||
for i, j in T.grid(16, 16):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("TT", [i, j]) # error
|
||||
T.evaluate(1.0)
|
||||
|
||||
def error_remap_value() -> None:
|
||||
for i, j in T.grid(16, 16):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i + j, j]) # error
|
||||
T.evaluate(1.0)
|
||||
|
||||
check_error(error_remap_type, 4)
|
||||
check_error(error_remap_value, 4)
|
||||
|
||||
|
||||
def test_invalid_block_axes():
|
||||
def invalid_block_axes(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, (16, 16), "float32")
|
||||
for i, j in T.grid(16, 16):
|
||||
with T.sblock():
|
||||
vi = T.axis.S(i, A) # error
|
||||
T.evaluate(1.0)
|
||||
|
||||
check_error(invalid_block_axes, 5)
|
||||
|
||||
|
||||
def test_duplicate_block_axes():
|
||||
def duplicate_block_axes() -> None:
|
||||
for i, j in T.grid(16, 16):
|
||||
with T.sblock():
|
||||
vi = T.axis.S(16, i)
|
||||
vi = T.axis.S(16, j) # error
|
||||
T.evaluate(1.0)
|
||||
|
||||
def duplicate_block_axes_remap() -> None:
|
||||
for i, j in T.grid(16, 16):
|
||||
with T.sblock():
|
||||
vi, vi = T.axis.remap("SS", [i, j]) # error
|
||||
T.evaluate(1.0)
|
||||
|
||||
check_error(duplicate_block_axes, 5)
|
||||
check_error(duplicate_block_axes_remap, 4)
|
||||
|
||||
|
||||
def test_miss_block_bind():
|
||||
def miss_block_bind_value() -> None:
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock():
|
||||
vi = T.axis.S(i) # error
|
||||
T.evaluate(1.0)
|
||||
|
||||
check_error(miss_block_bind_value, 4)
|
||||
|
||||
|
||||
def test_invalid_loop_var():
|
||||
def invalid_loop_var() -> None:
|
||||
for i, j in range(0, 16): # error
|
||||
T.evaluate(1.0)
|
||||
|
||||
check_error(invalid_loop_var, 2)
|
||||
|
||||
|
||||
def test_inconsistent_grid():
|
||||
def inconsistent_grid(A: T.Buffer(16)) -> None:
|
||||
for i in T.grid(16, 16): # valid, i is a tuple (iter0, iter1)
|
||||
T.evaluate(A[i]) # error
|
||||
|
||||
check_error(inconsistent_grid, 3)
|
||||
|
||||
|
||||
def test_invalid_match_buffer_region():
|
||||
def invalid_match_buffer_region() -> None:
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
A = T.match_buffer(vi) # error
|
||||
T.evaluate(1.0)
|
||||
|
||||
check_error(invalid_match_buffer_region, 5)
|
||||
|
||||
|
||||
def test_duplicate_buffer():
|
||||
def duplicate_buffer() -> None:
|
||||
A = T.sblock_alloc_buffer((128, 128), "float32")
|
||||
A = T.sblock_alloc_buffer((128, 128), "float32") # error
|
||||
|
||||
check_error(duplicate_buffer, 3)
|
||||
|
||||
|
||||
def test_duplicate_block_signature():
|
||||
def duplicate_reads() -> None:
|
||||
A = T.sblock_alloc_buffer((128, 128), "float32")
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
T.reads(A[0:8, 0:8])
|
||||
T.reads(A[0:16, 0:16]) # error
|
||||
T.evaluate(1.0)
|
||||
|
||||
def duplicate_writes() -> None:
|
||||
A = T.sblock_alloc_buffer((128, 128), "float32")
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
T.writes(A[0:8, 0:8])
|
||||
T.writes(A[0:16, 0:16]) # error
|
||||
T.evaluate(1.0)
|
||||
|
||||
def duplicate_predicate() -> None:
|
||||
for i, j in T.grid(16, 16):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
T.where(1)
|
||||
T.where(0) # error
|
||||
|
||||
def duplicate_init() -> None:
|
||||
for i, j in T.grid(16, 16):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
with T.init():
|
||||
T.evaluate(1.0)
|
||||
with T.init(): # error
|
||||
T.evaluate(1.0)
|
||||
|
||||
def duplicate_axes() -> None:
|
||||
for i, j in T.grid(16, 16):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
vi = T.axis.S(i, 16) # error
|
||||
T.evaluate(1.0)
|
||||
|
||||
def duplicate_sblock_attrs_with_same_key_diff_value() -> None:
|
||||
for i, j in T.grid(16, 16):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
T.sblock_attr({"key1": "block1"})
|
||||
T.sblock_attr({"key1": "block2"}) # error
|
||||
T.evaluate(1.0)
|
||||
|
||||
check_error(duplicate_reads, 7)
|
||||
check_error(duplicate_writes, 7)
|
||||
check_error(duplicate_predicate, 6)
|
||||
check_error(duplicate_init, 7)
|
||||
check_error(duplicate_axes, 5)
|
||||
check_error(duplicate_sblock_attrs_with_same_key_diff_value, 6)
|
||||
|
||||
|
||||
def test_opaque_access_during_complete():
|
||||
def opaque_access_during_complete(a: T.handle) -> None: # error
|
||||
A = T.match_buffer(a, (16, 16), "float32")
|
||||
for i, j in T.grid(16, 16):
|
||||
with T.sblock():
|
||||
T.evaluate(T.call_extern("dummy_extern_function", A.data, dtype="int32"))
|
||||
|
||||
check_error(opaque_access_during_complete, None)
|
||||
|
||||
|
||||
def test_convert_slice_to_bufferload():
|
||||
def convert_slice_to_bufferload() -> None:
|
||||
A = T.sblock_alloc_buffer((128, 128), "float32")
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
A[vi, vj] = A[vi : vi + 2, vj] + 1 # error
|
||||
|
||||
check_error(convert_slice_to_bufferload, 6)
|
||||
|
||||
|
||||
def test_tvm_exception_catch_from_special_stmt():
|
||||
def special_stmt_except() -> None:
|
||||
A = T.sblock_alloc_buffer("(128, 128)", "float32") # error
|
||||
T.evaluate(1.0)
|
||||
|
||||
check_error(special_stmt_except, 2)
|
||||
|
||||
|
||||
def test_tvm_exception_catch_from_scope_handler():
|
||||
def scope_handler_except() -> None:
|
||||
for i in T.serial("1", "1"): # error
|
||||
T.evaluate(1)
|
||||
|
||||
check_error(scope_handler_except, 2)
|
||||
|
||||
|
||||
def test_tvm_exception_catch_from_bare_intrin():
|
||||
def intrin_except_unassign(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, (16, 16), "float32")
|
||||
T.evaluate(A) # error
|
||||
|
||||
check_error(intrin_except_unassign, 3)
|
||||
|
||||
|
||||
def test_tvm_exception_catch_from_assigned_intrin():
|
||||
def intrin_except_assign(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, (16, 16), "float32")
|
||||
A[0, 0] = A[A] # error
|
||||
|
||||
check_error(intrin_except_assign, 3)
|
||||
|
||||
|
||||
def test_match_buffer_shape_mismatch():
|
||||
def buffer_shape_mismatch(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, (8, 8))
|
||||
for i, j in T.grid(8, 2):
|
||||
with T.sblock():
|
||||
T.reads([])
|
||||
T.writes([A[i, j * 4 : j * 4 + 4]])
|
||||
sub_A = T.match_buffer(
|
||||
A[i, j * 4 : j * 4 + 4], (5)
|
||||
) # error: shape mismatched between 4 and 5
|
||||
for jj in range(0, 4):
|
||||
sub_A[i, j * 4 + jj] = 1
|
||||
|
||||
check_error(buffer_shape_mismatch, 7)
|
||||
|
||||
|
||||
def test_high_dim_store():
|
||||
def high_dim_store() -> None:
|
||||
with T.sblock("root"):
|
||||
B = T.alloc_buffer((256,), "float32")
|
||||
for i, j in T.grid(16, 16):
|
||||
B[i, j] = 1.0 # error: Store is only allowed with one index
|
||||
|
||||
check_error(high_dim_store, 5)
|
||||
|
||||
|
||||
def test_block_has_option_vars():
|
||||
def block_has_option_vars() -> None:
|
||||
with T.sblock("root") as x: # error: block does not support option_vars
|
||||
T.evaluate(0.0)
|
||||
|
||||
check_error(block_has_option_vars, 2)
|
||||
|
||||
|
||||
def test_implicit_root_has_attrs():
|
||||
def implicit_root_has_read():
|
||||
T.reads([]) # error: implicit root does not support reads
|
||||
T.evaluate(0.0)
|
||||
|
||||
def implicit_root_has_write():
|
||||
T.writes([]) # error: implicit root does not support writes
|
||||
T.evaluate(0.0)
|
||||
|
||||
def implicit_root_has_attrs():
|
||||
T.sblock_attr({}) # error: implicit root does not support sblock_attr
|
||||
T.evaluate(0.0)
|
||||
|
||||
def implicit_root_has_predicate():
|
||||
T.where(True) # error: implicit root does not support predicate
|
||||
T.evaluate(0.0)
|
||||
|
||||
def implicit_root_has_axes():
|
||||
v = T.axis.S(0, 0) # error: implicit root does not support axis define
|
||||
T.evaluate(0.0)
|
||||
|
||||
check_error(implicit_root_has_read, 2)
|
||||
check_error(implicit_root_has_write, 2)
|
||||
check_error(implicit_root_has_attrs, 2)
|
||||
check_error(implicit_root_has_predicate, 2)
|
||||
check_error(implicit_root_has_axes, 2)
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def elementwise_not_affine(a: T.handle, b: T.handle) -> None:
|
||||
A = T.match_buffer(a, (128, 128, 128, 128))
|
||||
B = T.match_buffer(b, (128, 128, 128, 128))
|
||||
for i, j, k, l in T.grid(128, 128, 128, 8):
|
||||
with T.sblock("B"):
|
||||
vi, vj, vk = T.axis.remap("SSS", [i, j, k])
|
||||
vl = T.axis.S(128, l * 16)
|
||||
B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * 2.0
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def elementwise_non_single_branch(a: T.handle, b: T.handle) -> None:
|
||||
A = T.match_buffer(a, (128, 128, 128))
|
||||
C = T.sblock_alloc_buffer((128, 128, 128))
|
||||
B = T.match_buffer(b, (128, 128, 128))
|
||||
for i, j in T.grid(128, 128):
|
||||
for k in T.serial(0, 128):
|
||||
with T.sblock("C"):
|
||||
vi, vj, vk = T.axis.remap("SSS", [i, j, k])
|
||||
C[vi, vj, vk] = A[vi, vj, vk] * 2.0
|
||||
for k in T.serial(0, 128):
|
||||
with T.sblock("B"):
|
||||
vi, vj, vk = T.axis.remap("SSS", [i, j, k])
|
||||
B[vi, vj, vk] = C[vi, vj, vk] * 2.0
|
||||
|
||||
|
||||
def test_reorder_fail_block():
|
||||
sch = tvm.s_tir.Schedule(elementwise_not_affine, debug_mask="all")
|
||||
block_b = sch.get_sblock("B")
|
||||
i, j, k, l = sch.get_loops(block_b)
|
||||
with pytest.raises(tvm.s_tir.ScheduleError) as execinfo:
|
||||
sch.reorder(l, i)
|
||||
expected_sub_error_message = (
|
||||
" # tirx.SBlock#0\n"
|
||||
' with T.sblock("B"):\n'
|
||||
" ^^^^^^^^^^^^^^^^^^^\n"
|
||||
)
|
||||
assert expected_sub_error_message in str(execinfo.value)
|
||||
|
||||
|
||||
def test_reorder_fail_nested_loop_inner():
|
||||
sch = tvm.s_tir.Schedule(elementwise_non_single_branch, debug_mask="all")
|
||||
block_b = sch.get_sblock("B")
|
||||
i, j, k = sch.get_loops(block_b)
|
||||
with pytest.raises(tvm.s_tir.ScheduleError) as execinfo:
|
||||
sch.reorder(k, i)
|
||||
expected_sub_error_message = (
|
||||
" for i in range(128):\n"
|
||||
" # tirx.For#0\n"
|
||||
" for j in range(128):\n"
|
||||
" ^^^^^^^^^^^^^^^^^^^^\n"
|
||||
)
|
||||
assert expected_sub_error_message in str(execinfo.value)
|
||||
|
||||
|
||||
def test_fuse_fail_nested_loop_outer():
|
||||
sch = tvm.s_tir.Schedule(elementwise_non_single_branch, debug_mask="all")
|
||||
block_b = sch.get_sblock("B")
|
||||
i, j, k = sch.get_loops(block_b)
|
||||
with pytest.raises(tvm.s_tir.ScheduleError) as execinfo:
|
||||
sch.fuse(k, i)
|
||||
expected_sub_error_message = (
|
||||
" # tirx.For#1\n"
|
||||
" for i in range(128):\n"
|
||||
" ^^^^^^^^^^^^^^^^^^^^\n"
|
||||
" for j in range(128):\n"
|
||||
)
|
||||
assert expected_sub_error_message in str(execinfo.value)
|
||||
|
||||
|
||||
def test_report_error_root_block():
|
||||
sch = tvm.s_tir.Schedule(elementwise_non_single_branch, debug_mask="all")
|
||||
root = sch.get_sblock("root")
|
||||
with pytest.raises(tvm.s_tir.ScheduleError) as execinfo:
|
||||
sch.compute_inline(root)
|
||||
expected_sub_error_message = (
|
||||
' # tirx.SBlock#0\n with T.sblock("root"):\n ^^^^^^^^^^^^^^^^^^^^^^\n'
|
||||
)
|
||||
assert expected_sub_error_message in str(execinfo.value)
|
||||
|
||||
|
||||
def test_load_var():
|
||||
def load_var_multiple() -> None:
|
||||
d = T.float32()
|
||||
d[2] = d[2, 1] # error cannot provide two indices to load
|
||||
|
||||
check_error(load_var_multiple, 3)
|
||||
|
||||
|
||||
def test_store_var():
|
||||
def store_var_multiple() -> None:
|
||||
d = T.float32()
|
||||
d[2, 1] = d[1] # error cannot provide two indices to store
|
||||
|
||||
check_error(store_var_multiple, 3)
|
||||
|
||||
|
||||
def test_load_handle():
|
||||
def load_handle(h: T.handle) -> None:
|
||||
h_ = T.match_buffer(h, [1])
|
||||
h_[0] = h[0] # error cannot load from handle
|
||||
|
||||
check_error(load_handle, 3)
|
||||
|
||||
|
||||
def test_store_handle():
|
||||
def store_handle(h: T.handle) -> None:
|
||||
h_ = T.match_buffer(h, [1])
|
||||
h[0] = h_[0] # error cannot store to handle
|
||||
|
||||
check_error(store_handle, 3)
|
||||
|
||||
|
||||
def test_binop_bad_ast_type():
|
||||
def binop_bad_ast_type(h: T.handle):
|
||||
h_ = T.match_buffer(h, [1])
|
||||
h_[0] = h + [2] # error rhs should be a primexpr
|
||||
|
||||
check_error(binop_bad_ast_type, 3)
|
||||
|
||||
|
||||
def test_binop_bad_type():
|
||||
def binop_bad_type(h: T.handle):
|
||||
h_ = T.match_buffer(h, [1])
|
||||
h_[0] = h + 2 # error lhs and rhs should be the same type
|
||||
|
||||
check_error(binop_bad_type, 3)
|
||||
|
||||
|
||||
def test_non_integer_typed_block_iter():
|
||||
def non_integer_typed_block_iter():
|
||||
with T.sblock():
|
||||
i = T.axis.S(0.1, 0.1) # error IterVar requires an integer dtype
|
||||
|
||||
check_error(non_integer_typed_block_iter, 3)
|
||||
|
||||
|
||||
def test_illegal_buffer_slice():
|
||||
def strided_buffer_region(A: T.handle):
|
||||
# do not allow stride in buffer region
|
||||
A = T.match_buffer((128, 128), "int32")
|
||||
with T.sblock():
|
||||
T.reads([])
|
||||
T.writes([A[0:128:2, 0:128:3]]) # error
|
||||
T.evaluate(T.call_extern("strided_compute", dtype=""))
|
||||
|
||||
def access_reversed_slice(A: T.handle):
|
||||
# do not allow reversed slice step
|
||||
A = T.match_buffer((128,), "int32")
|
||||
A[0:128:-1] = T.broadcast(1, 128) # error
|
||||
|
||||
def access_non_const_slice_length(A: T.handle):
|
||||
# do not allow non-constant slice length
|
||||
A = T.match_buffer((128,), "int32")
|
||||
for i in range(4):
|
||||
T.evaluate(A[0:i:1]) # error
|
||||
|
||||
check_error(strided_buffer_region, 3)
|
||||
check_error(access_reversed_slice, 3)
|
||||
check_error(access_non_const_slice_length, 3)
|
||||
|
||||
|
||||
def test_syntax_sugar_fail():
|
||||
def loop_syntax_sugar_fail(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, (128,))
|
||||
for i in T.thread_binding(128, 128):
|
||||
A[i] = A[i] * 2.0
|
||||
|
||||
check_error(loop_syntax_sugar_fail, 3)
|
||||
|
||||
|
||||
def test_multi_line_error_report():
|
||||
"""A parse error whose offending AST node spans several physical source
|
||||
lines must render ALL spanned lines (each with its own gutter line number
|
||||
and an underline covering the span), not just the first line."""
|
||||
|
||||
# The offending call (`T.axis.remap(...)`) is deliberately split across
|
||||
# four physical lines so its AST node spans lineno..end_lineno > lineno.
|
||||
source_code = "\n".join(
|
||||
[
|
||||
"@T.prim_func(s_tir=True)",
|
||||
"def f() -> None:",
|
||||
" for i, j in T.grid(16, 16):",
|
||||
" vi, vj = T.axis.remap(",
|
||||
' "S",',
|
||||
" [i, j],",
|
||||
" ) # error",
|
||||
" T.evaluate(1.0)",
|
||||
]
|
||||
)
|
||||
|
||||
with pytest.raises(tvm.error.DiagnosticError) as execinfo:
|
||||
from_source(source_code)
|
||||
err_str = str(execinfo.value)
|
||||
|
||||
# All four spanned source lines must appear in the rendered snippet.
|
||||
assert "T.axis.remap(" in err_str, err_str
|
||||
assert '"S",' in err_str, err_str
|
||||
assert "[i, j]," in err_str, err_str
|
||||
# The trailing `)` closing line is also part of the span.
|
||||
rendered_lines = err_str.splitlines()
|
||||
assert any(" 7 " in line and ")" in line for line in rendered_lines), err_str
|
||||
# The underline carets must be present on more than one line (multi-line).
|
||||
marker_lines = [line for line in rendered_lines if "^" in line]
|
||||
assert len(marker_lines) >= 2, err_str
|
||||
# The gutter must show distinct line numbers for the spanned lines.
|
||||
assert " 4 " in err_str and " 5 " in err_str and " 6 " in err_str, err_str
|
||||
|
||||
|
||||
def test_format_source_snippet_multi_line():
|
||||
"""Unit-level check that _format_source_snippet renders every line in a
|
||||
multi-line span, with the underline covering start-col..EOL on the first
|
||||
line, full interior lines, and col-1..end-col on the last line."""
|
||||
from tvm.script.parser.core.diagnostics import _format_source_snippet
|
||||
|
||||
source_lines = [
|
||||
"first ignored line\n",
|
||||
" foo(bar,\n",
|
||||
" baz,\n",
|
||||
" qux)\n",
|
||||
"last ignored line\n",
|
||||
]
|
||||
# Span lines 2..4 (1-based), starting at col 5 ('foo'), ending at col 13
|
||||
# (exclusive) on line 4.
|
||||
snippet = _format_source_snippet(
|
||||
source_lines, lineno=2, col_offset=5, end_lineno=4, end_col_offset=13
|
||||
)
|
||||
lines = snippet.splitlines()
|
||||
# All three spanned source lines must be present.
|
||||
assert any("foo(bar," in line for line in lines), snippet
|
||||
assert any("baz," in line for line in lines), snippet
|
||||
assert any("qux)" in line for line in lines), snippet
|
||||
# Underline carets present on the first line under 'foo(bar,'.
|
||||
assert "^" in snippet, snippet
|
||||
# The line numbers 2, 3, 4 appear in the gutter.
|
||||
assert " 2 |" in snippet and " 3 |" in snippet and " 4 |" in snippet, snippet
|
||||
|
||||
|
||||
def test_format_source_snippet_single_line_unchanged():
|
||||
"""A single-line span (end_lineno == lineno) underlines only the
|
||||
[col_offset, end_col_offset) columns on that one line."""
|
||||
from tvm.script.parser.core.diagnostics import _format_source_snippet
|
||||
|
||||
source_lines = ["ignored\n", " abc + def\n", "ignored\n"]
|
||||
# Underline just 'abc' (cols 5..8 exclusive) on line 2.
|
||||
snippet = _format_source_snippet(
|
||||
source_lines, lineno=2, col_offset=5, end_lineno=2, end_col_offset=8
|
||||
)
|
||||
lines = snippet.splitlines()
|
||||
# Exactly one source-text line and one marker line (plus the leading gutter).
|
||||
text_lines = [line for line in lines if "abc + def" in line]
|
||||
assert len(text_lines) == 1, snippet
|
||||
marker_line = next(line for line in lines if "^" in line)
|
||||
assert marker_line.count("^") == 3, snippet
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,44 @@
|
||||
# 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.
|
||||
"""Unittests for tvm.script.ir_builder.base"""
|
||||
|
||||
import pytest
|
||||
|
||||
from tvm.script.ir_builder import IRBuilder
|
||||
|
||||
|
||||
def test_ir_builder_scope():
|
||||
with IRBuilder() as ib: # pylint: disable=invalid-name
|
||||
assert IRBuilder.current() == ib
|
||||
|
||||
|
||||
def test_ir_builder_multi_scope():
|
||||
with IRBuilder() as ib: # pylint: disable=invalid-name
|
||||
with IRBuilder() as ib2: # pylint: disable=invalid-name
|
||||
assert IRBuilder.current() == ib2
|
||||
assert IRBuilder.current() == ib
|
||||
|
||||
|
||||
def test_ir_builder_no_scope():
|
||||
with pytest.raises(ValueError):
|
||||
IRBuilder.current()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_ir_builder_scope()
|
||||
test_ir_builder_multi_scope()
|
||||
test_ir_builder_no_scope()
|
||||
@@ -0,0 +1,44 @@
|
||||
# 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: F401
|
||||
"""Unittests for tvm.script.ir_builder.ir"""
|
||||
|
||||
import pytest
|
||||
|
||||
import tvm.testing
|
||||
from tvm import ir
|
||||
from tvm.ir.base import assert_structural_equal
|
||||
from tvm.script.ir_builder import IRBuilder
|
||||
from tvm.script.ir_builder import ir as I
|
||||
|
||||
|
||||
def test_ir_builder_irmodule():
|
||||
with IRBuilder() as ib: # pylint: disable=invalid-name
|
||||
with I.ir_module():
|
||||
pass
|
||||
|
||||
# the ir_module generated by IRBuilder
|
||||
ir_module_actual = ib.get()
|
||||
|
||||
# the expected prim_func
|
||||
ir_module_expected = ir.IRModule(None, None)
|
||||
|
||||
assert_structural_equal(ir_module_actual, ir_module_expected, map_free_vars=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,510 @@
|
||||
# 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=invalid-name, missing-docstring
|
||||
# ruff: noqa: F401, F841
|
||||
"""Unittests for tvm.script.ir_builder.tirx"""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.runtime
|
||||
import tvm.testing
|
||||
from tvm import tirx
|
||||
from tvm.ir.base import assert_structural_equal
|
||||
from tvm.script.ir_builder import IRBuilder
|
||||
from tvm.script.ir_builder import tirx as T
|
||||
|
||||
|
||||
def test_ir_builder_tir_primfunc_base():
|
||||
with IRBuilder() as ib:
|
||||
with T.prim_func(s_tir=True):
|
||||
T.evaluate(0)
|
||||
|
||||
# the prim_func generated by IRBuilder
|
||||
prim_func_actual = ib.get()
|
||||
|
||||
# the expected prim_func
|
||||
prim_func_expected = tirx.PrimFunc(
|
||||
params=[],
|
||||
body=tirx.Evaluate(0),
|
||||
ret_type=None,
|
||||
buffer_map=None,
|
||||
attrs=tvm.ir.make_node("ir.DictAttrs", s_tir=True),
|
||||
)
|
||||
|
||||
# Check if the generated ir is expected
|
||||
assert_structural_equal(prim_func_actual, prim_func_expected, map_free_vars=True)
|
||||
|
||||
|
||||
def test_ir_builder_tir_primfunc_complete():
|
||||
with IRBuilder() as ib:
|
||||
with T.prim_func(s_tir=True):
|
||||
T.arg("a", T.handle())
|
||||
T.arg("b", T.int64())
|
||||
T.arg("c", T.Buffer((128, 128), "float32"))
|
||||
d = T.arg("d", T.handle())
|
||||
e = T.arg("e", T.Buffer((1024,), "int8"))
|
||||
T.func_attr({"key": "value"})
|
||||
T.func_ret(tvm.ir.PrimType("int64"))
|
||||
buffer_d = T.match_buffer(d, (64, 64), "int64")
|
||||
T.evaluate(0)
|
||||
|
||||
# the prim_func generated by IRBuilder
|
||||
prim_func_actual = ib.get()
|
||||
|
||||
# the expected prim_func
|
||||
c_handle, c_buffer = (
|
||||
tirx.Var("c_handle", tvm.ir.PointerType(tvm.ir.PrimType("void"))),
|
||||
tirx.decl_buffer((128, 128), "float32", name="c", layout=None),
|
||||
)
|
||||
d_handle, d_buffer = (
|
||||
tirx.Var("d", tvm.ir.PointerType(tvm.ir.PrimType("void"))),
|
||||
tirx.decl_buffer((64, 64), "int64", name="d", layout=None),
|
||||
)
|
||||
e_handle, e_buffer = (
|
||||
tirx.Var("e_handle", tvm.ir.PointerType(tvm.ir.PrimType("void"))),
|
||||
tirx.decl_buffer((1024,), "int8", name="e", layout=None),
|
||||
)
|
||||
prim_func_expected = tirx.PrimFunc(
|
||||
params=[
|
||||
tirx.Var("a", tvm.ir.PointerType(tvm.ir.PrimType("void"))),
|
||||
tirx.Var("b", "int64"),
|
||||
c_handle,
|
||||
d_handle,
|
||||
e_handle,
|
||||
],
|
||||
body=tirx.Evaluate(0),
|
||||
ret_type=tvm.ir.PrimType("int64"),
|
||||
buffer_map={c_handle: c_buffer, d_handle: d_buffer, e_handle: e_buffer},
|
||||
attrs=tvm.ir.make_node("ir.DictAttrs", key="value", s_tir=True),
|
||||
)
|
||||
|
||||
# Check if the generated ir is expected
|
||||
assert_structural_equal(prim_func_actual, prim_func_expected, map_free_vars=True)
|
||||
|
||||
|
||||
def test_ir_builder_tir_block_base():
|
||||
with IRBuilder() as ib:
|
||||
with T.sblock("block"):
|
||||
T.evaluate(0)
|
||||
|
||||
# the block generated by IRBuilder
|
||||
block_realize_actual = ib.get()
|
||||
|
||||
# the expected block
|
||||
block_expected = tirx.SBlock(
|
||||
iter_vars=[],
|
||||
reads=[],
|
||||
writes=[],
|
||||
name_hint="block",
|
||||
body=tirx.Evaluate(0),
|
||||
alloc_buffers=None,
|
||||
match_buffers=None,
|
||||
annotations={"tirx.script_parsing_detect_access": tirx.IntImm("int64", 3)},
|
||||
)
|
||||
block_realize_expected = tirx.SBlockRealize(
|
||||
iter_values=[],
|
||||
predicate=True,
|
||||
block=block_expected,
|
||||
)
|
||||
|
||||
# Check if the generated ir is expected
|
||||
assert_structural_equal(block_realize_actual, block_realize_expected, map_free_vars=True)
|
||||
|
||||
|
||||
def test_ir_builder_tir_block_complete():
|
||||
with IRBuilder() as ib:
|
||||
a = T.int64()
|
||||
b = T.Buffer((128, 128), "float32")
|
||||
c = T.Buffer((128, 128), "float32")
|
||||
d = T.int32()
|
||||
e = T.Buffer((128, 128), "float32")
|
||||
f = T.int32()
|
||||
with T.sblock("block"):
|
||||
T.where(a > 1)
|
||||
T.reads(b[0:16, 0:16])
|
||||
T.writes(c[d:128, d:128])
|
||||
T.sblock_attr({"key": "value"})
|
||||
T.sblock_alloc_buffer((128, 128), "float32")
|
||||
T.match_buffer(e[0:32, 0:32], (32, 32), "float32")
|
||||
T.axis.spatial(128, f)
|
||||
T.evaluate(0)
|
||||
|
||||
# the block generated by IRBuilder
|
||||
block_realize_actual = ib.get()
|
||||
|
||||
# the expected block
|
||||
var_a = tirx.Var("a", "int64")
|
||||
buffer_b = tirx.decl_buffer((128, 128), "float32", name="b")
|
||||
buffer_c = tirx.decl_buffer((128, 128), "float32", name="c")
|
||||
var_d = tirx.Var("d", "int32")
|
||||
buffer_e = tirx.decl_buffer((128, 128), "float32", name="c")
|
||||
var_f = tirx.Var("f", "int32")
|
||||
block_expected = tirx.SBlock(
|
||||
iter_vars=[tirx.IterVar((0, 128), tirx.Var("", "int32"), iter_type=tirx.IterVar.DataPar)],
|
||||
reads=[buffer_b[0:16, 0:16]],
|
||||
writes=[buffer_c[var_d:128, var_d:128]],
|
||||
name_hint="block",
|
||||
body=tirx.Evaluate(0),
|
||||
alloc_buffers=[tirx.decl_buffer((128, 128), "float32")],
|
||||
match_buffers=[
|
||||
tirx.MatchBufferRegion(tirx.decl_buffer((32, 32), "float32"), buffer_e[0:32, 0:32])
|
||||
],
|
||||
annotations={"key": "value"},
|
||||
)
|
||||
block_realize_expected = tirx.SBlockRealize(
|
||||
iter_values=[var_f],
|
||||
predicate=var_a > 1,
|
||||
block=block_expected,
|
||||
)
|
||||
|
||||
# Check if the generated ir is expected
|
||||
assert_structural_equal(block_realize_actual, block_realize_expected, map_free_vars=True)
|
||||
|
||||
|
||||
def test_ir_builder_tir_axis():
|
||||
with IRBuilder() as ib:
|
||||
a = T.int32()
|
||||
b = T.int32()
|
||||
c = T.int32()
|
||||
d = T.int32()
|
||||
with T.sblock("block"):
|
||||
T.axis.spatial(8, a)
|
||||
T.axis.reduce(16, b)
|
||||
T.axis.scan(32, c)
|
||||
T.axis.opaque(64, d)
|
||||
T.evaluate(0)
|
||||
|
||||
# the block generated by IRBuilder
|
||||
block_realize_actual = ib.get()
|
||||
|
||||
# the expected block
|
||||
var_a = tirx.Var("a", "int32")
|
||||
var_b = tirx.Var("b", "int32")
|
||||
var_c = tirx.Var("c", "int32")
|
||||
var_d = tirx.Var("d", "int32")
|
||||
block_expected = tirx.SBlock(
|
||||
iter_vars=[
|
||||
tirx.IterVar((0, 8), tirx.Var("", "int32"), iter_type=tirx.IterVar.DataPar),
|
||||
tirx.IterVar((0, 16), tirx.Var("", "int32"), iter_type=tirx.IterVar.CommReduce),
|
||||
tirx.IterVar((0, 32), tirx.Var("", "int32"), iter_type=tirx.IterVar.Ordered),
|
||||
tirx.IterVar((0, 64), tirx.Var("", "int32"), iter_type=tirx.IterVar.Opaque),
|
||||
],
|
||||
reads=[],
|
||||
writes=[],
|
||||
name_hint="block",
|
||||
body=tirx.Evaluate(0),
|
||||
annotations={"tirx.script_parsing_detect_access": tirx.IntImm("int64", 3)},
|
||||
)
|
||||
block_realize_expected = tirx.SBlockRealize(
|
||||
iter_values=[var_a, var_b, var_c, var_d],
|
||||
predicate=True,
|
||||
block=block_expected,
|
||||
)
|
||||
|
||||
# Check if the generated ir is expected
|
||||
assert_structural_equal(block_realize_actual, block_realize_expected, map_free_vars=True)
|
||||
|
||||
|
||||
def test_ir_builder_tir_for():
|
||||
with IRBuilder() as ib:
|
||||
with T.serial(128) as a:
|
||||
with T.parallel(64) as b:
|
||||
with T.vectorized(32) as c:
|
||||
with T.unroll(16) as d:
|
||||
with T.thread_binding(8, thread="threadIdx.x") as e:
|
||||
T.evaluate(0)
|
||||
|
||||
# the for generated by IRBuilder
|
||||
for_actual = ib.get()
|
||||
|
||||
# the expected for
|
||||
thread_binding_expected = tirx.For(
|
||||
loop_var=tirx.Var("", "int32"),
|
||||
min=0,
|
||||
extent=8,
|
||||
kind=tirx.ForKind.THREAD_BINDING,
|
||||
body=tirx.Evaluate(0),
|
||||
thread_binding=tirx.IterVar(
|
||||
None, tirx.Var("", "int32"), tirx.IterVar.ThreadIndex, "threadIdx.x"
|
||||
),
|
||||
)
|
||||
unroll_expected = tirx.For(
|
||||
loop_var=tirx.Var("", "int32"),
|
||||
min=0,
|
||||
extent=16,
|
||||
kind=tirx.ForKind.UNROLLED,
|
||||
body=thread_binding_expected,
|
||||
)
|
||||
vectorized_expected = tirx.For(
|
||||
loop_var=tirx.Var("", "int32"),
|
||||
min=0,
|
||||
extent=32,
|
||||
kind=tirx.ForKind.VECTORIZED,
|
||||
body=unroll_expected,
|
||||
)
|
||||
parallel_expected = tirx.For(
|
||||
loop_var=tirx.Var("", "int32"),
|
||||
min=0,
|
||||
extent=64,
|
||||
kind=tirx.ForKind.PARALLEL,
|
||||
body=vectorized_expected,
|
||||
)
|
||||
for_expected = tirx.For(
|
||||
loop_var=tirx.Var("", "int32"),
|
||||
min=0,
|
||||
extent=128,
|
||||
kind=tirx.ForKind.SERIAL,
|
||||
body=parallel_expected,
|
||||
)
|
||||
|
||||
# Check if the generated ir is expected
|
||||
assert_structural_equal(for_actual, for_expected, map_free_vars=True)
|
||||
|
||||
|
||||
def test_ir_builder_tir_for_uint():
|
||||
with IRBuilder() as ib:
|
||||
with T.serial(tirx.const(128, "uint32")) as a:
|
||||
T.evaluate(0)
|
||||
|
||||
# the for generated by IRBuilder
|
||||
for_actual = ib.get()
|
||||
|
||||
for_expected = tirx.For(
|
||||
loop_var=tirx.Var("", "uint32"),
|
||||
min=tirx.const(0, "uint32"),
|
||||
extent=tirx.const(128, "uint32"),
|
||||
kind=tirx.ForKind.SERIAL,
|
||||
body=tirx.Evaluate(0),
|
||||
)
|
||||
|
||||
# Check if the generated ir is expected
|
||||
assert_structural_equal(for_actual, for_expected, map_free_vars=True)
|
||||
|
||||
|
||||
def test_ir_builder_tir_assert():
|
||||
with IRBuilder() as ib:
|
||||
with T.Assert(T.int32() == 0, message="a is 0"):
|
||||
T.evaluate(0)
|
||||
# the assert generated by IRBuilder
|
||||
assert_actual = ib.get()
|
||||
|
||||
# AssertStmt is a leaf. The frame emits the assert and then the body stmts as siblings.
|
||||
assert_expected = tirx.SeqStmt(
|
||||
[
|
||||
tirx.AssertStmt(
|
||||
T.int32() == 0,
|
||||
tirx.StringImm("RuntimeError"),
|
||||
[tirx.StringImm("a is 0")],
|
||||
),
|
||||
tirx.Evaluate(0),
|
||||
]
|
||||
)
|
||||
|
||||
# Check if the generated ir is expected
|
||||
assert_structural_equal(assert_actual, assert_expected, map_free_vars=True)
|
||||
|
||||
|
||||
def test_ir_builder_tir_bind():
|
||||
# Test that T.bind emits a flat Bind statement and returns the Var.
|
||||
with IRBuilder() as ib:
|
||||
v = T.bind(tirx.IntImm("int32", 2))
|
||||
# the let binding generated by IRBuilder
|
||||
let_actual = ib.get()
|
||||
|
||||
# Bind is now flat (no body), so a single Bind stmt is emitted.
|
||||
let_expected = tirx.Bind(T.int32(), tirx.IntImm("int32", 2))
|
||||
|
||||
# Check if the generated ir is expected
|
||||
assert_structural_equal(let_actual, let_expected, map_free_vars=True)
|
||||
# Check that the returned value is a Var
|
||||
assert isinstance(v, tirx.Var)
|
||||
|
||||
|
||||
def test_ir_builder_tir_thread():
|
||||
with IRBuilder() as ib:
|
||||
with T.prim_func(s_tir=True):
|
||||
brow = T.env_thread("blockIdx.y")
|
||||
with T.launch_thread(brow, 1):
|
||||
T.evaluate(0)
|
||||
|
||||
# the prim_func generated by IRBuilder
|
||||
ir_actual = ib.get()
|
||||
|
||||
# the expected prim_func
|
||||
iter_var = tirx.IterVar((0, 1), "v", iter_type=1, thread_tag="blockIdx.y")
|
||||
attr_stmt = tirx.AttrStmt(iter_var, "thread_extent", 1, tirx.Evaluate(0))
|
||||
func = tirx.PrimFunc([], attr_stmt).with_attr("s_tir", True)
|
||||
|
||||
# Check if the generated ir is expected
|
||||
assert_structural_equal(ir_actual, func, map_free_vars=True)
|
||||
|
||||
|
||||
def test_ir_builder_tir_allocate():
|
||||
with IRBuilder() as ib:
|
||||
with T.prim_func(s_tir=True):
|
||||
T.func_name("test")
|
||||
buf = T.alloc_buffer([10], "float32", scope="local")
|
||||
T.evaluate(1)
|
||||
|
||||
# the allocate generated by IRBuilder
|
||||
ir_actual = ib.get()
|
||||
body = ir_actual.body
|
||||
|
||||
# AllocBuffer is flat: body should be a SeqStmt with [AllocBuffer, Evaluate(1)]
|
||||
assert isinstance(body, tirx.SeqStmt), f"Expected SeqStmt but got {type(body)}"
|
||||
assert len(body) == 2
|
||||
assert isinstance(body[0], tirx.AllocBuffer)
|
||||
assert isinstance(body[1], tirx.Evaluate)
|
||||
assert body[1].value.value == 1
|
||||
|
||||
|
||||
def test_ir_builder_tir_while():
|
||||
with IRBuilder() as ib:
|
||||
with T.While(T.int32() > 0):
|
||||
T.evaluate(0)
|
||||
|
||||
# the while generated by IRBuilder
|
||||
ir_actual = ib.get()
|
||||
|
||||
# the expected while
|
||||
ir_expected = tirx.While(tirx.Var("x", "int32") > 0, tirx.Evaluate(0))
|
||||
|
||||
# Check if the generated ir is expected
|
||||
assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)
|
||||
|
||||
|
||||
def test_ir_builder_tir_if_then_else():
|
||||
with IRBuilder() as ib:
|
||||
with T.If(T.int32() < 12):
|
||||
with T.Then():
|
||||
T.evaluate(T.int32(0))
|
||||
with T.Else():
|
||||
T.evaluate(T.int32(1))
|
||||
|
||||
# the if_then_else generated by IRBuilder
|
||||
ir_actual = ib.get()
|
||||
|
||||
# the expected if_then_else
|
||||
ir_expected = tirx.IfThenElse(
|
||||
tirx.Var("c", "int32") < 12,
|
||||
tirx.Evaluate(tirx.IntImm("int32", 0)),
|
||||
tirx.Evaluate(tirx.IntImm("int32", 1)),
|
||||
)
|
||||
|
||||
# Check if the generated ir is expected
|
||||
assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)
|
||||
|
||||
|
||||
def test_ir_builder_tir_buffer_store():
|
||||
buffer_a = T.Buffer((10, 10), "float32")
|
||||
i = T.int32()
|
||||
with IRBuilder() as ib:
|
||||
T.buffer_store(buffer_a, 0.1, [0, i])
|
||||
|
||||
# the buffer store generated by IRBuilder
|
||||
ir_actual = ib.get()
|
||||
|
||||
# the expected buffer store
|
||||
ir_expected = tirx.BufferStore(buffer_a, 0.1, [0, i])
|
||||
|
||||
# Check if the generated ir is expected
|
||||
assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)
|
||||
|
||||
|
||||
def test_ir_builder_tir_buffer_store_scalable_vec():
|
||||
buffer_a = T.Buffer((30,), "float32")
|
||||
value = T.broadcast(0.11, 4 * tvm.tirx.vscale())
|
||||
index = T.ramp(0, 1, 4 * tvm.tirx.vscale())
|
||||
|
||||
with IRBuilder() as ib:
|
||||
T.buffer_store(buffer_a, value, [index])
|
||||
|
||||
# the buffer store generated by IRBuilder
|
||||
ir_actual = ib.get()
|
||||
|
||||
# the expected buffer store
|
||||
ir_expected = tirx.BufferStore(buffer_a, value, [index])
|
||||
|
||||
# Check if the generated ir is expected
|
||||
assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)
|
||||
|
||||
|
||||
def test_ir_builder_tir_buffer_store_predicate():
|
||||
buffer_a = T.Buffer((30,), "float32")
|
||||
value = T.broadcast(0.11, T.vscale() * 4)
|
||||
index = T.ramp(0, 1, T.vscale() * 4)
|
||||
predicate = T.broadcast(T.bool(True), T.vscale() * 4)
|
||||
|
||||
with IRBuilder() as ib:
|
||||
T.buffer_store(buffer_a, value, [index], predicate)
|
||||
|
||||
ir_actual = ib.get()
|
||||
ir_expected = tirx.BufferStore(buffer_a, value, [index], predicate)
|
||||
assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)
|
||||
|
||||
|
||||
def test_ir_builder_tir_evaluate():
|
||||
with IRBuilder() as ib:
|
||||
T.evaluate(0)
|
||||
# the evaluate generated by IRBuilder
|
||||
eval_actual = ib.get()
|
||||
|
||||
# the expected evaluate
|
||||
eval_expected = tirx.Evaluate(0)
|
||||
|
||||
# Check if the generated ir is expected
|
||||
assert_structural_equal(eval_actual, eval_expected, map_free_vars=True)
|
||||
|
||||
|
||||
def test_ir_builder_tir_decl_buffer():
|
||||
with IRBuilder() as ib:
|
||||
with T.prim_func(s_tir=True):
|
||||
T.func_name("test")
|
||||
buf = T.decl_buffer([128, 128], "float32")
|
||||
T.evaluate(1)
|
||||
|
||||
# the decl_buffer generated by IRBuilder
|
||||
ir_actual = ib.get()
|
||||
body = ir_actual.body
|
||||
|
||||
# decl_buffer without data emits AllocBuffer (flat): body should be SeqStmt
|
||||
assert isinstance(body, tirx.SeqStmt), f"Expected SeqStmt but got {type(body)}"
|
||||
assert len(body) == 2
|
||||
assert isinstance(body[0], tirx.AllocBuffer)
|
||||
assert isinstance(body[1], tirx.Evaluate)
|
||||
assert body[1].value.value == 1
|
||||
|
||||
|
||||
def test_ir_builder_tir_inline():
|
||||
with IRBuilder() as ib:
|
||||
m, n = T.meta_var(1), T.meta_var(2)
|
||||
a, b = T.meta_var([3, 4])
|
||||
T.evaluate(m.value + n.value + a.value + b.value)
|
||||
# the evaluate generated by IRBuilder
|
||||
eval_actual = ib.get()
|
||||
|
||||
# the expected evaluate
|
||||
eval_expected = tirx.Evaluate(10)
|
||||
|
||||
# Check if the generated ir is expected
|
||||
assert_structural_equal(eval_actual, eval_expected, map_free_vars=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,86 @@
|
||||
# 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 tvm
|
||||
from tvm.script import tirx as T
|
||||
|
||||
|
||||
def test_meta_programming_matmul():
|
||||
def matmul_generator(M: int, N: int, K: int, dtype: str):
|
||||
@T.prim_func(s_tir=True)
|
||||
def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [M, K], dtype=dtype)
|
||||
B = T.match_buffer(b, [N, K], dtype=dtype)
|
||||
C = T.match_buffer(c, [M, N], dtype=dtype)
|
||||
|
||||
for i, j, k in T.grid(M, N, K):
|
||||
with T.sblock():
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = T.float32(0)
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
|
||||
|
||||
return matmul
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def matmul_128_128_128_fp16(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128], dtype="float16")
|
||||
B = T.match_buffer(b, [128, 128], dtype="float16")
|
||||
C = T.match_buffer(c, [128, 128], dtype="float16")
|
||||
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
with T.sblock():
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = T.float32(0)
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
|
||||
|
||||
f = matmul_generator(128, 128, 128, "float16").with_attr("global_symbol", "main")
|
||||
tvm.ir.assert_structural_equal(f, matmul_128_128_128_fp16.with_attr("global_symbol", "main"))
|
||||
|
||||
|
||||
def test_meta_programming_uncaptured_var():
|
||||
def generate_erf(dtype):
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((1,), dtype), C: T.Buffer((1,), dtype)):
|
||||
for i in range(1):
|
||||
with T.sblock("C"):
|
||||
C[i] = T.erf(A[i])
|
||||
|
||||
return main
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def fp32(A: T.Buffer((1,), "float32"), C: T.Buffer((1,), "float32")):
|
||||
for i in range(1):
|
||||
with T.sblock("C"):
|
||||
C[i] = T.erf(A[i])
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def fp16(A: T.Buffer((1,), "float16"), C: T.Buffer((1,), "float16")):
|
||||
for i in range(1):
|
||||
with T.sblock("C"):
|
||||
C[i] = T.erf(A[i])
|
||||
|
||||
f1 = generate_erf("float32").with_attr("global_symbol", "main")
|
||||
tvm.ir.assert_structural_equal(f1, fp32.with_attr("global_symbol", "main"))
|
||||
f2 = generate_erf("float16").with_attr("global_symbol", "main")
|
||||
tvm.ir.assert_structural_equal(f2, fp16.with_attr("global_symbol", "main"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_meta_programming_matmul()
|
||||
test_meta_programming_uncaptured_var()
|
||||
@@ -0,0 +1,267 @@
|
||||
# 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 numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import tirx as T
|
||||
from tvm.testing import env
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def get_valid_counts(
|
||||
data: T.handle,
|
||||
valid_count: T.handle,
|
||||
out: T.handle,
|
||||
out_indices: T.handle,
|
||||
score_threshold: T.float32,
|
||||
id_index: T.int32,
|
||||
score_index: T.int32,
|
||||
) -> None:
|
||||
data_buf = T.match_buffer(data, (1, 2500, 6), "float32")
|
||||
valid_count_buf = T.match_buffer(valid_count, (1,), "int32")
|
||||
out_buf = T.match_buffer(out, (1, 2500, 6), "float32")
|
||||
out_indices_buf = T.match_buffer(out_indices, (1, 2500), "int32")
|
||||
|
||||
with T.sblock("init"):
|
||||
vi = T.axis.S(1, 0)
|
||||
valid_count_buf[vi] = T.int32(0)
|
||||
for j in range(2500):
|
||||
with T.sblock("update"):
|
||||
vj = T.axis.S(2500, j)
|
||||
T.reads([data_buf[vi, vj, 6]])
|
||||
T.writes([valid_count_buf[vi], out_indices_buf[vi, vj], out_buf[vi, vj, 6]])
|
||||
if (data_buf[vi, vj, score_index] > score_threshold) and (
|
||||
(id_index < 0) or (data_buf[vi, vj, id_index] >= T.float32(0))
|
||||
):
|
||||
for k in T.serial(0, 6):
|
||||
out_buf[vi, valid_count_buf[vi], k] = data_buf[vi, vj, k]
|
||||
out_indices_buf[vi, valid_count_buf[vi]] = vj
|
||||
valid_count_buf[vi] = valid_count_buf[vi] + 1
|
||||
if vj >= valid_count_buf[vi]:
|
||||
for k in T.serial(0, 6):
|
||||
out_buf[vi, vj, k] = T.float32(-1)
|
||||
out_indices_buf[vi, vj] = T.int32(-1)
|
||||
|
||||
|
||||
def _check_get_valid_counts_with_numpy(f, dshape, score_threshold, id_index, score_index):
|
||||
dtype = "float32"
|
||||
ctx = tvm.cpu()
|
||||
batch_size, num_anchor, elem_length = dshape
|
||||
np_data = np.random.uniform(low=-2, high=2, size=dshape).astype(dtype)
|
||||
np_out1 = np.zeros(shape=(batch_size,), dtype="int32")
|
||||
np_out2 = np.zeros(shape=dshape).astype(dtype)
|
||||
np_out3 = np.zeros(shape=(batch_size, num_anchor), dtype="int32")
|
||||
for i in range(batch_size):
|
||||
np_out1[i] = 0
|
||||
inter_idx = 0
|
||||
for j in range(num_anchor):
|
||||
score = np_data[i, j, score_index]
|
||||
if score > score_threshold and (id_index < 0 or np_data[i, j, id_index] >= 0):
|
||||
for k in range(elem_length):
|
||||
np_out2[i, inter_idx, k] = np_data[i, j, k]
|
||||
np_out1[i] += 1
|
||||
np_out3[i, inter_idx] = j
|
||||
inter_idx += 1
|
||||
if j >= np_out1[i]:
|
||||
for k in range(elem_length):
|
||||
np_out2[i, j, k] = -1.0
|
||||
np_out3[i, j] = -1
|
||||
|
||||
in_data = tvm.runtime.tensor(np_data, ctx)
|
||||
out1 = tvm.runtime.tensor(np_out1, ctx)
|
||||
out2 = tvm.runtime.tensor(np_out2, ctx)
|
||||
out3 = tvm.runtime.tensor(np_out3, ctx)
|
||||
f(in_data, out1, out2, out3, score_threshold, id_index, score_index)
|
||||
tvm.testing.assert_allclose(out1.numpy(), np_out1, rtol=1e-5)
|
||||
tvm.testing.assert_allclose(out2.numpy(), np_out2, rtol=1e-5)
|
||||
tvm.testing.assert_allclose(out3.numpy(), np_out3, rtol=1e-5)
|
||||
print("test get_valid_counts end")
|
||||
|
||||
|
||||
def test_get_valid_counts_script_func():
|
||||
device = "llvm"
|
||||
# check lowering
|
||||
print(get_valid_counts.script())
|
||||
mod = tvm.ir.IRModule({"get_valid_counts": get_valid_counts})
|
||||
print(mod.script())
|
||||
# check building
|
||||
f = tvm.compile(mod["get_valid_counts"], target=device)
|
||||
_check_get_valid_counts_with_numpy(f, (1, 2500, 6), 0.0, 0, 1)
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def alloc_zero_dim_buffer(a: T.handle, b: T.handle) -> None:
|
||||
A = T.match_buffer(a, [], dtype="float32")
|
||||
B = T.match_buffer(b, [], dtype="float32")
|
||||
# body
|
||||
# tirx.with block("root")
|
||||
C = T.sblock_alloc_buffer([], dtype="float32")
|
||||
A[()] = T.float32(2)
|
||||
C[()] = A[()] + B[()]
|
||||
B[()] = C[()]
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def alloc_zero_dim_buffer_block(a: T.handle, b: T.handle) -> None:
|
||||
A = T.match_buffer(a, (), "float32")
|
||||
B = T.match_buffer(b, (), "float32")
|
||||
with T.sblock("root"):
|
||||
T.reads([])
|
||||
T.writes([])
|
||||
C = T.sblock_alloc_buffer((), "float32")
|
||||
A[()] = T.float32(2)
|
||||
C[()] = A[()] + B[()]
|
||||
B[()] = C[()]
|
||||
|
||||
|
||||
def _check_alloc_zero_dim_buffer(f):
|
||||
dtype = "float32"
|
||||
ctx = tvm.cpu()
|
||||
|
||||
np_data = np.zeros(shape=()).astype(dtype)
|
||||
np_out = np.zeros(shape=()).astype(dtype)
|
||||
tvm_data = tvm.runtime.tensor(np_data, ctx)
|
||||
tvm_out = tvm.runtime.tensor(np_out, ctx)
|
||||
|
||||
# np func exection
|
||||
np_inter = np.array(1)
|
||||
np_data[()] = 2.0
|
||||
np_inter[()] = np_data[()] + np_out[()]
|
||||
np_out[()] = np_inter[()]
|
||||
|
||||
# tvm func execution
|
||||
f(tvm_data, tvm_out)
|
||||
tvm.testing.assert_allclose(tvm_out.numpy(), np_out, rtol=1e-5)
|
||||
|
||||
|
||||
def test_alloc_zero_dim_buffer_round_trip():
|
||||
func = alloc_zero_dim_buffer
|
||||
func_with_block = alloc_zero_dim_buffer_block
|
||||
rt_func = tvm.script.from_source(func.script())
|
||||
rt_func_with_block = tvm.script.from_source(func_with_block.script())
|
||||
rt_mod = tvm.compile(rt_func, "llvm")
|
||||
rt_mod_with_block = tvm.compile(rt_func_with_block, "llvm")
|
||||
tvm.ir.assert_structural_equal(
|
||||
func.with_attr("global_symbol", "main"), func_with_block.with_attr("global_symbol", "main")
|
||||
)
|
||||
tvm.ir.assert_structural_equal(
|
||||
rt_func.with_attr("global_symbol", "main"),
|
||||
rt_func_with_block.with_attr("global_symbol", "main"),
|
||||
)
|
||||
_check_alloc_zero_dim_buffer(rt_mod)
|
||||
_check_alloc_zero_dim_buffer(rt_mod_with_block)
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def ceildiv_test(A: T.Buffer(16, "int32")):
|
||||
for i in range(16):
|
||||
A[i] = T.ceildiv(A[i], 4)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
||||
def test_ceildiv():
|
||||
f = tvm.compile(ceildiv_test, "llvm")
|
||||
a = tvm.runtime.tensor(np.arange(16).astype("int32"))
|
||||
f(a)
|
||||
ref = (np.arange(16) + 3) // 4
|
||||
tvm.testing.assert_allclose(a.numpy(), ref)
|
||||
|
||||
|
||||
try:
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def slice_op_test(
|
||||
A: T.Buffer((10,), "float32"), B: T.Buffer((10,), "float32"), C: T.Buffer((10,), "uint32")
|
||||
):
|
||||
B[0:5] = A[0:5] + B[0:5]
|
||||
B[0:5] = A[0:5] - B[0:5]
|
||||
B[0:5] = A[0:5] * B[0:5]
|
||||
B[0:5] = A[0:5] / B[0:5]
|
||||
C[0:5] = C[0:5] % T.broadcast(T.uint32(5), 5)
|
||||
B[0:5] = -B[0:5]
|
||||
C[0:5] = C[0:5] >> 4
|
||||
C[0:5] = C[0:5] << 4
|
||||
C[0:5] = C[0:5] << C[0:5]
|
||||
C[0:5] = C[0:5] >> C[0:5]
|
||||
T.evaluate(A[0:5] > B[0:5])
|
||||
T.evaluate(A[0:5] > 5)
|
||||
T.evaluate(A[0:5] >= B[0:5])
|
||||
T.evaluate(A[0:5] >= 5)
|
||||
T.evaluate(A[0:5] < B[0:5])
|
||||
T.evaluate(A[0:5] < 5)
|
||||
T.evaluate(A[0:5] <= B[0:5])
|
||||
T.evaluate(A[0:5] <= 5)
|
||||
T.evaluate(A[0:5] == B[0:5])
|
||||
T.evaluate(A[0:5] == 5)
|
||||
T.evaluate(A[0:5] != B[0:5])
|
||||
T.evaluate(A[0:5] != 5)
|
||||
T.evaluate((A[0:5] > 0) and (B[0:5] > 0))
|
||||
T.evaluate((A[0:5] > 0) or (B[0:5] > 0))
|
||||
T.evaluate((A[0:5] < 0) and (1 > 0))
|
||||
T.evaluate((A[0:5] > 0) or (1 > 0))
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def slice_op_test_ref(
|
||||
A: T.Buffer((10,), "float32"), B: T.Buffer((10,), "float32"), C: T.Buffer((10,), "uint32")
|
||||
):
|
||||
B[0:5] = A[0:5] + B[0:5]
|
||||
B[0:5] = A[0:5] - B[0:5]
|
||||
B[0:5] = A[0:5] * B[0:5]
|
||||
B[0:5] = A[0:5] / B[0:5]
|
||||
C[0:5] = C[0:5] % T.Broadcast(T.uint32(5), 5)
|
||||
B[0:5] = B[0:5] * T.Broadcast(T.float32(-1), 5)
|
||||
C[0:5] = T.shift_right(C[0:5], T.Broadcast(T.uint32(4), 5))
|
||||
C[0:5] = T.shift_left(C[0:5], T.Broadcast(T.uint32(4), 5))
|
||||
C[0:5] = T.shift_left(C[0:5], C[0:5])
|
||||
C[0:5] = T.shift_right(C[0:5], C[0:5])
|
||||
T.evaluate(A[0:5] > B[0:5])
|
||||
T.evaluate(A[0:5] > T.Broadcast(T.float32(5), 5))
|
||||
T.evaluate(A[0:5] >= B[0:5])
|
||||
T.evaluate(A[0:5] >= T.Broadcast(T.float32(5), 5))
|
||||
T.evaluate(A[0:5] < B[0:5])
|
||||
T.evaluate(A[0:5] < T.Broadcast(T.float32(5), 5))
|
||||
T.evaluate(A[0:5] <= B[0:5])
|
||||
T.evaluate(A[0:5] <= T.Broadcast(T.float32(5), 5))
|
||||
T.evaluate(A[0:5] == B[0:5])
|
||||
T.evaluate(A[0:5] == T.Broadcast(T.float32(5), 5))
|
||||
T.evaluate(A[0:5] != B[0:5])
|
||||
T.evaluate(A[0:5] != T.Broadcast(T.float32(5), 5))
|
||||
T.bitwise_and(A[0:5] > T.Broadcast(T.float32(0), 5), B[0:5] > T.Broadcast(T.float32(0), 5))
|
||||
T.bitwise_or(A[0:5] > T.Broadcast(T.float32(0), 5), B[0:5] > T.Broadcast(T.float32(0), 5))
|
||||
T.bitwise_and(A[0:5] < T.Broadcast(T.float32(0), 5), T.Broadcast(T.bool(1), 5))
|
||||
T.bitwise_or(A[0:5] > T.Broadcast(T.float32(0), 5), T.Broadcast(T.bool(1), 5))
|
||||
except tvm.error.DiagnosticError:
|
||||
slice_op_test = None
|
||||
slice_op_test_ref = None
|
||||
|
||||
|
||||
def test_slice_op():
|
||||
if slice_op_test is None:
|
||||
pytest.skip("slice arithmetic on BufferRegion is not defined")
|
||||
tvm.ir.assert_structural_equal(
|
||||
slice_op_test.with_attr("global_symbol", "main"),
|
||||
slice_op_test_ref.with_attr("global_symbol", "main"),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_get_valid_counts_script_func()
|
||||
test_alloc_zero_dim_buffer_round_trip()
|
||||
test_slice_op()
|
||||
@@ -0,0 +1,66 @@
|
||||
# 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: F401
|
||||
"""Unittests for tvm.script.parser.evaluator"""
|
||||
|
||||
import pytest
|
||||
|
||||
import tvm.testing
|
||||
from tvm.script.parser.core.diagnostics import Source
|
||||
from tvm.script.parser.core.evaluator import ExprEvaluator
|
||||
|
||||
|
||||
def _calc(expr, extra_vars=None):
|
||||
if extra_vars is None:
|
||||
extra_vars = {}
|
||||
source = Source(expr)
|
||||
mod_ast = source.as_ast()
|
||||
mod_body_ast = mod_ast.body
|
||||
expr_stmt_ast = mod_body_ast[0]
|
||||
expr_ast = expr_stmt_ast.value
|
||||
return ExprEvaluator.eval(None, extra_vars, expr_ast)
|
||||
|
||||
|
||||
def test_evaluator_basic():
|
||||
assert _calc("1, 3.14, True, 'str'") == (1, 3.14, True, "str")
|
||||
|
||||
|
||||
def test_evaluator_op():
|
||||
assert _calc("1 + 2, 1 - 2, 1 * 2, 1 / 2") == (3, -1, 2, 0.5)
|
||||
|
||||
|
||||
def test_evaluator_value_table():
|
||||
res = _calc("a + b, a - b, a * b, a / b", {"a": 1, "b": 2})
|
||||
a, b = 1, 2
|
||||
assert res == (a + b, a - b, a * b, a / b)
|
||||
|
||||
|
||||
def test_evaluator_func_call():
|
||||
def func(a, b):
|
||||
return a + b, a - b, a * b, a / b
|
||||
|
||||
assert _calc("func(1, 2)", {"func": func}) == func(1, 2)
|
||||
|
||||
|
||||
def test_evaluator_slice():
|
||||
res = _calc("a, a[1:], a[:5], a[1: 5], a[1: 5: 2]", {"a": [1, 2, 3, 4, 5, 6]})
|
||||
a = [1, 2, 3, 4, 5, 6]
|
||||
assert res == (a, a[1:], a[:5], a[1:5], a[1:5:2])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,39 @@
|
||||
# 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: F401
|
||||
"""Unittests for tvm.script.parser.ir"""
|
||||
|
||||
import inspect
|
||||
|
||||
import pytest
|
||||
|
||||
import tvm.testing
|
||||
from tvm.ir import IRModule
|
||||
from tvm.script.parser import ir_module
|
||||
|
||||
|
||||
def test_ir_base():
|
||||
@ir_module
|
||||
class BlankIRModule:
|
||||
pass
|
||||
|
||||
assert isinstance(BlankIRModule, IRModule) and len(BlankIRModule.functions.items()) == 0
|
||||
assert BlankIRModule.__name__ == "BlankIRModule"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,105 @@
|
||||
# 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: F401
|
||||
"""Unittests for tvm.script.parser.core"""
|
||||
|
||||
import inspect
|
||||
|
||||
import pytest
|
||||
|
||||
import tvm.testing
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.parser.core import doc_core as doc
|
||||
from tvm.script.parser.core.diagnostics import Source
|
||||
|
||||
|
||||
def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
B = T.match_buffer(b, [128, 128])
|
||||
C = T.match_buffer(c, [128, 128])
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
with T.sblock("update"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
|
||||
|
||||
|
||||
def test_source_base():
|
||||
source = Source(matmul)
|
||||
assert (
|
||||
source.source_name == inspect.getsourcefile(matmul)
|
||||
and source.start_line is not None
|
||||
and source.start_column == 0
|
||||
and source.source == inspect.getsource(matmul)
|
||||
and source.full_source == inspect.getsource(inspect.getmodule(matmul))
|
||||
)
|
||||
|
||||
|
||||
def test_source_ast():
|
||||
source = Source(matmul)
|
||||
mod = source.as_ast()
|
||||
assert isinstance(mod, doc.Module)
|
||||
func_def = mod.body[0]
|
||||
assert isinstance(func_def, doc.FunctionDef)
|
||||
assert func_def.name == "matmul"
|
||||
func_args = func_def.args
|
||||
assert (
|
||||
len(func_args.args) == 3
|
||||
and func_args.args[0].arg == "a"
|
||||
and func_args.args[1].arg == "b"
|
||||
and func_args.args[2].arg == "c"
|
||||
)
|
||||
func_body = func_def.body
|
||||
assert len(func_body) == 4
|
||||
func_assigns = func_body[:3]
|
||||
assert (
|
||||
isinstance(func_assigns[0], doc.Assign)
|
||||
and func_assigns[0].targets[0].id == "A"
|
||||
and isinstance(func_assigns[1], doc.Assign)
|
||||
and func_assigns[1].targets[0].id == "B"
|
||||
and isinstance(func_assigns[2], doc.Assign)
|
||||
and func_assigns[2].targets[0].id == "C"
|
||||
)
|
||||
func_for = func_body[3]
|
||||
assert (
|
||||
len(func_for.target.elts) == 3
|
||||
and func_for.target.elts[0].id == "i"
|
||||
and func_for.target.elts[1].id == "j"
|
||||
and func_for.target.elts[2].id == "k"
|
||||
)
|
||||
for_body = func_for.body
|
||||
assert len(for_body) == 1
|
||||
for_block = for_body[0]
|
||||
assert isinstance(for_block, doc.With) and len(for_block.body) == 2
|
||||
|
||||
|
||||
def test_nesting_parsing():
|
||||
class dummy:
|
||||
pass
|
||||
|
||||
for i in range(1):
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Module:
|
||||
@T.prim_func(s_tir=True)
|
||||
def impl(
|
||||
A: T.Buffer((12, 196, 64), "float32"),
|
||||
) -> None:
|
||||
T.evaluate(0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,709 @@
|
||||
# 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.
|
||||
"""Unittests for tvm.script.parser.tirx"""
|
||||
|
||||
import pytest
|
||||
import tvm_ffi
|
||||
|
||||
import tvm.testing
|
||||
from tvm import ir, tirx
|
||||
from tvm.script.parser import tirx as T
|
||||
|
||||
|
||||
def test_tir_buffer_proxy():
|
||||
buffer_0 = T.Buffer((128, 128), "float32")
|
||||
assert (
|
||||
isinstance(buffer_0, tirx.Buffer)
|
||||
and list(buffer_0.shape) == [128, 128]
|
||||
and buffer_0.dtype == ir.PrimType("float32")
|
||||
)
|
||||
|
||||
buffer_1 = T.Buffer((64, 64, 64), "int32")
|
||||
assert (
|
||||
isinstance(buffer_1, tirx.Buffer)
|
||||
and list(buffer_1.shape) == [64, 64, 64]
|
||||
and buffer_1.dtype == ir.PrimType("int32")
|
||||
)
|
||||
|
||||
|
||||
def test_tir_ptr_proxy():
|
||||
ptr_0 = T.handle("int32", "global")
|
||||
assert (
|
||||
isinstance(ptr_0, tirx.Var)
|
||||
and isinstance(ptr_0.ty, ir.PointerType)
|
||||
and ptr_0.ty.element_type == ir.PrimType("int32")
|
||||
and ptr_0.ty.storage_scope == "global"
|
||||
)
|
||||
|
||||
ptr_1 = T.handle("float32", "shared")
|
||||
assert (
|
||||
isinstance(ptr_1, tirx.Var)
|
||||
and isinstance(ptr_1.ty, ir.PointerType)
|
||||
and ptr_1.ty.element_type == ir.PrimType("float32")
|
||||
and ptr_1.ty.storage_scope == "shared"
|
||||
)
|
||||
|
||||
|
||||
def test_tir_func_name():
|
||||
@T.prim_func(s_tir=True)
|
||||
def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
B = T.match_buffer(b, [128, 128])
|
||||
C = T.match_buffer(c, [128, 128])
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
with T.sblock("update"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
|
||||
|
||||
assert matmul.__name__ == "matmul"
|
||||
assert matmul.attrs["global_symbol"] == "matmul"
|
||||
|
||||
|
||||
def test_tir_func_private_attrs():
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
T.func_attr({"attr": "value"})
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
B = T.match_buffer(b, [128, 128])
|
||||
C = T.match_buffer(c, [128, 128])
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
with T.sblock("update"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
|
||||
|
||||
assert "global_symbol" not in matmul.attrs
|
||||
|
||||
|
||||
def test_tir_func_private_manual_global_symbol_fail():
|
||||
with pytest.raises(tvm.error.DiagnosticError):
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
T.func_attr({"global_symbol": "matmul"})
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
B = T.match_buffer(b, [128, 128])
|
||||
C = T.match_buffer(c, [128, 128])
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
with T.sblock("update"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
|
||||
|
||||
# should not execute
|
||||
assert matmul.__name__ == "matmul"
|
||||
|
||||
|
||||
def test_tir_macro_decorator_signature():
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def evaluate0():
|
||||
T.evaluate(0)
|
||||
|
||||
# Ok, no parentheses
|
||||
@T.inline
|
||||
def func1():
|
||||
T.evaluate(0)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def use1():
|
||||
func1()
|
||||
|
||||
tvm.ir.assert_structural_equal(use1, evaluate0)
|
||||
|
||||
# Ok, empty parentheses
|
||||
@T.inline()
|
||||
def func2():
|
||||
T.evaluate(0)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def use2():
|
||||
func2()
|
||||
|
||||
tvm.ir.assert_structural_equal(use1, evaluate0)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
# Wrong: non-keyword argument
|
||||
@T.inline(True)
|
||||
def func3():
|
||||
T.evaluate()
|
||||
|
||||
|
||||
def test_tir_macro_signature():
|
||||
@T.inline
|
||||
def assign(i, *args, t1, **kwargs):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, args[0], args[1]])
|
||||
kwargs["t3"][vi, vj] = kwargs["t3"][vi, vj] + t1[vi, vk] * kwargs["t2"][vj, vk]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def matmul_w_macro(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
B = T.match_buffer(b, [128, 128])
|
||||
C = T.match_buffer(c, [128, 128])
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
with T.sblock("update"):
|
||||
assign(i, j, k, t1=A, t2=B, t3=C)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def matmul_no_macro(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
B = T.match_buffer(b, [128, 128])
|
||||
C = T.match_buffer(c, [128, 128])
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
with T.sblock("update"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
|
||||
|
||||
tvm.ir.assert_structural_equal(matmul_no_macro, matmul_w_macro)
|
||||
|
||||
|
||||
def test_tir_macro_hygienic():
|
||||
x_value = 128
|
||||
|
||||
@T.inline
|
||||
def static_capture(A, B):
|
||||
B[()] = A[x_value]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def use_hygienic(A: T.Buffer((1024,), "int32"), B: T.Buffer((), "int32")) -> None:
|
||||
for x_value in T.serial(10):
|
||||
static_capture(A, B)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected_hygienic(A: T.Buffer((1024,), "int32"), B: T.Buffer((), "int32")) -> None:
|
||||
for x_value in range(10):
|
||||
B[()] = A[128]
|
||||
|
||||
tvm.ir.assert_structural_equal(use_hygienic, expected_hygienic)
|
||||
|
||||
|
||||
def test_tir_inline_late_binding():
|
||||
"""Inline defined inside prim_func uses LEGB late binding:
|
||||
it sees the current value of variables from its enclosing scope at call time."""
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def use_late_binding(A: T.Buffer((1024,), "int32"), B: T.Buffer((), "int32")) -> None:
|
||||
for x_value in T.serial(10):
|
||||
|
||||
@T.inline
|
||||
def capture(A, B):
|
||||
B[()] = A[x_value]
|
||||
|
||||
capture(A, B)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(A: T.Buffer((1024,), "int32"), B: T.Buffer((), "int32")) -> None:
|
||||
for x_value in range(10):
|
||||
B[()] = A[x_value]
|
||||
|
||||
tvm.ir.assert_structural_equal(use_late_binding, expected)
|
||||
|
||||
|
||||
def test_tir_macro_in_class():
|
||||
class Object:
|
||||
def __init__(self, x: T.Buffer):
|
||||
self.local_x = T.sblock_alloc_buffer(x.shape, x.dtype)
|
||||
|
||||
@T.inline
|
||||
def load(self, x: T.Buffer):
|
||||
N, M = T.meta_var(self.local_x.shape)
|
||||
for i, j in T.grid(N, M):
|
||||
with T.sblock("update"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
self.local_x[vi, vj] = x[vi, vj]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def func_w_macro(a: T.handle):
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
o1 = T.meta_var(Object(A))
|
||||
o1.load(A)
|
||||
o2 = T.meta_var(Object(A))
|
||||
o2.load(o1.local_x)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def func_no_macro(a: T.handle):
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
local_a = T.sblock_alloc_buffer([128, 128])
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock("update"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
local_a[vi, vj] = A[vi, vj]
|
||||
local_b = T.sblock_alloc_buffer([128, 128])
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock("update"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
local_b[vi, vj] = local_a[vi, vj]
|
||||
|
||||
tvm.ir.assert_structural_equal(func_no_macro, func_w_macro)
|
||||
|
||||
|
||||
def test_tir_starred_expression():
|
||||
dims = (128, 128)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def starred(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, *dims], "int32")
|
||||
for i, j, k in T.grid(128, *dims):
|
||||
A[i, j, k] = T.int32(1)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def non_starred(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128, 128], "int32")
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
A[i, j, k] = T.int32(1)
|
||||
|
||||
tvm.ir.assert_structural_equal(starred, non_starred)
|
||||
|
||||
|
||||
def test_tir_starred_shape_expression():
|
||||
dims = (128, 128)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def starred(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, *dims], "int32")
|
||||
for i, j, k in T.grid(*A.shape):
|
||||
A[i, j, k] = T.int32(1)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def non_starred(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128, 128], "int32")
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
A[i, j, k] = T.int32(1)
|
||||
|
||||
tvm.ir.assert_structural_equal(starred, non_starred)
|
||||
|
||||
|
||||
def test_tir_dynamic_for_loop():
|
||||
dims = (128, 128)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def starred(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, *dims], "int32")
|
||||
for iters in T.grid(*A.shape):
|
||||
A[iters] = T.int32(1)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def non_starred(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128, 128], "int32")
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
A[i, j, k] = T.int32(1)
|
||||
|
||||
tvm.ir.assert_structural_equal(starred, non_starred)
|
||||
|
||||
|
||||
def test_tir_starred_for_loop():
|
||||
dims = (128, 128)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def starred(a: T.handle, b: T.handle):
|
||||
A = T.match_buffer(a, [*dims, 128], "int32")
|
||||
B = T.match_buffer(b, dims, "int32")
|
||||
for *spatial, reduction in T.grid(*A.shape):
|
||||
with T.sblock("reduce"):
|
||||
with T.init():
|
||||
B[spatial] = T.int32(0)
|
||||
B[spatial] = B[spatial] + A[(*spatial, reduction)]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def non_starred(a: T.handle, b: T.handle):
|
||||
A = T.match_buffer(a, [128, 128, 128], "int32")
|
||||
B = T.match_buffer(b, [128, 128], "int32")
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
with T.sblock("reduce"):
|
||||
with T.init():
|
||||
B[i, j] = T.int32(0)
|
||||
B[i, j] = B[i, j] + A[i, j, k]
|
||||
|
||||
tvm.ir.assert_structural_equal(starred, non_starred)
|
||||
|
||||
|
||||
def test_tir_loop_steps():
|
||||
N = T.Var("N", "int32")
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def loop_with_steps(
|
||||
A: T.Buffer((N,)), B: T.Buffer((N,)), C: T.Buffer((N,)), tid: T.int32, v: T.int32
|
||||
):
|
||||
for i in T.serial(tid, N, step=2):
|
||||
C[i] = A[i] + B[i]
|
||||
for i in T.unroll(tid, N, step=3):
|
||||
C[i] = A[i] + B[i]
|
||||
for i in T.vectorized(tid, N, step=4):
|
||||
C[i] = A[i] + B[i]
|
||||
for i in T.parallel(tid, N, step=5):
|
||||
C[i] = A[i] + B[i]
|
||||
for i in T.serial(tid, N, step=v):
|
||||
C[i] = A[i] + B[i]
|
||||
|
||||
stmts = loop_with_steps.body.seq
|
||||
assert stmts[0].step == 2
|
||||
assert stmts[1].step == 3
|
||||
assert stmts[2].step == 4
|
||||
assert stmts[3].step == 5
|
||||
assert stmts[4].step.name == "v"
|
||||
|
||||
|
||||
def test_tir_empty_tuple_index():
|
||||
@T.inline
|
||||
def bar(val):
|
||||
T.evaluate(val)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def func_with_empty_tuple(A: T.Buffer((), "int32"), B: T.Buffer((), "int32")):
|
||||
bar(val=A[()])
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(A: T.Buffer((), "int32"), B: T.Buffer((), "int32")):
|
||||
T.evaluate(A[()])
|
||||
|
||||
tvm.ir.assert_structural_equal(func_with_empty_tuple, expected)
|
||||
|
||||
|
||||
def test_tir_builtin_expression():
|
||||
dims = (128, 128)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def with_builtin(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, [len(dims), *dims], "int32")
|
||||
for i, j, k in T.grid(*A.shape):
|
||||
A[i, j, k] = T.int32(1 + len(A.shape))
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def evaluated(A: T.Buffer((2, 128, 128), "int32")):
|
||||
for i, j, k in T.grid(2, 128, 128):
|
||||
A[i, j, k] = 4
|
||||
|
||||
tvm.ir.assert_structural_equal(with_builtin, evaluated)
|
||||
|
||||
|
||||
def test_thread_binding_dtype():
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def func(A: T.Buffer((128, 128)), B: T.Buffer((128, 128))):
|
||||
for i in T.thread_binding(T.int64(128), "threadIdx.x"):
|
||||
for j in T.thread_binding(128, "threadIdx.y"):
|
||||
B[i, j] = A[i, j]
|
||||
|
||||
loop_i = func.body
|
||||
loop_j = loop_i.body
|
||||
assert loop_i.loop_var.ty.dtype == "int64"
|
||||
assert loop_i.thread_binding.var.ty.dtype == "int64"
|
||||
assert loop_j.loop_var.ty.dtype == "int32"
|
||||
assert loop_j.thread_binding.var.ty.dtype == "int32"
|
||||
|
||||
|
||||
def test_inferred_ty_with_prim_args():
|
||||
"""A PrimFunc may have inferred Type"""
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func(M: T.int32, N: T.int32) -> T.int32:
|
||||
T.ret(M * N)
|
||||
|
||||
expected = tvm.relax.FuncType(
|
||||
[
|
||||
tvm.ir.PrimType("int32"),
|
||||
tvm.ir.PrimType("int32"),
|
||||
],
|
||||
tvm.ir.PrimType("int32"),
|
||||
purity=True,
|
||||
)
|
||||
tvm.ir.assert_structural_equal(func.ty, expected)
|
||||
|
||||
|
||||
def test_inferred_ty_with_buffer_args():
|
||||
"""PrimFunc buffer arguments are inferred as R.Tensor"""
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func(A: T.Buffer([16, 16], "float32"), B: T.Buffer([256], "int32")) -> T.float32:
|
||||
T.ret(T.float32(42.0))
|
||||
|
||||
expected = tvm.relax.FuncType(
|
||||
[
|
||||
tvm.relax.TensorType([16, 16], "float32"),
|
||||
tvm.relax.TensorType([256], "int32"),
|
||||
],
|
||||
tvm.ir.PrimType("float32"),
|
||||
purity=True,
|
||||
)
|
||||
tvm.ir.assert_structural_equal(func.ty, expected)
|
||||
|
||||
|
||||
def test_inferred_ty_with_internal_allocation():
|
||||
"""A pure function may still write to internal allocations.
|
||||
|
||||
Whether a function writes to internal allocations is not a visible
|
||||
effect, and does not impact the purity of a function.
|
||||
"""
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func(A: T.Buffer([16, 16], "float32")) -> T.float32:
|
||||
Sum = T.decl_buffer([], "float32")
|
||||
Sum[()] = 0.0
|
||||
for i, j in T.grid(16, 16):
|
||||
Sum[()] = Sum[()] + A[i, j]
|
||||
|
||||
T.ret(Sum[()])
|
||||
|
||||
expected = tvm.relax.FuncType(
|
||||
[
|
||||
tvm.relax.TensorType([16, 16], "float32"),
|
||||
],
|
||||
tvm.ir.PrimType("float32"),
|
||||
purity=True,
|
||||
)
|
||||
tvm.ir.assert_structural_equal(func.ty, expected)
|
||||
|
||||
|
||||
def test_inferred_ty_with_output_buffer():
|
||||
"""A pure function may not write to an argument buffer
|
||||
|
||||
If an argument buffer is written to, the function must be impure.
|
||||
"""
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func(A: T.Buffer(16, "float32"), B: T.Buffer(16, "float32")):
|
||||
for i in range(16):
|
||||
B[i] = A[i]
|
||||
|
||||
expected = tvm.relax.FuncType(
|
||||
[
|
||||
tvm.relax.TensorType([16], "float32"),
|
||||
tvm.relax.TensorType([16], "float32"),
|
||||
],
|
||||
tvm.relax.TupleType([]),
|
||||
purity=False,
|
||||
)
|
||||
tvm.ir.assert_structural_equal(func.ty, expected)
|
||||
|
||||
|
||||
def test_inferred_ty_with_dynamic_buffer():
|
||||
"""The inferred Type may contain dynamic shapes"""
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func(a_handle: T.handle, b_handle: T.handle):
|
||||
M = T.int64()
|
||||
N = T.int64()
|
||||
A = T.match_buffer(a_handle, [M, N], "float32")
|
||||
B = T.match_buffer(b_handle, [M * N], "float32")
|
||||
for i, j in T.grid(M, N):
|
||||
B[i * N + j] = A[i, j]
|
||||
|
||||
M = tvm.tirx.Var("M", "int64")
|
||||
N = tvm.tirx.Var("N", "int64")
|
||||
expected = tvm.relax.FuncType(
|
||||
[
|
||||
tvm.relax.TensorType([M, N], "float32"),
|
||||
tvm.relax.TensorType([M * N], "float32"),
|
||||
],
|
||||
tvm.relax.TupleType([]),
|
||||
purity=False,
|
||||
)
|
||||
tvm.ir.assert_structural_equal(func.ty, expected)
|
||||
|
||||
|
||||
def test_reinterpret_nop():
|
||||
"""Test builtin reinterpret op"""
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func(A: T.Buffer((32,), "float32"), B: T.Buffer((32,), "float32")) -> None:
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
for i in T.serial(0, 32):
|
||||
with T.sblock():
|
||||
vi = T.axis.remap("S", [i])
|
||||
B[vi] = T.reinterpret("float32", A[vi])
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def expected(A: T.Buffer((32,), "float32"), B: T.Buffer((32,), "float32")) -> None:
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
for i in T.serial(0, 32):
|
||||
with T.sblock():
|
||||
vi = T.axis.remap("S", [i])
|
||||
B[vi] = A[vi]
|
||||
|
||||
tvm.ir.assert_structural_equal(func, expected)
|
||||
|
||||
|
||||
def test_launch_thread_i64():
|
||||
"""Test launching thread with int64"""
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func() -> None:
|
||||
blockIdx_x = T.launch_thread("blockIdx.x", T.int64(1))
|
||||
if blockIdx_x == T.int64(0):
|
||||
T.evaluate(T.int64(0))
|
||||
else:
|
||||
T.evaluate(T.int64(1))
|
||||
|
||||
assert func.body.node.dom.min.ty.dtype == "int64"
|
||||
assert func.body.node.dom.extent.ty.dtype == "int64"
|
||||
|
||||
|
||||
def test_deterministic_branch():
|
||||
"""Test deterministic branch"""
|
||||
|
||||
def create_func(predicate: bool):
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def func() -> None:
|
||||
if predicate:
|
||||
T.evaluate(0)
|
||||
else:
|
||||
T.evaluate(1)
|
||||
|
||||
return func
|
||||
|
||||
def create_expected(value):
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected() -> None:
|
||||
T.evaluate(value)
|
||||
|
||||
return expected
|
||||
|
||||
tvm.ir.assert_structural_equal(create_func(True), create_expected(0))
|
||||
tvm.ir.assert_structural_equal(create_func(False), create_expected(1))
|
||||
|
||||
|
||||
def test_block_annotation_merge():
|
||||
def _to_dict(anno: tvm_ffi.container.Map):
|
||||
result = {}
|
||||
for k, v in anno.items():
|
||||
result[k] = _to_dict(v) if isinstance(v, tvm_ffi.container.Map) else v
|
||||
return result
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func0():
|
||||
with T.sblock():
|
||||
T.sblock_attr({"key1": "block1"})
|
||||
T.sblock_attr({"key2": "block2"})
|
||||
T.evaluate(0)
|
||||
|
||||
assert _to_dict(func0.body.block.annotations) == {"key1": "block1", "key2": "block2"}
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func1():
|
||||
with T.sblock():
|
||||
T.sblock_attr({"key": {"key1": "block1"}})
|
||||
T.sblock_attr({"key": {"key2": "block2"}})
|
||||
T.evaluate(0)
|
||||
|
||||
assert _to_dict(func1.body.block.annotations) == {"key": {"key1": "block1", "key2": "block2"}}
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func2():
|
||||
with T.sblock():
|
||||
T.sblock_attr({"key1": "block1"})
|
||||
T.sblock_attr({"key1": "block1"})
|
||||
T.evaluate(0)
|
||||
|
||||
assert _to_dict(func2.body.block.annotations) == {"key1": "block1"}
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func3():
|
||||
with T.sblock():
|
||||
T.sblock_attr({"key1": "block1"})
|
||||
T.sblock_attr({"key1": "block2"})
|
||||
T.evaluate(0)
|
||||
|
||||
|
||||
def test_alloc_inside_block():
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def func() -> None:
|
||||
with T.sblock():
|
||||
A = T.sblock_alloc_buffer([10], "float32")
|
||||
for i in T.serial(0, 10):
|
||||
B = T.sblock_alloc_buffer([10], "float32")
|
||||
for j in T.serial(0, 10):
|
||||
B[j] = T.float32(j)
|
||||
A[i] += B[j]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected() -> None:
|
||||
with T.sblock():
|
||||
A = T.sblock_alloc_buffer([10], "float32")
|
||||
B = T.sblock_alloc_buffer([10], "float32")
|
||||
for i, j in T.grid(10, 10):
|
||||
B[j] = T.float32(j)
|
||||
A[i] += B[j]
|
||||
|
||||
tvm.ir.assert_structural_equal(func, expected)
|
||||
|
||||
|
||||
def test_tir_macro_block_name_suffix():
|
||||
@T.inline
|
||||
def operation(A, idx):
|
||||
with T.sblock("op"):
|
||||
v = T.axis.remap("S", [idx])
|
||||
A[v] = A[v] * T.float32(2)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def func_w_macro(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, [10])
|
||||
for i in T.serial(0, 10):
|
||||
operation(A, i)
|
||||
operation(A, i)
|
||||
operation(A, i)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, [10])
|
||||
for i in T.serial(0, 10):
|
||||
with T.sblock("op"):
|
||||
v = T.axis.remap("S", [i])
|
||||
A[v] = A[v] * T.float32(2)
|
||||
with T.sblock("op_1"):
|
||||
v = T.axis.remap("S", [i])
|
||||
A[v] = A[v] * T.float32(2)
|
||||
with T.sblock("op_2"):
|
||||
v = T.axis.remap("S", [i])
|
||||
A[v] = A[v] * T.float32(2)
|
||||
|
||||
tvm.ir.assert_structural_equal(func_w_macro, expected)
|
||||
|
||||
|
||||
def test_ifexp():
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def func(A: T.buffer((128, 128), "float32")):
|
||||
for i, j in T.grid(128, 128):
|
||||
A[i, j] = i if i < j else j
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(A: T.buffer((128, 128), "float32")):
|
||||
for i, j in T.grid(128, 128):
|
||||
A[i, j] = T.if_then_else(i < j, i, j)
|
||||
|
||||
tvm.ir.assert_structural_equal(func, expected)
|
||||
|
||||
|
||||
def test_sequence_compare():
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def tir_func(A: T.Buffer((128, 128), "float32")):
|
||||
for i, j in T.grid(128, 128):
|
||||
if 0 < i < 128 and 0 < j < 128:
|
||||
A[i, j] = 1
|
||||
else:
|
||||
A[i, j] = 0
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(A: T.buffer((128, 128), "float32")):
|
||||
for i, j in T.grid(128, 128):
|
||||
if (0 < i and i < 128) and (0 < j and j < 128):
|
||||
A[i, j] = 1
|
||||
else:
|
||||
A[i, j] = 0
|
||||
|
||||
tvm.ir.assert_structural_equal(tir_func, expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,158 @@
|
||||
# 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.
|
||||
"""Test TVMScript with PEP 563 (from __future__ import annotations).
|
||||
|
||||
IMPORTANT: The `from __future__ import annotations` import below is the
|
||||
test condition itself, because we need to test compatibility with it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import tirx as T
|
||||
|
||||
|
||||
def _normalize(func):
|
||||
"""Strip the global_symbol so function names do not affect structural equality."""
|
||||
return func.with_attr("global_symbol", "")
|
||||
|
||||
|
||||
def test_prim_func_closure_shape():
|
||||
"""Closure variable used in Buffer shape annotation."""
|
||||
|
||||
def f(M=16):
|
||||
@T.prim_func(s_tir=True)
|
||||
def func(A: T.Buffer((M,), "float32")):
|
||||
T.evaluate(0)
|
||||
|
||||
return func
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def expected_16(A: T.Buffer((16,), "float32")):
|
||||
T.evaluate(0)
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def expected_32(A: T.Buffer((32,), "float32")):
|
||||
T.evaluate(0)
|
||||
|
||||
tvm.ir.assert_structural_equal(_normalize(f(16)), _normalize(expected_16))
|
||||
tvm.ir.assert_structural_equal(_normalize(f(32)), _normalize(expected_32))
|
||||
|
||||
|
||||
def test_prim_func_closure_dtype():
|
||||
"""Closure variable used as Buffer dtype."""
|
||||
|
||||
def f(dtype="float32"):
|
||||
@T.prim_func(s_tir=True)
|
||||
def func(A: T.Buffer((16,), dtype)):
|
||||
T.evaluate(0)
|
||||
|
||||
return func
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def expected_f32(A: T.Buffer((16,), "float32")):
|
||||
T.evaluate(0)
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def expected_f16(A: T.Buffer((16,), "float16")):
|
||||
T.evaluate(0)
|
||||
|
||||
tvm.ir.assert_structural_equal(_normalize(f("float32")), _normalize(expected_f32))
|
||||
tvm.ir.assert_structural_equal(_normalize(f("float16")), _normalize(expected_f16))
|
||||
|
||||
|
||||
def test_prim_func_nested_closure():
|
||||
"""Variables from enclosing scope active on the call stack (grandparent frame fallback).
|
||||
|
||||
With PEP 563, closure-only variables are missing from __closure__ unless they
|
||||
appear in the function body. The ChainMap fallback walks the live call stack,
|
||||
so this works when the enclosing frames are still active (outer calls middle
|
||||
which applies the decorator, keeping outer's frame alive on the stack).
|
||||
"""
|
||||
|
||||
def outer(M=16):
|
||||
def middle(N=8):
|
||||
@T.prim_func(s_tir=True)
|
||||
def func(A: T.Buffer((M, N), "float32")):
|
||||
T.evaluate(0)
|
||||
|
||||
return func
|
||||
|
||||
return middle()
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def expected_16_8(A: T.Buffer((16, 8), "float32")):
|
||||
T.evaluate(0)
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def expected_32_8(A: T.Buffer((32, 8), "float32")):
|
||||
T.evaluate(0)
|
||||
|
||||
tvm.ir.assert_structural_equal(_normalize(outer(16)), _normalize(expected_16_8))
|
||||
tvm.ir.assert_structural_equal(_normalize(outer(32)), _normalize(expected_32_8))
|
||||
|
||||
|
||||
def test_ir_module_closure():
|
||||
"""Closure variable in @I.ir_module class method."""
|
||||
|
||||
def f(M=16):
|
||||
@I.ir_module
|
||||
class Mod:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((M,), "float32")):
|
||||
T.evaluate(0)
|
||||
|
||||
return Mod
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def expected_16(A: T.Buffer((16,), "float32")):
|
||||
T.evaluate(0)
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def expected_32(A: T.Buffer((32,), "float32")):
|
||||
T.evaluate(0)
|
||||
|
||||
tvm.ir.assert_structural_equal(_normalize(f(16)["main"]), _normalize(expected_16))
|
||||
tvm.ir.assert_structural_equal(_normalize(f(32)["main"]), _normalize(expected_32))
|
||||
|
||||
|
||||
def test_mixed_closure_usage():
|
||||
"""Closure var used in both annotation AND body -- regression check."""
|
||||
|
||||
def f(M=16):
|
||||
@T.prim_func(s_tir=True)
|
||||
def func(A: T.Buffer((M,), "float32")):
|
||||
T.evaluate(M)
|
||||
|
||||
return func
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def expected_16(A: T.Buffer((16,), "float32")):
|
||||
T.evaluate(16)
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def expected_32(A: T.Buffer((32,), "float32")):
|
||||
T.evaluate(32)
|
||||
|
||||
tvm.ir.assert_structural_equal(_normalize(f(16)), _normalize(expected_16))
|
||||
tvm.ir.assert_structural_equal(_normalize(f(32)), _normalize(expected_32))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,118 @@
|
||||
# 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: F401
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
from tvm_ffi.access_path import AccessPath
|
||||
|
||||
from tvm.script import tirx as T
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def _func():
|
||||
T.evaluate(-1)
|
||||
T.evaluate(1)
|
||||
T.evaluate(2)
|
||||
T.evaluate(3)
|
||||
T.evaluate(4)
|
||||
T.evaluate(5)
|
||||
T.evaluate(6)
|
||||
T.evaluate(7)
|
||||
|
||||
|
||||
def test_annotation_multi_access_paths():
|
||||
result = _func.with_attr("global_symbol", "main").script(
|
||||
path_to_annotate={
|
||||
AccessPath.root().attr("body").attr("seq").array_item(1): "annotation 1",
|
||||
AccessPath.root().attr("body").attr("seq").array_item(3): "annotation 3",
|
||||
AccessPath.root().attr("body").attr("seq").array_item(5): "annotation 5",
|
||||
AccessPath.root().attr("body").attr("seq").array_item(7): "annotation 7",
|
||||
}
|
||||
)
|
||||
assert (
|
||||
result
|
||||
== """# from tvm.script import tirx as T
|
||||
# from tvm.tirx.layout import Axis
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def main():
|
||||
T.evaluate(-1)
|
||||
T.evaluate(1) # annotation 1
|
||||
T.evaluate(2)
|
||||
T.evaluate(3) # annotation 3
|
||||
T.evaluate(4)
|
||||
T.evaluate(5) # annotation 5
|
||||
T.evaluate(6)
|
||||
T.evaluate(7) # annotation 7"""
|
||||
)
|
||||
|
||||
|
||||
def test_annotate_from_multi_obj():
|
||||
result = _func.with_attr("global_symbol", "main").script(
|
||||
obj_to_annotate={
|
||||
_func.body.seq[1]: "annotation 1",
|
||||
_func.body.seq[3]: "annotation 3",
|
||||
_func.body.seq[5]: "annotation 5",
|
||||
_func.body.seq[7]: "annotation 7",
|
||||
}
|
||||
)
|
||||
assert (
|
||||
result
|
||||
== """# from tvm.script import tirx as T
|
||||
# from tvm.tirx.layout import Axis
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def main():
|
||||
T.evaluate(-1)
|
||||
T.evaluate(1) # annotation 1
|
||||
T.evaluate(2)
|
||||
T.evaluate(3) # annotation 3
|
||||
T.evaluate(4)
|
||||
T.evaluate(5) # annotation 5
|
||||
T.evaluate(6)
|
||||
T.evaluate(7) # annotation 7"""
|
||||
)
|
||||
|
||||
|
||||
def test_disable_concise_scoping_when_scope_annotated():
|
||||
@T.prim_func(s_tir=True)
|
||||
def _func():
|
||||
x = 1
|
||||
y = x + 1
|
||||
T.evaluate(y - 1)
|
||||
|
||||
# In fork, each bare `x = expr` lowers to AllocBuffer + BufferStore (local_scalar);
|
||||
# the printer fuses each pair into a single `y: T.int32 = x + 1` line. Annotate the
|
||||
# AllocBuffer that originates this fused line.
|
||||
result = _func.with_attr("global_symbol", "main").script(
|
||||
obj_to_annotate={
|
||||
_func.body.seq[2]: "annotation 1",
|
||||
}
|
||||
)
|
||||
assert (
|
||||
result
|
||||
== """# from tvm.script import tirx as T
|
||||
# from tvm.tirx.layout import Axis
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def main():
|
||||
x: T.int32 = 1
|
||||
y: T.int32 = x + 1 # annotation 1
|
||||
T.evaluate(y - 1)"""
|
||||
)
|
||||
@@ -0,0 +1,563 @@
|
||||
# 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: F841
|
||||
"""
|
||||
In this test file, we want to make sure the Python code can construct
|
||||
Doc objects, then access and modify their attributes correctly.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from tvm_ffi.access_path import AccessPath
|
||||
|
||||
import tvm
|
||||
from tvm.script.printer.doc import (
|
||||
AssertDoc,
|
||||
AssignDoc,
|
||||
AttrAccessDoc,
|
||||
CallDoc,
|
||||
ClassDoc,
|
||||
CommentDoc,
|
||||
DictDoc,
|
||||
DocStringDoc,
|
||||
ExprStmtDoc,
|
||||
ForDoc,
|
||||
FunctionDoc,
|
||||
IdDoc,
|
||||
IfDoc,
|
||||
IndexDoc,
|
||||
LambdaDoc,
|
||||
ListDoc,
|
||||
LiteralDoc,
|
||||
OperationDoc,
|
||||
OperationKind,
|
||||
ReturnDoc,
|
||||
ScopeDoc,
|
||||
SliceDoc,
|
||||
StmtBlockDoc,
|
||||
TupleDoc,
|
||||
WhileDoc,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value",
|
||||
[None, "test", 0, 1, -2, 0.0, 1.5, -1.3, True, False],
|
||||
)
|
||||
def test_literal_doc_construction(value):
|
||||
doc = LiteralDoc(value)
|
||||
|
||||
if isinstance(value, float):
|
||||
# FloatImm cannot be compared with Python's float directly
|
||||
assert float(doc.value) == pytest.approx(value)
|
||||
else:
|
||||
assert doc.value == value
|
||||
|
||||
|
||||
def test_id_doc():
|
||||
doc = IdDoc("name")
|
||||
|
||||
assert doc.name == "name"
|
||||
|
||||
|
||||
def test_attr_access_doc():
|
||||
target = IdDoc("x")
|
||||
|
||||
doc = AttrAccessDoc(target, "attribute")
|
||||
|
||||
assert doc.value == target
|
||||
assert doc.name == "attribute"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"indices",
|
||||
[
|
||||
[],
|
||||
[LiteralDoc(1)],
|
||||
[LiteralDoc(2), IdDoc("x")],
|
||||
[SliceDoc(LiteralDoc(1), LiteralDoc(2))],
|
||||
[SliceDoc(LiteralDoc(1)), IdDoc("y")],
|
||||
],
|
||||
)
|
||||
def test_index_doc(indices):
|
||||
target = IdDoc("x")
|
||||
|
||||
doc = IndexDoc(target, indices)
|
||||
|
||||
assert doc.value == target
|
||||
assert list(doc.indices) == indices
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"args, kwargs",
|
||||
[
|
||||
([], {}),
|
||||
([LiteralDoc("arg")], {}),
|
||||
([LiteralDoc("arg"), IdDoc("x")], {}),
|
||||
([], {"x": LiteralDoc("x")}),
|
||||
([], {"x": LiteralDoc("x"), "y": LiteralDoc("y")}),
|
||||
([LiteralDoc("arg")], {"x": LiteralDoc("x"), "y": LiteralDoc("y")}),
|
||||
([LiteralDoc("arg"), IdDoc("x")], {"x": LiteralDoc("x"), "y": LiteralDoc("y")}),
|
||||
],
|
||||
)
|
||||
def test_call_doc(args, kwargs):
|
||||
target = IdDoc("x")
|
||||
|
||||
doc = CallDoc(target, *args, **kwargs)
|
||||
|
||||
assert doc.callee == target
|
||||
assert list(doc.args) == args
|
||||
assert dict(zip(doc.kwargs_keys, doc.kwargs_values)) == kwargs
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"operands",
|
||||
[
|
||||
[],
|
||||
[LiteralDoc(1)],
|
||||
[LiteralDoc(2), IdDoc("x")],
|
||||
[LiteralDoc(2), IdDoc("x"), LiteralDoc("y")],
|
||||
],
|
||||
)
|
||||
def test_operation_doc(operands):
|
||||
# Here we just test the contructor and attr visitor of OperationDoc
|
||||
# so the choice of OperationKind doesn't matter
|
||||
operator = OperationKind.Add
|
||||
|
||||
doc = OperationDoc(OperationKind.Add, operands)
|
||||
|
||||
assert doc.kind == operator
|
||||
assert list(doc.operands) == operands
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"args",
|
||||
[
|
||||
[],
|
||||
[IdDoc("x")],
|
||||
[IdDoc("x"), IdDoc("y")],
|
||||
],
|
||||
)
|
||||
def test_lambda_doc(args):
|
||||
body = LiteralDoc(1)
|
||||
|
||||
doc = LambdaDoc(args, body)
|
||||
|
||||
assert doc.body == body
|
||||
assert list(doc.args) == args
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"elements",
|
||||
[
|
||||
[],
|
||||
[IdDoc("x")],
|
||||
[IdDoc("x"), IdDoc("y")],
|
||||
],
|
||||
)
|
||||
def test_tuple_doc(elements):
|
||||
doc = TupleDoc(elements)
|
||||
|
||||
assert list(doc.elements) == elements
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"elements",
|
||||
[
|
||||
[],
|
||||
[IdDoc("x")],
|
||||
[IdDoc("x"), IdDoc("y")],
|
||||
],
|
||||
)
|
||||
def test_list_doc(elements):
|
||||
doc = ListDoc(elements)
|
||||
|
||||
assert list(doc.elements) == elements
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"content",
|
||||
[
|
||||
{},
|
||||
{LiteralDoc("k"): IdDoc("v")},
|
||||
{LiteralDoc("k"): IdDoc("v"), LiteralDoc("k2"): IdDoc("v2")},
|
||||
],
|
||||
)
|
||||
def test_dict_doc(content):
|
||||
doc = DictDoc(content)
|
||||
|
||||
assert dict(zip(doc.keys, doc.values)) == content
|
||||
|
||||
|
||||
@pytest.mark.parametrize("start", [LiteralDoc(1), None])
|
||||
@pytest.mark.parametrize("stop", [LiteralDoc(2), None])
|
||||
@pytest.mark.parametrize("step", [LiteralDoc(3), None])
|
||||
def test_slice_doc(start, stop, step):
|
||||
doc = SliceDoc(start, stop)
|
||||
|
||||
assert doc.start == start
|
||||
assert doc.stop == stop
|
||||
|
||||
|
||||
def test_expr_doc_attr_access():
|
||||
target = IdDoc("x")
|
||||
attr = "test"
|
||||
|
||||
doc = target.attr(attr)
|
||||
|
||||
assert doc.value == target
|
||||
assert doc.name == attr
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"indices",
|
||||
[
|
||||
(),
|
||||
LiteralDoc(1),
|
||||
SliceDoc(LiteralDoc(1), LiteralDoc(2)),
|
||||
(LiteralDoc(1),),
|
||||
(LiteralDoc(2), IdDoc("x")),
|
||||
(SliceDoc(LiteralDoc(1), LiteralDoc(2)),),
|
||||
(SliceDoc(LiteralDoc(1)), IdDoc("y")),
|
||||
],
|
||||
)
|
||||
def test_expr_doc_get_item(indices):
|
||||
target = IdDoc("x")
|
||||
|
||||
doc = target[indices]
|
||||
|
||||
assert doc.value == target
|
||||
if not isinstance(indices, tuple):
|
||||
indices = (indices,)
|
||||
assert tuple(doc.indices) == indices
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"args, kwargs",
|
||||
[
|
||||
([], {}),
|
||||
([LiteralDoc("arg")], {}),
|
||||
([LiteralDoc("arg"), IdDoc("x")], {}),
|
||||
([], {"x": LiteralDoc("x")}),
|
||||
([], {"x": LiteralDoc("x"), "y": LiteralDoc("y")}),
|
||||
([LiteralDoc("arg")], {"x": LiteralDoc("x"), "y": LiteralDoc("y")}),
|
||||
([LiteralDoc("arg"), IdDoc("x")], {"x": LiteralDoc("x"), "y": LiteralDoc("y")}),
|
||||
],
|
||||
)
|
||||
def test_expr_doc_call_with(args, kwargs):
|
||||
target = IdDoc("x")
|
||||
|
||||
doc = target.call(*args, **kwargs)
|
||||
|
||||
assert doc.callee == target
|
||||
assert list(doc.args) == args
|
||||
assert dict(zip(doc.kwargs_keys, doc.kwargs_values)) == kwargs
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"stmts",
|
||||
[
|
||||
[],
|
||||
[ExprStmtDoc(IdDoc("x"))],
|
||||
[ExprStmtDoc(IdDoc("x")), ExprStmtDoc(IdDoc("y"))],
|
||||
],
|
||||
)
|
||||
def test_stmt_block_doc(stmts):
|
||||
doc = StmtBlockDoc(stmts)
|
||||
|
||||
assert list(doc.stmts) == stmts
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"lhs, rhs, annotation",
|
||||
[
|
||||
(IdDoc("x"), IdDoc("y"), None),
|
||||
(IdDoc("x"), None, IdDoc("int")),
|
||||
(IdDoc("x"), IdDoc("y"), IdDoc("int")),
|
||||
],
|
||||
)
|
||||
def test_assign_doc(lhs, rhs, annotation):
|
||||
doc = AssignDoc(lhs, rhs, annotation)
|
||||
|
||||
assert doc.lhs == lhs
|
||||
assert doc.rhs == rhs
|
||||
assert doc.annotation == annotation
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"lhs, rhs, annotation",
|
||||
[
|
||||
(IdDoc("x"), None, None),
|
||||
(TupleDoc([IdDoc("x"), IdDoc("y")]), None, IdDoc("int")),
|
||||
(TupleDoc([IdDoc("x"), IdDoc("y")]), IdDoc("u"), IdDoc("int")),
|
||||
],
|
||||
)
|
||||
def test_invalid_assign_doc(lhs, rhs, annotation):
|
||||
with pytest.raises(ValueError) as e:
|
||||
AssignDoc(lhs, rhs, annotation)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"else_branch",
|
||||
[
|
||||
[],
|
||||
[ExprStmtDoc(IdDoc("x"))],
|
||||
[ExprStmtDoc(IdDoc("x")), ExprStmtDoc(IdDoc("y"))],
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"then_branch",
|
||||
[
|
||||
[],
|
||||
[ExprStmtDoc(IdDoc("x"))],
|
||||
[ExprStmtDoc(IdDoc("x")), ExprStmtDoc(IdDoc("y"))],
|
||||
],
|
||||
)
|
||||
def test_if_doc(then_branch, else_branch):
|
||||
predicate = IdDoc("x")
|
||||
|
||||
if not then_branch and not else_branch:
|
||||
with pytest.raises(ValueError) as e:
|
||||
IfDoc(predicate, then_branch, else_branch)
|
||||
return
|
||||
else:
|
||||
doc = IfDoc(predicate, then_branch, else_branch)
|
||||
|
||||
assert doc.predicate == predicate
|
||||
assert list(doc.then_branch) == then_branch
|
||||
assert list(doc.else_branch) == else_branch
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"body",
|
||||
[
|
||||
[],
|
||||
[ExprStmtDoc(IdDoc("x"))],
|
||||
[ExprStmtDoc(IdDoc("x")), ExprStmtDoc(IdDoc("y"))],
|
||||
],
|
||||
)
|
||||
def test_while_doc(body):
|
||||
predicate = IdDoc("x")
|
||||
|
||||
doc = WhileDoc(predicate, body)
|
||||
|
||||
assert doc.predicate == predicate
|
||||
assert list(doc.body) == body
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"body",
|
||||
[
|
||||
[],
|
||||
[ExprStmtDoc(IdDoc("x"))],
|
||||
[ExprStmtDoc(IdDoc("x")), ExprStmtDoc(IdDoc("y"))],
|
||||
],
|
||||
)
|
||||
def test_for_doc(body):
|
||||
lhs = IdDoc("x")
|
||||
rhs = IdDoc("y")
|
||||
|
||||
doc = ForDoc(lhs, rhs, body)
|
||||
|
||||
assert doc.lhs == lhs
|
||||
assert doc.rhs == rhs
|
||||
assert list(doc.body) == body
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"lhs",
|
||||
[
|
||||
None,
|
||||
IdDoc("x"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"body",
|
||||
[
|
||||
[],
|
||||
[ExprStmtDoc(IdDoc("x"))],
|
||||
[ExprStmtDoc(IdDoc("x")), ExprStmtDoc(IdDoc("y"))],
|
||||
],
|
||||
)
|
||||
def test_scope_doc(lhs, body):
|
||||
rhs = IdDoc("y")
|
||||
|
||||
doc = ScopeDoc(lhs, rhs, body)
|
||||
|
||||
assert doc.lhs == lhs
|
||||
assert doc.rhs == rhs
|
||||
assert list(doc.body) == body
|
||||
|
||||
|
||||
def test_expr_stmt_doc():
|
||||
expr = IdDoc("x")
|
||||
|
||||
doc = ExprStmtDoc(expr)
|
||||
|
||||
assert doc.expr == expr
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"msg",
|
||||
[
|
||||
None,
|
||||
LiteralDoc("msg"),
|
||||
],
|
||||
)
|
||||
def test_assert_doc(msg):
|
||||
test = IdDoc("x")
|
||||
|
||||
doc = AssertDoc(test, msg)
|
||||
|
||||
assert doc.test == test
|
||||
assert doc.msg == msg
|
||||
|
||||
|
||||
def test_return_doc():
|
||||
value = IdDoc("x")
|
||||
|
||||
doc = ReturnDoc(value)
|
||||
|
||||
assert doc.value == value
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"args",
|
||||
[
|
||||
[],
|
||||
[AssignDoc(IdDoc("x"), None, IdDoc("int"))],
|
||||
[
|
||||
AssignDoc(IdDoc("x"), None, IdDoc("int")),
|
||||
AssignDoc(IdDoc("y"), LiteralDoc(1), IdDoc("int")),
|
||||
],
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"decorators",
|
||||
[
|
||||
[],
|
||||
[IdDoc("test")],
|
||||
[IdDoc("test"), IdDoc("test2")],
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"return_type",
|
||||
[
|
||||
None,
|
||||
LiteralDoc(None),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"body",
|
||||
[
|
||||
[],
|
||||
[ExprStmtDoc(IdDoc("x"))],
|
||||
[ExprStmtDoc(IdDoc("x")), ExprStmtDoc(IdDoc("y"))],
|
||||
],
|
||||
)
|
||||
def test_function_doc(args, decorators, return_type, body):
|
||||
name = IdDoc("name")
|
||||
|
||||
doc = FunctionDoc(name, args, decorators, return_type, body)
|
||||
|
||||
assert doc.name == name
|
||||
assert list(doc.args) == args
|
||||
assert list(doc.decorators) == decorators
|
||||
assert doc.return_type == return_type
|
||||
assert list(doc.body) == body
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"decorators",
|
||||
[
|
||||
[],
|
||||
[IdDoc("test")],
|
||||
[IdDoc("test"), IdDoc("test2")],
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"body",
|
||||
[
|
||||
[],
|
||||
[ExprStmtDoc(IdDoc("x"))],
|
||||
[ExprStmtDoc(IdDoc("x")), ExprStmtDoc(IdDoc("y"))],
|
||||
],
|
||||
)
|
||||
def test_class_doc(decorators, body):
|
||||
name = IdDoc("name")
|
||||
|
||||
doc = ClassDoc(name, decorators, body)
|
||||
|
||||
assert doc.name == name
|
||||
assert list(doc.decorators) == decorators
|
||||
assert list(doc.body) == body
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"comment",
|
||||
[
|
||||
"",
|
||||
"test comment 1",
|
||||
"test comment 1\ntest comment 1",
|
||||
],
|
||||
)
|
||||
def test_comment_doc(comment):
|
||||
doc = CommentDoc(comment)
|
||||
assert doc.comment == comment
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"comment",
|
||||
[
|
||||
"",
|
||||
"test comment 1",
|
||||
"test comment 1\ntest comment 1",
|
||||
],
|
||||
)
|
||||
def test_doc_string_doc(comment):
|
||||
doc = DocStringDoc(comment)
|
||||
assert doc.comment == comment
|
||||
|
||||
|
||||
def test_stmt_doc_comment():
|
||||
doc = ExprStmtDoc(IdDoc("x"))
|
||||
assert doc.comment is None
|
||||
|
||||
comment = "test comment"
|
||||
doc.comment = comment
|
||||
# Make sure the previous statement doesn't set attribute
|
||||
# as if it's an ordinary Python object (__slots__ enforces this).
|
||||
assert not hasattr(doc, "__dict__") or "comment" not in doc.__dict__
|
||||
assert doc.comment == comment
|
||||
|
||||
|
||||
def test_doc_source_paths():
|
||||
doc = IdDoc("x")
|
||||
assert len(doc.source_paths) == 0
|
||||
|
||||
source_paths = [AccessPath.root(), AccessPath.root().attr("x")]
|
||||
|
||||
doc.source_paths = source_paths
|
||||
# This should triggers the __getattr__ and gets a tvm_ffi.Array
|
||||
assert not isinstance(doc.source_paths, list)
|
||||
assert list(doc.source_paths) == source_paths
|
||||
|
||||
doc.source_paths = []
|
||||
assert len(doc.source_paths) == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,55 @@
|
||||
# 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: F401
|
||||
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.highlight import _format, cprint
|
||||
|
||||
|
||||
def test_highlight_script():
|
||||
@tvm.script.ir_module
|
||||
class Module:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main( # type: ignore
|
||||
a: T.handle,
|
||||
b: T.handle,
|
||||
c: T.handle,
|
||||
) -> None: # pylint: disable=no-self-argument
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
A = T.match_buffer(a, [16, 128, 128])
|
||||
B = T.match_buffer(b, [16, 128, 128])
|
||||
C = T.match_buffer(c, [16, 128, 128])
|
||||
for n, i, j, k in T.grid(16, 128, 128, 128):
|
||||
with T.sblock("matmul"):
|
||||
vn, vi, vj, vk = T.axis.remap("SSSR", [n, i, j, k])
|
||||
with T.init():
|
||||
C[vn, vi, vj] = 0.0 # type: ignore
|
||||
C[vn, vi, vj] = C[vn, vi, vj] + A[vn, vi, vk] * B[vn, vj, vk]
|
||||
|
||||
Module.show()
|
||||
Module["main"].show()
|
||||
Module["main"].show(style="light")
|
||||
Module["main"].show(style="dark")
|
||||
Module["main"].show(style="ansi")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,68 @@
|
||||
# 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.
|
||||
|
||||
from tvm_ffi.access_path import AccessPath
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.runtime.script_printer import PrinterConfig, _script
|
||||
|
||||
|
||||
def test_render_invisible_path_info_defaults_to_true():
|
||||
config = PrinterConfig()
|
||||
assert config.render_invisible_path_info is True
|
||||
|
||||
result = tvm.ir.PrimType("int32").script(path_to_underline=[AccessPath.root().attr("dtype")])
|
||||
assert result == (
|
||||
"Access path: <root>.dtype\n"
|
||||
"Note: The underlined object is the nearest visible parent of this path.\n\n"
|
||||
"T.int32\n"
|
||||
"^^^^^^^"
|
||||
)
|
||||
|
||||
|
||||
def test_render_invisible_path_info_without_target_path_is_unchanged():
|
||||
assert tvm.ir.PrimType("int32").script() == "T.int32"
|
||||
|
||||
|
||||
def test_render_invisible_path_info_explicit_false_uses_legacy_output():
|
||||
result = tvm.ir.PrimType("int32").script(
|
||||
path_to_underline=[AccessPath.root().attr("dtype")],
|
||||
extra_config={"render_invisible_path_info": False},
|
||||
)
|
||||
assert result == "T.int32\n^^^^^^^"
|
||||
|
||||
|
||||
def test_render_invisible_path_info_exact_visible_path_omits_note():
|
||||
result = tvm.ir.PrimType("int32").script(path_to_underline=[AccessPath.root()])
|
||||
assert result == "Access path: <root>\n\nT.int32\n^^^^^^^"
|
||||
|
||||
|
||||
def test_render_invisible_path_info_without_visible_path():
|
||||
result = _script(
|
||||
tvm.runtime.ShapeTuple([1, 2]),
|
||||
PrinterConfig(path_to_underline=[AccessPath.root().attr("missing")]),
|
||||
)
|
||||
assert result == (
|
||||
"Access path: <root>.missing\n"
|
||||
"Note: No visible object for this path is rendered in TVMScript.\n\n"
|
||||
"Shape(1, 2)"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,68 @@
|
||||
# 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-docstring
|
||||
|
||||
import pytest
|
||||
|
||||
from tvm import IRModule
|
||||
from tvm.script.ir_builder import IRBuilder
|
||||
from tvm.script.ir_builder import ir as I
|
||||
from tvm.script.ir_builder import tirx as T
|
||||
|
||||
|
||||
def _assert_print(obj, expected):
|
||||
assert str(obj).strip() == expected.strip()
|
||||
assert repr(obj).strip() == expected.strip()
|
||||
if isinstance(obj, IRModule):
|
||||
assert obj.script().strip() == expected.strip()
|
||||
|
||||
|
||||
def test_ir_module():
|
||||
with IRBuilder() as ib: # pylint: disable=invalid-name
|
||||
with I.ir_module():
|
||||
with T.prim_func(s_tir=True):
|
||||
T.func_name("foo")
|
||||
mod = ib.get()
|
||||
_assert_print(
|
||||
mod,
|
||||
"""
|
||||
# from tvm.script import ir as I
|
||||
# from tvm.script import tirx as T
|
||||
# from tvm.tirx.layout import Axis
|
||||
|
||||
@I.ir_module
|
||||
class Module:
|
||||
@T.prim_func(s_tir=True)
|
||||
def foo():
|
||||
T.evaluate(0)""",
|
||||
)
|
||||
|
||||
|
||||
def test_failed_invalid_prefix():
|
||||
with IRBuilder() as ib: # pylint: disable=invalid-name
|
||||
with I.ir_module():
|
||||
with T.prim_func():
|
||||
T.func_name("foo")
|
||||
mod = ib.get()
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
mod.script(ir_prefix="2I")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_ir_module()
|
||||
test_failed_invalid_prefix()
|
||||
@@ -0,0 +1,48 @@
|
||||
# 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-docstring
|
||||
# ruff: noqa: F841
|
||||
import tvm.testing
|
||||
from tvm.script.parser import ir as I
|
||||
from tvm.script.parser import tirx as T
|
||||
|
||||
|
||||
def test_str_metadata():
|
||||
# This test is to check we reuse the existing metadata element for the same tirx.StringImm
|
||||
# So metadata["tirx.StringImm"][0] will occur in the printed script for three times
|
||||
str_imm = T.StringImm("aaa\nbbb\n")
|
||||
|
||||
@I.ir_module
|
||||
class Module:
|
||||
@T.prim_func(s_tir=True)
|
||||
def foo() -> None:
|
||||
A = str_imm
|
||||
B = str_imm
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def foo1() -> None:
|
||||
A = str_imm
|
||||
|
||||
printed_str = Module.script(verbose_expr=True)
|
||||
assert (
|
||||
printed_str.count('metadata["tirx.StringImm"][0]') == 3
|
||||
and printed_str.count('metadata["tirx.StringImm"][1]') == 0
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,174 @@
|
||||
# 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: F841
|
||||
|
||||
import pytest
|
||||
from tvm_ffi.access_path import AccessPath
|
||||
|
||||
import tvm
|
||||
from tvm.ir import assert_structural_equal
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import tirx as T
|
||||
|
||||
|
||||
def _error_message(exception):
|
||||
return str(exception)
|
||||
|
||||
|
||||
def _expected_result(func1, func2, objpath1, objpath2):
|
||||
return f"""StructuralEqual check failed, caused by lhs at {objpath1}:
|
||||
{func1.script(path_to_underline=[objpath1], syntax_sugar=False)}
|
||||
and rhs at {objpath2}:
|
||||
{func2.script(path_to_underline=[objpath2], syntax_sugar=False)}"""
|
||||
|
||||
|
||||
def test_prim_type_hidden_path_exact_message():
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
assert_structural_equal(tvm.ir.PrimType("int32"), tvm.ir.PrimType("float32"))
|
||||
|
||||
assert str(exc_info.value) == (
|
||||
"StructuralEqual check failed, caused by lhs at <root>.dtype:\n"
|
||||
"Access path: <root>.dtype\n"
|
||||
"Note: The underlined object is the nearest visible parent of this path.\n\n"
|
||||
"T.int32\n"
|
||||
"^^^^^^^\n"
|
||||
"and rhs at <root>.dtype:\n"
|
||||
"Access path: <root>.dtype\n"
|
||||
"Note: The underlined object is the nearest visible parent of this path.\n\n"
|
||||
"T.float32\n"
|
||||
"^^^^^^^^^"
|
||||
)
|
||||
|
||||
|
||||
def test_prim_func_buffer_map():
|
||||
@T.prim_func(s_tir=True)
|
||||
def func1(a: T.handle, b: T.handle):
|
||||
A = T.match_buffer(a, (128, 128))
|
||||
B = T.match_buffer(b, (128, 128))
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func2(a: T.handle, b: T.handle):
|
||||
A = T.match_buffer(a, (128, 128))
|
||||
B = T.match_buffer(b, (128, 256))
|
||||
|
||||
func1 = func1.with_attr("global_symbol", "main")
|
||||
func2 = func2.with_attr("global_symbol", "main")
|
||||
|
||||
with pytest.raises(ValueError) as ve:
|
||||
assert_structural_equal(func1, func2)
|
||||
assert _error_message(ve.value) == _expected_result(
|
||||
func1,
|
||||
func2,
|
||||
AccessPath.root()
|
||||
.attr("buffer_map")
|
||||
.map_item(func1.params[1])
|
||||
.attr("shape")
|
||||
.array_item(1)
|
||||
.attr("value"),
|
||||
AccessPath.root()
|
||||
.attr("buffer_map")
|
||||
.map_item(func2.params[1])
|
||||
.attr("shape")
|
||||
.array_item(1)
|
||||
.attr("value"),
|
||||
)
|
||||
|
||||
|
||||
def test_evaluate():
|
||||
@I.ir_module(s_tir=True)
|
||||
class module1:
|
||||
@T.prim_func(s_tir=True)
|
||||
def func():
|
||||
T.evaluate(0)
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class module2:
|
||||
@T.prim_func(s_tir=True)
|
||||
def func():
|
||||
T.evaluate(1)
|
||||
|
||||
with pytest.raises(ValueError) as ve:
|
||||
assert_structural_equal(module1, module2)
|
||||
assert _error_message(ve.value) == _expected_result(
|
||||
module1,
|
||||
module2,
|
||||
AccessPath.root()
|
||||
.attr("functions")
|
||||
.map_item(module1.get_global_var("func"))
|
||||
.attr("body")
|
||||
.attr("value")
|
||||
.attr("value"),
|
||||
AccessPath.root()
|
||||
.attr("functions")
|
||||
.map_item(module2.get_global_var("func"))
|
||||
.attr("body")
|
||||
.attr("value")
|
||||
.attr("value"),
|
||||
)
|
||||
|
||||
|
||||
def test_allocate():
|
||||
@T.prim_func(s_tir=True)
|
||||
def func1():
|
||||
a = T.alloc_buffer((128, 128), dtype="float32")
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func2():
|
||||
a = T.alloc_buffer((256, 128), dtype="float32")
|
||||
|
||||
func1 = func1.with_attr("global_symbol", "main")
|
||||
func2 = func2.with_attr("global_symbol", "main")
|
||||
|
||||
with pytest.raises(ValueError) as ve:
|
||||
assert_structural_equal(func1, func2)
|
||||
|
||||
assert _error_message(ve.value) == _expected_result(
|
||||
func1,
|
||||
func2,
|
||||
AccessPath.root().attr("body").attr("buffer").attr("shape").array_item(0).attr("value"),
|
||||
AccessPath.root().attr("body").attr("buffer").attr("shape").array_item(0).attr("value"),
|
||||
)
|
||||
|
||||
|
||||
def test_for():
|
||||
@T.prim_func(s_tir=True)
|
||||
def func1():
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock():
|
||||
pass
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func2():
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
with T.sblock():
|
||||
pass
|
||||
|
||||
func1 = func1.with_attr("global_symbol", "main")
|
||||
func2 = func2.with_attr("global_symbol", "main")
|
||||
|
||||
with pytest.raises(ValueError) as ve:
|
||||
assert_structural_equal(func1, func2)
|
||||
assert _error_message(ve.value) == _expected_result(
|
||||
func1,
|
||||
func2,
|
||||
AccessPath.root().attr("body").attr("block").attr("body").attr("body").attr("body"),
|
||||
AccessPath.root().attr("body").attr("block").attr("body").attr("body").attr("body"),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,590 @@
|
||||
# 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
|
||||
from tvm_ffi.access_path import AccessPath
|
||||
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.printer.doc import (
|
||||
ExprStmtDoc,
|
||||
IdDoc,
|
||||
OperationDoc,
|
||||
OperationKind,
|
||||
StmtBlockDoc,
|
||||
)
|
||||
from tvm.script.printer.doc_printer import to_python_script
|
||||
|
||||
|
||||
def make_path(name: str) -> AccessPath:
|
||||
return AccessPath.root().attr(name)
|
||||
|
||||
|
||||
def make_id_doc(name: str, path_name: str | None = None) -> IdDoc:
|
||||
if path_name is None:
|
||||
path_name = name
|
||||
doc = IdDoc(name)
|
||||
doc.source_paths = [make_path(path_name)]
|
||||
return doc
|
||||
|
||||
|
||||
def format_script(s: str) -> str:
|
||||
"""
|
||||
Remove leading and trailing blank lines, and make the minimum idention 0
|
||||
"""
|
||||
s = s.strip("\n")
|
||||
|
||||
non_empty_lines = [line for line in s.splitlines() if line and not line.isspace()]
|
||||
if not non_empty_lines:
|
||||
# no actual content
|
||||
return "\n"
|
||||
|
||||
line_indents = [len(line) - len(line.lstrip(" ")) for line in non_empty_lines]
|
||||
spaces_to_remove = min(line_indents)
|
||||
|
||||
cleaned_lines = "\n".join(line[spaces_to_remove:] for line in s.splitlines())
|
||||
if not cleaned_lines.endswith("\n"):
|
||||
cleaned_lines += "\n"
|
||||
return cleaned_lines.strip()
|
||||
|
||||
|
||||
def format_script_with_path_info(s: str, *path_info: str) -> str:
|
||||
return "\n".join(path_info) + "\n\n" + format_script(s)
|
||||
|
||||
|
||||
def test_underline_basic():
|
||||
doc = StmtBlockDoc(
|
||||
[
|
||||
ExprStmtDoc(make_id_doc("foo")),
|
||||
ExprStmtDoc(OperationDoc(OperationKind.Add, [make_id_doc("bar"), make_id_doc("baz")])),
|
||||
ExprStmtDoc(make_id_doc("qux")),
|
||||
]
|
||||
)
|
||||
assert to_python_script(
|
||||
doc, path_to_underline=[make_path("baz")]
|
||||
) == format_script_with_path_info(
|
||||
"""
|
||||
foo
|
||||
bar + baz
|
||||
^^^
|
||||
qux
|
||||
""",
|
||||
"Access path: <root>.baz",
|
||||
)
|
||||
|
||||
|
||||
def test_underline_multiple_spans():
|
||||
doc = StmtBlockDoc(
|
||||
[
|
||||
ExprStmtDoc(make_id_doc("foo")),
|
||||
ExprStmtDoc(make_id_doc("bar")),
|
||||
ExprStmtDoc(OperationDoc(OperationKind.Add, [make_id_doc("foo"), make_id_doc("foo")])),
|
||||
]
|
||||
)
|
||||
assert to_python_script(
|
||||
doc, path_to_underline=[make_path("foo")]
|
||||
) == format_script_with_path_info(
|
||||
"""
|
||||
foo
|
||||
^^^
|
||||
bar
|
||||
foo + foo
|
||||
^^^ ^^^
|
||||
""",
|
||||
"Access path: <root>.foo",
|
||||
)
|
||||
|
||||
|
||||
def test_underline_multiple_spans_with_line_numbers():
|
||||
doc = StmtBlockDoc(
|
||||
[
|
||||
ExprStmtDoc(make_id_doc("foo")),
|
||||
ExprStmtDoc(make_id_doc("bar")),
|
||||
ExprStmtDoc(OperationDoc(OperationKind.Add, [make_id_doc("foo"), make_id_doc("foo")])),
|
||||
]
|
||||
)
|
||||
assert to_python_script(
|
||||
doc, print_line_numbers=True, path_to_underline=[make_path("foo")]
|
||||
) == format_script_with_path_info(
|
||||
"""
|
||||
1 foo
|
||||
^^^
|
||||
2 bar
|
||||
3 foo + foo
|
||||
^^^ ^^^
|
||||
""",
|
||||
"Access path: <root>.foo",
|
||||
)
|
||||
|
||||
|
||||
def test_underline_multiline():
|
||||
doc = StmtBlockDoc(
|
||||
[
|
||||
ExprStmtDoc(IdDoc("foo")),
|
||||
ExprStmtDoc(IdDoc("bar")),
|
||||
]
|
||||
)
|
||||
doc.source_paths = [make_path("whole_doc")]
|
||||
|
||||
assert to_python_script(
|
||||
doc, path_to_underline=[make_path("whole_doc")]
|
||||
) == format_script_with_path_info(
|
||||
"""
|
||||
foo
|
||||
^^^
|
||||
bar
|
||||
^^^
|
||||
""",
|
||||
"Access path: <root>.whole_doc",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"to_underline, expected_text",
|
||||
[
|
||||
(
|
||||
[0],
|
||||
"""
|
||||
x0
|
||||
^^
|
||||
x1
|
||||
x2
|
||||
(... 7 lines skipped ...)
|
||||
""",
|
||||
),
|
||||
(
|
||||
[1],
|
||||
"""
|
||||
x0
|
||||
x1
|
||||
^^
|
||||
x2
|
||||
x3
|
||||
(... 6 lines skipped ...)
|
||||
""",
|
||||
),
|
||||
(
|
||||
[3],
|
||||
"""
|
||||
x0
|
||||
x1
|
||||
x2
|
||||
x3
|
||||
^^
|
||||
x4
|
||||
x5
|
||||
(... 4 lines skipped ...)
|
||||
""",
|
||||
),
|
||||
(
|
||||
[4],
|
||||
"""
|
||||
(... 2 lines skipped ...)
|
||||
x2
|
||||
x3
|
||||
x4
|
||||
^^
|
||||
x5
|
||||
x6
|
||||
(... 3 lines skipped ...)
|
||||
""",
|
||||
),
|
||||
(
|
||||
[6],
|
||||
"""
|
||||
(... 4 lines skipped ...)
|
||||
x4
|
||||
x5
|
||||
x6
|
||||
^^
|
||||
x7
|
||||
x8
|
||||
x9
|
||||
""",
|
||||
),
|
||||
(
|
||||
[9],
|
||||
"""
|
||||
(... 7 lines skipped ...)
|
||||
x7
|
||||
x8
|
||||
x9
|
||||
^^
|
||||
""",
|
||||
),
|
||||
(
|
||||
[0, 9],
|
||||
"""
|
||||
x0
|
||||
^^
|
||||
x1
|
||||
x2
|
||||
(... 4 lines skipped ...)
|
||||
x7
|
||||
x8
|
||||
x9
|
||||
^^
|
||||
""",
|
||||
),
|
||||
(
|
||||
[0, 3, 9],
|
||||
"""
|
||||
x0
|
||||
^^
|
||||
x1
|
||||
x2
|
||||
x3
|
||||
^^
|
||||
x4
|
||||
x5
|
||||
x6
|
||||
x7
|
||||
x8
|
||||
x9
|
||||
^^
|
||||
""",
|
||||
),
|
||||
(
|
||||
[0, 6, 9],
|
||||
"""
|
||||
x0
|
||||
^^
|
||||
x1
|
||||
x2
|
||||
x3
|
||||
x4
|
||||
x5
|
||||
x6
|
||||
^^
|
||||
x7
|
||||
x8
|
||||
x9
|
||||
^^
|
||||
""",
|
||||
),
|
||||
(
|
||||
[33],
|
||||
"""
|
||||
x0
|
||||
x1
|
||||
x2
|
||||
x3
|
||||
x4
|
||||
x5
|
||||
x6
|
||||
x7
|
||||
x8
|
||||
x9
|
||||
""",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_print_two_context_lines(to_underline, expected_text):
|
||||
doc = StmtBlockDoc(
|
||||
[ExprStmtDoc(make_id_doc(f"x{i}", "yes" if i in to_underline else "no")) for i in range(10)]
|
||||
)
|
||||
result = to_python_script(doc, num_context_lines=2, path_to_underline=[make_path("yes")])
|
||||
path_info = ["Access path: <root>.yes"]
|
||||
if to_underline == [33]:
|
||||
path_info.append("Note: No visible object for this path is rendered in TVMScript.")
|
||||
assert result == format_script_with_path_info(expected_text, *path_info)
|
||||
|
||||
|
||||
def test_underline_and_print_line_numbers():
|
||||
doc = StmtBlockDoc([ExprStmtDoc(make_id_doc(f"line{i + 1}")) for i in range(12)])
|
||||
result = to_python_script(doc, print_line_numbers=True, path_to_underline=[make_path("line6")])
|
||||
assert result == "Access path: <root>.line6\n\n " + format_script(
|
||||
"""
|
||||
1 line1
|
||||
2 line2
|
||||
3 line3
|
||||
4 line4
|
||||
5 line5
|
||||
6 line6
|
||||
^^^^^
|
||||
7 line7
|
||||
8 line8
|
||||
9 line9
|
||||
10 line10
|
||||
11 line11
|
||||
12 line12
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def test_underline_multi_access_paths():
|
||||
doc = StmtBlockDoc([ExprStmtDoc(make_id_doc(f"line{i + 1}")) for i in range(10)])
|
||||
result = to_python_script(
|
||||
doc,
|
||||
path_to_underline=[
|
||||
make_path("line1"),
|
||||
make_path("line3"),
|
||||
make_path("line5"),
|
||||
make_path("line7"),
|
||||
make_path("line9"),
|
||||
],
|
||||
)
|
||||
assert result == format_script_with_path_info(
|
||||
"""
|
||||
line1
|
||||
^^^^^
|
||||
line2
|
||||
line3
|
||||
^^^^^
|
||||
line4
|
||||
line5
|
||||
^^^^^
|
||||
line6
|
||||
line7
|
||||
^^^^^
|
||||
line8
|
||||
line9
|
||||
^^^^^
|
||||
line10
|
||||
""",
|
||||
"Access path: <root>.line1",
|
||||
"Access path: <root>.line3",
|
||||
"Access path: <root>.line5",
|
||||
"Access path: <root>.line7",
|
||||
"Access path: <root>.line9",
|
||||
)
|
||||
|
||||
|
||||
def test_underline_and_print_line_numbers_with_context():
|
||||
doc = StmtBlockDoc([ExprStmtDoc(make_id_doc(f"line{i + 1}")) for i in range(12)])
|
||||
result = to_python_script(
|
||||
doc, print_line_numbers=True, num_context_lines=2, path_to_underline=[make_path("line8")]
|
||||
)
|
||||
assert result == format_script_with_path_info(
|
||||
"""
|
||||
(... 5 lines skipped ...)
|
||||
6 line6
|
||||
7 line7
|
||||
8 line8
|
||||
^^^^^
|
||||
9 line9
|
||||
10 line10
|
||||
(... 2 lines skipped ...)
|
||||
""",
|
||||
"Access path: <root>.line8",
|
||||
)
|
||||
|
||||
|
||||
def test_underline_based_on_path_prefix():
|
||||
doc = StmtBlockDoc([ExprStmtDoc(make_id_doc("foo")), ExprStmtDoc(make_id_doc("bar"))])
|
||||
result = to_python_script(doc, path_to_underline=[make_path("foo").attr("x").attr("y")])
|
||||
# There is no document that matches the desired path exactly,
|
||||
# but path of "foo" is a prefix of the desired path, and thus should be underlined.
|
||||
assert result == format_script_with_path_info(
|
||||
"""
|
||||
foo
|
||||
^^^
|
||||
bar
|
||||
""",
|
||||
"Access path: <root>.foo.x.y",
|
||||
"Note: The underlined object is the nearest visible parent of this path.",
|
||||
)
|
||||
|
||||
|
||||
def test_longer_prefix_must_win():
|
||||
foo_x = IdDoc("foo_x")
|
||||
foo_x.source_paths = [make_path("foo").attr("x")]
|
||||
|
||||
doc = StmtBlockDoc(
|
||||
[ExprStmtDoc(make_id_doc("foo")), ExprStmtDoc(make_id_doc("bar")), ExprStmtDoc(foo_x)]
|
||||
)
|
||||
result = to_python_script(doc, path_to_underline=[make_path("foo").attr("x").attr("y")])
|
||||
# "foo" should not be underlined because there is a document with a more specific path prefix
|
||||
assert result == format_script_with_path_info(
|
||||
"""
|
||||
foo
|
||||
bar
|
||||
foo_x
|
||||
^^^^^
|
||||
""",
|
||||
"Access path: <root>.foo.x.y",
|
||||
"Note: The underlined object is the nearest visible parent of this path.",
|
||||
)
|
||||
|
||||
|
||||
def test_underline_from_obj():
|
||||
@T.prim_func(s_tir=True)
|
||||
def func(a: T.int32, b: T.int32):
|
||||
T.evaluate(a)
|
||||
T.evaluate(b)
|
||||
T.evaluate(a)
|
||||
T.evaluate(b)
|
||||
T.evaluate(a)
|
||||
T.evaluate(b)
|
||||
|
||||
result = func.with_attr("global_symbol", "main").script(
|
||||
obj_to_underline=[func.params[0]],
|
||||
extra_config={"render_invisible_path_info": False},
|
||||
)
|
||||
assert result == format_script(
|
||||
"""
|
||||
# from tvm.script import tirx as T
|
||||
# from tvm.tirx.layout import Axis
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.int32, b: T.int32):
|
||||
T.evaluate(a)
|
||||
^
|
||||
T.evaluate(b)
|
||||
T.evaluate(a)
|
||||
^
|
||||
T.evaluate(b)
|
||||
T.evaluate(a)
|
||||
^
|
||||
T.evaluate(b)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def test_underline_from_multi_obj():
|
||||
@T.prim_func(s_tir=True)
|
||||
def func():
|
||||
T.evaluate(-1)
|
||||
T.evaluate(1)
|
||||
T.evaluate(2)
|
||||
T.evaluate(3)
|
||||
T.evaluate(4)
|
||||
T.evaluate(5)
|
||||
T.evaluate(6)
|
||||
T.evaluate(7)
|
||||
|
||||
result = func.with_attr("global_symbol", "main").script(
|
||||
obj_to_underline=[
|
||||
func.body.seq[1],
|
||||
func.body.seq[3],
|
||||
func.body.seq[5],
|
||||
func.body.seq[7],
|
||||
],
|
||||
extra_config={"render_invisible_path_info": False},
|
||||
)
|
||||
assert result == format_script(
|
||||
"""
|
||||
# from tvm.script import tirx as T
|
||||
# from tvm.tirx.layout import Axis
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def main():
|
||||
T.evaluate(-1)
|
||||
T.evaluate(1)
|
||||
^^^^^^^^^^^^^
|
||||
T.evaluate(2)
|
||||
T.evaluate(3)
|
||||
^^^^^^^^^^^^^
|
||||
T.evaluate(4)
|
||||
T.evaluate(5)
|
||||
^^^^^^^^^^^^^
|
||||
T.evaluate(6)
|
||||
T.evaluate(7)
|
||||
^^^^^^^^^^^^^
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def test_underline_func():
|
||||
@T.prim_func(s_tir=True)
|
||||
def func():
|
||||
T.evaluate(0)
|
||||
|
||||
result = func.with_attr("global_symbol", "main").script(
|
||||
path_to_underline=[
|
||||
AccessPath.root(),
|
||||
],
|
||||
extra_config={"render_invisible_path_info": False},
|
||||
)
|
||||
assert result == format_script(
|
||||
"""
|
||||
# from tvm.script import tirx as T
|
||||
# from tvm.tirx.layout import Axis
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
def main():
|
||||
^^^^^^^^^^^
|
||||
T.evaluate(0)
|
||||
^^^^^^^^^^^^^
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def test_underline_func_in_irmodule():
|
||||
@I.ir_module
|
||||
class irmodule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def func():
|
||||
T.evaluate(0)
|
||||
|
||||
result = irmodule.script(
|
||||
path_to_underline=[
|
||||
AccessPath.root().attr("functions").map_item(irmodule.get_global_var("func")),
|
||||
],
|
||||
extra_config={"render_invisible_path_info": False},
|
||||
)
|
||||
assert result == format_script(
|
||||
"""
|
||||
# from tvm.script import ir as I
|
||||
# from tvm.script import tirx as T
|
||||
# from tvm.tirx.layout import Axis
|
||||
|
||||
@I.ir_module
|
||||
class Module:
|
||||
@T.prim_func(s_tir=True)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
def func():
|
||||
^^^^^^^^^^^
|
||||
T.evaluate(0)
|
||||
^^^^^^^^^^^^^
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def test_underline_irmodule():
|
||||
@I.ir_module
|
||||
class irmodule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def func():
|
||||
T.evaluate(0)
|
||||
|
||||
result = irmodule.script(
|
||||
path_to_underline=[
|
||||
AccessPath.root(),
|
||||
],
|
||||
extra_config={"render_invisible_path_info": False},
|
||||
)
|
||||
assert result == format_script(
|
||||
"""
|
||||
# from tvm.script import ir as I
|
||||
# from tvm.script import tirx as T
|
||||
# from tvm.tirx.layout import Axis
|
||||
|
||||
@I.ir_module
|
||||
^^^^^^^^^^^^
|
||||
class Module:
|
||||
^^^^^^^^^^^^^
|
||||
@T.prim_func(s_tir=True)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
def func():
|
||||
^^^^^^^^^^^
|
||||
T.evaluate(0)
|
||||
^^^^^^^^^^^^^
|
||||
"""
|
||||
)
|
||||
@@ -0,0 +1,94 @@
|
||||
# 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: F841
|
||||
import numpy
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import tirx as T
|
||||
|
||||
# This numpy array is used to test the comparison between the global objects and the
|
||||
# `tvm.script.tirx` submodule.
|
||||
np_array = numpy.array([0, 1, 2, 3])
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
B = T.match_buffer(b, [128, 128])
|
||||
C = T.match_buffer(c, [128, 128])
|
||||
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
with T.sblock("update"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = T.float32(0)
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
|
||||
|
||||
|
||||
def test_multi_element_array_in_outmost_namespace():
|
||||
func = matmul
|
||||
rt_func = tvm.script.from_source(func.script())
|
||||
tvm.ir.assert_structural_equal(func, rt_func)
|
||||
|
||||
|
||||
def test_different_dtype_assignment_to_var():
|
||||
@T.prim_func(s_tir=True)
|
||||
def test_case():
|
||||
a = T.sblock_alloc_buffer((10, 10), dtype="int8")
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func_ref():
|
||||
a = T.sblock_alloc_buffer([10, 10], dtype="int8")
|
||||
T.evaluate(0)
|
||||
|
||||
tvm.ir.assert_structural_equal(
|
||||
test_case.with_attr("global_symbol", "main"), func_ref.with_attr("global_symbol", "main")
|
||||
)
|
||||
|
||||
|
||||
def test_var_capturing_order():
|
||||
b = 2
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def test_case():
|
||||
k: T.let[T.int32] = b
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func_ref():
|
||||
k: T.let[T.int32] = 2
|
||||
T.evaluate(0)
|
||||
|
||||
tvm.ir.assert_structural_equal(
|
||||
test_case.with_attr("global_symbol", "main"), func_ref.with_attr("global_symbol", "main")
|
||||
)
|
||||
|
||||
|
||||
def test_tir_buffer_region_extent_correct_dtype():
|
||||
@T.prim_func(s_tir=True)
|
||||
def func(A: T.Buffer((T.int64(16), T.int64(1)), "float32")):
|
||||
for i in T.grid(T.int64(16)):
|
||||
with T.sblock("block"):
|
||||
vi = T.axis.remap("S", [i])
|
||||
T.reads(A[vi, T.int64(0) : T.int64(1)])
|
||||
T.evaluate(0)
|
||||
|
||||
assert func.body.block.body.body.block.reads[0].region[0].extent.ty.dtype == "int64"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,545 @@
|
||||
# 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,invalid-name,pointless-string-statement
|
||||
# ruff: noqa: E741, F401, F841
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
import tvm.testing
|
||||
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
|
||||
from tvm.script import from_source
|
||||
from tvm.script import tirx as T
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def transformed_matmul_no_syntax_sugar(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
B = T.match_buffer(b, [128, 128])
|
||||
C = T.match_buffer(c, [128, 128])
|
||||
|
||||
for i0, i1, i2_outer, i2_inner_outer, i2_inner_inner in T.grid(128, 128, 4, 8, 4):
|
||||
with T.sblock("update"):
|
||||
vi, vj = T.axis.remap("SS", [i0, i1])
|
||||
vk = T.axis.R(128, i2_outer * 32 + i2_inner_outer * 4 + i2_inner_inner)
|
||||
T.reads([C[vi, vj], A[vi, vk], B[vj, vk]])
|
||||
T.writes([C[vi, vj], A[vi, vk]])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
A[vi, vk] = A[vi, vk] + B[vj, vk]
|
||||
C[vi, vj] = C[vi, vj] + (A[vi, vk] * B[vj, vk])
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def transformed_matmul_syntax_sugar(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
B = T.match_buffer(b, [128, 128])
|
||||
C = T.match_buffer(c, [128, 128])
|
||||
|
||||
for i0, i1, i2_outer, i2_inner_outer, i2_inner_inner in T.grid(128, 128, 4, 8, 4):
|
||||
with T.sblock("update"):
|
||||
vi, vj = T.axis.remap("SS", [i0, i1])
|
||||
vk = T.axis.R(128, i2_outer * 32 + i2_inner_outer * 4 + i2_inner_inner)
|
||||
T.reads(C[vi, vj], A[vi, vk], B[vj, vk])
|
||||
T.writes(C[vi, vj], A[vi, vk])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
A[vi, vk] = A[vi, vk] + B[vj, vk]
|
||||
C[vi, vj] = C[vi, vj] + (A[vi, vk] * B[vj, vk])
|
||||
|
||||
|
||||
def test_reads_writes_syntax_sugar():
|
||||
assert_structural_equal_ignore_global_symbol(
|
||||
transformed_matmul_no_syntax_sugar, transformed_matmul_syntax_sugar
|
||||
)
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def loop_no_syntax_sugar(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, (128, 128, 128, 128))
|
||||
for i in T.serial(0, 128):
|
||||
for j in T.parallel(0, 128):
|
||||
for k in T.vectorized(0, 128):
|
||||
for x in T.unroll(0, 128):
|
||||
for y in T.thread_binding(0, 128, thread="threadIdx.x"):
|
||||
for z in T.thread_binding(0, 128, thread="threadIdx.x"):
|
||||
A[i, j, k, x] = A[i, j, k, x] * 2.0
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def loop_syntax_sugar(a: T.handle) -> None:
|
||||
A = T.match_buffer(a, (128, 128, 128, 128))
|
||||
for i in T.serial(128):
|
||||
for j in T.parallel(128):
|
||||
for k in T.vectorized(128):
|
||||
for x in T.unroll(128):
|
||||
for y in T.thread_binding(128, "threadIdx.x"):
|
||||
for z in T.thread_binding(128, thread="threadIdx.x"):
|
||||
A[i, j, k, x] = A[i, j, k, x] * 2.0
|
||||
|
||||
|
||||
def test_loop_syntax_sugar():
|
||||
assert_structural_equal_ignore_global_symbol(loop_no_syntax_sugar, loop_syntax_sugar)
|
||||
|
||||
|
||||
# match buffer - use kwargs
|
||||
@T.prim_func(s_tir=True)
|
||||
def elementwise_handle(
|
||||
a: T.handle,
|
||||
b: T.handle,
|
||||
) -> None:
|
||||
A = T.match_buffer(a, (128, 128, 128, 128))
|
||||
B = T.match_buffer(b, (128, 128, 128, 128))
|
||||
for i, j, k, l in T.grid(128, 128, 128, 128):
|
||||
with T.sblock("B"):
|
||||
vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
|
||||
B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * 2.0
|
||||
|
||||
|
||||
# match buffer - use buffer with kwargs
|
||||
@T.prim_func(s_tir=True)
|
||||
def elementwise_buffer_kwargs(
|
||||
a: T.Buffer(shape=(128, 128, 128, 128), dtype="float32"),
|
||||
b: T.Buffer(shape=(128, 128, 128, 128), dtype="float32"),
|
||||
) -> None:
|
||||
for i, j, k, l in T.grid(128, 128, 128, 128):
|
||||
with T.sblock("B"):
|
||||
vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
|
||||
b[vi, vj, vk, vl] = a[vi, vj, vk, vl] * 2.0
|
||||
|
||||
|
||||
# match buffer - use buffer without kwargs
|
||||
@T.prim_func(s_tir=True)
|
||||
def elementwise_buffer_no_kwargs(
|
||||
a: T.Buffer((128, 128, 128, 128), "float32"),
|
||||
b: T.Buffer((128, 128, 128, 128), "float32"),
|
||||
) -> None:
|
||||
for i, j, k, l in T.grid(128, 128, 128, 128):
|
||||
with T.sblock("B"):
|
||||
vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
|
||||
b[vi, vj, vk, vl] = a[vi, vj, vk, vl] * 2.0
|
||||
|
||||
|
||||
def test_match_buffer_syntax_sugar():
|
||||
# with kwargs
|
||||
assert_structural_equal_ignore_global_symbol(elementwise_handle, elementwise_buffer_kwargs)
|
||||
# without kwargs
|
||||
assert_structural_equal_ignore_global_symbol(elementwise_handle, elementwise_buffer_no_kwargs)
|
||||
|
||||
|
||||
def test_match_buffer_1d():
|
||||
@T.prim_func(s_tir=True)
|
||||
def func_no_sugar(a: T.handle):
|
||||
A = T.match_buffer(a, shape=(16,))
|
||||
for i in T.serial(16):
|
||||
A[i] = 0.0
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func_with_sugar(A: T.Buffer(16, "float32")):
|
||||
for i in T.serial(16):
|
||||
A[i] = 0.0
|
||||
|
||||
assert_structural_equal_ignore_global_symbol(func_no_sugar, func_with_sugar)
|
||||
|
||||
|
||||
# dynamic shape gemm
|
||||
@T.prim_func(s_tir=True)
|
||||
def gemm_dyn_shape(a: T.handle, b: T.handle, c: T.handle):
|
||||
N = T.int32()
|
||||
M = T.int32()
|
||||
K = T.int32()
|
||||
A = T.match_buffer(a, (N, K), "float32")
|
||||
B = T.match_buffer(b, (K, M), "float32")
|
||||
C = T.match_buffer(c, (N, M), "float32")
|
||||
for i, j, k in T.grid(N, M, K):
|
||||
with T.sblock("gemm"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
|
||||
def test_dynamic_shape_gemm():
|
||||
gemm_dyn_shape_roundtrip = from_source(gemm_dyn_shape.script())
|
||||
assert_structural_equal_ignore_global_symbol(gemm_dyn_shape, gemm_dyn_shape_roundtrip)
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def match_buffer_int64(a: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, (T.int64(128), T.int64(128)), dtype="float32")
|
||||
B = T.sblock_alloc_buffer((T.int64(128), T.int64(128)), dtype="float32")
|
||||
C = T.match_buffer(c, (T.int64(128), T.int64(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(T.int64(128), T.int64(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 match_buffer_int64_after_roundtrip(
|
||||
A: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
||||
C: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
||||
) -> None:
|
||||
B = T.sblock_alloc_buffer((T.int64(128), T.int64(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(T.int64(128), T.int64(128)):
|
||||
with T.sblock("C"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
C[vi, vj] = B[vi, vj] + 1.0
|
||||
|
||||
|
||||
def test_match_buffer_int64():
|
||||
original = match_buffer_int64
|
||||
after_roundtrip = match_buffer_int64_after_roundtrip
|
||||
assert_structural_equal_ignore_global_symbol(original, after_roundtrip, True)
|
||||
|
||||
|
||||
def test_match_buffer_region_has_implicit_shape_dtype():
|
||||
@T.prim_func(s_tir=True)
|
||||
def explicit_shape_dtype(A: T.Buffer((16, 64), "int32")):
|
||||
with T.sblock():
|
||||
B = T.match_buffer(A[8:16, 32:64], shape=(8, 32), dtype="int32")
|
||||
T.evaluate(0)
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def implicit_shape_dtype(A: T.Buffer((16, 64), "int32")):
|
||||
with T.sblock():
|
||||
B = T.match_buffer(A[8:16, 32:64])
|
||||
T.evaluate(0)
|
||||
|
||||
assert_structural_equal_ignore_global_symbol(explicit_shape_dtype, implicit_shape_dtype)
|
||||
|
||||
|
||||
def test_match_buffer_input_requires_shape_arg():
|
||||
with pytest.raises(tvm.error.DiagnosticError):
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func(a: T.handle):
|
||||
A = T.match_buffer(a, dtype="int32")
|
||||
T.evaluate(0)
|
||||
|
||||
|
||||
def test_bind_bufferload_without_type_annotation():
|
||||
# Variable assignment of Expr types uses the dtype of the
|
||||
# Expr to determine the variable's dtype. Parsing of
|
||||
# buf[indices] is done by generating a BufferSlice object, which
|
||||
# handles both store and load cases. BufferSlice is not a
|
||||
# Expr, and implements BufferSlice.dtype explicitly.
|
||||
|
||||
# Failure occurred during parsing of the tvmscript.
|
||||
@T.prim_func(s_tir=True)
|
||||
def func_without_type_annotation(A: T.Buffer((1,), "int32")):
|
||||
x = A[0]
|
||||
T.evaluate(x)
|
||||
|
||||
|
||||
def test_bind_with_constant():
|
||||
@T.prim_func(s_tir=True)
|
||||
def constant_binds():
|
||||
x = T.meta_var(1)
|
||||
y = T.meta_var(42.0)
|
||||
T.evaluate(T.cast(x, "float32") + y)
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def constant_binds_wrapped():
|
||||
x = T.meta_var(T.int32(1))
|
||||
y = T.meta_var(T.float32(42.0))
|
||||
T.evaluate(T.cast(x, "float32") + y)
|
||||
|
||||
assert_structural_equal_ignore_global_symbol(constant_binds, constant_binds_wrapped)
|
||||
|
||||
|
||||
def test_func_call():
|
||||
def shared_16x16_to_ldmatrix_32x8_layout(i, j):
|
||||
thread_id = (i % 8) * 4 + (j % 8) // 2
|
||||
return T.meta_var((thread_id, (j // 8) * 4 + (i // 8) * 2 + (j % 2)))
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def mma_sync_m16n16k16_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, (32, 8), "float16", align=64, offset_factor=16, scope="warp")
|
||||
B = T.match_buffer(b, (32, 8), "float16", align=64, offset_factor=16, scope="warp")
|
||||
C = T.match_buffer(c, (32, 8), "float16", align=64, offset_factor=16, scope="warp")
|
||||
|
||||
with T.sblock("root"):
|
||||
T.reads(C[0:32, 0:8], A[0:32, 0:8], B[0:32, 0:8])
|
||||
T.writes(C[0:32, 0:8])
|
||||
for i, j, k in T.grid(16, 16, 16):
|
||||
with T.sblock("C"):
|
||||
i, j, k = T.axis.remap("SSR", [i, j, k])
|
||||
thread_id_C, local_id_C = shared_16x16_to_ldmatrix_32x8_layout(i, j)
|
||||
thread_id_A, local_id_A = shared_16x16_to_ldmatrix_32x8_layout(i, k)
|
||||
thread_id_B, local_id_B = shared_16x16_to_ldmatrix_32x8_layout(k, j)
|
||||
|
||||
T.reads(
|
||||
C[thread_id_C, local_id_C],
|
||||
A[thread_id_A, local_id_A],
|
||||
B[thread_id_B, local_id_B],
|
||||
)
|
||||
T.writes(C[thread_id_C, local_id_C])
|
||||
|
||||
C[thread_id_C, local_id_C] += (
|
||||
A[thread_id_A, local_id_A] * B[thread_id_B, local_id_B]
|
||||
)
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def mma_sync_m16n16k16_desc_manual(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, (32, 8), "float16", align=64, offset_factor=16, scope="warp")
|
||||
B = T.match_buffer(b, (32, 8), "float16", align=64, offset_factor=16, scope="warp")
|
||||
C = T.match_buffer(c, (32, 8), "float16", align=64, offset_factor=16, scope="warp")
|
||||
|
||||
with T.sblock("root"):
|
||||
T.reads(C[0:32, 0:8], A[0:32, 0:8], B[0:32, 0:8])
|
||||
T.writes(C[0:32, 0:8])
|
||||
for i, j, k in T.grid(16, 16, 16):
|
||||
with T.sblock("C"):
|
||||
i, j, k = T.axis.remap("SSR", [i, j, k])
|
||||
T.reads(
|
||||
C[i % 8 * 4 + j % 8 // 2, j // 8 * 4 + i // 8 * 2 + j % 2],
|
||||
A[i % 8 * 4 + k % 8 // 2, k // 8 * 4 + i // 8 * 2 + k % 2],
|
||||
B[k % 8 * 4 + j % 8 // 2, j // 8 * 4 + k // 8 * 2 + j % 2],
|
||||
)
|
||||
T.writes(C[i % 8 * 4 + j % 8 // 2, j // 8 * 4 + i // 8 * 2 + j % 2])
|
||||
C[i % 8 * 4 + j % 8 // 2, j // 8 * 4 + i // 8 * 2 + j % 2] = (
|
||||
C[i % 8 * 4 + j % 8 // 2, j // 8 * 4 + i // 8 * 2 + j % 2]
|
||||
+ A[i % 8 * 4 + k % 8 // 2, k // 8 * 4 + i // 8 * 2 + k % 2]
|
||||
* B[k % 8 * 4 + j % 8 // 2, j // 8 * 4 + k // 8 * 2 + j % 2]
|
||||
)
|
||||
|
||||
assert_structural_equal_ignore_global_symbol(
|
||||
mma_sync_m16n16k16_desc, mma_sync_m16n16k16_desc_manual
|
||||
)
|
||||
|
||||
# The following is an example of an error message from calling an invalid function
|
||||
|
||||
# error: Error occurred when invoking the function sqrt:
|
||||
# loop of ufunc does not support argument 0 of type Var which has no callable sqrt method
|
||||
# --> test_tvmscript_syntax_sugar.py:334:19
|
||||
# |
|
||||
# 334 | ind = sqrt(i)
|
||||
# | ^^^^^^^
|
||||
# note: run with `TVM_BACKTRACE=1` environment variable to display a backtrace.
|
||||
|
||||
# Uncomment to see the error above.
|
||||
# def sqrt(x):
|
||||
# import numpy as np
|
||||
# return np.sqrt(x)
|
||||
|
||||
# @T.prim_func
|
||||
# def loop(a: T.handle) -> None:
|
||||
# A = T.match_buffer(a, (128,))
|
||||
# for i in T.serial(128):
|
||||
# ind = sqrt(i)
|
||||
# A[i] = A[ind]
|
||||
|
||||
|
||||
def test_int64_loop():
|
||||
@T.prim_func(s_tir=True)
|
||||
def int64_grid(
|
||||
A: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
||||
B: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
||||
) -> None:
|
||||
for i, j in T.grid(T.int64(128), T.int64(128)):
|
||||
with T.sblock("C"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
B[vi, vj] = A[vi, vj] + 1.0
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def int64_grid_expanded(
|
||||
A: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
||||
B: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
||||
) -> None:
|
||||
for i in range(T.int64(0), T.int64(128)):
|
||||
for j in range(T.int64(0), T.int64(128)):
|
||||
with T.sblock("C"):
|
||||
vi = T.axis.spatial(T.int64(128), i)
|
||||
vj = T.axis.spatial(T.int64(128), j)
|
||||
B[vi, vj] = A[vi, vj] + 1.0
|
||||
|
||||
assert_structural_equal_ignore_global_symbol(int64_grid, int64_grid_expanded)
|
||||
|
||||
|
||||
def test_implicit_evaluate_assume():
|
||||
@T.prim_func(s_tir=True)
|
||||
def explicit(A: T.Buffer(1, "int32")):
|
||||
T.evaluate(T.assume(A[0] == 5))
|
||||
A[0] = 10
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def implicit(A: T.Buffer(1, "int32")):
|
||||
T.assume(A[0] == 5)
|
||||
A[0] = 10
|
||||
|
||||
assert_structural_equal_ignore_global_symbol(implicit, explicit)
|
||||
|
||||
|
||||
def test_implicit_evaluate_call_extern():
|
||||
@T.prim_func(s_tir=True)
|
||||
def explicit(A: T.Buffer(1, "int32")):
|
||||
T.evaluate(T.call_extern("extern_func", A.data, dtype="int32"))
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def implicit(A: T.Buffer(1, "int32")):
|
||||
T.call_extern("extern_func", A.data, dtype="int32")
|
||||
|
||||
assert_structural_equal_ignore_global_symbol(implicit, explicit)
|
||||
|
||||
|
||||
def test_preserve_trivial_let_binding():
|
||||
"""Trivial `T.let[...]` annotations survive the parser as LetStmt and are not inlined.
|
||||
|
||||
In fork, bare `j = i` lowers to a local_scalar (AllocBuffer + BufferStore); the
|
||||
LetStmt form is opt-in via `T.let[T.dtype]`. Both the explicit `T.bind(..., var=j)`
|
||||
builder API and the `j: T.let[T.dtype]` annotation produce the same LetStmt IR.
|
||||
"""
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def explicit(i: T.int32):
|
||||
j = T.int32()
|
||||
T.bind(i, var=j)
|
||||
T.evaluate(j)
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def implicit(i: T.int32):
|
||||
j: T.let[T.int32] = i
|
||||
T.evaluate(j)
|
||||
|
||||
assert_structural_equal_ignore_global_symbol(implicit, explicit)
|
||||
|
||||
|
||||
def test_preserve_trivial_let_binding_of_value():
|
||||
"""Same as test_preserve_trivial_let_binding but with a constant RHS."""
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def explicit(i: T.int32):
|
||||
j = T.int32()
|
||||
T.bind(42, var=j)
|
||||
T.evaluate(j)
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def implicit(i: T.int32):
|
||||
j: T.let[T.int32] = 42
|
||||
T.evaluate(j)
|
||||
|
||||
assert_structural_equal_ignore_global_symbol(implicit, explicit)
|
||||
|
||||
|
||||
def test_preserve_parameter_name():
|
||||
@T.prim_func(s_tir=True)
|
||||
def func(i: T.int32):
|
||||
j = i
|
||||
T.evaluate(j)
|
||||
|
||||
param_name = func.params[0].name
|
||||
assert param_name == "i"
|
||||
|
||||
|
||||
def test_preserve_variable_name():
|
||||
"""Use variable name when generating tirx::Bind / AllocBuffer"""
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func():
|
||||
for i in T.serial(16):
|
||||
j = i // 4
|
||||
T.evaluate(j)
|
||||
|
||||
# In fork, bare `j = i // 4` lowers to AllocBuffer (local_scalar) in the for-body
|
||||
# SeqStmt; the variable name lives on the underlying buffer.
|
||||
var_name = func.body.body.seq[0].buffer.name
|
||||
assert var_name == "j"
|
||||
|
||||
|
||||
def test_boolean_constant():
|
||||
"""Python booleans should become T.Bool objects"""
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def explicit():
|
||||
T.evaluate(T.bool(True))
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def implicit():
|
||||
T.evaluate(True)
|
||||
|
||||
assert_structural_equal_ignore_global_symbol(implicit, explicit)
|
||||
|
||||
|
||||
def test_foldable_boolean_in_assert():
|
||||
"""Foldable booleans T.Bool objects
|
||||
|
||||
The condition of an assert statement should be a boolean
|
||||
expression. Previously, this test failed because the FFI does not
|
||||
distinguish between integer primitives and boolean primitives.
|
||||
"""
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def explicit():
|
||||
assert T.bool(False), "Message"
|
||||
T.evaluate(0)
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def implicit():
|
||||
assert 0 == 1, "Message"
|
||||
T.evaluate(0)
|
||||
|
||||
assert_structural_equal_ignore_global_symbol(implicit, explicit)
|
||||
|
||||
|
||||
def test_return_statement():
|
||||
"""A python `return` statement uses `T.ret`"""
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def explicit():
|
||||
T.evaluate(T.ret(5))
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def implicit():
|
||||
return 5
|
||||
|
||||
assert_structural_equal_ignore_global_symbol(implicit, explicit)
|
||||
|
||||
|
||||
def test_loop_jump_statement():
|
||||
"""`break` and `continue` evaluates to TIR intrinsics"""
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def explicit():
|
||||
for i in range(16):
|
||||
if i % 2 == 0:
|
||||
T.evaluate(T.continue_loop())
|
||||
if i < 15:
|
||||
T.evaluate(T.break_loop())
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def implicit():
|
||||
for i in range(16):
|
||||
if i % 2 == 0:
|
||||
continue
|
||||
if i < 15:
|
||||
break
|
||||
|
||||
assert_structural_equal_ignore_global_symbol(implicit, explicit)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,182 @@
|
||||
# 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,invalid-name,pointless-string-statement
|
||||
from tvm.script import tirx as T
|
||||
|
||||
"""
|
||||
This prim func include necessary buffer types that need to be checked
|
||||
e.g. reads/writes, match_buffer/alloc_buffer, serial/block etc.
|
||||
"""
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def element_wise_storage_align(a: T.handle, c: T.handle) -> None:
|
||||
C = T.match_buffer(c, [128, 128], elem_offset=0, align=64, offset_factor=1)
|
||||
A = T.match_buffer(a, [128, 128], elem_offset=0, align=64, offset_factor=1)
|
||||
# body
|
||||
with T.sblock("root"):
|
||||
T.reads([])
|
||||
T.writes([])
|
||||
B = T.sblock_alloc_buffer([128, 128], elem_offset=0, align=64, offset_factor=1)
|
||||
for i0 in T.serial(0, 128):
|
||||
for ax1 in T.serial(0, 128):
|
||||
with T.sblock("B"):
|
||||
vi = T.axis.S(128, i0)
|
||||
vj = T.axis.S(128, ax1)
|
||||
T.reads([A[vi, vj]])
|
||||
T.writes([B[vi, vj]])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 0, 128, 127]]})
|
||||
B[vi, vj] = A[vi, vj] * T.float32(2)
|
||||
for i1 in T.serial(0, 128):
|
||||
with T.sblock("C"):
|
||||
vi_1, vj_1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads([B[vi_1, vj_1]])
|
||||
T.writes([C[vi_1, vj_1]])
|
||||
C[vi_1, vj_1] = B[vi_1, vj_1] + T.float32(1)
|
||||
|
||||
|
||||
"""
|
||||
This prim func include necessary thread types that need to be checked
|
||||
e.g. env_thread, launch_thread, thread_binding etc.
|
||||
"""
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def element_wise_env_thread_x(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
j1_0 = T.env_thread("threadIdx.x")
|
||||
j0_0 = T.env_thread("threadIdx.x")
|
||||
i = T.env_thread("blockIdx.x")
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
B = T.match_buffer(b, [128, 128])
|
||||
C = T.match_buffer(c, [128, 128])
|
||||
T.launch_thread(i, 128)
|
||||
T.launch_thread(j0_0, 4)
|
||||
T.launch_thread(j1_0, 4)
|
||||
|
||||
for blockIdx_x in T.thread_binding(0, 128, "blockIdx.x"):
|
||||
for threadIdx_x in T.thread_binding(0, 4, "threadIdx.x"):
|
||||
for j0_1 in T.serial(0, 32):
|
||||
with T.sblock(""):
|
||||
B[blockIdx_x, threadIdx_x * 32 + j0_1] = (
|
||||
A[blockIdx_x, threadIdx_x * 32 + j0_1] * 2.0
|
||||
)
|
||||
for j1_1 in T.serial(0, 32):
|
||||
with T.sblock(""):
|
||||
C[blockIdx_x, threadIdx_x * 32 + j1_1] = (
|
||||
B[blockIdx_x, threadIdx_x * 32 + j1_1] + 1.0
|
||||
)
|
||||
|
||||
|
||||
"""
|
||||
This test case is added to test T.grid
|
||||
"""
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def loop_split(a: T.handle, b: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128], dtype="float32")
|
||||
B = T.match_buffer(b, [128], dtype="float32")
|
||||
for i, ko in T.grid(128, 4):
|
||||
for ki in T.thread_binding(0, 32, thread="threadIdx.x"):
|
||||
with T.sblock("B"):
|
||||
vi = T.axis.S(128, i)
|
||||
vk = T.axis.R(128, ko * 32 + ki)
|
||||
T.reads([B[vi], A[vi, vk]])
|
||||
T.writes([B[vi]])
|
||||
with T.init():
|
||||
B[vi] = T.float32(0)
|
||||
B[vi] = B[vi] + A[vi, vk]
|
||||
|
||||
|
||||
"""
|
||||
This test case is added to test T.comm_reducer, T.reinterpret, T.tvm_thread_allreduce
|
||||
"""
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def lowered_loop_split(a: T.handle, b: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128], dtype="float32")
|
||||
B = T.match_buffer(b, [128], dtype="float32")
|
||||
reduce_temp0 = T.sblock_alloc_buffer([1], dtype="float32", strides=[1], scope="local")
|
||||
normal_reduce_temp0 = T.sblock_alloc_buffer([1], dtype="float32", strides=[1], scope="local")
|
||||
for i in T.serial(0, 128):
|
||||
for ki in T.thread_binding(0, 32, thread="threadIdx.x"):
|
||||
normal_reduce_temp0[0] = T.float32(0)
|
||||
for ko in T.serial(0, 4):
|
||||
with T.sblock("B_normal_reduction"):
|
||||
vi = T.axis.S(128, i)
|
||||
vk = T.axis.R(128, ko * 32 + ki)
|
||||
T.reads([A[vi, vk], normal_reduce_temp0[0]])
|
||||
T.writes([normal_reduce_temp0[0]])
|
||||
normal_reduce_temp0[0] = normal_reduce_temp0[0] + A[vi, vk]
|
||||
with T.sblock("B_cross_thread_reduction"):
|
||||
T.reads([normal_reduce_temp0[0]])
|
||||
T.writes([reduce_temp0[0]])
|
||||
T.attr(
|
||||
T.comm_reducer(lambda x, y: x + y, [T.float32(0)]),
|
||||
"reduce_scope",
|
||||
T.int32(0),
|
||||
)
|
||||
T.evaluate(
|
||||
T.tvm_thread_allreduce(
|
||||
T.uint32(1),
|
||||
normal_reduce_temp0[0],
|
||||
True,
|
||||
reduce_temp0.data,
|
||||
ki,
|
||||
dtype="handle",
|
||||
)
|
||||
)
|
||||
with T.sblock("B_write_back"):
|
||||
vi = T.axis.S(128, i)
|
||||
T.reads([reduce_temp0[0]])
|
||||
T.writes([B[vi]])
|
||||
B[vi] = reduce_temp0[0]
|
||||
|
||||
|
||||
"""
|
||||
This test case is added to test T.Buffer with slice as argument and T.exp
|
||||
"""
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def different_access_indices(a: T.handle, b: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128, 128], dtype="float32")
|
||||
B = T.match_buffer(b, [128, 128], dtype="float32")
|
||||
for i, j in T.grid(128, 128):
|
||||
for k in T.thread_binding(0, 128, thread="threadIdx.x"):
|
||||
with T.sblock("B"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
T.reads([B[vi, vj], A[vi, vj, vk]])
|
||||
T.writes(
|
||||
[
|
||||
B[
|
||||
T.min(vj, vi) : T.min(vj, vi) # type: ignore[misc]
|
||||
+ (T.max(vj, vi) + 1 - T.min(vj, vi)),
|
||||
T.min(vi, vj) : T.min(vi, vj) # type: ignore[misc]
|
||||
+ (T.max(vi, vj) + 1 - T.min(vi, vj)),
|
||||
]
|
||||
]
|
||||
)
|
||||
with T.init():
|
||||
B[vj, vi] = T.exp(B[vj, vi], dtype="float32")
|
||||
B[vi, vj] = B[vi, vj] + A[vi, vj, vk]
|
||||
|
||||
|
||||
# Not running any test as we only want to type-check here
|
||||
if __name__ == "__main__":
|
||||
pass
|
||||
Reference in New Issue
Block a user