Tile kernels¶

These GPU kernels work on tiles rather than threads: you write what a program does to a tile, and the compiler writes the threads, the shared memory, and the shuffles. The primitives are few:

tile.arange(256) is the lane index.
tile.load gathers a tile, and tile.store scatters one back.
Arithmetic is lane by lane.
tile.sum and its kind reduce a tile to one number.

The compiler gives each program a block of threads that each hold a slice of every tile, and reduces across the block on its own.

Run it¶

python  tiles.ppy
ppy run tiles.ppy
ppy emit cuda tiles.ppy

What it prints¶

python tiles.ppy

1.0 1999.0
[100.0, 98.0, 100.0, 100.0]
49.8008 852.8236
# kernels compiled for a device here: False

ppy run tiles.ppy

1.0 1999.0
[100.0, 98.0, 100.0, 100.0]
49.8008 852.8236
# kernels compiled for a device here: True

ppy emit cuda tiles.ppy

386 lines

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

typedef struct ppy_v1i64 {
    int64_t lanes[1];
} ppy_v1i64;

typedef struct ppy_v1b {
    bool lanes[1];
} ppy_v1b;

typedef struct ppy_v1f64 {
    double lanes[1];
} ppy_v1f64;

typedef struct ppy_v4i64 {
    int64_t lanes[4];
} ppy_v4i64;

typedef struct ppy_v4b {
    bool lanes[4];
} ppy_v4b;

typedef struct ppy_v4f64 {
    double lanes[4];
} ppy_v4f64;

static inline __host__ __device__ ppy_v1i64 ppy_v1i64_splat(int64_t x) {
    ppy_v1i64 r;
    for (int i = 0; i < 1; i++) {
        r.lanes[i] = x;
    }
    return r;
}

static inline __host__ __device__ ppy_v1i64 ppy_v1i64_add(ppy_v1i64 a, ppy_v1i64 b) {
    ppy_v1i64 r;
    for (int i = 0; i < 1; i++) {
        r.lanes[i] = ((int64_t)(((uint64_t)(a.lanes[i])) + ((uint64_t)(b.lanes[i]))));
    }
    return r;
}

static inline __host__ __device__ ppy_v1b ppy_v1i64_cmp_lt(ppy_v1i64 a, ppy_v1i64 b) {
    ppy_v1b r;
    for (int i = 0; i < 1; i++) {
        r.lanes[i] = a.lanes[i] < b.lanes[i];
    }
    return r;
}

static inline __host__ __device__ ppy_v1f64 ppy_v1f64_splat(double x) {
    ppy_v1f64 r;
    for (int i = 0; i < 1; i++) {
        r.lanes[i] = x;
    }
    return r;
}

static inline __host__ __device__ ppy_v1f64 ppy_v1f64_insert(ppy_v1f64 v, double x, int64_t i) {
    v.lanes[i] = x;
    return v;
}

static inline __host__ __device__ ppy_v1f64 ppy_v1f64_mul(ppy_v1f64 a, ppy_v1f64 b) {
    ppy_v1f64 r;
    for (int i = 0; i < 1; i++) {
        r.lanes[i] = (a.lanes[i] * b.lanes[i]);
    }
    return r;
}

static inline __host__ __device__ ppy_v1f64 ppy_v1f64_add(ppy_v1f64 a, ppy_v1f64 b) {
    ppy_v1f64 r;
    for (int i = 0; i < 1; i++) {
        r.lanes[i] = (a.lanes[i] + b.lanes[i]);
    }
    return r;
}

static inline __host__ __device__ double ppy_v1f64_reduce_max(ppy_v1f64 v) {
    double acc = v.lanes[0];
    for (int i = 1; i < 1; i++) {
        acc = v.lanes[i] > acc ? v.lanes[i] : acc;
    }
    return acc;
}

static inline __host__ __device__ ppy_v4i64 ppy_v4i64_splat(int64_t x) {
    ppy_v4i64 r;
    for (int i = 0; i < 4; i++) {
        r.lanes[i] = x;
    }
    return r;
}

static inline __host__ __device__ ppy_v4i64 ppy_v4i64_insert(ppy_v4i64 v, int64_t x, int64_t i) {
    v.lanes[i] = x;
    return v;
}

static inline __host__ __device__ ppy_v4i64 ppy_v4i64_add(ppy_v4i64 a, ppy_v4i64 b) {
    ppy_v4i64 r;
    for (int i = 0; i < 4; i++) {
        r.lanes[i] = ((int64_t)(((uint64_t)(a.lanes[i])) + ((uint64_t)(b.lanes[i]))));
    }
    return r;
}

static inline __host__ __device__ ppy_v4b ppy_v4i64_cmp_lt(ppy_v4i64 a, ppy_v4i64 b) {
    ppy_v4b r;
    for (int i = 0; i < 4; i++) {
        r.lanes[i] = a.lanes[i] < b.lanes[i];
    }
    return r;
}

static inline __host__ __device__ ppy_v4f64 ppy_v4f64_splat(double x) {
    ppy_v4f64 r;
    for (int i = 0; i < 4; i++) {
        r.lanes[i] = x;
    }
    return r;
}

static inline __host__ __device__ ppy_v4f64 ppy_v4f64_insert(ppy_v4f64 v, double x, int64_t i) {
    v.lanes[i] = x;
    return v;
}

static inline __host__ __device__ double ppy_v4f64_reduce_add(ppy_v4f64 v) {
    double acc = v.lanes[0];
    for (int i = 1; i < 4; i++) {
        acc = (acc + v.lanes[i]);
    }
    return acc;
}

static inline __host__ __device__ int64_t ppy_v4i64_reduce_add(ppy_v4i64 v) {
    int64_t acc = v.lanes[0];
    for (int i = 1; i < 4; i++) {
        acc = ((int64_t)(((uint64_t)(acc)) + ((uint64_t)(v.lanes[i]))));
    }
    return acc;
}

static inline __host__ __device__ ppy_v4f64 ppy_v4f64_sub(ppy_v4f64 a, ppy_v4f64 b) {
    ppy_v4f64 r;
    for (int i = 0; i < 4; i++) {
        r.lanes[i] = (a.lanes[i] - b.lanes[i]);
    }
    return r;
}

static inline __host__ __device__ ppy_v4f64 ppy_v4f64_mul(ppy_v4f64 a, ppy_v4f64 b) {
    ppy_v4f64 r;
    for (int i = 0; i < 4; i++) {
        r.lanes[i] = (a.lanes[i] * b.lanes[i]);
    }
    return r;
}

static inline __host__ __device__ ppy_v4f64 ppy_v4f64_select(ppy_v4b m, ppy_v4f64 a, ppy_v4f64 b) {
    ppy_v4f64 r;
    for (int i = 0; i < 4; i++) {
        r.lanes[i] = m.lanes[i] ? a.lanes[i] : b.lanes[i];
    }
    return r;
}

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" {

__global__ void ppy_tiles_saxpy(int64_t n, double a, const double *x, double *y);
__global__ void ppy_tiles_block_max(const double *x, double *out);
__global__ void ppy_tiles_row_stats(int64_t n, const double *x, double *out);
int32_t ppy_tiles_run(int64_t n, double a, const double *x, double *y, int64_t *out);

__global__ void ppy_tiles_saxpy(int64_t n, double a, const double *x, double *y) {
    int64_t t1 = (int64_t)blockIdx.x;
    int64_t t2 = (int64_t)threadIdx.x;
    ppy_v1i64 t3 = ppy_v1i64_splat(t2);
    ppy_v1i64 t4 = ppy_v1i64_splat(int64_t(uint64_t(t1) * 256u));
    ppy_v1i64 offsets = ppy_v1i64_add(t4, t3);
    ppy_v1i64 t5 = ppy_v1i64_splat(n);
    ppy_v1b mask = ppy_v1i64_cmp_lt(offsets, t5);
    ppy_v1f64 t6 = ppy_v1f64_splat(0.0);
    int64_t t7 = offsets.lanes[0];
    bool t8 = mask.lanes[0];
    ppy_v1f64 xs = ppy_v1f64_insert(t6, (t8 ? x[t8 ? t7 : 0] : 0.0), 0);
    ppy_v1f64 t9 = ppy_v1f64_splat(0.0);
    int64_t t10 = offsets.lanes[0];
    bool t11 = mask.lanes[0];
    ppy_v1f64 ys = ppy_v1f64_insert(t9, (t11 ? y[t11 ? t10 : 0] : 0.0), 0);
    ppy_v1f64 t12 = ppy_v1f64_splat(a);
    ppy_v1f64 t13 = ppy_v1f64_mul(t12, xs);
    ppy_v1f64 t14 = ppy_v1f64_add(t13, ys);
    int64_t t15 = offsets.lanes[0];
    double *t16 = y + t15;
    double t17 = t14.lanes[0];
    bool t18 = mask.lanes[0];
    if (t18) {
        *t16 = t17;
    }
}

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

    int64_t pid = (int64_t)blockIdx.x;
    int64_t t1 = (int64_t)threadIdx.x;
    ppy_v1i64 t2 = ppy_v1i64_splat(t1);
    ppy_v1i64 t3 = ppy_v1i64_splat(int64_t(uint64_t(pid) * 64u));
    ppy_v1i64 t4 = ppy_v1i64_add(t3, t2);
    ppy_v1f64 t5 = ppy_v1f64_splat(0.0);
    int64_t t6 = t4.lanes[0];
    ppy_v1f64 values = ppy_v1f64_insert(t5, x[t6], 0);
    double t7 = ppy_v1f64_reduce_max(values);
    double t8 = __shfl_xor_sync(0xffffffffu, t7, (int)16);
    double t9 = t7 > t8 ? t7 : t8;
    double t10 = __shfl_xor_sync(0xffffffffu, t9, (int)8);
    double t11 = t9 > t10 ? t9 : t10;
    double t12 = __shfl_xor_sync(0xffffffffu, t11, (int)4);
    double t13 = t11 > t12 ? t11 : t12;
    double t14 = __shfl_xor_sync(0xffffffffu, t13, (int)2);
    double t15 = t13 > t14 ? t13 : t14;
    double t16 = __shfl_xor_sync(0xffffffffu, t15, (int)1);
    double t17 = t15 > t16 ? t15 : t16;
    int64_t t18 = (int64_t)threadIdx.x;
    int64_t t19 = t18 / 32 - (t18 % 32 < 0);
    __syncthreads();
    if ((t18 % 32 + 32) % 32 == 0) {
        shared[t19] = t17;
    }
    __syncthreads();
    double t20 = *shared;
    double t21 = shared[1];
    double t22 = t20 > t21 ? t20 : t21;
    double *t23 = out + pid;
    int64_t t24 = (int64_t)threadIdx.x;
    if (t24 == 0) {
        *t23 = t22;
    }
}

__global__ void ppy_tiles_row_stats(int64_t n, const double *x, double *out) {
    __shared__ double shared[8];
    __shared__ int64_t shared_2[8];
    __shared__ double shared_3[8];

    int64_t pid = (int64_t)blockIdx.x;
    int64_t t1 = (int64_t)threadIdx.x;
    ppy_v4i64 t2 = ppy_v4i64_splat(t1);
    ppy_v4i64 t3 = ppy_v4i64_splat(0);
    ppy_v4i64 t4 = ppy_v4i64_insert(t3, 256, 1);
    ppy_v4i64 t5 = ppy_v4i64_insert(t4, 512, 2);
    ppy_v4i64 t6 = ppy_v4i64_insert(t5, 768, 3);
    ppy_v4i64 t7 = ppy_v4i64_add(t2, t6);
    ppy_v4i64 t8 = ppy_v4i64_splat(int64_t(uint64_t(pid) * 1024u));
    ppy_v4i64 offsets = ppy_v4i64_add(t8, t7);
    ppy_v4i64 t9 = ppy_v4i64_splat(n);
    ppy_v4b mask = ppy_v4i64_cmp_lt(offsets, t9);
    ppy_v4f64 t10 = ppy_v4f64_splat(0.0);
    int64_t t11 = offsets.lanes[0];
    bool t12 = mask.lanes[0];
    ppy_v4f64 t13 = ppy_v4f64_insert(t10, (t12 ? x[t12 ? t11 : 0] : 0.0), 0);
    int64_t t14 = offsets.lanes[1];
    bool t15 = mask.lanes[1];
    ppy_v4f64 t16 = ppy_v4f64_insert(t13, (t15 ? x[t15 ? t14 : 0] : 0.0), 1);
    int64_t t17 = offsets.lanes[2];
    bool t18 = mask.lanes[2];
    ppy_v4f64 t19 = ppy_v4f64_insert(t16, (t18 ? x[t18 ? t17 : 0] : 0.0), 2);
    int64_t t20 = offsets.lanes[3];
    bool t21 = mask.lanes[3];
    ppy_v4f64 xs = ppy_v4f64_insert(t19, (t21 ? x[t21 ? t20 : 0] : 0.0), 3);
    double t22 = ppy_v4f64_reduce_add(xs);
    double t23 = __shfl_xor_sync(0xffffffffu, t22, (int)16);
    double t24 = t22 + t23;
    double t25 = __shfl_xor_sync(0xffffffffu, t24, (int)8);
    double t26 = t24 + t25;
    double t27 = __shfl_xor_sync(0xffffffffu, t26, (int)4);
    double t28 = t26 + t27;
    double t29 = __shfl_xor_sync(0xffffffffu, t28, (int)2);
    double t30 = t28 + t29;
    double t31 = __shfl_xor_sync(0xffffffffu, t30, (int)1);
    double t32 = t30 + t31;
    int64_t t33 = (int64_t)threadIdx.x;
    int64_t t34 = t33 / 32 - (t33 % 32 < 0);
    __syncthreads();
    if ((t33 % 32 + 32) % 32 == 0) {
        shared[t34] = t32;
    }
    __syncthreads();
    double total = *shared + shared[1] + shared[2] + shared[3] + shared[4] + shared[5] + shared[6] + shared[7];
    ppy_v4i64 t35 = ppy_v4i64_splat(0);
    bool t36 = mask.lanes[0];
    ppy_v4i64 t37 = ppy_v4i64_insert(t35, (t36 ? 1 : 0), 0);
    bool t38 = mask.lanes[1];
    ppy_v4i64 t39 = ppy_v4i64_insert(t37, (t38 ? 1 : 0), 1);
    bool t40 = mask.lanes[2];
    ppy_v4i64 t41 = ppy_v4i64_insert(t39, (t40 ? 1 : 0), 2);
    bool t42 = mask.lanes[3];
    ppy_v4i64 t43 = ppy_v4i64_insert(t41, (t42 ? 1 : 0), 3);
    int64_t t44 = ppy_v4i64_reduce_add(t43);
    int64_t t45 = __shfl_xor_sync(0xffffffffu, t44, (int)16);
    int64_t t46 = int64_t(uint64_t(t44) + uint64_t(t45));
    int64_t t47 = __shfl_xor_sync(0xffffffffu, t46, (int)8);
    int64_t t48 = int64_t(uint64_t(t46) + uint64_t(t47));
    int64_t t49 = __shfl_xor_sync(0xffffffffu, t48, (int)4);
    int64_t t50 = int64_t(uint64_t(t48) + uint64_t(t49));
    int64_t t51 = __shfl_xor_sync(0xffffffffu, t50, (int)2);
    int64_t t52 = int64_t(uint64_t(t50) + uint64_t(t51));
    int64_t t53 = __shfl_xor_sync(0xffffffffu, t52, (int)1);
    int64_t t54 = int64_t(uint64_t(t52) + uint64_t(t53));
    int64_t t55 = (int64_t)threadIdx.x;
    int64_t t56 = t55 / 32 - (t55 % 32 < 0);
    __syncthreads();
    if ((t55 % 32 + 32) % 32 == 0) {
        shared_2[t56] = t54;
    }
    __syncthreads();
    int64_t count = int64_t(uint64_t(*shared_2) + uint64_t(shared_2[1]) + uint64_t(shared_2[2]) + uint64_t(shared_2[3]) + uint64_t(shared_2[4]) + uint64_t(shared_2[5]) + uint64_t(shared_2[6]) + uint64_t(shared_2[7]));
    double mean = total / double(count);
    ppy_v4f64 t57 = ppy_v4f64_splat(mean);
    ppy_v4f64 t58 = ppy_v4f64_sub(xs, t57);
    ppy_v4f64 t59 = ppy_v4f64_splat(mean);
    ppy_v4f64 t60 = ppy_v4f64_sub(xs, t59);
    ppy_v4f64 t61 = ppy_v4f64_mul(t58, t60);
    ppy_v4f64 t62 = ppy_v4f64_splat(0.0);
    ppy_v4f64 t63 = ppy_v4f64_select(mask, t61, t62);
    double t64 = ppy_v4f64_reduce_add(t63);
    double t65 = __shfl_xor_sync(0xffffffffu, t64, (int)16);
    double t66 = t64 + t65;
    double t67 = __shfl_xor_sync(0xffffffffu, t66, (int)8);
    double t68 = t66 + t67;
    double t69 = __shfl_xor_sync(0xffffffffu, t68, (int)4);
    double t70 = t68 + t69;
    double t71 = __shfl_xor_sync(0xffffffffu, t70, (int)2);
    double t72 = t70 + t71;
    double t73 = __shfl_xor_sync(0xffffffffu, t72, (int)1);
    double t74 = t72 + t73;
    int64_t t75 = (int64_t)threadIdx.x;
    int64_t t76 = t75 / 32 - (t75 % 32 < 0);
    __syncthreads();
    if ((t75 % 32 + 32) % 32 == 0) {
        shared_3[t76] = t74;
    }
    __syncthreads();
    double spread = *shared_3 + shared_3[1] + shared_3[2] + shared_3[3] + shared_3[4] + shared_3[5] + shared_3[6] + shared_3[7];
    double *t77 = out + int64_t(uint64_t(pid) * 2u);
    int64_t t78 = (int64_t)threadIdx.x;
    if (t78 == 0) {
        *t77 = mean;
    }
    double t79 = spread / double(count);
    double *t80 = out + int64_t(uint64_t(pid) * 2u + 1u);
    int64_t t81 = (int64_t)threadIdx.x;
    if (t81 == 0) {
        *t80 = t79;
    }
}

int32_t ppy_tiles_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_tiles_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" */

A program owns a tile¶

@tile.kernel
def saxpy(n: int, a: float, x: native.const_ptr[float], y: native.ptr[float]) -> None:
    offsets = tile.program_id() * 256 + tile.arange(256)
    mask = offsets < n
    xs = tile.load(x, offsets, mask)
    ys = tile.load(y, offsets, mask)
    tile.store(y, offsets, a * xs + ys, mask)

tile.launch(saxpy, (n + 255) // 256, n, a, x, y) runs one program per 256 elements.

offsets < n is a tile of bools. A masked load reads nothing off the end and gives 0 (or the other you pass) in those lanes. A masked store writes nothing there.
A scalar beside a tile is broadcast: a * xs scales every lane.

The same file runs on CPython, where a launch runs the programs one after another and a tile is a list per lane. That is the reference the device is held to.

A reduction is one call¶

@tile.kernel
def block_max(x: native.const_ptr[float], out: native.ptr[float]) -> None:
    pid = tile.program_id()
    values = tile.load(x, pid * 64 + tile.arange(64))
    tile.store(out, pid, tile.max(values))

tile.max(values) is the block's max in one call. Natively, the program is a block of threads, as many as the tile has lanes, up to 256 (so a 1024 tile is four lanes per thread). Each thread reduces its own lanes, a shuffle tree reduces the warp, the warps meet in shared memory, and every thread ends with the same number. tile.store(out, pid, ...) of one element is written once.

row_stats does a mean and a variance the same way, with tile.where(mask, ..., 0.0) keeping the padding out of the sum. What the CUDA example spells with cuda.shared, cuda.syncthreads, and cuda.shfl_xor is here what the compiler writes.

What a tile kernel may hold¶

A tile has int, float, or bool lanes.
One kernel has one block size, a power of two of at least 32, named by its tile.arange.
The operators are + - * / // %, the comparisons, & | ^ ~ on ints and bools, tile.where, and tile.sum, tile.max, tile.min.
E1644 names a misuse.

ppy emit cuda writes a tile kernel as CUDA C++ like any other, and ppy emit ptx as PTX.

Compared with Triton and Taichi¶

The same two kernels over sixteen million doubles, written as programs over tiles in the three tools that have them. They are in compare/: tiles_bench.ppy, tiles_triton.py, tiles_taichi.py. Times are in milliseconds, best of warm launches with the device synchronized, over five processes. The thread-level ports of the same kernels (CuPy, Numba, Mojo, CUDA C) are in the CUDA example's comparison.

Here is the block max as each tool spells it.

PPy: a program loads its tile and reduces it. The same file runs on CPython.

@tile.kernel
def block_max(x: native.const_ptr[float], out: native.ptr[float]) -> None:
    pid = tile.program_id()
    values = tile.load(x, pid * 64 + tile.arange(64))
    tile.store(out, pid, tile.max(values))

Triton: the same shape, with tl.program_id, tl.arange, tl.load, tl.max. It is launched over CuPy memory through a six-line data_ptr() wrapper, and its compiler tiles the work across a warp group.

@triton.jit
def block_max(x_ptr, out_ptr, BLOCK: tl.constexpr):
    offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    tl.store(out_ptr + tl.program_id(0), tl.max(tl.load(x_ptr + offsets), axis=0))

Taichi: no program either. Arrays are ti.fields, the outer loop of a kernel is parallel over the blocks, and the max of each block is an inner serial loop. ti.init(arch=ti.cuda, default_fp=ti.f64) makes the answers match, and on WSL the driver goes on the library path by hand.

@ti.kernel
def block_max():
    for b in out:
        best = values[b * 64]
        for k in range(1, 64):
            candidate = values[b * 64 + k]
            best = candidate if candidate > best else best
        out[b] = best

	PPy `tile.launch`	Triton	Taichi
saxpy, arrays on the device	0.72 ± 0.09	0.77 ± 0.10	0.89 ± 0.11
block max, arrays on the device	0.41 ± 0.02	0.81 ± 0.02	0.31 ± 0.03
saxpy, arrays copied in and out per launch	41.11 ± 0.67	51.73 ± 6.92	163.66 ± 7.49
block max, array copied in per launch	14.87 ± 0.40	11.38 ± 0.36	29.90 ± 2.21

On saxpy, the three run at the memory bandwidth of the device.

On the block max they are three lowerings of one idea:

PPy gives the 64-lane tile a block of 64 threads and reduces it with a shuffle tree.
Triton gives it a warp group and its own reduction.
Taichi gives each block one thread that loops over the 64, with no exchange at all. That is why it is the fastest here, and why it would not be on a wider tile.

The copying rows are the other memory model: an array sent in and brought back on every launch. PPy's driver reads the host array in place, Triton goes through CuPy's asarray, and Taichi's from_numpy copies twice.

NVIDIA GeForce RTX 5080 Laptop GPU, driver 610.71, CUDA 13.3; Triton 3.8.0, Taichi 1.7.4 on CPython 3.12.13; PPy on CPython 3.13.13.

`44_tile/tiles.ppy`¶

from ppy import native, tile


@tile.kernel
def saxpy(n: int, a: float, x: native.const_ptr[float], y: native.ptr[float]) -> None:
    offsets = tile.program_id() * 256 + tile.arange(256)
    mask = offsets < n
    xs = tile.load(x, offsets, mask)
    ys = tile.load(y, offsets, mask)
    tile.store(y, offsets, a * xs + ys, mask)


@tile.kernel
def block_max(x: native.const_ptr[float], out: native.ptr[float]) -> None:
    pid = tile.program_id()
    values = tile.load(x, pid * 64 + tile.arange(64))
    tile.store(out, pid, tile.max(values))


@tile.kernel
def row_stats(n: int, x: native.const_ptr[float], out: native.ptr[float]) -> None:
    pid = tile.program_id()
    offsets = pid * 1024 + tile.arange(1024)
    mask = offsets < n
    xs = tile.load(x, offsets, mask)
    total = tile.sum(xs)
    count = tile.sum(mask)
    mean = total / count
    spread = tile.sum(tile.where(mask, (xs - mean) * (xs - mean), 0.0))
    tile.store(out, pid * 2, mean)
    tile.store(out, pid * 2 + 1, spread / count)


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


def main() -> None:
    n = 1000
    x = native.stack_alloc[float](n)
    y = native.stack_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)
    print(native.load(native.offset(y, 0)), native.load(native.offset(y, 999)))
    values = native.stack_alloc[float](256)
    out = native.stack_alloc[float](4)
    for i in range(256):
        native.store(native.offset(values, i), float((i * 37) % 101))
    tile.launch(block_max, 4, values, out)
    print([native.load(native.offset(out, i)) for i in range(4)])
    stats = native.stack_alloc[float](2)
    tile.launch(row_stats, 1, 256, values, stats)
    print(f"{native.load(stats):.4f} {native.load(native.offset(stats, 1)):.4f}")
    print(f"# kernels compiled for a device here: {tile.compiled(saxpy)}")


main()

Counterpart programs¶

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

tiles_bench.ppy (PPy)

"""saxpy over sixteen million doubles and a per-block max as tile kernels: PPY over device
memory and over host arrays."""

import time

from ppy import cuda, native, tile

N = 1 << 24


@tile.kernel
def saxpy(n: int, a: float, x: native.const_ptr[float], y: native.ptr[float]) -> None:
    offsets = tile.program_id() * 256 + tile.arange(256)
    mask = offsets < n
    xs = tile.load(x, offsets, mask)
    ys = tile.load(y, offsets, mask)
    tile.store(y, offsets, a * xs + ys, mask)


@tile.kernel
def block_max(x: native.const_ptr[float], out: native.ptr[float]) -> None:
    pid = tile.program_id()
    values = tile.load(x, pid * 64 + tile.arange(64))
    tile.store(out, pid, tile.max(values))


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


def run_block_max(blocks: int, x: native.const_ptr[float], out: native.ptr[float]) -> None:
    tile.launch(block_max, blocks, 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: {tile.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")
    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()

tiles_taichi.py (Python)

"""saxpy over sixteen million doubles and a per-block max: Taichi on its CUDA backend."""

import contextlib
import os
import sys
import time

import numpy as np

os.environ.setdefault("LD_LIBRARY_PATH", "/usr/lib/wsl/lib")  # where WSL keeps libcuda
with contextlib.redirect_stdout(sys.stderr):  # the banners are not answers
    import taichi as ti

    ti.init(arch=ti.cuda, default_fp=ti.f64, default_ip=ti.i64)

N = 1 << 24
x = ti.field(ti.f64, shape=N)
y = ti.field(ti.f64, shape=N)
values = ti.field(ti.f64, shape=N)
out = ti.field(ti.f64, shape=N // 64)


@ti.func
def fma(a, x, y):
    return a * x + y


@ti.kernel
def saxpy(a: ti.f64):
    for i in x:
        y[i] = fma(a, x[i], y[i])


@ti.kernel
def block_max():
    # Taichi has no thread, block, or shared memory to name: the outer loop is
    # parallel over blocks, and the max of each block is an inner serial loop.
    for b in out:
        best = values[b * 64]
        for k in range(1, 64):
            candidate = values[b * 64 + k]
            best = candidate if candidate > best else best
        out[b] = best


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


def main():
    hx = np.arange(N, dtype=np.float64)
    hy = np.ones(N, dtype=np.float64)
    hv = np.array([float((i * 37) % 101) for i in range(N)])
    x.from_numpy(hx)
    y.from_numpy(hy)
    values.from_numpy(hv)
    saxpy(2.0)
    block_max()
    print(float(y.to_numpy().sum()), float(out.to_numpy().max()))
    timed("saxpy", lambda: saxpy(2.0))
    timed("block_max", block_max)
    hy = y.to_numpy()

    def saxpy_copying():
        x.from_numpy(hx)
        y.from_numpy(hy)
        saxpy(2.0)
        hy[:] = y.to_numpy()

    def block_max_copying():
        values.from_numpy(hv)
        block_max()
        return out.to_numpy()

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


main()

tiles_triton.py (Python)

"""saxpy over sixteen million doubles and a per-block max: Triton programs over CuPy memory."""

import time

import cupy as cp
import triton
import triton.language as tl

N = 1 << 24


class Pointer:
    """What a Triton launch needs of an array: its device pointer and its dtype."""

    def __init__(self, array):
        self.array = array
        self.dtype = array.dtype

    def data_ptr(self):
        return self.array.data.ptr


@triton.jit
def fma(a, x, y):
    return a * x + y


@triton.jit
def saxpy(n, a, x_ptr, y_ptr, BLOCK: tl.constexpr):
    offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    mask = offsets < n
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(y_ptr + offsets, fma(a, x, y), mask=mask)


@triton.jit
def block_max(x_ptr, out_ptr, BLOCK: tl.constexpr):
    # One program per block of 64: the block's values are a vector, and the
    # max is a reduction over it -- there is no thread, no shared memory, no shuffle.
    offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    values = tl.load(x_ptr + offsets)
    tl.store(out_ptr + tl.program_id(0), tl.max(values, axis=0))


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 run_saxpy(x, y):
    saxpy[((N + 255) // 256,)](N, 2.0, Pointer(x), Pointer(y), BLOCK=256)


def run_block_max(values, out):
    block_max[(N // 64,)](Pointer(values), Pointer(out), BLOCK=64)


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)
    run_saxpy(x, y)
    run_block_max(values, out)
    print(float(y.sum()), float(out.max()))
    timed("saxpy", lambda: run_saxpy(x, y))
    timed("block_max", lambda: run_block_max(values, out))
    hx, hy, hv = x.get(), y.get(), values.get()

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

    def block_max_copying():
        dv = cp.asarray(hv)
        run_block_max(dv, out)
        return out.get()

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


main()

Source: examples/44_tile.