diff --git a/tinygrad/callify.py b/tinygrad/callify.py index 7280ab4aad..4a5211d4f2 100644 --- a/tinygrad/callify.py +++ b/tinygrad/callify.py @@ -61,7 +61,7 @@ def _make_buffer_view(src:UOp) -> UOp|None: buf = buf.src[0] if byte_offset % buf.dtype.itemsize != 0: return None offset = byte_offset // buf.dtype.itemsize - return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(dtypes.weakint, offset)), src.numel()) + return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(None, offset)), src.numel()) def contiguous_mops_to_view(c:UOp, src:UOp): """MOPS(BUFFER) → SLICE when movement ops collapse to a contiguous range.""" diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index ac3dba750d..f690245bba 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -125,7 +125,7 @@ def do_devectorize(b:UOp): if not all(x.shape == b.shape or x.base.arg is Invalid for x in b.src): return None src = [] for idx in itertools.product(*[range(x) for x in b.shape]): - idx_c = [UOp.const(dtypes.weakint, i) for i in idx] + idx_c = [UOp.const(None, i) for i in idx] src.append(b.replace(src=tuple(x.base if x.base.arg is Invalid else x.index(*idx_c) for x in b.src))) return UOp.stack(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src) @@ -135,7 +135,7 @@ def do_stack_wmma(u:UOp): src = [] for b in u.src: if b.op != Ops.STACK: - src.append(UOp.stack(*[b.index(UOp.const(dtypes.weakint, i)) for i in range(b.max_numel())])) + src.append(UOp.stack(*[b.index(UOp.const(None, i)) for i in range(b.max_numel())])) else: src.append(b) return u.replace(src=tuple(src)) @@ -161,7 +161,7 @@ devectorizer2 = mop_cleanup+pm_mops+PatternMatcher([ # RESHAPE a void is removed (hack for AFTER) (UPat(Ops.RESHAPE, dtype=dtypes.void, name="x"), lambda x: x.src[0]), # reshape of a single element shaped value to scalar is an index - (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(UOp.const(dtypes.weakint, 0)) if x.marg == () and x.src[0].shape == (1,) else None), + (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(UOp.const(None, 0)) if x.marg == () and x.src[0].shape == (1,) else None), # EXPAND on scalar -> STACK (UPat(Ops.EXPAND, src=(UPat.var("x"), UPat()), name="out"), lambda x,out: UOp.stack(*([x]*out.max_numel())) if x.shape == () and out.shape == (out.max_numel(),) else None), diff --git a/tinygrad/codegen/late/coalesce.py b/tinygrad/codegen/late/coalesce.py index 24429c44d6..8830c643b4 100644 --- a/tinygrad/codegen/late/coalesce.py +++ b/tinygrad/codegen/late/coalesce.py @@ -141,12 +141,12 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp: grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])] for full_grp in grouped_offsets: while len(full_grp): - offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(dtypes.weakint, full_grp[0]) + offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(None, full_grp[0]) length = [l for l in lengths if l <= len(full_grp) and (not must_divide or offset.divides(l) is not None)][0] grp = full_grp[:length] # NOTE: we apply the valid again after we determine the length offset = offset.valid(valid) if valid is not None else offset - idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(dtypes.weakint, len(grp)))) if len(grp) > 1 else buf.index(offset) + idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(None, len(grp)))) if len(grp) > 1 else buf.index(offset) if op == Ops.STORE: datas = [] for i,g in enumerate(grp): @@ -158,7 +158,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp: ld = idx.load() for i,g in enumerate(grp): for oo in offsets[g]: - replacements[oo] = ld.index(UOp.const(dtypes.weakint, i)) if len(grp) > 1 else ld + replacements[oo] = ld.index(UOp.const(None, i)) if len(grp) > 1 else ld full_grp = full_grp[length:] # apply diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index c5612df1e8..77bd5a1448 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -348,7 +348,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(dtypes.weakint, offs.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases} + subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(None, offs.get(b, 0))), 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)]) diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index 0a574529b4..1ef2d8fee0 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -53,7 +53,7 @@ class IndexingContext: def new_range(self, s:sint, axistype:AxisType=AxisType.LOOP) -> UOp: if isinstance(s, UOp) and s.op is Ops.RANGE: return s # if a range has a 1 src, it's the same as UOp.const(dtypes.weakint, 0) - return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(dtypes.weakint, 0) + return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(None, 0) def broadcast_rngs(x:UOp, src:UOp, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]: if x.op not in GroupOp.Broadcastable: return rngs @@ -100,7 +100,7 @@ def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp): def convert_pad_to_where_to_keep_behavior_local(ctx:IndexingContext, x:UOp): if x not in ctx.range_map: return None bx = create_bufferize_and_index_based_on_ranges(ctx, x) - valid: UOp = UOp.const(dtypes.bool, True).uprod([r.get_valid() for r in ctx.range_map[x][0]]) + valid: UOp = UOp.const(None, True).uprod([r.get_valid() for r in ctx.range_map[x][0]]) return valid.where(bx.src[0], UOp.const(x.dtype, 0)) def convert_reduce_to_reduce_with_ranges(ctx:IndexingContext, x:UOp): @@ -148,7 +148,7 @@ def _apply_reshape(in_shape:tuple[sint,...], out_shape:tuple[sint, ...], urngs:U for s,src in list(zip(out_shape, urngs.src))[::-1]: axes_in.append(acc*src) acc *= s - combined_axes = UOp.const(dtypes.weakint, 0).usum(axes_in) + combined_axes = UOp.const(None, 0).usum(axes_in) axes_out:list[UOp] = [] for s in in_shape[::-1]: axes_out.append(combined_axes % s) @@ -248,7 +248,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: # we compare the ranges without their valids if all_all_same or (PCONTIG and all_same(local_rngs)): # the new valid is the OR of all the children valids - minimum_valid = UOp.const(dtypes.bool, False).usum(valids) + minimum_valid = UOp.const(None, False).usum(valids) _out_rngs.append(graph_rewrite(local_rngs[0].valid(minimum_valid), symbolic, name="minimum_valid")) else: _out_rngs.append(rctx.new_range(x.shape[i])) diff --git a/tinygrad/schedule/memory.py b/tinygrad/schedule/memory.py index f88f60675a..3f163617d9 100644 --- a/tinygrad/schedule/memory.py +++ b/tinygrad/schedule/memory.py @@ -56,7 +56,7 @@ def memory_plan_rewrite(linear:UOp, held_bufs:set[UOp]|None=None) -> UOp: 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(dtypes.weakint, offset)), buf_uop.max_numel()) + replace_map[buf_uop] = UOp(Ops.SLICE, buf_uop.dtype, (arenas[_key(buf_uop)], UOp.const(None, offset)), buf_uop.max_numel()) 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") diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 250423abf3..d019326401 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -78,7 +78,7 @@ def split_reduceop(reduce:UOp, x:UOp): # split is moved to the end to provide maximum locality for the second phase reduce. # get expanded by rangeifying the UOp x - indexed = x.index(*[UOp.range(s, i) if resolve(s>1) else UOp.const(dtypes.weakint, 0) for i,s in enumerate(x.shape)]) + indexed = x.index(*[UOp.range(s, i) if resolve(s>1) else UOp.const(None, 0) for i,s in enumerate(x.shape)]) range_nums = [y.arg[0] for y in indexed.substitute({x.base:UOp(Ops.NOOP, x.base.dtype)}, extra_pm=pm_mops).ranges] is_expanded = [i not in range_nums for i in range(len(x.shape))] diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index c1675d2199..bd264c4c95 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -71,9 +71,9 @@ class Tensor(RandMixin): # create a UOp from the different types of inputs if data is None: - data = UOp.const(_dtype or dtypes.weakfloat, 0.0) + data = UOp.const(_dtype, 0.0) elif isinstance(data, get_args(ConstType)): - data = UOp.const(_dtype or dtypes.from_py(data), data) + data = UOp.const(_dtype, data) elif is_numpy_ndarray(data) and data.shape == (): data = UOp.const(_dtype or _from_np_dtype(data.dtype), data.item()) elif not isinstance(data, UOp): diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 6b48d94b07..6a506f7b6b 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -602,7 +602,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass): for idx in itertools.product(*[range(int(r.vmax)+1) for r in rngs])]) def alu(self, op, *src:UOp, **kwargs): return UOp(op, src=(self, *src), **kwargs) @staticmethod - def const(dtype:DType, b:ConstLike, shape:tuple[sint, ...]|None=None): + def const(dtype:DType|None, b:ConstLike, shape:tuple[sint, ...]|None=None): + if dtype is None: dtype = dtypes.from_py(b) if isinstance(b, UOp): return b.cast(dtype) # NOTE: it always has to be STACK now, even if they are all the same if isinstance(b, tuple): @@ -633,7 +634,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): ret = UOp(Ops.REDUCE, src=(self.permute(perm),), arg=(op, len(reduce_axis))) return ret.reshape(tuple(s for i,s in enumerate(self.shape) if i not in axis)) if axis != reduce_axis else ret @staticmethod - def invalid(): return UOp.const(dtypes.bool, Invalid) + def invalid(): return UOp.const(None, Invalid) def valid(self, cond): return cond.where(self, self.const_like(Invalid)) def get_idx(self) -> UOp: @@ -641,7 +642,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return self.src[1] if self.op is Ops.WHERE and self.src[2].arg is Invalid else self def get_valid(self) -> UOp: if self.op is Ops.STACK: return UOp.stack(*(x.get_valid() for x in self.src)) - return self.src[0] if self.op is Ops.WHERE and self.src[2].arg is Invalid else UOp.const(dtypes.bool, self.arg is not Invalid) + return self.src[0] if self.op is Ops.WHERE and self.src[2].arg is Invalid else UOp.const(None, self.arg is not Invalid) def reduce(self, *src:UOp, **kwargs): arg = kwargs.pop('arg', None) if isinstance(arg, Ops): arg = (arg, 0) @@ -1774,8 +1775,8 @@ pm_unbind = PatternMatcher([(UPat(Ops.BIND, name="x"), do_unbind)]) # ctx is source UOp for which we are finding a contiguous view for. used in contiguous_view_offset pm_contiguous_view_offset = PatternMatcher([ - (UPat(Ops.INDEX, src=(UPat(),)), lambda: UOp.const(dtypes.weakint, 0)), - (UPat(Ops.INDEX, src=(UPat(), UPat(Ops.RANGE))), lambda: UOp.const(dtypes.weakint, 0)), + (UPat(Ops.INDEX, src=(UPat(),)), lambda: UOp.const(None, 0)), + (UPat(Ops.INDEX, src=(UPat(), UPat(Ops.RANGE))), lambda: UOp.const(None, 0)), (UPat(Ops.INDEX, src=(UPat(), UPat(Ops.RANGE)+UPat.cvar('c'))), lambda c: c), (UPat(Ops.INDEX, src=(UPat(), UPat.cvar('c'))), lambda ctx, c: c if resolve(ctx.numel() == 1, False) else None), ]) diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index c1f6897db9..49baafcac2 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -11,7 +11,7 @@ def validate_index(uidx:UOp, gate:UOp|None=None): if len(uidx.src) != 2: return True # skip for non final index. TODO: check more complex index with shape buf,idx = uidx.src if idx.op is Ops.CONST and idx.arg is Invalid: return True - if gate is None: gate = UOp.const(dtypes.bool, True) + if gate is None: gate = UOp.const(None, True) # TODO: check for overflow if not CHECK_OOB or is_image_shape(buf._shape): return True diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 6ab28c4b71..1d9a7c658d 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -383,7 +383,7 @@ def reduce_mul_chain(r:UOp) -> UOp|None: def drop_and_clauses(cond:UOp, x:UOp, i:UOp) -> UOp|None: keep, drop = partition(cond.split_uop(Ops.AND), lambda c: any(r in x.ranges for r in c.ranges)) - return UOp.const(dtypes.bool, True).uprod(*keep).where(x, i) if drop else None + return UOp.const(None, True).uprod(*keep).where(x, i) if drop else None pm_drop_and_clauses = PatternMatcher([(invalid_gate, drop_and_clauses)]) # move conditions from where to load's valid, drop clauses already in load @@ -398,7 +398,7 @@ def where_on_load(cond:UOp, buf:UOp, idx:UOp, or_cast:UOp) -> UOp|None: if len(keep) == len(where_clauses): return None idx = buf.index(idx.get_idx().valid(load_valid.uprod(*moved))) ret_idx = idx.cast(or_cast.dtype) if or_cast.op is Ops.CAST else idx - return UOp.const(dtypes.bool, True).uprod(*keep).where(ret_idx, ret_idx.const_like(0)) + return UOp.const(None, True).uprod(*keep).where(ret_idx, ret_idx.const_like(0)) # where after gated load becomes alt value, TODO: this is sort of duplicated with rules in devectorizer pm_move_where_on_load = PatternMatcher([