This is an automated email from the ASF dual-hosted git repository.
bohan pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 34e87ae7f8 [Unity] Fix ConvertLayout on binary elemwise ops involving
scalar input (#14961)
34e87ae7f8 is described below
commit 34e87ae7f8ad17eab5c4556e69d19cf859720113
Author: masahi <[email protected]>
AuthorDate: Wed May 31 04:07:12 2023 +0900
[Unity] Fix ConvertLayout on binary elemwise ops involving scalar input
(#14961)
---
src/relax/op/tensor/binary.cc | 9 +++++
src/relax/transform/convert_layout.cc | 2 +-
.../python/relax/test_transform_convert_layout.py | 45 +++++++++++++++++++++-
3 files changed, 54 insertions(+), 2 deletions(-)
diff --git a/src/relax/op/tensor/binary.cc b/src/relax/op/tensor/binary.cc
index 96d1f01e8a..d44b6da629 100644
--- a/src/relax/op/tensor/binary.cc
+++ b/src/relax/op/tensor/binary.cc
@@ -92,10 +92,19 @@ InferLayoutOutput InferLayoutBinaryEwise(const Call& call,
<< "Unknown dim tensors should not be handled by this function";
if (x1_sinfo->ndim <= x2_sinfo->ndim) {
+ if (x1_sinfo->ndim == 0) {
+ LayoutDecision out_layout = layout2;
+ return InferLayoutOutput({LayoutDecision(""), layout2}, {out_layout},
Attrs(call->attrs));
+ }
LayoutDecision out_layout = FollowDecision(layout1, x2_sinfo->ndim);
return InferLayoutOutput({layout1, out_layout}, {out_layout},
Attrs(call->attrs));
} else {
+ if (x2_sinfo->ndim == 0) {
+ LayoutDecision out_layout = layout1;
+ return InferLayoutOutput({layout1, LayoutDecision("")}, {out_layout},
Attrs(call->attrs));
+ }
LayoutDecision out_layout = FollowDecision(layout2, x1_sinfo->ndim);
+
return InferLayoutOutput({out_layout, layout2}, {out_layout},
Attrs(call->attrs));
}
}
diff --git a/src/relax/transform/convert_layout.cc
b/src/relax/transform/convert_layout.cc
index 4f36cfbc0f..91dcd5d8e8 100644
--- a/src/relax/transform/convert_layout.cc
+++ b/src/relax/transform/convert_layout.cc
@@ -90,7 +90,7 @@ class LayoutConvertMutator : public ExprMutator {
Expr RewriteExpr(const Expr& expr, const NLayout& to) {
auto fvisitleaf = [&](const Expr& expr, std::array<NLayout, 2> layouts) ->
Expr {
NLayout from = layouts[0], to = layouts[1];
- if (NLayoutEqual()(from, to)) return expr;
+ if (NLayoutEqual()(from, to) || layouts[0].LeafValue()->layout->name ==
"") return expr;
// If not both from and to are unknown, then none of them can be unknown.
ICHECK(!NLayoutEqual()(from, LayoutDecision::InitUnknownDim()) &&
!NLayoutEqual()(to, LayoutDecision::InitUnknownDim()))
diff --git a/tests/python/relax/test_transform_convert_layout.py
b/tests/python/relax/test_transform_convert_layout.py
index 5187ab30b7..570a53b48f 100644
--- a/tests/python/relax/test_transform_convert_layout.py
+++ b/tests/python/relax/test_transform_convert_layout.py
@@ -24,7 +24,6 @@ from tvm.script.parser import ir as I, relax as R, tir as T
def verify(input, expected):
mod = ConvertLayout({"relax.nn.conv2d": ["NHWC", "OHWI"]})(input)
mod = Normalize()(mod)
- print(mod.script())
tvm.ir.assert_structural_equal(mod, expected)
@@ -1402,5 +1401,49 @@ def test_binary_broadcast():
verify(Input, Expected)
+def test_binary_ewise_scalar():
+ @I.ir_module
+ class Input:
+ @R.function
+ def main(
+ x: R.Tensor((2, 3, 28, 28), "float32"), w: R.Tensor((4, 3, 3, 3),
"float32")
+ ) -> R.Tensor(None, "float32", ndim=4):
+ with R.dataflow():
+ gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w,
out_dtype="float32")
+ gv2: R.Tensor((2, 4, 26, 26), "float32") = R.add(gv,
R.const(1, "float32"))
+ R.output(gv2)
+ return gv2
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(
+ x: R.Tensor((2, 3, 28, 28), dtype="float32"), w: R.Tensor((4, 3,
3, 3), dtype="float32")
+ ) -> R.Tensor((2, 4, 26, 26), dtype="float32"):
+ with R.dataflow():
+ lv: R.Tensor((2, 28, 28, 3), dtype="float32") =
R.permute_dims(x, axes=[0, 2, 3, 1])
+ lv1: R.Tensor((4, 3, 3, 3), dtype="float32") =
R.permute_dims(w, axes=[0, 2, 3, 1])
+ gv: R.Tensor((2, 26, 26, 4), dtype="float32") = R.nn.conv2d(
+ lv,
+ lv1,
+ strides=[1, 1],
+ padding=[0, 0, 0, 0],
+ dilation=[1, 1],
+ groups=1,
+ data_layout="NHWC",
+ kernel_layout="OHWI",
+ out_layout="NHWC",
+ out_dtype="float32",
+ )
+ lv2: R.Tensor((2, 26, 26, 4), dtype="float32") = R.add(gv,
R.const(1, "float32"))
+ gv2: R.Tensor((2, 4, 26, 26), dtype="float32") =
R.permute_dims(
+ lv2, axes=[0, 3, 1, 2]
+ )
+ R.output(gv2)
+ return gv2
+
+ verify(Input, Expected)
+
+
if __name__ == "__main__":
tvm.testing.main()
