nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

remove dtypes base (#16931)

Browse Source 
* remove dtypes base

* find/replace bug

* vcount is junk too
This commit is contained in:
George Hotz
2026-07-08 11:48:11 -07:00
committed by GitHub
parent 63cb1369cb
commit fdffc6c0c8
34 changed files with 92 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files
AGENTS.md
+5
View File
@@ -0,0 +1,5 @@
# Notes
- Run tests with `-n12` for speed (e.g. `python -m pytest test/null/test_dtype.py -x -q -n12`)
- Run `python -m mypy tinygrad/` to typecheck
- Run `python -m ruff check .` to lint
examples/mlperf/model_train.py
+1 -1
View File
@@ -2006,7 +2006,7 @@ def train_stable_diffusion():
# move to CPU first so more GPU bufs aren't created (can trigger OOM)
for k,v in ckpt.items(): ckpt[k] = v.detach().to("CPU")
Tensor.realize(*[v for v in ckpt.values()])
for k,v in ckpt.items(): ckpt[k] = v.cast(v.dtype.base).contiguous()
for k,v in ckpt.items(): ckpt[k] = v.cast(v.dtype).contiguous()
Tensor.realize(*[v for v in ckpt.values()])
return ckpt
extra/gemm/amd_copy_matmul.py
+2 -2
View File
@@ -46,8 +46,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
# -- GLOBAL -> LOCAL --
# wmma: spatial outer, k inner (k contiguous for vectorized WMMA tile loads)
# gemm: k outer, spatial inner
A_local = UOp.placeholder((BLOCK_M, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_M), a.dtype.base, slot=0, addrspace=AddrSpace.LOCAL)
B_local = UOp.placeholder((BLOCK_N, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_N), b.dtype.base, slot=1, addrspace=AddrSpace.LOCAL)
A_local = UOp.placeholder((BLOCK_M, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_M), a.dtype, slot=0, addrspace=AddrSpace.LOCAL)
B_local = UOp.placeholder((BLOCK_N, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_N), b.dtype, slot=1, addrspace=AddrSpace.LOCAL)
a = a.reshape(K // BLOCK_K, BLOCK_K, BLOCK_M)
b = b.reshape(K // BLOCK_K, BLOCK_K, BLOCK_N)
extra/gemm/cdna_asm_gemm.py
+1 -1
View File
@@ -118,7 +118,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(dtypes.weakint, K)+k))*
B.flatten().index((k*UOp.const(dtypes.weakint, N)+n))).cast(dtypes.float32)
red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype.base)
red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype)
store = C.flatten().index((m*UOp.const(dtypes.weakint, N)+n)).store(red).end(m, n)
return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
extra/hcq2/hcq2.py
+6 -6
View File
@@ -27,7 +27,7 @@ class HCQ2Compiled(Compiled):
(UPat(Ops.PARAM, tag="timeline_value"), lambda ctx: ctx.timeline_value()),
(UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx.timeline_signal("sentinel", (1 << 64) - 1)),
(UPat(Ops.PARAM, name="b"), lambda ctx, b:
Buffer(ctx.device, b.max_numel(), b.dtype.base, options=BufferSpec(host=False, uncached=True, cpu_access=True, nolru=True))
Buffer(ctx.device, b.max_numel(), b.dtype, options=BufferSpec(host=False, uncached=True, cpu_access=True, nolru=True))
if b.tag is not None else None), # TODO: remove nolru
])
@@ -152,7 +152,7 @@ def make_placeholder(devs, size:int, dtype, name=None, unique=True) -> UOp:
return UOp.param(next(UOp.unique_num) if unique else 0, dtype, shape=(size,), device=devs).rtag(name or "buf")
def make_patch(buf:UOp, off:sint, val:UOp, dtype=None) -> UOp:
return buf.index(UOp.const(dtypes.int, off//buf.dtype.base.itemsize)).store(val.cast(dtype or buf.dtype.base))
return buf.index(UOp.const(dtypes.int, off//buf.dtype.itemsize)).store(val.cast(dtype or buf.dtype))
def make_cmdbuf(lin, devs):
blob, patches = b'', []
@@ -199,7 +199,7 @@ def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS) and not all_d
def stage_copy(dst:UOp, src:UOp) -> UOp|None:
if not (_need_staging(src, dst) or _need_staging(dst, src)): return None
stage = UOp.new_buffer("CPU", src.max_numel() * src.dtype.base.itemsize, dtypes.uint8)
stage = UOp.new_buffer("CPU", src.max_numel() * src.dtype.itemsize, dtypes.uint8)
return UOp(Ops.LINEAR, dtypes.void, (src.copy_to_device("CPU").call(stage, src), stage.copy_to_device(dst.device).call(dst, stage)))
pm_insert_copy_staging = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy)])
@@ -232,7 +232,7 @@ pm_tag_hcq_calls = PatternMatcher([(UPat(Ops.LINEAR, name="linear"),
class HCQDepsTracker(DepsTracker):
@staticmethod
def _key(buf:Any) -> tuple[Any, int, int]:
return (buf.arg.slot, 0, buf.max_numel() * buf.dtype.base.itemsize) if isinstance(buf, UOp) else DepsTracker._key(buf)
return (buf.arg.slot, 0, buf.max_numel() * buf.dtype.itemsize) if isinstance(buf, UOp) else DepsTracker._key(buf)
def make_deps(u:UOp, dep_lanes:list[tuple[UOp, int, int]], nlanes:int) -> UOp:
deps:dict[UOp, list[int|None]] = collections.defaultdict(lambda: [None]*nlanes)
@@ -414,7 +414,7 @@ def _make_getaddrs_sub(call:UOp, gaddrs:list[UOp], name:str):
b = make_placeholder(call.arg.aux.device, len(order), dtypes.uint64, name)
sub = {g: b.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(dtypes.int, order.index(gr))).load() for g,gr in bare.items()}
return sub, (b.after(*[make_patch(b, i * b.dtype.base.itemsize, gr) for i,gr in enumerate(order)]),) if order else ()
return sub, (b.after(*[make_patch(b, i * b.dtype.itemsize, gr) for i,gr in enumerate(order)]),) if order else ()
def rm_rt_getaddrs(call:UOp) -> UOp|None:
if not (gaddrs:=[u for u in call.src[0].toposort() if u.op is Ops.GETADDR]): return None
@@ -534,7 +534,7 @@ def fold_blob_store(buf:UOp, blob:UOp) -> UOp:
def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp:
for b, v in zip((bs:=mb.bufs if isinstance((mb:=buf.buffer), MultiBuffer) else (mb,)), val.src if val.op is Ops.STACK else (val,)*len(bs)):
struct.pack_into(f'<{v.dtype.fmt}', b.ensure_allocated()._buf.cpu_view().mv.cast('B'), off.arg * buf.dtype.base.itemsize, truncate[v.dtype](v.arg))
struct.pack_into(f'<{v.dtype.fmt}', b.ensure_allocated()._buf.cpu_view().mv.cast('B'), off.arg * buf.dtype.itemsize, truncate[v.dtype](v.arg))
return UOp(Ops.NOOP)
def resolve_getaddr(buf:UOp, g:UOp) -> UOp:
extra/hcq2/ops_amd2.py
+1 -1
View File
@@ -176,7 +176,7 @@ class SDMAOps(FastEnum): COPY = auto(); POLL_REGMEM = auto(); FENCE = auto(); TR
def sdma_copy(ctx, call):
dst, src = call.src[1], call.src[2]
sz = src.max_numel() * src.dtype.base.itemsize
sz = src.max_numel() * src.dtype.itemsize
src_addr, dst_addr = make_getaddr(src, ctx.devs), make_getaddr(dst, ctx.devs)
return UOp(Ops.LINEAR, dtypes.void, tuple([make_ins(SDMAOps.COPY,
ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR),
extra/hip_large_kernel.py
+1 -1
View File
@@ -18,7 +18,7 @@ prg = dev.runtime("write_ones", mbin)
prg(buf0._buf, global_size=(1,65537,1), local_size=(1,1,1), wait=True)
import numpy as np
def to_np(buf): return np.frombuffer(buf.as_memoryview().cast(buf.dtype.base.fmt), dtype=_to_np_dtype(buf.dtype.base))
def to_np(buf): return np.frombuffer(buf.as_memoryview().cast(buf.dtype.fmt), dtype=_to_np_dtype(buf.dtype))
big = to_np(buf0)
print(big)
extra/llama_kernels/fused_ce/__init__.py
+1 -1
View File
@@ -36,7 +36,7 @@ def _custom_fused_ce_loss_bwd(d_logits:UOp, logits:UOp, lse:UOp, targets:UOp, sc
smooth = label_smoothing / vocab
grad = (prob - target - smooth) * scale[0]
return d_logits[b, s, v].store(grad.cast(d_logits.dtype.base)).end(v, row).sink(arg=KernelInfo(f"fused_ce_loss_bwd_{rows}_{vocab}"))
return d_logits[b, s, v].store(grad.cast(d_logits.dtype)).end(v, row).sink(arg=KernelInfo(f"fused_ce_loss_bwd_{rows}_{vocab}"))
def _fused_ce_loss_bwd(gradient:UOp, kernel:UOp, label_smoothing:float):
# NOTE: forward inputs are (loss_out, max_out, lse_out, logits, targets)
extra/llama_kernels/fused_silu_mul_quantize_mxfp8/__init__.py
+3 -3
View File
@@ -41,7 +41,7 @@ def _custom_silu_mul_quantize_mxfp8(fp8_out:UOp, e8_out:UOp, si_out:UOp, x_w1:UO
scaled = (act * qscale).maximum(-FP8_MAX).minimum(FP8_MAX)
e8u8 = e8f.cast(dtypes.uint8)
fp8_store = fp8_out[idx].store(scaled.cast(fp8_out.dtype.base)).end(lane)
fp8_store = fp8_out[idx].store(scaled.cast(fp8_out.dtype)).end(lane)
e8_store = e8_out.after(fp8_store)[super_idx * PACK + sb].store(e8u8)
packed = (e8u8.cast(dtypes.uint32) << (sb.cast(dtypes.uint32) * 8)).reduce(sb, arg=Ops.ADD)
row, col4 = super_idx // sk4, super_idx % sk4
@@ -72,8 +72,8 @@ def _custom_silu_mul_bwd_mxfp8(gx1_out:UOp, gx3_out:UOp, x_w1:UOp, x_w3:UOp, gra
sig = (1.0 + (w1 * -LOG2E).exp2()).reciprocal()
s = w1 * sig
sprime = sig * (1.0 + w1 * (1.0 - sig))
gx1 = gx1_out[idx].store((ga * sprime * w3).cast(gx1_out.dtype.base))
gx3 = gx3_out.after(gx1)[idx].store((ga * s).cast(gx3_out.dtype.base))
gx1 = gx1_out[idx].store((ga * sprime * w3).cast(gx1_out.dtype))
gx3 = gx3_out.after(gx1)[idx].store((ga * s).cast(gx3_out.dtype))
return gx3.end(lane, tid, wg).sink(arg=KernelInfo(f"silu_mul_bwd_mxfp8_{n_elems}", opts_to_apply=()))
def _silu_mul_quantize_mxfp8_bwd(gradient:UOp, kernel:UOp):
extra/llama_kernels/quantize_fp8_delayed/__init__.py
+2 -2
View File
@@ -27,7 +27,7 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_partial:UOp, x:UOp, amax_st
abs_x = (x_f < 0.0).where(-x_f, x_f)
scaled = (x_f * scale).maximum(-FP8_MAX).minimum(FP8_MAX)
fp8_store = fp8_out[idx].store(scaled.cast(fp8_out.dtype.base)).end(lane)
fp8_store = fp8_out[idx].store(scaled.cast(fp8_out.dtype)).end(lane)
lane_max = abs_x.reduce(lane, arg=Ops.MAX)
lmax = UOp.placeholder((1,), dtypes.float, slot=1, addrspace=AddrSpace.REG)
@@ -56,7 +56,7 @@ def _custom_quantize_fp8_scalar(fp8_out:UOp, x:UOp, amax_state:UOp) -> UOp:
x_f = x.reshape(n_elems)[i].cast(dtypes.float)
scale = FP8_MAX / (amax_state[0].cast(dtypes.float) + 1e-8)
store = fp8_out.reshape(n_elems)[i].store((x_f * scale).cast(fp8_out.dtype.base))
store = fp8_out.reshape(n_elems)[i].store((x_f * scale).cast(fp8_out.dtype))
return store.end(i).sink(arg=KernelInfo(f"quantize_fp8_scalar_{n_elems}"))
extra/llama_kernels/quantize_mxfp8_fused/__init__.py
+1 -1
View File
@@ -38,7 +38,7 @@ def _custom_quantize_mxfp8(fp8_out:UOp, e8_out:UOp, si_out:UOp, x:UOp) -> UOp:
scaled = (x_f * qscale).maximum(-FP8_MAX).minimum(FP8_MAX)
e8u8 = e8f.cast(dtypes.uint8)
fp8_store = fp8_out[idx].store(scaled.cast(fp8_out.dtype.base)).end(lane)
fp8_store = fp8_out[idx].store(scaled.cast(fp8_out.dtype)).end(lane)
e8_store = e8_out.after(fp8_store)[super_idx * PACK + sb].store(e8u8)
# pack the 4 e8 of this super-block into one uint32 (little-endian: byte sb), write transposed (sk4, row)
extra/thunder/tiny/tk/group.py
+22 -22
View File
@@ -47,8 +47,8 @@ class Group:
rngs_for_shape = tuple(self.ker.raw_range(dim) for dim in dst.shape)
src_load = src[*rngs_for_shape]
if src.dtype.base != dst.dtype.base:
src_load = src_load.cast(dst.dtype.base)
if src.dtype != dst.dtype:
src_load = src_load.cast(dst.dtype)
dst_store = dst[*rngs_for_shape].store(src_load).end(*rngs_for_shape)
self.ker.push_store(dst_store, dst)
@@ -62,8 +62,8 @@ class Group:
for width in self.ker.range(src.shape[-2], track=False):
for inner in self.ker.range(src.shape[-1], track=False):
src_load = src[height, width, inner]
if src.dtype.base != dst.dtype.base:
src_load = src_load.cast(dst.dtype.base)
if src.dtype != dst.dtype:
src_load = src_load.cast(dst.dtype)
dst_store = dst[width, height, inner].store(src_load).end(height, width, inner)
self.ker.push_store(dst_store, dst)
@@ -209,8 +209,8 @@ class Group:
vec, src = cast(UOp, vec), cast(UOp, src)
assert self.warps == 1
red_local = self.ker.alloc((self.group_threads,), src.dtype.base, AddrSpace.LOCAL)
red_reg = self.ker.alloc((1,), src.dtype.base, AddrSpace.REG)
red_local = self.ker.alloc((self.group_threads,), src.dtype, AddrSpace.LOCAL)
red_reg = self.ker.alloc((1,), src.dtype, AddrSpace.REG)
for height in self.ker.range(src.shape[-3], track=False):
i = self.ker.raw_range(red_reg.size)
@@ -243,8 +243,8 @@ class Group:
vec, src = cast(UOp, vec), cast(UOp, src)
assert self.warps == 1
red_local = self.ker.alloc((self.group_threads,), src.dtype.base, AddrSpace.LOCAL)
red_reg = self.ker.alloc((1,), src.dtype.base, AddrSpace.REG)
red_local = self.ker.alloc((self.group_threads,), src.dtype, AddrSpace.LOCAL)
red_reg = self.ker.alloc((1,), src.dtype, AddrSpace.REG)
for width in self.ker.range(src.shape[-2], track=False):
i = self.ker.raw_range(red_reg.size)
@@ -306,8 +306,8 @@ class Group:
srow, scol = cast(ST, src).swizzle(row, col)
src_load = src[*idxs[:-2], sheight, swidth, srow, scol]
if src.dtype.base != dst.dtype.base:
src_load = src_load.cast(dst.dtype.base)
if src.dtype != dst.dtype:
src_load = src_load.cast(dst.dtype)
dst_store = dst[*dst_idxs, height, width, inner].store(src_load)
dst_store = dst_store.end(height, width, inner)
elif dst.addrspace == AddrSpace.LOCAL and src.addrspace == AddrSpace.GLOBAL:
@@ -340,8 +340,8 @@ class Group:
src_i += row * row_stride + col
src_load = srcf[src_i]
if src.dtype.base != dst.dtype.base:
src_load = src_load.cast(dst.dtype.base)
if src.dtype != dst.dtype:
src_load = src_load.cast(dst.dtype)
dst_store = dst[*dst_idxs, height, width, srow, scol].store(src_load)
dst_store = dst_store.end(height, width, outer, inner).barrier()
elif dst.addrspace == AddrSpace.REG and src.addrspace == AddrSpace.GLOBAL and isinstance(dst, RT):
@@ -374,8 +374,8 @@ class Group:
src_i += srow * row_stride + scol
src_load = srcf[src_i]
if src.dtype.base != dst.dtype.base:
src_load = src_load.cast(dst.dtype.base)
if src.dtype != dst.dtype:
src_load = src_load.cast(dst.dtype)
dst_store = dst[*dst_idxs, height, width, inner].store(src_load).end(height, width, inner)
elif dst.addrspace == AddrSpace.REG and src.addrspace == AddrSpace.GLOBAL and isinstance(dst, RV):
srcf = src.flatten()
@@ -394,8 +394,8 @@ class Group:
src_i += outer * reductions + (laneid % reductions)
src_load = srcf[src_i]
if src.dtype.base != dst.dtype.base:
src_load = src_load.cast(dst.dtype.base)
if src.dtype != dst.dtype:
src_load = src_load.cast(dst.dtype)
dst_store = dst[outer, 0].store(src_load).end(outer)
else:
raise NotImplementedError(f"load from {src.addrspace} to {dst.addrspace} not implemented for {type(dst)=}")
@@ -423,8 +423,8 @@ class Group:
srow, scol = cast(ST, dst).swizzle(row, col)
src_load = src[*src_idxs, height, width, inner]
if src.dtype.base != dst.dtype.base:
src_load = src_load.cast(dst.dtype.base)
if src.dtype != dst.dtype:
src_load = src_load.cast(dst.dtype)
dst_store = dst[*idxs[:-2], height, width, srow, scol].store(src_load)
dst_store = dst_store.end(height, width, inner)
elif src.addrspace == AddrSpace.REG and dst.addrspace == AddrSpace.GLOBAL and isinstance(src, RT):
@@ -457,8 +457,8 @@ class Group:
dst_i += srow * row_stride + scol
src_load = src[*src_idxs, height, width, inner]
if src.dtype.base != dst.dtype.base:
src_load = src_load.cast(dst.dtype.base)
if src.dtype != dst.dtype:
src_load = src_load.cast(dst.dtype)
dst_store = dstf[dst_i].store(src_load).end(height, width, inner)
elif src.addrspace == AddrSpace.REG and dst.addrspace == AddrSpace.GLOBAL and isinstance(src, RV):
dstf = dst.flatten()
@@ -477,8 +477,8 @@ class Group:
dst_i += outer * reductions + (laneid % reductions)
src_load = src[outer, 0]
if src.dtype.base != dst.dtype.base:
src_load = src_load.cast(dst.dtype.base)
if src.dtype != dst.dtype:
src_load = src_load.cast(dst.dtype)
dst_store = dstf[dst_i].store(src_load).end(outer)
else:
raise NotImplementedError(f"store from {src.addrspace} to {dst.addrspace} not implemented for {type(src)=}")
extra/thunder/tiny/tk/tiles.py
+1 -1
View File
@@ -209,7 +209,7 @@ class ST:
return cls(uop, rows, cols, layout, base_shape, ker)
def swizzle(self, row, col):
swizzled_offset = self.base_shape.swizzle(row, col, self._uop.dtype.base.scalar())
swizzled_offset = self.base_shape.swizzle(row, col, self._uop.dtype.scalar())
row = swizzled_offset // self.base_shape.cols
col = swizzled_offset % self.base_shape.cols
test/amd/test_asm_kernel.py
+2 -2
View File
@@ -16,7 +16,7 @@ from extra.gemm.amd_asm_matmul import Kernel
def custom_add_one(A:UOp) -> UOp:
A = A.flatten()
assert dtypes.is_float(A.dtype.base), f"buffer dtype must be float32, got {A.dtype}"
assert dtypes.is_float(A.dtype), f"buffer dtype must be float32, got {A.dtype}"
threads = UOp.special(A.numel(), "lidx0")
insts = [
s_load_b64(s[0:1], s[0:1], soffset=NULL),
@@ -34,7 +34,7 @@ def custom_add_one(A:UOp) -> UOp:
def custom_add_var(A:UOp, B:UOp) -> UOp:
A,B = A.flatten(), B.flatten()
assert A.dtype.base == dtypes.uint32, f"buffer dtype must be uint32, got {A.dtype}"
assert A.dtype == dtypes.uint32, f"buffer dtype must be uint32, got {A.dtype}"
threads = UOp.special(A.numel(), "lidx0")
var = UOp.param(2, dtypes.weakint, vmin_vmax=(0, 10), name="var", addrspace=AddrSpace.ALU)
insts = [
test/backend/test_custom_kernel.py
+5 -5
View File
@@ -7,12 +7,12 @@ from tinygrad.uop.ops import KernelInfo, AxisType, Ops
def custom_arange_kernel(C:UOp) -> UOp:
i = UOp.range(C.shape[0], 0)
return C[i].store(i.cast(C.dtype.base)).end(i).sink(arg=KernelInfo(name=f"custom_arange_{C.shape[0]}"))
return C[i].store(i.cast(C.dtype)).end(i).sink(arg=KernelInfo(name=f"custom_arange_{C.shape[0]}"))
def custom_eye_kernel(C:UOp) -> UOp:
i = UOp.range(C.shape[0], 0)
j = UOp.range(C.shape[1], 1)
return C[i, j].store((i.eq(j)).cast(C.dtype.base)).end(i, j).sink(arg=KernelInfo(name=f"custom_eye_{C.numel()}"))
return C[i, j].store((i.eq(j)).cast(C.dtype)).end(i, j).sink(arg=KernelInfo(name=f"custom_eye_{C.numel()}"))
def custom_add_one_kernel(B:UOp, A:UOp) -> UOp:
A,B = A.flatten(), B.flatten()
@@ -57,7 +57,7 @@ def flip_contract_kernel(dest:UOp, src:UOp):
def slice_sum_kernel(dest:UOp, src:UOp):
G = UOp.range(src.shape[0], 0)
slice_src = src[G, :]
reg = UOp.placeholder((1,), dest.dtype.base, 0, addrspace=AddrSpace.REG)
reg = UOp.placeholder((1,), dest.dtype, 0, addrspace=AddrSpace.REG)
reg = reg.after(G)[0].set(0)
R = UOp.range(src.shape[1], 1, AxisType.REDUCE)
reg = reg[0].set(reg.after(R)[0] + slice_src[R], end=R)
@@ -73,12 +73,12 @@ def simple_qkv_kernel(O:UOp, Q:UOp, K:UOp, V:UOp) -> UOp:
j = UOp.range(N, 2, axis_type=AxisType.REDUCE)
k_inner = UOp.range(d, 3, axis_type=AxisType.REDUCE)
qk_acc = UOp.placeholder((1,), Q.dtype.base, 0, addrspace=AddrSpace.REG)
qk_acc = UOp.placeholder((1,), Q.dtype, 0, addrspace=AddrSpace.REG)
qk_acc = qk_acc.after(i, j)[0].set(0.0)
qk_acc = qk_acc[0].set(qk_acc.after(k_inner)[0] + Q[i, k_inner] * K[j, k_inner], end=k_inner)
qk_score = qk_acc[0] / (d ** 0.5)
out_acc = UOp.placeholder((1,), Q.dtype.base, 1, addrspace=AddrSpace.REG)
out_acc = UOp.placeholder((1,), Q.dtype, 1, addrspace=AddrSpace.REG)
out_acc = out_acc.after(i, d_out)[0].set(0.0)
out_acc = out_acc[0].set(out_acc.after(j)[0] + qk_score * V[j, d_out], end=j)
test/backend/test_linearizer.py
+2 -2
View File
@@ -211,14 +211,14 @@ class TestLinearizer(unittest.TestCase):
realized_ast = a.schedule_linear().src[-1].src[0]
program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
local = [uop for uop in tuple(program.src[1].src) if uop.op is Ops.BUFFER and uop.addrspace in (AddrSpace.LOCAL, AddrSpace.REG)]
assert local[0].dtype.base == acc_dtype
assert local[0].dtype == acc_dtype
def test_arg_acc_dtype(self):
def helper_arg_acc_dtype(c: Tensor, expected_dtype:DType):
realized_ast = c.schedule_linear().src[-1].src[0]
program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
local = [uop for uop in tuple(program.src[1].src) if uop.op is Ops.BUFFER and uop.addrspace in (AddrSpace.LOCAL, AddrSpace.REG)]
self.assertEqual(local[0].dtype.base, expected_dtype)
self.assertEqual(local[0].dtype, expected_dtype)
tests = (
(dtypes.float16, None, dtypes.float),
test/mockgpu/amd/pcode.py
+1 -1
View File
@@ -829,7 +829,7 @@ class Parser:
adt = dtypes.uint64 if addr.dtype == dtypes.uint64 else dtypes.uint32
active = self.vars.get('_active')
def mindex(idx:UOp): return mem.index(idx.valid(active) if active is not None else idx)
byte_mem = mem.dtype.base == dtypes.uint8
byte_mem = mem.dtype == dtypes.uint8
if byte_mem:
idx = addr
if dt in (dtypes.uint64, dtypes.int64, dtypes.float64):
test/null/test_tensor.py
+1 -1
View File
@@ -69,7 +69,7 @@ class TestIdxUpcast(unittest.TestCase):
if not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)):
assert idx.op is Ops.INDEX
idx_val = idx.src[1]
self.assertFalse(idx_val.overflows(idx_val.dtype.base.scalar()))
self.assertFalse(idx_val.overflows(idx_val.dtype.scalar()))
# use expand to generate kernel that uses large idx
def do_op_then_assert(self, dtype: DType, dim1, dim2, dim3):
tinygrad/codegen/opt/heuristic.py
+1 -1
View File
@@ -50,7 +50,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
# upcast float4 images, this must be early so we don't accidentally add locals before the upcast
if IMAGE:
for buf_index,buf in enumerate(k.bufs):
if image_valid_dims(buf.src[0].dtype.base, buf.src[0].max_numel(), k.ren.target.arch):
if image_valid_dims(buf.src[0].dtype, buf.src[0].max_numel(), k.ren.target.arch):
# part of is_expanded
unit_stride_axes_mul_4 = [k.rngs.index(c) for c in k.bufs[buf_index].src[1].get_idx().split_uop(Ops.ADD) if
c.op is Ops.RANGE and (c.vmax+1)%4 == 0]
tinygrad/codegen/opt/postrange.py
+1 -1
View File
@@ -332,7 +332,7 @@ class Scheduler:
def bufs_from_ast(ast:UOp, dname:str) -> list[Buffer]:
glbls = sorted([x for x in ast.backward_slice if x.op is Ops.PARAM and x.arg.slot >= 0], key=lambda x: x.arg.slot)
return [Buffer(dname, x.max_numel(), x.dtype.base) for x in glbls]
return [Buffer(dname, x.max_numel(), x.dtype) for x in glbls]
def apply_opts(ast:UOp, ren:Renderer, beam:int=0) -> UOp:
if ast.tag is not None: return ast
tinygrad/device.py
+2 -2
View File
@@ -202,8 +202,8 @@ class Buffer:
return self.copyout(memoryview(bytearray(self.nbytes)))
def numpy(self) -> 'np.ndarray': # type: ignore [name-defined] # noqa: F821
import numpy as np
assert _to_np_dtype(self.dtype.base) is not None, f"no np dtype for {self.dtype.base}"
return np.frombuffer(self.as_memoryview(), dtype=_to_np_dtype(self.dtype.base))
assert _to_np_dtype(self.dtype) is not None, f"no np dtype for {self.dtype}"
return np.frombuffer(self.as_memoryview(), dtype=_to_np_dtype(self.dtype))
def copyin(self, mv:memoryview):
mv = flat_mv(mv)
assert len(mv) == self.nbytes, f"size mismatch, {len(mv)=} != {self.dtype=} {self.size=}"
tinygrad/dtype.py
-4
View File
@@ -69,10 +69,6 @@ class DType(metaclass=DTypeMetaClass):
def __reduce__(self): return type(self), tuple(getattr(self, f.name) for f in fields(self))
def __repr__(self): return f"dtypes.{INVERSE_DTYPES_DICT[self.scalar().name]}"+(f".vec({self.count})" if self.count != 1 else "")
def __lt__(self, o:DType): return (self.priority, self.bitsize, self.name, self.fmt, self.count) < (o.priority, o.bitsize, o.name, o.fmt, o.count)
@property
def base(self): return self
@property
def vcount(self): return self.count
@functools.cache # pylint: disable=method-cache-max-size-none
def vec(self, sz:int) -> DType:
assert self.count == 1, f"can't vectorize {self} with size {sz}"
tinygrad/engine/realize.py
+1 -1
View File
@@ -231,7 +231,7 @@ pm_flatten_linear = PatternMatcher([
def _validate(call:UOp, sink:UOp) -> UOp:
params = get_call_arg_uops(call)
shadows = tuple(UOp.new_buffer(("CPU",)*len(p.device) if isinstance(p.device, tuple) else "CPU", prod(p.max_shape), p.dtype.base) for p in params)
shadows = tuple(UOp.new_buffer(("CPU",)*len(p.device) if isinstance(p.device, tuple) else "CPU", prod(p.max_shape), p.dtype) for p in params)
copies = tuple(p.copy_to_device(s.device).call(s, p) for s, p in zip(shadows, params))
return UOp(Ops.LINEAR, src=copies + (call, UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(sink,), arg="validate").call(*shadows, *params)))
pm_validate = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.SINK, name="sink"),), name="call", allow_any_len=True), _validate)]) + pm_flatten_linear
tinygrad/mixin/dtype.py
+1 -1
View File
@@ -67,7 +67,7 @@ class DTypeMixin:
print(t.is_floating_point())
```
"""
return dtypes.is_float(self.dtype.base)
return dtypes.is_float(self.dtype)
def float(self) -> Self:
"""
tinygrad/renderer/__init__.py
+1 -1
View File
@@ -35,7 +35,7 @@ class Estimates:
while len(buf.src) and buf.op is not Ops.PARAM: buf = buf.src[0]
if buf.op is Ops.PARAM:
# u.src[0] is INDEX, cap at buffer size for re-reads (e.g. matmul)
accessed = mem.get((buf, u.op), 0) + u.src[0].max_numel() * u.src[0].dtype.base.scalar().itemsize * mults
accessed = mem.get((buf, u.op), 0) + u.src[0].max_numel() * u.src[0].dtype.scalar().itemsize * mults
mem[(buf, u.op)] = smin(accessed, buf.max_numel() * buf.dtype.scalar().itemsize)
if u.op is Ops.RANGE:
mult_stack.append(mults)
tinygrad/renderer/cstyle.py
+2 -2
View File
@@ -227,7 +227,7 @@ class CStyleLanguage(Renderer):
assert l is not None, f"failed to render {u.op} {u.dtype} {[(x.op,x.dtype) for x in u.src]} {u.arg}"
if u.op in {Ops.ENDIF, Ops.END}: depth -= 1
if (u.op is not Ops.CAST or u.dtype.vcount == 1) and (u.op in {Ops.CONST, Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \
if (u.op is not Ops.CAST or u.dtype.count == 1) and (u.op in {Ops.CONST, Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \
(u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG) or \
(u.op is Ops.CAST and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)) or \
(u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))):
@@ -323,7 +323,7 @@ class OpenCLRenderer(CStyleLanguage):
]) + base_rewrite
def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str:
if any(uop.dtype.base == dtypes.half for uop in uops): prefix = (["#pragma OPENCL EXTENSION cl_khr_fp16 : enable"] + (prefix or []))
if any(uop.dtype == dtypes.half for uop in uops): prefix = (["#pragma OPENCL EXTENSION cl_khr_fp16 : enable"] + (prefix or []))
return super().render_kernel(function_name, kernel, bufs, uops, prefix)
def aux(self, uops:list[UOp]):
tinygrad/renderer/llvmir.py
+2 -2
View File
@@ -153,12 +153,12 @@ class LLVMRenderer(Renderer):
r[u] = f"%{'local' if u.addrspace == AddrSpace.LOCAL else 'reg'}_{str(u.arg.slot)}"
size = u.max_numel()
if u.addrspace == AddrSpace.REG:
kernel.append(f" {r[u]} = alloca [{size} x {ldt(u.dtype.base)}]")
kernel.append(f" {r[u]} = alloca [{size} x {ldt(u.dtype)}]")
elif self.has_local:
local_args.append(f"@{r[u][1:]} = internal unnamed_addr addrspace(3) global [{size} x {ldt(u.dtype)}] undef, align 16")
kernel.append(f" {r[u]} = addrspacecast [{size} x {ldt(u.dtype)}] addrspace(3)* @{r[u][1:]} to [{size} x {ldt(u.dtype)}]*")
else:
kernel.append(f" {r[u]} = alloca [{size} x {ldt(u.dtype.base)}], align 16")
kernel.append(f" {r[u]} = alloca [{size} x {ldt(u.dtype)}], align 16")
elif u.op is Ops.CONST: r[u] = lconst(u.arg, u.dtype)
elif u.op is Ops.CAST and ldt(u.dtype) == ldt(u.src[0].dtype):
r[u] = r[u.src[0]] # cast from signed to unsigned of the same size is a noop, or pointer cast
tinygrad/renderer/ptx.py
+2 -2
View File
@@ -175,7 +175,7 @@ class PTXRenderer(Renderer):
def ssa(prefix:str, u:UOp|None=None, dtype:str|None=None) -> str:
nonlocal c
prefix += f"_{dtype if dtype is not None else self.types[unwrap(u).dtype.base]}_"
prefix += f"_{dtype if dtype is not None else self.types[unwrap(u).dtype]}_"
c[prefix] += 1
return f"%{prefix}{c[prefix]-1}"
@@ -192,7 +192,7 @@ class PTXRenderer(Renderer):
r[u] = [cast(str,r[x]) for x in u.src]
continue
if u.op is Ops.BUFFER and u.addrspace == AddrSpace.REG:
r[u] = [ssa("reg", u, self.types[u.dtype.base.scalar()]) for _ in range(u.max_numel())]
r[u] = [ssa("reg", u, self.types[u.dtype.scalar()]) for _ in range(u.max_numel())]
continue
if u.op in {Ops.INDEX, Ops.SHRINK, Ops.LOAD} and u.src[0].addrspace in (AddrSpace.REG, AddrSpace.ALU):
# on REG, INDEX/SHRINK pick the register (must be CONST) and LOAD is a noop
tinygrad/renderer/wgsl.py
+2 -2
View File
@@ -95,14 +95,14 @@ class WGSLRenderer(CStyleLanguage):
def render_cast(self, u:UOp, val: str) -> str: return f"{self.type_map[u.dtype]}({val})"
def _render_dtype(self, dtype:DType, sz:int=1, addrspace=AddrSpace.REG, mutable=True, override_ptr=False, shape=None): return "var"
def render_load(self, x:str, u:UOp) -> str: return f"atomicLoad(&{x})" if is_packed(u) else x
def buf_map(self, u:UOp) -> str: return "atomic<u32>" if is_packed(u) else self.type_map[u.dtype.base]
def buf_map(self, u:UOp) -> str: return "atomic<u32>" if is_packed(u) else self.type_map[u.dtype]
def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str,tuple[UOp,bool]]], uops:list[UOp], prefix=None) -> str:
local_size = [u.src[0].ssimplify() for u in sorted([u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == 'l'], key=lambda u: u.arg)]
if not local_size: local_size = [1]
bind_it = iter(range(len(bufs)))
external_local_bufs = [line.lstrip() for line in kernel if "var<workgroup>" in line]
kernel[:] = [line for line in kernel if "var<workgroup>" not in line]
prg = "enable f16;\n" if any(uop.dtype.base == dtypes.half for uop in uops) else ""
prg = "enable f16;\n" if any(uop.dtype == dtypes.half for uop in uops) else ""
prg += "fn nan() -> f32 { let bits = 0xffffffffu; return bitcast<f32>(bits); }\n"
prg += "@group(0) @binding(0)\nvar<uniform> INFINITY : f32;\n"
prg += "\n".join((external_local_bufs or [])+[f"@group(0) @binding({next(bind_it)+1})" +
tinygrad/runtime/ops_python.py
+1 -1
View File
@@ -87,7 +87,7 @@ class PythonProgram:
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)
storage_fmt = storage_fmt_for_dtype(u.dtype)
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:
tinygrad/tensor.py
+8 -7
View File
@@ -237,7 +237,7 @@ class Tensor(RandMixin):
if capturing and not getenv("UNSAFE_ALLOW_JIT_BUFFER"):
from tinygrad.engine.jit import JitError
raise JitError("cannot access tensor data during JIT capture, the value will be baked in")
x = self.cast(self.dtype.base).contiguous()
x = self.cast(self.dtype).contiguous()
if self.uop.device is None or isinstance(self.device, tuple): x = x.clone("CPU")
return cast(Buffer, x.realize().uop.buffer).ensure_allocated()
@@ -252,11 +252,12 @@ class Tensor(RandMixin):
print(np.frombuffer(t.data(), dtype=np.int32))
```
"""
if 0 in self.shape: return memoryview(bytearray(0)).cast(self.dtype.base.fmt)
if 0 in self.shape: return memoryview(bytearray(0)).cast(self.dtype.fmt) # type: ignore[arg-type,return-value]
assert all_int(self.shape), f"no data if shape is symbolic, {self.shape=}"
assert self.dtype.base.fmt is not None, f"no fmt dtype for {self.dtype.base}"
assert self.dtype.base.fmt != "e" or sys.version_info >= (3, 12)
return self._data().cast(self.dtype.base.fmt, self.shape)
fmt = self.dtype.fmt
assert fmt is not None, f"no fmt dtype for {self.dtype}"
assert fmt != "e" or sys.version_info >= (3, 12)
return self._data().cast(fmt, self.shape) # type: ignore[arg-type,return-value]
# NOTE: list[Any] because return type is recursive (list[list[...]] for higher dimensions)
def tolist(self) -> PyConst|list[Any]:
@@ -288,8 +289,8 @@ class Tensor(RandMixin):
"""
assert all_int(self.shape), f"no data if shape is symbolic, {self.shape=}"
import numpy as np
if self.dtype.base in { dtypes.bfloat16, *dtypes.fp8s }: return self.float().numpy()
if 0 in self.shape: return np.empty(self.shape, dtype=_to_np_dtype(self.dtype.base))
if self.dtype in { dtypes.bfloat16, *dtypes.fp8s }: return self.float().numpy()
if 0 in self.shape: return np.empty(self.shape, dtype=_to_np_dtype(self.dtype))
return self._buffer().numpy().reshape(self.shape)
def clone(self, device:str|tuple[str, ...]|None=None) -> Tensor:
tinygrad/uop/ops.py
+6 -6
View File
@@ -453,7 +453,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), (self,)+tuple(new_srcs), **kwargs)
return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype), (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)
@@ -473,10 +473,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
@classmethod
def _wrap_uop(cls, u:UOp) -> UOp: return u
def const_like(self, b:ConstLike, dtype:DType|None=None):
return UOp.const(dtype or self.dtype.base, b, shape=self._shape)
return UOp.const(dtype or self.dtype, b, shape=self._shape)
def vconst_like(self, b:ConstLike, dtype:DType|None=None):
# for use after movement ops have been removed
ret = UOp.const(dtype or self.dtype.base, b)
ret = UOp.const(dtype or self.dtype, b)
if self.shape == (): return ret
if len(self.shape) == 1: return UOp(Ops.STACK, ret.dtype, (ret,)*self.max_numel())
raise RuntimeError(f"vconst_like only works on 0 or 1D shapes, not {self.shape}")
@@ -495,7 +495,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
def bitcast(self, dtype:DTypeLike):
dtype = to_dtype(dtype)
return self if self.dtype == dtype else UOp(Ops.BITCAST, dtype, (self,))
def load(self, *src:UOp, **kwargs): return UOp(Ops.LOAD, dtype=kwargs.pop("dtype", self.dtype.base), src=(self,)+src, **kwargs)
def load(self, *src:UOp, **kwargs): return UOp(Ops.LOAD, dtype=kwargs.pop("dtype", self.dtype), src=(self,)+src, **kwargs)
def store(self, src:UOp|ConstType, gate:UOp|None=None, **kwargs):
srcs = (self, self.const_like(src) if not isinstance(src, UOp) else src) + ((gate,) if gate is not None else ())
return UOp(Ops.STORE, dtypes.void, srcs, **kwargs)
@@ -512,7 +512,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
@staticmethod
def wmma(a:UOp, b:UOp, acc:UOp, arg:tuple[tuple[int, int, int], str, int]):
dims, device, threads = arg
dtype_in, dtype_out = a.dtype.base, acc.dtype.base
dtype_in, dtype_out = a.dtype, acc.dtype
tc_upcast_axes = tuple(((i, s.shape[-1]),) for i,s in enumerate((a, b, acc)))
name = f"WMMA_{'_'.join(map(str, dims))}_{dtype_in.name}_{dtype_out.name}"
return UOp(Ops.WMMA, dtype_out, (a, b, acc), arg=(name, dims, dtype_in, dtype_out, device, threads, tc_upcast_axes, ()))
@@ -828,7 +828,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
return ret
assert self.op is Ops.BUFFER, f"must be BUFFER {self.op}"
if (cret:=buffers.get(self)) is not None: return cret
rdtype = self.dtype.base
rdtype = self.dtype
if isinstance(self.device, tuple): ret = MultiBuffer(self.device, self.max_numel(), rdtype).ref(1)
else: ret = Buffer(self.device, self.max_numel(), rdtype).ref(1)
buffers[self] = ret
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 != x.dtype else None),
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__
tinygrad/uop/spec.py
+3 -3
View File
@@ -48,7 +48,7 @@ def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher):
# these ops can be used in the tensor graph and programs
spec_shared = PatternMatcher([
# no vec dtypes allowed
(UPat(GroupOp.All, name="x"), lambda x: False if x.dtype.vcount > 1 else None),
(UPat(GroupOp.All, name="x"), lambda x: False if x.dtype.count > 1 else None),
# NOTE: for testing, we let sinks be anything
(UPat(Ops.SINK, dtypes.void), lambda: True),
@@ -65,11 +65,11 @@ spec_shared = PatternMatcher([
# ALUs: most ALUs have all matching dtypes, except CMPLT, CMPNE, and WHERE
(UPat(Ops.WHERE, name="w", src=(UPat(dtype=dtypes.bool), UPat.var("x"), UPat.var("y"))), lambda w,x,y: w.dtype == x.dtype == y.dtype),
(UPat((Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ), dtype=dtypes.bool, src=(UPat.var("x"), UPat.var("y"))), lambda x,y: x.dtype.base == y.dtype.base),
(UPat((Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ), dtype=dtypes.bool, src=(UPat.var("x"), UPat.var("y"))), lambda x,y: x.dtype == y.dtype),
# and SHL/SHR, the shift distance can be an int
(UPat((Ops.SHL, Ops.SHR), src=(UPat.var("x"), UPat.var("y")), name="a"), lambda a,x,y: a.dtype == x.dtype and y.dtype in (x.dtype, dtypes.uint)),
(UPat((Ops.CDIV, Ops.CMOD, Ops.FLOORDIV, Ops.FLOORMOD), name="x"), lambda x: None if dtypes.is_int(x.dtype) else False),
(UPat(GroupOp.ALU, name="x"), lambda x: all(x.dtype.base == y.dtype.base for y in x.src)),
(UPat(GroupOp.ALU, name="x"), lambda x: all(x.dtype == y.dtype for y in x.src)),
# CAST
(UPat((Ops.BITCAST, Ops.CAST), src=(UPat(),), name="x"), lambda x: x.arg is None),