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

IQ.Pilot Release Commit @ bec7652

Browse Source
This commit is contained in:
IQ.Lvbs history cleanup
2026-08-22 23:42:41 -05:00
commit 58039e647c
4603 changed files with 1236178 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
artifacts/package_sources/tinygrad/test/backend/__init__.py Normal file
View File

269
artifacts/package_sources/tinygrad/test/backend/test_arange.py Normal file
View File

@@ -0,0 +1,269 @@
import unittest
import numpy as np
from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
from tinygrad.helpers import Context, getenv, DEV
from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear
from tinygrad.renderer.ptx import PTXRenderer
from test.helpers import needs_second_gpu, check_schedule, assert_kernel_count, KernelCountException
class TestArange(unittest.TestCase):
def _get_flops(self, tensor, desired):
GlobalCounters.reset()
linear = compile_linear(tensor.schedule_linear())
if len(linear.src) != 1: raise KernelCountException(1, len(linear.src))
run_linear(linear)
np.testing.assert_equal(tensor.numpy(), desired)
return estimate_uop(linear.src[-1]).ops
def test_arange_complexity(self):
self.assertLess(self._get_flops(Tensor.arange(256).clone(), np.arange(256)), 256*4)
self.assertLess(self._get_flops(Tensor.arange(2560).clone(), np.arange(2560)), 2560*4)
def test_cat_complexity(self):
x = Tensor.arange(2**10) + Tensor.empty((), dtype=dtypes.uint32)
out = x.cat(x).cat(Tensor.empty(1, dtype=dtypes.uint32))
linear = compile_linear(out.schedule_linear())
self.assertLessEqual(estimate_uop(linear.src[-1]).ops, out.numel()*20)
@unittest.skipIf(Device.DEFAULT == "CL", "flaky in CI")
def test_arange_cumsum(self):
np.testing.assert_equal(Tensor.arange(513).cumsum(0).numpy(), np.arange(513).cumsum())
def test_arange_cat(self):
t = Tensor.arange(2, dtype=dtypes.int)+Tensor([3])
self.assertEqual(t.cat(t).tolist(), [3, 4, 3, 4])
def test_eye_complexity(self):
with Context(NOOPT=1):
# NOTE: not every backend supports CMPEQ
self.assertLessEqual(self._get_flops(Tensor.eye(2560).clone(), np.eye(2560)), 2*2560*2560)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX indexing is weird")
def test_tri_complexity(self):
with Context(NOOPT=1):
t = Tensor.ones(256, 256).contiguous().realize()
linear = compile_linear(t.triu().schedule_linear())
self.assertLessEqual(estimate_uop(linear.src[-1]).ops, 4 * 256 * 256)
DSET, DDIM = 2048, 32
class TestIndexing(unittest.TestCase):
def test_arange_2_reduce(self):
needle = Tensor.zeros(16384, dtype=dtypes.int).contiguous()
needle[1337] = 1
needle.realize()
with Context(NOOPT=1):
GlobalCounters.reset()
out = ((Tensor.arange(1,16385)-1)*needle).sum()
linear, var_vals = check_schedule(out, 1)
run_linear(linear, var_vals)
self.assertEqual(out.item(), 1337)
def test_manual_index(self):
dataset = Tensor.rand(DSET, DDIM).realize()
idxs = Tensor([0,3,5,6]).realize()
real_index = dataset.numpy()[idxs.numpy()]
print("*** indexing ***")
with Context(NOOPT=1):
GlobalCounters.reset()
rng = Tensor.arange(DSET, dtype=dtypes.int).reshape(1, 1, DSET, 1).expand(4, DDIM, DSET, 1)
idxs = idxs.reshape(4,1,1,1).expand(4, DDIM, DSET, 1)
reshape_dataset = dataset.T.reshape(1, DDIM, DSET, 1).expand(4, DDIM, DSET, 1)
full = (rng==idxs).where(reshape_dataset, Tensor.zeros(4, DDIM, DSET, 1, buffer=False))
X = full.sum(axis=(2,3))
linear, var_vals = check_schedule(X, 1)
run_linear(linear, var_vals)
assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}"
np.testing.assert_allclose(real_index, X.numpy())
def test_index_variable(self):
dataset = Tensor.rand(DSET, DDIM).realize()
v = Variable("v", 0, DDIM-1)
with Context(NOOPT=1):
GlobalCounters.reset()
vb = Tensor(v.bind(12))
comp = dataset[vb].numpy()
# no global ops because they are all indexing
self.assertLess(GlobalCounters.global_ops, 1000)
np.testing.assert_allclose(comp, dataset.numpy()[12])
def test_index(self):
dataset = Tensor.rand(DSET, DDIM).realize()
idxs = Tensor([0,3,5,6]).realize()
real_index = dataset.numpy()[idxs.numpy()]
print("*** indexing ***")
with Context(NOOPT=1):
GlobalCounters.reset()
X = dataset[idxs]
assert X.shape == (4,DDIM)
linear, var_vals = check_schedule(X, 1)
run_linear(linear, var_vals)
assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}"
np.testing.assert_allclose(real_index, X.numpy())
def test_index_fused(self, noopt=1):
dataset = Tensor.rand(DSET, DDIM).realize()
idxs = Tensor([0,3,5,6]).realize()
real_index = dataset.numpy()[idxs.numpy()]
print("*** indexing ***")
with Context(NOOPT=noopt):
GlobalCounters.reset()
X = dataset[idxs]
assert X.shape == (4,DDIM)
linear, var_vals = check_schedule(X, 1)
run_linear(linear, var_vals)
assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops} != {4*DSET}"
np.testing.assert_allclose(real_index, X.numpy())
@unittest.skip("not ready")
def test_index_fused_opt(self): self.test_index_fused(0)
@unittest.skipIf(Device.DEFAULT == "CL", "rusticl/llvmpipe bug: https://gitlab.freedesktop.org/mesa/mesa/-/work_items/15667")
def test_index_fused_out_of_bounds(self):
dataset = Tensor.rand(256, 256).realize()
idxs = Tensor([-19238, -257, 256, 495, 10982377]).realize()
with Context(NOOPT=1):
X = dataset[idxs]
np.testing.assert_equal(X.numpy(), 0)
def test_index_mnist(self, noopt=1, op_limit=512*784*13, split_reduceop=0):
# WEBGPU generates more ops due to bitpacking of < 4-byte dtypes
if Device.DEFAULT == "WEBGPU": op_limit *= 15
# from tinygrad.nn.datasets import mnist
X_train, Y_train = Tensor.randint(DSET, 1, 28, 28, dtype='uchar').realize(), Tensor.randint(DSET, dtype='uchar').realize()
with Context(NOOPT=noopt, SPLIT_REDUCEOP=split_reduceop):
samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0]).realize()
GlobalCounters.reset()
x = X_train[samples].numpy()
y = Y_train[samples].numpy()
assert GlobalCounters.global_ops < op_limit, f"too many ops {GlobalCounters.global_ops} != {op_limit}"
np.testing.assert_allclose(X_train.numpy()[samples.numpy()], x)
np.testing.assert_allclose(Y_train.numpy()[samples.numpy()], y)
def test_index_mnist_opt(self): self.test_index_mnist(0)
def test_index_mnist_split(self): self.test_index_mnist(1, split_reduceop=1)
def test_index_mnist_opt_split(self): self.test_index_mnist(0, split_reduceop=1)
def test_llama_embedding(self, noopt=1, op_limit=65536):
# llama3 is 128256
vocab_size, embed_size = (10, 3)
emb = nn.Embedding(vocab_size, embed_size)
emb_w = emb.weight.numpy()
x = Tensor([1,2,3,4])
with Context(NOOPT=noopt):
GlobalCounters.reset()
z = emb(x).realize()
self.assertLessEqual(GlobalCounters.global_ops, op_limit)
assert_kernel_count(2)
if getenv("CHECK", 1):
import torch
with torch.no_grad():
torch_emb = torch.nn.Embedding(vocab_size, embed_size).eval()
torch_emb.weight[:] = torch.tensor(emb_w, dtype=torch.float32)
torch_z = torch_emb(torch.tensor(x.numpy()))
# TODO: reshape to match torch, should we do this in nn?
np.testing.assert_allclose(z.numpy().reshape(4, embed_size), torch_z.detach().numpy(), atol=1e-8, rtol=1e-8)
# at least the arange is being fused
def test_llama_embedding_opt(self): self.test_llama_embedding(0, 1_736_704_000)
# NOTE: call doesn't work with SPEC=2
@unittest.skipIf(Device.DEFAULT not in ("CPU", "AMD"), "atomics only on AMD/CPU")
@Context(USE_ATOMICS=1, SPEC=1)
def test_llama_8b_embedding_backward(self):
from tinygrad.renderer.cstyle import CStyleLanguage
if Device.DEFAULT == "CPU" and not isinstance(Device["CPU"].renderer, CStyleLanguage): self.skipTest("CPU needs Clang renderer")
vocab_size, embed_size = 1000, 128
bs, seqlen = 4, 256
idx = Tensor.randint(bs, seqlen, high=vocab_size)
emb = nn.Embedding(vocab_size, embed_size)
emb.weight = Tensor.ones(vocab_size, embed_size)
gt = Tensor.zeros(bs, seqlen, embed_size)
Tensor.realize(idx, emb.weight, gt)
GlobalCounters.reset()
loss = (emb(idx)-gt).square().sum()
loss.backward()
emb.weight.grad.realize()
bwd_ops = GlobalCounters.global_ops
print(f"embedding bwd: {GlobalCounters.kernel_count} kernels, {bwd_ops:,} ops")
self.assertLess(bwd_ops, bs*seqlen*embed_size*20, f"backward ops {bwd_ops:,} should be less than 20 per with atomic scatter-add")
# correctness check
expected_grad = np.zeros((vocab_size, embed_size), dtype=np.float32)
for i in idx.flatten().numpy(): expected_grad[i] += 2
np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)
@unittest.skipIf(Device.DEFAULT not in ("CPU", "AMD"), "atomics only on AMD/CPU")
@Context(USE_ATOMICS=1, SPEC=1)
def test_embedding_backward_padded_embed(self):
from tinygrad.renderer.cstyle import CStyleLanguage
if Device.DEFAULT == "CPU" and not isinstance(Device["CPU"].renderer, CStyleLanguage): self.skipTest("CPU needs Clang renderer")
vocab_size, embed_size = 1000, 300
bs, seqlen = 4, 256
idx = Tensor.randint(bs, seqlen, high=vocab_size)
emb = nn.Embedding(vocab_size, embed_size)
emb.weight = Tensor.ones(vocab_size, embed_size)
gt = Tensor.zeros(bs, seqlen, embed_size)
Tensor.realize(idx, emb.weight, gt)
loss = (emb(idx)-gt).square().sum()
loss.backward()
emb.weight.grad.realize()
# correctness check
expected_grad = np.zeros((vocab_size, embed_size), dtype=np.float32)
for i in idx.flatten().numpy(): expected_grad[i] += 2
np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)
@needs_second_gpu
@unittest.skipIf(Device.DEFAULT not in ("CPU", "AMD"), "atomics only on AMD/CPU")
@Context(USE_ATOMICS=1, SPEC=1)
def test_embedding_backward_vocab_sharded(self):
from tinygrad.renderer.cstyle import CStyleLanguage
if Device.DEFAULT == "CPU" and not isinstance(Device["CPU"].renderer, CStyleLanguage): self.skipTest("CPU needs Clang renderer")
devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
vocab_size, embed_size = 1000, 128
bs, seqlen = 4, 256
idx = Tensor.randint(bs, seqlen, high=vocab_size)
emb = nn.Embedding(vocab_size, embed_size)
emb.weight = Tensor.ones(vocab_size, embed_size)
gt = Tensor.zeros(bs, seqlen, embed_size)
Tensor.realize(idx, emb.weight, gt)
# compute expected grad on single device
expected_grad = np.zeros((vocab_size, embed_size), dtype=np.float32)
for i in idx.flatten().numpy(): expected_grad[i] += 2
# now shard the embedding weight on vocab axis and recompute
emb.weight = Tensor.ones(vocab_size, embed_size)
emb.weight.shard_(devices, axis=0)
idx = idx.shard(devices, axis=None)
gt = gt.shard(devices, axis=None)
Tensor.realize(idx, emb.weight, gt)
loss = (emb(idx)-gt).square().sum()
loss.backward()
np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)
@unittest.skipUnless(Device.DEFAULT == "AMD" or (Device.DEFAULT == "NULL" and DEV.arch.startswith("gfx")), "tests AMD bf16 cast overhead")
def base_test_llama_8b_rope_backward(self, dtype, ops_scale=1):
from extra.models.llama import precompute_freqs_cis, apply_rotary_emb
bs, seqlen, dim, n_heads = 1, 512, 256, 4
head_dim = dim // n_heads
x = Tensor.randn(bs, seqlen, dim, dtype=dtype)
wq = Tensor.randn(dim, dim, dtype=dtype)
freqs_cis = precompute_freqs_cis(head_dim, seqlen).cast(dtype)
Tensor.realize(x, wq, freqs_cis)
xq = (x @ wq.T)
# main llama does not fuse it
#xq = xq.contiguous_backward()
xq = xq.reshape(bs, seqlen, n_heads, head_dim)
xq_rope, _ = apply_rotary_emb(xq, xq, freqs_cis)
xq_rope.sum().backward()
linear = compile_linear(wq.grad.schedule_linear())
if len(linear.src) != 1: raise KernelCountException(1, len(linear.src))
bwd_ops = estimate_uop(linear.src[0]).ops
expected_ops = bs*seqlen*dim*dim*ops_scale
print(f"rope matmul bwd ({dtype}): {GlobalCounters.kernel_count} kernels, {bwd_ops:,} ops")
self.assertLess(bwd_ops, expected_ops, f"rope bwd ops {bwd_ops:,} should be < {ops_scale} per (got {bwd_ops/(bs*seqlen*dim*dim):.1f})")
def test_llama_8b_rope_backward_f16(self):
self.base_test_llama_8b_rope_backward(dtypes.float16, ops_scale=2)
# bfloat16 on non CDNA4 has ~10x ops overhead because of the software emulation
def test_llama_8b_rope_backward_bf16(self):
self.base_test_llama_8b_rope_backward(dtypes.bfloat16, ops_scale=2 if Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950") else 25)
if __name__ == "__main__":
unittest.main()

426
artifacts/package_sources/tinygrad/test/backend/test_asm_gemm.py Normal file
View File

@@ -0,0 +1,426 @@
import unittest
from tinygrad import Tensor, Device, dtypes, Context
from tinygrad.helpers import getenv, system, DEV
from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm
from test.helpers import needs_second_gpu
from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX
# On non CDNA4 it will only validate the Tensor.custom_kernel integration
# Use DEV=NULL:HIP:gfx950 to also test the assembly
def is_cdna4(): return Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950")
def has_hipcc():
try: system("hipcc --version")
except Exception: return False
return True
def run_asm_gemm(a_shape, b_shape, dtype=dtypes.bfloat16, a_shard=None, b_shard=None, gpus:int=1) -> None:
Tensor.manual_seed(0)
input_dtype = dtypes.bfloat16 if dtype == FP8_DTYPE else dtype
a_rand = Tensor.randn(a_shape, dtype=dtypes.float).sub(0.5).cast(input_dtype)
b_rand = Tensor.randn(b_shape, dtype=dtypes.float).sub(0.5).cast(input_dtype)
with Context(DEBUG=0):
Tensor.realize(a_rand, b_rand)
devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(gpus)) if (multi:=gpus>1) else None
if dtype == FP8_DTYPE:
x_scale = Tensor.full((), FP8_MAX, dtype=dtypes.float32, device=devs).contiguous()
a_rand, _, _ = quantize_fp8(a_rand.shard(devs, axis=a_shard) if multi else a_rand, amax_state=x_scale)
b_rand, w_scale, _ = quantize_fp8(b_rand.T.contiguous())
if multi: b_rand, w_scale = b_rand.shard(devs, axis=None if b_shard is None else 1-b_shard), w_scale.to(devs).contiguous()
grad_amax_state = Tensor.full((), FP8_MAX, dtype=dtypes.float32, device=devs).contiguous()
next_grad_amax_state = Tensor.empty((), dtype=dtypes.float32, device=devs)
with Context(DEBUG=0):
Tensor.realize(a_rand, x_scale, b_rand, w_scale, grad_amax_state, next_grad_amax_state)
# clone all inputs before any backward: a clone copies the source's current .grad
a, b = a_rand.clone(), b_rand.clone()
if dtype == FP8_DTYPE:
a_ref, b_ref = a_rand.detach().cast(dtypes.bfloat16), b_rand.detach().cast(dtypes.bfloat16)
else:
a_ref, b_ref = a_rand.clone(), b_rand.clone()
if multi and isinstance(a.device, str): a, b = a.shard(devs, axis=a_shard), b.shard(devs, axis=b_shard)
if dtype == FP8_DTYPE:
tst = asm_gemm(a, b.T, x_scale=x_scale, w_scale=w_scale, grad_amax_state=grad_amax_state,
next_grad_amax_state=next_grad_amax_state)
else:
tst = asm_gemm(a, b)
tst.sum().backward()
Tensor.realize(tst, a.grad, b.grad)
if multi and isinstance(a_ref.device, str): a_ref, b_ref = a_ref.shard(devs, axis=a_shard), b_ref.shard(devs, axis=b_shard)
if dtype == FP8_DTYPE:
ref = ((a_ref @ b_ref.T) * ((x_scale.float() + 1e-8) / FP8_MAX) * w_scale).cast(dtypes.bfloat16)
else:
ref = a_ref @ b_ref
ref.sum().backward()
Tensor.realize(ref, a_ref.grad, b_ref.grad)
# no validation on the NULL device
if Device.DEFAULT.startswith("NULL"): return None
atol, rtol = (2e-1, 1e-2) if dtype == dtypes.bfloat16 else (256, 1e-2) if dtype == FP8_DTYPE else (1e-2, 1e-3)
# allow more rtol for multi because of ALLREDUCE_CAST
grad_atol, grad_rtol = (16895, 0.125) if dtype == FP8_DTYPE else (atol, 2e-2 if multi else rtol)
with Context(DEBUG=0):
# enable for debugging, slow for larger gemms
if getenv("USE_NPY"):
import numpy as np
np.testing.assert_allclose(tst.numpy(), ref.numpy(), atol=atol, rtol=rtol)
np.testing.assert_allclose(a.grad.numpy(), a_ref.grad.numpy(), atol=grad_atol, rtol=grad_rtol)
np.testing.assert_allclose(b.grad.numpy(), b_ref.grad.numpy(), atol=grad_atol, rtol=grad_rtol)
assert tst.allclose(ref, atol=atol, rtol=rtol).item(), "forward mismatch"
assert a.grad.allclose(a_ref.grad, atol=grad_atol, rtol=grad_rtol).item(), "grad_a mismatch"
assert b.grad.allclose(b_ref.grad, atol=grad_atol, rtol=grad_rtol).item(), "grad_b mismatch"
def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.bfloat16, gpus:int=1) -> None:
run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=0, b_shard=None, gpus=gpus)
def verify_asm_gemm_k_sharded(M:int, N:int, K:int, dtype=dtypes.bfloat16, gpus:int=8) -> None:
run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=1, b_shard=0, gpus=gpus)
def verify_asm_gemm_n_sharded(batch:int, M:int, N:int, K:int, dtype=dtypes.bfloat16, gpus:int=2) -> None:
run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=None, b_shard=1, gpus=gpus)
def verify_asm_gemm_m_sharded(M:int, N:int, K:int, dtype=dtypes.bfloat16, gpus:int=2) -> None:
run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=0, b_shard=None, gpus=gpus)
def verify_asm_gemm_n_sharded_2d(M:int, N:int, K:int, dtype=dtypes.bfloat16, gpus:int=2) -> None:
run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=None, b_shard=1, gpus=gpus)
def verify_asm_gemm_k_sharded_3d(batch:int, M:int, N:int, K:int, dtype=dtypes.bfloat16, gpus:int=2) -> None:
run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=2, b_shard=0, gpus=gpus)
# 128x smaller than usual
# uses the UOp GEMM, runs on non CDNA4 and CI
@unittest.skipUnless(dtypes.bfloat16 in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
class TestGemm(unittest.TestCase):
def setUp(self):
if is_cdna4(): self.skipTest("shapes are too small for the assembly GEMM")
def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 32), N, N, dtype=dtypes.bfloat16)
def test_gemm(self): verify_asm_gemm(1, 64, 32, 112)
def test_gemm_batched(self): verify_asm_gemm(2, 64, 32, 32)
@needs_second_gpu
def test_gemm_multi(self): verify_asm_gemm(2, 64, 32, 32, gpus=2)
@needs_second_gpu
def test_gemm_k_sharded(self): verify_asm_gemm_k_sharded(64, 64, 2*64, gpus=2)
@needs_second_gpu
def test_gemm_m_sharded(self): verify_asm_gemm_m_sharded(2*64, 64, 32, gpus=2)
@needs_second_gpu
def test_gemm_n_sharded(self): verify_asm_gemm_n_sharded(1, 64, 64, 32, gpus=2)
@needs_second_gpu
def test_gemm_n_sharded_2d(self): verify_asm_gemm_n_sharded_2d(64, 2*64, 32, gpus=2)
@needs_second_gpu
def test_gemm_k_sharded_3d(self): verify_asm_gemm_k_sharded_3d(1, 64, 32, 2*64, gpus=2)
# uses the smallest size for the cdna assembly gemm
class TestAsmGEMM(unittest.TestCase):
def setUp(self):
if not is_cdna4() or not has_hipcc():
self.skipTest("assembly gemm is only for cdna4")
def test_tiny(self): verify_asm_gemm(1, 256, 256, 256)
def test_verify_with_numpy(self):
import numpy as np
M, N, K = 256, 256, 256
rng = np.random.default_rng(0)
a_np = (rng.random((M, K), dtype=np.float32) - 0.5).astype(np.float32)
b_np = (rng.random((K, N), dtype=np.float32) - 0.5).astype(np.float32)
c_np = (a_np.astype(np.float32) @ b_np.astype(np.float32)).astype(np.float32)
Tensor.manual_seed(0)
a, b = Tensor(a_np).cast(dtypes.bfloat16), Tensor(b_np).cast(dtypes.bfloat16)
c = asm_gemm(a, b)
c.realize()
# no validation on the NULL device
if a.device.startswith("NULL"): return None
np.testing.assert_allclose(c.numpy(), c_np, atol=2e-1, rtol=1e-2)
def test_unsupported_batch(self):
with self.assertRaisesRegex(AssertionError, "batch size"):
verify_asm_gemm(3, 256, 256, 256)
def test_unsupported_k(self):
with self.assertRaisesRegex(AssertionError, "not a multiple"):
verify_asm_gemm(1, 1024, 1024, 100)
def test_unsupported_m(self):
with self.assertRaisesRegex(AssertionError, "not a multiple"):
verify_asm_gemm(1, 1000, 256, 256)
def test_unsupported_n(self):
with self.assertRaisesRegex(AssertionError, "not a multiple"):
verify_asm_gemm(1, 256, 1000, 256)
class TestMXFP4(unittest.TestCase):
def setUp(self):
if not is_cdna4() or DEV.interface.startswith("MOCK"):
self.skipTest("requires real amd machine")
def test_quantize(self):
import numpy as np
from extra.llama_kernels.quantize_mxfp4 import quantize_mxfp4
rng = np.random.default_rng(0)
x = np.triu(rng.standard_normal((256, 256), dtype=np.float32))
x += np.triu(x, 1).T
x[:32, :32] = 0
row, row_scale, col, col_scale = quantize_mxfp4(Tensor(x, dtype=dtypes.bfloat16))
Tensor.realize(row, row_scale, col, col_scale)
row, row_scale = row.numpy(), row_scale.numpy()
col, col_scale = col.numpy(), col_scale.numpy()
np.testing.assert_array_equal(row, col)
np.testing.assert_array_equal(row_scale, col_scale)
self.assertTrue(row.any())
self.assertTrue((row_scale == 127).any())
self.assertTrue((row_scale != 127).any())
def test_correctness(self):
import numpy as np
M = N = K = 256
rng = np.random.default_rng(1)
a = Tensor(rng.standard_normal((M, K), dtype=np.float32), dtype=dtypes.bfloat16)
b = Tensor(rng.standard_normal((N, K), dtype=np.float32), dtype=dtypes.bfloat16)
out = asm_gemm(a, b.T, mxfp4=True).realize().numpy().astype(np.float32)
ref = a.numpy().astype(np.float32) @ b.numpy().astype(np.float32).T
self.assertLess(np.linalg.norm(out-ref) / np.linalg.norm(ref), 0.2)
def test_empty(self):
M, N, K = getenv("M", 16384), getenv("N", 4096), getenv("K", 14336)
a = Tensor.empty(M, K, dtype=dtypes.bfloat16)
b = Tensor.empty(N, K, dtype=dtypes.bfloat16)
asm_gemm(a, b.T, mxfp4=True).realize()
# test the Asm GEMM with Llama shapes, only run on the real machine for speed
@unittest.skipUnless(has_hipcc(), "requires hipcc to compile")
class TestGemmLlama(unittest.TestCase):
dtype = FP8_DTYPE
def setUp(self):
if not is_cdna4() or DEV.interface.startswith("MOCK"):
self.skipTest("very slow on non mi350x")
def test_empty(self): asm_gemm(Tensor.empty(N:=getenv("N", 4096), N, dtype=self.dtype), Tensor.empty(N, N, dtype=self.dtype)).realize()
def test_empty_bw(self):
x = Tensor.empty(1, N:=getenv("N", 4096), N, dtype=self.dtype)
y = Tensor.empty((N, N), dtype=self.dtype)
if self.dtype == FP8_DTYPE:
x_scale = Tensor.empty((), dtype=dtypes.float32)
w_scale = Tensor.empty((), dtype=dtypes.float32)
grad_amax_state = Tensor.empty((), dtype=dtypes.float32).contiguous()
next_grad_amax_state = Tensor.empty((), dtype=dtypes.float32)
z = asm_gemm(x, y, x_scale=x_scale, w_scale=w_scale, grad_amax_state=grad_amax_state,
next_grad_amax_state=next_grad_amax_state)
else:
z = asm_gemm(x, y)
z.sum().backward()
Tensor.realize(z, x.grad, y.grad)
# FP8 GEMM stores bf16 output and its backward produces bf16 gradients.
grad_dtype = dtypes.bfloat16 if self.dtype == FP8_DTYPE else self.dtype
assert z.dtype == dtypes.bfloat16
assert x.grad.dtype == y.grad.dtype == grad_dtype
def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 4096), N, N, dtype=self.dtype)
def test_gemm(self): verify_asm_gemm(1, 8192, 4096, 14336, dtype=self.dtype)
def test_gemm_batched(self): verify_asm_gemm(2, 8192, 4096, 4096, dtype=self.dtype)
def test_gemm1(self): verify_asm_gemm(8, 8192, 4096, 14336, dtype=self.dtype, gpus=8)
def test_gemm2(self): verify_asm_gemm(8, 8192, 128256, 4096, dtype=self.dtype, gpus=8)
def test_gemm3(self): verify_asm_gemm(8, 8192, 14336, 4096, dtype=self.dtype, gpus=8)
def test_gemm4(self): verify_asm_gemm(8, 4096, 14336, 4096, dtype=self.dtype, gpus=8)
def test_gemm5(self): verify_asm_gemm(8, 4096, 4096, 14336, dtype=self.dtype, gpus=8)
def test_gemm6(self): verify_asm_gemm(16, 4096, 4096, 14336, dtype=self.dtype, gpus=8)
def test_gemm7(self): verify_asm_gemm(1, 8192, 128256, 4096, dtype=self.dtype)
def test_gemm8(self): verify_asm_gemm(1, 4096, 14336, 8192, dtype=self.dtype)
def test_gemm9(self): verify_asm_gemm(8, 4096, 14336, 8192, dtype=self.dtype, gpus=8)
def test_gemm10(self): verify_asm_gemm(1, 4096, 8192, 4096, dtype=self.dtype)
def test_gemm11(self): verify_asm_gemm(8, 1024, 1024, 4096, dtype=self.dtype, gpus=8)
def test_k_sharded_1(self): verify_asm_gemm_k_sharded(14336, 4096, 8*8192, dtype=self.dtype, gpus=8)
def test_k_sharded_2(self): verify_asm_gemm_k_sharded(4096, 14336, 8*8192, dtype=self.dtype, gpus=8)
def test_k_sharded_3(self): verify_asm_gemm_k_sharded(4096, 4096, 8*8192, dtype=self.dtype, gpus=8)
# M-sharded 2D
def test_m_sharded_1(self): verify_asm_gemm_m_sharded(8*8192, 4096, 4096, dtype=self.dtype, gpus=8)
def test_m_sharded_2(self): verify_asm_gemm_m_sharded(8*4096, 14336, 4096, dtype=self.dtype, gpus=8)
# N-sharded 2D
def test_n_sharded_2d_1(self): verify_asm_gemm_n_sharded_2d(8192, 8*4096, 4096, dtype=self.dtype, gpus=8)
def test_n_sharded_2d_2(self): verify_asm_gemm_n_sharded_2d(4096, 8*14336, 4096, dtype=self.dtype, gpus=8)
# tensor parallel shapes (Llama 8B, MP=8)
def test_tp_n_sharded_wq(self): verify_asm_gemm_n_sharded(1, 8192, 4096, 4096, dtype=self.dtype, gpus=8)
def test_tp_n_sharded_w1(self): verify_asm_gemm_n_sharded(1, 8192, 14336, 4096, dtype=self.dtype, gpus=8)
def test_tp_k_sharded_wo(self): verify_asm_gemm_k_sharded_3d(1, 8192, 4096, 4096, dtype=self.dtype, gpus=8)
def test_tp_k_sharded_w2(self): verify_asm_gemm_k_sharded_3d(1, 8192, 4096, 14336, dtype=self.dtype, gpus=8)
# more shapes: vary M, N, K independently
def test_shape_small_square(self): verify_asm_gemm(1, 256, 256, 256, dtype=self.dtype)
def test_shape_small_rect_m(self): verify_asm_gemm(1, 512, 256, 256, dtype=self.dtype)
def test_shape_small_rect_n(self): verify_asm_gemm(1, 256, 512, 256, dtype=self.dtype)
def test_shape_small_rect_k(self): verify_asm_gemm(1, 256, 256, 512, dtype=self.dtype)
def test_shape_tall(self): verify_asm_gemm(1, 2048, 256, 256, dtype=self.dtype)
def test_shape_wide(self): verify_asm_gemm(1, 256, 2048, 256, dtype=self.dtype)
def test_shape_deep(self): verify_asm_gemm(1, 256, 256, 4096, dtype=self.dtype)
def test_shape_non_square(self): verify_asm_gemm(1, 1024, 2048, 512, dtype=self.dtype)
def test_shape_batched_small(self): verify_asm_gemm(2, 256, 256, 256, dtype=self.dtype)
def test_shape_batched_rect(self): verify_asm_gemm(2, 512, 1024, 256, dtype=self.dtype)
# K edge cases: change iters to exercise different loop paths, k big enough for hk kernel
def test_shape_k256(self): verify_asm_gemm(1, 256, 256, 256, dtype=self.dtype)
def test_shape_k512(self): verify_asm_gemm(1, 256, 256, 512, dtype=self.dtype)
def test_shape_k768(self): verify_asm_gemm(1, 256, 256, 768, dtype=self.dtype)
def test_llama3_out1(self): verify_asm_gemm(1, 8192, 128256, 4096, dtype=self.dtype)
def test_llama3_out2(self): verify_asm_gemm(1, 8192, 4096, 128256, dtype=self.dtype)
def test_llama3_out3(self): verify_asm_gemm(1, 4096, 128256, 8192, dtype=self.dtype)
# mxfp8: 1x32 block scaling along K, e8m0 scales packed iteration-major (K/128, dim) uint32
def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
rows, K = x.shape
scale_K, k_iters = K // 32, K // 128
xb = x.reshape(rows, scale_K, 32).float()
amax = xb.abs().max(axis=-1)
e8 = (amax.log2().floor() + 127).clamp(0, 254)
e8 = (amax == 0).where(Tensor.zeros_like(e8), e8).cast(dtypes.uint8)
xq = (xb * (127.0 - e8.cast(dtypes.float32)).exp2().reshape(rows, scale_K, 1)).cast(FP8_DTYPE).reshape(rows, K)
packed = e8.reshape(rows, k_iters, 4).bitcast(dtypes.uint32).reshape(rows, k_iters).permute(1, 0)
return xq.contiguous(), e8, packed.contiguous()
def dequant_mxfp8(xq:Tensor, e8:Tensor) -> Tensor:
rows, K = xq.shape
scale = (e8.cast(dtypes.float32) - 127.0).exp2()
return (xq.float().reshape(rows, K // 32, 32) * scale.reshape(rows, K // 32, 1)).reshape(rows, K)
def run_mxfp8_gemm(M:int, N:int, K:int) -> None:
import functools
from extra.gemm.cdna_asm_gemm import custom_hk_mxfp8_gemm
Tensor.manual_seed(0)
a = (Tensor.randn(M, K, dtype=dtypes.float) * 0.5).realize()
b = (Tensor.randn(N, K, dtype=dtypes.float) * 0.5).realize()
a_q, a_e8, a_si = quantize_mxfp8(a)
b_q, b_e8, b_si = quantize_mxfp8(b)
Tensor.realize(a_q, a_e8, a_si, b_q, b_e8, b_si)
out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device)
tst = out.custom_kernel(a_q.unsqueeze(0), b_q, a_si, b_si, fxn=functools.partial(custom_hk_mxfp8_gemm, dname=a.device))[0].squeeze(0)
ref_mx = dequant_mxfp8(a_q, a_e8) @ dequant_mxfp8(b_q, b_e8).T
ref = a @ b.T
Tensor.realize(tst, ref_mx, ref)
if a.device.startswith("NULL"): return
err_mx = ((tst.float() - ref_mx).abs().mean() / ref_mx.abs().mean()).item()
err = ((tst.float() - ref).abs().mean() / ref.abs().mean()).item()
assert err_mx < 1e-2, f"kernel vs mxfp8 reference rel err {err_mx}"
assert err < 6e-2, f"kernel vs fp32 rel err {err}"
def run_mx_gemm_bw(M:int, N:int, K:int, w_post:bool=False) -> None:
Tensor.manual_seed(0)
a_rand = (Tensor.randn(M, K, dtype=dtypes.float) * 0.5).cast(dtypes.bfloat16).realize()
b_rand = (Tensor.randn(N, K, dtype=dtypes.float) * 0.5).cast(dtypes.bfloat16).realize()
w_post_scale = (Tensor.rand(N, dtype=dtypes.float) + 0.5).realize() if w_post else None
a, b, a_ref, b_ref = a_rand.clone(), b_rand.clone(), a_rand.clone(), b_rand.clone()
tst = asm_gemm(a, b.T, mx=True, w_post_scale=w_post_scale)
tst.sum().backward()
Tensor.realize(tst, a.grad, b.grad)
a_grad, b_grad = a.grad.float().contiguous().realize(), b.grad.float().contiguous().realize()
ref = a_ref.float() @ b_ref.float().T
if w_post is not None and w_post_scale is not None: ref = ref * w_post_scale.reshape(1, -1)
ref.sum().backward()
ref_b_grad = b_ref.grad / w_post_scale.reshape(-1, 1) if w_post_scale is not None else b_ref.grad
Tensor.realize(ref, a_ref.grad, b_ref.grad)
if a.device.startswith("NULL"): return
for name, t, r in [("fw", tst, ref), ("grad_a", a_grad, a_ref.grad), ("grad_b", b_grad, ref_b_grad)]:
err = ((t.float() - r.float()).abs().mean() / (r.float().abs().mean() + 1e-8)).item()
assert err < 6e-2, f"{name} rel err {err}"
def run_mx_gemm_multi(M:int, N:int, K:int, x_shard, w_shard, g_shard, gpus:int=2) -> None:
Tensor.manual_seed(0)
devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(gpus))
x_r = (Tensor.randn(M, K, dtype=dtypes.float) * 0.5).cast(dtypes.bfloat16).realize()
w_r = (Tensor.randn(N, K, dtype=dtypes.float) * 0.5).cast(dtypes.bfloat16).realize()
def run(shard):
x = (x_r.shard(devs, axis=x_shard) if shard else x_r.clone())
w = (w_r.shard(devs, axis=w_shard) if shard else w_r.clone())
out = asm_gemm(x, w.T, mx=True)
gmul = Tensor.ones(M, N).cast(dtypes.bfloat16)
(out.float() * (gmul.shard(devs, axis=g_shard) if shard else gmul).float()).sum().backward()
Tensor.realize(out, x.grad, w.grad)
to = (lambda t: t.to(Device.DEFAULT)) if shard else (lambda t: t)
return to(out).float().numpy(), to(x.grad).float().numpy(), to(w.grad).float().numpy()
ref = run(False)
if Device.DEFAULT.startswith("NULL"): return
got = run(True)
for name, g, r in zip(("fw", "grad_x", "grad_w"), got, ref):
err = ((abs(g - r)).mean() / (abs(r).mean() + 1e-8))
assert err < 2e-2, f"{name} sharded vs single rel err {err}"
def run_mx_prequant(M:int, N:int, K:int) -> None:
from extra.gemm.cdna_asm_gemm import quantize_mxfp8
Tensor.manual_seed(0)
x_rand = (Tensor.randn(M, K, dtype=dtypes.float) * 0.5).cast(dtypes.bfloat16).realize()
w_rand = (Tensor.randn(N, K, dtype=dtypes.float) * 0.5).cast(dtypes.bfloat16).realize()
x, w = x_rand.clone(), w_rand.clone()
x_q, x_e8, x_si = quantize_mxfp8(x)
w_q, w_e8, w_si = quantize_mxfp8(w)
out = asm_gemm(x_q, w_q.T, mx=True, mx_scales=(x_si, x_e8, w_si, w_e8))
out.sum().backward()
Tensor.realize(out, x.grad, w.grad)
if Device.DEFAULT.startswith("NULL"): return
ref_out, gx = x_rand.float() @ w_rand.float().T, w_rand.float().sum(0)
gw = x_rand.float().sum(0).reshape(1, K).expand(N, K)
for name, t, r in [("fw", out, ref_out), ("grad_x", x.grad, gx), ("grad_w", w.grad, gw)]:
err = ((t.float() - r.float()).abs().mean() / (r.float().abs().mean() + 1e-8)).item()
assert err < 6e-2, f"{name} prequant vs analytic rel err {err}"
@unittest.skipUnless(has_hipcc(), "requires hipcc to compile")
class TestGemmMXFP8(unittest.TestCase):
def setUp(self):
if not is_cdna4() or DEV.interface.startswith("MOCK"): self.skipTest("mxfp8 gemm is only for cdna4")
def test_prequant_simple(self): run_mx_prequant(256, 256, 256)
def test_prequant_rect(self): run_mx_prequant(512, 256, 512)
def test_simple(self): run_mxfp8_gemm(N:=getenv("N", 256), N, 2*128)
def test_rect(self): run_mxfp8_gemm(512, 256, 512)
def test_llama_ffn(self): run_mxfp8_gemm(8192, 14336, 4096)
def test_llama_ffn2(self): run_mxfp8_gemm(8192, 4096, 14336)
def test_llama_qkv(self): run_mxfp8_gemm(8192, 4096, 4096)
def test_general_n_fw(self):
for N in (256, 1792, 2048, 8192): run_mxfp8_gemm(8192, N, 4096)
# backward needs all dims tile-aligned (dgrad reduces N, wgrad reduces M)
def test_bw_simple(self): run_mx_gemm_bw(256, 256, 256)
def test_bw_rect(self): run_mx_gemm_bw(512, 256, 512)
def test_bw_w_post(self): run_mx_gemm_bw(256, 256, 256, w_post=True)
def test_bw_llama_qkv(self): run_mx_gemm_bw(8192, 4096, 4096)
def test_general_n_bw(self):
for N in (2048, 8192, 14336): run_mx_gemm_bw(8192, N, 4096)
# MP sharding: col-parallel (w on out axis), row-parallel (x,w on in axis)
@needs_second_gpu
def test_multi_col_parallel(self): run_mx_gemm_multi(512, 512, 512, x_shard=None, w_shard=0, g_shard=1)
@needs_second_gpu
def test_multi_row_parallel(self): run_mx_gemm_multi(512, 512, 512, x_shard=1, w_shard=1, g_shard=None)
@needs_second_gpu
def test_multi_data_parallel(self): run_mx_gemm_multi(512, 512, 512, x_shard=0, w_shard=None, g_shard=0)
def run_atb_gemm(rows, M, N, a_shard=None, b_shard=None, gpus=1, atol=1.0, rtol=3e-2) -> None:
import numpy as np
Tensor.manual_seed(0)
a = Tensor.randn(1, rows, M, dtype=dtypes.float).cast(dtypes.bfloat16)
b = Tensor.randn(1, rows, N, dtype=dtypes.float).cast(dtypes.bfloat16)
with Context(DEBUG=0): Tensor.realize(a, b)
ref = (a[0].float().transpose(0, 1) @ b[0].float()).realize() # [M, N]
if gpus > 1:
devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(gpus))
a, b = a.shard(devs, axis=a_shard), b.shard(devs, axis=b_shard)
out = hk_bf16_atb_gemm(a, b)
np.testing.assert_allclose(out.float().numpy(), ref.numpy(), atol=atol, rtol=rtol)
@unittest.skipUnless(has_hipcc(), "requires hipcc to compile")
class TestHkBf16AtbGemm(unittest.TestCase):
def setUp(self):
if not is_cdna4(): self.skipTest("hk bf16 atb gemm is cdna4 only")
def test_single(self): run_atb_gemm(256, 256, 256)
@needs_second_gpu
def test_k_sharded(self): run_atb_gemm(512, 256, 256, a_shard=1, b_shard=1, gpus=2)
@needs_second_gpu
def test_n_sharded(self): run_atb_gemm(256, 256, 512, a_shard=None, b_shard=2, gpus=2)
@needs_second_gpu
def test_m_sharded(self): run_atb_gemm(256, 512, 256, a_shard=2, b_shard=None, gpus=2)
if __name__ == "__main__":
unittest.main()

39
artifacts/package_sources/tinygrad/test/backend/test_call.py Normal file
View File

@@ -0,0 +1,39 @@
import unittest, ctypes
from tinygrad import Tensor, UOp
from tinygrad.device import Device
from tinygrad.dtype import dtypes
from tinygrad.renderer.cstyle import CStyleLanguage
from tinygrad.uop.ops import KernelInfo
def call_out_kernel(F:UOp, C:UOp) -> UOp:
call = F[0].load().call(UOp.const(3).cast(dtypes.int), C[0], ret_dtype=dtypes.void)
return C.after(call)[1].store(C.after(call)[0].load() + 1).sink(arg=KernelInfo(name="call_out"))
def call_ret_kernel(F:UOp, C:UOp) -> UOp:
val = F[0].load().call(UOp.const(21).cast(dtypes.int), ret_dtype=dtypes.int)
return C[0].store(val * 2).sink(arg=KernelInfo(name="call_ret"))
@unittest.skipUnless(isinstance(Device["CPU"].renderer, CStyleLanguage), "TODO: CALL is rendered in C style only")
class TestCall(unittest.TestCase):
def test_call_out_param(self):
called = []
@ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.POINTER(ctypes.c_int))
def fxn(n, out):
called.append(n)
out[0] = n * 2
f = Tensor([ctypes.cast(fxn, ctypes.c_void_p).value], dtype=dtypes.uint64, device="CPU")
c = Tensor.empty(2, dtype=dtypes.int, device="CPU")
c = Tensor.custom_kernel(f, c, fxn=call_out_kernel)[1]
self.assertEqual(c.tolist(), [6, 7])
self.assertEqual(called, [3])
def test_call_ret(self):
@ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int)
def fxn(n): return n + 1
f = Tensor([ctypes.cast(fxn, ctypes.c_void_p).value], dtype=dtypes.uint64, device="CPU")
c = Tensor.empty(1, dtype=dtypes.int, device="CPU")
c = Tensor.custom_kernel(f, c, fxn=call_ret_kernel)[1]
c.realize()
self.assertEqual(c.item(), 44)
if __name__ == "__main__": unittest.main()

211
artifacts/package_sources/tinygrad/test/backend/test_const_folding.py Normal file
View File

@@ -0,0 +1,211 @@
import unittest, math
from tinygrad import Tensor, Device, dtypes
from tinygrad.dtype import DTYPES_DICT
from tinygrad.uop.ops import Ops, UOp
from tinygrad.codegen.decomp.op import threefry2x32
import numpy as np
from test.helpers import not_support_multi_device
def _check_ast_count(desired_count:int, t:Tensor):
# NOTE: this has side effect because everything can be scheduled only once
schedule = t.schedule_linear()
asts = [s for s in schedule.src if s.src[0].op is Ops.SINK]
len(asts)
# NOT SUPPORTED ANYMORE
#assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
class TestMovedConstFolding(unittest.TestCase):
def test_contiguous_deviceless_const(self):
t = Tensor(UOp.const(2.0, dtypes.float)).contiguous()
self.assertIs(t.uop.op, Ops.CONST)
self.assertIsNone(t.uop.device)
def test_add_shrunk_zero(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(6).shrink(((1, 5),)))
def test_add_padded_zero(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(2).pad(((1, 1),)))
def test_mul_shrunk_one(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.ones(6).shrink(((1, 5),)))
def test_add_padded_one(self):
_check_ast_count(1, Tensor([1.0, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),)))
def test_copy_padded_const(self):
schedule = Tensor.ones(4, buffer=False).pad(((1, 1),)).to("CPU:1").schedule_linear()
assert not any(si.src[0].op is Ops.COPY for si in schedule.src), "const copy should be folded"
np.testing.assert_equal(Tensor.ones(4, buffer=False).pad(((1, 1),)).to("CPU:1").numpy(), [0, 1, 1, 1, 1, 0])
def test_cast_padded(self):
# NOTE: it's always 1 kernel when calling .numpy, limitation of _check_ast_count
_check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16))
np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16).numpy(), [0, 1, 1, 1, 1, 0])
_check_ast_count(1, Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16))
np.testing.assert_equal(Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16).numpy(), [0, 65535, 65535, 65535, 65535, 0])
# folded
_check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64))
np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64).numpy(), [0, 1, 1, 1, 1, 0])
class TestReduceOpsConstFolding(unittest.TestCase):
def test_const_sum(self):
_check_ast_count(0, Tensor.ones(4, 5, 6).sum())
np.testing.assert_equal(Tensor.ones(4, 5, 6).sum().numpy(), 4 * 5 * 6)
_check_ast_count(0, Tensor.ones(4, 5, 6).sum(axis=0))
np.testing.assert_equal(Tensor.ones(4, 5, 6).sum(axis=0).numpy(), np.full((5, 6), 4))
_check_ast_count(0, Tensor(4).sum())
np.testing.assert_equal(Tensor(4).sum().numpy(), 4)
def test_padded_const_sum(self):
_check_ast_count(0, Tensor.ones(4).pad(((1, 1),)).sum())
np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).sum().numpy(), 4)
# NOTE: cannot just count the non-padded area because some Ops f do not have f(0) = 0.
_check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).exp().sum())
np.testing.assert_allclose(Tensor.ones(4).pad(((1, 1),)).exp().sum().numpy(), 4 * math.e + 2)
def test_bool_zero_max(self):
_check_ast_count(0, Tensor.full((1, 2), True).shrink(((0, 1), (0, 0))).max((1, 0)))
np.testing.assert_equal(Tensor.full((1, 2), True).shrink(((0, 1), (0, 0))).max((1, 0)).numpy(), False)
def test_zero_size_ops(self):
for reduceop in [lambda x:x.prod(), lambda x:x.sum()]: # lambda x:x.max() NOTE: numpy gives "reduction operation maximum which has no identity"
_check_ast_count(0, reduceop(Tensor.empty(1, 0)))
np.testing.assert_equal(reduceop(Tensor.empty(shape:=(1, 0))).numpy(), reduceop(np.empty(shape)))
def test_zero_size_ops_view(self):
for reduceop in [lambda x:x.prod(), lambda x:x.sum()]:
_check_ast_count(0, reduceop(Tensor.empty(1, 0, 4).permute((1, 2, 0)).contiguous()))
np.testing.assert_equal(reduceop(Tensor.empty(shape:=(1, 0))).numpy(), reduceop(np.empty((shape))))
def test_zero_size_ops_realized(self):
for reduceop in [lambda x:x.prod(), lambda x:x.sum()]:
_check_ast_count(0, reduceop((Tensor.randn(0, 1)+1).realize()))
np.testing.assert_equal(reduceop((Tensor.randn(shape:=(0, 1))+1).realize()).numpy(), reduceop(np.empty(shape)))
def test_zero_size_realize_folded(self):
# non contiguous folded output doesn't realize
_check_ast_count(0, Tensor.empty(1, 0).sum())
# contiguous folded const can still schedule
a = Tensor.empty(1, 0).sum().contiguous()
_check_ast_count(2, a+2)
self.assertIs(a.uop.base.op, Ops.BUFFER)
np.testing.assert_equal((Tensor.empty(1, 0).sum().contiguous()+2).numpy(), 2)
# otherwise we just fuse it
_check_ast_count(1, (Tensor.empty(1, 0).sum()+2).contiguous())
np.testing.assert_equal((Tensor.empty(1, 0).sum()+2).numpy(), 2)
def test_const_prod(self):
_check_ast_count(0, Tensor.full((2, 3), fill_value=2).prod())
np.testing.assert_equal(Tensor.full((2, 3), fill_value=2).prod().numpy(), 2**(2*3))
_check_ast_count(0, Tensor.full((4, 5, 6), fill_value=2).prod(axis=0))
np.testing.assert_equal(Tensor.full((4, 5, 6), fill_value=2).prod(axis=0).numpy(), np.full((5, 6), 2**4))
_check_ast_count(0, Tensor(4).prod())
np.testing.assert_equal(Tensor(4).prod().numpy(), 4)
def test_const_max(self):
_check_ast_count(0, Tensor.ones(4, 5, 6).max())
np.testing.assert_equal(Tensor.ones(4, 5, 6).max().numpy(), 1)
_check_ast_count(0, Tensor(4).max())
np.testing.assert_equal(Tensor(4).max().numpy(), 4)
def test_sum_output_dtype(self):
# sum output dtype can be different from input
for dt in DTYPES_DICT.values():
if dt in Device[Device.DEFAULT].renderer.supported_dtypes():
t = Tensor.ones(16, dtype=dt).reshape(4, 4)
assert t.sum().dtype == t.contiguous().sum().dtype
@unittest.skipIf(not_support_multi_device() or True, "no multi, RANGEIFY doesn't support multi const folding")
class TestMultiConstFolding(unittest.TestCase):
def test_multi_const_folding_literal(self):
ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
t = Tensor.arange(16).float().clone().to(ds).realize()
# non const folding case creates one ast on each shard
_check_ast_count(4, t + 1)
_check_ast_count(4, 1 + t)
_check_ast_count(4, t * 2)
_check_ast_count(4, 2 * t)
# const folded
_check_ast_count(0, t + 0)
_check_ast_count(0, 0 + t)
_check_ast_count(0, t * 0)
_check_ast_count(0, 0 * t)
_check_ast_count(0, t * 1)
_check_ast_count(0, 1 * t)
np.testing.assert_equal((t + 0).numpy(), np.arange(16))
np.testing.assert_equal((t * 0).numpy(), [0] * 16)
np.testing.assert_equal((t * 1).numpy(), np.arange(16))
_check_ast_count(0, t ** 0)
_check_ast_count(0, t ** 1)
_check_ast_count(0, 1 ** t)
def test_multi_const_folding_tensor(self):
ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
t = Tensor.arange(16).float().clone().to(ds).realize()
zero = Tensor.zeros(16).to(ds).realize()
one = Tensor.ones(16).to(ds).realize()
# const folded
_check_ast_count(0, t + zero)
_check_ast_count(0, zero + t)
_check_ast_count(0, t * zero)
_check_ast_count(0, zero * t)
_check_ast_count(0, t * one)
_check_ast_count(0, one * t)
np.testing.assert_equal((t + zero).numpy(), np.arange(16))
np.testing.assert_equal((t * zero).numpy(), [0] * 16)
np.testing.assert_equal((t * one).numpy(), np.arange(16))
_check_ast_count(0, t ** zero)
_check_ast_count(0, t ** one)
_check_ast_count(0, one ** t)
np.testing.assert_equal((t ** zero).numpy(), [1] * 16)
np.testing.assert_equal((t ** one).numpy(), np.arange(16))
np.testing.assert_equal((one ** t).numpy(), [1] * 16)
class TestThreefryConstFolding(unittest.TestCase):
def test_threefry(self):
# THREEFRY(const,const) folds to a const once decomposed
x = threefry2x32(UOp.const(5, dtypes.uint64), UOp.const(10, dtypes.uint64))
self.assertIs(x.simplify().op, Ops.CONST)
class TestTautologicalCompare(unittest.TestCase):
# without const folding, these would have triggered -Wtautological-compare in clang
def test_lt_false(self):
# bool < False is always false
np.testing.assert_equal((Tensor([True, False]) < False).numpy(), [False, False])
def test_true_lt(self):
# True < bool is always false
np.testing.assert_equal((True < Tensor([True, False])).numpy(), [False, False])
def test_truth_table(self):
np.testing.assert_equal((Tensor(False) < Tensor(False)).numpy(), False)
np.testing.assert_equal((Tensor(False) < Tensor(True)).numpy(), True)
np.testing.assert_equal((Tensor(True) < Tensor(False)).numpy(), False)
np.testing.assert_equal((Tensor(True) < Tensor(True)).numpy(), False)
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support NaN comparison correctly")
def test_a_eq_a(self):
# self eq is always true for int or bool
a = Tensor([1, 2, 3])
np.testing.assert_equal((a == a).numpy(), [True, True, True])
# not true for nan
a = Tensor([math.nan, 1.0, 2.0])
np.testing.assert_equal((a == a).numpy(), [False, True, True])
def test_a_ne_a(self):
# self not eq is always false for int or bool
a = Tensor([1, 2, 3])
np.testing.assert_equal((a != a).numpy(), [False, False, False])
# not true for nan
a = Tensor([math.nan, 1.0, 2.0])
np.testing.assert_equal((a != a).numpy(), [True, False, False])
if __name__ == '__main__':
unittest.main()

697
artifacts/package_sources/tinygrad/test/backend/test_custom_kernel.py Normal file
View File

@@ -0,0 +1,697 @@
import unittest
from tinygrad import Tensor, UOp, GlobalCounters, Context, Device
import numpy as np
from tinygrad.dtype import AddrSpace, dtypes, Invalid
from tinygrad.uop.ops import KernelInfo, AxisType, Ops
from tinygrad.renderer.ptx import PTXRenderer
from test.helpers import assert_kernel_count
# **** kernels ****
def custom_arange_kernel(C:UOp) -> UOp:
i = UOp.range(C.shape[0], 0)
return C[i].store(i.cast(C.dtype)).end(i).sink(arg=KernelInfo(name=f"custom_arange_{C.shape[0]}"))
def custom_eye_kernel(C:UOp) -> UOp:
i = UOp.range(C.shape[0], 0)
j = UOp.range(C.shape[1], 1)
return C[i, j].store((i.eq(j)).cast(C.dtype)).end(i, j).sink(arg=KernelInfo(name=f"custom_eye_{C.numel()}"))
def custom_add_one_kernel(B:UOp, A:UOp) -> UOp:
A,B = A.flatten(), B.flatten()
assert B.numel() == A.numel()
i = UOp.range(A.numel(), 0)
return B[i].store(A[i] + 1).end(i).sink(arg=KernelInfo(name=f"add_one_{A.numel()}"))
def custom_elementwise_add_kernel(C:UOp, A:UOp, B:UOp) -> UOp:
C,A,B = C.flatten(), A.flatten(), B.flatten()
i = UOp.range(C.numel(), 0)
return C[i].store(A[i]+B[i]).end(i).sink(arg=KernelInfo(name=f"custom_add_kernel_{C.numel()}")).simplify()
def custom_elementwise_addmul_kernel(C:UOp, D:UOp, A:UOp, B:UOp) -> UOp:
C,D,A,B = C.flatten(), D.flatten(), A.flatten(), B.flatten()
assert C.numel() == D.numel()
i = UOp.range(C.numel(), 0)
store_c = C[i].store(A[i]+B[i])
store_d = D[i].store(A[i]*B[i])
return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name=f"custom_addmul_kernel_{C.numel()}")).simplify()
def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
assert A.shape[1] == B.shape[0]
i, j, k = UOp.range(C.shape[0], 0), UOp.range(C.shape[1], 1), UOp.range(A.shape[1], 2, axis_type=AxisType.REDUCE)
C = C[i, j].set(0.0)
C = C[i, j].set(C.after(k)[i, j] + A[i, k] * B[k, j], end=k)
prog = C.end(i, j)
return prog.sink(arg=KernelInfo(name=f"custom_gemm_{C.shape[0]}_{C.shape[1]}_{A.shape[1]}", opts_to_apply=()))
def custom_sum(B:UOp, A:UOp) -> UOp:
i = UOp.range(A.shape[0], 0, axis_type=AxisType.REDUCE)
B = B[0].set(0.0)
B = B[0].set(B.after(i)[0] + A[i], end=i)
return B.sink(arg=KernelInfo(name=f"custom_sum_{A.shape[0]}", opts_to_apply=()))
def flip_contract_kernel(dest:UOp, src:UOp):
i = UOp.range(dest.shape[0], 0)
j = UOp.range(dest.shape[1], 1, AxisType.UPCAST)
vec = src[i, j].contract(j)
store = UOp.group(*[dest[i, k].store(vec.index(3-k)) for k in range(4)])
return store.end(i, j).sink(arg=KernelInfo(name=f"flip_contract_{dest.numel()}", opts_to_apply=()))
def slice_sum_kernel(dest:UOp, src:UOp):
G = UOp.range(src.shape[0], 0, dtype=dtypes.int)
slice_src = src[G, :]
reg = UOp.placeholder((1,), dest.dtype, 0, addrspace=AddrSpace.REG)
reg = reg.after(G)[0].set(0)
R = UOp.range(src.shape[1], 1, AxisType.REDUCE)
reg = reg[0].set(reg.after(R)[0] + slice_src[R], end=R)
ast = dest[G].set(reg[0], end=G)
return ast.sink(arg=KernelInfo(name=f"slice_sum_{src.shape[0]}_{src.shape[1]}", opts_to_apply=()))
def simple_qkv_kernel(O:UOp, Q:UOp, K:UOp, V:UOp) -> UOp:
# attention without softmax
N, d = Q.shape[0], Q.shape[1]
i = UOp.range(N, 0) # output row
d_out = UOp.range(d, 1) # output column
j = UOp.range(N, 2, axis_type=AxisType.REDUCE)
k_inner = UOp.range(d, 3, axis_type=AxisType.REDUCE)
qk_acc = UOp.placeholder((1,), Q.dtype, 0, addrspace=AddrSpace.REG)
qk_acc = qk_acc.after(i, j)[0].set(0.0)
qk_acc = qk_acc[0].set(qk_acc.after(k_inner)[0] + Q[i, k_inner] * K[j, k_inner], end=k_inner)
qk_score = qk_acc[0] / (d ** 0.5)
out_acc = UOp.placeholder((1,), Q.dtype, 1, addrspace=AddrSpace.REG)
out_acc = out_acc.after(i, d_out)[0].set(0.0)
out_acc = out_acc[0].set(out_acc.after(j)[0] + qk_score * V[j, d_out], end=j)
store = O[i, d_out].store(out_acc[0])
return store.end(d_out).end(i).sink(arg=KernelInfo(name=f"simple_qkv_{N}_{d}", opts_to_apply=()))
# **** backward callbacks ****
def backward_gemm(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
out, a, b = kernel.src[1:]
grad_a = (Tensor(gradient) @ Tensor(b).T).uop
grad_b = (Tensor(a).T @ Tensor(gradient)).uop
return (None, grad_a, grad_b)
def backward_gemm_custom(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
out, a, b = kernel.src[1:]
grad_a = Tensor.empty_like(Tensor(a)).custom_kernel(Tensor(gradient), Tensor(b).T, fxn=custom_gemm)[0].uop
grad_b = Tensor.empty_like(Tensor(b)).custom_kernel(Tensor(a).T, Tensor(gradient), fxn=custom_gemm)[0].uop
return (None, grad_a, grad_b)
# **** tests ****
class TestCustomKernel(unittest.TestCase):
def test_empty(self):
a = Tensor.empty(1)
a = Tensor.custom_kernel(a, fxn=lambda _: UOp.sink(arg=KernelInfo()))[0]
a.realize()
def test_simple(self):
a = Tensor.ones(16, 16).contiguous()
b = Tensor.ones(16, 16).contiguous()
c = Tensor.empty(16, 16)
c = Tensor.custom_kernel(c,a,b, fxn=custom_elementwise_add_kernel)[0]
out = c.flatten().tolist()
assert all(x == 2 for x in out), "all 2"
def test_duplicate_call_arg(self):
x = Tensor.arange(4).clone().realize()
x = Tensor.custom_kernel(x, x, fxn=custom_add_one_kernel)[0]
# webgpu silently errors when a kernel has duplicate buffer args, so the list stays the same.
# https://gpuweb.github.io/gpuweb/#abstract-opdef-encoder-bind-groups-alias-a-writable-resource
self.assertEqual(x.tolist(), [1, 2, 3, 4] if Device.DEFAULT != "WEBGPU" else [0, 1, 2, 3])
def test_simple_sharded(self):
devs = ("CPU:0", "CPU:1")
a = Tensor.ones(16, 16).contiguous().shard(devs, axis=0)
b = Tensor.ones(16, 16).contiguous().shard(devs, axis=0)
# ugly construction to get a sharded empty tensor
c = Tensor(Tensor.empty(8, 16, device=devs).uop.unshard(0), device=devs)
c = Tensor.custom_kernel(c,a,b, fxn=custom_elementwise_add_kernel)[0]
out = c.flatten().tolist()
assert all(x == 2 for x in out), "all 2"
def test_sharded_add_one(self):
# PYTHON backend explicitly checks for OOB access for wrong multi shape regression
devs = ("PYTHON:0", "PYTHON:1")
a = Tensor.ones(4, 4).contiguous().shard(devs, axis=0)
c = Tensor(Tensor.empty(2, 4, device=devs).uop.unshard(0), device=devs)
c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
assert (c == 2).all().item()
def test_multioutput(self):
a = Tensor.full((16, 16), 3.).contiguous()
b = Tensor.full((16, 16), 3.).contiguous()
c = Tensor.empty(16, 16)
d = Tensor.empty(16, 16)
c,d = Tensor.custom_kernel(c,d,a,b, fxn=custom_elementwise_addmul_kernel)[:2]
Tensor.realize(c,d)
assert all(x == 6 for x in c.flatten().tolist()), "all 6"
assert all(x == 9 for x in d.flatten().tolist()), "all 9"
def test_arange(self):
ref = Tensor.arange(100)
tst = Tensor.empty_like(ref)
tst = tst.custom_kernel(fxn=custom_arange_kernel)[0]
self.assertTrue((ref == tst).all().item())
def test_eye(self):
ref = Tensor.eye(1024).clone().realize()
tst = Tensor.empty_like(ref)
tst = tst.custom_kernel(fxn=custom_eye_kernel)[0]
self.assertTrue((ref == tst).all().item())
@unittest.skip("contract shouldn't be supported here")
def test_flip_contract(self):
a = Tensor.randn(10,4)
b = Tensor.empty_like(a)
b = b.custom_kernel(a, fxn=flip_contract_kernel)[0]
self.assertTrue((a.flip(1) == b).all().item())
def test_noncontig(self):
a = Tensor.ones(16, 16).contiguous()
tst = Tensor.empty_like(a)
b = a+1
b_p1 = Tensor.custom_kernel(tst, b, fxn=custom_add_one_kernel)[0]
self.assertTrue((b_p1 == 3).all().item())
def test_sum(self):
a = Tensor([1.0, 2, 3, 4, 5])
tst = Tensor.empty(1)
b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
self.assertEqual(b.item(), 15)
def test_sum_outside(self):
a = Tensor([1.0, 2, 3, 4, 5])+1
tst = Tensor.empty(1)
b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
self.assertEqual(b.item(), 20)
def test_sum_int(self):
a = Tensor([1, 2, 3, 4, 5])
tst = Tensor.empty(1, dtype=a.dtype)
b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
self.assertEqual(b.item(), 15)
def test_slice_sum(self):
A = Tensor.randn(16, 16).contiguous()
B = Tensor.empty(16)
B = Tensor.custom_kernel(B, A, fxn=slice_sum_kernel)[0]
self.assertTrue(B.allclose(A.sum(1)).item())
def test_gemm(self):
N = 16
a = Tensor.randn(N, N)
b = Tensor.randn(N, N)
c = Tensor.empty(N, N)
tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
self.assertTrue(tst.allclose(a@b, atol=1e-3).item())
def test_gemm_multi(self):
devs = ("CPU:0", "CPU:1")
N = 16
a = Tensor.randn(N, N).shard_(devs, axis=0)
b = Tensor.randn(N, N).to(devs)
c = Tensor(Tensor.empty(N//2, N, device=devs).uop.unshard(0), device=devs)
tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
self.assertTrue(tst.allclose(a@b, atol=1e-3).item())
def test_gemm_backward_custom(self): self.test_gemm_backward(True)
# NOTE: grad_fxn doesn't work with pyrender
def test_gemm_backward(self, custom_backward_gemm=False):
N = 4
a_rand = Tensor.randn(N, 8)
b_rand = Tensor.randn(8, N)
Tensor.realize(a_rand, b_rand)
a, b = Tensor(a_rand.numpy()), Tensor(b_rand.numpy())
c = Tensor.empty(N, N)
tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm, grad_fxn=backward_gemm_custom if custom_backward_gemm else backward_gemm)[0]
tst.sum().backward()
grad_a, grad_b = a.grad, b.grad
Tensor.realize(tst, grad_a, grad_b)
a, b = Tensor(a_rand.numpy()), Tensor(b_rand.numpy())
ref = (a@b)
ref.sum().backward()
real_grad_a, real_grad_b = a.grad, b.grad
Tensor.realize(ref, real_grad_a, real_grad_b)
self.assertTrue(tst.allclose(ref, atol=1e-3).item())
self.assertTrue(grad_a.allclose(real_grad_a, atol=1e-3).item())
self.assertTrue(grad_b.allclose(real_grad_b, atol=1e-3).item())
def test_simple_qkv(self):
N, d = 8, 4
Q = Tensor.randn(N, d)
K = Tensor.randn(N, d)
V = Tensor.randn(N, d)
O = Tensor.empty(N, d)
O_custom = Tensor.custom_kernel(O, Q, K, V, fxn=lambda o,q,k,v: simple_qkv_kernel(o,q,k,v))[0]
O_ref = ((Q @ K.T) / (d ** 0.5)) @ V
Tensor.realize(O_custom, O_ref)
self.assertTrue(O_custom.allclose(O_ref, atol=1e-3).item())
def test_gemm_qkv(self):
B, N, K_DIM, H_KV, REP, D = 2, 7, 6, 2, 2, 6
H, QKV = H_KV * REP, H_KV * (REP + 2) * D
x = Tensor.empty(B*N, K_DIM)
w = Tensor.empty(K_DIM, QKV)
qkv = Tensor.empty(B*N, QKV)
qkv = Tensor.custom_kernel(qkv, x, w, fxn=custom_gemm)[0]
qkv = qkv.reshape(B, N, H_KV, REP + 2, D)
q = qkv[:, :, :, :REP, :].reshape(B, N, H, D).transpose(1, 2)
k = qkv[:, :, :, REP, :].transpose(1, 2)
v = qkv[:, :, :, REP + 1, :].transpose(1, 2)
out = q.scaled_dot_product_attention(k, v, enable_gqa=True)
GlobalCounters.reset()
out.realize()
assert_kernel_count(5)
def test_simple_reshape(self):
a = Tensor.ones(2,3,4).realize()
b = Tensor.custom_kernel(Tensor.empty_like(a), a, fxn=custom_add_one_kernel)[0]
b2 = b.reshape(2,12)
c = Tensor.custom_kernel(Tensor.empty_like(b2), b2, fxn=custom_add_one_kernel)[0]
GlobalCounters.reset()
c.realize()
assert all(i == 3. for i in c.flatten().tolist()), f"all 3 {c.tolist()}"
assert_kernel_count(2)
def test_multi_after_schedule_order(self):
"""Test correct scheduling order when custom_kernel has multiple outputs.
custom_kernel with 4 arguments creates 4 AFTERs from the same kernel.
The custom_kernel depends on both A2 and B2, so it must be scheduled after both.
E only depends on A2, so E can run before custom_kernel finishes waiting for B2.
Expected schedule order: [A2, B2, E, custom_addmul, final_sum]
The custom_addmul kernel should be at index 3.
"""
A, B = Tensor.empty(4, 4), Tensor.empty(4, 4)
A2 = (A + 1).contiguous() # kernel 0: depends on A
B2 = (B * 2).contiguous() # kernel 1: depends on B
C, D = Tensor.empty(4, 4), Tensor.empty(4, 4)
C, D, _, _ = Tensor.custom_kernel(C, D, A2, B2, fxn=custom_elementwise_addmul_kernel) # depends on A2 AND B2
E = (A2 * 3).contiguous() # kernel 2: depends only on A2
result = (C + D + E).sum() # kernel 3: custom_addmul, then kernel 4: sum
schedule = result.schedule_linear().src
# Find the custom_addmul kernel position
custom_idx = next((i for i, item in enumerate(schedule)
if hasattr(item.src[0], "arg") and hasattr(item.src[0].arg, "name")
and "custom_addmul" in item.src[0].arg.name), None)
self.assertIsNotNone(custom_idx, "custom_addmul kernel not found in schedule")
self.assertEqual(custom_idx, 3, f"custom_addmul should be at index 3, got {custom_idx}")
def test_invalids_into_custom_kernel_no_empty_kernel(self):
from tinygrad.engine.realize import compile_linear
a = Tensor.full((4, 4), 3.).contiguous()
b = Tensor.full((4, 4), 2.).contiguous()
Tensor.realize(a, b)
out = Tensor.invalids(*a.shape, dtype=a.dtype)
out, *_ = Tensor.custom_kernel(out, a, b, fxn=custom_elementwise_add_kernel)
compiled = compile_linear(out.schedule_linear())
for call in compiled.src:
prg = call.src[0]
if prg.op is not Ops.PROGRAM: continue
self.assertTrue(len(prg.arg.globals) > 0, f"empty kernel compiled (no globals): name={prg.arg.name}")
def test_multi_invalids_custom_kernel_no_copy(self):
devs = ("CPU:0", "CPU:1")
a = Tensor.ones(4, 4).shard(devs, axis=0).realize()
c = Tensor(Tensor.invalids(2, 4, dtype=dtypes.float, device=devs).uop.unshard(0), device=devs)
c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
GlobalCounters.reset()
c.realize()
assert_kernel_count(len(devs))
self.assertTrue((c == 2).all().item())
def test_partial_invalid_store_keeps_uncovered_reads(self):
x = Tensor([10., 20., 30., 40.])
after = x.uop.after(x.uop.shrink(((0, 2),)).store(Invalid))
self.assertEqual(Tensor(after).contiguous().tolist(), [10., 20., 30., 40.])
def test_multi_after_invalid_store_dep_removed(self):
x = Tensor.empty(4).uop
self.assertEqual(Tensor(x.after(x.store(5), x.store(Invalid))).tolist(), [5]*4)
def test_expand_view_invalid_assign_keeps_uncovered_reads(self):
x = Tensor([[10., 11., 12., 13.], [20., 21., 22., 23.], [30., 31., 32., 33.], [40., 41., 42., 43.]]).realize()
v = x[:1, :].expand(4, 4)
v.assign(Tensor.invalids(4, 4, dtype=dtypes.float))
self.assertEqual(v.contiguous().tolist(), [[10., 11., 12., 13.]]*4)
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "kernel timing not supported")
def test_invalids_into_custom_kernel_with_beam(self):
a = Tensor.full((4, 4), 3.).contiguous()
b = Tensor.full((4, 4), 2.).contiguous()
Tensor.realize(a, b)
with Context(BEAM=1, IGNORE_BEAM_CACHE=1):
out = Tensor.invalids(*a.shape, dtype=a.dtype)
out, *_ = Tensor.custom_kernel(out, a, b, fxn=custom_elementwise_add_kernel)
result = out.flatten().tolist()
self.assertTrue(all(x == 5 for x in result), f"expected all 5.0, got {result}")
@unittest.skip("what are anonymous buffers?")
def test_anonymous_buffers_in_function(self):
"""Test that custom kernels with anonymous output buffers work inside @function."""
a = Tensor.full((4, 4), 3.).contiguous()
b = Tensor.full((4, 4), 2.).contiguous()
Tensor.realize(a, b)
def custom_add_with_tmp(o1:UOp, o2:UOp, A:UOp, B:UOp) -> UOp:
o1,o2,A,B = o1.flatten(), o2.flatten(), A.flatten(), B.flatten()
i = UOp.range(o1.numel(), 0)
store_o1 = o1[i].store(A[i]+B[i])
store_o2 = o2[i].store(A[i]+B[i]+2)
return UOp.group(store_o1, store_o2).end(i).sink(arg=KernelInfo(name=f"add_with_tmp_{o1.numel()}")).simplify()
from tinygrad import function
@function(precompile=True)
def run(x:Tensor, w:Tensor) -> Tensor:
out = Tensor.invalids(*x.shape, dtype=x.dtype)
tmp = Tensor.invalids(*x.shape, dtype=x.dtype)
out, tmp = Tensor.custom_kernel(out, tmp, x, w, fxn=custom_add_with_tmp)[:2]
return out+tmp
result = run(a, b).flatten().tolist()
expected = (3+2)*2+2
assert all(x == expected for x in result), f"expected all {expected}, got {result}"
def test_custom_kernel_sched(self, use_custom=False):
x = Tensor.arange(32).reshape(8, 4).clone().realize()
y = Tensor.empty_like(x)
y = Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
if use_custom:
z = Tensor.empty_like(x)
z = Tensor.custom_kernel(z, y.T.T, fxn=custom_add_one_kernel)[0]
else: z = y.T.T+1
GlobalCounters.reset()
z.realize()
assert_kernel_count(2)
self.assertEqual(z.tolist(), x.add(2).tolist())
def test_custom_kernel_sched_copy(self): self.test_custom_kernel_sched(use_custom=True)
def test_sliced_buffer_function(self):
x = Tensor.arange(32).reshape(8, 4).clone().realize()
from tinygrad import function
@function(precompile=True)
def run(x:Tensor) -> Tensor:
y = Tensor.invalids(*x.shape, dtype=x.dtype)
return Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
GlobalCounters.reset()
y = run(x[0]).realize()
# it's copying the input and the output
# TODO: subbuffer usage has runtime specific behavior, this will be fixed after the removal of SLICE.
assert_kernel_count(2 if y.device in ("CL", "WEBGPU") else 1)
self.assertEqual(y.tolist(), [1, 2, 3, 4])
@Context(DEV="CPU")
def test_simple_from_source(self):
a = Tensor.arange(4).clone().realize()
src = "void test_src(int* restrict a) { a[0] = 1; }"
# TODO: it currently requires a compiler for Ops.BINARY
from tinygrad.device import Device
binary = Device[a.device].renderer.compiler.compile(src)
def custom_src_kernel(A:UOp, B:UOp) -> UOp:
sink = UOp.sink(A, arg=KernelInfo(name="test_src"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
a = Tensor.custom_kernel(a.reshape(2, 2).clone(), a.reshape(2, 2).T, fxn=custom_src_kernel)[0]
self.assertEqual(a.tolist(), [[1, 1], [2, 3]])
@Context(DEV="CPU")
def test_simple_from_source_alt(self):
a = Tensor.arange(4).clone().realize()
src = "void copy(int* restrict out, int* restrict in) { for (int i = 0; i < 4; i++) out[i] = in[i]; }"
from tinygrad.device import Device
binary = Device[a.device].renderer.compiler.compile(src)
def custom_src_kernel(out:UOp, inp:UOp) -> UOp:
sink = UOp.sink(out, inp, arg=KernelInfo(name="copy"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
out = Tensor.custom_kernel(Tensor.empty_like(a), a+1, fxn=custom_src_kernel)[0]
GlobalCounters.reset()
out.realize()
assert_kernel_count(2)
self.assertEqual(out.tolist(), [1, 2, 3, 4])
@unittest.skip("this shouldn't be expected to work")
def test_inplace_transpose(self):
def custom_assign_row_max_kernel(A:UOp) -> UOp:
row = UOp.range(A.shape[0], 0)
col = UOp.range(A.shape[1], 1)
return A[row, col].store(A[row].max(axis=0)).end(col).end(row).sink(arg=KernelInfo(name=f"assign_row_max_{A.numel()}"))
a = Tensor.arange(4).clone().realize()
a = Tensor.custom_kernel(a.reshape(2, 2).T, fxn=custom_assign_row_max_kernel)[0]
self.assertEqual(a.flatten().tolist(), [2, 2, 3, 3])
self.assertEqual(a.shape, (2, 2))
class TestCustomKernelInput(unittest.TestCase):
def _test_mop(self, mop_fxn, max_kernels):
# default: input is BUFFER
x = mop_fxn(Tensor.arange(32).clone("CPU").realize())
y = Tensor.custom_kernel(Tensor.empty_like(x), x, fxn=custom_add_one_kernel)[0]
GlobalCounters.reset()
y.realize()
kernel_count = GlobalCounters.kernel_count
self.assertEqual(y.tolist(), x.add(1).tolist())
self.assertLessEqual(kernel_count, max_kernels)
# same test with @function, input is PARAM
from tinygrad import function
x0 = Tensor.arange(32).clone("CPU").realize()
@function(precompile=True)
def run(a:Tensor) -> Tensor:
xv = mop_fxn(a)
y = Tensor.invalids(*xv.shape, dtype=xv.dtype, device=a.device)
return Tensor.custom_kernel(y, xv, fxn=custom_add_one_kernel)[0]
GlobalCounters.reset()
y = run(x0).realize()
kernel_count = GlobalCounters.kernel_count
self.assertEqual(y.tolist(), mop_fxn(x0).add(1).tolist())
self.assertLessEqual(kernel_count, max_kernels)
def test_reshape(self): self._test_mop(lambda x: x.reshape(16, 2), max_kernels=2)
def test_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T, max_kernels=3)
def test_double_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T.T, max_kernels=2)
def test_shrink(self): self._test_mop(lambda x: x[:4], max_kernels=1)
def test_pad(self): self._test_mop(lambda x: x[:4].pad(((0, 4),)), max_kernels=2)
def test_flip(self): self._test_mop(lambda x: x.flip(0), max_kernels=2)
def test_offset_shrink(self): self._test_mop(lambda x: x[4:8], max_kernels=2)
def test_2d_shrink(self): self._test_mop(lambda x: x.reshape(4, 8)[:, 2:6], max_kernels=3)
def test_expand(self): self._test_mop(lambda x: x.reshape(16, 2)[:, :1].expand(16, 2), max_kernels=3)
class TestUnshardIndex(unittest.TestCase):
"""Regression tests for INDEX on UNSHARD (fragment) resolution in schedule/multi.py.
A fragment is a per-thread REG buffer wrapped in UNSHARD over LOCAL thread ranges.
index_multi must resolve an INDEX on the UNSHARD view into an INDEX on the per-thread
shard. Two ownership patterns must work:
contiguous: idx = rng*shard_sz + local (thread rng owns [rng*shard_sz, ...))
strided: idx = rng + ir*shard_sz (thread rng owns {rng, rng+shard_sz, ...})
"""
def _run(self, kernel, shape=(8, 8)):
c = Tensor.empty(*shape)
out = Tensor.custom_kernel(c, fxn=kernel)[0]
try: return out.numpy()
except RuntimeError as e:
if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and "dynamic register indexing" in str(e):
self.skipTest("PTX does not support dynamic register indexing")
raise
@unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
def test_contiguous_fragment_index(self):
# thread ty owns rows [ty*8, ty*8+8) of a 64-row fragment -- contiguous ownership.
# This is the pre-existing case that index_multi always handled.
def kernel(C:UOp) -> UOp:
ty = UOp.range(8, 0, AxisType.LOCAL)
ir = UOp.range(8, 1, AxisType.LOOP)
j = UOp.range(8, 2, AxisType.LOOP)
# 8x8 fragment, 8 threads -> 64x8 full tile. thread ty owns rows [ty*8, ty*8+8).
frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
return C[ty*8 + ir, j].store(frag[ty*8 + ir, j]).end(j, ir, ty).sink(arg=KernelInfo(name="contig_frag"))
out = self._run(kernel, (64, 8))
assert out.shape == (64, 8)
@unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
def test_strided_fragment_index(self):
# thread ty owns rows {ty, ty+8, ty+16, ty+24, ..., ty+56} of a 64-row fragment --
# strided ownership. idx = ty + ir*8 where shard_sz=8 (8 threads, shard rows=8).
# The contiguous check (idx - rng*shard_sz) fails; the strided check
# (idx-rng) % shard_sz == 0 must succeed. This is the pattern the index_multi fix adds.
def kernel(C:UOp) -> UOp:
ty = UOp.range(8, 0, AxisType.LOCAL)
ir = UOp.range(8, 1, AxisType.LOOP)
j = UOp.range(8, 2, AxisType.LOOP)
# 8x8 fragment, 8 threads -> 64x8 full tile. thread ty owns rows {ty, ty+8, ..., ty+56}.
frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
return C[ty + ir*8, j].store(frag[ty + ir*8, j]).end(j, ir, ty).sink(arg=KernelInfo(name="strided_frag"))
out = self._run(kernel, (64, 8))
assert out.shape == (64, 8)
def test_fragment_index_cannot_shard(self):
# thread ty indexing rows [ty, ty+8) overlaps with other threads' rows -- this matches neither
# the contiguous nor the strided ownership pattern, so index_multi must raise.
def kernel(C:UOp) -> UOp:
ty = UOp.range(8, 0, AxisType.LOCAL)
ir = UOp.range(8, 1, AxisType.LOOP)
j = UOp.range(8, 2, AxisType.LOOP)
frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
return C[ty + ir, j].store(frag[ty + ir, j]).end(j, ir, ty).sink(arg=KernelInfo(name="bad_frag"))
with self.assertRaisesRegex(RuntimeError, "cannot shard index"):
self._run(kernel, (64, 8))
def _run_fragment_kernel(testcase, kernel, out_shape, inputs=()):
c = Tensor.empty(*out_shape)
out = Tensor.custom_kernel(c, *inputs, fxn=kernel)[0]
try: return out.numpy()
except RuntimeError as e:
if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and "dynamic register indexing" in str(e):
testcase.skipTest("PTX does not support dynamic register indexing")
raise
class TestUnshardAlu(unittest.TestCase):
"""Tests for ALU on (fragment) UNSHARD values in schedule/multi.py's alu_multi.
An ALU with UNSHARD srcs lowers to per-shard ops when every src is one of:
same sharding: peel the UNSHARD, keep the layout
scalar: broadcast to every shard
whole unsharded same-shape value: takes its per-shard sub-view (shard_subview)
"""
@unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
def test_alu_scalar_broadcast(self):
# scalar srcs broadcast to every shard: frag*2.0 where frag is 1.5 per thread -> 3.0 everywhere
def kernel(C:UOp) -> UOp:
ty = UOp.range(8, 0, AxisType.LOCAL)
# 8 values per thread, 8 threads -> 64-value full view
frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,))
v = frag.after(frag.store(1.5)) * 2.0
return C.store(v).end(ty).sink(arg=KernelInfo(name="alu_scalar", opts_to_apply=()))
out = _run_fragment_kernel(self, kernel, (64,))
np.testing.assert_allclose(out, 3.0)
@unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
def test_alu_whole_value_subview(self):
# UNSHARD + whole unsharded same-shape value: each shard adds its own sub-view of A.
def kernel(C:UOp, A:UOp) -> UOp:
ty = UOp.range(8, 0, AxisType.LOCAL)
frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,))
v = frag.after(frag.store(0.0)) + A
return C.store(v).end(ty).sink(arg=KernelInfo(name="alu_subview", opts_to_apply=()))
a = Tensor(np.arange(64, dtype=np.float32))
out = _run_fragment_kernel(self, kernel, (64,), inputs=(a,))
np.testing.assert_allclose(out, a.numpy(), atol=1e-4)
class TestUnshardStore(unittest.TestCase):
"""Tests for STORE of a sharded value into an unsharded dest (store_value_multi in schedule/multi.py).
Every shard stores its value into its own contiguous sub-view of the dest, one SHRINK per sharded axis.
"""
@unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
def test_store_unshard_value(self):
# single-axis: 8 threads each own 8 values of the 64-value output tile
def kernel(C:UOp) -> UOp:
ty = UOp.range(8, 0, AxisType.LOCAL)
frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,))
v = frag.after(frag.store(0.0)) + 2.5
return C.store(v).end(ty).sink(arg=KernelInfo(name="store_unshard", opts_to_apply=()))
out = _run_fragment_kernel(self, kernel, (64,))
np.testing.assert_allclose(out, 2.5)
@unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
def test_store_unshard_value_2axis(self):
# two sharded axes (the gemm fragment layout): thread (ty, tx) owns the (2, 1, 1, 2) sub-view of the
# (2, 4, 2, 2) output tile; the store must SHRINK dest on both sharded axes
def kernel(C:UOp, A:UOp) -> UOp:
ty = UOp.range(4, 0, AxisType.LOCAL)
tx = UOp.range(2, 1, AxisType.LOCAL)
frag = UOp.placeholder((2, 1, 1, 2), dtypes.float32, 0, AddrSpace.REG).unshard((1, 2), (ty, tx))
v = frag.after(frag.store(0.0)) + A
return C.store(v).end(tx, ty).sink(arg=KernelInfo(name="store_unshard_2axis", opts_to_apply=()))
a = Tensor(np.arange(32, dtype=np.float32).reshape(2, 4, 2, 2))
out = _run_fragment_kernel(self, kernel, (2, 4, 2, 2), inputs=(a,))
np.testing.assert_allclose(out, a.numpy(), atol=1e-4)
class TestUOpReduce(unittest.TestCase):
def test_uop_sum(self):
a = Tensor([1.0, 2, 3, 4, 5])
self.assertAlmostEqual(Tensor(a.uop.sum(axis=0)).item(), 15.0)
def test_uop_sum_2d(self):
a = Tensor.arange(6).reshape(2, 3).float()
result = Tensor(a.uop.sum(axis=1)).numpy()
assert result[0] == 3 and result[1] == 12
def test_uop_sum_all(self):
a = Tensor.arange(6).reshape(2, 3).float()
self.assertAlmostEqual(Tensor(a.uop.sum()).item(), 15.0)
def test_uop_sum_keepdim(self):
a = Tensor.arange(6).reshape(2, 3).float()
result = Tensor(a.uop.sum(axis=1, keepdim=True))
assert result.shape == (2, 1)
def test_uop_sum_negative_axis(self):
a = Tensor.arange(6).reshape(2, 3).float()
result = Tensor(a.uop.sum(axis=-1)).numpy()
assert result[0] == 3 and result[1] == 12
def test_uop_sum_multi_axis(self):
a = Tensor.arange(24).reshape(2, 3, 4).float()
ref = a.sum(axis=(0, 2)).numpy()
result = Tensor(a.uop.sum(axis=(0, 2))).numpy()
for i in range(3): self.assertAlmostEqual(result[i], ref[i])
def test_uop_sum_dtype(self):
a = Tensor([1.0, 2, 3], dtype=dtypes.float16)
result = Tensor(a.uop.sum(axis=0, dtype=dtypes.float32))
self.assertEqual(result.dtype, dtypes.float)
self.assertAlmostEqual(result.item(), 6.0, places=2)
def test_uop_prod(self):
a = Tensor([1.0, 2, 3, 4, 5])
self.assertAlmostEqual(Tensor(a.uop.prod(axis=0)).item(), 120.0)
def test_uop_max(self):
a = Tensor([1.0, 5, 3, 2, 4])
self.assertAlmostEqual(Tensor(a.uop.max(axis=0)).item(), 5.0)
def test_uop_max_2d(self):
a = Tensor([[1, 5, 3], [4, 2, 6]]).float()
result = Tensor(a.uop.max(axis=0)).numpy()
assert result[0] == 4 and result[1] == 5 and result[2] == 6
def test_uop_std(self):
a = Tensor([2.0, 4, 4, 4, 5, 5, 7, 9])
self.assertAlmostEqual(Tensor(a.uop.std()).item(), a.std().item(), places=5)
class TestUOpWhere(unittest.TestCase):
def test_uop_where_both_const(self):
cond = Tensor([True, False, True])
result = Tensor(cond.uop.where(1, 0))
self.assertEqual(result.tolist(), [1, 0, 1])
result = Tensor(cond.uop.where(1.5, 0))
self.assertEqual(result.tolist(), [1.5, 0, 1.5])
if __name__ == '__main__':
unittest.main()

438
artifacts/package_sources/tinygrad/test/backend/test_dtype.py Normal file
View File

@@ -0,0 +1,438 @@
import contextlib, unittest, math
import numpy as np
import torch
from typing import Any, List
from tinygrad.helpers import getenv, DEBUG, EMULATED_DTYPES, DEV
from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8, _to_np_dtype, _to_torch_dtype, truncate
from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.nir import NIRRenderer
from tinygrad import Context, Device, Tensor, dtypes
from hypothesis import given, settings, strategies as strat
from test.helpers import rand_for_dtype, min_normal
from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX, FP8E4M3FNUZ_MAX, FP8E5M2FNUZ_MAX
import pytest
pytestmark = pytest.mark.filterwarnings("ignore")
settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
settings.load_profile("my_profile")
supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
def get_available_cast_dtypes(dtype: DType) -> List[DType]:
dts = [v for k, v in DTYPES_DICT.items() if v != dtype and v in supported_dtypes or v in dtypes.fp8s+(dtypes.half,dtypes.bfloat16,dtypes.long)]
if dtype in (dtypes.long, dtypes.ulong) and (dtype not in supported_dtypes or dtypes.long in EMULATED_DTYPES.tolist(dtypes)):
return [dt for dt in dts if dt != dtypes.double] # can't bitcast with no 64-bit support
if dtype not in supported_dtypes and dtype not in dtypes.fp8s+(dtypes.half,dtypes.bfloat16): return []
return dts
def _to_torch_storage(a:Tensor) -> torch.Tensor:
# tolist() of an fp8 Tensor gives floats, so convert and store in uint8
if a.dtype in dtypes.fp8s: return torch.tensor([float_to_fp8(x, a.dtype) for x in a.flatten().tolist()], dtype=torch.uint8).reshape(a.shape)
return torch.tensor(a.tolist(), dtype=_to_torch_dtype(a.dtype))
def _test_to_np(a:Tensor, np_dtype, target):
if DEBUG >= 2: print(a)
na = a.numpy()
if DEBUG >= 2: print(na, na.dtype, a.uop.base.realized)
try:
assert na.dtype == np_dtype
np.testing.assert_allclose(na, target)
except AssertionError as e:
raise AssertionError(f"\ntensor {a.numpy()} does not match target {target} with np_dtype {np_dtype}") from e
def _test_op(fxn, target_dtype:DType, target):
_assert_eq(fxn(), target_dtype, target)
def _test_cast(a:Tensor, target_dtype:DType):
if a.is_floating_point() and dtypes.is_unsigned(target_dtype):
# converting negative float to unsigned integer is undefined
a = a.abs()
if a.is_floating_point() and dtypes.is_float(target_dtype) and (mn:=min_normal(target_dtype)) >= min_normal(a.dtype):
# subnormals are zero, so an input below the target's min normal casts to 0
a = (a.abs() < mn).where(0, a)
expected = list(a.numpy().astype(_to_np_dtype(target_dtype)))
if target_dtype in dtypes.fp8s: expected = [truncate[target_dtype](x) for x in expected]
_test_op(lambda: a.cast(target_dtype), target_dtype, expected)
def _test_bitcast(a:Tensor, target_dtype:DType, target=None):
expected = _to_torch_storage(a).view(_to_torch_dtype(target_dtype)).tolist()
if target_dtype in dtypes.fp8s: expected = [fp8_to_float(x, target_dtype) for x in expected]
_test_op(lambda: a.bitcast(target_dtype), target_dtype, target or expected)
class TestDType(unittest.TestCase):
DTYPE: Any = None
DATA: Any = None
@classmethod
def setUpClass(cls):
if cls.DTYPE is None: raise unittest.SkipTest("base class")
cls.DATA = rand_for_dtype(cls.DTYPE, 0x10, allow_subnormal=cls.DTYPE in supported_dtypes and cls.DTYPE not in dtypes.fp8s)
def test_to_np(self):
a = Tensor(self.DATA, dtype=self.DTYPE)
self.assertEqual(a.dtype, self.DTYPE)
_test_to_np(a, _to_np_dtype(self.DTYPE), np.array(self.DATA, dtype=_to_np_dtype(self.DTYPE)))
def test_casts_to(self):
for dtype in get_available_cast_dtypes(self.DTYPE):
_test_cast(Tensor(self.DATA, dtype=dtype), self.DTYPE)
def test_casts_from(self):
for dtype in get_available_cast_dtypes(self.DTYPE):
_test_cast(Tensor(self.DATA, dtype=self.DTYPE), dtype)
def test_same_size_ops(self):
for dtype in get_available_cast_dtypes(self.DTYPE):
if dtype.itemsize == self.DTYPE.itemsize:
_test_ops(a_dtype=self.DTYPE, b_dtype=dtype)
def test_upcast_ops(self):
for dtype in get_available_cast_dtypes(self.DTYPE):
if dtype.itemsize > self.DTYPE.itemsize:
_test_ops(a_dtype=self.DTYPE, b_dtype=dtype)
def test_upcast_to_ops(self):
for dtype in get_available_cast_dtypes(self.DTYPE):
if dtype.itemsize < self.DTYPE.itemsize:
_test_ops(a_dtype=dtype, b_dtype=self.DTYPE)
def test_bitcast(self):
if self.DTYPE == dtypes.bool: raise unittest.SkipTest("no bools in bitcast")
for dtype in get_available_cast_dtypes(self.DTYPE):
if dtype != dtypes.bool:
_test_bitcast(Tensor(self.DATA[:8], dtype=self.DTYPE), dtype)
@unittest.skipIf(Device.DEFAULT == "PYTHON", "skip for now")
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "skip for now")
def test_uint_overflow(self):
if not dtypes.is_unsigned(self.DTYPE): raise unittest.SkipTest("only for unsigned")
v = self.DTYPE.max
_test_to_np(Tensor(v, dtype=self.DTYPE)+2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))+2)
_test_to_np(Tensor(v, dtype=self.DTYPE)*2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))*2)
def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None):
target_dtype = target_dtype or least_upper_dtype(a_dtype, b_dtype)
if a_dtype == dtypes.bool or b_dtype == dtypes.bool: return
_assert_eq(Tensor([1,2,3,4], dtype=a_dtype)+Tensor([1,2,3,4], dtype=b_dtype), target_dtype, [2,4,6,8])
_assert_eq((Tensor([1], dtype=a_dtype).cast(b_dtype)+Tensor([1], dtype=a_dtype).cast(b_dtype)).cast(a_dtype), a_dtype, [2])
_assert_eq(Tensor([1,2,3,4], dtype=a_dtype)*Tensor([1,2,3,4], dtype=b_dtype), target_dtype, [1,4,9,16])
_assert_eq(Tensor([[1,2],[3,4]], dtype=a_dtype)@Tensor.eye(2, dtype=b_dtype), target_dtype, [[1,2],[3,4]])
_assert_eq(Tensor([1,1,1,1], dtype=a_dtype)+Tensor.ones((4,4), dtype=b_dtype), target_dtype, 2*np.ones((4,4)))
_assert_eq(Tensor([1,1,1,1], dtype=a_dtype)+Tensor.ones((4,4), dtype=b_dtype).clone(), target_dtype, 2*np.ones((4,4)))
_assert_eq(Tensor.ones((4,4), dtype=b_dtype).clone(), b_dtype, np.ones((4,4)))
class TestFp8sConversions(unittest.TestCase):
@given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3_MAX, max_value=FP8E4M3_MAX))
def test_float_to_fp8e4m3(self, x):
np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.float8_e4m3fn).view(torch.uint8).item())
@unittest.skip("fp8 overflow semantics are inconsistent")
def test_float_to_fp8e4m3_extreme_values(self):
for x in [FP8E4M3_MAX, FP8E4M3_MAX*1.01, -FP8E4M3_MAX, -FP8E4M3_MAX*1.01, math.inf, -math.inf, math.nan, -math.nan]:
np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.float8_e4m3fn).view(torch.uint8).item())
@given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E5M2_MAX, max_value=FP8E5M2_MAX))
def test_float_to_fp8e5m2(self, x):
np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.float8_e5m2).view(torch.uint8).item())
def test_float_to_fp8e5m2_extreme_values(self):
for x in [FP8E5M2_MAX, FP8E5M2_MAX*1.01, -FP8E5M2_MAX, -FP8E5M2_MAX*1.01, math.inf, -math.inf, math.nan, -math.nan]:
np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.float8_e5m2).view(torch.uint8).item())
@given(strat.integers(min_value=0, max_value=255))
def test_fp8e4m3_to_float(self, x):
np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e4m3fn).float().item())
@given(strat.integers(min_value=0, max_value=255))
def test_fp8e5m2_to_float(self, x):
np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2).float().item())
@given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3FNUZ_MAX, max_value=FP8E4M3FNUZ_MAX))
def test_float_to_fp8e4m3fnuz(self, x):
np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.float8_e4m3fnuz).view(torch.uint8).item())
def test_float_to_fp8e4m3fnuz_extreme_values(self):
for x in [FP8E4M3FNUZ_MAX, FP8E4M3FNUZ_MAX*1.01, -FP8E4M3FNUZ_MAX, -FP8E4M3FNUZ_MAX*1.01, math.inf, -math.inf, math.nan, 0.0, -0.0]:
np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.float8_e4m3fnuz).view(torch.uint8).item())
@given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E5M2FNUZ_MAX, max_value=FP8E5M2FNUZ_MAX))
def test_float_to_fp8e5m2fnuz(self, x):
np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.float8_e5m2fnuz).view(torch.uint8).item())
def test_float_to_fp8e5m2fnuz_extreme_values(self):
for x in [FP8E5M2FNUZ_MAX, FP8E5M2FNUZ_MAX*1.01, -FP8E5M2FNUZ_MAX, -FP8E5M2FNUZ_MAX*1.01, math.inf, -math.inf, math.nan, 0.0, -0.0]:
np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.float8_e5m2fnuz).view(torch.uint8).item())
@given(strat.integers(min_value=0, max_value=255))
def test_fp8e4m3fnuz_to_float(self, x):
np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e4m3fnuz).float().item())
@given(strat.integers(min_value=0, max_value=255))
def test_fp8e5m2fnuz_to_float(self, x):
np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2fnuz).float().item())
def test_fp8e5m2fnuz_to_float_smallest_normals(self):
# fnuz bias exceeds half's, so exp-1 normals land below half's normal range: they flush to zero like denormals
if dtypes.half not in supported_dtypes or dtypes.half in EMULATED_DTYPES.tolist(dtypes) or dtypes.fp8e5m2fnuz in supported_dtypes:
self.skipTest("needs the emulated fp8 with a native half intermediate")
vals = Tensor([0x04, 0x05, 0x06, 0x07], dtype=dtypes.uint8).bitcast(dtypes.fp8e5m2fnuz).float().numpy()
np.testing.assert_equal(vals, [0., 0., 0., 0.])
class TestBFloat16DType(unittest.TestCase):
def test_bf16_to_float(self):
_test_cast(Tensor([100000], dtype=dtypes.bfloat16), dtypes.float32)
def test_float_to_bf16(self):
_test_cast(Tensor([100000], dtype=dtypes.float32), dtypes.bfloat16)
def test_bf16(self):
t = Tensor([10000, -1, -1000, -10000, 20]).cast(dtypes.bfloat16)
t.realize()
back = t.cast(dtypes.float32)
assert tuple(back.numpy().tolist()) == (9984., -1, -1000, -9984, 20)
class TestBFloat16DTypeCast(unittest.TestCase):
def test_f16_to_bf16_conversion(self):
original_tensor = Tensor([1.0, 2.0, 3.0], dtype=dtypes.float16)
converted_tensor = original_tensor.cast(dtypes.bfloat16)
self.assertEqual(converted_tensor.dtype, dtypes.bfloat16)
back_to_float32 = converted_tensor.cast(dtypes.float32)
original_to_float32 = original_tensor.cast(dtypes.float32)
np.testing.assert_allclose(back_to_float32.numpy(), original_to_float32.numpy(), rtol=1e-2, atol=1e-3)
def test_f16_to_bf16_edge_cases(self):
edge_cases = Tensor([0.0, -0.0, float('inf'), float('-inf'), float('nan')], dtype=dtypes.float16)
converted = edge_cases.cast(dtypes.bfloat16).cast(dtypes.float32)
np.testing.assert_equal(converted.numpy(), edge_cases.cast(dtypes.float32).numpy())
def test_f16_to_bf16_range_precision(self):
large_value = Tensor([65504.0], dtype=dtypes.float16) # Max representable in float16
small_value = Tensor([6.1035e-5], dtype=dtypes.float16) # Smallest positive normal float16
large_converted = large_value.cast(dtypes.bfloat16).cast(dtypes.float32)
small_converted = small_value.cast(dtypes.bfloat16).cast(dtypes.float32)
np.testing.assert_allclose(large_converted.numpy(), large_value.cast(dtypes.float32).numpy(), rtol=1e-2, atol=1e-3)
np.testing.assert_equal(small_converted.numpy(), small_value.cast(dtypes.float32).numpy())
def test_f16_to_bf16_randomized(self):
np.random.seed(42) # For reproducibility
random_values = Tensor(np.random.uniform(-65504, 65504, 1000), dtype=dtypes.float16)
converted = random_values.cast(dtypes.bfloat16).cast(dtypes.float32)
np.testing.assert_allclose(converted.numpy(), random_values.cast(dtypes.float32).numpy(), rtol=1e-2, atol=1e-3)
class TestHalfDType(TestDType): DTYPE = dtypes.half
class TestEmulatedHalf(TestHalfDType):
@classmethod
def setUpClass(cls):
cls.stack = contextlib.ExitStack()
cls.stack.enter_context(Context(EMULATED_DTYPES="half"))
cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
@classmethod
def tearDownClass(cls): cls.stack.close()
class TestFloatDType(TestDType):
DTYPE = dtypes.float
def test_float_to_uint(self):
_test_op(lambda: Tensor([-0.9, -0.3, 1.2], dtype=dtypes.float32).cast(dtypes.uint32), dtypes.uint32,
[0, 0, 1])
@unittest.skipUnless(dtypes.double in supported_dtypes, f"no double on {Device.DEFAULT}")
class TestDoubleDType(TestDType):
DTYPE = dtypes.double
@unittest.skipIf((DEV.interface.startswith("MOCK") and Device.DEFAULT in {"CUDA", "NV"}) or \
isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "conversion not supported on CI CUDA, PTX, and NIR") # TODO: why not?
def test_float64_increased_precision(self):
for func in [
lambda t: t.exp(),
lambda t: t.exp2(),
lambda t: t.log(),
lambda t: t.log2(),
lambda t: t.sqrt(),
lambda t: t.rsqrt(),
lambda t: t.sin(),
lambda t: t.cos(),
lambda t: t.tan(),
lambda t: t.sigmoid(),
]:
a = [2, 3, 4]
np.testing.assert_allclose(func(Tensor(a, dtype=self.DTYPE)).numpy(), func(torch.tensor(a, dtype=torch.float64)), rtol=1e-12, atol=1e-12)
def test_float64_to_float32_cast_inf(self):
_test_op(lambda: Tensor([3.4e40, 3.4e38, 1, 0], dtype=dtypes.float64).cast(dtypes.float32),
dtypes.float32, [float('inf'), 3.4e38, 1, 0])
class TestInt8DType(TestDType):
DTYPE = dtypes.int8
@unittest.skipIf(Device.DEFAULT == "CUDA" or isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "cuda saturation works differently")
def test_int8_to_uint8_negative(self):
_test_op(lambda: Tensor([-1, -2, -3, -4], dtype=dtypes.int8).cast(dtypes.uint8), dtypes.uint8, [255, 254, 253, 252])
def test_int8_to_uint16_negative(self):
_test_op(lambda: Tensor([-1, -2, -3, -4], dtype=dtypes.int8).cast(dtypes.uint16), dtypes.uint16, [2**16-1, 2**16-2, 2**16-3, 2**16-4])
def test_bitcast_alt(self):
a = Tensor([72, -90, 27, 40, -53, 70, 96, 51], dtype=dtypes.int8).bitcast(dtypes.short)
self.assertListEqual(a.tolist(), [-22968, 10267, 18123, 13152])
class TestUint8DType(TestDType):
DTYPE = dtypes.uint8
@unittest.skipIf(Device.DEFAULT == "CUDA" or isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "cuda saturation works differently")
def test_uint8_to_int8_overflow(self):
_test_op(lambda: Tensor([255, 254, 253, 252], dtype=dtypes.uint8).cast(dtypes.int8), dtypes.int8, [-1, -2, -3, -4])
class TestBitCast(unittest.TestCase):
@given(strat.sampled_from(dtype_ints + dtype_floats), strat.sampled_from(dtype_ints + dtype_floats))
def test_shape_change_bitcast(self, dt1, dt2):
data = rand_for_dtype(dt1, 32).reshape(2, 2, 8)
a = Tensor(data, dtype=dt1)
expected = _to_torch_storage(a).view(_to_torch_dtype(dt2))
if dt2 in dtypes.fp8s:
expected = torch.tensor([fp8_to_float(x, dt2) for x in expected.view(-1).tolist()]).view_as(expected)
_test_op(lambda: a.bitcast(dt2), dt2, expected.tolist())
def test_shape_change_bitcast_exceptions(self):
with self.assertRaises(RuntimeError):
# should fail because 3 int8 is 3 bytes but float16 is two and 3 isn't a multiple of 2
Tensor.empty((3,), dtype=dtypes.int8).bitcast(dtypes.float16).shape
def test_bitcast_float_to_int32(self):
a = Tensor([1.,2,3])
b = a.bitcast(dtypes.int32)
assert b.numpy()[0] == 0x3f800000
def test_bitcast_upcasted(self):
a = Tensor.zeros(100, 4, dtype=dtypes.int32).contiguous() + 0x3f800000
b = a.bitcast(dtypes.float32)
assert b.numpy()[0,0] == 1.
def test_bitcast_bf16_from_cast(self):
# a bfloat16 from a cast holds bfloat16 bits. 1.0 is 0x3f80 in bfloat16, which is 1.875 in half
a = Tensor([1.0], dtype=dtypes.float32).cast(dtypes.bfloat16)
assert a.bitcast(dtypes.half).numpy()[0] == 1.875
class TestInt16DType(TestDType): DTYPE = dtypes.int16
class TestUint16DType(TestDType):
DTYPE = dtypes.uint16
def test_uint16_to_int8_overflow(self):
_test_op(lambda: Tensor([2**16-1, 2**16-2, 1, 0], dtype=dtypes.uint16).cast(dtypes.int8), dtypes.int8, [-1, -2, 1, 0])
class TestInt32DType(TestDType): DTYPE = dtypes.int32
class TestUint32DType(TestDType): DTYPE = dtypes.uint32
class TestInt64DType(TestDType): DTYPE = dtypes.int64
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
class TestEmulatedInt64DType(TestInt64DType):
@classmethod
def setUpClass(cls):
cls.stack = contextlib.ExitStack()
cls.stack.enter_context(Context(EMULATED_DTYPES="long"))
cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
@classmethod
def tearDownClass(cls): cls.stack.close()
class TestUint64DType(TestDType):
DTYPE = dtypes.uint64
def test_uint64_load(self):
assert Tensor(2**64 - 1, dtype=dtypes.uint64).numpy() == 2**64 - 1
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
class TestEmulatedUInt64DType(TestUint64DType):
@classmethod
def setUpClass(cls):
cls.stack = contextlib.ExitStack()
cls.stack.enter_context(Context(EMULATED_DTYPES="long"))
cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
@classmethod
def tearDownClass(cls): cls.stack.close()
class TestBoolDType(TestDType): DTYPE = dtypes.bool
class TestBFloat16Type(TestDType): DTYPE = dtypes.bfloat16
class TestEmulatedBFloat16Type(TestBFloat16Type):
@classmethod
def setUpClass(cls):
cls.stack = contextlib.ExitStack()
cls.stack.enter_context(Context(EMULATED_DTYPES="bfloat16"))
cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
@classmethod
def tearDownClass(cls): cls.stack.close()
class TestFp8e4m3(TestDType): DTYPE = dtypes.fp8e4m3
class TestEmulatedFp8e4m3(TestFp8e4m3):
@classmethod
def setUpClass(cls):
cls.stack = contextlib.ExitStack()
cls.stack.enter_context(Context(EMULATED_DTYPES="fp8e4m3"))
cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
@classmethod
def tearDownClass(cls): cls.stack.close()
class TestFp8e5m2(TestDType): DTYPE = dtypes.fp8e5m2
class TestEmulatedFp8e5m2(TestFp8e5m2):
@classmethod
def setUpClass(cls):
cls.stack = contextlib.ExitStack()
cls.stack.enter_context(Context(EMULATED_DTYPES="fp8e5m2"))
cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
@classmethod
def tearDownClass(cls): cls.stack.close()
class TestImplicitFunctionTypeChange(unittest.TestCase):
def test_functions(self):
result = []
for func in [
lambda t: t.exp(),
lambda t: t.exp2(),
lambda t: t.log(),
lambda t: t.log2(),
lambda t: t.sqrt(),
lambda t: t.sin(),
]:
t = func(Tensor([4.0, 3.0])).max() == func(Tensor([4.0, 3.0]))
result.append(t.numpy().sum())
assert all(result)
class TestTensorMethod(unittest.TestCase):
@given(strat.sampled_from(core_dtypes))
def test_abs_diff(self, dt):
if dt == dtypes.bool or dt not in supported_dtypes: return
a, b = Tensor([2], dtype=dt), Tensor([1], dtype=dt)
ret = (a - b).abs()
np.testing.assert_allclose(ret.numpy(), np.abs(a.numpy()-b.numpy()))
class TestDtypeUsage(unittest.TestCase):
def test_max_w_alu(self):
for d in dtypes.ints:
if d in supported_dtypes:
t = Tensor([[1, 2], [3, 4]], dtype=d)
(t*t).max().item()
@unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}")
class TestOpsBFloat16(unittest.TestCase):
def test_cast(self):
# TODO: helper_test_op breaks in unrelated part
data = [60000.0, 70000.0, 80000.0]
np.testing.assert_allclose(Tensor(data).cast("bfloat16").numpy(), torch.tensor(data).type(torch.bfloat16).float().numpy())
# some CPUs there is no native bfloat16 sqrt
@unittest.skipIf(Device.DEFAULT == "CPU", "no approximation")
def test_no_approximation(self):
data = [326.0, 339.0, 10603200512.0]
expected = torch.tensor(data, dtype=torch.bfloat16).sqrt().float().numpy()
np.testing.assert_allclose(Tensor(data, dtype=dtypes.bfloat16).sqrt().numpy(), expected)
if __name__ == '__main__':
unittest.main()

411
artifacts/package_sources/tinygrad/test/backend/test_dtype_alu.py Normal file
View File

@@ -0,0 +1,411 @@
import unittest, operator, math
from tinygrad import Context, Tensor, dtypes, Device
from tinygrad.dtype import DType, truncate, fp8_to_float
from tinygrad.helpers import EMULATED_DTYPES, DEV, getenv
from tinygrad.tensor import _to_np_dtype
from tinygrad.runtime.ops_python import from_storage_scalar
from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.nir import NIRRenderer
from tinygrad.uop import Ops
import numpy as np
import pytest
from hypothesis import assume, given, strategies as strat, settings
pytestmark = pytest.mark.filterwarnings("ignore")
settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
settings.load_profile("my_profile")
print(settings.default)
dtypes_float = (dtypes.float16, dtypes.float32, dtypes.float64)
dtypes_int = (dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64, dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64)
dtypes_bool = (dtypes.bool,)
binary_operations = [operator.add, operator.sub, operator.mul, operator.lt, operator.eq]
integer_binary_operations = binary_operations + [(Tensor.bitwise_xor, np.bitwise_xor), (Tensor.bitwise_and, np.bitwise_and),
(Tensor.bitwise_or, np.bitwise_or), (Tensor.maximum, np.maximum), operator.mod]
integer_unary_operations = [operator.neg]
unary_operations = [(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.sin),
(Tensor.sqrt, np.sqrt), (Tensor.reciprocal, np.reciprocal), (Tensor.cos, np.cos)]
# TODO: enable this (this is a dtype issue)
#binary_operations.append(operator.truediv)
# TODO: CI CUDA segfaults on sin, WEBGPU and NIR sines are not precise enough for large numbers
if ((DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"})
or Device.DEFAULT == "WEBGPU" or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer)):
unary_operations.remove((Tensor.sin, np.sin))
unary_operations.remove((Tensor.cos, np.cos))
# transcendental isn't accurate enough
if Ops.SQRT not in Device[Device.DEFAULT].renderer.code_for_op: unary_operations.remove((Tensor.sqrt, np.sqrt))
supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
class ht:
float64 = strat.floats(width=64, allow_subnormal=False)
float32 = strat.floats(width=32, allow_subnormal=False)
float16 = strat.floats(width=16, allow_subnormal=False)
uint8 = strat.integers(0, 255)
uint16 = strat.integers(0, 65535)
uint32 = strat.integers(0, 2**32-1)
uint64 = strat.integers(0, 2**64-1)
int8 = strat.integers(-128, 127)
int16 = strat.integers(-32768, 32767)
int32 = strat.integers(-2147483648, 2147483647)
int64 = strat.integers(-9223372036854775808, 9223372036854775807)
bool = strat.booleans()
ht.bfloat16 = ht.uint16.filter(lambda x: ((x >> 7) & 0xFF) != 0) # filter subnormal bfloat16
ht.fp8e4m3 = ht.uint8
ht.fp8e5m2 = ht.uint8
ht.fp8e4m3fnuz = ht.uint8
ht.fp8e5m2fnuz = ht.uint8
def universal_test(a, b, dtype, op):
if not isinstance(op, tuple): op = (op, op)
if op[0] == operator.mod and b == 0: return
# lt and max with nan is undefined in tinygrad
if op[0] in (operator.lt, Tensor.maximum) and (math.isnan(a) or math.isnan(b)): return
ta, tb = Tensor([a], dtype=dtype), Tensor([b], dtype=dtype)
if dtype in dtypes.fp8s and op[0] not in (operator.lt, operator.eq):
tensor_value = fp8_to_float((op[0](ta.realize(), tb.realize())).bitcast(dtypes.uint8).item(), dtype)
numpy_value = truncate[dtype](op[1](ta.numpy(), tb.numpy()).item())
else: tensor_value, numpy_value = (op[0](ta, tb)).numpy(), op[1](ta.numpy(), tb.numpy())
if dtype in dtypes.floats:
if dtype not in supported_dtypes or dtype in EMULATED_DTYPES.tolist(dtypes): # denormals are zero
fe, fm = dtypes.finfo(dtype)
atol, rtol = 2 ** (2 - (1 << (fe - 1))), 2 ** (-fm)
else: atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1),
dtypes.fp8e4m3fnuz:(1e-1, 1e-1), dtypes.fp8e5m2fnuz:(5e-1, 5e-1)}.get(dtype, (1e-10, 1e-7))
np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
else: np.testing.assert_equal(tensor_value, numpy_value)
def universal_test_unary(a, dtype, op):
if not isinstance(op, tuple): op = (op, op)
ta = Tensor([a], dtype=dtype)
# TODO: cos does not match for large input
if op[0] == Tensor.cos and abs(a) > 30: return
if op[0] == Tensor.log and a <= 0: return
if dtype in dtypes.fp8s:
# denormals are zero
if (dtype in EMULATED_DTYPES.tolist(dtypes) or dtype not in supported_dtypes
and abs(ta.numpy().item()) < 0.015625): return
tensor_value = fp8_to_float(op[0](ta.realize()).bitcast(dtypes.uint8).item(), dtype)
numpy_value = truncate[dtype](v:=op[1](ta.numpy()).item())
# cuda cast f32 inf to f8 MAX, amd cast it to nan(E4M3)/inf(E5M2)
if math.isinf(v): return
else: tensor_value, numpy_value = op[0](ta).numpy(), op[1](ta.numpy())
if dtype in dtypes.floats:
atol, rtol = { dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 2e-2),
dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2: (1.0, 5e-1),
dtypes.fp8e4m3fnuz:(1e-1, 1e-1), dtypes.fp8e5m2fnuz: (5e-1, 5e-1)}.get(dtype, (1e-6, 1e-5))
np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
else: np.testing.assert_equal(tensor_value, numpy_value)
def universal_test_cast(a, in_dtype, dtype):
tensor_value = Tensor([a], dtype=in_dtype).cast(dtype)
numpy_value = np.array([a], dtype=_to_np_dtype(in_dtype)).astype(_to_np_dtype(dtype))
np.testing.assert_equal(tensor_value.numpy(), numpy_value)
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "Inf and nan cases are wrong on WebGPU")
def universal_test_midcast(a, b, c, op1, op2, d1:DType, d2:DType):
if not isinstance(op1, tuple): op1 = (op1, op1)
if not isinstance(op2, tuple): op2 = (op2, op2)
if op1[0] == operator.mod and b == 0: return
# lt and max with nan is undefined in tinygrad
if op1[0] in (operator.lt, Tensor.maximum) and (math.isnan(a) or math.isnan(b)): return
if op2[0] in (operator.lt, Tensor.maximum) and math.isnan(c): return
at, bt, ct = Tensor([a], dtype=d1), Tensor([b], dtype=d1), Tensor([c], dtype=d2)
an, bn, cn = np.array([a]).astype(_to_np_dtype(d1)), np.array([b]).astype(_to_np_dtype(d1)), np.array([c]).astype(_to_np_dtype(d2))
tensor_value = op2[0](op1[0](at, bt).cast(d2), ct).numpy()
numpy_value = op2[1](op1[1](an, bn).astype(_to_np_dtype(d2)), cn)
np.testing.assert_allclose(tensor_value, numpy_value, rtol=1e-6 if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) else 1e-7)
class TestDTypeALU(unittest.TestCase):
@unittest.skipUnless(dtypes.float64 in supported_dtypes, f"no float64 on {Device.DEFAULT}")
@given(ht.float64, ht.float64, strat.sampled_from(binary_operations))
def test_float64(self, a, b, op): universal_test(a, b, dtypes.float64, op)
@given(ht.float32, ht.float32, strat.sampled_from(binary_operations))
def test_float32(self, a, b, op): universal_test(a, b, dtypes.float32, op)
@unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
@given(ht.float16, ht.float16, strat.sampled_from(binary_operations))
def test_float16(self, a, b, op): universal_test(a, b, dtypes.float16, op)
@given(ht.float16, ht.float16, strat.sampled_from(binary_operations))
@Context(EMULATED_DTYPES="half")
def test_emulated_float16(self, a, b, op): universal_test(a, b, dtypes.float16, op)
@unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}")
@given(ht.bfloat16, ht.bfloat16, strat.sampled_from(binary_operations))
def test_bfloat16(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(b, dtypes.bfloat16), dtypes.bfloat16, op)
@given(ht.bfloat16, ht.bfloat16, strat.sampled_from(binary_operations))
@Context(EMULATED_DTYPES="bfloat16")
def test_emulated_bfloat16(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(b, dtypes.bfloat16), dtypes.bfloat16, op)
@unittest.skipUnless(dtypes.fp8e4m3 in supported_dtypes, f"no fp8e4m3 on {Device.DEFAULT}")
@given(ht.fp8e4m3, ht.fp8e4m3, strat.sampled_from(binary_operations))
def test_fp8e4m3(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.fp8e4m3), from_storage_scalar(b, dtypes.fp8e4m3), dtypes.fp8e4m3, op)
@given(ht.fp8e4m3, ht.fp8e4m3, strat.sampled_from(binary_operations))
@Context(EMULATED_DTYPES="fp8e4m3")
def test_emulated_fp8e4m3(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.fp8e4m3), from_storage_scalar(b, dtypes.fp8e4m3), dtypes.fp8e4m3, op)
@unittest.skipUnless(dtypes.fp8e5m2 in supported_dtypes, f"no fp8e5m2 on {Device.DEFAULT}")
@given(ht.fp8e5m2, ht.fp8e5m2, strat.sampled_from(binary_operations))
def test_fp8e5m2(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.fp8e5m2), from_storage_scalar(b, dtypes.fp8e5m2), dtypes.fp8e5m2, op)
@given(ht.fp8e5m2, ht.fp8e5m2, strat.sampled_from(binary_operations))
@Context(EMULATED_DTYPES="fp8e5m2")
def test_emulated_fp8e5m2(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.fp8e5m2), from_storage_scalar(b, dtypes.fp8e5m2), dtypes.fp8e5m2, op)
@unittest.skipUnless(dtypes.fp8e4m3fnuz in supported_dtypes, f"no fp8e4m3fnuz on {Device.DEFAULT}")
@given(ht.fp8e4m3fnuz, ht.fp8e4m3fnuz, strat.sampled_from(binary_operations))
def test_fp8e4m3fnuz(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.fp8e4m3fnuz), from_storage_scalar(b, dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
@unittest.skipUnless(dtypes.fp8e5m2fnuz in supported_dtypes, f"no fp8e5m2fnuz on {Device.DEFAULT}")
@given(ht.fp8e5m2fnuz, ht.fp8e5m2fnuz, strat.sampled_from(binary_operations))
def test_fp8e5m2fnuz(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.fp8e5m2fnuz), from_storage_scalar(b, dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
@given(ht.fp8e4m3fnuz, ht.fp8e4m3fnuz, strat.sampled_from(binary_operations))
@Context(EMULATED_DTYPES="fp8e4m3fnuz")
def test_emulated_fp8e4m3fnuz(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.fp8e4m3fnuz), from_storage_scalar(b, dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
@given(ht.fp8e5m2fnuz, ht.fp8e5m2fnuz, strat.sampled_from(binary_operations))
@Context(EMULATED_DTYPES="fp8e5m2fnuz")
def test_emulated_fp8e5m2fnuz(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.fp8e5m2fnuz), from_storage_scalar(b, dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
@given(ht.float32, strat.sampled_from(unary_operations))
def test_float32_unary(self, a, op): universal_test_unary(a, dtypes.float32, op)
@unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
@given(ht.float16, strat.sampled_from(unary_operations))
def test_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op)
@given(ht.float16, strat.sampled_from(unary_operations))
@Context(EMULATED_DTYPES="half")
def test_emulated_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op)
@unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}")
@given(ht.bfloat16, strat.sampled_from(unary_operations))
def test_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)
@given(ht.bfloat16, strat.sampled_from(unary_operations))
@Context(EMULATED_DTYPES="bfloat16")
def test_emulated_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)
@unittest.skipUnless(dtypes.fp8e4m3 in supported_dtypes, f"no fp8e4m3 on {Device.DEFAULT}")
@given(ht.fp8e4m3, strat.sampled_from(unary_operations))
def test_fp8e4m3_unary(self, a, op):
if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3) != 0.0)
universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3), dtypes.fp8e4m3, op)
@given(ht.fp8e4m3, strat.sampled_from(unary_operations))
@Context(EMULATED_DTYPES="fp8e4m3")
def test_emulated_fp8e4m3_unary(self, a, op):
if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3) != 0.0)
universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3), dtypes.fp8e4m3, op)
@unittest.skipUnless(dtypes.fp8e5m2 in supported_dtypes, f"no fp8e5m2 on {Device.DEFAULT}")
@given(ht.fp8e5m2, strat.sampled_from(unary_operations))
def test_fp8e5m2_unary(self, a, op):
if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2) != 0.0)
universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2), dtypes.fp8e5m2, op)
@given(ht.fp8e5m2, strat.sampled_from(unary_operations))
@Context(EMULATED_DTYPES="fp8e5m2")
def test_emulated_fp8e5m2_unary(self, a, op):
if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2) != 0.0)
universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2), dtypes.fp8e5m2, op)
@unittest.skipUnless(dtypes.fp8e4m3fnuz in supported_dtypes, f"no fp8e4m3fnuz on {Device.DEFAULT}")
@given(ht.fp8e4m3fnuz, strat.sampled_from(unary_operations))
def test_fp8e4m3fnuz_unary(self, a, op):
if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz) != 0.0)
universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
@unittest.skipUnless(dtypes.fp8e5m2fnuz in supported_dtypes, f"no fp8e5m2fnuz on {Device.DEFAULT}")
@given(ht.fp8e5m2fnuz, strat.sampled_from(unary_operations))
def test_fp8e5m2fnuz_unary(self, a, op):
if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz) != 0.0)
universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
@given(ht.fp8e4m3fnuz, strat.sampled_from(unary_operations))
@Context(EMULATED_DTYPES="fp8e4m3fnuz")
def test_emulated_fp8e4m3fnuz_unary(self, a, op):
if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz) != 0.0)
universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
@given(ht.fp8e5m2fnuz, strat.sampled_from(unary_operations))
@Context(EMULATED_DTYPES="fp8e5m2fnuz")
def test_emulated_fp8e5m2fnuz_unary(self, a, op):
if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz) != 0.0)
universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
@given(ht.uint8, ht.uint8, strat.sampled_from(integer_binary_operations))
def test_uint8(self, a, b, op): universal_test(a, b, dtypes.uint8, op)
@unittest.skipUnless(dtypes.uint16 in supported_dtypes, f"no uint16 on {Device.DEFAULT}")
@given(ht.uint16, ht.uint16, strat.sampled_from(integer_binary_operations))
def test_uint16(self, a, b, op): universal_test(a, b, dtypes.uint16, op)
@unittest.skipUnless(dtypes.uint32 in supported_dtypes, f"no uint32 on {Device.DEFAULT}")
@given(ht.uint32, ht.uint32, strat.sampled_from(integer_binary_operations))
def test_uint32(self, a, b, op): universal_test(a, b, dtypes.uint32, op)
@unittest.skipUnless(dtypes.uint64 in supported_dtypes, f"no uint64 on {Device.DEFAULT}")
@given(ht.uint64, ht.uint64, strat.sampled_from(integer_binary_operations))
def test_uint64(self, a, b, op): universal_test(a, b, dtypes.uint64, op)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
@given(ht.uint64, ht.uint64, strat.sampled_from(integer_binary_operations))
@Context(EMULATED_DTYPES="long")
def test_emulated_uint64(self, a, b, op): universal_test(a, b, dtypes.uint64, op)
@given(ht.int8, ht.int8, strat.sampled_from(integer_binary_operations))
def test_int8(self, a, b, op): universal_test(a, b, dtypes.int8, op)
@given(ht.int16, ht.int16, strat.sampled_from(integer_binary_operations))
def test_int16(self, a, b, op): universal_test(a, b, dtypes.int16, op)
@given(ht.int32, ht.int32, strat.sampled_from(integer_binary_operations))
def test_int32(self, a, b, op): universal_test(a, b, dtypes.int32, op)
@given(ht.int64, ht.int64, strat.sampled_from(integer_binary_operations))
def test_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
@given(ht.int64, ht.int64, strat.sampled_from(integer_binary_operations))
@Context(EMULATED_DTYPES="long")
def test_emulated_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op)
def _test_shl(self):
for dtype, values, distances in ((dtypes.int64, [-0x1234, 0x80000001, -1, 0x1234, 1], [0, 5, 31, 32, 62]),
(dtypes.uint64, [0x80000001, 0x80000001, 1, 0xFEDC, 1], [0, 5, 31, 32, 62]),
(dtypes.int8, [-3, 1, 7, -2, 1], [0, 1, 3, 5, 6]),
(dtypes.uint16, [3, 1, 0xFF, 7, 1], [0, 1, 7, 12, 15])):
with self.subTest(dtype=dtype):
result = Tensor(values, dtype=dtype) << Tensor(distances, dtype=dtype)
np.testing.assert_equal(result.numpy(), [x << d for x, d in zip(values, distances)])
def _test_shr(self):
for dtype, values, distances in ((dtypes.int64, [-(2**40), -1, -(2**50), -(2**40), 0x123456789ABCDEF], [0, 5, 31, 32, 63]),
(dtypes.uint64, [0xFEDCBA9876543210] * 5, [0, 5, 31, 32, 63]),
(dtypes.int8, [-128, -1, 64, -37, 1], [0, 1, 3, 5, 7]),
(dtypes.uint16, [0xFFFF] * 5, [0, 1, 8, 13, 15])):
with self.subTest(dtype=dtype):
result = Tensor(values, dtype=dtype) >> Tensor(distances, dtype=dtype)
np.testing.assert_equal(result.numpy(), [x >> d for x, d in zip(values, distances)])
def test_shl(self): self._test_shl()
def test_shr(self): self._test_shr()
@Context(EMULATED_DTYPES="long")
def test_emulated_shl(self): self._test_shl()
@Context(EMULATED_DTYPES="long")
def test_emulated_shr(self): self._test_shr()
@given(ht.uint8, strat.sampled_from(integer_unary_operations))
def test_uint8_unary(self, a, op): universal_test_unary(a, dtypes.uint8, op)
@unittest.skipUnless(dtypes.uint16 in supported_dtypes, f"no uint16 on {Device.DEFAULT}")
@given(ht.uint16, strat.sampled_from(integer_unary_operations))
def test_uint16_unary(self, a, op): universal_test_unary(a, dtypes.uint16, op)
@unittest.skipUnless(dtypes.uint32 in supported_dtypes, f"no uint32 on {Device.DEFAULT}")
@given(ht.uint32, strat.sampled_from(integer_unary_operations))
def test_uint32_unary(self, a, op): universal_test_unary(a, dtypes.uint32, op)
@unittest.skipUnless(dtypes.uint64 in supported_dtypes, f"no uint64 on {Device.DEFAULT}")
@given(ht.uint64, strat.sampled_from(integer_unary_operations))
def test_uint64_unary(self, a, op): universal_test_unary(a, dtypes.uint64, op)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
@given(ht.uint64, strat.sampled_from(integer_unary_operations))
@Context(EMULATED_DTYPES="long")
def test_emulated_uint64_unary(self, a, op): universal_test_unary(a, dtypes.uint64, op)
@given(ht.int8, strat.sampled_from(integer_unary_operations))
def test_int8_unary(self, a, op): universal_test_unary(a, dtypes.int8, op)
@given(ht.int16, strat.sampled_from(integer_unary_operations))
def test_int16_unary(self, a, op): universal_test_unary(a, dtypes.int16, op)
@given(ht.int32, strat.sampled_from(integer_unary_operations))
def test_int32_unary(self, a, op): universal_test_unary(a, dtypes.int32, op)
@given(ht.int64, strat.sampled_from(integer_unary_operations))
def test_int64_unary(self, a, op): universal_test_unary(a, dtypes.int64, op)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
@given(ht.int64, strat.sampled_from(integer_unary_operations))
@Context(EMULATED_DTYPES="long")
def test_emulated_int64_unary(self, a, op): universal_test_unary(a, dtypes.int64, op)
@given(ht.bool, ht.bool, strat.sampled_from(((operator.add, operator.add), (operator.mul, operator.mul))))
def test_bool(self, a, b, op): universal_test(a, b, dtypes.bool, op)
@given(ht.int32, ht.int32, ht.float32, strat.sampled_from(integer_binary_operations), strat.sampled_from(binary_operations))
def test_int32_midcast_float(self, a, b, c, op1, op2): universal_test_midcast(a, b, c, op1, op2, dtypes.int32, dtypes.float32)
# Metal and (MOCK)CUDA and HIP and NIR behave differently than numpy for overflows
skip_overflow = ((DEV.interface.startswith("MOCK") and Device.DEFAULT in {"AMD", "NV", "CUDA"})
or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer))
@given(strat.floats(width=32, min_value=0, max_value=10.0) if skip_overflow else ht.float32,
strat.floats(width=32, min_value=0, max_value=10.0) if skip_overflow else ht.float32,
ht.int32, strat.sampled_from(binary_operations), strat.sampled_from(integer_binary_operations))
@unittest.skipIf(Device.DEFAULT == "PYTHON", "TODO: fix cast inf to int32 in PYTHON")
@unittest.skip("broken on Mac")
def test_float_midcast_int32(self, a, b, c, op1, op2): universal_test_midcast(a, b, c, op1, op2, dtypes.float32, dtypes.int32)
@unittest.skip("broken. TODO: fix it")
@given(ht.float32, strat.sampled_from(dtypes_float+dtypes_int+dtypes_bool))
def test_float_cast(self, a, dtype): universal_test_cast(a, dtypes.float32, dtype)
@unittest.skip("broken. TODO: fix it")
@given(ht.int32, strat.sampled_from(dtypes_float+dtypes_int+dtypes_bool))
def test_int32_cast(self, a, dtype): universal_test_cast(a, dtypes.int32, dtype)
@given(strat.floats(width=32, min_value=1.0, max_value=254.0, allow_subnormal=False),
strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
def test_float_cast_to_unsigned(self, a, float_dtype, unsigned_dtype):
if float_dtype not in supported_dtypes: float_dtype = dtypes.float32
universal_test_cast(a, float_dtype, unsigned_dtype)
@unittest.skip("relied on hacks")
@given(strat.floats(width=32, min_value=256.0, max_value=65000.0, allow_subnormal=False),
strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
def test_float_cast_to_unsigned_overflow(self, a, float_dtype, unsigned_dtype):
if float_dtype not in supported_dtypes: float_dtype = dtypes.float32
universal_test_cast(a, float_dtype, unsigned_dtype)
@unittest.skip("relied on hacks")
@given(strat.floats(width=32, min_value=-65000.0, max_value=-1.0, allow_subnormal=False),
strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
def test_float_cast_to_unsigned_underflow(self, a, float_dtype, unsigned_dtype):
if float_dtype not in supported_dtypes: float_dtype = dtypes.float32
universal_test_cast(a, float_dtype, unsigned_dtype)
def test_unsafe_cast_float_to_int(self):
# the value is off the float32 grid but rounds in-range: the buffer and const-fold paths must agree
# (out-of-range float->int cast stays undefined: hardware may saturate where the fold wraps)
val = 2147483000.0
t1 = Tensor([val], dtype=dtypes.float32).cast(dtypes.int32)
t2 = Tensor(val, dtype=dtypes.float32).cast(dtypes.int32)
np.testing.assert_equal(t1.item(), t2.item())
if __name__ == '__main__':
unittest.main()

246
artifacts/package_sources/tinygrad/test/backend/test_edgecases.py Normal file
View File

@@ -0,0 +1,246 @@
# end to end tests of tinygrad that you think might be edge cases.
# using the documentation, write code you think should work.
# you can compare the outputs to torch or numpy, or just tinygrad assert/raise while doing things that should be valid
# i'm not interested in tests that currently pass, i'm only interested in tests that you think should pass but don't.
# mark them with @unittest.expectedFailure
# all the tests in here didn't pass until bugs were fixed
# get creative! think about things that failed in pytorch or tensorflow for a long time until they were fixed.
# every test should surface a unique bug. if tinygrad throws an error saying something is not supported, this is probably not a bug.
# the tests don't have to test the same parts of the code that these current ones test, more diversity is better
# focus on making tinygrad throw runtime errors or assertions for valid things, or find clear numerical mismatches from pytorch
# confirm any bugs found are valid by doing the same thing in pytorch in the test.
# for any failing tests, explain in a comment why tinygrad is wrong and what the desired behavior should be.
# don't worry about running mypy or linters. focus on writing more of these tests and running them to confirm broken behavior.
# surface level bugs, like issues with empty tensors, input validation, or nans, are not that interesting.
# focus on bugs that would frustrate real users.
# these are not bugs, these are desired behavior. don't add failing tests for them:
# tinygrad only accepts tinygrad dtypes or strings of the tinygrad dtype.
# boolean indexing, or anything with unknown output shape of tensor at compile time isn't supported.
# invalid indexing in things like gather and one_hot is not an error in tinygrad. nothing that depends on the value is
# repeat_interleave doesn't support a tensor as the dim. check tinygrad type signature before claiming something is a bug
import unittest
import numpy as np
import torch
from tinygrad import Tensor, dtypes, nn, Context
from tinygrad.device import Device
from tinygrad.helpers import DEV
from tinygrad.renderer.nir import NIRRenderer
MOCKGPU = DEV.interface.startswith("MOCK")
class TestNaNEdgeCases(unittest.TestCase):
# we don't need more of these. it's unclear if torch's behavior is desired here
@unittest.expectedFailure
def test_max_nan(self):
# Reductions with NaN should propagate NaN like PyTorch.
arr = [1.0, float('nan'), 3.0]
torch_out = torch.tensor(arr).max().item()
out = Tensor(arr).max().numpy()
if np.isnan(torch_out):
self.assertTrue(np.isnan(out))
else:
np.testing.assert_equal(out, torch_out)
@unittest.skip("passes on webgpu")
@unittest.expectedFailure
def test_argmax_nan(self):
# PyTorch returns the index of the NaN, tinygrad returns the index of the maximum value.
arr = [1.0, float('nan'), 3.0]
torch_idx = torch.tensor(arr).argmax().item()
idx = Tensor(arr).argmax().item()
self.assertEqual(idx, torch_idx)
@unittest.expectedFailure
def test_sort_with_nan(self):
# Sorting a tensor containing NaN should keep NaN at the end like PyTorch.
arr = [1.0, float('nan'), 3.0]
torch_vals, torch_idxs = torch.tensor(arr).sort()
vals, idxs = Tensor(arr).sort()
np.testing.assert_equal(vals.numpy(), torch_vals.numpy())
np.testing.assert_equal(idxs.numpy(), torch_idxs.numpy().astype(np.int32))
class TestEmptyTensorEdgeCases(unittest.TestCase):
# we don't need more of these
def test_sort_empty(self):
# Sorting an empty tensor works in PyTorch and should return empty
# values and indices. tinygrad raises an error instead.
torch_vals, torch_idxs = torch.tensor([]).sort()
values, indices = Tensor([]).sort()
np.testing.assert_equal(values.numpy(), torch_vals.numpy())
np.testing.assert_equal(indices.numpy(), torch_idxs.numpy().astype(np.int32))
@unittest.expectedFailure
def test_max_empty(self):
# Max on an empty tensor should also raise an error.
with self.assertRaises(RuntimeError):
torch.tensor([]).max()
with self.assertRaises(RuntimeError):
Tensor([]).max()
@unittest.expectedFailure
def test_argmax_empty(self):
# Argmax on an empty tensor should raise an error like torch does.
with self.assertRaises(RuntimeError):
torch.tensor([]).argmax()
with self.assertRaises(RuntimeError):
Tensor([]).argmax()
def test_masked_select_empty(self):
# Masked select on empty tensors should return an empty tensor.
torch_out = torch.tensor([], dtype=torch.float32).masked_select(torch.tensor([], dtype=torch.bool))
out = Tensor([], dtype=dtypes.float32).masked_select(Tensor([], dtype=dtypes.bool))
np.testing.assert_equal(out.numpy(), torch_out.numpy())
class TestDropoutProbabilityEdgeCases(unittest.TestCase):
# we don't need more of these
def test_dropout_rate_one(self):
with Context(TRAINING=1):
out = Tensor.ones(100).dropout(1.0)
np.testing.assert_allclose(out.numpy(), np.zeros(100))
def test_dropout_invalid_prob(self):
with self.assertRaises(ValueError):
torch.nn.functional.dropout(torch.ones(10), -0.1, True)
with self.assertRaises(ValueError):
with Context(TRAINING=1):
Tensor.ones(10).dropout(-0.1)
class TestInputValidation(unittest.TestCase):
# we don't need more of these, input validation bugs are not very interesting, many are WONTFIX
@unittest.expectedFailure
def test_repeat_negative(self):
# repeating with a negative value should error like PyTorch
with self.assertRaises(RuntimeError):
torch.tensor([1, 2, 3]).repeat(-1, 2)
with self.assertRaises(RuntimeError):
Tensor([1, 2, 3]).repeat(-1, 2)
def test_negative_weight_decay(self):
with self.assertRaises(ValueError):
torch.optim.AdamW([torch.tensor([1.], requires_grad=True)], lr=0.1, weight_decay=-0.1)
with self.assertRaises(ValueError):
nn.optim.AdamW([Tensor([1.])], lr=0.1, weight_decay=-0.1)
def test_negative_lr(self):
with self.assertRaises(ValueError):
torch.optim.SGD([torch.tensor([1.], requires_grad=True)], lr=-0.1)
with self.assertRaises(ValueError):
nn.optim.SGD([Tensor([1.])], lr=-0.1)
def test_negative_momentum(self):
with self.assertRaises(ValueError):
torch.optim.SGD([torch.tensor([1.], requires_grad=True)], lr=0.1, momentum=-0.1)
with self.assertRaises(ValueError):
nn.optim.SGD([Tensor([1.])], lr=0.1, momentum=-0.1)
class TestZeroFolding(unittest.TestCase):
# we don't need more of these
# folding rules treat x/x, x//x and x%x as constants even when x can be zero
@unittest.expectedFailure
def test_divide_by_self_with_zero(self):
x = Tensor([0.0, 1.0])
torch_out = torch.tensor([0.0, 1.0]) / torch.tensor([0.0, 1.0])
out = (x / x).numpy()
np.testing.assert_allclose(out, torch_out.numpy(), equal_nan=True)
@unittest.expectedFailure
def test_floordiv_by_self_with_zero(self):
x = Tensor([0])
with self.assertRaises(RuntimeError):
torch.tensor([0]) // torch.tensor([0])
with self.assertRaises(RuntimeError):
(x // x).numpy()
@unittest.expectedFailure
def test_mod_by_self_with_zero(self):
x = Tensor([0])
with self.assertRaises(RuntimeError):
torch.tensor([0]) % torch.tensor([0])
with self.assertRaises(RuntimeError):
(x % x).numpy()
class TestAssignIssues(unittest.TestCase):
# these are good failures. i'm not sure we need more, but we need to fix these.
def test_assign_permuted_view_constant(self):
# assigning to a permuted view should modify the underlying tensor
arr = np.arange(6).reshape(2, 3).astype(np.float32)
torch_tensor = torch.tensor(arr)
torch_tensor.t().copy_(torch.tensor([[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]))
t = Tensor(arr).contiguous().realize()
t.permute(1, 0).assign(Tensor([[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]))
np.testing.assert_allclose(t.numpy(), torch_tensor.numpy())
def test_assign_shrink_view_constant(self):
# assigning to a shrunk view should update the base tensor
arr = np.arange(9).reshape(3, 3).astype(np.float32)
torch_tensor = torch.tensor(arr)
torch_tensor[1:3, 1:3] = torch.ones(2, 2)
t = Tensor(arr).contiguous().realize()
t.shrink(((1, 3), (1, 3))).assign(Tensor.ones(2, 2))
np.testing.assert_allclose(t.numpy(), torch_tensor.numpy())
@unittest.expectedFailure
def test_assign_broadcast(self):
# broadcasting during assign should behave like PyTorch
# NOTE: we don't want implicit dtype casting (int64 -> float32 loses precision), so this fails
torch_tensor = torch.zeros(3, 5)
torch_tensor[:] = torch.arange(5)
t = Tensor.zeros(3, 5)
t.assign(Tensor.arange(5))
np.testing.assert_allclose(t.numpy(), torch_tensor.numpy())
class TestUOpValidationIssue(unittest.TestCase):
# these fail with UOp verification error.
# we want more of these with diverse errors!
@unittest.skipIf(MOCKGPU or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer), "hangs gpuocelot, NIR cannot render")
def test_tensor_index_overflow(self):
val = Tensor([1])
big = val.expand(2**31 + 3)
idx = Tensor([0, 2**31 + 2])
np.testing.assert_equal(big[idx].numpy(), np.array([1, 1]))
def test_float_floordiv_scalar(self):
(Tensor.arange(4, dtype=dtypes.float32) // 2).realize()
def test_float_floordiv_tensor(self):
(Tensor.arange(4, dtype=dtypes.float32) // Tensor.ones(4, dtype=dtypes.float32)).realize()
class TestEdgeCases(unittest.TestCase):
# add tests exposing new and diverse kinds of bugs that might impact real users here
def test_circular_pad_negative(self):
# negative pads with circular mode should wrap like PyTorch
arr = np.arange(9).reshape(1, 1, 3, 3).astype(np.float32)
torch_out = torch.nn.functional.pad(torch.tensor(arr), (1, -1, 1, -1), mode='circular')
out = Tensor(arr).pad((1, -1, 1, -1), mode='circular')
np.testing.assert_equal(out.numpy(), torch_out.numpy())
def test_arange_float_step(self):
# float steps should match PyTorch exactly
torch_out = torch.arange(0, 2, 0.3).numpy()
out = Tensor.arange(0, 2, 0.3).numpy()
np.testing.assert_allclose(out, torch_out, atol=1e-7)
@unittest.skip("this is flaky")
@unittest.expectedFailure
def test_topk_ties_indices(self):
# topk should match PyTorch tie-breaking behavior when values are equal
arr = [1.0, 1.0, 1.0, 1.0]
_, ti = torch.tensor(arr).topk(2)
_, i = Tensor(arr).topk(2)
np.testing.assert_equal(i.numpy(), ti.numpy().astype(np.int32))
if __name__ == "__main__":
unittest.main()

151
artifacts/package_sources/tinygrad/test/backend/test_encodings.py Normal file
View File

@@ -0,0 +1,151 @@
import unittest
from tinygrad import Device
from tinygrad.uop.ops import UOp, Ops
from tinygrad.dtype import dtypes
from tinygrad.renderer.isa.x86 import X86Ops, X86Renderer, RBP, RDI, RSP, RSI, RAX, RDX, XMM, GPR, imm, def_reg
def ins(op, dt, src, tag=None): return UOp(Ops.INS, arg=op, dtype=dt, src=src, tag=tag)
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only on x86")
class TestEncodingsX86(unittest.TestCase):
# NOTE: x86 supports a single displacement as memory address and index without base memory address
# these have no use cases so they aren't supported
def encode(self, u:UOp): return Device[Device.DEFAULT].renderer.render([u])
# displacement of 0 isn't emitted
def test_base_address(self):
load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RDI), UOp(Ops.NOOP), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), RDI)
# mov edi, dword ptr [rdi]
self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 3F"))
# rsp/r12 require a sib byte when used as base memory address
def test_rsp_base_address(self):
load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), RSP)
# mov esp, dword ptr [rsp]
self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 24 24"))
# rbp/r13 require a displacement when used as base memory address
def test_rbp_base_address(self):
load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RBP), UOp(Ops.NOOP), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), RBP)
# mov ebp, dword ptr [rbp + 0]
self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 6D 00"))
# test [base + index*scale]
def test_base_index_address(self):
load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RAX), def_reg(dtypes.int32, RDX), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), RAX)
# mov eax, dword ptr [rax + rdx*4]
self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 04 90"))
# rsp as index means no index
def test_rsp_index_address(self):
load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RAX), def_reg(dtypes.int32, RSP), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), RAX)
# mov eax, dword ptr [rax]
self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 00"))
# however r12 is a valid index
def test_r12_index_address(self):
load = ins(X86Ops.MOV, dtypes.int32,
(def_reg(dtypes.uint64, RAX), def_reg(dtypes.int32, GPR[12]), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), RAX)
# mov eax, dword ptr [rax + r12*4]
self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("42 8B 04 A0"))
# test [base + index*scale + 8bit disp]
def test_complex_address_8bit_disp(self):
load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10), imm(dtypes.uint8, 4)), RDI)
# mov edi, dword ptr [rdi + rsi*4 + 0xa]
self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 7C B7 0A"))
# test [base + index*scale + 32bit disp]
def test_complex_address_32bit_disp(self):
load = ins(X86Ops.MOV, dtypes.int32,
(def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int32, 10000), imm(dtypes.uint8, 4)), RDI)
# mov edi, dword ptr [rdi + rsi*4 + 0x2710]
self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B BC B7 10 27 00 00"))
# 8bit variants of legacy instructions subtract 1 from opcode
def test_8bit_legacy_encoding(self):
cast = ins(X86Ops.MOVSX, dtypes.int32, (def_reg(dtypes.int8, RDX),), RAX)
# movsx eax, dl
self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("0F BE C2"))
# accessing lower 8 bits of rsp, rbp, rsi, rdi requires rex prefix
def test_lower_8bits_reg(self):
cast = ins(X86Ops.MOVSX, dtypes.int32, (def_reg(dtypes.int8, RDI),), RAX)
# movsx eax, dil
self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("40 0F BE C7"))
# test 16 bit variant of legacy instruction
def test_16bit_legacy_encoding(self):
cast = ins(X86Ops.MOVSX, dtypes.int16, (def_reg(dtypes.int8, RDX),), RAX)
# movsx ax, dl
self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("66 0F BE C2"))
# test 64 bit variant of legacy instruction
def test_64bit_legacy_encoding(self):
cast = ins(X86Ops.MOVSX, dtypes.int64, (def_reg(dtypes.int8, RDX),), RAX)
# movsx rax, dl
self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("48 0F BE C2"))
# test compact vex encoding
def test_compact_vex_encoding(self):
xmm0, xmm1 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[1])
add = ins(X86Ops.VADDSS, dtypes.float32, (xmm0, xmm1), XMM[0])
# vaddss xmm0, xmm0, xmm1
self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C5 FA 58 C1"))
# test long vex encoding
def test_long_vex_encoding(self):
xmm0, xmm8 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[8])
add = ins(X86Ops.VADDSS, dtypes.float32, (xmm0, xmm8), XMM[0])
# vaddss xmm0, xmm0, xmm8
self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C4 C1 7A 58 C0"))
# test ymm encoding
def test_ymm_encoding(self):
xmm0, xmm1 = def_reg(dtypes._uint256, XMM[0]), def_reg(dtypes._uint256, XMM[1])
add = ins(X86Ops.VADDPS, dtypes._uint256, (xmm0, xmm1), XMM[0])
# vaddps ymm0, ymm0, ymm1
self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C5 FC 58 C1"))
# test encoding where register is in the immediate field
def test_reg_in_imm_field(self):
xmm0, xmm1, xmm2 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[1]), def_reg(dtypes.float32, XMM[2])
blend = ins(X86Ops.VBLENDVPS, dtypes.float32, (xmm0, xmm1, xmm2), XMM[0])
# vblendvps xmm0, xmm0, xmm1, xmm2
self.assertEqual(bytes.fromhex(self.encode(blend)), bytes.fromhex("C4 E3 79 4A C1 20"))
# when writting to mem the uop takes the store form where dtype is void and there's no definition
def test_write_mem(self):
address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10), imm(dtypes.uint8, 4))
xmm0 = def_reg(dtypes.float32, XMM[0])
extr = ins(X86Ops.VPEXTRD, dtypes.void, address + (xmm0, imm(dtypes.uint8, 0)))
# vpextrd dword ptr [rdi + rsi*4 + 0xa], xmm0, 0
self.assertEqual(bytes.fromhex(self.encode(extr)), bytes.fromhex("C4 E3 79 16 44 B7 0A 00"))
# test two address instruction with fused load works
def test_two_address_load(self):
address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10), imm(dtypes.uint8, 4))
cmove = ins(X86Ops.CMOVE, dtypes.int32, address, RAX)
# cmove eax, dword ptr [rdi + rsi*4 + 0xa]
self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 44 B7 0A"))
# test instruction where displacement and imm have the same value
def test_disp_imm_same_value(self):
address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int8, RSI), imm(dtypes.int8, 10), imm(dtypes.uint8, 1))
mov = ins(X86Ops.MOVi, dtypes.void, address + (imm(dtypes.int8, 10),))
# mov byte ptr [rdi + rsi + 0xa], 0xa
self.assertEqual(bytes.fromhex(self.encode(mov)), bytes.fromhex("40 C6 44 37 0A 0A"))
address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int32, 10), imm(dtypes.uint8, 4))
imul = ins(X86Ops.IMULi, dtypes.int32, address + (imm(dtypes.int32, 10),), RDI)
# imul edi, dword ptr [rdi + rsi*4 + 0xa], 0xa
self.assertEqual(bytes.fromhex(self.encode(imul)), bytes.fromhex("69 BC B7 0A 00 00 00 0A 00 00 00"))
# cmoves have the cmp as the last src even though it is not explicitly used, the cmp doesn't define a reg and is ignored in the encoding
def test_cmove_ignore_cmp(self):
cmove = ins(X86Ops.CMOVE, dtypes.int32, (def_reg(dtypes.int32, RAX), UOp(Ops.INS, arg=X86Ops.CMP)), RDX)
# cmove edx, eax
self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 D0"))
if __name__ == "__main__":
unittest.main()

319
artifacts/package_sources/tinygrad/test/backend/test_graph.py Normal file
View File

@@ -0,0 +1,319 @@
import numpy as np
import functools, unittest
from tinygrad.device import Device, Buffer
from tinygrad.tensor import Tensor
from tinygrad.helpers import Context
from tinygrad.dtype import dtypes
from tinygrad.engine.jit import MultiGraphRunner
from tinygrad.engine.realize import run_linear, compile_linear
from tinygrad.uop.ops import UOp, Ops, buffers
from test.helpers import needs_second_gpu
np.random.seed(1337)
Tensor.manual_seed(1337)
BUF_SIZE = 4096
RUN_CNT = 5
# cache AST by (device, num_inputs)
cached_asts: dict[tuple[str, int], UOp] = {}
def get_ast(device:str, num_inputs:int) -> UOp:
if (device, num_inputs) not in cached_asts:
with Context(DEBUG=0):
fst = [Tensor.randn(BUF_SIZE, dtype=dtypes.int).realize() for _ in range(num_inputs)]
s = fst[0]
for i in range(1, num_inputs): s = s.bitwise_xor(fst[i])
cached_asts[(device, num_inputs)] = s.schedule_linear().src[-1].src[0]
return cached_asts[(device, num_inputs)]
def make_buffer(device, size=BUF_SIZE, fill=False):
buf = Buffer(device, size, dtypes.int).ensure_allocated()
if fill:
with Context(DEBUG=0):
buf.copy_from(Tensor(np.random.randint(-10000, 10000, size=size, dtype=np.int32)).realize().uop.base.realized)
return buf
def make_view(base, offset_elems, size_elems):
return Buffer(base.device, size_elems, base.dtype, base=base, offset=offset_elems * base.dtype.itemsize).ensure_allocated()
def get_buf_uop(buf:Buffer, cache:dict[Buffer,UOp]) -> UOp:
if buf not in cache:
cache[buf] = u = UOp.new_buffer(buf.device, buf.size, buf.dtype)
buffers[u] = buf
return cache[buf]
def copy_call(dst:Buffer, src:Buffer, c:dict[Buffer,UOp]) -> UOp:
return get_buf_uop(src,c).copy_to_device(dst.device).call(get_buf_uop(dst,c), get_buf_uop(src,c))
def make_graph(graph_cls, calls:list[UOp]):
linear = compile_linear(UOp(Ops.LINEAR, src=tuple(calls)))
cf = UOp(Ops.CUSTOM_FUNCTION, src=(linear,), arg="graph")
return graph_cls(cf, [])
def run_schedule(calls:list[UOp]):
run_linear(UOp(Ops.LINEAR, src=tuple(calls)))
def zero_bufs(bufs):
for b in bufs: b.copy_from(Buffer("PYTHON", b.size, b.dtype, opaque=memoryview(bytearray(b.nbytes))))
@unittest.skipUnless(Device[Device.DEFAULT].graph is not None, "graph support required")
class TestGraph(unittest.TestCase):
def skip_if_no_offset(self):
if Device.DEFAULT in {"WEBGPU", "CL"}: self.skipTest("device does not support _offset")
def skip_if_not_multigraph(self):
graph = g.func if isinstance(g:=(d:=Device[Device.DEFAULT]).graph, functools.partial) else g
if not issubclass(graph, MultiGraphRunner): self.skipTest("graph is not supported (not MultiGraphRunner)")
if not hasattr(d.allocator, '_transfer') or not d.allocator.supports_transfer: self.skipTest("device is not supported (no transfers)")
def test_order_2_writes_to_same_buf(self):
d0 = Device.DEFAULT
b = [make_buffer(d0, fill=True) for _ in range(5)]
c: dict[Buffer,UOp] = {}
calls = [
get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c)),
get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[3],c), get_buf_uop(b[4],c)),
]
zero_bufs([b[0]])
run_schedule(calls)
expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
for _ in range(RUN_CNT):
zero_bufs([b[0]])
make_graph(Device[d0].graph, calls)([], {})
for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
def test_order_read_write_same_buf(self):
d0 = Device.DEFAULT
b = [make_buffer(d0, fill=True) for _ in range(5)]
c: dict[Buffer,UOp] = {}
calls = [
get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c)),
get_ast(d0, 2).call(get_buf_uop(b[1],c), get_buf_uop(b[3],c), get_buf_uop(b[4],c)),
]
zero_bufs([b[0], b[1]])
run_schedule(calls)
expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
for _ in range(RUN_CNT):
zero_bufs([b[0], b[1]])
make_graph(Device[d0].graph, calls)([], {})
for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
def test_order_write_read_same_buf(self):
d0 = Device.DEFAULT
b = [make_buffer(d0, fill=True) for _ in range(5)]
c: dict[Buffer,UOp] = {}
calls = [
get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c)),
get_ast(d0, 2).call(get_buf_uop(b[1],c), get_buf_uop(b[0],c), get_buf_uop(b[4],c)),
]
zero_bufs([b[0], b[1]])
run_schedule(calls)
expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
for _ in range(RUN_CNT):
zero_bufs([b[0], b[1]])
make_graph(Device[d0].graph, calls)([], {})
for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
def test_order_copy_writed(self):
self.skip_if_not_multigraph()
d0 = Device.DEFAULT
b = [make_buffer(d0, fill=True) for _ in range(4)]
c: dict[Buffer,UOp] = {}
calls = [
get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c)),
copy_call(b[3], b[0], c),
]
zero_bufs([b[0], b[3]])
run_schedule(calls)
expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
for _ in range(RUN_CNT):
zero_bufs([b[0], b[3]])
make_graph(Device[d0].graph, calls)([], {})
for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
def test_order_copy_then_read(self):
self.skip_if_not_multigraph()
d0 = Device.DEFAULT
b = [make_buffer(d0, fill=True) for _ in range(4)]
c: dict[Buffer,UOp] = {}
calls = [
copy_call(b[1], b[0], c),
get_ast(d0, 2).call(get_buf_uop(b[3],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c)),
]
zero_bufs([b[1], b[3]])
run_schedule(calls)
expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
for _ in range(RUN_CNT):
zero_bufs([b[1], b[3]])
make_graph(Device[d0].graph, calls)([], {})
for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
def test_read_write_several_graphs(self):
d0 = Device.DEFAULT
b = [make_buffer(d0, fill=True) for _ in range(8)]
c: dict[Buffer,UOp] = {}
calls1 = [get_ast(d0, 2).call(get_buf_uop(b[3],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c))]
calls2 = [get_ast(d0, 2).call(get_buf_uop(b[4],c), get_buf_uop(b[1],c), get_buf_uop(b[3],c))]
calls3 = [get_ast(d0, 2).call(get_buf_uop(b[5],c), get_buf_uop(b[4],c), get_buf_uop(b[2],c))]
out = [b[3], b[4], b[5]]
zero_bufs(out)
run_schedule(calls1 + calls2 + calls3)
expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
for _ in range(RUN_CNT):
zero_bufs(out)
make_graph(Device[d0].graph, calls1)([], {})
make_graph(Device[d0].graph, calls2)([], {})
make_graph(Device[d0].graph, calls3)([], {})
for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
@needs_second_gpu
def test_copies_2_devs(self):
self.skip_if_not_multigraph()
d0, d1 = Device.DEFAULT, f"{Device.DEFAULT}:1"
b0 = [make_buffer(d0, fill=True) for _ in range(3)]
b1 = [make_buffer(d1, fill=True)]
c: dict[Buffer,UOp] = {}
calls = [
copy_call(b1[0], b0[0], c),
get_ast(d0, 2).call(get_buf_uop(b0[2],c), get_buf_uop(b0[0],c), get_buf_uop(b0[1],c)),
]
out = [b1[0], b0[2]]
zero_bufs(out)
run_schedule(calls)
expected = {buf: np.frombuffer(buf.as_memoryview(), np.int32).copy() for buf in b0 + b1}
for _ in range(RUN_CNT):
zero_bufs(out)
make_graph(Device[d0].graph, calls)([], {})
for buf in b0 + b1: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
def test_graph_offset_bufs(self):
self.skip_if_not_multigraph()
self.skip_if_no_offset()
d0 = Device.DEFAULT
b0 = make_buffer(d0, fill=True)
b1 = make_view(b0, 0, b0.size)
b2 = make_view(b0, 0, b0.size)
c: dict[Buffer,UOp] = {}
calls = [
copy_call(b0, b2, c),
get_ast(d0, 2).call(get_buf_uop(b1,c), get_buf_uop(b0,c), get_buf_uop(b2,c)),
]
zero_bufs([b0])
run_schedule(calls)
expected = np.frombuffer(b0.as_memoryview(), np.int32).copy()
for _ in range(RUN_CNT):
zero_bufs([b0])
make_graph(Device[d0].graph, calls)([], {})
np.testing.assert_equal(expected, np.frombuffer(b0.as_memoryview(), np.int32))
def test_partial_write_preserves_write_dep(self):
self.skip_if_not_multigraph()
self.skip_if_no_offset()
d0 = Device.DEFAULT
base = make_buffer(d0, BUF_SIZE * 2, fill=True)
copy_src_full = make_buffer(d0, BUF_SIZE * 2, fill=True)
copy_src_lo = make_buffer(d0, fill=True)
v_lo, v_hi = make_view(base, 0, BUF_SIZE), make_view(base, BUF_SIZE, BUF_SIZE)
a, out = make_buffer(d0, fill=True), make_buffer(d0, fill=True)
c: dict[Buffer,UOp] = {}
calls = [
copy_call(base, copy_src_full, c),
copy_call(v_lo, copy_src_lo, c),
get_ast(d0, 2).call(get_buf_uop(out,c), get_buf_uop(v_hi,c), get_buf_uop(a,c)),
]
zero_bufs([base, out])
run_schedule(calls)
expected = {base: np.frombuffer(base.as_memoryview(), np.int32).copy(), out: np.frombuffer(out.as_memoryview(), np.int32).copy()}
for _ in range(RUN_CNT):
zero_bufs([base, out])
make_graph(Device[d0].graph, calls)([], {})
for buf in [base, out]: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
def test_partial_write_preserves_read_dep(self):
self.skip_if_not_multigraph()
self.skip_if_no_offset()
d0 = Device.DEFAULT
base = make_buffer(d0, BUF_SIZE * 2, fill=True)
copy_dst = make_buffer(d0, BUF_SIZE * 2, fill=True)
copy_src_lo = make_buffer(d0, fill=True)
v_lo, v_hi = make_view(base, 0, BUF_SIZE), make_view(base, BUF_SIZE, BUF_SIZE)
a, b = make_buffer(d0, fill=True), make_buffer(d0, fill=True)
c: dict[Buffer,UOp] = {}
calls = [
copy_call(copy_dst, base, c),
copy_call(v_lo, copy_src_lo, c),
get_ast(d0, 2).call(get_buf_uop(v_hi,c), get_buf_uop(a,c), get_buf_uop(b,c)),
]
zero_bufs([copy_dst, base])
run_schedule(calls)
expected = {copy_dst: np.frombuffer(copy_dst.as_memoryview(), np.int32).copy(), base: np.frombuffer(base.as_memoryview(), np.int32).copy()}
for _ in range(RUN_CNT):
zero_bufs([copy_dst, base])
make_graph(Device[d0].graph, calls)([], {})
for buf in [copy_dst, base]: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
def test_middle_write_splits_write_dep(self):
self.skip_if_not_multigraph()
self.skip_if_no_offset()
d0 = Device.DEFAULT
base = make_buffer(d0, BUF_SIZE * 3, fill=True)
copy_src_full = make_buffer(d0, BUF_SIZE * 3, fill=True)
copy_src_mid = make_buffer(d0, fill=True)
v_lo, v_mid, v_hi = make_view(base, 0, BUF_SIZE), make_view(base, BUF_SIZE, BUF_SIZE), make_view(base, BUF_SIZE * 2, BUF_SIZE)
a, out1, out2 = make_buffer(d0, fill=True), make_buffer(d0, fill=True), make_buffer(d0, fill=True)
c: dict[Buffer,UOp] = {}
calls = [
copy_call(base, copy_src_full, c),
copy_call(v_mid, copy_src_mid, c),
get_ast(d0, 2).call(get_buf_uop(out1,c), get_buf_uop(v_lo,c), get_buf_uop(a,c)),
get_ast(d0, 2).call(get_buf_uop(out2,c), get_buf_uop(v_hi,c), get_buf_uop(a,c)),
]
outs = [base, out1, out2]
zero_bufs(outs)
run_schedule(calls)
expected = {buf: np.frombuffer(buf.as_memoryview(), np.int32).copy() for buf in outs}
for _ in range(RUN_CNT):
zero_bufs(outs)
make_graph(Device[d0].graph, calls)([], {})
for buf in outs: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
if __name__ == '__main__':
unittest.main()

52
artifacts/package_sources/tinygrad/test/backend/test_interop.py Normal file
View File

@@ -0,0 +1,52 @@
#!/usr/bin/env python
import unittest, os
import torch
import numpy as np
from tinygrad.helpers import DEV
from tinygrad.tensor import Tensor
from tinygrad.device import Device
from tinygrad.dtype import _from_torch_dtype, _to_torch_dtype
MOCKGPU = DEV.interface.startswith("MOCK")
@unittest.skipIf(Device.DEFAULT not in ["METAL", "CUDA"] or MOCKGPU, f"no support on {Device.DEFAULT}")
class TestInterop(unittest.TestCase):
def setUp(self):
if Device.DEFAULT == "CUDA": self.torch_device = "cuda"
elif Device.DEFAULT == "METAL": self.torch_device = "mps"
def test_torch_interop(self):
inp = torch.rand(2, 2, 3, device=torch.device(self.torch_device))
if self.torch_device == "mps": torch.mps.synchronize()
else: torch.cuda.synchronize()
tg_data = Tensor.from_blob(inp.data_ptr(), inp.shape, dtype=_from_torch_dtype(inp.dtype))
tg_out = tg_data[:, :, 0] * 0.2989 + tg_data[:, :, 1] * 0.5870 + tg_data[:, :, 2] * 0.1140
tg_res = tg_out.numpy()
if self.torch_device == "mps" and os.getenv("CI", "") != "":
# MPS backend out of memory: https://discuss.pytorch.org/t/mps-back-end-out-of-memory-on-github-action/189773
# Calculate expected value on cpu.
inp = inp.cpu()
torch_out = inp[:, :, 0] * 0.2989 + inp[:, :, 1] * 0.5870 + inp[:, :, 2] * 0.1140
np.testing.assert_allclose(tg_res, torch_out.cpu().numpy(), atol=1e-5, rtol=1e-5)
def test_torch_interop_write(self):
tg_data = Tensor.randn((4, 4), device=Device.DEFAULT)
out = torch.empty(4, 4, device=torch.device(self.torch_device), dtype=_to_torch_dtype(tg_data.dtype))
tg_out = Tensor.from_blob(out.data_ptr(), out.shape, dtype=_from_torch_dtype(out.dtype))
tg_out.assign(tg_data).realize()
Device[Device.DEFAULT].synchronize()
torch_out_np = out.cpu().numpy()
np.testing.assert_allclose(tg_data.numpy(), torch_out_np, atol=1e-5, rtol=1e-5)
if __name__ == '__main__':
unittest.main()

55
artifacts/package_sources/tinygrad/test/backend/test_isel.py Normal file
View File

@@ -0,0 +1,55 @@
import unittest
from typing import cast
from tinygrad import Device
from tinygrad.uop import Ops
from tinygrad.uop.ops import UOp, dtypes, graph_rewrite
from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
from tinygrad.renderer.isa import IselContext
# INDEX on a register value with a constant index extracts a single element (the old GEP)
def lane(y:UOp, i:int) -> UOp: return y.index(UOp.const(i, dtypes.int), dtype=y.dtype.scalar())
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86")
class TestIselX86(unittest.TestCase):
def isel_rewrite(self, x:UOp):
return graph_rewrite(x, cast(X86Renderer, Device[Device.DEFAULT].renderer).isel_matcher, IselContext(x), bottom_up=True)
def _check_op(self, dt_op, expr):
nargs = expr.__code__.co_argcount
for dt,op in dt_op:
with self.subTest(dtype=dt):
v = [UOp.variable(str(i), 0, 0, dt) for i in range(nargs)]
n = self.isel_rewrite(expr(*v))
self.assertIs(n.arg, op)
def test_cmove(self):
a = UOp.variable("a", 0, 0, dtypes.int32)
b = UOp.variable("b", 0, 0, dtypes.int32)
c = (a < b).where(a, b)
d = (a != b).where(a, b)
f = c + d
n = self.isel_rewrite(f)
self.assertTrue(n.src[0].arg is X86Ops.CMOVL and n.src[1].arg is X86Ops.CMOVNE)
# both comparisons become the same instruction
self.assertTrue(n.src[0].src[2] == n.src[1].src[2] and n.src[0].src[2].arg is X86Ops.CMP)
def test_vinsertps(self):
a = UOp.variable("a", 0, 0, dtypes.float32)
b = UOp.variable("b", 0, 0, dtypes.float32)
c = UOp.variable("c", 0, 0, dtypes.float32)
d = UOp.variable("e", 0, 0, dtypes.float32)
valid = [UOp.stack(lane(a, 0), lane(b, 1), lane(a, 2), lane(b, 3)),
UOp.stack(lane(a, 3), lane(b, 2), lane(c, 1), d)]
for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VINSERTPS)
# complex address is [base + index*scale + displacement]
def test_complex_address(self):
a = UOp.variable("a", 0, 0, dtypes.int32)
load = UOp.param(0, dtypes.int32, (16,)).index(a + 1).load()
n = self.isel_rewrite(load)
# displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32
self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].val == 4)
if __name__ == "__main__":
unittest.main()

450
artifacts/package_sources/tinygrad/test/backend/test_jit.py Normal file
View File

@@ -0,0 +1,450 @@
#!/usr/bin/env python
import unittest
import numpy as np
from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu, KernelCountException
from test.unit.test_jit import _simple_test
from tinygrad import Tensor, Variable, TinyJit, Device, dtypes
from tinygrad.engine.jit import graph_class
from tinygrad.helpers import JIT, DEV, GlobalCounters
from tinygrad.uop.ops import Ops
from tinygrad.renderer.isa.x86 import X86Renderer
class TestJit(unittest.TestCase):
def test_simple_jit(self):
@TinyJit
def add(a, b): return (a+b).realize()
_simple_test(add)
@unittest.skipUnless(Device.DEFAULT == "CPU", "core_id is a CPU runtimevar")
def test_hcq_core_id_runtimevar_merge(self):
N = 262144
@TinyJit
def f(x, st):
y = (x + 1).contiguous().realize()
z = x.shrink(((st, st + N),)).contiguous().realize()
return y, z
x = Tensor.arange(2*N).clone().realize()
for _ in range(3): y, z = f(x, Variable("a", 0, N).bind(0))
self.assertEqual(y.shape, (2*N,))
self.assertEqual(z.shape, (N,))
def test_jit_input_view(self):
@TinyJit
def f(x): return (x[2:5].contiguous() + 1).realize()
for i in range(5):
x = (Tensor.arange(10).float() + i * 10).clone().realize()
np.testing.assert_allclose(f(x).numpy(), x.numpy()[2:5] + 1)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "estimates are wrong for x86")
def test_global_counters_jit(self):
@TinyJit
def f(a, b):
c = (a + b).realize()
d = (c * 2).realize()
return (d - a).realize()
a, b = Tensor.randn(64, 64).realize(), Tensor.randn(64, 64).realize()
for _ in range(4):
GlobalCounters.reset()
f(a, b)
Device[a.device].synchronize()
self.assertGreater(GlobalCounters.global_mem, 0)
self.assertGreater(GlobalCounters.global_ops, 0)
def test_jit_assign(self, dtype=dtypes.float32):
@TinyJit
def add(a):
a += 1
a.realize()
a = Tensor.zeros(1, dtype=dtype).contiguous().realize()
for _ in range(5): add(a)
self.assertEqual(a.item(), 5)
def test_jit_assign_int8(self): self.test_jit_assign(dtypes.int8)
def test_jit_copyin(self):
@TinyJit
def f(a):
return a + Tensor([1,2,3])
for _ in range(5):
b = Tensor.randn(3)
c = f(b)
np.testing.assert_allclose(c.numpy(), b.numpy()+[1,2,3], atol=1e-4, rtol=1e-5)
def test_jit_batch_split(self):
if Device[Device.DEFAULT].graph is None or JIT >= 2: raise unittest.SkipTest("only test graphs")
# Create long jit with 83 kernels.
def f(a, b, c, d, e):
for _ in range(80):
a = (a+b).realize()
y = (a*c).realize()
z = (y*d).realize()
w = (z*e)
return w.realize()
a = Tensor.randn(10, 10).realize()
b = Tensor.randn(10, 10).realize()
c = Tensor.randn(10, 10).realize()
d = Tensor.randn(10, 10).realize()
e = Tensor.randn(10, 10).realize()
jf = TinyJit(f)
prev = None
for _ in range(5):
o = jf(a, b, c, d, e).numpy()
if prev is not None: np.testing.assert_allclose(o, prev, atol=1e-4, rtol=1e-5)
prev = o
# Checking that 2 graphs are inited.
if len(jf.captured.linear.src) != 2: raise KernelCountException(2, len(jf.captured.linear.src))
for si in jf.captured.linear.src:
assert call_is_graph(si)
def test_jitted_clone(self):
def f(a): return a.clone().realize()
jf = TinyJit(f)
for _ in range(5):
a = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
ja = jf(a)
np.testing.assert_allclose(a.numpy(), ja.numpy(), atol=1e-4, rtol=1e-5)
@needs_second_gpu
@unittest.skipIf(not_support_multi_device(), "no multi")
def test_jitted_transfers(self):
d0, d1 = f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"
def f(a, b):
x = a.to(d1)
y = b.to(d1)
return x.realize(), y.realize()
jf = TinyJit(f)
for _ in range(5):
a = Tensor.randn(10, 10, device=d0).realize()
b = Tensor.randn(10, 10, device=d0).realize()
xc, yc = jf(a, b)
np.testing.assert_allclose(a.numpy(), xc.numpy(), atol=1e-4, rtol=1e-5)
np.testing.assert_allclose(b.numpy(), yc.numpy(), atol=1e-4, rtol=1e-5)
def test_jit_several_devs(self):
d0, d1 = f"{Device.DEFAULT}:0", "CPU"
def f(a, b):
x = a.to(d0).realize()
y = b.to(d0).realize()
return x+y.realize(), x*y.realize()
jf = TinyJit(f)
for _ in range(5):
a = Tensor.randn(10, 10, device=d1).realize()
b = Tensor.randn(10, 10, device=d1).realize()
zc, wc = jf(a, b)
np.testing.assert_allclose((a.numpy()+b.numpy()), zc.numpy(), atol=1e-4, rtol=1e-5)
np.testing.assert_allclose((a.numpy()*b.numpy()), wc.numpy(), atol=1e-4, rtol=1e-5)
@needs_second_gpu
@unittest.skipIf(not_support_multi_device(), "no multi")
def test_jitted_view(self):
d0, d1 = f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"
def f(a):
x1 = a.sum(axis=(1,))
x = (x1 + 5).bitcast(dtypes.int32)
y = x.to(d1)
return y.realize()
jf = TinyJit(f)
for _ in range(5):
a = Tensor.randn(10, 1000, device=d0).realize()
xc = jf(a)
np.testing.assert_allclose((a.numpy().sum(axis=(1,)) + 5).view(np.int32), xc.numpy(), atol=1e-4, rtol=5e-5)
@unittest.skip("Pending multioutput implementation #3607")
class TestMultioutputJit(unittest.TestCase):
def _test(self, f):
for _ in range(5):
a, b = Tensor.randn(10, 10), Tensor.randn(10, 10)
out0, out1, out2 = f(a, b)
np.testing.assert_allclose(out0.numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
np.testing.assert_allclose(out1.numpy(), a.numpy()-b.numpy(), atol=1e-4, rtol=1e-5)
np.testing.assert_allclose(out2.numpy(), a.numpy()*b.numpy(), atol=1e-4, rtol=1e-5)
def test_jit_multioutput_realize(self):
@TinyJit
def fxn(a, b): return (a+b).realize(), (a-b).realize(), (a*b).realize()
self._test(fxn)
assert_jit_cache_len(fxn, 3)
def test_jit_multioutput_norealize(self):
@TinyJit
def fxn(a, b): return a+b, a-b, a*b
self._test(fxn)
assert_jit_cache_len(fxn, 1)
def test_jit_multioutput_mix(self):
@TinyJit
def fxn(a, b): return a+b, a-b, (a*b).realize()
self._test(fxn)
assert_jit_cache_len(fxn, 2)
class TestCopyInsideJit(unittest.TestCase):
def test_copy_inside_jit(self):
@TinyJit
def add(x,y) -> Tensor: return x.to(Device.DEFAULT)+y
for _ in range(5):
# create a Tensor on CPU
a = Tensor.rand(16,16,device="CPU").realize()
b = Tensor.rand(16,16).realize()
out = add(a,b)
np.testing.assert_allclose(out.flatten().tolist(), [x+y for x,y in zip(a.flatten().tolist(), b.flatten().tolist())])
class TestJitPrune(unittest.TestCase):
def test_prune_w_copy_correct(self):
weights = Tensor.rand(16).realize()
def w2(x) -> Tensor: return (weights*2).contiguous() + x.to(Device.DEFAULT)
w2_noprune = TinyJit(w2)
w2_prune = TinyJit(w2, prune=True)
for _ in range(3):
a = Tensor.rand(16, device="CPU").realize()
out = w2_noprune(a)
np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
for _ in range(3):
a = Tensor.rand(16, device="CPU").realize()
out = w2_prune(a)
np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
def test_prune_w_independent_copy_correct(self):
weights = Tensor.rand(16, device="CPU").realize()
def w2(x) -> Tensor: return (weights*2).contiguous().to(Device.DEFAULT) + x
w2_noprune = TinyJit(w2)
w2_prune = TinyJit(w2, prune=True)
for _ in range(3):
a = Tensor.rand(16).realize()
out = w2_noprune(a)
np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
for _ in range(3):
a = Tensor.rand(16).realize()
out = w2_prune(a)
np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
assert_jit_cache_len(w2_prune, 1)
class TestJitFree(unittest.TestCase):
def test_free_intermediates(self):
ext_tensor = Tensor([1,24,23,45,1])
@TinyJit
def fxn(x:Tensor):
t1 = (x * 2).contiguous().realize()
t2 = (t1 + ext_tensor).contiguous().realize()
out = (t2.sum()).contiguous().realize()
return out
for i in range(5):
out = fxn(inp:=Tensor([i,1,2,3,4]))
self.assertEqual(out.item(), 114+2*i)
pre_free = GlobalCounters.mem_used
fxn.captured.free_intermediates()
savings_after_free = pre_free - GlobalCounters.mem_used
expected_savings = (len(inp) * inp.dtype.itemsize * 2) + dtypes.float32.itemsize # (t1 and t2) + out
self.assertGreaterEqual(savings_after_free, expected_savings)
out = fxn(Tensor([11,1,2,3,4]))
self.assertEqual(out.item(), 136)
# Try one more time...
pre_free = GlobalCounters.mem_used
fxn.captured.free_intermediates()
fxn.captured.free_intermediates() # 2nd time to validate
savings_after_free = pre_free - GlobalCounters.mem_used
self.assertGreaterEqual(savings_after_free, expected_savings)
out = fxn(Tensor([11,1,2,3,4]))
self.assertEqual(out.item(), 136)
def test_updated_not_freed(self):
x = Tensor([1]).realize()
@TinyJit
def fxn(y):
nonlocal x
x += y
return x
for _ in range(5): fxn(Tensor([1]))
self.assertEqual(x.item(), 6)
pre_free = GlobalCounters.mem_used
fxn.captured.free_intermediates()
savings_after_free = pre_free - GlobalCounters.mem_used
self.assertEqual(savings_after_free, 0)
fxn(Tensor([2]))
self.assertEqual(x.item(), 8)
class TestJitGraphSplit(unittest.TestCase):
def compute(self, device, inp):
assert inp.device == device, f"Input device {inp.device} does not match expected {device}"
return (inp + 1.0).contiguous().realize()
def copy(self, device, to_device, inp):
assert inp.device == device, f"Input device {inp.device} does not match expected {device}"
return inp.to(to_device).realize()
def expect(self, f, *args, graph=None, multigraph=None, hcqgraph=None):
def _numpies(tpl): return tpl.numpy() if tpl.__class__ is Tensor else tuple([t.numpy() for t in tpl])
expected = _numpies(f(*args))
for i in range(4):
res = _numpies(f(*args))
np.testing.assert_allclose(res, expected, atol=1e-4, rtol=1e-5)
dev = Device[Device.DEFAULT]
graph_t = graph_class(dev)
if graph_t is None: return
got = f.captured.linear.src
from tinygrad.runtime.graph.hcq import HCQGraph
from tinygrad.engine.jit import MultiGraphRunner
if graph_t is HCQGraph:
validate = hcqgraph
elif issubclass(graph_t, MultiGraphRunner):
validate = multigraph
else:
validate = graph
assert len(got) == len(validate), f"Expected {len(validate)} operations, got {len(got)}"
for expected, si in zip(validate, got):
ast = si.src[0]
if expected["type"] == "graph":
assert call_is_graph(si), f"Expected graph, got {ast.op}"
inner_cnt = len(ast.src[0].src)
assert inner_cnt == expected["cnt"], f"Expected {expected['cnt']} operations in graph, got {inner_cnt}"
elif expected["type"] == "comp":
assert ast.op in (Ops.SINK, Ops.PROGRAM), f"Expected kernel, got {ast.op}"
elif expected["type"] in ("copy", "xfer"):
assert ast.op is Ops.COPY, f"Expected COPY, got {ast.op}"
def ji_graph(self, cnt): return {"type": "graph", "cnt": cnt}
def ji_comp(self): return {"type": "comp"}
def ji_copy(self): return {"type": "copy"}
def ji_xfer(self): return {"type": "xfer"}
def test_jit_split_simple(self):
@TinyJit
def f(inp):
op0 = self.compute(Device.DEFAULT, inp)
op1 = self.compute(Device.DEFAULT, op0)
op2 = self.compute(Device.DEFAULT, op1)
return op2
inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
self.expect(f, inp,
graph=[self.ji_graph(3)],
multigraph=[self.ji_graph(3)],
hcqgraph=[self.ji_graph(3)])
def test_jit_cpu_simple(self):
if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
@TinyJit
def f(inp, inp_cpu):
op0 = self.compute(Device.DEFAULT, inp)
op1 = self.compute(Device.DEFAULT, op0)
op2 = self.compute("CPU", inp_cpu)
op3 = self.compute(Device.DEFAULT, op1)
return op2, op3
inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
inp_cpu = Tensor.randn(10, 10, device="CPU").realize()
self.expect(f, inp, inp_cpu,
graph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
multigraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
hcqgraph=[self.ji_graph(4)])
def test_jit_cpu_several(self):
if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
@TinyJit
def f(inp, inp_cpu):
op0 = self.compute(Device.DEFAULT, inp)
op1 = self.compute(Device.DEFAULT, op0)
op2 = self.compute("CPU", inp_cpu)
op3 = self.compute("CPU", op2)
op4 = self.compute(Device.DEFAULT, op1)
return op3, op4
inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
inp_cpu = Tensor.randn(10, 10, device="CPU").realize()
self.expect(f, inp, inp_cpu,
graph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
multigraph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
hcqgraph=[self.ji_graph(5)])
def test_jit_multidev(self):
if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
try: Device[f"{Device.DEFAULT}:1"]
except Exception: raise unittest.SkipTest("no multidevice")
@TinyJit
def f(inp, inp_d1):
op0 = self.compute(Device.DEFAULT, inp)
op1 = self.compute(Device.DEFAULT, op0)
op2 = self.compute(f"{Device.DEFAULT}:1", inp_d1)
op3 = self.compute(f"{Device.DEFAULT}:1", op2)
op4 = self.compute(Device.DEFAULT, op1)
return op3, op4
inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
inp_d1 = Tensor.randn(10, 10, device=f"{Device.DEFAULT}:1").realize()
self.expect(f, inp, inp_d1,
graph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
multigraph=[self.ji_graph(5)],
hcqgraph=[self.ji_graph(5)])
def test_jit_multidev_xfer(self):
if Device.DEFAULT in {"CPU"}: raise unittest.SkipTest("CPU is not a valid default device for this test (zero-copies)")
if Device.DEFAULT == "METAL": raise unittest.SkipTest("Metal is flaky, with multidevice (same as metal llama 4gpu?)")
try: Device[f"{Device.DEFAULT}:1"]
except Exception: raise unittest.SkipTest("no multidevice")
@TinyJit
def f(inp, inp_d1):
op0 = self.compute(Device.DEFAULT, inp)
op1 = self.compute(Device.DEFAULT, op0)
op2 = self.compute(f"{Device.DEFAULT}:1", inp_d1)
op3 = self.copy(f"{Device.DEFAULT}:1", Device.DEFAULT, op2)
op4 = self.compute(f"{Device.DEFAULT}:1", op2)
op5 = self.compute(Device.DEFAULT, op3)
return op1, op4, op5
inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
inp_d1 = Tensor.randn(10, 10, device=f"{Device.DEFAULT}:1").realize()
self.expect(f, inp, inp_d1,
graph=[self.ji_graph(2), self.ji_comp(), self.ji_xfer(), self.ji_comp(), self.ji_comp()],
multigraph=[self.ji_graph(6)],
hcqgraph=[self.ji_graph(6)])
@unittest.skip("this fails if you don't have SDMA or are using AMD_DISABLE_SDMA=1")
@unittest.skipIf(DEV.interface.startswith("MOCK"), "MockGPU does not support parallel copies")
def test_jit_multidev_copy(self):
if Device.DEFAULT in {"CPU"}: raise unittest.SkipTest("CPU/LLVM is not a valid default device for this test (zero-copies)")
@TinyJit
def f(inp):
op0 = self.compute(Device.DEFAULT, inp)
op1 = self.compute(Device.DEFAULT, op0)
op2 = self.copy(Device.DEFAULT, "CPU", op1)
op3 = self.compute("CPU", op2)
return op3
inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
self.expect(f, inp,
graph=[self.ji_graph(2), self.ji_copy(), self.ji_comp()],
multigraph=[self.ji_graph(2), self.ji_copy(), self.ji_comp()],
hcqgraph=[self.ji_graph(4)])
if __name__ == '__main__':
unittest.main()

469
artifacts/package_sources/tinygrad/test/backend/test_linearizer.py Normal file
View File

@@ -0,0 +1,469 @@
import numpy as np
import unittest
from tinygrad.codegen.opt import Opt, OptOps
from tinygrad.uop.ops import UOp, Ops, GroupOp, AxisType, buffers
from tinygrad.device import Device, Buffer
from tinygrad.tensor import Tensor, _to_np_dtype
from tinygrad.engine.realize import run_linear
from tinygrad.codegen import to_program
from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT, DEV
from tinygrad.dtype import DType, dtypes, AddrSpace
from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.cstyle import CUDARenderer
from tinygrad.renderer.isa import ISARenderer
from test.helpers import replace_opts, check_schedule
from test.backend.test_softmax_fusion import single_kernel_softmax
MOCKGPU = DEV.interface.startswith("MOCK")
from tinygrad.uop.render import print_uops # noqa: F401 # pylint: disable=unused-import
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, ISARenderer), "isa backends don't preserve the op spec when lowering")
class TestLinearizer(unittest.TestCase):
def test_arg_dedup(self):
# NOTE: this realize exists because Tensor.numpy calls .contiguous() internally
# without contiguous folding, rand.to("CPU") and rand.contiguous().to("CPU") are different UOps.
# this test asserts they are the identical Buffer
# having different buffers is fine for correctness, because the outputs match.
a, b = Tensor.randn(4).realize(), Tensor.randn(4).realize()
np_a, np_b = a.numpy(), b.numpy()
c = ((a.shrink(((0, 2),)) - a.shrink(((2, 4),))) - (b.shrink(((0, 2),)) - b.shrink(((2, 4),))))
linear = c.schedule_linear()
run_linear(linear)
rawbufs = [s.buffer for s in linear.src[-1].src[1:] if s.op is not Ops.BIND]
assert len(rawbufs) == 3 and set(rawbufs[1:]) == {a.uop.base.realized, b.uop.base.realized}
np_c = (np_a[:2] - np_a[2:]) - (np_b[:2] - np_b[2:])
np.testing.assert_allclose(np_c, c.numpy(), atol=1e-4, rtol=1e-4)
def test_load_removed(self):
a = Tensor.rand(1).realize()
b = Tensor.rand(1).realize()
ta = Tensor.where(Tensor(True), a, b).numpy()
tb = Tensor.where(Tensor(False), a, b).numpy()
np.testing.assert_equal(a.numpy(), ta)
np.testing.assert_equal(b.numpy(), tb)
@unittest.skip("TODO: some backends insert more casts")
def test_cast_there_and_back(self):
tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
out = tst.neg().cast(dtypes.char).cast(dtypes.int).cast(dtypes.char) * 2
ast = helper_linearizer_opt(out)
uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)
@unittest.expectedFailure
def test_cast_back_and_there(self):
tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
out = tst.neg().cast(dtypes.char).cast(dtypes.int) * 2
ast = helper_linearizer_opt(out)
uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx")
def test_late_bias_load(self):
img = Tensor.empty(1, 3, 16, 16)
w = Tensor.empty(16, 3, 3, 3)
b = Tensor.empty(16)
out = img.conv2d(w, b)
ast = helper_linearizer_opt(out)
uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
# slice at the last loop end
uslice = [i for i,u in enumerate(uops) if u.op == Ops.END][-1]
# only valid test if outermost range is the reduce
if uops[uslice].src[-1].arg[-1] == AxisType.REDUCE:
load_idxs = [u.src[0] for u in uops[uslice+1:] if u.op == Ops.LOAD]
# assert that there is a global load after the reduce ends
assert any(u.addrspace == AddrSpace.GLOBAL for u in load_idxs)
def _test_no_nested_ranges(self, lins, skip=None):
for l in lins:
range_in_acc = flatten([[x for x in u.src if x.op is Ops.RANGE] for u in l.uops if u.op is Ops.BUFFER and u.addrspace is AddrSpace.REG])
ranges = [u.op for u in l.uops if (u.op is Ops.RANGE and u in range_in_acc) or (u.op is Ops.END and u.src[0] in range_in_acc)]
for i,u in enumerate(ranges):
if skip and i in skip: continue
assert ranges[i-1] != u, f"multireduce nested the ranges! {ranges[i-1], {u}}"
def test_two_nested_range(self):
a = Tensor.randn(2, ).realize()
out = a.reshape(2, 1).expand(2, 3).sum()
ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)).sum()])
uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
assert len(ranges) == 1 # NOTE: it collapses now
def test_three_nested_range(self):
a = Tensor.randn(2, ).realize()
out = a.reshape(2, 1).expand(2, 3).expand(2, 2, 3).sum()
ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)), (2, 2, 3)).sum()])
uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
assert len(ranges) == 1 # NOTE: it collapses now
def test_two_nested_range_alt_indexing(self):
a = Tensor([2, 2]).realize()
out = a.reshape(2, 1).pad(((1, 1), (1, 1)), value=2).sum()
ast = helper_linearizer_opt(out, wanna_output=[24])
uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
# RANGE -> ALU -> RANGE -> ALU + LOAD -> STORE
assert any(x.op in GroupOp.ALU for x in uops[ranges[0]:ranges[1]])
# the index of the load doesnt depend on the second range
assert any(x.op is Ops.LOAD for x in uops[ranges[0]:ranges[1]])
assert any(x.op in {*GroupOp.ALU, Ops.LOAD} for x in uops[ranges[1]:])
def test_range_outer_op_before_phi(self):
a = Tensor.randn(4, 1).realize()
b = Tensor.randn(1, 1).realize()
out = (a + b[0]).sum() + b[0]
ast = helper_linearizer_opt(out, wanna_output=[(a.numpy()+b.numpy()[0]).sum()+b.numpy()])
uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
# LOAD -> RANGE -> LOAD -> STORE
assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD]) == 1
def test_range_outer_op_before_phi_nested_range(self):
a = Tensor.randn(2, ).realize()
b = Tensor.randn(1, 1).realize()
out = (a.reshape(2, 1).expand(2, 3) + b[0]).sum() + b[0]
ast = helper_linearizer_opt(out, wanna_output=[(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)) + b.numpy()[0]).sum() + b.numpy()])
uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
assert len(ranges) == 1 # NOTE: it collapses now
def test_load_dedup(self):
# for different leaves in the AST, the same loads may occur.
a = Tensor.randn(4).realize()
# these are of size 3 to avoid float4 coalesce
r = a[:-1] + a[1:]
uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
renderer=Device[Device.DEFAULT].renderer).src[1].src)
num_loads = len([uop for uop in uops if uop.op is Ops.LOAD])
assert num_loads <= 4, "more load uops than needed"
assert num_loads >= 1, "expected at least one load uop"
@unittest.skip("this is handled at higher level now")
def test_upcast_cse(self):
# when upcasting, within a subtree, there may be common expressions.
a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
r = a.expand([2]) + b.expand([2])
uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
renderer=Device[Device.DEFAULT].renderer).src[1].src)
num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
assert num_ops <= 1, "more alu uops than needed"
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
def test_reduce_upcast(self):
x, w = Tensor.randn((1,1,3)).realize(), Tensor.randn((1,1,2)).realize()
r = Tensor.conv2d(x,w,padding=1).relu()
uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0],
[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).src[1].src)
accs = [u for u in uops if u.op is Ops.BUFFER and u.addrspace is AddrSpace.REG]
stores = [u for u in uops if u.op is Ops.STORE]
assert len(accs) == 0 # it's removed now
assert len(stores) == 1
# NOTE: can reenable, it does work. it just makes BEAM slow
@unittest.expectedFailure
@unittest.skipUnless(Device.DEFAULT == "CPU", "test only for CPU")
def test_upcast_with_locals_cpu(self):
out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous()
prg = to_program(replace_opts(out.schedule_linear().src[-1].src[0], [Opt(OptOps.LOCAL, axis=0, arg=4)]),
renderer=Device[Device.DEFAULT].renderer)
self.assertEqual(len(prg.src[2].arg.split("for")), 5)
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx for some reason")
def test_upcast_with_locals(self):
x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
r = (x@y).relu()
opts_to_apply = [Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
program = to_program(replace_opts(r.schedule_linear().src[-1].src[0], opts_to_apply), renderer=Device[Device.DEFAULT].renderer)
stores = [u for u in tuple(program.src[1].src) if u.op is Ops.STORE and u.src[0].addrspace != AddrSpace.REG]
# the first store is to lds and can be upcasted
assert stores[0].src[1].max_numel() == 4
assert any(x.addrspace is AddrSpace.LOCAL for x in stores[0].toposort())
# the second store is to gds with no upcasts
assert stores[1].src[1].max_numel() == 1
assert stores[1].src[1].dtype == dtypes.float
assert any(x.op is Ops.PARAM for x in stores[1].toposort())
def test_zero_fold(self):
a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
r = Tensor.stack(a, b)
uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
renderer=Device[Device.DEFAULT].renderer).src[1].src)
num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
assert num_ops == 0, "more alu uops than needed"
def test_sum_acc_dtype(self):
for tensor_dtype, acc_dtype in (
(dtypes.bool, dtypes.int), (dtypes.int16, dtypes.int), (dtypes.float16, dtypes.float), (dtypes.bfloat16, dtypes.float)):
if tensor_dtype in (dts:=Device[Device.DEFAULT].renderer.supported_dtypes()) and acc_dtype in dts:
a = Tensor([1, 2, 3], dtype=tensor_dtype).sum()
realized_ast = a.schedule_linear().src[-1].src[0]
program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
local = [uop for uop in tuple(program.src[1].src) if uop.op is Ops.BUFFER and uop.addrspace in (AddrSpace.LOCAL, AddrSpace.REG)]
assert local[0].dtype == acc_dtype
def test_arg_acc_dtype(self):
def helper_arg_acc_dtype(c: Tensor, expected_dtype:DType):
realized_ast = c.schedule_linear().src[-1].src[0]
program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
local = [uop for uop in tuple(program.src[1].src) if uop.op is Ops.BUFFER and uop.addrspace in (AddrSpace.LOCAL, AddrSpace.REG)]
self.assertEqual(local[0].dtype, expected_dtype)
tests = (
(dtypes.float16, None, dtypes.float),
(dtypes.bfloat16, None, dtypes.float),
(dtypes.float, None, dtypes.float),
(dtypes.float16, dtypes.float16, dtypes.float16),
(dtypes.bfloat16, dtypes.bfloat16, dtypes.bfloat16),
(dtypes.float, dtypes.float16, dtypes.float16),
)
for tensor_dtype, acc_dtype, expected_dtype in tests:
if tensor_dtype in (dts:=Device[Device.DEFAULT].renderer.supported_dtypes()) and acc_dtype in dts and expected_dtype in dts:
a, b = Tensor.rand(8, 8, dtype=tensor_dtype), Tensor.rand(8, 8, dtype=tensor_dtype)
helper_arg_acc_dtype(a.sum(dtype=acc_dtype), expected_dtype)
helper_arg_acc_dtype(a.matmul(b, dtype=acc_dtype), expected_dtype)
helper_arg_acc_dtype(Tensor.einsum("ki,ij->kj", a, b, dtype=acc_dtype), expected_dtype)
d, w = Tensor.rand(4, 8, 8, 8, dtype=tensor_dtype), Tensor.rand(8, 8, 2, 2, dtype=tensor_dtype)
helper_arg_acc_dtype(d.conv2d(w, dtype=acc_dtype), expected_dtype)
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
def test_simple_unroll_no_between_phi_dependencies(self):
x, y = Tensor.empty(64, 64), Tensor.empty(64, 64)
r = (x@y).relu()
opt = [Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4)]
ast = helper_linearizer_opt(r, [opt])
# the uops graph is reg BUFFER -> 4x STORE 0.0 -> RANGE -> 4x ALU -> 4x STORE -> ENDRANGE
uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[1].src)
begin_range = [i for i, x in enumerate(uops) if x.op is Ops.RANGE][-1]
end_range = [i for i, x in enumerate(uops) if x.op is Ops.END][0]
for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype)
for u in uops:
if u.op is Ops.STORE and u.src[0].addrspace is AddrSpace.REG:
if uops.index(u) < begin_range:
assert u.src[1].op is Ops.CONST
else:
assert u.src[1].op in GroupOp.ALU
assert begin_range < uops.index(u) < end_range
# children of END are placed after ENDRANGE
if any(x.op is Ops.END and x.src[1].op in GroupOp.ALU for x in u.src):
assert end_range < uops.index(u)
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
def test_default_global_reversed(self):
# shrink so that the dims do not collapse
t = Tensor.ones(5, 6, 7).contiguous().realize().shrink(((0, 4), (0, 5), (0, 6)))
ast = helper_linearizer_opt(t+1)
uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
idxs = sorted(idxs, key=lambda uop: uop.arg)
assert (idxs[0].arg, idxs[0].src[0].val) == ('gidx0', 6), idxs[0]
assert (idxs[1].arg, idxs[1].src[0].val) == ('gidx1', 5), idxs[1].arg
assert (idxs[2].arg, idxs[2].src[0].val) == ('gidx2', 4), idxs[2].arg
def test_sum_collapse(self):
t = Tensor([2]).reshape(1, 1).expand(256, 256).sum()
sched = [si for si in t.schedule_linear().src if si.src[0].op is Ops.SINK]
# sum_collapse is a full collapse now
assert len(sched) == 1
assert not any(u.op is Ops.REDUCE and u.arg[1] > 0 for u in sched[0].src[0].toposort()), "found reduce in sum collapse"
#lin = Kernel(sched[0].ast)
#assert not any(u.op is Ops.RANGE for u in lin.linearize().uops), "found loop in sum collapse"
def test_assign_fold(self):
a = Tensor.ones(4, 4).contiguous().realize()
m = Tensor.ones(4, 4).shrink(((1, 2), None)).pad(((1, 2), None))
a.assign(a+m)
a.realize()
np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])
@unittest.skipIf(MOCKGPU and isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, CUDARenderer)), "PTX indexes differently. might be ok?")
def test_where_fold(self):
a = Tensor.ones(4, 4).contiguous().realize()
b = a.shrink(((1, 2), None)).pad(((1, 2), None)).bool()
a.assign(b.where(2, a))
linear, var_vals = check_schedule(a, 1)
run_linear(linear, var_vals)
np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])
program = to_program(replace_opts(linear.src[-1].src[0], []), renderer=Device[Device.DEFAULT].renderer)
assert not any(u.op == Ops.WHERE for u in tuple(program.src[1].src)), "found where where where should be folded"
def test_phi_simplification(self):
def helper(t, max_ops=0):
ast = helper_linearizer_opt(t)
uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[1].src)
# ignore kernel optimized IF statements for now
if if_op:=next((u for u in uops if u.op is Ops.IF), None):
uops = uops[:uops.index(if_op)]
assert len(set([u.op for u in uops if u.op in {Ops.RANGE, Ops.SPECIAL}])) == 1, "has either specials or ranges, not both"
reg_stores = [u for u in uops if u.op is Ops.STORE and u.src[0].addrspace == AddrSpace.REG]
assert len(reg_stores) == 0, "STORE to reg should have been simplified"
assert len([u for u in uops if u.op is Ops.MAX]) <= max_ops, "no unnecessary MAX ops"
helper(Tensor.arange(5.5, (3.5*300), 3.5).clone(), max_ops=2)
helper(Tensor.arange(-1, -100, -5).clone(), max_ops=2)
# NOTE: both of these split the reduce (this just wasn't tracked before)
#helper(Tensor.arange(-3.2, 6.7, 0.64), max_ops=2)
#helper(Tensor.arange(256), max_ops=2)
helper(Tensor.arange(255).clone(), max_ops=2)
@unittest.skip("test implicitly depends on certain optimizations")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx for some reason")
def test_grouped_store_phis(self):
"""
float4 acc0 = float4(0.0,0.0,0.0,0.0);
{
acc0 = // ...
}
*((device float4*)(data0+alu2)) = float4(acc0.x,acc0.y,acc0.z,acc0.w);
simplifies to:
*((device float4*)(data0+alu2)) = acc0;
"""
x, y = Tensor.empty(64,64), Tensor.empty(64,64)
out = x.matmul(y)
with Context(TC=0):
ast = helper_linearizer_opt(out)
uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[1].src)
# check that the float4 cast collapses
store_vals = [u.src[1] for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
for val in store_vals:
assert val.dtype == dtypes.float # and val.op is not Ops.VECTORIZE
@unittest.skip("test implicitly depends on certain optimizations")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
def test_grouped_store_values(self):
x = Tensor.randn((4,3,6,6)).realize()
out = x.flip((0,1)).contiguous()
ast = helper_linearizer_opt(out)
store_val = [u.src[1] for u in tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[1].src) if u.op is Ops.STORE][0]
assert store_val.dtype == dtypes.float and store_val.op is not Ops.STACK
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
def test_grouped_store_locals_and_globals(self):
x, y = Tensor.empty(64, 64), Tensor.empty(64, 64)
out = x@y
opt = [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.GROUPTOP, 0, 8),
Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 2)] # upcast accs in both reduces
ast = helper_linearizer_opt(out, opts=[opt])
def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) for v in uop.src])
uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[1].src)
local_stores = [u for u in uops if u.op is Ops.STORE and u.addrspace == AddrSpace.LOCAL]
global_stores = [u for u in uops if u.op is Ops.STORE and u.addrspace == AddrSpace.GLOBAL]
barrier = [u for u in uops if u.op is Ops.BARRIER]
assert len(barrier) == 1
# check that the float4 cast collapses for all stores
for store in local_stores+global_stores:
assert store.src[1].max_numel() > 1, f"store shape {store.src[1].shape} on {store.addrspace}"
# # check the children's vins
# TODO: src ALU are not the same, should it?
# assert barrier.src == tuple(local_stores)
assert len([u for u in uops if u.op is Ops.IF])
@unittest.skip("test implicitly depends on certain optimizations")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx for some reason")
def test_grouped_store_local_only(self):
x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
r = (x@y).relu()
ast = helper_linearizer_opt(r)
uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[1].src)
stores = [u for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
# the float4 value stores directly in lds and we skip upcast
self.assertEqual(stores[0].src[1].dtype, dtypes.float)
#assert stores[0].src[-1].op is not Ops.VECTORIZE
# the global store doesn't change
assert stores[1].src[1].dtype == dtypes.float
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
def test_two_grouped_stores_local(self):
# GROUP on both reduces puts two LOCAL buffers in one kernel, and the store to each needs its own barrier
a = Tensor.rand(32, 32).realize()
opts = [Opt(OptOps.GROUP, 1, 4), Opt(OptOps.GROUP, 2, 4)]
ast = helper_linearizer_opt(single_kernel_softmax(a), [opts])
uops = to_program(replace_opts(ast, opts), renderer=Device[Device.DEFAULT].renderer).src[1].src
self.assertEqual(len([u for u in uops if u.op is Ops.BARRIER]), 2)
# *** helpers ***
def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
if isinstance(r, Tensor): r = [r]
linear, var_vals = Tensor.linear_with_vars(*r)
run_linear(UOp(Ops.LINEAR, src=linear.src[:-1]), var_vals) # run all kernels except the last one
last_call = linear.src[-1]
ast = last_call.src[0]
assert ast.op is Ops.SINK, f"helper_realized_ast expects a SINK {last_call}"
last_bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]
# now all input buffers in last_call should be realized
# create fresh buffers for the outputs
bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(ast.src) else x for i,x in enumerate(last_bufs)]
# ensure buffers are allocated
for b in bufs: b.ensure_allocated()
return ast, bufs
def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs):
assert isinstance(ast, UOp), "ast must be UOp"
inbufs = [x.uop.base.buffer for x in inputs]
outbufs = [Buffer(inbufs[-1].device if inbufs else Device.DEFAULT, out.size, out.src[1].dtype).allocate() for out in ast.src]
_helper_linearizer_opt_ast(ast, outbufs+inbufs, *args, **kwargs)
def helper_linearizer_opt(r:Tensor|list[Tensor], *args, **kwargs):
realized_ast, real_bufs = helper_realized_ast(r)
_helper_linearizer_opt_ast(realized_ast, real_bufs, *args, **kwargs)
return realized_ast
def copyout_outputs(outbufs:list[Buffer]) -> list[np.ndarray]:
return [np.frombuffer(x.as_memoryview(), _to_np_dtype(x.dtype)) for x in outbufs]
def reset_bufs(bufs:list[Buffer]):
for buf in bufs: buf.copy_from(Buffer("PYTHON", buf.size, buf.dtype, opaque=memoryview(bytearray(buf.nbytes))))
def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[],
apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]):
outbufs = real_bufs[:len(realized_ast.src)]
wanna_output = [np.array(x).flatten() for x in wanna_output]
buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in real_bufs]
for u,b in zip(buf_uops, real_bufs): buffers[u] = b
def run_prg(opts):
ast = realized_ast if opts is None else replace_opts(realized_ast, list(opts))
run_linear(UOp(Ops.LINEAR, src=(ast.call(*buf_uops),)))
def check_opt(opts):
reset_bufs(outbufs)
run_prg(opts)
for x,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(x, want, atol=atol, rtol=rtol)
# Get baseline if it is not provided, which is not optimized at all.
run_prg(opts=())
if len(wanna_output) == 0: wanna_output = copyout_outputs(outbufs)
else:
for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol)
# Check correctness of handcoded optimiztions.
reset_bufs(outbufs)
run_prg(opts=None)
for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol)
for x in opts: # Check custom transformations if any.
check_opt(([Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, 1))] if apply_tc else [])+x)
if __name__ == '__main__':
unittest.main()

31
artifacts/package_sources/tinygrad/test/backend/test_linearizer_dumb.py Normal file
View File

@@ -0,0 +1,31 @@
# ruff: noqa: E501
# tests where the Linearizer is doing something dumb
# like test_linearizer_failures, but they don't have to fail
import unittest
from tinygrad import Device, dtypes
from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
from tinygrad.codegen.opt.search import Opt, OptOps
from tinygrad.codegen import to_program
class TestLinearizerFailure(unittest.TestCase):
@unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
def test_failure_beam_mnist(self):
c0 = UOp.param(0, dtypes.uchar, (4014080,))
c1 = UOp.range(UOp.const(512), 0, AxisType.GLOBAL)
c2 = UOp.range(UOp.const(784), 1, AxisType.GLOBAL)
c3 = UOp.range(UOp.const(10), 3, AxisType.GLOBAL)
c4 = UOp.param(1, dtypes.int, (512,))
c5 = c4.index(c1.valid(UOp.const(True)))
c6 = UOp.range(UOp.const(6000), 1004, AxisType.REDUCE)
c7 = UOp.range(UOp.const(3750), 2006, AxisType.REDUCE)
c8 = UOp.range(UOp.const(16), 2007, AxisType.GROUP_REDUCE)
c9 = UOp.param(2, dtypes.uchar, (47040000,))
c10 = c9.index((((c3*UOp.const(4704000))+c2)+(c6*UOp.const(784))).valid(UOp.const(True)))
c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(6000))+c6)+((c7*UOp.const(16))+c8)).alu(Ops.CMPLT, UOp.const(59999)).where(UOp.const(0).cast(dtypes.int), UOp.const(1).cast(dtypes.int)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(-1).cast(dtypes.int))).where(UOp.const(0).cast(dtypes.uchar), c10).reduce(c6, arg=Ops.ADD)
c12 = c0.index((((c1*UOp.const(7840))+(c2*UOp.const(10)))+c3).valid(UOp.const(True))).store(c11).end(c1, c2, c3)
ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None))
_ = to_program(ast, Device["METAL"].renderer)
if __name__ == '__main__':
unittest.main()

165
artifacts/package_sources/tinygrad/test/backend/test_llama_kernels.py Normal file
View File

@@ -0,0 +1,165 @@
import unittest, functools
from tinygrad import Tensor, Device, dtypes, Context, GlobalCounters
from tinygrad.helpers import getenv
from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8
from extra.llama_kernels.fused_ce import fused_ce_loss
from extra.llama_kernels import local_abs_max
from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed, quantize_fp8_scalar
from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
from extra.thunder.amd.fa import custom_fused_qkv_rope_backward, fused_qkv_rope
from test.helpers import needs_second_gpu, assert_kernel_count
from test.backend.test_asm_gemm import has_hipcc
def run_fused_ce(bs:int, seqlen:int, vocab:int, label_smoothing:float=0.0) -> None:
Tensor.manual_seed(0)
logits_rand = Tensor.randn(bs, seqlen, vocab).cast(dtypes.bfloat16)
targets = Tensor.randint(bs, seqlen, high=vocab, dtype=dtypes.int32)
logits, logits_ref = logits_rand.clone(), logits_rand.detach().float().contiguous()
with Context(DEBUG=0):
Tensor.realize(logits, logits_ref, targets)
loss = fused_ce_loss(logits, targets, label_smoothing=label_smoothing)
loss.backward()
Tensor.realize(loss, logits.grad)
ref = logits_ref.sparse_categorical_crossentropy(targets, label_smoothing=label_smoothing)
ref.backward()
Tensor.realize(ref, logits_ref.grad)
assert logits.grad.shape == (bs, seqlen, vocab)
with Context(DEBUG=0):
assert loss.allclose(ref, atol=2e-3, rtol=2e-3).item(), "forward mismatch"
assert logits.grad.allclose(logits_ref.grad, atol=2e-3, rtol=2e-3).item(), "grad mismatch"
class TestFusedCE(unittest.TestCase):
def setUp(self):
if dtypes.bfloat16 not in Device[Device.DEFAULT].renderer.supported_dtypes(): self.skipTest("need bfloat16")
def test_fused_ce_1_2_16(self): run_fused_ce(1, 2, 16, label_smoothing=0.2)
def test_fused_ce_2_16_128(self): run_fused_ce(2, 16, 128)
def test_fused_ce_4_128_1024(self): run_fused_ce(4, 128, 1024, label_smoothing=0.2)
# note: this is the shape used in llama 8b
#def test_fused_ce_smoothing_16_1024_128256(self): run_fused_ce(16, 1024, 128256, label_smoothing=0.2)
def run_quantize_fp8(shape:tuple[int, ...], delayed:bool=True) -> None:
Tensor.manual_seed(0)
x = Tensor.randn(*shape).cast(dtypes.bfloat16).contiguous()
amax_state = Tensor.full((), 2.0, dtype=dtypes.float32).contiguous()
with Context(DEBUG=0): Tensor.realize(x, amax_state)
if delayed:
amax_out = Tensor.zeros((), dtype=dtypes.float32, device=x.device).realize()
fp8, inv_scale = quantize_fp8_delayed(x, amax_state, amax_out, FP8_DTYPE)
ref_fp8, ref_inv_scale, ref_new_amax = quantize_fp8(x, amax_state=amax_state)
Tensor.realize(fp8, inv_scale)
Tensor.realize(ref_fp8, ref_inv_scale, ref_new_amax)
else:
fp8 = quantize_fp8_scalar(x, amax_state, FP8_DTYPE)
ref_fp8, _, _ = quantize_fp8(x, amax_state=amax_state)
Tensor.realize(fp8)
Tensor.realize(ref_fp8)
with Context(DEBUG=0):
assert fp8.cast(dtypes.float).allclose(ref_fp8.cast(dtypes.float), atol=0, rtol=0).item(), "fp8 mismatch"
if delayed:
assert inv_scale.allclose(ref_inv_scale, atol=0, rtol=0).item(), "inv_scale mismatch"
assert amax_out.allclose(ref_new_amax, atol=0, rtol=0).item(), \
f"amax mismatch: got={amax_out.item()} ref={ref_new_amax.item()} diff={abs(amax_out.item()-ref_new_amax.item())}"
@unittest.skipUnless(Device.DEFAULT == "AMD", "requires atomic max")
class TestQuantizeFP8(unittest.TestCase):
def setUp(self):
ren = Device[Device.DEFAULT].renderer
if dtypes.bfloat16 not in ren.supported_dtypes(): self.skipTest("need bfloat16")
if not ren.has_local or not ren.has_shared: self.skipTest("need local/shared")
def test_scalar(self): run_quantize_fp8((getenv("N", 1024), 32), delayed=False)
def test_delayed(self): run_quantize_fp8((getenv("N", 2048), 1024))
@needs_second_gpu
def test_multi(self):
devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(8))
x = Tensor.empty(2048*8, 1024, dtype=dtypes.bfloat16, device=devs).uop.unshard(0)
x = Tensor(x, device=devs)
amax_state = Tensor.full((), 2.0, dtype=dtypes.float32, device=devs).contiguous()
amax_out = Tensor.zeros((), dtype=dtypes.float32, device=devs).realize()
fp8, _ = quantize_fp8_delayed(x, amax_state, amax_out, FP8_DTYPE)
Tensor.realize(fp8)
assert fp8.uop.shape == x.uop.shape
assert amax_out.shape == ()
class TestLocalAmax(unittest.TestCase):
def test_multi_tensor_local_shard_amax(self):
devices = ("CPU:0", "CPU:1")
x = Tensor.arange(16).reshape(4, 4).cast(dtypes.float).clone(devices[0]).realize().shard(devices, axis=0).realize()
GlobalCounters.reset()
out = (x * local_abs_max(x)).clone().realize()
assert_kernel_count(2)
self.assertEqual(out.tolist(), [[0., 7., 14., 21.], [28., 35., 42., 49.], [120., 135., 150., 165.], [180., 195., 210., 225.]])
@unittest.skipUnless(has_hipcc() and Device.DEFAULT == "AMD", "requires hipcc to compile and amd device to run")
class TestFusedQKVRoPE(unittest.TestCase):
SHAPE = (2, 8192, 32, 8, 128)
def rand_bf16(self, *shape:int) -> Tensor:
return (Tensor.randn(*shape) * 0.1).cast(dtypes.bfloat16).contiguous().realize()
def freqs_cis(self) -> Tensor:
_, N, _, _, D = self.SHAPE
return precompute_freqs_cis(D, N * 2).cast(dtypes.bfloat16).clone().realize()
def test_llama31_8b_forward(self):
Tensor.manual_seed(0)
B, N, H, H_KV, D = self.SHAPE
GROUP = H // H_KV
freqs_cis = self.freqs_cis()
x = self.rand_bf16(B, N, H_KV * (GROUP + 2) * D)
q, k, v = fused_qkv_rope(x, freqs_cis, H, H_KV, D)
Tensor.realize(q, k, v)
packed_ref = x.reshape(B, N, H_KV, GROUP + 2, D)
q_ref = packed_ref[:, :, :, :GROUP].reshape(B, N, H, D)
k_ref, v_ref = packed_ref[:, :, :, GROUP], packed_ref[:, :, :, GROUP+1]
q_ref, k_ref = apply_rotary_emb(q_ref, k_ref, freqs_cis[:, :N])
q_ref, k_ref, v_ref = q_ref.cast(dtypes.bfloat16), k_ref.cast(dtypes.bfloat16), v_ref.cast(dtypes.bfloat16)
Tensor.realize(q_ref, k_ref, v_ref)
with Context(DEBUG=0):
self.assertTrue(q.allclose(q_ref, atol=2e-2, rtol=0).item(), "Q forward mismatch")
self.assertTrue(k.allclose(k_ref, atol=2e-2, rtol=0).item(), "K forward mismatch")
self.assertTrue(v.allclose(v_ref, atol=0, rtol=0).item(), "V forward mismatch")
def test_llama31_8b_backward(self):
Tensor.manual_seed(1)
B, N, H, H_KV, D = self.SHAPE
PARTIALS = 2
GROUP = H // H_KV
freqs_cis = self.freqs_cis()
dq = self.rand_bf16(B, N, H, D)
dk_partial = self.rand_bf16(B * PARTIALS, N, H_KV, D)
dv_partial = self.rand_bf16(B * PARTIALS, N, H_KV, D)
# Invert Flash Attention's dQ layout transform to reproduce its native buffer.
dq_native = dq.transpose(1, 2).reshape(B, H, N//16, 4, 4, 4, 2, D//32, 2, 2) \
.permute(0, 1, 2, 5, 6, 8, 7, 3, 4, 9).reshape(B, H, N, D).contiguous().realize()
dx = Tensor.empty(B, N, H_KV * (GROUP + 2) * D, dtype=dtypes.bfloat16)
arch = Device[Device.DEFAULT].renderer.target.arch
fxn = functools.partial(custom_fused_qkv_rope_backward, device=Device.DEFAULT, arch=arch,
B=B, N=N, H=H, H_KV=H_KV, D=D)
dx = Tensor.custom_kernel(dx, dq_native, dk_partial, dv_partial, freqs_cis, fxn=fxn)[0].realize()
def inverse_rope(x:Tensor) -> Tensor:
x = x.reshape(*x.shape[:-1], D//2, 2).float()
cs = freqs_cis[:, :N].float()
return Tensor.stack(x[..., 0] * cs[..., 0] + x[..., 1] * cs[..., 1],
-x[..., 0] * cs[..., 1] + x[..., 1] * cs[..., 0], dim=-1).flatten(-2).cast(dtypes.bfloat16)
dq_ref = inverse_rope(dq).reshape(B, N, H_KV, GROUP, D)
dk_ref = inverse_rope(dk_partial.float().reshape(B, PARTIALS, N, H_KV, D).sum(1).cast(dtypes.bfloat16)).unsqueeze(3)
dv_ref = dv_partial.float().reshape(B, PARTIALS, N, H_KV, D).sum(1).cast(dtypes.bfloat16).unsqueeze(3)
ref = Tensor.cat(dq_ref, dk_ref, dv_ref, dim=3).reshape(*dx.shape).realize()
with Context(DEBUG=0): self.assertTrue(dx.allclose(ref, atol=2e-2, rtol=2e-2).item(), "backward mismatch")
if __name__ == '__main__':
unittest.main()

536
artifacts/package_sources/tinygrad/test/backend/test_multitensor.py Normal file
View File

@@ -0,0 +1,536 @@
import unittest, random
from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable
from tinygrad.uop.ops import Ops, UOp, AxisType
from tinygrad.helpers import getenv, prod, Context
from tinygrad.nn.state import get_parameters
from tinygrad.engine.realize import run_linear, compile_linear
import numpy as np
from hypothesis import given, strategies as strat, settings
from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count
settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
settings.load_profile("my_profile")
d0 = f"{Device.DEFAULT}:0"
d1 = f"{Device.DEFAULT}:1"
d2 = f"{Device.DEFAULT}:2"
d3 = f"{Device.DEFAULT}:3"
d4 = f"{Device.DEFAULT}:4"
d5 = f"{Device.DEFAULT}:5"
devices_2 = (d1, d2)
devices_3 = (d1, d2, d3)
devices_4 = (d1, d2, d3, d4)
N = 128
# shard_x is "data parallel"
# shard_w is "model parallel"
def _test_allreduce(t:Tensor):
aa = (t[0:64] + t[64:128] + t[128:192] + t[192:256]).repeat([4,1]).realize()
ts = t.shard(devices_4, 0).realize()
b = Tensor(UOp.allreduce(ts.uop, Ops.ADD, ts.device))
b.realize()
return aa, b
@unittest.skipIf(not_support_multi_device(), "no multi")
class TestMultiTensor(unittest.TestCase):
@needs_second_gpu
def setUp(self): pass
def test_to(self):
X = Tensor.ones(256).contiguous().realize()
X.to_(devices_2)
assert X.shape == (256,)
(X + X).realize()
def test_gradient(self):
X = Tensor.ones(256).contiguous().realize()
X.to_(devices_2)
grad = X.sum().gradient(X)[0]
grad.realize()
def test_shard(self):
X = Tensor.ones(256).contiguous().realize()
X.shard_(devices_2, 0)
assert X.uop.src[0].shape == (128,)
# the MULTI carries and ends the DEVICE range as its second src
assert X.uop.src[1].op is Ops.RANGE and X.uop.src[1].arg[-1] is AxisType.DEVICE
assert X.uop.ended_ranges == X.uop.src[1:]
(X + X).realize()
@unittest.expectedFailure # TODO: fix
def test_shard_empty(self):
GlobalCounters.reset()
X = Tensor.empty(256).shard(devices_2, 0).realize()
assert_kernel_count(0)
(X + X).realize()
# TODO: fix this to not copy on the src device
@unittest.expectedFailure
def test_shard_no_recompile(self):
X = Tensor.ones(256).contiguous().realize()
X.shard_(devices_2, 0)
out = (X + X)
linear = compile_linear(out.schedule_linear())
names = [call.src[0].src[0].arg.name for call in linear.src if call.src[0].op is Ops.PROGRAM]
run_linear(linear)
self.assertEqual(len(set(names)), 1, "function was relinearized")
def test_shard_beam(self):
cpu_2 = ("CPU:1", "CPU:2")
src = Tensor.ones(16).shard(cpu_2, 0).realize()
pad = src.to(cpu_2[::-1]).schedule_linear().src[0]
with Context(BEAM=1, IGNORE_BEAM_CACHE=1): prg = compile_linear(UOp(Ops.LINEAR, src=(pad,))).src[0].src[0]
self.assertNotEqual(prg.src[0].arg.applied_opts, ())
def test_shard_same_device(self):
X = Tensor.ones(256).contiguous().realize()
X.shard_((d1, X.device), 0)
(X + X).realize()
def test_numpy(self):
X = Tensor.ones(256)
X.shard_((d1, d2), 0)
np.testing.assert_allclose(X.numpy(), 1)
def test_four_add(self):
X = Tensor.ones(256, 256).contiguous().realize()
W = Tensor.ones(256, 256).contiguous().realize()
X.shard_(devices_4, 1)
W.shard_(devices_4, None)
O = X + W
np.testing.assert_allclose(O.numpy(), 2)
def test_elementwise_dtype(self):
Tensor.manual_seed(0)
X = Tensor.randn(8, 8).realize()
W = Tensor.randn(8, 8).realize()
X.shard_(devices_4, 0)
W.shard_(devices_4, 0)
O = X.shrink(((0, 2), None)) * W.shrink(((0, 2), None)) < 2
np.testing.assert_allclose(O.numpy(), X.numpy()[0:2]*W.numpy()[0:2] < 2)
def test_shrink_on_shard_axis(self):
X = Tensor.arange(4*4).reshape(4,4).clone().realize()
X_np = X.numpy()
X.shard_(devices_2, 0)
# only shrink on the device that owns the shard, this is enabled by the mselect simplifier
for i in range(2):
xt = X[i*2:i*2+2].contiguous()
linear, var_vals = xt.linear_with_vars()
#kernels = [call for call in linear.src if call.src[0].op is Ops.SINK]
#self.assertEqual(len(kernels), 1)
#self.assertEqual(kernels[0].src[1].buffer.device, devices_2[i])
run_linear(linear, var_vals)
np.testing.assert_equal(xt.numpy(), X_np[i*2:i*2+2])
@given(strat.sampled_from((devices_2, devices_3)),
strat.sampled_from((Ops.ADD, Ops.MUL, Ops.MAX)),
strat.sampled_from((None, 0, 1)), strat.sampled_from((None, 0, 1)))
def test_simple_reduce(self, devices, rop, shard_axis, reduce_axis):
N = 4 * len(devices)
X = (Tensor.rand(N*N)-1).reshape(N, N).shard_(devices, shard_axis)
n = X.numpy()
f = {Ops.ADD: lambda x: x.sum(reduce_axis), Ops.MUL: lambda x: x.prod(reduce_axis), Ops.MAX: lambda x: x.max(reduce_axis)}[rop]
fX = f(X)
fn = f(n)
np.testing.assert_allclose(fX.numpy(), fn, rtol=1e-6, atol=1e-6)
def test_stack(self):
X = Tensor.rand(4, 4).shard_(devices_2, 0)
Y = Tensor.rand(4, 4).shard_(devices_2, 0)
Z = Tensor.rand(4, 4).shard_(devices_2, 1) # mismatched shard axis gets resharded
for dim in (0, 1):
np.testing.assert_allclose(Tensor.stack(X, Y, Z, dim=dim).numpy(), np.stack([X.numpy(), Y.numpy(), Z.numpy()], axis=dim))
grad = Tensor.stack(X, Y).sum().gradient(X)[0]
np.testing.assert_allclose(grad.numpy(), 1)
def test_allreduce_naive(self):
with Context(RING=0):
a,b = _test_allreduce(Tensor.rand(256, 256))
np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
def test_allreduce_ring(self):
with Context(RING=2):
a,b = _test_allreduce(Tensor.rand(256, 256))
np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
def test_allreduce_all2all(self):
with Context(ALL2ALL=2):
a,b = _test_allreduce(Tensor.rand(256, 256))
np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
def test_copy_jit(self):
@TinyJit
def copy_tensor(x:Tensor): return (x.to(f"{x.device.split(':')[0]}:1") + 1)
for _ in range(5):
t = Tensor.rand(256).realize()
x = copy_tensor(t)
np.testing.assert_equal((t+1).numpy(), x.numpy())
def test_allreduce_naive_jit(self):
with Context(RING=0):
jit_allreduce = TinyJit(_test_allreduce)
for _ in range(5):
a,b = jit_allreduce(Tensor.rand(256, 256))
np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
def test_allreduce_ring_jit(self):
with Context(RING=2):
jit_allreduce = TinyJit(_test_allreduce)
for _ in range(5):
a,b = jit_allreduce(Tensor.rand(256, 256))
np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
def test_multitensor_jit_input(self):
@TinyJit
def f(x): return (x+1).contiguous().sum()
for _ in range(5):
tt = Tensor.arange(0, 4).clone().realize().shard((d1,d2), 0).realize()
out = f(tt)
assert out.item() == 1+2+3+4
def test_multitensor_inside_jit(self):
@TinyJit
def f(x): return (x.shard((d1,d2), 0)+1).contiguous().sum()
for _ in range(5):
tt = Tensor.arange(0, 4).clone().realize()
out = f(tt)
assert out.item() == 1+2+3+4
def test_fuzz_allreduce(self):
random.seed(41)
for it in range(2):
for n in range(2, 4+1):
shape = tuple([(n if i == 0 else 1) * random.randint(1, 10) for i in range(random.randint(1, 4))])
t = Tensor.rand(shape).shard_(tuple([d0, d1, d2, d3][:n]), 0)
with Context(RING=0):
a = Tensor(UOp.allreduce(t.uop, Ops.ADD, t.device))
with Context(RING=2):
b = Tensor(UOp.allreduce(t.uop, Ops.ADD, t.device))
diff = a - b
mean_err = diff.reshape((prod(diff.shape),)).abs().mean().numpy()
max_err = diff.reshape((prod(diff.shape),)).abs().max().numpy()
assert mean_err < 1e-6, f"big mean error, iteration {it}_{n}"
assert max_err < 1e-6, f"big max error, iteration {it}_{n}"
def _test_model_train_step(self, m, fake_image, labels):
from tinygrad.nn.optim import LARS
optimizer = LARS(get_parameters(m), 0.1)
optimizer.zero_grad()
m.load_from_pretrained()
output = m(fake_image).sparse_categorical_crossentropy(labels, label_smoothing=0.1)
output.backward()
grad = m.conv1.weight.grad.numpy()
fake_image_sharded = fake_image.shard(devices_2, axis=0)
labels_sharded = labels.shard(devices_2, axis=0)
for p in get_parameters(m): p.shard_(devices_2).realize()
GlobalCounters.reset()
optimizer.zero_grad()
shard_output = m(fake_image_sharded).sparse_categorical_crossentropy(labels_sharded, label_smoothing=0.1)
shard_output.backward()
shard_grad = m.conv1.weight.grad.numpy()
# sometimes there is zeros in these grads... why?
np.testing.assert_allclose(grad, shard_grad, atol=1e-5, rtol=1e-5)
@slow
def test_data_parallel_resnet_train_step(self):
from extra.models.resnet import ResNet18
fake_image = Tensor.rand((2, 3, 224//16, 224//16))
labels = Tensor.randint(2, low=0, high=1000)
m = ResNet18()
self._test_model_train_step(m, fake_image, labels)
def test_data_parallel_simple_train_step(self):
class Model:
def __init__(self): self.conv1 = nn.Linear(128,128)
def __call__(self, x): return self.conv1(x)
def load_from_pretrained(self): pass
fake_image = Tensor.rand((128,))
labels = Tensor.randint(2, low=0, high=127)
m = Model()
self._test_model_train_step(m, fake_image, labels)
def test_assign_kv_cache_multi(self):
bsz, max_context = 2, 8
class Attn:
@TinyJit
def __call__(self, xk:Tensor, start_pos:UOp):
seqlen = xk.shape[1]
if not hasattr(self, "cache_k"):
self.cache_k = Tensor.zeros(bsz, max_context, 1, 1).shard(devices_2).contiguous().realize()
keys = self.cache_k.shrink((None, (0, start_pos), None, None)).cat(xk, dim=1).contiguous() if start_pos > 0 else xk
self.cache_k.assign(keys.pad((None,(0,max_context-start_pos-seqlen),None,None)).contiguous()).realize()
attn = Attn()
xk = Tensor.ones(bsz, 3, 1, 1).shard(devices_2).contiguous()
attn(xk, 0)
for i in range(3,6):
# copied from LLaMA
start_pos = Variable("start_pos", 1, max_context).bind(i)
xk = Tensor.ones(bsz, 1, 1, 1).shard(devices_2).contiguous()
attn(xk, start_pos)
out = attn.cache_k.flatten().numpy()
np.testing.assert_allclose(out, [1.,1.,1.,1.,1.,1.,0.,0.,1.,1.,1.,1.,1.,1.,0.,0.])
def test_multi_tensor_jit_graph_assign_updates_each_shard(self):
@TinyJit
def jf(out: Tensor) -> Tensor:
tmp = (Tensor.arange(4, dtype=dtypes.float).clone().shard(devices_2, 0) + 1).contiguous().realize()
out.assign((tmp + 1).contiguous()).realize()
return out
out = Tensor.full((4,), -1.0).shard(devices_2, 0).contiguous().realize()
expected = np.arange(4, dtype=np.float32) + 2
for _ in range(5):
out.assign(Tensor.full((4,), -1.0).shard(devices_2, 0).contiguous()).realize()
jf(out)
np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-5)
assert jf.captured is not None
@unittest.skip("test broken")
def test_multi_device_jit_graph(self):
if Device[d0].graph is None or Device[d1].graph is None: raise unittest.SkipTest("only test graphs")
@TinyJit
def jf(a: Tensor, b: Tensor, c: Tensor, d:Tensor):
# Create 80 entries on device 0: 2 batches.
for _ in range(40):
a = ((a + b).realize() + (a * b).realize()).realize()
# Create 80 entries on device 1: 2 batches.
for _ in range(40):
c = ((c + d).realize() + (c * d).realize()).realize()
# Create a copy from device 0 to 1: 1 entry.
a = a.to(d1).realize()
# Creates one last entry on device 1: 1 batch.
return (a + c).realize()
a = Tensor.randn(10, 10, device=d0).realize()
b = Tensor.randn(10, 10, device=d0).realize()
c = Tensor.randn(10, 10, device=d1).realize()
d = Tensor.randn(10, 10, device=d1).realize()
ref = jf(a, b, c, d).numpy()
for _ in range(5):
o = jf(a, b, c, d).numpy()
np.testing.assert_allclose(ref, o, atol=1e-4, rtol=1e-5)
# Checking that 2 graphs per device, 1 copy and 1 last graph on device 1 are created.
sis = jf.captured.linear.src
assert len(sis) == 6
for si in (sis[0], sis[1], sis[2], sis[3], sis[5]):
assert call_is_graph(si)
assert sis[4].src[0].op is Ops.COPY
def test_rand_on_multiple_devices(self):
# different devices generate different rand
d0_rand = Tensor.rand(256, device=d0).realize()
d1_rand = Tensor.rand(256, device=d1).realize()
assert not np.allclose(d0_rand.numpy(), d1_rand.numpy())
def test_rand_on_multiple_devices_manual_seed(self):
Tensor.manual_seed(123)
d0_rand = Tensor.rand(2, device=d0).tolist()
d1_rand = Tensor.rand(2, device=d1).tolist()
# manual_seed again gives the same values
Tensor.manual_seed(123)
d0_rand2 = Tensor.rand(2, device=d0).tolist()
d1_rand2 = Tensor.rand(2, device=d1).tolist()
self.assertEqual(d0_rand, d0_rand2)
self.assertEqual(d1_rand, d1_rand2)
# device seed is only determined by init order, so flipping init order flips rands
Tensor.manual_seed(123)
d1_rand_flip = Tensor.rand(2, device=d1).tolist()
d0_rand_flip = Tensor.rand(2, device=d0).tolist()
self.assertEqual(d0_rand, d1_rand_flip)
self.assertEqual(d1_rand, d0_rand_flip)
def test_const_like_shrink_on_shard_axis(self):
t = Tensor.ones(16, 16, dtype=dtypes.int).shard(devices_2, axis=0)
out = t.const_like(2)[:, :8]
linear, var_vals = check_schedule(out, 0)
run_linear(linear, var_vals)
self.assertEqual(out.tolist(), [[2]*8]*16)
@unittest.skipIf(not_support_multi_device(), "no multi")
class TestHandleData(unittest.TestCase):
@needs_second_gpu
def test_copied_to_device(self):
device = (d0, d1, d2, d3)
t = Tensor([1, 2, 3, 4]).shard(device).realize()
not_covered = t.to(d5)
sched = not_covered.schedule_linear().src
assert len(sched) == 1
# setup again because create_schedule has side effect
t = Tensor([1, 2, 3, 4]).shard(device).realize()
not_covered = t.to(d5)
assert not_covered.realize().tolist() == [1, 2, 3, 4]
for d in device:
t = Tensor([1, 2, 3, 4]).shard(device).realize()
covered = t.to(d)
sched = covered.schedule_linear().src
# TODO: this isn't optimized out anymore
#assert len(sched) == 0
# setup again because create_schedule has side effect
t = Tensor([1, 2, 3, 4]).shard(device).realize()
covered = t.to(d)
assert covered.realize().tolist() == [1, 2, 3, 4]
@unittest.skipIf(not_support_multi_device(), "need multi")
class TestMultiBufferView(unittest.TestCase):
@needs_second_gpu
def setUp(self): pass
def _check(self, a_ref:Tensor, a_multi:Tensor, view_fn):
b_ref = view_fn(a_ref)
b_multi = view_fn(a_multi).contiguous()
linear, var_vals = b_multi.linear_with_vars()
if all(not d.startswith(("WEBGPU", "CL")) for d in b_multi.device):
compiled = [call for call in linear.src if call.src[0].op is Ops.SINK]
self.assertEqual(len(compiled), 0, f"expected zero compiled kernels, got {len(compiled)}")
run_linear(linear, var_vals)
np.testing.assert_equal(b_multi.numpy(), b_ref.numpy())
@unittest.skip("flaky on LLVM")
def test_shrink_non_shard_axis(self):
ref = Tensor.arange(8*4*10).reshape(8, 4, 10).clone().realize()
a = Tensor.arange(8*4*10).reshape(8, 4, 10).clone().shard(devices_2, axis=1).realize()
self._check(ref, a, lambda t: t[3])
def test_shrink_2d(self):
ref = Tensor.arange(6*4).reshape(6, 4).clone().realize()
a = Tensor.arange(6*4).reshape(6, 4).clone().shard(devices_2, axis=1).realize()
self._check(ref, a, lambda t: t.shrink(((1, 4), None)))
def test_reshape_then_shrink(self):
ref = Tensor.arange(8*6).reshape(8, 6).clone().realize()
a = Tensor.arange(8*6).reshape(8, 6).clone().shard(devices_2, axis=1).realize()
self._check(ref, a, lambda t: t.reshape(4, 2, 6)[1])
def test_chained_shrink(self):
ref = Tensor.arange(10*8).reshape(10, 8).clone().realize()
a = Tensor.arange(10*8).reshape(10, 8).clone().shard(devices_2, axis=1).realize()
self._check(ref, a, lambda t: t.shrink(((2, 8), None)).shrink(((1, 4), None)))
def test_4_devices(self):
ref = Tensor.arange(8*12).reshape(8, 12).clone().realize()
a = Tensor.arange(8*12).reshape(8, 12).clone().shard(devices_4, axis=1).realize()
out = a[5].contiguous()
linear, var_vals = out.linear_with_vars()
if all(not d.startswith(("WEBGPU", "CL")) for d in out.device):
compiled = [call for call in linear.src if call.src[0].op is Ops.SINK]
self.assertEqual(len(compiled), 0)
run_linear(linear, var_vals)
np.testing.assert_equal(out.numpy(), ref[5].numpy())
@unittest.skipIf(not_support_multi_device(), "need multi")
class Test2DShard(unittest.TestCase):
def setUp(self):
self.devices_4 = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
self.rng = UOp.range(4, -1, AxisType.DEVICE)
self.rng0, self.rng1 = self.rng // 2, self.rng % 2
def _shard_2d(self, t:Tensor) -> Tensor:
u = t.uop.copy_to_device(self.devices_4)._shard(0, self.rng0)._shard(1, self.rng1).unshard((0, 1), (self.rng0, self.rng1))
return Tensor(u)
def test_2d_shard_basic(self):
ref = Tensor.arange(16).reshape(4, 4).contiguous().realize()
t = self._shard_2d(ref)
out = t.contiguous().realize()
np.testing.assert_equal(out.numpy(), ref.numpy())
def test_2d_shard_elementwise(self):
ref = Tensor.arange(16).reshape(4, 4).contiguous().realize()
t = self._shard_2d(ref)
out = (t + 1).contiguous().realize()
np.testing.assert_equal(out.numpy(), ref.numpy() + 1)
def test_2d_shard_sum_all(self):
ref = Tensor.arange(16).reshape(4, 4).contiguous().realize()
t = self._shard_2d(ref)
out = t.sum().contiguous().realize()
np.testing.assert_equal(out.numpy(), np.array(ref.numpy().sum()))
def test_2d_shard_sum_non_sharded_axis(self):
ref = Tensor.arange(4*4*2).reshape(4, 4, 2).contiguous().realize()
t = self._shard_2d(ref)
out = t.sum(axis=2).contiguous().realize()
np.testing.assert_equal(out.numpy(), ref.numpy().sum(axis=2))
def test_2d_shard_matmul(self):
a = Tensor.arange(16).reshape(4, 4).contiguous().realize()
b = Tensor.arange(16).reshape(4, 4).contiguous().realize()
a_s = self._shard_2d(a)
b_s = self._shard_2d(b)
out = (a_s @ b_s).contiguous().realize()
np.testing.assert_equal(out.numpy(), a.numpy() @ b.numpy())
@unittest.skipIf(not_support_multi_device(), "need multi")
class TestMultiTransformer(unittest.TestCase):
@needs_second_gpu
def test_transformer(self):
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
from extra.models.llama import Transformer
args = {"dim": 32, "n_heads": 1, "n_kv_heads": 1, "n_layers": 2, "norm_eps": 1e-5, "rope_theta": 500000, "vocab_size": 1024,
"hidden_dim": 32, "max_context": 12}
real_model = Transformer(**args)
shard_model = Transformer(**args)
# copy state
nn.state.load_state_dict(shard_model, nn.state.get_state_dict(real_model))
# shard
for k,v in nn.state.get_state_dict(shard_model).items():
if 'scale' in k: v.shard_(device, axis=None) # from quantized
elif '.attention.' in k: v.shard_(device, axis=-1)
elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
elif '.feed_forward.' in k: v.shard_(device, axis=-1)
elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
elif 'output.weight' in k: v.shard_(device, axis=0)
else: v.shard_(device, axis=None)
last_tok = 0
for i in range(5):
real_tok = real_model(Tensor([[last_tok]], device=Device.DEFAULT), i).item()
shard_tok = shard_model(Tensor([[last_tok]], device=device), i).item()
# test kv cache
kv1 = real_model.layers[0].attention.cache_kv.numpy()
kv2 = shard_model.layers[0].attention.cache_kv.numpy()
#print(np.concatenate([kv1[:, :, :, :, 0:1], kv2[:, :, :, :, 0:1]], axis=4))
np.testing.assert_allclose(kv1, kv2, atol=1e-5, rtol=1e-5, err_msg=f"issue at token {i}")
# test token
self.assertEqual(real_tok, shard_tok, f"issue at token {i}")
last_tok = real_tok
@unittest.skip("super slow")
def test_llama1b_full(self):
from tinygrad.helpers import fetch
fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model", "tokenizer.model", subdir="llama3-1b-instruct")
model = fetch("https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q6_K.gguf",
"Llama-3.2-1B-Instruct-Q6_K.gguf", subdir="llama3-1b-instruct")
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
from examples.llama3 import build_transformer
real_model = build_transformer(model, model_size="1B", device=Device.DEFAULT)
shard_model = build_transformer(model, model_size="1B", device=device)
last_tok = 0
real_tok = real_model(Tensor([[last_tok]], device=Device.DEFAULT), 0)
shard_tok = shard_model(Tensor([[last_tok]], device=device), 0)
self.assertEqual(real_tok.item(), shard_tok.item())
if __name__ == '__main__':
unittest.main()

620
artifacts/package_sources/tinygrad/test/backend/test_nn.py Normal file
View File

@@ -0,0 +1,620 @@
#!/usr/bin/env python
import unittest
import numpy as np
import torch
from tinygrad import Tensor, Device, TinyJit, dtypes
from tinygrad.helpers import GlobalCounters, Context
from tinygrad.nn import Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, Linear, Embedding
from tinygrad.nn import BatchNorm, LayerNorm, LayerNorm2d, GroupNorm, InstanceNorm, RMSNorm, LSTMCell
from tinygrad.nn.state import load_state_dict
from test.helpers import check_schedule
from tinygrad.engine.realize import run_linear
from test.helpers import not_support_multi_device, needs_second_gpu, slow
@slow
class TestNN(unittest.TestCase):
def test_batchnorm2d(self, training=False, threed=False, track_running_stats=True):
with Context(TRAINING=training):
szs = [4, 8, 16, 32]
for sz in szs:
# create in tinygrad
bn = BatchNorm(sz, eps=1e-5, track_running_stats=track_running_stats)
bn.weight = Tensor.randn(sz)
bn.bias = Tensor.randn(sz)
if track_running_stats:
bn.running_mean = Tensor.randn(sz)
bn.running_var = Tensor.randn(sz)
bn.running_var.numpy()[bn.running_var.numpy() < 0] = 0
# create in torch
with torch.no_grad():
if threed:
tbn = torch.nn.BatchNorm3d(sz, track_running_stats=track_running_stats).eval()
else:
tbn = torch.nn.BatchNorm2d(sz, track_running_stats=track_running_stats).eval()
tbn.training = training
tbn.weight[:] = torch.tensor(bn.weight.numpy())
tbn.bias[:] = torch.tensor(bn.bias.numpy())
if track_running_stats:
tbn.running_mean[:] = torch.tensor(bn.running_mean.numpy())
tbn.running_var[:] = torch.tensor(bn.running_var.numpy())
if track_running_stats:
np.testing.assert_allclose(bn.running_mean.numpy(), tbn.running_mean.detach().numpy(), rtol=1e-5, atol=1e-6)
np.testing.assert_allclose(bn.running_var.numpy(), tbn.running_var.detach().numpy(), rtol=1e-5, atol=1e-6)
# trial
if threed:
inn = Tensor.randn(2, sz, 3, 3, 3)
else:
inn = Tensor.randn(2, sz, 3, 3)
# in tinygrad
outt = bn(inn)
# in torch
toutt = tbn(torch.tensor(inn.numpy()))
# close
np.testing.assert_allclose(outt.numpy(), toutt.detach().numpy(), rtol=5e-4, atol=1e-6)
if track_running_stats:
np.testing.assert_allclose(bn.running_mean.numpy(), tbn.running_mean.detach().numpy(), rtol=1e-5, atol=1e-6)
np.testing.assert_allclose(bn.running_var.numpy(), tbn.running_var.detach().numpy(), rtol=1e-5, atol=1e-6)
def test_batchnorm2d_training(self): self.test_batchnorm2d(True, False, True)
def test_batchnorm2d_no_running_stats(self): self.test_batchnorm2d(False, False, False)
def test_batchnorm2d_training_no_running_stats(self): self.test_batchnorm2d(True, False, False)
def test_batchnorm3d(self): self.test_batchnorm2d(False, True, True)
def test_batchnorm3d_training(self): self.test_batchnorm2d(True, True, True)
def test_batchnorm3d_no_running_stats(self): self.test_batchnorm2d(False, True, False)
def test_batchnorm3d_training_no_running_stats(self): self.test_batchnorm2d(True, True, False)
def test_batchnorm_axis(self):
sz = (2, 4, 3, 2, 2)
x = Tensor.randn(sz)
weight = Tensor.randn(2, 3)
bias = Tensor.randn(2, 3)
mean = Tensor.randn(2, 3)
invstd = Tensor.randn(2, 3)
a = (x.batchnorm(weight, bias, mean, invstd, axis=(0, 2))
.permute(1, 0, 2, 3, 4).reshape(4, 6, 2, 2))
b = (x.permute(1, 0, 2, 3, 4).reshape(4, 6, 2, 2)
.batchnorm(weight.flatten(), bias.flatten(), mean.flatten(), invstd.flatten()))
t_x = torch.tensor(x.permute(1, 0, 2, 3, 4).reshape(4, 6, 2, 2).numpy())
t_weight, t_bias = torch.tensor(weight.flatten().numpy()), torch.tensor(bias.flatten().numpy())
t_mean, t_invstd = torch.tensor(mean.flatten().numpy()), torch.tensor(invstd.flatten().numpy())
torch.nn.functional.batch_norm(t_x, t_mean, 1.0 / t_invstd**2, t_weight, t_bias)
np.testing.assert_allclose(a.numpy(), b.numpy())
def test_linear(self):
def _test_linear(x, in_dim, out_dim):
# create in tinygrad
model = Linear(in_dim, out_dim)
z = model(x)
# create in torch
with torch.no_grad():
torch_layer = torch.nn.Linear(in_dim, out_dim).eval()
torch_layer.weight[:] = torch.tensor(model.weight.numpy(), dtype=torch.float32)
torch_layer.bias[:] = torch.tensor(model.bias.numpy(), dtype=torch.float32)
torch_x = torch.tensor(x.numpy(), dtype=torch.float32)
torch_z = torch_layer(torch_x)
# test
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
BS, T, in_dim, out_dim = 4, 2, 8, 16
_test_linear(Tensor.randn(BS, in_dim), in_dim, out_dim)
_test_linear(Tensor.randn(BS, T, in_dim), in_dim, out_dim) # test with more dims
def _test_conv(self, tiny_conv, torch_conv, BS, C1, DIMS, C2, K, S, P, D=1):
# create in tinygrad
layer = tiny_conv(C1, C2, kernel_size=K, stride=S, padding=P, dilation=D)
# create in torch
with torch.no_grad():
torch_layer = torch_conv(C1, C2, kernel_size=K, stride=S, padding=P, dilation=D).eval()
torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
torch_layer.bias[:] = torch.tensor(layer.bias.numpy(), dtype=torch.float32)
# test
x = Tensor.uniform(BS, C1, *DIMS)
z = layer(x)
torch_x = torch.tensor(x.numpy())
torch_z = torch_layer(torch_x)
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
def test_conv1d(self): self._test_conv(Conv1d, torch.nn.Conv1d, BS=4, C1=16, DIMS=[224//4], C2=64, K=7, S=2, P=1)
def test_conv2d(self): self._test_conv(Conv2d, torch.nn.Conv2d, BS=4, C1=16, DIMS=[224//4, 224//4], C2=64, K=7, S=2, P=1)
def test_conv1d_same_padding(self):
self._test_conv(Conv1d, torch.nn.Conv1d, BS=8, C1=3, DIMS=[32], C2=16, K=3, S=1, P='same')
def test_conv2d_same_padding_odd_input(self):
self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=16, DIMS=[29, 31], C2=32, K=5, S=1, P='same')
def test_conv2d_same_padding_large_kernel(self):
self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=16, DIMS=[28, 33], C2=32, K=9, S=1, P='same')
def test_conv2d_same_padding_with_dilation(self):
self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=3, DIMS=[28, 28], C2=32, K=3, S=1, P='same', D=3)
def test_conv2d_same_padding_invalid_stride(self):
self.assertRaises(ValueError, Conv2d, in_channels=16, out_channels=32, kernel_size=2, stride=2, padding='same')
def test_conv2d_same_padding_invalid_padding_str(self):
self.assertRaises(ValueError, Conv2d, in_channels=16, out_channels=32, kernel_size=2, stride=1, padding='not_same')
@unittest.skip("Takes too long to compile for Compiled backends")
def test_conv2d_winograd(self):
BS, C1, H, W = 2, 8, 16, 16
C2, K, S, P = 8, 3, 1, 1
# create in tinygrad
layer = Conv2d(C1, C2, kernel_size=K, stride=S, padding=P)
# create in torch
torch_layer = torch.nn.Conv2d(C1, C2, kernel_size=K, stride=S, padding=P).eval()
torch_layer.weight = torch.nn.Parameter(torch.tensor(layer.weight.numpy(), dtype=torch.float32))
torch_layer.bias = torch.nn.Parameter(torch.tensor(layer.bias.numpy(), dtype=torch.float32))
# test
x = Tensor.uniform(BS, C1, H, W)
with Context(WINO=1):
z = layer(x)
m = z.mean()
m.backward()
torch_x = torch.tensor(x.numpy(), requires_grad=True)
torch_z = torch_layer(torch_x)
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
gw = layer.weight.grad.realize()
gb = layer.bias.grad.realize()
gx = x.grad.realize()
torch_z.mean().backward()
np.testing.assert_allclose(gw.numpy(), torch_layer.weight.grad.numpy(), atol=5e-4, rtol=1e-5)
np.testing.assert_allclose(gb.numpy(), torch_layer.bias.grad.numpy(), atol=5e-4, rtol=1e-5)
np.testing.assert_allclose(gx.numpy(), torch_x.grad.numpy(), atol=5e-4, rtol=1e-5)
def test_conv_transpose1d(self):
self._test_conv(ConvTranspose1d, torch.nn.ConvTranspose1d, BS=4, C1=16, DIMS=[224//4], C2=64, K=7, S=2, P=1)
def test_conv_transpose2d(self):
self._test_conv(ConvTranspose2d, torch.nn.ConvTranspose2d, BS=4, C1=16, DIMS=[224//4, 224//4], C2=64, K=7, S=2, P=1)
def test_groupnorm(self):
BS, H, W, C, G = 20, 10, 10, 6, 3
# create in torch
torch_layer = torch.nn.GroupNorm(G, C).eval()
# create in tinygrad
layer = GroupNorm(G, C)
layer.weight = Tensor(torch_layer.weight.detach().numpy())
layer.bias = Tensor(torch_layer.bias.detach().numpy())
for _ in range(10):
# forward
x = Tensor.randn(BS, C, H, W)
z = layer(x)
z.sum().backward()
torch_x = torch.tensor(x.numpy(), requires_grad=True)
torch_z = torch_layer(torch_x)
torch_z.sum().backward()
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
def test_layernorm_forward(self):
N, C, H, W = 20, 5, 10, 10
# create in torch
torch_layer = torch.nn.LayerNorm([H, W]).eval()
# create in tinygrad
layer = LayerNorm([H, W])
layer.weight = Tensor(torch_layer.weight.detach().numpy())
layer.bias = Tensor(torch_layer.bias.detach().numpy())
x = Tensor.empty(N, C, H, W)
z = layer(x)
z.realize()
torch_x = torch.tensor(x.numpy(), requires_grad=True)
torch_z = torch_layer(torch_x)
torch_z.sum().backward()
# TODO: why is torch numbers all 0?
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=5e-6)
def test_layernorm(self):
N, C, H, W = 20, 5, 10, 10
# create in torch
torch_layer = torch.nn.LayerNorm([H, W]).eval()
# create in tinygrad
layer = LayerNorm([H, W])
layer.weight = Tensor(torch_layer.weight.detach().numpy())
layer.bias = Tensor(torch_layer.bias.detach().numpy())
for _ in range(10):
# forward
x = Tensor.randn(N, C, H, W)
z = layer(x)
z.sum().backward()
torch_x = torch.tensor(x.numpy(), requires_grad=True)
torch_z = torch_layer(torch_x)
torch_z.sum().backward()
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
def test_layernorm_2d(self):
N, C, H, W = 20, 5, 10, 10
# create in torch
torch_layer = torch.nn.LayerNorm([C]).eval()
# create in tinygrad
layer = LayerNorm2d(C)
layer.weight = Tensor(torch_layer.weight.detach().numpy())
layer.bias = Tensor(torch_layer.bias.detach().numpy())
for _ in range(10):
# forward
x = Tensor.randn(N, C, H, W)
z = layer(x)
z.sum().backward()
torch_x = torch.tensor(x.numpy(), requires_grad=True)
torch_z = torch_layer(torch_x.permute(0,2,3,1)).permute(0,3,1,2)
torch_z.sum().backward()
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
def test_instancenorm_2d(self):
N, C, H, W = 20, 10, 10, 10
# create in torch
torch_layer = torch.nn.InstanceNorm2d(C, affine=True).eval()
# create in tinygrad
layer = InstanceNorm(C)
layer.weight = Tensor(torch_layer.weight.detach().numpy())
layer.bias = Tensor(torch_layer.bias.detach().numpy())
for _ in range(10):
# forward
x = Tensor.randn(N, C, H, W)
z = layer(x)
z.sum().backward()
torch_x = torch.tensor(x.numpy(), requires_grad=True)
torch_z = torch_layer(torch_x)
torch_z.sum().backward()
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
def test_instancenorm_3d(self):
N, C, D, H, W = 20, 10, 10, 10, 10
# create in torch
torch_layer = torch.nn.InstanceNorm3d(C, affine=True).eval()
# create in tinygrad
layer = InstanceNorm(C)
layer.weight = Tensor(torch_layer.weight.detach().numpy())
layer.bias = Tensor(torch_layer.bias.detach().numpy())
for _ in range(10):
# forward
x = Tensor.randn(N, C, D, H, W)
z = layer(x)
z.sum().backward()
torch_x = torch.tensor(x.numpy(), requires_grad=True)
torch_z = torch_layer(torch_x)
torch_z.sum().backward()
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
# TODO: is this numerical issue or a bug? RANGEIFY big reduce kernel amplifies numerical issue
np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=1e-2, rtol=1e-3)
np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
def test_rmsnorm(self):
class TorchRMSNorm(torch.nn.Module):
# https://github.com/meta-llama/llama/blob/be327c427cc5e89cc1d3ab3d3fec4484df771245/llama/model.py#L34C1-L77C36
def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True):
super().__init__()
self.eps = eps
self.elementwise_affine = elementwise_affine
self.weight = torch.nn.Parameter(torch.ones(dim)) if elementwise_affine else None
def _norm(self, x):
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def forward(self, x):
output = self._norm(x.float()).type_as(x)
return output if self.weight is None else output * self.weight
B, T, embed_size = 4, 10, 20
torch_layer = TorchRMSNorm(embed_size)
layer = RMSNorm(embed_size)
for _ in range(10):
# forward
x = Tensor.randn(B, T, embed_size)
z = layer(x)
z.sum().backward()
torch_x = torch.tensor(x.numpy(), requires_grad=True)
torch_z = torch_layer(torch_x)
torch_z.sum().backward()
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=2e-3, rtol=1e-3)
torch_layer = TorchRMSNorm(embed_size, elementwise_affine=False)
layer = RMSNorm(embed_size, elementwise_affine=False)
for _ in range(10):
# forward
x = Tensor.randn(B, T, embed_size)
z = layer(x)
z.sum().backward()
torch_x = torch.tensor(x.numpy(), requires_grad=True)
torch_z = torch_layer(torch_x)
torch_z.sum().backward()
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
def test_embedding(self):
B, T, embed_size, vocab_size = 4, 10, 20, 28
# create in tinygrad
layer = Embedding(vocab_size, embed_size)
with torch.no_grad():
torch_layer = torch.nn.Embedding(vocab_size, embed_size).eval()
torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
# test
x = Tensor(np.random.randint(0, vocab_size, (B, T), dtype=np.int32))
z = layer(x)
torch_x = torch.tensor(x.numpy())
torch_z = torch_layer(torch_x)
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=1e-8, rtol=1e-8)
# test with empty input length
x = Tensor(np.random.randint(0, vocab_size, (B, 0), dtype=np.int32))
z = layer(x)
torch_x = torch.tensor(x.numpy())
torch_z = torch_layer(torch_x)
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=1e-8, rtol=1e-8)
# test with jit enabled
@TinyJit
def layer_jit(x):
return layer(x).realize()
for _ in range(3):
x = Tensor(np.random.randint(0, vocab_size, (B, T), dtype=np.int32))
z = layer_jit(x)
torch_x = torch.tensor(x.numpy())
torch_z = torch_layer(torch_x)
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=1e-8, rtol=1e-8)
def test_embedding_one_kernel(self, ops=612000, kcount=2):
GlobalCounters.reset()
layer = Embedding(20, 30)
layer.weight = Tensor.zeros_like(layer.weight).contiguous()
a = Tensor([[1, 5, 9, 11],
[12, 19, 8, 1]])
result = layer(a)
linear, var_vals = check_schedule(result, kcount)
run_linear(linear, var_vals)
b = Tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
result = layer(b)
linear, var_vals = check_schedule(result, 1)
run_linear(linear, var_vals)
print(f"Embedding used {GlobalCounters.global_ops} ops")
self.assertLessEqual(GlobalCounters.global_ops, ops)
# TODO: fused with opts uses more ops
def test_embedding_one_kernel_fused(self):
with Context(NOOPT=0):
self.test_embedding_one_kernel(ops=612_000, kcount=2)
def test_embedding_one_kernel_fused_noopt(self):
with Context(NOOPT=1):
self.test_embedding_one_kernel(ops=0, kcount=2)
def test_embedding_shape(self):
vocab_size, embed_size = 10, 16
layer = Embedding(vocab_size, embed_size)
for rank in range(5):
shp = (1,) * rank
a = Tensor([3]).reshape(shp)
result = layer(a)
self.assertEqual(result.shape, shp + (embed_size,))
def test_embedding_regression(self):
# used to fail bounds check
embedding = Embedding(100, 1024)
input_ids = Tensor.empty(16, 16, dtype=dtypes.int)
embedding(input_ids).realize()
def test_load_state_dict(self):
layer = Conv2d(3, 5, kernel_size=3)
state_dict = {
'weight': Tensor.randn(5, 3, 3, 3),
'bias': Tensor.randn(5),
}
load_state_dict(layer, state_dict)
np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())
#https://github.com/pytorch/pytorch/blob/d38164a545b4a4e4e0cf73ce67173f70574890b6/torch/nn/modules/module.py#L2425
def test_load_conv_num_batches_tracked(self):
layer = BatchNorm(sz=1, track_running_stats=False)
state_dict = {
'weight': Tensor.ones(1),
'bias': Tensor.ones(1),
'num_batches_tracked': Tensor.ones(1),
}
load_state_dict(layer, state_dict)
state_dict['num_batches_tracked'] = Tensor.empty()
load_state_dict(layer, state_dict)
layer.num_batches_tracked = Tensor.ones(1)
load_state_dict(layer, state_dict)
@needs_second_gpu
@unittest.skipIf(not_support_multi_device(), "no multi")
def test_load_state_dict_sharded_model(self):
devices = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3")
layer = Conv2d(3, 5, kernel_size=3)
layer.weight.shard_(devices, 3)
layer.bias.shard_(devices, None)
state_dict = {
'weight': Tensor.randn(5, 3, 3, 3).realize(),
'bias': Tensor.randn(5).realize(),
}
load_state_dict(layer, state_dict)
# sharded model shards the state_dict
self.assertEqual(layer.weight.device, devices)
self.assertEqual(layer.weight.uop.axis, 3)
self.assertEqual(layer.bias.device, devices)
self.assertEqual(layer.bias.uop.axis, None)
np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())
@unittest.skipIf(not_support_multi_device, "no multi")
def test_load_state_dict_sharded_dict(self):
devices = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3")
layer = Conv2d(3, 5, kernel_size=3)
state_dict = {
'weight': Tensor.randn(5, 3, 3, 3).shard(devices, 3),
'bias': Tensor.randn(5).shard(devices, None),
}
load_state_dict(layer, state_dict)
# NOTE: model is not sharded, still not sharded after load_state_dict
self.assertEqual(layer.weight.device, Device.DEFAULT)
self.assertEqual(layer.bias.device, Device.DEFAULT)
np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())
@needs_second_gpu
@unittest.skipIf(not_support_multi_device(), "no multi")
def test_load_state_dict_sharded_model_dict_same_axis(self):
devices = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3")
layer = Conv2d(3, 5, kernel_size=3)
layer.weight.shard_(devices, 3)
layer.bias.shard_(devices, None)
state_dict = {
'weight': Tensor.randn(5, 3, 3, 3).shard(devices, 3),
'bias': Tensor.randn(5).shard(devices, None),
}
load_state_dict(layer, state_dict)
self.assertEqual(layer.weight.device, devices)
self.assertEqual(layer.weight.uop.axis, 3)
self.assertEqual(layer.bias.device, devices)
self.assertEqual(layer.bias.uop.axis, None)
np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())
@unittest.skipIf(not_support_multi_device, "no multi")
def test_load_state_dict_sharded_model_dict_different_axis(self):
devices = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3")
devices5 = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3", f"{Device.DEFAULT}:4", f"{Device.DEFAULT}:5")
layer = Conv2d(3, 5, kernel_size=3)
layer.weight.shard_(devices, 3)
layer.bias.shard_(devices, None)
# different shard axis
state_dict = {
'weight': Tensor.randn(5, 3, 3, 3).shard(devices, None),
'bias': Tensor.randn(5).shard(devices5, 0),
}
load_state_dict(layer, state_dict)
# NOTE: model and state_dict shard differently, use the state_dict sharding # TODO: revisit this?
self.assertEqual(layer.weight.device, devices)
self.assertEqual(layer.weight.uop.axis, None)
self.assertEqual(layer.bias.device, devices5)
self.assertEqual(layer.bias.uop.axis, 0)
np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())
def test_load_state_dict_shape_mismatch(self):
d1, d2 = 2, 4
layer = Linear(d1, d1, bias=False)
state_dict = {'weight': Tensor.randn(d2, d2)}
with self.assertRaisesRegex(ValueError, r'Shape mismatch in layer `weight`: Expected shape \(2, 2\), but found \(4, 4\) in state dict.'):
load_state_dict(layer, state_dict)
def test_lstm_cell(self):
layer = LSTMCell(32, 16)
with torch.no_grad():
torch_layer = torch.nn.LSTMCell(32, 16)
layer.weight_hh.assign(torch_layer.weight_hh.numpy())
layer.weight_ih.assign(torch_layer.weight_ih.numpy())
layer.bias_hh.assign(torch_layer.bias_hh.numpy())
layer.bias_ih.assign(torch_layer.bias_ih.numpy())
inp = Tensor.randn(1, 32)
out_h, out_c = layer(inp)
torch_out_h, torch_out_c = torch_layer(torch.tensor(inp.numpy()))
np.testing.assert_allclose(out_h.numpy(), torch_out_h.numpy(), atol=1e-6)
np.testing.assert_allclose(out_c.numpy(), torch_out_c.numpy(), atol=1e-6)
out_h, out_c = layer(inp, (out_h, out_c))
torch_out_h, torch_out_c = torch_layer(torch.tensor(inp.numpy()), (torch_out_h, torch_out_c))
np.testing.assert_allclose(out_h.numpy(), torch_out_h.numpy(), atol=1e-6)
np.testing.assert_allclose(out_c.numpy(), torch_out_c.numpy(), atol=1e-6)
def test_lstm_cell_no_bias(self):
layer = LSTMCell(32, 16, bias=False)
inp = Tensor.randn(1, 32)
out_h, out_c = layer(inp)
out_h.realize()
out_c.realize()
h = Tensor.randn(1, 16)
c = Tensor.randn(1, 16)
out_h, out_c = layer(inp, (h, c))
out_h.realize()
out_c.realize()
assert layer.bias_hh is None
assert layer.bias_ih is None
if __name__ == '__main__':
unittest.main()

3429
artifacts/package_sources/tinygrad/test/backend/test_ops.py Normal file
View File

File diff suppressed because it is too large Load Diff

46
artifacts/package_sources/tinygrad/test/backend/test_opt_gemm.py Normal file
View File

@@ -0,0 +1,46 @@
import numpy as np
import unittest
from tinygrad import Tensor
from tinygrad.helpers import get_single_element
from tinygrad.codegen.opt import Opt, OptOps
from tinygrad.engine.realize import run_linear
from tinygrad.uop.ops import Ops, UOp
from test.helpers import replace_opts
class TestOptGemm(unittest.TestCase):
@classmethod
def setUpClass(cls):
N = 64
cls.a = Tensor.randn(N, N).contiguous().realize()
cls.b = Tensor.randn(N, N).contiguous().realize()
with np.errstate(all='ignore'):
cls.res = cls.a.T.numpy() @ cls.b.T.numpy()
def _test_gemm_unrolled_permute_l(self, opts=[]):
t = self.a.T @ self.b.T
# TODO: this should be a generic test helper
call = get_single_element(t.schedule_linear().src)
new_call = call.replace(src=(replace_opts(call.src[0], opts), *call.src[1:]))
run_linear(UOp(Ops.LINEAR, src=(new_call,)))
test = call.src[1].buffer.numpy().reshape(self.res.shape)
np.testing.assert_allclose(self.res, test, atol=1e-4)
def test_gemm_unrolled_permute_l_44(self):
opts = [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=4)]
self._test_gemm_unrolled_permute_l(opts)
def test_gemm_unrolled_permute_l_424(self):
# was failing with LLVM
opts = [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
self._test_gemm_unrolled_permute_l(opts)
def test_gemm_unrolled_permute_l_42(self):
opts = [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]
self._test_gemm_unrolled_permute_l(opts)
def test_gemm_unrolled_permute_l_22(self):
opts = [Opt(op=OptOps.UPCAST, axis=0, arg=2), Opt(op=OptOps.UPCAST, axis=1, arg=2)]
self._test_gemm_unrolled_permute_l(opts)
if __name__ == '__main__':
unittest.main()

189
artifacts/package_sources/tinygrad/test/backend/test_optim.py Normal file
View File

@@ -0,0 +1,189 @@
import numpy as np
import torch
import unittest
from tinygrad import Tensor, Device, dtypes
from tinygrad.nn.optim import Adam, SGD, AdamW, Muon, LAMB
from tinygrad.helpers import Context
from test.helpers import needs_second_gpu, slow
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)
def _param(tensor, val):
return tensor(val, requires_grad=True) if tensor is torch.tensor else tensor(val)
class TeenyNet:
def __init__(self, tensor):
self.x = _param(tensor, x_init.copy())
self.W = _param(tensor, W_init.copy())
def forward(self):
return (self.x * self.W).sum()
class TinyNet:
def __init__(self, tensor):
self.x = _param(tensor, x_init.copy())
self.W = _param(tensor, W_init.copy())
self.m = tensor(m_init.copy())
def forward(self):
out = self.x.matmul(self.W).relu()
# print(out.detach().numpy())
out = out.log_softmax(1)
out = out.mul(self.m).add(self.m).sum()
return out
def step(tensor, optim, steps=1, teeny=False, **kwargs):
net = TeenyNet(tensor) if teeny else TinyNet(tensor)
optim = optim([net.x, net.W], **kwargs)
for _ in range(steps):
out = net.forward()
optim.zero_grad()
out.backward()
optim.step()
return net.x.detach().numpy(), net.W.detach().numpy()
@slow
class TestOptim(unittest.TestCase):
def setUp(self): self.enterContext(Context(TRAINING=1))
def _test_optim(self, tinygrad_optim, torch_optim, steps, opts, atol, rtol):
for x,y in zip(step(Tensor, tinygrad_optim, steps, **opts),
step(torch.tensor, torch_optim, steps, **opts)):
np.testing.assert_allclose(x, y, atol=atol, rtol=rtol)
def _test_sgd(self, steps, opts, atol, rtol): self._test_optim(SGD, torch.optim.SGD, steps, opts, atol, rtol)
def _test_adam(self, steps, opts, atol, rtol): self._test_optim(Adam, torch.optim.Adam, steps, opts, atol, rtol)
def _test_adamw(self, steps, opts, atol, rtol): self._test_optim(AdamW, torch.optim.AdamW, steps, opts, atol, rtol)
def _test_muon(self, steps, opts, atol, rtol): self._test_optim(Muon, torch.optim.Muon, steps, opts, atol, rtol)
def test_multistep_sgd_high_lr_teeny(self): self._test_sgd(2, {'lr': 1.1, 'teeny': True}, 1e-6, 1e-5)
def test_multistep_adam_high_lr_teeny(self): self._test_adam(2, {'lr': 1.1, 'teeny': True}, 2e-4, 5e-4)
def test_multistep_muon_high_lr_teeny(self): self._test_muon(2, {'lr': 1.1, 'teeny': True}, 1e-2, 5e-4)
def test_sgd(self): self._test_sgd(1, {'lr': 0.001}, 1e-6, 0)
def test_sgd_high_lr(self): self._test_sgd(1, {'lr': 10}, 1e-6, 1e-5)
def test_sgd_wd(self): self._test_sgd(1, {'lr': 0.001, 'weight_decay': 0.1}, 1e-6, 0)
def test_sgd_high_lr_wd(self): self._test_sgd(1, {'lr': 10, 'weight_decay': 0.1}, 1e-6, 1e-5)
def test_multistep_sgd(self): self._test_sgd(10, {'lr': 0.001}, 1e-6, 0)
def test_multistep_sgd_high_lr(self): self._test_sgd(10, {'lr': 10}, 1e-6, 3e-4)
def test_multistep_sgd_wd(self): self._test_sgd(10, {'lr': 0.001, 'weight_decay': 0.1}, 1e-6, 0)
def test_multistep_sgd_high_lr_wd(self): self._test_sgd(10, {'lr': 9, 'weight_decay': 0.1}, 1e-6, 3e-4)
def test_multistep_sgd_momentum(self): self._test_sgd(10, {'lr': 0.001, 'momentum': 0.9}, 1e-6, 0)
def test_multistep_sgd_high_lr_momentum(self): self._test_sgd(10, {'lr': 10, 'momentum': 0.9}, 1e-5, 3e-4)
def test_multistep_sgd_momentum_wd(self): self._test_sgd(10, {'lr': 0.001, 'momentum': 0.9, 'weight_decay': 0.1}, 1e-6, 0)
def test_multistep_sgd_high_lr_momentum_wd(self): self._test_sgd(10, {'lr': 10, 'momentum': 0.9, 'weight_decay': 0.1}, 1e-5, 3e-4)
def test_multistep_sgd_nesterov_momentum(self): self._test_sgd(10, {'lr': 0.001, 'momentum': 0.9, 'nesterov': True}, 1e-5, 0)
def test_multistep_sgd_high_lr_nesterov_momentum(self): self._test_sgd(10, {'lr': 10, 'momentum': 0.9, 'nesterov': True}, 1e-5, 3e-4)
def test_multistep_sgd_nesterov_momentum_wd(self):
self._test_sgd(10, {'lr': 0.001, 'momentum': 0.9, 'nesterov': True, 'weight_decay': 0.1}, 1e-5, 0)
def test_multistep_sgd_high_lr_nesterov_momentum_wd(self):
self._test_sgd(10, {'lr': 9, 'momentum': 0.9, 'nesterov': True, 'weight_decay': 0.1}, 1e-5, 3e-4)
def test_muon(self): self._test_muon(1, {'lr': 0.001}, 1e-3, 0)
# TODO: disabled due to big atol
# def test_muon_high_lr(self): self._test_muon(1, {'lr': 10}, 1e-6, 3e-4)
def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 0.01}, 1e-3, 3e-4)
# TODO: disabled due to big atol
# def test_muon_high_lr_wd(self): self._test_muon(1, {'lr': 10, 'weight_decay': 0.01}, 1e-6, 5e-4)
# NOTE: momentum set to 0.95 by default, nesterov set to True by default
def test_multistep_muon_momentum_wd(self): self._test_muon(10, {'lr': 0.001, 'weight_decay': 0.01}, 3e-3, 0)
# ns defaults are numerically unstable, but it is tolerable in real training (see nsteps/nparam tests)
# TODO: disabled due to big atol
# def test_multistep_muon_high_lr_momentum_wd(self): self._test_muon(10, {'lr': 10, 'weight_decay': 0.01}, 1e-1, 3e-4)
def test_multistep_muon_no_nesterov_momentum(self): self._test_muon(10, {'lr': 0.001, 'nesterov': False}, 1e-3, 0)
# TODO: disabled due to big atol
# def test_multistep_muon_high_lr_no_nesterov_momentum(self): self._test_muon(10, {'lr': 10, 'nesterov': False}, 5e-2, 1e-1)
def test_muon_ns_steps(self): self._test_muon(1, {'lr': 0.001, 'ns_steps': 3}, 1e-4, 0)
# TODO: disabled due to big atol
# def test_muon_high_lr_ns_steps(self): self._test_muon(1, {'lr': 10, 'ns_steps': 3}, 1e-5, 3e-4)
def test_muon_ns_coefficients(self): self._test_muon(1, {'lr': 0.001,'ns_coefficients': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
# TODO: disabled due to big atol
# def test_muon_high_lr_ns_coefficients(self): self._test_muon(1, {'lr': 10,'ns_coefficients': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
def test_muon_momentum_wd_ns_steps_ns_coefficients(self):
self._test_muon(10, {'lr': 0.001, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_coefficients': (2.0,-1.5,0.5)}, 1e-4, 0)
# TODO: disabled due to big atol
# def test_multistep_muon_high_lr_momentum_wd_ns_steps_ns_coefficients(self):
# self._test_muon(10, {'lr': 10, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_coefficients': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
def test_adam(self): self._test_adam(1, {'lr': 0.001}, 1e-5, 0)
def test_adam_high_lr(self): self._test_adam(1, {'lr': 10}, 1e-4, 1e-4)
def test_adamw(self): self._test_adamw(1, {'lr': 0.001}, 1e-5, 0)
def test_adamw_high_lr(self): self._test_adamw(1, {'lr': 10}, 1e-4, 1e-4)
def test_multistep_adam(self): self._test_adam(10, {'lr': 0.001}, 1e-5, 0)
def test_multistep_adam_high_lr(self): self._test_adam(10, {'lr': 10}, 2e-3, 5e-4)
def test_multistep_adamw(self): self._test_adamw(10, {'lr': 0.001}, 1e-5, 0)
def test_multistep_adamw_high_lr(self): self._test_adamw(10, {'lr': 10}, 5e-4, 2e-3)
def test_duped_weights(self):
for Opt in [Adam, AdamW, SGD]:
losses = []
for i in range(2):
w = Tensor(x_init.copy())
opt = Opt([w], lr=0.1) if i == 0 else Opt([w, w], lr=0.1)
loss = None
for _ in range(3):
loss = w.sum()
opt.zero_grad()
loss.backward()
opt.step()
losses.append(loss.numpy())
np.testing.assert_allclose(losses[0], losses[1], atol=1e-4, rtol=0)
@unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
def test_mixed_precision(self):
self.enterContext(Context(DEFAULT_FLOAT=dtypes.half))
# weight update would overflow without upcasting
self._test_sgd(10, {'lr': 1e10}, 1e-6, 3e-4)
self._test_adam(1, {'lr': 1e10}, 1e-4, 1e-4)
self._test_adamw(1, {'lr': 1e10}, 1e-4, 1e-4)
def test_assert_tensor_train(self):
t = Tensor.ones((1,1))
optimizer = Adam([t])
optimizer.zero_grad()
t.sum().backward()
with Context(TRAINING=0):
self.assertRaises(RuntimeError, optimizer.step)
with Context(TRAINING=1):
optimizer.step()
def test_lamb_cpu_offload(self):
# test that LAMB works when optimizer params (m, v, b1_t, b2_t) are moved to CPU
t = Tensor(x_init.copy())
opt = LAMB([t])
# move optimizer state to CPU
for p in opt.m + opt.v + [opt.b1_t, opt.b2_t]: p.to_("CPU")
# run a step
t.sum().backward()
opt.step()
self.assertEqual(t.device, Device.DEFAULT)
self.assertEqual(opt.m[0].device, "CPU")
@needs_second_gpu
def test_lamb_cpu_offload_multi(self):
ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
t = Tensor(x_init.copy()).shard(ds, axis=1)
ds = t.device
opt = LAMB([t])
# move optimizer state to CPU
for p in opt.m + opt.v + [opt.b1_t, opt.b2_t]: p.to_("CPU")
# run a step
t.sum().backward()
opt.step()
self.assertEqual(t.device, ds)
self.assertEqual(opt.m[0].device, "CPU")
if __name__ == '__main__':
unittest.main()

196
artifacts/package_sources/tinygrad/test/backend/test_pickle.py Normal file
View File

@@ -0,0 +1,196 @@
import unittest, pickle, types, tracemalloc
import numpy as np
from tinygrad import Tensor, Device, TinyJit, Variable, dtypes
from tinygrad.helpers import GlobalCounters, ContextVar, Context, DEV
from tinygrad.uop.ops import PatternMatcher, UPat, UOp
class TestPickle(unittest.TestCase):
def test_pickle_code_object(self):
y = lambda x: x*2 # noqa: E731
code_str = pickle.dumps(y.__code__)
fxn = types.FunctionType(pickle.loads(code_str), globals())
self.assertEqual(fxn(2), 4)
def test_pickle_pattern_matcher(self):
pm = PatternMatcher([(UPat.cvar('x'), lambda x: x*2)])
sink = UOp.const(2)
tt = pm.rewrite(sink)
pm_str = pickle.dumps(pm)
pm2 = pickle.loads(pm_str)
self.assertEqual(pm2.rewrite(sink).key, tt.key)
def test_pickle_main_pattern_matcher(self):
from tinygrad.uop.symbolic import sym
ssym = pickle.dumps(sym)
dsym = pickle.loads(ssym)
self.assertEqual(dsym.patterns[0][0].location, sym.patterns[0][0].location)
def test_pickle_realized_tensor(self):
print("** init")
t = Tensor.rand(10, 10).realize()
st = pickle.dumps(t)
t_values = t.numpy()
del t # free buffers
print("** post pickle")
GlobalCounters.reset()
t2:Tensor = pickle.loads(st)
np.testing.assert_equal(t_values, t2.numpy())
# expect at most one COPY kernel
self.assertLessEqual(GlobalCounters.kernel_count, 1)
def test_pickle_realized_tensor_alt(self):
print("** init")
t = Tensor.rand(10, 10).to("CPU").realize()
st = pickle.dumps(t)
t_values = t.numpy()
del t # free buffers
print("** post pickle")
t2:Tensor = pickle.loads(st)
assert t2.uop.is_realized
np.testing.assert_equal(t_values, t2.numpy())
def test_pickle_realized_tensor_alt2(self):
print("** init")
t = Tensor.rand(10, 10).to("CPU").realize()
tensor_uop = t.uop
assert tensor_uop.is_realized, f"expected {tensor_uop} to be realized"
t_values = t.numpy()
# pickle
st = pickle.dumps(t)
# free buffers
del t
del tensor_uop
print("** post pickle")
t2:Tensor = pickle.loads(st)
assert t2.uop.is_realized, f"expected {t2.uop} to be realized"
np.testing.assert_equal(t_values, t2.numpy())
# NOTE: currently Buffer exists on the uop, not tensor
def test_pickle_buffer_uop(self):
t = Tensor.arange(4).clone().realize()
a = t.uop
assert a.is_realized
self.assertIsNotNone(buffer:=a.base.realized)
s = pickle.dumps(a)
# free buffers
del a
del buffer
a2:UOp = pickle.loads(s)
self.assertListEqual(a2.base.realized.as_memoryview().cast("I").tolist(), [0, 1, 2, 3])
@unittest.skipIf(DEV.interface.startswith("MOCK"), "mock device buffers live in host RAM, not VRAM")
def test_pickle_oob_ram(self):
N, M = 8, 10**6
ts = [Tensor.rand(M, dtype='float32').realize() for _ in range(N)]
tracemalloc.start()
st = pickle.dumps(ts, protocol=5, buffer_callback=lambda pb: pb.release())
self.assertLess(tracemalloc.get_traced_memory()[1], N*M*4)
tracemalloc.reset_peak()
def make_fake_buffers():
for _ in range(N):
Device[Device.DEFAULT].synchronize()
yield pickle.PickleBuffer(bytearray(M*4))
pickle.loads(st, buffers=make_fake_buffers())
self.assertLess(tracemalloc.get_traced_memory()[1], N*M*4)
tracemalloc.stop()
def test_pickle_unrealized_tensor(self):
t = Tensor.ones(10, 10)
st = pickle.dumps(t)
t2:Tensor = pickle.loads(st)
np.testing.assert_equal(t.numpy(), t2.numpy())
def test_pickle_variable(self):
v = Variable("i", 1, 20).bind(10)
t1 = Tensor.ones(10, v).contiguous()
t2 = Tensor.ones(10, v).contiguous()
ret = (t1+t2).sum(1)
st = pickle.dumps(ret)
del ret
vt2 = pickle.loads(st)
np.testing.assert_equal(vt2.numpy(), 20)
def test_pickle_buffer_view(self):
t = Tensor.arange(10).clone(device="CPU").realize()
vt = t[3:5].contiguous().realize()
assert hasattr(vt.uop.buffer, 'base')
ref_value = vt.tolist()
st = pickle.dumps(vt)
del t, vt
vt2 = pickle.loads(st)
assert hasattr(vt2.uop.buffer, 'base')
assert ref_value == vt2.tolist()
def test_pickle_numpy(self):
t = Tensor(np.array([1,2,3,4.]), dtype=dtypes.float32)
st = pickle.dumps(t)
t2:Tensor = pickle.loads(st)
np.testing.assert_equal(t.numpy(), t2.numpy())
def test_pickle_jit(self):
@TinyJit
def add(a, b): return a.sum()+b+1
for _ in range(3): add(Tensor.rand(10, 10), Tensor.rand(10, 10))
st = pickle.dumps(add)
del add
add_fxn = pickle.loads(st)
x = Tensor.ones(10, 10).contiguous().realize()
y = Tensor.ones(10, 10).contiguous().realize()
print("post jit")
out = add_fxn(x, y)
np.testing.assert_equal(out.numpy(), 102)
def test_pickle_jit_no_del(self):
@TinyJit
def fn(x): return x + 1.0
for _ in range(3): fn(Tensor.randn(4))
loaded = pickle.loads(pickle.dumps(fn))
self.assertEqual(loaded(Tensor([1.0,2.0,3.0,4.0])).tolist(), [2.0,3.0,4.0,5.0])
def test_pickle_context_var(self):
v = ContextVar("test_var", 0)
with Context(test_var=1):
vs = pickle.dumps(v)
v2 = pickle.loads(vs)
self.assertEqual(v2.value, 1)
def test_pickle_schedule(self):
a = Tensor([1,2])
out = a + 2
sched = out.schedule_linear()
pk = pickle.dumps(sched)
sched_pk = pickle.loads(pk)
self.assertEqual(sched_pk.src[-1].src[0], sched.src[-1].src[0])
def test_pickle_renderer(self):
from tinygrad.device import Device
pk = pickle.dumps(Device.default.renderer)
pickle.loads(pk)
class TestPickleJIT(unittest.TestCase):
@classmethod
def setUpClass(cls):
N = 10
@TinyJit
def add(a, b): return a.sum()+b+1
for _ in range(3): add(Tensor.rand(N, N), Tensor.rand(N, N))
cls.st = pickle.dumps(add)
del add
def test_inspect(self):
import io
class FakeClass:
def __init__(self, *args, **kwargs):
print(self.module, self.name)
class InspectUnpickler(pickle.Unpickler):
def find_class(self, module, name): return type("SpecializedFakeClass", (FakeClass,), {"name": name, "module": module})
InspectUnpickler(io.BytesIO(self.st)).load()
@unittest.skip("we are still saving intermediate buffers")
def test_size(self):
# confirm no intermediate buffers are saved
self.assertLess(len(self.st), 1_000_000)
if __name__ == '__main__':
unittest.main()

227
artifacts/package_sources/tinygrad/test/backend/test_profiler.py Normal file
View File

@@ -0,0 +1,227 @@
import unittest, struct, contextlib, statistics, gc
from tinygrad import Device, Tensor, dtypes, TinyJit
from tinygrad.helpers import DEV, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileDeviceEvent, ProfileGraphEvent
from tinygrad.runtime.support.hcq import HCQCompiled
from tinygrad.engine.realize import get_runtime
from tinygrad.codegen import to_program
MOCKGPU = DEV.interface.startswith("MOCK")
def _dev_base(d):
p = d.split(":")
return p[0] if len(p) < 2 or not p[1].isdigit() else f"{p[0]}:{p[1]}"
@contextlib.contextmanager
def helper_collect_profile(*devs):
for dev in devs: dev.synchronize()
saved = [x for x in Compiled.profile_events if isinstance(x, ProfileDeviceEvent) and x.device.startswith("METAL")]
Compiled.profile_events.clear()
for x in saved: Compiled.profile_events.append(x)
cpu_events.clear()
profile_list = []
with Context(PROFILE=1):
yield profile_list
for dev in devs: dev.synchronize()
for dev in devs: dev._at_profile_finalize()
for x in Compiled.profile_events: profile_list.append(x)
profile_list.extend(cpu_events)
def helper_profile_filter_device(profile, device:str):
assert any(getattr(x, "device", None) == device and isinstance(x, ProfileDeviceEvent) for x in profile), f"device {device} is not registred"
dev_events = [x for x in profile if getattr(x, "device", None) == device and isinstance(x, ProfileDeviceEvent)]
assert len(dev_events) == 1, "only one device registration event is expected"
return [x for x in profile if getattr(x, "device", None) == device], dev_events[0]
# TODO: support in HCQCompiled
is_cpu_hcq = Device.DEFAULT in {"CPU"}
@unittest.skipUnless((issubclass(type(Device[Device.DEFAULT]), HCQCompiled) and not is_cpu_hcq) or Device.DEFAULT in {"METAL"}, "Dev not supported")
class TestProfiler(unittest.TestCase):
@classmethod
def setUpClass(self):
TestProfiler.d0 = Device[Device.DEFAULT]
TestProfiler.a = Tensor([0.,1.], device=Device.DEFAULT).realize()
TestProfiler.b = self.a + 1
si = self.b.schedule_linear().src[-1]
TestProfiler.prg = to_program(si.src[0], TestProfiler.d0.renderer)
TestProfiler.runtime = get_runtime(TestProfiler.d0.device, TestProfiler.prg)
TestProfiler.b.uop.buffer.allocate()
def test_profile_kernel_run(self, wait=False):
runner_name = TestProfiler.runtime.name
with helper_collect_profile(TestProfiler.d0) as profile:
gs, ls = TestProfiler.prg.arg.launch_dims({})
TestProfiler.runtime(TestProfiler.b.uop.buffer._buf, TestProfiler.a.uop.buffer._buf, global_size=gs, local_size=ls, wait=wait)
profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent)]
assert len(kernel_runs) == 1, "one kernel run is expected"
assert kernel_runs[0].name == runner_name, "kernel name is not correct"
assert _dev_base(kernel_runs[0].device) == kernel_runs[0].device, "kernel should not be on a sub-device"
def test_profile_kernel_run_wait(self):
self.test_profile_kernel_run(wait=True)
def test_profile_copyin(self):
buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
with helper_collect_profile(TestProfiler.d0) as profile:
buf1.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 1)))))
kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith((TestProfiler.d0.device, "PYTHON"))]
assert len(kernel_runs) == 1, "one kernel run is expected"
def test_profile_multiops(self):
runner_name = TestProfiler.runtime.name
buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
with helper_collect_profile(TestProfiler.d0) as profile:
buf1.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 1)))))
gs, ls = TestProfiler.prg.arg.launch_dims({})
TestProfiler.runtime(buf1._buf, TestProfiler.a.uop.buffer._buf, global_size=gs, local_size=ls)
buf1.as_memoryview()
evs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith((TestProfiler.d0.device, "PYTHON"))]
assert len(evs) == 3, "3 kernel runs are expected"
# NOTE: order of events does not matter, the tool is responsible for sorting them
prg_events = [e for e in evs if e.device == TestProfiler.d0.device]
assert any(e.name == runner_name for e in prg_events), "kernel name is not correct"
#for i in range(1, 3):
# assert evs[i].st > evs[i-1].en, "timestamp not aranged"
def test_profile_multidev(self):
try: d1 = Device[f"{Device.DEFAULT}:1"]
except Exception as e: self.skipTest(f"second device not available {e}")
buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(f"{Device.DEFAULT}:1", 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
with helper_collect_profile(TestProfiler.d0, d1) as profile:
buf1.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 1)))))
buf2.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 1)))))
for dev in [TestProfiler.d0.device, d1.device]:
evs = [x for x in profile if isinstance(x, ProfileRangeEvent) and _dev_base(x.device) == dev]
assert len(evs) == (0 if hasattr(TestProfiler.d0.allocator, '_as_buffer') else 1), "one kernel runs are expected"
def test_profile_multidev_transfer(self):
try: d1 = Device[f"{Device.DEFAULT}:1"]
except Exception as e: self.skipTest(f"second device not available {e}")
buf1 = Tensor.randn(10, 10, device=f"{Device.DEFAULT}:0").realize()
with helper_collect_profile(TestProfiler.d0, d1) as profile:
buf1.to(f"{Device.DEFAULT}:1").realize()
kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith(TestProfiler.d0.device)]
assert len(kernel_runs) == 1, "one kernel run is expected"
@unittest.skipIf(Device.DEFAULT in "METAL" or (MOCKGPU and Device.DEFAULT == "AMD"), "AMD mockgpu does not support queue wait interrupts")
def test_profile_graph(self):
try: d1 = Device[f"{Device.DEFAULT}:1"]
except Exception as e: self.skipTest(f"second device not available {e}")
def f(a):
x = (a + 1).realize()
return x, x.to(d1.device).realize()
a = Tensor.randn(10, 10, device=TestProfiler.d0.device).realize()
with helper_collect_profile(TestProfiler.d0, d1) as profile:
jf = TinyJit(f)
for _ in range(3): jf(a)
del jf
graph_evs = [x for x in profile if isinstance(x, ProfileGraphEvent)]
_, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
_, _ = helper_profile_filter_device(profile, d1.device)
assert len(graph_evs) == 2, "2 graph events are expected"
assert len(graph_evs[0].ents) == 2, "two entities are expected"
@unittest.skipIf(MOCKGPU, "skip MOCKGPU")
@unittest.skipUnless(issubclass(type(Device[Device.DEFAULT]), HCQCompiled), "must be HCQ")
def test_dev_jitter_matrix(self):
dev_cnt = 6
try: devs = [Device[f"{Device.DEFAULT}:{i}"] for i in range(dev_cnt)]
except Exception as e: self.skipTest(f"multiple devices not available {e}")
for dev in devs: dev.synchronize()
for dev in devs: dev._at_profile_finalize()
def _sync_d2d(d1:HCQCompiled, d2:HCQCompiled):
d1.hw_compute_queue_t().signal(d1.timeline_signal, d1.timeline_value).wait(d2.timeline_signal, d2.timeline_value) \
.timestamp(d1.timeline_signal).signal(d1.timeline_signal, d1.timeline_value+1).submit(d1)
d2.hw_compute_queue_t().signal(d2.timeline_signal, d2.timeline_value).wait(d1.timeline_signal, d1.timeline_value) \
.timestamp(d2.timeline_signal).signal(d2.timeline_signal, d2.timeline_value+1).submit(d2)
d1.timeline_value += 2
d2.timeline_value += 2
d1.timeline_signal.wait(d1.timeline_value - 1)
d2.timeline_signal.wait(d2.timeline_value - 1)
return d2.timeline_signal.timestamp - d1.timeline_signal.timestamp
# then test it by timing the GPU to GPU times
dev_evs = {x.device:x for x in Compiled.profile_events if isinstance(x, ProfileDeviceEvent)}
jitter_matrix = [[float('nan')] * len(devs) for _ in range(len(devs))]
pairs = [(p1, p2) for p1 in enumerate(devs) for p2 in enumerate(devs) if p1 != p2]
for (i1, d1), (i2, d2) in pairs:
cpu_diff = dev_evs[d1.device].tdiff - dev_evs[d2.device].tdiff
jitter_matrix[i1][i2] = statistics.median(_sync_d2d(d1, d2) - _sync_d2d(d2, d1) for _ in range(20)) / 2 - cpu_diff
print("pairwise clock jitter matrix (us):\n" + '\n'.join([''.join([f'{float(item):8.3f}' for item in row]) for row in jitter_matrix]))
for (i1, d1), (i2, d2) in pairs:
assert abs(jitter_matrix[i1][i2]) < 0.5, "jitter should be less than 0.5us"
def test_cpu_profile(self):
def test_fxn(err=False):
if err: raise Exception()
with helper_collect_profile(dev:=TestProfiler.d0) as profile:
with cpu_profile("test_1", dev):
test_fxn(err=False)
with self.assertRaises(Exception):
with cpu_profile("test_2", dev):
test_fxn(err=True)
range_events = [p for p in profile if isinstance(p, ProfileRangeEvent) and p.device == dev]
self.assertEqual(len(range_events), 2)
@unittest.skip("this test is flaky")
@unittest.skipUnless(Device[Device.DEFAULT].graph is not None, "graph support required")
def test_graph(self):
from test.backend.test_graph import helper_alloc_rawbuffer, helper_exec_op, helper_test_graphs
device = TestProfiler.d0.device
bufs = [helper_alloc_rawbuffer(device, fill=True) for _ in range(5)]
graphs = [[helper_exec_op(device, bufs[0], [bufs[1], bufs[2]]), helper_exec_op(device, bufs[0], [bufs[3], bufs[4]]),]]
with helper_collect_profile(dev:=TestProfiler.d0) as profile:
helper_test_graphs(dev.graph, graphs, runs:=2)
# NOTE: explicitly trigger deletion of all graphs
graphs.clear()
gc.collect()
graphs = [e for e in profile if isinstance(e, ProfileGraphEvent)]
self.assertEqual(len(graphs), runs)
for ge in graphs:
self.assertEqual(len(ge.ents), len(graphs))
@unittest.skip("this test is flaky")
def test_trace_metadata(self):
with Context(TRACEMETA=1):
a = Tensor.empty(1)+2
b = Tensor.empty(1)+2
with helper_collect_profile(TestProfiler.d0) as profile:
Tensor.realize(a, b)
profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
exec_points = [e for e in profile if isinstance(e, ProfilePointEvent) and e.name == "exec"]
range_events = [e for e in profile if isinstance(e, ProfileRangeEvent) and _dev_base(e.device) == e.device]
self.assertEqual(len(exec_points), len(range_events), 2)
self.assertEqual(len(dedup(e.arg['name'] for e in exec_points)), 1)
self.assertEqual(len(dedup(e.arg['metadata'] for e in exec_points)), 1)
if __name__ == "__main__":
unittest.main()

239
artifacts/package_sources/tinygrad/test/backend/test_quantize_onnx.py Normal file
View File

@@ -0,0 +1,239 @@
# ruff: noqa: E501
import numpy as np
import tempfile, unittest
from tinygrad import Tensor, Context, Device, dtypes, UOp
from tinygrad.uop.ops import Ops
from tinygrad.dtype import AddrSpace
from tinygrad.codegen.opt import Opt, OptOps
from tinygrad.engine.realize import run_linear
from tinygrad.codegen import to_program
from test.helpers import replace_opts
N = 512
def create_gemm_model(model_path:str, batch_size=N, in_size=N, out_size=N, bias=False):
import onnx
from onnx import helper, numpy_helper, TensorProto
# Define input and output
input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT, [batch_size, in_size])
output_tensor = helper.make_tensor_value_info("output", TensorProto.FLOAT, [batch_size, out_size])
# Create random weights and bias
W_data = np.random.randn(in_size, out_size).astype(np.float32)
W_init = numpy_helper.from_array(W_data, name="W")
if bias:
B_data = np.random.randn(out_size).astype(np.float32)
B_init = numpy_helper.from_array(B_data, name="B")
gemm_node = helper.make_node("Gemm", inputs=["input", "W", "B"], outputs=["output"], alpha=1.0, beta=1.0, transB=0)
graph_def = helper.make_graph([gemm_node], "SingleGemmGraph", [input_tensor], [output_tensor], initializer=[W_init, B_init])
else:
gemm_node = helper.make_node("Gemm", inputs=["input", "W"], outputs=["output"], alpha=1.0, beta=1.0, transB=0)
graph_def = helper.make_graph([gemm_node], "SingleGemmGraph", [input_tensor], [output_tensor], initializer=[W_init])
# Create and save the model
#model_def = helper.make_model(graph_def, producer_name="single_gemm_example")
# TODO remove this once ORT supports 1.18.0
model_def = helper.make_model(graph_def, producer_name="single_gemm_example", ir_version=10, opset_imports=[helper.make_opsetid("", 22)])
onnx.save_model(model_def, model_path)
return model_path
def sexec(out:Tensor, opts:list[Opt], replace_src=None, run_count=3):
linear = out.schedule_linear()
call = linear.src[-1]
prg = to_program(replace_opts(call.src[0], opts), renderer=Device[Device.DEFAULT].renderer)
if replace_src is not None:
old_name = prg.src[2].arg.split("__attribute__((noinline)) void ")[1].split("(")[0]
new_src = replace_src + "/* DSP boilerplate */" + prg.src[2].arg.split("/* DSP boilerplate */")[1].replace(old_name, "fxn")
# drop BINARY and replace SOURCE so run_linear recompiles
prg = prg.replace(src=prg.src[:2] + (UOp(Ops.SOURCE, arg=new_src),))
linear = linear.replace(src=linear.src[:-1] + (call.replace(src=(prg, *call.src[1:])),))
for _ in range(run_count): run_linear(linear)
def get_quantized_model(sz):
from onnxruntime.quantization import quantize_static, QuantFormat, QuantType, CalibrationDataReader
class FakeDataReader(CalibrationDataReader):
def __init__(self): self.cnt = 0
def get_next(self) -> dict:
self.cnt += 1
if self.cnt == 100: return None
return {"input": np.random.uniform(size=(sz, sz)).astype(np.float32)}
tmpdir = tempfile.mkdtemp()
out_file = f"{tmpdir}/test_out.onnx"
quantize_static(create_gemm_model(f"{tmpdir}/test_in.onnx", sz, sz, sz), out_file,
FakeDataReader(), quant_format=QuantFormat.QDQ, per_channel=False, reduce_range=False,
activation_type=QuantType.QUInt8, weight_type=QuantType.QInt8,
extra_options={"ActivationSymmetric": False})
return out_file
@unittest.skip("this is broken")
@unittest.skipIf(Device.DEFAULT != "CPU", "only tests for CPU")
class TestQuantizeOnnxCPU(unittest.TestCase):
def test_quant_128(self, sz=128):
try:
import onnx # noqa: F401 # pylint: disable=unused-import
except ImportError:
raise unittest.SkipTest()
from tinygrad.nn.onnx import OnnxRunner
out_file = get_quantized_model(sz)
run_onnx = OnnxRunner(out_file)
inp = Tensor(np.random.uniform(size=(sz, sz)).astype(np.float32))
with Context(QUANTIZE=1):
linear = run_onnx({"input":inp})["output"].schedule_linear()
prg = to_program(linear.src[-2].src[0], renderer=Device[Device.DEFAULT].renderer)
daccs = [u for u in tuple(prg.src[1].src) if u.op is Ops.BUFFER and u.addrspace is AddrSpace.REG]
assert all(u.dtype.scalar() is dtypes.int for u in daccs)
@unittest.skipIf(Device.DEFAULT != "DSP", "only tests for DSP")
class TestQuantizeOnnx(unittest.TestCase):
def test_quant_128(self): self.test_quant(128)
def test_quant(self, sz=512):
from examples.benchmark_onnx import load_onnx_model
# divide is ~1500-2000 without reduce_range, 750-900 with it
out_file = get_quantized_model(sz)
run_onnx_jit, _ = load_onnx_model(out_file)
run_onnx_jit(input=Tensor(np.random.uniform(size=(sz, sz)).astype(np.float32)))
def test_prequant_conv2d_1x1(self):
X = Tensor(np.random.uniform(0, 255, size=(1, 32, 128, 128)).astype(np.uint8))
W = Tensor(np.random.uniform(0, 255, size=(64, 32, 1, 1)).astype(np.uint8))
out = X.conv2d(W, dtype=X.dtype)
opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
sexec(out, opts)
def test_prequant_gemm(self):
N = 512
X = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(np.uint8))
W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(np.uint8))
out = X.matmul(W, dtype=X.dtype)
opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
sexec(out, opts)
# TODO: this has to work
def test_prequant_gemm_intacc_early(self, xi=np.int8, wi=np.int8):
N = 512
X = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(xi))
W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(wi))
# this divide is interesting and forces the accumulator to actually be an int
out = (X.cast("int").matmul(W.cast("int"))//1000).cast("int8")
opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
sexec(out, opts)
def test_prequant_gemm_handcode(self):
src = """typedef int int128 __attribute__((aligned(512),vector_size(512)));
typedef int int32 __attribute__((aligned(128),vector_size(128)));
typedef int int64 __attribute__((aligned(256),vector_size(256)));
typedef unsigned char unsigned_char4 __attribute__((aligned(4),vector_size(4)));
typedef signed char signed_char128 __attribute__((aligned(128),vector_size(128)));
typedef unsigned char unsigned_char128 __attribute__((aligned(128),vector_size(128)));
typedef unsigned char unsigned_char256 __attribute__((aligned(256),vector_size(256)));
union V256 {
unsigned_char256 vec256;
struct {
unsigned_char128 lo128;
unsigned_char128 hi128;
};
};
__attribute__((noinline)) void fxn(unsigned char* restrict __attribute__((align_value(128))) data0,
unsigned char* restrict __attribute__((align_value(128))) data1,
signed char* restrict __attribute__((align_value(128))) data2) {
for (int ridx0 = 0; ridx0 < 512; ridx0++) {
int alu0 = (ridx0<<9);
for (int ridx1 = 0; ridx1 < 4; ridx1++) {
int alu1 = (ridx1<<7);
int32 acc0 = __builtin_HEXAGON_V6_vd0_128B();
int32 acc1 = __builtin_HEXAGON_V6_vd0_128B();
int32 acc2 = __builtin_HEXAGON_V6_vd0_128B();
int32 acc3 = __builtin_HEXAGON_V6_vd0_128B();
for (int ridx2 = 0; ridx2 < 128; ridx2++) {
unsigned_char4 val0 = *((unsigned_char4*)((data1+(alu0+(ridx2<<2)))));
int alu2 = (alu1+(ridx2<<11));
signed_char128 x0 = *((signed_char128*)((data2+alu2)));
signed_char128 x1 = *((signed_char128*)((data2+(alu2+512))));
signed_char128 x2 = *((signed_char128*)((data2+(alu2+1024))));
signed_char128 x3 = *((signed_char128*)((data2+(alu2+1536))));
union V256 ss01;
// ss01.lo128 = (x0[0], x1[0], x0[2], x1[2], x0[4], x1[4], ...)
// ss01.hi128 = (x0[1], x1[1], x0[3], x1[3], x0[5], x1[5], ...)
ss01.vec256 = __builtin_HEXAGON_V6_vshufoeb_128B(x1, x0);
union V256 ss23;
// ss23.lo128 = (x2[0], x3[0], x2[2], x3[2], x2[4], x3[4], ...)
// ss23.hi128 = (x2[1], x3[1], x2[3], x3[3], x2[5], x3[5], ...)
ss23.vec256 = __builtin_HEXAGON_V6_vshufoeb_128B(x3, x2);
union V256 sslo;
// sslo.lo128 = (x0[0], x1[0], x2[0], x3[0], x0[4], x1[4], ...)
// sslo.hi128 = (x0[2], x1[2], x2[2], x3[2], x0[6], x1[6], ...)
sslo.vec256 = __builtin_HEXAGON_V6_vdealvdd_128B(ss23.lo128, ss01.lo128, 2);
union V256 sshi;
// sshi.lo128 = (x0[1], x1[1], x2[1], x3[1], x0[5], x1[5], ...)
// sshi.hi128 = (x0[3], x1[3], x2[3], x3[3], x0[7], x1[7], ...)
sshi.vec256 = __builtin_HEXAGON_V6_vdealvdd_128B(ss23.hi128, ss01.hi128, 2);
//unsigned_char128 w0 = (unsigned_char128){val0[0],val0[1],val0[2],val0[3],val0[0],val0[1],val0[2],val0[3],...
unsigned_char128 w0 = __builtin_HEXAGON_V6_lvsplatw_128B(*((unsigned int*)&val0));
acc0 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc0, w0, sslo.lo128);
acc1 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc1, w0, sshi.lo128);
acc2 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc2, w0, sslo.hi128);
acc3 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc3, w0, sshi.hi128);
}
acc0 /= 1000;
acc1 /= 1000;
acc2 /= 1000;
acc3 /= 1000;
// ','.join([f"acc{j}[{i}]" for i in range(32) for j in range(4)])
// acc0[0], acc0[1], acc0[2], ..... acc3[30], acc3[31]
unsigned_char128 packed = __builtin_HEXAGON_V6_vpackhub_sat_128B(__builtin_HEXAGON_V6_vpackwh_sat_128B(acc3, acc2),
__builtin_HEXAGON_V6_vpackwh_sat_128B(acc1, acc0));
packed = __builtin_HEXAGON_V6_vshuffb_128B(packed);
packed = __builtin_HEXAGON_V6_vshuffb_128B(packed);
// acc0[0], acc1[0], acc2[0], ..... acc2[31], acc3[31]
*((unsigned_char128*)((data0+(alu0+alu1)))) = packed;
}
}
}"""
self.test_prequant_gemm_intacc(np.uint8, np.int8, src)
def test_prequant_gemm_intacc_32(self):
opts = [Opt(op=OptOps.UPCAST, axis=1, arg=0), Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UNROLL, axis=0, arg=0)]
self.test_prequant_gemm_intacc(np.uint8, np.int8, N=32, opts=opts)
def test_prequant_gemm_intacc_128(self): self.test_prequant_gemm_intacc(np.uint8, np.int8, N=128)
def test_prequant_gemm_intacc_256(self): self.test_prequant_gemm_intacc(np.uint8, np.int8, N=256)
def test_prequant_gemm_intacc(self, xi=np.uint8, wi=np.uint8, replace_src=None, N=512, clip=True, opts=None):
X = Tensor(m1:=(np.random.uniform(0, 255, size=(N,N)).astype(xi))).realize()
W = Tensor(m2:=(np.random.uniform(0, 255, size=(N,N)).astype(wi))).realize()
tg_dtype = dtypes.int8 if xi == np.int8 else dtypes.uint8
out = (X.int().matmul(W.int())//1000)
if clip: out = out.clip(tg_dtype.min, tg_dtype.max)
out = out.cast(tg_dtype)
opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)] if opts is None else opts
sexec(out, opts, replace_src, run_count=1)
tout = out.numpy()
mout = ((m1.astype(np.int32) @ m2.astype(np.int32)) // 1000)
if clip: mout = mout.clip(tg_dtype.min, tg_dtype.max)
mout = mout.astype(xi)
print(tout)
print(mout)
np.testing.assert_equal(tout, mout)
def test_prequant_gemm_intacc_wi(self): self.test_prequant_gemm_intacc(wi=np.int8)
def test_prequant_gemm_intacc_xiwi(self): self.test_prequant_gemm_intacc(xi=np.int8, wi=np.int8)
def test_prequant_gemm_intacc_xiwi_noclip(self): self.test_prequant_gemm_intacc(xi=np.int8, wi=np.int8, clip=False)
def test_prequant_gemv(self):
N = 2048
X = Tensor(np.random.uniform(0, 255, size=(1,N)).astype(np.uint8)).realize()
W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(np.uint8)).realize()
#out = X.cast(dtypes.int) @ W.cast(dtypes.int)
#out = X @ W
out = X.matmul(W, dtype=X.dtype)
opts = [Opt(op=OptOps.UPCAST, axis=0, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
sexec(out, opts)
if __name__ == "__main__":
unittest.main()

260
artifacts/package_sources/tinygrad/test/backend/test_randomness.py Normal file
View File

@@ -0,0 +1,260 @@
import unittest, math
from tinygrad import dtypes, Tensor, Device
from tinygrad.helpers import getenv, DEV, Context
from tinygrad.codegen import to_program
from tinygrad.uop.ops import Ops
from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.nir import NIRRenderer
from tinygrad.renderer.isa.x86 import X86Renderer
from test.helpers import not_support_multi_device, needs_second_gpu
from test.unit.test_randomness import equal_distribution, normal_test
import numpy as np
import torch
from hypothesis import given, settings, strategies as strat
settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
settings.load_profile("my_profile")
class TestRandomness(unittest.TestCase):
def test_rand(self):
self.assertFalse(normal_test(Tensor.rand))
self.assertTrue(equal_distribution(Tensor.rand, torch.rand, lambda x: np.random.rand(*x)))
def test_rand_is_lazy(self):
Tensor.manual_seed(0)
r1 = Tensor.rand(10)
self.assertFalse(r1.uop.is_realized, "rand should be lazy - tensor should not be realized")
counter = Tensor._device_rng_counters[Device.DEFAULT]
self.assertFalse(counter.uop.is_realized, "rand should be lazy - counter should not be realized")
# second rand triggers assign path
r2 = Tensor.rand(10)
self.assertFalse(r2.uop.is_realized, "rand should be lazy - tensor should not be realized after second rand")
self.assertFalse(counter.uop.is_realized, "rand should be lazy - counter should not be realized after second rand")
Tensor.realize(r1, r2)
self.assertTrue(r1.uop.is_realized, "tensor should be realized after .realize()")
self.assertTrue(r2.uop.is_realized, "tensor should be realized after .realize()")
@unittest.skipUnless(dtypes.float16 in Device[Device.DEFAULT].renderer.supported_dtypes(), "need float16 support")
def test_rand_float16(self):
N = 128
x = Tensor.rand((2, N, N), dtype=dtypes.float16)
assert x.dtype == dtypes.float16
nx = x.numpy()
# seed dependant, check output range is [0, 1)
assert nx[nx == 1].size == 0
assert nx[nx == 0].size > 0
equal_distribution(lambda *x: Tensor.rand(*x, dtype=dtypes.float16), torch.rand, lambda x: np.random.rand(*x), shape=(2, N, N))
@unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}, "gpuocelot doesn't support certain ops needed for threefry")
def test_threefry_against_reference(self):
Tensor.manual_seed(1337)
# reference generated using
"""
key0 = 1337
key1 = 0
values = jax.extend.random.threefry_2x32((np.uint32(key1), np.uint32(key0)), np.arange(20, dtype=np.uint32))
print(f"[{', '.join(f'{v}' for v in values)}]")
"""
jr = np.array([2221762175, 1752107825, 653745012, 1967534793, 1395205442, 3840423848, 2159346757,
603508235, 3319473678, 3363866483, 3544324138, 1436466838, 2169858556, 2570072943,
2387150698, 3678370550, 2911697663, 403244401, 2560861638, 1692360114])
counts = Tensor.arange(20, dtype=dtypes.uint32)
counts0, counts1 = counts.chunk(2)
r = Tensor._threefry_random_bits(Tensor([0, 1337], dtype='uint32'), counts0, counts1).numpy()
np.testing.assert_allclose(jr, r)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "PTX and NIR use pointer arithmetic")
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "X86 callee saved registers have ulong dtype")
def test_threefry_doesnt_use_long(self):
linear = Tensor.rand(20).schedule_linear()
for call in linear.src:
ast = call.src[0]
if ast.op is Ops.SINK:
prg = to_program(ast, renderer=Device[Device.DEFAULT].renderer)
for u in tuple(prg.src[1].src):
self.assertNotIn(u.dtype, {dtypes.long, dtypes.ulong}, msg=f"long found in {prg.arg.name}")
def test_threefry_against_reference_full(self):
Tensor.manual_seed(1337)
# reference generated using
"""
key0 = 1337
key1 = int.from_bytes(hashlib.sha256(int(0).to_bytes(4)).digest(), "big") & 0xffffffff
# derive new key for the counter offset (c_low=0, c_high=0 for first call)
new_key_values = jax.extend.random.threefry_2x32((np.uint32(key1), np.uint32(key0)), np.array([0, 0], dtype=np.uint32))
new_key = (np.uint32(new_key_values[0]), np.uint32(new_key_values[1]))
values = jax.extend.random.threefry_2x32(new_key, np.arange(20, dtype=np.uint32))
values = (values >> (32 - 23)) | np.array(1, dtype=np.float32).view(np.uint32)
values = values.view(np.float32) - 1
print(f"[{', '.join(f'{v}' for v in values)}]")
"""
jr = np.array([0.45735931396484375, 0.6311527490615845, 0.15571284294128418, 0.8149417638778687, 0.7862188816070557,
0.8008807897567749, 0.568588376045227, 0.9852620363235474, 0.42314577102661133, 0.9811755418777466,
0.38059568405151367, 0.09186363220214844, 0.9497315883636475, 0.5826880931854248, 0.3796330690383911,
0.5610522031784058, 0.16122901439666748, 0.3732343912124634, 0.9795231819152832, 0.3280656337738037], dtype=np.float32)
r = Tensor.rand(20).numpy()
np.testing.assert_allclose(r, jr, atol=1e-5, rtol=1e-5)
# next 20 (c_low=20, c_high=0)
jr = np.array([0.09199333190917969, 0.9130761623382568, 0.7048608064651489, 0.22254979610443115, 0.0014830827713012695,
0.37023448944091797, 0.7790107727050781, 0.7484984397888184, 0.7524604797363281, 0.19875383377075195,
0.48537540435791016, 0.10002851486206055, 0.5369305610656738, 0.3294715881347656, 0.5246957540512085,
0.7659651041030884, 0.7949080467224121, 0.34988296031951904, 0.9798505306243896, 0.2599533796310425], dtype=np.float32)
r = Tensor.rand(20).numpy()
np.testing.assert_allclose(r, jr, atol=1e-5, rtol=1e-5)
# next 10 (c_low=40, c_high=0)
jr = np.array([0.3198714256286621, 0.7984923124313354, 0.320881724357605, 0.4716068506240845, 0.7323365211486816,
0.9663800001144409, 0.13873648643493652, 0.16062307357788086, 0.49300849437713623, 0.10077548027038574], dtype=np.float32)
r = Tensor.rand(10).numpy()
np.testing.assert_allclose(r, jr, atol=1e-5, rtol=1e-5)
@needs_second_gpu
@unittest.skipIf(not_support_multi_device(), "no multi")
def test_threefry_tensors_cnt(self):
Tensor.manual_seed(1337)
Tensor.rand(20).realize()
assert len(Tensor._device_rng_counters) == 1
assert len(Tensor._device_seeds) == 1
Tensor.rand(20, device=f"{Device.DEFAULT}:1").realize()
assert len(Tensor._device_rng_counters) == 2
assert len(Tensor._device_seeds) == 2
Tensor.manual_seed(2)
assert len(Tensor._device_rng_counters) == 0
assert len(Tensor._device_seeds) == 0
@needs_second_gpu
@unittest.skipIf(not_support_multi_device(), "no multi")
def test_threefry_same_kernels(self):
Tensor.manual_seed(0)
Tensor.rand(1).realize()
s = Tensor.rand(20).schedule_linear().src
s2 = Tensor.rand(20).schedule_linear().src
assert len(s) == len(s2), f"{len(s)} != {len(s2)}"
for x,y in zip(s, s2):
if not (x.src[0] == y.src[0]):
print(f"{x.src[0]} != {y.src[0]}")
Tensor.rand(1, device=f"{Device.DEFAULT}:1").realize()
s3 = Tensor.rand(20, device=f"{Device.DEFAULT}:1").schedule_linear().src
s4 = Tensor.rand(20, device=f"{Device.DEFAULT}:1").schedule_linear().src
assert len(s3) == len(s4), f"{len(s3)} != {len(s4)}"
assert len(s2) == len(s4), f"{len(s)} != {len(s3)}"
for x,y in zip(s3, s4):
if not (x.src[0] == y.src[0]):
print(f"{x.src[0]} != {y.src[0]}")
@unittest.skipUnless(dtypes.bfloat16 in Device[Device.DEFAULT].renderer.supported_dtypes(), "need bfloat16 support")
def test_rand_bfloat16(self):
N = 128
x = Tensor.rand((2, N, N), dtype=dtypes.bfloat16)
assert x.dtype == dtypes.bfloat16
nx = x.numpy()
assert nx[nx == 1].size == 0
assert nx[nx == 0].size > 0
equal_distribution(lambda *x: Tensor.rand(*x, dtype=dtypes.bfloat16).float(), torch.rand, lambda x: np.random.rand(*x), shape=(2, N, N))
def test_rand_like(self):
empty = Tensor.empty((80, 44))
rand = Tensor.rand_like(empty)
assert rand.shape == empty.shape
assert rand.dtype == empty.dtype
assert rand.device == empty.device
def test_randn_like(self):
empty = Tensor.empty((80, 44))
rand = Tensor.randn_like(empty)
assert rand.shape == empty.shape
assert rand.dtype == empty.dtype
assert rand.device == empty.device
def test_rand_like_zero_shape(self):
empty = Tensor.empty(0, 20)
rand = Tensor.rand_like(empty)
assert rand.shape == empty.shape
assert rand.dtype == empty.dtype
assert rand.device == empty.device
def test_rand_like_more_dims(self):
empty = Tensor.empty((1, 2, 3, 4, 5, 6))
rand = Tensor.rand_like(empty)
assert rand.shape == empty.shape
assert rand.dtype == empty.dtype
assert rand.device == empty.device
def test_rand_like_dtype(self):
empty = Tensor.empty((80, 44), dtype=dtypes.float16)
rand = Tensor.rand_like(empty)
assert rand.shape == empty.shape
assert rand.dtype == empty.dtype
assert rand.device == empty.device
empty = Tensor.empty((80, 44))
rand = Tensor.rand_like(empty, dtype=dtypes.float16)
assert rand.shape == empty.shape
assert rand.dtype == dtypes.float16
assert rand.device == empty.device
def test_randn_like_dtype(self):
empty = Tensor.empty((80, 44), dtype=dtypes.float16)
rand = Tensor.randn_like(empty)
assert rand.shape == empty.shape
assert rand.dtype == empty.dtype
assert rand.device == empty.device
empty = Tensor.empty((80, 44))
rand = Tensor.randn_like(empty, dtype=dtypes.float16)
assert rand.shape == empty.shape
assert rand.dtype == dtypes.float16
assert rand.device == empty.device
def test_randn_device(self):
self.assertEqual(Tensor.randn(3,3,device="CPU").device, "CPU")
@given(strat.sampled_from([dtypes.float, dtypes.float16, dtypes.bfloat16]))
def test_randn_finite(self, default_float):
if default_float not in Device[Device.DEFAULT].renderer.supported_dtypes(): return
# low precision can result in inf from randn
self.enterContext(Context(DEFAULT_FLOAT=default_float))
t = Tensor.randn(64, 64)
mx = t.max().numpy().item()
mn = t.min().numpy().item()
print(f"testing with {default_float=}")
assert math.isfinite(mx), mx
assert math.isfinite(mn), mn
def test_random_counter_overflow(self):
device = Device.DEFAULT
Tensor.manual_seed(1337)
Tensor.rand(1).realize()
Tensor._device_rng_counters[device].assign(Tensor([dtypes.uint32.max - 5, 0], device=device, dtype=dtypes.uint32)).realize()
Tensor.rand(10).realize()
c = Tensor._device_rng_counters[device].numpy()
np.testing.assert_allclose(c, [4, 1])
Tensor.rand(10).realize()
c = Tensor._device_rng_counters[device].numpy()
np.testing.assert_allclose(c, [14, 1])
if __name__ == "__main__":
unittest.main()

234
artifacts/package_sources/tinygrad/test/backend/test_rangeify.py Normal file
View File

@@ -0,0 +1,234 @@
import unittest
from tinygrad import Tensor, nn, Device, dtypes, Variable
from tinygrad.helpers import Context, GlobalCounters, getenv, PCONTIG, DEBUG
from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops, UOp
from tinygrad.codegen.opt import OptOps, Opt
from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.nir import NIRRenderer
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "broken in LVP and PTX")
class TestDoubleMatmul(unittest.TestCase):
def setUp(self):
with Context(DEBUG=0):
self.a, self.b, self.c = [Tensor.randn(16, 16).contiguous().realize() for _ in range(3)]
self.ref = (self.a @ self.b @ self.c).realize()
def _test(self, opts):
with Context(PCONTIG=2, DEBUG=max(2, DEBUG.value)):
out = (self.a @ self.b @ self.c).contiguous(arg=opts).realize()
with Context(DEBUG=0):
err = (out-self.ref).square()
self.assertLess(err.max().item(), 1e-4)
self.assertLess(err.mean().item(), 1e-6)
def test_baseline(self): self._test(())
def test_upcast_0(self): self._test((Opt(OptOps.UPCAST, 0, 4),))
def test_upcast_1(self): self._test((Opt(OptOps.UPCAST, 1, 4),))
def test_upcast_2(self): self._test((Opt(OptOps.UPCAST, 2, 4),))
def test_upcast_01(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4)))
def test_upcast_01_mismatch(self): self._test((Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UPCAST, 1, 4)))
def test_upcast_02(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 2, 4)))
def test_upcast_12(self): self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 2, 4)))
def test_unroll_0(self): self._test((Opt(OptOps.UNROLL, 0, 4),))
def test_unroll_1(self): self._test((Opt(OptOps.UNROLL, 1, 4),))
def test_unroll_01(self): self._test((Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
def test_upcast_0_unroll_0(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 0, 4)))
def test_upcast_1_unroll_0(self): self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)))
def test_upcast_2_unroll_0(self): self._test((Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UNROLL, 0, 4)))
def test_upcast_0_unroll_1(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
def test_upcast_1_unroll_1(self): self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 1, 4)))
def test_upcast_2_unroll_1(self): self._test((Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UNROLL, 1, 4)))
def test_upcast_1_unroll_1_small(self): self._test((Opt(OptOps.UPCAST, 1, 2), Opt(OptOps.UNROLL, 1, 2)))
def test_upcast_1_unroll_1_rev(self): self._test((Opt(OptOps.UNROLL, 1, 2), Opt(OptOps.UPCAST, 1, 2)))
def test_upcast_01_unroll_01(self):
self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
def test_upcast_12_unroll_01(self):
self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
class TestRangeifyAssign(unittest.TestCase):
def test_assign_permuted(self):
A = Tensor.empty(4, 4, dtype='int')
B = Tensor.arange(16).reshape(4,4)
ret = A.permute(1,0).assign(B)
lst = ret.tolist()
lst2 = A.tolist()
lst3 = B.tolist()
print(lst)
print(lst2)
print(lst3)
self.assertListEqual(lst, lst3)
self.assertListEqual(lst2, B.permute(1, 0).tolist())
class TestRangeifyEdgeCase(unittest.TestCase):
def test_variable_stack_data(self):
# a bound-Variable STACK used as data gets ranges from its graph position
v = Variable("v", 0, 10).bind(3)
t = Tensor(UOp.stack(v, v+1).cast(dtypes.int)) + Tensor.arange(2)
self.assertListEqual(t.tolist(), [3, 5])
def test_variable_data_and_shape(self):
# the same Variable has a data edge through CAST and a structural edge through SHRINK
v = Variable("shared_v", 1, 10).bind(3)
t = Tensor.ones(10)[:v] * Tensor(v.cast(dtypes.float))
self.assertEqual(t.sum().item(), 9)
def test_matmul_relu_cat(self):
a = Tensor.ones(100, 512).contiguous().realize()
c = Tensor.ones(1, 512).contiguous().realize()
cm = Tensor.ones(512, 512)
c = c @ cm
c = c.relu()
res = Tensor.cat(a, c, dim=0)
self.assertEqual(res.numpy()[-1, :16].tolist(), [512] * 16)
def test_pcontig_multi_gather(self):
# regression test: local bufferize must have device set for const_like to work
with Context(PCONTIG=2):
# NOTE: with uint type, this will become a long and fail on WEBGPU
forest = Tensor(list(range(8)), dtype='int')
idx = Tensor([0, 0], dtype='int')
node_val = forest.gather(0, idx)
idx2 = idx * 2 + 1
node_val2 = forest.gather(0, idx2)
result = (node_val + node_val2).numpy()
self.assertEqual(result.tolist(), [1, 1])
if getenv("BIG") > 2:
# llama 8B (8192)
BS, HEADS, SEQLEN, EMB = 4, 32, 8192, 128
elif getenv("BIG") > 1:
# llama 8B
BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128
elif getenv("BIG") > 0:
# bigger
BS, HEADS, SEQLEN, EMB = 4, 32, 128, 128
else:
BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
def fa():
Tensor.manual_seed(1337)
with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
GlobalCounters.reset()
return q.scaled_dot_product_attention(k, v)
def fa_bw():
Tensor.manual_seed(1337)
with Context(DEBUG=0):
q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
attn_output = nn.Linear(HEADS*EMB, HEADS*EMB, bias=False)
attn_output.weight.realize()
target = Tensor.rand(BS, SEQLEN, HEADS*EMB).contiguous().realize()
GlobalCounters.reset()
attn = q.scaled_dot_product_attention(k, v).contiguous().contiguous_backward()
attn = attn.transpose(1, 2).reshape(BS, SEQLEN, -1)
out = attn_output(attn)
loss = (out - target).square().mean()
loss.backward()
#ret = [out, Tensor.stack(q.grad, k.grad, v.grad, dim=-1)]
#ret = [out, Tensor.stack(q.grad, k.grad, dim=-1), v.grad]
ret = [out, q.grad, k.grad, v.grad]
Tensor.realize(*ret)
return ret
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "broken in LVP and PTX")
class TestPcontig(unittest.TestCase):
def test_flash_attention_bw(self):
with Context(PCONTIG=max(2, PCONTIG.value), DEBUG=2):
grads = fa_bw()
print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
with Context(PCONTIG=0, DEBUG=2):
cmp_grads = fa_bw()
print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
with Context(DEBUG=0):
mses = [((x-y)**2).sum().item() for x,y in zip(grads, cmp_grads)]
mse = sum(mses)
print(f"mse: {mse}")
self.assertLessEqual(mse, 1e-6)
def test_flash_attention(self, opts=None):
with Context(PCONTIG=2, DEBUG=max(2, DEBUG.value)):
ret = fa().realize() if opts is None else fa().contiguous(arg=opts).realize()
print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
with Context(DEBUG=2):
cmp = fa().realize()
print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
with Context(DEBUG=0):
mse = ((cmp-ret)**2).sum().item()
print(f"mse: {mse}")
self.assertLessEqual(mse, 1e-6)
def test_flash_attention_opt(self):
opts = ()
# columns in top matrix
opts += (Opt(OptOps.UPCAST, 0, 4),)
# columns in bottom matrix
opts += (Opt(OptOps.UPCAST, 3, 4),)
# rows in all the matrix
opts += (Opt(OptOps.UPCAST, 4, 4),)
self.test_flash_attention(opts)
# contiguous + reduce can support ranges?
@unittest.skip("pm_rangeify no longer exists. test this in a different way")
class TestRangeifyPM(unittest.TestCase):
def setUp(self): self.base = Tensor.empty(10*10).reshape(10, 10).contiguous()
def assert_same(self, a, b):
def run_pm_rangeify(t:Tensor):
from tinygrad.schedule.rangeify import pm_rangeify, RangeifyContext
sink = t.uop.sink()
pm_realize = PatternMatcher([(UPat(Ops.CONTIGUOUS, name="x"), lambda x: x.replace(op=Ops.REALIZE))])
sink = graph_rewrite(sink, pm_realize)
return graph_rewrite(sink, pm_rangeify, ctx=RangeifyContext())
self.assertIs(run_pm_rangeify(a.contiguous()), run_pm_rangeify(b.contiguous()))
def test_nothing_match(self):
a = self.base.pad(((0,0),(0,1)))
b = self.base.pad(((0,0),(0,1)))
self.assert_same(a, b)
def test_reshape_match(self):
a = self.base
b = self.base.reshape(100).reshape(10, 10)
self.assert_same(a, b)
def test_permute_reshape_match(self):
a = self.base
b = self.base.permute(1,0).reshape(100).reshape(10, 10).permute(1,0)
self.assert_same(a, b)
def test_padded_permute_match(self):
a = self.base.pad(((0,0),(0,1)))
b = self.base.permute(1,0).pad(((0,1),(0,0))).permute(1,0)
self.assert_same(a, b)
@unittest.expectedFailure
def test_padded_reshape_match(self):
a = self.base.pad(((0,0),(0,1)))
b = self.base.reshape(100).reshape(10, 10).pad(((0,0),(0,1)))
self.assert_same(a, b)
@unittest.expectedFailure
def test_padded_permute_reshape_match(self):
a = self.base.pad(((0,0),(0,1)))
b = self.base.permute(1,0).reshape(100).reshape(10, 10).pad(((0,1),(0,0))).permute(1,0)
self.assert_same(a, b)
# why is this failing?
@unittest.expectedFailure
def test_cross_pad_match(self):
a = self.base.pad(((0,0),(0,1))).pad(((0,1),(0,0)))
b = self.base.pad(((0,1),(0,0))).pad(((0,0),(0,1)))
self.assert_same(a, b)
if __name__ == '__main__':
unittest.main()

123
artifacts/package_sources/tinygrad/test/backend/test_renderer_failures.py Normal file
View File

@@ -0,0 +1,123 @@
import unittest
import numpy as np
from tinygrad.device import Device, Buffer
from tinygrad.dtype import dtypes, ConstType
from tinygrad.engine.realize import run_linear
from tinygrad.codegen import to_program
from tinygrad.helpers import prod
from tinygrad.renderer.cstyle import CStyleLanguage
from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.wgsl import WGSLRenderer
from test.helpers import check_schedule
from tinygrad.runtime.ops_python import PythonRenderer
from tinygrad.uop.ops import UOp, Ops, KernelInfo, python_alu
from tinygrad.tensor import Tensor
def _test_uop_result(inputs:list[Tensor], sink:UOp, local_size=None):
for x in inputs: x.realize()
sz = 1 if local_size is None else prod(local_size)
outs = [UOp.new_buffer(Device.DEFAULT, sz, u.src[1].dtype) for u in sink.src if u.op is Ops.STORE]
for u in outs: u.buffer.allocate().copy_from(Buffer("PYTHON", sz, u.dtype, opaque=memoryview(bytearray(u.buffer.nbytes))))
run_linear(UOp(Ops.LINEAR, src=(sink.call(*outs, *(x.uop.base for x in inputs)),)))
return [u.buffer.numpy() for u in outs]
def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
dtype = alu_src_uops[0].dtype
a = UOp.param(0, dtype, (1,))
b = UOp.param(1, dtype, (1,))
idx = UOp.const(0)
ld = b.index(idx).load()
alu = ld.alu(alu_op, *alu_src_uops)
store = UOp.store(a.index(idx), alu)
return _test_uop_result([Tensor([input_val])], UOp(Ops.SINK, src=(store,), arg=KernelInfo()))[0]
class TestRendererFailures(unittest.TestCase):
@unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
def test_gated_store_with_alu(self):
a = UOp.param(0, dtypes.int, (4,))
gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0.valid(gate_alu)), UOp.const(1).cast(dtypes.int)))
sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo())
ret = _test_uop_result([], sink, local_size=[4, 1, 1])[0]
np.testing.assert_equal(ret, [0, 1, 1, 1])
@unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
def test_gated_store_with_alu_2d(self):
a = UOp.param(0, dtypes.int, (8,))
gate_alu_0 = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
gate_alu_1 = (lidx1:=UOp.special(2, 'lidx1')).ne(0)
gated_alu_store = UOp(Ops.STORE, src=(a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(1).cast(dtypes.int)))
sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo())
ret = _test_uop_result([], sink, local_size=[4, 2, 1])[0]
np.testing.assert_equal(ret, [0, 0, 0, 0, 0, 1, 1, 1])
@unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage), "uops are for cstyle")
class TestCStyleFailures(unittest.TestCase):
def test_inline_const_alu(self):
# CPU doesn't use the max function
ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int.min+1).cast(dtypes.int))
self.assertEqual(ret[0], 1)
def _test_src_strip_paren(self, op: Ops, should_strip_paren:bool=True):
dtype = "bool" if op in (Ops.OR, Ops.XOR, Ops.AND) else None
ret = Tensor.empty(1, dtype=dtype)
for _ in range(5): ret = python_alu[op](ret, Tensor.empty(1, dtype=dtype))
linear, _ = check_schedule(ret, 1)
src = to_program(linear.src[0].src[0], Device[Device.DEFAULT].renderer).src[2].arg
self.assertEqual("("*5 not in src, should_strip_paren)
def test_repeat_add(self): self._test_src_strip_paren(Ops.ADD)
def test_repeat_mul(self): self._test_src_strip_paren(Ops.MUL)
def test_repeat_xor(self): self._test_src_strip_paren(Ops.XOR)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "wgsl ends up with '(' * 5")
def test_repeat_or(self): self._test_src_strip_paren(Ops.OR)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "wgsl ends up with '(' * 5")
def test_repeat_and(self): self._test_src_strip_paren(Ops.AND)
def test_repeat_sub(self): self._test_src_strip_paren(Ops.SUB, should_strip_paren=False)
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "tests for wgsl renderer")
class TestWGSLFailures(unittest.TestCase):
def test_multiply_infinity(self):
# multiplying a positive constant by infinity should return infinity
# WGSL pipelines do not handle this reliably, some of which return zero, unless infinity always comes from a read on a dynamic buffer
ret = _setup_and_test_alu(Ops.MUL, 5.0, UOp.const(float("inf")).cast(dtypes.float32))
self.assertEqual(ret[0], float("inf"))
# WGSL has a specific select(alt, val, gate) ternary operator instead of gate?val:alt
def test_gated_load(self):
a = UOp.param(0, dtypes.int, (4,))
b = UOp.param(1, dtypes.int, (4,))
c = UOp.param(2, dtypes.int, (4,))
lidx0 = UOp.special(4, "lidx0")
gate = lidx0.ne(0)
alt = c.index(lidx0).load()
ld = UOp.load(b.index(lidx0.valid(gate)))
alt_load = gate.where(ld, alt)
store = UOp.store(a.index(lidx0), alt_load)
sink = UOp(Ops.SINK, src=(store,), arg=KernelInfo())
ret = _test_uop_result([Tensor([0,1,2,3], dtype=dtypes.int), Tensor([4,5,6,7], dtype=dtypes.int)], sink, local_size=[4])[0]
np.testing.assert_equal(ret, [4,1,2,3])
@unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "tests for ptx renderer")
class TestPTXFailures(unittest.TestCase):
@unittest.skip("INDEX can only have a gate ALU parent, not an IF")
def test_gated_store_with_if(self):
a = UOp.param(0, dtypes.int, (4,))
gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
val = UOp.const(1).cast(dtypes.int)
if_uop = UOp(Ops.IF, src=(gate_alu,))
gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0, if_uop), val))
sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo())
ret = _test_uop_result([], sink, local_size=[4, 1, 1])[0]
np.testing.assert_equal(ret, [0, 1, 1, 1])
@unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
def test_gated_define_acc_with_half_dtype(self):
a = Tensor.randn(32, 32, dtype=dtypes.half).realize()
b = Tensor.randn(34, 32, dtype=dtypes.half).realize()
result = a.pad((1,1)).matmul(b, dtype=dtypes.half).numpy()
reference = a.pad((1,1)).matmul(b, dtype=dtypes.float).numpy()
np.testing.assert_allclose(result, reference, atol=1e-2, rtol=1e-2)
if __name__ == '__main__':
unittest.main()

484
artifacts/package_sources/tinygrad/test/backend/test_schedule.py Normal file
View File

@@ -0,0 +1,484 @@
# this will be the new test_ops for the next level
# schedule confirms the right things are capable of fusing
# NOTE: this has overlap with external_test_opt.py
import unittest, time
import numpy as np
from tinygrad import nn, dtypes, Device, Tensor, Variable
from tinygrad.uop.ops import Ops, UPat
from tinygrad.helpers import DEV, GlobalCounters, Context, all_same, temp
from tinygrad.engine.realize import run_linear
from test.helpers import check_schedule, assert_kernel_count
supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
def _realize_weights(m):
for p in nn.state.get_parameters(m): p.realize()
class TestSchedule(unittest.TestCase):
def setUp(self):
self.ctx = Context(SPLIT_REDUCEOP=0)
self.ctx.__enter__()
def tearDown(self):
self.ctx.__exit__(None, None, None)
@unittest.skip("no longer supported")
def test_double_from(self):
x = Tensor([1,2,3,4])
out = x.to('python')
check_schedule(out, 0, filter_sink=False)
def test_example_matmul_same(self):
x = Tensor.eye(64).clone().realize()
z = x.matmul(x).sum()
z.backward()
out = x.grad.contiguous()
run_linear(*check_schedule(out, 1))
# NOTE: the gradient flows twice
np.testing.assert_allclose(out.numpy(), 2*np.ones((64,64)))
def test_pad_reduce_scope_collision(self):
b = Tensor.rand(4, 3).realize()
s1 = b.pad(((1, 1), (0, 0))).sum(axis=1)
s2 = b.pad(((1, 2), (0, 0))).shrink(((0, 6), (0, 3))).sum(axis=1)
out = s1 + s2
run_linear(*check_schedule(out, 1))
np.testing.assert_allclose(out.numpy(), 2*np.pad(b.numpy(), ((1, 1), (0, 0))).sum(axis=1), rtol=1e-6)
def test_cumsum_parallel_reduce_fused(self):
# two-stage cumsum + ops triggers parallel REDUCEs in one kernel that must share an END (same nesting context = should merge)
step, num_steps = 513, 10
t = Tensor.arange(step).float().realize()
phase = t.cumsum()
tiled = phase.repeat((num_steps,)).reshape(num_steps, step)
pattern = Tensor([1,0,0,1,0,0,0,0,1,0]).reshape(num_steps, 1)
out = (tiled * pattern).flatten()
expected = np.tile(np.arange(step).astype(np.float32).cumsum(), num_steps).reshape(num_steps, step)
expected = (expected * np.array([1,0,0,1,0,0,0,0,1,0]).reshape(num_steps, 1)).flatten()
np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-4)
@unittest.skipIf(Device.DEFAULT == "CL", "TODO: fails on CI CL")
def test_reduce_different_nesting_depth(self):
# two REDUCEs sharing the same RANGE at different nesting depths must NOT merge
x = Tensor.arange(768).reshape(3, 256).float()
np.testing.assert_allclose((x.sum(axis=1) + x.sum(axis=1).sum()).numpy(), x.numpy().sum(axis=1) + x.numpy().sum(axis=1).sum())
def test_fuse_assign_contiguous(self):
x = Tensor.zeros(4, 4, dtype=dtypes.int).contiguous().realize()
a = Tensor.arange(8).reshape(4, 2)
run_linear(*check_schedule(x.shrink((None, (0, 2))).assign(a.clone()), 2))
np.testing.assert_equal(x.numpy(), [[0, 1, 0, 0], [2, 3, 0, 0], [4, 5, 0, 0], [6, 7, 0, 0]])
def test_assign_non_contiguous_alt(self): self.test_assign_non_contiguous(alt=True)
def test_assign_non_contiguous(self, alt=False):
x = (Tensor.arange(16)-100).reshape(4,4).clone().realize()
xref = x.numpy()
if alt:
y = Tensor.randint(2, 4).contiguous().realize()
a = Tensor.arange(8).reshape(2, 4)+y
tst = x.shrink(((0, 2), None)).assign(a).realize()
xref[:2, :] = np.arange(8).reshape(2, 4)+y.numpy()
else:
y = Tensor.randint(4, 2).contiguous().realize()
a = Tensor.arange(8).reshape(4, 2)+y
tst = x.shrink((None, (0, 2))).assign(a).realize()
xref[:, :2] = np.arange(8).reshape(4, 2)+y.numpy()
np.testing.assert_equal(x.numpy(), xref)
np.testing.assert_equal(tst.numpy(), a.numpy())
def test_setitem_sched(self, mop=lambda x:x, expected_kcount=1):
a = Tensor.arange(16).reshape(4, 4).clone(device="CPU").realize()
a2 = mop(a)
expected = (a+a2).tolist()
a.assign(a+a2)
linear, var_vals = check_schedule(a, expected_kcount)
run_linear(linear, var_vals)
self.assertListEqual(a.tolist(), expected)
def test_setitem_permuted_sched(self): self.test_setitem_sched(lambda x: x.T, 2)
def test_setitem_paddded_sched(self): self.test_setitem_sched(lambda x: x.shrink_to(4, 1).pad_to(4, 4), 1)
def test_setitem_const_fused(self):
# https://github.com/tinygrad/tinygrad/issues/10690
a = Tensor.arange(16).clone().realize()
GlobalCounters.reset()
a[4] = 3
assert_kernel_count(0)
a.realize()
assert_kernel_count(1)
self.assertListEqual(a.tolist(), [0, 1, 2, 3, 3, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15])
def test_no_extra_contiguous_on_setitem_assign_back(self):
# pattern: contiguous copy, advanced setitem, assign back (e.g. torch backend _view_write)
base = Tensor.arange(16).reshape(4, 4).clone()
flat_base = base.reshape(16).contiguous()
idx = Tensor([1,2,5,6], dtype=dtypes.int32)
flat_base[idx] = Tensor([99,99,99,99])
base.assign(flat_base.reshape(4, 4))
sched = check_schedule(base, 4)
run_linear(*sched)
expected = list(range(16))
for i, v in zip([1,2,5,6], [99,99,99,99]): expected[i] = v
np.testing.assert_equal(base.reshape(16).numpy(), expected)
def test_const_folding_alt(self):
t = Tensor.full((2,), 1.)
lt = (t < 0.)
a = Tensor.empty(2).assign(t*lt.where(-1., 0.))
b = Tensor.empty(2, dtype=dtypes.bool).assign(lt)
Tensor.realize(a, b)
self.assertEqual(a.tolist(), [0., 0.])
self.assertEqual(b.tolist(), [False, False])
def test_self_assign_no_empty_kernel(self):
for shape in [(3, 3), (4, 4)]:
a = Tensor.ones(*shape).contiguous().realize()
a.assign(a / 1)
run_linear(*check_schedule(a, 0, filter_sink=False))
self.assertListEqual(a.tolist(), [[1.]*shape[1]]*shape[0])
def test_deviceless_materialize_localizes_to_target(self):
dev = "CPU" if Device.DEFAULT != "CPU" else "CPU:1"
t = Tensor.arange(Variable("s", 1, 128).bind(64)).cumsum().clone(dev)
self.assertEqual(t.device, dev)
np.testing.assert_equal(t[:64].numpy(), np.arange(64).cumsum())
def test_copy_multi_scalar(self):
devs = ("CPU:0", "CPU:1")
x = Tensor.ones(2, device="CPU").shard(devs, axis=0).realize()
out = (x.sum()*2).reshape(1).to("CPU")
run_linear(*check_schedule(out, 5))
np.testing.assert_equal(out.numpy(), [4.])
class TestLimitBufs(unittest.TestCase):
@unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV", "crashes in ocelot")
def test_limit_bufs_with_var(self):
N = 31
with Context(TRACK_MATCH_STATS=0, DEBUG=0):
bufs = [Tensor([1]*10).contiguous().realize() for i in range(N)]
vi = Variable("i", 0, 9).bind(1)
vj = Variable("j", 0, 9).bind(2)
root = bufs[0][vi] + bufs[0][vj]
for X in range(1,N): root = root + bufs[X][vi] + bufs[X][vj]
self.assertEqual(root.item(), N * 2)
def test_limit_bufs_arange_condition(self):
# WHERE with arange-based condition (pure index math, no device) and many buffer loads should not crash limit_bufs
with Context(MAX_KERNEL_BUFFERS=8):
N = 8
idx = Tensor.arange(N)
base = Tensor.zeros(N)
for i in range(4):
a, b = Tensor.rand(N).realize(), Tensor.rand(N).realize()
base = (idx >= i).where(a + b, base)
assert all(x > 0 for x in base.tolist())
def test_limit_bufs_linear_scaling(self):
def sched_time(n):
with Context(TRACK_MATCH_STATS=0, DEBUG=0):
bufs = [Tensor.ones(16).contiguous().realize() for _ in range(4)]
root = bufs[0]
for i in range(n): root = root + bufs[i % 4]
with Context(MAX_KERNEL_BUFFERS=8, SCACHE=0):
st = time.perf_counter()
root.schedule_linear()
return time.perf_counter() - st
sched_time(400)
t1, t2 = min(sched_time(400) for _ in range(3)), min(sched_time(1600) for _ in range(3))
self.assertLess(t2/t1, 8, f"{t1*1e3:.1f}ms -> {t2*1e3:.1f}ms")
class TestSwizzle(unittest.TestCase):
def test_swizzle_simple(self):
Tensor.manual_seed(0)
with Context(DEBUG=0, TRACK_MATCH_STATS=0):
a = Tensor.randint(32, 32).realize()
r = (a+a).sum(1).sum(0)
# double reduce collapses to a single reduce
run_linear(*check_schedule(r, 1))
self.assertEqual(r.numpy(), (a.numpy()+a.numpy()).sum(1).sum(0))
def test_single_swizzle(self):
Tensor.manual_seed(0)
with Context(DEBUG=0, TRACK_MATCH_STATS=0):
a = Tensor.randint(4, 1).realize()
b = Tensor.ones((1, 1), dtype=a.dtype).contiguous().realize()
# ADD(REDUCE(RESHAPE(LOAD)), LOAD) to ADD(REDUCE(RESHAPE(LOAD))), RESHAPE(LOAD)
r = a.sum(0)+b
run_linear(*check_schedule(r, 1))
self.assertEqual(r.numpy(), a.numpy().sum(0)+1)
def test_double_swizzle_possible(self):
Tensor.manual_seed(0)
with Context(DEBUG=0, TRACK_MATCH_STATS=0):
a = Tensor.randint(4,).realize()
b = Tensor.randint(4,).realize()
# parallel reduce!
add = a.sum(0)+b.sum(0)
run_linear(*check_schedule(add, 1))
self.assertEqual(add.numpy(), a.numpy().sum(0)+b.numpy().sum(0))
def test_swizzle_reduceop(self):
Tensor.manual_seed(0)
x = Tensor.randn(4,4).realize()
y = Tensor.randn(4,4,4).realize()
out = x.reshape(4,4,1).expand(4,4,4).sum(axis=(1,))+y
run_linear(*check_schedule(out, 2)) # TODO: 1?
np.testing.assert_allclose(out.numpy(), np.tile(x.numpy().reshape(4,4,1), (1,1,4)).sum(axis=1)+y.numpy())
def test_permute_rewrite(self):
x = Tensor.randn(4, 4, 16).realize()
y = Tensor.randn(4, 1, 16).realize()
z = Tensor.randn(4, 4, 1).realize()
t = (x*y).sum(axis=(0, 2)).reshape(1, 4, 1).permute(0, 2, 1)+z
run_linear(*check_schedule(t, 2)) # TODO: 1?
t_np = (x.numpy()*y.numpy()).sum(axis=(0, 2)).reshape(1, 4, 1).transpose(0, 2, 1)+z.numpy()
np.testing.assert_allclose(t.numpy(), t_np, atol=1e-6, rtol=1e-3)
@unittest.skip("TODO: this swizzle isn't resolvable when there's a mask")
def test_swizzle_failure_permute(self):
a = Tensor.empty(45,65).T.reshape(65,1,45).pad((None,None,(0,45))).expand(65,45,90)
b = Tensor.empty(45,65)
a_reduce = a.sum(axis=(2,), keepdim=True).sum(axis=(1,))
b_reduce = b.sum(axis=(0,))
t = a_reduce+b_reduce
run_linear(*check_schedule(t, 1))
def test_parallel_reduce_possible(self):
Tensor.manual_seed(0)
x = Tensor.randn(4, 2, 2).realize()
y = Tensor.randn(4, 2, 2).realize()
t = x.sum(axis=1)+y.sum(axis=1)
run_linear(*check_schedule(t, 1))
np.testing.assert_allclose(t.numpy(), x.numpy().sum(axis=1)+y.numpy().sum(axis=1), atol=1e-6, rtol=1e-3)
# kernels can only have 1 or n in each dim
def test_dont_parallelize_different_n(self):
Tensor.manual_seed(0)
x = Tensor.randn(4, 2, 2).realize()
y = Tensor.randn(4, 3, 2).realize()
t = x.sum(axis=1)+y.sum(axis=1)
run_linear(*check_schedule(t, 1))
np.testing.assert_allclose(t.numpy(), x.numpy().sum(axis=1)+y.numpy().sum(axis=1), atol=1e-6, rtol=1e-3)
def test_unsafe_pad(self):
x = Tensor.full((2,2), 1.0).contiguous()
y = x*x.sum((1,)).reciprocal()
t = y.pad(((0,1),None))
run_linear(*check_schedule(t, 3))
np.testing.assert_equal(t.numpy(), [[0.5, 0.5], [0.5, 0.5], [0., 0.]])
zero_pm = UPat(Ops.CONST, arg=0)
class TestView(unittest.TestCase):
def test_all_masked_out(self):
# start with non CONST Ops
a = Tensor.rand(10, 10).realize()
# all masked out, degrades to const 0
b = a.pad(((0, 10), None))[10:]
sched = check_schedule(b.contiguous(), 1)
run_linear(*sched)
np.testing.assert_equal(b.numpy(), 0)
def test_mask_dim_1(self):
# mask out dim = 1 works too
a = Tensor.rand(10, 10).realize()
b = a.pad((None, (0, 10)))[:, 10:]
assert b.shape == (10, 10)
sched = check_schedule(b.contiguous(), 1)
run_linear(*sched)
np.testing.assert_equal(b.numpy(), 0)
def test_partial_mask(self):
# partial masked out does not degrade into CONST
a = Tensor.rand(10, 10).realize()
b = a.pad(((0, 5), None))[5:]
assert b.shape == (10, 10)
sched = check_schedule(b.contiguous(), 1)
run_linear(*sched)
np.testing.assert_allclose(b.numpy(), np.pad(a.numpy(), ((0, 5), (0, 0)))[5:])
# a*VIEW(x), where VIEW(x) = 0
# x collapses along with its children
def test_parent_view_collapses(self):
a = Tensor([1, 2])
b = Tensor.arange(3).clone()
bv = b.pad(((0, 2),))[-2:]
# this becomes a late a*0
late_mul = a*bv
run_linear(*check_schedule(late_mul, 2))
# the arange doesn't realize
#self.assertIsNone(b.uop.base.realized)
# mul doesn't realize
#self.assertIsNone(late_mul.uop.base.realized)
self.assertEqual(late_mul.tolist(), [0, 0])
# SINK has two branches:
# a*VIEW(x), where VIEW(x) = 0
# x+2
# as long as one child realizes, x does not collapse
def test_parent_multiple_children_no_collapse(self):
a = Tensor([1, 2])
b = Tensor.arange(3).clone()
bv = b.pad(((0, 2),))[-2:]
late_mul = a*bv
other_child = b+2
s = check_schedule([late_mul, other_child], 3)
# the arange becomes a BUFFER
self.assertIs(b.uop.base.op, Ops.BUFFER)
# NOTE: no longer checked
# mul still collapses
#self.assertIs(late_mul.uop.base.op, Ops.CONST)
run_linear(*s)
self.assertEqual(other_child.tolist(), [2, 3, 4])
@unittest.skipIf(Device.DEFAULT == "CPU", "tests copy from another device to cpu")
class TestCopyFolding(unittest.TestCase):
def test_const_copy_is_free(self):
b = Tensor(1).to("CPU") * 4
run_linear(*check_schedule(b, 0, filter_sink=False))
assert b.item() == 4
def test_one_hot_with_copy(self):
y = Tensor([1, 2, 3]).to("CPU")
x = y.one_hot(10).int()
check_schedule(x, 3, filter_sink=False)
@unittest.skip("no longer supported")
def test_late_const_copy_folding(self):
a = Tensor.arange(3).clone().realize()
zeros = Tensor.zeros(3, buffer=False).realize()
b = (a*zeros).to("CPU") + 1
run_linear(*check_schedule(b, 1, filter_sink=False))
self.assertListEqual(b.tolist(), [1, 1, 1])
self.assertEqual(b.device, "CPU")
def test_alu_after_copy(self):
a = Tensor.ones((4,)).to("CPU")
b = Tensor.empty(4, device="CPU")
add = a+b
assert all_same([x.device for x in add.uop.src]), f"ALU has different devices! {[x.device for x in add.src]}"
add.schedule_linear()
def test_alu_before_copy(self):
buf = Tensor.ones(1).contiguous().realize()
a = buf+1
b = a.to("CPU")
self.assertListEqual(b.tolist(), [2.])
def test_copy_to_same_device(self):
a = Tensor.empty(4).uop
b = a.copy_to_device(a.device)
check_schedule(b, 1, filter_sink=False) # TODO: 0?
def test_copy_to_same_device_alt(self):
a = Tensor.empty(4, 4).uop
b = a.copy_to_device(a.device)
check_schedule(b, 1, filter_sink=False) # TODO: 0?
def test_copy_to_same_device_sched(self):
a = Tensor.ones(4).contiguous().realize().uop.buf_uop
t = Tensor(a.copy_to_device(a.device))
linear, var_vals = t.linear_with_vars()
assert len([call for call in linear.src if call.src[0].op is Ops.COPY]) == 0
run_linear(linear, var_vals)
assert t.uop.is_realized, f"didn't realize Tensor {t}"
self.assertListEqual(t.tolist(), [1.,1.,1.,1.])
@unittest.skip("same-device copies are no-ops")
def test_self_assign_same_device_copy(self):
a = Tensor.ones(4, 4).contiguous().realize()
# use copy_to_device to bypass Tensor.to() shortcircuit and force a real same-device COPY in the graph
a.assign(Tensor(a.uop.copy_to_device(a.device), a.device))
run_linear(*check_schedule(a, 2, filter_sink=False))
self.assertListEqual(a.tolist(), [[1.]*4]*4)
def test_clone(self):
a = Tensor.empty(4)
check_schedule(a.clone(), 1, filter_sink=False)
def test_shrink_copy(self):
a = Tensor.arange(4)
view = a.shrink(((0, 2),))
b = view.clone()
run_linear(*check_schedule(b, 1, filter_sink=False))
self.assertEqual(b.uop.base.buffer.size, 2)
self.assertEqual(b.uop.numel(), 2)
self.assertListEqual(b.tolist(), [0, 1])
def test_expanded_copy(self):
a = Tensor.arange(2)
view = a.reshape(2, 1).expand(2, 2)
b = view.clone()
run_linear(*check_schedule(b, 1, filter_sink=False))
self.assertEqual(b.uop.base.buffer.size, 4)
self.assertEqual(b.uop.numel(), 4)
self.assertListEqual(b.tolist(), [[0, 0], [1, 1]])
def test_permuted_copy(self):
a = Tensor.arange(4)
b = a.reshape(2, 2).permute(1, 0)
b.realize()
self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
def test_permute_on_disk(self):
with open(temp('dt_arange_4_permute'), "wb") as f: f.write(Tensor.arange(4).clone().realize().uop.base.buffer.as_memoryview())
a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute')}")
b = a.reshape(2, 2).permute(1, 0).to("CPU")
b.realize()
self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
def test_permute_on_disk_contiguous(self):
with open(temp('dt_arange_4_permute_contig'), "wb") as f: f.write(Tensor.arange(4).clone().realize().uop.base.buffer.as_memoryview())
a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute_contig')}")
b = a.reshape(2, 2).permute(1, 0).contiguous().to("CPU")
b.realize()
self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
def test_permute_after_shrink(self):
a = Tensor.arange(5)
b = a.shrink(((0, 4),)).reshape(2, 2).permute(1, 0).to("CPU")
b.realize()
self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
# NOTE: disk permute must come after COPY
def test_permute_after_shrink_on_disk(self):
with open(temp('dt_arange_5_permute'), "wb") as f: f.write(Tensor.arange(5).clone().realize().uop.base.buffer.as_memoryview())
a = Tensor.empty(5, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_5_permute')}")
b = a.shrink(((0, 4),)).reshape(2, 2).permute(1, 0).to("CPU")
b.realize()
self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
def test_permute_copy_to_device(self):
b = Tensor([[0, 1, 2, 3], [4, 5, 6, 7]], device="CPU").permute(1, 0).to("PYTHON")
self.assertListEqual(b.tolist(), [[0, 4], [1, 5], [2, 6], [3, 7]])
def test_flip_copy_to_device(self):
b = Tensor([0, 1, 2, 3], device="CPU").flip(0).to("PYTHON")
self.assertListEqual(b.tolist(), [3, 2, 1, 0])
class TestUOpBecome(unittest.TestCase):
def test_setitem_offset(self):
a = Tensor.full((16,), 0.).contiguous().realize()
b = Tensor.full((16,), 1.).contiguous().realize()
a_view = a[4:].reshape(3, 4).shrink(((0,2),(0,2))).reshape((4,))
b.shrink(((0,4),)).assign(a_view).realize()
self.assertListEqual(b.tolist(), [0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0])
class TestFusionOp(unittest.TestCase):
def test_contiguous_add(self):
def test(contig=False):
bt = Tensor(np.arange(16), dtype=dtypes.float32).reshape(4,4)
x = bt.permute(1,0)
if contig: x = x.contiguous()
return (x.permute(1,0) + bt).data()
assert test() == test(True)
def test_expand_fuse(self):
bt = Tensor(np.ones((10, 1)), dtype=dtypes.float32)
out = (bt*2).expand(10,10).sum(1)
run_linear(*out.linear_with_vars())
outd = out.tolist()
assert all(x == 20.0 for x in outd)
if __name__ == '__main__':
unittest.main(verbosity=2)

378
artifacts/package_sources/tinygrad/test/backend/test_setitem.py Normal file
View File

@@ -0,0 +1,378 @@
import unittest
from tinygrad import Tensor, TinyJit, Variable, dtypes, Device
from tinygrad.helpers import Context
import numpy as np
class TestSetitem(unittest.TestCase):
def test_simple_setitem(self):
cases = (
((6,6), (slice(2,4), slice(3,5)), Tensor.ones(2,2)),
((6,6), (slice(2,4), slice(3,5)), Tensor([1.,2.])),
((6,6), (slice(2,4), slice(3,5)), 1.0),
((6,6), (3, 4), 1.0),
((6,6), (3, None, 4, None), 1.0),
((4,4,4,4), (Ellipsis, slice(1,3), slice(None)), Tensor(4.0)),
((4,4,4,4), (Ellipsis, slice(1,3)), 4),
((4,4,4,4), (2, slice(1,3), None, 1), 4),
((4,4,4,4), (slice(1,3), slice(None), slice(0,4,2)), 4),
((4,4,4,4), (slice(1,3), slice(None), slice(None), slice(0,3)), 4),
((6,6), (slice(1,5,2), slice(0,5,3)), 1.0),
((6,6), (slice(5,1,-2), slice(5,0,-3)), 1.0),
)
for shp, slc, val in cases:
t = Tensor.zeros(shp).contiguous()
t[slc] = val
n = np.zeros(shp)
n[slc] = val.numpy() if isinstance(val, Tensor) else val
np.testing.assert_allclose(t.numpy(), n)
def test_padded_setitem(self):
t = Tensor.arange(10)
t[4:1:-2] = 11
self.assertListEqual(t.tolist(), [0, 1, 11, 3, 11, 5, 6, 7, 8, 9])
def test_setitem_inplace_mul(self):
t = Tensor.arange(10).clone().realize()
t[:3] *= 10
self.assertListEqual(t.tolist(), [0, 10, 20, 3, 4, 5, 6, 7, 8, 9])
@unittest.skip("crashed in LLVM CI")
def test_setitem_fancy_on_unrealized_view(self):
# fancy indexing setitem on unrealized SHRINK view (triggered infinite loop in graph_rewrite)
base = Tensor.arange(20, dtype=dtypes.float).reshape(4, 5).clone().realize()
sub = base[1:3]
flat = sub.reshape(sub.numel()).contiguous()
idx = Tensor([0, 3, 7, 9])
flat[idx] = Tensor([99, 98, 97, 96], dtype=dtypes.float)
sub.assign(flat.reshape(2, 5))
np.testing.assert_allclose(sub.numpy(), [[99, 6, 7, 98, 9], [10, 11, 97, 13, 96]])
def test_setitem_dtype(self):
for dt in (dtypes.int, dtypes.float, dtypes.bool):
for v in (5., 5, True):
t = Tensor.ones(6,6, dtype=dt).contiguous()
t[1] = v
self.assertEqual(t.dtype, dt)
def test_setitem_dtype_mismatch(self):
t = Tensor.zeros(6, dtype=dtypes.float).contiguous().realize()
with self.assertRaises(RuntimeError): t[2:4] = Tensor([1, 2], dtype=dtypes.int)
def test_setitem_chained_indexing(self):
# N[i][j] must work the same as N[i, j]
N1 = Tensor.zeros((3, 3)).contiguous().realize()
N1[1, 2] = 5
N2 = Tensor.zeros((3, 3)).contiguous().realize()
N2[1][2] = 5
np.testing.assert_equal(N1.numpy(), N2.numpy())
def test_setitem_detach(self):
# setitem on detached tensor should work
t = Tensor.zeros((3, 3)).contiguous().realize()
t.detach()[1, 2] = 5
self.assertEqual(t[1, 2].item(), 5.0)
def test_setitem_permute(self):
# setitem on permuted tensor should modify original
t = Tensor.zeros((2, 3)).contiguous().realize()
t.T[1, 0] = 5 # t.T is (3, 2), so [1, 0] maps to t[0, 1]
self.assertEqual(t[0, 1].item(), 5.0)
def test_setitem_flip(self):
# setitem on flipped tensor should modify original
t = Tensor.zeros((3,)).contiguous().realize()
t[::-1][0] = 5 # flip, then set first element (which is last in original)
self.assertEqual(t[2].item(), 5.0)
def test_setitem_inplace_operator(self):
t = Tensor.arange(4).reshape(2, 2).contiguous()
t[1] += 2
np.testing.assert_allclose(t.numpy(), [[0, 1], [4, 5]])
t = Tensor.arange(4).reshape(2, 2).contiguous()
t[1] -= 1
np.testing.assert_allclose(t.numpy(), [[0, 1], [1, 2]])
t = Tensor.arange(4).reshape(2, 2).contiguous()
t[1] *= 2
np.testing.assert_allclose(t.numpy(), [[0, 1], [4, 6]])
# NOTE: have to manually cast setitem target to least_upper_float for div
t = Tensor.arange(4, dtype=dtypes.float).reshape(2, 2).contiguous()
t[1] /= 2
np.testing.assert_allclose(t.numpy(), [[0, 1], [1, 1.5]])
t = Tensor.arange(4).reshape(2, 2).contiguous()
t[1] **= 2
np.testing.assert_allclose(t.numpy(), [[0, 1], [4, 9]])
t = Tensor.arange(4).reshape(2, 2).contiguous()
t[1] ^= 5
np.testing.assert_allclose(t.numpy(), [[0, 1], [7, 6]])
def test_setitem_consecutive_inplace_operator(self):
t = Tensor.arange(4).reshape(2, 2).contiguous()
t[1] += 2
t[1] -= 1
np.testing.assert_allclose(t.numpy(), [[0, 1], [3, 4]])
def test_setitem_overlapping_indices(self):
t = Tensor([1,2,3,4])
# regular overlapping indices
t[[1,1]] = Tensor([5,6])
np.testing.assert_allclose(t.numpy(), [1,6,3,4])
# overlapping indices with zero value overlapped
t[[1,1]] = Tensor([0,1])
np.testing.assert_allclose(t.numpy(), [1,1,3,4])
def test_setitem_overlapping_indices_with_0(self):
t = Tensor([1,2,3,4])
t[[1,1]] = Tensor([1,0])
np.testing.assert_allclose(t.numpy(), [1,0,3,4])
def test_setitem_with_1_in_shape(self):
t = Tensor([[1],[2],[3]])
t[[0,0]] = Tensor([[1],[2]])
np.testing.assert_allclose(t.numpy(), [[2],[2],[3]])
def test_fancy_setitem(self):
t = Tensor.zeros(6,6).contiguous()
t[[1,2], [3,2]] = 3
n = np.zeros((6,6))
n[[1,2], [3,2]] = 3
np.testing.assert_allclose(t.numpy(), n)
def test_simple_jit_setitem(self):
@TinyJit
def f(t:Tensor, a:Tensor):
t[2:4, 3:5] = a
# NOTE: without return t or an explicit realize, it's lazy and not captured
return t
for i in range(1, 6):
t = Tensor.zeros(6, 6).contiguous().realize()
a = Tensor.full((2, 2), fill_value=i, dtype=dtypes.float).contiguous()
f(t, a)
n = np.zeros((6, 6))
n[2:4, 3:5] = np.full((2, 2), i)
np.testing.assert_allclose(t.numpy(), n)
def test_jit_setitem_variable_offset(self):
with Context(CHECK_OOB=0):
@TinyJit
def f(t:Tensor, a:Tensor, v:Variable):
t.shrink(((v,v+1), None)).assign(a).realize()
t = Tensor.zeros(6, 6).contiguous().realize()
n = np.zeros((6, 6))
for i in range(6):
v = Variable("v", 0, 6).bind(i)
a = Tensor.full((1, 6), fill_value=i+1, dtype=dtypes.float).contiguous()
n[i, :] = i+1
f(t, a, v)
np.testing.assert_allclose(t.numpy(), n)
np.testing.assert_allclose(t.numpy(), [[1,1,1,1,1,1],[2,2,2,2,2,2],[3,3,3,3,3,3],[4,4,4,4,4,4],[5,5,5,5,5,5],[6,6,6,6,6,6]])
def test_setitem_overlapping_inplace1(self):
t = Tensor([[3.0], [2.0], [1.0]]).contiguous()
t[1:] = t[:-1]
self.assertEqual(t.tolist(), [[3.0], [3.0], [2.0]])
def test_setitem_overlapping_inplace2(self):
t = Tensor([[3.0], [2.0], [1.0]]).contiguous()
t[:-1] = t[1:]
self.assertEqual(t.tolist(), [[2.0], [1.0], [1.0]])
# TODO: WEBGPU pipeline validation error. this generates (1==gidx0)|(2==gidx0)|(3==gidx0)|(4==gidx0)|(5==gidx0) ...
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU pipeline validation error")
def test_setitem_big(self):
idx_size, val = 256, 4
t = Tensor.arange(0, idx_size+1)
idx = Tensor.arange(0, idx_size)
t[idx] = val
self.assertEqual(t.tolist(), [val]*idx_size+[idx_size])
def test_setitem_advanced_indexing(self):
# Example from https://numpy.org/doc/stable/user/basics.indexing.html#combining-advanced-and-basic-indexing
t = Tensor.zeros(10,20,30,40,50, dtype=dtypes.int).contiguous()
ind_1 = Tensor([5,3,7,8])
ind_2 = Tensor([[[0],[1],[2]],[[3],[4],[5]]])
v = Tensor.arange(2*3*4*10*30*50).reshape(2,3,4,10,30,50)
t[:, ind_1, :, ind_2, :] = v
n = np.zeros((10,20,30,40,50), dtype=np.int32)
n[:, ind_1.numpy(), :, ind_2.numpy(), :] = v.numpy()
np.testing.assert_equal(t.numpy(), n)
def test_setitem_tensor_int_indexing(self):
t = Tensor.zeros(4, 3, dtype=dtypes.int).contiguous()
t[Tensor([0, 2]), 0] = Tensor([99, 88], dtype=dtypes.int)
n = np.zeros((4, 3), dtype=np.int32)
n[[0, 2], 0] = [99, 88]
np.testing.assert_equal(t.numpy(), n)
def test_setitem_tensor_slice_indexing(self):
t = Tensor.zeros(4, 3, dtype=dtypes.int).contiguous()
t[Tensor([0, 2]), :2] = Tensor([[10, 20], [30, 40]], dtype=dtypes.int)
n = np.zeros((4, 3), dtype=np.int32)
n[[0, 2], :2] = [[10, 20], [30, 40]]
np.testing.assert_equal(t.numpy(), n)
def test_setitem_2d_tensor_indexing(self):
t = Tensor.zeros(2, dtype=dtypes.int).contiguous()
index = Tensor([[0, 1], [1,0]])
v = Tensor.arange(2*2).reshape(2, 2).contiguous()
t[index] = v
n = np.zeros((2,), dtype=np.int32)
n[index.numpy()] = v.numpy()
np.testing.assert_equal(t.numpy(), n)
def test_setitem_swap_rows(self):
t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).clone().realize()
tmp = t[0]
t[0] = t[1]
t[2] = tmp
# NOTE: not [[2, 3], [2, 3], [0, 1]], same with eager
np.testing.assert_allclose(t.numpy(), [[2, 3], [2, 3], [2, 3]])
# eager version
t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).clone().realize()
tmp = t[0].realize()
t[0] = t[1].realize()
t[2] = tmp.realize()
np.testing.assert_allclose(t.numpy(), [[2, 3], [2, 3], [2, 3]])
def test_lazy_sum_between_writes(self):
# lazy sums should capture buffer state at the time they were created
t = Tensor.zeros(6).contiguous().realize()
s0 = t.sum()
t[:3].assign(1.0)
s1 = t.sum()
t[3:].assign(2.0)
s2 = t.sum()
try:
np.testing.assert_allclose([s0.item(), s1.item(), s2.item()], [0.0, 3.0, 9.0])
except AssertionError:
# TODO: broken now, lazy sums all see final buffer state
np.testing.assert_allclose([s0.item(), s1.item(), s2.item()], [9.0, 9.0, 9.0])
# eager version
t = Tensor.zeros(6).contiguous().realize()
s0 = t.sum().realize()
t[:3].assign(1.0).realize()
s1 = t.sum().realize()
t[3:].assign(2.0).realize()
s2 = t.sum().realize()
np.testing.assert_allclose([s0.item(), s1.item(), s2.item()], [0.0, 3.0, 9.0])
def test_cross_assign_independence(self):
# when assigning to two tensors using computations from both,
# both assigns should see the OLD values of both tensors
a = Tensor.arange(4, dtype=dtypes.float).clone().realize()
b = Tensor.arange(4, 8, dtype=dtypes.float).clone().realize()
new_a = a + b # [4, 6, 8, 10]
new_b = a * 2 # [0, 2, 4, 6] -- should use OLD a
a.assign(new_a)
b.assign(new_b)
np.testing.assert_allclose(a.numpy(), [4, 6, 8, 10])
try:
np.testing.assert_allclose(b.numpy(), [0, 2, 4, 6])
except AssertionError:
# TODO: broken now, new_b sees mutated a
np.testing.assert_allclose(b.numpy(), [8, 12, 16, 20])
# eager version
a = Tensor.arange(4, dtype=dtypes.float).clone().realize()
b = Tensor.arange(4, 8, dtype=dtypes.float).clone().realize()
new_a = (a + b).realize()
new_b = (a * 2).realize()
a.assign(new_a).realize()
b.assign(new_b).realize()
np.testing.assert_allclose(a.numpy(), [4, 6, 8, 10])
np.testing.assert_allclose(b.numpy(), [0, 2, 4, 6])
def test_setitem_multiple_disjoint_on_invalid(self):
z = Tensor.invalids(10, dtype="int").realize()
z[2:5] = 2
z[6:7] = 3
z.realize()
self.assertListEqual(z[2:5].tolist(), [2, 2, 2])
self.assertListEqual(z[6:7].tolist(), [3])
class TestWithGrad(unittest.TestCase):
def test_basic_setitem_works(self):
z = Tensor.rand(8, 8)
x = Tensor.rand(8)
z[:3] = x
def test_set_backward(self):
z = Tensor.ones(8, 8)
x = Tensor.rand(8, 8)
z[:] = x
z.sum().backward()
np.testing.assert_allclose(x.grad.numpy(), np.ones((8, 8)))
def test_set_nonleaf_backward(self):
x = Tensor([1.0, 2.0, 3.0, 4.0])
z = x * 2
z[:2] = Tensor([10.0, 20.0])
z.sum().backward()
np.testing.assert_allclose(x.grad.numpy(), [0, 0, 2, 2])
def test_set_overlapping_backward(self):
z = Tensor.zeros(6)
x = Tensor.ones(4).contiguous()
y = Tensor.ones(4) * 2
z[:4] = x
z[2:] = y
z.sum().backward()
np.testing.assert_allclose(x.grad.numpy(), [1, 1, 0, 0])
np.testing.assert_allclose(y.grad.numpy(), np.ones(4))
def test_set_iadd_backward(self):
z = Tensor([1.0, 2.0, 3.0, 4.0])
x = Tensor([10.0, 20.0])
z[:2] += x
z.sum().backward()
np.testing.assert_allclose(z.grad.numpy(), np.ones(4))
np.testing.assert_allclose(x.grad.numpy(), np.ones(2))
def test_set_used_before_setitem(self):
z = Tensor([1.0, 2.0, 3.0, 4.0])
_ = z.sum()
with self.assertRaises(RuntimeError):
z[:2] = Tensor([0.0, 0.0])
def test_setitem_raises_with_unrealized_downstream(self):
x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
_y = x * 2.0
with self.assertRaises(RuntimeError):
x[0] = 99.0
def test_setitem_raises_on_unrealized_compute_base(self):
# y has a compute (unrealized) base; tmp is a view of y. eager: tmp would follow y's mutation. lazy: tmp keeps the old MUL graph.
x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
y = x * 2.0
_tmp = y[:1]
with self.assertRaises(RuntimeError):
y[0] = 99.0
def test_setitem_raises_on_aliased_uop(self):
# two Tensor objects sharing the exact same unrealized uop. setitem on one updates its uop, the other keeps the stale graph reference.
x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
y = x * 2.0
_z = Tensor(y.uop)
with self.assertRaises(RuntimeError):
y[0] = 99.0
class TestSetitemLoop(unittest.TestCase):
def test_arange(self):
N = 10
cmp = Tensor.empty(N)
for i in range(N): cmp[i] = i
self.assertListEqual(Tensor.arange(N).tolist(), cmp.tolist())
if __name__ == '__main__':
unittest.main()

206
artifacts/package_sources/tinygrad/test/backend/test_softmax_fusion.py Normal file
View File

@@ -0,0 +1,206 @@
import unittest
import numpy as np
from tinygrad import Tensor, GlobalCounters, Context, Device
from tinygrad.dtype import DTypeLike, dtypes
from tinygrad.engine.realize import run_linear
from tinygrad.helpers import DEBUG, get_single_element
from test.helpers import check_schedule
def single_kernel_softmax(x_in:Tensor, axis=-1, dtype:DTypeLike|None=None) -> Tensor:
# only support axis =-1
x = x_in.reshape(-1, x_in.shape[-1])
nr_dim, r_dim = x.shape
inp = x.reshape(nr_dim, 1, 1, r_dim).expand(nr_dim, r_dim, 1, r_dim)
imx = x.reshape(nr_dim, 1, r_dim, 1).expand(nr_dim, r_dim, r_dim, r_dim).max(axis=-2, keepdim=True)
m = inp - imx.detach()
if dtype is not None: m = m.cast(dtype)
e = m.exp()
ss = e.sum(axis=-1, keepdim=True)
inp = x.reshape(nr_dim, r_dim, 1, 1)
imx = x.reshape(nr_dim, 1, r_dim, 1).expand(nr_dim, r_dim, r_dim, 1).max(axis=-2, keepdim=True)
m = inp - imx.detach()
if dtype is not None: m = m.cast(dtype)
e = m.exp()
out = e.div(ss).reshape(x_in.shape)
return out
def run_one_schedule_item(out):
linear = out.schedule_linear()
get_single_element(linear.src)
run_linear(linear)
class TestFuse(unittest.TestCase):
def _test_fuse(self, fxn, *args, atol=1e-6, allow_multiple=False, **kwargs):
GlobalCounters.reset()
out_single = fxn(*args, **kwargs)
if not allow_multiple: run_one_schedule_item(out_single)
np_single = out_single.numpy()
GlobalCounters.reset()
np_multi = fxn(*args, **kwargs).numpy()
np.testing.assert_allclose(np_single, np_multi, atol=atol)
@unittest.skip("needs RANGEIFY>1")
def test_fuse_norm(self):
a = Tensor.rand(50,50).realize()
self._test_fuse(lambda a: a / a.mean(axis=1), a)
@unittest.skip("needs RANGEIFY>1")
def test_fuse_argmax(self):
a = Tensor.rand(50,50).realize()
self._test_fuse(lambda a: a.argmax(axis=-1), a)
@unittest.skip("needs RANGEIFY>1")
def test_fuse_softmax(self):
a = Tensor.rand(50,50).realize()
self._test_fuse(lambda a: a.softmax(axis=-1), a)
def test_fuse_gemm_softmax(self):
a = Tensor.rand(50,50).realize()
b = Tensor.rand(50,50).realize()
self._test_fuse(lambda a,b: ((a@b).relu()+a).contiguous().softmax(axis=-1), a,b, allow_multiple=True)
@unittest.skipUnless(dtypes.float16 in Device[Device.DEFAULT].renderer.supported_dtypes(), f"no float16 on {Device.DEFAULT}")
@unittest.skip("needs RANGEIFY>1")
def test_fuse_softmax_dtype(self):
a = Tensor.rand(50,50).realize()
self._test_fuse(lambda a: a.softmax(axis=-1, dtype='half'), a, atol=3e-4)
def test_fuse_arange_eye(self):
self._test_fuse(lambda: (Tensor.arange(10).reshape(10,1).expand(10,10) == Tensor.arange(10).reshape(1,10).expand(10,10)).clone())
@unittest.skip("needs RANGEIFY>1")
def test_double_gemm(self):
N = 32
with Context(TRACK_MATCH_STATS=0, DEBUG=0):
a = (Tensor.rand(N,N)-0.5).realize()
b = (Tensor.rand(N,N)-0.5).realize()
c = (Tensor.rand(N,N)-0.5).realize()
self._test_fuse(lambda a,b,c: a@b@c, a, b, c, atol=1e-5)
def test_embedding(self):
with Context(TRACK_MATCH_STATS=0, DEBUG=0):
vocab_sz = 123
embed_sz = 16
weight = (Tensor.rand(vocab_sz, embed_sz)-0.5).realize()
a = Tensor([1, 1, 2, 3]).realize()
def embedding(idx:Tensor):
arange = Tensor.arange(vocab_sz).unsqueeze(-1)
big_shp = idx.shape + (vocab_sz, embed_sz)
arange, vals = arange.expand(big_shp), weight.expand(big_shp)
idx = idx.reshape(idx.shape+(1, 1)).expand(big_shp)
return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
self._test_fuse(embedding, a, atol=1e-5)
@unittest.skip("needs RANGEIFY>1")
def test_attention_kernel_count(self):
wq = Tensor.empty(32, 32)
wk = Tensor.empty(32, 32)
wv = Tensor.empty(32, 32)
x = Tensor.empty(2, 100, 32)
q = (x @ wq).contiguous()
k = (x @ wk).contiguous()
v = (x @ wv).contiguous()
attn = q.scaled_dot_product_attention(k, v)
check_schedule(attn, 4) # 3 matmul and 1 attention
@unittest.skip("needs RANGEIFY>1")
def test_flash_attention(self):
BS = 4
HEADS = 2
MATDIM = 16
EMB = 8
with Context(TRACK_MATCH_STATS=0, DEBUG=0):
q = Tensor.randn(BS, HEADS, MATDIM, EMB).realize()
k = Tensor.randn(BS, HEADS, MATDIM, EMB).realize()
v = Tensor.randn(BS, HEADS, MATDIM, EMB).realize()
# TODO: OPT is breaking things. NOOPT isn't linearizing
with Context(NOOPT=1):
self._test_fuse(Tensor.scaled_dot_product_attention, q, k, v, atol=1e-5)
def test_mismatch_reduce(self):
a = Tensor.ones(16, 10).contiguous().realize()
b = Tensor.ones(16, 20).contiguous().realize()
c = (a.sum(axis=1) + b.sum(axis=1))
self.assertListEqual(c.tolist(), [30]*16)
@unittest.skipUnless(Device.DEFAULT == "METAL", "METAL TC")
def test_fuse_and_tc_opt(self):
A = Tensor.randn(8, 8).realize()
B = Tensor.randn(8, 8).realize()
C = Tensor.ones(1, 8, 8).pad(((1,1), None, None),).sum(0)
out = (C + (A @ B))
out.realize()
class TestSoftmaxFusion(unittest.TestCase):
@classmethod
def setUpClass(cls):
with Context(TRACK_MATCH_STATS=0): cls.test = Tensor.rand(32, 10).contiguous().realize()
def setUp(self):
GlobalCounters.reset()
def test_norm(self):
print("*** norm ***")
with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
# NOTE: there's an implied expand on the mean here
sout = self.test / self.test.mean(-1, keepdim=True)
sout.realize()
print("*** single kernel norm ***")
with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
inp = self.test.reshape(32, 10, 1)
div = self.test.reshape(32, 1, 10).expand(32, 10, 10).mean(axis=-1, keepdim=True)
out = (inp / div).reshape(32, 10)
out.realize()
np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
def test_softmax(self):
# this is the softmax from scaled_dot_product_attention
# it becomes 3 kernels
print("*** softmax ***")
with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
sout = self.test.softmax(-1)
sout.realize()
print("*** single kernel softmax ***")
with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
out = single_kernel_softmax(self.test)
out.realize()
np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
@unittest.skip("needs RANGEIFY>1")
def test_auto_softmax(self):
print("*** softmax ***")
with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
sout = self.test.softmax(-1)
sout.realize()
print("*** auto single kernel softmax ***")
with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
out = self.test.contiguous().softmax(-1)
run_one_schedule_item(out)
np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
def test_softmax_bw(self):
print("*** softmax bw ***")
with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
self.test.softmax(-1).sum().backward()
sg = self.test.grad.realize()
self.test.grad = None
print("*** single kernel softmax bw ***")
with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
single_kernel_softmax(self.test).sum().backward()
g = self.test.grad.realize()
np.testing.assert_allclose(sg.numpy(), g.numpy(), atol=1e-7)
if __name__ == '__main__':
unittest.main()

59
artifacts/package_sources/tinygrad/test/backend/test_stunning.py Normal file
View File

@@ -0,0 +1,59 @@
import unittest
from tinygrad import nn, Tensor, Variable, Context, Device
from tinygrad.helpers import trange
class Model:
def __init__(self): self.layer = nn.Linear(28*28, 10)
def __call__(self, x:Tensor) -> Tensor: return self.layer(x.flatten(1))
class TestStunning(unittest.TestCase):
def test_indexing_variable(self):
a = Tensor.arange(100*10).reshape(100, 10).contiguous()
# index without variable
nv = a[12].tolist()
# index with variable
vi = Variable('i', 0, a.shape[0]-1)
wv = a[vi.bind(12)].tolist()
self.assertListEqual(nv, wv)
def test_indexing_two_bind(self):
a = Tensor.arange(100*10).reshape(100, 10).contiguous()
nv = a[12].cat(a[76]).tolist()
vi = Variable('i', 0, a.shape[0]-1)
with self.assertRaisesRegex(RuntimeError, "bind mismatch on"):
wv = a[vi.bind(12)].cat(a[vi.bind(76)]).tolist()
self.assertListEqual(nv, wv)
@unittest.skipIf(Device.DEFAULT in {"WEBGPU", "NV", "CUDA"}, "Too many buffers / too slow")
@unittest.skip("This is binding a Variable to two different values")
def test_simple_train(self, steps=6, bs=4, adam=True):
X_train, Y_train, _, _ = nn.datasets.mnist()
model = Model()
if adam: opt = nn.optim.Adam(nn.state.get_parameters(model))
else: opt = nn.optim.SGD(nn.state.get_parameters(model), momentum=0.1)
samples = Tensor.randint(steps, bs, high=X_train.shape[0])
Y_train = Y_train.one_hot(10)
X_samp, Y_samp = X_train[samples], Y_train[samples]
vi = Variable('i', 0, samples.shape[0]-1)
with Context(SPLIT_REDUCEOP=0):
with Context(TRAINING=1):
losses = []
for i in range(samples.shape[0]):
vib = vi.bind(i)
opt.zero_grad()
pred = model(X_samp[vib].realize())
loss = (pred - Y_samp[vib]).square().mean()
losses.append(loss.backward())
opt.schedule_step()
#losses = Tensor.stack(*losses)
# run
for i in (t:=trange(len(losses))): t.set_description(f"loss: {losses[i].item():6.2f}")
if __name__ == '__main__':
unittest.main()

183
artifacts/package_sources/tinygrad/test/backend/test_subbuffer.py Normal file
View File

@@ -0,0 +1,183 @@
import unittest
from tinygrad import Device, dtypes, Tensor
from tinygrad.device import Buffer
from tinygrad.helpers import Context, DEV
from test.helpers import needs_second_gpu
@unittest.skipIf(Device.DEFAULT in {"WEBGPU", "CL"}, "subbuffer not supported")
class TestSubBuffer(unittest.TestCase):
def setUp(self):
self.buf = Buffer(Device.DEFAULT, 10, dtypes.uint8, initial_value=bytes(range(10)))
self.buf_unalloc = Buffer(Device.DEFAULT, 10, dtypes.uint8)
def test_subbuffer(self):
vbuf = self.buf.view(2, dtypes.uint8, offset=3).ensure_allocated()
tst = vbuf.as_memoryview().tolist()
assert tst == [3, 4]
def test_subbuffer_cast(self):
# NOTE: bitcast depends on endianness
vbuf = self.buf.view(2, dtypes.uint16, offset=3).ensure_allocated()
tst = vbuf.as_memoryview().cast("H").tolist()
assert tst == [3|(4<<8), 5|(6<<8)]
def test_subbuffer_double(self):
vbuf = self.buf.view(4, dtypes.uint8, offset=3).ensure_allocated()
vvbuf = vbuf.view(2, dtypes.uint8, offset=1).ensure_allocated()
tst = vvbuf.as_memoryview().tolist()
assert tst == [4, 5]
def test_subbuffer_len(self):
vbuf = self.buf.view(5, dtypes.uint8, 2).ensure_allocated()
mv = vbuf.as_memoryview()
assert len(mv) == 5
mv = vbuf.as_memoryview(allow_zero_copy=True)
assert len(mv) == 5
def test_subbuffer_used(self):
t = Tensor.arange(0, 10, dtype=dtypes.uint8).clone().realize()
vt = t[2:4].realize()
out = (vt + 100).tolist()
assert out == [102, 103]
@needs_second_gpu
@unittest.skipIf(Device.DEFAULT not in {"CUDA", "NV", "AMD"} or DEV.interface.startswith("MOCK"), "only NV, AMD, CUDA")
def test_subbuffer_transfer(self):
t = Tensor.arange(0, 10, dtype=dtypes.uint8).clone().realize()
vt = t[2:5].contiguous().realize()
out = vt.to(f"{Device.DEFAULT}:1").realize().tolist()
assert out == [2, 3, 4]
def test_subbuffer_deallocate(self):
with Context(LRU=0):
vbuf = self.buf.view(2, dtypes.uint8, offset=3).ensure_allocated()
self.buf.deallocate()
vbuf.deallocate()
# Allocate a fake one on the same place
_ = Buffer(Device.DEFAULT, 10, dtypes.uint8).ensure_allocated()
self.buf.ensure_allocated()
self.buf.copy_from(Buffer("PYTHON", 10, dtypes.uint8, opaque=memoryview(bytearray(range(10, 20)))))
vbuf.ensure_allocated()
tst = vbuf.as_memoryview().tolist()
assert tst == [13, 14]
def test_subbuffer_is_allocated(self):
buf = self.buf_unalloc
sub_buf = buf.view(3, dtypes.uint8, offset=4)
self.assertFalse(buf.is_allocated())
self.assertFalse(buf.is_initialized())
self.assertFalse(sub_buf.is_allocated())
self.assertFalse(sub_buf.is_initialized())
# base buffer alloc
buf.allocate()
self.assertTrue(buf.is_allocated())
self.assertTrue(buf.is_initialized())
self.assertTrue(sub_buf.is_allocated())
self.assertFalse(sub_buf.is_initialized())
# sub buffer alloc
sub_buf.allocate()
self.assertTrue(sub_buf.is_initialized())
# sub buffer dealloc
sub_buf.deallocate()
self.assertTrue(buf.is_allocated())
self.assertTrue(buf.is_initialized())
self.assertTrue(sub_buf.is_allocated())
self.assertFalse(sub_buf.is_initialized())
# base buffer dealloc
buf.deallocate()
self.assertFalse(buf.is_allocated())
self.assertFalse(buf.is_initialized())
self.assertFalse(sub_buf.is_allocated())
self.assertFalse(sub_buf.is_initialized())
# sub buffer alloc
sub_buf.ensure_allocated()
self.assertTrue(buf.is_allocated())
self.assertTrue(buf.is_initialized())
self.assertTrue(sub_buf.is_allocated())
self.assertTrue(sub_buf.is_initialized())
def test_subbuffer_copy_in_out(self):
sub_buf = self.buf.view(3, dtypes.uint8, offset=3).ensure_allocated() # [3:6]
data_out_sub = bytearray([0]*3)
data_out_sub[:] = sub_buf.as_memoryview()
assert data_out_sub == bytearray(range(3, 6))
sub_buf.copy_from(Buffer("PYTHON", 3, dtypes.uint8, opaque=memoryview(bytearray(range(3)))))
assert sub_buf.as_memoryview().tolist() == list(range(3))
assert self.buf.as_memoryview().tolist()[3:6] == list(range(3))
data_out_sub[:] = sub_buf.as_memoryview()
assert data_out_sub == bytearray(range(3))
data_out_base = bytearray([0]*10)
data_out_base[:] = self.buf.as_memoryview()
assert data_out_base[0:3] == bytearray(range(0, 3))
assert data_out_base[3:6] == data_out_sub
assert data_out_base[6:10] == bytearray(range(6, 10))
def test_subbuffer_copy_in_out_view_of_view(self):
view1 = self.buf.view(7, dtypes.uint8, offset=2).ensure_allocated() # [2:9]
view2 = view1.view(3, dtypes.uint8, offset=2).ensure_allocated() # [4:7]
self.assertTrue(view1.is_allocated())
self.assertTrue(view2.is_allocated())
data_in = bytearray([7, 8, 9])
view2.copy_from(Buffer("PYTHON", 3, view2.dtype, opaque=memoryview(data_in)))
data_out_v2 = bytearray([0]*3)
data_out_v2[:] = view2.as_memoryview()
assert data_in == data_out_v2
expected_base_data = memoryview(bytearray(range(10)))
expected_base_data[4:7] = data_in
data_out_base = bytearray([0]*10)
data_out_base[:] = self.buf.as_memoryview()
assert expected_base_data == data_out_base
def test_subbuffer_alloc(self):
sub_buf = self.buf.view(4, dtypes.int8, offset=3)
sub_buf.allocate()
sub_buf.copy_from(Buffer("PYTHON", 4, dtypes.int8, opaque=memoryview(bytearray(range(10, 14)))))
assert self.buf.as_memoryview().tolist()[3:7] == sub_buf.as_memoryview().tolist()
sub_buf = self.buf_unalloc.view(4, dtypes.int8, offset=3)
sub_buf.allocate()
sub_buf.copy_from(Buffer("PYTHON", 4, dtypes.int8, opaque=memoryview(bytearray(range(10, 14)))))
assert self.buf_unalloc.as_memoryview().tolist()[3:7] == sub_buf.as_memoryview().tolist()
def test_subbuffer_dealloc(self):
sub_buf = self.buf.view(4, dtypes.int8, offset=3).ensure_allocated()
sub_buf.deallocate()
assert self.buf.as_memoryview().tolist() == list(range(10))
def test_subbuffer_double_dealloc(self):
sub_buf = self.buf.view(3, dtypes.uint8, offset=4).ensure_allocated()
self.buf.deallocate()
with self.assertRaises(AssertionError):
self.buf.deallocate()
sub_buf.deallocate()
with self.assertRaises(AssertionError):
sub_buf.deallocate()
def test_subbuffer_uaf(self):
sub_buf = self.buf.view(4, dtypes.int8, offset=3).ensure_allocated()
assert self.buf.as_memoryview().tolist(), list(range(10))
sub_buf.deallocate()
with self.assertRaises(AssertionError):
sub_buf.as_memoryview().tolist()
assert self.buf.as_memoryview().tolist(), list(range(10))
sub_buf = self.buf.view(4, dtypes.int8, offset=3).ensure_allocated()
assert sub_buf.as_memoryview().tolist(), list(range(3, 7))
self.buf.deallocate()
with self.assertRaises(AssertionError):
sub_buf.as_memoryview().tolist()
if __name__ == '__main__':
unittest.main()

332
artifacts/package_sources/tinygrad/test/backend/test_symbolic_jit.py Normal file
View File

@@ -0,0 +1,332 @@
import unittest
from test.helpers import assert_jit_cache_len
from tinygrad import Variable, Tensor, TinyJit
from tinygrad.engine.jit import JitError
import numpy as np
class TestSymbolicJit(unittest.TestCase):
def test_plus1(self):
def f(a): return (a+1).realize()
jf = TinyJit(f)
a = Tensor.rand(3, 10)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = jf(a[:, :vi])[:3, :i].numpy()
expected = f(a[:, :i]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_plus1_pad(self):
# TODO: without contiguous, the pad is not captured in jit
def f(a): return (a+1).pad((None, (0, 10-a.shape[1]))).contiguous().realize()
jf = TinyJit(f)
a = Tensor.rand(3, 10)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = jf(a[:, :vi]).numpy()
expected = f(a[:, :i]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_add(self):
def f(a, b): return (a+b).realize()
jf = TinyJit(f)
a = Tensor.rand(3, 10)
b = Tensor.rand(3, 10)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = jf(a[:, :vi], b[:, :vi])
symbolic = symbolic[:3, :i].numpy()
expected = f(a[:, :i], b[:, :i]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_matmul(self):
def f(a, b): return (a@b).realize()
jf = TinyJit(f)
a = Tensor.rand(3, 10)
b = Tensor.rand(10, 5)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = jf(a[:, :vi], b[:vi, :]).numpy()
expected = f(a[:, :i], b[:i, :]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_mixed_with_no_symbol_kernel(self):
def f(a, b):
s = (a@b).realize()
s = (s+s).realize() # this one does not have symbols in input
return s
jf = TinyJit(f)
a = Tensor.rand(3, 10)
b = Tensor.rand(10, 5)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = jf(a[:, :vi], b[:vi, :]).numpy()
expected = f(a[:, :i], b[:i, :]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 2)
def test_attention(self):
def f(q, k, v): return Tensor.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).realize()
jf = TinyJit(f)
q = Tensor.rand(2, 1, 4, 8)
k = Tensor.rand(2, 10, 4, 8)
v = Tensor.rand(2, 10, 4, 8)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = jf(q, k[:, :vi], v[:, :vi])[:2, :4, :1, :8].numpy()
expected = f(q, k[:, :i], v[:, :i]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 5)
def test_cat_dim0(self):
def f(a, b): return a.cat(b, dim=0).realize()
jf = TinyJit(f)
a = Tensor.rand(10, 3)
b = Tensor.rand(2, 3)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = jf(a[:vi], b)[:i+2, :3].numpy()
expected = f(a[:i], b).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_cat_dim1(self):
def f(a, b): return a.cat(b, dim=1).realize()
jf = TinyJit(f)
a = Tensor.rand(3, 10)
b = Tensor.rand(3, 2)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = jf(a[:, :vi], b)[:3, :i+2].numpy()
expected = f(a[:, :i], b).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_cat_dim0_two_vars(self):
def f(a, b): return a.cat(b, dim=0).realize()
jf = TinyJit(f)
a = Tensor.rand(10, 3)
b = Tensor.rand(10, 3)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
symbolic = jf(a[:vi], b[:vj])[:i+j, :3].numpy()
expected = f(a[:i], b[:j]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_cat_dim1_two_vars(self):
def f(a, b): return a.cat(b, dim=1).realize()
jf = TinyJit(f)
a = Tensor.rand(3, 10)
b = Tensor.rand(3, 10)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
symbolic = jf(a[:, :vi], b[:, :vj])[:3, :i+j].numpy()
expected = f(a[:, :i], b[:, :j]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_two_vars_plus1_ij(self):
def f(a, b): return (a@b+1).realize()
jf = TinyJit(f)
a = Tensor.rand(10, 3)
b = Tensor.rand(3, 10)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
symbolic = jf(a[:vi, :], b[:, :vj])[:i, :j].numpy()
expected = f(a[:i, :], b[:, :j]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_two_vars_plus1_ji(self):
def f(a, b): return (a@b+1).realize()
jf = TinyJit(f)
a = Tensor.rand(10, 3)
b = Tensor.rand(3, 10)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
symbolic = jf(a[:vj, :], b[:, :vi])[:j, :i].numpy()
expected = f(a[:j, :], b[:, :i]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_jit_symbolic_shape_mismatch(self):
@TinyJit
def add(a, b): return (a+b).realize()
a = Tensor.rand(3, 10)
b = Tensor.rand(3, 10)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
add(a[:, :vi], b[:, :vi])
vi2 = Variable("i", 1, 10).bind(7)
a = Tensor.rand(3, 7)[:, :vi2]
bad = Tensor.rand(4, 7)[:, :vi2]
with self.assertRaises(JitError):
add(a, bad)
def test_shrink(self):
# shrink is a movement, so we pair it with a simple function to test the JIT interaction
def f(a): return (a+1).realize()
jf = TinyJit(f)
a = Tensor.rand(7, 11)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = a.shrink(((3,5),(vi,vi+2)))
symbolic = jf(symbolic).numpy()
expected = f(a.shrink(((3,5),(i,i+2)))).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_slice(self):
# slice is a movement, so we pair it with a simple function to test the JIT interaction
def f(a): return (a+1).realize()
jf = TinyJit(f)
a = Tensor.rand(7, 11)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = a[3:5, vi:vi+2]
symbolic = jf(symbolic).numpy()
expected = f(a[3:5, i:i+2]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_slice_var_shape(self):
def f(a): return (a+1).realize()
jf = TinyJit(f)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
a = Tensor.ones(vi, 11).contiguous()
symbolic = a[:, 1:2]
symbolic = jf(symbolic)[:i, :1].numpy()
expected = f(a[:i, :][:, 1:2]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_ones_sum(self):
def f(a): return a.sum().realize()
jf = TinyJit(f)
t = Tensor.ones(10).contiguous()
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = jf(t[:vi]).item()
expected = f(t[:i]).item()
np.testing.assert_equal(symbolic, expected)
def test_mean(self):
def f(a): return a.mean().realize()
def f0(a): return a.mean(0).realize()
def f1(a): return a.mean(1).realize()
jf = TinyJit(f)
jf0 = TinyJit(f0)
jf1 = TinyJit(f1)
a = Tensor.rand(10, 3)
b = Tensor.rand(10, 3)
c = Tensor.rand(10, 3)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
# axis = None
symbolic = jf(a[:vi]).numpy()
expected = a[:i].mean().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# axis = 0
symbolic = jf0(b[:vi]).numpy()
expected = b[:i].mean(0).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# axis = 1
symbolic = jf1(c[:vi])[:i].numpy()
expected = c[:i].mean(1).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_mean_2d(self):
def f(a): return a.mean().realize()
def f0(a): return a.mean(0).realize()
def f1(a): return a.mean(1).realize()
jf = TinyJit(f)
jf0 = TinyJit(f0)
jf1 = TinyJit(f1)
a = Tensor.rand(10, 10)
b = Tensor.rand(10, 10)
c = Tensor.rand(10, 10)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
# axis = None
symbolic = jf(a[:vi, :vj]).numpy()
expected = a[:i, :j].mean().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# axis = 0
symbolic = jf0(b[:vi, :vj])[:j].numpy()
expected = b[:i, :j].mean(0).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# axis = 1
symbolic = jf1(c[:vi, :vj])[:i].numpy()
expected = c[:i, :j].mean(1).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_var(self):
def f(a): return a.var().realize()
def f0(a): return a.var(0).realize()
def f1(a): return a.var(1).realize()
jf = TinyJit(f)
jf0 = TinyJit(f0)
jf1 = TinyJit(f1)
a = Tensor.rand(10, 3)
b = Tensor.rand(10, 3)
c = Tensor.rand(10, 3)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
# axis = None
symbolic = jf(a[:vi]).numpy()
expected = a[:i].var().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# axis = 0
symbolic = jf0(b[:vi]).numpy()
expected = b[:i].var(0).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# axis = 1
symbolic = jf1(c[:vi])[:i].numpy()
expected = c[:i].var(1).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_var_2d(self):
def f(a): return a.var().realize()
def f0(a): return a.var(0).realize()
def f1(a): return a.var(1).realize()
jf = TinyJit(f)
jf0 = TinyJit(f0)
jf1 = TinyJit(f1)
a = Tensor.rand(10, 10)
b = Tensor.rand(10, 10)
c = Tensor.rand(10, 10)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
# axis = None
symbolic = jf(a[:vi, :vj]).numpy()
expected = a[:i, :j].var().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# axis = 0
symbolic = jf0(b[:vi, :vj])[:j].numpy()
expected = b[:i, :j].var(0).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# axis = 1
symbolic = jf1(c[:vi, :vj])[:i].numpy()
expected = c[:i, :j].var(1).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
if __name__ == '__main__':
unittest.main()

367
artifacts/package_sources/tinygrad/test/backend/test_symbolic_ops.py Normal file
View File

@@ -0,0 +1,367 @@
import unittest
from tinygrad import Tensor, Variable, GlobalCounters, Context
from tinygrad.uop.ops import sym_infer
from tinygrad.dtype import dtypes
from examples.gpt2 import Attention
import numpy as np
class TestSymbolicOps(unittest.TestCase):
def test_negative_slice(self):
a = Tensor.rand(3, 10, 4)
for i in range(3, 10):
vi = Variable("i", 1, 10).bind(i)
# negative int bounds against a symbolic dim must resolve against the size, like slice.indices
np.testing.assert_allclose(a[:, :vi][:, -3:-1].numpy(), a[:, :i][:, -3:-1].numpy(), atol=1e-6, rtol=1e-6)
np.testing.assert_allclose(a[:, :vi][:, -1:].numpy(), a[:, :i][:, -1:].numpy(), atol=1e-6, rtol=1e-6)
np.testing.assert_allclose(a[:, :vi][:, -1].numpy(), a[:, :i][:, -1].numpy(), atol=1e-6, rtol=1e-6)
def test_plus1(self):
def f(a): return (a+1).realize()
a = Tensor.rand(3, 10)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = f(a[:, :vi])[:3, :i].numpy()
expected = f(a[:, :i]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_plus1_pad(self):
def f(a): return (a+1).pad((None, (0, 10-a.shape[1]))).realize()
a = Tensor.rand(3, 10)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = f(a[:, :vi]).numpy()
expected = f(a[:, :i]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_add(self):
def f(a, b): return (a+b).realize()
a = Tensor.rand(3, 10)
b = Tensor.rand(3, 10)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = f(a[:, :vi], b[:, :vi])[:, :i].numpy()
expected = f(a[:, :i], b[:, :i]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_matmul(self):
def f(a, b): return (a@b).realize()
a = Tensor.rand(3, 10)
b = Tensor.rand(10, 5)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = f(a[:, :vi], b[:vi, :]).numpy()
expected = f(a[:, :i], b[:i, :]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_attention(self, dropout_p=0.0, imin=1, imax=5, use_symbolic=True):
def f(q, k, v): return Tensor.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), dropout_p=dropout_p).realize()
q = Tensor.rand(2, 1, 4, 8)
k = Tensor.rand(2, 10, 4, 8)
v = Tensor.rand(2, 10, 4, 8)
for i in range(imin, imax):
vi = Variable("i", 1, 10).bind(i) if use_symbolic else i
Tensor.realize(q, k, v)
GlobalCounters.reset()
symbolic = f(q, k[:, :vi, :, :], v[:, :vi, :, :])[:2, :4, :1, :8].numpy()
expected = f(q, k[:, :i, :, :], v[:, :i, :, :]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_attention_cmp_symbolic(self):
# symbolic isn't seeing if i == i, so it's not putting them on the same axis
self.test_attention(imin=4, imax=5, use_symbolic=False)
self.test_attention(imin=4, imax=5, use_symbolic=True)
def test_attention_training(self):
with Context(TRAINING=1):
self.test_attention(dropout_p=0.0)
with self.assertRaises(ValueError):
# symbolic shape dropout is not supported
self.test_attention(dropout_p=0.5)
def test_sdpa_symbolic_seq_len(self):
# symbolic seq_len on all of q/k/v (dim -2 after transpose)
q = Tensor.rand(2, 10, 4, 8)
k = Tensor.rand(2, 10, 4, 8)
v = Tensor.rand(2, 10, 4, 8)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
Tensor.realize(q, k, v)
symbolic = q[:, :vi].transpose(1, 2).scaled_dot_product_attention(
k[:, :vi].transpose(1, 2), v[:, :vi].transpose(1, 2)).realize()[:2, :4, :i, :8].numpy()
expected = q[:, :i].transpose(1, 2).scaled_dot_product_attention(
k[:, :i].transpose(1, 2), v[:, :i].transpose(1, 2)).realize().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_sdpa_symbolic_seq_len_query_only(self):
# symbolic seq_len on query only (dim -2 after transpose)
q = Tensor.rand(2, 10, 4, 8)
k = Tensor.rand(2, 5, 4, 8)
v = Tensor.rand(2, 5, 4, 8)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
Tensor.realize(q, k, v)
symbolic = q[:, :vi].transpose(1, 2).scaled_dot_product_attention(
k.transpose(1, 2), v.transpose(1, 2)).realize()[:2, :4, :i, :8].numpy()
expected = q[:, :i].transpose(1, 2).scaled_dot_product_attention(
k.transpose(1, 2), v.transpose(1, 2)).realize().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_sdpa_symbolic_batch(self):
# symbolic batch dim (dim 0)
q = Tensor.rand(10, 4, 3, 8)
k = Tensor.rand(10, 4, 3, 8)
v = Tensor.rand(10, 4, 3, 8)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
Tensor.realize(q, k, v)
symbolic = q[:vi].scaled_dot_product_attention(k[:vi], v[:vi]).realize()[:i, :4, :3, :8].numpy()
expected = q[:i].scaled_dot_product_attention(k[:i], v[:i]).realize().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_sdpa_symbolic_heads(self):
# symbolic heads dim (dim -3)
q = Tensor.rand(2, 10, 3, 8)
k = Tensor.rand(2, 10, 3, 8)
v = Tensor.rand(2, 10, 3, 8)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
Tensor.realize(q, k, v)
symbolic = q[:, :vi].scaled_dot_product_attention(k[:, :vi], v[:, :vi]).realize()[:2, :i, :3, :8].numpy()
expected = q[:, :i].scaled_dot_product_attention(k[:, :i], v[:, :i]).realize().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_attention_pos_0_sz_0(self):
Attention(128, 8)(Tensor.ones(1, 0, 128), Variable("start_pos", 0, 128).bind(0), None)
def test_attention_pos_0_sz_1(self):
Attention(128, 8)(Tensor.ones(1, 1, 128), Variable("start_pos", 0, 128).bind(0), None)
def test_attention_pos_0_sz_2(self):
Attention(128, 8)(Tensor.ones(1, 2, 128), Variable("start_pos", 0, 128).bind(0), None)
def test_cat_dim0(self):
def f(a, b): return a.cat(b, dim=0).realize()
a = Tensor.rand(10, 3)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
b = Tensor.rand(2, 3)
symbolic = f(a[:vi, :], b)[:i+2, :3].numpy()
expected = f(a[:i, :], b).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_cat_dim1(self):
def f(a, b): return a.cat(b, dim=1).realize()
a = Tensor.rand(3, 10)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
b = Tensor.rand(3, 2)
symbolic = f(a[:, :vi], b)[:3, :i+2].numpy()
expected = f(a[:, :i], b).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_cat_dim0_two_vars(self):
def f(a, b): return a.cat(b, dim=0).realize()
a = Tensor.rand(10, 3)
b = Tensor.rand(10, 3)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
symbolic = f(a[:vi, :], b[:vj, :])[:i+j, :3].numpy()
expected = f(a[:i, :], b[:j, :]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_cat_dim1_two_vars(self):
def f(a, b): return a.cat(b, dim=1).realize()
a = Tensor.rand(3, 10)
b = Tensor.rand(3, 10)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
symbolic = f(a[:, :vi], b[:, :vj])[:3, :i+j].numpy()
expected = f(a[:, :i], b[:, :j]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_two_vars_plus1_ij(self):
def f(a, b): return (a@b+1).realize()
a = Tensor.rand(10, 3).realize()
b = Tensor.rand(3, 10).realize()
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
symbolic = f(a[:vi, :], b[:, :vj])[:i, :j].numpy()
expected = f(a[:i, :], b[:, :j]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_two_vars_plus1_ji(self):
# reverse the order of variables
def f(a, b): return (a@b+1).realize()
a = Tensor.rand(10, 3).realize()
b = Tensor.rand(3, 10).realize()
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
symbolic = f(a[:vj, :], b[:, :vi])[:j, :i].numpy()
expected = f(a[:j, :], b[:, :i]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_invalid_symbolic_reshape(self):
a = Tensor.rand(30)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
# Cannot reshape into symbolic from non-symbolic
with self.assertRaises(ValueError): a.reshape((3, vi))
def test_shrink(self):
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
a = Tensor.rand(7, 11)
symbolic = a.shrink(((3,5),(vi,vi+2)))
symbolic = symbolic.numpy()
expected = a.shrink(((3,5),(i,i+2))).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_slice(self):
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
a = Tensor.rand(7, 11)
symbolic = a[3:5, vi:vi+2]
symbolic = symbolic.numpy()
expected = a[3:5, i:i+2].numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_slice_no_start(self):
a = Tensor.rand(7, 11)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = a[3:5, :vi:1][:2, :i].numpy()
expected = a[3:5, :i:1].numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_expand_padded(self):
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
a = Tensor(1).unsqueeze(0).pad((0, 1)).unsqueeze(0)
symbolic = a.expand(vi, 2)[:i, :2].numpy()
expected = a.expand(i, 2).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_slice_var_shape(self):
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
a = Tensor.ones(vi, 11).contiguous()
symbolic = a[:, 1:2][:i, :1].numpy()
expected = Tensor.ones(i, 11)[:, 1:2].numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_ones_sum(self):
t = Tensor.ones(10).contiguous()
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = t[:vi].sum().item()
expected = t[:i].sum().item()
np.testing.assert_equal(symbolic, expected)
def test_mean(self):
a = Tensor.rand(10, 3)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
for axis in [None, 0, 1]:
expected = a[:i].mean(axis).numpy()
symbolic = a[:vi].mean(axis)
if axis is None:
symbolic = symbolic.numpy()
else:
symbolic = symbolic[:expected.shape[0]].numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_mean_2d(self):
a = Tensor.rand(10, 10)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
for axis in [None, 0, 1]:
expected = a[:i, :j].mean(axis).numpy()
symbolic = a[:vi, :vj].mean(axis)
if axis is None:
symbolic = symbolic.numpy()
else:
symbolic = symbolic[:expected.shape[0]].numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_var(self):
a = Tensor.rand(10, 3)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
for axis in [None, 0, 1]:
expected = a[:i].var(axis).numpy()
symbolic = a[:vi].var(axis)
if axis is None:
symbolic = symbolic.numpy()
else:
symbolic = symbolic[:expected.shape[0]].numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_var_2d(self):
a = Tensor.rand(10, 10)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
for axis in [None, 0, 1]:
expected = a[:i, :j].var(axis).numpy()
symbolic_result = a[:vi, :vj].var(axis)
if axis is None:
symbolic = symbolic_result.numpy()
else:
symbolic = symbolic_result[:expected.shape[0]].numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_bitcast_down(self):
a = Tensor.rand(10, 3)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
expected = a[:i].bitcast(dtypes.uint8).numpy()
symbolic_result = a[:vi].bitcast(dtypes.uint8)
if len(expected.shape) == 2:
symbolic = symbolic_result[:expected.shape[0], :expected.shape[1]].numpy()
else:
symbolic = symbolic_result[:].numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=0)
def test_bitcast_up(self):
a = Tensor.rand(10, 4)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
expected = a[:i].bitcast(dtypes.uint64).numpy()
symbolic_result = a[:vi].bitcast(dtypes.uint64)
if len(expected.shape) == 2:
symbolic = symbolic_result[:expected.shape[0], :expected.shape[1]].numpy()
else:
symbolic = symbolic_result[:].numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=0)
def test_conv2d_ceildiv_edge_case(self):
# tests symbolic ceildiv in conv2d output shape calculation
# val=79 triggers the edge case where old ceildiv simplifies incorrectly: old gives floor=12, correct ceildiv=13
v = Variable('v', 11, 100)
val = 79
x_full = Tensor.randn(1, 8, 100)
weight = Tensor.randn(16, 8, 12)
# symbolic version
result = x_full[:, :, :v.bind(val)].conv2d(weight=weight, groups=1, stride=6, dilation=1, padding=(3, 3))
var_val = {v.expr: val}
shape = tuple(sym_infer(s, var_val) for s in result.shape)
self.assertEqual(shape, (1, 16, 13))
# concrete version for comparison
expected = x_full[:, :, :val].conv2d(weight=weight, groups=1, stride=6, dilation=1, padding=(3, 3))
np.testing.assert_allclose(result[:, :, :13].numpy(), expected.numpy(), atol=1e-5, rtol=1e-5)
if __name__ == '__main__':
unittest.main()

777
artifacts/package_sources/tinygrad/test/backend/test_tensor.py Normal file
View File

@@ -0,0 +1,777 @@
import numpy as np
import torch
import unittest, copy, mmap, random, math, array
from tinygrad import Tensor, Device, dtypes, nn, Context
from tinygrad.helpers import getenv, temp, mv_address
from extra.gradcheck import numerical_jacobian, jacobian, gradcheck
from hypothesis import given, settings, strategies as strat
from tinygrad.dtype import DTYPES_DICT
from tinygrad.uop.ops import UOp
settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
settings.load_profile("my_profile")
x_init = np.random.randn(1,3).astype(np.float32)
U_init = np.random.randn(3,3).astype(np.float32)
V_init = np.random.randn(3,3).astype(np.float32)
W_init = np.random.randn(3,3).astype(np.float32)
m_init = np.random.randn(1,3).astype(np.float32)
gradient = np.random.randn(1,3).astype(np.float32)
class TestTinygrad(unittest.TestCase):
def test_zerodim_initialization(self):
self.assertEqual(Tensor(55).shape, ())
self.assertEqual(Tensor(3.14).shape, ())
def test_deviceless_const_construct_device_repr(self):
t = Tensor(UOp.const(2.0).cast(dtypes.float))
self.assertIsNone(t.uop.device)
self.assertIsNone(t.device)
self.assertIn("<UOp None", repr(t))
def test_deviceless_const_realize_noop(self):
t = Tensor(UOp.const(2.0).cast(dtypes.float))
uop = t.uop
t.realize()
self.assertIs(t.uop, uop)
self.assertIsNone(t.uop.device)
def test_plus_equals(self):
a = Tensor.randn(10,10)
b = Tensor.randn(10,10)
c = a + b
val1 = c.numpy()
a += b
val2 = a.numpy()
np.testing.assert_allclose(val1, val2)
def test_backward_pass(self):
def test_tinygrad():
x = Tensor(x_init)
W = Tensor(W_init)
m = Tensor(m_init)
out = x.dot(W).relu()
out = out.log_softmax()
out = out.mul(m).add(m).sum()
out.backward()
return out.numpy(), x.grad.numpy(), W.grad.numpy()
def test_pytorch():
x = torch.tensor(x_init, requires_grad=True)
W = torch.tensor(W_init, requires_grad=True)
m = torch.tensor(m_init)
out = x.matmul(W).relu()
out = torch.nn.functional.log_softmax(out, dim=1)
out = out.mul(m).add(m).sum()
out.backward()
return out.detach().numpy(), x.grad, W.grad
for x,y in zip(test_tinygrad(), test_pytorch()):
np.testing.assert_allclose(x, y, atol=1e-5)
# A simple test is to check that we can accumulate gradients (run backward twice or more times)
def test_accumulate_gradients(self):
x = Tensor(x_init)
W = Tensor(W_init)
m = Tensor(m_init)
out = x.dot(W).relu()
out = out.log_softmax()
out = out.mul(m).add(m).sum()
out.backward()
xgrad, wgrad = x.grad.numpy(), W.grad.numpy()
out.backward()
xgrad2, wgrad2 = x.grad.numpy(), W.grad.numpy()
out.backward() # no need to retain again since we will not re-run backward
xgrad3, wgrad3 = x.grad.numpy(), W.grad.numpy()
np.testing.assert_allclose(xgrad3, xgrad * 3., atol=1e-6)
np.testing.assert_allclose(wgrad3, wgrad * 3., atol=1e-6)
np.testing.assert_allclose(xgrad2, xgrad * 2., atol=1e-6)
np.testing.assert_allclose(wgrad2, wgrad * 2., atol=1e-6)
def test_second_order_backward_pass(self):
def test_pytorch():
x_val = torch.tensor([2.0], requires_grad=True)
f = x_val**3
first_derivative = torch.autograd.grad(outputs=f, inputs=x_val, create_graph=True)[0]
second_derivative = torch.autograd.grad(outputs=first_derivative, inputs=x_val)[0]
# d^2f/dx^2 = 6x = 6*2 = 12
return second_derivative.numpy()
def test_tinygrad():
x_val = Tensor([2.0])
f = x_val**3
first_derivative = f.sum().gradient(x_val)[0]
second_derivative = first_derivative.sum().gradient(x_val)[0]
return second_derivative.numpy()
np.testing.assert_allclose(test_tinygrad(), test_pytorch(), atol=1e-5)
# passing `gradient` to backward
def test_backward_pass_vjp(self):
def test_tinygrad():
x = Tensor(x_init)
W = Tensor(W_init)
m = Tensor(m_init)
out = x.dot(W).relu()
out = out.log_softmax()
out = out.mul(m).add(m)
out.backward(Tensor(gradient))
return out.numpy(), x.grad.numpy(), W.grad.numpy()
def test_pytorch():
x = torch.tensor(x_init, requires_grad=True)
W = torch.tensor(W_init, requires_grad=True)
m = torch.tensor(m_init)
out = x.matmul(W).relu()
out = torch.nn.functional.log_softmax(out, dim=1)
out = out.mul(m).add(m)
out.backward(torch.tensor(gradient))
return out.detach().numpy(), x.grad, W.grad
for x,y in zip(test_tinygrad(), test_pytorch()):
np.testing.assert_allclose(x, y, atol=1e-5)
def test_backward_pass_diamond_model(self):
def test_tinygrad():
u = Tensor(U_init)
v = Tensor(V_init)
w = Tensor(W_init)
x = u.mul(v).relu()
y = u.mul(w).relu()
out = x.add(y).mul(y).relu()
out = out.log_softmax()
out = out.sum()
out.backward()
return out.numpy(), u.grad.numpy(), v.grad.numpy(), w.grad.numpy()
def test_pytorch():
u = torch.tensor(U_init, requires_grad=True)
v = torch.tensor(V_init, requires_grad=True)
w = torch.tensor(W_init, requires_grad=True)
x = u.mul(v).relu()
y = u.mul(w).relu()
out = x.add(y).mul(y).relu()
out = torch.nn.functional.log_softmax(out, dim=1)
out = out.sum()
out.backward()
return out.detach().numpy(), u.grad, v.grad, w.grad
for x,y in zip(test_tinygrad(), test_pytorch()):
np.testing.assert_allclose(x, y, atol=1e-5, rtol=1e-6)
def test_const_backward_pass(self):
init = 3.5
def test_pytorch():
w1 = torch.tensor(init, requires_grad=True)
w2 = torch.tensor(init, requires_grad=True)
out = w1.add(w2)
out.backward()
return w1.grad, w2.grad
def test_tinygrad():
w1 = Tensor(init).clone()
w2 = Tensor(init).clone()
out = w1.add(w2)
out.backward()
return w1.grad.numpy(), w2.grad.numpy()
for x, y in zip(test_tinygrad(), test_pytorch()):
np.testing.assert_allclose(x, y, atol=1e-5)
def test_const_backward_pass_optimizer(self):
init = 3.5
def test_pytorch():
w1 = torch.tensor(init, requires_grad=True)
w2 = torch.tensor(init, requires_grad=True)
out = w1.add(w2)
out.backward()
return w1.grad.numpy(), w2.grad.numpy()
def test_tinygrad():
w1 = Tensor(init).clone()
w2 = Tensor(init).clone()
assert w1.is_param is True and w2.is_param is True
nn.optim.SGD([w1, w2], lr=0.01)
assert w1.is_param is True and w2.is_param is True
out = w1.add(w2)
out.backward()
return w1.grad.numpy(), w2.grad.numpy()
for x, y in zip(test_tinygrad(), test_pytorch()):
np.testing.assert_allclose(x, y, atol=1e-5)
def test_dropout(self):
with Context(TRAINING=1):
n, rate = 1_000_000, 0.1
w = Tensor.ones(n).dropout(rate)
non_zeros = np.count_nonzero(w.numpy())
expected = n * (1 - rate)
np.testing.assert_allclose(non_zeros, expected, rtol=2e-3)
def test_jacobian(self):
W = np.random.RandomState(42069).random((10, 5)).astype(np.float32)
x = np.random.RandomState(69420).random((1, 10)).astype(np.float32)
torch_x = torch.tensor(x, requires_grad=True)
torch_W = torch.tensor(W, requires_grad=True)
def torch_func(x): return torch.nn.functional.log_softmax(x.matmul(torch_W).relu(), dim=1)
PJ = torch.autograd.functional.jacobian(torch_func, torch_x).squeeze().numpy()
tiny_x = Tensor(x)
tiny_W = Tensor(W)
def tiny_func(x): return x.dot(tiny_W).relu().log_softmax()
J = jacobian(tiny_func, tiny_x)
NJ = numerical_jacobian(tiny_func, tiny_x)
np.testing.assert_allclose(PJ, J, atol = 1e-5)
np.testing.assert_allclose(PJ, NJ, atol = 1e-3)
def test_gradcheck(self):
W = np.random.RandomState(1337).random((10, 5)).astype(np.float32)
x = np.random.RandomState(7331).random((1, 10)).astype(np.float32)
tiny_x = Tensor(x)
tiny_W = Tensor(W)
def tiny_func(x): return x.dot(tiny_W).relu().log_softmax()
self.assertTrue(gradcheck(tiny_func, tiny_x, eps = 1e-3))
# coarse approx. since a "big" eps and the non-linearities of the model
self.assertFalse(gradcheck(tiny_func, tiny_x, eps = 1e-5))
def test_random_fns_are_deterministic_with_seed(self):
for random_fn in [Tensor.randn, Tensor.normal, Tensor.uniform, Tensor.scaled_uniform, Tensor.glorot_uniform, Tensor.kaiming_normal]:
with self.subTest(msg=f"Tensor.{random_fn.__name__}"):
Tensor.manual_seed(1337)
a = random_fn(10,10).realize()
Tensor.manual_seed(1337)
b = random_fn(10,10).realize()
np.testing.assert_allclose(a.numpy(), b.numpy())
def test_randperm(self):
Tensor.manual_seed(0)
a = Tensor.randperm(10).realize()
np.testing.assert_equal(a.numpy(), [8, 9, 4, 3, 6, 1, 7, 5, 2, 0])
b = Tensor.randperm(1000).realize()
np.testing.assert_equal(set(b.numpy()), set(range(1000)))
def test_rand_rejects_unknown_kwargs(self):
with self.assertRaises(TypeError): Tensor.rand(5, generator="foo")
def test_randn_isnt_inf_on_zero(self):
# simulate failure case of rand handing a zero to randn
original_rand, Tensor.rand = Tensor.rand, Tensor.zeros
try: self.assertNotIn(np.inf, Tensor.randn(16).numpy())
except: raise
finally: Tensor.rand = original_rand
def test_zeros_like_has_same_dtype_and_shape(self):
for datatype in [dtypes.float16, dtypes.float32, dtypes.int8, dtypes.int32, dtypes.int64, dtypes.uint8]:
a = Tensor([1, 2, 3], dtype=datatype)
b = Tensor.zeros_like(a)
assert a.dtype == b.dtype, f"dtype mismatch {a.dtype=} != {b.dtype}"
assert a.shape == b.shape, f"shape mismatch {a.shape} != {b.shape}"
a = Tensor([1, 2, 3])
b = Tensor.zeros_like(a, dtype=dtypes.int8)
assert a.dtype == dtypes.default_int and b.dtype == dtypes.int8, "a.dtype should be int and b.dtype should be char"
assert a.shape == b.shape, f"shape mismatch {a.shape} != {b.shape}"
def test_ones_like_has_same_dtype_and_shape(self):
for datatype in [dtypes.float16, dtypes.float32, dtypes.int8, dtypes.int32, dtypes.int64, dtypes.uint8]:
a = Tensor([1, 2, 3], dtype=datatype)
b = Tensor.ones_like(a)
assert a.dtype == b.dtype, f"dtype mismatch {a.dtype=} != {b.dtype}"
assert a.shape == b.shape, f"shape mismatch {a.shape} != {b.shape}"
a = Tensor([1, 2, 3])
b = Tensor.ones_like(a, dtype=dtypes.int8)
assert a.dtype == dtypes.default_int and b.dtype == dtypes.int8, "a.dtype should be int and b.dtype should be char"
assert a.shape == b.shape, f"shape mismatch {a.shape} != {b.shape}"
def test_rand_like_device(self):
a = Tensor.ones(3, 3, device="CPU")
b = Tensor.rand_like(a)
self.assertEqual(b.device, a.device)
def test_ndim(self):
assert Tensor(1).ndim == 0
assert Tensor.randn(1).ndim == 1
assert Tensor.randn(2,2,2).ndim == 3
assert Tensor.randn(1,1,1,1,1,1).ndim == 6
def test_argfix(self):
for f in [Tensor.zeros, Tensor.ones, Tensor.rand, Tensor.randn, Tensor.empty]:
self.assertEqual(f().shape, ())
self.assertEqual(f(1).shape, (1,))
self.assertEqual(f(10,20,40).shape, (10,20,40))
self.assertEqual(f([]).shape, ())
self.assertEqual(f([1]).shape, (1,))
self.assertEqual(f([10,20,40]).shape, (10,20,40))
self.assertEqual(f(()).shape, ())
self.assertEqual(f((1,)).shape, (1,))
self.assertEqual(f((10,20,40)).shape, (10,20,40))
with self.assertRaises(ValueError): f((2, 2), 2, 2)
with self.assertRaises(ValueError): f((2, 2), (2, 2))
with self.assertRaises(ValueError): f((128, 128), 0.0, 0.01)
def test_numel(self):
assert Tensor.randn(10, 10).numel() == 100
assert Tensor.randn(1,2,5).numel() == 10
assert Tensor.randn(1,1,1,1,1,1).numel() == 1
assert Tensor([]).numel() == 0
assert Tensor.randn(1,0,2,5).numel() == 0
assert Tensor(3).numel() == 1
def test_len(self):
assert len(torch.zeros(7)) == len(Tensor.zeros(7))
assert len(torch.zeros(10,20)) == len(Tensor.zeros(10,20))
assert len(torch.zeros(10,20)) == len(Tensor.zeros(10,20,30))
assert len(torch.zeros(1).flatten()) == len(Tensor.zeros(1).flatten())
with self.assertRaises(TypeError): len(Tensor(3))
def test_size(self):
t1, t2 = torch.zeros(10,20), Tensor.zeros(10,20)
assert t1.size() == t2.size()
assert t1.size(0) == t2.size(0)
assert t1.size(1) == t2.size(1)
assert t1.size(-1) == t2.size(-1)
assert t1.size(-2) == t2.size(-2)
with self.assertRaises(IndexError): t2.size(2)
def test_tolist(self):
# NOTE: float16 Tensor.tolist() requires python 3.12
for arr in [[1,2,3], [1.5,2,3], [[1,2,3], [4,5,6]], 3]:
assert Tensor(arr).tolist() == torch.tensor(arr).tolist() == arr
def test_element_size(self):
for _, dtype in DTYPES_DICT.items():
assert dtype.itemsize == Tensor.randn(3, dtype=dtype).element_size(), f"Tensor.element_size() not matching Tensor.dtype.itemsize for {dtype}"
def test_deepwalk_ctx_check(self):
layer = Tensor.uniform(1, 1)
x = Tensor.randn(1, 1, 1)
x.dot(layer).mean().backward()
x = Tensor.randn(1, 1, 1)
x.dot(layer).mean().backward()
def test_zerosized_tensors(self):
np.testing.assert_equal(Tensor([]).numpy(), np.array([]))
np.testing.assert_equal(Tensor(None).numpy(), np.array([]))
def test_tensor_ndarray_dtype(self):
arr = np.array([1]) # where dtype is implicitly int64
assert Tensor(arr).dtype == dtypes.int64
assert Tensor(arr, dtype=dtypes.float32).dtype == dtypes.float32 # check if ndarray correctly casts to Tensor dtype
assert Tensor(arr, dtype=dtypes.float64).dtype == dtypes.float64 # check that it works for something else
def test_tensor_from_blob(self):
x = memoryview(bytearray(16)).cast('I')
t = Tensor.from_blob(mv_address(x), (4,), dtype=dtypes.int, device="CPU")
z = (t+1)
np.testing.assert_equal(z.numpy(), [1, 1, 1, 1])
x[:] = array.array('I', [0, 1, 2, 3])
z = (t+1)
np.testing.assert_equal(z.numpy(), [1, 2, 3, 4])
def test_tensor_list_dtype(self):
for arr in ([1], [[[1]]], [[1,1],[1,1]], [[[1,1],[1,1]],[[1,1],[1,1]]]):
assert Tensor(arr).dtype == dtypes.default_int
assert Tensor(arr, dtype=dtypes.float32).dtype == dtypes.float32
assert Tensor(arr, dtype=dtypes.float64).dtype == dtypes.float64
for arr in ([True], [[[False]]], [[True,False],[True,False]], [[[False,True],[False,False]],[[True,True],[False,True]]]):
assert Tensor(arr).dtype == dtypes.bool
assert Tensor(arr, dtype=dtypes.float32).dtype == dtypes.float32
assert Tensor(arr, dtype=dtypes.float64).dtype == dtypes.float64
# empty tensor defaults
for arr in ([], [[[]]], [[],[]]):
t = Tensor(arr)
assert t.dtype == dtypes.default_float
np.testing.assert_allclose(t.numpy(), np.array(arr))
# mixture of bool and int
for arr in ([True, 3], [[True],[3]], [[[True]], [[3]]], [[True, 3], [3, True]]):
t = Tensor(arr)
assert t.dtype == dtypes.default_int
np.testing.assert_allclose(t.numpy(), np.array(arr))
# mixture of bool, int and float
for arr in ([[True,True],[3.,True]], [[0,1],[3.,4]], [[[0],[1]],[[3.],[4]]], [[[True],[1]],[[3.],[4]]]):
t = Tensor(arr)
assert t.dtype == dtypes.default_float
np.testing.assert_allclose(t.numpy(), np.array(arr))
def test_tensor_list_shapes(self):
self.assertEqual(Tensor([[[]]]).shape, (1,1,0))
self.assertEqual(Tensor([[],[]]).shape, (2,0))
self.assertEqual(Tensor([[[[]],[[]]], [[[]],[[]]], [[[]],[[]]]]).shape, (3,2,1,0))
def test_tensor_list_errors(self):
# inhomogeneous shape
with self.assertRaises(ValueError): Tensor([[],[[]]])
with self.assertRaises(ValueError): Tensor([[1],[]])
with self.assertRaises(ValueError): Tensor([[1],[1],1])
with self.assertRaises(ValueError): Tensor([[[1,1,1],[1,1]]])
with self.assertRaises(ValueError): Tensor([[1,1,1],[[1,1,1]]])
def test_tensor_mixed_list_tuple(self):
def _list_or_tuple(): return list if random.random() < 0.5 else tuple
def _generate_data(depth):
if depth == 0: return _list_or_tuple()()
if depth == 1: return _list_or_tuple()([random.random(), random.random()])
return _list_or_tuple()([_generate_data(depth-1), _generate_data(depth-1)])
for depth in range(7):
for _ in range(20):
data = _generate_data(depth)
np.testing.assert_allclose(Tensor(data).numpy(), np.array(data))
def test_tensor_list_implicit_cast(self):
data = [True, False]
np.testing.assert_equal(Tensor(data, dtype=dtypes.int).numpy(), torch.tensor(data, dtype=torch.int).numpy())
np.testing.assert_equal(Tensor(data, dtype=dtypes.uint8).numpy(), torch.tensor(data, dtype=torch.uint8).numpy())
np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())
data = [-1, 0, 1, 2, 3]
np.testing.assert_equal(Tensor(data, dtype=dtypes.int).numpy(), torch.tensor(data, dtype=torch.int).numpy())
np.testing.assert_equal(Tensor(data, dtype=dtypes.uint8).numpy(), torch.tensor(data, dtype=torch.uint8).numpy())
np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())
data = [-3.5, -2.5, -1.5, 0, 1.5, 2.5, 3.5]
np.testing.assert_equal(Tensor(data, dtype=dtypes.int).numpy(), torch.tensor(data, dtype=torch.int).numpy())
# NOTE: torch and jax raise OverflowError: Python integer -3 out of bounds for uint8
# np.testing.assert_equal(Tensor(data, dtype=dtypes.uint8).numpy(), torch.tensor(data, dtype=torch.uint8).numpy())
np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())
def test_tensor_list_special_values(self):
if dtypes.float16 in Device[Device.DEFAULT].renderer.supported_dtypes():
data = [math.nan, -math.inf, 65504, 65519, 65519.999, 65520, 65520.1]
data = data + [-x for x in data]
with np.errstate(over='ignore'): np.testing.assert_allclose(Tensor(data, dtype=dtypes.float16).numpy(), np.array(data).astype(np.float16))
# uint32
data = [1 << 33, 1 << 32, 1 << 32 - 1, 1]
data = data + [-x for x in data]
np.testing.assert_allclose(Tensor(data, dtype=dtypes.uint32).numpy(), np.array(data).astype(np.uint32))
# int32
data = [1 << 33, 1 << 32, 1 << 32 - 1, 1]
data = data + [-x for x in data]
np.testing.assert_allclose(Tensor(data, dtype=dtypes.int32).numpy(), np.array(data).astype(np.int32))
@unittest.skip("list elements are python scalars: numpy values in a list have no inferred dtype, use Tensor(np.array(data))")
def test_tensor_list_ndarray(self):
data = [np.array([1, 2, 3]), np.array([1, 2, 3]), np.array([1, 2, 3])]
np.testing.assert_equal(Tensor(data).numpy(), np.array(data))
data = [np.array([1.0, 2.0, 3.0]), np.array([1, 2, 3]), np.array([1, 2, 3])]
np.testing.assert_equal(Tensor(data).numpy(), np.array(data))
data = [np.array(1.0), np.array(2.0), np.array(3.0)]
np.testing.assert_equal(Tensor(data).numpy(), np.array(data))
def test_tensor_dtype_errors(self):
with self.assertRaises(AttributeError): Tensor([3], dtype="typo")
with self.assertRaises(AttributeError): Tensor([3], dtype=(dtypes.int,))
def test_tensor_bytes(self):
data = b"abc123"
t = Tensor(data)
assert t.dtype == dtypes.uint8
assert t.shape == (6,)
np.testing.assert_equal(t.numpy(), list(data))
def test_tensor_copy(self):
x = copy.deepcopy(Tensor.ones((3,3,3)))
np.testing.assert_allclose(x.numpy(), np.ones((3,3,3)))
def test_copy_from_disk(self):
t = Tensor.randn(30).to(f"disk:{temp('test_copy_from_disk')}")
a = t[10:20]
dev = a.to(Device.DEFAULT)
np.testing.assert_allclose(a.numpy(), dev.numpy())
def test_copy_from_numpy_dtype(self):
data = np.array([1.0, 2, 3], dtype=np.float32)
t = Tensor(data, dtype=dtypes.bfloat16)
try:
# TODO: fix dtype in tinygrad space
assert t.dtype == dtypes.bfloat16
except AssertionError:
assert t.dtype == dtypes.float32
np.testing.assert_equal(t.tolist(), data)
np.testing.assert_equal((t+1).tolist(), data+1)
# Regression test for https://github.com/tinygrad/tinygrad/issues/1751
def test_copy_from_numpy_unaligned(self):
# 2**15 is the minimum for repro
arr = np.random.randn(2**15).astype(np.float32)
fn = temp('test_copy_from_numpy_unaligned')
with open(fn, 'wb') as f: f.write(b't' + arr.tobytes())
with open(fn, "a+b") as f: memview = memoryview(mmap.mmap(f.fileno(), arr.nbytes + 1))
ua_arr = np.frombuffer(memview[1:], dtype=arr.dtype, count=arr.shape[0])
np.testing.assert_allclose(arr, ua_arr)
assert not ua_arr.flags.aligned
# force device copy - to() is opt'd away - Tensor(dev)/1 is ignored
np.testing.assert_allclose(ua_arr, (Tensor(ua_arr)/Tensor(1)).numpy())
def test_item_to_tensor_to_item(self):
for a in [0, 1, 2, 3, -1, -100, 100, -101.1, 2.345, 100.1, True, False]:
item = Tensor(a).item()
assert type(item) is type(a), a
np.testing.assert_allclose(item, a), a
buffered_item = Tensor([a]).item()
assert type(buffered_item) is type(a), a
np.testing.assert_allclose(buffered_item, a), a
reshaped_item = Tensor([a]).reshape((1, 1, 1, 1, 1)).item()
assert type(reshaped_item) is type(a), a
np.testing.assert_allclose(reshaped_item, a), a
def test_no_bool(self):
with self.assertRaises(TypeError):
if Tensor(3):
print("hi")
with self.assertRaises(TypeError):
_a = Tensor([3]) in [Tensor([3]), Tensor([4]), Tensor([5])]
def test_repr_with_grad(self):
a = Tensor([1.0])
b = Tensor([1])
c = (a + b).sum().backward()
print(a)
print(c)
def test_no_attributeerror_after_apply_uop_exception(self):
try:
Tensor.arange(4).reshape(3,2)
except ValueError:
Tensor.zeros(2, 2).realize()
def test_shrink(self):
t = Tensor.arange(32).clone().realize()
self.assertListEqual(t[16:20].tolist(), [16,17,18,19])
self.assertListEqual(t.shrink_to(16).tolist(), list(range(16)))
t = t.reshape(4, 8).contiguous().realize()
self.assertListEqual(t.shrink_to(2, 2).tolist(), [[0, 1], [8, 9]])
self.assertListEqual(t.shrink_to(None, 2).tolist(), t.shrink_to(4, 2).tolist())
with self.assertRaises(ValueError): t.shrink_to(2)
with self.assertRaises(ValueError): t.shrink_to(2, 2, 2)
@unittest.skip("this test is just flaky, sync issue")
class TestMoveTensor(unittest.TestCase):
d0, d1 = f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"
@given(strat.sampled_from([d0, d1]), strat.sampled_from([d0, d1]),
strat.sampled_from([dtypes.float16, dtypes.float32]), strat.sampled_from([True, False]))
def test_to_preserves(self, src, dest, dtype, is_param):
if dtype not in Device[Device.DEFAULT].renderer.supported_dtypes():
return
s = Tensor([1, 2, 3], device=src, dtype=dtype).is_param_(is_param)
if is_param: s.sum().backward()
t = s.to(dest)
np.testing.assert_equal(s.numpy(), t.numpy())
assert s.dtype == t.dtype
assert s.is_param == t.is_param
if is_param:
np.testing.assert_equal(s.grad.numpy(), t.grad.numpy())
@given(strat.sampled_from([dtypes.float16, dtypes.float32]), strat.sampled_from([True, False]))
def test_shard_preserves(self, dtype, is_param):
s = Tensor([1, 2, 3], dtype=dtype).is_param_(is_param)
t = s.shard((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"))
np.testing.assert_equal(s.numpy(), t.numpy())
assert s.dtype == t.dtype
assert s.is_param == t.is_param
@given(strat.sampled_from([d0, d1]))
def test_same_dev(self, dev):
x = Tensor([1,2,3], device=dev)
y = x.to(dev)
assert x is y
def test_to_grad(self):
x = Tensor.eye(3).clone(self.d0)
y = Tensor([[2.0,0,-2.0]], device=self.d0)
z = y.matmul(x).to(self.d1).sum()
z.backward()
np.testing.assert_equal(x.grad.numpy(), [[2,2,2],[0,0,0],[-2,-2,-2]])
class TestZeroShapeTensor(unittest.TestCase):
def test_rand(self):
t = Tensor.rand(3, 2, 0)
assert t.shape == (3, 2, 0)
np.testing.assert_equal(t.numpy(), np.zeros((3, 2, 0)))
t = Tensor.rand(0)
assert t.shape == (0,)
np.testing.assert_equal(t.numpy(), np.zeros((0,)))
t = Tensor.rand(0, 0, 0)
assert t.shape == (0, 0, 0)
np.testing.assert_equal(t.numpy(), np.zeros((0, 0, 0)))
def test_full(self):
t = Tensor.zeros(3, 2, 0)
assert t.shape == (3, 2, 0)
np.testing.assert_equal(t.numpy(), np.zeros((3, 2, 0)))
t = Tensor.full((3, 2, 0), 12)
assert t.shape == (3, 2, 0)
np.testing.assert_equal(t.numpy(), np.full((3, 2, 0), 12))
def test_reshape(self):
t = Tensor.zeros(3, 2, 0)
a = t.reshape(7, 0)
assert a.shape == (7, 0)
np.testing.assert_equal(a.numpy(), np.zeros((7, 0)))
a = t.reshape(0)
assert a.shape == (0,)
np.testing.assert_equal(a.numpy(), np.zeros((0,)))
with self.assertRaises(ValueError):
# cannot reshape from size 0 to size 1
a = t.reshape(())
def test_expand(self):
t = Tensor.full((1, 2, 0), 12).expand((6, 2, 0))
assert t.shape == (6, 2, 0)
np.testing.assert_equal(t.numpy(), np.full((6, 2, 0), 12))
def test_pad(self):
t = Tensor.rand(3, 2, 0).pad((None, None, (1, 1)), value=1)
self.assertEqual(t.shape, (3, 2, 2))
np.testing.assert_equal(t.numpy(), np.ones((3, 2, 2)))
t = Tensor.rand(3, 2, 0).pad((None, (1, 1), None), value=1)
self.assertEqual(t.shape, (3, 4, 0))
np.testing.assert_equal(t.numpy(), np.ones((3, 4, 0)))
t = Tensor.rand(3, 2, 0).pad(((1, 1), None, None), value=1)
self.assertEqual(t.shape, (5, 2, 0))
np.testing.assert_equal(t.numpy(), np.ones((5, 2, 0)))
np.testing.assert_equal(Tensor([1, 2]).pad_to(4).numpy(), [1, 2, 0, 0])
np.testing.assert_equal(Tensor([[1, 2]]).pad_to(2, 3).numpy(), [[1, 2, 0], [0, 0, 0]])
np.testing.assert_equal(Tensor([[1, 2]]).pad_to(1, 3).numpy(), [[1, 2, 0]])
np.testing.assert_equal(Tensor([[1, 2]]).pad_to(None, 3).numpy(), [[1, 2, 0]])
with self.assertRaises(ValueError): Tensor([1, 2]).pad_to(2, 3)
with self.assertRaises(ValueError): Tensor([[1, 2]]).pad_to(3)
def test_shrink_into_zero(self):
t = Tensor.rand(3, 4).realize()
assert t.shrink((None, (2, 2))).realize().shape == (3, 0)
assert t.shrink(((2, 2), None)).realize().shape == (0, 4)
assert t.shrink(((2, 2), (2, 2))).realize().shape == (0, 0)
def test_cat(self):
a = Tensor.rand(3, 2, 2)
b = Tensor.rand(3, 2, 0)
t = a.cat(b, dim=2)
assert t.shape == (3, 2, 2)
np.testing.assert_equal(t.numpy(), a.numpy())
t = b.cat(a, dim=2)
assert t.shape == (3, 2, 2)
np.testing.assert_equal(t.numpy(), a.numpy())
t = b.cat(b, dim=0)
assert t.shape == (6, 2, 0)
np.testing.assert_equal(t.numpy(), np.zeros((6, 2, 0)))
t = b.cat(b, dim=1)
assert t.shape == (3, 4, 0)
np.testing.assert_equal(t.numpy(), np.zeros((3, 4, 0)))
t = b.cat(b, dim=2)
assert t.shape == (3, 2, 0)
np.testing.assert_equal(t.numpy(), np.zeros((3, 2, 0)))
def test_elementwise(self):
a = Tensor.rand(3, 2, 0)
a_exp = a.exp()
assert a_exp.shape == (3, 2, 0)
np.testing.assert_equal(a_exp.numpy(), np.exp(a.numpy()))
b = Tensor.rand(3, 2, 0)
assert b.shape == (3, 2, 0)
ab = a * b
assert ab.shape == (3, 2, 0)
np.testing.assert_equal(ab.numpy(), a.numpy() * b.numpy())
mask = (Tensor.rand(3, 2, 0) > 0.5)
assert mask.shape == (3, 2, 0)
c = mask.where(a, b)
assert c.shape == (3, 2, 0)
np.testing.assert_equal(c.numpy(), np.where(mask.numpy(), a.numpy(), b.numpy()))
def test_reduce_over_non_zero(self):
a = Tensor.ones(3, 2, 0).sum(axis=1)
assert a.shape == (3, 0)
np.testing.assert_equal(a.numpy(), np.sum(np.zeros((3, 2, 0)), axis=1))
def test_reduce_over_zero(self):
a = Tensor.ones(3, 2, 0).sum(axis=2)
assert a.shape == (3, 2)
np.testing.assert_equal(a.numpy(), np.sum(np.zeros((3, 2, 0)), axis=2))
a = Tensor.ones(3, 2, 0).sum(axis=2, keepdim=True)
assert a.shape == (3, 2, 1)
np.testing.assert_equal(a.numpy(), np.sum(np.zeros((3, 2, 0)), axis=2, keepdims=True))
def test_clone(self):
a = Tensor.rand(16, 16).realize()
b = a.clone()
np.testing.assert_allclose(a.numpy(), b.numpy())
self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)
a = Tensor.rand(16, 16).mul(5.0).add(5.0).realize()
b = a.clone()
np.testing.assert_allclose(a.numpy(), b.numpy())
self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)
def test_clone_deviceless_const(self):
t = Tensor(UOp.const(2.0).cast(dtypes.float)).clone()
np.testing.assert_equal(t.numpy(), 2.0)
self.assertTrue(t.uop.has_buffer_identity())
def test_numpy_deviceless_const(self):
np.testing.assert_equal(Tensor(UOp.const(2.0).cast(dtypes.float)).numpy(), 2.0)
def test_clone_with_shrink(self):
a = Tensor.rand(16, 16)
b = a.shrink(((2, 10), None)).clone()
b.realize()
self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)
def test_clone_with_shrink_realized(self):
a = Tensor.rand(16, 16).realize()
b = a.shrink(((2, 10), None)).clone()
b.realize()
self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)
def test_clone_with_grad(self):
a = Tensor.rand(16, 16)
a.mul(5.0).add(5.0).mean().backward()
b = a.clone()
assert a.grad is not None
assert b.grad is not None
np.testing.assert_allclose(a.grad.numpy(), b.grad.numpy())
def test_clone_deviceless_const_to_cpu(self):
t = Tensor(UOp.const(2.0).cast(dtypes.float)).clone(device="CPU")
self.assertEqual(t.device, "CPU")
np.testing.assert_equal(t.numpy(), 2.0)
def test_reduce_default(self):
np.testing.assert_equal(Tensor([]).max().numpy(), -float("inf"))
np.testing.assert_equal(Tensor([]).min().numpy(), float("inf"))
np.testing.assert_equal(Tensor([]).sum().numpy(), 0)
np.testing.assert_equal(Tensor([]).mean().numpy(), float("nan"))
class TestTensorCreationDevice(unittest.TestCase):
# test auxiliary tensors are created on the same device
def test_one_hot(self):
y = Tensor([1, 2, 3]).to("CPU")
x = y.one_hot(10)
x.realize()
if __name__ == '__main__':
unittest.main()

221
artifacts/package_sources/tinygrad/test/backend/test_tensor_variable.py Normal file
View File

@@ -0,0 +1,221 @@
import unittest
import numpy as np
from tinygrad import Device, Tensor, Variable, TinyJit, dtypes
from tinygrad.helpers import CHECK_OOB
class TestTensorVariable(unittest.TestCase):
def test_add_tvar(self):
vv = Variable("a", 0, 10).bind(1)
ret = (Tensor(vv) + 3).item()
assert ret == 4
def test_variable_mul_tensor(self):
vv = Variable("a", 1, 10).bind(2)
t = Tensor.ones(3, dtype=dtypes.int8)
self.assertListEqual((t * vv).tolist(), [2, 2, 2])
# TODO: fix
try:
self.assertListEqual((vv * t).tolist(), [2, 2, 2])
except RuntimeError: pass
@unittest.skipUnless(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires long support")
def test_large_range_variable(self):
self.assertEqual(Tensor(Variable("b", 0, 2**40, dtype=dtypes.long).bind(2**35)).clone(Device.DEFAULT).item(), 2**35)
@unittest.skipUnless(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires long support")
def test_large_range_variable_jit(self):
@TinyJit
def f(a,b): return (Tensor(a+b).clone(Device.DEFAULT) * 2).realize()
for i in range(3):
a = Variable("a", 0, 2**10, dtype=dtypes.int).bind(i)
b = Variable("b", 0, 2**40, dtype=dtypes.long).bind(2**35)
self.assertEqual(f(a,b).item(), (2**35 + i) * 2)
def test_variable_defers_like_a_literal(self):
vv = Variable("a", 1, 10).bind(2)
self.assertEqual(Tensor(vv).dtype, dtypes.weakint)
self.assertEqual((Tensor(vv) + Tensor([1], dtype=dtypes.int8)).dtype, dtypes.int8) # takes the concrete side, no widening
self.assertEqual(Tensor(vv).item(), 2) # a read commits at default_int
def test_variable_tensor_dtype_arg(self):
vv = Variable("a", 1, 10).bind(2)
t = Tensor(vv, dtype=dtypes.float32)
self.assertEqual(t.dtype, dtypes.float32)
self.assertEqual(t.item(), 2.0)
def test_unbound_variable_tensor(self):
# an unbound variable schedules fine, but can't execute
with self.assertRaisesRegex(RuntimeError, "unbound"): Tensor(Variable("u", 1, 10)).item()
with self.assertRaisesRegex(RuntimeError, "unbound"): (Tensor(Variable("u", 1, 10)) + 1).item()
# bound variables in an expression are fine
self.assertEqual(Tensor(Variable("u", 1, 10).bind(2) + 1).item(), 3)
def test_shrink_beyond_buffer_variable(self):
# TODO: shrink by a variable whose vmax exceeds the dim should fail at build, today only CHECK_OOB=1 rejects it
t = Tensor.ones(3).contiguous()[:Variable("a", 1, 10).bind(5)]
if CHECK_OOB: self.assertRaises(RuntimeError, t.sum().item)
else: t.sum().item() # silent OOB: reads 2 elements past the buffer, result depends on the allocator
def test_symbolic_shape_mul_variable_tensor(self):
# NOTE: the buffer dim must cover the variable's vmax
vv = Variable("a", 1, 10).bind(2)
self.assertEqual((Tensor.ones(10).contiguous()[:vv] * Tensor(vv)).sum().item(), 4.0)
# a vmin=0 symbolic dim broadcasts too
v0 = Variable("z", 0, 10).bind(2)
self.assertEqual((Tensor.ones(10).contiguous()[:v0] * Tensor(v0)).sum().item(), 4.0)
def test_inner_tvar_node(self):
vv = Variable("w", 0, 10).bind(2)
ret = Tensor(vv * 4).item()
assert ret == 8
def test_inner_tvar_mul(self):
vv = Variable("w", 0, 10).bind(2)
assert (Tensor(3) * vv).item() == 6
def test_inner_tvar_mul_node(self):
vv = Variable("w", 0, 10).bind(2)
assert (Tensor(3) * (vv * 4)).item() == 24
def test_symbolic_mean(self):
vv = Variable("a", 1, 10).bind(2)
t = Tensor.ones(2, 10).contiguous()[:, :vv]
ret = t.mean().item()
assert ret == 1
def test_symbolic_mean_2d(self):
vv = Variable("a", 1, 10).bind(2)
vv2 = Variable("b", 1, 10).bind(2)
t = Tensor.ones(10, 10).contiguous()[:vv2, :vv]
ret = t.mean().item()
assert ret == 1
def test_symbolic_mean_2d_axis_1(self):
vv = Variable("a", 1, 10).bind(2)
vv2 = Variable("b", 1, 10).bind(2)
t = Tensor.ones(10, 10).contiguous()[:vv2, :vv]
ret = t.mean(axis=1)[:2].reshape(2, 1).numpy()
assert np.all(ret == 1)
def test_symbolic_mean_2d_add(self):
add_term = Variable("c", 0, 10).bind(1)
vv = Variable("a", 1, 10).bind(1)
vv2 = Variable("b", 1, 10).bind(1)
t = Tensor.ones(20, 20).contiguous()[:vv2+add_term, :vv+add_term]
ret = t.mean().item()
assert ret == 1
def test_symbolic_var(self):
vv = Variable("a", 1, 10).bind(2)
t = Tensor.ones(2, 10).contiguous()[:, :vv]
ret = t.var().item()
assert ret == 0
def test_symbolic_pad(self):
vv = Variable("a", 1, 10).bind(2)
t = Tensor.ones(2, 2).contiguous()
t = t.pad([vv, vv, vv, vv]).mean()
ones = 4
zeros = 6+6+4+4+6+6
self.assertAlmostEqual(t.item(), ones/(ones+zeros))
def test_symbolic_arange(self):
vv = Variable("a", 1, 10)
ret = Tensor.arange(0, vv.bind(4))
self.assertListEqual(ret[:4].tolist(), [0,1,2,3])
def test_symbolic_arange_sym_start(self):
vv = Variable("a", 1, 6)
ret = Tensor.arange(vv.bind(4), 7)
self.assertListEqual(ret[:3].tolist(), [4,5,6])
def test_symbolic_arange_sym_step(self):
vv = Variable("step", 1, 3)
ret = Tensor.arange(0, 10, vv.bind(2))
self.assertListEqual(ret[:5].tolist(), [0,2,4,6,8])
def test_symbolic_arange_two_vars(self):
begin = Variable("b", 1, 5)
end = Variable("e", 6, 10)
ret = Tensor.arange(begin.bind(4), end.bind(7))
self.assertListEqual(ret[:3].tolist(), [4,5,6])
def test_symbolic_arange_three_vars(self):
begin = Variable("b", 0, 5)
end = Variable("e", 10, 20)
step = Variable("s", 1, 3)
ret = Tensor.arange(begin.bind(2), end.bind(14), step.bind(3))
self.assertListEqual(ret[:4].tolist(), [2,5,8,11])
def test_symbolic_full(self):
vv = Variable("x", 1, 10).bind(5)
t = Tensor.full((3,), vv)
self.assertListEqual(t.tolist(), [5,5,5])
def test_variable_empty(self):
v = Variable("i", 1, 10)
# TODO: Tensor creation from unbound variable should assert
# with self.assertRaises(AssertionError): t = Tensor.empty(3, v)
vb = v.bind(3)
t = Tensor.empty(3, vb)
assert t.uop.base.buffer.size == 30
assert t.uop.shape == (3, vb)
def test_symbolic_chunk(self):
# chunk should work when split dimension is concrete, even if other dims are symbolic
vv = Variable("a", 1, 10).bind(4)
t = Tensor.ones(10, 8).contiguous()[:vv, :] # shape (vv, 8)
chunks = t.chunk(2, dim=-1) # split along concrete dim 8
assert len(chunks) == 2
assert chunks[0].shape[1] == 4
assert chunks[1].shape[1] == 4
# verify the values by shrinking to concrete shape first
np.testing.assert_equal(chunks[0].shrink(((0, 4), (0, 4))).numpy(), np.ones((4, 4)))
np.testing.assert_equal(chunks[1].shrink(((0, 4), (0, 4))).numpy(), np.ones((4, 4)))
def test_symbolic_split(self):
# split should work when split dimension is concrete, even if other dims are symbolic
vv = Variable("a", 1, 10).bind(3)
t = Tensor.arange(30).reshape(10, 3).contiguous()[:, :vv] # shape (10, vv)
splits = t.split(5, dim=0) # split along concrete dim 10
assert len(splits) == 2
assert splits[0].shape[0] == 5
assert splits[1].shape[0] == 5
# verify the values by shrinking to concrete shape first
np.testing.assert_equal(splits[0].shrink(((0, 5), (0, 3))).numpy(), np.arange(30).reshape(10, 3)[:5, :3])
np.testing.assert_equal(splits[1].shrink(((0, 5), (0, 3))).numpy(), np.arange(30).reshape(10, 3)[5:, :3])
def test_symbolic_chunk_error_on_symbolic_dim(self):
# chunk should fail when trying to split along a symbolic dimension
vv = Variable("a", 1, 10).bind(4)
t = Tensor.ones(10, 8).contiguous()[:vv, :] # shape (vv, 8)
with self.assertRaises(AssertionError):
t.chunk(2, dim=0) # can't split along symbolic dim
def test_symbolic_var_sum(self, var_name="u"):
t = Variable("t", 1, 10).bind(4)
v = Variable(var_name, 1, 5).bind(1)
mask = (Tensor.full((1, 1, t, v+t), 1) + 1).contiguous()
mask.shrink(((0, 1), (0, 1), (0, 4), (0, 4))).numpy()
def test_symbolic_var_sum_alt_name(self): self.test_symbolic_var_sum("s")
def test_symbolic_triu(self):
t = Variable("t", 1, 10).bind(4)
for start_pos in (0, 1, 3):
var_start_pos = Variable("start_pos", 0, 5).bind(start_pos)
mask = Tensor.full((1, 1, t, var_start_pos+t), float("-inf")).triu(var_start_pos+1)
out = mask.shrink(((0, 1), (0, 1), (0, 4), (0, start_pos+4))).numpy()
expected = np.triu(np.full((1, 1, 4, start_pos+4), float("-inf")), k=start_pos+1)
np.testing.assert_equal(out, expected)
def test_symbolic_tril(self):
t = Variable("t", 1, 10).bind(4)
for start_pos in (0, 1, 3):
var_start_pos = Variable("start_pos", 0, 5).bind(start_pos)
mask = Tensor.full((1, 1, t, var_start_pos+t), float("-inf")).tril(var_start_pos+1)
out = mask.shrink(((0, 1), (0, 1), (0, 4), (0, start_pos+4))).numpy()
expected = np.tril(np.full((1, 1, 4, start_pos+4), float("-inf")), k=start_pos+1)
np.testing.assert_equal(out, expected)
if __name__ == '__main__':
unittest.main()

17
artifacts/package_sources/tinygrad/test/backend/test_to_numpy.py Normal file
View File

@@ -0,0 +1,17 @@
from tinygrad.tensor import Tensor
import numpy as np
import pickle
import unittest
class TestToNumpy(unittest.TestCase):
def test_numpy_is_numpy(self):
output = Tensor.ones((1, 3, 4096)).realize().numpy()
new = np.copy(output)
print(type(new))
serialized = pickle.dumps(new)
out = pickle.loads(serialized)
assert out.shape == (1,3,4096)
assert (out==1).all()
if __name__ == '__main__':
unittest.main()

205
artifacts/package_sources/tinygrad/test/backend/test_transcendental.py Normal file
View File

@@ -0,0 +1,205 @@
import unittest
from tinygrad import Tensor, Device, dtypes
from tinygrad.tensor import _to_np_dtype
from tinygrad.helpers import Context, getenv, DEV, OSX
from test.helpers import check_schedule
from test.backend.test_dtype_alu import ht, dtypes_float
import numpy as np
import math
from hypothesis import given, settings, strategies as strat
settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
settings.load_profile("my_profile")
supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
class TestTranscendentalMath(unittest.TestCase):
@unittest.skipUnless(dtypes.float64 in supported_dtypes, f"no float64 on {Device.DEFAULT}")
@unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}, "crashed")
@given(ht.float64, strat.sampled_from([(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.sin)]))
def test_float64(self, x, op):
if op[0] == Tensor.sin:
# TODO: reduction does not work # 536870912.125 # 2914593.01171875 # 134217728.03125 # 230581075.65625 # 139216373.71875
if abs(x) > 100_000_000: return
with Context(TRANSCENDENTAL=2), np.errstate(all='ignore'):
np.testing.assert_allclose(op[0](Tensor([x], dtype=dtypes.float64)).numpy(),
op[1](np.array([x], dtype=_to_np_dtype(dtypes.float64))),
atol=3e-2, rtol=1e-5) # sin can have bigger atol for very big x
@unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}, "crashed")
@given(ht.float32, strat.sampled_from([(Tensor.exp, np.exp),(Tensor.log, np.log)] +
([(Tensor.sin, np.sin)] if dtypes.ulong in supported_dtypes else [])))
def test_float32(self, x, op):
# wrong nan behavior on Vulkan
if (math.isnan(x) or (x < 0 and op[0] == Tensor.log)) and Device.DEFAULT == "WEBGPU" and not OSX: return
with Context(TRANSCENDENTAL=2), np.errstate(all='ignore'):
np.testing.assert_allclose(op[0](Tensor([x], dtype=dtypes.float32)).numpy(),
op[1](np.array([x], dtype=_to_np_dtype(dtypes.float32))),
atol=2e-5, rtol=1e-5)
@unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
@given(ht.float16, strat.sampled_from([(Tensor.exp, np.exp),(Tensor.log, np.log)] +
([(Tensor.sin, np.sin)] if dtypes.ulong in supported_dtypes else [])))
def test_float16(self, x, op):
# wrong nan behavior on Vulkan
if (math.isnan(x) or (x < 0 and op[0] == Tensor.log)) and Device.DEFAULT == "WEBGPU" and not OSX: return
with Context(TRANSCENDENTAL=2), np.errstate(all='ignore'):
np.testing.assert_allclose(op[0](Tensor([x], dtype=dtypes.float16)).numpy(),
op[1](np.array([x], dtype=_to_np_dtype(dtypes.float16))),
atol=1e-2, rtol=5e-3) # exp can have bigger rtol
# TODO: WEBGPU produces incorrect values near infinity
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU incorrect values near inf")
@given(strat.sampled_from([(dtypes.float64, 709.5), (dtypes.float32, 88.7), (dtypes.float16, 11)]))
def test_exp_near_inf(self, dtype_x):
# reordering compute might return inf
dtype, x = dtype_x
if dtype not in supported_dtypes: return
with Context(TRANSCENDENTAL=2):
y = Tensor([x], dtype=dtype).exp().numpy()
expected = np.exp(np.array([x], dtype=_to_np_dtype(dtype)))
np.testing.assert_allclose(y, expected, rtol=5e-3)
class TestFromFuzzer(unittest.TestCase):
@given(strat.sampled_from(dtypes_float))
@unittest.skipUnless(dtypes.ulong in supported_dtypes, "Needs ulong")
def test_sin(self, dtype):
if dtype not in supported_dtypes: return
if dtype == dtypes.float64:
# crashes in CI CUDA
if DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}: return
def _test_value(n: float, unit: float=1.0):
next_float = np.nextafter(1.0, 2.0, dtype=_to_np_dtype(dtype))
ulp = next_float - 1.0
ulp = unit * ulp
with Context(TRANSCENDENTAL=2):
np.testing.assert_allclose(Tensor([n], dtype=dtype).sin().numpy(), np.sin(np.array([n], dtype=_to_np_dtype(dtype))), atol=ulp, rtol=1e-5)
_test_value(-35.0)
_test_value(-25.0)
_test_value(25.0)
_test_value(30.0) # 30.0 == switch_over
_test_value(35.0)
_test_value(0.0)
_test_value(np.pi / 2)
# worst case of ulp 1.5
_test_value(np.pi * 2, unit=1.5)
@given(strat.sampled_from(dtypes_float))
def test_log2(self, dtype):
if dtype not in supported_dtypes: return
if dtype == dtypes.float64:
# crashes in CI CUDA
if DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}: return
def _test_value(n: float, unit: float=1.0):
next_float = np.nextafter(1.0, 2.0, dtype=_to_np_dtype(dtype))
ulp = next_float - 1.0
ulp = unit * ulp
with Context(TRANSCENDENTAL=2):
np.testing.assert_allclose(Tensor([n], dtype=dtype).log2().numpy(), np.log2(np.array([n], dtype=_to_np_dtype(dtype))), atol=ulp, rtol=1e-5)
fmin = np.finfo(_to_np_dtype(dtype)).tiny
for scale in [1.0, 1e10, 1e20, 1e30]:
_test_value(fmin * scale)
_test_value(-fmin * scale)
_test_value(0)
_test_value(0.0000009)
class TestFloat16Log2(unittest.TestCase):
"""Tests for native float16 log2 implementation (no float32 cast)"""
@unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
def test_float16_log2_basic(self):
# basic values
test_values = [1.0, 2.0, 4.0, 0.5, 0.25, 10.0, 100.0, 1000.0]
with Context(TRANSCENDENTAL=2):
for val in test_values:
result = Tensor([val], dtype=dtypes.float16).log2().numpy()[0]
expected = np.log2(np.float16(val))
np.testing.assert_allclose(result, expected, rtol=1e-3, err_msg=f"log2({val})")
@unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
@unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "Nan handling differs on Vulkan")
def test_float16_log2_special(self):
# special values: inf, -inf, nan, 0, negative
with Context(TRANSCENDENTAL=2), np.errstate(all='ignore'):
# log2(inf) = inf
assert np.isinf(Tensor([np.inf], dtype=dtypes.float16).log2().numpy()[0])
# log2(0) = -inf
assert Tensor([0.0], dtype=dtypes.float16).log2().numpy()[0] == -np.inf
# log2(negative) = nan
assert np.isnan(Tensor([-1.0], dtype=dtypes.float16).log2().numpy()[0])
# log2(nan) = nan
assert np.isnan(Tensor([np.nan], dtype=dtypes.float16).log2().numpy()[0])
@unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
def test_float16_log2_denormal(self):
# test values near and below float16 min normal (6.1e-5)
# these exercise the denormal handling path with 2^10 scaling
test_values = [1e-4, 6e-5, 1e-5]
with Context(TRANSCENDENTAL=2):
for val in test_values:
result = Tensor([val], dtype=dtypes.float16).log2().numpy()[0]
expected = np.log2(np.float16(val))
# denormals have lower precision due to float16 limitations
np.testing.assert_allclose(result, expected, rtol=5e-2, err_msg=f"log2({val})")
class TestTranscendentalSchedule(unittest.TestCase):
@unittest.skipUnless(dtypes.ulong in supported_dtypes, "Needs ulong")
def test_transcendental_sin_fusion(self):
with Context(TRANSCENDENTAL=2):
a = Tensor.empty(10)
b = Tensor.empty(10)
c = a.sin() + b.sin()
c = c.sin()
check_schedule(c, 1)
def test_transcendental_log2_fusion(self):
with Context(TRANSCENDENTAL=2):
a = Tensor.empty(10)
b = Tensor.empty(10)
c = a.log2() + b.log2()
c = c.log2()
check_schedule(c, 1)
def test_transcendental_exp2_fusion(self):
with Context(TRANSCENDENTAL=2):
a = Tensor.empty(10)
b = Tensor.empty(10)
c = a.exp2() + b.exp2()
c = c.exp2()
check_schedule(c, 1)
class TestTranscendentalVectorized(unittest.TestCase):
def _vectorized_data(self, low, high, vec_size):
np_data = np.linspace(low, high, num=(128 // vec_size) * vec_size, dtype=np.float32).reshape(-1, vec_size)
data = Tensor(np_data, dtype=dtypes.float32)
return data, np_data
def _test_vectorized_op(self, fxn, np_fxn, data_range, vec_size, param_range=None):
data, np_data = self._vectorized_data(data_range[0], data_range[1], vec_size)
if param_range:
param, np_param = self._vectorized_data(param_range[0], param_range[1], vec_size)
out, np_out = fxn(data, param), np_fxn(np_data, np_param)
else:
out, np_out = fxn(data), np_fxn(np_data)
np.testing.assert_allclose(out.numpy(), np_out, rtol=1e-4)
def test_exp2_vectorized(self):
for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.exp2, np.exp2, (-100, 100), vec_size)
def test_log2_vectorized(self):
for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.log2, np.log2, (0.001, 200), vec_size)
@unittest.skipIf(Device.DEFAULT == "DSP", "requires int division")
@unittest.skipIf(DEV.renderer == "NAK", "MUFU.SIN is not accurate enough")
@unittest.skipIf(Device.DEFAULT == "WEBGPU" and OSX, "WEBGPU Metal backend is not accurate enough")
def test_sin_vectorized(self):
for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.sin, np.sin, (-100, 100), vec_size)
def test_pow_vectorized(self):
# np.pow returns nan for negative values raised to a non-integral power
for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.pow, np.pow, (0.001, 200), vec_size, param_range=(-10, 10))
def test_sqrt_vectorized(self):
for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.sqrt, np.sqrt, (0, 100), vec_size)
if __name__ == '__main__':
unittest.main()

345
artifacts/package_sources/tinygrad/test/backend/test_uops.py Normal file
View File

@@ -0,0 +1,345 @@
from typing import Optional, Any
import unittest, math
import numpy as np
from tinygrad.tensor import Tensor, _to_np_dtype
from tinygrad.helpers import Context
from tinygrad.dtype import dtypes, DType, AddrSpace, ConstFloat # noqa: F401
from tinygrad.device import Buffer, Device
from tinygrad.uop.ops import Ops, UOp, KernelInfo, AxisType, buffers
from tinygrad.renderer.cstyle import CStyleLanguage
from tinygrad.engine.realize import run_linear
from tinygrad.codegen import to_program
from tinygrad.codegen.opt import Opt, OptOps
from tinygrad.renderer.ptx import PTXRenderer
from test.helpers import to_uops_list
def run_uops(uops_list:list[UOp], bufs:list[Buffer]):
buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in bufs]
for u,b in zip(buf_uops, bufs): buffers[u] = b
run_linear(UOp(Ops.LINEAR, src=(UOp.sink(*uops_list, arg=KernelInfo()).call(*buf_uops),)))
def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
if op is Ops.CONST: uops.append(UOp.const(arg).cast(dtype))
elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, shape=(1,)))
else: uops.append(UOp(op, dtype, tuple(src), arg))
return uops[-1]
def _test_single_value(vals, op, dts):
uops = []
output_dtype = dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else dts[-1]
buf_store = uop(uops, Ops.PARAM, output_dtype, (), 0)
buf_loads = [uop(uops, Ops.PARAM, dtype, (), i+1) for i,dtype in enumerate(dts)]
loads = (buf_loads[i].index(uop(uops, Ops.CONST, dtypes.int32, (), 0)) for i, dtype in enumerate(dts))
alu = uop(uops, op, output_dtype, loads)
out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), alu))
buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
buf2 = [Buffer(Device.DEFAULT, 1, dtype, initial_value=np.array([a], dtype=_to_np_dtype(dtype)).tobytes()) for a,dtype in zip(vals, dts)]
run_uops([out], [buf]+buf2)
return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]
def _test_single_value_const(vals, op, dts):
uops = []
output_dtype = dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else dts[-1]
buf_store = uop(uops, Ops.PARAM, output_dtype, (), 0)
loads = (uop(uops, Ops.CONST, dtype, [], a) for a,dtype in zip(vals, dts))
alu = uop(uops, op, output_dtype, loads)
out = buf_store[UOp.const(0).cast(dtypes.int32)].store(alu)
buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
run_uops([out], [buf])
return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]
def _test_uops_result(output_dtype, uops, res):
# uops = []
buf_store = uop(uops, Ops.PARAM, output_dtype, (), 0)
# res = output_fn(uops)
out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), res))
buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
run_uops([out], [buf])
return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]
class TestUOps(unittest.TestCase):
def _equal(self, v1, v2):
assert isinstance(v2, (float, int, bool))
if isinstance(v2, float):
np.testing.assert_allclose(v1, v2, rtol=2e-7)
else:
np.testing.assert_equal(v1, v2)
def _test_uop_fxn(self, op, fxn, dts=(dtypes.float32, )):
for f in [_test_single_value, _test_single_value_const]:
for a in [-2.0, 0.0, 1.0]:
a = dts[0].const(a)
self._equal(f([a], op, dts), fxn(a))
def _test_bop_fxn(self, op, fxn, dts=(dtypes.float32, )*2, no_b_zero=False, no_b_neg=False):
for f in [_test_single_value, _test_single_value_const]:
for a in [-2.0, 0.0, 1.0]:
for b in [-3.0, 1.0] + ([] if no_b_zero else [0.0]):
a = dts[0].const(a)
b = dts[1].const(abs(b) if no_b_neg else b)
self._equal(f([a,b], op, dts), fxn(a,b))
def _test_top_fxn(self, op, fxn, dts=(dtypes.float32, )*3):
for f in [_test_single_value, _test_single_value_const]:
for a in [-2.0, 0, 1]:
for b in [-3.0, 3.0]:
for c in [-4.0, 4.0]:
a = dts[0].const(a)
b = dts[1].const(b)
c = dts[2].const(c)
self._equal(f([a,b,c], op, dts), fxn(a,b,c))
class TestFloatUOps(TestUOps):
@unittest.skipIf(Device.DEFAULT == "CPU", 'not supported as uop')
def test_exp2(self): self._test_uop_fxn(Ops.EXP2, lambda a: np.exp2(a))
@unittest.skipIf(Device.DEFAULT == "CPU", 'not supported as uop')
def test_log2(self): self._test_uop_fxn(Ops.LOG2, lambda a: math.log2(a) if a > 0 else float('-inf' if a==0 else 'nan'))
@unittest.skipIf(Device.DEFAULT == "CPU", 'not supported as uop')
def test_sin(self): self._test_uop_fxn(Ops.SIN, lambda a: math.sin(a))
def test_recip(self): self._test_uop_fxn(Ops.RECIPROCAL, lambda a: 1/a if a != 0 else float('inf'))
def test_sqrt(self): self._test_uop_fxn(Ops.SQRT, lambda a: math.sqrt(a) if a >= 0 else float('nan'))
def test_add(self): self._test_bop_fxn(Ops.ADD, lambda a,b: a+b)
def test_mul(self): self._test_bop_fxn(Ops.MUL, lambda a,b: a*b)
def test_max(self): self._test_bop_fxn(Ops.MAX, lambda a,b: max(a,b))
def test_cmplt(self): self._test_bop_fxn(Ops.CMPLT, lambda a,b: a<b)
def test_cmpne(self): self._test_bop_fxn(Ops.CMPNE, lambda a,b: a!=b)
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support NaN comparison correctly")
def test_cmpne_nan(self): # NaN != x for any x (IEEE 754)
for a, b in [(math.nan, 1.0), (1.0, math.nan), (math.nan, math.nan)]:
self.assertTrue(_test_single_value(
[dtypes.float32.const(a), dtypes.float32.const(b)],
Ops.CMPNE, (dtypes.float32, dtypes.float32)))
# MOD isn't tested on floats
def test_where(self):
self._test_top_fxn(Ops.WHERE, lambda a,b,c: b if a!=0 else c, (dtypes.bool, dtypes.float, dtypes.float))
@unittest.skipUnless(Device.DEFAULT == "PYTHON", "only python supports MULACC")
def test_mulacc(self):
self._test_top_fxn(Ops.MULACC, lambda a,b,c: a*b+c, (dtypes.float, dtypes.float, dtypes.float))
class TestNonFloatUOps(TestUOps):
def test_add_int32(self): self._test_bop_fxn(Ops.ADD, lambda a,b: int(a)+int(b), (dtypes.int32, dtypes.int32))
def test_mul_int32(self): self._test_bop_fxn(Ops.MUL, lambda a,b: int(a)*int(b), (dtypes.int32, dtypes.int32))
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, CStyleLanguage)), "only ptx and cstyle use bitshifts")
def test_shr_int32(self): self._test_bop_fxn(Ops.SHR, lambda a,b: int(a)>>int(b), (dtypes.int32, dtypes.int32), no_b_neg=True)
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, CStyleLanguage)), "only ptx and cstyle use bitshifts")
def test_shl_int32(self): self._test_bop_fxn(Ops.SHL, lambda a,b: int(a)<<int(b), (dtypes.int32, dtypes.int32), no_b_neg=True)
def test_div_int32(self):
self._test_bop_fxn(Ops.CDIV, lambda a,b: int(a/b), (dtypes.int32, dtypes.int32), no_b_zero=True)
def test_and_int32(self): self._test_bop_fxn(Ops.AND, lambda a,b: int(a)&int(b), (dtypes.int32, dtypes.int32))
def test_or_int32(self): self._test_bop_fxn(Ops.OR, lambda a,b: int(a)|int(b), (dtypes.int32, dtypes.int32))
def test_mod_int32(self):
self._test_bop_fxn(Ops.CMOD,
lambda a,b: abs(int(a))%abs(int(b))*(1,-1)[a<0], (dtypes.int32, dtypes.int32), no_b_zero=True)
def test_cmplt_int32(self): self._test_bop_fxn(Ops.CMPLT, lambda a,b: int(a)<int(b), (dtypes.int32, dtypes.int32))
def test_cmpne_int32(self): self._test_bop_fxn(Ops.CMPNE, lambda a,b: int(a)!=int(b), (dtypes.int32, dtypes.int32))
def test_mul_bool(self): self._test_bop_fxn(Ops.MUL, lambda a,b: bool(a) and bool(b), (dtypes.bool, dtypes.bool))
def test_where_float16(self):
self._test_top_fxn(Ops.WHERE, lambda a,b,c: b if a!=0 else c, (dtypes.bool, dtypes.float16, dtypes.float16))
class TestBoolUOps(TestUOps):
def _test_uop_bool_fxn(self, op, fxn):
for f in [_test_single_value, _test_single_value_const]:
for a in [False, True]:
self._equal(f([a], op, (dtypes.bool, )*1), fxn(a))
def _test_bop_bool_fxn(self, op, fxn):
for f in [_test_single_value, _test_single_value_const]:
for a in [False, True]:
for b in [False, True]:
self._equal(f([a,b], op, (dtypes.bool, )*2), fxn(a,b))
def _test_top_bool_fxn(self, op, fxn):
for f in [_test_single_value, _test_single_value_const]:
for a in [False, True]:
for b in [False, True]:
for c in [False, True]:
self._equal(f([a,b,c], op, (dtypes.bool, )*3), fxn(a,b,c))
def test_add_bool(self): self._test_bop_bool_fxn(Ops.ADD, lambda a,b: a or b)
def test_mul_bool(self): self._test_bop_bool_fxn(Ops.MUL, lambda a,b: a and b)
def test_xor_bool(self): self._test_bop_bool_fxn(Ops.XOR, lambda a,b: a != b)
def test_and_bool(self): self._test_bop_bool_fxn(Ops.AND, lambda a,b: a & b)
def test_or_bool(self): self._test_bop_bool_fxn(Ops.OR, lambda a,b: a | b)
def test_cmpne_bool(self): self._test_bop_bool_fxn(Ops.CMPNE, lambda a,b: a != b)
def test_cmplt_bool(self): self._test_bop_bool_fxn(Ops.CMPLT, lambda a,b: a < b)
def test_where_bool(self): self._test_top_bool_fxn(Ops.WHERE, lambda a,b,c: b if a else c)
class TestLocalAccess(unittest.TestCase):
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared memory")
def test_local_basic(self):
uops = []
smem = UOp.placeholder((16,), dtypes.float32, slot=0, addrspace=AddrSpace.LOCAL)
uops.append(smem)
st = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), uop(uops, Ops.CONST, dtypes.float32, (), 42.0)))
barr = uop(uops, Ops.BARRIER, dtypes.void, (st,))
sres = uop(uops, Ops.LOAD, dtypes.float32, (smem.after(barr).index(uop(uops, Ops.CONST, dtypes.int32, (), 0)),))
self.assertEqual(_test_uops_result(dtypes.float32, uops, sres), 42)
# NOTE: webgpu specific, since only webgpu performs bitpacking
@unittest.skipUnless(Device.DEFAULT == "WEBGPU", "Test local access with packed data type")
def test_local_packed(self):
uops = []
smem = UOp.placeholder((16,), dtypes.uint8, slot=0, addrspace=AddrSpace.LOCAL)
uops.append(smem)
st = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), uop(uops, Ops.CONST, dtypes.uint8, (), 42)))
barr = uop(uops, Ops.BARRIER, dtypes.void, (st,))
sres = smem.after(barr).index(uop(uops, Ops.CONST, dtypes.int32, (), 0))
self.assertEqual(_test_uops_result(dtypes.uint8, uops, sres), 42)
# NOTE: webgpu specific, since only webgpu performs bitpacking
@unittest.skipUnless(Device.DEFAULT == "WEBGPU", "Test local memory size for packed data types")
def test_packed_smem_size(self):
_dtypes = [dtypes.char, dtypes.uchar, dtypes.short, dtypes.ushort, dtypes.half]
size = 16
for dtype in _dtypes:
temp = UOp.placeholder((size,), dtype, slot=0, addrspace=AddrSpace.LOCAL)
uops = to_uops_list([temp], ren=Device[Device.DEFAULT].renderer)
out = Device[Device.DEFAULT].renderer.render(uops)
# half is supported in wgsl, so it doesn't have to be packed
corrected_size = size//(4//dtype.itemsize) if dtype != dtypes.half else size
# temp0: array<{Device[Device.DEFAULT].renderer.buf_map(dtype)},{corrected_size}>;
self.assertIn(f",{corrected_size}>;", out)
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared memory")
@unittest.skip("tinygrad doesn't support this behavior")
def test_local_indirect(self):
uops = []
smem = UOp.placeholder((16,), dtypes.int32, slot=0, addrspace=AddrSpace.LOCAL)
uops.append(smem)
st1 = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 1)), uop(uops, Ops.CONST, dtypes.int32, (), 2)))
st2 = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 2)), uop(uops, Ops.CONST, dtypes.int32, (), 42)))
barr = uop(uops, Ops.BARRIER, dtypes.void, (st1,st2))
ofs = uop(uops, Ops.LOAD, dtypes.int32, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 1)), barr))
sres = uop(uops, Ops.LOAD, dtypes.int32, (smem.index(ofs),))
self.assertEqual(_test_uops_result(dtypes.int32, uops, sres), 42)
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "This only tests assembly backends")
class TestAssembly(unittest.TestCase):
def test_bitshift_left(self):
g1 = UOp.param(0, dtypes.int32, shape=(3,))
out = UOp.param(1, dtypes.int32, shape=(2,))
c1 = UOp.const(2)
c2 = UOp.const(3)
l1 = g1.index(c1)
a1 = UOp(Ops.MUL, src=(l1, c1))
a2 = UOp(Ops.MUL, src=(l1, c2))
uops = to_uops_list([out.index(UOp.const(0)).store(a1), out.index(UOp.const(1)).store(a2)],
ren=Device[Device.DEFAULT].renderer)
Device[Device.DEFAULT].renderer.render(uops)
ops = [x.op for x in uops]
self.assertIn(Ops.SHL, ops)
self.assertIn(Ops.MUL, ops)
@unittest.skip("this is a questionable microoptimization i won't enforce")
def test_mulacc_unrolled(self):
# test that acc = acc + a0*b0 + a1*b1 + a2*b2 + a3*b3
# is not acc = acc + (a0*b0 + a1*b1 + a2*b2 + a3*b3)
a = Tensor.empty(1024)
b = Tensor.empty(1024)
c = (a*b).sum()
ast = c.schedule_linear().src[-1].src[0]
opts_to_apply = [Opt(OptOps.UNROLL, 0, 4)]
ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
program = to_program(ast, Device[Device.DEFAULT].renderer)
uops = tuple(program.src[1].src)
self.assertGreaterEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4)
def test_mulacc_shl(self):
g1 = UOp.param(0, dtypes.int32, shape=(2,))
c1 = UOp.const(0)
c2 = UOp.const(1)
expr = g1.index(c1) * UOp.const(4096) + g1.index(c2)
uops = to_uops_list([expr], ren=Device[Device.DEFAULT].renderer)
Device[Device.DEFAULT].renderer.render(uops)
self.assertIn(Ops.MULACC, [x.op for x in uops])
def test_use_cmpeq(self):
g = UOp.param(0, dtypes.uint32, shape=(8,))
c = UOp.const(7)
comp = g.index(c).ne(c).ne(True)
uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer)
Device[Device.DEFAULT].renderer.render(uops)
ops = [x.op for x in uops]
self.assertIn(Ops.CMPEQ, ops)
self.assertNotIn(Ops.CMPNE, ops)
class TestZeroRange(unittest.TestCase):
def test_reduce_variable(self):
for i in range(3,-1,-1):
v = UOp.variable("i", 0, 5).bind(i)
out = Tensor.ones(10, dtype=dtypes.int).contiguous().shrink(((0,v),)).sum()
self.assertEqual(out.item(), i)
class TestUOpPrograms(unittest.TestCase):
def _run(self, prog:UOp, *tensors:Tensor):
run_linear(UOp(Ops.LINEAR, src=(prog.call(*[t.uop.buf_uop for t in tensors]),)), update_stats=False)
def test_simple(self):
out = Tensor.empty(10,10,dtype=dtypes.int)
ptr = UOp.placeholder(out.shape, out.dtype, slot=0)
i, j = UOp.range(10, axis_id=0), UOp.range(10, axis_id=1)
prog = ptr[i,j].set(42).end(i,j)
self._run(prog.sink(arg=KernelInfo()), out)
with Context(DEBUG=0): self.assertTrue((out == 42).all().item())
def test_matmul(self):
a = Tensor.randn(10,10)
b = Tensor.randn(10,10)
c = Tensor.empty(10,10)
ref = (a@b)
with Context(DEBUG=0): Tensor.realize(a, b, c, ref)
# C[i,j] = sum_k A[i,k] * B[k,j]
# Shapes: A[M,K], B[K,N], C[M,N]
M = N = K = 10
DT = dtypes.float32
# Placeholders (bind slots explicitly)
A = UOp.placeholder((M, K), DT, slot=0)
B = UOp.placeholder((K, N), DT, slot=1)
C = UOp.placeholder((M, N), DT, slot=2)
# Axes: i,j are spatial; k is a reduction axis over the shared dim K
i = UOp.range(M, axis_id=0) # rows of A/C
j = UOp.range(N, axis_id=1) # cols of B/C
k = UOp.range(K, axis_id=2, axis_type=AxisType.REDUCE) # reduction over K
# Zero-init: write a scalar 0 to each (i,j).
C = C[i, j].set(0.0)
# Accumulate: C_after(k) enforces the dependency along the reduction axis
C = C[i, j].set(C.after(k)[i, j] + A[i, k] * B[k, j])
# Finalize the loop nest / schedule in (i, j, k) order
prog = C.end(i, j, k)
# run program
# TODO: make this work with opts_to_apply
self._run(prog.sink(arg=KernelInfo(opts_to_apply=())), a, b, c)
with Context(DEBUG=0): self.assertLessEqual((c-ref).square().mean().item(), 1e-6)
def test_matmul_relu(self):
a, b, c = Tensor.randn(10,10), Tensor.randn(10,10), Tensor.empty(10,10)
ref = (a@b).relu()
with Context(DEBUG=0): Tensor.realize(a, b, c, ref)
A, B, C = a.uop.placeholder_like(0), b.uop.placeholder_like(1), c.uop.placeholder_like(2)
i, j, k = UOp.range(10, 0), UOp.range(10, 1), UOp.range(10, 2, axis_type=AxisType.REDUCE)
C = C[i, j].set(0.0)
C = C[i, j].set(C.after(k)[i, j] + A[i, k] * B[k, j], end=k)
C = C[i, j].set(C[i, j].maximum(0.0))
prog = C.end(i, j)
self._run(prog.sink(arg=KernelInfo(opts_to_apply=())), a, b, c)
with Context(DEBUG=0): self.assertLessEqual((c-ref).square().mean().item(), 1e-6)
if __name__ == '__main__':
unittest.main(verbosity=2)

127
artifacts/package_sources/tinygrad/test/backend/test_wait_loop.py Normal file
View File

@@ -0,0 +1,127 @@
import unittest, threading
from tinygrad import Tensor, UOp
from tinygrad.device import Device, Buffer, BufferSpec
from tinygrad.dtype import AddrSpace, dtypes
from tinygrad.engine.realize import run_linear
from tinygrad.uop.ops import Ops, KernelInfo
def wait_loop_kernel(C:UOp) -> UOp:
N = 10
# a RANGE with no src is a bound-less loop header: a jump target with no induction variable.
# the compare and conditional backedge are expanded by the renderers from the loop RANGE/END
l = UOp.loop(0)
i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG)
# i = 0
i = i.after(i[0].store(0))
# i + 1, read loop-carried through after(l)
inc = i.after(l)[0].load() + 1
# i = inc; END(store, l, cond): conditional backedge, loop again while inc < N (do-while)
# NOTE: the cond uses the computed value, not a reload of the register
st = i[0].store(inc)
i = i.after(st.end(l, inc < N))
return C[0].store(i[0].load()).sink(arg=KernelInfo(name="wait_loop"))
def nested_loop_kernel(C:UOp) -> UOp:
r = UOp.range(4, 0)
l = UOp.loop(1)
i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG)
i = i.after(i[0].store(0))
inc = i.after(l, r)[0].load() + 1
st = i[0].store(inc)
lend = st.end(l, inc < (r.cast(dtypes.int)+1)*3)
i = i.after(lend.end(r))
return C[0].store(i[0].load()).sink(arg=KernelInfo(name="nested_loop", opts_to_apply=()))
def wait_ext_kernel() -> UOp:
sig = UOp.param(0, dtypes.int, (1,), volatile=True)
l = UOp.loop(0)
v = sig.after(l)[0].load()
e = v.end(l, v < 1)
return e.sink(arg=KernelInfo(name="wait_ext"))
def two_loops_kernel(C:UOp) -> UOp:
# two sequential loops on the same counter: ++ until 10, then ++ until 25
l1, l2 = UOp.loop(0), UOp.loop(1)
i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG)
i = i.after(i[0].store(0))
inc1 = i.after(l1)[0].load() + 1
i = i.after(i[0].store(inc1).end(l1, inc1 < 10))
inc2 = i.after(l2)[0].load() + 1
i = i.after(i[0].store(inc2).end(l2, inc2 < 25))
return C[0].store(i[0].load()).sink(arg=KernelInfo(name="two_loops", opts_to_apply=()))
def loop_in_loop_kernel(C:UOp) -> UOp:
# outer loop while i < 12, inner loop increments until i % 4 == 0 -> 12
l1, l2 = UOp.loop(0), UOp.loop(1)
i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG)
i = i.after(i[0].store(0))
inc = i.after(l1, l2)[0].load() + 1
st = i[0].store(inc)
# the outer END closes the inner END, and its cond reloads the register after the inner loop (in scope at the outer level)
e2 = st.end(l2, inc % 4 != 0)
oc = i.after(e2)[0].load()
i = i.after(e2.end(l1, oc < 12))
return C[0].store(i[0].load()).sink(arg=KernelInfo(name="loop_in_loop", opts_to_apply=()))
class TestWaitLoop(unittest.TestCase):
def test_wait_loop(self):
c = Tensor.empty(1, dtype=dtypes.int)
c = Tensor.custom_kernel(c, fxn=wait_loop_kernel)[0]
c.realize()
self.assertEqual(c.item(), 10)
def test_nested_loop_in_range(self):
c = Tensor.empty(1, dtype=dtypes.int)
c = Tensor.custom_kernel(c, fxn=nested_loop_kernel)[0]
c.realize()
self.assertEqual(c.item(), 12)
def test_two_sequential_loops(self):
c = Tensor.empty(1, dtype=dtypes.int)
c = Tensor.custom_kernel(c, fxn=two_loops_kernel)[0]
c.realize()
self.assertEqual(c.item(), 25)
def test_loop_in_loop(self):
c = Tensor.empty(1, dtype=dtypes.int)
c = Tensor.custom_kernel(c, fxn=loop_in_loop_kernel)[0]
c.realize()
self.assertEqual(c.item(), 12)
@unittest.skipUnless(Device.DEFAULT in ("CPU", "AMD", "NV"), "need proper uncached=True handling")
class TestVolatileLoops(unittest.TestCase):
def test_async_wait_ext(self):
sig_buf = Buffer(Device.DEFAULT, 1, dtypes.int, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
try: sig_view = sig_buf.as_memoryview(force_zero_copy=True).cast('i')
except (AssertionError, NotImplementedError): self.skipTest(f"{Device.DEFAULT} does not support host-visible buffers")
sig_view[0] = 0
def set_signal():
threading.Event().wait(0.3)
sig_view[0] = 1
sync = threading.Thread(target=set_signal, daemon=True)
sync.start()
run_linear(UOp(Ops.LINEAR, src=(wait_ext_kernel().call(UOp.from_buffer(sig_buf)),)), wait=True)
sync.join(timeout=3)
if __name__ == "__main__": unittest.main()

27
artifacts/package_sources/tinygrad/test/backend/test_zero_copy.py Normal file
View File

@@ -0,0 +1,27 @@
import unittest
from tinygrad import Tensor, Device
import time
def time_tensor_numpy(out:Tensor):
times = []
for _ in range(5):
st = time.perf_counter()
out.uop.base.realized.as_memoryview(allow_zero_copy=True)
et = time.perf_counter() - st
times.append(et)
return min(times)
N = 4096
class TestZeroCopy(unittest.TestCase):
@unittest.skipIf(Device.DEFAULT not in {"CPU", "METAL"}, "device isn't zero copy")
def test_zero_copy_from_default_to_cpu(self):
demo = Tensor.rand(1).realize()
t1 = time_tensor_numpy(demo)
out = Tensor.rand(N, N).realize()
t2 = time_tensor_numpy(out)
gbps = out.nbytes()*1e-9/max(t2-t1, 1e-10)
print(f"time(base): {t1*1e3:.2f} ms, time(copy): {t2*1e3:.2f} ms : copy speed {gbps:.2f} GB/s")
self.assertGreater(gbps, 600) # more than 600 GB/s = no copy
if __name__ == '__main__':
unittest.main(verbosity=2)