mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-10-28 09:03:46 +08:00
Daniel Koepping
296 lines
13 KiB
Python
296 lines
13 KiB
Python
import unittest
|
|
from tinygrad import dtypes
|
|
from tinygrad.dtype import AddrSpace
|
|
from tinygrad.helpers import Context
|
|
from tinygrad.uop.ops import Ops, UOp, AxisType
|
|
from tinygrad.uop.validate import validate_index_with_z3
|
|
from test.helpers import to_uops_list
|
|
|
|
def Variable(name, nmin, nmax): return UOp.variable(name, nmin, nmax)
|
|
|
|
class TestValidateOOB(unittest.TestCase):
|
|
"""Test z3 validation of index bounds for different ALU ops and patterns."""
|
|
|
|
# basic index patterns
|
|
def test_const_index(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 16)
|
|
to_uops_list([buf.index(UOp.const(0)).load()]) # valid
|
|
to_uops_list([buf.index(UOp.const(15)).load()]) # valid (last element)
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(UOp.const(16)).load()]) # off by one
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(UOp.const(42)).load()]) # way out
|
|
|
|
def test_variable_index(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 16)
|
|
to_uops_list([buf.index(Variable("i", 0, 15)).load()]) # valid
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(Variable("i", 0, 20)).load()]) # oob
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(Variable("i", -5, 10)).load()]) # negative
|
|
|
|
def test_range_with_mask(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 16)
|
|
r = UOp.range(42, 0, AxisType.GLOBAL)
|
|
to_uops_list([buf.index(r.valid(r < 16)).load()]) # valid
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(r.valid(r < 17)).load()]) # oob
|
|
|
|
def test_variable_with_mask(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 16)
|
|
v = Variable("v", -5, 80)
|
|
to_uops_list([buf.index(v.valid((v >= 0) & (v < 16))).load()]) # valid
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(v.valid(v < 20)).load()]) # negative not masked
|
|
|
|
def test_gated_store(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 16)
|
|
v = Variable("v", 0, 20)
|
|
to_uops_list([buf.index(v.valid(v < 16)).store(0)]) # valid
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(v.valid(v < 20)).store(0)]) # oob
|
|
|
|
# ALU ops in index
|
|
def test_floordiv(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 16)
|
|
to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2).load()]) # 0..15 valid
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(UOp.range(34, 0, AxisType.GLOBAL) // 2).load()]) # 0..16 oob
|
|
|
|
def test_mod(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 16)
|
|
r = UOp.range(100, 0, AxisType.GLOBAL)
|
|
to_uops_list([buf.index(r % 16).load()]) # 0..15 valid
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(r % 20).load()]) # 0..19 oob
|
|
|
|
def test_shr(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 16)
|
|
to_uops_list([buf.index(UOp.range(64, 0, AxisType.GLOBAL) >> 2).load()]) # 0..15 valid
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(UOp.range(128, 0, AxisType.GLOBAL) >> 2).load()]) # 0..31 oob
|
|
|
|
def test_shl(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 64)
|
|
r = UOp.range(8, 0, AxisType.GLOBAL)
|
|
to_uops_list([buf.index(r << 2).load()]) # 0..28 valid
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(r << 4).load()]) # 0..112 oob
|
|
|
|
def test_and(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 16)
|
|
r = UOp.range(100, 0, AxisType.GLOBAL)
|
|
to_uops_list([buf.index(r & 15).load()]) # 0..15 valid
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(r & 31).load()]) # 0..31 oob
|
|
# align masks round down to a multiple of 2^k
|
|
to_uops_list([buf.index((r & -4).valid(r < 16)).load()]) # 0..12 valid
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(r & -2).load()]) # 0..100 oob
|
|
# any mask is modeled now, and r&21 can hit 16..21
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(r & 21).load()])
|
|
|
|
def test_or(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 512)
|
|
r = UOp.range(256, 0, AxisType.GLOBAL)
|
|
# disjoint bits: a byte | high bits << 8
|
|
to_uops_list([buf.index((r & 255) | ((r & 1) << 8)).load()]) # 0..511 valid
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index((r & 255) | ((r & 3) << 8)).load()]) # 0..1023 oob
|
|
# overlapping bits
|
|
to_uops_list([buf.index(r | (r << 1)).load()]) # 0..511 valid
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(r | (r << 3)).load()]) # 0..2047 oob
|
|
|
|
def test_or_negative(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 16)
|
|
v, u = Variable("v", -16, -1), Variable("u", 0, 15)
|
|
# or with a negative side is negative, the mask is never true
|
|
w = v | u
|
|
to_uops_list([buf.index(w.valid(w >= 0)).load()])
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(w).load()])
|
|
|
|
def test_or_in_mask(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 32)
|
|
r = UOp.range(64, 0, AxisType.GLOBAL)
|
|
to_uops_list([buf.index(r.valid((r < 8) | ((r >= 16) & (r < 32)))).load()]) # 0..7,16..31 valid
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(r.valid((r < 8) | (r >= 56))).load()]) # 8..55 unmasked
|
|
|
|
def test_max(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 16)
|
|
to_uops_list([buf.index(Variable("v", -10, 15).maximum(0)).load()]) # 0..15 valid
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(Variable("v2", -10, 20).maximum(0)).load()]) # 0..20 oob
|
|
|
|
def test_xor_in_mask(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 16)
|
|
r = UOp.range(32, 0, AxisType.GLOBAL)
|
|
to_uops_list([buf.index(r.valid((r < 8) ^ ((r >= 8) & (r < 16)))).load()]) # 0..15 valid
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(r.valid((r < 10) ^ (r >= 20))).load()]) # 0..9,20..31 oob
|
|
|
|
# cast patterns
|
|
def test_float_cast_in_index(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 16)
|
|
r = UOp.range(20, 0)
|
|
i = (r.cast(dtypes.float) * 0.68).trunc().cast(dtypes.int)
|
|
to_uops_list([buf.index(i.valid((i >= 0) & (i < 16))).load()])
|
|
# a float entirely out of the int range has no value, not an empty one
|
|
f = UOp.variable("f", 3e9, 4e9, dtypes.float32).cast(dtypes.int)
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(f).load()])
|
|
|
|
def test_float_cast_in_mask(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 1)
|
|
r = UOp.range(20, 0)
|
|
unknown = r.cast(dtypes.float).cast(dtypes.bool) # a bool from a float is unconstrained
|
|
to_uops_list([buf.index(r.valid((r < 1) & unknown)).load()])
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(r.valid(unknown)).load()])
|
|
|
|
def test_bitcast_in_index(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 16)
|
|
r = UOp.range(16, 0)
|
|
# the WEBGPU shift: int -> uint, shift, back to int
|
|
i = (r.cast(dtypes.int).bitcast(dtypes.uint) << UOp.const(1).cast(dtypes.uint)).bitcast(dtypes.int)
|
|
to_uops_list([buf.index(i.valid(i < 16)).load()])
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(i).load()]) # 0..30 oob
|
|
# a negative char reads as a large uchar
|
|
c = Variable("c", -128, -113).cast(dtypes.char)
|
|
to_uops_list([UOp.param(1, dtypes.int, 144).index(c.bitcast(dtypes.uchar).cast(dtypes.int)).load()]) # 128..143 valid
|
|
# the bits of a float are any int
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf.index(r.cast(dtypes.float).bitcast(dtypes.int)).load()])
|
|
|
|
def test_bool_cast_in_mask(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf = UOp.param(0, dtypes.int, 1)
|
|
r = UOp.range(20, 0)
|
|
to_uops_list([buf.index(r.valid(r.cast(dtypes.bool).logical_not())).load()]) # only r=0 valid
|
|
|
|
# load result as index/mask
|
|
def test_load_as_index(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf0 = UOp.param(0, dtypes.int, 16)
|
|
buf1 = UOp.param(1, dtypes.int, 64)
|
|
r = UOp.range(42, 0, AxisType.GLOBAL)
|
|
ld0 = buf0.index(r.valid(r < 8)).load().cast(dtypes.weakint)
|
|
to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32))).load()]) # valid
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 64))).load()]) # oob
|
|
|
|
def test_load_from_shrink_as_index(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf0 = UOp.param(0, dtypes.int, 16)
|
|
buf1 = UOp.param(1, dtypes.int, 64)
|
|
shrink = UOp(Ops.SHRINK, src=(buf0, UOp.const(0, dtypes.int), UOp.const(4)))
|
|
ld0 = shrink.load().index(0)
|
|
to_uops_list([buf1.index(ld0.valid((ld0 >= 0) & (ld0 < 64))).load()])
|
|
|
|
def test_load_bool_as_mask(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
buf_bool = UOp.param(0, dtypes.bool, 16)
|
|
buf_int = UOp.param(1, dtypes.int, 8)
|
|
gidx = UOp(Ops.SPECIAL, src=(UOp.const(16),), arg="gidx0")
|
|
ld_bool = buf_bool.index(gidx).load()
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([buf_int.index(gidx.valid(ld_bool)).load()]) # gidx 0..15, buf_int size 8
|
|
|
|
# local memory
|
|
def test_gated_local(self):
|
|
with Context(CHECK_OOB=1, SPEC=2):
|
|
gbuf = UOp.param(0, dtypes.uint, 400)
|
|
sbuf = UOp.placeholder((8,), dtypes.uint, slot=0, addrspace=AddrSpace.LOCAL)
|
|
gidx = UOp(Ops.SPECIAL, src=(UOp.const(416),), arg="gidx0")
|
|
lidx = UOp(Ops.SPECIAL, src=(UOp.const(10),), arg="lidx0")
|
|
store = sbuf.index(lidx.valid(lidx < 8)).store(UOp.const(1))
|
|
load = sbuf.after(store).index(lidx.valid(lidx < 8)).load()
|
|
to_uops_list([gbuf.index(gidx.valid(gidx < 400)).store(load)]) # valid: local store and load gated to 8, global store gated to 400
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([gbuf.index(gidx.valid(gidx < 400)).store(sbuf.after(store).index(lidx).load())]) # lidx 0..9 into 8
|
|
with self.assertRaises(RuntimeError):
|
|
to_uops_list([gbuf.index(gidx).store(load)]) # gidx 0..415 into 400
|
|
|
|
class TestShiftBounds(unittest.TestCase):
|
|
def _check_max_index(self, idx, maximum):
|
|
self.assertTrue(validate_index_with_z3(maximum+1, idx, UOp.const(True)))
|
|
self.assertFalse(validate_index_with_z3(maximum, idx, UOp.const(True)))
|
|
|
|
def test_constant_shifts(self):
|
|
for count in (0, 1, 31, 64, 129):
|
|
with self.subTest(count=count):
|
|
self._check_max_index(Variable("a", 0, 15) << count, 15 << count)
|
|
self._check_max_index(Variable("b", 0, 15 << count) >> count, 15)
|
|
|
|
def test_symbolic_shifts(self):
|
|
n = Variable("n", 0, 3)
|
|
for dtype in (*dtypes.ints, dtypes.weakint):
|
|
with self.subTest(dtype=str(dtype)):
|
|
self._check_max_index(UOp.const(127, dtype) >> n, 127)
|
|
self._check_max_index(UOp.const(1, dtype) << n, 8)
|
|
|
|
def test_nonzero_shift_minimum(self):
|
|
n = Variable("n", 3, 6)
|
|
self._check_max_index(UOp.const(511) >> n, 63)
|
|
self._check_max_index(UOp.const(1) << n, 64)
|
|
self._check_max_index((UOp.const(-129) >> n) + 17, 14)
|
|
|
|
def test_conditional_shift_count(self):
|
|
n = Variable("n", 0, 3)
|
|
shift = n.eq(0).where(0, n.eq(1).where(3, n.eq(2).where(6, 9)))
|
|
self._check_max_index(UOp.const(65535) >> shift, 65535)
|
|
self._check_max_index(UOp.const(1) << shift, 512)
|
|
|
|
def test_signed_right_shift(self):
|
|
n = Variable("n", 0, 3)
|
|
for dtype in (*dtypes.sints, dtypes.weakint):
|
|
with self.subTest(dtype=str(dtype)):
|
|
# -9 >> n is -9, -5, -3, -2: arithmetic shifts round down, not toward zero.
|
|
self._check_max_index((UOp.const(-9, dtype) >> n) + 9, 7)
|
|
|
|
def test_lowered_shift_count(self):
|
|
n = UOp.variable("n", 0, 63, dtypes.uint32)
|
|
self._check_max_index(UOp.const(65535, dtypes.uint64).alu(Ops.SHR, n), 65535)
|
|
|
|
def test_weak_shifts_beyond_64_bits(self):
|
|
n = Variable("n", 64, 70)
|
|
self._check_max_index(UOp.const(2**70) >> n, 64)
|
|
self._check_max_index((UOp.const(1) << n) // 2**64, 64)
|
|
|
|
def test_loaded_shift(self):
|
|
shift = UOp.param(0, dtypes.uint32, 1).index(0).load() & 31
|
|
self._check_max_index(UOp.const(256, dtypes.uint32) >> shift, 256)
|
|
|
|
def test_negative_shift_is_not_assumed_safe(self):
|
|
n = Variable("n", -1, 3)
|
|
for op in (Ops.SHL, Ops.SHR):
|
|
for count in (UOp.const(-1), n):
|
|
idx = UOp.const(1, dtypes.int32).alu(op, count)
|
|
self.assertFalse(validate_index_with_z3(9, idx, UOp.const(True)))
|
|
self.assertTrue(validate_index_with_z3(9, idx, count >= 0))
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|