diff --git a/benchmarks/agx_isa_probe.py b/benchmarks/agx_isa_probe.py index 15f9ae4..cd7702c 100644 --- a/benchmarks/agx_isa_probe.py +++ b/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)): diff --git a/metile/target/agx_isa.py b/metile/target/agx_isa.py index b651a02..e481111 100644 --- a/metile/target/agx_isa.py +++ b/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,7 +182,9 @@ 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 @@ -172,10 +192,14 @@ def encode_fma(register, multiplier, addend=None, last=False, negate_product=Fal 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 @@ -183,19 +207,20 @@ def encode_fma(register, multiplier, addend=None, last=False, negate_product=Fal 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, diff --git a/tests/test_agx_isa.py b/tests/test_agx_isa.py index 7606d81..c8da19c 100644 --- a/tests/test_agx_isa.py +++ b/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 +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