Python-hosted GPU DSL for Codon kernels

[BI71317]

Summary

This Discussion proposes a Python-hosted GPU DSL path for Codon:

import codon
import numpy as np

@codon.gpu
def add_kernel(a: np.ndarray[int, 1], b: np.ndarray[int, 1], c: np.ndarray[int, 1]):
    i = gpu.thread.x
    c[i] = a[i] + b[i]

add_kernel(a, b, c, grid=1, block=n)

The goal for this discussion is to let Python users define and launch raw Codon GPU kernels directly from Python while preserving a familiar decorator API similar to codon.jit.

Motivation

Codon already has native GPU support through import gpu and @gpu.kernel in Codon source.

The missing piece is a Python DSL path where Python code can define raw GPU kernels and launch them without writing a separate .codon file.

This is useful for:

• rapid GPU kernel prototyping from Python
• NumPy-backed host data interop
• testing Codon GPU lowering through Python-facing workflows
• eventually building higher-level Python libraries on top of Codon GPU kernels

User Model

@codon.gpu
def scale_kernel(a: np.ndarray[int, 1], c: np.ndarray[int, 1], factor: int):
    i = gpu.thread.x
    c[i] = a[i] * factor

scale_kernel(a, c, 4, grid=1, block=n)