diff --git a/tools/hw-oracle/hw_oracle.py b/tools/hw-oracle/hw_oracle.py index 6573c9f54..df15c544a 100644 --- a/tools/hw-oracle/hw_oracle.py +++ b/tools/hw-oracle/hw_oracle.py @@ -1,6 +1,7 @@ import argparse import hashlib import os +import operator import shlex import shutil import struct @@ -93,7 +94,16 @@ def assemble(body, work, wave64, *, ieee, denorm32, denorm16, dx10_clamp, round3 def run(body, rows, extra=b"", wave64=False, coarse=False, *, ieee, denorm32, denorm16, dx10_clamp, round32, round16, fp16_overflow, lds=DEFAULT_LDS): - rows = [tuple(r) for r in rows] + validated = [] + for index, row in enumerate(rows, 1): + try: + values = tuple(operator.index(value) for value in row) + except TypeError as error: + raise ValueError(f"row {index} needs 4 unsigned 32-bit integers") from error + if len(values) != 4 or any(value < 0 or value > 0xffffffff for value in values): + raise ValueError(f"row {index} needs 4 unsigned 32-bit integers") + validated.append(values) + rows = validated wave = 64 if wave64 else 32 padded = rows + [(0, 0, 0, 0)] * (-len(rows) % wave) env = dict(os.environ) @@ -115,6 +125,13 @@ def run(body, rows, extra=b"", wave64=False, coarse=False, *, ieee, denorm32, de return out[:len(rows)] +def parse_word(value): + try: + return int(value, 0) + except ValueError: + return int(value, 10) + + def main(): parser = argparse.ArgumentParser(description="Run an RDNA kernel body on the local AMD GPU, one input row per lane.") parser.add_argument("body", type=Path, help="assembly inserted into template.s") @@ -132,7 +149,7 @@ def main(): parser.add_argument("--extra", type=Path, help="bytes appended after the rows") parser.add_argument("--outs", type=int, choices=range(1, 17), default=16, metavar="N", help="print v10..v(10+N-1)") args = parser.parse_args() - rows = [tuple(int(v, 0) for v in line.split()) for line in args.rows.read_text().splitlines() if line.strip()] + rows = [tuple(parse_word(v) for v in line.split()) for line in args.rows.read_text().splitlines() if line.strip()] if any(len(r) != 4 for r in rows): sys.exit("every row needs 4 values") extra = args.extra.read_bytes() if args.extra else b"" diff --git a/tools/hw-oracle/test_hw_oracle.py b/tools/hw-oracle/test_hw_oracle.py index 26c05c624..bdf6940a7 100644 --- a/tools/hw-oracle/test_hw_oracle.py +++ b/tools/hw-oracle/test_hw_oracle.py @@ -76,6 +76,45 @@ class OracleCacheTests(unittest.TestCase): self.assertEqual([file for file in cache.rglob("*") if file.is_file()], [binary]) +class InputRowsTests(unittest.TestCase): + def test_invalid_rows_fail_before_assembler_or_gpu_lookup(self): + for row in ((1, 2, 3), (1, 2, 3, 4, 5), (-1, 0, 0, 0), + (0x100000000, 0, 0, 0), (1.5, 0, 0, 0), ("1", 0, 0, 0)): + with self.subTest(row=row), patch.object(hw_oracle, "assemble", side_effect=AssertionError("assembler called")) as assemble: + with self.assertRaisesRegex(ValueError, "row 1"): + hw_oracle.run("s_nop 0", [row], **MODES) + assemble.assert_not_called() + + def test_unsigned_bits_and_indexable_values_are_preserved(self): + class Indexable: + def __index__(self): + return 0xffffffff + + def dispatch(command, **kwargs): + data = Path(command[3]).read_bytes() + self.assertEqual(data[:16], hw_oracle.struct.pack("<4I", 0, 0xffffffff, 0x80000000, 1)) + self.assertEqual(data[16:32 * 16], bytes(31 * 16)) + self.assertEqual(data[32 * 16:], b"extra") + Path(command[4]).write_bytes(bytes(32 * 64)) + + with patch.object(hw_oracle, "assemble", return_value="kernel.co"), \ + patch.object(hw_oracle, "oracle", return_value="oracle"), \ + patch.object(hw_oracle.subprocess, "run", side_effect=dispatch): + results = hw_oracle.run("s_nop 0", [(0, Indexable(), 0x80000000, True)], b"extra", **MODES) + self.assertEqual(results, [tuple([0] * 16)]) + + def test_cli_accepts_zero_padded_decimal_rows(self): + with tempfile.TemporaryDirectory() as tmp: + body, rows = Path(tmp) / "body.s", Path(tmp) / "rows.txt" + body.write_text("s_nop 0") + rows.write_text("0001 08 0xFFFFFFFF 0b10\n") + argv = ["hw_oracle.py", str(body), str(rows)] + argv += [item for name, value in MODES.items() for item in ["--" + name.replace("_", "-"), str(value)]] + with patch("sys.argv", argv), patch.object(hw_oracle, "run", return_value=[]) as run: + hw_oracle.main() + self.assertEqual(run.call_args.args[1], [(1, 8, 0xffffffff, 2)]) + + class FloatModeTests(unittest.TestCase): def test_python_requires_each_mode(self): for name in MODES: