Relay reshape reshape_like compute and schedule (#2159)

python/tvm/relay/op/_transform.py

@@ -53,3 +53,11 @@ _reg.register_schedule("strided_slice", schedule_injective)
 # slice_like
 _reg.register_schedule("slice_like", schedule_injective)
 _reg.register_pattern("slice_like", OpPattern.INJECTIVE)
+
+# reshape
+_reg.register_schedule("reshape", schedule_injective)
+_reg.register_pattern("reshape", OpPattern.INJECTIVE)
+
+# reshape_like
+_reg.register_schedule("reshape_like", schedule_injective)
+_reg.register_pattern("reshape_like", OpPattern.INJECTIVE)

src/relay/op/tensor/transform.cc

@@ -376,7 +376,15 @@ Example::
 .set_attrs_type_key("relay.attrs.ReshapeAttrs")
 .add_argument("data", "Tensor", "The input tensor.")
 .set_support_level(3)
-.add_type_rel("Reshape", ReshapeRel);
+.add_type_rel("Reshape", ReshapeRel)
+.set_attr<FTVMCompute>("FTVMCompute", [](const Attrs& attrs,
+                                         const Array<Tensor>& inputs,
+                                         const Type& out_type,
+                                         const Target& target) {
+  const auto* param = attrs.as<ReshapeAttrs>();
+  CHECK(param != nullptr);
+  return Array<Tensor>{ topi::reshape(inputs[0], param->newshape) };
+});
 
 
 /*!
@@ -431,7 +439,13 @@ the input array into an output array with the same shape as the second input arr
 .add_argument("data", "Tensor", "The input tensor.")
 .add_argument("shape_like", "Tensor", "Shape tensor.")
 .set_support_level(3)
-.add_type_rel("ReshapeLike", ReshapeLikeRel);
+.add_type_rel("ReshapeLike", ReshapeLikeRel)
+.set_attr<FTVMCompute>("FTVMCompute", [](const Attrs& attrs,
+                                         const Array<Tensor>& inputs,
+                                         const Type& out_type,
+                                         const Target& target) {
+  return Array<Tensor>{ topi::reshape(inputs[0], inputs[1]->shape) };
+});
 
 
 // Take

tests/python/relay/test_op_level3.py

@@ -123,8 +123,28 @@ def test_reshape_infer_type():
     assert yy.checked_type == relay.TensorType(
         (n, t, 2000), "float32")
 
+def test_reshape():
+    def verify_reshape(shape, oshape):
+        x_data = np.random.uniform(low=-1, high=1, size=shape).astype("float32")
+        ref_res = np.reshape(x_data, oshape)
 
-def test_reshape_like():
+        x = relay.var("x", relay.TensorType(shape, "float32"))
+        z = relay.reshape(x, newshape=ref_res.shape)
+        zz = relay.ir_pass.infer_type(z)
+        assert "newshape=" in z.astext()
+        assert zz.checked_type == relay.ty.TensorType(oshape, "float32")
+
+        func = relay.Function([x], z)
+
+        for target, ctx in ctx_list():
+            for kind in ["graph", "debug"]:
+                intrp = relay.create_executor(kind, ctx=ctx, target=target)
+                op_res = intrp.evaluate(func)(x_data)
+                tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
+    verify_reshape((2, 3, 4), (8, 3))
+    verify_reshape((4, 7), (2, 7, 2))
+
+def test_reshape_like_infer_type():
     # concrete shape
     x = relay.var("x", relay.TensorType((1, 2, 3), "float32"))
     y = relay.var("y", relay.TensorType((1,6), "float32"))
@@ -141,6 +161,29 @@ def test_reshape_like():
     assert zz.checked_type == relay.TensorType((1, 8, 8), "float32")
 
 
+def test_reshape_like():
+    def verify_reshape_like(shape, oshape):
+        x_data = np.random.uniform(low=-1, high=1, size=shape).astype("float32")
+        y_data = np.random.uniform(low=-1, high=1, size=oshape).astype("float32")
+        ref_res = np.reshape(x_data, y_data.shape)
+
+        x = relay.var("x", relay.TensorType(shape, "float32"))
+        y = relay.var("x", relay.TensorType(oshape, "float32"))
+        z = relay.reshape_like(x, y)
+        zz = relay.ir_pass.infer_type(z)
+        assert zz.checked_type == relay.ty.TensorType(ref_res.shape, "float32")
+
+        func = relay.Function([x, y], z)
+
+        for target, ctx in ctx_list():
+            for kind in ["graph", "debug"]:
+                intrp = relay.create_executor(kind, ctx=ctx, target=target)
+                op_res = intrp.evaluate(func)(x_data, y_data)
+                tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
+
+    verify_reshape_like((2, 3, 4), (1, 8, 3))
+    verify_reshape_like((4, 7), (2, 7, 2))
+
 def test_take_infer_type():
     def verify_take(dshape, indices_shape, oshape, axis=None):
         x = relay.var("x", relay.TensorType(dshape, "float32"))
@@ -318,6 +361,8 @@ if __name__ == "__main__":
     test_clip()
     test_transpose_infer_type()
     test_reshape_infer_type()
+    test_reshape()
+    test_reshape_like_infer_type()
     test_reshape_like()
     test_take_infer_type()
     test_full()