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

final removal of PtrDType (glm) (#16913)

Browse Source 
* final removal of PtrDType (glm)

* junk
This commit is contained in:
George Hotz
2026-07-07 13:37:33 -07:00
committed by GitHub
parent 27de6c6db0
commit 2fc7e5341b
14 changed files with 36 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/gemm/mi350x_uop_matmul_2.py
+1 -1
View File
@@ -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, :]
extra/hcq2/hcq2.py
+1 -1
View File
@@ -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))
extra/thunder/tiny/tk/group.py
+10 -14
View File
@@ -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)
tinygrad/callify.py
+2 -2
View File
@@ -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),
tinygrad/device.py
+2 -2
View File
@@ -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:
tinygrad/dtype.py
+4 -19
View File
@@ -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
tinygrad/mixin/__init__.py
+2 -2
View File
@@ -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:
tinygrad/renderer/cstyle.py
+2 -2
View File
@@ -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]]]]:
tinygrad/renderer/llvmir.py
+2 -2
View File
@@ -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
tinygrad/runtime/ops_qcom.py
+1 -1
View File
@@ -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"))
tinygrad/schedule/rangeify.py
+2 -5
View File
@@ -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),
tinygrad/uop/ops.py
+4 -15
View File
@@ -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)
tinygrad/uop/spec.py
+2 -2
View File
@@ -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
tinygrad/uop/validate.py
+1 -1
View File
@@ -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] = {}