ripb200node.mojo

ripb200node.mojo

from math import ceildiv
from sys import has_amd_gpu_accelerator, has_nvidia_gpu_accelerator

from gpu import global_idx
from gpu.host import DeviceContext
from layout import Layout, LayoutTensor

alias float_dtype = DType.float32
alias VECTOR_WIDTH = 10
alias BLOCK_SIZE = 5
alias layout = Layout.row_major(VECTOR_WIDTH)


def main():
    constrained[has_nvidia_gpu_accelerator()]()

    # Get context for the attached GPU
    var ctx0 = DeviceContext(0)
    var ctx1 = DeviceContext(1)

    debug_assert(ctx0.can_access(ctx1))
    ctx0.enable_peer_access(ctx1)

    # Allocate data on the GPU address space
    var src_buffer = ctx1.enqueue_create_buffer[float_dtype](VECTOR_WIDTH)
    var dst_buffer = ctx0.enqueue_create_buffer[float_dtype](VECTOR_WIDTH)

    # Fill in values across the entire width
    _ = src_buffer.enqueue_fill(1.25)

    # Wrap the device buffers in tensors
    var src_tensor = LayoutTensor[float_dtype, layout](src_buffer)
    var dst_tensor = LayoutTensor[float_dtype, layout](dst_buffer)

    # Calculate the number of blocks needed to cover the vector
    var grid_dim = ceildiv(VECTOR_WIDTH, BLOCK_SIZE)

    # Launch the vector_addition function as a GPU kernel
    ctx0.enqueue_function[cpy](
        src_tensor,
        dst_tensor,
        VECTOR_WIDTH,
        grid_dim=grid_dim,
        block_dim=BLOCK_SIZE,
    )

    # Map to host so that values can be printed from the CPU
    with dst_buffer.map_to_host() as host_buffer:
        var host_tensor = LayoutTensor[float_dtype, layout](host_buffer)
        print("Resulting vector:", host_tensor)


fn cpy(
    src_tensor: LayoutTensor[float_dtype, layout, MutableAnyOrigin],
    dst_tensor: LayoutTensor[float_dtype, layout, MutableAnyOrigin],
    size: Int,
):
    var global_tid = global_idx.x
    if global_tid < size:
        dst_tensor[global_tid] = src_tensor[global_tid]