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/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 a1c593a..322d56d 100644 --- a/metile/target/__init__.py +++ b/metile/target/__init__.py @@ -6,27 +6,49 @@ """ 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, + THREADGROUP_CONFLICT_STRIDE_BYTES, + THREADGROUP_GBPS_BY_STRIDE, + THREADGROUP_OVER_RESIDENT, + THREADGROUP_PEAK_GBPS, Unavailable, ilp_headroom, inspect, machine_code, + read_bandwidth_gbps, + resident, spills, + threadgroup_conflicts, + 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", + "THREADGROUP_CONFLICT_STRIDE_BYTES", + "THREADGROUP_GBPS_BY_STRIDE", + "THREADGROUP_OVER_RESIDENT", + "THREADGROUP_PEAK_GBPS", "Unavailable", "ilp_headroom", "inspect", "machine_code", + "read_bandwidth_gbps", + "resident", "spills", + "threadgroup_conflicts", + "tiling_gain", ] diff --git a/metile/target/agx.py b/metile/target/agx.py index ee56d9e..763091b 100644 --- a/metile/target/agx.py +++ b/metile/target/agx.py @@ -37,6 +37,76 @@ 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 + +# 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.""" @@ -52,6 +122,53 @@ 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 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. + + 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..d7bd242 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,96 @@ 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) + + +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"