Rust compiled to WASM slower than equivalent JS

25 points 101arrowz

I'm trying to rewrite some JS image processing code in Rust from a project where I'm trying to create a document scanning app with zero dependencies (article series in my profile if you're interested). I was hoping WebAssembly would provide a performance boost, and while WASM is ahead for small images (under 400x400 or so), somehow JS is way faster for large images on both desktop and mobile.

There are two components to this code: grayscale and downsampling. For simplicity's sake, I've removed some logic that makes the downsampling correct (right now it leaves a thin black border on the edges) but the hot loop, which is the main performance concern, is still there.

Here's the downsampling code in TypeScript:

const downscale = (src: Float32Array, width: number, height: number, by: number, dst?: Float32Array) => {
  const dw = Math.floor(width / by);
  const dh = Math.floor(height / by);
  if (!dst) {
    dst = new Float32Array(dh * dw);
  }
  const by2 = by * by;
  const right = 1, below = dw, diag = right + below;
  for (let i = 0; i < height; ++i) {
    for (let j = 0; j < width; ++j) {
      const val = src[i * width + j] / by2;
      const si = i / by;
      const sii = Math.floor(si);
      const sirr = si - sii;
      const sir = 1.0 - sirr;
      const sj = j / by;
      const sji = Math.floor(sj)
      const sjrr = sj - sji
      const sjr = 1.0 - sjrr;
      const di = sii * dw + sji;
      dst[di] += val * sir * sjr;
      dst[di + right] += val * sirr * sjr;
      dst[di + below] += val * sir * sjrr;
      dst[di + diag] += val * sirr * sjrr;
    }
  }
  return dst;
}

And here it is in Rust:

pub struct Image {
    pub data: Vec<f32>,
    pub width: usize,
    pub height: usize
}

pub fn downscale(source: &Image, by: f32) -> Image {
    assert!(by > 1.0);
    let width = (source.width as f32 / by) as usize;
    let height = (source.height as f32 / by) as usize;
    let by2 = by * by;
    let mut data = vec![0.0; width * height];
    let right = 1;
    let below = width;
    let diag = right + below;
    let sub = (by * 2.0) as usize;
    let mw = source.width - sub;
    let mh = source.height - sub;
    for i in 0..=mw {
        for j in 0..=mh {
            // Forgot to use get_unchecked here but I doubt it would've helped
            let val = source.data[j * source.width + i] / by2;
            let si = (i as f32) / by;
            let sii = si as usize;
            let sirr = si.fract();
            let sir = 1.0 - sirr;
            let sj = (j as f32) / by;
            let sji = sj as usize;
            let sjrr = sj.fract();
            let sjr = 1.0 - sjrr;
            let di = sji * width + sii;
            unsafe {
                *data.get_unchecked_mut(di) += val * sir * sjr;
                *data.get_unchecked_mut(di + right) += val * sirr * sjr;
                *data.get_unchecked_mut(di + below) += val * sir * sjrr;
                *data.get_unchecked_mut(di + diag) += val * sirr * sjrr;
            }
        }
    }
    Image { data, width, height }
}

The JS version is substantially faster.

I managed to get the grayscaling code around the same in performance for both, but honestly I was hoping for better from Rust. Here's the JS:

// RGBA -> grayscale (intensity in range 0-1)
// I know this reuses the buffer for the destination array by default,
// but I made sure to use a freshly allocated array for dst when calling it.
const grayscale = (src: Uint8ClampedArray, dst = new Float32Array(src.buffer, src.byteOffset, src.byteLength >> 2)) => {
  for (let px = 0; px < dst.length; ++px) {
    const pos = px << 2;
    dst[px] = src[pos] * 0.00116796875 + src[pos + 1] * 0.00229296875 + src[pos + 2] * 0.0004453125;
  }
  return dst;
}

And here's the Rust:

pub struct RGBAImage {
    pub data: Vec<u8>,
    pub width: usize,
    pub height: usize
}

pub fn grayscale(source: &RGBAImage) -> Image {
    let &RGBAImage { data: ref source, width, height } = source;
    Image {
        data: source.chunks_exact(4).map(|rgba| unsafe {
            (*rgba.get_unchecked(0) as f32) * 0.00116796875 +
            (*rgba.get_unchecked(1) as f32) * 0.00229296875 +
            (*rgba.get_unchecked(2) as f32) * 0.0004453125
        }).collect(),
        width,
        height
    }
}

I'm running with -O3, release mode, and using wasm-opt. I'm not familiar with SIMD but when I tried to enable it, autovectorization did not occur for any of the hot loops.

By the way, I'm comparing performance without sending the grayscaled, downsampled image from WASM back to JS to avoid incurring the marshalling overhead, but it's still slower. When I do send the image back, I get another substantial performance penalty on top of everything else.

This performance deficit is especially frustrating because I never even tried to optimize the JS but I've had to go through many iterations on the Rust code to get the performance even close, and I've had to use unsafe. Admittedly I'm a much better JS programmer than Rust, but I would expect terrible Rust code to thrash fine-tuned JS.

Any suggestions on optimizing?