148 lines
7.0 KiB
Python
148 lines
7.0 KiB
Python
import argparse
|
|
import hashlib
|
|
import os
|
|
import shlex
|
|
import shutil
|
|
import struct
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from functools import cache
|
|
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():
|
|
for root in (os.environ.get("ROCM_PATH"), "/opt/rocm"):
|
|
if root and (Path(root) / "include" / "hsa" / "hsa.h").exists():
|
|
return Path(root)
|
|
return None
|
|
|
|
|
|
def tool(name):
|
|
found = shutil.which(name)
|
|
if found:
|
|
return found
|
|
root = rocm_root()
|
|
if root and (root / "llvm" / "bin" / name).exists():
|
|
return str(root / "llvm" / "bin" / name)
|
|
sys.exit(f"{name} not found: install LLVM with the AMDGPU target or set ROCM_PATH")
|
|
|
|
|
|
def oracle():
|
|
source = HERE / "oracle.c"
|
|
compiler = shlex.split(os.environ.get("CC", "cc"))
|
|
if not compiler:
|
|
raise ValueError("CC must name a compiler command")
|
|
root = rocm_root()
|
|
if root is not None:
|
|
root = root.resolve()
|
|
configuration = repr((compiler, str(root))).encode("utf-8")
|
|
key = hashlib.sha256(source.read_bytes() + configuration).hexdigest()
|
|
directory = CACHE / key
|
|
binary = directory / "oracle"
|
|
if binary.is_file():
|
|
return binary
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
with tempfile.TemporaryDirectory(dir=directory) as tmp:
|
|
output = Path(tmp) / "oracle"
|
|
command = [*compiler, "-O1", str(source), "-o", str(output)]
|
|
if root:
|
|
lib = root / "lib"
|
|
command += [f"-I{root / 'include'}", f"-L{lib}", f"-Wl,-rpath,{lib}"]
|
|
subprocess.run(command + ["-lhsa-runtime64"], check=True)
|
|
output.replace(binary)
|
|
return binary
|
|
|
|
|
|
|
|
|
|
@cache
|
|
def detected_target(binary):
|
|
return subprocess.run([binary, "--target"], capture_output=True, text=True, check=True).stdout.strip()
|
|
|
|
|
|
def target():
|
|
return os.environ.get("HW_ORACLE_TARGET") or detected_target(oracle())
|
|
|
|
|
|
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))), ("@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)
|
|
subprocess.run([tool("ld.lld"), "-m", "elf64_amdgpu", "-shared", str(work / "k.o"), "-o", str(work / "k.co")], check=True)
|
|
return work / "k.co"
|
|
|
|
|
|
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)
|
|
env = dict(os.environ)
|
|
env.pop("COARSE", None)
|
|
if coarse:
|
|
env["COARSE"] = "1"
|
|
out = []
|
|
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, lds=lds)
|
|
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()
|
|
out += [struct.unpack_from("<16I", data, lane * 64) for lane in range(len(chunk))]
|
|
return out[:len(rows)]
|
|
|
|
|
|
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")
|
|
parser.add_argument("rows", type=Path, help="one row per line: 4 u32 values (decimal or 0x hex)")
|
|
parser.add_argument("--wave64", action="store_true")
|
|
parser.add_argument("--ieee", type=int, choices=(0, 1), required=True)
|
|
parser.add_argument("--denorm32", type=int, choices=range(4), required=True)
|
|
parser.add_argument("--denorm16", type=int, choices=range(4), required=True)
|
|
parser.add_argument("--dx10-clamp", type=int, choices=(0, 1), required=True)
|
|
parser.add_argument("--round32", type=int, choices=range(4), required=True)
|
|
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()
|
|
rows = [tuple(int(v, 0) 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""
|
|
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, lds=args.lds):
|
|
print(" ".join(f"{v:08x}" for v in result[:args.outs]))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|