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

remove the dtype.scalar calls everywhere [pr] (#16930)

Browse Source 
* remove the dtype.scalar calls everywhere

* remove base everywhere

* Revert "remove base everywhere"

This reverts commit 0dd0a5a243.
This commit is contained in:
George Hotz
2026-07-08 10:10:03 -07:00
committed by GitHub
parent e4d0d634d4
commit 6e05cbdbcb
12 changed files with 23 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/null/test_uop_symbolic.py
+2 -2
View File
@@ -1355,10 +1355,10 @@ class TestGatedUopGivenValid(unittest.TestCase):
idx0 = (r0 + uconst(-1)) // uconst(3)
idx1 = r0 % uconst(3)
idx:UOp = (r0 < 3).where(UOp(Ops.STACK, dtypes.weakint.vec(2), (idx0, idx1)), UOp.invalid())
idx:UOp = (r0 < 3).where(UOp(Ops.STACK, dtypes.weakint, (idx0, idx1)), UOp.invalid())
idx = graph_rewrite(idx, pm_simplify_valid)
# independent simplification: (r0-1)//3 -> (r0+2)//3 - 1, and r0%3 -> r0 when r0 in [0,2]
expected_vec = UOp(Ops.STACK, dtypes.weakint.vec(2), ((r0 + uconst(2)) // uconst(3) + uconst(-1), r0))
expected_vec = UOp(Ops.STACK, dtypes.weakint, ((r0 + uconst(2)) // uconst(3) + uconst(-1), r0))
self.assertEqual(idx, (r0 < 3).where(expected_vec, UOp.invalid()))
class TestRangeSplitting(unittest.TestCase):
tinygrad/codegen/__init__.py
+1 -1
View File
@@ -224,7 +224,7 @@ def reduce_ranges_to_acc(ctx:ReduceContext, r:UOp):
topo = r.src[0].toposort()
ended_ranges = flatten([x.ended_ranges for x in topo if x.op is Ops.END])
input_ranges = tuple(x for x in topo if x.op is Ops.RANGE and x not in r.src[1:] and x not in ended_ranges)
acc_init = acc.after(*input_ranges).store(identity_element(r.arg[0], r.dtype.scalar()))
acc_init = acc.after(*input_ranges).store(identity_element(r.arg[0], r.dtype))
acc_initted = acc.after(acc_init, *r.src[1:])
inp = r.src[0].reduce(arg=r.arg) if r.arg[1] else r.src[0]
acc_out = acc_initted.store(acc_initted.alu(r.arg[0], inp)).end(*r.src[1:]).rtag("mergeable")
tinygrad/codegen/opt/postrange.py
+1 -1
View File
@@ -228,7 +228,7 @@ class Scheduler:
raise KernelOptError(f"invalid tensor core choice {tc_select}")
for tc in tensor_cores:
if self.ren.target.device in ("CUDA", "NV") and tc.dtype_in == dtypes.float and not ALLOW_TF32: continue
if tc.dtype_in == in0.dtype.scalar() and tc.dtype_in == in1.dtype.scalar() and tc.dtype_out == reduceop.dtype.scalar():
if tc.dtype_in == in0.dtype and tc.dtype_in == in1.dtype and tc.dtype_out == reduceop.dtype:
# tensor cores have three ranges. X, Y, and REDUCE
in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: x.arg[0], reverse=True)
in1_ranges = sorted([u for u in in1.ranges if u not in in0.ranges], key=lambda x: x.arg[0], reverse=True)
tinygrad/mixin/elementwise.py
+1 -1
View File
@@ -65,7 +65,7 @@ class ElementwiseMixin(CreationMixin):
print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).neg().numpy())
```
"""
return self.logical_not() if self.dtype.scalar() == dtypes.bool else self * (-1)
return self.logical_not() if self.dtype == dtypes.bool else self * (-1)
def _check_dtype(self) -> None:
if not (dtypes.is_bool(self.dtype) or dtypes.is_int(self.dtype)):
tinygrad/runtime/ops_python.py
+4 -4
View File
@@ -81,13 +81,13 @@ class PythonProgram:
store_gate = exec_masks[-1]
for j,val in enumerate(src_values[1] if u.max_numel() > 1 else [src_values[1]]):
for (m,o),v,g in zip(src_values[0], val, store_gate):
if g: _store(m, o+j, v, src_dtypes[1].scalar())
if g: _store(m, o+j, v, src_dtypes[1])
i += 1
continue
if u.op is Ops.AFTER: values[u] = src_values[0]
elif u.op is Ops.PARAM and u.addrspace is AddrSpace.ALU: values[u] = [pvals.pop(0)] * warp_size
elif u.op in {Ops.PARAM, Ops.BUFFER}:
storage_fmt = storage_fmt_for_dtype(u.dtype.base.scalar())
storage_fmt = storage_fmt_for_dtype(u.dtype.base)
if storage_fmt is None: raise RuntimeError(f"dtype={u.dtype} is not supported")
if TYPE_CHECKING or sys.version_info < (3, 12): assert storage_fmt != "e"
if u.addrspace == AddrSpace.REG:
@@ -129,13 +129,13 @@ class PythonProgram:
if (load_sz := u.max_numel()) > 1:
# buf and gate are not vecs
values[u] = [load([src_values[k] if k in [0,2] else src_values[k][j] \
for k in range(len(src_values))], j, u.dtype.scalar()) for j in range(load_sz)]
for k in range(len(src_values))], j, u.dtype) for j in range(load_sz)]
else:
values[u] = load(src_values, 0, u.dtype)
elif u.op is Ops.WMMA:
first_src_dtype = u.src[0].dtype
assert isinstance(first_src_dtype, DType) # mypy
dims, dtype_in, device, threads = u.arg[1], first_src_dtype.scalar(), u.arg[4], u.arg[5]
dims, dtype_in, device, threads = u.arg[1], first_src_dtype, u.arg[4], u.arg[5]
wmma_helper = functools.partial(generic_wmma_helper, src_values, warp_size)
# TODO: refactor these to a shared TensorCoreLayout
if device == "METAL":
tinygrad/schedule/indexing.py
+1 -1
View File
@@ -171,7 +171,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
# treat MSTACK/MSELECT like SINK
if x.op in {Ops.MSTACK, Ops.MSELECT}: continue
if x.dtype.scalar() == dtypes.weakint: continue # TODO: why do I need this?
if x.dtype == dtypes.weakint: continue # TODO: why do I need this?
ending_ranges[x] = sum([ending_ranges.get(u, []) for u in consumer_map[x]], [])
# *** the ranges on the output are
tinygrad/schedule/multi.py
+1 -1
View File
@@ -70,7 +70,7 @@ def reduce_multi(root:UOp, multi:UOp):
if multi.axis is not None and multi.axis < num_axes:
local = multi.src[0]._rop(op, tuple(range(num_axes)))
# allreduce in pre-cast dtype when sum_acc_dtype promoted from bf16/half
if ALLREDUCE_CAST and multi.src[0].op is Ops.CAST and multi.src[0].src[0].dtype.scalar() in (dtypes.bfloat16, dtypes.half):
if ALLREDUCE_CAST and multi.src[0].op is Ops.CAST and multi.src[0].src[0].dtype in (dtypes.bfloat16, dtypes.half):
orig_dtype = multi.src[0].src[0].dtype
return local.cast(orig_dtype).allreduce(op, multi.device).cast(local.dtype)
return local.allreduce(op, multi.device)
tinygrad/uop/ops.py
+5 -7
View File
@@ -452,7 +452,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
def index(self, *srcs:UOp|int|None, **kwargs):
new_srcs: list[UOp] = [UOp.const(dtypes.weakint, x) if isinstance(x, int) else x for x in srcs if x is not None]
if len(new_srcs) == 1 and new_srcs[0].op is Ops.CONST and self.op is Ops.STACK: return self.src[new_srcs[0].arg]
return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype.base.scalar()), (self,)+tuple(new_srcs), **kwargs)
return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype.base), (self,)+tuple(new_srcs), **kwargs)
def __getitem__(self, idx):
# buffers index into INDEX UOps (scalar lookup); everything else uses the shared mixin view path
if self.addrspace in (None, AddrSpace.ALU) or self.device is not None: return super(UOp, self).__getitem__(idx)
@@ -485,7 +485,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
if dtypes.is_float(self.dtype) or (dtypes.is_int(self.dtype) and isinstance(x, (int, InvalidType))): return self.const_like(x)
return self.const_like(x, dtypes.from_py(x))
def broadcast(self, count:int):
assert self.dtype.vcount == 1
if count == 1: return self
return UOp(Ops.STACK, self.dtype, (self,)*count)
def cast(self, dtype:DTypeLike):
@@ -528,7 +527,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
@staticmethod
def const(dtype:DType, b:ConstLike, shape:tuple[sint, ...]|None=None):
if isinstance(b, UOp): return b.cast(dtype)
assert dtype.vcount == 1, f"const dtype must be scalar, got {dtype}"
# NOTE: it always has to be STACK now, even if they are all the same
if isinstance(b, tuple):
stk = [UOp(Ops.CONST, dtype, arg=dtype.const(c), src=()) for c in b]
@@ -1147,16 +1145,16 @@ python_alu: dict[Ops, Callable] = {
def exec_alu(op:Ops, dtype:DType, operands, truncate_output=True):
if any(isinstance(x, tuple) for x in operands):
count = max(len(x) for x in operands if isinstance(x, tuple))
return tuple([exec_alu(op, dtype.scalar(), [x[i] if isinstance(x, tuple) else x for x in operands]) for i in range(count)])
return tuple([exec_alu(op, dtype, [x[i] if isinstance(x, tuple) else x for x in operands]) for i in range(count)])
if dtype==dtypes.weakint and op in GroupOp.Binary and Invalid in operands: return Invalid
alu = python_alu[op](*operands)
return truncate.get(dtype, lambda x: x)(alu) if truncate_output else alu
def bitcast(x, in_dtype:DType, out_dtype:DType):
assert in_dtype.itemsize == out_dtype.itemsize, "bitcast itemsize mismatch"
packed = struct.pack(storage_fmt_for_dtype(in_dtype.scalar()), to_storage_scalar(x, in_dtype.scalar()))
out_val = struct.unpack(storage_fmt_for_dtype(out_dtype.scalar()), packed)[0]
return from_storage_scalar(out_val, out_dtype.scalar())
packed = struct.pack(storage_fmt_for_dtype(in_dtype), to_storage_scalar(x, in_dtype))
out_val = struct.unpack(storage_fmt_for_dtype(out_dtype), packed)[0]
return from_storage_scalar(out_val, out_dtype)
# ***** pattern matcher *****
tinygrad/uop/render.py
+1 -1
View File
@@ -92,7 +92,7 @@ pm_pyrender_extra = PatternMatcher([
# 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.base.scalar() != x.dtype else None),
f"dtype={x.dtype})" if x.src[0].dtype.base != 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__
tinygrad/uop/spec.py
+1 -1
View File
@@ -129,7 +129,7 @@ def valid_gettuple(g:UOp, t:UOp):
spec_tensor = PatternMatcher([
# BUFFER
(UPat(Ops.BUFFER, src=(UPat(),), name="buf"), lambda buf:
(isinstance(buf.dtype, DType) and buf.src[0].dtype.scalar() == dtypes.weakint and is_device(buf.arg.device))
(isinstance(buf.dtype, DType) and buf.src[0].dtype == dtypes.weakint and is_device(buf.arg.device))
if isinstance(buf.arg, ParamArg) and buf.addrspace is AddrSpace.GLOBAL else None),
# Tensor variable bindings
tinygrad/uop/symbolic.py
+4 -4
View File
@@ -20,7 +20,7 @@ def simplify_pow(x:UOp, c:UOp) -> UOp|None:
return None
def fold_bitcast(root:UOp, c:UOp) -> UOp|None:
if (from_fmt:=c.dtype.scalar().fmt) is None or (to_fmt:=root.dtype.scalar().fmt) is None: return None
if (from_fmt:=c.dtype.fmt) is None or (to_fmt:=root.dtype.fmt) is None: return None
if c.dtype.itemsize != root.dtype.itemsize: return None
def convert(v:ConstType) -> ConstType: return struct.unpack(to_fmt, struct.pack(from_fmt, v))[0]
return root.const_like(convert(c.arg))
@@ -101,7 +101,7 @@ propagate_invalid = pm_index_invalid + pm_data_invalid
# NOTE: this happens in padded WMMA, so rewrite to 0
pm_remove_invalid = PatternMatcher([
(invalid_pat, lambda i: i.const_like(0) if i.dtype.scalar() is not dtypes.weakint else None),
(invalid_pat, lambda i: i.const_like(0) if i.dtype is not dtypes.weakint else None),
])
symbolic_simple = propagate_invalid + PatternMatcher([
@@ -397,7 +397,7 @@ pm_move_where_on_load = PatternMatcher([
])
def gated_given_valid(cond:UOp, x:UOp, i:UOp) -> UOp|None:
if x.dtype.scalar() is not dtypes.weakint: return None
if x.dtype is not dtypes.weakint: return None
# Skip if x contains DIV/MOD AND IMAGE mode is enabled -> image index e.g. openpilot
if IMAGE.value > 0 and x.op_in_backward_slice_with_self(Ops.CDIV, Ops.CMOD, Ops.FLOORDIV, Ops.FLOORMOD): return None
return cond.where(uop_given_valid(cond, x, try_simplex=False), i)
@@ -430,7 +430,7 @@ pm_clean_up_group_sink = PatternMatcher([
sym = symbolic+pm_simplify_valid+PatternMatcher([
# reorder ALU/VECTORIZE
(UPat(GroupOp.ALU, src=(UPat(Ops.STACK, src=UPat(name='x')), UPat(Ops.STACK, src=UPat(name='y'))), name='alu'),
lambda x,y,alu: UOp(Ops.STACK, alu.dtype, (UOp(alu.op, alu.dtype.scalar(), (x,y)),))),
lambda x,y,alu: UOp(Ops.STACK, alu.dtype, (UOp(alu.op, alu.dtype, (x,y)),))),
# ** where **
# # fold nested where with same condition: in cond.where(t,f), cond.where(a,b)->a in t, ->b in f
# (UPat.var("cond").where(UPat.var("t"), UPat.var("f")), fold_where_closure),
tinygrad/uop/validate.py
+1 -1
View File
@@ -53,7 +53,7 @@ z3_renderer = PatternMatcher([
def uops_to_z3(solver:z3.Solver, *uops: UOp) -> list[z3.ExprRef]:
# gate on upstream AFTER/BUFFER, but keep INDEX as an unknown LOAD
lst = list(UOp.sink(*uops).toposort(gate=lambda x: x.op not in {Ops.AFTER, Ops.BUFFER} and \
(x.dtype.scalar() in dtypes.ints+(dtypes.bool, dtypes.weakint) or x.op is Ops.SINK)))[:-1]
(x.dtype in dtypes.ints+(dtypes.bool, dtypes.weakint) or x.op is Ops.SINK)))[:-1]
z3map: dict[UOp, z3.ExprRef] = {}
for u in lst:
# NOTE: we skip STACK here, it can't actually be accessed