import unittest from tinygrad import Tensor, Context, Device from tinygrad.engine.realize import get_program from tinygrad.renderer import Opt, OptOps from tinygrad.uop.ops import KernelInfo class TestLinearizerRewrite(unittest.TestCase): def test_reduction(self): t = Tensor.ones((64,64), device="NULL").contiguous().realize() out = (t*2).sum(axis=1) with Context(SPLIT_REDUCEOP=0, DEVECTORIZE=0): si = out.schedule()[-1] opts_to_apply = [] opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4)) opts_to_apply.append(Opt(OptOps.UNROLL, 0, 4)) ast = si.ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply))) prg = get_program(ast, Device["CPU"].renderer) print(prg.src) def test_arange(self): out = Tensor.arange(32, device="NULL") with Context(SPLIT_REDUCEOP=0, DEVECTORIZE=0): si = out.schedule()[-1] opts_to_apply = [] opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4)) opts_to_apply.append(Opt(OptOps.UNROLL, 0, 4)) ast = si.ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply))) prg = get_program(ast, Device["CPU"].renderer) print(prg.src) def test_kernel_info(self): out = Tensor.arange(4, device="NULL") si = out.schedule()[-1] ast = si.ast.replace(arg=KernelInfo(opts_to_apply=())) prg = get_program(ast, Device["CPU"].renderer) assert prg.applied_opts == (), f"expected no opts, got {prg}" prg = get_program(ast.replace(arg=None), Device["CPU"].renderer) assert prg.applied_opts != (), f"expected opts to apply, got {prg.applied_opts}" prg = get_program(ast.replace(arg=KernelInfo(name="custom")), Device["CPU"].renderer) self.assertEqual(prg.name, "custom") if __name__ == '__main__': unittest.main()