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

# ################################################################################
#
# This example demonstrates how to use cuda.core to compile and launch a
# simple vector addition kernel.
#
# ################################################################################

# /// script
# dependencies = ["cuda_bindings", "cuda_core", "nvidia-cuda-nvrtc", "cupy-cuda13x"]
# ///

import cupy as cp

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

# compute c = a + b
code = """
template<typename T>
__global__ void vector_add(const T* A,
                           const T* B,
                           T* C,
                           size_t N) {
    const unsigned int tid = threadIdx.x + blockIdx.x * blockDim.x;
    for (size_t i=tid; i<N; i+=gridDim.x*blockDim.x) {
        C[i] = A[i] + B[i];
    }
}
"""


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

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

        # run in single precision
        kernel = mod.get_kernel("vector_add<float>")
        dtype = cp.float32

        # prepare input/output
        size = 50000
        rng = cp.random.default_rng()
        a = rng.random(size, dtype=dtype)
        b = rng.random(size, dtype=dtype)
        c = cp.empty_like(a)

        # cupy runs on a different stream from stream, so sync before accessing
        dev.sync()

        # prepare launch
        block = 256
        grid = (size + block - 1) // block
        config = LaunchConfig(grid=grid, block=block)

        # launch kernel on stream
        launch(stream, config, kernel, a.data.ptr, b.data.ptr, c.data.ptr, cp.uint64(size))
        stream.sync()

        # check result
        assert cp.allclose(c, a + b)
    finally:
        stream.close()


if __name__ == "__main__":
    main()
