geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions 7 Activity

remove Ops.SLICE (#17492)

Browse Source
This commit is contained in:
sirhcm
2026-08-11 18:50:04 -04:00
committed by GitHub
parent 2b5018e86a
commit 479ffb0cda
12 changed files with 52 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/null/test_memory_planner.py
+8 -3
View File
@@ -27,10 +27,15 @@ def _make_linear(buffer_lists, copies=None):
calls.append(UOp(Ops.CALL, src=(src0, *bufs)))
return UOp(Ops.LINEAR, src=tuple(calls))
def _get_planned_view(buf:UOp) -> tuple[UOp, int, int]|None:
view = buf.src[0] if buf.op is Ops.BITCAST else buf
if view.op is not Ops.SHRINK or view.src[0].op is not Ops.BUFFER: return None
return (arena:=view.src[0]), view.src[1].val * arena.dtype.itemsize, view.src[2].val * arena.dtype.itemsize
def _get_arena(buf, linear, result):
for orig_si, new_si in zip(linear.src, result.src):
for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
if orig is buf and new.op is Ops.SLICE: return new.src[0]
if orig is buf and (planned:=_get_planned_view(new)) is not None: return planned[0]
return None
def check_assign(buffer_lists, copies=None):
@@ -41,8 +46,8 @@ def check_assign(buffer_lists, copies=None):
replace_map: dict[int, tuple[UOp, int, int]] = {}
for orig_si, new_si in zip(linear.src, result.src):
for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
if new.op is Ops.SLICE and id(orig) not in replace_map:
replace_map[id(orig)] = (new.src[0], new.src[1].val * new.src[0].dtype.itemsize, new.arg * new.dtype.itemsize)
if (planned:=_get_planned_view(new)) is not None and id(orig) not in replace_map:
replace_map[id(orig)] = planned
# verify pinned buffers are not planned
for buf in held_bufs:
test/unit/test_metal_graph.py
+15 -25
View File
@@ -1,5 +1,4 @@
import unittest
from unittest.mock import MagicMock
from tinygrad import Device
from tinygrad.uop.ops import Ops, UOp
from tinygrad.dtype import dtypes
@@ -11,36 +10,27 @@ class TestMetalGraph(unittest.TestCase):
self.MetalGraph = MetalGraph
self.dev = Device[Device.DEFAULT]
def metal_buf(self, offset):
buf = MagicMock()
if offset > 0:
buf.op = Ops.SLICE
src = MagicMock()
src.dtype = dtypes.uint8
buf.src = (src, UOp.const(offset))
buf.dtype = dtypes.uint8
else:
buf.op = Ops.BUFFER
buf.device = Device.DEFAULT
return buf
def metal_buf(self, offset, bitcast=False):
size = 4 if bitcast else 1
buf = UOp.new_buffer(Device.DEFAULT, offset+size, dtypes.uint8)
if offset: buf = buf[offset:offset+size]
return buf.bitcast(dtypes.float32) if bitcast else buf
def call(self, *bufs):
c = MagicMock()
c.src = (MagicMock(op=Ops.PROGRAM),) + tuple(bufs)
return c
def supports_uop(self, *bufs):
return self.MetalGraph.supports_uop([self.dev], UOp(Ops.PROGRAM, src=(UOp.sink(),)).call(*bufs))
def test_supports_uop_normal_offset(self):
assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(100), self.metal_buf(0xFFFFFFFF))) is True
assert self.supports_uop(self.metal_buf(0), self.metal_buf(100), self.metal_buf(0xFFFFFFFF)) is True
def test_supports_uop_overflow_offset(self):
assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(0x100000000))) is False
assert self.supports_uop(self.metal_buf(0), self.metal_buf(0x100000000)) is False
def test_supports_uop_nonmetal_buf(self):
# non-SLICE ops should not be checked for offset
buf = MagicMock()
buf.op = Ops.BUFFER
buf.device = Device.DEFAULT
self.MetalGraph.supports_uop([self.dev], self.call(buf))
def test_supports_uop_non_view_buf(self):
assert self.supports_uop(self.metal_buf(0)) is True
def test_supports_uop_bitcast(self):
assert self.supports_uop(self.metal_buf(0xFFFFFFFF, bitcast=True)) is True
assert self.supports_uop(self.metal_buf(0x100000000, bitcast=True)) is False
if __name__ == "__main__":
unittest.main()
tinygrad/engine/jit.py
+1 -3
View File
@@ -44,8 +44,6 @@ def graph_split_rewrite(linear:UOp, max_batch_size:int=0) -> UOp:
current_batch, current_batch_devs = [], []
for si in linear.src:
if si.src[0].op is Ops.SLICE: continue
devs = dedup([Device[x] for b in si.src[1:] if b.op is not Ops.BIND for x in (b.device if isinstance(b.device, tuple) else (b.device,))])
graph_t = graph_class(devs[0]) if devs[0].graph is not None else None
@@ -180,7 +178,7 @@ class CapturedJit(Generic[ReturnType]):
if call.op is not Ops.CALL: continue
arg_uops = get_call_arg_uops(call)
outs, ins = get_call_outs_ins(call)
out |= {arg_uops[k] for k in set(outs) - set(ins) if arg_uops[k].op in (Ops.BUFFER, Ops.SLICE)}
out |= {b for k in set(outs) - set(ins) if (b:=u if (cv:=(u:=arg_uops[k]).contiguous_view()) is None else cv[0]).op is Ops.BUFFER}
return out
def __call__(self, input_uops:list[UOp], var_vals:dict[str, int]) -> ReturnType:
tinygrad/engine/realize.py
+3 -14
View File
@@ -4,7 +4,7 @@ import time, random, itertools, math, contextlib, weakref, array
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 BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, wait_cond
from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, buffers, graph_rewrite
from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite
from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
from tinygrad.renderer import Estimates
from tinygrad.codegen import to_program
@@ -17,7 +17,7 @@ def get_call_arg_uops(call:UOp) -> tuple[UOp, ...]: return tuple(s for s in call
def get_call_outs_ins(call:UOp) -> tuple[tuple[int, ...], tuple[int, ...]]:
ast = call.src[0]
if ast.op is Ops.PROGRAM: return tuple(ast.arg.outs), tuple(ast.arg.ins)
if ast.op in (Ops.COPY, Ops.SLICE): return (0,), (1,)
if ast.op is Ops.COPY: return (0,), (1,)
if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return (0,), tuple(range(1, len(get_call_arg_uops(call))))
return (), ()
@@ -27,9 +27,6 @@ def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|N
ast, arg_uops = call.src[0], get_call_arg_uops(call)
if ast.op is Ops.PROGRAM: return ast.arg.name
if ast.op is Ops.SLICE:
offset = ast.src[1].val * arg_uops[1].dtype.itemsize
return colored(f"view {_uop_sz_to_str(arg_uops[0]):>10} @ {offset:<10d}", "yellow")
if ast.op is Ops.COPY: return colored(f"copy {_uop_sz_to_str(arg_uops[0]):>10}, {_dev_str(bufs[0]):>7s} <- {_dev_str(bufs[1]):7s}", "yellow")
if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return colored(f"enc/dec {_uop_sz_to_str(arg_uops[0])}", "yellow")
if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return colored(f"batched {len(ast.src[0].src)}", "cyan")
@@ -140,7 +137,7 @@ class ExecContext:
cache: bool = True
def _resolve(b:UOp, inputs:tuple[UOp, ...]) -> UOp:
if b.op in (Ops.SLICE, Ops.MSELECT, Ops.SHRINK) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg.slot], *b.src[1:]))
if b.op in (Ops.MSELECT, Ops.SHRINK) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg.slot], *b.src[1:]))
if b.op is Ops.MSTACK: return b.replace(src=tuple(_resolve(x, inputs) for x in b.src))
return inputs[b.arg.slot] if b.op is Ops.PARAM else b
def resolve_params(call:UOp, inputs:tuple[UOp, ...]) -> list[UOp]: return [_resolve(b, inputs) for b in get_call_arg_uops(call)]
@@ -154,13 +151,6 @@ def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], d
for x in call.src[0].toposort())
for j, per_dev in enumerate(zip(*[cast(MultiBuffer, b).bufs for b in bufs])): yield list(per_dev), {"_device_num": j} if has_dnum else {}
def exec_view(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
resolved = resolve_params(call, ctx.input_uops)
bufs = [cast(Buffer, b.buffer) for b in resolved]
bv = bufs[1].view(resolved[0].max_numel(), ast.dtype, ast.src[1].val*bufs[1].dtype.itemsize)
with track_stats(ctx, call, bv.device, [bv, bufs[1]], ctx.var_vals): buffers[resolved[0]] = bv
return None
def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)):
dest, src = bufs[0].ensure_allocated(), bufs[1].ensure_allocated()
@@ -264,7 +254,6 @@ pm_optimize_local_size = PatternMatcher([
])
pm_exec = PatternMatcher([
(UPat(Ops.CALL, src=(UPat(Ops.SLICE, name="ast"),), name="call", allow_any_len=True), exec_view),
(UPat(Ops.CALL, src=(UPat(Ops.COPY, name="ast"),), name="call", allow_any_len=True), exec_copy),
(UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="ast"),), name="call", allow_any_len=True), exec_kernel),
(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="encdec", name="ast"),), name="call", allow_any_len=True), exec_encdec),
tinygrad/runtime/graph/metal.py
+2 -1
View File
@@ -113,5 +113,6 @@ class MetalGraph(GraphRunner):
@staticmethod
def supports_uop(batch_devs, new_call:UOp) -> bool:
# Metal ICB replay encodes offsets as uint32; reject if any Metal buffer offset exceeds 32-bit range.
if any(b.op in {Ops.SLICE, Ops.SHRINK} and b.src[1].val * b.src[0].dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False
for shrink in [s for src in new_call.src[1:] if (s:=src.src[0] if src.op is Ops.BITCAST else src).op is Ops.SHRINK]:
if shrink.src[1].val * shrink.src[0].dtype.itemsize > 0xFFFFFFFF: return False
return GraphRunner.supports_uop(batch_devs, new_call)
tinygrad/runtime/support/hcq2.py
+8 -7
View File
@@ -39,7 +39,7 @@ def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for
def unwrap_mstack(u):
if u.op is Ops.MSTACK: return tuple(x for s in u.src for x in unwrap_mstack(s))
return unwrap_mstack(u.src[0]) if u.op in {Ops.MSELECT, Ops.SLICE} else (u,)
return unwrap_mstack(u.src[0]) if u.op is Ops.MSELECT else (u,)
def is_value_known_at_link(val:UOp) -> bool:
runtime_reads = [u for u in val.toposort() if u.op in (Ops.LOAD, Ops.INDEX)]
@@ -377,14 +377,15 @@ pm_replace_params = PatternMatcher([
# *****************
def resolve_getaddr_slice(bv:UOp, g:UOp) -> UOp:
def resolve_getaddr_view(bv:UOp, g:UOp) -> UOp:
base = bv.src[0].after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ())
itemsize = bv.src[0].dtype.itemsize if bv.src[0].without_after.op in (Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize
if bv.op is Ops.BITCAST: return UOp(Ops.GETADDR, src=(base,), arg=g.arg)
itemsize = bv.src[0].dtype.itemsize if bv.src[0].without_after.op in (Ops.BUFFER, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize
return UOp(Ops.GETADDR, src=(base,), arg=g.arg) + UOp.const(bv.src[1].val * itemsize, dtypes.uint64)
pm_early_simplify = PatternMatcher([
(UPat(Ops.GETADDR, src=(UPat.any(sl:=UPat((Ops.SLICE, Ops.SHRINK), name="bv"), sl.after(allow_any_len=True)),), name="g"), resolve_getaddr_slice),
(UPat(Ops.INDEX, src=(UPat(Ops.SLICE, name="bv"),), allow_any_len=True, name="x"),
(UPat(Ops.GETADDR, src=(UPat((Ops.SHRINK, Ops.BITCAST), name="bv").or_after(),), name="g"), resolve_getaddr_view),
(UPat(Ops.INDEX, src=(UPat(Ops.SHRINK, name="bv"),), allow_any_len=True, name="x"),
lambda bv,x: x.replace(src=(bv.src[0], x.src[1] + bv.src[1].cast(x.src[1].dtype), *x.src[2:]))),
])
@@ -402,7 +403,7 @@ def pack_hcq_placeholders(call:UOp) -> UOp|None:
sizes[b.tag] = offs[b] + b.max_numel()
counts = collections.Counter(b.tag for b in bufs)
bases = {b.tag:UOp.placeholder((sizes[b.tag],), b.dtype, next(UOp.unique_num), device=b.device).rtag(b.tag) for b in bufs if counts[b.tag] > 1}
subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(offs.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases}
subs = {b:bases[b.tag][(off:=offs.get(b, 0)):off+b.max_numel()] for b in bufs if b.tag in bases}
return call.replace(src=(call.src[0].substitute(subs, walk=True), *call.src[1:])) if subs else None
pm_pack_placeholders = PatternMatcher([
(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), pack_hcq_placeholders)])
@@ -493,7 +494,7 @@ pm_resolve_patches = PatternMatcher([
(UPat(name="buf").index(UPat(Ops.RANGE), allow_any_len=True)
.store(UPat.any(UPat(Ops.BINARY, name="blob"), UPat(Ops.BINARY, name="blob").bitcast()).index(UPat(Ops.RANGE), allow_any_len=True).load())
.end(UPat(Ops.RANGE)), fold_binary),
(UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf").index(UPat(Ops.STACK, name="off")).store(UPat(Ops.STACK, name="val")), fold_const_store),
(UPat({Ops.BUFFER, Ops.MSTACK}, name="buf").index(UPat(Ops.STACK, name="off")).store(UPat(Ops.STACK, name="val")), fold_const_store),
])
pm_assert_no_afters = PatternMatcher([(UPat(Ops.AFTER, name="a"), lambda a: panic(RuntimeError, f"AFTER left at hcq_link: {a.src[0].op}"))])
tinygrad/schedule/indexing.py
+3 -7
View File
@@ -23,7 +23,7 @@ class IndexingContext:
return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(0)
ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.BUFFER, Ops.SLICE,
ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.BUFFER,
Ops.CONST, Ops.BIND, Ops.MSELECT, Ops.MSTACK, Ops.PARAM,
Ops.LOAD, Ops.CALL, Ops.FUNCTION}
@@ -34,10 +34,6 @@ def realize_srcs(ctx:IndexingContext, rb:UOp) -> None:
if s.base.op not in ALWAYS_CONTIGUOUS: ctx.realize_map[s] = None
def realize_store_after_src(ctx:IndexingContext, dest:UOp, src:UOp):
# don't realize SLICE when it's the direct source of STORE+AFTER — the target buffer is the output
if src.op is Ops.SLICE and src in ctx.realize_map \
and not dest.op_in_backward_slice_with_self(Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD):
del ctx.realize_map[src]
# you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce
if dest.base in src.backward_slice_with_self: ctx.realize_map[src] = None
@@ -74,7 +70,7 @@ def broadcast_rngs(x:UOp, src:UOp, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]:
# TODO: srcs contain (real data srcs, something else, ranges) and the boundary is confusing. see range_start
def data_srcs(op:Ops, src:tuple[UOp, ...]) -> tuple[UOp, ...]:
if op in {Ops.PARAM, Ops.BUFFER, Ops.RANGE, Ops.SPECIAL, Ops.BIND}: return ()
if op in GroupOp.Movement|{Ops.INDEX, Ops.SLICE, Ops.STAGE, Ops.REDUCE, Ops.AFTER, Ops.END}: return src[:1]
if op in GroupOp.Movement|{Ops.INDEX, Ops.STAGE, Ops.REDUCE, Ops.AFTER, Ops.END}: return src[:1]
return src
def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]:
@@ -84,7 +80,7 @@ def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]:
for i, s in enumerate(x.src):
new_src = s
src_rngs = broadcast_rngs(x, s, ctx.range_map[x][0]) if x in ctx.range_map else ()
if s.op in {Ops.PARAM, Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
if s.op in {Ops.PARAM, Ops.BUFFER, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
if x in ctx.range_map and i < data_src_count: new_src = new_src.index(*src_rngs)
elif s in ctx.realize_map:
realized_ranges = ctx.realize_map[s]
tinygrad/schedule/memory.py
+2 -4
View File
@@ -52,11 +52,9 @@ def memory_plan_rewrite(linear:UOp, held_bufs:set[UOp]|None=None) -> UOp:
peaks[_key(buf)] = (max(peaks[_key(buf)][0], offsets[buf] + buf.max_numel() * buf.dtype.itemsize), peaks[_key(buf)][1])
arena_sizes = {key: round_up(peak, block_size) for key, (peak, _) in peaks.items()}
# build replace_map: each buffer becomes a SLICE into a shared per-device-lane arena
# build replace_map: each buffer becomes a SHRINK/BITCAST into a shared per-device-lane arena
arenas = {key: UOp.new_buffer(key[0], sz, dtypes.int8) for key, sz in arena_sizes.items()}
replace_map:dict[UOp, UOp] = {}
for buf_uop, offset in offsets.items():
replace_map[buf_uop] = UOp(Ops.SLICE, buf_uop.dtype, (arenas[_key(buf_uop)], UOp.const(offset)), buf_uop.max_numel())
replace_map = {buf_uop:arenas[_key(buf_uop)][offset:offset+buf_uop.nbytes()].bitcast(buf_uop.dtype) for buf_uop, offset in offsets.items()}
if DEBUG >= 1 and (omem:=sum(nbytes.values()) / 1e6) != (nmem:=sum(arena_sizes.values()) / 1e6):
print(f"memory reduced from {omem:.2f} MB -> {nmem:.2f} MB, {len(first_appearance)} -> {len(arenas)} bufs")
tinygrad/uop/__init__.py
+1 -1
View File
@@ -93,7 +93,7 @@ class Ops(FastEnum):
CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto()
# buffer ops
STAGE = auto(); COPY = auto(); SLICE = auto(); MSELECT = auto(); MSTACK = auto(); CUSTOM_FUNCTION = auto()
STAGE = auto(); COPY = auto(); MSELECT = auto(); MSTACK = auto(); CUSTOM_FUNCTION = auto()
# the core 6 movement ops! these only exist in the tensor graph
RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); FLIP = auto()
tinygrad/uop/ops.py
+8 -25
View File
@@ -45,8 +45,7 @@ axis_colors = {AxisType.DEVICE: "green", AxisType.GLOBAL: "blue", AxisType.THREA
axis_to_pos = {AxisType.DEVICE: -2, AxisType.WEAK: -1, AxisType.LOOP: -1, AxisType.THREAD: 0, 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.FUNCTION: 1,
Ops.SLICE: 2, Ops.LINEAR: 0}
range_start = {Ops.STAGE: 1, Ops.REDUCE: 1, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.FUNCTION: 1, Ops.LINEAR: 0}
# https://en.wikipedia.org/wiki/Identity_element
def identity_element(op:Ops, dt:DType) -> PyConst: return dt.const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dt.min}[op])
@@ -171,9 +170,6 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None:
return arg.dtype
case Ops.BINARY:
return dtypes.uint8
case Ops.SLICE:
# TODO: slice just shouldn't exist
return None
case Ops.CAST | Ops.BITCAST:
assert isinstance(arg, DType), f"CAST/BITCAST arg must be DType, got {arg}"
return arg
@@ -221,7 +217,7 @@ class UOpMetaClass(type):
return created
# some uops map to other stuff
buffers:weakref.WeakKeyDictionary[UOp, Buffer|MultiBuffer] = weakref.WeakKeyDictionary() # this maps BUFFER/SLICE uops to their device Buffers
buffers:weakref.WeakKeyDictionary[UOp, Buffer|MultiBuffer] = weakref.WeakKeyDictionary() # this maps BUFFER/view uops to their device Buffers
all_metadata:weakref.WeakKeyDictionary[UOp, tuple[Metadata, ...]] = weakref.WeakKeyDictionary() # TODO: should this be here?
# recursive_property replaces functools.cached_property in recursive UOp functions to prevent RecursionError
@@ -386,10 +382,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
case Ops.BUFFER:
if len(self.src): return self.src[0].as_shape
return ()
case Ops.SLICE:
# HACK: SLICE is used inside kernels, so we set the shape to () if it's on an INDEX
if self.src[0].op is Ops.INDEX: return ()
return (self.arg,)
case Ops.CUSTOM | Ops.CUSTOMI:
if self.dtype is dtypes.void: return None
input_shapes = [x._shape for x in self.src if x._shape is not None]
@@ -822,7 +814,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
unique_num = itertools.count(0)
def getaddr(self, device=None) -> UOp:
if self.without_after.op not in {Ops.BUFFER, Ops.SLICE, Ops.SHRINK, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM}: return self
if self.without_after.op not in {Ops.BUFFER, Ops.SHRINK, Ops.BITCAST, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM}: return self
return UOp(Ops.GETADDR, src=(self,), arg=device or to_tuple(self.device)[0])
@staticmethod
def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None):
@@ -924,7 +916,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
# TODO: this is confusing because UOp.variable('v', 0, 1, dtypes.weakfloat) is True for jit to work, but it doesn't have a buffer
if self.op in {Ops.RESHAPE, Ops.UNSHARD, Ops.MSELECT}: return self.src[0].has_buffer_identity(after_ok)
if after_ok and self.op == Ops.AFTER: return self.src[0].has_buffer_identity(after_ok)
return self.op in {Ops.BUFFER, Ops.SLICE, Ops.PARAM}
return self.op in {Ops.BUFFER, Ops.PARAM}
def _base_buffer_is_realized(self) -> bool:
"""Walk through AFTER chain to find if the underlying buffer is realized (has allocated memory)."""
@@ -937,25 +929,16 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
if self.op in {Ops.CONTIGUOUS, Ops.CONTIGUOUS_BACKWARD, Ops.RESHAPE, Ops.UNSHARD, Ops.DETACH, Ops.AFTER}: return self.src[0].buffer
# this buffer can process disk tensors and simple movement ops
if self is not self.base or self.op is Ops.BITCAST:
if (cret:=buffers.get(self)) is not None: return cret
if (cv := self.contiguous_view()) is None: raise RuntimeError(f"non-contiguous view is not supported for {self.device} buffer")
buf, offset = (b:=cv[0]).base.buffer, cv[1]
if isinstance(buf, MultiBuffer):
mbuf = MultiBuffer.__new__(MultiBuffer)
mbuf.bufs = [x.view(prod(self.max_shape), self.dtype, offset*b.dtype.itemsize) for x in buf.bufs]
return mbuf
return buf.view(prod(self.max_shape), self.dtype, offset*b.dtype.itemsize)
if self.op is Ops.SLICE:
if (cret:=buffers.get(self)) is not None: return cret
buf = self.src[0].buffer
offset = self.src[1].val
if isinstance(buf, MultiBuffer):
mbuf = MultiBuffer.__new__(MultiBuffer)
mbuf.bufs = [b.view(self.arg, self.dtype, offset * self.src[0].dtype.itemsize) for b in buf.bufs]
buffers[self] = mbuf
return mbuf
assert isinstance(buf, Buffer), "must be a Buffer for SLICE"
buffers[self] = bv = buf.view(self.arg, self.dtype, offset * self.src[0].dtype.itemsize)
return bv
buffers[self] = buf.view(prod(self.max_shape), self.dtype, offset*b.dtype.itemsize)
return buffers[self]
if self.op is Ops.MSELECT:
ret = self.src[0].buffer
assert isinstance(ret, MultiBuffer)
@@ -1181,7 +1164,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
def custom_function(name:str, *src:UOp) -> UOp: return UOp(Ops.CUSTOM_FUNCTION, src=src, arg=name)
# opaque bodies stay as Ops.CALL; value-producing bodies become Ops.FUNCTION (wrapped in TUPLE)
_OPAQUE_CALL_BODIES = {Ops.SINK, Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.SLICE, Ops.CUSTOM_FUNCTION}
_OPAQUE_CALL_BODIES = {Ops.SINK, Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.CUSTOM_FUNCTION}
def call(self, *srcs:UOp, ret_dtype:DType|None=None, grad_fxn:Callable|None=None,
name:str|None=None, precompile:bool=False, precompile_backward:bool=False, aux:Any=None) -> UOp:
if ret_dtype is not None: return UOp(Ops.CALL, ret_dtype, src=(self,)+srcs)
tinygrad/uop/spec.py
-7
View File
@@ -233,13 +233,6 @@ spec_hcq = PatternMatcher([
spec_full = PatternMatcher([
(UPat(Ops.REWRITE_ERROR, dtypes.void, name="x"), lambda x: isinstance(x.arg, str)),
# SLICE on BUFFER is allowed if BUFFER is
(UPat(Ops.SLICE, src=(UPat(GroupOp.Movement.union({Ops.BUFFER, Ops.PARAM, Ops.STAGE, Ops.AFTER})),
UPat(Ops.CONST, dtype=dtypes.weakint)), allow_any_len=True, name="bv"),
lambda bv: isinstance(bv.arg, int)),
(UPat(Ops.CALL, dtypes.void, src=(UPat((Ops.SLICE,)),), allow_any_len=True), lambda: True),
# codegen may end ranges after gpudims has replaced RANGE with SPECIAL.
(UPat(Ops.END, src=(UPat(), UPat()), allow_any_len=True), lambda: True),
tinygrad/viz/serve.py
+1 -1
View File
@@ -50,7 +50,7 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
Ops.INDEX: "#CEF9B7", Ops.STACK: "#D8F9E4",
Ops.WMMA: "#efefc0", Ops.UNSHARD: "#f6ccff", Ops.INS: "#eec4ff",
**{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
Ops.SLICE: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",
Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",
Ops.CALL: "#00B7C8", Ops.FUNCTION: "#C07788", Ops.PARAM: "#14686F", Ops.SOURCE: "#c0c0c0", Ops.BINARY: "#404040",
Ops.LINEAR: "#7DF4FF",
Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",