[Fix][CI] Yield explicitly in _reset_dynamo autouse fixture

Co-Authored-By: Ibuki 🍃 — a wind born from GPTs <Ibuki-wind@users.noreply.github.com>

tests/conftest.py

@@ -47,6 +47,24 @@ def setup() -> None:
         torch.cuda.manual_seed_all(1235)
 
 
+@pytest.fixture
+def isolated_dynamo():
+    """Reset torch._dynamo state around a test that calls ``torch.compile``.
+
+    Dynamo's recompile cache is keyed per code object, and every
+    ``torch.compile``-d plain callable (any non-``nn.Module``, e.g. a TileOps
+    ``Op`` instance) shares torch's single wrapper frame. Each compiled op
+    instance therefore consumes one slot of that frame's shared
+    ``cache_size_limit`` (default 8) for the whole pytest process, so compile
+    tests pollute each other's cache and later ``fullgraph=True`` tests fail
+    with ``FailOnRecompileLimitHit``. Request this fixture from every test
+    that calls ``torch.compile``.
+    """
+    torch._dynamo.reset()
+    yield
+    torch._dynamo.reset()
+
+
 NON_RUNTIME_OPS_TIER_FILES = {
     "tests/ops/test_elementwise_caching_autotune.py",
     "tests/ops/test_elementwise_compile.py",

tests/ops/test_elementwise_compile.py

@@ -81,11 +81,9 @@ from tileops.ops.elementwise import (
 
 
 @pytest.fixture(autouse=True)
-def _reset_dynamo():
-    """Reset torch._dynamo before each test to avoid recompile limit."""
-    torch._dynamo.reset()
+def _reset_dynamo(isolated_dynamo):
+    """Isolate dynamo state for every compile test in this module."""
     yield
-    torch._dynamo.reset()
 
 
 # ---------------------------------------------------------------------------

tests/ops/test_norm_ops.py

@@ -57,8 +57,7 @@ class TestBatchNormFwdValidation:
 # ---------------------------------------------------------------------------
 
 @pytest.mark.smoke
-
-
+@pytest.mark.usefixtures("isolated_dynamo")
 class TestBatchNormCustomOp:
 
     def test_fwd_torch_compile_smoke(self):

tests/ops/test_pool.py

@@ -1386,6 +1386,7 @@ def test_max_pool1d_dynamic_shape_kernel_cache_and_roofline(
 
 @pytest.mark.smoke
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.usefixtures("isolated_dynamo")
 @pytest.mark.parametrize(
     ("op_cls", "return_indices"),
     [
@@ -1962,6 +1963,7 @@ def test_max_pool3d_dynamic_shape_kernel_cache_and_roofline(
 
 @pytest.mark.smoke
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.usefixtures("isolated_dynamo")
 @pytest.mark.parametrize(
     ("op_cls", "return_indices"),
     [
@@ -2559,6 +2561,7 @@ def test_max_pool2d_dynamic_shape_kernel_cache_and_roofline(
 
 @pytest.mark.smoke
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.usefixtures("isolated_dynamo")
 @pytest.mark.parametrize(
     ("op_cls", "return_indices"),
     [

tests/test_compile.py

@@ -20,6 +20,7 @@ class MhaCompileFixture(FixtureBase):
 
 
 @pytest.mark.full
+@pytest.mark.usefixtures("isolated_dynamo")
 @MhaCompileFixture
 def test_mha_kernel_compile(B: int, S: int, H: int, D: int, causal: bool, dtype: torch.dtype):
     test = MhaFwdTest(B, H, S, D, causal, dtype)