mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 05:23:41 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
112 lines
6.3 KiB
Python
112 lines
6.3 KiB
Python
# minimal amdgpu elf packer
|
|
import ctypes
|
|
from tinygrad.dtype import AddrSpace
|
|
from tinygrad.helpers import ceildiv, round_up
|
|
from tinygrad.uop.ops import UOp, Ops
|
|
from tinygrad.runtime.autogen import amdgpu_kd, hsa, libc
|
|
from tinygrad.renderer.amd.dsl import Reg, FixedBitField
|
|
from tinygrad.runtime.autogen.amd.common import OpType
|
|
|
|
# instructions used for padding
|
|
from tinygrad.runtime.autogen.amd.rdna3.ins import s_code_end # same encoding as RDNA4
|
|
from tinygrad.runtime.autogen.amd.cdna.ins import s_nop as s_nop_cdna
|
|
|
|
_arch_map = {"gfx9": "cdna", "gfx10": "rdna3", "gfx11": "rdna3", "gfx12": "rdna4"}
|
|
def assemble_linear(prg:UOp, lin:UOp, arch:str) -> bytes:
|
|
insts = [u.arg[0] for u in lin.src]
|
|
|
|
# ** scan for max vgpr/sgpr/accvgpr
|
|
max_vgpr, max_sgpr, max_accvgpr = 0, 0, 0
|
|
_ACCVGPR_TYPES = {OpType.OPR_ACCVGPR, OpType.OPR_SRC_ACCVGPR}
|
|
for inst in insts:
|
|
# build set of field names that are AccVGPR for this instruction
|
|
accvgpr_fields: set[str] = set()
|
|
for opr_name, (_, _, opr_type) in inst.operands.items():
|
|
if opr_type in _ACCVGPR_TYPES: accvgpr_fields.add(opr_name)
|
|
elif opr_type in {OpType.OPR_VGPR_OR_ACCVGPR, OpType.OPR_SRC_VGPR_OR_ACCVGPR, OpType.OPR_SRC_VGPR_OR_ACCVGPR_OR_CONST}:
|
|
if getattr(inst, 'acc_cd', 0) == 1: accvgpr_fields.add(opr_name)
|
|
for name, field in inst._fields:
|
|
if isinstance(field, FixedBitField): continue
|
|
val = getattr(inst, name)
|
|
if not isinstance(val, Reg): continue
|
|
if 256 <= val.offset < 512:
|
|
if name in accvgpr_fields: max_accvgpr = max(max_accvgpr, (val.offset - 256) + val.sz)
|
|
else: max_vgpr = max(max_vgpr, (val.offset - 256) + val.sz)
|
|
elif val.offset < 106: max_sgpr = max(max_sgpr, val.offset + val.sz)
|
|
|
|
# ** scan sink for metadata
|
|
sink, param_sizes, lds_size, gids = prg.src[0], {}, 0, set()
|
|
for u in sink.toposort():
|
|
if u.op is Ops.PARAM: param_sizes[u.arg.slot] = u.dtype.itemsize if u.addrspace is AddrSpace.ALU else 8
|
|
elif u.op is Ops.BUFFER and u.addrspace is AddrSpace.LOCAL: lds_size += u.max_numel() * u.dtype.itemsize
|
|
elif u.op is Ops.SPECIAL and u.arg.startswith("gidx"): gids.add(int(u.arg[-1]))
|
|
code_bytes = b"".join(inst.to_bytes() for inst in insts)
|
|
arch = next(v for k, v in _arch_map.items() if arch.startswith(k))
|
|
is_cdna, is_rdna4 = arch == "cdna", arch == "rdna4"
|
|
|
|
# ** pad text to ISA alignment
|
|
padding_inst = (s_nop_cdna(0) if is_cdna else s_code_end()).to_bytes()
|
|
text = code_bytes + padding_inst * ((hsa.AMD_ISA_ALIGN_BYTES - len(code_bytes) % hsa.AMD_ISA_ALIGN_BYTES) % hsa.AMD_ISA_ALIGN_BYTES)
|
|
text_offset = round_up(ctypes.sizeof(libc.Elf64_Ehdr), hsa.AMD_ISA_ALIGN_BYTES)
|
|
|
|
# ** pack kernel descriptor (rodata)
|
|
# CDNA: total VGPRs = regular VGPRs + AccVGPRs, each rounded to granularity of 4
|
|
accum_offset = round_up(max_vgpr, 4) if max_accvgpr > 0 else 0
|
|
next_free_vgpr = round_up(accum_offset + max_accvgpr, 8) if max_accvgpr > 0 else round_up(max_vgpr, 8)
|
|
next_free_sgpr = round_up(max_sgpr, 8)
|
|
vgpr_granule = max(0, (next_free_vgpr + 7) // 8 - 1)
|
|
# CDNA: add 6 for VCC(2) + FLAT_SCRATCH(2) + XNACK_MASK(2), next_free_sgpr is unused in RDNA.
|
|
sgpr_granule = max(0, ceildiv(next_free_sgpr + 6, 8) - 1) if is_cdna else 0
|
|
desc = amdgpu_kd.llvm_amdhsa_kernel_descriptor_t()
|
|
desc.group_segment_fixed_size = lds_size
|
|
for sz in (param_sizes[i] for i in sorted(param_sizes)): desc.kernarg_size = round_up(desc.kernarg_size, sz) + sz
|
|
desc.kernel_code_entry_byte_offset = -len(text)
|
|
|
|
# https://llvm.org/docs/AMDGPUUsage.html#amdgpu-amdhsa-compute-pgm-rsrc1-gfx6-gfx12-table
|
|
# NOTE: CU mode is the default
|
|
desc.compute_pgm_rsrc1 = (vgpr_granule << amdgpu_kd.COMPUTE_PGM_RSRC1_GRANULATED_WORKITEM_VGPR_COUNT_SHIFT |
|
|
sgpr_granule << amdgpu_kd.COMPUTE_PGM_RSRC1_GRANULATED_WAVEFRONT_SGPR_COUNT_SHIFT |
|
|
3 << amdgpu_kd.COMPUTE_PGM_RSRC1_FLOAT_DENORM_MODE_16_64_SHIFT |
|
|
(0 if is_rdna4 else 1) << amdgpu_kd.COMPUTE_PGM_RSRC1_GFX6_GFX11_ENABLE_DX10_CLAMP_SHIFT |
|
|
(0 if is_rdna4 else 1) << amdgpu_kd.COMPUTE_PGM_RSRC1_GFX6_GFX11_ENABLE_IEEE_MODE_SHIFT |
|
|
(0 if is_cdna else 1) << amdgpu_kd.COMPUTE_PGM_RSRC1_GFX10_PLUS_MEM_ORDERED_SHIFT)
|
|
desc.compute_pgm_rsrc2 = (2 << amdgpu_kd.COMPUTE_PGM_RSRC2_USER_SGPR_COUNT_SHIFT |
|
|
int(0 in gids) << amdgpu_kd.COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_X_SHIFT |
|
|
int(1 in gids) << amdgpu_kd.COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_Y_SHIFT |
|
|
int(2 in gids) << amdgpu_kd.COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_Z_SHIFT)
|
|
desc.kernel_code_properties = (1 << amdgpu_kd.KERNEL_CODE_PROPERTY_ENABLE_SGPR_KERNARG_SEGMENT_PTR_SHIFT |
|
|
(0 if is_cdna else 1) << amdgpu_kd.KERNEL_CODE_PROPERTY_ENABLE_WAVEFRONT_SIZE32_SHIFT)
|
|
if is_cdna and max_accvgpr > 0:
|
|
desc.compute_pgm_rsrc3 = max(0, accum_offset // 4 - 1) << amdgpu_kd.COMPUTE_PGM_RSRC3_GFX90A_ACCUM_OFFSET_SHIFT
|
|
rodata = bytes(desc)
|
|
|
|
# ** pack ELF
|
|
sh_names:list[int] = []
|
|
strtab = bytearray(b"\x00")
|
|
for name in [".text", ".rodata", ".strtab"]:
|
|
sh_names.append(len(strtab))
|
|
strtab += name.encode("ascii") + b"\x00"
|
|
|
|
rodata_offset = round_up(text_offset + (text_size := len(text)), hsa.AMD_KERNEL_CODE_ALIGN_BYTES)
|
|
strtab_offset = rodata_offset + (rodata_size := len(rodata))
|
|
shdr_offset = strtab_offset + (strtab_size := len(strtab))
|
|
|
|
sections = [(libc.SHT_PROGBITS, libc.SHF_ALLOC | libc.SHF_EXECINSTR, text_offset, text_offset, text_size),
|
|
(libc.SHT_PROGBITS, libc.SHF_ALLOC, rodata_offset, rodata_offset, rodata_size),
|
|
(libc.SHT_STRTAB, 0, 0, strtab_offset, strtab_size)]
|
|
shdrs = (libc.Elf64_Shdr * len(sections))()
|
|
for i, s in enumerate(sections): shdrs[i] = libc.Elf64_Shdr(sh_names[i], *s)
|
|
|
|
ehdr = libc.Elf64_Ehdr()
|
|
ehdr.e_ident[:5], ehdr.e_shoff, ehdr.e_shnum, ehdr.e_shstrndx = b"\x7FELF\x02", shdr_offset, len(sections), 2
|
|
|
|
elf = bytearray(shdr_offset + ctypes.sizeof(shdrs))
|
|
elf[0:ctypes.sizeof(ehdr)] = bytes(ehdr)
|
|
elf[text_offset:text_offset+text_size] = text
|
|
elf[rodata_offset:rodata_offset+rodata_size] = rodata
|
|
elf[strtab_offset:strtab_offset+strtab_size] = strtab
|
|
elf[shdr_offset:shdr_offset+ctypes.sizeof(shdrs)] = bytes(shdrs)
|
|
binary = bytes(elf)
|
|
|
|
return binary
|