From d4edd7cd0b7bda060a4f106abb719e0c788a6e28 Mon Sep 17 00:00:00 2001 From: Jared Roesch Date: Sat, 1 Jun 2019 00:57:05 -0700 Subject: [PATCH 1/5] Fix calling convention for packed functions + tuples --- src/relay/backend/vm/compiler.cc | 48 ++++++++++++++++++++++++-------- src/relay/op/tensor/transform.cc | 6 ++-- tests/python/relay/test_vm.py | 1 + 3 files changed, 41 insertions(+), 14 deletions(-) diff --git a/src/relay/backend/vm/compiler.cc b/src/relay/backend/vm/compiler.cc index 602e92759624..621dfe0ad895 100644 --- a/src/relay/backend/vm/compiler.cc +++ b/src/relay/backend/vm/compiler.cc @@ -334,15 +334,42 @@ struct VMCompiler : ExprFunctor { return Instruction::AllocTensor(last_register, dltype, NewRegister()); } - void EmitInvokePrimitive(const Function& func, std::vector args_registers, + void EmitInvokePrimitive(const Function& func, + const std::vector& args_registers, const Type& ret_type) { + std::vector unpacked_arg_regs; std::vector allocs; - size_t return_num = 0; + + // Arity calculation must flatten tuples. + size_t arity = 0; + CHECK_EQ(func->params.size(), args_registers.size()); + for (size_t i = 0; i < func->params.size(); i++) { + auto ty = func->params[i]->checked_type(); + if (auto tensor_ty = ty.as()) { + unpacked_arg_regs.push_back(args_registers[i]); + arity += 1; + } else if (auto tuple_ty = ty.as()) { + for (size_t f = 0; f < tuple_ty->fields.size(); f++) { + const auto& field = tuple_ty->fields[f]; + CHECK(field.as()) + << "only supports non-nested tuples currently " + << "found " << field; + auto dst = NewRegister(); + Emit(Instruction::GetField(args_registers[i], f, dst)); + unpacked_arg_regs.push_back(dst); + } + arity += tuple_ty->fields.size(); + } else { + LOG(FATAL) << "unsupported parameter type " << ty; + } + } + + size_t return_val_count = 0; if (const TensorTypeNode* ttype = ret_type.as()) { // Allocate space for the return tensor. auto alloc = AllocTensorFromType(ttype); allocs.push_back(alloc); - return_num = 1; + return_val_count = 1; } else if (const TupleTypeNode* ttype = ret_type.as()) { std::vector fields_registers; @@ -352,14 +379,15 @@ struct VMCompiler : ExprFunctor { allocs.push_back(AllocTensorFromType(f_type)); fields_registers.push_back(allocs.back().dst); } - return_num = ttype->fields.size(); + return_val_count = ttype->fields.size(); } else { LOG(FATAL) << "Unsupported return value type"; } + arity += return_val_count; for (auto& alloc : allocs) { Emit(alloc); - args_registers.push_back(alloc.dst); + unpacked_arg_regs.push_back(alloc.dst); } // Next generate the invoke instruction. @@ -378,17 +406,15 @@ struct VMCompiler : ExprFunctor { op_index = seen_funcs[cfunc->funcs[0]]; } - // If Tensor, 1 - // If Tuple, size of tuple - size_t arity = func->params.size() + return_num; - Emit(Instruction::InvokePacked(op_index, arity, return_num, args_registers)); - if (return_num > 1) { + Emit(Instruction::InvokePacked(op_index, arity, return_val_count, unpacked_arg_regs)); + + if (return_val_count > 1) { // return value is a tuple, we need to create a tuple std::vector fields_registers; for (size_t i = func->params.size(); i < arity; ++i) { fields_registers.push_back(args_registers[i]); } - Emit(Instruction::AllocDatatype(0, return_num, fields_registers, NewRegister())); + Emit(Instruction::AllocDatatype(0, return_val_count, fields_registers, NewRegister())); } } diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index 873e75d9660b..16fb8bdca55f 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -6,9 +6,9 @@ * 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 @@ -292,7 +292,7 @@ RELAY_REGISTER_OP("concatenate") .add_type_rel("Concatenate", ConcatenateRel) .set_attr("FInferCorrectLayout", ConcatenateLayout) .set_attr("FTVMCompute", ConcatenateCompute) -.set_attr("TOpPattern", kInjective); +.set_attr("TOpPattern", kOpaque); TVM_REGISTER_NODE_TYPE(StackAttrs); diff --git a/tests/python/relay/test_vm.py b/tests/python/relay/test_vm.py index bc99418d5da4..db47a981cd46 100644 --- a/tests/python/relay/test_vm.py +++ b/tests/python/relay/test_vm.py @@ -259,6 +259,7 @@ def test_closure(): test_tuple_second() test_let_scalar() test_let_tensor() + test_split() # TODO(@jroesch): restore when match is supported # test_list_constructor() test_closure() From 2642915bfee5ddd9a5bea7b19ed50e42c0284351 Mon Sep 17 00:00:00 2001 From: Jared Roesch Date: Tue, 4 Jun 2019 02:34:55 -0700 Subject: [PATCH 2/5] Kill warning --- src/relay/backend/vm/compiler.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/relay/backend/vm/compiler.cc b/src/relay/backend/vm/compiler.cc index 621dfe0ad895..093b46d5774a 100644 --- a/src/relay/backend/vm/compiler.cc +++ b/src/relay/backend/vm/compiler.cc @@ -345,7 +345,7 @@ struct VMCompiler : ExprFunctor { CHECK_EQ(func->params.size(), args_registers.size()); for (size_t i = 0; i < func->params.size(); i++) { auto ty = func->params[i]->checked_type(); - if (auto tensor_ty = ty.as()) { + if (ty.as()) { unpacked_arg_regs.push_back(args_registers[i]); arity += 1; } else if (auto tuple_ty = ty.as()) { From 213245fa36c040571111688fe962abb5e52ae7c3 Mon Sep 17 00:00:00 2001 From: Jared Roesch Date: Tue, 4 Jun 2019 11:48:49 -0700 Subject: [PATCH 3/5] Reset fusion changes --- src/relay/op/tensor/transform.cc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index 16fb8bdca55f..873e75d9660b 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -6,9 +6,9 @@ * 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 @@ -292,7 +292,7 @@ RELAY_REGISTER_OP("concatenate") .add_type_rel("Concatenate", ConcatenateRel) .set_attr("FInferCorrectLayout", ConcatenateLayout) .set_attr("FTVMCompute", ConcatenateCompute) -.set_attr("TOpPattern", kOpaque); +.set_attr("TOpPattern", kInjective); TVM_REGISTER_NODE_TYPE(StackAttrs); From ef99ce6d74053f3894edad317458c5e7634a38e6 Mon Sep 17 00:00:00 2001 From: Jared Roesch Date: Tue, 4 Jun 2019 11:53:47 -0700 Subject: [PATCH 4/5] Add test --- tests/python/relay/test_vm.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/python/relay/test_vm.py b/tests/python/relay/test_vm.py index db47a981cd46..d727e776cbcd 100644 --- a/tests/python/relay/test_vm.py +++ b/tests/python/relay/test_vm.py @@ -49,6 +49,17 @@ def test_split(): res = veval(f, x_data) tvm.testing.assert_allclose(res.asnumpy(), np.split(x_data, 3, axis=0)[0]) +def test_split_no_fuse(): + x = relay.var('x', shape=(12,)) + y = relay.split(x, 3, axis=0).astuple() + z = relay.concatenate([relay.TupleGetItem(y, 0)], axis=0) + z = relay.annotation.stop_fusion(z) + f = relay.Function([x], z) + x_data = np.random.rand(12,).astype('float32') + res = veval(f, x_data) + tvm.testing.assert_allclose(res.asnumpy(), np.split(x_data, 3, axis=0)[0]) + + def test_id(): x = relay.var('x', shape=(10, 10)) f = relay.Function([x], x) @@ -260,6 +271,7 @@ def test_closure(): test_let_scalar() test_let_tensor() test_split() + test_split_no_fuse() # TODO(@jroesch): restore when match is supported # test_list_constructor() test_closure() From b12a0174e575c401e4e3eb7f8f85c60a4b090050 Mon Sep 17 00:00:00 2001 From: Wei Chen Date: Wed, 5 Jun 2019 17:25:03 +0800 Subject: [PATCH 5/5] Fix return field register --- src/relay/backend/vm/compiler.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/relay/backend/vm/compiler.cc b/src/relay/backend/vm/compiler.cc index 093b46d5774a..db98a9a9d3fd 100644 --- a/src/relay/backend/vm/compiler.cc +++ b/src/relay/backend/vm/compiler.cc @@ -411,8 +411,8 @@ struct VMCompiler : ExprFunctor { if (return_val_count > 1) { // return value is a tuple, we need to create a tuple std::vector fields_registers; - for (size_t i = func->params.size(); i < arity; ++i) { - fields_registers.push_back(args_registers[i]); + for (size_t i = arity - return_val_count; i < arity; ++i) { + fields_registers.push_back(unpacked_arg_regs[i]); } Emit(Instruction::AllocDatatype(0, return_val_count, fields_registers, NewRegister())); }