[REFACTOR][TEST] Replace CompareBeforeAfter for pytest compact (#18711)

This PR refactors test infrastructure by removing the CompareBeforeAfter
base class from tvm.testing and converting all dependent tests to use a
simpler, more explicit pattern.

We need this change as latest pytest do not allow calling fixture as
inner patterns which the previous CompareBeforeAfter depend on.

Author: tqchen

python/tvm/testing/utils.py

@@ -73,7 +73,6 @@ def test_something():
 import pickle
 import platform
 import sys
-import textwrap
 import time
 import shutil
 
@@ -1939,248 +1938,3 @@ def strtobool(val):
 def main():
     test_file = inspect.getsourcefile(sys._getframe(1))
     sys.exit(pytest.main([test_file] + sys.argv[1:]))
-
-
-class CompareBeforeAfter:
-    """Utility for comparing before/after of TIR transforms
-
-    A standard framework for writing tests that take a TIR PrimFunc as
-    input, apply a transformation, then either compare against an
-    expected output or assert that the transformation raised an error.
-    A test should subclass CompareBeforeAfter, defining class members
-    `before` / `Before`, `transform`, and `expected` / `Expected`.  CompareBeforeAfter will
-    then use these members to define a test method and test fixture.
-
-    `transform` may be one of the following.
-
-    - An instance of `tvm.ir.transform.Pass`
-
-    - A method that takes no arguments and returns a `tvm.ir.transform.Pass`
-
-    - A pytest fixture that returns a `tvm.ir.transform.Pass`
-
-    `before` / `Before` may be any one of the following.
-
-    - An instance of `tvm.tir.PrimFunc`.  This is allowed, but is not
-      the preferred method, as any errors in constructing the
-      `PrimFunc` occur while collecting the test, preventing any other
-      tests in the same file from being run.
-
-    - An TVMScript function, without the ``@T.prim_func`` decoration.
-      The ``@T.prim_func`` decoration will be applied when running the
-      test, rather than at module import.
-
-    - A method that takes no arguments and returns a `tvm.tir.PrimFunc`
-
-    - A pytest fixture that returns a `tvm.tir.PrimFunc`
-
-    `expected` / `Expected` may be any one of the following.  The type of
-    `expected` / `Expected` defines the test being performed.  If `expected`
-    provides a `tvm.tir.PrimFunc`, the result of the transformation
-    must match `expected`.  If `expected` is an exception, then the
-    transformation must raise that exception type.
-
-    - Any option supported for `before` / `Before`.
-
-    - The `Exception` class object, or a class object that inherits
-      from `Exception`.
-
-    - A method that takes no arguments and returns `Exception` or a
-      class object that inherits from `Exception`.
-
-    - A pytest fixture that returns `Exception` or an class object
-      that inherits from `Exception`.
-
-    Examples
-    --------
-
-    .. python::
-
-        class TestRemoveIf(tvm.testing.CompareBeforeAfter):
-            transform = tvm.tir.transform.Simplify()
-
-            def before(A: T.Buffer(1, "int32")):
-                if True:
-                    A[0] = 42
-                else:
-                    A[0] = 5
-
-            def expected(A: T.Buffer(1, "int32")):
-                A[0] = 42
-
-    """
-
-    check_well_formed: bool = True
-
-    def __init_subclass__(cls):
-        assert len([getattr(cls, name) for name in ["before", "Before"] if hasattr(cls, name)]) <= 1
-        assert (
-            len([getattr(cls, name) for name in ["expected", "Expected"] if hasattr(cls, name)])
-            <= 1
-        )
-        for name in ["before", "Before"]:
-            if hasattr(cls, name):
-                cls.before = cls._normalize_before(getattr(cls, name))
-                break
-        for name in ["expected", "Expected"]:
-            if hasattr(cls, name):
-                cls.expected = cls._normalize_expected(getattr(cls, name))
-                break
-        if hasattr(cls, "transform"):
-            cls.transform = cls._normalize_transform(cls.transform)
-
-    @classmethod
-    def _normalize_ir_module(cls, func):
-        if isinstance(func, (tvm.tir.PrimFunc, tvm.IRModule)):
-
-            def inner(self):
-                # pylint: disable=unused-argument
-                return func
-
-        elif cls._is_method(func):
-
-            def inner(self):
-                # pylint: disable=unused-argument
-                return func(self)
-
-        elif inspect.isclass(func):
-
-            def inner(self):
-                # pylint: disable=unused-argument
-                func_dict = {}
-                for name, method in func.__dict__.items():
-                    if name.startswith("_"):
-                        pass
-                    elif isinstance(method, tvm.ir.function.BaseFunc):
-                        func_dict[name] = method.with_attr("global_symbol", name)
-                    else:
-                        source_code = "@T.prim_func\n" + textwrap.dedent(inspect.getsource(method))
-                        prim_func = tvm.script.from_source(
-                            source_code, check_well_formed=self.check_well_formed
-                        )
-                        func_dict[name] = prim_func.with_attr("global_symbol", name)
-                return tvm.IRModule(func_dict)
-
-        else:
-
-            def inner(self):
-                # pylint: disable=unused-argument
-                source_code = "@T.prim_func\n" + textwrap.dedent(inspect.getsource(func))
-                return tvm.script.from_source(source_code, check_well_formed=self.check_well_formed)
-
-        return pytest.fixture(inner)
-
-    @classmethod
-    def _normalize_before(cls, func):
-        if hasattr(func, "_pytestfixturefunction"):
-            return func
-        else:
-            return cls._normalize_ir_module(func)
-
-    @classmethod
-    def _normalize_expected(cls, func):
-        if hasattr(func, "_pytestfixturefunction"):
-            return func
-
-        elif inspect.isclass(func) and issubclass(func, Exception):
-
-            def inner(self):
-                # pylint: disable=unused-argument
-                return func
-
-            return pytest.fixture(inner)
-
-        else:
-            return cls._normalize_ir_module(func)
-
-    @classmethod
-    def _normalize_transform(cls, transform):
-        def apply(module_transform):
-            def inner(obj):
-                if isinstance(obj, tvm.IRModule):
-                    return module_transform(obj)
-                elif isinstance(obj, tvm.tir.PrimFunc):
-                    mod = tvm.IRModule({"main": obj})
-                    mod = module_transform(mod)
-                    return mod["main"]
-                else:
-                    raise TypeError(f"Expected IRModule or PrimFunc, but received {type(obj)}")
-
-            return inner
-
-        if hasattr(transform, "_pytestfixturefunction"):
-            if not hasattr(cls, "_transform_orig"):
-                cls._transform_orig = transform
-
-            def inner(self, _transform_orig):
-                # pylint: disable=unused-argument
-                return apply(_transform_orig)
-
-        elif isinstance(transform, tvm.ir.transform.Pass):
-
-            def inner(self):
-                # pylint: disable=unused-argument
-                return apply(transform)
-
-        elif cls._is_method(transform):
-
-            def inner(self):
-                # pylint: disable=unused-argument
-                return apply(transform(self))
-
-        else:
-            raise TypeError(
-                "Expected transform to be a tvm.ir.transform.Pass, or a method returning a Pass"
-            )
-
-        return pytest.fixture(inner)
-
-    @staticmethod
-    def _is_method(func):
-        return callable(func) and "self" in inspect.signature(func).parameters
-
-    def test_compare(self, before, expected, transform):
-        """Unit test to compare the expected TIR PrimFunc to actual"""
-
-        if inspect.isclass(expected) and issubclass(expected, Exception):
-            with pytest.raises(expected):
-                after = transform(before)
-
-                # This portion through pytest.fail isn't strictly
-                # necessary, but gives a better error message that
-                # includes the before/after.
-                before_str = before.script(name="before")
-                after_str = after.script(name="after")
-
-                pytest.fail(
-                    msg=(
-                        f"Expected {expected.__name__} to be raised from transformation, "
-                        f"instead received TIR\n:{before_str}\n{after_str}"
-                    )
-                )
-
-        elif isinstance(expected, (tvm.tir.PrimFunc, tvm.ir.IRModule)):
-            after = transform(before)
-
-            try:
-                # overwrite global symbol so it doesn't come up in the comparison
-                if isinstance(after, tvm.tir.PrimFunc):
-                    after = after.with_attr("global_symbol", "main")
-                    expected = expected.with_attr("global_symbol", "main")
-                tvm.ir.assert_structural_equal(after, expected)
-            except ValueError as err:
-                before_str = before.script(name="before")
-                after_str = after.script(name="after")
-                expected_str = expected.script(name="expected")
-                raise ValueError(
-                    f"TIR after transformation did not match expected:\n"
-                    f"{before_str}\n{after_str}\n{expected_str}"
-                ) from err
-
-        else:
-            raise TypeError(
-                f"tvm.testing.CompareBeforeAfter requires the `expected` fixture "
-                f"to return either `Exception`, an `Exception` subclass, "
-                f"or an instance of `tvm.tir.PrimFunc`.  "
-                f"Instead, received {type(expected)}."
-            )

tests/python/dlight/test_cpu_gemv.py

@@ -23,20 +23,9 @@
 from tvm.target import Target
 
 
-class BaseBeforeAfter(tvm.testing.CompareBeforeAfter):
-    @pytest.fixture
-    def transform(self):
-        def transform(mod):
-            with Target("llvm"):
-                return dl.ApplyDefaultSchedule(dl.cpu.GEMV())(mod)
-
-        return transform
-
-
-class TestGEMV(BaseBeforeAfter):
+def test_gemv_basic():
     # fmt: off
-
-    @T.prim_func
+    @T.prim_func(private=True)
     def before(lv1637: T.Buffer((1, 32, 1, 128), "float16"), p_lv1638: T.handle, p_lv1614: T.handle, p_output0: T.handle):
         T.func_attr({"tir.noalias": True})
         n = T.int32()
@@ -81,7 +70,7 @@ def before(lv1637: T.Buffer((1, 32, 1, 128), "float16"), p_lv1638: T.handle, p_l
                 T.writes(var_compute_intermediate[v_i0, v_i1, v_i2, v_i3])
                 var_compute_intermediate[v_i0, v_i1, v_i2, v_i3] = T.Cast("float32", var_T_minimum_intermediate[v_i0, v_i1, v_i2, v_i3])
 
-    @T.prim_func
+    @T.prim_func(private=True)
     def expected(lv1637: T.Buffer((1, 32, 1, 128), "float16"), p_lv1638: T.handle, p_lv1614: T.handle, p_output0: T.handle):
         T.func_attr({"tir.is_scheduled": True, "tir.noalias": True})
         n = T.int32()
@@ -114,6 +103,11 @@ def expected(lv1637: T.Buffer((1, 32, 1, 128), "float16"), p_lv1638: T.handle, p
 
     # fmt: on
 
+    mod = tvm.IRModule({"main": before})
+    with Target("llvm"):
+        mod = dl.ApplyDefaultSchedule(dl.cpu.GEMV())(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
 
 def test_decode_gemv_256_threads():
     # fmt: off

tests/python/dlight/test_gpu_conv.py

@@ -15,28 +15,16 @@
 # specific language governing permissions and limitations
 # under the License.
 # pylint: disable=missing-docstring
-import pytest
-
+import tvm
 import tvm.testing
 from tvm import dlight as dl
 from tvm.script import tir as T
 from tvm.target import Target
 
 
-class BaseBeforeAfter(tvm.testing.CompareBeforeAfter):
-    @pytest.fixture
-    def transform(self):
-        def transform(mod):
-            with Target("nvidia/geforce-gtx-1080-ti"):
-                # Use Matmul rule for Conv for now
-                return dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
-
-        return transform
-
-
-class TestConv3d(BaseBeforeAfter):
+def test_conv3d():
     # fmt: off
-    @T.prim_func
+    @T.prim_func(private=True)
     def before(
         A: T.Buffer((14308, 3, 2, 14, 14), "float16"),
         W: T.Buffer((1280, 3, 2, 14, 14), "float16"),
@@ -54,7 +42,7 @@ def before(
                     C[v_nn, v_ff, v_yy, v_xx, v_zz] = T.float16(0.0)
                 C[v_nn, v_ff, v_yy, v_xx, v_zz] += pad_A[v_nn, v_rc, v_yy * 2 + v_ry, v_xx * 14 + v_rx, v_zz * 14 + v_rz]* W[v_ff, v_rc, v_ry, v_rx, v_rz]
 
-    @T.prim_func
+    @T.prim_func(private=True)
     def expected(A: T.Buffer((14308, 3, 2, 14, 14), "float16"), W: T.Buffer((1280, 3, 2, 14, 14), "float16"), C: T.Buffer((14308, 1280, 1, 1, 1), "float16")):
         T.func_attr({"tir.is_scheduled": True})
         # with T.sblock("root"):
@@ -113,6 +101,11 @@ def expected(A: T.Buffer((14308, 3, 2, 14, 14), "float16"), W: T.Buffer((1280, 3
                                             C[v1, v2, 0, 0, 0] = C_reindex_pad_local[v0, v1, v2]
     # fmt: on
 
+    mod = tvm.IRModule({"main": before})
+    with Target("nvidia/geforce-gtx-1080-ti"):
+        mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
+    tvm.ir.assert_structural_equal(mod["main"], expected)
+
 
 if __name__ == "__main__":
     tvm.testing.main()