From 8d30c27589abade41b33d6e852be9ea09fa82e06 Mon Sep 17 00:00:00 2001 From: ale Date: Fri, 9 Oct 2026 02:24:02 -0300 Subject: [PATCH] feat(hw-oracle): make the group segment size an option The template allocated a fixed 4 KiB of LDS, so kernels with immediate offsets past it could not run. --lds (Python lds) sets the group segment size in bytes, up to 64 KiB, with 4096 as the default. Co-Authored-By: Claude Fable 5.1 --- docs/dev/HW_ORACLE.md | 7 ++++-- tools/hw-oracle/hw_oracle.py | 15 ++++++++---- tools/hw-oracle/template.s | 4 ++-- tools/hw-oracle/test_hw_oracle.py | 39 +++++++++++++++++++++++++++++-- 4 files changed, 54 insertions(+), 11 deletions(-) diff --git a/docs/dev/HW_ORACLE.md b/docs/dev/HW_ORACLE.md index 464033fa2..9e6475c5c 100644 --- a/docs/dev/HW_ORACLE.md +++ b/docs/dev/HW_ORACLE.md @@ -26,9 +26,12 @@ from the GPU (`HW_ORACLE_TARGET` overrides it). - `v4`-`v7` hold the row's 4 input dwords. `v10`-`v25` start at 0 and are stored as the result (16 dwords). - Leave `v0` (lane id), `v1` (input offset), `v2` (output offset) and `s[4:7]` (input and output addresses) unchanged. -- 4 KiB of LDS. Rows run in workgroups of up to 1024 lanes, padded to whole waves with zero rows. +- 4 KiB of LDS, or the group segment size `--lds` (Python `lds`) gives in bytes, up to 64 KiB. Rows run in workgroups of up to + 1024 lanes, padded to whole waves with zero rows. - Bytes passed as `extra` follow the rows of each dispatch, at `s[4:5] + 16 * rows`: buffer contents, texels, etc. - Build buffer and image descriptors in SGPRs from that address. + Build buffer and image descriptors in SGPRs from that address. A kernel that needs more than 4 input dwords per lane + takes the rest from there too, for example `v_lshlrev_b32 v40, 3, v0` / `v_add_nc_u32 v40, 16 * rows, v40` / + `global_load_dwordx2 v[8:9], v40, s[4:5]` for two more dwords per lane, with `rows` the padded count. Every run must explicitly set all floating-point controls, including runs of integer instructions. The CLI flags use hyphens; Python keywords use underscores. Record these settings alongside the GPU and measured results. diff --git a/tools/hw-oracle/hw_oracle.py b/tools/hw-oracle/hw_oracle.py index b3d6e4ed7..6cb1c4738 100644 --- a/tools/hw-oracle/hw_oracle.py +++ b/tools/hw-oracle/hw_oracle.py @@ -11,6 +11,8 @@ from pathlib import Path HERE = Path(__file__).resolve().parent CACHE = Path(os.environ.get("HW_ORACLE_CACHE") or Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache") / "anyps5-hw-oracle") ROWS_PER_DISPATCH = 1024 +DEFAULT_LDS = 4096 +MAX_LDS = 65536 def rocm_root(): @@ -53,18 +55,20 @@ def target(): return os.environ.get("HW_ORACLE_TARGET") or subprocess.run([oracle(), "--target"], capture_output=True, text=True, check=True).stdout.strip() -def assemble(body, work, wave64, *, ieee, denorm32, denorm16, dx10_clamp, round32, round16, fp16_overflow): +def assemble(body, work, wave64, *, ieee, denorm32, denorm16, dx10_clamp, round32, round16, fp16_overflow, lds=DEFAULT_LDS): modes = (("ieee", ieee, 2), ("denorm32", denorm32, 4), ("denorm16", denorm16, 4), ("dx10_clamp", dx10_clamp, 2), ("round32", round32, 4), ("round16", round16, 4), ("fp16_overflow", fp16_overflow, 2)) for name, value, limit in modes: if not isinstance(value, int) or value not in range(limit): raise ValueError(f"{name} must be an integer from 0 to {limit - 1}") + if not isinstance(lds, int) or lds not in range(MAX_LDS + 1): + raise ValueError(f"lds must be an integer from 0 to {MAX_LDS}") text = (HERE / "template.s").read_text() for key, value in (("@TARGET@", target()), ("@WAVE32@", "0" if wave64 else "1"), ("@WAVESIZE@", "64" if wave64 else "32"), ("@DENORM@", str(int(denorm32))), ("@DENORM16@", str(int(denorm16))), ("@IEEE@", str(int(ieee))), ("@DX10_CLAMP@", str(int(dx10_clamp))), ("@ROUND32@", str(int(round32))), ("@ROUND16@", str(int(round16))), - ("@FP16_OVERFLOW@", str(int(fp16_overflow))), ("@BODY@", body)): + ("@FP16_OVERFLOW@", str(int(fp16_overflow))), ("@LDS@", str(lds)), ("@BODY@", body)): text = text.replace(key, value) (work / "k.s").write_text(text) subprocess.run([tool("clang"), "-x", "assembler", "-target", "amdgcn-amd-amdhsa", f"-mcpu={target()}", *(["-mwavefrontsize64"] if wave64 else []), "-c", str(work / "k.s"), "-o", str(work / "k.o")], check=True) @@ -72,7 +76,7 @@ def assemble(body, work, wave64, *, ieee, denorm32, denorm16, dx10_clamp, round3 return work / "k.co" -def run(body, rows, extra=b"", wave64=False, coarse=False, *, ieee, denorm32, denorm16, dx10_clamp, round32, round16, fp16_overflow): +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] wave = 64 if wave64 else 32 padded = rows + [(0, 0, 0, 0)] * (-len(rows) % wave) @@ -84,7 +88,7 @@ def run(body, rows, extra=b"", wave64=False, coarse=False, *, ieee, denorm32, de with tempfile.TemporaryDirectory() as tmp: 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) + 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] (work / "in.bin").write_bytes(b"".join(struct.pack("<4I", *r) for r in chunk) + extra) @@ -107,6 +111,7 @@ def main(): parser.add_argument("--round16", type=int, choices=range(4), required=True) parser.add_argument("--fp16-overflow", type=int, choices=(0, 1), required=True) parser.add_argument("--coarse", action="store_true", help="input and output in coarse-grained GPU memory") + parser.add_argument("--lds", type=int, default=DEFAULT_LDS, metavar="BYTES", help=f"group segment size, 0 to {MAX_LDS} (default {DEFAULT_LDS})") 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() @@ -117,7 +122,7 @@ def main(): for result in run(args.body.read_text(), rows, extra, args.wave64, args.coarse, ieee=args.ieee, denorm32=args.denorm32, denorm16=args.denorm16, dx10_clamp=args.dx10_clamp, round32=args.round32, round16=args.round16, - fp16_overflow=args.fp16_overflow): + fp16_overflow=args.fp16_overflow, lds=args.lds): print(" ".join(f"{v:08x}" for v in result[:args.outs])) diff --git a/tools/hw-oracle/template.s b/tools/hw-oracle/template.s index 915409b61..f14a96c03 100644 --- a/tools/hw-oracle/template.s +++ b/tools/hw-oracle/template.s @@ -40,7 +40,7 @@ probe: .amdhsa_kernel probe .amdhsa_user_sgpr_kernarg_segment_ptr 1 .amdhsa_kernarg_size 16 - .amdhsa_group_segment_fixed_size 4096 + .amdhsa_group_segment_fixed_size @LDS@ .amdhsa_next_free_vgpr 128 .amdhsa_next_free_sgpr 48 .amdhsa_wavefront_size32 @WAVE32@ @@ -60,7 +60,7 @@ amdhsa.kernels: - .name: probe .symbol: probe.kd .kernarg_segment_size: 16 - .group_segment_fixed_size: 4096 + .group_segment_fixed_size: @LDS@ .private_segment_fixed_size: 0 .kernarg_segment_align: 8 .wavefront_size: @WAVESIZE@ diff --git a/tools/hw-oracle/test_hw_oracle.py b/tools/hw-oracle/test_hw_oracle.py index 15a33fd05..0b8ad0ea9 100644 --- a/tools/hw-oracle/test_hw_oracle.py +++ b/tools/hw-oracle/test_hw_oracle.py @@ -112,14 +112,49 @@ class FloatModeTests(unittest.TestCase): 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() - run.assert_called_once_with("s_nop 0", [(1, 2, 3, 4)], b"", True, True, **modes) + run.assert_called_once_with("s_nop 0", [(1, 2, 3, 4)], b"", True, True, **modes, lds=4096) def test_run_forwards_modes_to_assembler(self): with patch.object(hw_oracle, "assemble") as assemble: hw_oracle.run("s_nop 0", [], wave64=True, **MODES) self.assertEqual(assemble.call_args.args[0], "s_nop 0") self.assertTrue(assemble.call_args.args[2]) - self.assertEqual(assemble.call_args.kwargs, MODES) + self.assertEqual(assemble.call_args.kwargs, MODES | {"lds": 4096}) + + +class GroupSegmentTests(unittest.TestCase): + def test_cli_forwards_lds(self): + with tempfile.TemporaryDirectory() as tmp: + body = Path(tmp) / "body.s" + rows = Path(tmp) / "rows.txt" + body.write_text("s_nop 0") + rows.write_text("1 2 3 4\n") + argv = ["hw_oracle.py", str(body), str(rows), "--lds", "8192"] + 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.kwargs["lds"], 8192) + + def test_run_forwards_lds_to_assembler(self): + with patch.object(hw_oracle, "assemble") as assemble: + hw_oracle.run("s_nop 0", [], lds=65536, **MODES) + self.assertEqual(assemble.call_args.kwargs["lds"], 65536) + + def test_invalid_lds_fails_before_tool_lookup(self): + with tempfile.TemporaryDirectory() as tmp, patch.object(hw_oracle, "target") as target: + for value in (-1, 65537, 0.5, "4096"): + with self.subTest(value=value), self.assertRaisesRegex(ValueError, "lds"): + hw_oracle.assemble("s_nop 0", Path(tmp), False, lds=value, **MODES) + target.assert_not_called() + + def test_template_takes_lds(self): + with tempfile.TemporaryDirectory() as tmp, patch.object(hw_oracle, "target", return_value="gfx1036"), \ + patch.object(hw_oracle.subprocess, "run"), patch.object(hw_oracle, "tool", return_value="clang"): + hw_oracle.assemble("s_nop 0", Path(tmp), False, lds=8192, **MODES) + text = (Path(tmp) / "k.s").read_text() + self.assertIn(".amdhsa_group_segment_fixed_size 8192", text) + self.assertIn(".group_segment_fixed_size: 8192", text) + self.assertNotIn("@LDS@", text) class AssemblyTests(unittest.TestCase):