Correct the addend flag: it picks immediate or register, not present or absent

benchmarks/agx_isa_probe.py

@@ -183,7 +183,7 @@ def rewrite(original, m=multiplier, a=addend, where=tuple(chosen)):
     for flag, clear, step, label in (
         (agx_isa.PRODUCT_NEGATE, False, lambda v: -v * 2.0 + 1.0, "product negated: -a*m+d"),
         (agx_isa.ADDEND_NEGATE, False, lambda v: v * 2.0 - 1.0, "addend negated: a*m-d"),
-        (agx_isa.ADDEND_ENABLE, True, lambda v: v * 2.0, "addend dropped: a*m"),
+        (agx_isa.ADDEND_IMMEDIATE, True, lambda v: v * 2.0, "addend becomes a zero register: a*m"),
     ):
 
         def rewrite(original, f=flag, c=clear, where=tuple(compact)):

metile/target/agx_isa.py

@@ -138,7 +138,15 @@ def __repr__(self):
 # fma(a, m, d) with every flag clear computes a * m + d.
 PRODUCT_NEGATE = FmaFlag(2, 0x10, "negate the product", "-a*m + d")
 ADDEND_NEGATE = FmaFlag(4, 0x10, "negate the addend", "a*m - d")
-ADDEND_ENABLE = FmaFlag(4, 0x20, "include the addend at all", "a*m + d, clear gives a*m")
+# Named for what it selects, after an earlier name got it wrong. It was called ADDEND_ENABLE and
+# described as including the addend or not, because clearing it turned a*m+d into a*m. That reading
+# survived a prediction on four inputs and was still incomplete: clearing it switches the addend
+# slot from an immediate to a *register*, and the byte left in the slot happened to name register
+# 88, outside the sixteen the field can reach, which reads zero. The addend was never absent, it
+# was zero.
+ADDEND_IMMEDIATE = FmaFlag(
+    4, 0x20, "the addend slot holds an immediate", "immediate; clear means it names a register"
+)
 INSTRUCTION_DISABLE = FmaFlag(6, 0x20, "retire the instruction", "no effect, like a nop")
 # Not exported: setting 0x01 in byte 2 made `a*2+1` compute `a*a`, giving 225 from an accumulator
 # holding 15, so some operand slot is being redirected to the accumulator. Which one was never
@@ -154,6 +162,16 @@ def __repr__(self):
 FMA_REGISTER_BYTE = 1
 FMA_MAX_REGISTER = 15
 
+# The addend slot uses the same shape when it names a register: `r << 1`, low bit ignored. Verified
+# by rewriting instructions to `rd = rd * m + rs` for every ordered pair of three live registers at
+# two multipliers, predicting all four threads each: eighteen rewrites, seventy-two exact values.
+# This is the register-plus-register add, reached as `rd * 1 + rs`.
+#
+# Indices at or above sixteen read zero, which is what makes a zero addend expressible at all,
+# since the immediate field's smallest value is 2**-11 and it has no encoding for zero.
+ARCHITECTURAL_REGISTERS = 16
+_ZERO_REGISTER_SLOT = 0x58  # names register 44, above the sixteen reachable, so it reads zero
+
 # 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
@@ -164,38 +182,45 @@ def __repr__(self):
 _FMA_MODE_BASE = 0x0E
 
 
-def encode_fma(register, multiplier, addend=None, last=False, negate_product=False):
+def encode_fma(
+    register, multiplier, addend=None, last=False, negate_product=False, addend_register=None
+):
     """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.
+    The addend has three forms. A float is an immediate; a negative one is encoded by setting the
+    negate flag, which is how the slot reaches values its unsigned field cannot. `addend_register`
+    names a register instead, which is how `rd * 1 + rs` becomes a register-plus-register add.
+    Passing neither leaves a zero addend, using a register index above the sixteen the field can
+    reach because those read zero and the immediate field has no encoding for zero.
     """
+    if addend is not None and addend_register is not None:
+        raise EncodingError("an addend is either an immediate or a register, not both")
     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)
+    # With the immediate bit clear the slot names a register, so what goes in it is a register
+    # index and not a leftover. Writing 0x00 there was tried first and the kernel computed a*m + a,
+    # eight-fold growth per step where seven was predicted, because index zero was the accumulator.
+    control = _FMA_ADDEND_CONTROL & ~ADDEND_IMMEDIATE.mask
+    addend_byte = _ZERO_REGISTER_SLOT
     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)
+    elif addend_register is not None:
+        if not 0 <= addend_register < ARCHITECTURAL_REGISTERS:
+            raise EncodingError(f"addend register {addend_register} is outside the field")
+        addend_byte = addend_register << 1
     return bytes(
         (
             (register << 4) | FMA_OPCODE_NIBBLE,

tests/test_agx_isa.py

@@ -163,7 +163,7 @@ def _compact_offsets(text):
     (
         (agx_isa.PRODUCT_NEGATE, False, lambda v: -v * 2.0 + 1.0),
         (agx_isa.ADDEND_NEGATE, False, lambda v: v * 2.0 - 1.0),
-        (agx_isa.ADDEND_ENABLE, True, lambda v: v * 2.0),
+        (agx_isa.ADDEND_IMMEDIATE, True, lambda v: v * 2.0),
     ),
 )
 def test_each_arithmetic_flag_does_what_it_claims(flag, clear, step):
@@ -197,7 +197,7 @@ def rewrite(original):
 def test_flags_read_back_the_way_they_were_written():
     instruction = bytes.fromhex("0901 2ec1 21b0 0202")
     assert not agx_isa.read_flag(instruction, 0, agx_isa.PRODUCT_NEGATE)
-    assert agx_isa.read_flag(instruction, 0, agx_isa.ADDEND_ENABLE)
+    assert agx_isa.read_flag(instruction, 0, agx_isa.ADDEND_IMMEDIATE)
     negated = agx_isa.write_flag(instruction, 0, agx_isa.PRODUCT_NEGATE, True)
     assert agx_isa.read_flag(negated, 0, agx_isa.PRODUCT_NEGATE)
     assert len(negated) == len(instruction)
@@ -263,8 +263,8 @@ def test_dropping_the_addend_does_not_write_zero_into_its_slot():
     """
     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)
+    assert not agx_isa.read_flag(built, 0, agx_isa.ADDEND_IMMEDIATE)
+    assert agx_isa.read_flag(agx_isa.encode_fma(0, 7.0, 1.0), 0, agx_isa.ADDEND_IMMEDIATE)
 
 
 def test_a_negative_addend_uses_the_negate_flag():
@@ -311,3 +311,106 @@ def rewrite(original):
         predicted.append(running)
 
     assert agx_isa.execute(CHAIN, "probe", inputs, rewrite=rewrite) == predicted
+
+
+def test_the_addend_slot_can_name_a_register():
+    """`rd * 1 + rs` is the register-plus-register add, which the immediate form cannot express."""
+    built = agx_isa.encode_fma(2, 1.0, addend_register=3)
+    assert not agx_isa.read_flag(built, 0, agx_isa.ADDEND_IMMEDIATE)
+    assert built[agx_isa.FMA_ADDEND_BYTE] >> 1 == 3
+    with pytest.raises(agx_isa.EncodingError, match="not both"):
+        agx_isa.encode_fma(2, 1.0, addend=1.0, addend_register=3)
+    with pytest.raises(agx_isa.EncodingError, match="addend register"):
+        agx_isa.encode_fma(2, 1.0, addend_register=agx_isa.ARCHITECTURAL_REGISTERS)
+
+
+def test_a_zero_addend_names_a_register_beyond_the_reachable_range():
+    """There is no zero immediate, so a zero addend has to come from a register that reads zero.
+
+    Asserted rather than left implicit because it is the one place the encoder relies on an index
+    the field cannot reach. If a future part reduced that range, this would start adding whatever
+    a now-reachable register holds, and silently.
+    """
+    built = agx_isa.encode_fma(2, 1.0)
+    assert not agx_isa.read_flag(built, 0, agx_isa.ADDEND_IMMEDIATE)
+    assert built[agx_isa.FMA_ADDEND_BYTE] >> 1 >= agx_isa.ARCHITECTURAL_REGISTERS
+
+
+THREE_CHAINS = (
+    """#include <metal_stdlib>
+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 ((2.0, "a"), (3.0, "b"), (5.0, "c"))
+        for _ in range(3)
+    )
+    + """
+    out[gid] = a + b + c;
+}
+"""
+)
+
+
+def test_a_register_addend_adds_that_register_on_the_gpu():
+    """Rewrite one instruction to read another chain's register, and predict all four threads.
+
+    Four threads rather than one, because thread `gid` reads x[gid], x[gid+1] and x[gid+2]: an
+    earlier version of this prediction used thread 0's inputs for every thread and mispredicted
+    three of four while the encoding was perfectly correct.
+    """
+    from metile.target import Unavailable, machine_code
+
+    multipliers = {"a": 2.0, "b": 3.0, "c": 5.0}
+    try:
+        text = machine_code(THREE_CHAINS, "probe")
+    except Unavailable as error:
+        pytest.skip(f"no Metal toolchain: {error}")
+
+    byname = {agx_isa.encode_immediate(m, low_bit=1): name for name, m in multipliers.items()}
+    compact = [
+        (offset, byname[text[offset + agx_isa.FMA_MULTIPLIER_BYTE]])
+        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 byname
+    ]
+    registers = {chain: text[offset] >> 4 for offset, chain in compact}
+    assert len(registers) == 3
+
+    inputs = [1.0, 2.0, 3.0, 5.0, 7.0, 11.0]
+    threads = 4
+    victim = [offset for offset, chain in compact if chain == "a"][-1]
+    source = registers["b"]
+    by_register = {value: name for name, value in registers.items()}
+
+    def rewrite(original):
+        patched = bytearray(original)
+        patched[victim : victim + agx_isa.FMA_LENGTH] = agx_isa.encode_fma(
+            registers["a"], 1.0, addend_register=source, last=False
+        )
+        return bytes(patched)
+
+    predicted = []
+    for gid in range(threads):
+        state = {"a": inputs[gid], "b": inputs[gid + 1], "c": inputs[gid + 2]}
+        for chain in state:
+            hidden = 3 - sum(1 for _, owner in compact if owner == chain)
+            for _ in range(hidden):
+                state[chain] = state[chain] * multipliers[chain] + 1.0
+        for offset, chain in compact:
+            if offset == victim:
+                state[chain] = state[chain] * 1.0 + state[by_register[source]]
+            else:
+                state[chain] = state[chain] * multipliers[chain] + 1.0
+        predicted.append(sum(state.values()))
+
+    got = agx_isa.execute(THREE_CHAINS, "probe", inputs, rewrite=rewrite)[:threads]
+    assert got == predicted