wehub-resource-sync
250 lines
7.9 KiB
Python
250 lines
7.9 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.
|
|
|
|
import numpy as np
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm.script import ir as I
|
|
from tvm.script import tirx as T
|
|
from tvm.support import utils
|
|
|
|
|
|
def test_add():
|
|
nn = 1024
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def test_fadd(
|
|
A: T.Buffer((1024,), "float32"),
|
|
B: T.Buffer((1024,), "float32"),
|
|
C: T.Buffer((1024,), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0 in range(1024):
|
|
with T.sblock("C"):
|
|
v_i0 = T.axis.spatial(1024, i0)
|
|
T.reads(A[v_i0], B[v_i0])
|
|
T.writes(C[v_i0])
|
|
C[v_i0] = A[v_i0] + B[v_i0]
|
|
|
|
def check_c():
|
|
mhost = tvm.compile(Module, target="c")
|
|
temp = utils.tempdir()
|
|
path_dso = temp.relpath("temp.so")
|
|
mhost.export_library(path_dso)
|
|
m = tvm.runtime.load_module(path_dso)
|
|
fadd = m["test_fadd"]
|
|
dev = tvm.cpu(0)
|
|
n = nn
|
|
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)
|
|
fadd(a, b, c)
|
|
tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())
|
|
|
|
check_c()
|
|
|
|
|
|
def test_reinterpret():
|
|
nn = 1024
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def test_reinterpret(
|
|
A: T.Buffer((1024,), "int32"),
|
|
B: T.Buffer((1024,), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0 in range(1024):
|
|
with T.sblock("B"):
|
|
v_i0 = T.axis.spatial(1024, i0)
|
|
T.reads(A[v_i0])
|
|
T.writes(B[v_i0])
|
|
B[v_i0] = T.reinterpret("float32", A[v_i0] + 2)
|
|
|
|
def check_c():
|
|
mhost = tvm.compile(Module, target="c")
|
|
temp = utils.tempdir()
|
|
path_dso = temp.relpath("temp.so")
|
|
mhost.export_library(path_dso)
|
|
m = tvm.runtime.load_module(path_dso)
|
|
fadd = m["test_reinterpret"]
|
|
dev = tvm.cpu(0)
|
|
n = nn
|
|
a = tvm.runtime.tensor(np.random.randint(-(2**30), 2**30, size=n).astype("int32"), dev)
|
|
b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
|
|
fadd(a, b)
|
|
tvm.testing.assert_allclose(b.numpy(), (2 + a.numpy()).view("float32"))
|
|
|
|
check_c()
|
|
|
|
|
|
def test_ceil():
|
|
nn = 1024
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def test_ceil(
|
|
A: T.Buffer((1024,), "float32"),
|
|
B: T.Buffer((1024,), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0 in range(1024):
|
|
with T.sblock("B"):
|
|
v_i0 = T.axis.spatial(1024, i0)
|
|
T.reads(A[v_i0])
|
|
T.writes(B[v_i0])
|
|
B[v_i0] = T.ceil(A[v_i0])
|
|
|
|
def check_c():
|
|
mhost = tvm.compile(Module, target="c")
|
|
temp = utils.tempdir()
|
|
path_dso = temp.relpath("temp.so")
|
|
mhost.export_library(path_dso)
|
|
m = tvm.runtime.load_module(path_dso)
|
|
fceil = m["test_ceil"]
|
|
dev = tvm.cpu(0)
|
|
n = nn
|
|
a = tvm.runtime.tensor(np.random.rand(n).astype("float32"), dev)
|
|
b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
|
|
fceil(a, b)
|
|
tvm.testing.assert_allclose(b.numpy(), (np.ceil(a.numpy()).view("float32")))
|
|
|
|
check_c()
|
|
|
|
|
|
def test_floor():
|
|
nn = 1024
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def test_floor(
|
|
A: T.Buffer((1024,), "float32"),
|
|
B: T.Buffer((1024,), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0 in range(1024):
|
|
with T.sblock("B"):
|
|
v_i0 = T.axis.spatial(1024, i0)
|
|
T.reads(A[v_i0])
|
|
T.writes(B[v_i0])
|
|
B[v_i0] = T.floor(A[v_i0])
|
|
|
|
def check_c():
|
|
mhost = tvm.compile(Module, target="c")
|
|
temp = utils.tempdir()
|
|
path_dso = temp.relpath("temp.so")
|
|
mhost.export_library(path_dso)
|
|
m = tvm.runtime.load_module(path_dso)
|
|
ffloor = m["test_floor"]
|
|
dev = tvm.cpu(0)
|
|
n = nn
|
|
a = tvm.runtime.tensor(np.random.rand(n).astype("float32"), dev)
|
|
b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
|
|
ffloor(a, b)
|
|
tvm.testing.assert_allclose(b.numpy(), (np.floor(a.numpy()).view("float32")))
|
|
|
|
check_c()
|
|
|
|
|
|
def test_round():
|
|
nn = 1024
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def test_round(
|
|
A: T.Buffer((1024,), "float32"),
|
|
B: T.Buffer((1024,), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0 in range(1024):
|
|
with T.sblock("B"):
|
|
v_i0 = T.axis.spatial(1024, i0)
|
|
T.reads(A[v_i0])
|
|
T.writes(B[v_i0])
|
|
B[v_i0] = T.round(A[v_i0])
|
|
|
|
def check_c():
|
|
mhost = tvm.compile(Module, target="c")
|
|
temp = utils.tempdir()
|
|
path_dso = temp.relpath("temp.so")
|
|
mhost.export_library(path_dso)
|
|
m = tvm.runtime.load_module(path_dso)
|
|
fround = m["test_round"]
|
|
dev = tvm.cpu(0)
|
|
n = nn
|
|
a = tvm.runtime.tensor(np.random.rand(n).astype("float32"), dev)
|
|
b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
|
|
fround(a, b)
|
|
tvm.testing.assert_allclose(b.numpy(), (np.round(a.numpy()).view("float32")))
|
|
|
|
check_c()
|
|
|
|
|
|
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(private=True, s_tir=True)
|
|
def subroutine(A_data: T.handle("float32")):
|
|
A = T.decl_buffer(1, dtype="float32", data=A_data)
|
|
A[0] = 42.0
|
|
|
|
built = tvm.tirx.build(Module, target="c")
|
|
|
|
source = built.inspect_source()
|
|
assert source.count("__tvm_ffi_main(void*") == 2, (
|
|
"Expected two occurrences, for forward-declaration and definition"
|
|
)
|
|
assert source.count("subroutine(float*") == 2, (
|
|
"Expected two occurrences, for forward-declaration and definition"
|
|
)
|
|
assert source.count("subroutine(") == 3, (
|
|
"Expected three occurrences, for forward-declaration, definition, and call from main."
|
|
)
|
|
|
|
|
|
def test_workspace_allocation_cast():
|
|
@I.ir_module
|
|
class Module:
|
|
@T.prim_func
|
|
def main(A: T.Buffer((256,), "float32")):
|
|
workspace = T.alloc_buffer((256,), "float32", scope="global")
|
|
for i in range(256):
|
|
workspace[i] = A[i]
|
|
for i in range(256):
|
|
A[i] = workspace[i]
|
|
|
|
built = tvm.tirx.build(Module, target="c")
|
|
assert "((float*)TVMBackendAllocWorkspace(" in built.inspect_source()
|
|
|
|
temp = utils.tempdir()
|
|
built.export_library(temp.relpath("workspace.so"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|