[Fix][CI] Isolate dynamo state in torch.compile tests (#1743)

Closes #1741

## Summary

- Add a shared `isolated_dynamo` fixture in `tests/conftest.py` that
runs `torch._dynamo.reset()` before and after every `torch.compile`
test.
- Request the fixture from all compile tests (`test_compile.py`,
`test_elementwise_compile.py`, `test_norm_ops.py`, `test_pool.py`),
replacing the module-local reset fixture in
`test_elementwise_compile.py`.
- Fixes `FailOnRecompileLimitHit` on all three `test_mha_kernel_compile`
cases in the push-tier gpu-smoke run (cross-test dynamo cache pollution,
not intrinsic guard churn).

## Test plan

- [x] Modified files pass unit tests (297/297 nodes across the 4 touched
test files at 8de7df1e; pre-commit green)
- [x] All three `test_mha_kernel_compile` parametrizations pass under
the push-tier selection `-m "smoke or full" -k compile` with push-tier
ordering on H200 (11/11; 3 failed pre-fix). CI push-tier confirmation
lands on the post-merge run — PR-event gpu-smoke runs `-m smoke` only.
- [x] This PR documents the flip commit/mechanism and the layer choice
(see Regression)
- [ ] Subsequent push-to-main GPU Smoke run is green (verifiable only
after merge)

## Regression

- **Flip commit**: c9229cab (#1731) added
`test_max_pool1d/3d_compile_fullgraph` (2 params each, smoke tier),
raising pre-MHA push-tier compiles from 4 (norm 2 + pool2d 2) to exactly
8.
- **Mechanism**: dynamo's recompile cache is keyed per code object.
`torch.compile` of any plain callable (every TileOps `Op` instance via
`Op.__call__`) shares torch's single wrapper frame `inner`
(`torch/_dynamo/external_utils.py:66`), so each compiled op instance in
one pytest process consumes one of the 8 `cache_size_limit` slots of
that shared frame. The first `test_mha_kernel_compile` case was the 9th
entry → `FailOnRecompileLimitHit` under `fullgraph=True`; all three MHA
parametrizations failed.
- **Reproduction**: `pytest tests/ops/test_norm_ops.py
tests/ops/test_pool.py tests/test_compile.py -m "smoke or full" -k
compile` — 3 failed pre-fix, 11 passed post-fix on H200.
`tests/test_compile.py` alone passed pre-fix, confirming cross-test
cache pollution rather than intrinsic guard churn.
- **Why test-isolation layer (not op wrapper, not a limit raise)**: the
exhausted code object is torch-owned and shared across ALL compiled
plain callables, so no `tileops/ops/` wrapper change can prevent
distinct op instances consuming per-code-object cache slots; each
individual compile is legitimate (one entry per instance, no intra-test
recompilation). Fix: shared `isolated_dynamo` fixture in
`tests/conftest.py` (`torch._dynamo.reset()` before/after) requested by
every `torch.compile` test. Raising `cache_size_limit` was rejected — it
would mask the next capacity flip instead of removing cross-test
coupling.

## Test node delta

```
File                                     Base    HEAD    Delta
--------------------------------------------------------------
tests/conftest.py                           0       0        0
tests/ops/test_elementwise_compile.py       0      84      +84
tests/ops/test_norm_ops.py                  5       5        0
tests/ops/test_pool.py                    205     205        0
tests/test_compile.py                       3       3        0
--------------------------------------------------------------
TOTAL                                     213     297      +84
```

**Justification:** no real growth. The +84 on
`tests/ops/test_elementwise_compile.py` is a base-side pytest collection
failure in the delta tool (warning emitted); the file has the same 84
nodes at base and HEAD — the diff only rewires an autouse fixture (2
insertions, 5 deletions, no test functions added).

---------

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)