pytorch - 💡(How to fix) Fix torch.compile + TMA path : Illegal Memory Access

"""
Minimal reproducer: Inductor TMA codegen IMA at >= ~43M elements.

Environment:
  - PyTorch: built from source (main)
  - GPU: NVIDIA GB200
  - Inductor config: triton.use_tensor_descriptor=True, assume_aligned_inputs=True

Usage:
  python repro_inductor_tma_ima.py
  CUDA_LAUNCH_BLOCKING=1 python repro_inductor_tma_ima.py
"""

import torch
import torch._inductor.config as inductor_config


def test_compile_tma(M, N):
    """Test torch.compile + TMA at a given shape."""
    inductor_config.triton.use_tensor_descriptor = True
    inductor_config.assume_aligned_inputs = True

    torch._dynamo.reset()
    fn = torch.compile(
        lambda x, r: x + r,
        fullgraph=True,
    )

    x = torch.randn(M, N, dtype=torch.bfloat16, device="cuda")
    r = torch.randn(M, N, dtype=torch.bfloat16, device="cuda")
    n_elem = M * N

    try:
        for _ in range(3):
            out = fn(x, r)
        torch.cuda.synchronize()
        ref = (x.float() + r.float()).bfloat16()
        torch.testing.assert_close(out, ref, rtol=1e-2, atol=2e-2)
        print(f"  ({M:>5}, {N:>5}) = {n_elem:>12,} elements: OK")
        return True
    except (torch.AcceleratorError, RuntimeError) as e:
        print(f"  ({M:>5}, {N:>5}) = {n_elem:>12,} elements: FAIL (IMA)", e)
        return False


def main():

    for M, N in [
        (4096, 4096),   # 16M elements -- OK
        (4096, 8192),   # 32M elements -- OK
        (8192, 4096),   # 32M elements -- OK
        (4096, 12288),  # ~49M elements -- FAIL
        (8192, 8192),   # 64M elements -- FAIL
    ]:
        ok = test_compile_tma(M, N)
        if not ok:
            break

if __name__ == "__main__":
    main()