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

# ################################################################################
#
# This example demonstrates CUDA graphs to capture and execute multiple
# kernel launches with minimal overhead. The graph performs a sequence of
# vector operations: add, multiply, and subtract.
#
# ################################################################################

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

import sys
import time

import cupy as cp

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


def main():
    # CUDA kernels for vector operations
    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];
        }
    }

    template<typename T>
    __global__ void vector_multiply(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];
        }
    }

    template<typename T>
    __global__ void vector_subtract(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];
        }
    }
    """

    # Initialize device and stream
    dev = Device()
    dev.set_current()
    stream = dev.create_stream()
    graph_builder = None
    graph = None
    # tell CuPy to use our stream as the current stream:
    cp.cuda.ExternalStream(int(stream.handle)).use()

    try:
        # Compile the 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>", "vector_multiply<float>", "vector_subtract<float>")
        )

        # Get kernel functions
        add_kernel = mod.get_kernel("vector_add<float>")
        multiply_kernel = mod.get_kernel("vector_multiply<float>")
        subtract_kernel = mod.get_kernel("vector_subtract<float>")

        # Prepare data
        size = 1000000
        dtype = cp.float32

        # Create input arrays
        rng = cp.random.default_rng(42)  # Fixed seed for reproducibility
        a = rng.random(size, dtype=dtype)
        b = rng.random(size, dtype=dtype)
        c = rng.random(size, dtype=dtype)

        # Create output arrays
        result1 = cp.empty_like(a)
        result2 = cp.empty_like(a)
        result3 = cp.empty_like(a)

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

        # Sync before graph capture
        dev.sync()

        print("Building CUDA graph...", file=sys.stderr)

        # Build the graph
        graph_builder = stream.create_graph_builder()
        graph_builder.begin_building()

        # Add multiple kernel launches to the graph
        # Kernel 1: result1 = a + b
        launch(graph_builder, config, add_kernel, a.data.ptr, b.data.ptr, result1.data.ptr, cp.uint64(size))

        # Kernel 2: result2 = result1 * c
        launch(graph_builder, config, multiply_kernel, result1.data.ptr, c.data.ptr, result2.data.ptr, cp.uint64(size))

        # Kernel 3: result3 = result2 - a
        launch(graph_builder, config, subtract_kernel, result2.data.ptr, a.data.ptr, result3.data.ptr, cp.uint64(size))

        # Complete the graph
        graph = graph_builder.end_building().complete()

        # Upload the graph to the stream
        graph.upload(stream)

        # Execute the entire graph with a single launch
        print("Executing graph...", file=sys.stderr)
        start_time = time.time()
        graph.launch(stream)
        stream.sync()
        end_time = time.time()

        graph_execution_time = end_time - start_time
        if graph_execution_time == 0.0:
            print("Graph execution time is too fast to measure accurately.")
            graph_execution_time = 1e-9  # Assign a small value to avoid division by zero in speedup calculation
        print(f"Graph execution time: {graph_execution_time:.6f} seconds")

        # Verify results
        expected_result1 = a + b
        expected_result2 = expected_result1 * c
        expected_result3 = expected_result2 - a

        assert cp.allclose(result1, expected_result1, rtol=1e-5, atol=1e-5), "Result 1 mismatch"
        assert cp.allclose(result2, expected_result2, rtol=1e-5, atol=1e-5), "Result 2 mismatch"
        assert cp.allclose(result3, expected_result3, rtol=1e-5, atol=1e-5), "Result 3 mismatch"

        # Demonstrate performance benefit by running the same operations without graph
        print("\nRunning same operations without graph for comparison...", file=sys.stderr)

        # Reset results
        result1.fill(0)
        result2.fill(0)
        result3.fill(0)

        start_time = time.time()

        # Individual kernel launches
        launch(stream, config, add_kernel, a.data.ptr, b.data.ptr, result1.data.ptr, cp.uint64(size))
        launch(stream, config, multiply_kernel, result1.data.ptr, c.data.ptr, result2.data.ptr, cp.uint64(size))
        launch(stream, config, subtract_kernel, result2.data.ptr, a.data.ptr, result3.data.ptr, cp.uint64(size))

        stream.sync()
        end_time = time.time()

        individual_execution_time = end_time - start_time
        print(f"Individual kernel execution time: {individual_execution_time:.6f} seconds")

        # Calculate speedup
        speedup = individual_execution_time / graph_execution_time
        print(f"Graph provides {speedup:.2f}x speedup")

        # Verify results again
        assert cp.allclose(result1, expected_result1, rtol=1e-5, atol=1e-5), "Result 1 mismatch"
        assert cp.allclose(result2, expected_result2, rtol=1e-5, atol=1e-5), "Result 2 mismatch"
        assert cp.allclose(result3, expected_result3, rtol=1e-5, atol=1e-5), "Result 3 mismatch"

    finally:
        if graph is not None:
            graph.close()
        if graph_builder is not None:
            graph_builder.close()
        stream.close()
        cp.cuda.Stream.null.use()  # reset CuPy's current stream to the null stream


if __name__ == "__main__":
    main()
