mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-10-28 00:53:45 +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.
46 lines
1.7 KiB
Python
46 lines
1.7 KiB
Python
import unittest
|
|
from tinygrad import Tensor, Device, Context
|
|
from tinygrad.codegen import do_to_program
|
|
from tinygrad.codegen.opt import Opt, OptOps
|
|
from tinygrad.uop.ops import AxisType
|
|
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.SPLIT, 0, (4, AxisType.UPCAST))]
|
|
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)
|