Files
tinygrad/test/unit/test_cpu.py
nimlgen d3378010ee schedule() -> schedule_linear() in tests (batch 1) (#15915)
* schedule_with_vars -> linear_with_vars in tests

* tests batch 1

* batch 2

* estimate_uop

* simpler

* rm
2026-04-24 23:40:53 +03:00

25 lines
1.0 KiB
Python

import unittest, io
from contextlib import redirect_stdout
from tinygrad import Tensor, Device
from tinygrad.helpers import Target
from tinygrad.renderer.nir import LVPRenderer
from tinygrad.engine.realize import get_program
@unittest.skipIf(Device.DEFAULT != "CPU", "only run on CPU")
class TestCPU(unittest.TestCase):
def test_arch_feats(self):
ast = (Tensor.empty(16) + Tensor.empty(16)).schedule_linear().src[-1].src[0]
for ren in Device[Device.DEFAULT].renderers:
for arch, expect_vmov in [("x86_64,x86-64,avx", True), ("x86_64,x86-64,-avx", False)]:
with self.subTest(arch=arch):
if ren is LVPRenderer: continue # LVP does not play nice with cross compilation
r = ren(Target(device="CPU", arch=arch))
p = get_program(ast, r)
lib = r.compiler.compile(p.src)
out = io.StringIO()
with redirect_stdout(out): r.compiler.disassemble(lib)
self.assertEqual("vmov" in out.getvalue(), expect_vmov, out.getvalue())
if __name__ == '__main__':
unittest.main()