mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 10:56:08 +00:00
dtype is not a UOp field anymore (#17821)
This commit is contained in:
chenyu
@@ -215,7 +215,7 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
|
||||
k = UOp.range(K, 0, AxisType.REDUCE)
|
||||
mul = (A.flatten().index((m*UOp.const(K)+k))*
|
||||
B.flatten().index((k*UOp.const(N)+n))).cast(dtypes.float32)
|
||||
red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype)
|
||||
red = mul.reduce(k, arg=Ops.ADD).cast(C.dtype)
|
||||
store = C.flatten().index((m*UOp.const(N)+n)).store(red).end(m, n)
|
||||
return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
|
||||
|
||||
|
||||
@@ -21,7 +21,7 @@ def run_uops(uops_list:list[UOp], bufs:list[Buffer]):
|
||||
def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
|
||||
if op is Ops.CONST: uops.append(UOp.const(arg).cast(dtype))
|
||||
elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, 1))
|
||||
else: uops.append(UOp(op, dtype, tuple(src), arg))
|
||||
else: uops.append(UOp(op, tuple(src), arg))
|
||||
return uops[-1]
|
||||
|
||||
def _test_single_value(vals, op, dts):
|
||||
|
||||
+2
-2
@@ -7,8 +7,8 @@ BENCHMARK_OPS = {Ops.INDEX, Ops.STAGE}
|
||||
|
||||
@functools.cache
|
||||
def create_uop(a:int) -> UOp:
|
||||
op, dtype, src, arg, *rest = trace.uop_fields[a]
|
||||
return UOp(op, dtype, tuple(create_uop(s) for s in src), arg, *rest)
|
||||
op, src, arg, *rest = trace.uop_fields[a]
|
||||
return UOp(op, tuple(create_uop(s) for s in src), arg, *rest)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# load rewrite trace
|
||||
|
||||
@@ -40,7 +40,7 @@ class TestPatternMatcher(unittest.TestCase):
|
||||
assert len(x.src) == 0
|
||||
return UOp(Ops.CONST, src=(UOp(Ops.CONST),))
|
||||
matcher = PatternMatcher([(UPat(Ops.CONST, src=(), name="x"), fxn)])
|
||||
c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
|
||||
c1 = UOp(Ops.CONST, arg=1.0)
|
||||
# second rewrite shouldn't match anything
|
||||
c1 = matcher.rewrite(c1)
|
||||
c1 = matcher.rewrite(c1)
|
||||
|
||||
@@ -4,21 +4,21 @@ from tinygrad import UOp
|
||||
class TestUOpRepr(unittest.TestCase):
|
||||
def test_simple_const(self):
|
||||
a = UOp.const(42)
|
||||
self.assertEqual(repr(a), "UOp(Ops.CONST, dtypes.weakint, arg=42, src=())")
|
||||
self.assertEqual(repr(a), "UOp(Ops.CONST, arg=42, src=())")
|
||||
def test_different_consts(self):
|
||||
a, b = UOp.const(42), UOp.const(3)
|
||||
expected = (
|
||||
"UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
|
||||
" UOp(Ops.CONST, dtypes.weakint, arg=42, src=()),\n" +
|
||||
" UOp(Ops.CONST, dtypes.weakint, arg=3, src=()),))"
|
||||
"UOp(Ops.ADD, arg=None, src=(\n" +
|
||||
" UOp(Ops.CONST, arg=42, src=()),\n" +
|
||||
" UOp(Ops.CONST, arg=3, src=()),))"
|
||||
)
|
||||
self.assertEqual(repr(a+b), expected)
|
||||
def test_walrus_operator_indentation(self):
|
||||
# The reference should have the same indentation as the definition
|
||||
a = UOp.const(42)
|
||||
expected = (
|
||||
"UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
|
||||
" x0:=UOp(Ops.CONST, dtypes.weakint, arg=42, src=()),\n" +
|
||||
"UOp(Ops.ADD, arg=None, src=(\n" +
|
||||
" x0:=UOp(Ops.CONST, arg=42, src=()),\n" +
|
||||
" x0,))"
|
||||
)
|
||||
self.assertEqual(repr(a+a), expected)
|
||||
@@ -26,9 +26,9 @@ class TestUOpRepr(unittest.TestCase):
|
||||
# Ensure indentation is consistent at multiple levels
|
||||
b = (a:=UOp.const(1)) + a
|
||||
expected = (
|
||||
"UOp(Ops.MUL, dtypes.weakint, arg=None, src=(\n" +
|
||||
" x0:=UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
|
||||
" x1:=UOp(Ops.CONST, dtypes.weakint, arg=1, src=()),\n" +
|
||||
"UOp(Ops.MUL, arg=None, src=(\n" +
|
||||
" x0:=UOp(Ops.ADD, arg=None, src=(\n" +
|
||||
" x1:=UOp(Ops.CONST, arg=1, src=()),\n" +
|
||||
" x1,)),\n" +
|
||||
" x0,))"
|
||||
)
|
||||
|
||||
@@ -38,15 +38,9 @@ class TestDTypeFromUOp(unittest.TestCase):
|
||||
self.assertEqual(UOp(Ops.CONST, arg=ConstFloat(3.0)).dtype, dtypes.weakfloat)
|
||||
self.assertEqual(UOp(Ops.CONST, arg=True).dtype, dtypes.bool)
|
||||
self.assertEqual(UOp(Ops.CONST, arg=Invalid).dtype, dtypes.bool)
|
||||
# an explicit (strong) const dtype is legal until the field is removed
|
||||
# UOp.const at a strong dtype builds the CAST that carries it
|
||||
self.assertEqual(UOp.const(3, dtypes.int32).dtype, dtypes.int32)
|
||||
|
||||
def test_invalid_stated_dtype(self):
|
||||
# UOp.const normalizes a stated dtype away (const_like/full pass their position's); the core constructor does not,
|
||||
# and the spec is what rejects a non-bool Invalid
|
||||
self.assertIs(UOp.const(Invalid, dtypes.float32), UOp.invalid())
|
||||
with self.assertRaises(RuntimeError): type_verify(UOp(Ops.CONST, dtypes.float32, arg=Invalid), spec_shared)
|
||||
|
||||
def test_invalid_dtype_and_consumers(self):
|
||||
invalid = UOp.invalid()
|
||||
self.assertIs(invalid.dtype, dtypes.bool)
|
||||
|
||||
@@ -191,7 +191,7 @@ class TestWeakPromotion(unittest.TestCase):
|
||||
# float bitwise builds, the spec rejects it
|
||||
with Context(SPEC=1):
|
||||
f32, wf = UOp.const(1.0, dtypes.float32), UOp.const(1.0)
|
||||
for bad in (f32.alu(Ops.AND, f32), UOp(Ops.AND, dtypes.float32, (f32, f32)), UOp(Ops.AND, dtypes.int32, (wf, wf))):
|
||||
for bad in (f32.alu(Ops.AND, f32), UOp(Ops.AND, (f32, f32)), UOp(Ops.AND, (wf, wf))):
|
||||
with self.assertRaises(RuntimeError): type_verify([bad], spec_shared)
|
||||
|
||||
def test_integer_values(self):
|
||||
|
||||
@@ -116,7 +116,7 @@ def _amd_load(ptr:UOp, lanes:int|None=None) -> UOp:
|
||||
if lanes is None: return ptr.load(arg="nontemporal")
|
||||
buf, coords = ptr.src[0], ptr.src[1:]
|
||||
idx = sum((coord*math.prod(buf.shape[i+1:]) for i,coord in enumerate(coords)), UOp.const(0))
|
||||
return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes))).load(dtype=ptr.dtype)
|
||||
return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes))).load()
|
||||
|
||||
def _load_byte(raw:UOp, base:UOp, offset:UOp) -> UOp: return (raw[base + offset//4] >> ((offset&3)*8).cast(dtypes.uint32)) & 255
|
||||
def _half(value:UOp) -> UOp: return value.cast(dtypes.uint16).bitcast(dtypes.float16).float()
|
||||
|
||||
@@ -266,7 +266,7 @@ def store_dest_multi(root:UOp, multi:UOp):
|
||||
|
||||
def passthrough_multi(root:UOp, multi:UOp):
|
||||
new_src = (multi.src[0],)+tuple(x.src[0] if x.op is Ops.UNSHARD else x for x in root.src[1:])
|
||||
return UOp(root.op, root.dtype, src=new_src, arg=root.arg).unshard(multi.arg, multi.src[1:])
|
||||
return UOp(root.op, src=new_src, arg=root.arg).unshard(multi.arg, multi.src[1:])
|
||||
|
||||
def rewrite_into_function(call:UOp):
|
||||
if call.arg.precompile: return None
|
||||
|
||||
+15
-16
@@ -193,14 +193,13 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType:
|
||||
|
||||
class UOpMetaClass(type):
|
||||
ucache:dict[tuple, weakref.ReferenceType[UOp]] = {}
|
||||
def __call__(cls, op:Ops, dtype:DType|None=None, src:tuple[UOp,...]=tuple(), arg:Any=None, tag:Any=None,
|
||||
def __call__(cls, op:Ops, src:tuple[UOp,...]=tuple(), arg:Any=None, tag:Any=None,
|
||||
metadata:tuple[Metadata,...]|None=None, _buffer:Buffer|None=None):
|
||||
if dtype is None: dtype = dtype_from_uop(op, src, arg)
|
||||
# TODO: delete this once the dtype field is removed, for now it just re-implements spec.py
|
||||
elif SPEC == 2 and (expected_dtype:=dtype_from_uop(op, src, arg)) != dtype:
|
||||
raise RuntimeError(f"bad dtype {dtype}, expected {expected_dtype} on {op}")
|
||||
if (wret:=UOpMetaClass.ucache.get(key:=(op, dtype, src, arg, tag), None)) is not None and (ret:=wret()) is not None: return ret
|
||||
UOpMetaClass.ucache[key] = weakref.ref(created:=super().__call__(*key))
|
||||
# NOTE: the key must separate nodes of different dtype: a CONST's dtype is the type of its arg, and True == 1 as dict keys
|
||||
if (wret:=UOpMetaClass.ucache.get(key:=(op, src, arg, tag, type(arg)), None)) is not None and (ret:=wret()) is not None: return ret
|
||||
UOpMetaClass.ucache[key] = weakref.ref(created:=super().__call__(op, src, arg, tag))
|
||||
# derive at construction: bottom up, so no recursion, and a bad node fails where it is built
|
||||
created.__dict__["dtype"] = dtype_from_uop(op, src, arg)
|
||||
if metadata is not None: all_metadata[created] = metadata
|
||||
# NOTE: this value is set by pickle when pickling a realized tensor
|
||||
if _buffer is not None:
|
||||
@@ -241,24 +240,25 @@ from tinygrad.mixin.rand import RandMixin
|
||||
@dataclass(eq=False, slots=True)
|
||||
class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
op:Ops
|
||||
dtype:DType = dtypes.void
|
||||
src:tuple[UOp, ...] = tuple()
|
||||
arg:Any = None
|
||||
tag:Any = None
|
||||
@functools.cached_property
|
||||
def dtype(self) -> DType: return dtype_from_uop(self.op, self.src, self.arg)
|
||||
def __del__(self):
|
||||
# NOTE: getattr because this object may be partially constructed (e.g. if __init__ raised, like the BEAM timeout SIGALRM)
|
||||
if Ops is not None and getattr(self, 'op', None) is Ops.BUFFER and (buffer:=buffers.get(self)) is not None: buffer.ref(-1)
|
||||
try: del UOpMetaClass.ucache[(self.op, self.dtype, self.src, self.arg, self.tag)]
|
||||
try: del UOpMetaClass.ucache[(self.op, self.src, self.arg, self.tag, type(self.arg))]
|
||||
except (AttributeError, KeyError): pass
|
||||
def __reduce__(self):
|
||||
args = [self.op, self.dtype, self.src, self.arg, self.tag, self.metadata]
|
||||
args = [self.op, self.src, self.arg, self.tag, self.metadata]
|
||||
if self.op is Ops.BUFFER and self.realized is not None: args.append(self.realized)
|
||||
return UOp, tuple(args)
|
||||
def replace(self, **kwargs) -> UOp:
|
||||
new_args = (kwargs.pop("op", self.op), kwargs.pop("src", self.src), kwargs.pop("arg", self.arg), kwargs.pop("tag", self.tag))
|
||||
assert len(kwargs) == 0, f"unused kwargs in replace {list(kwargs)}"
|
||||
if (self.op, self.src, self.arg, self.tag) == new_args: return self
|
||||
return UOp(new_args[0], src=new_args[1], arg=new_args[2], tag=new_args[3])
|
||||
return UOp(*new_args)
|
||||
def rtag(self, tag=True): return self.replace(tag=tag)
|
||||
@property
|
||||
def val(self):
|
||||
@@ -544,7 +544,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
@recursive_property
|
||||
def trace_num(self):
|
||||
num = next(ucount)
|
||||
uop_fields[num] = (self.op, self.dtype, tuple(s.trace_num for s in self.src), self.arg, self.tag)+((self.metadata,) if TRACEMETA>=2 else ())
|
||||
uop_fields[num] = (self.op, tuple(s.trace_num for s in self.src), self.arg, self.tag)+((self.metadata,) if TRACEMETA>=2 else ())
|
||||
return num
|
||||
|
||||
# *** uop syntactic sugar ***
|
||||
@@ -603,8 +603,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
@property
|
||||
def without_after(self) -> UOp: return self.src[0] if self.op is Ops.AFTER else self
|
||||
def barrier(self, *src:UOp): return UOp(Ops.BARRIER, src=(self,)+src)
|
||||
def ins(self, arg, **kwargs):
|
||||
return UOp(Ops.INS, src=kwargs.pop("src", self.src), arg=(arg, kwargs.pop("dtype", self.dtype)), tag=kwargs.pop("tag", self.tag))
|
||||
def ins(self, arg, **kwargs): return UOp(Ops.INS, kwargs.pop("src", self.src), (arg, kwargs.pop("dtype", self.dtype)), kwargs.pop("tag", self.tag))
|
||||
def contract(self, *rngs:UOp):
|
||||
assert all(x.arg[-1] == AxisType.UPCAST for x in rngs), "all contract ranges must be upcast"
|
||||
return UOp.stack(*[self.substitute(dict(zip(rngs, [r.const_like(i) for r,i in zip(rngs, idx)])))
|
||||
@@ -1689,7 +1688,7 @@ class RewriteContext:
|
||||
else:
|
||||
# rebuild node with rewritten srcs
|
||||
new_src = tuple(self.replace.get(x, x) for x in n.src)
|
||||
new_n = UOp(n.op, src=new_src, arg=n.arg, tag=n.tag) if new_src != n.src else n
|
||||
new_n = UOp(n.op, new_src, n.arg, n.tag) if new_src != n.src else n
|
||||
# top-down: try pm on rebuilt node, use result as-is (no re-traversal)
|
||||
if self.pm is not None and (rewritten:=self.pm_rewrite(new_n)) is not None: new_n = rewritten
|
||||
self.replace[n] = new_n
|
||||
@@ -1748,7 +1747,7 @@ class RewriteContext:
|
||||
continue
|
||||
else:
|
||||
# if srcs changed from rewrites, construct a new UOp with the new srcs
|
||||
new_src_n = UOp(new_n.op, src=new_src, arg=new_n.arg, tag=new_n.tag)
|
||||
new_src_n = UOp(new_n.op, new_src, new_n.arg, new_n.tag)
|
||||
# trigger a rewrite of new_src_n, then after that rewrite is done, link it back to n
|
||||
stack.append((n, 2, new_src_n))
|
||||
stack.append((new_src_n, 0, new_src_n))
|
||||
|
||||
@@ -11,7 +11,7 @@ def pretty_print(x:UOp, cache=None, d=0)->str:
|
||||
if cache is None: dfs(x, cache:={})
|
||||
if (cx:=cache.setdefault(x, [0,0,False]))[2]: return f"{' '*d}x{cx[0]}"
|
||||
cx[2], srcs = True, (''.join(f'\n{pretty_print(s, cache, d+2)},' for s in x.src))
|
||||
return f"{' '*d}{f'x{cx[0]}:=' * (cx[1]>1)}{type(x).__name__}({x.op}, {x.dtype}, arg={x.argstr()}{x.tagstr()}, src=({srcs}))"
|
||||
return f"{' '*d}{f'x{cx[0]}:=' * (cx[1]>1)}{type(x).__name__}({x.op}, arg={x.argstr()}{x.tagstr()}, src=({srcs}))"
|
||||
|
||||
# ***** uop helpers *****
|
||||
|
||||
@@ -94,10 +94,6 @@ pm_pyrender_extra = PatternMatcher([
|
||||
(UPat(Ops.RANGE, src=(UPat(Ops.CONST, name="c"),), allow_any_len=True, name="x"), lambda ctx,x,c:
|
||||
"UOp.range("+', '.join([str(c.val)] + [repr(y) for y in x.arg])+
|
||||
(f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+")"),
|
||||
# TODO: index shouldn't mismatch dtype
|
||||
(UPat(Ops.INDEX, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda ctx,x:
|
||||
f"{ctx[x.src[0]]}.index({ctx[x.src[1]]}, "+''.join([f"{ctx[xx]}, " for xx in x.src[2:]])+
|
||||
f"dtype={x.dtype})" if x.src[0].dtype != x.dtype else None),
|
||||
# TODO: movement ops simplify stuff, this can break SPEC=2
|
||||
#(UPat(GroupOp.Movement, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({render_marg(ctx,x)})"),
|
||||
# NOTE: CMPNE doesn't work cause there's no __rne__
|
||||
@@ -117,7 +113,7 @@ pm_pyrender_extra = PatternMatcher([
|
||||
|
||||
# NOTE: you can remove pm_pyrender_extra and it'll still be correct
|
||||
pm_pyrender = pm_pyrender_extra+PatternMatcher([
|
||||
(UPat(GroupOp.All, name="u"), lambda ctx,u: f"UOp({u.op}, {u.dtype}, {srcs(ctx,u.src)}"+(f", {repr(u.arg)})" if u.arg is not None else ")")),
|
||||
(UPat(GroupOp.All, name="u"), lambda ctx,u: f"UOp({u.op}, {srcs(ctx,u.src)}"+(f", {repr(u.arg)})" if u.arg is not None else ")")),
|
||||
])
|
||||
|
||||
def _render_with_splits(lst:list[UOp], pm:PatternMatcher, to_render:set[UOp], split_depth:int=100) -> dict[str, str]:
|
||||
|
||||
@@ -54,7 +54,7 @@ spec_shared = PatternMatcher([
|
||||
(UPat(Ops.NOOP), lambda: True),
|
||||
|
||||
# CONST is everywhere; Invalid is a bool const
|
||||
(UPat(Ops.CONST, src=(), name="x"), lambda x: x.dtype is dtypes.bool if x.is_invalid else type(x.val) is type(x.dtype.const(x.val))),
|
||||
(UPat(Ops.CONST, src=(), name="x"), lambda x: x.is_invalid or type(x.val) is type(x.dtype.const(x.val))),
|
||||
|
||||
# STACK is everywhere too
|
||||
(UPat(Ops.STACK, dtype=dtypes.void, src=()), lambda: True),
|
||||
|
||||
@@ -171,9 +171,9 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]:
|
||||
|
||||
def _reconstruct(data:VizData, a:int, depth:int|None=None):
|
||||
if depth is None and a in data.all_uops: return data.all_uops[a]
|
||||
op, dtype, src, arg, *rest = data.trace.uop_fields[a]
|
||||
if depth is not None and depth <= 0: return UOp(op, dtype, (), arg, *rest)
|
||||
ret = UOp(op, dtype, tuple(_reconstruct(data, s, None if depth is None else depth-1) for s in src), arg, *rest)
|
||||
op, src, arg, *rest = data.trace.uop_fields[a]
|
||||
if depth is not None and depth <= 0: return UOp(op, (), arg, *rest)
|
||||
ret = UOp(op, tuple(_reconstruct(data, s, None if depth is None else depth-1) for s in src), arg, *rest)
|
||||
if depth is None: data.all_uops[a] = ret
|
||||
return ret
|
||||
|
||||
|
||||
Reference in New Issue
Block a user