diff --git a/tools/hw-oracle/hw_oracle.py b/tools/hw-oracle/hw_oracle.py index 6cb1c4738..fbecb317e 100644 --- a/tools/hw-oracle/hw_oracle.py +++ b/tools/hw-oracle/hw_oracle.py @@ -89,8 +89,9 @@ def run(body, rows, extra=b"", wave64=False, coarse=False, *, ieee, denorm32, de work = Path(tmp) code = assemble(body, work, wave64, ieee=ieee, denorm32=denorm32, denorm16=denorm16, dx10_clamp=dx10_clamp, round32=round32, round16=round16, fp16_overflow=fp16_overflow, lds=lds) - for start in range(0, len(padded), ROWS_PER_DISPATCH): - chunk = padded[start:start + ROWS_PER_DISPATCH] + rows_per_dispatch = ROWS_PER_DISPATCH // 2 if wave64 else ROWS_PER_DISPATCH + for start in range(0, len(padded), rows_per_dispatch): + chunk = padded[start:start + rows_per_dispatch] (work / "in.bin").write_bytes(b"".join(struct.pack("<4I", *r) for r in chunk) + extra) subprocess.run([oracle(), str(code), str(len(chunk)), str(work / "in.bin"), str(work / "out.bin"), "16"], check=True, env=env) data = (work / "out.bin").read_bytes() diff --git a/tools/hw-oracle/test_hw_oracle.py b/tools/hw-oracle/test_hw_oracle.py index 0b8ad0ea9..606517706 100644 --- a/tools/hw-oracle/test_hw_oracle.py +++ b/tools/hw-oracle/test_hw_oracle.py @@ -123,6 +123,25 @@ class FloatModeTests(unittest.TestCase): class GroupSegmentTests(unittest.TestCase): + def test_wave64_dispatches_at_most_512_rows(self): + rows = [(index, 0, 0, 0) for index in range(1025)] + for wave64, expected in ((False, [1024, 32]), (True, [512, 512, 64])): + with self.subTest(wave64=wave64), tempfile.TemporaryDirectory() as tmp: + dispatch_sizes = [] + + def dispatch(command, *, check, env): + self.assertTrue(check) + dispatch_sizes.append(int(command[2])) + Path(command[4]).write_bytes(bytes(int(command[2]) * 64)) + + with patch.object(hw_oracle, "assemble", return_value=Path(tmp) / "kernel.co"), \ + patch.object(hw_oracle, "oracle", return_value="oracle"), \ + patch.object(hw_oracle.subprocess, "run", side_effect=dispatch): + result = hw_oracle.run("s_nop 0", rows, wave64=wave64, **MODES) + + self.assertEqual(dispatch_sizes, expected) + self.assertEqual(len(result), len(rows)) + def test_cli_forwards_lds(self): with tempfile.TemporaryDirectory() as tmp: body = Path(tmp) / "body.s"