Mark Shields
c8a6089073
* This continues the work outlined in the RFC https://discuss.tvm.apache.org/t/rfc-relay-tecompiler-rewrite-existing-compile-engine-to-match-updated-compiler-flow/9233 This gets about halfway there for the Interpreter: * Remove direct access to TECompiler from interpreter, and instead call tec::LowerTEExpr when 'preparing' a module and expression for evaluation. * Make clear there's no phase distinction between create_interpreter and evaluate on the Python side -- both must be prepared together as a single IRModule. * But in return make sure the result of evaluate on the Python side is a packed func ready to directly apply 'simple' arguments to an already interpreted closure. * The interpreter builds and caches primitive TIR functions (and their corresponding dynamic shape functions) as packed funcs as they are encountered. * Cleanup uses of interpreter for constant folding on the C++ side. Future work: * Fold LoweredModule into IRModule so tec::LowerTEExpr is just another pass. * Get rid of the implicit caching of lowered functions in TECompiler. * Make calling convention from Relay to TIR explicit, and remove all the function attribute hackery currently needed so the interpreter can correctly invoke lowered functions as it encounters them. * Make TECompiler private. Though could do this now it will make migrating the VM and AOT uses of CompilerEngine harder. Force a gc between sphinx-gallery items to reclaim GPU memory. (#8722) GPU memory is only released once the PackedFunc for evaling the model is gced by Python. In CI we're noticing intermittent 'CUDA: Out of memory' failures while processing the tutorials, and tracing showed there was no gc happening between items. Not confident this will solve the problem but worth a try. * Get rid of logs spam.
786 lines
31 KiB
Python
786 lines
31 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 pytest
|
|
import tvm
|
|
from tvm import relay
|
|
from tvm.relay import testing
|
|
from tvm.relay.backend.interpreter import ConstructorValue
|
|
from tvm.relay import create_executor
|
|
from tvm.relay.prelude import Prelude, StaticTensorArrayOps
|
|
from tvm.relay.testing import count as count_, make_nat_value, make_nat_expr
|
|
|
|
import numpy as np
|
|
|
|
|
|
def vmobj_to_list(mod, o, dtype="float32"):
|
|
_, tensor_nil, _, _, _, _, _, _, _ = mod.get_type(f"tensor_{dtype}_t")
|
|
if isinstance(o, tvm.nd.NDArray):
|
|
return [o.numpy().tolist()]
|
|
elif isinstance(o, tvm.runtime.container.ADT):
|
|
if len(o) == 0:
|
|
if tensor_nil.tag == o.tag:
|
|
return [0]
|
|
return []
|
|
|
|
result = []
|
|
for f in o:
|
|
result.extend(vmobj_to_list(mod, f, dtype))
|
|
return result
|
|
elif isinstance(o, tvm.relay.backend.interpreter.ConstructorValue):
|
|
if o.constructor.name_hint == "Cons":
|
|
tl = vmobj_to_list(mod, o.fields[1], dtype)
|
|
hd = vmobj_to_list(mod, o.fields[0], dtype)
|
|
hd.extend(tl)
|
|
return hd
|
|
elif o.constructor.name_hint == "Nil":
|
|
return []
|
|
elif "tensor_nil" in o.constructor.name_hint:
|
|
return [0]
|
|
elif "tensor" in o.constructor.name_hint:
|
|
return [o.fields[0].numpy()]
|
|
else:
|
|
raise RuntimeError("Unknown object type: %s" % o.constructor.name_hint)
|
|
else:
|
|
raise RuntimeError("Unknown object type: %s" % type(o))
|
|
|
|
|
|
def check_tensor_array(ta_mod, ref_res, *args, dtype="float32", rtol=1e-5):
|
|
for kind in ["debug", "vm"]:
|
|
for target, dev in [("llvm", tvm.cpu(0))]: # testing.enabled_targets():
|
|
if kind == "debug" and dev.device_type != tvm.cpu().device_type:
|
|
continue
|
|
result = relay.create_executor(kind, mod=ta_mod, device=dev, target=target).evaluate()(
|
|
*args
|
|
)
|
|
got = vmobj_to_list(ta_mod, result, dtype)
|
|
tvm.testing.assert_allclose(ref_res, got, rtol=rtol, atol=rtol)
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_tensor_expand_dims():
|
|
def run(dtype):
|
|
x = relay.var("x")
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
expand_dims_func = p.get_global_var("tensor_expand_dims", dtype)
|
|
tensor1 = p.get_tensor_ctor("tensor1", dtype)
|
|
mod["main"] = relay.Function([x], expand_dims_func(tensor1(x)))
|
|
x_np = np.random.uniform(low=0.0, high=8.0, size=(1,)).astype(dtype)
|
|
expected = [np.expand_dims(x_np, axis=0)]
|
|
check_tensor_array(mod, expected, x_np)
|
|
|
|
run("float32")
|
|
run("int32")
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_tensor_array_constructor():
|
|
def run(dtype):
|
|
x = relay.var("x")
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
tensor_array = p.get_global_var("tensor_array", dtype)
|
|
mod["main"] = relay.Function([x], tensor_array(x))
|
|
expected = np.array([0, 0, 0, 0, 0])
|
|
check_tensor_array(mod, expected, 5, dtype=dtype)
|
|
|
|
run("float32")
|
|
run("int32")
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_tensor_array_read():
|
|
def run(dtype):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
l = relay.var("l")
|
|
i = relay.var("i")
|
|
read_func = p.get_global_var("tensor_array_read", dtype)
|
|
tensor_array = p.get_global_var("tensor_array", dtype)
|
|
mod["main"] = relay.Function([l, i], read_func(tensor_array(l), i))
|
|
expected = [0]
|
|
check_tensor_array(mod, expected, *(1, 0), dtype=dtype)
|
|
check_tensor_array(mod, expected, *(5, 1), dtype=dtype)
|
|
|
|
run("float32")
|
|
run("int32")
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_tensor_array_write():
|
|
def run(dtype):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
tensor_t = p.get_type("tensor_t", dtype)
|
|
v1 = relay.var("v1")
|
|
v2 = relay.var("v2")
|
|
tensor_array = p.get_global_var("tensor_array", dtype)
|
|
init_tensor_array = tensor_array(relay.const(2))
|
|
write_func = p.get_global_var("tensor_array_write", dtype)
|
|
tensor1 = p.get_tensor_ctor("tensor1", dtype)
|
|
tensor_array1 = write_func(init_tensor_array, relay.const(0), tensor1(v1))
|
|
tensor_array2 = write_func(tensor_array1, relay.const(1), tensor1(v2))
|
|
mod["main"] = relay.Function([v1, v2], tensor_array2)
|
|
expected = [3, 7]
|
|
check_tensor_array(mod, expected, *(3, 7), dtype=dtype)
|
|
|
|
run("float32")
|
|
run("int32")
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_tensor_array_stack():
|
|
def run(dtype):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
tensor_t = p.get_type("tensor_t", dtype)
|
|
rlist = p.mod.get_global_type_var(f"List")
|
|
tensor_array = p.get_global_var("tensor_array", dtype)
|
|
tensor1 = p.get_tensor_ctor("tensor1", dtype)
|
|
write = p.get_global_var("tensor_array_write", dtype)
|
|
stack = p.get_global_var("tensor_array_stack", dtype)
|
|
# TODO extract test case from inference failures
|
|
# setting this wrong causes crashes
|
|
v = relay.var("v", shape=(1,), dtype=dtype)
|
|
init_tensor_array = tensor_array(relay.const(3))
|
|
tensor_array1 = write(init_tensor_array, relay.const(0), tensor1(v))
|
|
tensor_array2 = write(tensor_array1, relay.const(1), tensor1(v))
|
|
tensor_array3 = write(tensor_array2, relay.const(2), tensor1(v))
|
|
tensor_array4 = stack(tensor_array3)
|
|
mod["main"] = relay.Function([v], tensor_array4, tensor_t())
|
|
t = np.random.uniform(low=0.0, high=8.0, size=(1,)).astype(dtype)
|
|
expected = [np.stack([t, t, t])]
|
|
check_tensor_array(mod, expected, t, dtype=dtype)
|
|
|
|
run("float32")
|
|
run("int32")
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_tensor_array_unstack():
|
|
def run(dtype):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
unstack_tensor1 = p.get_global_var("tensor_array_unstack_tensor1", dtype)
|
|
v = relay.var("v")
|
|
mod["main"] = relay.Function([v], unstack_tensor1(v))
|
|
t = np.random.uniform(low=0.0, high=8.0, size=(1,)).astype(dtype)
|
|
check_tensor_array(mod, t, t, dtype=dtype)
|
|
|
|
run("float32")
|
|
run("int32")
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_tensor_take():
|
|
def run(dtype):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
take = p.get_global_var("tensor_take", dtype)
|
|
tensor2 = p.get_tensor_ctor("tensor2", dtype)
|
|
v = relay.var("v")
|
|
lower = relay.var("lower")
|
|
upper = relay.var("upper")
|
|
mod["main"] = relay.Function([v, lower, upper], take(tensor2(v), lower, upper))
|
|
v_data = np.random.uniform(low=0.0, high=8.0, size=(10, 10)).astype(dtype)
|
|
expected = [np.take(v_data, range(2, 5), axis=0)]
|
|
check_tensor_array(mod, expected, *(v_data, 2, 5), dtype=dtype)
|
|
expected = [np.take(v_data, range(0, 9), axis=0)]
|
|
check_tensor_array(mod, expected, *(v_data, 0, 9), dtype=dtype)
|
|
|
|
run("float32")
|
|
run("int32")
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_tensor_concatenate():
|
|
def run(dtype):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
concat = p.get_global_var("tensor_concatenate", dtype)
|
|
tensor1 = p.get_tensor_ctor("tensor1", dtype)
|
|
v1 = relay.var("v1", shape=(tvm.tir.Any(),), dtype=dtype)
|
|
v2 = relay.var("v2", shape=(tvm.tir.Any(),), dtype=dtype)
|
|
mod["main"] = relay.Function([v1, v2], concat(tensor1(v1), tensor1(v2)))
|
|
v1_data = np.random.uniform(low=0.0, high=8.0, size=(5,)).astype(dtype)
|
|
v2_data = np.random.uniform(low=0.0, high=8.0, size=(5,)).astype(dtype)
|
|
expected = [np.concatenate((v1_data, v2_data))]
|
|
check_tensor_array(mod, expected, *(v1_data, v2_data), dtype=dtype)
|
|
|
|
run("float32")
|
|
run("int32")
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_tensor_array_concat():
|
|
def run(dtype):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
v1 = relay.var("v1")
|
|
v2 = relay.var("v2")
|
|
tensor_array = p.get_global_var("tensor_array", dtype)
|
|
tensor_array1 = tensor_array(relay.const(2))
|
|
write_func = p.get_global_var("tensor_array_write", dtype)
|
|
concat_func = p.get_global_var("tensor_array_concat", dtype)
|
|
tensor1 = p.get_tensor_ctor("tensor2", dtype)
|
|
tensor_array1 = write_func(tensor_array1, relay.const(0), tensor1(v1))
|
|
tensor_array1 = write_func(tensor_array1, relay.const(1), tensor1(v2))
|
|
tensor_array_concat = concat_func(tensor_array1)
|
|
mod["main"] = relay.Function([v1, v2], tensor_array_concat)
|
|
v1_data = np.random.uniform(low=0.0, high=8.0, size=(2, 3)).astype(dtype)
|
|
v2_data = np.random.uniform(low=0.0, high=8.0, size=(1, 3)).astype(dtype)
|
|
expected = [np.concatenate((v1_data, v2_data), axis=0)]
|
|
check_tensor_array(mod, expected, *(v1_data, v2_data), dtype=dtype)
|
|
|
|
run("float32")
|
|
run("int32")
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_tensor_array_scatter():
|
|
def run(dtype):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
|
|
# tensor array
|
|
v1 = relay.var("v1")
|
|
v2 = relay.var("v2")
|
|
v3 = relay.var("v2")
|
|
tensor_array = p.get_global_var("tensor_array", dtype)
|
|
tensor_array1 = tensor_array(relay.const(3))
|
|
write_func = p.get_global_var("tensor_array_write", dtype)
|
|
scatter_func = p.get_global_var("tensor_array_scatter", dtype)
|
|
tensor2 = p.get_tensor_ctor("tensor2", dtype)
|
|
tensor_array1 = write_func(tensor_array1, relay.const(0), tensor2(v1))
|
|
tensor_array1 = write_func(tensor_array1, relay.const(1), tensor2(v2))
|
|
tensor_array1 = write_func(tensor_array1, relay.const(2), tensor2(v3))
|
|
|
|
# indices array
|
|
index = relay.var("index")
|
|
|
|
# values array
|
|
value_0 = relay.var("value_0")
|
|
value_1 = relay.var("value_1")
|
|
values_array = tensor_array(relay.const(2))
|
|
values_array = write_func(values_array, relay.const(0), tensor2(value_0))
|
|
values_array = write_func(values_array, relay.const(1), tensor2(value_1))
|
|
|
|
# create the scatter function
|
|
tensor_array_scatter = scatter_func(tensor_array1, index, values_array)
|
|
mod["main"] = relay.Function([v1, v2, v3, index, value_0, value_1], tensor_array_scatter)
|
|
|
|
# initialize and check
|
|
v1_data = np.random.uniform(low=0.0, high=8.0, size=(2, 3)).astype(dtype)
|
|
v2_data = np.random.uniform(low=0.0, high=8.0, size=(2, 3)).astype(dtype)
|
|
v3_data = np.random.uniform(low=0.0, high=8.0, size=(2, 3)).astype(dtype)
|
|
index_data = np.array([0, 1], dtype="int32")
|
|
val1_data = np.random.uniform(low=0.0, high=8.0, size=(2, 3)).astype(dtype)
|
|
val2_data = np.random.uniform(low=0.0, high=8.0, size=(2, 3)).astype(dtype)
|
|
expected = [val1_data, val2_data, v3_data]
|
|
check_tensor_array(
|
|
mod,
|
|
expected,
|
|
*(v1_data, v2_data, v3_data, index_data, val1_data, val2_data),
|
|
dtype=dtype,
|
|
)
|
|
|
|
run("float32")
|
|
run("int32")
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_tensor_array_split():
|
|
def run(dtype):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
|
|
# tensor array
|
|
v1 = relay.var("v1")
|
|
v2 = relay.var("v2")
|
|
v3 = relay.var("v2")
|
|
tensor_array = p.get_global_var("tensor_array", dtype)
|
|
tensor_array1 = tensor_array(relay.const(3))
|
|
write_func = p.get_global_var("tensor_array_write", dtype)
|
|
split_func = p.get_global_var("tensor_array_split", dtype)
|
|
tensor2 = p.get_tensor_ctor("tensor2", dtype)
|
|
tensor_array1 = write_func(tensor_array1, relay.const(0), tensor2(v1))
|
|
tensor_array1 = write_func(tensor_array1, relay.const(1), tensor2(v2))
|
|
tensor_array1 = write_func(tensor_array1, relay.const(2), tensor2(v3))
|
|
|
|
# value tensor
|
|
value = relay.var("value")
|
|
|
|
# lengths tensor
|
|
ta_len = relay.var("length")
|
|
|
|
# create the scatter function
|
|
tensor_array_split = split_func(tensor_array1, tensor2(value), ta_len)
|
|
mod["main"] = relay.Function([v1, v2, v3, value, ta_len], tensor_array_split)
|
|
|
|
# initialize and check
|
|
v1_data = np.random.uniform(low=0.0, high=8.0, size=(2, 3)).astype(dtype)
|
|
v2_data = np.random.uniform(low=0.0, high=8.0, size=(2, 3)).astype(dtype)
|
|
v3_data = np.random.uniform(low=0.0, high=8.0, size=(2, 3)).astype(dtype)
|
|
value_data = np.random.uniform(low=0.0, high=8.0, size=(4, 3)).astype(dtype)
|
|
length_data = np.array([2, 2], dtype="int32")
|
|
expected = np.concatenate([value_data, v3_data])
|
|
expected = np.split(expected, indices_or_sections=[2, 4])
|
|
check_tensor_array(
|
|
mod, expected, *(v1_data, v2_data, v3_data, value_data, length_data), dtype=dtype
|
|
)
|
|
|
|
run("float32")
|
|
run("int32")
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_static_tensor_take():
|
|
def run(dtype, shape):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
|
|
static_tensor_array_ops.register()
|
|
|
|
take = p.get_global_var_static("tensor_take", dtype, shape)
|
|
tensor_constructor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
|
|
v = relay.var("v")
|
|
lower = relay.var("lower")
|
|
upper = relay.var("upper")
|
|
mod["main"] = relay.Function([v, lower, upper], take(tensor_constructor(v), lower, upper))
|
|
v_data = np.random.uniform(low=0.0, high=8.0, size=shape).astype(dtype)
|
|
expected = [np.take(v_data, range(2, 5), axis=0)]
|
|
check_tensor_array(mod, expected, *(v_data, 2, 5), dtype=dtype)
|
|
expected = [np.take(v_data, range(0, 9), axis=0)]
|
|
check_tensor_array(mod, expected, *(v_data, 0, 9), dtype=dtype)
|
|
|
|
run("float32", [10, 10])
|
|
run("int32", [15, 11])
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_static_tensor_concatenate():
|
|
def run(dtype, shape):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
|
|
static_tensor_array_ops.register()
|
|
|
|
concat = p.get_global_var_static("tensor_concatenate", dtype, shape)
|
|
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
|
|
v1 = relay.var("v1")
|
|
v2 = relay.var("v2")
|
|
mod["main"] = relay.Function([v1, v2], concat(tensor(v1), tensor(v2)))
|
|
v1_data = np.random.uniform(low=0.0, high=8.0, size=shape).astype(dtype)
|
|
v2_data = np.random.uniform(low=0.0, high=8.0, size=shape).astype(dtype)
|
|
expected = [np.concatenate((v1_data, v2_data))]
|
|
check_tensor_array(mod, expected, *(v1_data, v2_data), dtype=dtype)
|
|
|
|
run(
|
|
"float32",
|
|
[
|
|
5,
|
|
],
|
|
)
|
|
run("int32", [2, 3])
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_static_tensor_expand_dims():
|
|
def run(dtype, shape):
|
|
x = relay.var("x")
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
|
|
static_tensor_array_ops.register()
|
|
|
|
expand_dims_func = p.get_global_var_static("tensor_expand_dims", dtype, shape)
|
|
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
|
|
mod["main"] = relay.Function([x], expand_dims_func(tensor(x)))
|
|
x_np = np.random.uniform(low=0.0, high=8.0, size=shape).astype(dtype)
|
|
expected = [np.expand_dims(x_np, axis=0)]
|
|
check_tensor_array(mod, expected, x_np)
|
|
|
|
run("float32", [])
|
|
run(
|
|
"int32",
|
|
[
|
|
2,
|
|
],
|
|
)
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_static_tensor_array_constructor():
|
|
def run(dtype, shape):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
|
|
static_tensor_array_ops.register()
|
|
tensor_constructor = p.get_name_static("tensor_constructor", dtype, shape)
|
|
assert tensor_constructor != None
|
|
|
|
run("float32", [1, 1])
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_static_tensor_array_read():
|
|
def run(dtype, shape):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
|
|
static_tensor_array_ops.register()
|
|
|
|
np_data_list = []
|
|
ta_length = 3
|
|
for _ in range(ta_length):
|
|
np_data_list.append(np.random.uniform(0, 10, size=shape).astype(dtype))
|
|
|
|
v0 = relay.var("v0")
|
|
v1 = relay.var("v1")
|
|
v2 = relay.var("v2")
|
|
n = relay.var("n")
|
|
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
|
|
tensor_array = p.get_global_var_static("tensor_array", dtype, shape)
|
|
init_tensor_array = tensor_array(relay.const(ta_length))
|
|
read_func = p.get_global_var_static("tensor_array_read", dtype, shape)
|
|
write_func = p.get_global_var_static("tensor_array_write", dtype, shape)
|
|
tensor_array0 = write_func(init_tensor_array, relay.const(0), tensor(v0))
|
|
tensor_array1 = write_func(tensor_array0, relay.const(1), tensor(v1))
|
|
tensor_array2 = write_func(tensor_array1, relay.const(2), tensor(v2))
|
|
|
|
mod["main"] = relay.Function([v0, v1, v2, n], read_func(tensor_array2, n))
|
|
expected = [np_data_list[0]]
|
|
check_tensor_array(mod, expected, *list(np_data_list + [0]), dtype=dtype)
|
|
expected = [np_data_list[1]]
|
|
check_tensor_array(mod, expected, *list(np_data_list + [1]), dtype=dtype)
|
|
expected = [np_data_list[2]]
|
|
check_tensor_array(mod, expected, *list(np_data_list + [2]), dtype=dtype)
|
|
|
|
run("float32", [])
|
|
run("int32", [2, 3])
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_static_tensor_array_write():
|
|
def run(dtype, shape):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
|
|
static_tensor_array_ops.register()
|
|
|
|
ta_length = 2
|
|
np_data_list = [
|
|
np.random.uniform(0, 10, size=shape).astype(dtype) for _ in range(ta_length)
|
|
]
|
|
|
|
v0 = relay.var("v0")
|
|
v1 = relay.var("v1")
|
|
tensor_array = p.get_global_var_static("tensor_array", dtype, shape)
|
|
init_tensor_array = tensor_array(relay.const(ta_length))
|
|
write_func = p.get_global_var_static("tensor_array_write", dtype, shape)
|
|
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
|
|
tensor_array0 = write_func(init_tensor_array, relay.const(0), tensor(v0))
|
|
tensor_array1 = write_func(tensor_array0, relay.const(1), tensor(v1))
|
|
mod["main"] = relay.Function([v0, v1], tensor_array1)
|
|
expected = np_data_list
|
|
check_tensor_array(mod, expected, *np_data_list, dtype=dtype)
|
|
|
|
run("float32", [])
|
|
run("int32", [2, 3])
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_static_tensor_array_unstack():
|
|
def run(dtype, shape):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
|
|
static_tensor_array_ops.register()
|
|
|
|
unstack_tensor = p.get_global_var_static("tensor_array_unstack", dtype, shape)
|
|
v = relay.var("v")
|
|
mod["main"] = relay.Function([v], unstack_tensor(v))
|
|
t = np.random.uniform(low=0, high=10, size=shape).astype(dtype)
|
|
(*expected,) = t
|
|
check_tensor_array(mod, expected, t, dtype=dtype)
|
|
|
|
run("float32", [4])
|
|
run("int32", [2, 3])
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_static_tensor_array_scatter():
|
|
def run(dtype, shape, indices_shape=None):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
|
|
static_tensor_array_ops.register()
|
|
if indices_shape is not None:
|
|
static_tensor_array_ops.define_tensor_array_scatter(indices_shape, True)
|
|
|
|
# tensor array
|
|
v1 = relay.var("v1")
|
|
v2 = relay.var("v2")
|
|
v3 = relay.var("v2")
|
|
tensor_array = p.get_global_var_static("tensor_array", dtype, shape)
|
|
tensor_array0 = tensor_array(relay.const(3))
|
|
write_func = p.get_global_var_static("tensor_array_write", dtype, shape)
|
|
scatter_func = p.get_global_var_static("tensor_array_scatter", dtype, shape)
|
|
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
|
|
tensor_array1 = write_func(tensor_array0, relay.const(0), tensor(v1))
|
|
tensor_array1 = write_func(tensor_array1, relay.const(1), tensor(v2))
|
|
tensor_array1 = write_func(tensor_array1, relay.const(2), tensor(v3))
|
|
|
|
# indices array
|
|
index = relay.var("index")
|
|
|
|
# values array
|
|
value_0 = relay.var("value_0")
|
|
value_1 = relay.var("value_1")
|
|
values_array = tensor_array(relay.const(2))
|
|
values_array = write_func(values_array, relay.const(0), tensor(value_0))
|
|
values_array = write_func(values_array, relay.const(1), tensor(value_1))
|
|
|
|
# create the scatter function
|
|
tensor_array_scatter = scatter_func(tensor_array1, index, values_array)
|
|
mod["main"] = relay.Function([v1, v2, v3, index, value_0, value_1], tensor_array_scatter)
|
|
|
|
# initialize and check
|
|
v1_data = np.random.uniform(low=0.0, high=8.0, size=shape).astype(dtype)
|
|
v2_data = np.random.uniform(low=0.0, high=8.0, size=shape).astype(dtype)
|
|
v3_data = np.random.uniform(low=0.0, high=8.0, size=shape).astype(dtype)
|
|
index_data = np.array([0, 1], dtype="int32")
|
|
val1_data = np.random.uniform(low=0.0, high=8.0, size=shape).astype(dtype)
|
|
val2_data = np.random.uniform(low=0.0, high=8.0, size=shape).astype(dtype)
|
|
expected = [val1_data, val2_data, v3_data]
|
|
check_tensor_array(
|
|
mod,
|
|
expected,
|
|
*(v1_data, v2_data, v3_data, index_data, val1_data, val2_data),
|
|
dtype=dtype,
|
|
)
|
|
|
|
run("float32", [2, 3])
|
|
run("int32", [2, 3])
|
|
run(
|
|
"float32",
|
|
[2, 3],
|
|
[
|
|
2,
|
|
],
|
|
)
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_static_tensor_array_split():
|
|
def run(dtype, shape, value_shape=None, lengths_shape=None):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
|
|
static_tensor_array_ops.register()
|
|
if value_shape is not None or lengths_shape is not None:
|
|
static_tensor_array_ops.define_tensor_array_split(value_shape, lengths_shape, False)
|
|
|
|
# tensor array
|
|
v1 = relay.var("v1")
|
|
v2 = relay.var("v2")
|
|
v3 = relay.var("v2")
|
|
|
|
adt_shape = [
|
|
relay.Any(),
|
|
] + shape[1:]
|
|
test_ops = StaticTensorArrayOps(p, dtype, adt_shape)
|
|
test_ops.register()
|
|
tensor_array = test_ops.get_global_var("tensor_array")
|
|
|
|
tensor_array1 = tensor_array(relay.const(3))
|
|
write_func = test_ops.get_global_var("tensor_array_write")
|
|
split_ops = StaticTensorArrayOps(p, dtype, shape)
|
|
split_ops.register()
|
|
split_func = split_ops.get_global_var("tensor_array_split")
|
|
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, test_ops.shape)
|
|
tensor_array1 = write_func(tensor_array1, relay.const(0), tensor(v1))
|
|
tensor_array1 = write_func(tensor_array1, relay.const(1), tensor(v2))
|
|
tensor_array1 = write_func(tensor_array1, relay.const(2), tensor(v3))
|
|
|
|
# value tensor
|
|
value = relay.var("value")
|
|
|
|
# lengths tensor
|
|
ta_len = relay.var("length")
|
|
|
|
# create the split function
|
|
if value_shape is None:
|
|
tensor1 = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
|
|
else:
|
|
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, value_shape)
|
|
static_tensor_array_ops.register()
|
|
tensor1 = p.get_tensor_ctor_static("tensor_constructor", dtype, test_ops.shape)
|
|
|
|
tensor_array_split = split_func(tensor_array1, tensor1(value), ta_len)
|
|
mod["main"] = relay.Function([v1, v2, v3, value, ta_len], tensor_array_split)
|
|
|
|
# initialize and check
|
|
v1_data = np.random.uniform(low=0.0, high=8.0, size=[2, 3]).astype(dtype)
|
|
v2_data = np.random.uniform(low=0.0, high=8.0, size=[2, 3]).astype(dtype)
|
|
v3_data = np.random.uniform(low=0.0, high=8.0, size=[2, 3]).astype(dtype)
|
|
value_data = np.random.uniform(low=0.0, high=8.0, size=value_shape or shape).astype(dtype)
|
|
length_data = np.array([2, 2], dtype="int32")
|
|
expected = np.concatenate([value_data, v3_data])
|
|
expected = np.split(expected, indices_or_sections=[2, 4])
|
|
check_tensor_array(
|
|
mod, expected, *(v1_data, v2_data, v3_data, value_data, length_data), dtype=dtype
|
|
)
|
|
|
|
run("float32", [4, 3])
|
|
run("int32", [4, 3])
|
|
run(
|
|
"int32",
|
|
[relay.Any(), 3],
|
|
[4, 3],
|
|
[
|
|
2,
|
|
],
|
|
)
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_static_tensor_array_concat():
|
|
def run(dtype, shape):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
|
|
static_tensor_array_ops.register()
|
|
|
|
v1 = relay.var("v1")
|
|
v2 = relay.var("v2")
|
|
tensor_array = p.get_global_var_static("tensor_array", dtype, shape)
|
|
tensor_array1 = tensor_array(relay.const(2))
|
|
write_func = p.get_global_var_static("tensor_array_write", dtype, shape)
|
|
concat_func = p.get_global_var_static("tensor_array_concat", dtype, shape)
|
|
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
|
|
tensor_array1 = write_func(tensor_array1, relay.const(0), tensor(v1))
|
|
tensor_array1 = write_func(tensor_array1, relay.const(1), tensor(v2))
|
|
tensor_array_concat = concat_func(tensor_array1)
|
|
mod["main"] = relay.Function([v1, v2], tensor_array_concat)
|
|
v1_data = np.random.uniform(low=0.0, high=8.0, size=(2, 3)).astype(dtype)
|
|
v2_data = np.random.uniform(low=0.0, high=8.0, size=(1, 3)).astype(dtype)
|
|
expected = [np.concatenate((v1_data, v2_data), axis=0)]
|
|
check_tensor_array(mod, expected, *(v1_data, v2_data), dtype=dtype)
|
|
|
|
run("float32", [relay.Any(), 3])
|
|
run("int32", [relay.Any(), 3])
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_static_tensor_array_gather():
|
|
def run(dtype, shape):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
|
|
static_tensor_array_ops.register()
|
|
|
|
tensor_array = p.get_global_var_static("tensor_array", dtype, shape)
|
|
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
|
|
write = p.get_global_var_static("tensor_array_write", dtype, shape)
|
|
gather = p.get_global_var_static("tensor_array_gather", dtype, shape)
|
|
v = relay.var("v")
|
|
indice = relay.var("indice")
|
|
init_tensor_array = tensor_array(relay.const(3))
|
|
tensor_array1 = write(init_tensor_array, relay.const(0), tensor(v))
|
|
tensor_array2 = write(tensor_array1, relay.const(1), tensor(v))
|
|
tensor_array3 = write(tensor_array2, relay.const(2), tensor(v))
|
|
out = gather(tensor_array3, indice)
|
|
mod["main"] = relay.Function([v, indice], out)
|
|
t = np.random.uniform(low=0.0, high=8.0, size=shape).astype(dtype)
|
|
indice_data = np.array([0, 2], dtype="int32")
|
|
expected = [np.stack([t, t])]
|
|
check_tensor_array(mod, expected, *(t, indice_data), dtype=dtype)
|
|
|
|
run("float32", [])
|
|
run("int32", [2, 3])
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_static_tensor_array_stack():
|
|
def run(dtype, shape):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
|
|
static_tensor_array_ops.register()
|
|
|
|
tensor_array = p.get_global_var_static("tensor_array", dtype, shape)
|
|
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
|
|
write = p.get_global_var_static("tensor_array_write", dtype, shape)
|
|
stack = p.get_global_var_static("tensor_array_stack", dtype, shape)
|
|
v = relay.var("v")
|
|
init_tensor_array = tensor_array(relay.const(3))
|
|
tensor_array1 = write(init_tensor_array, relay.const(0), tensor(v))
|
|
tensor_array2 = write(tensor_array1, relay.const(1), tensor(v))
|
|
tensor_array3 = write(tensor_array2, relay.const(2), tensor(v))
|
|
tensor_array4 = stack(tensor_array3)
|
|
mod["main"] = relay.Function([v], tensor_array4)
|
|
t = np.random.uniform(low=0.0, high=8.0, size=shape).astype(dtype)
|
|
expected = [np.stack([t, t, t])]
|
|
check_tensor_array(mod, expected, t, dtype=dtype)
|
|
|
|
run("float32", [])
|
|
run("int32", [2, 3])
|
|
|
|
|
|
@tvm.testing.uses_gpu
|
|
def test_static_tensor_get_data():
|
|
def run(dtype, shape):
|
|
mod = tvm.IRModule()
|
|
p = Prelude(mod)
|
|
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
|
|
static_tensor_array_ops.register()
|
|
|
|
np_data_list = []
|
|
ta_length = 3
|
|
for _ in range(ta_length):
|
|
np_data_list.append(np.random.uniform(0, 10, size=shape).astype(dtype))
|
|
|
|
v0 = relay.var("v0")
|
|
v1 = relay.var("v1")
|
|
v2 = relay.var("v2")
|
|
n = relay.var("n")
|
|
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
|
|
tensor_array = p.get_global_var_static("tensor_array", dtype, shape)
|
|
init_tensor_array = tensor_array(relay.const(ta_length))
|
|
read_func = p.get_global_var_static("tensor_array_read", dtype, shape)
|
|
write_func = p.get_global_var_static("tensor_array_write", dtype, shape)
|
|
get_data_func = p.get_global_var_static("tensor_get_data", dtype, shape)
|
|
tensor_array0 = write_func(init_tensor_array, relay.const(0), tensor(v0))
|
|
tensor_array1 = write_func(tensor_array0, relay.const(1), tensor(v1))
|
|
tensor_array2 = write_func(tensor_array1, relay.const(2), tensor(v2))
|
|
|
|
mod["main"] = relay.Function([v0, v1, v2, n], get_data_func(read_func(tensor_array2, n)))
|
|
expected = [np_data_list[0]]
|
|
check_tensor_array(mod, expected, *list(np_data_list + [0]), dtype=dtype)
|
|
expected = [np_data_list[1]]
|
|
check_tensor_array(mod, expected, *list(np_data_list + [1]), dtype=dtype)
|
|
expected = [np_data_list[2]]
|
|
check_tensor_array(mod, expected, *list(np_data_list + [2]), dtype=dtype)
|
|
|
|
run("float32", [])
|
|
run("int32", [2, 3])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__])
|