superrob/IQ.Pilot
1
0
Fork 0
forked from IQ.Lvbs/IQ.Pilot
Code Pull Requests Activity

IQ.Pilot Release Commit @ f2a861c

Browse Source
This commit is contained in:
IQ.Lvbs CI [bot]
2026-09-02 15:07:09 -05:00
parent b42569dbca
commit e8748fd704
5497 changed files with 316070 additions and 179848 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
artifacts/package_sources/tinygrad/test/external/external_benchmark_am.py vendored Normal file
View File

@@ -0,0 +1,9 @@
from tinygrad.helpers import Profiling
from tinygrad import Device
if __name__ == "__main__":
am = Device["AMD"]
# kfd is 0.55ms!
with Profiling("allocation 127.7mb"):
am.allocator.alloc(int(127.7*1024*1024))

19
artifacts/package_sources/tinygrad/test/external/external_benchmark_bert_matmuls.py vendored Normal file
View File

@@ -0,0 +1,19 @@
from tinygrad import Tensor, dtypes
from tinygrad.dtype import to_dtype
from tinygrad.helpers import getenv, Context
Context(DEFAULT_FLOAT=dtypes.float16).__enter__()
if __name__ == "__main__":
# matmuls in bert layers
BS = getenv("BS", 96//6)
acc_dtype = to_dtype(getenv("ACC_DTYPE", "half"))
tensors = [
(Tensor.empty(BS, 512, 1024), Tensor.empty(1024, 1024).T), # linear to get qkv
(Tensor.empty(BS, 512, 16, 64).permute(0,2,1,3), Tensor.empty(BS, 512, 16, 64).permute(0,2,3,1)), # q@k
(Tensor.empty(BS, 16, 512, 512), Tensor.empty(BS, 512, 16, 64).permute(0,2,1,3)), # qk@v
]
for t0, t1 in tensors:
print(f"{t0.shape=}, {t0.uop.st.is_expanded()=}, {t1.shape=}, {t1.uop.st.is_expanded()=}")
for _ in range(5):
t0.dot(t1, dtype=acc_dtype).realize()

18
artifacts/package_sources/tinygrad/test/external/external_benchmark_bert_softmax.py vendored Normal file
View File

@@ -0,0 +1,18 @@
from tinygrad import Tensor, dtypes, GlobalCounters
from tinygrad.dtype import to_dtype
from tinygrad.helpers import getenv, Context
from test.backend.test_softmax_fusion import single_kernel_softmax
Context(DEFAULT_FLOAT=dtypes.float16).__enter__()
if __name__ == "__main__":
# softmax in bert layers
BS = getenv("BS", 96//6)
acc_dtype = to_dtype(getenv("ACC_DTYPE", "half"))
t = Tensor.empty(BS, 16, 512, 512)
t.softmax(-1, dtype="half").realize()
# test single kernel softmax
GlobalCounters.reset()
single_kernel_softmax(t, -1, acc_dtype).realize()

8
artifacts/package_sources/tinygrad/test/external/external_benchmark_disk_raw.py vendored Normal file
View File

@@ -0,0 +1,8 @@
import pathlib
from tinygrad import Tensor, Device, Context
from tinygrad.helpers import getenv
if __name__ == "__main__":
with Context(DEBUG=2):
disk_llama = Tensor(pathlib.Path(getenv("TESTFILE", "/raid/weights/LLaMA-3/8B/consolidated.00.pth")))
device_llama = disk_llama.to(Device.DEFAULT).realize()

31
artifacts/package_sources/tinygrad/test/external/external_benchmark_hip_compile.py vendored Normal file
View File

@@ -0,0 +1,31 @@
import random, os
from tinygrad.helpers import Timing
from tinygrad.runtime.ops_hip import compile_hip, HIPDevice
from tinygrad.runtime.ops_cl import compile_cl, CLDevice
# OMP_NUM_THREADS=1 strace -tt -f -e trace=file python3 test/external/external_benchmark_hip_compile.py
# AMD_COMGR_REDIRECT_LOGS=stdout AMD_COMGR_EMIT_VERBOSE_LOGS=1 python3 test/external/external_benchmark_hip_compile.py
# issue is in https://github.com/ROCm-Developer-Tools/clr/
if __name__ == "__main__":
HIPDevice()
CLDevice()
# warmup
name = "none"+str(random.randint(0, 1000000))
compile_cl.__wrapped__(f"void {name}() {{}}")
print("compile cl warmed up")
compile_hip.__wrapped__(f"void {name}() {{}}")
print("compile hip warmed up")
print("**** benchmark ****")
name = "none"+str(random.randint(0, 1000000))
# this uses AMD_COMGR_ACTION_COMPILE_SOURCE_TO_BC, then it links the lib on the next step
with Timing("compile cl: "): compile_cl.__wrapped__(f"void {name}() {{}}")
# this uses AMD_COMGR_ACTION_COMPILE_SOURCE_WITH_DEVICE_LIBS_TO_BC, much slower
with Timing("compile hip: "): compile_hip.__wrapped__(f"void {name}() {{}}")
os._exit(0)

20
artifacts/package_sources/tinygrad/test/external/external_benchmark_keccak.py vendored Normal file
View File

@@ -0,0 +1,20 @@
from tinygrad import Tensor, dtypes
from tinygrad.engine.jit import TinyJit
from tinygrad.helpers import Timing, getenv
if __name__ == "__main__":
BS = getenv("BS", 2**14)
BLOCKSIZE = getenv("BLOCKSIZE", 4096)
HASHFN = getenv("HASHFN", "shake_128")
NRUNS = getenv("NRUNS", 5)
@TinyJit
def hasher(data: Tensor): return data.keccak(HASHFN)
t = Tensor.randn(BS, BLOCKSIZE, dtype=dtypes.uint8).realize()
ds_mib = t.nbytes() / 1024**2
print(f"--- benchmarking (hash: {HASHFN}, data size: {ds_mib} MiB, block size: {BLOCKSIZE} B, batch size: {BS})")
for i in range(NRUNS):
with Timing(f"run: {i+1}, elapsed time: ", (lambda et: f", throughput: {ds_mib / (et*1e-9):.2f} MiB/s")):
hasher(t).realize()

38
artifacts/package_sources/tinygrad/test/external/external_benchmark_kernel_launch.py vendored Normal file
View File

@@ -0,0 +1,38 @@
import time
from tinygrad import Tensor, TinyJit, Device, Context
from tinygrad.helpers import Profiling, Timing, GlobalCounters
# python3 test/speed/external_test_speed_v_torch.py TestSpeed.test_add_a
@TinyJit
def plus(a:Tensor, b:Tensor): return a+b
if __name__ == "__main__":
a = Tensor([1]).realize()
b = Tensor([1]).realize()
for i in range(5):
with Timing(prefix=f"{i}:"):
c = plus(a,b)
Device[c.device].synchronize()
assert c.item() == 2
for i in range(5):
st = time.perf_counter()
c = plus(a,b)
et = time.perf_counter() - st
Device[c.device].synchronize()
print(f"nosync {i}: {et*1e6:.2f} us")
for i in range(5):
st = time.perf_counter()
c = plus(a,b)
Device[c.device].synchronize()
et = time.perf_counter() - st
print(f"precise {i}: {et*1e6:.2f} us")
assert GlobalCounters.time_sum_s == 0
with Context(DEBUG=2):
st = time.perf_counter()
c = plus(a,b)
Device[c.device].synchronize()
et = time.perf_counter() - st
print(f"kernel {GlobalCounters.time_sum_s*1e3:.2f} ms / full {et*1e3:.2f} ms -- {et/(GlobalCounters.time_sum_s+1e-12):.2f} x")
with Profiling():
c = plus(a,b)

68
artifacts/package_sources/tinygrad/test/external/external_benchmark_llama_offload.py vendored Normal file
View File

@@ -0,0 +1,68 @@
from tinygrad import Tensor, dtypes, Device, TinyJit
from tinygrad.helpers import getenv, Timing, round_up
# llama 70b shapes
DIM, HIDDEN, N_HEADS, N_KV_HEADS, HEAD_DIM = 8192, 28672, 64, 8, 128
QKV_OUT = N_HEADS*HEAD_DIM + 2*N_KV_HEADS*HEAD_DIM
LAYER_SHAPES = {
"wqkv": ((QKV_OUT, DIM), 0),
"wo": ((DIM, N_HEADS*HEAD_DIM), 1),
"w1": ((HIDDEN, DIM), 0),
"w3": ((HIDDEN, DIM), 0),
"w2": ((DIM, HIDDEN), 1),
}
def build_params(layers:int, devs:tuple[str,...], dtype):
ts, nbytes = [], 0
ndev = len(devs)
for _ in range(layers):
for (a,b), ax in LAYER_SHAPES.values():
shape = (round_up(a, ndev) if ax==0 else a, round_up(b, ndev) if ax==1 else b)
t = Tensor.empty(*shape, dtype=dtype).shard(devs, axis=ax).realize()
ts.append(t)
nbytes += t.numel()*dtype.itemsize
return ts, nbytes
def sync():
Device["CPU"].synchronize()
Device[Device.DEFAULT].synchronize()
if __name__ == "__main__":
ngpus = getenv("GPUS", 8)
layers = getenv("LAYERS", 8)
devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(ngpus))
# to cpu
grads, nbytes = build_params(layers, devs, dtypes.bfloat16)
cpu_dst = [Tensor.rand(*g.shape, dtype=dtypes.bfloat16, device="CPU").contiguous().realize() for g in grads]
sync()
def to_cpu():
for d,g in zip(cpu_dst, grads):
d.assign(g.to("CPU"))
return Tensor.realize(*cpu_dst)
to_cpu_jit = TinyJit(to_cpu)
for _ in range(3):
to_cpu_jit()
sync()
for _ in range(5):
with Timing("to cpu ", on_exit=lambda ns: f" @ {nbytes/ns:.2f} GB/s"):
to_cpu_jit()
sync()
# to gpu
from extra.gemm.cdna_asm_gemm import FP8_DTYPE
w_gpu, nbytes = build_params(layers, devs, FP8_DTYPE)
cpu_src = [Tensor.rand(*w.shape, dtype=FP8_DTYPE, device="CPU").contiguous().realize() for w in w_gpu]
sync()
def to_gpu():
for w,c in zip(w_gpu, cpu_src):
w.assign(c.shard_like(w))
return Tensor.realize(*w_gpu)
to_gpu_jit = TinyJit(to_gpu)
for _ in range(3):
to_gpu_jit()
sync()
for _ in range(5):
with Timing("to gpu ", on_exit=lambda ns: f" @ {nbytes/ns:.2f} GB/s"):
to_gpu_jit()
sync()

34
artifacts/package_sources/tinygrad/test/external/external_benchmark_llama_schedule.py vendored Normal file
View File

@@ -0,0 +1,34 @@
from tinygrad import nn, Tensor, dtypes
from tinygrad.helpers import DEV, Timing, Context
from extra.models.llama import Transformer
from examples.llama3 import MODEL_PARAMS
if __name__ == "__main__":
DEV.value = "NULL"
#model_size = "8B"
model_size = "405B"
with Context(TRAINING=1), Timing("total "):
with Timing("***** create model in "):
model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=nn.Linear, embedding=nn.Embedding,
max_context=1024, jit=True, disable_kv_cache=True)
with Timing("***** fake state in "):
Tensor.realize(*[p.assign(Tensor.empty(*p.shape, device=p.device, dtype=p.dtype)) for p in nn.state.get_parameters(model)])
with Timing("***** create optim in "):
opt = nn.optim.AdamW(nn.state.get_parameters(model))
with Timing("***** run model in "):
toks = Tensor.empty(1, 1024, dtype=dtypes.int)
out = model(toks, 0, temperature=float('nan'))
with Timing("***** backward in "):
out.mean().backward()
with Timing("***** realize in "):
out.realize()
with Timing("***** step in "):
opt.step()

49
artifacts/package_sources/tinygrad/test/external/external_benchmark_multitensor_allreduce.py vendored Normal file
View File

@@ -0,0 +1,49 @@
from tinygrad import Tensor, Device, GlobalCounters, TinyJit, dtypes
from tinygrad.helpers import getenv, Context, DEBUG
def test(devs: list[str], N: int, iters:int = 10, name:str = "allreduce"):
@TinyJit
def f(t: Tensor) -> Tensor: t.sum(0).realize()
secs, gflops, gbs = 0, 0, 0
for i in range(-3, iters):
t = Tensor.empty((len(devs), N))
t = t.shard(devs, 0).realize()
GlobalCounters.reset()
f(t)
for d in devs: Device[d].synchronize()
if i < 0: continue # warm up jit
i_secs = GlobalCounters.time_sum_s
i_gflops = GlobalCounters.global_ops/i_secs/10**9
i_gbs = (N*4)/i_secs/10**9
print(f"{name} iter {i+1}/{iters}: {i_secs:.6f} sec {i_gflops:.2f} GFLOP/s {i_gbs:.2f} GB/s")
secs += i_secs
gflops += i_gflops
gbs += i_gbs
return (gflops/iters, gbs/iters, secs/iters)
def run(sz, n_gpus=6, iters=10, ring=0, all2all=0):
devs = tuple([f"{Device.DEFAULT}:{x}" for x in range(n_gpus)])
N = sz // dtypes.float32.itemsize
name = "all2all" if all2all else ("ring" if ring else "naive")
with Context(RING=(2 if ring else 0), ALL2ALL=(2 if all2all else 0), JIT_BATCH_SIZE=0, DEBUG=max(DEBUG.value, 2)):
return test(devs, N, iters=iters, name=name)
def main():
n_gpus = getenv("GPUS", 6)
iters = getenv("ITERS", 10)
sz = getenv("SZ", 1000) * 10**6 # size of data on each gpu
print(f"Using {sz/10**9:.2f} GB of numbers on each of {n_gpus} GPUs, {n_gpus*sz/10**9:.2f} GB total.")
results = {}
for name, kwargs in [("naive", {}), ("ring", {"ring": 2}), ("all2all", {"all2all": 2})]:
results[name] = run(sz, n_gpus=n_gpus, iters=iters, **kwargs)
print("\n=== RESULTS ===")
for name, (gflops, gbs, secs) in results.items():
print(f"{name.upper()}:\n {secs:.6f} seconds/iter\n {gflops:.2f} GFLOP/s\n {gbs:.2f} GB/s")
if __name__ == "__main__":
main()

102
artifacts/package_sources/tinygrad/test/external/external_benchmark_op_conv.py vendored Normal file
View File

@@ -0,0 +1,102 @@
# ruff: noqa: E501 E712 F401
from dataclasses import replace
from tinygrad import dtypes, Device
from tinygrad.uop.ops import UOp, AxisType, Ops, KernelInfo
from tinygrad.codegen.opt import Opt, OptOps # pylint: disable=unused-import
from tinygrad.engine.realize import get_runtime
from tinygrad.codegen import to_program
from tinygrad.helpers import dedup, getenv
from tinygrad.device import Buffer
from tinygrad.dtype import Invalid
# PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
def vision_conv_143():
c0 = UOp.param(0, dtypes.half, shape=(16, 1024, 4))
c2 = UOp.range(32, 3)
c5 = UOp.range(128, 4)
c8 = UOp.range(16, 2)
c16 = UOp.range(7, 0, AxisType.REDUCE)
c17 = c8*2+c16
c24 = ((c17<3)!=True)&(c17<35)
c26 = UOp.range(7, 1, AxisType.REDUCE)
c27 = c2*2+c26
c32 = ((c27<3)!=True)&(c27<67)
c34 = UOp.param(1, dtypes.half, shape=(32, 1024, 4))
c38 = c5//2
c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(Invalid, dtypes.weakint))
c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
c49 = UOp.param(2, dtypes.half, shape=(64, 49, 4))
c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
c63 = UOp.param(3, dtypes.float, (128,))
c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
c67 = c0.index((c2*128+c5+c8*4096)).store(c65).end(c8, c2, c5)
opts = None
# JITBEAM=2
# (Opt(op=OptOps.UPCAST, axis=2, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None), Opt(op=OptOps.UPCAST, axis=2, arg=2), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.SWAP, axis=1, arg=2))
return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
def vision_conv_153():
c0 = UOp.param(0, dtypes.half, shape=(8, 1024, 4))
c2 = UOp.range(16, 3)
c5 = UOp.range(256, 4)
c8 = UOp.range(8, 2)
c16 = UOp.range(7, 0, AxisType.REDUCE)
c17 = c8*2+c16
c24 = ((c17<3)!=True)&(c17<19)
c26 = UOp.range(7, 1, AxisType.REDUCE)
c27 = c2*2+c26
c32 = ((c27<3)!=True)&(c27<35)
c34 = UOp.param(1, dtypes.half, shape=(16, 1024, 4))
c38 = c5//2
c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(Invalid, dtypes.weakint))
c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
c49 = UOp.param(2, dtypes.half, shape=(128, 49, 4))
c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
c63 = UOp.param(3, dtypes.float, (256,))
c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
c67 = c0.index((c2*256+c5+c8*4096)).store(c65).end(c8, c2, c5)
opts = None
# JITBEAM=2
# (Opt(op=OptOps.UPCAST, axis=2, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None), Opt(op=OptOps.UPCAST, axis=2, arg=2), Opt(op=OptOps.SWAP, axis=1, arg=2))
return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
def dm_conv_172():
c0 = UOp.param(0, dtypes.half, shape=(1, 240, 4))
c2 = UOp.range(960, 4)
c5 = UOp.param(1, dtypes.half, shape=(8, 384, 4))
c7 = UOp.range(32, 0, AxisType.REDUCE)
c10 = UOp.range(4, 1, AxisType.REDUCE)
c13 = UOp.range(12, 3, AxisType.REDUCE)
c18 = UOp.range(8, 2, AxisType.REDUCE)
c23 = UOp.param(2, dtypes.half, shape=(240, 128, 4))
c35 = c5.index((c7*4+c10+c13*128+c18*1536))*c23.index((c10*4+c2%4+c7*16+c2//4*512))
c37 = UOp.param(3, dtypes.float, (960,))
c39 = c35.reduce(c7, c10, arg=Ops.ADD)+c37.index(c2)
c50 = (1.0+((c39+0.044708251953125*(c39*(c39*c39)))*-2.3021129851685216).exp2()).reciprocal()*c39
c53 = c50.reduce(c18, c13, arg=Ops.ADD)*0.010416666666666666
c55 = c0.index(c2).store(c53).end(c2)
opts = None
# JITBEAM=2
# (Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.GROUPTOP, axis=1, arg=32), Opt(op=OptOps.UNROLL, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.GROUP, axis=1, arg=0))
return c55.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
ast = {143: vision_conv_143, 153: vision_conv_153, 172: dm_conv_172}[getenv("NUM", 143)]()
renderer = Device.default.renderer
allocator = Device.default.allocator
ps = to_program(ast, renderer)
rt = get_runtime(Device.DEFAULT, ps)
gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.PARAM]), key=lambda u: u.arg)
# print(len(gs))
# print([g.dtype for g in gs])
bufs = [Buffer(ps.arg.device, g.max_numel(), g.dtype).ensure_allocated() for g in gs]
gsize, lsize = ps.arg.launch_dims({})
t = rt(*[b._buf for b in bufs], global_size=gsize, local_size=lsize, vals=ps.arg.vals({}), wait=True)
print(f"{t*1e6:.2f} us")

34
artifacts/package_sources/tinygrad/test/external/external_benchmark_pyrender.py vendored Normal file
View File

@@ -0,0 +1,34 @@
# benchmark speed of pyrender for all created UOps saved with TRACK_MATCH_STATS=2
import functools, pickle
from tinygrad.uop.ops import UOp, Ops
from tinygrad.helpers import tqdm, temp, time_to_str, cpu_profile
BENCHMARK_OPS = {Ops.INDEX, Ops.STAGE}
@functools.cache
def create_uop(a:int) -> UOp:
op, dtype, src, arg, *rest = trace.uop_fields[a]
return UOp(op, dtype, tuple(create_uop(s) for s in src), arg, *rest)
if __name__ == "__main__":
# load rewrite trace
with open(temp("rewrites.pkl", append_user=True), "rb") as f:
trace = pickle.load(f)
# benchmark
result:list[tuple[str, int]] = []
try:
for steps in tqdm(trace.rewrites):
for r in steps:
for _,yn,_,__ in r.matches:
y = create_uop(yn)
if y.op in BENCHMARK_OPS:
with cpu_profile("pyrender") as e:
try: ren = y.render()
except Exception: ren = "PYRENDER_ERR"
result.append((ren, float(e.en-e.st)/1e6))
finally:
N = 10
print(f"Slowst {N} renders from {len(result)} samples:")
for ren,tm in sorted(result, key=lambda x:x[1], reverse=True)[:N]:
print(f"{time_to_str(tm).strip():<10s} {ren}")

127
artifacts/package_sources/tinygrad/test/external/external_benchmark_resnet.py vendored Normal file
View File

@@ -0,0 +1,127 @@
import functools
import time
import unittest
from tinygrad import Tensor, TinyJit, GlobalCounters, Device
from tinygrad.helpers import getenv, Context
from tinygrad.nn.optim import SGD
from tinygrad.nn.state import get_parameters
from tinygrad.engine.realize import run_linear
from extra.models import resnet
from examples.mlperf.initializers import Conv2dHeNormal, Linear
from examples.hlb_cifar10 import UnsyncedBatchNorm
# benchmark memory or kernel count: DEFAULT_FLOAT=HALF python test/external/external_benchmark_resnet.py
# benchmark speed: BEAM=2 JITCNT=10 DEFAULT_FLOAT=HALF python test/external/external_benchmark_resnet.py
# benchmark only one layer: BEAM=2 DEFAULT_FLOAT=HALF python test/external/external_benchmark_resnet.py BenchmarkResnetTrain.test_layer1_2
# inspect: DEBUG=2 BEAM=2 DEFAULT_FLOAT=HALF python test/external/external_benchmark_resnet.py
# inspect 1x1 convs: DEBUG=2 BEAM=2 CONV=2 DEFAULT_FLOAT=HALF python test/external/external_benchmark_resnet.py
# inspect 3x3 convs: DEBUG=2 BEAM=2 CONV=2 DEFAULT_FLOAT=HALF python test/external/external_benchmark_resnet.py
# inspect 3x3 convs with batchnorm: DEBUG=2 BEAM=2 CONV=2 BN=1 DEFAULT_FLOAT=HALF python test/external/external_benchmark_resnet.py
# etc
# use ASSIGN=0 to disable batchnorm/optimizer assigns
# memory will be slightly high with JITCNT > 1
bs = getenv("BS", 64)
class BenchmarkResnetTrain(unittest.TestCase):
def _get_layer(self, layer_i, slice_i):
# isolate to conv, with or without BN
conv = getenv("CONV", 0)
bn = getenv("BN", 0)
if not hasattr(self, 'model'):
resnet.Conv2d = Conv2dHeNormal
resnet.Linear = Linear
if not getenv("SYNCBN"): resnet.BatchNorm = functools.partial(UnsyncedBatchNorm, num_devices=1)
self.model = resnet.ResNet50()
self.layers = [self.model.layer1, self.model.layer2, self.model.layer3, self.model.layer4]
layer = self.layers[layer_i][slice_i]
xy = 112 >> layer_i
xy >>= (1 if slice_i > 0 or layer_i == 0 else 0) # layer 1 is preceded by maxpool2d
name = f"layer{layer_i+1} slice{slice_i+1}"
# get specific conv
if conv:
convs = [layer.conv1, layer.conv2, layer.conv3] + ([layer.downsample[0]] if layer.downsample else [])
bns = [layer.bn1, layer.bn2, layer.bn3] + ([layer.downsample[1]] if layer.downsample else [])
f = [convs[conv-1]]
if bn: f.append(bns[conv-1])
f.append(Tensor.relu)
cin = f[0].in_channels
if conv == 3: xy //= convs[1].stride
return f"{name} conv{conv} x{str((bs, cin, xy, xy)):20s} k{str(f[0].weight.shape):20s}" + (" bn" if bn else ""), f, cin, xy
cin = layer.conv1.in_channels
return f"{name} x{(bs, cin, xy, xy)}", [layer], cin, xy
def _test_layer(self, name, layer, cin, xy):
optim = SGD(get_parameters(layer), bs / 128 * 1.0) # need sgd for some params but not consequential for benchmarking
with Context(TRACK_MATCH_STATS=0): Tensor.realize(*[t.assign(t.detach().contiguous()) for t in get_parameters(optim)])
JITCNT = getenv("JITCNT", 1)
@TinyJit
@Context(TRAINING=1)
def step(x):
optim.zero_grad()
x.grad = None
y = x.sequential(layer).contiguous().contiguous_backward()
y.sum().backward()
if getenv("ASSIGN", 1): linear, var_vals = Tensor.linear_with_vars(y, x.grad, *optim.schedule_step())
else: linear, var_vals = Tensor.linear_with_vars(y, x.grad, *[t.grad for t in optim.params])
for _ in range(JITCNT):
run_linear(linear, var_vals)
CNT = getenv("CNT", 5)
best_tm = None
flops, mem_used, mem, kernels = None, None, None, None
for i in range(CNT):
with Context(TRACK_MATCH_STATS=0): x = Tensor.randn(bs, cin, xy, xy).realize()
GlobalCounters.reset()
st = time.perf_counter()
step(x)
Device[Device.DEFAULT].synchronize()
et = time.perf_counter()
flops = GlobalCounters.global_ops / JITCNT
mem_used = GlobalCounters.mem_used # a little high with JITCNT > 1 fsr
mem = GlobalCounters.global_mem / JITCNT
if kernels is None: kernels = GlobalCounters.kernel_count // JITCNT
tm = (et-st) / JITCNT
if best_tm is None or tm < best_tm: best_tm = tm
print(f"\r{name:38s}: {best_tm * 1000:>9.2f} ms, {flops / 10**12 / best_tm:>6.2f} tflops, {mem / 10**9 / best_tm:>5.0f} GB/s, "
f"{mem_used / 10**9: 6.2f} GB used, {kernels:>5d} kernels")
return best_tm, flops, mem, kernels
def test_layer1_1(self): self._est(*self._test_layer(*self._get_layer(0, 0)), 1)
def test_layer1_2(self): self._est(*self._test_layer(*self._get_layer(0, 1)), 2)
def test_layer2_1(self): self._est(*self._test_layer(*self._get_layer(1, 0)), 1)
def test_layer2_2(self): self._est(*self._test_layer(*self._get_layer(1, 1)), 3)
def test_layer3_1(self): self._est(*self._test_layer(*self._get_layer(2, 0)), 1)
def test_layer3_2(self): self._est(*self._test_layer(*self._get_layer(2, 1)), 5)
def test_layer4_1(self): self._est(*self._test_layer(*self._get_layer(3, 0)), 1)
def test_layer4_2(self): self._est(*self._test_layer(*self._get_layer(3, 1)), 2)
est_tm, est_flops, est_mem, est_kernels = 0, 0, 0, 0
@classmethod
def _est(cls, tm, flops, mem, kernels, mult):
cls.est_tm += tm * mult
cls.est_flops += flops * mult
cls.est_mem += mem * mult
cls.est_kernels += kernels * mult
@classmethod
def tearDownClass(cls):
print(f"\restimated step tm: {cls.est_tm * 1000.0:.2f} ms, {cls.est_flops / 10 ** 12 / cls.est_tm:.3f} tflops, "
f"{cls.est_mem / 10 ** 9 / cls.est_tm:.2f} GB/s, {cls.est_kernels} kernels")
if __name__ == '__main__':
unittest.main()

45
artifacts/package_sources/tinygrad/test/external/external_benchmark_schedule.py vendored Normal file
View File

@@ -0,0 +1,45 @@
from extra.models.resnet import ResNet50
from tinygrad import Tensor, nn, Device
from tinygrad.helpers import Profiling, Timing, getenv
from tinygrad.uop.ops import Ops
from tinygrad.codegen import full_rewrite_to_sink
from tinygrad.codegen.late.linearizer import linearize
from tinygrad.uop.spec import type_verify, spec_program
if __name__ == "__main__":
mdl = ResNet50()
for p in nn.state.get_parameters(mdl): p.replace(Tensor.empty(p.shape))
img = Tensor.empty(64, 3, 224, 224)
PROFILE = getenv("PYPROFILE", 0)
FORWARD_ONLY = getenv("FORWARD_ONLY", 0)
SCHEDULE_ONLY = getenv("SCHEDULE_ONLY", 0)
LINEARIZE = bool(getenv("LINEARIZE", 1))
with Timing("all "):
with Timing("***** model tensor in "):
out = mdl(img)
if not FORWARD_ONLY:
with Timing("***** model schedule in "):
with Profiling(PROFILE >= 3):
linear = out.schedule_linear()
if not SCHEDULE_ONLY:
asts = list({call.src[0].key:call.src[0] for call in linear.src if call.src[0].op is Ops.SINK}.values())
if (restrict_kernel := getenv("RESTRICT_KERNEL", -1)) != -1: asts = asts[restrict_kernel:restrict_kernel+1]
with Profiling(PROFILE, fn="/tmp/rewrite.prof"):
with Timing("***** model rewrite in "):
rewritten_uops = []
for u in asts:
rewritten_uops.append(full_rewrite_to_sink(u, ren=Device.default.renderer))
if LINEARIZE:
with Timing("***** model linearize in "):
uops_line = []
for u in rewritten_uops:
uops_line.append(linearize(u))
with Timing("***** model verify in "):
for u in uops_line: type_verify(u, spec_program)
print(sum(len(u) for u in uops_line))

13
artifacts/package_sources/tinygrad/test/external/external_cl_half_max.py vendored Normal file
View File

@@ -0,0 +1,13 @@
from tinygrad.runtime.ops_cl import CLDevice, CLProgram, compile_cl
if __name__ == "__main__":
dev = CLDevice()
lib = compile_cl("""
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
__kernel void test(__global half *out, __global half *a, __global half *b) {
int gid = get_global_id(0);
out[gid] = max(a[gid], b[gid]);
}
""")
prg = CLProgram(dev, "test", lib)

109
artifacts/package_sources/tinygrad/test/external/external_fuzz_ampt.py vendored Normal file
View File

@@ -0,0 +1,109 @@
import random
from typing import Optional
from tinygrad.helpers import round_up
from tinygrad.runtime.support.am.amdev import AMPageTableTraverseContext
from test.external.external_test_am import helper_read_entry_components, FakeAM
class AMPTFuzzer:
def __init__(self, total_size):
self.total_size = total_size
self.alloc_payload = 0
self.d = FakeAM()
self.allocations: dict[int, tuple[int, int]] = {} # ptr -> (size, pattern)
self.min_alloc_size = 0x1000
self.max_alloc_size = int(total_size * 0.1)
self.alloc_probability = 0.7
def generate_pattern(self, ptr: int, size: int) -> int: return random.randint(0, 0xff)
def fill_memory(self, va, size: int, pattern: int):
ctx = AMPageTableTraverseContext(self.d, self.d.mm.root_page_table, va.va_addr)
pages = list(ctx.next(size))
for _offset, _pt, _pte_idx, _n_ptes, _pte_covers in pages:
_vaddr = va.va_addr + _offset
for i in range(_n_ptes):
pte = helper_read_entry_components(_pt.entries[_pte_idx + i])
self.d.vram[pte['paddr']] = pattern # Mark this page
assert pte['valid'] == 1
# If page has contiguous fragment, all range should be this valid memory
frags_cnt = pte['fragment']
contig_range = (1 << (frags_cnt + 12))
start_vaddr = _vaddr & ~(contig_range - 1)
start_paddr = pte['paddr'] - (_vaddr - start_vaddr)
contig_ptes = contig_range // _pte_covers
assert contig_ptes > 0
ctx = AMPageTableTraverseContext(self.d, self.d.mm.root_page_table, start_vaddr)
frags_l = list(ctx.next(contig_range))
for f_offset, f_pt, f_pte_idx, f_n_ptes, f_pte_covers in frags_l:
for j in range(f_n_ptes):
f_pte = helper_read_entry_components(f_pt.entries[f_pte_idx + j])
assert f_pte['valid'] == 1
assert f_pte['paddr'] == start_paddr+f_offset+j*f_pte_covers, f"paddr {f_pte['paddr']:#x} not {start_paddr+f_offset+j*f_pte_covers:#x}"
_vaddr += _pte_covers
_offset += _pte_covers
return pages
def verify_memory(self, pages, pattern: int) -> bool:
for _offset, _pt, _pte_idx, _n_ptes, _pte_covers in pages:
for i in range(_n_ptes):
pte = helper_read_entry_components(_pt.entries[_pte_idx + i])
if self.d.vram[pte['paddr']] != pattern: return False
if pte['valid'] == 0: return False
return True
def random_alloc(self) -> Optional[int]:
if self.total_size - self.alloc_payload < self.min_alloc_size: return None
size = random.randint(self.min_alloc_size, min(self.max_alloc_size, self.total_size - self.alloc_payload))
size = round_up(size, (2 << 20) if size > (4 << 20) else (4 << 10))
try: ptr = self.d.mm.valloc(size)
except MemoryError:
print(f"Failed to allocate {size} bytes. Payload size is {self.alloc_payload}, so fragmenation is {(size / self.total_size)*100.0:.2f}%")
return None
pattern = self.generate_pattern(ptr, size)
pages = self.fill_memory(ptr, size, pattern)
self.allocations[ptr.va_addr] = (size, pattern, pages, ptr)
self.alloc_payload += size
print(f"Allocated {size} bytes at {ptr.va_addr:x}, pattern: {pattern:02x}")
return ptr
def random_free(self) -> bool:
if not self.allocations: return False
ptr = random.choice(list(self.allocations.keys()))
size, pattern, pages, vm = self.allocations[ptr]
# Verify pattern before freeing
if not self.verify_memory(pages, pattern):
raise RuntimeError(f"Memory corruption detected at {vm.va_addr:x}!")
print(f"Freeing {size} bytes at {vm.va_addr:x}, pattern verified: {pattern:02x}")
self.alloc_payload -= size
self.d.mm.vfree(vm)
del self.allocations[ptr]
return True
def run(self):
for i in range(10000000):
if (random.random() < self.alloc_probability or not self.allocations): self.random_alloc()
else: self.random_free()
print("\nCleaning up remaining allocations...")
while self.allocations: self.random_free()
print("Fuzzing completed successfully!")
if __name__ == "__main__":
fuzzer = AMPTFuzzer(1 << 30)
fuzzer.run()

31
artifacts/package_sources/tinygrad/test/external/external_fuzz_beam_timeout_recovery.py vendored Normal file
View File

@@ -0,0 +1,31 @@
#!/usr/bin/env python3
"""
Stress test for beam timeout + device recovery on AM devices.
Usage:
DEV=AMD python test/external/external_fuzz_beam_timeout_recovery.py
"""
from tinygrad import Tensor, Device
from tinygrad.helpers import Context
from tinygrad.runtime.ops_amd import AMDDevice
if __name__ == "__main__":
dev = Device["AMD"]
assert isinstance(dev, AMDDevice) and dev.is_am(), "not am"
N = 10000
for i in range(N):
with Context(DEBUG=0, BEAM=0):
a = Tensor.rand(4096, 4096, device="AMD").contiguous().realize()
b = Tensor.rand(4096, 4096, device="AMD").contiguous().realize()
c = a.matmul(b)
c.realize()
try: dev.synchronize(timeout=1)
except RuntimeError as e: print(e)
with Context(DEBUG=0, BEAM=0):
a = Tensor.ones(512, 512, device="AMD").contiguous().realize()
b = Tensor.ones(512, 512, device="AMD").contiguous().realize()
result = a.matmul(b).realize()[0, 0].item()
assert result == 512.0, f"iter {i}: got {result}"
print(f" iter {i+1}/{N}: ok")
print(f"=== All {N} iterations passed ===")

44
artifacts/package_sources/tinygrad/test/external/external_fuzz_hcq_mp.py vendored Normal file
View File

@@ -0,0 +1,44 @@
import subprocess
import random
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from tinygrad.helpers import getenv
# checks that HCQ drivers can be killed during operation without causing issues
def run_test(i, full_run=False, force_ok=False):
print(f"\rRunning iteration {i}...", end=" ", flush=True)
p = subprocess.Popen(["python3", "test/test_tiny.py", "TestTiny.test_plus"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
if not full_run:
time.sleep(random.uniform(0, 1200) / 1000.0)
p.kill()
_, stderr = p.communicate()
else:
_, stderr = p.communicate()
stderr_text = stderr.decode()
assert ("Ran 1 test in" in stderr_text and "OK" in stderr_text) or (not force_ok and "Failed to take lock file" in stderr_text), stderr_text
if __name__ == "__main__":
max_workers = getenv("MAX_WORKERS", 4)
with ProcessPoolExecutor(max_workers=max_workers) as executor:
futures = []
for i in range(1000000):
if i % 100 == 0:
# wait for everything we launched so far
for f in as_completed(futures):
try:
f.result()
except Exception as e:
print(f"\nError in iteration: {e}")
futures = []
# do a full run in the main proc
run_test(i, True, force_ok=True)
else:
futures.append(executor.submit(run_test, i, bool(getenv("FULL_RUN", 0))))
# keep list small
if len(futures) > max_workers * 2:
futures = [f for f in futures if not f.done()]

31
artifacts/package_sources/tinygrad/test/external/external_fuzz_hcq_signals.py vendored Normal file
View File

@@ -0,0 +1,31 @@
import random
from tinygrad import Device
from tinygrad.helpers import getenv, DEBUG
def main():
seed = getenv("SEED", 1337)
n_gpus = getenv("GPUS", 3)
iters = getenv("ITERS", 10000000)
only_compute = bool(getenv("ONLY_COMPUTE", 0))
print(f"{n_gpus} GPUs for {iters} iterations, {only_compute=}, seed {seed}")
devs = tuple([Device[f"{Device.DEFAULT}:{x}"] for x in range(n_gpus)])
for i in range(iters):
dev = random.choice(devs)
q_t = random.choice([dev.hw_copy_queue_t, dev.hw_compute_queue_t] if not only_compute else [dev.hw_compute_queue_t])
deps_sigs = random.randint(0, len(devs))
wait_devs = random.sample(devs, deps_sigs)
q = q_t()
for d in wait_devs: q.wait(d.timeline_signal, d.timeline_value - 1)
q.wait(dev.timeline_signal, dev.timeline_value - 1).signal(dev.timeline_signal, dev.timeline_value).submit(dev)
dev.timeline_value += 1
if sync:=random.randint(0, 10) < 3: dev.synchronize()
if DEBUG >= 2: print(f"{i}: {q_t} {dev.device_id} timeline {dev.timeline_value}, wait for {[d.device_id for d in wait_devs]}, {sync=}")
elif i % 100 == 0: print(f"\rCompleted {i} iterations", end='')
if __name__ == "__main__":
main()

11
artifacts/package_sources/tinygrad/test/external/external_fuzz_qcom_cpu_cache.py vendored Normal file
View File

@@ -0,0 +1,11 @@
#!/usr/bin/env python3
from tinygrad.tensor import Tensor
import numpy as np
while True:
arr = np.ones(1000000, dtype=np.uint8)
print(f"numpy: {(arr + 1)[:10]}")
ptr = arr.ctypes.data
tensor = Tensor.from_blob(ptr, arr.shape, dtype='uint8', device='QCOM').realize() + 1
print(f"from_blob: {tensor.numpy()[:10]}")

55
artifacts/package_sources/tinygrad/test/external/external_fuzz_sdma_warm_start.py vendored Normal file
View File

@@ -0,0 +1,55 @@
import subprocess, sys, os, random
CHILD_SCRIPT = """
import os, random
import numpy as np
from tinygrad import Tensor, Device
from tinygrad.runtime.ops_amd import AMDDevice
dev = Device["AMD"]
for i in range({N}):
sz = random.randint(1, {MAX_SZ})
data = np.random.randint(0, 256, sz, dtype=np.uint8)
t = Tensor(data, device="AMD").contiguous().realize()
dev.synchronize()
result = t.numpy()
assert (result == data).all(), f"Data mismatch at iter {{i}}"
""".strip()
def run_child(n_ops, max_sz, timeout):
env = os.environ.copy()
env.setdefault("SDMA_RING_SIZE", "4096")
script = CHILD_SCRIPT.format(N=n_ops, MAX_SZ=max_sz)
p = subprocess.Popen([sys.executable, "-c", script], stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env)
try:
_, stderr = p.communicate(timeout=timeout)
return ("ok" if p.returncode == 0 else "fail"), stderr.decode(errors='replace')
except subprocess.TimeoutExpired:
p.kill()
p.communicate()
return "timeout", "TIMEOUT: SDMA ring likely stuck"
if __name__ == "__main__":
n_iters = int(os.environ.get("FUZZ_ITERS", "10000"))
timeout = int(os.environ.get("FUZZ_TIMEOUT", "10"))
max_sz = int(os.environ.get("FUZZ_MAX_SZ", "65536"))
timeouts = 0
failures = 0
for i in range(n_iters):
# Run child with many ops to stress the small sdma ring buffer across warm starts
n_ops = random.randint(20, 100)
status, stderr = run_child(n_ops=n_ops, max_sz=max_sz, timeout=timeout)
if status == "timeout":
timeouts += 1
print(f"\tstderr: {stderr[:500]}")
elif status == "fail":
failures += 1
print(f"\tstderr: {stderr[:500]}")
else:
print(f"iter {i}: ok (n_ops={n_ops})")
print(f"\n=== Results: {n_iters} iterations, {timeouts} timeouts, {failures} failures ===")

81
artifacts/package_sources/tinygrad/test/external/external_fuzz_tlsf.py vendored Normal file
View File

@@ -0,0 +1,81 @@
import random
from typing import Dict, Optional
from tinygrad.helpers import getenv
from tinygrad.runtime.support.memory import TLSFAllocator
class AllocatorFuzzer:
def __init__(self, total_size):
self.total_size = total_size
self.alloc_payload = 0
self.mv = memoryview(bytearray(total_size))
self.alloctor = TLSFAllocator(total_size, block_size=16)
self.allocations: Dict[int, tuple[int, int]] = {} # ptr -> (size, pattern)
self.min_alloc_size = 16
self.max_alloc_size = int(total_size * 0.3)
self.alloc_probability = 0.7
def generate_pattern(self, ptr: int, size: int) -> int: return (ptr * 31 + size * 17) & 0xFF
def fill_memory(self, ptr: int, size: int, pattern: int):
for i in range(min(size, 32)):
self.mv[ptr + i] = pattern
self.mv[ptr + (size - 1 - i)] = pattern
def verify_memory(self, ptr: int, size: int, pattern: int) -> bool:
for i in range(min(size, 32)):
assert self.mv[ptr + i] == pattern
assert self.mv[ptr + (size - 1 - i)] == pattern
return True
def random_alloc(self) -> Optional[int]:
if self.total_size - self.alloc_payload < self.min_alloc_size: return None
size = random.randint(self.min_alloc_size, min(self.max_alloc_size, self.total_size - self.alloc_payload))
try:
ptr = self.alloctor.alloc(size)
except MemoryError:
print(f"Failed to allocate {size} bytes. Payload size is {self.alloc_payload}, so fragmenation is {(size / self.total_size)*100.0:.2f}%")
return None
pattern = self.generate_pattern(ptr, size)
self.fill_memory(ptr, size, pattern)
self.allocations[ptr] = (size, pattern)
self.alloc_payload += size
print(f"Allocated {size} bytes at {ptr:x}, pattern: {pattern:02x}")
return ptr
def random_free(self) -> bool:
if not self.allocations: return False
ptr = random.choice(list(self.allocations.keys()))
size, pattern = self.allocations[ptr]
# Verify pattern before freeing
if not self.verify_memory(ptr, size, pattern):
raise RuntimeError(f"Memory corruption detected at {ptr:x}!")
print(f"Freeing {size} bytes at {ptr:x}, pattern verified: {pattern:02x}")
self.alloc_payload -= size
self.alloctor.free(ptr)
del self.allocations[ptr]
return True
def run(self):
for i in range(getenv("ITERS", 100000)):
if (random.random() < self.alloc_probability or not self.allocations): self.random_alloc()
else: self.random_free()
print("\nCleaning up remaining allocations...")
while self.allocations: self.random_free()
print("Fuzzing completed successfully!")
if __name__ == "__main__":
SEED = getenv("SEED", 42)
random.seed(SEED)
fuzzer = AllocatorFuzzer(1 << 30)
fuzzer.run()

16
artifacts/package_sources/tinygrad/test/external/external_gpu_fail_osx.py vendored Normal file
View File

@@ -0,0 +1,16 @@
# ugh, OS X OpenCL doesn't support half
from tinygrad.runtime.ops_cl import CLDevice, CLProgram, CLCompiler
src = """#pragma OPENCL EXTENSION cl_khr_fp16 : enable
__kernel void max_half(__global half* data0, const __global half* data1) {
int gidx0 = get_group_id(0);
data0[gidx0] = max(data1[gidx0], (half)0.0);
}"""
if __name__ == "__main__":
dev = CLDevice()
print("created device")
lib = CLCompiler(dev, "test").compile(src)
print("created lib", len(lib))
prg = CLProgram(dev, "max_half", lib)
print("created prg")

17
artifacts/package_sources/tinygrad/test/external/external_jit_failure.py vendored Normal file
View File

@@ -0,0 +1,17 @@
from tinygrad import Tensor, TinyJit, Device
import numpy as np
GPUS = 4
N = 128
ds = tuple([Device.canonicalize(f"{Device.DEFAULT}:{i}") for i in range(GPUS)])
t = Tensor.rand(N, N, N).shard(ds, 0)
n = t.numpy()
@TinyJit
def allreduce(t:Tensor) -> Tensor:
return t.sum(0) #.realize()
for i in range(10):
print(i)
tn = allreduce(t).numpy()
np.testing.assert_allclose(tn, n.sum(0), atol=1e-4, rtol=1e-4)

50
artifacts/package_sources/tinygrad/test/external/external_llm_eval.py vendored Normal file
View File

@@ -0,0 +1,50 @@
# eval for OpenAI API server
# uses Meta's exact ARC-Challenge prompt template from lm-evaluation-harness llama3 tasks
import argparse, re, pyarrow.parquet as pq
from openai import OpenAI
from tinygrad.helpers import fetch, colored
LABEL = ["A", "B", "C", "D"]
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--port", "-p", type=int, default=8000)
parser.add_argument("--limit", "-L", type=int, default=None)
parser.add_argument("--max_tokens", "-T", type=int, default=4096)
parser.add_argument("--offset", "-O", type=int, default=0)
parser.add_argument("--temperature", "-t", type=float, default=0.0)
parser.add_argument("--no_think", action="store_true", help="disable thinking (prefills empty think block via assistant message)")
parser.add_argument("--debug", action="store_true")
args = parser.parse_args()
client = OpenAI(base_url=f"http://127.0.0.1:{args.port}/v1", api_key="tinygrad")
dat = fetch("https://huggingface.co/datasets/allenai/ai2_arc/resolve/main/ARC-Challenge/test-00000-of-00001.parquet")
table = pq.read_table(dat)
num_correct, num_answered = 0, 0
# filter to 4-choice questions and normalize labels to A/B/C/D (matches Meta's eval)
rows = [(q, c, a) for q, c, a in zip(table["question"], table["choices"], table["answerKey"]) if len(c["label"]) == 4]
total_questions = min(len(rows), args.offset + args.limit) if args.limit else len(rows)
for question, choices, answer in rows[args.offset:total_questions]:
phrasing = "Given the following question and four candidate answers (A, B, C and D), choose the best answer.\n" +\
f"Question: {question}\n" + '\n'.join([f"{l}. {t}" for l, t in zip(LABEL, choices['text'])]) +\
'\nYour response should end with "The best answer is [the_answer_letter]"' +\
" where the [the_answer_letter] is one of A, B, C or D."
messages = [{"role": "user", "content": phrasing}]
if args.no_think: messages.append({"role": "assistant", "content": "<think>\n\n</think>\n\n"})
resp = client.chat.completions.create(model="test", messages=messages,
max_tokens=args.max_tokens, temperature=args.temperature)
# normalize answer key (some use 1/2/3/4 instead of A/B/C/D)
correct = answer.as_py().strip()
if correct not in LABEL: correct = LABEL[int(correct) - 1]
# extract answer: take last single capital letter A-D from response (prompt asks model to end with the answer)
text = resp.choices[0].message.content.strip()
if args.debug: print(f"\n--- PROMPT ---\n{phrasing}\n--- RESPONSE ---\n{text}\n---")
m = re.findall(r'\b([A-D])\b', text)
given = m[-1] if m else text[:1]
num_correct += correct == given
num_answered += 1
print(f"{num_answered:4d}/{total_questions:4d} "+\
f"Correct Answer: {correct} "+\
f"Given Answer: {colored(given, 'green' if correct==given else 'red')} "+\
f"Percent: {num_correct*100.0/num_answered:.2f}%")

230
artifacts/package_sources/tinygrad/test/external/external_metal_compile_fail.py vendored Normal file
View File

@@ -0,0 +1,230 @@
# AssertionError: Error Domain=AGXMetalG15X_B0 Code=3 "Compiler encountered an internal error"
src = """
#include <metal_stdlib>
using namespace metal;
kernel void r_64_32_8_16_4_6_6_4(device float* data0, const device float* data1,
uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {
int gidx0 = gid.x; /* 64 */
int lidx2 = lid.x; /* 8 */
int gidx1 = gid.y; /* 32 */
int lidx3 = lid.y; /* 16 */
int alu0 = ((gidx0*4096)+(gidx1*16)+(lidx2*512)+lidx3);
int alu1 = ((gidx0*147456)+(gidx1*576)+(lidx2*18432)+(lidx3*36));
float acc0 = 0.0f;
float acc1 = 0.0f;
float acc2 = 0.0f;
float acc3 = 0.0f;
float acc4 = 0.0f;
float acc5 = 0.0f;
float acc6 = 0.0f;
float acc7 = 0.0f;
float acc8 = 0.0f;
float acc9 = 0.0f;
float acc10 = 0.0f;
float acc11 = 0.0f;
float acc12 = 0.0f;
float acc13 = 0.0f;
float acc14 = 0.0f;
float acc15 = 0.0f;
float acc16 = 0.0f;
float acc17 = 0.0f;
float acc18 = 0.0f;
float acc19 = 0.0f;
float acc20 = 0.0f;
float acc21 = 0.0f;
float acc22 = 0.0f;
float acc23 = 0.0f;
float acc24 = 0.0f;
float acc25 = 0.0f;
float acc26 = 0.0f;
float acc27 = 0.0f;
float acc28 = 0.0f;
float acc29 = 0.0f;
float acc30 = 0.0f;
float acc31 = 0.0f;
float acc32 = 0.0f;
float acc33 = 0.0f;
float acc34 = 0.0f;
float acc35 = 0.0f;
for (int ridx0 = 0; ridx0 < 4; ridx0++) {
int alu2 = (ridx0*6);
int alu3 = (alu2+1);
int alu4 = (alu2+2);
int alu5 = (alu2+3);
int alu6 = (alu2+4);
int alu7 = (alu2+5);
int alu8 = (alu2%7);
int alu9 = ((alu8+1)%7);
int alu10 = ((alu8+2)%7);
int alu11 = ((alu8+3)%7);
int alu12 = ((alu8+4)%7);
int alu13 = ((alu8+5)%7);
int alu14 = ((((alu0+(alu3/21))%262144)*144)+(((alu3/7)%3)*3)+(alu9*36));
int alu15 = ((((alu0+(alu4/21))%262144)*144)+(((alu4/7)%3)*3)+(alu10*36));
int alu16 = ((((alu0+(alu5/21))%262144)*144)+(((alu5/7)%3)*3)+(alu11*36));
int alu17 = ((((alu0+(alu6/21))%262144)*144)+(((alu6/7)%3)*3)+(alu12*36));
int alu18 = ((((alu0+(alu7/21))%262144)*144)+(((alu7/7)%3)*3)+(alu13*36));
int alu19 = (alu8%7);
int alu20 = ((((alu0+(alu2/21))%262144)*144)+(((alu2/7)%3)*3)+(alu19*36));
bool alu21 = ((alu2<16)&(alu13<4));
bool alu22 = ((alu2<17)&(alu12<4));
bool alu23 = ((alu2<18)&(alu11<4));
bool alu24 = ((alu2<19)&(alu10<4));
bool alu25 = ((alu2<20)&(alu9<4));
bool alu26 = ((alu2<21)&(alu19<4));
float val0 = (alu25?*(data1+alu14+1):0.0f);
float val1 = (alu25?*(data1+alu14+2):0.0f);
float val2 = (alu25?*(data1+alu14+9):0.0f);
float val3 = (alu25?*(data1+alu14+10):0.0f);
float val4 = (alu25?*(data1+alu14+11):0.0f);
float val5 = (alu25?*(data1+alu14+18):0.0f);
float val6 = (alu25?*(data1+alu14+19):0.0f);
float val7 = (alu25?*(data1+alu14+20):0.0f);
float val8 = (alu25?*(data1+alu14+27):0.0f);
float val9 = (alu25?*(data1+alu14+28):0.0f);
float val10 = (alu25?*(data1+alu14+29):0.0f);
float val11 = (alu24?*(data1+alu15+1):0.0f);
float val12 = (alu24?*(data1+alu15+2):0.0f);
float val13 = (alu24?*(data1+alu15+9):0.0f);
float val14 = (alu24?*(data1+alu15+10):0.0f);
float val15 = (alu24?*(data1+alu15+11):0.0f);
float val16 = (alu24?*(data1+alu15+18):0.0f);
float val17 = (alu24?*(data1+alu15+19):0.0f);
float val18 = (alu24?*(data1+alu15+20):0.0f);
float val19 = (alu24?*(data1+alu15+27):0.0f);
float val20 = (alu24?*(data1+alu15+28):0.0f);
float val21 = (alu24?*(data1+alu15+29):0.0f);
float val22 = (alu23?*(data1+alu16+1):0.0f);
float val23 = (alu23?*(data1+alu16+2):0.0f);
float val24 = (alu23?*(data1+alu16+9):0.0f);
float val25 = (alu23?*(data1+alu16+10):0.0f);
float val26 = (alu23?*(data1+alu16+11):0.0f);
float val27 = (alu23?*(data1+alu16+18):0.0f);
float val28 = (alu23?*(data1+alu16+19):0.0f);
float val29 = (alu23?*(data1+alu16+20):0.0f);
float val30 = (alu23?*(data1+alu16+27):0.0f);
float val31 = (alu23?*(data1+alu16+28):0.0f);
float val32 = (alu23?*(data1+alu16+29):0.0f);
float val33 = (alu22?*(data1+alu17+1):0.0f);
float val34 = (alu22?*(data1+alu17+2):0.0f);
float val35 = (alu22?*(data1+alu17+9):0.0f);
float val36 = (alu22?*(data1+alu17+10):0.0f);
float val37 = (alu22?*(data1+alu17+11):0.0f);
float val38 = (alu22?*(data1+alu17+18):0.0f);
float val39 = (alu22?*(data1+alu17+19):0.0f);
float val40 = (alu22?*(data1+alu17+20):0.0f);
float val41 = (alu22?*(data1+alu17+27):0.0f);
float val42 = (alu22?*(data1+alu17+28):0.0f);
float val43 = (alu22?*(data1+alu17+29):0.0f);
float val44 = (alu21?*(data1+alu18+1):0.0f);
float val45 = (alu21?*(data1+alu18+2):0.0f);
float val46 = (alu21?*(data1+alu18+9):0.0f);
float val47 = (alu21?*(data1+alu18+10):0.0f);
float val48 = (alu21?*(data1+alu18+11):0.0f);
float val49 = (alu21?*(data1+alu18+18):0.0f);
float val50 = (alu21?*(data1+alu18+19):0.0f);
float val51 = (alu21?*(data1+alu18+20):0.0f);
float val52 = (alu21?*(data1+alu18+27):0.0f);
float val53 = (alu21?*(data1+alu18+28):0.0f);
float val54 = (alu21?*(data1+alu18+29):0.0f);
float val55 = (alu26?*(data1+alu20+1):0.0f);
float val56 = (alu26?*(data1+alu20+2):0.0f);
float val57 = (alu26?*(data1+alu20+9):0.0f);
float val58 = (alu26?*(data1+alu20+10):0.0f);
float val59 = (alu26?*(data1+alu20+11):0.0f);
float val60 = (alu26?*(data1+alu20+18):0.0f);
float val61 = (alu26?*(data1+alu20+19):0.0f);
float val62 = (alu26?*(data1+alu20+20):0.0f);
float val63 = (alu26?*(data1+alu20+27):0.0f);
float val64 = (alu26?*(data1+alu20+28):0.0f);
float val65 = (alu26?*(data1+alu20+29):0.0f);
float val66 = (alu25?*(data1+alu14):0.0f);
float val67 = (alu24?*(data1+alu15):0.0f);
float val68 = (alu23?*(data1+alu16):0.0f);
float val69 = (alu22?*(data1+alu17):0.0f);
float val70 = (alu21?*(data1+alu18):0.0f);
float val71 = (alu26?*(data1+alu20):0.0f);
acc0 = (acc0+val71);
acc1 = (acc1+val66);
acc2 = (acc2+val67);
acc3 = (acc3+val68);
acc4 = (acc4+val69);
acc5 = (acc5+val70);
acc6 = (acc6+val57+val55);
acc7 = (acc7+val2+val0);
acc8 = (acc8+val13+val11);
acc9 = (acc9+val24+val22);
acc10 = (acc10+val35+val33);
acc11 = (acc11+val46+val44);
acc12 = (acc12+val60+val58+val56);
acc13 = (acc13+val5+val3+val1);
acc14 = (acc14+val16+val14+val12);
acc15 = (acc15+val27+val25+val23);
acc16 = (acc16+val38+val36+val34);
acc17 = (acc17+val49+val47+val45);
acc18 = (acc18+val63+val61+val59);
acc19 = (acc19+val8+val6+val4);
acc20 = (acc20+val19+val17+val15);
acc21 = (acc21+val30+val28+val26);
acc22 = (acc22+val41+val39+val37);
acc23 = (acc23+val52+val50+val48);
acc24 = (acc24+val64+val62);
acc25 = (acc25+val9+val7);
acc26 = (acc26+val20+val18);
acc27 = (acc27+val31+val29);
acc28 = (acc28+val42+val40);
acc29 = (acc29+val53+val51);
acc30 = (acc30+val65);
acc31 = (acc31+val10);
acc32 = (acc32+val21);
acc33 = (acc33+val32);
acc34 = (acc34+val43);
acc35 = (acc35+val54);
}
*(data0+alu1+1) = acc6;
*(data0+alu1+2) = acc12;
*(data0+alu1+3) = acc18;
*(data0+alu1+4) = acc24;
*(data0+alu1+5) = acc30;
*(data0+alu1+6) = acc1;
*(data0+alu1+7) = acc7;
*(data0+alu1+8) = acc13;
*(data0+alu1+9) = acc19;
*(data0+alu1+10) = acc25;
*(data0+alu1+11) = acc31;
*(data0+alu1+12) = acc2;
*(data0+alu1+13) = acc8;
*(data0+alu1+14) = acc14;
*(data0+alu1+15) = acc20;
*(data0+alu1+16) = acc26;
*(data0+alu1+17) = acc32;
*(data0+alu1+18) = acc3;
*(data0+alu1+19) = acc9;
*(data0+alu1+20) = acc15;
*(data0+alu1+21) = acc21;
*(data0+alu1+22) = acc27;
*(data0+alu1+23) = acc33;
*(data0+alu1+24) = acc4;
*(data0+alu1+25) = acc10;
*(data0+alu1+26) = acc16;
*(data0+alu1+27) = acc22;
*(data0+alu1+28) = acc28;
*(data0+alu1+29) = acc34;
*(data0+alu1+30) = acc5;
*(data0+alu1+31) = acc11;
*(data0+alu1+32) = acc17;
*(data0+alu1+33) = acc23;
*(data0+alu1+34) = acc29;
*(data0+alu1+35) = acc35;
*(data0+alu1) = acc0;
}
"""
from tinygrad.runtime.ops_metal import MetalDevice, MetalCompiler, MetalProgram
if __name__ == "__main__":
dev = MetalDevice("METAL")
lib = MetalCompiler().compile(src)
prg = MetalProgram(dev, "r_64_32_8_16_4_6_6_4", lib)

139
artifacts/package_sources/tinygrad/test/external/external_model_benchmark.py vendored Normal file
View File

@@ -0,0 +1,139 @@
import csv, pathlib, time
import numpy as np
import torch
torch.set_num_threads(1)
import onnxruntime as ort
from onnx2torch import convert
from tinygrad.nn.onnx import OnnxRunner
from tinygrad.helpers import OSX, DEBUG, fetch, getenv
from tinygrad.dtype import _to_np_dtype
from tinygrad import Tensor, Device, Context, dtypes
MODELS = {
"resnet50": "https://github.com/onnx/models/raw/main/validated/vision/classification/resnet/model/resnet50-caffe2-v1-9.onnx",
"openpilot": "https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx",
"efficientnet": "https://github.com/onnx/models/raw/main/validated/vision/classification/efficientnet-lite4/model/efficientnet-lite4-11.onnx",
"shufflenet": "https://github.com/onnx/models/raw/main/validated/vision/classification/shufflenet/model/shufflenet-9.onnx",
# TODO: precision issue
# "commavq": "https://huggingface.co/commaai/commavq-gpt2m/resolve/main/gpt2m.onnx",
"dm": "https://github.com/commaai/openpilot/raw/ba7f840a06dbc8ae3c45b3b4976c88a21895aed0/selfdrive/modeld/models/dmonitoring_model.onnx",
# broken in torch MPS
# "zfnet": "https://github.com/onnx/models/raw/main/archive/vision/classification/zfnet-512/model/zfnet512-9.onnx",
# TypeError: BatchNormalization() got an unexpected keyword argument 'is_test'
# "densenet": "https://github.com/onnx/models/raw/main/archive/vision/classification/densenet-121/model/densenet-3.onnx",
# AssertionError: only onnx version >= 10 supported for slice
# "bert": "https://github.com/onnx/models/raw/main/archive/text/machine_comprehension/bert-squad/model/bertsquad-8.onnx",
# really slow
# "resnet18": "https://github.com/onnx/models/raw/main/archive/vision/classification/resnet/model/resnet18-v2-7.onnx",
}
half_models = ["openpilot", "commavq"]
CSV = {}
open_csv = None
def benchmark(mnm, nm, fxn):
tms = []
for _ in range(3):
st = time.perf_counter_ns()
ret = fxn()
tms.append(time.perf_counter_ns() - st)
print(f"{mnm:15s} {nm:25s} {min(tms)*1e-6:7.2f} ms")
CSV[nm] = min(tms)*1e-6
return min(tms), ret
#BASE = pathlib.Path(__file__).parents[2] / "weights" / "onnx"
BASE = pathlib.Path("/tmp/onnx")
def benchmark_model(m, devices, validate_outs=False):
torch.manual_seed(1)
global open_csv, CSV
CSV = {"model": m}
fn = fetch(MODELS[m])
runner = OnnxRunner(fn)
output_names = runner.graph_outputs
input_shapes = {name: tuple(s if isinstance(s, int) and s != 0 else 1 for s in spec.shape) for name, spec in runner.graph_inputs.items()}
input_types = {name: spec.dtype for name, spec in runner.graph_inputs.items()}
np_inputs = {k:torch.randn(shp).numpy().astype(_to_np_dtype(input_types[k])) for k,shp in input_shapes.items()}
assert len(input_shapes) < 30, f"too many input shapes {len(input_shapes)}"
# print input names
if DEBUG >= 2: print(list(runner.graph_inputs))
for device in devices:
with Context(DEV=device):
inputs = {k:Tensor(inp) for k,inp in np_inputs.items()}
tinygrad_model = runner.to(device)
benchmark(m, f"tinygrad_{device.lower()}_jitless", lambda: {k:v.numpy() for k,v in tinygrad_model(inputs).items()})
from tinygrad.engine.jit import TinyJit
tinygrad_jitted_model = TinyJit(lambda **kwargs: {k:v.realize() for k,v in tinygrad_model(kwargs).items()})
for _ in range(3): {k:v.numpy() for k,v in tinygrad_jitted_model(**inputs).items()}
benchmark(m, f"tinygrad_{device.lower()}_jit", lambda: {k:v.numpy() for k,v in tinygrad_jitted_model(**inputs).items()}) # noqa: F821
del inputs, tinygrad_model, tinygrad_jitted_model
# convert model to torch
try:
torch_model = convert(fn)
except Exception as e:
# model conversion failed
print(f"{m:16s}onnx2torch {type(e).__name__:>25}")
else:
torch_inputs = [torch.tensor(x) for x in np_inputs.values()]
try: benchmark(m, "torch_cpu", lambda: torch_model(*torch_inputs))
except Exception as e: print(f"{m:16s}torch_cpu {type(e).__name__:>25}")
torch_device = "mps" if OSX else "cuda"
torch_mps_model = torch_model.to(torch_device)
torch_mps_inputs = [x.to(torch_device) for x in torch_inputs]
try: benchmark(m, f"torch_{torch_device}", lambda: torch_mps_model(*torch_mps_inputs))
except Exception as e: print(f"{m:16s}torch_{torch_device} {type(e).__name__:>25}")
# bench onnxruntime
ort_options = ort.SessionOptions()
ort_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
ort_options.log_severity_level = 3 # no warnings
for backend in ["CPU", "CUDA" if not OSX else "CoreML"]: # https://onnxruntime.ai/docs/execution-providers/
provider = backend+"ExecutionProvider"
if provider not in ort.get_available_providers(): continue
ort_sess = ort.InferenceSession(str(fn), ort_options, [provider])
try:
benchmark(m, f"onnxruntime_{backend.lower()}", lambda: ort_sess.run(output_names, np_inputs))
except Exception as e: print(f"{m:16s}onnxruntime_{backend.lower()} {type(e).__name__:>25}")
del ort_sess
if validate_outs:
for device in devices:
rtol, atol = 2e-3, 2e-3 # tolerance for fp16 models
with Context(DEV=device):
# force half inputs to float for numerical stability when validating
# this will rely on automatic dtype promotion for converting half weights inside the graph
if m in half_models:
inputs = {k:Tensor(inp, dtype=dtypes.float32) if inp.dtype == np.float16 else Tensor(inp) for k,inp in np_inputs.items()}
else:
inputs = {k:Tensor(inp) for k,inp in np_inputs.items()}
tinygrad_model = runner.to(device)
tinygrad_out = tinygrad_model(inputs)
ort_sess = ort.InferenceSession(str(fn), ort_options, ["CPUExecutionProvider"])
onnx_out = ort_sess.run(output_names, np_inputs)
onnx_out = dict([*list(zip(output_names, onnx_out))])
assert_allclose(tinygrad_out, onnx_out, rtol=rtol, atol=atol)
print(f"{m:16s}outputs validated on {device=} with rtol={rtol:.1e}, atol={atol:.1e}")
if open_csv is None:
open_csv = csv.DictWriter(open('onnx_inference_speed.csv', 'w', newline=''), fieldnames=list(CSV.keys()))
open_csv.writeheader()
open_csv.writerow(CSV)
def assert_allclose(tiny_out:dict, onnx_out:dict, rtol, atol):
assert tiny_out.keys() == onnx_out.keys()
for k in tiny_out.keys():
tiny_v, onnx_v = tiny_out[k], onnx_out[k]
np.testing.assert_allclose(tiny_v.numpy(), onnx_v, rtol=rtol, atol=atol, err_msg=f"For tensor '{k}' in {tiny_out.keys()}")
if __name__ == "__main__":
devices = [Device.DEFAULT] if getenv("NOCLANG") else [Device.DEFAULT, "CPU"]
if (model:=getenv("MODEL", "")) != "": benchmark_model(model, devices, validate_outs=True)
else:
for m in MODELS: benchmark_model(m, devices, validate_outs=True)

79
artifacts/package_sources/tinygrad/test/external/external_multi_gpu.py vendored Normal file
View File

@@ -0,0 +1,79 @@
#!/usr/bin/env python3
# cd extra/disassemblers/ && git clone --recursive github.com:geohot/cuda_ioctl_sniffer.git
# LD_PRELOAD=$PWD/extra/disassemblers/cuda_ioctl_sniffer/out/sniff.so DEV=CL python3 test/external/external_multi_gpu.py
import numpy as np
from tinygrad.tensor import Tensor
from tinygrad.helpers import colored, Timing, getenv
from tinygrad.device import Device
d0, d1 = f'{Device.DEFAULT}:0', f'{Device.DEFAULT}:1'
def sync():
Device[d0].synchronize()
Device[d1].synchronize()
if __name__ == "__main__":
print("GPU devices", d0, d1)
sz = getenv("N", 1024*1024*256) # 1 GB
with Timing("GPU initial sync: "): sync()
with Timing("CPU creation: ", on_exit=lambda x: f", {(sz*4*2)/x:.2f} GB/sec"):
c0 = (Tensor.ones(sz, device="CPU")/2).realize()
c1 = (Tensor.ones(sz, device="CPU")/4).realize()
print(c0.uop.base.realized)
print(c1.uop.base.realized)
with Timing("CPU -> 0: ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
a0 = c0.to(d0).realize()
sync()
with Timing("CPU -> 1: ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
b1 = c1.to(d1).realize()
sync()
# cross copy. this is (sometimes) going through the CPU
with Timing("0 -> 1: ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
a1 = a0.to(d1).realize()
sync()
with Timing("1 -> 0: ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
b0 = b1.to(d0).realize()
sync()
# sum
with Timing("0+0 -> 0 (sum): ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
ab0 = (a0 + b0).realize()
sync()
with Timing("1+1 -> 1 (sum): ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
ab1 = (a1 + b1).realize()
sync()
# cross device sum (does this work?)
with Timing(colored("0+1 -> 0 (sum): ", "red"), on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
abx0 = (a0 + b1.to(d0)).realize()
sync()
with Timing(colored("1+0 -> 1 (sum): ", "red"), on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
abx1 = (b1 + a0.to(d1)).realize()
sync()
# copy back
# NOTE: half of this slowness is caused by allocating memory on the CPU
with Timing("0 -> CPU: ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
cc0 = ab0.numpy()
with Timing("1 -> CPU: ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
cc1 = ab1.numpy()
# same
print("testing")
np.testing.assert_allclose(cc0, cc1)
# same (cross)
print("testing (cross)")
np.testing.assert_allclose(cc0, abx0.numpy())
np.testing.assert_allclose(cc0, abx1.numpy())
# devices
print(ab0)
print(ab1)
print(abx0)
print(abx1)

38
artifacts/package_sources/tinygrad/test/external/external_nhwc_conv.py vendored Normal file
View File

@@ -0,0 +1,38 @@
from tinygrad import Tensor, nn, Context, GlobalCounters
if __name__ == "__main__":
conv = nn.Conv2d(64, 128, 3)
img = Tensor.randn((1,64,128,128))
with Context(DEBUG=0, BEAM=0):
Tensor.realize(img, conv.weight, conv.bias)
tst = conv(img).permute(0,2,3,1).realize()
print(tst.shape)
print("NEW")
img_perm = img.permute(0,2,3,1).contiguous()
print(img_perm.shape)
pp = img_perm.permute(0,3,1,2)._pool((3,3)).permute(0,2,3,4,5,1)
def hwio(pp, conv):
pp = pp.unsqueeze(-1)
weight = conv.weight.permute(2,3,1,0).contiguous()
print(pp.shape, weight.shape, (pp*weight).shape)
return (pp * weight).sum([-4,-3, -2])
def ohwi(pp, conv):
pp = pp.unsqueeze(-4)
weight = conv.weight.permute(0,2,3,1).contiguous()
print(pp.shape, weight.shape, (pp*weight).shape)
return (pp * weight).sum([-3,-2,-1])
for f in [hwio, ohwi]:
GlobalCounters.reset()
print("\n**************", f.__name__, "**************")
out = f(pp, conv)
out.realize()
print(out.shape)
with Context(DEBUG=0, BEAM=0):
err = (tst-out).square()
print(err.mean().item(), err.max().item())

55
artifacts/package_sources/tinygrad/test/external/external_openpilot_image_warp.py vendored Normal file
View File

@@ -0,0 +1,55 @@
import time
from tinygrad.tensor import Tensor, Device
MODEL_WIDTH = 512
MODEL_HEIGHT = 256
MODEL_FRAME_SIZE = MODEL_WIDTH * MODEL_HEIGHT * 3 // 2
IMG_INPUT_SHAPE = (1, 12, 128, 256)
def tensor_arange(end): return Tensor([float(i) for i in range(end)])
def tensor_round(tensor:Tensor): return (tensor + 0.5).floor()
h_src, w_src = 1208, 1928
h_dst, w_dst = MODEL_HEIGHT, MODEL_WIDTH
x = tensor_arange(w_dst).reshape(1, w_dst).expand(h_dst, w_dst)
y = tensor_arange(h_dst).reshape(h_dst, 1).expand(h_dst, w_dst)
ones = Tensor.ones_like(x)
dst_coords = x.reshape((1,-1)).cat(y.reshape((1,-1))).cat(ones.reshape((1,-1)))
def warp_perspective_tinygrad(src:Tensor, M_inv:Tensor) -> Tensor:
src_coords = M_inv @ dst_coords
src_coords = src_coords / src_coords[2:3, :]
x_src = src_coords[0].reshape(h_dst, w_dst)
y_src = src_coords[1].reshape(h_dst, w_dst)
x_nearest = tensor_round(x_src).clip(0, w_src - 1).cast('int')
y_nearest = tensor_round(y_src).clip(0, h_src - 1).cast('int')
# TODO: make 2d indexing fast
idx = y_nearest*src.shape[1] + x_nearest
dst = src.flatten()[idx]
return dst.reshape(h_dst, w_dst)
if __name__ == "__main__":
from tinygrad.engine.jit import TinyJit
update_img_jit = TinyJit(warp_perspective_tinygrad, prune=True)
step_times = []
for _ in range(10):
# regenerate inputs
inputs = [Tensor.randn(1928,1208), Tensor.randn(3,3)]
Tensor.realize(*inputs)
Device.default.synchronize()
# do the warp
st = time.perf_counter()
out = update_img_jit(*inputs)
mt = time.perf_counter()
val = out.contiguous().realize()
Device.default.synchronize()
et = time.perf_counter()
# measure the time
step_times.append((et-st)*1e3)
print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")

41
artifacts/package_sources/tinygrad/test/external/external_osx_profiling.py vendored Normal file
View File

@@ -0,0 +1,41 @@
from tinygrad.runtime.ops_cl import CLProgram, CL, CLBuffer
from tinygrad import dtypes
import time
N = 1000000
a = CLBuffer(N, dtypes.float32)
b = CLBuffer(N, dtypes.float32)
c = CLBuffer(N, dtypes.float32)
prg = CLProgram("test", """__kernel void test(__global float *a, __global float *b, __global float *c) {
int idx = get_global_id(0);
a[idx] = b[idx] + c[idx];
}""")
prg.clprgs[0](CL.cl_queue[0], [N,], None, a._buf, b._buf, c._buf)
t1 = time.monotonic_ns()
e1 = prg.clprgs[0](CL.cl_queue[0], [N,], None, a._buf, b._buf, c._buf)
CL.synchronize()
t2 = time.monotonic_ns()
time.sleep(3)
t3 = time.monotonic_ns()
e2 = prg.clprgs[0](CL.cl_queue[0], [N,], None, a._buf, b._buf, c._buf)
CL.synchronize()
t4 = time.monotonic_ns()
print(e1.profile.queued)
print(e1.profile.submit)
print(e1.profile.start)
print(e1.profile.end)
print(e1, e2)
print(t2-t1, e1.profile.end - e1.profile.start)
print(t4-t3, e2.profile.end - e2.profile.start)
print(t3-t2, e2.profile.queued-e1.profile.end)
print((t3-t2) / (e2.profile.start-e1.profile.end), "ratio")
print("ratio since boot", t1/e1.profile.start)
print(e1.profile.start)
print(e1.profile.end)
print(e2.profile.start)
print(e2.profile.end)

224
artifacts/package_sources/tinygrad/test/external/external_test_am.py vendored Normal file
View File

@@ -0,0 +1,224 @@
import unittest
from tinygrad.runtime.support.am.amdev import AMMemoryManager, AMPageTableEntry
from tinygrad.runtime.support.am.ip import AM_GMC
from tinygrad.runtime.support.hcq import MMIOInterface
from tinygrad.runtime.support.memory import PageTableTraverseContext, AddrSpace
from tinygrad.runtime.autogen.am import am
from tinygrad.helpers import mv_address
class FakeGMC(AM_GMC):
def __init__(self, adev):
self.adev = adev
self.vm_base = 0x0
self.address_space_mask = (1 << 44) - 1
def init_hw(self): pass
def flush_tlb(self, *args, **kwargs): pass
class FakePCIDev:
def __init__(self): self.regions = [(0xc12300000000, 0xc12400000000, 0x0)]
class FakeAM:
def __init__(self):
self.is_booting, self.smi_dev = True, False
self.pcidev = FakePCIDev()
self.vram_size = (512 << 20)
self.vram_mv = memoryview(bytearray(self.vram_size))
self.vram = MMIOInterface(mv_address(self.vram_mv), self.vram_mv.nbytes)
self.gmc = FakeGMC(self)
self.mm = AMMemoryManager(self, self.vram_size, boot_size=(32 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39], va_bits=48,
first_lv=am.AMDGPU_VM_PDB2, va_base=AMMemoryManager.va_allocator.base,
palloc_ranges=[(1 << (i + 12), (2 << 20) if i >= 9 else 0x1000) for i in range(9 * (3 - am.AMDGPU_VM_PDB2), -1, -1)])
self.is_booting = False
self.ip_ver = {am.GC_HWIP: (11, 0, 0)}
def paddr2cpu(self, paddr:int) -> int: return paddr + mv_address(self.vram)
def paddr2mc(self, paddr:int) -> int: return paddr
def paddr2xgmi(self, paddr:int) -> int: return paddr
def xgmi2paddr(self, xgmi_paddr:int) -> int: return xgmi_paddr
# * PTE format:
# * 63:59 reserved
# * 58:57 reserved
# * 56 F
# * 55 L
# * 54 reserved
# * 53:52 SW
# * 51 T
# * 50:48 mtype
# * 47:12 4k physical page base address
# * 11:7 fragment
# * 6 write
# * 5 read
# * 4 exe
# * 3 Z
# * 2 snooped
# * 1 system
# * 0 valid
def helper_read_entry_components(entry_val):
return {"paddr": entry_val & 0x0000FFFFFFFFF000, "fragment":(entry_val >> 7) & 0x1f, "valid": entry_val & 0x1,
"read": (entry_val >> 5) & 0x1, "write": (entry_val >> 6) & 0x1, "exec": (entry_val >> 4) & 0x1,
"mtype": (entry_val >> 48) & 0x7, "T": (entry_val >> 51) & 0x1, "L": (entry_val >> 55) & 0x1, "F": (entry_val >> 56) & 0x1}
def helper_va(va:int): return va + AMMemoryManager.va_allocator.base
class TestAMPageTable(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.d = [FakeAM() for _ in range(2)]
def test_page_table_walkers(self):
mm = self.d[0].mm
for va,sz in [(0x10000, 0x3000), (0x11000, 0x300000), (0x10000, 0x2000), (0x11000, 0x5000),
(0x2000000, 0x2000), (0x4000000, 0x4000000), (0x38000, 0x303000), (0x8000, 0x1000)]:
mm.map_range(vaddr=helper_va(va), size=sz, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
ctx = PageTableTraverseContext(self.d[0], mm.root_page_table, helper_va(va))
results = list(ctx.next(sz))
total_covered = 0
for tup in results:
_offset, _pt, _pte_idx, _n_ptes, _pte_covers = tup
total_covered += _n_ptes * _pte_covers
assert total_covered == sz, f"Expected total coverage {total_covered} to be {sz}"
for tup in results:
_offset, _pt, _pte_idx, _n_ptes, _pte_covers = tup
for i in range(_n_ptes):
pte = helper_read_entry_components(_pt.entries[_pte_idx + i])
assert pte['paddr'] == va + _offset + i * _pte_covers, f"Expected paddr {pte['paddr']:#x} to be {va + _offset + i * _pte_covers:#x}"
assert pte['valid'] == 1
mm.unmap_range(helper_va(va), sz)
for tup in results:
_offset, _pt, _pte_idx, _n_ptes, _pte_covers = tup
for i in range(_n_ptes):
pte = helper_read_entry_components(_pt.entries[_pte_idx + i])
assert pte['paddr'] == 0
assert pte['valid'] == 0
def test_map_notaligned(self):
mm0 = self.d[0].mm
for (va1,sz1),(va2,sz2) in [((0x10000, (0x1000)), (0x11000, (2 << 20)))]:
mm0.map_range(vaddr=helper_va(va1), size=sz1, paddrs=[(va1, sz1)], aspace=AddrSpace.PHYS)
mm0.map_range(vaddr=helper_va(va2), size=sz2, paddrs=[(va2, sz2)], aspace=AddrSpace.PHYS)
mm0.unmap_range(helper_va(va2), sz2)
mm0.unmap_range(helper_va(va1), sz1)
def test_double_map(self):
mm0 = self.d[0].mm
for va,sz in [(0x10000, 0x3000), (0x1000000, 0x1000000), (0x12000, 0x4000)]:
exteranl_va = helper_va(va)
mm0.map_range(vaddr=exteranl_va, size=sz, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
with self.assertRaises(AssertionError):
mm0.map_range(vaddr=exteranl_va, size=0x1000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
with self.assertRaises(AssertionError):
mm0.map_range(vaddr=exteranl_va, size=0x100000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
with self.assertRaises(AssertionError):
mm0.map_range(vaddr=exteranl_va + 0x1000, size=0x1000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
with self.assertRaises(AssertionError):
mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
mm0.unmap_range(vaddr=exteranl_va, size=sz)
# Finally can map and check paddrs
mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(0xdead0000, 0x1000), (0xdead1000, 0xff000)], aspace=AddrSpace.PHYS)
ctx = PageTableTraverseContext(self.d[0], mm0.root_page_table, exteranl_va + 0x2000)
for tup in ctx.next(0x100000):
_offset, _pt, _pte_idx, _n_ptes, _pte_covers = tup
for i in range(_n_ptes):
pte = helper_read_entry_components(_pt.entries[_pte_idx + i])
assert pte['paddr'] == 0xdead0000 + _offset + i * _pte_covers, f"paddr {pte['paddr']:#x} not {0xdead0000 + _offset + i * _pte_covers:#x}"
assert pte['valid'] == 1
mm0.unmap_range(vaddr=exteranl_va + 0x2000, size=0x100000)
def test_try_bad_unmap(self):
mm0 = self.d[0].mm
with self.assertRaises(AssertionError):
mm0.unmap_range(helper_va(0x10000), 0x3000)
mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)], aspace=AddrSpace.PHYS)
mm0.unmap_range(helper_va(0x10000), 0x3000)
with self.assertRaises(AssertionError):
mm0.unmap_range(helper_va(0x10000), 0x3000)
mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)], aspace=AddrSpace.PHYS)
mm0.unmap_range(helper_va(0x10000), 0x3000)
with self.assertRaises(AssertionError):
mm0.unmap_range(helper_va(0x10000), 0x3000)
def test_free_pt(self):
mm0 = self.d[0].mm
# offset from start
for off in [0, 0x3000, 0x10000]:
mm0.map_range(helper_va(0x1000000) + off, (2 << 20) - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000), aspace=AddrSpace.PHYS)
mm0.unmap_range(helper_va(0x1000000) + off, (2 << 20) - off)
mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)], aspace=AddrSpace.PHYS)
mm0.unmap_range(helper_va(0x1000000), 2 << 20)
# offset from end
for off in [0x1000, 0x20000]:
mm0.map_range(helper_va(0x1000000), (2 << 20) - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000), aspace=AddrSpace.PHYS)
mm0.unmap_range(helper_va(0x1000000), (2 << 20) - off)
mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)], aspace=AddrSpace.PHYS)
mm0.unmap_range(helper_va(0x1000000), 2 << 20)
def test_inspect_mode(self):
mm0 = self.d[0].mm
# Map a few disjoint ranges inside a larger region.
mappings = [(0x10000, 0x3000), (0x20000, 0x2000), (0x1000000, 2 << 20)]
for va, sz in mappings:
mm0.map_range(helper_va(va), sz, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
# Inspect over the whole region: should visit all mapped pages.
ctx = PageTableTraverseContext(self.d[0], mm0.root_page_table, helper_va(0x0), inspect=True)
visited = set()
for _off, pt, pte_idx, n_ptes, pte_covers in ctx.next(0x4000000):
for i in range(n_ptes):
pte = helper_read_entry_components(pt.entries[pte_idx + i])
if pte['valid']:
for p in range(0, pte_covers, 0x1000): visited.add(pte['paddr'] + p)
expected_pages = {va + off for va, sz in mappings for off in range(0, sz, 0x1000)}
assert visited == expected_pages
for va, sz in mappings:
mm0.unmap_range(helper_va(va), sz)
# Inspect after unmap: should find no valid entries.
ctx = PageTableTraverseContext(self.d[0], mm0.root_page_table, helper_va(0x0), inspect=True)
for _off, pt, pte_idx, n_ptes, pte_covers in ctx.next(0x4000000):
for i in range(n_ptes): assert not pt.valid(pte_idx + i)
def test_frag_size(self):
mm0 = self.d[0].mm
def must_cover_checker(va, sz):
ans = (1 << (mm0._frag_size(va=va, sz=sz, must_cover=True) + 12))
assert va % ans == 0 and sz % ans == 0 and (va % (2 * ans) != 0 or sz % (2 * ans) != 0), f"va {va:#x} sz {sz:#x} ans {ans:#x}"
def not_cover_checker(va, sz):
ans = (1 << (mm0._frag_size(va=va, sz=sz, must_cover=False) + 12))
assert va % ans == 0 and ans <= sz and (va % (2 * ans) != 0 or (2 * ans) > sz), f"va {va:#x} sz {sz:#x} ans {ans:#x}"
for va, sz in [(0x0, 0x1000), (0x1000, 0x2000), (0x1000, 0x3000), (0x2000, 0x2000), (0x4000, 0x8000), (0x8000, 0x4000), (0x10000, 0x4000),
(0x0, 0x4000), (0x10000, 0x4000), (0x10000, 0x40000), (0x10001000, 0x40000), (0x100001000, 0x3000)]:
must_cover_checker(va, sz)
not_cover_checker(va, sz)
if __name__ == "__main__":
unittest.main()

139
artifacts/package_sources/tinygrad/test/external/external_test_am_fault_recovery.py vendored Normal file
View File

@@ -0,0 +1,139 @@
# ruff: noqa: F405
import unittest, subprocess, os
from tinygrad.helpers import DEV
from tinygrad.runtime.autogen.amd.rdna3.ins import * # noqa: F403
from tinygrad.renderer.amd.dsl import s, v, Inst, NULL
def assemble_kernel(insts:list[Inst], name:str="test") -> str:
kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "wavefront_size32": 1, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8}
from test.amd.disasm import disasm as _disasm
disasm = "\n".join(_disasm(inst) for inst in insts)
hsasrc = f".text\n.globl {name}\n.p2align 8\n.type {name},@function\n{name}:\n{disasm}\n"
return hsasrc + f".rodata\n.p2align 6\n.amdhsa_kernel {name}\n" + "\n".join(f".amdhsa_{k} {v}" for k, v in kd.items()) + "\n.end_amdhsa_kernel"
def _run(code:str, timeout:float=15.0) -> subprocess.CompletedProcess:
# TODO: AM_RESET is required for now, so subprocesses
return subprocess.run(["python", "-c", code], env={**os.environ, "AMD": "1"}, capture_output=True, text=True, timeout=timeout)
def _run_asm(asm_src:str) -> subprocess.CompletedProcess:
return _run('from tinygrad.device import Device, TinyELF; from tinygrad.helpers import Target; '
'from tinygrad.runtime.support.compiler_amd import HIPCompiler; dev = Device["AMD"]; '
f'dev.runtime(TinyELF(HIPCompiler(dev.arch).compile("""{asm_src}"""), "test", Target("AMD", arch=dev.arch), ()))('
'dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True)')
def _verify_recovery() -> subprocess.CompletedProcess:
return _run('from tinygrad import Tensor; t = Tensor([1.0, 2.0], device="AMD").realize(); assert (t + 1).numpy().tolist() == [2.0, 3.0]')
_ILLEGAL_INST_ASM = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n.byte 0xff,0xff,0xff,0xff\ns_endpgm\n" \
".rodata\n.p2align 6\n.amdhsa_kernel test\n.amdhsa_next_free_vgpr 8\n.amdhsa_next_free_sgpr 8\n" \
".amdhsa_wavefront_size32 1\n.amdhsa_user_sgpr_kernarg_segment_ptr 1\n.amdhsa_kernarg_size 8\n.end_amdhsa_kernel"
@unittest.skipIf(DEV.device != "AMD" or not DEV.interface.startswith("MOCK"), "AMD with AM driver required")
class TestAMFaultRecovery(unittest.TestCase):
def _run_kernel(self, insts: list[Inst]) -> subprocess.CompletedProcess: return _run_asm(assemble_kernel(insts))
def _assert_fault_and_recovery(self, result:subprocess.CompletedProcess):
if result.stdout.strip(): print(f"\nstdout: {result.stdout.strip()}")
if result.stderr.strip(): print(f"\nstderr: {result.stderr.strip()}")
self.assertNotEqual(result.returncode, 0, f"Expected fault but succeeded: {result.stdout}")
self.assertEqual(_verify_recovery().returncode, 0)
class TestGlobalMemoryFaults(TestAMFaultRecovery):
def test_global_load_unmapped(self):
insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD),
global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
self._assert_fault_and_recovery(self._run_kernel(insts))
def test_global_store_unmapped(self):
insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 0x12345678),
global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
self._assert_fault_and_recovery(self._run_kernel(insts))
def test_global_null_ptr(self):
insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0),
global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
self._assert_fault_and_recovery(self._run_kernel(insts))
def test_global_misaligned_b64(self):
insts = [v_mov_b32_e32(v[0], 0xBEEF0001), v_mov_b32_e32(v[1], 0xDEAD),
global_load_b64(v[2:3], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
self._assert_fault_and_recovery(self._run_kernel(insts))
def test_global_misaligned_b128(self):
insts = [v_mov_b32_e32(v[0], 0xBEEF0004), v_mov_b32_e32(v[1], 0xDEAD),
global_load_b128(v[2:5], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
self._assert_fault_and_recovery(self._run_kernel(insts))
class TestSMEMFaults(TestAMFaultRecovery):
def test_smem_null_base(self):
insts = [s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(lgkmcnt=0), s_endpgm()]
self._assert_fault_and_recovery(self._run_kernel(insts))
def test_smem_unmapped_address(self):
insts = [s_mov_b32(s[2], 0xBEEF0000), s_mov_b32(s[3], 0xDEAD),
s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(lgkmcnt=0), s_endpgm()]
self._assert_fault_and_recovery(self._run_kernel(insts))
def test_smem_misaligned_b64(self):
insts = [s_mov_b32(s[2], 0xBEEF0004), s_mov_b32(s[3], 0xDEAD),
s_load_b64(s[4:5], s[2:3], 0, soffset=NULL), s_waitcnt(lgkmcnt=0), s_endpgm()]
self._assert_fault_and_recovery(self._run_kernel(insts))
def test_smem_misaligned_b128(self):
insts = [s_mov_b32(s[2], 0xBEEF0004), s_mov_b32(s[3], 0xDEAD),
s_load_b128(s[4:7], s[2:3], 0, soffset=NULL), s_waitcnt(lgkmcnt=0), s_endpgm()]
self._assert_fault_and_recovery(self._run_kernel(insts))
class TestIllegalInstruction(TestAMFaultRecovery):
def test_malformed_encoding(self):
self._assert_fault_and_recovery(_run_asm(_ILLEGAL_INST_ASM))
class TestFlatFaults(TestAMFaultRecovery):
def test_flat_load_unmapped(self):
insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD),
flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0, lgkmcnt=0), s_endpgm()]
self._assert_fault_and_recovery(self._run_kernel(insts))
def test_flat_store_unmapped(self):
insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 0x12345678),
flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(vmcnt=0, lgkmcnt=0), s_endpgm()]
self._assert_fault_and_recovery(self._run_kernel(insts))
class TestAtomicFaults(TestAMFaultRecovery):
def test_global_atomic_unmapped(self):
insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 1),
global_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
self._assert_fault_and_recovery(self._run_kernel(insts))
def test_flat_atomic_unmapped(self):
insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 1),
flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(vmcnt=0, lgkmcnt=0), s_endpgm()]
self._assert_fault_and_recovery(self._run_kernel(insts))
class TestRecovery(TestAMFaultRecovery):
def test_recovery_after_memviol(self):
insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD),
global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
self.assertNotEqual(self._run_kernel(insts).returncode, 0)
self.assertEqual(_verify_recovery().returncode, 0)
def test_recovery_after_illegal_inst(self):
self.assertNotEqual(_run_asm(_ILLEGAL_INST_ASM).returncode, 0)
self.assertEqual(_verify_recovery().returncode, 0)
def test_multiple_faults_recovery(self):
insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD),
global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
for _ in range(3):
self.assertNotEqual(self._run_kernel(insts).returncode, 0)
self.assertEqual(_verify_recovery().returncode, 0)
if __name__ == "__main__":
unittest.main()

28
artifacts/package_sources/tinygrad/test/external/external_test_amd.py vendored Normal file
View File

@@ -0,0 +1,28 @@
import unittest
from tinygrad import Device, Tensor
from tinygrad.schedule import create_schedule
from tinygrad.runtime.ops_amd import AMDDevice
class TestAMD(unittest.TestCase):
@classmethod
def setUpClass(self):
TestAMD.d0: AMDDevice = Device["AMD"]
TestAMD.a = Tensor([0.,1.], device="AMD").realize()
TestAMD.b = self.a + 1
si = create_schedule([self.b.uop])[-1]
TestAMD.d0_runner = TestAMD.d0.get_runner(*si.ast)
TestAMD.b.uop.buffer.allocate()
def test_amd_ring_64bit_doorbell(self):
TestAMD.d0.pm4_write_pointer[0] = TestAMD.d0.pm4_write_pointer[0] + (2 << 32) - TestAMD.d0.pm4_ring.size // 4
for _ in range(2000):
TestAMD.d0_runner.clprg(TestAMD.b.uop.buffer._buf, TestAMD.a.uop.buffer._buf,
global_size=TestAMD.d0_runner.global_size, local_size=TestAMD.d0_runner.local_size)
TestAMD.d0_runner.clprg(TestAMD.a.uop.buffer._buf, TestAMD.b.uop.buffer._buf,
global_size=TestAMD.d0_runner.global_size, local_size=TestAMD.d0_runner.local_size)
val = TestAMD.a.uop.buffer.as_memoryview().cast("f")[0]
assert val == 4000.0, f"got val {val}"
if __name__ == "__main__":
unittest.main()

172
artifacts/package_sources/tinygrad/test/external/external_test_datasets.py vendored Normal file
View File

@@ -0,0 +1,172 @@
from extra.datasets.kits19 import iterate, preprocess
from examples.mlperf.dataloader import batch_load_unet3d, batch_load_retinanet
from test.external.mlperf_retinanet.coco_utils import get_openimages
from test.external.mlperf_retinanet.openimages import postprocess_targets
from test.external.mlperf_retinanet.presets import DetectionPresetTrain, DetectionPresetEval
from test.external.mlperf_retinanet.model.transform import GeneralizedRCNNTransform
from test.external.mlperf_unet3d.kits19 import PytTrain, PytVal
from tinygrad.helpers import temp
from pathlib import Path
from pycocotools.coco import COCO
import json
import nibabel as nib
import numpy as np
import os
import PIL
import random
import tempfile
import torch
import unittest
class ExternalTestDatasets(unittest.TestCase):
def _set_seed(self):
np.random.seed(42)
random.seed(42)
torch.manual_seed(42)
class TestKiTS19Dataset(ExternalTestDatasets):
def _create_samples(self, val, num_samples=2):
self._set_seed()
img, lbl = np.random.rand(190, 392, 392).astype(np.float32), np.random.randint(0, 100, size=(190, 392, 392)).astype(np.uint8)
img, lbl = nib.Nifti1Image(img, np.eye(4)), nib.Nifti1Image(lbl, np.eye(4))
dataset = "val" if val else "train"
preproc_pth = Path(tempfile.gettempdir() + f"/{dataset}")
for i in range(num_samples):
os.makedirs(tempfile.gettempdir() + f"/case_000{i}", exist_ok=True)
nib.save(img, temp(f"case_000{i}/imaging.nii.gz"))
nib.save(lbl, temp(f"case_000{i}/segmentation.nii.gz"))
preproc_img, preproc_lbl = preprocess(Path(tempfile.gettempdir()) / f"case_000{i}")
preproc_img_pth, preproc_lbl_pth = temp(f"{dataset}/case_000{i}_x.npy"), temp(f"{dataset}/case_000{i}_y.npy")
os.makedirs(preproc_pth, exist_ok=True)
np.save(preproc_img_pth, preproc_img, allow_pickle=False)
np.save(preproc_lbl_pth, preproc_lbl, allow_pickle=False)
return preproc_pth, list(preproc_pth.glob("*_x.npy")), list(preproc_pth.glob("*_y.npy"))
def _create_ref_dataloader(self, preproc_img_pths, preproc_lbl_pths, val):
if val:
dataset = PytVal(preproc_img_pths, preproc_lbl_pths)
else:
dataset = PytTrain(preproc_img_pths, preproc_lbl_pths, patch_size=(128, 128, 128), oversampling=0.4)
return iter(dataset)
def _create_tinygrad_dataloader(self, preproc_pth, val, batch_size=1, shuffle=False, seed=42, use_old_dataloader=False):
if use_old_dataloader:
dataset = iterate(list(Path(tempfile.gettempdir()).glob("case_*")), preprocessed_dir=preproc_pth, val=val, shuffle=shuffle, bs=batch_size)
else:
dataset = batch_load_unet3d(preproc_pth, batch_size=batch_size, val=val, shuffle=shuffle, seed=seed)
return iter(dataset)
@unittest.skip("flaky")
def test_training_set(self):
preproc_pth, preproc_img_pths, preproc_lbl_pths = self._create_samples(False)
ref_dataset = self._create_ref_dataloader(preproc_img_pths, preproc_lbl_pths, False)
tinygrad_dataset = self._create_tinygrad_dataloader(preproc_pth, False)
for ref_sample, tinygrad_sample in zip(ref_dataset, tinygrad_dataset):
self._set_seed()
np.testing.assert_equal(tinygrad_sample[0][:, 0].numpy(), ref_sample[0])
np.testing.assert_equal(tinygrad_sample[1][:, 0].numpy(), ref_sample[1])
def test_validation_set(self):
preproc_pth, preproc_img_pths, preproc_lbl_pths = self._create_samples(True)
ref_dataset = self._create_ref_dataloader(preproc_img_pths, preproc_lbl_pths, True)
tinygrad_dataset = self._create_tinygrad_dataloader(preproc_pth, True, use_old_dataloader=True)
for ref_sample, tinygrad_sample in zip(ref_dataset, tinygrad_dataset):
np.testing.assert_equal(tinygrad_sample[0][:, 0], ref_sample[0])
np.testing.assert_equal(tinygrad_sample[1], ref_sample[1])
class TestOpenImagesDataset(ExternalTestDatasets):
@classmethod
def setUpClass(cls):
cls.img_mean = [0.485, 0.456, 0.406]
cls.img_std = [0.229, 0.224, 0.225]
cls.img_size = (800, 800)
def _create_samples(self, subset):
os.makedirs((base_dir := Path(tempfile.gettempdir() + "/openimages")) / f"{subset}/data", exist_ok=True)
os.makedirs(base_dir / Path(f"{subset}/labels"), exist_ok=True)
lbls, img_size = ["cls_1", "cls_2"], (447, 1024)
cats = [{"id": i, "name": c, "supercategory": None} for i, c in enumerate(lbls)]
imgs = [
{
"id": i, "file_name": f"img_{i}.jpg",
"height": img_size[0], "width": img_size[1],
"subset": subset, "license": None, "coco_url": None
}
for i in range(len(lbls))
]
annots = [
{
"id": i, "image_id": i,
"category_id": 0, "bbox": [23.217183744, 31.75409775, 964.1241282560001, 326.09017434000003],
"area": 314391.4050683996, "IsOccluded": 0,
"IsInside": 0, "IsDepiction": 0,
"IsTruncated": 0, "IsGroupOf": 0,
"iscrowd": 0
}
for i in range(len(lbls))
]
info = {"dataset": "openimages_mlperf", "version": "v6"}
coco_annotations = {"info": info, "licenses": [], "categories": cats, "images": imgs, "annotations": annots}
with open(ann_file:=base_dir / Path(f"{subset}/labels/openimages-mlperf.json"), "w") as fp:
json.dump(coco_annotations, fp)
for i in range(len(lbls)):
img = PIL.Image.new("RGB", img_size[::-1])
img.save(base_dir / Path(f"{subset}/data/img_{i}.jpg"))
return base_dir, ann_file
def _create_ref_dataloader(self, base_dir, ann_file, subset):
self._set_seed()
transforms = DetectionPresetTrain("hflip") if subset == "train" else DetectionPresetEval()
return iter(get_openimages(ann_file.stem, base_dir, subset, transforms))
def _create_tinygrad_dataloader(self, base_dir, ann_file, subset, batch_size=1, seed=42):
return iter(batch_load_retinanet(COCO(ann_file), subset == "validation", base_dir, batch_size=batch_size, shuffle=False, seed=seed))
def _normalize_img(self, img):
return ((img / 255.0) - np.array(self.img_mean)) / np.array(self.img_std)
def test_training_set(self):
base_dir, ann_file = self._create_samples((subset:="train"))
transform = GeneralizedRCNNTransform(self.img_size, self.img_mean, self.img_std)
anchors = torch.ones((120087, 4))
tinygrad_dataloader = self._create_tinygrad_dataloader(base_dir, ann_file, subset)
ref_dataloader = self._create_ref_dataloader(base_dir, ann_file, subset)
for ((tinygrad_img, tinygrad_boxes, tinygrad_labels, _, _, _), (ref_img, ref_tgt)) in zip(tinygrad_dataloader, ref_dataloader):
ref_img, ref_tgt = transform(ref_img.unsqueeze(0), [ref_tgt])
ref_tgt = postprocess_targets(ref_tgt, anchors.unsqueeze(0))
ref_boxes, ref_labels = ref_tgt[0]["boxes"], ref_tgt[0]["labels"]
np.testing.assert_allclose(self._normalize_img(tinygrad_img.numpy()), ref_img.tensors.transpose(1, 3).numpy(), rtol=1e-6)
np.testing.assert_equal(tinygrad_boxes[0].numpy(), ref_boxes.numpy())
np.testing.assert_equal(tinygrad_labels[0].numpy(), ref_labels.numpy())
def test_validation_set(self):
base_dir, ann_file = self._create_samples((subset:="validation"))
transform = GeneralizedRCNNTransform(self.img_size, self.img_mean, self.img_std)
tinygrad_dataloader = self._create_tinygrad_dataloader(base_dir, ann_file, subset)
ref_dataloader = self._create_ref_dataloader(base_dir, ann_file, "val")
for ((tinygrad_img, _, _, _), (ref_img, _)) in zip(tinygrad_dataloader, ref_dataloader):
ref_img, _ = transform(ref_img.unsqueeze(0))
np.testing.assert_allclose(self._normalize_img(tinygrad_img.numpy()), ref_img.tensors.transpose(1, 3).numpy(), rtol=1e-6)
if __name__ == '__main__':
unittest.main()

26
artifacts/package_sources/tinygrad/test/external/external_test_dev_var.py vendored Normal file
View File

@@ -0,0 +1,26 @@
import subprocess, unittest, os, sys
from tinygrad.device import Device
class TestTinygradSlow(unittest.TestCase):
def test_env_overwrite_default_device(self):
subprocess.run([f'DEV={Device.DEFAULT} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
shell=True, check=True)
if Device.DEFAULT != "CPU":
# setting device via DEV
subprocess.run([f'DEV={Device.DEFAULT.capitalize()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
shell=True, check=True)
subprocess.run([f'DEV={Device.DEFAULT.lower()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
shell=True, check=True)
subprocess.run([f'DEV={Device.DEFAULT.upper()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
shell=True, check=True)
class TestRunAsModule(unittest.TestCase):
def test_module_runs(self):
p = subprocess.run([sys.executable, "-m", "tinygrad.device"],stdout=subprocess.PIPE, stderr=subprocess.PIPE,
env={**os.environ, "DEBUG": "1"}, timeout=40,)
out = (p.stdout + p.stderr).decode()
self.assertEqual(p.returncode, 0, msg=out)
if __name__ == '__main__':
unittest.main()

74
artifacts/package_sources/tinygrad/test/external/external_test_example.py vendored Normal file
View File

@@ -0,0 +1,74 @@
import unittest, sys
from tinygrad import Device
from tinygrad.tensor import Tensor
from tinygrad.helpers import getenv, OSX
def multidevice_test(fxn):
exclude_devices = getenv("EXCLUDE_DEVICES", "").split(",")
def ret(self):
for device in Device._devices:
# broken on OSX USB AMD, why?
if device in ["DISK", "NPY", "FAKE", "DSP", "NULL"] or (OSX and device in ["AMD"]): continue
if sys.stdout.isatty(): print(device)
if device in exclude_devices:
if sys.stdout.isatty(): print(f"WARNING: {device} test is excluded")
continue
with self.subTest(device=device):
try:
Device[device]
except Exception:
if sys.stdout.isatty(): print(f"WARNING: {device} test isn't running")
continue
fxn(self, device)
return ret
class TestExample(unittest.TestCase):
@multidevice_test
def test_convert_to_cpu(self, device):
a = Tensor([[1,2],[3,4]], device=device)
assert a.numpy().shape == (2,2)
b = a.to("CPU")
assert b.numpy().shape == (2,2)
@multidevice_test
def test_2_plus_3(self, device):
a = Tensor([2], device=device)
b = Tensor([3], device=device)
result = a + b
print(f"{a.numpy()} + {b.numpy()} = {result.numpy()}")
assert result.numpy()[0] == 5.
@multidevice_test
def test_example_readme(self, device):
x = Tensor.eye(3).clone().to(device)
y = Tensor([[2.0,0,-2.0]], device=device)
z = y.matmul(x).sum()
z.backward()
x.grad.numpy() # dz/dx
y.grad.numpy() # dz/dy
assert x.grad.device == device
assert y.grad.device == device
@multidevice_test
def test_example_matmul(self, device):
try:
Device[device]
except Exception:
print(f"WARNING: {device} test isn't running")
return
x = Tensor.eye(8).clone().to(device)
y = Tensor.eye(8).clone().to(device)
z = y.matmul(x).sum()
z.backward()
x.grad.numpy() # dz/dx
y.grad.numpy() # dz/dy
assert x.grad.device == device
assert y.grad.device == device
if __name__ == '__main__':
unittest.main()

128
artifacts/package_sources/tinygrad/test/external/external_test_gpu_crash.py vendored Normal file
View File

@@ -0,0 +1,128 @@
# ruff: noqa: F405
"""Tests for GPU crash scenarios using AMD assembly to trigger invalid operations.
These tests intentionally cause GPU faults to verify error handling.
Run with: DEV=AMD python -m pytest test/external/external_test_gpu_crash.py -v
"""
import unittest, re, importlib
from tinygrad import Device, UOp, dtypes
from tinygrad.engine.realize import run_linear
from tinygrad.uop.ops import Ops, KernelInfo
from tinygrad.renderer.amd.dsl import s, v, Inst, NULL
RDNA3_CDNA3_MAP = {"v_mov_b32_e32": "v_mov_b32_e32", "s_mov_b32": "s_mov_b32", "s_waitcnt": "s_waitcnt", "s_endpgm": "s_endpgm",
"global_load_b32": "global_load_dword", "global_store_b32": "global_store_dword",
"global_atomic_add_u32": "global_atomic_add", "flat_load_b32": "flat_load_dword",
"flat_store_b32": "flat_store_dword", "flat_atomic_add_u32": "flat_atomic_add", "s_load_b32": "s_load_dword"}
@unittest.skipIf(Device.DEFAULT != "AMD", "AMD required")
class TestGPUCrash(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.dev = Device["AMD"]
cls.is_cdna = cls.dev.target[0] < 10
ins = importlib.import_module('tinygrad.runtime.autogen.amd.' + ('cdna' if cls.is_cdna else 'rdna3') + '.ins')
for rdna3_name, cdna3_name in RDNA3_CDNA3_MAP.items():
setattr(cls, rdna3_name, getattr(ins, cdna3_name if cls.is_cdna else rdna3_name))
def setUp(self):
# Verify device works before each test
from tinygrad import Tensor
try:
t = Tensor([1.0, 2.0], device="AMD").realize()
assert (t + 1).numpy().tolist() == [2.0, 3.0]
except Exception:
self.fail("Device not working before test")
def _run_insts(self, insts: list[Inst]):
buf = UOp.new_buffer("AMD", 64, dtypes.uint8)
sink = UOp.sink(UOp.param(0, dtypes.uint8, (64,), device="AMD"), UOp.special(1, "lidx0"), arg=KernelInfo("test"))
prg = UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=i) for i in insts))))
run_linear(UOp(Ops.LINEAR, src=(prg.call(buf),)), wait=True)
def _assert_gpu_fault(self, func):
"""Assert that func raises a RuntimeError indicating a GPU fault (not a setup error)."""
with self.assertRaises(RuntimeError) as cm:
func()
err_msg = str(cm.exception).lower()
# Verify it's a GPU fault, not a setup/device initialization error
self.assertTrue(
re.search(r'fault|hang|timeout|illegal|memviol', err_msg),
f"Expected GPU fault error, got: {cm.exception}"
)
class TestOutOfBoundsMemoryAccess(TestGPUCrash):
"""Tests for out-of-bounds memory accesses."""
def test_global_load_null_ptr(self):
"""Global load from NULL pointer."""
insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0),
self.global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
def test_global_store_null_ptr(self):
"""Global store to NULL pointer."""
insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 0xDEADBEEF),
self.global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
def test_global_load_unmapped_high_address(self):
"""Global load from high unmapped address (0xDEAD00000000)."""
insts = [self.v_mov_b32_e32(v[0], 0x00000000), self.v_mov_b32_e32(v[1], 0xDEAD),
self.global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
def test_global_store_unmapped_high_address(self):
"""Global store to high unmapped address."""
insts = [self.v_mov_b32_e32(v[0], 0x00000000), self.v_mov_b32_e32(v[1], 0xDEAD), self.v_mov_b32_e32(v[2], 0x12345678),
self.global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
def test_global_atomic_unmapped(self):
"""Atomic operation on unmapped memory."""
insts = [self.v_mov_b32_e32(v[0], 0xBEEF0000), self.v_mov_b32_e32(v[1], 0xDEAD), self.v_mov_b32_e32(v[2], 1),
self.global_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
class TestSMEMFaults(TestGPUCrash):
"""Tests for scalar memory (SMEM) faults."""
def test_smem_load_null(self):
"""SMEM load from NULL base."""
insts = [self.s_mov_b32(s[2], 0), self.s_mov_b32(s[3], 0),
self.s_load_b32(s[4], s[2:3], 0, soffset=NULL), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
def test_smem_load_unmapped(self):
"""SMEM load from unmapped address."""
insts = [self.s_mov_b32(s[2], 0xBEEF0000), self.s_mov_b32(s[3], 0xDEAD),
self.s_load_b32(s[4], s[2:3], 0, soffset=NULL), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
class TestFlatMemoryFaults(TestGPUCrash):
"""Tests for FLAT memory instruction faults."""
def test_flat_load_null(self):
"""FLAT load from NULL address."""
insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0),
self.flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
def test_flat_store_null(self):
"""FLAT store to NULL address."""
insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 0xDEADBEEF),
self.flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
def test_flat_atomic_null(self):
"""FLAT atomic on NULL address."""
insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 1),
self.flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
if __name__ == "__main__":
unittest.main()

328
artifacts/package_sources/tinygrad/test/external/external_test_hcq.py vendored Normal file
View File

@@ -0,0 +1,328 @@
import unittest, ctypes, struct, time, array
from tinygrad import Device, Tensor, dtypes
from tinygrad.helpers import to_mv, DEV
from tinygrad.device import Buffer, BufferSpec
from tinygrad.engine.realize import get_runtime
from tinygrad.codegen import to_program
def _time_queue(q, d):
st = time.perf_counter()
q.signal(d.timeline_signal, d.timeline_value)
q.submit(d)
d._wait_signal(d.timeline_signal, d.timeline_value)
d.timeline_value += 1
return time.perf_counter() - st
@unittest.skipUnless(Device.DEFAULT in ["NV", "AMD"], "Runs only on NV or AMD")
class TestHCQ(unittest.TestCase):
@classmethod
def setUpClass(self):
TestHCQ.d0 = Device[Device.DEFAULT]
#TestHCQ.d1: AMDDevice = Device["AMD:1"]
TestHCQ.a = Tensor([0.,1.], device=Device.DEFAULT).realize()
TestHCQ.b = self.a + 1
linear = self.b.schedule_linear()
TestHCQ.prg = to_program(linear.src[-1].src[0], TestHCQ.d0.renderer)
TestHCQ.runtime = get_runtime(TestHCQ.d0.device, TestHCQ.prg)
TestHCQ.b.uop.buffer.allocate()
# wow that's a lot of abstraction layers
TestHCQ.addr = struct.pack("QQ", TestHCQ.b.uop.buffer._buf, TestHCQ.a.uop.buffer._buf)
TestHCQ.addr2 = struct.pack("QQ", TestHCQ.a.uop.buffer._buf, TestHCQ.b.uop.buffer._buf)
TestHCQ.kernargs_off = TestHCQ.runtime.kernargs_offset
TestHCQ.kernargs_size = TestHCQ.runtime.kernargs_alloc_size
ctypes.memmove(TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_off, TestHCQ.addr, len(TestHCQ.addr))
ctypes.memmove(TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size+TestHCQ.kernargs_off, TestHCQ.addr2, len(TestHCQ.addr2))
if Device.DEFAULT == "AMD":
from tinygrad.runtime.ops_amd import HWQueue, HWPM4Queue
TestHCQ.compute_queue = HWPM4Queue
TestHCQ.copy_queue = HWQueue
elif Device.DEFAULT == "NV":
from tinygrad.runtime.ops_nv import HWQueue, HWQueue
# nv need to copy constbuffer there as well
to_mv(TestHCQ.d0.kernargs_ptr, 0x160).cast('I')[:] = array.array('I', TestHCQ.runtime.constbuffer_0)
to_mv(TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, 0x160).cast('I')[:] = array.array('I', TestHCQ.runtime.constbuffer_0)
TestHCQ.compute_queue = HWQueue
TestHCQ.copy_queue = HWQueue
def setUp(self):
TestHCQ.d0.synchronize()
TestHCQ.a.uop.buffer.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 1)))))
TestHCQ.b.uop.buffer.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 0)))))
TestHCQ.d0.synchronize() # wait for copyins to complete
def test_run_1000_times_one_submit(self):
temp_signal, temp_value = TestHCQ.d0._alloc_signal(value=0), 0
q = TestHCQ.compute_queue()
for _ in range(1000):
q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
q.signal(temp_signal, temp_value + 1).wait(temp_signal, temp_value + 1)
temp_value += 1
q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
q.signal(temp_signal, temp_value + 1).wait(temp_signal, temp_value + 1)
temp_value += 1
q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
q.submit(TestHCQ.d0)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
assert val == 2000.0, f"got val {val}"
def test_run_1000_times(self):
temp_signal = TestHCQ.d0._alloc_signal(value=0)
q = TestHCQ.compute_queue()
q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
q.signal(temp_signal, 2).wait(temp_signal, 2)
q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.prg.arg.global_size,
TestHCQ.prg.arg.local_size)
for _ in range(1000):
TestHCQ.d0._set_signal(temp_signal, 1)
q.submit(TestHCQ.d0)
TestHCQ.compute_queue().signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
assert val == 2000.0, f"got val {val}"
def test_run_to_3(self):
temp_signal = TestHCQ.d0._alloc_signal(value=0)
q = TestHCQ.compute_queue()
q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
q.signal(temp_signal, 1).wait(temp_signal, 1)
q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
q.signal(temp_signal, 2).wait(temp_signal, 2)
q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
assert val == 3.0, f"got val {val}"
def test_update_exec(self):
q = TestHCQ.compute_queue()
exec_cmd_idx = len(q)
q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
q.update_exec(exec_cmd_idx, (1,1,1), (1,1,1))
q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
assert val == 1.0, f"got val {val}"
val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
assert val == 0.0, f"got val {val}, should not be updated"
@unittest.skipUnless(Device.DEFAULT == "NV", "Only NV supports bind")
def test_bind_run(self):
temp_signal = TestHCQ.d0._alloc_signal(value=0)
q = TestHCQ.compute_queue()
q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
q.signal(temp_signal, 2).wait(temp_signal, 2)
q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.prg.arg.global_size,
TestHCQ.prg.arg.local_size)
q.bind(TestHCQ.d0)
for _ in range(1000):
TestHCQ.d0._set_signal(temp_signal, 1)
q.submit(TestHCQ.d0)
TestHCQ.compute_queue().signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
assert val == 2000.0, f"got val {val}"
@unittest.skipUnless(Device.DEFAULT == "NV", "Only NV supports bind")
def test_update_exec_binded(self):
q = TestHCQ.compute_queue()
exec_ptr = q.ptr()
q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
q.bind(TestHCQ.d0)
q.update_exec(exec_ptr, (1,1,1), (1,1,1))
q.submit(TestHCQ.d0)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
assert val == 1.0, f"got val {val}"
val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
assert val == 0.0, f"got val {val}, should not be updated"
@unittest.skipIf(DEV.interface.startswith("MOCK"), "Can't handle async update on CPU")
def test_wait_signal(self):
temp_signal = TestHCQ.d0._alloc_signal(value=0)
TestHCQ.compute_queue().wait(temp_signal, value=1).signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
with self.assertRaises(RuntimeError):
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value, timeout=50)
# clean up
TestHCQ.d0._set_signal(temp_signal, 1)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value, timeout=100)
TestHCQ.d0.timeline_value += 1
@unittest.skipIf(DEV.interface.startswith("MOCK"), "Can't handle async update on CPU")
def test_wait_copy_signal(self):
temp_signal = TestHCQ.d0._alloc_signal(value=0)
TestHCQ.copy_queue().wait(temp_signal, value=1).signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
with self.assertRaises(RuntimeError):
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value, timeout=50)
# clean up
TestHCQ.d0._set_signal(temp_signal, 1)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value, timeout=100)
TestHCQ.d0.timeline_value += 1
def test_run_normal(self):
q = TestHCQ.compute_queue()
q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
assert val == 1.0, f"got val {val}"
def test_submit_empty_queues(self):
TestHCQ.compute_queue().submit(TestHCQ.d0)
TestHCQ.copy_queue().submit(TestHCQ.d0)
def test_signal_timeout(self):
with self.assertRaises(RuntimeError):
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value, timeout=50)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value + 122, timeout=50)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1, timeout=50)
def test_signal(self):
new_timeline_value = TestHCQ.d0.timeline_value + 0xff
TestHCQ.compute_queue().signal(TestHCQ.d0.timeline_signal, new_timeline_value).submit(TestHCQ.d0)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, new_timeline_value)
TestHCQ.d0.timeline_value = new_timeline_value + 1 # update to not break runtime
def test_copy_signal(self):
new_timeline_value = TestHCQ.d0.timeline_value + 0xff
TestHCQ.copy_queue().signal(TestHCQ.d0.timeline_signal, new_timeline_value).submit(TestHCQ.d0)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, new_timeline_value)
TestHCQ.d0.timeline_value = new_timeline_value + 1 # update to not break runtime
def test_run_signal(self):
q = TestHCQ.compute_queue()
q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
q.submit(TestHCQ.d0)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
assert val == 1.0, f"got val {val}"
def test_copy_1000_times(self):
q = TestHCQ.copy_queue()
q.copy(TestHCQ.a.uop.buffer._buf, TestHCQ.b.uop.buffer._buf, 8)
q.copy(TestHCQ.b.uop.buffer._buf, TestHCQ.a.uop.buffer._buf, 8)
for _ in range(1000):
q.submit(TestHCQ.d0)
TestHCQ.copy_queue().signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
# confirm the signal didn't exceed the put value
with self.assertRaises(RuntimeError):
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value + 1, timeout=50)
val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
assert val == 0.0, f"got val {val}"
def test_copy(self):
q = TestHCQ.copy_queue()
q.copy(TestHCQ.b.uop.buffer._buf, TestHCQ.a.uop.buffer._buf, 8)
q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
q.submit(TestHCQ.d0)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
assert val == 1.0, f"got val {val}"
@unittest.skipUnless(Device.DEFAULT == "NV", "Only NV supports bind")
def test_bind_copy(self):
q = TestHCQ.copy_queue()
q.copy(TestHCQ.a.uop.buffer._buf, TestHCQ.b.uop.buffer._buf, 8)
q.copy(TestHCQ.b.uop.buffer._buf, TestHCQ.a.uop.buffer._buf, 8)
q.bind(TestHCQ.d0)
for _ in range(1000):
q.submit(TestHCQ.d0)
TestHCQ.copy_queue().signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
# confirm the signal didn't exceed the put value
with self.assertRaises(RuntimeError):
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value + 1, timeout=50)
val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
assert val == 0.0, f"got val {val}"
def test_copy_bandwidth(self):
# THEORY: the bandwidth is low here because it's only using one SDMA queue. I suspect it's more stable like this at least.
SZ = 2_000_000_000
a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
b = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
q = TestHCQ.copy_queue()
q.copy(a._buf, b._buf, SZ)
et = _time_queue(q, TestHCQ.d0)
gb_s = (SZ/1e9)/et
print(f"same device copy: {et*1e3:.2f} ms, {gb_s:.2f} GB/s")
assert 0.3 <= gb_s <= 1000
def test_cross_device_copy_bandwidth(self):
SZ = 2_000_000_000
b = Buffer(f"{Device.DEFAULT}:1", SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
TestHCQ.d0._gpu_map(b._buf)
q = TestHCQ.copy_queue()
q.copy(a._buf, b._buf, SZ)
et = _time_queue(q, TestHCQ.d0)
gb_s = (SZ/1e9)/et
print(f"cross device copy: {et*1e3:.2f} ms, {gb_s:.2f} GB/s")
assert 0.3 <= gb_s <= 50
def test_interleave_compute_and_copy(self):
q = TestHCQ.compute_queue()
qc = TestHCQ.copy_queue()
q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size) # b = [1, 2]
q.signal(sig:=TestHCQ.d0._alloc_signal(value=0), value=1)
qc.wait(sig, value=1)
qc.copy(TestHCQ.a.uop.buffer._buf, TestHCQ.b.uop.buffer._buf, 8)
qc.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
qc.submit(TestHCQ.d0)
time.sleep(0.02) # give it time for the wait to fail
q.submit(TestHCQ.d0)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
assert val == 1.0, f"got val {val}"
def test_cross_device_signal(self):
d1 = Device[f"{Device.DEFAULT}:1"]
q1 = TestHCQ.compute_queue()
q2 = TestHCQ.compute_queue()
q1.signal(sig:=TestHCQ.d0._alloc_signal(value=0), value=0xfff)
q2.wait(sig, value=0xfff)
q2.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
q2.submit(TestHCQ.d0)
q1.signal(d1.timeline_signal, d1.timeline_value)
q1.submit(d1)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
d1._wait_signal(d1.timeline_signal, d1.timeline_value)
d1.timeline_value += 1
def test_timeline_signal_rollover(self):
# NV 64bit, AMD 32bit
TestHCQ.d0.timeline_value = (1 << 64) - 20 if Device.DEFAULT == "NV" else (1 << 32) - 20 # close value to reset
TestHCQ.compute_queue().signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1).submit(TestHCQ.d0)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1)
for _ in range(40):
q = TestHCQ.compute_queue()
q.wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1)
q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
assert val == 1.0, f"got val {val}"
if __name__ == "__main__":
unittest.main()

20
artifacts/package_sources/tinygrad/test/external/external_test_hcq_open.py vendored Normal file
View File

@@ -0,0 +1,20 @@
import os
if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
import unittest, time
from tinygrad import Device
class TestOpen(unittest.TestCase):
def generate_test_open(n):
def test(self):
dev = Device[Device.DEFAULT]
for i in range(10):
dev.allocator.alloc(10 << 20)
time.sleep(0.5)
test.__name__ = f'test_open_{n}'
return test
for i in range(64): locals()[f'test_open_{i}'] = generate_test_open(i)
if __name__ == '__main__':
unittest.main()

39
artifacts/package_sources/tinygrad/test/external/external_test_hip_compile.py vendored Normal file
View File

@@ -0,0 +1,39 @@
import time, unittest
from tinygrad.runtime.support.hip_comgr import compile_hip
from tinygrad import Tensor
from tinygrad.device import Device
from tinygrad.schedule import create_schedule
from tinygrad.codegen.opt.kernel import Kernel
class TestHIPCompileSpeed(unittest.TestCase):
@unittest.skipIf(Device.DEFAULT != "HIP", "only run on HIP")
def test_hip_compile(self):
a, b = Tensor([1,2,3,4,5]), Tensor([1,2,3,4,5])
out = a + b
lin = Kernel(create_schedule([out.uop])[-1].ast[0])
lin.to_program()
reference = """
#include <hip/hip_common.h>
typedef long unsigned int size_t;
extern "C" __attribute__((device)) __attribute__((const)) size_t __ockl_get_local_id(unsigned int);
extern "C" __attribute__((device)) __attribute__((const)) size_t __ockl_get_group_id(unsigned int);
extern "C" __attribute__((device)) __attribute__((const)) size_t __ockl_get_local_size(unsigned int);
extern "C" __attribute__((global))void {name}(int* data0, const int* data1, const int* data2) {{
int gidx0 = __ockl_get_group_id(0); /* 5 */
int val0 = data1[gidx0];
int val1 = data2[gidx0];
data0[gidx0] = (val0+val1);
}}
"""
def time_compile(code):
st = time.perf_counter()
compile_hip(code)
return (time.perf_counter() - st) * 1000
tinygrad_tm = min([time_compile(Device[Device.DEFAULT].renderer.render(f"test{i}", lin.uops)) for i in range(10)])
ref_tm = min([time_compile(reference.format(name=f"test{i}")) for i in range(10)])
print(f"tinygrad {tinygrad_tm:6.2f} ms")
print(f"reference {ref_tm:6.2f} ms")
assert (tinygrad_tm - ref_tm) <= 10

43
artifacts/package_sources/tinygrad/test/external/external_test_hive_reset.py vendored Executable file
View File

@@ -0,0 +1,43 @@
#!/usr/bin/env python3
import subprocess, sys
from tinygrad.helpers import getenv
LOOPS = getenv("LOOPS", 50)
BROKEN = getenv("BROKEN", 0)
ONLY_RESET = getenv("ONLY_RESET", 0)
BROKEN_KERNEL_SCRIPT = """
from tinygrad.device import Device
from tinygrad.runtime.ops_amd import AMDProgram, AMDDevice
from tinygrad.runtime.support.compiler_amd import compile_hip
dev = Device["AMD"]
assert isinstance(dev, AMDDevice) and dev.is_am(), "Need AM driver (not KFD)"
broken_src = '''
extern "C" __attribute__((global)) void broken(int* dummy) {
volatile int* bad_ptr = (volatile int*)0xDEAD00000000ULL;
*bad_ptr = 0x42;
}
'''
broken_lib = compile_hip(broken_src, dev.arch)
broken_prg = AMDProgram(dev, "broken", broken_lib)
buf = dev.allocator.alloc(64)
try:
broken_prg(buf, global_size=(1,1,1), local_size=(1,1,1), wait=True)
print(" ERROR: Kernel did not fault!")
except RuntimeError as e:
print(f" Got expected error: {e}")
"""
for i in range(LOOPS):
print(f"=== Running hive_reset.py ({i+1}/{LOOPS}) ===")
subprocess.run([sys.executable, "extra/amdpci/hive_reset.py"], check=True)
print("=== hive_reset complete ===")
if BROKEN:
print(f"=== Running broken kernel ({i+1}/{LOOPS}) ===")
ret = subprocess.run([sys.executable, "-c", BROKEN_KERNEL_SCRIPT])
print(f"=== broken kernel exited with code {ret.returncode} ===")
elif not ONLY_RESET:
print(f"=== Running test_tiny.py ({i+1}/{LOOPS}) ===")
ret = subprocess.run([sys.executable, "test/test_tiny.py", "TestTiny.test_plus"])
print(f"=== test_tiny.py exited with code {ret.returncode} ===")

51
artifacts/package_sources/tinygrad/test/external/external_test_image.py vendored Normal file
View File

@@ -0,0 +1,51 @@
#!/usr/bin/env python
import os
import unittest
import numpy as np
if 'IMAGE' not in os.environ:
os.environ['IMAGE'] = '2'
os.environ['CL'] = '1'
os.environ['OPT'] = '2'
from tinygrad.tensor import Tensor
from tinygrad.nn import Conv2d
class TestImage(unittest.TestCase):
def test_create_image(self):
t = Tensor.ones(128, 128, 1)
t = t.reshape(128, 32, 4) + 3
t.realize()
np.testing.assert_array_equal(t.numpy(), np.ones((128,32,4))*4)
def test_sum_image(self):
t1 = Tensor.ones(16, 16, 1).reshape(16, 4, 4) + 3
t1.realize()
t1 = t1.sum()
t1.realize()
assert t1.numpy() == 16*4*4*4, f"got {t1.numpy()}"
def test_add_image(self):
t1 = Tensor.ones(16, 16, 1).reshape(16, 4, 4) + 3
t2 = Tensor.ones(16, 16, 1).reshape(16, 4, 4) + 4
t1.realize()
t2.realize()
t3 = t1 + t2
t3.realize()
np.testing.assert_array_equal(t3.numpy(), np.ones((16,4,4))*9)
def test_padded_conv(self):
bs, in_chans, out_chans = 1,12,32
tiny_conv = Conv2d(in_chans, out_chans, 3, bias=None, padding=1)
tiny_dat = Tensor.ones(bs, 12, 64, 128)
tiny_conv(tiny_dat).realize()
def test_op_conv(self):
bs, in_chans, out_chans = 1,12,32
tiny_conv = Conv2d(in_chans, out_chans, 3, bias=None, padding=1)
tiny_dconv = Conv2d(out_chans, out_chans, 1, bias=None, padding=0)
tiny_dat = Tensor.ones(bs, 12, 64, 128)
p2 = tiny_conv(tiny_dat).relu()
p2 = tiny_dconv(p2)
p2.realize()
if __name__ == '__main__':
unittest.main()

40
artifacts/package_sources/tinygrad/test/external/external_test_jit_on_models.py vendored Normal file
View File

@@ -0,0 +1,40 @@
#!/usr/bin/env python
import unittest
import numpy as np
from tinygrad import Tensor, dtypes
from tinygrad.engine.jit import TinyJit
from tinygrad.helpers import Context
from test.helpers import derandomize_model
from examples.llama import Transformer
def helper_test_jitted_correctness(gen, train, train_jit):
nojit = train(*gen()).numpy()
for _ in range(5): jit = train_jit(*gen()).numpy()
np.testing.assert_allclose(nojit, jit, rtol=1e-3, atol=1e-5)
class TestJittedModels(unittest.TestCase):
def test_jitted_tiny_llama(self):
self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
args_tiny = {"dim": 1024, "hidden_dim": 1024, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-05, "vocab_size": 1000}
model = Transformer(**args_tiny)
derandomize_model(model)
def test(t): return model(t, 0).realize()
@TinyJit
def test_jit(t): return model(t, 0).realize()
helper_test_jitted_correctness(lambda: (Tensor([[1,]]),), test, test_jit)
def test_jitted_stable_diffusion(self):
from examples.stable_diffusion import UNetModel, unet_params
model = UNetModel(**unet_params)
derandomize_model(model)
def test(t, t2): return model(t, 801, t2).realize()
@TinyJit
def test_jit(t, t2): return model(t, 801, t2).realize()
helper_test_jitted_correctness(lambda: (Tensor.randn(1, 4, 16, 16),Tensor.randn(1, 77, 768)), test, test_jit)
if __name__ == "__main__":
unittest.main()

31
artifacts/package_sources/tinygrad/test/external/external_test_keccak.py vendored Normal file
View File

@@ -0,0 +1,31 @@
import unittest, zipfile, re
from tinygrad import Tensor
from tinygrad.helpers import fetch, tqdm
SHA3_URL = "https://csrc.nist.gov/CSRC/media/Projects/Cryptographic-Algorithm-Validation-Program/documents/sha3/sha-3bytetestvectors.zip"
SHAKE_URL = "https://csrc.nist.gov/CSRC/media/Projects/Cryptographic-Algorithm-Validation-Program/documents/sha3/shakebytetestvectors.zip"
class TestExternalKeccak(unittest.TestCase):
def test_sha3_224(self): self.check_nist_vectors(SHA3_URL, ["SHA3_224LongMsg.rsp", "SHA3_224ShortMsg.rsp"], "sha3_224")
def test_sha3_256(self): self.check_nist_vectors(SHA3_URL, ["SHA3_256LongMsg.rsp", "SHA3_256ShortMsg.rsp"], "sha3_256")
def test_shake_128(self): self.check_nist_vectors(SHAKE_URL, ["SHAKE128LongMsg.rsp", "SHAKE128ShortMsg.rsp"], "shake_128")
def check_nist_vectors(self, url: str, filenames: list[str], preset: str):
pattern = r"Len\s*=\s*(?P<Len>\d+)\s+Msg\s*=\s*(?P<Msg>[0-9a-fA-F\s]+)\s+(MD|Output)\s*=\s*(?P<Output>[0-9a-fA-F]+)"
vecs_zip = fetch(url)
for filename in filenames:
vecs = zipfile.ZipFile(vecs_zip).open(filename).read().decode()
vectors = [ (l, bytes.fromhex(match["Msg"].lower()), bytes.fromhex(match["Output"].lower()))
for match in re.finditer(pattern, vecs) if (l:=int(match["Len"])) < 8192 ]
self.assertTrue(len(vectors) > 0)
print("file", filename)
for data_len, data, output in tqdm(vectors):
tinyout = bytes(Tensor(data[:data_len//8]).keccak(preset).data())
self.assertEqual(tinyout, output)
if __name__ == '__main__':
unittest.main()

12
artifacts/package_sources/tinygrad/test/external/external_test_llama3_ff.py vendored Normal file
View File

@@ -0,0 +1,12 @@
#!/usr/bin/env python3
from tinygrad import Tensor, TinyJit, nn
from extra.models.llama import FeedForward
if __name__ == "__main__":
model = FeedForward(4096, 14336)
for x in nn.state.get_parameters(model): x.replace(x.half()).realize()
jrun = TinyJit(model)
for i in range(5):
print(f"*** run {i}")
jrun(Tensor.rand(1, 4096))

24
artifacts/package_sources/tinygrad/test/external/external_test_llama3_layer.py vendored Normal file
View File

@@ -0,0 +1,24 @@
#!/usr/bin/env python3
from tinygrad import Tensor, TinyJit, nn, dtypes
from tinygrad.helpers import getenv
from extra.models.llama import TransformerBlock, precompute_freqs_cis
BS = getenv("BS", 1)
SEQLEN = getenv("SEQLEN", 128)
# DEFAULT_FLOAT=bfloat16 SEQLEN=8192 ASM_GEMM=1 HK_FLASH_ATTENTION=1 DEV=NULL:HIP:gfx950 DEBUG=2 VIZ=1 PYTHONPATH="."
# python test/external/external_test_llama3_layer.py
if __name__ == "__main__":
dim, hidden_dim, n_heads, n_kv_heads, norm_eps = 4096, 14336, 32, 8, 1e-5
layer = TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, max_context=0)
for x in nn.state.get_parameters(layer): x.replace(x.cast(dtypes.default_float)).realize()
freqs_cis = precompute_freqs_cis(dim // n_heads, SEQLEN, theta=500000.0).contiguous().realize()
@TinyJit
def run(t): return layer(t, 0, freqs_cis, None)
for i in range(5):
print(f"*** run {i}")
run(Tensor.rand(BS, SEQLEN, dim, dtype=dtypes.default_float).realize())

10
artifacts/package_sources/tinygrad/test/external/external_test_lm_head.py vendored Normal file
View File

@@ -0,0 +1,10 @@
from tinygrad import Tensor, nn
if __name__ == "__main__":
vocab_size = 50257
n_embd = 768
lm_head = nn.Linear(n_embd, vocab_size, bias=False)
bs = 4
seq_len = 1024
x = Tensor.rand(bs, seq_len, n_embd)
ret = lm_head(x).realize()

40
artifacts/package_sources/tinygrad/test/external/external_test_losses.py vendored Normal file
View File

@@ -0,0 +1,40 @@
from tinygrad import Tensor
from test.external.mlperf_retinanet.focal_loss import sigmoid_focal_loss as ref_sigmoid_focal_loss
from test.external.mlperf_unet3d.dice import DiceCELoss
from examples.mlperf.losses import dice_ce_loss, sigmoid_focal_loss, l1_loss
import numpy as np
import torch
import unittest
class ExternalTestLosses(unittest.TestCase):
def setUp(self):
np.random.seed(42)
def _assert_loss(self, pred, tgt, tinygrad_metrics, ref_metrics, rtol=1e-07, atol=0, **kwargs):
tinygrad_metrics_res = tinygrad_metrics(Tensor(pred), Tensor(tgt), **kwargs)
ref_metrics_res = ref_metrics(torch.from_numpy(pred), torch.from_numpy(tgt), **kwargs)
np.testing.assert_allclose(tinygrad_metrics_res.numpy(), ref_metrics_res.numpy(), rtol=rtol, atol=atol)
def test_dice_ce_loss(self):
pred, label = np.random.rand(1, 3, 128, 128, 128).astype(np.float32), np.ones((1, 1, 128, 128, 128)).astype(np.uint8)
tinygrad_metrics_res, ref_metrics_res = dice_ce_loss, DiceCELoss(True, True, "NCDHW", False)
self._assert_loss(pred, label, tinygrad_metrics_res, ref_metrics_res, atol=1e-4)
def test_sigmoid_focal_loss(self):
def _apply_logit(p): return np.log(p / (1 - p))
pred, tgt = _apply_logit(np.random.rand(5,2).astype(np.float32)), np.random.randint(0, 2, (5, 2)).astype(np.float32)
for reduction in ["mean", "sum", "none"]:
for alpha, gamma in zip([-1, 0.58], [0, 2]):
self._assert_loss(pred, tgt, sigmoid_focal_loss, ref_sigmoid_focal_loss, rtol=1e-4, alpha=alpha, gamma=gamma, reduction=reduction)
def test_l1_loss(self):
N, C, H, W = 3, 4, 5, 6
shapes = ((N, C), (N, C, H), (N, C, H, W))
for reduction in ["mean", "sum", "none"]:
for shape in shapes:
pred, tgt = np.random.randint(shape).astype(np.float32), np.random.randint(shape)
self._assert_loss(pred, tgt, l1_loss, torch.nn.functional.l1_loss, reduction=reduction)
if __name__ == '__main__':
unittest.main()

67
artifacts/package_sources/tinygrad/test/external/external_test_lr_schedule.py vendored Normal file
View File

@@ -0,0 +1,67 @@
import unittest, math
import numpy as np
from tinygrad.tensor import Tensor
from tinygrad.nn.optim import AdamW
from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup, LambdaLR, LambdaLinearScheduler
np.random.seed(1337)
x_init = np.random.randn(1,4).astype(np.float32)
W_init = np.random.randn(4,4).astype(np.float32)
m_init = np.random.randn(1,4).astype(np.float32)
class TinyNet:
def __init__(self):
self.x = Tensor(x_init.copy())
self.W = Tensor(W_init.copy())
self.m = Tensor(m_init.copy())
def forward(self):
out = self.x.matmul(self.W).relu()
out = out.log_softmax(1)
out = out.mul(self.m).add(self.m).sum()
return out
class TestCosineAnnealingLRWithWarmup(unittest.TestCase):
# only tests the lr
def _test_lr(self, base_lr, end_lr, warmup_steps, decay_steps):
net = TinyNet()
optim = AdamW([net.W], lr=0.0)
tiny_lr = CosineAnnealingLRWithWarmup(optim, base_lr, end_lr, warmup_steps, decay_steps)
lr = []
for _ in range(warmup_steps+decay_steps):
lr.append(optim.lr.item())
tiny_lr.step()
# reimplemented in python
expected = []
for i in range(warmup_steps): expected.append((i+1)/warmup_steps*base_lr)
for i in range(decay_steps): expected.append(end_lr+(base_lr-end_lr)*(1+math.cos((i+1)/decay_steps*math.pi))/2)
np.testing.assert_allclose(lr, expected, rtol=1e-5)
def test_lr_0(self): self._test_lr(3e-4, 8e-5, 3, 5)
def test_lr_1(self): self._test_lr(3e-4, 8e-5, 10, 20)
def test_lr_llama3(self): self._test_lr(8e-5, 8e-7, 20, 100)
class TestLambdaLRLinearWarmup(unittest.TestCase):
def test_linear_lr_warmup(self):
BS, BASE_LR = 304, 2.5e-7
lr = BS * BASE_LR
# Use a dummy Tensor parameter for optimizer because the lr_scheduler only needs the optimizer's device and lr, the params aren't touched.
optimizer = AdamW([Tensor([1.])])
lambda_lr_callback = LambdaLinearScheduler(1000, 1.0, 1.0, 1e-06, 10000000000000).schedule
lr_scheduler = LambdaLR(optimizer, Tensor(lr, device=optimizer.device), lambda_lr_callback)
lrs = {}
# with above settings, optimizer.lr should warm up to lr over 1000 steps linearly
for i in range(1200):
lr_scheduler.step()
if i in {0, 499, 998, 999, 1000, 1199}:
lrs[i] = optimizer.lr.item()
np.testing.assert_allclose(lr, lrs[999], rtol=0, atol=1e-11)
np.testing.assert_equal(lrs[999], lrs[1000])
np.testing.assert_equal(lrs[999], lrs[1199])
np.testing.assert_allclose(lrs[999] / lrs[0], 1000, rtol=0, atol=1)
np.testing.assert_allclose(lrs[999] / lrs[499], 2, rtol=0, atol=1e-5)
if __name__ == '__main__':
unittest.main()

24
artifacts/package_sources/tinygrad/test/external/external_test_mamba.py vendored Normal file
View File

@@ -0,0 +1,24 @@
import unittest
from test.helpers import slow
from examples.mamba import Mamba, generate
from transformers import AutoTokenizer
PROMPT = 'Why is gravity '
TOKENIZER = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
@slow
class TestMamba(unittest.TestCase):
def test_mamba_130M(self):
OUT_130M = '''Why is gravity \nnot a good idea?\n\nA:'''
model = Mamba.from_pretrained('130m')
tinyoutput = generate(model, TOKENIZER, PROMPT, n_tokens_to_gen=10)
self.assertEqual(OUT_130M, tinyoutput)
del model
def test_mamba_370M(self):
OUT_370M = '''Why is gravity \nso important?\nBecause it's the only'''
model = Mamba.from_pretrained('370m')
tinyoutput = generate(model, TOKENIZER, PROMPT, n_tokens_to_gen=10)
self.assertEqual(OUT_370M, tinyoutput)
del model
if __name__ == '__main__':
unittest.main()

42
artifacts/package_sources/tinygrad/test/external/external_test_metrics.py vendored Normal file
View File

@@ -0,0 +1,42 @@
from tinygrad import Tensor
from test.external.mlperf_unet3d.dice import DiceScore
from examples.mlperf.metrics import dice_score, log_perplexity
import numpy as np
import torch
import unittest, math
class ExternalTestMetrics(unittest.TestCase):
def _test_metrics(self, tinygrad_metrics, orig_metrics, pred, label, atol=1e-8, rtol=1e-7):
tinygrad_metrics_res = tinygrad_metrics(Tensor(pred), Tensor(label)).squeeze().numpy()
orig_metrics_res = orig_metrics(torch.from_numpy(pred), torch.from_numpy(label)).numpy()
np.testing.assert_allclose(tinygrad_metrics_res, orig_metrics_res, atol=atol, rtol=rtol)
def test_dice(self):
pred, label = np.random.rand(1, 3, 128, 128, 128).astype(np.float32), np.ones((1, 1, 128, 128, 128)).astype(np.uint8)
self._test_metrics(dice_score, DiceScore(), pred, label)
def test_log_perplexity(self):
# equally likely
np.testing.assert_allclose(log_perplexity(Tensor([[[1.0, 1, 1, 1]]]), Tensor([[2]])).numpy(), math.log(4))
np.testing.assert_allclose(log_perplexity(Tensor([[[1.0]*256]*32]), Tensor([[2]*32])).numpy(), math.log(256), rtol=1e-6)
# pretty correct and incorrect
np.testing.assert_allclose(log_perplexity(Tensor([[[10000., 0, 0, 0]]]), Tensor([[0]])).numpy(), 0)
np.testing.assert_allclose(log_perplexity(Tensor([[[0.0, 10000, 10000, 10000]]]), Tensor([[0]])).numpy(), 10000, rtol=1e-3)
# higher logit -> lower loss
x = Tensor([[[4.0, 3, 2, 1]]])
for i in range(x.numel()-1): self.assertLess(log_perplexity(x, Tensor([[i]])).item(), log_perplexity(x, Tensor([[i+1]])).item())
# torch eval examples
np.testing.assert_allclose(
log_perplexity(Tensor([[[0.3659, 0.7025, 0.3104], [0.0097, 0.6577, 0.1947]]]), Tensor([[2, 1]])).exp().numpy(),
2.7593, rtol=1e-5)
np.testing.assert_allclose(
log_perplexity(Tensor([[[0.3, 0.7, 0.3, 0.1], [0.5, 0.4, 0.1, 0.4],[0.1, 0.1, 0.2, 0.5]],
[[0.1, 0.6, 0.1, 0.5], [0.3, 0.7, 0.3, 0.4], [0.3, 0.7, 0.3, 0.4]]]), Tensor([[2, 1, 3], [1, 0, 1]])).exp().numpy(),
3.6216, rtol=1e-5)
np.testing.assert_allclose(
log_perplexity(Tensor([[[0.3659, 0.7025, 0.3104], [0.0097, 0.6577, 0.1947]]]), Tensor([[2, 1]]), ignore_index=1).exp().numpy(),
3.5372, rtol=1e-4)
if __name__ == '__main__':
unittest.main()

8
artifacts/package_sources/tinygrad/test/external/external_test_mnist_data_select.py vendored Normal file
View File

@@ -0,0 +1,8 @@
#!/bin/bash
from tinygrad import Tensor
from extra.datasets import fetch_mnist
if __name__ == "__main__":
X_train, Y_train, X_test, Y_test = fetch_mnist(tensors=True)
samples = Tensor.randint(512, high=X_train.shape[0])
select = X_train[samples].realize()

210
artifacts/package_sources/tinygrad/test/external/external_test_onnx_backend.py vendored Normal file
View File

@@ -0,0 +1,210 @@
import unittest
from typing import Any, Tuple
from onnx.backend.base import Backend, BackendRep
import onnx.backend.test
import numpy as np
from tinygrad import Tensor, Device, dtypes
from tinygrad.helpers import getenv, OSX
from tinygrad.nn.onnx import OnnxRunner
# pip3 install tabulate
pytest_plugins = 'onnx.backend.test.report',
class TinygradModel(BackendRep):
def __init__(self, run_onnx, input_names):
super().__init__()
self.fxn = run_onnx
self.input_names = input_names
def run(self, inputs: Any, **kwargs: Any) -> Tuple[Any, ...]:
real_inputs = dict(zip(self.input_names, inputs))
ret = self.fxn(real_inputs, debug=2)
return tuple(x.numpy() if isinstance(x, Tensor) else [i.numpy() for i in x] if isinstance(x, list) else np.array(x) for x in ret.values())
class TinygradBackend(Backend):
@classmethod
def prepare(cls, model: onnx.ModelProto, device):
input_all = [x.name for x in model.graph.input]
input_initializer = [x.name for x in model.graph.initializer]
net_feed_input = [x for x in input_all if x not in input_initializer]
print("prepare", cls, device, net_feed_input)
model = Tensor(model.SerializeToString(), device="PYTHON")
run_onnx = OnnxRunner(model)
return TinygradModel(run_onnx, net_feed_input)
@classmethod
def supports_device(cls, device: str) -> bool:
# NOTE: this is onnx CPU
return device == "CPU"
backend_test = onnx.backend.test.BackendTest(TinygradBackend, __name__)
# BUG: buggy onnx tests
backend_test.exclude('test_adam_multiple_cpu')
# BUG: ORT fails these with runtime error
backend_test.exclude('test_PReLU_1d_multiparam_cpu')
backend_test.exclude('test_PReLU_2d_multiparam_cpu')
backend_test.exclude('test_PReLU_3d_multiparam_cpu')
# BUG: we don't match ORT here due to some div inaccuracy with floats
backend_test.exclude('test_dynamicquantizelinear_cpu')
backend_test.exclude('test_dynamicquantizelinear_expanded_cpu')
# BUG: ORT fails these with numerical error but we match ORT numerically
# see: https://onnx.ai/backend-scoreboard/onnxruntime_details_stable.html
# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_qlinearmatmul_2D_int8_float16
backend_test.exclude('test_qlinearmatmul_2D_int8_float16_cpu')
# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_qlinearmatmul_3D_int8_float16
backend_test.exclude('test_qlinearmatmul_3D_int8_float16_cpu')
# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_qlinearmatmul_2D_int8_float32
backend_test.exclude('test_qlinearmatmul_2D_int8_float32_cpu')
# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_qlinearmatmul_3D_int8_float32
backend_test.exclude('test_qlinearmatmul_3D_int8_float32_cpu')
# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_maxunpool_export_with_output_shape
backend_test.exclude('test_maxunpool_export_with_output_shape_cpu')
# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_resize_downsample_scales_linear_align_corners
backend_test.exclude('test_resize_downsample_scales_linear_align_corners_cpu')
# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_resize_downsample_scales_cubic_align_corners
backend_test.exclude('test_resize_downsample_scales_cubic_align_corners_cpu')
# about different dtypes
if dtypes.float64 not in Device[Device.DEFAULT].renderer.supported_dtypes():
backend_test.exclude('float64')
backend_test.exclude('DOUBLE')
# these have float64 inputs
backend_test.exclude('test_eyelike_with_dtype_cpu')
backend_test.exclude('test_reduce_log_sum_exp*')
backend_test.exclude('test_operator_add*')
backend_test.exclude('test_einsum_*')
backend_test.exclude('test_cumsum_*')
if dtypes.float16 not in Device[Device.DEFAULT].renderer.supported_dtypes():
backend_test.exclude('float16')
backend_test.exclude('FLOAT16')
# dtype cast
backend_test.exclude('STRING')
backend_test.exclude('FLOAT8')
backend_test.exclude('INT4')
backend_test.exclude('UINT4')
backend_test.exclude('BFLOAT16') # not supported in numpy
backend_test.exclude('FLOAT4E2M1')
backend_test.exclude('test_dequantizelinear_int4_cpu')
backend_test.exclude('test_dequantizelinear_uint4_cpu')
backend_test.exclude('test_quantizelinear_int4_cpu')
backend_test.exclude('test_quantizelinear_uint4_cpu')
# no support for FLOAT8
backend_test.exclude('test_quantizelinear_e4m3fn_cpu')
backend_test.exclude('test_quantizelinear_e5m2_cpu')
backend_test.exclude('test_quantizelinear_e4m3fn_cpu')
backend_test.exclude('test_quantizelinear_e5m2_cpu')
backend_test.exclude('test_quantizelinear_float4e2m1_cpu')
backend_test.exclude('test_dequantizelinear_e4m3fn_cpu')
backend_test.exclude('test_dequantizelinear_e4m3fn_zero_point_cpu')
backend_test.exclude('test_dequantizelinear_e4m3fn_float16_cpu')
backend_test.exclude('test_dequantizelinear_e5m2_cpu')
backend_test.exclude('test_dequantizelinear_float4e2m1_cpu')
# we don't support indexes
# no support for int pow
backend_test.exclude('test_pow_types_int32_int32_cpu')
backend_test.exclude('test_pow_types_int64_int64_cpu')
# no boolean ops (2d, 3d, 4d)
backend_test.exclude('test_bitshift_*')
# no string ops
backend_test.exclude('string')
backend_test.exclude('test_strnorm_*')
backend_test.exclude('test_regex_*')
# no rnn
backend_test.exclude('test_gru_*')
backend_test.exclude('test_rnn_*')
backend_test.exclude('test_lstm_*')
backend_test.exclude('test_simple_rnn_*')
# no control flow
# control flow uses AttributeProto.GRAPH
backend_test.exclude('test_loop*')
backend_test.exclude('test_range_float_type_positive_delta_expanded_cpu') # requires loop
backend_test.exclude('test_affine_grid_2d_align_corners_expanded_cpu')
backend_test.exclude('test_affine_grid_2d_expanded_cpu')
backend_test.exclude('test_affine_grid_3d_align_corners_expanded_cpu')
backend_test.exclude('test_affine_grid_3d_expanded_cpu')
backend_test.exclude('test_range_int32_type_negative_delta_expanded_cpu')
# unsupported (strange) ops
backend_test.exclude('test_bernoulli_*')
backend_test.exclude('test_det_*')
backend_test.exclude('test_col2im_*')
backend_test.exclude('test_gridsample_*')
backend_test.exclude('test_dft_*')
backend_test.exclude('test_unique_*')
backend_test.exclude('test_sequence_*')
backend_test.exclude('test_nonmaxsuppression_*')
backend_test.exclude('test_reversesequence_*')
backend_test.exclude('test_roialign_*')
backend_test.exclude('test_tfidfvectorizer_*')
backend_test.exclude('test_stft_*')
backend_test.exclude('test_melweightmatrix_*')
# more strange ops
backend_test.exclude('test_basic_deform_conv_*')
backend_test.exclude('test_deform_conv_*')
backend_test.exclude('test_lppool_*')
backend_test.exclude('test_scan_*')
backend_test.exclude('test_split_to_sequence_*')
backend_test.exclude('test_ai_onnx_ml_tree_ensemble_*') # https://github.com/onnx/onnx/blob/main/onnx/reference/ops/aionnxml/op_tree_ensemble.py#L121
backend_test.exclude('test_attention_4d_diff_heads_mask4d_padded_kv_cpu') # needs nonpad_kv_seqlen handling
backend_test.exclude('test_attention_4d_fp16_cpu') # fp16 numerical issues
backend_test.exclude('test_attention_4d_fp16_expanded_cpu') # fp16 numerical issues
backend_test.exclude('test_attention_4d_gqa_with_past_and_present_fp16_cpu') # fp16 numerical issues
backend_test.exclude('test_attention_4d_gqa_with_past_and_present_fp16_expanded_cpu') # fp16 numerical issues
# rest of the failing tests
backend_test.exclude('test_resize_tf_crop_and_resize_cpu') # tf_crop_and_resize not implemented
backend_test.exclude('test_resize_tf_crop_and_resize_axes_2_3_cpu') # tf_crop_and_resize not implemented
backend_test.exclude('test_resize_tf_crop_and_resize_axes_3_2_cpu') # tf_crop_and_resize not implemented
backend_test.exclude('test_resize_tf_crop_and_resize_extrapolation_value_cpu') # tf_crop_and_resize value not implemented
backend_test.exclude('test_resize_downsample_scales_linear_antialias_cpu') # antialias not implemented
backend_test.exclude('test_resize_downsample_sizes_linear_antialias_cpu') # antialias not implemented
backend_test.exclude('test_resize_downsample_scales_cubic_antialias_cpu') # antialias not implemented
backend_test.exclude('test_resize_downsample_sizes_cubic_antialias_cpu') # antialias not implemented
backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_value_only_mapping_cpu') # bad data type string
backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_mapping_cpu') # bad data type string
backend_test.exclude('test_if_opt_cpu') # ValueError: 13 is not a valid AttributeType
backend_test.exclude('test_if_seq_cpu') # NotImplementedError: op='SequenceConstruct' is not supported
if Device.DEFAULT == "METAL" or (OSX and Device.DEFAULT == "CL"):
# numerical inaccuracy
backend_test.exclude('test_mish_cpu')
backend_test.exclude('test_mish_expanded_cpu')
# disable model tests for now since they are slow
if not getenv("MODELTESTS"):
for x in backend_test.test_suite:
if 'OnnxBackendRealModelTest' in str(type(x)):
backend_test.exclude(str(x).split(" ")[0])
else:
# model tests all pass!
backend_test.include('test_resnet50')
backend_test.include('test_inception_v1')
backend_test.include('test_inception_v2')
backend_test.include('test_densenet121')
backend_test.include('test_shufflenet')
backend_test.include('test_squeezenet')
backend_test.include('test_bvlc_alexnet')
backend_test.include('test_zfnet512')
backend_test.include('test_vgg19')
globals().update(backend_test.enable_report().test_cases)
if __name__ == '__main__':
unittest.main()

622
artifacts/package_sources/tinygrad/test/external/external_test_onnx_ops.py vendored Normal file
View File

@@ -0,0 +1,622 @@
# inputs, attributes, and outputs for tests are found here:
# https://github.com/onnx/onnx/blob/main/docs/Operators.md
# https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md
from typing import Any
import unittest, onnx, tempfile
from tinygrad import dtypes, Tensor, Context
from tinygrad.nn.onnx import OnnxRunner
import numpy as np
from extra.onnx_helpers import validate
from onnx.defs import ONNX_DOMAIN, AI_ONNX_PREVIEW_TRAINING_DOMAIN
MICROSOFT_CONTRIB_OPS_DOMAIN = "com.microsoft"
# TODO: remove this once ORT supports 1.18.0
from onnx.helper import VERSION_TABLE
VERSION_MAP = {row[0]: row[1:] for row in VERSION_TABLE}
IR_VERSION, ai_onnx, ai_onnx_ml, ai_onnx_training = VERSION_MAP["1.17.0"]
class TestOnnxOps(unittest.TestCase):
DOMAIN = None
def helper_build_model(self, op:str, inps:dict[str, np.ndarray], opts:dict[str, Any], outs:list[str]):
onnx_inputs = [onnx.helper.make_tensor_value_info(name, onnx.helper.np_dtype_to_tensor_dtype(arr.dtype), arr.shape) for name, arr in inps.items()]
onnx_outputs = [onnx.helper.make_empty_tensor_value_info(name) for name in outs]
nodes = [onnx.helper.make_node(op, list(inps), list(outs), domain=self.DOMAIN, **opts)]
graph = onnx.helper.make_graph(nodes, f"test_{op.lower()}", onnx_inputs, onnx_outputs)
#model = onnx.helper.make_model(graph, producer_name=f"test_{op.lower()}")
# TODO: remove this once ORT supports 1.18.0
opset_id = None
if type(self).__name__ == "TestMainOnnxOps": opset_id = ai_onnx
if type(self).__name__ == "TestTrainingOnnxOps": opset_id = ai_onnx_training
if type(self).__name__ == "TestContribOnnxOps": opset_id = 1
model = onnx.helper.make_model(graph, producer_name=f"test_{op.lower()}", ir_version=IR_VERSION,
opset_imports=[onnx.helper.make_opsetid(self.DOMAIN, opset_id)])
return model
def helper_test_single_op(self, op:str, inps:dict[str, np.ndarray], opts:dict[str, Any], outs:list[str], rtol=1e-3, atol=1e-6):
model = self.helper_build_model(op, inps, opts, outs)
with tempfile.NamedTemporaryFile() as tmp:
onnx.save(model, tmp.name)
validate(tmp.name, inps, rtol, atol)
class TestMainOnnxOps(TestOnnxOps):
DOMAIN = ONNX_DOMAIN
def test_reshape(self):
inputs = {"in": np.arange(6, dtype=np.float32), "shape": np.array([2,3], dtype=np.int64)}
attributes = {}
outputs = ["out"]
self.helper_test_single_op("Reshape", inputs, attributes, outputs)
def test_squeeze(self):
# axes is None
inputs = {"data": np.random.randn(1, 3, 1, 1).astype(np.float32)}
attributes = {}
outputs = ["squeezed"]
self.helper_test_single_op("Squeeze", inputs, attributes, outputs)
def test_conv(self):
# test VALID auto_pad
inputs = {
"x": np.random.randn(1, 3, 384, 384).astype(np.float32),
"w": np.random.randn(1152, 3, 14, 14).astype(np.float32),
"b": np.random.randn(1152).astype(np.float32)
}
attributes = {'auto_pad': 'VALID', 'dilations': (1, 1), 'group': 1, 'kernel_shape': (14, 14), 'strides': (14, 14)}
outputs = ["y"]
self.helper_test_single_op("Conv", inputs, attributes, outputs, atol=1e-4)
def test_pad_constant_value_zero(self):
from tinygrad.nn.onnx import onnx_ops
Pad = onnx_ops["Pad"]
x = Tensor.arange(4).reshape(1, 1, 2, 2).float()
pads = [0, 0, 1, 1, 0, 0, 1, 1]
out = Pad(x, pads, constant_value=0, value=3)
expected = x.pad((pads[3], pads[7], pads[2], pads[6], pads[1], pads[5], pads[0], pads[4]), value=0)
self.assertEqual(out.tolist(), expected.tolist())
def test_gather(self):
# test const negative indices
inputs = {
"input": np.random.randn(1, 3, 3).astype(np.float32),
"indices": np.array(-2, dtype=np.int64),
}
attributes = {'axis': 1}
outputs = ["y"]
self.helper_test_single_op("Gather", inputs, attributes, outputs)
def test_gather_jit_different_indices(self):
# Gather should not assume indices is const when it can change at runtime
from tinygrad import TinyJit
from tinygrad.nn.onnx import onnx_ops
Gather = onnx_ops["Gather"]
x = Tensor([10, 20, 30, 40, 50])
indices_list = [[0, 1], [2, 3], [4, 0]]
expected = [[10, 20], [30, 40], [50, 10]]
# without JIT: correct
self.assertEqual([Gather(x, Tensor(idx)).tolist() for idx in indices_list], expected)
@TinyJit
def gather_jit(x, indices): return Gather(x, indices)
self.assertEqual([gather_jit(x, Tensor(idx)).tolist() for idx in indices_list], expected)
def test_gather_jit_const_zero_index(self):
# Gather with const index=0 (falsy in Python) should work with JIT cache
from tinygrad import TinyJit
# Create model: y = Gather(x, 0) + x where 0 is from initializer
# The Add ensures there's a kernel to JIT
x_input = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, (5,))
y_output = onnx.helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, (5,))
idx_init = onnx.numpy_helper.from_array(np.array(0, dtype=np.int64), name="idx")
graph = onnx.helper.make_graph([
onnx.helper.make_node("Gather", ["x", "idx"], ["g"], axis=0),
onnx.helper.make_node("Add", ["g", "x"], ["y"])],
"test_gather_zero", [x_input], [y_output], [idx_init])
model = onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 13)])
with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as tmp:
onnx.save(model, tmp.name)
runner = OnnxRunner(tmp.name)
@TinyJit
def run_gather(x): return runner({"x": x})["y"]
# Run multiple times - JIT capture should use cached index=0 correctly
for val in [[10, 20, 30, 40, 50], [100, 200, 300, 400, 500], [1, 2, 3, 4, 5]]:
result = run_gather(Tensor(val, dtype=dtypes.float32))
np.testing.assert_equal(result.numpy(), np.array(val) + val[0])
# NOTE: resize OP is sensitive to numerical errors
def _test_resize_scales(self, scale_values, **kwargs):
for sc in scale_values:
for ct_mode in ["half_pixel", "align_corners", "asymmetric", "pytorch_half_pixel", "half_pixel_symmetric"]:
with self.subTest(coordinate_transformation_mode=ct_mode, scale=sc, **kwargs):
X = np.array([[[[1, 2, 3, 4],
[5, 6, 7, 8],
[9,10,11,12]]]], dtype=np.float32)
scales = np.array([1.0, 1.0, sc, sc], dtype=np.float32)
inputs = {"X": X, "roi": np.array([], dtype=np.float32), "scales": scales}
attributes = {"coordinate_transformation_mode": ct_mode, **kwargs}
outputs = ["out"]
self.helper_test_single_op("Resize", inputs, attributes, outputs)
def test_resize_linear_mode(self):
self._test_resize_scales([0.01, 0.25, 0.5, 0.51, 0.6, 1.0, 1.5, 2.0, 3.5, 20.0], mode="linear")
def test_resize_nearest_mode(self):
# excluded 3.5 because some values divide into slight numerical differences, which when rounded gives wrong results
self._test_resize_scales([0.01, 0.25, 0.5, 0.51, 0.6, 1.0, 1.5, 2.0, 20.0], mode="nearest")
def test_resize_cubic_mode(self):
self._test_resize_scales([0.01, 0.25, 0.5, 0.51, 0.6, 1.0, 1.5, 2.0, 3.5, 20.0], mode="cubic", exclude_outside=1)
self._test_resize_scales([0.01, 0.25, 0.5, 0.51, 0.6, 1.0, 1.5, 2.0, 3.5, 20.0], mode="cubic", exclude_outside=0)
def _test_if(self, then_value, else_value):
then_out = onnx.helper.make_tensor_value_info("res", onnx.TensorProto.FLOAT, then_value.shape)
else_out = onnx.helper.make_tensor_value_info("res", onnx.TensorProto.FLOAT, else_value.shape)
then_const_node = onnx.helper.make_node("Constant", inputs=[], outputs=["res"], value=onnx.numpy_helper.from_array(then_value))
else_const_node = onnx.helper.make_node("Constant", inputs=[], outputs=["res"], value=onnx.numpy_helper.from_array(else_value))
then_body = onnx.helper.make_graph([then_const_node], "then_body", [], [then_out])
else_body = onnx.helper.make_graph([else_const_node], "else_body", [], [else_out])
self.helper_test_single_op("If", {"cond": np.array(False).astype(bool)}, {"then_branch": then_body, "else_branch": else_body}, ["res"])
self.helper_test_single_op("If", {"cond": np.array(True).astype(bool)}, {"then_branch": then_body, "else_branch": else_body}, ["res"])
def test_if_different_shapes_broadcastable(self):
self._test_if(np.array([[1], [2]]).astype(np.float32), np.array([[6, 5, 4, 3, 2, 1]]).astype(np.float32))
def test_if_different_shapes_not_broadcastable(self):
self._test_if(np.array([[1, 2, 3], [4, 5, 6]]).astype(np.float32), np.array([[6, 5, 4, 3, 2, 1]]).astype(np.float32))
def test_if_jit_different_shapes(self):
# When shapes differ, Python selection evaluates condition at graph build time, breaking JIT
from tinygrad import TinyJit
from tinygrad.engine.jit import JitError
# then: x+1 shape (3,), else: x[:2]+1 shape (2,)
x_input = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, (3,))
then_out = onnx.helper.make_tensor_value_info("res", onnx.TensorProto.FLOAT, (3,))
then_body = onnx.helper.make_graph([
onnx.helper.make_node("Constant", [], ["one"], value=onnx.numpy_helper.from_array(np.array(1, dtype=np.float32))),
onnx.helper.make_node("Add", ["x", "one"], ["res"])], "then_body", [x_input], [then_out])
else_out = onnx.helper.make_tensor_value_info("res", onnx.TensorProto.FLOAT, (2,))
else_body = onnx.helper.make_graph([
onnx.helper.make_node("Constant", [], ["starts"], value=onnx.numpy_helper.from_array(np.array([0], dtype=np.int32))),
onnx.helper.make_node("Constant", [], ["ends"], value=onnx.numpy_helper.from_array(np.array([2], dtype=np.int32))),
onnx.helper.make_node("Constant", [], ["one"], value=onnx.numpy_helper.from_array(np.array(1, dtype=np.float32))),
onnx.helper.make_node("Slice", ["x", "starts", "ends"], ["x2"]),
onnx.helper.make_node("Add", ["x2", "one"], ["res"])], "else_body", [x_input], [else_out])
cond_input = onnx.helper.make_tensor_value_info("cond", onnx.TensorProto.BOOL, (1,))
main_x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, (3,))
graph = onnx.helper.make_graph([onnx.helper.make_node("If", ["cond"], ["res"], then_branch=then_body, else_branch=else_body)],
"test", [cond_input, main_x], [onnx.helper.make_empty_tensor_value_info("res")])
model = onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 22)])
with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as tmp:
onnx.save(model, tmp.name)
runner = OnnxRunner(tmp.name)
@TinyJit
def run_if(cond, x): return runner({"cond": cond, "x": x})["res"]
x = Tensor([1.0, 2.0, 3.0])
with self.assertRaises(JitError):
for _ in range(3):
run_if(Tensor([True]), x)
def test_resize_downsample_scales_linear_align_corners(self):
# https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-131
X = np.array([[[[1, 2, 3, 4], [5, 6, 7, 8]]]], dtype=np.float32)
scales = np.array([1.0, 1.0, 0.6, 0.6], dtype=np.float32)
inputs = {"X": X, "roi": np.array([], dtype=np.float32), "scales": scales}
attributes = {"mode": "linear", "coordinate_transformation_mode": "align_corners"}
outputs = ["out"]
self.helper_test_single_op("Resize", inputs, attributes, outputs)
def test_resize_downsample_scales_cubic_align_corners(self):
# https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-131
X = np.array([[[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]]], dtype=np.float32)
scales = np.array([1.0, 1.0, 0.8, 0.8], dtype=np.float32)
inputs = {"X": X, "roi": np.array([], dtype=np.float32), "scales": scales}
attributes = {"mode": "cubic", "coordinate_transformation_mode": "align_corners"}
outputs = ["out"]
self.helper_test_single_op("Resize", inputs, attributes, outputs)
def test_maxunpool_export_with_output_shape(self):
# https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-91
xT = np.array([[[[5, 6], [7, 8]]]], dtype=np.float32)
xI = np.array([[[[5, 7], [13, 15]]]], dtype=np.int64)
output_shape = np.array((1, 1, 5, 5), dtype=np.int64)
inputs = {"x": xT, "indices": xI, "output_shape": output_shape}
attributes = {"kernel_shape": [2, 2], "strides": [2, 2]}
outputs = ["y"]
self.helper_test_single_op("MaxUnpool", inputs, attributes, outputs)
def test_averagepool_3d_dilations_large_count_include_pad_is_1_ceil_mode_is_True(self):
# https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-13
inputs = {"x": np.random.randn(1, 1, 32, 32, 32).astype(np.float32)}
attributes = {"kernel_shape": (5, 5, 5), "strides": (3, 3, 3), "dilations": (2, 2, 2), "count_include_pad": 1, "ceil_mode": True}
outputs = ["y"]
self.helper_test_single_op("AveragePool", inputs, attributes, outputs)
def test_isinf(self):
# https://github.com/onnx/onnx/blob/main/docs/Operators.md#isinf
# attributes are int but output expects bool
x = np.array([-1.2, np.nan, np.inf, 2.8, -np.inf, np.inf], dtype=np.float32)
inputs = {"x": x}
attributes = {"detect_negative":1, "detect_positive":1}
outputs = ["y"]
model = self.helper_build_model("IsInf", inputs, attributes, outputs)
runner = OnnxRunner(Tensor(model.SerializeToString(), device="PYTHON"))
outputs = runner(inputs)
assert outputs["y"].dtype is dtypes.bool
def test_quantize_linear(self):
test_cases = [
{"test_case": "round_half_to_even", "qdtype": np.int8, "qzero_point": 0, "x": [-1.5, -0.5, 0.5, 1.5], "scale": 1.0},
{"test_case": "round_to_even_before_add_zero_point", "qdtype": np.uint8, "qzero_point": 1, "x": [0.5, 1.5], "scale": 1.0},
]
for case in test_cases:
with self.subTest(test_case=case["test_case"]):
inputs = {
"x": np.array([case["x"]], dtype=np.float32),
"y_scale": np.array(case["scale"], dtype=np.float32),
"y_zero_point": np.array(case["qzero_point"], dtype=case["qdtype"])
}
self.helper_test_single_op("QuantizeLinear", inputs, {}, ["y"])
def test_dynamic_quantize_linear(self):
test_cases = [
{"name": "round_half_to_even", "x": np.array([0, 0.5, 1.5, 255], dtype=np.float32)},
{"name": "round_zero_point_half_down_to_even", "x": np.array([-1, 509], dtype=np.float32)},
{"name": "round_zero_point_half_up_to_even", "x": np.array([-11, 499], dtype=np.float32)},
# other tests from https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-45
{"name": "max_adjusted", "x": np.array([-1.0, -2.1, -1.3, -2.5, -3.34, -4.0], dtype=np.float32)},
{"name": "min_adjusted", "x": np.array([1, 2.1, 1.3, 2.5, 3.34, 4.0, 1.5, 2.6, 3.9, 4.0, 3.0, 2.345], dtype=np.float32).reshape((3, 4))},
]
for case in test_cases:
with self.subTest(test_case=case["name"]):
self.helper_test_single_op("DynamicQuantizeLinear", {"x": case["x"]}, {}, ["y", "y_scale", "y_zero_point"])
def test_qlinear_conv(self):
for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]:
for b in (np.ones([32], dtype=np.int32), np.zeros([32], dtype=np.int32)):
for channel_shape in [(), (32,)]:
with self.subTest(dtype=dtype.__name__, zero_point=zero_point, channel_shape=channel_shape):
dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max
inputs = {
"x": np.random.randint(dtype_min, dtype_max + 1, [1, 3, 224, 224], dtype=dtype),
"x_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
"x_zero_point": np.array(zero_point, dtype=dtype),
"w": np.random.randint(dtype_min, dtype_max + 1, [32, 3, 3, 3], dtype=dtype),
"w_scale": np.random.uniform(0.01, 0.1, channel_shape).astype(np.float32),
"w_zero_point": np.full(channel_shape, zero_point, dtype=dtype),
"y_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
"y_zero_point": np.array(zero_point, dtype=dtype),
"b": b
}
attributes = {'auto_pad': 'NOTSET', 'dilations': (1, 1), 'group': 1, 'kernel_shape': (3, 3), 'pads': (1, 1, 1, 1), 'strides': (2, 2)}
outputs = ["out"]
self.helper_test_single_op("QLinearConv", inputs, attributes, outputs, atol=1) # occasionally inaccurate
def test_qlinear_matmul(self):
for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]:
with self.subTest(dtype=dtype.__name__, zero_point=zero_point):
dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max
inputs = {
"A": np.random.randint(dtype_min, dtype_max + 1, [10, 10], dtype=dtype),
"A_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
"A_zero_point": np.array(zero_point, dtype=dtype),
"B": np.random.randint(dtype_min, dtype_max + 1, [10, 10], dtype=dtype),
"B_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
"B_zero_point": np.array(zero_point, dtype=dtype),
"Y_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
"Y_zero_point": np.array(zero_point, dtype=dtype)
}
attributes = {}
outputs = ["Y"]
self.helper_test_single_op("QLinearMatMul", inputs, attributes, outputs)
for name,val in (("round_half_down_to_even", 1), ("round_half_up_to_even", 3)):
with self.subTest(test_case=name, val=val):
inputs = {
"A": np.array([val], dtype=np.int8),
"A_scale": np.array(0.5, dtype=np.float32),
"A_zero_point": np.array(0, dtype=np.int8),
"B": np.array([1], dtype=np.int8),
"B_scale": np.array(1, dtype=np.float32),
"B_zero_point": np.array(0, dtype=np.int8),
"Y_scale": np.array(1, dtype=np.float32),
"Y_zero_point": np.array(0, dtype=np.int8)
}
attributes = {}
outputs = ["Y"]
self.helper_test_single_op("QLinearMatMul", inputs, attributes, outputs)
def _run_qlinearmatmul_test(self, quant_type, dtype, dims):
# https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-111
if dims == 2:
a = np.array([[208, 236, 0, 238], [3, 214, 255, 29]])
b = np.array([[152, 51, 244], [60, 26, 255], [0, 127, 246], [127, 254, 247]])
else:
a = np.array([[[208, 236, 0, 238], [3, 214, 255, 29]], [[208, 236, 0, 238], [3, 214, 255, 29]]])
b = np.array([[[152, 51, 244], [60, 26, 255], [0, 127, 246], [127, 254, 247]], [[152, 51, 244], [60, 26, 255], [0, 127, 246], [127, 254, 247]]])
a_zero_point = np.array([113])
b_zero_point = np.array([114])
y_zero_point = np.array([118])
if quant_type == np.int8:
a, b, a_zero_point, b_zero_point, y_zero_point = (x - 127 for x in (a, b, a_zero_point, b_zero_point, y_zero_point))
a, b, a_zero_point, b_zero_point, y_zero_point = (x.astype(quant_type) for x in (a, b, a_zero_point, b_zero_point, y_zero_point))
inputs = {
"a": a, "a_scale": np.array([0.0066], dtype=dtype), "a_zero_point": a_zero_point,
"b": b, "b_scale": np.array([0.00705], dtype=dtype), "b_zero_point": b_zero_point,
"y_scale": np.array([0.0107], dtype=dtype), "y_zero_point": y_zero_point
}
self.helper_test_single_op("QLinearMatMul", inputs, {}, ["y"],)
def test_qlinearmatmul_2D_int8_float16(self): self._run_qlinearmatmul_test(np.int8, np.float16, 2)
def test_qlinearmatmul_3D_int8_float16(self): self._run_qlinearmatmul_test(np.int8, np.float16, 3)
def test_qlinearmatmul_2D_int8_float32(self): self._run_qlinearmatmul_test(np.int8, np.float32, 2)
def test_qlinearmatmul_3D_int8_float32(self): self._run_qlinearmatmul_test(np.int8, np.float32, 3)
def test_reduce_l2_half(self):
inputs = {"data": np.random.randn(1, 1, 32, 32, 32).astype(np.half)*100}
self.helper_test_single_op("ReduceL2", inputs, {}, ["reduced"])
def test_same_device_as_input(self):
from tinygrad.nn.onnx import onnx_ops
EyeLike = onnx_ops["EyeLike"]
Shape = onnx_ops["Shape"]
Compress = onnx_ops["Compress"]
with Context(DEV="CPU"):
x = Tensor.arange(4).clone("PYTHON").reshape(2,2)
self.assertIsNone(EyeLike(x).device)
self.assertEqual(Shape(x).device, x.device)
out = Compress(x, [True, False, True, False])
self.assertEqual(out.device, x.device)
self.assertEqual(out.tolist(), [0, 2])
class TestTrainingOnnxOps(TestOnnxOps):
# NOTE: ORT doesn't actually support training ops on cpu so we test using functions provided by onnx
DOMAIN = AI_ONNX_PREVIEW_TRAINING_DOMAIN
def _validate_training(self, op:str, onnx_fxn, inps:dict[str, np.ndarray], opts:dict[str, Any], outs:list[str]):
model = self.helper_build_model(op, inps, opts, outs)
if op == "Momentum": del opts['mode']
runner = OnnxRunner(Tensor(model.SerializeToString(), device="PYTHON"))
tiny_out = runner(inps)
onnx_out = onnx_fxn(**inps, **opts)
for (nm, t_out), o_out in zip(tiny_out.items(), onnx_out):
np.testing.assert_allclose(t_out.numpy(), o_out, rtol=1e-6, atol=1e-6, err_msg=f"{nm} failed")
def test_adagrad_t(self):
from onnx.backend.test.case.node.adagrad import apply_adagrad
for t in [0, 1, 3, 100]:
inputs = {
"r": np.array(0.01, dtype=np.float32),
"t": np.array(t, dtype=np.int32),
"x": np.random.randn(3, 3).astype(np.float32),
"g": np.random.randn(3, 3).astype(np.float32),
"h": np.random.randn(3, 3).astype(np.float32),
}
attributes = {"decay_factor": 0.1, "epsilon": 1e-6, "norm_coefficient": 0.01}
outputs = ["X_out", "H_out"]
self._validate_training("Adagrad", apply_adagrad, inputs, attributes, outputs)
def test_momentum(self):
from onnx.backend.test.case.node.momentum import apply_momentum, apply_nesterov
for onnx_fxn, mode in ((apply_momentum, "standard"), (apply_nesterov, "nesterov")):
for t in [0, 1, 3, 100]:
inputs = {
"r": np.array(0.01, dtype=np.float32),
"t": np.array(t, dtype=np.int32),
"x": np.random.randn(3, 3).astype(np.float32),
"g": np.random.randn(3, 3).astype(np.float32),
"v": np.random.randn(3, 3).astype(np.float32),
}
attributes = {"alpha": 0.9, "beta": 0.1, "mode": mode, "norm_coefficient": 0.01}
outputs = ["X_out", "V_out"]
self._validate_training("Momentum", onnx_fxn, inputs, attributes, outputs)
def test_adam(self):
from onnx.backend.test.case.node.adam import apply_adam
for t in [0, 1, 3, 100]:
inputs = {
"r": np.array(0.01, dtype=np.float32),
"t": np.array(t, dtype=np.int32),
"x": np.random.randn(3, 3).astype(np.float32),
"g": np.random.randn(3, 3).astype(np.float32),
"v": np.random.randn(3, 3).astype(np.float32),
"h": np.random.randn(3, 3).astype(np.float32),
}
attributes = { "alpha": 0.9, "beta": 0.999, "epsilon": 1e-8, "norm_coefficient": 0.01, "norm_coefficient_post": 0.02 }
outputs = ["X_new", "V_new", "H_new"]
self._validate_training("Adam", apply_adam, inputs, attributes, outputs)
class TestContribOnnxOps(TestOnnxOps):
DOMAIN = MICROSOFT_CONTRIB_OPS_DOMAIN
def test_attention(self):
batch_size, seq_len, input_hidden_size = 2, 8, 256
num_heads, head_size = 4, 64
hidden_size = num_heads * head_size
v_hidden_size = hidden_size
# for mask_index
right_padding_mask = np.random.randint(1, seq_len + 1, size=(batch_size,), dtype=np.int32)
end_positions = np.random.randint(1, seq_len + 1, size=(batch_size,), dtype=np.int32)
start_positions = np.array([np.random.randint(0, end) for end in end_positions], dtype=np.int32)
left_padding_mask = np.concatenate([end_positions, start_positions])
base_inps = {
"input": np.random.randn(batch_size, seq_len, input_hidden_size).astype(np.float32),
"weights": np.random.randn(input_hidden_size, hidden_size * 3).astype(np.float32),
# bias is required in ORT (segfaults otherwise), eventhough docs says it's optional
"bias": np.random.randn(hidden_size * 2 + v_hidden_size).astype(np.float32),
}
base_opts = {"num_heads": num_heads}
test_cases = [
({}, {}),
({}, {"scale": 0.1}),
({}, {"scale": 1.0}),
({}, {"unidirectional": 1}),
({"mask_index": right_padding_mask}, {}),
({"mask_index": left_padding_mask}, {}),
({"mask_index": np.random.randint(0, seq_len, size=(batch_size, seq_len), dtype=np.int32)}, {"mask_filter_value": -5000.0}),
({"mask_index": np.random.randint(0, seq_len, size=(batch_size, seq_len, seq_len), dtype=np.int32)}, {"mask_filter_value": -np.inf}),
# BUG: when `mask_index` is used with `unidirectional`, the first value must be True
# otherwise this will trigger a different ORT behavior where start consecutive Falses will be turned True
# e.g. mask_index = [[0, 0, 1, 0, 1, 1, 1, 1], [0, 0, 1, 0, 1, 1, 1, 1]]
# will need mask[:, :, 0:1, 0:1] = True
({"mask_index": np.array([[1, 0, 1, 0, 1, 1, 1, 1], [1, 0, 1, 0, 1, 1, 1, 1]], dtype=np.int32)}, {"unidirectional": 1}),
({ "weights": np.random.randn(input_hidden_size, hidden_size + hidden_size + 128).astype(np.float32),
"bias": np.random.randn(hidden_size + hidden_size + 128).astype(np.float32)},
{"qkv_hidden_sizes": [hidden_size, hidden_size, 128]}),
# TODO: past is not tested. ORT gives type error for input
]
for i, (extra_inps, extra_opts) in enumerate(test_cases):
with self.subTest(f"test_attention_{i}"):
inps = {**base_inps, **extra_inps}
opts = {**base_opts, **extra_opts}
outputs = ["output", "present"] if "past" in inps else ["output"]
self.helper_test_single_op("Attention", inps, opts, outputs, atol=1e-4)
def test_skip_layer_normalization(self):
shape = (2, 8, 32)
for has_beta in [True, False]:
for has_bias in [True, False]:
with self.subTest(has_beta=has_beta, has_bias=has_bias):
hidden_size = shape[-1]
inputs = {
"input": np.random.randn(*shape).astype(np.float32),
"skip": np.random.randn(*shape).astype(np.float32),
"gamma": np.random.randn(hidden_size).astype(np.float32),
}
if has_beta: inputs["beta"] = np.random.randn(hidden_size).astype(np.float32)
if has_bias: inputs["bias"] = np.random.randn(hidden_size).astype(np.float32)
attributes = {"epsilon": 1e-12}
outputs = ["output", "mean", "inv_std_var", "input_skip_bias_sum"]
self.helper_test_single_op("SkipLayerNormalization", inputs, attributes, outputs)
def test_bias_gelu(self):
shape = (2,3,4)
inputs = {
"A": np.random.randn(*shape).astype(np.float32),
"B": np.random.randn(shape[-1]).astype(np.float32)
}
attributes = {}
outputs = ["C"]
self.helper_test_single_op("BiasGelu", inputs, attributes, outputs)
def test_qlinear_add(self):
for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]:
with self.subTest(dtype=dtype.__name__, zero_point=zero_point):
dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max
inputs = {
"A": np.random.randint(dtype_min, dtype_max + 1, [10, 10], dtype=dtype),
"A_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
"A_zero_point": np.array(zero_point, dtype=dtype),
"B": np.random.randint(dtype_min, dtype_max + 1, [10, 10], dtype=dtype),
"B_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
"B_zero_point": np.array(zero_point, dtype=dtype),
"C_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
"C_zero_point": np.array(zero_point, dtype=dtype)
}
attributes = {}
outputs = ["C"]
self.helper_test_single_op("QLinearAdd", inputs, attributes, outputs, atol=1) # TODO: look into why this is inaccurate
def test_qlinear_add_round_half_to_even(self):
with self.subTest(test_case="round_half_to_even"):
inputs = {
"A": np.array([1, 1, 1, 1], dtype=np.int8),
"A_scale": np.array(1, dtype=np.float32),
"A_zero_point": np.array(0, dtype=np.int8),
"B": np.array([1, 5, -3, -7], dtype=np.int8),
"B_scale": np.array(1, dtype=np.float32),
"B_zero_point": np.array(0, dtype=np.int8),
"C_scale": np.array(4, dtype=np.float32),
"C_zero_point": np.array(0, dtype=np.int8)
}
attributes = {}
outputs = ["C"]
self.helper_test_single_op("QLinearAdd", inputs, attributes, outputs, atol=1) # TODO: look into why this is inaccurate
def test_qlinear_mul(self):
for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]:
with self.subTest(dtype=dtype.__name__, zero_point=zero_point):
dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max
inputs = {
"A": np.random.randint(dtype_min, dtype_max + 1, [10, 10], dtype=dtype),
"A_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
"A_zero_point": np.array(zero_point, dtype=dtype),
"B": np.random.randint(dtype_min, dtype_max + 1, [10, 10], dtype=dtype),
"B_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
"B_zero_point": np.array(zero_point, dtype=dtype),
"C_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
"C_zero_point": np.array(zero_point, dtype=dtype)
}
attributes = {}
outputs = ["C"]
self.helper_test_single_op("QLinearMul", inputs, attributes, outputs)
with self.subTest(test_case="round_half_to_even"):
inputs = {
"A": np.array([1, 1, 1, 1], dtype=np.int8),
"A_scale": np.array(1, dtype=np.float32),
"A_zero_point": np.array(0, dtype=np.int8),
"B": np.array([2, 6, -2, -6], dtype=np.int8),
"B_scale": np.array(1, dtype=np.float32),
"B_zero_point": np.array(0, dtype=np.int8),
"C_scale": np.array(4, dtype=np.float32),
"C_zero_point": np.array(0, dtype=np.int8)
}
attributes = {}
outputs = ["C"]
self.helper_test_single_op("QLinearMul", inputs, attributes, outputs)
def test_qlinear_global_average_pool(self):
for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]:
for channels_last in [0, 1]:
with self.subTest(dtype=dtype.__name__, zero_point=zero_point, channels_last=channels_last):
dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max
# NCHW for channels_last=0, NHWC for channels_last=1
shape = [1, 3, 32, 32] if channels_last == 0 else [1, 32, 32, 3]
inputs = {
"X": np.random.randint(dtype_min, dtype_max + 1, shape, dtype=dtype),
"x_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
"x_zero_point": np.array(zero_point, dtype=dtype),
"y_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
"y_zero_point": np.array(zero_point, dtype=dtype)
}
attributes = {"channels_last": channels_last}
outputs = ["C"]
self.helper_test_single_op("QLinearGlobalAveragePool", inputs, attributes, outputs)
def test_same_device_as_input(self):
from tinygrad.nn.onnx import onnx_ops, OpSetId, Domain
EmbedLayerNormalization = onnx_ops["EmbedLayerNormalization"]
Attention = onnx_ops["Attention"]
with Context(DEV="CPU"):
input_ids = Tensor([[1, 2]], device="PYTHON", dtype=dtypes.int32)
segment_ids = Tensor([[0, 0]], device="PYTHON", dtype=dtypes.int32)
word = Tensor.ones(4, 3, device="PYTHON")
pos = Tensor.ones(5, 3, device="PYTHON")
seg = Tensor.ones(1, 3, device="PYTHON")
gamma, beta = Tensor.ones(3, device="PYTHON"), Tensor.zeros(3, device="PYTHON")
out, _, _ = EmbedLayerNormalization(input_ids, segment_ids, word, pos, seg, gamma, beta)
self.assertEqual(out.device, input_ids.device)
out.realize()
attn = Attention[OpSetId(Domain.MICROSOFT_CONTRIB_OPS, 1)]
x = Tensor.ones(1, 2, 4, device="PYTHON")
w = Tensor.ones(4, 12, device="PYTHON")
mask = Tensor([2, 0], device="PYTHON", dtype=dtypes.int32)
out, _ = attn(x, w, mask_index=mask, num_heads=1, unidirectional=1)
self.assertEqual(out.device, x.device)
out.realize()
if __name__ == "__main__":
unittest.main()

165
artifacts/package_sources/tinygrad/test/external/external_test_onnx_runner.py vendored Normal file
View File

@@ -0,0 +1,165 @@
import unittest, onnx, tempfile, pathlib
import numpy as np
from tinygrad import Tensor
from tinygrad.uop.ops import Ops
from typing import Any
from tinygrad.nn.onnx import OnnxRunner, OnnxPBParser, OnnxDataType
from hypothesis import given, strategies as st
# copied from test_const_folding.py
def _check_ast_count(desired_count:int, t:Tensor):
# NOTE: this has side effect because everything can be scheduled only once
linear = t.schedule_linear()
asts = [call for call in linear.src if call.src[0].op is Ops.SINK]
assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
def build_onnx(nodes, from_disk:bool=True, **kwargs):
"""Helper to build and return an OnnxRunner from ONNX nodes."""
graph = onnx.helper.make_graph(nodes, 'test', kwargs.get('inputs', []), kwargs.get('outputs', []), kwargs.get('initializers', []))
model = onnx.helper.make_model(graph)
if from_disk:
with tempfile.TemporaryDirectory() as tmpdir:
tmp_path = pathlib.Path(tmpdir)
model_path = tmp_path / "model.onnx"
onnx.save(model, model_path)
runner = OnnxRunner(model_path)
else:
# use the in-memory method
runner = OnnxRunner(Tensor(model.SerializeToString(), device="PYTHON"))
return runner
class TestOnnxRunner(unittest.TestCase):
def _test_const_fold_unary_op(self, from_disk:bool):
runner = build_onnx(
nodes=[
onnx.helper.make_node('Expand', ['inp', 'shape'], ['expanded']),
onnx.helper.make_node('Exp', ['expanded'], ['output'])
],
outputs=[onnx.helper.make_tensor_value_info('output', onnx.TensorProto.FLOAT, (5,))],
initializers=[
onnx.helper.make_tensor('inp', onnx.TensorProto.FLOAT, (), [1.0]),
onnx.helper.make_tensor('shape', onnx.TensorProto.INT64, (1,), [5])
],
from_disk=from_disk)
output = runner({'inp': Tensor([1.0])})['output']
_check_ast_count(0, output)
def _test_const_fold_binary_op(self, from_disk:bool):
runner = build_onnx(
nodes=[onnx.helper.make_node('Add', ['inp', 'const'], ['output'])],
outputs=[onnx.helper.make_tensor_value_info('output', onnx.TensorProto.FLOAT, (4,))],
initializers=[
onnx.helper.make_tensor('inp', onnx.TensorProto.FLOAT, (4,), [1, 2, 3, 4]),
onnx.helper.make_tensor('const', onnx.TensorProto.FLOAT, (), [0])
],
from_disk=from_disk)
output = runner({'inp': Tensor([1, 2, 3, 4])})['output']
_check_ast_count(0, output)
@unittest.skip("const folding is removed")
def test_const_fold_from_disk(self):
self._test_const_fold_unary_op(True)
self._test_const_fold_binary_op(True)
@unittest.skip("const folding is removed")
def test_const_fold_from_memory(self):
self._test_const_fold_unary_op(False)
# TODO: understand this and fix this, bitcast related
# self._test_const_fold_binary_op(False)
def test_external_data_loading(self):
weights = np.arange(4, dtype=np.float32)
tensor_with_data = onnx.helper.make_tensor('weights', onnx.TensorProto.FLOAT, weights.shape, weights.tobytes(), raw=True)
graph = onnx.helper.make_graph(
nodes=[onnx.helper.make_node('Add', ['inp', 'weights'], ['output'])],
name='test_external',
inputs=[onnx.helper.make_tensor_value_info('inp', onnx.TensorProto.FLOAT, (1,))],
outputs=[onnx.helper.make_tensor_value_info('output', onnx.TensorProto.FLOAT, weights.shape)],
initializer=[tensor_with_data]
)
model = onnx.helper.make_model(graph)
with tempfile.TemporaryDirectory() as tmpdir:
tmp_path = pathlib.Path(tmpdir)
model_path = tmp_path / "model.onnx"
onnx.save_model(model, model_path, save_as_external_data=True, all_tensors_to_one_file=True, size_threshold=0, location="weights.onnx_data")
runner = OnnxRunner(model_path)
output = runner({'inp': Tensor([1])})['output']
np.testing.assert_equal(output.numpy(), weights + 1)
all_dtypes = list(OnnxDataType)
class TestOnnxRunnerDtypes(unittest.TestCase):
"""
Internal tensors (initializers, attributes) fallback to default dtype if unsupported by device.
External tensors (inputs) preserve their original dtype - user must ensure compatibility with device.
"""
def _get_expected_dtype(self, onnx_dtype: int, is_input: bool): return OnnxDataType(onnx_dtype).to_dtype()
@given(onnx_dtype=st.sampled_from(all_dtypes))
def test_input_dtype(self, onnx_dtype: int):
expected_dtype = self._get_expected_dtype(onnx_dtype, True)
runner = build_onnx(
nodes=[onnx.helper.make_node('Identity', ['input'], ['output'])],
inputs=[onnx.helper.make_tensor_value_info('input', onnx_dtype, ())],
outputs=[onnx.helper.make_tensor_value_info('output', onnx_dtype, ())],
from_disk=False)
self.assertEqual(runner.graph_inputs['input'].dtype, expected_dtype)
@given(onnx_dtype=st.sampled_from(all_dtypes))
def test_initializer_dtype(self, onnx_dtype: int):
expected_dtype = self._get_expected_dtype(onnx_dtype, False)
runner = build_onnx(
nodes=[onnx.helper.make_node('Identity', ['initializer'], ['output'])],
outputs=[onnx.helper.make_tensor_value_info('output', onnx_dtype, (2,))],
initializers=[onnx.helper.make_tensor('initializer', onnx_dtype, (2,), [1, 2])],
from_disk=False)
self.assertEqual(runner.graph_values['initializer'].dtype, expected_dtype)
@given(onnx_dtype=st.sampled_from(all_dtypes))
def test_node_attribute_dtype(self, onnx_dtype: int):
expected_dtype = self._get_expected_dtype(onnx_dtype, False)
value_tensor = onnx.helper.make_tensor('value', onnx_dtype, (2,), [1, 2])
runner = build_onnx(
nodes=[onnx.helper.make_node('Constant', [], ['output'], value=value_tensor)],
outputs=[onnx.helper.make_tensor_value_info('output', onnx_dtype, (2,))],
from_disk=False)
self.assertEqual(runner.graph_nodes[0].opts['value'].dtype, expected_dtype)
# from openpilot selfdrive/modeld/get_model_metadata.py
class MetadataOnnxPBParser(OnnxPBParser):
def _parse_ModelProto(self) -> dict:
obj: dict[str, Any] = {"graph": {"input": [], "output": []}, "metadata_props": []}
for fid, wire_type in self._parse_message(self.reader.len):
match fid:
case 7: obj["graph"] = self._parse_GraphProto()
case 14: obj["metadata_props"].append(self._parse_StringStringEntryProto())
case _: self.reader.skip_field(wire_type)
return obj
class TestOnnxMetadata(unittest.TestCase):
def test_metadata_props(self):
graph = onnx.helper.make_graph(
nodes=[onnx.helper.make_node('Identity', ['input'], ['output'])],
name='test',
inputs=[onnx.helper.make_tensor_value_info('input', onnx.TensorProto.FLOAT, (1, 3))],
outputs=[onnx.helper.make_tensor_value_info('output', onnx.TensorProto.FLOAT, (1, 3))],
)
model = onnx.helper.make_model(graph)
model.metadata_props.append(onnx.StringStringEntryProto(key="model_checkpoint", value="v1.0"))
model.metadata_props.append(onnx.StringStringEntryProto(key="output_slices", value="dGVzdA=="))
with tempfile.TemporaryDirectory() as tmpdir:
model_path = pathlib.Path(tmpdir) / "model.onnx"
onnx.save(model, model_path)
parsed = MetadataOnnxPBParser(model_path).parse()
# metadata_props should be accessible as dicts with "key" and "value"
self.assertEqual(len(parsed["metadata_props"]), 2)
self.assertEqual(parsed["metadata_props"][0]["key"], "model_checkpoint")
self.assertEqual(parsed["metadata_props"][0]["value"], "v1.0")
self.assertEqual(parsed["metadata_props"][1]["key"], "output_slices")
self.assertEqual(parsed["metadata_props"][1]["value"], "dGVzdA==")
if __name__ == '__main__':
unittest.main()

236
artifacts/package_sources/tinygrad/test/external/external_test_opt.py vendored Normal file
View File

@@ -0,0 +1,236 @@
#!/usr/bin/env python
import gc, unittest
import numpy as np
import torch
from tinygrad import GlobalCounters, Tensor, Device
from tinygrad.helpers import getenv, Context
from tinygrad.nn.state import get_parameters
from tinygrad.engine.realize import capturing, run_linear
from tinygrad.tensor import _to_np_dtype
class CLCache:
def __init__(self, allowed=None, strict=False, preclear=True, var_vals=None):
self.allowed, self.strict, self.preclear, self.var_vals = allowed, strict, preclear, var_vals if var_vals is not None else {}
self.count = 0
def add_linear(self, linear, var_vals):
self.count += len(linear.src)
run_linear(linear, var_vals)
def __enter__(self):
if self.preclear:
gc.collect()
for x in [x for x in gc.get_objects() if isinstance(x, Tensor)]:
x.realize()
GlobalCounters.reset()
capturing.append(self)
print("cache: entering")
return self
def __exit__(self, _type, value, traceback):
capturing.clear()
print(f"cache: exiting with size {self.count}", f"allowed {self.allowed}" if self.allowed is not None else "")
if self.allowed is not None:
assert self.count <= self.allowed, f"{self.count} > {self.allowed}"
from extra.models.convnext import ConvNeXt
from extra.models.efficientnet import EfficientNet
from extra.models.resnet import ResNet18
from extra.models.vit import ViT
@unittest.skipUnless(Device.DEFAULT == "CL", "Not Implemented")
class TestInferenceMinKernels(unittest.TestCase):
def setUp(self): self.enterContext(Context(TRAINING=0))
def test_convnext(self):
model = ConvNeXt()
for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype)))
img = Tensor.randn(1, 3, 224, 224)
with CLCache(143):
model(img).realize()
def test_enet(self):
model = EfficientNet(getenv("ENET_NUM", 0), has_se=False)
for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype)))
img = Tensor.randn(1, 3, 224, 224)
with CLCache(51):
model.forward(img).realize()
def test_enet_se(self):
model = EfficientNet(getenv("ENET_NUM", 0), has_se=True)
for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype)))
img = Tensor.randn(1, 3, 224, 224)
# TODO: this seems very high
with CLCache(115):
model.forward(img).realize()
def test_resnet(self):
model = ResNet18()
for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype)))
img = Tensor.randn(1, 3, 224, 224)
with CLCache(23):
model.forward(img).realize()
def test_vit(self):
model = ViT(embed_dim=192, num_heads=3)
for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype)))
img = Tensor.randn(1, 3, 224, 224)
with CLCache(209) as cache: # NOTE: this is way too high
out = model.forward(img)
assert cache.count == 0, "ViT prerealized?"
out.realize()
@unittest.skip("llama is fp16 but CI does not have fp16")
def test_llama(self):
from examples.llama import Transformer
args_tiny = {"dim": 512, "hidden_dim": 1024, "n_heads": 8, "n_layers": 4, "norm_eps": 1e-05, "vocab_size": 1000}
model = Transformer(**args_tiny)
for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype)))
inp = Tensor([[1,2,3,4]])
with CLCache(100):
model(inp, 0).realize()
@unittest.skipUnless(Device.DEFAULT == "CL", "Not Implemented")
class TestOptBinOp(unittest.TestCase):
def _test_no_binop_rerun(self, f1, f2=None, allowed=1):
a = Tensor.randn(16, 16)
b = Tensor.randn(16, 16)
with CLCache() as cache:
c = f1(a, b)
if f2 is not None: d = f2(a, b)
c.realize()
if f2 is not None: d.realize()
assert cache.count == allowed, "binop was rerun!"
if f2 is not None: np.testing.assert_allclose(c.numpy().ravel(), d.numpy().ravel(), rtol=1e-3, atol=1e-5)
def test_no_binop_rerun(self): return self._test_no_binop_rerun(lambda a,b: a*b, lambda a,b: (a*b).reshape(16, 16, 1))
def test_no_binop_rerun_alt(self): return self._test_no_binop_rerun(lambda a,b: (a*b).reshape(16, 16, 1), lambda a,b: a*b)
def test_no_binop_rerun_reduce_broadcast(self):
return self._test_no_binop_rerun(lambda a,b: a.sum()+b, lambda a,b: a.sum().reshape(1,1)+b, allowed=2)
@unittest.skip("this test started failing with the new change, based movementop issue")
def test_no_binop_rerun_transposed(self): return self._test_no_binop_rerun(lambda a,b: (a.T*b.T).T, lambda a,b: a*b)
def test_no_binop_rerun_mid_reshape(self): return self._test_no_binop_rerun(lambda a,b: (a*b).reshape(256)+a.reshape(256))
# currently non working tests
# def test_no_binop_rerun_preshape(self): return self._test_no_binop_rerun(lambda a,b: a.reshape(16, 16, 1)*b.reshape(16, 16, 1), lambda a,b: a*b)
#def test_no_binop_rerun_reduce(self): return self._test_no_binop_rerun(lambda a,b: (a*b).sum(), lambda a,b: (a*b).reshape(16, 16, 1).sum())
#def test_no_binop_rerun_reduce_alt(self): return self._test_no_binop_rerun(lambda a,b: a.sum(1)+b[0], lambda a,b: a.sum(1).reshape(1,16)+b[0])
@unittest.skipUnless(Device.DEFAULT == "CL", "Not Implemented")
class TestOptReduceLoop(unittest.TestCase):
def test_loop_left(self):
a = Tensor.randn(16, 16)
b = Tensor.randn(16, 16)
with CLCache() as cache:
t = a.sum(0)
b = t.reshape(16,1).expand(16,16).sum(0)
c = (t+b)
c.realize()
assert cache.count == 2, "loop left fusion broken"
def test_loop_right(self):
a = Tensor.randn(16, 16)
b = Tensor.randn(16, 16)
with CLCache() as cache:
t = a.sum(0)
b = t.reshape(16,1).expand(16,16).sum(0)
c = (b+t)
c.realize()
assert cache.count == 2, "loop right fusion broken"
@unittest.skipUnless(Device.DEFAULT == "CL", "Not Implemented")
class TestOptWChild(unittest.TestCase):
@unittest.skip("this no longer happens, use realize")
def test_unrealized_child(self):
a = Tensor.randn(16, 16)
b = Tensor.randn(16, 16)
with CLCache() as cache:
c = (a*b).sum()
d = c+1
e = c+2 # noqa: F841
d.realize()
assert cache.count == 2, "don't fuse if you have children"
@unittest.skipUnless(Device.DEFAULT == "CL", "Not Implemented")
class TestOpt(unittest.TestCase):
def test_muladd(self):
a,b,c = [Tensor.randn(2,2).realize() for _ in range(3)]
na,nb,nc = a.numpy(),b.numpy(),c.numpy()
with CLCache(allowed=1):
d = a * b + c
d.realize()
np.testing.assert_allclose(d.numpy(), na*nb+nc, rtol=1e-5, atol=1e-7)
def test_permute_was_pushed(self):
a = Tensor.randn(16, 16, 16)
with CLCache(1):
c = a.sum(2)
d = c.permute(1,0).contiguous()
d.realize()
np.testing.assert_allclose(a.numpy().sum(2).transpose(1,0), d.numpy(), rtol=1e-3, atol=1e-5)
def test_permute_was_pushed_through_contract_reshape(self):
a = Tensor.randn(4, 4, 4, 4, 4)
with CLCache(1):
c = a.sum(-1)
d = c.reshape(16,16).permute(1,0).contiguous()
d.realize()
np.testing.assert_allclose(a.numpy().sum(-1).reshape(16,16).transpose(1,0), d.numpy(), rtol=1e-3, atol=1e-5)
def test_permute_was_pushed_through_contractw1s_reshape(self):
a = Tensor.randn(4, 4, 4, 4, 4)
with CLCache(1):
c = a.sum(-1)
d = c.reshape(16,1,16).permute(2,1,0).contiguous()
d.realize()
np.testing.assert_allclose(a.numpy().sum(-1).reshape(16,1,16).transpose(2,1,0), d.numpy(), rtol=1e-3, atol=1e-5)
def test_permute_was_pushed_through_expand_reshape(self):
a = Tensor.randn(16, 16, 16)
with CLCache(1):
c = a.sum(2)
d = c.reshape(4,4,4,4).permute(2,3,0,1).contiguous()
d.realize()
np.testing.assert_allclose(a.numpy().sum(2).transpose(1,0).reshape(4,4,4,4), d.numpy(), rtol=1e-3, atol=1e-5)
def test_no_reduceop_rerun(self):
a = Tensor.randn(16, 16, 16)
with CLCache() as cache:
c = a.sum(2)
d = a.sum(2).permute(1,0)
c.realize()
d.realize()
cache_len = cache.count
np.testing.assert_allclose(c.numpy().transpose(1,0), d.numpy(), rtol=1e-3, atol=1e-5)
assert cache_len == 1, "reduceop was rerun!"
def test_no_reduceop_rerun_alt(self):
a = Tensor.randn(16, 16, 16)
with CLCache() as cache:
c = a.sum(2).permute(1,0)
d = a.sum(2)
c.realize()
d.realize()
cache_len = cache.count
np.testing.assert_allclose(c.numpy(), d.numpy().transpose(1,0), rtol=1e-3, atol=1e-5)
assert cache_len == 1, "reduceop was rerun!"
def test_expand_reduce_is_folded_on_same_axis(self):
for axis in [0, 1]:
for n in [4, 8, 16]:
b = torch.ones(n, n).sum(axis).reshape(n, 1).expand(n, n).sum(axis)
with CLCache(allowed=3):
a = Tensor.ones(n, n).contiguous().sum(axis).reshape(n, 1).expand(n, n).sum(axis)
a.realize()
np.testing.assert_allclose(a.numpy(), b.numpy(), rtol=1e-3, atol=1e-5)
def test_expand_reduce_is_folded_on_different_axes(self):
axis1, axis2 = 0, 1
for n in [4, 8, 16]:
b = torch.ones(n, n).sum(axis1).reshape(n, 1).expand(n, n).sum(axis2)
with CLCache(allowed=3):
a = Tensor.ones(n, n).contiguous().sum(axis1).reshape(n, 1).expand(n, n).sum(axis2)
a.realize()
np.testing.assert_allclose(a.numpy(), b.numpy(), rtol=1e-3, atol=1e-5)
if __name__ == '__main__':
unittest.main()

174
artifacts/package_sources/tinygrad/test/external/external_test_optim.py vendored Normal file
View File

@@ -0,0 +1,174 @@
#!/usr/bin/env python
import unittest
import numpy as np
import tensorflow as tf
from tensorflow.keras.optimizers import Lamb
from tensorflow.python.ops import math_ops
from extra.lr_scheduler import LRSchedulerGroup
from tinygrad.tensor import Tensor
from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup
from tinygrad.helpers import Context
from test.external.mlperf_resnet.lars_optimizer import LARSOptimizer
from examples.mlperf.lr_schedulers import PolynomialDecayWithWarmup
from test.external.mlperf_resnet.lars_util import PolynomialDecayWithWarmup as PolynomialDecayWithWarmup_tf
np.random.seed(1337)
x_init = np.random.randn(1,4).astype(np.float32)
W_init = np.random.randn(4,4).astype(np.float32)
m_init = np.random.randn(1,4).astype(np.float32)
class TinyNet:
def __init__(self):
self.x = Tensor(x_init.copy())
self.W = Tensor(W_init.copy())
self.m = Tensor(m_init.copy())
def forward(self):
out = self.x.matmul(self.W).relu()
out = out.log_softmax(1)
out = out.mul(self.m).add(self.m).sum()
return out
class TinyNetTF:
def __init__(self):
self.x = tf.Variable(x_init.copy(), trainable=True, name="x")
self.W = tf.Variable(W_init.copy(), trainable=True, name="W")
self.m = tf.constant(m_init.copy())
def forward(self):
out = tf.matmul(self.x, self.W)
out = tf.nn.relu(out)
out = tf.nn.log_softmax(out, axis=1)
out = tf.multiply(out, self.m) + self.m
out = tf.reduce_sum(out)
return out
def step(optim, steps=1, kwargs={}, scheduler=None, schedopts=None, do_optim=True):
net = TinyNet()
optim = optim([net.x, net.W], **kwargs)
if scheduler is not None: scheduler = scheduler(optim, **schedopts)
lrs = []
for _ in range(steps):
if do_optim:
out = net.forward()
optim.zero_grad()
out.backward()
lrs.append(optim.lr.item() if not isinstance(optim, OptimizerGroup) else optim.optimizers[0].lr.item())
if do_optim: optim.step()
if scheduler is not None: scheduler.step()
return lrs, net.x.detach().numpy(), net.W.detach().numpy()
def step_tf(optim, steps=1, kwargs={}, scheduler=None, schedopts=None, do_optim=True):
net = TinyNetTF()
if scheduler is not None: kwargs['lr'] = scheduler(**schedopts)
optim = optim(**kwargs)
lrs = []
for _ in range(steps):
if do_optim:
with tf.GradientTape() as tape:
out = net.forward()
lr_t = optim.learning_rate
# refer to test/external/mlperf_resnet/lars_optimizer.py:_prepare_local
if callable(lr_t): lr_t = lr_t(math_ops.cast(optim.iterations, tf.float32))
lrs.append(lr_t)
if do_optim:
grads = tape.gradient(out, [net.x, net.W])
optim.apply_gradients(zip(grads, [net.x, net.W]))
# optim calls scheduler in tf
else:
optim._iterations.assign_add(1)
return lrs, net.x.numpy(), net.W.numpy()
# skip list is skipping W
def create_tiny_lars(params, lr, skip_list=False):
if skip_list: return OptimizerGroup(LARS([params[0]], lr), SGD([params[1]], lr, classic=True, weight_decay=0., momentum=.9))
return LARS(params, lr)
def create_tf_lars(lr, skip_list=False): return LARSOptimizer(lr, skip_list=["W"] if skip_list else None)
def create_tf_lamb(lr=0.001, b1=0.9, b2=0.999, eps=1e-7, weight_decay=0.0):
return Lamb(learning_rate=float(lr), beta_1=b1, beta_2=b2, epsilon=eps, weight_decay=weight_decay)
def create_tiny_polylr(optim, initial_lr, end_lr, train_steps, warmup, power=2, skip_list=False):
assert power == 2
if skip_list: return LRSchedulerGroup(
PolynomialDecayWithWarmup(optim[0], initial_lr, end_lr, train_steps, warmup, power),
PolynomialDecayWithWarmup(optim[1], initial_lr, end_lr, train_steps, warmup, power))
return PolynomialDecayWithWarmup(optim, initial_lr, end_lr, train_steps, warmup, power)
def create_tf_polylr(initial_lr, end_lr, train_steps, warmup, power=2, skip_list=False):
assert power == 2
return PolynomialDecayWithWarmup_tf(1, 1, train_steps,
initial_learning_rate=initial_lr, end_learning_rate=end_lr, warmup_epochs=warmup)
class ExternalTestOptim(unittest.TestCase):
def setUp(self): self.enterContext(Context(TRAINING=1))
def _test_optim(self, tinygrad_optim, tensorflow_optim, steps, opts, atol, rtol, tiny_sched=None, tf_sched=None, schedopts=None, do_optim=True):
for x,y in zip(step(tinygrad_optim, steps=steps, kwargs=opts, scheduler=tiny_sched, schedopts=schedopts, do_optim=do_optim),
step_tf(tensorflow_optim, steps=steps, kwargs=opts, scheduler=tf_sched, schedopts=schedopts, do_optim=do_optim)):
np.testing.assert_allclose(x, y, atol=atol, rtol=rtol)
def _test_lamb(self, steps, opts, atol, rtol): self._test_optim(LAMB, create_tf_lamb, steps, opts, atol, rtol)
def _test_lars(self, steps, opts, atol, rtol): self._test_optim(create_tiny_lars, create_tf_lars, steps, opts, atol, rtol)
def _test_lars_polylr(self, steps, opts, schedopts, atol, rtol, do_optim=True):
self._test_optim(create_tiny_lars, create_tf_lars, steps, opts, atol, rtol,
tiny_sched=create_tiny_polylr, tf_sched=create_tf_polylr, schedopts=schedopts, do_optim=do_optim)
def test_lamb(self): self._test_lamb(1, {'lr': 0.001}, 1e-5, 0)
def test_lamb_high_lr(self): self._test_lamb(1, {'lr': 10}, 1e-5, 1e-5)
def test_multistep_lamb(self): self._test_lamb(10, {'lr': 0.001}, 1e-5, 0)
def test_multistep_lamb_high_lr(self): self._test_lamb(10, {'lr': 10}, 1e-5, 3e-4)
def test_lars(self): self._test_lars(1, {'lr': 0.01}, 1e-5, 0)
def test_lars_high_lr(self): self._test_lars(1, {'lr': 10}, 1e-5, 1e-5)
def test_multistep_lars(self): self._test_lars(10, {'lr': 0.001}, 1e-5, 0)
def test_multistep_lars_high_lr(self): self._test_lars(10, {'lr': 10}, 1e-5, 3e-4)
def test_lars_skip(self): self._test_lars(10, {'lr': 10, 'skip_list': True}, 1e-5, 3e-4)
def test_lars_skip_high_lr(self): self._test_lars(1, {'lr': 10, 'skip_list': True}, 1e-5, 1e-5)
def test_lars_skip_multistep(self): self._test_lars(10, {'lr': 0.001, 'skip_list': True}, 1e-5, 0)
def test_lars_skip_multistep_high_lr(self): self._test_lars(10, {'lr': 10, 'skip_list': True}, 1e-5, 3e-4)
def test_lars_polylr(self):
self._test_lars_polylr(10, {'lr': 1.0}, {
'initial_lr': 1.0,
'end_lr': 1e-4,
'train_steps': 10,
'warmup': 3
}, 1e-5, 1e-5)
def test_lars_polylr_large(self):
self._test_lars_polylr(100, {'lr': 10.0}, {
'initial_lr': 10.0,
'end_lr': 1e-5,
'train_steps': 100,
'warmup': 43
}, 1e-5, 1e-5, do_optim=False)
def test_lars_polylr_skip(self):
self._test_lars_polylr(10, {'lr': 1.0, 'skip_list': True}, {
'initial_lr': 1.0,
'end_lr': 1e-4,
'train_steps': 10,
'warmup': 3,
'skip_list': True
}, 1e-5, 1e-5)
@unittest.skip("slow, but you can run this locally to check")
def test_lars_polylr_resnet(self):
train_files = 1_281_167
BS = 624
steps_per_epoch = train_files // BS
epochs = 45
warmup_epochs = 5
self._test_lars_polylr(steps_per_epoch * epochs, {'lr': 10.4}, {
'initial_lr': 10.4,
'end_lr': 1e-4,
# step counts for BS=624 EPOCHS=45 resnet
'train_steps': steps_per_epoch * epochs,
'warmup': steps_per_epoch * warmup_epochs,
}, 1e-5, 1e-5, do_optim=False)
if __name__ == '__main__':
unittest.main()

57
artifacts/package_sources/tinygrad/test/external/external_test_process_count.py vendored Normal file
View File

@@ -0,0 +1,57 @@
import os, sys, time, multiprocessing
N = int(os.environ.get("NPROC", str(os.cpu_count())))
DEVICE = os.environ.get("DEV", "AMD")
# this tests the total number of processes that can be running tinygrad at a time
def proc(i, device, stop_evt):
from tinygrad import Tensor
try:
a = Tensor.ones(2, device=device).contiguous()
b = Tensor.ones(2, device=device).contiguous()
c = (a + b).realize()
assert c.tolist() == [2, 2]
except Exception as e:
# fail if it fails
print(f"[child {i:2d}] tinygrad op failed: {e}", file=sys.stderr)
# non-zero exit code propagated back to parent
sys.exit(1)
# TODO: wait here for global exit if success. fail if it fails
# -> We wait on a global Event shared from the parent.
print(f"[child {i:2d}] success")
stop_evt.wait()
# Normal successful exit
sys.exit(0)
if __name__ == "__main__":
print(f"testing {N} concurrent tinygrad processes")
# global exit event, shared by all children
stop_evt = multiprocessing.Event()
procs = []
# launch n proc of proc 1 per 200 ms
for i in range(N):
p = multiprocessing.Process(target=proc, args=(i, DEVICE, stop_evt), name=f"tinygrad-proc-{i}")
p.start()
procs.append(p)
time.sleep(0.1) # 100 ms between launches
# signal global exit
time.sleep(0.5)
stop_evt.set()
# join all children
for p in procs: p.join()
# check for failures
failed = [p for p in procs if p.exitcode != 0]
if failed:
print(f"{len(failed)} / {len(procs)} processes failed "
f"with exit codes: {[p.exitcode for p in failed]}", file=sys.stderr)
sys.exit(1)
print(f"All {len(procs)} tinygrad processes ran successfully")
sys.exit(0)

134
artifacts/package_sources/tinygrad/test/external/external_test_schedule_scaling.py vendored Normal file
View File

@@ -0,0 +1,134 @@
import unittest, time
from tinygrad import Tensor
class TestScheduleScaling(unittest.TestCase):
"""Test that .schedule() scales linearly with graph size (no O(n^2) behavior)."""
def _assert_linear(self, fn, n_small=200, n_large=1000):
"""Assert schedule time scales at most ~linearly: time(n_large)/time(n_small) should be close to n_large/n_small."""
fn(n_small).schedule_linear() # warmup
t_small = min(self._time_schedule(fn, n) for n in [n_small]*3)
t_large = min(self._time_schedule(fn, n) for n in [n_large]*3)
size_ratio = n_large / n_small # 5.0
time_ratio = t_large / t_small
# O(n) -> time_ratio ~ 5, O(n^2) -> time_ratio ~ 25. threshold at 10 catches n^2 with margin.
self.assertLess(time_ratio / size_ratio, 2.0,
f"schedule appears superlinear: n={n_small} {t_small*1e3:.1f}ms, n={n_large} {t_large*1e3:.1f}ms "
f"(time grew {time_ratio:.1f}x for {size_ratio:.0f}x size, per-node ratio {time_ratio/size_ratio:.2f})")
@staticmethod
def _time_schedule(fn, n) -> float:
st = time.perf_counter()
fn(n).schedule_linear()
return time.perf_counter() - st
# *** rangeify: ending_ranges accumulation and consumer merge ***
# ending_ranges accumulation via sum([], []) and nested scan in run_rangeify.
# this creates reduce ops whose ending_ranges lists grow with graph depth, causing O(n^2) list copies.
def test_multi_reduce_scaling(self):
def multi_reduce(n):
x = Tensor.empty(256, 256)
for _ in range(n):
s = x.sum(axis=-1, keepdim=True)
x = x + s + s
return x
self._assert_linear(multi_reduce)
# reduce+elementwise chain stresses ending_ranges propagation and post-rangeify rewrites
def test_wide_reduce_scaling(self):
def wide_reduce(n):
x = Tensor.empty(256, 256)
for _ in range(n):
x = x + x.sum(axis=-1, keepdim=True)
return x
self._assert_linear(wide_reduce)
# expand ops inject into ending_ranges via the EXPAND path in run_rangeify
def test_expand_reduce_scaling(self):
def expand_reduce(n):
x = Tensor.empty(256, 1)
for _ in range(n):
y = x.expand(256, 256)
x = (y + y).sum(axis=-1, keepdim=True)
return x
self._assert_linear(expand_reduce)
# *** graph_rewrite: multi-consumer DAG patterns ***
# multi-consumer diamond pattern (fan-out/fan-in) stresses consumer_rngs merge in run_rangeify
def test_diamond_scaling(self):
def diamond(n):
x = Tensor.empty(256, 256)
for _ in range(n):
a = x + 1
b = x + 2
x = a + b
return x
self._assert_linear(diamond)
# elementwise chain baseline — should be trivially O(n)
def test_chain_scaling(self):
def chain(n):
x = Tensor.empty(256, 256)
for _ in range(n): x = x + 1
return x
self._assert_linear(chain)
# softmax has multi-consumer structure (x used for max, exp, and sum), stresses graph_rewrite on DAGs
def test_softmax_scaling(self):
def softmax_chain(n):
x = Tensor.empty(64, 256)
for _ in range(n): x = x.softmax(axis=-1)
return x
self._assert_linear(softmax_chain)
# *** post-rangeify: symbolic rewrites, kernel splitting ***
# matmul chain stresses symbolic+reduce_collapse and split_store
def test_matmul_scaling(self):
def matmul_chain(n):
xs = [Tensor.empty(32, 32) for _ in range(n + 1)]
result = xs[0]
for i in range(n): result = result @ xs[i + 1]
return result
self._assert_linear(matmul_chain)
# contiguous chain stresses remove_bufferize callbacks (toposort per BUFFERIZE node)
def test_contiguous_scaling(self):
def contiguous_chain(n):
x = Tensor.empty(256, 256)
for _ in range(n): x = (x + 1).contiguous()
return x
self._assert_linear(contiguous_chain)
# *** schedule: AFTER handling, assign ***
# assign chain stresses AFTER cycle detection (toposort inside toposort loop in get_rangeify_map)
def test_assign_scaling(self):
def assign_chain(n):
x = Tensor.empty(256, 256).realize()
for _ in range(n): x.assign(x + 1)
return x
self._assert_linear(assign_chain)
# layernorm has multi-consumer reduces (mean reused in variance), stresses consumer_rngs merge and symbolic rewrites
def test_layernorm_scaling(self):
def layernorm_chain(n):
x = Tensor.empty(64, 256)
for _ in range(n):
mean = x.mean(axis=-1, keepdim=True)
var = ((x - mean) ** 2).mean(axis=-1, keepdim=True)
x = (x - mean) / (var + 1e-5).sqrt()
return x
self._assert_linear(layernorm_chain)
# concat chain stresses MSTACK/MSELECT handling and wide SINK construction
def test_concat_scaling(self):
def concat_chain(n):
parts = [Tensor.empty(4, 256) + i for i in range(n)]
return parts[0].cat(*parts[1:])
self._assert_linear(concat_chain)
if __name__ == '__main__':
unittest.main(verbosity=2)

50
artifacts/package_sources/tinygrad/test/external/external_test_simple_tokenizer.py vendored Normal file
View File

@@ -0,0 +1,50 @@
import functools, multiprocessing
from transformers import AutoTokenizer
from datasets import load_dataset
from tinygrad.llm.cli import SimpleTokenizer
from tinygrad.helpers import tqdm, getenv, partition
@functools.cache
def get_tokenizers():
print("getting tokenizers")
base_tokenizer = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B-Instruct")
special_tokens, normal_tokens = partition(((t, tid) for t, tid in base_tokenizer.vocab.items()), lambda e: e[1] in base_tokenizer.all_special_ids)
simple_tokenizer = SimpleTokenizer(dict(normal_tokens), dict(special_tokens))
return base_tokenizer, simple_tokenizer
def test_tokenize(samp) -> bool:
base_tokenizer, simple_tokenizer = get_tokenizers()
idx, txt = samp
try: simple_tokens = tuple(simple_tokenizer.encode(txt))
except RuntimeError: simple_tokens = ()
base_tokens = tuple(base_tokenizer.encode(txt, add_special_tokens=False))
if simple_tokens != base_tokens:
print(f"tokens mismatch at index: {idx}.\n")
color_codes = [91, 92, 94, 93, 95]
def color_tokens(tids):
return "".join(f"\033[{color_codes[i%len(color_codes)]}m{base_tokenizer.decode([t])}" for i, t in enumerate(tids)) + "\033[0m"
print("simple: ", color_tokens(simple_tokens))
print("official:", color_tokens(base_tokens) + "\n")
return False
if simple_tokenizer.decode(simple_tokens) != txt:
print(f"decode mismatch at {idx}")
return False
return True
# use ALLOW_FAILED=-1 to go over the entire dataset without printing.
if __name__ == "__main__":
print("loading datasets")
ds = load_dataset("OpenAssistant/oasst1")
loaded_ds = [(idx, el["text"]) for idx, el in enumerate(ds["train"])]
print(f"loaded {len(loaded_ds)}")
allow_failed = getenv("ALLOW_FAILED", 10)
fail_count, total = 0, 0
with multiprocessing.Pool(16) as pool:
for good in tqdm(pool.imap_unordered(test_tokenize, loaded_ds), total=len(loaded_ds)):
total += 1
if not good:
fail_count += 1
allow_failed -= 1
if allow_failed == 0: break
print(f"{fail_count}/{total} samples are inconsistent with the official tokenizer.")

58
artifacts/package_sources/tinygrad/test/external/external_test_speed_llama.py vendored Normal file
View File

@@ -0,0 +1,58 @@
# NOTE: this only tests the speed of the LLaMA codegen, it doesn't actually run the net
import unittest, time
from examples.llama import Transformer, MODEL_PARAMS
from tinygrad.tensor import Tensor
from tinygrad import Device
from tinygrad.nn.state import get_state_dict
from tinygrad.device import Allocator, Compiled
from tinygrad.codegen import to_program_cache
from tinygrad.helpers import Profiling
class FakeProgram:
def __init__(self, name:str, lib:bytes, *args, **kwargs): pass
def __call__(self, *bufs, global_size, local_size, vals=(), wait=False, **kw): pass
class FakeAllocator(Allocator[Compiled]):
def _alloc(self, sz, options): return None
def _copyin(self, dest, src:memoryview): pass
class TestLLaMASpeed(unittest.TestCase):
def test_llama_compile(self):
backup_program = Device[Device.DEFAULT].runtime
backup_allocator = Device[Device.DEFAULT].allocator
backup_compiler = Device[Device.DEFAULT].compiler.compile_cached
Device[Device.DEFAULT].runtime = FakeProgram
Device[Device.DEFAULT].allocator = FakeAllocator(Device.default)
print("testing llama python run time")
model = Transformer(**MODEL_PARAMS["1"]["7B"]["args"])
print("built model")
# assign fake tensors to the values
for v in get_state_dict(model).values(): v.assign(Tensor.empty(*v.shape, dtype=v.dtype))
print("assigned empty tensors, doing warmup")
def run_llama(st, empty_cache=True):
if empty_cache: to_program_cache.clear()
tms = [time.perf_counter()]
for i in range(5):
model(Tensor([[1,2,3,4]]), i).realize()
tms.append(time.perf_counter())
timings = [(tms[i+1]-tms[i])*1000 for i in range(len(tms)-1)]
print(f"{st:15s} mean runtime: {sum(timings)/len(timings):7.2f}ms, runs: ", ", ".join(f'{x:7.2f}' for x in timings))
run_llama("codegen(0)")
run_llama("codegen(1)")
# test no compiler use for this
Device[Device.DEFAULT].compiler.compile_cached = None
run_llama("methodcache", False)
with Profiling(sort='time', frac=0.1, fn="/tmp/llama.prof", ts=5):
run_llama("profile", False)
Device[Device.DEFAULT].runtime = backup_program
Device[Device.DEFAULT].allocator = backup_allocator
Device[Device.DEFAULT].compiler.compile_cached = backup_compiler
if __name__ == '__main__':
TestLLaMASpeed().test_llama_compile()
#unittest.main()

48
artifacts/package_sources/tinygrad/test/external/external_test_speed_theoretical.py vendored Normal file
View File

@@ -0,0 +1,48 @@
import time, unittest
from tinygrad import Tensor, TinyJit, Device, dtypes
from tinygrad.helpers import getenv, GlobalCounters
SZMAX = getenv("SZMAX", 10)
SZMIN = min(SZMAX, getenv("SZMIN", 10))
def _test(tcount, fxn, dtype=dtypes.float):
print(f"**** testing {fxn.__name__} {dtype}")
allgbs = []
for sz in range(SZMIN, SZMAX+1):
jfxn = TinyJit(fxn)
ts = [Tensor.zeros((2**sz)*1024*1024, dtype=dtype).contiguous().realize() for _ in range(tcount)]
tms = []
for _ in range(10):
ts = [(x+1).realize() for x in ts]
Device.default.synchronize()
GlobalCounters.global_ops = 0
GlobalCounters.global_mem = 0
st = time.perf_counter()
jfxn(*ts).nbytes()
Device.default.synchronize()
tms.append(time.perf_counter() - st)
ops, mem = GlobalCounters.global_ops, GlobalCounters.global_mem
gflops = ops*1e-9/min(tms)
gbs = mem*1e-9/min(tms)
print(f"{ts[0].nbytes()/(1024*1024):10.0f} MB, {min(tms)*1e3:6.2f} ms {gbs:10.2f} GB/s {gflops:10.2f} GFLOPS {str(ts[0].shape):20s}")
allgbs.append(gbs)
return max(allgbs)
MEMBW = getenv("MEMBW", 10)
class TestRamBandwidth(unittest.TestCase):
def test_add(self): self.assertGreater(_test(2, Tensor.add), MEMBW)
def test_exp(self): self.assertGreater(_test(1, Tensor.exp), MEMBW)
def test_sum(self): self.assertGreater(_test(1, Tensor.sum), MEMBW)
# ratio between MEM and FLOPS < 1000
# NOTE: On AMD, (x*x)+1 gets ~30 TFLOPS, (x*x)+3 gets ~60 TFLOPS
def flopsmax(x):
for _ in range(500): x = (x*x)+3
return x
class TestFlops(unittest.TestCase):
def test_flops_int8(self): _test(1, flopsmax, dtypes.int8)
def test_flops_fp16(self): _test(1, flopsmax, dtypes.half)
def test_flops_fp32(self): _test(1, flopsmax)
if __name__ == '__main__':
unittest.main()

78
artifacts/package_sources/tinygrad/test/external/external_test_tlsf.py vendored Normal file
View File

@@ -0,0 +1,78 @@
import unittest
from tinygrad.runtime.support.memory import TLSFAllocator
class TestTLSFAllocator(unittest.TestCase):
def setUp(self):
self.allocator = TLSFAllocator(1024, block_size=16)
def test_basic_alloc_free(self):
addr1 = self.allocator.alloc(32)
self.assertEqual(addr1, 0)
addr2 = self.allocator.alloc(64)
self.assertEqual(addr2, 32)
self.allocator.free(addr1)
addr3 = self.allocator.alloc(32)
self.assertEqual(addr3, 0)
def test_block_size_alignment(self):
addr1 = self.allocator.alloc(20)
addr2 = self.allocator.alloc(35)
self.assertEqual(addr1 % 16, 0)
self.assertEqual(addr2 % 16, 0)
def test_merge_blocks(self):
addr1 = self.allocator.alloc(32)
addr2 = self.allocator.alloc(32)
self.allocator.alloc(32)
self.allocator.free(addr1)
self.allocator.free(addr2)
addr4 = self.allocator.alloc(64)
self.assertEqual(addr4, addr1)
def test_split_blocks(self):
addr1 = self.allocator.alloc(128)
self.allocator.free(addr1)
addr2 = self.allocator.alloc(32)
self.assertEqual(addr2, addr1)
addr3 = self.allocator.alloc(32)
self.assertEqual(addr3, addr1 + 32)
def test_out_of_memory(self):
with self.assertRaises(MemoryError):
self.allocator.alloc(2048)
def test_fragmentation_handling(self):
addrs = []
for _ in range(5):
addrs.append(self.allocator.alloc(32))
# Free alternate blocks
for i in range(0, len(addrs), 2):
self.allocator.free(addrs[i])
def test_custom_start_address(self):
allocator = TLSFAllocator(1024, start_addr=1000)
addr1 = allocator.alloc(32)
self.assertEqual(addr1, 1000)
addr2 = allocator.alloc(64)
self.assertEqual(addr2, 1032)
def test_block_tracking(self):
addr1 = self.allocator.alloc(32)
addr2 = self.allocator.alloc(64)
self.assertTrue(addr1 in [addr - self.allocator.start_addr for addr in self.allocator.blocks])
self.assertTrue(addr2 in [addr - self.allocator.start_addr for addr in self.allocator.blocks])
self.allocator.free(addr1)
self.assertTrue(addr1 in [addr - self.allocator.start_addr for addr in self.allocator.blocks])
if __name__ == '__main__':
unittest.main()

38
artifacts/package_sources/tinygrad/test/external/external_test_usb_asm24.py vendored Normal file
View File

@@ -0,0 +1,38 @@
import unittest
from tinygrad.helpers import Timing, getenv
from tinygrad import Tensor, Device
import numpy as np
class TestDevCopySpeeds(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.sz = getenv("SIZE", 2e6)
cls.dev = Device["AMD"]
if not cls.dev.is_usb(): raise unittest.SkipTest("only test this on USB devices")
def testCopyCPUtoDefault(self):
for _ in range(10):
t = Tensor.ones(self.sz, device="CPU", dtype='uchar').contiguous().realize()
with Timing(f"copyin of {t.nbytes()/1e6:.2f} MB: ", on_exit=lambda ns: f" @ {t.nbytes()/ns * 1e3:.2f} MB/s"): # noqa: F821
t.to(Device.DEFAULT).realize()
Device[Device.DEFAULT].synchronize()
del t
def testCopyDefaulttoCPU(self):
t = Tensor.ones(self.sz, dtype='uchar').contiguous().realize()
for _ in range(10):
with Timing(f"copyout of {t.nbytes()/1e6:.2f} MB: ", on_exit=lambda ns: f" @ {t.nbytes()/ns * 1e3:.2f} MB/s"):
t.to('CPU').realize()
def testValidateCopies(self):
t = Tensor.randn(self.sz, device="CPU", dtype='uchar').contiguous().realize()
x = t.to(Device.DEFAULT).realize()
Device[Device.DEFAULT].synchronize()
y = x.to('CPU').realize()
np.testing.assert_equal(t.numpy(), y.numpy())
del x, y, t
if __name__ == "__main__":
unittest.main()

83
artifacts/package_sources/tinygrad/test/external/external_test_whisper_librispeech.py vendored Normal file
View File

@@ -0,0 +1,83 @@
import unittest
import torch
import tqdm
import torchaudio
import pathlib
import jiwer
import os
import numpy as np
from whisper.normalizers import EnglishTextNormalizer
from examples.whisper import init_whisper, transcribe_waveform
class TestWhisperLibriSpeech(unittest.TestCase):
# reference WERs determined by running https://github.com/openai/whisper/blob/main/notebooks/LibriSpeech.ipynb
# the values should be consistent with the paper D.1.1 https://cdn.openai.com/papers/whisper.pdf#page=22
# tinygrad WERs do not perfectly match due to what seem to be precision differences vs torch
def test_en_tiny(self):
run_evaluation("tiny.en", 0.056629001883239174, 0.05655609406528749)
def test_tiny(self):
run_evaluation("tiny", 0.0771121409407306, 0.07558413638335187)
def test_en_base(self):
run_evaluation("base.en", 0.041412520064205455, 0.04271408904897505)
def test_en_small(self):
run_evaluation("small.en", 0.03369011117172363, 0.030531615969223228)
def run_evaluation(model_name, tinygrad_expected_wer, reference_wer):
dataset = LibriSpeech()
batch_size=16
loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size)
model, enc = init_whisper(model_name, batch_size=batch_size)
hypotheses = []
references = []
for audio, texts in tqdm.tqdm(loader):
transcriptions = transcribe_waveform(model, enc, audio.numpy(), truncate=True)
hypotheses.extend(transcriptions)
references.extend(texts)
normalizer = EnglishTextNormalizer()
normalized_hypotheses = [normalizer(text) for text in hypotheses]
normalized_references = [normalizer(text) for text in references]
wer = jiwer.wer(normalized_hypotheses, normalized_references)
np.testing.assert_almost_equal(wer, tinygrad_expected_wer)
print(f'tinygrad WER {wer} vs reference WER {reference_wer}')
del model, enc
class LibriSpeech(torch.utils.data.Dataset):
def __init__(self):
folder = pathlib.Path(__file__).parent.parent.parent / "extra" / "datasets" / "librispeech"
if not os.path.exists(folder):
os.makedirs(folder)
self.dataset = torchaudio.datasets.LIBRISPEECH(
root=folder,
url="test-clean",
download=True,
)
def __len__(self):
return len(self.dataset)
def __getitem__(self, item):
audio, sample_rate, text, _, _, _ = self.dataset[item]
assert sample_rate == 16000
return pad_or_trim_tensor(audio[0]), text
def pad_or_trim_tensor(tensor, target_len=480000):
curr_len = len(tensor)
if curr_len == target_len:
return tensor
elif curr_len < target_len:
return torch.cat((tensor, torch.zeros(target_len - curr_len)))
else:
return tensor[:target_len]
if __name__ == '__main__':
unittest.main()

73
artifacts/package_sources/tinygrad/test/external/external_test_yolov8.py vendored Normal file
View File

@@ -0,0 +1,73 @@
import numpy as np
from examples.yolov8 import YOLOv8, get_variant_multiples, preprocess, label_predictions, postprocess
from tinygrad import Tensor
import unittest
import io, cv2
import onnxruntime as ort
import ultralytics
from tinygrad.nn.state import safe_load, load_state_dict
from tinygrad.helpers import fetch
class TestYOLOv8(unittest.TestCase):
def test_all_load_weights(self):
for variant in ['n', 's', 'm', 'l', 'x']:
depth, width, ratio = get_variant_multiples(variant)
TinyYolov8 = YOLOv8(w=width, r=ratio, d=depth, num_classes=80)
state_dict = safe_load(fetch(f'https://gitlab.com/r3sist/yolov8_weights/-/raw/master/yolov8{variant}.safetensors'))
load_state_dict(TinyYolov8, state_dict)
print(f'successfully loaded weights for yolov{variant}')
def test_predictions(self):
test_image_urls = ['https://raw.githubusercontent.com/ultralytics/yolov5/master/data/images/bus.jpg', 'https://www.aljazeera.com/wp-content/uploads/2022/10/2022-04-28T192650Z_1186456067_UP1EI4S1I0P14_RTRMADP_3_SOCCER-ENGLAND-MUN-CHE-REPORT.jpg']
variant = 'n'
depth, width, ratio = get_variant_multiples(variant)
TinyYolov8 = YOLOv8(w=width, r=ratio, d=depth, num_classes=80)
state_dict = safe_load(fetch(f'https://gitlab.com/r3sist/yolov8_weights/-/raw/master/yolov8{variant}.safetensors'))
load_state_dict(TinyYolov8, state_dict)
for i in range(len(test_image_urls)):
img = cv2.imdecode(np.frombuffer(fetch(test_image_urls[i]).read_bytes(), np.uint8), 1)
test_image = preprocess([img])
predictions = TinyYolov8(test_image)
labels = label_predictions(predictions.numpy())
assert labels == {5: 1, 0: 4, 11: 1} if i == 0 else labels == {0: 12, 29: 1, 32: 1}
def test_forward_pass_torch_onnx(self):
variant = 'n'
weights_location = fetch(f'https://gitlab.com/r3sist/yolov8_weights/-/raw/master/yolov8{variant}.safetensors')
weights_location_pt = fetch(f'https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8{variant}.pt', name=f"yolov8{variant}.pt") # it needs the pt extension # noqa: E501
weights_location_onnx = weights_location_pt.parent / f"yolov8{variant}.onnx"
# the ultralytics export prints a lot of unneccesary things
if not weights_location_onnx.is_file():
model = ultralytics.YOLO(model=weights_location_pt, task='Detect')
model.export(format="onnx",imgsz=[640, 480])
depth, width, ratio = get_variant_multiples(variant)
TinyYolov8 = YOLOv8(w=width, r=ratio, d=depth, num_classes=80)
state_dict = safe_load(weights_location)
load_state_dict(TinyYolov8, state_dict)
image_location = [np.frombuffer(io.BytesIO(fetch('https://raw.githubusercontent.com/ultralytics/yolov5/master/data/images/bus.jpg').read_bytes()).read(), np.uint8)] # noqa: E501
orig_image = [cv2.imdecode(image_location[0], 1)]
input_image = preprocess(orig_image)
onnx_session = ort.InferenceSession(weights_location_onnx)
onnx_input_name = onnx_session.get_inputs()[0].name
onnx_output_name = onnx_session.get_outputs()[0].name
onnx_output = onnx_session.run([onnx_output_name], {onnx_input_name: input_image.numpy()})
tiny_output = TinyYolov8(input_image).numpy()
onnx_output = postprocess(Tensor(onnx_output[0])).numpy()
#invalid boxes are multiplied by zero in postprocess
onnx_output = onnx_output[onnx_output[:, 4] != 0]
tiny_output = tiny_output[tiny_output[:, 4] != 0]
# currently rtol is 0.025 because there is a 1-2% difference in our predictions
# because of the zero padding in SPPF module (line 280) maxpooling layers rather than the -infinity in torch.
# This difference does not make a difference "visually".
np.testing.assert_allclose(onnx_output, tiny_output, atol=5e-4, rtol=0.01)
if __name__ == '__main__':
unittest.main()

85
artifacts/package_sources/tinygrad/test/external/external_uop_gc.py vendored Normal file
View File

@@ -0,0 +1,85 @@
import gc
from tinygrad import Tensor, UOp, Device, nn
from tinygrad.schedule import schedule_cache
from tinygrad.codegen import to_program, to_program_cache
from tinygrad.schedule.indexing import apply_movement_op, _apply_reshape
from tinygrad.uop.divandmod import fold_divmod_general
from test.test_tiny import TestTiny
def uops_allocated(): return sum([isinstance(x, UOp) for x in gc.get_objects()])
def print_uops():
for x in gc.get_objects():
if isinstance(x, UOp): print(x)
def start(): pass
def single_tensor(): Tensor([2])
def two_plus_two(): Tensor([2])+Tensor([2])
def two_plus_two_schedule(): (Tensor([2])+Tensor([2])).schedule_linear()
def two_plus_two_kernel():
linear = (Tensor([2])+Tensor([2])).schedule_linear()
to_program(linear.src[-1].src[0], Device.default.renderer)
def two_plus_two_linearize():
linear = (Tensor([2])+Tensor([2])).schedule_linear()
to_program(linear.src[-1].src[0], Device.default.renderer)
def two_plus_two_realize(): (Tensor([2])+Tensor([2])).realize()
def two_plus_two_item(): (Tensor([2])+Tensor([2])).item()
def gradient_test():
x = Tensor.eye(3)
y = Tensor([[2.0,0,-2.0]])
z = y.matmul(x).sum()
z.backward()
def realized_eye():
Tensor.eye(3).clone().realize()
def realized_list():
Tensor([[2.0,0,-2.0]]).realize()
def kernel_matmul():
x = Tensor.eye(3)
y = Tensor([[2.0,0,-2.0]])
z = y.matmul(x)
linear = z.schedule_linear()
to_program(linear.src[-1].src[0], Device.default.renderer)
def realized_matmul():
x = Tensor.eye(3)
y = Tensor([[2.0,0,-2.0]])
z = y.matmul(x)
Tensor.realize(z)
def realized_gradient():
x = Tensor.eye(3)
y = Tensor([[2.0,0,-2.0]])
z = y.matmul(x).sum()
z.backward()
Tensor.realize(x, y, z, x.grad, y.grad)
def nn_batchnorm(): nn.BatchNorm(64)
def nn_conv2d(): nn.Conv2d(64, 64, 3)
def plus(): TestTiny().test_plus()
def mnist(): TestTiny().test_mnist()
def mnist_backward(): TestTiny().test_mnist_backward()
tests = [start, single_tensor, two_plus_two, two_plus_two_schedule, two_plus_two_kernel,
two_plus_two_linearize, two_plus_two_realize, two_plus_two_item, gradient_test,
realized_eye, realized_list, kernel_matmul, realized_matmul, realized_gradient,
nn_batchnorm, nn_conv2d, plus, mnist, mnist_backward]
if __name__ == "__main__":
gc.disable()
start_uops = uops_allocated()
# there's a few consts created as default values
print_uops()
for t in tests:
t()
# these caches will keep uops alive
schedule_cache.clear()
to_program_cache.clear()
apply_movement_op.cache_clear()
_apply_reshape.cache_clear()
fold_divmod_general.cache_clear()
Tensor._device_seeds.clear()
Tensor._device_rng_counters.clear()
new_uops = uops_allocated()
gc.collect()
new_uops_gc = uops_allocated()
print(f"{t.__name__:30s}: {new_uops:3d} -> {new_uops_gc:3d}")
if new_uops != start_uops: print_uops()
assert new_uops == start_uops

25
artifacts/package_sources/tinygrad/test/external/fuzz_fast_idiv.py vendored Normal file
View File

@@ -0,0 +1,25 @@
import random
import z3
from tinygrad import dtypes, Device
from tinygrad.uop.validate import uops_to_z3, z3_cdiv
from tinygrad.uop.ops import UOp
from tinygrad.codegen.decomp.op import fast_idiv
random.seed(42)
powers_of_two = [2**i for i in range(64)]
if __name__ == "__main__":
for i in range(10_000):
if i % 1000 == 0:
print(f"Progress: {i}")
dt = random.choice(dtypes.ints)
u = UOp.variable('x', random.randint(dt.min, 0), random.randint(1, dt.max), dtype=dt)
d = random.randint(1, max(1, u.vmax)*2)
if d in powers_of_two: continue
expr = fast_idiv(Device[Device.DEFAULT].renderer, u, d)
if expr is None: continue
solver = z3.Solver()
z3_expr, x =uops_to_z3(solver, expr, u)
if solver.check(z3_expr != z3_cdiv(x, d)) == z3.sat:
assert False, f"Failed: {expr.render()} != x//{d} at x={solver.model()}\nx={u}\nd={d}\n{z3_expr=}\n{x/d=}"

127
artifacts/package_sources/tinygrad/test/external/fuzz_graph.py vendored Normal file
View File

@@ -0,0 +1,127 @@
import random
import numpy as np
from tinygrad.device import Buffer, Device
from tinygrad.helpers import Context, getenv
from tinygrad.dtype import dtypes
from tinygrad.tensor import Tensor, _to_np_dtype
from tinygrad.engine.realize import BufferXfer, get_runner, ExecItem
from tinygrad.uop.ops import UOp, Ops
from tinygrad.engine.jit import apply_graph_to_jit
BUF_LEN = getenv("BUF_LEN", 128)
cached_prgs = {}
def gen_prg(device, inputs_cnt):
if (device, inputs_cnt) in cached_prgs: return cached_prgs[(device, inputs_cnt)]
with Context(DEBUG=0):
fst = [Tensor.randn(BUF_LEN, dtype=dtypes.int).realize() for i in range(inputs_cnt)]
s = fst[0]
for i in range(1, inputs_cnt): s = s.bitwise_xor(fst[i])
linear = s.schedule_linear()
prg = get_runner(device, linear.src[-1].src[0])
cached_prgs[(device, inputs_cnt)] = prg
return prg
def alloc_rawbuffer(device, fill=False):
rawbuf = Buffer(device, BUF_LEN, dtypes.int).ensure_allocated()
if fill:
with Context(DEBUG=0):
data = np.random.randint(-10000, 10000, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
rawbuf.copy_from(Tensor(data).realize().uop.base.realized)
return rawbuf
def gen_kernel_ji(device, deps):
assert len(deps) >= 2
out = alloc_rawbuffer(device)
prg = gen_prg(device, len(deps))
return ExecItem(UOp(Ops.NOOP), [out] + deps, prg=prg)
def gen_copy_ji(device, deps):
assert len(deps) == 1
out = alloc_rawbuffer(device)
prg = BufferXfer(deps[0].nbytes, device, deps[0].device)
return ExecItem(UOp(Ops.NOOP), [out] + deps, prg=prg)
def gen_graph():
input_buffers = []
all_buffers = []
jis = []
last_n_deps = getenv("LAST_N_DEPS", 0)
kernel_count = random.randint(2, getenv("MAX_KERNELS", 128))
for i in range(kernel_count):
target_device_id = random.randint(0, getenv("MAX_DEVICES", 6) - 1)
target_device = f"{Device.DEFAULT}:{target_device_id}"
is_copy = random.randint(0, 10) < 3
if is_copy:
deps_pool = [buf for buf in all_buffers[-last_n_deps:] if buf.device != target_device]
if len(deps_pool) == 0: deps = []
else: deps = random.sample(deps_pool, 1)
else:
deps_pool = [buf for buf in all_buffers[-last_n_deps:] if buf.device == target_device]
deps_count = random.randint(0, min(getenv("MAX_DEPS_COUNT", 6), len(deps_pool)))
if deps_count == 0: deps = []
else: deps = random.sample(deps_pool, deps_count)
if len(deps) == 0 or (not is_copy and len(deps) < 2):
buf = alloc_rawbuffer(target_device, fill=True)
input_buffers.append(buf)
all_buffers.append(buf)
elif is_copy:
jis.append(gen_copy_ji(target_device, deps))
all_buffers.append(jis[-1].bufs[0])
else:
jis.append(gen_kernel_ji(target_device, deps))
all_buffers.append(jis[-1].bufs[0])
return jis, all_buffers, input_buffers
def run_jit(jis, all_buffers, input_buffers, var_vals):
with Context(DEBUG=0):
for rawbuf in all_buffers:
if rawbuf in input_buffers: continue
rawbuf.copy_from(Buffer("PYTHON", rawbuf.size, rawbuf.dtype, opaque=memoryview(bytearray(rawbuf.nbytes))))
for ei in jis: ei.run(var_vals, jit=True)
with Context(DEBUG=0):
res_buffers = []
for rawbuf in all_buffers: res_buffers.append(rawbuf.as_memoryview())
return res_buffers
def fuzz_graph(jis, all_buffers, input_buffers):
ground_thruth_bufs = run_jit(jis, input_buffers, all_buffers, {})
ground_truth_np = [np.frombuffer(x, _to_np_dtype(all_buffers[i].dtype)) for i,x in enumerate(ground_thruth_bufs)]
for _ in range(getenv("FUZZ_GRAPH_SPLIT_RUNS", 64)):
max_split_points = len(jis) // 3
split_points = random.randint(0, min(max_split_points, getenv("FUZZ_GRAPH_MAX_SPLITS", 8)))
split = [0]
for i in range(split_points - 1):
split.append(random.randint(split[-1] + 2, len(jis) - 2 * (max_split_points - i)))
split.append(len(jis))
graphed_jit = []
for sp in range(len(split)-1):
graphed_jit += apply_graph_to_jit(jis[split[sp]:split[sp+1]], [], {})
for _ in range(getenv("FUZZ_GRAPH_SPLIT_RETRY_RUNS", 4)):
test_bufs = run_jit(graphed_jit, input_buffers, all_buffers, {})
test_bufs_np = [np.frombuffer(x, _to_np_dtype(all_buffers[i].dtype)) for i,x in enumerate(test_bufs)]
for i in range(len(ground_thruth_bufs)): np.testing.assert_equal(ground_truth_np[i], test_bufs_np[i])
if __name__ == "__main__":
SEED = getenv("SEED", 42)
random.seed(SEED)
np.random.seed(SEED)
next_graph_id = 0
for i in range(getenv("ITERS", 1000)):
print("Running graph", next_graph_id)
jis, all_buffers, input_buffers = gen_graph()
fuzz_graph(jis, all_buffers, input_buffers)
next_graph_id += 1

28
artifacts/package_sources/tinygrad/test/external/fuzz_kfd.py vendored Normal file
View File

@@ -0,0 +1,28 @@
#!/usr/bin/env python3
import random
from tqdm import trange
from typing import List
from tinygrad import Device
from tinygrad.runtime.ops_amd import AMDDevice, HWQueue
if __name__ == "__main__":
dev: List[AMDDevice] = [Device[f"KFD:{i}"] for i in range(6)]
print(f"got {len(dev)} devices")
buffers = [(rd:=random.choice(dev), rd.allocator.alloc(random.randint(1, 10000))) for i in range(100)]
for _ in trange(100000):
d1, b1 = random.choice(buffers)
d2, b2 = random.choice(buffers)
d1._gpu_map(b2)
q = HWQueue()
q.signal(sig:=AMDDevice._alloc_signal(10))
qc = HWQueue()
qc.wait(sig)
qc.copy(b1, b2, min(b1.size, b2.size))
d1.completion_signal.value = 1
qc.signal(d1.completion_signal)
qc.submit(d1)
q.wait(d1.completion_signal)
q.submit(d1)
AMDDevice._wait_on(d1.completion_signal.event_id)

87
artifacts/package_sources/tinygrad/test/external/fuzz_shape_ops.py vendored Normal file
View File

@@ -0,0 +1,87 @@
from __future__ import annotations
import unittest
from math import prod
from hypothesis import assume, given, settings, strategies as st
from hypothesis.extra import numpy as stn
import numpy as np
import torch
from tinygrad import Tensor
from tinygrad.helpers import getenv
settings.register_profile(__file__, settings.default,
max_examples=100, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
# torch wraparound for large numbers
st_int32 = st.integers(-2147483648, 2147483647)
@st.composite
def st_shape(draw) -> tuple[int, ...]:
s = draw(stn.array_shapes(min_dims=0, max_dims=6,
min_side=0, max_side=128))
assume(prod(s) <= 1024 ** 2)
assume(prod([d for d in s if d]) <= 1024 ** 4)
return s
def tensors_for_shape(s:tuple[int, ...]) -> tuple[torch.tensor, Tensor]:
x = np.arange(prod(s)).reshape(s)
return torch.from_numpy(x), Tensor(x)
def apply(tor, ten, tor_fn, ten_fn=None):
ok = True
try: tor = tor_fn(tor)
except: tor, ok = None, not ok # noqa: E722
try: ten = ten_fn(ten) if ten_fn is not None else tor_fn(ten)
except: ten, ok = None, not ok # noqa: E722
return tor, ten, ok
class TestShapeOps(unittest.TestCase):
@settings.get_profile(__file__)
@given(st_shape(), st_int32, st.one_of(st_int32, st.lists(st_int32)))
def test_split(self, s:tuple[int, ...], dim:int, sizes:int|list[int]):
tor, ten = tensors_for_shape(s)
tor, ten, ok = apply(tor, ten, lambda t: t.split(sizes, dim))
assert ok
if tor is None and ten is None: return
assert len(tor) == len(ten)
assert all([np.array_equal(tor.numpy(), ten.numpy()) for (tor, ten) in zip(tor, ten)])
@settings.get_profile(__file__)
@given(st_shape(), st_int32, st_int32)
def test_chunk(self, s:tuple[int, ...], dim:int, num:int):
# chunking on a 0 dim is cloning and leads to OOM if done unbounded.
assume((0 <= (actual_dim := len(s)-dim if dim < 0 else dim) < len(s) and s[actual_dim] > 0) or
(num < 16))
tor, ten = tensors_for_shape(s)
tor, ten, ok = apply(tor, ten, lambda t: t.chunk(num, dim))
assert ok
if tor is None and ten is None: return
assert len(tor) == len(ten)
assert all([np.array_equal(tor.numpy(), ten.numpy()) for (tor, ten) in zip(tor, ten)])
@settings.get_profile(__file__)
@given(st_shape(), st_int32)
def test_squeeze(self, s:tuple[int, ...], dim:int):
tor, ten = tensors_for_shape(s)
tor, ten, ok = apply(tor, ten, lambda t: t.squeeze(dim))
assert ok
if tor is None and ten is None: return
assert np.array_equal(tor.numpy(), ten.numpy())
@settings.get_profile(__file__)
@given(st_shape(), st_int32)
def test_unsqueeze(self, s:tuple[int, ...], dim:int):
tor, ten = tensors_for_shape(s)
tor, ten, ok = apply(tor, ten, lambda t: t.unsqueeze(dim))
assert ok
if tor is None and ten is None: return
assert np.array_equal(tor.numpy(), ten.numpy())
if __name__ == '__main__':
unittest.main()

92
artifacts/package_sources/tinygrad/test/external/fuzz_symbolic.py vendored Normal file
View File

@@ -0,0 +1,92 @@
# NOTE: z3-solver 4.15.4 segfaults (exit code 139) when creating many z3.Context() with complex expressions.
# Reproduces consistently with seed=74 around iteration 1767. Versions <=4.15.3 are fine.
# Workaround: reuse a single z3.Context, or pin z3-solver<4.15.4 (see pyproject.toml).
# To repro: pip install z3-solver==4.15.4.0 && python test/external/fuzz_symbolic.py 74
import random, operator, sys
import z3
from tinygrad import Variable, dtypes
from tinygrad.uop.ops import UOp
from tinygrad.uop.validate import uops_to_z3
from tinygrad.helpers import DEBUG
seed = int(sys.argv[1]) if len(sys.argv) > 1 else random.randint(0, 100)
print(f"Seed: {seed}", flush=True)
random.seed(seed)
unary_ops = [lambda a:a+random.randint(-4, 4), lambda a: a*random.randint(-4, 4),
lambda a: a//random.choice(list(range(-9, 0)) + list(range(1, 10))), lambda a: a%random.randint(1, 9),
lambda a:a.maximum(random.randint(-10, 10)), lambda a:a.minimum(random.randint(-10, 10))]
binary_ops = [lambda a,b: a+b, lambda a,b: a*b, lambda a,b:a.maximum(b), lambda a,b:a.minimum(b)]
comp_ops = [operator.lt, operator.le, operator.gt, operator.ge]
def random_or_sub_expression_int(depth, expr):
sub_expr = random.choice([e for e in expr.toposort() if e.dtype not in (dtypes.bool, dtypes.void)])
return random.choice([random_int_expr(depth-1), sub_expr])
def random_int_expr(depth=10):
if depth <= 0: return random.choice(v)
expr1 = random_int_expr(depth-1)
# we give more weight to arithmatic ops than to minimum and maximum
ops = [
lambda: random.choices(unary_ops, weights=[4, 4, 4, 4, 1, 1])[0](expr1),
# for the second operand its either another random exprssion or some subexpression of the first operand
lambda: random.choices(binary_ops, [8, 1, 1, 1])[0](expr1, random_or_sub_expression_int(depth-1, expr1)),
lambda: random_bool_expr(3, random_or_sub_expression_int(depth-1, expr1)).where(expr1, random_or_sub_expression_int(depth-1, expr1)),
]
# we give weight proportional to the amount of ops in each branch
return random.choices(ops, weights=[6, 4, 1])[0]()
def random_bool_expr(depth=10, expr1=None):
if depth == 0: return True
if expr1 is None: expr1 = random_int_expr(depth-1)
expr2 = random.choice([random_or_sub_expression_int(depth-1, expr1), UOp.const(random.randint(-10, 10))])
return random.choice(comp_ops)(expr1, expr2)
if __name__ == "__main__":
skipped = 0
for i in range(10000):
if i % 1000 == 0:
print(f"Running test {i}")
upper_bounds = [*list(range(1, 10)), 16, 32, 64, 128, 256]
u1 = Variable("v1", 0, random.choice(upper_bounds))
u2 = Variable("v2", 0, random.choice(upper_bounds))
u3 = Variable("v3", 0, random.choice(upper_bounds))
v = [u1,u2,u3]
expr = random_int_expr(6)
simplified_expr = expr.simplify()
solver = z3.Solver(ctx=z3.Context())
solver.set(timeout=5000) # some expressions take very long verify, but its very unlikely they actually return sat
z3_expr, z3_simplified_expr, v1, v2, v3 = uops_to_z3(solver, expr, simplified_expr, u1, u2, u3)
check = solver.check(z3_simplified_expr != z3_expr)
if check == z3.unknown and DEBUG>=1:
skipped += 1
print("Skipped due to timeout or interrupt:\n" +
f"v1=Variable(\"{u1.arg[0]}\", {u1.arg[1]}, {u1.arg[2]})\n" +
f"v2=Variable(\"{u2.arg[0]}\", {u2.arg[1]}, {u2.arg[2]})\n" +
f"v3=Variable(\"{u3.arg[0]}\", {u3.arg[1]}, {u3.arg[2]})\n" +
f"expr = {expr.render(simplify=False)}\n")
elif check == z3.sat:
m = solver.model()
n1, n2, n3 = m[v1], m[v2], m[v3]
u1_val, u2_val, u3_val = u1.const_like(n1.as_long()), u2.const_like(n2.as_long()), u3.const_like(n3.as_long())
num = expr.simplify().substitute({u1:u1_val, u2:u2_val, u3:u3_val}).ssimplify()
rn = expr.substitute({u1:u1_val, u2:u2_val, u3:u3_val}).ssimplify()
if num==rn: print("z3 found a mismatch but the expressions are equal!!")
assert False, f"mismatched {expr.render()} at v1={m[v1]}; v2={m[v2]}; v3={m[v3]} = {num} != {rn}\n" +\
"Reproduce with:\n" +\
f"v1=Variable(\"{u1.arg[0]}\", {u1.arg[1]}, {u1.arg[2]})\n" +\
f"v2=Variable(\"{u2.arg[0]}\", {u2.arg[1]}, {u2.arg[2]})\n" +\
f"v3=Variable(\"{u3.arg[0]}\", {u3.arg[1]}, {u3.arg[2]})\n" +\
f"expr = {expr}\n" +\
f"v1_val, v2_val, v3_val = UOp.const({n1.as_long()}), UOp.const({n2.as_long()})," +\
f"UOp.const({n3.as_long()})\n" +\
"num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()\n" +\
"rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()\n" +\
"assert num==rn, f\"{num} != {rn}\"\n"
if DEBUG >= 2: print(f"validated {expr.render()}")
print(f"Skipped {skipped} expressions due to timeout")

65
artifacts/package_sources/tinygrad/test/external/fuzz_symbolic_div.py vendored Normal file
View File

@@ -0,0 +1,65 @@
import random, sys
import z3
from tinygrad.uop.ops import UOp, Ops
from tinygrad.uop.validate import uops_to_z3
from tinygrad.helpers import DEBUG, colored
seed = int(sys.argv[1]) if len(sys.argv) > 1 else random.randint(0, 100)
print(f"Seed: {seed}", flush=True)
random.seed(seed)
def get_random_term(ranges, factors):
# 10% chance of nesting
if random.randint(0,9) == 0: return get_random_expr(ranges, factors)
return random.choice(ranges)*random.choice(factors)*random.choice([1, 1, 1, -1])
def get_random_expr(ranges, factors):
num_terms = random.randint(2,4)
x = UOp.usum(*[get_random_term(ranges, factors) for _ in range(num_terms)])
return x.alu(random.choice([Ops.CDIV, Ops.CMOD]), x.ufix(random.choice(factors)*random.choice([1, 1, 1, -1])))
if __name__ == "__main__":
skipped = 0
for i in range(700):
if i % 100 == 0:
print(f"Running test {i}")
upper_bounds = [*list(range(1, 4)), 16, 33, 53, 64, 256]
variable_names = ["i", "j", "k"]
variables = [UOp.variable(s, 1, random.choice(upper_bounds)) for s in variable_names]
factors = variables+upper_bounds
# add some products
for _ in range(2): factors.append(random.choice(variables)*random.choice(variables))
# add some adds
for _ in range(2): factors.append(random.choice(variables)+random.choice(factors))
num_ranges = 4
ranges = [UOp.range(random.choice(factors), i) for i in range(num_ranges)]
variable_names += [f"r{i}" for i in range(num_ranges)]
expr = get_random_expr(ranges, factors)
simplified_expr = expr.simplify()
if DEBUG>=1:
print(expr.render(simplify=False), " --> ", simplified_expr.render(simplify=False))
solver = z3.Solver()
solver.set(timeout=1000) # some expressions take very long verify, but its very unlikely they actually return sat
z3_expr, z3_simplified_expr, *z3_vars = uops_to_z3(solver, expr, simplified_expr, *variables, *ranges)
check = solver.check(z3_simplified_expr != z3_expr)
if check == z3.unknown:
skipped += 1
if DEBUG>=1: print("skipped z3 verification due to timeout")
elif check == z3.sat:
print(colored("simplify INCORRECT!", "red"))
print(solver.model())
var_vals = {s:solver.model()[z] for s,z in zip(variable_names, z3_vars)}
print("reproduce with:")
print("var_vals = ", var_vals)
print("globals = var_vals|{'cdiv':cdiv,'cmod':cmod}")
print("expr = ast.simplify()")
print("assert eval(ast.render(pm=renderer_infer, simplify=False),globals) == eval(expr.render(pm=renderer_infer, simplify=False),globals)")
print()
assert False
if DEBUG >= 2: print(f"validated {expr.render()}")
print(f"Skipped {skipped} expressions due to timeout")

102
artifacts/package_sources/tinygrad/test/external/mlperf_bert/external_benchmark_bert.py vendored Normal file
View File

@@ -0,0 +1,102 @@
import unittest, time
from tinygrad import Tensor, TinyJit, GlobalCounters, Device
from tinygrad.helpers import getenv, Context
from tinygrad.nn.optim import LAMB
from tinygrad.nn.state import get_parameters
from tinygrad.engine.realize import run_linear
from extra.models import bert
bs = getenv("BS", 16)
seq_len = getenv("SEQ_LEN", 512)
class BenchmarkBertTrain(unittest.TestCase):
def _get_layer(self, layer_id):
if not hasattr(self, "model"):
dropout_prob = 0.0 if getenv("DISABLE_DROPOUT") else 0.1
self.model = bert.BertForPretraining(attention_probs_dropout_prob=dropout_prob, hidden_dropout_prob=dropout_prob)
hidden_size = self.model.bert.embeddings.word_embeddings.embed_sz
intermediate_size = self.model.bert.encoder.layer[0].intermediate.dense.weight.shape[0]
layer_map = {
"embedding": self.model.bert.embeddings,
"attention_self": self.model.bert.encoder.layer[0].attention.self,
"attention_output": self.model.bert.encoder.layer[0].attention.output,
"intermediate": self.model.bert.encoder.layer[0].intermediate,
"output": self.model.bert.encoder.layer[0].output
}
input_shapes = {
"embedding": [(bs, seq_len), (bs, seq_len)],
"attention_self": [(bs, seq_len, hidden_size), (bs, 1, 1, seq_len)],
"attention_output": [(bs, seq_len, hidden_size), (bs, seq_len, 1)],
"intermediate": [(bs, seq_len, hidden_size)],
"output": [(bs, seq_len, intermediate_size), (bs, seq_len, 1)]
}.get(layer_id)
return f"{layer_id}-layer, Input: {input_shapes}", layer_map.get(layer_id), input_shapes
def _test_layer(self, name, layer, input_shapes):
optim = LAMB(get_parameters(layer))
with Context(TRACK_MATCH_STATS=0): Tensor.realize(*[t.assign(t.detach().contiguous()) for t in get_parameters(optim)])
JITCNT = getenv("JITCNT", 1)
@TinyJit
@Context(TRAINING=1)
def step(inputs):
optim.zero_grad()
for i in inputs: i.grad = None
y = layer(*inputs).contiguous().contiguous_backward()
y.sum().backward()
if getenv("ASSIGN", 1): linear, var_vals = Tensor.linear_with_vars(y, *list(inputs), *optim.schedule_step())
else: linear, var_vals = Tensor.linear_with_vars(y, *list(inputs), *[t.grad for t in optim.params])
for _ in range(JITCNT):
run_linear(linear, var_vals)
CNT = getenv("CNT", 5)
best_tm = None
flops, mem_used, mem, kernels = None, None, None, None
for _ in range(CNT):
with Context(TRACK_MATCH_STATS=0): inputs = [Tensor.randn(*shape).realize() for shape in input_shapes]
GlobalCounters.reset()
st = time.perf_counter()
step(inputs)
Device[Device.DEFAULT].synchronize()
et = time.perf_counter()
flops = GlobalCounters.global_ops / JITCNT
mem_used = GlobalCounters.mem_used
mem = GlobalCounters.global_mem / JITCNT
if kernels is None: kernels = GlobalCounters.kernel_count // JITCNT
tm = (et-st) / JITCNT
if best_tm is None or tm < best_tm: best_tm = tm
print(f"\r{name:70s}: {best_tm * 1000:>9.2f} ms, {flops / 10**12 / best_tm:>6.2f} TFLOPS, {mem / 10**9 / best_tm:>5.0f} GB/s, "
f"{mem_used / 10**9: 6.2f} GB used, {kernels:>5d} kernels")
return best_tm, flops, mem, kernels
def test_embedding_layer(self): self._est(*self._test_layer(*self._get_layer("embedding")), 1)
def test_attention_self_layer(self): self._est(*self._test_layer(*self._get_layer("attention_self")), 24) # Assumes BERT-large
def test_attention_output_layer(self): self._est(*self._test_layer(*self._get_layer("attention_output")), 24)
def test_intermediate_layer(self): self._est(*self._test_layer(*self._get_layer("intermediate")), 24)
def test_output_layer(self): self._est(*self._test_layer(*self._get_layer("output")), 24)
est_tm, est_flops, est_mem, est_kernels = 0, 0, 0, 0
@classmethod
def _est(cls, tm, flops, mem, kernels, mult):
cls.est_tm += tm * mult
cls.est_flops += flops * mult
cls.est_mem += mem * mult
cls.est_kernels += kernels * mult
@classmethod
def tearDownClass(cls):
print(f"\restimated step tm: {cls.est_tm * 1000.0:.2f} ms, {cls.est_flops / 10 ** 12 / cls.est_tm:.3f} tflops, "
f"{cls.est_mem / 10 ** 9 / cls.est_tm:.2f} GB/s, {cls.est_kernels} kernels")
if __name__ == "__main__":
unittest.main()

53
artifacts/package_sources/tinygrad/test/external/mlperf_bert/external_test_checkpoint_loading.py vendored Normal file
View File

@@ -0,0 +1,53 @@
# Test whether pretrained weights from first BERT pretraining phase have been loaded correctly
# Usage:
# 1. Download the BERT checkoints with `wikipedia_download.py`
# Command: BASEDIR=/path/to/wiki python3 wikipedia_download.py
# 2. Run this script. (Adjust EVAL_BS and GPUS as needed)
# Command: EVAL_BEAM=4 DEFAULT_FLOAT=half GPUS=6 BASEDIR=/path/to/wiki python3 test/external/mlperf_bert/external_test_checkpoint_loading.py
import os
from tqdm import tqdm
from tinygrad.tensor import Tensor
from tinygrad.device import Device
from tinygrad.helpers import getenv
from tinygrad.nn.state import get_state_dict
from examples.mlperf.helpers import get_mlperf_bert_model, get_data_bert
from examples.mlperf.dataloader import batch_load_val_bert
from examples.mlperf.model_train import eval_step_bert
if __name__ == "__main__":
BASEDIR = os.environ["BASEDIR"] = getenv("BASEDIR", "/raid/datasets/wiki")
INIT_CKPT_DIR = getenv("INIT_CKPT_DIR", BASEDIR)
GPUS = [f"{Device.DEFAULT}:{i}" for i in range(getenv("GPUS", 1))]
EVAL_BS = getenv("EVAL_BS", 4 * len(GPUS))
max_eval_steps = (10000 + EVAL_BS - 1) // EVAL_BS
for i in range(10):
assert os.path.exists(os.path.join(BASEDIR, "eval", f"{i}.pkl")), \
f"File {i}.pkl does not exist in {os.path.join(BASEDIR, 'eval')}"
required_files = ["checkpoint", "model.ckpt-28252.data-00000-of-00001", "model.ckpt-28252.index", "model.ckpt-28252.meta"]
assert all(os.path.exists(os.path.join(INIT_CKPT_DIR, f)) for f in required_files), \
f"Missing checkpoint files in INIT_CKPT_DIR: {required_files}"
model = get_mlperf_bert_model(INIT_CKPT_DIR)
for _, x in get_state_dict(model).items():
x.realize().to_(GPUS)
eval_accuracy = []
eval_it = iter(batch_load_val_bert(EVAL_BS))
for _ in tqdm(range(max_eval_steps), desc="Evaluating", total=max_eval_steps):
eval_data = get_data_bert(GPUS, eval_it)
eval_result: dict[str, Tensor] = eval_step_bert(model, eval_data["input_ids"], eval_data["segment_ids"], eval_data["input_mask"], \
eval_data["masked_lm_positions"], eval_data["masked_lm_ids"], \
eval_data["masked_lm_weights"], eval_data["next_sentence_labels"])
mlm_accuracy = eval_result["masked_lm_accuracy"].numpy().item()
eval_accuracy.append(mlm_accuracy)
total_lm_accuracy = sum(eval_accuracy) / len(eval_accuracy)
assert total_lm_accuracy >= 0.34, "Checkpoint loaded incorrectly. Accuracy should be very close to 0.34085 as per MLPerf BERT README."
print(f"Checkpoint loaded correctly. Accuracy of {total_lm_accuracy*100:.3f}% achieved. (Reference: 34.085%)")

435
artifacts/package_sources/tinygrad/test/external/mlperf_bert/preprocessing/create_pretraining_data.py vendored Normal file
View File

@@ -0,0 +1,435 @@
# https://github.com/mlcommons/training/blob/1c8a098ae3e70962a4f7422c0b0bd35ae639e357/language_model/tensorflow/bert/cleanup_scripts/create_pretraining_data.py
# NOTE: This is a direct copy of the original script
# NOTE: With python 3.7.12, pip install tensorflow=1.15.5
"""Create masked LM/next sentence masked_lm TF examples for BERT."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
os.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python' # NOTE: This is a workaround for protobuf issue
import collections
import random
import tokenization
import tensorflow as tf
flags = tf.flags
FLAGS = flags.FLAGS
flags.DEFINE_string("input_file", None,
"Input raw text file (or comma-separated list of files).")
flags.DEFINE_string(
"output_file", None,
"Output TF example file (or comma-separated list of files).")
flags.DEFINE_string("vocab_file", None,
"The vocabulary file that the BERT model was trained on.")
flags.DEFINE_bool(
"do_lower_case", True,
"Whether to lower case the input text. Should be True for uncased "
"models and False for cased models.")
flags.DEFINE_integer("max_seq_length", 128, "Maximum sequence length.")
flags.DEFINE_integer("max_predictions_per_seq", 20,
"Maximum number of masked LM predictions per sequence.")
flags.DEFINE_integer("random_seed", 12345, "Random seed for data generation.")
flags.DEFINE_integer(
"dupe_factor", 10,
"Number of times to duplicate the input data (with different masks).")
flags.DEFINE_float("masked_lm_prob", 0.15, "Masked LM probability.")
flags.DEFINE_float(
"short_seq_prob", 0.1,
"Probability of creating sequences which are shorter than the "
"maximum length.")
class TrainingInstance(object):
"""A single training instance (sentence pair)."""
def __init__(self, tokens, segment_ids, masked_lm_positions, masked_lm_labels,
is_random_next):
self.tokens = tokens
self.segment_ids = segment_ids
self.is_random_next = is_random_next
self.masked_lm_positions = masked_lm_positions
self.masked_lm_labels = masked_lm_labels
def __str__(self):
s = ""
s += "tokens: %s\n" % (" ".join(
[tokenization.printable_text(x) for x in self.tokens]))
s += "segment_ids: %s\n" % (" ".join([str(x) for x in self.segment_ids]))
s += "is_random_next: %s\n" % self.is_random_next
s += "masked_lm_positions: %s\n" % (" ".join(
[str(x) for x in self.masked_lm_positions]))
s += "masked_lm_labels: %s\n" % (" ".join(
[tokenization.printable_text(x) for x in self.masked_lm_labels]))
s += "\n"
return s
def __repr__(self):
return self.__str__()
def write_instance_to_example_files(instances, tokenizer, max_seq_length,
max_predictions_per_seq, output_files):
"""Create TF example files from `TrainingInstance`s."""
writers = []
for output_file in output_files:
writers.append(tf.python_io.TFRecordWriter(output_file))
writer_index = 0
total_written = 0
for (inst_index, instance) in enumerate(instances):
input_ids = tokenizer.convert_tokens_to_ids(instance.tokens)
input_mask = [1] * len(input_ids)
segment_ids = list(instance.segment_ids)
assert len(input_ids) <= max_seq_length
while len(input_ids) < max_seq_length:
input_ids.append(0)
input_mask.append(0)
segment_ids.append(0)
assert len(input_ids) == max_seq_length
assert len(input_mask) == max_seq_length
assert len(segment_ids) == max_seq_length
masked_lm_positions = list(instance.masked_lm_positions)
masked_lm_ids = tokenizer.convert_tokens_to_ids(instance.masked_lm_labels)
masked_lm_weights = [1.0] * len(masked_lm_ids)
while len(masked_lm_positions) < max_predictions_per_seq:
masked_lm_positions.append(0)
masked_lm_ids.append(0)
masked_lm_weights.append(0.0)
next_sentence_label = 1 if instance.is_random_next else 0
features = collections.OrderedDict()
features["input_ids"] = create_int_feature(input_ids)
features["input_mask"] = create_int_feature(input_mask)
features["segment_ids"] = create_int_feature(segment_ids)
features["masked_lm_positions"] = create_int_feature(masked_lm_positions)
features["masked_lm_ids"] = create_int_feature(masked_lm_ids)
features["masked_lm_weights"] = create_float_feature(masked_lm_weights)
features["next_sentence_labels"] = create_int_feature([next_sentence_label])
tf_example = tf.train.Example(features=tf.train.Features(feature=features))
writers[writer_index].write(tf_example.SerializeToString())
writer_index = (writer_index + 1) % len(writers)
total_written += 1
if inst_index < 20:
tf.logging.info("*** Example ***")
tf.logging.info("tokens: %s" % " ".join(
[tokenization.printable_text(x) for x in instance.tokens]))
for feature_name in features.keys():
feature = features[feature_name]
values = []
if feature.int64_list.value:
values = feature.int64_list.value
elif feature.float_list.value:
values = feature.float_list.value
tf.logging.info(
"%s: %s" % (feature_name, " ".join([str(x) for x in values])))
for writer in writers:
writer.close()
tf.logging.info("Wrote %d total instances", total_written)
def create_int_feature(values):
feature = tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))
return feature
def create_float_feature(values):
feature = tf.train.Feature(float_list=tf.train.FloatList(value=list(values)))
return feature
def create_training_instances(input_files, tokenizer, max_seq_length,
dupe_factor, short_seq_prob, masked_lm_prob,
max_predictions_per_seq, rng):
"""Create `TrainingInstance`s from raw text."""
all_documents = [[]]
# Input file format:
# (1) One sentence per line. These should ideally be actual sentences, not
# entire paragraphs or arbitrary spans of text. (Because we use the
# sentence boundaries for the "next sentence prediction" task).
# (2) Blank lines between documents. Document boundaries are needed so
# that the "next sentence prediction" task doesn't span between documents.
for input_file in input_files:
with tf.gfile.GFile(input_file, "r") as reader:
while True:
line = tokenization.convert_to_unicode(reader.readline())
if not line:
break
line = line.strip()
# Empty lines are used as document delimiters
if not line:
all_documents.append([])
tokens = tokenizer.tokenize(line)
if tokens:
all_documents[-1].append(tokens)
# Remove empty documents
all_documents = [x for x in all_documents if x]
rng.shuffle(all_documents)
vocab_words = list(tokenizer.vocab.keys())
instances = []
for _ in range(dupe_factor):
for document_index in range(len(all_documents)):
instances.extend(
create_instances_from_document(
all_documents, document_index, max_seq_length, short_seq_prob,
masked_lm_prob, max_predictions_per_seq, vocab_words, rng))
rng.shuffle(instances)
return instances
def create_instances_from_document(
all_documents, document_index, max_seq_length, short_seq_prob,
masked_lm_prob, max_predictions_per_seq, vocab_words, rng):
"""Creates `TrainingInstance`s for a single document."""
document = all_documents[document_index]
# Account for [CLS], [SEP], [SEP]
max_num_tokens = max_seq_length - 3
# We *usually* want to fill up the entire sequence since we are padding
# to `max_seq_length` anyways, so short sequences are generally wasted
# computation. However, we *sometimes*
# (i.e., short_seq_prob == 0.1 == 10% of the time) want to use shorter
# sequences to minimize the mismatch between pre-training and fine-tuning.
# The `target_seq_length` is just a rough target however, whereas
# `max_seq_length` is a hard limit.
target_seq_length = max_num_tokens
if rng.random() < short_seq_prob:
target_seq_length = rng.randint(2, max_num_tokens)
# We DON'T just concatenate all of the tokens from a document into a long
# sequence and choose an arbitrary split point because this would make the
# next sentence prediction task too easy. Instead, we split the input into
# segments "A" and "B" based on the actual "sentences" provided by the user
# input.
instances = []
current_chunk = []
current_length = 0
i = 0
while i < len(document):
segment = document[i]
current_chunk.append(segment)
current_length += len(segment)
if i == len(document) - 1 or current_length >= target_seq_length:
if current_chunk:
# `a_end` is how many segments from `current_chunk` go into the `A`
# (first) sentence.
a_end = 1
if len(current_chunk) >= 2:
a_end = rng.randint(1, len(current_chunk) - 1)
tokens_a = []
for j in range(a_end):
tokens_a.extend(current_chunk[j])
tokens_b = []
# Random next
is_random_next = False
if len(current_chunk) == 1 or rng.random() < 0.5:
is_random_next = True
target_b_length = target_seq_length - len(tokens_a)
# This should rarely go for more than one iteration for large
# corpora. However, just to be careful, we try to make sure that
# the random document is not the same as the document
# we're processing.
for _ in range(10):
random_document_index = rng.randint(0, len(all_documents) - 1)
if random_document_index != document_index:
break
random_document = all_documents[random_document_index]
random_start = rng.randint(0, len(random_document) - 1)
for j in range(random_start, len(random_document)):
tokens_b.extend(random_document[j])
if len(tokens_b) >= target_b_length:
break
# We didn't actually use these segments so we "put them back" so
# they don't go to waste.
num_unused_segments = len(current_chunk) - a_end
i -= num_unused_segments
# Actual next
else:
is_random_next = False
for j in range(a_end, len(current_chunk)):
tokens_b.extend(current_chunk[j])
truncate_seq_pair(tokens_a, tokens_b, max_num_tokens, rng)
assert len(tokens_a) >= 1
assert len(tokens_b) >= 1
tokens = []
segment_ids = []
tokens.append("[CLS]")
segment_ids.append(0)
for token in tokens_a:
tokens.append(token)
segment_ids.append(0)
tokens.append("[SEP]")
segment_ids.append(0)
for token in tokens_b:
tokens.append(token)
segment_ids.append(1)
tokens.append("[SEP]")
segment_ids.append(1)
(tokens, masked_lm_positions,
masked_lm_labels) = create_masked_lm_predictions(
tokens, masked_lm_prob, max_predictions_per_seq, vocab_words, rng)
instance = TrainingInstance(
tokens=tokens,
segment_ids=segment_ids,
is_random_next=is_random_next,
masked_lm_positions=masked_lm_positions,
masked_lm_labels=masked_lm_labels)
instances.append(instance)
current_chunk = []
current_length = 0
i += 1
return instances
MaskedLmInstance = collections.namedtuple("MaskedLmInstance",
["index", "label"])
def create_masked_lm_predictions(tokens, masked_lm_prob,
max_predictions_per_seq, vocab_words, rng):
"""Creates the predictions for the masked LM objective."""
cand_indexes = []
for (i, token) in enumerate(tokens):
if token == "[CLS]" or token == "[SEP]":
continue
cand_indexes.append(i)
rng.shuffle(cand_indexes)
output_tokens = list(tokens)
num_to_predict = min(max_predictions_per_seq,
max(1, int(round(len(tokens) * masked_lm_prob))))
masked_lms = []
covered_indexes = set()
for index in cand_indexes:
if len(masked_lms) >= num_to_predict:
break
if index in covered_indexes:
continue
covered_indexes.add(index)
masked_token = None
# 80% of the time, replace with [MASK]
if rng.random() < 0.8:
masked_token = "[MASK]"
else:
# 10% of the time, keep original
if rng.random() < 0.5:
masked_token = tokens[index]
# 10% of the time, replace with random word
else:
masked_token = vocab_words[rng.randint(0, len(vocab_words) - 1)]
output_tokens[index] = masked_token
masked_lms.append(MaskedLmInstance(index=index, label=tokens[index]))
masked_lms = sorted(masked_lms, key=lambda x: x.index)
masked_lm_positions = []
masked_lm_labels = []
for p in masked_lms:
masked_lm_positions.append(p.index)
masked_lm_labels.append(p.label)
return (output_tokens, masked_lm_positions, masked_lm_labels)
def truncate_seq_pair(tokens_a, tokens_b, max_num_tokens, rng):
"""Truncates a pair of sequences to a maximum sequence length."""
while True:
total_length = len(tokens_a) + len(tokens_b)
if total_length <= max_num_tokens:
break
trunc_tokens = tokens_a if len(tokens_a) > len(tokens_b) else tokens_b
assert len(trunc_tokens) >= 1
# We want to sometimes truncate from the front and sometimes from the
# back to add more randomness and avoid biases.
if rng.random() < 0.5:
del trunc_tokens[0]
else:
trunc_tokens.pop()
def main(_):
tf.logging.set_verbosity(tf.logging.INFO)
tokenizer = tokenization.FullTokenizer(
vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case)
input_files = []
for input_pattern in FLAGS.input_file.split(","):
input_files.extend(tf.gfile.Glob(input_pattern))
tf.logging.info("*** Reading from input files ***")
for input_file in input_files:
tf.logging.info(" %s", input_file)
rng = random.Random(FLAGS.random_seed)
instances = create_training_instances(
input_files, tokenizer, FLAGS.max_seq_length, FLAGS.dupe_factor,
FLAGS.short_seq_prob, FLAGS.masked_lm_prob, FLAGS.max_predictions_per_seq,
rng)
output_files = FLAGS.output_file.split(",")
tf.logging.info("*** Writing to output files ***")
for output_file in output_files:
tf.logging.info(" %s", output_file)
write_instance_to_example_files(instances, tokenizer, FLAGS.max_seq_length,
FLAGS.max_predictions_per_seq, output_files)
if __name__ == "__main__":
flags.mark_flag_as_required("input_file")
flags.mark_flag_as_required("output_file")
flags.mark_flag_as_required("vocab_file")
tf.app.run()

79
artifacts/package_sources/tinygrad/test/external/mlperf_bert/preprocessing/external_test_preprocessing_part.py vendored Normal file
View File

@@ -0,0 +1,79 @@
# USAGE:
# 1. Download raw text data with `wikipedia_download.py`
# 2. Install python==3.7.12 and tensorflow==1.15.5
# Run `create_pretraining_data.py` to create TFRecords on specific part (This will take some time)
# Command: python3 create_pretraining_data.py --input_file=/path/to/part-00XXX-of-00500 --vocab_file=/path/to/vocab.txt \
# --output_file=/path/to/output.tfrecord --max_seq_length=512 --max_predictions_per_seq=76
#
# 2.1 For eval: --input_file=/path/to/eval.txt and
# Command: python3 pick_eval_samples.py --input_tfrecord=/path/to/eval.tfrecord --output_tfrecord=/path/to/output_eval.tfrecord
# 3. Run `wikipedia.py` to preprocess the data with tinygrad (Use python > 3.7)
# Command: BASEDIR=/path/to/basedir python3 wikipedia.py pre-train X (NOTE: part number needs to match part of step 2)
# This will output to /path/to/basedir/train/X.pkl
#
# 3.1 For eval:
# Command: BASEDIR=/path/to/basedir python3 wikipedia.py pre-eval
# This will output to /path/to/basedir/eval.pkl
# 4. Run this script to verify the correctness of the preprocessing script for specific part
# Command: python3 external_test_preprocessing_part.py --preprocessed_part=/path/to/basedir/train/X.pkl --tf_records=/path/to/output.tfrecord
import os, argparse, pickle
from tqdm import tqdm
# This is a workaround for protobuf issue
os.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
tf.compat.v1.enable_eager_execution()
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
def _parse_function(proto, max_seq_length, max_predictions_per_seq):
feature_description = {
'input_ids': tf.io.FixedLenFeature([max_seq_length], tf.int64),
'input_mask': tf.io.FixedLenFeature([max_seq_length], tf.int64),
'segment_ids': tf.io.FixedLenFeature([max_seq_length], tf.int64),
'masked_lm_positions': tf.io.FixedLenFeature([max_predictions_per_seq], tf.int64),
'masked_lm_ids': tf.io.FixedLenFeature([max_predictions_per_seq], tf.int64),
'masked_lm_weights': tf.io.FixedLenFeature([max_predictions_per_seq], tf.float32),
'next_sentence_labels': tf.io.FixedLenFeature([1], tf.int64),
}
return tf.io.parse_single_example(proto, feature_description)
def load_dataset(file_path, max_seq_length=512, max_predictions_per_seq=76):
dataset = tf.data.TFRecordDataset(file_path)
parse_function = lambda proto: _parse_function(proto, max_seq_length, max_predictions_per_seq) # noqa: E731
return dataset.map(parse_function)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Verify the correctness of the preprocessing script for specific part",
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("--preprocessed_part", type=str, default=None,
help="Path to preprocessed samples file from `wikipedia.py`")
parser.add_argument("--tf_records", type=str, default=None,
help="Path to TFRecords file from `create_pretraining_data.py` (Reference implementation)")
parser.add_argument("--max_seq_length", type=int, default=512, help="Max sequence length. For MLPerf keep it at 512")
parser.add_argument("--max_predictions_per_seq", type=int, default=76, help="Max predictions per sequence. For MLPerf keep it at 76")
parser.add_argument("--is_eval", type=bool, default=False, help="Whether to run eval or train preprocessing")
args = parser.parse_args()
assert os.path.isfile(args.preprocessed_part), f"The specified file {args.preprocessed_part} does not exist."
assert os.path.isfile(args.tf_records), f"The specified TFRecords file {args.tf_records} does not exist."
with open(args.preprocessed_part, 'rb') as f:
preprocessed_samples = pickle.load(f)
dataset = load_dataset(args.tf_records, args.max_seq_length, args.max_predictions_per_seq)
tf_record_count = sum(1 for _ in dataset)
assert tf_record_count == len(preprocessed_samples), f"Samples in reference: {tf_record_count} != Preprocessed samples: {len(preprocessed_samples)}"
print(f"Total samples in the part: {tf_record_count}")
for i, (reference_example, preprocessed_sample) in tqdm(enumerate(zip(dataset, preprocessed_samples)), desc="Checking samples", total=len(preprocessed_samples)): # noqa: E501
feature_keys = ["input_ids", "input_mask", "segment_ids", "masked_lm_positions", "masked_lm_ids", "masked_lm_weights", "next_sentence_labels"]
for key in feature_keys:
reference_example_feature = reference_example[key].numpy()
assert (reference_example_feature == preprocessed_sample[key]).all(), \
f"{key} are not equal at index {i}\nReference: {reference_example_feature}\nPreprocessed: {preprocessed_sample[key]}"

127
artifacts/package_sources/tinygrad/test/external/mlperf_bert/preprocessing/pick_eval_samples.py vendored Normal file
View File

@@ -0,0 +1,127 @@
# https://github.com/mlcommons/training/blob/1c8a098ae3e70962a4f7422c0b0bd35ae639e357/language_model/tensorflow/bert/cleanup_scripts/pick_eval_samples.py
# NOTE: This is a direct copy of the original script
"""Script for picking certain number of sampels.
"""
import argparse
import time
import logging
import collections
import tensorflow as tf
parser = argparse.ArgumentParser(
description="Eval sample picker for BERT.")
parser.add_argument(
'--input_tfrecord',
type=str,
default='',
help='Input tfrecord path')
parser.add_argument(
'--output_tfrecord',
type=str,
default='',
help='Output tfrecord path')
parser.add_argument(
'--num_examples_to_pick',
type=int,
default=10000,
help='Number of examples to pick')
parser.add_argument(
'--max_seq_length',
type=int,
default=512,
help='The maximum number of tokens within a sequence.')
parser.add_argument(
'--max_predictions_per_seq',
type=int,
default=76,
help='The maximum number of predictions within a sequence.')
args = parser.parse_args()
max_seq_length = args.max_seq_length
max_predictions_per_seq = args.max_predictions_per_seq
logging.basicConfig(level=logging.INFO)
def decode_record(record):
"""Decodes a record to a TensorFlow example."""
name_to_features = {
"input_ids":
tf.FixedLenFeature([max_seq_length], tf.int64),
"input_mask":
tf.FixedLenFeature([max_seq_length], tf.int64),
"segment_ids":
tf.FixedLenFeature([max_seq_length], tf.int64),
"masked_lm_positions":
tf.FixedLenFeature([max_predictions_per_seq], tf.int64),
"masked_lm_ids":
tf.FixedLenFeature([max_predictions_per_seq], tf.int64),
"masked_lm_weights":
tf.FixedLenFeature([max_predictions_per_seq], tf.float32),
"next_sentence_labels":
tf.FixedLenFeature([1], tf.int64),
}
example = tf.parse_single_example(record, name_to_features)
# tf.Example only supports tf.int64, but the TPU only supports tf.int32.
# So cast all int64 to int32.
for name in list(example.keys()):
t = example[name]
if t.dtype == tf.int64:
t = tf.to_int32(t)
example[name] = t
return example
def create_int_feature(values):
feature = tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))
return feature
def create_float_feature(values):
feature = tf.train.Feature(float_list=tf.train.FloatList(value=list(values)))
return feature
if __name__ == '__main__':
tic = time.time()
tf.enable_eager_execution()
d = tf.data.TFRecordDataset(args.input_tfrecord)
num_examples = 0
records = []
for record in d:
records.append(record)
num_examples += 1
writer = tf.python_io.TFRecordWriter(args.output_tfrecord)
i = 0
pick_ratio = num_examples / args.num_examples_to_pick
num_examples_picked = 0
for i in range(args.num_examples_to_pick):
example = decode_record(records[int(i * pick_ratio)])
features = collections.OrderedDict()
features["input_ids"] = create_int_feature(
example["input_ids"].numpy())
features["input_mask"] = create_int_feature(
example["input_mask"].numpy())
features["segment_ids"] = create_int_feature(
example["segment_ids"].numpy())
features["masked_lm_positions"] = create_int_feature(
example["masked_lm_positions"].numpy())
features["masked_lm_ids"] = create_int_feature(
example["masked_lm_ids"].numpy())
features["masked_lm_weights"] = create_float_feature(
example["masked_lm_weights"].numpy())
features["next_sentence_labels"] = create_int_feature(
example["next_sentence_labels"].numpy())
tf_example = tf.train.Example(features=tf.train.Features(feature=features))
writer.write(tf_example.SerializeToString())
num_examples_picked += 1
writer.close()
toc = time.time()
logging.info("Picked %d examples out of %d samples in %.2f sec",
num_examples_picked, num_examples, toc - tic)

415
artifacts/package_sources/tinygrad/test/external/mlperf_bert/preprocessing/tokenization.py vendored Normal file
View File

@@ -0,0 +1,415 @@
# https://github.com/mlcommons/training/blob/1c8a098ae3e70962a4f7422c0b0bd35ae639e357/language_model/tensorflow/bert/cleanup_scripts/tokenization.py
# NOTE: This is a direct copy of the original script
"""Tokenization classes."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import collections
import re
import unicodedata
from absl import flags
import six
import tensorflow.compat.v1 as tf
FLAGS = flags.FLAGS
flags.DEFINE_bool(
"preserve_unused_tokens", False,
"If True, Wordpiece tokenization will not be applied to words in the vocab."
)
_UNUSED_TOKEN_RE = re.compile("^\\[unused\\d+\\]$")
def preserve_token(token, vocab):
"""Returns True if the token should forgo tokenization and be preserved."""
if not FLAGS.preserve_unused_tokens:
return False
if token not in vocab:
return False
return bool(_UNUSED_TOKEN_RE.search(token))
def validate_case_matches_checkpoint(do_lower_case, init_checkpoint):
"""Checks whether the casing config is consistent with the checkpoint name."""
# The casing has to be passed in by the user and there is no explicit check
# as to whether it matches the checkpoint. The casing information probably
# should have been stored in the bert_config.json file, but it's not, so
# we have to heuristically detect it to validate.
if not init_checkpoint:
return
m = re.match("^.*?([A-Za-z0-9_-]+)/bert_model.ckpt", init_checkpoint)
if m is None:
return
model_name = m.group(1)
lower_models = [
"uncased_L-24_H-1024_A-16", "uncased_L-12_H-768_A-12",
"multilingual_L-12_H-768_A-12", "chinese_L-12_H-768_A-12"
]
cased_models = [
"cased_L-12_H-768_A-12", "cased_L-24_H-1024_A-16",
"multi_cased_L-12_H-768_A-12"
]
is_bad_config = False
if model_name in lower_models and not do_lower_case:
is_bad_config = True
actual_flag = "False"
case_name = "lowercased"
opposite_flag = "True"
if model_name in cased_models and do_lower_case:
is_bad_config = True
actual_flag = "True"
case_name = "cased"
opposite_flag = "False"
if is_bad_config:
raise ValueError(
"You passed in `--do_lower_case=%s` with `--init_checkpoint=%s`. "
"However, `%s` seems to be a %s model, so you "
"should pass in `--do_lower_case=%s` so that the fine-tuning matches "
"how the model was pre-training. If this error is wrong, please "
"just comment out this check." % (actual_flag, init_checkpoint,
model_name, case_name, opposite_flag))
def convert_to_unicode(text):
"""Converts `text` to Unicode (if it's not already), assuming utf-8 input."""
if six.PY3:
if isinstance(text, str):
return text
elif isinstance(text, bytes):
return text.decode("utf-8", "ignore")
else:
raise ValueError("Unsupported string type: %s" % (type(text)))
elif six.PY2:
if isinstance(text, str):
return text.decode("utf-8", "ignore")
elif isinstance(text, unicode): # noqa: F821
return text
else:
raise ValueError("Unsupported string type: %s" % (type(text)))
else:
raise ValueError("Not running on Python2 or Python 3?")
def printable_text(text):
"""Returns text encoded in a way suitable for print or `tf.logging`."""
# These functions want `str` for both Python2 and Python3, but in one case
# it's a Unicode string and in the other it's a byte string.
if six.PY3:
if isinstance(text, str):
return text
elif isinstance(text, bytes):
return text.decode("utf-8", "ignore")
else:
raise ValueError("Unsupported string type: %s" % (type(text)))
elif six.PY2:
if isinstance(text, str):
return text
elif isinstance(text, unicode): # noqa: F821
return text.encode("utf-8")
else:
raise ValueError("Unsupported string type: %s" % (type(text)))
else:
raise ValueError("Not running on Python2 or Python 3?")
def load_vocab(vocab_file):
"""Loads a vocabulary file into a dictionary."""
vocab = collections.OrderedDict()
with tf.gfile.GFile(vocab_file, "r") as reader:
while True:
token = convert_to_unicode(reader.readline())
if not token:
break
token = token.strip()
if token not in vocab:
vocab[token] = len(vocab)
return vocab
def convert_by_vocab(vocab, items):
"""Converts a sequence of [tokens|ids] using the vocab."""
output = []
for item in items:
output.append(vocab[item])
return output
def convert_tokens_to_ids(vocab, tokens):
return convert_by_vocab(vocab, tokens)
def convert_ids_to_tokens(inv_vocab, ids):
return convert_by_vocab(inv_vocab, ids)
def whitespace_tokenize(text):
"""Runs basic whitespace cleaning and splitting on a piece of text."""
text = text.strip()
if not text:
return []
tokens = text.split()
return tokens
class FullTokenizer(object):
"""Runs end-to-end tokenziation."""
def __init__(self, vocab_file, do_lower_case=True):
self.vocab = load_vocab(vocab_file)
self.inv_vocab = {v: k for k, v in self.vocab.items()}
self.basic_tokenizer = BasicTokenizer(
do_lower_case=do_lower_case, vocab=self.vocab)
self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab)
def tokenize(self, text):
split_tokens = []
for token in self.basic_tokenizer.tokenize(text):
if preserve_token(token, self.vocab):
split_tokens.append(token)
continue
for sub_token in self.wordpiece_tokenizer.tokenize(token):
split_tokens.append(sub_token)
return split_tokens
def convert_tokens_to_ids(self, tokens):
return convert_by_vocab(self.vocab, tokens)
def convert_ids_to_tokens(self, ids):
return convert_by_vocab(self.inv_vocab, ids)
class BasicTokenizer(object):
"""Runs basic tokenization (punctuation splitting, lower casing, etc.)."""
def __init__(self, do_lower_case=True, vocab=tuple()):
"""Constructs a BasicTokenizer.
Args:
do_lower_case: Whether to lower case the input.
vocab: A container of tokens to not mutate during tokenization.
"""
self.do_lower_case = do_lower_case
self.vocab = vocab
def tokenize(self, text):
"""Tokenizes a piece of text."""
text = convert_to_unicode(text)
text = self._clean_text(text)
# This was added on November 1st, 2018 for the multilingual and Chinese
# models. This is also applied to the English models now, but it doesn't
# matter since the English models were not trained on any Chinese data
# and generally don't have any Chinese data in them (there are Chinese
# characters in the vocabulary because Wikipedia does have some Chinese
# words in the English Wikipedia.).
text = self._tokenize_chinese_chars(text)
orig_tokens = whitespace_tokenize(text)
split_tokens = []
for token in orig_tokens:
if preserve_token(token, self.vocab):
split_tokens.append(token)
continue
if self.do_lower_case:
token = token.lower()
token = self._run_strip_accents(token)
split_tokens.extend(self._run_split_on_punc(token))
output_tokens = whitespace_tokenize(" ".join(split_tokens))
return output_tokens
def _run_strip_accents(self, text):
"""Strips accents from a piece of text."""
text = unicodedata.normalize("NFD", text)
output = []
for char in text:
cat = unicodedata.category(char)
if cat == "Mn":
continue
output.append(char)
return "".join(output)
def _run_split_on_punc(self, text):
"""Splits punctuation on a piece of text."""
chars = list(text)
i = 0
start_new_word = True
output = []
while i < len(chars):
char = chars[i]
if _is_punctuation(char):
output.append([char])
start_new_word = True
else:
if start_new_word:
output.append([])
start_new_word = False
output[-1].append(char)
i += 1
return ["".join(x) for x in output]
def _tokenize_chinese_chars(self, text):
"""Adds whitespace around any CJK character."""
output = []
for char in text:
cp = ord(char)
if self._is_chinese_char(cp):
output.append(" ")
output.append(char)
output.append(" ")
else:
output.append(char)
return "".join(output)
def _is_chinese_char(self, cp):
"""Checks whether CP is the codepoint of a CJK character."""
# This defines a "chinese character" as anything in the CJK Unicode block:
# https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
#
# Note that the CJK Unicode block is NOT all Japanese and Korean characters,
# despite its name. The modern Korean Hangul alphabet is a different block,
# as is Japanese Hiragana and Katakana. Those alphabets are used to write
# space-separated words, so they are not treated specially and handled
# like all of the other languages.
if ((cp >= 0x4E00 and cp <= 0x9FFF) or #
(cp >= 0x3400 and cp <= 0x4DBF) or #
(cp >= 0x20000 and cp <= 0x2A6DF) or #
(cp >= 0x2A700 and cp <= 0x2B73F) or #
(cp >= 0x2B740 and cp <= 0x2B81F) or #
(cp >= 0x2B820 and cp <= 0x2CEAF) or
(cp >= 0xF900 and cp <= 0xFAFF) or #
(cp >= 0x2F800 and cp <= 0x2FA1F)): #
return True
return False
def _clean_text(self, text):
"""Performs invalid character removal and whitespace cleanup on text."""
output = []
for char in text:
cp = ord(char)
if cp == 0 or cp == 0xfffd or _is_control(char):
continue
if _is_whitespace(char):
output.append(" ")
else:
output.append(char)
return "".join(output)
class WordpieceTokenizer(object):
"""Runs WordPiece tokenziation."""
def __init__(self, vocab, unk_token="[UNK]", max_input_chars_per_word=200):
self.vocab = vocab
self.unk_token = unk_token
self.max_input_chars_per_word = max_input_chars_per_word
def tokenize(self, text):
"""Tokenizes a piece of text into its word pieces.
This uses a greedy longest-match-first algorithm to perform tokenization
using the given vocabulary.
For example:
input = "unaffable"
output = ["un", "##aff", "##able"]
Args:
text: A single token or whitespace separated tokens. This should have
already been passed through `BasicTokenizer.
Returns:
A list of wordpiece tokens.
"""
text = convert_to_unicode(text)
output_tokens = []
for token in whitespace_tokenize(text):
chars = list(token)
if len(chars) > self.max_input_chars_per_word:
output_tokens.append(self.unk_token)
continue
is_bad = False
start = 0
sub_tokens = []
while start < len(chars):
end = len(chars)
cur_substr = None
while start < end:
substr = "".join(chars[start:end])
if start > 0:
substr = "##" + substr
if substr in self.vocab:
cur_substr = substr
break
end -= 1
if cur_substr is None:
is_bad = True
break
sub_tokens.append(cur_substr)
start = end
if is_bad:
output_tokens.append(self.unk_token)
else:
output_tokens.extend(sub_tokens)
return output_tokens
def _is_whitespace(char):
"""Checks whether `chars` is a whitespace character."""
# \t, \n, and \r are technically control characters but we treat them
# as whitespace since they are generally considered as such.
if char == " " or char == "\t" or char == "\n" or char == "\r":
return True
cat = unicodedata.category(char)
if cat == "Zs":
return True
return False
def _is_control(char):
"""Checks whether `chars` is a control character."""
# These are technically control characters but we count them as whitespace
# characters.
if char == "\t" or char == "\n" or char == "\r":
return False
cat = unicodedata.category(char)
if cat in ("Cc", "Cf"):
return True
return False
def _is_punctuation(char):
"""Checks whether `chars` is a punctuation character."""
cp = ord(char)
# We treat all non-letter/number ASCII as punctuation.
# Characters such as "^", "$", and "`" are not in the Unicode
# Punctuation class but we treat them as punctuation anyways, for
# consistency.
if ((cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64) or
(cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126)):
return True
cat = unicodedata.category(char)
if cat.startswith("P"):
return True
return False

219
artifacts/package_sources/tinygrad/test/external/mlperf_resnet/lars_optimizer.py vendored Normal file
View File

@@ -0,0 +1,219 @@
# https://github.com/mlcommons/training/blob/e3769c8dcf88cd21e1001dd2f894b40a1513ec5d/image_classification/tensorflow2/lars_optimizer.py
# changes: don't call lr_t if it's not a schedule
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Layer-wise Adaptive Rate Scaling optimizer for large-batch training."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import tensorflow as tf
# from tf2_common.training import optimizer_v2modified
from tensorflow.python.framework import ops
from tensorflow.python.keras import backend_config
from tensorflow.python.keras.optimizer_v2 import optimizer_v2
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import linalg_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import state_ops
# class LARSOptimizer(optimizer_v2modified.OptimizerV2Modified):
class LARSOptimizer(optimizer_v2.OptimizerV2):
"""Layer-wise Adaptive Rate Scaling for large batch training.
Introduced by "Large Batch Training of Convolutional Networks" by Y. You,
I. Gitman, and B. Ginsburg. (https://arxiv.org/abs/1708.03888)
Implements the LARS learning rate scheme presented in the paper above. This
optimizer is useful when scaling the batch size to up to 32K without
significant performance degradation. It is recommended to use the optimizer
in conjunction with:
- Gradual learning rate warm-up
- Linear learning rate scaling
- Poly rule learning rate decay
Note, LARS scaling is currently only enabled for dense tensors. Sparse tensors
use the default momentum optimizer.
"""
def __init__(
self,
learning_rate,
momentum=0.9,
weight_decay=0.0001,
# The LARS coefficient is a hyperparameter
eeta=0.001,
epsilon=0.0,
name="LARSOptimizer",
# Enable skipping variables from LARS scaling.
# TODO(sameerkm): Enable a direct mechanism to pass a
# subset of variables to the optimizer.
skip_list=None,
use_nesterov=False,
**kwargs):
"""Construct a new LARS Optimizer.
Args:
learning_rate: A `Tensor`, floating point value, or a schedule that is a
`tf.keras.optimizers.schedules.LearningRateSchedule`, or a callable
that takes no arguments and returns the actual value to use. The
learning rate.
momentum: A floating point value. Momentum hyperparameter.
weight_decay: A floating point value. Weight decay hyperparameter.
eeta: LARS coefficient as used in the paper. Dfault set to LARS
coefficient from the paper. (eeta / weight_decay) determines the highest
scaling factor in LARS.
epsilon: Optional epsilon parameter to be set in models that have very
small gradients. Default set to 0.0.
name: Optional name prefix for variables and ops created by LARSOptimizer.
skip_list: List of strings to enable skipping variables from LARS scaling.
If any of the strings in skip_list is a subset of var.name, variable
'var' is skipped from LARS scaling. For a typical classification model
with batch normalization, the skip_list is ['batch_normalization',
'bias']
use_nesterov: when set to True, nesterov momentum will be enabled
**kwargs: keyword arguments.
Raises:
ValueError: If a hyperparameter is set to a non-sensical value.
"""
if momentum < 0.0:
raise ValueError("momentum should be positive: %s" % momentum)
if weight_decay < 0.0:
raise ValueError("weight_decay should be positive: %s" % weight_decay)
super(LARSOptimizer, self).__init__(name=name, **kwargs)
self._set_hyper("learning_rate", learning_rate)
# When directly using class members, instead of
# _set_hyper and _get_hyper (such as learning_rate above),
# the values are fixed after __init(), and not being
# updated during the training process.
# This provides better performance but less flexibility.
self.momentum = momentum
self.weight_decay = weight_decay
self.eeta = eeta
self.epsilon = epsilon or backend_config.epsilon()
self._skip_list = skip_list
self.use_nesterov = use_nesterov
def _prepare_local(self, var_device, var_dtype, apply_state):
lr_t = self._get_hyper("learning_rate", var_dtype)
local_step = math_ops.cast(self.iterations, var_dtype)
if callable(lr_t): lr_t = math_ops.cast(lr_t(local_step), var_dtype)
learning_rate_t = array_ops.identity(lr_t)
apply_state[(var_device, var_dtype)].update(
dict(
learning_rate=learning_rate_t,
))
def _create_slots(self, var_list):
for v in var_list:
self.add_slot(v, "momentum")
def compute_lr(self, grad, var, coefficients):
scaled_lr = coefficients["learning_rate"]
if self._skip_list is None or not any(v in var.name
for v in self._skip_list):
w_norm = linalg_ops.norm(var, ord=2)
g_norm = linalg_ops.norm(grad, ord=2)
trust_ratio = array_ops.where(
math_ops.greater(w_norm, 0),
array_ops.where(
math_ops.greater(g_norm, 0),
(self.eeta * w_norm /
(g_norm + self.weight_decay * w_norm + self.epsilon)), 1.0), 1.0)
scaled_lr = coefficients["learning_rate"] * trust_ratio
# Add the weight regularization gradient
grad = grad + self.weight_decay * var
return scaled_lr, grad
def _apply_dense(self, grad, var, apply_state=None):
return self._resource_apply_dense(grad, var, apply_state)
def _resource_apply_dense(self, grad, var, apply_state=None):
var_device, var_dtype = var.device, var.dtype.base_dtype
coefficients = ((apply_state or {}).get((var_device, var_dtype))
or self._fallback_apply_state(var_device, var_dtype))
scaled_lr, grad = self.compute_lr(grad, var, coefficients)
mom = self.get_slot(var, "momentum")
# Use ApplyKerasMomentum instead of ApplyMomentum
# training_ops.resource_apply_keras_momentum(
# var.handle,
# mom.handle,
# scaled_lr,
# grad,
# coefficients["momentum"],
# use_locking=False,
# use_nesterov=self.use_nesterov)
mom_t = mom * self.momentum - grad * scaled_lr
mom_t = state_ops.assign(mom, mom_t, use_locking=False)
if self.use_nesterov:
var_t = var + mom_t * self.momentum - grad * scaled_lr
else:
var_t = var + mom_t
return state_ops.assign(var, var_t, use_locking=False).op
# Fallback to momentum optimizer for sparse tensors
def _apply_sparse(self, grad, var, apply_state=None):
var_device, var_dtype = var.device, var.dtype.base_dtype
coefficients = ((apply_state or {}).get((var_device, var_dtype))
or self._fallback_apply_state(var_device, var_dtype))
mom = self.get_slot(var, "momentum")
return tf.raw_ops.SparseApplyMomentum(
var=var,
accum=mom,
lr=coefficients["learning_rate"],
grad=grad.values,
indices=grad.indices,
momentum=self.momentum,
use_locking=False,
use_nesterov=self.use_nesterov)
def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
var_device, var_dtype = var.device, var.dtype.base_dtype
coefficients = ((apply_state or {}).get((var_device, var_dtype))
or self._fallback_apply_state(var_device, var_dtype))
mom = self.get_slot(var, "momentum")
return tf.raw_ops.ResourceSparseApplyKerasMomentum(
var=var.handle,
accum=mom.handle,
lr=coefficients["learning_rate"],
grad=grad,
indices=indices,
momentum=self.momentum,
use_locking=False,
use_nesterov=self.use_nesterov)
def get_config(self):
config = super(LARSOptimizer, self).get_config()
config.update({
"learning_rate": self._serialize_hyperparameter("learning_rate"),
"momentum": self.momentum,
"weight_decay": self.weight_decay,
"eeta": self.eeta,
"epsilon": self.epsilon,
"use_nesterov": self.use_nesterov,
})
return config

179
artifacts/package_sources/tinygrad/test/external/mlperf_resnet/lars_util.py vendored Normal file
View File

@@ -0,0 +1,179 @@
# https://github.com/mlcommons/training/blob/e237206991d10449d9675d95606459a3cb6c21ad/image_classification/tensorflow2/lars_util.py
# changes: commented out logging
# changes: convert_to_tensor_v2 -> convert_to_tensor
# changes: extend from tf.python.keras.optimizer_v2.learning_rate_schedule.LearningRateScheduler
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Enable Layer-wise Adaptive Rate Scaling optimizer in ResNet."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from absl import flags
import tensorflow as tf
#from tf2_common.utils.mlp_log import mlp_log
from tensorflow.python.eager import context
from tensorflow.python.framework import ops
from tensorflow.python.ops import math_ops
from tensorflow.python.keras.optimizer_v2 import learning_rate_schedule
FLAGS = flags.FLAGS
def define_lars_flags():
"""Defines flags needed by LARS optimizer."""
flags.DEFINE_float(
'end_learning_rate', default=None,
help=('Polynomial decay end learning rate.'))
flags.DEFINE_float(
'lars_epsilon', default=0.0,
help=('Override autoselected LARS epsilon.'))
flags.DEFINE_float(
'warmup_epochs', default=None,
help=('Override autoselected polynomial decay warmup epochs.'))
flags.DEFINE_float(
'momentum',
default=0.9,
help=('Momentum parameter used in the MomentumOptimizer.'))
class PolynomialDecayWithWarmup(learning_rate_schedule.LearningRateSchedule):
"""A LearningRateSchedule that uses a polynomial decay with warmup."""
def __init__(
self,
batch_size,
steps_per_epoch,
train_steps,
initial_learning_rate=None,
end_learning_rate=None,
warmup_epochs=None,
compute_lr_on_cpu=False,
name=None):
"""Applies a polynomial decay to the learning rate with warmup."""
super(PolynomialDecayWithWarmup, self).__init__()
self.batch_size = batch_size
self.steps_per_epoch = steps_per_epoch
self.train_steps = train_steps
self.name = name
self.learning_rate_ops_cache = {}
self.compute_lr_on_cpu = compute_lr_on_cpu
if batch_size < 16384:
self.initial_learning_rate = 10.0
warmup_epochs_ = 5
elif batch_size < 32768:
self.initial_learning_rate = 25.0
warmup_epochs_ = 5
else:
self.initial_learning_rate = 31.2
warmup_epochs_ = 25
# Override default poly learning rate and warmup epochs
if initial_learning_rate:
self.initial_learning_rate = initial_learning_rate
if end_learning_rate:
self.end_learning_rate = end_learning_rate
else:
self.end_learning_rate = 0.0001
if warmup_epochs is not None:
warmup_epochs_ = warmup_epochs
self.warmup_epochs = warmup_epochs_
"""
opt_name = FLAGS.optimizer.lower()
mlp_log.mlperf_print('opt_name', opt_name)
if opt_name == 'lars':
mlp_log.mlperf_print('{}_epsilon'.format(opt_name), FLAGS.lars_epsilon)
mlp_log.mlperf_print('{}_opt_weight_decay'.format(opt_name),
FLAGS.weight_decay)
mlp_log.mlperf_print('{}_opt_base_learning_rate'.format(opt_name),
self.initial_learning_rate)
mlp_log.mlperf_print('{}_opt_learning_rate_warmup_epochs'.format(opt_name),
warmup_epochs_)
mlp_log.mlperf_print('{}_opt_end_learning_rate'.format(opt_name),
self.end_learning_rate)
"""
warmup_steps = warmup_epochs_ * steps_per_epoch
self.warmup_steps = tf.cast(warmup_steps, tf.float32)
self.decay_steps = train_steps - warmup_steps + 1
"""
mlp_log.mlperf_print('{}_opt_learning_rate_decay_steps'.format(opt_name),
int(self.decay_steps))
mlp_log.mlperf_print(
'{}_opt_learning_rate_decay_poly_power'.format(opt_name), 2.0)
mlp_log.mlperf_print('{}_opt_momentum'.format(opt_name), FLAGS.momentum)
"""
self.poly_rate_scheduler = tf.keras.optimizers.schedules.PolynomialDecay(
initial_learning_rate=self.initial_learning_rate,
decay_steps=self.decay_steps,
end_learning_rate=self.end_learning_rate,
power=2.0)
def __call__(self, step):
if tf.executing_eagerly():
return self._get_learning_rate(step)
# In an eager function or graph, the current implementation of optimizer
# repeatedly call and thus create ops for the learning rate schedule. To
# avoid this, we cache the ops if not executing eagerly.
graph = tf.compat.v1.get_default_graph()
if graph not in self.learning_rate_ops_cache:
if self.compute_lr_on_cpu:
with tf.device('/device:CPU:0'):
self.learning_rate_ops_cache[graph] = self._get_learning_rate(step)
else:
self.learning_rate_ops_cache[graph] = self._get_learning_rate(step)
return self.learning_rate_ops_cache[graph]
def _get_learning_rate(self, step):
with ops.name_scope_v2(self.name or 'PolynomialDecayWithWarmup') as name:
initial_learning_rate = ops.convert_to_tensor(
self.initial_learning_rate, name='initial_learning_rate')
warmup_steps = ops.convert_to_tensor(
self.warmup_steps, name='warmup_steps')
warmup_rate = (
initial_learning_rate * step / warmup_steps)
poly_steps = math_ops.subtract(step, warmup_steps)
poly_rate = self.poly_rate_scheduler(poly_steps)
decay_rate = tf.where(step <= warmup_steps,
warmup_rate, poly_rate, name=name)
return decay_rate
def get_config(self):
return {
'batch_size': self.batch_size,
'steps_per_epoch': self.steps_per_epoch,
'train_steps': self.train_steps,
'initial_learning_rate': self.initial_learning_rate,
'end_learning_rate': self.end_learning_rate,
'warmup_epochs': self.warmup_epochs,
'name': self.name,
}

89
artifacts/package_sources/tinygrad/test/external/mlperf_retinanet/coco_utils.py vendored Normal file
View File

@@ -0,0 +1,89 @@
# Copied from https://github.com/mlcommons/training/blob/637c82f9e699cd6caf108f92efb2c1d446b630e0/single_stage_detector/ssd/coco_utils.py
import os
import torch
import torchvision
from test.external.mlperf_retinanet import transforms as T
class ConvertCocoPolysToMask(object):
def __init__(self, filter_iscrowd=True):
self.filter_iscrowd = filter_iscrowd
def __call__(self, image, target):
w, h = image.size
image_id = target["image_id"]
image_id = torch.tensor([image_id])
anno = target["annotations"]
if self.filter_iscrowd:
anno = [obj for obj in anno if obj['iscrowd'] == 0]
boxes = [obj["bbox"] for obj in anno]
# guard against no boxes via resizing
boxes = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4)
boxes[:, 2:] += boxes[:, :2]
boxes[:, 0::2].clamp_(min=0, max=w)
boxes[:, 1::2].clamp_(min=0, max=h)
classes = [obj["category_id"] for obj in anno]
classes = torch.tensor(classes, dtype=torch.int64)
keypoints = None
if anno and "keypoints" in anno[0]:
keypoints = [obj["keypoints"] for obj in anno]
keypoints = torch.as_tensor(keypoints, dtype=torch.float32)
num_keypoints = keypoints.shape[0]
if num_keypoints:
keypoints = keypoints.view(num_keypoints, -1, 3)
keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
boxes = boxes[keep]
classes = classes[keep]
target = {}
target["boxes"] = boxes
target["labels"] = classes
target["image_id"] = image_id
# for conversion to coco api
area = torch.tensor([obj["area"] for obj in anno])
iscrowd = torch.tensor([obj["iscrowd"] for obj in anno])
target["area"] = area
target["iscrowd"] = iscrowd
return image, target
class CocoDetection(torchvision.datasets.CocoDetection):
def __init__(self, img_folder, ann_file, transforms):
super(CocoDetection, self).__init__(img_folder, ann_file)
self._transforms = transforms
def __getitem__(self, idx):
img, target = super(CocoDetection, self).__getitem__(idx)
image_id = self.ids[idx]
target = dict(image_id=image_id, annotations=target)
if self._transforms is not None:
img, target = self._transforms(img, target)
return img, target
def get_openimages(name, root, image_set, transforms):
PATHS = {
"train": os.path.join(root, "train"),
"val": os.path.join(root, "validation"),
}
t = [ConvertCocoPolysToMask(filter_iscrowd=False)]
if transforms is not None:
t.append(transforms)
transforms = T.Compose(t)
img_folder = os.path.join(PATHS[image_set], "data")
ann_file = os.path.join(PATHS[image_set], "labels", f"{name}.json")
dataset = CocoDetection(img_folder, ann_file, transforms=transforms)
return dataset

51
artifacts/package_sources/tinygrad/test/external/mlperf_retinanet/focal_loss.py vendored Normal file
View File

@@ -0,0 +1,51 @@
# Copied from https://github.com/mlcommons/training/blob/cdd928d4596c142c15a7d86b2eeadbac718c8da2/single_stage_detector/ssd/model/focal_loss.py
import torch
import torch.nn.functional as F
def sigmoid_focal_loss(
inputs: torch.Tensor,
targets: torch.Tensor,
alpha: float = 0.25,
gamma: float = 2,
reduction: str = "none",
):
"""
Original implementation from https://github.com/facebookresearch/fvcore/blob/master/fvcore/nn/focal_loss.py .
Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002.
Args:
inputs: A float tensor of arbitrary shape.
The predictions for each example.
targets: A float tensor with the same shape as inputs. Stores the binary
classification label for each element in inputs
(0 for the negative class and 1 for the positive class).
alpha: (optional) Weighting factor in range (0,1) to balance
positive vs negative examples or -1 for ignore. Default = 0.25
gamma: Exponent of the modulating factor (1 - p_t) to
balance easy vs hard examples.
reduction: 'none' | 'mean' | 'sum'
'none': No reduction will be applied to the output.
'mean': The output will be averaged.
'sum': The output will be summed.
Returns:
Loss tensor with the reduction option applied.
"""
p = torch.sigmoid(inputs)
ce_loss = F.binary_cross_entropy_with_logits(
inputs, targets, reduction="none"
)
p_t = p * targets + (1 - p) * (1 - targets)
loss = ce_loss * ((1 - p_t) ** gamma)
if alpha >= 0:
alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
loss = alpha_t * loss
if reduction == "mean":
loss = loss.mean()
elif reduction == "sum":
loss = loss.sum()
return loss

65
artifacts/package_sources/tinygrad/test/external/mlperf_retinanet/model/boxes.py vendored Normal file
View File

@@ -0,0 +1,65 @@
# Copied from https://github.com/mlcommons/training/blob/637c82f9e699cd6caf108f92efb2c1d446b630e0/single_stage_detector/ssd/model/boxes.py
import torch
from torch import Tensor
from typing import Tuple
def _upcast(t: Tensor) -> Tensor:
# Protects from numerical overflows in multiplications by upcasting to the equivalent higher type
if t.is_floating_point():
return t if t.dtype in (torch.float32, torch.float64) else t.float()
else:
return t if t.dtype in (torch.int32, torch.int64) else t.int()
def box_area(boxes: Tensor) -> Tensor:
"""
Computes the area of a set of bounding boxes, which are specified by their
(x1, y1, x2, y2) coordinates.
Args:
boxes (Tensor[N, 4]): boxes for which the area will be computed. They
are expected to be in (x1, y1, x2, y2) format with
``0 <= x1 < x2`` and ``0 <= y1 < y2``.
Returns:
Tensor[N]: the area for each box
"""
boxes = _upcast(boxes)
return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
# implementation from https://github.com/kuangliu/torchcv/blob/master/torchcv/utils/box.py
# with slight modifications
def _box_inter_union(boxes1: Tensor, boxes2: Tensor) -> Tuple[Tensor, Tensor]:
area1 = box_area(boxes1)
area2 = box_area(boxes2)
lt = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2]
rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2]
wh = _upcast(rb - lt).clamp(min=0) # [N,M,2]
inter = wh[:, :, 0] * wh[:, :, 1] # [N,M]
union = area1[:, None] + area2 - inter
return inter, union
def box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
"""
Return intersection-over-union (Jaccard index) between two sets of boxes.
Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with
``0 <= x1 < x2`` and ``0 <= y1 < y2``.
Args:
boxes1 (Tensor[N, 4]): first set of boxes
boxes2 (Tensor[M, 4]): second set of boxes
Returns:
Tensor[N, M]: the NxM matrix containing the pairwise IoU values for every element in boxes1 and boxes2
"""
inter, union = _box_inter_union(boxes1, boxes2)
iou = inter / union
return iou

27
artifacts/package_sources/tinygrad/test/external/mlperf_retinanet/model/image_list.py vendored Normal file
View File

@@ -0,0 +1,27 @@
# Copied from https://github.com/mlcommons/training/blob/637c82f9e699cd6caf108f92efb2c1d446b630e0/single_stage_detector/ssd/model/image_list.py
import torch
from torch import Tensor
from typing import List, Tuple
class ImageList(object):
"""
Structure that holds a list of images (of possibly
varying sizes) as a single tensor.
This works by padding the images to the same size,
and storing in a field the original sizes of each image
"""
def __init__(self, tensors: Tensor, image_sizes: List[Tuple[int, int]]):
"""
Args:
tensors (tensor)
image_sizes (list[tuple[int, int]])
"""
self.tensors = tensors
self.image_sizes = image_sizes
def to(self, device: torch.device) -> 'ImageList':
cast_tensor = self.tensors.to(device)
return ImageList(cast_tensor, self.image_sizes)

163
artifacts/package_sources/tinygrad/test/external/mlperf_retinanet/model/transform.py vendored Normal file
View File

@@ -0,0 +1,163 @@
# Copied from https://github.com/mlcommons/training/blob/637c82f9e699cd6caf108f92efb2c1d446b630e0/single_stage_detector/ssd/model/transform.py
import torch
from torch import nn, Tensor
from typing import List, Tuple, Dict, Optional
from test.external.mlperf_retinanet.model.image_list import ImageList
@torch.jit.unused
def _get_shape_onnx(image: Tensor) -> Tensor:
from torch.onnx import operators
return operators.shape_as_tensor(image)[-2:]
def _resize_image_and_masks(image: Tensor,
target: Optional[Dict[str, Tensor]] = None,
image_size: Optional[Tuple[int, int]] = None,
) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
image = torch.nn.functional.interpolate(image[None], size=image_size, scale_factor=None, mode='bilinear',
recompute_scale_factor=None, align_corners=False)[0]
if target is None:
return image, target
if "masks" in target:
mask = target["masks"]
mask = torch.nn.functional.interpolate(mask[:, None].float(), size=image_size, scale_factor=None,
recompute_scale_factor=None)[:, 0].byte()
target["masks"] = mask
return image, target
class GeneralizedRCNNTransform(nn.Module):
"""
Performs input / target transformation before feeding the data to a GeneralizedRCNN
model.
The transformations it perform are:
- input normalization (mean subtraction and std division)
- input / target resizing to match image_size
It returns a ImageList for the inputs, and a List[Dict[Tensor]] for the targets
"""
def __init__(self, image_size: Optional[Tuple[int, int]],
image_mean: List[float], image_std: List[float],):
super(GeneralizedRCNNTransform, self).__init__()
self.image_size = image_size
self.image_mean = image_mean
self.image_std = image_std
def forward(self,
images: List[Tensor],
targets: Optional[List[Dict[str, Tensor]]] = None
) -> Tuple[ImageList, Optional[List[Dict[str, Tensor]]]]:
images = list(img for img in images)
if targets is not None:
# make a copy of targets to avoid modifying it in-place
# once torchscript supports dict comprehension
# this can be simplified as follows
# targets = [{k: v for k,v in t.items()} for t in targets]
targets_copy: List[Dict[str, Tensor]] = []
for t in targets:
data: Dict[str, Tensor] = {}
for k, v in t.items():
data[k] = v
targets_copy.append(data)
targets = targets_copy
for i in range(len(images)):
image = images[i]
target_index = targets[i] if targets is not None else None
if image.dim() != 3:
raise ValueError("images is expected to be a list of 3d tensors "
"of shape [C, H, W], got {}".format(image.shape))
image = self.normalize(image)
image, target_index = self.resize(image, target_index)
images[i] = image
if targets is not None and target_index is not None:
targets[i] = target_index
image_sizes = [img.shape[-2:] for img in images]
images = torch.stack(images)
image_sizes_list: List[Tuple[int, int]] = []
for image_size in image_sizes:
assert len(image_size) == 2
image_sizes_list.append((image_size[0], image_size[1]))
image_list = ImageList(images, image_sizes_list)
return image_list, targets
def normalize(self, image: Tensor) -> Tensor:
if not image.is_floating_point():
raise TypeError(
f"Expected input images to be of floating type (in range [0, 1]), "
f"but found type {image.dtype} instead"
)
dtype, device = image.dtype, image.device
mean = torch.as_tensor(self.image_mean, dtype=dtype, device=device)
std = torch.as_tensor(self.image_std, dtype=dtype, device=device)
return (image - mean[:, None, None]) / std[:, None, None]
def torch_choice(self, k: List[int]) -> int:
"""
Implements `random.choice` via torch ops so it can be compiled with
TorchScript. Remove if https://github.com/pytorch/pytorch/issues/25803
is fixed.
"""
index = int(torch.empty(1).uniform_(0., float(len(k))).item())
return k[index]
def resize(self,
image: Tensor,
target: Optional[Dict[str, Tensor]] = None,
) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
h, w = image.shape[-2:]
image, target = _resize_image_and_masks(image, target, self.image_size)
if target is None:
return image, target
bbox = target["boxes"]
bbox = resize_boxes(bbox, (h, w), image.shape[-2:])
target["boxes"] = bbox
return image, target
def postprocess(self,
result: List[Dict[str, Tensor]],
image_shapes: List[Tuple[int, int]],
original_image_sizes: List[Tuple[int, int]]
) -> List[Dict[str, Tensor]]:
if self.training:
return result
for i, (pred, im_s, o_im_s) in enumerate(zip(result, image_shapes, original_image_sizes)):
boxes = pred["boxes"]
boxes = resize_boxes(boxes, im_s, o_im_s)
result[i]["boxes"] = boxes
return result
def __repr__(self) -> str:
format_string = self.__class__.__name__ + '('
_indent = '\n '
format_string += "{0}Normalize(mean={1}, std={2})".format(_indent, self.image_mean, self.image_std)
format_string += "{0}Resize(height={1}, width={2}, mode='bilinear')".format(_indent, self.image_size[0],
self.image_size[1])
format_string += '\n)'
return format_string
def resize_boxes(boxes: Tensor, original_size: List[int], new_size: List[int]) -> Tensor:
ratios = [
torch.tensor(s, dtype=torch.float32, device=boxes.device) /
torch.tensor(s_orig, dtype=torch.float32, device=boxes.device)
for s, s_orig in zip(new_size, original_size)
]
ratio_height, ratio_width = ratios
xmin, ymin, xmax, ymax = boxes.unbind(1)
xmin = xmin * ratio_width
xmax = xmax * ratio_width
ymin = ymin * ratio_height
ymax = ymax * ratio_height
res = torch.stack((xmin, ymin, xmax, ymax), dim=1)
return res

123
artifacts/package_sources/tinygrad/test/external/mlperf_retinanet/model/utils.py vendored Normal file
View File

@@ -0,0 +1,123 @@
# Copied from https://github.com/mlcommons/training/blob/637c82f9e699cd6caf108f92efb2c1d446b630e0/single_stage_detector/ssd/model/utils.py
import torch
class Matcher(object):
"""
This class assigns to each predicted "element" (e.g., a box) a ground-truth
element. Each predicted element will have exactly zero or one matches; each
ground-truth element may be assigned to zero or more predicted elements.
Matching is based on the MxN match_quality_matrix, that characterizes how well
each (ground-truth, predicted)-pair match. For example, if the elements are
boxes, the matrix may contain box IoU overlap values.
The matcher returns a tensor of size N containing the index of the ground-truth
element m that matches to prediction n. If there is no match, a negative value
is returned.
"""
BELOW_LOW_THRESHOLD = -1
BETWEEN_THRESHOLDS = -2
__annotations__ = {
'BELOW_LOW_THRESHOLD': int,
'BETWEEN_THRESHOLDS': int,
}
def __init__(self, high_threshold, low_threshold, allow_low_quality_matches=False):
# type: (float, float, bool) -> None
"""
Args:
high_threshold (float): quality values greater than or equal to
this value are candidate matches.
low_threshold (float): a lower quality threshold used to stratify
matches into three levels:
1) matches >= high_threshold
2) BETWEEN_THRESHOLDS matches in [low_threshold, high_threshold)
3) BELOW_LOW_THRESHOLD matches in [0, low_threshold)
allow_low_quality_matches (bool): if True, produce additional matches
for predictions that have only low-quality match candidates. See
set_low_quality_matches_ for more details.
"""
self.BELOW_LOW_THRESHOLD = -1
self.BETWEEN_THRESHOLDS = -2
assert low_threshold <= high_threshold
self.high_threshold = high_threshold
self.low_threshold = low_threshold
self.allow_low_quality_matches = allow_low_quality_matches
def __call__(self, match_quality_matrix):
"""
Args:
match_quality_matrix (Tensor[float]): an MxN tensor, containing the
pairwise quality between M ground-truth elements and N predicted elements.
Returns:
matches (Tensor[int64]): an N tensor where N[i] is a matched gt in
[0, M - 1] or a negative value indicating that prediction i could not
be matched.
"""
if match_quality_matrix.numel() == 0:
# empty targets or proposals not supported during training
if match_quality_matrix.shape[0] == 0:
raise ValueError(
"No ground-truth boxes available for one of the images "
"during training")
raise ValueError(
"No proposal boxes available for one of the images "
"during training")
# match_quality_matrix is M (gt) x N (predicted)
# Max over gt elements (dim 0) to find best gt candidate for each prediction
matched_vals, matches = match_quality_matrix.max(dim=0)
if self.allow_low_quality_matches:
all_matches = matches.clone()
else:
all_matches = None
# Assign candidate matches with low quality to negative (unassigned) values
below_low_threshold = matched_vals < self.low_threshold
between_thresholds = (matched_vals >= self.low_threshold) & (
matched_vals < self.high_threshold
)
matches[below_low_threshold] = self.BELOW_LOW_THRESHOLD
matches[between_thresholds] = self.BETWEEN_THRESHOLDS
if self.allow_low_quality_matches:
assert all_matches is not None
self.set_low_quality_matches_(matches, all_matches, match_quality_matrix)
return matches
def set_low_quality_matches_(self, matches, all_matches, match_quality_matrix):
"""
Produce additional matches for predictions that have only low-quality matches.
Specifically, for each ground-truth find the set of predictions that have
maximum overlap with it (including ties); for each prediction in that set, if
it is unmatched, then match it to the ground-truth with which it has the highest
quality value.
"""
# For each gt, find the prediction with which it has highest quality
highest_quality_foreach_gt, _ = match_quality_matrix.max(dim=1)
# Find highest quality match available, even if it is low, including ties
gt_pred_pairs_of_highest_quality = torch.where(
match_quality_matrix == highest_quality_foreach_gt[:, None]
)
# Example gt_pred_pairs_of_highest_quality:
# tensor([[ 0, 39796],
# [ 1, 32055],
# [ 1, 32070],
# [ 2, 39190],
# [ 2, 40255],
# [ 3, 40390],
# [ 3, 41455],
# [ 4, 45470],
# [ 5, 45325],
# [ 5, 46390]])
# Each row is a (gt index, prediction index)
# Note how gt items 1, 2, 3, and 5 each have two ties
pred_inds_to_update = gt_pred_pairs_of_highest_quality[1]
matches[pred_inds_to_update] = all_matches[pred_inds_to_update]

25
artifacts/package_sources/tinygrad/test/external/mlperf_retinanet/openimages.py vendored Normal file
View File

@@ -0,0 +1,25 @@
from test.external.mlperf_retinanet.model.boxes import box_iou
from test.external.mlperf_retinanet.model.utils import Matcher
import torch
# This applies the filtering in https://github.com/mlcommons/training/blob/cdd928d4596c142c15a7d86b2eeadbac718c8da2/single_stage_detector/ssd/model/retinanet.py#L117
# and https://github.com/mlcommons/training/blob/cdd928d4596c142c15a7d86b2eeadbac718c8da2/single_stage_detector/ssd/model/retinanet.py#L203
# to match with tinygrad's dataloader implementation.
def postprocess_targets(targets, anchors):
proposal_matcher, matched_idxs = Matcher(0.5, 0.4, allow_low_quality_matches=True), []
for anchors_per_image, targets_per_image in zip(anchors, targets):
if targets_per_image['boxes'].numel() == 0:
matched_idxs.append(torch.full((anchors_per_image.size(0),), -1, dtype=torch.int64,
device=anchors_per_image.device))
continue
match_quality_matrix = box_iou(targets_per_image['boxes'], anchors_per_image)
matched_idxs.append(proposal_matcher(match_quality_matrix))
for targets_per_image, matched_idxs_per_image in zip(targets, matched_idxs):
foreground_idxs_per_image = matched_idxs_per_image >= 0
targets_per_image["boxes"] = targets_per_image["boxes"][matched_idxs_per_image[foreground_idxs_per_image]]
targets_per_image["labels"] = targets_per_image["labels"][matched_idxs_per_image[foreground_idxs_per_image]]
return targets

24
artifacts/package_sources/tinygrad/test/external/mlperf_retinanet/presets.py vendored Normal file
View File

@@ -0,0 +1,24 @@
# Copied from https://github.com/mlcommons/training/blob/637c82f9e699cd6caf108f92efb2c1d446b630e0/single_stage_detector/ssd/presets.py
from test.external.mlperf_retinanet import transforms as T
class DetectionPresetTrain:
def __init__(self, data_augmentation, hflip_prob=0.5, mean=(123., 117., 104.)):
if data_augmentation == 'hflip':
self.transforms = T.Compose([
T.RandomHorizontalFlip(p=hflip_prob),
T.ToTensor(),
])
else:
raise ValueError(f'Unknown data augmentation policy "{data_augmentation}"')
def __call__(self, img, target):
return self.transforms(img, target)
class DetectionPresetEval:
def __init__(self):
self.transforms = T.ToTensor()
def __call__(self, img, target):
return self.transforms(img, target)

77
artifacts/package_sources/tinygrad/test/external/mlperf_retinanet/transforms.py vendored Normal file
View File

@@ -0,0 +1,77 @@
# Copied from https://github.com/mlcommons/training/blob/637c82f9e699cd6caf108f92efb2c1d446b630e0/single_stage_detector/ssd/transforms.py
import torch
import torchvision
from torch import nn, Tensor
from torchvision.transforms import functional as F
from torchvision.transforms import transforms as T
from typing import List, Tuple, Dict, Optional
from PIL import Image
Image.MAX_IMAGE_PIXELS = None
from typing import Any
try:
import accimage
except ImportError:
accimage = None
@torch.jit.unused
def _is_pil_image(img: Any) -> bool:
if accimage is not None:
return isinstance(img, (Image.Image, accimage.Image))
else:
return isinstance(img, Image.Image)
def get_image_size_tensor(img: Tensor) -> List[int]:
# Returns (w, h) of tensor image
torchvision.transforms._functional_tensor._assert_image_tensor(img)
return [img.shape[-1], img.shape[-2]]
@torch.jit.unused
def get_image_size_pil(img: Any) -> List[int]:
if _is_pil_image(img):
return list(img.size)
raise TypeError("Unexpected type {}".format(type(img)))
def get_image_size(img: Tensor) -> List[int]:
"""Returns the size of an image as [width, height].
Args:
img (PIL Image or Tensor): The image to be checked.
Returns:
List[int]: The image size.
"""
if isinstance(img, torch.Tensor):
return get_image_size_tensor(img)
return get_image_size_pil(img)
class Compose(object):
def __init__(self, transforms):
self.transforms = transforms
def __call__(self, image, target):
for t in self.transforms:
image, target = t(image, target)
return image, target
class RandomHorizontalFlip(T.RandomHorizontalFlip):
def forward(self, image: Tensor,
target: Optional[Dict[str, Tensor]] = None) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
if torch.rand(1) < self.p:
image = F.hflip(image)
if target is not None:
width, _ = get_image_size(image)
target["boxes"][:, [0, 2]] = width - target["boxes"][:, [2, 0]]
if "masks" in target:
target["masks"] = target["masks"].flip(-1)
return image, target
class ToTensor(nn.Module):
def forward(self, image: Tensor,
target: Optional[Dict[str, Tensor]] = None) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
image = F.to_tensor(image)
return image, target

58
artifacts/package_sources/tinygrad/test/external/mlperf_stable_diffusion/external_test_eval.py vendored Normal file
View File

@@ -0,0 +1,58 @@
import unittest, os
import numpy as np
from pathlib import Path
from tempfile import TemporaryDirectory
from tinygrad import Device, Tensor
from tinygrad.helpers import getenv, Context
from tinygrad.nn.state import safe_save, torch_load, get_parameters
from examples.mlperf.model_eval import eval_stable_diffusion, vae_decode
from examples.stable_diffusion import AutoencoderKL
def set_eval_params():
# override these as needed from cli
for k,v in {"MODEL": "stable_diffusion", "GPUS": "8", "EVAL_SAMPLES": "600", "CONTEXT_BS": "816", "DENOISE_BS": "600", "DECODE_BS": "384",
"INCEPTION_BS": "560", "CLIP_BS": "240", "DATADIR": "/raid/datasets/stable_diffusion", "CKPTDIR": "/raid/weights/stable_diffusion"}.items():
os.environ[k] = getenv(k, v)
class TestEval(unittest.TestCase):
def test_eval_ckpt(self):
set_eval_params()
with TemporaryDirectory(prefix="test-eval") as tmp:
os.environ["EVAL_CKPT_DIR"] = tmp
# NOTE Although this checkpoint has the original fully trained model from StabilityAI, we are using mlperf code that uses different
# GroupNorm num_groups. Therefore, eval results may not reflect eval results on the original model.
# The purpose of using this checkpoint is to have reproducible eval outputs.
# Eval code expects file and weight names in a specific format, as .safetensors (not .ckpt), which is why we resave the checkpoint
sd_v2 = torch_load(Path(getenv("CKPTDIR", "")) / "sd" / "512-base-ema.ckpt")["state_dict"]
sd_v2 = {k.replace("model.diffusion_model.", "", 1): v for k,v in sd_v2.items() if k.startswith("model.diffusion_model.")}
safe_save(sd_v2, f"{tmp}/0.safetensors")
clip, fid, ckpt = eval_stable_diffusion()
assert ckpt == 0
if Device.DEFAULT == "NULL":
assert clip == 0
assert fid > 0 and fid < 1000
else:
# observed:
# clip=0.08369670808315277, fid=301.05236173709545 (if SEED=12345, commit=c01b2c93076e80ae6d1ebca64bb8e83a54dadba6)
# clip=0.08415728807449341, fid=300.3710877072948 (if SEED=12345, commit=179c7fcfe132f1a6344b57c9d8cef4eded586867)
# clip=0.0828116238117218, fid=301.241909555543 (if SEED=98765, commit=c01b2c93076e80ae6d1ebca64bb8e83a54dadba6)
np.testing.assert_allclose(fid, 301.147, rtol=0.1, atol=0)
np.testing.assert_allclose(clip, 0.08325, rtol=0.1, atol=0)
# only tested on 8xMI300x system
@unittest.skipUnless(getenv("HANG_OK"), "expected to hang")
def test_decoder_beam_hang(self):
set_eval_params()
for k,v in {"BEAM": "2", "HCQDEV_WAIT_TIMEOUT_MS": "300000", "BEAM_UOPS_MAX": "8000", "BEAM_UPCAST_MAX": "256", "BEAM_LOCAL_MAX": "1024",
"BEAM_MIN_PROGRESS": "5", "IGNORE_JIT_FIRST_BEAM": "1"}.items():
os.environ[k] = getenv(k, v)
with Context(BEAM=int(os.environ["BEAM"])): # necessary because helpers.py has already set BEAM=0 and cached getenv for "BEAM"
GPUS = [f"{Device.DEFAULT}:{i}" for i in range(getenv("GPUS", 8))]
vae = AutoencoderKL()
for p in get_parameters(vae): p.to_(GPUS).realize()
x = Tensor.zeros(48,4,64,64).contiguous().to(GPUS).realize()
x.uop = x.uop.unshard(0)
for _ in range(2): vae_decode(x, vae)
if __name__=="__main__":
unittest.main()

114
artifacts/package_sources/tinygrad/test/external/mlperf_stable_diffusion/external_test_models.py vendored Normal file
View File

@@ -0,0 +1,114 @@
import unittest
import numpy as np
from pathlib import Path
from tinygrad import Tensor, dtypes, Device
from tinygrad.helpers import getenv
from tinygrad.nn.state import get_parameters
from extra.models import clip
from examples.mlperf.initializers import gelu_erf, init_stable_diffusion, attn_f32_softmax
from typing import Literal
clip_params = {"dims": 1024, "n_heads": 16, "layers": 24, "return_pooled": False, "ln_penultimate": True}
def get_cond_stage_model(GPUS:list[str]|None=None) -> clip.FrozenOpenClipEmbedder:
clip.gelu = gelu_erf
model = clip.FrozenOpenClipEmbedder(**clip_params)
if GPUS and len(GPUS) > 1:
for p in get_parameters(model): p.to_(GPUS)
return model
def get_tokens(BS:int) -> Tensor: return Tensor([0] * 77 * BS, dtype=dtypes.int32).reshape(-1, 77)
class TestOpenClip(unittest.TestCase):
def test_tokenizer(self):
prompt = "Beautiful is better than ugly.\nExplicit is better than implicit.\nSimple is better than complex.\nComplex is better than complicated."
model = get_cond_stage_model()
tokens = model.tokenizer.encode(prompt, pad_with_zeros=True)
expected = [49406, 1215, 533, 1539, 1126, 36070, 269, 33228, 533, 1539, 1126, 15269, 529, 269, 4129, 533, 1539, 1126, 16099, 269, 6324, 533,
1539, 1126, 11460, 14589, 269, 49407] + [0]*49
self.assertEqual(tokens, expected)
def test_clip_gelu_init(self):
for resblock in get_cond_stage_model().model.transformer.resblocks:
self.assertEqual(resblock.mlp.gelu, gelu_erf)
def test_multigpu_clip_embed(self):
BS = 304
GPUS = [f"{Device.DEFAULT}:{i}" for i in range(8)]
model = get_cond_stage_model(GPUS)
tokens = get_tokens(BS)
embeds = model.embed_tokens(tokens.shard(GPUS, axis=0)).realize()
self.assertEqual(embeds.shape, (BS, 77, 1024))
self.assertEqual(embeds.dtype, dtypes.float32)
def test_multigpu_clip_score(self):
BS = 240
GPUS = [f"{Device.DEFAULT}:{i}" for i in range(8)]
vision_cfg = {'width': 1280, 'layers': 32, 'd_head': 80, 'image_size': 224, 'patch_size': 14}
text_cfg = {'width': 1024, 'n_heads': 16, 'layers': 24, 'vocab_size': 49408, 'ctx_length': 77}
clip.gelu = gelu_erf
clip_encoder = clip.OpenClipEncoder(1024, text_cfg, vision_cfg)
for p in get_parameters(clip_encoder): p.to_(GPUS)
tokens = get_tokens(BS)
imgs = Tensor.zeros(BS,3,224,224).contiguous()
scores = clip_encoder.get_clip_score(tokens.shard(GPUS, axis=0), imgs.shard(GPUS, axis=0)).realize()
self.assertEqual(scores.shape, (BS,))
self.assertEqual(scores.dtype, dtypes.float32)
class TestInitStableDiffusion(unittest.TestCase):
def setUp(self):
# NOTE: set env variable based on where checkpoints are on the system
self.CKPTDIR = Path(getenv("CKPTDIR", "/raid/weights/stable_diffusion"))
def helper_test_init(self, version:Literal["v2-mlperf-train", "v2-mlperf-eval"]):
model, unet, sqrt_acp, sqrt_omacp = init_stable_diffusion(version, self.CKPTDIR / "sd" / "512-base-ema.ckpt", ["CPU"])
with self.subTest("test that StableDiffusion has correct models"):
self.assertEqual(model.model.diffusion_model, unet)
has_encoder = True if version=="v2-mlperf-eval" else False
self.assertEqual(hasattr(model, "first_stage_model"), has_encoder, "only the eval model uses the encoder")
self.assertTrue(isinstance(model.cond_stage_model, clip.FrozenOpenClipEmbedder))
with self.subTest("test for mlperf unique attributes"):
self.assertEqual(model.cond_stage_model.tokenizer.version, 'sd_mlperf_v5_0')
self.assertEqual(unet.out[0].num_groups, 16)
self.assertEqual(unet.input_blocks[1][1].norm.eps, 1e-6)
self.assertEqual(unet.input_blocks[1][1].transformer_blocks[0].attn1.attn, attn_f32_softmax)
with self.subTest("test loaded clip parameters"):
sample = model.cond_stage_model.model.transformer.resblocks[8].mlp.c_fc.bias.flatten()[42:46].numpy()
expected = np.array([-0.49812260270118713, -0.3039605915546417, -0.40284937620162964, -0.45069342851638794], dtype=np.float32)
np.testing.assert_allclose(sample, expected, rtol=1e-7, atol=0, err_msg="loaded clip parameters are incorrect")
if version=="v2-mlperf-train":
with self.subTest("test that zero_module worked"):
self.assertTrue((unet.out[2].weight == 0).all().item(), "expected all zeroes")
self.assertTrue((unet.out[2].bias == 0).all().item(), "expected all zeroes")
elif version=="v2-mlperf-eval":
with self.subTest("test loaded vae parameters"):
sample = model.first_stage_model.decoder.up[0]['block'][1].conv2.weight.flatten()[42:46].numpy()
expected = np.array([0.08192943036556244, 0.040095631033182144, 0.07541035860776901, 0.1475081741809845], dtype=np.float32)
np.testing.assert_allclose(sample, expected, rtol=1e-7, atol=0, err_msg="loaded vae parameters are incorrect")
with self.subTest("check schedules"):
expected = np.array([0.9995748996734619, 0.06826484948396683], dtype=np.float32)
np.testing.assert_allclose(sqrt_acp[[0,-1]].numpy(), expected, rtol=1e-7, atol=0, err_msg="sqrt_acp is incorrect")
expected = np.array([0.029155133292078972, 0.9976672530174255], dtype=np.float32)
np.testing.assert_allclose(sqrt_omacp[[0,-1]].numpy(), expected, rtol=1e-7, atol=0, err_msg="sqrt_omacp is incorrect")
with self.subTest("check mixed precision"):
out = unet.input_blocks[2][1].proj_in(Tensor.randn(320, dtype=dtypes.float32))
self.assertEqual(out.dtype, dtypes.bfloat16, "expected float32 to be downcast to bfloat16 by Linear")
out = unet.out[2](Tensor.randn(304,320,64,64, dtype=dtypes.float32))
self.assertEqual(out.dtype, dtypes.bfloat16, "expected float32 to be downcast to bfloat16 by Conv2d")
out = unet.input_blocks[1][1].transformer_blocks[0].norm1(Tensor.randn(320, dtype=dtypes.bfloat16))
self.assertEqual(out.dtype, dtypes.float32, "expected bfloat16 to be upcast to float32 by LayerNorm")
out = unet.input_blocks[5][0].in_layers[0](Tensor.randn(304, 640, dtype=dtypes.bfloat16))
self.assertEqual(out.dtype, dtypes.float32, "expected bfloat16 to be upcast to float32 by GroupNorm")
def test_train_model(self):
self.helper_test_init("v2-mlperf-train")
def test_eval_model(self):
self.helper_test_init("v2-mlperf-eval")
if __name__=="__main__":
unittest.main()

23
artifacts/package_sources/tinygrad/test/external/mlperf_stable_diffusion/external_test_train.py vendored Normal file
View File

@@ -0,0 +1,23 @@
import unittest, os
from tempfile import TemporaryDirectory
from tinygrad import Context
from tinygrad.helpers import getenv
from examples.mlperf.model_train import train_stable_diffusion
class TestTrain(unittest.TestCase):
def test_train_to_ckpt(self):
# train for num_steps, save checkpoint, and stop training
num_steps = 42
os.environ.update({"MODEL": "stable_diffusion", "TOTAL_CKPTS": "1", "CKPT_STEP_INTERVAL": str(num_steps), "GPUS": "8", "BS": "304"})
# NOTE: update these based on where data/checkpoints are on your system
if not getenv("DATADIR", ""): os.environ["DATADIR"] = "/raid/datasets/stable_diffusion"
if not getenv("CKPTDIR", ""): os.environ["CKPTDIR"] = "/raid/weights/stable_diffusion"
with TemporaryDirectory(prefix="test-train") as tmp:
os.environ["UNET_CKPTDIR"] = tmp
with Context(TRAINING=1):
saved_ckpts = train_stable_diffusion()
expected_ckpt = f"{tmp}/{num_steps}.safetensors"
assert len(saved_ckpts) == 1 and saved_ckpts[0] == expected_ckpt
if __name__=="__main__":
unittest.main()

94
artifacts/package_sources/tinygrad/test/external/mlperf_unet3d/dice.py vendored Normal file
View File

@@ -0,0 +1,94 @@
# https://github.com/mlcommons/training/blob/master/image_segmentation/pytorch/model/losses.py
import torch
import torch.nn as nn
import torch.nn.functional as F
class Dice:
def __init__(self,
to_onehot_y: bool = True,
to_onehot_x: bool = False,
use_softmax: bool = True,
use_argmax: bool = False,
include_background: bool = False,
layout: str = "NCDHW"):
self.include_background = include_background
self.to_onehot_y = to_onehot_y
self.to_onehot_x = to_onehot_x
self.use_softmax = use_softmax
self.use_argmax = use_argmax
self.smooth_nr = 1e-6
self.smooth_dr = 1e-6
self.layout = layout
def __call__(self, prediction, target):
if self.layout == "NCDHW":
channel_axis = 1
reduce_axis = list(range(2, len(prediction.shape)))
else:
channel_axis = -1
reduce_axis = list(range(1, len(prediction.shape) - 1))
num_pred_ch = prediction.shape[channel_axis]
if self.use_softmax:
prediction = torch.softmax(prediction, dim=channel_axis)
elif self.use_argmax:
prediction = torch.argmax(prediction, dim=channel_axis)
if self.to_onehot_y:
target = to_one_hot(target, self.layout, channel_axis)
if self.to_onehot_x:
prediction = to_one_hot(prediction, self.layout, channel_axis)
if not self.include_background:
assert num_pred_ch > 1, \
f"To exclude background the prediction needs more than one channel. Got {num_pred_ch}."
if self.layout == "NCDHW":
target = target[:, 1:]
prediction = prediction[:, 1:]
else:
target = target[..., 1:]
prediction = prediction[..., 1:]
assert (target.shape == prediction.shape), \
f"Target and prediction shape do not match. Target: ({target.shape}), prediction: ({prediction.shape})."
intersection = torch.sum(target * prediction, dim=reduce_axis)
target_sum = torch.sum(target, dim=reduce_axis)
prediction_sum = torch.sum(prediction, dim=reduce_axis)
return (2.0 * intersection + self.smooth_nr) / (target_sum + prediction_sum + self.smooth_dr)
def to_one_hot(array, layout, channel_axis):
if len(array.shape) >= 5:
array = torch.squeeze(array, dim=channel_axis)
array = F.one_hot(array.long(), num_classes=3)
if layout == "NCDHW":
array = array.permute(0, 4, 1, 2, 3).float()
return array
class DiceCELoss(nn.Module):
def __init__(self, to_onehot_y, use_softmax, layout, include_background):
super(DiceCELoss, self).__init__()
self.dice = Dice(to_onehot_y=to_onehot_y, use_softmax=use_softmax, layout=layout,
include_background=include_background)
self.cross_entropy = nn.CrossEntropyLoss()
def forward(self, y_pred, y_true):
cross_entropy = self.cross_entropy(y_pred, torch.squeeze(y_true, dim=1).long())
dice = torch.mean(1.0 - self.dice(y_pred, y_true))
return (dice + cross_entropy) / 2
class DiceScore:
def __init__(self, to_onehot_y: bool = True, use_argmax: bool = True, layout: str = "NCDHW",
include_background: bool = False):
self.dice = Dice(to_onehot_y=to_onehot_y, to_onehot_x=True, use_softmax=False,
use_argmax=use_argmax, layout=layout, include_background=include_background)
def __call__(self, y_pred, y_true):
return torch.mean(self.dice(y_pred, y_true), dim=0)

165
artifacts/package_sources/tinygrad/test/external/mlperf_unet3d/kits19.py vendored Normal file
View File

@@ -0,0 +1,165 @@
# copied from https://github.com/mlcommons/training/blob/5c08ce57e7f582cc4558035d8324a2bf4c8ca225/image_segmentation/pytorch/data_loading/pytorch_loader.py
import random
import numpy as np
import scipy.ndimage
from torch.utils.data import Dataset
from torchvision import transforms
def get_train_transforms():
rand_flip = RandFlip()
cast = Cast(types=(np.float32, np.uint8))
rand_scale = RandomBrightnessAugmentation(factor=0.3, prob=0.1)
rand_noise = GaussianNoise(mean=0.0, std=0.1, prob=0.1)
train_transforms = transforms.Compose([rand_flip, cast, rand_scale, rand_noise])
return train_transforms
class RandBalancedCrop:
def __init__(self, patch_size, oversampling):
self.patch_size = patch_size
self.oversampling = oversampling
def __call__(self, data):
image, label = data["image"], data["label"]
if random.random() < self.oversampling:
image, label, cords = self.rand_foreg_cropd(image, label)
else:
image, label, cords = self._rand_crop(image, label)
data.update({"image": image, "label": label})
return data
@staticmethod
def randrange(max_range):
return 0 if max_range == 0 else random.randrange(max_range)
def get_cords(self, cord, idx):
return cord[idx], cord[idx] + self.patch_size[idx]
def _rand_crop(self, image, label):
ranges = [s - p for s, p in zip(image.shape[1:], self.patch_size)]
cord = [self.randrange(x) for x in ranges]
low_x, high_x = self.get_cords(cord, 0)
low_y, high_y = self.get_cords(cord, 1)
low_z, high_z = self.get_cords(cord, 2)
image = image[:, low_x:high_x, low_y:high_y, low_z:high_z]
label = label[:, low_x:high_x, low_y:high_y, low_z:high_z]
return image, label, [low_x, high_x, low_y, high_y, low_z, high_z]
def rand_foreg_cropd(self, image, label):
def adjust(foreg_slice, patch_size, label, idx):
diff = patch_size[idx - 1] - (foreg_slice[idx].stop - foreg_slice[idx].start)
sign = -1 if diff < 0 else 1
diff = abs(diff)
ladj = self.randrange(diff)
hadj = diff - ladj
low = max(0, foreg_slice[idx].start - sign * ladj)
high = min(label.shape[idx], foreg_slice[idx].stop + sign * hadj)
diff = patch_size[idx - 1] - (high - low)
if diff > 0 and low == 0:
high += diff
elif diff > 0:
low -= diff
return low, high
cl = np.random.choice(np.unique(label[label > 0]))
foreg_slices = scipy.ndimage.find_objects(scipy.ndimage.measurements.label(label==cl)[0])
foreg_slices = [x for x in foreg_slices if x is not None]
slice_volumes = [np.prod([s.stop - s.start for s in sl]) for sl in foreg_slices]
slice_idx = np.argsort(slice_volumes)[-2:]
foreg_slices = [foreg_slices[i] for i in slice_idx]
if not foreg_slices:
return self._rand_crop(image, label)
foreg_slice = foreg_slices[random.randrange(len(foreg_slices))]
low_x, high_x = adjust(foreg_slice, self.patch_size, label, 1)
low_y, high_y = adjust(foreg_slice, self.patch_size, label, 2)
low_z, high_z = adjust(foreg_slice, self.patch_size, label, 3)
image = image[:, low_x:high_x, low_y:high_y, low_z:high_z]
label = label[:, low_x:high_x, low_y:high_y, low_z:high_z]
return image, label, [low_x, high_x, low_y, high_y, low_z, high_z]
class RandFlip:
def __init__(self):
self.axis = [1, 2, 3]
self.prob = 1 / len(self.axis)
def flip(self, data, axis):
data["image"] = np.flip(data["image"], axis=axis).copy()
data["label"] = np.flip(data["label"], axis=axis).copy()
return data
def __call__(self, data):
for axis in self.axis:
if random.random() < self.prob:
data = self.flip(data, axis)
return data
class Cast:
def __init__(self, types):
self.types = types
def __call__(self, data):
data["image"] = data["image"].astype(self.types[0])
data["label"] = data["label"].astype(self.types[1])
return data
class RandomBrightnessAugmentation:
def __init__(self, factor, prob):
self.prob = prob
self.factor = factor
def __call__(self, data):
image = data["image"]
if random.random() < self.prob:
factor = np.random.uniform(low=1.0-self.factor, high=1.0+self.factor, size=1)
image = (image * (1 + factor)).astype(image.dtype)
data.update({"image": image})
return data
class GaussianNoise:
def __init__(self, mean, std, prob):
self.mean = mean
self.std = std
self.prob = prob
def __call__(self, data):
image = data["image"]
if random.random() < self.prob:
scale = np.random.uniform(low=0.0, high=self.std)
noise = np.random.normal(loc=self.mean, scale=scale, size=image.shape).astype(image.dtype)
data.update({"image": image + noise})
return data
class PytTrain(Dataset):
def __init__(self, images, labels, **kwargs):
self.images, self.labels = images, labels
self.train_transforms = get_train_transforms()
patch_size, oversampling = kwargs["patch_size"], kwargs["oversampling"]
self.patch_size = patch_size
self.rand_crop = RandBalancedCrop(patch_size=patch_size, oversampling=oversampling)
def __len__(self):
return len(self.images)
def __getitem__(self, idx):
data = {"image": np.load(self.images[idx]), "label": np.load(self.labels[idx])}
data = self.rand_crop(data)
data = self.train_transforms(data)
return data["image"], data["label"]
class PytVal(Dataset):
def __init__(self, images, labels):
self.images, self.labels = images, labels
def __len__(self):
return len(self.images)
def __getitem__(self, idx):
return np.load(self.images[idx]), np.load(self.labels[idx])

BIN
artifacts/package_sources/tinygrad/test/external/openpilot/b1ab7897cbfa35981e1636fe551e4ce5.npy vendored Normal file
View File

Binary file not shown.

19
artifacts/package_sources/tinygrad/test/external/process_replay/README.md vendored Normal file
View File

@@ -0,0 +1,19 @@
# Process replay tests
Process replay is a tool for creating a diff of generated kernels between two commits. By default, process replay doesn't assert kernel diffs.
Refactor and speedup PRs must enable the assert by including `[pr]` in the pull request title.
Note that process replay [early stops when over 20% of kernels change, for speed.](https://github.com/tinygrad/tinygrad/pull/5480).
## Running locally
To run process replay locally:
(optional: clear previous process replay runs with `test/external/process_replay/reset.py`)
1. Run tests with `CAPTURE_PROCESS_REPLAY=1` in your branch. This will pickle process inputs to CACHEDB.
2. Checkout master
3. Run `test/external/process_replay/process_replay.py`
For reference, see `test/external/process_replay/local.sh`.

10
artifacts/package_sources/tinygrad/test/external/process_replay/local.sh vendored Executable file
View File

@@ -0,0 +1,10 @@
#!/bin/bash
set -e
HEAD=$(git rev-parse --abbrev-ref HEAD)
python test/external/process_replay/reset.py
CAPTURE_PROCESS_REPLAY=1 python test/backend/test_ops.py TestOps.test_add
git checkout master
git checkout $HEAD -- test/external/process_replay/process_replay.py
ASSERT_PROCESS_REPLAY=${ASSERT_PROCESS_REPLAY:-1} python test/external/process_replay/process_replay.py
git checkout $HEAD

128
artifacts/package_sources/tinygrad/test/external/process_replay/process_replay.py vendored Executable file
View File

@@ -0,0 +1,128 @@
#!/usr/bin/env python3
# compare kernels created by HEAD against master
import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, functools, base64, codecs
from dataclasses import replace
from typing import Callable, Any
ASSERT_DIFF = int((flag:="[PR]") in os.getenv("COMMIT_MESSAGE", flag) or flag in os.getenv("PR_TITLE", flag))
if not int(os.getenv("ASSERT_PROCESS_REPLAY", "1")): ASSERT_DIFF = 0
try:
from tinygrad.renderer import Renderer
from tinygrad.codegen import to_program
from tinygrad.uop.ops import UOp, Ops
from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm
except ImportError as e:
print(repr(e))
exit(int(ASSERT_DIFF))
# *** process replay settings
# internal
PAGE_SIZE = getenv("PAGE_SIZE", 100)
REF = os.getenv("GITHUB_REF_NAME", "")
MAX_DIFF_PCT = getenv("PROCESS_REPLAY_MAX_DIFF_PCT", 20)
TABLE_NAME = f"process_replay_{VERSION}"
os.environ["CAPTURE_PROCESS_REPLAY"] = "0"
early_stop = multiprocessing.Event()
logging.basicConfig(level=logging.INFO, format="%(message)s")
MAX_LINES = 500
def trunc_log(x):
if len(lines:=(x if isinstance(x, str) else repr(x)).splitlines()) > MAX_LINES:
lines = lines[:MAX_LINES]+[f"WARN: truncated string with {len(lines)} lines"]
logging.info("\n".join(lines))
# user config
SKIP_PROCESS_REPLAY = (k:="[skip_process_replay]") in os.getenv("COMMIT_MESSAGE", "") or k in os.getenv("PR_TITLE", "")
# uncomment this to disable by default
#SKIP_PROCESS_REPLAY = not ASSERT_DIFF and not ((k:="[p]") in os.getenv("COMMIT_MESSAGE", "") or k in os.getenv("PR_TITLE", ""))
if REF == "master": SKIP_PROCESS_REPLAY = True
class ProcessReplayWarning(Warning): pass
# *** replay the function and convert return values to string
def replay_to_program(p:UOp, ast:UOp, renderer:Renderer) -> tuple[str, str, tuple[Any, ...]]:
if ast.op is Ops.PROGRAM: input_ast = ast
else:
sink_arg = ast.arg
if sink_arg.beam: sink_arg = replace(sink_arg, opts_to_apply=p.src[0].arg.applied_opts)
input_ast = ast.replace(arg=replace(sink_arg, name=p.src[0].arg.name))
p2 = to_program(input_ast, renderer=renderer)
device = renderer.target.device
def to_str(ret:UOp) -> str:
src = next(x.arg for x in ret.src if x.op is Ops.SOURCE)
# PYTHON renderer pickles UOps, first unpickle and decode here
if device.startswith("PYTHON"): return "\n".join([str(x) for x in pickle.loads(base64.b64decode(src))])
return src
# properly color the name arg
ast_repr = codecs.decode(str(input_ast), "unicode_escape")
return to_str(p2), to_str(p), (ast_repr, renderer)
replayers: dict[str, Callable[..., tuple[str, str, tuple[Any, ...]]]] = {}
replayers["do_to_program"] = replay_to_program
# *** run replayers on captured rows and print diffs
def diff(offset:int, fxns:dict[str, Callable[..., tuple|None]]) -> None:
if ASSERT_DIFF: warnings.filterwarnings("error", category=ProcessReplayWarning)
if early_stop.is_set(): return None
conn = db_connection()
cur = conn.cursor()
cur.execute(f"SELECT val FROM '{TABLE_NAME}' LIMIT ? OFFSET ?", (PAGE_SIZE, offset))
changed = 0
for row in cur.fetchall():
if changed > MAX_DIFF_PCT:
warnings.warn(f"detected changes in over {MAX_DIFF_PCT}%. skipping further diff generation.", ProcessReplayWarning)
early_stop.set()
break
name, loc = "", ""
try:
name, args, kwargs, ctx_vals, loc, ret = pickle.loads(row[0])
ctx_vars = {k:v.value for k,v in ctx_vals.items() if k not in ("DEBUG", "CAPTURE_PROCESS_REPLAY")
and (var:=ContextVar._cache.get(k)) is not None and var.value != v.value}
if (replayer:=fxns.get(name)) is None: continue
with Context(**ctx_vars):
if (ret:=replayer(ret, *args, **kwargs)) is None: continue
good, compare, metadata = ret
if good != compare:
for m in metadata: trunc_log(m)
logging.info(loc)
for line in difflib.unified_diff(good.splitlines(), compare.splitlines()):
logging.info(colored(line, "red" if line.startswith("-") else "green" if line.startswith("+") else None))
if ctx_vars: logging.info(ctx_vars)
warnings.warn("PROCESS REPLAY DETECTED CHANGE", ProcessReplayWarning)
except Exception as e:
changed += 1
warnings.warn(f"{name=} {loc=} {e=}", ProcessReplayWarning)
cur.close()
# *** generic runner to map rows of a table to a function in parallel
def _pmap(fxns:dict[str, Callable]) -> None:
conn = db_connection()
cur = conn.cursor()
try: row_count = cur.execute(f"select count(*) from '{TABLE_NAME}'").fetchone()[0]
except sqlite3.OperationalError:
raise RuntimeError(f"{TABLE_NAME} isn't accessible in master, did DB_VERSION change?")
finally:
cur.close()
with multiprocessing.get_context("spawn").Pool(multiprocessing.cpu_count()) as pool:
bar = tqdm(total=row_count)
for _ in pool.imap_unordered(functools.partial(diff, fxns=fxns), range(0, row_count, s:=min(PAGE_SIZE, row_count))): bar.update(s)
pool.close()
pool.join()
pool.terminate()
# *** main loop
if __name__ == "__main__":
if SKIP_PROCESS_REPLAY:
logging.info("skipping process replay.")
exit(0)
logging.info(f"running process replay with {ASSERT_DIFF=}")
try: _pmap(replayers)
except Exception as e:
logging.info(f"process replay err: {e}")
exit(int(ASSERT_DIFF))

Some files were not shown because too many files have changed in this diff Show More