# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0

# ################################################################################
#
# This example demonstrates how to use cuda.core to compile a CUDA kernel
# and launch it using PyTorch tensors as inputs. Requires PyTorch with CUDA.
#
# ################################################################################

# /// script
# dependencies = ["cuda_bindings", "cuda_core", "torch"]
# ///

import sys

import torch

from cuda.core import Device, LaunchConfig, Program, ProgramOptions, launch

# SAXPY kernel - passing a as a pointer to avoid any type issues
code = """
template<typename T>
__global__ void saxpy_kernel(const T* a, const T* x, const T* y, T* out, size_t N) {
 const unsigned int tid = threadIdx.x + blockIdx.x * blockDim.x;
 if (tid < N) {
   // Dereference a to get the scalar value
   out[tid] = (*a) * x[tid] + y[tid];
 }
}
"""


def torch_version() -> tuple[int, int, int]:
    v = torch.__version__
    v = v.split("+")[0].split(".")
    return tuple(int(x) for x in v)


def main():
    dev = Device()
    dev.set_current()

    print(f"PyTorch version: {torch_version()}", file=sys.stderr)
    if torch_version() < (2, 10, 0):
        # Need support for the __cuda_stream__ method.
        print("PyTorch version 2.10 or later required, skipping example", file=sys.stderr)
        return

    pt_stream = torch.cuda.current_stream()
    print(f"PyTorch stream: {pt_stream}", file=sys.stderr)

    stream = dev.create_stream(pt_stream)

    try:
        # prepare program
        program_options = ProgramOptions(std="c++11", arch=f"sm_{dev.arch}")
        prog = Program(code, code_type="c++", options=program_options)
        mod = prog.compile(
            "cubin",
            logs=sys.stdout,
            name_expressions=("saxpy_kernel<float>", "saxpy_kernel<double>"),
        )

        # Run in single precision
        kernel = mod.get_kernel("saxpy_kernel<float>")
        dtype = torch.float32

        # prepare input/output
        size = 64
        # Use a single element tensor for 'a'
        a = torch.tensor([10.0], dtype=dtype, device="cuda")
        x = torch.rand(size, dtype=dtype, device="cuda")
        y = torch.rand(size, dtype=dtype, device="cuda")
        out = torch.empty_like(x)

        # prepare launch
        block = 32
        grid = int((size + block - 1) // block)
        config = LaunchConfig(grid=grid, block=block)
        kernel_args = (a.data_ptr(), x.data_ptr(), y.data_ptr(), out.data_ptr(), size)

        # launch kernel on our stream
        launch(stream, config, kernel, *kernel_args)

        # check result
        assert torch.allclose(out, a.item() * x + y)

        # let's repeat again with double precision
        kernel = mod.get_kernel("saxpy_kernel<double>")
        dtype = torch.float64

        # prepare input
        size = 128
        # Use a single element tensor for 'a'
        a = torch.tensor([42.0], dtype=dtype, device="cuda")
        x = torch.rand(size, dtype=dtype, device="cuda")
        y = torch.rand(size, dtype=dtype, device="cuda")

        # prepare output
        out = torch.empty_like(x)

        # prepare launch
        block = 64
        grid = int((size + block - 1) // block)
        config = LaunchConfig(grid=grid, block=block)
        kernel_args = (a.data_ptr(), x.data_ptr(), y.data_ptr(), out.data_ptr(), size)

        # launch kernel on PyTorch's stream
        launch(stream, config, kernel, *kernel_args)

        # check result
        assert torch.allclose(out, a * x + y)
    finally:
        stream.close()


if __name__ == "__main__":
    main()
