diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 6b1d495f97..2a9ff72172 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -648,6 +648,6 @@ jobs: - name: Checkout Code uses: actions/checkout@v6 - name: Speed Test - run: DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py + run: DEV=CPU:LLVM python3 test/speed/external_test_speed_v_torch.py - name: Speed Test (BEAM=2) - run: BEAM=2 DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py + run: BEAM=2 DEV=CPU:LLVM python3 test/speed/external_test_speed_v_torch.py diff --git a/examples/mlperf/dataloader.py b/examples/mlperf/dataloader.py index 275bcd4aef..3f5db0c74a 100644 --- a/examples/mlperf/dataloader.py +++ b/examples/mlperf/dataloader.py @@ -5,7 +5,7 @@ from multiprocessing import Queue, Process, shared_memory, connection, Lock import numpy as np from tinygrad import dtypes, Tensor -from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX, NUM_CPU_THREADS +from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX, CPU_COUNT from tinygrad.nn.state import TensorIO ### ResNet @@ -131,7 +131,7 @@ def batch_load_resnet(batch_size=64, val=False, shuffle=True, seed=None, pad_fir else: X = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name}") Y = [None] * (batch_size*BATCH_COUNT) - for _ in range(NUM_CPU_THREADS.value): + for _ in range(CPU_COUNT): p = Process(target=loader_process, args=(q_in, q_out, X, seed)) p.daemon = True p.start() @@ -212,7 +212,7 @@ def batch_load_train_bert(BS:int, seed:int|None=None): rng.shuffle(fs) train_files.append(fs.pop(0)) - cycle_length = min(NUM_CPU_THREADS.value, len(train_files)) + cycle_length = min(CPU_COUNT, len(train_files)) assert cycle_length > 0, "cycle_length must be greater than 0" dataset = InterleavedDataset(train_files, cycle_length) @@ -301,7 +301,7 @@ def batch_load_unet3d(preprocessed_dataset_dir:Path, batch_size:int=6, val:bool= X = Tensor.empty(*sz, dtype=dtypes.float32, device=f"disk:/dev/shm/{shm_name_x}") Y = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name_y}") - for _ in range(NUM_CPU_THREADS.value): + for _ in range(CPU_COUNT): proc = Process(target=load_unet3d_data, args=(preprocessed_dataset_dir, seed, queue_in, queue_out, X, Y)) proc.daemon = True proc.start() @@ -437,7 +437,7 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh dataset_iter = iter(image_ids) try: - for _ in range(NUM_CPU_THREADS.value): + for _ in range(CPU_COUNT): proc = Process( target=load_retinanet_data, args=(base_dir, val, queue_in, queue_out, imgs, boxes, labels), diff --git a/extra/export_model.py b/extra/export_model.py index f5d5a09a84..a6baf20bc7 100644 --- a/extra/export_model.py +++ b/extra/export_model.py @@ -241,8 +241,7 @@ export default {model_name}; def export_model(model, target:str, *inputs, model_name: Optional[str] = "model", stream_weights=False): assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"only {', '.join(EXPORT_SUPPORTED_DEVICE)} are supported" - # NOTE: NUM_CPU_THREADS=1, since export does not support threading - with Context(JIT=2, NUM_CPU_THREADS=1): linear, output_bufs = jit_model(model, *inputs) + with Context(JIT=2): linear, output_bufs = jit_model(model, *inputs) functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs) state = get_state_dict(model) weight_names = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None} diff --git a/extra/onnx_helpers.py b/extra/onnx_helpers.py index c11a8d41f8..a0a0363d4b 100644 --- a/extra/onnx_helpers.py +++ b/extra/onnx_helpers.py @@ -1,12 +1,12 @@ from tinygrad import Tensor -from tinygrad.helpers import NUM_CPU_THREADS +import os from tinygrad.tensor import _to_np_dtype from tinygrad.nn.onnx import OnnxRunner, OnnxValue import numpy as np import onnxruntime as ort ort_options = ort.SessionOptions() ort_options.log_severity_level = 3 -ort_options.intra_op_num_threads = NUM_CPU_THREADS.value +ort_options.intra_op_num_threads = os.cpu_count() or 1 def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}): """ diff --git a/spec/tinyspec.pdf b/spec/tinyspec.pdf index 5e8ebe58b3..7dec94f0b0 100644 Binary files a/spec/tinyspec.pdf and b/spec/tinyspec.pdf differ diff --git a/spec/tinyspec.tex b/spec/tinyspec.tex index c541657a65..1a31a03eb8 100644 --- a/spec/tinyspec.tex +++ b/spec/tinyspec.tex @@ -23,7 +23,6 @@ \definecolor{axblue}{HTML}{1565C0} % GLOBAL \definecolor{axcyan}{HTML}{00838F} % LOCAL \definecolor{axbrcyan}{HTML}{00ACC1} % WARP -\definecolor{axbrblue}{HTML}{42A5F5} % THREAD \definecolor{axwhite}{HTML}{616161} % LOOP (gray on white paper) \definecolor{axred}{HTML}{C62828} % REDUCE \definecolor{axbrred}{HTML}{E53935} % GROUP_REDUCE @@ -307,7 +306,6 @@ Each kernel's iteration space is a set of \op{Range} axes. Every range has an \t {\color{axblue}\texttt{GLOBAL}} & \texttt{g} & --- & --- & GPU global workgroup dimension. \\ {\color{axcyan}\texttt{LOCAL}} & \texttt{l} & g, L & inner & Workgroup local dimension (shared memory). \\ {\color{axbrcyan}\texttt{WARP}} & \texttt{w} & \multicolumn{2}{l}{(created by \op{TC})} & Warp-level lanes for tensor cores. \\ -{\color{axbrblue}\texttt{THREAD}} & \texttt{t} & g & outer & CPU thread parallelism. \\ {\color{axwhite}\texttt{LOOP}} & \texttt{L} & --- & --- & Generic sequential loop (initial state). \\ {\color{axred}\texttt{REDUCE}} & \texttt{R} & --- & --- & Reduction axis. \\ {\color{axbrred}\texttt{GROUP\_REDUCE}} & \texttt{G} & R & inner/outer & Shared-memory group reduction. \\ diff --git a/test/backend/test_jit.py b/test/backend/test_jit.py index a9c34ad678..920cfa2f3d 100644 --- a/test/backend/test_jit.py +++ b/test/backend/test_jit.py @@ -4,7 +4,7 @@ import numpy as np from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu, KernelCountException from test.unit.test_jit import _simple_test -from tinygrad import Tensor, Variable, TinyJit, Device, dtypes +from tinygrad import Tensor, TinyJit, Device, dtypes from tinygrad.engine.jit import graph_class from tinygrad.helpers import JIT, DEV, GlobalCounters, HCQ2 from tinygrad.uop.ops import Ops @@ -16,19 +16,6 @@ class TestJit(unittest.TestCase): def add(a, b): return (a+b).realize() _simple_test(add) - @unittest.skipUnless(Device.DEFAULT == "CPU", "core_id is a CPU runtimevar") - def test_hcq_core_id_runtimevar_merge(self): - N = 262144 - @TinyJit - def f(x, st): - y = (x + 1).contiguous().realize() - z = x.shrink(((st, st + N),)).contiguous().realize() - return y, z - x = Tensor.arange(2*N).clone().realize() - for _ in range(3): y, z = f(x, Variable("a", 0, N).bind(0)) - self.assertEqual(y.shape, (2*N,)) - self.assertEqual(z.shape, (N,)) - def test_jit_input_view(self): @TinyJit def f(x): return (x[2:5].contiguous() + 1).realize() diff --git a/test/device/test_hcq2.py b/test/device/test_hcq2.py index 364d02f0ec..002be54572 100644 --- a/test/device/test_hcq2.py +++ b/test/device/test_hcq2.py @@ -12,6 +12,16 @@ class TestHCQ2(unittest.TestCase): with patch.object(Device[Device.DEFAULT], "has_copy_queue", False): np.testing.assert_equal(Tensor(np.arange(61, dtype=np.float32)).to(Device.DEFAULT).contiguous().realize().numpy(), np.arange(61)) + @unittest.skipIf(Device.DEFAULT == "CPU", "ping-pong needs a non-CPU hcq2 device") + def test_cpu_device_ping_pong(self): + # CPU submits run inline, so alternating dependencies must be submitted in schedule order to avoid blocking the host submitter. + x = Tensor.ones(16, device="CPU").contiguous().realize() + a = (x + 1).contiguous() + b = (a.to(Device.DEFAULT).contiguous() + 1).contiguous() + c = (b.to("CPU").contiguous() + 1).contiguous() + out = (c.to(Device.DEFAULT).contiguous() + 1).contiguous().realize() + np.testing.assert_equal(out.numpy(), np.full(16, 5)) + @unittest.skipIf(Device.DEFAULT == "CPU", "staged copies need a non-CPU hcq2 device") def test_staged_copy_slot_reuse(self): # chunks of a staged copy rotate through the staging buffer slots, many rotations must stay bit-exact in both directions diff --git a/test/opt/test_kernel_opts.py b/test/opt/test_kernel_opts.py index 0c347bef21..44150cb8f6 100644 --- a/test/opt/test_kernel_opts.py +++ b/test/opt/test_kernel_opts.py @@ -333,20 +333,6 @@ class TestKernelOpts(unittest.TestCase): #[Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.UNROLL, axis=1, arg=4)], # noqa: E501 ]) - @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_threads, "test requires threads") - @unittest.skipUnless(Device[Device.DEFAULT].renderer.global_max is not None and - Device[Device.DEFAULT].renderer.global_max[0] > 1, "test requires multicore") - def test_thread_opts(self): - a = Tensor.rand(4, 4, 4, 4) - b = Tensor.rand(4, 4, 4) - r = (b.sqrt() + ((a+1).sum(axis=3).exp())) - helper_linearizer_opt(r, [ - [Opt(OptOps.THREAD, 0, 2)], - [Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.THREAD, 0, 2)], - [Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.THREAD, 0, 2), Opt(OptOps.UNROLL, 0, 2)], - ] + [[Opt(OptOps.THREAD, 0, 4)] if Device[Device.DEFAULT].renderer.global_max[0] >= 4 else []] - + [[Opt(OptOps.THREAD, 0, 8)] if Device[Device.DEFAULT].renderer.global_max[0] >= 8 else []]) - def test_double_sum_group(self): a = Tensor.rand(4, 4, 4) r = a.sum((1, 2)).sum() diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 4f01a8eca5..d555254dca 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -1,7 +1,7 @@ from dataclasses import replace, dataclass import itertools, functools from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC -from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, NUM_CPU_THREADS, TC_SELECT, TC_OPT, TracingKey, Context, panic +from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, TC_SELECT, TC_OPT, TracingKey, Context, panic from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, Ops, UPat, rewrite_group, KernelInfo, ProgramInfo, GroupOp, AxisType from tinygrad.uop.weak import pm_lower_weak, pm_commit_weak, pm_cast_const from tinygrad.uop.render import pyrender @@ -497,7 +497,7 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp: # config affects generated programs and cache keys; context also carries compile-only behavior to workers to_program_config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32, - DEFAULT_FLOAT, DEFAULT_INT, NUM_CPU_THREADS, TC_SELECT, TC_OPT) + DEFAULT_FLOAT, DEFAULT_INT, TC_SELECT, TC_OPT) to_program_context = (*to_program_config, SPEC, DEBUG) def to_program_key(ast:UOp, renderer:Renderer) -> tuple: return (ast.key, type(renderer), renderer.target, *[x.value for x in to_program_config]) diff --git a/tinygrad/codegen/gpudims.py b/tinygrad/codegen/gpudims.py index 10f45d89c8..849061f4b0 100644 --- a/tinygrad/codegen/gpudims.py +++ b/tinygrad/codegen/gpudims.py @@ -1,6 +1,6 @@ import math from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType -from tinygrad.dtype import dtypes, AddrSpace +from tinygrad.dtype import AddrSpace from tinygrad.renderer import Renderer def _dim_max(d:sint) -> int: return d if isinstance(d, int) else int(d.vmax) @@ -47,7 +47,7 @@ def add_gpudims(ctx:Renderer, s:UOp): all_ranges = {x.arg[0:-1]:x for x in s_topo if x.op is Ops.RANGE} # extract global/local dims - global_dims = sorted([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] in (AxisType.GLOBAL, AxisType.THREAD)]) + global_dims = sorted([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] is AxisType.GLOBAL]) local_dims = sorted([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] in (AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE)]) if not global_dims and not local_dims: return None @@ -57,8 +57,7 @@ def add_gpudims(ctx:Renderer, s:UOp): # get the idxs ki: KernelInfo = s.arg - if ctx.has_threads: idxs = [UOp.variable("core_id", 0, int(global_shape[0])-1, dtypes.int, param=True).cast(dtypes.weakint)] - elif ki.dont_use_locals: + if ki.dont_use_locals: assert not local_dims, "can't use locals if there's no local dims" idxs = get_grouped_dims("idx", global_shape, ctx.global_max, reverse=True) else: diff --git a/tinygrad/codegen/opt/__init__.py b/tinygrad/codegen/opt/__init__.py index fb4b84ddee..0618cb4d27 100644 --- a/tinygrad/codegen/opt/__init__.py +++ b/tinygrad/codegen/opt/__init__.py @@ -4,7 +4,7 @@ from enum import Enum, auto from dataclasses import dataclass class OptOps(Enum): - TC = auto(); UPCAST = auto(); UNROLL = auto(); LOCAL = auto(); THREAD = auto() # noqa: E702 + TC = auto(); UPCAST = auto(); UNROLL = auto(); LOCAL = auto() # noqa: E702 GROUP = auto(); GROUPTOP = auto(); NOLOCALS = auto(); PADTO = auto(); SWAP = auto() # noqa: E702 def __lt__(self, x:OptOps): return self.value < x.value diff --git a/tinygrad/codegen/opt/heuristic.py b/tinygrad/codegen/opt/heuristic.py index 9ff87f9c09..c0e70443cb 100644 --- a/tinygrad/codegen/opt/heuristic.py +++ b/tinygrad/codegen/opt/heuristic.py @@ -191,17 +191,4 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: k.apply_opt(Opt(OptOps.LOCAL, axis, local_sz)) if will_delete_shape: deleted_shape += 1 - # **** threading **** - - if k.ren.has_threads and k.ren.global_max is not None: - for threads in [32,16,12,8,6,5,4,3,2]: - # Skip if too many threads. Heuristic: use about 128K ops per thread - if threads > k.ren.global_max[0] or resolve(prod(k.full_shape) // (128 << 10) < threads): continue - for axis in k.axes_of(AxisType.WEAK): - if k.full_shape[axis] % threads == 0: - try: k.apply_opt(Opt(OptOps.THREAD, axis, threads)) - except KernelOptError: pass - break - if k.applied_opts and k.applied_opts[-1].op is OptOps.THREAD: break - return k diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index cda176f7e7..f15cad181c 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -133,7 +133,7 @@ class Scheduler: opt_to_at = { OptOps.LOCAL: AxisType.LOCAL, OptOps.UPCAST: AxisType.UPCAST, OptOps.UNROLL: AxisType.UNROLL, OptOps.GROUP: AxisType.GROUP_REDUCE, - OptOps.GROUPTOP: AxisType.GROUP_REDUCE, OptOps.THREAD: AxisType.THREAD} + OptOps.GROUPTOP: AxisType.GROUP_REDUCE} ret = None if opt.op in opt_to_at: @@ -160,16 +160,11 @@ class Scheduler: if opt.op is OptOps.LOCAL: check(not self.dont_use_locals, "can't use locals") check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.WEAK}, "local is for globals") - if opt.op is OptOps.THREAD: - check(self.ren is not None and self.ren.has_threads, "target does not support threads") - check(self.ren is not None and self.ren.global_max is not None and amt <= self.ren.global_max[0], "too many threads") - check(all(x is not AxisType.THREAD for x in self.axis_types), "already threaded") - check(rng in self._globalizable_rngs(), "can't apply range to this dim") if opt.op in {OptOps.GROUP, OptOps.GROUPTOP}: check(all(x.op is not OptOps.TC for x in self.applied_opts), "no grouping with tensor cores") # TODO: why is this wrong? check(not self.dont_use_locals, "can't use locals") check(rng.arg[-1] == AxisType.REDUCE, "group is for reduce") - ret = self.shift_to(rng, amt, opt_to_at[opt.op], top=opt.op in {OptOps.GROUPTOP, OptOps.THREAD}) + ret = self.shift_to(rng, amt, opt_to_at[opt.op], top=opt.op is OptOps.GROUPTOP) elif opt.op is OptOps.TC: check(len(self.applied_opts) == 0, "tensor core opts must be first") # TODO: remove the need for this by having warps check(opt.axis is not None, "tensor core opts must have an axis") @@ -183,7 +178,6 @@ class Scheduler: elif opt.op is OptOps.PADTO: check(rng.src[0].op is Ops.CONST, "only pad const axes") check(rng.arg[-1] not in {AxisType.UPCAST, AxisType.UNROLL}, "cannot pad upcasted") # TODO: why is this wrong? - check(rng.arg[-1] is not AxisType.THREAD, "cannot pad thread") new_sz = round_up(int(rng.vmax+1), cast(int, opt.arg)) check(rng.vmax+1 > new_sz//4, "pad adds more than quadruple the work") replaced_rng = UOp.range(new_sz, *rng.arg, dtype=rng.dtype) diff --git a/tinygrad/codegen/opt/search.py b/tinygrad/codegen/opt/search.py index cd7bc59284..9d3d7d3753 100644 --- a/tinygrad/codegen/opt/search.py +++ b/tinygrad/codegen/opt/search.py @@ -22,7 +22,6 @@ actions += [Opt(op=OptOps.TC, axis=0, arg=(-1, 0, getenv("TC", 1)))] # covers resnet kernels (3 global * 3 reduce) actions += [Opt(op=OptOps.TC, axis=axis, arg=(-1, getenv("TC_OPT", 2), getenv("TC", 1))) for axis in range(9)] actions += [Opt(op=OptOps.SWAP, axis=axis_0, arg=axis_1) for axis_0 in range(5) for axis_1 in range(axis_0+1, 5)] -actions += [Opt(op=OptOps.THREAD, axis=axis, arg=amt) for amt in [2,3,4,5,8,12,16,24,32,64] for axis in range(3)] if getenv("NOLOCALS"): actions += [Opt(op=OptOps.NOLOCALS)] def get_test_global_size(global_size, max_global_size, var_vals): diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index bfb7f29d35..cbb32b62cc 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -106,13 +106,13 @@ class GraphRunner: def is_sym_dim(dim) -> bool: return not all(isinstance(d, (int, float)) for d in dim) crs = [(j, self.calls[j][1].arg, self.calls[j][3]) for j in range(len(self.calls)) if self.calls[j][1].op is Ops.PROGRAM] - self.vars = sorted({v.expr for _,p,dv in crs for v in p.vars if v.expr not in dv | p.runtimevars}) + self.vars = sorted({v.expr for _,p,dv in crs for v in p.vars if v.expr not in dv}) self.symbolic_dims = dedup(tuple(d) for _,p,_ in crs for d in (p.local_size, p.global_size) if d and is_sym_dim(d)) def find_symbolic_dim(dim): return self.symbolic_dims.index(tuple(dim)) if dim is not None and tuple(dim) in self.symbolic_dims else None for j,p,dv in crs: - if (replace:=[(i, self.vars.index(v.expr)) for i, v in enumerate(p.vars) if v.expr not in dv | p.runtimevars]): + if (replace:=[(i, self.vars.index(v.expr)) for i, v in enumerate(p.vars) if v.expr not in dv]): self.var_vals_replace[j] = replace global_dim_idx, local_dim_idx = find_symbolic_dim(p.global_size), find_symbolic_dim(p.local_size) if global_dim_idx is not None or local_dim_idx is not None: diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index b3e7fdded8..a318d5cf8c 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -260,13 +260,13 @@ def _get_cpu_count() -> int: if quota != "max": count = min(count, max(1, int(quota) // int(period))) except (FileNotFoundError, ValueError, ZeroDivisionError): pass return count -NUM_CPU_THREADS = ContextVar("NUM_CPU_THREADS", _get_cpu_count()) +CPU_COUNT = _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) +PARALLEL = ContextVar("PARALLEL", CPU_COUNT // 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 diff --git a/tinygrad/renderer/__init__.py b/tinygrad/renderer/__init__.py index 40682e7ccf..15fde7b193 100644 --- a/tinygrad/renderer/__init__.py +++ b/tinygrad/renderer/__init__.py @@ -50,7 +50,6 @@ class Estimates: mults = mults.substitute({x:x.const_like(0) for x in mults.toposort() if x.op is Ops.SPECIAL}) if isinstance(mults, UOp) else mults elif u.op is Ops.END: mults = mult_stack.pop(-1) elif u.op is Ops.SPECIAL: mults *= cast(sint, u.src[0].ssimplify()) # NOTE: we don't push to the mult_stack here, you can't end these - elif u.op is Ops.PARAM and u.arg.addrspace == AddrSpace.ALU and u.expr == 'core_id': mults *= int(u.vmax) + 1 elif u.op is Ops.LOAD and u.src[0].addrspace != AddrSpace.REG: lds += u.max_numel() * u.dtype.itemsize * mults elif u.op is Ops.STORE and u.src[0].addrspace != AddrSpace.REG: @@ -67,7 +66,6 @@ class Renderer: # TODO: make this generic with a list of supported types supports_float4: bool = True has_local: bool = True - has_threads: bool = False has_shared: bool = True # NOTE: these two should be in (x,y,z) order to match the max_sizes argument in get_grouped_dims global_max: tuple[int, ...]|None = (0x8FFFFFFF,) * (3) # TODO: Ops.SPECIAL int32 indexes right now diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 680bd4608d..c62a69d14a 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -4,7 +4,7 @@ from collections import defaultdict, Counter from tinygrad.codegen.opt import tc from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str, axis_letters from tinygrad.uop.weak import commit_weak_consts -from tinygrad.helpers import strip_parens, getenv, prod, dedup, Target, NUM_CPU_THREADS, IMAGE, FLOAT16, is_image_shape +from tinygrad.helpers import strip_parens, getenv, prod, dedup, Target, IMAGE, FLOAT16, is_image_shape from tinygrad.dtype import dtypes, DType, AddrSpace, truncate, float_to_bf16 from tinygrad.renderer import Renderer @@ -263,9 +263,7 @@ class ClangRenderer(CStyleLanguage): float4_style = ('{', '}') gep_arr_threshold = 0 has_local = False - has_threads = bool(getenv("THREADS", 1)) - @property - def global_max(self): return (NUM_CPU_THREADS.value, 0, 0) # type: ignore[override] + global_max = (1, 0, 0) infinity = "__builtin_inff()" nan = '__builtin_nanf("")' diff --git a/tinygrad/renderer/isa/x86.py b/tinygrad/renderer/isa/x86.py index 042486d3d7..16fcff26cb 100644 --- a/tinygrad/renderer/isa/x86.py +++ b/tinygrad/renderer/isa/x86.py @@ -7,7 +7,7 @@ from tinygrad.dtype import dtypes, DType, truncate, AddrSpace from tinygrad.uop import FastEnum, auto, Ops, GroupOp from tinygrad.uop.ops import UOp, UPat, PatternMatcher, promo_dtype from tinygrad.renderer.isa import ISARenderer, IselContext, Register, PreRegAllocContext, greg -from tinygrad.helpers import getenv, NUM_CPU_THREADS, unwrap, Target +from tinygrad.helpers import unwrap, Target # ***** X86 Ops ***** @@ -791,9 +791,7 @@ encodings = { class X86Renderer(ISARenderer): device = "CPU" has_local = False - has_threads = bool(getenv("THREADS", 1)) - @property - def global_max(self): return (NUM_CPU_THREADS.value, 0, 0) # type: ignore[override] + global_max = (1, 0, 0) 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 2b47bb8ea9..0664d79a17 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -5,7 +5,7 @@ from tinygrad.renderer.cstyle import HIPRenderer, create_non_native_float_pats, from tinygrad.codegen.decomp.transcendental import xexp2, xlog2 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp, range_str from tinygrad.dtype import dtypes, float_to_fp8, DType, truncate, AddrSpace -from tinygrad.helpers import prod, Target, NUM_CPU_THREADS, getenv, OSX +from tinygrad.helpers import prod, Target, OSX def is_volatile(u:UOp) -> bool: return (buf:=u.buf_uop).op is Ops.PARAM and buf.arg.volatile @@ -203,9 +203,7 @@ class LLVMRenderer(Renderer): class CPULLVMRenderer(LLVMRenderer): has_local = False - has_threads = bool(getenv("THREADS", 1)) - @property - def global_max(self): return (NUM_CPU_THREADS.value, 0, 0) # type: ignore[override] + global_max = (1, 0, 0) 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))) diff --git a/tinygrad/runtime/graph/hcq.py b/tinygrad/runtime/graph/hcq.py index 39e68a4430..3ce29acde6 100644 --- a/tinygrad/runtime/graph/hcq.py +++ b/tinygrad/runtime/graph/hcq.py @@ -96,7 +96,6 @@ class HCQGraph(MultiGraphRunner): # set any fixedvars on the device self.device_vars[enqueue_dev] = merge_dicts([self.device_vars.get(enqueue_dev, {}), device_vars]) - if runtime is not None: self.device_vars[enqueue_dev] = merge_dicts([self.device_vars[enqueue_dev], {k: 0 for k in ast.arg.runtimevars}]) if runtime is not None: enqueue_queue = self.comp_queues[enqueue_dev] diff --git a/tinygrad/runtime/ops_cpu.py b/tinygrad/runtime/ops_cpu.py index 048c23ce11..ae9967a3d5 100644 --- a/tinygrad/runtime/ops_cpu.py +++ b/tinygrad/runtime/ops_cpu.py @@ -1,6 +1,6 @@ from __future__ import annotations -import platform, sys, os, ctypes, functools, mmap, threading, array, struct, time -from dataclasses import dataclass, replace +import platform, sys, ctypes, functools, mmap, array, struct, time +from dataclasses import replace from typing import cast, Callable from tinygrad.helpers import to_mv, from_mv, OSX, WIN, Context, mv_address, suppress_finalizing, unwrap, data64_le from tinygrad.device import Buffer, BufferSpec, TinyELF, Program, Device @@ -19,10 +19,10 @@ from tinygrad import UOp, dtypes from tinygrad.dtype import AddrSpace from tinygrad.uop.ops import KernelInfo, Ops, UPat, PatternMatcher -MAX_ARGS, CMD_SIZE, RING_SLOTS, FUNCS = 63, 64, (16 << 10), (() if WIN else ('clock_gettime', 'sem_wait', 'sem_post')) +MAX_ARGS, CMD_SIZE, FUNCS = 63, 64, (() if WIN else ('clock_gettime',)) # ***************** -# 1. workers +# 1. signal programs def signal_prog(): val = UOp.param(1, dtypes.int, (), vmin_vmax=(0, dtypes.int.max), name="value", addrspace=AddrSpace.ALU) @@ -40,21 +40,6 @@ def timestamp_prog(): val = ts.after(call)[0].load() * 1_000_000_000 + ts.after(call)[1].load() return UOp.param(0, dtypes.uint64, 1)[0].store(val) -def worker_prog(): - ring = UOp.param(0, dtypes.uint64, RING_SLOTS * CMD_SIZE, volatile=True) - wait, done = UOp.param(1, dtypes.uint64, 1, volatile=True), UOp.param(2, dtypes.uint64, 1, volatile=True) - sem, cur = UOp.param(3, dtypes.uint64, 1), UOp.range(2**64-1, 0, dtype=dtypes.uint64) # sem is unused on windows, it has to come last - - # spin on windows, sem_wait to sleep on posix - if WIN: ready = (v:=wait.after(lw:=UOp.loop(1), cur)[0].load()).end(lw, v <= cur) - else: ready = (rv:=wait.after(lw:=UOp.loop(1), cur)[0].load().call(sem.after(cur)[0], ret_dtype=dtypes.int)).end(lw, rv != 0) - - entry = [ring.after(ready).index((cur % RING_SLOTS) * CMD_SIZE + i).load() for i in range(CMD_SIZE)] - return done.after(entry[0].call(*entry[1:], ret_dtype=dtypes.void)).index(0).store(cur + 1).end(cur) - -@dataclass -class CPUWorker: ring:Buffer; put:Buffer; sem:Buffer; sys:Buffer; done:Buffer; thread:threading.Thread # noqa: E702 - # ***************** # 2. queue encoders @@ -67,10 +52,7 @@ def cpu_cmd(devs:tuple[str, ...], prog, *args:UOp) -> UOp: def cpu_exec(ctx, call:UOp, prg:UOp) -> UOp: devs = ctx.devs args = [get_call_arg_uops(call)[i].getaddr(devs) for i in prg.arg.globals] + [v.cast(dtypes.uint64) for v in get_call_var_uops(call, prg)] - if (core:=prg.arg.runtimevars.get('core_id')) is None: return cpu_cmd(devs, prg, *args) - - la = [cpu_cmd(devs,prg,*args[:(cid:=(len(prg.arg.globals)+core))],UOp.const(t,dtypes.uint64),*args[cid+1:]) for t in range(prg.arg.global_size[0])] - return UOp(Ops.LINEAR, src=tuple(la)) + return cpu_cmd(devs, prg, *args) pm_cpu_opsel = PatternMatcher([ (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), cpu_exec), @@ -85,25 +67,11 @@ pm_cpu_opsel = PatternMatcher([ ]) def cpu_submit(ctx, cmdbuf:UOp) -> UOp: - # copy the cmd entries into the worker ring and post the semaphore once per entry - assert ctx.nbytes % (CMD_SIZE * 8) == 0 and ctx.nbytes // (CMD_SIZE * 8) < RING_SLOTS, f"submit of {ctx.nbytes} bytes doesn't fit the ring" - devs, cnt, cb = ctx.devs, hcq_size_var(cmdbuf) // (CMD_SIZE * 8), cmdbuf.bitcast(dtypes.uint64) - ring, put, done, sem = (make_buf(devs, tag=f"{ctx.queue}_{n}") for n in ("ring", "put", "done", "sem")) - - # submits are serialized on the submitter, so they can bump put without atomics - ran = done.after(l:=UOp.loop(10)).index(0).load() - room = ran.end(l, put.index(0).load() - ran > (RING_SLOTS - cnt).cast(ran.dtype)) # wait until cnt entries fit in the ring - base = ((put.after(room, cmdbuf).index(0).load() % RING_SLOTS) * CMD_SIZE).cast(dtypes.int) - e = UOp.range(cnt, 11, dtype=dtypes.int, src=(cmdbuf, ring)) - # the slot is a multiple of CMD_SIZE, so a word can never wrap on its own: take the modulo once per entry, not per word - slot = (base + e*CMD_SIZE) % (RING_SLOTS * CMD_SIZE) - w = UOp.range(CMD_SIZE, 12, dtype=dtypes.int, src=(cmdbuf, ring)) - copy = ring.index(slot + w).store(cb.index(e*CMD_SIZE + w).load()).end(w) - - bumped = put.after(copy.end(e)).index(0).store(put.index(0).load() + cnt.cast(dtypes.uint64)) - if WIN: return make_buf(devs, tag=f"{ctx.queue}_sys").after(bumped).index(0).store(put.after(bumped).index(0).load()) - e = UOp.range(cnt, 13, dtype=dtypes.int, src=(bumped,)) - return make_buf(devs, tag="func:sem_post").after(e).index(0).load().call(sem.after(e).index(0), ret_dtype=dtypes.void).end(e) + # run the cmd entries inline on the submitting thread, the cpu has no worker threads + cb, cnt = cmdbuf.bitcast(dtypes.uint64), hcq_size_var(cmdbuf) // (CMD_SIZE * 8) + e = UOp.range(cnt, 10, dtype=dtypes.int, src=(cmdbuf,)) + entry = [cb.index(e*CMD_SIZE + i).load() for i in range(CMD_SIZE)] + return entry[0].call(*entry[1:], ret_dtype=dtypes.void).end(e) pm_cpu_submit = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(name="cmdbuf"),)), cpu_submit)]) @@ -119,7 +87,6 @@ class CPUProgram(Program['CPUDevice']): def __init__(self, dev:CPUDevice, obj:TinyELF): self.dev, self.name, self.signature = dev, obj.name, obj.signature - self.runtimevars = {name:slot for name,slot,*_ in obj.signature if name == 'core_id'} self.lvp = obj.target.renderer == "LVP" if sys.platform == "win32": # mypy doesn't understand when WIN is used here @@ -165,9 +132,7 @@ class CPUProgram(Program['CPUDevice']): else: args = [*[cast(int, b.va_addr) for b in bufs], *cast(tuple[int, ...], vals)] assert len(args) <= MAX_ARGS, f"CPU programs support at most {MAX_ARGS} arguments, got {len(args)}" - for tid in range(global_size[0]): - if 'core_id' in self.runtimevars: args[self.runtimevars['core_id']] = tid - self.fxn(*[ctypes.c_uint64(x) for x in args]) + self.fxn(*[ctypes.c_uint64(x) for x in args]) return time.perf_counter() - st if wait else None @suppress_finalizing @@ -198,48 +163,22 @@ class CPUDevice(HCQ2Compiled): pm_encode, pm_lower = {"COMPUTE": pm_cpu_opsel, "SUBMIT": pm_cpu_opsel}, {"COMPUTE": pm_cpu_submit, "SUBMIT": pm_cpu_submit} def __init__(self, device:str=""): - self.workers:list[CPUWorker] = [] super().__init__(device, CPUAllocator(self), [ClangRenderer, CPULLVMRenderer, LVPRenderer, X86Renderer], CPUProgram, arch={'amd64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine().lower(), m)+",native") self.pm_bufferize = PatternMatcher( - [(UPat(Ops.PARAM, tag=f"{q}_{n}"), lambda ctx, q=q, n=n: getattr(ctx[0].worker(q), n)) - for q in ("COMPUTE:0", "SUBMIT:0") for n in ("ring","put","sem","sys","done")] + [(UPat(Ops.PARAM, tag=f"func:{f}"), lambda ctx, f=f: ctx[0].func_ptr(f)) for f in FUNCS]) + self.pm_bufferize with Context(EMULATED_DTYPES="", TRACK_MATCH_STATS=0): clang = ClangRenderer(replace(self.renderer.target, renderer="CLANG")) self.prgs:dict[Callable, CPUProgram] = {f: CPUProgram(self, do_to_program(f().sink(arg=KernelInfo(f.__name__), tag=1), clang).to_elf()) - for f in (signal_prog, wait_prog, timestamp_prog, worker_prog)} + for f in (signal_prog, wait_prog, timestamp_prog)} def func_ptr(self, name:str) -> Buffer: return self.func_table.view(1, dtypes.uint64, FUNCS.index(name)*8).ensure_allocated() - def synchronize(self, timeout:int|None=None): - for worker in self.workers: - put, done = (getattr(worker, x)._buf.cpu_view().view(fmt='Q') for x in ("put", "done")) - while done[0] < put[0]: self._wait_signal(done, put[0], timeout) - super().synchronize(timeout) - @functools.cached_property def func_table(self) -> Buffer: lib = ctypes.windll.kernel32 if sys.platform == "win32" else libc.dll # type: ignore[attr-defined] (ft:=Buffer(self.device, len(FUNCS), dtypes.uint64, preallocate=True))._buf.cpu_view().view(fmt='Q')[:] = \ array.array('Q', [unwrap(ctypes.cast(getattr(lib, f), ctypes.c_void_p).value) for f in FUNCS]) return ft - - @functools.cache - def worker(self, queue:str) -> CPUWorker: - ring, put, sysbuf, done = (Buffer(self.device, sz, dtypes.uint64, preallocate=True) for sz in (RING_SLOTS*CMD_SIZE, 1, 1, 1)) - addr, hsem = 0, None - - # sem are posix-only - if not WIN: - hsem = libc.sem_open(nm:=f"/tinygrad-{os.getpid()}-{id(ring):x}".encode(), os.O_CREAT|os.O_EXCL, 0o600, 0) # type: ignore[call-arg] - if (addr:=unwrap(ctypes.cast(hsem, ctypes.c_void_p).value)) == ctypes.c_void_p(-1).value or libc.sem_unlink(nm): - raise OSError(ctypes.get_errno(), "semaphore") - sem = Buffer(self.device, 1, dtypes.uint64, options=BufferSpec(external_ptr=addr), preallocate=True) - - worker_args = [ring._buf.va_addr, sysbuf._buf.va_addr if WIN else self.func_ptr('sem_wait')._buf.va_addr, done._buf.va_addr, addr] - (thread:=threading.Thread(target=self.prgs[worker_prog].fxn, daemon=True, args=[ctypes.c_uint64(x) for x in worker_args])).start() - self.workers.append(worker:=CPUWorker(ring, put, sem, sysbuf, done, thread)) - return worker diff --git a/tinygrad/runtime/ops_dsp.py b/tinygrad/runtime/ops_dsp.py index e487f666e7..2de37e791c 100644 --- a/tinygrad/runtime/ops_dsp.py +++ b/tinygrad/runtime/ops_dsp.py @@ -10,7 +10,6 @@ from tinygrad.runtime.autogen import libc, qcom_dsp if getenv("IOCTL"): import extra.dsp.run # noqa: F401 # pylint: disable=unused-import class DSPRenderer(ClangRenderer): - has_threads = False buffer_suffix = " restrict __attribute__((align_value(128)))" kernel_typedef = "__attribute__((noinline)) void" type_map = { **ClangRenderer.type_map, dtypes.uint64: "unsigned long long", dtypes.int64: "long long" } diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index af3ad32a63..5c8519b549 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -167,6 +167,11 @@ def _merge_submits(calls:list[UOp]) -> UOp: estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates()).simplify())) def _merge_queues(submits:list[UOp]) -> list[UOp]: + # CPU submits run inline and can block on another queue. Keep multi-queue CPU work in schedule order so every + # producer queue is submitted before a CPU wait; merging by queue can otherwise deadlock alternating dependencies. + keys = [unwrap(get_submit(call)).src[0].arg for call in submits] + if len(set(keys)) > 1 and any(any(d.split(":")[0] == "CPU" for d in devs) for devs, _ in keys): return submits + merged:list[UOp] = [] opened:dict[tuple[tuple[str, ...], str], list[UOp]] = {} # (devs, queue) -> hcq calls in submit order limits:dict[tuple[tuple[str, ...], str], int] = collections.defaultdict(lambda: JIT_BATCH_SIZE.value) @@ -412,8 +417,8 @@ pm_lower_hcq = PatternMatcher([ (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq", src=(UPat(Ops.SINK),)),), name="call", allow_any_len=True), lower_hcq_call)]) # ***************** -# 6. batch: adjacent hcq calls fold into one submitter on the host SUBMIT:0 ring: a submit whose cmds call the -# compiled piece programs, so the worker runs the batch in fifo order and the python exec is one ring push +# 6. batch: adjacent hcq calls fold into one submitter on the host SUBMIT:0 queue: a submit whose cmds call the +# compiled piece programs, so the batch runs in fifo order and the python exec is one submitter call def _lane_arg(a:UOp, lane:int) -> UOp: return a.mselect(lane) if len(to_tuple(a.device)) > 1 else a diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 196c085133..486f1dbf37 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -17,7 +17,7 @@ if TYPE_CHECKING: class AxisType(Enum): def __repr__(self): return str(self) DEVICE = auto(); GLOBAL = auto(); WARP = auto(); LOCAL = auto(); WEAK = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto() # noqa: E702 - UNROLL = auto(); THREAD = auto(); PLACEHOLDER = auto(); LOOP = auto() # noqa: E702 + UNROLL = auto(); PLACEHOLDER = auto(); LOOP = auto() # noqa: E702 @dataclass(frozen=True, order=True) class ParamArg: @@ -39,14 +39,14 @@ class ParamArg: args = [repr(self.slot), repr(self.dtype)] + ([repr(self.size)] if self.size is not None else []) + \ [f"{k}={v!r}" for k,default in fields if (v:=getattr(self, k)) != default] return f"ParamArg({', '.join(args)})" -axis_letters = {AxisType.DEVICE: "d", AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.WEAK: "L", +axis_letters = {AxisType.DEVICE: "d", AxisType.GLOBAL: "g", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.WEAK: "L", AxisType.LOOP: "L", AxisType.UPCAST: "u", AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"} -axis_colors = {AxisType.DEVICE: "green", AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", +axis_colors = {AxisType.DEVICE: "green", AxisType.GLOBAL: "blue", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.WEAK: "WHITE", AxisType.LOOP: "WHITE", AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"} # NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters -axis_to_pos = {AxisType.DEVICE: -2, AxisType.WEAK: -1, AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, +axis_to_pos = {AxisType.DEVICE: -2, AxisType.WEAK: -1, AxisType.LOOP: -1, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3, AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5} range_start = {Ops.STAGE: 1, Ops.REDUCE: 1, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.LINEAR: 0} @@ -1268,16 +1268,13 @@ class ProgramInfo: @property def function_name(self): return to_function_name(self.name) - @property - def runtimevars(self) -> dict[str, int]: return {v.expr: i for i, v in enumerate(self.vars) if v.expr == 'core_id'} - def launch_dims(self, var_vals:dict[str, int]) -> tuple[tuple[int, ...], tuple[int, ...]|None]: global_size = tuple([sym_infer(sz, var_vals) for sz in self.global_size]) # type: ignore[arg-type] local_size = tuple([sym_infer(sz, var_vals) for sz in self.local_size]) if self.local_size is not None else None return global_size, local_size - def vals(self, var_vals:dict[str, int]) -> tuple[int|None, ...]: - try: return tuple(var_vals[k.expr] if k.expr not in self.runtimevars else None for k in self.vars) + def vals(self, var_vals:dict[str, int]) -> tuple[int, ...]: + try: return tuple(var_vals[k.expr] for k in self.vars) except KeyError as e: raise RuntimeError(f"unbound Variable {e} used by {self.function_name}") from None @staticmethod @@ -1298,7 +1295,6 @@ class ProgramInfo: if u.arg[0] == 'i': local_size = None special_size = local_size if u.arg[0] == 'l' else global_size if special_size is not None: special_size[int(u.arg[-1])] = cast(int, u.src[0].ssimplify()) - if u.op is Ops.PARAM and u in _vars and u.expr == 'core_id': global_size[0] = int(u.vmax) + 1 return ProgramInfo(sink.arg.name if isinstance(sink.arg, KernelInfo) else "test", tuple(global_size), tuple(local_size) if local_size is not None else None, tuple(sorted(dedup(_vars), key=lambda v: v.arg.slot)), tuple(sorted(dedup(_globals))), tuple(sorted(dedup(outs))), tuple(sorted(dedup(ins))), target)