diff --git a/extra/gemm/mi350x_uop_matmul_2.py b/extra/gemm/mi350x_uop_matmul_2.py index 41c51f6e74..921dde021b 100644 --- a/extra/gemm/mi350x_uop_matmul_2.py +++ b/extra/gemm/mi350x_uop_matmul_2.py @@ -72,7 +72,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp: K_outer_loop = UOp.range(K//BLOCK_K, 0, AxisType.REDUCE) # split out the globals into blocks - C = C.src[0].cast(dtypes.float.vec(4).ptr(C.ptrdtype.size)).reshape((M//BLOCK_M, BLOCK_M, N//BLOCK_N, BLOCK_N)) + C = C.src[0].cast(dtypes.float.vec(4)).reshape((M//BLOCK_M, BLOCK_M, N//BLOCK_N, BLOCK_N)) A = A.reshape((M//BLOCK_M, BLOCK_M, K//BLOCK_K, BLOCK_K))[gx, :, K_outer_loop, :] B = B.reshape((K//BLOCK_K, BLOCK_K, N//BLOCK_N, BLOCK_N))[K_outer_loop, :, gy, :] diff --git a/extra/hcq2/hcq2.py b/extra/hcq2/hcq2.py index be672dc587..2c62ce832f 100644 --- a/extra/hcq2/hcq2.py +++ b/extra/hcq2/hcq2.py @@ -149,7 +149,7 @@ def make_ins(op, *srcs): return UOp(Ops.INS, dtypes.void, tuple(UOp.const(dtypes.uint32, s) if isinstance(s, int) else s.cast(dtypes.uint32) for s in srcs), op) def make_placeholder(devs, size:int, dtype, name=None, unique=True) -> UOp: - return UOp.param(next(UOp.unique_num) if unique else 0, dtype.ptr(size), device=devs).rtag(name or "buf") + 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)) diff --git a/extra/thunder/tiny/tk/group.py b/extra/thunder/tiny/tk/group.py index cb76fe507d..797ff101c5 100644 --- a/extra/thunder/tiny/tk/group.py +++ b/extra/thunder/tiny/tk/group.py @@ -2,7 +2,7 @@ import math from typing import cast, Callable from tinygrad import dtypes from tinygrad.uop.ops import AxisType, UOp, Ops -from tinygrad.dtype import AddrSpace, PtrDType +from tinygrad.dtype import AddrSpace from tinygrad.helpers import prod from extra.thunder.tiny.tk import WARP_THREADS @@ -277,9 +277,7 @@ class Group: def load(self, dst:ALL_TILES, src:ALL_TILES, dst_idxs:tuple[UOp|int,...]=(), idxs:tuple[UOp|int,...]=(), axis:int=0): dst, src = cast(UOp, dst), cast(UOp, src) - assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType) - dst_dtype, src_dtype = dst.dtype, src.dtype - if dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.LOCAL: + if dst.addrspace == AddrSpace.REG and src.addrspace == AddrSpace.LOCAL: laneid = self.ker.laneid rt, st = cast(RT, dst), cast(ST, src) elements_per_thread = rt.base_shape.elements_per_thread @@ -312,7 +310,7 @@ class Group: src_load = src_load.cast(dst.dtype.base) dst_store = dst[*dst_idxs, height, width, inner].store(src_load) dst_store = dst_store.end(height, width, inner) - elif dst_dtype.addrspace == AddrSpace.LOCAL and src_dtype.addrspace == AddrSpace.GLOBAL: + elif dst.addrspace == AddrSpace.LOCAL and src.addrspace == AddrSpace.GLOBAL: srcf = src.flatten() row_stride = prod(src.shape[axis+1:]) @@ -346,7 +344,7 @@ class Group: src_load = src_load.cast(dst.dtype.base) dst_store = dst[*dst_idxs, height, width, srow, scol].store(src_load) dst_store = dst_store.end(height, width, outer, inner).barrier() - elif dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.GLOBAL and isinstance(dst, RT): + elif dst.addrspace == AddrSpace.REG and src.addrspace == AddrSpace.GLOBAL and isinstance(dst, RT): srcf = src.flatten() row_stride = prod(src.shape[axis+1:]) @@ -379,7 +377,7 @@ class Group: if src.dtype.base != dst.dtype.base: src_load = src_load.cast(dst.dtype.base) dst_store = dst[*dst_idxs, height, width, inner].store(src_load).end(height, width, inner) - elif dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.GLOBAL and isinstance(dst, RV): + elif dst.addrspace == AddrSpace.REG and src.addrspace == AddrSpace.GLOBAL and isinstance(dst, RV): srcf = src.flatten() row_stride = prod(src.shape[axis+1:]) @@ -400,16 +398,14 @@ class Group: src_load = src_load.cast(dst.dtype.base) dst_store = dst[outer, 0].store(src_load).end(outer) else: - raise NotImplementedError(f"load from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented for {type(dst)=}") + raise NotImplementedError(f"load from {src.addrspace} to {dst.addrspace} not implemented for {type(dst)=}") self.ker.push_store(dst_store, dst) return dst.after(dst_store).reshape(dst.shape) def store(self, dst:ALL_TILES, src:ALL_TILES, idxs:tuple[UOp|int,...]=(), src_idxs:tuple[UOp|int,...]=(), axis:int=0): dst, src = cast(UOp, dst), cast(UOp, src) - assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType) - dst_dtype, src_dtype = dst.dtype, src.dtype - if src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.LOCAL: + if src.addrspace == AddrSpace.REG and dst.addrspace == AddrSpace.LOCAL: laneid = self.ker.laneid st, rt = cast(ST, dst), cast(RT, src) elements_per_thread = rt.base_shape.elements_per_thread @@ -431,7 +427,7 @@ class Group: src_load = src_load.cast(dst.dtype.base) dst_store = dst[*idxs[:-2], height, width, srow, scol].store(src_load) dst_store = dst_store.end(height, width, inner) - elif src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.GLOBAL and isinstance(src, RT): + elif src.addrspace == AddrSpace.REG and dst.addrspace == AddrSpace.GLOBAL and isinstance(src, RT): dstf = dst.flatten() row_stride = prod(dst.shape[axis+1:]) @@ -464,7 +460,7 @@ class Group: if src.dtype.base != dst.dtype.base: src_load = src_load.cast(dst.dtype.base) dst_store = dstf[dst_i].store(src_load).end(height, width, inner) - elif src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.GLOBAL and isinstance(src, RV): + elif src.addrspace == AddrSpace.REG and dst.addrspace == AddrSpace.GLOBAL and isinstance(src, RV): dstf = dst.flatten() row_stride = prod(dst.shape[axis+1:]) @@ -485,7 +481,7 @@ class Group: src_load = src_load.cast(dst.dtype.base) dst_store = dstf[dst_i].store(src_load).end(outer) else: - raise NotImplementedError(f"store from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented for {type(src)=}") + raise NotImplementedError(f"store from {src.addrspace} to {dst.addrspace} not implemented for {type(src)=}") self.ker.push_store(dst_store, dst) return dst.after(dst_store).reshape(dst.shape) diff --git a/tinygrad/callify.py b/tinygrad/callify.py index 6b20673ec9..4732c53393 100644 --- a/tinygrad/callify.py +++ b/tinygrad/callify.py @@ -1,5 +1,5 @@ from dataclasses import dataclass, field -from tinygrad.dtype import dtypes, AddrSpace, PtrDType, ImageDType +from tinygrad.dtype import dtypes, AddrSpace from tinygrad.uop.ops import UOp, UPat, PatternMatcher, Ops, GroupOp, ParamArg, graph_rewrite, track_rewrites from tinygrad.helpers import VIZ, pluralize, all_int @@ -177,7 +177,7 @@ def replace_input_buffer(ctx:AllocCtx, b:UOp): ctx.replacements.append(b) return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b.device, b._min_max if b.op is Ops.BIND else None, b.src[0].expr if b.op is Ops.BIND else None, - b.addrspace if isinstance(b.dtype, (PtrDType, ImageDType)) else AddrSpace.GLOBAL) + b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL) pm_finalize_call = PatternMatcher([ (UPat(Ops.AFTER, name="x"), finalize_after), diff --git a/tinygrad/device.py b/tinygrad/device.py index 047b50e62f..fba2e554ee 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -6,7 +6,7 @@ import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickl from tinygrad.helpers import LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize -from tinygrad.dtype import DType, PtrDType, _to_np_dtype +from tinygrad.dtype import DType, _to_np_dtype if TYPE_CHECKING: from tinygrad.renderer import Renderer # **************** Device **************** @@ -102,7 +102,7 @@ class Buffer: profile_events:list[ProfileEvent] = [] def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:BufferSpec|None=None, initial_value:bytes|None=None, uop_refcount=0, base:Buffer|None=None, offset:int=0, preallocate=False): - assert isinstance(dtype, DType) and not isinstance(dtype, PtrDType) + assert isinstance(dtype, DType) self.device, self.size, self.dtype, self.options, self.offset, self.allocated_views = device, size, dtype, options, offset, 0 self._bufs: dict[str, Any] = {} if base is None: diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py index ab5958eb28..40f4e0d27d 100644 --- a/tinygrad/dtype.py +++ b/tinygrad/dtype.py @@ -78,10 +78,7 @@ class DType(metaclass=DTypeMetaClass): assert self.count == 1, f"can't vectorize {self} with size {sz}" if sz == 1 or self == dtypes.void: return self # void doesn't vectorize, and sz=1 is scalar return DType(self.priority, self.bitsize*sz, f"{INVERSE_DTYPES_DICT[self.name]}{sz}", None, sz, self) - def ptr(self, size=-1, addrspace=AddrSpace.GLOBAL) -> PtrDType: - return PtrDType(self.priority, self.bitsize, self.name, self.fmt, self.count, None, self, addrspace, 1, size) def scalar(self) -> DType: return self._scalar if self._scalar is not None else self - def nbytes(self) -> int: raise RuntimeError("only ptr types have nbytes") @functools.cached_property def min(self): if dtypes.is_int(self): return 0 if dtypes.is_unsigned(self) else -2**(self.scalar().bitsize-1) @@ -101,36 +98,24 @@ class DType(metaclass=DTypeMetaClass): return ConstFloat(float(val)) if dtypes.is_float(self) else bool(val) if dtypes.is_bool(self) else int(val) @dataclass(frozen=True, eq=False) -class PtrDType(DType): +class ImageDType(DType): _base: DType addrspace: AddrSpace v: int size: int = -1 # -1 is unlimited size + shape: tuple[int, ...] = () # shape of the Image @property def base(self): return self._base @functools.cache # pylint: disable=method-cache-max-size-none def vec(self, sz:int) -> DType: - assert self.v == 1, f"can't vectorize ptr {self} with size {sz}" + assert self.v == 1, f"can't vectorize image {self} with size {sz}" if sz == 1: return self # sz=1 is a scalar - if isinstance(self, ImageDType): - return ImageDType(self.priority, self.bitsize, self.name, self.fmt, self.count, self, self._base, self.addrspace, sz, self.size, self.shape) - return type(self)(self.priority, self.bitsize, self.name, self.fmt, self.count, self, self._base, self.addrspace, sz, self.size) - def ptr(self, size=-1, addrspace=AddrSpace.GLOBAL) -> PtrDType: raise RuntimeError("can't make a pointer from a pointer") + return ImageDType(self.priority, self.bitsize, self.name, self.fmt, self.count, self, self._base, self.addrspace, sz, self.size, self.shape) def nbytes(self) -> int: if self.size == -1: raise RuntimeError("can't get nbytes of a pointer with unlimited size") return self.size*self.itemsize @property def vcount(self): return self.v - def __repr__(self): - return f"{self.base.__repr__()}.ptr({self.size}{', '+str(self.addrspace) if self.addrspace != AddrSpace.GLOBAL else ''})" + \ - (f'.vec({self.v})' if self.v != 1 else '') - -@dataclass(frozen=True, eq=False) -class ImageDType(PtrDType): - shape: tuple[int, ...] = () # shape of the Image - def ptr(self, size=-1, addrspace=AddrSpace.GLOBAL) -> PtrDType: - assert addrspace == AddrSpace.GLOBAL, "images can't be local" - return self def __repr__(self): return f"dtypes.{self.name}({self.shape})" + (f'.vec({self.v})' if self.v != 1 else '') # for 1d images on macos, we need to round pitch up to 256 pixels to make CL happy diff --git a/tinygrad/mixin/__init__.py b/tinygrad/mixin/__init__.py index ff94401589..d7ad9b813c 100644 --- a/tinygrad/mixin/__init__.py +++ b/tinygrad/mixin/__init__.py @@ -6,7 +6,7 @@ from tinygrad.mixin.movement import MovementMixin from tinygrad.mixin.reduce import ReduceMixin from tinygrad.uop import Ops from tinygrad.uop.ops import _broadcast_shape, resolve, smax, smin, identity_element -from tinygrad.dtype import ConstType, DType, DTypeLike, Invalid, PtrDType, PyConst, dtypes, least_upper_dtype, sum_acc_dtype, to_dtype +from tinygrad.dtype import ConstType, DType, DTypeLike, Invalid, ImageDType, PyConst, dtypes, least_upper_dtype, sum_acc_dtype, to_dtype from tinygrad.helpers import all_int, argfix, argsort, ceildiv, flatten, flat_to_grouped, fully_flatten, get_shape, make_tuple, merge_dicts, prod from tinygrad.helpers import resolve_pool_pads, round_up, IMAGE, FLOAT16, WINO @@ -366,7 +366,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): x, y = x._broadcast_to(out_shape), y._broadcast_to(out_shape) except (RuntimeError, ValueError): pass # ptr dtypes aren't in the promo lattice - if x.dtype == y.dtype or any(isinstance(d, PtrDType) for d in (x.dtype, y.dtype)): return x, y + if x.dtype == y.dtype or any(isinstance(d, ImageDType) for d in (x.dtype, y.dtype)): return x, y return x.cast(out_dtype := least_upper_dtype(x.dtype, y.dtype)), y.cast(out_dtype) def dot(self, w:Self, dtype:DTypeLike|None=None) -> Self: diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 6351a3e61f..0071e57cac 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -4,7 +4,7 @@ from collections import defaultdict, Counter from tinygrad.codegen.opt import tc from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str, axis_letters from tinygrad.helpers import strip_parens, getenv, prod, dedup, Target, CPU_COUNT, IMAGE, FLOAT16 -from tinygrad.dtype import ImageDType, dtypes, DType, PtrDType, AddrSpace, truncate, float_to_bf16 +from tinygrad.dtype import ImageDType, dtypes, DType, AddrSpace, truncate, float_to_bf16 from tinygrad.renderer import Renderer @@ -185,7 +185,7 @@ class CStyleLanguage(Renderer): # LEGACY def render_dtype(self, dt:DType, mutable=True) -> str: - return self._render_dtype(dt, dt.count, dt.addrspace if isinstance(dt, PtrDType) else AddrSpace.REG) + return self._render_dtype(dt, dt.count, dt.addrspace if isinstance(dt, ImageDType) else AddrSpace.REG) def __getitem__(self, key): return self.r[key] # hacky helper def _render(self, uops:list[UOp]) -> tuple[str, list[str], list[tuple[str,tuple[UOp,bool]]]]: diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 262da899ac..0d2d381417 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -4,7 +4,7 @@ from tinygrad.renderer import Renderer from tinygrad.renderer.cstyle import HIPRenderer, create_non_native_float_pats, pm_manual_bf16_cast from tinygrad.codegen.decomp.transcendental import xexp2, xlog2 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp, range_str -from tinygrad.dtype import dtypes, float_to_fp8, DType, PtrDType, truncate, AddrSpace +from tinygrad.dtype import dtypes, float_to_fp8, DType, truncate, AddrSpace from tinygrad.helpers import prod, Target, CPU_COUNT, getenv, OSX def ldt(dt:DType, count=1, ptr=False): @@ -160,7 +160,7 @@ class LLVMRenderer(Renderer): else: kernel.append(f" {r[u]} = alloca [{size} x {ldt(u.dtype.base)}], 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) or isinstance(u.dtype, PtrDType)): + 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 else: # if it's an assign target, it's already preallocated diff --git a/tinygrad/runtime/ops_qcom.py b/tinygrad/runtime/ops_qcom.py index 2c8df47f71..dac94ffe8f 100644 --- a/tinygrad/runtime/ops_qcom.py +++ b/tinygrad/runtime/ops_qcom.py @@ -20,7 +20,7 @@ BUFTYPE_BUF, BUFTYPE_TEX, BUFTYPE_IBO = 0, 1, 2 def dcache_flush(): from tinygrad.uop.ops import UOp, Ops, KernelInfo from tinygrad.codegen import to_program - buf, n = UOp.param(0, dtypes.uint8.ptr(1)), UOp.param(1, dtypes.int, shape=(1,), name="n", addrspace=None) + buf, n = UOp.param(0, dtypes.uint8, shape=(1,)), UOp.param(1, dtypes.int, shape=(1,), name="n", addrspace=None) i = UOp.range(n, 0, dtype=dtypes.int) flush = UOp(Ops.CUSTOM, dtypes.void, (buf.index(i * 64),), arg='__asm__ volatile("dc cvac, %0" :: "r"({0}) : "memory");') sink = UOp.sink(flush.end(i), UOp(Ops.CUSTOM, dtypes.void, (), arg='__asm__ volatile("dsb sy" ::: "memory");'), arg=KernelInfo(name="dcache_flush")) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 0569086d0a..5d4af5fb9b 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -1,7 +1,7 @@ from dataclasses import dataclass, field, replace from typing import cast import itertools -from tinygrad.dtype import dtypes, PtrDType, AddrSpace, Invalid +from tinygrad.dtype import dtypes, AddrSpace, Invalid from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, identity_element from tinygrad.uop.symbolic import symbolic @@ -18,9 +18,6 @@ import sys sys.setrecursionlimit(10000) pm_syntactic_sugar = PatternMatcher([ - # INDEX on ptr INDEX concats them - (UPat(Ops.INDEX, name="i1").f(Ops.INDEX, name="i2", allow_any_len=True), - lambda i1,i2: i2.replace(src=i1.src+i2.src[1:]) if isinstance(i1.dtype, PtrDType) and not isinstance(i2.dtype, PtrDType) else None), # early rangeify (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise | {Ops.CONST}, name="x"),), allow_any_len=True, name="idx"), lambda idx,x: x.replace(src=tuple([s.index(*idx.src[1:]) for s in x.src]))), @@ -515,7 +512,7 @@ to_define_global = PatternMatcher([ # this renumbers the params (UPat(Ops.PARAM, name="buf"), lambda ctx, buf: - None if buf.tag != () or isinstance(buf.dtype, PtrDType) or buf.arg.name is not None or buf._shape is None else debuf(ctx, buf)), + None if buf.tag != () or buf.arg.name is not None or buf._shape is None else debuf(ctx, buf)), # ALU params are scalar symbolic values, not buffers. (UPat(Ops.INDEX, src=(UPat(Ops.PARAM, name="v"),)), lambda v: v if v.addrspace == AddrSpace.ALU else None), diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index c6e4d6d47c..af237c520a 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -4,7 +4,7 @@ import sys, time, functools, itertools, math, operator, hashlib, os, types, pick from dataclasses import dataclass from enum import Enum, auto from tinygrad.uop import Ops, GroupOp -from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, DTypeLike, to_dtype, truncate, PtrDType, least_upper_dtype, Invalid, AddrSpace +from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, DTypeLike, to_dtype, truncate, least_upper_dtype, Invalid, AddrSpace from tinygrad.dtype import ConstFloat, PyConst, InvalidType, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar from tinygrad.device import Buffer, MultiBuffer, canonicalize_device from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA @@ -215,11 +215,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass): def tuplize(self:UOp) -> tuple: return (self.op.value, self.arg, self.dtype,)+tuple([x.tuplize for x in self.src]) - @property - def ptrdtype(self) -> PtrDType: - if not isinstance(self.dtype, PtrDType): raise RuntimeError(f"ptrdtype called on UOp with type {self.dtype}") - return self.dtype - # *** uop shape stuff *** @recursive_property @@ -256,11 +251,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): case Ops.STACK: if len(self.src) == 0: return () - if isinstance(self.dtype, PtrDType): - # TODO: this is broken - return self.src[0].shape - else: - return (len(self.src),) + self.src[0].shape + return (len(self.src),) + self.src[0].shape case Ops.CONST: return (self.dtype.count,) if self.dtype.count > 1 else () @@ -270,7 +261,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass): case Ops.BINARY: return (len(self.arg),) case Ops.BUFFER: if len(self.src): return self.src[0].as_shape - if isinstance(self.dtype, PtrDType): return (self.ptrdtype.size, self.dtype.count) if self.dtype.count > 1 else (self.ptrdtype.size,) return (self.dtype.count,) if self.dtype.count > 1 else () case Ops.SLICE: # HACK: SLICE is used inside kernels, so we set the shape to () if it's on an INDEX @@ -286,7 +276,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return tuple([int(r.vmax+1) for r in self.src[1:]])+self.src[0].shape case Ops.PARAM: if isinstance(self.dtype, ImageDType): return self.dtype.shape - if isinstance(self.dtype, PtrDType): return (self.ptrdtype.size,) return self.src[0].as_shape if len(self.src) >= 1 else None # wmma output shape = accumulator shape (src[2]) @@ -975,7 +964,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if self.op is Ops.BIND: return self.src[0]._min_max # ignore the bound value if self.op is Ops.STACK: return min(x.vmin for x in self.src), max(x.vmax for x in self.src) if self.op is Ops.CONST and self.arg is not Invalid: return self.arg, self.arg - if self.op is Ops.INDEX and not isinstance(self.src[0].dtype, PtrDType): return self.src[0]._min_max + if self.op is Ops.INDEX: return self.src[0]._min_max # TODO: CAST to bool/unsigned is not monotone, still some case can be simplified if self.op is Ops.CAST and self.dtype in dtypes.floats+dtypes.sints+(dtypes.weakint,): return max(self.dtype.min, self.src[0].vmin), min(self.src[0].vmax, self.dtype.max) @@ -1038,7 +1027,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): src: tuple[UOp, ...] = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),) return UOp(Ops.PARAM, dtype, src, arg=ParamArg(slot, vmin_vmax, name, addrspace, axis, device)) def param_like(self, slot:int): - addrspace = self.addrspace if isinstance(self.dtype, (PtrDType, ImageDType)) else AddrSpace.GLOBAL + addrspace = self.addrspace if self.addrspace is not None else AddrSpace.GLOBAL if self.op is Ops.BIND: return UOp.param(slot, self.dtype, self._shape, self.device, cast(tuple[int, int], self._min_max), self.src[0].expr, addrspace) return UOp.param(slot, self.dtype, self.shard_shape if self.axis is not None else self._shape, self.device, addrspace=addrspace, axis=self.axis) diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index b3cdc61110..f0058a88a2 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -2,7 +2,7 @@ import math from typing import Any from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, AxisType, KernelInfo, ParamArg from tinygrad.uop.render import print_uops, pyrender -from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid, ConstFloat +from tinygrad.dtype import DType, ImageDType, dtypes, AddrSpace, Invalid, ConstFloat from tinygrad.helpers import DEBUG, Context, SPEC, Metadata, panic, CHECK_OOB, all_same # ***** uop helpers ***** @@ -26,7 +26,7 @@ def validate_index(uidx:UOp, gate:UOp|None=None): # VECTORIZE can't be properly modeled in z3 since it doesn't support vectors # don't descend into PARAM shape metadata; only the PARAM value participates in index arithmetic for x in idx.toposort(gate=lambda x: x.op is not Ops.PARAM) | gate.toposort(gate=lambda x: x.op is not Ops.PARAM): - if x.op in {Ops.BITCAST, Ops.STACK} or (x.op is Ops.CAST and isinstance(x.src[0].dtype, PtrDType)): return True + if x.op in {Ops.BITCAST, Ops.STACK}: return True # if all is good and CHECK_OOB=1, validate with z3 from tinygrad.uop.validate import validate_index_with_z3 diff --git a/tinygrad/uop/validate.py b/tinygrad/uop/validate.py index 0fe4db31d7..73f3ddd206 100644 --- a/tinygrad/uop/validate.py +++ b/tinygrad/uop/validate.py @@ -51,7 +51,7 @@ z3_renderer = PatternMatcher([ ]) def uops_to_z3(solver:z3.Solver, *uops: UOp) -> list[z3.ExprRef]: - # gate on upstream AFTER/BUFFER as a replacement for PtrDType, but keep INDEX as an unknown LOAD + # 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] z3map: dict[UOp, z3.ExprRef] = {}