Wunk

vpternlog: Signed Saturation

Thu, 04 Dec 2025 23:34:50 +0000

vpternlog{d,q} is one of those instructions that a time traveler should have told Intel to put into the original x86 ISA years ago once you understand what it does.

vpternlog{d,q} is an AVX512F (and ideally with AVX512VL) instruction that implements ternary-logic across its three operands, overwriting one of its inputs.

Note that this is a bit different from the understanding of ternary-logic that operates on trits, but it is ternary in the sense that it is doing three-operand logic much like the ternary conditional operator.

An intuition to have about this instruction is that it will take one lane of bits from each of its three operands to create a 3-bit value, which further indexes a singular bit within an 8-bit look up table value alongside the instruction. This 8-bit look-up operation has the ability to generalize many different three-operand binary operations such as combining several XOR, AND, NOT and OR operations into a singular instruction.

The instruction itself is bitwise across the entire register, though the d(32-bit) and q(64-bit) specifiers of vpternlog{d,q} provide additional element masking with a mask-register:

Synopsis
__m128i _mm_ternarylogic_epi32 (__m128i a, __m128i b, __m128i c, int imm8)
#include <immintrin.h>
Instruction: vpternlogd xmm, xmm, xmm, imm8
CPUID Flags: AVX512F + AVX512VL
Description
Bitwise ternary logic that provides the capability to implement any three-operand binary function; the specific binary function is specified by value in imm8. For each bit in each packed 32-bit integer, the corresponding bit from a, b, and c are used according to imm8, and the result is written to the corresponding bit in dst.

DEFINE TernaryOP(imm8, a, b, c) {
	CASE imm8[7:0] OF
	0: dst[0] := 0                   // imm8[7:0] := 0
	1: dst[0] := NOT (a OR b OR c)   // imm8[7:0] := NOT (_MM_TERNLOG_A OR _MM_TERNLOG_B OR _MM_TERNLOG_C)
	// ...
	254: dst[0] := a OR b OR c       // imm8[7:0] := _MM_TERNLOG_A OR _MM_TERNLOG_B OR _MM_TERNLOG_C
	255: dst[0] := 1                 // imm8[7:0] := 1
	ESAC
}
imm8[7:0] = LogicExp(_MM_TERNLOG_A, _MM_TERNLOG_B, _MM_TERNLOG_C)
FOR j := 0 to 3
	i := j*32
	FOR h := 0 to 31
		dst[i+h] := TernaryOP(imm8[7:0], a[i+h], b[i+h], c[i+h])
	ENDFOR
ENDFOR
dst[MAX:128] := 0

Intel intrinsics guide vers. 3.6.9

Based on https://www.sandpile.org/x86/opc_3.htm, there was apparently an intent to implement vpternlogb and vpternlogw variants for 16-bit and 8-bit elements. sandpile.org also provides a helpful table of ternary logic operations with their matching vpternlog LUT values: https://www.sandpile.org/x86/ternlog.htm

Nvidia’s PTX instruction set implements a similar instruction called LOP3.LUT for their GPUs.

Intel’s Virtual Instruction Set Architecture, used for their GPUs, implements a similar instruction called BFN

Compile time constants

Rather than spending the time to fumble up the look up table bits yourself, you can actually express these look-up bits directly in high-level code and generate the LUT bits at compile time!

GCC and Clang both provide the enum variables _MM_TERNLOG_{A,B,C} in avx512fintrin.h (included by immintrin.h) which allow you to directly express your binary operation in such a way that directly generates the 8-bit LUT value:

const std::uint8_t TernLut = (_MM_TERNLOG_A | ~_MM_TERNLOG_B) & _MM_TERNLOG_C;
//						   = 0b10100010
//						   = 0xa2

// later...
... = _mm_ternarylogic_epi32(a, b, c, TernLut);
// vpternlogd a, b, c, 0xa2

const std::uint8_t TernLut = ~(_MM_TERNLOG_A ^ _MM_TERNLOG_B) & _MM_TERNLOG_C;
//    					   = 0b10000010
//    					   = 0x82

// later...
... = _mm_ternarylogic_epi32(a, b, c, TernLut);
// vpternlogd a, b, c, 0x82

MSVC does not currently provide these constants, though. You are probably better off implementing these constants yourself for proper portability across compilers. It’s just three integer constants!

I did something like this for both Dynarmic and Xenia when emitting the vpternlog{d,q} instructions for their x64 JIT backends:

namespace Tern {
constexpr std::uint8_t a = 0b11110000;
constexpr std::uint8_t b = 0b11001100;
constexpr std::uint8_t c = 0b10101010;
}  // namespace Tern

With these constants, you can express your ternary operation directly within your C or C++ code, and the resulting compile-time constant will be the appropriate 8-bit look up table value for vpternlog{d,q} to do the same operation upon its operands!

//// This implements "(a | ~b) & c" as a single instruction!
const std::uint8_t TernLut = (Tern::a | ~Tern::b) & Tern::c;
//		    			   = 0b10100010
//		    			   = 0xa2

... = _mm_ternarylogic_epi32(a, b, c, TernLut);
// vpternlogd a, b, c, 0xa2

//// This implements "~(a ^ b) & c" as a single instruction!
const std::uint8_t TernLut = ~(Tern::a ^ Tern::b) & Tern::c;
//    					   = 0b10000010
//    					   = 0x82

... = _mm_ternarylogic_epi32(a, b, c, TernLut);
// vpternlogd a, b, c, 0x82

Intel’s uses a similar mechanism of compile-time constants when explicitly emitting the bfn instruction with SYCL:

https://www.intel.com/content/www/us/en/docs/oneapi/optimization-guide-gpu/2024-2/optimizing-explicit-simd-kernels.html#MATH-FUNCTIONS

Signed Saturated arithmetic

Saturated arithmetic is pretty commonplace in tasks relating to image processing, digital signal processing, fixed point arithmetic, safety critical practices, control algorithms, and anywhere else where having predictable underflow/overflow results is better than modulo arithmetic’s wrap-around results.

Wrapping vs. Saturation on an underflowing/overflowing input signal

Alright so where are the SIMD instructions…

On ARM, there are NEON instructions for SIMD vector addition and subtraction with the {s,u}qadd and {s,u}qsub instructions. It supports signed and unsigned 8 bit, 16 bit, 32 bit, and 64 bit elements!

Learn the architecture - Neon programmers’ guide: C.7. List of saturating instructions

On RISC-V, saturated vector operations are also supported by the V vector extensions. Specifically, the vsadd{,u} and vssub{,u} instructions:

RISC-V V vector extension specification 1.0: Vector Single-Width Saturating Add and Subtract

How about on x64?

There is 8-bit and 16-bit saturation, but no 32-bit or 64-bit saturation! With the absence of native support for 32 and 64-bit saturated SIMD instructions, we are left to implement it ourselves.

There actually seemed to be a intent to release an AVX512_BITALG2 instruction set extension that adds the VPADDSD, VPADDSQ, VPADDUSD and VPADDUSQ instructions. Notably, there was not a subtraction (SUBS) version in what is known about AVX512_BITALG2’s intended design.

Section EVEX + F3h: https://www.sandpile.org/x86/opc_3.htm

Saturated arithmetic depends on doing a regular integer addition/subtraction and then determining if saturation needs to occur based on the sign bit (the most significant bit) of both the input operands and the intermediate addition/subtraction result.

Using vpternlog{d,q}, it can be concisely determined if a saturation needs to occur by utilizing the sign bits of each SIMD vector’s elements. The rest of the bits of each element will be along for the ride. The resulting signed saturation values can be either an overflow (0x7FFFFF...) or underflow (0x80000...).

In the case of addition:

Signed saturation occurs if…

The sign bits of the two inputs is the same.

The sign bit of resulting sum has a different sign from the inputs.

Signed Saturated Add: Saturation happens if...
a = Operand A sign bit
b = Operand B sign bit
c = Sum sign bit

Interpreted as binary operations:
 "Input sign bits are the same"                   = ~(a ^ b)
 "result has a different sign bit from the input" =  (a ^ c)

Saturation needed  = ~(a ^ b) & (a ^ c)

TernlogLut = (~(_MM_TERNLOG_A ^ _MM_TERNLOG_B)) & (_MM_TERNLOG_A ^ _MM_TERNLOG_C)

A(`OpA` sign bit)	B(`OpB` sign bit)	C(`Addition` sign bit)	LUT(Saturation is needed)
0	0	0	0
0	0	1	1 (If Pos + Pos = Neg, a signed-overflow occurred)
0	1	0	0
0	1	1	0
1	0	0	0
1	0	1	0
1	1	0	1 (If Neg + Neg = Pos, a signed-underflow occurred)
1	1	1	0

In the case of subtraction:

Signed saturation occurs if…

The sign bits of the two inputs are different.

The resulting subtraction has a different sign from the first operand.

// Signed Saturated Subtract saturation happens if...
a = Operand A sign
b = Operand B sign
c = Subtraction sign bit

Interpreted as binary operations:
 "Input signs are different"                          = (a ^ b)
 "Result has a different sign from the first operand" = (a ^ c)

Saturation needed = (a ^ b) & (a ^ c)

TernlogLut = (_MM_TERNLOG_A ^ _MM_TERNLOG_B) & (_MM_TERNLOG_A ^ _MM_TERNLOG_C)

A(`OpA` sign bit)	B(`OpB` sign bit)	C(`Subtraction` sign bit)	LUT(Saturation is needed)
0	0	0	0
0	0	1	0
0	1	0	0
0	1	1	1 (If Pos - Neg = Neg, a signed-overflow occurred)
1	0	0	1 (If Neg - Pos = Pos, a signed-underflow occurred)
1	0	1	0
1	1	0	0
1	1	1	0

This results in the vpternlog{d,q} constants 0b01000010 for detecting the need to saturate lanes during addition and 0b00011000 for detecting the need to saturate lanes during subtraction.

// Detect if a signed overflow/underflow has occurred by using ternary
// logic on the sign bit of each element, other bits can be ignored.

// Addition
... = _mm_ternarylogic_epi32(OpA, OpB, Addition,    0b01000010);
// vpternlogd OpA, OpB, Addition, 0b01000010

// Subtraction:
... = _mm_ternarylogic_epi32(OpA, OpB, Subtraction, 0b00011000);
// vpternlogd OpA, OpB, Subtraction, 0b00011000

With AVX512, this value can then directly be converted into an execution-mask by using _mm_movepi32_mask(vpmovd2m), which will convert the most significant bit (the sign bit) of each element into a mask-register. With this write-mask of where underflows/overflows have occurred, we now need to write the correct overflow/underflow values into these lanes.

Either an underflow (0x80000000) or overflow (0x7FFFFFFF) value is written depending on the sign bit of the sum. This can be achieved in several ways depending on how much you want to touch memory or what execution ports you want to use among neighboring work. I’ll show one way that attempts to emit a small amount of instructions that tries to avoid additional masking or reading constants from memory:

The sign bit of the resulting addition/subtraction we calculated before can be repurposed directly to create the desired overflow/underflow values. Each lane will have either a 0xxxxxx... or a 1xxxxxx... and these values can be turned into either an 0x80000000(under) or an 0x7FFFFFFF(overflow) respectively.

We do this with an arithmetic shift to the right by 31 bits to “spread” the sign bit to all other bits. Then, we can flip the sign bit again to get either a 0x80000000/0x7FFFFFFF value.

Arithmetic Shift the Addition/Subtraction value to the right by 31 bits
- Creates either a 0x00000000 or 0xFFFFFFFF value in each lane
Invert the sign bit
- Creates either 0x80000000 or 0x7FFFFFFF value in each lane

Both saturated addition and subtraction implementations are done similarly, the only difference is the ternary look up table constant used and the intermediate result uses either a 32 bit addition or subtraction with _mm_{add,sub}_epi32. The 64-bit version is implemented the same as well, with careful attention to the data-types of the masks and arithmetic operations and constants.

Anyways you’re probably here to copy-paste some code:

// #include <immintrin.h>

__m128i _mm_adds_epi32(__m128i a, __m128i b)
{
	__m128i Result = _mm_add_epi32(a, b);

	// Resolves to 0b01000010(0x42)
	const std::uint32_t Lut = 0b01000010;
		// Not supported on MSVC
		// = (~(_MM_TERNLOG_A ^ _MM_TERNLOG_B)) // OpA and OpB are the same
		// & (_MM_TERNLOG_A ^ _MM_TERNLOG_C);   // OpA and OpC are different

	// Each element's sign bit is `1` if an overflow/underflow has occurred.
	// Other bits are just along for the ride
	const __m128i  SaturationCheck = _mm_ternarylogic_epi32(a, b, Result, Lut);
	const __mmask8 SaturationMask  = _mm_movepi32_mask(SaturationCheck);

	// "Spread" the sign bit to the rest of the bits, creating either
	// a `0x00000000` or a `0xFFFFFFFF`
	Result = _mm_mask_srai_epi32(Result, SaturationMask, Result, 31);

	// Flip the top-most bit creating either a `0x80000000` or a `0x7FFFFFFF`
	Result = _mm_mask_xor_epi32(
		Result, SaturationMask, Result, _mm_set1_epi32(0x80000000u)
	);

	return Result;
}

__m128i _mm_subs_epi32(__m128i a, __m128i b)
{
	__m128i Result = _mm_sub_epi32(a, b);

	// Resolves to 0b00011000(0x18)
	const std::uint32_t Lut = 0b00011000;
		// Not supported on MSVC
		// = (_MM_TERNLOG_A ^ _MM_TERNLOG_B)  // OpA and OpB are different
		// & (_MM_TERNLOG_A ^ _MM_TERNLOG_C); // OpA and OpC are different

	// Each element's sign bit is `1` if an overflow/underflow has occurred.
	// Other bits are just along for the ride
	const __m128i  SaturationCheck = _mm_ternarylogic_epi32(a, b, Result, Lut);
	const __mmask8 SaturationMask  = _mm_movepi32_mask(SaturationCheck);

	// "Spread" the sign bit to the rest of the bits, creating either
	// a `0x00000000` or a `0xFFFFFFFF`
	Result = _mm_mask_srai_epi32(Result, SaturationMask, Result, 31);

	// Flip the top-most bit creating either a `0x80000000` or a `0x7FFFFFFF`
	Result = _mm_mask_xor_epi32(
		Result, SaturationMask, Result, _mm_set1_epi32(0x80000000u)
	);

	return Result;
}

__m128i _mm_adds_epi64(__m128i a, __m128i b)
{
	// Same general idea as _mm_adds_epi32
	__m128i Result = _mm_add_epi64(a, b);

	const __m128i  SaturationCheck = _mm_ternarylogic_epi64(a, b, Result, 0b01000010);
	const __mmask8 SaturationMask  = _mm_movepi64_mask(SaturationCheck);

	Result = _mm_mask_srai_epi64(Result, SaturationMask, Result, 63);
	Result = _mm_mask_xor_epi64(
		Result, SaturationMask, Result, _mm_set1_epi64x(0x8000000000000000)
	);

	return Result;
}

__m128i _mm_subs_epi64(__m128i a, __m128i b)
{
	// Same general idea as _mm_subs_epi32
	__m128i Result = _mm_sub_epi64(a, b);

	const __m128i  SaturationCheck = _mm_ternarylogic_epi64(a, b, Result, 0b00011000);
	const __mmask8 SaturationMask  = _mm_movepi64_mask(SaturationCheck);

	Result = _mm_mask_srai_epi64(Result, SaturationMask, Result, 63);
	Result = _mm_mask_xor_epi64(
		Result, SaturationMask, Result, _mm_set1_epi64x(0x8000000000000000)
	);

	return Result;
}

GCC seems to emit the best assembly for the 32-bit case. GCC generates the 0x80000000-constant within the SIMD registers while clang and MSVC involve reaching out to memory for the 0x80000000-constant. For the 64-bit case, all compilers will read the constant from memory. Depending on where this operation lands, you might want to use a different method to generate these constants.

g++ 16.0.0 20251203

-march=skylake-avx512 -O2

"_mm_adds_epi32(long long __vector(2), long long __vector(2))":
        vpaddd  xmm2, xmm0, xmm1
        vpternlogd      xmm0, xmm1, xmm2, 66
        vpmovd2m        k1, xmm0
        vpcmpeqd        xmm0, xmm0, xmm0
        vpslld  xmm0, xmm0, 31
        vpsrad  xmm2{k1}, xmm2, 31
        vpxord  xmm2{k1}, xmm2, xmm0
        vmovdqa xmm0, xmm2
        ret
"_mm_subs_epi32(long long __vector(2), long long __vector(2))":
        vpsubd  xmm2, xmm0, xmm1
        vpternlogd      xmm0, xmm1, xmm2, 24
        vpmovd2m        k1, xmm0
        vpcmpeqd        xmm0, xmm0, xmm0
        vpslld  xmm0, xmm0, 31
        vpsrad  xmm2{k1}, xmm2, 31
        vpxord  xmm2{k1}, xmm2, xmm0
        vmovdqa xmm0, xmm2
        ret
"_mm_adds_epi64(long long __vector(2), long long __vector(2))":
        vpaddq  xmm2, xmm0, xmm1
        vpternlogq      xmm0, xmm1, xmm2, 66
        vpmovq2m        k1, xmm0
        vpsraq  xmm2{k1}, xmm2, 63
        vpxorq  xmm2{k1}, xmm2, XMMWORD PTR .LC1[rip]
        vmovdqa xmm0, xmm2
        ret
"_mm_subs_epi64(long long __vector(2), long long __vector(2))":
        vpsubq  xmm2, xmm0, xmm1
        vpternlogq      xmm0, xmm1, xmm2, 24
        vpmovq2m        k1, xmm0
        vpsraq  xmm2{k1}, xmm2, 63
        vpxorq  xmm2{k1}, xmm2, XMMWORD PTR .LC1[rip]
        vmovdqa xmm0, xmm2
        ret
.LC1:
        .quad   -9223372036854775808
        .quad   -9223372036854775808

clang version 22.0.0

-march=skylake-avx512 -O2

.LCPI0_0:
        .long   2147483647
.LCPI0_1:
        .long   2147483648
_mm_adds_epi32(long long vector[2], long long vector[2]):
        vpaddd  xmm2, xmm1, xmm0
        vpternlogd      xmm0, xmm1, xmm2, 66
        vpmovd2m        k1, xmm0
        vpcmpeqd        xmm0, xmm0, xmm0
        vpcmpgtd        k2, xmm2, xmm0
        vpbroadcastd    xmm0, dword ptr [rip + .LCPI0_0]
        vpbroadcastd    xmm0 {k2}, dword ptr [rip + .LCPI0_1]
        vmovdqa32       xmm2 {k1}, xmm0
        vmovdqa xmm0, xmm2
        ret

.LCPI1_0:
        .long   2147483647
.LCPI1_1:
        .long   2147483648
_mm_subs_epi32(long long vector[2], long long vector[2]):
        vpsubd  xmm2, xmm0, xmm1
        vpternlogd      xmm0, xmm1, xmm2, 24
        vpmovd2m        k1, xmm0
        vpcmpeqd        xmm0, xmm0, xmm0
        vpcmpgtd        k2, xmm2, xmm0
        vpbroadcastd    xmm0, dword ptr [rip + .LCPI1_0]
        vpbroadcastd    xmm0 {k2}, dword ptr [rip + .LCPI1_1]
        vmovdqa32       xmm2 {k1}, xmm0
        vmovdqa xmm0, xmm2
        ret

.LCPI2_0:
        .quad   9223372036854775807
.LCPI2_1:
        .quad   -9223372036854775808
_mm_adds_epi64(long long vector[2], long long vector[2]):
        vpaddq  xmm2, xmm1, xmm0
        vpternlogq      xmm0, xmm1, xmm2, 66
        vpmovq2m        k1, xmm0
        vpcmpeqd        xmm0, xmm0, xmm0
        vpcmpgtq        k2, xmm2, xmm0
        vpbroadcastq    xmm0, qword ptr [rip + .LCPI2_0]
        vpbroadcastq    xmm0 {k2}, qword ptr [rip + .LCPI2_1]
        vmovdqa64       xmm2 {k1}, xmm0
        vmovdqa xmm0, xmm2
        ret

.LCPI3_0:
        .quad   9223372036854775807
.LCPI3_1:
        .quad   -9223372036854775808
_mm_subs_epi64(long long vector[2], long long vector[2]):
        vpsubq  xmm2, xmm0, xmm1
        vpternlogq      xmm0, xmm1, xmm2, 24
        vpmovq2m        k1, xmm0
        vpcmpeqd        xmm0, xmm0, xmm0
        vpcmpgtq        k2, xmm2, xmm0
        vpbroadcastq    xmm0, qword ptr [rip + .LCPI3_0]
        vpbroadcastq    xmm0 {k2}, qword ptr [rip + .LCPI3_1]
        vmovdqa64       xmm2 {k1}, xmm0
        vmovdqa xmm0, xmm2
        ret

MSVC 19.44.35219

/arch:AVX512 /Ox

__xmm@80000000800000008000000080000000 DB 00H, 00H, 00H, 080H, 00H, 00H, 00H
        DB      080H, 00H, 00H, 00H, 080H, 00H, 00H, 00H, 080H
__xmm@80000000000000008000000000000000 DB 00H, 00H, 00H, 00H, 00H, 00H, 00H
        DB      080H, 00H, 00H, 00H, 00H, 00H, 00H, 00H, 080H

a$ = 8
b$ = 16
__m128i _mm_adds_epi32(__m128i,__m128i) PROC          ; _mm_adds_epi32
        vmovdqu xmm1, XMMWORD PTR [rdx]
        vmovdqu xmm2, XMMWORD PTR [rcx]
        vpaddd  xmm0, xmm1, xmm2
        vpternlogd xmm2, xmm1, xmm0, 66       ; 00000042H
        vpmovd2m k1, xmm2
        vpsrad  xmm0 {k1}, xmm0, 31
        vpxord  xmm0 {k1}, xmm0, XMMWORD PTR __xmm@80000000800000008000000080000000
        ret     0
__m128i _mm_adds_epi32(__m128i,__m128i) ENDP          ; _mm_adds_epi32

a$ = 8
b$ = 16
__m128i _mm_subs_epi32(__m128i,__m128i) PROC          ; _mm_subs_epi32
        vmovdqu xmm1, XMMWORD PTR [rdx]
        vmovdqu xmm2, XMMWORD PTR [rcx]
        vpsubd  xmm0, xmm2, xmm1
        vpternlogd xmm2, xmm1, xmm0, 24
        vpmovd2m k1, xmm2
        vpsrad  xmm0 {k1}, xmm0, 31
        vpxord  xmm0 {k1}, xmm0, XMMWORD PTR __xmm@80000000800000008000000080000000
        ret     0
__m128i _mm_subs_epi32(__m128i,__m128i) ENDP          ; _mm_subs_epi32

a$ = 8
b$ = 16
__m128i _mm_adds_epi64(__m128i,__m128i) PROC          ; _mm_adds_epi64
        vmovdqu xmm0, XMMWORD PTR [rdx]
        vmovdqu xmm1, XMMWORD PTR [rcx]
        vpaddq  xmm2, xmm0, xmm1
        vpternlogq xmm1, xmm0, xmm2, 66       ; 00000042H
        vpmovq2m k1, xmm1
        mov     eax, 63                             ; 0000003fH
        vmovd   xmm1, eax
        vpsraq  xmm2 {k1}, xmm2, xmm1
        vpxorq  xmm2 {k1}, xmm2, XMMWORD PTR __xmm@80000000000000008000000000000000
        vmovdqu xmm0, xmm2
        ret     0
__m128i _mm_adds_epi64(__m128i,__m128i) ENDP          ; _mm_adds_epi64

a$ = 8
b$ = 16
__m128i _mm_subs_epi64(__m128i,__m128i) PROC          ; _mm_subs_epi64
        vmovdqu xmm0, XMMWORD PTR [rdx]
        vmovdqu xmm1, XMMWORD PTR [rcx]
        vpsubq  xmm2, xmm1, xmm0
        vpternlogq xmm1, xmm0, xmm2, 24
        vpmovq2m k1, xmm1
        mov     eax, 63                             ; 0000003fH
        vmovd   xmm1, eax
        vpsraq  xmm2 {k1}, xmm2, xmm1
        vpxorq  xmm2 {k1}, xmm2, XMMWORD PTR __xmm@80000000000000008000000000000000
        vmovdqu xmm0, xmm2
        ret     0
__m128i _mm_subs_epi64(__m128i,__m128i) ENDP          ; _mm_subs_epi64

vecint: Average Color

Mon, 30 Sep 2024 04:41:58 +0000

In a previous post, I used Intel’s AMX instructions intended for AI/ML use-cases to take the average color of an image. This was primarily a proof-of-concept since pedestrians like me generally don’t have access to Intel’s’ AMX-enabled hardware. The cost-of-entry for Intel’s Sapphire Rapids chips is pretty steep too. Maybe some day it will be ubiquitous in consumer-hardware and share a similar story as AVX-512.

Pedestrians like me do have access to an Apple M2 Mini though, after some frustration with trying to sustain a development-environment with a MacOS VM:

Apple M2 with 8-core CPU, 10-core GPU, 16‑core Neural Engine, 8GB unified memory

My “build server” for porting software to MacOS, since a VM wasn’t cutting it anymore.

Liquid Death for scale

AMX

These chips, confusingly, also have an instruction set referred to as “AMX” for AI/ML use-cases, sharing the same name-space as Intel’s AMX instructions.

This instruction set has been reverse-engineered as an open-source effort from the Apple-Clang compiler symbols, the iOS SDK, and other sources.

https://github.com/corsix/amx

Similar to Intel’s AMX, the Apple AMX state is composed of very large registers addressed in two dimensions. In Apple’s AMX, these are the X, Y, and Z register-pools. A single register within these register-pools addresses a 64-byte “row” of elements. The X and Y register-pool has 8 of these rows, and the Z register has 64 rows. The entire contents of the Z register alone is 4KiB!

// A 64 byte "row" of elements
union amx_reg
{
	std::uint8_t	u8 [64];
	std::uint16_t	u16[32];
	std::uint32_t	u32[16];

	std::int8_t		i8 [64];
	std::int16_t	i16[32];
	std::int32_t	i32[16];

	std::float16_t	f16[32];
	std::float32_t	f32[16];
	std::float64_t	f64[ 8];
};

struct amx_state
{
	amx_reg	x[ 8]; // 64 bytes *  8 rows = 512  bytes
	amx_reg	y[ 8]; // 64 bytes *  8 rows = 512  bytes
	amx_reg	z[64]; // 64 bytes * 64 rows = 4096 bytes
};

The instructions themselves will encode how the rows bytes are to be interpreted.

The actual AMX hardware itself can almost be thought of as a “co-processor” shared among the regular cores. Each E-Core and P-Core cluster gets one “AMX core” each. This limits the degree of parallelism you might anticipate depending on how your AMX workload is actually scheduled. You’ll find some interesting consequences from the way this is designed, as it is another “client” of the L2 cache among the other cores.

M2 die shot by Locuza

Average Color

Again, like with Intel’s AMX instructions, I’ll try to find a way to utilize these instructions to help determine the average color of an image and see how much it improves over a few other implementations.

Generic implementation

Here is a generic implementation: Each RGBA8 pixel is unpacked, and then it’s byte-values are added into a much larger 64-bit sum, generally safe from over-flow issues.

// For illustration purposes it is assumed you are providing an image in
// linear color-space. If your image is in an sRGB color-space be sure to
// convert it to linear!
std::uint32_t AverageColorRGBA8(
	const std::uint32_t Pixels[],
	std::size_t Count
)
{
	std::uint64_t RedSum, GreenSum, BlueSum, AlphaSum;
	RedSum = GreenSum = BlueSum = AlphaSum = 0;
	for( std::size_t i = 0; i < Count; ++i )
	{
		const std::uint32_t& CurColor = Pixels[i];
		AlphaSum += static_cast<std::uint8_t>( CurColor >> 24 );
		BlueSum  += static_cast<std::uint8_t>( CurColor >> 16 );
		GreenSum += static_cast<std::uint8_t>( CurColor >>  8 );
		RedSum   += static_cast<std::uint8_t>( CurColor >>  0 );
	}
	RedSum   /= Count;
	GreenSum /= Count;
	BlueSum  /= Count;
	AlphaSum /= Count;

	return
		(static_cast<std::uint32_t>( (std::uint8_t)AlphaSum ) << 24 ) |
		(static_cast<std::uint32_t>( (std::uint8_t) BlueSum ) << 16 ) |
		(static_cast<std::uint32_t>( (std::uint8_t)GreenSum ) <<  8 ) |
		(static_cast<std::uint32_t>( (std::uint8_t)  RedSum ) <<  0 );
}

On clang 18.1.0 with -O2 and an increasing amount of mega-pixels, this generic code performs pretty okay! It’s a clean linear relationship where every 4 megapixels takes about 1 millisecond.

Algorithm	Speed
Generic	~4 megapixels/ms

`UDOT`/`UADDW{,2}`

As a small aside, I also have two other implementations where I utilize UDOT and UADDW{,2} to try and speed up this generic code on ARMv8. They principally follow the idea of using SIMD to accumulate lanes of bytes into larger 16/32-bit sums, and then summing those sums into a larger sum every now-and-then to protect against integer over-flow.

The UDOT implementation can be seen here.

The UADDW{,2} implementation can be seen here.

Algorithm	Speed
Generic	~4 megapixels/ms
`UDOT`/`UADDW`	~13.33 megapixels/ms

Against the generic code, the speedup is pretty substantial. It’s about three times faster! Both the UDOT and UADDW implementations perform about the same and likely saturate the same execution units. At around 6 megapixels, around 22.89 MiB, the timing also seems to get very noisy. This is very likely a sign of this algorithm being memory-bound or an artifact of being scheduled off to other cores after a certain amount of execution time. Even with the noise, it is still consistently much faster than the generic code.

This is probably the kind of implementation you want to be using in your code base when targeting 64-bit ARM.

But! I’m trying to be quirky here and use AI instructions for something that isn’t strictly AI-related. If this was just about making this algorithm faster, the blog post would have ended here.

`vecint`(AMX)

This is the main-event of this post. I’ll step through the implementation details.

If you scroll through the list of AMX instructions, I was looking for something that involves large 8-bit additions across many elements at once at some point and ideally sums into a larger data-type without any lossy overflow issues. There is also the matter of unpacking the individual RGBA channels at some point, which ideally is not a part of the main loop.

The one that stood out to me the most for this particular use-case is vecint. Like most AMX instructions, the vecint instruction just takes a singular 64-bit General-Purpose-Register.

vecint will generally do an z[_][i] ±= f(x[i], y[i])-operation. The f(...)-operation itself and data-types and such is further configured by the 64 bits of the input register.

I’ve marked up the parts that we care about:

Operand bitfields

Used Bit Width Meaning Notes

✅ (47=4) 63 1 Z is signed (1) or unsigned (0)

✅ (47≠4) 63 1 X is signed (1) or unsigned (0)

✅ 58 5 Right shift amount Ignored when ALU mode in {5, 6}

- 57 1 Ignored

- 54 3 Must be zero No-op otherwise

- 53 1 Indexed load (1) or regular load (0)

- (53=1) 52 1 Ignored

- (53=1) 49 3 Register to index into

- (53=1) 48 1 Indices are 4 bits (1) or 2 bits (0)

- (53=1) 47 1 Indexed load of Y (1) or of X (0)

✅ (53=0) 47 6 ALU mode

- 46 1 Ignored

✅ 42 4 Lane width mode Meaning dependent upon ALU mode

- 41 1 Ignored

- (31=1) 35 6 Ignored

- (31=1) 32 3 Broadcast mode

- (31=0) 38 3 Write enable or broadcast mode

- (31=0) 32 6 Write enable value or broadcast lane index Meaning dependent upon associated mode

✅ 31 1 Perform operation for multiple vectors (1)
or just one vector (0) M2 only (always reads as 0 on M1)

- (47=4) 30 1 Saturate Z (1) or truncate Z (0)

- (47=4) 29 1 Right shift is rounding (1) or truncating (0)

- (47≠4) 29 2 X shuffle

- 27 2 Y shuffle

- (47=4) 26 1 Z saturation is signed (1) or unsigned (0)

✅ (47≠4) 26 1 Y is signed (1) or unsigned (0)

✅ (31=1) 25 1 “Multiple” means four vectors (1)
or two vectors (0) Top two bits of Z row ignored if operating on four vectors

- 20 6 Z row Low bits ignored in some lane width modes
When 31=1, top bit or top two bits ignored

- 19 1 Ignored

- 10 9 X offset (in bytes)

- 9 1 Ignored

- 0 9 Y offset (in bytes)

Marked up table from corsix/amx/vecint.md

Used	Bit	Width	Meaning	Notes
✅	(47=4) 63	1	Z is signed (`1`) or unsigned (`0`)
✅	(47≠4) 63	1	X is signed (`1`) or unsigned (`0`)
✅	58	5	Right shift amount	Ignored when ALU mode in {5, 6}
-	57	1	Ignored
-	54	3	Must be zero	No-op otherwise
-	53	1	Indexed load (`1`) or regular load (`0`)
-	(53=1) 52	1	Ignored
-	(53=1) 49	3	Register to index into
-	(53=1) 48	1	Indices are 4 bits (`1`) or 2 bits (`0`)
-	(53=1) 47	1	Indexed load of Y (`1`) or of X (`0`)
✅	(53=0) 47	6	ALU mode
-	46	1	Ignored
✅	42	4	Lane width mode	Meaning dependent upon ALU mode
-	41	1	Ignored
-	(31=1) 35	6	Ignored
-	(31=1) 32	3	Broadcast mode
-	(31=0) 38	3	Write enable or broadcast mode
-	(31=0) 32	6	Write enable value or broadcast lane index	Meaning dependent upon associated mode
✅	31	1	Perform operation for multiple vectors (`1`) or just one vector (`0`)	M2 only (always reads as `0` on M1)
-	(47=4) 30	1	Saturate Z (`1`) or truncate Z (`0`)
-	(47=4) 29	1	Right shift is rounding (`1`) or truncating (`0`)
-	(47≠4) 29	2	X shuffle
-	27	2	Y shuffle
-	(47=4) 26	1	Z saturation is signed (`1`) or unsigned (`0`)
✅	(47≠4) 26	1	Y is signed (`1`) or unsigned (`0`)
✅	(31=1) 25	1	“Multiple” means four vectors (`1`) or two vectors (`0`)	Top two bits of Z row ignored if operating on four vectors
-	20	6	Z row	Low bits ignored in some lane width modes When 31=1, top bit or top two bits ignored
-	19	1	Ignored
-	10	9	X offset (in bytes)
-	9	1	Ignored
-	0	9	Y offset (in bytes)

For the Lane-Width mode: There is an option that will treat the input X and Y operands as unsigned 8-bit integers, and the values in Z as 32-bit integers. Bits 63 and 26 indicate that these values are to be considered unsigned.

Used X Y Z 42

- i16 or u16 i16 or u16 i32 or u32 (two rows, interleaved pair) 3

✅ i8 or u8 i8 or u8 i32 or u32 (four rows, interleaved quartet) 10

- i8 or u8 i8 or u8 i16 or u16 (two rows, interleaved pair) 11

- i8 or u8 i16 or u16 (each lane used twice) i32 or u32 (four rows, interleaved quartet) 12

- i16 or u16 (each lane used twice) i8 or u8 i32 or u32 (four rows, interleaved quartet) 13

- i16 or u16 i16 or u16 i16 or u16 (one row) anything else

Marked up table from corsix/amx/vecint.md

Used	X	Y	Z	42
-	i16 or u16	i16 or u16	i32 or u32 (two rows, interleaved pair)	`3`
✅	i8 or u8	i8 or u8	i32 or u32 (four rows, interleaved quartet)	`10`
-	i8 or u8	i8 or u8	i16 or u16 (two rows, interleaved pair)	`11`
-	i8 or u8	i16 or u16 (each lane used twice)	i32 or u32 (four rows, interleaved quartet)	`12`
-	i16 or u16 (each lane used twice)	i8 or u8	i32 or u32 (four rows, interleaved quartet)	`13`
-	i16 or u16	i16 or u16	i16 or u16 (one row)	anything else

For the ALU mode, the actual operation is configured. With the previously defined bits, we have two u8 operands from X and Y, and a u32 value in Z.

Used Integer operation 47 Notes

- z+((x*y)>>s) 0

- z-((x*y)>>s) 1

✅ z+((x+y)>>s) 2 Particular write enable mode can skip x or y

- z-((x+y)>>s) 3 Particular write enable mode can skip x or y

- z>>s or sat(z>>s) 4 Shift can be rounding, saturation is optional

- sat(z+((x*y*2)>>16)) 5 Shift is rounding, saturation is signed

- sat(z-((x*y*2)>>16)) 6 Shift is rounding, saturation is signed

- (x*y)>>s 10 M2 only

- z+(x>>s) 11 M2 only (on M1, consider 47=2 with skipped y)

- z+(y>>s) 12 M2 only (on M1, consider 47=2 with skipped x)

- no-op anything else

Marked up table from corsix/amx/vecint.md

Used	Integer operation	47	Notes
-	`z+((x*y)>>s)`	`0`
-	`z-((x*y)>>s)`	`1`
✅	`z+((x+y)>>s)`	`2`	Particular write enable mode can skip `x` or `y`
-	`z-((x+y)>>s)`	`3`	Particular write enable mode can skip `x` or `y`
-	`z>>s` or `sat(z>>s)`	`4`	Shift can be rounding, saturation is optional
-	`sat(z+((xy2)>>16))`	`5`	Shift is rounding, saturation is signed
-	`sat(z-((xy2)>>16))`	`6`	Shift is rounding, saturation is signed
-	`(x*y)>>s`	`10`	M2 only
-	`z+(x>>s)`	`11`	M2 only (on M1, consider 47=2 with skipped `y`)
-	`z+(y>>s)`	`12`	M2 only (on M1, consider 47=2 with skipped `x`)
-	no-op	anything else

We effectively want to do Z += X + Y which maps closest to z+((x+y)>>s). The s argument is an additional right-shift that can be applied from the 5-bit integer at bit-position 58. In our case though, it is left at 0, so this effectively maps to z+((x+y)>>0) = z+(x+y).

The way this operation is configured now, each byte from a row X and Y will be added together and get its own unique 32-bit sum to accumulate to in Z. Because each byte is effectively expanded and added to a four-byte sum the results in Z have to be spread out across multiple rows to accommodate the data-type expansion. So one row of bytes in X and Y will map to four rows of 32-bit sums in Z.

It also de-interleaves every four bytes across these rows, so we get unpacked RGBA channels for free!

Row 0 of inputs X/Y will map to rows 0..3 of Z
- Each byte of each RGBA-element is spread across each of the four rows: 0:R 1:G 2:B 3:A

So with each iteration of vecint, we process process 16-pixels in X and 16-pixels in Y for a total of 32-pixels of input and mapping to 16 32-bit sums in Z of output. In just one instruction!

But, we have 8 rows in X and Y and 64 rows in Z that could be utilized more.

We could always just emit the instruction several times and index into the additional rows of X and Y with the 9-bit integers at bit-positions 0 and 10. But, the M2 chip in particular supports some additional bulk-processing configurations to processes even more data at a time in just one instruction:

Operand bitfields

Used Bit Width Meaning Notes

- … … …

✅ 31 1 Perform operation for multiple vectors (1)
or just one vector (0) M2 only (always reads as 0 on M1)

- … … …

✅ (31=1) 25 1 “Multiple” means four vectors (1)
or two vectors (0) Top two bits of Z row ignored if operating on four vectors

- … … …

Marked up table from corsix/amx/vecint.md

Used	Bit	Width	Meaning	Notes
-	…	…	…
✅	31	1	Perform operation for multiple vectors (`1`) or just one vector (`0`)	M2 only (always reads as `0` on M1)
-	…	…	…
✅	(31=1) 25	1	“Multiple” means four vectors (`1`) or two vectors (`0`)	Top two bits of Z row ignored if operating on four vectors
-	…	…	…

Setting bit 31 indicates that this one instruction will now process two rows of input rather than just one.

The M2 in particular adds bit 25 which further configures this instruction to process four rows of input at a time rather than just two! Each of the “multiple”-operation will separate the rows of output by 16 rows:

Row 0 of inputs X/Y will map to rows 0..3 of Z
- Each byte of each RGBA-element is spread across each of the four rows: 0:R 1:G 2:B 3:A
Row 1 of inputs X/Y will map to rows 16..19 of Z
- Each byte of each RGBA-element is spread across each of the four rows: 16:R 17:G 18:B 19:A
Row 2 of inputs X/Y will map to rows 32..35 of Z
- Each byte of each RGBA-element is spread across each of the four rows: 32:R 33:G 34:B 35:A
Row 3 of inputs X/Y will map to rows 48..51 of Z
- Each byte of each RGBA-element is spread across each of the four rows: 48:R 49:G 50:B 51:A

So at the end of our loop, there will be a one-time tango that has to be done to unpack these RGBA 32-bit sums into our larger 64-bit sums.

So with each iteration of vecint, we process process 64-pixels in X and 64-pixels in Y for a total of 128-pixels of input and mapping to 64 32-bit sums in Z of output. In just one instruction!

The instructions to actually load data into X and Y also support these additional bulk-processing configurations. The M2 in particular supports loading four rows of data at a time, while the baseline instruction supports just one or two.

For ldx / ldy:

Used Bit Width Meaning

- 63 1 Ignored

✅ 62 1 Load multiple registers (1) or single register (0)

- 61 1 On M1/M2: Ignored (loads are always to consecutive registers)
On M3: Load to non-consecutive registers (1) or to consecutive registers > (0)

✅ 60 1 On M1: Ignored (“multiple” always means two registers)
On M2/M3: “Multiple” means four registers (1) or two registers (0)

- 59 1 Ignored

- 56 3 X / Y register index

- 0 56 Pointer

Marked up table from corsix/amx/ldst.md

Used	Bit	Width	Meaning
-	63	1	Ignored
✅	62	1	Load multiple registers (`1`) or single register (`0`)
-	61	1	On M1/M2: Ignored (loads are always to consecutive registers) On M3: Load to non-consecutive registers (`1`) or to consecutive registers > (`0`)
✅	60	1	On M1: Ignored (“multiple” always means two registers) On M2/M3: “Multiple” means four registers (`1`) or two registers (`0`)
-	59	1	Ignored
-	56	3	X / Y register index
-	0	56	Pointer

With this, the main loop is just a very tight three instructions:

Load 64 pixels into X
Load 64 pixels into Y
Perform vecint(Z += X + Y)

At the end of the loop, and every now and then to protect from overflow, the sums in Z have to be written to memory and unpacked and added into the larger 64-bit sums. There is only one bit of configuration for storing two rows of Z at a time.

For ldz / stz:

Bit Width Meaning

63 1 Ignored

62 1 Load / store pair of registers (1) or single register (0)

56 6 Z row

0 56 Pointer

Marked up table from corsix/amx/ldst.md

Bit	Width	Meaning
63	1	Ignored
62	1	Load / store pair of registers (`1`) or single register (`0`)
56	6	Z row
0	56	Pointer

Once the rows are stored, there is some additional work to be done to add them to the outer 64-bit sums, while protecting from integer-overflow. This part of the process is a little ugly, but it’s also very infrequent. You either do this once at the end of your loop, or every 16'843'009 iterations(~17 mega-pixels) to protect against overflow.

Thankfully, since vecint already de-interleaves the RGBA channels, each row will only have R, G, B, or A-sums within it! So just one more horizontal-addition of 32-bit sums needs to be done.

This can be accelerated with pair-wise widening-additions like with vpaddlq_u32, vpadalq_u32 and a final vaddvq_u64 to add safely add up all of the 32-bit sums into a larger 64-bit one without any worry of overflow.

The final implementation:

std::uint32_t AverageColorRGBA8(const std::uint32_t Pixels[], std::size_t Count)
{
	std::size_t i = 0;

	std::uint64_t RedSum   = 0ULL;
	std::uint64_t GreenSum = 0ULL;
	std::uint64_t BlueSum  = 0ULL;
	std::uint64_t AlphaSum = 0ULL;

	// 128 pixels at a time!
	for( std::size_t j = i / 128; j < Count / 128; j++ )
	{
		// Required before any AMX instructions
		AMX_SET();

		// In the worst case, where all the bytes are just 0xFF being summed
		// into a 32-bit accumulator, the 32-bit sum may overflow unless we
		// ensure all 32-bit overflow-hazards are protected against.
		// In this case: `(0xFFFFFFFF / 0xFF == 0x1010101` is the max amount of
		// bytes we could ever safely accumulate into a 32-bit integer before we
		// have to flush it into a larger data-type
		constexpr std::size_t LocalSumOverflowMax = (0xFFFFFFFF / 0xFF);
		for( std::size_t k = 0; (k < LocalSumOverflowMax) && (j < Count / 128);
			k++, j++, i += 128 )
		{
			// Load 64 pixels into X
			AMX_LDX(
				reinterpret_cast<std::uintptr_t>((const uint8_t*)&Pixels[i + 0])
				| (1ULL << 62) // Load multiple times
				| (1ULL << 60) // Load four times
			);
			// Load another 64 pixels into Y
			AMX_LDY(
				reinterpret_cast<std::uintptr_t>((const uint8_t*)&Pixels[i + 64]
				)
				| (1ULL << 62) // Load multiple times
				| (1ULL << 60) // Load four times
			);

			constexpr std::uint64_t VecIntOp =
				// ALU mode
				//  2 : f = Z+(X + Y) >> s
				((2ULL) << 47) |
				// Lane width mode:
				// 10: Z.u32[i] += f(X.u8[i], Y.u8[i])
				// Produces 64 32-bit integers, requiring 256 bytes of data
				// total! four rows of Z: interleaved quartet(perfect for
				// RGBA!):
				((10ULL) << 42) |

				// Iterate multiple times (M2 only)
				((1ULL) << 31) |
				// Iterate x4 times (M2 only)
				((1ULL) << 25);

			// Add each 64-bit value into a 32-bit sum across four rows of Z
			// Z0 + Iter * 16: RSum32, RSum32, RSum32, RSum32, RSum32...
			// Z1 + Iter * 16: GSum32, GSum32, GSum32, GSum32, GSum32...
			// Z2 + Iter * 16: BSum32, BSum32, BSum32, BSum32, BSum32...
			// Z3 + Iter * 16: ASum32, ASum32, ASum32, ASum32, ASum32...
			AMX_VECINT(VecIntOp);
		}

		// Each row is full of 32-bit sums for the R, G, B, or A channel
		uint32x4x4_t ZMat[4 * 4] = {{}};

		for( std::size_t IterationIndex = 0; IterationIndex < 4;
			++IterationIndex )
		{
			for( std::size_t ZRow = 0; ZRow < 4; ZRow += 2 )
			{
				// Store pairs-of-rows of Z(64 bytes x 2)
				AMX_STZ(
					reinterpret_cast<std::uintptr_t>(&ZMat[ZRow + 4 * IterationIndex])
					| static_cast<std::uint64_t>(ZRow + 16 * IterationIndex) << 56
					| 1ULL << 62
				);
			}
		}

		// Should be used when done with any AMX processing.
		// Think of it like `vzeroupper` from AVX
		AMX_CLR();

		for( std::size_t IterationIndex = 0; IterationIndex < 4;
			++IterationIndex )
		{
			const std::size_t IterationOffset = IterationIndex * 4;
			// Widening pair-wise sums are used to ensure safety from overflow
			RedSum += vaddvq_u64(vpadalq_u32(
				vpadalq_u32(
					vpadalq_u32(
						vpaddlq_u32(ZMat[0 + IterationOffset].val[0]),
						ZMat[0 + IterationOffset].val[1]
					),
					ZMat[0 + IterationOffset].val[2]
				),
				ZMat[0 + IterationOffset].val[3]
			));
			GreenSum += vaddvq_u64(vpadalq_u32(
				vpadalq_u32(
					vpadalq_u32(
						vpaddlq_u32(ZMat[1 + IterationOffset].val[0]),
						ZMat[1 + IterationOffset].val[1]
					),
					ZMat[1 + IterationOffset].val[2]
				),
				ZMat[1 + IterationOffset].val[3]
			));
			BlueSum += vaddvq_u64(vpadalq_u32(
				vpadalq_u32(
					vpadalq_u32(
						vpaddlq_u32(ZMat[2 + IterationOffset].val[0]),
						ZMat[2 + IterationOffset].val[1]
					),
					ZMat[2 + IterationOffset].val[2]
				),
				ZMat[2 + IterationOffset].val[3]
			));
			AlphaSum += vaddvq_u64(vpadalq_u32(
				vpadalq_u32(
					vpadalq_u32(
						vpaddlq_u32(ZMat[3 + IterationOffset].val[0]),
						ZMat[3 + IterationOffset].val[1]
					),
					ZMat[3 + IterationOffset].val[2]
				),
				ZMat[3 + IterationOffset].val[3]
			));
		}
	}

	// ...Generic implementation
}

Benchmarked on my M2 Mac Mini:

While it’s not a huge improvement over the UDOT/UADDW implementation, it is still consistently faster than it by about 3 megapixels/ms.

Algorithm	Speed
Generic	~4 megapixels/ms
`UDOT`/`UADDW`	~13.33 megapixels/ms
`vecint`(AMX)	~16.80 megapixels/ms

We’re processing 128 pixels at a time… Why aren’t see seeing something even close to a theoretical x128-speedup over over a serial implementation?

This could be for a lot of reasons ~~that I didn’t care to look too deeply into since this is just a proof-of-concept for an undocumented instruction-set that is probably going to be deprecated~~.

It’s likely that these accelerators were primarily just intended as an efficiency-measure more than anything else.

Maybe having a slower but highly-efficient accelerator is worth sacrificing a bit of throughput? Why tie up a whole core with a highly redundant matrix/vector task when you can free-up the core for other tasks and move that workload over to an accelerator? What if the thread running this AMX code gets scheduled and pinned to the E-core implementation of AMX and not the P-Core?

Or it’s just memory-bound? 16.80 megapixels/ms translates to about 67.2 GB/s(62.58 GiB/s) of data being processed in this tight loop which is a good chunk of the 100 GB/s bandwidth available on the M2 considering the amount of bandwidth available for the rest of the system. On the M2 Pro/Max/Ultra and M3 chips, this chart would certainly look very different!

This proof-of-concept does not even utilize the 16 “Neural Engine” cores, but there is no obvious way to utilize these for non-AI tasks.

SME

As of the newer M4 chips of 2024, Apple seems to have begun implementing the more standardized SME instruction set designed by ARM themselves. So far, only the latest iPad Pro features the M4 chip as of this writing.

Some others have already begun to analyze its characteristics:

This likely indicates an intent to remove or repurpose the AMX die-space now that there is a standard instruction-set that even developers can confidently ship into their software.

AMX instructions were never intended to be used by developers in any shipped software. These instructions are just a hidden-implementation detail of Apple’s private library-code. So ripping these instructions out of the actual hardware and making compatibility adjustments to the affected library-code would just be a bit of a quiet tragic event while clients of that library code would be none the wiser of anything happening.

Maybe my next post will be about another quirky Proof Of Concept using SVE/SME instructions now that pedestrians like me can have access to these instructions!

GPU Debug Scopes

Tue, 17 Sep 2024 02:01:29 +0000

Rendering APIs these days tend to capture their gpu workloads into a serialized form such as a command-buffer or command-list to be dispatched at a later time into a work-queue.

Diagnostic tooling such as RenderDoc or Nsight-Graphics allows the disecting of these command-buffers, but it’s not very obvious to determine what is happening at a high level from the list of API commands alone:

RenderDoc(Before)

Nsight-Graphics(Before)

Without any additional debugging information, RenderDoc and Nsight will show a flat list of command-buffer API-calls and will provide some filtering and categorization of these commands to help track down the ones that you care about. This process is slow, especially when working with multiple captures and need to draw some kind of comparisons between them.

What if you want to ensure that some host-code ran?
What if you want to ensure Step 1 and Step 2 ran before some issue at Step 3?
Where did this extra API call come from?
How do I make sure my cool optimization ran here?
What if your host code made an opportunistic early-return and skipped some API calls that you were expecting?

It’s difficult to capture these kind of contexts with a flat list of API calls.

Thankfully, rendering APIs tend to allow the attaching of diagnostic data to both command-buffers and objects to provide valuable diagnostic information to your captures:

RenderDoc(After)

Nsight-Graphics(After)

After adding some object names, and debugging scopes, both RenderDoc and Nsight will interpret this data to have both readable object-names and allows groups of command-buffer API-calls to be grouped and even colorized to your liking. Above, I generated a color for the pipeline-syncing API-call based on bits of the hash of the graphics-pipeline itself so I can identify if a pipeline is being utilized repetitively at a glance.

I’ll talk about Vulkan and OpenGL’s particular implementation of such features and how to utilize RAII(Resource Acquisition Is Initialization) patterns to automatically maintain nested scopes to create a sort of call stack within your command-buffers.

This is a pattern I’ve found myself utilizing quite a lot to help with debugging, diagnosing issues, and profiling.

RAII

Before we get to the Graphics-API specific implementation, here’s a quick rundown on how RAII looks like in C++, C#, and Rust as well.

The final implementations will be provided in C++, but can be mapped to C# and Rust through their various OpenGL and Vulkan bindings.

C++

In C++, one simply has to implement code into the constructor and deconstructor to achieve a RAII pattern.

class DebugScope
{
public:
	// Constructor
	DebugScope(const char* ScopeName)
	{
		// Begin scope
		// Graphics API code
	}

	// Deconstructor
	~DebugScope()
	{
		// End Scope
		// Graphics API code
	}
};

Usage:

void Work();
{
	DebugScope Scope("Work"); // Constructor called

	EvenMoreWork();

	if( error )
	{
		return; // Deconstructor called
	}

} // Deconstructor called

void ScopeTest()
{
	DebugScope Scope("ScopeTest"); // Constructor called

	Work();

	if( error )
	{
		return; // Deconstructor called
	}

	MoreWork();

} // Deconstructor called

CSharp

C# gets a bit more tricky. C# as a language is garbage-collected so the lifetime of an object is determined by the scheduling of the garbage collector. So you cannot deterministically know when a class’s Deconstructor gets called or resources get released. Some additional work must be done to get C++’s behavior where the deconstructor automatically gets called upon leaving the scope. We’re trying to avoid having to manually call functions here!

C# exposes the IDisposable interface for classes to implement for the releasing of unmanaged resources such as file-objects, GPU-objects, Native-types, other IDisposable-types, or any other type that is not handled by the garbage collector. In this case though our needs are simpler. We aren’t actually freeing any GPU resources, we just want to call some Graphics-API calls for some automatic scope-management.

using System;

class DebugScope : IDisposable
{
	// Constructor
	public DebugScope(string ScopeName)
	{
		// Begin scope
		// Graphics API code
	}

	// Implement IDisposable
	public void Dispose()
	{
		// End Scope
		// Graphics API code
	}
}

The using keyword will also ensure that an object is valid during the scope of the using-block and will automatically call Dispose upon leaving the scope. There’s no clean way to enforce a class to only be utilized within a using-block though, to ensure proper RAII-behavior. This pattern will have to be a discretion of the code-base.

Usage:

static void Work()
{
	// Constructor called
	using( DebugScope Scope = new DebugScope("Work") )
	{
		EvenMoreWork();

		if( error )
		{
			return; // Dispose called
		}
	} // Dispose called
}

static void ScopeTest()
{
	// Constructor called
	using( DebugScope Scope = new DebugScope("ScopeTest") )
	{
		Work();

		if( error )
		{
			return; // Dispose called
		}

		MoreWork();
	} // Dispose called
}

I have an old PR for Ryujinx that implements this for their Vulkan backend!

Rust

Rust implements RAII by implementing the Drop trait. By implementing fn drop(&mut self);, code can now be ran when the struct leaves a scope.

pub struct DebugScope;

impl DebugScope {
    // Constructor
    pub fn new(scope_name: &str) -> Self {
        // Begin scope
        // Graphics API code
        return DebugScope {};
    }
}

// Implement Drop trait
impl Drop for DebugScope {
    // Deconstructor
    fn drop(&mut self) {
        // End scope
        // Graphics API code
    }
}

Similar to C++, just defining the object is enough for our code to run when defined within a scope and upon leaving the scope. Since scope-objects don’t usually have to be touched after they are defined, the object can be named with an underscore(_) before its name to avoid any “unused variable”-warnings.

Usage:

pub fn work() {
    // Constructor called
    let _scope = DebugScope::new("Work");

    even_more_work();

    if (error) {
        return; // Drop called
    }
} // Drop called

pub fn scope_test() {
    // Constructor called
    let _scope = DebugScope::new("ScopeTest");

    work();

    if (error) {
        return; // Drop called
    }

    more_work();
} // Drop called

Implementations

Vulkan

Vulkan provides the VK_EXT_debug_utils extension to allow attaching names to objects as well as colored labels to spans of command-buffer commands and queue-operations.

The vkCmd{Begin,End,Insert}DebugUtilsLabelEXT-functions are utilized to group and label particular spans of command buffer operations with a VkDebugUtilsLabelEXT structure. This allows both a plaintext name(const char*) an RGBA floating-point color(float[4]) to be correlated with command buffer operations:

// Provided by VK_EXT_debug_utils
typedef struct VkDebugUtilsLabelEXT {
	VkStructureType    sType;
	const void*        pNext;
	const char*        pLabelName;
	float              color[4];
} VkDebugUtilsLabelEXT;

vkCmdInsertDebugUtilsLabelEXT additionally allows the insertion of additional one-off labels within a command buffer as well. An additional function or operator-overload may be added to the DebugScope object to insert these additional labels.

A minimally viable Vulkan implementation, ready for you to copy-paste, could look like this:

class DebugScope {
private:
	// Keep this command buffer around so that the deconstructor can properly
	// end the debug-scope
	const VkCommandBuffer commandBuffer;

public:
	// Upon construction, begin the debug-scope
	DebugScope(
		VkCommandBuffer targetCommandBuffer,
		const char* scopeName, std::span<const float, 4> scopeColor
	) : commandBuffer(targetCommandBuffer)
	{
		VkDebugUtilsLabelEXT label = {};
		label.sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_LABEL_EXT;
		label.pLabelName = scopeName;
		std::copy_n(scopeColor.begin(), 4, label.color);
		vkCmdBeginDebugUtilsLabelEXT(commandBuffer, &label);
	}

	// A bonus operator to insert plain labels within the command-buffer
	void operator()(const char* scopeName, std::span<const float, 4> scopeColor) const
	{
		VkDebugUtilsLabelEXT label = {};
		label.sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_LABEL_EXT;
		label.pLabelName = scopeName;
		std::copy_n(scopeColor.begin(), 4, label.color);
		vkCmdInsertDebugUtilsLabelEXT(commandBuffer, &label);
	}

	// Upon deconstruction, begin the debug-scope
	~DebugScope()
	{
		vkCmdEndDebugUtilsLabelEXT(commandBuffer);
	}
};

// Usage
void DoThing(VkCommandBuffer commandBuffer)
{
	static float thingColor[4] = {1.0f, 1.0f, 0.0f, 1.0f};
	DebugScope scope(commandBuffer, "DoThing", thingColor);

	scope("Step1", thingColor);

	vkCmd...(commandBuffer);

	scope("Step2", thingColor);

	vkCmd...(commandBuffer);

	scope("Step3", thingColor);

	vkCmd...(commandBuffer);
}

The basic pattern can be extended further to add even more conveniences such as utilizing Vulkan-Hpp to help make the code more concise and expressive or utilizing fmt to aid in scope and label name generation. You could even put __FILE__ or __LINE__ or the calling function-name itself into the debug-scope name to be able to more easily “blame” each command buffer scope to the exact host-code that emitted it.

To avoid additional overhead, you might choose to use something like VK_EXT_tooling_info (Core in Vulkan 1.3) to only conditionally insert these API commands if it detects that RenderDoc or Nsight is attached to the Vulkan instance.

Scope and label coloration can either be manually decided at each call-site or it can be generated to your choosing.

A simple one is to maintain a static depth-integer within the DebugScope-object that increments/decrements in the ctor/dtor. Knowing what depth each scope is at allows for procedural color-selection such as utilizing something like Inigo Quilez’s procedural color-palettes. A contribution I made to DuckStation utilized this pattern in particular. This would work fine if you only ever operated upon a single recycled command buffer. In a multi-threaded environment, you will probably want to maintain this depth-variable in a per-command-buffer abstraction as opposed to having a globally shared depth-variable between all command buffers.

Before

After

Another option for coloration is to group certain operations by colors such as making all “transfer” workloads yellow, all “graphics” workloads green, all “compute” workloads orange, and all “present” operations magenta.

Some code-bases might further decide to color the labels based on the specific operation, such as coloring a label for a vkCmdClearColorImage operation with the clear-color itself. A contribution I made to Panda3DS utilized this pattern.

With this additional data in your command buffer, debug callbacks will also be able to interpret this additional context within the VkDebugUtilsMessengerCallbackDataEXT structure.

Each of the originally defined VkDebugUtilsLabelEXT structures for each label can be derived by iterating with the pCmdBufLabels and cmdBufLabelCount variables.

These labels are sorted from oldest to newest. So pCmdBufLabels[0] would be the oldest label that was set leading into the current debug message, and pCmdBufLabels[cmdBufLabelCount - 1] would be the most recent label.

This could provide valuable context around particular error messages to help diagnose an issue.

VKAPI_ATTR VkBool32 VKAPI_CALL DebugMessageCallback(
	VkDebugUtilsMessageSeverityFlagBitsEXT      MessageSeverity,
	VkDebugUtilsMessageTypeFlagsEXT             MessageType,
	const VkDebugUtilsMessengerCallbackDataEXT* CallbackData, void* UserData
)
{
	// Loop through all labels for this particular message
	for( std::uint32_t i = 0; i < CallbackData->cmdBufLabelCount; ++i )
	{
		const VkDebugUtilsLabelEXT& CurLabel = CallbackData->pCmdBufLabels[i];
		std::fprintf(stderr, "%u [%s]\n", i, CurLabel.pLabelName);
	}

	switch( vk::DebugUtilsMessageSeverityFlagBitsEXT(MessageSeverity) )
	{
	case vk::DebugUtilsMessageSeverityFlagBitsEXT::eError:
	case vk::DebugUtilsMessageSeverityFlagBitsEXT::eWarning:
	{
		// Something bad happened!
		const char* Message = CallbackData->pMessage;
		std::puts(Message);
		assert(0);
	}
	...
	}
	return VK_FALSE;
}

In this example output, I’ve artifically doubled the size of a pipeline barrier within a scope named Upload Data:

0 [Download]
1 [Rendering]
2 [Upload Data]        <<< This is the most-recent label reached before this error!
Validation Error: [ VUID-VkBufferMemoryBarrier-size-01189 ] Object 0: handle = 0xdb05ecf0, type = VK_OBJECT_TYPE_COMMAND_BUFFER; | MessageID = 0xb63479f2 | vkCmdPipelineBarrier(): pBufferMemoryBarriers[0].size VkBuffer 0xab64de0000000020[] has offset 0x0 and size 0x2f4400 whose sum is greater than total size 0x17a200. The Vulkan spec states: If size is not equal to VK_WHOLE_SIZE, size must be less than or equal to than the size of buffer minus offset (https://vulkan.lunarg.com/doc/view/1.3.268.0/windows/1.3-extensions/vkspec.html#VUID-VkBufferMemoryBarrier-size-01189)

With this, I can now know exactly what part of the code-base to start investigating the issue in.

OpenGL

OpenGL provides the GL_KHR_debug extension for attaching diagnostic information to the rendering context.

Since OpenGL API calls operate upon a global-state, gl{Push,Pop}DebugGroup will group together API-calls at a global-scope.

I have yet to see any GPU tooling utilize the id parameter of glPushDebugGroup, but I’ve assigned it to the global scope-depth to try and keep this code future-facing to any diagnostic tooling that may eventually decide to do something with it. You could just statically provide it 0 if you wanted to.

class DebugScope {
	inline static GLuint GlobalScopeDepth = 0;
	const GLuint ScopeDepth;

	public:
	DebugScope(std::string_view ScopeName)
		: ScopeDepth(GlobalScopeDepth++)
	{
		glPushDebugGroup(GL_DEBUG_SOURCE_APPLICATION, ScopeDepth, ScopeName.size(), ScopeName.data());
	}

	~DebugScope()
	{
		glPopDebugGroup();
		GlobalScopeDepth--;
	}
};

An implementation within Lua provided for Love2D

Wunk

vpternlog: Signed Saturation

Compile time constants

Signed Saturated arithmetic

vecint: Average Color

AMX

Average Color

Generic implementation

UDOT/UADDW{,2}

vecint(AMX)

Operand bitfields

Operand bitfields

SME

GPU Debug Scopes

RAII

If you already know what RAII is, you can just skip to the implementation

C++

CSharp

Rust

Implementations

Vulkan

OpenGL

`UDOT`/`UADDW{,2}`

`vecint`(AMX)