mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 05:23:41 +08:00
Daniel Koepping
32 lines
1.5 KiB
Python
32 lines
1.5 KiB
Python
import unittest
|
|
from tinygrad import Tensor, UOp, Context, GlobalCounters
|
|
from tinygrad.uop.ops import Ops, KernelInfo
|
|
from tinygrad.renderer.cstyle import ClangRenderer
|
|
from test.helpers import assert_kernel_count
|
|
|
|
class TestCustomKernel(unittest.TestCase):
|
|
@Context(DEV="CPU")
|
|
def test_simple_from_source(self):
|
|
a = Tensor.arange(4).clone().realize()
|
|
src = ClangRenderer.kernel_typedef + " test_src(int* restrict a) { a[0] = 1; }"
|
|
def custom_src_kernel(A:UOp, B:UOp) -> UOp:
|
|
sink = UOp.sink(A, arg=KernelInfo(name="test_src"))
|
|
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())), UOp(Ops.SOURCE, arg=src),))
|
|
a = Tensor.custom_kernel(a.reshape(2, 2).clone(), a.reshape(2, 2).T, fxn=custom_src_kernel)[0]
|
|
self.assertEqual(a.tolist(), [[1, 1], [2, 3]])
|
|
|
|
@Context(DEV="CPU")
|
|
def test_simple_from_source_alt(self):
|
|
a = Tensor.arange(4).clone().realize()
|
|
src = ClangRenderer.kernel_typedef + " copy(int* restrict out, int* restrict in) { for (int i = 0; i < 4; i++) out[i] = in[i]; }"
|
|
def custom_src_kernel(out:UOp, inp:UOp) -> UOp:
|
|
sink = UOp.sink(out, inp, arg=KernelInfo(name="copy"))
|
|
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())), UOp(Ops.SOURCE, arg=src),))
|
|
out = Tensor.custom_kernel(Tensor.empty_like(a), a+1, fxn=custom_src_kernel)[0]
|
|
GlobalCounters.reset()
|
|
out.realize()
|
|
assert_kernel_count(2)
|
|
self.assertEqual(out.tolist(), [1, 2, 3, 4])
|
|
|
|
if __name__ == "__main__": unittest.main()
|