"""Run with Python 3.10+ and PyTorch 2.x: python compile_basics.py."""
import platform
import time
import torch
from torch.utils.benchmark import Timer


def transform(x):
    return torch.relu(x * 2.0 + 1.0)


def main():
    torch.manual_seed(0)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print({'torch': torch.__version__, 'python': platform.python_version(),
           'device': device, 'cuda': torch.version.cuda})
    if device == 'cuda':
        print('GPU:', torch.cuda.get_device_name())
    torch.set_num_threads(1)
    x = torch.randn(1024 * 1024, device=device)

    def sync():
        if device == 'cuda':
            torch.cuda.synchronize()

    # This is the actual optimizing compiler, not the inspection backend below.
    torch.compiler.reset()
    compiled = torch.compile(transform, backend='inductor', fullgraph=True)
    expected = transform(x)
    sync()
    start = time.perf_counter()
    actual = compiled(x)
    sync()
    print(f'First compiled invocation: {time.perf_counter() - start:.3f} s')
    print('Includes execution and any compilation/cache-loading work; disk caches may be warm.')
    torch.testing.assert_close(actual, expected)
    for _ in range(5):
        compiled(x)
    sync()
    for name, fn in [('eager', transform), ('compiled', compiled)]:
        measurement = Timer(stmt='fn(x)', globals={'fn': fn, 'x': x}).blocked_autorange(min_run_time=0.5)
        print(f'{name}: {measurement.median * 1e6:.2f} us median')

    # Inspect Dynamo capture separately. Returning gm.forward does NOT optimize kernels.
    torch.compiler.reset()
    captures = []
    def inspect_backend(gm, example_inputs):
        captures.append(gm.code)
        print(f'\nCaptured graph {len(captures)}:\n{gm.code}')
        return gm.forward

    traced = torch.compile(transform, backend=inspect_backend, fullgraph=True, dynamic=False)
    for size in [8, 8, 16, 8]:
        probe = torch.randn(size, device=device)
        torch.testing.assert_close(traced(probe), transform(probe))
        print(f'shape=({size},), captured graphs so far={len(captures)}')
    assert len(captures) == 2, f'Expected two static-shape captures, got {len(captures)}'
    print('Output checks and static-shape reuse checks passed.')


if __name__ == '__main__':
    main()