[Feature Request] Support for Direct GPU Kernel Definitions and Execution in Mojo

[guna-sd]

Current Behavior:

Writing GPU kernels directly in Mojo is not yet natively supported. As a workaround, it is possible to write the GPU kernel in C or C++ (e.g., by specifying global, device), compile it into a ptx and use it by dynamically loading cuda lib in Mojo using a DLLHandle. The kernel can then be executed via cuLaunchKernel or similar APIs for GPU invocation.

example

from mccl.cuda import CudaLib, cudaDeviceProp
from memory import UnsafePointer

fn main():
    var cudaLib = CudaLib()
    var get_device_count = cudaLib.load_function[fn (UnsafePointer[Int32]) -> Int32]("cudaGetDeviceCount")
    var get_device_properties = cudaLib.load_function[fn (UnsafePointer[cudaDeviceProp], Int32) -> Int32]("cudaGetDeviceProperties")
    
    var deviceCount: Int32 = 0
    _ = get_device_count(UnsafePointer.address_of(deviceCount))
    
    print("Number of CUDA devices:", deviceCount)
    
    for i in range(deviceCount):
        var deviceProp = cudaDeviceProp()
        _ = get_device_properties(UnsafePointer.address_of(deviceProp), i)
        print("Device %d: %s\n", i, deviceProp.name)

Proposed Solution:

Mojo could introduce a new keyword or decorator (e.g., @kernel or @global) to explicitly declare a function as a GPU kernel.
This would allow Mojo to handle such functions differently from regular ones, generating the corresponding low-level LLVM-IR.

example:

@kernel
fn addveckernel(A: UnsafePointer[Float32], B: UnsafePointer[Float32], C: UnsafePointer[Float32], n:Int):
    var tid = threadIdx.x
    var bid = blockIdx.x
    var block_size = blockDim.x

    var global_idx = tid + bid * block_size
    
    if global_idx < n:
        var a_val = A[global_idx].load()
        var b_val = B[global_idx].load()
        var c_val = a_val + b_val
        
        C[global_idx].store(c_val)
    return

could possibly translated to below IR

module {
  target_attr = "nvptx64-nvidia-cuda"
  target_datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64"

  func @llvm.nvvm.read.ptx.sreg.tid.x() -> i32 attributes {llvm.linkage = "external", llvm.target_triple = "nvptx64-nvidia-cuda"}  
  func @llvm.nvvm.read.ptx.sreg.ctaid.x() -> i32 attributes {llvm.linkage = "external", llvm.target_triple = "nvptx64-nvidia-cuda"}  
  func @llvm.nvvm.read.ptx.sreg.ntid.x() -> i32 attributes {llvm.linkage = "external", llvm.target_triple = "nvptx64-nvidia-cuda"}  

  func @addveckernel(%A: memref<*xf32>, %B: memref<*xf32>, %C: memref<*xf32>, %n: i32) -> () {
    %tid = call @llvm.nvvm.read.ptx.sreg.tid.x() : () -> i32
    %bid = call @llvm.nvvm.read.ptx.sreg.ctaid.x() : () -> i32
    %block_size = call @llvm.nvvm.read.ptx.sreg.ntid.x() : () -> i32

    %global_idx = addi %tid, muli %bid, %block_size : i32

    %in_bounds = cmpf slt, %global_idx, %n : i32
    cond_br %in_bounds, ^compute, ^exit

  ^compute:
    %ptrA = getelementptr %A[%global_idx] : memref<*xf32>
    %ptrB = getelementptr %B[%global_idx] : memref<*xf32>
    %ptrC = getelementptr %C[%global_idx] : memref<*xf32>

    %a_val = load %ptrA : memref<*xf32>
    %b_val = load %ptrB : memref<*xf32>
    %c_val = addf %a_val, %b_val : f32
    store %c_val, %ptrC : memref<*xf32>

    br ^exit

  ^exit:
    return
  }

  attribute #llvm.nvvm.annotations = !{!0}
  !0 = !{@addveckernel, !"kernel", i32 1}
}

Reference

Checkout the official mojo blogs for details

Modular Blog about mojo, LLVM, MLIR

Modular Blog about gpu module in mojo

Kernel Launch Mechanism:

Mojo could introduce a LaunchKernel function to launch a GPU kernel, passing the necessary grid and block dimensions. The kernel function would be compiled to PTX code and executed on the GPU directly. The call might look like this:
LaunchKernel(addvecKernel, gridDim=dim3(1, 1, 1), blockDim=dim3(256, 1, 1), args=(a, b, c, n))
or direct launch similar to that of <<>> in c/cpp implementing kernel call in mojo, the call might look like this:
addvecKernel[dim3(1, 1, 1), dim3(256, 1, 1)](a, b, c, n)

Conclusion:

The proposed change to Mojo support for direct GPU kernel definitions and execution would significantly enhance the language's capabilities for high-performance parallel computing.

Related Issues

Beta Give feedback

1. GPU support (plans/roadmap) #3567
   +0