diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000000..e514039334 --- /dev/null +++ b/AGENTS.md @@ -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 diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index 78c331883c..c97bcaed8d 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -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 diff --git a/extra/gemm/amd_copy_matmul.py b/extra/gemm/amd_copy_matmul.py index 830e040319..af0ccb58a5 100644 --- a/extra/gemm/amd_copy_matmul.py +++ b/extra/gemm/amd_copy_matmul.py @@ -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) diff --git a/extra/gemm/cdna_asm_gemm.py b/extra/gemm/cdna_asm_gemm.py index 90459aff3b..1611ba6f23 100644 --- a/extra/gemm/cdna_asm_gemm.py +++ b/extra/gemm/cdna_asm_gemm.py @@ -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}')) diff --git a/extra/hcq2/hcq2.py b/extra/hcq2/hcq2.py index 2c62ce832f..c6e4637a6b 100644 --- a/extra/hcq2/hcq2.py +++ b/extra/hcq2/hcq2.py @@ -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: diff --git a/extra/hcq2/ops_amd2.py b/extra/hcq2/ops_amd2.py index ad0d610cfc..dfbbe18023 100644 --- a/extra/hcq2/ops_amd2.py +++ b/extra/hcq2/ops_amd2.py @@ -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), diff --git a/extra/hip_large_kernel.py b/extra/hip_large_kernel.py index 1b67384f3d..5eef8cbfd9 100644 --- a/extra/hip_large_kernel.py +++ b/extra/hip_large_kernel.py @@ -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) diff --git a/extra/llama_kernels/fused_ce/__init__.py b/extra/llama_kernels/fused_ce/__init__.py index 0527d5781b..362796c223 100644 --- a/extra/llama_kernels/fused_ce/__init__.py +++ b/extra/llama_kernels/fused_ce/__init__.py @@ -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) diff --git a/extra/llama_kernels/fused_silu_mul_quantize_mxfp8/__init__.py b/extra/llama_kernels/fused_silu_mul_quantize_mxfp8/__init__.py index f441f7e119..85126c2825 100644 --- a/extra/llama_kernels/fused_silu_mul_quantize_mxfp8/__init__.py +++ b/extra/llama_kernels/fused_silu_mul_quantize_mxfp8/__init__.py @@ -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): diff --git a/extra/llama_kernels/quantize_fp8_delayed/__init__.py b/extra/llama_kernels/quantize_fp8_delayed/__init__.py index dca8fa3164..f6eab12b7f 100644 --- a/extra/llama_kernels/quantize_fp8_delayed/__init__.py +++ b/extra/llama_kernels/quantize_fp8_delayed/__init__.py @@ -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}")) diff --git a/extra/llama_kernels/quantize_mxfp8_fused/__init__.py b/extra/llama_kernels/quantize_mxfp8_fused/__init__.py index 86236e8ef5..b176e77459 100644 --- a/extra/llama_kernels/quantize_mxfp8_fused/__init__.py +++ b/extra/llama_kernels/quantize_mxfp8_fused/__init__.py @@ -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) diff --git a/extra/thunder/tiny/tk/group.py b/extra/thunder/tiny/tk/group.py index 797ff101c5..389f564097 100644 --- a/extra/thunder/tiny/tk/group.py +++ b/extra/thunder/tiny/tk/group.py @@ -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)=}") diff --git a/extra/thunder/tiny/tk/tiles.py b/extra/thunder/tiny/tk/tiles.py index 2538d84e14..98cff727c7 100644 --- a/extra/thunder/tiny/tk/tiles.py +++ b/extra/thunder/tiny/tk/tiles.py @@ -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 diff --git a/test/amd/test_asm_kernel.py b/test/amd/test_asm_kernel.py index 59b715f7ff..5ba12765cd 100644 --- a/test/amd/test_asm_kernel.py +++ b/test/amd/test_asm_kernel.py @@ -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 = [ diff --git a/test/backend/test_custom_kernel.py b/test/backend/test_custom_kernel.py index 3599bed385..09e347d630 100644 --- a/test/backend/test_custom_kernel.py +++ b/test/backend/test_custom_kernel.py @@ -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) diff --git a/test/backend/test_linearizer.py b/test/backend/test_linearizer.py index d0cf729223..ac8fb3a846 100644 --- a/test/backend/test_linearizer.py +++ b/test/backend/test_linearizer.py @@ -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), diff --git a/test/mockgpu/amd/pcode.py b/test/mockgpu/amd/pcode.py index 41e8026564..11201504e2 100644 --- a/test/mockgpu/amd/pcode.py +++ b/test/mockgpu/amd/pcode.py @@ -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): diff --git a/test/null/test_tensor.py b/test/null/test_tensor.py index 7d5497b852..9090eccc5f 100644 --- a/test/null/test_tensor.py +++ b/test/null/test_tensor.py @@ -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): diff --git a/tinygrad/codegen/opt/heuristic.py b/tinygrad/codegen/opt/heuristic.py index a7d428f6bf..32bf1d29fd 100644 --- a/tinygrad/codegen/opt/heuristic.py +++ b/tinygrad/codegen/opt/heuristic.py @@ -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] diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index 84ae4166ee..40e18af9a7 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -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 diff --git a/tinygrad/device.py b/tinygrad/device.py index fba2e554ee..637089cdd2 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -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=}" diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py index 2653ae6e91..87431f3e5e 100644 --- a/tinygrad/dtype.py +++ b/tinygrad/dtype.py @@ -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}" diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index 370e331209..6346774dda 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -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 diff --git a/tinygrad/mixin/dtype.py b/tinygrad/mixin/dtype.py index 518d9994e5..09c14cd67c 100644 --- a/tinygrad/mixin/dtype.py +++ b/tinygrad/mixin/dtype.py @@ -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: """ diff --git a/tinygrad/renderer/__init__.py b/tinygrad/renderer/__init__.py index 761590becf..227217fdea 100644 --- a/tinygrad/renderer/__init__.py +++ b/tinygrad/renderer/__init__.py @@ -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) diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index f429340346..e48c6f8b2e 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -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]): diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 0d2d381417..c05e52c4b7 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -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 diff --git a/tinygrad/renderer/ptx.py b/tinygrad/renderer/ptx.py index a7ddfd0039..d798e5e00c 100644 --- a/tinygrad/renderer/ptx.py +++ b/tinygrad/renderer/ptx.py @@ -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 diff --git a/tinygrad/renderer/wgsl.py b/tinygrad/renderer/wgsl.py index 88a160dffe..677153d557 100644 --- a/tinygrad/renderer/wgsl.py +++ b/tinygrad/renderer/wgsl.py @@ -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" if is_packed(u) else self.type_map[u.dtype.base] + def buf_map(self, u:UOp) -> str: return "atomic" 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" in line] kernel[:] = [line for line in kernel if "var" 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(bits); }\n" prg += "@group(0) @binding(0)\nvar INFINITY : f32;\n" prg += "\n".join((external_local_bufs or [])+[f"@group(0) @binding({next(bind_it)+1})" + diff --git a/tinygrad/runtime/ops_python.py b/tinygrad/runtime/ops_python.py index bc6429af62..a3c5b53589 100644 --- a/tinygrad/runtime/ops_python.py +++ b/tinygrad/runtime/ops_python.py @@ -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: diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 0383ebb59f..5ada5b299f 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -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: diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 4086d8e3b1..f0b254b370 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -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 diff --git a/tinygrad/uop/render.py b/tinygrad/uop/render.py index 969c0f537e..0712c246fd 100644 --- a/tinygrad/uop/render.py +++ b/tinygrad/uop/render.py @@ -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__ diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 654e022d2a..c7ff560649 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -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),