mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 13:33:42 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
674 lines
25 KiB
Python
674 lines
25 KiB
Python
import numpy as np
|
|
import unittest
|
|
from tinygrad.function import function
|
|
from tinygrad import Tensor, GlobalCounters, Device
|
|
from tinygrad.dtype import Invalid
|
|
from tinygrad.uop.ops import UOp, Ops, KernelInfo, ProgramInfo
|
|
from test.helpers import assert_kernel_count, KernelCountException
|
|
|
|
class TestFunction(unittest.TestCase):
|
|
def test_simple(self):
|
|
@function
|
|
def f(a:Tensor, b:Tensor) -> Tensor: return a+b
|
|
|
|
a = Tensor([1,2,3])
|
|
b = Tensor([4,5,6])
|
|
np.testing.assert_equal(f(a,b).numpy(), [5,7,9])
|
|
|
|
def test_two_return(self, precompile=False):
|
|
@function(precompile=precompile)
|
|
def f(a:Tensor, b:Tensor) -> tuple[Tensor, Tensor]:
|
|
return (a+b, (a+b)*2)
|
|
a = Tensor([1,2,3])
|
|
b = Tensor([4,5,6])
|
|
c = f(a,b)
|
|
np.testing.assert_equal((c[0]+c[1]).numpy(), [5*3,7*3,9*3])
|
|
def test_two_return_precompiled(self): self.test_two_return(True)
|
|
|
|
def test_simple_same(self):
|
|
@function
|
|
def f(a:Tensor, b:Tensor) -> Tensor: return a+b
|
|
|
|
a = Tensor([1,2,3])
|
|
np.testing.assert_equal(f(a,a).numpy(), [2,4,6])
|
|
|
|
def test_depth_restored_on_exception(self):
|
|
from tinygrad.function import _function
|
|
@function
|
|
def f(a:Tensor) -> Tensor: raise ValueError("error")
|
|
with self.assertRaises(ValueError): f(Tensor([1]))
|
|
self.assertEqual(_function.depth, 0)
|
|
|
|
def test_implicit(self):
|
|
inp = Tensor([7,8,9])
|
|
@function(allow_implicit=True)
|
|
def f(a:Tensor, b:Tensor) -> Tensor: return a+b+inp
|
|
|
|
a = Tensor([1,2,3])
|
|
b = Tensor([4,5,6])
|
|
np.testing.assert_equal(f(a,b).numpy(), [12,15,18])
|
|
|
|
def test_implicit_same_as_input(self):
|
|
inp = Tensor([7,8,9])
|
|
@function(allow_implicit=True)
|
|
def f(a:Tensor, b:Tensor) -> Tensor: return a+b+inp
|
|
|
|
a = Tensor([1,2,3])
|
|
np.testing.assert_equal(f(a, inp).numpy(), [15,18,21])
|
|
|
|
def test_implicit_2(self):
|
|
inp = Tensor([7,8,9])
|
|
@function(allow_implicit=True)
|
|
def f(a:Tensor, b:Tensor) -> Tensor:
|
|
return a+b+inp
|
|
inp2 = Tensor([7,8,10])
|
|
@function(allow_implicit=True)
|
|
def g(a:Tensor, b:Tensor) -> Tensor:
|
|
return a+b+inp2
|
|
|
|
a = Tensor([1,2,3])
|
|
b = Tensor([4,5,6])
|
|
c = f(a,b)
|
|
d = g(a,b)
|
|
c.realize(d)
|
|
np.testing.assert_equal(c.numpy(), [12,15,18])
|
|
np.testing.assert_equal(d.numpy(), [12,15,19])
|
|
|
|
def test_implicit_unrealized(self):
|
|
inp = Tensor([1,2,3]) + Tensor([4,5,6])
|
|
@function(allow_implicit=True)
|
|
def f(a:Tensor) -> Tensor: return a + inp
|
|
|
|
np.testing.assert_equal(f(Tensor([10,20,30])).numpy(), [15,27,39])
|
|
|
|
def test_detach(self):
|
|
@function
|
|
def f(a:Tensor, b:Tensor) -> Tensor: return a.detach() + b
|
|
|
|
a = Tensor([1,2,3])
|
|
b = Tensor([4,5,6])
|
|
np.testing.assert_equal(f(a, b).numpy(), [5,7,9])
|
|
|
|
def test_contiguous_backward(self):
|
|
@function
|
|
def f(a:Tensor, b:Tensor) -> Tensor: return (a + b).contiguous_backward()
|
|
|
|
a = Tensor([1,2,3])
|
|
b = Tensor([4,5,6])
|
|
np.testing.assert_equal(f(a, b).numpy(), [5,7,9])
|
|
|
|
def test_method(self):
|
|
class Foo:
|
|
def __init__(self): self.w = Tensor([10,20,30])
|
|
@function
|
|
def __call__(self, x:Tensor) -> Tensor: return x + self.w
|
|
|
|
foo = Foo()
|
|
np.testing.assert_equal(foo(Tensor([1,2,3])).numpy(), [11,22,33])
|
|
|
|
def test_grad_gemm(self):
|
|
@function
|
|
def f(a:Tensor, b:Tensor) -> Tensor: return a @ b
|
|
|
|
a = Tensor([[1.,2.],[3.,4.]])
|
|
b = Tensor([[5.,6.],[7.,8.]])
|
|
(f(a, b).contiguous() * b).sum().backward()
|
|
Tensor.realize(a, b, a.grad, b.grad)
|
|
# L = sum((a@b) * b), dL/d(a@b) = b, dL/da = b @ b^T, dL/db = a^T @ b + (a@b)
|
|
na, nb = a.numpy(), b.numpy()
|
|
np.testing.assert_allclose(a.grad.numpy(), nb @ nb.T)
|
|
np.testing.assert_allclose(b.grad.numpy(), na.T @ nb + na @ nb)
|
|
|
|
def test_grad_implicit(self):
|
|
w = Tensor([1., 2., 3.])
|
|
w.realize() # TODO: this is required
|
|
@function(allow_implicit=True)
|
|
def f(x:Tensor) -> Tensor: return x * w
|
|
|
|
x = Tensor([4., 5., 6.])
|
|
f(x).sum().backward()
|
|
np.testing.assert_allclose(w.grad.numpy(), [4., 5., 6.])
|
|
|
|
def test_symbolic_index(self):
|
|
table = Tensor([10,20,30,40]).contiguous().realize()
|
|
@function(allow_implicit=True)
|
|
def f(x:Tensor, start_pos:int|UOp) -> Tensor:
|
|
return x + table[start_pos]
|
|
|
|
v = UOp.variable("start_pos", 0, 3)
|
|
np.testing.assert_equal(f(Tensor([1,2,3]), v.bind(0)).numpy(), [11,12,13])
|
|
|
|
def test_symbolic_shape_input(self):
|
|
table = Tensor([10,20,30,40]).contiguous().realize()
|
|
@function
|
|
def f(x:Tensor) -> Tensor: return x * 2
|
|
sz = UOp.variable("sz", 1, 3)
|
|
slic = table[:sz.bind(2)]
|
|
np.testing.assert_equal(f(slic)[:2].numpy(), [20,40])
|
|
|
|
def test_nested_calls(self):
|
|
w = Tensor([10., 20., 30.])
|
|
@function(allow_implicit=True)
|
|
def f(a:Tensor) -> Tensor: return a + w
|
|
@function(allow_implicit=True)
|
|
def g(a:Tensor) -> Tensor: return a * w
|
|
|
|
a = Tensor([1., 2., 3.])
|
|
np.testing.assert_allclose(g(f(a)).numpy(), [110., 440., 990.])
|
|
|
|
def test_nested_calls_backward(self):
|
|
w = Tensor([[1., 2.], [3., 4.]]).contiguous().realize()
|
|
@function(allow_implicit=True)
|
|
def inner(x:Tensor) -> Tensor: return x + w
|
|
@function(allow_implicit=True)
|
|
def outer(a:Tensor, b:Tensor) -> Tensor: return inner(a.reshape(1,2) + b.reshape(1,2))
|
|
|
|
a = Tensor([1., 2.])
|
|
b = Tensor([3., 4.])
|
|
outer(a, b).sum().backward()
|
|
np.testing.assert_allclose(a.grad.numpy(), [2., 2.])
|
|
np.testing.assert_allclose(b.grad.numpy(), [2., 2.])
|
|
|
|
def test_unused_param_backward(self):
|
|
@function
|
|
def f(a:Tensor, b:Tensor, c:Tensor) -> Tensor: return a + c # b is unused
|
|
|
|
a = Tensor([1., 2., 3.])
|
|
b = Tensor([4., 5., 6.])
|
|
c = Tensor([7., 8., 9.])
|
|
f(a, b, c).sum().backward()
|
|
np.testing.assert_allclose(a.grad.numpy(), [1., 1., 1.])
|
|
np.testing.assert_allclose(b.grad.numpy(), [0., 0., 0.])
|
|
np.testing.assert_allclose(c.grad.numpy(), [1., 1., 1.])
|
|
|
|
def test_name(self):
|
|
@function
|
|
def f(a:Tensor) -> Tensor: return a + 1
|
|
assert f(Tensor([1])).uop.src[1].arg.name.endswith("f")
|
|
|
|
def test_method_name(self):
|
|
class Foo:
|
|
@function
|
|
def __call__(self, x:Tensor) -> Tensor: return x + 1
|
|
assert Foo()(Tensor([1])).uop.src[1].arg.name.endswith("Foo.__call__")
|
|
|
|
def test_callable_instance(self):
|
|
class Foo:
|
|
def __init__(self): self.w = Tensor([10,20,30])
|
|
def __call__(self, x:Tensor) -> Tensor: return x + self.w
|
|
foo = Foo()
|
|
f = function(foo, allow_implicit=True)
|
|
np.testing.assert_equal(f(Tensor([1,2,3])).numpy(), [11,22,33])
|
|
assert f(Tensor([1,2,3])).uop.src[1].arg.name.endswith("Foo")
|
|
|
|
def test_iadd(self):
|
|
@function
|
|
def f(x:Tensor) -> Tensor:
|
|
x += 1
|
|
return x
|
|
|
|
a = Tensor([1,2,3]).realize()
|
|
np.testing.assert_equal(f(a).numpy(), [2,3,4])
|
|
np.testing.assert_equal(a.numpy(), [3,4,5]) # TODO: should be [1,2,3]
|
|
|
|
def test_implicit_assign(self):
|
|
a = Tensor([1,2,3])
|
|
a += 1
|
|
c = Tensor([2,2,2]).contiguous()
|
|
@function
|
|
def f(b:Tensor) -> Tensor: return a+b+c
|
|
b = Tensor([10,20,30]).realize()
|
|
np.testing.assert_equal(f(b).numpy(), [14,25,36])
|
|
|
|
def test_assign_input(self):
|
|
@function
|
|
def f(a:Tensor, b:Tensor) -> Tensor:
|
|
a.assign(b+1)
|
|
return a
|
|
|
|
a = Tensor([1,2,3]).realize()
|
|
b = Tensor([10,20,30]).realize()
|
|
np.testing.assert_equal(f(a,b).numpy(), [11,21,31])
|
|
np.testing.assert_equal(a.numpy(), [11,21,31]) # TODO: should be [1,2,3]
|
|
np.testing.assert_equal(b.numpy(), [10,20,30])
|
|
|
|
def test_view_assign_explicit_buffer(self):
|
|
"""view assign on an explicit param's buffer should not create implicit inputs."""
|
|
class State:
|
|
def __init__(self): self.buf = Tensor.zeros(2, 4).contiguous().realize()
|
|
@function(allow_implicit=False)
|
|
def __call__(self, x:Tensor) -> Tensor:
|
|
self.buf[:, 0:2].assign(x)
|
|
return self.buf[:, 0:2]
|
|
s = State()
|
|
np.testing.assert_equal(s(Tensor([[5., 6.], [7., 8.]])).numpy(), [[5., 6.], [7., 8.]])
|
|
|
|
def test_single_after_store(self):
|
|
"""AFTER(buf, STORE(view, data)) should write data through the view into buf, same as the double-after pattern."""
|
|
@function
|
|
def f(buf:Tensor, x:Tensor, start_pos:int|UOp) -> Tensor:
|
|
slice_uop = buf[:, start_pos:start_pos+1].uop
|
|
assigned = Tensor(buf.uop.after(slice_uop.store(x.uop)))
|
|
return assigned
|
|
|
|
buf = Tensor.zeros(2, 8).contiguous().realize()
|
|
x = Tensor([[1.], [2.]]).realize()
|
|
v = UOp.variable("sp", 0, 7)
|
|
r0 = f(buf, x, v.bind(0)).numpy()
|
|
np.testing.assert_equal(r0, [[1.,0.,0.,0.,0.,0.,0.,0.], [2.,0.,0.,0.,0.,0.,0.,0.]])
|
|
|
|
def test_single_after_store_precompile(self):
|
|
"""precompiled AFTER(buf, STORE(view, data)) should return buf after the store."""
|
|
@function(precompile=True)
|
|
def f(buf:Tensor, x:Tensor, start_pos:int|UOp) -> Tensor:
|
|
slice_uop = buf[:, start_pos:start_pos+1].uop
|
|
assigned = Tensor(buf.uop.after(slice_uop.store(x.uop)))
|
|
return assigned
|
|
|
|
x = Tensor([[1.], [2.]]).realize()
|
|
v = UOp.variable("sp", 0, 7)
|
|
for sp in (0, 2):
|
|
with self.subTest(sp=sp):
|
|
buf = Tensor.zeros(2, 8).clone().realize()
|
|
expected = np.zeros((2, 8), dtype=np.float32)
|
|
expected[:, sp] = [1., 2.]
|
|
np.testing.assert_equal(f(buf, x, v.bind(sp)).numpy(), expected)
|
|
np.testing.assert_equal(buf.numpy(), expected)
|
|
|
|
@unittest.expectedFailure
|
|
def test_assign_slice(self):
|
|
@function
|
|
def f(a:Tensor, b:Tensor) -> Tensor:
|
|
a[1:] = b[1:]+1
|
|
return a
|
|
|
|
a = Tensor([1,2,3]).realize()
|
|
b = Tensor([10,20,30]).realize()
|
|
np.testing.assert_equal(f(a,b).numpy(), [1,21,31])
|
|
np.testing.assert_equal(a.numpy(), [1,2,3])
|
|
np.testing.assert_equal(b.numpy(), [10,20,30])
|
|
|
|
class TestFunctionMulti(unittest.TestCase):
|
|
devices_2 = ("CPU:0", "CPU:1")
|
|
|
|
def test_simple_multi(self):
|
|
@function
|
|
def f(a:Tensor, b:Tensor) -> Tensor: return a+b
|
|
|
|
a = Tensor([1,2,3,4]).shard(self.devices_2, axis=None)
|
|
b = Tensor([10,20,30,40]).shard(self.devices_2, axis=None)
|
|
np.testing.assert_equal(f(a,b).numpy(), [11,22,33,44])
|
|
|
|
def test_simple_multi_sharded(self):
|
|
@function
|
|
def f(a:Tensor, b:Tensor) -> Tensor: return a+b
|
|
|
|
a = Tensor([1,2,3,4]).shard(self.devices_2, axis=0)
|
|
b = Tensor([10,20,30,40]).shard(self.devices_2, axis=0)
|
|
np.testing.assert_equal(f(a,b).numpy(), [11,22,33,44])
|
|
|
|
def test_data_parallel_multi(self):
|
|
@function
|
|
def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
|
|
|
|
x = Tensor([[1.,2.],[3.,4.],[5.,6.],[7.,8.]]).shard(self.devices_2, axis=0)
|
|
w = Tensor([[1.,0.],[0.,1.]]).shard(self.devices_2, axis=None)
|
|
np.testing.assert_allclose(f(x, w).numpy(), [[1.,2.],[3.,4.],[5.,6.],[7.,8.]])
|
|
|
|
def test_grad_implicit_multi(self):
|
|
w = Tensor([1., 2., 3., 4.]).shard(self.devices_2, axis=None)
|
|
w.realize()
|
|
@function(allow_implicit=True)
|
|
def f(x:Tensor) -> Tensor: return x * w
|
|
|
|
x = Tensor([4., 5., 6., 7.]).shard(self.devices_2, axis=None)
|
|
f(x).sum().backward()
|
|
np.testing.assert_allclose(w.grad.numpy(), [4., 5., 6., 7.])
|
|
|
|
def test_call_axis(self):
|
|
@function
|
|
def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
|
|
|
|
x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]]).shard(self.devices_2, axis=0)
|
|
w = Tensor([[1.,2.],[3.,4.]]).shard(self.devices_2, axis=None)
|
|
result = f(x, w)
|
|
# CALL output should inherit axis=0 from the sharded input
|
|
self.assertEqual(result.uop.axis, 0)
|
|
# reduce on the sharded axis should remove it
|
|
self.assertIsNone(result.sum().uop.axis)
|
|
|
|
def test_call_axis_shard_inside(self):
|
|
@function
|
|
def f(x:Tensor, w:Tensor) -> Tensor:
|
|
return x.shard(self.devices_2, axis=0) @ w.shard(self.devices_2, axis=None)
|
|
|
|
x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]])
|
|
w = Tensor([[1.,2.],[3.,4.]])
|
|
result = f(x, w)
|
|
self.assertEqual(result.uop.axis, 0)
|
|
np.testing.assert_allclose(result.numpy(), x.numpy() @ w.numpy())
|
|
|
|
def test_data_parallel_backward(self):
|
|
@function
|
|
def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
|
|
|
|
x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]]).shard(self.devices_2, axis=0)
|
|
w = Tensor([[1.,2.],[3.,4.]]).shard(self.devices_2, axis=None)
|
|
w.realize()
|
|
f(x, w).sum().backward()
|
|
# d/dx = ones @ w^T = [[1,3],[1,3],[1,3],[1,3]], but sum so ones(4,2) @ w^T? no:
|
|
# L = sum(x @ w), dL/dx = ones(4,2) @ w^T... actually dL/d(xw) = ones(4,2), dL/dx = ones(4,2) @ w^T
|
|
np.testing.assert_allclose(x.grad.numpy(), np.ones((4,2)) @ np.array([[1,3],[2,4]]))
|
|
|
|
def test_data_parallel_backward_4(self):
|
|
devices_4 = tuple(f"CPU:{i}" for i in range(4))
|
|
@function
|
|
def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
|
|
|
|
x = Tensor(np.arange(16).reshape(8,2).astype(np.float32)).shard(devices_4, axis=0)
|
|
w = Tensor([[1.,2.],[3.,4.]]).shard(devices_4, axis=None)
|
|
w.realize()
|
|
f(x, w).sum().backward()
|
|
np.testing.assert_allclose(x.grad.numpy(), np.ones((8,2)) @ np.array([[1,3],[2,4]]))
|
|
|
|
def test_data_parallel_backward_implicit(self):
|
|
devices_4 = tuple(f"CPU:{i}" for i in range(4))
|
|
w = Tensor([[1.,2.],[3.,4.]]).shard(devices_4, axis=None)
|
|
w.realize()
|
|
@function(allow_implicit=True)
|
|
def f(x:Tensor) -> Tensor: return x @ w
|
|
|
|
x = Tensor(np.arange(16).reshape(8,2).astype(np.float32)).shard(devices_4, axis=0)
|
|
f(x).sum().backward()
|
|
np.testing.assert_allclose(x.grad.numpy(), np.ones((8,2)) @ np.array([[1,3],[2,4]]))
|
|
|
|
def test_data_parallel_backward_twice(self):
|
|
devices_4 = tuple(f"CPU:{i}" for i in range(4))
|
|
w = Tensor([[1.,2.],[3.,4.]]).shard(devices_4, axis=None)
|
|
w.realize()
|
|
# pre-init grads like the training loop does
|
|
w.grad = w.zeros_like().contiguous().realize()
|
|
@function(allow_implicit=True)
|
|
def f(x:Tensor) -> Tensor: return x @ w
|
|
|
|
expected = np.ones((8,2)) @ np.array([[1,3],[2,4]])
|
|
for _ in range(2):
|
|
x = Tensor(np.arange(16).reshape(8,2).astype(np.float32)).shard(devices_4, axis=0)
|
|
f(x).sum().backward()
|
|
np.testing.assert_allclose(x.grad.numpy(), expected)
|
|
|
|
class TestFunctionTuple(unittest.TestCase):
|
|
def test_tuple(self, precompile=False):
|
|
x = Tensor.ones(3).contiguous()
|
|
@function(precompile=precompile)
|
|
def f(t:Tensor): return (t+1, t+2)
|
|
t1, t2 = f(x)
|
|
t1.realize(t2)
|
|
print(t1.tolist(), t2.tolist())
|
|
assert t1.tolist() == [2,2,2]
|
|
assert t2.tolist() == [3,3,3]
|
|
def test_tuple_precompile(self): self.test_tuple(True)
|
|
|
|
def test_grad_tuple(self, precompile=False):
|
|
x = Tensor.ones(3).contiguous()
|
|
y = Tensor.ones(3).contiguous()
|
|
@function(precompile=precompile)
|
|
def f(u1:Tensor, u2:Tensor): return (u1+1, u2+2)
|
|
t1, t2 = f(x,y)
|
|
(t1+t2).sum().backward()
|
|
x.grad.realize(y.grad)
|
|
def test_grad_tuple_precompile(self): self.test_grad_tuple(True)
|
|
|
|
def test_grad_fxn_tuple(self):
|
|
# grad_fxn for tuple: receives one gradient per output as positional args
|
|
def grad_fxn(d_out0:UOp, d_out1:UOp, call:UOp):
|
|
# f(u1, u2) = (u1+1, u2+2)
|
|
# df/du1 = d_out0, df/du2 = d_out1
|
|
return (d_out0, d_out1)
|
|
|
|
x = Tensor.ones(3).contiguous()
|
|
y = Tensor.ones(3).contiguous()
|
|
@function(grad_fxn=grad_fxn)
|
|
def f(u1:Tensor, u2:Tensor): return (u1+1, u2+2)
|
|
t1, t2 = f(x, y)
|
|
(t1+t2).sum().backward()
|
|
np.testing.assert_allclose(x.grad.numpy(), [1., 1., 1.])
|
|
np.testing.assert_allclose(y.grad.numpy(), [1., 1., 1.])
|
|
|
|
def test_grad_fxn_more_outputs_than_inputs(self):
|
|
def grad_fxn(grad:UOp, call:UOp): return (grad,)
|
|
|
|
x = Tensor([2.]).contiguous()
|
|
@function(grad_fxn=grad_fxn)
|
|
def f(x:Tensor): return (x+1, x+2)
|
|
_, y = f(x)
|
|
self.assertEqual(y.sum().gradient(x)[0].item(), 1.0)
|
|
|
|
def test_grad_unused_tuple_output_recursive(self):
|
|
# only one output is used
|
|
@function(precompile=True, precompile_backward=True)
|
|
def f(x:Tensor, w:Tensor):
|
|
a = x @ w
|
|
b = (x @ w) * 2 # shares x@w with a
|
|
return (a, b)
|
|
|
|
x = Tensor([[1., 2.], [3., 4.]]).contiguous()
|
|
w = Tensor([[1., 0.], [0., 1.]]).contiguous()
|
|
Tensor.realize(x, w)
|
|
t1, _ = f(x, w)
|
|
t1.sum().backward()
|
|
Tensor.realize(x.grad, w.grad)
|
|
# only t1 = x @ w flows to loss; dL/dw = x.T @ ones(2,2)
|
|
np.testing.assert_allclose(w.grad.numpy(), np.array([[1., 2.], [3., 4.]]).T @ np.ones((2, 2)))
|
|
np.testing.assert_allclose(x.grad.numpy(), np.ones((2, 2)) @ np.array([[1., 0.], [0., 1.]]).T)
|
|
|
|
def test_custom_kernel_save_unused_output(self):
|
|
def my_kernel(C:UOp, D:UOp, A:UOp) -> UOp:
|
|
i = UOp.range(A.shape[0], 0)
|
|
j = UOp.range(D.shape[0], 1)
|
|
store_c = C[i].store(A[i] * 2.0).end(i)
|
|
store_d = D[j].store(A[j]).end(j)
|
|
return UOp.sink(store_c, store_d, arg=KernelInfo(name="my_kernel"))
|
|
|
|
def my_grad(d_c:UOp, call:UOp):
|
|
a_input = call.src[3]
|
|
return (None, None, (Tensor(d_c) * 2.0 + Tensor(a_input) * 0).uop)
|
|
|
|
@function(precompile=True, precompile_backward=True)
|
|
def f(a:Tensor):
|
|
c = Tensor.invalids(*a.shape, dtype=a.dtype, device=a.device)
|
|
d = Tensor.invalids(3, dtype=a.dtype, device=a.device)
|
|
c, d = Tensor.custom_kernel(c, d, a, fxn=my_kernel, grad_fxn=my_grad)[:2]
|
|
return c, d
|
|
|
|
a = Tensor([1., 2., 3., 4.]).contiguous()
|
|
Tensor.realize(a)
|
|
c, _ = f(a)
|
|
c.sum().backward()
|
|
Tensor.realize(a.grad)
|
|
np.testing.assert_allclose(a.grad.numpy(), [2., 2., 2., 2.])
|
|
|
|
def test_custom_kernel_both_outputs_used(self):
|
|
def my_kernel(C:UOp, D:UOp, A:UOp) -> UOp:
|
|
i = UOp.range(A.shape[0], 0)
|
|
store_c = C[i].store(A[i] * 2.0)
|
|
store_d = D[i].store(A[i] * 3.0)
|
|
return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name="my_kernel"))
|
|
|
|
def my_grad(d_c:UOp, d_d:UOp, call:UOp):
|
|
return (None, None, (Tensor(d_c) + Tensor(d_d)).uop)
|
|
|
|
@function(precompile=True, precompile_backward=True)
|
|
def f(a:Tensor):
|
|
c = Tensor.invalids(*a.shape, dtype=a.dtype, device=a.device)
|
|
d = Tensor.invalids(*a.shape, dtype=a.dtype, device=a.device)
|
|
c, d = Tensor.custom_kernel(c, d, a, fxn=my_kernel, grad_fxn=my_grad)[:2]
|
|
return (c, d)
|
|
|
|
a = Tensor([1., 2., 3., 4.]).contiguous()
|
|
Tensor.realize(a)
|
|
c, d = f(a)
|
|
(c.sum() + d.sum()).backward() # dL/da = (1 + 1) since grad_fxn passes d_combined through
|
|
Tensor.realize(a.grad)
|
|
np.testing.assert_allclose(a.grad.numpy(), [2., 2., 2., 2.])
|
|
|
|
def test_custom_kernel_precompile_no_copy_kernel(self):
|
|
def my_kernel(C:UOp, A:UOp) -> UOp:
|
|
i = UOp.range(A.shape[0], 0)
|
|
return C[i].store(A[i] * 2.0).end(i).sink(arg=KernelInfo(name="my_kernel"))
|
|
|
|
def my_grad(d_c:UOp, call:UOp):
|
|
return (None, (Tensor(d_c) * 2.0).uop)
|
|
|
|
@function(precompile=True, precompile_backward=True)
|
|
def f(a:Tensor):
|
|
c = Tensor.invalids(*a.shape, dtype=a.dtype, device=a.device)
|
|
c = Tensor.custom_kernel(c, a, fxn=my_kernel, grad_fxn=my_grad)[0]
|
|
return c
|
|
|
|
def count_kernels(t:Tensor):
|
|
linear, _ = t.linear_with_vars()
|
|
return sum((len(call.device) if isinstance(call.device, tuple) else 1)
|
|
for call in linear.src if call.src[0].op is Ops.SINK)
|
|
|
|
a = Tensor([1., 2., 3., 4.]).contiguous()
|
|
Tensor.realize(a)
|
|
c = f(a)
|
|
|
|
if count_kernels(c) != 1: raise KernelCountException(1, count_kernels(c))
|
|
|
|
c.sum().backward()
|
|
Tensor.realize(a.grad)
|
|
np.testing.assert_allclose(a.grad.numpy(), [2., 2., 2., 2.])
|
|
|
|
def test_custom_kernel_precompile_multidevice(self):
|
|
# a custom_kernel output placeholder (invalids) under multi-device @function(precompile=True) must return the
|
|
# kernel's computed result. read it back through .numpy() so the cross-device gather reads the output buffer
|
|
devs = ("CPU:0", "CPU:1")
|
|
def double_kernel(C:UOp, A:UOp) -> UOp:
|
|
C, A = C.flatten(), A.flatten()
|
|
i = UOp.range(A.numel(), 0)
|
|
return C[i].store(A[i] * 2.0).end(i).sink(arg=KernelInfo(name="double_kernel"))
|
|
def double_grad(d_c:UOp, call:UOp): return (None, (Tensor(d_c) * 2.0).uop)
|
|
|
|
a = Tensor.full((4, 4), 7.0).contiguous().shard(devs, axis=0).realize()
|
|
|
|
@function(precompile=True, precompile_backward=True)
|
|
def f(a:Tensor):
|
|
c = Tensor(Tensor.invalids(a.shape[0]//len(devs), a.shape[1], dtype=a.dtype, device=devs).uop.unshard(0), device=devs)
|
|
return Tensor.custom_kernel(c, a, fxn=double_kernel, grad_fxn=double_grad)[0]
|
|
|
|
np.testing.assert_allclose(f(a).numpy(), 14.0)
|
|
|
|
# g is f with empty output instead of invalids
|
|
@function(precompile=True, allow_implicit=True)
|
|
def g(a:Tensor):
|
|
c = Tensor(Tensor.empty(a.shape[0]//len(devs), a.shape[1], dtype=a.dtype, device=devs).uop.unshard(0), device=devs)
|
|
return Tensor.custom_kernel(c, a, fxn=double_kernel, grad_fxn=double_grad)[0]
|
|
|
|
np.testing.assert_allclose(g(a).numpy(), 14.0)
|
|
|
|
def test_custom_kernel_inplace_output_is_implicit(self):
|
|
# a custom_kernel output the kernel also READS (in-place add) is not write-only, so it must be captured as an input
|
|
def inplace_add(C:UOp, A:UOp) -> UOp:
|
|
i = UOp.range(A.shape[0], 0)
|
|
return C[i].store(C[i].load() + A[i]).end(i).sink(arg=KernelInfo(name="inplace_add"))
|
|
@function(precompile=True, allow_implicit=False)
|
|
def f(a:Tensor): return Tensor.custom_kernel(Tensor.empty(*a.shape, dtype=a.dtype, device=a.device), a, fxn=inplace_add)[0]
|
|
with self.assertRaisesRegex(RuntimeError, "implicit buffer"): f(Tensor([1., 2., 3., 4.]).contiguous().realize())
|
|
|
|
def test_custom_kernel_write_only_persistent_output_is_implicit(self):
|
|
# a write-only custom_kernel output that is a realized buffer must be captured
|
|
def write(C:UOp, A:UOp) -> UOp:
|
|
i = UOp.range(A.shape[0], 0)
|
|
return C[i].store(A[i] * 2.0).end(i).sink(arg=KernelInfo(name="write"))
|
|
state = Tensor([100., 200., 300., 400.], device="CPU").contiguous().realize()
|
|
@function(precompile=True, allow_implicit=True)
|
|
def f(a:Tensor): return Tensor.custom_kernel(state, a, fxn=write)[0]
|
|
f(Tensor([1., 2., 3., 4.], device="CPU").contiguous().realize()).realize()
|
|
np.testing.assert_allclose(state.numpy(), [2., 4., 6., 8.])
|
|
|
|
def test_custom_kernel_program_invalids_not_captured(self):
|
|
# llama FP8 kernels are PROGRAM with bare-buffer sinks (no analyzable stores), so the invalids scratch
|
|
# still must not be captured as an input -- else it is read before the kernel writes it
|
|
src = "void k(float* restrict data0, float* restrict data1) { for (int i=0;i<4;i++) data0[i]=data1[i]*2.0f; }"
|
|
lib = Device["CPU"].compiler.compile(src)
|
|
def prog(C:UOp, A:UOp) -> UOp:
|
|
sink = UOp.sink(C.base, A.base, arg=KernelInfo(name="k"))
|
|
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
|
|
UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)),
|
|
arg=ProgramInfo(global_size=(1, 1, 1), local_size=(1, 1, 1), globals=(0, 1)))
|
|
|
|
@function(precompile=True)
|
|
def f(a:Tensor):
|
|
c = Tensor.invalids(*a.shape, dtype=a.dtype, device=a.device)
|
|
return Tensor.custom_kernel(c, a, fxn=prog)[0]
|
|
|
|
a = Tensor([1., 2., 3., 4.], device="CPU").contiguous().realize()
|
|
np.testing.assert_allclose(f(a).numpy(), [2., 4., 6., 8.])
|
|
|
|
def test_invalid_store_into_realized_buffer_is_captured(self):
|
|
# only fresh invalids() scratch is skipped; a realized buffer is a real input even if an Invalid store
|
|
# writes into part of it (its other elements must be preserved), so it is still captured
|
|
state = Tensor([10., 20., 30., 40.], device="CPU").contiguous().realize()
|
|
@function(precompile=True, allow_implicit=True)
|
|
def f(a:Tensor):
|
|
after = state.uop.after(state.uop.shrink(((0, 2),)).store(Invalid))
|
|
return Tensor(after).contiguous() + a
|
|
out = f(Tensor([1., 1., 1., 1.], device="CPU").contiguous().realize())
|
|
np.testing.assert_allclose(out.numpy(), [11., 21., 31., 41.])
|
|
|
|
def test_custom_kernel_precompile_further_compute(self, multi=False, kernel_count:int=2):
|
|
devs = ("CPU:0", "CPU:1")
|
|
def my_kernel(C:UOp, A:UOp) -> UOp:
|
|
C, A = C.flatten(), A.flatten()
|
|
i = UOp.range(A.numel(), 0)
|
|
return C[i].store(A[i] * 2.0).end(i).sink(arg=KernelInfo(name="my_kernel"))
|
|
|
|
@function(precompile=True)
|
|
def f(a:Tensor):
|
|
c = Tensor.invalids(*a.uop.shard_shape, dtype=a.dtype, device=a.device)
|
|
if multi: c = Tensor(c.uop.unshard(a.uop.axis), device=a.device)
|
|
c = Tensor.custom_kernel(c, a, fxn=my_kernel)[0]
|
|
return c + 1
|
|
|
|
a = Tensor([1., 2., 3., 4.]).contiguous()
|
|
if multi: a = a.shard(devs, axis=0)
|
|
a.realize()
|
|
out = f(a)
|
|
GlobalCounters.reset()
|
|
out.realize()
|
|
assert_kernel_count(kernel_count)
|
|
np.testing.assert_allclose(out.numpy(), [3., 5., 7., 9.])
|
|
|
|
def test_custom_kernel_precompile_further_compute_multi(self): self.test_custom_kernel_precompile_further_compute(multi=True, kernel_count=4)
|
|
|
|
class TestFunctionGrad(unittest.TestCase):
|
|
def test_function_grad_ops(self, precompile=False, precompile_backward=False):
|
|
N = 64
|
|
x = Tensor.ones(N,N).contiguous()
|
|
w1 = Tensor.ones(N,N).contiguous()
|
|
w2 = Tensor.ones(N,N).contiguous()
|
|
w3 = Tensor.ones(N,N).contiguous()
|
|
ref = Tensor.ones(N,N).contiguous()
|
|
Tensor.realize(x, w1, w2, w3, ref)
|
|
@function(precompile=precompile, precompile_backward=precompile_backward)
|
|
def f(x, w1, w2, w3) -> tuple[Tensor, ...]:
|
|
p1 = x@w1
|
|
p2 = p1@w2
|
|
p3 = p2@w3
|
|
return p1, p2, p3, p3.contiguous()
|
|
ret = f(x, w1, w2, w3)[-1]
|
|
loss = (ret-ref).square().mean().backward()
|
|
print("RESET")
|
|
GlobalCounters.reset()
|
|
loss.realize(w1.grad, w2.grad, w3.grad)
|
|
print(GlobalCounters.global_ops, GlobalCounters.global_mem)
|
|
self.assertLessEqual(GlobalCounters.global_ops, 5000000)
|
|
def test_function_grad_ops_precompile(self): self.test_function_grad_ops(precompile=True)
|
|
def test_function_grad_ops_precompile_backward(self):
|
|
self.test_function_grad_ops(precompile=True, precompile_backward=True)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|