The current sparse encoding of shuffles blows up code size in graphics-heavy and audio-heavy code

[dead-claudia]

I'm resurrecting the shuffle part of WebAssembly/simd#255. (Broadcasts have *.splat.)

v128.shuffle instructions are always 18 bytes long, which is excessive in most cases. For comparison, native shuffle instructions on x86-64 are at most five bytes long, even though they are more restricted in shuffle patterns.

For an example here's an implementation of 4x4 matrix transpose in WebAssembly SIMD and in SSE intrinsics:

#include <wasm_simd128.h>

void transpose4x4(float matrix[16]) {
	const v128_t row0 = wasm_v128_load(matrix);
	const v128_t row1 = wasm_v128_load(matrix + 4);
	const v128_t row2 = wasm_v128_load(matrix + 8);
	const v128_t row3 = wasm_v128_load(matrix + 12);

	const v128_t tmp0 = wasm_v32x4_shuffle(row0, row1, 0, 4, 1, 5);
	const v128_t tmp1 = wasm_v32x4_shuffle(row2, row3, 0, 4, 1, 5);
	const v128_t tmp2 = wasm_v32x4_shuffle(row0, row1, 2, 6, 3, 7);
	const v128_t tmp3 = wasm_v32x4_shuffle(row2, row3, 2, 6, 3, 7);

	const v128_t col0 = wasm_v32x4_shuffle(tmp0, tmp1, 0, 1, 4, 5);
	const v128_t col1 = wasm_v32x4_shuffle(tmp0, tmp1, 2, 3, 6, 7);
	const v128_t col2 = wasm_v32x4_shuffle(tmp2, tmp3, 0, 1, 4, 5);
	const v128_t col3 = wasm_v32x4_shuffle(tmp2, tmp3, 2, 3, 6, 7);

	wasm_v128_store(matrix, col0);
	wasm_v128_store(matrix + 4, col1);
	wasm_v128_store(matrix + 8, col2);
	wasm_v128_store(matrix + 12, col3);
}

#include <xmmintrin.h>

void transpose(float matrix[16]) {
	const __m128 row0 = _mm_loadu_ps(matrix);
	const __m128 row1 = _mm_loadu_ps(matrix + 4);
	const __m128 row2 = _mm_loadu_ps(matrix + 8);
	const __m128 row3 = _mm_loadu_ps(matrix + 12);

	const __m128 tmp0 = _mm_unpacklo_ps(row0, row1);
	const __m128 tmp2 = _mm_unpacklo_ps(row2, row3);
	const __m128 tmp1 = _mm_unpackhi_ps(row0, row1);
	const __m128 tmp3 = _mm_unpackhi_ps(row2, row3);
	
	const __m128 col0 = _mm_movelh_ps(tmp0, tmp2);
	const __m128 col1 = _mm_movehl_ps(tmp2, tmp0);
	const __m128 col2 = _mm_movelh_ps(tmp1, tmp3);
	const __m128 col3 = _mm_movehl_ps(tmp3, tmp1);

	_mm_storeu_ps(matrix, col0);
	_mm_storeu_ps(matrix + 4, col1);
	_mm_storeu_ps(matrix + 8, col2);
	_mm_storeu_ps(matrix + 12, col3);
}

The x86-64 SSE version compiles to 67 bytes of code, but the WebAssembly version compiles to 231 bytes of code, 3.5x larger, according to the linked issue.

Admittedly, this is a bad example. That should be using contiguous loads and strides stores. But it works well enough for show.

I propose the following new opcodes:

• Permute:
  • i8x16.permute lanes:(byte)^8 | (v128) -> (v128), with 8-bit lane indices
  • i32x4.permute lanes:byte | (v128) -> (v128), with 2-bit lane indices
  • No i64x2.permute as it could be done in terms of this.
• Lane shifts:
  • i8x16.lane_shift_up count:8 | (v128) -> (v128), zero-fills from first lane
  • i8x16.lane_shift_down count:8 | (v128) -> (v128), zero-fills from last lane
• Compact shuffles:
  • i32x4.shuffle_pairs lanes:byte | (v128, v128) -> (v128), with 2-bit lane indices
  • i64x2.shuffle lanes:byte | (v128, v128) -> (v128), with 2-bit lane indices
  • i32x4.shuffle4 lanes:bytes^2 | (v128, v128, v128, v128) -> (v128), with four 4-bit lane indices
• Multi-byte swizzle:
  • i16x8.swizzle | (v128, v128, v128) -> (v128)
  • i32x4.swizzle | (v128, v128, v128) -> (v128)
  • i64x2.swizzle | (v128, v128, v128) -> (v128)

All of these are exceptionally common, common enough to merit opcodes to reduce generated code size.

• The permutes replace local.tee $v; local.get $v; v8x16.shuffle lanes... and local.get $v; local.get $v; v8x16.shuffle lanes... sequences.
• The compact shuffles and lane shifts reduce code size requirements for very common cases.
• Lane shuffles have somewhat broad support. Worst case, shift the value by lane size and add 0/1/2/3/... to resolve byte indices.
• 4x4 32-bit matrix transpose is extremely common. This can be fully represented in a single opcode.

And a few ISAs even have optimized sequences for these in particular.

• x86-64 has direct equivalents for most of this:
  • i32x4.permute -> shufps a, a, imm8
  • i8x16.lane_shift_up -> pslldq a, imm
  • i8x16.lane_shift_down -> psrldqa, imm
  • i32x4.shuffle_pairs -> shufps a, b, imm8
  • i64x2.shuffle -> shufpd a, b, imm8
  • i16x8.swizzle -> pshufw a, b
  • i32x4.swizzle -> pshufd a, b
  • i64x2.swizzle -> pshufq a, b
• 32-bit ARM and AArch64 both have efficient ways of most of them, too. Notably, it can do VTBL/TBL and VLD4/LD4 to implement i32x4.shuffle4 in two instructions.
• x86-64, 32-bit ARM, and AArch64 all can build a 128-bit vector from two scalar doubles in a single opcode.
• While RISC-V's V extension has no opcode for simple 32x4-bit mixes or shuffles (my biggest criticism of the V extension's design), every single shuffle and permute of 8 registers or less can be fused into a single vsetvli; vrgather.vv sequence with a constant lane selector operand for vrgather.vv, as permutes index multi-register group as a single contiguous array when LMUL>1. And even processors with just Zve64d can implement v32x4.shuffle4 efficiently by doing this.

[tlively]

Do you have a sense for how much code size savings we could get on the scale of a full Wasm module? I think the assumption has always been that SIMD instructions would be relatively uncommon, so the code bloat wouldn't be a problem in the big picture, but I'd be eager to see a library where that's not the case.

[dead-claudia]

Edit: add percentages

Do you have a sense for how much code size savings we could get on the scale of a full Wasm module? I think the assumption has always been that SIMD instructions would be relatively uncommon, so the code bloat wouldn't be a problem in the big picture, but I'd be eager to see a library where that's not the case.

I scrapped together a quick opcode scraping script and ran it on Tensorflow's SIMD WebAssembly bundle, a 424594-byte file. Of that, 888 18-byte i8x16.shuffle opcodes, occupying 15984 bytes (or about 3.8% of the file). Here's how my above proposal could save bytes on those opcodes and lead to a significant savings in bundle size.

• 744 could be represented with my proposed (likely 11-byte) i8x16.permute. Accounting for the saved dummy op, I'd expect 5952 bytes saved in the raw bundle, or about 1.4%.
• 743 could be represented with my proposed (likely 4-byte) i32x4.shuffle_pairs. I'd expect 10402 bytes saved in the raw bundle, or about 2.4%.
• 395 could be represented with my proposed (likely 4-byte) i64x2.shuffle. Accounting for the saved dummy op, I'd expect 5530 bytes saved in the raw bundle, or about 1.3%.
• 363 could be represented with my proposed (likely 4-byte) i32x4.permute. Accounting for the saved dummy op, I'd expect 5445 bytes saved in the raw bundle, or about 1.3%.
• 2 could be represented with my proposed (likely 4-byte) i8x16.lane_shift_up. Accounting for the saved dummy op, I'd expect 45 bytes saved in the raw bundle.
• 2 could be represented with my proposed (likely 4-byte) i8x16.lane_shift_down. Accounting for the saved dummy op, I'd expect 45 bytes saved in the raw bundle.

Due to the way I scraped the bundle, it being a one-pass scraper that doesn't maintain any stack, I don't have any ready stats for i32x4.shuffle4, i16x8.swizzle, i32x4.swizzle, or i64x2.swizzle. Those would require a lot deeper inspection. I also don't have stats on how these replacements would impact gzip (though I expect some significant savings due to more structural context being available in compression windows).

[tlively]

Awesome, thanks for the detailed analysis.

[dead-claudia]

Here's the script I cobbled together, if you're curious.