wehub-resource-sync
1304 lines
46 KiB
Python
1304 lines
46 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.
|
|
# ruff: noqa: E501, E731, E741, F841
|
|
import math
|
|
import re
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm.script import ir as I
|
|
from tvm.script import tirx as T
|
|
from tvm.support import clang, utils
|
|
from tvm.target.codegen import llvm_get_intrinsic_name, llvm_lookup_intrinsic_id
|
|
from tvm.testing import env
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_duplicate_primfunc_global_symbol_diagnostic():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def first_unique_key(A: T.Buffer((1,), "float32")):
|
|
T.func_attr({"global_symbol": "dup_symbol", "tirx.noalias": True})
|
|
A[0] = T.float32(1)
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def second_unique_key(A: T.Buffer((1,), "float32")):
|
|
T.func_attr({"global_symbol": "dup_symbol", "tirx.noalias": True})
|
|
A[0] = T.float32(2)
|
|
|
|
with pytest.raises(
|
|
tvm.error.InternalError, match="Duplicate PrimFunc global_symbol 'dup_symbol'"
|
|
) as err:
|
|
tvm.compile(Module, target="llvm")
|
|
assert "first_unique_key" in str(err.value)
|
|
assert "second_unique_key" in str(err.value)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_unique_primfunc_global_symbols_compile():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def first_unique_key(A: T.Buffer((1,), "float32")):
|
|
T.func_attr({"global_symbol": "dup_symbol_a", "tirx.noalias": True})
|
|
A[0] = T.float32(1)
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def second_unique_key(A: T.Buffer((1,), "float32")):
|
|
T.func_attr({"global_symbol": "dup_symbol_b", "tirx.noalias": True})
|
|
A[0] = T.float32(2)
|
|
|
|
tvm.compile(Module, target="llvm")
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_intrin():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.handle("float32")):
|
|
A_buf = T.decl_buffer((4,), "float32", data=A)
|
|
T.evaluate(T.Call("tirx.prefetch", [T.address_of(A_buf[0]), 0, 3, 1], ret_ty="void"))
|
|
|
|
fcode = tvm.compile(Module)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_void_intrin():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.handle("uint8")):
|
|
# Create an intrinsic that returns void.
|
|
T.call_llvm_intrin("", "llvm.assume", T.bool(True))
|
|
|
|
fcode = tvm.compile(Module)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_intrinsic_id():
|
|
orig_name = "llvm.x86.sse2.pmadd.wd"
|
|
intrin_id = llvm_lookup_intrinsic_id(orig_name)
|
|
name = llvm_get_intrinsic_name(intrin_id)
|
|
assert orig_name == name
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_overloaded_intrin():
|
|
# Name lookup for overloaded intrinsics in LLVM 4- requires a name
|
|
# that includes the overloaded types.
|
|
if tvm.target.codegen.llvm_version_major() < 5:
|
|
return
|
|
|
|
# int1 is the type for the is_zero_undef parameter
|
|
int1_zero = tvm.tirx.const(0, "int1")
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((1, 1), "int32"), C: T.Buffer((1, 1), "int32")):
|
|
with T.sblock("C"):
|
|
T.reads()
|
|
T.writes()
|
|
C[0, 0] = T.call_llvm_pure_intrin("int32", "llvm.ctlz", A[0, 0], int1_zero)
|
|
|
|
f = tvm.compile(Module, target="llvm")
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_lookup_intrin():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.handle("uint8x8")):
|
|
A_buf = T.decl_buffer((1,), "uint8x8", data=A)
|
|
T.evaluate(T.call_llvm_pure_intrin("uint8x8", "llvm.ctpop.v8i8", T.uint32(1), A_buf[0]))
|
|
|
|
fcode = tvm.compile(Module, None)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_large_uintimm():
|
|
value = (1 << 63) + 123
|
|
large_val = tvm.tirx.const(value, "uint64")
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((), "uint64")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
with T.sblock("A"):
|
|
vi = T.axis.spatial(1, 0)
|
|
T.reads()
|
|
T.writes(A[()])
|
|
A[()] = large_val + T.uint64(3)
|
|
|
|
f = tvm.compile(Module, target="llvm")
|
|
dev = tvm.cpu(0)
|
|
a = tvm.runtime.empty((), dtype="uint64", device=dev)
|
|
f(a)
|
|
assert a.numpy() == value + 3
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_multi_parallel():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((128,), "float32"), C: T.Buffer((128,), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
B = T.sblock_alloc_buffer((128,))
|
|
for i0_0_0 in T.parallel(1):
|
|
for ax0 in range(128):
|
|
with T.sblock("B"):
|
|
v_i0 = T.axis.spatial(128, ax0)
|
|
T.reads(A[v_i0])
|
|
T.writes(B[v_i0])
|
|
B[v_i0] = A[v_i0] + T.float32(1.0)
|
|
for i0_0_1 in range(16):
|
|
for i0_1 in T.parallel(8):
|
|
with T.sblock("C"):
|
|
v_i0 = T.axis.spatial(128, i0_0_0 * 128 + i0_0_1 * 8 + i0_1)
|
|
T.reads(B[v_i0])
|
|
T.writes(C[v_i0])
|
|
C[v_i0] = T.sqrt(B[v_i0]) * T.float32(2.0) + T.float32(2.0)
|
|
|
|
n = 128
|
|
f = tvm.compile(Module, target="llvm")
|
|
dev = tvm.cpu(0)
|
|
a = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
|
|
c = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
|
|
f(a, c)
|
|
tvm.testing.assert_allclose(c.numpy(), np.sqrt(a.numpy() + 1) * 2 + 2, rtol=1e-5)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_flip_pipeline():
|
|
def check_llvm(nn, base):
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((nn + base,), "float32"), C: T.Buffer((nn,), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i_0 in T.parallel((nn + 3) // 4):
|
|
for i_1 in T.vectorized(4):
|
|
with T.sblock("C"):
|
|
v_i = T.axis.spatial(nn, i_0 * 4 + i_1)
|
|
T.where(i_0 * 4 + i_1 < nn)
|
|
T.reads(A[nn + base - 1 - v_i])
|
|
T.writes(C[v_i])
|
|
C[v_i] = A[nn + base - 1 - v_i]
|
|
|
|
f = tvm.compile(Module, target="llvm")
|
|
dev = tvm.cpu(0)
|
|
a = tvm.runtime.tensor(np.random.uniform(size=(nn + base)).astype("float32"), dev)
|
|
c = tvm.runtime.tensor(np.zeros(nn, dtype="float32"), dev)
|
|
f(a, c)
|
|
tvm.testing.assert_allclose(c.numpy(), a.numpy()[::-1][:nn])
|
|
|
|
check_llvm(4, 0)
|
|
check_llvm(128, 8)
|
|
check_llvm(3, 0)
|
|
check_llvm(128, 1)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_vadd_pipeline():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
n = T.int32()
|
|
A = T.match_buffer(var_A, (n,))
|
|
B = T.match_buffer(var_B, (n,))
|
|
C = T.match_buffer(var_C, (n,))
|
|
for i_0 in range((n + 3) // 4):
|
|
for i_1 in T.vectorized(4):
|
|
with T.sblock("C"):
|
|
v_i = T.axis.spatial(n, i_0 * 4 + i_1)
|
|
T.where(i_0 * 4 + i_1 < n)
|
|
T.reads(A[v_i], B[v_i])
|
|
T.writes(C[v_i])
|
|
C[v_i] = A[v_i] + B[v_i]
|
|
|
|
f = tvm.compile(Module, target="llvm")
|
|
dev = tvm.cpu(0)
|
|
n = 128
|
|
a = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
|
|
b = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
|
|
c = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
|
|
f(a, b, c)
|
|
tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_madd_pipeline():
|
|
def check_llvm(nn, base, stride):
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(
|
|
A: T.Buffer((nn + base, stride), "float32"),
|
|
C: T.Buffer((nn, stride), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i_0 in T.parallel((nn + 3) // 4):
|
|
for i_1 in T.vectorized(4):
|
|
for j in range(stride):
|
|
with T.sblock("C"):
|
|
v_i = T.axis.spatial(nn, i_0 * 4 + i_1)
|
|
v_j = T.axis.spatial(stride, j)
|
|
T.where(i_0 * 4 + i_1 < nn)
|
|
T.reads(A[v_i + base, v_j])
|
|
T.writes(C[v_i, v_j])
|
|
C[v_i, v_j] = A[v_i + base, v_j] + T.float32(1.0)
|
|
|
|
f = tvm.compile(Module, target="llvm")
|
|
dev = tvm.cpu(0)
|
|
a = tvm.runtime.tensor(np.random.uniform(size=(nn + base, stride)).astype("float32"), dev)
|
|
c = tvm.runtime.tensor(np.zeros((nn, stride), dtype="float32"), dev)
|
|
f(a, c)
|
|
tvm.testing.assert_allclose(c.numpy(), a.numpy()[base:] + 1)
|
|
|
|
check_llvm(64, 0, 2)
|
|
check_llvm(4, 0, 1)
|
|
|
|
with tvm.transform.PassContext(config={"tirx.noalias": False}):
|
|
check_llvm(4, 0, 3)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_temp_space():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((1024,), "float32"), C: T.Buffer((1024,), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
B = T.sblock_alloc_buffer((1024,))
|
|
for i in range(1024):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(1024, i)
|
|
T.reads(A[v_i])
|
|
T.writes(B[v_i])
|
|
B[v_i] = A[v_i] + T.float32(1.0)
|
|
for i in range(1024):
|
|
with T.sblock("C"):
|
|
v_i = T.axis.spatial(1024, i)
|
|
T.reads(B[v_i])
|
|
T.writes(C[v_i])
|
|
C[v_i] = B[v_i] + T.float32(1.0)
|
|
|
|
nn = 1024
|
|
f = tvm.compile(Module, target="llvm")
|
|
dev = tvm.cpu(0)
|
|
a = tvm.runtime.tensor(np.random.uniform(size=nn).astype("float32"), dev)
|
|
c = tvm.runtime.tensor(np.zeros(nn, dtype="float32"), dev)
|
|
f(a, c)
|
|
tvm.testing.assert_allclose(c.numpy(), a.numpy() + 1 + 1)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_multiple_func():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def fadd1(var_A: T.handle, var_B: T.handle, var_C: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
n = T.int32()
|
|
A = T.match_buffer(var_A, (n,))
|
|
B = T.match_buffer(var_B, (n,))
|
|
C = T.match_buffer(var_C, (n,))
|
|
for i in range(n):
|
|
with T.sblock("C"):
|
|
v_i = T.axis.spatial(n, i)
|
|
T.reads(A[v_i], B[v_i])
|
|
T.writes(C[v_i])
|
|
C[v_i] = A[v_i] + B[v_i]
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def fadd2(var_A: T.handle, var_B: T.handle, var_C: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
n = T.int32()
|
|
A = T.match_buffer(var_A, (n,))
|
|
B = T.match_buffer(var_B, (n,))
|
|
C = T.match_buffer(var_C, (n,))
|
|
for i in range(n):
|
|
with T.sblock("C"):
|
|
v_i = T.axis.spatial(n, i)
|
|
T.reads(A[v_i], B[v_i])
|
|
T.writes(C[v_i])
|
|
C[v_i] = A[v_i] + B[v_i]
|
|
|
|
f = tvm.compile(Module, target="llvm")
|
|
dev = tvm.cpu(0)
|
|
n = 10
|
|
a = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
|
|
b = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
|
|
c = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
|
|
|
|
f["fadd1"](a, b, c)
|
|
tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())
|
|
f["fadd2"](a, b, c)
|
|
tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_condition():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((64,), "float32"), C: T.Buffer((64,), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i in range(64):
|
|
with T.sblock("C"):
|
|
v_i = T.axis.spatial(64, i)
|
|
T.reads(A[v_i])
|
|
T.writes(C[v_i])
|
|
C[v_i] = T.if_then_else(8 <= v_i, A[v_i], T.float32(0.0))
|
|
|
|
n = 64
|
|
offset = 8
|
|
f = tvm.compile(Module, target="llvm")
|
|
dev = tvm.cpu(0)
|
|
a = tvm.runtime.tensor(np.random.uniform(size=(n,)).astype("float32"), dev)
|
|
c = tvm.runtime.empty((n,), "float32", dev)
|
|
f(a, c)
|
|
c_np = a.numpy()
|
|
c_np[:offset] = 0
|
|
tvm.testing.assert_allclose(c.numpy(), c_np)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_bool():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((64,), "int32"), C: T.Buffer((64,), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i in range(64):
|
|
with T.sblock("C"):
|
|
v_i = T.axis.spatial(64, i)
|
|
T.reads(A[v_i])
|
|
T.writes(C[v_i])
|
|
C[v_i] = T.Cast("float32", A[v_i] == 1)
|
|
|
|
n = 64
|
|
f = tvm.compile(Module, target="llvm")
|
|
dev = tvm.cpu(0)
|
|
a = tvm.runtime.tensor(np.random.randint(0, 2, size=(n,)).astype("int32"), dev)
|
|
c = tvm.runtime.empty((n,), "float32", dev)
|
|
f(a, c)
|
|
c_np = a.numpy() == 1
|
|
tvm.testing.assert_allclose(c.numpy(), c_np)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_cast_float_to_bool():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((4,), "float32"), C: T.Buffer((4,), "bool")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i in range(4):
|
|
with T.sblock("C"):
|
|
v_i = T.axis.spatial(4, i)
|
|
T.reads(A[v_i])
|
|
T.writes(C[v_i])
|
|
C[v_i] = T.Cast("bool", A[v_i])
|
|
|
|
n = 4
|
|
f = tvm.compile(Module, target="llvm")
|
|
dev = tvm.cpu(0)
|
|
a = tvm.runtime.tensor(np.array([0.0, 1.0, np.nan, np.inf], dtype="float32"), dev)
|
|
c = tvm.runtime.empty((n,), dtype="bool", device=dev)
|
|
f(a, c)
|
|
c_np = np.array([False, True, True, True], dtype="bool")
|
|
tvm.testing.assert_allclose(c.numpy(), c_np)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_rank_zero():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(
|
|
A: T.Buffer((64,), "float32"),
|
|
scale: T.Buffer((), "float32"),
|
|
compute: T.Buffer((), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
C = T.sblock_alloc_buffer(())
|
|
for k in range(64):
|
|
with T.sblock("C"):
|
|
v_k = T.axis.reduce(64, k)
|
|
T.reads(A[v_k], scale[()])
|
|
T.writes(C[()])
|
|
with T.init():
|
|
C[()] = T.float32(0.0)
|
|
C[()] = C[()] + A[v_k] * scale[()]
|
|
with T.sblock("compute"):
|
|
vi = T.axis.spatial(1, 0)
|
|
T.reads(C[()])
|
|
T.writes(compute[()])
|
|
compute[()] = C[()] + T.float32(1.0)
|
|
|
|
n = 64
|
|
f = tvm.compile(Module, target="llvm")
|
|
dev = tvm.cpu(0)
|
|
a = tvm.runtime.tensor(np.random.randint(0, 2, size=(n,)).astype("float32"), dev)
|
|
sc = tvm.runtime.tensor(np.random.randint(0, 2, size=()).astype("float32"), dev)
|
|
d = tvm.runtime.empty((), "float32", dev)
|
|
f(a, sc, d)
|
|
d_np = np.sum(a.numpy()) * sc.numpy() + 1
|
|
tvm.testing.assert_allclose(d.numpy(), d_np)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_rank_zero_bound_checkers():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(
|
|
A: T.Buffer((64,), "float32"),
|
|
scale: T.Buffer((), "float32"),
|
|
compute: T.Buffer((), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
C = T.sblock_alloc_buffer(())
|
|
for k in range(64):
|
|
with T.sblock("C"):
|
|
v_k = T.axis.reduce(64, k)
|
|
T.reads(A[v_k], scale[()])
|
|
T.writes(C[()])
|
|
with T.init():
|
|
C[()] = T.float32(0.0)
|
|
C[()] = C[()] + A[v_k] * scale[()]
|
|
with T.sblock("compute"):
|
|
vi = T.axis.spatial(1, 0)
|
|
T.reads(C[()])
|
|
T.writes(compute[()])
|
|
compute[()] = C[()] + T.float32(1.0)
|
|
|
|
n = 64
|
|
with tvm.transform.PassContext(config={"tirx.instrument_bound_checkers": True}):
|
|
f = tvm.compile(Module, target="llvm")
|
|
dev = tvm.cpu(0)
|
|
a = tvm.runtime.tensor(np.random.randint(0, 2, size=(n,)).astype("float32"), dev)
|
|
sc = tvm.runtime.tensor(np.random.randint(0, 2, size=()).astype("float32"), dev)
|
|
d = tvm.runtime.empty((), "float32", dev)
|
|
f(a, sc, d)
|
|
d_np = np.sum(a.numpy()) * sc.numpy() + 1
|
|
tvm.testing.assert_allclose(d.numpy(), d_np)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_alignment():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def test_alignment(A: T.Buffer((1024,), "float32"), B: T.Buffer((1024,), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i_0 in range(128):
|
|
for i_1 in T.vectorized(8):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(1024, i_0 * 8 + i_1)
|
|
T.reads(A[v_i])
|
|
T.writes(B[v_i])
|
|
B[v_i] = A[v_i] * T.float32(3.0)
|
|
|
|
f = tvm.tirx.build(Module, target="llvm")
|
|
|
|
lines = f.inspect_source().split("\n")
|
|
|
|
# Check alignment on load/store.
|
|
for l in lines:
|
|
if "align" in l and "4 x float" in l:
|
|
assert "align 32" in l
|
|
|
|
# Check parameter alignment. This looks for the definition of the
|
|
# outlined "compute_" function to see if there is an "align" attribute
|
|
# listed there.
|
|
def has_param_alignment():
|
|
for l in lines:
|
|
if re.search(r"test_alignment_compute_\([^(]*align [0-9]", l):
|
|
return True
|
|
return False
|
|
|
|
if tvm.target.codegen.llvm_version_major() >= 5:
|
|
assert has_param_alignment()
|
|
|
|
# Check for assume intrinsics. This isn't 100% accurate, since it just
|
|
# checks if the llvm.assume is there, but detailed check would require
|
|
# a much more detailed analysis of the LLVM IR.
|
|
def has_call_to_assume():
|
|
for l in lines:
|
|
if re.search(r"call.*llvm.assume", l):
|
|
return True
|
|
return False
|
|
|
|
assert has_call_to_assume()
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_div():
|
|
"""Check that the semantics of div and mod is correct"""
|
|
|
|
def check(start, end, dstart, dend, dtype, floor_div=False):
|
|
a_size = end - start + 1
|
|
b_size = dend - dstart + 1
|
|
|
|
div_fn = tvm.tirx.floordiv if floor_div else tvm.tirx.truncdiv
|
|
mod_fn = tvm.tirx.floormod if floor_div else tvm.tirx.truncmod
|
|
|
|
# Build clipping helpers — capture TIR const values from env
|
|
_start = tvm.tirx.const(start, dtype)
|
|
_end = tvm.tirx.const(end, dtype)
|
|
_dstart = tvm.tirx.const(dstart, dtype)
|
|
_dend = tvm.tirx.const(dend, dtype)
|
|
|
|
if start == end:
|
|
clipa = lambda x: _start
|
|
else:
|
|
clipa = lambda x: T.min(_end, T.max(_start, x))
|
|
|
|
if dstart == dend:
|
|
clipb = lambda x: _dstart
|
|
else:
|
|
clipb = lambda x: T.min(_dend, T.max(_dstart, x))
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(
|
|
A: T.Buffer((a_size,), dtype),
|
|
B: T.Buffer((b_size,), dtype),
|
|
D: T.Buffer((a_size, b_size), dtype),
|
|
M: T.Buffer((a_size, b_size), dtype),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i, j in T.grid(a_size, b_size):
|
|
with T.sblock("D"):
|
|
v_i, v_j = T.axis.remap("SS", [i, j])
|
|
T.reads(A[v_i], B[v_j])
|
|
T.writes(D[v_i, v_j])
|
|
D[v_i, v_j] = div_fn(clipa(A[v_i]), clipb(B[v_j]))
|
|
with T.sblock("M"):
|
|
v_i, v_j = T.axis.remap("SS", [i, j])
|
|
T.reads(A[v_i], B[v_j])
|
|
T.writes(M[v_i, v_j])
|
|
M[v_i, v_j] = mod_fn(clipa(A[v_i]), clipb(B[v_j]))
|
|
|
|
f = tvm.compile(Module, target="llvm")
|
|
|
|
# Fill input arrays with values
|
|
A_arr = tvm.runtime.empty((a_size,), dtype)
|
|
B_arr = tvm.runtime.empty((b_size,), dtype)
|
|
A_arr.copyfrom(np.arange(start, end + 1, dtype=dtype))
|
|
B_np = np.arange(dstart, dend + 1, dtype=dtype)
|
|
# If the range of the divisor contains 0, replace it with 1 to avoid division by zero
|
|
if dend >= 0 and dstart <= 0:
|
|
B_np[-dstart] = 1
|
|
B_arr.copyfrom(B_np)
|
|
D_arr = tvm.runtime.empty((a_size, b_size), dtype)
|
|
M_arr = tvm.runtime.empty((a_size, b_size), dtype)
|
|
|
|
# Run the function and convert the results to numpy
|
|
f(A_arr, B_arr, D_arr, M_arr)
|
|
D_arr = D_arr.numpy()
|
|
M_arr = M_arr.numpy()
|
|
|
|
# This helper just prints additional info on failure
|
|
def _show_info():
|
|
print(f"dtype: {dtype}")
|
|
print(f"dividend range: [{start}, {end}]")
|
|
print(f"divisor range: [{dstart}, {dend}]")
|
|
|
|
# Check that the computed values are correct
|
|
for i in range(start, end + 1):
|
|
for j in range(dstart, dend + 1):
|
|
if j == 0:
|
|
continue
|
|
|
|
if floor_div:
|
|
dref = i // j
|
|
mref = i % j
|
|
else:
|
|
dref = int(float(i) / j)
|
|
mref = int(math.fmod(i, j))
|
|
|
|
if D_arr[i - start, j - dstart] != dref:
|
|
_show_info()
|
|
raise AssertionError(
|
|
f"Incorrect division result: {div_fn.__name__}({i}, {j}) is {D_arr[i - start, j - dstart]} "
|
|
f"but should be {dref}"
|
|
)
|
|
if M_arr[i - start, j - dstart] != mref:
|
|
_show_info()
|
|
raise AssertionError(
|
|
f"Incorrect modulo result: {mod_fn.__name__}({i}, {j}) is {M_arr[i - start, j - dstart]} "
|
|
f"but should be {mref}"
|
|
)
|
|
|
|
# Try different ranges to cover different cases
|
|
for start, end in [
|
|
(-12, -12),
|
|
(-11, -1),
|
|
(-11, 0),
|
|
(0, 0),
|
|
(12, 12),
|
|
(1, 11),
|
|
(0, 11),
|
|
(-11, 11),
|
|
]:
|
|
for dstart, dend in [
|
|
(-11, -1),
|
|
(-11, 1),
|
|
(-4, -4),
|
|
(-2, -2),
|
|
(1, 11),
|
|
(0, 11),
|
|
(4, 4),
|
|
(2, 2),
|
|
(-11, 11),
|
|
]:
|
|
if end < start or dend < dstart or (dend == 0 and dstart == 0) or dend == 0:
|
|
continue
|
|
check(start, end, dstart, dend, "int32", floor_div=False)
|
|
check(start, end, dstart, dend, "int32", floor_div=True)
|
|
check(start, end, dstart, dend, "int8", floor_div=False)
|
|
check(start, end, dstart, dend, "int8", floor_div=True)
|
|
if start >= 0 and dstart >= 0:
|
|
check(start, end, dstart, dend, "uint32", floor_div=False)
|
|
check(start, end, dstart, dend, "uint32", floor_div=True)
|
|
|
|
# Additional tests for uint8
|
|
for dstart, dend in [(0, 11), (1, 11), (2, 2), (4, 4)]:
|
|
check(123, 133, dstart, dend, "uint8", floor_div=False)
|
|
check(123, 133, dstart, dend, "uint8", floor_div=True)
|
|
check(0, 255, dstart, dend, "uint8", floor_div=False)
|
|
check(0, 255, dstart, dend, "uint8", floor_div=True)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_fp_math():
|
|
@I.ir_module(s_tir=True)
|
|
class RecipModule:
|
|
@T.prim_func(s_tir=True)
|
|
def main(var_A: T.handle, var_B: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
n = T.int32()
|
|
A = T.match_buffer(var_A, (n,))
|
|
B = T.match_buffer(var_B, (n,))
|
|
for i in range(n):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(n, i)
|
|
T.reads(A[v_i])
|
|
T.writes(B[v_i])
|
|
B[v_i] = T.float32(1.0) / (
|
|
T.float32(9999999999999999538762658202121142272.0) * A[v_i]
|
|
)
|
|
|
|
f_recip = tvm.compile(RecipModule, target="llvm")
|
|
|
|
for n in [4, 8, 16]:
|
|
a = tvm.runtime.tensor(np.full((n,), 100, "float32"))
|
|
b = tvm.runtime.empty((n,), "float32")
|
|
f_recip(a, b)
|
|
tvm.testing.assert_allclose(b.numpy(), np.zeros((n,), "float32"))
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class SigmoidModule:
|
|
@T.prim_func(s_tir=True)
|
|
def main(var_A: T.handle, var_B: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
n = T.int32()
|
|
A = T.match_buffer(var_A, (n,))
|
|
B = T.match_buffer(var_B, (n,))
|
|
for i in range(n):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(n, i)
|
|
T.reads(A[v_i])
|
|
T.writes(B[v_i])
|
|
B[v_i] = T.sigmoid(A[v_i])
|
|
|
|
f_sigmoid = tvm.compile(SigmoidModule, target="llvm")
|
|
|
|
for n in [4, 8, 16]:
|
|
a = tvm.runtime.tensor(np.full((n,), -1000, "float32"))
|
|
b = tvm.runtime.empty((n,), "float32")
|
|
f_sigmoid(a, b)
|
|
tvm.testing.assert_allclose(b.numpy(), np.zeros((n,), "float32"))
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_dwarf_debug_information():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(
|
|
A: T.Buffer((1024,), "float32"),
|
|
B: T.Buffer((1024,), "float32"),
|
|
C: T.Buffer((1024,), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0_0 in T.parallel(256):
|
|
for i0_1 in T.vectorized(4):
|
|
with T.sblock("C"):
|
|
v_i0 = T.axis.spatial(1024, i0_0 * 4 + i0_1)
|
|
T.reads(A[v_i0], B[v_i0])
|
|
T.writes(C[v_i0])
|
|
C[v_i0] = A[v_i0] + B[v_i0]
|
|
|
|
def check_llvm_object():
|
|
if tvm.target.codegen.llvm_version_major() < 5:
|
|
return
|
|
if tvm.target.codegen.llvm_version_major() > 6:
|
|
return
|
|
# build two functions
|
|
mod = tvm.IRModule(
|
|
{
|
|
"fadd1": Module["main"].with_attr("global_symbol", "fadd1"),
|
|
"fadd2": Module["main"].with_attr("global_symbol", "fadd2"),
|
|
}
|
|
)
|
|
m = tvm.compile(mod, target="llvm")
|
|
temp = utils.tempdir()
|
|
o_path = temp.relpath("temp.o")
|
|
m.write_to_file(o_path)
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
|
|
# Try the dwarfdump utility (OS X)
|
|
if shutil.which("dwarfdump"):
|
|
output = subprocess.check_output(["dwarfdump", o_path])
|
|
assert re.search(r"""DW_AT_name\\t\("fadd1"\)""", str(output))
|
|
assert re.search(r"""DW_AT_name\\t\("fadd2"\)""", str(output))
|
|
|
|
# Try gobjdump (OS X)
|
|
if shutil.which("gobjdump"):
|
|
output = subprocess.check_output(["gobjdump", "--dwarf", o_path])
|
|
assert re.search(r"""DW_AT_name.*fadd1""", str(output))
|
|
assert re.search(r"""DW_AT_name.*fadd2""", str(output))
|
|
|
|
# Try objdump (Linux) - Darwin objdump has different DWARF syntax.
|
|
if shutil.which("objdump") and sys.platform != "darwin":
|
|
output = subprocess.check_output(["objdump", "--dwarf", o_path])
|
|
assert re.search(r"""DW_AT_name.*fadd1""", str(output))
|
|
assert re.search(r"""DW_AT_name.*fadd2""", str(output))
|
|
|
|
def check_llvm_ir():
|
|
if tvm.target.codegen.llvm_version_major() < 5:
|
|
return
|
|
if tvm.target.codegen.llvm_version_major() > 6:
|
|
return
|
|
# build two functions
|
|
mod = tvm.IRModule(
|
|
{
|
|
"fadd1": Module["main"].with_attr("global_symbol", "fadd1"),
|
|
"fadd2": Module["main"].with_attr("global_symbol", "fadd2"),
|
|
}
|
|
)
|
|
m = tvm.tirx.build(mod, target={"kind": "llvm", "mtriple": "aarch64-linux-gnu"})
|
|
ll = m.inspect_source("ll")
|
|
|
|
# On non-Darwin OS, don't explicitly specify DWARF version.
|
|
import re
|
|
|
|
assert not re.search(r""""Dwarf Version""" "", ll)
|
|
assert re.search(r"""llvm.dbg.value""", ll)
|
|
|
|
# Try Darwin, require DWARF-2
|
|
m = tvm.tirx.build(mod, target={"kind": "llvm", "mtriple": "x86_64-apple-darwin-macho"})
|
|
ll = m.inspect_source("ll")
|
|
assert re.search(r"""i32 4, !"Dwarf Version", i32 2""", ll)
|
|
assert re.search(r"""llvm.dbg.value""", ll)
|
|
|
|
check_llvm_object()
|
|
check_llvm_ir()
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_bf16():
|
|
def dotest(do_vectorize):
|
|
loop_kind = T.vectorized if do_vectorize else T.serial
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(
|
|
A: T.Buffer((32,), "bfloat16"),
|
|
B: T.Buffer((32,), "bfloat16"),
|
|
D: T.Buffer((32,), "bfloat16"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for x in loop_kind(32):
|
|
with T.sblock("D"):
|
|
v_x = T.axis.spatial(32, x)
|
|
T.reads(A[v_x], B[v_x])
|
|
T.writes(D[v_x])
|
|
D[v_x] = A[v_x] + B[v_x]
|
|
|
|
np.random.seed(122)
|
|
module = tvm.compile(Module, target="llvm")
|
|
npa = np.random.rand(32).astype("bfloat16")
|
|
npb = np.random.rand(32).astype("bfloat16")
|
|
res = npa + npb
|
|
a_ = tvm.runtime.tensor(npa)
|
|
b_ = tvm.runtime.tensor(npb)
|
|
c_ = tvm.runtime.empty((32,), "bfloat16")
|
|
module(a_, b_, c_)
|
|
# Note: directly compare without casting to float32 should work with the
|
|
# latest numpy version.
|
|
tvm.testing.assert_allclose(c_.numpy().astype("float32"), res.astype("float32"))
|
|
|
|
dotest(True)
|
|
dotest(False)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_crt_static_lib():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(
|
|
A: T.Buffer((32,), "bfloat16"),
|
|
B: T.Buffer((32,), "bfloat16"),
|
|
C: T.Buffer((32,), "bfloat16"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for x in range(32):
|
|
with T.sblock("compute"):
|
|
v_x = T.axis.spatial(32, x)
|
|
T.reads(A[v_x], B[v_x])
|
|
T.writes(C[v_x])
|
|
C[v_x] = A[v_x] + B[v_x]
|
|
|
|
module = tvm.tirx.build(
|
|
Module.with_attr("system_lib_prefix", ""),
|
|
target=tvm.target.Target("llvm"),
|
|
)
|
|
module.inspect_source()
|
|
with utils.tempdir() as temp:
|
|
module.write_to_file(temp.relpath("test.o"))
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_order_functions():
|
|
"""Check that functions in the LLVM module are ordered alphabetically."""
|
|
|
|
# Note: the order is alphabetical because that's a predictable ordering. Any predictable
|
|
# ordering will work fine, but if the ordering changes, this test will need to be updated.
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def Danny(v: T.float32) -> T.float32:
|
|
T.ret(T.call_extern("float32", "Dave", v))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def Sammy(v: T.float32) -> T.float32:
|
|
T.ret(T.call_extern("float32", "Eve", v))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def Kirby(v: T.float32) -> T.float32:
|
|
T.ret(T.call_extern("float32", "Fred", v))
|
|
|
|
ir_text = tvm.tirx.build(Module, target="llvm").inspect_source("ll")
|
|
# Skip functions whose names start with _.
|
|
matches = re.findall(r"^define[^@]*@([a-zA-Z][a-zA-Z0-9_]*)", ir_text, re.MULTILINE)
|
|
assert matches == sorted(matches)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
@tvm.testing.skip_if_32bit
|
|
def test_llvm_import():
|
|
"""all-platform-minimal-test: check shell dependent clang behavior."""
|
|
# extern "C" is necessary to get the correct signature
|
|
cc_code = """
|
|
extern "C" float my_add(float x, float y) {
|
|
return x + y;
|
|
}
|
|
"""
|
|
|
|
def check_llvm(use_file):
|
|
if not clang.find_clang(required=False):
|
|
print("skip because clang is not available")
|
|
return
|
|
temp = utils.tempdir()
|
|
ll_path = temp.relpath("temp.ll")
|
|
ll_code = clang.create_llvm(cc_code, output=ll_path)
|
|
import_val = ll_path if use_file else ll_code
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((10,), "float32"), B: T.Buffer((10,), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i in T.serial(10, annotations={"pragma_import_llvm": import_val}):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(10, i)
|
|
T.reads(A[v_i])
|
|
T.writes(B[v_i])
|
|
B[v_i] = T.call_pure_extern("float32", "my_add", A[v_i], T.float32(1.0))
|
|
|
|
f = tvm.compile(Module, target="llvm")
|
|
dev = tvm.cpu(0)
|
|
a = tvm.runtime.tensor(np.random.uniform(size=10).astype("float32"), dev)
|
|
b = tvm.runtime.tensor(np.random.uniform(size=10).astype("float32"), dev)
|
|
f(a, b)
|
|
tvm.testing.assert_allclose(b.numpy(), a.numpy() + 1.0)
|
|
|
|
check_llvm(use_file=True)
|
|
check_llvm(use_file=False)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_scalar_concat():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(x: T.int32, y: T.int32, buffer: T.Buffer((1,), "int32x2")):
|
|
buffer[0] = T.Shuffle([x, y], [0, 1])
|
|
|
|
# This will crash in LLVM codegen if CodeGenLLVM::CreateVecConcat doesn't convert
|
|
# scalars to single-lane LLVM vectors.
|
|
with tvm.transform.PassContext(config={"tirx.disable_assert": True}):
|
|
m = tvm.compile(Module, target="llvm")
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_raise_exception_during_codegen():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((4, 4), "float32"), B: T.Buffer((4, 4), "float32")) -> None:
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i in T.parallel(4):
|
|
for j in T.parallel(4):
|
|
B[i, j] = A[i, j] * 2.0
|
|
|
|
with pytest.raises(RuntimeError) as e:
|
|
tvm.compile(Module, target="llvm")
|
|
msg = str(e)
|
|
assert msg.find("Nested parallel loop is not supported") != -1
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_target_attributes():
|
|
"""Check that when LLVM codegen creates new functions, they get the same target
|
|
attributes as the original function.
|
|
"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def test_func(var_A: T.handle, var_B: T.handle, var_C: T.handle, tindex: T.int32):
|
|
T.func_attr({"tirx.noalias": True})
|
|
A = T.match_buffer(var_A, (tindex,))
|
|
B = T.match_buffer(var_B, (tindex,))
|
|
C = T.match_buffer(var_C, (tindex,))
|
|
for i in range(tindex):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(tindex, i)
|
|
T.reads(A[v_i])
|
|
T.writes(B[v_i])
|
|
B[v_i] = A[v_i]
|
|
for i_0 in T.parallel(2):
|
|
for i_1 in range((tindex + 1) // 2):
|
|
with T.sblock("C"):
|
|
v_i = T.axis.spatial(tindex, i_0 * ((tindex + 1) // 2) + i_1)
|
|
T.where(i_0 * ((tindex + 1) // 2) + i_1 < tindex)
|
|
T.reads(B[v_i])
|
|
T.writes(C[v_i])
|
|
C[v_i] = B[v_i] + T.float32(1.0)
|
|
|
|
target_llvm = {
|
|
"kind": "llvm",
|
|
"mtriple": "x86_64-linux-gnu",
|
|
"mcpu": "skylake",
|
|
"mattr": ["+avx512f"],
|
|
}
|
|
target = tvm.target.Target(target_llvm, host=target_llvm)
|
|
module = tvm.tirx.build(Module, target=target)
|
|
|
|
llvm_ir = module.inspect_source()
|
|
llvm_ir_lines = llvm_ir.split("\n")
|
|
|
|
attribute_definitions = dict()
|
|
function_attr_map = dict()
|
|
|
|
for line in llvm_ir_lines:
|
|
func_def = re.match(
|
|
"define.* @(?P<func_name>[^(]*)[(].* #(?P<attr_num>[0-9]+) (!.* |){$", line
|
|
)
|
|
if func_def:
|
|
function_attr_map[func_def.group("func_name")] = func_def.group("attr_num")
|
|
continue
|
|
attr_def = re.match("attributes #(?P<attr_num>[0-9]+) = {(?P<attr_list>.*)}", line)
|
|
if attr_def:
|
|
attribute_definitions[attr_def.group("attr_num")] = attr_def.group("attr_list")
|
|
|
|
expected_functions = [
|
|
"__tvm_ffi_test_func",
|
|
"__tvm_parallel_lambda",
|
|
]
|
|
for n in expected_functions:
|
|
assert n in function_attr_map, f"Expected function {n} not found in LLVM IR"
|
|
attr_num = function_attr_map[n]
|
|
assert re.match('.*"target-cpu"="skylake".*', attribute_definitions[attr_num])
|
|
assert re.match('.*"target-features"=".*[+]avx512f.*".*', attribute_definitions[attr_num])
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_llvm_assume():
|
|
"""
|
|
Check that LLVM does not error out when generating code with tirx.assume.
|
|
Verifying for llvm.assume being generated is not easy as the intrinsic and its
|
|
related instructions get removed during optimizations
|
|
"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((4, 4), "int32"), B: T.Buffer((14,), "int32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
A_1 = T.decl_buffer((16,), "int32", data=A.data)
|
|
for axis0, axis1 in T.grid(4, 4):
|
|
T.assume(axis0 < 3 or axis1 < 2 or A_1[axis0 * 4 + axis1] == 0)
|
|
for i in range(14):
|
|
B_1 = T.decl_buffer((14,), "int32", data=B.data)
|
|
B_1[i] = A_1[i] * 2
|
|
|
|
m = tvm.compile(Module, target="llvm")
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_debug_symbol_for_float64():
|
|
"""Check that LLVM can define DWARF debug type for float64
|
|
|
|
In previous versions, only specific data types could exist in the
|
|
function signature. In this test, the "calling_conv" attribute
|
|
prevents lowering to the PackedFunc API.
|
|
"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(a: T.handle("float64"), b: T.handle("float64"), n: T.int64):
|
|
T.func_attr({"calling_conv": 2})
|
|
A = T.decl_buffer(16, "float64", data=a)
|
|
B = T.decl_buffer(16, "float64", data=b)
|
|
for i in range(n):
|
|
B[i] = A[i]
|
|
|
|
tvm.compile(Module, target="llvm")
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_subroutine_call():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer(1, dtype="float32")):
|
|
Module.subroutine(A.data)
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def subroutine(A_data: T.handle("float32")):
|
|
# The calling_conv parameter is to prevent MakePackedAPI
|
|
# from changing the call signature of the subroutine.
|
|
T.func_attr({"calling_conv": -1})
|
|
A = T.decl_buffer(1, dtype="float32", data=A_data)
|
|
A[0] = 42.0
|
|
|
|
target = "llvm"
|
|
dev = tvm.cpu()
|
|
|
|
built = tvm.compile(Module)
|
|
|
|
arr = tvm.runtime.tensor(np.zeros([1], "float32"), device=dev)
|
|
built["main"](arr)
|
|
assert arr.numpy()[0] == 42.0
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_call_packed_returning_void():
|
|
"""Allow codegen of PackedFunc calls returning void
|
|
|
|
The LLVM codegen uses the CallNode's dtype to cast the return type
|
|
of the PackedFunc into the appropriate LLVM output type. However,
|
|
there is no runtime dtype value for a void return. When the return type
|
|
of a PackedFunc is void, the generated code should not attempt to
|
|
read the return value.
|
|
|
|
While `T.call_packed()` will produce a CallNode with an output
|
|
dtype of "int32", the use of other return types is valid in TIR.
|
|
This test case uses `T.Call` directly to allow an explicit dtype
|
|
for the packed function call.
|
|
"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main():
|
|
T.Call(
|
|
tvm.ir.Op.get("tirx.tvm_call_packed"),
|
|
["dummy_function_name"],
|
|
ret_ty="void",
|
|
)
|
|
|
|
# Error occurred during build, as part of
|
|
# CodeGenCPU::MakeCallPackedLowered.
|
|
built = tvm.compile(Module, target="llvm")
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_call_packed_without_string_arg():
|
|
"""The first argument to tvm_call_packed must be a string
|
|
|
|
Even if the invalid TIR is constructed, this should throw an
|
|
exception to exit cleanly. Previously, use of
|
|
`args[0].as<StringImmNode>()` without a null check resulted in
|
|
a segfault during codegen.
|
|
"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer(1, "float32")):
|
|
T.Call(tvm.ir.Op.get("tirx.tvm_call_packed"), [A.data], ret_ty="int32")
|
|
|
|
with pytest.raises(RuntimeError):
|
|
built = tvm.compile(Module, target="llvm")
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_call_extern_returning_void():
|
|
"""Like test_call_packed_returning_void, but for call_extern"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main():
|
|
T.Call(tvm.ir.Op.get("tirx.call_extern"), ["dummy_function_name"], ret_ty="void")
|
|
|
|
built = tvm.compile(Module, target="llvm")
|
|
|
|
|
|
def test_invalid_volatile_masked_buffer_load():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(b: T.handle):
|
|
B = T.match_buffer(b, [4])
|
|
A = T.alloc_buffer((4,), annotations={"tirx.volatile": True})
|
|
B[0:4] = A.vload([T.Ramp(0, 1, 4)], predicate=T.Broadcast(T.bool(True), 4))
|
|
|
|
err_msg = "The masked load intrinsic does not support declaring load as volatile."
|
|
with pytest.raises(RuntimeError, match=err_msg):
|
|
with tvm.target.Target("llvm"):
|
|
tvm.compile(Module)
|
|
|
|
|
|
def test_invalid_volatile_masked_buffer_store():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main():
|
|
A = T.alloc_buffer((4,), annotations={"tirx.volatile": True})
|
|
A.vstore(
|
|
[T.Ramp(0, 1, 4)],
|
|
T.Broadcast(0.0, 4),
|
|
predicate=T.Broadcast(T.bool(True), 4),
|
|
)
|
|
|
|
err_msg = "The masked store intrinsic does not support declaring store as volatile."
|
|
with pytest.raises(RuntimeError, match=err_msg):
|
|
with tvm.target.Target("llvm"):
|
|
tvm.compile(Module)
|
|
|
|
|
|
def test_int_parameter():
|
|
"""Boolean may be passed to functions accepting int"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(arg: T.int32) -> T.int32:
|
|
T.func_attr({"target": T.target("llvm")})
|
|
if arg > 0:
|
|
return 10
|
|
else:
|
|
return 20
|
|
|
|
built = tvm.compile(Module)
|
|
output = built(True)
|
|
assert output == 10
|
|
|
|
output = built(False)
|
|
assert output == 20
|
|
|
|
|
|
def test_bool_parameter():
|
|
"""Integers may be passed to functions accepting bool"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(arg: T.bool) -> T.int32:
|
|
T.func_attr({"target": T.target("llvm")})
|
|
if arg:
|
|
return 10
|
|
else:
|
|
return 20
|
|
|
|
built = tvm.compile(Module)
|
|
output = built(1)
|
|
assert output == 10
|
|
|
|
output = built(2)
|
|
assert output == 10
|
|
|
|
output = built(0)
|
|
assert output == 20
|
|
|
|
|
|
def test_bool_return_value():
|
|
"""Booleans may be returned from a PrimFunc"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(value: T.int32) -> T.bool:
|
|
T.func_attr({"target": T.target("llvm")})
|
|
return value < 10
|
|
|
|
built = tvm.compile(Module)
|
|
assert isinstance(built(0), bool)
|
|
assert built(0)
|
|
|
|
assert isinstance(built(15), bool)
|
|
assert not built(15)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|