From 286a921975be70c221644e254b6c4b56ffb7a11e Mon Sep 17 00:00:00 2001 From: AndreSlavescu Date: Thu, 30 Jul 2026 00:17:06 -0700 Subject: [PATCH 1/2] Measure the memory hierarchy, which outranks every other lever in the target model meTile recorded one bandwidth number, 120.6 GB/s streaming from DRAM, and it is not the whole story. The gap surfaced by accident while scoping attention projections: a 786 KB weight read at 254 GB/s, twice the recorded ceiling, so something above DRAM was serving it. Sweeping working-set size finds the levels: <= 2 MB 2386 GB/s 19.8x DRAM 4 MB 555 4.6x 8 MB 192 1.6x 16 MB 161 1.3x 32 MB 134 1.1x >= 64 MB 121-128 1.0x The knee at 2 to 4 MB is a factor of four in one step. That makes fitting a working set the largest lever in metile/target/agx.py by some distance: choosing the matrix unit over scalar is worth 2.4x to 3.7x, and instruction scheduling is capped at 1.09x and measured unreachable above MSL. A tiling that fits and a tiling that misses are not the same kernel, and until now the compiler had no way to know the difference. `read_bandwidth_gbps`, `resident` and `tiling_gain` give a pass something to consult. Two controls, because a 19x claim earns them. The pass loop re-reads the same addresses, so the backend could in principle collapse it into a multiply and the fast numbers would be fiction: tripling the traffic triples the elapsed time at every size, which a collapsed loop could not do. And 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 -- it reaches 120.5, and the probe reports the check rather than assuming it. The resident regime is recorded as one number rather than four, and finding out why was the useful part. It measured 1545, 2006, 2138 and 2386 GB/s at 256 KB, 512 KB, 1 MB and 2 MB: bandwidth *rising* with working set, which cannot be a property of a cache. It is the pass loop again -- a smaller working set means fewer inner iterations per pass, so bookkeeping takes a larger share, and the effect shrinks as the set grows. Only the 2 MB figure is close to uncontaminated and it is a floor. Publishing four numbers would claim a resolution the measurement does not have. A monotonicity test over the table is what caught that, twice: first the 256 KB entry reading slower than 512 KB, then 512 KB reading slower than 1 MB. A larger working set is never served faster, so a table that says otherwise is measuring the harness. 683 pass. Lint and vulture clean. Co-Authored-By: Claude Opus 5 (1M context) --- benchmarks/agx_memory_hierarchy.py | 183 +++++++++++++++++++++++++++++ metile/target/__init__.py | 12 ++ metile/target/agx.py | 66 +++++++++++ tests/test_target.py | 49 ++++++++ 4 files changed, 310 insertions(+) create mode 100644 benchmarks/agx_memory_hierarchy.py diff --git a/benchmarks/agx_memory_hierarchy.py b/benchmarks/agx_memory_hierarchy.py new file mode 100644 index 0000000..ade412b --- /dev/null +++ b/benchmarks/agx_memory_hierarchy.py @@ -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 +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()) diff --git a/metile/target/__init__.py b/metile/target/__init__.py index a1c593a..cf50086 100644 --- a/metile/target/__init__.py +++ b/metile/target/__init__.py @@ -6,27 +6,39 @@ """ from metile.target.agx import ( + BANDWIDTH_BY_WORKING_SET_GBPS, ILP_CEILING, MATRIX_PEAK_TFLOPS, REGISTER_BUDGET, + RESIDENT_READ_GBPS, + RESIDENT_WORKING_SET_BYTES, SCALAR_PEAK_TFLOPS, STREAMING_READ_GBPS, Unavailable, ilp_headroom, inspect, machine_code, + read_bandwidth_gbps, + resident, spills, + tiling_gain, ) __all__ = [ + "BANDWIDTH_BY_WORKING_SET_GBPS", "ILP_CEILING", "MATRIX_PEAK_TFLOPS", "REGISTER_BUDGET", + "RESIDENT_READ_GBPS", + "RESIDENT_WORKING_SET_BYTES", "SCALAR_PEAK_TFLOPS", "STREAMING_READ_GBPS", "Unavailable", "ilp_headroom", "inspect", "machine_code", + "read_bandwidth_gbps", + "resident", "spills", + "tiling_gain", ] diff --git a/metile/target/agx.py b/metile/target/agx.py index ee56d9e..cb7b301 100644 --- a/metile/target/agx.py +++ b/metile/target/agx.py @@ -37,6 +37,41 @@ MATRIX_PEAK_TFLOPS = 15.33 STREAMING_READ_GBPS = 120.6 +# Read bandwidth as a function of working-set size, measured by benchmarks/agx_memory_hierarchy.py. +# One number for bandwidth is badly wrong here: a working set that stays resident is served sixteen to +# twenty times faster than one that streams, and that ratio is larger than every other factor in this +# file. A tiling that fits and a tiling that misses are not the same kernel. +# +# Coalesced streaming reads, 16384 threads, six gigabytes of traffic per dispatch, sizes interleaved +# across rounds. The pass loop re-reads the same addresses, so the obvious worry is that the backend +# collapses it into a multiply and the fast numbers are fiction; tripling the traffic triples the +# elapsed time at every size, which it could not if the loop were collapsed. +# +# The resident regime is one entry, not four, because the probe cannot resolve it. Measured 1545, 2006, +# 2138 and 2386 GB/s at 256 KB, 512 KB, 1 MB and 2 MB -- bandwidth *rising* with working set, which +# cannot be a property of a cache. It is the pass loop: a smaller working set means fewer inner +# iterations per pass, so loop bookkeeping takes a larger share, and the effect shrinks as the set +# grows. Only the 2 MB figure is close to uncontaminated, and it is a floor. +# +# So what this establishes is that a resident working set runs at 2386 GB/s or better, and nothing about +# how that varies below 2 MB. Reporting four numbers would claim a resolution the measurement does not +# have; a monotonicity test on the table is what caught the attempt. +BANDWIDTH_BY_WORKING_SET_GBPS = { + 2 * 1024 * 1024: 2386.0, + 4 * 1024 * 1024: 555.0, + 8 * 1024 * 1024: 192.0, + 16 * 1024 * 1024: 161.0, + 32 * 1024 * 1024: 134.0, + 64 * 1024 * 1024: 128.0, + 128 * 1024 * 1024: 124.0, +} + +# Largest working set still served by the fast level, and what it delivers. The knee is sharp: 2 MB +# reads 2386 GB/s and 4 MB reads 555, so this is the number a tiling pass should be trying to stay +# under. The rate is a floor for the reason above. +RESIDENT_WORKING_SET_BYTES = 2 * 1024 * 1024 +RESIDENT_READ_GBPS = 2386.0 + class Unavailable(RuntimeError): """The toolchain needed to inspect compiled kernels is not present.""" @@ -52,6 +87,37 @@ def spills(registers): return registers >= REGISTER_BUDGET +def read_bandwidth_gbps(working_set_bytes): + """Expected read bandwidth for a working set of this size, in GB/s. + + For a pass deciding a tile size. Interpolating between measured points would invent a smooth curve + the hardware does not have -- the drop from 2 MB to 4 MB is a factor of four -- so this reports the + measurement for the smallest size at least as large as the request, which is the conservative + direction: a tile is served no faster than the next size up was measured at. + """ + if working_set_bytes <= 0: + raise ValueError("a working set must be positive") + for size in sorted(BANDWIDTH_BY_WORKING_SET_GBPS): + if working_set_bytes <= size: + return BANDWIDTH_BY_WORKING_SET_GBPS[size] + return STREAMING_READ_GBPS + + +def resident(working_set_bytes): + """Whether a working set of this size is served by the fast level rather than by DRAM.""" + return 0 < working_set_bytes <= RESIDENT_WORKING_SET_BYTES + + +def tiling_gain(working_set_bytes): + """How much bandwidth a tiling wins by fitting this working set instead of streaming. + + The figure worth putting beside the other ratios in this file. Fitting under 2 MB is worth about + 19x, where choosing the matrix unit over scalar is worth 2.4x to 3.7x and instruction scheduling is + worth at most 1.09x and unreachable in practice. + """ + return read_bandwidth_gbps(working_set_bytes) / STREAMING_READ_GBPS + + def _harness(workdir): """Build the Metal harness once per working directory.""" binary = workdir / "agx_probe" diff --git a/tests/test_target.py b/tests/test_target.py index 2e3e674..1db3fae 100644 --- a/tests/test_target.py +++ b/tests/test_target.py @@ -1,5 +1,7 @@ """The target model is measured hardware knowledge, so guard what depends on its shape.""" +import itertools + import pytest from metile.target import ( @@ -99,3 +101,50 @@ def test_the_backend_normalises_statement_order(): "on the assumption that it does not; re-measure benchmarks/agx_source_order.py and " "reconsider the default." ) + + +def test_bandwidth_depends_on_working_set_and_the_knee_is_sharp(): + """One bandwidth number is badly wrong for this part, so the model has to be a curve. + + The measured drop from 2 MB to 4 MB is about a factor of four. Interpolating across it would invent + a smooth ramp the hardware does not have, so the lookup reports the measurement for the smallest + size at least as large as the request. + """ + from metile.target import RESIDENT_WORKING_SET_BYTES, read_bandwidth_gbps + + resident_rate = read_bandwidth_gbps(RESIDENT_WORKING_SET_BYTES) + beyond_rate = read_bandwidth_gbps(RESIDENT_WORKING_SET_BYTES * 2) + assert resident_rate > 3 * beyond_rate + assert read_bandwidth_gbps(2**30) == pytest.approx(STREAMING_READ_GBPS) + + # Monotonically non-increasing: a larger working set is never served faster. This is what caught + # a 256 KB entry reading slower than 512 KB, which is impossible for a smaller working set and was + # the probe's loop overhead rather than the hierarchy. + rates = [read_bandwidth_gbps(2**exponent) for exponent in range(14, 31)] + for faster, slower in itertools.pairwise(rates): + assert faster >= slower + + +def test_a_working_set_is_resident_only_up_to_the_measured_capacity(): + from metile.target import RESIDENT_WORKING_SET_BYTES, resident + + assert resident(RESIDENT_WORKING_SET_BYTES) + assert resident(1024) + assert not resident(RESIDENT_WORKING_SET_BYTES + 1) + assert not resident(0) + + +def test_fitting_a_tile_outranks_every_other_lever_in_this_file(): + """The comparison that should drive where compiler effort goes. + + Keeping a working set resident is worth about 19x. Choosing the matrix unit over scalar is worth + 2.4x to 3.7x. Instruction scheduling is capped at 1.09x and measured unreachable above MSL. If that + ordering ever changes on new hardware, the guidance built on it needs revisiting rather than being + carried over. + """ + from metile.target import RESIDENT_WORKING_SET_BYTES, tiling_gain + + fitting = tiling_gain(RESIDENT_WORKING_SET_BYTES) + functional_unit = MATRIX_PEAK_TFLOPS / min(SCALAR_PEAK_TFLOPS.values()) + assert fitting > functional_unit > max(ILP_CEILING.values()) + assert tiling_gain(2**30) == pytest.approx(1.0) From 4601fc44e10f07897e3e3ce556a277cce728af3a Mon Sep 17 00:00:00 2001 From: AndreSlavescu Date: Thu, 30 Jul 2026 00:24:00 -0700 Subject: [PATCH 2/2] Measure threadgroup memory against the cache, and confirm the pad calculation was right The shared-memory passes -- cooperative loads, padding, swizzling, double buffering -- all assume threadgroup memory is the fast place to put data. On this part that needed checking, because device memory is not slow when resident: 2386 GB/s under 2 MB against DRAM's 121. The question is not whether threadgroup memory is fast but whether it beats the cache that would have served the same bytes. It does, barely. 3361 GB/s contiguous against 2749 for the same read from resident device memory, so 1.22x. That is a much smaller margin than the usual assumption about scratchpad memory and it means no pass can justify staging on bandwidth alone. What it gains in peak it gives back in fragility. Across per-lane strides the threadgroup arm spreads 7.69x and the device arm 1.80x, which inverts the usual habit of treating shared memory as forgiving and device memory as the thing needing careful access. Here device memory is the forgiving one. The collapses are bank aliasing, and the sweep separates that from any size effect: stride shared device 16 B 3361 2749 128 B 1216 2079 <- power of two 144 B 2322 2341 256 B 605 2004 <- power of two 512 B 437 1641 <- power of two 144 bytes is one vector larger than 128 and reads nearly twice as fast, so this is alignment rather than distance. Thirty-two banks of four bytes puts every lane on the same bank at 128. That confirms `_optimal_pad`, which was written from a stated model of 32 four-byte banks and never measured. It pads to an odd stride, odd strides do not collapse, so the pass was right for the reason it claimed; a test now ties the two together so the measurement and the pass cannot drift apart. Getting here took a wrong diagnosis worth recording. A per-thread span of eight vectors measured about half speed, which looked like a 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. Padding is not a direction, it is an arithmetic result, and plus one is not automatically safe. The sweep replaced the guess. Controls throughout: elapsed time scales with the pass count at every stride, min 1.88x max 2.00x on doubling, so no loop is being collapsed. 686 pass. Lint and vulture clean. Co-Authored-By: Claude Opus 5 (1M context) --- benchmarks/agx_threadgroup_bandwidth.py | 178 ++++++++++++++++++++++++ metile/target/__init__.py | 10 ++ metile/target/agx.py | 51 +++++++ tests/test_target.py | 46 ++++++ 4 files changed, 285 insertions(+) create mode 100644 benchmarks/agx_threadgroup_bandwidth.py diff --git a/benchmarks/agx_threadgroup_bandwidth.py b/benchmarks/agx_threadgroup_bandwidth.py new file mode 100644 index 0000000..78ae1e4 --- /dev/null +++ b/benchmarks/agx_threadgroup_bandwidth.py @@ -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 +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()) diff --git a/metile/target/__init__.py b/metile/target/__init__.py index cf50086..322d56d 100644 --- a/metile/target/__init__.py +++ b/metile/target/__init__.py @@ -14,6 +14,10 @@ RESIDENT_WORKING_SET_BYTES, SCALAR_PEAK_TFLOPS, STREAMING_READ_GBPS, + THREADGROUP_CONFLICT_STRIDE_BYTES, + THREADGROUP_GBPS_BY_STRIDE, + THREADGROUP_OVER_RESIDENT, + THREADGROUP_PEAK_GBPS, Unavailable, ilp_headroom, inspect, @@ -21,6 +25,7 @@ read_bandwidth_gbps, resident, spills, + threadgroup_conflicts, tiling_gain, ) @@ -33,6 +38,10 @@ "RESIDENT_WORKING_SET_BYTES", "SCALAR_PEAK_TFLOPS", "STREAMING_READ_GBPS", + "THREADGROUP_CONFLICT_STRIDE_BYTES", + "THREADGROUP_GBPS_BY_STRIDE", + "THREADGROUP_OVER_RESIDENT", + "THREADGROUP_PEAK_GBPS", "Unavailable", "ilp_headroom", "inspect", @@ -40,5 +49,6 @@ "read_bandwidth_gbps", "resident", "spills", + "threadgroup_conflicts", "tiling_gain", ] diff --git a/metile/target/agx.py b/metile/target/agx.py index cb7b301..763091b 100644 --- a/metile/target/agx.py +++ b/metile/target/agx.py @@ -72,6 +72,41 @@ RESIDENT_WORKING_SET_BYTES = 2 * 1024 * 1024 RESIDENT_READ_GBPS = 2386.0 +# Threadgroup memory, measured by benchmarks/agx_threadgroup_bandwidth.py against a resident device read +# over the same 32 KB, so this compares staging with a cache hit rather than with DRAM. +# +# It is faster, but only just, and only when read contiguously: 3361 GB/s against the device arm's 2749, +# so 1.22x. That is a far smaller margin than the usual assumption about scratchpad memory, and it means +# a pass cannot justify staging on bandwidth alone. +# +# What it buys in peak it gives back in sensitivity. Across strides the threadgroup arm spreads 7.69x +# and the device arm 1.80x, so threadgroup memory is the more fragile of the two -- the opposite of the +# habit of treating shared memory as forgiving scratch and device memory as the thing needing careful +# access. +THREADGROUP_PEAK_GBPS = 3361.0 +THREADGROUP_OVER_RESIDENT = 1.22 + +# GB/s by per-lane stride in bytes, contiguous first. The collapses are the power-of-two strides from 128 +# bytes up; 144 bytes reads 2322 while 128 reads 1216, which is bank aliasing and not a size effect. +THREADGROUP_GBPS_BY_STRIDE = { + 16: 3361.0, + 32: 2658.0, + 48: 2978.0, + 64: 2032.0, + 80: 2737.0, + 96: 2510.0, + 112: 2468.0, + 128: 1216.0, + 144: 2322.0, + 192: 2032.0, + 256: 605.0, + 512: 437.0, +} + +# Smallest per-lane stride at which a power of two collapses threadgroup bandwidth. 32 banks of four +# bytes, so 128 bytes puts every lane on the same bank. +THREADGROUP_CONFLICT_STRIDE_BYTES = 128 + class Unavailable(RuntimeError): """The toolchain needed to inspect compiled kernels is not present.""" @@ -108,6 +143,22 @@ def resident(working_set_bytes): return 0 < working_set_bytes <= RESIDENT_WORKING_SET_BYTES +def threadgroup_conflicts(stride_bytes): + """Whether this per-lane stride puts threadgroup memory into bank conflict. + + Power-of-two strides from 128 bytes collapse it: 128 reads 1216 GB/s against 3361 contiguous, 256 + reads 605 and 512 reads 437, while 144 bytes -- one vector larger than 128 -- reads 2322. Device + memory shows nothing comparable, so this is a hazard staging introduces rather than one it avoids. + + `metile.compiler.passes._optimal_pad` already pads to an odd stride, which this confirms is the + right direction; the docstring's reasoning about 32 four-byte banks was never measured until now. + """ + if stride_bytes <= 0: + raise ValueError("a stride must be positive") + power_of_two = stride_bytes & (stride_bytes - 1) == 0 + return power_of_two and stride_bytes >= THREADGROUP_CONFLICT_STRIDE_BYTES + + def tiling_gain(working_set_bytes): """How much bandwidth a tiling wins by fitting this working set instead of streaming. diff --git a/tests/test_target.py b/tests/test_target.py index 1db3fae..d7bd242 100644 --- a/tests/test_target.py +++ b/tests/test_target.py @@ -148,3 +148,49 @@ def test_fitting_a_tile_outranks_every_other_lever_in_this_file(): functional_unit = MATRIX_PEAK_TFLOPS / min(SCALAR_PEAK_TFLOPS.values()) assert fitting > functional_unit > max(ILP_CEILING.values()) assert tiling_gain(2**30) == pytest.approx(1.0) + + +def test_threadgroup_memory_is_barely_faster_than_a_resident_device_read(): + """The margin matters because passes are built on the assumption that it is much faster. + + Measured 3361 GB/s contiguous against 2749 for the same bytes read from resident device memory. A + pass that stages data has to justify itself on something other than 1.22x, and this is the number + that says so. + """ + from metile.target import RESIDENT_READ_GBPS, THREADGROUP_OVER_RESIDENT, THREADGROUP_PEAK_GBPS + + assert 1.0 < THREADGROUP_OVER_RESIDENT < 1.5 + assert THREADGROUP_PEAK_GBPS > RESIDENT_READ_GBPS + + +def test_power_of_two_strides_from_128_bytes_are_flagged_as_conflicting(): + """The hazard staging introduces, which device memory does not have. + + 128 bytes reads 1216 GB/s against 3361 contiguous, 256 reads 605 and 512 reads 437, while 144 bytes + reads 2322. Odd strides are safe, which is what `_optimal_pad` already produces. + """ + from metile.target import threadgroup_conflicts + + for stride in (128, 256, 512, 1024): + assert threadgroup_conflicts(stride) + for stride in (16, 48, 112, 144, 192): + assert not threadgroup_conflicts(stride) + # Below the conflict stride a power of two is still fine: a float4 already spans four banks. + for stride in (16, 32, 64): + assert not threadgroup_conflicts(stride) + + +def test_the_existing_pad_calculation_avoids_the_measured_hazard(): + """Ties the measurement to the pass it validates. + + `_optimal_pad` was written from a stated model of 32 four-byte banks and never measured. It pads to + an odd stride, and the sweep confirms odd strides do not collapse, so the pass was right for the + reason it claimed. + """ + from metile.compiler.passes import _optimal_pad + from metile.target import threadgroup_conflicts + + for stride in (8, 16, 32, 64, 96, 128, 256): + padded = stride + _optimal_pad(stride) + assert padded % 2 == 1, f"stride {stride} padded to {padded}, which is even" + assert not threadgroup_conflicts(padded * 4), f"padded stride {padded} still conflicts"