TRITON LAB #02 — triton-shared's output on a CPU: 0.6 GB/s → 17 GB/s, found with a profiler

TRACK: triton | DIFFICULTY: 4 | NEEDS: mlir-opt, clang, perf, qemu FILES: vec_add.linalg.mlir (triton-shared shape), vec_add.inline_copy.mlir (B), vec_add.punted.mlir (C), launcher.c

What this teaches

This is the lab closest to your actual job. It takes Triton's vec_add from TTIR (lab 01) in the form triton-shared produces, runs it through a CPU pipeline with a launcher that behaves like the real one, profiles it, and fixes what the profile shows. Correctness is checked on x86 and AArch64.

Honest scope note

Building triton-shared itself (it pins a specific Triton + LLVM) takes hours and isn't done here. Instead, vec_add.linalg.mlir is hand-written in the exact shape triton-shared emits, checked against the real triton-shared output that a vendor MLIR backend ships as a test input: triton-shared’s own lit test inputs (unranked memref<*xf32> args, 6 trailing i32 grid args, reinterpret_cast views, to_tensor → linalg.generic → materialize_in_destination). The masked-load pattern (alloc + partial memref.copy) is how triton-shared handles tl.load(..., mask=...). Once you have a production MLIR backend access, regenerate this file from its real ttsharedir stage (MLIR_ENABLE_DUMP=1) and rerun everything.

STEP 1 — The calling convention is the launcher's contract

$ mlir-opt vec_add.linalg.mlir --pass-pipeline="$LOWER_STD" | mlir-translate --mlir-to-llvmir -o out/vec_add.ll && grep '^define' out/vec_add.ll
define void @vec_add(i64 %0, ptr %1, i64 %2, ptr %3, i64 %4, ptr %5, i32 %6, i32 %7, i32 %8, i32 %9, i32 %10, i32 %11, i32 %12) {

Each memref<*xf32> is (i64 rank, ptr descriptor), then n, then num_programs x,y,z, then program_id x,y,z. That's exactly the signature a vendor MLIR backend's launcher calls (the generated Triton launcher, the num_programs_X ... pid_Z closure fields) and launcher.c reproduces: build descriptors, then loop over the grid calling the kernel once per program.

STEP 2 — Measure it

clang -O2 -march=native -c out/vec_add.ll -o out/vec_add.o -Wno-override-module
clang -O2 launcher.c out/vec_add.o -o out/va_shared -L$MLIR_LIB -lmlir_c_runner_utils -Wl,-rpath,$MLIR_LIB
../tools/pin.sh -q ./out/va_shared 10 triton-shared-shape

triton-shared-shape    grid=16385  min  331.61 ms  p50  443.78 ms    0.6 GB/s  OK

0.6 GB/s for an add. Lab 04/02's fused kernel does 16 GB/s. (Note the first link attempt fails with undefined reference to 'memrefCopy'. That symbol is the clue.)

STEP 3 — Profile, don't guess

perf record -q -o out/perf.data ./out/va_shared 5 x >/dev/null
perf report -i out/perf.data --stdio | grep -E '^\s+[0-9]' | head -6

    59.37%  va_shared  libmlir_c_runner_utils.so.24.0git  [.] memrefCopy
    29.42%  va_shared  libc.so.6                          [.] __memmove_avx_unaligned_erms
     3.95%  va_shared  [unknown]                          [k] 0xffffffff9f001280
     1.58%  va_shared  libmlir_c_runner_utils.so.24.0git  [.] memcpy@plt
     1.29%  va_shared  libc.so.6                          [.] _int_malloc
     1.02%  va_shared  va_shared                          [.] vec_add

The kernel is 1% of the time. 90% is copying:

Each masked tl.load became memref.alloc + memref.copy of the valid part. The copy is on memref<?xf32, strided<[1], offset: ?>> (dynamic size, dynamic offset), which finalize-memref-to-llvm can't prove contiguous, so it emits a call to the generic runtime memrefCopy. That function walks unranked descriptors and calls memcpy per element (the 29% memmove). Read it: $LLVM_SRC/mlir/lib/ExecutionEngine/CRunnerUtils.cpp, the loop at line ~107 (// Copy over the element, byte by byte.).
It happens 3× per program (2 loads + the store), 16 385 programs per call.

STEP 4 — Fix B: make the copies inline loops

Rewrite memref.copy → linalg.copy, which lowers to a loop LLVM vectorizes:

sed -E 's/memref\.copy (%[a-z0-9_]+), (%[a-z0-9_]+) : (.*) to (.*)$/linalg.copy ins(\1 : \3) outs(\2 : \4)/' vec_add.linalg.mlir > vec_add.inline_copy.mlir

inline_copy            grid=16385  min  152.85 ms  p50  359.71 ms    1.3 GB/s  OK

About 2×, but memrefCopy is still at the top of the profile (45%). Who put it back?

$ ../tools/who_did_it.sh 'memref.copy' mlir-opt vec_add.inline_copy.mlir --pass-pipeline='builtin.module(one-shot-bufferize{bufferize-function-boundaries function-boundary-type-conversion=identity-layout-map}, func.func(convert-linalg-to-loops), finalize-memref-to-llvm)'
order    after pass                                                   matches
0        (input)                                                      0
1        builtin_module_no-symbol-name/0_one-shot-bufferize.mlir      1
2        func_func_vec_add/0_0_convert-linalg-to-loops.mlir           1
3        builtin_module_no-symbol-name/1_finalize-memref-to-llvm.mlir 0

one-shot-bufferize creates a new memref.copy when it bufferizes the masked store (materialize_in_destination into a dynamic-length view). Our sed fixed the loads; the store's copy is born later, inside bufferization. (One more reason to check the IR after the pass you think you fixed.)

STEP 5 — Fix C: don't copy full blocks at all

All but one of the 16 385 programs have a full block (len == 1024); only the last is partial. For a full block the mask is all-true, so the temporaries are pointless: compute straight from the input views into the output view. vec_add.punted.mlir adds that fast path (scf.if %len == 1024) and keeps the triton-shared shape only for the tail. That's the idea behind a vendor MLIR backend's a copy-forwarding pass ("erase copies and forward buffer directly to user").

triton-shared-shape    grid=16385  min  331.61 ms  p50  443.78 ms    0.6 GB/s  OK
inline_copy            grid=16385  min  152.85 ms  p50  359.71 ms    1.3 GB/s  OK
punted                 grid=16385  min   11.54 ms  p50   12.36 ms   17.4 GB/s  OK

29×, same results, and now at the same bandwidth as the ideal fused kernel from lab 04/02. Page faults tell the same story (/usr/bin/time -f %R: 182 k for inline_copy vs 51 k for punted: the per-program temporaries are gone).

STEP 6 — Same three variants on AArch64 (correctness under qemu)

for v in "" .inline_copy .punted; do
  clang --target=aarch64-linux-gnu -O2 -mcpu=neoverse-v1 -c out/vec_add$v.ll -o out/a64$v.o -Wno-override-module
  a64cc -O2 launcher.c out/a64$v.o ../tools/a64_runtime.c -o out/va_a64$v -lm
  a64run ./out/va_a64$v 1 "aarch64${v:-.shared}" | grep -oE '^\S+|grid=\S+|OK|WRONG' | paste -sd' '
done

aarch64.shared grid=16385 OK
aarch64.inline_copy grid=16385 OK
aarch64.punted grid=16385 OK

(tools/a64_runtime.c supplies memrefCopy for AArch64, since the real runtime library is C++ and the sysroot has no libstdc++.)

What this means for a production MLIR backend

This is what a "per-kernel optimization plan" item looks like:

Kernel: vec_add (and every masked elementwise kernel) Symptom: 0.6 GB/s; 90% of time in memrefCopy / memmove. Cause: masked tl.load/tl.store → alloc + dynamic-size memref.copy per program, lowered to the generic runtime copy. Fix candidates: (1) full-block fast path (29×, measured); (2) lower dynamic contiguous copies inline; (3) masked vector loads (vector.maskedload, SVE predication) instead of copies. Test: a FileCheck test that no memrefCopy call remains for full blocks (lab 06/03).

Try this yourself

Add a 4th variant: replace the tail copies with vector.transfer_read ... mask (masked vector loads). This is the SVE-friendly answer to masks.
Change BLOCK to 4096 (in .mlir and launcher.c). How does the triton-shared shape scale? Why? (Per-program overhead ÷ 4.)
Point tools/ablate.py at vec_add.inline_copy.mlir with the passes you add while trying (1). Which one is your X?