From 3fdbb82bfebfb4dc8add74508ead469b89917240 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Thu, 20 Aug 2026 16:26:02 -0700 Subject: [PATCH] add ansipad and PARALLEL contextvar (#17642) --- test/backend/test_schedule.py | 2 +- test/null/test_viz.py | 2 +- tinygrad/engine/realize.py | 4 ++-- tinygrad/helpers.py | 8 ++++++-- tinygrad/renderer/cstyle.py | 3 ++- tinygrad/renderer/isa/x86.py | 3 ++- tinygrad/renderer/llvmir.py | 3 ++- 7 files changed, 16 insertions(+), 9 deletions(-) diff --git a/test/backend/test_schedule.py b/test/backend/test_schedule.py index dd2b9e5d71..4f4bffa10a 100644 --- a/test/backend/test_schedule.py +++ b/test/backend/test_schedule.py @@ -176,7 +176,7 @@ class TestLimitBufs(unittest.TestCase): def test_limit_bufs_linear_scaling(self): def sched_time(n): - with Context(TRACK_MATCH_STATS=0, DEBUG=0): + with Context(TRACK_MATCH_STATS=0, DEBUG=0, PARALLEL=0): bufs = [Tensor.ones(16).contiguous().realize() for _ in range(4)] root = bufs[0] for i in range(n): root = root + bufs[i % 4] diff --git a/test/null/test_viz.py b/test/null/test_viz.py index 87ecf4951a..5bfb663bae 100644 --- a/test/null/test_viz.py +++ b/test/null/test_viz.py @@ -43,7 +43,7 @@ def save_viz(): Buffer.profile_events.clear() cpu_events.clear() viz = VizTrace() - with Context(VIZ=-1, TRACK_MATCH_STATS=2, PROFILE=1): + with Context(VIZ=-1, TRACK_MATCH_STATS=2, PROFILE=1, PARALLEL=0): yield viz viz.set_data() diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index 17d94cff20..ce48d256f3 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -2,7 +2,7 @@ from __future__ import annotations from typing import cast, Iterator, Any, Sequence import random, itertools, math, weakref, array, decimal from dataclasses import dataclass, replace, field -from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple +from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansipad, all_int, prod, flatten, Context, getenv, to_tuple from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry @@ -89,7 +89,7 @@ def track_stats(ctx:ExecContext, call:UOp, st:decimal.Decimal, ets:list[float|No mem_str = f"{membw*1e-9:4.0f}|{ldsbw*1e-9:<6.0f} GB/s" if membw < 1e13 and ldsbw < 1e15 else \ colored(f"{membw*1e-12:4.0f}|{ldsbw*1e-12:<6.0f} TB/s", 'green') print(f"{colored(f'*** {device[:7]:7s} {GlobalCounters.kernel_count:4d}', header_color)}"+ - f" {display_name+' '*(46-ansilen(display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:6.2f} GB"+ + f" {ansipad(display_name, 46)} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:6.2f} GB"+ ("" if et is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})")) first_run_cache.add(kcall.src[0].key) diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index d291c7eae2..36f153261f 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -44,6 +44,7 @@ def time_to_str(t:float, w=8) -> str: return next((f"{t * d:{w}.2f}{pr}" for d,p def size_to_str(s:int) -> str: return next((f"{s / d:.2f} {pr}" for d,pr in [(1<<30, "GB"),(1<<20, "MB"),(1<<10, "KB")] if s >= d), f"{s} B") def ansistrip(s:str): return re.sub('\x1b\\[(K|.*?m)', '', s) def ansilen(s:str): return len(ansistrip(s)) +def ansipad(s:str, w:int): return s+' '*max(w-ansilen(s), 0) def make_tuple(x:int|Sequence[int], cnt:int) -> tuple[int, ...]: return (x,)*cnt if isinstance(x, int) else tuple(x) def to_tuple(x:T|tuple[T, ...]) -> tuple[T, ...]: return x if isinstance(x, tuple) else (x,) def flatten(l:Iterable[Iterable[T]]): return [item for sublist in l for item in sublist] @@ -263,6 +264,9 @@ NUM_CPU_THREADS = ContextVar("NUM_CPU_THREADS", _get_cpu_count()) NULL_ALLOW_COPYOUT = ContextVar("NULL_ALLOW_COPYOUT", 0) # VIZ implies PROFILE, but you can run PROFILE without VIZ VIZ = ContextVar("VIZ", 0) +# this PARALLEL is for BEAM and compilation, it's currently disabled if you are using VIZ +# pytest-xdist workers share the CPU budget, explicit PARALLEL still overrides this default +PARALLEL = ContextVar("PARALLEL", NUM_CPU_THREADS.value // max(1, getenv("PYTEST_XDIST_WORKER_COUNT", 1)) if VIZ == 0 else 0) PROFILE = ContextVar("PROFILE", abs(VIZ.value)) SPEC = ContextVar("SPEC", 1) # TODO: disable by default due to speed @@ -598,9 +602,9 @@ class tqdm(Generic[T]): est_text = f'<{HMS(elapsed/prog-elapsed) if self.n else "?"}' if self.t else '' it_text = (SI(self.n/elapsed) if self.unit_scale else f"{self.n/elapsed:5.2f}") if self.n else "?" suf = f'{prog_text} [{HMS(elapsed)}{est_text}, {it_text}{self.unit}/s]' - sz = max(ncols-len(self.desc)-3-2-2-len(suf), 1) + sz = max(ncols-ansilen(self.desc)-3-2-2-len(suf), 1) bar = '\r' + self.desc + (f'{100*prog:3.0f}%|{("█"*int(num:=sz*prog)+" ▏▎▍▌▋▊▉"[int(8*num)%8].strip()).ljust(sz," ")}| ' if self.t else '') + suf - print(bar[:ncols+1], flush=True, end='\n'*close, file=sys.stderr) + print(bar, flush=True, end='\n'*close, file=sys.stderr) @classmethod def write(cls, s:str): print(f"\r\033[K{s}", flush=True, file=sys.stderr) diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 4bc9602037..8cd584f63e 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -258,7 +258,8 @@ class ClangRenderer(CStyleLanguage): gep_arr_threshold = 0 has_local = False has_threads = bool(getenv("THREADS", 1)) - global_max = (NUM_CPU_THREADS.value, 0, 0) + @property + def global_max(self): return (NUM_CPU_THREADS.value, 0, 0) # type: ignore[override] infinity = "__builtin_inff()" nan = '__builtin_nanf("")' diff --git a/tinygrad/renderer/isa/x86.py b/tinygrad/renderer/isa/x86.py index 0510fa9949..e4d2e7471b 100644 --- a/tinygrad/renderer/isa/x86.py +++ b/tinygrad/renderer/isa/x86.py @@ -810,7 +810,8 @@ class X86Renderer(ISARenderer): device = "CPU" has_local = False has_threads = bool(getenv("THREADS", 1)) - global_max = (NUM_CPU_THREADS.value, 0, 0) + @property + def global_max(self): return (NUM_CPU_THREADS.value, 0, 0) # type: ignore[override] extra_matcher = extra_matcher pre_isel_matcher = pre_isel_matcher isel_matcher = isel_matcher diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index a6efae940f..2b47bb8ea9 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -204,7 +204,8 @@ class LLVMRenderer(Renderer): class CPULLVMRenderer(LLVMRenderer): has_local = False has_threads = bool(getenv("THREADS", 1)) - global_max = (NUM_CPU_THREADS.value, 0, 0) + @property + def global_max(self): return (NUM_CPU_THREADS.value, 0, 0) # type: ignore[override] abi = 'win64cc' if sys.platform == 'win32' else None string_rewrite = base_rewrite def render(self, uops: list[UOp]) -> str: return "\n".join((k:=self._render_kernel(uops))[0] + (k[1], self._render_footer(uops)))