mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 05:23:41 +08:00
Daniel Koepping
108 lines
4.8 KiB
Python
108 lines
4.8 KiB
Python
import unittest, math, itertools
|
|
from tinygrad import Tensor, Device, Context, dtypes
|
|
from tinygrad.dtype import DTYPES_DICT, DType, ConstType
|
|
from tinygrad.uop.ops import Ops, UOp, GroupOp
|
|
from tinygrad.codegen.decomp.op import threefry2x32
|
|
from test.helpers import full_rewrite
|
|
|
|
def _check_ast_count(desired_count:int, t:Tensor):
|
|
# NOTE: this has side effect because everything can be scheduled only once
|
|
schedule = t.schedule_linear()
|
|
asts = [s for s in schedule.src if s.src[0].op is Ops.SINK]
|
|
len(asts)
|
|
# NOT SUPPORTED ANYMORE
|
|
#assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
|
|
|
|
class TestWeakConstFolding(unittest.TestCase):
|
|
def test_weakint_math(self):
|
|
out = (UOp.const(2**40) + UOp.const(2**40)).simplify()
|
|
self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.weakint, 2**41))
|
|
|
|
def test_float_unaries(self):
|
|
for op in (Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL):
|
|
out = UOp.const(4.0).alu(op).simplify()
|
|
self.assertEqual((out.op, out.dtype), (Ops.CONST, dtypes.weakfloat))
|
|
|
|
def test_weakfloat_math(self):
|
|
out = (UOp.const(1.25) + UOp.const(2.5)).simplify()
|
|
self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.weakfloat, 3.75))
|
|
|
|
def test_nan_compare(self):
|
|
nan = UOp.const(math.nan)
|
|
self.assertTrue(nan.ne(nan).simplify().val)
|
|
|
|
def test_invalid_poison(self):
|
|
self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid)
|
|
|
|
class TestBitcastConstFolding(unittest.TestCase):
|
|
def test_out_of_range_source_value(self):
|
|
for val, src_dt, dst_dt, bits in ((3000000000, dtypes.int32, dtypes.uint32, 3000000000),
|
|
(70000, dtypes.int16, dtypes.uint16, 4464),
|
|
(-5, dtypes.uint32, dtypes.int32, -5)):
|
|
self.assertIs(UOp.const(val, src_dt).bitcast(dst_dt).simplify(), UOp.const(bits, dst_dt))
|
|
|
|
def test_scalar_bitcast(self):
|
|
def t(cases: dict[DType, ConstType]):
|
|
for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()):
|
|
if not math.isnan(from_v):
|
|
r = UOp.const(from_v, from_dt).bitcast(to_dt).simplify()
|
|
self.assertIs(r, UOp.const(to_v, to_dt), f"{from_dt} -> {to_dt} ({from_v} -> {to_v})")
|
|
|
|
t({dtypes.int8: 0, dtypes.uint8: 0, dtypes.bool: False})
|
|
t({dtypes.int8: 1, dtypes.uint8: 1, dtypes.bool: True})
|
|
|
|
t({dtypes.int8: -1, dtypes.uint8: 2**8-1})
|
|
t({dtypes.int16: -1, dtypes.uint16: 2**16-1, dtypes.float16: float('nan')})
|
|
t({dtypes.int32: -1, dtypes.uint32: 2**32-1, dtypes.float32: float('nan')})
|
|
t({dtypes.int64: -1, dtypes.uint64: 2**64-1, dtypes.float64: float('nan')})
|
|
|
|
t({dtypes.int8: -2**7, dtypes.uint8: 2**7})
|
|
t({dtypes.int16: -2**15, dtypes.uint16: 2**15})
|
|
t({dtypes.int32: -2**31, dtypes.uint32: 2**31})
|
|
t({dtypes.int64: -2**63, dtypes.uint64: 2**63})
|
|
|
|
t({dtypes.int16: 13496, dtypes.uint16: 13496, dtypes.float16: 0.294921875})
|
|
t({dtypes.int32: 1050081145, dtypes.uint32: 1050081145, dtypes.float32: 0.29485681653022766})
|
|
t({dtypes.int64: 4598983288165178391, dtypes.uint64: 4598983288165178391, dtypes.float64: 0.29485681936461233})
|
|
|
|
def test_vec_bitcast(self):
|
|
with Context(SPEC=0):
|
|
result = full_rewrite(UOp.const((-1, -2**31, 75), dtypes.int32).bitcast(dtypes.uint32).sink())
|
|
expected = full_rewrite(UOp.const((2**32-1, 2**31, 75), dtypes.uint32).sink())
|
|
self.assertEqual(result.src, expected.src)
|
|
|
|
class TestMovedConstFolding(unittest.TestCase):
|
|
def test_contiguous_deviceless_const(self):
|
|
t = Tensor(UOp.const(2.0, dtypes.float)).contiguous()
|
|
self.assertIs(t.uop, UOp.const(2.0, dtypes.float))
|
|
self.assertIsNone(t.uop.device)
|
|
|
|
def test_add_shrunk_zero(self):
|
|
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(6).shrink(((1, 5),)))
|
|
|
|
def test_add_padded_zero(self):
|
|
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(2).pad(((1, 1),)))
|
|
|
|
def test_mul_shrunk_one(self):
|
|
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.ones(6).shrink(((1, 5),)))
|
|
|
|
def test_add_padded_one(self):
|
|
_check_ast_count(1, Tensor([1.0, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),)))
|
|
|
|
class TestReduceOpsConstFolding(unittest.TestCase):
|
|
def test_sum_output_dtype(self):
|
|
# sum output dtype can be different from input
|
|
for dt in DTYPES_DICT.values():
|
|
if dt in Device[Device.DEFAULT].renderer.supported_dtypes():
|
|
t = Tensor.ones(16, dtype=dt).reshape(4, 4)
|
|
assert t.sum().dtype == t.contiguous().sum().dtype
|
|
|
|
class TestThreefryConstFolding(unittest.TestCase):
|
|
def test_threefry(self):
|
|
# THREEFRY(const,const) folds to a const once decomposed
|
|
x = threefry2x32(UOp.const(5, dtypes.uint64), UOp.const(10, dtypes.uint64)).simplify()
|
|
self.assertEqual([u.op for u in x.toposort() if u.op in GroupOp.ALU], [])
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|