chore: import upstream snapshot with attribution

tests/python/relax/test_transform_decompose_ops.py

@@ -0,0 +1,389 @@
+# 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.
+# ruff: noqa: F401, F841
+
+from typing import Union
+
+import tvm
+import tvm.script
+import tvm.testing
+from tvm import IRModule, relax
+from tvm.relax import Function
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.script import tirx as T
+
+
+def test_batch_norm_inference():
+    @I.ir_module
+    class Before:
+        @R.function
+        def main(
+            x: R.Tensor((1, 64, 112, 112), "float32"),
+            gamma: R.Tensor((64,), "float32"),
+            beta: R.Tensor((64,), "float32"),
+            moving_mean: R.Tensor((64,), "float32"),
+            moving_var: R.Tensor((64,), "float32"),
+        ):
+            with R.dataflow():
+                bn = R.nn.batch_norm(
+                    x,
+                    gamma,
+                    beta,
+                    moving_mean,
+                    moving_var,
+                    axis=1,
+                    epsilon=1e-5,
+                    center=True,
+                    scale=True,
+                )
+                gv = bn[0]
+                R.output(gv)
+            return gv
+
+    @I.ir_module
+    class Expected:
+        @R.function
+        def main(
+            x: R.Tensor((1, 64, 112, 112), dtype="float32"),
+            gamma: R.Tensor((64,), dtype="float32"),
+            beta: R.Tensor((64,), dtype="float32"),
+            moving_mean: R.Tensor((64,), dtype="float32"),
+            moving_var: R.Tensor((64,), dtype="float32"),
+        ) -> R.Tensor((1, 64, 112, 112), dtype="float32"):
+            with R.dataflow():
+                lv: R.Tensor((1, 64, 1, 1), dtype="float32") = R.expand_dims(
+                    moving_mean, axis=[0, 2, 3]
+                )
+                lv1: R.Tensor((1, 64, 112, 112), dtype="float32") = R.subtract(x, lv)
+                lv2: R.Tensor((1, 64, 1, 1), dtype="float32") = R.expand_dims(
+                    moving_var, axis=[0, 2, 3]
+                )
+                lv3: R.Tensor((1, 64, 1, 1), dtype="float32") = R.add(
+                    lv2, R.const(9.9999997473787516e-06, "float32")
+                )
+                lv4: R.Tensor((1, 64, 1, 1), dtype="float32") = R.sqrt(lv3)
+                lv5: R.Tensor((1, 64, 112, 112), dtype="float32") = R.divide(lv1, lv4)
+                lv6: R.Tensor((1, 64, 1, 1), dtype="float32") = R.expand_dims(gamma, axis=[0, 2, 3])
+                lv7: R.Tensor((1, 64, 112, 112), dtype="float32") = R.multiply(lv5, lv6)
+                lv8: R.Tensor((1, 64, 1, 1), dtype="float32") = R.expand_dims(beta, axis=[0, 2, 3])
+                lv9: R.Tensor((1, 64, 112, 112), dtype="float32") = R.add(lv7, lv8)
+                bn: R.Tuple(
+                    R.Tensor((1, 64, 112, 112), dtype="float32"),
+                    R.Tensor((64,), dtype="float32"),
+                    R.Tensor((64,), dtype="float32"),
+                ) = (lv9, moving_mean, moving_var)
+                gv: R.Tensor((1, 64, 112, 112), dtype="float32") = bn[0]
+                R.output(gv)
+            return gv
+
+    After = relax.transform.DecomposeOpsForInference("main")(Before)
+    tvm.ir.assert_structural_equal(Expected, After)
+
+
+def test_batch_norm_training():
+    @I.ir_module
+    class Before:
+        @R.function
+        def main(
+            x: R.Tensor((1, 64, 112, 112), "float32"),
+            gamma: R.Tensor((64,), "float32"),
+            beta: R.Tensor((64,), "float32"),
+            moving_mean: R.Tensor((64,), "float32"),
+            moving_var: R.Tensor((64,), "float32"),
+        ):
+            with R.dataflow():
+                bn = R.nn.batch_norm(
+                    x,
+                    gamma,
+                    beta,
+                    moving_mean,
+                    moving_var,
+                    axis=1,
+                    epsilon=1e-5,
+                    center=True,
+                    scale=True,
+                    momentum=0.1,
+                )
+                gv0 = bn[0]
+                gv1 = bn[1]
+                gv2 = bn[2]
+                R.output(gv0, gv1, gv2)
+            return gv0, gv1, gv2
+
+    @I.ir_module
+    class Expected:
+        @R.function
+        def main(
+            x: R.Tensor((1, 64, 112, 112), dtype="float32"),
+            gamma: R.Tensor((64,), dtype="float32"),
+            beta: R.Tensor((64,), dtype="float32"),
+            moving_mean: R.Tensor((64,), dtype="float32"),
+            moving_var: R.Tensor((64,), dtype="float32"),
+        ) -> R.Tuple(
+            R.Tensor((1, 64, 112, 112), dtype="float32"),
+            R.Tensor((64,), dtype="float32"),
+            R.Tensor((64,), dtype="float32"),
+        ):
+            with R.dataflow():
+                # This portion is training-specific, computing the
+                # mean/variance of the dataset.
+                lv = R.mean(x, axis=[0, 2, 3], keepdims=False)
+                lv3 = R.variance(x, axis=[0, 2, 3], keepdims=False)
+
+                # This portion is identical to the batch_norm run during inference
+                lv1 = R.expand_dims(lv, axis=[0, 2, 3])
+                lv2 = R.subtract(x, lv1)
+                lv4 = R.expand_dims(lv3, axis=[0, 2, 3])
+                lv5 = R.add(lv4, R.const(9.9999997473787516e-06, "float32"))
+                lv6 = R.sqrt(lv5)
+                lv7 = R.divide(lv2, lv6)
+                lv8 = R.expand_dims(gamma, axis=[0, 2, 3])
+                lv9 = R.multiply(lv7, lv8)
+                lv10 = R.expand_dims(beta, axis=[0, 2, 3])
+                lv11 = R.add(lv9, lv10)
+                inner_tuple = (lv11, lv, lv3)
+                # This is the result that would be returned from a
+                # batch_norm at inference.
+
+                # However, at training we need to update the moving
+                # mean/variance, and to return those updated values.
+                inner_res = inner_tuple[0]
+                lv12 = R.multiply(R.const(0.89999997615814209, "float32"), moving_mean)
+                lv13 = R.multiply(R.const(0.10000000149011612, "float32"), lv)
+                lv14 = R.add(lv12, lv13)
+                lv15 = R.multiply(R.const(0.89999997615814209, "float32"), moving_var)
+                lv16 = R.multiply(R.const(0.10000000149011612, "float32"), lv3)
+                lv17 = R.add(lv15, lv16)
+                bn = (inner_res, lv14, lv17)
+                gv0 = bn[0]
+                gv1 = bn[1]
+                gv2 = bn[2]
+                R.output(gv0, gv1, gv2)
+            return (gv0, gv1, gv2)
+
+    After = relax.transform.DecomposeOpsForTraining("main")(Before)
+    tvm.ir.assert_structural_equal(Expected, After)
+
+
+def test_batch_norm_multiple_functions():
+    @I.ir_module
+    class Before:
+        @R.function
+        def main(
+            x: R.Tensor((1, 64, 112, 112), "float32"),
+            gamma: R.Tensor((64,), "float32"),
+            beta: R.Tensor((64,), "float32"),
+            moving_mean: R.Tensor((64,), "float32"),
+            moving_var: R.Tensor((64,), "float32"),
+        ):
+            with R.dataflow():
+                bn = R.nn.batch_norm(
+                    x,
+                    gamma,
+                    beta,
+                    moving_mean,
+                    moving_var,
+                    axis=1,
+                    epsilon=1e-5,
+                    center=True,
+                    scale=True,
+                )
+                gv = bn[0]
+                R.output(gv)
+            return gv
+
+        @R.function
+        def main1(
+            x: R.Tensor((1, 64, 112, 112), "float32"),
+            gamma: R.Tensor((64,), "float32"),
+            beta: R.Tensor((64,), "float32"),
+            moving_mean: R.Tensor((64,), "float32"),
+            moving_var: R.Tensor((64,), "float32"),
+        ):
+            with R.dataflow():
+                bn = R.nn.batch_norm(
+                    x,
+                    gamma,
+                    beta,
+                    moving_mean,
+                    moving_var,
+                    axis=1,
+                    epsilon=1e-5,
+                    center=True,
+                    scale=True,
+                )
+                gv = bn[0]
+                R.output(gv)
+            return gv
+
+    @I.ir_module
+    class Expected:
+        @R.function
+        def main1(
+            x: R.Tensor((1, 64, 112, 112), dtype="float32"),
+            gamma: R.Tensor((64,), dtype="float32"),
+            beta: R.Tensor((64,), dtype="float32"),
+            moving_mean: R.Tensor((64,), dtype="float32"),
+            moving_var: R.Tensor((64,), dtype="float32"),
+        ) -> R.Tensor((1, 64, 112, 112), dtype="float32"):
+            with R.dataflow():
+                lv10: R.Tensor((1, 64, 1, 1), dtype="float32") = R.expand_dims(
+                    moving_mean, axis=[0, 2, 3]
+                )
+                lv11: R.Tensor((1, 64, 112, 112), dtype="float32") = R.subtract(x, lv10)
+                lv12: R.Tensor((1, 64, 1, 1), dtype="float32") = R.expand_dims(
+                    moving_var, axis=[0, 2, 3]
+                )
+                lv13: R.Tensor((1, 64, 1, 1), dtype="float32") = R.add(
+                    lv12, R.const(9.9999997473787516e-06, "float32")
+                )
+                lv14: R.Tensor((1, 64, 1, 1), dtype="float32") = R.sqrt(lv13)
+                lv15: R.Tensor((1, 64, 112, 112), dtype="float32") = R.divide(lv11, lv14)
+                lv16: R.Tensor((1, 64, 1, 1), dtype="float32") = R.expand_dims(
+                    gamma, axis=[0, 2, 3]
+                )
+                lv17: R.Tensor((1, 64, 112, 112), dtype="float32") = R.multiply(lv15, lv16)
+                lv18: R.Tensor((1, 64, 1, 1), dtype="float32") = R.expand_dims(beta, axis=[0, 2, 3])
+                lv19: R.Tensor((1, 64, 112, 112), dtype="float32") = R.add(lv17, lv18)
+                bn: R.Tuple(
+                    R.Tensor((1, 64, 112, 112), dtype="float32"),
+                    R.Tensor((64,), dtype="float32"),
+                    R.Tensor((64,), dtype="float32"),
+                ) = (lv19, moving_mean, moving_var)
+                gv: R.Tensor((1, 64, 112, 112), dtype="float32") = bn[0]
+                R.output(gv)
+            return gv
+
+        @R.function
+        def main(
+            x: R.Tensor((1, 64, 112, 112), dtype="float32"),
+            gamma: R.Tensor((64,), dtype="float32"),
+            beta: R.Tensor((64,), dtype="float32"),
+            moving_mean: R.Tensor((64,), dtype="float32"),
+            moving_var: R.Tensor((64,), dtype="float32"),
+        ) -> R.Tensor((1, 64, 112, 112), dtype="float32"):
+            with R.dataflow():
+                lv: R.Tensor((1, 64, 1, 1), dtype="float32") = R.expand_dims(
+                    moving_mean, axis=[0, 2, 3]
+                )
+                lv1: R.Tensor((1, 64, 112, 112), dtype="float32") = R.subtract(x, lv)
+                lv2: R.Tensor((1, 64, 1, 1), dtype="float32") = R.expand_dims(
+                    moving_var, axis=[0, 2, 3]
+                )
+                lv3: R.Tensor((1, 64, 1, 1), dtype="float32") = R.add(
+                    lv2, R.const(9.9999997473787516e-06, "float32")
+                )
+                lv4: R.Tensor((1, 64, 1, 1), dtype="float32") = R.sqrt(lv3)
+                lv5: R.Tensor((1, 64, 112, 112), dtype="float32") = R.divide(lv1, lv4)
+                lv6: R.Tensor((1, 64, 1, 1), dtype="float32") = R.expand_dims(gamma, axis=[0, 2, 3])
+                lv7: R.Tensor((1, 64, 112, 112), dtype="float32") = R.multiply(lv5, lv6)
+                lv8: R.Tensor((1, 64, 1, 1), dtype="float32") = R.expand_dims(beta, axis=[0, 2, 3])
+                lv9: R.Tensor((1, 64, 112, 112), dtype="float32") = R.add(lv7, lv8)
+                bn: R.Tuple(
+                    R.Tensor((1, 64, 112, 112), dtype="float32"),
+                    R.Tensor((64,), dtype="float32"),
+                    R.Tensor((64,), dtype="float32"),
+                ) = (lv9, moving_mean, moving_var)
+                gv: R.Tensor((1, 64, 112, 112), dtype="float32") = bn[0]
+                R.output(gv)
+            return gv
+
+    After = relax.transform.DecomposeOpsForInference()(Before)
+    tvm.ir.assert_structural_equal(Expected, After)
+
+
+def test_layer_norm():
+    @I.ir_module
+    class Before:
+        @R.function
+        def main(
+            x: R.Tensor((4, 64, 112, 112), "float32"),
+            gamma: R.Tensor((112, 112), "float32"),
+            beta: R.Tensor((112, 112), "float32"),
+        ):
+            with R.dataflow():
+                ln = R.nn.layer_norm(
+                    x,
+                    gamma,
+                    beta,
+                    axes=[-2, -1],
+                    epsilon=1e-5,
+                    center=True,
+                    scale=True,
+                )
+                R.output(ln)
+            return ln
+
+    @I.ir_module
+    class Expected:
+        @R.function
+        def main(
+            x: R.Tensor((4, 64, 112, 112), dtype="float32"),
+            gamma: R.Tensor((112, 112), dtype="float32"),
+            beta: R.Tensor((112, 112), dtype="float32"),
+        ) -> R.Tensor((4, 64, 112, 112), dtype="float32"):
+            with R.dataflow():
+                lv: R.Tensor((4, 64, 1, 1), dtype="float32") = R.mean(
+                    x, axis=[-2, -1], keepdims=True
+                )
+                lv1: R.Tensor((4, 64, 112, 112), dtype="float32") = R.subtract(x, lv)
+                lv2: R.Tensor((4, 64, 1, 1), dtype="float32") = R.variance(
+                    x, axis=[-2, -1], keepdims=True
+                )
+                lv3: R.Tensor((4, 64, 1, 1), dtype="float32") = R.add(
+                    lv2, R.const(9.9999997473787516e-06, "float32")
+                )
+                lv4: R.Tensor((4, 64, 1, 1), dtype="float32") = R.sqrt(lv3)
+                lv5: R.Tensor((4, 64, 112, 112), dtype="float32") = R.divide(lv1, lv4)
+                lv6: R.Tensor((4, 64, 112, 112), dtype="float32") = R.multiply(lv5, gamma)
+                ln: R.Tensor((4, 64, 112, 112), dtype="float32") = R.add(lv6, beta)
+                R.output(ln)
+            return ln
+
+    After = relax.transform.DecomposeOpsForTraining()(Before)
+    tvm.ir.assert_structural_equal(Expected, After)
+
+
+def test_op_tensor_to_shape():
+    @I.ir_module
+    class Before:
+        @R.function
+        def main(t: R.Tensor([3], dtype="int64")):
+            gv: R.Shape(ndim=3) = R.tensor_to_shape(t)
+            return gv
+
+    @I.ir_module
+    class Expected:
+        @R.function
+        def main(t: R.Tensor([3], dtype="int64")) -> R.Shape(ndim=3):
+            x = T.int64()
+            x_1 = T.int64()
+            x_2 = T.int64()
+            gv: R.Shape(ndim=3) = R.call_pure_packed(
+                "vm.builtin.tensor_to_shape", t, ty_args=(R.Shape(ndim=3),)
+            )
+            y: R.Shape([x, x_1, x_2]) = R.match_cast(gv, R.Shape([x, x_1, x_2]))
+            gv_1: R.Shape([x, x_1, x_2]) = R.shape([x, x_1, x_2])
+            return gv_1
+
+    After = relax.transform.DecomposeOpsForInference()(Before)
+    tvm.ir.assert_structural_equal(Expected, After)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()