Assemble G17 fma instructions from scratch, registers included

benchmarks/agx_isa_probe.py

@@ -11,6 +11,7 @@
     3  immediates         read the constant field out of kernels compiled with known constants
     4  encode             synthesise constants the compiler never emitted and predict the answer
     5  flags              set one arithmetic bit at a time and predict the answer again
+    6  assemble           build instructions from scratch and predict the answer once more
 
 Stage 4 is the one that matters. Stages 1 to 3 could all be satisfied by a field map that is
 merely consistent with what the compiler happens to emit; only predicting the result of bytes
@@ -31,6 +32,9 @@
 
 from metile.target import agx, agx_isa
 
+# The bit the compiler sets on every instruction of a run except the last.
+_CONTINUES = agx_isa.FmaFlag(2, 0x20, "more instructions follow in this run", "not the last")
+
 # Four dependent fmas of a*2+1. Dependent so the backend cannot reorder them, and with distinct
 # results per step so a removed instruction is visible in the output rather than absorbed.
 CHAIN = """#include <metal_stdlib>
@@ -200,11 +204,51 @@ def rewrite(original, f=flag, c=clear, where=tuple(compact)):
         shown = ", ".join(f"{value:g}" for value in predicted)
         print(f"   {label:<30}{shown:>28}   {'MATCHES' if agree else f'got {got}'}")
 
+    print("\n6. Can whole instructions be assembled? Reproduce the compiler, then go beyond it.")
+    for offset in compact:
+        original = bytes(text[offset : offset + agx_isa.FMA_LENGTH])
+        built = agx_isa.encode_fma(
+            original[0] >> 4, 2.0, 1.0, last=not agx_isa.read_flag(original, 0, _CONTINUES)
+        )
+        same = built == original
+        checks.append(same)
+        print(
+            f"   0x{offset:04x}  compiler {original.hex(' ')}  "
+            f"encoder {built.hex(' ')}  {'IDENTICAL' if same else 'DIFFERS'}"
+        )
+
+    print(f"\n   {'synthesised form':<24}{'predicted':>30}   verdict")
+    for fields, step, label in (
+        ({"multiplier": 3.0, "addend": 0.5}, lambda v: v * 3.0 + 0.5, "a*3+0.5"),
+        ({"multiplier": 1.5, "addend": -2.0}, lambda v: v * 1.5 - 2.0, "a*1.5-2"),
+        ({"multiplier": 7.0, "addend": None}, lambda v: v * 7.0, "a*7, no addend"),
+    ):
+
+        def rewrite(original, f=fields, where=tuple(compact)):
+            patched = bytearray(original)
+            for index, offset in enumerate(where):
+                patched[offset : offset + agx_isa.FMA_LENGTH] = agx_isa.encode_fma(
+                    original[offset] >> 4, last=(index == len(where) - 1), **f
+                )
+            return bytes(patched)
+
+        predicted = []
+        for value in inputs:
+            running = value * 2.0 + 1.0
+            for _ in compact:
+                running = step(running)
+            predicted.append(running)
+        got = agx_isa.execute(CHAIN, "probe", inputs, rewrite=rewrite, workdir=work)
+        agree = got == predicted
+        checks.append(agree)
+        shown = ", ".join(f"{value:g}" for value in predicted)
+        print(f"   {label:<24}{shown:>30}   {'MATCHES' if agree else f'got {got}'}")
+
     print()
     if all(checks) and all(observed):
-        print("Both the constant field and the arithmetic flags are encodable, not merely")
-        print("readable: every prediction made before running matched the GPU exactly, for")
-        print("bytes no Metal compiler produced.")
+        print("The compact fma is assemblable, not merely readable. The encoder reproduces the")
+        print("compiler's own bytes exactly, and every prediction made before running matched the")
+        print("GPU, for instructions no Metal compiler produced.")
     else:
         print("A prediction missed. The field map in metile/target/agx_isa.py is wrong here,")
         print("or this toolchain encodes constants differently; re-derive before relying on it.")

metile/target/agx_isa.py

@@ -25,6 +25,18 @@
                    rewriting three fmas to `a*6+7` gives 949 from x=1, and to `a*1.25+1.5`
                    gives 11.578125. Both were predicted before running.
 
+    registers      the register index appears twice in the compact form, as byte 0's high
+                   nibble and as `(r << 1) | 1` in byte 1, and the two agreed across every
+                   instruction examined. Confirmed by redirecting an instruction onto another
+                   chain's register and predicting the whole kernel's output: three redirects,
+                   three exact matches.
+
+    assembly       with registers, constants and flags all measured, `encode_fma` assembles the
+                   form from scratch. It reproduces the compiler's own bytes exactly for the
+                   cases the compiler emits, and four synthesised forms the compiler never
+                   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
@@ -134,6 +146,69 @@ def __repr__(self):
 # put it on the same footing as the flags above, which were each predicted on four inputs.
 
 
+# Register selection, established across three independent fma chains and then confirmed by
+# redirecting an instruction onto another chain's register and predicting the whole kernel's
+# output. Three redirects, three exact matches. The index appears twice: as byte 0's high nibble
+# and as `(r << 1) | 1` in byte 1, and the two agreed in all eight instructions examined.
+FMA_REGISTER_HIGH_BYTE = 0
+FMA_REGISTER_BYTE = 1
+FMA_MAX_REGISTER = 15
+
+# The remaining constant bytes of the form, taken as the compiler writes them. Byte 7 was 0x22 or
+# 0x42 on the first fma of a chain and 0x02 everywhere else, which looks like a dependency or wait
+# field; 0x02 is what a synthesised instruction uses, and synthesised instructions run correctly
+# with it.
+_FMA_ADDEND_CONTROL = 0x21
+_FMA_TAIL = bytes((0x02, 0x02))
+_FMA_NOT_LAST = 0x20
+_FMA_MODE_BASE = 0x0E
+
+
+def encode_fma(register, multiplier, addend=None, last=False, negate_product=False):
+    """Assemble a complete compact fma: `register = register * multiplier (+/- addend)`.
+
+    Every field comes from a measurement that was checked by prediction, so this is an encoder
+    rather than a template with holes: pass the register, the constants and the signs and the
+    eight bytes come out. Reproduces the compiler's own encoding exactly for the cases it emits,
+    which is the cheapest available check that the assembly is right.
+
+    `addend` of None drops the addend and leaves a plain multiply. A negative addend is encoded
+    by setting the negate flag, which is how the field reaches values its unsigned immediate
+    cannot.
+    """
+    if not 0 <= register <= FMA_MAX_REGISTER:
+        raise EncodingError(f"register {register} is outside the field")
+    mode = _FMA_MODE_BASE
+    if not last:
+        mode |= _FMA_NOT_LAST
+    if negate_product:
+        mode |= PRODUCT_NEGATE.mask
+    # Dropping the addend clears bit 0x20 of the control byte and leaves the immediate slot
+    # holding an ordinary encoded constant, because zero there is not inert. Writing 0x00 was
+    # tried and the kernel computed a*m + a: an eight-fold growth per step where seven was
+    # predicted, from a synthesised `a*7`. Zero selects a register operand rather than meaning
+    # "no operand", and register 0 happened to be the accumulator. The value written here is
+    # ignored once the control bit is clear, which is the configuration the flag scan verified.
+    control = _FMA_ADDEND_CONTROL & ~ADDEND_ENABLE.mask
+    addend_byte = encode_immediate(1.0, low_bit=0)
+    if addend is not None:
+        control = _FMA_ADDEND_CONTROL
+        if addend < 0:
+            control |= ADDEND_NEGATE.mask
+        addend_byte = encode_immediate(abs(addend), low_bit=0)
+    return bytes(
+        (
+            (register << 4) | FMA_OPCODE_NIBBLE,
+            (register << 1) | 1,
+            mode,
+            encode_immediate(multiplier, low_bit=1),
+            control,
+            addend_byte,
+            *_FMA_TAIL,
+        )
+    )
+
+
 def read_flag(text, offset, flag):
     """Whether one flag is set on the instruction at `offset`."""
     return bool(text[offset + flag.byte] & flag.mask)

tests/test_agx_isa.py

@@ -230,3 +230,84 @@ def by_nop(original):
     nopped = agx_isa.execute(CHAIN, "probe", inputs, rewrite=by_nop)
     assert flagged == nopped
     assert flagged[0] == 15.0  # three fmas of a*2+1 from x=1 instead of four
+
+
+# What the compiler emitted for `a = fma(a, 2.0f, 1.0f)` on register 0, read out of a compiled
+# kernel: the same instruction twice with the continue flag set, and once with it clear because it
+# ends the run.
+COMPILER_FMA = bytes.fromhex("09012ec121b00202")
+COMPILER_FMA_LAST = bytes.fromhex("09010ec121b00202")
+
+
+def test_the_assembler_reproduces_the_compilers_own_bytes():
+    """The cheapest check that assembly is right: agree with the only other assembler available."""
+    assert agx_isa.encode_fma(0, 2.0, 1.0, last=False) == COMPILER_FMA
+    assert agx_isa.encode_fma(0, 2.0, 1.0, last=True) == COMPILER_FMA_LAST
+    # Registers 2 and 3 as the compiler wrote them for two other accumulators.
+    assert agx_isa.encode_fma(2, 2.0, 1.0)[:2] == bytes((0x29, 0x05))
+    assert agx_isa.encode_fma(3, 3.0, 1.0)[:2] == bytes((0x39, 0x07))
+
+
+def test_the_assembler_refuses_a_register_outside_the_field():
+    for register in (-1, 16, 255):
+        with pytest.raises(agx_isa.EncodingError, match="register"):
+            agx_isa.encode_fma(register, 2.0, 1.0)
+
+
+def test_dropping_the_addend_does_not_write_zero_into_its_slot():
+    """Zero in the immediate slot is not inert, and getting this wrong is silent.
+
+    Writing 0x00 there produced a kernel that computed a*m + a: a synthesised `a*7` grew
+    eight-fold per step. Zero selects a register operand rather than meaning "no operand", so the
+    slot keeps an ordinary encoded constant and the control bit is what disables it.
+    """
+    built = agx_isa.encode_fma(0, 7.0, addend=None)
+    assert built[agx_isa.FMA_ADDEND_BYTE] != 0x00
+    assert not agx_isa.read_flag(built, 0, agx_isa.ADDEND_ENABLE)
+    assert agx_isa.read_flag(agx_isa.encode_fma(0, 7.0, 1.0), 0, agx_isa.ADDEND_ENABLE)
+
+
+def test_a_negative_addend_uses_the_negate_flag():
+    """The immediate field is unsigned, so the sign has to live in the control byte."""
+    built = agx_isa.encode_fma(0, 1.5, -2.0)
+    assert agx_isa.read_flag(built, 0, agx_isa.ADDEND_NEGATE)
+    assert agx_isa.decode_immediate(built[agx_isa.FMA_ADDEND_BYTE]) == 2.0
+
+
+@pytest.mark.parametrize(
+    ("fields", "step"),
+    (
+        ({"multiplier": 3.0, "addend": 0.5}, lambda v: v * 3.0 + 0.5),
+        ({"multiplier": 1.5, "addend": -2.0}, lambda v: v * 1.5 - 2.0),
+        ({"multiplier": 7.0, "addend": None}, lambda v: v * 7.0),
+        ({"multiplier": 2.0, "addend": 1.0, "negate_product": True}, lambda v: -v * 2.0 + 1.0),
+    ),
+)
+def test_synthesised_instructions_run_as_predicted(fields, step):
+    """Assemble instructions from scratch, overwrite real ones, and predict the GPU.
+
+    This is the strongest claim in the file. The bytes are not a compiler's output with a field
+    edited; every one of the eight is chosen from the measured field map, and the arithmetic they
+    produce was worked out before they were assembled.
+    """
+    text = _machine_code()
+    offsets = _compact_offsets(text)
+    assert len(offsets) == 3
+
+    def rewrite(original):
+        patched = bytearray(original)
+        for index, offset in enumerate(offsets):
+            patched[offset : offset + agx_isa.FMA_LENGTH] = agx_isa.encode_fma(
+                original[offset] >> 4, last=(index == len(offsets) - 1), **fields
+            )
+        return bytes(patched)
+
+    inputs = [1.0, 2.0, 3.0, 5.0]
+    predicted = []
+    for value in inputs:
+        running = value * 2.0 + 1.0  # the long-form fma, left alone
+        for _ in offsets:
+            running = step(running)
+        predicted.append(running)
+
+    assert agx_isa.execute(CHAIN, "probe", inputs, rewrite=rewrite) == predicted