wehub-resource-sync
1116 lines
28 KiB
Python
1116 lines
28 KiB
Python
# 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: E501, F401, F841
|
|
|
|
import re
|
|
|
|
import pytest
|
|
|
|
import tvm.testing
|
|
from tvm import ir, tirx
|
|
from tvm.ir import Range
|
|
from tvm.script.ir_builder import IRBuilder
|
|
from tvm.script.ir_builder import tirx as T
|
|
|
|
|
|
def _assert_print(obj, expected):
|
|
assert obj.script(verbose_expr=True).strip() == expected.strip()
|
|
|
|
|
|
def test_prim_func():
|
|
a = tirx.Var("a", "handle")
|
|
b = tirx.Var("b", "handle")
|
|
func = (
|
|
tirx.PrimFunc(
|
|
params=[a, b],
|
|
ret_type=None,
|
|
buffer_map={
|
|
a: tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A"),
|
|
b: tirx.decl_buffer(shape=[256, 256], dtype="float32", name="B"),
|
|
},
|
|
body=tirx.Evaluate(0),
|
|
)
|
|
.with_attr("global_symbol", "main")
|
|
.with_attr("s_tir", True)
|
|
)
|
|
_assert_print(
|
|
func,
|
|
expected="""
|
|
# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")):
|
|
T.evaluate(0)""",
|
|
)
|
|
|
|
|
|
def test_prim_func_no_sugar_inlined_buffer():
|
|
a = tirx.Var("a", "handle")
|
|
b = tirx.Var("b", "handle")
|
|
func = (
|
|
tirx.PrimFunc(
|
|
params=[a, b],
|
|
ret_type=None,
|
|
buffer_map={
|
|
a: tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A"),
|
|
b: tirx.decl_buffer(shape=[256, 256], dtype="float32", name="B"),
|
|
},
|
|
body=tirx.Evaluate(a),
|
|
)
|
|
.with_attr("global_symbol", "main")
|
|
.with_attr("s_tir", True)
|
|
)
|
|
_assert_print(
|
|
func,
|
|
expected="""
|
|
# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main(a: T.handle, B: T.Buffer((256, 256), "float32")):
|
|
A = T.match_buffer(a, (128, 128))
|
|
T.evaluate(a)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_prim_func_no_sugar_shared_buffer_data():
|
|
a = tirx.Var("a", "handle")
|
|
b = tirx.Var("b", "handle")
|
|
buffer_data = tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A").data
|
|
func = (
|
|
tirx.PrimFunc(
|
|
params=[a, b],
|
|
ret_type=None,
|
|
buffer_map={
|
|
a: tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A", data=buffer_data),
|
|
b: tirx.decl_buffer(shape=[256, 256], dtype="float32", name="B", data=buffer_data),
|
|
},
|
|
body=tirx.Evaluate(0),
|
|
)
|
|
.with_attr("global_symbol", "main")
|
|
.with_attr("s_tir", True)
|
|
)
|
|
_assert_print(
|
|
func,
|
|
expected="""
|
|
# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main(a: T.handle, b: T.handle):
|
|
A = T.match_buffer(a, (128, 128))
|
|
B = T.match_buffer(b, (256, 256), data=A.data)
|
|
T.evaluate(0)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_block_realize():
|
|
i = tirx.Var("i", "int32")
|
|
j = tirx.Var("j", "int32")
|
|
k = tirx.Var("k", "int32")
|
|
with IRBuilder() as ib:
|
|
with T.sblock(name="block", no_realize=False):
|
|
vi = ib.name("vi", T.axis.spatial(128, i))
|
|
vj = ib.name("vj", T.axis.spatial(64, j))
|
|
vk = ib.name("vk", T.axis.reduce(32, k))
|
|
T.reads()
|
|
T.writes()
|
|
T.evaluate(0)
|
|
obj = ib.get()
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
i = T.int32()
|
|
j = T.int32()
|
|
k = T.int32()
|
|
with T.sblock("block"):
|
|
vi = T.axis.spatial(128, i)
|
|
vj = T.axis.spatial(64, j)
|
|
vk = T.axis.reduce(32, k)
|
|
T.reads()
|
|
T.writes()
|
|
T.evaluate(0)""",
|
|
)
|
|
|
|
|
|
def test_block():
|
|
i = tirx.Var("i", "int32")
|
|
j = tirx.Var("j", "int32")
|
|
k = tirx.Var("k", "int32")
|
|
with IRBuilder() as ib:
|
|
with T.sblock(name="block", no_realize=False):
|
|
vi = ib.name("vi", T.axis.spatial(128, i))
|
|
vj = ib.name("vj", T.axis.spatial(64, j))
|
|
vk = ib.name("vk", T.axis.reduce(32, k))
|
|
T.reads()
|
|
T.writes()
|
|
T.evaluate(0)
|
|
obj = ib.get().block
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
with T.sblock("block", no_realize=True):
|
|
vi = T.axis.spatial(128)
|
|
vj = T.axis.spatial(64)
|
|
vk = T.axis.reduce(32)
|
|
T.reads()
|
|
T.writes()
|
|
T.evaluate(0)""",
|
|
)
|
|
|
|
|
|
def test_match_buffer_region():
|
|
src = tirx.decl_buffer((128, 128), "float32", name="src")
|
|
tgt = tirx.decl_buffer((64, 64), "float32", name="tgt")
|
|
obj = tirx.MatchBufferRegion(
|
|
tgt,
|
|
tirx.BufferRegion(
|
|
src,
|
|
[
|
|
Range(64, 128),
|
|
Range(64, 128),
|
|
],
|
|
),
|
|
)
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
src = T.Buffer((128, 128))
|
|
tgt = T.match_buffer(src[64:128, 64:128], (64, 64))
|
|
""",
|
|
)
|
|
|
|
|
|
def test_buffer():
|
|
a = tirx.decl_buffer((128, 128), "float16", name="A")
|
|
_assert_print(
|
|
a,
|
|
"""A = T.Buffer((128, 128), "float16")
|
|
A""",
|
|
)
|
|
|
|
|
|
def test_buffer_region():
|
|
src = tirx.decl_buffer((128, 128), "float32", name="src")
|
|
obj = tirx.BufferRegion(
|
|
src,
|
|
[
|
|
Range(64, 128),
|
|
Range(64, 128),
|
|
],
|
|
)
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
src = T.Buffer((128, 128))
|
|
src[64:128, 64:128]
|
|
""",
|
|
)
|
|
|
|
|
|
def test_buffer_load():
|
|
a = tirx.decl_buffer((128, 128), "float16", name="A")
|
|
obj = tirx.BufferLoad(a, [128, 128])
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
A = T.Buffer((128, 128), "float16")
|
|
A[128, 128]
|
|
""",
|
|
)
|
|
|
|
|
|
def test_buffer_store():
|
|
a = tirx.decl_buffer((128, 128), "float16", name="A")
|
|
with IRBuilder() as ib:
|
|
T.buffer_store(a, a[128, 128] + 1, [128, 128])
|
|
obj = ib.get()
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
A = T.Buffer((128, 128), "float16")
|
|
A[128, 128] = A[128, 128] + T.float16(1.0)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_for():
|
|
with IRBuilder() as ib:
|
|
with T.grid(128, 128, 128) as (i, j, k):
|
|
ib.name_many(["i", "j", "k"], [i, j, k])
|
|
T.evaluate(0)
|
|
obj = ib.get()
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
for i, j, k in T.grid(128, 128, 128):
|
|
T.evaluate(0)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_bind():
|
|
with IRBuilder() as ib:
|
|
with T.prim_func(s_tir=True):
|
|
v = T.bind(T.float32(10))
|
|
ib.name("v", v)
|
|
T.evaluate(1)
|
|
obj = ib.get()
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def main():
|
|
v: T.let[T.float32] = T.float32(10.0)
|
|
T.evaluate(1)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_attr_stmt():
|
|
with IRBuilder() as ib:
|
|
with T.attr("pragma", "unroll", 1):
|
|
T.evaluate(0)
|
|
obj = ib.get()
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
with T.attr("pragma", "unroll", 1):
|
|
T.evaluate(0)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_assert_stmt():
|
|
with IRBuilder() as ib:
|
|
with T.Assert(True, "assertion"):
|
|
T.evaluate(0)
|
|
obj = ib.get()
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
assert T.bool(True), ("RuntimeError", ["assertion"])
|
|
T.evaluate(0)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_while():
|
|
with IRBuilder() as ib:
|
|
x = T.int32()
|
|
with T.While(x < 10):
|
|
T.evaluate(0)
|
|
obj = ib.get()
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
v = T.int32()
|
|
while v < 10:
|
|
T.evaluate(0)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_allocate():
|
|
with IRBuilder() as ib:
|
|
with T.prim_func():
|
|
T.func_name("test")
|
|
buf = T.alloc_buffer([128, 128], "float32")
|
|
T.evaluate(1)
|
|
obj = ib.get()
|
|
_assert_print(
|
|
obj.body,
|
|
"""
|
|
buffer = T.alloc_buffer((128, 128))
|
|
T.evaluate(1)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_allocate_with_decl_buffer_sugar():
|
|
# AllocBuffer and DeclBuffer are flat siblings
|
|
with IRBuilder() as ib:
|
|
with T.prim_func():
|
|
T.func_name("test")
|
|
buf = T.alloc_buffer([128, 128], "float32")
|
|
buf2 = T.decl_buffer([128, 128], "float32", data=buf.data)
|
|
T.evaluate(1)
|
|
obj = ib.get()
|
|
_assert_print(
|
|
obj.body,
|
|
"""
|
|
buffer = T.alloc_buffer((128, 128))
|
|
buffer_1 = T.decl_buffer((128, 128), data=buffer.data)
|
|
T.evaluate(1)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_allocate_with_decl_buffer_sugar_multi_usage():
|
|
# AllocBuffer and DeclBuffer are flat siblings
|
|
with IRBuilder() as ib:
|
|
with T.prim_func():
|
|
T.func_name("test")
|
|
buf = T.alloc_buffer([128, 128], "float32")
|
|
buf2 = T.decl_buffer([128, 128], "float32", data=buf.data)
|
|
T.evaluate(buf.data)
|
|
obj = ib.get()
|
|
_assert_print(
|
|
obj.body,
|
|
"""
|
|
buffer = T.alloc_buffer((128, 128))
|
|
buffer_1 = T.decl_buffer((128, 128), data=buffer.data)
|
|
T.evaluate(buffer.data)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_allocate_with_decl_buffer_no_sugar_mismatch():
|
|
with IRBuilder() as ib:
|
|
with T.prim_func():
|
|
T.func_name("test")
|
|
buf = T.alloc_buffer([128, 128], "float32")
|
|
buf2 = T.decl_buffer([256, 256], "float32", data=buf.data)
|
|
T.evaluate(buf.data)
|
|
obj = ib.get()
|
|
_assert_print(
|
|
obj.body,
|
|
"""
|
|
buffer = T.alloc_buffer((128, 128))
|
|
buffer_1 = buffer.view(256, 256)
|
|
T.evaluate(buffer.data)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_decl_buffer():
|
|
# DeclBuffer is flat: we need a frame to hold multiple stmts
|
|
with IRBuilder() as ib:
|
|
with T.prim_func():
|
|
T.func_name("test")
|
|
buf = T.decl_buffer((10, 10), data=T.ptr("float32"))
|
|
T.evaluate(1)
|
|
obj = ib.get()
|
|
# Print only the body (skip PrimFunc wrapper)
|
|
_assert_print(
|
|
obj.body,
|
|
"""
|
|
v = T.handle("float32", "global")
|
|
buffer = T.decl_buffer((10, 10), data=v)
|
|
T.evaluate(1)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_seq_stmt():
|
|
with IRBuilder() as ib:
|
|
with T.serial(10):
|
|
T.evaluate(1)
|
|
T.evaluate(2)
|
|
obj = ib.get().body
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
T.evaluate(1)
|
|
T.evaluate(2)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_if_then_else():
|
|
with IRBuilder() as ib:
|
|
with T.If(T.int32() == 1):
|
|
with T.Then():
|
|
T.evaluate(0)
|
|
|
|
obj = ib.get()
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
v = T.int32()
|
|
if v == 1:
|
|
T.evaluate(0)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_evaluate():
|
|
with IRBuilder() as ib:
|
|
T.evaluate(0)
|
|
obj = ib.get()
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
T.evaluate(0)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_var():
|
|
a = tirx.Var("a", "float32")
|
|
_assert_print(
|
|
a,
|
|
"""
|
|
a = T.float32()
|
|
a""",
|
|
)
|
|
|
|
a = tirx.Var("a", "handle")
|
|
_assert_print(
|
|
a,
|
|
"""
|
|
a = T.handle()
|
|
a""",
|
|
)
|
|
|
|
a = tirx.Var("a", ir.PointerType(ir.PrimType("void"), "shared"))
|
|
_assert_print(
|
|
a,
|
|
"""
|
|
a = T.handle(storage_scope="shared")
|
|
a""",
|
|
)
|
|
|
|
|
|
def test_iter_var():
|
|
a = tirx.IterVar((0, 8), "a", iter_type=tirx.IterVar.DataPar)
|
|
_assert_print(
|
|
a,
|
|
"""
|
|
a = T.int32()
|
|
T.iter_var(a, T.Range(0, 8), "DataPar", "")
|
|
""",
|
|
)
|
|
|
|
|
|
def test_string_imm():
|
|
s = tirx.StringImm("str")
|
|
_assert_print(s, '"str"')
|
|
|
|
|
|
def test_cast():
|
|
obj = tirx.Cast("float64", tirx.Var("a", "float32"))
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
a = T.float32()
|
|
T.Cast("float64", a)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_llvm_intrin_imm():
|
|
a = tirx.call_llvm_intrin("int32x4", "llvm.donothing")
|
|
_assert_print(a, 'T.call_llvm_intrin("int32x4", "llvm.donothing")')
|
|
a = tirx.call_llvm_pure_intrin("int32x4", "llvm.donothing")
|
|
_assert_print(a, 'T.call_llvm_pure_intrin("int32x4", "llvm.donothing")')
|
|
|
|
|
|
def test_binary_arith():
|
|
a = tirx.Var("a", "int32")
|
|
b = tirx.Var("b", "int32")
|
|
for op, sign in [
|
|
(tirx.Add, "+"),
|
|
(tirx.Sub, "-"),
|
|
(tirx.Mul, "*"),
|
|
(tirx.Mod, "truncmod"),
|
|
(tirx.FloorDiv, "//"),
|
|
(tirx.FloorMod, "%"),
|
|
(tirx.LT, "<"),
|
|
(tirx.LE, "<="),
|
|
(tirx.EQ, "=="),
|
|
(tirx.NE, "!="),
|
|
(tirx.GT, ">"),
|
|
(tirx.GE, ">="),
|
|
]:
|
|
obj = op(a, b)
|
|
if sign.isalpha():
|
|
expected = f"""
|
|
a = T.int32()
|
|
b = T.int32()
|
|
T.{sign}(a, b)"""
|
|
else:
|
|
expected = f"""
|
|
a = T.int32()
|
|
b = T.int32()
|
|
a {sign} b"""
|
|
_assert_print(obj, expected)
|
|
|
|
|
|
def test_binary_arith_const():
|
|
a = tirx.IntImm("int64", 3)
|
|
b = tirx.IntImm("int64", 4)
|
|
for op, name in [
|
|
(tirx.Add, "Add"),
|
|
(tirx.Sub, "Sub"),
|
|
(tirx.Mul, "Mul"),
|
|
(tirx.Div, "Div"),
|
|
(tirx.Mod, "truncmod"),
|
|
(tirx.FloorDiv, "FloorDiv"),
|
|
(tirx.FloorMod, "FloorMod"),
|
|
(tirx.LT, "LT"),
|
|
(tirx.LE, "LE"),
|
|
(tirx.EQ, "EQ"),
|
|
(tirx.NE, "NE"),
|
|
(tirx.GT, "GT"),
|
|
(tirx.GE, "GE"),
|
|
]:
|
|
obj = op(a, b)
|
|
expected = f"""
|
|
T.{name}({a!s}, {b!s})"""
|
|
_assert_print(obj, expected)
|
|
|
|
|
|
def test_int_div():
|
|
a = tirx.Var("a", "int32")
|
|
b = tirx.Var("b", "int32")
|
|
_assert_print(
|
|
tirx.Div(a, b),
|
|
"""
|
|
a = T.int32()
|
|
b = T.int32()
|
|
T.Div(a, b)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_logical():
|
|
a = tirx.Var("a", "bool")
|
|
b = tirx.Var("b", "bool")
|
|
_assert_print(
|
|
tirx.And(a, b),
|
|
"""
|
|
a = T.bool()
|
|
b = T.bool()
|
|
a and b
|
|
""",
|
|
)
|
|
_assert_print(
|
|
tirx.Or(a, b),
|
|
"""
|
|
a = T.bool()
|
|
b = T.bool()
|
|
a or b
|
|
""",
|
|
)
|
|
_assert_print(
|
|
tirx.Not(a),
|
|
"""
|
|
a = T.bool()
|
|
not a
|
|
""",
|
|
)
|
|
|
|
|
|
def test_select():
|
|
obj = tirx.Select(True, 0, 2)
|
|
_assert_print(
|
|
obj,
|
|
"""T.Select(T.bool(True), 0, 2)
|
|
""",
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"lanes, scripted_lanes", [(32, "32"), (tvm.tirx.vscale() * 8, "T.vscale() * 8")]
|
|
)
|
|
def test_ramp(lanes, scripted_lanes):
|
|
a = tirx.Var("a", "int32")
|
|
obj = tirx.Ramp(a, 1, lanes)
|
|
_assert_print(
|
|
obj,
|
|
f"""
|
|
a = T.int32()
|
|
T.Ramp(a, 1, {scripted_lanes})
|
|
""",
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"lanes, scripted_lanes", [(4, "4"), (tvm.tirx.vscale() * 4, "T.vscale() * 4")]
|
|
)
|
|
def test_broadcast(lanes, scripted_lanes):
|
|
obj = tirx.Broadcast(0, lanes)
|
|
_assert_print(
|
|
obj,
|
|
f"""
|
|
T.Broadcast(0, {scripted_lanes})
|
|
""",
|
|
)
|
|
|
|
|
|
def test_let_expr():
|
|
x = tirx.Var("x", "int32")
|
|
obj = tirx.Let(x, 1, x + 1)
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
x = T.int32()
|
|
T.Let(x + 1, where={x: 1})
|
|
""",
|
|
)
|
|
|
|
|
|
def test_call():
|
|
obj = tirx.atan(T.float32(1.0))
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
T.atan(T.float32(1.0))
|
|
""",
|
|
)
|
|
|
|
|
|
def test_comm_reducer():
|
|
obj = T.comm_reducer(lambda x, y: x + y, identity=[T.float32(0)])
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
T.comm_reducer(lambda x, y: x + y, [T.float32(0.0)])
|
|
""",
|
|
)
|
|
|
|
|
|
def test_int_imm():
|
|
obj = T.int16(1)
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
T.int16(1)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_float_imm():
|
|
obj = T.float16(1)
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
T.float16(1.0)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_range():
|
|
obj = Range(0, 10)
|
|
_assert_print(
|
|
obj,
|
|
"""
|
|
I.Range(0, 10)
|
|
""",
|
|
)
|
|
|
|
|
|
def test_prim_type():
|
|
obj = ir.PrimType("float32")
|
|
_assert_print(obj, "T.float32")
|
|
|
|
|
|
def test_pointer_type():
|
|
obj = ir.PointerType(ir.PrimType("int32"), "global")
|
|
_assert_print(obj, 'T.handle("int32", "global")')
|
|
|
|
obj = ir.PointerType(ir.PrimType("void"))
|
|
_assert_print(obj, "T.handle")
|
|
|
|
obj = ir.PointerType(ir.PrimType("void"), "shared")
|
|
_assert_print(obj, 'T.handle(storage_scope="shared")')
|
|
|
|
|
|
def test_tuple_type():
|
|
obj = ir.TupleType([ir.PrimType("float32"), ir.PrimType("int32")])
|
|
_assert_print(obj, "T.Tuple(T.float32, T.int32)")
|
|
|
|
|
|
def test_remap():
|
|
from tvm.script import tirx as T
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def block_with_remap_implicitly():
|
|
for i0, i1, i2, i3, i4, i5 in T.grid(128, 128, 128, 128, 128, 128):
|
|
with T.sblock("update"):
|
|
v0 = T.axis.spatial(128, i0 + 1)
|
|
v1 = T.axis.spatial(128, i1)
|
|
v2 = T.axis.reduce(128, i2)
|
|
v3 = T.axis.spatial(128, i3 - 1)
|
|
v4 = T.axis.reduce(128, i4)
|
|
v5 = T.axis.spatial(128, i5)
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def block_with_remap_explicitly():
|
|
for i0, i1, i2, i3, i4, i5 in T.grid(128, 128, 128, 128, 128, 128):
|
|
with T.sblock("update"):
|
|
v0 = T.axis.spatial(128, i0 + 1)
|
|
v1, v2 = T.axis.remap("SR", [i1, i2])
|
|
v3 = T.axis.spatial(128, i3 - 1)
|
|
v4, v5 = T.axis.remap("RS", [i4, i5])
|
|
|
|
expected_output = """
|
|
# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main():
|
|
# with T.sblock("root"):
|
|
for i0, i1, i2, i3, i4, i5 in T.grid(128, 128, 128, 128, 128, 128):
|
|
with T.sblock("update"):
|
|
v0 = T.axis.spatial(128, i0 + 1)
|
|
v1, v2 = T.axis.remap("SR", [i1, i2])
|
|
v3 = T.axis.spatial(128, i3 - 1)
|
|
v4, v5 = T.axis.remap("RS", [i4, i5])
|
|
T.reads()
|
|
T.writes()
|
|
T.evaluate(0)"""
|
|
_assert_print(block_with_remap_explicitly.with_attr("global_symbol", "main"), expected_output)
|
|
_assert_print(block_with_remap_implicitly.with_attr("global_symbol", "main"), expected_output)
|
|
|
|
|
|
def test_root_block():
|
|
from tvm.script import tirx as T
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def root_block_implicitly():
|
|
a = T.sblock_alloc_buffer([128, 128])
|
|
for i, j in T.grid(128, 128):
|
|
with T.sblock():
|
|
T.evaluate(0)
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def root_block_explicitly():
|
|
with T.sblock("root"):
|
|
a = T.sblock_alloc_buffer([128, 128])
|
|
for i, j in T.grid(128, 128):
|
|
with T.sblock():
|
|
T.evaluate(0)
|
|
|
|
expected_output = """
|
|
# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main():
|
|
# with T.sblock("root"):
|
|
a = T.sblock_alloc_buffer((128, 128))
|
|
for i, j in T.grid(128, 128):
|
|
with T.sblock(""):
|
|
T.reads()
|
|
T.writes()
|
|
T.evaluate(0)
|
|
"""
|
|
_assert_print(root_block_implicitly.with_attr("global_symbol", "main"), expected_output)
|
|
_assert_print(root_block_explicitly.with_attr("global_symbol", "main"), expected_output)
|
|
|
|
|
|
def test_private_primfunc():
|
|
from tvm.script import tirx as T
|
|
|
|
a = tirx.Var("a", "handle")
|
|
b = tirx.Var("b", "handle")
|
|
func = tirx.PrimFunc(
|
|
params=[a, b],
|
|
ret_type=None,
|
|
buffer_map={
|
|
a: tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A"),
|
|
b: tirx.decl_buffer(shape=[256, 256], dtype="float32", name="B"),
|
|
},
|
|
body=tirx.Evaluate(0),
|
|
).with_attr("s_tir", True)
|
|
_assert_print(
|
|
func,
|
|
expected="""
|
|
# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")):
|
|
T.evaluate(0)""",
|
|
)
|
|
|
|
|
|
def test_prim_func_different_symbol():
|
|
from tvm.script import tirx as T
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")):
|
|
T.func_attr({"global_symbol": "func"})
|
|
T.evaluate(0)
|
|
|
|
expected_output = """
|
|
# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def func(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")):
|
|
T.evaluate(0)
|
|
"""
|
|
_assert_print(main, expected_output)
|
|
|
|
|
|
def test_variable_with_cpp_address():
|
|
"""The show_object_address option displays the C++ addressess
|
|
|
|
Because the C++ address may vary with each execution, the output
|
|
produced with this option cannot be compared to a fixed string.
|
|
Instead, this test uses the normal script output to generate a
|
|
regular expression against with the test output must match. The
|
|
regular expression validates that all names have been appended
|
|
with "_0x" followed by a hexadecimal number, and that the address
|
|
is the same for each variable.
|
|
"""
|
|
from tvm.script import tirx as T
|
|
|
|
# The test function has all named objects suffixed with "_name",
|
|
# to avoid spurious replacement when generating the expected
|
|
# regex.
|
|
@T.prim_func(s_tir=True)
|
|
def func(a_name: T.handle):
|
|
N_name = T.int64()
|
|
A_name = T.match_buffer(a_name, N_name, "float32")
|
|
for i_name in range(N_name):
|
|
A_name[i_name] = A_name[i_name] + 1.0
|
|
|
|
without_address = func.script(show_object_address=False)
|
|
script = func.script(show_object_address=True)
|
|
|
|
expected_regex = re.escape(without_address)
|
|
for name in ["a_name", "A_name", "N_name", "i_name"]:
|
|
# Replace all occurrences with a backref to an earlier match
|
|
expected_regex = expected_regex.replace(name, rf"(?P={name})")
|
|
# Then replace the first such backref with a capturing group.
|
|
expected_regex = expected_regex.replace(
|
|
rf"(?P={name})", rf"(?P<{name}>{name}_0x[A-Fa-f0-9]+)", 1
|
|
)
|
|
|
|
assert re.match(expected_regex, script)
|
|
|
|
|
|
def test_return_statement():
|
|
from tvm.script import tirx as T
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def func():
|
|
T.evaluate(T.ret(5))
|
|
|
|
expected_output = """
|
|
# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def func():
|
|
return 5
|
|
"""
|
|
_assert_print(func, expected_output)
|
|
|
|
|
|
CUSTOM_FLOAT_DTYPES = [
|
|
# Float8 variants
|
|
"float8_e3m4",
|
|
"float8_e4m3",
|
|
"float8_e4m3b11fnuz",
|
|
"float8_e4m3fn",
|
|
"float8_e4m3fnuz",
|
|
"float8_e5m2",
|
|
"float8_e5m2fnuz",
|
|
"float8_e8m0fnu",
|
|
# Float6 variants
|
|
"float6_e2m3fn",
|
|
"float6_e3m2fn",
|
|
# Float4 variant
|
|
"float4_e2m1fn",
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", CUSTOM_FLOAT_DTYPES)
|
|
def test_custom_float_types(dtype):
|
|
from tvm.script import tirx as T
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def func():
|
|
T.evaluate(getattr(T, dtype)(0.0))
|
|
|
|
expected_output = f"""
|
|
# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def func():
|
|
T.evaluate(T.{dtype}(0.0))
|
|
"""
|
|
_assert_print(func, expected_output)
|
|
|
|
|
|
def test_predicated_load_store():
|
|
from tvm.script import tirx as T
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main(a: T.handle, b: T.handle):
|
|
A = T.match_buffer(a, (128, 128), "float32")
|
|
B = T.match_buffer(b, (256, 256), "float32")
|
|
T.func_attr({"global_symbol": "func"})
|
|
a_load = T.meta_var(A.vload([0, T.Ramp(0, 4, 4)], predicate=T.Broadcast(T.bool(False), 4)))
|
|
A.vstore([0, T.Ramp(0, 2, 4)], a_load, predicate=T.Broadcast(T.bool(False), 4))
|
|
|
|
expected_output = """
|
|
# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def func(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")):
|
|
A.vstore([0, T.Ramp(0, 2, 4)], A.vload([0, T.Ramp(0, 4, 4)], predicate=T.Broadcast(T.bool(False), 4)), predicate=T.Broadcast(T.bool(False), 4))
|
|
"""
|
|
_assert_print(main, expected_output)
|
|
|
|
|
|
def test_predicated_buffer_load_store():
|
|
a = tirx.Var("a", "handle")
|
|
b = tirx.Var("b", "handle")
|
|
buffer_map = {
|
|
a: tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A"),
|
|
b: tirx.decl_buffer(shape=[256, 256], dtype="float32", name="B"),
|
|
}
|
|
buffer_load = tirx.BufferLoad(
|
|
buffer=buffer_map[b],
|
|
indices=[0, tirx.Ramp(0, 4, 4)],
|
|
predicate=tirx.Broadcast(tirx.IntImm("bool", 0), 4),
|
|
)
|
|
body = tirx.BufferStore(
|
|
buffer=buffer_map[a],
|
|
value=buffer_load,
|
|
indices=[0, tirx.Ramp(0, 2, 4)],
|
|
predicate=tirx.Broadcast(tirx.IntImm("bool", 0), 4),
|
|
)
|
|
func = tirx.PrimFunc(
|
|
params=[a, b],
|
|
ret_type=None,
|
|
buffer_map=buffer_map,
|
|
body=body,
|
|
).with_attr("s_tir", True)
|
|
|
|
expected_output = """
|
|
# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")):
|
|
A.vstore([0, T.Ramp(0, 2, 4)], B.vload([0, T.Ramp(0, 4, 4)], predicate=T.Broadcast(T.bool(False), 4)), predicate=T.Broadcast(T.bool(False), 4))
|
|
"""
|
|
_assert_print(func, expected_output)
|
|
|
|
|
|
def test_predicated_scalable_load_store():
|
|
from tvm.script import tirx as T
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main(a: T.handle, b: T.handle):
|
|
A = T.match_buffer(a, (128, 128), "float32")
|
|
B = T.match_buffer(b, (256, 256), "float32")
|
|
T.func_attr({"global_symbol": "func"})
|
|
mask = T.meta_var(T.get_active_lane_mask("uint1xvscalex4", 0, 13))
|
|
a_load = T.meta_var(A.vload([0, T.Ramp(0, 4, T.vscale() * 4)], predicate=mask))
|
|
A.vstore([0, T.Ramp(0, 2, T.vscale() * 4)], a_load, predicate=mask)
|
|
|
|
expected_output = """
|
|
# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def func(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")):
|
|
A.vstore([0, T.Ramp(0, 2, T.vscale() * 4)], A.vload([0, T.Ramp(0, 4, T.vscale() * 4)], predicate=T.get_active_lane_mask("uint1xvscalex4", 0, 13)), predicate=T.get_active_lane_mask("uint1xvscalex4", 0, 13))
|
|
"""
|
|
_assert_print(main, expected_output)
|
|
|
|
|
|
def test_vload_with_explicit_scalable_data_type():
|
|
from tvm.script import tirx as T
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main(a: T.handle, b: T.handle):
|
|
A = T.match_buffer(a, (128,), "float32")
|
|
B = T.match_buffer(b, (128,), "float32")
|
|
B[0 : T.vscale() * 4] = A.vload([T.Ramp(0, 1, T.vscale() * 4)], dtype="float32xvscalex4")
|
|
|
|
expected_output = """
|
|
# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")):
|
|
B[0:T.vscale() * 4] = A[0:T.vscale() * 4]
|
|
"""
|
|
_assert_print(main, expected_output)
|
|
|
|
|
|
def test_vectorize_llvm_pure_intrin():
|
|
from tvm.script import tirx as T
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main(a: T.handle, b: T.handle):
|
|
A = T.match_buffer(a, (4,), "float32")
|
|
B = T.match_buffer(b, (4,), "float32")
|
|
A[T.Ramp(0, 1, 4)] = T.call_llvm_pure_intrin("float32x4", "llvm.sqrt", B[T.Ramp(0, 1, 4)])
|
|
|
|
expected_output = """
|
|
# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((4,), "float32"), B: T.Buffer((4,), "float32")):
|
|
A[0:4] = T.call_llvm_pure_intrin("float32x4", "llvm.sqrt", B[0:4])
|
|
"""
|
|
_assert_print(main, expected_output)
|
|
|
|
|
|
def test_func_with_loop_jumps():
|
|
from tvm.script import tirx as T
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main(a: T.handle, b: T.handle):
|
|
A = T.match_buffer(a, (4,), "float32")
|
|
B = T.match_buffer(b, (4,), "float32")
|
|
for i in range(1000):
|
|
if i % 13 == 0:
|
|
A[1] = A[1] + 1
|
|
continue
|
|
if A[0] >= B[0]:
|
|
break
|
|
|
|
expected_output = """
|
|
# from tvm.script import tirx as T
|
|
# from tvm.tirx.layout import Axis
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((4,), "float32"), B: T.Buffer((4,), "float32")):
|
|
for i in range(1000):
|
|
if i % 13 == 0:
|
|
A[1] = A[1] + T.float32(1.0)
|
|
T.continue_loop()
|
|
if A[0] >= B[0]:
|
|
T.break_loop()
|
|
"""
|
|
_assert_print(main, expected_output)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|