From 74266eaee5fcb9a8d3f0b0879a538ca87b5dcf3b Mon Sep 17 00:00:00 2001 From: George Hotz Date: Thu, 25 Dec 2025 16:45:12 -0500 Subject: [PATCH] more handwritten --- extra/assembly/rdna3/asm.py | 51 ++++++++++++------- extra/assembly/rdna3/lib.py | 14 +++-- extra/assembly/rdna3/test/test_handwritten.py | 15 ++++++ extra/assembly/rdna3/test/test_llvm.py | 23 ++++++--- 4 files changed, 73 insertions(+), 30 deletions(-) diff --git a/extra/assembly/rdna3/asm.py b/extra/assembly/rdna3/asm.py index 2aaba134fe..1f0ea98da1 100644 --- a/extra/assembly/rdna3/asm.py +++ b/extra/assembly/rdna3/asm.py @@ -265,6 +265,11 @@ def disasm(inst: Inst) -> str: is_shift64 = 'rev' in op_name and '64' in op_name and op_name.startswith('v_') # v_ldexp_f64: 64-bit src0 (mantissa), 32-bit src1 (exponent) is_ldexp64 = op_name == 'v_ldexp_f64' + # SAD/QSAD/MQSAD instructions have mixed sizes + # v_qsad_pk_u16_u8, v_mqsad_pk_u16_u8: 64-bit dst/src0/src2, 32-bit src1 + # v_mqsad_u32_u8: 128-bit (4 reg) dst/src2, 64-bit src0, 32-bit src1 + is_sad64 = any(x in op_name for x in ('qsad_pk', 'mqsad_pk')) + is_mqsad_u32 = 'mqsad_u32' in op_name # Detect conversion ops: v_cvt_{dst_type}_{src_type} - each side may have different size # Also handle v_cvt_pk_* which packs two values into one if 'cvt_pk' in op_name: @@ -282,31 +287,32 @@ def disasm(inst: Inst) -> str: is_f16_src = True # src0 and src1 are 16-bit is_f16_src2 = False # src2 is 32-bit else: - # 16-bit ops need .h/.l suffix, but packed ops (dot2, pk) don't - is_16bit_op = ('f16' in op_name or 'i16' in op_name or 'u16' in op_name or 'b16' in op_name) and 'dot2' not in op_name + # 16-bit ops need .h/.l suffix, but packed ops (dot2, pk_, sad, msad, qsad, mqsad) don't + is_16bit_op = ('f16' in op_name or 'i16' in op_name or 'u16' in op_name or 'b16' in op_name) and not any(x in op_name for x in ('dot2', 'pk_', 'sad', 'msad', 'qsad', 'mqsad')) is_f16_dst = is_f16_src = is_f16_src2 = is_16bit_op - def fmt_vop3_src(v, neg_bit, abs_bit, hi_bit=False, force_64=False, is_16=False): + def fmt_vop3_src(v, neg_bit, abs_bit, hi_bit=False, reg_cnt=1, is_16=False): s = fmt_src(v) - # Add register pair for f64, or .h suffix for f16 VGPRs with opsel - if force_64 and v >= 256: s = _vreg(v - 256, 2) - elif force_64 and v <= 105: s = _sreg(v, 2) - elif force_64 and v == 106: s = "vcc" - elif force_64 and v == 126: s = "exec" - elif force_64 and 108 <= v <= 123: s = f"ttmp[{v-108}:{v-108+1}]" + # Add register pair/quad for 64/128-bit, or .h suffix for f16 VGPRs with opsel + if reg_cnt > 1 and v >= 256: s = _vreg(v - 256, reg_cnt) + elif reg_cnt > 1 and v <= 105: s = _sreg(v, reg_cnt) + elif reg_cnt == 2 and v == 106: s = "vcc" + elif reg_cnt == 2 and v == 126: s = "exec" + elif reg_cnt > 1 and 108 <= v <= 123: s = f"ttmp[{v-108}:{v-108+reg_cnt-1}]" elif is_16 and v >= 256: s = f"v{v - 256}.h" if hi_bit else f"v{v - 256}.l" if abs_bit: s = f"|{s}|" if neg_bit: s = f"-{s}" return s - # Determine which sources are 64-bit - src0_64 = is_f64 and not is_shift64 # shift ops have 32-bit shift amount - src1_64 = is_f64 and not is_class and not is_ldexp64 # class/ldexp ops have 32-bit src1 - src2_64 = is_f64 - src0_str = fmt_vop3_src(src0, neg & 1, abs_ & 1, opsel & 1, src0_64, is_f16_src) - src1_str = fmt_vop3_src(src1, neg & 2, abs_ & 2, opsel & 2, src1_64, is_f16_src) - src2_str = fmt_vop3_src(src2, neg & 4, abs_ & 4, opsel & 4, src2_64, is_f16_src2) + # Determine register count for each source + src0_cnt = 2 if ((is_f64 and not is_shift64) or is_sad64 or is_mqsad_u32) else 1 + src1_cnt = 2 if (is_f64 and not is_class and not is_ldexp64) else 1 + src2_cnt = 4 if is_mqsad_u32 else 2 if (is_f64 or is_sad64) else 1 + src0_str = fmt_vop3_src(src0, neg & 1, abs_ & 1, opsel & 1, src0_cnt, is_f16_src) + src1_str = fmt_vop3_src(src1, neg & 2, abs_ & 2, opsel & 2, src1_cnt, is_f16_src) + src2_str = fmt_vop3_src(src2, neg & 4, abs_ & 4, opsel & 4, src2_cnt, is_f16_src2) # Format destination - for 16-bit ops, use .h/.l suffix - if is_f64: - dst_str = _vreg(vdst, 2) + dst_cnt = 4 if is_mqsad_u32 else 2 if (is_f64 or is_sad64) else 1 + if dst_cnt > 1: + dst_str = _vreg(vdst, dst_cnt) elif is_f16_dst: dst_str = f"v{vdst}.h" if (opsel & 8) else f"v{vdst}.l" else: @@ -537,6 +543,15 @@ def asm(text: str) -> Inst: parts = text.replace(',', ' ').split() if not parts: raise ValueError("empty instruction") mnemonic, op_str = parts[0].lower(), text[len(parts[0]):].strip() + # Handle s_waitcnt specially before operand parsing + if mnemonic == 's_waitcnt': + vmcnt, expcnt, lgkmcnt = 0x3f, 0x7, 0x3f + for part in op_str.replace(',', ' ').split(): + if m := re.match(r'vmcnt\((\d+)\)', part): vmcnt = int(m.group(1)) + elif m := re.match(r'expcnt\((\d+)\)', part): expcnt = int(m.group(1)) + elif m := re.match(r'lgkmcnt\((\d+)\)', part): lgkmcnt = int(m.group(1)) + elif re.match(r'^0x[0-9a-f]+$|^\d+$', part): return autogen.s_waitcnt(simm16=int(part, 0)) + return autogen.s_waitcnt(simm16=waitcnt(vmcnt, expcnt, lgkmcnt)) operands, current, depth, in_pipe = [], "", 0, False for ch in op_str: if ch == '[': depth += 1 diff --git a/extra/assembly/rdna3/lib.py b/extra/assembly/rdna3/lib.py index a34609d549..e585c001d2 100644 --- a/extra/assembly/rdna3/lib.py +++ b/extra/assembly/rdna3/lib.py @@ -81,7 +81,11 @@ class Inst: annotations.update(getattr(cls, '__annotations__', {})) # Type check and encode values for name, val in list(self._values.items()): - if name == 'encoding' or isinstance(val, RawImm): continue + if name == 'encoding': continue + # For RawImm, only process RAW_FIELDS to unwrap to int + if isinstance(val, RawImm): + if name in RAW_FIELDS: self._values[name] = val.val + continue ann = annotations.get(name) # Type validation if ann is SGPR: @@ -98,9 +102,11 @@ class Inst: if encoded == 255 and self._literal is None and isinstance(val, int) and not isinstance(val, IntEnum): self._literal = val # Encode raw register fields for consistent repr - elif name in RAW_FIELDS and isinstance(val, Reg): - encoded = (108 + val.idx) if isinstance(val, TTMP) else (val.idx | (0x80 if val.hi else 0)) - self._values[name] = encoded + elif name in RAW_FIELDS: + if isinstance(val, Reg): + self._values[name] = (108 + val.idx) if isinstance(val, TTMP) else (val.idx | (0x80 if val.hi else 0)) + elif hasattr(val, 'value'): # IntEnum like SrcEnum.NULL + self._values[name] = val.value # Encode sbase (divided by 2) and srsrc/ssamp (divided by 4) elif name == 'sbase' and isinstance(val, Reg): self._values[name] = val.idx // 2 diff --git a/extra/assembly/rdna3/test/test_handwritten.py b/extra/assembly/rdna3/test/test_handwritten.py index b849f6b381..1f73d1adfe 100644 --- a/extra/assembly/rdna3/test/test_handwritten.py +++ b/extra/assembly/rdna3/test/test_handwritten.py @@ -55,5 +55,20 @@ class TestIntegration(unittest.TestCase): def test_simple_int_to_v(self): self.inst = v_mov_b32_e32(v[0], 1) + def test_three_add(self): + self.inst = v_add_co_ci_u32_e32(v[3], s[7], v[3]) + + def test_three_add_v(self): + self.inst = v_add_co_ci_u32_e32(v[3], v[7], v[3]) + + def test_three_add_const(self): + self.inst = v_add_co_ci_u32_e32(v[3], 2.0, v[3]) + + def test_swaitcnt_lgkm(self): self.inst = s_waitcnt(0xfc07) + def test_swaitcnt_vm(self): self.inst = s_waitcnt(0x03f7) + + def test_vmad(self): + self.inst = v_mad_u64_u32(v[1:2], NULL, s[2], 3, v[1:2]) + if __name__ == "__main__": unittest.main() diff --git a/extra/assembly/rdna3/test/test_llvm.py b/extra/assembly/rdna3/test/test_llvm.py index 3f6513a7d6..94d178c628 100644 --- a/extra/assembly/rdna3/test/test_llvm.py +++ b/extra/assembly/rdna3/test/test_llvm.py @@ -116,15 +116,22 @@ def _make_disasm_test(name): for asm_text, data in self.tests.get(name, []): if len(data) > fmt_cls._size(): continue # skip literals (need different handling) try: - decoded = fmt_cls.from_bytes(data) - op_val = decoded._values.get('op', 0) - op_val = op_val.val if hasattr(op_val, 'val') else op_val - # VOP3SD test uses VOP3 file - skip non-VOP3SD instructions - if name == 'vop3sd' and op_val not in vop3sd_opcodes: continue - # VOP3 and VOP3SD share encoding - validate with appropriate enum - if fmt_cls.__name__ == 'VOP3' and op_val in vop3sd_opcodes: - VOP3SDOp(op_val) # validate as VOP3SD + # VOP3 and VOP3SD share encoding - peek at opcode to determine which class to use + if fmt_cls.__name__ in ('VOP3', 'VOP3SD'): + temp = VOP3.from_bytes(data) + op_val = temp._values.get('op', 0) + op_val = op_val.val if hasattr(op_val, 'val') else op_val + is_vop3sd = op_val in vop3sd_opcodes + decoded = VOP3SD.from_bytes(data) if is_vop3sd else VOP3.from_bytes(data) + # Validate opcode with appropriate enum + if is_vop3sd: + VOP3SDOp(op_val) + else: + VOP3Op(op_val) else: + decoded = fmt_cls.from_bytes(data) + op_val = decoded._values.get('op', 0) + op_val = op_val.val if hasattr(op_val, 'val') else op_val op_enum(op_val) # validate opcode if decoded.to_bytes()[:len(data)] != data: failed += 1; failures.append(f"decode roundtrip failed for {data.hex()}"); continue