From 301df56de0b54ef1689729b93e217b9a118eb74c Mon Sep 17 00:00:00 2001 From: Xingyu Zhou Date: Wed, 21 Jul 2021 17:28:39 +0000 Subject: [PATCH 1/4] bug fix and add tensorarray with partition pass test case --- src/relay/transforms/partition_graph.cc | 2 +- tests/python/relay/test_tensor_array.py | 27 +++++++++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/src/relay/transforms/partition_graph.cc b/src/relay/transforms/partition_graph.cc index f2486356cce1..b48fbe44bd11 100644 --- a/src/relay/transforms/partition_graph.cc +++ b/src/relay/transforms/partition_graph.cc @@ -509,7 +509,7 @@ class NameMangleExtFuncs : public MixedModeMutator { // Walk the tree and mangle the functions. Then replace compiler functions // with mangled functions in the module - IRModule new_module; + IRModule new_module = IRModule({}, module_->type_definitions, module_->Imports()); for (const auto& pair : glob_funcs) { if (auto* fn = pair.second.as()) { auto func = GetRef(fn); diff --git a/tests/python/relay/test_tensor_array.py b/tests/python/relay/test_tensor_array.py index e93831bef95f..b127537e4b81 100644 --- a/tests/python/relay/test_tensor_array.py +++ b/tests/python/relay/test_tensor_array.py @@ -780,5 +780,32 @@ def run(dtype, shape): run("int32", [2, 3]) +@tvm.testing.uses_gpu +def test_static_tensor_array_gather_partition(): + 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) + from tvm.relay.op.contrib.tensorrt import partition_for_tensorrt + + mod, config = partition_for_tensorrt(mod, params=None, remove_no_mac_subgraphs=True) + + run("float32", [2, 3]) + + if __name__ == "__main__": pytest.main([__file__]) From a95ce8a581964b770634c40910a36384238ced9e Mon Sep 17 00:00:00 2001 From: Xingyu Zhou Date: Wed, 21 Jul 2021 20:24:44 +0000 Subject: [PATCH 2/4] change test function location and address comments --- .../python/relay/test_pass_partition_graph.py | 29 +++++++++++++++++++ tests/python/relay/test_tensor_array.py | 27 ----------------- 2 files changed, 29 insertions(+), 27 deletions(-) diff --git a/tests/python/relay/test_pass_partition_graph.py b/tests/python/relay/test_pass_partition_graph.py index 55b150d948c1..9795c70e83b3 100644 --- a/tests/python/relay/test_pass_partition_graph.py +++ b/tests/python/relay/test_pass_partition_graph.py @@ -1439,6 +1439,34 @@ def Optimize(mod): tvm.testing.assert_allclose(t0.body.data.numpy(), expected, rtol=1e-5, atol=1e-5) +# Test to make sure type definition and imports are preserved during the BYOC pipeline +def test_static_tensor_array_gather_partition(): + from tvm.relay.prelude import Prelude, StaticTensorArrayOps + + 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) + mod = transform.RemoveUnusedFunctions()(mod) + mod = transform.PartitionGraph()(mod) + + run("float32", [2, 3]) + + if __name__ == "__main__": test_multi_node_compiler() test_extern_ccompiler_single_op() @@ -1460,3 +1488,4 @@ def Optimize(mod): test_flatten_tuple_output() test_tuple_output_exec() test_extern_opt() + test_static_tensor_array_gather_partition() diff --git a/tests/python/relay/test_tensor_array.py b/tests/python/relay/test_tensor_array.py index b127537e4b81..e93831bef95f 100644 --- a/tests/python/relay/test_tensor_array.py +++ b/tests/python/relay/test_tensor_array.py @@ -780,32 +780,5 @@ def run(dtype, shape): run("int32", [2, 3]) -@tvm.testing.uses_gpu -def test_static_tensor_array_gather_partition(): - 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) - from tvm.relay.op.contrib.tensorrt import partition_for_tensorrt - - mod, config = partition_for_tensorrt(mod, params=None, remove_no_mac_subgraphs=True) - - run("float32", [2, 3]) - - if __name__ == "__main__": pytest.main([__file__]) From 65100426ddb8f548fd91638bfb653ac6c35ee1eb Mon Sep 17 00:00:00 2001 From: Xingyu Zhou Date: Wed, 21 Jul 2021 13:47:34 -0700 Subject: [PATCH 3/4] Update tests/python/relay/test_pass_partition_graph.py Co-authored-by: Cody Yu --- tests/python/relay/test_pass_partition_graph.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/python/relay/test_pass_partition_graph.py b/tests/python/relay/test_pass_partition_graph.py index 9795c70e83b3..0e0ae57b989d 100644 --- a/tests/python/relay/test_pass_partition_graph.py +++ b/tests/python/relay/test_pass_partition_graph.py @@ -1439,8 +1439,8 @@ def Optimize(mod): tvm.testing.assert_allclose(t0.body.data.numpy(), expected, rtol=1e-5, atol=1e-5) -# Test to make sure type definition and imports are preserved during the BYOC pipeline -def test_static_tensor_array_gather_partition(): +def test_preserve_type_import(): + """Test to make sure type definition and imports are preserved during the BYOC pipeline""" from tvm.relay.prelude import Prelude, StaticTensorArrayOps def run(dtype, shape): From 2cc8b9801802707d4f2c8832e6a57695ac370f35 Mon Sep 17 00:00:00 2001 From: Cody Yu Date: Fri, 23 Jul 2021 09:48:23 -0700 Subject: [PATCH 4/4] trigger CI --- tests/python/relay/test_pass_partition_graph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/python/relay/test_pass_partition_graph.py b/tests/python/relay/test_pass_partition_graph.py index 0e0ae57b989d..29d420def184 100644 --- a/tests/python/relay/test_pass_partition_graph.py +++ b/tests/python/relay/test_pass_partition_graph.py @@ -1440,7 +1440,7 @@ def Optimize(mod): def test_preserve_type_import(): - """Test to make sure type definition and imports are preserved during the BYOC pipeline""" + """Test to make sure type definition and imports are preserved during the BYOC pipeline.""" from tvm.relay.prelude import Prelude, StaticTensorArrayOps def run(dtype, shape):