How to encode JPEG in a Vulkan compute shader — parallel Huffman via restart markers

Post 51. JPEG encoding looks serial — the Huffman bitstream chains every block — until you notice DRI restart markers: declaring “restart every MCU” resets the DC predictor and the bitstream at each 16×16 macroblock, so each MCU encodes completely independently. That makes baseline JPEG embarrassingly parallel: one compute invocation per MCU, each writing a byte-aligned segment into its own SSBO slot. The host just concatenates the segments with RSTm markers between them.

$ ./app out.jpg
device: AMD Radeon RX 7800 XT (RADV NAVI32)
wrote out.jpg (25443 bytes, 1200 MCUs)

640×480 → a valid JFIF that any decoder opens — pixels never left the GPU until the final SSBO readback of the compressed bytes (~25 KB, not ~900 KB of RGBA).

The trick that makes it parallel

Baseline JPEG’s serial dependencies: DC coefficients are delta-encoded block-to-block and the Huffman bitstream is a single continuous stream. Insert FF DD (DRI, “restart interval = 1”) in the header, and the decoder expects FF D0+i%8 markers between MCUs — each MCU’s bitstream restarts byte-aligned with pred = 0. Now every work item is independent:

layout(local_size_x = 64) in;
void main() {
    uint mcu = gl_GlobalInvocationID.x;
    if (mcu >= pc.total) return;
    segBase = pc.total + mcu * pc.segWords;   // this MCU's slot
    // ... encode 4 luma + Cb + Cr blocks into buf[segBase..]
    buf[mcu] = byteCnt;                        // segment byte count
}

Per-MCU work: fetchC/fetchY (RGBA→YCbCr on the fly), float DCT-II as two separable 8-point transforms (CST[u*8+x] = cos((2x+1)uπ/16) baked in by the host string-format), quantize, zigzag, canonical Huffman with 0xFF-stuffing — all straight-line shader code.

Host side: JFIF skeleton + table upload

The shader needs six tables in a read-only SSBO — the two quality- scaled quant tables (natural order; the file’s DQT is zigzag) and four Huffman code tables as len<<16 | code words built from Annex K bits+vals:

scaleQ(kLumaQ, 85, tab);          // quality 1..100, Annex-K formula
buildHuff(kDcLB, kDcLV, 12, tab + 128); // canonical codes on CPU once

Assembly is mechanical: SOI, APP0/JFIF, DQT×2, SOF0 (0x22 sampling on Y), DHT×4, DRI=1, SOS, then for each MCU segment bytes + RSTm marker, then EOI. Byte counts come from buf[mcu] — the first total words of the SSBO are the size array.

for (uint32_t i = 0; i < total; ++i) {
    auto* seg = (const uint8_t*)(words + total + i * kSegWords);
    j.insert(j.end(), seg, seg + std::min(words[i], kSegWords*4));
    if (i + 1 < total) { u8(j,0xff); u8(j,0xd0 + (i % 8)); }
}

Why bother?

• Latency: the frame’s pixels are already in a VkImage (decoding happened on-GPU, see post 49). Downloading them to libjpeg costs a PCIe roundtrip; the compute path reads back ~30 KB of JPEG instead.
• VAAPI can’t encode JPEG — VCN has no JPEG encode entrypoint, decode only. This shader fills the gap (it’s how a camera app can serve a /api/screenshot endpoint and per-camera stills).
• Throughput: 1200 MCUs in one dispatch — a dGPU finishes in tens of µs.

Gotchas

• JFIF APP0 has a literal NUL — std::string_view("JFIF\0") truncates at the NUL; emit the byte explicitly.
• Bit-order: Huffman codes go MSB-first, bytes accumulate LSB-first into u32 words, and every emitted 0xFF data byte must be followed by a stuffed 0x00 (that’s what makes FF Dn markers unambiguous).
• DC predictor resets per MCU — don’t forget pred = 0 or the image decodes with banding.
• Pad each MCU’s segment buffer generously (16 KB here); a pathological MCU at quality 100 can outgrow tight estimates — the shader writes min(byteCnt, segBytes) and the size word tells the truth either way.
• 4:2:0 only — Y gets 2×2 sampling in SOF0; if you need 4:4:4 the MCU becomes one block per component (different loop, same idea).

The full program

Example: VulkanExamples/p51_gpujpeg