mawei/StarPilot
1
0
Fork 0
mirror of https://github.com/firestar5683/StarPilot.git synced 2026-11-08 13:33:42 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
63243d7e2d768514b61b5cd5d357200ecab7e208
StarPilot/tinygrad_repo/test/amd/hw/test_ds.py
T
firestar5683 22a8575199 Vendor dependencies and resources for standalone builds 
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.
2026-10-03 21:01:18 +00:00

984 lines
41 KiB
Python
Raw Blame History

"""Tests for DS instructions - data share (LDS) operations.
Includes: ds_store_b32, ds_load_b32, ds_store_2addr_*, ds_load_2addr_*,
ds_add_*, ds_max_*, ds_min_*, ds_and_*, ds_or_*, ds_xor_*,
ds_inc_*, ds_dec_*, ds_cmpstore_*, ds_storexchg_*
"""
import unittest
from test.amd.hw.helpers import *
class TestDSSwizzle(unittest.TestCase):
def test_modes_and_overlapping_registers(self):
for offset in (0x041f, 0x401f, 0x7c1f, 0x00a0, 0x801b, 0xc020, 0xc420, 0xc021, 0xe000, 0xe010, 0xe01f):
for dst in (0, 1):
with self.subTest(offset=hex(offset), dst=dst):
st = run_program([
v_add_nc_u32_e32(v[0], 1, v[255]),
ds_swizzle_b32(vdst=v[dst], addr=v[0], offset0=offset & 255, offset1=offset >> 8),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
], n_lanes=32)
self.assertEqual(sorted(st.vgpr[i][dst] for i in range(32)), [6]*32 if offset == 0x00a0 else list(range(1, 33)))
def test_inactive_sources_and_destinations(self):
st = run_program([
v_add_nc_u32_e32(v[0], 1, v[255]),
v_mov_b32_e32(v[1], 99),
s_mov_b32(EXEC_LO, 0x55555555),
ds_swizzle_b32(vdst=v[1], addr=v[0], offset0=0x1f, offset1=4),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(EXEC_LO, 0xffffffff),
], n_lanes=32)
self.assertEqual([st.vgpr[i][1] for i in range(32)], [0, 99]*16)
class TestDS2Addr(unittest.TestCase):
"""Tests for DS_*_2ADDR instructions."""
def test_ds_store_load_2addr_b32(self):
"""DS_STORE_2ADDR_B32 and DS_LOAD_2ADDR_B32 with offset * 4."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2:3], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
def test_ds_store_load_2addr_b64(self):
"""DS_STORE_2ADDR_B64 and DS_LOAD_2ADDR_B64."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0x12345678),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x9ABCDEF0),
v_mov_b32_e32(v[3], s[0]),
DS(DSOp.DS_STORE_2ADDR_B64, addr=v[10], data0=v[0:1], data1=v[2:3], vdst=v[0], offset0=0, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[10], vdst=v[4:7], offset0=0, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE)
self.assertEqual(st.vgpr[0][6], 0x12345678)
self.assertEqual(st.vgpr[0][7], 0x9ABCDEF0)
class TestDS2AddrMore(unittest.TestCase):
"""Additional DS_*_2ADDR tests."""
def test_ds_store_load_2addr_b32_nonzero_offsets(self):
"""DS_STORE_2ADDR_B32 with non-zero offsets (offset*4 scaling)."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0x11111111),
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x22222222),
v_mov_b32_e32(v[1], s[2]),
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=2, offset1=5),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2:3], offset0=2, offset1=5),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should have value from offset 8 (2*4)")
self.assertEqual(st.vgpr[0][3], 0x22222222, "v3 should have value from offset 20 (5*4)")
def test_ds_2addr_b64_no_overlap(self):
"""DS_LOAD_2ADDR_B64 with adjacent offsets should not overlap."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0x11111111),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_mov_b32(s[2], 0x22222222),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=4),
s_mov_b32(s[2], 0x33333333),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=8),
s_mov_b32(s[2], 0x44444444),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=12),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[10], vdst=v[4:7], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should be 0x11111111")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should be 0x22222222")
self.assertEqual(st.vgpr[0][6], 0x33333333, "v6 should be 0x33333333")
self.assertEqual(st.vgpr[0][7], 0x44444444, "v7 should be 0x44444444")
def test_ds_load_2addr_b32_no_overwrite(self):
"""DS_LOAD_2ADDR_B32 should only write 2 VGPRs."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[2]),
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 0xDEADBEEF),
v_mov_b32_e32(v[4], s[2]), # Sentinel
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2:3], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF, "v4 should be untouched")
def test_ds_load_2addr_b64_addr_overlaps_vdst(self):
"""DS_LOAD_2ADDR_B64 where addr register overlaps vdst range.
Hardware reads the address before writing any results, so addr=v[4]
with vdst=v[4:7] must load all 4 dwords using the original v[4] value.
"""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_mov_b32(s[2], 0xBBBBBBBB),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=4),
s_mov_b32(s[2], 0xCCCCCCCC),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=8),
s_mov_b32(s[2], 0xDDDDDDDD),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=12),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# addr=v[4] overlaps vdst=v[4:7]
v_mov_b32_e32(v[4], 0),
DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[4], vdst=v[4:7], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA, "v4 = LDS[0:4]")
self.assertEqual(st.vgpr[0][5], 0xBBBBBBBB, "v5 = LDS[4:8]")
self.assertEqual(st.vgpr[0][6], 0xCCCCCCCC, "v6 = LDS[8:12]")
self.assertEqual(st.vgpr[0][7], 0xDDDDDDDD, "v7 = LDS[12:16]")
def test_ds_load_2addr_b32_addr_overlaps_vdst(self):
"""DS_LOAD_2ADDR_B32 where addr register overlaps vdst range."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_mov_b32(s[2], 0xBBBBBBBB),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=4),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# addr=v[2] overlaps vdst=v[2:3]
v_mov_b32_e32(v[2], 0),
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[2], vdst=v[2:3], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA, "v2 = LDS[0:4]")
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB, "v3 = LDS[4:8]")
def test_ds_load_b64_no_overwrite(self):
"""DS_LOAD_B64 should only write 2 VGPRs."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xDEADBEEF),
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0xCAFEBABE),
v_mov_b32_e32(v[1], s[2]),
ds_store_b64(addr=v[10], data0=v[0:1], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 0x12345678),
v_mov_b32_e32(v[4], s[2]), # Sentinel
ds_load_b64(addr=v[10], vdst=v[2:3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][3], 0xCAFEBABE)
self.assertEqual(st.vgpr[0][4], 0x12345678, "v4 should be untouched")
class TestDSB96(unittest.TestCase):
"""Tests for DS_STORE_B96 and DS_LOAD_B96 (96-bit / 3 dwords)."""
def test_ds_store_load_b96(self):
"""DS_STORE_B96 stores 3 VGPRs, DS_LOAD_B96 loads them back."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0x33333333),
v_mov_b32_e32(v[2], s[0]),
ds_store_b96(addr=v[10], data0=v[0:2]),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b96(addr=v[10], vdst=v[4:6]),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have first dword")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have second dword")
self.assertEqual(st.vgpr[0][6], 0x33333333, "v6 should have third dword")
def test_ds_store_b96_with_offset(self):
"""DS_STORE_B96 with non-zero offset."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0xCCCCCCCC),
v_mov_b32_e32(v[2], s[0]),
DS(DSOp.DS_STORE_B96, addr=v[10], data0=v[0:2], offset0=12),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_B96, addr=v[10], vdst=v[4:6], offset0=12),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][5], 0xBBBBBBBB)
self.assertEqual(st.vgpr[0][6], 0xCCCCCCCC)
class TestDSB128(unittest.TestCase):
"""Tests for DS_STORE_B128 and DS_LOAD_B128 (128-bit / 4 dwords)."""
def test_ds_store_load_b128(self):
"""DS_STORE_B128 stores 4 VGPRs, DS_LOAD_B128 loads them back."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0x33333333),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x44444444),
v_mov_b32_e32(v[3], s[0]),
ds_store_b128(addr=v[10], data0=v[0:3]),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b128(addr=v[10], vdst=v[4:7]),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have first dword")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have second dword")
self.assertEqual(st.vgpr[0][6], 0x33333333, "v6 should have third dword")
self.assertEqual(st.vgpr[0][7], 0x44444444, "v7 should have fourth dword")
def test_ds_store_b128_with_offset(self):
"""DS_STORE_B128 with non-zero offset."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0xCCCCCCCC),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xDDDDDDDD),
v_mov_b32_e32(v[3], s[0]),
DS(DSOp.DS_STORE_B128, addr=v[10], data0=v[0:3], offset0=16),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_B128, addr=v[10], vdst=v[4:7], offset0=16),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][5], 0xBBBBBBBB)
self.assertEqual(st.vgpr[0][6], 0xCCCCCCCC)
self.assertEqual(st.vgpr[0][7], 0xDDDDDDDD)
class TestDSAtomic(unittest.TestCase):
"""Tests for DS atomic operations."""
def test_ds_max_rtn_u32(self):
"""DS_MAX_RTN_U32: atomically store max and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[1], s[2]),
ds_max_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 100, "v2 should have old value (100)")
self.assertEqual(st.vgpr[0][3], 200, "v3 should have max(100, 200) = 200")
def test_ds_min_rtn_u32(self):
"""DS_MIN_RTN_U32: atomically store min and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[1], s[2]),
ds_min_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 200)
self.assertEqual(st.vgpr[0][3], 100)
def test_ds_and_rtn_b32(self):
"""DS_AND_RTN_B32: atomically AND and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xFF00FF00),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 0xFFFF0000),
v_mov_b32_e32(v[1], s[2]),
ds_and_rtn_b32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xFF00FF00)
self.assertEqual(st.vgpr[0][3], 0xFF000000)
def test_ds_or_rtn_b32(self):
"""DS_OR_RTN_B32: atomically OR and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0x00FF0000),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 0x000000FF),
v_mov_b32_e32(v[1], s[2]),
ds_or_rtn_b32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x00FF0000)
self.assertEqual(st.vgpr[0][3], 0x00FF00FF)
def test_ds_xor_rtn_b32(self):
"""DS_XOR_RTN_B32: atomically XOR and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 0xFFFFFFFF),
v_mov_b32_e32(v[1], s[2]),
ds_xor_rtn_b32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][3], 0x55555555)
def test_ds_inc_rtn_u32(self):
"""DS_INC_RTN_U32: increment with wrap."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 5),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 10), # limit
v_mov_b32_e32(v[1], s[2]),
ds_inc_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 5)
self.assertEqual(st.vgpr[0][3], 6)
def test_ds_dec_rtn_u32(self):
"""DS_DEC_RTN_U32: decrement with wrap."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 5),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 10), # limit
v_mov_b32_e32(v[1], s[2]),
ds_dec_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 5)
self.assertEqual(st.vgpr[0][3], 4)
def test_ds_cmpstore_b32_match(self):
"""DS_CMPSTORE_B32: conditional store when compare matches."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[1], s[2]), # new value
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[2], s[2]), # compare = 100 (matches)
ds_cmpstore_b32(addr=v[10], data0=v[1], data1=v[2], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[4], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 200)
def test_ds_cmpstore_b32_no_match(self):
"""DS_CMPSTORE_B32: no store when compare doesn't match."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[1], s[2]), # new value
s_mov_b32(s[2], 50),
v_mov_b32_e32(v[2], s[2]), # compare = 50 (doesn't match)
ds_cmpstore_b32(addr=v[10], data0=v[1], data1=v[2], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[4], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 100)
def test_ds_max_u32_no_rtn(self):
"""DS_MAX_U32 (no RTN): atomically store max, no return value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[1], s[2]),
ds_max_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][3], 200, "v3 should have max(100, 200) = 200")
def test_ds_add_u32_no_rtn_preserves_vdst(self):
"""DS_ADD_U32 (no RTN) should NOT write to vdst."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[2]), # sentinel
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 50),
v_mov_b32_e32(v[1], s[2]),
ds_add_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xDEADBEEF, "v2 should preserve sentinel")
self.assertEqual(st.vgpr[0][3], 150, "v3 should have 100 + 50 = 150")
def test_ds_add_rtn_u32_writes_vdst(self):
"""DS_ADD_RTN_U32 should write old value to vdst."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[2]), # sentinel
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 50),
v_mov_b32_e32(v[1], s[2]),
ds_add_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 100, "v2 should have old value (100)")
self.assertEqual(st.vgpr[0][3], 150, "v3 should have 100 + 50 = 150")
def test_ds_dec_rtn_u32_wrap(self):
"""DS_DEC_RTN_U32: decrement wraps when value is 0 or > limit."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0), # Start at 0
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 10), # limit
v_mov_b32_e32(v[1], s[2]),
ds_dec_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0, "v2 should have old value (0)")
# When mem == 0 or mem > limit, result = limit
self.assertEqual(st.vgpr[0][3], 10, "v3 should wrap to limit (10)")
class TestDSStorexchg(unittest.TestCase):
"""Tests for DS_STOREXCHG instructions."""
def test_ds_storexchg_rtn_b32(self):
"""DS_STOREXCHG_RTN_B32: exchange value and return old."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STOREXCHG_RTN_B32, addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
class TestDSRegisterWidth(unittest.TestCase):
"""Regression tests: DS loads should only write correct number of VGPRs."""
def test_ds_load_b32_no_overwrite(self):
"""DS_LOAD_B32 should only write 1 VGPR."""
instructions = [
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[2], s[0]), # sentinel
ds_store_b32(addr=v[0], data0=v[1], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[0], vdst=v[1], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should be untouched")
class TestDS2AddrStride64(unittest.TestCase):
"""Tests for DS_*_2ADDR_STRIDE64 (offset * 256 for B32, offset * 512 for B64)."""
def test_ds_store_load_2addr_stride64_b32(self):
"""DS_STORE_2ADDR_STRIDE64_B32: stores at ADDR + offset*256."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B32, addr=v[10], vdst=v[2:3], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA, "v2 from addr 256")
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB, "v3 from addr 512")
def test_ds_store_load_2addr_stride64_b64(self):
"""DS_STORE_2ADDR_STRIDE64_B64: stores at ADDR + offset*512."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0x12345678),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x9ABCDEF0),
v_mov_b32_e32(v[3], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B64, addr=v[10], data0=v[0:1], data1=v[2:3], vdst=v[0], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B64, addr=v[10], vdst=v[4:7], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE)
self.assertEqual(st.vgpr[0][6], 0x12345678)
self.assertEqual(st.vgpr[0][7], 0x9ABCDEF0)
def test_ds_storexchg_2addr_rtn_b32(self):
"""DS_STOREXCHG_2ADDR_RTN_B32: exchange at two addresses."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[0]),
DS(DSOp.DS_STOREXCHG_2ADDR_RTN_B32, addr=v[10], data0=v[2], data1=v[3], vdst=v[4:5], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[6:7], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "old val 0")
self.assertEqual(st.vgpr[0][5], 0x22222222, "old val 1")
self.assertEqual(st.vgpr[0][6], 0xAAAAAAAA, "new val 0")
self.assertEqual(st.vgpr[0][7], 0xBBBBBBBB, "new val 1")
def test_ds_storexchg_rtn_b64(self):
"""DS_STOREXCHG_RTN_B64: exchange 64-bit value and return old."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[0], s[0]), # initial low
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[1], s[0]), # initial high
DS(DSOp.DS_STORE_B64, addr=v[10], data0=v[0:1], vdst=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 0x12345678),
v_mov_b32_e32(v[2], s[0]), # new low
s_mov_b32(s[0], 0x9ABCDEF0),
v_mov_b32_e32(v[3], s[0]), # new high
DS(DSOp.DS_STOREXCHG_RTN_B64, addr=v[10], data0=v[2:3], vdst=v[4:5], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_B64, addr=v[10], vdst=v[6:7], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF, "v4 should have old low dword")
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE, "v5 should have old high dword")
self.assertEqual(st.vgpr[0][6], 0x12345678, "v6 should have new low dword")
self.assertEqual(st.vgpr[0][7], 0x9ABCDEF0, "v7 should have new high dword")
def test_ds_store_load_2addr_stride64_b64_roundtrip(self):
"""DS_STORE_2ADDR_STRIDE64_B64 followed by DS_LOAD_2ADDR_STRIDE64_B64 works correctly."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B64, addr=v[10], data0=v[0:1], data1=v[0:1], vdst=v[0], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B64, addr=v[10], vdst=v[2:5], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should have val1 low")
self.assertEqual(st.vgpr[0][3], 0x22222222, "v3 should have val1 high")
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have val2 low")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have val2 high")
def test_ds_storexchg_2addr_stride64_rtn_b32(self):
"""DS_STOREXCHG_2ADDR_STRIDE64_RTN_B32: exchange at two addresses (offset*256)."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[0]),
DS(DSOp.DS_STOREXCHG_2ADDR_STRIDE64_RTN_B32, addr=v[10], data0=v[2], data1=v[3], vdst=v[4:5], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B32, addr=v[10], vdst=v[6:7], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have old value")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have old value")
self.assertEqual(st.vgpr[0][6], 0xAAAAAAAA, "v6 should have new value")
self.assertEqual(st.vgpr[0][7], 0xBBBBBBBB, "v7 should have new value")
def test_ds_storexchg_2addr_stride64_rtn_b64_returns_old(self):
"""DS_STOREXCHG_2ADDR_STRIDE64_RTN_B64: returns old values correctly."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B64, addr=v[10], data0=v[0:1], data1=v[0:1], vdst=v[0], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[6], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[7], s[0]),
DS(DSOp.DS_STOREXCHG_2ADDR_STRIDE64_RTN_B64, addr=v[10], data0=v[6:7], data1=v[6:7], vdst=v[8:11], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][8], 0x11111111, "v8 should have old val1 low")
self.assertEqual(st.vgpr[0][9], 0x22222222, "v9 should have old val1 high")
self.assertEqual(st.vgpr[0][10], 0x11111111, "v10 should have old val2 low")
self.assertEqual(st.vgpr[0][11], 0x22222222, "v11 should have old val2 high")
class TestAtomicOrdering(unittest.TestCase):
"""Tests for atomic operation return values and ordering."""
def test_ds_add_rtn_sequence(self):
"""DS_ADD_RTN returns correct old values in sequence."""
instructions = [
v_mov_b32_e32(v[10], 0),
v_mov_b32_e32(v[0], 100),
DS(DSOp.DS_STORE_B32, addr=v[10], data0=v[0], vdst=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[1], 25),
DS(DSOp.DS_ADD_RTN_U32, addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_ADD_RTN_U32, addr=v[10], data0=v[1], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_B32, addr=v[10], vdst=v[4], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 100, "First add should return 100")
self.assertEqual(st.vgpr[0][3], 125, "Second add should return 125")
self.assertEqual(st.vgpr[0][4], 150, "Final value should be 150")
class TestDsPermute(unittest.TestCase):
"""Tests for DS_PERMUTE_B32 and DS_BPERMUTE_B32 instructions."""
def test_ds_permute_b32_identity(self):
"""DS_PERMUTE_B32 with identity permutation (lane 0 sends to lane 0)."""
# For simplicity, test with single lane
instructions = [
v_mov_b32_e32(v[0], 0), # addr = 0 (lane 0)
v_mov_b32_e32(v[1], 0xDEADBEEF), # data
ds_permute_b32(v[2], v[0], v[1]),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
# Lane 0 sends to lane 0, so lane 0 gets 0xDEADBEEF
self.assertEqual(st.vgpr[0][2], 0xDEADBEEF)
def test_ds_bpermute_b32_identity(self):
"""DS_BPERMUTE_B32 with identity permutation (each lane reads from itself)."""
instructions = [
v_mov_b32_e32(v[0], 0), # addr = 0 (read from lane 0)
v_mov_b32_e32(v[1], 0xCAFEBABE), # data in lane 0
ds_bpermute_b32(v[2], v[0], v[1]),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
# Lane 0 reads from lane 0's v[1]
self.assertEqual(st.vgpr[0][2], 0xCAFEBABE)
def test_ds_permute_b32_broadcast(self):
"""DS_PERMUTE_B32 broadcast - all lanes send to lane 0."""
# With 4 lanes, all sending to lane 0, highest lane wins
instructions = [
v_mov_b32_e32(v[0], 0), # All lanes send to addr 0 (lane 0)
v_mov_b32_e32(v[1], 0x11111111), # All lanes send same data
ds_permute_b32(v[2], v[0], v[1]),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=4)
# Lane 0 receives data (highest numbered active lane wins)
self.assertEqual(st.vgpr[0][2], 0x11111111)
def test_ds_bpermute_b32_xor_swap(self):
"""DS_BPERMUTE_B32 with XOR-1 pattern — each lane reads from lane^1.
This is the pattern used by warp_shfl_xor in flash attention for reduce_max/reduce_sum.
Each lane has a unique value (lane_id + 100), and reads from the adjacent lane.
"""
instructions = [
# v[0] = (lane_id ^ 1) * 4 (byte offset for bpermute)
v_xor_b32_e32(v[0], 1, v[255]),
v_lshlrev_b32_e32(v[0], 2, v[0]),
# v[1] = lane_id + 100 (unique per-lane value)
s_mov_b32(s[0], 100),
v_add_nc_u32_e32(v[1], s[0], v[255]),
# ds_bpermute: v[2] = v[1] from lane (lane_id ^ 1)
ds_bpermute_b32(vdst=v[2], addr=v[0], data0=v[1]),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=32)
for lane in range(32):
src_lane = lane ^ 1
expected = src_lane + 100
self.assertEqual(st.vgpr[lane][2], expected, f"lane {lane}: expected v[1] from lane {src_lane} = {expected}, got {st.vgpr[lane][2]}")
class TestDSSubDword(unittest.TestCase):
"""Tests for sub-dword DS operations (ds_store_b16, ds_store_b16_d16_hi)."""
def test_ds_store_b16_and_d16_hi(self):
"""DS_STORE_B16 stores low 16 bits, DS_STORE_B16_D16_HI stores high 16 bits to adjacent LDS half-words."""
instructions = [
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0xBEEF1234),
DS(DSOp.DS_STORE_B16, addr=v[0], data0=v[1], offset0=0),
DS(DSOp.DS_STORE_B16_D16_HI, addr=v[0], data0=v[1], offset0=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(vdst=v[2], addr=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xBEEF1234, "lo=0x1234 at byte 0, hi=0xBEEF at byte 2")
class TestDSLargeOffset(unittest.TestCase):
"""Tests for DS instructions with offsets > 255 (offset1 > 0).
The DS offset is a 16-bit value encoded as (offset1 << 8) | offset0.
These tests verify that offset1 is used correctly, not just offset0.
"""
def test_ds_store_load_b32_offset_256(self):
"""DS_STORE_B32/DS_LOAD_B32 with offset=256 (offset0=0, offset1=1)."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[0], s[0]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0, offset1=1), # offset = 256
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[1], offset0=0, offset1=1), # offset = 256
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0xDEADBEEF)
def test_ds_store_load_b32_offset_300(self):
"""DS_STORE_B32/DS_LOAD_B32 with offset=300 (offset0=44, offset1=1)."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[0], s[0]),
ds_store_b32(addr=v[10], data0=v[0], offset0=44, offset1=1), # offset = 300
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[1], offset0=44, offset1=1), # offset = 300
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0xCAFEBABE)
def test_ds_store_load_b64_offset_512(self):
"""DS_STORE_B64/DS_LOAD_B64 with offset=512 (offset0=0, offset1=2)."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
ds_store_b64(addr=v[10], data0=v[0:1], offset0=0, offset1=2), # offset = 512
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b64(addr=v[10], vdst=v[2:3], offset0=0, offset1=2), # offset = 512
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x11111111)
self.assertEqual(st.vgpr[0][3], 0x22222222)
def test_ds_large_offset_distinct_from_small(self):
"""Verify offset=256 and offset=0 address different LDS locations."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
# Store 0xAAAAAAAA at offset=0, 0xBBBBBBBB at offset=256
ds_store_b32(addr=v[10], data0=v[0], offset0=0, offset1=0), # offset = 0
ds_store_b32(addr=v[10], data0=v[1], offset0=0, offset1=1), # offset = 256
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# Read back both
ds_load_b32(addr=v[10], vdst=v[2], offset0=0, offset1=0), # offset = 0
ds_load_b32(addr=v[10], vdst=v[3], offset0=0, offset1=1), # offset = 256
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA, "offset=0 should read 0xAAAAAAAA")
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB, "offset=256 should read 0xBBBBBBBB")
def test_ds_store_load_b32_offset_448(self):
"""DS_STORE_B32/DS_LOAD_B32 with offset=448 (offset0=192, offset1=1) - matches matmul B tile."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x12345678),
v_mov_b32_e32(v[0], s[0]),
ds_store_b32(addr=v[10], data0=v[0], offset0=192, offset1=1), # offset = 448
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[1], offset0=192, offset1=1), # offset = 448
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0x12345678)
def test_ds_load_b64_offset_392(self):
"""DS_LOAD_B64 with offset=392 (offset0=136, offset1=1) - matches matmul B tile load."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAABBCCDD),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x11223344),
v_mov_b32_e32(v[1], s[0]),
ds_store_b64(addr=v[10], data0=v[0:1], offset0=136, offset1=1), # offset = 392
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b64(addr=v[10], vdst=v[2:3], offset0=136, offset1=1), # offset = 392
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAABBCCDD)
self.assertEqual(st.vgpr[0][3], 0x11223344)
if __name__ == '__main__':
unittest.main()
Reference in New Issue View Git Blame Copy Permalink