diff --git a/benchmarks/agx_isa_probe.py b/benchmarks/agx_isa_probe.py new file mode 100644 index 0000000..aa9658d --- /dev/null +++ b/benchmarks/agx_isa_probe.py @@ -0,0 +1,215 @@ +"""Re-derive what meTile knows about G17 machine code, from nothing, on this machine. + +Everything in `metile.target.agx_isa` came out of this procedure, and running it is how to port +that knowledge to different hardware or a different toolchain. Nothing is read from a table; each +stage measures, and the last stage checks the result by predicting GPU output in advance and +comparing. + + 1 edit patch a byte range to nops and confirm the edit runs, which is the + premise for everything after it + 2 boundaries find where instructions actually start, behaviourally + 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 + +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 +nobody has compiled shows the encoding is understood. + +usage: + python benchmarks/agx_isa_probe.py + python benchmarks/agx_isa_probe.py --verbose # show every alignment the scan rejected +""" + +import argparse +import itertools +import sys +from pathlib import Path + +_root = str(Path(__file__).resolve().parent.parent) +sys.path.insert(0, _root) + +from metile.target import agx, agx_isa + +# 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 +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; +} +""" +INPUTS = [1.0, 2.0, 3.0, 4.0] +INTACT = 31.0 # ((((1*2+1)*2+1)*2+1)*2+1) +ONE_REMOVED = 15.0 # three fmas instead of four + + +# Rebuilt from a template rather than by substituting into CHAIN: replacing "2.0f" and then +# "1.0f" corrupts the source whenever one constant is textually the other. +def _constants(multiplier, addend): + step = f" a = fma(a, {multiplier}f, {addend}f);" + return CHAIN.replace(" a = fma(a, 2.0f, 1.0f);", step) + + +def _arguments(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--verbose", action="store_true") + parser.add_argument("--workdir", type=Path, default=Path(".metile-agx")) + return parser.parse_args() + + +def main(): + arguments = _arguments() + work = arguments.workdir + + try: + text = agx.machine_code(CHAIN, "probe", work) + except agx.Unavailable as error: + print(f"cannot read compiled kernels: {error}") + return 1 + + print(f"Kernel is {len(text)} bytes of machine code.") + nops = text.count(agx_isa.NOP) + print( + f"Contains {nops} occurrences of the two-byte nop and " + f"{text.count(agx_isa.BLOCK_END)} block terminators; " + f"code resumes at each {agx_isa.BLOCK_ALIGNMENT}-byte boundary.\n" + ) + + print("1. Does an edited archive actually run?") + baseline = agx_isa.execute(CHAIN, "probe", INPUTS, workdir=work) + print(f" unedited: x=1 -> {baseline[0]:g}, expected {INTACT:g}") + if baseline[0] != INTACT: + print(" the kernel does not compute what this probe assumes; stopping") + return 1 + + print("\n2. Where do instructions start? Nop out eight bytes at every even offset.") + print(f" {INTACT:g} means nothing was removed, {ONE_REMOVED:g} means exactly one fma was.") + region = (0x50, len(text) - 12, agx_isa.FMA_LENGTH) + offsets, _ = agx_isa.boundaries( + CHAIN, "probe", region, INPUTS, INTACT, ONE_REMOVED, workdir=work + ) + print(f" instruction starts: {[hex(offset) for offset in offsets]}") + strides = {b - a for a, b in itertools.pairwise(offsets)} + print(f" stride between them: {strides or 'n/a'}") + + # A confirmed boundary is not the same as a form whose fields are known. All four fmas sit on + # an eight-byte stride, but the compiler encodes the first differently from the rest, and + # reading the constant field out of that one yields nonsense. Being a boundary is measured; + # being the mapped form is what the opcode nibble decides. + compact = [offset for offset in offsets if text[offset] & 0x0F == agx_isa.FMA_OPCODE_NIBBLE] + for offset in offsets: + form = "compact, fields mapped" if offset in compact else "other form, not mapped" + print(f" 0x{offset:04x} {bytes(text[offset : offset + 8]).hex(' ')} {form}") + if arguments.verbose: + candidates = agx_isa.find_fma(text) + print(f" pattern match alone would have proposed: {[hex(c) for c in candidates]}") + if not compact: + print(" no instruction in the mapped form; the remaining stages depend on one") + return 1 + + print("\n3. What does the constant field look like? Compile known constants and read it.") + print(f" {'multiplier':>11}{'addend':>8}{'byte 3':>9}{'byte 5':>9} decoded") + observed = [] + for multiplier, addend in ( + (2.0, 1.0), + (3.0, 1.0), + (4.0, 1.0), + (8.0, 1.0), + (2.0, 3.0), + (2.0, 5.0), + ): + variant = agx.machine_code(_constants(multiplier, addend), "probe", work) + instruction = variant[compact[0] : compact[0] + agx_isa.FMA_LENGTH] + mul_byte = instruction[agx_isa.FMA_MULTIPLIER_BYTE] + add_byte = instruction[agx_isa.FMA_ADDEND_BYTE] + pair = (agx_isa.decode_immediate(mul_byte), agx_isa.decode_immediate(add_byte)) + agree = pair == (multiplier, addend) + observed.append(agree) + print( + f" {multiplier:>11}{addend:>8}{f'0x{mul_byte:02x}':>9}{f'0x{add_byte:02x}':>9}" + f" {pair[0]:g}, {pair[1]:g} {'ok' if agree else 'MISREAD'}" + ) + print(f" the derived format reads {sum(observed)} of {len(observed)} correctly") + + print("\n4. Can constants the compiler never emitted be synthesised? Predict, then run.") + print(f" {'rewrite':<26}{'predicted':>12}{'measured':>12} verdict") + checks = [] + for multiplier, addend, count in ( + (6.0, 7.0, 1), + (6.0, 7.0, len(compact)), + (1.25, 1.5, len(compact)), + ): + chosen = compact[:count] + + def rewrite(original, m=multiplier, a=addend, where=tuple(chosen)): + patched = original + for offset in where: + patched = agx_isa.rewrite_fma_immediates(patched, offset, m, a) + return patched + + # Predict by replaying the arithmetic, applying the new constants at exactly the steps + # that were patched. Which steps those are has to come from the offsets, not from a + # count: the compact instructions are the second, third and fourth fma here, so + # assuming the patched ones are the trailing ones predicted 97 where the GPU said 103. + rewritten = {offsets.index(offset) for offset in chosen} + value = 1.0 + for step in range(len(offsets)): + if step in rewritten: + value = value * multiplier + addend + else: + value = value * 2.0 + 1.0 + got = agx_isa.execute(CHAIN, "probe", INPUTS, rewrite=rewrite, workdir=work)[0] + agree = got == value + checks.append(agree) + label = f"{count} fma -> a*{multiplier:g}+{addend:g}" + print(f" {label:<26}{value:>12g}{got:>12g} {'MATCHES' if agree else 'differs'}") + + print("\n5. Do the arithmetic flags mean what they claim? One bit at a time, four inputs.") + print(f" {'flag':<30}{'predicted':>28} verdict") + inputs = [1.0, 2.0, 3.0, 5.0] + 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"), + ): + + def rewrite(original, f=flag, c=clear, where=tuple(compact)): + patched = original + for offset in where: + patched = agx_isa.write_flag(patched, offset, f, not c) + return 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:<30}{shown:>28} {'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.") + 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.") + return 0 if all(checks) and all(observed) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/metile/target/agx_execute.swift b/metile/target/agx_execute.swift new file mode 100644 index 0000000..0d4bcba --- /dev/null +++ b/metile/target/agx_execute.swift @@ -0,0 +1,95 @@ +// Run a kernel from a binary archive, including one whose machine code has been edited. +// +// usage: agx_execute +// +// The archive is the authority. Metal is given the same MSL it was built from, because a +// pipeline descriptor needs a function object, but `failOnBinaryArchiveMiss` makes the archive +// the only permitted source of machine code: if the driver cannot find a matching entry it +// errors instead of quietly recompiling the source. That is what makes an edited archive +// observable rather than silently ignored — without the flag a patched kernel and an unpatched +// one produce the same answer and the experiment proves nothing. +// +// Inputs and outputs are raw float32 files so the caller can drive this from anywhere. + +import Foundation +import Metal + +let arguments = CommandLine.arguments +guard arguments.count == 7 else { + FileHandle.standardError.write( + "usage: agx_execute \n" + .data(using: .utf8)!) + exit(2) +} + +guard let device = MTLCreateSystemDefaultDevice() else { + FileHandle.standardError.write("no Metal device\n".data(using: .utf8)!) + exit(1) +} + +func fail(_ message: String) -> Never { + FileHandle.standardError.write("error: \(message)\n".data(using: .utf8)!) + exit(1) +} + +do { + let source = try String(contentsOfFile: arguments[2], encoding: .utf8) + let library = try device.makeLibrary(source: source, options: MTLCompileOptions()) + guard let function = library.makeFunction(name: arguments[3]) else { + fail("no function \(arguments[3])") + } + + let descriptor = MTLBinaryArchiveDescriptor() + descriptor.url = URL(fileURLWithPath: arguments[1]) + let archive = try device.makeBinaryArchive(descriptor: descriptor) + + let pipelineDescriptor = MTLComputePipelineDescriptor() + pipelineDescriptor.computeFunction = function + pipelineDescriptor.binaryArchives = [archive] + let pipeline = try device.makeComputePipelineState( + descriptor: pipelineDescriptor, + options: .failOnBinaryArchiveMiss, + reflection: nil) + + let inputData = try Data(contentsOf: URL(fileURLWithPath: arguments[4])) + guard let threads = Int(arguments[6]), threads > 0 else { fail("threads must be positive") } + let elements = max(inputData.count / 4, threads) + + guard + let input = device.makeBuffer(length: elements * 4, options: .storageModeShared), + let output = device.makeBuffer(length: elements * 4, options: .storageModeShared), + let scalar = device.makeBuffer(length: 4, options: .storageModeShared) + else { fail("could not allocate buffers") } + + inputData.withUnsafeBytes { raw in + input.contents().copyMemory(from: raw.baseAddress!, byteCount: inputData.count) + } + memset(output.contents(), 0, elements * 4) + scalar.contents().bindMemory(to: UInt32.self, capacity: 1)[0] = UInt32(elements) + + guard + let queue = device.makeCommandQueue(), + let buffer = queue.makeCommandBuffer(), + let encoder = buffer.makeComputeCommandEncoder() + else { fail("could not create a command encoder") } + + encoder.setComputePipelineState(pipeline) + encoder.setBuffer(input, offset: 0, index: 0) + encoder.setBuffer(output, offset: 0, index: 1) + encoder.setBuffer(scalar, offset: 0, index: 2) + let width = min(pipeline.threadExecutionWidth, threads) + encoder.dispatchThreads( + MTLSize(width: threads, height: 1, depth: 1), + threadsPerThreadgroup: MTLSize(width: width, height: 1, depth: 1)) + encoder.endEncoding() + buffer.commit() + buffer.waitUntilCompleted() + + if let error = buffer.error { fail("dispatch failed: \(error)") } + + let produced = Data(bytes: output.contents(), count: elements * 4) + try produced.write(to: URL(fileURLWithPath: arguments[5])) + print("ran \(threads) threads from the archive") +} catch { + fail("\(error)") +} diff --git a/metile/target/agx_isa.py b/metile/target/agx_isa.py new file mode 100644 index 0000000..a4bb6d4 --- /dev/null +++ b/metile/target/agx_isa.py @@ -0,0 +1,326 @@ +"""What meTile has established about G17 machine code, and the method that established it. + +The instruction set is undocumented, so everything here was measured, and the method matters as +much as the results. Reading bytes and guessing produces plausible field maps that are wrong. +Every claim below was instead put to the GPU: patch the machine code inside a binary archive, +run it, and check the answer against a prediction made in advance. A hypothesis that survives +that is not an interpretation. + +Established, in descending order of confidence: + + execution A binary archive whose __text has been edited runs the edited code. The + driver does not re-derive it from the AIR it also carries, and does not + validate it. `execute` below is the primitive everything else rests on. + + boundaries Overwriting a byte range with nops and running the result finds instruction + boundaries behaviourally. On a chain of four `a = fma(a, 2, 1)` steps, only + the offsets 0x62, 0x6a and 0x72 yield 15 rather than 31, so exactly one fma + was removed at each and a compact f32 fma is eight bytes on an eight-byte + stride. Every other alignment gives 0, 1, 7 or a rejected kernel. + + immediates The eight-bit float operand field is `(e << 4) | (m << 1) | low`, holding + `(1 + m/8) * 2**(e - 11)`. Round-trips against every constant the compiler + was observed to emit, and — the part that makes it an encoder rather than a + table — correctly predicts results for constants the compiler never emitted: + 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. + + 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. + + 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 + opcode nibble is the one edit that makes the driver reject the kernel rather + than return a wrong answer, which is what desynchronising the stream would do + — but that does not generalise into a length field, see below. + +Deliberately not claimed: a general disassembler, and there is now evidence for why rather than +just caution. Instruction length looked like the low nibble of byte 0 — nop 0x06 at two bytes, +block end 0x0e at four, fma 0x09 at eight — and a table fitted to one kernel walked it exactly. +The same table then walked none of eight other kernels exactly. Lengths here come one form at a +time from the behavioural finder in `boundaries`, which cannot be fooled that way. + +Negative immediates likewise have no known encoding in this field, and `encode_immediate` refuses +them rather than returning the nearest byte, which would corrupt a kernel silently. + +`benchmarks/agx_isa_probe.py` re-derives all of it from scratch, which is how to port this to +new hardware. +""" + +import math +import struct +import subprocess +from pathlib import Path + +from metile.target import agx + +# A two-byte instruction with no effect, and a four-byte one that ends a block. Both read +# straight off the padding the compiler emits between blocks; both confirmed by patching them +# over real instructions and seeing exactly that instruction's contribution disappear. +NOP = bytes.fromhex("0600") +BLOCK_END = bytes.fromhex("0e000000") +BLOCK_ALIGNMENT = 64 + +# The compact f32 fma. Byte roles, each established by patching that byte and running: +# +# 0 opcode. Low nibble 9. Of all 256 values only the four with nibble 6 are rejected +# outright; the rest run and return a wrong answer. +# 1 register selection. Every alternative tried sent the chain's result somewhere the +# final store did not read. +# 2 operand mode and sign. Bit 0x10 negates the product. Bit 0x20 marks all but the last +# 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. +# 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. +FMA_LENGTH = 8 +FMA_OPCODE_NIBBLE = 0x09 +FMA_MULTIPLIER_BYTE = 3 +FMA_ADDEND_BYTE = 5 + +# Instruction length is NOT a function of the low nibble of byte 0, and this is the one place a +# plausible shortcut was tried and failed, so it is recorded to stop it being tried again. +# +# The nibble looks like a length field from one kernel: nop is 0x06 and two bytes, the block +# terminator 0x0e and four, the compact fma 0x09 and eight, and setting an eight-byte +# instruction's nibble to 6 is the only edit the driver rejects outright rather than running with +# a wrong answer, which is what desynchronising the stream would do. A table extended to fit one +# kernel walked it end to end and covered all four behaviourally confirmed fma boundaries. +# +# It then walked none of eight other kernels exactly: copy, two fma chains, a reduction loop, +# integer and half-precision mixes, a branch, and a sqrt. Five overran the end of the stream and +# the rest left unknown nibbles behind. Fitting sixteen free values to one 134-byte kernel simply +# is not evidence. Lengths here come from the behavioural finder in `boundaries` instead, one form +# at a time. + + +class FmaFlag: + """One bit of a compact fma whose meaning was established by patching and running. + + Each was found by scanning all 256 values of its byte, grouping the outputs by what + arithmetic they expressed, and then predicting the result of setting the bit across a chain + of instructions on four different inputs. All twelve predictions were exact. + """ + + def __init__(self, byte, mask, meaning, set_means): + self.byte = byte + self.mask = mask + self.meaning = meaning + self.set_means = set_means + + def __repr__(self): + return ( + f"FmaFlag(byte={self.byte}, mask=0x{self.mask:02x}, " + f"{self.meaning}, set gives {self.set_means})" + ) + + +# 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") +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 +# pinned down and it was never checked across several inputs, so it stays a note. Naming it would +# put it on the same footing as the flags above, which were each predicted on four inputs. + + +def read_flag(text, offset, flag): + """Whether one flag is set on the instruction at `offset`.""" + return bool(text[offset + flag.byte] & flag.mask) + + +def write_flag(text, offset, flag, value): + """Return `text` with one flag of the instruction at `offset` set or cleared.""" + patched = bytearray(text) + if value: + patched[offset + flag.byte] |= flag.mask + else: + patched[offset + flag.byte] &= ~flag.mask & 0xFF + return bytes(patched) + + +_IMMEDIATE_BIAS = 11 +_IMMEDIATE_MANTISSA_STEPS = 8 + + +class EncodingError(ValueError): + """A value has no encoding in this field, or none that has been verified.""" + + +def encode_immediate(value, low_bit=1): + """Encode a positive float into the eight-bit operand field. + + `low_bit` is the bit the field shares with its neighbour: set for the multiplier slot, + clear for the addend slot, both as the compiler emits them. It is passed rather than + inferred because what that bit belongs to has not been established, and guessing it would + silently corrupt the adjacent field. + """ + if value <= 0 or not math.isfinite(value): + raise EncodingError( + f"{value} has no verified encoding: the field carries no sign and no special values" + ) + exponent = math.floor(math.log2(value)) + mantissa = round((value / 2.0**exponent - 1.0) * _IMMEDIATE_MANTISSA_STEPS) + if mantissa == _IMMEDIATE_MANTISSA_STEPS: # rounded up to the next power of two + exponent, mantissa = exponent + 1, 0 + biased = exponent + _IMMEDIATE_BIAS + if not 0 <= biased <= 15: + raise EncodingError(f"{value} is outside the field's exponent range") + if decode_immediate((biased << 4) | (mantissa << 1) | low_bit) != value: + raise EncodingError(f"{value} is not exactly representable in this field") + return (biased << 4) | (mantissa << 1) | low_bit + + +def decode_immediate(byte): + """The float an operand byte stands for.""" + exponent = (byte >> 4) - _IMMEDIATE_BIAS + mantissa = (byte >> 1) & 0x07 + return (1.0 + mantissa / _IMMEDIATE_MANTISSA_STEPS) * 2.0**exponent + + +def find_fma(text): + """Offsets of compact f32 fma instructions, as a candidate list. + + Pattern matching on an undocumented encoding, so treat these as candidates: confirm one by + nopping it and checking the arithmetic changed the way removing that operation would. + `boundaries` does exactly that. + """ + found = [] + offset = 0 + while offset + FMA_LENGTH <= len(text): + if text[offset] & 0x0F == FMA_OPCODE_NIBBLE and text[offset + 2] & 0x0F in (0x0E,): + found.append(offset) + offset += FMA_LENGTH + else: + offset += 2 + return found + + +def rewrite_fma_immediates(text, offset, multiplier=None, addend=None): + """Return `text` with one fma's constants replaced. + + Specialises a compiled kernel's constants without recompiling it. The instruction keeps its + length, so nothing downstream shifts. + """ + patched = bytearray(text) + if multiplier is not None: + patched[offset + FMA_MULTIPLIER_BYTE] = encode_immediate(multiplier, low_bit=1) + if addend is not None: + patched[offset + FMA_ADDEND_BYTE] = encode_immediate(addend, low_bit=0) + return bytes(patched) + + +def _harness(workdir): + """Build the executor once per working directory.""" + binary = Path(workdir) / "agx_execute" + if binary.exists(): + return binary + source = Path(__file__).resolve().parent / "agx_execute.swift" + Path(workdir).mkdir(parents=True, exist_ok=True) + built = subprocess.run( + ["swiftc", "-O", str(source), "-o", str(binary)], capture_output=True, text=True + ) + if built.returncode != 0: + raise agx.Unavailable(f"could not build the executor: {built.stderr.strip()[:300]}") + return binary + + +def execute(source, function, inputs, rewrite=None, workdir=".metile-agx"): + """Compile, optionally rewrite the machine code, run, and return the outputs. + + `rewrite` takes the kernel's __text and returns replacement bytes of the same length. It is + applied in place inside the serialized archive, and the archive is then the only source the + driver is allowed to use, so what runs is what was written. + + Same length is required rather than merely advised. The bytes are patched at the offset + where they were found in the archive file, and a different length would shift everything + after them, including the metadata the driver reads to set the kernel up. + """ + workdir = Path(workdir) + prober = agx._harness(workdir) + executor = _harness(workdir) + metal = workdir / "isa.metal" + archive = workdir / "isa.bin" + metal.write_text(source) + built = subprocess.run( + [str(prober), str(metal), function, str(archive)], capture_output=True, text=True + ) + if built.returncode != 0: + raise RuntimeError(built.stderr.strip()[:300]) + + raw = archive.read_bytes() + target = archive + if rewrite is not None: + text = agx.machine_code(source, function, workdir) + replacement = rewrite(text) + if len(replacement) != len(text): + raise EncodingError( + f"a rewrite must keep the length: {len(replacement)} bytes for {len(text)}" + ) + offset = raw.find(bytes(text)) + if offset < 0 or raw.count(bytes(text)) != 1: + raise RuntimeError("could not locate __text uniquely inside the archive") + target = workdir / "isa_patched.bin" + target.write_bytes(raw[:offset] + replacement + raw[offset + len(text) :]) + + values = list(inputs) + inputs_file, outputs_file = workdir / "isa_in.f32", workdir / "isa_out.f32" + inputs_file.write_bytes(struct.pack(f"<{len(values)}f", *values)) + ran = subprocess.run( + [ + str(executor), + str(target), + str(metal), + function, + str(inputs_file), + str(outputs_file), + str(len(values)), + ], + capture_output=True, + text=True, + ) + if ran.returncode != 0: + raise RuntimeError(ran.stderr.strip().splitlines()[-1][:200] if ran.stderr else "no output") + produced = outputs_file.read_bytes()[: 4 * len(values)] + return list(struct.unpack(f"<{len(values)}f", produced)) + + +def boundaries(source, function, region, inputs, intact, removed, workdir=".metile-agx"): + """Offsets in `region` where nopping `stride` bytes removes exactly one operation. + + The behavioural instruction finder. `intact` is what the kernel returns untouched and + `removed` is what it returns with one of the operations gone, both worked out from the + kernel's arithmetic beforehand. An offset qualifies only when nopping there produces + `removed` exactly, which is a far stronger signal than the kernel merely still running: + wrong alignments do run, and return wrong answers rather than failing. + """ + text = agx.machine_code(source, function, workdir) + if execute(source, function, inputs, workdir=workdir)[0] != intact: + raise RuntimeError("the unpatched kernel does not return the expected value") + + start, stop, stride = region + found = [] + for offset in range(start, stop, 2): + + def rewrite(original, at=offset): + patched = bytearray(original) + patched[at : at + stride] = NOP * (stride // len(NOP)) + return bytes(patched) + + try: + got = execute(source, function, inputs, rewrite=rewrite, workdir=workdir) + except RuntimeError: + continue + if got and got[0] == removed: + found.append(offset) + return found, text diff --git a/tests/test_agx_isa.py b/tests/test_agx_isa.py new file mode 100644 index 0000000..9b2e0ef --- /dev/null +++ b/tests/test_agx_isa.py @@ -0,0 +1,232 @@ +"""The parts of the G17 encoding meTile claims to understand. + +Two kinds of test. The immediate format is arithmetic and is checked directly against the bytes +the Metal compiler was observed to emit, which needs no GPU. Whether an edited archive actually +runs the edited code is not arithmetic and cannot be argued, so that one compiles, patches, +dispatches, and compares against a number worked out beforehand. +""" + +import itertools + +import pytest + +from metile.target import agx_isa + +# Constants and the operand bytes the compiler emitted for them, read out of compiled kernels by +# benchmarks/agx_isa_probe.py. The multiplier slot carries a set low bit and the addend slot a +# clear one; what that bit belongs to is not established, which is why it is passed in rather +# than inferred. +OBSERVED_MULTIPLIERS = ((2.0, 0xC1), (3.0, 0xC9), (4.0, 0xD1), (8.0, 0xE1)) +OBSERVED_ADDENDS = ((1.0, 0xB0), (3.0, 0xC8), (5.0, 0xD4)) + + +@pytest.mark.parametrize(("value", "byte"), OBSERVED_MULTIPLIERS) +def test_the_multiplier_field_matches_what_the_compiler_emits(value, byte): + assert agx_isa.encode_immediate(value, low_bit=1) == byte + assert agx_isa.decode_immediate(byte) == value + + +@pytest.mark.parametrize(("value", "byte"), OBSERVED_ADDENDS) +def test_the_addend_field_matches_what_the_compiler_emits(value, byte): + assert agx_isa.encode_immediate(value, low_bit=0) == byte + assert agx_isa.decode_immediate(byte) == value + + +def test_every_representable_value_round_trips(): + """The format is (1 + m/8) * 2**(e - 11), so enumerate it and check both directions. + + Enumerating beats sampling here: the field is 128 values wide, so exhaustive is cheap, and a + sampled test would miss an off-by-one at an exponent boundary. + """ + for exponent in range(16): + for mantissa in range(8): + byte = (exponent << 4) | (mantissa << 1) + value = agx_isa.decode_immediate(byte) + assert agx_isa.encode_immediate(value, low_bit=0) == byte + + +def test_values_outside_the_field_are_refused_rather_than_approximated(): + """Silently encoding the nearest representable value would corrupt a kernel quietly. + + There is no sign bit in this field and no encoding for zero or the specials, and no exponent + reaches 2**5. A caller asking for one of those has a wrong model of the field and should be + told, not handed the closest byte. + """ + for value in (-1.0, 0.0, float("inf"), float("nan"), 1.1, 2.0**6): + with pytest.raises(agx_isa.EncodingError): + agx_isa.encode_immediate(value) + + +def test_rewriting_immediates_keeps_the_instruction_length(): + """Length has to be preserved: the patch lands at a fixed offset inside the archive.""" + instruction = bytes.fromhex("0901 2ec1 21b0 0202") + rewritten = agx_isa.rewrite_fma_immediates(instruction, 0, multiplier=6.0, addend=7.0) + assert len(rewritten) == len(instruction) + assert agx_isa.decode_immediate(rewritten[agx_isa.FMA_MULTIPLIER_BYTE]) == 6.0 + assert agx_isa.decode_immediate(rewritten[agx_isa.FMA_ADDEND_BYTE]) == 7.0 + + +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 _machine_code(): + from metile.target import Unavailable, machine_code + + try: + return machine_code(CHAIN, "probe") + except Unavailable as error: + pytest.skip(f"no Metal toolchain: {error}") + + +def test_an_edited_archive_runs_the_edited_code(): + """The premise the whole ISA effort rests on, so it is asserted rather than assumed. + + Four dependent fmas of a*2+1 turn x=1 into 31. Rewriting the constants of the compact ones + to a*6+7 must give 949, and that number was computed from the arithmetic before the bytes + were ever assembled. If the driver were recompiling from the AIR it also carries, or + ignoring the archive, the answer would still be 31 and nothing here would work. + """ + text = _machine_code() + compact = [ + 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 + ] + assert len(compact) == 3, f"expected three compact fmas, found {[hex(o) for o in compact]}" + + assert agx_isa.execute(CHAIN, "probe", [1.0, 2.0])[0] == 31.0 + + def rewrite(original): + patched = original + for offset in compact: + patched = agx_isa.rewrite_fma_immediates(patched, offset, 6.0, 7.0) + return patched + + # 1*2+1 = 3, then three steps of a*6+7: 25, 157, 949. + assert agx_isa.execute(CHAIN, "probe", [1.0, 2.0], rewrite=rewrite)[0] == 949.0 + + +def test_a_rewrite_that_changes_length_is_refused(): + """Shifting the bytes after the patch would move the metadata the driver reads.""" + _machine_code() + with pytest.raises(agx_isa.EncodingError, match="keep the length"): + agx_isa.execute(CHAIN, "probe", [1.0], rewrite=lambda text: text[:-2]) + + +def test_nopping_an_instruction_removes_exactly_its_effect(): + """The behavioural boundary finder, on the case that established the fma length. + + A wrong alignment does not fail; it runs and returns a wrong answer. That is why the test + asserts the value is exactly what dropping one fma gives, and that the offsets found sit on + an eight-byte stride, rather than merely checking the kernel survived. + """ + text = _machine_code() + offsets, _ = agx_isa.boundaries( + CHAIN, + "probe", + (0x50, len(text) - 12, agx_isa.FMA_LENGTH), + [1.0, 2.0], + intact=31.0, + removed=15.0, + ) + assert len(offsets) == 4, f"expected four fmas, found {[hex(o) for o in offsets]}" + assert {b - a for a, b in itertools.pairwise(offsets)} == {agx_isa.FMA_LENGTH} + + +def _compact_offsets(text): + return [ + 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 + ] + + +@pytest.mark.parametrize( + ("flag", "clear", "step"), + ( + (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), + ), +) +def test_each_arithmetic_flag_does_what_it_claims(flag, clear, step): + """Set the bit on every compact fma, then check the GPU against arithmetic, not a table. + + Four inputs rather than one. A single input can agree by coincidence -- negating the product + and negating the addend both happen to move the result by an even amount -- and a flag that + only holds for x=1 is not understood. + """ + text = _machine_code() + offsets = _compact_offsets(text) + assert len(offsets) == 3 + + def rewrite(original): + patched = original + for offset in offsets: + patched = agx_isa.write_flag(patched, offset, flag, not clear) + return 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 + + +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) + 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) + assert agx_isa.write_flag(negated, 0, agx_isa.PRODUCT_NEGATE, False) == instruction + + +def test_disabling_an_instruction_by_flag_matches_nopping_it(): + """The flag and the nop should be indistinguishable, and both equal dropping the operation. + + Worth asserting because it ties the two capabilities together: the behavioural boundary + finder works by overwriting bytes, and this flag achieves the same effect without touching + the instruction's length or its neighbours. + """ + text = _machine_code() + offsets = _compact_offsets(text) + last = offsets[-1] + + def by_flag(original): + return agx_isa.write_flag(original, last, agx_isa.INSTRUCTION_DISABLE, True) + + def by_nop(original): + patched = bytearray(original) + patched[last : last + agx_isa.FMA_LENGTH] = agx_isa.NOP * ( + agx_isa.FMA_LENGTH // len(agx_isa.NOP) + ) + return bytes(patched) + + inputs = [1.0, 2.0] + flagged = agx_isa.execute(CHAIN, "probe", inputs, rewrite=by_flag) + 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