mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 12:56:07 +00:00
prevent UB in long decomp and more emulated tests (#14447)
This commit is contained in:
@@ -771,9 +771,6 @@ jobs:
|
||||
run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
|
||||
- name: Run TRANSCENDENTAL math
|
||||
run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
|
||||
- name: Test dtype with emulated long
|
||||
if: matrix.backend != 'lvp'
|
||||
run: EMULATED_DTYPES=long python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
|
||||
@@ -1004,7 +1004,7 @@ def _get_runner(inst_bytes: bytes):
|
||||
canonical_name = f"{_op_name(inst).lower()}_{base.to_bytes(size, 'little').hex()}"
|
||||
sink = sink.replace(arg=KernelInfo(name=canonical_name)).rtag(1)
|
||||
|
||||
with Context(NOOPT=1, CHECK_OOB=0, TUPLE_ORDER=0):
|
||||
with Context(NOOPT=1, CHECK_OOB=0, TUPLE_ORDER=0, EMULATED_DTYPES=""):
|
||||
runner = get_runner('CPU', sink)
|
||||
_canonical_runner_cache.append((base, mask, size, runner))
|
||||
return runner, True
|
||||
|
||||
+27
-7
@@ -1,4 +1,4 @@
|
||||
import unittest, math
|
||||
import contextlib, unittest, math
|
||||
import numpy as np
|
||||
import torch
|
||||
from typing import Any, List
|
||||
@@ -7,7 +7,8 @@ from tinygrad.helpers import getenv, DEBUG, CI, EMULATED_DTYPES
|
||||
from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8, _to_np_dtype, _to_torch_dtype, truncate
|
||||
from tinygrad.renderer.ptx import PTXRenderer
|
||||
from tinygrad.renderer.nir import NIRRenderer
|
||||
from tinygrad import Device, Tensor, dtypes
|
||||
from tinygrad import Context, Device, Tensor, dtypes
|
||||
from tinygrad.uop import Ops
|
||||
from hypothesis import given, settings, strategies as strat
|
||||
from test.helpers import rand_for_dtype
|
||||
from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX
|
||||
@@ -336,18 +337,37 @@ class TestUint16DType(TestDType):
|
||||
class TestInt32DType(TestDType): DTYPE = dtypes.int32
|
||||
class TestUint32DType(TestDType): DTYPE = dtypes.uint32
|
||||
|
||||
class TestInt64DType(TestDType):
|
||||
DTYPE = dtypes.int64
|
||||
class TestInt64DType(TestDType): DTYPE = dtypes.int64
|
||||
|
||||
@unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "long decomp requires bitshift")
|
||||
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
|
||||
class TestEmulatedInt64DType(TestInt64DType):
|
||||
@classmethod
|
||||
def setUpClass(cls): cls.DATA = rand_for_dtype(cls.DTYPE, 10)
|
||||
def setUpClass(cls):
|
||||
cls.stack = contextlib.ExitStack()
|
||||
cls.stack.enter_context(Context(EMULATED_DTYPES="long"))
|
||||
cls.DATA = rand_for_dtype(cls.DTYPE, 10)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls): cls.stack.close()
|
||||
|
||||
class TestUint64DType(TestDType):
|
||||
@classmethod
|
||||
def setUpClass(cls): cls.DATA = rand_for_dtype(cls.DTYPE, 10)
|
||||
DTYPE = dtypes.uint64
|
||||
def test_uint64_load(self):
|
||||
assert Tensor(2**64 - 1, dtype=dtypes.uint64).numpy() == 2**64 - 1
|
||||
|
||||
@unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "long decomp requires bitshift")
|
||||
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
|
||||
class TestEmulatedUInt64DType(TestUint64DType):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.stack = contextlib.ExitStack()
|
||||
cls.stack.enter_context(Context(EMULATED_DTYPES="long"))
|
||||
cls.DATA = rand_for_dtype(cls.DTYPE, 10)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls): cls.stack.close()
|
||||
|
||||
class TestBoolDType(TestDType): DTYPE = dtypes.bool
|
||||
|
||||
class TestBFloat16Type(TestDType): DTYPE = dtypes.bfloat16
|
||||
|
||||
+30
-1
@@ -1,5 +1,5 @@
|
||||
import unittest, operator, math
|
||||
from tinygrad import Tensor, dtypes, Device
|
||||
from tinygrad import Context, Tensor, dtypes, Device
|
||||
from tinygrad.dtype import DType, truncate
|
||||
from tinygrad.helpers import CI, getenv
|
||||
from tinygrad.tensor import _to_np_dtype
|
||||
@@ -7,6 +7,7 @@ from tinygrad.device import is_dtype_supported
|
||||
from tinygrad.runtime.ops_python import from_storage_scalar
|
||||
from tinygrad.renderer.ptx import PTXRenderer
|
||||
from tinygrad.renderer.nir import NIRRenderer
|
||||
from tinygrad.uop import Ops
|
||||
import numpy as np
|
||||
import pytest
|
||||
from hypothesis import assume, given, strategies as strat, settings
|
||||
@@ -165,9 +166,16 @@ class TestDTypeALU(unittest.TestCase):
|
||||
@given(ht.uint32, ht.uint32, strat.sampled_from(integer_binary_operations))
|
||||
def test_uint32(self, a, b, op): universal_test(a, b, dtypes.uint32, op)
|
||||
|
||||
@unittest.skipUnless(is_dtype_supported(dtypes.uint64), f"no uint64 on {Device.DEFAULT}")
|
||||
@given(ht.uint64, ht.uint64, strat.sampled_from(integer_binary_operations))
|
||||
def test_uint64(self, a, b, op): universal_test(a, b, dtypes.uint64, op)
|
||||
|
||||
@unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "long decomp requires bitshift")
|
||||
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
|
||||
@given(ht.uint64, ht.uint64, strat.sampled_from(integer_binary_operations))
|
||||
@Context(EMULATED_DTYPES="long")
|
||||
def test_emulated_uint64(self, a, b, op): universal_test(a, b, dtypes.uint64, op)
|
||||
|
||||
@given(ht.int8, ht.int8, strat.sampled_from(integer_binary_operations))
|
||||
def test_int8(self, a, b, op): universal_test(a, b, dtypes.int8, op)
|
||||
|
||||
@@ -177,9 +185,16 @@ class TestDTypeALU(unittest.TestCase):
|
||||
@given(ht.int32, ht.int32, strat.sampled_from(integer_binary_operations))
|
||||
def test_int32(self, a, b, op): universal_test(a, b, dtypes.int32, op)
|
||||
|
||||
@unittest.skipUnless(is_dtype_supported(dtypes.int64), f"no int64 on {Device.DEFAULT}")
|
||||
@given(ht.int64, ht.int64, strat.sampled_from(integer_binary_operations))
|
||||
def test_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op)
|
||||
|
||||
@unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "long decomp requires bitshift")
|
||||
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
|
||||
@given(ht.int64, ht.int64, strat.sampled_from(integer_binary_operations))
|
||||
@Context(EMULATED_DTYPES="long")
|
||||
def test_emulated_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op)
|
||||
|
||||
@given(ht.uint8, strat.sampled_from(integer_unary_operations))
|
||||
def test_uint8_unary(self, a, op): universal_test_unary(a, dtypes.uint8, op)
|
||||
|
||||
@@ -191,9 +206,16 @@ class TestDTypeALU(unittest.TestCase):
|
||||
@given(ht.uint32, strat.sampled_from(integer_unary_operations))
|
||||
def test_uint32_unary(self, a, op): universal_test_unary(a, dtypes.uint32, op)
|
||||
|
||||
@unittest.skipUnless(is_dtype_supported(dtypes.uint64), f"no uint64 on {Device.DEFAULT}")
|
||||
@given(ht.uint64, strat.sampled_from(integer_unary_operations))
|
||||
def test_uint64_unary(self, a, op): universal_test_unary(a, dtypes.uint64, op)
|
||||
|
||||
@unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "long decomp requires bitshift")
|
||||
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
|
||||
@given(ht.uint64, strat.sampled_from(integer_unary_operations))
|
||||
@Context(EMULATED_DTYPES="long")
|
||||
def test_emulated_uint64_unary(self, a, op): universal_test_unary(a, dtypes.uint64, op)
|
||||
|
||||
@given(ht.int8, strat.sampled_from(integer_unary_operations))
|
||||
def test_int8_unary(self, a, op): universal_test_unary(a, dtypes.int8, op)
|
||||
|
||||
@@ -203,9 +225,16 @@ class TestDTypeALU(unittest.TestCase):
|
||||
@given(ht.int32, strat.sampled_from(integer_unary_operations))
|
||||
def test_int32_unary(self, a, op): universal_test_unary(a, dtypes.int32, op)
|
||||
|
||||
@unittest.skipUnless(is_dtype_supported(dtypes.int64), f"no int64 on {Device.DEFAULT}")
|
||||
@given(ht.int64, strat.sampled_from(integer_unary_operations))
|
||||
def test_int64_unary(self, a, op): universal_test_unary(a, dtypes.int64, op)
|
||||
|
||||
@unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "long decomp requires bitshift")
|
||||
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
|
||||
@given(ht.int64, strat.sampled_from(integer_unary_operations))
|
||||
@Context(EMULATED_DTYPES="long")
|
||||
def test_emulated_int64_unary(self, a, op): universal_test_unary(a, dtypes.int64, op)
|
||||
|
||||
@given(ht.bool, ht.bool, strat.sampled_from(((operator.add, operator.add), (operator.mul, operator.mul))))
|
||||
def test_bool(self, a, b, op): universal_test(a, b, dtypes.bool, op)
|
||||
|
||||
|
||||
@@ -353,8 +353,9 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
|
||||
case Ops.IDIV | Ops.MOD:
|
||||
# TAOCP Algorithm 4.3.1D could be faster here, but must be parameterized over the width of b
|
||||
if dt == dtypes.int:
|
||||
a0, a1 = (a_neg:=a1 < zero).where((n:=l2i(Ops.NEG, dt, a0, a1))[0], a0).bitcast(dtypes.uint), a_neg.where(n[1], a1).bitcast(dtypes.uint)
|
||||
b0, b1 = (b_neg:=b1 < zero).where((n:=l2i(Ops.NEG, dt, b0, b1))[0], b0).bitcast(dtypes.uint), b_neg.where(n[1], b1).bitcast(dtypes.uint)
|
||||
ua0, ua1, ub0, ub1 = a0.bitcast(dtypes.uint), a1.bitcast(dtypes.uint), b0.bitcast(dtypes.uint), b1.bitcast(dtypes.uint)
|
||||
a0, a1 = (a_neg:=a1 < zero).where((n:=l2i(Ops.NEG, dtypes.uint, ua0, ua1))[0], ua0), a_neg.where(n[1], ua1)
|
||||
b0, b1 = (b_neg:=b1 < zero).where((n:=l2i(Ops.NEG, dtypes.uint, ub0, ub1))[0], ub0), b_neg.where(n[1], ub1)
|
||||
q, r = (z:=UOp.const(dtypes.uint, 0), z), (z, z)
|
||||
for i in range(63, -1, -1):
|
||||
r = l2i(Ops.SHL, dtypes.uint, *r, UOp.const(dtypes.uint, 1), z)
|
||||
@@ -364,8 +365,9 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
|
||||
q = ((q[0] | cond.cast(dtypes.uint) << (i % 32), q[1]) if i < 32 else (q[0], q[1] | cond.cast(dtypes.uint) << (i % 32)))
|
||||
r = l2i(Ops.WHERE, dtypes.uint, cond, *diff, *r)
|
||||
if dt == dtypes.int:
|
||||
nq, nr = l2i(Ops.NEG, dt, q0:=q[0].bitcast(dt), q1:=q[1].bitcast(dt)), l2i(Ops.NEG, dt, r0:=r[0].bitcast(dt), r1:=r[1].bitcast(dt))
|
||||
return (a_neg.where(nr[0], r0), a_neg.where(nr[1], r1)) if op == Ops.MOD else ((a_neg^b_neg).where(nq[0], q0), (a_neg^b_neg).where(nq[1], q1))
|
||||
(nq0, nq1), (nr0, nr1) = l2i(Ops.BITCAST, dt, *l2i(Ops.NEG, dtypes.uint, *q)), l2i(Ops.BITCAST, dt, *l2i(Ops.NEG, dtypes.uint, *r))
|
||||
(q0, q1), (r0, r1) = l2i(Ops.BITCAST, dt, *q), l2i(Ops.BITCAST, dt, *r)
|
||||
return (a_neg.where(nr0, r0), a_neg.where(nr1, r1)) if op == Ops.MOD else ((a_neg^b_neg).where(nq0, q0), (a_neg^b_neg).where(nq1, q1))
|
||||
return (r[0].bitcast(dt), r[1].bitcast(dt)) if op == Ops.MOD else (q[0].bitcast(dt), q[1].bitcast(dt))
|
||||
case Ops.CMPLT: return (a1 < b1) | ((a1.eq(b1)) & (a0.bitcast(dtypes.uint) < b0.bitcast(dtypes.uint)))
|
||||
case Ops.CMPEQ: return a0.eq(b0) & a1.eq(b1)
|
||||
|
||||
Reference in New Issue
Block a user