From 5c33ff65362253befa4a1ffde50eca37f4a52d5c Mon Sep 17 00:00:00 2001 From: George Hotz Date: Wed, 28 Jan 2026 12:48:42 +0800 Subject: [PATCH] correct fix + tests --- extra/assembly/amd/emu.py | 38 +++--------- extra/assembly/amd/test/hw/test_sop.py | 84 +++++++++++++++++++++++--- 2 files changed, 84 insertions(+), 38 deletions(-) diff --git a/extra/assembly/amd/emu.py b/extra/assembly/amd/emu.py index c43e32adbc..65d4927b08 100644 --- a/extra/assembly/amd/emu.py +++ b/extra/assembly/amd/emu.py @@ -349,8 +349,12 @@ class _Ctx: if bits == 64: sgpr_val = _u64(sgpr_lo, self.rsgpr_dyn(off + _c(1))) - # Float constants: cast F32 to F64; integer inline: duplicate lo - inline = is_float_const.where(sgpr_lo.bitcast(dtypes.float32).cast(dtypes.float64).bitcast(dtypes.uint64), _u64(sgpr_lo, sgpr_lo)) + # Integer inline constants: sign-extend 32-bit value from buffer to 64-bit + # Float constants: cast F32 to F64 + int_inline = sgpr_lo.cast(dtypes.int32).cast(dtypes.int64) + float_inline = sgpr_lo.bitcast(dtypes.float32).cast(dtypes.float64) + # compute inline + inline = is_float_const.where(float_inline.bitcast(dtypes.uint64), int_inline.bitcast(dtypes.uint64)) if literal is not None: inline = off.eq(_c(255)).where(literal.cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32), inline) scalar_val = (off < _c(128)).where(sgpr_val, inline) else: @@ -510,32 +514,6 @@ def _compile_sop(inst: SOP1 | SOP2 | SOPC | SOPK, ctx: _Ctx) -> UOp: bits = inst.canonical_op_bits literal = ctx.inst_field(type(inst).literal) if hasattr(type(inst), 'literal') else None - def rsrc_dyn_scalar(off: UOp, is_64bit: bool) -> UOp: - """Read scalar source with dynamic offset (SGPR or inline constant). - For SOP, off is always 0-255 (SGPR or inline constant, never VGPR). - SGPR buffer has 260 entries: 0-127=SGPRs, 128-255=inline constants, 256-259=special.""" - is_sgpr = off < _c(128) - # For 64-bit: read SGPR pair if off < 128, else compute inline constant as 64-bit - # (can't just read from buffer since buffer has 32-bit values) - if is_64bit: - sgpr_val = _u64(ctx.rsgpr_dyn(off), ctx.rsgpr_dyn(off + _c(1))) - # Build inline constant: 128-192 = 0-64, 193-208 = -1 to -16 - inline_val = (off - _c(128)).cast(dtypes.uint64) # positive inline 0-64 - neg_val = (_c(192) - off).cast(dtypes.int64).cast(dtypes.uint64) # negative -1 to -16 - lit_val = literal.cast(dtypes.uint64) if literal is not None else UOp.const(dtypes.uint64, 0) - # Select between sgpr, positive inline, negative inline, or literal - is_neg_inline = (off >= _c(193)) & (off < _c(209)) - is_literal = off.eq(_c(255)) if literal is not None else UOp.const(dtypes.bool, False) - val = is_sgpr.where(sgpr_val, is_neg_inline.where(neg_val, is_literal.where(lit_val, inline_val))) - return val - # 32-bit: read from SGPR buffer (inline constants 128-255 are pre-populated) - # off is always 0-255 for SOP, all valid SGPR indices - sgpr_val = ctx.rsgpr_dyn(off) - # Handle literal (255) - literal value overrides the pre-populated 0 - if literal is not None: - sgpr_val = off.eq(_c(255)).where(literal, sgpr_val) - return sgpr_val - if isinstance(inst, SOPK): sdst_off = ctx.inst_field(SOPK.sdst) simm16 = ctx.inst_field(SOPK.simm16) @@ -546,9 +524,7 @@ def _compile_sop(inst: SOP1 | SOP2 | SOPC | SOPK, ctx: _Ctx) -> UOp: elif isinstance(inst, SOP1): sdst_off = ctx.inst_field(SOP1.sdst) ssrc0_off = ctx.inst_field(SOP1.ssrc0) - srcs = {'S0': rsrc_dyn_scalar(ssrc0_off, bits['s0'] == 64)} - # TODO: this is broken - #srcs = {'S0': ctx.rsrc_dyn(ssrc0_off, None, bits['s0'], literal)} + srcs = {'S0': ctx.rsrc_dyn(ssrc0_off, None, bits['s0'], literal)} dst_off, dst_size = sdst_off, bits['d'] // 32 elif isinstance(inst, SOP2): sdst_off = ctx.inst_field(SOP2.sdst) diff --git a/extra/assembly/amd/test/hw/test_sop.py b/extra/assembly/amd/test/hw/test_sop.py index 856cf63178..7785ec35d2 100644 --- a/extra/assembly/amd/test/hw/test_sop.py +++ b/extra/assembly/amd/test/hw/test_sop.py @@ -722,20 +722,90 @@ class TestNullRegister(unittest.TestCase): class Test64BitSOP1InlineConstants(unittest.TestCase): """Tests for 64-bit SOP1 instructions with inline constants. - Regression tests for bug where rsrc_dyn with lane=None didn't properly - handle 64-bit inline constants, causing incorrect values to be read. + Regression tests for bug where rsrc_dyn didn't properly handle 64-bit + inline constants, incorrectly duplicating lo bits to hi instead of + zero/sign-extending. """ - def test_s_mov_b64_inline_16(self): - """S_MOV_B64 with inline constant 16 should set lo=16, hi=0.""" + def test_s_mov_b64_inline_0(self): + """S_MOV_B64 with inline constant 0.""" instructions = [ - s_mov_b64(s[0:1], 16), # inline constant 16 + s_mov_b64(s[0:1], 0), v_mov_b32_e32(v[0], s[0]), v_mov_b32_e32(v[1], s[1]), ] st = run_program(instructions, n_lanes=1) - self.assertEqual(st.vgpr[0][0], 16, "lo should be 16") - self.assertEqual(st.vgpr[0][1], 0, "hi should be 0") + self.assertEqual(st.vgpr[0][0], 0) + self.assertEqual(st.vgpr[0][1], 0) + + def test_s_mov_b64_inline_16(self): + """S_MOV_B64 with inline constant 16 should set lo=16, hi=0.""" + instructions = [ + s_mov_b64(s[0:1], 16), + v_mov_b32_e32(v[0], s[0]), + v_mov_b32_e32(v[1], s[1]), + ] + st = run_program(instructions, n_lanes=1) + self.assertEqual(st.vgpr[0][0], 16) + self.assertEqual(st.vgpr[0][1], 0) + + def test_s_mov_b64_inline_64(self): + """S_MOV_B64 with inline constant 64 (max positive).""" + instructions = [ + s_mov_b64(s[0:1], 64), + v_mov_b32_e32(v[0], s[0]), + v_mov_b32_e32(v[1], s[1]), + ] + st = run_program(instructions, n_lanes=1) + self.assertEqual(st.vgpr[0][0], 64) + self.assertEqual(st.vgpr[0][1], 0) + + def test_s_mov_b64_inline_neg1(self): + """S_MOV_B64 with inline constant -1 should sign-extend.""" + instructions = [ + s_mov_b64(s[0:1], -1), + v_mov_b32_e32(v[0], s[0]), + v_mov_b32_e32(v[1], s[1]), + ] + st = run_program(instructions, n_lanes=1) + self.assertEqual(st.vgpr[0][0], 0xFFFFFFFF) + self.assertEqual(st.vgpr[0][1], 0xFFFFFFFF) + + def test_s_mov_b64_inline_neg16(self): + """S_MOV_B64 with inline constant -16 should sign-extend.""" + instructions = [ + s_mov_b64(s[0:1], -16), + v_mov_b32_e32(v[0], s[0]), + v_mov_b32_e32(v[1], s[1]), + ] + st = run_program(instructions, n_lanes=1) + self.assertEqual(st.vgpr[0][0], 0xFFFFFFF0) + self.assertEqual(st.vgpr[0][1], 0xFFFFFFFF) + + def test_s_or_b64_inline_constant(self): + """S_OR_B64 with 64-bit inline constant.""" + instructions = [ + s_mov_b64(s[0:1], 0), + s_or_b64(s[2:3], s[0:1], 16), + v_mov_b32_e32(v[0], s[2]), + v_mov_b32_e32(v[1], s[3]), + ] + st = run_program(instructions, n_lanes=1) + self.assertEqual(st.vgpr[0][0], 16) + self.assertEqual(st.vgpr[0][1], 0) + + def test_s_and_b64_inline_constant(self): + """S_AND_B64 with 64-bit inline constant.""" + instructions = [ + s_mov_b32(s[0], 0xFFFFFFFF), + s_mov_b32(s[1], 0xFFFFFFFF), + s_and_b64(s[2:3], s[0:1], 16), + v_mov_b32_e32(v[0], s[2]), + v_mov_b32_e32(v[1], s[3]), + ] + st = run_program(instructions, n_lanes=1) + self.assertEqual(st.vgpr[0][0], 16) + self.assertEqual(st.vgpr[0][1], 0) if __name__ == '__main__':