diff --git a/3rdparty/HalideIR b/3rdparty/HalideIR index e4a4c02764d37..2f3ecdfdedf3e 160000 --- a/3rdparty/HalideIR +++ b/3rdparty/HalideIR @@ -1 +1 @@ -Subproject commit e4a4c02764d37c9c3db0d64c4996651a3ef9513c +Subproject commit 2f3ecdfdedf3efa7e45a3945dca63a25856c4674 diff --git a/docs/langref/relay_op.rst b/docs/langref/relay_op.rst index 0b937f6636bfd..1455acba04af6 100644 --- a/docs/langref/relay_op.rst +++ b/docs/langref/relay_op.rst @@ -114,6 +114,17 @@ This level enables additional math and transform operators. tvm.relay.image.resize +**Level 10: Temporary Operators** + +This level support backpropagation of broadcast operators. It is temporary. + +.. autosummary:: + :nosignatures: + + tvm.relay.broadcast_to_like + tvm.relay.collapse_sum_like + + Level 1 Definitions ------------------- .. autofunction:: tvm.relay.log @@ -181,6 +192,13 @@ Level 4 Definitions .. autofunction:: tvm.relay.where + Level 5 Definitions ------------------- .. autofunction:: tvm.relay.image.resize + + +Level 10 Definitions +-------------------- +.. autofunction:: tvm.relay.broadcast_to_like +.. autofunction:: tvm.relay.collapse_sum_like diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py index c2036f5091332..84e2398f0a9ed 100644 --- a/python/tvm/relay/op/transform.py +++ b/python/tvm/relay/op/transform.py @@ -242,3 +242,41 @@ def where(condition, x, y): Note that the shape of condition, x, and y needs to be the same. """ return _make.where(condition, x, y) + + +def broadcast_to_like(data, broadcast_type): + """Return an scalar value array with the same shape and type as the input array. + + Parameters + ---------- + data : relay.Expr + The input tensor. + + broadcast_type : relay.Expr + Provide the type to broadcast to. + + Returns + ------- + result : relay.Expr + The resulting tensor. + """ + return _make.broadcast_to_like(data, broadcast_type) + + +def collapse_sum_like(data, collapse_type): + """Return an scalar value array with the same shape and type as the input array. + + Parameters + ---------- + data : relay.Expr + The input tensor. + + collapse_type : relay.Expr + Provide the type to collapse to. + + Returns + ------- + result : relay.Expr + The resulting tensor. + """ + return _make.collapse_sum_like(data, collapse_type) diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index 956883476d094..0b0b3fb5f4acf 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -701,5 +701,66 @@ RELAY_REGISTER_OP("squeeze") .set_support_level(3) .add_type_rel("Squeeze", SqueezeRel); +// Have no idea how to assert the constraint. +// CollapseSumLike: -> B where BroadCast(A, B) = A +bool CollapseSumLikeRel(const Array& types, + int num_inputs, + const Attrs& attrs, + const TypeReporter& reporter) { + CHECK_EQ(types.size(), 3); + reporter->Assign(types[2], types[1]); + return true; +} + +Expr MakeCollapseSumLike(Expr data, + Expr collapse_type) { + static const Op& op = Op::Get("collapse_sum_like"); + return CallNode::make(op, {data, collapse_type}, Attrs(), {}); +} + +TVM_REGISTER_API("relay.op._make.collapse_sum_like") +.set_body([](const TVMArgs& args, TVMRetValue* rv) { + runtime::detail::unpack_call(MakeCollapseSumLike, args, rv); + }); + +RELAY_REGISTER_OP("collapse_sum_like") +.describe(R"code(Collapse the first input to match the shape of the second input. +)code" TVM_ADD_FILELINE) +.set_num_inputs(2) +.add_argument("data", "Tensor", "The input tensor.") +.add_argument("collapse_type", "Tensor", "Provide the type to collapse to.") +.set_support_level(10) +.add_type_rel("CollapseSumLike", CollapseSumLikeRel); + +// BroadCastToLike: -> B where BroadCast(A, B) = B +bool BroadCastToLikeRel(const Array& types, + int num_inputs, + const Attrs& attrs, + const TypeReporter& reporter) { + CHECK_EQ(types.size(), 3); + reporter->Assign(types[2], types[1]); + return true; +} + +Expr MakeBroadCastToLike(Expr data, + Expr broadcast_type) { + static const Op& op = Op::Get("broadcast_to_like"); + return CallNode::make(op, {data, broadcast_type}, Attrs(), {}); +} + +TVM_REGISTER_API("relay.op._make.broadcast_to_like") +.set_body([](const TVMArgs& args, TVMRetValue* rv) { + runtime::detail::unpack_call(MakeBroadCastToLike, args, rv); + }); + +RELAY_REGISTER_OP("broadcast_to_like") +.describe(R"code(Broadcast the first input to match the shape of the second input. +)code" TVM_ADD_FILELINE) +.set_num_inputs(2) +.add_argument("data", "Tensor", "The input tensor.") +.add_argument("broadcast_type", "Tensor", "Provide the type to broadcast to.") +.set_support_level(10) +.add_type_rel("BroadCastToLike", BroadCastToLikeRel); + } // namespace relay } // namespace tvm diff --git a/tests/python/relay/test_op_level10.py b/tests/python/relay/test_op_level10.py new file mode 100644 index 0000000000000..270133b25accf --- /dev/null +++ b/tests/python/relay/test_op_level10.py @@ -0,0 +1,31 @@ +""" Support level10 operator test cases. +""" +import tvm +from tvm import relay + +def test_collapse_sum_like(): + ib = relay.ir_builder.IRBuilder() + x = ib.param("x", relay.ty.TensorType((3, 4, 5, 6), "int8")) + y = ib.param("y", relay.ty.TensorType((4, 1, 6), "int8")) + with ib.function(x, y) as func: + ib.ret(relay.collapse_sum_like(x.var, y.var)) + ib.ret(func) + func = relay.ir_pass.infer_type(ib.env, func.to_func()) + ftype = func.checked_type + assert ftype.ret_type == relay.ty.TensorType((4, 1, 6), "int8") + + +def test_broadcast_to_like(): + ib = relay.ir_builder.IRBuilder() + x = ib.param("x", relay.ty.TensorType((3, 4, 5, 6), "int8")) + y = ib.param("x", relay.ty.TensorType((4, 1, 6), "int8")) + with ib.function(x, y) as func: + ib.ret(relay.broadcast_to_like(y.var, x.var)) + ib.ret(func) + func = relay.ir_pass.infer_type(ib.env, func.to_func()) + ftype = func.checked_type + assert ftype.ret_type == relay.ty.TensorType((3, 4, 5, 6), "int8") + +if __name__ == "__main__": + test_collapse_sum_like() + test_broadcast_to_like()