Softmax lowering
RVV has no vector exponential: exp becomes a scalar libm call per element, so softmax leaves the vector path at exactly that op.
One op, five
linalg.softmax is documented as an aggregate operation that reduces to a small DAG of structured ops: the max along the chosen dimension, the subtraction, the exponential, the sum along the same dimension, the division. On this path the DAG becomes vector.multi_reduction maxnumf, arith.subf, math.exp, vector.multi_reduction add and arith.divf, and the first and last of those are the machinery from the previous node.
Tile the row, not the row’s columns
Tiling the row dimension is free. Tiling the softmax dimension is not: the maximum has to cover the whole row, so a tile on that dimension normalises a slice and returns a result that is finite, sums to one, and answers a different question. The op’s own documentation warns that the tiling implementation does not check the dimension against the semantics, so the check is yours to make.
The softmax dimension stays whole because the maximum needs the whole row, not because the row fits in a register group. At VLEN=256 an m4 group holds 32 f32, so a 1024-column row is 32 groups and is streamed twice: once for the max, once for the exponential and the sum.
The exponential is the hole in the path
RVV has no exponential. math.exp lowers to llvm.exp, and the RISC-V backend expands that into a libcall: one scalar expf per element, with the vector unit idle while the call runs. Two ways out. Fold log2(e) into the logits and use math.exp2, or write the polynomial yourself, a reduced-range minimax with vfmacc. Both change the bits, so the check against the scalar reference is a tolerance on the softmax output rather than equality on intermediates.
Things that break
Tiling the softmax dimension, which fails silently and leaves a plausible table behind. A benchmark that charges the libm calls to no one, so the vector kernel looks slower for a reason the writeup never names. Dividing per element instead of multiplying by one reciprocal. A fully masked row producing NaN, which the reduction section already covers and which fusion makes easier to hit.
Sources
- MLIR linalg dialect documentationdocs · LLVM 23
- MLIR math dialect documentationdocs · LLVM 23
- User Guide for RISC-V Targetdocs · LLVM 23
A version records the documentation the claim was checked against, not a house pin. Confirm flags with --help on the release you installed.
Check yourself
`linalg.softmax` is one op. What does it lower to?
Which dimension can be tiled without changing the result?
Why does the vector softmax spend most of its time outside the vector unit?