diff --git a/metile/compiler/agx_schedule.py b/metile/compiler/agx_schedule.py new file mode 100644 index 0000000..f1d0ab6 --- /dev/null +++ b/metile/compiler/agx_schedule.py @@ -0,0 +1,345 @@ +"""Optimising G17 machine code directly, where reordering is not undone. + +meTile's other scheduler works on Metal IR and is measured to be inert: Apple's backend rebuilds +the schedule from the dataflow, so two source orders compile to byte-identical instructions and +statement order is a suggestion it declines. That is not a failure of the pass, it is where the +boundary of control sits when the output is MSL. + +This pass is on the other side of that boundary. It reads the instructions the backend produced, +rewrites them, and puts them back, so nothing downstream gets to re-derive anything. It is the +same two jobs as the IR pass — dependence-preserving reordering, and the arithmetic rewriting +that instruction-level parallelism needs — done where they survive. + +Register dependences are what make it possible at all. `metile.target.agx_isa` establishes that a +compact fma names its register in byte 0's high nibble and again as `(r << 1) | 1` in byte 1, so +which instructions actually depend on each other is readable rather than guessed, and two fmas on +different registers are known to be independent. + +Two transformations, both bit-exact by construction, because neither changes any arithmetic: + + simplify retire instructions that compute nothing. `a * 1` with no addend is the identity, + and the flag that retires an instruction was verified to be indistinguishable + from nopping it. + reorder move independent instructions, preserving every register dependence. + + collapse fold a run of identical fmas into one, which is what instruction-level + parallelism is reaching for taken to its limit: instead of shortening a dependent + chain of k steps, remove it. k steps of `a*m + d` equal one + `a*m**k + d*(m**(k-1) + ... + 1)`, and both constants are often in the immediate + field, so three instructions become one. Reassociation, so off by default. + +Collapsing is the ILP-family transform this level can express, and it took two attempts to find. +Splitting a chain across registers is the obvious approach and needs a free register, which cannot +be proved free by decoding — a register touched only by an instruction this file cannot read is +indistinguishable from an unused one, and instruction lengths cannot be recovered in general. The +metadata's register count would settle it, but the transform still needs the chain to be a sum, +and the only chains decodable here are multiply-accumulate. Collapsing needs neither: no new +register, no new instruction, just arithmetic on the constants already there. + +Off by default because collapsing reassociates, so results change in the last bits, the same trade +its IR counterpart makes. It also self-limits: `encode_immediate` refuses constants outside the +field, so a run whose collapsed constants do not land in it is left alone rather than approximated. +Whether retiring the folded instructions saves time is measurable and not assumed here; what it +certainly saves is a dependent chain of length k. + +Scope is deliberately narrow and self-enforcing. Only the compact f32 fma is decoded, and every +byte that is not part of one decoded instruction is a barrier nothing crosses. That is not +caution for its own sake: the instruction stream cannot be walked reliably, since the length +field theory that would allow it was measured wrong on eight of eight kernels, so an undecoded +region might be one instruction or twenty and moving anything across it is unsound. +""" + +from metile.target import agx_isa + +# The bit the compiler sets on every instruction of a run but the last. Reordering has to maintain +# it: it is positional, not a property of the operation, so moving an instruction to the end of a +# run without clearing it, or away from the end without setting it, changes the code's meaning +# rather than its order. +_CONTINUES = agx_isa.FmaFlag(2, 0x20, "more instructions follow in this run", "not the last") + + +class Fma: + """One decoded compact fma: `register = register * multiplier + addend`. + + The addend is one of three things, because the slot is: `addend` holds an immediate, + `addend_register` names a register, and a register index at or above the reachable range reads + zero, which is the only way a zero addend is expressible at all. + """ + + def __init__(self, offset, register, multiplier, addend, addend_register, negate_product, last): + self.offset = offset + self.register = register + self.multiplier = multiplier + self.addend = addend + self.addend_register = addend_register + self.negate_product = negate_product + self.last = last + + def addend_is_zero(self): + """Whether the addend contributes nothing, which needs the register range to decide.""" + return ( + self.addend is None + and self.addend_register is not None + and self.addend_register >= agx_isa.ARCHITECTURAL_REGISTERS + ) + + def is_identity(self): + """Whether this instruction leaves its register unchanged. + + `a * 1` plus something that reads zero. Deciding it needs the addend's register index, not + merely its absence: an fma whose addend slot names a live register adds that register, and + retiring it would drop a real term. An earlier version checked for `a * 1 + 0` as an + immediate, which cannot exist — the field holds `(1 + m/8) * 2**(e - 11)`, whose smallest + value is 2**-11, so it has no encoding for zero. + + A negated product is not the identity even at these constants, because it flips the sign. + """ + return not self.negate_product and self.multiplier == 1.0 and self.addend_is_zero() + + def encode(self, last=None): + register = self.addend_register + if self.addend is not None or self.addend_is_zero(): + register = None + return agx_isa.encode_fma( + self.register, + self.multiplier, + self.addend, + last=self.last if last is None else last, + negate_product=self.negate_product, + addend_register=register, + ) + + def __repr__(self): + if self.addend is not None: + addend = f" + {self.addend:g}" + elif self.addend_is_zero(): + addend = "" + else: + addend = f" + r{self.addend_register}" + sign = "-" if self.negate_product else "" + return f"Fma(0x{self.offset:04x}: r{self.register} = {sign}r{self.register} * {self.multiplier:g}{addend})" + + +def decode(text, offsets): + """Decode the compact fmas at `offsets`, which must have been confirmed behaviourally. + + Offsets are required rather than discovered. Pattern matching proposes candidates and is + sometimes right, but a wrong offset here does not fail loudly: it decodes some other + instruction's bytes as an fma and re-encoding them corrupts the kernel silently. Callers get + their offsets from `agx_isa.boundaries`, which confirms each one by checking that nopping it + removes exactly one operation's contribution. + """ + found = [] + for offset in offsets: + window = text[offset : offset + agx_isa.FMA_LENGTH] + if len(window) < agx_isa.FMA_LENGTH: + raise ValueError(f"offset 0x{offset:04x} runs past the end of the code") + if window[0] & 0x0F != agx_isa.FMA_OPCODE_NIBBLE: + raise ValueError(f"offset 0x{offset:04x} is not a compact fma") + register = window[0] >> 4 + if window[1] != (register << 1) | 1: + raise ValueError( + f"offset 0x{offset:04x} disagrees with itself about its register: " + f"byte 0 says {register}, byte 1 says {window[1] >> 1}" + ) + addend = None + addend_register = None + if agx_isa.read_flag(window, 0, agx_isa.ADDEND_IMMEDIATE): + addend = agx_isa.decode_immediate(window[agx_isa.FMA_ADDEND_BYTE]) + if agx_isa.read_flag(window, 0, agx_isa.ADDEND_NEGATE): + addend = -addend + else: + addend_register = window[agx_isa.FMA_ADDEND_BYTE] >> 1 + found.append( + Fma( + offset=offset, + register=register, + multiplier=agx_isa.decode_immediate(window[agx_isa.FMA_MULTIPLIER_BYTE]), + addend=addend, + addend_register=addend_register, + negate_product=agx_isa.read_flag(window, 0, agx_isa.PRODUCT_NEGATE), + last=not agx_isa.read_flag(window, 0, _CONTINUES), + ) + ) + return found + + +def _runs(instructions): + """Group instructions into contiguous runs, split wherever undecoded bytes intervene. + + Two decoded fmas are in the same run only when they are adjacent in the code. Anything + between them is unidentified, and since instruction lengths cannot be recovered in general + there is no way to know what it does, so it bounds the region a reordering may touch. + """ + groups = [] + current = [] + for instruction in instructions: + if current and instruction.offset != current[-1].offset + agx_isa.FMA_LENGTH: + groups.append(current) + current = [] + current.append(instruction) + if current: + groups.append(current) + return groups + + +def simplify(text, offsets): + """Retire instructions that compute nothing. Returns (code, count). + + Uses the disable flag rather than nops so lengths and every neighbouring byte stay exactly as + they were; the flag was verified to produce results indistinguishable from nopping. + """ + patched = text + retired = 0 + for instruction in decode(text, offsets): + if not instruction.is_identity(): + continue + patched = agx_isa.write_flag(patched, instruction.offset, agx_isa.INSTRUCTION_DISABLE, True) + retired += 1 + return patched, retired + + +def collapse(text, offsets): + """Fold runs of identical fmas into one instruction each. Returns (code, folded count). + + Reassociates, so it changes results in the last bits and is off by default in `optimize`. + + A run qualifies when its instructions are adjacent, target the same register, carry the same + immediate multiplier and addend, and do not negate. Anything else is left alone: an addend that + names a register is excluded because its value can change between the steps, which is exactly + what makes the closed form invalid. + """ + instructions = decode(text, offsets) + patched = bytearray(text) + folded = 0 + + for run in _runs(instructions): + start = 0 + while start < len(run): + head = run[start] + stop = start + 1 + while stop < len(run) and _foldable(head, run[stop]): + stop += 1 + length = stop - start + if length < 2 or head.addend is None: + start = stop + continue + multiplier = head.multiplier**length + addend = head.addend * sum(head.multiplier**step for step in range(length)) + try: + built = agx_isa.encode_fma( + head.register, multiplier, addend, last=run[stop - 1].last + ) + except agx_isa.EncodingError: + # The closed form left the immediate field. Approximating would change the result + # by more than reassociation does, so the run keeps its instructions. + start = stop + continue + patched[head.offset : head.offset + agx_isa.FMA_LENGTH] = built + for absorbed in run[start + 1 : stop]: + patched = bytearray( + agx_isa.write_flag( + bytes(patched), absorbed.offset, agx_isa.INSTRUCTION_DISABLE, True + ) + ) + folded += 1 + start = stop + return bytes(patched), folded + + +def _foldable(head, other): + """Whether `other` continues a run that folds into a single instruction with `head`.""" + return ( + other.register == head.register + and other.multiplier == head.multiplier + and other.addend is not None + and head.addend is not None + and other.addend == head.addend + and not other.negate_product + and not head.negate_product + ) + + +def reorder(text, offsets): + """Reorder independent instructions within each run. Returns (code, moved count). + + Bit-exact by construction: instructions are moved, never rewritten, so every register still + sees the same operations in an order that respects every dependence. Two fmas touching + different registers are independent, and two touching the same one are not, which is the whole + dependence relation for this instruction form — it both reads and writes exactly one register. + + Ties break towards the original position, so a run with no freedom comes back byte-identical + rather than churned into an equivalent ordering. + """ + instructions = decode(text, offsets) + patched = bytearray(text) + moved = 0 + + for run in _runs(instructions): + # Stable grouping by register: a register's own instructions keep their relative order, + # which is what preserves the dependences, while whole registers may interleave. Emitting + # one register's chain at a time is the schedule that shortens no dependence but also + # breaks none, and it is the only reordering this form permits without renaming. + by_register = {} + for instruction in run: + by_register.setdefault(instruction.register, []).append(instruction) + if len(by_register) < 2: + continue + + ordered = [] + while by_register: + for register in sorted(by_register): + ordered.append(by_register[register].pop(0)) + by_register = {register: rest for register, rest in by_register.items() if rest} + + for position, instruction in enumerate(ordered): + offset = run[position].offset + if instruction.offset != offset: + moved += 1 + patched[offset : offset + agx_isa.FMA_LENGTH] = instruction.encode( + last=(position == len(ordered) - 1 and run[-1].last) + ) + return bytes(patched), moved + + +def optimize(text, offsets, simplify_identities=True, reorder_independent=True, fold_runs=False): + """Run the machine-level passes in order. Returns (code, report). + + The single entry point, so a caller asks for optimisation rather than for a list of + transformations. Simplification runs first: an instruction it retires is one reordering does + not have to place, and retiring changes no offsets, so the second pass sees the same layout. + + Folding runs comes first when asked for, because it changes which instructions exist and the + two passes after it should see the result. It is off by default: it reassociates, and the model + tests assert bit-exact logits. + + `metile.compiler.scheduling` is the IR-level half. Both are native to the compiler; they differ + in which side of the MSL boundary they act on, and only this side survives it. + """ + report = {"retired": 0, "moved": 0, "folded": 0} + if fold_runs: + text, report["folded"] = collapse(text, offsets) + if simplify_identities: + text, report["retired"] = simplify(text, offsets) + if reorder_independent: + text, report["moved"] = reorder(text, offsets) + return text, report + + +def summarise(text, offsets): + """What the pass can see and what it would do, without doing it. + + For reporting and for deciding whether a kernel is worth rewriting at all. Reports the runs + it found rather than a single count, because a run of one instruction offers nothing to + reorder however many such runs there are. + """ + instructions = decode(text, offsets) + runs = _runs(instructions) + registers = {instruction.register for instruction in instructions} + return { + "instructions": len(instructions), + "runs": [len(run) for run in runs], + "registers": sorted(registers), + "identities": sum(1 for instruction in instructions if instruction.is_identity()), + "reorderable": sum(len(run) for run in runs if len({i.register for i in run}) > 1), + } diff --git a/metile/target/agx.py b/metile/target/agx.py index 2cd8878..ee56d9e 100644 --- a/metile/target/agx.py +++ b/metile/target/agx.py @@ -143,7 +143,14 @@ def _compiled(source, function, workdir): [str(binary), str(metal), function, str(archive)], capture_output=True, text=True ) if built.returncode != 0: - raise RuntimeError(built.stderr.strip()[:300]) + message = built.stderr.strip() + # A device that will not serialize a binary archive cannot be inspected at all, which is a + # property of the machine and not a fault in the kernel. Reporting it as Unavailable, the + # way a missing swiftc is reported, lets callers skip; raising RuntimeError made every + # machine-code test fail on a CI runner rather than opt out of a capability it lacks. + if "MTLBinaryArchive" in message or "eligible to be serialized" in message: + raise Unavailable(f"this device does not serialize binary archives: {message[:200]}") + raise RuntimeError(message[:300]) subprocess.run( ["xcrun", "metal-lipo", str(archive), "-thin", _gpu_arch(archive), "-output", str(thin)], diff --git a/metile/target/agx_isa.py b/metile/target/agx_isa.py index e481111..5b22a74 100644 --- a/metile/target/agx_isa.py +++ b/metile/target/agx_isa.py @@ -37,12 +37,18 @@ emitted — a*3+0.5, a*1.5-2, a*7 with no addend, and -a*2+1 — each ran exactly as predicted on four inputs. - flags three bits of the compact fma control its arithmetic, each found by scanning - all 256 values of its byte and then predicting the result of setting it - across a chain on four inputs. Twelve predictions, all exact: byte 2 bit - 0x10 negates the product, byte 4 bit 0x10 negates the addend, byte 4 bit - 0x20 drops the addend entirely and leaves a multiply. So a compiled fma can - be rewritten into `-a*m+d`, `a*m-d` or `a*m` without recompiling it. + flags bits of the compact fma control its arithmetic, each found by scanning all 256 + values of its byte and then predicting the result of setting it across a chain + on four inputs. Byte 2 bit 0x10 negates the product and byte 4 bit 0x10 + negates the addend, twelve exact predictions between them. + + Byte 4 bit 0x20 was first read as "include the addend", because clearing it + turned a*m+d into a*m. That reading survived a prediction and was still wrong: + it selects whether the addend slot is an immediate or a *register*, and the + constant left in the slot named a register beyond the reachable range, which + reads zero. The addend was never absent, it was zero. Naming the register form + gives the register-plus-register add as `rd * 1 + rs`, verified across eighteen + register pairs on four threads each. structure `06 00` is a two-byte nop and `0e 00 00 00` ends a block. Blocks are padded with nops to a 64-byte boundary. Giving an eight-byte instruction the nop's @@ -87,9 +93,11 @@ # instruction of a run, so the compiler emits 0x2e throughout and 0x0e at the end. Bit # 0x01 turned `a*2+1` into `a*a`, producing 225 from an accumulator holding 15. # 3 multiplier immediate, in the format above. -# 4 addend control. Bit 0x20 includes the addend, and clearing it leaves a plain multiply. -# Bit 0x10 negates it. -# 5 addend immediate, same format with the low bit clear. +# 4 addend control. Bit 0x20 selects an immediate in the addend slot; clear selects a register. +# Bit 0x10 negates the addend. +# 5 the addend: an immediate in the format above with the low bit clear, or a register as +# `r << 1` when bit 0x20 of byte 4 is clear. There is no zero immediate, so a zero addend is +# a register index beyond the reachable range, which reads zero. # 6 bit 0x20 retires the instruction. Bits 0x40 and 0x80 change where the result goes; the # low four bits made no difference to the result at all. # 7 register selection, like byte 1. @@ -356,7 +364,15 @@ def execute(source, function, inputs, rewrite=None, workdir=".metile-agx"): [str(prober), str(metal), function, str(archive)], capture_output=True, text=True ) if built.returncode != 0: - raise RuntimeError(built.stderr.strip()[:300]) + message = built.stderr.strip() + # Same distinction `agx.machine_code` draws: a device that will not serialize an archive + # cannot run one either, and that is a capability the machine lacks rather than an error in + # the kernel, so callers get Unavailable and can skip. + if "MTLBinaryArchive" in message or "eligible to be serialized" in message: + raise agx.Unavailable( + f"this device does not serialize binary archives: {message[:200]}" + ) + raise RuntimeError(message[:300]) raw = archive.read_bytes() target = archive diff --git a/tests/test_agx_isa.py b/tests/test_agx_isa.py index c8da19c..944aa4f 100644 --- a/tests/test_agx_isa.py +++ b/tests/test_agx_isa.py @@ -122,7 +122,12 @@ def rewrite(original): def test_a_rewrite_that_changes_length_is_refused(): - """Shifting the bytes after the patch would move the metadata the driver reads.""" + """Shifting the bytes after the patch would move the metadata the driver reads. + + `_machine_code` is called for its skip, not its value: a device that will not serialize a + binary archive cannot run one either, and this test would otherwise fail there rather than + opt out. That is exactly what happened on CI. + """ _machine_code() with pytest.raises(agx_isa.EncodingError, match="keep the length"): agx_isa.execute(CHAIN, "probe", [1.0], rewrite=lambda text: text[:-2]) @@ -144,7 +149,12 @@ def test_nopping_an_instruction_removes_exactly_its_effect(): intact=31.0, removed=15.0, ) - assert len(offsets) == 4, f"expected four fmas, found {[hex(o) for o in offsets]}" + assert len(offsets) == 4, ( + f"expected four fmas, found {[hex(o) for o in offsets]}. The scan skips an offset whose " + f"patched kernel fails to dispatch, and heavy concurrent GPU work makes that happen to " + f"offsets that are perfectly valid, so run this without other GPU load before treating a " + f"short list as an ISA change." + ) assert {b - a for a, b in itertools.pairwise(offsets)} == {agx_isa.FMA_LENGTH} diff --git a/tests/test_agx_schedule.py b/tests/test_agx_schedule.py new file mode 100644 index 0000000..30c7bfe --- /dev/null +++ b/tests/test_agx_schedule.py @@ -0,0 +1,290 @@ +"""Optimising machine code has to be verified by running it, so that is what these do. + +The IR scheduler can be checked structurally: its output is IR, and a hazard either was or was +not respected. This pass emits instructions, and an unsound rewrite there produces a kernel that +compiles, dispatches, and returns a wrong number. So the two transformations that claim to be +bit-exact are asserted bit-exact against a GPU run, not against a model of one. +""" + +import pytest + +from metile.compiler import agx_schedule +from metile.target import agx_isa + +CHAINS = ((2.0, "a"), (3.0, "b"), (5.0, "c")) +SOURCE = ( + """#include +using namespace metal; +kernel void probe(device const float* x [[buffer(0)]], + device float* out [[buffer(1)]], + constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + float a = x[gid + 0]; + float b = x[gid + 1]; + float c = x[gid + 2]; +""" + + "\n".join(f" {v} = fma({v}, {m}f, 1.0f);" for m, v in CHAINS for _ in range(3)) + + """ + out[gid] = a + b + c; +} +""" +) +INPUTS = [1.0, 2.0, 3.0, 4.0] + + +def _code_and_offsets(): + from metile.target import Unavailable, machine_code + + try: + text = machine_code(SOURCE, "probe") + except Unavailable as error: + pytest.skip(f"no Metal toolchain: {error}") + wanted = {agx_isa.encode_immediate(m, low_bit=1) for m, _ in CHAINS} + offsets = [ + offset + for offset in range(0, len(text) - agx_isa.FMA_LENGTH, 2) + if text[offset] & 0x0F == agx_isa.FMA_OPCODE_NIBBLE + and text[offset + 2] & 0x0F == 0x0E + and text[offset + agx_isa.FMA_MULTIPLIER_BYTE] in wanted + ] + return text, offsets + + +def test_decoding_recovers_the_arithmetic_the_source_asked_for(): + """Multipliers and registers have to come back, or nothing downstream means anything.""" + text, offsets = _code_and_offsets() + decoded = agx_schedule.decode(text, offsets) + assert decoded, "no compact fmas were found" + multipliers = {instruction.multiplier for instruction in decoded} + assert multipliers <= {m for m, _ in CHAINS} + # One register per chain, and each instruction reads and writes the same one. + assert len({instruction.register for instruction in decoded}) == len(CHAINS) + for instruction in decoded: + assert instruction.addend == 1.0 + assert not instruction.negate_product + + +def test_decoding_rejects_an_offset_that_is_not_a_compact_fma(): + """A wrong offset must fail here, because downstream it silently corrupts the kernel.""" + text, offsets = _code_and_offsets() + with pytest.raises(ValueError, match="not a compact fma"): + agx_schedule.decode(text, [offsets[0] + 1]) + + +def test_reordering_moves_instructions_and_changes_no_result(): + """The claim the pass exists to make, checked on the GPU rather than argued. + + Three chains on three registers, so there is real freedom, and distinct multipliers so a + misordering shows up in the arithmetic instead of hiding behind identical operations. + """ + text, offsets = _code_and_offsets() + baseline = agx_isa.execute(SOURCE, "probe", INPUTS) + + reordered, moved = agx_schedule.reorder(text, offsets) + assert moved > 0, "nothing moved, so this asserts nothing about soundness" + assert reordered != text + + assert agx_isa.execute(SOURCE, "probe", INPUTS, rewrite=lambda _: reordered) == baseline + + +def test_reordering_a_single_register_run_is_a_no_op(): + """No freedom means no churn, and no chance to break something for nothing.""" + text, offsets = _code_and_offsets() + one_register = [ + instruction.offset + for instruction in agx_schedule.decode(text, offsets) + if instruction.register == agx_schedule.decode(text, offsets)[0].register + ] + reordered, moved = agx_schedule.reorder(text, one_register) + assert moved == 0 + assert reordered == text + + +def test_retiring_an_identity_instruction_changes_no_result(): + """`a * 1` computes nothing, so removing it must be invisible. + + The identity is planted rather than waited for, because the compiler does not emit one. That + is the point of the transformation: it exists for code another pass produced. + """ + text, offsets = _code_and_offsets() + planted = bytearray(text) + victim = offsets[1] + planted[victim : victim + agx_isa.FMA_LENGTH] = agx_isa.encode_fma( + text[victim] >> 4, + 1.0, + None, + last=not agx_isa.read_flag(text, victim, agx_schedule._CONTINUES), + ) + planted = bytes(planted) + + assert agx_schedule.summarise(planted, offsets)["identities"] == 1 + simplified, retired = agx_schedule.simplify(planted, offsets) + assert retired == 1 + + before = agx_isa.execute(SOURCE, "probe", INPUTS, rewrite=lambda _: planted) + after = agx_isa.execute(SOURCE, "probe", INPUTS, rewrite=lambda _: simplified) + assert after == before + + +def test_an_identity_is_only_an_identity_without_a_sign_flip(): + """`-a * 1` is a negation, not a no-op, and retiring it would change the answer.""" + negated = agx_isa.encode_fma(0, 1.0, None, negate_product=True) + assert not agx_schedule.decode(negated, [0])[0].is_identity() + assert agx_schedule.decode(agx_isa.encode_fma(0, 1.0, None), [0])[0].is_identity() + + +def test_summarise_reports_runs_rather_than_a_single_count(): + """A run of one offers nothing to reorder however many of them there are.""" + text, offsets = _code_and_offsets() + summary = agx_schedule.summarise(text, offsets) + assert summary["instructions"] == len(offsets) + assert sum(summary["runs"]) == len(offsets) + assert len(summary["registers"]) == len(CHAINS) + + +def test_optimize_runs_both_passes_and_reports_what_it_did(): + """One entry point, and its report has to match what the individual passes claim.""" + text, offsets = _code_and_offsets() + baseline = agx_isa.execute(SOURCE, "probe", INPUTS) + + optimised, report = agx_schedule.optimize(text, offsets) + assert report["moved"] == agx_schedule.reorder(text, offsets)[1] + assert report["retired"] == agx_schedule.simplify(text, offsets)[1] + assert agx_isa.execute(SOURCE, "probe", INPUTS, rewrite=lambda _: optimised) == baseline + + +def test_optimize_with_everything_switched_off_returns_the_code_untouched(): + text, offsets = _code_and_offsets() + untouched, report = agx_schedule.optimize( + text, offsets, simplify_identities=False, reorder_independent=False + ) + assert untouched == text + assert report == {"retired": 0, "moved": 0, "folded": 0} + + +def test_an_fma_adding_a_live_register_is_not_an_identity(): + """`a * 1 + r3` adds a real term, so retiring it would drop it. + + This is the property the addend-flag correction bought. When the flag was read as + "addend present or absent", `a * 1` with anything in the slot looked like a no-op, and this + instruction would have been retired and a term silently lost. Deciding it needs the register + index, because only an index beyond the reachable range reads zero. + """ + live = agx_isa.encode_fma(2, 1.0, addend_register=3) + decoded = agx_schedule.decode(live, [0])[0] + assert decoded.addend_register == 3 + assert not decoded.addend_is_zero() + assert not decoded.is_identity() + + zero = agx_schedule.decode(agx_isa.encode_fma(2, 1.0), [0])[0] + assert zero.addend_is_zero() + assert zero.is_identity() + + +def test_decoding_round_trips_every_addend_form(): + """Re-encoding a decoded instruction has to reproduce it, or a rewrite silently changes it.""" + for original in ( + agx_isa.encode_fma(2, 2.0, 1.0), + agx_isa.encode_fma(2, 2.0, -1.0), + agx_isa.encode_fma(2, 1.5, addend_register=3), + agx_isa.encode_fma(2, 1.0), + agx_isa.encode_fma(2, 2.0, 1.0, negate_product=True), + agx_isa.encode_fma(2, 2.0, 1.0, last=True), + ): + assert agx_schedule.decode(original, [0])[0].encode() == original + + +CONSTANT_CHAIN = """#include +using namespace metal; +kernel void probe(device const float* x [[buffer(0)]], + device float* out [[buffer(1)]], + constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + float a = x[gid]; + a = fma(a, 2.0f, 1.0f); + a = fma(a, 2.0f, 1.0f); + a = fma(a, 2.0f, 1.0f); + a = fma(a, 2.0f, 1.0f); + out[gid] = a; +} +""" + + +def _constant_chain_offsets(): + from metile.target import Unavailable, machine_code + + try: + text = machine_code(CONSTANT_CHAIN, "probe") + except Unavailable as error: + pytest.skip(f"no Metal toolchain: {error}") + offsets = [ + offset + for offset in range(0, len(text) - agx_isa.FMA_LENGTH, 2) + if text[offset] & 0x0F == agx_isa.FMA_OPCODE_NIBBLE + and text[offset + 2] & 0x0F == 0x0E + and text[offset + agx_isa.FMA_MULTIPLIER_BYTE] == 0xC1 + ] + return text, offsets + + +def test_collapsing_a_run_folds_it_into_one_instruction(): + """Three steps of a*2+1 are one a*8+7, and the GPU has to agree. + + The closed form for k steps is a*m**k + d*(m**(k-1) + ... + 1). Here the constants are powers + of two so the fold is exact and can be compared against the untouched kernel as well as against + the prediction; in general it reassociates, which is why it is off by default. + """ + text, offsets = _constant_chain_offsets() + assert len(offsets) == 3 + baseline = agx_isa.execute(CONSTANT_CHAIN, "probe", [1.0, 2.0, 3.0, 5.0]) + + folded_code, folded = agx_schedule.collapse(text, offsets) + assert folded == 1 + + remaining = agx_schedule.decode(folded_code, offsets) + assert remaining[0].multiplier == 8.0 + assert remaining[0].addend == 7.0 + assert all( + agx_isa.read_flag(folded_code, instruction.offset, agx_isa.INSTRUCTION_DISABLE) + for instruction in remaining[1:] + ) + + got = agx_isa.execute( + CONSTANT_CHAIN, "probe", [1.0, 2.0, 3.0, 5.0], rewrite=lambda _: folded_code + ) + assert got == baseline + + +def test_collapsing_leaves_a_run_alone_when_the_closed_form_escapes_the_field(): + """Approximating would change the result by more than reassociating does. + + Six steps of a*2+1 close to a*64 + 63, and 63 is not `(1 + m/8) * 2**e` for any m under eight, + so the fold is declined rather than rounded. + """ + instruction = agx_isa.encode_fma(0, 2.0, 1.0) + run = instruction * 6 + offsets = [index * agx_isa.FMA_LENGTH for index in range(6)] + with pytest.raises(agx_isa.EncodingError): + agx_isa.encode_immediate(63.0, low_bit=0) + unchanged, folded = agx_schedule.collapse(run, offsets) + assert folded == 0 + assert unchanged == run + + +def test_collapsing_declines_a_register_addend(): + """A register's value can change between steps, so the closed form does not hold.""" + instruction = agx_isa.encode_fma(0, 2.0, addend_register=3) + run = instruction * 3 + offsets = [index * agx_isa.FMA_LENGTH for index in range(3)] + unchanged, folded = agx_schedule.collapse(run, offsets) + assert folded == 0 + assert unchanged == run + + +def test_optimize_leaves_folding_off_unless_asked(): + """It reassociates, and the model tests assert bit-exact logits.""" + text, offsets = _constant_chain_offsets() + _, report = agx_schedule.optimize(text, offsets) + assert report["folded"] == 0 + _, asked = agx_schedule.optimize(text, offsets, fold_runs=True) + assert asked["folded"] == 1