from std.gpu import block_dim, block_idx, grid_dim, thread_idx
from std.gpu.host import DeviceBuffer, DeviceContext


@__name("oom_repro_spin")
def _spin_kernel(
    ptr: UnsafePointer[Scalar[DType.float32], MutAnyOrigin],
    count: Int,
    rounds: Int,
):
    var i = Int(block_idx.x) * Int(block_dim.x) + Int(thread_idx.x)
    var stride = Int(grid_dim.x) * Int(block_dim.x)
    while i < count:
        var acc = ptr[i]
        for _ in range(rounds):
            acc = acc * 1.000000119 + 0.25
        ptr[i] = acc
        i += stride


def main() raises:
    comptime GiB = 1024 * 1024 * 1024
    comptime MiB = 1024 * 1024
    with DeviceContext() as ctx:
        # Training-heap shape: big transients interleaved with small
        # long-lived tensors.
        var keep = List[DeviceBuffer[DType.uint8]]()
        var drop = List[DeviceBuffer[DType.float32]]()
        var i = 0
        try:
            while True:
                drop.append(
                    ctx.enqueue_create_buffer[DType.float32](GiB // 4)
                )
                keep.append(ctx.enqueue_create_buffer[DType.uint8](MiB))
                i += 1
        except e:
            _ = e
            pass
        print("pairs before first OOM:", i)

        # Keep the stream busy with kernels READING the big buffers, then
        # free them: the stream-ordered frees stay pending behind the
        # kernels, exactly like async training where the host runs ahead.
        for j in range(len(drop)):
            ctx.enqueue_function[_spin_kernel](
                drop[j].unsafe_ptr().as_unsafe_any_origin(),
                GiB // 4,
                2000,
                grid_dim=(1024,),
                block_dim=(256,),
            )
        var freed = len(drop)
        drop.clear()  # frees enqueue behind the spin kernels — all pending

        # Request 8 GiB NOW, while every free is still in flight.
        try:
            var big = ctx.enqueue_create_buffer[DType.uint8](8 * GiB)
            ctx.synchronize()
            print("OK: 8 GiB allocation succeeded with frees pending")
            _ = big^
        except e2:
            print(
                "BUG: 8 GiB allocation failed;",
                freed,
                "x 1 GiB freed but pending behind in-flight kernels",
            )
            print("error:", e2)
        ctx.synchronize()
        _ = len(keep)