A signed saturating add through NAK

I was browsing Mesa issues when add.sat.s32 caught my eye. It's signed addition with saturation instead of wraparound. NVIDIA's compiler knew what to do with it; NAK didn't have a native path. The issue already had a five-instruction solution, which looked like a manageable thing to investigate. Famous last words.

Faith Ekstrand had already posted that five-instruction sequence in the Mesa issue. I wasn't discovering the arithmetic. I wanted to know how to make NAK express those .SIGN operands, encode them, and keep the later passes from undoing it.

Some names that will come up:

PTX    NVIDIA's portable assembly language
SASS   the machine instructions executed by an NVIDIA GPU
NIR    Mesa's shared compiler intermediate representation
NAK    Mesa's NVIDIA shader backend: NIR → SASS
GPR    a 32-bit general-purpose register, one value per GPU lane
Pred   a one-bit predicate register used for conditions

Five instructions. What could go wrong?

The issue contained this SASS:

IADD3 R0, R2, R5, RZ
PLOP3.LUT P0, PT, R5.SIGN, R2.SIGN, R0.SIGN, 0x2,  0x0
PLOP3.LUT P1, PT, R5.SIGN, R2.SIGN, R0.SIGN, 0x40, 0x0
SEL R9, R0, 0x7fffffff, !P0
SEL R9, R9, 0x80000000, !P1

An add, two Boolean lookups, two selects. Fine. But look at R5.SIGN. That's a GPR sign bit, not an ordinary predicate. NAK could already encode predicate PLOP3; this was a different operand form. The disassembler knew its name. Did NAK know how to represent it?

Why NAK never saw the instruction

NIR already had nir_op_iadd_sat. NAK requested generic lowering unconditionally: 

lower_uadd_sat: dev.sm < 70,
lower_usub_sat: dev.sm < 70,
lower_iadd_sat: true, // TODO

So the operation disappeared before my encoder could possibly see it. I let it survive on SM70+, then followed it through NIR translation, NAK IR, legalization, and encoding. Five machine instructions had already become a compiler-wide tour.

What are those two LUT values?

The two LUT constants looked like magic numbers until I wrote down the three signs: x, y, and the wrapping sum. Two positive inputs with a negative sum mean positive overflow; two negative inputs with a nonnegative sum mean negative overflow. Only two rows of the truth table matter. 

x.sign  y.sign  sum.sign  result
   0       0        1     positive overflow → INT_MAX
   1       1        0     negative overflow → INT_MIN

Those rows are 001 and 110 in the chosen input order. Bits 1 and 6: 0x02 and 0x40. Nice. Now for the less tidy question: what bytes does the GPU expect? 

sum          = IADD3(x, y, 0)
pos_overflow = PLOP3(x.sign, y.sign, sum.sign, LUT=0x02)
neg_overflow = PLOP3(x.sign, y.sign, sum.sign, LUT=0x40)
tmp          = SEL(pos_overflow, INT_MAX, sum)
dst          = SEL(neg_overflow, INT_MIN, tmp)

Let the existing compiler show its hand

I let NVIDIA's tools answer first. A tiny PTX kernel through ptxas, then nvdisasm to see the machine instructions and their encoding. No need to guess from the mnemonic. 

.version 8.8
.target sm_89
.address_size 64

.visible .entry test_iadd_sat(
    .param .u64 out_ptr,
    .param .u64 in_ptr)
{
    .reg .b64 %rd<4>;
    .reg .s32 %r<4>;

    ld.param.u64 %rd1, [out_ptr];
    ld.param.u64 %rd2, [in_ptr];
    add.u64 %rd3, %rd2, 4;
    ld.global.s32 %r1, [%rd2];
    ld.global.s32 %r2, [%rd3];
    add.sat.s32 %r3, %r1, %r2;
    st.global.u32 [%rd1], %r3;
    ret;
}

The corresponding ptxas and nvdisasm commands are: 

ptxas -arch=sm_89 iadd_sat.ptx -o iadd_sat.cubin
nvdisasm --print-code --print-instruction-encoding iadd_sat.cubin |
  rg 'IADD3|PLOP3|SEL'

The relevant output was the same five-instruction core shown in the issue: 

IADD3 R2, R0, R3, RZ
PLOP3.LUT P0, PT, R0.SIGN, R3.SIGN, R2.SIGN, 0x2,  0x0
PLOP3.LUT P1, PT, R0.SIGN, R3.SIGN, R2.SIGN, 0x40, 0x0
SEL R7, R2, 0x7fffffff, !P0
SEL R7, R7, 0x80000000, !P1

The five-instruction core matched. I briefly considered reconstructing the whole 128-bit PLOP3 format, but most of it was already in NAK: the ordinary predicate variant, common ALU fields, and destinations. I compared the known form with the CUDA-produced .SIGN form and filled in what was missing.

The first encoder version reused NAK's common three-source ALU encoding. The comparison left only the PLOP3 form, LUT, and two predicate destinations to fill in: 

e.encode_alu(0x01f, None,
             Some(&self.srcs[0]),
             Some(&self.srcs[1]),
             Some(&self.srcs[2]));
e.set_field(72..80, self.ops[0].lut);
e.set_pred_dst(81..84, &self.dsts[0]);
e.set_pred_dst(84..87, &self.dsts[1]);

I then made NAK's disassembly test encode the IR and ask NVIDIA's decoder for the expected text: 

export NAK_TEST=/path/to/mesa-build/src/nouveau/compiler/nak

PATH=/path/to/cuda/bin:$PATH \
  "$NAK_TEST" \
  --test --exact nvdisasm_tests::test_plop3 --test-threads 1

nvdisasm could read the new encoding. Good. It still hadn't executed anything, though. For that I needed an actual GPU and inputs near signed overflow.

A model that was just enough

My first IR model was deliberately small: one flag for either all-predicate sources or all-GPR-sign sources. Both overflow lookups needed three signs, so it covered the case in front of me. I didn't yet know whether mixed inputs needed their own representation.

The patch connected the whole vertical slice:

NIR iadd_sat
  → NAK builder: IADD3 + two PLOP3 + two SEL
  → OpPLop3 with sign-mode GPR sources
  → legalization
  → SM70+ encoding
  → AD107

The local hardware test stores the GPU result and compares it with Rust's saturating_add for boundary inputs and 100 deterministic random pairs: 

PATH=/path/to/cuda/bin:$PATH \
  "$NAK_TEST" \
  --test --exact hw_tests::test_iadd_sat --test-threads 1

The GPU test passed. CTS disagreed.

The focused GPU test passed. So did the decoder test. Then I ran CTS: 152 failures, all at a compiler assertion. Not exactly the victory lap I had in mind. 

844 pass
152 fail
failure: NAK assertion during pipeline compilation

At least the GPU hadn't produced 152 mysterious wrong answers. It hadn't executed those shaders at all.

My small test used ordinary per-lane GPRs. CTS brought in values that NAK could put in uniform registers, and NAK still considered every PLOP3 uniform-capable. My .SIGN encoding only accepted GPRs. The two assumptions met at the encoder assertion.

I left the assertion alone and fixed the path leading to it: .SIGN PLOP3 was marked non-uniform, inputs were copied to GPRs when needed, and predicate optimizations stopped treating its sources as predicates. The CTS filter then passed all 996 cases. 

// Ordinary predicate PLOP3 may be uniform; .SIGN PLOP3 may not.
Op::PLop3(op) => !op.src_sign,

// The encoder's .SIGN form consumes normal GPRs.
if self.src_sign {
    for src in &mut self.srcs {
        b.copy_alu_src_if_not_reg(src, RegFile::GPR, SrcType::GPR);
    }
}
uniform value
    → copy to a normal GPR in each lane
    → warp PLOP3 R?.SIGN, R?.SIGN, R?.SIGN
export MESA_PREFIX=/path/to/patched-mesa-install
export VK_CTS=/path/to/VK-GL-CTS/build/external/vulkancts/modules/vulkan/deqp-vk

env VK_DRIVER_FILES="$MESA_PREFIX/share/vulkan/icd.d/nouveau_icd.x86_64.json" \
  LD_LIBRARY_PATH="$MESA_PREFIX/lib" \
  MESA_SHADER_CACHE_DISABLE=true \
  "$VK_CTS" -n \
  'dEQP-VK.spirv_assembly.instruction.compute.op*dotaccsatkhr.*' \
  --deqp-log-filename=iadd_sat.qpa \
  --deqp-log-images=disable
Passed: 996/996
Failed: 0

Is .SIGN really an instruction-wide switch?

Then I read Faith's proposed direction and Máté's work-in-progress. They treated .SIGN as a modifier on each source, not one switch for the whole instruction. That made my first model look rather narrow. Fair enough; it was built for one operation.

With three sources, each predicate (P) or sign bit (S), there are eight logical layouts. I knew the all-predicate and all-sign forms. The mixed ones were the problem. Did the hardware really have a separate encoding for every layout?

I kept a known instruction fixed, varied its 12-bit opcode field through all 4,096 possibilities, and kept the candidates that nvdisasm decoded as PLOP3. Brute force, but over twelve bits it is the polite kind. 

PPP  PPS  PSP  PSS  SPP  SPS  SSP  SSS

The scan only changed those 12 bits, so it could tell me about decoded opcode forms for that template—not every possible encoding or its runtime semantics. 

0x81c        PPP: predicate-only PLOP3
0x21d        PSP: one .SIGN source
0x21e        PSS: two .SIGN sources
0x21f        SSS: three .SIGN sources
0xa1d–0xa1f  constant-buffer source variants
0x89c        uniform UPLOP3

The scanner kept the other 116 bits fixed, replaced bits 0–11, and retained anything nvdisasm recognized as PLOP3: 

let base = [0x0200721f_u32, 0x00000003,
            0x00700204,     0x000fc000];

for opcode in 0_u32..0x1000 {
    let code = [
        (base[0] & !0xfff) | opcode,
        base[1], base[2], base[3],
    ];
    // Write code as little-endian bytes, then:
    nvdisasm -b SM89 --print-raw candidate.bin
}

For example, the controlled template changes like this:

logical form   four 32-bit words                         low 12 bits
PPP            0200781c 00000003 00700204 000fc000       0x81c
SSS            0200721f 00000003 00700204 000fc000       0x21f
                       ^ the scanned field is here

The scan gave me decoded forms, not execution semantics. Time to put them on the AD107.

Only four useful canonical forms turned up, not eight. The other logical layouts could be reordered into those forms. There's a catch: PLOP3 is a truth table, so swapping sources without permuting the LUT changes the function. A perfectly valid instruction can compute the wrong Boolean expression.

For a concrete reorder, suppose the logical operation is SPP(a, b, c, lut): a is a GPR sign bit, while b and c are predicates. The hardware form puts the one sign source in the middle slot: 

logical:   SPP(a, b, c, lut)
physical: PSP(b, a, c, permute_xy(lut))

Only swapping a and b would make LUT input x receive the old y. Applying the same permutation to the LUT preserves all eight rows of the original boolean function.

Now make the GPU check our reordering

I supplied Booleans as predicates or as GPR sign bits using 0 and INT_MIN, checked all eight layouts, and chose asymmetric LUTs (0x1b and 0xe4) so a mistaken swap wouldn't slip past. That's 128 result checks. 

8 source layouts × 8 truth-table rows × 2 LUTs = 128 result checks

I also tried persuading ptxas to generate mixed-source PLOP3 itself. It kept getting in the way, optimizing simple probes into ISETP. Adding PTX vote got me a PLOP3, but ptxas materialized the sign as a predicate first. Still no mixed-source example. 

sx = (x < 0)
out = LUT(sx, p, q)
sign test + predicate logic  → ISETP
opaque predicates via vote  → predicate materialization → PPP PLOP3

I didn't preserve those exact exploratory PTX files, so I'm deliberately showing the idea rather than pretending these are reproducible original inputs. And this tells us only what ptxas chose for my probes, not which mixed layouts the hardware supports. The decoder experiment and hardware truth-table test are the stronger evidence for that question.

Where I stopped

The narrow iadd_sat path passed the AD107 hardware test, decoder test, and 996 CTS cases. I also tried a Fossilize corpus: no shader-output changes, because none of those shaders contained a surviving nir_op_iadd_sat. Useful coverage information, not a performance result.

The broader per-source model isn't ready to merge. Its 128 truth-table checks passed, but two of 62 NAK tests still failed. I posted the findings on the issue and asked the maintainers about the representation and patch split. No MR yet; I'd rather resolve those failures than present a tidy ending I haven't earned.

References