Chapter 3: Memory Layout Strategies — Array-of-Structs vs Struct-of-Arrays¶
"The question a memory layout answers is never 'does the data fit?' It's 'when I ask for the ten bytes I actually need, how many other bytes does the hardware make me pay for anyway?' Get that answer wrong across a billion elements, and it doesn't matter how clever the algorithm sitting on top of it is."
What you will understand by the end of this chapter:
- Why bulk numerical code is usually limited by how many bytes move between RAM and the CPU, not by how many arithmetic instructions it runs — and how to compute exactly what fraction of those moved bytes were actually useful for a given operation
- The Array-of-Structs (AoS) layout: one struct per object, laid out one after another, and precisely which operations it's the right choice for
- The Struct-of-Arrays (SoA) layout: one contiguous array per field, and why it is what lets a single SIMD instruction (Chapter 1.4) load four, eight, or sixteen values in one shot
- How to hand-trace the exact same computation — a particle system's kinetic energy — through both layouts and confirm they produce the identical number, because the layout is an engineering decision, never a mathematical one
- Why this book's
Tensortype, introduced in Part 1, is built as SoA rather than AoS, traced down to the exact byte-stride difference it makes for a single SIMD load
What you need to know first:
- Chapter 1.4 (SIMD as a first-class type) — this chapter explains why SIMD needs the layout it needs, building directly on the lane model already traced there
- Chapter 2.1 and 2.4 (struct field layout,
UnsafePointer,__init__/__del__) — AoS and SoA are both just particular ways of arranging structs and heap arrays that Chapter 2 already introduced individually
3.1 The Memory Bus: Every Byte You Move Costs Bandwidth, Whether You Use It or Not [FOUNDATIONAL]¶
Intuition¶
Imagine hiring movers who only ever carry pre-packed boxes, never individual items. If your kitchen box has forks, plates, and pots all packed together, and you only actually need the forks at the new place, the movers still have to carry the whole box — the plates and pots ride along whether you wanted them there or not, and you pay for their weight regardless. If instead everything had been pre-sorted into a "forks box," a "plates box," and a "pots box," the movers would only need to carry the one box you actually asked for.
A struct laid out in memory is that mixed kitchen box: reading one field out of it means the surrounding fields — packed into the very same contiguous region — ride along for the trip from RAM to the CPU whether the computation needs them or not. This isn't a matter of the program being inefficient; it's a physical fact about how memory is organized when fields are interleaved.
Background¶
Modern CPUs move data between RAM and the processor over a memory bus, and for numerical code that scans large arrays — exactly the kind of code an automatic differentiation framework spends most of its time running — the speed limit is usually how many bytes cross that bus, not how many arithmetic instructions execute once the data has arrived. This is the defining fact of high-performance numerical computing: such code is typically memory-bandwidth-bound, not compute-bound.
Define bus utilization for a given operation as the fraction of bytes actually moved across the bus that the computation actually uses:
A utilization of 100% means every byte that crossed the bus was used; a utilization of, say, 50% means the hardware moved twice as much data as the computation could use, because unneeded bytes were interleaved with needed ones and couldn't be skipped without extra machinery.
Worked Example 3.1.1 — total_kinetic_energy on an Array-of-Structs¶
Take the Particle struct this chapter's code builds around:
struct Particle:
var x: Float32 # offset 0, 4 bytes
var y: Float32 # offset 4, 4 bytes
var z: Float32 # offset 8, 4 bytes
var vx: Float32 # offset 12, 4 bytes
var vy: Float32 # offset 16, 4 bytes
var vz: Float32 # offset 20, 4 bytes
var mass: Float32 # offset 24, 4 bytes
# total size: 28 bytes
Following Chapter 2.1's fixed-offset reasoning directly: seven Float32 fields at 4 bytes each give a fixed instance size of 7 × 4 = 28 bytes, with x at offset 0 and mass at offset 24. Now trace total_kinetic_energy(), which reads only vx, vy, vz, and mass for every particle — 4 of the 7 fields, 16 of the 28 bytes per particle.
Because those 16 useful bytes are interleaved with the 12 bytes of x, y, z inside every single particle record, there is no way to stream past the array collecting only vx/vy/vz/mass without physically passing over the x/y/z bytes sitting between them — they're part of the same contiguous 28-byte block. For an array of N = 1000 particles:
total bytes moved = 1000 particles × 28 bytes/particle = 28,000 bytes
bytes actually used = 1000 particles × 16 bytes/particle = 16,000 bytes
bus utilization = 16,000 / 28,000 ≈ 57.1%
Nearly 43% of every byte that crossed the memory bus for this computation was x, y, or z data the kinetic-energy formula never touches.
Worked Example 3.1.2 — The same computation on a Struct-of-Arrays¶
Now trace the SoA version, total_kinetic_energy_simd, where vx, vy, vz, and mass each live in their own separate, tightly-packed array:
var vx: UnsafePointer[Float32] # 1000 x 4 bytes = 4,000 bytes, contiguous
var vy: UnsafePointer[Float32] # 1000 x 4 bytes = 4,000 bytes, contiguous
var vz: UnsafePointer[Float32] # 1000 x 4 bytes = 4,000 bytes, contiguous
var mass: UnsafePointer[Float32] # 1000 x 4 bytes = 4,000 bytes, contiguous
# x, y, z live in three further separate arrays -- never allocated near
# vx/vy/vz/mass, and never touched by this computation at all
total bytes moved = 4 arrays x 4,000 bytes = 16,000 bytes
bytes actually used = 16,000 bytes (identical to Example 3.1.1)
bus utilization = 16,000 / 16,000 = 100%
Same formula, same 1000 particles, same 16,000 bytes of genuinely useful data — but the SoA version moves exactly that much and not one byte more, because x, y, and z sit in entirely separate allocations that this computation simply never opens.
ASCII Diagram — interleaved vs. separated¶
AoS, one particle's 28 bytes (x,y,z ride along unused):
[ x ][ y ][ z ][ vx ][ vy ][ vz ][ mass ]
used? no no no YES YES YES YES <- 16 of 28 bytes useful
SoA, the same particle's data spread across separate arrays:
vx array: [ vx0 ][ vx1 ][ vx2 ]... <- only useful bytes, contiguous
vy array: [ vy0 ][ vy1 ][ vy2 ]... <- only useful bytes, contiguous
vz array: [ vz0 ][ vz1 ][ vz2 ]... <- only useful bytes, contiguous
mass array: [ m0 ][ m1 ][ m2 ]... <- only useful bytes, contiguous
x, y, z arrays: never opened by this computation at all
[COMMON TRAP]It's tempting to think "the code just doesn't referenceself.particles[i].x, so it's free" — as if not naming a field in an expression means the hardware skips its bytes. It doesn't: in AoS,x,y, andzare physically wedged between the fields the loop does want, in the very same 28-byte block, so the memory system has no way to leave them behind without extra, expensive gather-style instructions. "Unused" and "not physically fetched" are only the same thing when the layout keeps unused data in a separate place — which is precisely what SoA does and AoS does not.
3.2 Array-of-Structs: The Object-Oriented Default [FOUNDATIONAL]¶
Intuition¶
A filing cabinet organized by employee — one folder per person, holding that person's ID copy, timesheets, and pay stubs all together — is exactly the right organization for "pull everything about employee #42." One folder, one trip to the cabinet, everything needed is right there. It is the wrong organization for "list every employee's current pay stub," which now means opening all forty folders one at a time.
AoS is the per-employee folder: ParticleSystemAoS allocates one contiguous array where each element is a complete Particle — all seven fields together, one trip to memory for the whole thing.
Background¶
struct ParticleSystemAoS:
var particles: UnsafePointer[Particle]
var count: Int
fn __init__(out self, count: Int):
self.count = count
self.particles = UnsafePointer[Particle].alloc(count)
# ...
UnsafePointer[Particle].alloc(count) reserves count × 28 contiguous bytes, exactly as Chapter 2.4's DynamicArray[dtype] reserved count × 4 bytes for int32 — the only difference is that each "element" here is itself a 28-byte struct rather than a single 4-byte scalar. self.particles[i] is Chapter 1.5's address arithmetic again, just with a wider stride: it means "the 28 bytes starting at base + i × 28," and self.particles[i].x adds Chapter 2.1's fixed-field-offset reasoning on top of that: "the 4 bytes starting at base + i × 28 + 0."
| Array-of-Structs | Struct-of-Arrays | |
|---|---|---|
| One element's fields | Contiguous with each other | Scattered across separate arrays |
| Same field across all elements | Scattered (strided by struct size) | Contiguous with each other |
| Best for | Operations touching most/all fields of one object | Operations touching one field across many objects |
Worked Example 3.2.1 — update_position on a single particle¶
update_position reads and writes x, y, z, vx, vy, vz — six of the seven fields, 24 of the 28 bytes (only mass, 4 bytes, rides along unused this time) — a utilization of 24/28 ≈ 85.7%, sharply better than Example 3.1.1's 57.1%, because this operation genuinely needs almost the entire record. This is AoS's home turf: when an operation needs most of an object's fields, "fetch the whole object in one contiguous block" wastes almost nothing.
Trace particle i = 2 from the initializer's pattern (x = i×0.1, y = i×0.2, z = i×0.3, and constant vx=1.0, vy=2.0, vz=3.0, mass=1.0 for every particle) through update_position(dt=0.1):
before: x=0.2, y=0.4, z=0.6, vx=1.0, vy=2.0, vz=3.0
x += vx * dt -> 0.2 + 1.0*0.1 = 0.3
y += vy * dt -> 0.4 + 2.0*0.1 = 0.6
z += vz * dt -> 0.6 + 3.0*0.1 = 0.9
after: x=0.3, y=0.6, z=0.9
ASCII Diagram — one struct, one contiguous fetch¶
particles[2], all 28 bytes fetched together in one trip:
+0 +4 +8 +12 +16 +20 +24
[0.2 ][0.4 ][0.6 ][1.0 ][2.0 ][3.0 ][1.0 ]
x y z vx vy vz mass
update_position touches the first 6 fields (24 bytes) --
mass (4 bytes) rode along in the same fetch, unused this time.
3.3 Struct-of-Arrays: The Performance-Optimized Layout [FOUNDATIONAL]¶
Intuition¶
Reorganize that same filing cabinet by document type instead of by employee: one drawer holds every employee's ID, a second drawer holds every timesheet, a third holds every pay stub. "List every employee's current pay stub" is now one trip to one drawer. "Pull everything about employee #42" is now the worse job — seven different drawers, one visit each.
SoA is that reorganized cabinet: ParticleSystemSoA allocates seven separate arrays, one per field, each internally contiguous.
Background¶
struct ParticleSystemSoA:
var x: UnsafePointer[Float32]
var y: UnsafePointer[Float32]
var z: UnsafePointer[Float32]
var vx: UnsafePointer[Float32]
var vy: UnsafePointer[Float32]
var vz: UnsafePointer[Float32]
var mass: UnsafePointer[Float32]
var count: Int
Because every field of the same kind now sits contiguously — all x values back-to-back, all vx values back-to-back — a single SIMD[DType.float32, 4].load(ptr) instruction (Chapter 1.4) can pull four consecutive x values into one vector register in a single memory transaction. That single instruction is only legal because the four values it loads are guaranteed to be neighbors in memory; in AoS, four consecutive x values are 28 bytes apart from each other, not 4, and no single contiguous SIMD load can reach them.
Worked Example 3.3.1 — Tracing the first SIMD group of update_positions_simd¶
var simd_count = (self.count // 4) * 4
var dt_vec = SIMD[DType.float32, 4](dt)
for i in range(0, simd_count, 4):
# loads x[i..i+3] and vx[i..i+3] as two contiguous SIMD vectors
var new_x = x_vals + vx_vals * dt_vec
For the first group (i = 0..3), the initializer's pattern gives x = [0.0, 0.1, 0.2, 0.3] (that's i × 0.1 for i = 0,1,2,3) and vx = [1.0, 1.0, 1.0, 1.0] (constant for every particle). With dt = 0.1, reusing Chapter 1.4's lane-wise multiply-then-add:
lane 0: 0.0 + 1.0*0.1 = 0.1
lane 1: 0.1 + 1.0*0.1 = 0.2
lane 2: 0.2 + 1.0*0.1 = 0.3
lane 3: 0.3 + 1.0*0.1 = 0.4
new_x = [0.1, 0.2, 0.3, 0.4] — and notice lane 2's result, 0.3, is exactly the updated x value Worked Example 3.2.1 computed by hand for particle i=2 through the scalar AoS path. Two completely different memory layouts, two completely different code paths — one identical answer, because both are computing x + vx·dt for the same particle.
The same lane-wise pattern gives y and z for this group: y = [0.0, 0.2, 0.4, 0.6] with constant vy=2.0 gives new_y = [0.2, 0.4, 0.6, 0.8]; z = [0.0, 0.3, 0.6, 0.9] with constant vz=3.0 gives new_z = [0.3, 0.6, 0.9, 1.2] — and new_z's lane 2, 0.9, again matches Worked Example 3.2.1's scalar result for particle i=2 exactly.
ASCII Diagram — one SIMD load, four contiguous values¶
x array (SoA), first 16 bytes:
+0 +4 +8 +12
[0.0 ][0.1 ][0.2 ][0.3 ]
x0 x1 x2 x3
<-- one SIMD[float32,4].load() reaches all four in a single instruction
The same four x-values in AoS would sit at byte offsets 0, 28, 56, 84
inside the particle array -- 28 bytes apart, not 4 -- which is why no
single contiguous SIMD load can reach them there.
3.4 Kinetic Energy: The Same Computation, Two Layouts, One Answer [FOUNDATIONAL]¶
Intuition¶
Two accountants working from differently-organized filing systems — one filing by employee, one filing by document type — must still arrive at exactly the same total on the company's books at the end of the quarter. The filing system is an operational choice about how fast each of them works; it has no power to change what the numbers add up to.
Background¶
Kinetic energy is KE = 0.5 × mass × (vx² + vy² + vz²), computed identically by total_kinetic_energy (AoS, scalar) and total_kinetic_energy_simd (SoA, vectorized) in this chapter's code. Both are implementations of the same formula; the layout changes how the bytes for vx, vy, vz, mass get to the CPU (Section 3.1's bandwidth argument) but cannot change what comes out of the formula once they've arrived.
Worked Example 3.4.1 — Four particles, verified two ways¶
The initializer gives every particle the same velocity and mass (vx=1.0, vy=2.0, vz=3.0, mass=1.0, regardless of i — only position varies with i), so every particle's kinetic energy is identical:
For N = 4 particles, the AoS scalar loop accumulates total = 0.0, then adds 7.0 four times: 0.0 → 7.0 → 14.0 → 21.0 → 28.0. The SoA SIMD path packs all four particles' kinetic energies into one vector, energy_sum = [7.0, 7.0, 7.0, 7.0], and the final reduction sums the four lanes: 7.0 + 7.0 + 7.0 + 7.0 = 28.0. Both paths land on 28.0, as they must — Section 3.1 argued the SoA path moves fewer wasted bytes to get there, never that it computes a different formula. This is the same "trust but verify" discipline Chapter 12.4's gradient_check automates for backward rules: before trusting that a performance-motivated layout change is safe, confirm by hand that it still produces the reference answer.
3.5 Why This Book's Tensor Is SoA [FOUNDATIONAL]¶
Intuition¶
Chapter 1.4 pictured SIMD as four workers stapling in lockstep — but that picture quietly assumed the four sheets each worker needs were already stacked side by side, ready to be handed out in one motion. If instead every fourth sheet in the stack were a stapler-cleaning cloth, no worker could just "grab four sheets" anymore; someone would have to pick through the stack first. SoA is what keeps the sheets pre-sorted so the four-workers-at-once picture stays true.
Background¶
Every Tensor this book builds from Part 1 onward stores its .data (and, when needed, its .grad) as one contiguous SoA-style buffer — every element of one field, packed together — rather than as an array of small per-element structs bundling a value with its own gradient. A hypothetical alternative design makes the cost concrete:
struct Element:
var value: Float32 # offset 0, 4 bytes
var grad: Float32 # offset 4, 4 bytes
# total size: 8 bytes
var elements: UnsafePointer[Element] # AoS: [value,grad][value,grad][value,grad]...
Worked Example 3.5.1 — A SIMD load that can't happen¶
Reuse Chapter 1.4's exact SIMD load, SIMD[DType.float32, 4].load(ptr), and ask where the four values it needs would sit under each layout. Under this book's actual SoA Tensor.data buffer, four consecutive values sit at byte offsets 0, 4, 8, 12 — stride 4, perfectly contiguous, exactly what a single SIMD[float32,4] load requires. Under the hypothetical AoS Element array above, four consecutive value fields sit at byte offsets 0, 8, 16, 24 — stride 8, because each element's own grad field (4 bytes) sits between one value and the next. A contiguous SIMD load physically cannot reach values that are 8 bytes apart while treating them as though they were 4 bytes apart; the hardware would need four separate scalar loads, or a slower gather instruction built exactly for this scattered case — either way, forfeiting the one-instruction, four-values-at-once behavior Chapter 1.4 relied on.
ASCII Diagram — contiguous vs. strided, for the exact same SIMD instruction¶
This book's Tensor.data (SoA), stride 4 bytes -- one SIMD load reaches all 4:
+0 +4 +8 +12
[v0 ][v1 ][v2 ][v3 ]
<---- SIMD[float32,4].load() ---->
Hypothetical Element[] (AoS), stride 8 bytes -- no single load reaches all 4:
+0 +8 +16 +24
[v0 ][g0 ][v1 ][g1 ][v2 ][g2 ][v3 ][g3 ]
^^ wanted (skip) ^^ wanted (skip) ^^ wanted (skip) ^^ wanted
[COMMON TRAP]None of this makes SoA universally "the better layout" — Section 3.2 already showed AoS winning decisively (85.7% vs. a single-field AoS scan's far lower utilization) whenever an operation needs most of one object's fields at once. This book chooses SoA forTensorspecifically because automatic differentiation's dominant access pattern is exactly the opposite: bulk, vectorized operations sweeping across every element's value or gradient at once — Chapter 8's gradient accumulation, Chapter 10's SIMD kernels, Part 5's GPU coalescing — the same pattern Section 3.1 and 3.3 traced by hand fortotal_kinetic_energyandupdate_positions_simd. A framework built for per-object physics simulation with heavy single-particle logic might reasonably choose AoS instead; the right layout always follows from the operations the data actually needs to support, not from a rule that one layout universally wins.
3.6 Complete Runnable Code¶
File: 11_aos_pattern.mojo — Section 3.2¶
Execution: pixi run mojo 11_aos_pattern.mojo
from memory import UnsafePointer
struct Particle:
"""Array of Structs (AoS) pattern - traditional object-oriented approach."""
var x: Float32
var y: Float32
var z: Float32
var vx: Float32
var vy: Float32
var vz: Float32
var mass: Float32
fn __init__(out self, x: Float32, y: Float32, z: Float32,
vx: Float32, vy: Float32, vz: Float32, mass: Float32):
self.x = x
self.y = y
self.z = z
self.vx = vx
self.vy = vy
self.vz = vz
self.mass = mass
fn __copyinit__(out self, existing: Self):
"""Copy constructor required for struct copying."""
self.x = existing.x
self.y = existing.y
self.z = existing.z
self.vx = existing.vx
self.vy = existing.vy
self.vz = existing.vz
self.mass = existing.mass
fn update_position(mut self, dt: Float32):
"""Update particle position based on velocity."""
self.x += self.vx * dt
self.y += self.vy * dt
self.z += self.vz * dt
fn kinetic_energy(self) -> Float32:
"""Calculate kinetic energy of particle."""
var v_squared = self.vx * self.vx + self.vy * self.vy + self.vz * self.vz
return 0.5 * self.mass * v_squared
struct ParticleSystemAoS:
"""Particle system using Array of Structs layout."""
var particles: UnsafePointer[Particle]
var count: Int
fn __init__(out self, count: Int):
self.count = count
self.particles = UnsafePointer[Particle].alloc(count)
# Initialize particles with sample data
for i in range(count):
var x = Float32(i) * 0.1
var y = Float32(i) * 0.2
var z = Float32(i) * 0.3
var vx = Float32(1.0)
var vy = Float32(2.0)
var vz = Float32(3.0)
var mass = Float32(1.0)
self.particles[i] = Particle(x, y, z, vx, vy, vz, mass)
fn __del__(owned self):
self.particles.free()
fn update_all_positions(self, dt: Float32):
"""Update all particle positions - demonstrates scattered memory access."""
print("AoS: Updating", self.count, "particle positions")
for i in range(self.count):
self.particles[i].update_position(dt)
fn total_kinetic_energy(self) -> Float32:
"""Calculate total kinetic energy - demonstrates mixed memory access."""
var total: Float32 = 0.0
for i in range(self.count):
total += self.particles[i].kinetic_energy()
return total
fn print_sample(self, num_samples: Int):
"""Print sample of particles."""
var samples = min(num_samples, self.count)
print("AoS Sample particles:")
for i in range(samples):
# Access fields directly to avoid copying
print(" Particle", i, ": pos(", self.particles[i].x, ",", self.particles[i].y, ",", self.particles[i].z,
") vel(", self.particles[i].vx, ",", self.particles[i].vy, ",", self.particles[i].vz, ") mass:", self.particles[i].mass)
fn min(a: Int, b: Int) -> Int:
return a if a < b else b
fn aos_demo():
"""Demonstrate Array of Structs memory layout."""
print("=== Array of Structs (AoS) Pattern ===")
var system = ParticleSystemAoS(1000)
# Memory layout visualization
print("Memory Layout: [x,y,z,vx,vy,vz,mass][x,y,z,vx,vy,vz,mass][...]")
print("Characteristics:")
print(" - Good: Object-oriented, cache-friendly for single particle operations")
print(" - Bad: Poor vectorization, scattered access for bulk operations")
system.print_sample(3)
# Simulate operations
var dt: Float32 = 0.1
system.update_all_positions(dt)
var total_energy = system.total_kinetic_energy()
print("Total kinetic energy:", total_energy)
fn main():
"""Main function for AoS demonstration."""
aos_demo()
File: 12_soa_pattern.mojo — Section 3.3¶
Execution: pixi run mojo 12_soa_pattern.mojo
from memory import UnsafePointer
struct ParticleSystemSoA:
"""Particle system using Struct of Arrays layout for optimal SIMD performance."""
# Separate arrays for each component (SoA pattern)
var x: UnsafePointer[Float32]
var y: UnsafePointer[Float32]
var z: UnsafePointer[Float32]
var vx: UnsafePointer[Float32]
var vy: UnsafePointer[Float32]
var vz: UnsafePointer[Float32]
var mass: UnsafePointer[Float32]
var count: Int
fn __init__(out self, count: Int):
self.count = count
# Allocate separate contiguous arrays for each component
self.x = UnsafePointer[Float32].alloc(count)
self.y = UnsafePointer[Float32].alloc(count)
self.z = UnsafePointer[Float32].alloc(count)
self.vx = UnsafePointer[Float32].alloc(count)
self.vy = UnsafePointer[Float32].alloc(count)
self.vz = UnsafePointer[Float32].alloc(count)
self.mass = UnsafePointer[Float32].alloc(count)
# Initialize with sample data
for i in range(count):
self.x[i] = Float32(i) * 0.1
self.y[i] = Float32(i) * 0.2
self.z[i] = Float32(i) * 0.3
self.vx[i] = Float32(1.0)
self.vy[i] = Float32(2.0)
self.vz[i] = Float32(3.0)
self.mass[i] = Float32(1.0)
fn __del__(owned self):
self.x.free()
self.y.free()
self.z.free()
self.vx.free()
self.vy.free()
self.vz.free()
self.mass.free()
fn update_all_positions(self, dt: Float32):
"""Update all positions - demonstrates optimal vectorized memory access."""
print("SoA: Updating", self.count, "particle positions with SIMD potential")
# This loop is highly vectorizable because we access contiguous memory
for i in range(self.count):
self.x[i] += self.vx[i] * dt
self.y[i] += self.vy[i] * dt
self.z[i] += self.vz[i] * dt
fn update_positions_simd(self, dt: Float32):
"""Vectorized position update using SIMD operations."""
print("SoA: SIMD-optimized position update")
# Process in chunks of 4 for SIMD operations
var simd_count = (self.count // 4) * 4
var dt_vec = SIMD[DType.float32, 4](dt)
for i in range(0, simd_count, 4):
# Manual load from memory using loop
var x_vals = SIMD[DType.float32, 4](0)
var vx_vals = SIMD[DType.float32, 4](0)
for j in range(4):
x_vals[j] = self.x[i + j]
vx_vals[j] = self.vx[i + j]
# Vectorized computation: x += vx * dt
var new_x = x_vals + vx_vals * dt_vec
# Manual store back to memory
for j in range(4):
self.x[i + j] = new_x[j]
# Repeat for y and z
var y_vals = SIMD[DType.float32, 4](0)
var vy_vals = SIMD[DType.float32, 4](0)
for j in range(4):
y_vals[j] = self.y[i + j]
vy_vals[j] = self.vy[i + j]
var new_y = y_vals + vy_vals * dt_vec
for j in range(4):
self.y[i + j] = new_y[j]
var z_vals = SIMD[DType.float32, 4](0)
var vz_vals = SIMD[DType.float32, 4](0)
for j in range(4):
z_vals[j] = self.z[i + j]
vz_vals[j] = self.vz[i + j]
var new_z = z_vals + vz_vals * dt_vec
for j in range(4):
self.z[i + j] = new_z[j]
# Handle remaining elements
for i in range(simd_count, self.count):
self.x[i] += self.vx[i] * dt
self.y[i] += self.vy[i] * dt
self.z[i] += self.vz[i] * dt
fn total_kinetic_energy_simd(self) -> Float32:
"""Calculate total kinetic energy using SIMD operations."""
var total: Float32 = 0.0
var simd_count = (self.count // 4) * 4
# SIMD processing
var energy_sum = SIMD[DType.float32, 4](0.0)
var half_vec = SIMD[DType.float32, 4](0.5)
for i in range(0, simd_count, 4):
# Manual load
var vx_vals = SIMD[DType.float32, 4](0)
var vy_vals = SIMD[DType.float32, 4](0)
var vz_vals = SIMD[DType.float32, 4](0)
var mass_vals = SIMD[DType.float32, 4](0)
for j in range(4):
vx_vals[j] = self.vx[i + j]
vy_vals[j] = self.vy[i + j]
vz_vals[j] = self.vz[i + j]
mass_vals[j] = self.mass[i + j]
var v_squared = vx_vals * vx_vals + vy_vals * vy_vals + vz_vals * vz_vals
var kinetic = half_vec * mass_vals * v_squared
energy_sum += kinetic
# Sum the SIMD vector
for i in range(4):
total += energy_sum[i]
# Handle remaining elements
for i in range(simd_count, self.count):
var v_sq = self.vx[i] * self.vx[i] + self.vy[i] * self.vy[i] + self.vz[i] * self.vz[i]
total += 0.5 * self.mass[i] * v_sq
return total
fn print_sample(self, num_samples: Int):
"""Print sample of particles."""
var samples = min(num_samples, self.count)
print("SoA Sample particles:")
for i in range(samples):
print(" Particle", i, ": pos(", self.x[i], ",", self.y[i], ",", self.z[i],
") vel(", self.vx[i], ",", self.vy[i], ",", self.vz[i], ") mass:", self.mass[i])
fn min(a: Int, b: Int) -> Int:
return a if a < b else b
fn soa_demo():
"""Demonstrate Struct of Arrays memory layout."""
print("=== Struct of Arrays (SoA) Pattern ===")
var system = ParticleSystemSoA(1000)
# Memory layout visualization
print("Memory Layout: x:[x0,x1,x2...] y:[y0,y1,y2...] vx:[vx0,vx1,vx2...]")
print("Characteristics:")
print(" - Good: Excellent vectorization, cache-friendly for bulk operations")
print(" - Bad: Less intuitive, scattered access for single particle operations")
system.print_sample(3)
# Simulate operations
var dt: Float32 = 0.1
# Compare regular vs SIMD operations
print("\nRegular update:")
system.update_all_positions(dt)
print("SIMD-optimized update:")
system.update_positions_simd(dt)
var total_energy = system.total_kinetic_energy_simd()
print("Total kinetic energy (SIMD):", total_energy)
fn main():
"""Main function for SoA demonstration."""
soa_demo()
File: 13_performance_comparison.mojo — a GPU-shaped tensor built as SoA (Section 3.5)¶
Execution: pixi run mojo 13_performance_comparison.mojo
from memory import UnsafePointer
# Tensor-like structure using SoA for optimal GPU performance
struct TensorSoA[dtype: DType]:
"""GPU-optimized tensor using Struct of Arrays layout."""
var data: UnsafePointer[Scalar[dtype]]
var gradients: UnsafePointer[Scalar[dtype]]
var shape: UnsafePointer[Int]
var strides: UnsafePointer[Int]
var ndim: Int
var size: Int
var requires_grad: Bool
fn __init__(out self, shape_list: List[Int], requires_grad: Bool = False):
self.ndim = len(shape_list)
self.requires_grad = requires_grad
# Allocate shape and stride arrays
self.shape = UnsafePointer[Int].alloc(self.ndim)
self.strides = UnsafePointer[Int].alloc(self.ndim)
# Calculate size and strides
self.size = 1
for i in range(self.ndim):
self.shape[i] = shape_list[i]
self.size *= shape_list[i]
# Calculate strides (row-major order)
if self.ndim > 0:
self.strides[self.ndim - 1] = 1
for i in range(self.ndim - 2, -1, -1):
self.strides[i] = self.strides[i + 1] * self.shape[i + 1]
# Allocate data arrays
self.data = UnsafePointer[Scalar[dtype]].alloc(self.size)
if self.requires_grad:
self.gradients = UnsafePointer[Scalar[dtype]].alloc(self.size)
# Initialize gradients to zero
for i in range(self.size):
self.gradients[i] = Scalar[dtype](0)
else:
self.gradients = UnsafePointer[Scalar[dtype]]()
fn __copyinit__(out self, existing: Self):
"""Copy constructor for tensor copying."""
self.ndim = existing.ndim
self.size = existing.size
self.requires_grad = existing.requires_grad
# Allocate and copy shape
self.shape = UnsafePointer[Int].alloc(self.ndim)
self.strides = UnsafePointer[Int].alloc(self.ndim)
for i in range(self.ndim):
self.shape[i] = existing.shape[i]
self.strides[i] = existing.strides[i]
# Allocate and copy data
self.data = UnsafePointer[Scalar[dtype]].alloc(self.size)
for i in range(self.size):
self.data[i] = existing.data[i]
# Allocate and copy gradients if needed
if self.requires_grad:
self.gradients = UnsafePointer[Scalar[dtype]].alloc(self.size)
for i in range(self.size):
self.gradients[i] = existing.gradients[i]
else:
self.gradients = UnsafePointer[Scalar[dtype]]()
fn __del__(owned self):
self.data.free()
if self.requires_grad:
self.gradients.free()
self.shape.free()
self.strides.free()
fn fill(self, value: Scalar[dtype]):
"""Fill tensor with a constant value."""
for i in range(self.size):
self.data[i] = value
fn add_elementwise(self, other: TensorSoA[dtype]) -> TensorSoA[dtype]:
"""Element-wise addition optimized for vectorization."""
# Create result tensor
var shape_list = List[Int]()
for i in range(self.ndim):
shape_list.append(self.shape[i])
var result = TensorSoA[dtype](shape_list, self.requires_grad or other.requires_grad)
# Vectorized addition using manual SIMD operations
var simd_size = 4
var simd_count = (self.size // simd_size) * simd_size
# Process in SIMD chunks
for i in range(0, simd_count, simd_size):
# Manual load from memory
var a_vals = SIMD[dtype, 4](0)
var b_vals = SIMD[dtype, 4](0)
for j in range(4):
a_vals[j] = self.data[i + j]
b_vals[j] = other.data[i + j]
# SIMD addition
var result_vals = a_vals + b_vals
# Manual store to memory
for j in range(4):
result.data[i + j] = result_vals[j]
# Handle remaining elements
for i in range(simd_count, self.size):
result.data[i] = self.data[i] + other.data[i]
return result
fn print_info(self):
"""Print tensor information."""
print("Tensor info:")
print(" Shape: [", end="")
for i in range(self.ndim):
print(self.shape[i], end="")
if i < self.ndim - 1:
print(", ", end="")
print("]")
print(" Size:", self.size)
print(" Requires grad:", self.requires_grad)
print(" First few elements:", self.data[0], self.data[1], self.data[2])
fn memory_layout_comparison():
"""Compare different memory layouts for tensor operations."""
print("=== Memory Layout Performance Comparison ===")
# Create tensors with different patterns
var shape_list = List[Int]()
shape_list.append(1000)
shape_list.append(100)
print("\nCreating tensors for performance comparison...")
# SoA-style tensor optimized for vectorization
var tensor_a = TensorSoA[DType.float32](shape_list, True)
var tensor_b = TensorSoA[DType.float32](shape_list, True)
# Fill with test data
tensor_a.fill(2.5)
tensor_b.fill(1.5)
print("Tensor A:")
tensor_a.print_info()
print("Tensor B:")
tensor_b.print_info()
# Perform vectorized operations
print("\nPerforming element-wise addition...")
var result = tensor_a.add_elementwise(tensor_b)
print("Result tensor:")
result.print_info()
print("\nMemory Layout Analysis:")
print("SoA Layout Benefits:")
print(" [+] Excellent cache locality for bulk operations")
print(" [+] SIMD vectorization friendly")
print(" [+] GPU memory coalescing")
print(" [+] Optimal for parallel processing")
print("SoA Layout Considerations:")
print(" - More complex indexing for individual elements")
print(" - Requires careful memory management")
print(" - Less intuitive than object-oriented approach")
fn main():
"""Main function for performance comparison."""
memory_layout_comparison()
File: 14_layout_patterns_complete.mojo — AoS and SoA side by side¶
Execution: pixi run mojo 14_layout_patterns_complete.mojo
from memory import UnsafePointer
# AoS Pattern Example
struct Point3D:
var x: Float32
var y: Float32
var z: Float32
fn __init__(out self, x: Float32, y: Float32, z: Float32):
self.x = x
self.y = y
self.z = z
# SoA Pattern Example
struct Points3DSoA:
var x: UnsafePointer[Float32]
var y: UnsafePointer[Float32]
var z: UnsafePointer[Float32]
var count: Int
fn __init__(out self, count: Int):
self.count = count
self.x = UnsafePointer[Float32].alloc(count)
self.y = UnsafePointer[Float32].alloc(count)
self.z = UnsafePointer[Float32].alloc(count)
fn __del__(owned self):
self.x.free()
self.y.free()
self.z.free()
fn set_point(self, index: Int, x: Float32, y: Float32, z: Float32):
self.x[index] = x
self.y[index] = y
self.z[index] = z
fn distance_sum_simd(self) -> Float32:
"""Calculate sum of distances from origin using SIMD."""
var total: Float32 = 0.0
var simd_count = (self.count // 4) * 4
var sum_vec = SIMD[DType.float32, 4](0.0)
for i in range(0, simd_count, 4):
# Manual load from memory
var x_vals = SIMD[DType.float32, 4](0)
var y_vals = SIMD[DType.float32, 4](0)
var z_vals = SIMD[DType.float32, 4](0)
for j in range(4):
x_vals[j] = self.x[i + j]
y_vals[j] = self.y[i + j]
z_vals[j] = self.z[i + j]
var dist_sq = x_vals * x_vals + y_vals * y_vals + z_vals * z_vals
sum_vec += dist_sq
# Sum SIMD vector elements
for i in range(4):
total += sum_vec[i]
# Handle remaining elements
for i in range(simd_count, self.count):
var dist_sq = self.x[i] * self.x[i] + self.y[i] * self.y[i] + self.z[i] * self.z[i]
total += dist_sq
return total
fn complete_layout_demo():
"""Complete demonstration of memory layout patterns."""
print("=== Complete Memory Layout Patterns Demo ===")
var num_points = 1000
# AoS approach
print("Array of Structs (AoS):")
print(" Memory: [x,y,z][x,y,z][x,y,z]...")
print(" Use case: Individual point operations")
var aos_points = UnsafePointer[Point3D].alloc(num_points)
for i in range(num_points):
aos_points[i] = Point3D(Float32(i), Float32(i) * 2, Float32(i) * 3)
print(" Sample AoS point:", aos_points[0].x, aos_points[0].y, aos_points[0].z)
aos_points.free()
# SoA approach
print("\nStruct of Arrays (SoA):")
print(" Memory: x:[x0,x1,x2...] y:[y0,y1,y2...] z:[z0,z1,z2...]")
print(" Use case: Bulk vectorized operations")
var soa_points = Points3DSoA(num_points)
for i in range(num_points):
soa_points.set_point(i, Float32(i), Float32(i) * 2, Float32(i) * 3)
print(" Sample SoA point:", soa_points.x[0], soa_points.y[0], soa_points.z[0])
var distance_sum = soa_points.distance_sum_simd()
print(" SIMD distance sum:", distance_sum)
print("\nPerformance Characteristics:")
print("AoS Benefits:")
print(" + Cache-friendly for single object operations")
print(" + Intuitive object-oriented design")
print(" + Good spatial locality for related data")
print("SoA Benefits:")
print(" + Excellent vectorization (SIMD/GPU)")
print(" + Cache-friendly for bulk operations")
print(" + Memory bandwidth optimization")
print(" + Parallel processing friendly")
print("\nTensor Framework Choice: SoA")
print("Reason: Automatic differentiation requires bulk operations on gradients,")
print(" making SoA optimal for GPU acceleration and SIMD vectorization.")
fn main():
"""Complete demonstration of memory layout patterns."""
complete_layout_demo()
Expected Output for 14_layout_patterns_complete.mojo¶
=== Complete Memory Layout Patterns Demo ===
Array of Structs (AoS):
Memory: [x,y,z][x,y,z][x,y,z]...
Use case: Individual point operations
Sample AoS point: 0.0 0.0 0.0
Struct of Arrays (SoA):
Memory: x:[x0,x1,x2...] y:[y0,y1,y2...] z:[z0,z1,z2...]
Use case: Bulk vectorized operations
Sample SoA point: 0.0 0.0 0.0
SIMD distance sum: [calculated value]
Performance Characteristics:
AoS Benefits:
+ Cache-friendly for single object operations
+ Intuitive object-oriented design
+ Good spatial locality for related data
SoA Benefits:
+ Excellent vectorization (SIMD/GPU)
+ Cache-friendly for bulk operations
+ Memory bandwidth optimization
+ Parallel processing friendly
Tensor Framework Choice: SoA
Reason: Automatic differentiation requires bulk operations on gradients,
making SoA optimal for GPU acceleration and SIMD vectorization.
The sample point being 0.0 0.0 0.0 in both layouts is Point3D/Points3DSoA index 0 — consistent with this chapter's initializer pattern, where element i scales with i, and i=0 scales everything to zero either way.
Chapter Summary¶
Bulk numerical code is usually limited by how many bytes move across the memory bus, not by how many instructions execute once those bytes have arrived, which makes the layout of the data — not just its algorithm — a first-class performance decision. Array-of-Structs packs every field of one object together, which is exactly right when an operation needs most or all of that object's fields at once (Section 3.2's update_position, at 85.7% bus utilization) and exactly wrong when an operation needs only one field across many objects, because the unwanted fields are physically wedged in between the wanted ones and ride along regardless (Section 3.1's total_kinetic_energy, at 57.1%). Struct-of-Arrays inverts the layout, packing every instance of one field together instead, which drives that same kinetic-energy computation to 100% bus utilization and — more fundamentally for this book — is the only layout a single contiguous SIMD load (Chapter 1.4) can actually use, since that instruction requires the values it loads to already be neighbors in memory. Both layouts compute identical answers for identical formulas (Section 3.4's kinetic-energy cross-check landed on 28.0 from AoS and SoA alike); the choice between them changes speed, never correctness. This book builds its Tensor type as SoA specifically because automatic differentiation's dominant workload — bulk operations across every element's value or gradient — is squarely SoA's strength, not because SoA is a universally superior layout.
Self-Check Questions¶
total_kinetic_energy(AoS) has a bus utilization of57.1%whileupdate_position(AoS, on the same struct) has85.7%. Both operate on the identicalParticlelayout — explain why their utilization numbers differ so much.- Why can a single
SIMD[DType.float32, 4].load()instruction pull four consecutivexvalues from aParticleSystemSoA, but not from aParticleSystemAoSholding the same 1000 particles? - Worked Example 3.4.1 found that both the AoS scalar loop and the SoA SIMD loop compute a total kinetic energy of
28.0for 4 particles. What conclusion is and is not licensed by the two layouts agreeing? - A hypothetical
Elementstruct bundles onevalue: Float32and onegrad: Float32together, and a program allocatesUnsafePointer[Element].alloc(n). What byte stride separates consecutivevaluefields in that array, and why does that number make a single contiguous SIMD load of fourvalues impossible? - Give one realistic operation where AoS would outperform SoA, and explain — in terms of bus utilization — why.
Where We Go Next¶
Chapter 4 moves from how data should be laid out in ordinary RAM to how to get that data onto a GPU in the first place — the device abstraction, memory transfer, and kernel-launch model that Part 1's Tensor (built as SoA, per this chapter) is designed to work with directly.
Worked Solutions¶
1. update_position reads and writes six of Particle's seven fields (x, y, z, vx, vy, vz — everything but mass), so of the 28 bytes the hardware must fetch for one particle, 24 are actually used: 24/28 ≈ 85.7%. total_kinetic_energy reads only four of the seven fields (vx, vy, vz, mass), so of the same 28 fetched bytes, only 16 are used: 16/28 ≈ 57.1%. The struct's layout and total size never change between the two operations — only how many of its fields each specific operation actually needs, which is what utilization measures.
2. A contiguous SIMD load requires the values it loads to sit at consecutive byte addresses with no gaps, which is exactly what a dedicated x array in ParticleSystemSoA provides — four x values back-to-back, 4 bytes apart. In ParticleSystemAoS, consecutive particles' x fields are 28 bytes apart (the full size of one Particle), because y, z, vx, vy, vz, and mass for that same particle sit physically between one x and the next — a stride no single contiguous SIMD load instruction can span.
3. It licenses the conclusion that the layout choice did not change the mathematical result of the kinetic-energy formula — both are valid, correct implementations of KE = 0.5·m·(vx²+vy²+vz²). It does not license any conclusion about speed: agreement on the answer says nothing about which one moved fewer bytes across the bus to get there — that comparison was made separately, in Section 3.1, by counting bytes, not by comparing final answers.
4. Element is 4 + 4 = 8 bytes (one Float32 value, one Float32 grad), so consecutive value fields in the array sit 8 bytes apart, not 4. A SIMD[DType.float32, 4] load reads 16 contiguous bytes and interprets them as four adjacent 4-byte lanes; at stride 8, the second value doesn't begin until byte 8, meaning the load would actually capture value0, grad0, value1, grad1 as its four lanes — the wrong four numbers entirely, not simply an inefficient fetch. Getting the actual four values requires either four separate scalar loads (no vectorization gained) or a gather instruction built for non-contiguous access (slower than a plain contiguous load).
5. Any operation that needs nearly all of one object's fields at once — Section 3.2's update_position is the book's own example, reading and writing six of seven fields per particle. In AoS, that entire near-complete record arrives in one contiguous 28-byte fetch, at 85.7% utilization. In SoA, the same operation would require six separate memory streams (one each for x, y, z, vx, vy, vz, each in its own array) to be advanced in lockstep for every particle — more separate memory streams active at once, and no bandwidth advantage to offset that complexity, since nearly all the fetched bytes were going to be used either way.