From e300451f3ad1c746ce4efad6aa027f85a668eb48 Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 30 Jul 2025 16:34:21 -0700 Subject: [PATCH 01/15] update llama3 (#11446) `LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 FUSE_ARANGE=1 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py` trained to 7 --- examples/mlperf/model_train.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index 032cc29592..d292c5720a 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -1294,9 +1294,12 @@ def train_llama3(): grad_acc = config["GRADIENT_ACC_STEPS"] = getenv("GRADIENT_ACC_STEPS", 1) GBS = config["GLOBAL_BATCH_SIZE"] = BS * grad_acc SEED = config["SEED"] = getenv("SEED", 5760) - SAMPLES = config["SAMPLES"] = getenv("SAMPLES", 1_200_000) SEQLEN = config["SEQLEN"] = getenv("SEQLEN", 8192) TRAIN_ON_VAL = config["TRAIN_ON_VAL"] = getenv("TRAIN_ON_VAL", 0) + SAMPLES = config["SAMPLES"] = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000) + + # LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 FUSE_ARANGE=1 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py + # trains to 7 opt_adamw_beta_1 = 0.9 opt_adamw_beta_2 = 0.95 @@ -1350,17 +1353,13 @@ def train_llama3(): for tokens in tqdm(iter, total=SAMPLES//BS): GlobalCounters.reset() loss, lr = train_step(model, tokens) - # BS=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=8B WARMUP_STEPS=2 DECAY_STEPS=300 PYTHONPATH=. AMD=1 MODEL=llama3 python3 examples/mlperf/model_train.py - # uses 43% ~= 83GB - # 8B bf16 = 16GB. model + grad + optim m and v = 64GB - # TODO: this OOM - # BS=1 SEQLEN=4000 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=8B WARMUP_STEPS=2 DECAY_STEPS=300 PYTHONPATH=. AMD=1 MODEL=llama3 python3 examples/mlperf/model_train.py # above as tqdm.write f-string tqdm.write(f"{loss.item():.4f} loss, {lr.item():.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used") - with open("loss.txt", "a") as f: - f.write(f"{i} {loss.item():.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n") + if (fname:=getenv("LOSS_FILE", "")): + with open(fname, "a") as f: + f.write(f"{i} {loss.item():.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n") - if i % 200 == 0 or i == 10: + if getenv("CKPT") and (i % 200 == 0 or i == 10): tqdm.write("saving checkpoint") if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir) fn = f"{ckpt_dir}/{i}.safe" From 6252f7770ee8889eec933bebb9509bf3ea03b4f6 Mon Sep 17 00:00:00 2001 From: wozeparrot Date: Wed, 30 Jul 2025 17:18:20 -0700 Subject: [PATCH 02/15] feat: fake data (#11447) --- examples/mlperf/model_train.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index d292c5720a..deaa4271f8 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -1346,8 +1346,14 @@ def train_llama3(): loss.realize(lr) return loss, lr - from examples.mlperf.dataloader import batch_load_llama3 - iter = batch_load_llama3(GBS, SAMPLES, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=bool(TRAIN_ON_VAL)) + if getenv("FAKEDATA", 0): + def fake_data(): + for _ in range(SAMPLES // GBS): + yield Tensor.randint(GBS, SEQLEN + 1, low=0, high=32000, dtype=dtypes.int32, device=Device.DEFAULT) + iter = fake_data() + else: + from examples.mlperf.dataloader import batch_load_llama3 + iter = batch_load_llama3(GBS, SAMPLES, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=bool(TRAIN_ON_VAL)) i = 0 for tokens in tqdm(iter, total=SAMPLES//BS): From cba3655de5e3287dd4739f58f09e9c23fc97c821 Mon Sep 17 00:00:00 2001 From: Eitan Turok <150733043+eitanturok@users.noreply.github.com> Date: Thu, 31 Jul 2025 05:03:41 +0300 Subject: [PATCH 03/15] Add Test for Setitem (#10559) * init * update * better * failing test * works * Delete test file * clean * lint * simplify variable name * rm contigious, rm int dtype, and add assertEqual --------- Co-authored-by: chenyu --- test/test_setitem.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/test/test_setitem.py b/test/test_setitem.py index b1fdcf3dd6..967acc29f1 100644 --- a/test/test_setitem.py +++ b/test/test_setitem.py @@ -158,6 +158,13 @@ class TestSetitem(unittest.TestCase): t[:-1] = t[1:] self.assertEqual(t.tolist(), [[2.0], [1.0], [1.0]]) + 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]) + class TestWithGrad(unittest.TestCase): def test_no_requires_grad_works(self): z = Tensor.rand(8, 8) From c3cfcb50cbb3a85751dc685c93cd24efe5b4762e Mon Sep 17 00:00:00 2001 From: kevvz <92408564+kevvz@users.noreply.github.com> Date: Wed, 30 Jul 2025 19:04:44 -0700 Subject: [PATCH 04/15] Add linalg_det and test for torch backend (#11405) * add linalg_det and test * space --------- Co-authored-by: chenyu --- extra/torch_backend/backend.py | 6 ++++++ extra/torch_backend/test.py | 5 +++++ 2 files changed, 11 insertions(+) diff --git a/extra/torch_backend/backend.py b/extra/torch_backend/backend.py index 80fc038c76..eba99816ef 100644 --- a/extra/torch_backend/backend.py +++ b/extra/torch_backend/backend.py @@ -128,6 +128,12 @@ def _linalg_eigh(self, UPLO: str = 'U'): w, v = torch.linalg.eigh(self.cpu(), UPLO=UPLO) return w.tiny(), v.tiny() +@torch.library.impl("aten::_linalg_det", "privateuseone") +# TODO: move to tinygrad +def _linalg_det(self: torch.Tensor): + result = aten._linalg_det(self.cpu()) + return result[0].tiny(), result[1].tiny(), result[2].tiny() + def upsample_backward(grad_out, output_size, input_size, *args, f=None): return f(grad_out.cpu(), output_size, input_size, *args).tiny() for i in [ diff --git a/extra/torch_backend/test.py b/extra/torch_backend/test.py index 10c077001d..087f944204 100644 --- a/extra/torch_backend/test.py +++ b/extra/torch_backend/test.py @@ -198,6 +198,11 @@ class TestTorchBackend(unittest.TestCase): recon = (v @ torch.diag(w) @ v.T).cpu().numpy() np.testing.assert_allclose(recon, a.cpu().numpy(), atol=1e-6) + def test_linalg_det(self): + a = torch.diag(torch.tensor([1,2,3,4,5], dtype = torch.float32, device=device)) + b = torch.linalg.det(a) + np.testing.assert_equal(b.cpu().numpy(), 120.0) + def test_scalar_assign(self): a = torch.tensor([1, 2, 3], device=device) a[1] = 4 From 24dd0d52edfc32ab6f887f22752145255d8524dc Mon Sep 17 00:00:00 2001 From: wozeparrot Date: Wed, 30 Jul 2025 20:18:56 -0700 Subject: [PATCH 05/15] feat: test remove to cpu (#11444) --- tinygrad/tensor.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 53a7790a62..307c1ae17b 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -282,7 +282,7 @@ class Tensor(MathTrait): # TODO: this is a hack for writing to DISK. remove with working assign if isinstance(self.device, str) and self.device.startswith("DISK"): if x.__class__ is not Tensor: x = Tensor(x, device="CPU", dtype=self.dtype) - cast(Buffer, self.contiguous().realize().uop.base.buffer).ensure_allocated().copyin(x._data()) + self._buffer().copyin(x._data()) return self if x.__class__ is not Tensor: x = Tensor(x, device=self.device, dtype=self.dtype) if self.uop is x.uop: return self # a self assign is a NOOP @@ -299,7 +299,10 @@ class Tensor(MathTrait): """ return Tensor(self.uop.detach(), device=self.device, requires_grad=False) - def _buffer(self) -> Buffer: return cast(Buffer, self.cast(self.dtype.base).contiguous().to("CPU").realize().uop.base.buffer) + def _buffer(self) -> Buffer: + x = self.cast(self.dtype.base).contiguous() + if isinstance(self.device, tuple): x = x.to("CPU") + return cast(Buffer, x.realize().uop.base.buffer).ensure_allocated() def _data(self) -> memoryview: return self._buffer().as_buffer() def data(self) -> memoryview: From 75c2c42def014b786b110b415dbcf9fb79f8298d Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Thu, 31 Jul 2025 13:57:12 +0300 Subject: [PATCH 06/15] suppress exceptions only during finalization (#11451) * suppress exceptions only during finalization * fix * fix typing * fix more warns * fix * better? * Revert "better?" This reverts commit a068aa57936f180e68f837247758c48cac8a344f. * mm? * no as e --- tinygrad/helpers.py | 7 +++++++ tinygrad/runtime/ops_amd.py | 11 +++++------ tinygrad/runtime/ops_cuda.py | 7 +++---- tinygrad/runtime/ops_gpu.py | 7 +++---- tinygrad/runtime/ops_nv.py | 9 ++++----- tinygrad/runtime/ops_webgpu.py | 7 +++---- 6 files changed, 25 insertions(+), 23 deletions(-) diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index e60840dd22..34306e22ac 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -81,6 +81,13 @@ def word_wrap(x, wrap=80): while len(ansistrip(x[:i])) < wrap and i < len(x): i += 1 return x[:i] + "\n" + word_wrap(x[i:], wrap) +def suppress_finalizing(func): + def wrapper(*args, **kwargs): + try: return func(*args, **kwargs) + except (AttributeError, TypeError, ImportError): + if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing + return wrapper + def pluralize(st:str, cnt:int): return f"{cnt} {st}"+('' if cnt == 1 else 's') class LazySeq(Generic[T]): # NOTE: Mapping requires __iter__ and __len__, Sequence requires supporting __len__ and slicing in __getitem__ diff --git a/tinygrad/runtime/ops_amd.py b/tinygrad/runtime/ops_amd.py index 3c440768ab..abb4b603e3 100644 --- a/tinygrad/runtime/ops_amd.py +++ b/tinygrad/runtime/ops_amd.py @@ -7,7 +7,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, H from tinygrad.runtime.support.hcq import MMIOInterface from tinygrad.uop.ops import sint from tinygrad.device import Compiled, DMAFdRef, BufferSpec -from tinygrad.helpers import getenv, to_mv, round_up, data64_le, all_same, flatten, DEBUG, AMD_LLVM, PROFILE, ProfileEvent +from tinygrad.helpers import getenv, to_mv, round_up, data64_le, all_same, flatten, DEBUG, AMD_LLVM, PROFILE, ProfileEvent, suppress_finalizing from tinygrad.renderer.cstyle import AMDRenderer from tinygrad.renderer.llvmir import AMDLLVMRenderer from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt @@ -473,11 +473,10 @@ class AMDAllocator(HCQAllocator['AMDDevice']): def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer: return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access) + @suppress_finalizing def _free(self, opaque, options:BufferSpec): - try: - self.dev.synchronize() - self.dev.iface.free(opaque) - except AttributeError: pass + self.dev.synchronize() + self.dev.iface.free(opaque) def _map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf) @@ -593,7 +592,7 @@ class KFDIface: def free(self, mem): if len(mem.mapped_devs) > 0: - gpus = (ctypes.c_int32 * len(mem.mapped_devs))(*[x.gpu_id for x in mem.mapped_devs]) + gpus = (ctypes.c_int32 * len(mem.mapped_devs))(*[x.iface.gpu_id for x in mem.mapped_devs]) stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=len(gpus)) assert stm.n_success == len(gpus) if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size) diff --git a/tinygrad/runtime/ops_cuda.py b/tinygrad/runtime/ops_cuda.py index 574c42e9fc..326ae84f01 100644 --- a/tinygrad/runtime/ops_cuda.py +++ b/tinygrad/runtime/ops_cuda.py @@ -1,6 +1,6 @@ from __future__ import annotations import ctypes, ctypes.util, functools -from tinygrad.helpers import DEBUG, getenv, mv_address, init_c_var, init_c_struct_t +from tinygrad.helpers import DEBUG, getenv, mv_address, init_c_var, init_c_struct_t, suppress_finalizing from tinygrad.device import Compiled, BufferSpec, LRUAllocator from tinygrad.renderer.cstyle import CUDARenderer from tinygrad.renderer.ptx import PTXRenderer @@ -45,9 +45,8 @@ class CUDAProgram: self.prg = prg if self.smem > 0: check(cuda.cuFuncSetAttribute(self.prg, cuda.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, self.smem)) - def __del__(self): - try: check(cuda.cuModuleUnload(self.module)) - except AttributeError: pass + @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): check(cuda.cuCtxSetCurrent(self.dev.context)) diff --git a/tinygrad/runtime/ops_gpu.py b/tinygrad/runtime/ops_gpu.py index 3b9b1d6965..c9ebb338c6 100644 --- a/tinygrad/runtime/ops_gpu.py +++ b/tinygrad/runtime/ops_gpu.py @@ -2,7 +2,7 @@ from __future__ import annotations from typing import cast import ctypes, functools, hashlib from tinygrad.runtime.autogen import opencl as cl -from tinygrad.helpers import init_c_var, to_char_p_p, from_mv, OSX, DEBUG, getenv, mv_address +from tinygrad.helpers import init_c_var, to_char_p_p, from_mv, OSX, DEBUG, getenv, mv_address, suppress_finalizing from tinygrad.renderer.cstyle import OpenCLRenderer, IntelRenderer from tinygrad.device import BufferSpec, LRUAllocator, Compiled, Compiler, CompileError @@ -69,9 +69,8 @@ class CLAllocator(LRUAllocator['CLDevice']): cl.cl_image_format(cl.CL_RGBA, {2: cl.CL_HALF_FLOAT, 4: cl.CL_FLOAT}[options.image.itemsize]), options.image.shape[1], options.image.shape[0], 0, None, status := ctypes.c_int32()), status), options) return (checked(cl.clCreateBuffer(self.dev.context, cl.CL_MEM_READ_WRITE, size, None, status := ctypes.c_int32()), status), options) - def _free(self, opaque:tuple[ctypes._CData, BufferSpec], options:BufferSpec): - try: check(cl.clReleaseMemObject(opaque[0])) - except AttributeError: pass + @suppress_finalizing + def _free(self, opaque:tuple[ctypes._CData, BufferSpec], options:BufferSpec): check(cl.clReleaseMemObject(opaque[0])) def _copyin(self, dest:tuple[ctypes._CData, BufferSpec], src:memoryview): if dest[1].image is not None: check(cl.clEnqueueWriteImage(self.dev.queue, dest[0], False, (ctypes.c_size_t * 3)(0,0,0), diff --git a/tinygrad/runtime/ops_nv.py b/tinygrad/runtime/ops_nv.py index ebb459bdf2..81dbbd3715 100644 --- a/tinygrad/runtime/ops_nv.py +++ b/tinygrad/runtime/ops_nv.py @@ -7,7 +7,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, H from tinygrad.runtime.support.hcq import MMIOInterface, FileIOInterface, MOCKGPU from tinygrad.uop.ops import sint from tinygrad.device import BufferSpec -from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, to_mv, hi32, lo32 +from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, to_mv, hi32, lo32, suppress_finalizing from tinygrad.renderer.ptx import PTXRenderer from tinygrad.renderer.cstyle import NVRenderer from tinygrad.runtime.support.compiler_cuda import CUDACompiler, PTXCompiler, PTX, NVPTXCompiler, NVCompiler @@ -276,11 +276,10 @@ class NVAllocator(HCQAllocator['NVDevice']): def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer: return self.dev.iface.alloc(size, cpu_access=options.cpu_access, host=options.host) + @suppress_finalizing def _free(self, opaque:HCQBuffer, options:BufferSpec): - try: - self.dev.synchronize() - self.dev.iface.free(opaque) - except AttributeError: pass + self.dev.synchronize() + self.dev.iface.free(opaque) def _map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf) diff --git a/tinygrad/runtime/ops_webgpu.py b/tinygrad/runtime/ops_webgpu.py index 1386c963cf..eef0a5cc30 100644 --- a/tinygrad/runtime/ops_webgpu.py +++ b/tinygrad/runtime/ops_webgpu.py @@ -1,7 +1,7 @@ import functools, struct from tinygrad.device import Compiled, Allocator, Compiler, BufferSpec from tinygrad.renderer.wgsl import WGSLRenderer -from tinygrad.helpers import round_up +from tinygrad.helpers import round_up, suppress_finalizing from tinygrad.runtime.autogen import webgpu from typing import List, Any, TypeAlias import ctypes @@ -188,9 +188,8 @@ class WebGpuAllocator(Allocator['WGPUDevPtr']): def _copyout(self, dest:memoryview, src:WGPUBufPtr): buffer_data = read_buffer(self.dev, src) dest[:] = buffer_data[:dest.nbytes] if webgpu.wgpuBufferGetSize(src) > dest.nbytes else buffer_data - def _free(self, opaque:WGPUBufPtr, options:BufferSpec): - try: webgpu.wgpuBufferDestroy(opaque) - except AttributeError: pass + @suppress_finalizing + def _free(self, opaque:WGPUBufPtr, options:BufferSpec): webgpu.wgpuBufferDestroy(opaque) class WebGpuDevice(Compiled): def __init__(self, device:str): From e847677e8af423f0e755d9778aa9c3e6336fa4d5 Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 31 Jul 2025 10:07:33 -0700 Subject: [PATCH 07/15] use AxisType in search instead of colors (#11452) --- tinygrad/opt/search.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tinygrad/opt/search.py b/tinygrad/opt/search.py index 22fc4713fe..be914cb703 100644 --- a/tinygrad/opt/search.py +++ b/tinygrad/opt/search.py @@ -2,7 +2,7 @@ from typing import cast, Callable import itertools, functools, random, math, time, multiprocessing, traceback, signal, atexit from collections import defaultdict from dataclasses import replace -from tinygrad.uop.ops import UOp, Ops, Variable, sym_infer +from tinygrad.uop.ops import UOp, Ops, Variable, sym_infer, AxisType from tinygrad.device import Device, Buffer, Compiler from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str from tinygrad.helpers import IGNORE_BEAM_CACHE, TC_SEARCH_OVER_SHAPE @@ -128,9 +128,9 @@ def get_kernel_actions(lin:Kernel, include_0=True, candidates:list[Opt]|None=Non try: lin2.apply_opt(a) up, lcl, tc_up = 1, 1, prod(tc.dims)//tc.threads if (tc:=lin2.tensor_core) else 1 - for s,c in zip(lin2.full_shape, lin2.colors()): - if c in {"magenta", "yellow"}: up *= s - elif c in {"cyan", "green", "white"}: lcl *= s + for s,c in zip(lin2.full_shape, lin2.axis_types): + if c in (AxisType.UPCAST, AxisType.UNROLL): up *= s + elif c in (AxisType.LOCAL, AxisType.GROUP_REDUCE): lcl *= s if up//tc_up > max_up or lcl > max_lcl: if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too many upcast/local. {up//tc_up=}, {max_up=}, {lcl=}, {max_lcl=}") continue From bad3cf573114f577ac24c46aa17ef276ad798d2b Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Fri, 1 Aug 2025 01:59:26 +0800 Subject: [PATCH 08/15] viz: add LLVM machine code analysis (#11421) * start * works everywhere * add viz api * utilization table * reg pressure ui * use llvm-mca * llvm-mca ui * work * cleanup * cycle through, defaults are enough * x86 pending * x86 nops * get mcpu/mtriple from autogen * cleanup server diff * move parser to python * normalize to pct of max * segments legend * imports * also monospace * max comes from the total per instruction * base on the value --- tinygrad/viz/index.html | 38 +++++++++++++++++++++++++++++-- tinygrad/viz/js/index.js | 48 ++++++++++++++++++++++++++++++++++++---- tinygrad/viz/serve.py | 27 +++++++++++++++++++--- 3 files changed, 104 insertions(+), 9 deletions(-) diff --git a/tinygrad/viz/index.html b/tinygrad/viz/index.html index f25e65bcdc..f4cc3527ec 100644 --- a/tinygrad/viz/index.html +++ b/tinygrad/viz/index.html @@ -239,6 +239,8 @@ padding: 0 8px; width: 100%; height: 100%; + max-height: 100vh; + overflow-x: auto; } .raw-text code { max-height: none !important; @@ -247,8 +249,6 @@ width: 100%; border-collapse: separate; border-spacing: 0; - border-radius: 8px; - overflow: hidden; background-color: #1a1b26; color: #f0f0f5; font-size: 0.95em; @@ -269,6 +269,40 @@ tr.main-row > td, tr.sub-row > td { padding: 8px 12px; } + tr.code-row > td:first-child { + font-family: monospace; + } + td.pct-row > div { + height: 12px; + width: 100%; + display: flex; + } + td.pct-row > div > div { + height: 100%; + } + thead { + position: sticky; + top: 0; + z-index: 10; + background-color: #20222e; + } + thead th { + text-align: left; + padding: 10px 12px; + font-weight: 600; + border-bottom: 1px solid #3a3d52; + font-size: 0.95em; + letter-spacing: 0.03em; + } + .legend { + display: flex; + align-items: center; + } + .legend > div { + width: 0.95em; + height: 0.95em; + margin-right: 4px; + } diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js index 1a8a1a16d3..cb6ff462e0 100644 --- a/tinygrad/viz/js/index.js +++ b/tinygrad/viz/js/index.js @@ -114,6 +114,7 @@ const devColors = {"TINY":["rgb(27 87 69)", "rgb(53 79 82)", "rgb(53 79 82)", "r const bufColors = ["#3A57B7","#5066C1","#6277CD","#7488D8","#8A9BE3","#A3B4F2"]; const lighten = (rgb, depth, step=0.08) => rgb.replace(/\d+/g, n => Math.round(parseInt(n)+(255-parseInt(n)) * Math.min(1, depth*step))); +const segmentColors = ["#ff8080", "#F4A261", "#C8F9D4", "#8D99AE", "#F4A261", "#ffffa2", "#ffffc0", "#87CEEB"]; var profileRet, focusedDevice, canvasZoom, zoomLevel = d3.zoomIdentity; async function renderProfiler() { @@ -377,11 +378,16 @@ function codeBlock(st, language, { loc, wrap }={}) { return ret; } -function appendRow(table, name, value, unit, cls) { +function appendTd(tr, value, unit=null) { + const fmt = (typeof value === "number" && !Number.isInteger(value)) ? value.toFixed(2) : value; + tr.appendChild(document.createElement("td")).innerText = unit == "us" ? formatTime(value) : fmt+(unit ?? ""); +} + +function appendRow(table, name, value, unit=null, cls="main-row") { const tr = table.appendChild(document.createElement("tr")); tr.className = cls; tr.appendChild(document.createElement("td")).innerText = name; - tr.appendChild(document.createElement("td")).innerText = unit === "us" ? formatTime(value) : value.toFixed(2)+(unit != null ? " "+unit : "%"); + appendTd(tr, value, unit); return tr; } @@ -495,10 +501,44 @@ async function main() { if (ckey.startsWith("/disasm")) { if (!(ckey in cache)) cache[ckey] = ret = await (await fetch(ckey)).json(); displayGraph("profiler"); - document.querySelector(".metadata").innerHTML = ""; const root = document.createElement("div"); root.className = "raw-text"; - root.appendChild(codeBlock(ret.src, "x86asm")); + const metadata = document.querySelector(".metadata"); + metadata.innerHTML = ""; + // detailed assembly view + if (ret.cols != null) { + const asm = root.appendChild(document.createElement("table")); + const thead = asm.appendChild(document.createElement("thead")); + const usage = {}; + for (const c of ret.cols) thead.appendChild(document.createElement("th")).innerText = c; + for (const r of ret.rows) { + const tr = asm.appendChild(document.createElement("tr")); + tr.className = "main-row code-row"; + for (const d of Object.values(r.data)) appendTd(tr, d); + const segmentsTd = tr.appendChild(document.createElement("td")); + segmentsTd.className = "pct-row"; + const usageBar = segmentsTd.appendChild(document.createElement("div")); + for (const [k, {width, value}] of Object.entries(r.segs)) { + const seg = usageBar.appendChild(document.createElement("div")); + seg.style.width = width+"%"; + seg.title = `${ret.segments[k]} ${value}`; + seg.style.background = segmentColors[parseInt(k)%segmentColors.length]; + if (!(k in usage)) usage[k] = 0; + usage[k] += value; + } + } + const summary = metadata.appendChild(document.createElement("table")); + for (const [i,s] of ret.segments.entries()) { + const tr = summary.appendChild(document.createElement("tr")); + tr.className = "main-row"; + const td = tr.appendChild(document.createElement("td")); + const div = td.appendChild(document.createElement("div")); + div.className = "legend"; + div.appendChild(document.createElement("div")).style.background = segmentColors[i%segmentColors.length]; + div.appendChild(document.createElement("p")).textContent = s; + appendTd(tr, usage[i] ?? 0); + } + } else root.appendChild(codeBlock(ret.src, "x86asm")); return document.querySelector(".profiler").replaceChildren(root); } // ** UOp view (default) diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index e0c60f1d37..023d1c2fd0 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -1,5 +1,6 @@ #!/usr/bin/env python3 import multiprocessing, pickle, difflib, os, threading, json, time, sys, webbrowser, socket, argparse, socketserver, functools, codecs, io +import subprocess, ctypes from contextlib import redirect_stdout from decimal import Decimal from http.server import BaseHTTPRequestHandler @@ -187,9 +188,29 @@ def get_runtime_stats(key) -> list[dict]: def get_disassembly(ctx:list[str]): if not isinstance(prg:=contexts[0][int(ctx[0])].ret, ProgramSpec): return - lib = Device[prg.device].compiler.compile(prg.src) - with redirect_stdout(buf:=io.StringIO()): Device[prg.device].compiler.disassemble(lib) - return json.dumps({"src":buf.getvalue()}).encode() + lib = (compiler:=Device[prg.device].compiler).compile(prg.src) + with redirect_stdout(buf:=io.StringIO()): compiler.disassemble(lib) + disasm_str = buf.getvalue() + from tinygrad.runtime.ops_llvm import llvm, LLVMCompiler + if isinstance(compiler, LLVMCompiler): + mtriple = ctypes.string_at(llvm.LLVMGetTargetMachineTriple(tm:=compiler.target_machine)).decode() + mcpu = ctypes.string_at(llvm.LLVMGetTargetMachineCPU(tm)).decode() + # NOTE: llvm-objdump may contain headers, skip if llvm-mca can't parse those lines + data = json.loads(subprocess.check_output(["llvm-mca", f"-mtriple={mtriple}", f"-mcpu={mcpu}", "-skip-unsupported-instructions=parse-failure", + "--json", "-"], input=disasm_str.encode())) + cr = data["CodeRegions"][0] + instrs:list = [{"data":[rep], "segs":{}} for rep in cr["Instructions"]] + for i,info in enumerate(cr["InstructionInfoView"]["InstructionList"]): instrs[i]["data"].append(info["Latency"]) + for d in cr["ResourcePressureView"]["ResourcePressureInfo"]: + i, r = d["InstructionIndex"], d["ResourceIndex"] + if i>len(instrs)-1: continue + instrs[i]["segs"][r] = instrs[i]["segs"].get(r, 0)+d["ResourceUsage"] + # rescale segment width to 0-100 + if instrs: + hi = max([sum(ins["segs"].values()) for ins in instrs]) + for n in instrs: n["segs"] = {k:{"width":v/hi*100, "value":v} for k,v in n["segs"].items()} + return json.dumps({"rows":instrs, "cols":["Opcode", "Latency", "HW Resources"], "segments":data["TargetInfo"]["Resources"]}).encode() + return json.dumps({"src":disasm_str}).encode() # ** HTTP server From e5b6149dfb9adc0e39e276329b4979b4e22d736b Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Thu, 31 Jul 2025 23:26:33 +0300 Subject: [PATCH 09/15] more typing in drivers (#11454) * more typing in drivers * rm --- tinygrad/runtime/support/am/ip.py | 16 ++++++------ tinygrad/runtime/support/nv/ip.py | 39 ++++++++++++++-------------- tinygrad/runtime/support/nv/nvdev.py | 10 +++---- 3 files changed, 33 insertions(+), 32 deletions(-) diff --git a/tinygrad/runtime/support/am/ip.py b/tinygrad/runtime/support/am/ip.py index 4878f8582b..15236c047d 100644 --- a/tinygrad/runtime/support/am/ip.py +++ b/tinygrad/runtime/support/am/ip.py @@ -169,12 +169,12 @@ class AM_SMU(AM_IP): self._send_msg(self.smu_mod.PPSMC_MSG_SetSoftMinByFreq, clck << 16 | (vals[level])) self._send_msg(self.smu_mod.PPSMC_MSG_SetSoftMaxByFreq, clck << 16 | (vals[level])) - def _smu_cmn_send_msg(self, msg, param=0, debug=False): + def _smu_cmn_send_msg(self, msg:int, param=0, debug=False): (self.adev.mmMP1_SMN_C2PMSG_90 if not debug else self.adev.mmMP1_SMN_C2PMSG_54).write(0) # resp reg (self.adev.mmMP1_SMN_C2PMSG_82 if not debug else self.adev.mmMP1_SMN_C2PMSG_53).write(param) (self.adev.mmMP1_SMN_C2PMSG_66 if not debug else self.adev.mmMP1_SMN_C2PMSG_75).write(msg) - def _send_msg(self, msg, param, read_back_arg=False, timeout=10000, debug=False): # 10s + def _send_msg(self, msg:int, param:int, read_back_arg=False, timeout=10000, debug=False): # default timeout is 10 seconds self._smu_cmn_send_msg(msg, param, debug=debug) wait_cond(lambda: (self.adev.mmMP1_SMN_C2PMSG_90 if not debug else self.adev.mmMP1_SMN_C2PMSG_54).read(), value=1, timeout_ms=timeout, msg=f"SMU msg {msg:#x} timeout") @@ -414,12 +414,12 @@ class AM_PSP(AM_IP): def _wait_for_bootloader(self): wait_cond(lambda: self.adev.reg(f"{self.reg_pref}_35").read() & 0x80000000, value=0x80000000, msg="BL not ready") - def _prep_msg1(self, data): + def _prep_msg1(self, data:memoryview): assert len(data) <= self.msg1_view.nbytes, f"msg1 buffer is too small {len(data):#x} > {self.msg1_view.nbytes:#x}" self.msg1_view[:len(data)+4] = bytes(data) + b'\x00' * 4 self.adev.gmc.flush_hdp() - def _bootloader_load_component(self, fw, compid): + def _bootloader_load_component(self, fw:int, compid:int): if fw not in self.adev.fw.sos_fw: return 0 self._wait_for_bootloader() @@ -458,7 +458,7 @@ class AM_PSP(AM_IP): wait_cond(lambda: self.adev.reg(f"{self.reg_pref}_64").read() & 0x8000FFFF, value=0x80000000, msg="sOS ring not created") - def _ring_submit(self, cmd): + def _ring_submit(self, cmd:am.struct_psp_gfx_cmd_resp) -> am.struct_psp_gfx_cmd_resp: msg = am.struct_psp_gfx_rb_frame(fence_value=(prev_wptr:=self.adev.reg(f"{self.reg_pref}_67").read()), cmd_buf_addr_lo=lo32(self.adev.paddr2mc(self.cmd_paddr)), cmd_buf_addr_hi=hi32(self.adev.paddr2mc(self.cmd_paddr)), fence_addr_lo=lo32(self.adev.paddr2mc(self.fence_paddr)), fence_addr_hi=hi32(self.adev.paddr2mc(self.fence_paddr))) @@ -477,7 +477,7 @@ class AM_PSP(AM_IP): return resp - def _load_ip_fw_cmd(self, fw_types, fw_bytes): + def _load_ip_fw_cmd(self, fw_types:list[int], fw_bytes:memoryview): self._prep_msg1(fw_bytes) for fw_type in fw_types: if DEBUG >= 2: print(f"am {self.adev.devfmt}: loading fw: {am.psp_gfx_fw_type__enumvalues[fw_type]}") @@ -487,7 +487,7 @@ class AM_PSP(AM_IP): cmd.cmd.cmd_load_ip_fw.fw_type = fw_type self._ring_submit(cmd) - def _tmr_load_cmd(self): + def _tmr_load_cmd(self) -> am.struct_psp_gfx_cmd_resp: cmd = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_SETUP_TMR) cmd.cmd.cmd_setup_tmr.buf_phy_addr_hi, cmd.cmd.cmd_setup_tmr.buf_phy_addr_lo = data64(self.adev.paddr2mc(self.tmr_paddr)) cmd.cmd.cmd_setup_tmr.system_phy_addr_hi, cmd.cmd.cmd_setup_tmr.system_phy_addr_lo = data64(self.tmr_paddr) @@ -495,7 +495,7 @@ class AM_PSP(AM_IP): cmd.cmd.cmd_setup_tmr.buf_size = self.tmr_size return self._ring_submit(cmd) - def _load_toc_cmd(self, toc_size): + def _load_toc_cmd(self, toc_size:int) -> am.struct_psp_gfx_cmd_resp: cmd = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_TOC) cmd.cmd.cmd_load_toc.toc_phy_addr_hi, cmd.cmd.cmd_load_toc.toc_phy_addr_lo = data64(self.msg1_addr) cmd.cmd.cmd_load_toc.toc_size = toc_size diff --git a/tinygrad/runtime/support/nv/ip.py b/tinygrad/runtime/support/nv/ip.py index bae8e0cbe9..eb4867dc38 100644 --- a/tinygrad/runtime/support/nv/ip.py +++ b/tinygrad/runtime/support/nv/ip.py @@ -1,6 +1,6 @@ from __future__ import annotations import ctypes, time, array, struct, itertools, dataclasses -from typing import cast +from typing import cast, Any from tinygrad.runtime.autogen.nv import nv from tinygrad.helpers import to_mv, lo32, hi32, DEBUG, round_up, round_down, mv_address, fetch, wait_cond from tinygrad.runtime.support.system import System @@ -8,7 +8,7 @@ from tinygrad.runtime.support.elf import elf_loader from tinygrad.runtime.autogen import nv_gpu @dataclasses.dataclass(frozen=True) -class GRBufDesc: size:int; v:int; p:int; lc:int=0 # noqa: E702 +class GRBufDesc: size:int; virt:bool; phys:bool; local:bool=False # noqa: E702 class NV_IP: def __init__(self, nvdev): self.nvdev = nvdev @@ -26,13 +26,13 @@ class NVRpcQueue: self.gsp, self.va, self.queue_va, self.seq = gsp, va, va + self.tx.entryOff, 0 self.queue_mv = to_mv(self.queue_va, self.tx.msgSize * self.tx.msgCount) - def _checksum(self, data): + def _checksum(self, data:bytes): if (pad_len:=(-len(data)) % 8): data += b'\x00' * pad_len checksum = 0 for offset in range(0, len(data), 8): checksum ^= struct.unpack_from('Q', data, offset)[0] return hi32(checksum) ^ lo32(checksum) - def send_rpc(self, func, msg, wait=False): + def send_rpc(self, func:int, msg:bytes, wait=False): header = nv.rpc_message_header_v(signature=nv.NV_VGPU_MSG_SIGNATURE_VALID, rpc_result=nv.NV_VGPU_MSG_RESULT_RPC_PENDING, rpc_result_private=nv.NV_VGPU_MSG_RESULT_RPC_PENDING, header_version=(3<<24), function=func, length=len(msg) + 0x20) @@ -49,7 +49,7 @@ class NVRpcQueue: self.seq += 1 self.gsp.nvdev.NV_PGSP_QUEUE_HEAD[0].write(0x0) - def wait_resp(self, cmd) -> memoryview: + def wait_resp(self, cmd:int) -> memoryview: while True: System.memory_barrier() if self.rx.readPtr == self.tx.writePtr: continue @@ -177,7 +177,7 @@ class NV_FLCN(NV_IP): self.nvdev.NV_PFALCON_FALCON_OS.with_base(self.falcon).write(0x0) assert self.nvdev.NV_PRISCV_RISCV_CPUCTL.with_base(self.falcon).read_bitfields()['active_stat'] == 1, "GSP Core is not active" - def execute_dma(self, base, cmd, dest, mem_off, sysmem, size): + def execute_dma(self, base:int, cmd:int, dest:int, mem_off:int, sysmem:int, size:int): wait_cond(lambda: self.nvdev.NV_PFALCON_FALCON_DMATRFCMD.with_base(base).read_bitfields()['full'], value=0, msg="DMA does not progress") self.nvdev.NV_PFALCON_FALCON_DMATRFBASE.with_base(base).write(lo32(sysmem >> 8)) @@ -194,7 +194,7 @@ class NV_FLCN(NV_IP): wait_cond(lambda: self.nvdev.NV_PFALCON_FALCON_DMATRFCMD.with_base(base).read_bitfields()['idle'], msg="DMA does not complete") - def start_cpu(self, base): + def start_cpu(self, base:int): if self.nvdev.NV_PFALCON_FALCON_CPUCTL.with_base(base).read_bitfields()['alias_en'] == 1: self.nvdev.wreg(base + self.nvdev.NV_PFALCON_FALCON_CPUCTL_ALIAS, 0x2) else: self.nvdev.NV_PFALCON_FALCON_CPUCTL.with_base(base).write(startcpu=1) @@ -232,11 +232,11 @@ class NV_FLCN(NV_IP): if mailbox is not None: return self.nvdev.NV_PFALCON_FALCON_MAILBOX0.with_base(base).read(), self.nvdev.NV_PFALCON_FALCON_MAILBOX1.with_base(base).read() - def disable_ctx_req(self, base): + def disable_ctx_req(self, base:int): self.nvdev.NV_PFALCON_FBIF_CTL.with_base(base).update(allow_phys_no_ctx=1) self.nvdev.NV_PFALCON_FALCON_DMACTL.with_base(base).write(0x0) - def reset(self, base, riscv=False): + def reset(self, base:int, riscv=False): engine_reg = self.nvdev.NV_PGSP_FALCON_ENGINE if base == self.falcon else self.nvdev.NV_PSEC_FALCON_ENGINE engine_reg.write(reset=1) time.sleep(0.1) @@ -408,10 +408,10 @@ class NV_GSP(NV_IP): assert self.nvdev.flcn.frts_offset == m.frtsOffset, f"FRTS mismatch: {self.nvdev.flcn.frts_offset} != {m.frtsOffset}" self.wpr_meta, self.wpr_meta_sysmem = self.nvdev._alloc_boot_struct(m) - def promote_ctx(self, client, subdevice, obj, ctxbufs, bufs=None, virt=None, phys=None): + def promote_ctx(self, client:int, subdevice:int, obj:int, ctxbufs:dict[int, GRBufDesc], bufs=None, virt=None, phys=None): res, prom = {}, nv_gpu.NV2080_CTRL_GPU_PROMOTE_CTX_PARAMS(entryCount=len(ctxbufs), engineType=0x1, hChanClient=client, hObject=obj) for i,(buf,desc) in enumerate(ctxbufs.items()): - use_v, use_p = (desc.v if virt is None else virt), (desc.p if phys is None else phys) + use_v, use_p = (desc.virt if virt is None else virt), (desc.phys if phys is None else phys) x = (bufs or {}).get(buf, self.nvdev.mm.valloc(desc.size, contiguous=True)) # allocate buffers prom.promoteEntry[i] = nv_gpu.NV2080_CTRL_GPU_PROMOTE_CTX_BUFFER_ENTRY(bufferId=buf, gpuVirtAddr=x.va_addr if use_v else 0, bInitialize=use_p, gpuPhysAddr=x.paddrs[0][0] if use_p else 0, size=desc.size if use_p else 0, physAttr=0x4 if use_p else 0, bNonmapped=(use_p and not use_v)) @@ -449,10 +449,11 @@ class NV_GSP(NV_IP): gr_size = _ctx_info(nv_gpu.NV0080_CTRL_FIFO_GET_ENGINE_CONTEXT_PROPERTIES_ENGINE_ID_GRAPHICS, add=0x40000) patch_size = _ctx_info(nv_gpu.NV0080_CTRL_FIFO_GET_ENGINE_CONTEXT_PROPERTIES_ENGINE_ID_GRAPHICS_PATCH) cfgs_sizes = {x: _ctx_info(x + 14, align=(2 << 20) if x == 5 else None) for x in range(3, 11)} # indices 3–10 are mapped to 17–24 - self.grctx_bufs = {0: GRBufDesc(gr_size, p=1, v=1), 1: GRBufDesc(patch_size, p=1, v=1, lc=1), 2: GRBufDesc(patch_size, p=1, v=1), - **{x: GRBufDesc(cfgs_sizes[x], p=0, v=1) for x in range(3, 7)}, 9: GRBufDesc(cfgs_sizes[9], p=1, v=1), - 10: GRBufDesc(cfgs_sizes[10], p=1, v=0), 11: GRBufDesc(cfgs_sizes[10], p=1, v=1)} # NOTE: 11 reuses cfgs_sizes[10] - self.promote_ctx(self.priv_root, subdev, ch_gpfifo, {k:v for k, v in self.grctx_bufs.items() if v.lc == 0}) + self.grctx_bufs = {0: GRBufDesc(gr_size, phys=True, virt=True), 1: GRBufDesc(patch_size, phys=True, virt=True, local=True), + 2: GRBufDesc(patch_size, phys=True, virt=True), **{x: GRBufDesc(cfgs_sizes[x], phys=False, virt=True) for x in range(3, 7)}, + 9: GRBufDesc(cfgs_sizes[9], phys=True, virt=True), 10: GRBufDesc(cfgs_sizes[10], phys=True, virt=False), + 11: GRBufDesc(cfgs_sizes[10], phys=True, virt=True)} # NOTE: 11 reuses cfgs_sizes[10] + self.promote_ctx(self.priv_root, subdev, ch_gpfifo, {k:v for k, v in self.grctx_bufs.items() if not v.local}) self.rpc_rm_alloc(hParent=ch_gpfifo, hClass=self.compute_class, params=None) self.rpc_rm_alloc(hParent=ch_gpfifo, hClass=self.dma_class, params=None) @@ -473,7 +474,7 @@ class NV_GSP(NV_IP): ### RPCs - def rpc_rm_alloc(self, hParent, hClass, params, client=None) -> int: + def rpc_rm_alloc(self, hParent:int, hClass:int, params:Any, client=None) -> int: if hClass == self.gpfifo_class: ramfc_alloc = self.nvdev.mm.valloc(0x1000, contiguous=True) params.ramfcMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=ramfc_alloc.paddrs[0][0], size=0x200, addressSpace=2, cacheAttrib=0) @@ -499,7 +500,7 @@ class NV_GSP(NV_IP): self.promote_ctx(client, self.subdevice, hParent, {k:v for k,v in self.grctx_bufs.items() if k in [0, 1, 2]}, phys_gr_ctx, phys=False) return obj if hClass != nv_gpu.NV1_ROOT else client - def rpc_rm_control(self, hObject, cmd, params, client=None): + def rpc_rm_control(self, hObject:int, cmd:int, params:Any, client=None): control_args = nv.rpc_gsp_rm_control_v(hClient=(client:=client or self.priv_root), hObject=hObject, cmd=cmd, flags=0x0, paramsSize=ctypes.sizeof(params) if params is not None else 0x0) self.cmd_q.send_rpc(nv.NV_VGPU_MSG_FUNCTION_GSP_RM_CONTROL, bytes(control_args) + (bytes(params) if params is not None else b'')) @@ -511,7 +512,7 @@ class NV_GSP(NV_IP): cast(nv_gpu.NVC36F_CTRL_CMD_GPFIFO_GET_WORK_SUBMIT_TOKEN_PARAMS, st).workSubmitToken |= (1 << 30) return st - def rpc_set_page_directory(self, device, hVASpace, pdir_paddr, client=None, pasid=0xffffffff): + def rpc_set_page_directory(self, device:int, hVASpace:int, pdir_paddr:int, client=None, pasid=0xffffffff): params = nv.struct_NV0080_CTRL_DMA_SET_PAGE_DIRECTORY_PARAMS_v1E_05(physAddress=pdir_paddr, numEntries=self.nvdev.mm.pte_cnt[0], flags=0x8, hVASpace=hVASpace, pasid=pasid, subDeviceId=1, chId=0) # flags field is all channels. alloc_args = nv.rpc_set_page_directory_v(hClient=client or self.priv_root, hDevice=device, pasid=pasid, params=params) @@ -544,7 +545,7 @@ class NV_GSP(NV_IP): header = nv.PACKED_REGISTRY_TABLE(size=hdr_size + len(entries_bytes) + len(data_bytes), numEntries=len(table)) self.cmd_q.send_rpc(nv.NV_VGPU_MSG_FUNCTION_SET_REGISTRY, bytes(header) + entries_bytes + data_bytes) - def run_cpu_seq(self, seq_buf): + def run_cpu_seq(self, seq_buf:memoryview): hdr = nv.rpc_run_cpu_sequencer_v17_00.from_address(mv_address(seq_buf)) cmd_iter = iter(seq_buf[ctypes.sizeof(nv.rpc_run_cpu_sequencer_v17_00):].cast('I')[:hdr.cmdIndex]) diff --git a/tinygrad/runtime/support/nv/nvdev.py b/tinygrad/runtime/support/nv/nvdev.py index d5cd11382a..ce25760964 100644 --- a/tinygrad/runtime/support/nv/nvdev.py +++ b/tinygrad/runtime/support/nv/nvdev.py @@ -71,7 +71,7 @@ class NVMemoryManager(MemoryManager): def on_range_mapped(self): self.dev.NV_VIRTUAL_FUNCTION_PRIV_MMU_INVALIDATE.write((1 << 0) | (1 << 1) | (1 << 6) | (1 << 31)) class NVDev(PCIDevImplBase): - def __init__(self, devfmt, mmio:MMIOInterface, vram:MMIOInterface, venid:int, subvenid:int, rev:int, bars:dict): + def __init__(self, devfmt:str, mmio:MMIOInterface, vram:MMIOInterface, venid:int, subvenid:int, rev:int, bars:dict): self.devfmt, self.mmio, self.vram, self.venid, self.subvenid, self.rev, self.bars = devfmt, mmio, vram, venid, subvenid, rev, bars self.lock_fd = System.flock_acquire(f"nv_{self.devfmt}.lock") @@ -101,10 +101,10 @@ class NVDev(PCIDevImplBase): for ip in [self.gsp, self.flcn]: ip.fini_hw() def reg(self, reg:str) -> NVReg: return self.__dict__[reg] - def wreg(self, addr, value): + def wreg(self, addr:int, value:int): self.mmio[addr // 4] = value if NV_DEBUG >= 4: print(f"wreg: {hex(addr)} = {hex(value)}") - def rreg(self, addr): return self.mmio[addr // 4] + def rreg(self, addr:int) -> int: return self.mmio[addr // 4] def _early_init(self): self.reg_names:set[str] = set() @@ -134,12 +134,12 @@ class NVDev(PCIDevImplBase): self.vram_size = self.reg("NV_PGC6_AON_SECURE_SCRATCH_GROUP_42").read() << 20 - def _alloc_boot_struct(self, struct): + def _alloc_boot_struct(self, struct:ctypes.Structure) -> tuple[ctypes.Structure, int]: va, paddrs = System.alloc_sysmem(sz:=ctypes.sizeof(type(struct)), contiguous=True) to_mv(va, sz)[:] = bytes(struct) return type(struct).from_address(va), paddrs[0] - def _download(self, file) -> str: + def _download(self, file:str) -> str: url = f"https://raw.githubusercontent.com/NVIDIA/open-gpu-kernel-modules/8ec351aeb96a93a4bb69ccc12a542bf8a8df2b6f/{file}" return fetch(url, subdir="defines").read_text() From 056dabda5a920da342f6b1fe0a65b430d7ae8fbc Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Fri, 1 Aug 2025 05:17:50 +0800 Subject: [PATCH 10/15] viz: refactor to color scheme (#11455) --- tinygrad/viz/js/index.js | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js index cb6ff462e0..a390a83662 100644 --- a/tinygrad/viz/js/index.js +++ b/tinygrad/viz/js/index.js @@ -109,12 +109,11 @@ function formatTime(ts, dur=ts) { } const formatUnit = (d, unit="") => d3.format(".3~s")(d)+unit; -const devColors = {"TINY":["rgb(27 87 69)", "rgb(53 79 82)", "rgb(53 79 82)", "rgb(70 172 194)", "rgb(29, 46, 98)"], - "DEFAULT":["rgb(29,31,42)","rgb(42,45,61)","rgb(55,59,79)","rgb(68,72,98)","rgb(18,19,26)","rgb(47,50,68)","rgb(59,63,84)","rgb(74,78,101)","rgb(24,26,35)","rgb(35,37,50)","rgb(49,53,72)","rgb(64,68,89)"],} -const bufColors = ["#3A57B7","#5066C1","#6277CD","#7488D8","#8A9BE3","#A3B4F2"]; - -const lighten = (rgb, depth, step=0.08) => rgb.replace(/\d+/g, n => Math.round(parseInt(n)+(255-parseInt(n)) * Math.min(1, depth*step))); -const segmentColors = ["#ff8080", "#F4A261", "#C8F9D4", "#8D99AE", "#F4A261", "#ffffa2", "#ffffc0", "#87CEEB"]; +const colorScheme = {TINY:["#1b5745", "#354f52", "#354f52", "#46acc2", "#1d2e62"], + DEFAULT:["#2b2e39", "#2c2f3a", "#31343f", "#323544", "#2d303a", "#2e313c", "#343746", "#353847", "#3c4050", "#404459", "#444862", "#4a4e65"], + BUFFER:["#3A57B7","#5066C1","#6277CD","#7488D8","#8A9BE3","#A3B4F2"], + CATEGORICAL:["#ff8080", "#F4A261", "#C8F9D4", "#8D99AE", "#F4A261", "#ffffa2", "#ffffc0", "#87CEEB"],} +const cycleColors = (lst, i) => lst[i%lst.length]; var profileRet, focusedDevice, canvasZoom, zoomLevel = d3.zoomIdentity; async function renderProfiler() { @@ -154,10 +153,10 @@ async function renderProfiler() { for (const e of timeline.shapes) { if (e.depth === 0) colorKey = e.cat ?? e.name; if (!colorMap.has(colorKey)) { - const colors = devColors[k] ?? devColors.DEFAULT; + const colors = colorScheme[k] ?? colorScheme.DEFAULT; colorMap.set(colorKey, colors[colorMap.size%colors.length]); } - const fillColor = lighten(colorMap.get(colorKey), e.depth); + const fillColor = d3.color(colorMap.get(colorKey)).brighter(e.depth).toString(); const label = parseColors(e.name).map(({ color, st }) => ({ color, st, width:ctx.measureText(st).width })); if (e.ref != null) ref = {ctx:e.ref, step:0}; else if (ref != null) { @@ -185,7 +184,7 @@ async function renderProfiler() { const y0 = e.y.map(yscale); const y1 = e.y.map(y => yscale(y+e.arg.nbytes)); const arg = { tooltipText:`${e.arg.dtype} len:${formatUnit(e.arg.sz)}\n${formatUnit(e.arg.nbytes, "B")}` }; - data.shapes.push({ x, y0, y1, arg, fillColor:bufColors[i%bufColors.length] }); + data.shapes.push({ x, y0, y1, arg, fillColor:cycleColors(colorScheme.BUFFER, i) }); } // lastly, adjust device rect by number of levels div.style.height = `${Math.max(levelHeight*timeline.maxDepth, baseHeight)+area+padding}px`; @@ -522,7 +521,7 @@ async function main() { const seg = usageBar.appendChild(document.createElement("div")); seg.style.width = width+"%"; seg.title = `${ret.segments[k]} ${value}`; - seg.style.background = segmentColors[parseInt(k)%segmentColors.length]; + seg.style.background = cycleColors(colorScheme.CATEGORICAL, parseInt(k)); if (!(k in usage)) usage[k] = 0; usage[k] += value; } @@ -534,7 +533,7 @@ async function main() { const td = tr.appendChild(document.createElement("td")); const div = td.appendChild(document.createElement("div")); div.className = "legend"; - div.appendChild(document.createElement("div")).style.background = segmentColors[i%segmentColors.length]; + div.appendChild(document.createElement("div")).style.background = cycleColors(colorScheme.CATEGORICAL, i); div.appendChild(document.createElement("p")).textContent = s; appendTd(tr, usage[i] ?? 0); } From fa66d9772de4abadc8e03e732408d848ce663a26 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Fri, 1 Aug 2025 06:01:58 +0800 Subject: [PATCH 11/15] viz: show const node when it's root (#11456) --- test/unit/test_viz.py | 10 ++++++++++ tinygrad/viz/serve.py | 2 +- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/test/unit/test_viz.py b/test/unit/test_viz.py index 7b9ca46b73..e6a4ddefd6 100644 --- a/test/unit/test_viz.py +++ b/test/unit/test_viz.py @@ -138,6 +138,16 @@ class TestViz(BaseTestViz): nop = UOp(Ops.NOOP, arg="infinite loop in fixed_point_rewrite") self.assertEqual(graphs[2], uop_to_json(nop)[id(nop)]) + def test_const_node_visibility(self): + a = UOp.variable("a", 0, 10) + z = UOp.const(dtypes.int, 0) + alu = a*z + exec_rewrite(alu, [sym]) + graphs = [x["graph"] for x in get_details(tracked_ctxs[0][0])] + # embed const in the parent node when possible + self.assertEqual(list(graphs[0]), [id(a), id(alu)]) + self.assertEqual(list(graphs[1]), [id(z)]) + # VIZ displays nested graph_rewrites in a tree view def leaf_rewrite(x:UOp): return x.rtag(1) if x.tag is None else None diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 023d1c2fd0..0cb536962d 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -55,7 +55,7 @@ def uop_to_json(x:UOp) -> dict[int, dict]: excluded: set[UOp] = set() for u in (toposort:=x.toposort()): # always exclude DEVICE/CONST/UNIQUE - if u.op in {Ops.DEVICE, Ops.CONST, Ops.UNIQUE}: excluded.add(u) + if u.op in {Ops.DEVICE, Ops.CONST, Ops.UNIQUE} and u is not x: excluded.add(u) # only exclude CONST VIEW source if it has no other children in the graph if u.op is Ops.CONST and len(u.src) != 0 and all(cr.op is Ops.CONST for c in u.src[0].children if (cr:=c()) is not None and cr in toposort): excluded.update(u.src) From 474ee9daa5fc25aefa2c19df16a871daa1ff3821 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Thu, 31 Jul 2025 15:06:19 -0700 Subject: [PATCH 12/15] hotfix: add contiguous_backward to llama --- extra/models/llama.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/extra/models/llama.py b/extra/models/llama.py index ecee88f64a..09907b964e 100644 --- a/extra/models/llama.py +++ b/extra/models/llama.py @@ -111,7 +111,7 @@ class TransformerBlock: def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]): h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask) - return (h + self.feed_forward(self.ffn_norm(h))).contiguous() + return (h + self.feed_forward(self.ffn_norm(h))).contiguous().contiguous_backward() # standard openai sampling def sample(logits: Tensor, temp: float, k: int, p: float, af: float, ap: float): From 3f742a5a7cd8e934d4295bb20b90d2633b3c28f2 Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 31 Jul 2025 16:06:18 -0700 Subject: [PATCH 13/15] comma space lab models benchmark (#11461) --- .github/workflows/benchmark.yml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 8285aa0270..a64c726f3a 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -617,6 +617,10 @@ jobs: run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/e8bea2c78ffa92685ece511e9b554122aaf1a79d/selfdrive/modeld/models/supercombo.onnx - name: openpilot dmonitoring compile3 0.9.7 run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/dmonitoring_model.onnx + - name: openpilot compile3 Space Lab policy + vision + run: | + PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29 + PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b - name: benchmark MobileNetV2 on DSP run: | # generate quantized weights From 719827b95d4e620e2fa5424d1d0aeca3803a8e8c Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Fri, 1 Aug 2025 07:12:30 +0800 Subject: [PATCH 14/15] viz: add flops / mem bw to device programs (#11459) * viz: add flops / mem bw to device programs * better spacing style --- tinygrad/viz/index.html | 1 + tinygrad/viz/js/index.js | 2 +- tinygrad/viz/serve.py | 9 ++++++--- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/tinygrad/viz/index.html b/tinygrad/viz/index.html index f4cc3527ec..28a05fdaa3 100644 --- a/tinygrad/viz/index.html +++ b/tinygrad/viz/index.html @@ -223,6 +223,7 @@ pointer-events: none; display: none; font-size: 10px; + white-space: pre; } #device-list > div { min-height: 32px; diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js index a390a83662..5ed3bfc7f8 100644 --- a/tinygrad/viz/js/index.js +++ b/tinygrad/viz/js/index.js @@ -164,7 +164,7 @@ async function renderProfiler() { const stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name); ref = stepIdx === -1 ? null : {ctx:ref.ctx, step:stepIdx}; } - const arg = { tooltipText:formatTime(e.dur), ...ref }; + const arg = { tooltipText:formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...ref }; // offset y by depth data.shapes.push({x:e.st-st, y:offsetY+levelHeight*e.depth, width:e.dur, height:levelHeight, arg, label, fillColor }); } diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 0cb536962d..6b3dc37177 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -126,12 +126,15 @@ def timeline_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict: depth = next((i for i,level_et in enumerate(levels) if st>=level_et), len(levels)) if depth < len(levels): levels[depth] = et else: levels.append(et) - name, cat = e.name, None - if (ref:=ref_map.get(name)) is not None: name = ctxs[ref]["name"] + name, cat, info = e.name, None, None + if (ref:=ref_map.get(name)) is not None: + name = ctxs[ref]["name"] + if isinstance(p:=contexts[0][ref].ret, ProgramSpec): + info = f"{p.estimates.ops/(t:=dur*1e3):.2f} GFLOPS {p.estimates.mem/t:4.1f}|{p.estimates.lds/t:.1f} GB/s" elif isinstance(e.name, TracingKey): name, cat = e.name.display_name, e.name.cat ref = next((v for k in e.name.keys if (v:=ref_map.get(k)) is not None), None) - shapes.append({"name":name, "ref":ref, "st":st, "dur":dur, "depth":depth, "cat":cat}) + shapes.append({"name":name, "ref":ref, "st":st, "dur":dur, "depth":depth, "cat":cat, "info":info}) return {"shapes":shapes, "maxDepth":len(levels)} def mem_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict: From 8ff03806e81970cf589c0aa3f3f9e90254caefa7 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Thu, 31 Jul 2025 16:28:04 -0700 Subject: [PATCH 15/15] add llama layers (#11460) * add llama layers * add contig bw for speed --- examples/mlperf/model_train.py | 4 +++- extra/models/llama.py | 4 +++- tinygrad/engine/realize.py | 2 +- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index deaa4271f8..e49f3071ee 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -1314,7 +1314,9 @@ def train_llama3(): # TODO: confirm weights are in bf16 # vocab_size from the mixtral tokenizer - model = Transformer(**(MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}), max_context=SEQLEN, jit=False, disable_kv_cache=True) + params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000} + if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers + model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True) optim = AdamW(get_parameters(model), lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay) diff --git a/extra/models/llama.py b/extra/models/llama.py index 09907b964e..b942115816 100644 --- a/extra/models/llama.py +++ b/extra/models/llama.py @@ -99,7 +99,9 @@ class FeedForward: self.w3 = linear(dim, hidden_dim, bias=False) # the gate in Gated Linear Unit def __call__(self, x:Tensor) -> Tensor: - return self.w2(self.w1(x).silu() * self.w3(x)) # SwiGLU [arxiv/2002.05202, eq (5)] + w1 = self.w1(x).silu() + w3 = self.w3(x.contiguous_backward()) # this fixes a strange fusion that makes tensor cores miss + return self.w2(w1 * w3) class TransformerBlock: def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int, norm_eps:float, max_context:int, linear=nn.Linear, diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index fd343d0e70..3c233092fa 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -156,7 +156,7 @@ class ExecItem: lds_est = sym_infer(self.prg.estimates.lds, var_vals) mem_est = min(mem_est, lds_est) # there can't be more memory accessed than loads/stores. remove this when symbolic is fixed ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else "" - print(f"{colored(f'*** {self.prg.device[:7]:7s} {GlobalCounters.kernel_count:4d}', 'magenta' if jit else ('green' if self.prg.first_run else None))} {self.prg.display_name+' '*(41-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:5.2f} GB " + # noqa: E501 + print(f"{colored(f'*** {self.prg.device[:7]:7s} {GlobalCounters.kernel_count:4d}', 'magenta' if jit else ('green' if self.prg.first_run else None))} {self.prg.display_name+' '*(44-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:5.2f} GB " + # noqa: E501 (str() if et is None else f"tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({op_est/((et or 1e-20)*1e9):9.2f} GFLOPS {mem_est/((et or 1e-20)*1e9):6.1f}|{lds_est/((et or 1e-20)*1e9):<7.1f} GB/s)" + # noqa: E501 f" {[repr(m) if TRACEMETA >= 2 else str(m) for m in self.metadata] if self.metadata else ''}")) self.prg.first_run = False