VariableBuilder(tx, AttrSource(...))._wrap(user_obj)

---

def __call__(self, value):
    if value in self.tx.output.side_effects:
        # returns existing tracked variable
        ...

---

#!/usr/bin/env python3
"""
Reproducer for PyTorch Dynamo bug with triton.autotune + prune_configs_by.

Tested with:
    - torch 2.10.0+rocm7.1
    - triton 3.6.0
    - AMD MI300X (ROCm)
    Should also reproduce on CUDA with equivalent torch/triton versions.
"""

import sys
import traceback

import torch
import triton
import triton.language as tl


# --- Version info ---

def print_versions():
    print(f"Python:  {sys.version.split()[0]}")
    print(f"PyTorch: {torch.__version__}")
    print(f"Triton:  {triton.__version__}")
    if torch.cuda.is_available():
        print(f"GPU:     {torch.cuda.get_device_name(0)}")
    else:
        print("GPU:     not available")
    print()


# --- Prune function (no-op, just passes configs through) ---

def noop_prune(configs, named_args, **kwargs):
    """A no-op prune function. The bug is triggered by the prune_configs_by
    mechanism itself, not by any particular pruning logic."""
    return configs


# --- Kernel definitions ---

# Kernel WITH prune_configs_by (triggers the bug on second call)
@triton.autotune(
    configs=[
        triton.Config({"BLOCK": 128}, num_warps=4),
        triton.Config({"BLOCK": 256}, num_warps=4),
    ],
    key=["N"],
    prune_configs_by={"early_config_prune": noop_prune},
)
@triton.jit
def add_kernel_with_prune(
    x_ptr, y_ptr, out_ptr, N,
    BLOCK: tl.constexpr,
):
    pid = tl.program_id(0)
    offsets = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offsets < N
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x + y, mask=mask)


# Kernel WITHOUT prune_configs_by (works fine, even called twice)
@triton.autotune(
    configs=[
        triton.Config({"BLOCK": 128}, num_warps=4),
        triton.Config({"BLOCK": 256}, num_warps=4),
    ],
    key=["N"],
)
@triton.jit
def add_kernel_no_prune(
    x_ptr, y_ptr, out_ptr, N,
    BLOCK: tl.constexpr,
):
    pid = tl.program_id(0)
    offsets = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offsets < N
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x + y, mask=mask)


# --- Wrapper functions ---

def call_kernel(kernel, x, y):
    """Call the given autotuned kernel."""
    out = torch.empty_like(x)
    N = x.numel()
    grid = lambda meta: (triton.cdiv(N, meta["BLOCK"]),)
    kernel[grid](x, y, out, N)
    return out


def fn_with_prune_called_once(x, y):
    """Calls the pruned kernel once. Should work."""
    return call_kernel(add_kernel_with_prune, x, y)


def fn_with_prune_called_twice(x, y):
    """Calls the pruned kernel twice. Triggers the bug."""
    a = call_kernel(add_kernel_with_prune, x, y)
    b = call_kernel(add_kernel_with_prune, a, y)
    return b


def fn_no_prune_called_twice(x, y):
    """Calls the non-pruned kernel twice. Should work."""
    a = call_kernel(add_kernel_no_prune, x, y)
    b = call_kernel(add_kernel_no_prune, a, y)
    return b


# --- Test runner ---

def run_test(name, fn, x, y, expect_fail=False):
    """Run a single test case under torch.compile."""
    compiled_fn = torch.compile(fn, fullgraph=True, backend="eager")
    print(f"  {name}: ", end="", flush=True)
    try:
        result = compiled_fn(x, y)
        # Verify correctness
        expected = fn(x, y)
        if torch.allclose(result, expected):
            if expect_fail:
                print("UNEXPECTED PASS (bug may be fixed!)")
            else:
                print("PASS")
        else:
            print(f"FAIL (wrong result, max diff={torch.max(torch.abs(result - expected)).item():.6f})")
    except AssertionError as e:
        if "already tracked for mutation" in str(e):
            if expect_fail:
                print(f"EXPECTED FAIL: {e}")
            else:
                print(f"UNEXPECTED FAIL: {e}")
        else:
            print(f"FAIL (unexpected assertion): {e}")
            traceback.print_exc()
    except Exception as e:
        print(f"FAIL (unexpected error): {type(e).__name__}: {e}")
        traceback.print_exc()


def main():
    print("=" * 72)
    print("Dynamo + triton.autotune + prune_configs_by bug reproducer")
    print("=" * 72)
    print()
    print_versions()

    if not torch.cuda.is_available():
        print("ERROR: No GPU available. This reproducer requires a GPU.")
        sys.exit(1)

    # Reset Dynamo state between tests to ensure isolation
    device = "cuda"
    N = 1024
    x = torch.randn(N, device=device)
    y = torch.randn(N, device=device)

    print("Test 1: autotune WITHOUT prune_configs_by, kernel called twice")
    print("  Expected: PASS (prune_configs_by is the trigger)")
    torch._dynamo.reset()
    run_test("no_prune_twice", fn_no_prune_called_twice, x, y, expect_fail=False)
    print()

    print("Test 2: autotune WITH prune_configs_by, kernel called ONCE")
    print("  Expected: PASS (bug only triggers on second call)")
    torch._dynamo.reset()
    run_test("prune_once", fn_with_prune_called_once, x, y, expect_fail=False)
    print()

    print("Test 3: autotune WITH prune_configs_by, kernel called TWICE")
    print("  Expected: FAIL with 'already tracked for mutation'")
    print("  This is the bug.")
    torch._dynamo.reset()
    run_test("prune_twice", fn_with_prune_called_twice, x, y, expect_fail=True)
    print()

if __name__ == "__main__":
    main()

---

ListVariable(length=2) is already tracked for mutation. This could be because you are not using VariableBuilder to construct the variable tracker. Source of new object: AttrSource(base=GlobalSource(global_name='add_kernel_with_prune'), member='configs'). Source of previously tracked object: AttrSource(base=GlobalSource(global_name='add_kernel_with_prune'), member='configs').

from user code:
   File "dynamo_prune_configs_bug.py", line 105, in fn_with_prune_called_twice
    b = call_kernel(add_kernel_with_prune, a, y)
  File "dynamo_prune_configs_bug.py", line 93, in call_kernel
    kernel[grid](x, y, out, N)

---

Collecting environment information...
PyTorch version: 2.10.0+rocm7.1
Is debug build: False
CUDA used to build PyTorch: N/A
ROCM used to build PyTorch: 7.1.25424

OS: Ubuntu 24.04.3 LTS (x86_64)
GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04) 13.3.0
Clang version: Could not collect
CMake version: version 3.28.3
Libc version: glibc-2.39

Python version: 3.12.3 (main, Jan  8 2026, 11:30:50) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-6.8.0-85-generic-x86_64-with-glibc2.39
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to:
GPU models and configuration: AMD Instinct MI300X VF (gfx942:sramecc+:xnack-)
Nvidia driver version: Could not collect
cuDNN version: Could not collect
Is XPU available: False
HIP runtime version: 7.1.25424
MIOpen runtime version: 3.5.1
Is XNNPACK available: True
Caching allocator config: N/A

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Address sizes:                        46 bits physical, 57 bits virtual
Byte Order:                           Little Endian
CPU(s):                               160
On-line CPU(s) list:                  0-159
Vendor ID:                            GenuineIntel
Model name:                           INTEL(R) XEON(R) PLATINUM 8568Y+
CPU family:                           6
Model:                                207
Thread(s) per core:                   1
Core(s) per socket:                   80
Socket(s):                            2
Stepping:                             2
BogoMIPS:                             4600.00
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon pebs bts rep_good nopl xtopology cpuid tsc_known_freq pni pclmulqdq dtes64 vmx ssse3 fma cx16 pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch cpuid_fault ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx_vnni avx512_bf16 wbnoinvd arat vnmi avx512vbmi umip pku ospke waitpkg avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid bus_lock_detect cldemote movdiri movdir64b fsrm md_clear serialize tsxldtrk avx512_fp16 arch_capabilities
Virtualization:                       VT-x
Hypervisor vendor:                    KVM
Virtualization type:                  full
L1d cache:                            5 MiB (160 instances)
L1i cache:                            5 MiB (160 instances)
L2 cache:                             640 MiB (160 instances)
NUMA node(s):                         2
NUMA node0 CPU(s):                    0-79
NUMA node1 CPU(s):                    80-159
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Unknown: No mitigations
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Not affected
Vulnerability Spec rstack overflow:   Not affected
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; Enhanced / Automatic IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI SW loop, KVM SW loop
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Mitigation; TSX disabled

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] pytorch-lightning==2.5.0
[pip3] torch==2.10.0+rocm7.1
[pip3] torchmetrics==1.8.2
[pip3] torchvision==0.25.0+rocm7.1
[pip3] triton==3.6.0
[conda] Could not collect