Zig Density Plot

Note: This was originally published as an Observable notebook.

Experimenting with Zig for faster visualizations. This notebook contains a prototype implementation of a simple one-dimensional density plot.

See the Zig code ( lines) 
const std = @import("std");
const math = std.math;

// Struct storing colors for generating canvas image data
const Rgba = packed struct { r: u8, g: u8, b: u8, a: u8 };

// Accumulate density from points into bins. Bins are not zeroed before accumulation.
export fn density1d(bins: [*]f64, lenBins: usize, xs: [*]f64, lenPoints: usize) void {
  // Compute min and max
  var min = math.inf(f64);
  var max = -math.inf(f64);
  for (xs[0..lenPoints]) |x| {
      min = math.min(x, min);
      max = math.max(x, max);
  }
  // handle equal endpoints
  if (max == min) max += 1;
  // Compute bins with our min and max
  density1dCore(bins[0..lenBins], xs[0..lenPoints], min, max, true);
}

// Compute the binning with a presupplied min and max, e.g. to zoom into a particular sub-range of data values.
export fn density1dMinMax(bins: [*]f64, lenBins: usize, xs: [*]f64, lenPoints: usize, min: f64, max: f64) void {
  density1dCore(bins[0..lenBins], xs[0..lenPoints], min, max, false);
}

// Core density processing
fn density1dCore(bins: []f64, xs: []f64, min: f64, max: f64, comptime allPointsKnownInBounds: bool) void {
  const scale = binScale(bins.len, min, max);
  for (xs) |x| {
      // Ensure out-of-bounds points don't contribute any weight.
      // When called with [min, max] computed from the data, we can 
      // avoid this check at comptime since all points are in bounds.
      if (allPointsKnownInBounds or (min <= x and x <= max)) {
          bins[binIndex(x, min, scale)] += 1;
      }
  }
}

// Compute canvas image data, given bins and a linear color ramp.
// Nonlinear color ramps can be done by pre-transforming the bin values.
export fn binsToImgData(imgData: [*]Rgba, bins: [*]f64, len: usize, ramp: [*]Rgba, rampLen: usize) void {
  binsToImgDataCore(imgData[0..len], bins[0..len], ramp[0..rampLen]);
}

// Core image data processing
fn binsToImgDataCore(imgData: []Rgba, bins: []f64, ramp: []Rgba) void {
  // For now, map [0, max] to the color range.
  const min = 0;
  var max = -math.inf(f64);
  for (bins) |value| {
      max = math.max(value, max);
  }
  // handle equal endpoints
  if (max == min) max += 1;
  // Compute a color for each bin by looking up the appropriate ramp value
  const scale = binScale(ramp.len, min, max);
  for (bins) |value, i| {
      imgData[i] = ramp[binIndex(value, min, scale)];
  }
}

// Helper functions

// Scale factor from [min, max] to [0..bins.len - eps]. The epsilon ensures
// that maximum point gets mapped to the largest bin (instead of beyond it)
fn binScale(len: usize, min: f64, max: f64) f64 {
  // This eps value correctly rounds f64 values when there are less than 32768^2 bins.
  const eps64 = 1e-7; 
  return (@intToFloat(f64, len) - eps64) / (max - min);
}

// Given a value, the min extent, and a bin scale factor, compute the bin index.
fn binIndex(value: f64, min: f64, scale: f64) usize {
  const pc = (value - min) * scale;
  return @floatToInt(usize, pc);
}

The core function computes a histogram with uniform bins, which is a visualization primitive for bar charts, density plots, and contours.

So far, the Zig version outperforms all JavaScript variants I've seen, including my own.

For an additional performance boost the code also computes bin colors with a user-supplied color ramp, populating a byte buffer that can be copied directly into canvas image data.

I'm a week into learning Zig, so please let me know if you spot a bug or think the code can be improved.

Things I Learned

You can manage WASM memory from JS. I initially used Zig allocators to manage WASM memory, but the code became simpler and faster when I allowed JavaScript to manage heap memory using a simple arena allocator instead.

My first implementation also had a bug where a const global value stored in WASM memory was being overwritten by my JS-side allocations, and I eventually realized that I need to avoid writing to the memory pages that were instantiated by Zig on load, and allocate only within additional pages. This amounted to setting the "floor" of the bump allocator to the initial memory size (see alloc below).

Exported function signatures can use struct types. I was initially using @ptrCast, @alignCast, and @alignOf to cast between [*]u8 and [*]Rgba pointer types, and was not confident that my code was handling alignment correctly.

Then I learned that you can use just directly use the target type in the signature of an exported function, which expresses the intent much more clearly:

const Rgba = packed struct { r: u8, g: u8, b: u8, a: u8 };

export fn binsToImgData(imgData: [*]Rgba) void { ... }

Zig binaries are small and fast even when compiled in safe mode. The .wasm file for this demo is just 1,316 bytes when compiled with Zig 0.9.1 and the following command:

$ zig build-lib density.zig -target wasm32-freestanding -dynamic -OReleaseSafe

Compiling with ReleaseFast produces a 1,564-byte file, while compiling with ReleaseSmall produces an even smaller 739-byte file.

References

Thanks to Fil for comments on an earlier draft of this post.

Implementation

const density = (data, xBins, colorRamp, min, max) => {
  // Start the timer. Faster with the console closed, by the way – at least in Chrome.
  const timeLabel = `${comma(numPoints)} data points`;
  console.time(timeLabel);

  // Free all pre-existing WASM-side allocations
  alloc.reset();

  // Allocate a buffer for bins and initialize it to zero.
  const bins = alloc(Float64Array, xBins).fill(0);

  // Allocate a buffer for the data and initialize it.
  const xs = alloc(Float64Array, data.length);
  xs.set(data);

  if (min !== undefined && max !== undefined) {
    // Bin the points, counting the number of points per bin in the range (min, max).
    // prettier-ignore
    zig.density1dMinMax(bins.byteOffset, bins.length, xs.byteOffset, xs.length, min, max );
  } else {
    // Bin the points, computing the min and max from the data instead:
    zig.density1d(bins.byteOffset, bins.length, xs.byteOffset, xs.length);
  }

  // Allocate a buffer for what will become the canvas image data.
  const imgData = alloc(Uint8ClampedArray, 4 * bins.length);

  // Allocate a buffer for the color ramp and initialize it.
  const ramp = alloc(Uint8ClampedArray, colorRamp.length);
  ramp.set(colorRamp);

  // Compute the image data based on the bins and linear ramp.
  zig.binsToImgData(
    imgData.byteOffset,
    bins.byteOffset,
    bins.length,
    ramp.byteOffset,
    ramp.length / 4
  );

  // Transfer the data to the image and update the canvas.
  image.data.set(imgData);
  ctx.putImageData(image, 0, 0);

  // Provide the raw bins as the value of this view.
  // Note that these are in WASM memory.
  ctx.canvas.value = bins;

  console.timeEnd(timeLabel);
  return ctx.canvas;
}

  const request = fetch(await FileAttachment("density@1.wasm").url());
  const prom = WebAssembly.instantiateStreaming
    ? WebAssembly.instantiateStreaming(request)
    : request
        .then((res) => res.arrayBuffer())
        .then((bytes) => WebAssembly.instantiate(bytes));
  const results = await prom;
    // Ensure the WASM memory is large enough to store all of the data.
    const instance = results.instance;
    const memory = instance.exports.memory;
    const bytesPerPage = 65536;
    const numPages = memory.grow(0);
    const initialBytes = numPages * bytesPerPage;
    const desiredBytes =
      Float64Array.BYTES_PER_ELEMENT * 1e6 * (maxMillions + 1);
    const additionalPages = Math.ceil(
      (desiredBytes - initialBytes) / bytesPerPage
    );
    memory.grow(additionalPages);
    instance.initialBytes = initialBytes;

const zig = instance.exports

const dataset = Float64Array.from({ length: maxMillions * 1e6 }, d3.randomNormal())

const data = dataset.subarray(0, numPoints)

// canvas

  const canvas = DOM_canvas(xBins, 1, 1);
  const ctx = canvas.getContext("2d");
  canvas.style.width = `${width}px`;
  canvas.style.height = `${70}px`;
  canvas.style.imageRendering = "pixelated";
  ctx.imageSmoothingEnabled = false;

const image = ctx.getImageData(0, 0, xBins, 1)

// utilities

const rampLen = 256;
const ramp = new Uint8Array(4 * rampLen);
const scale = interpolator;
for (let i = 0; i < rampLen; i++) {
  const { r, g, b } = d3.rgb(scale(i / (rampLen - 1)));
  ramp[4 * i] = r;
  ramp[4 * i + 1] = g;
  ramp[4 * i + 2] = b;
  ramp[4 * i + 3] = 255;
}
const colorRamp = ramp;

// A simple memory allocator that allocates from the WebAssembly memory buffer.
  // Pointer to the address of the next allocation.
  // To prevent overwriting memory we don't control, we never allocate in
  // the initial memory pages loaded by the runtime.
  let ptr = instance.initialBytes;
  const alloc = Object.assign(
    (T, len) => {
      // Align allocations to an 8-byte boundary.
      if (ptr % 8 > 0) ptr = ptr + 8 - (ptr % 8);
      const numBytes = T.BYTES_PER_ELEMENT * len;
      const buffer = new T(instance.exports.memory.buffer, ptr, len);
      ptr += numBytes;
      return buffer;
    },
    {
      reset: () => (ptr = instance.initialBytes)
    }
  );

const comma = d3.format(",")

const numPoints = 1e6 * millions

const maxMillions = 10

const wasmMemoryMb = zig.memory.buffer.byteLength / (1024 * 1024)

const details = (summary, contents, isOpen) => {
  return html`<details ${isOpen ? "open" : ""}>
    ${contents}
    <summary>${summary}</summary>
  </details>`;
}

Note: The computed bins can be used with Plot: