GPU.js Non-Square and Irregular Thread Grids

This article explains how gpu.js executes compute kernels on non-square and irregular multi-dimensional thread grids. It details how the library maps arbitrary dimensions via its output configuration, translates multi-dimensional coordinates to low-level WebGL textures, and handles irregular or non-uniform data structures through rectangular bounding boxes and index flattening.

Defining Non-Square Dimensions

In gpu.js, thread grids are defined using the setOutput() method. The library natively supports non-square multi-dimensional rectangular grids by accepting an array of integers representing up to three dimensions: [width], [width, height], or [width, height, depth].

const gpu = new GPU();
const kernel = gpu.createKernel(function() {
    return this.thread.x + this.thread.y;
}).setOutput([500, 200]); // 500 threads wide (X), 200 threads tall (Y)

Inside the kernel function, the current execution coordinate is accessed through this.thread.x, this.thread.y, and this.thread.z. gpu.js maps these values strictly according to the bounds assigned in setOutput(), meaning non-square rectangular domains require no special mapping code from the developer.

WebGL Texture Mapping

Under the hood, gpu.js compiles JavaScript functions into WebGL fragment shaders. WebGL represents compute data using 2D textures. When a non-square or 3D thread grid is defined:

  1. 2D Non-Square Grids: gpu.js allocates an off-screen frame buffer and texture matching the exact width (\(X\)) and height (\(Y\)) of the output array, provided they do not exceed the hardware's MAX_TEXTURE_SIZE.
  2. 3D Grids: Because WebGL 1 and 2 fragment-shader compute targets are two-dimensional, gpu.js flattens 3D thread grids (\(X \times Y \times Z\)) into a tiled 2D texture layout. It automatically computes the mathematical offsets required to resolve this.thread.z relative to the 2D texture coordinates.

Handling Irregular or Ragged Grids

GPU hardware architectures and graphics pipelines require uniform Cartesian grids. gpu.js does not natively support "ragged" arrays (arrays where rows or columns have variable lengths). To compute over irregular or non-uniform shapes, you must employ specific architectural patterns:

1. The Bounding Box Strategy (Padding)

The standard approach involves allocating a rectangular output space that matches the maximum possible dimensions of the dataset, then using conditional logic to skip irrelevant cells.

const kernel = gpu.createKernel(function(rowLengths) {
    const currentRowLength = rowLengths[this.thread.y];
    
    // Check if the current thread falls outside the irregular shape
    if (this.thread.x >= currentRowLength) {
        return 0; // Padding value / early discard
    }
    
    return computeValue(this.thread.x, this.thread.y);
}).setOutput([maxWidth, totalRows]);

Threads outside the active data boundaries execute minimal logic and write dummy or sentinel values, which can later be trimmed on the CPU.

2. 1D Linear Flattening

For highly irregular structures (such as sparse graphs or triangular matrices) where bounding boxes cause excessive memory waste, data can be flattened into a single 1D thread array.

Execution Constraints and Divergence

When managing non-square or irregular grids in gpu.js, keep these performance factors in mind: