delete stale tests (#17596)

This commit is contained in:
chenyu
2026-08-19 11:18:05 -04:00
committed by GitHub
parent 9550378704
commit bd6e70ac15
3 changed files with 2 additions and 345 deletions
+1 -109
View File
@@ -1,39 +1,10 @@
import unittest, itertools, math import unittest, itertools, math
from tinygrad import Tensor, dtypes, Context from tinygrad import dtypes, Context
from tinygrad.dtype import DType, ConstType from tinygrad.dtype import DType, ConstType
from tinygrad.uop.ops import Ops, UOp from tinygrad.uop.ops import Ops, UOp
from test.helpers import full_rewrite from test.helpers import full_rewrite
import numpy as np import numpy as np
def _check_ast_count(desired_count:int, t:Tensor):
# NOTE: this has side effect because everything can be scheduled only once
linear = t.schedule_linear()
asts = [s for s in linear.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 TestUnaryOpsConstFolding(unittest.TestCase):
def test_all_consts_ops(self):
_check_ast_count(0, Tensor.ones(4).exp())
_check_ast_count(0, Tensor.ones(4).sqrt())
_check_ast_count(0, Tensor.ones(4) + Tensor.ones(4))
_check_ast_count(0, Tensor.ones(4) / Tensor.ones(4))
def test_cast(self):
_check_ast_count(0, Tensor.ones(4).cast(dtypes.int16))
_check_ast_count(0, Tensor.full(4, fill_value=-1).cast(dtypes.uint16))
def test_neg_folding(self):
_check_ast_count(0, Tensor([1, 2, 3]).mul(-1).neg())
_check_ast_count(0, Tensor([1, 2, 3]).neg().mul(-1))
_check_ast_count(0, Tensor([1, 2, 3]).neg().neg())
def test_neg_realized_no_fold(self):
x = Tensor.randn(32, 32)
x = x.clip(0, 1).realize()
_check_ast_count(1, x.neg())
class TestWeakConstFolding(unittest.TestCase): class TestWeakConstFolding(unittest.TestCase):
def test_weakint_math(self): def test_weakint_math(self):
out = (UOp.const(2**40) + UOp.const(2**40)).simplify() out = (UOp.const(2**40) + UOp.const(2**40)).simplify()
@@ -51,70 +22,6 @@ class TestWeakConstFolding(unittest.TestCase):
def test_invalid_poison(self): def test_invalid_poison(self):
self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid) self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid)
class TestBinaryOpsConstFolding(unittest.TestCase):
def test_add_literal_zero(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) + 0)
def test_add_tensor_zero(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(4))
def test_literal_zero_add(self):
_check_ast_count(0, 0 + Tensor([1.0, 2, 3, 4]))
def test_tensor_zero_add(self):
_check_ast_count(0, Tensor.zeros(4) + Tensor([1.0, 2, 3, 4]))
def test_sub_literal_zero(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) - 0)
def test_sub_tensor_zero(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) - Tensor.zeros(4))
def test_mul_literal_zero(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) * 0)
def test_mul_tensor_zero(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.zeros(4))
def test_literal_zero_mul(self):
_check_ast_count(0, 0 * Tensor([1.0, 2, 3, 4]) * 0)
def test_tensor_zero_mul(self):
_check_ast_count(0, Tensor.zeros(4) * Tensor([1.0, 2, 3, 4]))
def test_mul_literal_one(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) * 1)
def test_mul_tensor_one(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.ones(4))
def test_literal_one_mul(self):
_check_ast_count(0, 1 * Tensor([1.0, 2, 3, 4]))
def test_tensor_one_mul(self):
_check_ast_count(0, Tensor.ones(4) * Tensor([1.0, 2, 3, 4]))
def test_bool_tensor_mul_bool(self):
_check_ast_count(0, Tensor([True, False]) * True)
_check_ast_count(0, Tensor([True, False]) * False)
def test_bool_mul_bool_tensor(self):
_check_ast_count(0, True * Tensor([True, False]))
_check_ast_count(0, False * Tensor([True, False]))
def test_div_literal_one(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) / 1)
def test_div_tensor_one(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) / Tensor.ones(4))
def test_floordiv_literal_one(self):
_check_ast_count(0, Tensor([1, 2, 3, 4]) // 1)
def test_floordiv_tensor_one(self):
_check_ast_count(0, Tensor([1, 2, 3, 4]) // Tensor.ones(4, dtype=dtypes.int32))
def test_pow_literal_zero(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** 0)
def test_pow_tensor_zero(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** Tensor.zeros(4))
def test_pow_literal_one(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** 1)
def test_pow_tensor_one(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** Tensor.ones(4))
def test_literal_one_pow(self):
_check_ast_count(0, 1 ** Tensor([1.0, 2, 3, 4]))
def test_tensor_one_pow(self):
_check_ast_count(0, Tensor.ones(4) ** Tensor([1.0, 2, 3, 4]))
class TestBitcastConstFolding(unittest.TestCase): class TestBitcastConstFolding(unittest.TestCase):
def test_scalar_bitcast(self): def test_scalar_bitcast(self):
def t(cases: dict[DType, ConstType]): def t(cases: dict[DType, ConstType]):
@@ -148,20 +55,5 @@ class TestBitcastConstFolding(unittest.TestCase):
expected = full_rewrite(UOp.const((2**32-1, 2**31, 75), dtypes.uint32).sink()) expected = full_rewrite(UOp.const((2**32-1, 2**31, 75), dtypes.uint32).sink())
self.assertEqual(result.src, expected.src) self.assertEqual(result.src, expected.src)
# folds advance indexing into basic indexing
class TestIndexingConstFolding(unittest.TestCase):
def test_scalar_index(self):
t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize()
_check_ast_count(1, t[:,:,Tensor(1),:])
_check_ast_count(1, t[:,:,Tensor(1)+2,:])
_check_ast_count(1, t[:,:,Tensor(1),Tensor(0)])
def test_const_tensor_index(self):
# TODO: these can be 0, implement const tensor folded indexing
t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize()
_check_ast_count(1, t[:,:,Tensor.ones(2,1,dtype=dtypes.int),:])
_check_ast_count(1, t[:,:,Tensor.ones(1,2,dtype=dtypes.int)+2,:])
_check_ast_count(1, t[:,:,Tensor.ones(1,1,dtype=dtypes.int),Tensor.zeros(2,1,2,dtype=dtypes.int)])
if __name__ == '__main__': if __name__ == '__main__':
unittest.main() unittest.main()
+1 -158
View File
@@ -1,8 +1,7 @@
import unittest, math import unittest, math
from tinygrad import dtypes from tinygrad import dtypes
from tinygrad.dtype import AddrSpace
from tinygrad.helpers import all_same, Context from tinygrad.helpers import all_same, Context
from tinygrad.uop.ops import GroupOp, UOp, Ops, exec_alu, PatternMatcher, TrackedPatternMatcher, UPat from tinygrad.uop.ops import GroupOp, UOp, Ops, PatternMatcher, TrackedPatternMatcher, UPat
from test.helpers import full_rewrite from test.helpers import full_rewrite
from hypothesis import given, strategies as strat from hypothesis import given, strategies as strat
@@ -16,108 +15,7 @@ def const_value(uop:UOp):
assert uop.op is Ops.CONST assert uop.op is Ops.CONST
return uop.val return uop.val
def evaluate_uop(uop, variables):
if uop.op == Ops.CONST:
return uop.val
elif uop.op == Ops.PARAM and uop.arg.addrspace is AddrSpace.ALU:
return variables[uop.expr]
elif uop.op is Ops.CAST:
return uop.dtype.const(evaluate_uop(uop.src[0], variables))
elif uop.op in GroupOp.ALU:
src_values = [evaluate_uop(src, variables) for src in uop.src]
return exec_alu(uop.op, uop.dtype, src_values)
else:
raise NotImplementedError(f"Unsupported UOp {uop.op}")
class TestArithmeticSimplifications(unittest.TestCase):
def test_full_graph_rewrite_division_by_zero(self):
optimized_div_uop = apply_rewrite(UOp.const(10.0) / UOp.const(0.0))
value = const_value(optimized_div_uop)
self.assertTrue(math.isinf(value) or math.isnan(value))
def test_full_graph_rewrite_redundant_operations(self):
optimized_uop = apply_rewrite((UOp.const(10.0) + UOp.const(0.0)) * UOp.const(1.0))
self.assertIs(optimized_uop, apply_rewrite(UOp.const(10.0)))
def test_full_graph_rewrite_large_graph(self):
prev_uop = UOp.const(0)
for i in range(1, 101):
prev_uop += UOp.const(i)
optimized_uop = apply_rewrite(prev_uop)
self.assertIs(optimized_uop, apply_rewrite(UOp.const(sum(range(1, 101)))))
def test_full_graph_rewrite_division_by_one(self):
optimized_uop = apply_rewrite(UOp.const(42.0) / UOp.const(1.0))
self.assertIs(optimized_uop, apply_rewrite(UOp.const(42.0)))
def test_full_graph_rewrite_modulo_by_one(self):
optimized_uop = apply_rewrite(UOp.const(42) % UOp.const(1))
self.assertIs(optimized_uop, apply_rewrite(UOp.const(0, dtypes.int)))
class TestFoldingAndReduction(unittest.TestCase):
@unittest.skip("reduce is removed now")
def test_full_graph_rewrite_constant_reduction_folding(self):
const1 = UOp.const(5)
const2 = UOp.const(10)
const3 = UOp.const(20)
optimized_sink = apply_rewrite((const1 + const2 + const3).reduce(Ops.ADD))
expected_sum = 5 + 10 + 20
self.assertEqual(optimized_sink.val, expected_sum)
@unittest.skip("reduce is removed now")
def test_full_graph_rewrite_reduction_with_unused_range(self):
const1 = UOp.const(15)
const2 = UOp.const(25)
rng = UOp.range(10, idx=0)
optimized_sink = apply_rewrite((const1 + const2).reduce(Ops.ADD, rng))
expected_sum = 10 * (15 + 25)
self.assertEqual(optimized_sink.val, expected_sum)
@unittest.skip("currently failing")
def test_full_graph_rewrite_range_reduction(self):
simple_range = UOp.range(5, idx=0)
optimized_sink = apply_rewrite(simple_range.reduce(Ops.ADD, simple_range))
expected_sum = sum(range(5))
self.assertEqual(optimized_sink.val, expected_sum)
@unittest.skip("currently failing")
def test_full_graph_rewrite_simple_reduction_folding(self):
simple_range = UOp.range(4, idx=0)
add_uop = simple_range + UOp.const(1)
optimized_sink = apply_rewrite(add_uop.reduce(Ops.ADD, simple_range))
expected_sum = sum(i + 1 for i in range(4))
self.assertEqual(optimized_sink.val, expected_sum)
@unittest.skip("currently failing")
def test_full_graph_rewrite_nested_loop_collapse(self):
outer_range = UOp.range(8, 0)
inner_range = UOp.range(4, 1)
expr = (outer_range * 10) + inner_range
optimized_reduce_uop = apply_rewrite(expr.reduce(Ops.ADD, outer_range, inner_range))
self.assertEqual(optimized_reduce_uop.op, Ops.CONST)
self.assertEqual(optimized_reduce_uop.val, sum((i * 10) + j for i in range(8) for j in range(4)))
class TestModuloAndDivisionFolding(unittest.TestCase): class TestModuloAndDivisionFolding(unittest.TestCase):
def test_full_graph_rewrite_modulo_folding_with_define_var(self):
# index dtype because div-mod rules only work on index
x_var_uop = UOp.variable('x', 0, 100).cast(dtypes.weakint)
optimized_mod_uop = apply_rewrite(((x_var_uop * 4) + 2) % 4)
self.assertIs(optimized_mod_uop, apply_rewrite(UOp.const(2, dtypes.int)))
def test_full_graph_rewrite_division_folding_with_define_var(self):
# index dtype because div-mod rules only work on index
n_var_uop = UOp.variable('n', 1, 1000).cast(dtypes.weakint)
optimized_div_uop = apply_rewrite((n_var_uop * 6) // 3)
self.assertIs(optimized_div_uop, apply_rewrite(n_var_uop * 2))
def test_full_graph_rewrite_complex_mod_div_folding(self):
# index dtype because div-mod rules only work on index
k_var_uop = UOp.variable('k', 0, 50).cast(dtypes.weakint)
optimized_div_uop = apply_rewrite(((k_var_uop * 12 + 8) % 6) // 2)
self.assertIs(optimized_div_uop, apply_rewrite(UOp.const(1, dtypes.int)))
def test_graph_rewrite_div_folding_bug(self): def test_graph_rewrite_div_folding_bug(self):
lhs = UOp.stack(*(UOp.special(32, 'lidx0'),)*4) + UOp.const((0, 256, 512, 768)) lhs = UOp.stack(*(UOp.special(32, 'lidx0'),)*4) + UOp.const((0, 256, 512, 768))
rhs = UOp.const((2,)*4) rhs = UOp.const((2,)*4)
@@ -127,26 +25,6 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
print(opt) print(opt)
if opt.op is Ops.STACK: self.assertFalse(all_same(opt.src)) if opt.op is Ops.STACK: self.assertFalse(all_same(opt.src))
def test_full_graph_rewrite_modulo_large_divisor(self):
# index dtype because div-mod rules only work on index
x_var_uop = UOp.variable('x', 1, 5)
self.assertIs(apply_rewrite(x_var_uop.cast(dtypes.weakint) % 10).render(simplify=False), x_var_uop.render(simplify=False))
def test_full_graph_rewrite_division_with_remainder(self):
x_var_uop = UOp.variable('x', 7, 9, param=True)
optimized_sink = apply_rewrite(x_var_uop // 2)
for x_value in range(7, 10):
self.assertEqual(x_value // 2, evaluate_uop(optimized_sink, {'x': x_value}))
def test_full_graph_rewrite_complex_mod_div_expression(self):
x_var_uop = UOp.variable('x', 1, 10, param=True)
optimized_sink = apply_rewrite(((x_var_uop * 5) % 3) // 2)
for x_value in range(1, 11):
original_result = ((x_value * 5) % 3) // 2
optimized_result = evaluate_uop(optimized_sink, {'x': x_value})
self.assertEqual(original_result, optimized_result)
class TestEdgeCasesAndSpecialOperations(unittest.TestCase): class TestEdgeCasesAndSpecialOperations(unittest.TestCase):
def test_full_graph_rewrite_transcendental_edge_cases(self): def test_full_graph_rewrite_transcendental_edge_cases(self):
optimized_sink = full_rewrite(UOp.const(-1.0).log2().sink(UOp.const(0.0).reciprocal())) optimized_sink = full_rewrite(UOp.const(-1.0).log2().sink(UOp.const(0.0).reciprocal()))
@@ -155,20 +33,6 @@ class TestEdgeCasesAndSpecialOperations(unittest.TestCase):
self.assertTrue(math.isnan(log2_neg), f"Expected NaN for log2(-1.0), got {log2_neg}") self.assertTrue(math.isnan(log2_neg), f"Expected NaN for log2(-1.0), got {log2_neg}")
self.assertTrue(math.isinf(recip_zero) and recip_zero > 0, f"Expected +inf for reciprocal(0.0), got {recip_zero}") self.assertTrue(math.isinf(recip_zero) and recip_zero > 0, f"Expected +inf for reciprocal(0.0), got {recip_zero}")
@unittest.skip("broken")
def test_full_graph_rewrite_modulo_negative_dividend(self):
x_var_uop = UOp.variable('x', -5, -1)
optimized_sink = full_rewrite((x_var_uop % 3).sink())
for x_value in range(-5, 0):
self.assertEqual(x_value % 3, evaluate_uop(optimized_sink.src[0], {'x': x_value}))
@unittest.skip("broken")
def test_full_graph_rewrite_division_negative_divisor(self):
x_var_uop = UOp.variable('x', 1, 5)
optimized_sink = full_rewrite((x_var_uop // -2).sink())
for x_value in range(1, 6):
self.assertEqual(x_value // -2, evaluate_uop(optimized_sink.src[0], {'x': x_value}))
class TestGEPAndVectorizeRewrite(unittest.TestCase): class TestGEPAndVectorizeRewrite(unittest.TestCase):
def test_gep_single_element_extraction(self): def test_gep_single_element_extraction(self):
# GEP on a vector dtype to extract a single element # GEP on a vector dtype to extract a single element
@@ -181,17 +45,6 @@ class TestGEPAndVectorizeRewrite(unittest.TestCase):
self.assertIs(apply_rewrite(UOp.stack(*[base_vector.index(i) for i in (2, 3)])), self.assertIs(apply_rewrite(UOp.stack(*[base_vector.index(i) for i in (2, 3)])),
apply_rewrite(UOp.stack(base_vector.src[2], base_vector.src[3]))) apply_rewrite(UOp.stack(base_vector.src[2], base_vector.src[3])))
def test_gep_on_const_stack(self):
# GEP on a const STACK to extract a single element
const_stack = UOp.const((1.0, 2.0, 3.0, 4.0))
self.assertIs(apply_rewrite(const_stack.index(2)), apply_rewrite(const_stack.src[2]))
def test_gep_tuple_on_const_stack(self):
# GEP on a const STACK using a tuple to extract multiple elements
const_stack = UOp.const((7.0, 8.0, 9.0, 10.0))
self.assertIs(apply_rewrite(UOp.stack(*[const_stack.index(i) for i in (1, 3)])),
apply_rewrite(UOp.stack(const_stack.src[1], const_stack.src[3])))
def test_vectorize_multiple_elements(self): def test_vectorize_multiple_elements(self):
# Vectorizing multiple elements using GEP # Vectorizing multiple elements using GEP
base_vector = UOp.const((5.0, 10.0, 15.0, 20.0)) base_vector = UOp.const((5.0, 10.0, 15.0, 20.0))
@@ -248,16 +101,6 @@ class TestSubstitute(unittest.TestCase):
ret = substitute(ret, {a.sin():b}) ret = substitute(ret, {a.sin():b})
self.assertIs(ret, b.sin()) self.assertIs(ret, b.sin())
# broken due to infinite recursion
# NOTE: VIZ hangs and doesn't recover if you click this one
@unittest.skip("recursion error no longer raised")
def test_assert_inf_recurse(self):
a = UOp.variable('a', 0, 10)
n1 = a.sin()
ret = n1
with self.assertRaises(RecursionError):
ret = substitute(ret, {n1:n1.sqrt()})
def test_sin_to_sqrt(self): def test_sin_to_sqrt(self):
a = UOp.variable('a', 0, 10, dtype=dtypes.float) a = UOp.variable('a', 0, 10, dtype=dtypes.float)
n1 = a.sin() n1 = a.sin()
-78
View File
@@ -1,7 +1,6 @@
import unittest, pytest import unittest, pytest
from tinygrad import dtypes, Variable, Device from tinygrad import dtypes, Variable, Device
from tinygrad.dtype import AddrSpace from tinygrad.dtype import AddrSpace
from tinygrad.helpers import DEBUG
from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType, broadcast_axes, KernelInfo from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType, broadcast_axes, KernelInfo
from tinygrad.uop.symbolic import sym from tinygrad.uop.symbolic import sym
from test.helpers import full_rewrite, to_uops_list from test.helpers import full_rewrite, to_uops_list
@@ -34,13 +33,6 @@ class TestGraphRewriteConst(unittest.TestCase):
self.assertEqual(ret.op, Ops.STACK) self.assertEqual(ret.op, Ops.STACK)
self.assertEqual(const_values(ret), (5,7,9)) self.assertEqual(const_values(ret), (5,7,9))
def test_add_const_lose_v(self):
v1 = UOp.const((0,1,2))
v2 = UOp.const((2,1,0))
ret = graph_rewrite(v1+v2, sym)
self.assertEqual(ret.op, Ops.STACK)
self.assertEqual(const_values(ret), (2,2,2))
def xfail_broken_const_wraparound(fn): def xfail_broken_const_wraparound(fn):
fn = pytest.mark.xfail(reason="const folding does not properly implement modular arithmetic")(fn) fn = pytest.mark.xfail(reason="const folding does not properly implement modular arithmetic")(fn)
return unittest.expectedFailure(fn) return unittest.expectedFailure(fn)
@@ -190,12 +182,6 @@ class TestGraphRewrite(unittest.TestCase):
self.assertEqual(len([x for x in sink.toposort() if x.op is Ops.CONST]), 1) self.assertEqual(len([x for x in sink.toposort() if x.op is Ops.CONST]), 1)
class TestUOpGraph(unittest.TestCase): class TestUOpGraph(unittest.TestCase):
def test_add_constant_fold(self):
c1 = UOp.const(1.0, dtypes.float)
c2 = UOp.const(2.0, dtypes.float)
out = c1+c2
self.assertIs(out.simplify(), UOp.const(3.0, dtypes.float))
def test_where_same_fold(self): def test_where_same_fold(self):
v = UOp.variable('tmp', 0, 1) v = UOp.variable('tmp', 0, 1)
c0 = UOp.const(0) c0 = UOp.const(0)
@@ -216,11 +202,6 @@ class TestUOpGraph(unittest.TestCase):
out = bf.cast(dtypes.int) out = bf.cast(dtypes.int)
self.assertIs(full_rewrite(out.sink()).src[0], full_rewrite(UOp.const(0, dtypes.int).sink()).src[0]) self.assertIs(full_rewrite(out.sink()).src[0], full_rewrite(UOp.const(0, dtypes.int).sink()).src[0])
def test_const_bitcast(self):
bf = UOp.const(1.0, dtypes.float)
out = bf.bitcast(dtypes.uint32)
self.assertIs(out.simplify(), UOp.const(0x3F800000, dtypes.uint32))
def test_devectorize_derives_lane_dtype(self): def test_devectorize_derives_lane_dtype(self):
from tinygrad.codegen import do_devectorize from tinygrad.codegen import do_devectorize
# an Invalid lane derives bool while the value lane derives float: the lane rebuild must derive, not inherit # an Invalid lane derives bool while the value lane derives float: the lane rebuild must derive, not inherit
@@ -229,58 +210,6 @@ class TestUOpGraph(unittest.TestCase):
invalid_lane_mul = next(u for u in out.src[0].toposort() if u.op is Ops.MUL) invalid_lane_mul = next(u for u in out.src[0].toposort() if u.op is Ops.MUL)
self.assertIs(invalid_lane_mul.dtype, dtypes.bool) self.assertIs(invalid_lane_mul.dtype, dtypes.bool)
@unittest.skip("this test isn't valid uops")
def test_noop_vectorize_fold(self):
d0 = UOp.param(0, dtypes.float, (1,))
idx = UOp.const(0)
ld = d0.load(idx, dtype=dtypes.float)
vec = UOp.stack(ld)
x = vec.index(0)
alu = x.sqrt()
out = d0.index(idx).store(alu)
uops = to_uops_list([out])
self.assertEqual(len([x for x in uops if x.op is Ops.STACK]), 0)
@unittest.skip("this test isn't valid uops")
def test_gep_vec_fold(self):
d0 = UOp.param(0, dtypes.float, (1,))
d1 = UOp.param(1, dtypes.float, (1,))
d2 = UOp.param(2, dtypes.float, (1,))
idx = UOp.const(0)
def _test_vec(geps, count=4):
vec = UOp.stack(*geps)
out = d0.index(idx).store(vec)
rewritten = full_rewrite(out.sink())
if DEBUG >= 4:
from tinygrad import Device
print(Device[Device.DEFAULT].renderer.render(rewritten.toposort()))
return rewritten.src[0].src[1]
# possible
val = d1.index(idx).load(dtype=dtypes.float)
xyzw = tuple(val.index(i) for i in range(4))
self.assertIs(_test_vec(xyzw).op, Ops.LOAD)
# unaligned
val = d1.index(idx).load(dtype=dtypes.float)
wzyx = tuple(val.index(i) for i in reversed(range(4)))
self.assertIs(_test_vec(wzyx).op, Ops.STACK)
# different_size
val = d1.index(idx).load(dtype=dtypes.float)
xy = tuple(val.index(i) for i in range(2))
self.assertIs(_test_vec(xy+xy).op, Ops.STACK)
val = d1.index(idx).load(dtype=dtypes.float)
xy = tuple(val.index(i) for i in range(2))
self.assertIs(_test_vec(xy, count=2).op, Ops.STACK)
# different vals
val1 = d1.index(idx).load(dtype=dtypes.float)
val2 = d2.index(idx).load(dtype=dtypes.float)
xy1 = tuple(val1.index(i) for i in range(2))
xy2 = tuple(val2.index(i) for i in range(2))
self.assertIs(_test_vec(xy1+xy2).op, Ops.STACK)
def test_gep_vec_const_fold(self): def test_gep_vec_const_fold(self):
for vec_size in [2, 4, 8]: for vec_size in [2, 4, 8]:
consts = [UOp.const(float(i), dtypes.float) for i in range(vec_size)] consts = [UOp.const(float(i), dtypes.float) for i in range(vec_size)]
@@ -468,13 +397,6 @@ class TestUOpGraph(unittest.TestCase):
# only the second store happens # only the second store happens
self.assertEqual(len([u for u in uops if u.op is Ops.STORE]), 1) self.assertEqual(len([u for u in uops if u.op is Ops.STORE]), 1)
@unittest.skip("this is a uop type error")
def test_asserts_bad_gate(self):
glbl0 = UOp.param(0, dtypes.int, (1,))
idx = UOp.const(0)
bad_gate = UOp.const(1)
with self.assertRaises(AssertionError): to_uops_list([glbl0.index(idx).store(UOp.const(42), bad_gate)])
def test_after_end(self): def test_after_end(self):
r = UOp.range(10, 0) r = UOp.range(10, 0)