[tvm] branch main updated: [TVMC] Apply constant folding when converting layout (#13216)

[mo...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

mousius pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 4ecf303695 [TVMC] Apply constant folding when converting layout (#13216)
4ecf303695 is described below

commit 4ecf3036951f153a578fbf4685542a47b0068733
Author: Luke Hutton <lu...@arm.com>
AuthorDate: Wed Nov 2 14:15:00 2022 +0000

    [TVMC] Apply constant folding when converting layout (#13216)
    
    This commit ensures that constant folding is applied when a desired
    layout is selected during compilation. It ensures that
    `layout_transform` operations are removed where possible so that
    pattern matching for BYOC backends can work effectively.
    
    A test has been added to check this regression.
---
 python/tvm/driver/tvmc/transform.py        |  1 +
 tests/python/driver/tvmc/test_transform.py | 57 ++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/python/tvm/driver/tvmc/transform.py b/python/tvm/driver/tvmc/transform.py
index 3f77765778..51c9e52f21 100644
--- a/python/tvm/driver/tvmc/transform.py
+++ b/python/tvm/driver/tvmc/transform.py
@@ -50,6 +50,7 @@ def convert_graph_layout(mod, desired_layout):
         [
             relay.transform.RemoveUnusedFunctions(),
             relay.transform.ConvertLayout(desired_layouts),
+            relay.transform.FoldConstant(),
         ]
     )
 
diff --git a/tests/python/driver/tvmc/test_transform.py b/tests/python/driver/tvmc/test_transform.py
new file mode 100644
index 0000000000..98a0210a1b
--- /dev/null
+++ b/tests/python/driver/tvmc/test_transform.py
@@ -0,0 +1,57 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# 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
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import pytest
+import numpy as np
+
+import tvm
+from tvm import relay
+from tvm.driver.tvmc.transform import convert_graph_layout
+
+
+def test_layout_transform():
+    """
+    Test layout is correctly transformed and constant folding is applied.
+    """
+    dtype = "int8"
+    iinfo = np.iinfo(dtype)
+    data_min = iinfo.min
+    data_max = iinfo.max
+
+    x = relay.var("x", shape=(1, 4, 2, 2), dtype=dtype)
+    weight = relay.const(
+        np.random.randint(data_min, data_max, size=(2, 4, 2, 2), dtype=dtype), dtype=dtype
+    )
+    x = relay.nn.conv2d(x, weight)
+    func = relay.Function(relay.analysis.free_vars(x), x)
+    mod = tvm.IRModule.from_expr(func)
+
+    desired_layout = "NHWC"
+    mod = convert_graph_layout(mod, desired_layout)
+
+    main_expr = mod["main"].body
+    conv = main_expr.args[0]
+    assert conv.op.name == "nn.conv2d"
+    assert conv.attrs["data_layout"] == "NHWC"
+    assert conv.attrs["kernel_layout"] == "HWIO"
+
+    # Ensure transform has been folded into the constant
+    weights = conv.args[1]
+    assert isinstance(weights, relay.expr.Constant)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()