Thinking in lanes: SIMD in Mojo 1.0

SIMD has a reputation for being a specialist performance technique. The machinery can become specialised, but the first useful idea is much smaller: stop asking what happens to one value and ask what happens to a packet of values.

This is a Mojo companion to Mitchell Hashimoto's Everyone Should Know SIMD. His article is the foundation: it explains SIMD through a real Zig and Ghostty example, then reduces the technique to a five-step pattern. I am not repeating that article here. Read it for the original argument and example; use this companion to make the lanes visible and see how the same way of thinking is expressed in Mojo 1.0.

What I want to add is a runnable Mojo 1.0 example, then connect that example to the choices in the ray tracer: scalar CPU code, explicit CPU SIMD, and a GPU megakernel.

Start with the data, not the loop

A scalar loop makes one value feel like the natural unit of work. That is useful, but it also hides a question: are the neighbouring values independent, and do we perform the same operation on them?

For this example the answer is yes. Changing the brightness of pixel seven does not require pixel six or pixel eight. Every output depends on one input pixel and two constants. That independence is what makes the operation a clean fit for CPU SIMD. It is also what would let a GPU assign one pixel to one thread.

Begin with a precise image operation

We will process one scanline from a grayscale image. Each pixel is a Float32 value in the inclusive range 0.0...1.0, where 0.0 is black and 1.0 is display white.

The job has three parts:

1. Increase each pixel's exposure by 50%.
2. Clamp values above display white back to 1.0.
3. Preserve the number and order of the pixels.

For a single pixel, the mathematical operation is:

output_pixel = min(input_pixel * 1.5, 1.0)

This is not meant to be a complete colour-management or tone-mapping pipeline. It is a small, real image operation with a result we can inspect. More importantly, there is no dependency between pixels.

Establish the scalar reference first

Before introducing lanes, write the operation for one pixel:

exposed = pixel * 1.5
if exposed > 1.0:
    return 1.0
return exposed

The Mojo version says exactly that:

23def expose_pixel(pixel: Float32) -> Float32:24    """Increase one normalised pixel value and clamp it to display white."""25    var exposed = pixel * 1.526    return 1.0 if exposed > 1.0 else exposed

This function is our statement of correctness. It is also the fallback for an unsupported target and the implementation we use for pixels left over at the end. SIMD does not replace this reasoning. It changes how many independent pixels we present to the CPU at once.

Change the unit of work from a pixel to a packet

Mojo represents a packet as SIMD[dtype, width]. Each position in the packet is a lane. In this example, each lane owns one pixel. Arithmetic applies to all lanes and a comparison returns one Boolean per lane.

simd_width_of[DType.float32]() asks for the native Float32 width of the target CPU. It is four lanes on the M1 Ultra used here. A different target can select a different width from the same source.

The input packet below contains four neighbouring image pixels. They are adjacent in memory, but their calculations are independent.

This is the perspective shift. The scalar question is, "What happens to this pixel?" The SIMD question is, "Can I perform the same operation on this packet of pixels?"

Mitchell's five steps, expressed in Mojo

Mitchell's article makes the common structure explicit: prepare the vectors, loop over full packets, perform the SIMD operation, reduce or store the result, and finish with a scalar tail. We will follow those same steps, but apply them to image data and Mojo.

His Ghostty example scans code points and needs to locate where a printable run ends. Our image operation needs one transformed output for every input pixel. That difference matters in step four: Mitchell reduces information from the lanes to find a position; we store all of our lanes because every output pixel matters.

Step 1: choose the lane type and prepare the constants

We choose Float32 because that is how our normalised pixels are represented. We ask Mojo for the CPU's native lane count rather than hard-coding four.

pixel_type = Float32
lane_count = native SIMD width for Float32
exposure   = 1.5 in every lane
white      = 1.0 in every lane

In Mojo, scalar constants used with a SIMD value are applied lane by lane. For select, we explicitly construct a SIMD value containing 1.0 in every lane.

Step 2: walk the image one complete packet at a time

The scanline is a flat, contiguous list of pixels. The vector loop runs only while at least LANES pixels remain. It then loads adjacent pixels into the packet.

while cursor + lane_count <= pixel_count:
    packet = load pixels[cursor .. cursor + lane_count]

72    # Process every complete packet. Each lane owns one independent pixel.73    while cursor + Int(LANES) <= count:74        var packet = SIMD[DTYPE, LANES](0.0)75        for lane in range(Int(LANES)):76            packet[lane] = input_pixels[cursor + lane]

This teaching programme loads the lanes one at a time so the mapping is visible. A production implementation would normally use a contiguous vector load after verifying its alignment and memory assumptions.

Step 3: perform the image operation across every lane

Now the scalar operation becomes a packet operation:

exposed    = pixel_lanes * 1.5
over_white = exposed > 1.0
result     = select(over_white, 1.0, exposed)

29def expose_packet[30    width: SIMDLength,31](values: SIMD[DTYPE, width]) -> SIMD[DTYPE, width]:32    """Apply the same image operation to every pixel lane."""33    var exposed = values * 1.534    var over_white = exposed.gt(1.0)35    return over_white.select(SIMD[DTYPE, width](1.0), exposed)

The multiply brightens all pixels. The comparison does not produce one answer for the packet. It produces a Boolean mask with one answer per lane. over_white.select(white, exposed) chooses white in the true lanes and the exposed value in the false lanes.

For the sample packet, 0.75 becomes 1.125 and 1.00 becomes 1.50. Those lanes are above display white, so they become 1.00. The 0.25 and 0.50 pixels become 0.375 and 0.75 and remain there.

The programme prints those intermediate values directly:

38def demonstrate_one_register():39    print("1. One SIMD value contains several independent image pixels")40    print("   Native Float32 lane count:", Int(LANES))41 42    var input = SIMD[DTYPE, LANES](0.0)43    for lane in range(Int(LANES)):44        input[lane] = Float32(lane + 1) * 0.2545 46    var exposed = input * 1.547    var over_white = exposed.gt(1.0)48    var output = over_white.select(SIMD[DTYPE, LANES](1.0), exposed)49 50    show_packet("   input pixels:  ", input)51    show_packet("   after exposure:", exposed)52    print("   over white?:   ", over_white)53    show_packet("   output pixels: ", output)54    print()

Step 4: store each lane in its output pixel

This algorithm does not need a reduction. There is one output pixel for every input pixel, so each result lane is stored at the corresponding image position.

for lane in packet:
    output_pixels[cursor + lane] = result[lane]
cursor += lane_count

83        # Store each transformed lane in its corresponding output pixel.84        for lane in range(Int(LANES)):85            output_pixels[cursor + lane] = result[lane]

The distinction between storing and reducing is worth holding onto. A transform such as exposure usually stores every lane. A sum, count, or search may need to combine lane results into one scalar answer.

Step 5: finish the incomplete packet with scalar code

An image width will not always be an exact multiple of the native SIMD width. If three pixels remain and the CPU packet holds four, we do not read a fourth pixel that does not exist. We finish those three with the scalar function we started with.

while cursor < pixel_count:
    output_pixels[cursor] = expose_pixel(input_pixels[cursor])
    cursor += 1

 90    # Fewer than LANES pixels remain, so finish with the scalar reference. 91    print("   scalar tail starts at index", cursor) 92    while cursor < count: 93        output_pixels[cursor] = expose_pixel(input_pixels[cursor]) 94        print( 95            "      index", 96            cursor, 97            ":", 98            input_pixels[cursor], 99            "->",100            output_pixels[cursor],101        )102        cursor += 1

The scalar tail is not a failure of SIMD. It is the simplest correct way to finish an incomplete packet, and it keeps the scalar reference exercised.

Put the steps together

The programme processes two complete packets and deliberately leaves three pixels for the tail. Run it with pixi run simd-demo:

SIMD in Mojo 1.0
================
One instruction, several values. Each value occupies one lane.
Example: raise normalised brightness by 50%, then clamp to white (1.0).


One SIMD value contains several independent image pixels
Native Float32 lane count: 4
input pixels:   [0.25, 0.5, 0.75, 1.0]
after exposure: [0.375, 0.75, 1.125, 1.5]
over white?:    [False, False, True, True]
output pixels:  [0.375, 0.75, 1.0, 1.0]


Process one grayscale image scanline in SIMD packets
packet 0 covers indices 0 to 3
input:   [0.125, 0.5, 0.875, 0.25]
output:  [0.1875, 0.75, 1.0, 0.375]
packet 1 covers indices 4 to 7
input:   [0.625, 1.0, 0.375, 0.75]
output:  [0.9375, 1.0, 0.5625, 1.0]
scalar tail starts at index 8
index 8 : 0.125 -> 0.1875
index 9 : 0.5 -> 0.75
index 10 : 0.875 -> 1.0
input scanline:  [0.125, 0.5, 0.875, 0.25, 0.625, 1.0, 0.375, 0.75, 0.125, 0.5, 0.875]
output scanline: [0.1875, 0.75, 1.0, 0.375, 0.9375, 1.0, 0.5625, 1.0, 0.1875, 0.75, 1.0]


The perspective shift
Scalar question: what happens to this value?
SIMD question:  what happens to this packet of values?
Masks preserve lane-by-lane decisions without splitting the packet.

The lane-by-lane loads and stores are intentionally obvious here. The point is to expose the method before optimising the memory operations. Once the result is correct, the next question is whether a profile justifies moving to contiguous loads and stores.

This pixel operation is unusually regular. A blur is different because each output reads neighbouring pixels. It may still benefit from SIMD, but the data access, edge handling and lane boundaries need more thought. Independence makes the first example clear; it is not a claim that every image algorithm is equally simple.

When a GPU megakernel is enough

The ray tracer maps one pixel to one GPU thread. Each thread follows one path through its samples and bounces, then writes one independent output. The GPU groups those scalar threads in hardware.

For this scene, a megakernel is enough. The state fits with the thread, the scene is read-only, paths do not exchange partial results, and one launch can carry a path to completion without intermediate queues.

Adding explicit multi-lane CPU-style SIMD inside each GPU thread would usually make that kernel worse. It would increase per-thread state and register pressure while duplicating parallelism the GPU already supplies across threads.

That changes when the paths become too divergent or their state becomes too large. A wavefront renderer can split intersection, materials and shadow work into queues, trading extra memory traffic and launches for more coherent execution.

When CPU SIMD is worth considering

SIMD is a good candidate when a CPU hot loop scans, compares, counts or transforms a large contiguous collection. The operation should repeat often enough for packet setup, reduction and tail handling to disappear into the useful work.

It is also useful when the CPU is the product rather than a fallback: low-latency work, small batches that do not justify a device launch, machines without a usable GPU, or code that must remain close to CPU-owned memory.

SIMD and multicore execution solve different layers. Workers distribute packets across cores. SIMD advances several values inside each worker. Our benchmark showed that the two gains compose, although they do not multiply perfectly.

When to stay away from it

Stay scalar when the loop is cold, the input is tiny, or the compiler already produces good vector code. SIMD is an implementation cost. If the cost cannot be seen in a profile, adding it is difficult to defend.

Irregular pointer chasing, frequent cross-lane dependencies, scattered memory and deeply divergent branches are also warning signs. You may spend more work gathering data and reconciling lanes than the vector operation saves.

The useful rule is not that SIMD is good or bad. It is narrower: use SIMD where the data is regular, the operation is shared, and the amount of work is large enough to measure the difference.

The challenges do not disappear

Once we think in packets, we have to handle packet width, tails, masked lanes, reductions and memory layout. We also have to test the scalar and vector forms against one another.

Floating-point reductions may associate values differently. A four-lane result can be numerically valid without being bit-for-bit identical to a scalar loop. That distinction belongs in the test plan, not as a surprise after optimisation.

Lane width is a compile-time parameter in Mojo. That gives us portable source, but it does not make every width equally efficient. Use the native width as the starting point and benchmark on the actual target.

Where this leaves us

Start with the scalar algorithm because it is the clearest statement of correctness. Then look for the packet. If several independent values receive the same work, make the lanes explicit and preserve the scalar tail.

If the work is already expressed as thousands of independent GPU threads, a megakernel may be the simpler answer. If the work remains on the CPU, SIMD is another layer of parallelism you should at least know how to recognise.

The real skill is not memorising vector syntax. It is learning to see data as groups of independent values, then choosing the execution model that matches those groups without pretending every problem has the same answer.

The complete source is simd_walkthrough.mojo. For the API details, keep the official Mojo SIMD reference nearby.