Skip to content

GPU kernels

A saxpy and a block reduction with shared memory and a warp shuffle, written once in ppy.cuda. The same file runs three ways:

  • under CPython, through the reference launch
  • under ppy run, as PTX through the CUDA driver where a device is present
  • as CUDA or HIP source from ppy emit

The printed totals are identical.

Run it

python  saxpy.ppy
ppy run saxpy.ppy
ppy emit cuda saxpy.ppy
ppy emit hip saxpy.ppy
ppy emit ptx saxpy.ppy
ppy inspect saxpy.ppy --stage gpu

The outputs of each command are under What it prints.

Writing a kernel

Kernels and device functions

@cuda.kernel
def saxpy(n: int, a: float, x: native.const_ptr[float], y: native.ptr[float]) -> None:
    i = cuda.global_id()
    if i < n:
        slot = native.offset(y, i)
        native.store(slot, fma(a, native.load(native.offset(x, i)), native.load(slot)))

@cuda.kernel marks a function of scalars and pointers that returns nothing and runs once per thread of a launch. @cuda.device marks fma, a function a kernel calls.

Where a thread is

thread_id, block_id, block_dim, and global_id say where a thread is.

Launching and device memory

cuda.launch(kernel, grid, block, *args) runs the kernel over grid blocks of block threads and waits. A native pointer's whole array goes to the device and, when the pointer is mutable, comes back, so a launch means what the reference launch means.

x and y are made with cuda.device_alloc[float](n) instead. That is memory that lives on the device between launches, filled and read through the same native.store and native.load, so the launch passes an address and copies nothing.

Shared memory, a barrier, a shuffle

block_max parks each thread's value in cuda.shared[float, 64](), waits at syncthreads, trades with its neighbour through shfl_xor(mine, 1), and lets thread 0 finish the reduction. Its code is in How each tool writes the block max.

Under CPython the launch runs a block's threads together, each a Python thread that knows its position, so the barrier and the shuffle are real.

Rules inside device code

Inside device code int arithmetic wraps and nothing guards, as on the device. The gpu dialect's verifier refuses what has no device form:

  • a list or buffer parameter
  • a returned value
  • a call to a host function

Where it ran

cuda.compiled(saxpy) says whether PTX ran. Without a driver, a device, or the NVPTX backend, the reference launch runs and W2008 says why. The line that prints it starts with #, the mark for output that may differ by machine.

  • PPY_CUDA_ARCH picks the architecture the PTX is written for (sm_70 unless set).
  • A built artifact carries its kernels. ppy build stages each PTX beside the manifest, and the launcher binds it without the compiler.

Compared with CuPy, Numba, Mojo, and CUDA C

The same two kernels over sixteen million doubles, written as thread-level kernels the way each tool spells them, in compare/: saxpy_bench.ppy, saxpy_cupy.py, saxpy_numba.py, saxpy.mojo, saxpy.cu. Times are milliseconds, best of warm launches with the device synchronized, over five processes.

Tile-level tools such as Triton and Taichi program a block as one vector and never write the shared-memory exchange, which makes it a different kernel. They are compared with PPy's tile kernels instead of these.

How each tool writes the block max

PPy uses a Python function with cuda.global_id(), shared memory, a barrier, and a shuffle. The same file runs on CPython through the reference launch.

@cuda.kernel
def block_max(x: native.const_ptr[float], out: native.ptr[float]) -> None:
    parked = cuda.shared[float, 64]()
    tid = cuda.thread_id()
    native.store(native.offset(parked, tid), native.load(native.offset(x, cuda.global_id())))
    cuda.syncthreads()
    mine = native.load(native.offset(parked, tid))
    other = cuda.shfl_xor(mine, 1)
    ...

Numba CUDA reads the same way, with @cuda.jit, cuda.grid(1), cuda.shared.array, cuda.syncthreads(), and cuda.shfl_xor_sync:

@cuda.jit
def block_max(x, out):
    parked = cuda.shared.array(64, dtype=np.float64)
    tid = cuda.threadIdx.x
    parked[tid] = x[cuda.grid(1)]
    cuda.syncthreads()
    mine = parked[tid]
    other = cuda.shfl_xor_sync(0xFFFFFFFF, mine, 1)
    ...

CuPy writes saxpy in one line, an ElementwiseKernel, and the block max as CUDA C in a string handed to RawKernel.

CUDA C is the kernel the others approximate, timed with events around the launch alone.

Mojo writes the kernel as a def with global_idx, thread_idx, a stack_allocation in AddressSpace.SHARED, barrier() from MAX's max.gpu.sync, and shuffle_xor. Some details differ from the others:

  • shuffle_xor has no Float64 form, so the double crosses as its bits.
  • Arguments must be fixed-width (Int64, not Int).
  • out is a reserved parameter name.
  • The launch is DeviceContext.enqueue_function with grid_dim and block_dim.
def block_max(x: UnsafePointer[Float64, MutAnyOrigin], result: UnsafePointer[Float64, MutAnyOrigin]):
    var parked = stack_allocation[64, Float64, address_space = AddressSpace.SHARED]()
    var tid = Int(thread_idx.x)
    parked[tid] = x[Int(global_idx.x)]
    barrier()
    var mine = parked[tid]
    var other = bitcast[DType.float64, 1](shuffle_xor(mine.to_bits[DType.uint64](), 1))
    ...

Results

PPy cuda.launch CuPy Numba CUDA Mojo CUDA C
saxpy, arrays on the device 0.77 ± 0.08 0.78 ± 0.08 0.72 ± 0.06 0.55 ± 0.05 0.50 ± 0.00
block max, arrays on the device 2.03 ± 0.05 1.98 ± 0.02 2.02 ± 0.03 1.98 ± 0.02 1.88 ± 0.00
saxpy, arrays copied in and out per launch 39.19 ± 1.60 48.92 ± 2.05 36.79 ± 1.50 28.43 ± 1.45 30.46 ± 0.74
block max, array copied in per launch 15.66 ± 0.43 12.33 ± 0.34 14.11 ± 0.42 11.21 ± 0.33 11.26 ± 0.20

With the arrays on the device, a launch is the kernel. The five tools run the same block max in the same time. On saxpy the Python-hosted ones sit a tenth of a millisecond above Mojo and CUDA C, the cost of a launch through the interpreter.

The copying rows use the other memory model: a native.stack_alloc array sent in and brought back on every launch. There the driver reads the host array in place, and the traffic costs what it costs every tool. A program that launches more than once keeps its data on the device by allocating it there, with cuda.device_alloc.

NVIDIA GeForce RTX 5080 Laptop GPU, driver 610.71, CUDA 13.3; CuPy 14.2.0, Numba 0.67.0 on CPython 3.12.13; Mojo 1.0.0 with MAX 26.5; nvcc 13.3; PPy on CPython 3.13.13.

What it prints

python saxpy.ppy

90000.0
100.0 98.0
# kernels compiled for a device here: False

ppy run saxpy.ppy

90000.0
100.0 98.0
# kernels compiled for a device here: True

ppy emit cuda saxpy.ppy

79 lines 
/* saxpy: generated by ppy, CUDA C++ */
#include <cuda_runtime.h>
#include <cmath>
#include <cstdint>
#include <cstdlib>

static inline __host__ __device__ int ppy_ovf_add_i64(int64_t a, int64_t b, int64_t *out) {
#if defined(__GNUC__) || defined(__clang__)
    return __builtin_add_overflow(a, b, out);
#else
    if ((b > 0 && a > INT64_MAX - b) || (b < 0 && a < INT64_MIN - b)) {
        return 1;
    }
    *out = a + b;
    return 0;
#endif
}

extern "C" {

static __device__ double ppy_saxpy_fma(double a, double x, double y);
__global__ void ppy_saxpy_saxpy(int64_t n, double a, const double *x, double *y);
__global__ void ppy_saxpy_block_max(const double *x, double *out);
int32_t ppy_saxpy_run(int64_t n, double a, const double *x, double *y, int64_t *out);

static __device__ double ppy_saxpy_fma(double a, double x, double y) {
    return a * x + y;
}

__global__ void ppy_saxpy_saxpy(int64_t n, double a, const double *x, double *y) {
    int64_t t1 = (int64_t)blockIdx.x;
    int64_t t2 = (int64_t)blockDim.x;
    int64_t t3 = (int64_t)threadIdx.x;
    int64_t i = int64_t(uint64_t(t1) * uint64_t(t2) + uint64_t(t3));
    if (i < n) {
        double *slot = y + i;
        *slot = ppy_saxpy_fma(a, x[i], *slot);
    }
}

__global__ void ppy_saxpy_block_max(const double *x, double *out) {
    __shared__ double shared[64];

    double *parked = shared;
    int64_t tid = (int64_t)threadIdx.x;
    int64_t t1 = (int64_t)blockIdx.x;
    int64_t t2 = (int64_t)blockDim.x;
    int64_t t3 = (int64_t)threadIdx.x;
    parked[tid] = x[int64_t(uint64_t(t1) * uint64_t(t2) + uint64_t(t3))];
    __syncthreads();
    double mine = parked[tid];
    double other = __shfl_xor_sync(0xffffffffu, mine, (int)1);
    mine = other > mine ? other : mine;
    parked[tid] = mine;
    __syncthreads();
    if (tid == 0) {
        double best = *parked;
        int64_t t4 = (int64_t)blockDim.x;
        int64_t k = 1;
        while (k < t4) {
            double candidate = parked[k];
            best = candidate > best ? candidate : best;
            k = int64_t(uint64_t(k) + 1u);
        }
        int64_t t5 = (int64_t)blockIdx.x;
        out[t5] = best;
    }
}

int32_t ppy_saxpy_run(int64_t n, double a, const double *x, double *y, int64_t *out) {
    int64_t t1;
    if (ppy_ovf_add_i64(n, 255, &t1)) return 1; /* arith.ok */
    ppy_saxpy_saxpy<<<dim3((unsigned)(t1 / 256 - (t1 % 256 < 0)), (unsigned)1, (unsigned)1), dim3((unsigned)256, (unsigned)1, (unsigned)1)>>>(n, a, x, y);
    if (cudaDeviceSynchronize() != cudaSuccess) return 1; /* launch.ok */
    *out = 0;
    return 0;
}

} /* extern "C" */

ppy emit hip saxpy.ppy

79 lines 
/* saxpy: generated by ppy, HIP C++ */
#include <hip/hip_runtime.h>
#include <cmath>
#include <cstdint>
#include <cstdlib>

static inline __host__ __device__ int ppy_ovf_add_i64(int64_t a, int64_t b, int64_t *out) {
#if defined(__GNUC__) || defined(__clang__)
    return __builtin_add_overflow(a, b, out);
#else
    if ((b > 0 && a > INT64_MAX - b) || (b < 0 && a < INT64_MIN - b)) {
        return 1;
    }
    *out = a + b;
    return 0;
#endif
}

extern "C" {

static __device__ double ppy_saxpy_fma(double a, double x, double y);
__global__ void ppy_saxpy_saxpy(int64_t n, double a, const double *x, double *y);
__global__ void ppy_saxpy_block_max(const double *x, double *out);
int32_t ppy_saxpy_run(int64_t n, double a, const double *x, double *y, int64_t *out);

static __device__ double ppy_saxpy_fma(double a, double x, double y) {
    return a * x + y;
}

__global__ void ppy_saxpy_saxpy(int64_t n, double a, const double *x, double *y) {
    int64_t t1 = (int64_t)blockIdx.x;
    int64_t t2 = (int64_t)blockDim.x;
    int64_t t3 = (int64_t)threadIdx.x;
    int64_t i = int64_t(uint64_t(t1) * uint64_t(t2) + uint64_t(t3));
    if (i < n) {
        double *slot = y + i;
        *slot = ppy_saxpy_fma(a, x[i], *slot);
    }
}

__global__ void ppy_saxpy_block_max(const double *x, double *out) {
    __shared__ double shared[64];

    double *parked = shared;
    int64_t tid = (int64_t)threadIdx.x;
    int64_t t1 = (int64_t)blockIdx.x;
    int64_t t2 = (int64_t)blockDim.x;
    int64_t t3 = (int64_t)threadIdx.x;
    parked[tid] = x[int64_t(uint64_t(t1) * uint64_t(t2) + uint64_t(t3))];
    __syncthreads();
    double mine = parked[tid];
    double other = __shfl_xor(mine, (int)1);
    mine = other > mine ? other : mine;
    parked[tid] = mine;
    __syncthreads();
    if (tid == 0) {
        double best = *parked;
        int64_t t4 = (int64_t)blockDim.x;
        int64_t k = 1;
        while (k < t4) {
            double candidate = parked[k];
            best = candidate > best ? candidate : best;
            k = int64_t(uint64_t(k) + 1u);
        }
        int64_t t5 = (int64_t)blockIdx.x;
        out[t5] = best;
    }
}

int32_t ppy_saxpy_run(int64_t n, double a, const double *x, double *y, int64_t *out) {
    int64_t t1;
    if (ppy_ovf_add_i64(n, 255, &t1)) return 1; /* arith.ok */
    ppy_saxpy_saxpy<<<dim3((unsigned)(t1 / 256 - (t1 % 256 < 0)), (unsigned)1, (unsigned)1), dim3((unsigned)256, (unsigned)1, (unsigned)1)>>>(n, a, x, y);
    if (hipDeviceSynchronize() != hipSuccess) return 1; /* launch.ok */
    *out = 0;
    return 0;
}

} /* extern "C" */

ppy emit ptx saxpy.ppy

146 lines 
//
// Generated by LLVM NVPTX Back-End
//

.version 7.0
.target sm_70
.address_size 64

    // .globl   ppy_saxpy_saxpy
// ppy_saxpy_block_max_shared1 has been demoted

.visible .entry ppy_saxpy_saxpy(
    .param .u64 ppy_saxpy_saxpy_param_0,
    .param .f64 ppy_saxpy_saxpy_param_1,
    .param .u64 .ptr .align 1 ppy_saxpy_saxpy_param_2,
    .param .u64 .ptr .align 1 ppy_saxpy_saxpy_param_3
)
{
    .reg .pred  %p<2>;
    .reg .b32   %r<4>;
    .reg .b64   %rd<17>;

    ld.param.b64    %rd5, [ppy_saxpy_saxpy_param_0];
    ld.param.b64    %rd6, [ppy_saxpy_saxpy_param_3];
    cvta.to.global.u64  %rd1, %rd6;
    ld.param.b64    %rd7, [ppy_saxpy_saxpy_param_2];
    cvta.to.global.u64  %rd2, %rd7;
    mov.u32     %r1, %ctaid.x;
    mov.u32     %r2, %ntid.x;
    mul.wide.u32    %rd8, %r1, %r2;
    mov.u32     %r3, %tid.x;
    cvt.u64.u32     %rd9, %r3;
    add.s64     %rd3, %rd8, %rd9;
    setp.ge.s64     %p1, %rd3, %rd5;
    @%p1 bra    $L__BB0_2;
    ld.param.b64    %rd4, [ppy_saxpy_saxpy_param_1];
    shl.b64     %rd10, %rd3, 3;
    add.s64     %rd11, %rd1, %rd10;
    add.s64     %rd12, %rd2, %rd10;
    ld.global.b64   %rd13, [%rd12];
    ld.global.b64   %rd14, [%rd11];
    mul.rn.f64  %rd15, %rd4, %rd13;
    add.rn.f64  %rd16, %rd15, %rd14;
    st.global.b64   [%rd11], %rd16;
$L__BB0_2:
    ret;

}
    // .globl   ppy_saxpy_block_max
.visible .entry ppy_saxpy_block_max(
    .param .u64 .ptr .align 1 ppy_saxpy_block_max_param_0,
    .param .u64 .ptr .align 1 ppy_saxpy_block_max_param_1
)
{
    .reg .pred  %p<13>;
    .reg .b32   %r<10>;
    .reg .b64   %rd<40>;
    // demoted variable
    .shared .align 8 .b8 ppy_saxpy_block_max_shared1[512];
    ld.param.b64    %rd6, [ppy_saxpy_block_max_param_0];
    cvta.to.global.u64  %rd7, %rd6;
    ld.param.b64    %rd8, [ppy_saxpy_block_max_param_1];
    cvta.to.global.u64  %rd1, %rd8;
    mov.u32     %r1, %tid.x;
    mul.wide.u32    %rd9, %r1, 8;
    mov.b64     %rd10, ppy_saxpy_block_max_shared1;
    add.s64     %rd11, %rd10, %rd9;
    mov.u32     %r2, %ctaid.x;
    mov.u32     %r3, %ntid.x;
    mul.wide.u32    %rd12, %r2, %r3;
    shl.b64     %rd13, %rd12, 3;
    add.s64     %rd14, %rd7, %rd13;
    add.s64     %rd15, %rd14, %rd9;
    ld.global.b64   %rd16, [%rd15];
    st.shared.b64   [%rd11], %rd16;
    bar.sync    0;
    ld.shared.b64   %rd17, [%rd11];
    cvt.u32.u64     %r4, %rd17;
    shfl.sync.bfly.b32  %r5, %r4, 1, 31, -1;
    { .reg .b32 tmp; mov.b64 {tmp, %r6}, %rd17; }
    shfl.sync.bfly.b32  %r7, %r6, 1, 31, -1;
    cvt.u64.u32     %rd18, %r7;
    shl.b64     %rd19, %rd18, 32;
    cvt.u64.u32     %rd20, %r5;
    or.b64  %rd21, %rd19, %rd20;
    setp.lt.f64     %p1, %rd17, %rd21;
    selp.f64    %rd22, %rd21, %rd17, %p1;
    st.shared.b64   [%rd11], %rd22;
    bar.sync    0;
    setp.ne.b32     %p2, %r1, 0;
    @%p2 bra    $L__BB1_9;
    cvt.u64.u32     %rd2, %r2;
    cvt.u64.u32     %rd3, %r3;
    cvt.u32.u64     %r8, %rd3;
    ld.shared.b64   %rd39, [ppy_saxpy_block_max_shared1];
    setp.lt.u32     %p3, %r8, 2;
    @%p3 bra    $L__BB1_8;
    add.s64     %rd4, %rd3, -1;
    and.b64     %rd37, %rd4, 3;
    add.s32     %r9, %r8, -2;
    setp.lt.u32     %p4, %r9, 3;
    mov.b64     %rd36, 1;
    @%p4 bra    $L__BB1_6;
    add.s64     %rd34, %rd10, 16;
    and.b64     %rd5, %rd4, -4;
    mov.b64     %rd35, 0;
$L__BB1_4:
    ld.shared.b64   %rd23, [%rd34+-8];
    setp.gt.f64     %p5, %rd23, %rd39;
    selp.f64    %rd24, %rd23, %rd39, %p5;
    ld.shared.b64   %rd25, [%rd34];
    setp.gt.f64     %p6, %rd25, %rd24;
    selp.f64    %rd26, %rd25, %rd24, %p6;
    ld.shared.b64   %rd27, [%rd34+8];
    setp.gt.f64     %p7, %rd27, %rd26;
    selp.f64    %rd28, %rd27, %rd26, %p7;
    ld.shared.b64   %rd29, [%rd34+16];
    setp.gt.f64     %p8, %rd29, %rd28;
    selp.f64    %rd39, %rd29, %rd28, %p8;
    add.s64     %rd35, %rd35, 4;
    add.s64     %rd34, %rd34, 32;
    setp.ne.b64     %p9, %rd5, %rd35;
    @%p9 bra    $L__BB1_4;
    setp.eq.b64     %p10, %rd37, 0;
    add.s64     %rd36, %rd35, 1;
    @%p10 bra   $L__BB1_8;
$L__BB1_6:
    shl.b64     %rd30, %rd36, 3;
    add.s64     %rd38, %rd10, %rd30;
$L__BB1_7:
    .pragma "nounroll";
    ld.shared.b64   %rd31, [%rd38];
    setp.gt.f64     %p11, %rd31, %rd39;
    selp.f64    %rd39, %rd31, %rd39, %p11;
    add.s64     %rd38, %rd38, 8;
    add.s64     %rd37, %rd37, -1;
    setp.ne.b64     %p12, %rd37, 0;
    @%p12 bra   $L__BB1_7;
$L__BB1_8:
    shl.b64     %rd32, %rd2, 3;
    add.s64     %rd33, %rd1, %rd32;
    st.global.b64   [%rd33], %rd39;
$L__BB1_9:
    ret;

}

ppy inspect saxpy.ppy --stage gpu

165 lines 
; ---- saxpy [gpu] ----
func @saxpy_fma(%a: f64, %x: f64, %y: f64) -> f64 attrs {effects = [], gpu.kind = "device", ppy.abi = "ppy", ppy.qualname = "saxpy.fma", ppy.releases_gil = true, ppy.symbol = "ppy_saxpy_fma"} loc("examples/38_cuda/saxpy.ppy":5:0) {
^entry:
    %a_addr = core.alloca : ptr<f64, stack> loc("examples/38_cuda/saxpy.ppy":5:0)
    core.store %a, %a_addr
    %x_addr = core.alloca : ptr<f64, stack>
    core.store %x, %x_addr
    %y_addr = core.alloca : ptr<f64, stack>
    core.store %y, %y_addr
    %0 = core.load %a_addr : f64 loc("examples/38_cuda/saxpy.ppy":6:4)
    %1 = core.load %x_addr : f64
    %2 = core.mul %0, %1 : f64
    %3 = core.load %y_addr : f64
    %4 = core.add %2, %3 : f64
    core.ret %4
}

func @saxpy_saxpy(%n: i64, %a: f64, %x: ptr<f64, generic, const>, %y: ptr<f64>) -> () attrs {effects = ["read_memory", "write_memory"], gpu.kind = "kernel", ppy.abi = "ppy", ppy.qualname = "saxpy.saxpy", ppy.releases_gil = true, ppy.symbol = "ppy_saxpy_saxpy"} loc("examples/38_cuda/saxpy.ppy":10:0) {
^entry:
    %n_addr = core.alloca : ptr<i64, stack> loc("examples/38_cuda/saxpy.ppy":10:0)
    core.store %n, %n_addr
    %a_addr = core.alloca : ptr<f64, stack>
    core.store %a, %a_addr
    %x_addr = core.alloca : ptr<ptr<f64, generic, const>, stack>
    core.store %x, %x_addr
    %y_addr = core.alloca : ptr<ptr<f64>, stack>
    core.store %y, %y_addr
    %0 = gpu.block_id.x : index loc("examples/38_cuda/saxpy.ppy":11:4)
    %1 = gpu.block_dim.x : index
    %2 = core.mul %0, %1 {overflow = "wrap"} : index
    %3 = gpu.thread_id.x : index
    %4 = core.add %2, %3 {overflow = "wrap"} : index
    %5 = core.cast %4 : i64
    %i_addr = core.alloca : ptr<i64, stack>
    core.store %5, %i_addr
    %6 = core.load %i_addr : i64 loc("examples/38_cuda/saxpy.ppy":12:4)
    %7 = core.load %n_addr : i64
    %n_entry = core.load %n_addr : i64
    %8 = core.cmp.lt %6, %7 : bool
    %slot_addr = core.alloca : ptr<ptr<f64>, stack>
    core.cond_br %8, ^then1, ^else2 loc("examples/38_cuda/saxpy.ppy":12:4)
^then1:
    %9 = core.load %y_addr : ptr<f64> loc("examples/38_cuda/saxpy.ppy":13:8)
    %10 = core.load %i_addr : i64
    %11 = core.ptr_offset %9, %10 : ptr<f64>
    core.store %11, %slot_addr
    %12 = core.load %slot_addr : ptr<f64> loc("examples/38_cuda/saxpy.ppy":14:8)
    %13 = core.load %a_addr : f64
    %14 = core.load %x_addr : ptr<f64, generic, const>
    %15 = core.load %i_addr : i64
    %16 = core.ptr_offset %14, %15 : ptr<f64, generic, const>
    %17 = core.load %16 : f64
    %18 = core.load %slot_addr : ptr<f64>
    %19 = core.load %18 : f64
    %20 = core.call %13, %17, %19 {callee = @saxpy_fma} : f64
    core.store %20, %12
    core.br ^endif3
^else2:
    core.br ^endif3 loc("examples/38_cuda/saxpy.ppy":14:8)
^endif3:
    core.ret loc("examples/38_cuda/saxpy.ppy":14:8)
}

func @saxpy_block_max(%x: ptr<f64, generic, const>, %out: ptr<f64>) -> () attrs {effects = ["alloc", "may_raise", "read_memory", "sync", "write_memory"], gpu.kind = "kernel", ppy.abi = "ppy", ppy.qualname = "saxpy.block_max", ppy.releases_gil = true, ppy.symbol = "ppy_saxpy_block_max"} loc("examples/38_cuda/saxpy.ppy":18:0) {
^entry:
    %x_addr = core.alloca : ptr<ptr<f64, generic, const>, stack> loc("examples/38_cuda/saxpy.ppy":18:0)
    core.store %x, %x_addr
    %out_addr = core.alloca : ptr<ptr<f64>, stack>
    core.store %out, %out_addr
    %0 = gpu.shared_alloc {count = 64} : ptr<f64, shared> loc("examples/38_cuda/saxpy.ppy":19:4)
    %parked_addr = core.alloca : ptr<ptr<f64, shared>, stack>
    core.store %0, %parked_addr
    %1 = gpu.thread_id.x : index loc("examples/38_cuda/saxpy.ppy":20:4)
    %2 = core.cast %1 : i64
    %tid_addr = core.alloca : ptr<i64, stack>
    core.store %2, %tid_addr
    %3 = core.load %parked_addr : ptr<f64, shared> loc("examples/38_cuda/saxpy.ppy":21:4)
    %4 = core.load %tid_addr : i64
    %5 = core.ptr_offset %3, %4 : ptr<f64, shared>
    %6 = core.load %x_addr : ptr<f64, generic, const>
    %7 = gpu.block_id.x : index
    %8 = gpu.block_dim.x : index
    %9 = core.mul %7, %8 {overflow = "wrap"} : index
    %10 = gpu.thread_id.x : index
    %11 = core.add %9, %10 {overflow = "wrap"} : index
    %12 = core.cast %11 : i64
    %13 = core.ptr_offset %6, %12 : ptr<f64, generic, const>
    %14 = core.load %13 : f64
    core.store %14, %5
    gpu.barrier loc("examples/38_cuda/saxpy.ppy":22:4)
    %15 = core.load %parked_addr : ptr<f64, shared> loc("examples/38_cuda/saxpy.ppy":23:4)
    %16 = core.load %tid_addr : i64
    %17 = core.ptr_offset %15, %16 : ptr<f64, shared>
    %18 = core.load %17 : f64
    %mine_addr = core.alloca : ptr<f64, stack>
    core.store %18, %mine_addr
    %19 = core.load %mine_addr : f64 loc("examples/38_cuda/saxpy.ppy":24:4)
    %20 = core.const 1 : i64
    %21 = gpu.subgroup_shuffle.xor %19, %20 : f64
    %other_addr = core.alloca : ptr<f64, stack>
    core.store %21, %other_addr
    %22 = core.load %other_addr : f64 loc("examples/38_cuda/saxpy.ppy":25:4)
    %23 = core.load %mine_addr : f64
    %24 = core.cmp.gt %22, %23 : bool
    %25 = core.load %other_addr : f64
    %26 = core.load %mine_addr : f64
    %27 = core.select %24, %25, %26 : f64
    core.store %27, %mine_addr
    %28 = core.load %parked_addr : ptr<f64, shared> loc("examples/38_cuda/saxpy.ppy":26:4)
    %29 = core.load %tid_addr : i64
    %30 = core.ptr_offset %28, %29 : ptr<f64, shared>
    %31 = core.load %mine_addr : f64
    core.store %31, %30
    gpu.barrier loc("examples/38_cuda/saxpy.ppy":27:4)
    %32 = core.load %tid_addr : i64 loc("examples/38_cuda/saxpy.ppy":28:4)
    %33 = core.const 0 : i64
    %34 = core.cmp.eq %32, %33 : bool
    %best_addr = core.alloca : ptr<f64, stack>
    %35 = core.const 1 : i64
    %36 = core.const 1 : i64
    %k_addr = core.alloca : ptr<i64, stack>
    %candidate_addr = core.alloca : ptr<f64, stack>
    core.cond_br %34, ^then1, ^else2 loc("examples/38_cuda/saxpy.ppy":28:4)
^then1:
    %37 = core.load %parked_addr : ptr<f64, shared> loc("examples/38_cuda/saxpy.ppy":29:8)
    %38 = core.load %37 : f64
    core.store %38, %best_addr
    %39 = gpu.block_dim.x : index loc("examples/38_cuda/saxpy.ppy":30:8)
    %40 = core.cast %39 : i64
    core.store %35, %k_addr
    core.br ^for.head6
^else2:
    core.br ^endif3 loc("examples/38_cuda/saxpy.ppy":33:8)
^endif3:
    core.ret loc("examples/38_cuda/saxpy.ppy":33:8)
^for.head6:
    %41 = core.load %k_addr : i64 loc("examples/38_cuda/saxpy.ppy":30:8)
    %42 = core.cmp.lt %41, %40 : bool
    core.cond_br %42, ^for.body7, ^for.end9
^for.body7:
    %43 = core.load %parked_addr : ptr<f64, shared> loc("examples/38_cuda/saxpy.ppy":31:12)
    %44 = core.load %k_addr : i64
    %45 = core.ptr_offset %43, %44 : ptr<f64, shared>
    %46 = core.load %45 : f64
    core.store %46, %candidate_addr
    %47 = core.load %candidate_addr : f64 loc("examples/38_cuda/saxpy.ppy":32:12)
    %48 = core.load %best_addr : f64
    %49 = core.cmp.gt %47, %48 : bool
    %50 = core.load %candidate_addr : f64
    %51 = core.load %best_addr : f64
    %52 = core.select %49, %50, %51 : f64
    core.store %52, %best_addr
    %53 = core.load %k_addr : i64
    %54 = core.add %53, %36 {overflow = "wrap"} : i64
    core.store %54, %k_addr
    core.br ^for.head6
^for.end9:
    %55 = core.load %out_addr : ptr<f64> loc("examples/38_cuda/saxpy.ppy":33:8)
    %56 = gpu.block_id.x : index
    %57 = core.cast %56 : i64
    %58 = core.ptr_offset %55, %57 : ptr<f64>
    %59 = core.load %best_addr : f64
    core.store %59, %58
    core.br ^endif3
}

Read on: GPU kernels · The IR: the gpu dialect

saxpy.ppy is hand-written; there is no .py source and no conversion step.

38_cuda/saxpy.ppy

from ppy import cuda, native


@cuda.device
def fma(a: float, x: float, y: float) -> float:
    return a * x + y


@cuda.kernel
def saxpy(n: int, a: float, x: native.const_ptr[float], y: native.ptr[float]) -> None:
    i = cuda.global_id()
    if i < n:
        slot = native.offset(y, i)
        native.store(slot, fma(a, native.load(native.offset(x, i)), native.load(slot)))


@cuda.kernel
def block_max(x: native.const_ptr[float], out: native.ptr[float]) -> None:
    parked = cuda.shared[float, 64]()
    tid = cuda.thread_id()
    native.store(native.offset(parked, tid), native.load(native.offset(x, cuda.global_id())))
    cuda.syncthreads()
    mine = native.load(native.offset(parked, tid))
    other = cuda.shfl_xor(mine, 1)
    mine = other if other > mine else mine
    native.store(native.offset(parked, tid), mine)
    cuda.syncthreads()
    if tid == 0:
        best = native.load(parked)
        for k in range(1, cuda.block_dim()):
            candidate = native.load(native.offset(parked, k))
            best = candidate if candidate > best else best
        native.store(native.offset(out, cuda.block_id()), best)


def run(n: int, a: float, x: native.const_ptr[float], y: native.ptr[float]) -> None:
    cuda.launch(saxpy, (n + 255) // 256, 256, n, a, x, y)


def main() -> None:
    n = 300
    x = cuda.device_alloc[float](n)
    y = cuda.device_alloc[float](n)
    for i in range(n):
        native.store(native.offset(x, i), float(i))
        native.store(native.offset(y, i), 1.0)
    run(n, 2.0, x, y)
    total = 0.0
    for i in range(n):
        total += native.load(native.offset(y, i))
    print(total)
    values = native.stack_alloc[float](128)
    for i in range(128):
        native.store(native.offset(values, i), float((i * 37) % 101))
    out = native.stack_alloc[float](2)
    cuda.launch(block_max, 2, 64, values, out)
    print(native.load(out), native.load(native.offset(out, 1)))
    print(f"# kernels compiled for a device here: {cuda.compiled(saxpy)}")


main()

Counterpart programs

The programs the comparison above measured, each written the way its tool expects. The PPy one is first.

saxpy_bench.ppy (PPy) 
"""saxpy over sixteen million doubles and a per-block max: PPY kernels over device memory
and over host arrays."""

import time

from ppy import cuda, native

N = 1 << 24


@cuda.device
def fma(a: float, x: float, y: float) -> float:
    return a * x + y


@cuda.kernel
def saxpy(n: int, a: float, x: native.const_ptr[float], y: native.ptr[float]) -> None:
    i = cuda.global_id()
    if i < n:
        slot = native.offset(y, i)
        native.store(slot, fma(a, native.load(native.offset(x, i)), native.load(slot)))


@cuda.kernel
def block_max(x: native.const_ptr[float], out: native.ptr[float]) -> None:
    parked = cuda.shared[float, 64]()
    tid = cuda.thread_id()
    native.store(native.offset(parked, tid), native.load(native.offset(x, cuda.global_id())))
    cuda.syncthreads()
    mine = native.load(native.offset(parked, tid))
    other = cuda.shfl_xor(mine, 1)
    mine = other if other > mine else mine
    native.store(native.offset(parked, tid), mine)
    cuda.syncthreads()
    if tid == 0:
        best = native.load(parked)
        for k in range(1, cuda.block_dim()):
            candidate = native.load(native.offset(parked, k))
            best = candidate if candidate > best else best
        native.store(native.offset(out, cuda.block_id()), best)


def run_saxpy(n: int, a: float, x: native.const_ptr[float], y: native.ptr[float]) -> None:
    cuda.launch(saxpy, (n + 255) // 256, 256, n, a, x, y)


def run_block_max(blocks: int, x: native.const_ptr[float], out: native.ptr[float]) -> None:
    cuda.launch(block_max, blocks, 64, x, out)


def main() -> None:
    x = cuda.device_alloc[float](N)
    y = cuda.device_alloc[float](N)
    values = cuda.device_alloc[float](N)
    out = cuda.device_alloc[float](N // 64)
    for i in range(N):
        native.store(native.offset(x, i), float(i))
        native.store(native.offset(y, i), 1.0)
        native.store(native.offset(values, i), float((i * 37) % 101))
    run_saxpy(N, 2.0, x, y)
    run_block_max(N // 64, values, out)
    total = 0.0
    for i in range(N):
        total += native.load(native.offset(y, i))
    best = 0.0
    for i in range(N // 64):
        candidate = native.load(native.offset(out, i))
        best = candidate if candidate > best else best
    print(total, best)
    print(f"# kernels compiled for a device here: {cuda.compiled(saxpy)}")
    best_saxpy = 1e9
    best_block = 1e9
    for _ in range(5):
        started = time.perf_counter()
        run_saxpy(N, 2.0, x, y)
        best_saxpy = min(best_saxpy, (time.perf_counter() - started) * 1000.0)
        started = time.perf_counter()
        run_block_max(N // 64, values, out)
        best_block = min(best_block, (time.perf_counter() - started) * 1000.0)
    print(f"# saxpy: {best_saxpy:.3f} ms")
    print(f"# block_max: {best_block:.3f} ms")
    # The same launches over host arrays: every launch copies the arrays in, the mutable ones back.
    hx = native.stack_alloc[float](N)
    hy = native.stack_alloc[float](N)
    hv = native.stack_alloc[float](N)
    hout = native.stack_alloc[float](N // 64)
    for i in range(N):
        native.store(native.offset(hx, i), float(i))
        native.store(native.offset(hy, i), 1.0)
        native.store(native.offset(hv, i), float((i * 37) % 101))
    best_saxpy = 1e9
    best_block = 1e9
    for _ in range(5):
        started = time.perf_counter()
        run_saxpy(N, 2.0, hx, hy)
        best_saxpy = min(best_saxpy, (time.perf_counter() - started) * 1000.0)
        started = time.perf_counter()
        run_block_max(N // 64, hv, hout)
        best_block = min(best_block, (time.perf_counter() - started) * 1000.0)
    print(f"# saxpy with copies: {best_saxpy:.3f} ms")
    print(f"# block_max with copies: {best_block:.3f} ms")


main()
saxpy.cu (CUDA C) 
// saxpy over sixteen million doubles and a per-block max: CUDA C, timed with events.
#include <cstdio>
#include <cuda_runtime.h>

#define N (1 << 24)

__device__ double fma3(double a, double x, double y) { return a * x + y; }

__global__ void saxpy(int n, double a, const double *x, double *y) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) y[i] = fma3(a, x[i], y[i]);
}

__global__ void block_max(const double *x, double *out) {
    __shared__ double parked[64];
    int tid = threadIdx.x;
    parked[tid] = x[blockIdx.x * blockDim.x + threadIdx.x];
    __syncthreads();
    double mine = parked[tid];
    double other = __shfl_xor_sync(0xffffffffu, mine, 1);
    mine = other > mine ? other : mine;
    parked[tid] = mine;
    __syncthreads();
    if (tid == 0) {
        double best = parked[0];
        for (int k = 1; k < blockDim.x; k++) best = parked[k] > best ? parked[k] : best;
        out[blockIdx.x] = best;
    }
}

static float timed(void (*run)(), const char *label) {
    cudaEvent_t start, stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);
    float best = 1e9f;
    run();
    cudaDeviceSynchronize();
    for (int r = 0; r < 10; r++) {
        cudaEventRecord(start);
        run();
        cudaEventRecord(stop);
        cudaEventSynchronize(stop);
        float ms = 0;
        cudaEventElapsedTime(&ms, start, stop);
        if (ms < best) best = ms;
    }
    printf("# %s: %.3f ms\n", label, best);
    return best;
}

static double *dx, *dy, *dv, *dout, *hx, *hy, *hv, *hout;
static void run_saxpy() { saxpy<<<(N + 255) / 256, 256>>>(N, 2.0, dx, dy); }
static void run_block_max() { block_max<<<N / 64, 64>>>(dv, dout); }

/* The same launches over host arrays: the inputs copied in, the results copied out. */
static void run_saxpy_copying() {
    cudaMemcpy(dx, hx, N * sizeof(double), cudaMemcpyHostToDevice);
    cudaMemcpy(dy, hy, N * sizeof(double), cudaMemcpyHostToDevice);
    run_saxpy();
    cudaMemcpy(hy, dy, N * sizeof(double), cudaMemcpyDeviceToHost);
}

static void run_block_max_copying() {
    cudaMemcpy(dv, hv, N * sizeof(double), cudaMemcpyHostToDevice);
    run_block_max();
    cudaMemcpy(hout, dout, N / 64 * sizeof(double), cudaMemcpyDeviceToHost);
}

int main() {
    hx = new double[N]; hy = new double[N]; hv = new double[N]; hout = new double[N / 64];
    for (int i = 0; i < N; i++) { hx[i] = (double)i; hy[i] = 1.0; hv[i] = (double)((i * 37) % 101); }
    cudaMalloc(&dx, N * sizeof(double)); cudaMalloc(&dy, N * sizeof(double));
    cudaMalloc(&dv, N * sizeof(double)); cudaMalloc(&dout, N / 64 * sizeof(double));
    cudaMemcpy(dx, hx, N * sizeof(double), cudaMemcpyHostToDevice);
    cudaMemcpy(dy, hy, N * sizeof(double), cudaMemcpyHostToDevice);
    cudaMemcpy(dv, hv, N * sizeof(double), cudaMemcpyHostToDevice);
    run_saxpy();
    run_block_max();
    cudaMemcpy(hy, dy, N * sizeof(double), cudaMemcpyDeviceToHost);
    cudaMemcpy(hout, dout, N / 64 * sizeof(double), cudaMemcpyDeviceToHost);
    double total = 0.0, best = 0.0;
    for (int i = 0; i < N; i++) total += hy[i];
    for (int i = 0; i < N / 64; i++) best = hout[i] > best ? hout[i] : best;
    printf("%.1f %.1f\n", total, best);
    timed(run_saxpy, "saxpy");
    timed(run_block_max, "block_max");
    cudaMemcpy(hy, dy, N * sizeof(double), cudaMemcpyDeviceToHost);
    timed(run_saxpy_copying, "saxpy with copies");
    timed(run_block_max_copying, "block_max with copies");
    return 0;
}
saxpy.mojo (Mojo) 
"""saxpy over sixteen million doubles and a per-block max: Mojo 1.0 kernels through MAX's DeviceContext."""

from max.gpu.host import DeviceContext
from max.gpu.sync import barrier
from std.gpu import block_dim, block_idx, global_idx, thread_idx
from std.gpu.primitives.warp import shuffle_xor
from std.memory import AddressSpace, UnsafePointer, bitcast, stack_allocation
from std.time import perf_counter_ns

comptime N = 1 << 24


def fma(a: Float64, x: Float64, y: Float64) -> Float64:
    return a * x + y


def saxpy(n: Int64, a: Float64, x: UnsafePointer[Float64, MutAnyOrigin], y: UnsafePointer[Float64, MutAnyOrigin]):
    var i = Int(global_idx.x)
    if Int64(i) < n:
        y[i] = fma(a, x[i], y[i])


def block_max(x: UnsafePointer[Float64, MutAnyOrigin], result: UnsafePointer[Float64, MutAnyOrigin]):
    var parked = stack_allocation[64, Float64, address_space = AddressSpace.SHARED]()
    var tid = Int(thread_idx.x)
    parked[tid] = x[Int(global_idx.x)]
    barrier()
    var mine = parked[tid]
    # The shuffle has no Float64 form here: the bits go across as an Int64.
    var other = bitcast[DType.float64, 1](shuffle_xor(mine.to_bits[DType.uint64](), 1))
    mine = other if other > mine else mine
    parked[tid] = mine
    barrier()
    if tid == 0:
        var best = parked[0]
        for k in range(1, Int(block_dim.x)):
            best = parked[k] if parked[k] > best else best
        result[Int(block_idx.x)] = best


def main() raises:
    var ctx = DeviceContext()
    var hx = List[Float64](length=N, fill=0.0)
    var hy = List[Float64](length=N, fill=1.0)
    var hv = List[Float64](length=N, fill=0.0)
    var hout = List[Float64](length=N // 64, fill=0.0)
    for i in range(N):
        hx[i] = Float64(i)
        hv[i] = Float64((i * 37) % 101)
    var x = ctx.enqueue_create_buffer[DType.float64](N)
    var y = ctx.enqueue_create_buffer[DType.float64](N)
    var v = ctx.enqueue_create_buffer[DType.float64](N)
    var result = ctx.enqueue_create_buffer[DType.float64](N // 64)
    ctx.enqueue_copy(x, hx.unsafe_ptr())
    ctx.enqueue_copy(y, hy.unsafe_ptr())
    ctx.enqueue_copy(v, hv.unsafe_ptr())
    ctx.enqueue_function[saxpy](Int64(N), 2.0, x.unsafe_ptr(), y.unsafe_ptr(), grid_dim=(N + 255) // 256, block_dim=256)
    ctx.enqueue_function[block_max](v.unsafe_ptr(), result.unsafe_ptr(), grid_dim=N // 64, block_dim=64)
    ctx.enqueue_copy(hy.unsafe_ptr(), y)
    ctx.enqueue_copy(hout.unsafe_ptr(), result)
    ctx.synchronize()
    var total = 0.0
    for i in range(N):
        total += hy[i]
    var best = 0.0
    for i in range(N // 64):
        best = hout[i] if hout[i] > best else best
    print(total, best)
    var best_saxpy = 1e18
    var best_block = 1e18
    for _ in range(10):
        var started = perf_counter_ns()
        ctx.enqueue_function[saxpy](Int64(N), 2.0, x.unsafe_ptr(), y.unsafe_ptr(), grid_dim=(N + 255) // 256, block_dim=256)
        ctx.synchronize()
        var took = Float64(perf_counter_ns() - started) / 1e6
        best_saxpy = took if took < best_saxpy else best_saxpy
        started = perf_counter_ns()
        ctx.enqueue_function[block_max](v.unsafe_ptr(), result.unsafe_ptr(), grid_dim=N // 64, block_dim=64)
        ctx.synchronize()
        took = Float64(perf_counter_ns() - started) / 1e6
        best_block = took if took < best_block else best_block
    print("# saxpy:", best_saxpy, "ms")
    print("# block_max:", best_block, "ms")
    best_saxpy = 1e18
    best_block = 1e18
    for _ in range(10):
        var started = perf_counter_ns()
        ctx.enqueue_copy(x, hx.unsafe_ptr())
        ctx.enqueue_copy(y, hy.unsafe_ptr())
        ctx.enqueue_function[saxpy](Int64(N), 2.0, x.unsafe_ptr(), y.unsafe_ptr(), grid_dim=(N + 255) // 256, block_dim=256)
        ctx.enqueue_copy(hy.unsafe_ptr(), y)
        ctx.synchronize()
        var took = Float64(perf_counter_ns() - started) / 1e6
        best_saxpy = took if took < best_saxpy else best_saxpy
        started = perf_counter_ns()
        ctx.enqueue_copy(v, hv.unsafe_ptr())
        ctx.enqueue_function[block_max](v.unsafe_ptr(), result.unsafe_ptr(), grid_dim=N // 64, block_dim=64)
        ctx.enqueue_copy(hout.unsafe_ptr(), result)
        ctx.synchronize()
        took = Float64(perf_counter_ns() - started) / 1e6
        best_block = took if took < best_block else best_block
    print("# saxpy with copies:", best_saxpy, "ms")
    print("# block_max with copies:", best_block, "ms")
saxpy_cupy.py (Python) 
"""saxpy over sixteen million doubles and a per-block max: CuPy; timed warm, device-synchronized."""

import time

import cupy as cp

N = 1 << 24
block_max = cp.RawKernel(
    r"""
extern "C" __global__ void block_max(const double *x, double *out) {
    __shared__ double parked[64];
    int tid = threadIdx.x;
    parked[tid] = x[blockIdx.x * blockDim.x + threadIdx.x];
    __syncthreads();
    double mine = parked[tid];
    double other = __shfl_xor_sync(0xffffffffu, mine, 1);
    mine = other > mine ? other : mine;
    parked[tid] = mine;
    __syncthreads();
    if (tid == 0) {
        double best = parked[0];
        for (int k = 1; k < blockDim.x; k++) {
            best = parked[k] > best ? parked[k] : best;
        }
        out[blockIdx.x] = best;
    }
}
""",
    "block_max",
)
saxpy = cp.ElementwiseKernel("float64 a, float64 x, float64 y", "float64 z", "z = a * x + y", "saxpy")


def timed(label, run, repeats=10):
    run()
    cp.cuda.Device().synchronize()
    best = 1e9
    for _ in range(repeats):
        started = time.perf_counter()
        run()
        cp.cuda.Device().synchronize()
        best = min(best, (time.perf_counter() - started) * 1000.0)
    print(f"# {label}: {best:.3f} ms")


def main():
    x = cp.arange(N, dtype=cp.float64)
    y = cp.ones(N, dtype=cp.float64)
    values = cp.asarray([float((i * 37) % 101) for i in range(N)])
    out = cp.zeros(N // 64, dtype=cp.float64)
    saxpy(2.0, x, y, y)
    block_max((N // 64,), (64,), (values, out))
    print(float(y.sum()), float(out.max()))
    timed("saxpy", lambda: saxpy(2.0, x, y, y))
    timed("block_max", lambda: block_max((N // 64,), (64,), (values, out)))
    # The same launches the way a host-array launch means them: arrays copied in, results out.
    import numpy as np

    hx, hy, hv = x.get(), y.get(), values.get()

    def saxpy_copying():
        dx, dy = cp.asarray(hx), cp.asarray(hy)
        saxpy(2.0, dx, dy, dy)
        hy[:] = dy.get()

    def block_max_copying():
        dv = cp.asarray(hv)
        block_max((N // 64,), (64,), (dv, out))
        return out.get()

    timed("saxpy with copies", saxpy_copying)
    timed("block_max with copies", block_max_copying)


main()
saxpy_numba.py (Python) 
"""saxpy over sixteen million doubles and a per-block max: Numba's CUDA target; timed warm."""

import time

import numpy as np
from numba import cuda

N = 1 << 24


@cuda.jit(device=True)
def fma(a, x, y):
    return a * x + y


@cuda.jit
def saxpy(n, a, x, y):
    i = cuda.grid(1)
    if i < n:
        y[i] = fma(a, x[i], y[i])


@cuda.jit
def block_max(x, out):
    parked = cuda.shared.array(64, dtype=np.float64)
    tid = cuda.threadIdx.x
    parked[tid] = x[cuda.grid(1)]
    cuda.syncthreads()
    mine = parked[tid]
    other = cuda.shfl_xor_sync(0xFFFFFFFF, mine, 1)
    mine = other if other > mine else mine
    parked[tid] = mine
    cuda.syncthreads()
    if tid == 0:
        best = parked[0]
        for k in range(1, cuda.blockDim.x):
            best = parked[k] if parked[k] > best else best
        out[cuda.blockIdx.x] = best


def timed(label, run, repeats=10):
    run()
    cuda.synchronize()
    best = 1e9
    for _ in range(repeats):
        started = time.perf_counter()
        run()
        cuda.synchronize()
        best = min(best, (time.perf_counter() - started) * 1000.0)
    print(f"# {label}: {best:.3f} ms")


def main():
    x = cuda.to_device(np.arange(N, dtype=np.float64))
    y = cuda.to_device(np.ones(N, dtype=np.float64))
    values = cuda.to_device(np.array([float((i * 37) % 101) for i in range(N)]))
    out = cuda.device_array(N // 64, dtype=np.float64)
    saxpy[(N + 255) // 256, 256](N, 2.0, x, y)
    block_max[N // 64, 64](values, out)
    print(float(y.copy_to_host().sum()), float(out.copy_to_host().max()))
    timed("saxpy", lambda: saxpy[(N + 255) // 256, 256](N, 2.0, x, y))
    timed("block_max", lambda: block_max[N // 64, 64](values, out))
    hx, hy, hv = x.copy_to_host(), y.copy_to_host(), values.copy_to_host()

    def saxpy_copying():
        dx, dy = cuda.to_device(hx), cuda.to_device(hy)
        saxpy[(N + 255) // 256, 256](N, 2.0, dx, dy)
        dy.copy_to_host(hy)

    def block_max_copying():
        dv = cuda.to_device(hv)
        block_max[N // 64, 64](dv, out)
        return out.copy_to_host()

    timed("saxpy with copies", saxpy_copying)
    timed("block_max with copies", block_max_copying)


main()

Source: examples/38_cuda.