wehub-resource-sync
221 lines
8.2 KiB
Python
221 lines
8.2 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: F811
|
|
import numpy as np
|
|
|
|
import tvm
|
|
from tvm import te
|
|
|
|
|
|
def test_trace_default_action():
|
|
n = 2
|
|
x = te.placeholder((n, n, n), name="X", dtype="float32")
|
|
y = te.compute(x.shape, lambda i, j, k: tvm.tirx.trace([i, j, k, x[i][j][k]]))
|
|
f = tvm.compile(te.create_prim_func([x, y]), target="llvm")
|
|
xnd = tvm.runtime.tensor(np.ones((n, n, n), dtype=x.dtype))
|
|
ynd = tvm.runtime.tensor(np.zeros((n, n, n), dtype=y.dtype))
|
|
f(xnd, ynd)
|
|
|
|
|
|
def test_trace_expr_assign():
|
|
@tvm.register_global_func("tvm.tirx.trace_callback2")
|
|
def trace_buffer(x):
|
|
return
|
|
|
|
def check_assign(dtype):
|
|
n = 4
|
|
x = te.placeholder((n, n, n), name="X", dtype=dtype)
|
|
y = te.compute(
|
|
x.shape, lambda i, j, k: tvm.tirx.trace([x[i][j][k]], "tvm.tirx.trace_callback2")
|
|
)
|
|
z = te.compute(
|
|
x.shape, lambda i, j, k: tvm.tirx.trace([y[i][j][k]], "tvm.tirx.trace_callback2")
|
|
)
|
|
f = tvm.compile(te.create_prim_func([x, y, z]), "llvm")
|
|
|
|
xnd = tvm.runtime.tensor(np.ones((n, n, n), dtype=x.dtype))
|
|
ynd = tvm.runtime.tensor(np.zeros((n, n, n), dtype=y.dtype))
|
|
znd = tvm.runtime.tensor(np.zeros((n, n, n), dtype=z.dtype))
|
|
f(xnd, ynd, znd)
|
|
|
|
assert np.array_equal(xnd.numpy(), np.ones((n, n, n)))
|
|
assert np.array_equal(ynd.numpy(), np.ones((n, n, n)))
|
|
assert np.array_equal(znd.numpy(), np.ones((n, n, n)))
|
|
|
|
for t in ["float64", "float32", "int64", "int32"]:
|
|
check_assign(t)
|
|
|
|
|
|
def test_trace_expr_sum_generated():
|
|
@tvm.register_global_func("tvm.tirx.trace_callback3")
|
|
def trace_buffer(x):
|
|
return
|
|
|
|
def check_expr_sum(dtype):
|
|
n = 4
|
|
a = te.placeholder((n, n, n), name="a", dtype=dtype)
|
|
b = te.placeholder((n, n, n), name="b", dtype=dtype)
|
|
c = te.compute(
|
|
a.shape,
|
|
lambda i, j, k: (
|
|
tvm.tirx.trace([a[i][j][k]], "tvm.tirx.trace_callback3")
|
|
+ tvm.tirx.trace([b[i][j][k]], "tvm.tirx.trace_callback3")
|
|
),
|
|
)
|
|
f = tvm.compile(te.create_prim_func([a, b, c]))
|
|
xnd = tvm.runtime.tensor(np.array(np.ones((n, n, n), dtype=a.dtype)))
|
|
ynd = tvm.runtime.tensor(np.array(np.ones((n, n, n), dtype=b.dtype)))
|
|
znd = tvm.runtime.tensor(np.zeros((n, n, n), dtype=c.dtype))
|
|
f(xnd, ynd, znd)
|
|
assert np.array_equal(znd.numpy(), xnd.numpy() + ynd.numpy())
|
|
|
|
for t in ["float64", "float32", "int64", "int32"]:
|
|
check_expr_sum(t)
|
|
|
|
|
|
def test_trace_expr_sum_args():
|
|
@tvm.register_global_func("tvm.tirx.trace_silent")
|
|
def silent(*args):
|
|
return
|
|
|
|
def check_expr_sum(dtype):
|
|
n = 4
|
|
a = te.placeholder((n, n, n), name="a", dtype=dtype)
|
|
b = te.placeholder((n, n, n), name="b", dtype=dtype)
|
|
e = te.placeholder((n, n, n), name="e", dtype=dtype)
|
|
d = te.placeholder((n, n, n), name="d", dtype=dtype)
|
|
|
|
c = te.compute(
|
|
a.shape,
|
|
lambda i, j, k: (
|
|
tvm.tirx.trace([i, j, k, a[i][j][k]], "tvm.tirx.trace_silent")
|
|
+ tvm.tirx.trace([i, j, k, b[i][j][k]], "tvm.tirx.trace_silent")
|
|
+ tvm.tirx.trace([i, j, k, d[i][j][k]], "tvm.tirx.trace_silent")
|
|
+ tvm.tirx.trace([i, j, k, e[i][j][k]], "tvm.tirx.trace_silent")
|
|
),
|
|
)
|
|
f = tvm.compile(te.create_prim_func([a, b, d, e, c]))
|
|
a_nd = tvm.runtime.tensor(np.array(np.ones((n, n, n), dtype=a.dtype)))
|
|
b_nd = tvm.runtime.tensor(np.array(np.ones((n, n, n), dtype=b.dtype)))
|
|
d_nd = tvm.runtime.tensor(np.array(np.ones((n, n, n), dtype=d.dtype)))
|
|
e_nd = tvm.runtime.tensor(np.array(np.ones((n, n, n), dtype=e.dtype)))
|
|
c_nd = tvm.runtime.tensor(np.zeros((n, n, n), dtype=c.dtype))
|
|
f(a_nd, b_nd, d_nd, e_nd, c_nd)
|
|
assert np.array_equal(
|
|
c_nd.numpy(), a_nd.numpy() + b_nd.numpy() + d_nd.numpy() + e_nd.numpy()
|
|
)
|
|
|
|
for t in ["float64", "float32", "int64", "int32"]:
|
|
check_expr_sum(t)
|
|
|
|
|
|
def test_trace_expr_sum_custom():
|
|
@tvm.register_global_func("tvm.tirx.trace_callback4")
|
|
def trace_buffer(x):
|
|
return
|
|
|
|
def check_expr_sum_custom(dtype):
|
|
n = 4
|
|
a = te.placeholder((n, n), name="a", dtype=dtype)
|
|
b = te.placeholder((n, n), name="b", dtype=dtype)
|
|
c = te.compute(
|
|
a.shape,
|
|
lambda i, j: (
|
|
tvm.tirx.trace([a[i][j]], "tvm.tirx.trace_callback4")
|
|
+ tvm.tirx.trace([b[i][j]], "tvm.tirx.trace_callback4")
|
|
),
|
|
)
|
|
f = tvm.compile(te.create_prim_func([a, b, c]))
|
|
npa = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], dtype=a.dtype)
|
|
npb = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], dtype=a.dtype)
|
|
xnd = tvm.runtime.tensor(npa)
|
|
ynd = tvm.runtime.tensor(npb)
|
|
znd = tvm.runtime.tensor(np.zeros((n, n), dtype=c.dtype))
|
|
f(xnd, ynd, znd)
|
|
assert np.array_equal(znd.numpy(), npa + npb)
|
|
|
|
for t in ["float64", "float32", "int64", "int32"]:
|
|
check_expr_sum_custom(t)
|
|
|
|
|
|
def test_trace_can_change_traced_value_int():
|
|
@tvm.register_global_func("tvm.tirx.trace_change_int_first")
|
|
def trace_buffer(x):
|
|
return 13
|
|
|
|
@tvm.register_global_func("tvm.tirx.trace_change_int_second")
|
|
def trace_buffer(x):
|
|
return 14
|
|
|
|
def check_assign(dtype):
|
|
n = 4
|
|
x = te.placeholder((n,), name="X", dtype=dtype)
|
|
y = te.compute(x.shape, lambda i: tvm.tirx.trace([x[i]], "tvm.tirx.trace_change_int_first"))
|
|
z = te.compute(
|
|
x.shape, lambda i: tvm.tirx.trace([y[i]], "tvm.tirx.trace_change_int_second")
|
|
)
|
|
f = tvm.compile(te.create_prim_func([x, y, z]))
|
|
|
|
xnd = tvm.runtime.tensor(np.ones((n,), dtype=x.dtype))
|
|
ynd = tvm.runtime.tensor(np.zeros((n,), dtype=y.dtype))
|
|
znd = tvm.runtime.tensor(np.zeros((n,), dtype=z.dtype))
|
|
f(xnd, ynd, znd)
|
|
check_array_first = np.array([13, 13, 13, 13])
|
|
check_array_second = np.array([14, 14, 14, 14])
|
|
assert np.array_equal(ynd.numpy(), check_array_first)
|
|
assert np.array_equal(znd.numpy(), check_array_second)
|
|
|
|
for t in ["int64", "int32"]:
|
|
check_assign(t)
|
|
|
|
|
|
def test_trace_can_change_traced_value_float():
|
|
@tvm.register_global_func("tvm.tirx.trace_change_float_first")
|
|
def trace_buffer(x):
|
|
return 13.0
|
|
|
|
@tvm.register_global_func("tvm.tirx.trace_change_float_second")
|
|
def trace_buffer(x):
|
|
return 14.0
|
|
|
|
def check_assign(dtype):
|
|
n = 4
|
|
x = te.placeholder((n,), name="X", dtype=dtype)
|
|
y = te.compute(
|
|
x.shape, lambda i: tvm.tirx.trace([x[i]], "tvm.tirx.trace_change_float_first")
|
|
)
|
|
z = te.compute(
|
|
x.shape, lambda i: tvm.tirx.trace([y[i]], "tvm.tirx.trace_change_float_second")
|
|
)
|
|
f = tvm.compile(te.create_prim_func([x, y, z]), target="llvm")
|
|
|
|
xnd = tvm.runtime.tensor(np.ones((n,), dtype=x.dtype))
|
|
ynd = tvm.runtime.tensor(np.zeros((n,), dtype=y.dtype))
|
|
znd = tvm.runtime.tensor(np.zeros((n,), dtype=z.dtype))
|
|
f(xnd, ynd, znd)
|
|
check_array_first = np.array([13.0, 13.0, 13.0, 13.0])
|
|
check_array_second = np.array([14.0, 14.0, 14.0, 14.0])
|
|
assert np.array_equal(ynd.numpy(), check_array_first)
|
|
assert np.array_equal(znd.numpy(), check_array_second)
|
|
|
|
for t in ["float64", "float32"]:
|
|
check_assign(t)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|