mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 05:23:41 +08:00
Daniel Koepping
23 lines
933 B
Python
23 lines
933 B
Python
import unittest
|
|
from tinygrad import Tensor, dtypes, Device
|
|
from tinygrad.helpers import Context
|
|
from tinygrad.engine.realize import estimate_uop, compile_linear
|
|
from tinygrad.renderer.ptx import PTXRenderer
|
|
|
|
class TestArange(unittest.TestCase):
|
|
def test_cat_complexity(self):
|
|
x = Tensor.arange(2**10) + Tensor.empty((), dtype=dtypes.uint32)
|
|
out = x.cat(x).cat(Tensor.empty(1, dtype=dtypes.uint32))
|
|
linear = compile_linear(out.schedule_linear())
|
|
self.assertLessEqual(estimate_uop(linear.src[-1]).ops, out.numel()*20)
|
|
|
|
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX indexing is weird")
|
|
def test_tri_complexity(self):
|
|
with Context(NOOPT=1):
|
|
t = Tensor.ones(256, 256).contiguous().realize()
|
|
linear = compile_linear(t.triu().schedule_linear())
|
|
self.assertLessEqual(estimate_uop(linear.src[-1]).ops, 4 * 256 * 256)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|