45 lines
1.2 KiB
Python
45 lines
1.2 KiB
Python
#!/usr/bin/env python
|
|
import unittest
|
|
from tinygrad.tensor import Tensor
|
|
from tinygrad.codegen.linearizer import Linearizer
|
|
from tinygrad.renderer.opencl import OpenCLRenderer
|
|
from tinygrad.graph import graph_uops
|
|
from tinygrad.nn import Conv2d
|
|
|
|
class TestUopsGraph(unittest.TestCase):
|
|
def test_matmul(self):
|
|
N = 1024
|
|
a = Tensor.rand(N,N)
|
|
b = Tensor.rand(N,N)
|
|
si = (a@b).lazydata.schedule()[-1]
|
|
lin = Linearizer(si.ast)
|
|
lin.hand_coded_optimizations()
|
|
print(lin.colored_shape())
|
|
uops = lin.linearize().uops
|
|
graph_uops(uops)
|
|
for u in uops: print(u)
|
|
print(OpenCLRenderer("matmul", uops)[0])
|
|
|
|
def test_reduce(self):
|
|
a = Tensor.rand(1024*1024)
|
|
si = a.sum().lazydata.schedule()[-1]
|
|
lin = Linearizer(si.ast)
|
|
lin.hand_coded_optimizations()
|
|
uops = lin.linearize().uops
|
|
graph_uops(uops)
|
|
#print(OpenCLRenderer("reduce", uops)[0])
|
|
|
|
def test_conv(self):
|
|
x = Tensor.rand(1,3,16,16)
|
|
c = Conv2d(3, 16, (3,3))
|
|
si = c(x).elu().lazydata.schedule()[-1]
|
|
lin = Linearizer(si.ast)
|
|
lin.hand_coded_optimizations()
|
|
uops = lin.linearize().uops
|
|
graph_uops(uops)
|
|
print(lin.colored_shape())
|
|
print(OpenCLRenderer("conv", uops)[0])
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|