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 <noreply@anthropic.com>
This commit is contained in:
ale
parent
d70b899894
commit
8d30c27589
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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]))
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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@
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Reference in New Issue