Flatten Multi-Dimensional Thread Index in gpu.js

This article explains how to convert multi-dimensional thread coordinates into a single 1D memory offset when writing GPU kernels using gpu.js. In GPU programming, tasks often run across 2D or 3D coordinate grids, but underlying memory buffers and arrays are typically stored linearly in 1D. You will learn the mathematical formulas and practical code patterns required to map 2D and 3D thread coordinates (this.thread.x, this.thread.y, and this.thread.z) to standard 1D linear array indices.

Understanding gpu.js Thread Indexing

When defining a kernel in gpu.js with multi-dimensional outputs, gpu.js provides built-in thread properties:

Because JavaScript TypedArrays and flat buffers are sequential 1D structures, each multi-dimensional thread must compute its exact position along that single axis.

Flattening a 2D Thread Index

In a standard row-major layout, the X coordinate moves fastest (along the row), and the Y coordinate steps from row to row.

Mathematical Formula

\[\text{Index}_{1D} = x + (y \times \text{width})\]

Implementation Example

Pass the row width directly into the kernel function or declare it within kernel constants.

const { GPU } = require('gpu.js');
const gpu = new GPU();

const width = 512;
const height = 512;

const kernel2D = gpu.createKernel(function(flatInput, width) {
  // Compute the 1D index from 2D coordinates
  const index = this.thread.x + (this.thread.y * width);

  // Access flat 1D memory
  return flatInput[index] * 2;
}).setOutput([width, height]);

Flattening a 3D Thread Index

For 3D data volumes, the index calculation extends to incorporate the depth slice (\(Z\)). Each step along the Z-axis spans an entire 2D slice of size \((\text{width} \times \text{height})\).

Mathematical Formula

\[\text{Index}_{1D} = x + (y \times \text{width}) + (z \times \text{width} \times \text{height})\]

Implementation Example

const { GPU } = require('gpu.js');
const gpu = new GPU();

const width = 64;
const height = 64;
const depth = 64;

const kernel3D = gpu.createKernel(function(flatInput, width, height) {
  // Compute the 1D index from 3D coordinates
  const sliceSize = width * height;
  const index = this.thread.x + (this.thread.y * width) + (this.thread.z * sliceSize);

  // Read or process 1D buffer data
  return flatInput[index] + 1.0;
}).setOutput([width, height, depth]);

Using Kernel Constants for Optimization

Instead of passing dimensions as dynamic arguments on every kernel execution, define them inside constants. This allows the GPU compiler to optimize index calculations with fixed numeric values:

const width = 1024;
const height = 768;

const fastKernel = gpu.createKernel(function(flatData) {
  const index = this.thread.x + (this.thread.y * this.constants.width);
  return flatData[index];
}, {
  constants: { width: width },
  output: [width, height]
});

Handling Multi-Channel Data (e.g., RGBA)

If the flat buffer contains multiple interleaved elements per thread (such as RGBA pixel values), multiply the final 1D index by the channel count:

const channels = 4; // R, G, B, A
const baseIndex = (this.thread.x + (this.thread.y * width)) * channels;

const r = flatPixels[baseIndex];
const g = flatPixels[baseIndex + 1];
const b = flatPixels[baseIndex + 2];
const a = flatPixels[baseIndex + 3];