wehub-resource-sync
268 lines
9.7 KiB
Python
268 lines
9.7 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 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()
|