From 792ff1ab546b5d15a5d8810f379deabf128f4a51 Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton@arm.com>
Date: Wed, 26 Oct 2022 20:15:13 +0000
Subject: [PATCH] [TVMC] Apply constant folding when converting layout

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.

Change-Id: Ie7646a2c8fdb5ec17dcbeb0876f969f8b8218b6e
---
 python/tvm/driver/tvmc/transform.py        |  1 +
 tests/python/driver/tvmc/test_transform.py | 57 ++++++++++++++++++++++
 2 files changed, 58 insertions(+)
 create mode 100644 tests/python/driver/tvmc/test_transform.py

diff --git a/python/tvm/driver/tvmc/transform.py b/python/tvm/driver/tvmc/transform.py
index 3f7776577876..51c9e52f21d6 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 000000000000..98a0210a1bb6
--- /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()