forked from tinygrad/tinygrad
delete CPU threading support (#17866)
* delete CPU threading support * more removals * fixes
This commit is contained in:
George Hotz
@@ -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
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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}
|
||||
|
||||
@@ -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={}):
|
||||
"""
|
||||
|
||||
Binary file not shown.
@@ -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. \\
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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])
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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:
|
||||
|
||||
+2
-2
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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("")'
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)))
|
||||
|
||||
@@ -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]
|
||||
|
||||
+12
-73
@@ -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
|
||||
|
||||
@@ -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" }
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
+6
-10
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user