CUDA.jl Tasking

Write custom GPU kernels in Julia using CUDA.jl and execute them through the Legate distributed runtime. Your kernels automatically benefit from Legate's data partitioning, dependency tracking, and multi-GPU scheduling.

Experimental Feature

CUDA.jl tasking is experimental. You must opt in before using @cuda_task or @launch:

cuNumeric.Experimental(true)

The interface has two steps: 2. Compile & Register - @cuda_task JIT-compiles a kernel to PTX and registers it with Legate.

3. Launch - @launch submits the kernel with grid dimensions, inputs, outputs, and scalars.

NDArray arguments are automatically mapped to their CUDA equivalents (NDArray{T,1} → CuDeviceVector{T,1}, etc.). Scalar arguments are passed through by copy.

Inputs vs. outputs

Correctly separating inputs and outputs is critical for Legate's dependency analysis. If an array is both read and written, list it as an output.

Array sizes

Mismatched array sizes are automatically padded to the largest shape. To address this, we plan to add support for other Legate constraints in the future (more information here).

Example

using cuNumeric
using CUDA
import CUDA: i32

# Enable experimental features
cuNumeric.Experimental(true)

# 1. Write a standard CUDA.jl kernel
function kernel_sin(a, b, N)
    i = (blockIdx().x - 1i32) * blockDim().x + threadIdx().x
    if i <= N
        @inbounds b[i] = sin(a[i])
    end
    return nothing
end

N = 1024
threads = 256
blocks = cld(N, threads)

a = cuNumeric.fill(1.0f0, N)
b = cuNumeric.zeros(Float32, N)

# 2. Compile and register (args are used only for type inference)
task = cuNumeric.@cuda_task kernel_sin(a, b, UInt32(1))

# 3. Launch through Legate
cuNumeric.@launch task=task threads=threads blocks=blocks inputs=a outputs=b scalars=UInt32(N)

allowscalar() do
    println("sin(1) = ", b[:][1])  # ≈ 0.8414709
end

See examples/custom_cuda.jl for a more complete example with multiple kernels.

API Reference

cuNumeric.@cuda_task Macro

@cuda_task(f(args...))

Compile a Julia GPU kernel to PTX, register it with the Legate runtime, and return a CUDATask object for later launch.

Arguments

• f — The name of the Julia CUDA.jl GPU kernel function to compile.

• args... — Example arguments to the kernel, used to determine the argument type signature when generating PTX.

Description

This macro automates the process of: 2. Inferring the CUDA argument types for the given args using map_ndarray_cuda_types.

3. Using CUDA.code_ptx to compile the specified GPU kernel (f) into raw PTX text for the inferred types.

4. Extracting the kernel's function symbol name from the PTX using extract_kernel_name.

5. Registering the compiled PTX and kernel name with the Legate runtime via ptx_task, making it available for GPU execution.

6. Returning a CUDATask struct that stores the kernel name and type signature, which can be used to configure and launch the kernel later.

Notes

• The args... are not executed; they are used solely for type inference.

• This macro is intended for use with the Legate runtime and assumes a CUDA context is available.

• Make sure your kernel code is GPU-compatible and does not rely on unsupported Julia features.

Example

mytask = @cuda_task my_kernel(A, B, C)

source

cuNumeric.@launch Macro

@launch(; task, blocks=(1,), threads=(256,), inputs=(), outputs=(), scalars=())

Launch a GPU kernel (previously registered via @cuda_task) through the Legate runtime.

Keywords

• task — A CUDATask object, typically returned by @cuda_task.

• blocks — Tuple or single element specifying the CUDA grid dimensions. Defaults to (1,).

• threads — Tuple or single element specifying the CUDA block dimensions. Defaults to (256,).

• inputs — Tuple or single element of input NDArray objects.

• outputs — Tuple or single element of output NDArray objects.

• scalars — Tuple or single element of scalar values.

Description

The @launch macro validates the provided keywords, ensuring only the allowed set (:task, :blocks, :threads, :inputs, :outputs, :scalars) are present. It then expands to a call to cuNumeric.launch, passing the given arguments to the Legate runtime for execution.

This macro is meant to provide a concise, declarative syntax for launching GPU kernels, separating kernel compilation (via @cuda_task) from execution configuration.

Notes

• task must be a kernel registered with the runtime, usually from @cuda_task.

• All keyword arguments must be specified as assignments, e.g. blocks=(2,2) not positional arguments.

• Defaults are chosen for single-block, 256-thread 1D launches.

• The macro escapes its body so that the values of inputs/outputs/scalars are captured from the surrounding scope at macro expansion time.

Example

mytask = @cuda_task my_kernel(A, B, C)

@launch task=mytask blocks=(8,8) threads=(32,32) inputs=(A, B) outputs=(C)

source