Skip to content

Measure the memory hierarchy and threadgroup memory, which reorders the target model #31

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
AndreSlavescu merged 2 commits into from
Jul 30, 2026
Merged

Measure the memory hierarchy and threadgroup memory, which reorders the target model #31

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
286a921
Measure the memory hierarchy, which outranks every other lever in the…
AndreSlavescu Jul 30, 2026
4601fc4
Measure threadgroup memory against the cache, and confirm the pad cal…
AndreSlavescu Jul 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
183 changes: 183 additions & 0 deletions benchmarks/agx_memory_hierarchy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,183 @@
"""Where the memory hierarchy changes speed, which is what tiling has to be sized against.

meTile records one bandwidth number, 120.6 GB/s, measured streaming from DRAM. That number is not the
whole story and the gap showed up by accident: a 786 KB weight matrix read at 254 GB/s, more than twice
the recorded ceiling, which means something above DRAM is serving it. A compiler choosing tile sizes
needs to know where that changes, because a tile that fits the fast level and a tile that misses it are
not the same kernel.

The method is a coalesced streaming read over a working set of `size` bytes, looped until a fixed total
of traffic has been read. Bandwidth is total bytes over elapsed time, so a working set that stays
resident reports the resident level's bandwidth and one that does not reports DRAM's. Plateaus are
levels; the knees between them are capacities.

Three things make the numbers trustworthy rather than suggestive:

saturation the sweep has to reach the known 120.6 GB/s at large sizes, or the thread count is too
low to saturate and every number is a lower bound on something else. Printed as a check
rather than assumed.
amortisation each dispatch reads gigabytes and runs for tens of milliseconds, so launch overhead is
far below the noise. Timing small kernels through a host round trip is what produced
three fabricated results earlier in this project.
interleaving sizes are measured in rotating order across rounds, because a sweep that walks from
small to large measures thermal drift as much as it measures the hierarchy.

One limit of this probe is worth knowing before reading its output. Inside the resident regime bandwidth
*rises* with working set -- 1545, 2006, 2138 and 2386 GB/s at 256 KB, 512 KB, 1 MB and 2 MB -- which
cannot be a property of a cache. A smaller working set means fewer inner iterations per pass, so the
pass loop's bookkeeping takes a larger share, and the effect shrinks as the set grows. The probe
therefore establishes that a resident set runs at 2386 GB/s or better and says nothing reliable about
how that varies below 2 MB. `metile.target.agx` records it as one number for that reason.

usage:
python benchmarks/agx_memory_hierarchy.py
python benchmarks/agx_memory_hierarchy.py --traffic 8 --rounds 7
"""

import argparse
import itertools
import statistics
import sys
import time
from pathlib import Path

import numpy as np

_root = str(Path(__file__).resolve().parent.parent)
sys.path.insert(0, _root)

import metile
from metile.target import agx

THREADGROUP = 256
# Enough threads to saturate DRAM, and few enough that a 256 KB working set still gives every thread
# something to do: the inner loop strides by the total thread count, so a working set smaller than that
# leaves threads idle and measures parallelism instead of bandwidth.
THREADGROUPS = 64
VECTOR_BYTES = 16 # float4

SIZES_KB = (256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072, 262144)


def _arguments():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--traffic", type=float, default=4.0, help="GB read per dispatch")
parser.add_argument("--rounds", type=int, default=5)
parser.add_argument("--threadgroups", type=int, default=THREADGROUPS)
return parser.parse_args()


READ_KERNEL = """#include <metal_stdlib>
using namespace metal;

kernel void probe(device const float4* data [[buffer(0)]],
device float4* out [[buffer(1)]],
constant uint& vectors [[buffer(2)]],
constant uint& passes [[buffer(3)]],
constant uint& stride [[buffer(4)]],
uint gid [[thread_position_in_grid]]) {
float4 total = float4(0.0f);
for (uint pass = 0; pass < passes; ++pass) {
for (uint index = gid; index < vectors; index += stride) {
total += data[index];
}
}
// Written unconditionally so nothing above can be discarded as dead.
out[gid] = total;
}
"""


def main():
arguments = _arguments()
from metile.runtime.metal_device import MetalDevice

device = MetalDevice.get()
threads = arguments.threadgroups * THREADGROUP
grid, block = (threads, 1, 1), (THREADGROUP, 1, 1)
pipeline = device.compile_msl(READ_KERNEL, "probe")

print(f"device: {device.name}")
print(f"{threads} threads, {arguments.traffic:g} GB read per dispatch")
print(f"recorded streaming ceiling: {agx.STREAMING_READ_GBPS} GB/s\n")

target_bytes = arguments.traffic * 1e9
cases = []
largest = max(SIZES_KB) * 1024
data = metile.Buffer(data=np.ones(largest // 4, dtype=np.float32))
out = metile.Buffer(data=np.zeros(threads * 4, dtype=np.float32))

for size_kb in SIZES_KB:
size = size_kb * 1024
vectors = size // VECTOR_BYTES
if vectors < threads:
continue
passes = max(1, int(target_bytes // size))
buffers = [
data.metal_buffer,
out.metal_buffer,
metile.Buffer(data=np.array([vectors], dtype=np.uint32)).metal_buffer,
metile.Buffer(data=np.array([passes], dtype=np.uint32)).metal_buffer,
metile.Buffer(data=np.array([threads], dtype=np.uint32)).metal_buffer,
]
cases.append((size, vectors, passes, buffers))

def measure(buffers):
started = time.perf_counter_ns()
device.dispatch_kernel(pipeline, buffers, grid, block)
device.sync()
return (time.perf_counter_ns() - started) / 1e9

for _, _, _, buffers in cases:
measure(buffers)

samples = {size: [] for size, _, _, _ in cases}
for index in range(arguments.rounds):
ordered = cases[index % len(cases) :] + cases[: index % len(cases)]
for size, _, _, buffers in ordered:
samples[size].append(measure(buffers))

print(f"{'working set':>13}{'passes':>8}{'ms':>9}{'GB/s':>9}{'vs DRAM':>9}")
results = []
for size, vectors, passes, _ in cases:
seconds = statistics.median(samples[size])
gbps = (vectors * VECTOR_BYTES * passes) / seconds / 1e9
results.append((size, gbps))
label = f"{size // 1024} KB" if size < 1024 * 1024 else f"{size // (1024 * 1024)} MB"
print(
f"{label:>13}{passes:>8}{seconds * 1e3:>9.1f}{gbps:>9.1f}"
f"{gbps / agx.STREAMING_READ_GBPS:>8.2f}x"
)

dram = min(gbps for _, gbps in results)
peak = max(gbps for _, gbps in results)
print(f"\nslowest {dram:.1f} GB/s, fastest {peak:.1f} GB/s, ratio {peak / dram:.2f}x")
saturated = dram >= 0.9 * agx.STREAMING_READ_GBPS
print(
f"saturation check: largest working sets reach {dram:.1f} GB/s against a recorded "
f"{agx.STREAMING_READ_GBPS} GB/s -- {'ok' if saturated else 'TOO LOW, raise --threadgroups'}"
)
if not saturated:
print(
"Until that passes, every number here is a lower bound and the knees may be artefacts."
)
return 1

# Report the knees rather than leaving them to be eyeballed: a level boundary is where bandwidth
# drops materially between adjacent sizes.
print("\ntransitions, where bandwidth drops by more than 15% between adjacent sizes:")
found = False
for (small, fast), (large, slow) in itertools.pairwise(results):
if slow < 0.85 * fast:
found = True
print(
f" between {small // 1024} KB and {large // 1024} KB: "
f"{fast:.1f} -> {slow:.1f} GB/s"
)
if not found:
print(" none; the hierarchy looks flat over this range")
return 0


if __name__ == "__main__":
raise SystemExit(main())
178 changes: 178 additions & 0 deletions benchmarks/agx_threadgroup_bandwidth.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,178 @@
"""What threadgroup memory is actually for on this part, and which strides it punishes.

meTile has several passes built around threadgroup memory -- cooperative loads, shared-memory padding
and swizzling, double buffering -- and they all assume it is the fast place to put data. On M5 that
assumption needs checking, because device memory is not slow when it is resident: a working set under
2 MB reads at 2386 GB/s against DRAM's 121. So the question is not whether threadgroup memory is fast,
but whether it beats the cache that would have served the same bytes anyway.

Two things get measured, both over a 32 KB working set per threadgroup so the device arm is resident and
the comparison is staging against a cache hit rather than against DRAM.

bandwidth a coalesced read from threadgroup memory against the same read from device memory.
stride the same total bytes with consecutive lanes a fixed stride apart, swept, for both spaces.
This separates how much each cares about indexing, which is the property padding and
swizzling exist to manage.

The stride sweep is the interesting half, and it corrected a wrong diagnosis on the way here. A
per-thread span of eight vectors made the threadgroup arm about twice as slow, which looked like a bank
conflict, so an arm was added that padded the span by one to break it up. Padding made it worse -- the
unpadded span was seven, and adding one moved it *onto* 128 bytes rather than off it. Padding is not a
direction, it is an arithmetic result, and plus one is not automatically safe.

Controls, for the reasons three fabricated results earlier in this project established: elapsed time has
to scale with the pass count or the loop is being collapsed, and each dispatch runs for milliseconds so
launch overhead sits far below the noise.
"""

import argparse
import statistics
import sys
import time
from pathlib import Path

import numpy as np

_root = str(Path(__file__).resolve().parent.parent)
sys.path.insert(0, _root)

import metile
from metile.target import agx

THREADGROUP = 256
THREADGROUPS = 64
VECTOR_BYTES = 16
# A power of two so the sweep can wrap with a mask rather than a modulo, which would put a divide in the
# inner loop and measure that instead.
TILE_VECTORS = 2048
TILE_BYTES = TILE_VECTORS * VECTOR_BYTES
# Vectors each thread reads per pass, constant across strides so every stride moves the same bytes.
PER_THREAD = 8
STRIDES = (1, 2, 3, 4, 5, 6, 7, 8, 9, 12, 16, 32)


def _arguments():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--passes", type=int, default=6000)
parser.add_argument("--rounds", type=int, default=3)
return parser.parse_args()


def _source(space, stride):
"""One arm: read PER_THREAD vectors per pass from `space`, lanes `stride` vectors apart."""
staging = (
f" threadgroup float4 tile[{TILE_VECTORS}];\n"
f" for (uint index = lid; index < {TILE_VECTORS}; index += {THREADGROUP}) {{\n"
f" tile[index] = data[index];\n"
f" }}\n"
f" threadgroup_barrier(mem_flags::mem_threadgroup);\n"
if space == "threadgroup"
else ""
)
read = "tile" if space == "threadgroup" else "data"
return f"""#include <metal_stdlib>
using namespace metal;

kernel void probe(device const float4* data [[buffer(0)]],
device float4* out [[buffer(1)]],
constant uint& passes [[buffer(2)]],
uint gid [[thread_position_in_grid]],
uint lid [[thread_position_in_threadgroup]]) {{
{staging}
float4 total = float4(0.0f);
for (uint pass = 0; pass < passes; ++pass) {{
for (uint step = 0; step < {PER_THREAD}; ++step) {{
total += {read}[(lid * {stride} + step) & {TILE_VECTORS - 1}];
}}
}}
out[gid] = total;
}}
"""


def main():
arguments = _arguments()
from metile.runtime.metal_device import MetalDevice

device = MetalDevice.get()
threads = THREADGROUPS * THREADGROUP
grid, block = (threads, 1, 1), (THREADGROUP, 1, 1)

data = metile.Buffer(data=np.ones(TILE_BYTES // 4, dtype=np.float32))
out = metile.Buffer(data=np.zeros(threads * 4, dtype=np.float32))

print(f"device: {device.name}")
print(f"{THREADGROUPS} threadgroups x {THREADGROUP} threads, {TILE_BYTES // 1024} KB per group")
print(f"resident device read: {agx.RESIDENT_READ_GBPS} GB/s, DRAM: {agx.STREAMING_READ_GBPS}\n")

arms = {}
for space in ("threadgroup", "device"):
for stride in STRIDES:
arms[space, stride] = device.compile_msl(_source(space, stride), "probe")

def measure(pipeline, passes):
buffers = [
data.metal_buffer,
out.metal_buffer,
metile.Buffer(data=np.array([passes], dtype=np.uint32)).metal_buffer,
]
started = time.perf_counter_ns()
device.dispatch_kernel(pipeline, buffers, grid, block)
device.sync()
return (time.perf_counter_ns() - started) / 1e9

for pipeline in arms.values():
measure(pipeline, arguments.passes)

def sweep(passes):
keys = list(arms)
samples = {key: [] for key in keys}
for index in range(arguments.rounds):
ordered = keys[index % len(keys) :] + keys[: index % len(keys)]
for key in ordered:
samples[key].append(measure(arms[key], passes))
return {key: statistics.median(values) for key, values in samples.items()}

per_pass = THREADGROUPS * THREADGROUP * PER_THREAD * VECTOR_BYTES
medians = sweep(arguments.passes)

def rate(key):
return per_pass * arguments.passes / medians[key] / 1e9

print(f"{'stride':>7}{'bytes':>7}{'shared GB/s':>13}{'device GB/s':>13}{'shared/device':>15}")
for stride in STRIDES:
shared, plain = rate(("threadgroup", stride)), rate(("device", stride))
print(
f"{stride:>7}{stride * VECTOR_BYTES:>7}{shared:>13.0f}{plain:>13.0f}"
f"{shared / plain:>14.2f}x"
)

best_shared = max(rate(("threadgroup", stride)) for stride in STRIDES)
worst_shared = min(rate(("threadgroup", stride)) for stride in STRIDES)
best_device = max(rate(("device", stride)) for stride in STRIDES)
worst_device = min(rate(("device", stride)) for stride in STRIDES)
print(
f"\nspread across strides: threadgroup {best_shared / worst_shared:.2f}x, "
f"device {best_device / worst_device:.2f}x"
)
print(f"best threadgroup {best_shared:.0f} GB/s vs best device {best_device:.0f} GB/s")

penalised = [
stride
for stride in STRIDES
if rate(("threadgroup", stride)) < 0.7 * best_shared
and rate(("device", stride)) > 0.7 * best_device
]
print(f"strides threadgroup memory punishes and device memory does not: {penalised or 'none'}")

doubled = sweep(arguments.passes * 2)
ratios = [doubled[key] / medians[key] for key in arms]
print(f"\ncontrol, doubling passes: min {min(ratios):.2f}x, max {max(ratios):.2f}x, want ~2.00")
if not all(1.7 <= ratio <= 2.3 for ratio in ratios):
print("Time does not scale with work; the loop is being optimised and these are fiction.")
return 1
return 0


if __name__ == "__main__":
raise SystemExit(main())
Loading
Loading