IQ.Pilot Release Commit @ bec7652
This commit is contained in:
IQ.Lvbs history cleanup
@@ -0,0 +1,35 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor, dtypes, TinyJit, UOp
|
||||
from tinygrad.llm.model import apply_rope as apply_rope_new, precompute_freqs_cis
|
||||
from test.helpers import assert_jit_cache_len, check_schedule
|
||||
|
||||
def apply_rope(x:Tensor, start_pos:int):
|
||||
B, H, T, Hd = x.shape
|
||||
precompute_freqs_cis.cache_clear()
|
||||
freqs_cis = precompute_freqs_cis(Hd, start_pos+T)[start_pos:start_pos+T]
|
||||
return apply_rope_new(x, freqs_cis)
|
||||
|
||||
class TestAttention(unittest.TestCase):
|
||||
def test_half_qkv_buffers(self):
|
||||
BS, seqlen, dim = 10, 4, 100
|
||||
q = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
|
||||
k = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
|
||||
v = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
|
||||
attn = q.scaled_dot_product_attention(k, v)
|
||||
# attention has 4 kernels now
|
||||
check_schedule(attn, 4)
|
||||
|
||||
def test_apply_rope_jit_prune(self):
|
||||
def rope_fn(x_in, pos): return apply_rope(x_in, pos)
|
||||
rope_noprune = TinyJit(rope_fn)
|
||||
rope_prune = TinyJit(rope_fn, prune=True)
|
||||
|
||||
v_pos = UOp.variable("start_pos", 0, 100)
|
||||
for _ in range(3):
|
||||
rope_noprune(Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32), v_pos.bind(1))
|
||||
rope_prune(Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32), v_pos.bind(1))
|
||||
assert_jit_cache_len(rope_prune, 1)
|
||||
assert_jit_cache_len(rope_noprune, 3)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
535
artifacts/package_sources/tinygrad/test/null/test_autogen.py
Normal file
535
artifacts/package_sources/tinygrad/test/null/test_autogen.py
Normal file
@@ -0,0 +1,535 @@
|
||||
import ctypes, struct, subprocess, tempfile, unittest
|
||||
from tinygrad.helpers import OSX, WIN
|
||||
from tinygrad.runtime.support.c import DLL, record, Field
|
||||
from tinygrad.runtime.support import c
|
||||
from tinygrad.runtime.support.autogen import gen
|
||||
|
||||
@unittest.skipIf(WIN, "doesn't compile on windows")
|
||||
class TestC(unittest.TestCase):
|
||||
def compile(self, src):
|
||||
with tempfile.NamedTemporaryFile(suffix=".so") as f:
|
||||
subprocess.check_output(('clang', '-x', 'c', '-fPIC', '-shared', '-', '-o', f.name), input=src.encode())
|
||||
return DLL("test", f.name)
|
||||
|
||||
def test_struct_array_init(self):
|
||||
@record
|
||||
class Foo(c.Struct):
|
||||
SIZE = 12
|
||||
a = Field(ctypes.c_int * 3, 0)
|
||||
|
||||
f = Foo((1,2,3))
|
||||
assert f.a[0] == 1
|
||||
assert f.a[1] == 2
|
||||
assert f.a[2] == 3
|
||||
f = Foo((ctypes.c_int * 3)(1,2,3))
|
||||
assert f.a[0] == 1
|
||||
assert f.a[1] == 2
|
||||
assert f.a[2] == 3
|
||||
|
||||
def test_field_ranges(self):
|
||||
@record
|
||||
class Foo(c.Struct):
|
||||
SIZE = 2
|
||||
s = Field(ctypes.c_int8, 0)
|
||||
u = Field(ctypes.c_uint8, 1)
|
||||
|
||||
f = Foo()
|
||||
f.s = -1
|
||||
f.u = -1
|
||||
assert f.s == -1
|
||||
assert f.u == 255
|
||||
|
||||
# this syntax is inherited from ctypes, but it seems a bit nonsensical?
|
||||
def test_voidp_none(self):
|
||||
@record
|
||||
class Foo(c.Struct):
|
||||
SIZE = 8
|
||||
p = Field(ctypes.c_void_p, 0)
|
||||
|
||||
f = Foo(None)
|
||||
assert f.p is None
|
||||
f.p = ctypes.c_void_p(0xDEADBEEF)
|
||||
assert f.p == 0xDEADBEEF
|
||||
f.p = None
|
||||
assert f.p is None
|
||||
|
||||
def test_packed_struct(self):
|
||||
@record
|
||||
class Baz(c.Struct):
|
||||
SIZE = 8
|
||||
a = Field(ctypes.c_uint, 0, 30)
|
||||
b = Field(ctypes.c_uint, 3, 30, 6)
|
||||
c = Field(ctypes.c_uint, 7, 2, 4)
|
||||
d = Field(ctypes.c_uint, 7, 2, 6)
|
||||
|
||||
b = Baz(0x3AAADEAD, 0xBEEF, 1, 0)
|
||||
assert b.a == 0x3AAADEAD
|
||||
assert b.b == 0xBEEF
|
||||
assert b.c == 1
|
||||
assert b.d == 0
|
||||
|
||||
b.a = 0xCAFE
|
||||
assert b.a == 0xCAFE
|
||||
assert b.b == 0xBEEF
|
||||
assert b.c == 1
|
||||
assert b.d == 0
|
||||
|
||||
def test_packed_struct_interop(self):
|
||||
@record
|
||||
class Baz(c.Struct):
|
||||
SIZE = 8
|
||||
a = Field(ctypes.c_int, 0, 30)
|
||||
b = Field(ctypes.c_int, 3, 30, 6)
|
||||
c = Field(ctypes.c_int, 7, 2, 4)
|
||||
d = Field(ctypes.c_int, 7, 2, 6)
|
||||
|
||||
src = '''
|
||||
struct __attribute__((packed)) baz {
|
||||
int a:30;
|
||||
int b:30;
|
||||
int c:2;
|
||||
int d:2;
|
||||
};
|
||||
|
||||
int test(struct baz x) {
|
||||
return x.a + x.b + x.c + x.d;
|
||||
}
|
||||
'''
|
||||
dll = self.compile(src)
|
||||
b = Baz(0xAA000, 0x00BB0, 0, 1)
|
||||
@dll.bind(ctypes.c_int, Baz)
|
||||
def test(x:Baz) -> ctypes.c_int: ...
|
||||
self.assertEqual(test(b), b.a + b.b + b.c + b.d)
|
||||
|
||||
# https://github.com/python/cpython/issues/90914
|
||||
def test_bitfield_interop(self):
|
||||
@record
|
||||
class Baz(c.Struct):
|
||||
SIZE = 1
|
||||
a = Field(ctypes.c_bool, 0, 1, 0)
|
||||
b = Field(ctypes.c_bool, 0, 1, 1)
|
||||
c = Field(ctypes.c_bool, 0, 1, 2)
|
||||
d = Field(ctypes.c_bool, 0, 1, 3)
|
||||
e = Field(ctypes.c_bool, 0, 1, 4)
|
||||
f = Field(ctypes.c_bool, 0, 1, 5)
|
||||
g = Field(ctypes.c_bool, 0, 1, 6)
|
||||
h = Field(ctypes.c_bool, 0, 1, 7)
|
||||
src = '''#include <stdbool.h>
|
||||
struct baz {
|
||||
bool a:1, b:1, c:1, d:1, e:1, f:1, g:1, h:1;
|
||||
};
|
||||
|
||||
int test(struct baz x) {
|
||||
return x.c;
|
||||
}
|
||||
'''
|
||||
dll = self.compile(src)
|
||||
@dll.bind(ctypes.c_int, Baz)
|
||||
def test(x:Baz) -> ctypes.c_int: ...
|
||||
for i in range(8): self.assertEqual(test(Baz(*(j==i for j in range(8)))), i==2)
|
||||
|
||||
def test_struct_interop(self):
|
||||
@record
|
||||
class Baz(c.Struct):
|
||||
SIZE = 32
|
||||
a = Field(ctypes.c_int, 0)
|
||||
b = Field(ctypes.c_int, 4)
|
||||
c = Field(ctypes.c_int, 8)
|
||||
d = Field(ctypes.c_int, 12)
|
||||
e = Field(ctypes.c_int, 16)
|
||||
f = Field(ctypes.c_int, 20)
|
||||
g = Field(ctypes.c_int, 24)
|
||||
h = Field(ctypes.c_int, 28)
|
||||
src = '''#include <stdio.h>
|
||||
struct baz {
|
||||
int a, b, c, d, e, f, g, h;
|
||||
};
|
||||
|
||||
struct baz test(struct baz x) {
|
||||
return (struct baz){x.h, x.g, x.f, x.e, x.d, x.c, x.b, x.a};
|
||||
}
|
||||
'''
|
||||
dll = self.compile(src)
|
||||
@dll.bind(Baz, Baz)
|
||||
def test(x:Baz) -> Baz: ...
|
||||
self.assertEqual(bytes(test(Baz(*range(8)))), struct.pack("8i", *range(7, -1, -1)))
|
||||
|
||||
def test_aos_interop(self):
|
||||
@record
|
||||
class Item(c.Struct):
|
||||
SIZE = 4
|
||||
val = Field(ctypes.c_int, 0)
|
||||
src = """
|
||||
struct item { int val; };
|
||||
int test(struct item arr[3]) {
|
||||
int ret = 0;
|
||||
for (int i = 0; i < 3; i++) ret += arr[i].val;
|
||||
return ret;
|
||||
}
|
||||
"""
|
||||
dll = self.compile(src)
|
||||
@dll.bind(ctypes.c_int, Item * 3)
|
||||
def test(arr:(Item * 3)) -> ctypes.c_int: ...
|
||||
self.assertEqual(test((Item * 3)(Item(10), Item(20), Item(30))), 60)
|
||||
|
||||
def test_soa_interop(self):
|
||||
@record
|
||||
class Row(c.Struct):
|
||||
SIZE = 16
|
||||
data = Field(ctypes.c_int * 3, 0)
|
||||
src = """
|
||||
struct row { int data[3]; };
|
||||
struct row test(struct row x) {
|
||||
return (struct row){{ x.data[2], x.data[1], x.data[0] }};
|
||||
}
|
||||
"""
|
||||
dll = self.compile(src)
|
||||
@dll.bind(Row, Row)
|
||||
def test(x:Row) -> Row: ...
|
||||
r = test(Row((ctypes.c_int * 3)(10, 20, 30)))
|
||||
self.assertIsInstance(r, Row)
|
||||
self.assertEqual(r.data[0], 30)
|
||||
self.assertEqual(r.data[1], 20)
|
||||
self.assertEqual(r.data[2], 10)
|
||||
|
||||
def test_soa_ptr_interop(self):
|
||||
@record
|
||||
class Row(c.Struct):
|
||||
SIZE = 8
|
||||
data = Field(c.POINTER[ctypes.c_int], 0)
|
||||
src = """
|
||||
struct row { int *data; };
|
||||
int test(struct row x) {
|
||||
return x.data[2] + x.data[1] + x.data[0];
|
||||
}
|
||||
"""
|
||||
dll = self.compile(src)
|
||||
@dll.bind(ctypes.c_int, Row)
|
||||
def test(x:Row) -> ctypes.c_int: ...
|
||||
assert test(Row((ctypes.c_int * 3)(10, 20, 30))) == 60
|
||||
|
||||
def test_nested_struct_interop(self):
|
||||
@record
|
||||
class Inner(c.Struct):
|
||||
SIZE = 4
|
||||
a = Field(ctypes.c_int, 0)
|
||||
@record
|
||||
class Outer(c.Struct):
|
||||
SIZE = 8
|
||||
inner = Field(Inner, 0)
|
||||
b = Field(ctypes.c_int, 4)
|
||||
src = """
|
||||
struct i { int a; };
|
||||
struct o { struct i i; int b; };
|
||||
struct o test(struct o x) {
|
||||
return (struct o){(struct i){ x.b }, x.i.a };
|
||||
}
|
||||
"""
|
||||
dll = self.compile(src)
|
||||
@dll.bind(Outer, Outer)
|
||||
def test(x:Outer) -> Outer: ...
|
||||
o = test(Outer(Inner(10), 20))
|
||||
self.assertEqual(o.inner.a, 20)
|
||||
self.assertEqual(o.b, 10)
|
||||
|
||||
def test_struct_pointer_interop(self):
|
||||
@record
|
||||
class Foo(c.Struct):
|
||||
SIZE = 8
|
||||
a = Field(ctypes.c_int, 0)
|
||||
b = Field(ctypes.c_int, 4)
|
||||
src = """
|
||||
struct foo { int a, b; };
|
||||
struct foo *test(struct foo *f) {
|
||||
int x = f->a;
|
||||
f->a = f->b;
|
||||
f->b = x;
|
||||
return f;
|
||||
}
|
||||
"""
|
||||
dll = self.compile(src)
|
||||
@dll.bind(ctypes.POINTER(Foo), ctypes.POINTER(Foo))
|
||||
def test(f:ctypes.POINTER(Foo)) -> ctypes.POINTER(Foo): ...
|
||||
inp = ctypes.pointer(Foo(10, 20))
|
||||
out = test(inp)
|
||||
self.assertEqual(out.contents.a, 20)
|
||||
self.assertEqual(out.contents.b, 10)
|
||||
|
||||
def test_pointer_field_roundtrip(self):
|
||||
# This tests storing a pointer in a record struct field and passing it to C
|
||||
# Mimics how mesa.struct_lp_build_tgsi_params.mask is used
|
||||
from tinygrad.runtime.support.c import POINTER
|
||||
@record
|
||||
class Inner(c.Struct):
|
||||
SIZE = 8
|
||||
value = Field(ctypes.c_int, 0)
|
||||
flag = Field(ctypes.c_int, 4)
|
||||
@record
|
||||
class Outer(c.Struct):
|
||||
SIZE = 16
|
||||
x = Field(ctypes.c_int, 0)
|
||||
inner_ptr = Field(POINTER[Inner], 8)
|
||||
|
||||
src = """
|
||||
struct inner { int value; int flag; };
|
||||
struct outer { int x; struct inner *inner_ptr; };
|
||||
int test(struct inner *p) {
|
||||
return p->value + p->flag;
|
||||
}
|
||||
"""
|
||||
dll = self.compile(src)
|
||||
@dll.bind(ctypes.c_int, ctypes.POINTER(Inner))
|
||||
def test(p:POINTER[Inner]) -> ctypes.c_int: ...
|
||||
|
||||
inner = Inner(value=42, flag=10)
|
||||
outer = Outer(x=1, inner_ptr=ctypes.pointer(inner))
|
||||
# Retrieve pointer from struct field and pass to C
|
||||
self.assertEqual(test(outer.inner_ptr), 52)
|
||||
|
||||
def test_pointer_field_loses_reference(self):
|
||||
# BUG: When a pointer is stored in a record struct field, only the address bytes are saved.
|
||||
# The pointer's _objects dict (which prevents GC of the pointed-to object) is lost.
|
||||
# This causes the pointed-to object to be garbage collected, leading to use-after-free.
|
||||
from tinygrad.runtime.support.c import POINTER
|
||||
@record
|
||||
class MaskContext(c.Struct):
|
||||
SIZE = 16
|
||||
value = Field(ctypes.c_int, 0)
|
||||
initialized = Field(ctypes.c_int, 4)
|
||||
ptr = Field(ctypes.c_void_p, 8)
|
||||
@record
|
||||
class Params(c.Struct):
|
||||
SIZE = 16
|
||||
x = Field(ctypes.c_int, 0)
|
||||
mask = Field(POINTER[MaskContext], 8)
|
||||
|
||||
src = """
|
||||
struct mask_ctx { int value; int initialized; void *ptr; };
|
||||
void mask_begin(struct mask_ctx *m, int val) { m->value = val; m->initialized = 1; }
|
||||
int mask_end(struct mask_ctx *m) { return m->value + m->initialized; }
|
||||
"""
|
||||
dll = self.compile(src)
|
||||
@dll.bind(None, ctypes.POINTER(MaskContext), ctypes.c_int)
|
||||
def mask_begin(m:POINTER[MaskContext], val:ctypes.c_int) -> None: ...
|
||||
@dll.bind(ctypes.c_int, ctypes.POINTER(MaskContext))
|
||||
def mask_end(m:POINTER[MaskContext]) -> ctypes.c_int: ...
|
||||
|
||||
# When MaskContext() is created inline, it gets garbage collected after the pointer
|
||||
# is stored because only the address bytes are saved, not the _objects reference.
|
||||
params = Params(x=1, mask=ctypes.pointer(MaskContext()))
|
||||
mask_begin(params.mask, 42)
|
||||
result = mask_end(params.mask)
|
||||
self.assertEqual(result, 43) # 42 + 1
|
||||
|
||||
@unittest.skipIf(OSX and ('MTLCompiler' in DLL._loaded_ or 'llvm' in DLL._loaded_), "libclang can't be loaded after MTLCompiler or llvm on OSX")
|
||||
@unittest.skipIf(WIN, "doesn't compile on windows")
|
||||
class TestAutogen(unittest.TestCase):
|
||||
def run_gen(self, contents):
|
||||
with tempfile.NamedTemporaryFile(mode='w', suffix='.h') as f:
|
||||
f.write(contents)
|
||||
f.flush()
|
||||
|
||||
generated_code = gen(name="test_header", dll=None, files=[f.name])
|
||||
|
||||
namespace = {}
|
||||
exec(generated_code, namespace)
|
||||
return namespace
|
||||
|
||||
def test_packed_structs(self):
|
||||
ns = self.run_gen("""
|
||||
typedef unsigned NvU32;
|
||||
typedef unsigned long NvU64;
|
||||
|
||||
typedef struct
|
||||
{
|
||||
NvU32 version;
|
||||
NvU32 size;
|
||||
NvU64 gfwImageOffset;
|
||||
NvU32 gfwImageSize;
|
||||
NvU32 flags;
|
||||
} __attribute__((packed)) FWSECLIC_READ_VBIOS_DESC;
|
||||
|
||||
#define FWSECLIC_READ_VBIOS_STRUCT_FLAGS (2)
|
||||
|
||||
typedef struct
|
||||
{
|
||||
NvU32 version;
|
||||
NvU32 size;
|
||||
NvU32 frtsRegionOffset4K;
|
||||
NvU32 frtsRegionSize;
|
||||
NvU32 frtsRegionMediaType;
|
||||
} __attribute__((packed)) FWSECLIC_FRTS_REGION_DESC;
|
||||
|
||||
#define FWSECLIC_FRTS_REGION_MEDIA_FB (2)
|
||||
#define FWSECLIC_FRTS_REGION_SIZE_1MB_IN_4K (0x100)
|
||||
|
||||
typedef struct
|
||||
{
|
||||
FWSECLIC_READ_VBIOS_DESC readVbiosDesc;
|
||||
FWSECLIC_FRTS_REGION_DESC frtsRegionDesc;
|
||||
} __attribute__((packed)) FWSECLIC_FRTS_CMD;
|
||||
""")
|
||||
|
||||
FWSECLIC_READ_VBIOS_DESC = ns['FWSECLIC_READ_VBIOS_DESC']
|
||||
FWSECLIC_FRTS_REGION_DESC = ns['FWSECLIC_FRTS_REGION_DESC']
|
||||
FWSECLIC_FRTS_CMD = ns['FWSECLIC_FRTS_CMD']
|
||||
|
||||
read_vbios_desc = FWSECLIC_READ_VBIOS_DESC(version=0x1, size=ctypes.sizeof(FWSECLIC_READ_VBIOS_DESC), flags=2)
|
||||
frst_reg_desc = FWSECLIC_FRTS_REGION_DESC(version=0x1, size=ctypes.sizeof(FWSECLIC_FRTS_REGION_DESC),
|
||||
frtsRegionOffset4K=0xdead, frtsRegionSize=0x100, frtsRegionMediaType=2)
|
||||
frts_cmd = FWSECLIC_FRTS_CMD(readVbiosDesc=read_vbios_desc, frtsRegionDesc=frst_reg_desc)
|
||||
assert int.from_bytes(frts_cmd, 'little') == 0x2000001000000dead0000001400000001000000020000000000000000000000000000001800000001
|
||||
assert int.from_bytes(frts_cmd.readVbiosDesc, 'little') == int.from_bytes(read_vbios_desc, 'little')
|
||||
assert int.from_bytes(frts_cmd.frtsRegionDesc, 'little') == int.from_bytes(frst_reg_desc, 'little')
|
||||
assert frts_cmd.readVbiosDesc.__class__ is FWSECLIC_READ_VBIOS_DESC
|
||||
assert frts_cmd.frtsRegionDesc.__class__ is FWSECLIC_FRTS_REGION_DESC
|
||||
|
||||
def test_gen_from_header(self):
|
||||
namespace = self.run_gen("""
|
||||
typedef struct {
|
||||
int x;
|
||||
int y;
|
||||
} Point;
|
||||
|
||||
typedef enum {
|
||||
RED = 0,
|
||||
GREEN = 1,
|
||||
BLUE = 2
|
||||
} Color;
|
||||
|
||||
typedef struct {
|
||||
Point origin;
|
||||
int width;
|
||||
int height;
|
||||
Color color;
|
||||
} Rectangle;
|
||||
|
||||
int add_points(Point a, Point b);""")
|
||||
|
||||
self.assertIn('Point', namespace)
|
||||
self.assertIn('Color', namespace)
|
||||
self.assertIn('Rectangle', namespace)
|
||||
self.assertIn('RED', namespace)
|
||||
self.assertIn('GREEN', namespace)
|
||||
self.assertIn('BLUE', namespace)
|
||||
|
||||
self.assertEqual(namespace['RED'], 0)
|
||||
self.assertEqual(namespace['GREEN'], 1)
|
||||
self.assertEqual(namespace['BLUE'], 2)
|
||||
|
||||
Point = namespace['Point']
|
||||
p = Point()
|
||||
self.assertTrue(hasattr(p, 'x'))
|
||||
self.assertTrue(hasattr(p, 'y'))
|
||||
|
||||
Rectangle = namespace['Rectangle']
|
||||
rect = Rectangle()
|
||||
self.assertTrue(hasattr(rect, 'origin'))
|
||||
self.assertTrue(hasattr(rect, 'width'))
|
||||
self.assertTrue(hasattr(rect, 'height'))
|
||||
self.assertTrue(hasattr(rect, 'color'))
|
||||
|
||||
p2 = Point(10, 20)
|
||||
self.assertEqual(p2.x, 10)
|
||||
self.assertEqual(p2.y, 20)
|
||||
|
||||
def test_struct_ordering(self):
|
||||
namespace = self.run_gen("""
|
||||
struct A;
|
||||
struct C;
|
||||
typedef struct A A;
|
||||
|
||||
struct B {
|
||||
struct C *c_ptr;
|
||||
};
|
||||
|
||||
struct C {
|
||||
struct A *a_ptr;
|
||||
};
|
||||
|
||||
struct A {
|
||||
int x;
|
||||
struct B *b_ptr;
|
||||
};""")
|
||||
|
||||
self.assertIn('struct_A', namespace)
|
||||
self.assertIn('struct_B', namespace)
|
||||
self.assertIn('struct_C', namespace)
|
||||
A, B, C = namespace['A'], namespace['struct_B'], namespace['struct_C']
|
||||
a, b, c = A(), B(), C()
|
||||
self.assertTrue(hasattr(a, 'x'))
|
||||
self.assertTrue(hasattr(a, 'b_ptr'))
|
||||
self.assertTrue(hasattr(b, 'c_ptr'))
|
||||
self.assertTrue(hasattr(c, 'a_ptr'))
|
||||
|
||||
def test_anonymous_children(self):
|
||||
namespace = self.run_gen("""
|
||||
struct foo {
|
||||
struct {
|
||||
int a,b;
|
||||
} bar;
|
||||
};
|
||||
""")
|
||||
|
||||
self.assertIn('struct_foo', namespace)
|
||||
self.assertIn('struct_foo_bar', namespace)
|
||||
|
||||
def test_enums(self):
|
||||
namespace = self.run_gen("""
|
||||
enum Foo { A, B, C };
|
||||
enum Bar { X, Y, Z };
|
||||
""")
|
||||
|
||||
assert namespace["A"] == 0
|
||||
assert namespace["B"] == 1
|
||||
assert namespace["C"] == 2
|
||||
assert namespace["X"] == 0
|
||||
assert namespace["Y"] == 1
|
||||
assert namespace["Z"] == 2
|
||||
assert namespace["enum_Foo"].get(0) == "A"
|
||||
assert namespace["enum_Foo"].get(1) == "B"
|
||||
assert namespace["enum_Foo"].get(2) == "C"
|
||||
assert namespace["enum_Bar"].get(0) == "X"
|
||||
assert namespace["enum_Bar"].get(1) == "Y"
|
||||
assert namespace["enum_Bar"].get(2) == "Z"
|
||||
|
||||
@unittest.skipIf(OSX, "can't find stdint?")
|
||||
def test_packed_fields(self):
|
||||
ns = self.run_gen("""#include <stdint.h>
|
||||
typedef struct die_info
|
||||
{
|
||||
uint16_t die_id;
|
||||
uint16_t die_offset; /* Points to the corresponding die_header structure */
|
||||
} die_info;
|
||||
|
||||
typedef struct ip_discovery_header
|
||||
{
|
||||
uint32_t signature; /* Table Signature */
|
||||
uint16_t version; /* Table Version */
|
||||
uint16_t size; /* Table Size */
|
||||
uint32_t id; /* Table ID */
|
||||
uint16_t num_dies; /* Number of Dies */
|
||||
die_info die_info[16]; /* list die information for up to 16 dies */
|
||||
union {
|
||||
uint16_t padding[1]; /* version <= 3 */
|
||||
struct { /* version == 4 */
|
||||
uint8_t base_addr_64_bit : 1; /* ip structures are using 64 bit base address */
|
||||
uint8_t reserved : 7;
|
||||
uint8_t reserved2;
|
||||
};
|
||||
};
|
||||
} ip_discovery_header;
|
||||
""")
|
||||
|
||||
ip_discovery_header = ns['ip_discovery_header']
|
||||
|
||||
hdr = b'IPDS\x04\x00|\x1d\x80\x1a\xffd\x01\x00\x00\x00\x8c\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00' # noqa: E501
|
||||
ihdr = ip_discovery_header.from_buffer_copy(hdr)
|
||||
|
||||
assert ctypes.sizeof(ihdr) == 80
|
||||
assert ihdr.signature == 0x53445049
|
||||
assert ihdr.version == 0x0004
|
||||
assert ihdr.num_dies == 1
|
||||
assert ihdr.base_addr_64_bit == 1
|
||||
|
||||
if __name__ == "__main__": unittest.main()
|
||||
@@ -0,0 +1,30 @@
|
||||
import unittest, io
|
||||
from contextlib import redirect_stdout
|
||||
from tinygrad import Tensor, dtypes, Device
|
||||
from tinygrad.helpers import OSX
|
||||
from tinygrad.engine.realize import compile_linear
|
||||
from tinygrad.codegen import to_program
|
||||
|
||||
class TestCompileFailures(unittest.TestCase):
|
||||
def compile(self, out:Tensor):
|
||||
compile_linear(out.schedule_linear())
|
||||
|
||||
def test_interpolate_atari(self):
|
||||
self.compile(Tensor.empty(210, 160, dtype='uint8').interpolate((64, 64)))
|
||||
|
||||
def test_add_max_uchar(self):
|
||||
self.compile((Tensor.empty(1024, dtype='uint8') + Tensor.empty(1024, dtype='uint8')).max())
|
||||
|
||||
class TestDisassembly(unittest.TestCase):
|
||||
@unittest.skipUnless(Device.DEFAULT == "CPU" and OSX, "m series cpus support fp16 arithmetic")
|
||||
def test_float16_alu(self):
|
||||
c = Tensor([1], dtype=dtypes.float16) + Tensor([1], dtype=dtypes.float16)
|
||||
s = c.schedule_linear().src[-1]
|
||||
p = to_program(s.src[0], Device[Device.DEFAULT].renderer)
|
||||
lib = Device[Device.DEFAULT].compiler.compile(p.src[2].arg)
|
||||
out = io.StringIO()
|
||||
with redirect_stdout(out): Device[Device.DEFAULT].compiler.disassemble(lib)
|
||||
assert "fcvt" not in out.getvalue()
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,178 @@
|
||||
import unittest, itertools, math
|
||||
from tinygrad import Tensor, dtypes, Context
|
||||
from tinygrad.dtype import DType, ConstType, truncate
|
||||
from tinygrad.uop.ops import Ops, UOp
|
||||
from test.helpers import full_rewrite
|
||||
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):
|
||||
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_invalid_poison(self):
|
||||
self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid)
|
||||
|
||||
def test_cast_commits_to_dtype_grid(self):
|
||||
# committing a weak const to a stated width puts the value on that width's grid, same as storage packing and native compilers
|
||||
v = 1/123008 # not representable in float16
|
||||
out = UOp.const(v).cast(dtypes.half).simplify()
|
||||
self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.half, truncate[dtypes.half](v)))
|
||||
self.assertNotEqual(out.val, v)
|
||||
# the grid commit preserves the sign of zero
|
||||
self.assertEqual(math.copysign(1, UOp.const(-0.0).cast(dtypes.half).simplify().val), -1)
|
||||
# observable at tensor level: the const-folded comparison agrees with the committed value
|
||||
self.assertTrue((Tensor(-3.2).cast(dtypes.float32) <= truncate[dtypes.float32](-3.2)).item())
|
||||
|
||||
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):
|
||||
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 = full_rewrite(UOp.const(from_v, from_dt).bitcast(to_dt).sink()).src[0]
|
||||
self.assertEqual(r.op, Ops.CONST, msg:=f"{from_dt} -> {to_dt} ({from_v} -> {to_v})")
|
||||
self.assertEqual(r.dtype, to_dt, msg)
|
||||
np.testing.assert_equal(r.val, to_v, msg)
|
||||
|
||||
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):
|
||||
srcs = full_rewrite(UOp.const((-1, -2**31, 75), dtypes.int32).bitcast(dtypes.uint32).sink()).src
|
||||
self.assertTrue(all(r.op is Ops.CONST and r.dtype == dtypes.uint32 for r in srcs))
|
||||
self.assertEqual(tuple(x.val for x in srcs), (2**32-1, 2**31, 75))
|
||||
|
||||
# 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__':
|
||||
unittest.main()
|
||||
196
artifacts/package_sources/tinygrad/test/null/test_device.py
Normal file
196
artifacts/package_sources/tinygrad/test/null/test_device.py
Normal file
@@ -0,0 +1,196 @@
|
||||
#!/usr/bin/env python
|
||||
import unittest, os, subprocess
|
||||
from unittest.mock import patch
|
||||
from tinygrad import Tensor
|
||||
from tinygrad.device import Device, Compiler, enumerate_devices_str
|
||||
from tinygrad.helpers import diskcache_get, diskcache_put, getenv, Context, Target, WIN, OSX, DEV
|
||||
from tinygrad.runtime.support.c import DLL
|
||||
|
||||
class TestDevice(unittest.TestCase):
|
||||
def test_canonicalize(self):
|
||||
self.assertEqual(Device.canonicalize(None), Device.DEFAULT)
|
||||
self.assertEqual(Device.canonicalize("CPU"), "CPU")
|
||||
self.assertEqual(Device.canonicalize("cpu"), "CPU")
|
||||
self.assertEqual(Device.canonicalize("CL"), "CL")
|
||||
self.assertEqual(Device.canonicalize("CL:0"), "CL")
|
||||
self.assertEqual(Device.canonicalize("cl:0"), "CL")
|
||||
self.assertEqual(Device.canonicalize("CL:1"), "CL:1")
|
||||
self.assertEqual(Device.canonicalize("cl:1"), "CL:1")
|
||||
self.assertEqual(Device.canonicalize("CL:2"), "CL:2")
|
||||
self.assertEqual(Device.canonicalize("disk:/dev/shm/test"), "DISK:/dev/shm/test")
|
||||
self.assertEqual(Device.canonicalize("disk:000.txt"), "DISK:000.txt")
|
||||
|
||||
def test_getitem_not_exist(self):
|
||||
with self.assertRaises(ModuleNotFoundError):
|
||||
Device["TYPO"]
|
||||
|
||||
@unittest.skipIf(Device.DEFAULT != "CPU", "only run on CPU")
|
||||
def test_nonexistent_renderer(self):
|
||||
with self.assertRaisesRegex(RuntimeError, "has no renderer"):
|
||||
with Context(DEV="CPU:TYPO"): Device[Device.DEFAULT].renderer
|
||||
with self.assertRaisesRegex(RuntimeError, "did you mean: 'CLANG'"):
|
||||
with Context(DEV="CPU:CLANGJIT"): Device[Device.DEFAULT].renderer
|
||||
|
||||
@unittest.skipIf(Device.DEFAULT != "AMD", "only run on AMD")
|
||||
def test_nonexistent_iface(self):
|
||||
result = subprocess.run(['python3', '-c', 'from tinygrad import Device; Device[Device.DEFAULT].iface'],
|
||||
env={**os.environ, "DEV":"USA+AMD"}, capture_output=True)
|
||||
self.assertNotEqual(result.returncode, 0)
|
||||
self.assertIn(b"did you mean: 'USB'", result.stderr)
|
||||
|
||||
@unittest.skipIf(Device.DEFAULT != "AMD", "only run on AMD")
|
||||
def test_dev_id_out_of_range(self):
|
||||
result = subprocess.run(['python3', '-c', 'from tinygrad import Device; Device[Device.DEFAULT]'],
|
||||
env={**os.environ, "DEV":":99+AMD"}, capture_output=True)
|
||||
self.assertNotEqual(result.returncode, 0)
|
||||
self.assertIn(b"invalid visibility filter", result.stderr)
|
||||
|
||||
def test_lowercase_canonicalizes(self):
|
||||
device = Device.DEFAULT
|
||||
with Context(DEV=device.lower()):
|
||||
self.assertEqual(Device.canonicalize(None), device)
|
||||
|
||||
def test_set_device_default_raises(self):
|
||||
with self.assertRaisesRegex(AttributeError, "setting Device.DEFAULT is deprecated"):
|
||||
Device.DEFAULT = "CPU"
|
||||
|
||||
def test_old_device_env_raises(self):
|
||||
result = subprocess.run(['python3', '-c', 'from tinygrad import Device; Device.DEFAULT'],
|
||||
env={**os.environ, "CPU": "1", "DEV": ""}, capture_output=True)
|
||||
self.assertNotEqual(result.returncode, 0)
|
||||
self.assertIn(b"deprecated", result.stderr)
|
||||
|
||||
def test_old_renderer_env_raises(self):
|
||||
result = subprocess.run(['python3', '-c', 'from tinygrad import Device; Device[Device.DEFAULT].renderer'],
|
||||
env={**os.environ, "DEV": "CPU", "CPU_LLVM": "1"}, capture_output=True)
|
||||
self.assertNotEqual(result.returncode, 0)
|
||||
self.assertIn(b"deprecated", result.stderr)
|
||||
|
||||
@unittest.skipIf(WIN, "skipping windows test") # TODO: subprocess causes memory violation?
|
||||
def test_env_overwrite_default_compiler(self):
|
||||
if Device.DEFAULT == "CPU":
|
||||
from tinygrad.runtime.support.compiler_cpu import ClangCompiler
|
||||
from tinygrad.runtime.support.compiler_llvm import CPULLVMCompiler
|
||||
try: _, _ = CPULLVMCompiler(), ClangCompiler()
|
||||
except Exception as e: self.skipTest(f"skipping compiler test: not all compilers: {e}")
|
||||
|
||||
imports = ("from tinygrad import Device; from tinygrad.runtime.support.compiler_cpu import ClangCompiler; "
|
||||
"from tinygrad.runtime.support.compiler_llvm import CPULLVMCompiler")
|
||||
subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, CPULLVMCompiler)"'],
|
||||
shell=True, check=True, env={**os.environ, "DEV": "CPU:LLVM"})
|
||||
subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, ClangCompiler)"'],
|
||||
shell=True, check=True, env={**os.environ, "DEV": "CPU"})
|
||||
subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, ClangCompiler)"'],
|
||||
shell=True, check=True, env={**os.environ, "DEV": "CPU:CLANG"})
|
||||
elif Device.DEFAULT == "AMD":
|
||||
from tinygrad.runtime.support.compiler_amd import HIPCompiler
|
||||
from tinygrad.runtime.support.compiler_llvm import AMDLLVMCompiler
|
||||
try: _, _ = HIPCompiler(Device[Device.DEFAULT].arch), AMDLLVMCompiler(Device[Device.DEFAULT].arch)
|
||||
except Exception as e: self.skipTest(f"skipping compiler test: not all compilers: {e}")
|
||||
|
||||
imports = ("from tinygrad import Device; from tinygrad.runtime.support.compiler_amd import HIPCompiler; "
|
||||
"from tinygrad.runtime.support.compiler_amd import AMDLLVMCompiler")
|
||||
subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, AMDLLVMCompiler)"'],
|
||||
shell=True, check=True, env={**os.environ, "DEV": "AMD:LLVM"})
|
||||
subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, HIPCompiler)"'],
|
||||
shell=True, check=True, env={**os.environ, "DEV": "AMD"})
|
||||
subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, HIPCompiler)"'],
|
||||
shell=True, check=True, env={**os.environ, "DEV": "AMD:HIP"})
|
||||
else: self.skipTest("only run on CPU/AMD")
|
||||
|
||||
@unittest.skipIf(WIN, "skipping windows test")
|
||||
def test_env_online(self):
|
||||
from tinygrad.runtime.support.compiler_cpu import ClangCompiler
|
||||
from tinygrad.runtime.support.compiler_llvm import CPULLVMCompiler
|
||||
try: _, _ = CPULLVMCompiler(), ClangCompiler()
|
||||
except Exception as e: self.skipTest(f"skipping compiler test: not all compilers: {e}")
|
||||
|
||||
with Context(DEV="CPU:LLVM"):
|
||||
inst = Device["CPU"].compiler
|
||||
self.assertIsInstance(Device["CPU"].compiler, CPULLVMCompiler)
|
||||
with Context(DEV="CPU"):
|
||||
self.assertIsInstance(Device["CPU"].compiler, ClangCompiler)
|
||||
with Context(DEV="CPU:LLVM"):
|
||||
self.assertIsInstance(Device["CPU"].compiler, CPULLVMCompiler)
|
||||
assert inst is Device["CPU"].compiler # cached
|
||||
|
||||
@unittest.skipIf(Device.DEFAULT != "CPU", "only run on CPU")
|
||||
def test_compiler_autodetect_fallback(self):
|
||||
from tinygrad.runtime.support.compiler_llvm import CPULLVMCompiler
|
||||
|
||||
try: CPULLVMCompiler()
|
||||
except Exception as e: self.skipTest(f"skipping: LLVM not available: {e}")
|
||||
|
||||
dev = Device["CPU"]
|
||||
dev.cached_renderer.clear()
|
||||
with patch("tinygrad.renderer.cstyle.ClangRenderer.__init__", side_effect=RuntimeError("broken")):
|
||||
self.assertIsInstance(dev.renderer.compiler, CPULLVMCompiler)
|
||||
|
||||
def test_dev_contextvar(self):
|
||||
orig_dev = Device.DEFAULT
|
||||
with Context(DEV="CPU"): self.assertEqual(Tensor.empty(1).device, "CPU")
|
||||
with Context(DEV="NULL"): self.assertEqual(Tensor.empty(1).device, "NULL")
|
||||
self.assertEqual(Tensor.empty(1).device, orig_dev)
|
||||
|
||||
class TestDevVar(unittest.TestCase):
|
||||
def test_parse(self):
|
||||
for d, t in [("AMD", Target(device="AMD", renderer="")), ("AMD:LLVM", Target(device="AMD", renderer="LLVM")),
|
||||
(":LLVM", Target(device="", renderer="LLVM")), ("AMD::gfx1100", Target(device="AMD", arch="gfx1100")),
|
||||
("AMD:LLVM:gfx1100", Target(device="AMD", renderer="LLVM", arch="gfx1100")), ("::gfx1100", Target(arch="gfx1100")),
|
||||
("USB+", Target(interface="USB")), ("USB+AMD", Target(device="AMD", interface="USB")),
|
||||
("PCI:0+AMD", Target(device="AMD", interface="PCI", indices="0")), (":0+AMD", Target(device="AMD", indices="0")),
|
||||
("PCI:0,1+AMD", Target(device="AMD", interface="PCI", indices="0,1")),
|
||||
("QCOM;USB+AMD", [Target(device="QCOM"), Target(device="AMD", interface="USB")])]:
|
||||
with Context(DEV=d):
|
||||
self.assertEqual(DEV.value, t if isinstance(t, list) else [t])
|
||||
self.assertEqual(str(DEV), d)
|
||||
|
||||
def test_target(self):
|
||||
with Context(DEV="CPU"): self.assertEqual(DEV.target("CPU"), Target("CPU"))
|
||||
with Context(DEV="CPU:LLVM"): self.assertEqual(DEV.target("CPU"), Target("CPU", "LLVM"))
|
||||
with Context(DEV=":LLVM"): self.assertEqual(DEV.target("CPU"), Target("CPU", "LLVM"))
|
||||
with Context(DEV="AMD:LLVM"): self.assertEqual(DEV.target("CPU"), Target("CPU"))
|
||||
with Context(DEV=""): self.assertEqual(DEV.target("CPU"), Target("CPU"))
|
||||
with Context(DEV="QCOM:IR3;AMD:LLVM"):
|
||||
self.assertEqual(DEV.target("QCOM"), Target("QCOM", "IR3"))
|
||||
self.assertEqual(DEV.target("AMD"), Target("AMD", "LLVM"))
|
||||
self.assertEqual(DEV.target("CPU"), Target("CPU"))
|
||||
|
||||
def test_dev_arch_override(self):
|
||||
with Context(DEV="NULL::gfx1100"):
|
||||
self.assertEqual(Device["NULL"].renderer.target.arch, "gfx1100")
|
||||
|
||||
class MockCompiler(Compiler):
|
||||
def __init__(self, key): super().__init__(key)
|
||||
def compile(self, src) -> bytes: return src.encode()
|
||||
|
||||
class TestCompiler(unittest.TestCase):
|
||||
def test_compile_cached(self):
|
||||
diskcache_put("key", "123", None) # clear cache
|
||||
getenv.cache_clear()
|
||||
with Context(CCACHE=1):
|
||||
self.assertEqual(MockCompiler("key").compile_cached("123"), str.encode("123"))
|
||||
self.assertEqual(diskcache_get("key", "123"), str.encode("123"))
|
||||
|
||||
def test_compile_cached_disabled(self):
|
||||
diskcache_put("disabled_key", "123", None) # clear cache
|
||||
getenv.cache_clear()
|
||||
with Context(CCACHE=0):
|
||||
self.assertEqual(MockCompiler("disabled_key").compile_cached("123"), str.encode("123"))
|
||||
self.assertIsNone(diskcache_get("disabled_key", "123"))
|
||||
|
||||
def test_device_compile(self):
|
||||
getenv.cache_clear()
|
||||
with Context(CCACHE=0):
|
||||
a = Tensor([0.,1.], device=Device.DEFAULT).realize()
|
||||
(a + 1).realize()
|
||||
|
||||
@unittest.skip("this test is broken if you have tinymesa installed")
|
||||
@unittest.skipIf(OSX and 'libclang' in DLL._loaded_, "MTLCompiler can't be loaded after libclang on OSX")
|
||||
class TestRunAsModule(unittest.TestCase):
|
||||
def test_module_runs(self):
|
||||
cpu_line = [l for l in enumerate_devices_str() if "CPU" in l][0]
|
||||
self.assertIn("PASS", cpu_line, f"expected CPU to PASS, got: {cpu_line}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
108
artifacts/package_sources/tinygrad/test/null/test_disk_cache.py
Normal file
108
artifacts/package_sources/tinygrad/test/null/test_disk_cache.py
Normal file
@@ -0,0 +1,108 @@
|
||||
import unittest
|
||||
import pickle
|
||||
from tinygrad.helpers import diskcache_get, diskcache_put, diskcache, diskcache_clear
|
||||
|
||||
def remote_get(table,q,k): q.put(diskcache_get(table, k))
|
||||
def remote_put(table,k,v): diskcache_put(table, k, v)
|
||||
|
||||
class DiskCache(unittest.TestCase):
|
||||
def test_putget(self):
|
||||
table = "test_putget"
|
||||
diskcache_put(table, "hello", "world")
|
||||
self.assertEqual(diskcache_get(table, "hello"), "world")
|
||||
diskcache_put(table, "hello", "world2")
|
||||
self.assertEqual(diskcache_get(table, "hello"), "world2")
|
||||
|
||||
def test_putcomplex(self):
|
||||
table = "test_putcomplex"
|
||||
diskcache_put(table, "k", ("complex", 123, "object"))
|
||||
ret = diskcache_get(table, "k")
|
||||
self.assertEqual(ret, ("complex", 123, "object"))
|
||||
|
||||
def test_getotherprocess(self):
|
||||
table = "test_getotherprocess"
|
||||
from multiprocessing import Process, Queue
|
||||
diskcache_put(table, "k", "getme")
|
||||
q = Queue()
|
||||
p = Process(target=remote_get, args=(table,q,"k"))
|
||||
p.start()
|
||||
p.join()
|
||||
self.assertEqual(q.get(), "getme")
|
||||
|
||||
def test_putotherprocess(self):
|
||||
table = "test_putotherprocess"
|
||||
from multiprocessing import Process
|
||||
p = Process(target=remote_put, args=(table,"k", "remote"))
|
||||
p.start()
|
||||
p.join()
|
||||
self.assertEqual(diskcache_get(table, "k"), "remote")
|
||||
|
||||
def test_no_table(self):
|
||||
self.assertIsNone(diskcache_get("faketable", "k"))
|
||||
|
||||
def test_ret(self):
|
||||
table = "test_ret"
|
||||
self.assertEqual(diskcache_put(table, "key", ("vvs",)), ("vvs",))
|
||||
|
||||
def test_non_str_key(self):
|
||||
table = "test_non_str_key"
|
||||
diskcache_put(table, 4, 5)
|
||||
self.assertEqual(diskcache_get(table, 4), 5)
|
||||
self.assertEqual(diskcache_get(table, "4"), 5)
|
||||
|
||||
def test_decorator(self):
|
||||
calls = 0
|
||||
@diskcache
|
||||
def hello(x):
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
return "world"+x
|
||||
self.assertEqual(hello("bob"), "worldbob")
|
||||
self.assertEqual(hello("billy"), "worldbilly")
|
||||
kcalls = calls
|
||||
self.assertEqual(hello("bob"), "worldbob")
|
||||
self.assertEqual(hello("billy"), "worldbilly")
|
||||
self.assertEqual(kcalls, calls)
|
||||
|
||||
def test_dict_key(self):
|
||||
table = "test_dict_key"
|
||||
fancy_key = {"hello": "world", "goodbye": 7, "good": True, "pkl": pickle.dumps("cat")}
|
||||
fancy_key2 = {"hello": "world", "goodbye": 8, "good": True, "pkl": pickle.dumps("cat")}
|
||||
fancy_key3 = {"hello": "world", "goodbye": 8, "good": True, "pkl": pickle.dumps("dog")}
|
||||
diskcache_put(table, fancy_key, 5)
|
||||
self.assertEqual(diskcache_get(table, fancy_key), 5)
|
||||
diskcache_put(table, fancy_key2, 8)
|
||||
self.assertEqual(diskcache_get(table, fancy_key2), 8)
|
||||
self.assertEqual(diskcache_get(table, fancy_key), 5)
|
||||
self.assertEqual(diskcache_get(table, fancy_key3), None)
|
||||
|
||||
def test_table_name(self):
|
||||
table = "test_gfx1010:xnack-"
|
||||
diskcache_put(table, "key", "test")
|
||||
self.assertEqual(diskcache_get(table, "key"), "test")
|
||||
|
||||
@unittest.skip("disabled by default because this drops cache table")
|
||||
def test_clear_cache(self):
|
||||
# clear cache to start
|
||||
diskcache_clear()
|
||||
tables = [f"test_clear_cache:{i}" for i in range(3)]
|
||||
for table in tables:
|
||||
# check no entries
|
||||
self.assertIsNone(diskcache_get(table, "k"))
|
||||
for table in tables:
|
||||
diskcache_put(table, "k", "test")
|
||||
# check insertion
|
||||
self.assertEqual(diskcache_get(table, "k"), "test")
|
||||
|
||||
diskcache_clear()
|
||||
for table in tables:
|
||||
# check no entries again
|
||||
self.assertIsNone(diskcache_get(table, "k"))
|
||||
|
||||
# calling multiple times is fine
|
||||
diskcache_clear()
|
||||
diskcache_clear()
|
||||
diskcache_clear()
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
69
artifacts/package_sources/tinygrad/test/null/test_dtype.py
Normal file
69
artifacts/package_sources/tinygrad/test/null/test_dtype.py
Normal file
@@ -0,0 +1,69 @@
|
||||
import unittest, pickle
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.dtype import dtypes, DType, to_dtype, Invalid, InvalidType
|
||||
|
||||
class TestEqStrDType(unittest.TestCase):
|
||||
def test_strs(self):
|
||||
self.assertEqual(str(dtypes.float32), "dtypes.float")
|
||||
|
||||
class TestToDtype(unittest.TestCase):
|
||||
def test_dtype_to_dtype(self):
|
||||
dtype = dtypes.int32
|
||||
res = to_dtype(dtype)
|
||||
self.assertIsInstance(res, DType)
|
||||
self.assertEqual(res, dtypes.int32)
|
||||
|
||||
def test_str_to_dtype(self):
|
||||
dtype = "int32"
|
||||
res = to_dtype(dtype)
|
||||
self.assertIsInstance(res, DType)
|
||||
self.assertEqual(res, dtypes.int32)
|
||||
|
||||
class TestCastConvenienceMethod(unittest.TestCase):
|
||||
def test_method(self):
|
||||
for input_dtype in (dtypes.float, dtypes.int):
|
||||
t = Tensor([1, 2], dtype=input_dtype)
|
||||
self.assertEqual(t.dtype, input_dtype)
|
||||
self.assertEqual(t.bool().dtype, dtypes.bool)
|
||||
self.assertEqual(t.short().dtype, dtypes.short)
|
||||
self.assertEqual(t.int().dtype, dtypes.int)
|
||||
self.assertEqual(t.long().dtype, dtypes.long)
|
||||
self.assertEqual(t.half().dtype, dtypes.half)
|
||||
self.assertEqual(t.bfloat16().dtype, dtypes.bfloat16)
|
||||
self.assertEqual(t.float().dtype, dtypes.float)
|
||||
self.assertEqual(t.double().dtype, dtypes.double)
|
||||
|
||||
class TestDtypeTolist(unittest.TestCase):
|
||||
def test_bfloat16(self):
|
||||
self.assertEqual(Tensor([-60000, 1.5, 3.1, 60000], device="PYTHON", dtype=dtypes.bfloat16).tolist(), [-59904.0, 1.5, 3.09375, 59904.0])
|
||||
def test_fp8(self):
|
||||
# 448
|
||||
self.assertEqual(Tensor([-30000, 1.5, 3.1, 30000], device="PYTHON", dtype=dtypes.fp8e4m3).tolist(), [-448.0, 1.5, 3.0, 448.0])
|
||||
# 57344
|
||||
self.assertEqual(Tensor([-30000, 1.5, 3.1, 30000], device="PYTHON", dtype=dtypes.fp8e5m2).tolist(), [-28672.0, 1.5, 3.0, 28672.0])
|
||||
|
||||
class TestCanLosslessCast(unittest.TestCase):
|
||||
def test_can_lossless_cast(self):
|
||||
from tinygrad.dtype import can_lossless_cast
|
||||
# signed -> unsigned is NOT lossless (negative values wrap)
|
||||
self.assertFalse(can_lossless_cast(dtypes.int8, dtypes.uint64))
|
||||
self.assertFalse(can_lossless_cast(dtypes.int32, dtypes.uint32))
|
||||
# unsigned -> larger signed is lossless
|
||||
self.assertTrue(can_lossless_cast(dtypes.uint8, dtypes.int16))
|
||||
self.assertTrue(can_lossless_cast(dtypes.uint32, dtypes.int64))
|
||||
# large ints don't fit in floats
|
||||
self.assertFalse(can_lossless_cast(dtypes.int32, dtypes.float))
|
||||
self.assertFalse(can_lossless_cast(dtypes.int64, dtypes.double))
|
||||
# half has more mantissa bits
|
||||
self.assertTrue(can_lossless_cast(dtypes.int8, dtypes.half))
|
||||
self.assertFalse(can_lossless_cast(dtypes.int8, dtypes.bfloat16))
|
||||
|
||||
class TestInvalidSingleton(unittest.TestCase):
|
||||
def test_singleton(self):
|
||||
self.assertIs(InvalidType(), InvalidType())
|
||||
self.assertIs(InvalidType(), Invalid)
|
||||
def test_pickle(self):
|
||||
self.assertIs(pickle.loads(pickle.dumps(Invalid)), Invalid)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
451
artifacts/package_sources/tinygrad/test/null/test_dtype_spec.py
Normal file
451
artifacts/package_sources/tinygrad/test/null/test_dtype_spec.py
Normal file
@@ -0,0 +1,451 @@
|
||||
import unittest, math, struct, operator
|
||||
from tinygrad import Tensor, Device
|
||||
from tinygrad.dtype import DTYPES_DICT, dtypes, Invalid, truncate, float_to_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float
|
||||
|
||||
from tinygrad.helpers import getenv, Context
|
||||
from hypothesis import given, settings, strategies as strat
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
settings.register_profile("my_profile", max_examples=50, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
|
||||
settings.load_profile("my_profile")
|
||||
|
||||
core_dtypes = list(DTYPES_DICT.values())
|
||||
dtype_ints = [dt for dt in core_dtypes if dtypes.is_int(dt) and dt in Device[Device.DEFAULT].renderer.supported_dtypes()]
|
||||
dtype_floats = [dt for dt in core_dtypes if dtypes.is_float(dt) and dt in Device[Device.DEFAULT].renderer.supported_dtypes()]
|
||||
|
||||
FP8E4M3_MAX = 448.0
|
||||
FP8E5M2_MAX = 57344.0
|
||||
|
||||
def u32_to_f32(u): return struct.unpack('f', struct.pack('I', u))[0]
|
||||
def f32_to_u32(f): return struct.unpack('I', struct.pack('f', f))[0]
|
||||
|
||||
class TestHelpers(unittest.TestCase):
|
||||
signed_ints = (dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64)
|
||||
uints = (dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64)
|
||||
floats = (dtypes.float16, dtypes.float32, dtypes.float64)
|
||||
|
||||
@given(strat.sampled_from(signed_ints+uints))
|
||||
def test_is_int(self, dtype):
|
||||
assert dtypes.is_int(dtype)
|
||||
assert not dtypes.is_float(dtype)
|
||||
|
||||
@given(strat.sampled_from(uints))
|
||||
def test_is_unsigned_uints(self, dtype):
|
||||
assert dtypes.is_unsigned(dtype)
|
||||
|
||||
@given(strat.sampled_from(signed_ints))
|
||||
def test_is_unsigned_signed_ints(self, dtype):
|
||||
assert not dtypes.is_unsigned(dtype)
|
||||
|
||||
@given(strat.sampled_from(floats))
|
||||
def test_is_float(self, dtype):
|
||||
assert dtypes.is_float(dtype)
|
||||
assert not dtypes.is_int(dtype)
|
||||
assert not dtypes.is_unsigned(dtype)
|
||||
|
||||
def test_bf16_is_float(self):
|
||||
assert dtypes.is_float(dtypes.bfloat16)
|
||||
|
||||
def test_fp8s_are_float(self):
|
||||
assert dtypes.is_float(dtypes.fp8e4m3)
|
||||
assert dtypes.is_float(dtypes.fp8e5m2)
|
||||
|
||||
@given(strat.sampled_from([d for d in DTYPES_DICT.values() if dtypes.is_float(d) or dtypes.is_int(d)]))
|
||||
def test_scalar(self, dtype):
|
||||
assert dtype.scalar() == dtype
|
||||
|
||||
def test_from_py(self):
|
||||
assert dtypes.from_py(True) == dtypes.bool
|
||||
assert dtypes.from_py(Invalid) == dtypes.bool
|
||||
assert dtypes.from_py(2) == dtypes.weakint
|
||||
assert dtypes.from_py(3.0) == dtypes.weakfloat
|
||||
assert dtypes.from_py([]) == dtypes.default_float
|
||||
assert dtypes.from_py(()) == dtypes.default_float
|
||||
assert dtypes.from_py([True]) == dtypes.bool
|
||||
assert dtypes.from_py([True, 2]) == dtypes.default_int
|
||||
assert dtypes.from_py([True, 3.0]) == dtypes.default_float
|
||||
assert dtypes.from_py([2, 3.0]) == dtypes.default_float
|
||||
assert dtypes.from_py([True, 2, 3.0]) == dtypes.default_float
|
||||
with self.assertRaises(RuntimeError): dtypes.from_py(None)
|
||||
with self.assertRaises(RuntimeError): dtypes.from_py([None])
|
||||
with self.assertRaises(RuntimeError): dtypes.from_py({})
|
||||
with self.assertRaises(RuntimeError): dtypes.from_py(set())
|
||||
|
||||
def test_dtype_range(self):
|
||||
for dt in core_dtypes:
|
||||
if dtypes.is_float(dt):
|
||||
np.testing.assert_equal(dt.min, -math.inf)
|
||||
np.testing.assert_equal(dt.max, math.inf)
|
||||
np.testing.assert_equal(dt.min, -math.inf)
|
||||
np.testing.assert_equal(dt.max, math.inf)
|
||||
elif dtypes.is_int(dt):
|
||||
info = np.iinfo(_to_np_dtype(dt))
|
||||
np.testing.assert_equal(dt.min, info.min)
|
||||
np.testing.assert_equal(dt.max, info.max)
|
||||
np.testing.assert_equal(dt.min, info.min)
|
||||
np.testing.assert_equal(dt.max, info.max)
|
||||
else:
|
||||
assert dt == dtypes.bool, dt
|
||||
np.testing.assert_equal(dt.min, False)
|
||||
np.testing.assert_equal(dt.max, True)
|
||||
np.testing.assert_equal(dt.min, False)
|
||||
np.testing.assert_equal(dt.max, True)
|
||||
|
||||
def test_dtype_range_vec(self):
|
||||
for dt in core_dtypes:
|
||||
self.assertEqual(dt.min, dt.min)
|
||||
self.assertEqual(dt.max, dt.max)
|
||||
|
||||
def test_float_to_fp16(self):
|
||||
self.assertEqual(float_to_fp16(1), 1)
|
||||
self.assertEqual(float_to_fp16(65504), 65504)
|
||||
self.assertEqual(float_to_fp16(65519.999), 65504)
|
||||
self.assertEqual(float_to_fp16(65520), math.inf)
|
||||
self.assertEqual(float_to_fp16(1e-8), 0.0)
|
||||
self.assertEqual(float_to_fp16(-65504), -65504)
|
||||
self.assertEqual(float_to_fp16(-65519.999), -65504)
|
||||
self.assertEqual(float_to_fp16(-65520), -math.inf)
|
||||
self.assertTrue(math.isnan(float_to_fp16(math.nan)))
|
||||
|
||||
def test_float_to_bf16(self):
|
||||
max_bf16 = torch.finfo(torch.bfloat16).max
|
||||
for a in [1, 1.1, 1234, 23456, -777.777, max_bf16, max_bf16 * 1.00001, -max_bf16, -max_bf16 * 1.00001, math.inf, -math.inf]:
|
||||
self.assertEqual(float_to_bf16(a), torch.tensor([a], dtype=torch.bfloat16).item())
|
||||
self.assertTrue(math.isnan(float_to_bf16(math.nan)))
|
||||
|
||||
def test_float_to_bf16_nan(self):
|
||||
patterns = [0x7FC00001, 0xFFC00001, 0x7F800001, 0xFF800001, 0x7FFFFFFF, 0xFFFFFFFF]
|
||||
for u in patterns:
|
||||
x = u32_to_f32(u)
|
||||
y = float_to_bf16(x)
|
||||
t = torch.tensor([x], dtype=torch.bfloat16).item()
|
||||
self.assertTrue(math.isnan(y))
|
||||
self.assertTrue(math.isnan(t))
|
||||
|
||||
def test_float_to_bf16_round(self):
|
||||
uppers = [0x3f800000, 0x41230000, 0xC1460000]
|
||||
for upper in uppers:
|
||||
base = upper & 0xFFFF0000
|
||||
base_f32 = u32_to_f32(base)
|
||||
base_f32_round_up = u32_to_f32(base + 0x00010000)
|
||||
|
||||
x = u32_to_f32(base | 0x00007000)
|
||||
self.assertEqual(float_to_bf16(x), base_f32)
|
||||
self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32)
|
||||
|
||||
x = u32_to_f32(base | 0x0000C000)
|
||||
self.assertEqual(float_to_bf16(x), base_f32_round_up)
|
||||
self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32_round_up)
|
||||
|
||||
if ((upper >> 16) & 1) == 0:
|
||||
x = u32_to_f32(base | 0x00008000)
|
||||
self.assertEqual(float_to_bf16(x), base_f32)
|
||||
self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32)
|
||||
else:
|
||||
x = u32_to_f32(base | 0x00008000)
|
||||
self.assertEqual(float_to_bf16(x), base_f32_round_up)
|
||||
self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32_round_up)
|
||||
|
||||
def test_float_to_bf16_boundary(self):
|
||||
base = 0x7F7F0000
|
||||
inf_u32 = 0x7F800000
|
||||
|
||||
x = u32_to_f32(base | 0x00007FFF)
|
||||
self.assertEqual(f32_to_u32(float_to_bf16(x)), base)
|
||||
self.assertEqual(f32_to_u32(torch.tensor([x], dtype=torch.bfloat16).item()), base)
|
||||
|
||||
x = u32_to_f32(base | 0x0000C000)
|
||||
self.assertEqual(f32_to_u32(float_to_bf16(x)), inf_u32)
|
||||
self.assertEqual(f32_to_u32(torch.tensor([x], dtype=torch.bfloat16).item()), inf_u32)
|
||||
|
||||
x = u32_to_f32(base | 0x00008000)
|
||||
self.assertEqual(f32_to_u32(float_to_bf16(x)), inf_u32)
|
||||
self.assertEqual(f32_to_u32(torch.tensor([x], dtype=torch.bfloat16).item()), inf_u32)
|
||||
|
||||
@given(strat.floats(width=32, allow_subnormal=True, allow_nan=True, allow_infinity=True))
|
||||
def test_truncate_fp8e4m3(self, x):
|
||||
if math.isnan(x): np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), x)
|
||||
elif math.isinf(x): np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), math.copysign(math.nan, x))
|
||||
elif x > FP8E4M3_MAX: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), FP8E4M3_MAX)
|
||||
elif x < -FP8E4M3_MAX: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), -FP8E4M3_MAX)
|
||||
else: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), torch.tensor(x, dtype=torch.float8_e4m3fn).float().item())
|
||||
|
||||
@given(strat.floats(width=32, allow_subnormal=True, allow_nan=True, allow_infinity=True))
|
||||
def test_truncate_fp8e5m2(self, x):
|
||||
if math.isnan(x): np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), x)
|
||||
elif math.isinf(x): np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), x)
|
||||
elif x > FP8E5M2_MAX: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), FP8E5M2_MAX)
|
||||
elif x < -FP8E5M2_MAX: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), -FP8E5M2_MAX)
|
||||
else: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), torch.tensor(x, dtype=torch.float8_e5m2).float().item())
|
||||
|
||||
def test_finfo(self):
|
||||
for dt in [dtypes.float16, dtypes.float32, dtypes.float64]:
|
||||
info = np.finfo(_to_np_dtype(dt))
|
||||
self.assertEqual(info.bits, dt.bitsize)
|
||||
self.assertEqual((info.nexp, info.nmant), dtypes.finfo(dt))
|
||||
|
||||
class TestTypePromotion(unittest.TestCase):
|
||||
@given(strat.sampled_from(core_dtypes))
|
||||
def test_self_promo_to_self(self, dtype):
|
||||
assert least_upper_dtype(dtype) == dtype
|
||||
assert least_upper_dtype(dtype, dtype) == dtype
|
||||
assert least_upper_dtype(dtype, dtype, dtype) == dtype
|
||||
|
||||
@given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
|
||||
def test_promo_resulted_higher_than_inputs(self, dtype1, dtype2):
|
||||
result = least_upper_dtype(dtype1, dtype2)
|
||||
assert not (result < dtype1) and not (result < dtype2)
|
||||
|
||||
def test_dtype_promo(self):
|
||||
assert least_upper_dtype(dtypes.bool, dtypes.int8) == dtypes.int8
|
||||
assert least_upper_dtype(dtypes.int8, dtypes.uint8) == dtypes.int16
|
||||
assert least_upper_dtype(dtypes.uint8, dtypes.int16) == dtypes.int16
|
||||
assert least_upper_dtype(dtypes.int16, dtypes.uint16) == dtypes.int32
|
||||
assert least_upper_dtype(dtypes.uint16, dtypes.int32) == dtypes.int32
|
||||
assert least_upper_dtype(dtypes.int32, dtypes.uint32) == dtypes.int64
|
||||
assert least_upper_dtype(dtypes.uint32, dtypes.int64) == dtypes.int64
|
||||
# uint64 has no common integer supertype with any signed int (JAX JEP), they all defer up to weakfloat
|
||||
for st in dtypes.sints: assert least_upper_dtype(st, dtypes.uint64) == dtypes.weakfloat
|
||||
assert least_upper_dtype(dtypes.float16, dtypes.float32) == dtypes.float32
|
||||
assert least_upper_dtype(dtypes.float32, dtypes.float64) == dtypes.float64
|
||||
|
||||
assert least_upper_dtype(dtypes.bool, dtypes.float32) == dtypes.float32
|
||||
assert least_upper_dtype(dtypes.bool, dtypes.float64) == dtypes.float64
|
||||
assert least_upper_dtype(dtypes.float16, dtypes.int64) == dtypes.float16
|
||||
assert least_upper_dtype(dtypes.float16, dtypes.uint64) == dtypes.float16
|
||||
assert least_upper_dtype(dtypes.fp8e4m3, dtypes.fp8e5m2) == dtypes.half
|
||||
assert least_upper_dtype(dtypes.fp8e4m3, dtypes.bfloat16) == dtypes.bfloat16
|
||||
assert least_upper_dtype(dtypes.fp8e5m2, dtypes.bfloat16) == dtypes.bfloat16
|
||||
assert least_upper_dtype(dtypes.fp8e4m3, dtypes.float16) == dtypes.float16
|
||||
assert least_upper_dtype(dtypes.fp8e5m2, dtypes.float16) == dtypes.float16
|
||||
assert least_upper_dtype(dtypes.fp8e4m3, dtypes.int64) == dtypes.fp8e4m3
|
||||
assert least_upper_dtype(dtypes.fp8e4m3, dtypes.uint64) == dtypes.fp8e4m3
|
||||
assert least_upper_dtype(dtypes.fp8e5m2, dtypes.int64) == dtypes.fp8e5m2
|
||||
assert least_upper_dtype(dtypes.fp8e5m2, dtypes.uint64) == dtypes.fp8e5m2
|
||||
|
||||
def test_weakint_promo(self):
|
||||
assert least_upper_dtype(dtypes.weakint, dtypes.weakint) == dtypes.weakint
|
||||
assert least_upper_dtype(dtypes.bool, dtypes.weakint) == dtypes.weakint
|
||||
assert least_upper_dtype(dtypes.weakint, dtypes.int8) == dtypes.int8
|
||||
|
||||
def test_weakfloat_promo(self):
|
||||
# weakfloat is a float, but is not one of dtypes.floats
|
||||
assert dtypes.is_float(dtypes.weakfloat) and dtypes.weakfloat not in dtypes.floats
|
||||
# weakfloat with itself is weakfloat
|
||||
assert least_upper_dtype(dtypes.weakfloat, dtypes.weakfloat) == dtypes.weakfloat
|
||||
# weakfloat is above bool and any concrete int (they defer up to it)
|
||||
assert least_upper_dtype(dtypes.weakfloat, dtypes.bool) == dtypes.weakfloat
|
||||
assert least_upper_dtype(dtypes.weakfloat, dtypes.int32) == dtypes.weakfloat
|
||||
assert least_upper_dtype(dtypes.weakfloat, dtypes.uint64) == dtypes.weakfloat
|
||||
# weakfloat defers to any concrete float type
|
||||
assert least_upper_dtype(dtypes.weakfloat, dtypes.fp8e4m3) == dtypes.fp8e4m3
|
||||
assert least_upper_dtype(dtypes.weakfloat, dtypes.float16) == dtypes.float16
|
||||
assert least_upper_dtype(dtypes.weakfloat, dtypes.float32) == dtypes.float32
|
||||
assert least_upper_dtype(dtypes.weakfloat, dtypes.float64) == dtypes.float64
|
||||
|
||||
class TestTypeSpec(unittest.TestCase):
|
||||
def test_set_dtype_default(self):
|
||||
for default_int in [dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64]:
|
||||
with Context(DEFAULT_INT=default_int):
|
||||
assert dtypes.default_int == default_int
|
||||
|
||||
for default_float in [*dtypes.fp8s, dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]:
|
||||
with Context(DEFAULT_FLOAT=default_float):
|
||||
assert dtypes.default_float == default_float
|
||||
|
||||
@given(strat.sampled_from(core_dtypes), strat.sampled_from([operator.gt, operator.ge, operator.le, operator.lt, operator.eq, operator.ne]))
|
||||
def test_bool_ops(self, dtype, op):
|
||||
assert op(Tensor.ones(4, 4, dtype=dtype), Tensor.ones(4, 4, dtype=dtype)).dtype == dtypes.bool
|
||||
|
||||
@given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
|
||||
def test_functions_return_index(self, dtype, default_int, default_float):
|
||||
self.enterContext(Context(DEFAULT_INT=default_int, DEFAULT_FLOAT=default_float))
|
||||
assert Tensor([0, 1], dtype=dtype).argmax().dtype == dtypes.int32
|
||||
assert Tensor([0, 1], dtype=dtype).argmin().dtype == dtypes.int32
|
||||
assert Tensor([0, 1], dtype=dtype).multinomial().dtype == dtypes.int32
|
||||
|
||||
@given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints))
|
||||
def test_tensor_indexing_returns_same_dtype(self, data_dtype, indices_dtype):
|
||||
X_data = Tensor.ones(60000, 1, 28, 28, dtype=data_dtype)
|
||||
indices = Tensor.randint(512, high=X_data.shape[0]).cast(indices_dtype)
|
||||
assert X_data[indices].dtype == X_data.dtype
|
||||
|
||||
@given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints))
|
||||
def test_gather_returns_same_dtype(self, data_dtype, indices_dtype):
|
||||
X_data = Tensor([[1, 0], [0, 1]], dtype=data_dtype)
|
||||
indices = Tensor([[0, 0], [1, 0]], dtype=indices_dtype)
|
||||
assert X_data.gather(0, indices).dtype == X_data.dtype
|
||||
assert X_data.gather(1, indices).dtype == X_data.dtype
|
||||
|
||||
@given(strat.sampled_from(dtype_floats), strat.sampled_from(dtype_floats))
|
||||
def test_attention_returns_same_dtype(self, data_dtype, default_float):
|
||||
self.enterContext(Context(DEFAULT_FLOAT=default_float))
|
||||
query = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
|
||||
key = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
|
||||
value = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
|
||||
mask = (Tensor.rand(32, 8, 128, 128) < 0.5)
|
||||
assert query.scaled_dot_product_attention(key, value, is_causal=True).dtype == data_dtype
|
||||
assert query.scaled_dot_product_attention(key, value, is_causal=True, dropout_p=0.3).dtype == data_dtype
|
||||
assert query.scaled_dot_product_attention(key, value, is_causal=False).dtype == data_dtype
|
||||
assert query.scaled_dot_product_attention(key, value, attn_mask=mask).dtype == data_dtype
|
||||
|
||||
class TestAutoCastType(unittest.TestCase):
|
||||
@given(strat.sampled_from(dtype_floats), strat.sampled_from(dtype_floats))
|
||||
def test_least_upper_float_input_is_float(self, input_dtype, default_float):
|
||||
self.enterContext(Context(DEFAULT_FLOAT=default_float))
|
||||
self.assertEqual(least_upper_float(input_dtype), input_dtype)
|
||||
|
||||
@given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
|
||||
def test_least_upper_float_input_is_int(self, input_dtype, default_float):
|
||||
self.enterContext(Context(DEFAULT_FLOAT=default_float))
|
||||
self.assertEqual(least_upper_float(input_dtype), default_float)
|
||||
|
||||
@given(strat.sampled_from(core_dtypes))
|
||||
def test_broadcast_scalar(self, dt):
|
||||
assert (Tensor.ones(4, 4, dtype=dt) + 2.3).dtype == (dt if dtypes.is_float(dt) else dtypes.weakfloat)
|
||||
assert (Tensor.ones(4, 4, dtype=dt) + 2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.weakint)
|
||||
assert (Tensor.ones(4, 4, dtype=dt) + True).dtype == dt
|
||||
|
||||
@given(strat.sampled_from(core_dtypes))
|
||||
def test_pad_scalar(self, dt):
|
||||
t = Tensor.ones(4, dtype=dt)
|
||||
assert t.pad(((1, 1),), value=2.3).dtype == (dt if dtypes.is_float(dt) else dtypes.weakfloat)
|
||||
assert t.pad(((1, 1),), value=2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.weakint)
|
||||
assert t.pad(((1, 1),), value=True).dtype == dt
|
||||
|
||||
@given(strat.sampled_from(core_dtypes))
|
||||
def test_sort(self, dt):
|
||||
# sort pads with dtype.min/max, a scalar of its own dtype
|
||||
assert Tensor([3, 1, 2], dtype=dt).sort()[0].dtype == dt
|
||||
assert Tensor([3, 1, 2], dtype=dt).sort(descending=True)[0].dtype == dt
|
||||
|
||||
@given(strat.sampled_from(dtype_floats))
|
||||
def test_int_div_int(self, default_float):
|
||||
self.enterContext(Context(DEFAULT_FLOAT=default_float))
|
||||
self.assertEqual(Tensor([1]).div(Tensor([2])).dtype, default_float)
|
||||
|
||||
def test_sum(self):
|
||||
assert (Tensor([0, 1], dtype=dtypes.bool)).sum().dtype == dtypes.int32
|
||||
assert (Tensor([0, 1], dtype=dtypes.int8)).sum().dtype == dtypes.int32
|
||||
assert (Tensor([0, 1], dtype=dtypes.int16)).sum().dtype == dtypes.int32
|
||||
assert (Tensor([0, 1], dtype=dtypes.int32)).sum().dtype == dtypes.int32
|
||||
assert (Tensor([0, 1], dtype=dtypes.int64)).sum().dtype == dtypes.int64
|
||||
assert (Tensor([0, 1], dtype=dtypes.uint8)).sum().dtype == dtypes.uint32
|
||||
assert (Tensor([0, 1], dtype=dtypes.uint16)).sum().dtype == dtypes.uint32
|
||||
assert (Tensor([0, 1], dtype=dtypes.uint32)).sum().dtype == dtypes.uint32
|
||||
assert (Tensor([0, 1], dtype=dtypes.uint64)).sum().dtype == dtypes.uint64
|
||||
assert (Tensor([0, 1], dtype=dtypes.fp8e4m3)).sum().dtype == dtypes.fp8e4m3
|
||||
assert (Tensor([0, 1], dtype=dtypes.fp8e5m2)).sum().dtype == dtypes.fp8e5m2
|
||||
assert (Tensor([0, 1], dtype=dtypes.float16)).sum().dtype == dtypes.float16
|
||||
assert (Tensor([0, 1], dtype=dtypes.bfloat16)).sum().dtype == dtypes.bfloat16
|
||||
assert (Tensor([0, 1], dtype=dtypes.float32)).sum().dtype == dtypes.float32
|
||||
assert (Tensor([0, 1], dtype=dtypes.float64)).sum().dtype == dtypes.float64
|
||||
|
||||
def test_mean(self):
|
||||
assert (Tensor([0, 1], dtype=dtypes.bool)).mean().dtype == dtypes.float32
|
||||
assert (Tensor([0, 1], dtype=dtypes.int8)).mean().dtype == dtypes.float32
|
||||
assert (Tensor([0, 1], dtype=dtypes.int16)).mean().dtype == dtypes.float32
|
||||
assert (Tensor([0, 1], dtype=dtypes.int32)).mean().dtype == dtypes.float32
|
||||
assert (Tensor([0, 1], dtype=dtypes.int64)).mean().dtype == dtypes.float32
|
||||
assert (Tensor([0, 1], dtype=dtypes.uint8)).mean().dtype == dtypes.float32
|
||||
assert (Tensor([0, 1], dtype=dtypes.uint16)).mean().dtype == dtypes.float32
|
||||
assert (Tensor([0, 1], dtype=dtypes.uint32)).mean().dtype == dtypes.float32
|
||||
assert (Tensor([0, 1], dtype=dtypes.uint64)).mean().dtype == dtypes.float32
|
||||
assert (Tensor([0, 1], dtype=dtypes.fp8e4m3)).mean().dtype == dtypes.fp8e4m3
|
||||
assert (Tensor([0, 1], dtype=dtypes.fp8e5m2)).mean().dtype == dtypes.fp8e5m2
|
||||
assert (Tensor([0, 1], dtype=dtypes.float16)).mean().dtype == dtypes.float16
|
||||
assert (Tensor([0, 1], dtype=dtypes.bfloat16)).mean().dtype == dtypes.bfloat16
|
||||
assert (Tensor([0, 1], dtype=dtypes.float32)).mean().dtype == dtypes.float32
|
||||
assert (Tensor([0, 1], dtype=dtypes.float64)).mean().dtype == dtypes.float64
|
||||
|
||||
def test_cumsum(self):
|
||||
assert (Tensor([0, 1], dtype=dtypes.bool)).cumsum(0).dtype == dtypes.int32
|
||||
assert (Tensor([0, 1], dtype=dtypes.int8)).cumsum(0).dtype == dtypes.int32
|
||||
assert (Tensor([0, 1], dtype=dtypes.int16)).cumsum(0).dtype == dtypes.int32
|
||||
assert (Tensor([0, 1], dtype=dtypes.int32)).cumsum(0).dtype == dtypes.int32
|
||||
assert (Tensor([0, 1], dtype=dtypes.int64)).cumsum(0).dtype == dtypes.int64
|
||||
assert (Tensor([0, 1], dtype=dtypes.uint8)).cumsum(0).dtype == dtypes.uint32
|
||||
assert (Tensor([0, 1], dtype=dtypes.uint16)).cumsum(0).dtype == dtypes.uint32
|
||||
assert (Tensor([0, 1], dtype=dtypes.uint32)).cumsum(0).dtype == dtypes.uint32
|
||||
assert (Tensor([0, 1], dtype=dtypes.uint64)).cumsum(0).dtype == dtypes.uint64
|
||||
assert (Tensor([0, 1], dtype=dtypes.fp8e4m3)).cumsum(0).dtype == dtypes.fp8e4m3
|
||||
assert (Tensor([0, 1], dtype=dtypes.fp8e5m2)).cumsum(0).dtype == dtypes.fp8e5m2
|
||||
assert (Tensor([0, 1], dtype=dtypes.float16)).cumsum(0).dtype == dtypes.float16
|
||||
assert (Tensor([0, 1], dtype=dtypes.bfloat16)).cumsum(0).dtype == dtypes.bfloat16
|
||||
assert (Tensor([0, 1], dtype=dtypes.float32)).cumsum(0).dtype == dtypes.float32
|
||||
assert (Tensor([0, 1], dtype=dtypes.float64)).cumsum(0).dtype == dtypes.float64
|
||||
|
||||
def test_cumsum_empty(self):
|
||||
# empty cumsum dtype must match non-empty
|
||||
for d in (dtypes.bool, dtypes.int8, dtypes.uint8, dtypes.float16, dtypes.float32):
|
||||
self.assertEqual(Tensor([], dtype=d).cumsum(0).dtype, Tensor([0, 1], dtype=d).cumsum(0).dtype)
|
||||
|
||||
@given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
|
||||
def test_matmul(self, dt1, dt2, acc_dt):
|
||||
t1 = Tensor([0, 1], dtype=dt1)
|
||||
t2 = Tensor([0, 1], dtype=dt2)
|
||||
self.assertEqual(t1.matmul(t2).dtype, least_upper_dtype(t1.dtype, t2.dtype))
|
||||
self.assertEqual(t1.matmul(t2, dtype=acc_dt).dtype, acc_dt)
|
||||
|
||||
@given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
|
||||
def test_linear(self, dt1, dt2, dt3, acc_dt):
|
||||
x = Tensor([0, 1], dtype=dt1)
|
||||
w = Tensor([0, 1], dtype=dt2)
|
||||
b = Tensor([0, 1], dtype=dt3)
|
||||
self.assertEqual(x.linear(w).dtype, least_upper_dtype(x.dtype, w.dtype))
|
||||
self.assertEqual(x.linear(w, b).dtype, least_upper_dtype(least_upper_dtype(x.dtype, w.dtype), b.dtype))
|
||||
self.assertEqual(x.linear(w, dtype=acc_dt).dtype, acc_dt)
|
||||
self.assertEqual(x.linear(w, b, dtype=acc_dt).dtype, acc_dt)
|
||||
|
||||
@staticmethod
|
||||
def check_where_alternate_input_other(input_, other, data_type):
|
||||
assert (Tensor([True, False]).where(input_, other)).dtype == data_type
|
||||
assert (Tensor([True, False]).where(other, input_)).dtype == data_type
|
||||
|
||||
@given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
|
||||
def test_where_no_scalar(self, dt1, dt2):
|
||||
self.check_where_alternate_input_other(Tensor(2, dtype=dt1), Tensor(3, dtype=dt2), least_upper_dtype(dt1, dt2))
|
||||
|
||||
@given(strat.sampled_from(core_dtypes))
|
||||
def test_where_one_scalar(self, dt):
|
||||
t = Tensor(2, dtype=dt)
|
||||
self.check_where_alternate_input_other(t, 3.2, (dt if dtypes.is_float(dt) else dtypes.weakfloat))
|
||||
self.check_where_alternate_input_other(t, 3, (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.weakint))
|
||||
self.check_where_alternate_input_other(t, True, dt)
|
||||
|
||||
def test_where_two_scalars(self):
|
||||
self.check_where_alternate_input_other(3.1, 3.2, dtypes.weakfloat)
|
||||
self.check_where_alternate_input_other(3.1, 3, dtypes.weakfloat)
|
||||
self.check_where_alternate_input_other(3.1, True, dtypes.weakfloat)
|
||||
self.check_where_alternate_input_other(3, 2, dtypes.weakint)
|
||||
self.check_where_alternate_input_other(3, True, dtypes.weakint)
|
||||
|
||||
def test_where_non_bool_cond_raises(self):
|
||||
with self.assertRaises(RuntimeError): Tensor([1, 0, 2]).where(1, 0)
|
||||
self.check_where_alternate_input_other(False, True, dtypes.bool)
|
||||
|
||||
@given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
|
||||
def test_maximum(self, dt1, dt2):
|
||||
assert Tensor([0, 1, 2], dtype=dt1).maximum(Tensor([2, 0, 5], dtype=dt2)).dtype == least_upper_dtype(dt1, dt2)
|
||||
|
||||
@given(strat.sampled_from(core_dtypes))
|
||||
def test_maximum_const(self, dt):
|
||||
assert Tensor([1, 2], dtype=dt).maximum(3.1).dtype == (dt if dtypes.is_float(dt) else dtypes.weakfloat)
|
||||
assert Tensor([1, 2], dtype=dt).maximum(3).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.weakint)
|
||||
assert Tensor([1, 2], dtype=dt).maximum(True).dtype == dt
|
||||
|
||||
def test_div(self):
|
||||
assert (Tensor([1, 2], dtype=dtypes.int32) / Tensor([2, 2], dtype=dtypes.int32)).dtype == dtypes.default_float
|
||||
assert (Tensor([1, 2], dtype=dtypes.int16) / Tensor([2, 2], dtype=dtypes.int32)).dtype == dtypes.default_float
|
||||
assert (Tensor([1, 2], dtype=dtypes.float32) / Tensor([2, 2], dtype=dtypes.float16)).dtype == dtypes.float32
|
||||
assert (Tensor([1, 2], dtype=dtypes.int32) / Tensor([2, 2], dtype=dtypes.float16)).dtype == dtypes.float16
|
||||
|
||||
def test_div_const(self):
|
||||
assert (Tensor([1, 2], dtype=dtypes.int32) / 2).dtype == dtypes.default_float
|
||||
assert (Tensor([1, 2], dtype=dtypes.int32) / 2.0).dtype == dtypes.weakfloat
|
||||
assert (Tensor([1, 2], dtype=dtypes.float16) / 2).dtype == dtypes.float16
|
||||
assert (Tensor([1, 2], dtype=dtypes.float16) / 2.0).dtype == dtypes.float16
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
38
artifacts/package_sources/tinygrad/test/null/test_elf.py
Normal file
38
artifacts/package_sources/tinygrad/test/null/test_elf.py
Normal file
@@ -0,0 +1,38 @@
|
||||
import unittest, subprocess, platform
|
||||
from tinygrad.runtime.support.compiler_cpu import ClangCompiler
|
||||
from tinygrad.runtime.support.elf import elf_loader
|
||||
|
||||
class TestElfLoader(unittest.TestCase):
|
||||
def test_load_clang_jit_strtab(self):
|
||||
src = '''
|
||||
int something; // will be a load from a relocation (needed for .rela.text to exist)
|
||||
int test(int x) {
|
||||
return something + x;
|
||||
}
|
||||
'''
|
||||
args = ('-x', 'c', '-c', '-target', f'{platform.machine()}-none-unknown-elf', '-march=native', '-fPIC', '-O2', '-ffreestanding', '-nostdlib')
|
||||
obj = subprocess.check_output(('clang',) + args + ('-', '-o', '-'), input=src.encode('utf-8'))
|
||||
_, sections, _ = elf_loader(obj)
|
||||
section_names = [sh.name for sh in sections]
|
||||
assert '.text' in section_names and '.rela.text' in section_names, str(section_names)
|
||||
def test_clang_jit_compiler_external_raise(self):
|
||||
src = '''
|
||||
int evil_external_function(int);
|
||||
int test(int x) {
|
||||
return evil_external_function(x+2)*2;
|
||||
}
|
||||
'''
|
||||
with self.assertRaisesRegex(RuntimeError, 'evil_external_function'):
|
||||
ClangCompiler([{'AMD64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine(), m), "native"]).compile(src)
|
||||
def test_link(self):
|
||||
src = '''
|
||||
float powf(float, float); // from libm
|
||||
float test(float x, float y) { return powf(x, y); }
|
||||
'''
|
||||
args = ('-x', 'c', '-c', '-target', f'{platform.machine()}-none-unknown-elf', '-march=native', '-fPIC', '-O2', '-ffreestanding', '-nostdlib')
|
||||
obj = subprocess.check_output(('clang',) + args + ('-', '-o', '-'), input=src.encode())
|
||||
with self.assertRaisesRegex(RuntimeError, 'powf'): elf_loader(obj)
|
||||
elf_loader(obj, link_libs=['m'])
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
120
artifacts/package_sources/tinygrad/test/null/test_gc.py
Normal file
120
artifacts/package_sources/tinygrad/test/null/test_gc.py
Normal file
@@ -0,0 +1,120 @@
|
||||
#!/usr/bin/env python
|
||||
import gc, inspect
|
||||
import unittest
|
||||
import numpy as np
|
||||
from tinygrad.device import Buffer
|
||||
from tinygrad.engine.realize import run_linear
|
||||
from tinygrad.uop.ops import UOp
|
||||
from tinygrad.tensor import Tensor
|
||||
|
||||
def _allocations_of_type(t):
|
||||
ret = 0
|
||||
for x in gc.get_objects():
|
||||
try:
|
||||
if isinstance(x, t): ret += 1
|
||||
except ReferenceError:
|
||||
pass
|
||||
return ret
|
||||
|
||||
def tensors_allocated():
|
||||
gc.collect()
|
||||
return _allocations_of_type(Tensor)
|
||||
|
||||
def bufs_allocated():
|
||||
gc.collect()
|
||||
return _allocations_of_type(Buffer)
|
||||
|
||||
class TestGC(unittest.TestCase):
|
||||
|
||||
def test_gc(self):
|
||||
Tensor.manual_seed(0)
|
||||
base = tensors_allocated()
|
||||
a = Tensor.rand(4, 4)
|
||||
b = Tensor.zeros(4, 4)
|
||||
(a*b).mean().backward()
|
||||
assert (tensors_allocated()-base > 0)
|
||||
del a,b
|
||||
assert (tensors_allocated()-base == 2) # one for Tensor._device_rng_counters, and one for Tensor._device_seeds
|
||||
Tensor.manual_seed(0)
|
||||
|
||||
def test_gc_complex(self):
|
||||
Tensor.manual_seed(0)
|
||||
base = tensors_allocated()
|
||||
a = Tensor(np.zeros((4, 4), dtype=np.float32))
|
||||
b = Tensor.rand(4, 4)
|
||||
assert (tensors_allocated()-base == 4)
|
||||
(a*b).mean().backward()
|
||||
assert (tensors_allocated()-base == 6)
|
||||
del b
|
||||
assert (tensors_allocated()-base == 4)
|
||||
b = Tensor(np.zeros((4, 4), dtype=np.float32))
|
||||
print(tensors_allocated())
|
||||
(a*b).mean().backward()
|
||||
print(tensors_allocated())
|
||||
assert (tensors_allocated()-base == 6)
|
||||
del b
|
||||
assert (tensors_allocated()-base == 4)
|
||||
Tensor.manual_seed(0)
|
||||
|
||||
def test_schedule_gc(self):
|
||||
init = bufs_allocated()
|
||||
x = Tensor.ones(256).contiguous().realize()
|
||||
y = Tensor.ones(5, 5).contiguous()
|
||||
y.schedule_linear()
|
||||
del x
|
||||
del y
|
||||
self.assertEqual(bufs_allocated()-init, 0)
|
||||
|
||||
def test_schedule_gc_with_inputs(self):
|
||||
init = bufs_allocated()
|
||||
x = Tensor.ones(256).contiguous().realize()
|
||||
y = x+Tensor.ones(256).contiguous()
|
||||
del x
|
||||
run_linear(*y.linear_with_vars())
|
||||
self.assertEqual(bufs_allocated()-init, 1)
|
||||
del y
|
||||
self.assertEqual(bufs_allocated()-init, 0)
|
||||
|
||||
def test_toposort_blocks_gc(self):
|
||||
init = bufs_allocated()
|
||||
x = Tensor.ones(4,4).contiguous().realize()+1
|
||||
self.assertEqual(bufs_allocated()-init, 1)
|
||||
# try commenting this part out, it's green!
|
||||
x.uop.toposort()
|
||||
del x
|
||||
if bufs_allocated()-init != 0:
|
||||
print(inspect.getclosurevars(UOp.toposort().fget))
|
||||
raise AssertionError(f"never gced {[x for x in gc.get_objects() if isinstance(x, Buffer)]}")
|
||||
|
||||
def test_buffer_refcount(self):
|
||||
init = bufs_allocated()
|
||||
a = Tensor.empty(10)
|
||||
self.assertEqual(bufs_allocated()-init, 0)
|
||||
a.realize()
|
||||
real_buf = a.uop.buffer
|
||||
# after the Tensor UOp is deleted there shouldn't be any references on the Buffer
|
||||
self.assertEqual(real_buf.uop_refcount, 1)
|
||||
self.assertEqual(bufs_allocated()-init, 1)
|
||||
del a.uop
|
||||
self.assertEqual(real_buf.uop_refcount, 0)
|
||||
self.assertEqual(bufs_allocated()-init, 1) # keep the buffer alive
|
||||
del real_buf
|
||||
self.assertEqual(bufs_allocated()-init, 0)
|
||||
|
||||
def test_assign_refcount(self):
|
||||
init = bufs_allocated()
|
||||
a = Tensor.full((4,), 1.).contiguous()
|
||||
a.realize()
|
||||
real_buf = a.uop.buffer
|
||||
self.assertEqual(real_buf.uop_refcount, 1)
|
||||
a.assign(Tensor.full((4,), 2.))
|
||||
self.assertIs(a.uop.src[0].buffer, real_buf)
|
||||
# NOTE: this is still 1, we don't count the ASSIGN
|
||||
self.assertEqual(real_buf.uop_refcount, 1)
|
||||
a.realize()
|
||||
del a
|
||||
self.assertEqual(real_buf.uop_refcount, 0) # no UOps for this Buffer
|
||||
self.assertEqual(bufs_allocated()-init, 1) # Buffer is alive
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
121
artifacts/package_sources/tinygrad/test/null/test_gpudims.py
Normal file
121
artifacts/package_sources/tinygrad/test/null/test_gpudims.py
Normal file
@@ -0,0 +1,121 @@
|
||||
import unittest, math
|
||||
import z3
|
||||
from tinygrad.codegen.gpudims import get_grouped_dims, add_gpudims
|
||||
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
|
||||
from tinygrad.uop.validate import uops_to_z3
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.renderer import Renderer
|
||||
from tinygrad.helpers import flatten, dedup, Target
|
||||
|
||||
class TestGroupedDims(unittest.TestCase):
|
||||
def _check_grouped_dims(self, prefix, dims, max_sizes, reverse, expected_sizes, assert_same_length=True):
|
||||
idxs = get_grouped_dims(prefix, dims, max_sizes, reverse)
|
||||
loop_idxs = dedup(flatten([[y for y in x.toposort() if y.op is Ops.SPECIAL] for x in idxs]))
|
||||
loop_idxs = sorted(loop_idxs, key=lambda uop: uop.arg)
|
||||
sizes = [x.src[0].val for x in loop_idxs]
|
||||
assert len(idxs) == len(dims), f"expected idxs to have same length as dims {len(dims)}, got {len(idxs)}"
|
||||
if assert_same_length:
|
||||
assert len(loop_idxs) == min(len(sizes), len(dims)), f"expected idxs to have length {min(len(sizes), len(dims))}, got {len(loop_idxs)}"
|
||||
assert sizes == expected_sizes, f"expected sizes={expected_sizes}, got {sizes=}"
|
||||
self._verify_indices_z3(idxs, dims)
|
||||
|
||||
def _verify_indices_z3(self, idxs, dims):
|
||||
"""Use z3 to prove bijectivity: bounds (0 <= flat < total) + injectivity (different inputs => different flat)."""
|
||||
total = math.prod(dims)
|
||||
specials = sorted(dedup(flatten([[y for y in x.toposort() if y.op is Ops.SPECIAL] for x in idxs])), key=lambda u: u.arg)
|
||||
# build flat index and primed flat (same expression with renamed SPECIALs)
|
||||
flat = UOp.const(0)
|
||||
for i, idx in enumerate(idxs):
|
||||
flat = flat + idx * int(math.prod(dims[i+1:]))
|
||||
flat_p = flat.substitute({s: UOp(Ops.SPECIAL, src=s.src, arg=s.arg+"_p") for s in specials})
|
||||
solver = z3.Solver()
|
||||
[z3_flat, z3_flat_p] = uops_to_z3(solver, flat, flat_p)
|
||||
# bounds
|
||||
self.assertEqual(solver.check(z3_flat < 0), z3.unsat, f"flat can be negative: {dims=}")
|
||||
self.assertEqual(solver.check(z3_flat >= total), z3.unsat, f"flat can be >= {total}: {dims=}")
|
||||
# injectivity: flat == flat' but inputs differ => unsat
|
||||
inputs_differ = z3.Or(*[z3.Int(s.arg) != z3.Int(s.arg+"_p") for s in specials])
|
||||
self.assertEqual(solver.check(z3.And(z3_flat == z3_flat_p, inputs_differ)), z3.unsat, f"not injective: {dims=}")
|
||||
|
||||
def test_grouped_dims(self):
|
||||
# no-op
|
||||
self._check_grouped_dims("gidx", (2,), (16,16,16), False, [2])
|
||||
self._check_grouped_dims("gidx", (2,3), (16,16,16), False, [2,3])
|
||||
|
||||
# check reverse dims
|
||||
self._check_grouped_dims("gidx", (2,3), (16,16,16), True, [3,2])
|
||||
self._check_grouped_dims("gidx", (2,3,4), (16,16,16), False, [2,3,4])
|
||||
|
||||
# test splitting globals: len(dims) == len(max)
|
||||
self._check_grouped_dims("gidx", (64,3,4), (16,16,16), False, [16,12,4])
|
||||
self._check_grouped_dims("gidx", (64,3,4), (16,4,16), False, [16,3,16])
|
||||
self._check_grouped_dims("gidx", (64,3,4), (16,16,16), True, [16,3,16])
|
||||
self._check_grouped_dims("gidx", (128,3,4), (16,4,256), False, [16,3,32])
|
||||
self._check_grouped_dims("gidx", (4,4,512), (16,4,256), False, [8,4,256])
|
||||
self._check_grouped_dims("gidx", (5,12,7), (8,4,16), False, [10,3,14])
|
||||
|
||||
# prefer group_dim strategy when possible
|
||||
self._check_grouped_dims("gidx", (512,4,2), (8192,2,2), False, [2048,2])
|
||||
|
||||
# test splitting globals: len(dims) < len(max)
|
||||
# len(dim) -> len(limited)
|
||||
# 1 -> 2
|
||||
self._check_grouped_dims("gidx", (128,), (16,16,256), False, [16,8], False)
|
||||
# 1 -> 3
|
||||
self._check_grouped_dims("gidx", (65536,), (16,16,256), False, [16,16,256], False)
|
||||
# 2 -> 2
|
||||
self._check_grouped_dims("gidx", (65536,2), (65535,65535,65535), False, [32768,4], False)
|
||||
# test when the only divisor is the square root of dim
|
||||
self._check_grouped_dims("gidx", (121,), (12,12,12), False, [11,11], False)
|
||||
# 2 -> 3
|
||||
self._check_grouped_dims("gidx", (128,128), (16,16,256), False, [16,16,64], False)
|
||||
|
||||
# collapse on onto the left most axis
|
||||
self._check_grouped_dims("gidx", (2,3,4,5), (16,16,16), False, [6,4,5])
|
||||
self._check_grouped_dims("gidx", (2,3,4,5), (32,16,16), True, [20,3,2])
|
||||
|
||||
# collapse on left-most available axis (the left most is too small)
|
||||
self._check_grouped_dims("gidx", (2,3,4,5), (4,16,16), False, [2,12,5])
|
||||
self._check_grouped_dims("gidx", (2,3,4,5), (16,16,16), True, [5,12,2])
|
||||
|
||||
# dim too large and not factorable
|
||||
with self.assertRaises(RuntimeError):
|
||||
get_grouped_dims("gidx", (23,), (16,16,16), False,)
|
||||
with self.assertRaises(RuntimeError):
|
||||
get_grouped_dims("gidx", (128,3,4), (16,2,2), False,)
|
||||
|
||||
# too large for sizes
|
||||
with self.assertRaises(RuntimeError):
|
||||
get_grouped_dims("gidx", (2,3,4,5,6), (16,16,16))
|
||||
|
||||
def test_grouped_direct_dims_are_special(self):
|
||||
# when (2,3) are merged into 6, the unmerged dims (4,5) should map directly to SPECIAL ops (no div/mod)
|
||||
idxs = get_grouped_dims("gidx", (2,3,4,5), (16,16,16), False)
|
||||
assert idxs[2].op is Ops.SPECIAL, f"expected SPECIAL for direct-mapped dim, got {idxs[2].op}"
|
||||
assert idxs[3].op is Ops.SPECIAL, f"expected SPECIAL for direct-mapped dim, got {idxs[3].op}"
|
||||
|
||||
def test_grouped_dims_high_rank(self):
|
||||
# 4D collapsed onto 2 axes
|
||||
self._check_grouped_dims("gidx", (4,4,4,4), (16,16), False, [16,16])
|
||||
# 4D untouched
|
||||
self._check_grouped_dims("gidx", (2,3,4,5), None, False, [2,3,4,5])
|
||||
idxs = get_grouped_dims("gidx", (2,3,4,5), None, False)
|
||||
assert all(u.op is Ops.SPECIAL for u in idxs), f"expected all-SPECIAL when untouched, got {[u.op for u in idxs]}"
|
||||
# 5D and 6D collapsed onto 3 axes
|
||||
self._check_grouped_dims("gidx", (2,2,2,2,2), (4,4,4), False, [4,4,2])
|
||||
self._check_grouped_dims("gidx", (2,2,2,2,2,2), (8,8,8), False, [8,4,2])
|
||||
|
||||
def test_global_prod_max(self):
|
||||
g, l = UOp.range(256, 0, AxisType.GLOBAL), UOp.range(256, 1, AxisType.LOCAL)
|
||||
sink = UOp.param(0, dtypes.float, (512,)).index(g + l).store(UOp.const(1.0)).end(g, l).sink(arg=KernelInfo())
|
||||
class R(Renderer): global_max, local_max, global_prod_max = (256, 256, 256), (128, 128, 128), (128, 128, 128)
|
||||
specials = [u for u in add_gpudims(R(Target()), sink).toposort() if u.op is Ops.SPECIAL]
|
||||
self.assertGreater(len([s for s in specials if "lidx" in s.arg]), 1)
|
||||
self.assertGreater(len([s for s in specials if "gidx" in s.arg]), 1)
|
||||
|
||||
def test_max_sizes_none(self):
|
||||
self._check_grouped_dims("gidx", (2,3,4), None, False, [2,3,4])
|
||||
self._check_grouped_dims("gidx", (100,), None, False, [100])
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,83 @@
|
||||
from typing import Callable
|
||||
import unittest, math
|
||||
import torch
|
||||
from tinygrad import Tensor
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.uop.ops import UOp
|
||||
from tinygrad.mixin.gradient import compute_gradient
|
||||
|
||||
class TestGradient(unittest.TestCase):
|
||||
def _cmp_nan_okay(self, x, y):
|
||||
if math.isnan(x) and math.isnan(y): return
|
||||
self.assertAlmostEqual(x, y, places=5)
|
||||
|
||||
def _test_one_input_function(self, f:Callable, jf:Callable|None=None):
|
||||
if jf is None: jf = f
|
||||
x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float)
|
||||
gx = compute_gradient(f(x), UOp.const(1.0), set([x]))[x]
|
||||
|
||||
for val in [-5., -2.0, 0.0, 2.0, 5.]:
|
||||
tg_out = gx.substitute({x: UOp.const(val)}).ssimplify()
|
||||
tx = torch.tensor([val], dtype=torch.float, requires_grad=True)
|
||||
torch_out = torch.autograd.grad(jf(tx), tx)[0].item()
|
||||
self._cmp_nan_okay(tg_out, torch_out)
|
||||
|
||||
def _test_two_input_function(self, f:Callable, jf:Callable|None=None):
|
||||
if jf is None: jf = f
|
||||
x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float)
|
||||
y = UOp.variable('y', -math.inf, math.inf, dtype=dtypes.float)
|
||||
grads = compute_gradient(f(x, y), UOp.const(1.0), set([x, y]))
|
||||
gx, gy = grads[x], grads[y]
|
||||
|
||||
for valx in [-5., -2.0, 0.0, 2.0, 5.]:
|
||||
for valy in [-5., -2.0, 0.0, 2.0, 5.]:
|
||||
# Substitute the values into the gradient expressions
|
||||
substitutions = {x: UOp.const(valx), y: UOp.const(valy)}
|
||||
tg_out_x = gx.substitute(substitutions).ssimplify()
|
||||
tg_out_y = gy.substitute(substitutions).ssimplify()
|
||||
|
||||
tx = torch.tensor([valx], dtype=torch.float, requires_grad=True)
|
||||
ty = torch.tensor([valy], dtype=torch.float, requires_grad=True)
|
||||
torch_grad = torch.autograd.grad(jf(tx, ty), [tx, ty])
|
||||
torch_out_x, torch_out_y = [x.item() for x in torch_grad]
|
||||
|
||||
self._cmp_nan_okay(tg_out_x, torch_out_x)
|
||||
self._cmp_nan_okay(tg_out_y, torch_out_y)
|
||||
|
||||
# unary ops unit
|
||||
def test_recip(self): self._test_one_input_function(lambda x: 1.0/x)
|
||||
def test_sin(self): self._test_one_input_function(lambda x: x.sin())
|
||||
def test_sqrt(self): self._test_one_input_function(lambda x: x.sqrt())
|
||||
def test_log2(self): self._test_one_input_function(lambda x: x.log2())
|
||||
def test_exp2(self): self._test_one_input_function(lambda x: x.exp2())
|
||||
|
||||
# binary ops unit
|
||||
def test_add(self): self._test_two_input_function(lambda x,y: x+y)
|
||||
def test_mul(self): self._test_two_input_function(lambda x,y: x*y)
|
||||
|
||||
# chain rule
|
||||
def test_chain(self): self._test_one_input_function(lambda x: x.sin().sqrt())
|
||||
def test_chain_binop(self): self._test_two_input_function(lambda x,y: (x*y)+x*y)
|
||||
def test_big_add_sin(self): self._test_two_input_function(lambda x,y: x.sin()+3.0/y)
|
||||
def test_big_chain(self): self._test_two_input_function(lambda x,y: (1.0/x*y)+x*y)
|
||||
def test_where(self): self._test_two_input_function(lambda x,y: (x<y).where(x,y), lambda x,y: torch.where(x<y,x,y))
|
||||
|
||||
class TestRealizeMeansRealize(unittest.TestCase):
|
||||
def test_randn_realizes(self):
|
||||
x = Tensor.randn(2, 3, 64, 64).realize()
|
||||
assert x.uop is not x.uop.base
|
||||
assert x.uop.is_realized
|
||||
|
||||
def test_uniform_realizes(self):
|
||||
x = Tensor.uniform(16, 3, 3, 3).realize()
|
||||
print(x.uop)
|
||||
assert x.uop is not x.uop.base
|
||||
assert x.uop.is_realized
|
||||
|
||||
def test_uniform_gradient(self):
|
||||
x = Tensor.uniform(16, 3, 3, 3).realize()
|
||||
y = x * 2
|
||||
y.sum().gradient(x)[0].realize()
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,531 @@
|
||||
import unittest, math
|
||||
from tinygrad import dtypes
|
||||
from tinygrad.dtype import AddrSpace
|
||||
from tinygrad.helpers import all_same, Context
|
||||
from tinygrad.uop.ops import GroupOp, UOp, Ops, exec_alu, PatternMatcher, TrackedPatternMatcher, UPat
|
||||
from test.helpers import full_rewrite
|
||||
from hypothesis import given, strategies as strat
|
||||
|
||||
# Helper function to apply the graph rewrite
|
||||
@Context(SPEC=0)
|
||||
def apply_rewrite(expr):
|
||||
return full_rewrite(expr.sink()).src[0]
|
||||
|
||||
@Context(SPEC=0)
|
||||
def apply_rewrite_values(expr):
|
||||
srcs = full_rewrite(expr.sink()).src
|
||||
if len(srcs) == 1:
|
||||
if srcs[0].op is Ops.CONST: return (srcs[0].val,)
|
||||
if srcs[0].op is Ops.STACK: return tuple(s.val for s in srcs[0].src)
|
||||
return tuple(s.val for s in srcs)
|
||||
|
||||
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 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))
|
||||
self.assertEqual(optimized_div_uop.op, Ops.CONST)
|
||||
self.assertTrue(math.isinf(optimized_div_uop.val) or math.isnan(optimized_div_uop.val))
|
||||
|
||||
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.assertEqual(optimized_uop.op, Ops.CONST)
|
||||
self.assertEqual(optimized_uop.val, 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.assertEqual(optimized_uop.op, Ops.CONST)
|
||||
self.assertEqual(optimized_uop.val, 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.assertEqual(optimized_uop.op, Ops.CONST)
|
||||
self.assertEqual(optimized_uop.val, 42.0)
|
||||
|
||||
def test_full_graph_rewrite_modulo_by_one(self):
|
||||
optimized_uop = apply_rewrite(UOp.const(42) % UOp.const(1))
|
||||
self.assertEqual(optimized_uop.op, Ops.CONST)
|
||||
self.assertEqual(optimized_uop.val, 0)
|
||||
|
||||
|
||||
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):
|
||||
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.assertEqual(optimized_mod_uop.op, Ops.CONST)
|
||||
self.assertEqual(optimized_mod_uop.val, 2)
|
||||
|
||||
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.assertEqual(optimized_div_uop.op, Ops.MUL)
|
||||
self.assertEqual(optimized_div_uop.src[1].val, 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.assertEqual(optimized_div_uop.op, Ops.CONST)
|
||||
self.assertEqual(optimized_div_uop.val, 1)
|
||||
|
||||
def test_graph_rewrite_div_folding_bug(self):
|
||||
lhs = UOp(Ops.ADD, src=(
|
||||
UOp(Ops.STACK, arg=None, src=(UOp(Ops.SPECIAL, src=(UOp.const(32),), arg='lidx0'),)*4),
|
||||
UOp.const((0, 256, 512, 768))))
|
||||
rhs = UOp.const((2,)*4)
|
||||
unopt = lhs<rhs
|
||||
opt = apply_rewrite(unopt)
|
||||
print(unopt)
|
||||
print(opt)
|
||||
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)
|
||||
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)
|
||||
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):
|
||||
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_log2_neg, optimized_recip_zero = optimized_sink.src
|
||||
self.assertTrue(math.isnan(optimized_log2_neg.val), f"Expected NaN for log2(-1.0), got {optimized_log2_neg.val}")
|
||||
self.assertTrue(math.isinf(optimized_recip_zero.val) and optimized_recip_zero.val > 0,
|
||||
f"Expected +inf for reciprocal(0.0), got {optimized_recip_zero.val}")
|
||||
|
||||
@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):
|
||||
def test_gep_single_element_extraction(self):
|
||||
# GEP on a vector dtype to extract a single element
|
||||
base_vector = UOp.const((1.0, 2.0, 3.0, 4.0))
|
||||
self.assertEqual(apply_rewrite(base_vector.index(2)).val, 3.0)
|
||||
|
||||
def test_gep_tuple_extraction(self):
|
||||
# GEP on a vector dtype to extract multiple elements as a vector
|
||||
base_vector = UOp.const((1.0, 2.0, 3.0, 4.0))
|
||||
self.assertEqual(list(apply_rewrite_values(UOp.stack(*[base_vector.index(i) for i in (2, 3)]))), [3.0, 4.0])
|
||||
|
||||
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.assertEqual(apply_rewrite(const_stack.index(2)).val, 3.0)
|
||||
|
||||
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.assertEqual(list(apply_rewrite_values(UOp.stack(*[const_stack.index(i) for i in (1, 3)]))), [8.0, 10.0])
|
||||
|
||||
def test_vectorize_multiple_elements(self):
|
||||
# Vectorizing multiple elements using GEP
|
||||
base_vector = UOp.const((5.0, 10.0, 15.0, 20.0))
|
||||
vectorized_uop = UOp(Ops.STACK, src=tuple(base_vector.index(i) for i in range(4)))
|
||||
self.assertEqual(list(apply_rewrite_values(vectorized_uop)), [5.0, 10.0, 15.0, 20.0])
|
||||
|
||||
|
||||
import inspect
|
||||
from tinygrad.uop.ops import graph_rewrite, _substitute, rewrite_group
|
||||
from tinygrad.uop.symbolic import symbolic_simple
|
||||
|
||||
class TestBottomUpRewrite(unittest.TestCase):
|
||||
def test_const_folding(self):
|
||||
a = UOp.const(5)
|
||||
ret = (a*3) + (a*7)
|
||||
gt = graph_rewrite(ret, symbolic_simple)
|
||||
ret = graph_rewrite(ret, symbolic_simple, bottom_up=True)
|
||||
self.assertIs(gt, ret)
|
||||
|
||||
# normally .substitute would be fine, but it's not tracked
|
||||
@rewrite_group()
|
||||
def named_substitute(name:str, uop:UOp, rel:dict[UOp, UOp]): return graph_rewrite(uop, _substitute, rel, bottom_up=True)
|
||||
def substitute(uop:UOp, rel:dict[UOp, UOp]): return named_substitute(inspect.stack()[1].function, uop, rel)
|
||||
|
||||
class TestSubstitute(unittest.TestCase):
|
||||
# these work because the substituted things don't have parents
|
||||
def test_simple(self):
|
||||
a = UOp.variable('a', 0, 10)
|
||||
b = UOp.variable('b', 0, 10)
|
||||
ret = a + 4
|
||||
ret = substitute(ret, {a:b})
|
||||
self.assertIs(ret, b+4)
|
||||
|
||||
def test_double(self):
|
||||
a = UOp.variable('a', 0, 10)
|
||||
b = UOp.variable('b', 0, 10)
|
||||
c = UOp.variable('c', 0, 10)
|
||||
ret = (a + 4) + b
|
||||
ret = substitute(ret, {a:c, b:c})
|
||||
self.assertIs(ret, (c + 4) + c)
|
||||
|
||||
def test_diamond(self):
|
||||
a = UOp.variable('a', 0, 10)
|
||||
b = UOp.variable('b', 0, 10)
|
||||
ret = (a + 4) + (a + 5)
|
||||
ret = substitute(ret, {a:b})
|
||||
self.assertIs(ret, (b + 4) + (b + 5))
|
||||
|
||||
# this works because there's nothing above the substituted node
|
||||
def test_sin(self):
|
||||
a = UOp.variable('a', 0, 10, dtype=dtypes.float)
|
||||
b = UOp.variable('b', 0, 10, dtype=dtypes.float)
|
||||
ret = a.sin().sin()
|
||||
ret = substitute(ret, {a.sin():b})
|
||||
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):
|
||||
a = UOp.variable('a', 0, 10, dtype=dtypes.float)
|
||||
n1 = a.sin()
|
||||
ret = n1.sin()
|
||||
ret = substitute(ret, {a.sin():a.sqrt()})
|
||||
self.assertIs(ret, a.sqrt().sin())
|
||||
|
||||
def test_double_sin_to_sqrt(self):
|
||||
a = UOp.variable('a', 0, 10, dtype=dtypes.float)
|
||||
n1 = a.sin()
|
||||
ret = n1.sin()
|
||||
# NOTE: this would work if it had gone in the opposite order
|
||||
ret = substitute(ret, {a.sin():a.sqrt(), n1.sin():n1.sqrt()})
|
||||
self.assertIs(ret, a.sqrt().sqrt())
|
||||
|
||||
def test_tagged_replace(self):
|
||||
a = UOp.variable('a', 0, 10)
|
||||
b = UOp.variable('b', 0, 10)
|
||||
ret = (a+4).replace(tag=1)
|
||||
ret = substitute(ret, {a:b})
|
||||
# the srcs are rewritten but we keep tag
|
||||
self.assertIs(ret, (b+4).replace(tag=1))
|
||||
|
||||
matchers = strat.sampled_from([PatternMatcher, TrackedPatternMatcher])
|
||||
|
||||
class TestRecurse(unittest.TestCase):
|
||||
@given(matchers)
|
||||
def test_no_inf_loop(self, PatternMatcher):
|
||||
a = UOp.variable('a', 0, 10)
|
||||
pm = PatternMatcher([(UPat(Ops.PARAM, name="x"), lambda x: x)])
|
||||
graph_rewrite(a, pm)
|
||||
|
||||
@given(matchers)
|
||||
def test_no_inf_loop_bottom_up(self, PatternMatcher):
|
||||
a = UOp.variable('a', 0, 10)
|
||||
pm = PatternMatcher([(UPat(Ops.PARAM, name="x"), lambda x: x)])
|
||||
graph_rewrite(a, pm, bottom_up=True)
|
||||
|
||||
def test_inf_loop(self):
|
||||
a = UOp.const(3)
|
||||
pm = PatternMatcher([
|
||||
(UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)),
|
||||
(UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)),
|
||||
])
|
||||
with self.assertRaises(RuntimeError):
|
||||
graph_rewrite(a, pm)
|
||||
|
||||
def test_inf_loop_bottom_up(self):
|
||||
a = UOp.const(3)
|
||||
pm = PatternMatcher([
|
||||
(UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)),
|
||||
(UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)),
|
||||
])
|
||||
with self.assertRaises(RuntimeError):
|
||||
graph_rewrite(a, pm, bottom_up=True)
|
||||
|
||||
def bidir_append(ctx, x, b): ctx.append((x.val if x.op is Ops.CONST else "+", b))
|
||||
class TestBidirectional(unittest.TestCase):
|
||||
def test_simple(self):
|
||||
a = UOp.const(1)
|
||||
b = UOp.const(2)
|
||||
c = a + b
|
||||
pm = PatternMatcher([ (UPat(GroupOp.All, name="x"), lambda ctx,x: bidir_append(ctx, x, False)) ])
|
||||
bpm = PatternMatcher([ (UPat(GroupOp.All, name="x"), lambda ctx,x: bidir_append(ctx, x, True)) ])
|
||||
ctx_list = []
|
||||
graph_rewrite(c, pm, ctx=ctx_list, bpm=bpm)
|
||||
self.assertListEqual(ctx_list, [('+', True), (1, True), (1, False), (2, True), (2, False), ('+', False)])
|
||||
|
||||
class TestStopEarly(unittest.TestCase):
|
||||
def test_stop_early(self):
|
||||
a = UOp.const(3)
|
||||
b = UOp.const(4)
|
||||
c = a+b
|
||||
cn = UOp.const(7)
|
||||
d = UOp.const(2)
|
||||
def visit_const(c:UOp):
|
||||
print(f"visit {c.val}")
|
||||
assert c.val not in (3,4)
|
||||
pm_cvisit = PatternMatcher([(UPat(Ops.CONST, name="c"), visit_const),])
|
||||
ret = (c+d).substitute({c:cn}, extra_pm=pm_cvisit)
|
||||
assert ret == cn+d
|
||||
|
||||
class TestWalkRewrite(unittest.TestCase):
|
||||
"""Tests for graph_rewrite with walk=True (MLIR Walk Pattern Rewrite Driver semantics).
|
||||
walk=True gives a single-pass traversal that does NOT revisit or re-traverse into rewritten subtrees.
|
||||
Supports both top-down (default) and bottom-up (bottom_up=True) modes."""
|
||||
|
||||
# *** top-down walk (default): process children first, then try pm on rebuilt node ***
|
||||
|
||||
def test_walk_topdown_simple_substitute(self):
|
||||
a = UOp.variable('a', 0, 10)
|
||||
b = UOp.variable('b', 0, 10)
|
||||
ret = graph_rewrite(a + 4, _substitute, {a:b}, walk=True)
|
||||
self.assertIs(ret, b+4)
|
||||
|
||||
def test_walk_topdown_does_not_traverse_into_replacement(self):
|
||||
"""Top-down walk: replacement subtrees are NOT re-entered."""
|
||||
a = UOp.variable('a', 0, 10)
|
||||
b = UOp.variable('b', 0, 10)
|
||||
c = UOp.variable('c', 0, 10)
|
||||
d = UOp.variable('d', 0, 10)
|
||||
# a is replaced by b+c, but b inside the replacement is NOT further substituted to d
|
||||
ret_walk = graph_rewrite(a + 4, _substitute, {a:b+c, b:d}, walk=True)
|
||||
self.assertIs(ret_walk, (b+c)+4)
|
||||
# contrast: greedy bottom_up WOULD replace b inside the replacement
|
||||
ret_greedy = graph_rewrite(a + 4, _substitute, {a:b+c, b:d}, bottom_up=True)
|
||||
self.assertIs(ret_greedy, (d+c)+4)
|
||||
|
||||
def test_walk_topdown_no_fixed_point(self):
|
||||
"""A bouncing pattern applies once and stops instead of looping."""
|
||||
a = UOp.const(3)
|
||||
pm = PatternMatcher([
|
||||
(UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)),
|
||||
(UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)),
|
||||
])
|
||||
with self.assertRaises(RuntimeError):
|
||||
graph_rewrite(a, pm, bottom_up=True)
|
||||
ret = graph_rewrite(a, pm, walk=True)
|
||||
self.assertIs(ret, UOp.const(4))
|
||||
|
||||
def test_walk_topdown_rewrites_children(self):
|
||||
a = UOp.variable('a', 0, 10)
|
||||
b = UOp.variable('b', 0, 10)
|
||||
c = UOp.variable('c', 0, 10)
|
||||
ret = graph_rewrite((a + 4) + (b + 5), _substitute, {a:c, b:c}, walk=True)
|
||||
self.assertIs(ret, (c + 4) + (c + 5))
|
||||
|
||||
def test_walk_topdown_diamond(self):
|
||||
a = UOp.variable('a', 0, 10)
|
||||
b = UOp.variable('b', 0, 10)
|
||||
ret = graph_rewrite((a + 4) + (a + 5), _substitute, {a:b}, walk=True)
|
||||
self.assertIs(ret, (b + 4) + (b + 5))
|
||||
|
||||
def test_walk_topdown_children_rewritten_before_parent(self):
|
||||
"""Top-down walk processes children first: child substitution changes the rebuilt parent."""
|
||||
a = UOp.variable('a', 0, 10, dtype=dtypes.float)
|
||||
n1 = a.sin() # sin(a)
|
||||
ret = n1.sin() # sin(sin(a))
|
||||
# sin(a)->sqrt(a) fires first (child), parent rebuilds to sin(sqrt(a)), which doesn't match sin(sin(a)) in dvars
|
||||
ret_walk = graph_rewrite(ret, _substitute, {a.sin():a.sqrt(), n1.sin():n1.sqrt()}, walk=True)
|
||||
self.assertIs(ret_walk, a.sqrt().sin())
|
||||
|
||||
def test_walk_topdown_self_referential_replacement(self):
|
||||
"""Replacement containing the replaced node works without infinite recursion."""
|
||||
a = UOp.variable('a', 0, 10, dtype=dtypes.float)
|
||||
ret = graph_rewrite(a.sin() + 4, _substitute, {a.sin(): a.sin().sqrt()}, walk=True)
|
||||
self.assertIs(ret, a.sin().sqrt() + 4)
|
||||
|
||||
def test_walk_topdown_visit_order(self):
|
||||
"""Top-down walk fires pm after children are processed (post-order)."""
|
||||
visited = []
|
||||
def track_visit(ctx, x):
|
||||
ctx.append(x.val if x.op is Ops.CONST else x.op)
|
||||
return None
|
||||
pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)])
|
||||
a = UOp.const(1)
|
||||
b = UOp.const(2)
|
||||
graph_rewrite(a + b, pm, ctx=visited, walk=True)
|
||||
self.assertEqual(visited, [1, 2, Ops.ADD])
|
||||
|
||||
# *** bottom-up walk: try bpm on node first, skip children if it matches ***
|
||||
|
||||
def test_walk_bottomup_simple_substitute(self):
|
||||
a = UOp.variable('a', 0, 10)
|
||||
b = UOp.variable('b', 0, 10)
|
||||
ret = graph_rewrite(a + 4, _substitute, {a:b}, bottom_up=True, walk=True)
|
||||
self.assertIs(ret, b+4)
|
||||
|
||||
def test_walk_bottomup_does_not_traverse_into_replacement(self):
|
||||
"""Bottom-up walk: replacement subtrees are NOT entered."""
|
||||
a = UOp.variable('a', 0, 10)
|
||||
b = UOp.variable('b', 0, 10)
|
||||
c = UOp.variable('c', 0, 10)
|
||||
d = UOp.variable('d', 0, 10)
|
||||
ret = graph_rewrite(a + 4, _substitute, {a:b+c, b:d}, bottom_up=True, walk=True)
|
||||
self.assertIs(ret, (b+c)+4)
|
||||
|
||||
def test_walk_bottomup_parent_match_skips_children(self):
|
||||
"""Bottom-up walk matches parent first: if it matches, children are never visited."""
|
||||
a = UOp.variable('a', 0, 10, dtype=dtypes.float)
|
||||
n1 = a.sin()
|
||||
ret = n1.sin() # sin(sin(a))
|
||||
# sin(sin(a)) matches n1.sin()->n1.sqrt() immediately, children never visited, sin(a) inside replacement untouched
|
||||
ret_walk = graph_rewrite(ret, _substitute, {a.sin():a.sqrt(), n1.sin():n1.sqrt()}, bottom_up=True, walk=True)
|
||||
self.assertIs(ret_walk, a.sin().sqrt())
|
||||
|
||||
def test_walk_bottomup_no_fixed_point(self):
|
||||
"""Bottom-up walk also applies once per node, no fixed-point iteration."""
|
||||
a = UOp.const(3)
|
||||
pm = PatternMatcher([
|
||||
(UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)),
|
||||
(UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)),
|
||||
])
|
||||
ret = graph_rewrite(a, pm, bottom_up=True, walk=True)
|
||||
self.assertIs(ret, UOp.const(4))
|
||||
|
||||
def test_walk_bottomup_visit_order(self):
|
||||
"""Bottom-up walk fires bpm before descending (pre-order)."""
|
||||
visited = []
|
||||
def track_visit(ctx, x):
|
||||
ctx.append(x.val if x.op is Ops.CONST else x.op)
|
||||
return None
|
||||
pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)])
|
||||
a = UOp.const(1)
|
||||
b = UOp.const(2)
|
||||
graph_rewrite(a + b, pm, ctx=visited, bottom_up=True, walk=True)
|
||||
# bpm fires on each node before children: +, 1, 2
|
||||
self.assertEqual(visited, [Ops.ADD, 1, 2])
|
||||
|
||||
def test_walk_bottomup_unmatched_falls_through_to_children(self):
|
||||
"""Bottom-up walk: if bpm doesn't match a node, its children are still processed."""
|
||||
a = UOp.variable('a', 0, 10)
|
||||
b = UOp.variable('b', 0, 10)
|
||||
c = UOp.variable('c', 0, 10)
|
||||
# only a is in dvars, not a+4. bpm won't match a+4, so it descends and finds a.
|
||||
ret = graph_rewrite((a + 4) + (b + 5), _substitute, {a:c, b:c}, bottom_up=True, walk=True)
|
||||
self.assertIs(ret, (c + 4) + (c + 5))
|
||||
|
||||
# *** bidirectional walk: bpm fires before children, pm fires after rebuild ***
|
||||
|
||||
def test_walk_bidirectional_visit_order(self):
|
||||
"""Bidirectional walk: bpm fires pre-order, pm fires post-order."""
|
||||
visited = []
|
||||
def bpm_visit(ctx, x):
|
||||
ctx.append((x.val if x.op is Ops.CONST else x.op, "bpm"))
|
||||
return None
|
||||
def pm_visit(ctx, x):
|
||||
ctx.append((x.val if x.op is Ops.CONST else x.op, "pm"))
|
||||
return None
|
||||
bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_visit)])
|
||||
pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_visit)])
|
||||
a = UOp.const(1)
|
||||
b = UOp.const(2)
|
||||
graph_rewrite(a + b, pm, ctx=visited, bpm=bpm, walk=True)
|
||||
# bpm fires pre-order, pm fires post-order
|
||||
self.assertEqual(visited, [
|
||||
(Ops.ADD, "bpm"), (1, "bpm"), (1, "pm"), (2, "bpm"), (2, "pm"), (Ops.ADD, "pm"),
|
||||
])
|
||||
|
||||
def test_walk_bidirectional_bpm_short_circuits(self):
|
||||
"""If bpm matches, children are skipped and pm never fires on that node."""
|
||||
visited = []
|
||||
def bpm_match(ctx, x):
|
||||
ctx.append((x.val if x.op is Ops.CONST else x.op, "bpm"))
|
||||
# rewrite const(1) -> const(10), short-circuiting its subtree
|
||||
if x.op is Ops.CONST and x.val == 1: return UOp.const(10, x.dtype)
|
||||
return None
|
||||
def pm_match(ctx, x):
|
||||
ctx.append((x.val if x.op is Ops.CONST else x.op, "pm"))
|
||||
return None
|
||||
bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_match)])
|
||||
pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_match)])
|
||||
a = UOp.const(1)
|
||||
b = UOp.const(2)
|
||||
ret = graph_rewrite(a + b, pm, ctx=visited, bpm=bpm, walk=True)
|
||||
# bpm matches const(1) and short-circuits it, so pm never fires on const(1)
|
||||
self.assertNotIn((1, "pm"), visited)
|
||||
# but pm still fires on const(2) and the rebuilt ADD
|
||||
self.assertIn((2, "pm"), visited)
|
||||
self.assertIs(ret, UOp.const(10) + b)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
105
artifacts/package_sources/tinygrad/test/null/test_hcq_iface.py
Normal file
105
artifacts/package_sources/tinygrad/test/null/test_hcq_iface.py
Normal file
@@ -0,0 +1,105 @@
|
||||
import unittest, array, time
|
||||
from tinygrad.helpers import mv_address
|
||||
from tinygrad.runtime.support.hcq import MMIOInterface
|
||||
from tinygrad.runtime.support.usb import USBMMIOInterface
|
||||
from test.mockgpu.usb import MockUSB
|
||||
|
||||
class TestHCQIface(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.size = 4 << 10
|
||||
self.buffer = bytearray(self.size)
|
||||
self.mv = memoryview(self.buffer).cast('I')
|
||||
self.mmio = MMIOInterface(mv_address(self.mv), self.size, fmt='I')
|
||||
|
||||
def test_getitem_setitem(self):
|
||||
self.mmio[1] = 0xdeadbeef
|
||||
self.assertEqual(self.mmio[1], 0xdeadbeef)
|
||||
values = array.array('I', [10, 20, 30, 40])
|
||||
self.mmio[2:6] = values
|
||||
read_slice = self.mmio[2:6]
|
||||
# self.assertIsInstance(read_slice, array.array)
|
||||
self.assertEqual(read_slice, values.tolist())
|
||||
self.assertEqual(self.mv[2:6].tolist(), values.tolist())
|
||||
|
||||
def test_view(self):
|
||||
full = self.mmio.view()
|
||||
self.assertEqual(len(full), len(self.mmio))
|
||||
self.mmio[0] = 0x12345678
|
||||
self.assertEqual(full[0], 0x12345678)
|
||||
|
||||
# offset-only view
|
||||
self.mmio[1] = 0xdeadbeef
|
||||
off = self.mmio.view(offset=4)
|
||||
self.assertEqual(off[0], 0xdeadbeef)
|
||||
|
||||
# offset + size view: write into sub-view and confirm underlying buffer
|
||||
values = array.array('I', [11, 22, 33])
|
||||
sub = self.mmio.view(offset=8, size=12)
|
||||
sub[:] = values
|
||||
self.assertEqual(sub[:], values.tolist())
|
||||
self.assertEqual(self.mv[2:5].tolist(), values.tolist())
|
||||
|
||||
def test_speed(self):
|
||||
start = time.perf_counter()
|
||||
for i in range(10000):
|
||||
self.mmio[3:100] = array.array('I', [i] * 97)
|
||||
_ = self.mmio[3:100]
|
||||
end = time.perf_counter()
|
||||
|
||||
mvstart = time.perf_counter()
|
||||
for i in range(10000):
|
||||
self.mv[3:100] = array.array('I', [i] * 97)
|
||||
_ = self.mv[3:100].tolist()
|
||||
mvend = time.perf_counter()
|
||||
print(f"speed: hcq {end - start:.6f}s vs plain mv {mvend - mvstart:.6f}s")
|
||||
|
||||
class TestUSBMMIOInterface(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.size = 256
|
||||
self.buffer = bytearray(self.size)
|
||||
self.usb = MockUSB(self.buffer)
|
||||
self.mmio = USBMMIOInterface(self.usb, 0, self.size, fmt='B', pcimem=False)
|
||||
|
||||
def test_getitem_setitem_byte(self):
|
||||
self.mmio[1] = 0xAB
|
||||
self.assertEqual(self.mmio[1], 0xAB)
|
||||
self.assertEqual(self.usb.mem[1], 0xAB)
|
||||
|
||||
def test_slice_getitem_setitem(self):
|
||||
values = [1, 2, 3, 4]
|
||||
self.mmio[10:14] = values
|
||||
raw = self.mmio[10:14]
|
||||
self.assertIsInstance(raw, bytes)
|
||||
self.assertEqual(list(raw), values)
|
||||
self.assertEqual(list(self.usb.mem[10:14]), values)
|
||||
|
||||
def test_view(self):
|
||||
self.mmio[0] = 5
|
||||
view = self.mmio.view(offset=1, size=3)
|
||||
self.assertEqual(view[0], self.usb.mem[1])
|
||||
view[:] = [7, 8, 9]
|
||||
self.assertEqual(list(self.usb.mem[1:4]), [7, 8, 9])
|
||||
full_view = self.mmio.view()
|
||||
self.assertEqual(len(full_view), len(self.mmio))
|
||||
self.mmio[2] = 0xFE
|
||||
self.assertEqual(full_view[2], 0xFE)
|
||||
|
||||
def test_pcimem_dword(self):
|
||||
usb2 = MockUSB(bytearray(self.size))
|
||||
mmio_pci = USBMMIOInterface(usb2, 0, self.size, fmt='I', pcimem=True)
|
||||
mmio_pci[3] = 0x11223344
|
||||
self.assertEqual(mmio_pci[3], 0x11223344)
|
||||
self.assertEqual(usb2.mem[12:16], b'\x44\x33\x22\x11')
|
||||
|
||||
def test_pcimem_slice(self):
|
||||
usb3 = MockUSB(bytearray(self.size))
|
||||
mmio_pci = USBMMIOInterface(usb3, 0, self.size, fmt='B', pcimem=True)
|
||||
values = [2, 3, 4, 5]
|
||||
mmio_pci[4:8] = values
|
||||
raw = mmio_pci[4:8]
|
||||
self.assertIsInstance(raw, bytes)
|
||||
self.assertEqual(list(raw), values)
|
||||
self.assertEqual(list(usb3.mem[4:8]), values)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
413
artifacts/package_sources/tinygrad/test/null/test_helpers.py
Normal file
413
artifacts/package_sources/tinygrad/test/null/test_helpers.py
Normal file
@@ -0,0 +1,413 @@
|
||||
import ctypes, gzip, unittest, timeit, pickle
|
||||
from tinygrad import Variable
|
||||
from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, mv_address, count, all_same
|
||||
from tinygrad.tensor import is_numpy_ndarray
|
||||
from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
|
||||
from tinygrad.helpers import ceildiv, ansistrip, get_shape
|
||||
from tinygrad.tensor import Tensor
|
||||
import numpy as np
|
||||
|
||||
VARIABLE = ContextVar("VARIABLE", 0)
|
||||
|
||||
class TestContextVars(unittest.TestCase):
|
||||
# Ensuring that the test does not modify variables outside the tests.
|
||||
ctx = Context()
|
||||
def setUp(self): TestContextVars.ctx.__enter__()
|
||||
def tearDown(self): TestContextVars.ctx.__exit__()
|
||||
|
||||
def test_initial_value_is_set(self):
|
||||
_TMP = ContextVar("_TMP", 5)
|
||||
self.assertEqual(_TMP.value, 5)
|
||||
|
||||
def test_cannot_recreate(self):
|
||||
_TMP2 = ContextVar("_TMP2", 1)
|
||||
with self.assertRaises(RuntimeError):
|
||||
_TMP2 = ContextVar("_TMP2", 2)
|
||||
|
||||
def test_new_var_inside_context(self):
|
||||
with Context(VARIABLE=1):
|
||||
_TMP3 = ContextVar("_TMP3", 1)
|
||||
with self.assertRaises(RuntimeError):
|
||||
_TMP3 = ContextVar("_TMP3", 2)
|
||||
|
||||
def test_value_across_modules(self):
|
||||
# Mocking module import by invoking the code but not in our globals().
|
||||
exec('from tinygrad.helpers import ContextVar;C = ContextVar("C", 13)', {}) # pylint:disable=exec-used
|
||||
# It should not matter that the first creation was in another module.
|
||||
with self.assertRaises(RuntimeError):
|
||||
_C = ContextVar("C", 0)
|
||||
|
||||
def test_assignment_across_modules(self):
|
||||
B = ContextVar("B", 1)
|
||||
# local assignment
|
||||
B.value = 2
|
||||
self.assertEqual(B.value, 2)
|
||||
with self.assertRaises(RuntimeError):
|
||||
# Assignment in another module.
|
||||
exec('from tinygrad.helpers import ContextVar;B = ContextVar("B", 0);B.value = 3;', {}) # pylint:disable=exec-used
|
||||
|
||||
def test_context_assignment(self):
|
||||
with Context(VARIABLE=1):
|
||||
self.assertEqual(VARIABLE.value, 1)
|
||||
self.assertEqual(VARIABLE.value, 0)
|
||||
|
||||
def test_unknown_param_to_context(self):
|
||||
with self.assertRaises(KeyError):
|
||||
with Context(SOMETHING_ELSE=1):
|
||||
pass
|
||||
|
||||
def test_nested_context(self):
|
||||
with Context(VARIABLE=1):
|
||||
with Context(VARIABLE=2):
|
||||
MORE = ContextVar("MORE", 2)
|
||||
with Context(VARIABLE=3, MORE=3):
|
||||
self.assertEqual(VARIABLE.value, 3)
|
||||
self.assertEqual(MORE.value, 3)
|
||||
self.assertEqual(VARIABLE.value, 2)
|
||||
self.assertEqual(MORE.value, 2)
|
||||
self.assertEqual(VARIABLE.value, 1)
|
||||
self.assertEqual(MORE.value, 2) # TODO: should this raise?
|
||||
self.assertEqual(VARIABLE.value, 0)
|
||||
|
||||
def test_decorator(self):
|
||||
@Context(VARIABLE=1, DEBUG=4)
|
||||
def test():
|
||||
self.assertEqual(VARIABLE.value, 1)
|
||||
|
||||
self.assertEqual(VARIABLE.value, 0)
|
||||
test()
|
||||
self.assertEqual(VARIABLE.value, 0)
|
||||
|
||||
def test_context_exit_reverts_updated_values(self):
|
||||
D = ContextVar("D", 1)
|
||||
D.value = 2
|
||||
with Context(D=3):
|
||||
...
|
||||
assert D.value == 2, f"Expected D to be 2, but was {D.value}. Indicates that Context.__exit__ did not restore to the correct value."
|
||||
|
||||
class TestAllSame(unittest.TestCase):
|
||||
def test_empty(self): self.assertTrue(all_same([]))
|
||||
def test_single(self): self.assertTrue(all_same([1]))
|
||||
def test_same(self): self.assertTrue(all_same([1, 1, 1]))
|
||||
def test_different(self): self.assertFalse(all_same([1, 2, 1]))
|
||||
|
||||
class TestMergeDicts(unittest.TestCase):
|
||||
def test_merge_dicts(self):
|
||||
a = {"a": 1, "b": 2}
|
||||
b = {"a": 1, "c": 3}
|
||||
c = {}
|
||||
d = {"a": 2, "b": 2}
|
||||
assert merge_dicts([a, b]) == {"a": 1, "b": 2, "c": 3}
|
||||
assert merge_dicts([a, c]) == a
|
||||
assert merge_dicts([a, b, c]) == {"a": 1, "b": 2, "c": 3}
|
||||
with self.assertRaises(RuntimeError):
|
||||
merge_dicts([a, d])
|
||||
|
||||
class TestStripParens(unittest.TestCase):
|
||||
def test_simple(self): self.assertEqual("1+2", strip_parens("(1+2)"))
|
||||
def test_nested(self): self.assertEqual("1+(2+3)", strip_parens("(1+(2+3))"))
|
||||
def test_casted_no_strip(self): self.assertEqual("(int)(1+2)", strip_parens("(int)(1+2)"))
|
||||
def test_unmatched_parens(self): self.assertEqual("((c35+c39>>23&255)+-127).cast(dtypes.float)",
|
||||
strip_parens("((c35+c39>>23&255)+-127).cast(dtypes.float)"))
|
||||
def test_single_paren_left(self): self.assertEqual("(abc", strip_parens("(abc"))
|
||||
def test_single_paren_right(self): self.assertEqual("abc)", strip_parens("abc)"))
|
||||
def test_parens_at_different_depths(self): self.assertEqual("(a+(b))*(c)", strip_parens("(a+(b))*(c)"))
|
||||
|
||||
class TestProd(unittest.TestCase):
|
||||
def test_empty(self): self.assertEqual(1, prod(tuple()))
|
||||
def test_ints(self): self.assertEqual(30, prod((2, 3, 5)))
|
||||
def test_variable(self): self.assertEqual("(a*12)", prod((Variable("a", 1, 5), 3, 4)).render())
|
||||
def test_variable_order(self): self.assertEqual("(a*12)", prod((3, 4, Variable("a", 1, 5))).render())
|
||||
|
||||
class TestRoundUp(unittest.TestCase):
|
||||
def test_round_up(self):
|
||||
self.assertEqual(round_up(-3,4), 0)
|
||||
self.assertEqual(round_up(-4,4), -4)
|
||||
self.assertEqual(round_up(6,4), 8)
|
||||
self.assertEqual(round_up(8,4), 8)
|
||||
self.assertEqual(round_up(232, 24984), 24984)
|
||||
self.assertEqual(round_up(24984, 232), 25056)
|
||||
|
||||
class TestCeilDiv(unittest.TestCase):
|
||||
def test_int(self):
|
||||
self.assertEqual(ceildiv(10, 3), 4)
|
||||
self.assertEqual(ceildiv(9, 3), 3)
|
||||
self.assertEqual(ceildiv(0, 5), 0)
|
||||
self.assertEqual(ceildiv(1, 5), 1)
|
||||
def test_symbolic(self):
|
||||
# tests that ceildiv with UOp uses (num + amt - 1) // amt formula for non-negative num
|
||||
v = Variable('v', 0, 100)
|
||||
result = ceildiv(v, 6)
|
||||
self.assertEqual(result.render(), "((v+5)//6)")
|
||||
def test_symbolic_negative_offset(self):
|
||||
# tests ceildiv(v-5, 6) which is used in conv2d output shape
|
||||
# old implementation incorrectly simplified -(x//-y) to ((v+1)//6-1) for v-5
|
||||
# new implementation uses (v-5+5)//6 = v//6 which is correct
|
||||
v = Variable('v', 11, 100)
|
||||
result = ceildiv(v - 5, 6)
|
||||
self.assertEqual(result.render(), "(v//6)")
|
||||
|
||||
class TestCount(unittest.TestCase):
|
||||
def test_count_basic(self):
|
||||
c = count(3)
|
||||
self.assertEqual(next(c), 3)
|
||||
self.assertEqual(next(c), 4)
|
||||
|
||||
def test_count_step_pickle(self):
|
||||
c = count(1, 2)
|
||||
self.assertEqual(next(c), 1)
|
||||
c2 = pickle.loads(pickle.dumps(c))
|
||||
self.assertEqual(next(c2), 3)
|
||||
|
||||
@unittest.skip("no fetch tests because they need internet")
|
||||
class TestFetch(unittest.TestCase):
|
||||
def test_fetch_bad_http(self):
|
||||
self.assertRaises(Exception, fetch, 'http://www.google.com/404', allow_caching=False)
|
||||
|
||||
def test_fetch_small(self):
|
||||
assert (len(fetch('https://google.com', allow_caching=False).read_bytes())>0)
|
||||
|
||||
def test_fetch_img(self):
|
||||
from PIL import Image
|
||||
img = fetch("https://avatars.githubusercontent.com/u/132956020", allow_caching=False)
|
||||
with Image.open(img) as pimg:
|
||||
assert pimg.size == (77, 77), pimg.size
|
||||
|
||||
def test_fetch_subdir(self):
|
||||
from PIL import Image
|
||||
img = fetch("https://avatars.githubusercontent.com/u/132956020", allow_caching=False, subdir="images")
|
||||
with Image.open(img) as pimg:
|
||||
assert pimg.size == (77, 77), pimg.size
|
||||
assert img.parent.name == "images"
|
||||
|
||||
def test_fetch_gunzip_valid(self):
|
||||
# compare fetch(gunzip=True) to fetch(gunzip=False) plus decompressing afterwards
|
||||
gzip_url: str = 'https://ftp.gnu.org/gnu/gzip/gzip-1.13.tar.gz'
|
||||
fp_gz = fetch(gzip_url, gunzip=True)
|
||||
fp_no_gz = fetch(gzip_url, gunzip=False)
|
||||
with open(fp_gz, 'rb') as f: content_gz = f.read()
|
||||
with open(fp_no_gz, 'rb') as f: content_no_gz = gzip.decompress(f.read())
|
||||
assert fp_gz.stat().st_size > fp_no_gz.stat().st_size
|
||||
assert isinstance(content_gz, bytes) and isinstance(content_no_gz, bytes)
|
||||
assert len(content_gz) == len(content_no_gz)
|
||||
assert content_gz == content_no_gz
|
||||
|
||||
def test_fetch_gunzip_invalid(self):
|
||||
# given a non-gzipped file, fetch(gunzip=True) fails
|
||||
no_gzip_url: str = 'https://ftp.gnu.org/gnu/gzip/gzip-1.13.zip'
|
||||
with self.assertRaises(gzip.BadGzipFile):
|
||||
fetch(no_gzip_url, gunzip=True)
|
||||
|
||||
def test_fetch_user_agent(self):
|
||||
fetch("https://csrc.nist.gov/CSRC/media/Projects/lightweight-cryptography/documents/finalist-round/updated-submissions/sparkle.zip",
|
||||
allow_caching=False)
|
||||
|
||||
def test_fetch_half_and_full_file(self):
|
||||
x = fetch("https://csrc.nist.gov/CSRC/media/Projects/lightweight-cryptography/documents/finalist-round/updated-submissions/sparkle.zip",
|
||||
headers={"Range": "bytes=0-10"}).read_bytes()
|
||||
assert len(x) == 11, f"{len(x) != 11}"
|
||||
x = fetch("https://csrc.nist.gov/CSRC/media/Projects/lightweight-cryptography/documents/finalist-round/updated-submissions/sparkle.zip",
|
||||
headers={"Range": "bytes=0-100"}).read_bytes()
|
||||
assert len(x) == 101, f"{len(x) != 101}"
|
||||
|
||||
def test_fetch_sha(self):
|
||||
self.assertRaises(Exception, fetch, "https://ftp.gnu.org/gnu/gzip/gzip-1.13.tar.gz", allow_caching=False, sha256="a")
|
||||
fetch("https://ftp.gnu.org/gnu/gzip/gzip-1.13.tar.gz", allow_caching=False,
|
||||
sha256="20fc818aeebae87cdbf209d35141ad9d3cf312b35a5e6be61bfcfbf9eddd212a")
|
||||
|
||||
class TestFullyFlatten(unittest.TestCase):
|
||||
def test_fully_flatten(self):
|
||||
self.assertEqual(fully_flatten([[1, 3], [1, 2]]), [1, 3, 1, 2])
|
||||
self.assertEqual(fully_flatten(((1, 3), (1, 2))), [1, 3, 1, 2])
|
||||
self.assertEqual(fully_flatten([[[1], [3]], [[1], [2]]]), [1, 3, 1, 2])
|
||||
self.assertEqual(fully_flatten([[[[1], 2], 3], 4]), [1, 2, 3, 4])
|
||||
self.assertEqual(fully_flatten([[1, 2, [3, 4]], [5, 6], 7]), [1, 2, 3, 4, 5, 6, 7])
|
||||
self.assertEqual(fully_flatten([[1, "ab"], [True, None], [3.14, [5, "b"]]]), [1, "ab", True, None, 3.14, 5, "b"])
|
||||
|
||||
def test_fully_flatten_numpy(self):
|
||||
self.assertEqual(fully_flatten([np.array([])]), [])
|
||||
self.assertEqual(fully_flatten([np.array(3)]), [3])
|
||||
self.assertEqual(fully_flatten([np.array([3])]), [3])
|
||||
self.assertEqual(fully_flatten([np.array([[3]])]), [3])
|
||||
self.assertEqual(fully_flatten([np.array([1, 3]), np.array([1, 2])]), [1, 3, 1, 2])
|
||||
self.assertEqual(fully_flatten((np.array([1, 3]), np.array([1, 2]))), [1, 3, 1, 2])
|
||||
self.assertEqual(fully_flatten([np.array([[1], [3]]), np.array([[1], [2]])]), [1, 3, 1, 2])
|
||||
self.assertEqual(fully_flatten([[1, "ab"], [True, None], np.array([[3.14], [6.28]])]), [1, "ab", True, None, 3.14, 6.28])
|
||||
|
||||
class TestMemoryview(unittest.TestCase):
|
||||
def test_from_mv_to_mv(self):
|
||||
base = memoryview(bytearray(b"\x11\x22\x33"*40))
|
||||
ct = from_mv(base)
|
||||
mv = to_mv(ctypes.addressof(ct), len(base))
|
||||
mv[0] = 2
|
||||
assert base[0] == 2
|
||||
|
||||
@unittest.skip("allocates tons of memory")
|
||||
def test_to_mv(self):
|
||||
sizes = [
|
||||
(16, "16 B"),
|
||||
(64, "64 B"),
|
||||
(256, "256 B"),
|
||||
(1024, "1 KB"),
|
||||
(4 * 1024, "4 KB"),
|
||||
(16 * 1024, "16 KB"),
|
||||
(64 * 1024, "64 KB"),
|
||||
(256 * 1024, "256 KB"),
|
||||
(1 * 1024 * 1024, "1 MB"),
|
||||
(10 * 1024 * 1024, "10 MB"),
|
||||
(200 * 1024 * 1024, "200 MB"),
|
||||
]
|
||||
|
||||
for sz, label in sizes:
|
||||
buf = np.random.randint(0, 256, sz, dtype=np.uint8)
|
||||
ptr = buf.ctypes.data
|
||||
|
||||
iters = 100_000
|
||||
t_us = timeit.timeit(lambda: to_mv(ptr, sz), number=iters) * 1e6 / iters
|
||||
print(f"Size {label:>9} | Time: {t_us:8.3f} µs")
|
||||
|
||||
def test_speed_from_mv_vs_mv_address(self):
|
||||
x = memoryview(bytearray(1))
|
||||
|
||||
iters = 100000
|
||||
fmv_us = timeit.timeit(lambda: from_mv(x), number=iters) * 1e6 / iters
|
||||
mva_us = timeit.timeit(lambda: mv_address(x), number=iters) * 1e6 / iters
|
||||
print(f"from_mv vs mv_address: {fmv_us:8.3f} µs vs {mva_us:8.3f} µs")
|
||||
|
||||
class TestGetShape(unittest.TestCase):
|
||||
def test_get_shape(self):
|
||||
assert get_shape(2) == ()
|
||||
assert get_shape([]) == (0,)
|
||||
assert get_shape([[]]) == (1, 0)
|
||||
assert get_shape([[1, 2]]) == (1, 2)
|
||||
assert get_shape([[1, 2], (3, 4)]) == (2, 2)
|
||||
|
||||
def test_inhomogeneous_shape(self):
|
||||
with self.assertRaises(ValueError): get_shape([[], [1]])
|
||||
with self.assertRaises(ValueError): get_shape([[1, [2]], [1]])
|
||||
|
||||
class TestPolyN(unittest.TestCase):
|
||||
def test_float(self):
|
||||
np.testing.assert_allclose(polyN(1.0, [1.0, -2.0, 1.0]), 0.0)
|
||||
np.testing.assert_allclose(polyN(2.0, [1.0, -2.0, 1.0]), 1.0)
|
||||
np.testing.assert_allclose(polyN(3.0, [1.0, -2.0, 1.0]), 4.0)
|
||||
np.testing.assert_allclose(polyN(4.0, [1.0, -2.0, 1.0]), 9.0)
|
||||
|
||||
def test_uop(self):
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.uop.ops import UOp
|
||||
from test.helpers import eval_uop
|
||||
np.testing.assert_allclose(eval_uop(polyN(UOp.const(1.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 0.0)
|
||||
np.testing.assert_allclose(eval_uop(polyN(UOp.const(2.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 1.0)
|
||||
np.testing.assert_allclose(eval_uop(polyN(UOp.const(3.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 4.0)
|
||||
np.testing.assert_allclose(eval_uop(polyN(UOp.const(4.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 9.0)
|
||||
|
||||
class TestTimeToStr(unittest.TestCase):
|
||||
def test_seconds(self): self.assertEqual(" 10.01s ", time_to_str(10.01))
|
||||
def test_boundary_sec_ms(self): self.assertEqual("10000.00ms", time_to_str(10))
|
||||
def test_milliseconds(self): self.assertEqual(" 500.00ms", time_to_str(0.5))
|
||||
def test_boundary_ms_us(self): self.assertEqual("10000.00us", time_to_str(0.01))
|
||||
def test_microseconds(self): self.assertEqual(" 100.00us", time_to_str(0.0001))
|
||||
def test_zero(self): self.assertEqual(" 0.00us", time_to_str(0))
|
||||
def test_width_formatting(self): self.assertEqual(" 10.01s ", time_to_str(10.01, w=6))
|
||||
|
||||
class TestCStyleDivMod(unittest.TestCase):
|
||||
def test_div_pos(self):
|
||||
self.assertEqual(cdiv(-9, 5), -1)
|
||||
self.assertEqual(cdiv(-4, 5), 0)
|
||||
self.assertEqual(cdiv(0, 5), 0)
|
||||
self.assertEqual(cdiv(4, 5), 0)
|
||||
self.assertEqual(cdiv(9, 5), 1)
|
||||
def test_div_neg(self):
|
||||
self.assertEqual(cdiv(-9, -5), 1)
|
||||
self.assertEqual(cdiv(-4, -5), 0)
|
||||
self.assertEqual(cdiv(0, -5), 0)
|
||||
self.assertEqual(cdiv(4, -5), 0)
|
||||
self.assertEqual(cdiv(9, -5), -1)
|
||||
def test_mod_pos(self):
|
||||
self.assertEqual(cmod(-9, 5), -4)
|
||||
self.assertEqual(cmod(-4, 5), -4)
|
||||
self.assertEqual(cmod(0, 5), 0)
|
||||
self.assertEqual(cmod(4, 5), 4)
|
||||
self.assertEqual(cmod(9, 5), 4)
|
||||
def test_mod_neg(self):
|
||||
self.assertEqual(cmod(-9, -5), -4)
|
||||
self.assertEqual(cmod(-4, -5), -4)
|
||||
self.assertEqual(cmod(0, -5), 0)
|
||||
self.assertEqual(cmod(4, -5), 4)
|
||||
self.assertEqual(cmod(9, -5), 4)
|
||||
|
||||
class TestGetBits(unittest.TestCase):
|
||||
def test_low_bits(self):
|
||||
self.assertEqual(getbits(0b11010110, 0, 3), 0b0110)
|
||||
|
||||
def test_high_bits(self):
|
||||
self.assertEqual(getbits(0b11010110, 4, 7), 0b1101)
|
||||
|
||||
def test_middle_bits(self):
|
||||
self.assertEqual(getbits(0b11010110, 3, 5), 0b010)
|
||||
|
||||
def test_full_range(self):
|
||||
self.assertEqual(getbits(0b11010110, 0, 7), 0b11010110)
|
||||
|
||||
def test_single_bit(self):
|
||||
self.assertEqual(getbits(0b100000000, 8, 8), 1)
|
||||
|
||||
class TestArgFix(unittest.TestCase):
|
||||
def test_none(self):
|
||||
self.assertEqual(argfix(None), (None, ))
|
||||
self.assertEqual(argfix(None, None), (None, None))
|
||||
def test_positional_arguments(self):
|
||||
self.assertEqual(argfix(1, 2, 3), (1, 2, 3))
|
||||
def test_tuple(self):
|
||||
self.assertEqual(argfix((1., 2., 3.)), (1., 2., 3.))
|
||||
def test_list(self):
|
||||
self.assertEqual(argfix([True, False]), (True, False))
|
||||
|
||||
class TestWordWrap(unittest.TestCase):
|
||||
def test_wrap_simple(self):
|
||||
wrap = 10
|
||||
st = "x"*wrap*2
|
||||
st2 = word_wrap(st, wrap)
|
||||
self.assertEqual(len(st2.splitlines()), 2)
|
||||
|
||||
def test_wrap_colored(self):
|
||||
wrap = 10
|
||||
st = colored("x"*wrap*2, "red")
|
||||
st2 = word_wrap(st, wrap=wrap)
|
||||
self.assertEqual(len(st2.splitlines()), 2)
|
||||
|
||||
def test_wrap_colored_at_boundary(self):
|
||||
wrap = 10
|
||||
st = "x"*(wrap-2) + colored("yyy", "red")
|
||||
st2 = word_wrap(st, wrap=wrap)
|
||||
self.assertEqual(ansistrip(st2), "x"*(wrap-2)+"yy\ny")
|
||||
self.assertNotIn("\x1b[\n", st2)
|
||||
|
||||
def test_wrap_explicit_newline(self):
|
||||
wrap = 10
|
||||
st = "\n".join(["x"*wrap, "x"*wrap, "x"*wrap])
|
||||
st2 = word_wrap(st, wrap=wrap)
|
||||
self.assertEqual(len(st2.splitlines()), len(st.splitlines()))
|
||||
|
||||
st = "\n".join(["x"*(wrap+1), "x"*wrap, "x"*wrap])
|
||||
st2 = word_wrap(st, wrap=wrap)
|
||||
self.assertEqual(len(st2.splitlines()), len(st.splitlines())+1)
|
||||
|
||||
st = "\n".join(["x"*(wrap+1), "x"*(wrap+1), "x"*(wrap+1)])
|
||||
st2 = word_wrap(st, wrap=wrap)
|
||||
self.assertEqual(len(st2.splitlines()), len(st.splitlines())+3)
|
||||
|
||||
class TestIsNumpyNdarray(unittest.TestCase):
|
||||
def test_ndarray(self):
|
||||
self.assertTrue(is_numpy_ndarray(np.array([1, 2, 3])))
|
||||
def test_ndarray_tolist(self):
|
||||
self.assertFalse(is_numpy_ndarray(np.array([1, 2, 3]).tolist()))
|
||||
def test_list(self):
|
||||
self.assertFalse(is_numpy_ndarray([1, 2, 3]))
|
||||
def test_tensor(self):
|
||||
self.assertFalse(is_numpy_ndarray(Tensor([1, 2, 3])))
|
||||
self.assertFalse(is_numpy_ndarray(Tensor(np.array([1, 2, 3]))))
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
100
artifacts/package_sources/tinygrad/test/null/test_indexing.py
Normal file
100
artifacts/package_sources/tinygrad/test/null/test_indexing.py
Normal file
@@ -0,0 +1,100 @@
|
||||
# test cases are modified from pytorch test_indexing.py
|
||||
|
||||
import unittest
|
||||
|
||||
from tinygrad import Tensor
|
||||
|
||||
class TestIndexing(unittest.TestCase):
|
||||
def test_single_int(self):
|
||||
v = Tensor.randn(5, 7, 3)
|
||||
self.assertEqual(v[4].shape, (7, 3))
|
||||
|
||||
def test_multiple_int(self):
|
||||
v = Tensor.randn(5, 7, 3)
|
||||
self.assertEqual(v[4].shape, (7, 3))
|
||||
self.assertEqual(v[4, :, 1].shape, (7,))
|
||||
|
||||
def test_none(self):
|
||||
v = Tensor.randn(5, 7, 3)
|
||||
self.assertEqual(v[None].shape, (1, 5, 7, 3))
|
||||
self.assertEqual(v[:, None].shape, (5, 1, 7, 3))
|
||||
self.assertEqual(v[:, None, None].shape, (5, 1, 1, 7, 3))
|
||||
self.assertEqual(v[..., None].shape, (5, 7, 3, 1))
|
||||
|
||||
def test_int_indices(self):
|
||||
v = Tensor.randn(5, 7, 3)
|
||||
self.assertEqual(v[[0, 4, 2]].shape, (3, 7, 3))
|
||||
self.assertEqual(v[:, [0, 4, 2]].shape, (5, 3, 3))
|
||||
self.assertEqual(v[:, [[0, 1], [4, 3]]].shape, (5, 2, 2, 3))
|
||||
|
||||
def test_index_src_datatype(self):
|
||||
src = Tensor.ones(3, 2, 4)
|
||||
# test index
|
||||
res = src[[0, 2, 1], :, :]
|
||||
self.assertEqual(res.shape, src.shape)
|
||||
|
||||
def test_empty_slice(self):
|
||||
x = Tensor.randn(2, 3, 4, 5)
|
||||
y = x[:, :, :, 1]
|
||||
z = y[:, 1:1, :]
|
||||
self.assertEqual((2, 0, 4), z.shape)
|
||||
|
||||
def test_invalid_index(self):
|
||||
x = Tensor.arange(0, 16).reshape(4, 4)
|
||||
self.assertRaises(TypeError, lambda: x["0":"1"])
|
||||
|
||||
def test_out_of_bound_index(self):
|
||||
x = Tensor.arange(0, 100).reshape(2, 5, 10)
|
||||
self.assertRaises(IndexError, lambda: x[0, 5])
|
||||
self.assertRaises(IndexError, lambda: x[4, 5])
|
||||
self.assertRaises(IndexError, lambda: x[0, 1, 15])
|
||||
self.assertRaises(IndexError, lambda: x[:, :, 12])
|
||||
|
||||
class TestNumpy(unittest.TestCase):
|
||||
def test_index_no_floats(self):
|
||||
a = Tensor([[[5.]]])
|
||||
|
||||
self.assertRaises(IndexError, lambda: a[0.0])
|
||||
self.assertRaises(IndexError, lambda: a[0, 0.0])
|
||||
self.assertRaises(IndexError, lambda: a[0.0, 0])
|
||||
self.assertRaises(IndexError, lambda: a[0.0, :])
|
||||
self.assertRaises(IndexError, lambda: a[:, 0.0])
|
||||
self.assertRaises(IndexError, lambda: a[:, 0.0, :])
|
||||
self.assertRaises(IndexError, lambda: a[0.0, :, :])
|
||||
self.assertRaises(IndexError, lambda: a[0, 0, 0.0])
|
||||
self.assertRaises(IndexError, lambda: a[0.0, 0, 0])
|
||||
self.assertRaises(IndexError, lambda: a[0, 0.0, 0])
|
||||
self.assertRaises(IndexError, lambda: a[-1.4])
|
||||
self.assertRaises(IndexError, lambda: a[0, -1.4])
|
||||
self.assertRaises(IndexError, lambda: a[-1.4, 0])
|
||||
self.assertRaises(IndexError, lambda: a[-1.4, :])
|
||||
self.assertRaises(IndexError, lambda: a[:, -1.4])
|
||||
self.assertRaises(IndexError, lambda: a[:, -1.4, :])
|
||||
self.assertRaises(IndexError, lambda: a[-1.4, :, :])
|
||||
self.assertRaises(IndexError, lambda: a[0, 0, -1.4])
|
||||
self.assertRaises(IndexError, lambda: a[-1.4, 0, 0])
|
||||
self.assertRaises(IndexError, lambda: a[0, -1.4, 0])
|
||||
# these two trigger slice internal type verification first
|
||||
self.assertRaises(TypeError, lambda: a[0.0:, 0.0])
|
||||
self.assertRaises(TypeError, lambda: a[0.0:, 0.0,:])
|
||||
|
||||
def test_none_index(self):
|
||||
# `None` index adds newaxis
|
||||
a = Tensor([1, 2, 3])
|
||||
self.assertEqual(a[None].ndim, a.ndim+1)
|
||||
|
||||
def test_everything_returns_views(self):
|
||||
# Before `...` would return a itself.
|
||||
a = Tensor([5])
|
||||
|
||||
self.assertIs(a, a[()])
|
||||
self.assertIs(a, a[...])
|
||||
self.assertIs(a, a[:])
|
||||
|
||||
def test_broaderrors_indexing(self):
|
||||
a = Tensor.zeros(5, 5)
|
||||
self.assertRaises(IndexError, a.__getitem__, ([0, 1], [0, 1, 2]))
|
||||
self.assertRaises(IndexError, a.contiguous().__setitem__, ([0, 1], [0, 1, 2]), 0)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,25 @@
|
||||
# ruff: noqa: E501
|
||||
import unittest
|
||||
from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.codegen import to_program
|
||||
|
||||
class TestLinearizerFailures(unittest.TestCase):
|
||||
def test_fail_1(self):
|
||||
c0 = UOp.param(0, dtypes.float, (64,))
|
||||
c1 = UOp.range(UOp.const(2), 1, AxisType.WEAK)
|
||||
c2 = UOp.range(UOp.const(32), 2, AxisType.WEAK)
|
||||
c3 = ((c1*UOp.const(32))+c2)
|
||||
c4 = UOp.param(1, dtypes.float, (163840,))
|
||||
c5 = UOp.range(UOp.const(2560), 0, AxisType.REDUCE)
|
||||
c6 = c4.index(((((((c5//UOp.const(8))%UOp.const(8))*UOp.const(8))+(c5%UOp.const(8)))+(((c2*UOp.const(40))+(c5//UOp.const(64)))*UOp.const(64)))+(c1*UOp.const(81920))))
|
||||
c7 = UOp.param(2, dtypes.float, (64,))
|
||||
c8 = c7.index(c3)
|
||||
c9 = ((((c6+(c8*UOp.const(-1.0)))*(c6+(c8*UOp.const(-1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(0.000390625))+UOp.const(1e-05)).sqrt().reciprocal()
|
||||
c10 = c0.index(c3).store(c9).end(c1, c2)
|
||||
ast = c10.sink(arg=KernelInfo())
|
||||
to_program(ast, renderer=Device[Device.DEFAULT].renderer)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,45 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor, Context, Device
|
||||
from tinygrad.codegen import to_program
|
||||
from tinygrad.codegen.opt import Opt, OptOps
|
||||
from tinygrad.uop.ops import KernelInfo
|
||||
|
||||
class TestLinearizerRewrite(unittest.TestCase):
|
||||
def test_reduction(self):
|
||||
t = Tensor.ones((64,64), device="NULL").contiguous().realize()
|
||||
out = (t*2).sum(axis=1)
|
||||
with Context(SPLIT_REDUCEOP=0):
|
||||
si = out.schedule_linear().src[-1]
|
||||
opts_to_apply = []
|
||||
opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4))
|
||||
opts_to_apply.append(Opt(OptOps.UNROLL, 0, 4))
|
||||
ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
|
||||
prg = to_program(ast, Device["CPU"].renderer)
|
||||
print(prg.src[2].arg)
|
||||
|
||||
def test_arange(self):
|
||||
out = Tensor.arange(32).clone("NULL")
|
||||
with Context(SPLIT_REDUCEOP=0):
|
||||
si = out.schedule_linear().src[-1]
|
||||
opts_to_apply = []
|
||||
opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4))
|
||||
ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
|
||||
prg = to_program(ast, Device["CPU"].renderer)
|
||||
print(prg.src[2].arg)
|
||||
|
||||
def test_kernel_info(self):
|
||||
out = Tensor.arange(4).clone("NULL")
|
||||
si = out.schedule_linear().src[-1]
|
||||
|
||||
ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=()))
|
||||
prg = to_program(ast, Device["CPU"].renderer)
|
||||
assert prg.src[0].arg.applied_opts == (), f"expected no opts, got {prg}"
|
||||
|
||||
#prg = to_program(ast.replace(arg=KernelInfo()), Device["CPU"].renderer)
|
||||
#assert prg.src[0].arg.applied_opts != (), f"expected opts to apply, got {prg.src[0].arg.applied_opts}"
|
||||
|
||||
prg = to_program(ast.replace(arg=KernelInfo(name="custom")), Device["CPU"].renderer)
|
||||
self.assertEqual(prg.arg.name, "custom")
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
290
artifacts/package_sources/tinygrad/test/null/test_llm_server.py
Normal file
290
artifacts/package_sources/tinygrad/test/null/test_llm_server.py
Normal file
@@ -0,0 +1,290 @@
|
||||
import unittest, threading, time, json
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
class TestLLMServer(unittest.TestCase):
|
||||
"""Integration tests using the real OpenAI client."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.mock_tok = Mock()
|
||||
cls.mock_tok.encode = Mock(return_value=[200, 201, 202])
|
||||
cls.mock_tok.decode = Mock(return_value="Hello")
|
||||
cls.mock_tok.stream_decoder = Mock(return_value=lambda tid=None: "Hello" if tid is not None else "")
|
||||
cls.mock_tok.preset = "llama3"
|
||||
cls.mock_tok.bos_id = 1
|
||||
cls.mock_tok.eos_id = 999
|
||||
cls.mock_tok.eot_id = None
|
||||
cls.mock_tok.is_end = Mock(side_effect=lambda tid: tid in (999,))
|
||||
|
||||
cls.mock_model = Mock()
|
||||
cls.mock_model.max_context = 4
|
||||
cls.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 999]))
|
||||
cls.mock_model.get_start_pos = Mock(return_value=0)
|
||||
|
||||
from tinygrad.llm.cli import FallbackTemplate
|
||||
from tinygrad.llm.serve import LLMServer
|
||||
|
||||
cls.server = LLMServer(('127.0.0.1', 0), cls.mock_model, "test-model", cls.mock_tok, FallbackTemplate(cls.mock_tok))
|
||||
cls.port = cls.server.server_address[1]
|
||||
cls.server_thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
|
||||
cls.server_thread.start()
|
||||
time.sleep(0.1)
|
||||
|
||||
from openai import OpenAI
|
||||
cls.client = OpenAI(base_url=f"http://127.0.0.1:{cls.port}/v1", api_key="test")
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
cls.server.shutdown()
|
||||
cls.server.server_close()
|
||||
|
||||
def test_chat_completion_stream(self):
|
||||
stream = self.client.chat.completions.create(
|
||||
model="test",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
stream=True
|
||||
)
|
||||
|
||||
chunks = list(stream)
|
||||
self.assertGreater(len(chunks), 0)
|
||||
self.assertEqual(chunks[0].choices[0].delta.role, "assistant")
|
||||
self.assertEqual(chunks[-1].choices[0].finish_reason, "stop")
|
||||
|
||||
def test_openai_response_structure(self):
|
||||
stream = self.client.chat.completions.create(
|
||||
model="test-model",
|
||||
messages=[{"role": "user", "content": "Test"}],
|
||||
stream=True
|
||||
)
|
||||
|
||||
for chunk in stream:
|
||||
self.assertTrue(chunk.id.startswith("chatcmpl-"))
|
||||
self.assertEqual(chunk.object, "chat.completion.chunk")
|
||||
self.assertIsNotNone(chunk.choices)
|
||||
self.assertIsNotNone(chunk.created)
|
||||
self.assertIsInstance(chunk.created, int)
|
||||
self.assertEqual(chunk.model, "test-model")
|
||||
|
||||
def test_stream_with_usage(self):
|
||||
def generate(ids, **kwargs):
|
||||
for token in (300, 301, 999):
|
||||
ids.append(token)
|
||||
yield token
|
||||
with patch.object(self.mock_model, "generate", side_effect=generate):
|
||||
chunks = list(self.client.chat.completions.create(
|
||||
model="test", messages=[{"role": "user", "content": "Hello"}], stream=True, stream_options={"include_usage": True}))
|
||||
last_chunk = chunks[-1]
|
||||
|
||||
self.assertEqual(last_chunk.usage.prompt_tokens, 3)
|
||||
self.assertEqual(last_chunk.usage.completion_tokens, 2)
|
||||
self.assertEqual(last_chunk.usage.total_tokens, 5)
|
||||
|
||||
def test_multi_turn_conversation(self):
|
||||
stream = self.client.chat.completions.create(
|
||||
model="test",
|
||||
messages=[
|
||||
{"role": "system", "content": "You are helpful."},
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "assistant", "content": "Hi!"},
|
||||
{"role": "user", "content": "How are you?"}
|
||||
],
|
||||
stream=True
|
||||
)
|
||||
|
||||
chunks = list(stream)
|
||||
self.assertGreater(len(chunks), 0)
|
||||
self.assertEqual(chunks[-1].choices[0].finish_reason, "stop")
|
||||
|
||||
def test_content_is_streamed(self):
|
||||
stream = self.client.chat.completions.create(
|
||||
model="test",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
stream=True
|
||||
)
|
||||
|
||||
contents = []
|
||||
for chunk in stream:
|
||||
if chunk.choices and chunk.choices[0].delta.content:
|
||||
contents.append(chunk.choices[0].delta.content)
|
||||
|
||||
self.assertGreater(len(contents), 0)
|
||||
|
||||
def test_interrupted_stream_logs_tokens(self):
|
||||
with patch.object(self.mock_model, "generate", side_effect=lambda ids, **kwargs: iter([300, 301, 999])), \
|
||||
patch("tinygrad.llm.serve.stderr_log") as log, patch("tinygrad.llm.serve.colored", side_effect=lambda text, color: text) as color:
|
||||
stream = self.server.RequestHandlerClass.run_model(Mock(server=self.server), [200, 201, 202], "test")
|
||||
next(stream)
|
||||
next(stream)
|
||||
stream.close()
|
||||
interrupt = log.call_args.args[0]
|
||||
self.assertFalse(interrupt.startswith("\n"))
|
||||
self.assertTrue(interrupt.endswith("\n"))
|
||||
self.assertIn("gen:", interrupt)
|
||||
self.assertIn("out: 1", interrupt)
|
||||
self.assertTrue(any(args[0].startswith("total:") and args[1] == "red" for args, _ in color.call_args_list))
|
||||
|
||||
def test_stream_disconnect_closes_source(self):
|
||||
from tinygrad.llm.serve import Handler
|
||||
source, handler = Mock(), Mock()
|
||||
source.__iter__ = Mock(return_value=iter([{}]))
|
||||
handler.wfile.write.side_effect = BrokenPipeError
|
||||
Handler.stream_json(handler, source)
|
||||
source.close.assert_called_once()
|
||||
|
||||
def test_non_streaming(self):
|
||||
resp = self.client.chat.completions.create(
|
||||
model="test-model",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
stream=False
|
||||
)
|
||||
|
||||
self.assertTrue(resp.id.startswith("chatcmpl-"))
|
||||
self.assertEqual(resp.object, "chat.completion")
|
||||
self.assertEqual(resp.model, "test-model")
|
||||
self.assertIsNotNone(resp.created)
|
||||
self.assertEqual(len(resp.choices), 1)
|
||||
self.assertEqual(resp.choices[0].message.role, "assistant")
|
||||
self.assertIsNotNone(resp.choices[0].message.content)
|
||||
self.assertEqual(resp.choices[0].finish_reason, "stop")
|
||||
self.assertIsNotNone(resp.usage)
|
||||
self.assertIsNotNone(resp.usage.prompt_tokens)
|
||||
self.assertIsNotNone(resp.usage.completion_tokens)
|
||||
|
||||
def test_context_length_error(self):
|
||||
from openai import BadRequestError
|
||||
self.mock_tok.encode.return_value = [200, 201, 202, 203]
|
||||
try:
|
||||
with self.assertRaises(BadRequestError) as err:
|
||||
self.client.chat.completions.create(model="test-model", messages=[{"role":"user", "content":"too long"}])
|
||||
self.assertEqual(err.exception.code, "context_length_exceeded")
|
||||
finally:
|
||||
self.mock_tok.encode.return_value = [200, 201, 202]
|
||||
|
||||
def test_max_tokens_streaming(self):
|
||||
self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 302, 303, 999]))
|
||||
stream = self.client.chat.completions.create(
|
||||
model="test", messages=[{"role": "user", "content": "Hello"}], stream=True, max_tokens=2
|
||||
)
|
||||
chunks = list(stream)
|
||||
content_chunks = [c for c in chunks if c.choices and c.choices[0].delta.content]
|
||||
self.assertEqual(len(content_chunks), 2)
|
||||
self.assertEqual(chunks[-1].choices[0].finish_reason, "length")
|
||||
|
||||
def test_max_tokens_non_streaming(self):
|
||||
self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 302, 303, 999]))
|
||||
resp = self.client.chat.completions.create(
|
||||
model="test", messages=[{"role": "user", "content": "Hello"}], stream=False, max_tokens=2
|
||||
)
|
||||
self.assertEqual(resp.choices[0].finish_reason, "length")
|
||||
self.assertEqual(resp.usage.completion_tokens, 2)
|
||||
|
||||
def test_models_endpoint(self):
|
||||
import requests as req
|
||||
resp = req.get(f"http://127.0.0.1:{self.port}/v1/models")
|
||||
self.assertEqual(resp.status_code, 200)
|
||||
data = resp.json()
|
||||
self.assertEqual(data["object"], "list")
|
||||
self.assertEqual(len(data["data"]), 1)
|
||||
self.assertEqual(data["data"][0]["id"], "test-model")
|
||||
self.assertEqual(data["data"][0]["object"], "model")
|
||||
|
||||
class TestLLMToolCalls(unittest.TestCase):
|
||||
"""Tool calling through the OpenAI-compatible HTTP API."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.mock_tok = Mock()
|
||||
cls.mock_tok.encode = Mock(return_value=[200, 201, 202])
|
||||
cls.mock_tok.decode = Mock(return_value="")
|
||||
cls.mock_tok.preset = "qwen2"
|
||||
cls.mock_tok.bos_id, cls.mock_tok.eos_id, cls.mock_tok.eot_id = None, 999, None
|
||||
cls.mock_tok.is_end = Mock(return_value=False)
|
||||
|
||||
cls.mock_model = Mock()
|
||||
cls.mock_model.max_context = 4
|
||||
cls.mock_model.get_start_pos = Mock(return_value=0)
|
||||
|
||||
from tinygrad.llm.serve import LLMServer
|
||||
import jinja2
|
||||
# .items() matches tool-aware templates and ensures OpenAI JSON argument strings are normalized before rendering the next turn.
|
||||
template = jinja2.Template("""{% for m in messages %}{{ m.content or '' }}{% for tc in m.tool_calls or [] %}
|
||||
{% for key, value in tc.function.arguments.items() %}{{ key }}={{ value }}{% endfor %}{% endfor %}{% endfor %}""")
|
||||
cls.server = LLMServer(('127.0.0.1', 0), cls.mock_model, "tool-model", cls.mock_tok, template)
|
||||
cls.port = cls.server.server_address[1]
|
||||
cls.server_thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
|
||||
cls.server_thread.start()
|
||||
time.sleep(0.1)
|
||||
|
||||
from openai import OpenAI
|
||||
cls.client = OpenAI(base_url=f"http://127.0.0.1:{cls.port}/v1", api_key="test")
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
cls.server.shutdown()
|
||||
cls.server.server_close()
|
||||
|
||||
def set_output(self, text:str):
|
||||
pieces = dict(enumerate(text, 1))
|
||||
self.mock_tok.stream_decoder = Mock(return_value=lambda tid=None: pieces[tid] if tid is not None else "")
|
||||
self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter(pieces))
|
||||
|
||||
@staticmethod
|
||||
def tools():
|
||||
return [{"type":"function", "function":{"name":"read", "description":"Read a file",
|
||||
"parameters":{"type":"object", "properties":{"path":{"type":"string"}}, "required":["path"]}}}]
|
||||
|
||||
def test_streaming_tool_call(self):
|
||||
self.set_output('before<tool_call>{"name":"read","arguments":{"path":"README.md"}}</tool_call>')
|
||||
chunks = list(self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Read README.md"}],
|
||||
tools=self.tools(), stream=True))
|
||||
self.assertEqual("".join(c.choices[0].delta.content or "" for c in chunks if c.choices), "before")
|
||||
calls = [tc for c in chunks if c.choices for tc in c.choices[0].delta.tool_calls or []]
|
||||
self.assertEqual(len(calls), 1)
|
||||
self.assertEqual(calls[0].function.name, "read")
|
||||
self.assertEqual(json.loads(calls[0].function.arguments), {"path":"README.md"})
|
||||
self.assertEqual(chunks[-1].choices[0].finish_reason, "tool_calls")
|
||||
|
||||
def test_multiple_xml_tool_calls(self):
|
||||
self.set_output("<tool_call><function=read><parameter=path>\"a\"</parameter></function></tool_call>"
|
||||
"<tool_call><function=read><parameter=path>\"b\"</parameter></function></tool_call>")
|
||||
response = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Read a and b"}],
|
||||
tools=self.tools())
|
||||
self.assertEqual([json.loads(tc.function.arguments)["path"] for tc in response.choices[0].message.tool_calls], ["a", "b"])
|
||||
self.assertEqual(response.choices[0].finish_reason, "tool_calls")
|
||||
|
||||
def test_multiline_tool_argument_preserves_trailing_newline(self):
|
||||
self.set_output("<tool_call>\n<function=write>\n<parameter=content>\nfirst\nsecond\n\n</parameter>\n"
|
||||
"<parameter=filePath>\nout.txt\n</parameter>\n</function>\n</tool_call>")
|
||||
response = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Write out.txt"}], tools=self.tools())
|
||||
args = json.loads(response.choices[0].message.tool_calls[0].function.arguments)
|
||||
self.assertEqual(args, {"content":"first\nsecond\n", "filePath":"out.txt"})
|
||||
|
||||
def test_invalid_tool_call_becomes_content(self):
|
||||
self.set_output("<tool_call>not a call</tool_call>")
|
||||
response = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Hello"}], tools=self.tools())
|
||||
self.assertEqual(response.choices[0].message.content, "<tool_call>not a call</tool_call>")
|
||||
self.assertIsNone(response.choices[0].message.tool_calls)
|
||||
self.assertEqual(response.choices[0].finish_reason, "stop")
|
||||
|
||||
def test_tool_call_in_reasoning_is_not_executed(self):
|
||||
self.set_output('<think>draft <tool_call>{"name":"wrong","arguments":{}}</tool_call></think>answer')
|
||||
response = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Hello"}], tools=self.tools())
|
||||
self.assertEqual(response.choices[0].message.content, "answer")
|
||||
self.assertIsNone(response.choices[0].message.tool_calls)
|
||||
self.assertEqual(response.choices[0].finish_reason, "stop")
|
||||
|
||||
def test_tool_result_round_trip(self):
|
||||
self.set_output('<tool_call>{"name":"read","arguments":{"path":"README.md"}}</tool_call>')
|
||||
first = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Read README.md"}], tools=self.tools())
|
||||
call = first.choices[0].message.tool_calls[0]
|
||||
self.set_output("done")
|
||||
second = self.client.chat.completions.create(model="tool-model", messages=[
|
||||
{"role":"user", "content":"Read README.md"},
|
||||
{"role":"assistant", "content":None, "tool_calls":[call.model_dump()]},
|
||||
{"role":"tool", "tool_call_id":call.id, "content":"file contents"},
|
||||
], tools=self.tools())
|
||||
self.assertEqual(second.choices[0].message.content, "done")
|
||||
self.assertEqual(second.choices[0].finish_reason, "stop")
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,108 @@
|
||||
import unittest, base64, functools, re, sys, time, unicodedata
|
||||
from tinygrad.llm.cli import SimpleTokenizer, FallbackTemplate
|
||||
from tinygrad.helpers import fetch
|
||||
|
||||
@unittest.skipIf(sys.platform == 'win32', "fetch race condition on Windows")
|
||||
class TestLLMTokenizer(unittest.TestCase):
|
||||
@functools.cached_property
|
||||
def llama_tok(self):
|
||||
# from https://github.com/tinygrad/tinygrad/blob/e0106b6b257ebc003eb3694144e3e198f7d8cc37/examples/llama3.py#L14
|
||||
model_file = fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model")
|
||||
with open(model_file, "rt") as fd:
|
||||
str_vocab = [line.split(maxsplit=1) for line in fd.read().splitlines() if line]
|
||||
|
||||
# https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
|
||||
bs = [*range(33, 127), *range(161, 173), *range(174, 256)] # bytes that map to themselves
|
||||
_byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)}
|
||||
_byte_encoder = {v:k for k,v in _byte_decoder.items()}
|
||||
normal_tokens = {''.join([_byte_encoder[x] for x in base64.b64decode(stok)]): int(srank) for stok, srank in str_vocab}
|
||||
|
||||
special_tokens = [
|
||||
"<|begin_of_text|>",
|
||||
"<|end_of_text|>",
|
||||
"<|reserved_special_token_0|>",
|
||||
"<|reserved_special_token_1|>",
|
||||
"<|reserved_special_token_2|>",
|
||||
"<|reserved_special_token_3|>",
|
||||
"<|start_header_id|>",
|
||||
"<|end_header_id|>",
|
||||
"<|reserved_special_token_4|>",
|
||||
"<|eot_id|>",
|
||||
] + [ f"<|reserved_special_token_{i}|>" for i in range(5, 256 - 5) ]
|
||||
return SimpleTokenizer(normal_tokens, {token: len(normal_tokens) + i for i, token in enumerate(special_tokens)})
|
||||
|
||||
def _test_coding(self, tok: SimpleTokenizer, text: str, expected_tokens: list[int]):
|
||||
self.assertEqual(tok.encode(text), expected_tokens)
|
||||
self.assertEqual(tok.decode(expected_tokens), text)
|
||||
|
||||
# NOTE: the correct tokenization for this can only be found by looking up the text chunk in the vocab, not by applying merges
|
||||
def test_llama_early_tokenize(self): self._test_coding(self.llama_tok, " например", [ 111797 ])
|
||||
|
||||
def test_llama_basic(self): self._test_coding(self.llama_tok, "hello world", [ 15339, 1917 ])
|
||||
def test_llama_control_char(self): self._test_coding(self.llama_tok, " \x850", [ 220, 116360, 15 ])
|
||||
def test_llama_bytes(self): self._test_coding(self.llama_tok, " \xec\x8b\xa4\xed", [ 1717, 105, 116174, 82638, 2483 ])
|
||||
def test_llama_special1(self): self._test_coding(self.llama_tok, "hello <|end_of_text|>", [ 15339, 220, 128001 ])
|
||||
def test_llama_special2(self): self._test_coding(self.llama_tok, "<|start_header_id|>user<|end_header_id|>\n\n", [ 128006, 882, 128007, 271 ])
|
||||
def test_llama_repeat(self): self._test_coding(self.llama_tok, "00000000000000000", [ 931, 931, 931, 931, 931, 410 ])
|
||||
def test_llama_pat(self): self._test_coding(self.llama_tok, "today\n \n", [ 31213, 14211 ])
|
||||
|
||||
def test_split_regex_matches_naive_listing(self):
|
||||
# the compacted codepoint ranges must match the same text as listing every codepoint
|
||||
def naive(pre): return "".join(re.escape(chr(cp)) for cp in range(0x323b0) if unicodedata.category(chr(cp)).startswith(pre))
|
||||
r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + naive("Z"), naive("N"), naive("L")
|
||||
naive_re = re.compile("(?i:'s|'t|'re|'ve|'m|'ll|'d)|" +
|
||||
f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+")
|
||||
sample = "hello world 한국어 中文 текст ١٢٣ 123 😊\n \ttoday\n'équivalent ²³№ "
|
||||
self.assertEqual(SimpleTokenizer({}, {})._split_to_word.findall(sample), naive_re.findall(sample))
|
||||
|
||||
def test_split_regex_speed(self):
|
||||
# the naive listing compiles a 429KB pattern that takes 10+s to match a 225KB prompt; ranges keep it small and fast
|
||||
tok = SimpleTokenizer({}, {})
|
||||
self.assertLess(len(tok._split_to_word.pattern), 100_000)
|
||||
text = "The quick brown fox jumps over the lazy dog. " * 5000
|
||||
tok._split_to_word.findall(text) # warmup
|
||||
tms = []
|
||||
for _ in range(5):
|
||||
st = time.perf_counter()
|
||||
words = tok._split_to_word.findall(text)
|
||||
tms.append(time.perf_counter() - st)
|
||||
self.assertLess(min(tms), 4) # best-of-5 is robust to CI scheduling pauses; new code takes ~60ms
|
||||
self.assertEqual(len(words), 50001)
|
||||
|
||||
def test_llama_continued_conversation(self):
|
||||
self._test_coding(self.llama_tok, "hello <|eot_id|>world", [15339, 220, 128009, 14957])
|
||||
self._test_coding(self.llama_tok, "hello <|eot_id|>world again", [15339, 220, 128009, 14957, 1578])
|
||||
self._test_coding(self.llama_tok, "hello changed <|eot_id|>world again", [15339, 5614, 220, 128009, 14957, 1578])
|
||||
|
||||
def test_long_cached_prompt_matches_fresh_tokenization(self):
|
||||
prefix = "system tools\n" * 700 + "<|eot_id|>"
|
||||
first, changed = prefix + "run tower of hanoi", prefix + "run ls /"
|
||||
expected = self.llama_tok.encode(changed)
|
||||
self.llama_tok.encode(first)
|
||||
self.assertEqual(self.llama_tok.encode(changed), expected)
|
||||
|
||||
def test_tekken_from_gguf_kv(self):
|
||||
kv = {
|
||||
"tokenizer.ggml.tokens": ["<unk>", "<s>", "</s>", "[INST]", "[/INST]", "hello"],
|
||||
"tokenizer.ggml.token_type": [3, 3, 3, 3, 3, 1],
|
||||
"tokenizer.ggml.pre": "tekken",
|
||||
"tokenizer.ggml.eos_token_id": 2,
|
||||
}
|
||||
tok = SimpleTokenizer.from_gguf_kv(kv)
|
||||
template = FallbackTemplate(tok)
|
||||
self.assertEqual(template.role("user"), "[INST]")
|
||||
self.assertEqual(tok.encode("hello"), [5])
|
||||
self.assertEqual(template.end_turn(), "[/INST]")
|
||||
self.assertEqual(template.role("assistant"), "")
|
||||
|
||||
def test_stream_decoder(self):
|
||||
"""stream_decoder buffers incomplete UTF-8: token 25677 has 3/4 of emoji, token 138 completes it."""
|
||||
bs = [*range(33, 127), *range(161, 173), *range(174, 256)]
|
||||
be = {b: chr(b) for b in bs} | {b: chr(256+i) for i,b in enumerate(b for b in range(256) if b not in bs)}
|
||||
token_bytes = {25677: b'\x20\xf0\x9f\x98', 138: b'\x8a'} # ' ' + 3/4 emoji | 1/4 emoji (qwen3.5)
|
||||
tok = SimpleTokenizer({"".join(be[b] for b in v): k for k, v in token_bytes.items()}, {})
|
||||
dec = tok.stream_decoder()
|
||||
self.assertEqual(dec(25677) + dec(138) + dec(), " 😊")
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,251 @@
|
||||
import unittest
|
||||
from tinygrad import dtypes
|
||||
from tinygrad.uop.ops import UOp, Ops
|
||||
from tinygrad.schedule.memory import memory_plan_rewrite
|
||||
|
||||
global_map = {}
|
||||
held_bufs: set[UOp] = set()
|
||||
def b(i, base=None, offset=0, pin=False, size=16):
|
||||
global global_map
|
||||
if i in global_map: return global_map[i]
|
||||
if base is not None:
|
||||
global_map[i] = global_map[base]
|
||||
return global_map[i]
|
||||
global_map[i] = UOp.new_buffer("NULL", size, dtypes.int8)
|
||||
if pin: held_bufs.add(global_map[i])
|
||||
return global_map[i]
|
||||
|
||||
def _make_linear(buffer_lists, copies=None):
|
||||
copy_pairs = {frozenset((id(dst), id(src))) for dst, src in copies} if copies else set()
|
||||
calls = []
|
||||
for bufs in buffer_lists:
|
||||
is_copy = len(bufs) == 2 and frozenset((id(bufs[0]), id(bufs[1]))) in copy_pairs
|
||||
if is_copy:
|
||||
src0 = bufs[0].copy_to_device(bufs[1].device)
|
||||
else:
|
||||
src0 = UOp(Ops.SINK, src=tuple(bufs))
|
||||
calls.append(UOp(Ops.CALL, src=(src0, *bufs)))
|
||||
return UOp(Ops.LINEAR, src=tuple(calls))
|
||||
|
||||
def _get_arena(buf, linear, result):
|
||||
for orig_si, new_si in zip(linear.src, result.src):
|
||||
for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
|
||||
if orig is buf and new.op is Ops.SLICE: return new.src[0]
|
||||
return None
|
||||
|
||||
def check_assign(buffer_lists, copies=None):
|
||||
linear = _make_linear(buffer_lists, copies)
|
||||
result = memory_plan_rewrite(linear, held_bufs)
|
||||
|
||||
# build mapping: original buf -> (arena, offset_bytes, nbytes) from the result
|
||||
replace_map: dict[int, tuple[UOp, int, int]] = {}
|
||||
for orig_si, new_si in zip(linear.src, result.src):
|
||||
for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
|
||||
if new.op is Ops.SLICE and id(orig) not in replace_map:
|
||||
replace_map[id(orig)] = (new.src[0], new.src[1].val * new.src[0].dtype.itemsize, new.arg * new.dtype.itemsize)
|
||||
|
||||
# verify pinned buffers are not planned
|
||||
for buf in held_bufs:
|
||||
assert id(buf) not in replace_map, "pinned buffer was planned"
|
||||
|
||||
# compute lifetimes
|
||||
first_appearance, last_appearance = {}, {}
|
||||
for i, bufs in enumerate(buffer_lists):
|
||||
for buf in bufs:
|
||||
if buf in held_bufs: continue
|
||||
if id(buf) not in first_appearance: first_appearance[id(buf)] = i
|
||||
last_appearance[id(buf)] = i
|
||||
|
||||
# verify non-overlapping: no two live buffers share the same arena region
|
||||
taken_parts: set[tuple[int, int, int, int]] = set() # (id(arena), offset, nbytes, id(buf))
|
||||
for i, bufs in enumerate(buffer_lists):
|
||||
for buf in bufs:
|
||||
if buf in held_bufs or id(buf) not in replace_map: continue
|
||||
arena, off, nb = replace_map[id(buf)]
|
||||
for part in taken_parts:
|
||||
assert id(buf) == part[3] or part[0] != id(arena) or part[1] + part[2] <= off or part[1] >= off + nb, \
|
||||
f"overlap at step {i}: [{off}, {off+nb}) conflicts with [{part[1]}, {part[1]+part[2]})"
|
||||
if first_appearance.get(id(buf)) == i: taken_parts.add((id(arena), off, nb, id(buf)))
|
||||
if last_appearance.get(id(buf)) == i: taken_parts.discard((id(arena), off, nb, id(buf)))
|
||||
|
||||
class TestMemoryPlanner(unittest.TestCase):
|
||||
def setUp(self):
|
||||
global global_map
|
||||
held_bufs.clear()
|
||||
global_map = {}
|
||||
|
||||
def test_simple_buffer(self):
|
||||
bs = [
|
||||
[b(0), b(1), b(2)],
|
||||
[b(1), b(2), b(3)],
|
||||
[b(4), b(3)],
|
||||
[b(5), b(2)],
|
||||
]
|
||||
check_assign(bs)
|
||||
|
||||
def test_simple_pinned(self):
|
||||
bs = [
|
||||
[b(0, pin=True), b(1), b(2, pin=True)],
|
||||
[b(1), b(2), b(3)],
|
||||
[b(4), b(3)],
|
||||
[b(5), b(2)],
|
||||
]
|
||||
check_assign(bs)
|
||||
|
||||
def test_all_pinned(self):
|
||||
bs = [
|
||||
[b(0, pin=True), b(1, pin=True)],
|
||||
[b(1), b(2, pin=True)],
|
||||
[b(4, pin=True), b(3, pin=True)],
|
||||
]
|
||||
check_assign(bs)
|
||||
|
||||
def test_simple_buffer_offset(self):
|
||||
bs = [
|
||||
[b(0, pin=True), b(1, base=0, offset=1, size=8), b(2)],
|
||||
[b(1), b(2), b(3, base=0, offset=1, size=8)],
|
||||
[b(4), b(3)],
|
||||
]
|
||||
check_assign(bs)
|
||||
|
||||
def test_buffer_offset(self):
|
||||
bs = [
|
||||
[b(0, pin=True), b(1, base=0, offset=1, size=8), b(2)],
|
||||
[b(1), b(2), b(3, base=0, offset=1, size=8)],
|
||||
[b(4), b(3)],
|
||||
[b(5, base=2, offset=2, size=8), b(3)],
|
||||
[b(6), b(5), b(0)],
|
||||
[b(7), b(8, pin=True)],
|
||||
[b(8), b(9, base=2, offset=2, size=8)],
|
||||
[b(9), b(3), b(5)],
|
||||
]
|
||||
check_assign(bs)
|
||||
|
||||
def test_buffer_offset2(self):
|
||||
bs = [
|
||||
[b(0, pin=True), b(1), b(2)],
|
||||
[b(1), b(2), b(3)],
|
||||
[b(4), b(3)],
|
||||
[b(5), b(3)],
|
||||
[b(6), b(5), b(0)],
|
||||
[b(7), b(8, pin=True)],
|
||||
[b(8), b(9)],
|
||||
[b(9), b(3), b(5)],
|
||||
[b(11), b(0)],
|
||||
[b(11), b(10), b(5)],
|
||||
[b(12), b(11), b(0)],
|
||||
[b(6), b(12), b(7)],
|
||||
[b(13), b(6), b(11)],
|
||||
]
|
||||
check_assign(bs)
|
||||
|
||||
def test_all_offsets_of_one(self):
|
||||
bs = [
|
||||
[b(0, pin=True), b(1)],
|
||||
[b(3, base=1, offset=0, size=8), b(2, base=0, offset=0, size=8)],
|
||||
[b(5, base=1, offset=8, size=8), b(4, base=0, offset=8, size=8)],
|
||||
[b(7, base=1, offset=4, size=8), b(6, base=0, offset=4, size=8)],
|
||||
|
||||
[b(4), b(5), b(2)],
|
||||
[b(3), b(7)],
|
||||
[b(10), b(6), b(7)],
|
||||
[b(11), b(3), b(2)],
|
||||
[b(12), b(5), b(4), b(3), b(2)],
|
||||
[b(13), b(6), b(12), b(7)],
|
||||
]
|
||||
check_assign(bs)
|
||||
|
||||
def test_very_small_buffers(self):
|
||||
bs = [
|
||||
[b(0, pin=True), b(1, size=32)],
|
||||
[b(3, size=4), b(4, size=6)],
|
||||
]
|
||||
check_assign(bs)
|
||||
|
||||
def test_very_big_buffers(self):
|
||||
bs = [
|
||||
[b(0, pin=True), b(1, size=34359738368000)],
|
||||
[b(3, size=1 << 128), b(4, size=1 << 64)],
|
||||
]
|
||||
check_assign(bs)
|
||||
|
||||
def test_copy_bufs_separate_from_compute(self):
|
||||
bs = [
|
||||
[b(0), b(1)],
|
||||
[b(1), b(2)],
|
||||
[b(3), b(2)],
|
||||
]
|
||||
linear = _make_linear(bs, copies=[(b(1), b(0))])
|
||||
result = memory_plan_rewrite(linear)
|
||||
r1_arena, r2_arena = _get_arena(b(1), linear, result), _get_arena(b(2), linear, result)
|
||||
assert r1_arena is not None and r2_arena is not None
|
||||
assert r1_arena is not r2_arena
|
||||
|
||||
def test_copy_bufs_reuse_among_copies(self):
|
||||
bs = [
|
||||
[b(0), b(1)],
|
||||
[b(2), b(1)],
|
||||
[b(3), b(2)],
|
||||
]
|
||||
linear = _make_linear(bs, copies=[(b(1), b(0)), (b(2), b(1))])
|
||||
result = memory_plan_rewrite(linear)
|
||||
r1_arena, r2_arena = _get_arena(b(1), linear, result), _get_arena(b(2), linear, result)
|
||||
assert r1_arena is not None and r2_arena is not None
|
||||
assert r1_arena is r2_arena
|
||||
|
||||
def test_compute_bufs_reuse_among_compute(self):
|
||||
bs = [
|
||||
[b(0), b(1)],
|
||||
[b(2), b(1)],
|
||||
[b(3), b(2)],
|
||||
[b(4), b(3)],
|
||||
]
|
||||
linear = _make_linear(bs, copies=[(b(1), b(0))])
|
||||
result = memory_plan_rewrite(linear)
|
||||
r2_arena, r3_arena = _get_arena(b(2), linear, result), _get_arena(b(3), linear, result)
|
||||
assert r2_arena is not None and r3_arena is not None
|
||||
assert r2_arena is r3_arena
|
||||
|
||||
def test_copy_and_compute_no_cross_reuse(self):
|
||||
bs = [
|
||||
[b(0), b(1)],
|
||||
[b(2), b(1)],
|
||||
[b(3), b(2)],
|
||||
]
|
||||
linear = _make_linear(bs, copies=[(b(2), b(1))])
|
||||
result = memory_plan_rewrite(linear)
|
||||
r0_arena, r2_arena = _get_arena(b(0), linear, result), _get_arena(b(2), linear, result)
|
||||
assert r0_arena is not None and r2_arena is not None
|
||||
assert r0_arena is not r2_arena
|
||||
|
||||
def test_multiple_copy_bufs_with_offsets(self):
|
||||
bs = [
|
||||
[b(0, pin=True), b(1), b(2)],
|
||||
[b(3, base=0, offset=1, size=8), b(1), b(2)],
|
||||
[b(4), b(3)],
|
||||
[b(5), b(4)],
|
||||
]
|
||||
check_assign(bs, copies=[(b(1), b(0)), (b(2), b(0))])
|
||||
|
||||
def test_copy_bufs_pinned_mixed(self):
|
||||
bs = [
|
||||
[b(0, pin=True), b(1), b(2)],
|
||||
[b(1), b(3), b(2)],
|
||||
[b(4), b(3)],
|
||||
[b(5), b(4), b(0)],
|
||||
]
|
||||
check_assign(bs, copies=[(b(1), b(0)), (b(3), b(1))])
|
||||
|
||||
def test_deferred_copy_frees_chain(self):
|
||||
bs = []
|
||||
copies = []
|
||||
for i in range(6):
|
||||
copy_buf, compute_buf = b(i * 2 + 1), b(i * 2 + 2)
|
||||
bs.append([copy_buf, b(0, pin=True)])
|
||||
bs.append([compute_buf, copy_buf])
|
||||
copies.append((copy_buf, b(0, pin=True)))
|
||||
bs.append([b(100, pin=True)])
|
||||
check_assign(bs, copies=copies)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,48 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor, Device, Variable
|
||||
from examples.gpt2 import Transformer
|
||||
from tinygrad.nn.state import get_state_dict
|
||||
|
||||
class TestMethodCache(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.backup_compiler = Device[Device.DEFAULT].compiler.compile_cached
|
||||
def tearDown(self):
|
||||
Device[Device.DEFAULT].compiler.compile_cached = self.backup_compiler
|
||||
|
||||
def test_simple_methodcache(self):
|
||||
a = Tensor([1])
|
||||
b = Tensor([2])
|
||||
c = Tensor([3])
|
||||
d = Tensor([4])
|
||||
(a+b).realize()
|
||||
Device[Device.DEFAULT].compiler.compile_cached = None
|
||||
(c+d).realize()
|
||||
|
||||
def test_nested_methodcache(self):
|
||||
a,b,c,d = Tensor([1]), Tensor([2]), Tensor([3]), Tensor([4])
|
||||
((a+b)+(a+b)).realize()
|
||||
Device[Device.DEFAULT].compiler.compile_cached = None
|
||||
((c+d)+(c+d)).realize()
|
||||
|
||||
def test_nested_methodcache_swap(self):
|
||||
a,b,c,d = Tensor([1]), Tensor([2]), Tensor([3]), Tensor([4])
|
||||
((a+b)+(c+d)).realize()
|
||||
Device[Device.DEFAULT].compiler.compile_cached = None
|
||||
((c+d)+(a+b)).realize()
|
||||
|
||||
@unittest.skip("incorrect use of transformer")
|
||||
def test_small_transformer(self):
|
||||
args_tiny = {"dim": 16, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-05, "vocab_size": 10}
|
||||
model = Transformer(**args_tiny)
|
||||
for v in get_state_dict(model).values(): v.assign(Tensor.empty(*v.shape, dtype=v.dtype).realize())
|
||||
# NOTE: you have to do this twice due to the k-v cache
|
||||
for i in range(3): model(Tensor([[1,2,3,4]]), Variable("start_pos", 0, 10).bind(i)).realize()
|
||||
for i in range(3): model(Tensor([[1,2,3,4]]), Variable("start_pos", 0, 10).bind(i)).realize()
|
||||
Device[Device.DEFAULT].compiler.compile_cached = None
|
||||
for i in range(3): model(Tensor([[1,2,3,4]]), Variable("start_pos", 0, 10).bind(i)).realize()
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,100 @@
|
||||
import unittest, time
|
||||
from tinygrad import Tensor, UOp, getenv
|
||||
from tinygrad.helpers import Profiling
|
||||
|
||||
PYPROFILE = getenv("PYPROFILE")
|
||||
class TestBench(unittest.TestCase):
|
||||
@staticmethod
|
||||
def setUpClass():
|
||||
# no fixed cost
|
||||
Tensor.empty(10,10)
|
||||
Tensor.randn(10,10)
|
||||
|
||||
def start_time(self): self.st = time.perf_counter()
|
||||
def setUp(self):
|
||||
# it's about 1 ms per 1k UOps on M3
|
||||
if PYPROFILE:
|
||||
self.prof = Profiling()
|
||||
self.prof.__enter__()
|
||||
else:
|
||||
self.prof = None
|
||||
self.N = 10000
|
||||
self.start_time()
|
||||
|
||||
def tearDown(self):
|
||||
et = (time.perf_counter() - self.st)
|
||||
if self.prof is not None: self.prof.__exit__()
|
||||
print(f"{self._testMethodName:30s} {et*1e6/self.N:.2f} us")
|
||||
|
||||
def test_uop_instant_creation(self):
|
||||
for i in range(self.N): UOp.const(100+i)
|
||||
|
||||
def test_uop_list_creation(self):
|
||||
[UOp.const(100+i) for i in range(self.N)]
|
||||
|
||||
def test_uop_add_2n(self):
|
||||
a = UOp.const(2)
|
||||
for _ in range(self.N): a = a + a
|
||||
|
||||
def test_uop_toposort(self):
|
||||
a = UOp.const(0)
|
||||
for i in range(self.N): a = a + UOp.const(100+i)
|
||||
self.start_time()
|
||||
self.assertEqual(len(a.toposort()), 2*self.N+1)
|
||||
|
||||
def test_uop_toposort_2n(self):
|
||||
a = UOp.const(0)
|
||||
for _ in range(self.N): a = a + a
|
||||
self.start_time()
|
||||
self.assertEqual(len(a.toposort()), self.N+1)
|
||||
|
||||
def test_uop_simplify(self):
|
||||
a = UOp.const(2)
|
||||
for _ in range(self.N): (a+a).simplify()
|
||||
|
||||
def test_uop_simplify_complex(self):
|
||||
self.N //= 10 # this test is slow
|
||||
x = UOp.variable("x", 0, 10)
|
||||
y = UOp.variable("y", 0, 10)
|
||||
expr = (x*2)+5+(x*4)+(y*2)+y
|
||||
for _ in range(self.N): expr.simplify()
|
||||
|
||||
def test_uop_simplify_div(self):
|
||||
self.N //= 10 # this test is slow
|
||||
x = UOp.variable("x", 0, 10)
|
||||
y = UOp.variable("y", 0, 10)
|
||||
z = UOp.variable("z", 0, 10)
|
||||
expr = (x*4+y*8)//(z*2)
|
||||
for _ in range(self.N): expr.simplify()
|
||||
|
||||
def test_uop_chain_free(self):
|
||||
a = UOp.const(2)
|
||||
for _ in range(self.N): a = a + a
|
||||
self.start_time()
|
||||
del a
|
||||
|
||||
def test_tensor_zeros(self):
|
||||
self.N //= 10 # this test is slow
|
||||
for _ in range(self.N): Tensor.zeros(10, 10)
|
||||
|
||||
def test_tensor_add(self):
|
||||
self.N //= 10 # this test is slow
|
||||
a = Tensor.zeros(10, 10)
|
||||
b = Tensor.zeros(10, 10)
|
||||
for _ in range(self.N): a+b
|
||||
|
||||
def test_tensor_empty(self):
|
||||
self.N //= 10 # this test is slow
|
||||
for _ in range(self.N): Tensor.empty(10, 10)
|
||||
|
||||
def test_tensor_rand(self):
|
||||
self.N //= 100 # this test is very slow
|
||||
for _ in range(self.N): Tensor.rand(10, 10)
|
||||
|
||||
def test_tensor_randn(self):
|
||||
self.N //= 100 # this test is very slow
|
||||
for _ in range(self.N): Tensor.randn(10, 10)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,14 @@
|
||||
import unittest
|
||||
from tinygrad.helpers import GlobalCounters
|
||||
from tinygrad.nn.datasets import mnist
|
||||
|
||||
class TestDataset(unittest.TestCase):
|
||||
def test_dataset_is_realized(self):
|
||||
X_train, _, _, _ = mnist()
|
||||
X_train[0].contiguous().realize()
|
||||
GlobalCounters.reset()
|
||||
X_train[0].contiguous().realize()
|
||||
self.assertLessEqual(GlobalCounters.kernel_count, 1) # 0 if SLICE (zero-copy), 1 otherwise
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
227
artifacts/package_sources/tinygrad/test/null/test_multitensor.py
Normal file
227
artifacts/package_sources/tinygrad/test/null/test_multitensor.py
Normal file
@@ -0,0 +1,227 @@
|
||||
import gc, unittest
|
||||
from tinygrad import Tensor, UOp, GlobalCounters, dtypes
|
||||
from tinygrad.engine.jit import TinyJit
|
||||
from tinygrad.helpers import Context
|
||||
|
||||
class TestMultiRamUsage(unittest.TestCase):
|
||||
def setUp(self):
|
||||
gc.collect()
|
||||
self.baseline = GlobalCounters.mem_used
|
||||
self.baseline_per_device = dict(GlobalCounters.mem_used_per_device)
|
||||
self.N = 100
|
||||
def assertUsed(self, amt, strict=True):
|
||||
gc.collect()
|
||||
used = GlobalCounters.mem_used - self.baseline
|
||||
print(f"used {used} bytes")
|
||||
if strict: self.assertEqual(used, amt)
|
||||
else: self.assertLessEqual(used, amt)
|
||||
def assertDeviceUsed(self, expected:dict[str, int]):
|
||||
gc.collect()
|
||||
for dev, amt in expected.items():
|
||||
used = GlobalCounters.mem_used_per_device[dev] - self.baseline_per_device.get(dev, 0)
|
||||
self.assertEqual(used, amt, f"device {dev}: expected {amt} bytes used, got {used}")
|
||||
|
||||
def test_zeros(self):
|
||||
_ = Tensor.zeros(self.N, self.N).contiguous().realize()
|
||||
self.assertUsed(self.N*self.N*4)
|
||||
|
||||
def test_zeros_del(self):
|
||||
_ = Tensor.zeros(self.N, self.N).contiguous().realize()
|
||||
del _
|
||||
self.assertUsed(0)
|
||||
|
||||
def test_zeros_copy(self):
|
||||
devices_2 = ("NULL:1", "NULL:2")
|
||||
_ = Tensor.zeros(self.N, self.N).contiguous().to(devices_2).realize()
|
||||
# NOTE: the first one on the DEFAULT device should be freed
|
||||
self.assertUsed(self.N*self.N*4*2)
|
||||
|
||||
def test_zeros_shard(self, devices=("NULL:1", "NULL:2")):
|
||||
_ = Tensor.zeros(self.N, self.N).contiguous().shard(devices, axis=0).realize()
|
||||
self.assertUsed(self.N*self.N*4) # sharding should not increase total ram usage
|
||||
def test_zeros_shard_self(self): self.test_zeros_shard(("NULL:0", "NULL:1"))
|
||||
|
||||
def test_zeros_contiguous_shard(self):
|
||||
devices_2 = ("NULL:1", "NULL:2")
|
||||
_ = Tensor.zeros(self.N, self.N).contiguous().shard(devices_2, axis=0).contiguous().realize()
|
||||
self.assertUsed(self.N*self.N*4) # sharding should not increase total ram usage
|
||||
|
||||
def test_sharded_memory_replicated(self):
|
||||
devices_4 = tuple(f"NULL:{i+1}" for i in range(4))
|
||||
X = Tensor.ones(256).contiguous().realize()
|
||||
self.assertUsed(256 * 4)
|
||||
X.shard_(devices_4).realize()
|
||||
self.assertUsed(256 * 4 * 4)
|
||||
|
||||
def test_sharded_memory_replicated_const(self):
|
||||
devices_4 = tuple(f"NULL:{i+1}" for i in range(4))
|
||||
X = Tensor.ones(256, buffer=False).realize()
|
||||
self.assertUsed(0)
|
||||
X.shard_(devices_4).realize()
|
||||
self.assertUsed(0)
|
||||
|
||||
def test_sharded_memory_axis_const(self):
|
||||
devices_4 = tuple(f"NULL:{i+1}" for i in range(4))
|
||||
X = Tensor.ones(256, buffer=False).realize()
|
||||
self.assertUsed(0)
|
||||
X.shard_(devices_4, axis=0).realize()
|
||||
self.assertUsed(0)
|
||||
|
||||
def test_zeros_per_device(self):
|
||||
_ = Tensor.zeros(self.N, self.N, device="NULL").contiguous().realize()
|
||||
self.assertDeviceUsed({"NULL": self.N*self.N*4})
|
||||
|
||||
def test_zeros_del_per_device(self):
|
||||
_ = Tensor.zeros(self.N, self.N, device="NULL").contiguous().realize()
|
||||
del _
|
||||
self.assertDeviceUsed({"NULL": 0})
|
||||
|
||||
def test_zeros_copy_per_device(self):
|
||||
devices_2 = ("NULL:1", "NULL:2")
|
||||
_ = Tensor.zeros(self.N, self.N).contiguous().to(devices_2).realize()
|
||||
self.assertDeviceUsed({"NULL:1": self.N*self.N*4, "NULL:2": self.N*self.N*4})
|
||||
|
||||
def test_zeros_shard_per_device(self):
|
||||
devices_2 = ("NULL:1", "NULL:2")
|
||||
_ = Tensor.zeros(self.N, self.N).contiguous().shard(devices_2, axis=0).realize()
|
||||
self.assertDeviceUsed({"NULL:1": self.N*(self.N//2)*4, "NULL:2": self.N*(self.N//2)*4})
|
||||
|
||||
def test_sharded_memory_replicated_per_device(self):
|
||||
devices_4 = tuple(f"NULL:{i+1}" for i in range(4))
|
||||
X = Tensor.ones(256, device="NULL").contiguous().realize()
|
||||
self.assertDeviceUsed({"NULL": 256*4})
|
||||
X.shard_(devices_4).realize()
|
||||
for d in devices_4:
|
||||
self.assertDeviceUsed({d: 256*4})
|
||||
|
||||
def _test_matmul_half(self, dev_count:int):
|
||||
N = 32
|
||||
total_mem = {}
|
||||
devs = tuple(f"NULL:{i}" for i in range(dev_count))
|
||||
for dtype in {dtypes.float, dtypes.half}:
|
||||
GlobalCounters.reset()
|
||||
a = Tensor.empty((N, N), dtype=dtype, device=devs[0]).shard(devs, axis=0)
|
||||
b = Tensor.empty((N, N), dtype=dtype, device=devs[0]).shard(devs, axis=None)
|
||||
(a @ b).realize()
|
||||
total_mem[dtype] = GlobalCounters.global_mem
|
||||
self.assertEqual(total_mem[dtypes.half], total_mem[dtypes.float] // 2)
|
||||
|
||||
def test_matmul_half(self): self._test_matmul_half(dev_count=2)
|
||||
def test_matmul_half_alt(self): self._test_matmul_half(dev_count=4)
|
||||
|
||||
def test_multi_layer_allreduce(self):
|
||||
N = 32
|
||||
devices_2 = ("NULL:1", "NULL:2")
|
||||
|
||||
def make_inp():
|
||||
x = Tensor.zeros(N, N).contiguous().shard(devices_2, axis=None).realize()
|
||||
w1 = Tensor.zeros(N, N).contiguous().shard(devices_2, axis=1).realize()
|
||||
w2 = Tensor.zeros(N, N).contiguous().shard(devices_2, axis=0).realize()
|
||||
return x, w1, w2
|
||||
|
||||
def run_layers(n_layers):
|
||||
GlobalCounters.reset()
|
||||
|
||||
@TinyJit
|
||||
def f(x, w1, w2):
|
||||
for _ in range(n_layers):
|
||||
x = (x @ w1 @ w2)
|
||||
return x.contiguous()
|
||||
|
||||
for _ in range(3):
|
||||
a = make_inp()
|
||||
r = f(*a)
|
||||
del a, r
|
||||
|
||||
gc.collect()
|
||||
return GlobalCounters.mem_used
|
||||
|
||||
mem_2 = run_layers(2)
|
||||
mem_4 = run_layers(4)
|
||||
self.assertEqual(mem_2, mem_4, f"graph memory should not grow with layers: 2 layers={mem_2}, 4 layers={mem_4}")
|
||||
|
||||
def test_allreduce_cast_dtype_memory(self):
|
||||
N = 32
|
||||
devices_2 = ("NULL:1", "NULL:2")
|
||||
mem = {}
|
||||
for allreduce_cast in (0, 1):
|
||||
GlobalCounters.reset()
|
||||
with Context(ALLREDUCE_CAST=allreduce_cast, SCACHE=0):
|
||||
x = Tensor.empty((N, N), dtype=dtypes.bfloat16, device="NULL:1").shard(devices_2, axis=0)
|
||||
x.sum(0).realize()
|
||||
mem[allreduce_cast] = GlobalCounters.global_mem
|
||||
# with ALLREDUCE_CAST, allreduce copies happen in bf16 (2 bytes) instead of fp32 (4 bytes)
|
||||
self.assertLess(mem[1], mem[0])
|
||||
|
||||
class TestMultiScalarALU(unittest.TestCase):
|
||||
"""Test that tuple-device scalars work correctly in ALU with MULTI tensors (_shard scalar fix)."""
|
||||
def test_multi_times_replicated_scalar(self):
|
||||
devices = ("NULL:0", "NULL:1")
|
||||
x = Tensor.ones(4).contiguous().shard(devices, axis=0)
|
||||
s = Tensor(2.0).to(devices)
|
||||
result = x * s
|
||||
self.assertEqual(result.shape, (4,))
|
||||
self.assertEqual(result.uop.axis, 0)
|
||||
|
||||
def test_multi_add_replicated_scalar(self):
|
||||
devices = ("NULL:0", "NULL:1")
|
||||
x = Tensor.ones(4).contiguous().shard(devices, axis=0)
|
||||
s = Tensor(1.0).to(devices)
|
||||
result = x + s
|
||||
self.assertEqual(result.shape, (4,))
|
||||
self.assertEqual(result.uop.axis, 0)
|
||||
|
||||
def test_multi_times_call_scalar(self):
|
||||
"""Per-device scalar from a CALL (like FP8 local amax) used in ALU with MULTI."""
|
||||
import functools
|
||||
from tinygrad.uop.ops import Ops
|
||||
devices = ("NULL:0", "NULL:1")
|
||||
x = Tensor.ones(4, 4).contiguous().shard(devices, axis=0)
|
||||
# simulate per-device scalar via CALL (strips MULTI from param body → no allreduce)
|
||||
@functools.cache
|
||||
def _fxn(x_p, device):
|
||||
t = Tensor(x_p, device=device)
|
||||
inner = Tensor(t.uop.src[0]) if t.uop.op is Ops.UNSHARD else t
|
||||
return (inner.sum(),)
|
||||
param = x.as_param(0)
|
||||
fxn = _fxn(param.uop, x.device)
|
||||
per_dev_scalar = Tensor(fxn[0].uop.call(x.uop).gettuple(0))
|
||||
result = x * per_dev_scalar
|
||||
self.assertEqual(result.shape, (4, 4))
|
||||
self.assertEqual(result.uop.axis, 0)
|
||||
result.realize()
|
||||
|
||||
class TestMultiAxis(unittest.TestCase):
|
||||
def test_reshape_shard_invalid(self):
|
||||
devices = ("NULL:0", "NULL:1")
|
||||
t = Tensor.ones(4, 3).shard(devices, axis=0)
|
||||
with self.assertRaises(RuntimeError, msg="reshape cannot move items between shards"):
|
||||
t.reshape(3, 4).uop.axis
|
||||
|
||||
def test_reshape_shard_valid(self):
|
||||
devices = ("NULL:0", "NULL:1")
|
||||
t = Tensor.ones(4, 8).shard(devices, axis=0)
|
||||
self.assertEqual(t.reshape(2, 16).uop.axis, 0)
|
||||
self.assertEqual(t.reshape(2, 2, 8).uop.axis, 0)
|
||||
|
||||
def test_uop_shard_axis_none(self):
|
||||
devices = ("NULL:0", "NULL:1")
|
||||
u = Tensor.ones(8).contiguous().realize().uop
|
||||
self.assertIsNone(u.shard(devices).axis)
|
||||
self.assertEqual(u.shard(devices, 0).axis, 0)
|
||||
|
||||
def test_empty_like_sharded(self):
|
||||
t = Tensor.ones(4, 8).shard(("NULL:0", "NULL:1"), axis=0)
|
||||
e = t.empty_like()
|
||||
self.assertEqual(e.shape, t.shape)
|
||||
self.assertEqual(e.device, t.device)
|
||||
self.assertEqual(e.uop.axis, 0)
|
||||
self.assertTrue(e.uop.has_buffer_identity())
|
||||
|
||||
def test_symbolic_reshape_shard_axis(self):
|
||||
rows = UOp.variable("rows", 1, 4).bind(3)
|
||||
x = Tensor.empty(4, 2).shard(("NULL:1", "NULL:2"), axis=1)[:rows]
|
||||
self.assertEqual(x.reshape(rows, 1, 2).uop.axis, 2)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,216 @@
|
||||
import unittest, itertools
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.uop.ops import Ops, UOp, GroupOp # noqa: F401
|
||||
from tinygrad.uop.ops import PatternMatcher, UPat
|
||||
|
||||
class TestPatternMatcher(unittest.TestCase):
|
||||
def test_simple_match(self):
|
||||
matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.weakfloat), lambda x: x.rtag())])
|
||||
c1 = UOp.const(1.0)
|
||||
c2 = UOp.const(1)
|
||||
self.assertEqual(matcher.rewrite(c1), c1.rtag())
|
||||
self.assertEqual(matcher.rewrite(c2), None)
|
||||
|
||||
def test_upat_any(self):
|
||||
def test(a, x=None, y=None, z=None):
|
||||
#print(x,y,z)
|
||||
if y is not None: return (a+y).rtag()
|
||||
matcher = PatternMatcher([
|
||||
(UPat.var("a")+UPat.any(UPat.var("x"), UPat.var("y"), UPat.var("z")), test),
|
||||
])
|
||||
v1 = UOp.variable("a", 0, 10)
|
||||
v2 = UOp.variable("b", 0, 10)
|
||||
c1 = v1+v2
|
||||
self.assertEqual(matcher.rewrite(c1), c1.rtag())
|
||||
|
||||
def test_minimum_len(self):
|
||||
matcher = PatternMatcher([
|
||||
(UPat(Ops.NOOP, src=(UPat(Ops.NOOP), UPat(Ops.NOOP)), allow_any_len=True), lambda: True),
|
||||
])
|
||||
self.assertTrue(matcher.rewrite(UOp(Ops.NOOP, src=(UOp(Ops.NOOP), UOp(Ops.NOOP), UOp(Ops.NOOP),))))
|
||||
self.assertTrue(matcher.rewrite(UOp(Ops.NOOP, src=(UOp(Ops.NOOP), UOp(Ops.NOOP)))))
|
||||
self.assertIsNone(matcher.rewrite(UOp(Ops.NOOP, src=(UOp(Ops.NOOP),))))
|
||||
|
||||
@unittest.skip("closures aren't supported on pattern matchers")
|
||||
def test_match_sz_0(self):
|
||||
match_cnt = 0
|
||||
def fxn(x):
|
||||
nonlocal match_cnt
|
||||
match_cnt += 1
|
||||
assert len(x.src) == 0
|
||||
return UOp(Ops.CONST, src=(UOp(Ops.CONST),))
|
||||
matcher = PatternMatcher([(UPat(Ops.CONST, src=(), name="x"), fxn)])
|
||||
c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
|
||||
# second rewrite shouldn't match anything
|
||||
c1 = matcher.rewrite(c1)
|
||||
c1 = matcher.rewrite(c1)
|
||||
self.assertEqual(match_cnt, 1)
|
||||
|
||||
def test_match_sz_0_ctx(self):
|
||||
def fxn(ctx, x):
|
||||
ctx.append(True)
|
||||
assert len(x.src) == 0
|
||||
return x.replace(src=(UOp(Ops.NOOP),))
|
||||
matcher = PatternMatcher([(UPat(Ops.NOOP, src=(), name="x"), fxn)])
|
||||
c1 = UOp(Ops.NOOP)
|
||||
# second rewrite shouldn't match anything
|
||||
ctx = []
|
||||
c1 = matcher.rewrite(c1, ctx)
|
||||
c1 = matcher.rewrite(c1, ctx)
|
||||
self.assertEqual(len(ctx), 1)
|
||||
|
||||
def test_uop(self):
|
||||
matcher = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.rtag())])
|
||||
c1 = UOp.const(1.0)
|
||||
c2 = UOp(Ops.ADD, src=(c1, c1))
|
||||
self.assertEqual(matcher.rewrite(c1), c1.rtag())
|
||||
self.assertEqual(matcher.rewrite(c2), None)
|
||||
|
||||
def test_uop_set(self):
|
||||
matcher = PatternMatcher([(UPat((Ops.CONST, Ops.CAST), name="x"), lambda x: x.rtag())])
|
||||
c1 = UOp.const(False)
|
||||
c2 = UOp(Ops.CAST, arg=dtypes.int, src=(c1,))
|
||||
c3 = UOp.const(1.0)
|
||||
c4 = UOp(Ops.ADD, src=(c3, c3))
|
||||
self.assertEqual(matcher.rewrite(c1), c1.rtag())
|
||||
self.assertEqual(matcher.rewrite(c2), c2.rtag())
|
||||
self.assertEqual(matcher.rewrite(c4), None)
|
||||
|
||||
def test_arg(self):
|
||||
matcher = PatternMatcher([
|
||||
(UPat(Ops.CONST, arg=0, name="x"), lambda x: x.rtag()),
|
||||
(UPat(Ops.CONST, arg=False, name="x"), lambda x: x.rtag()),
|
||||
(UPat(Ops.MAX, name="x"), lambda x: x.rtag()),
|
||||
])
|
||||
c1 = UOp.const(0.0)
|
||||
c2 = UOp.const(False)
|
||||
c3 = UOp(Ops.MAX, src=(c1, c1))
|
||||
c4 = UOp(Ops.MUL, src=(c1, c1))
|
||||
c5 = UOp.const(-1)
|
||||
self.assertEqual(matcher.rewrite(c1), c1.rtag())
|
||||
self.assertEqual(matcher.rewrite(c2), c2.rtag())
|
||||
self.assertEqual(matcher.rewrite(c3), c3.rtag())
|
||||
self.assertEqual(matcher.rewrite(c4), None)
|
||||
self.assertEqual(matcher.rewrite(c5), None)
|
||||
|
||||
def test_filter_arg(self):
|
||||
matcher = PatternMatcher([
|
||||
(UPat(Ops.MUL, src=[UPat(Ops.CONST, name="c"), UPat(Ops.CONST, arg=2)], name="x"),
|
||||
lambda x,c: x.rtag() if c.val in {1, -1} else None)
|
||||
])
|
||||
y1 = UOp.const(1)
|
||||
y2 = UOp.const(2)
|
||||
y3 = UOp.const(-1)
|
||||
c1 = UOp(Ops.MUL, src=(y1, y2))
|
||||
c2 = UOp(Ops.MUL, src=(y2, y2))
|
||||
c3 = UOp(Ops.MUL, src=(y3, y2))
|
||||
c4 = UOp(Ops.MUL, src=(y2, y1))
|
||||
c5 = UOp(Ops.MUL, src=(y2, y3))
|
||||
self.assertEqual(matcher.rewrite(c1), c1.rtag())
|
||||
self.assertEqual(matcher.rewrite(c2), None)
|
||||
self.assertEqual(matcher.rewrite(c3), c3.rtag())
|
||||
self.assertEqual(matcher.rewrite(c4), c4.rtag())
|
||||
self.assertEqual(matcher.rewrite(c5), c5.rtag())
|
||||
|
||||
def test_dup_name(self):
|
||||
matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST, name="y"), UPat(Ops.CONST, name="y"))), lambda x, y: x.rtag())])
|
||||
y1 = UOp.const(1.0)
|
||||
y2 = UOp.const(1.0)
|
||||
c1 = UOp(Ops.ADD, src=(y1, y1))
|
||||
c2 = UOp(Ops.ADD, src=(y1, y2))
|
||||
self.assertEqual(matcher.rewrite(c1), c1.rtag())
|
||||
self.assertEqual(matcher.rewrite(c2), c1.rtag())
|
||||
|
||||
def test_dtype(self):
|
||||
# a concrete const dtype lives on the pair's CAST
|
||||
matcher = PatternMatcher([(UPat(Ops.CAST, name="x", dtype=dtypes.float32), lambda x: x.rtag())])
|
||||
c1 = UOp.const(1.0).cast(dtypes.float32)
|
||||
c2 = UOp.const(1.0).cast(dtypes.float64)
|
||||
self.assertEqual(matcher.rewrite(c1), c1.rtag())
|
||||
self.assertEqual(matcher.rewrite(c2), None)
|
||||
|
||||
def test_dtype_set(self):
|
||||
matcher = PatternMatcher([(UPat(Ops.CAST, name="x", dtype={dtypes.float32, dtypes.float64}), lambda x: x.rtag())])
|
||||
c1 = UOp.const(1.0).cast(dtypes.float32)
|
||||
c2 = UOp.const(1.0).cast(dtypes.float64)
|
||||
c3 = UOp.const(1.0).cast(dtypes.float16)
|
||||
c4 = UOp.const(1).cast(dtypes.int)
|
||||
self.assertEqual(matcher.rewrite(c1), c1.rtag())
|
||||
self.assertEqual(matcher.rewrite(c2), c2.rtag())
|
||||
self.assertEqual(matcher.rewrite(c3), None)
|
||||
self.assertEqual(matcher.rewrite(c4), None)
|
||||
|
||||
def test_src_one(self):
|
||||
matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST), UPat(Ops.CONST))), lambda x: x.rtag())])
|
||||
c1 = UOp.const(1.0)
|
||||
c2 = UOp.const(2.0)
|
||||
c3 = UOp(Ops.ADD, src=(c1,c2))
|
||||
self.assertEqual(matcher.rewrite(c3), c3.rtag())
|
||||
self.assertEqual(matcher.rewrite(c2), None)
|
||||
# that CONST/ALU -> ALU/CONST rewrite is now instant
|
||||
"""
|
||||
matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST), UPat(GroupOp.ALU))), lambda x: x)])
|
||||
c4 = UOp(Ops.ADD, src=(c1,c3))
|
||||
c5 = UOp(Ops.ADD, src=(c3,c1))
|
||||
self.assertEqual(matcher.rewrite(c3), None)
|
||||
self.assertEqual(matcher.rewrite(c4), c4)
|
||||
self.assertEqual(matcher.rewrite(c5), None)
|
||||
"""
|
||||
|
||||
def test_src_permutations(self):
|
||||
matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=[UPat(Ops.CONST), UPat(GroupOp.ALU)]), lambda x: x.rtag())])
|
||||
c1 = UOp.const(1.0)
|
||||
c2 = UOp.const(2.0)
|
||||
c3 = UOp(Ops.ADD, src=(c1,c2))
|
||||
c4 = UOp(Ops.ADD, src=(c3,c2))
|
||||
c5 = UOp(Ops.ADD, src=(c2,c3))
|
||||
c6 = UOp(Ops.ADD, src=(c3,c4))
|
||||
self.assertEqual(matcher.rewrite(c3), None)
|
||||
self.assertEqual(matcher.rewrite(c4), c4.rtag())
|
||||
self.assertEqual(matcher.rewrite(c5), c5.rtag())
|
||||
self.assertEqual(matcher.rewrite(c6), None)
|
||||
|
||||
def test_src_repeat(self):
|
||||
matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=UPat(Ops.CONST)), lambda x: x.rtag())])
|
||||
c1 = UOp.const(1.0)
|
||||
c2 = UOp.const(2.0)
|
||||
c3 = UOp(Ops.ADD, src=(c1,c2))
|
||||
c4 = UOp(Ops.ADD, src=(c2,c3))
|
||||
self.assertEqual(matcher.rewrite(c3), c3.rtag())
|
||||
self.assertEqual(matcher.rewrite(c4), None)
|
||||
|
||||
def test_allow_len(self):
|
||||
matcher = PatternMatcher([(UPat(Ops.MULACC, name="x", src=(UPat(Ops.CONST),), allow_any_len=True), lambda x: x.rtag())])
|
||||
c1 = UOp.const(1.0)
|
||||
c2 = UOp.const(2.0)
|
||||
c3 = UOp.const(3.0)
|
||||
c4 = UOp(Ops.EXP2, src=(c1,))
|
||||
c5 = UOp(Ops.ADD, src=(c1,c2))
|
||||
c6 = UOp(Ops.MULACC, src=(c1,c2,c3))
|
||||
self.assertEqual(matcher.rewrite(c4), None)
|
||||
self.assertEqual(matcher.rewrite(c5), None)
|
||||
self.assertEqual(matcher.rewrite(c6), c6.rtag())
|
||||
|
||||
def test_deep_src_permutations(self):
|
||||
c1 = UOp.const(1.0)
|
||||
c2 = UOp.const(2.0)
|
||||
u1 = (c1 + c2) + c1
|
||||
u2 = (c2 + c1) + c1
|
||||
matcher = PatternMatcher([
|
||||
(UPat(GroupOp.ALU, src=[UPat(GroupOp.ALU, src=[UPat(name='a'), UPat(name='b')]), UPat(name='b')]), lambda a,b: b)
|
||||
])
|
||||
self.assertIsNotNone(matcher.rewrite(u1))
|
||||
self.assertIsNotNone(matcher.rewrite(u2))
|
||||
|
||||
def _assert_eq_upat(self, a:UPat, b:UPat):
|
||||
assert (sorted(map(str,a.op)) if a.op else [] == (sorted(map(str,b.op)) if b.op else []))
|
||||
assert (sorted(a.match_dtype) if a.match_dtype else [] == (sorted(b.match_dtype) if b.match_dtype else []))
|
||||
assert (a.name, type(a.src)) == (b.name, type(b.src))
|
||||
def simple_src(u:UPat):
|
||||
if u.src is None: return []
|
||||
if isinstance(u.src, itertools.repeat): return next(u.src[0])
|
||||
return u.src[0]
|
||||
for a,b in zip(simple_src(a), simple_src(b)): self._assert_eq_upat(a, b)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main(verbosity=2)
|
||||
@@ -0,0 +1,44 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor, Device, Context
|
||||
from tinygrad.codegen import do_to_program
|
||||
from tinygrad.codegen.opt import Opt, OptOps
|
||||
from test.external.process_replay.process_replay import replay_to_program
|
||||
from test.helpers import replace_opts
|
||||
|
||||
N = 16
|
||||
class TestProcessReplay(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.ast = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule_linear().src[-1].src[0]
|
||||
cls.renderer = Device[Device.DEFAULT].renderer
|
||||
|
||||
def test_replay_no_opts(self):
|
||||
# opts=None means use default heuristic path
|
||||
p = do_to_program(self.ast, self.renderer)
|
||||
good, compare, _ = replay_to_program(p, self.ast, self.renderer)
|
||||
self.assertEqual(good, compare)
|
||||
|
||||
def test_replay_empty_opts(self):
|
||||
# opts=[] means explicitly apply zero opts (unoptimized)
|
||||
ast = replace_opts(self.ast, [])
|
||||
p = do_to_program(ast, self.renderer)
|
||||
good, compare, _ = replay_to_program(p, ast, self.renderer)
|
||||
self.assertEqual(good, compare)
|
||||
|
||||
def test_replay_with_opt(self):
|
||||
# opts=[Opt(...)] means apply a specific opt
|
||||
opts = [Opt(OptOps.UPCAST, 0, 4)]
|
||||
ast = replace_opts(self.ast, opts)
|
||||
p = do_to_program(ast, self.renderer)
|
||||
good, compare, _ = replay_to_program(p, ast, self.renderer)
|
||||
self.assertEqual(good, compare)
|
||||
|
||||
def test_beam(self):
|
||||
with Context(BEAM=1):
|
||||
ast = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule_linear().src[-1].src[0]
|
||||
p = do_to_program(ast, self.renderer)
|
||||
good, compare, _ = replay_to_program(p, ast, self.renderer)
|
||||
self.assertEqual(good, compare)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main(verbosity=2)
|
||||
186
artifacts/package_sources/tinygrad/test/null/test_real_world.py
Normal file
186
artifacts/package_sources/tinygrad/test/null/test_real_world.py
Normal file
@@ -0,0 +1,186 @@
|
||||
import unittest, time, gc
|
||||
import numpy as np
|
||||
from tinygrad.nn import optim
|
||||
from tinygrad.nn.state import get_parameters
|
||||
from tinygrad.engine.jit import TinyJit
|
||||
from tinygrad import Tensor, Device, GlobalCounters, dtypes, Variable
|
||||
from tinygrad.helpers import Context
|
||||
from test.helpers import slow, jit_cache_count, KernelCountException
|
||||
from extra.lr_scheduler import OneCycleLR
|
||||
from test.helpers import derandomize_model
|
||||
|
||||
from examples.gpt2 import Transformer as GPT2Transformer
|
||||
from examples.hlb_cifar10 import SpeedyResNet, hyp
|
||||
from extra.models.llama import Transformer as LLaMaTransformer
|
||||
from examples.stable_diffusion import UNetModel, unet_params
|
||||
from extra.models.unet import ResBlock
|
||||
from extra.models.bert import BertForPretraining
|
||||
|
||||
global_mem_used = 0
|
||||
def helper_test(nm, gen, model, max_memory_allowed, max_kernels_allowed, all_jitted=False):
|
||||
with Context(JIT=2):
|
||||
tms = []
|
||||
for _ in range(2):
|
||||
early_gen = [x.realize() if isinstance(x, Tensor) else x for x in gen()]
|
||||
GlobalCounters.reset()
|
||||
Device[Device.DEFAULT].synchronize()
|
||||
st = time.perf_counter_ns()
|
||||
model(*early_gen)
|
||||
Device[Device.DEFAULT].synchronize()
|
||||
tms.append(time.perf_counter_ns() - st)
|
||||
mem_used = (GlobalCounters.mem_used - global_mem_used) / 1e9
|
||||
|
||||
kernels_used = jit_cache_count(model.captured.linear) if getattr(model, "captured", None) is not None else None
|
||||
print(f"{nm}: used {mem_used/1e9:.2f} GB and {kernels_used} kernels in {min(tms)/1e6:.2f} ms")
|
||||
assert mem_used < max_memory_allowed, f"{nm} used more than {max_memory_allowed:.3f} GB - {mem_used:.3} GB used"
|
||||
assert (max_memory_allowed - mem_used) / max_memory_allowed < 0.2, f"{max_memory_allowed:.3f} GB is too far from {mem_used:.3} GB used"
|
||||
if kernels_used:
|
||||
if kernels_used > max_kernels_allowed: raise KernelCountException(max_kernels_allowed, kernels_used)
|
||||
if (max_kernels_allowed - kernels_used) / max_kernels_allowed >= 0.2:
|
||||
raise KernelCountException(max_kernels_allowed, kernels_used)
|
||||
if all_jitted:
|
||||
assert kernels_used > 0 and kernels_used == GlobalCounters.kernel_count or (kernels_used <= GlobalCounters.kernel_count and getattr(Device[Device.DEFAULT], "graph", None)), f"only {kernels_used} out of {GlobalCounters.kernel_count} were jitted" # noqa: E501
|
||||
|
||||
supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
|
||||
|
||||
class TestRealWorld(unittest.TestCase):
|
||||
def setUp(self):
|
||||
gc.collect()
|
||||
global global_mem_used
|
||||
global_mem_used = GlobalCounters.mem_used
|
||||
np.random.seed(2002)
|
||||
|
||||
@slow
|
||||
@unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
|
||||
def test_stable_diffusion(self):
|
||||
params = unet_params
|
||||
params["model_ch"] = 8
|
||||
params["ctx_dim"] = 8
|
||||
params["num_res_blocks"] = 1
|
||||
params["n_heads"] = 2
|
||||
model = UNetModel(**params)
|
||||
derandomize_model(model)
|
||||
@TinyJit
|
||||
def test(t, t2): return model(t, Tensor([801]), t2).realize()
|
||||
helper_test("test_sd", lambda: (Tensor.randn(1, 4, 32, 32), Tensor.randn(1, 77, params["ctx_dim"])), test, 0.011, 515)
|
||||
|
||||
def test_unet_resblock(self):
|
||||
model = [ResBlock(16, 24, 16) for _ in range(4)]
|
||||
derandomize_model(model)
|
||||
@TinyJit
|
||||
def test(t, t2):
|
||||
for l in model: t = l(t, t2)
|
||||
return t.realize()
|
||||
|
||||
# TODO: support _offset on CL to get mem down to 0.0002
|
||||
exp_mem = 0.00037 if Device.DEFAULT == "CL" else 0.0002
|
||||
helper_test("test_unet_resblock", lambda: (Tensor.empty(4, 16, 8, 8), Tensor.empty(1, 24)), test, exp_mem, 37)
|
||||
|
||||
@unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
|
||||
def test_llama(self):
|
||||
self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
|
||||
|
||||
args_tiny = {"dim": 1024, "hidden_dim": 2048, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-05, "vocab_size": 1000}
|
||||
model = LLaMaTransformer(**args_tiny)
|
||||
derandomize_model(model)
|
||||
@TinyJit
|
||||
def test(t): return model(t, 0).realize()
|
||||
# TODO: test first token vs rest properly
|
||||
helper_test("test_llama", lambda: (Tensor([[1,2,3,4]]),), test, 0.23, 118, all_jitted=True)
|
||||
|
||||
@unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
|
||||
def test_gpt2(self):
|
||||
self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
|
||||
|
||||
args_tiny = {"dim": 1024, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-5, "vocab_size": 1000}
|
||||
model = GPT2Transformer(**args_tiny)
|
||||
derandomize_model(model)
|
||||
@TinyJit
|
||||
def test(t, v):
|
||||
with Context(JIT=0): return model(t, v).realize()
|
||||
helper_test("test_gpt2", lambda: (Tensor([[1,]]),Variable("pos", 1, 100).bind(1)), test, 0.23, 168, all_jitted=True)
|
||||
|
||||
@slow
|
||||
def test_train_mnist(self):
|
||||
from examples.beautiful_mnist import Model
|
||||
with Context(TRAINING=1):
|
||||
model = Model()
|
||||
optimizer = optim.Adam(get_parameters(model))
|
||||
BS = 32
|
||||
|
||||
@TinyJit
|
||||
def train(X):
|
||||
out = model(X)
|
||||
loss = out.mean()
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.017, 103)
|
||||
|
||||
@slow
|
||||
def test_forward_cifar(self):
|
||||
BS = 32
|
||||
# with training batchnorm still though
|
||||
with Context(TRAINING=1):
|
||||
model = SpeedyResNet(Tensor.ones((12,3,2,2)))
|
||||
@TinyJit
|
||||
def run(X): return model(X)
|
||||
helper_test("forward_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), run, 0.033, 27)
|
||||
|
||||
@slow
|
||||
def test_train_cifar(self):
|
||||
with Context(TRAINING=1):
|
||||
model = SpeedyResNet(Tensor.ones((12,3,2,2)))
|
||||
optimizer = optim.SGD(get_parameters(model), lr=0.01, momentum=0.8, nesterov=True, weight_decay=0.15)
|
||||
BS = 32
|
||||
|
||||
@TinyJit
|
||||
def train(X):
|
||||
out = model(X)
|
||||
loss = out.mean()
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
helper_test("train_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), train, 0.12, 126)
|
||||
|
||||
@unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
|
||||
def test_train_cifar_hyp(self):
|
||||
self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
|
||||
with Context(TRAINING=1):
|
||||
model = SpeedyResNet(Tensor.ones((12,3,2,2)))
|
||||
optimizer = optim.SGD(get_parameters(model), lr=0.01, momentum=hyp['opt']['momentum'], nesterov=True, weight_decay=hyp['opt']['bias_decay'])
|
||||
initial_div_factor = hyp['opt']['initial_div_factor']
|
||||
final_lr_ratio = hyp['opt']['final_lr_ratio']
|
||||
pct_start = hyp['opt']['percent_start']
|
||||
lr_scheduler = OneCycleLR(optimizer, max_lr=hyp['opt']['bias_lr'], pct_start=pct_start, div_factor=initial_div_factor,
|
||||
final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=4)
|
||||
assert not np.isnan(lr_scheduler.min_lr), "lr too small or initial_div_facotr too big for half"
|
||||
|
||||
@slow
|
||||
def test_bert(self):
|
||||
with Context(TRAINING=1):
|
||||
args_tiny = {"attention_probs_dropout_prob": 0.0, "hidden_dropout_prob": 0.0, "vocab_size": 30522, "type_vocab_size": 2,
|
||||
"max_position_embeddings": 512, "hidden_size": 128, "intermediate_size": 512, "num_attention_heads": 2, "num_hidden_layers": 2}
|
||||
model = BertForPretraining(**args_tiny)
|
||||
optimizer = optim.LAMB(get_parameters(model))
|
||||
|
||||
@TinyJit
|
||||
def train(input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor,
|
||||
masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor):
|
||||
lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
|
||||
loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
from examples.mlperf.helpers import get_fake_data_bert
|
||||
data = get_fake_data_bert(BS=4)
|
||||
for v in data.values(): v.to_(Device.DEFAULT)
|
||||
|
||||
helper_test("train_bert", lambda: (data["input_ids"], data["segment_ids"], data["input_mask"], data["masked_lm_positions"], \
|
||||
data["masked_lm_ids"], data["masked_lm_weights"], data["next_sentence_labels"]), train, 0.31, 400)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,171 @@
|
||||
# modified from
|
||||
# https://github.com/arogozhnikov/einops/blob/master/tests/test_examples.py
|
||||
# https://github.com/arogozhnikov/einops/blob/master/tests/test_ops.py
|
||||
# https://github.com/arogozhnikov/einops/blob/master/tests/test_parsing.py
|
||||
|
||||
import numpy as np
|
||||
import unittest
|
||||
from tinygrad import Tensor
|
||||
|
||||
|
||||
class test_rearrange_examples(unittest.TestCase):
|
||||
def test1(self):
|
||||
# transpose
|
||||
x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
|
||||
y = x.rearrange("b c h w -> b h w c")
|
||||
assert tuple(y.shape) == (10, 30, 40, 20)
|
||||
|
||||
def test2(self):
|
||||
# view / reshape
|
||||
x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
|
||||
y = x.rearrange("b c h w -> b (c h w)")
|
||||
assert tuple(y.shape) == (10, 20 * 30 * 40)
|
||||
|
||||
def test3(self):
|
||||
# depth-to-space
|
||||
x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
|
||||
y = x.rearrange("b (c h1 w1) h w -> b c (h h1) (w w1)", h1=2, w1=2)
|
||||
assert tuple(y.shape) == (10, 5, 30 * 2, 40 * 2)
|
||||
|
||||
def test4(self):
|
||||
# space-to-depth
|
||||
x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
|
||||
y = x.rearrange("b c (h h1) (w w1) -> b (h1 w1 c) h w", h1=2, w1=2)
|
||||
assert tuple(y.shape) == (10, 20 * 4, 30 // 2, 40 // 2)
|
||||
|
||||
def test5(self):
|
||||
# simple transposition
|
||||
x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
|
||||
y = x.rearrange("b1 sound b2 letter -> b1 b2 sound letter")
|
||||
assert tuple(y.shape) == (10, 30, 20, 40)
|
||||
|
||||
def test6(self):
|
||||
# parsing parameters
|
||||
x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
|
||||
t = x.rearrange("b c h w -> (b h w) c")
|
||||
t = t[:, ::2] # replacement for dot-product, just changes size of second axis
|
||||
assert tuple(t.shape) == (10 * 30 * 40, 10)
|
||||
|
||||
def test7(self):
|
||||
x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
|
||||
# split of embedding into groups
|
||||
y1, y2 = x.rearrange("b (c g) h w -> g b c h w", g=2)
|
||||
assert tuple(y1.shape) == (10, 10, 30, 40)
|
||||
assert tuple(y2.shape) == (10, 10, 30, 40)
|
||||
|
||||
def test8(self):
|
||||
x = Tensor(np.arange(10 * 20 * 1 * 1, dtype=np.int32).reshape([10, 20, 1, 1]))
|
||||
# squeeze - unsqueeze
|
||||
y = x.rearrange("b c () () -> b c")
|
||||
assert tuple(y.shape) == (10, 20)
|
||||
y = y.rearrange("b c -> c b () ()")
|
||||
assert tuple(y.shape) == (20, 10, 1, 1)
|
||||
|
||||
def test9(self):
|
||||
x = Tensor(np.arange(10 * 20 * 1 * 1, dtype=np.int32).reshape([10, 20, 1, 1]))
|
||||
# squeeze - unsqueeze
|
||||
y = x.rearrange("b c 1 1 -> b c")
|
||||
assert tuple(y.shape) == (10, 20)
|
||||
y = y.rearrange("b1 c -> c b1 1 1")
|
||||
assert tuple(y.shape) == (20, 10, 1, 1)
|
||||
|
||||
|
||||
class test_rearrange_ops(unittest.TestCase):
|
||||
def test_rearrange_errors(self):
|
||||
x = Tensor.zeros([1, 1, 1, 1, 1])
|
||||
x.rearrange("a b c d ... -> a b c ... d")
|
||||
bad_patterns = [
|
||||
"a b c d (...) -> a b c ... d", # collapsed ellipsis on input
|
||||
"a b (c d ... -> a b c ... d", # unbalanced brackets
|
||||
"a b* c d ... -> a b c ... d", # not alphanumeric
|
||||
"a b c d -> a b c d -> a b c d", # two "->"
|
||||
"a ... c ... -> ... a ... c", # two "..."
|
||||
"a b c d e -> f b c d e", # name mismatch
|
||||
]
|
||||
for pattern in bad_patterns:
|
||||
with self.assertRaises(AssertionError):
|
||||
x.rearrange(pattern)
|
||||
|
||||
x.rearrange("... -> (...)")
|
||||
with self.assertRaises(AssertionError):
|
||||
x.rearrange("(...) -> (...)")
|
||||
|
||||
y = Tensor.zeros([8, 1])
|
||||
y.rearrange("(a1 a2 a3) b -> b a3 a2 a1", a1=2, a2=2)
|
||||
with self.assertRaises(RuntimeError):
|
||||
## should fail as not enough dimensions specified
|
||||
y.rearrange("(a1 a2 a3) b -> b a3 a2 a1", a1=2)
|
||||
with self.assertRaises(ValueError):
|
||||
## should fail as 6 does not divide 8
|
||||
y.rearrange("(a1 a2 a3) b -> b a3 a2 a1", a1=3, a2=2)
|
||||
with self.assertRaises(AssertionError):
|
||||
## incorrect dimension provided for an axis that is only permuted
|
||||
y.rearrange("(a1 a2 a3) b -> b a3 a2 a1", a1=2, a2=2, b=2)
|
||||
with self.assertRaises(AssertionError):
|
||||
## unused axis provided
|
||||
y.rearrange("(a b c) d -> a b c d", b=2, c=2, e=2)
|
||||
|
||||
|
||||
class test_rearrange_parsing(unittest.TestCase):
|
||||
def test_elementary_axis_name(self):
|
||||
for name in [
|
||||
"a",
|
||||
"b",
|
||||
"h",
|
||||
"dx",
|
||||
"h1",
|
||||
"zz",
|
||||
"i9123",
|
||||
"somelongname",
|
||||
"Alex",
|
||||
"camelCase",
|
||||
"u_n_d_e_r_score",
|
||||
"unreasonablyLongAxisName",
|
||||
]:
|
||||
Tensor.ones((1,)).rearrange(f"{name} -> {name}")
|
||||
|
||||
for name in ["2b", "12", "_startWithUnderscore", "endWithUnderscore_", "_"]:
|
||||
with self.assertRaises(AssertionError):
|
||||
Tensor.ones((1,)).rearrange(f"{name} -> {name}")
|
||||
|
||||
with self.assertRaises(RuntimeError):
|
||||
Tensor.ones((1,)).rearrange(" -> ")
|
||||
|
||||
def test_invalid_expressions(self):
|
||||
# double ellipsis should raise an error
|
||||
def _test_expression(expression: str):
|
||||
Tensor.ones((2, 3, 4, 5, 6)).rearrange(f"{expression} -> {expression}")
|
||||
|
||||
_test_expression("... a b c d")
|
||||
with self.assertRaises(AssertionError):
|
||||
_test_expression("... a b c d ...")
|
||||
with self.assertRaises(AssertionError):
|
||||
_test_expression("... a b c (d ...)")
|
||||
with self.assertRaises(AssertionError):
|
||||
_test_expression("(... a) b c (d ...)")
|
||||
|
||||
# double/missing/enclosed parenthesis
|
||||
Tensor.ones((2, 3, 4, 5, 6)).rearrange("a b c d ... -> (a) b c (d ...)")
|
||||
with self.assertRaises(AssertionError):
|
||||
_test_expression("(a)) b c (d ...)")
|
||||
with self.assertRaises(AssertionError):
|
||||
_test_expression("(a b c (d ...)")
|
||||
with self.assertRaises(AssertionError):
|
||||
_test_expression("(a) (()) b c (d ...)")
|
||||
with self.assertRaises(AssertionError):
|
||||
_test_expression("(a) ((b c) (d ...))")
|
||||
|
||||
# invalid identifiers
|
||||
_test_expression("camelCase under_scored cApiTaLs ß ...")
|
||||
with self.assertRaises(AssertionError):
|
||||
_test_expression("1a")
|
||||
with self.assertRaises(AssertionError):
|
||||
_test_expression("_pre")
|
||||
with self.assertRaises(AssertionError):
|
||||
_test_expression("...pre")
|
||||
with self.assertRaises(AssertionError):
|
||||
_test_expression("pre...")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
21
artifacts/package_sources/tinygrad/test/null/test_resnet.py
Normal file
21
artifacts/package_sources/tinygrad/test/null/test_resnet.py
Normal file
@@ -0,0 +1,21 @@
|
||||
import unittest
|
||||
from extra.models import resnet
|
||||
|
||||
class TestResnet(unittest.TestCase):
|
||||
def test_model_load(self):
|
||||
model = resnet.ResNet18()
|
||||
model.load_from_pretrained()
|
||||
|
||||
model = resnet.ResNeXt50_32X4D()
|
||||
model.load_from_pretrained()
|
||||
|
||||
def test_model_load_no_fc_layer(self):
|
||||
model = resnet.ResNet18(num_classes=None)
|
||||
model.load_from_pretrained()
|
||||
|
||||
model = resnet.ResNeXt50_32X4D(num_classes=None)
|
||||
model.load_from_pretrained()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,28 @@
|
||||
import unittest
|
||||
from tinygrad.uop.ops import PatternMatcher, UOp, graph_rewrite, Ops, UPat, BottomUpGate
|
||||
|
||||
def assert_not_reached(): assert False, "This function should not be reached"
|
||||
def gate(): raise BottomUpGate
|
||||
|
||||
class TestBottomUpGate(unittest.TestCase):
|
||||
def test_basic_bottom_up_gate(self):
|
||||
"""Test that BottomUpGate stops bottom-up"""
|
||||
pm = PatternMatcher([
|
||||
(UPat(Ops.ADD), gate),
|
||||
(UPat(Ops.MUL), assert_not_reached)
|
||||
])
|
||||
|
||||
a,b,c = UOp.variable("a",0,10), UOp.variable("b",0,10), UOp.variable("c",0,10)
|
||||
graph_rewrite((a*a)+(b*c), pm, bottom_up=True)
|
||||
|
||||
def test_bottom_up_gate_with_rewriting(self):
|
||||
pm = PatternMatcher([
|
||||
(UPat.var("a")+UPat.var("a"), lambda a: 2*a),
|
||||
(UPat(Ops.MUL), gate),
|
||||
(UPat(Ops.CONST), assert_not_reached)
|
||||
])
|
||||
a = UOp.variable("a",0,10)
|
||||
graph_rewrite(a+a, pm, bottom_up=True)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
2003
artifacts/package_sources/tinygrad/test/null/test_schedule.py
Normal file
2003
artifacts/package_sources/tinygrad/test/null/test_schedule.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,41 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor, Variable, Context
|
||||
from tinygrad.helpers import cpu_events
|
||||
from tinygrad.schedule import schedule_cache
|
||||
|
||||
def schedule_one():
|
||||
Tensor([1]).schedule_linear()
|
||||
|
||||
class TestScheduleCache(unittest.TestCase):
|
||||
def test_bound_variable_var_vals(self):
|
||||
v = Variable('pos', 1, 100)
|
||||
x = Tensor.ones(10).contiguous().realize()
|
||||
|
||||
t = x + Tensor(v.bind(42))
|
||||
_, var_vals = t.linear_with_vars()
|
||||
self.assertEqual(var_vals, {'pos': 42})
|
||||
|
||||
def test_disable_schedule_cache(self):
|
||||
schedule_cache.clear()
|
||||
|
||||
# test write
|
||||
with Context(SCACHE=0): schedule_one()
|
||||
self.assertEqual(len(schedule_cache), 0)
|
||||
with Context(SCACHE=1):
|
||||
schedule_one()
|
||||
schedule_one()
|
||||
self.assertEqual(len(schedule_cache), 1)
|
||||
|
||||
# test read
|
||||
with Context(PROFILE=1):
|
||||
cpu_events.clear()
|
||||
with Context(SCACHE=0): schedule_one()
|
||||
num_events_no_cache = len(cpu_events)
|
||||
|
||||
cpu_events.clear()
|
||||
with Context(SCACHE=1): schedule_one()
|
||||
num_events_cache = len(cpu_events)
|
||||
self.assertLess(num_events_cache, num_events_no_cache)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,596 @@
|
||||
import unittest, itertools
|
||||
|
||||
from tinygrad.codegen.late.coalesce import indexing_simplify
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.uop.ops import UOp, Ops, graph_rewrite
|
||||
from tinygrad.uop.weak import pm_lower_index_dtype
|
||||
from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load
|
||||
from tinygrad.helpers import Context
|
||||
from test.helpers import full_rewrite
|
||||
from test.null.test_uop_symbolic import check_uop_against_string
|
||||
|
||||
# symbolic-only idx + valid simplification (no late lowering of FLOORDIV/FLOORMOD)
|
||||
def simplify_valid_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load, name="simplify_valid_idx")
|
||||
# image-aware idx + valid simplification: adds the codegen-layer matcher that drops provably in-bounds gates
|
||||
def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+indexing_simplify, name="simplify_image_idx")
|
||||
|
||||
def get_gated_load_uop(valid:UOp, idx:UOp):
|
||||
return UOp(Ops.LOAD, src=(
|
||||
UOp.param(0, dtypes.float, (1024,)).index(idx.valid(valid)),
|
||||
))
|
||||
|
||||
def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
|
||||
return UOp(Ops.LOAD, src=(
|
||||
UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)),
|
||||
))
|
||||
|
||||
def Special(expr, nmax): return UOp(Ops.SPECIAL, src=(UOp.const(nmax),), arg=expr)
|
||||
def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax)
|
||||
def Range(n, nmax): return UOp.range(nmax, n)
|
||||
|
||||
class TestValidIdxSimplification(unittest.TestCase):
|
||||
def check(self, load, sidx, svalid, extra=()):
|
||||
load = simplify_valid_idx(UOp.sink(load, *extra)).src[0]
|
||||
off = load.src[0].src[1]
|
||||
check_uop_against_string(self, off.get_idx(), sidx)
|
||||
check_uop_against_string(self, off.get_valid(), svalid)
|
||||
|
||||
def test_cumsum(self):
|
||||
gidx0 = Special("gidx0", 5)
|
||||
lidx0 = Special("lidx0", 4)
|
||||
gate = (gidx0*4+lidx0<19).ne(True)
|
||||
idx = gidx0*4+lidx0-19
|
||||
load = get_gated_load_uop(gate, idx)
|
||||
self.check(load,
|
||||
"0",
|
||||
"(((lidx0+(gidx0*4))<19)!=True)")
|
||||
|
||||
def test_simplify_within_valid1(self):
|
||||
ridx0 = Range(0, 4)
|
||||
ridx1 = Range(1, 4)
|
||||
ridx2 = Range(2, 4)
|
||||
ridx3 = Range(3, 4)
|
||||
valid = ((ridx0*3+ridx1)<8) & ((((ridx0*3+ridx1)//8+ridx2*3+ridx3)%4)<2)
|
||||
idx = ridx0+ridx1+ridx2+ridx3
|
||||
load = get_gated_load_uop(valid, idx)
|
||||
self.check(load,
|
||||
"(((r0+r1)+r2)+r3)",
|
||||
"((((r0*3)+r1)<8)&((((r2*3)+r3)%4)<2))")
|
||||
|
||||
def test_simplify_within_valid2(self):
|
||||
gidx0 = Special("gidx0", 56)
|
||||
ridx0 = Range(0, 3)
|
||||
alu0 = gidx0+ridx0
|
||||
valid = (alu0 < 57) & (alu0 >= 1)
|
||||
self.assertIsNone(simplify_valid(valid))
|
||||
|
||||
def test_valid_order_matters1(self):
|
||||
ridx0 = Range(0, 2)
|
||||
v0 = ridx0<1
|
||||
v1 = ((ridx0*5+1)%6)<5
|
||||
self.assertEqual(simplify_valid(v0&v1).render(), "(r0<1)")
|
||||
self.assertEqual(simplify_valid(v1&v0).render(), "(r0<1)")
|
||||
|
||||
def test_valid_order_matters2(self):
|
||||
gidx0 = Special("gidx0", 13)
|
||||
gidx1 = Special("gidx1", 13)
|
||||
ridx0 = Range(0, 4)
|
||||
alu0 = (gidx1+(ridx0*13))
|
||||
v0 = (gidx0+11)%14<11
|
||||
v1 = (alu0+((gidx0+39)//42))%14<11
|
||||
v2 = gidx0<3
|
||||
v3 = alu0<42
|
||||
|
||||
for v in itertools.permutations([v0,v1,v2,v3]):
|
||||
self.assertEqual(simplify_valid(v[0]&v[1]&v[2]&v[3]).render(), "False")
|
||||
|
||||
def test_simplify_valid_from_div(self):
|
||||
x = Variable("x", -100, 100)
|
||||
valid = ((x<0)&((100%x).cast(dtypes.bool)))
|
||||
# NOTE: this simplifies the (100%x) part somehow, still has two clauses
|
||||
self.assertIsNotNone(simplify_valid(valid))
|
||||
self.assertEqual(len(list(valid.split_uop(Ops.AND))), 2)
|
||||
|
||||
@unittest.expectedFailure # TODO: fix
|
||||
def test_from_merge_views(self):
|
||||
# taken from test_merges_from_fuzzer1
|
||||
# generated by
|
||||
# v0 = View(shape=(2, 4), strides=(2, 1), offset=-2, mask=((0, 2), (2, 4)), contiguous=False)
|
||||
# v1 = View(shape=(2, 4, 2, 2), strides=(4, 0, -2, -1), offset=3, mask=None, contiguous=False)
|
||||
# s = ShapeTracker((v0, v1))
|
||||
# idx, valid = s.to_indexed_uops()
|
||||
# print(f"{idx.render()=}")
|
||||
# print(f"{valid.render()=}")
|
||||
|
||||
# s = ShapeTracker((View(shape=(2, 4, 2, 2), strides=(2, 0, 0, -1), offset=1, mask=((0, 2), (0, 4), (0, 1), (0, 2)), contiguous=False),))
|
||||
# idx, valid = s.to_indexed_uops()
|
||||
# print(f"{idx.render()=}")
|
||||
# print(f"{valid.render()=}")
|
||||
ridx0 = Range(0, 2)
|
||||
ridx2 = Range(2, 2)
|
||||
ridx3 = Range(3, 2)
|
||||
idx = (((ridx0*2)+((((ridx2*2)+(ridx3*3))+3)%4))+-2)
|
||||
valid = ((((((ridx2*2)+(ridx3*3))+3)%4)<2)!=True) # noqa: E712
|
||||
load = get_gated_load_uop(valid, idx)
|
||||
self.check(load,
|
||||
"(((r0*2)+(r3*-1))+1)",
|
||||
"(r2<1)")
|
||||
|
||||
def test_load_in_valid(self):
|
||||
# from FUSE_ARANGE=1 python test/test_ops.py TestOps.test_scatter_add
|
||||
# can lead to OOB
|
||||
ridx2 = Range(2, 4)
|
||||
lidx0 = Special("lidx0", 3)
|
||||
gidx0 = Special("gidx0", 2)
|
||||
idx=(((lidx0+(gidx0*3))+(ridx2*5))+40)
|
||||
valid = (lidx0+(gidx0*3)) < 5
|
||||
val7 = get_gated_load_uop(valid, idx)
|
||||
valid2 = valid & val7.cast(dtypes.bool).logical_not()
|
||||
self.assertIsNone(simplify_valid(valid2))
|
||||
|
||||
def test_valid_becomes_const1(self):
|
||||
# from DSP mobilenetv2
|
||||
ridx0 = Range(0, 30)
|
||||
ridx1 = Range(1, 7)
|
||||
ridx2 = Range(2, 2)
|
||||
alu11 = (ridx1+ridx2)
|
||||
alu15 = ((alu11+1)//7)
|
||||
idx = (alu15*-31)+(((((alu11+218)//224)+ridx0)%30)*1568)
|
||||
valid = (ridx2<1)&(ridx1<6)
|
||||
load = get_gated_load_uop(valid, idx)
|
||||
# prevent ridx1 and ridx2 from being shrunk
|
||||
red = load.reduce(ridx1, ridx2, arg=Ops.ADD)
|
||||
self.check(load,
|
||||
"(r0*1568)",
|
||||
"((r2<1)&(r1<6))",
|
||||
extra=(red,))
|
||||
|
||||
def test_valid_becomes_const1_z3(self):
|
||||
from z3 import Ints, Solver, And, If, Not, unsat
|
||||
ridx0, ridx1, ridx2, alu11, alu15 = Ints('ridx0 ridx1 ridx2 alu11 alu15')
|
||||
alu11 = (ridx1+ridx2)
|
||||
alu15 = ((alu11+1)/7)
|
||||
idx = (alu15*-31)+(((((alu11+218)/224)+ridx0)%30)*1568)
|
||||
valid = (ridx2<1)&(ridx1<6)
|
||||
load = If(valid, idx, 0)
|
||||
|
||||
# correct simplification
|
||||
s = Solver()
|
||||
s.add(And(0<=ridx0, ridx0<30, 0<=ridx1, ridx1<7, 0<=ridx2, ridx2<2))
|
||||
simplifed_idx = (ridx0*1568)
|
||||
simplifed_load = If(valid, simplifed_idx, 0)
|
||||
s.add(Not(load == simplifed_load)) # Check if they are NOT equivalent
|
||||
assert s.check() == unsat, f"The expressions are not equivalent. {s.model()=}"
|
||||
|
||||
# new solver for a wrong simplified expression
|
||||
s = Solver()
|
||||
s.add(And(0<=ridx0, ridx0<30, 0<=ridx1, ridx1<7, 0<=ridx2, ridx2<2))
|
||||
wrong_simplifed_idx = (ridx0*1567)+ridx1
|
||||
wrong_simplifed_load = If(valid, wrong_simplifed_idx, 0)
|
||||
s.add(Not(load == wrong_simplifed_load)) # Check if they are NOT equivalent
|
||||
assert s.check() != unsat, "The expressions are equivalent??"
|
||||
print("The expressions are not equivalent.")
|
||||
print(s.model())
|
||||
|
||||
def test_valid_becomes_const2(self):
|
||||
ridx0 = Range(0, 4)
|
||||
ridx1 = Range(1, 4)
|
||||
ridx2 = Range(2, 4)
|
||||
ridx3 = Range(3, 4)
|
||||
# TODO: this should also work without the extra nesting
|
||||
idx = (((ridx0+ridx1)+(ridx2+ridx3)+28)//30)
|
||||
valid = ((ridx0+ridx1)<1).ne(True) & ((ridx2+ridx3)<1).ne(True)
|
||||
load = get_gated_load_uop(valid, idx)
|
||||
self.check(load,
|
||||
"1",
|
||||
"((((r0+r1)<1)!=True)&(((r2+r3)<1)!=True))")
|
||||
|
||||
def test_valid_with_non_const_rhs(self):
|
||||
ridx0 = Range(0, 1024)
|
||||
ridx1 = Range(1, 4)
|
||||
ridx2 = Range(2, 4)
|
||||
valid = (ridx0<(ridx1*4 + ridx2))&(ridx0<-1).ne(True)
|
||||
idx = ridx0
|
||||
load = get_gated_load_uop(valid, idx)
|
||||
self.check(load,
|
||||
"r0",
|
||||
"(r0<((r1*4)+r2))")
|
||||
|
||||
class TestImageSimplification(unittest.TestCase):
|
||||
def check(self, load, svalid, sidx0, sidx1):
|
||||
load = simplify_image_idx(load.sink()).src[0]
|
||||
off = load.src[0]
|
||||
self.assertEqual(len(off.src), 3)
|
||||
idx0, idx1 = off.src[2].get_idx(), off.src[1].get_idx()
|
||||
check_uop_against_string(self, idx0, sidx0)
|
||||
check_uop_against_string(self, idx1, sidx1)
|
||||
self.assertEqual(off.src[1].get_valid(), off.src[2].get_valid())
|
||||
if svalid is not None:
|
||||
check_uop_against_string(self, off.src[1].get_valid(), svalid)
|
||||
else:
|
||||
self.assertEqual(off.src[1].get_valid(), UOp.const(True), "svalid is None but valid is not True")
|
||||
|
||||
def test_idx_gt_c(self):
|
||||
# (idx1 < c+1).ne(True) ? (..., idx1-1+c) : 0 can drop the valid
|
||||
# (idx1 < c+1).ne(True) -> idx > c
|
||||
gidx0 = Special("gidx0", 32)
|
||||
gidx1 = Special("gidx1", 32)
|
||||
shape = (10, 10, 4)
|
||||
load = get_load_image_uop(shape, (gidx1<1).ne(True), (gidx0, gidx1-1))
|
||||
self.check(load, None, "gidx0", "(gidx1+-1)")
|
||||
load = get_load_image_uop(shape, (gidx1<1).ne(True), (gidx0, gidx1-2))
|
||||
self.check(load, None, "gidx0", "(gidx1+-2)")
|
||||
|
||||
# should match any one of the AND clause and drop the matched statement from valid
|
||||
valid = (gidx0<1).ne(True) & (gidx1<1).ne(True)
|
||||
load = get_load_image_uop(shape, valid, (gidx0+1, gidx1-1))
|
||||
self.check(load, "((gidx0<1)!=True)", "(gidx0+1)", "(gidx1+-1)")
|
||||
|
||||
valid = (gidx0<1).ne(True) & (gidx1<1).ne(True)
|
||||
load = get_load_image_uop(shape, valid, (gidx0, gidx1-1))
|
||||
self.check(load, "((gidx0<1)!=True)", "gidx0", "(gidx1+-1)")
|
||||
|
||||
def test_idx_lt_bound(self):
|
||||
# (idx1 < image_bound) ? (..., idx1) : 0 can drop the valid
|
||||
gidx0 = Special("gidx0", 32)
|
||||
gidx1 = Special("gidx1", 32)
|
||||
load = get_load_image_uop((10, 10, 4), gidx1<10, (gidx0, gidx1))
|
||||
self.check(load, None, "gidx0", "gidx1")
|
||||
|
||||
# same thing, valid has a div
|
||||
load = get_load_image_uop((10, 10, 4), gidx1//2<5, (gidx0, gidx1))
|
||||
self.check(load, None, "gidx0", "gidx1")
|
||||
|
||||
# 10x20 image, not out of bound
|
||||
load = get_load_image_uop((20, 10, 4), gidx1<10, (gidx0, gidx1))
|
||||
self.check(load, "(gidx1<10)", "gidx0", "gidx1")
|
||||
|
||||
def test_generic_idx_lt_bound(self):
|
||||
# (idx1 < image_bound - c) ? (..., idx1 + c) : 0 can drop the valid
|
||||
gidx0 = Special("gidx0", 32)
|
||||
gidx1 = Special("gidx1", 32)
|
||||
shape = (10, 10, 4)
|
||||
load = get_load_image_uop(shape, (gidx1<8), (gidx0, gidx1+2))
|
||||
self.check(load, None, "gidx0", "(gidx1+2)")
|
||||
|
||||
load = get_load_image_uop(shape, (gidx1<5), (gidx0, gidx1+5))
|
||||
self.check(load, None, "gidx0", "(gidx1+5)")
|
||||
|
||||
def test_valid_empty_set(self):
|
||||
gidx0 = Special("gidx0", 32)
|
||||
gidx1 = Special("gidx1", 32)
|
||||
shape = (32, 32, 4)
|
||||
idx = (gidx0%2, gidx1+2)
|
||||
# not empty
|
||||
load = get_load_image_uop(shape, gidx0<8, idx)
|
||||
self.check(load, "(gidx0<8)", "(gidx0%2)", "(gidx1+2)")
|
||||
|
||||
# empty -> invalid
|
||||
load = get_load_image_uop(shape, (gidx0<8) & (gidx0<8).ne(True), idx)
|
||||
with Context(NOOPT=1, SPEC=0):
|
||||
load = full_rewrite(load.sink()).src[0]
|
||||
self.assertFalse(load.op_in_backward_slice_with_self(Ops.LOAD))
|
||||
|
||||
def test_openpilot_conv1(self):
|
||||
# first conv in openpilot
|
||||
# kernel in tinygrad ae5d1407ee844a97a52ad3756835d38e7e2b9e1b https://gist.github.com/chenyuxyz/39c2d4e9a076b46731c67d345ff066b6
|
||||
idx1 = Special("idx1", 32)
|
||||
idx2 = Special("idx2", 64)
|
||||
# ridx0 = Variable("ridx0", 0, 5)
|
||||
# ridx1 = Variable("ridx1", 0, 2)
|
||||
# ridx2 = Variable("ridx2", 0, 2)
|
||||
ridx0 = Range(0, 6)
|
||||
ridx1 = Range(1, 3)
|
||||
ridx2 = Range(2, 3)
|
||||
|
||||
alu1 = ((idx2*2)+ridx1)
|
||||
alu4 = ((idx1*48)+(ridx2*6)+ridx0)
|
||||
|
||||
valid = ((((idx2*2)+(ridx1))<1).ne(True))&((((idx1*8)+(ridx2))<1).ne(True))
|
||||
shape = (128, 1536, 4)
|
||||
idx = ((alu4+1530)%1536, alu1+((idx1+((ridx2+7)//8)+31)//32)+(-2))
|
||||
|
||||
load = get_load_image_uop(shape, valid, idx)
|
||||
self.check(load, None, "((((idx1*48)+(r2*6))+r0)+-6)", "(((idx2*2)+r1)+-1)")
|
||||
|
||||
def test_openpilot_conv2(self):
|
||||
# conv in test/external/external_test_valid_remove.py
|
||||
idx1 = Special("idx1", 32)
|
||||
idx2 = Special("idx2", 64)
|
||||
# ridx0 = Variable("ridx0", 0, 2)
|
||||
# ridx1 = Variable("ridx1", 0, 2)
|
||||
# ridx2 = Variable("ridx2", 0, 2)
|
||||
ridx0 = Range(0, 3)
|
||||
ridx1 = Range(1, 3)
|
||||
ridx2 = Range(2, 3)
|
||||
|
||||
alu1 = ((idx2*2)+ridx1)
|
||||
alu3 = ((idx1*24)+(ridx2*3)+ridx0)
|
||||
|
||||
valid = ((((idx2*2)+ridx1)<1).ne(True))&((((idx1*8)+ridx2)<1).ne(True))
|
||||
shape = (128, 768, 4)
|
||||
idx = ((alu3+765)%768, alu1+((idx1+((ridx2+7)//8)+31)//32)+(-2))
|
||||
load = get_load_image_uop(shape, valid, idx)
|
||||
|
||||
self.check(load, None, "((((idx1*24)+(r2*3))+r0)+-3)", "(((idx2*2)+r1)+-1)")
|
||||
|
||||
def test_openpilot_conv3(self):
|
||||
# in openpilot 0.9.7
|
||||
idx0 = Special("idx0", 64)
|
||||
idx1 = Special("idx1", 2)
|
||||
idx2 = Special("idx2", 4)
|
||||
ridx0 = Range(0, 7)
|
||||
ridx1 = Range(1, 7)
|
||||
|
||||
alu2 = ((idx2*2)+ridx0)
|
||||
alu4 = ((idx1*8)+ridx1)
|
||||
alu6 = ((idx1*512)+(ridx1*64)+idx0)
|
||||
|
||||
valid = (alu2<11)&(alu4<3).ne(True)
|
||||
shape = (8, 1024, 4)
|
||||
idx = (((alu6+832)%1024),(alu2+((idx1+((ridx1+5)//8)+1)//2)+(-4)))
|
||||
|
||||
load = get_load_image_uop(shape, valid, idx)
|
||||
|
||||
self.check(load,
|
||||
"((((idx2*2)+r0)<11)&((((idx1*8)+r1)<3)!=True))",
|
||||
"(idx0+(idx1*512+r1*64)+-192)",
|
||||
"((((idx2*2)+r0)+(((idx1+((r1+5)//8))+1)//2))+-4)")
|
||||
|
||||
def test_simplify1(self):
|
||||
# idx has the form (A % m, A // m + k) and valid has (c0 < A) and (A < c1)
|
||||
gidx = Special("gidx", 512)
|
||||
valid = (gidx<488) & (gidx<480).ne(True)
|
||||
idx = ((gidx*3+18)%26, (gidx*3+18)//26-56)
|
||||
load = get_load_image_uop((1, 26, 4), valid, idx)
|
||||
self.check(load, None, "((gidx*3)+-1438)", "0")
|
||||
|
||||
def test_simplify2(self):
|
||||
# from DEV=CL DEBUG=4 FORWARD_ONLY=1 IMAGE=2 python3 test/test_ops.py TestOps.test_simple_padding_conv2d
|
||||
lidx = Special("lidx", 4)
|
||||
valid = (lidx<3) & (lidx<1).ne(True)
|
||||
idx = ((lidx+1)%2, (lidx+1)//2-1)
|
||||
load = get_load_image_uop((1, 2, 4), valid, idx)
|
||||
self.check(load, None, "(lidx+-1)", "0")
|
||||
|
||||
def test_simplify3(self):
|
||||
# from openpilot
|
||||
idx0 = Special("idx0", 265)
|
||||
valid = (idx0<201).ne(True)
|
||||
idx = ((idx0+55)%64, (idx0+55)//64-4)
|
||||
load = get_load_image_uop((1, 64, 4), valid, idx)
|
||||
self.check(load, None, "(idx0+-201)", "0")
|
||||
|
||||
def test_simplify4(self):
|
||||
idx0 = Special("idx0", 512)
|
||||
shape = (4, 64, 4)
|
||||
alu2 = ((idx0*4+1)%32)
|
||||
alu3 = ((idx0*4+2)%32)
|
||||
alu4 = ((idx0*4+3)%32)
|
||||
alu5 = (idx0*4%32)
|
||||
alu8 = (idx0//8%32//4)
|
||||
alu9 = idx0<256
|
||||
|
||||
load = get_load_image_uop(shape, alu9, (((alu8+(alu2*8))%64),(alu2//8)))
|
||||
self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32+8)", "(idx0//2%4)")
|
||||
|
||||
load = get_load_image_uop(shape, alu9, (((alu8+(alu3*8))%64),(alu3//8)))
|
||||
self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32+16)", "(idx0//2%4)")
|
||||
|
||||
load = get_load_image_uop(shape, alu9, (((alu8+(alu4*8))%64),(alu4//8)))
|
||||
self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32+24)", "(idx0//2%4)")
|
||||
|
||||
load = get_load_image_uop(shape, alu9, (((alu8+(alu5*8))%64),(alu5//8)))
|
||||
self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32)", "(idx0//2%4)")
|
||||
|
||||
def test_simplify5(self):
|
||||
# openpilot 0.9.7, chunk replacement to simplify
|
||||
shape = (10, 384, 4)
|
||||
idx0 = Special("idx0", 16)
|
||||
idx1 = Special("idx1", 24)
|
||||
alu0 = idx0*4
|
||||
alu1 = (idx1*256)+alu0
|
||||
alu2 = idx1//3
|
||||
alu3 = ((alu1+1)%768)
|
||||
idx = ((idx0+((((alu3//640)+alu2)%8)*16)+128),((alu3//64)%10))
|
||||
valid = alu3<640
|
||||
|
||||
load = get_load_image_uop(shape, valid, idx)
|
||||
self.check(load, None, "((idx0+((idx1//3)*16))+128)", "((idx1%3)*4)")
|
||||
|
||||
def test_simplify6(self):
|
||||
# from openpilot
|
||||
# the valid implies the numerator of the div/mod is positive and can be simplified with floordiv rules
|
||||
idx1 = Special("idx1", 16)
|
||||
idx2 = Special("idx2", 64)
|
||||
ridx3 = Range(3, 3)
|
||||
ridx4 = Range(4, 3)
|
||||
ridx5 = Range(5, 3)
|
||||
alu0 = ((idx2*1536)+(ridx4*768)+ridx3+(idx1*24)+(ridx5*3)+-771)%768
|
||||
alu1 = ((idx2*1536)+(ridx4*768)+ridx3+(idx1*24)+(ridx5*3)+-771)//768
|
||||
valid = (((idx2+ridx4)<1)!=1)&(((idx1+ridx5)<1)!=1)
|
||||
load = get_load_image_uop((128, 768, 4), valid, (alu0, alu1))
|
||||
self.check(load, None, "((((idx1*24)+(r5*3))+r3)+-3)", "(((idx2*2)+r4)+-1)")
|
||||
|
||||
def test_simplify7(self):
|
||||
# DEBUG=2 ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1397 ALLOWED_GATED_READ_IMAGE=94 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916 # noqa: E501
|
||||
# kernel 143
|
||||
gidx0 = Special("gidx0", 32)
|
||||
lidx0 = Special("lidx0", 16)
|
||||
lidx1 = Special("lidx1", 8)
|
||||
r0 = Range(0, 7)
|
||||
|
||||
# buf.render()='UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), arg=1, src=())'
|
||||
alu0 = ((gidx0*2+(lidx0*128+r0*64+lidx1*8+-183)%64*64+(lidx0*128+r0*64+lidx1*8+-183)//64%32*4096+1)//4%1024)
|
||||
alu1 = ((gidx0*2+(lidx0*128+r0*64+lidx1*8+-183)%64*64+(lidx0*128+r0*64+lidx1*8+-183)//64%32*4096+1)//4096)
|
||||
valid = ((lidx1<7)&((((lidx0*2+r0)<3)!=1)&((lidx0*2+r0)<35)))
|
||||
load = get_load_image_uop((32, 1024, 4), valid, (alu0, alu1))
|
||||
self.check(load, None, "(lidx1*128+gidx0//2+144)", "(lidx0*2+r0+-3)")
|
||||
|
||||
# same idx, written without the inline simplification of the inner div/mod
|
||||
alu0 = ((gidx0*2+lidx1*512+(lidx0*8192+r0*4096)+-11711)//4%1024)
|
||||
alu1 = (lidx0*2+r0+-3)
|
||||
valid = ((lidx1<7)&((((lidx0*2+r0)<3)!=1)&((lidx0*2+r0)<35)))
|
||||
load = get_load_image_uop((32, 1024, 4), valid, (alu0, alu1))
|
||||
self.check(load, None, "(lidx1*128+gidx0//2+144)", "(lidx0*2+r0+-3)")
|
||||
|
||||
def test_simplify8(self):
|
||||
# from openpilot compile3, kernel r_4_16_8_16_4_4_3_3n1
|
||||
# valid guarantees A >= 0, so divmod simplifies and gate is removed
|
||||
gidx0 = Special("gidx0", 16)
|
||||
gidx1 = Special("gidx1", 4)
|
||||
lidx0 = Special("lidx0", 8)
|
||||
lidx1 = Special("lidx1", 16)
|
||||
A = gidx0 + gidx1*8192 + lidx0*1024 + lidx1*64 - 1040
|
||||
valid = ((lidx1 < 1).ne(True)) & (((gidx1 + lidx0) < 1).ne(True))
|
||||
load = get_load_image_uop((32, 1024, 4), valid, (A % 1024, A // 1024))
|
||||
self.check(load, None, "(gidx0+lidx1*64+-16)", "(lidx0+gidx1*8+-1)")
|
||||
|
||||
def test_simplify9(self):
|
||||
# from openpilot compile3, kernel r_32_16_8_4_4_7_7 (image 1x16384)
|
||||
# valid guarantees A1 >= 0 and A1 < 512, gate should be removable
|
||||
gidx0 = Special("gidx0", 32)
|
||||
lidx0 = Special("lidx0", 16)
|
||||
lidx1 = Special("lidx1", 8)
|
||||
r0 = Range(0, 7)
|
||||
A1 = lidx0*32 + r0*32 + lidx1*4 - 99
|
||||
valid = ((lidx1 < 1).ne(True)) & ((lidx0 + r0) < 3).ne(True) & ((lidx0 + r0) < 19)
|
||||
alu0 = gidx0 + (A1 % 32)*32 + (A1 // 32 % 16)*1024
|
||||
load = get_load_image_uop((1, 16384, 4), valid, (alu0, UOp.const(0)))
|
||||
try:
|
||||
self.check(load, None, "(gidx0+lidx0*1024+r0*1024+lidx1*128+-3168)", "0")
|
||||
except AssertionError:
|
||||
# TODO: fold valid
|
||||
self.check(load, "(((lidx1<1)!=True)&(((lidx0+r0)<3)!=True)&((lidx0+r0)<19))",
|
||||
"(gidx0+lidx1*128+(lidx0*1024+r0*1024)+-3168)", "0")
|
||||
|
||||
def test_simplify10(self):
|
||||
# from openpilot compile3, kernel r_16_8_4_4_4_4_7_7 (image 1x8192)
|
||||
# valid guarantees A1 >= 0 and A1 < 128, gate should be removable
|
||||
gidx0 = Special("gidx0", 16)
|
||||
lidx0 = Special("lidx0", 8)
|
||||
lidx1 = Special("lidx1", 4)
|
||||
lidx2 = Special("lidx2", 4)
|
||||
r0 = Range(0, 7)
|
||||
A1 = lidx0*16 + r0*16 + lidx1*4 - 51
|
||||
valid = ((lidx1 < 1).ne(True)) & ((lidx0 + r0) < 3).ne(True) & ((lidx0 + r0) < 11)
|
||||
alu0 = lidx2 + gidx0*4 + (A1 % 16)*64 + (A1 // 16 % 8)*1024
|
||||
load = get_load_image_uop((1, 8192, 4), valid, (alu0, UOp.const(0)))
|
||||
try:
|
||||
self.check(load, None, "(lidx2+gidx0*4+lidx0*1024+r0*1024+lidx1*256+-3264)", "0")
|
||||
except AssertionError:
|
||||
# TODO: fold valid
|
||||
self.check(load, "(((lidx1<1)!=True)&(((lidx0+r0)<3)!=True)&((lidx0+r0)<11))",
|
||||
"(lidx2+gidx0*4+lidx1*256+(lidx0*1024+r0*1024)+-3264)", "0")
|
||||
|
||||
def test_drop_non_monotonic_window(self):
|
||||
# two-sided window valid (645 <= gidx0 < 653) on a non-monotonic index (lane split via %4 and //4):
|
||||
# gidx0 outside the window pushes idx_x out of the (1, 48) image, so the gate is dropped
|
||||
gidx0 = Special("gidx0", 1064)
|
||||
r12 = Range(12, 3)
|
||||
valid = ((gidx0 < 645).ne(True)) & (gidx0 < 653)
|
||||
idx = (r12*4 + (gidx0+3)%4 + (gidx0+3)//4*24 - 3888, UOp.const(0))
|
||||
load = get_load_image_uop((1, 48, 4), valid, idx)
|
||||
self.check(load, None, "(r12*4+(gidx0+3)%4+(gidx0+3)//4*24+-3888)", "0")
|
||||
|
||||
def test_drop_gate_committed_in_the_index_pass(self):
|
||||
# the fused index pass runs without symbolic, so committing a weak src must not leave a CAST that
|
||||
# symbolic later folds inside the index only: the gate's copy of the expression has to stay the same node
|
||||
f = UOp.variable("f", 0.0, 9.0, dtypes.float)
|
||||
idx_y = (f + UOp.const(1.0)).cast(dtypes.int)
|
||||
load = get_load_image_uop((10, 10, 4), (UOp.const(-1) < idx_y) & (idx_y < UOp.const(10)),
|
||||
(Special("gidx0", 10), idx_y))
|
||||
off = graph_rewrite(load.sink(), pm_lower_index_dtype+indexing_simplify, ctx={}).src[0].src[0]
|
||||
self.assertEqual(off.src[1].get_valid(), UOp.const(True))
|
||||
|
||||
class TestDropTrueGate(unittest.TestCase):
|
||||
def test_drop_true_gate_on_index(self):
|
||||
# test that INDEX with a constant True valid gets simplified to drop the valid
|
||||
from tinygrad.codegen.late.coalesce import indexing_simplify
|
||||
from tinygrad.uop.ops import graph_rewrite
|
||||
from tinygrad.uop.symbolic import sym
|
||||
buf = UOp.param(0, dtypes.int, (1,))
|
||||
idx = UOp.const(0)
|
||||
true_gate = UOp.const(True)
|
||||
index_with_gate = UOp(Ops.INDEX, src=(buf, idx.valid(true_gate)))
|
||||
# apply the optimization
|
||||
result = graph_rewrite(index_with_gate, sym+indexing_simplify)
|
||||
# the True valid should be dropped (INDEX should only have 2 sources)
|
||||
self.assertEqual(len(result.src), 2, "True valid should be dropped from INDEX")
|
||||
|
||||
class TestRangeShrink(unittest.TestCase):
|
||||
def get_ranges(self, sink):
|
||||
with Context(NOOPT=1, SPEC=0):
|
||||
result = full_rewrite(sink)
|
||||
return [u for u in result.toposort() if u.op is Ops.RANGE]
|
||||
|
||||
def test_range_shrink_single_guard(self):
|
||||
# range 0..203 guarded by r < 4 everywhere -> shrink to 0..3
|
||||
r = Range(0, 204)
|
||||
load = get_gated_load_uop(r < UOp.const(4), r)
|
||||
ranges = self.get_ranges(load.sink())
|
||||
self.assertEqual(len(ranges), 1)
|
||||
self.assertEqual(ranges[0].src[0].val, 4)
|
||||
|
||||
def test_range_shrink_picks_max_guard(self):
|
||||
# two loads guard the same range with r < 4 and r < 8 -> shrink to max(4, 8) = 8
|
||||
r = Range(0, 204)
|
||||
load1 = get_gated_load_uop(r < UOp.const(4), r)
|
||||
load2 = get_gated_load_uop(r < UOp.const(8), r)
|
||||
ranges = self.get_ranges(UOp.sink(load1, load2))
|
||||
self.assertEqual(len(ranges), 1)
|
||||
self.assertEqual(ranges[0].src[0].val, 8)
|
||||
|
||||
def test_range_no_shrink_guard_ge_max(self):
|
||||
# guard r < 300 with range max 204 -> no shrink (guard doesn't constrain)
|
||||
r = Range(0, 204)
|
||||
load = get_gated_load_uop(r < UOp.const(300), r)
|
||||
ranges = self.get_ranges(load.sink())
|
||||
self.assertEqual(len(ranges), 1)
|
||||
self.assertEqual(ranges[0].src[0].val, 204)
|
||||
|
||||
def test_range_no_shrink_when_unguarded_elsewhere(self):
|
||||
# one load guards r < 4, but another load uses r without a gate -> no shrink
|
||||
r = Range(0, 204)
|
||||
load1 = get_gated_load_uop(r < UOp.const(4), r)
|
||||
load2 = UOp(Ops.LOAD, src=(UOp.param(1, dtypes.float, (204,)).index(r),))
|
||||
ranges = self.get_ranges(UOp.sink(load1, load2))
|
||||
self.assertEqual(len(ranges), 1)
|
||||
self.assertEqual(ranges[0].src[0].val, 204)
|
||||
|
||||
def test_range_no_shrink_when_used_in_reduce(self):
|
||||
# range used in both a gated load AND directly in the reduce expression -> no shrink
|
||||
r = Range(0, 204)
|
||||
gated_load = get_gated_load_uop(r < UOp.const(4), r)
|
||||
red = (r.cast(dtypes.float) + gated_load).reduce(r, arg=Ops.ADD)
|
||||
ranges = self.get_ranges(red.sink())
|
||||
self.assertEqual(len(ranges), 1)
|
||||
self.assertEqual(ranges[0].src[0].val, 204)
|
||||
|
||||
def test_range_shrink_to_single_iteration(self):
|
||||
# guard r < 1 shrinks range to 1 -> single iteration, range eliminated entirely
|
||||
r = Range(0, 204)
|
||||
load = get_gated_load_uop(r < UOp.const(1), r)
|
||||
ranges = self.get_ranges(load.sink())
|
||||
self.assertEqual(len(ranges), 0)
|
||||
|
||||
def test_range_shrink_store_where_invalid(self):
|
||||
# emulates mask.where(x.pad_to(mask.shape), Invalid): range should shrink accordingly
|
||||
from tinygrad.dtype import Invalid
|
||||
r = Range(0, 204)
|
||||
x = (r < 4).where(UOp.const(1.0), Invalid)
|
||||
ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink())
|
||||
self.assertEqual(len(ranges), 1)
|
||||
self.assertEqual(ranges[0].src[0].val, 4)
|
||||
|
||||
def test_range_shrink_store_where_invalid_flipped(self):
|
||||
# above, but flipped
|
||||
from tinygrad.dtype import Invalid
|
||||
r = Range(0, 204)
|
||||
x = (r < 4).where(UOp.const(1.0), Invalid)
|
||||
ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink())
|
||||
self.assertEqual(len(ranges), 1)
|
||||
self.assertEqual(ranges[0].src[0].val, 4)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,121 @@
|
||||
import unittest
|
||||
from tinygrad import Variable
|
||||
from tinygrad.uop.ops import UOp
|
||||
|
||||
|
||||
class TestFuzzFailure(unittest.TestCase):
|
||||
def test_fuzz_failure1(self):
|
||||
v1=Variable('v1', 0, 8)
|
||||
v2=Variable('v2', 0, 2)
|
||||
v3=Variable('v3', 0, 1)
|
||||
expr = (((((((((((((((((((((((0//4)%2)//8)+-2)+-4)+-3)+v1)+-4)+v2)+-2)+v3)+v2)//3)%7)*1)//2)+v2)*-1)+2)+1)+0)+-3)+v3)
|
||||
v1_val, v2_val, v3_val = UOp.const(8), UOp.const(0), UOp.const(0)
|
||||
num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
self.assertEqual(num, rn)
|
||||
|
||||
def test_fuzz_failure2(self):
|
||||
v1=Variable('v1', 0, 16)
|
||||
v2=Variable('v2', 0, 5)
|
||||
v3=Variable('v3', 0, 3)
|
||||
expr = (((((((((((((((((((((((((0*4)//5)*2)*-1)*-2)+-4)*4)*2)*3)*4)+-4)*4)+v2)+v2)+v3)//3)+v2)+v1)//9)+3)+1)//1)+-4)//4)*2)
|
||||
expr = (((((v1+(v2+(((v3+(v2*2))+1)//3)))+4)//9)+-57)//(9*4))
|
||||
v1_val, v2_val, v3_val = UOp.const(6), UOp.const(0), UOp.const(0)
|
||||
num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
self.assertEqual(num, rn)
|
||||
|
||||
def test_fuzz_failure3(self):
|
||||
v1=Variable('v1', 0, 2)
|
||||
v2=Variable('v2', 0, 1)
|
||||
v3=Variable('v3', 0, 2)
|
||||
expr = (((((((((((((((((((0//2)//3)+v3)+0)+-4)*-2)*-2)+-1)+2)+3)+v3)+0)//8)*-3)+0)*-2)*-4)*-2)//5)
|
||||
v1_val, v2_val, v3_val = UOp.const(0), UOp.const(0), UOp.const(0)
|
||||
num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
self.assertEqual(num, rn)
|
||||
|
||||
def test_fuzz_failure4(self):
|
||||
v1=Variable('v1', 0, 2)
|
||||
v2=Variable('v2', 0, 3)
|
||||
v3=Variable('v3', 0, 4)
|
||||
expr = (((((((((((((((((((((((((((((0*-2)+0)*-1)//9)//6)//8)+v1)*-4)+v2)//4)//8)+4)*3)+v1)+v3)//8)//7)+4)+v3)*-4)+1)+v1)*3)+4)*2)//5)//2)//3)*-4)
|
||||
v1_val, v2_val, v3_val = UOp.const(2), UOp.const(0), UOp.const(2)
|
||||
num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
self.assertEqual(num, rn)
|
||||
|
||||
def test_fuzz_failure5(self):
|
||||
v1=Variable('v1', 0, 1)
|
||||
v2=Variable('v2', 0, 1)
|
||||
v3=Variable('v3', 0, 3)
|
||||
expr = ((((((((((((((0+v2)+v1)*0)+v2)//1)//7)+-2)+v2)+v1)*4)+-3)//5)+v2)+1)
|
||||
v1_val, v2_val, v3_val = UOp.const(0), UOp.const(0), UOp.const(0)
|
||||
num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
self.assertEqual(num, rn)
|
||||
|
||||
def test_fuzz_failure6(self):
|
||||
v1=Variable('v1', 0, 8)
|
||||
v2=Variable('v2', 0, 64)
|
||||
v3=Variable('v3', 0, 128)
|
||||
expr = (((((((((((((((((((((((((((((0//3)+4)+v1)//2)+-1)//1)*1)*-1)*4)//5)+v1)//6)+v1)*-1)+-4)+v2)+-2)*-3)+v3)+-4)+-2)*-1)//8)//4)*-4)+3)+v3)*
|
||||
-2)+v2)
|
||||
v1_val, v2_val, v3_val = UOp.const(8), UOp.const(3), UOp.const(2)
|
||||
num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
self.assertEqual(num, rn)
|
||||
|
||||
def test_fuzz_failure7(self):
|
||||
v1=Variable('v1', 0, 64)
|
||||
v2=Variable('v2', 0, 5)
|
||||
v3=Variable('v3', 0, 128)
|
||||
expr = (((((((((((((((((((((((((((((0+v2)*-4)+0)//9)+-4)*-2)*3)*4)//9)+v3)+v1)//4)+v1)+v3)+-1)*4)//4)+v2)//7)//3)+v1)+v2)+v3)+1)*2)//4)*3)+-1)*1)
|
||||
v1_val, v2_val, v3_val = UOp.const(0), UOp.const(2), UOp.const(65)
|
||||
num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
self.assertEqual(num, rn)
|
||||
|
||||
def test_fuzz_failure8(self):
|
||||
v1=Variable('v1', 0, 2)
|
||||
v2=Variable('v2', 0, 8)
|
||||
v3=Variable('v3', 0, 9)
|
||||
expr = (((((((0+-1)+2)+v1)*-2)//3)+v1)*-4)
|
||||
v1_val, v2_val, v3_val = UOp.const(0), UOp.const(0), UOp.const(0)
|
||||
num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
self.assertEqual(num, rn)
|
||||
|
||||
def test_fuzz_failure9(self):
|
||||
v1=Variable('v1', 0, 256)
|
||||
v2=Variable('v2', 0, 1)
|
||||
v3=Variable('v3', 0, 8)
|
||||
expr = (((((((((((((((((((((((((((((0*-2)//1)+3)*-2)+-3)*-4)*1)+v1)+0)%2)%8)%9)+v2)%9)+-4)//4)+-1)*-2)+0)+v1)+v1)+3)+v1)+4)+-4)+0)*2)+-3)%6)
|
||||
v1_val, v2_val, v3_val = UOp.const(0), UOp.const(1), UOp.const(0)
|
||||
num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
self.assertEqual(num, rn)
|
||||
|
||||
def test_fuzz_failure10(self):
|
||||
v1=Variable("v1", 0, 256)
|
||||
v2=Variable("v2", 0, 32)
|
||||
v3=Variable("v3", 0, 32)
|
||||
x5 = (v1 <= 9).where(v1 * -4 - 4, v1 // 9) // 9
|
||||
expr = ((x5 >= -4).where(x5, (v2 % 3 + v2) // 5) * -1).maximum(((v1 * -2) % 6 + v3 % 1) * -1) * -1
|
||||
v1_val, v2_val, v3_val = UOp.const(9), UOp.const(0), UOp.const(0)
|
||||
num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
self.assertEqual(num, rn)
|
||||
|
||||
def test_fuzz_failure11(self):
|
||||
v1=Variable("v1", 0, 16)
|
||||
v2=Variable("v2", 0, 128)
|
||||
v3=Variable("v3", 0, 5)
|
||||
expr = (((v2 * 0).maximum(8) - v2 * 2) % 5 + v1 // 6 + v1 + 5) % 5
|
||||
v1_val, v2_val, v3_val = UOp.const(0), UOp.const(7), UOp.const(0)
|
||||
num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
|
||||
self.assertEqual(num, rn)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,129 @@
|
||||
import unittest
|
||||
from tinygrad import Variable
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.uop.ops import _broadcast_shape
|
||||
|
||||
class TestBroadcastShape(unittest.TestCase):
|
||||
def test_symbolic(self):
|
||||
v = Variable("v", 1, 10)
|
||||
self.assertEqual(_broadcast_shape((v,), (1,)), (v,))
|
||||
self.assertEqual(_broadcast_shape((v,), ()), (v,))
|
||||
self.assertEqual(_broadcast_shape((v,), (v,)), (v,))
|
||||
with self.assertRaises(IndexError): _broadcast_shape((v,), (5,))
|
||||
|
||||
def test_symbolic_vmin_zero(self):
|
||||
# a symbolic dim that may be 0 still broadcasts against 1 to itself
|
||||
v0 = Variable("v0", 0, 10)
|
||||
self.assertEqual(_broadcast_shape((v0,), (1,)), (v0,))
|
||||
self.assertEqual(_broadcast_shape((v0,), ()), (v0,))
|
||||
self.assertEqual(_broadcast_shape((3, v0), (3, 1)), (3, v0))
|
||||
with self.assertRaises(IndexError): _broadcast_shape((v0,), (5,))
|
||||
|
||||
class TestSymbolic(unittest.TestCase):
|
||||
def assert_tuple_equal(self, x, y):
|
||||
for a,b in zip(x,y): self.assertFalse(a != b)
|
||||
|
||||
def test_cat_dim0_is_expanded(self):
|
||||
i = Variable("i", 1, 5).bind(3)
|
||||
j = Variable("j", 1, 5).bind(3)
|
||||
k = Variable("k", 1, 5).bind(3)
|
||||
t = Tensor.rand(5, 4)[:i].cat(Tensor.rand(5, 4)[:j], dim=0).cat(Tensor.rand(5, 4)[:k], dim=0)
|
||||
self.assert_tuple_equal(t.shape, (i+j+k, 4))
|
||||
t = Tensor.rand(5, 3)[:i].cat(Tensor.rand(5, 3)[:i], dim=0).cat(Tensor.rand(3, 3), dim=0)
|
||||
self.assert_tuple_equal(t.shape, (2*i+3, 3))
|
||||
|
||||
def test_cat_dim1_strides(self):
|
||||
i = Variable("i", 1, 5).bind(4)
|
||||
j = Variable("j", 1, 5).bind(4)
|
||||
k = Variable("k", 1, 5).bind(4)
|
||||
t = Tensor.rand(3, 5)[:, :i].cat(Tensor.rand(3, 5)[:, :j], dim=1).cat(Tensor.rand(3, 5)[:, :k], dim=1)
|
||||
self.assert_tuple_equal(t.shape, (3, i+j+k))
|
||||
|
||||
class TestSymbolicVarVals(unittest.TestCase):
|
||||
def assert_equal(self, x, y): self.assertFalse(x != y)
|
||||
|
||||
def test_shrink_unbind(self):
|
||||
v = Variable("v", 1, 100)
|
||||
bv = Variable("v", 1, 100).bind(2)
|
||||
t = Tensor.rand(3, 4).shrink(((0,bv),(0,4)))
|
||||
unbound_st, var_val = t.uop.unbind_all()
|
||||
assert var_val == {v: 2}
|
||||
t = Tensor.rand(3, 4).shrink(((bv, bv+1), (0, 4)))
|
||||
unbound_st, var_val = t.uop.unbind_all()
|
||||
assert var_val == {v: 2}
|
||||
|
||||
class TestSymbolicReshape(unittest.TestCase):
|
||||
def test_reshape(self):
|
||||
a = Tensor.rand(5, 4)
|
||||
b = Tensor.rand(5, 6)
|
||||
for i in range(1, 6):
|
||||
vi = Variable("i", 1, 5).bind(i)
|
||||
ret = a[:vi]
|
||||
ret = ret.reshape((vi, 4))
|
||||
assert ret.shape == (vi, 4)
|
||||
ret = b[:vi]
|
||||
ret = ret.reshape((vi, 2, 3))
|
||||
assert ret.shape == (vi, 2, 3)
|
||||
|
||||
def test_two_symbol_reshape(self):
|
||||
t = Tensor.rand(5, 5)
|
||||
for i in range(1, 6):
|
||||
for j in range(1, 6):
|
||||
vi = Variable("i", 1, 5).bind(i)
|
||||
vj = Variable("j", 1, 5).bind(j)
|
||||
ret = t[:vi, :vj]
|
||||
ret = ret.reshape(vj, vi)
|
||||
assert ret.shape == (vj, vi)
|
||||
ret = ret.reshape(vi, vj)
|
||||
assert ret.shape == (vi, vj)
|
||||
ret = ret.reshape(1, vi*vj)
|
||||
assert ret.shape == (1, vi*vj)
|
||||
|
||||
class TestSymbolicExpand(unittest.TestCase):
|
||||
def test_expand_into_symbols(self):
|
||||
vi = Variable("i", 1, 5).bind(3)
|
||||
vj = Variable("j", 1, 5).bind(3)
|
||||
a = Tensor([[1], [2], [3]]).expand((3, vi))
|
||||
assert a.shape == (3, vi)
|
||||
a = a.reshape(3, vi, 1).expand((3, vi, vj))
|
||||
assert a.shape == (3, vi, vj)
|
||||
|
||||
def test_plus_expands_constant(self):
|
||||
a = Tensor.rand(3, 5)
|
||||
for i in range(1, 6):
|
||||
vi = Variable("i", 1, 5).bind(i)
|
||||
ret = a[:, :vi]
|
||||
ret = ret + 1
|
||||
self.assertTupleEqual(ret.shape, (3, vi))
|
||||
|
||||
def test_pad_then_expand_into_symbols(self):
|
||||
vi = Variable("i", 1, 10).bind(3)
|
||||
a = Tensor(1).unsqueeze(0).pad((0, 24)).unsqueeze(0).expand((vi, 25))
|
||||
self.assertEqual(a.shape, (vi, 25))
|
||||
self.assertEqual(a.reshape(25*vi).shape, (vi*25,))
|
||||
self.assertEqual(a.reshape(vi*25).shape, (vi*25,))
|
||||
|
||||
class TestSymbolicShrink(unittest.TestCase):
|
||||
def test_shrink_symbols_simple(self):
|
||||
vi = Variable("i", 1, 5)
|
||||
t = Tensor.rand(5, 5).shrink(((0, 5),(0,vi)))
|
||||
assert t.shape == (5, vi)
|
||||
|
||||
def test_shrink_symbols(self):
|
||||
vi = Variable("i", 1, 5)
|
||||
t = Tensor.rand(3, 5).shrink(((0, 2), (vi, vi+1)))
|
||||
assert t.shape == (2, 1)
|
||||
|
||||
class TestSymbolicContiguousViewOffset(unittest.TestCase):
|
||||
def test_shrink_from_start(self):
|
||||
v = Variable("v", 1, 10).bind(5)
|
||||
t = Tensor.rand(10).realize().shrink(((0, v),))
|
||||
self.assertEqual(t.uop.contiguous_view_offset(), 0)
|
||||
|
||||
def test_shrink_with_offset(self):
|
||||
v = Variable("v", 1, 7).bind(4)
|
||||
t = Tensor.rand(10).realize().shrink(((3, 3+v),))
|
||||
self.assertEqual(t.uop.contiguous_view_offset(), 3)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
212
artifacts/package_sources/tinygrad/test/null/test_tensor.py
Normal file
212
artifacts/package_sources/tinygrad/test/null/test_tensor.py
Normal file
@@ -0,0 +1,212 @@
|
||||
# tensor tests that pass on NULL backend (no copyout needed)
|
||||
import numpy as np
|
||||
import unittest
|
||||
from tinygrad import Tensor, Device, dtypes
|
||||
from tinygrad.uop.ops import Ops, UOp
|
||||
from tinygrad.renderer.ptx import PTXRenderer
|
||||
from tinygrad.renderer.nir import NIRRenderer
|
||||
from tinygrad.codegen import to_program
|
||||
from tinygrad.dtype import DType
|
||||
|
||||
x_init = np.random.randn(1,3).astype(np.float32)
|
||||
W_init = np.random.randn(3,3).astype(np.float32)
|
||||
m_init = np.random.randn(1,3).astype(np.float32)
|
||||
|
||||
class TestInferenceMode(unittest.TestCase):
|
||||
def test_inference(self):
|
||||
x = Tensor(x_init)
|
||||
m = Tensor(m_init)
|
||||
W = Tensor(W_init)
|
||||
tmp = x.mul(m)
|
||||
mm = tmp.matmul(W)
|
||||
out = mm.relu()
|
||||
out = out.sum()
|
||||
#out.backward()
|
||||
assert x.grad is None
|
||||
assert m.grad is None
|
||||
assert tmp.grad is None
|
||||
assert mm.grad is None
|
||||
assert W.grad is None
|
||||
|
||||
def test_no_grad_mode_context_manager(self):
|
||||
x = Tensor(x_init)
|
||||
m = Tensor(m_init)
|
||||
W = Tensor(W_init)
|
||||
def f(x, m, W):
|
||||
tmp = x.mul(m)
|
||||
mm = tmp.matmul(W)
|
||||
out = mm.relu()
|
||||
out = out.sum()
|
||||
#out.backward()
|
||||
assert x.grad is None
|
||||
assert m.grad is None
|
||||
assert tmp.grad is None
|
||||
assert mm.grad is None
|
||||
assert W.grad is None
|
||||
f(x, m, W)
|
||||
|
||||
class TestIdxUpcast(unittest.TestCase):
|
||||
def _find_op(self, ast: UOp, op: Ops):
|
||||
if ast.op is op: return ast
|
||||
for src in ast.src:
|
||||
if (ret:=self._find_op(src, op)) is not None: return ret
|
||||
def _schedule_render(self, a: Tensor):
|
||||
linear, _ = a.linear_with_vars()
|
||||
for si in linear.src:
|
||||
ast = si.src[0]
|
||||
if ast.op is Ops.SINK:
|
||||
renderer = Device[si.src[1].buffer.device].renderer
|
||||
prg = to_program(ast, renderer)
|
||||
return tuple(prg.src[1].src)
|
||||
|
||||
def _assert(self, dtype: DType, a: Tensor):
|
||||
uops = self._schedule_render(a)
|
||||
# Assert the dtype of the INDEX value, This will need be updated if UOp spec changes
|
||||
store = next(uop for uop in uops if uop.op is Ops.STORE)
|
||||
assert store.op is Ops.STORE
|
||||
idx = self._find_op(store, Ops.INDEX)
|
||||
# PTX and NIR turn Ops.INDEX into pointer arithmetic earlier than cstyle, plus it's already cast to int64
|
||||
if not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)):
|
||||
assert idx.op is Ops.INDEX
|
||||
idx_val = idx.src[1]
|
||||
self.assertFalse(idx_val.overflows(idx_val.dtype.scalar()))
|
||||
|
||||
# use expand to generate kernel that uses large idx
|
||||
def do_op_then_assert(self, dtype: DType, dim1, dim2, dim3):
|
||||
self._assert(dtype, Tensor.empty(dim1, dim2, 1).expand(-1, -1, dim3).contiguous())
|
||||
|
||||
@unittest.skipUnless(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "int64 is supported")
|
||||
def test_overflow(self):
|
||||
# 2**11, 2**11, 2**11 -> 2**33 will overflow when indexed
|
||||
self.do_op_then_assert(dtypes.long, 2048, 2048, 2048)
|
||||
|
||||
@unittest.skipUnless(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "int64 is supported")
|
||||
def test_overflow_sym(self):
|
||||
self.do_op_then_assert(dtypes.long, 2048, 2048, UOp.variable("dim3", 1, 2048).bind(32))
|
||||
|
||||
def test_regular(self):
|
||||
self.do_op_then_assert(dtypes.int, 64, 64, 64)
|
||||
|
||||
def test_regular_sym(self):
|
||||
self.do_op_then_assert(dtypes.int, 256, 256, UOp.variable("dim3", 1, 64).bind(32))
|
||||
|
||||
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "PTX and NIR always converts Ops.INDEX to int64")
|
||||
def test_symfold(self):
|
||||
# This would cause an overflow, but after sym fold it's within int32
|
||||
a = Tensor.arange(65535).clone()
|
||||
uops = self._schedule_render(a)
|
||||
assert all(uop.dtype is not dtypes.long for uop in uops)
|
||||
|
||||
@unittest.skipIf(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "int64 is supported")
|
||||
def test_int64_unsupported_overflow_sym(self):
|
||||
with self.assertRaises((KeyError, RuntimeError)):
|
||||
self.do_op_then_assert(dtypes.long, 2048, 2048, UOp.variable("dim3", 1, 2048).bind(32))
|
||||
|
||||
@unittest.skipIf(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "int64 is supported")
|
||||
@unittest.expectedFailure # bug in gpu dims limiting
|
||||
def test_int64_unsupported_overflow(self):
|
||||
with self.assertRaises((KeyError, RuntimeError)):
|
||||
self.do_op_then_assert(dtypes.long, 2048, 2048, 2048)
|
||||
|
||||
@unittest.skip("This is kept for reference, it requires large memory to run")
|
||||
def test_overflow_kernel_run(self):
|
||||
# This creates a total of 2**31+10 elements, requiring at least 2147 MB memory to run
|
||||
# Modified example from issue 3271
|
||||
a = Tensor.empty(2**11, 2**11, 1, dtype=dtypes.int8).permute((2, 0, 1)).expand((2**9+10, -1, -1)).contiguous()
|
||||
a.realize()
|
||||
|
||||
class TestTensorUnique(unittest.TestCase):
|
||||
def test_empty_bufs_unique(self):
|
||||
a = Tensor.empty(10, 10).contiguous()
|
||||
b = Tensor.empty(10, 10).contiguous()
|
||||
Tensor.realize(a,b)
|
||||
self.assertIsNot(a.uop.buffer, b.uop.buffer)
|
||||
|
||||
def test_zeros_bufs_unique_sep(self):
|
||||
a = Tensor.zeros(10, 10).contiguous()
|
||||
Tensor.realize(a)
|
||||
b = Tensor.zeros(10, 10).contiguous()
|
||||
Tensor.realize(b)
|
||||
self.assertIsNot(a.uop.buffer, b.uop.buffer)
|
||||
|
||||
def test_zeros_bufs_unique(self):
|
||||
a = Tensor.zeros(10, 10).contiguous()
|
||||
b = Tensor.zeros(10, 10).contiguous()
|
||||
Tensor.realize(a,b)
|
||||
self.assertIsNot(a.uop.buffer, b.uop.buffer)
|
||||
|
||||
def test_times_2_not_unique(self):
|
||||
a = Tensor.zeros(10, 10).contiguous()
|
||||
b = a * 2
|
||||
c = a * 2
|
||||
Tensor.realize(b,c)
|
||||
self.assertIs(b.uop.buffer, c.uop.buffer)
|
||||
|
||||
class TestRand(unittest.TestCase):
|
||||
def test_rand_large_tensor(self):
|
||||
# large tensor rand (num > uint32.max) should not crash in frontend
|
||||
Tensor.manual_seed(0)
|
||||
Tensor.rand(2**17, 2**17).schedule_linear()
|
||||
Tensor.rand(2**17, 2**17).schedule_linear()
|
||||
Tensor.rand(2**17, 2**17).schedule_linear()
|
||||
|
||||
class TestTensorConstLike(unittest.TestCase):
|
||||
def test_const_like_shape(self):
|
||||
t = Tensor.ones(3, 4)
|
||||
c = t.const_like(0)
|
||||
self.assertEqual(c.shape, (3, 4))
|
||||
self.assertEqual(c.dtype, t.dtype)
|
||||
|
||||
def test_const_like_multi_device(self):
|
||||
devs = ("NULL:0", "NULL:1")
|
||||
t = Tensor.ones(8, 4).shard(devs, axis=0)
|
||||
c = t.const_like(5)
|
||||
self.assertEqual(c.shape, (8, 4))
|
||||
self.assertIsNone(c.device)
|
||||
out = t+c
|
||||
self.assertEqual(out.device, t.device)
|
||||
self.assertEqual(out.uop.axis, 0)
|
||||
|
||||
def test_full_like_device_on_multi_raises(self):
|
||||
t = Tensor.ones(8, 4).shard(("NULL:0", "NULL:1"), axis=0)
|
||||
with self.assertRaises(RuntimeError): t.full_like(5, device="NULL")
|
||||
|
||||
class TestTensorDevice(unittest.TestCase):
|
||||
def test_create_from_single_device_tuple(self):
|
||||
(Tensor([1.0], device=(Device.DEFAULT,)) + Tensor([2.0])).realize()
|
||||
|
||||
class TestTensorPad(unittest.TestCase):
|
||||
# padding int tensor with float-only value (like -inf) must promote dtype to fit value
|
||||
def test_pad_int_with_neg_inf(self):
|
||||
t = Tensor.arange(9).reshape(1, 1, 3, 3)
|
||||
self.assertEqual(t.dtype, dtypes.int)
|
||||
r = t.pad((1, 2, 0, -1), value=-float('inf'))
|
||||
self.assertEqual(r.dtype, dtypes.weakfloat)
|
||||
self.assertEqual(r.shape, (1, 1, 2, 6))
|
||||
|
||||
class TestTensorDeviceMismatch(unittest.TestCase):
|
||||
def test_gather(self):
|
||||
x = Tensor.empty(3, 4, device="NULL")
|
||||
idx = Tensor.zeros(3, 4, dtype=dtypes.int32, device="NULL:1")
|
||||
with self.assertRaises(RuntimeError): x.gather(0, idx)
|
||||
def test_scatter_index(self):
|
||||
x = Tensor.zeros(3, 4, device="NULL")
|
||||
idx = Tensor.zeros(3, 4, dtype=dtypes.int32, device="NULL:1")
|
||||
src = Tensor.ones(3, 4, device="NULL")
|
||||
with self.assertRaises(RuntimeError): x.scatter(0, idx, src)
|
||||
def test_scatter_src(self):
|
||||
x = Tensor.zeros(3, 4, device="NULL")
|
||||
idx = Tensor.zeros(3, 4, dtype=dtypes.int32, device="NULL")
|
||||
src = Tensor.ones(3, 4, device="NULL:1")
|
||||
with self.assertRaises(RuntimeError): x.scatter(0, idx, src)
|
||||
def test_getitem_tensor_index(self):
|
||||
x = Tensor.empty(4, 5, device="NULL")
|
||||
idx = Tensor([0, 1], dtype=dtypes.int32, device="NULL:1")
|
||||
with self.assertRaises(RuntimeError): x[idx]
|
||||
def test_sparse_categorical_crossentropy(self):
|
||||
x = Tensor.zeros(2, 3, device="NULL")
|
||||
Y = Tensor([0, 1], dtype=dtypes.int32, device="NULL:1")
|
||||
with self.assertRaises(RuntimeError): x.sparse_categorical_crossentropy(Y)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,24 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor, dtypes
|
||||
from tinygrad.nn.state import TensorIO
|
||||
|
||||
class TestTensorIO(unittest.TestCase):
|
||||
def test_create(self):
|
||||
with self.assertRaises(ValueError):
|
||||
TensorIO(Tensor(b"Hello World").reshape(1, -1))
|
||||
with self.assertRaises(ValueError):
|
||||
TensorIO(Tensor([], dtype=dtypes.int64).reshape(1, -1))
|
||||
|
||||
def test_seek(self):
|
||||
t = Tensor(b"Hello World!")
|
||||
fobj = TensorIO(t)
|
||||
self.assertEqual(fobj.tell(), 0)
|
||||
self.assertEqual(fobj.seek(1), 1)
|
||||
self.assertEqual(fobj.seek(-2, 2), len(t) - 2)
|
||||
self.assertEqual(fobj.seek(1, 1), len(t) - 1)
|
||||
self.assertEqual(fobj.seek(10, 1), len(t))
|
||||
self.assertEqual(fobj.seek(10, 2), len(t))
|
||||
self.assertEqual(fobj.seek(-10, 0), 0)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,127 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor, dtypes
|
||||
from tinygrad.tensor import _METADATA
|
||||
from tinygrad.engine.realize import capturing
|
||||
from tinygrad.helpers import Context
|
||||
|
||||
@unittest.skip("tensor metadata is no longer supported")
|
||||
class TestTensorMetadata(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
_METADATA.set(None)
|
||||
self._ctx = Context(SCACHE=0)
|
||||
self._ctx.__enter__()
|
||||
def tearDown(self) -> None:
|
||||
self._ctx.__exit__(None, None, None)
|
||||
|
||||
@unittest.skip("why would this be true?")
|
||||
def test_exclude_noop_metadata(self):
|
||||
a = Tensor.rand(4, 4)*1
|
||||
self.assertEqual(a.uop.metadata[0].name, "__mul__")
|
||||
k = a.schedule_linear().src[-1]
|
||||
self.assertEqual([m.name for m in k.arg.metadata], ["rand"])
|
||||
|
||||
@unittest.skip("metadata not reaching kernel schedule")
|
||||
def test_exclude_const_metadata(self):
|
||||
a = Tensor.arange(4)
|
||||
b = Tensor.full((4,), -1, dtype=dtypes.int).contiguous()
|
||||
sched = a.schedule_linear(b)
|
||||
self.assertEqual([m.name for m in sched.src[0].arg.metadata], ["arange"])
|
||||
self.assertEqual([m.name for m in sched.src[1].arg.metadata], ["contiguous"])
|
||||
|
||||
def test_matmul(self):
|
||||
x = Tensor.rand(3)
|
||||
W = Tensor.rand(3, 3)
|
||||
out = x.matmul(W)
|
||||
self.assertEqual(out.uop.metadata[0].name, "matmul")
|
||||
si = out.schedule_linear().src[-1]
|
||||
self.assertEqual(len(si.arg.metadata), 1)
|
||||
self.assertEqual(si.arg.metadata[0].name, "matmul")
|
||||
|
||||
def test_relu(self):
|
||||
x = Tensor.rand(3)
|
||||
out = x.relu()
|
||||
self.assertEqual(out.uop.metadata[0].name, "relu")
|
||||
si = out.schedule_linear().src[-1]
|
||||
self.assertEqual(len(si.arg.metadata), 1)
|
||||
self.assertEqual(si.arg.metadata[0].name, "relu")
|
||||
|
||||
@unittest.skip("assign metadata no longer captured")
|
||||
def test_assign(self):
|
||||
x = Tensor.empty(10, 10).realize()
|
||||
x.assign(Tensor.ones(10, 10).contiguous())
|
||||
si = x.schedule_linear().src[-1]
|
||||
self.assertEqual(len(si.arg.metadata), 1)
|
||||
self.assertEqual(si.arg.metadata[0].name, "assign")
|
||||
|
||||
def test_complex(self):
|
||||
x = Tensor.rand(3)
|
||||
y = Tensor.rand(3)
|
||||
out = x.relu() * y.sigmoid()
|
||||
self.assertEqual(out.uop.metadata[0].name, "__mul__")
|
||||
self.assertEqual(out.uop.src[0].metadata[0].name, "relu")
|
||||
self.assertEqual(out.uop.src[1].metadata[0].name, "sigmoid")
|
||||
si = out.schedule_linear().src[-1]
|
||||
self.assertEqual(len(si.arg.metadata), 3)
|
||||
self.assertEqual(set(m.name for m in si.arg.metadata), {"relu", "sigmoid", "__mul__"})
|
||||
|
||||
@unittest.skip("flaky")
|
||||
def test_complex_backward(self):
|
||||
x = Tensor.rand(3).realize()
|
||||
y = Tensor.rand(3).realize()
|
||||
out = (x.relu() * y.sigmoid()).sum()
|
||||
self.assertEqual(out.uop.metadata[0].name, "sum")
|
||||
out.backward()
|
||||
self.assertEqual(x.grad.uop.metadata[0].name, "relu")
|
||||
#self.assertTrue(x.grad.uop.metadata[0].backward) # TODO: backward flag is False
|
||||
self.assertEqual(y.grad.uop.metadata[0].name, "sigmoid")
|
||||
#self.assertTrue(y.grad.uop.metadata[0].backward) # TODO: backward flag is False
|
||||
si = out.schedule_linear(x.grad, y.grad).src[-1]
|
||||
#self.assertEqual(len(si.arg.metadata), 3, f"failed with {si.arg.metadata}")
|
||||
# skip numpy, this is schedule cache
|
||||
self.assertSetEqual(set(m.name for m in si.arg.metadata if m.name != "numpy"), {"sigmoid", "relu"})
|
||||
#bw = [m for m in si.metadata if m.backward]
|
||||
#self.assertEqual(len(bw), 1)
|
||||
#self.assertEqual(bw[0].name, "sigmoid")
|
||||
|
||||
def test_tracemeta_0(self):
|
||||
with Context(TRACEMETA=0):
|
||||
x = Tensor.rand(3)
|
||||
y = Tensor.rand(3)
|
||||
out = (x.relu() * y.sigmoid()).sum()
|
||||
self.assertIsNone(out.uop.metadata)
|
||||
self.assertIsNone(out.uop.src[0].metadata)
|
||||
si = out.schedule_linear().src[-1]
|
||||
self.assertEqual(si.arg.metadata, ())
|
||||
|
||||
def _has_metadata(self, h, name):
|
||||
linears = []
|
||||
capturing.append(type("", (), {"add_linear": lambda _, linear, var_vals: linears.append(linear)})())
|
||||
try: h.realize()
|
||||
finally: capturing.clear()
|
||||
calls = [call for linear in linears for call in linear.src]
|
||||
return any(m.name == name for call in calls for m in call.arg.metadata)
|
||||
|
||||
def test_metadata_survives_realize_pending_assign(self):
|
||||
shared = Tensor.rand(4)
|
||||
c = Tensor.zeros(8).contiguous().realize()
|
||||
c[:4].assign(shared)
|
||||
self.assertTrue(self._has_metadata(c[:4].relu(), "relu"))
|
||||
|
||||
@unittest.expectedFailure
|
||||
def test_metadata_lost_realize_pending_assign(self):
|
||||
shared = Tensor.rand(4)
|
||||
c = Tensor.zeros(8).contiguous().realize()
|
||||
c[:4].assign(shared)
|
||||
self.assertTrue(self._has_metadata((c[:4] + shared).relu(), "relu"))
|
||||
|
||||
class TestTraceMetaShutdown(unittest.TestCase):
|
||||
def test_tracemeta_del_no_shutdown_error(self):
|
||||
import subprocess, os
|
||||
result = subprocess.run(['python3', '-c', 'from tinygrad import Tensor\n'
|
||||
'x=Tensor.eye(3); (x@x).sum().backward()'],
|
||||
env={**os.environ, "TRACEMETA": "2"}, capture_output=True)
|
||||
self.assertEqual(result.returncode, 0)
|
||||
self.assertNotIn(b"Exception", result.stderr)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,544 @@
|
||||
import math, unittest
|
||||
from dataclasses import replace
|
||||
from tinygrad import Tensor, dtypes, Context
|
||||
from tinygrad.uop.ops import ParamArg, UOp, UPat, Ops, PatternMatcher, graph_rewrite
|
||||
|
||||
_strip_unique_pm = PatternMatcher([
|
||||
(UPat(Ops.BUFFER, name="b"), lambda b: b.replace(arg=replace(b.arg, slot=0)) if isinstance(b.arg, ParamArg) and b.arg.slot != 0 else None),
|
||||
])
|
||||
def _strip_unique(u: UOp) -> UOp: return graph_rewrite(u, _strip_unique_pm)
|
||||
|
||||
def _t(*shape):
|
||||
return Tensor.arange(math.prod(shape)).reshape(*shape)
|
||||
|
||||
def _ti(data): # int32 index tensor
|
||||
return Tensor(data, dtype=dtypes.int32)
|
||||
|
||||
# Tensor().func().uop should be the same as UOp.func()
|
||||
def _check(tc: unittest.TestCase, t: Tensor, fn):
|
||||
tc.assertIs(fn(t).uop, fn(t.uop), f"\ntensor.uop = {fn(t).uop}\nuop = {fn(t.uop)}")
|
||||
|
||||
class TestTensorUOpBinop(unittest.TestCase):
|
||||
# Tensor's binop upcasts mixed dtypes via least_upper_dtype + explicit CAST; UOp should match.
|
||||
def test_mul_float_int(self):
|
||||
t = _t(3).float()
|
||||
self.assertIs((t * Tensor.arange(3)).uop, t.uop * UOp.arange(3))
|
||||
def test_mul_bool_int(self):
|
||||
t = _t(3)
|
||||
self.assertIs((t.eq(1) * Tensor.arange(3)).uop, t.uop.eq(1) * UOp.arange(3))
|
||||
# Tensor's ufix picks float dtype when scalar is float and self is int; UOp should match.
|
||||
def test_add_scalar_float_on_int(self): _check(self, _t(3), lambda x: x + 1.5)
|
||||
# div: Tensor.div (default case) delegates to ElementwiseMixin.div; trees must match for Tensor and UOp.
|
||||
def test_div_tensor_by_tensor(self):
|
||||
a, b = _t(4).float(), _t(4).float() + 1
|
||||
self.assertIs((a/b).uop, a.uop/b.uop)
|
||||
def test_div_int_by_int(self): _check(self, _t(4), lambda x: x / 3)
|
||||
def test_div_sum_by_sum(self): _check(self, _t(4).float(), lambda x: x.sum() / (x + 1).sum())
|
||||
def test_div_broadcast_tensor_by_tensor(self):
|
||||
a, b = _t(3, 4).float(), _t(4).float() + 1
|
||||
self.assertIs((a/b).uop, a.uop/b.uop)
|
||||
# isclose used `self == other` which is Python identity on UOp (not elementwise); now uses .eq().
|
||||
def test_isclose(self):
|
||||
t = _t(4).float()
|
||||
self.assertIs(t.isclose(t).uop, t.uop.isclose(t.uop))
|
||||
# __floordiv__/mod/fmod and div(rounding_mode=...) dispatch on dtype in mixin
|
||||
def test_floordiv_int(self): _check(self, _t(4), lambda x: x // 3)
|
||||
def test_floordiv_float(self): _check(self, _t(4).float() + 1.5, lambda x: x // 2.0)
|
||||
def test_rfloordiv_int(self): _check(self, _t(4)+1, lambda x: 7 // x)
|
||||
def test_mod_int(self): _check(self, _t(4), lambda x: x % 3)
|
||||
def test_mod_float(self): _check(self, _t(4).float() + 1.5, lambda x: x % 2.0)
|
||||
def test_div_trunc_int(self): _check(self, _t(4), lambda x: x.div(3, rounding_mode="trunc"))
|
||||
def test_div_trunc_float(self):_check(self, _t(4).float() + 1.5, lambda x: x.div(2.0, rounding_mode="trunc"))
|
||||
def test_fmod_int(self): _check(self, _t(4), lambda x: x.fmod(3))
|
||||
def test_fmod_float(self): _check(self, _t(4).float() + 1.5, lambda x: x.fmod(2.0))
|
||||
def test_floordiv_bool(self): _check(self, _t(4).cast(dtypes.bool), lambda x: x // True)
|
||||
def test_mod_bool(self): _check(self, _t(4).cast(dtypes.bool), lambda x: x % True)
|
||||
def test_fmod_bool(self): _check(self, _t(4).cast(dtypes.bool), lambda x: x.fmod(True))
|
||||
|
||||
class TestTensorUOpClone(unittest.TestCase):
|
||||
def test_clone(self):
|
||||
t = _t(3, 4).float()
|
||||
self.assertIs(_strip_unique(t.clone().uop), _strip_unique(t.uop.clone()))
|
||||
def test_clone_deviceless_const(self):
|
||||
u = UOp.const(2.0)
|
||||
self.assertIs(_strip_unique(Tensor(u).clone().uop), _strip_unique(u.clone()))
|
||||
|
||||
class TestTensorUOpGradient(unittest.TestCase):
|
||||
def test_gradient(self):
|
||||
x = _t(3, 3).float()
|
||||
z = (x * 2).sum()
|
||||
(tg,) = z.gradient(x)
|
||||
(ug,) = z.uop.gradient(x.uop)
|
||||
self.assertIs(tg.uop, ug)
|
||||
|
||||
class TestTensorUOpGetitem(unittest.TestCase):
|
||||
# ---- pure slice patterns ----
|
||||
def test_slice_full(self): _check(self, _t(4), lambda x: x[slice(None)])
|
||||
def test_slice_positive(self): _check(self, _t(8), lambda x: x[1:5])
|
||||
def test_slice_open_start(self): _check(self, _t(8), lambda x: x[:5])
|
||||
def test_slice_open_stop(self): _check(self, _t(8), lambda x: x[3:])
|
||||
def test_slice_negative_start(self): _check(self, _t(8), lambda x: x[-3:])
|
||||
def test_slice_negative_stop(self): _check(self, _t(8), lambda x: x[:-2])
|
||||
def test_slice_both_negative(self): _check(self, _t(8), lambda x: x[-5:-1])
|
||||
|
||||
# ---- slice with stride ----
|
||||
def test_slice_stride(self): _check(self, _t(6), lambda x: x[::2])
|
||||
def test_slice_start_stop_stride(self): _check(self, _t(6), lambda x: x[1:5:2])
|
||||
def test_slice_reverse(self): _check(self, _t(6), lambda x: x[::-1])
|
||||
def test_slice_singleton_negative_step(self): _check(self, _t(8), lambda x: x[3:2:-1])
|
||||
|
||||
# ---- empty / out-of-bounds slice ----
|
||||
def test_slice_empty(self): _check(self, _t(6), lambda x: x[3:1])
|
||||
def test_slice_oob_stop(self): _check(self, _t(6), lambda x: x[0:100])
|
||||
|
||||
# ---- single int (reduces a dim) ----
|
||||
def test_int_positive(self): _check(self, _t(8), lambda x: x[3])
|
||||
def test_int_negative(self): _check(self, _t(8), lambda x: x[-1])
|
||||
|
||||
# ---- ellipsis ----
|
||||
def test_ellipsis_only(self): _check(self, _t(2, 3, 4), lambda x: x[...])
|
||||
def test_ellipsis_then_int(self): _check(self, _t(2, 3, 4), lambda x: x[..., -1])
|
||||
def test_ellipsis_then_slice(self): _check(self, _t(2, 3, 4), lambda x: x[..., 1:3])
|
||||
def test_ellipsis_then_none(self): _check(self, _t(2, 3), lambda x: x[..., None])
|
||||
|
||||
# ---- None (unsqueeze) ----
|
||||
def test_none_front(self): _check(self, _t(4), lambda x: x[None])
|
||||
def test_none_back(self): _check(self, _t(4), lambda x: x[:, None])
|
||||
def test_none_middle(self): _check(self, _t(2, 3), lambda x: x[:, None, :])
|
||||
def test_multiple_none(self): _check(self, _t(2, 3), lambda x: x[None, :, None])
|
||||
|
||||
# ---- mixed multi-dim ----
|
||||
def test_int_then_slice(self): _check(self, _t(2, 3), lambda x: x[1, :])
|
||||
def test_multi_int(self): _check(self, _t(2, 3, 4), lambda x: x[1, 2])
|
||||
def test_mixed_slice_int(self): _check(self, _t(2, 3, 4), lambda x: x[0:2, -1, 1:3])
|
||||
def test_mixed_slice_slice(self): _check(self, _t(3, 4, 5), lambda x: x[1:3, :, 0:2])
|
||||
def test_high_rank_combo(self): _check(self, _t(4, 5, 6), lambda x: x[1:3, :, -1, None])
|
||||
|
||||
# ---- advanced indexing: UOp index on UOp must match Tensor index on Tensor (same uop) ----
|
||||
def _check_adv(self, t, idx):
|
||||
# idx is a Tensor or a tuple mixing Tensor index arrays with slice/int/None/Ellipsis
|
||||
ui = idx.uop if isinstance(idx, Tensor) else (tuple(i.uop if isinstance(i,Tensor) else i for i in idx) if isinstance(idx, tuple) else idx)
|
||||
self.assertIs(t[idx].uop, t.uop[ui])
|
||||
|
||||
def test_adv_single(self): self._check_adv(_t(5), _ti([2,1,0,1,2]))
|
||||
def test_adv_negative(self): self._check_adv(_t(5), _ti([-1,-2,0]))
|
||||
def test_adv_out_of_bounds(self): self._check_adv(_t(5), _ti([4,7,2])) # oob -> 0
|
||||
def test_adv_2d_dim0(self): self._check_adv(_t(3,4), _ti([2,0,1]))
|
||||
def test_adv_2d_index_array(self): self._check_adv(_t(4,5), _ti([[0,1],[2,3]])) # shaped index
|
||||
def test_adv_two_consecutive(self): self._check_adv(_t(3,4), (_ti([2,0,1]), _ti([1,2,3]))) # linear path
|
||||
def test_adv_two_broadcast(self): self._check_adv(_t(3,4), (_ti([2,0,1]), _ti([[1],[2],[3]])))
|
||||
def test_adv_three_consecutive(self):self._check_adv(_t(3,4,5), (_ti([2,0]), _ti([1,3]), _ti([4,2])))
|
||||
def test_adv_after_slice(self): self._check_adv(_t(2,3,4), (slice(None), _ti([2,0,1])))
|
||||
def test_adv_non_consec_permute(self): self._check_adv(_t(2,3,4), (_ti([1,0]), slice(None), _ti([2,0])))
|
||||
def test_adv_idx_then_int(self): self._check_adv(_t(3,4), (_ti([2,0,1]), 2))
|
||||
def test_adv_int_then_idx(self): self._check_adv(_t(3,4), (1, _ti([2,0,1])))
|
||||
def test_adv_idx_then_none(self): self._check_adv(_t(3,4), (_ti([2,0,1]), None))
|
||||
def test_adv_none_then_idx(self): self._check_adv(_t(3,4), (None, _ti([2,0,1])))
|
||||
def test_adv_ellipsis_then_idx(self):self._check_adv(_t(2,3,4), (Ellipsis, _ti([2,0,1])))
|
||||
def test_adv_idx_slice_mix(self): self._check_adv(_t(4,5,6), (_ti([1,3]), slice(1,4), _ti([2,0])))
|
||||
|
||||
# bool index is unsupported
|
||||
def test_adv_bool_index_rejected(self):
|
||||
with self.assertRaises(IndexError): _t(5)[_t(5) > 2]
|
||||
with self.assertRaises(IndexError): _t(5).uop[(_t(5) > 2).uop]
|
||||
|
||||
# python list/tuple indices
|
||||
def test_adv_python_list(self):
|
||||
self.assertIs(_strip_unique(_t(5)[[2,1,0]].uop), _strip_unique(_t(5).uop[[2,1,0]]))
|
||||
def test_adv_python_list_negative(self):
|
||||
self.assertIs(_strip_unique(_t(5)[[-1,-2,0]].uop), _strip_unique(_t(5).uop[[-1,-2,0]]))
|
||||
def test_adv_python_list_nested(self):
|
||||
self.assertIs(_strip_unique(_t(3,4)[[0,1],[2,0]].uop), _strip_unique(_t(3,4).uop[[0,1],[2,0]]))
|
||||
def test_adv_python_list_2d_index(self):
|
||||
self.assertIs(_strip_unique(_t(4,5)[[[0,1],[2,3]]].uop), _strip_unique(_t(4,5).uop[[[0,1],[2,3]]]))
|
||||
|
||||
class TestTensorUOpCumalu(unittest.TestCase):
|
||||
def test_cumsum_1d(self): _check(self, _t(5), lambda x: x.cumsum())
|
||||
def test_cumsum_2d(self): _check(self, _t(3, 4), lambda x: x.cumsum(1))
|
||||
def test_cumsum_non_last(self): _check(self, _t(3, 4), lambda x: x.cumsum(0))
|
||||
def test_cumsum_large(self): _check(self, _t(600), lambda x: x.cumsum()) # exercises _split_cumalu
|
||||
def test_cumprod(self): _check(self, _t(4), lambda x: x.cumprod(0))
|
||||
|
||||
class TestTensorUOpCumMinMax(unittest.TestCase):
|
||||
def _check_pair(self, t, fn):
|
||||
vt, it = fn(t)
|
||||
vu, iu = fn(t.uop)
|
||||
self.assertIs(vt.uop, vu)
|
||||
self.assertIs(it.uop, iu)
|
||||
def test_cummax_1d(self): self._check_pair(_t(5), lambda x: x.cummax(0))
|
||||
def test_cummax_2d(self): self._check_pair(_t(3, 4), lambda x: x.cummax(1))
|
||||
def test_cummax_0d(self): self._check_pair(_t(1).reshape(()), lambda x: x.cummax(0))
|
||||
def test_cummin_1d(self): self._check_pair(_t(5), lambda x: x.cummin(0))
|
||||
def test_cummin_2d(self): self._check_pair(_t(3, 4), lambda x: x.cummin(1))
|
||||
|
||||
class TestTensorUOpArgMinMax(unittest.TestCase):
|
||||
def _check(self, t, fn): self.assertIs(fn(t).uop, fn(t.uop))
|
||||
def test_argmax(self): self._check(_t(3, 4), lambda x: x.argmax(axis=1))
|
||||
def test_argmax_flat(self): self._check(_t(3, 4), lambda x: x.argmax())
|
||||
def test_argmin(self): self._check(_t(3, 4), lambda x: x.argmin(axis=0))
|
||||
|
||||
class TestTensorUOpSequential(unittest.TestCase):
|
||||
def test_sequential(self): _check(self, _t(4), lambda x: x.sequential([lambda y: y * 2, lambda y: y + 1]))
|
||||
|
||||
class TestTensorUOpOneHot(unittest.TestCase):
|
||||
def test_one_hot(self):
|
||||
t = _t(5)
|
||||
self.assertIs(t.one_hot(5).uop, t.uop.one_hot(5))
|
||||
|
||||
class TestTensorUOpSort(unittest.TestCase):
|
||||
def _check(self, t, **kw):
|
||||
tv, ti = t.sort(**kw)
|
||||
uv, ui = t.uop.sort(**kw)
|
||||
self.assertIs(tv.uop, uv)
|
||||
self.assertIs(ti.uop, ui)
|
||||
def test_sort_1d(self): self._check(Tensor([0.5, 0.1, 0.3]).float())
|
||||
def test_sort_descending(self): self._check(Tensor([0.5, 0.1, 0.3]).float(), descending=True)
|
||||
def test_sort_2d(self): self._check(_t(2, 4).float())
|
||||
def test_sort_single(self): self._check(Tensor([1.0]).float())
|
||||
def test_argsort(self):
|
||||
t = Tensor([0.5, 0.1, 0.3]).float()
|
||||
self.assertIs(t.argsort().uop, t.uop.argsort())
|
||||
def test_topk(self):
|
||||
t = _t(2, 4).float()
|
||||
tv, ti = t.topk(2)
|
||||
uv, ui = t.uop.topk(2)
|
||||
self.assertIs(tv.uop, uv)
|
||||
self.assertIs(ti.uop, ui)
|
||||
|
||||
class TestTensorUOpAllclose(unittest.TestCase):
|
||||
def test_allclose(self):
|
||||
a, b = _t(4).float(), _t(4).float()
|
||||
self.assertIs(a.allclose(b).uop, a.uop.allclose(b.uop))
|
||||
|
||||
class TestTensorUOpCast(unittest.TestCase):
|
||||
def test_cast_str_dtype(self):
|
||||
t = _t(4)
|
||||
self.assertIs(t.cast("float32").uop, t.uop.cast("float32"))
|
||||
self.assertIs(t.uop.cast("float32").dtype, dtypes.float32)
|
||||
|
||||
class TestTensorUOpBitcast(unittest.TestCase):
|
||||
def test_bitcast_same_dtype(self): _check(self, _t(4).float(), lambda x: x.bitcast(dtypes.float32))
|
||||
def test_bitcast_str_dtype(self):
|
||||
t = _t(4)
|
||||
self.assertIs(t.bitcast("uint32").uop, t.uop.bitcast("uint32"))
|
||||
self.assertIs(t.uop.bitcast("uint32").dtype, dtypes.uint32)
|
||||
def test_bitcast_same_and_diff_size(self):
|
||||
_check(self, _t(4).float(), lambda x: x.bitcast(dtypes.uint32)) # same size
|
||||
_check(self, _t(4).cast(dtypes.uint8), lambda x: x.bitcast(dtypes.uint16)) # widen: uint8[4] -> uint16[2]
|
||||
_check(self, _t(4).cast(dtypes.uint16), lambda x: x.bitcast(dtypes.uint8)) # narrow: uint16[4] -> uint8[8]
|
||||
|
||||
class TestTensorUOpRand(unittest.TestCase):
|
||||
def test_random_bits(self):
|
||||
k = UOp.empty((2,), dtype=dtypes.uint32)
|
||||
c = UOp.zeros(2, dtype=dtypes.uint32)
|
||||
for num in (1, 4, 7, 1024):
|
||||
self.assertIs(Tensor.random_bits(Tensor(k), Tensor(c), num).uop, UOp.random_bits(k, c, num))
|
||||
def test_bits_to_rand_float32(self):
|
||||
bits_uop = UOp.empty((8,), dtype=dtypes.uint32)
|
||||
for shape in ((8,), (2, 4), (5,)):
|
||||
self.assertIs(Tensor._bits_to_rand(Tensor(bits_uop), shape, dtypes.float32).uop, UOp._bits_to_rand(bits_uop, shape, dtypes.float32))
|
||||
def test_threefry(self):
|
||||
t = _t(4).cast(dtypes.uint64)
|
||||
self.assertIs(t.threefry(t).uop, t.uop.threefry(t.uop))
|
||||
def test_threefry_random_bits(self):
|
||||
key, c0, c1 = UOp.empty((2,), dtype=dtypes.uint32), UOp.arange(4, dtype=dtypes.uint32), UOp.arange(4, dtype=dtypes.uint32)
|
||||
self.assertIs(Tensor._threefry_random_bits(Tensor(key), Tensor(c0), Tensor(c1)).uop, UOp._threefry_random_bits(key, c0, c1))
|
||||
def test_rand(self):
|
||||
k, c = UOp.empty((2,), dtype=dtypes.uint32), UOp.zeros(2, dtype=dtypes.uint32)
|
||||
self.assertIs(Tensor._rand(Tensor(k), Tensor(c), (2, 2), dtypes.float32).uop, UOp._rand(k, c, (2, 2), dtypes.float32))
|
||||
self.assertIs(Tensor._rand(Tensor(k), Tensor(c), (0, 3), dtypes.float32).uop, UOp._rand(k, c, (0, 3), dtypes.float32))
|
||||
|
||||
class TestTensorUOpGather(unittest.TestCase):
|
||||
def _check(self, t, dim, idx):
|
||||
self.assertIs(t.gather(dim, idx).uop, t.uop.gather(dim, idx.uop))
|
||||
def test_gather_1d(self): self._check(_t(5), 0, Tensor([2, 1, 0, 1, 2], dtype=dtypes.int32))
|
||||
def test_gather_dim0(self): self._check(_t(3, 4), 0, Tensor([[0, 1, 2, 0], [1, 2, 0, 1], [2, 0, 1, 2]], dtype=dtypes.int32))
|
||||
def test_gather_dim1(self): self._check(_t(3, 4), 1, Tensor([[0, 1, 2, 3], [1, 2, 3, 0], [2, 3, 0, 1]], dtype=dtypes.int32))
|
||||
|
||||
class TestTensorUOpInterpolate(unittest.TestCase):
|
||||
def _check(self, t, mode):
|
||||
self.assertIs(t.interpolate(size=(2, 2), mode=mode).uop, t.uop.interpolate(size=(2, 2), mode=mode))
|
||||
def test_interpolate_nearest(self): self._check(_t(1, 1, 4, 4).float(), "nearest")
|
||||
def test_interpolate_nearest_exact(self): self._check(_t(1, 1, 4, 4).float(), "nearest-exact")
|
||||
def test_interpolate_linear(self): self._check(_t(1, 1, 4, 4).float(), "linear")
|
||||
|
||||
class TestTensorUOpLoss(unittest.TestCase):
|
||||
def test_cross_entropy(self):
|
||||
t, Y = _t(2, 3).float(), Tensor([1, 2], dtype=dtypes.int32)
|
||||
self.assertIs(t.cross_entropy(Y).uop, t.uop.cross_entropy(Y.uop))
|
||||
def test_sparse_categorical_crossentropy(self):
|
||||
t, Y = _t(2, 3).float(), Tensor([1, 2], dtype=dtypes.int32)
|
||||
self.assertIs(t.sparse_categorical_crossentropy(Y).uop, t.uop.sparse_categorical_crossentropy(Y.uop))
|
||||
def test_sparse_categorical_crossentropy_ignore_index(self):
|
||||
t, Y = _t(2, 3).float(), Tensor([1, 2], dtype=dtypes.int32)
|
||||
self.assertIs(t.sparse_categorical_crossentropy(Y, ignore_index=0).uop, t.uop.sparse_categorical_crossentropy(Y.uop, ignore_index=0))
|
||||
def test_nll_loss(self):
|
||||
t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32)
|
||||
self.assertIs(t.nll_loss(Y).uop, t.uop.nll_loss(Y.uop))
|
||||
def test_nll_loss_weight(self):
|
||||
t, Y, w = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32), _t(3).float()
|
||||
self.assertIs(t.nll_loss(Y, weight=w).uop, t.uop.nll_loss(Y.uop, weight=w.uop))
|
||||
def test_nll_loss_ignore_index(self):
|
||||
t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32)
|
||||
self.assertIs(t.nll_loss(Y, ignore_index=1).uop, t.uop.nll_loss(Y.uop, ignore_index=1))
|
||||
def test_nll_loss_none_reduction(self):
|
||||
t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32)
|
||||
self.assertIs(t.nll_loss(Y, reduction="none").uop, t.uop.nll_loss(Y.uop, reduction="none"))
|
||||
def test_nll_loss_weight_ignore_index(self):
|
||||
t, Y, w = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32), _t(3).float()
|
||||
self.assertIs(t.nll_loss(Y, weight=w, ignore_index=1).uop, t.uop.nll_loss(Y.uop, weight=w.uop, ignore_index=1))
|
||||
|
||||
class TestTensorUOpScatter(unittest.TestCase):
|
||||
def test_scatter(self):
|
||||
x, idx, src = _t(3, 4).float(), Tensor([[0, 1, 2, 0]], dtype=dtypes.int32), _t(1, 4).float()
|
||||
self.assertIs(x.scatter(0, idx, src).uop, x.uop.scatter(0, idx.uop, src.uop))
|
||||
def test_scatter_scalar_src(self):
|
||||
x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32)
|
||||
self.assertIs(x.scatter(1, idx, 3.14).uop, x.uop.scatter(1, idx.uop, 3.14))
|
||||
# inf cannot be cast to int — this regresses if scalar src is routed through index.dtype first
|
||||
def test_scatter_inf_src(self):
|
||||
x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32)
|
||||
self.assertIs(x.scatter(1, idx, float("inf")).uop, x.uop.scatter(1, idx.uop, float("inf")))
|
||||
def test_scatter_add(self):
|
||||
x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32)
|
||||
self.assertIs(x.scatter(1, idx, 3.14, reduce="add").uop, x.uop.scatter(1, idx.uop, 3.14, reduce="add"))
|
||||
def test_scatter_multiply(self):
|
||||
x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32)
|
||||
self.assertIs(x.scatter(1, idx, 3.14, reduce="multiply").uop, x.uop.scatter(1, idx.uop, 3.14, reduce="multiply"))
|
||||
# tensor src with reduce hits the "elif reduce: raise" branch in both Tensor and UOp paths
|
||||
def test_scatter_tensor_src_with_reduce_raises(self):
|
||||
x, idx, src = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32), _t(1, 2).float()
|
||||
with self.assertRaises(TypeError): x.scatter(1, idx, src, reduce="add")
|
||||
with self.assertRaises(TypeError): x.uop.scatter(1, idx.uop, src.uop, reduce="add")
|
||||
|
||||
class TestTensorUOpScatterReduce(unittest.TestCase):
|
||||
def _check(self, x, idx, src, **kw):
|
||||
self.assertIs(x.scatter_reduce(0, idx, src, **kw).uop, x.uop.scatter_reduce(0, idx.uop, src.uop, **kw))
|
||||
def test_sum(self): self._check(_t(3, 4).float(), Tensor([[0, 1, 0, 1]]*3, dtype=dtypes.int32), Tensor.ones(3, 4).float(), reduce="sum")
|
||||
def test_prod(self): self._check(_t(3, 4).float(), Tensor([[0, 1, 0, 1]]*3, dtype=dtypes.int32), Tensor.ones(3, 4).float(), reduce="prod")
|
||||
def test_mean(self): self._check(_t(3, 4).float(), Tensor([[0, 1, 0, 1]]*3, dtype=dtypes.int32), Tensor.ones(3, 4).float(), reduce="mean")
|
||||
def test_amax(self): self._check(_t(3, 4).float(), Tensor([[0, 1, 0, 1]]*3, dtype=dtypes.int32), Tensor.ones(3, 4).float(), reduce="amax")
|
||||
def test_amin(self): self._check(_t(3, 4).float(), Tensor([[0, 1, 0, 1]]*3, dtype=dtypes.int32), Tensor.ones(3, 4).float(), reduce="amin")
|
||||
def test_mean_exclude_self(self):
|
||||
self._check(_t(3, 4).float(), Tensor([[0, 1, 0, 1]]*3, dtype=dtypes.int32), Tensor.ones(3, 4).float(), reduce="mean", include_self=False)
|
||||
|
||||
class TestTensorUOpMaskedSelect(unittest.TestCase):
|
||||
# only the fixed-size path is pure
|
||||
def _check(self, t, mask, **kw):
|
||||
self.assertIs(t.masked_select(mask, **kw).uop, t.uop.masked_select(mask.uop, **kw))
|
||||
def test_masked_select_1d(self): self._check(_t(6), Tensor([True, False, True, False, True, False]), size=4)
|
||||
def test_masked_select_2d(self):
|
||||
self._check(_t(3, 3), Tensor([[True, False, True], [False, True, False], [False, False, True]]), size=6, fill_value=-1)
|
||||
|
||||
class TestTensorUOpNonzero(unittest.TestCase):
|
||||
def _check(self, t, **kw): self.assertIs(t.nonzero(**kw).uop, t.uop.nonzero(**kw))
|
||||
def test_nonzero_1d(self): self._check(_t(5), size=3)
|
||||
def test_nonzero_2d(self): self._check(_t(2, 3), size=4)
|
||||
|
||||
class TestTensorUOpPool(unittest.TestCase):
|
||||
def test_avg_pool2d(self): _check(self, _t(1, 1, 5, 5).float(), lambda x: x.avg_pool2d())
|
||||
def test_avg_pool2d_padding(self): _check(self, _t(1, 1, 5, 5).float(), lambda x: x.avg_pool2d(padding=1))
|
||||
def test_avg_pool2d_ceil(self): _check(self, _t(1, 1, 5, 5).float(), lambda x: x.avg_pool2d(ceil_mode=True))
|
||||
def test_avg_pool2d_no_count_pad(self): _check(self, _t(1, 1, 5, 5).float(), lambda x: x.avg_pool2d(padding=1, count_include_pad=False))
|
||||
def test_max_pool2d(self): _check(self, _t(1, 1, 5, 5).float(), lambda x: x.max_pool2d())
|
||||
def test_max_pool2d_padding(self): _check(self, _t(1, 1, 5, 5).float(), lambda x: x.max_pool2d(padding=1))
|
||||
def test_max_pool2d_ceil(self): _check(self, _t(1, 1, 5, 5).float(), lambda x: x.max_pool2d(ceil_mode=True))
|
||||
def test_max_pool2d_return_indices(self):
|
||||
t = _t(1, 1, 5, 5).float()
|
||||
vt, it = t.max_pool2d(return_indices=True)
|
||||
vu, iu = t.uop.max_pool2d(return_indices=True)
|
||||
self.assertIs(vt.uop, vu)
|
||||
self.assertIs(it.uop, iu)
|
||||
def test_max_unpool2d(self):
|
||||
t = _t(1, 1, 4, 4).float()
|
||||
out, idx = t.max_pool2d(return_indices=True)
|
||||
self.assertIs(out.max_unpool2d(idx).uop, out.uop.max_unpool2d(idx.uop))
|
||||
|
||||
class TestTensorUOpCat(unittest.TestCase):
|
||||
def test_cat_dim0(self): _check(self, _t(2, 3), lambda x: x.cat(x, dim=0))
|
||||
def test_cat_dim1(self): _check(self, _t(2, 3), lambda x: x.cat(x, dim=1))
|
||||
def test_cat_3tensors(self): _check(self, _t(2, 3), lambda x: x.cat(x, x, dim=0))
|
||||
def test_cat_neg_dim(self): _check(self, _t(2, 3, 4), lambda x: x.cat(x, dim=-1))
|
||||
|
||||
class TestTensorUOpPad(unittest.TestCase):
|
||||
def test_pad_flat(self): _check(self, _t(4, 5), lambda x: x.pad((1, 2, 0, 3)))
|
||||
def test_pad_flat_negative(self): _check(self, _t(4, 5), lambda x: x.pad((1, -1, 0, 2), value=-1.0))
|
||||
def test_pad_grouped_none(self): _check(self, _t(4, 5), lambda x: x.pad((None, (0, 3))))
|
||||
def test_pad_circular(self): _check(self, _t(4, 5), lambda x: x.pad(((1, 2), (0, 3)), mode="circular"))
|
||||
def test_pad_circular_zero_after(self):_check(self, _t(4, 5), lambda x: x.pad(((1, 0), (2, 0)), mode="circular"))
|
||||
def test_pad_reflect(self): _check(self, _t(4, 5), lambda x: x.pad(((1, 2), (0, 3)), mode="reflect"))
|
||||
def test_pad_reflect_negative(self): _check(self, _t(4, 5), lambda x: x.pad(((1, -1), (0, 2)), mode="reflect"))
|
||||
def test_pad_replicate(self): _check(self, _t(4, 5), lambda x: x.pad(((1, 2), (0, 3)), mode="replicate"))
|
||||
def test_pad_replicate_negative(self): _check(self, _t(4, 5), lambda x: x.pad(((1, -1), (0, 2)), mode="replicate"))
|
||||
|
||||
class TestTensorUOpStack(unittest.TestCase):
|
||||
def test_stack_dim0(self): _check(self, _t(2, 3), lambda x: x.stack(x, dim=0))
|
||||
def test_stack_dim1(self): _check(self, _t(2, 3), lambda x: x.stack(x, dim=1))
|
||||
def test_stack_3tensors(self): _check(self, _t(2, 3), lambda x: x.stack(x, x, dim=0))
|
||||
def test_stack_new_last(self): _check(self, _t(2, 3), lambda x: x.stack(x, dim=-1))
|
||||
def test_stack_mixed_dtype(self):
|
||||
w = _t(2, 3).float()
|
||||
_check(self, _t(2, 3), lambda x: x.stack(w if isinstance(x, Tensor) else w.uop))
|
||||
self.assertIs(_t(2, 3).uop.stack(w.uop).dtype, dtypes.float32)
|
||||
def test_stack_index_dtype(self):
|
||||
# index is outside the promotion lattice, equal dtypes bypass promotion
|
||||
self.assertEqual(UOp.const(1).stack(UOp.const(2)).shape, (2,))
|
||||
|
||||
class TestTensorUOpConv2d(unittest.TestCase):
|
||||
def test_conv2d_basic(self):
|
||||
w = _t(1, 1, 2, 2).float()
|
||||
_check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop))
|
||||
def test_conv2d_padded(self):
|
||||
w = _t(1, 1, 2, 2).float()
|
||||
_check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop, padding=1))
|
||||
def test_conv2d_negative_padding(self):
|
||||
w = _t(1, 1, 3, 3).float()
|
||||
_check(self, _t(1, 1, 5, 5).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop, padding=(-1,-1,-1,-1)))
|
||||
def test_conv2d_multichannel_bias(self):
|
||||
w, b = _t(4, 2, 3, 3).float(), _t(4).float()
|
||||
_check(self, _t(2, 2, 5, 5).float(), lambda x: x.conv2d(*(y if isinstance(x, Tensor) else y.uop for y in (w, b))))
|
||||
def test_conv2d_stride_dilation(self):
|
||||
w = _t(2, 2, 2, 2).float()
|
||||
_check(self, _t(1, 2, 6, 6).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop, stride=2, dilation=2))
|
||||
def test_conv2d_groups(self):
|
||||
w = _t(4, 1, 2, 2).float()
|
||||
_check(self, _t(1, 4, 4, 4).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop, groups=4))
|
||||
def test_conv2d_3d(self):
|
||||
w = _t(1, 1, 2, 2, 2).float()
|
||||
_check(self, _t(1, 1, 3, 3, 3).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop))
|
||||
def test_conv2d_winograd(self):
|
||||
w, a = _t(2, 2, 3, 3).float(), _t(1, 2, 6, 6).float()
|
||||
with Context(WINO=0): direct = a.conv2d(w).uop
|
||||
with Context(WINO=1):
|
||||
self.assertIsNot(a.conv2d(w).uop, direct)
|
||||
_check(self, a, lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop))
|
||||
def test_conv2d_image(self):
|
||||
w, a = _t(4, 4, 3, 3).float(), _t(1, 4, 8, 8).float()
|
||||
with Context(IMAGE=0): direct = a.conv2d(w).uop
|
||||
with Context(IMAGE=1):
|
||||
self.assertIsNot(a.conv2d(w).uop, direct)
|
||||
_check(self, a, lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop))
|
||||
def test_dot_image(self):
|
||||
y, a = _t(4, 3).float(), _t(2, 4).float()
|
||||
with Context(IMAGE=0): direct = a.dot(y).uop
|
||||
with Context(IMAGE=1):
|
||||
self.assertIsNot(a.dot(y).uop, direct)
|
||||
_check(self, a, lambda x: x.dot(y if isinstance(x, Tensor) else y.uop))
|
||||
def test_conv_transpose2d_basic(self):
|
||||
w = _t(1, 1, 2, 2).float()
|
||||
_check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv_transpose2d(w if isinstance(x, Tensor) else w.uop))
|
||||
def test_conv_transpose2d_stride(self):
|
||||
w = _t(1, 1, 2, 2).float()
|
||||
_check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv_transpose2d(w if isinstance(x, Tensor) else w.uop, stride=2))
|
||||
|
||||
class TestTensorUOpHashing(unittest.TestCase):
|
||||
def test_keccak_sha3_256(self): _check(self, _t(8).cast(dtypes.uint8), lambda x: x.keccak())
|
||||
def test_keccak_shake_128(self): _check(self, _t(8).cast(dtypes.uint8), lambda x: x.keccak("shake_128"))
|
||||
|
||||
class TestTensorUOpEinsum(unittest.TestCase):
|
||||
def test_einsum_dot(self): _check(self, _t(2, 3), lambda x: type(x).einsum("ij,ij->", x, x))
|
||||
def test_einsum_transpose(self): _check(self, _t(2, 3), lambda x: type(x).einsum("ij->ji", x))
|
||||
|
||||
class TestTensorUOpSoftmax(unittest.TestCase):
|
||||
def test_softmax_default(self): _check(self, _t(2, 3).float(), lambda x: x.softmax())
|
||||
def test_softmax_axis0(self): _check(self, _t(2, 3).float(), lambda x: x.softmax(axis=0))
|
||||
def test_log_softmax_default(self): _check(self, _t(2, 3).float(), lambda x: x.log_softmax())
|
||||
def test_log_softmax_axis0(self): _check(self, _t(2, 3).float(), lambda x: x.log_softmax(axis=0))
|
||||
|
||||
class TestTensorUOpQR(unittest.TestCase):
|
||||
def _check(self, t):
|
||||
qt, rt = t.qr()
|
||||
qu, ru = t.uop.qr()
|
||||
self.assertIs(qt.uop, qu)
|
||||
self.assertIs(rt.uop, ru)
|
||||
def test_qr_square(self): self._check(_t(3, 3).float())
|
||||
def test_qr_tall(self): self._check(_t(4, 3).float())
|
||||
def test_qr_wide(self): self._check(_t(3, 4).float())
|
||||
def test_qr_zero_col(self): self._check(Tensor([[0.0, 1.0], [0.0, 2.0]]))
|
||||
def test_qr_batched(self): self._check(_t(2, 3, 3).float())
|
||||
|
||||
class TestTensorUOpSVD(unittest.TestCase):
|
||||
def _check(self, t, **kw):
|
||||
ut, st, vt = t.svd(**kw)
|
||||
uu, su, vu = t.uop.svd(**kw)
|
||||
self.assertIs(ut.uop, uu)
|
||||
self.assertIs(st.uop, su)
|
||||
self.assertIs(vt.uop, vu)
|
||||
def test_svd_square(self): self._check(_t(2, 2).float())
|
||||
def test_svd_tall(self): self._check(_t(3, 2).float())
|
||||
def test_svd_wide(self): self._check(_t(2, 3).float())
|
||||
def test_svd_odd_num(self): self._check(_t(3, 3).float()) # exercises odd-num runoff path
|
||||
def test_svd_batched(self): self._check(_t(2, 2, 2).float())
|
||||
def test_svd_nonfull(self): self._check(_t(3, 2).float(), full_matrices=False)
|
||||
|
||||
class TestUOpEmpty(unittest.TestCase):
|
||||
def test_empty_dtype_string(self):
|
||||
self.assertEqual(UOp.empty((3, 4), dtype="float32").dtype, dtypes.float32)
|
||||
|
||||
def test_empty_like_dtype_override(self):
|
||||
u = Tensor.ones(3, 4).uop.empty_like(dtype=dtypes.int8)
|
||||
self.assertEqual((u.shape, u.dtype), ((3, 4), dtypes.int8))
|
||||
self.assertTrue(u.has_buffer_identity())
|
||||
|
||||
def test_empty_like_sharded_to_single_device(self):
|
||||
# regression: sharded source, override to single device must yield full logical shape with no axis
|
||||
t = Tensor.ones(8, 4).shard(("NULL:0", "NULL:1"), axis=0)
|
||||
for dev in ("NULL:2", ("NULL:2",)): # singleton tuple also canonicalizes to single device
|
||||
u = t.uop.empty_like(device=dev, dtype=dtypes.int32)
|
||||
self.assertEqual((u.shape, u.device, u.dtype, u.axis), ((8, 4), "NULL:2", dtypes.int32, None))
|
||||
self.assertTrue(u.has_buffer_identity())
|
||||
|
||||
def test_empty_direct_singleton_tuple_device(self):
|
||||
u = UOp.empty((4,), dtype=dtypes.float32, device=("NULL:0",))
|
||||
self.assertEqual((u.shape, u.device, u.axis), ((4,), "NULL", None))
|
||||
|
||||
class TestTensorUOpCreation(unittest.TestCase):
|
||||
def test_empty(self):
|
||||
self.assertIs(_strip_unique(Tensor.empty(2, 3).uop), _strip_unique(UOp.empty(2, 3)))
|
||||
def test_full(self):
|
||||
self.assertIs(_strip_unique(Tensor.full((2, 3), 42).uop), _strip_unique(UOp.full((2, 3), 42)))
|
||||
def test_full_kwargs(self):
|
||||
self.assertIs(_strip_unique(Tensor.full((2, 3), 42, dtype=dtypes.int8, device="NULL").uop),
|
||||
_strip_unique(UOp.full((2, 3), 42, dtype=dtypes.int8, device="NULL")))
|
||||
def test_full_symbolic_fill(self):
|
||||
t = Tensor.full((2, 3), UOp.variable("x", 1, 10).bind(5))
|
||||
self.assertEqual(t.shape, (2, 3))
|
||||
def test_zeros(self):
|
||||
self.assertIs(_strip_unique(Tensor.zeros(2, 3).uop), _strip_unique(UOp.zeros(2, 3)))
|
||||
def test_ones(self):
|
||||
self.assertIs(_strip_unique(Tensor.ones(2, 3).uop), _strip_unique(UOp.ones(2, 3)))
|
||||
def test_invalids(self):
|
||||
self.assertIs(_strip_unique(Tensor.invalids(2, 3, dtype=dtypes.int8).uop), _strip_unique(UOp.invalids((2, 3), dtype=dtypes.int8)))
|
||||
def test_empty_like(self):
|
||||
t = Tensor.empty(2, 3, dtype=dtypes.int8)
|
||||
self.assertIs(_strip_unique(t.empty_like().uop), _strip_unique(t.uop.empty_like()))
|
||||
self.assertIs(_strip_unique(t.empty_like(dtype=dtypes.float, device="NULL").uop), _strip_unique(t.uop.empty_like(dtypes.float, "NULL")))
|
||||
def test_arange(self):
|
||||
self.assertIs(Tensor.arange(5).uop, UOp.arange(5))
|
||||
def test_arange_empty(self):
|
||||
self.assertIs(Tensor.arange(5, 5).uop, UOp.arange(5, 5))
|
||||
def test_arange_step(self):
|
||||
self.assertIs(Tensor.arange(5, 10, 2).uop, UOp.arange(5, 10, 2))
|
||||
def test_linspace(self):
|
||||
self.assertIs(Tensor.linspace(0, 10, 5).uop, UOp.linspace(0, 10, 5))
|
||||
def test_linspace_one_step(self):
|
||||
self.assertIs(Tensor.linspace(5, 10, 1).uop, UOp.linspace(5, 10, 1))
|
||||
def test_eye(self):
|
||||
self.assertIs(Tensor.eye(3).uop, UOp.eye(3))
|
||||
def test_eye_rect(self):
|
||||
self.assertIs(Tensor.eye(2, 4).uop, UOp.eye(2, 4))
|
||||
def test_triu(self):
|
||||
t = _t(3, 4)
|
||||
self.assertIs(t.triu().uop, t.uop.triu())
|
||||
def test_triu_diagonal(self):
|
||||
t = _t(3, 4)
|
||||
self.assertIs(t.triu(diagonal=1).uop, t.uop.triu(diagonal=1))
|
||||
def test_tril(self):
|
||||
t = _t(3, 4)
|
||||
self.assertIs(t.tril().uop, t.uop.tril())
|
||||
def test_tril_diagonal(self):
|
||||
t = _t(3, 4)
|
||||
self.assertIs(t.tril(diagonal=-1).uop, t.uop.tril(diagonal=-1))
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,60 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor
|
||||
from tinygrad.uop.ops import UPat, Ops, UOp
|
||||
|
||||
# NOTE: unlike before base for a realized tensor is always a BUFFER
|
||||
realized_pattern = UPat(Ops.BUFFER)
|
||||
def is_pattern_uop(u:UOp, pat:UPat): assert pat.match(u, {}), f"{u}\nis not\n{pat}"
|
||||
def is_pattern(ten:Tensor, pat:UPat): is_pattern_uop(ten.uop, pat)
|
||||
|
||||
class TestTensorMutates(unittest.TestCase):
|
||||
@unittest.skip("this doesn't mutate anymore")
|
||||
def test_mutate_add(self):
|
||||
a = Tensor([1,2,3])
|
||||
b = Tensor([4,5,6])
|
||||
ret = a+b
|
||||
pa = a.uop
|
||||
pb = b.uop
|
||||
pr = ret.uop
|
||||
ret.schedule_linear()
|
||||
self.assertIsNot(pa, a.uop)
|
||||
self.assertIsNot(pb, b.uop)
|
||||
self.assertIsNot(pr, ret.uop)
|
||||
for t in [a,b,ret]: is_pattern_uop(t.uop.base, realized_pattern)
|
||||
|
||||
def test_reshape_is_same_parent(self):
|
||||
a = Tensor([1,2,3])
|
||||
b = Tensor([4,5,6])
|
||||
c = a+b
|
||||
d = (a+b).reshape(3,1)
|
||||
d.realize()
|
||||
is_pattern_uop(d.uop.base, realized_pattern)
|
||||
is_pattern_uop(c.uop.base, realized_pattern)
|
||||
is_pattern_uop(c.uop.base, realized_pattern)
|
||||
assert d.uop is not d.uop.base
|
||||
|
||||
def test_reshape_is_same_child(self):
|
||||
a = Tensor([1,2,3])
|
||||
b = Tensor([4,5,6])
|
||||
c = a+b
|
||||
d = (a+b).reshape(3,1)
|
||||
c.realize()
|
||||
is_pattern_uop(c.uop.base, realized_pattern)
|
||||
is_pattern_uop(d.uop.base, realized_pattern)
|
||||
|
||||
class TestTensorUopRepresentation(unittest.TestCase):
|
||||
def test_realized(self):
|
||||
a = Tensor([1.,2,3]).realize()
|
||||
print(a.uop)
|
||||
is_pattern_uop(a.uop.base, realized_pattern)
|
||||
|
||||
def test_add_realized(self):
|
||||
a = Tensor([1.,2,3]).realize()
|
||||
b = Tensor([4.,5,6]).realize()
|
||||
c = a+b
|
||||
print(c.uop)
|
||||
is_pattern(c, UPat(Ops.ADD))
|
||||
for s in c.uop.src: is_pattern_uop(s.base, realized_pattern)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
332
artifacts/package_sources/tinygrad/test/null/test_tqdm.py
Normal file
332
artifacts/package_sources/tinygrad/test/null/test_tqdm.py
Normal file
@@ -0,0 +1,332 @@
|
||||
import time, random, unittest, itertools
|
||||
from unittest.mock import patch
|
||||
from io import StringIO
|
||||
from collections import namedtuple
|
||||
from tqdm import tqdm
|
||||
from tinygrad.helpers import tqdm as tinytqdm, trange
|
||||
import numpy as np
|
||||
|
||||
def _get_iter_per_second(raw:str) -> float:
|
||||
# raw might have unit scale
|
||||
if raw.endswith("k"): return float(raw[:-1])*1e3
|
||||
if raw.endswith("M"): return float(raw[:-1])*1e6
|
||||
return float(raw)
|
||||
|
||||
# TODO: _get_iter_per_second in test_unit_scale might fail if lower bound is too small
|
||||
NCOLS_RANGE = [80, 240]
|
||||
|
||||
class TestProgressBar(unittest.TestCase):
|
||||
def _compare_bars(self, bar1, bar2):
|
||||
prefix1, prog1, suffix1 = bar1.split("|")
|
||||
prefix2, prog2, suffix2 = bar2.split("|")
|
||||
|
||||
self.assertEqual(len(bar1), len(bar2))
|
||||
self.assertEqual(prefix1, prefix2)
|
||||
|
||||
def parse_timer(timer): return sum(int(x) * y for x, y in zip(timer.split(':')[::-1], (1, 60, 3600)))
|
||||
|
||||
if "?" not in suffix1 and "?" not in suffix2:
|
||||
# allow for few sec diff in timers (removes flakiness)
|
||||
timer1, rm1 = [parse_timer(timer) for timer in suffix1.split("[")[-1].split(",")[0].split("<")]
|
||||
timer2, rm2 = [parse_timer(timer) for timer in suffix2.split("[")[-1].split(",")[0].split("<")]
|
||||
np.testing.assert_allclose(timer1, timer2, atol=5, rtol=1e-2)
|
||||
np.testing.assert_allclose(rm1, rm2, atol=5, rtol=1e-2)
|
||||
|
||||
# get suffix without timers
|
||||
suffix1 = suffix1.split("[")[0] + suffix1.split(",")[1]
|
||||
suffix2 = suffix2.split("[")[0] + suffix2.split(",")[1]
|
||||
self.assertEqual(suffix1, suffix2)
|
||||
else:
|
||||
self.assertEqual(suffix1, suffix2)
|
||||
|
||||
diff = sum([c1 != c2 for c1, c2 in zip(prog1, prog2)]) # allow 1 char diff to be less flaky, but it should match
|
||||
assert diff <= 1, f"{diff=}\n{prog1=}\n{prog2=}"
|
||||
|
||||
@patch('sys.stderr', new_callable=StringIO)
|
||||
@patch('shutil.get_terminal_size')
|
||||
def test_tqdm_output_iter(self, mock_terminal_size, mock_stderr):
|
||||
for _ in range(10):
|
||||
total, ncols = random.randint(5, 30), random.randint(*NCOLS_RANGE)
|
||||
mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
|
||||
mock_stderr.truncate(0)
|
||||
|
||||
# compare bars at each iteration (only when tinytqdm bar has been updated)
|
||||
for n in (bar := tinytqdm(range(total), desc="Test")):
|
||||
if bar.i % bar.skip != 0: continue
|
||||
tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
|
||||
iters_per_sec = float(tinytqdm_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0
|
||||
elapsed = n/iters_per_sec if n>0 else 0
|
||||
tqdm_output = tqdm.format_meter(n=n, total=total, elapsed=elapsed, ncols=ncols, prefix="Test")
|
||||
self._compare_bars(tinytqdm_output, tqdm_output)
|
||||
|
||||
# compare final bars
|
||||
tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
|
||||
iters_per_sec = float(tinytqdm_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0
|
||||
elapsed = total/iters_per_sec if n>0 else 0
|
||||
tqdm_output = tqdm.format_meter(n=total, total=total, elapsed=elapsed, ncols=ncols, prefix="Test")
|
||||
self._compare_bars(tinytqdm_output, tqdm_output)
|
||||
|
||||
@unittest.skip("this is flaky")
|
||||
@patch('sys.stderr', new_callable=StringIO)
|
||||
@patch('shutil.get_terminal_size')
|
||||
def test_unit_scale(self, mock_terminal_size, mock_stderr):
|
||||
for unit_scale in [True, False]:
|
||||
# NOTE: numpy comparison raises TypeError if exponent > 22
|
||||
for exponent in range(1, 22, 3):
|
||||
low, high = 10 ** exponent, 10 ** (exponent+1)
|
||||
for _ in range(5):
|
||||
total, ncols = random.randint(low, high), random.randint(*NCOLS_RANGE)
|
||||
mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
|
||||
mock_stderr.truncate(0)
|
||||
|
||||
# compare bars at each iteration (only when tinytqdm bar has been updated)
|
||||
# setting high rate to make sure it does not skip
|
||||
for n in tinytqdm(range(total), desc="Test", total=total, unit_scale=unit_scale, rate=10**9):
|
||||
tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
|
||||
|
||||
if n:
|
||||
iters_per_sec = _get_iter_per_second(tinytqdm_output.split("it/s")[-2].split(" ")[-1])
|
||||
elapsed = n/iters_per_sec
|
||||
else:
|
||||
elapsed = 0
|
||||
tqdm_output = tqdm.format_meter(n=n, total=total, elapsed=elapsed, ncols=ncols, prefix="Test", unit_scale=unit_scale)
|
||||
self._compare_bars(tinytqdm_output, tqdm_output)
|
||||
if n > 3: break
|
||||
|
||||
@patch('sys.stderr', new_callable=StringIO)
|
||||
@patch('shutil.get_terminal_size')
|
||||
def test_unit_scale_exact(self, mock_terminal_size, mock_stderr):
|
||||
unit_scale = True
|
||||
ncols = 80
|
||||
mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
|
||||
mock_stderr.truncate(0)
|
||||
|
||||
total = 10
|
||||
with patch('time.perf_counter', side_effect=[0]+list(range(100))): # one more 0 for the init call
|
||||
# compare bars at each iteration (only when tinytqdm bar has been updated)
|
||||
for n in tinytqdm(range(total), desc="Test", total=total, unit_scale=unit_scale, rate=10**9):
|
||||
tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
|
||||
elapsed = n
|
||||
tqdm_output = tqdm.format_meter(n=n, total=total, elapsed=elapsed, ncols=ncols, prefix="Test", unit_scale=unit_scale)
|
||||
self._compare_bars(tinytqdm_output, tqdm_output)
|
||||
if n > 5: break
|
||||
|
||||
total = 10
|
||||
k=0.001000001
|
||||
# regression test for
|
||||
# E AssertionError: ' 1.00/10.0 1000it/s]' != ' 1.00/10.0 1.00kit/s]'
|
||||
# E - 1.00/10.0 1000it/s]
|
||||
# E ? ^
|
||||
# E + 1.00/10.0 1.00kit/s]
|
||||
# E ? + ^
|
||||
with patch('time.perf_counter', side_effect=[0, *[i*k for i in range(100)]]): # one more 0 for the init call
|
||||
# compare bars at each iteration (only when tinytqdm bar has been updated)
|
||||
for n in tinytqdm(range(total), desc="Test", total=total, unit_scale=unit_scale, rate=10**9):
|
||||
tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
|
||||
elapsed = n*k
|
||||
tqdm_output = tqdm.format_meter(n=n, total=total, elapsed=elapsed, ncols=ncols, prefix="Test", unit_scale=unit_scale)
|
||||
self._compare_bars(tinytqdm_output, tqdm_output)
|
||||
if n > 5: break
|
||||
|
||||
@patch('sys.stderr', new_callable=StringIO)
|
||||
@patch('shutil.get_terminal_size')
|
||||
def test_si_boundary(self, mock_terminal_size, mock_stderr):
|
||||
"""Test SI formatting at boundaries (e.g., 999.5 -> 1.00k, not 1000)"""
|
||||
ncols = 80
|
||||
mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
|
||||
|
||||
# Test rates at the boundary: 999 stays as "999", 999.5+ becomes "1.00k"
|
||||
for rate in [999, 999.4, 999.5, 1000, 1001]:
|
||||
mock_stderr.truncate(0)
|
||||
mock_stderr.seek(0)
|
||||
elapsed = 1.0 / rate
|
||||
# Need 3 perf_counter calls: init st, init update, final update
|
||||
with patch('time.perf_counter', side_effect=[0, 0, elapsed]):
|
||||
bar = tinytqdm(desc="Test", total=1, unit_scale=True, rate=10**9)
|
||||
bar.update(1, close=True)
|
||||
tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
|
||||
tqdm_output = tqdm.format_meter(n=1, total=1, elapsed=elapsed, ncols=ncols, prefix="Test", unit_scale=True)
|
||||
self._compare_bars(tinytqdm_output, tqdm_output)
|
||||
|
||||
@unittest.skip("this is flaky")
|
||||
@patch('sys.stderr', new_callable=StringIO)
|
||||
@patch('shutil.get_terminal_size')
|
||||
def test_set_description(self, mock_terminal_size, mock_stderr):
|
||||
for _ in range(10):
|
||||
total, ncols = random.randint(5, 30), random.randint(*NCOLS_RANGE)
|
||||
mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
|
||||
mock_stderr.truncate(0)
|
||||
|
||||
expected_prefix = "Test"
|
||||
# compare bars at each iteration (only when tinytqdm bar has been updated)
|
||||
for i,n in enumerate(bar := tinytqdm(range(total), desc="Test")):
|
||||
if bar.i % bar.skip != 0: continue
|
||||
tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
|
||||
iters_per_sec = float(tinytqdm_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0
|
||||
elapsed = n/iters_per_sec if n>0 else 0
|
||||
tqdm_output = tqdm.format_meter(n=n, total=total, elapsed=elapsed, ncols=ncols, prefix=expected_prefix)
|
||||
expected_prefix = desc = f"Test {i}" if i % 2 == 0 else ""
|
||||
bar.set_description(desc)
|
||||
self._compare_bars(tinytqdm_output, tqdm_output)
|
||||
|
||||
# compare final bars
|
||||
tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
|
||||
iters_per_sec = float(tinytqdm_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0
|
||||
elapsed = total/iters_per_sec if n>0 else 0
|
||||
tqdm_output = tqdm.format_meter(n=total, total=total, elapsed=elapsed, ncols=ncols, prefix=expected_prefix)
|
||||
self._compare_bars(tinytqdm_output, tqdm_output)
|
||||
|
||||
@patch('sys.stderr', new_callable=StringIO)
|
||||
@patch('shutil.get_terminal_size')
|
||||
def test_trange_output_iter(self, mock_terminal_size, mock_stderr):
|
||||
for _ in range(5):
|
||||
total, ncols = random.randint(5, 30), random.randint(*NCOLS_RANGE)
|
||||
mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
|
||||
mock_stderr.truncate(0)
|
||||
|
||||
# compare bars at each iteration (only when tinytqdm bar has been updated)
|
||||
for n in (bar := trange(total, desc="Test")):
|
||||
if bar.i % bar.skip != 0: continue
|
||||
tiny_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
|
||||
iters_per_sec = float(tiny_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0
|
||||
elapsed = n/iters_per_sec if n>0 else 0
|
||||
tqdm_output = tqdm.format_meter(n=n, total=total, elapsed=elapsed, ncols=ncols, prefix="Test")
|
||||
self._compare_bars(tiny_output, tqdm_output)
|
||||
|
||||
# compare final bars
|
||||
tiny_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
|
||||
iters_per_sec = float(tiny_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0
|
||||
elapsed = total/iters_per_sec if n>0 else 0
|
||||
tqdm_output = tqdm.format_meter(n=total, total=total, elapsed=elapsed, ncols=ncols, prefix="Test")
|
||||
self._compare_bars(tiny_output, tqdm_output)
|
||||
|
||||
@patch('sys.stderr', new_callable=StringIO)
|
||||
@patch('shutil.get_terminal_size')
|
||||
def test_tqdm_output_custom(self, mock_terminal_size, mock_stderr):
|
||||
for _ in range(10):
|
||||
total, ncols = random.randint(10000, 1000000), random.randint(*NCOLS_RANGE)
|
||||
mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
|
||||
mock_stderr.truncate(0)
|
||||
|
||||
# compare bars at each iteration (only when tinytqdm bar has been updated)
|
||||
bar = tinytqdm(total=total, desc="Test")
|
||||
n = 0
|
||||
while n < total:
|
||||
incr = (total // 100) + random.randint(0, 1000)
|
||||
if n + incr > total: incr = total - n
|
||||
bar.update(incr, close=n+incr==total)
|
||||
n += incr
|
||||
if bar.i % bar.skip != 0: continue
|
||||
|
||||
tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
|
||||
iters_per_sec = float(tinytqdm_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0
|
||||
elapsed = n/iters_per_sec if n>0 else 0
|
||||
tqdm_output = tqdm.format_meter(n=n, total=total, elapsed=elapsed, ncols=ncols, prefix="Test")
|
||||
self._compare_bars(tinytqdm_output, tqdm_output)
|
||||
|
||||
@patch('sys.stderr', new_callable=StringIO)
|
||||
@patch('shutil.get_terminal_size')
|
||||
def test_tqdm_output_custom_0_total(self, mock_terminal_size, mock_stderr):
|
||||
for _ in range(10):
|
||||
total, ncols = random.randint(10000, 100000), random.randint(*NCOLS_RANGE)
|
||||
mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
|
||||
mock_stderr.truncate(0)
|
||||
|
||||
# compare bars at each iteration (only when tinytqdm bar has been updated)
|
||||
bar = tinytqdm(total=0, desc="Test")
|
||||
n = 0
|
||||
while n < total:
|
||||
incr = (total // 10) + random.randint(0, 100)
|
||||
if n + incr > total: incr = total - n
|
||||
bar.update(incr, close=n+incr==total)
|
||||
n += incr
|
||||
if bar.i % bar.skip != 0: continue
|
||||
|
||||
tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
|
||||
iters_per_sec = float(tinytqdm_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0
|
||||
elapsed = n/iters_per_sec if n>0 else 0
|
||||
tqdm_output = tqdm.format_meter(n=n, total=0, elapsed=elapsed, ncols=ncols, prefix="Test")
|
||||
self.assertEqual(tinytqdm_output, tqdm_output)
|
||||
|
||||
@patch('sys.stderr', new_callable=StringIO)
|
||||
@patch('shutil.get_terminal_size')
|
||||
def test_tqdm_output_custom_nolen_total(self, mock_terminal_size, mock_stderr):
|
||||
for unit_scale in [True, False]:
|
||||
for _ in range(5):
|
||||
gen = itertools.count(0)
|
||||
ncols = random.randint(*NCOLS_RANGE)
|
||||
mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
|
||||
mock_stderr.truncate(0)
|
||||
|
||||
# compare bars at each iteration (only when tinytqdm bar has been updated)
|
||||
# setting high rate to make sure it does not skip
|
||||
for n,g in enumerate(tinytqdm(gen, desc="Test", unit_scale=unit_scale, rate=10**9)):
|
||||
assert g == n
|
||||
tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
|
||||
if n:
|
||||
iters_per_sec = _get_iter_per_second(tinytqdm_output.split("it/s")[-2].split(" ")[-1])
|
||||
elapsed = n/iters_per_sec
|
||||
else:
|
||||
elapsed = 0
|
||||
tqdm_output = tqdm.format_meter(n=n, total=0, elapsed=elapsed, ncols=ncols, prefix="Test", unit_scale=unit_scale)
|
||||
self.assertEqual(tinytqdm_output, tqdm_output)
|
||||
if n > 5: break
|
||||
|
||||
@patch('sys.stderr', new_callable=StringIO)
|
||||
@patch('shutil.get_terminal_size')
|
||||
def test_tqdm_write(self, mock_terminal_size, mock_stderr):
|
||||
for _ in range(5):
|
||||
ncols, tqdm_fp = random.randint(*NCOLS_RANGE), StringIO()
|
||||
mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
|
||||
mock_stderr.truncate(0)
|
||||
tqdm_fp.truncate(0)
|
||||
for i in tinytqdm(range(10)):
|
||||
time.sleep(0.01)
|
||||
tinytqdm.write(str(i))
|
||||
tqdm.write(str(i), file=tqdm_fp)
|
||||
tinytqdm_out, tqdm_out = mock_stderr.getvalue(), tqdm_fp.getvalue()
|
||||
self.assertEqual(tinytqdm_out.split("\r\033[K")[-1], tqdm_out.split(f"{i-1}\n")[-1])
|
||||
self.assertEqual(tinytqdm_out, tinytqdm_out)
|
||||
|
||||
@patch('sys.stderr', new_callable=StringIO)
|
||||
@patch('shutil.get_terminal_size')
|
||||
def test_tqdm_context_manager(self, mock_terminal_size, mock_stderr):
|
||||
for _ in range(10):
|
||||
total, ncols = random.randint(5, 30), random.randint(*NCOLS_RANGE)
|
||||
mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
|
||||
mock_stderr.truncate(0)
|
||||
|
||||
with tinytqdm(desc="Test", total=total) as bar:
|
||||
for _ in range(total):
|
||||
bar.update(1)
|
||||
|
||||
tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
|
||||
iters_per_sec = float(tinytqdm_output.split("it/s")[-2].split(" ")[-1])
|
||||
elapsed = total/iters_per_sec
|
||||
tqdm_output = tqdm.format_meter(n=total, total=total, elapsed=elapsed, ncols=ncols, prefix="Test")
|
||||
self._compare_bars(tinytqdm_output, tqdm_output)
|
||||
|
||||
def test_tqdm_perf(self):
|
||||
st = time.perf_counter()
|
||||
for _ in tqdm(range(100)): pass
|
||||
tqdm_time = time.perf_counter() - st
|
||||
|
||||
st = time.perf_counter()
|
||||
for _ in tinytqdm(range(100)): pass
|
||||
tinytqdm_time = time.perf_counter() - st
|
||||
|
||||
assert tinytqdm_time < 5 * tqdm_time
|
||||
|
||||
def test_tqdm_perf_high_iter(self):
|
||||
st = time.perf_counter()
|
||||
for _ in tqdm(range(10^7)): pass
|
||||
tqdm_time = time.perf_counter() - st
|
||||
|
||||
st = time.perf_counter()
|
||||
for _ in tinytqdm(range(10^7)): pass
|
||||
tinytqdm_time = time.perf_counter() - st
|
||||
|
||||
assert tinytqdm_time < 20 * tqdm_time
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,74 @@
|
||||
import unittest, math
|
||||
import numpy as np
|
||||
from tinygrad import dtypes
|
||||
from tinygrad.uop.ops import UOp
|
||||
from tinygrad.codegen.decomp.transcendental import payne_hanek_reduction, cody_waite_reduction, frexp, rintk, pow2if
|
||||
from test.helpers import eval_uop
|
||||
|
||||
class TestTranscendentalFunctions(unittest.TestCase):
|
||||
def test_payne_hanek_reduction(self):
|
||||
# TODO: Test constant input when constant folding is fixed (or maybe test both variants)
|
||||
# Load input value from a buffer to prevent constant folding
|
||||
input_buf = UOp.param(1, dtypes.double, (1,))
|
||||
loaded_value = input_buf.index(UOp.const(0)).load()
|
||||
def eval_payne_hanek_reduction(v:float) -> tuple[float, int]:
|
||||
return tuple(eval_uop(u, [(dtypes.float64, [v])]) for u in payne_hanek_reduction(loaded_value))
|
||||
|
||||
r, q = eval_payne_hanek_reduction(12 * math.pi + 0.1)
|
||||
np.testing.assert_allclose(r, 0.1 - math.pi / 2)
|
||||
np.testing.assert_equal(q, 1)
|
||||
|
||||
r, q = eval_payne_hanek_reduction(12 * math.pi)
|
||||
np.testing.assert_allclose(r, 0.0, atol=1e-8)
|
||||
np.testing.assert_equal(q, 4)
|
||||
|
||||
r, q = eval_payne_hanek_reduction(12 * math.pi - 0.1)
|
||||
np.testing.assert_allclose(r, -0.1)
|
||||
np.testing.assert_equal(q, 4)
|
||||
|
||||
def test_cody_waite_reduction(self):
|
||||
r, q = (eval_uop(u) for u in cody_waite_reduction(UOp.const(12 * math.pi + 0.1).cast(dtypes.float64)))
|
||||
np.testing.assert_allclose(r, 0.1)
|
||||
np.testing.assert_equal(q, 12)
|
||||
|
||||
def test_frexp(self):
|
||||
for x in (1, -1):
|
||||
mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(float(x)).cast(dtypes.float64)))
|
||||
np.testing.assert_equal(mantissa, 0.5)
|
||||
np.testing.assert_equal(exponent, 1)
|
||||
|
||||
for x in (2, -2):
|
||||
mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(2.0).cast(dtypes.float64)))
|
||||
np.testing.assert_equal(mantissa, 0.5)
|
||||
np.testing.assert_equal(exponent, 2)
|
||||
|
||||
mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(5.0).cast(dtypes.float64)))
|
||||
np.testing.assert_equal(mantissa, 0.625)
|
||||
np.testing.assert_equal(exponent, 3)
|
||||
|
||||
mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(1000.0).cast(dtypes.float64)))
|
||||
np.testing.assert_allclose(mantissa, 0.9765625)
|
||||
np.testing.assert_equal(exponent, 10)
|
||||
|
||||
def test_rintk(self):
|
||||
np.testing.assert_allclose(eval_uop(rintk(UOp.const(0.0).cast(dtypes.float))), 0)
|
||||
np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.0).cast(dtypes.float))), 5)
|
||||
np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.5).cast(dtypes.float))), 6)
|
||||
np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.999).cast(dtypes.float))), 6)
|
||||
np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.0).cast(dtypes.float))), -5)
|
||||
np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.5).cast(dtypes.float))), -6)
|
||||
np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.999).cast(dtypes.float))), -6)
|
||||
|
||||
def test_pow2if(self):
|
||||
np.testing.assert_allclose(eval_uop(pow2if(UOp.const(0).cast(dtypes.int), dtypes.float)), 1.0)
|
||||
np.testing.assert_allclose(eval_uop(pow2if(UOp.const(1).cast(dtypes.int), dtypes.float)), 2.0)
|
||||
np.testing.assert_allclose(eval_uop(pow2if(UOp.const(2).cast(dtypes.int), dtypes.float)), 4.0)
|
||||
np.testing.assert_allclose(eval_uop(pow2if(UOp.const(10).cast(dtypes.int), dtypes.float)), 1024.0)
|
||||
np.testing.assert_allclose(eval_uop(pow2if(UOp.const(63).cast(dtypes.int), dtypes.float)), 2**63)
|
||||
np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-1).cast(dtypes.int), dtypes.float)), 0.5)
|
||||
np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-2).cast(dtypes.int), dtypes.float)), 0.25)
|
||||
np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-10).cast(dtypes.int), dtypes.float)), 2**-10)
|
||||
np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-63).cast(dtypes.int), dtypes.float)), 2**-63)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
740
artifacts/package_sources/tinygrad/test/null/test_uop_graph.py
Normal file
740
artifacts/package_sources/tinygrad/test/null/test_uop_graph.py
Normal file
@@ -0,0 +1,740 @@
|
||||
import unittest, pytest
|
||||
from tinygrad import dtypes, Variable, Device
|
||||
from tinygrad.dtype import AddrSpace
|
||||
from tinygrad.helpers import DEBUG, Context
|
||||
from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType, broadcast_axes, KernelInfo
|
||||
from tinygrad.uop.symbolic import sym
|
||||
from test.helpers import to_uops_list
|
||||
from tinygrad.codegen import full_rewrite_to_sink
|
||||
|
||||
simple_pm = PatternMatcher([
|
||||
(UPat.cvar('x', dtypes.weakint), lambda x: UOp.const(1.0) + UOp.const(2.0)),
|
||||
(UPat.cvar('x') + UPat.cvar('y'), lambda x,y: UOp.const(x.val+y.val)),
|
||||
(UPat.cvar('x') * UPat.cvar('y') * UPat.cvar('z'), lambda x,y,z: UOp.const(x.val*y.val*z.val)),
|
||||
((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.val+c2.val)),
|
||||
])
|
||||
|
||||
def const_values(u:UOp):
|
||||
if u.op is Ops.CONST: return (u.val,)
|
||||
if u.op is Ops.STACK: return tuple(x.val for x in u.src)
|
||||
raise AssertionError(f"expected const-like UOp, got {u.op}")
|
||||
|
||||
class TestGraphRewriteConst(unittest.TestCase):
|
||||
def test_gep_const(self):
|
||||
v1 = UOp.const((0,1,2), dtypes.int)
|
||||
v2 = v1.index(1)
|
||||
ret = graph_rewrite(v2, sym)
|
||||
self.assertEqual(ret.dtype, dtypes.int)
|
||||
self.assertEqual(ret.val, 1)
|
||||
|
||||
def test_add_const(self):
|
||||
v1 = UOp.const((0,1,2))
|
||||
v2 = UOp.const((5,6,7))
|
||||
ret = graph_rewrite(v1+v2, sym)
|
||||
self.assertEqual(ret.op, Ops.STACK)
|
||||
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):
|
||||
fn = pytest.mark.xfail(reason="const folding does not properly implement modular arithmetic")(fn)
|
||||
return unittest.expectedFailure(fn)
|
||||
class TestModularWraparound(unittest.TestCase):
|
||||
def _test(self, uop:UOp, expected:int):
|
||||
results = to_uops_list([uop])
|
||||
self.assertEqual(len(results), 2) # +1 for SINK
|
||||
self.assertEqual(results[0].op, Ops.CONST)
|
||||
self.assertEqual(results[0].dtype, uop.dtype)
|
||||
self.assertEqual(results[0].val, expected)
|
||||
|
||||
@xfail_broken_const_wraparound
|
||||
def test_cast(self):
|
||||
t = self._test
|
||||
t(UOp.const(0xABCD17D6, dtypes.uint).cast(dtypes.uint8), 0xD6)
|
||||
t(UOp.const(0xABCD17D6, dtypes.uint).cast(dtypes.uint8).cast(dtypes.uint), 0xD6)
|
||||
|
||||
@xfail_broken_const_wraparound
|
||||
def test_mul(self):
|
||||
t = self._test
|
||||
t(UOp.const(0xABCD17D6, dtypes.uint) * 0xAABBCCDD, 1147018174)
|
||||
t(UOp.const(0xABCD17D6, dtypes.int) * 10, -1241321892)
|
||||
|
||||
@xfail_broken_const_wraparound
|
||||
def test_div(self):
|
||||
t = self._test
|
||||
t(UOp.const(0xABCD17D6, dtypes.uint) * 0xAABBCCDD // 11, 104274379)
|
||||
t(UOp.const(0xABCD17D6, dtypes.int) * 10 // 11, -112847444)
|
||||
|
||||
@xfail_broken_const_wraparound
|
||||
def test_neg(self):
|
||||
t = self._test
|
||||
t(-UOp.const(1, dtypes.uint8), 0xFF)
|
||||
t(-UOp.const(1, dtypes.uint16), 0xFFFF)
|
||||
t(-UOp.const(1, dtypes.uint32), 0xFFFFFFFF)
|
||||
t(-UOp.const(1, dtypes.uint64), 0xFFFFFFFFFFFFFFFF)
|
||||
|
||||
@xfail_broken_const_wraparound
|
||||
def test_neg_min_int(self):
|
||||
t = self._test
|
||||
t(-UOp.const(-2**7, dtypes.int8), -2**7)
|
||||
t(-UOp.const(-2**15, dtypes.int16), -2**15)
|
||||
t(-UOp.const(-2**31, dtypes.int32), -2**31)
|
||||
t(-UOp.const(-2**63, dtypes.int64), -2**63)
|
||||
|
||||
@xfail_broken_const_wraparound
|
||||
def test_payne_hanek_reduction_bug(self):
|
||||
t = self._test
|
||||
a = (UOp.const(43748177600, dtypes.uint).cast(dtypes.uint) | 36).cast(dtypes.ulong)
|
||||
b = 2536655455 * a + 4294967296 * UOp.const(25366554550, dtypes.ulong)
|
||||
c = (b + 2261737165) // 4611686018427387904
|
||||
t(c, 0)
|
||||
|
||||
class TestGraphRewrite(unittest.TestCase):
|
||||
def test_dedup(self):
|
||||
v1 = UOp.variable("v", 0, 1, dtypes.float)
|
||||
v2 = UOp.variable("v", 0, 1, dtypes.float)
|
||||
nout = graph_rewrite(v1+v2, PatternMatcher([]))
|
||||
self.assertIs(nout.src[0], nout.src[1])
|
||||
|
||||
# NOTE: this shows why we can't have a UOp in arg
|
||||
@unittest.expectedFailure
|
||||
def test_no_dedup_args(self):
|
||||
a1 = UOp.variable("a1", UOp.const(0), UOp.const(11), dtypes.int)
|
||||
a2 = UOp.variable("a2", UOp.const(0), UOp.const(11), dtypes.int)
|
||||
sink = a1.sink(a2)
|
||||
variables = [x for x in graph_rewrite(sink, PatternMatcher([])).toposort() if x.op is Ops.PARAM and x.addrspace is AddrSpace.ALU]
|
||||
self.assertEqual(len(variables), 1)
|
||||
|
||||
def test_simple(self):
|
||||
c1 = UOp.const(1.0)
|
||||
c2 = UOp.const(2.0)
|
||||
nout = graph_rewrite(c1+c2, simple_pm)
|
||||
self.assertEqual(nout.op, Ops.CONST)
|
||||
self.assertEqual(nout.val, 3.0)
|
||||
|
||||
def test_depth_2_late(self):
|
||||
c1 = UOp.const(1.0)
|
||||
c2 = UOp.const(2.0)
|
||||
c3 = UOp.const(3.0)
|
||||
nout = graph_rewrite(c1*c2*(c3+c3), simple_pm)
|
||||
self.assertEqual(nout.op, Ops.CONST)
|
||||
self.assertEqual(nout.val, 12.0)
|
||||
|
||||
def test_double(self):
|
||||
c1 = UOp.const(1.0)
|
||||
c2 = UOp.const(2.0)
|
||||
c3 = UOp.const(3.0)
|
||||
nout = graph_rewrite(c1+c2+c3, simple_pm)
|
||||
self.assertEqual(nout.op, Ops.CONST)
|
||||
self.assertEqual(nout.val, 6.0)
|
||||
|
||||
def test_triple(self):
|
||||
c1 = UOp.const(1.0)
|
||||
c2 = UOp.const(2.0)
|
||||
c3 = UOp.const(3.0)
|
||||
c4 = UOp.const(4.0)
|
||||
nout = graph_rewrite(c1+c2+c3+c4, simple_pm)
|
||||
self.assertEqual(nout.op, Ops.CONST)
|
||||
self.assertEqual(nout.val, 10.0)
|
||||
|
||||
def test_diamond(self):
|
||||
c1 = UOp.const(1.0)
|
||||
c2 = UOp.const(2.0)
|
||||
c3 = UOp.const(3.0)
|
||||
nout = graph_rewrite((c1+c2)+(c1+c3), simple_pm)
|
||||
self.assertEqual(nout.op, Ops.CONST)
|
||||
self.assertEqual(nout.val, 7.0)
|
||||
|
||||
def test_magic_4(self):
|
||||
c1 = UOp.const(4)
|
||||
nout = graph_rewrite(c1, simple_pm)
|
||||
self.assertEqual(nout.op, Ops.CONST)
|
||||
self.assertEqual(nout.val, 3.0)
|
||||
|
||||
def test_depth_2_fold(self):
|
||||
v = UOp.variable("v", 0, 1, dtypes.float)
|
||||
c1 = UOp.const(1.0)
|
||||
c2 = UOp.const(2.0)
|
||||
nout = graph_rewrite(v+c1+c2, simple_pm)
|
||||
self.assertEqual(nout.op, Ops.ADD)
|
||||
self.assertEqual(nout.src[0].op, Ops.PARAM)
|
||||
self.assertEqual(nout.src[1].op, Ops.CONST)
|
||||
self.assertEqual(nout.src[1].val, 3.0)
|
||||
|
||||
def test_commutative_work(self):
|
||||
a = UOp.variable('a', 0, 1)
|
||||
b = UOp.variable('b', 0, 1)
|
||||
self.assertIs((a+b).simplify(), (b+a).simplify())
|
||||
|
||||
def test_consts_go_last_right_away(self):
|
||||
a = UOp.variable('a', 0, 1)
|
||||
tst = (2+a).simplify()
|
||||
self.assertIs(tst.src[0], a)
|
||||
self.assertIs(tst.src[1], UOp.const(2))
|
||||
|
||||
def test_consts_go_last(self):
|
||||
a = UOp.variable('a', 0, 1)
|
||||
b = UOp.variable('b', 0, 1)
|
||||
c = UOp.variable('c', 0, 1)
|
||||
d = UOp.variable('d', 0, 1)
|
||||
outs = [2+a, 2+a+d+3+b+c+4, UOp.const(2)+a, (4+d)+c+(2+a)+b]
|
||||
for out in outs:
|
||||
sink = graph_rewrite(out, sym)
|
||||
print(sink.render())
|
||||
self.assertEqual(sink.op, Ops.ADD)
|
||||
self.assertEqual(sink.src[1].op, Ops.CONST)
|
||||
self.assertEqual(len([x for x in sink.toposort() if x.op is Ops.CONST]), 1)
|
||||
|
||||
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
|
||||
uops = to_uops_list([out])
|
||||
self.assertEqual(len(uops), 2) # +1 for SINK
|
||||
out = uops[-2]
|
||||
self.assertEqual(out.op, Ops.CONST)
|
||||
self.assertEqual(out.val, 3.0)
|
||||
|
||||
def test_where_same_fold(self):
|
||||
v = UOp.variable('tmp', 0, 1)
|
||||
c0 = UOp.const(0)
|
||||
vc = v != c0
|
||||
c1 = UOp.const(1.0, dtypes.float)
|
||||
out = vc.where(c1, c1)
|
||||
uops = to_uops_list([out])
|
||||
self.assertEqual(len(uops), 2) # +1 for SINK
|
||||
out = uops[-2]
|
||||
self.assertEqual(out.op, Ops.CONST)
|
||||
self.assertEqual(out.val, 1.0)
|
||||
|
||||
def test_where_const_fold(self):
|
||||
bf = UOp.const(False)
|
||||
c1 = UOp.const(1.0, dtypes.float)
|
||||
c2 = UOp.const(2.0, dtypes.float)
|
||||
out = bf.where(c1, c2)
|
||||
uops = to_uops_list([out])
|
||||
self.assertEqual(len(uops), 2) # +1 for SINK
|
||||
out = uops[-2]
|
||||
self.assertEqual(out.op, Ops.CONST)
|
||||
self.assertEqual(out.val, 2.0)
|
||||
|
||||
def test_const_cast(self):
|
||||
bf = UOp.const(False)
|
||||
out = bf.cast(dtypes.int)
|
||||
uops = to_uops_list([out])
|
||||
self.assertEqual(len(uops), 2) # +1 for SINK
|
||||
out = uops[-2]
|
||||
self.assertEqual(out.op, Ops.CONST)
|
||||
self.assertEqual(out.val, 0)
|
||||
|
||||
def test_const_bitcast(self):
|
||||
bf = UOp.const(1.0, dtypes.float)
|
||||
out = bf.bitcast(dtypes.uint32)
|
||||
uops = to_uops_list([out])
|
||||
self.assertEqual(len(uops), 2) # +1 for SINK
|
||||
out = uops[-2]
|
||||
self.assertEqual(out.op, Ops.CONST)
|
||||
self.assertEqual(out.val, 0x3F800000)
|
||||
|
||||
@unittest.expectedFailure
|
||||
def test_const_shape_change_bitcast(self):
|
||||
bf = UOp.const(0x3F).cast(dtypes.uint8)
|
||||
out = bf.bitcast(dtypes.half)
|
||||
uops = to_uops_list([out])
|
||||
self.assertEqual(len(uops), 2) # +1 for SINK
|
||||
|
||||
def test_devectorize_derives_lane_dtype(self):
|
||||
from tinygrad.codegen import do_devectorize
|
||||
# an Invalid lane derives bool while the value lane derives float: the lane rebuild must derive, not inherit
|
||||
lhs = UOp.stack(UOp.invalid(), UOp.const(1.0).cast(dtypes.float))
|
||||
out = do_devectorize(lhs * lhs)
|
||||
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)
|
||||
|
||||
@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(Ops.STACK, dtypes.float, (ld,))
|
||||
x = vec.index(0)
|
||||
alu = UOp(Ops.SQRT, src=(x, ))
|
||||
out = UOp(Ops.STORE, src=(d0, idx, 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(Ops.STACK, dtypes.float, geps)
|
||||
out = d0.index(idx).store(vec)
|
||||
uops = to_uops_list([out])
|
||||
if DEBUG >= 4:
|
||||
from tinygrad import Device
|
||||
print(Device[Device.DEFAULT].renderer.render(uops))
|
||||
return uops[-2].src[-1] # -2 to skip SINK
|
||||
|
||||
# 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):
|
||||
for vec_size in [2, 4, 8]:
|
||||
consts = [UOp.const(float(i), dtypes.float) for i in range(vec_size)]
|
||||
vec = UOp(Ops.STACK, src=tuple(consts))
|
||||
with Context(SPEC=0):
|
||||
uops = to_uops_list([vec.index(i) for i in range(vec_size)])
|
||||
for uop, const in zip(uops, consts):
|
||||
self.assertEqual(uop, const)
|
||||
|
||||
def test_cast_alu_fold(self):
|
||||
d0 = UOp.param(0, dtypes.bool, (1,))
|
||||
d1 = UOp.param(1, dtypes.int, (1,))
|
||||
idx = UOp.const(0)
|
||||
ld = d1.index(idx)
|
||||
alu = (ld<1).cast(dtypes.bool)
|
||||
out = d0.index(idx).store(alu)
|
||||
uops = to_uops_list([out])
|
||||
self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)
|
||||
|
||||
def test_double_cast_fold(self):
|
||||
d0 = UOp.param(0, dtypes.float, (1,))
|
||||
d1 = UOp.param(1, dtypes.int, (1,))
|
||||
idx = UOp.const(0, dtypes.int)
|
||||
ld = d1.index(idx)
|
||||
alu = ld.cast(dtypes.float).cast(dtypes.float)
|
||||
out = d0.index(idx).store(alu)
|
||||
uops = to_uops_list([out])
|
||||
self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)
|
||||
|
||||
def test_depth_2_const_fold(self):
|
||||
v = UOp.variable("tmp", 0, 1, dtypes.int)
|
||||
c2 = UOp.const(2, dtypes.int)
|
||||
c4 = UOp.const(4, dtypes.int)
|
||||
vc = v+c2
|
||||
out = vc+c4
|
||||
uops = to_uops_list([out])
|
||||
self.assertEqual(len(uops), 5) # +1 for SINK, +1 for the PARAM shape STACK
|
||||
out = uops[-2] # -2 to skip SINK
|
||||
self.assertEqual(out.op, Ops.ADD)
|
||||
self.assertEqual(out.src[1].op, Ops.CONST)
|
||||
self.assertEqual(out.src[1].val, 6)
|
||||
|
||||
def test_bitcast_to_same_dtype_fold(self):
|
||||
for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
|
||||
d0 = UOp.param(0, dt, (1,))
|
||||
v = d0.index(UOp.const(0))
|
||||
uops = to_uops_list([v.bitcast(dt)])
|
||||
self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST and x.dtype is dt]), 0, f"dtype = {dt}")
|
||||
|
||||
def test_sub_with_cast_folds(self):
|
||||
a = Variable("a", 0, 5)
|
||||
uops = to_uops_list([a.cast(dtypes.int)+(-a).cast(dtypes.int)])
|
||||
assert uops[0] == UOp.const(0, dtypes.int)
|
||||
assert uops[-1].op == Ops.SINK
|
||||
|
||||
def test_where_on_gated_load_fold(self):
|
||||
ridx0 = UOp.range(100, 0)
|
||||
d0 = UOp.param(0, dtypes.long, (100,))
|
||||
ld = d0.index(ridx0.valid(ridx0<50))
|
||||
w = (ridx0<50).where(ld, 5)
|
||||
out = UOp.param(1, dtypes.long, (100,))
|
||||
uops = to_uops_list([out.index(ridx0).store(w)])
|
||||
for u in uops:
|
||||
assert u.op is not Ops.WHERE
|
||||
if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val==5
|
||||
|
||||
def test_where_on_gated_load_folds_swapped_branches(self):
|
||||
ridx0 = UOp.range(100, 0)
|
||||
d0 = UOp.param(0, dtypes.long, (100,))
|
||||
ld = d0.index(ridx0.valid((ridx0<50).logical_not()))
|
||||
w = (ridx0<50).where(5, ld)
|
||||
uops = to_uops_list([w])
|
||||
for u in uops:
|
||||
assert u.op is not Ops.WHERE
|
||||
if u.op is Ops.LOAD: assert u.src[1].val==5
|
||||
|
||||
def test_where_on_gated_load_with_cast(self):
|
||||
ridx0 = UOp.range(100, 0)
|
||||
d0 = UOp.param(0, dtypes.int, (100,))
|
||||
gate_idx = ridx0.valid((ridx0<50))
|
||||
ld = d0.index(gate_idx).cast(dtypes.float)
|
||||
w = (ridx0<50).where(ld, 5.0)
|
||||
out = UOp.param(1, dtypes.float, (100,))
|
||||
uops = to_uops_list([out.index(ridx0).store(w)])
|
||||
for u in uops:
|
||||
assert u.op is not Ops.WHERE
|
||||
if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val == 5
|
||||
|
||||
def test_where_on_casted_gated_load_extra_cond(self):
|
||||
ridx0 = UOp.range(100, 0)
|
||||
d0 = UOp.param(0, dtypes.float, (100,))
|
||||
ld = d0.index(ridx0.valid(ridx0<50))
|
||||
w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(0.0)).cast(dtypes.half)
|
||||
out = UOp.param(1, dtypes.half, (100,))
|
||||
uops = to_uops_list([out.index(ridx0).store(w)])
|
||||
for u in uops:
|
||||
assert u.op is not Ops.WHERE
|
||||
|
||||
def test_where_on_casted_gated_load_extra_cond_swapped(self):
|
||||
ridx0 = UOp.range(100, 0)
|
||||
d0 = UOp.param(0, dtypes.float, (100,))
|
||||
ld = d0.index(ridx0.valid(ridx0<50))
|
||||
w = ((ridx0<50) & (ridx0>30)).where(UOp.const(0.0), ld).cast(dtypes.half)
|
||||
out = UOp.param(1, dtypes.half, (100,))
|
||||
uops = to_uops_list([out.index(ridx0).store(w)])
|
||||
for u in uops:
|
||||
assert u.op is not Ops.WHERE
|
||||
|
||||
def test_where_in_store_becomes_gate(self):
|
||||
ridx0 = UOp.range(100, 0)
|
||||
d0 = UOp.param(0, dtypes.long, (100,))
|
||||
idx = d0.index(ridx0)
|
||||
ld = idx.load()
|
||||
val = (ridx0<50).where(5, ld)
|
||||
st = idx.store(val).end(ridx0)
|
||||
uops = to_uops_list([st])
|
||||
for u in uops:
|
||||
assert u.op is not Ops.WHERE
|
||||
if u.op is Ops.STORE: assert u.src[1].val==5
|
||||
|
||||
def test_load_idx_becomes_int(self):
|
||||
# mnist indexing with split reduceop
|
||||
# Make sure we are not doign math on the loaded index, which would promote it to long
|
||||
c0 = UOp.param(0, dtypes.uchar, (128000,))
|
||||
c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK)
|
||||
c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK)
|
||||
c3 = UOp.param(1, dtypes.int, (512,))
|
||||
c4 = c3.index(c1)
|
||||
c5 = UOp.range(UOp.const(240), 0, AxisType.REDUCE)
|
||||
c6 = ((c2*UOp.const(240))+c5)
|
||||
c7 = UOp.param(2, dtypes.uchar, (60000,))
|
||||
c8 = c7.index(c6)
|
||||
c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD)
|
||||
c10 = c0.index(((c1*UOp.const(250))+c2)).store(c9).end(c1, c2)
|
||||
uops = to_uops_list([c10])
|
||||
for u in uops:
|
||||
self.assertNotEqual(u.dtype, dtypes.long)
|
||||
|
||||
def test_load_idx_no_math_on_loaded(self):
|
||||
# test the (x+y)<c pattern where x has loads - we shouldn't do math on loaded indices
|
||||
c0 = UOp.param(0, dtypes.uchar, (128000,))
|
||||
c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK)
|
||||
c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK)
|
||||
c3 = UOp.param(1, dtypes.int, (512,))
|
||||
c4 = c3.index(c1) # c4 is a load
|
||||
c5 = UOp.range(UOp.const(240), 0, AxisType.REDUCE)
|
||||
c6 = ((c2*UOp.const(240))+c5)
|
||||
c7 = UOp.param(2, dtypes.uchar, (60000,))
|
||||
c8 = c7.index(c6)
|
||||
# (loaded + range) < const pattern - loaded value shouldn't be promoted to long
|
||||
loaded_idx = c4.cast(dtypes.weakint)
|
||||
comparison = (loaded_idx + c5) < UOp.const(60000)
|
||||
c9 = comparison.where(c8.cast(dtypes.uint).cast(dtypes.uchar), 0).reduce(c5, arg=Ops.ADD)
|
||||
c10 = c0.index(((c1*UOp.const(250))+c2)).store(c9).end(c1, c2)
|
||||
uops = to_uops_list([c10])
|
||||
for u in uops:
|
||||
self.assertNotEqual(u.dtype, dtypes.long)
|
||||
|
||||
def test_fold_gated_load(self):
|
||||
glbl0 = UOp.param(0, dtypes.int, (1,))
|
||||
glbl1 = UOp.param(1, dtypes.int, (1,))
|
||||
glbl2 = UOp.param(2, dtypes.int, (1,))
|
||||
idx = UOp.const(0)
|
||||
ld0 = glbl1.index(UOp.invalid())
|
||||
ld1 = glbl2.index(idx.valid(UOp.const(True)))
|
||||
uops = to_uops_list([glbl0.index(idx).store(ld1+ld0)])
|
||||
# the gate and invalid value are deleted from ld1
|
||||
self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 1)
|
||||
|
||||
def test_fold_gated_load_local(self):
|
||||
glbl0 = UOp.param(0, dtypes.int, (16,))
|
||||
smem = UOp.placeholder((18,), dtypes.int, slot=0, addrspace=AddrSpace.LOCAL)
|
||||
lidx = UOp.special(16, "lidx0")
|
||||
st = smem.index(lidx).store(glbl0.index(lidx).load())
|
||||
barrier = st.barrier()
|
||||
ld0 = smem.after(barrier).index(UOp.invalid())
|
||||
ld1 = smem.after(barrier).index((lidx+2).valid(UOp.const(True)))
|
||||
uops = to_uops_list([glbl0.index(lidx).store(ld1+ld0)])
|
||||
|
||||
# the gate and invalid value are deleted from ld1
|
||||
self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 2)
|
||||
|
||||
def test_fold_gated_store(self):
|
||||
glbl = UOp.param(0, dtypes.int, (1,))
|
||||
idx0 = UOp.const(0)
|
||||
val = UOp.const(42)
|
||||
st0 = glbl.index(UOp.invalid()).store(val)
|
||||
st1 = glbl.index(idx0.valid(UOp.const(True))).store(val)
|
||||
uops = to_uops_list([st0, st1])
|
||||
# only the second store happens
|
||||
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([UOp(Ops.STORE, src=(glbl0, idx, UOp.const(42), bad_gate))])
|
||||
|
||||
def test_after_end(self):
|
||||
r = UOp.range(10, 0)
|
||||
|
||||
c = r + 1
|
||||
self.assertIn(r, c.ranges)
|
||||
|
||||
e = UOp.const(1).end(r)
|
||||
self.assertNotIn(r, e.ranges)
|
||||
|
||||
a = c.after(e)
|
||||
self.assertNotIn(r, a.ranges)
|
||||
|
||||
class TestReduceCollapse(unittest.TestCase):
|
||||
def test_multi_range_reduce_add(self):
|
||||
"""Test that (x + y).reduce(r1, r2) distributes over multiple ranges"""
|
||||
from tinygrad.codegen.simplify import pm_reduce_collapse
|
||||
# Create two ranges
|
||||
r1 = UOp.range(3, 0)
|
||||
r2 = UOp.range(4, 1)
|
||||
# Create x + y where x and y depend on different ranges
|
||||
x = r1.cast(dtypes.float)
|
||||
y = r2.cast(dtypes.float)
|
||||
# (x + y).reduce(r1, r2) should be rewritten
|
||||
red = (x + y).reduce(r1, r2, arg=Ops.ADD)
|
||||
self.assertEqual(len(red.src), 3) # value + 2 ranges
|
||||
result = graph_rewrite(red, pm_reduce_collapse, name='test')
|
||||
# Should become add of two separate reduces
|
||||
self.assertEqual(result.op, Ops.ADD)
|
||||
|
||||
def test_reduce_shapeless_const_unroll(self):
|
||||
"""a REDUCE over a shapeless CONST (e.g. x*0 folded late in codegen) must collapse before the expander"""
|
||||
out = UOp.param(0, dtypes.float, (1,))
|
||||
red = UOp.const(3.0).cast(dtypes.float).reduce(UOp.range(4, 0, AxisType.UNROLL), arg=(Ops.ADD, 0))
|
||||
ast = UOp.sink(out.index(UOp.const(0)).store(red)).replace(arg=KernelInfo())
|
||||
uops = full_rewrite_to_sink(ast, Device["CPU"].renderer, optimize=False).toposort()
|
||||
self.assertNotIn(Ops.REDUCE, [u.op for u in uops])
|
||||
self.assertIn(12.0, [u.val for u in uops if u.op is Ops.CONST])
|
||||
|
||||
class TestMovementOps(unittest.TestCase):
|
||||
def test_pm_mops_partial_reshape_index_removes_reshape(self):
|
||||
from tinygrad.schedule.rangeify import pm_mops
|
||||
src = UOp.param(0, dtypes.float, shape=(32, 4))
|
||||
r0, r1 = UOp.range(4, 0), UOp.range(8, 1)
|
||||
result = graph_rewrite(src.reshape((4, 8, 4)).index(r0, r1), pm_mops, name="test")
|
||||
self.assertEqual(result.op, Ops.INDEX)
|
||||
self.assertIs(result.src[0], src)
|
||||
self.assertEqual(result.shape, (4,))
|
||||
self.assertNotIn(Ops.RESHAPE, [u.op for u in result.toposort()])
|
||||
|
||||
def test_pm_mops_partial_reshape_index_suffix_mismatch_does_nothing(self):
|
||||
from tinygrad.schedule.rangeify import pm_mops
|
||||
src = UOp.param(0, dtypes.float, shape=(2, 6))
|
||||
result = graph_rewrite(src.reshape((2, 3, 2)).index(UOp.range(2, 0)), pm_mops, name="test")
|
||||
self.assertEqual(result.op, Ops.INDEX)
|
||||
self.assertEqual(result.src[0].op, Ops.RESHAPE)
|
||||
|
||||
class TestConstBufferize(unittest.TestCase):
|
||||
def test_const_bufferize_with_ranges(self):
|
||||
"""Test that CONST.BUFFERIZE with ranges is folded correctly.
|
||||
|
||||
BUFFERIZE can have ranges as additional sources beyond the value.
|
||||
The pattern at rangeify.py uses allow_any_len=True because
|
||||
CONST doesn't depend on ranges (constant is same value everywhere).
|
||||
"""
|
||||
from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts
|
||||
c = UOp.const(42.0)
|
||||
r1 = UOp.range(3, 0)
|
||||
bufferize_with_range = UOp(Ops.STAGE, src=(c, r1), arg=BufferizeOpts(device="CPU"))
|
||||
self.assertEqual(len(bufferize_with_range.src), 2) # const + 1 range
|
||||
|
||||
result = graph_rewrite(bufferize_with_range, pm_const_buffer_folding, name='test')
|
||||
# BUFFERIZE should be removed, result is const broadcast to shape
|
||||
self.assertNotEqual(result.op, Ops.STAGE)
|
||||
const_vals = [u.val for u in result.toposort() if u.op is Ops.CONST and u.dtype is dtypes.weakfloat]
|
||||
self.assertIn(42.0, const_vals)
|
||||
|
||||
def test_const_bufferize_with_multiple_ranges(self):
|
||||
"""Test CONST.BUFFERIZE with multiple ranges is also folded."""
|
||||
from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts
|
||||
c = UOp.const(3.14)
|
||||
r1 = UOp.range(3, 0)
|
||||
r2 = UOp.range(4, 1)
|
||||
bufferize_with_ranges = UOp(Ops.STAGE, src=(c, r1, r2), arg=BufferizeOpts(device="CPU"))
|
||||
self.assertEqual(len(bufferize_with_ranges.src), 3) # const + 2 ranges
|
||||
|
||||
result = graph_rewrite(bufferize_with_ranges, pm_const_buffer_folding, name='test')
|
||||
# BUFFERIZE should be removed
|
||||
self.assertNotEqual(result.op, Ops.STAGE)
|
||||
const_vals = [u.val for u in result.toposort() if u.op is Ops.CONST and u.dtype is dtypes.weakfloat]
|
||||
self.assertIn(3.14, const_vals)
|
||||
|
||||
class TestUOpTags(unittest.TestCase):
|
||||
def test_inc_by_one(self):
|
||||
g = UOp.const(1) + UOp.const(1)
|
||||
assert g.ssimplify() == 2
|
||||
pm_plus_1 = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: UOp.const(x.val+1, x.dtype).rtag(1) if x.tag is None else None)])
|
||||
pm_strip_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
|
||||
g = graph_rewrite(g, pm_plus_1)
|
||||
assert g.ssimplify() == 4
|
||||
g = graph_rewrite(g, pm_plus_1)
|
||||
assert g.ssimplify() == 4
|
||||
g = graph_rewrite(g, pm_strip_tags)
|
||||
assert g.ssimplify() == 4
|
||||
g = graph_rewrite(g, pm_plus_1)
|
||||
assert g.ssimplify() == 6
|
||||
|
||||
class TestUOpGetItem(unittest.TestCase):
|
||||
def _placeholder(self, shape, dtype=dtypes.half):
|
||||
return UOp.placeholder(shape, dtype, slot=0, addrspace=AddrSpace.LOCAL)
|
||||
|
||||
# full slices (no shrink)
|
||||
def test_full_slice(self):
|
||||
p = self._placeholder((64, 64))
|
||||
self.assertEqual(p[:, :].shape, (64, 64))
|
||||
def test_full_slice_explicit(self):
|
||||
p = self._placeholder((64, 64))
|
||||
self.assertEqual(p[0:64, 0:64].shape, (64, 64))
|
||||
|
||||
# partial slices (shrink)
|
||||
def test_shrink_cols(self):
|
||||
p = self._placeholder((64, 80))
|
||||
self.assertEqual(p[:, :64].shape, (64, 64))
|
||||
def test_shrink_rows(self):
|
||||
p = self._placeholder((80, 64))
|
||||
self.assertEqual(p[:64, :].shape, (64, 64))
|
||||
def test_shrink_both(self):
|
||||
p = self._placeholder((80, 80))
|
||||
self.assertEqual(p[:64, :64].shape, (64, 64))
|
||||
def test_shrink_start(self):
|
||||
p = self._placeholder((64, 64))
|
||||
self.assertEqual(p[8:, :].shape, (56, 64))
|
||||
def test_shrink_start_and_end(self):
|
||||
p = self._placeholder((64, 64))
|
||||
self.assertEqual(p[8:56, 4:60].shape, (48, 56))
|
||||
|
||||
# mixed slice and index
|
||||
def test_index_and_slice(self):
|
||||
p = self._placeholder((64, 80))
|
||||
r = UOp.range(64, 100)
|
||||
result = p[r, :64]
|
||||
self.assertEqual(result.shape, (64,))
|
||||
def test_slice_and_index(self):
|
||||
p = self._placeholder((80, 64))
|
||||
r = UOp.range(64, 100)
|
||||
result = p[:64, r]
|
||||
self.assertEqual(result.shape, (64,))
|
||||
def test_shrink_then_index(self):
|
||||
p = self._placeholder((64, 80))
|
||||
s = p[:, :64]
|
||||
r = UOp.range(64, 100)
|
||||
result = s[r]
|
||||
self.assertEqual(result.shape, (64,))
|
||||
|
||||
# integer index (no slice)
|
||||
def test_int_index(self):
|
||||
p = self._placeholder((64, 64))
|
||||
result = p[0]
|
||||
self.assertEqual(result.shape, (64,))
|
||||
|
||||
# ellipsis
|
||||
def test_ellipsis_all_slices(self):
|
||||
p = self._placeholder((64, 80))
|
||||
self.assertEqual(p[..., :64].shape, (64, 64))
|
||||
def test_ellipsis_with_int(self):
|
||||
p = self._placeholder((64, 80))
|
||||
r = UOp.range(64, 100)
|
||||
result = p[..., r]
|
||||
self.assertEqual(result.op, Ops.INDEX)
|
||||
def test_ellipsis_only(self):
|
||||
p = self._placeholder((64, 64))
|
||||
self.assertEqual(p[...].shape, (64, 64))
|
||||
|
||||
# all slices should not create a bare INDEX
|
||||
def test_all_slices_no_index(self):
|
||||
p = self._placeholder((64, 80))
|
||||
result = p[:, :64]
|
||||
self.assertNotEqual(result.op, Ops.INDEX)
|
||||
def test_all_full_slices_no_index(self):
|
||||
p = self._placeholder((64, 64))
|
||||
result = p[:, :]
|
||||
self.assertNotEqual(result.op, Ops.INDEX)
|
||||
|
||||
class TestUOpBroadcast(unittest.TestCase):
|
||||
def test_broadcast_row(self):
|
||||
a = UOp.const(1.0).expand((4, 8))
|
||||
b = UOp.const(2.0).expand((4, 1))
|
||||
c = a + b
|
||||
self.assertEqual(c.shape, (4, 8))
|
||||
self.assertEqual(c.op, Ops.ADD)
|
||||
|
||||
def test_broadcast_col(self):
|
||||
a = UOp.const(1.0).expand((4, 8))
|
||||
b = UOp.const(2.0).expand((1, 8))
|
||||
c = a + b
|
||||
self.assertEqual(c.shape, (4, 8))
|
||||
self.assertEqual(c.op, Ops.ADD)
|
||||
|
||||
def test_broadcast_lower_dim(self):
|
||||
a = UOp.const(1.0).expand((4, 8))
|
||||
b = UOp.const(2.0).expand((8,))
|
||||
c = a * b
|
||||
self.assertEqual(c.shape, (4, 8))
|
||||
self.assertEqual(c.op, Ops.MUL)
|
||||
|
||||
def test_broadcast_scalar(self):
|
||||
a = UOp.const(1.0).expand((4, 8))
|
||||
c = a * 2
|
||||
self.assertEqual(c.shape, (4, 8))
|
||||
self.assertEqual(c.op, Ops.MUL)
|
||||
|
||||
def test_broadcast_symbolic_same_shape(self):
|
||||
t = Variable("t", 1, 10)
|
||||
a = UOp.const(1.0).expand((1, 1, t))
|
||||
b = UOp.const(2.0).expand((1, 1, t))
|
||||
c = a + b
|
||||
self.assertEqual(c.op, Ops.ADD)
|
||||
|
||||
def test_broadcast_axes(self):
|
||||
t = Variable("t", 1, 10)
|
||||
self.assertEqual(broadcast_axes((4, 8), (4, 8)), ())
|
||||
self.assertEqual(broadcast_axes((8,), (4, 8)), (0,))
|
||||
self.assertEqual(broadcast_axes((), (4, 8)), (0, 1))
|
||||
self.assertEqual(broadcast_axes((3, 1), (4, 3, 8)), (0, 2))
|
||||
self.assertEqual(broadcast_axes((1, 8), (1, 8)), ())
|
||||
self.assertEqual(broadcast_axes((t, 8), (t, 8)), ())
|
||||
self.assertEqual(broadcast_axes((1, 8), (t, 8)), (0,))
|
||||
with self.assertRaises(RuntimeError): broadcast_axes((4, 8), (8,))
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main(verbosity=2)
|
||||
@@ -0,0 +1,38 @@
|
||||
import unittest
|
||||
from tinygrad import UOp
|
||||
|
||||
class TestUOpRepr(unittest.TestCase):
|
||||
def test_simple_const(self):
|
||||
a = UOp.const(42)
|
||||
self.assertEqual(repr(a), "UOp(Ops.CONST, dtypes.weakint, arg=42, src=())")
|
||||
def test_different_consts(self):
|
||||
a, b = UOp.const(42), UOp.const(3)
|
||||
expected = (
|
||||
"UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
|
||||
" UOp(Ops.CONST, dtypes.weakint, arg=42, src=()),\n" +
|
||||
" UOp(Ops.CONST, dtypes.weakint, arg=3, src=()),))"
|
||||
)
|
||||
self.assertEqual(repr(a+b), expected)
|
||||
def test_walrus_operator_indentation(self):
|
||||
# The reference should have the same indentation as the definition
|
||||
a = UOp.const(42)
|
||||
expected = (
|
||||
"UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
|
||||
" x0:=UOp(Ops.CONST, dtypes.weakint, arg=42, src=()),\n" +
|
||||
" x0,))"
|
||||
)
|
||||
self.assertEqual(repr(a+a), expected)
|
||||
def test_nested_walrus_indentation(self):
|
||||
# Ensure indentation is consistent at multiple levels
|
||||
b = (a:=UOp.const(1)) + a
|
||||
expected = (
|
||||
"UOp(Ops.MUL, dtypes.weakint, arg=None, src=(\n" +
|
||||
" x0:=UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
|
||||
" x1:=UOp(Ops.CONST, dtypes.weakint, arg=1, src=()),\n" +
|
||||
" x1,)),\n" +
|
||||
" x0,))"
|
||||
)
|
||||
self.assertEqual(repr(b*b), expected)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
132
artifacts/package_sources/tinygrad/test/null/test_uop_resolve.py
Normal file
132
artifacts/package_sources/tinygrad/test/null/test_uop_resolve.py
Normal file
@@ -0,0 +1,132 @@
|
||||
import unittest
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.uop.ops import UOp, resolve
|
||||
|
||||
class TestUOpResolve(unittest.TestCase):
|
||||
def test_simple_int(self):
|
||||
u = UOp.const(4, dtypes.int)
|
||||
self.assertEqual(int(u), 4)
|
||||
|
||||
def test_weak_const(self):
|
||||
self.assertEqual(int(UOp.const(5)), 5)
|
||||
self.assertEqual(float(UOp.const(1.5)), 1.5)
|
||||
|
||||
def test_int_add(self):
|
||||
u = UOp.const(4, dtypes.int) + 7
|
||||
self.assertEqual(int(u), 11)
|
||||
|
||||
def test_lt(self):
|
||||
u = UOp.const(4) < 7
|
||||
self.assertTrue(u)
|
||||
|
||||
def test_rfloordiv(self):
|
||||
u = 8 // UOp.const(4, dtypes.int)
|
||||
self.assertEqual(int(u), 2)
|
||||
|
||||
def test_rtruediv(self):
|
||||
u = 9 / UOp.const(4, dtypes.float)
|
||||
self.assertEqual(float(u), 2.25)
|
||||
|
||||
def test_leq(self):
|
||||
u = UOp.const(4) <= 4
|
||||
self.assertTrue(u)
|
||||
|
||||
def test_ne(self):
|
||||
u = UOp.const(4) != 7
|
||||
self.assertTrue(u)
|
||||
|
||||
def test_ne_f(self):
|
||||
u = UOp.const(4) != 4
|
||||
self.assertFalse(u)
|
||||
|
||||
def test_ngt(self):
|
||||
u = UOp.const(4) > 7
|
||||
self.assertFalse(u)
|
||||
|
||||
def test_ssimplify(self):
|
||||
self.assertEqual((8 % UOp.const(4)).ssimplify(), 0)
|
||||
self.assertEqual((8 * UOp.const(4)).ssimplify(), 32)
|
||||
|
||||
def test_ambiguous_less_than(self):
|
||||
u = UOp.variable("i", 1, 10)
|
||||
self.assertTrue(resolve(u < 4))
|
||||
self.assertFalse(resolve(u < 4, False))
|
||||
self.assertTrue(resolve(u < 11, False))
|
||||
self.assertFalse(resolve(u < -1, False))
|
||||
self.assertFalse(resolve(u < -1, True))
|
||||
|
||||
def test_float_direct(self):
|
||||
u = UOp.const(4.5, dtypes.float) + 7
|
||||
self.assertEqual(float(u), 11.5)
|
||||
|
||||
def test_var_cmp_t(self):
|
||||
u = UOp.variable("i", 1, 10) < 20
|
||||
self.assertTrue(u)
|
||||
|
||||
def test_var_cmp_t2(self):
|
||||
u = UOp.variable("i", 1, 10)//2 < 20
|
||||
self.assertTrue(u)
|
||||
|
||||
def test_var_cmp_f(self):
|
||||
u = UOp.variable("i", 1, 10) < 1
|
||||
self.assertFalse(u)
|
||||
|
||||
def test_var_cmp_f2(self):
|
||||
u = UOp.variable("i", 1, 10) > 11
|
||||
self.assertFalse(u)
|
||||
|
||||
def test_or_true(self):
|
||||
u = UOp.variable("b", False, True, dtypes.bool) | True
|
||||
self.assertTrue(u)
|
||||
|
||||
def test_or_false(self):
|
||||
with self.assertRaises(ValueError):
|
||||
u = UOp.variable("b", False, True, dtypes.bool) | False
|
||||
self.assertTrue(u)
|
||||
|
||||
def test_and_false(self):
|
||||
u = UOp.variable("b", False, True, dtypes.bool) & False
|
||||
self.assertFalse(u)
|
||||
|
||||
def test_max(self):
|
||||
x = UOp.variable("x", 1, 10)
|
||||
y = UOp.variable("y", 5, 10)
|
||||
u = x.maximum(y)
|
||||
self.assertTrue(u < 20)
|
||||
self.assertFalse(u < 3)
|
||||
|
||||
def test_x_lt_x(self):
|
||||
x = UOp.variable("i", 1, 10)
|
||||
self.assertFalse(x < x)
|
||||
|
||||
def test_x_lt_xp1(self):
|
||||
x = UOp.variable("i", 1, 10)
|
||||
u = x < (x+1)
|
||||
# TODO: improve
|
||||
with self.assertRaises(ValueError):
|
||||
bool(u)
|
||||
|
||||
def test_and_true(self):
|
||||
u = UOp.variable("b", False, True, dtypes.bool) & True
|
||||
with self.assertRaises(ValueError):
|
||||
bool(u)
|
||||
|
||||
def test_var_cmp_range(self):
|
||||
v = UOp.variable("i", 1, 10)
|
||||
u = (v > 4) | (v < 6)
|
||||
# TODO: improve
|
||||
with self.assertRaises(ValueError):
|
||||
bool(u)
|
||||
|
||||
def test_var_cmp_assert(self):
|
||||
with self.assertRaises(ValueError):
|
||||
u = UOp.variable("i", 1, 10) < 5
|
||||
self.assertFalse(u)
|
||||
|
||||
def test_plus_ordering_lt(self):
|
||||
i = UOp.variable("i", 1, 10)
|
||||
j = UOp.variable("j", 1, 10)
|
||||
self.assertFalse((i+j) < (j+i))
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
1487
artifacts/package_sources/tinygrad/test/null/test_uop_symbolic.py
Normal file
1487
artifacts/package_sources/tinygrad/test/null/test_uop_symbolic.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,432 @@
|
||||
import unittest, math
|
||||
from tinygrad.uop.ops import UOp, Ops
|
||||
from tinygrad.dtype import dtypes, Invalid, truncate
|
||||
|
||||
class TestVminVmaxProperties(unittest.TestCase):
|
||||
def test_vmin_vmax_constant(self):
|
||||
# vmin and vmax for a constant
|
||||
uop = UOp.const(42)
|
||||
self.assertEqual(uop.vmin, 42)
|
||||
self.assertEqual(uop.vmax, 42)
|
||||
|
||||
def test_vmin_vmax_cmpne(self):
|
||||
uop = UOp.const(42)
|
||||
def test_bool(u, x):
|
||||
self.assertEqual(u.vmin, x)
|
||||
self.assertEqual(u.vmax, x)
|
||||
test_bool(uop != 42, False)
|
||||
test_bool(uop != 43, True)
|
||||
test_bool(uop != 41, True)
|
||||
|
||||
def test_vmin_vmax_addition_with_variable(self):
|
||||
# vmin and vmax for addition with a variable
|
||||
x = UOp.variable('x', 10, 20)
|
||||
uop = x + 5
|
||||
self.assertEqual(uop.vmin, 15)
|
||||
self.assertEqual(uop.vmax, 25)
|
||||
|
||||
def test_vmin_vmax_subtraction_with_variable(self):
|
||||
x = UOp.variable('x', 10, 20)
|
||||
uop = x - 5
|
||||
self.assertEqual(uop.vmin, 5)
|
||||
self.assertEqual(uop.vmax, 15)
|
||||
uop = 5 - x
|
||||
self.assertEqual(uop.vmin, -15)
|
||||
self.assertEqual(uop.vmax, -5)
|
||||
|
||||
def test_vmin_vmax_and_with_variable(self):
|
||||
x = UOp.variable('x', 10, 20)
|
||||
uop = x & 5
|
||||
self.assertEqual(uop.vmin, 0)
|
||||
self.assertEqual(uop.vmax, 5)
|
||||
|
||||
uop = x & 15
|
||||
self.assertEqual(uop.vmin, 0)
|
||||
self.assertEqual(uop.vmax, 15)
|
||||
|
||||
# TODO: this can be improved
|
||||
uop = x & 32
|
||||
self.assertEqual(uop.vmin, 0)
|
||||
self.assertEqual(uop.vmax, 20) # shoud be 0
|
||||
|
||||
def test_vmin_vmax_and_with_negative_variable(self):
|
||||
# when mask doesn't have sign bit set, result is always non-negative
|
||||
x = UOp.variable('x', -100, 100, dtypes.int32)
|
||||
# 511 = 0x1FF, doesn't have sign bit set for int32
|
||||
uop = x & 511
|
||||
self.assertEqual(uop.vmin, 0)
|
||||
self.assertEqual(uop.vmax, 511)
|
||||
|
||||
# 0x7FFFFFFF is max positive int32, doesn't have sign bit
|
||||
uop = x & 0x7FFFFFFF
|
||||
self.assertEqual(uop.vmin, 0)
|
||||
self.assertEqual(uop.vmax, 0x7FFFFFFF)
|
||||
|
||||
# negative mask: x & -1 could be anything since -1 has all bits set
|
||||
uop = x & -1
|
||||
self.assertEqual(uop.vmin, dtypes.int32.min)
|
||||
self.assertEqual(uop.vmax, dtypes.int32.max)
|
||||
|
||||
def test_vmin_vmax_multiplication_with_variable(self):
|
||||
# vmin and vmax for multiplication with a variable
|
||||
x = UOp.variable('x', -3, 4)
|
||||
uop = x * 2
|
||||
self.assertEqual(uop.vmin, -6)
|
||||
self.assertEqual(uop.vmax, 8)
|
||||
|
||||
def test_vmin_vmax_variable_inside_special(self):
|
||||
uop = UOp(Ops.SPECIAL, arg='gidx0', src=(UOp.variable('i', 1, 10, dtypes.int),))
|
||||
self.assertEqual(uop.vmin, 0)
|
||||
self.assertEqual(uop.vmax, 9)
|
||||
|
||||
def test_vmin_vmax_multiplication_0_inf(self):
|
||||
# vmin and vmax for multiplication with a variable
|
||||
x = UOp.const(0.0)
|
||||
y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(0), dtype=dtypes.float)
|
||||
uop = x * y
|
||||
# TODO: these should be 0, but definitely should not be nan
|
||||
self.assertEqual(uop.vmin, -math.inf)
|
||||
self.assertEqual(uop.vmax, math.inf)
|
||||
|
||||
def test_vmin_vmax_with_negative_multiplication(self):
|
||||
# vmin and vmax when multiplying by a negative number
|
||||
x = UOp.variable('x', 2, 5)
|
||||
uop = x * -3
|
||||
self.assertEqual(uop.vmin, -15)
|
||||
self.assertEqual(uop.vmax, -6)
|
||||
|
||||
def test_vmin_vmax_with_negative_multiplication2(self):
|
||||
# vmin and vmax when multiplying by a negative number
|
||||
x = UOp.variable('x', -2, 5)
|
||||
uop = x * -3
|
||||
self.assertEqual(uop.vmin, -15)
|
||||
self.assertEqual(uop.vmax, 6)
|
||||
|
||||
def test_vmin_vmax_nested_min_max(self):
|
||||
# vmin and vmax with nested min/max operations
|
||||
x = UOp.variable('x', 0, 10)
|
||||
uop = x.maximum(5).minimum(8)
|
||||
self.assertEqual(uop.vmin, 5)
|
||||
self.assertEqual(uop.vmax, 8)
|
||||
|
||||
def test_vmin_vmax_where(self):
|
||||
x = UOp.variable('x', 0, 10)
|
||||
y = UOp.variable('y', 1, 11)
|
||||
z = UOp.variable('z', 2, 12)
|
||||
uop = (x<5).where(y, z)
|
||||
self.assertEqual(uop.vmin, 1)
|
||||
self.assertEqual(uop.vmax, 12)
|
||||
|
||||
def test_vmin_vmax_shl(self):
|
||||
x = UOp.variable('x', 0, 10) << 5
|
||||
self.assertEqual(x.vmin, 0)
|
||||
self.assertEqual(x.vmax, 10 << 5)
|
||||
|
||||
def test_vmin_vmax_shr(self):
|
||||
x = UOp.variable('x', 0, 10) >> 2
|
||||
self.assertEqual(x.vmin, 0)
|
||||
self.assertEqual(x.vmax, 10 >> 2)
|
||||
|
||||
def test_vmin_vmax_cast_unsigned(self):
|
||||
# a fitting source keeps exact bounds: no wrap can occur
|
||||
self.assertEqual(UOp.variable('x', 5, 10).cast(dtypes.uint8)._min_max, (5, 10))
|
||||
# a possibly-negative or too-large source can wrap: conservative
|
||||
self.assertEqual(UOp.variable('x', -1, 10).cast(dtypes.uint8)._min_max, (0, 255))
|
||||
self.assertEqual(UOp.variable('x', 250, 260).cast(dtypes.uint8)._min_max, (0, 255))
|
||||
|
||||
def test_vmin_vmax_xor_neg1(self):
|
||||
x = UOp.variable('x', 3, 7)
|
||||
uop = x ^ -1
|
||||
self.assertEqual(uop.vmin, ~7)
|
||||
self.assertEqual(uop.vmax, ~3)
|
||||
# negative range
|
||||
y = UOp.variable('y', -10, -3)
|
||||
uop2 = y ^ -1
|
||||
self.assertEqual(uop2.vmin, ~(-3))
|
||||
self.assertEqual(uop2.vmax, ~(-10))
|
||||
# range spanning zero
|
||||
z = UOp.variable('z', -5, 6)
|
||||
uop3 = z ^ -1
|
||||
self.assertEqual(uop3.vmin, ~6)
|
||||
self.assertEqual(uop3.vmax, ~(-5))
|
||||
|
||||
def test_vmin_vmax_cast(self):
|
||||
x = UOp.variable('x', -10, 10, dtypes.int)
|
||||
x_float = x.cast(dtypes.float)
|
||||
self.assertEqual(x_float.vmin, -10)
|
||||
self.assertEqual(x_float.vmax, 10)
|
||||
x_bool = x.cast(dtypes.bool)
|
||||
self.assertEqual(x_bool.vmin, False)
|
||||
self.assertEqual(x_bool.vmax, True)
|
||||
x_uint = x.cast(dtypes.uint)
|
||||
self.assertEqual(x_uint.vmin, dtypes.uint.min)
|
||||
self.assertEqual(x_uint.vmax, dtypes.uint.max)
|
||||
|
||||
def test_vmin_vmax_cast_float_to_int(self):
|
||||
self.assertEqual(UOp.variable('x', -4.5, 4.5, dtypes.float).cast(dtypes.int)._min_max, (-4, 4))
|
||||
self.assertEqual(UOp.const(4.5).cast(dtypes.float).cast(dtypes.int)._min_max, (4, 4))
|
||||
x = UOp.const(4.5).cast(dtypes.float)
|
||||
self.assertIs(x.ne(x.cast(dtypes.int).cast(dtypes.float)).simplify().arg, True)
|
||||
|
||||
def test_vmin_vmax_invalid(self):
|
||||
i = UOp.invalid()
|
||||
self.assertNotEqual(i.vmin, i.vmax)
|
||||
|
||||
def test_vmin_vmax_invalid_vconst(self):
|
||||
x = UOp.const((0, 4, Invalid, Invalid))
|
||||
self.assertEqual((x.vmin, x.vmax), (0, 4))
|
||||
|
||||
class TestVminVmaxDivMod(unittest.TestCase):
|
||||
def test_vmin_vmax_division_positive(self):
|
||||
# vmin and vmax for division of a variable by a positive constant
|
||||
x = UOp.variable('x', 10, 20)
|
||||
uop = x // 2
|
||||
self.assertEqual(uop.vmin, 5)
|
||||
self.assertEqual(uop.vmax, 10)
|
||||
|
||||
def test_vmin_vmax_division_negative(self):
|
||||
# floor division of a variable by a negative constant
|
||||
x = UOp.variable('x', 10, 20)
|
||||
uop = x // -2
|
||||
self.assertEqual(uop.vmin, -10)
|
||||
self.assertEqual(uop.vmax, -5)
|
||||
uop = x // -3
|
||||
self.assertEqual(uop.vmin, -7)
|
||||
self.assertEqual(uop.vmax, -4)
|
||||
|
||||
x = UOp.variable('x', -20, -10)
|
||||
uop = x // -2
|
||||
self.assertEqual(uop.vmin, 5)
|
||||
self.assertEqual(uop.vmax, 10)
|
||||
uop = x // -3
|
||||
self.assertEqual(uop.vmin, 3)
|
||||
self.assertEqual(uop.vmax, 6)
|
||||
|
||||
def test_vmin_vmax_floordiv_floormod(self):
|
||||
x = UOp.variable('x', -7, 7)
|
||||
floordiv = x.alu(Ops.FLOORDIV, UOp.const(3))
|
||||
self.assertEqual(floordiv.vmin, -3)
|
||||
self.assertEqual(floordiv.vmax, 2)
|
||||
floormod = x.alu(Ops.FLOORMOD, UOp.const(3))
|
||||
self.assertEqual(floormod.vmin, 0)
|
||||
self.assertEqual(floormod.vmax, 2)
|
||||
# negative const divisor: floormod range is [c+1, 0]
|
||||
floormod_neg = x.alu(Ops.FLOORMOD, UOp.const(-3))
|
||||
self.assertEqual(floormod_neg.vmin, -2)
|
||||
self.assertEqual(floormod_neg.vmax, 0)
|
||||
|
||||
# cross 0
|
||||
x = UOp.variable('x', -10, 10)
|
||||
uop = x // -2
|
||||
self.assertEqual(uop.vmin, -5)
|
||||
self.assertEqual(uop.vmax, 5)
|
||||
uop = x // -3
|
||||
self.assertEqual(uop.vmin, -4)
|
||||
self.assertEqual(uop.vmax, 3)
|
||||
|
||||
def test_vmin_vmax_floordiv_floormod_empty_range(self):
|
||||
# empty numerator range (vmin > vmax, e.g. RANGE with end=0) short-circuits to (0, 0)
|
||||
rng = UOp.range(0, 0)
|
||||
self.assertEqual(rng.vmin, 0)
|
||||
self.assertEqual(rng.vmax, -1)
|
||||
self.assertEqual((rng // 4).vmin, 0)
|
||||
self.assertEqual((rng // 4).vmax, 0)
|
||||
self.assertEqual((rng % 4).vmin, 0)
|
||||
self.assertEqual((rng % 4).vmax, 0)
|
||||
|
||||
def test_vmin_vmax_div_symbolic(self):
|
||||
x = UOp.variable('x', 1, 10)
|
||||
y = UOp.variable('y', 3, 5)
|
||||
self.assertEqual((x//y).vmin, 0)
|
||||
self.assertEqual((x//y).vmax, 3)
|
||||
self.assertEqual(((-x)//y).vmin, -4)
|
||||
self.assertEqual(((-x)//y).vmax, -1)
|
||||
self.assertEqual((x//(-y)).vmin, -4)
|
||||
self.assertEqual((x//(-y)).vmax, -1)
|
||||
self.assertEqual(((-x)//(-y)).vmin, 0)
|
||||
self.assertEqual(((-x)//(-y)).vmax, 3)
|
||||
|
||||
self.assertEqual((100//y).vmin, 20)
|
||||
self.assertEqual((100//y).vmax, 33)
|
||||
self.assertEqual(((-100)//y).vmin, -34)
|
||||
self.assertEqual(((-100)//y).vmax, -20)
|
||||
self.assertEqual((100//(-y)).vmin, -34)
|
||||
self.assertEqual((100//(-y)).vmax, -20)
|
||||
self.assertEqual(((-100)//(-y)).vmin, 20)
|
||||
self.assertEqual(((-100)//(-y)).vmax, 33)
|
||||
|
||||
def test_vmin_vmax_mod_positive(self):
|
||||
# floor mod with positive divisor: result in [0, c-1] regardless of dividend sign
|
||||
positive = UOp.variable('positive', 10, 20)
|
||||
uop = positive % 3
|
||||
self.assertEqual(uop.vmin, 0)
|
||||
self.assertEqual(uop.vmax, 2)
|
||||
|
||||
negative = UOp.variable('negative', -20, -10)
|
||||
uop = negative % 3
|
||||
self.assertEqual(uop.vmin, 0)
|
||||
self.assertEqual(uop.vmax, 2)
|
||||
|
||||
mixed = UOp.variable('mixed', -20, 20)
|
||||
uop = mixed % 3
|
||||
self.assertEqual(uop.vmin, 0)
|
||||
self.assertEqual(uop.vmax, 2)
|
||||
|
||||
def test_vmin_vmax_mod_negative(self):
|
||||
# floor mod with negative divisor: result in [c+1, 0] regardless of dividend sign
|
||||
positive = UOp.variable('positive', 10, 20)
|
||||
uop = positive % -3
|
||||
self.assertEqual(uop.vmin, -2)
|
||||
self.assertEqual(uop.vmax, 0)
|
||||
|
||||
negative = UOp.variable('negative', -20, -10)
|
||||
uop = negative % -3
|
||||
self.assertEqual(uop.vmin, -2)
|
||||
self.assertEqual(uop.vmax, 0)
|
||||
|
||||
mixed = UOp.variable('mixed', -20, 20)
|
||||
uop = mixed % -3
|
||||
self.assertEqual(uop.vmin, -2)
|
||||
self.assertEqual(uop.vmax, 0)
|
||||
|
||||
class TestVminVmaxVConst(unittest.TestCase):
|
||||
def test_vmin_vmax_vconst_single_element(self):
|
||||
# vmin and vmax for a single-element vector constant
|
||||
uop = UOp.const((42,))
|
||||
self.assertEqual(uop.vmin, 42)
|
||||
self.assertEqual(uop.vmax, 42)
|
||||
|
||||
def test_vmin_vmax_vconst_multiple_elements(self):
|
||||
# vmin and vmax for a multi-element vector constant
|
||||
uop = UOp.const((10, 20, -5, 7))
|
||||
self.assertEqual(uop.vmin, -5)
|
||||
self.assertEqual(uop.vmax, 20)
|
||||
|
||||
def test_vmin_vmax_vconst_all_equal(self):
|
||||
# vmin and vmax for a vector where all elements are equal
|
||||
uop = UOp.const((7, 7, 7))
|
||||
self.assertEqual(uop.vmin, 7)
|
||||
self.assertEqual(uop.vmax, 7)
|
||||
|
||||
def test_vmin_vmax_vconst_with_negative_values(self):
|
||||
# vmin and vmax for a vector constant containing negative values
|
||||
uop = UOp.const((-10, -20, -5, -15))
|
||||
self.assertEqual(uop.vmin, -20)
|
||||
self.assertEqual(uop.vmax, -5)
|
||||
|
||||
def test_vmin_vmax_vconst_with_floats(self):
|
||||
# vmin and vmax for a vector constant of float values
|
||||
uop = UOp.const((1.5, -3.2, 0.0))
|
||||
self.assertEqual(uop.vmin, truncate[dtypes.default_float](-3.2))
|
||||
self.assertEqual(uop.vmax, truncate[dtypes.default_float](1.5))
|
||||
|
||||
def test_vmin_vmax_vconst_with_bools(self):
|
||||
# vmin and vmax for a vector constant of bool values
|
||||
uop = UOp.const((True, False, False))
|
||||
self.assertIs(uop.vmin, False)
|
||||
self.assertIs(uop.vmax, True)
|
||||
|
||||
def test_vmin_vmax_vector_with_gep(self):
|
||||
# vmin and vmax for a vector constant of bool values
|
||||
d1 = UOp.param(1, dtypes.int, (1,))
|
||||
idx = UOp.const(0)
|
||||
val = UOp(Ops.LOAD, src=(d1.index(idx),))
|
||||
uop = (val // 32)
|
||||
self.assertEqual(uop.vmin, -67108864)
|
||||
self.assertEqual(uop.vmax, 67108863)
|
||||
|
||||
class TestConstFactor(unittest.TestCase):
|
||||
def test_const_factor_constant(self):
|
||||
# const_factor for a constant
|
||||
uop = UOp.const(42)
|
||||
self.assertEqual(uop.const_factor(), 42)
|
||||
|
||||
def test_const_factor_addition(self):
|
||||
# const_factor for an addition of constants
|
||||
uop = UOp.const(30) + UOp.const(12)
|
||||
self.assertEqual(uop.const_factor(), 6) # GCD(30, 12) = 6
|
||||
|
||||
def test_const_factor_multiplication(self):
|
||||
# const_factor for a multiplication of constants
|
||||
uop = UOp.const(5) * UOp.const(7)
|
||||
self.assertEqual(uop.const_factor(), 5) # For multiplication, it's one of the factors
|
||||
|
||||
def test_const_factor_with_variable(self):
|
||||
# const_factor for an expression involving a variable
|
||||
x = UOp.variable('x', 10, 20)
|
||||
uop = x * 3
|
||||
self.assertEqual(uop.const_factor(), 3)
|
||||
|
||||
def test_const_factor_division(self):
|
||||
# const_factor for an expression with division
|
||||
x = UOp.variable('x', 10, 20)
|
||||
uop = x // 4
|
||||
self.assertEqual(uop.const_factor(), 1) # Division reduces the const_factor to 1
|
||||
|
||||
def test_const_factor_multiplication_of_var_and_const(self):
|
||||
# const_factor for multiplication of a variable and a constant
|
||||
x = UOp.variable('x', 6, 18)
|
||||
uop = x * 4
|
||||
self.assertEqual(uop.const_factor(), 4) # Constant factor 4
|
||||
|
||||
@unittest.skip("broken")
|
||||
def test_const_factor_multiplication_of_consts_and_vars(self):
|
||||
# Multiplying constants and variables
|
||||
x = UOp.variable('x', 10, 20)
|
||||
uop = (x * 3) * 5
|
||||
self.assertEqual(uop.const_factor(), 15) # Constant multipliers are combined (3 * 5 = 15)
|
||||
|
||||
def test_const_factor_variable_multiple_of(self):
|
||||
x = UOp.variable('x', 16, 32, multiple_of=4)
|
||||
self.assertEqual(x.const_factor(), 4)
|
||||
|
||||
class TestDivides(unittest.TestCase):
|
||||
def test_divides_constant_exact(self):
|
||||
# Divides a constant by an exact divisor
|
||||
uop = UOp.const(42)
|
||||
result = uop.divides(7)
|
||||
self.assertIsNotNone(result)
|
||||
self.assertEqual(result.const_factor(), 6) # 42 / 7 = 6
|
||||
|
||||
def test_divides_constant_inexact(self):
|
||||
# Try to divide a constant by a non-exact divisor
|
||||
uop = UOp.const(42)
|
||||
result = uop.divides(5)
|
||||
self.assertIsNone(result) # 42 is not divisible by 5
|
||||
|
||||
@unittest.skip("broken")
|
||||
def test_divides_variable_and_constant(self):
|
||||
# Multiplying a variable by a constant, then dividing by the same constant
|
||||
x = UOp.variable('x', 10, 20)
|
||||
uop = x * 6
|
||||
result = uop.divides(6)
|
||||
self.assertIsNotNone(result)
|
||||
self.assertEqual(result, x) # (x * 6) / 6 = x
|
||||
|
||||
def test_divides_complex_expression(self):
|
||||
# Dividing a more complex expression
|
||||
x = UOp.variable('x', 10, 20)
|
||||
uop = (x * 6) + 18
|
||||
result = uop.divides(6)
|
||||
self.assertIsNotNone(result)
|
||||
self.assertEqual(result.const_factor(), 1) # (x + 3), const_factor is 1
|
||||
|
||||
def test_divides_with_inexact_factors(self):
|
||||
# Multiplying by a constant but dividing by a non-exact divisor
|
||||
x = UOp.variable('x', 15, 45)
|
||||
uop = x * 4
|
||||
result = uop.divides(3)
|
||||
self.assertIsNone(result) # Cannot divide by 3, since 4 is not divisible by 3
|
||||
|
||||
def test_divides_variable_multiple_of_exact(self):
|
||||
x = UOp.variable('x', 16, 32, multiple_of=4)
|
||||
result = x.divides(4)
|
||||
self.assertIsNotNone(result)
|
||||
|
||||
def test_divides_variable_multiple_of_factor(self):
|
||||
x = UOp.variable('x', 16, 32, multiple_of=4)
|
||||
result = x.divides(2)
|
||||
self.assertIsNotNone(result)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
497
artifacts/package_sources/tinygrad/test/null/test_uops.py
Normal file
497
artifacts/package_sources/tinygrad/test/null/test_uops.py
Normal file
@@ -0,0 +1,497 @@
|
||||
# uops tests that pass on NULL backend (no copyout needed)
|
||||
import math, unittest
|
||||
import numpy as np
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.helpers import Timing, Context, cdiv
|
||||
from tinygrad.dtype import dtypes, AddrSpace, ConstFloat, Invalid # noqa: F401
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.uop.ops import Ops, ParamArg, PatternMatcher, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite # noqa: F401 # ParamArg used by eval(str(uop)) roundtrip tests
|
||||
from tinygrad.uop.weak import pm_lower_index_dtype
|
||||
from tinygrad.uop.spec import spec_program, spec_shared, type_verify
|
||||
from tinygrad.uop.symbolic import sym, pm_remove_invalid
|
||||
from test.helpers import eval_uop, to_uops_list
|
||||
|
||||
class TestDTypeFromUOp(unittest.TestCase):
|
||||
def test_broadcastable_promotion(self):
|
||||
self.assertEqual(dtype_from_uop(Ops.ADD, (UOp.const(1.0).cast(dtypes.float32), UOp.const(1.0).cast(dtypes.float16)), None), dtypes.float32)
|
||||
self.assertEqual(dtype_from_uop(Ops.MUL, (UOp.const(1).cast(dtypes.int8), UOp.const(1).cast(dtypes.int32)), None), dtypes.int32)
|
||||
|
||||
def test_same_dtype_fast_path(self):
|
||||
src = (UOp.const(1), UOp.const(2))
|
||||
self.assertEqual(dtype_from_uop(Ops.ADD, src, None), dtypes.weakint)
|
||||
|
||||
def test_where_promotion(self):
|
||||
cond = UOp.const(True)
|
||||
srcs = (cond, UOp.const(1.0).cast(dtypes.float32), UOp.const(1.0).cast(dtypes.float16))
|
||||
self.assertEqual(dtype_from_uop(Ops.WHERE, srcs, None), dtypes.float32)
|
||||
idx = UOp.range(4, 0)
|
||||
self.assertEqual(idx.valid(idx < 4).dtype, dtypes.weakint)
|
||||
|
||||
def test_const_dtype_from_value(self):
|
||||
self.assertEqual(dtype_from_uop(Ops.CONST, (), True), dtypes.bool)
|
||||
self.assertEqual(dtype_from_uop(Ops.CONST, (), ConstFloat(3.0)), dtypes.weakfloat)
|
||||
self.assertEqual(dtype_from_uop(Ops.CONST, (), Invalid), dtypes.bool)
|
||||
self.assertRaises(TypeError, dtype_from_uop, Ops.CONST, (), (1, 2))
|
||||
|
||||
@Context(SPEC=2)
|
||||
def test_const_default_dtype_is_derived(self):
|
||||
self.assertEqual(UOp(Ops.CONST, arg=ConstFloat(3.0)).dtype, dtypes.weakfloat)
|
||||
self.assertEqual(UOp(Ops.CONST, arg=True).dtype, dtypes.bool)
|
||||
self.assertEqual(UOp(Ops.CONST, arg=Invalid).dtype, dtypes.bool)
|
||||
# an explicit (strong) const dtype is legal until the field is removed
|
||||
self.assertEqual(UOp.const(3, dtypes.int32).dtype, dtypes.int32)
|
||||
|
||||
def test_weak_dtype_rejected_by_program_spec(self):
|
||||
for weak, concrete, value in ((dtypes.weakint, dtypes.int32, 1), (dtypes.weakfloat, dtypes.float32, 1.0)):
|
||||
with self.assertRaises(RuntimeError): type_verify(UOp.const(value, weak).sink(), spec_program)
|
||||
type_verify(UOp.const(value, concrete).sink(), spec_program)
|
||||
|
||||
def test_invalid_stated_dtype(self):
|
||||
# UOp.const normalizes a stated dtype away (const_like/full pass their position's); the core constructor does not,
|
||||
# and the spec is what rejects a non-bool Invalid
|
||||
self.assertIs(UOp.const(Invalid, dtypes.float32), UOp.invalid())
|
||||
with self.assertRaises(RuntimeError): type_verify(UOp(Ops.CONST, dtypes.float32, arg=Invalid), spec_shared)
|
||||
|
||||
def test_invalid_dtype_and_consumers(self):
|
||||
invalid = UOp.invalid()
|
||||
self.assertIs(invalid.dtype, dtypes.bool)
|
||||
self.assertIs(UOp.const(Invalid, dtypes.float32), invalid)
|
||||
scratch = Tensor.invalids(4, dtype=dtypes.float32)
|
||||
self.assertEqual((scratch.dtype, next(u.dtype for u in scratch.uop.toposort() if u.op is Ops.BUFFER), next(u.dtype for u in scratch.uop.toposort()
|
||||
if u.is_invalid)), (dtypes.float32, dtypes.float32, dtypes.bool))
|
||||
invalid, value = UOp.invalid(), UOp.const(1, dtypes.float32)
|
||||
for u in (UOp.param(0, dtypes.bool, ()).where(value, invalid), value+invalid, UOp.stack(value, invalid)): self.assertIs(u.src[-1], invalid)
|
||||
for u in (UOp(Ops.STACK, dtypes.float32, src=(value, invalid)), UOp(Ops.ADD, dtypes.float32, src=(value, invalid)),
|
||||
UOp.const(True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)),
|
||||
UOp.param(0, dtypes.float32, (4,)).index(invalid)): type_verify(u, spec_shared)
|
||||
gate, value = UOp.param(0, dtypes.bool, ()), UOp.param(1, dtypes.float, ())
|
||||
self.assertIs((out:=graph_rewrite(gate.where(value, UOp.invalid()), pm_remove_invalid)).src[2], UOp.const(0, dtypes.float))
|
||||
type_verify(out.sink(), spec_program)
|
||||
|
||||
def test_remove_invalid_stack_lanes(self):
|
||||
stack = UOp(Ops.STACK, dtypes.half, (UOp.const(1, dtypes.half), UOp.invalid()))
|
||||
out = graph_rewrite(stack, pm_remove_invalid)
|
||||
self.assertEqual(out.src, (UOp.const(1, dtypes.half), UOp.const(0, dtypes.half)))
|
||||
type_verify(out.sink(), spec_program)
|
||||
|
||||
class TestLowerIndexDtype(unittest.TestCase):
|
||||
def test_gated_shrink_lowers_to_selected_width(self):
|
||||
# coalesce builds gated SHRINKs for masked vectorized loads; lowering must resolve them at the
|
||||
# width the offset bounds select (this one needs long)
|
||||
buf = UOp.param(0, dtypes.float, (2**31+64,))
|
||||
i = UOp.variable("i", 0, 2**28)
|
||||
shrink = UOp(Ops.SHRINK, src=(buf, (i*24).valid(i < 2**28), UOp.const(4)))
|
||||
lowered = graph_rewrite(shrink.sink(), pm_lower_index_dtype)
|
||||
self.assertTrue(all(u.dtype != dtypes.weakint for u in lowered.backward_slice_with_self), "lowering must resolve all weakint")
|
||||
sh = next(u for u in lowered.backward_slice_with_self if u.op is Ops.SHRINK)
|
||||
self.assertEqual(sh.src[1].dtype, dtypes.long)
|
||||
|
||||
def test_reg_buffer_size_lowers(self):
|
||||
reg = UOp.placeholder((4,), dtypes.float, 0, addrspace=AddrSpace.REG)
|
||||
self.assertEqual(reg.src[0].dtype, dtypes.weakint)
|
||||
lowered = graph_rewrite(reg.sink(), pm_lower_index_dtype)
|
||||
self.assertTrue(all(u.dtype != dtypes.weakint for u in lowered.backward_slice_with_self), "lowering must resolve all weakint")
|
||||
self.assertEqual(next(u for u in lowered.backward_slice_with_self if u.op is Ops.BUFFER).src[0].dtype, dtypes.int)
|
||||
|
||||
class TestSafeCast(unittest.TestCase):
|
||||
def test_cast_folds(self):
|
||||
a = UOp.variable("a", 1, 10, dtype=dtypes.int32)
|
||||
self.assertEqual(a.cast(dtypes.int64).cast(dtypes.int32).simplify(), a)
|
||||
self.assertEqual(a.cast(dtypes.double).cast(dtypes.int32).simplify(), a)
|
||||
a = UOp.variable("a", 1, 10, dtype=dtypes.uint8)
|
||||
self.assertEqual(a.cast(dtypes.int64).cast(dtypes.uint8).simplify(), a)
|
||||
self.assertEqual(a.cast(dtypes.uint32).cast(dtypes.uint8).simplify(), a)
|
||||
|
||||
def test_remove_intermediate_cast(self):
|
||||
a = UOp.variable("a", 0., 100., dtype=dtypes.half)
|
||||
self.assertEqual(a.cast(dtypes.double).cast(dtypes.float).simplify(), a.cast(dtypes.float))
|
||||
a = UOp.variable("a", 1, 10, dtype=dtypes.int32)
|
||||
# TODO: double preserves certain int dtypes
|
||||
self.assertEqual(a.cast(dtypes.double).cast(dtypes.float).simplify(), a.cast(dtypes.float))
|
||||
self.assertEqual(a.cast(dtypes.int64).cast(dtypes.int16).simplify(), a.cast(dtypes.int16))
|
||||
a = UOp.variable("a", 1, 10, dtype=dtypes.uint8)
|
||||
self.assertEqual(a.cast(dtypes.int64).cast(dtypes.int32).simplify(), a.cast(dtypes.int32))
|
||||
|
||||
def test_safe_cast_using_bounds(self):
|
||||
a = UOp.variable("a", 1, 10, dtype=dtypes.uint64)
|
||||
self.assertEqual(a.cast(dtypes.int16).cast(dtypes.int).simplify(), a.cast(dtypes.int))
|
||||
a = UOp.variable("a", -10, 10, dtype=dtypes.int32)
|
||||
self.assertEqual(a.cast(dtypes.int8).cast(dtypes.int64).simplify(), a.cast(dtypes.int64))
|
||||
self.assertEqual(a.cast(dtypes.int8).cast(dtypes.float).simplify(), a.cast(dtypes.float))
|
||||
|
||||
class TestConstFloatEq(unittest.TestCase):
|
||||
def test_nan_eq_ne_agree(self):
|
||||
nan = dtypes.float32.const(math.nan)
|
||||
self.assertTrue(nan == math.nan)
|
||||
self.assertFalse(nan != math.nan) # float.__ne__ would say True here
|
||||
self.assertFalse(nan == Invalid)
|
||||
self.assertTrue(nan != Invalid) # __ne__ must defer to the reflected eq, not swallow NotImplemented
|
||||
|
||||
def test_invalid_eq_defers_to_reflected(self):
|
||||
class HoldsInvalid: # a carrier that knows it holds Invalid. returning False for foreign types would silence its eq
|
||||
def __eq__(self, other): return other is Invalid
|
||||
self.assertTrue(Invalid == HoldsInvalid())
|
||||
self.assertFalse(Invalid != HoldsInvalid())
|
||||
|
||||
def test_matchers_agree_on_nan(self):
|
||||
n = UOp.const(math.nan, dtypes.float32)
|
||||
for compiled in (False, True):
|
||||
pm = PatternMatcher([(UPat(Ops.CONST, arg=math.nan), lambda: True)], compiled=compiled)
|
||||
self.assertTrue(pm.rewrite(n), f"{compiled=}")
|
||||
|
||||
class TestExecALU(unittest.TestCase):
|
||||
def test_sqrt(self):
|
||||
self.assertEqual(exec_alu(Ops.SQRT, dtypes.float, (0.0,)), 0.0)
|
||||
|
||||
def test_trunc_nonfinite(self):
|
||||
self.assertEqual(exec_alu(Ops.TRUNC, dtypes.float, (math.inf,)), math.inf)
|
||||
self.assertEqual(exec_alu(Ops.TRUNC, dtypes.float, (-math.inf,)), -math.inf)
|
||||
self.assertTrue(math.isnan(exec_alu(Ops.TRUNC, dtypes.float, (math.nan,))))
|
||||
|
||||
def test_invalid_poison(self):
|
||||
# Invalid poisons any binary op regardless of result dtype: a comparison must not fold to a boolean
|
||||
self.assertIs(exec_alu(Ops.CMPLT, dtypes.bool, (Invalid, 1)), Invalid)
|
||||
self.assertIs(exec_alu(Ops.CMPNE, dtypes.bool, (Invalid, 1)), Invalid)
|
||||
self.assertIs(exec_alu(Ops.ADD, dtypes.weakint, (Invalid, 1)), Invalid)
|
||||
|
||||
def test_div(self):
|
||||
self.assertEqual(exec_alu(Ops.CDIV, dtypes.int8, (8, 2)), 4)
|
||||
self.assertEqual(exec_alu(Ops.CDIV, dtypes.int8, (7, 3)), 2)
|
||||
self.assertEqual(exec_alu(Ops.CDIV, dtypes.int8, (7, -3)), -2)
|
||||
self.assertEqual(exec_alu(Ops.CDIV, dtypes.int8, (-50, 6)), -8)
|
||||
|
||||
def test_floordiv(self):
|
||||
self.assertEqual(exec_alu(Ops.FLOORDIV, dtypes.int8, (8, 2)), 4)
|
||||
self.assertEqual(exec_alu(Ops.FLOORDIV, dtypes.int8, (7, 3)), 2)
|
||||
self.assertEqual(exec_alu(Ops.FLOORDIV, dtypes.int8, (7, -3)), -3)
|
||||
self.assertEqual(exec_alu(Ops.FLOORDIV, dtypes.int8, (-7, 3)), -3)
|
||||
self.assertEqual(exec_alu(Ops.FLOORDIV, dtypes.int8, (-50, 6)), -9)
|
||||
|
||||
def test_floormod(self):
|
||||
self.assertEqual(exec_alu(Ops.FLOORMOD, dtypes.int8, (8, 2)), 0)
|
||||
self.assertEqual(exec_alu(Ops.FLOORMOD, dtypes.int8, (7, 3)), 1)
|
||||
self.assertEqual(exec_alu(Ops.FLOORMOD, dtypes.int8, (7, -3)), -2)
|
||||
self.assertEqual(exec_alu(Ops.FLOORMOD, dtypes.int8, (-7, 3)), 2)
|
||||
self.assertEqual(exec_alu(Ops.FLOORMOD, dtypes.int8, (-50, 6)), 4)
|
||||
|
||||
np.testing.assert_allclose(exec_alu(Ops.MUL, dtypes.float32, (7.0, exec_alu(Ops.RECIPROCAL, dtypes.float32, (3.0,)))), 2+(1.0/3.0))
|
||||
np.testing.assert_allclose(exec_alu(Ops.MUL, dtypes.float32, (7.0, exec_alu(Ops.RECIPROCAL, dtypes.float32, (-3.0,)))), -2-(1.0/3.0))
|
||||
|
||||
def test_recip(self):
|
||||
np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (8,)), 1/8)
|
||||
np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (7,)), 1/7)
|
||||
np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (-3,)), 1/-3)
|
||||
np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (-50,)), 1/-50)
|
||||
|
||||
np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, ((32+521+3),)), 1/(32+521+3))
|
||||
np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, ((34**2),)), 1/(34**2))
|
||||
np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (10,)), 1/10)
|
||||
|
||||
def test_bool_cmplt(self):
|
||||
self.assertEqual(exec_alu(Ops.CMPLT, dtypes.bool, (False, False)), False)
|
||||
self.assertEqual(exec_alu(Ops.CMPLT, dtypes.bool, (False, True)), True)
|
||||
self.assertEqual(exec_alu(Ops.CMPLT, dtypes.bool, (True, False)), False)
|
||||
self.assertEqual(exec_alu(Ops.CMPLT, dtypes.bool, (True, True)), False)
|
||||
|
||||
def test_bool_cmpne(self):
|
||||
self.assertEqual(exec_alu(Ops.CMPNE, dtypes.bool, (False, False)), False)
|
||||
self.assertEqual(exec_alu(Ops.CMPNE, dtypes.bool, (False, True)), True)
|
||||
self.assertEqual(exec_alu(Ops.CMPNE, dtypes.bool, (True, False)), True)
|
||||
self.assertEqual(exec_alu(Ops.CMPNE, dtypes.bool, (True, True)), False)
|
||||
|
||||
def test_bool_where(self):
|
||||
self.assertEqual(exec_alu(Ops.WHERE, dtypes.bool, (False, False, False)), False)
|
||||
self.assertEqual(exec_alu(Ops.WHERE, dtypes.int, (False, 2, 4)), 4)
|
||||
np.testing.assert_allclose(exec_alu(Ops.WHERE, dtypes.float, (False, 2.2, 4.5)), 4.5)
|
||||
|
||||
def test_overflow(self):
|
||||
self.assertEqual(exec_alu(Ops.ADD, dtypes.uint8, (250, 250)), 244)
|
||||
self.assertEqual(exec_alu(Ops.ADD, dtypes.uint8, (256, 0)), 0)
|
||||
self.assertEqual(exec_alu(Ops.ADD, dtypes.uint8, (0, -1)), 255)
|
||||
self.assertEqual(exec_alu(Ops.ADD, dtypes.uint8, (0, -1000)), 24)
|
||||
|
||||
self.assertEqual(exec_alu(Ops.ADD, dtypes.int8, (127, 0)), 127)
|
||||
self.assertEqual(exec_alu(Ops.ADD, dtypes.int8, (-128, 0)), -128)
|
||||
self.assertEqual(exec_alu(Ops.ADD, dtypes.int8, (-100, -100)), 56)
|
||||
self.assertEqual(exec_alu(Ops.ADD, dtypes.int8, (-1000, -0)), 24)
|
||||
self.assertEqual(exec_alu(Ops.ADD, dtypes.int8, (-130, -0)), 126)
|
||||
|
||||
self.assertEqual(exec_alu(Ops.ADD, dtypes.int8, (1, 1)), 2)
|
||||
self.assertEqual(exec_alu(Ops.ADD, dtypes.int8, (-128, 0)), -128)
|
||||
|
||||
# test no truncate
|
||||
self.assertEqual(exec_alu(Ops.ADD, dtypes.uint8, (250, 250), truncate_output=False), 500)
|
||||
|
||||
class TestGatedStoreRewrite(unittest.TestCase):
|
||||
def test_tiny_gate_store(self):
|
||||
gmem = UOp.param(0, dtypes.float, (8,))
|
||||
gidx0 = UOp.special(4, 'gidx0')
|
||||
gate = gidx0<UOp.const(1)
|
||||
idx = UOp(Ops.INDEX, src=(gmem, (gidx0 * UOp.const(2)).valid(gate)))
|
||||
val = UOp.const(42.0).cast(dtypes.float)
|
||||
store = UOp(Ops.STORE, src=(idx, val))
|
||||
uops = to_uops_list([store])
|
||||
if_uop = next(u for u in uops if u.op is Ops.IF)
|
||||
endif = next(u for u in uops if u.op is Ops.ENDIF)
|
||||
assert endif.src[0] is if_uop
|
||||
gated_uops = tuple(uops[uops.index(if_uop)+1:uops.index(endif)])
|
||||
self.assertEqual(len(gated_uops), 1)
|
||||
self.assertIs(gated_uops[-1].op, Ops.STORE)
|
||||
self.assertEqual(len(gated_uops[-1].src), 2)
|
||||
|
||||
def test_gate_some_stores(self):
|
||||
gmem0 = UOp.param(0, dtypes.float, (8,))
|
||||
gmem1 = UOp.param(1, dtypes.float, (8,))
|
||||
gidx0 = UOp.special(4, 'gidx0')
|
||||
idx = gidx0 * UOp.const(2)
|
||||
idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gidx0<UOp.const(1))))
|
||||
idx1 = UOp(Ops.INDEX, src=(gmem1, idx))
|
||||
val = UOp.const(42.0).cast(dtypes.float)
|
||||
stores = [UOp.store(idx0, val), UOp.store(idx1, val)]
|
||||
uops = to_uops_list(stores)
|
||||
if_uop = next(u for u in uops if u.op is Ops.IF)
|
||||
endif = next(u for u in uops if u.op is Ops.ENDIF)
|
||||
assert endif.src[0] is if_uop
|
||||
gated_uops = tuple(uops[uops.index(if_uop)+1:uops.index(endif)])
|
||||
self.assertEqual(len(gated_uops), 1)
|
||||
self.assertIs(gated_uops[-1].op, Ops.STORE)
|
||||
self.assertEqual(len(gated_uops[-1].src), 2)
|
||||
|
||||
# scaled down version of TestLinearizerDumb.test_unmerged_ifs
|
||||
@unittest.skip("we don't merge ifs anymore")
|
||||
def test_merge_ifs_alt(self):
|
||||
gmem0 = UOp.param(0, dtypes.float, (8,))
|
||||
gmem1 = UOp.param(1, dtypes.float, (8,))
|
||||
gidx0 = UOp.special(4, 'gidx0')
|
||||
idx = gidx0*UOp.const(2)
|
||||
gate = gidx0<UOp.const(1)
|
||||
idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gate)))
|
||||
idx1 = UOp(Ops.INDEX, src=(gmem1, idx.valid(gate)))
|
||||
val = UOp.const(42.0).cast(dtypes.float)
|
||||
stores = [UOp.store(idx0, val), UOp.store(idx1, val)]
|
||||
uops = to_uops_list(stores)
|
||||
ifs = [u for u in uops if u.op is Ops.IF]
|
||||
endifs = [u for u in uops if u.op is Ops.ENDIF]
|
||||
self.assertEqual(len(ifs), 1)
|
||||
self.assertEqual(len(endifs), 1)
|
||||
gated_uops = tuple(uops[uops.index(ifs[0])+1:uops.index(endifs[0])])
|
||||
self.assertEqual(len(gated_uops), 2)
|
||||
for x in gated_uops: self.assertIs(x.op, Ops.STORE)
|
||||
for x in gated_uops: self.assertEqual(len(x.src), 2)
|
||||
|
||||
@unittest.skipIf(Device.DEFAULT == "METAL", "compiler bug")
|
||||
@unittest.skipUnless(Ops.SHR in Device[Device.DEFAULT].renderer.code_for_op, "fast_idiv requires SHR")
|
||||
class TestFastIdiv(unittest.TestCase):
|
||||
def test_division_power_of_two(self):
|
||||
for dt in (dtypes.int32, dtypes.uint32):
|
||||
g = UOp.param(0, dt, (3,))
|
||||
c = UOp.const(2).cast(dt)
|
||||
l = g.index(c)
|
||||
a = UOp(Ops.CDIV, dt, (l, c))
|
||||
uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
|
||||
Device[Device.DEFAULT].renderer.render(uops)
|
||||
ops = [x.op for x in uops]
|
||||
self.assertIn(Ops.SHR, ops, f"For dtype={dt} divison by power of two did not simplify to shift")
|
||||
self.assertNotIn(Ops.CDIV, ops, f"For dtype={dt} divison by power of two did not simplify to shift")
|
||||
|
||||
def test_floormod_power_of_two(self):
|
||||
# FLOORMOD by a power of two lowers to AND (correct floor mod for any sign in two's complement)
|
||||
for dt in (dtypes.int32, dtypes.uint32):
|
||||
g = UOp.param(0, dt, (9,))
|
||||
c = UOp.const(8).cast(dt)
|
||||
a = UOp(Ops.FLOORMOD, dt, (g.index(c), c))
|
||||
uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
|
||||
ops = [x.op for x in uops]
|
||||
self.assertIn(Ops.AND, ops, f"For dtype={dt} FLOORMOD by pow2 did not simplify to AND")
|
||||
self.assertNotIn(Ops.CMOD, ops, f"For dtype={dt} FLOORMOD by pow2 left a MOD")
|
||||
self.assertNotIn(Ops.FLOORMOD, ops, f"For dtype={dt} FLOORMOD survived past late rewrite")
|
||||
|
||||
def test_floordiv_power_of_two_uint(self):
|
||||
# uint FLOORDIV by a power of two lowers to a shift, leaving no IDIV/FLOORDIV in the kernel
|
||||
for dt in (dtypes.uint32, dtypes.uint64):
|
||||
g = UOp.param(0, dt, (3,))
|
||||
c = UOp.const(2).cast(dt)
|
||||
a = UOp(Ops.FLOORDIV, dt, (g.index(c), c))
|
||||
uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
|
||||
ops = [x.op for x in uops]
|
||||
self.assertIn(Ops.SHR, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
|
||||
self.assertNotIn(Ops.CDIV, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
|
||||
self.assertNotIn(Ops.FLOORDIV, ops, f"For dtype={dt} FLOORDIV survived past late rewrite")
|
||||
|
||||
@Context(DISABLE_FAST_IDIV=0)
|
||||
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
|
||||
def test_fast_idiv_and_mod(self):
|
||||
g = UOp.param(0, dtypes.uint32, (4,))
|
||||
c = UOp.const(3).cast(dtypes.uint)
|
||||
l = g.index(c)
|
||||
a = UOp(Ops.CDIV, src=(l, c))
|
||||
uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
|
||||
Device[Device.DEFAULT].renderer.render(uops)
|
||||
ops = [x.op for x in uops]
|
||||
self.assertIn(Ops.SHR, ops)
|
||||
self.assertNotIn(Ops.CDIV, ops)
|
||||
|
||||
b = UOp(Ops.CMOD, src=(l, c))
|
||||
uops = to_uops_list([b], ren=Device[Device.DEFAULT].renderer)
|
||||
Device[Device.DEFAULT].renderer.render(uops)
|
||||
ops = [x.op for x in uops]
|
||||
self.assertIn(Ops.SHR, ops)
|
||||
self.assertNotIn(Ops.CMOD, ops)
|
||||
|
||||
@Context(DISABLE_FAST_IDIV=0)
|
||||
def test_fast_idiv_bounded_numerator_zero(self):
|
||||
x = UOp.variable("x", 0, 1, dtype=dtypes.int32)
|
||||
for val in range(2):
|
||||
self.assertEqual(eval_uop(x.alu(Ops.CDIV, UOp.const(3).cast(x.dtype)), vals=(val,)), cdiv(val, 3))
|
||||
|
||||
@Context(DISABLE_FAST_IDIV=0)
|
||||
def test_fast_idiv_remove_powers_of_two(self):
|
||||
ridx = UOp.range(2**20, 0)
|
||||
uops = to_uops_list([ridx//(7*64)], ren=Device[Device.DEFAULT].renderer)
|
||||
ops = [x.op for x in uops]
|
||||
# this requires shifting out the powers of two before doing fast_idiv
|
||||
# (((ridx0>>6)*18725)>>17) instead of (int)((((long)(ridx0)*1198373)>>29))
|
||||
self.assertNotIn(Ops.CAST, ops)
|
||||
|
||||
@unittest.expectedFailure
|
||||
def test_fast_idiv_overflow(self):
|
||||
# This will be possible with a slightly different method for fast_idiv
|
||||
g = UOp.param(0, dtypes.uint32, (8,))
|
||||
c = UOp.const(7).cast(dtypes.uint)
|
||||
l = UOp(Ops.LOAD, src=(g.index(c),))
|
||||
a = UOp(Ops.CDIV, src=(l, c))
|
||||
uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
|
||||
Device[Device.DEFAULT].renderer.render(uops)
|
||||
ops = [x.op for x in uops]
|
||||
self.assertIn(Ops.SHR, ops)
|
||||
self.assertNotIn(Ops.CDIV, ops)
|
||||
|
||||
def test_disable_fast_idiv(self):
|
||||
g = UOp.param(0, dtypes.uint32, (4,))
|
||||
c = UOp.const(3).cast(dtypes.uint)
|
||||
l = g.index(c)
|
||||
a = UOp(Ops.CDIV, src=(l, c))
|
||||
with Context(DISABLE_FAST_IDIV=1):
|
||||
uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
|
||||
ops = [x.op for x in uops]
|
||||
self.assertNotIn(Ops.SHR, ops)
|
||||
self.assertIn(Ops.CDIV, ops)
|
||||
|
||||
class TestUOpMethod(unittest.TestCase):
|
||||
@unittest.skip("uops lt no longer ordered")
|
||||
def test_compare_alu_same_src_different_arg(self):
|
||||
a = UOp.const(2.0)
|
||||
b = UOp.const(3.0)
|
||||
|
||||
add = UOp(Ops.ADD, src=(a, b))
|
||||
mul = UOp(Ops.MUL, src=(a, b))
|
||||
assert (add < mul) or (mul < add), "add and mul with same src should have an order"
|
||||
|
||||
def test_uop_variables(self):
|
||||
a = UOp.variable("a", 1, 10)
|
||||
uop_var = Tensor(a.bind(1))
|
||||
st_var = Tensor.empty((2, 10))[:, :a.bind(1)]
|
||||
_, var_vals = (uop_var+st_var).linear_with_vars()
|
||||
self.assertEqual(len(var_vals), 1)
|
||||
self.assertEqual(list(var_vals)[0], a.expr)
|
||||
|
||||
def test_const_factor(self):
|
||||
gidx0 = UOp(Ops.SPECIAL, src=(UOp.const(8),), arg='gidx0')
|
||||
self.assertEqual(UOp.const(17).const_factor(), 17)
|
||||
self.assertEqual(gidx0.const_factor(), 1)
|
||||
self.assertEqual((gidx0*3).const_factor(), 3)
|
||||
self.assertEqual((gidx0*3+6).const_factor(), 3)
|
||||
self.assertEqual((gidx0*3+1).const_factor(), 1)
|
||||
|
||||
def test_cmp_self_folding_multidim(self):
|
||||
for shape in ((), (3,), (2, 3), (2, 3, 4)):
|
||||
x = Tensor.empty(*shape, dtype=dtypes.int).uop
|
||||
self.assertIs((x < x).simplify(), x.const_like(False, dtypes.bool))
|
||||
self.assertIs((x != x).simplify(), x.const_like(False, dtypes.bool))
|
||||
|
||||
def test_replace(self):
|
||||
x = UOp.param(0, dtypes.int, (1,))
|
||||
self.assertEqual(x.replace(arg=UOp.param(1, dtypes.int, (1,)).arg).arg.slot, 1)
|
||||
with self.assertRaises(AssertionError): x.replace(field="a")
|
||||
|
||||
def test_const_zero_neg_zero_different(self):
|
||||
# -0.0 and 0.0 must be different UOps (for IEEE754 correctness, e.g. 1/-0.0 = -inf)
|
||||
pos_zero = UOp.const(0.0)
|
||||
neg_zero = UOp.const(-0.0)
|
||||
self.assertIsNot(pos_zero, neg_zero)
|
||||
self.assertNotEqual(hash(pos_zero.arg), hash(neg_zero.arg))
|
||||
|
||||
def test_const_nan_same(self):
|
||||
# nan constants should be deduplicated
|
||||
nan1 = UOp.const(float('nan'))
|
||||
nan2 = UOp.const(float('nan'))
|
||||
self.assertIs(nan1, nan2)
|
||||
|
||||
class TestUOpStr(unittest.TestCase):
|
||||
def test_uop_str(self):
|
||||
a = UOp.const(2.0) + UOp.const(3.0)
|
||||
for _ in range(20): a = a + a
|
||||
assert len(str(a)) < 10_000, "exponential string growth"
|
||||
assert str(eval(str(a))) == str(a)
|
||||
|
||||
def test_vectorized_str(self):
|
||||
vec = UOp(Ops.STACK, src=tuple(UOp.const(x) for x in range(4)))
|
||||
assert str(eval(str(vec))) == str(vec)
|
||||
|
||||
def test_reduceop_arg(self):
|
||||
sum_uop = Tensor.empty(32, 32).sum().uop
|
||||
assert str(eval(str(sum_uop))) == str(sum_uop)
|
||||
|
||||
class TestUPatHelpers(unittest.TestCase):
|
||||
def test_location(self):
|
||||
self.assertEqual(sym.patterns[-1][0].location[0].replace("\\", "/").split("/")[-1], "symbolic.py")
|
||||
self.assertEqual(spec_shared.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "spec.py")
|
||||
test_upat = UPat(Ops.CONST, dtypes.bool)
|
||||
self.assertEqual(test_upat.location[0].replace("\\", "/").split("/")[-1], __file__.replace("\\", "/").split("/")[-1])
|
||||
test_upat_named = test_upat.named("test_name")
|
||||
self.assertEqual(test_upat.location[0], test_upat_named.location[0])
|
||||
self.assertNotEqual(test_upat.location[1], test_upat_named.location[1])
|
||||
|
||||
class TestUopsObject(unittest.TestCase):
|
||||
def test_timing(self):
|
||||
with Timing("create 10k uops:"): ret = [UOp.const(10000000+i, dtypes.int) for i in range(10000)]
|
||||
assert len(ret) == 10000
|
||||
|
||||
def test_nested(self):
|
||||
a = UOp.new_buffer(Device.DEFAULT, 1, dtypes.char)
|
||||
for _ in range(10_000): a = a+a
|
||||
self.assertEqual(a.device, Device.DEFAULT)
|
||||
|
||||
class TestUOpRender(unittest.TestCase):
|
||||
def test_render_vectorize_empty(self):
|
||||
u = UOp(Ops.STACK, dtype=dtypes.void, src=())
|
||||
self.assertEqual(u.render(simplify=False), "{}")
|
||||
def test_render_vectorize_empty_simplified(self):
|
||||
u = UOp(Ops.STACK, dtype=dtypes.void, src=())
|
||||
self.assertEqual(u.render(), "{}")
|
||||
def test_render_vectorize_same(self):
|
||||
u = UOp(Ops.STACK, src=(UOp.const(0),)*3)
|
||||
self.assertEqual(u.render(simplify=False), "{0,0,0}")
|
||||
def test_render_vectorize_different(self):
|
||||
u = UOp(Ops.STACK, src=tuple(UOp.const(i) for i in range(3)))
|
||||
self.assertEqual(u.render(simplify=False), "{0,1,2}")
|
||||
def test_render_vectorize_same_simplified(self):
|
||||
u = UOp(Ops.STACK, src=(UOp.const(0),)*3)
|
||||
self.assertEqual(u.render(), "{0,0,0}")
|
||||
def test_render_vectorize_different_simplified(self):
|
||||
u = UOp(Ops.STACK, src=tuple(UOp.const(i) for i in range(3)))
|
||||
self.assertEqual(u.render(), "{0,1,2}")
|
||||
|
||||
class TestContiguousViewOffset(unittest.TestCase):
|
||||
def _check(self, u, expected): self.assertEqual(u.contiguous_view_offset(), expected)
|
||||
|
||||
def test_simple(self): self._check(UOp.empty(10), 0)
|
||||
def test_shrink(self): self._check(UOp.empty(10)[1:8], 1)
|
||||
def test_2d(self): self._check(UOp.empty(2,5)[1, 2:4], 7)
|
||||
def test_shrink_to_one(self): self._check(UOp.empty(10)[1], 1)
|
||||
def test_expand_is_none(self): self._check(UOp.empty(1).expand(2), None)
|
||||
def test_shrink_invalid(self): self._check(UOp.empty(4).pad((2,2))[0], None)
|
||||
def test_strided(self): self._check(UOp.empty(4)[::2], None)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
257
artifacts/package_sources/tinygrad/test/null/test_uops_stats.py
Normal file
257
artifacts/package_sources/tinygrad/test/null/test_uops_stats.py
Normal file
@@ -0,0 +1,257 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor
|
||||
from tinygrad.helpers import GlobalCounters
|
||||
from tinygrad.engine.realize import compile_linear, estimate_uop
|
||||
from tinygrad.codegen import to_program
|
||||
from tinygrad.renderer import Estimates
|
||||
from tinygrad.uop.ops import Ops, UOp
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.renderer.ptx import PTXRenderer
|
||||
from test.helpers import replace_opts
|
||||
|
||||
def flops_mem(uops, ignore_indexing=False):
|
||||
est = Estimates.from_uops(uops, ignore_indexing)
|
||||
return est.ops, est.lds
|
||||
|
||||
# **************** new FlopCounter ****************
|
||||
|
||||
def get_stats(x:Tensor):
|
||||
est = estimate_uop(compile_linear(x.schedule_linear()).src[-1])
|
||||
return est.ops, est.mem
|
||||
|
||||
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "webgpu does extra load/store for packed types")
|
||||
class TestMemoryCount(unittest.TestCase):
|
||||
def test_add(self):
|
||||
a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)
|
||||
b = Tensor.empty(1024, 1024, dtype=dtypes.uint8)
|
||||
_, mem = get_stats(a+b)
|
||||
self.assertEqual(mem, 1024*1024*3) # 2 reads + 1 write
|
||||
|
||||
def test_add_const(self):
|
||||
a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)
|
||||
_, mem = get_stats(a+3)
|
||||
self.assertEqual(mem, 1024*1024*2) # 1 read + 1 write
|
||||
|
||||
@unittest.skip("depends on subbuffer working")
|
||||
def test_add_slice(self):
|
||||
a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)[:512]
|
||||
_, mem = get_stats(a+3)
|
||||
self.assertEqual(mem, 512*1024*2) # 1 read + 1 write
|
||||
|
||||
def test_expanded(self):
|
||||
a = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)
|
||||
b = Tensor.empty(1024, 1024, dtype=dtypes.uint8)
|
||||
_, mem = get_stats(a+b)
|
||||
self.assertEqual(mem, 1024*1024*2 + 1024) # 1 full read + 1 lil read + 1 write
|
||||
|
||||
@unittest.skip("no longer supported")
|
||||
def test_both_expanded(self):
|
||||
# TODO: this probably should be a full write
|
||||
a = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)
|
||||
b = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)
|
||||
_, mem = get_stats(a+b)
|
||||
# rangeify is smart!
|
||||
self.assertEqual(mem, 1024 + 2*1024) # 2 lil reads + 1 lil write
|
||||
|
||||
def test_self_add(self):
|
||||
a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)
|
||||
_, mem = get_stats(a+a)
|
||||
self.assertEqual(mem, 1024*1024*2) # 1 read + 1 write
|
||||
|
||||
def test_self_add_transposed(self):
|
||||
a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)
|
||||
_, mem = get_stats(a+a.T)
|
||||
self.assertEqual(mem, 1024*1024*2) # 1 read + 1 write
|
||||
|
||||
def test_self_add_assign(self):
|
||||
a = Tensor.empty(1024, 1024, dtype=dtypes.uint8).realize()
|
||||
_, mem = get_stats(a.assign(a+a))
|
||||
self.assertEqual(mem, 1024*1024*2) # 1 read + 1 write
|
||||
|
||||
@unittest.skipIf(Device.DEFAULT == "CPU", "test copy to CPU from other device")
|
||||
def test_copyout(self):
|
||||
a = Tensor.empty(32, dtype=dtypes.uint8).to("CPU")
|
||||
_, mem = get_stats(a)
|
||||
self.assertEqual(mem, 32*1)
|
||||
a = Tensor.empty(32, dtype=dtypes.uint32).to("CPU")
|
||||
_, mem = get_stats(a)
|
||||
self.assertEqual(mem, 32*4)
|
||||
|
||||
# NOTE: this still isn't testing unroll using the acc
|
||||
@unittest.skipUnless(Device.DEFAULT == "PYTHON", "only run test on emulated tensor cores")
|
||||
class TestUOpsStatsMatmulHalf(unittest.TestCase):
|
||||
def test_simple_matmul_half(self, N=16):
|
||||
GlobalCounters.reset()
|
||||
a, b = Tensor.empty(N, N, dtype=dtypes.half), Tensor.empty(N, N, dtype=dtypes.half)
|
||||
c = a.matmul(b)
|
||||
c.realize()
|
||||
expected_ops = N ** 3 * 2
|
||||
self.assertEqual(expected_ops, GlobalCounters.global_ops)
|
||||
|
||||
def test_bigger_matmul_half(self): self.test_simple_matmul_half(64)
|
||||
|
||||
def test_batched_matmul_half(self, N=16):
|
||||
GlobalCounters.reset()
|
||||
a, b = Tensor.empty(4, N, N, dtype=dtypes.half), Tensor.empty(1, N, N, dtype=dtypes.half)
|
||||
c = a.matmul(b)
|
||||
c.realize()
|
||||
expected_ops = 4 * N ** 3 * 2
|
||||
self.assertEqual(expected_ops, GlobalCounters.global_ops)
|
||||
|
||||
class TestUOpsStats(unittest.TestCase):
|
||||
def test_simple_add(self):
|
||||
a = Tensor.empty(100,100)
|
||||
b = Tensor.empty(100,100)
|
||||
c = a+b
|
||||
ops, mem = get_stats(c)
|
||||
expected_ops = c.numel()
|
||||
expected_mem = a.nbytes() + b.nbytes() + c.nbytes()
|
||||
self.assertEqual(mem, expected_mem)
|
||||
# NOTE; ops also include indexing ops
|
||||
assert expected_ops <= ops and ops <= expected_ops * 2
|
||||
|
||||
def test_simple_add_sq(self):
|
||||
a = Tensor.empty(100,100)
|
||||
b = Tensor.empty(100,100)
|
||||
c = (a+b)*(a+b)
|
||||
ops, mem = get_stats(c)
|
||||
expected_ops = c.numel()*2
|
||||
expected_mem = a.nbytes() + b.nbytes() + c.nbytes()
|
||||
self.assertEqual(mem, expected_mem)
|
||||
# NOTE; ops also include indexing ops
|
||||
assert expected_ops <= ops and ops <= expected_ops * 2
|
||||
|
||||
def test_cat_equal_pieces(self):
|
||||
# concatenating equal-size pieces lowers to STACK: pure data movement, no arithmetic
|
||||
equal = [Tensor.empty(256, 128) for _ in range(4)]
|
||||
self.assertEqual(get_stats(Tensor.cat(*equal, dim=1))[0], 0)
|
||||
# a mismatched piece falls back to pad+usum, which sums N zero-padded copies and pays their adds
|
||||
unequal = equal[:3] + [Tensor.empty(256, 129)]
|
||||
self.assertGreater(get_stats(Tensor.cat(*unequal, dim=1))[0], 0)
|
||||
|
||||
def test_simple_matmul(self, M=1024, N=1024, K=1024):
|
||||
a = Tensor.empty(M,N)
|
||||
b = Tensor.empty(N,K)
|
||||
c = a@b
|
||||
ops, mem = get_stats(c)
|
||||
expected_ops = c.numel() * N * 2
|
||||
required_mem = a.nbytes() + b.nbytes() + c.nbytes()
|
||||
assert expected_ops <= ops and ops <= expected_ops * 1.2
|
||||
# NOTE: it's hard to assert on the memory here, all depends on caching
|
||||
assert required_mem <= mem
|
||||
|
||||
def test_simple_matmul_8192(self): self.test_simple_matmul(8192, 8192, 8192)
|
||||
|
||||
#MULACC should have the same stats as MUL + ADD
|
||||
def test_mulacc(self):
|
||||
globl = UOp.param(0, dtypes.int, (3,))
|
||||
o1 = UOp.const(1, dtypes.int)
|
||||
o2 = UOp.const(2, dtypes.int)
|
||||
u1 = globl.index(o1)
|
||||
u2 = globl.index(o2)
|
||||
u3 = UOp.const(3, dtypes.int)
|
||||
u4 = UOp(Ops.MUL, src=(u1,u2))
|
||||
u5 = UOp(Ops.ADD, src=(u4,u3))
|
||||
uops = tuple(u5.toposort())
|
||||
|
||||
globl = UOp.param(0, dtypes.int, (3,))
|
||||
o1 = UOp.const(1, dtypes.int)
|
||||
o2 = UOp.const(2, dtypes.int)
|
||||
u1 = globl.index(o1)
|
||||
u2 = globl.index(o2)
|
||||
u3 = UOp.const(3, dtypes.int)
|
||||
u4 = UOp(Ops.MULACC, src=(u1,u2,u3))
|
||||
uops_fma = tuple(u4.toposort())
|
||||
|
||||
self.assertEqual(flops_mem(uops), flops_mem(uops_fma))
|
||||
|
||||
N = 64
|
||||
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "wrong in PTX") # maybe?
|
||||
class TestStatsOptimized(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.ast_gemm = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule_linear().src[-1].src[0]
|
||||
cls.ast_gemm_half = (Tensor.empty(N, N, dtype=dtypes.half) @ Tensor.empty(N, N, dtype=dtypes.half)).schedule_linear().src[-1].src[0]
|
||||
cls.ast_reduce = (Tensor.empty(N*N).sum()).schedule_linear().src[-1].src[0]
|
||||
|
||||
def check_gemm(self, p:UOp, extra_flops=0, half=False):
|
||||
est = p.src[0].arg.estimates
|
||||
print(p.arg.name, est.ops, est.mem, est.lds)
|
||||
self.assertEqual(est.ops, 2*N*N*N + extra_flops) # N**3 mulaccs
|
||||
self.assertEqual(est.mem, 3*N*N*(2 if half else 4)) # 3 NxN mats with floats
|
||||
|
||||
def test_gemm(self):
|
||||
p = to_program(replace_opts(self.ast_gemm, []), renderer=Device[Device.DEFAULT].renderer)
|
||||
self.check_gemm(p)
|
||||
self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4 + 4*N*N)
|
||||
|
||||
@unittest.skip("fails locally on AMD")
|
||||
def test_gemm_tc_unroll_half(self):
|
||||
try:
|
||||
p = to_program(replace_opts(self.ast_gemm_half, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)]),
|
||||
renderer=Device[Device.DEFAULT].renderer)
|
||||
except KernelOptError:
|
||||
raise unittest.SkipTest("no tensor cores")
|
||||
print(p.src[2].arg)
|
||||
self.check_gemm(p, half=True)
|
||||
|
||||
def test_gemm_tc_unroll(self):
|
||||
try:
|
||||
p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)]),
|
||||
renderer=Device[Device.DEFAULT].renderer)
|
||||
except KernelOptError:
|
||||
raise unittest.SkipTest("no tensor cores")
|
||||
print(p.src[2].arg)
|
||||
self.check_gemm(p)
|
||||
|
||||
# this is a good lesson about why UPCASTing is a good idea
|
||||
|
||||
def test_gemm_one_upcasted(self):
|
||||
p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4)]), renderer=Device[Device.DEFAULT].renderer)
|
||||
self.check_gemm(p)
|
||||
self.assertEqual(p.src[0].arg.estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N)
|
||||
|
||||
def test_gemm_upcasted(self):
|
||||
p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)]),
|
||||
renderer=Device[Device.DEFAULT].renderer)
|
||||
self.check_gemm(p)
|
||||
self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
|
||||
|
||||
def test_gemm_upcasted_locals(self):
|
||||
try:
|
||||
p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 0, 4),
|
||||
Opt(OptOps.LOCAL, 1, 4)]), renderer=Device[Device.DEFAULT].renderer)
|
||||
except KernelOptError:
|
||||
raise unittest.SkipTest("no locals")
|
||||
self.check_gemm(p)
|
||||
self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
|
||||
|
||||
def test_gemm_group(self):
|
||||
try:
|
||||
p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.GROUP, 0, 4)]), renderer=Device[Device.DEFAULT].renderer)
|
||||
except KernelOptError:
|
||||
raise unittest.SkipTest("no locals")
|
||||
SZ = N*N*4
|
||||
# NOTE: these are sort of wrong. they aren't honoring the IF statement
|
||||
self.check_gemm(p, extra_flops=SZ*5)
|
||||
self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4 + SZ*4 + (SZ*4 + 4*N*N)*4)
|
||||
|
||||
def test_reduce(self):
|
||||
p = to_program(replace_opts(self.ast_reduce, []), renderer=Device[Device.DEFAULT].renderer)
|
||||
est = p.src[0].arg.estimates
|
||||
print(p.arg.name, est.ops, est.mem, est.lds)
|
||||
self.assertEqual(est.ops, N*N)
|
||||
self.assertEqual(est.mem, N*N*4 + 4)
|
||||
|
||||
def test_reduce_group(self):
|
||||
try:
|
||||
p = to_program(replace_opts(self.ast_reduce, [Opt(OptOps.GROUP, 0, 50)]), renderer=Device[Device.DEFAULT].renderer)
|
||||
except KernelOptError:
|
||||
raise unittest.SkipTest("no locals")
|
||||
est = p.src[0].arg.estimates
|
||||
print(p.arg.name, est.ops, est.mem, est.lds)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main(verbosity=2)
|
||||
@@ -0,0 +1,54 @@
|
||||
import unittest
|
||||
from tinygrad.helpers import DEBUG, Context
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.uop.ops import UPat, rewrite_group, GroupOp, Ops
|
||||
from tinygrad.uop.upat import _get_code, upat_compile
|
||||
import dis
|
||||
|
||||
@rewrite_group()
|
||||
def do_compile(up):
|
||||
print("\n***** COMPILE", up)
|
||||
match_code = _get_code(up, False)
|
||||
match = upat_compile(up, lambda **kwargs: None)
|
||||
print(match_code[0])
|
||||
if DEBUG >= 2: dis.dis(match)
|
||||
return match_code[0]
|
||||
|
||||
@Context(SPEC=0)
|
||||
class TestUPatCompile(unittest.TestCase):
|
||||
def test_double(self):
|
||||
up = UPat.var("x") * UPat.cvar("c0") + UPat.var("x") * UPat.cvar("c1")
|
||||
do_compile(up)
|
||||
|
||||
def test_single(self):
|
||||
up = UPat.var("x") + UPat.var("y")
|
||||
do_compile(up)
|
||||
|
||||
def test_xpx(self):
|
||||
up = UPat.var("x") + UPat.var("x")
|
||||
do_compile(up)
|
||||
|
||||
def test_xp0(self):
|
||||
up = UPat.var("x") + 0
|
||||
do_compile(up)
|
||||
|
||||
def test_bool(self):
|
||||
up = UPat.var('x', dtype=dtypes.bool) * UPat.var('y', dtype=dtypes.bool)
|
||||
do_compile(up)
|
||||
|
||||
def test_single_c(self):
|
||||
up = (UPat.var("x") + UPat.var("y")) * UPat.var("c")
|
||||
do_compile(up)
|
||||
|
||||
def test_const_folding(self):
|
||||
up = UPat(GroupOp.ALU-{Ops.THREEFRY}, name="a", src=UPat((Ops.CONST, Ops.STACK)))
|
||||
do_compile(up)
|
||||
|
||||
@unittest.skip("fix this")
|
||||
def test_range_named(self):
|
||||
# this should be one src, but this should also still work
|
||||
up = UPat(Ops.CAST, dtypes.float, UPat.var("x", dtypes.bfloat16))
|
||||
do_compile(up)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,194 @@
|
||||
import unittest
|
||||
from tinygrad import dtypes, Variable
|
||||
from tinygrad.dtype import AddrSpace
|
||||
from tinygrad.helpers import Context
|
||||
from tinygrad.uop.ops import Ops, UOp, AxisType
|
||||
from test.helpers import to_uops_list
|
||||
|
||||
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(dtype=dtypes.int)]) # valid
|
||||
to_uops_list([buf.index(UOp.const(15)).load(dtype=dtypes.int)]) # valid (last element)
|
||||
with self.assertRaises(RuntimeError):
|
||||
to_uops_list([buf.index(UOp.const(16)).load(dtype=dtypes.int)]) # off by one
|
||||
with self.assertRaises(RuntimeError):
|
||||
to_uops_list([buf.index(UOp.const(42)).load(dtype=dtypes.int)]) # 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(dtype=dtypes.int)]) # valid
|
||||
with self.assertRaises(RuntimeError):
|
||||
to_uops_list([buf.index(Variable("i", 0, 20)).load(dtype=dtypes.int)]) # oob
|
||||
with self.assertRaises(RuntimeError):
|
||||
to_uops_list([buf.index(Variable("i", -5, 10)).load(dtype=dtypes.int)]) # 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(dtype=dtypes.int)]) # valid
|
||||
with self.assertRaises(RuntimeError):
|
||||
to_uops_list([buf.index(r.valid(r < 17)).load(dtype=dtypes.int)]) # 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(dtype=dtypes.int)]) # valid
|
||||
with self.assertRaises(RuntimeError):
|
||||
to_uops_list([buf.index(v.valid(v < 20)).load(dtype=dtypes.int)]) # 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(dtype=dtypes.int)]) # 0..15 valid
|
||||
with self.assertRaises(RuntimeError):
|
||||
to_uops_list([buf.index(UOp.range(34, 0, AxisType.GLOBAL) // 2).load(dtype=dtypes.int)]) # 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(dtype=dtypes.int)]) # 0..15 valid
|
||||
with self.assertRaises(RuntimeError):
|
||||
to_uops_list([buf.index(r % 20).load(dtype=dtypes.int)]) # 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(dtype=dtypes.int)]) # 0..15 valid
|
||||
with self.assertRaises(RuntimeError):
|
||||
to_uops_list([buf.index(UOp.range(128, 0, AxisType.GLOBAL) >> 2).load(dtype=dtypes.int)]) # 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(dtype=dtypes.int)]) # 0..28 valid
|
||||
with self.assertRaises(RuntimeError):
|
||||
to_uops_list([buf.index(r << 4).load(dtype=dtypes.int)]) # 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(dtype=dtypes.int)]) # 0..15 valid
|
||||
with self.assertRaises(RuntimeError):
|
||||
to_uops_list([buf.index(r & 31).load(dtype=dtypes.int)]) # 0..31 oob
|
||||
# align masks round down to a multiple of 2^k
|
||||
to_uops_list([buf.index((r & -4).valid(r < 16)).load(dtype=dtypes.int)]) # 0..12 valid
|
||||
with self.assertRaises(RuntimeError):
|
||||
to_uops_list([buf.index(r & -2).load(dtype=dtypes.int)]) # 0..100 oob
|
||||
# other masks can't be modeled as mod
|
||||
with self.assertRaisesRegex(RuntimeError, "z3 int AND only supports"):
|
||||
to_uops_list([buf.index(r & 21).load(dtype=dtypes.int)])
|
||||
|
||||
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(dtype=dtypes.int)]) # 0..15 valid
|
||||
with self.assertRaises(RuntimeError):
|
||||
to_uops_list([buf.index(Variable("v2", -10, 20).maximum(0)).load(dtype=dtypes.int)]) # 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(dtype=dtypes.int)]) # 0..15 valid
|
||||
with self.assertRaises(RuntimeError):
|
||||
to_uops_list([buf.index(r.valid((r < 10) ^ (r >= 20))).load(dtype=dtypes.int)]) # 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(dtype=dtypes.int)])
|
||||
|
||||
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(dtype=dtypes.int)]) # 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(dtype=dtypes.int).cast(dtypes.weakint)
|
||||
to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32))).load(dtype=dtypes.int)]) # valid
|
||||
with self.assertRaises(RuntimeError):
|
||||
to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 64))).load(dtype=dtypes.int)]) # 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(dtype=dtypes.int).index(0)
|
||||
to_uops_list([buf1.index(ld0.valid((ld0 >= 0) & (ld0 < 64))).load(dtype=dtypes.int)])
|
||||
|
||||
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
|
||||
|
||||
# skipped tests (moved from test_uop_graph.py)
|
||||
@unittest.skip("if not allowed in graph")
|
||||
def test_in_bounds_access_gated_local(self):
|
||||
with Context(CHECK_OOB=1):
|
||||
# Define buffers
|
||||
gbuf = UOp.param(0, dtypes.uint, (400,))
|
||||
sbuf = UOp.placeholder((8,), dtypes.uint, slot=0, addrspace=AddrSpace.LOCAL)
|
||||
|
||||
# Define indices, valids and barrier
|
||||
gidx = UOp(Ops.SPECIAL, src=(UOp.const(416),), arg="gidx0")
|
||||
lidx = UOp(Ops.SPECIAL, src=(UOp.const(10),), arg="lidx0")
|
||||
|
||||
gate = (gidx<400) & (lidx<8)
|
||||
|
||||
local_store = sbuf.index(lidx.valid(lidx<8)).store(UOp.const(1))
|
||||
|
||||
barrier = UOp(Ops.BARRIER, src=(local_store,))
|
||||
if_barrier = UOp(Ops.IF, src=(gate, barrier))
|
||||
|
||||
# Load from local memory (after the IF/barrier)
|
||||
local_load = UOp(Ops.LOAD, src=(sbuf.index(lidx), if_barrier))
|
||||
|
||||
# Store to global memory
|
||||
global_store = UOp(Ops.STORE, src=(gbuf.index(gidx), local_load))
|
||||
to_uops_list([global_store])
|
||||
|
||||
@unittest.skip("Bool load is not supported yet")
|
||||
def test_load_mask(self):
|
||||
with Context(CHECK_OOB=1):
|
||||
glbl0 = UOp.param(0, dtypes.int, (16,))
|
||||
mask = UOp.param(0, dtypes.bool, (16,))
|
||||
ridx = UOp.range(20, 0)
|
||||
ld0 = UOp(Ops.LOAD, src=(glbl0.index(UOp.const(ridx<16&mask, ridx))))
|
||||
to_uops_list([ld0])
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
1121
artifacts/package_sources/tinygrad/test/null/test_viz.py
Normal file
1121
artifacts/package_sources/tinygrad/test/null/test_viz.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,57 @@
|
||||
import unittest, sys
|
||||
from tinygrad import Tensor, GlobalCounters, dtypes, Context
|
||||
from tinygrad.helpers import WINO
|
||||
from test.helpers import check_schedule
|
||||
|
||||
@unittest.skipIf(sys.platform.startswith("win"), "flaky on Windows")
|
||||
class TestWinograd(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.old = WINO.value
|
||||
WINO.value = 1
|
||||
def tearDown(self):
|
||||
WINO.value = self.old
|
||||
|
||||
def test_forward_kernels(self):
|
||||
x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize()
|
||||
out = Tensor.conv2d(x,w)
|
||||
check_schedule(out, 4)
|
||||
|
||||
def test_backward_counters(self):
|
||||
# contiguous_backward on the pooled input keeps the input-transform adjoint out of the overlap accumulation, so
|
||||
# winograd backward runs in a fraction of the direct-conv flops; NOOPT=1 keeps the raw flop ratio from drifting with the optimizer
|
||||
IC, OC, H = 64, 64, 28
|
||||
x,w = Tensor.empty(1,IC,H,H,device="NULL").realize(), Tensor.empty(OC,IC,3,3,device="NULL").realize()
|
||||
x.requires_grad = w.requires_grad = True
|
||||
def backward_ops(wino):
|
||||
x.grad = w.grad = None
|
||||
GlobalCounters.reset()
|
||||
with Context(NOOPT=1, WINO=wino):
|
||||
Tensor.conv2d(x,w,padding=1).mean().backward()
|
||||
Tensor.realize(x.grad, w.grad)
|
||||
return GlobalCounters.global_ops
|
||||
ops_wino, ops_normal = backward_ops(1), backward_ops(0)
|
||||
print(f"backward ops: normal {ops_normal} wino {ops_wino} ratio {ops_wino/ops_normal:.2f}")
|
||||
self.assertLess(ops_wino/ops_normal, 0.35)
|
||||
|
||||
def test_counters(self):
|
||||
IC, OC, H = 64, 64, 28
|
||||
x,w = Tensor.empty(1,IC,H,H,device="NULL").realize(), Tensor.empty(OC,IC,3,3,device="NULL").realize()
|
||||
GlobalCounters.reset()
|
||||
with Context(NOOPT=0, WINO=1): Tensor.conv2d(x,w).realize()
|
||||
ops_wino = GlobalCounters.global_ops
|
||||
GlobalCounters.reset()
|
||||
with Context(NOOPT=0, WINO=0): Tensor.conv2d(x,w).realize()
|
||||
ops_normal = GlobalCounters.global_ops
|
||||
print(f"ops: normal {ops_normal} wino {ops_wino} ratio {ops_wino/ops_normal:.2f}")
|
||||
self.assertLess(ops_wino/ops_normal, 0.6)
|
||||
|
||||
def test_dtype(self):
|
||||
IC, OC, X, Y = 4,4,9,9
|
||||
x,w = Tensor.empty(1,IC,Y,X), Tensor.empty(OC,IC,3,3)
|
||||
self.assertEqual(Tensor.conv2d(x,w).dtype, dtypes.default_float)
|
||||
|
||||
x,w = Tensor.empty(1,IC,Y,X,dtype=dtypes.half), Tensor.empty(OC,IC,3,3,dtype=dtypes.half)
|
||||
self.assertEqual(Tensor.conv2d(x,w).dtype, dtypes.half)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main(verbosity=2)
|
||||
Reference in New Issue
Block a user