tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 12:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'master' into new_x86_backend

Browse Source
This commit is contained in:
ttomsa
2026-03-02 21:14:07 +00:00
committed by GitHub
parent 82954c7ca4 14d1c5fdfd
commit 393e591f49
52 changed files with 617 additions and 230 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/actions/setup-tinygrad/action.yml
+1 -1
View File
@@ -233,7 +233,7 @@ runs:
shell: bash
run: |
sudo mkdir -p /usr/local/lib
curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/nimlgen/amdcomgr_dylib/releases/latest | \
curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/tinygrad/amdcomgr_dylib/releases/latest | \
jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \
sudo xargs curl -fL -o /usr/local/lib/libamd_comgr.dylib
cargo build --release --manifest-path ./extra/remu/Cargo.toml
examples/mlperf/model_train.py
+10 -6
View File
@@ -1336,11 +1336,13 @@ def train_llama3():
# vocab_size from the mixtral tokenizer
if not SMALL: model_params |= {"vocab_size": 32000}
real_vocab_size = model_params['vocab_size']
if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP)
vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size
if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
print(f"model parameters: {model_params}")
# pad vocab
if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP)
vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size
model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
params = get_parameters(model)
# weights are all bfloat16 for now
@@ -1409,7 +1411,7 @@ def train_llama3():
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
tokens = tokens.shard(device)
logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
loss = vocab_mask.where(-float("inf"), logits).sparse_categorical_crossentropy(tokens[:, 1:])
loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
loss.backward()
assert all(p.grad is g for p,g in zip(optim.params, grads))
Tensor.realize(loss, *grads)
@@ -1439,13 +1441,15 @@ def train_llama3():
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
tokens = tokens.shard(device)
logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
loss = vocab_mask.where(-float("inf"), logits).sparse_categorical_crossentropy(tokens[:, 1:])
loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
return loss.flatten().float().to("CPU")
# ** data iters **
def fake_data(bs, samples):
import numpy as np
for _ in range(samples // bs):
yield Tensor.randint(bs, SEQLEN + 1, low=0, high=model_params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
fake_data_np = np.random.randint(0, model_params["vocab_size"], size=(bs, SEQLEN + 1), dtype=np.int32)
yield Tensor(fake_data_np, device="NPY")
def get_train_iter():
if getenv("FAKEDATA", 0):
@@ -1550,7 +1554,7 @@ def train_llama3():
# run eval
eval_losses = []
eval_iter = get_eval_iter()
tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")
tqdm.write(f"evaluating {EVAL_SAMPLES//EVAL_BS} batches of {EVAL_BS} sequences")
for j,tokens in tqdm(enumerate(eval_iter), total=EVAL_SAMPLES//EVAL_BS):
eval_losses += eval_step(tokens).tolist()
examples/mlperf/optim.py
+15 -2
View File
@@ -12,7 +12,20 @@ class GradAccClipAdamW(Optimizer):
self.v = self._new_optim_param()
self.grad_acc, self.clip_norm = grad_acc, clip_norm
def fstep(self, grads:list[Tensor]):
if self.fused:
out, extra = self._step([], grads)
updates = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
else:
updates, extra = self._step([], grads)
for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i]))
to_realize = extra+self.params+self.buffers
Tensor.realize(*to_realize)
def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
grads = list(grads)
for i in range(len(grads)):
if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device)
@@ -33,13 +46,13 @@ class GradAccClipAdamW(Optimizer):
ret = []
self.b1_t *= self.b1
self.b2_t *= self.b2
for i, (t, g) in enumerate(zip(params, grads)):
for i, g in enumerate(grads):
self.m[i].assign((self.b1 * self.m[i] + (1.0 - self.b1) * g).cast(self.m[i].dtype))
self.v[i].assign((self.b2 * self.v[i] + (1.0 - self.b2) * (g * g)).cast(self.v[i].dtype))
m_hat = self.m[i] / (1.0 - self.b1_t)
v_hat = self.v[i] / (1.0 - self.b2_t)
up = m_hat / (v_hat.sqrt() + self.eps)
ret.append((self.lr * up).cast(t.dtype))
ret.append((self.lr * up).cast(g.dtype))
return ret, [self.b1_t, self.b2_t] + self.m + self.v
def _apply_update(self, t:Tensor, up:Tensor) -> Tensor:
examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
+1
View File
@@ -5,6 +5,7 @@ export DEV=${DEV:-AMD}
export EMULATE="AMD_CDNA4"
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
export DEVICE_IN_FUNCTION_BUG=1
export DEBUG=${DEBUG:-0}
export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh
+1 -1
View File
@@ -3,4 +3,4 @@ export BENCHMARK=5
export EVAL_BS=0
export VIZ=${VIZ:--1}
examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
PYTHONPATH="." extra/viz/cli.py --profile --device "AMD" --top 20
extra/viz/cli.py --profile --device "AMD" --top 20
extra/torch_backend/test_kernel_fusion.py
+6 -9
View File
@@ -1,7 +1,6 @@
# simple tests
import unittest
import torch
import warnings
from tinygrad.helpers import getenv, GlobalCounters
if getenv("TINY_BACKEND2"):
import extra.torch_backend.backend2
@@ -18,9 +17,7 @@ class TestKernelFusionRegression(unittest.TestCase):
torch.manual_seed(42)
GlobalCounters.reset()
fn().detach().cpu().numpy()
expectation = f"{GlobalCounters.kernel_count} vs {expected_kernels} expected."
if GlobalCounters.kernel_count < expected_kernels: warnings.warn(f"{expectation} Expectation can be lowered.", UserWarning)
self.assertLessEqual(GlobalCounters.kernel_count, expected_kernels, f"{expectation}")
self.assertEqual(GlobalCounters.kernel_count, expected_kernels)
def test_elementwise_fusion(self):
def fn():
@@ -34,7 +31,7 @@ class TestKernelFusionRegression(unittest.TestCase):
conv = torch.nn.Conv2d(3, 16, 3, padding=1).to(device)
with torch.no_grad():
return torch.nn.functional.relu(conv(x))
self._check_kernel_count(fn, 8)
self._check_kernel_count(fn, 6)
def test_batchnorm_fusion(self):
def fn():
@@ -44,7 +41,7 @@ class TestKernelFusionRegression(unittest.TestCase):
bn.eval()
with torch.no_grad():
return torch.nn.functional.relu(bn(conv(x)))
self._check_kernel_count(fn, 16)
self._check_kernel_count(fn, 10)
def test_reduce_fusion(self):
def fn():
@@ -92,7 +89,7 @@ class TestKernelFusionRegression(unittest.TestCase):
out = bn(conv(x))
out += identity
return torch.nn.functional.relu(out)
self._check_kernel_count(fn, 17)
self._check_kernel_count(fn, 12)
def test_multiple_inplace_ops_fusion(self):
def fn():
@@ -117,7 +114,7 @@ class TestKernelFusionRegression(unittest.TestCase):
bn.train()
with torch.no_grad():
return bn(x)
self._check_kernel_count(fn, 10)
self._check_kernel_count(fn, 8)
# this is a minimal extra/other_mnist/beautiful_mnist_torch.py to cover fusion for training with optimizer
def test_mnist_training_fusion(self):
@@ -138,7 +135,7 @@ class TestKernelFusionRegression(unittest.TestCase):
loss.backward()
optimizer.step()
return loss
self._check_kernel_count(fn, 28)
self._check_kernel_count(fn, 24)
if __name__ == "__main__":
unittest.main()
extra/usbgpu/tbgpu/installer/TinyGPUDriverExtension/TinyGPUDriver.NV.Release.entitlements
+17
View File
@@ -0,0 +1,17 @@
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
<key>com.apple.application-identifier</key>
<string>9YG3G8543N.org.tinygrad.tinygpu.edriver</string>
<key>com.apple.developer.driverkit</key>
<true/>
<key>com.apple.developer.driverkit.transport.pci</key>
<array>
<dict>
<key>IOPCIPrimaryMatch</key>
<string>0x000010de&amp;0x0000FFFF</string>
</dict>
</array>
</dict>
</plist>
extra/usbgpu/tbgpu/installer/build_and_sign_nv.sh
Executable
+33
View File
@@ -0,0 +1,33 @@
#!/bin/bash
set -e
xcodebuild clean build CODE_SIGN_IDENTITY="" CODE_SIGNING_REQUIRED=NO -alltargets -configuration Release build
cp "../profiles/edriver_rel_2.provisionprofile" "./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext/embedded.provisionprofile"
cp "../profiles/installer_provisioning.provisionprofile" "./build/Release/TinyGPU.app/Contents/embedded.provisionprofile"
codesign \
--sign "Developer ID Application: tinygrad, Corp. (9YG3G8543N)" \
--entitlements ./TinyGPUDriverExtension/TinyGPUDriver.NV.Release.entitlements \
--verbose \
--options runtime \
--timestamp \
--force \
./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext
codesign \
--sign "Developer ID Application: tinygrad, Corp. (9YG3G8543N)" \
--entitlements ./macOS/macOS.entitlements \
--options runtime \
--verbose \
--timestamp \
--force \
./build/Release/TinyGPU.app
codesign --verify --deep --strict --verbose=4 ./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext
codesign --verify --deep --strict --verbose=4 ./build/Release/TinyGPU.app
spctl -a -vv ./build/Release/TinyGPU.app
spctl -a -vv ./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext
extra/viz/README
+3 -3
View File
@@ -1,17 +1,17 @@
A command line tool for exploring the VIZ trace.
After running with VIZ=-1, use `PYTHONPATH=. extra/viz/cli.py` to explore the saved trace files.
After running with VIZ=-1, use `extra/viz/cli.py` to explore the saved trace files.
## Inspect runtime profiling
Use `PYTHONPATH=. extra/viz/cli.py --profile` to list all traced devices.
Use `extra/viz/cli.py --profile` to list all traced devices.
List top slowest kernels on a device: `--profile --device "AMD"`
List samples of a kernel on a device: `--profile --device "AMD" --kernel E_3`
## Inspect codegen and PatternMatcher
Use `PYTHONPATH=. extra/viz/cli.py --rewrites` to list all traced kernels.
Use `extra/viz/cli.py --rewrites` to list all traced kernels.
List all codegen steps for a kernel: `--rewrites --kernel E_3`
Get source code: `--rewrites --kernel E_3 --select "View Source"`
extra/viz/cli.py
+48 -19
View File
@@ -1,44 +1,73 @@
#!/usr/bin/env python3
import os
os.environ["VIZ"] = "0"
import argparse, pathlib
import argparse, pathlib, sys, struct, json
from typing import Iterator
from tinygrad.viz import serve as viz
from tinygrad.uop.ops import RewriteTrace
from tinygrad.helpers import temp, ansistrip, colored, time_to_str, ansilen
from test.null.test_viz import load_profile
# ** generic helpers
def optional_eq(val:dict, arg:str|None) -> bool: return arg is None or ansistrip(val["name"]) == arg
def print_data(data:dict) -> None:
if isinstance(data.get("value"), Iterator):
for m in data["value"]:
if m.get("uop"):
print("Input UOp:")
print(m["uop"])
if not m["diff"]: continue
print("Rewrites:")
fp = pathlib.Path(m["upat"][0][0])
print(f"{fp.parent.name}/{fp.name}:{m['upat'][0][1]}")
print(m["upat"][1])
for line in m["diff"]:
color = "red" if line.startswith("-") else "green" if line.startswith("+") else None
print(colored(line, color))
if m.get("uop"): print(f"Input UOp:\n{m['uop']}")
if m.get("diff"):
loc = pathlib.Path(m["upat"][0][0])
print(f"Rewrite at {loc.parent.name}/{loc.name}:{m['upat'][0][1]}\n{m['upat'][1]}")
for line in m["diff"]: print(colored(line, "red" if line.startswith("-") else "green" if line.startswith("+") else None))
if data.get("src") is not None: print(data["src"])
# ** Profiler trace decoder
# 0 means None, otherwise it's an enum value
def option(i:int) -> int|None: return None if i == 0 else i-1
def decode_profile(data:bytes) -> dict:
ret, off = data, 0
def u(fmt:str) -> tuple:
nonlocal off
vals = struct.unpack_from(fmt, ret, off)
off += struct.calcsize(fmt)
return vals
total_dur, global_peak, index_len, layout_len = u("<IQII")
strings, dtypes, markers = json.loads(ret[off:off+index_len]).values()
off += index_len
layout:dict[str, dict] = {}
for _ in range(layout_len):
klen = u("<B")[0]
k = ret[off:off+klen].decode()
off += klen
layout[k] = v = {"events":[]}
event_type, event_count = u("<BI")
if event_type == 0:
for _ in range(event_count):
name, ref, key, st, dur, fmt = u("<IIIIfI")
v["events"].append({"name":strings[name], "ref":option(ref), "key":option(key), "st":st, "dur":dur, "fmt":strings[fmt]})
else:
v["peak"] = u("<Q")[0]
for _ in range(event_count):
alloc, ts, key = u("<BII")
if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIIB") for _ in range(u("<I")[0])]}})
return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}
if __name__ == "__main__":
parser = argparse.ArgumentParser()
g_mode = parser.add_argument_group("mode")
g_mode.add_argument("--profile", action="store_true", help="View profile trace")
g_mode.add_argument("--rewrites", action="store_true", help="View rewrites trace")
g_common = parser.add_argument_group("common options")
g_common.add_argument("--kernel", type=str, default=None, metavar="NAME", help="Select a kernel by name (optional name, default: only list names)")
g_profile = parser.add_argument_group("profile options")
g_profile.add_argument("--device", type=str, default=None, metavar="NAME", help="Select a device (optional name, default: only list names)")
g_profile.add_argument("--top", type=int, default=10, metavar="N", help="Number of top kernels to show (-1 for all, default: 10)")
g_rewrites = parser.add_argument_group("rewrites options")
g_rewrites.add_argument("--select", type=str, default=None, metavar="NAME",
help="Select an item within the chosen kernel (optional name, default: only list names)")
g_common = parser.add_argument_group("common options")
g_common.add_argument("--kernel", type=str, default=None, metavar="NAME", help="Select a kernel by name (optional name, default: only list names)")
parser.add_argument("--profile-path", type=pathlib.Path, metavar="PATH", help="Path to profile (optional file, default: latest profile)",
default=pathlib.Path(temp("profile.pkl", append_user=True)))
parser.add_argument("--rewrites-path", type=pathlib.Path, metavar="PATH", help="Path to rewrites (optional file, default: latest rewrites)",
@@ -46,14 +75,14 @@ if __name__ == "__main__":
args = parser.parse_args()
if not args.profile and not args.rewrites:
parser.print_help()
exit(0)
sys.exit(0)
viz.trace = viz.load_pickle(args.rewrites_path, default=RewriteTrace([], [], {}))
viz.ctxs = viz.get_rewrites(viz.trace)
if args.profile:
from tabulate import tabulate
profile = load_profile(viz.load_pickle(args.profile_path, default=[]))
profile = decode_profile(viz.get_profile(viz.load_pickle(args.profile_path, default=[])))
agg, total, n = {}, 0, 0
if args.device is None: print("Select a device:")
for k,v in profile["layout"].items():
@@ -63,7 +92,7 @@ if __name__ == "__main__":
for e in v.get("events", []):
et = e["dur"]*1e-6
if args.kernel is not None:
if ansistrip(e["name"]) == args.kernel and n < 10:
if optional_eq(e, args.kernel) and n < 10:
ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
name = e["name"]+(" " * (46 - ansilen(e["name"])))
print(f"{name} {ptm}/{(et or 0)*1e3:9.2f}ms "+e['fmt'].replace('\n', ' | ')+" ")
@@ -81,7 +110,7 @@ if __name__ == "__main__":
other_t = total-sum(t for _, (t, _) in sel)
table.append([f"Other ({len(other)} unique)", time_to_str(other_t, w=9), sum(c for _,(_,c) in other), f"{other_t/total*100.0:.2f}%"])
print(tabulate(table, headers=["name", "total", "count", "pct"], tablefmt="github"))
exit(0)
sys.exit(0)
for k in viz.ctxs:
if not optional_eq(k, args.kernel): continue
test/backend/test_multitensor.py
+11 -11
View File
@@ -228,17 +228,17 @@ class TestMultiTensor(unittest.TestCase):
a,b = _test_allreduce(Tensor.rand(256, 256))
np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
def test_multiple_to_single_device_naive(self):
with Context(RING=0):
t = Tensor.arange(32).shard(devices_4, 0).to(Device.DEFAULT).realize()
self.assertEqual(t.device, Device.DEFAULT)
np.testing.assert_equal(t.numpy(), np.arange(32))
def test_multiple_to_single_device_ring(self):
with Context(RING=2):
t = Tensor.arange(32).shard(devices_4, 0).to(Device.DEFAULT).realize()
self.assertEqual(t.device, Device.DEFAULT)
np.testing.assert_equal(t.numpy(), np.arange(32))
def test_multiple_to_single_device(self):
kernel_counts = {}
for ring in (0, 2):
GlobalCounters.reset()
with Context(RING=ring, SCACHE=0):
t = Tensor.arange(32).contiguous().shard(devices_4, 0).to(Device.DEFAULT)
t.realize()
kernel_counts[ring] = GlobalCounters.kernel_count
self.assertEqual(t.device, Device.DEFAULT)
np.testing.assert_equal(t.numpy(), np.arange(32))
self.assertNotEqual(kernel_counts[0], kernel_counts[2])
def test_allreduce_all2all(self):
with Context(ALL2ALL=2):
test/backend/test_schedule.py
+32
View File
@@ -795,6 +795,38 @@ class TestSchedule(unittest.TestCase):
self.assertIsNotNone(out.uop.base.realized)
self.assertIsInstance(out.uop.base.realized.dtype, ImageDType)
@unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
@unittest.expectedFailure
def test_image_dot_f16_fusion(self):
with Context(FLOAT16=1):
def cnt():
x, y, z = Tensor.empty((64, 64), dtype='float'), Tensor.empty((64, 64), dtype='float'), Tensor.empty((64, 64), dtype='float')
a = (x @ y).relu()
sched = ((a @ z).relu() + a).schedule()
for si in sched: si.lower()
return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
with Context(IMAGE=1): cnt1 = cnt()
with Context(IMAGE=2): cnt2 = cnt()
self.assertEqual(cnt1, cnt2)
@unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
@unittest.expectedFailure
def test_image_conv_fusion(self):
def cnt():
x, y, z = Tensor.empty((1, 4, 3, 3)), Tensor.empty((4, 1, 3, 3)), Tensor.empty((4, 1, 7, 7))
a = x.conv2d(y, Tensor.empty(4), groups=4, padding=1)
b = a.conv2d(z, groups=4, padding=3)
sched = (a + b).schedule()
for si in sched: si.lower()
return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
with Context(IMAGE=1): cnt1 = cnt()
with Context(IMAGE=2): cnt2 = cnt()
self.assertEqual(cnt1, cnt2)
def _test_fusion(self, shapes, f, cnt):
with Context(DEBUG=0, TRACK_MATCH_STATS=0): args = [Tensor.randn(s).realize() for s in shapes]
run_schedule(check_schedule(compare:=f(*args), cnt))
test/backend/test_symbolic_ops.py
+52
View File
@@ -69,6 +69,58 @@ class TestSymbolicOps(unittest.TestCase):
# 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)
test/backend/test_tensor_variable.py
+24
View File
@@ -136,6 +136,30 @@ class TestTensorVariable(unittest.TestCase):
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()
test/external/external_fuzz_beam_timeout_recovery.py
Vendored
+31
View File
@@ -0,0 +1,31 @@
#!/usr/bin/env python3
"""
Stress test for beam timeout + device recovery on AM devices.
Usage:
AMD=1 python test/external/external_test_beam_timeout_recovery.py
"""
from tinygrad import Tensor, Device
from tinygrad.helpers import Context
from tinygrad.runtime.ops_amd import AMDDevice
if __name__ == "__main__":
dev = Device["AMD"]
assert isinstance(dev, AMDDevice) and dev.is_am(), "not am"
N = 10000
for i in range(N):
with Context(DEBUG=0, BEAM=0):
a = Tensor.rand(4096, 4096, device="AMD").contiguous().realize()
b = Tensor.rand(4096, 4096, device="AMD").contiguous().realize()
c = a.matmul(b)
c.realize()
try: dev.synchronize(timeout=1)
except RuntimeError as e: print(e)
with Context(DEBUG=0, BEAM=0):
a = Tensor.ones(512, 512, device="AMD").contiguous().realize()
b = Tensor.ones(512, 512, device="AMD").contiguous().realize()
result = a.matmul(b).realize()[0, 0].item()
assert result == 512.0, f"iter {i}: got {result}"
print(f" iter {i+1}/{N}: ok")
print(f"=== All {N} iterations passed ===")
test/external/external_test_gpu_crash.py
Vendored
+35 -25
View File
@@ -4,13 +4,19 @@
These tests intentionally cause GPU faults to verify error handling.
Run with: AMD=1 python -m pytest test/external/external_test_gpu_crash.py -v
"""
import unittest, re
import unittest, re, importlib
from tinygrad.device import Device
from tinygrad.runtime.autogen.amd.rdna3.ins import * # noqa: F403
from tinygrad.renderer.amd.dsl import s, v, Inst, NULL
def assemble(code:str, name:str="test") -> str:
kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "wavefront_size32": 1, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8}
RDNA3_CDNA3_MAP = {"v_mov_b32_e32": "v_mov_b32_e32", "s_mov_b32": "s_mov_b32", "s_waitcnt": "s_waitcnt", "s_endpgm": "s_endpgm",
"global_load_b32": "global_load_dword", "global_store_b32": "global_store_dword",
"global_atomic_add_u32": "global_atomic_add", "flat_load_b32": "flat_load_dword",
"flat_store_b32": "flat_store_dword", "flat_atomic_add_u32": "flat_atomic_add", "s_load_b32": "s_load_dword"}
def assemble(code:str, name:str="test", is_cdna:bool=False) -> str:
kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8}
if is_cdna: kd["accum_offset"] = 8
else: kd["wavefront_size32"] = 1
return f".text\n.globl {name}\n.p2align 8\n.type {name},@function\n{name}:\n{code}\n.rodata\n.p2align 6\n.amdhsa_kernel {name}\n" + \
"\n".join(f".amdhsa_{k} {v}" for k,v in kd.items()) + "\n.end_amdhsa_kernel"
@@ -21,6 +27,10 @@ class TestGPUCrash(unittest.TestCase):
from tinygrad.runtime.support.compiler_amd import HIPCompiler
cls.dev = Device["AMD"]
cls.compiler = HIPCompiler(cls.dev.arch)
cls.is_cdna = cls.dev.target[0] < 10
ins = importlib.import_module('tinygrad.runtime.autogen.amd.' + ('cdna' if cls.is_cdna else 'rdna3') + '.ins')
for rdna3_name, cdna3_name in RDNA3_CDNA3_MAP.items():
setattr(cls, rdna3_name, getattr(ins, cdna3_name if cls.is_cdna else rdna3_name))
def setUp(self):
# Verify device works before each test
@@ -33,7 +43,7 @@ class TestGPUCrash(unittest.TestCase):
def _run(self, code: str):
from tinygrad.runtime.ops_amd import AMDProgram
prg = AMDProgram(self.dev, "test", self.compiler.compile(assemble(code)))
prg = AMDProgram(self.dev, "test", self.compiler.compile(assemble(code, is_cdna=self.is_cdna)))
prg(self.dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True)
def _run_insts(self, insts: list[Inst]):
@@ -57,32 +67,32 @@ class TestOutOfBoundsMemoryAccess(TestGPUCrash):
def test_global_load_null_ptr(self):
"""Global load from NULL pointer."""
insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0),
global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0),
self.global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
def test_global_store_null_ptr(self):
"""Global store to NULL pointer."""
insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 0xDEADBEEF),
global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 0xDEADBEEF),
self.global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
def test_global_load_unmapped_high_address(self):
"""Global load from high unmapped address (0xDEAD00000000)."""
insts = [v_mov_b32_e32(v[0], 0x00000000), v_mov_b32_e32(v[1], 0xDEAD),
global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
insts = [self.v_mov_b32_e32(v[0], 0x00000000), self.v_mov_b32_e32(v[1], 0xDEAD),
self.global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
def test_global_store_unmapped_high_address(self):
"""Global store to high unmapped address."""
insts = [v_mov_b32_e32(v[0], 0x00000000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 0x12345678),
global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
insts = [self.v_mov_b32_e32(v[0], 0x00000000), self.v_mov_b32_e32(v[1], 0xDEAD), self.v_mov_b32_e32(v[2], 0x12345678),
self.global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
def test_global_atomic_unmapped(self):
"""Atomic operation on unmapped memory."""
insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 1),
global_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
insts = [self.v_mov_b32_e32(v[0], 0xBEEF0000), self.v_mov_b32_e32(v[1], 0xDEAD), self.v_mov_b32_e32(v[2], 1),
self.global_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
@@ -91,14 +101,14 @@ class TestSMEMFaults(TestGPUCrash):
def test_smem_load_null(self):
"""SMEM load from NULL base."""
insts = [s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(0), s_endpgm()]
insts = [self.s_mov_b32(s[2], 0), self.s_mov_b32(s[3], 0),
self.s_load_b32(s[4], s[2:3], 0, soffset=NULL), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
def test_smem_load_unmapped(self):
"""SMEM load from unmapped address."""
insts = [s_mov_b32(s[2], 0xBEEF0000), s_mov_b32(s[3], 0xDEAD),
s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(0), s_endpgm()]
insts = [self.s_mov_b32(s[2], 0xBEEF0000), self.s_mov_b32(s[3], 0xDEAD),
self.s_load_b32(s[4], s[2:3], 0, soffset=NULL), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
@@ -107,20 +117,20 @@ class TestFlatMemoryFaults(TestGPUCrash):
def test_flat_load_null(self):
"""FLAT load from NULL address."""
insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0),
flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0),
self.flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
def test_flat_store_null(self):
"""FLAT store to NULL address."""
insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 0xDEADBEEF),
flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 0xDEADBEEF),
self.flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
def test_flat_atomic_null(self):
"""FLAT atomic on NULL address."""
insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 1),
flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 1),
self.flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
self._assert_gpu_fault(lambda: self._run_insts(insts))
test/external/external_test_speed_llama.py
Vendored
+1 -1
View File
@@ -10,7 +10,7 @@ from tinygrad.helpers import Profiling
class FakeProgram:
def __init__(self, name:str, prg:bytes, **kwargs): pass
def __call__(self, *bufs, global_size, local_size, vals=(), wait=False): pass
def __call__(self, *bufs, global_size, local_size, vals=(), wait=False, **kw): pass
class FakeAllocator(Allocator[Compiled]):
def _alloc(self, sz, options): return None
test/models/test_onnx.py
+1 -1
View File
@@ -87,7 +87,7 @@ class TestHuggingFaceOnnxModels(unittest.TestCase):
"input_ids": np.random.randint(0, 250002, (1, 11), dtype=np.int64),
"attention_mask": np.ones((1, 11), dtype=np.int64),
}
self._validate(repo_id, model_file, custom_inputs)
self._validate(repo_id, model_file, custom_inputs, atol=1e-3)
if __name__ == "__main__":
unittest.main()
test/null/test_schedule.py
+12
View File
@@ -234,6 +234,18 @@ class TestSchedule(unittest.TestCase):
d = Tensor.empty(1).assign(c)
check_schedule(d, 1)
def test_detach_assign(self):
a = Tensor.ones(4, 4).contiguous().realize()
buf1, buf2 = Tensor.empty(4, 4).contiguous(), Tensor.empty(4, 4).contiguous()
r = buf2.assign(buf1.assign(a + 1.0) * 2.0)
check_schedule(r.detach().contiguous(), 2)
def test_contiguous_backward_assign(self):
a = Tensor.ones(4, 4).contiguous().realize()
buf1, buf2 = Tensor.empty(4, 4).contiguous(), Tensor.empty(4, 4).contiguous()
r = buf2.assign(buf1.assign(a + 1.0) * 2.0)
check_schedule(r.contiguous_backward().contiguous(), 2)
def test_mulacc_relu_fusion(self):
a = Tensor.empty(10)
b = Tensor.empty(10)
test/null/test_viz.py
+3 -35
View File
@@ -1,4 +1,4 @@
import unittest, decimal, json, struct, sys
import unittest, decimal, sys
from dataclasses import dataclass
from typing import Generator
@@ -357,41 +357,9 @@ class TestVizIntegration(BaseTestViz):
from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry
from tinygrad.viz.serve import get_profile
from extra.viz.cli import decode_profile
class TinyUnpacker:
def __init__(self, buf): self.buf, self.offset = buf, 0
def __call__(self, fmt:str) -> tuple:
ret = struct.unpack_from(fmt, self.buf, self.offset)
self.offset += struct.calcsize(fmt)
return ret
# 0 means None, otherwise it's an enum value
def option(i:int) -> int|None: return None if i == 0 else i-1
def load_profile(lst:list[ProfileEvent]) -> dict:
ret = get_profile(lst)
u = TinyUnpacker(ret)
total_dur, global_peak, index_len, layout_len = u("<IQII")
strings, dtypes, markers = json.loads(ret[u.offset:u.offset+index_len]).values()
u.offset += index_len
layout:dict[str, dict] = {}
for _ in range(layout_len):
klen = u("<B")[0]
k = ret[u.offset:u.offset+klen].decode()
u.offset += klen
layout[k] = v = {"events":[]}
event_type, event_count = u("<BI")
if event_type == 0:
for _ in range(event_count):
name, ref, key, st, dur, fmt = u("<IIIIfI")
v["events"].append({"name":strings[name], "ref":option(ref), "key":option(key), "st":st, "dur":dur, "fmt":strings[fmt]})
else:
v["peak"] = u("<Q")[0]
for _ in range(event_count):
alloc, ts, key = u("<BII")
if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIIB") for _ in range(u("<I")[0])]}})
return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}
def load_profile(lst:list[ProfileEvent]) -> dict: return decode_profile(get_profile(lst))
class TestVizProfiler(BaseTestViz):
def test_transfer_uses_copy_device(self):
test/unit/test_function.py
+105
View File
@@ -70,6 +70,14 @@ class TestFunction(unittest.TestCase):
b = Tensor([4,5,6])
np.testing.assert_equal(f(a, b).numpy(), [5,7,9])
def test_contiguous_backward(self):
@function
def f(a:Tensor, b:Tensor) -> Tensor: return (a + b).contiguous_backward()
a = Tensor([1,2,3])
b = Tensor([4,5,6])
np.testing.assert_equal(f(a, b).numpy(), [5,7,9])
def test_method(self):
class Foo:
def __init__(self): self.w = Tensor([10,20,30])
@@ -129,6 +137,31 @@ class TestFunction(unittest.TestCase):
a = Tensor([1., 2., 3.])
np.testing.assert_allclose(g(f(a)).numpy(), [110., 440., 990.])
def test_nested_calls_backward(self):
w = Tensor([[1., 2.], [3., 4.]]).contiguous().realize()
@function
def inner(x:Tensor) -> Tensor: return x + w
@function
def outer(a:Tensor, b:Tensor) -> Tensor: return inner(a.reshape(1,2) + b.reshape(1,2))
a = Tensor([1., 2.], requires_grad=True)
b = Tensor([3., 4.], requires_grad=True)
outer(a, b).sum().backward()
np.testing.assert_allclose(a.grad.numpy(), [2., 2.])
np.testing.assert_allclose(b.grad.numpy(), [2., 2.])
def test_unused_param_backward(self):
@function
def f(a:Tensor, b:Tensor, c:Tensor) -> Tensor: return a + c # b is unused
a = Tensor([1., 2., 3.], requires_grad=True)
b = Tensor([4., 5., 6.], requires_grad=True)
c = Tensor([7., 8., 9.], requires_grad=True)
f(a, b, c).sum().backward()
np.testing.assert_allclose(a.grad.numpy(), [1., 1., 1.])
np.testing.assert_allclose(b.grad.numpy(), [0., 0., 0.])
np.testing.assert_allclose(c.grad.numpy(), [1., 1., 1.])
def test_name(self):
@function
def f(a:Tensor) -> Tensor: return a + 1
@@ -230,5 +263,77 @@ class TestFunctionMulti(unittest.TestCase):
f(x).sum().backward()
np.testing.assert_allclose(w.grad.numpy(), [4., 5., 6., 7.])
def test_call_axis(self):
@function
def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]]).shard(self.devices_2, axis=0)
w = Tensor([[1.,2.],[3.,4.]]).shard(self.devices_2, axis=None)
result = f(x, w)
# CALL output should inherit axis=0 from the sharded input
self.assertEqual(result.uop.axis, 0)
# reduce on the sharded axis should remove it
self.assertIsNone(result.sum().uop.axis)
def test_call_axis_shard_inside(self):
@function
def f(x:Tensor, w:Tensor) -> Tensor:
return x.shard(self.devices_2, axis=0) @ w.shard(self.devices_2, axis=None)
x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]])
w = Tensor([[1.,2.],[3.,4.]])
result = f(x, w)
self.assertEqual(result.uop.axis, 0)
np.testing.assert_allclose(result.numpy(), x.numpy() @ w.numpy())
def test_data_parallel_backward(self):
@function
def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]], requires_grad=True).shard(self.devices_2, axis=0)
w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(self.devices_2, axis=None)
w.realize()
f(x, w).sum().backward()
# d/dx = ones @ w^T = [[1,3],[1,3],[1,3],[1,3]], but sum so ones(4,2) @ w^T? no:
# L = sum(x @ w), dL/dx = ones(4,2) @ w^T... actually dL/d(xw) = ones(4,2), dL/dx = ones(4,2) @ w^T
np.testing.assert_allclose(x.grad.numpy(), np.ones((4,2)) @ np.array([[1,3],[2,4]]))
def test_data_parallel_backward_4(self):
devices_4 = tuple(f"CPU:{i}" for i in range(4))
@function
def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
x = Tensor(np.arange(16).reshape(8,2).astype(np.float32), requires_grad=True).shard(devices_4, axis=0)
w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(devices_4, axis=None)
w.realize()
f(x, w).sum().backward()
np.testing.assert_allclose(x.grad.numpy(), np.ones((8,2)) @ np.array([[1,3],[2,4]]))
def test_data_parallel_backward_implicit(self):
devices_4 = tuple(f"CPU:{i}" for i in range(4))
w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(devices_4, axis=None)
w.realize()
@function
def f(x:Tensor) -> Tensor: return x @ w
x = Tensor(np.arange(16).reshape(8,2).astype(np.float32), requires_grad=True).shard(devices_4, axis=0)
f(x).sum().backward()
np.testing.assert_allclose(x.grad.numpy(), np.ones((8,2)) @ np.array([[1,3],[2,4]]))
def test_data_parallel_backward_twice(self):
devices_4 = tuple(f"CPU:{i}" for i in range(4))
w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(devices_4, axis=None)
w.realize()
# pre-init grads like the training loop does
w.grad = w.zeros_like().contiguous().realize()
@function
def f(x:Tensor) -> Tensor: return x @ w
expected = np.ones((8,2)) @ np.array([[1,3],[2,4]])
for _ in range(2):
x = Tensor(np.arange(16).reshape(8,2).astype(np.float32), requires_grad=True).shard(devices_4, axis=0)
f(x).sum().backward()
np.testing.assert_allclose(x.grad.numpy(), expected)
if __name__ == '__main__':
unittest.main()
test/unit/test_linalg.py
+9
View File
@@ -83,6 +83,15 @@ class TestLinAlg(unittest.TestCase):
s_diag = (S.unsqueeze(-2) * Tensor.eye(2))
reconstruction_helper([U, s_diag, V], a)
def test_svd_identity_4x4(self):
a = Tensor.eye(4)
U,S,V = a.svd()
assert not np.isnan(U.numpy()).any()
assert not np.isnan(S.numpy()).any()
assert not np.isnan(V.numpy()).any()
s_diag = (S.unsqueeze(-2) * Tensor.eye(4))
reconstruction_helper([U, s_diag, V], a)
def test_svd_rank1(self):
a = Tensor([[1.0, 1.0], [2.0, 2.0]]).realize()
U, S, V = a.svd()
tinygrad/apps/llm.py
+14 -10
View File
@@ -143,7 +143,8 @@ class TransformerBlock:
#v = self.cache_kv[1, :, :, 0:start_pos+T, :]
# NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).triu(int(start_pos)+1) if T > 1 else None
# TODO: this if statement should be removed and it shouldn't generate extra kernels
mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).triu(start_pos+1) if T > 1 else None
attn = q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True) # (B,H,T,Hd)
attn = attn.transpose(1, 2).reshape(B, T, -1) # back to (B,T,D)
attn = self.attn_output(attn)
@@ -263,7 +264,7 @@ CHAT_HTML = b'''<!DOCTYPE html><html><head><title>tinygrad chat</title><style>
</style></head><body><div id="chat"></div>
<textarea id="input" rows="1" placeholder="Ask anything"></textarea>
<script>
input.onkeydown = (e) => { if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); send() } }
input.onkeydown = (e) => { if (e.key === 'Enter' && !e.shiftKey && !e.isComposing) { e.preventDefault(); send() } }
const msgs = [];
async function send() {
if (!input.value.trim()) return;
@@ -344,20 +345,23 @@ if __name__ == "__main__":
args = parser.parse_args()
# load the model
model, kv = Transformer.from_gguf(Tensor.from_url(models[args.model]), args.max_context)
if DEBUG >= 1: print(f"using model {args.model}")
raw_model = Tensor.from_url(models[args.model])
model, kv = Transformer.from_gguf(raw_model, args.max_context)
if DEBUG >= 1 or args.benchmark:
print(f"using model {args.model} with {raw_model.nbytes():,} bytes and {sum(x.numel() for x in nn.state.get_parameters(model)):,} params")
del raw_model
# TODO: why this is required to free the RAM of the GGUF copy?
import gc
gc.collect()
# do benchmark
if args.benchmark:
param_bytes = sum(x.nbytes() for x in nn.state.get_parameters(model))
for b in model.blk:
if hasattr(b, 'ffn_gate_exps'):
expert_bytes = b.ffn_gate_exps.weight.nbytes() + b.ffn_up_exps.weight.nbytes() + b.ffn_down_exps.weight.nbytes()
param_bytes -= int(expert_bytes * (1 - b.num_experts_per_tok / b.ffn_gate_exps.weight.shape[0]))
gen = model.generate([0], 0)
for _ in range(args.benchmark):
GlobalCounters.reset()
with Timing(on_exit=lambda x: f", {1e9/x:6.2f} tok/s, {GlobalCounters.global_mem/x:7.2f} GB/s, param {param_bytes/x:7.2f} GB/s"): next(gen)
with Timing(on_exit=lambda x: f", {1e9/x:6.2f} tok/s, {GlobalCounters.global_mem/x:7.2f} GB/s,"
f" {GlobalCounters.global_mem//1000000}/{GlobalCounters.mem_used//1000000} MB"): next(gen)
exit(0)
# extract some metadata
tinygrad/codegen/late/devectorizer.py
+7 -7
View File
@@ -1,7 +1,7 @@
from typing import Any, cast
import functools, itertools
from collections import defaultdict
from dataclasses import dataclass, field
from dataclasses import dataclass
from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace, Invalid, PtrDType
from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, identity_element
from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
@@ -311,8 +311,6 @@ pm_render = PatternMatcher([
@dataclass
class ReduceContext:
acc_num: int = 0
# track ENDs by range for merging parallel reduces
range_to_ends: dict[tuple[UOp, ...], list[UOp]] = field(default_factory=dict)
def horizontal_reduce(inp:UOp, out_dtype:DType) -> list[UOp]:
# if this has a horizontal reduction component, do that first
@@ -338,13 +336,15 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp):
ctx.acc_num += 1
ret = functools.reduce(lambda x,y: x.alu(red.arg, y), lst)
if len(reduce_range) == 0: return ret
end = acc.index(UOp.const(dtypes.int, 0)).store(ret).end(*reduce_range)
ctx.range_to_ends.setdefault(reduce_range, []).append(end)
end = acc.index(UOp.const(dtypes.int, 0)).store(ret).end(*reduce_range).rtag("mergeable")
return acc.after(end).index(UOp.const(dtypes.int, 0))
def merge_reduce_ends(ctx:ReduceContext, sink:UOp):
# merge ENDs that share the same range
subs = {e: UOp.group(*(e.src[0] for e in ends)).end(*r) for r, ends in ctx.range_to_ends.items() if len(ends) > 1 for e in ends}
# merge ENDs that share the same range (only those created by reduce_to_acc)
range_to_ends: dict[tuple[UOp, ...], list[UOp]] = {}
for u in sink.backward_slice:
if u.op is Ops.END and u.tag == "mergeable": range_to_ends.setdefault(u.src[1:], []).append(u)
subs = {e: UOp.group(*(e.src[0] for e in ends)).end(*r) for r, ends in range_to_ends.items() if len(ends) > 1 for e in ends}
return sink.substitute(subs) if subs else None
pm_reduce = PatternMatcher([
tinygrad/codegen/opt/search.py
+5 -3
View File
@@ -36,7 +36,8 @@ def get_test_global_size(global_size, max_global_size, var_vals):
return test_global_size, input_size / prod(test_global_size)
def _time_program(p:ProgramSpec, lib:bytes, var_vals:dict[str, int], rawbufs:list[Buffer], early_stop:float|None=None,
allow_test_size:int=True, max_global_size:int|None=65536, clear_l2=False, cnt=3, name="test") -> list[float]:
allow_test_size:int=True, max_global_size:int|None=65536, clear_l2=False, cnt=3, name="test", dev_timeout=False) -> list[float]:
timeout = int(early_stop * 1e3) if dev_timeout and early_stop is not None and early_stop < math.inf else None
factor = 1
if allow_test_size and max_global_size is not None:
global_size, factor = get_test_global_size(p.global_size, max_global_size, var_vals)
@@ -50,7 +51,7 @@ def _time_program(p:ProgramSpec, lib:bytes, var_vals:dict[str, int], rawbufs:lis
if hasattr(dev:=Device[p.device], 'invalidate_caches'): dev.invalidate_caches()
else:
with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024,1024).contiguous().realize(do_update_stats=False)
tms.append(unwrap(car(input_bufs, var_vals, wait=True))*factor)
tms.append(unwrap(car(input_bufs, var_vals, wait=True, timeout=timeout))*factor)
if early_stop is not None and early_stop < min(tms): break
return tms
@@ -161,7 +162,8 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True
continue
seen_libs.add(lib)
try: tms = _time_program(p, lib, var_vals, rawbufs, early_stop=beam[0][1]*3 if len(beam) else 1.0,
allow_test_size=allow_test_size, clear_l2=hasattr(dev, 'invalidate_caches'))
allow_test_size=allow_test_size, clear_l2=hasattr(dev, 'invalidate_caches'),
dev_timeout=getenv("BEAM_DEV_TIMEOUT", 1))
except Exception as e:
if BEAM_DEBUG: print(f"BEAM failed for opts: {candidates[i].applied_opts}\n{e}")
if isinstance(e, RuntimeError): continue
tinygrad/engine/realize.py
+2 -2
View File
@@ -50,7 +50,7 @@ class CompiledRunner(Runner):
def __reduce__(self): return self.__class__, (self.p,)
def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int]|None=None, wait=False) -> float|None:
def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int]|None=None, wait=False, timeout:int|None=None) -> float|None:
if var_vals is None: var_vals = {}
global_size, local_size = self.p.launch_dims(var_vals)
if Device[self.p.device].renderer.has_local and local_size is None and all_int(self.p.global_size):
@@ -58,7 +58,7 @@ class CompiledRunner(Runner):
global_size = [g//l if g%l == 0 else g/l for g,l in zip(global_size, local_size)]
self.p = replace(self.p, global_size=global_size, local_size=local_size)
return self._prg(*[x._buf for x in rawbufs], global_size=tuple(global_size), local_size=tuple(local_size) if local_size else None,
vals=tuple(var_vals[k.expr] if k.expr not in self.p.runtimevars else None for k in self.p.vars), wait=wait)
vals=tuple(var_vals[k.expr] if k.expr not in self.p.runtimevars else None for k in self.p.vars), wait=wait, timeout=timeout)
class ViewOp(Runner):
def __init__(self, buf:Buffer): super().__init__(colored(f"view {buf.nbytes:8d} @ {buf.offset:<10d}", "yellow"), buf.device)
tinygrad/gradient.py
+5 -4
View File
@@ -17,12 +17,13 @@ def call_gradient(ctx:UOp, k:UOp) -> tuple[UOp|None, ...]:
if k.arg.grad_fxn is not None: return (None,) + k.arg.grad_fxn(ctx, k)
# auto-differentiate the function
fxn, args = k.src[0], k.src[1:]
params = sorted([x for x in fxn.toposort() if x.op == Ops.PARAM], key=lambda x: x.arg)
grads = compute_gradient(fxn, ctx.param_like(len(args)), set(params))
params = {x.arg:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM}
grads = compute_gradient(fxn, ctx.param_like(len(args)), set(params.values()))
ret: list[UOp|None] = [None]
for i,p in enumerate(params):
if p in grads:
for i in range(len(args)):
if (p:=params.get(i, None)) is not None and p in grads:
# TODO: compact the args and remove unused ones
assert not grads[p].op_in_backward_slice_with_self(Ops.BUFFER), "BUG: BUFFER in backward slice of grad"
ret.append(grads[p].call(*args, ctx, name=(k.arg.name or "")+f"_backward_{i}"))
else:
ret.append(None)
tinygrad/helpers.py
+2 -1
View File
@@ -172,7 +172,8 @@ class ContextVar(Generic[T]):
assert isinstance(self.value, str)
return [getattr(obj, x) if obj else x for x in self.value.split(',') if x]
DEBUG, IMAGE, BEAM, NOOPT = ContextVar("DEBUG", 0), ContextVar("IMAGE", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0)
DEBUG, BEAM, NOOPT = ContextVar("DEBUG", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0)
IMAGE, FLOAT16 = ContextVar("IMAGE", 0), ContextVar("FLOAT16", 0)
JIT, JIT_BATCH_SIZE = ContextVar("JIT", 2 if OSX and ARCH_X86 else 1), ContextVar("JIT_BATCH_SIZE", 32)
WINO, CAPTURING, TRACEMETA = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1)
USE_TC, TC_SELECT, TC_OPT, AMX = ContextVar("TC", 1), ContextVar("TC_SELECT", -1), ContextVar("TC_OPT", 0), ContextVar("AMX", 0)
tinygrad/nn/optim.py
+4 -4
View File
@@ -17,6 +17,7 @@ class Optimizer:
assert len(self.params) != 0, "optimizer must have at least one param"
self.buffers: list[Tensor] = dedup([x for x in params if not x.requires_grad]) # buffers are still realized
self.device = device or self.params[0].device
self.param_dtype = to_dtype(getenv("OPTIM_DTYPE", "float32"))
self.fused = fused
# store lr in at least float32 precision
self.lr = Tensor(lr if getenv("CONST_LR") else [lr], requires_grad=False, device=self.device,
@@ -24,10 +25,9 @@ class Optimizer:
if self.fused: self.pos_params = list(itertools.accumulate(self.params, lambda x,y: x+y.numel(), initial=0))
def _new_optim_param(self) -> list[Tensor]:
param_dtype = to_dtype(getenv("OPTIM_DTYPE", "float32"))
if self.fused: return [Tensor.zeros(self.pos_params[-1], dtype=param_dtype, device=self.device, requires_grad=False)]
if isinstance(self.device, tuple): return [Tensor.zeros_like(t, dtype=param_dtype, requires_grad=False) for t in self.params]
else: return [Tensor.zeros(t.shape, dtype=param_dtype, device=self.device, requires_grad=False) for t in self.params]
if self.fused: return [Tensor.zeros(self.pos_params[-1], dtype=self.param_dtype, device=self.device, requires_grad=False)]
if isinstance(self.device, tuple): return [Tensor.zeros_like(t, dtype=self.param_dtype, requires_grad=False) for t in self.params]
else: return [Tensor.zeros(t.shape, dtype=self.param_dtype, device=self.device, requires_grad=False) for t in self.params]
def zero_grad(self):
"""
tinygrad/nn/state.py
+2 -2
View File
@@ -304,7 +304,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
# native types
if (dtype := { 0: dtypes.float32, 1: dtypes.float16, 16: dtypes.int8, 17: dtypes.int16, 18: dtypes.int32 }.get(ggml_type)) is not None:
return t[:dtype.itemsize * n].bitcast(dtype)
return t[:dtype.itemsize * n].contiguous().bitcast(dtype)
def q_to_uint8(t: Tensor, b: int) -> Tensor:
# TODO: rewrite with arange?
@@ -313,7 +313,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
# map to (number of elements, number of bytes)
if (nelements_nbytes := { 2: (32, 18), 3: (32, 20), 8: (32, 34), 12: (256, 144), 14: (256, 210), 39: (32, 17) }.get(ggml_type)) is not None:
blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1]))
blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1])).contiguous()
if ggml_type == 2: return (q_to_uint8(blocks[:,2:], 4).bitcast(dtypes.int8) - 8) * blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32)
if ggml_type == 3:
d, m = (blocks[:,s:s+2].bitcast(dtypes.float16).cast(dtypes.float32) for s in [ 0, 2 ])
tinygrad/runtime/ops_amd.py
+6 -4
View File
@@ -598,9 +598,10 @@ class AMDProgram(HCQProgram):
base=self.lib_gpu.va_addr)
weakref.finalize(self, self._fini, self.dev, self.lib_gpu, buf_spec)
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int|None, ...]=(), wait=False):
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int|None, ...]=(),
wait=False, timeout:int|None=None):
if self.dev.sqtt_enabled: cast(AMDComputeQueue, self.dev.hw_compute_queue_t()).sqtt_start(self.dev.sqtt_buffers).submit(self.dev)
res = super().__call__(*bufs, global_size=global_size, local_size=local_size, vals=vals, wait=wait)
res = super().__call__(*bufs, global_size=global_size, local_size=local_size, vals=vals, wait=wait, timeout=timeout)
if self.dev.pmc_enabled:
cast(AMDComputeQueue, self.dev.hw_compute_queue_t()).pmc_read(self.dev.pmc_buffer, self.dev.pmc_sched) \
.signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
@@ -869,7 +870,7 @@ class PCIIface(PCIIfaceBase):
devs:list[AMDDevice] = [d for pg in HCQCompiled.peer_groups.values() for d in pg if isinstance(d, AMDDevice) and d.is_am()]
for d in devs:
d.iface.dev_impl.ih.interrupt_handler()
if reset and d.iface.dev_impl.recover():
if reset and d.iface.dev_impl.recover(force=d.error_state is not None):
d.compute_queue.put_value, _ = d.iface.dev_impl.gfx.setup_ring(*d.compute_queue.params)
d.compute_queue.read_ptr[0] = d.compute_queue.write_ptr[0] = d.compute_queue.put_value
d.timeline_signal.value = d.timeline_value - 1
@@ -977,7 +978,8 @@ class AMDDevice(HCQCompiled):
super().__init__(device, AMDAllocator(self), compilers, functools.partial(AMDProgram, self), AMDSignal,
functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self),
functools.partial(AMDCopyQueue, self, max_copy_size=self.max_copy_size) if self.has_sdma_queue else None,
kernargs_size=(8 << 10) if self.is_usb() else (16 << 20), sigalloc_size=0x100 if self.is_usb() else 0x1000)
kernargs_size=(8 << 10) if self.is_usb() else (16 << 20), sigalloc_size=0x100 if self.is_usb() else 0x1000,
can_recover=self.is_am())
# Scratch setup
self.max_private_segment_size = 0
tinygrad/runtime/ops_cl.py
+1 -1
View File
@@ -54,7 +54,7 @@ class CLProgram:
except (TypeError, AttributeError): pass
def __call__(self, *bufs:tuple[cl.cl_mem, BufferSpec], global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]|None=None,
vals:tuple[int, ...]=(), wait=False) -> float|None:
vals:tuple[int, ...]=(), wait=False, **kw) -> float|None:
i = 0
for i,(b,_) in enumerate(bufs):
for real_i, dt in self.arg_dtypes[i]:
tinygrad/runtime/ops_cuda.py
+1 -1
View File
@@ -51,7 +51,7 @@ class CUDAProgram:
@suppress_finalizing
def __del__(self): check(cuda.cuModuleUnload(self.module))
def __call__(self, *args, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
def __call__(self, *args, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
check(cuda.cuCtxSetCurrent(self.dev.context))
if not hasattr(self, "vargs"):
self.c_args, self.vargs = encode_args(args, vals)
tinygrad/runtime/ops_dsp.py
+2 -2
View File
@@ -84,7 +84,7 @@ class DSPProgram:
def __init__(self, dev:DSPDevice, name:str, lib:bytes, **kwargs):
self.dev, self.lib = dev, lib
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
if len(bufs) >= 16: raise RuntimeError(f"Too many buffers to execute: {len(bufs)}")
pra, fds, attrs, _ = rpc_prep_args(ins=[var_vals_mv:=memoryview(bytearray((len(bufs)+len(vals))*4)), off_mv:=memoryview(bytearray(len(bufs)*4))],
@@ -293,7 +293,7 @@ class MockDSPRenderer(DSPRenderer):
class MockDSPProgram:
def __init__(self, name:str, lib:bytes, **kwargs): self.lib = lib
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
with tempfile.NamedTemporaryFile(suffix=".out") as dsp_lib:
dsp_lib.write(self.lib)
dsp_lib.flush()
tinygrad/runtime/ops_hip.py
+1 -1
View File
@@ -32,7 +32,7 @@ class HIPProgram:
def __del__(self):
if hasattr(self, 'module'): check(hip.hipModuleUnload(self.module))
def __call__(self, *args, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
def __call__(self, *args, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
check(hip.hipSetDevice(self.dev.device_id))
if not hasattr(self, "vargs"):
fields = [(f'f{i}', hip.hipDeviceptr_t, i*8) for i in range(len(args))] + [(f'v{i}', ctypes.c_int, len(args)*8+i*4) for i in range(len(vals))]
tinygrad/runtime/ops_metal.py
+1 -1
View File
@@ -123,7 +123,7 @@ class MetalProgram:
# cache these msg calls
self.max_total_threads: int = self.pipeline_state.maxTotalThreadsPerThreadgroup()
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
if prod(local_size) > self.max_total_threads:
exec_width = self.pipeline_state.threadExecutionWidth()
memory_length = self.pipeline_state.staticThreadgroupMemoryLength()
tinygrad/runtime/ops_null.py
+1 -1
View File
@@ -15,7 +15,7 @@ class NullRenderer(CStyleLanguage):
class NullProgram:
def __init__(self, device:str, name:str, lib:bytes, *args, **kwargs): self.device, self.name = device, name
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
with cpu_profile(self.name, self.device): return 1e-3
class NullAllocator(Allocator['NullDevice']):
tinygrad/runtime/ops_nv.py
+3 -2
View File
@@ -312,12 +312,13 @@ class NVProgram(HCQProgram):
yield typ, param, sh.content[start_off+4:start_off+sz+4] if typ == 0x4 else sz
start_off += (sz if typ == 0x4 else 0) + 4
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int|None, ...]=(), wait=False):
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int|None, ...]=(),
wait=False, timeout:int|None=None):
if prod(local_size) > 1024 or self.max_threads < prod(local_size) or self.lcmem_usage > cast(NVDevice, self.dev).slm_per_thread:
raise RuntimeError(f"Too many resources requested for launch, {prod(local_size)=}, {self.max_threads=}")
if any(cur > mx for cur,mx in zip(global_size, [2147483647, 65535, 65535])) or any(cur > mx for cur,mx in zip(local_size, [1024, 1024, 64])):
raise RuntimeError(f"Invalid global/local dims {global_size=}, {local_size=}")
res = super().__call__(*bufs, global_size=global_size, local_size=local_size, vals=vals, wait=wait)
res = super().__call__(*bufs, global_size=global_size, local_size=local_size, vals=vals, wait=wait, timeout=timeout)
if self.dev.pma_enabled:
self.dev.synchronize()
if pma_blob:=self.dev._prof_readback():
tinygrad/runtime/ops_python.py
+1 -1
View File
@@ -41,7 +41,7 @@ def generic_wmma_helper(inp, warp_size, WARP_THREADS, K, NUM_A, NUM_B, NUM_C, a_
class PythonProgram:
def __init__(self, name:str, lib:bytes, **kwargs):
self.uops: list[tuple[Ops, DType, list[int], Any]] = pickle.loads(lib)
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
st = time.perf_counter()
warp = list(itertools.product(*[range(x) for x in local_size[::-1]]))
warp_size = len(warp)
tinygrad/runtime/ops_qcom.py
+2 -1
View File
@@ -266,7 +266,8 @@ class QCOMProgram(HCQProgram):
super().__init__(QCOMArgsState, self.dev, self.name, kernargs_alloc_size=kernargs_alloc_size)
weakref.finalize(self, self._fini, self.dev, self.lib_gpu, buf_spec)
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int|None, ...]=(), wait=False):
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
vals:tuple[int|None, ...]=(), wait=False, **kw):
if self.max_threads < prod(local_size): raise RuntimeError("Too many resources requested for launch")
if any(g*l>mx for g,l,mx in zip(global_size, local_size, [65536, 65536, 65536])) and any(l>mx for l,mx in zip(local_size, [1024, 1024, 1024])):
raise RuntimeError(f"Invalid global/local dims {global_size=}, {local_size=}")
tinygrad/runtime/ops_webgpu.py
+1 -1
View File
@@ -90,7 +90,7 @@ class WebGPUProgram:
self.name, self.lib, self.prg = name, lib, shader_module
def __call__(self, *bufs:WGPUBufPtr, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
vals:tuple[int, ...]=(), wait=False) -> float|None:
vals:tuple[int, ...]=(), wait=False, **kw) -> float|None:
wait = wait and self.timestamp_supported
tmp_bufs = [*bufs]
buf_patch = False
tinygrad/runtime/support/am/amdev.py
+6 -6
View File
@@ -225,13 +225,13 @@ class AMDev(PCIDevImplBase):
self.ih.interrupt_handler()
self.reg("regSCRATCH_REG6").write(self.is_err_state) # set finalized state.
def recover(self) -> bool:
if not self.is_err_state: return False
if DEBUG >= 2: print(f"am {self.devfmt}: Start recovery")
def recover(self, force=False) -> bool:
if not force and not self.is_err_state: return False
if DEBUG >= 3: print(f"am {self.devfmt}: Start recovery")
self.ih.interrupt_handler()
self.gfx.reset_mec()
self.is_err_state = False
if DEBUG >= 2: print(f"am {self.devfmt}: Recovery complete")
if DEBUG >= 3: print(f"am {self.devfmt}: Recovery complete")
return True
def is_hive(self) -> bool: return self.gmc.xgmi_seg_sz > 0
@@ -243,14 +243,14 @@ class AMDev(PCIDevImplBase):
def reg(self, reg:str) -> AMRegister: return self.__dict__[reg]
def rreg(self, reg:int) -> int:
val = self.indirect_rreg(reg) if reg > len(self.mmio) else self.mmio[reg]
val = self.indirect_rreg(reg) if reg >= len(self.mmio) else self.mmio[reg]
if AM_DEBUG >= 4 and getattr(self, '_prev_rreg', None) != (reg, val): print(f"am {self.devfmt}: Reading register {reg:#x} with value {val:#x}")
self._prev_rreg = (reg, val)
return val
def wreg(self, reg:int, val:int):
if AM_DEBUG >= 4: print(f"am {self.devfmt}: Writing register {reg:#x} with value {val:#x}")
if reg > len(self.mmio): self.indirect_wreg(reg, val)
if reg >= len(self.mmio): self.indirect_wreg(reg, val)
else: self.mmio[reg] = val
def wreg_pair(self, reg_base:str, lo_suffix:str, hi_suffix:str, val:int, inst:int=0):
tinygrad/runtime/support/am/ip.py
+12 -12
View File
@@ -25,7 +25,7 @@ class AM_SOC(AM_IP):
return {getattr(am, k): k[off+9:] for k in dir(am) if k.startswith(f'{pref}_{self.adev.ip_ver[hwip][0]}') and (off:=k.find('__SRCID__')) != -1}
gfx_srcs, sdma_srcs = _ih_srcs('GFX', am.GC_HWIP), _ih_srcs('SDMA0', am.SDMA0_HWIP)
self.ih_scrs_names:dict[int, dict[int, str]] = {**{k: gfx_srcs for k in self.gfx_ih_clients}, **{k: sdma_srcs for k in self.sdma_ih_clients}}
self.ih_srcs_names:dict[int, dict[int, str]] = {**{k: gfx_srcs for k in self.gfx_ih_clients}, **{k: sdma_srcs for k in self.sdma_ih_clients}}
def init_hw(self):
if self.adev.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)}:
@@ -240,7 +240,8 @@ class AM_GFX(AM_IP):
def init_hw(self):
# Wait for RLC autoload to complete
while self.adev.regCP_STAT.read() != 0 and self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read_bitfields()['bootload_complete'] != 0: pass
wait_cond(lambda: self.adev.regCP_STAT.read() == 0 or self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read_bitfields()['bootload_complete'] == 0,
value=True, msg="RLC autoload timeout")
self.adev.gmc.init_hub("GC", inst_cnt=self.xccs)
if self.adev.partial_boot: return self.reset_mec()
@@ -285,18 +286,16 @@ class AM_GFX(AM_IP):
self._enable_mec()
# Set 1 partition
if self.xccs > 1 and not self.adev.partial_boot: self.adev.psp._spatial_partition_cmd(1)
if self.xccs > 1: self.adev.psp._spatial_partition_cmd(1)
def fini_hw(self): self._dequeue_hqds()
def reset_mec(self):
self._dequeue_hqds(reset=True)
self._dequeue_hqds()
# issue a soft reset to reset aql sync counter on multixcc systems.
if self.xccs > 1:
for xcc in range(self.xccs): self.adev.regGRBM_SOFT_RESET.write(soft_reset_cp=1, soft_reset_gfx=1, inst=xcc)
time.sleep(0.05)
for xcc in range(self.xccs): self.adev.regGRBM_SOFT_RESET.write(0x0, inst=xcc)
for xcc in range(self.xccs): self.adev.regGRBM_SOFT_RESET.write(soft_reset_cp=1, soft_reset_cpc=1, inst=xcc)
time.sleep(0.05)
for xcc in range(self.xccs): self.adev.regGRBM_SOFT_RESET.write(0x0, inst=xcc)
self._config_mec()
self._enable_mec()
@@ -384,13 +383,13 @@ class AM_GFX(AM_IP):
if self.adev.ip_ver[am.GC_HWIP] >= (10,0,0):
_config_helper(eng_name="MEC", cntl_reg="MEC_RS64", eng_reg="MEC_RS64", pipe_cnt=1, me=1, xcc=xcc)
def _dequeue_hqds(self, reset=False):
def _dequeue_hqds(self):
for q in range(2):
for xcc in range(self.xccs):
self._grbm_select(me=1, pipe=0, queue=q, inst=xcc)
if self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1:
self.adev.regCP_HQD_DEQUEUE_REQUEST.write(0x2, inst=xcc) # 1 - DRAIN_PIPE; 2 - RESET_WAVES
if not reset: wait_cond(lambda: self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1, value=0, msg="HQD dequeue timeout")
if not self.adev.is_err_state: wait_cond(lambda: self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1, value=0, msg="HQD dequeue timeout")
self._grbm_select()
class AM_IH(AM_IP):
@@ -437,7 +436,7 @@ class AM_IH(AM_IP):
[getattr(am, f'SOC15_{n}_FROM_IH_ENTRY')(entry) for n in ['CLIENT_ID', 'SOURCE_ID', 'RING_ID', 'VMID', 'VMID_TYPE', 'PASID', 'NODEID']]
ctx = [getattr(am, f'SOC15_CONTEXT_ID{i}_FROM_IH_ENTRY')(entry) for i in range(4)]
src_name = self.adev.soc.ih_scrs_names.get(client, {}).get(src, '')
src_name = self.adev.soc.ih_srcs_names.get(client, {}).get(src, '')
print(f"am {self.adev.devfmt}: IH ({rptr:#x}/{wptr['offset']:#x}) client={self.adev.soc.ih_clients.get(client)} src={src_name}({src}) "
f"ring={ring_id} vmid={vmid}({vmid_type}) pasid={pasid} node={node} ctx=[{ctx[0]:#x}, {ctx[1]:#x}, {ctx[2]:#x}, {ctx[3]:#x}]")
@@ -451,6 +450,7 @@ class AM_IH(AM_IP):
bf = self.adev.reg(self.adev.gmc.pf_status_reg('GC')).read_bitfields()
va = (self.adev.reg('regGCVM_L2_PROTECTION_FAULT_ADDR_HI32').read()<<32) | self.adev.reg('regGCVM_L2_PROTECTION_FAULT_ADDR_LO32').read()
print(f"am {self.adev.devfmt}: GCVM_L2_PROTECTION_FAULT_STATUS: {bf} {va<<12:#x}")
self.adev.reg('regGCVM_L2_PROTECTION_FAULT_CNTL').update(clear_protection_fault_status_addr=1)
self.adev.is_err_state = True
else: self.adev.is_err_state = True
tinygrad/runtime/support/hcq.py
+10 -7
View File
@@ -253,7 +253,7 @@ class HCQSignal(Generic[HCQDeviceType]):
Raises RuntimeError if a fault is detected.
"""
def wait(self, value:int, timeout:int=getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000)):
def wait(self, value:int, timeout:int|None=None):
"""
Waits the signal is greater than or equal to a specific value.
@@ -261,6 +261,7 @@ class HCQSignal(Generic[HCQDeviceType]):
value: The value to wait for.
timeout: Maximum time to wait in milliseconds. Defaults to 30s.
"""
timeout = timeout or getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000)
start_time = int(time.perf_counter() * 1000)
while (not_passed:=(prev_value:=self.value) < value) and (cur_time:=int(time.perf_counter() * 1000)) - start_time < timeout:
self._sleep(cur_time - start_time)
@@ -325,7 +326,7 @@ class HCQProgram(Generic[HCQDeviceType]):
return self.args_state_t(argsbuf, self, bufs, vals=vals)
def __call__(self, *bufs:HCQBuffer, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
vals:tuple[int|None, ...]=(), wait:bool=False) -> float|None:
vals:tuple[int|None, ...]=(), wait:bool=False, timeout:int|None=None) -> float|None:
"""
Enqueues the program for execution with the given arguments and dimensions.
@@ -349,7 +350,7 @@ class HCQProgram(Generic[HCQDeviceType]):
q.signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
if wait: self.dev.synchronize()
if wait: self.dev.synchronize(timeout=timeout)
return (float(sig_en.timestamp - sig_st.timestamp) / 1e6) if wait else None
class HCQCompiled(Compiled, Generic[SignalType]):
@@ -362,7 +363,8 @@ class HCQCompiled(Compiled, Generic[SignalType]):
cpu_devices: list[HCQCompiled] = []
def __init__(self, device:str, allocator:HCQAllocatorBase, compilers:CompilerSet, runtime, signal_t:Type[SignalType],
comp_queue_t:Callable[..., HWQueue], copy_queue_t:Callable[..., HWQueue]|None=None, kernargs_size=(16 << 20), sigalloc_size=0x1000):
comp_queue_t:Callable[..., HWQueue], copy_queue_t:Callable[..., HWQueue]|None=None, kernargs_size=(16 << 20), sigalloc_size=0x1000,
can_recover:bool=False):
self.device_id:int = int(device.split(":")[1]) if ":" in device else 0
from tinygrad.runtime.graph.hcq import HCQGraph
@@ -386,22 +388,23 @@ class HCQCompiled(Compiled, Generic[SignalType]):
self.kernargs_buf:HCQBuffer = self.allocator.alloc(kernargs_size, BufferSpec(cpu_access=True))
self.kernargs_offset_allocator:BumpAllocator = BumpAllocator(self.kernargs_buf.size, wrap=True)
self.can_recover = can_recover # Whether the device can recover from faults or timeouts
self.error_state:Exception|None = None # Exception if error is unrecoverable and sync will always fail
if self._is_cpu(): HCQCompiled.cpu_devices.append(self)
def synchronize(self):
def synchronize(self, timeout:int|None=None):
if self.error_state is not None: raise self.error_state
# If we have any work on CPU devices, need to synchronize them. This is just an optimization to release GIL allowing to finish faster.
if not self._is_cpu():
for dev in HCQCompiled.cpu_devices: dev.synchronize()
try: self.timeline_signal.wait(self.timeline_value - 1)
try: self.timeline_signal.wait(self.timeline_value - 1, timeout=timeout if timeout is not None and self.can_recover else None)
except RuntimeError as e:
self.error_state = e
if hasattr(self, 'on_device_hang'): self.on_device_hang()
else: raise e
raise e
if self.timeline_value > (1 << 31): self._wrap_timeline_signal()
if PROFILE:
tinygrad/schedule/indexing.py
+1 -1
View File
@@ -153,7 +153,7 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO
rngs = tuple(r if (s == 0 and e == 0) else graph_rewrite((r >= s) & (r < (sh+s)),
symbolic+pm_simplify_valid, name="pad").where(r-s, UOp.invalid()) for r,sh,(s,e) in zip(rngs, in_shape, arg))
case Ops.RESHAPE:
sink = UOp.sink(*rngs)
sink = UOp.sink(*rngs).simplify() # NOTE: this applies any commutative flips to the rngs early
sub_array = {r:UOp.range(r.src[0], i, AxisType.PLACEHOLDER) for i,r in enumerate(sink.ranges)}
rngs = _apply_reshape(in_shape, arg, sink.substitute(sub_array)).substitute({v:k for k,v in sub_array.items()}).src
case _: raise RuntimeError(f"{op} is not a MovementOp")
tinygrad/schedule/multi.py
+9 -1
View File
@@ -164,10 +164,18 @@ def passthrough_multi(root:UOp, multi:UOp):
return UOp(root.op, root.dtype, (multi.src[0],)+tuple(x.src[0] if x.op is Ops.MULTI else x for x in root.src[1:]), root.arg).multi(multi.axis)
def rewrite_into_call(call:UOp):
return call.replace(src=(graph_rewrite(call.src[0], multi_pm, name="subcall"),)+call.src[1:]) if should_resolve_call(call) else None
if not should_resolve_call(call): return None
new_body = graph_rewrite(call.src[0], multi_pm, name="subcall")
new_args = tuple(a.src[0] if a.op is Ops.MULTI else a for a in call.src[1:])
return call.replace(src=(new_body,)+new_args)
def param_to_multi(p:UOp):
if p.axis is None: return None
return UOp.param(p.arg, p.dtype, p.shard_shape, p._device).multi(p.axis)
# NOTE: this is the same pattern as Ops.UNROLL
multi_pm = PatternMatcher([
(UPat(Ops.PARAM, name="p"), param_to_multi),
(UPat(GroupOp.ALU, name="root", custom_early_reject=set([Ops.MULTI])), alu_multi),
(UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.MULTI, name="multi"), ), name="root"), reduce_multi),
(UPat(Ops.RESHAPE, src=(UPat(Ops.MULTI, name="multi"), UPat()), name="root"), reshape_multi),
tinygrad/schedule/rangeify.py
+1
View File
@@ -90,6 +90,7 @@ def resolve_call(c:UOp, allow_param_mismatch=True) -> UOp|None:
dict_map = {x:args[x.arg] for x in params}
for i, (p, a) in enumerate(dict_map.items()):
if p.axis != a.axis: raise TypeError(f"arg {i} axis mismatch: expected {p.axis}, got {a.axis}")
if p.max_shape != a.max_shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}")
if p.dtype != a.dtype: raise TypeError(f"arg {i} dtype mismatch: expected {p.dtype}, got {a.dtype}")
return c.src[0].substitute(dict_map, walk=True)
tinygrad/tensor.py
+10 -13
View File
@@ -7,7 +7,7 @@ if TYPE_CHECKING: import numpy
from tinygrad.dtype import DType, DTypeLike, dtypes, ImageDType, ConstType, least_upper_float, least_upper_dtype, sum_acc_dtype, to_dtype, truncate
from tinygrad.dtype import _from_np_dtype, _to_np_dtype, PyConst
from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten
from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ASM_GEMM, ceildiv, fetch, is_numpy_ndarray, TracingKey, cpu_profile
from tinygrad.helpers import IMAGE, FLOAT16, WINO, Metadata, TRACEMETA, ASM_GEMM, ceildiv, fetch, is_numpy_ndarray, TracingKey, cpu_profile
from tinygrad.helpers import suppress_finalizing, disable_gc
from tinygrad.gradient import compute_gradient
from tinygrad.mixin import OpMixin
@@ -1750,7 +1750,7 @@ class Tensor(OpMixin):
print(t.all(axis=1, keepdim=True).numpy())
```
"""
return self.logical_not().any(axis, keepdim).logical_not()
return self.bool().min(axis, keepdim)
def isclose(self, other:Tensor, rtol:float=1e-05, atol:float=1e-08, equal_nan=False) -> Tensor:
"""
@@ -2565,11 +2565,10 @@ class Tensor(OpMixin):
return values._inverse(), indices
@staticmethod
def _tri(r:sint, c:sint, diagonal:int=0, device=None, requires_grad:bool|None=None) -> Tensor:
assert isinstance(r, int) and isinstance(c, int), f"does not support symbolic, getting {r=}, {c=}"
def _tri(r:sint, c:sint, diagonal=0, device=None, requires_grad:bool|None=None) -> Tensor:
return (Tensor.arange(r, device=device).unsqueeze(-1) + diagonal <= Tensor.arange(c, device=device)).requires_grad_(requires_grad)
def triu(self, diagonal:int=0) -> Tensor:
def triu(self, diagonal:sint=0) -> Tensor:
"""
Returns the upper triangular part of the tensor, the other elements are set to 0.
@@ -2592,7 +2591,7 @@ class Tensor(OpMixin):
"""
return Tensor._tri(self.shape[-2], self.shape[-1], diagonal=diagonal, device=self.device).where(self, self.zeros_like())
def tril(self, diagonal:int=0) -> Tensor:
def tril(self, diagonal:sint=0) -> Tensor:
"""
Returns the lower triangular part of the tensor, the other elements are set to 0.
@@ -3266,7 +3265,7 @@ class Tensor(OpMixin):
```
"""
if not dtypes.is_int(self.dtype): raise RuntimeError(f"expect integer dtype, getting {self.dtype=}")
if num_classes == -1: num_classes = int((self.max()+1).item())
if num_classes == -1: num_classes = int(self.max().item())+1
return self[..., None]._one_hot_along_dim(num_classes).where(1, 0)
def scaled_dot_product_attention(self, key:Tensor, value:Tensor, attn_mask:Tensor|None=None, dropout_p:float=0.0,
@@ -3284,9 +3283,6 @@ class Tensor(OpMixin):
print(q.scaled_dot_product_attention(k, v).numpy())
```
"""
# NOTE: it also works when `key` and `value` have symbolic shape.
assert all_int(self.shape), f"does not support symbolic shape {self.shape}"
if getenv("FLASH_ATTENTION"):
from extra.thunder.tiny.fa import flash_attention
return flash_attention(self, key, value, attn_mask=attn_mask, is_causal=is_causal)
@@ -3462,8 +3458,9 @@ class Tensor(OpMixin):
#preprocess the matrix
Q, R = (self.qr() if m >= n else self.transpose(-2, -1).qr())
num, q_num = min(m, n), max(m, n)
U = R.shrink(tuple([None] * len(b_shape) + [(0, num), (0, num)]))
V = Tensor.eye(num, dtype=self.dtype).reshape((1,) * len(b_shape) + (num, num)).expand(b_shape + (num, num))
# TODO: codegen infinite loop without contiguous
U = R.shrink(tuple([None] * len(b_shape) + [(0, num), (0, num)])).contiguous()
V = Tensor.eye(num, dtype=self.dtype).reshape((1,) * len(b_shape) + (num, num)).expand(b_shape + (num, num)).contiguous()
#prepare round robin pairing
permute, inverse_permute = Tensor.arange(0, num, dtype=dtypes.int), Tensor.zeros(num, dtype=dtypes.int)
permute[num//2:num] = permute[num//2:num].flip(0)
@@ -3601,7 +3598,7 @@ class Tensor(OpMixin):
return cx.image_conv2d(cw, groups=groups, dtype=dtype).reshape(out_shape_t).transpose(self.ndim-1, self.ndim-2)
def image_conv2d(self, weight:Tensor, bias:Tensor|None=None, groups=1, stride=1, dilation=1, padding=0, dtype=None) -> Tensor:
base_image_type, dtsz = (dtypes.imageh, 2) if (FLOAT16:=getenv("FLOAT16", 0)) else (dtypes.imagef, 4)
base_image_type, dtsz = (dtypes.imageh, 2) if FLOAT16 else (dtypes.imagef, 4)
(bs,_,iy,ix), (cout,cin,H,W) = self.shape, weight.shape
x, w = self, weight.reshape(groups, (rcout := cout//groups), cin, H, W)
tinygrad/uop/ops.py
+36 -18
View File
@@ -163,7 +163,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
# Check self first, then iterate backward_slice (avoids creating intermediate dict)
return self.op in ops or any(x.op in ops for x in self.backward_slice)
def toposort(self, gate:Callable|None=None) -> dict[UOp, None]:
def toposort(self, gate:Callable|None=None, enter_calls=True) -> dict[UOp, None]:
cache: dict[UOp, None] = {}
stack: list[tuple[UOp, bool]] = [(self, False)] # each stack entry is (node, visited_flag)
while stack:
@@ -172,7 +172,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
if not visited:
if gate is None or gate(node):
stack.append((node, True)) # push node back on stack to process after its srcs
for s in reversed(node.src): stack.append((s, False)) # push srcs on the stack
for s in reversed(node.src if enter_calls or node.op is not Ops.CALL else node.src[1:]):
stack.append((s, False)) # push srcs on the stack
else: cache[node] = None # second time i'm seeing this node, add it to returned toposort
return cache
@@ -253,6 +254,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
case Ops.RESHAPE:
if self.src[0]._shape is None: return self.marg
# MULTI marker (axis info in PARAM sources) has no shape
case Ops.MULTI if len(self.src) == 0: return None
# movement ops change the shape
# NOTE: ssimplify is required because the shape needs to be canonical for broadcasting and same shape checking
if self.op in GroupOp.Movement.union({Ops.MULTI, Ops.REDUCE_AXIS, Ops.WMMA}):
@@ -515,6 +519,11 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
# COPY removes axis. TODO: add more tests for this, and consider MSELECT/MSTACK
if self.op is Ops.COPY: return None
if self.op is Ops.MULTI: return self.arg
# PARAM: axis is stored as a MULTI source
if self.op is Ops.PARAM:
for s in self.src:
if s.op is Ops.MULTI: return s.arg
return None
# NOTE: they all have to share an axis, we always choose [-1]
if self.op in GroupOp.ALU: return axes[-1] if (axes := dedup([x.axis for x in self.src if x.axis is not None])) else None
if len(self.src) == 0: return None
@@ -649,10 +658,22 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
while len(s.src) and s.op not in {Ops.BUFFER, Ops.PARAM, Ops.BUFFERIZE, Ops.MSTACK}: s = s.src[0]
return s
def contiguous_view_offset(self) -> tuple[int, int]|None:
"""If movement ops on a BUFFER collapse to a contiguous range, return (size, offset) in elements. Otherwise None."""
from tinygrad.schedule.rangeify import pm_mops
from tinygrad.uop.symbolic import symbolic
out = graph_rewrite(self._mop(Ops.RESHAPE, (self.size,)).index(UOp.range(self.size, 0)), pm_mops+symbolic, name="contiguous_view_offset")
if out.op is not Ops.INDEX: return None
if out.src[1].op is Ops.CONST and self.size == 1: return (1, out.src[1].arg)
if out.src[1].op is Ops.RANGE: return (self.size, 0)
if out.src[1].op is Ops.ADD and out.src[1].src[0].op is Ops.RANGE and out.src[1].src[1].op is Ops.CONST:
return (self.size, out.src[1].src[1].arg)
return None
def has_buffer_identity(self):
"""Check if this UOp has a concrete buffer identity in the graph (RESHAPE/MULTI -> BUFFER chain)."""
if self.op in {Ops.RESHAPE, Ops.MULTI}: return self.src[0].has_buffer_identity()
return self.op in {Ops.BUFFER, Ops.PARAM}
return self.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.PARAM}
@property
def buffer(self) -> Buffer|MultiBuffer:
@@ -660,23 +681,20 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
if self.op in {Ops.CONTIGUOUS, Ops.RESHAPE}: return self.src[0].buffer
# this buffer can process disk tensors and simple movement ops
if self is not self.base:
from tinygrad.schedule.rangeify import pm_mops
from tinygrad.uop.symbolic import symbolic
out = graph_rewrite(self.flatten().index(UOp.range(self.size, 0)), pm_mops+symbolic)
buf = out.src[0].buffer
size_offset = self.contiguous_view_offset()
if size_offset is None: raise RuntimeError(f"cannot collapse movement ops on {self.base.op} to a contiguous view")
size, offset = size_offset
buf = self.base.buffer
assert isinstance(buf, Buffer), "must be a Buffer for movement ops"
assert out.op is Ops.INDEX, "couldn't collapse to a single INDEX"
if out.src[1].op is Ops.CONST:
return buf.view(1, out.dtype, out.src[1].arg*out.dtype.itemsize)
if out.src[1].op is Ops.RANGE:
return buf.view(self.size, out.dtype, 0)
if out.src[1].op is Ops.ADD and out.src[1].src[0].op is Ops.RANGE and out.src[1].src[1].op is Ops.CONST:
return buf.view(self.size, out.dtype, out.src[1].src[1].arg*out.dtype.itemsize)
raise RuntimeError(f"cannot collapse INDEX {out.pyrender()} to a single size/offset")
return buf.view(size, self.dtype, offset*self.dtype.itemsize)
if self.op is Ops.BITCAST:
buf = self.src[0].buffer
assert isinstance(buf, Buffer), "must be a Buffer for BITCAST"
return buf.view(self.size, self.dtype, 0)
if self.op is Ops.BUFFER_VIEW:
buf = self.src[0].buffer
assert isinstance(buf, Buffer), "must be a Buffer for BUFFER_VIEW"
return buf.view(self.size, self.dtype, self.arg[1] * self.dtype.itemsize)
if self.op is Ops.MSELECT:
ret = self.src[0].buffer
assert isinstance(ret, MultiBuffer)
@@ -868,9 +886,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
def param_like(self, slot:int):
if self.op is Ops.BIND:
return UOp.param(slot, self.dtype, self._shape, self._device, self._min_max, self.src[0].arg[0])
if self.axis is not None:
return UOp.param(slot, self.dtype, self.shard_shape, self._device).multi(self.axis)
return UOp.param(slot, self.dtype, self._shape, self._device)
p = UOp.param(slot, self.dtype, self._shape, self._device)
if self.axis is not None: p = p.replace(src=p.src + (UOp(Ops.MULTI, arg=self.axis),))
return p
def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=(), name:str|None=None) -> UOp:
# TODO: reenable this after ENCDEC is fixed
tinygrad/uop/spec.py
+3
View File
@@ -87,6 +87,9 @@ _tensor_spec = PatternMatcher([
(UPat(Ops.BUFFER, src=(UPat((Ops.LUNIQUE, Ops.UNIQUE)), UPat(Ops.DEVICE)), name="buf"),
lambda buf: isinstance(buf.arg, int) and isinstance(buf.dtype, (DType, ImageDType))),
# BUFFER_VIEW on BUFFER is allowed if BUFFER is
(UPat(Ops.BUFFER_VIEW, src=(UPat(Ops.BUFFER),)), lambda: True),
# KERNEL can attach to an AFTER to describe the compute required to realize a BUFFER
(UPat(Ops.CALL, src=UPat((Ops.BUFFER, Ops.AFTER, Ops.MSELECT, Ops.MSTACK, Ops.BIND))), lambda: True),
tinygrad/uop/symbolic.py
-2
View File
@@ -427,8 +427,6 @@ sym = symbolic+pm_simplify_valid+PatternMatcher([
(UPat((Ops.SINK, Ops.GROUP), name="root"),
lambda root: UOp(root.op, root.dtype, tuple(flatten(x.src if x.op in REMOVE_FROM_SINK_LIKE else (x,) for x in root.src)), root.arg)
if any(x.op in REMOVE_FROM_SINK_LIKE for x in root.src) else None),
# remove END with empty NOOP
(UPat(Ops.END, src=(UPat(Ops.NOOP, src=(), name="noop"),), allow_any_len=True), lambda noop:noop),
# ** combine terms (opinionated) **
(-1 * (UPat.var("x") + UPat.var("y")), lambda x,y: (-x)+(-y)), # -(x+y) -> -x + -y
# (x+y)*c -> x*c+y*c. only for int, float has inf*0=nan issue
tinygrad/viz/js/index.js
+17 -7
View File
@@ -253,6 +253,14 @@ const Modes = {0:'read', 1:'write', 2:'write+read'};
function setFocus(key) {
if (key !== focusedShape) {
saveToHistory({ shape:focusedShape });
// adjust zoom if the entire shape is off screen
const { eventType, e } = selectShape(key);
if (e != null) {
const [x0, x1] = eventType === EventTypes.EXEC ? [e.x, e.x+e.width] : [e.x[0], e.x.at(-1)];
const xscale = d3.scaleLinear().domain([data.first, data.dur]).range([0, document.getElementById("timeline").clientWidth]);
const [st, et] = xscale.range().map(zoomLevel.invertX, zoomLevel).map(xscale.invert, xscale);
if (x1 < st || x0 > et) zoomLevel = d3.zoomIdentity.translate(-xscale((x0+x1)/2-(et-st)/2)*zoomLevel.k, 0).scale(zoomLevel.k);
}
focusedShape = key; d3.select("#timeline").call(canvasZoom.transform, zoomLevel);
}
const { eventType, e } = selectShape(key);
@@ -312,7 +320,7 @@ async function renderProfiler(path, unit, opts) {
const u64 = () => { const ret = new Number(view.getBigUint64(offset, true)); offset += 8; return ret; }
const f32 = () => { const ret = view.getFloat32(offset, true); offset += 4; return ret; }
const optional = (i) => i === 0 ? null : i-1;
const dur = u32(), tracePeak = u64(), indexLen = u32(), layoutsLen = u32();
const dur = u32(), tracePeak = u64(), indexLen = u32(), layoutsLen = u32(); data.dur = dur;
const textDecoder = new TextDecoder("utf-8");
const { strings, dtypeSize, markers } = JSON.parse(textDecoder.decode(new Uint8Array(buf, offset, indexLen))); offset += indexLen;
// place devices on the y axis and set vertical positions
@@ -1050,13 +1058,15 @@ document.addEventListener("keydown", (event) => {
if (expandSteps && getSubrewrites(step).length) return step.children[0].click();
return setState({ expandSteps:!expandSteps });
}
// left and right go through rewrites in a single UOp
if (event.key == "ArrowLeft") {
event.preventDefault()
return setState({ currentRewrite:Math.max(0, currentRewrite-1) });
}
if (event.key == "ArrowRight") {
// left and right go through rewrites in a single UOp, in profiler go forward/backward in time
if (event.key == "ArrowLeft" || event.key == "ArrowRight") {
event.preventDefault()
if (profiler.style.display !== "none" && focusedShape != null) {
const [t, idx] = focusedShape.split("-");
const i = parseInt(idx), last = data.tracks.get(t).shapes.length-1;
return setFocus(`${t}-${event.key == "ArrowLeft" ? Math.max(0, i-1) : Math.min(last, i+1)}`);
}
if (event.key == "ArrowLeft") return setState({ currentRewrite:Math.max(0, currentRewrite-1) });
const totalRewrites = ret.length-1;
return setState({ currentRewrite:Math.min(totalRewrites, currentRewrite+1) });
}