forked from IQ.Lvbs/IQ.Pilot
IQ.Pilot Prebuilt Release @ ab07000
This commit is contained in:
355
tinygrad_repo/test/amd/hw/test_scratch.py
Normal file
355
tinygrad_repo/test/amd/hw/test_scratch.py
Normal file
@@ -0,0 +1,355 @@
|
||||
"""Tests for SCRATCH instructions - scratch (private) memory operations.
|
||||
|
||||
Includes: scratch_load_*, scratch_store_*
|
||||
"""
|
||||
import unittest
|
||||
from test.amd.hw.helpers import *
|
||||
|
||||
class TestScratchStore(unittest.TestCase):
|
||||
"""Tests for SCRATCH store instructions."""
|
||||
|
||||
def test_scratch_store_b32_basic(self):
|
||||
"""SCRATCH_STORE_B32 stores 32-bit value to scratch memory."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
|
||||
s_mov_b32(s[4], 0xDEADBEEF),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
# Store via scratch
|
||||
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
# Load back via scratch
|
||||
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], v[3]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 0xDEADBEEF)
|
||||
|
||||
def test_scratch_store_b64_basic(self):
|
||||
"""SCRATCH_STORE_B64 stores 64-bit value to scratch memory."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
|
||||
s_mov_b32(s[4], 0xDEADBEEF),
|
||||
s_mov_b32(s[5], 0xCAFEBABE),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
v_mov_b32_e32(v[3], s[5]),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
scratch_store_b64(addr=v[0], data=v[2:3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
scratch_load_b64(addr=v[0], vdst=v[4:5], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], v[4]),
|
||||
v_mov_b32_e32(v[1], v[5]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 0xDEADBEEF)
|
||||
self.assertEqual(st.vgpr[0][1], 0xCAFEBABE)
|
||||
|
||||
def test_scratch_store_b8_basic(self):
|
||||
"""SCRATCH_STORE_B8 stores single byte to scratch memory."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
|
||||
# First store full word
|
||||
s_mov_b32(s[4], 0xDEADBEEF),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
# Store single byte
|
||||
v_mov_b32_e32(v[2], 0x42),
|
||||
scratch_store_b8(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
# Load back
|
||||
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], v[3]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# Only byte 0 should change from 0xEF to 0x42
|
||||
self.assertEqual(st.vgpr[0][0], 0xDEADBE42)
|
||||
|
||||
def test_scratch_store_b16_basic(self):
|
||||
"""SCRATCH_STORE_B16 stores 16-bit value to scratch memory."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
|
||||
s_mov_b32(s[4], 0xDEADBEEF),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
s_mov_b32(s[4], 0xCAFE),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
scratch_store_b16(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], v[3]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 0xDEADCAFE)
|
||||
|
||||
|
||||
class TestScratchLoad(unittest.TestCase):
|
||||
"""Tests for SCRATCH load instructions."""
|
||||
|
||||
def test_scratch_load_b96(self):
|
||||
"""SCRATCH_LOAD_B96 loads 96-bit value correctly."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
s_mov_b32(s[4], 0xAAAAAAAA),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
s_mov_b32(s[4], 0xBBBBBBBB),
|
||||
v_mov_b32_e32(v[3], s[4]),
|
||||
s_mov_b32(s[4], 0xCCCCCCCC),
|
||||
v_mov_b32_e32(v[4], s[4]),
|
||||
scratch_store_b96(addr=v[0], data=v[2:4], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
scratch_load_b96(addr=v[0], vdst=v[5:7], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], v[5]),
|
||||
v_mov_b32_e32(v[1], v[6]),
|
||||
v_mov_b32_e32(v[2], v[7]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA)
|
||||
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB)
|
||||
self.assertEqual(st.vgpr[0][2], 0xCCCCCCCC)
|
||||
|
||||
def test_scratch_load_b128(self):
|
||||
"""SCRATCH_LOAD_B128 loads 128-bit value correctly."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
s_mov_b32(s[4], 0xDEADBEEF),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
s_mov_b32(s[4], 0xCAFEBABE),
|
||||
v_mov_b32_e32(v[3], s[4]),
|
||||
s_mov_b32(s[4], 0x12345678),
|
||||
v_mov_b32_e32(v[4], s[4]),
|
||||
s_mov_b32(s[4], 0x9ABCDEF0),
|
||||
v_mov_b32_e32(v[5], s[4]),
|
||||
scratch_store_b128(addr=v[0], data=v[2:5], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
scratch_load_b128(addr=v[0], vdst=v[6:9], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], v[6]),
|
||||
v_mov_b32_e32(v[1], v[7]),
|
||||
v_mov_b32_e32(v[2], v[8]),
|
||||
v_mov_b32_e32(v[3], v[9]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 0xDEADBEEF)
|
||||
self.assertEqual(st.vgpr[0][1], 0xCAFEBABE)
|
||||
self.assertEqual(st.vgpr[0][2], 0x12345678)
|
||||
self.assertEqual(st.vgpr[0][3], 0x9ABCDEF0)
|
||||
|
||||
def test_scratch_load_u8(self):
|
||||
"""SCRATCH_LOAD_U8 loads unsigned byte with zero extension."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
s_mov_b32(s[4], 0xDEADBEAB),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
scratch_load_u8(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], v[3]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 0xAB)
|
||||
|
||||
def test_scratch_load_i8(self):
|
||||
"""SCRATCH_LOAD_I8 loads signed byte with sign extension."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
s_mov_b32(s[4], 0x80), # -128 as signed byte
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
scratch_store_b8(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
scratch_load_i8(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], v[3]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 0xFFFFFF80)
|
||||
|
||||
def test_scratch_load_u16(self):
|
||||
"""SCRATCH_LOAD_U16 loads unsigned 16-bit with zero extension."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
s_mov_b32(s[4], 0xDEADCAFE),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
scratch_load_u16(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], v[3]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 0xCAFE)
|
||||
|
||||
def test_scratch_load_i16(self):
|
||||
"""SCRATCH_LOAD_I16 loads signed 16-bit with sign extension."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
s_mov_b32(s[4], 0x8000), # -32768 as signed 16-bit
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
scratch_store_b16(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
scratch_load_i16(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], v[3]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 0xFFFF8000)
|
||||
|
||||
|
||||
class TestScratchSVE(unittest.TestCase):
|
||||
"""Tests for SCRATCH SVE (Scratch VGPR Enable) bit behavior."""
|
||||
|
||||
def test_scratch_sve_zero_ignores_vaddr(self):
|
||||
"""With SVE=0, VADDR should be ignored in address calculation."""
|
||||
TEST_OFFSET = 256
|
||||
# Store a marker value at offset 256 (where SVE=0 should go)
|
||||
# Then set v[0] to a non-zero value (100) and store via scratch with SVE=0
|
||||
# If SVE=0 is handled correctly, the VADDR (100) should be IGNORED,
|
||||
# and the store should go to offset 256, not 256+100=356
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
|
||||
# First, store 0xAAAAAAAA at offset 256 with v[0]=0
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
s_mov_b32(s[4], 0xAAAAAAAA),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=0),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
# Now set v[0] to 100 (non-zero) and store 0xBBBBBBBB with SVE=0
|
||||
# With SVE=0, v[0] should be IGNORED, so this should overwrite offset 256
|
||||
v_mov_b32_e32(v[0], 100),
|
||||
s_mov_b32(s[4], 0xBBBBBBBB),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=0),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
# Load back from offset 256 (with v[0]=0) - should get 0xBBBBBBBB
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=0),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], v[3]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# If SVE=0 works correctly, v[0] should be 0xBBBBBBBB (the second store overwrote the first)
|
||||
# If SVE=0 is wrong (VADDR used), v[0] would be 0xAAAAAAAA (stores went to different locations)
|
||||
self.assertEqual(st.vgpr[0][0], 0xBBBBBBBB, "SVE=0 should ignore VADDR, both stores should go to same location")
|
||||
|
||||
def test_scratch_sve_one_uses_vaddr(self):
|
||||
"""With SVE=1, VADDR should be used as offset in address calculation."""
|
||||
TEST_OFFSET = 256
|
||||
# Store at offset 256 with v[0]=0, then store at offset 256 with v[0]=100 and SVE=1
|
||||
# With SVE=1, the second store should go to 256+100=356, not 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
|
||||
# First, store 0xAAAAAAAA at offset 256 with v[0]=0
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
s_mov_b32(s[4], 0xAAAAAAAA),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=1),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
# Now set v[0] to 100 and store 0xBBBBBBBB with SVE=1
|
||||
# With SVE=1, v[0] IS used, so this should go to offset 256+100=356
|
||||
v_mov_b32_e32(v[0], 100),
|
||||
s_mov_b32(s[4], 0xBBBBBBBB),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=1),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
# Load back from offset 256 (with v[0]=0) - should still be 0xAAAAAAAA
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=1),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], v[3]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# If SVE=1 works correctly, v[0] should be 0xAAAAAAAA (stores went to different locations)
|
||||
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, "SVE=1 should use VADDR, stores should go to different locations")
|
||||
|
||||
|
||||
class TestScratchMultiLane(unittest.TestCase):
|
||||
"""Tests for SCRATCH operations with multiple lanes."""
|
||||
|
||||
def test_scratch_store_load_multi_lane(self):
|
||||
"""SCRATCH store/load works correctly with multiple lanes (private per-lane memory)."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
|
||||
# Each lane stores its lane ID
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
v_mov_b32_e32(v[2], v[255]), # v[255] has packed workitem IDs, low 10 bits = x
|
||||
v_and_b32_e32(v[2], 0x3FF, v[2]), # extract lane ID
|
||||
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
# Load back
|
||||
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
|
||||
v_mov_b32_e32(v[0], v[3]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=4)
|
||||
# Each lane should have loaded its own lane ID
|
||||
for lane in range(4):
|
||||
self.assertEqual(st.vgpr[lane][0], lane, f"Lane {lane} should have value {lane}")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user