from __future__ import annotations

import collections
import contextlib
import functools
import importlib
import inspect
import io
import logging
import operator
import os
from pathlib import Path
import re
import sys
from typing import TYPE_CHECKING
from typing import Callable
from typing import Generator
from typing import Sequence
from typing import cast
import unittest

import pytest
import torch
import torch.distributed as dist
from torch.utils._pytree import tree_map

from ._compat import get_mtia_tunable_fragments
from ._compat import get_tensor_descriptor_fn_name
from ._compat import requires_torch_version
from ._compat import supports_amd_cdna_tunables
from ._compat import supports_tensor_descriptor
from ._utils import counters
from .autotuner.benchmarking import sync_object as sync_object
from .runtime.settings import _get_backend

if _get_backend() == "pallas":
    from .autotuner.benchmarking import compute_repeat_generic as compute_repeat
    from .autotuner.benchmarking import do_bench_generic as do_bench
    from .autotuner.benchmarking import interleaved_bench_generic as interleaved_bench
else:
    from .autotuner.benchmarking import compute_repeat
    from .autotuner.benchmarking import do_bench as do_bench
    from .autotuner.benchmarking import interleaved_bench

from .runtime.config import Config
from .runtime.ref_mode import is_ref_mode_enabled
from .runtime.settings import RefMode

if TYPE_CHECKING:
    import types

    from .runtime.kernel import Kernel


def _strip_launcher_args(value: str) -> str:
    strip_pairs = []
    if supports_amd_cdna_tunables():
        strip_pairs += [
            (r", waves_per_eu=\d+", ""),
            (r", matrix_instr_nonkdim=\d+", ""),
        ]
    if _get_backend() == "tileir":
        strip_pairs += [(r", num_ctas=\d+", ""), (r", occupancy=\d+", "")]
    for tunable in get_mtia_tunable_fragments():
        # Match tunable=value patterns:
        # - Quoted strings: 'value' or "value"
        # - Unquoted values (bools, numbers): stops at comma or closing paren
        strip_pairs.append((rf", {tunable}=(?:'[^']*'|\"[^\"]*\"|[^,)]+)", ""))
    for pattern, replacement in strip_pairs:
        value = re.sub(pattern, replacement, value)
    return value


def _get_triton_backend() -> str | None:
    try:
        import triton

        # pyrefly: ignore [missing-attribute]
        return triton.runtime.driver.active.get_current_target().backend
    except Exception:
        return None


def skipIfFn(
    cond_fn: Callable[[], bool], reason: str
) -> Callable[[Callable], Callable]:
    """Decorator that evaluates skip condition at test execution time.

    Unlike unittest.skipIf which evaluates at decoration time, this defers
    evaluation to e.g. avoid CUDA init during pytest-xdist collection.

    Works on both test methods and test classes. When applied to a class,
    wraps setUp to check the skip condition before each test runs.
    """

    def decorator(test_item: Callable) -> Callable:
        if isinstance(test_item, type):
            # For classes: wrap setUp to check skip condition at test execution time
            original_setUp = test_item.setUp  # pyrefly: ignore [missing-attribute]

            @functools.wraps(original_setUp)
            def new_setUp(self: object, *args: object, **kwargs: object) -> object:
                if cond_fn():
                    assert isinstance(self, unittest.TestCase)
                    self.skipTest(reason)
                return original_setUp(self, *args, **kwargs)

            test_item.setUp = new_setUp  # type: ignore[attr-defined]
            return test_item
        # For functions/methods

        @functools.wraps(test_item)
        def wrapper(*args: object, **kwargs: object) -> object:
            if cond_fn():
                # Use self.skipTest() when called on a TestCase method so that
                # RefEagerTestBase's patched skipTest counter is incremented.
                if args and isinstance(args[0], unittest.TestCase):
                    args[0].skipTest(reason)
                else:
                    raise unittest.SkipTest(reason)
            return test_item(*args, **kwargs)

        return wrapper

    return decorator


def xfailIfFn(
    cond_fn: Callable[[], bool], reason: str
) -> Callable[[Callable], Callable]:
    """Decorator that marks tests expected-failure when condition is true."""

    def decorator(test_item: Callable) -> Callable:
        if not cond_fn():
            return test_item

        if isinstance(test_item, type):
            for name, value in vars(test_item).items():
                if name.startswith("test") and callable(value):
                    setattr(test_item, name, unittest.expectedFailure(value))
            return test_item

        decorated = unittest.expectedFailure(test_item)
        decorated.__doc__ = (
            f"{decorated.__doc__ or ''}\n[xfailIfFn reason] {reason}".rstrip()
        )
        return decorated

    return decorator


def is_mtia() -> bool:
    """Return True if running on MTIA."""
    return _get_triton_backend() == "mtia"


def skipIfMTIA(reason: str) -> Callable[[Callable], Callable]:
    # Defers check to test execution time to avoid CUDA init during pytest-xdist collection.
    return skipIfFn(is_mtia, reason)


class _LogCapture(logging.Handler):
    """Simple logging handler to capture log records."""

    def __init__(self) -> None:
        super().__init__()
        self.records: list[logging.LogRecord] = []

    def emit(self, record: logging.LogRecord) -> None:
        self.records.append(record)

    def clear(self) -> None:
        self.records.clear()


class _OutputCapture:
    """Simple output capture class for stdout/stderr."""

    def __init__(self) -> None:
        self.stdout = io.StringIO()
        self.stderr = io.StringIO()

    def readouterr(self) -> tuple[str, str]:
        """Read and clear captured output, returning (stdout, stderr) tuple."""
        stdout_val = self.stdout.getvalue()
        stderr_val = self.stderr.getvalue()
        # Clear the buffers
        self.stdout.seek(0)
        self.stdout.truncate()
        self.stderr.seek(0)
        self.stderr.truncate()
        return (stdout_val, stderr_val)


def is_cuda() -> bool:
    """Return True if running on CUDA (NVIDIA GPU)."""
    return _get_triton_backend() == "cuda" and torch.cuda.is_available()


PROJECT_ROOT: Path = Path(__file__).parent.parent
EXAMPLES_DIR: Path = PROJECT_ROOT / "examples"
DEVICE = None


def _has_mtia_runtime() -> bool:
    try:
        # is_mtia() calls _get_triton_backend() which triggers CUDA init on CUDA devices,
        # so we first try importing MTIA lib to make sure we are in MTIA-available environment.
        import mtia.host_runtime.torch_mtia.dynamic_library  # noqa: F401  # pyrefly: ignore[missing-import]

        return is_mtia()
    except ImportError:
        return False


def _init_tpu_device() -> bool:
    """Try to initialize the TPU device. Returns True if successful."""
    try:
        import torch_tpu.api  # type: ignore[import-not-found]

        torch_tpu.api.tpu_device()
        return True
    except (ImportError, RuntimeError):
        return False


# Determine DEVICE without calling functions that initialize CUDA.
if _get_backend() == "pallas":
    _init_tpu_device()
    DEVICE = torch.device("tpu")
elif torch.xpu.is_available():
    DEVICE = torch.device("xpu")
elif _has_mtia_runtime():
    DEVICE = torch.device("mtia")
else:
    DEVICE = torch.device("cuda")

# Half-precision dtype: bfloat16 on TPU (float16 not supported), float16 elsewhere
if _get_backend() == "pallas":
    HALF_DTYPE = torch.bfloat16
else:
    HALF_DTYPE = torch.float16


def get_nvidia_gpu_model() -> str:
    """
    Retrieves the model of the NVIDIA GPU being used.
    Will return the name of the current device.
    Returns:
        str: The model of the NVIDIA GPU or empty str if not found.
    """
    if torch.cuda.is_available():
        props = torch.cuda.get_device_properties(torch.cuda.current_device())
        return getattr(props, "name", "")
    return ""


def skipIfRefEager(reason: str) -> Callable[[Callable], Callable]:
    """Skip test if running in ref eager mode (HELION_INTERPRET=1)."""
    return unittest.skipIf(os.environ.get("HELION_INTERPRET") == "1", reason)


def skipIfNormalMode(reason: str) -> Callable[[Callable], Callable]:
    """Skip test if running in normal mode (i.e. if HELION_INTERPRET=1 is not set)."""
    return unittest.skipIf(os.environ.get("HELION_INTERPRET") != "1", reason)


def skipIfRocm(reason: str) -> Callable[[Callable], Callable]:
    """Skip test if running with rocm"""
    return unittest.skipIf(torch.version.hip is not None, reason)


def skipIfTileIR(reason: str) -> Callable[[Callable], Callable]:
    """Skip test if running with tileir"""
    # Defers check to test execution time to avoid CUDA init during pytest-xdist collection.
    return skipIfFn(lambda: _get_backend() == "tileir", reason)


def skipIfPallas(reason: str) -> Callable[[Callable], Callable]:
    """Skip test if running with pallas"""
    # Defers check to test execution time to avoid CUDA init during pytest-xdist collection.
    return skipIfFn(lambda: _get_backend() == "pallas", reason)


def xfailIfPallas(reason: str) -> Callable[[Callable], Callable]:
    """Mark test as expected failure if running with pallas"""
    return xfailIfFn(lambda: _get_backend() == "pallas", reason)


def skipUnlessAMDCDNA(reason: str) -> Callable[[Callable], Callable]:
    """Skip test unless running on AMD CDNA architecture."""
    from helion._compat import supports_amd_cdna_tunables

    # Defers check to test execution time to avoid CUDA init during pytest-xdist collection.
    return skipIfFn(lambda: not supports_amd_cdna_tunables(), reason)


def skipUnlessMTIA(reason: str) -> Callable[[Callable], Callable]:
    """Skip test unless running on MTIA hardware."""
    from ._compat import supports_mtia_tunables

    # Defers check to test execution time to avoid CUDA init during pytest-xdist collection.
    return skipIfFn(lambda: not supports_mtia_tunables(), reason)


def skipUnlessTileIR(reason: str) -> Callable[[Callable], Callable]:
    """Skip test unless running on tileir"""
    # Defers check to test execution time to avoid CUDA init during pytest-xdist collection.
    return skipIfFn(lambda: _get_backend() != "tileir", reason)


@functools.cache
def _has_cute_dsl() -> bool:
    try:
        import cutlass.cute as _cute  # noqa: F401
    except ModuleNotFoundError:
        return False
    return True


def skipUnlessCuteAvailable(reason: str) -> Callable[[Callable], Callable]:
    """Skip test unless CUTLASS CuTe Python DSL is importable."""
    return skipIfFn(lambda: not _has_cute_dsl(), reason)


def xfailIfCute(reason: str) -> Callable[[Callable], Callable]:
    """Mark test xfail when CUTLASS CuTe backend is selected."""
    return xfailIfFn(lambda: _get_backend() == "cute", reason)


def skipIfNotTriton(reason: str) -> Callable[[Callable], Callable]:
    """Skip test when backend is not Triton (e.g. cute, pallas)."""
    return skipIfFn(lambda: _get_backend() != "triton", reason)


def onlyBackends(
    backends: Sequence[str],
) -> Callable[[type[unittest.TestCase]], type[unittest.TestCase]]:
    """Skip an entire test class unless `_get_backend() in backends`"""

    def wrapper(cls: type[unittest.TestCase]) -> type[unittest.TestCase]:
        backend = _get_backend()
        if backend in backends or (backend == "tileir" and "triton" in backends):
            return cls
        return unittest.skip(f"disabled for HELION_BACKEND={backend}")(cls)

    return wrapper


def skipUnlessTensorDescriptor(reason: str) -> Callable[[Callable], Callable]:
    """Skip test unless tensor descriptors are supported."""
    # Defers check to test execution time to avoid CUDA init during pytest-xdist collection.
    return skipIfFn(lambda: not supports_tensor_descriptor(), reason)


def skipUnlessTf32Supported(
    reason: str = "TF32 not supported on this GPU",
) -> Callable[[Callable], Callable]:
    """Skip test unless TF32 precision is supported (NVIDIA or AMD CDNA3 gfx942)."""
    from helion._compat import is_hip
    from helion._compat import supports_tf32_precision_on_amd

    # TF32 is supported on NVIDIA or on AMD GPUs that support it (gfx908-gfx942)
    tf32_supported = not is_hip() or supports_tf32_precision_on_amd()
    return unittest.skipUnless(tf32_supported, reason)


def get_test_dot_precision() -> str:
    """Get the appropriate dot precision for tests based on platform support.

    Returns 'tf32' if supported (NVIDIA or AMD gfx908-gfx942), otherwise 'ieee'.
    """
    from helion._compat import is_hip
    from helion._compat import supports_tf32_precision_on_amd

    if not is_hip():
        # NVIDIA - always supports tf32
        return "tf32"
    if supports_tf32_precision_on_amd():
        # AMD CDNA with TF32 support (gfx908-gfx942)
        return "tf32"
    # AMD without TF32 support (gfx950+)
    return "ieee"


def skipIfXPU(reason: str) -> Callable[[Callable], Callable]:
    """Skip test if running with Intel XPU"""
    return unittest.skipIf(torch.xpu.is_available(), reason)


def skipUnlessPallas(reason: str) -> Callable[[Callable], Callable]:
    """Skip test unless JAX Pallas TPU backend is available."""

    def _has_tpu_pallas() -> bool:
        try:
            from jax.experimental import pallas  # noqa: F401
            import torch_tpu.api  # type: ignore[import-not-found]

            torch_tpu.api.tpu_device()
            return True
        except Exception:
            return False

    return skipIfFn(lambda: not _has_tpu_pallas(), reason)


def skipIfA10G(reason: str) -> Callable[[Callable], Callable]:
    """Skip test if running on A10G GPU."""
    # Defers check to test execution time to avoid CUDA init during pytest-xdist collection.
    return skipIfFn(lambda: "A10G" in get_nvidia_gpu_model(), reason=reason)


def skipIfNotCUDA() -> Callable[[Callable], Callable]:
    """Skip test if not running on CUDA (NVIDIA GPU)."""
    # Defers check to test execution time to avoid CUDA init during pytest-xdist collection.
    return skipIfFn(
        lambda: not is_cuda(),
        reason="Test skipped: CUDA (NVIDIA GPU) is not available.",
    )


def skipIfCudaCapabilityLessThan(
    min_capability: tuple[int, int], *, reason: str | None = None
) -> Callable[[Callable], Callable]:
    """Skip test if not running on CUDA or capability is less than min_capability."""

    def cond() -> bool:
        if not is_cuda():
            return True
        return torch.cuda.get_device_capability() < min_capability

    # Defers check to test execution time to avoid CUDA init during pytest-xdist collection.
    return skipIfFn(
        cond,
        reason=reason
        or f"Requires CUDA capability >= {min_capability[0]}.{min_capability[1]}",
    )


def skipIfCudaSharedMemoryLessThan(
    min_shared_memory: int,
    *,
    reason: str | None = None,
) -> Callable[[Callable], Callable]:
    """Skip test if CUDA shared memory per block is below min_shared_memory."""

    def cond() -> bool:
        if not is_cuda():
            return False
        props = torch.cuda.get_device_properties(torch.cuda.current_device())
        default_shared = cast("int", props.shared_memory_per_block)
        optin_shared = cast(
            "int | None", getattr(props, "shared_memory_per_block_optin", None)
        )
        max_shared = default_shared if optin_shared is None else optin_shared
        return max_shared < min_shared_memory

    # Defers check to test execution time to avoid CUDA init during pytest-xdist collection.
    return skipIfFn(
        cond,
        reason=reason
        or f"Requires CUDA shared memory per block >= {min_shared_memory} bytes",
    )


def skipIfLowVRAM(
    reason: str = "Test requires high VRAM",
    *,
    required_bytes: int | None = None,
) -> Callable[[Callable], Callable]:
    """Skip test on systems with low GPU VRAM.

    When called with only a reason, returns a decorator that skips tests on GPUs
    with less than ~30 GiB total VRAM. When provided with required_bytes, uses
    that value as the threshold.
    """

    threshold_bytes = (
        int(30.0 * (1024**3)) if required_bytes is None else required_bytes
    )

    def is_low_vram() -> bool:
        try:
            if torch.cuda.is_available():
                device = torch.cuda.current_device()
                props = torch.cuda.get_device_properties(device)
                total_memory = int(getattr(props, "total_memory", 0))
                if total_memory < threshold_bytes:
                    return True
                # Also check free memory for shared GPU environments
                free_memory, _ = torch.cuda.mem_get_info(device)
                if free_memory < threshold_bytes:
                    return True
        except Exception:
            pass
        return False

    # Defers check to test execution time to avoid CUDA init during pytest-xdist collection.
    return skipIfFn(is_low_vram, reason=reason)


def skipIfPyTorchBaseVerLessThan(min_version: str) -> Callable[[Callable], Callable]:
    """Skip test if PyTorch base version is less than the specified version.

    Uses the base version for comparison, which ignores pre-release/dev/post suffixes.
    This allows development versions like "2.10.0.dev20251104" to pass when checking >= "2.10".

    Args:
        min_version: Minimum required PyTorch version (e.g., "2.10")

    Returns:
        Decorator that skips the test if PyTorch base version is below min_version
    """
    return unittest.skipIf(
        not requires_torch_version(min_version),
        f"PyTorch version {min_version} or higher required",
    )


@contextlib.contextmanager
def track_run_ref_calls() -> Generator[list[int], None, None]:
    """Context manager that tracks BoundKernel.run_ref calls.

    Yields:
        A list that will contain the count of run_ref calls.
    """
    from .runtime.kernel import BoundKernel

    original_run_ref = BoundKernel.run_ref
    run_ref_count = [0]

    def tracked_run_ref(self: BoundKernel, *args: object) -> object:
        run_ref_count[0] += 1
        return original_run_ref(self, *args)

    # pyrefly: ignore [bad-assignment]
    BoundKernel.run_ref = tracked_run_ref

    try:
        yield run_ref_count
    finally:
        BoundKernel.run_ref = original_run_ref


@contextlib.contextmanager
def assert_helion_ref_mode(
    ref_mode: RefMode = RefMode.OFF,
) -> Generator[None, None, None]:
    """Context manager that asserts Helion compilation behavior based on RefMode.

    - RefMode.OFF: expects compilation (run_ref should not be called)
    - RefMode.EAGER: expects no compilation (run_ref should be called)
    """
    with track_run_ref_calls() as run_ref_count:
        yield

        if ref_mode == RefMode.OFF:
            # In normal mode (RefMode.OFF), run_ref should not be called
            assert run_ref_count[0] == 0, (
                f"Expected run_ref to not be called in normal mode (RefMode.OFF), but got: run_ref={run_ref_count[0]}"
            )
        elif ref_mode == RefMode.EAGER:
            # In ref eager mode (RefMode.EAGER), run_ref should be called
            assert run_ref_count[0] > 0, (
                f"Expected run_ref to be called in ref eager mode (RefMode.EAGER), but got: run_ref={run_ref_count[0]}"
            )
        else:
            raise ValueError(f"Unknown RefMode: {ref_mode}")


assert_helion_compilation = functools.partial(
    assert_helion_ref_mode, ref_mode=RefMode.OFF
)

assert_ref_eager_mode = functools.partial(
    assert_helion_ref_mode, ref_mode=RefMode.EAGER
)


class RefEagerTestBase:
    """Base class for all ref eager mode test shards of normal Helion unit test files."""

    # Class-level tracking for assert_close counting
    _assert_close_count = 0
    _original_assert_close_func = None
    # Class-level tracking for assertRaises counting
    _assert_raises_count = 0
    _original_assert_raises_func = None
    # Class-level tracking for skipTest counting
    _skip_test_count = 0
    _original_skip_test_func = None
    # Class-level tracking for pytest.raises patching
    _original_pytest_raises = None
    # Class-level tracking for assertTrue/assertFalse/assertGreater
    _assert_true_count = 0
    _original_assert_true_func = None
    _assert_false_count = 0
    _original_assert_false_func = None
    _assert_greater_count = 0
    _original_assert_greater_func = None

    def setUp(self) -> None:
        """Common setup for all ref eager tests."""
        super().setUp()  # type: ignore[misc]

        # Check if HELION_INTERPRET is already set
        self._in_ref_eager_mode = os.environ.get("HELION_INTERPRET") == "1"

        # If not in ref eager mode, skip the setup
        if not self._in_ref_eager_mode:
            return

        # Reset assert_close counter for this test
        RefEagerTestBase._assert_close_count = 0
        # Reset assertRaises counter for this test
        RefEagerTestBase._assert_raises_count = 0
        # Reset skipTest counter for this test
        RefEagerTestBase._skip_test_count = 0
        # Reset assertTrue/assertFalse/assertGreater counters
        RefEagerTestBase._assert_true_count = 0
        RefEagerTestBase._assert_false_count = 0
        RefEagerTestBase._assert_greater_count = 0

        # Patch torch.testing.assert_close to count calls
        if RefEagerTestBase._original_assert_close_func is None:
            # pyrefly: ignore [bad-assignment]
            RefEagerTestBase._original_assert_close_func = torch.testing.assert_close

        def counting_assert_close(*args: object, **kwargs: object) -> None:
            RefEagerTestBase._assert_close_count += 1
            return RefEagerTestBase._original_assert_close_func(*args, **kwargs)  # type: ignore[misc]

        torch.testing.assert_close = counting_assert_close

        # Patch self.assertRaises to count calls
        if RefEagerTestBase._original_assert_raises_func is None:
            # pyrefly: ignore [bad-assignment]
            RefEagerTestBase._original_assert_raises_func = self.assertRaises

        def counting_assert_raises(*args: object, **kwargs: object) -> object:
            RefEagerTestBase._assert_raises_count += 1
            return RefEagerTestBase._original_assert_raises_func(*args, **kwargs)  # type: ignore[misc]

        self.assertRaises = counting_assert_raises

        # Patch self.skipTest to count calls
        if RefEagerTestBase._original_skip_test_func is None:
            # pyrefly: ignore [bad-assignment]
            RefEagerTestBase._original_skip_test_func = self.skipTest

        def counting_skip_test(*args: object, **kwargs: object) -> object:
            RefEagerTestBase._skip_test_count += 1
            return RefEagerTestBase._original_skip_test_func(*args, **kwargs)  # type: ignore[misc]

        self.skipTest = counting_skip_test

        # Store the tracking context manager instance so we can check counts in tearDown
        self._run_ref_tracker = track_run_ref_calls()
        self._run_ref_count = self._run_ref_tracker.__enter__()

        # Patch pytest.raises to count calls
        if RefEagerTestBase._original_pytest_raises is None:
            # pyrefly: ignore [bad-assignment]
            RefEagerTestBase._original_pytest_raises = pytest.raises

        def counting_pytest_raises(*args: object, **kwargs: object) -> object:
            """Wrapper for pytest.raises that counts calls but still runs the original logic."""
            RefEagerTestBase._assert_raises_count += 1
            assert RefEagerTestBase._original_pytest_raises is not None
            return RefEagerTestBase._original_pytest_raises(*args, **kwargs)

        pytest.raises = counting_pytest_raises  # type: ignore[assignment]

        # Patch self.assertTrue to count calls
        if RefEagerTestBase._original_assert_true_func is None:
            # pyrefly: ignore [bad-assignment]
            RefEagerTestBase._original_assert_true_func = self.assertTrue

        def counting_assert_true(*args: object, **kwargs: object) -> None:
            RefEagerTestBase._assert_true_count += 1
            return RefEagerTestBase._original_assert_true_func(*args, **kwargs)  # type: ignore[misc]

        self.assertTrue = counting_assert_true  # type: ignore[assignment]

        # Patch self.assertFalse to count calls
        if RefEagerTestBase._original_assert_false_func is None:
            # pyrefly: ignore [bad-assignment]
            RefEagerTestBase._original_assert_false_func = self.assertFalse

        def counting_assert_false(*args: object, **kwargs: object) -> None:
            RefEagerTestBase._assert_false_count += 1
            return RefEagerTestBase._original_assert_false_func(*args, **kwargs)  # type: ignore[misc]

        self.assertFalse = counting_assert_false  # type: ignore[assignment]

        # Patch self.assertGreater to count calls
        if RefEagerTestBase._original_assert_greater_func is None:
            # pyrefly: ignore [bad-assignment]
            RefEagerTestBase._original_assert_greater_func = self.assertGreater

        def counting_assert_greater(*args: object, **kwargs: object) -> None:
            RefEagerTestBase._assert_greater_count += 1
            return RefEagerTestBase._original_assert_greater_func(*args, **kwargs)  # type: ignore[misc]

        self.assertGreater = counting_assert_greater  # type: ignore[assignment]

    def tearDown(self) -> None:
        """Common teardown with assertion counting check."""
        # If not in ref eager mode, skip the teardown logic
        if not self._in_ref_eager_mode:
            super().tearDown()  # type: ignore[misc]
            return

        try:
            # Exit the run_ref tracker
            self._run_ref_tracker.__exit__(None, None, None)

            # Check if the test was skipped
            test_method = getattr(self, self._testMethodName, None)  # type: ignore[attr-defined]
            is_skipped = (
                test_method is not None
                and hasattr(test_method, "__unittest_skip__")
                and test_method.__unittest_skip__
            ) or RefEagerTestBase._skip_test_count > 0

            # Assert that either run_ref was called or the test was skipped
            if not is_skipped and self._run_ref_count[0] == 0:
                self.fail(  # type: ignore[attr-defined]
                    # pyrefly: ignore [missing-attribute]
                    f"Test {self._testMethodName} did not call run_ref and was not skipped"
                )

            if not is_skipped:
                # Check that either assert_close, assertRaises, skipTest, assertTrue, assertFalse, or assertGreater was called
                total_assertions = (
                    RefEagerTestBase._assert_close_count
                    + RefEagerTestBase._assert_raises_count
                    + RefEagerTestBase._skip_test_count
                    + RefEagerTestBase._assert_true_count
                    + RefEagerTestBase._assert_false_count
                    + RefEagerTestBase._assert_greater_count
                )
                # Need to use the original assertGreater to avoid recursion
                if RefEagerTestBase._original_assert_greater_func is not None:
                    RefEagerTestBase._original_assert_greater_func(  # type: ignore[misc]
                        total_assertions,
                        0,
                        f"Test {self._testMethodName} did not call torch.testing.assert_close, assertRaises, skipTest, assertTrue, assertFalse, or assertGreater",  # type: ignore[attr-defined]
                    )
                else:
                    # Fallback if original not available
                    assert total_assertions > 0, (
                        f"Test {self._testMethodName} did not call any assertion methods"  # type: ignore[attr-defined]
                    )
        finally:
            # Restore the original assert_close function
            if RefEagerTestBase._original_assert_close_func is not None:
                torch.testing.assert_close = (
                    RefEagerTestBase._original_assert_close_func
                )

            # Restore the original assertRaises function
            if RefEagerTestBase._original_assert_raises_func is not None:
                self.assertRaises = RefEagerTestBase._original_assert_raises_func

            # Restore the original skipTest function
            if RefEagerTestBase._original_skip_test_func is not None:
                self.skipTest = RefEagerTestBase._original_skip_test_func

            # Restore the original pytest.raises function
            if RefEagerTestBase._original_pytest_raises is not None:
                pytest.raises = RefEagerTestBase._original_pytest_raises

            # Restore the original assertTrue function
            if RefEagerTestBase._original_assert_true_func is not None:
                self.assertTrue = RefEagerTestBase._original_assert_true_func

            # Restore the original assertFalse function
            if RefEagerTestBase._original_assert_false_func is not None:
                self.assertFalse = RefEagerTestBase._original_assert_false_func

            # Restore the original assertGreater function
            if RefEagerTestBase._original_assert_greater_func is not None:
                self.assertGreater = RefEagerTestBase._original_assert_greater_func

            super().tearDown()  # type: ignore[misc]

    # NOTE: We no-op these methods because they commonly check behaviors that are not relevant in ref eager mode.
    # Instead, we solely rely on the unit test's `torch.testing.assert_close` and `assertRaises` checks to ensure ref eager mode's correctness.
    def assertExpectedJournal(self, value: str) -> None:
        if not self._in_ref_eager_mode:
            super().assertExpectedJournal(value)  # type: ignore[misc]

    def assertIn(
        self, member: object, container: object, msg: str | None = None
    ) -> None:
        if not self._in_ref_eager_mode:
            super().assertIn(member, container, msg)  # type: ignore[misc]

    def assertNotIn(
        self, member: object, container: object, msg: str | None = None
    ) -> None:
        if not self._in_ref_eager_mode:
            super().assertNotIn(member, container, msg)  # type: ignore[misc]

    def assertIs(self, expr1: object, expr2: object, msg: str | None = None) -> None:
        if not self._in_ref_eager_mode:
            super().assertIs(expr1, expr2, msg)  # type: ignore[misc]

    def assertIsNot(self, expr1: object, expr2: object, msg: str | None = None) -> None:
        if not self._in_ref_eager_mode:
            super().assertIsNot(expr1, expr2, msg)  # type: ignore[misc]

    def assertTrueIfInNormalMode(self, condition: bool, msg: str | None = None) -> None:
        if not self._in_ref_eager_mode:
            self.assertTrue(condition, msg)  # type: ignore[attr-defined]

    def assertEqualCode(self, first: str, second: str, msg: str | None = None) -> None:
        if not self._in_ref_eager_mode:
            super().assertEqual(first, second, msg)  # type: ignore[misc]

    def assertNotEqualCode(
        self, first: str, second: str, msg: str | None = None
    ) -> None:
        if not self._in_ref_eager_mode:
            super().assertNotEqual(first, second, msg)  # type: ignore[misc]

    def getUserDefinedTunable(
        self, user_defined_tunables: dict[str, object], key: str
    ) -> object | None:
        """Look up a specific value via key from user defined tunables. Returns None in ref mode."""
        if self._in_ref_eager_mode:
            return None
        return user_defined_tunables.get(key)

    def assertIsInstance(
        self, obj: object, cls: type | tuple[type, ...], msg: str | None = None
    ) -> None:
        if not self._in_ref_eager_mode:
            super().assertIsInstance(obj, cls, msg)  # type: ignore[misc]


def import_path(filename: Path) -> types.ModuleType:
    module_name = f"{__name__}.{filename.stem}"
    if module_name not in sys.modules:
        # pyrefly: ignore [implicit-import]
        spec = importlib.util.spec_from_file_location(module_name, filename)
        assert spec is not None
        # pyrefly: ignore [implicit-import]
        module = importlib.util.module_from_spec(spec)
        assert spec.loader is not None
        spec.loader.exec_module(module)
        sys.modules[module_name] = module
    return sys.modules[module_name]


def code_and_output(
    fn: Kernel,
    args: tuple[object, ...],
    **kwargs: object,
) -> tuple[str, object]:
    bound = fn.bind(args)
    if is_ref_mode_enabled(bound.kernel.settings):
        if kwargs:
            # pyrefly: ignore [bad-argument-type]
            config = Config(**kwargs)
            bound._config = config
        result = fn(*args)
        # Return the original kernel source code
        code = inspect.getsource(fn.fn)
        return code, result

    if kwargs:
        config = Config(
            # pyrefly: ignore [bad-argument-type]
            **kwargs
        )
    elif fn.configs:
        (config,) = fn.configs
    else:
        config = fn.bind(args).config_spec.default_config()
    # Strip config keys not supported by the current backend so that
    # tests with Triton-specific keys (num_warps, num_stages, indexing, etc.)
    # can run on other backends like Pallas/TPU.
    config_spec = fn.bind(args).config_spec
    for key in config_spec.unsupported_config_keys(config.config):
        config.config.pop(key, None)
    code = fn.bind(args).to_triton_code(config)
    compiled_kernel = fn.bind(args).compile_config(config)
    try:
        result = compiled_kernel(*args)
    except Exception:
        sys.stderr.write(f"Failed to run kernel:\n{code}\n")
        raise
    return code, result


def run_example(
    kernel_fn: Callable[..., torch.Tensor] | Kernel | dict[str, Kernel],
    baseline_fn: Callable[..., torch.Tensor] | dict[str, Callable[..., torch.Tensor]],
    args: tuple[object, ...],
    kernel_name: str = "helion",
    baseline_name: str = "torch",
    rtol: float = 1e-2,
    atol: float = 1e-1,
    bwd: bool = False,
) -> None:
    """Run complete example: correctness check + benchmark.

    Args:
        kernel_fn: Single kernel function, or dict of {name: function} for multiple kernel variants
        baseline_fn: Single baseline function or dict of {name: function} for multiple baselines
        args: Arguments to pass to all functions
        kernel_name: Name for single kernel in output (default: "helion")
        baseline_name: Name for single baseline in output (default: "torch")
        rtol: Relative tolerance for correctness check (default: 1e-2)
        atol: Absolute tolerance for correctness check (default: 1e-1)
        bwd: Whether to also test backward pass (default: False)
    """
    try:
        torch.backends.cuda.matmul.fp32_precision = "tf32"
        torch.backends.cudnn.conv.fp32_precision = "tf32"  # type: ignore[reportAttributeAccessIssue]
    except AttributeError:  # No cudnn available
        torch.set_float32_matmul_precision("high")  # older deprecated API

    # Normalize to dict format
    kernels = kernel_fn if isinstance(kernel_fn, dict) else {kernel_name: kernel_fn}
    baselines = (
        baseline_fn if isinstance(baseline_fn, dict) else {baseline_name: baseline_fn}
    )

    # Check correctness against first baseline
    first_baseline_name, first_baseline_func = next(iter(baselines.items()))
    expected = first_baseline_func(*args).clone()

    for name, func in {**kernels, **baselines}.items():
        if name != first_baseline_name:
            print(f"Testing {name} correctness...", file=sys.stderr)
            # Clone args to avoid buffer donation issues (e.g., Pallas/TPU)
            cloned_args = tuple(
                a.clone() if isinstance(a, torch.Tensor) else a for a in args
            )
            result = func(*cloned_args).clone()
            torch.testing.assert_close(
                result.to(torch.float32),
                expected.to(torch.float32),
                rtol=rtol,
                atol=atol,
            )

    # Test backward pass
    if bwd:
        # Find tensors that require gradients in args
        grad_tensors = []
        for arg in args:
            if isinstance(arg, torch.Tensor) and arg.requires_grad:
                grad_tensors.append(arg)

        # Ensure we have tensors to check
        assert len(grad_tensors) > 0, (
            "BWD: No tensors with requires_grad=True found. "
            "Check that input tensors have requires_grad set."
        )

        # Run baseline backward pass
        baseline_out = first_baseline_func(*args)
        grad_output = torch.randn_like(baseline_out)

        # Save original gradients
        baseline_out.backward(grad_output, retain_graph=True)
        baseline_grads = [
            t.grad.clone() if t.grad is not None else None for t in grad_tensors
        ]

        # Ensure at least one gradient was computed
        has_gradient = any(g is not None for g in baseline_grads)
        assert has_gradient, (
            "BWD: No gradients were computed. All tensors have grad=None. "
            "Check that backward was called and tensors require gradients."
        )

        # Clear gradients
        for t in grad_tensors:
            t.grad = None

        # Test each implementation
        for name, func in {**kernels, **baselines}.items():
            if name != first_baseline_name:
                # Run backward
                out = func(*args)
                out.backward(grad_output, retain_graph=True)

                # Collect implementation gradients
                impl_grads = [t.grad for t in grad_tensors]

                # Ensure same number of grad tensors
                assert len(impl_grads) == len(baseline_grads), (
                    f"BWD: Mismatch in number of grad tensors for {name}: "
                    f"{len(impl_grads)} vs {len(baseline_grads)}"
                )

                # Compare each tensor's gradient
                for i, (tensor, baseline_grad) in enumerate(
                    zip(grad_tensors, baseline_grads, strict=False)
                ):
                    # Check gradient existence
                    assert (baseline_grad is None) == (tensor.grad is None), (
                        f"BWD: Gradient existence mismatch for tensor {i} in {name}: "
                        f"baseline has grad={baseline_grad is not None}, "
                        f"impl has grad={tensor.grad is not None}"
                    )

                    if baseline_grad is not None:
                        torch.testing.assert_close(
                            tensor.grad.to(torch.float32),
                            baseline_grad.to(torch.float32),
                            rtol=rtol,
                            atol=atol,
                            msg=f"BWD: Gradient mismatch for tensor {i} with shape {tensor.shape} in {name}",
                        )

                # Clear gradients for next test
                for t in grad_tensors:
                    t.grad = None

    # Benchmark all functions — clone args to avoid buffer donation issues
    cloned_args = tuple(a.clone() if isinstance(a, torch.Tensor) else a for a in args)
    all_benchmarks = {**kernels, **baselines}
    bench_fns = [functools.partial(fn, *cloned_args) for fn in all_benchmarks.values()]
    repeat = compute_repeat(bench_fns[0])

    # For distributed workload, different rank may have slightly different
    # benchmarking result causing diverging `repeat` value.
    # Running different number of times on different ranks may cause
    # stuck processes.
    if dist.is_initialized():
        repeat = sync_object(repeat)

    # pyrefly: ignore [bad-argument-type]
    timings = interleaved_bench(bench_fns, repeat=repeat, desc="Benchmarking")
    all_times = dict(zip(all_benchmarks.keys(), timings, strict=True))
    best_baseline_time = min(all_times[name] for name in baselines)

    if not dist.is_initialized() or dist.get_rank() == 0:
        # Print results (on rank 0)
        print(f"\n{'=' * 65}\nBenchmark Results\n{'=' * 65}", file=sys.stderr)
        print(
            f"{'Implementation':<20} {'Time (ms)':<12} {'Speedup':<15}\n{'-' * 65}",
            file=sys.stderr,
        )

        for name, time in all_times.items():
            is_best_baseline = name in baselines and time == best_baseline_time
            speedup_str = (
                "1.00x (ref)"
                if is_best_baseline
                else f"{best_baseline_time / time:.2f}x"
            )
            print(f"{name:<20} {time:<12.4f} {speedup_str:<15}", file=sys.stderr)

        print(f"{'=' * 65}\n", file=sys.stderr)


def check_example(
    name: str,
    args: tuple[torch.Tensor, ...],
    expected: object,
    fn_name: str | None = None,
    skip_accuracy: bool = False,
    static_shapes: bool | None = None,
    atol: float = 1e-1,
    rtol: float = 1e-2,
    **kwargs: object,
) -> str:
    """Helper used in unit tests to run a single example kernel and check its output."""
    kernel_fn = getattr(import_path(EXAMPLES_DIR / f"{name}.py"), fn_name or name)
    if static_shapes is not None:
        assert static_shapes in (True, False)
        kernel_fn.settings.static_shapes = static_shapes

    code, result = code_and_output(
        kernel_fn,
        args,
        **kwargs,
    )

    if not skip_accuracy:
        # Use tree_map to apply assert_close to all tensor pairs
        def assert_close_fn(got: object, exp: object) -> None:
            # Skip if expected is None (i.e. we don't care what the actual value is)
            if exp is None:
                return
            # Both None is OK
            if got is None and exp is None:
                return
            assert isinstance(got, torch.Tensor) and isinstance(exp, torch.Tensor), (
                f"Type mismatch: got {type(got)}, expected {type(exp)}"
            )
            torch.testing.assert_close(
                got.to(torch.float32),
                exp.to(torch.float32),
                atol=atol,
                rtol=rtol,
            )

        tree_map(assert_close_fn, result, expected)
    return code


class AssertExpectedJournal:
    """
    Manages a <testfile>.expected file that contains expected output for TestCase.assertExpectedJournal() calls.

    This replaces the previous `expecttest` assertExpectedInline approach by storing expected output
    in external .expected files rather than inline strings in test files. This provides better
    organization and avoids cluttering test files with large code blocks.

    The .expected file format uses sections like:
    --- assertExpectedJournal(TestClass.test_method)
    expected output here

    --- assertExpectedJournal(TestClass.test_method)
    second expected output for same test

    Environment variable EXPECTTEST_ACCEPT=1 can be used to update expected outputs.
    """

    def expected_filename(self, basename: Path) -> Path:
        backend = _get_backend()
        if backend == "triton":
            return basename
        return Path(f"{basename}_{backend}")

    def __init__(self, cls: type[TestCase]) -> None:
        pyfile = os.path.abspath(inspect.getfile(cls))
        assert "/test/" in pyfile
        assert pyfile.endswith(".py")
        self._base_filename = basename = Path(pyfile[:-3] + ".expected")
        self.filename: Path = self.expected_filename(basename).resolve()
        self._cache: dict[str, list[str]] | None = None
        self._current_id: str | None = None
        self._current_index: int = 0

    @property
    def cache(self) -> dict[str, list[str]]:
        if self._cache is None:
            return self.reload()
        return self._cache

    def reload(self) -> dict[str, list[str]]:
        if self.filename.exists():
            data = self.filename.read_text()
        elif self.filename != self._base_filename and self._base_filename.exists():
            # use default expected file if specific one doesn't exist
            data = self._base_filename.read_text()
        else:
            data = ""
        result = collections.defaultdict(list)
        for name, expected in re.findall(
            r"--- assertExpectedJournal\(([^)]*)\)\n(.*?)(?=^--- assertExpectedJournal\(|\Z)",
            data,
            re.MULTILINE | re.DOTALL,
        ):
            result[name].append(expected.strip())
        self._cache = result
        return result

    def save(self) -> None:
        tmp = f"{self.filename}.tmp{os.getpid()}"
        with open(tmp, "w") as f:
            f.write(
                f"This file is automatically generated by assertExpectedJournal calls in {self.filename.stem}.py.\n"
                "Update expected outputs by running tests with the EXPECTTEST_ACCEPT=1 environment variable set.\n\n"
            )
            for name, expected_values in sorted(
                self.cache.items(), key=operator.itemgetter(0)
            ):
                f.writelines(
                    f"--- assertExpectedJournal({name})\n{expected}\n\n"
                    for expected in expected_values
                )
            # Remove the last newline to play nicer with some people's editors
            f.truncate(f.tell() - 1)
        os.rename(tmp, self.filename)

    @staticmethod
    def normalize_id(test_id: str) -> str:
        match = re.search(r"\b([^.]+\.[^.]+)$", test_id)
        assert match, f"Test ID '{test_id}' does not match expected format"
        return match.group(1)

    @staticmethod
    def normalize_tensor_descriptors(code: str) -> str:
        return code.replace(
            get_tensor_descriptor_fn_name(), "tl.make_tensor_descriptor"
        )

    @staticmethod
    def normalize_device_name(code: str) -> str:
        """
        convert device='cuda:0' or device(type='cuda', index=0) etc to device=DEVICE
        """
        # device='cuda:0'
        reg_pattern_for_device_str = r"device\s*=\s*['\"][^'\"]+['\"]"
        normalized_code = re.sub(reg_pattern_for_device_str, "device=DEVICE", code)
        # device(type='cuda', index=0)
        reg_pattern_for_torch_device = (
            r"device\s*\(type\s*=\s*['\"][^'\"]+['\"][^'\"\)]*\)"
        )
        return re.sub(reg_pattern_for_torch_device, "device=DEVICE", normalized_code)

    @staticmethod
    def normalize_codegen_variants(code: str) -> str:
        # TODO(oulgen): Remove when PyTorch 2.10 becomes stable

        # Remove libdevice import line if present
        code = re.sub(
            r"^\s*from torch\._inductor\.runtime\.triton_compat import libdevice\s*\n?",
            "",
            code,
            flags=re.MULTILINE,
        )

        # Normalize sqrt variants
        # libdevice.sqrt( -> tl.sqrt_rn(
        code = re.sub(r"\blibdevice\.sqrt\s*\(", "tl.sqrt_rn(", code)
        # tl.sqrt( -> tl.sqrt_rn(
        code = re.sub(r"\btl\.sqrt\s*\(", "tl.sqrt_rn(", code)

        # Normalize rsqrt variants
        # libdevice.rsqrt( -> tl.rsqrt(
        code = re.sub(r"\blibdevice\.rsqrt\s*\(", "tl.rsqrt(", code)

        total_num_triton_helpers_replacements = 0

        # Normalize maximum variants
        # tl.maximum(a, b, tl.PropagateNan.ALL) -> triton_helpers.maximum(a, b)
        code, num_replacements = re.subn(
            r"\btl\.maximum\s*\(([^,]+),\s*([^,]+),\s*tl\.PropagateNan\.ALL\s*\)",
            r"triton_helpers.maximum(\1, \2)",
            code,
        )
        total_num_triton_helpers_replacements += num_replacements

        # Normalize minimum variants
        # tl.minimum(a, b, tl.PropagateNan.ALL) -> triton_helpers.minimum(a, b)
        code, num_replacements = re.subn(
            r"\btl\.minimum\s*\(([^,]+),\s*([^,]+),\s*tl\.PropagateNan\.ALL\s*\)",
            r"triton_helpers.minimum(\1, \2)",
            code,
        )
        total_num_triton_helpers_replacements += num_replacements

        # Normalize tl.full scalar constants
        # tl.full([], VALUE, tl.float32) -> VALUE
        # tl.full([], VALUE, tl.float64) -> VALUE
        # tl.full([], VALUE, tl.int32) -> VALUE
        # etc.
        code = re.sub(
            r"\btl\.full\s*\(\s*\[\s*\]\s*,\s*([^,]+)\s*,\s*tl\.\w+\s*\)",
            r"\1",
            code,
        )

        triton_helpers_import = "from torch._inductor.runtime import triton_helpers"
        if (
            total_num_triton_helpers_replacements > 0
            and triton_helpers_import not in code
        ):
            # Insert right after `import triton.language as tl`
            code = re.sub(
                r"(^import triton\.language as tl$)",
                rf"\1\n{triton_helpers_import}",
                code,
                count=1,
                flags=re.MULTILINE,
            )

        return code

    @staticmethod
    def normalize_source_comment_structure(code: str) -> str:
        pattern = re.compile(
            r"^(?P<indent>\s*)# src\[(?P<prefix>[^:\]]+:)(?P<start>\d+|N)(?:-(?P<end>\d+|N))?]: (?P<text>.*?)(?P<newline>\r?\n|$)",
            flags=re.MULTILINE,
        )

        def replacer(match: re.Match[str]) -> str:
            text = match.group("text").rstrip()
            if not text.strip():
                return ""
            indent = match.group("indent")
            prefix = match.group("prefix")
            suffix = "N-N" if match.group("end") is not None else "N"
            newline = match.group("newline")
            return f"{indent}# src[{prefix}{suffix}]: {text}{newline}"

        # Normalize structured src comments
        code = pattern.sub(replacer, code)

        # Normalize file line refs: foo.py:123 -> foo.py:N
        return re.sub(r"(\b[^:\s]+\.py):(\d+)\b", r"\1:N", code)

    @classmethod
    def normalize_code(cls, code: str) -> str:
        code = cls.normalize_source_comment_structure(code)
        code = cls.normalize_tensor_descriptors(code)
        code = cls.normalize_device_name(code)
        code = cls.normalize_codegen_variants(code)
        return code.strip()

    def lookup(self, test_id: str, value: str) -> tuple[str, str]:
        test_id = self.normalize_id(test_id)
        if self._current_id != test_id:
            self._current_id = test_id
            self._current_index = 0

        expected_values = self.cache[test_id]
        if self._current_index < len(expected_values):
            expected = expected_values[self._current_index]
        else:
            assert self._current_index == len(expected_values)
            expected_values.append("")
            expected = ""

        # Normalize both actual and expected for robust comparisons
        value = self.normalize_code(value)
        expected = self.normalize_code(expected)

        if value != expected and os.environ.get("EXPECTTEST_ACCEPT", "0") not in {
            "0",
            "false",
            "False",
            "",
        }:
            expected_values[self._current_index] = value
            # Reload to play nicer with other processes
            self.reload()[test_id][:] = expected_values
            self.save()
            expected = value
            print(
                f"Expected output for {test_id} updated: {len(expected)} => {len(value)} bytes",
                file=sys.stderr,
            )
        self._current_index += 1
        return value, expected


class RefEagerTestDisabled:
    """Base class for test classes that should be skipped when ref eager mode is enabled."""

    def setUp(self) -> None:
        """Skip test if ref eager mode is enabled."""
        super().setUp()  # type: ignore[misc]
        if os.environ.get("HELION_INTERPRET") == "1":
            self.skipTest("Test class disabled in ref eager mode")  # type: ignore[attr-defined]


class TestCase(unittest.TestCase):
    maxDiff = 16384

    @classmethod
    def setUpClass(cls) -> None:
        cls._expected_journal = AssertExpectedJournal(cls)

        if is_mtia():
            # pyrefly: ignore [missing-import]
            import mtia.host_runtime.torch_mtia.dynamic_library  # noqa: F401

            # pyrefly: ignore [missing-import]
            from mtia.re.re_unittest_lib import MTIAUnittest

            # pyrefly: ignore [missing-import]
            from triton_mtia.python.mtia.eager import mtia_triton_launcher

            # Call MTIAUnittest.setUpClass for MTIA initialization
            MTIAUnittest.setUpClass.__func__(cls)
            # Initialize MTIA properly
            mtia_triton_launcher.init()
        super().setUpClass()

    @classmethod
    def tearDownClass(cls) -> None:
        super().tearDownClass()
        del cls._expected_journal

    def setUp(self) -> None:
        super().setUp()
        self._test_stack = contextlib.ExitStack()

        from torch._inductor.utils import fresh_cache

        self._test_stack.enter_context(fresh_cache())

        counters.clear()

    def tearDown(self) -> None:
        super().tearDown()
        self._test_stack.close()

    def assertExpectedJournal(self, value: str) -> None:
        """
        Assert that the given value matches the expected output stored in <testfile>.expected.

        This method replaces assertExpectedInline for code generation tests. Instead of storing
        expected output as inline strings in test files, it uses external .expected files for
        better organization.

        Args:
            value: The actual output to compare (usually generated Triton code)

        Raises:
            AssertionError: If value doesn't match expected output

        Note:
            Use EXPECTTEST_ACCEPT=1 environment variable to update expected outputs.
        """
        value = _strip_launcher_args(value)
        value, expected = self._expected_journal.lookup(self.id(), value)
        expected = _strip_launcher_args(expected)
        # Normalize input_precision for consistent test comparisons across GPUs
        value = re.sub(
            r"input_precision='(tf32|ieee)'", "input_precision='ieee'", value
        )
        expected = re.sub(
            r"input_precision='(tf32|ieee)'", "input_precision='ieee'", expected
        )
        self.assertMultiLineEqual(
            value,
            expected,
            msg="To accept the new output, re-run test with env EXPECTTEST_ACCEPT=1",
        )

    @contextlib.contextmanager
    def capture_logs(self) -> Generator[_LogCapture, None, None]:
        """Context manager to capture logs."""
        handler = _LogCapture()
        handler.setLevel(logging.DEBUG)
        logger = logging.getLogger()
        logger.addHandler(handler)
        try:
            yield handler
        finally:
            logger.removeHandler(handler)

    @contextlib.contextmanager
    def capture_output(self) -> Generator[_OutputCapture, None, None]:
        """Context manager to capture stdout/stderr."""
        capture = _OutputCapture()
        old_stdout, old_stderr = sys.stdout, sys.stderr
        sys.stdout, sys.stderr = capture.stdout, capture.stderr
        try:
            yield capture
        finally:
            sys.stdout, sys.stderr = old_stdout, old_stderr
