mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-04-29 03:00:14 -04:00
107 lines
4.5 KiB
Python
107 lines
4.5 KiB
Python
import os, time, struct, functools, unittest
|
|
from dataclasses import replace
|
|
from typing import Any, Callable
|
|
import numpy as np
|
|
from tinygrad import Tensor, dtypes, Device
|
|
from tinygrad.uop.ops import UOp, Ops, KernelInfo
|
|
from tinygrad.tensor import _to_np_dtype
|
|
from tinygrad.engine.realize import get_program
|
|
from tinygrad.dtype import DType
|
|
from tinygrad.nn.state import get_parameters
|
|
from tinygrad.helpers import T, CI, Target
|
|
from tinygrad.renderer import Renderer
|
|
from tinygrad.codegen import full_rewrite_to_sink, line_rewrite, pm_linearize_cleanups
|
|
from tinygrad.codegen.late.linearizer import linearize
|
|
|
|
# decorator to skip slow tests by default, run with RUN_SLOW=1 to include them
|
|
slow = unittest.skipUnless(os.getenv("RUN_SLOW"), "slow test, set RUN_SLOW=1 to run")
|
|
from tinygrad.runtime.ops_python import PythonProgram, PythonRenderer, PythonCompiler
|
|
|
|
def get_uops(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
|
|
"""Extract linearized UOps from a sink. Test helper that only does linearization (no render)."""
|
|
if ren is None: ren = Renderer(Target())
|
|
if sink.arg is None: sink = sink.replace(arg=KernelInfo())
|
|
full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
|
|
return line_rewrite(linearize(full_sink), pm_linearize_cleanups)
|
|
|
|
def replace_opts(ast:UOp, opts:list) -> UOp: return ast.replace(arg=replace(ast.arg, opts_to_apply=tuple(opts)))
|
|
|
|
def derandomize_model(model):
|
|
for p in get_parameters(model):
|
|
p.replace(Tensor.empty(p.shape, device=p.device, dtype=p.dtype))
|
|
p.realize()
|
|
|
|
def call_is_graph(call:UOp) -> bool:
|
|
ast = call.src[0]
|
|
return ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph"
|
|
|
|
def jit_cache_count(linear:UOp) -> int:
|
|
n = 0
|
|
for call in linear.src:
|
|
ast = call.src[0]
|
|
if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": n += jit_cache_count(ast.src[0])
|
|
else: n += 1
|
|
return n
|
|
|
|
def assert_jit_cache_len(fxn, expected_len):
|
|
linear = fxn.captured.linear if fxn.captured is not None else None
|
|
if linear is None or not linear.src:
|
|
assert expected_len == 0, expected_len
|
|
return
|
|
if call_is_graph(linear.src[0]):
|
|
assert len(linear.src) == 1, len(linear.src)
|
|
inner = linear.src[0].src[0].src[0] # LINEAR UOp inside CUSTOM_FUNCTION
|
|
assert len(inner.src) == expected_len, f"expected {expected_len}, got {len(inner.src)}"
|
|
else:
|
|
assert len(linear.src) == expected_len, f"expected {expected_len}, got {len(linear.src)}"
|
|
|
|
def rand_for_dtype(dt:DType, size:int, allow_subnormal=True):
|
|
if dtypes.is_unsigned(dt):
|
|
return np.random.randint(0, 100, size=size, dtype=_to_np_dtype(dt))
|
|
elif dtypes.is_int(dt):
|
|
return np.random.randint(-100, 100, size=size, dtype=_to_np_dtype(dt))
|
|
elif dt == dtypes.bool:
|
|
return np.random.choice([True, False], size=size)
|
|
ret = np.random.uniform(-10, 10, size=size).astype(_to_np_dtype(dt))
|
|
if not allow_subnormal:
|
|
min_normal = 2.0 ** (2 - (1 << (dtypes.finfo(dt)[0] - 1)))
|
|
ret = np.where(np.abs(ret) < min_normal, 0, ret)
|
|
return ret
|
|
|
|
def timeit(fxn:Callable[..., T], *args, **kwargs) -> tuple[T, float]:
|
|
st = time.perf_counter_ns()
|
|
ret = fxn(*args, **kwargs)
|
|
return ret, (time.perf_counter_ns()-st)*1e-6
|
|
|
|
def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None):
|
|
allocator = Device['PYTHON'].allocator
|
|
bufs = []
|
|
for buf_dt, data in inputs or []:
|
|
bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
|
|
allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
|
|
g = UOp(Ops.PARAM, uop.dtype.ptr(), arg=0, src=())
|
|
prg = get_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
|
|
prog = PythonProgram("run", PythonCompiler().compile(prg.src))
|
|
prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs)
|
|
return out_buf.cast(uop.dtype.fmt or "").tolist()[0]
|
|
|
|
def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:
|
|
sink = UOp.group(*u)
|
|
for r in sink.ranges: sink = sink.end(r)
|
|
ret = get_uops(sink.sink(arg=KernelInfo(opts_to_apply=())), ren)
|
|
assert ret[-1].op is Ops.SINK
|
|
return ret
|
|
|
|
def not_support_multi_device():
|
|
# CL and CUDA don't support multi device if in CI
|
|
return CI and Device.DEFAULT in ("CL", "CUDA")
|
|
|
|
def needs_second_gpu(fn):
|
|
@functools.wraps(fn)
|
|
def wrapper(self, *args, **kwargs):
|
|
# check if there's a second GPU, if not, skip multi tests
|
|
try: Tensor.zeros(10, device=f"{Device.DEFAULT}:1").contiguous().realize()
|
|
except Exception as e: self.skipTest(f"second device not available: {e}")
|
|
return fn(self, *args, **kwargs)
|
|
return wrapper
|