Merge pull request #2923 from MarcZX100/fix/hw-oracle-wave64-batch

fix(hw-oracle): batch wave64 probes into 512-row dispatches
This commit is contained in:
Dean 2026-10-11 09:16:41 +13:00 committed by GitHub
commit 56fb1b3b95
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 22 additions and 2 deletions

View File

@ -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()

View File

@ -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"