diff --git a/test/backend/test_uops.py b/test/backend/test_uops.py index 39860e24f8..9a454bf3c5 100644 --- a/test/backend/test_uops.py +++ b/test/backend/test_uops.py @@ -20,15 +20,15 @@ def run_uops(uops_list:list[UOp], bufs:list[Buffer]): def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp: if op is Ops.CONST: uops.append(UOp.const(dtype, arg)) - elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype).replace(src=())) + elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, shape=(1,))) else: uops.append(UOp(op, dtype, tuple(src), arg)) return uops[-1] def _test_single_value(vals, op, dts): uops = [] output_dtype = dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else dts[-1] - buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(1), (), 0) - buf_loads = [uop(uops, Ops.PARAM, dtype.ptr(1), (), i+1) for i,dtype in enumerate(dts)] + buf_store = uop(uops, Ops.PARAM, output_dtype, (), 0) + buf_loads = [uop(uops, Ops.PARAM, dtype, (), i+1) for i,dtype in enumerate(dts)] loads = (buf_loads[i].index(uop(uops, Ops.CONST, dtypes.int32, (), 0)) for i, dtype in enumerate(dts)) alu = uop(uops, op, output_dtype, loads) out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0), ptr=True), alu)) @@ -42,7 +42,7 @@ def _test_single_value(vals, op, dts): def _test_single_value_const(vals, op, dts): uops = [] output_dtype = dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else dts[-1] - buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(1), (), 0) + buf_store = uop(uops, Ops.PARAM, output_dtype, (), 0) loads = (uop(uops, Ops.CONST, dtype, [], a) for a,dtype in zip(vals, dts)) alu = uop(uops, op, output_dtype, loads) out = buf_store[UOp.const(dtypes.int32, 0)].store(alu) @@ -54,7 +54,7 @@ def _test_single_value_const(vals, op, dts): def _test_uops_result(output_dtype, uops, res): # uops = [] - buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(1), (), 0) + buf_store = uop(uops, Ops.PARAM, output_dtype, (), 0) # res = output_fn(uops) out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), res)) buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate() @@ -225,8 +225,8 @@ class TestLocalAccess(unittest.TestCase): @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "This only tests assembly backends") class TestAssembly(unittest.TestCase): def test_bitshift_left(self): - g1 = UOp.param(0, dtypes.int32.ptr(3)) - out = UOp.param(1, dtypes.int32.ptr(2)) + g1 = UOp.param(0, dtypes.int32, shape=(3,)) + out = UOp.param(1, dtypes.int32, shape=(2,)) c1 = UOp.const(dtypes.int, 2) c2 = UOp.const(dtypes.int, 3) l1 = g1.index(c1) @@ -254,7 +254,7 @@ class TestAssembly(unittest.TestCase): self.assertGreaterEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4) def test_mulacc_shl(self): - g1 = UOp.param(0, dtypes.int32.ptr(2)) + g1 = UOp.param(0, dtypes.int32, shape=(2,)) c1 = UOp.const(dtypes.int, 0) c2 = UOp.const(dtypes.int, 1) expr = g1.index(c1) * UOp.const(dtypes.int, 4096) + g1.index(c2) @@ -263,7 +263,7 @@ class TestAssembly(unittest.TestCase): self.assertIn(Ops.MULACC, [x.op for x in uops]) def test_use_cmpeq(self): - g = UOp.param(0, dtypes.uint32.ptr(8)) + g = UOp.param(0, dtypes.uint32, shape=(8,)) c = UOp.const(dtypes.uint, 7) comp = g.index(c).ne(c).ne(True) uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 489db49ece..81d48b7b0a 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -8,7 +8,7 @@ from tinygrad.uop.render import pyrender from tinygrad.uop.spec import type_verify, spec_tensor, spec_program from tinygrad.renderer import Renderer, Estimates from tinygrad.renderer.isa import ISARenderer, IselContext, PreRegAllocContext -from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace +from tinygrad.dtype import dtypes, AddrSpace # import all pattern matchers here from tinygrad.codegen.gpudims import pm_add_gpudims @@ -28,16 +28,6 @@ from tinygrad.helpers import all_same, flatten, argsort, partition from tinygrad.uop.ops import _align_left, _broadcast_shape, identity_element from tinygrad.schedule.rangeify import BufferizeOpts -pm_remove_vec_dtypes = PatternMatcher([ - # rewrite PARAM to non pointer - (UPat((Ops.PARAM, Ops.BUFFER), name="buf"), lambda buf: - buf.replace(dtype=buf.dtype.base, src=(UOp.const(dtypes.int, buf.ptrdtype.size),)) \ - if isinstance(buf.dtype, PtrDType) and not isinstance(buf.dtype, ImageDType) else None), - # remove pointer dtypes from non-PARAM/BUFFER ops - (UPat(GroupOp.All-{Ops.PARAM, Ops.BUFFER}, name="x"), - lambda x: x.replace(dtype=x.dtype.base) if isinstance(x.dtype, PtrDType) else None), -])+pm_clean_up_group_sink - def do_number_param(ctx:list[int], x:UOp): if x.arg.slot != -1: return None ctx[0] += 1 @@ -235,8 +225,7 @@ def merge_reduce_ends(sink:UOp): return sink.substitute(subs) if subs else None def reduce_ranges_to_acc(ctx:ReduceContext, r:UOp): - # TODO: remove this is_ptr when placeholder isn't ptr - acc = UOp.placeholder_like(r, ctx.acc_num, AddrSpace.REG, is_ptr=False) + acc = UOp.placeholder_like(r, ctx.acc_num, AddrSpace.REG) ctx.acc_num += 1 topo = r.src[0].toposort() ended_ranges = flatten([x.ended_ranges for x in topo if x.op is Ops.END]) @@ -267,8 +256,7 @@ pm_move_regs = PatternMatcher([ ]) def add_local_buffer(ctx, x:UOp): - # TODO: remove this is_ptr when placeholder isn't ptr - buf = UOp.placeholder(x.max_shape, x.dtype, slot=next(ctx), addrspace=x.arg.addrspace, is_ptr=False) + buf = UOp.placeholder(x.max_shape, x.dtype, slot=next(ctx), addrspace=x.arg.addrspace) return buf.after(buf.index(*x.src[1:]).store(x.src[0]).end(*x.src[1:]).barrier()) pm_add_local_buffers = PatternMatcher([ @@ -300,9 +288,6 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # do postrange optimization, BEAM or hand_coded_optimizations sink = apply_opts(sink, ren, beam=ast.arg.beam) - # this is new style (TODO: this should all be removed) - sink = graph_rewrite(sink, pm_remove_vec_dtypes, name="transform to new style") - # ** expander (expand_rewrite) ** sink = graph_rewrite(sink, sym+pm_move_where_on_load+pm_flatten_range, name="postopt symbolic") diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index e5270a1b4c..70df4dfd3f 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -31,7 +31,7 @@ def lower_shaped_wmma(ctx, x): name = f"WMMA_{'_'.join(map(str, dims))}_{dtype_in.name}_{dtype_out.name}" wmma_arg = (name, dims, dtype_in, dtype_out, device, threads, tc_upcast_axes, ()) wmma = UOp(Ops.WMMA, dtype_out, tuple(s[u].contract(u) for s, u in upcasts), arg=wmma_arg) - tmp = UOp.placeholder((x.src[2].shape[-1],), dtype_out, slot=next(ctx), addrspace=AddrSpace.REG, is_ptr=False) + tmp = UOp.placeholder((x.src[2].shape[-1],), dtype_out, slot=next(ctx), addrspace=AddrSpace.REG) return tmp.after(UOp.group(*[tmp[e].store(wmma.index(e)) for e in range(x.src[2].shape[-1])])) pm_store_ranges = PatternMatcher([ @@ -444,7 +444,7 @@ def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True): if allow_locals: # handle locals - buf = UOp.placeholder((size,), x.dtype, next(ctx), AddrSpace.LOCAL, is_ptr=False) + buf = UOp.placeholder((size,), x.dtype, next(ctx), AddrSpace.LOCAL) do_store = buf.broadcast(x.src[1].dtype.count).index(idx).store(x.src[0]).end(*rngs) return buf.after(do_store.barrier()) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 819ab147d6..b610026a96 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -465,8 +465,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return self.src[new_srcs[0].arg] return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype if ptr else self.dtype.base.scalar()), (self,)+tuple(new_srcs), **kwargs) def __getitem__(self, idx): - # pointers index into INDEX UOps (scalar lookup); everything else uses the shared mixin view path - if not isinstance(self.dtype, PtrDType): return super(UOp, self).__getitem__(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) idx = self._normalize_indices(list(argfix(idx))) if len(slice_idx:=[i for i,x in enumerate(idx) if isinstance(x, slice)]): # apply SHRINK for slices that aren't the full range @@ -999,20 +999,18 @@ class UOp(RandMixin, metaclass=UOpMetaClass): # *** uop high level syntactic sugar *** @staticmethod - def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL, is_ptr=True): + def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL): if addrspace is AddrSpace.GLOBAL: - # TODO: this should have a shape - ret = UOp(Ops.PARAM, dtype.ptr(prod(shape), addrspace) if is_ptr else dtype, arg=ParamArg(slot, addrspace=addrspace)) + ret = UOp(Ops.PARAM, dtype, src=(shape_to_shape_arg((prod(shape),)),), arg=ParamArg(slot, addrspace=addrspace)) else: assert addrspace in (AddrSpace.LOCAL, AddrSpace.REG) buf_shape = (prod(shape),) + ((dtype.count,) if dtype.count > 1 else ()) - ret = UOp(Ops.BUFFER, dtype.ptr(prod(shape), addrspace) if is_ptr else dtype, - src=(shape_to_shape_arg(buf_shape),), arg=ParamArg(slot, addrspace=addrspace)) + ret = UOp(Ops.BUFFER, dtype, src=(shape_to_shape_arg(buf_shape),), arg=ParamArg(slot, addrspace=addrspace)) if len(shape) > 1: ret = ret.reshape(shape + ((dtype.count,) if addrspace in (AddrSpace.LOCAL, AddrSpace.REG) and dtype.count > 1 else ())) return ret - def placeholder_like(self, slot:int, addrspace=AddrSpace.GLOBAL, is_ptr=True): + def placeholder_like(self, slot:int, addrspace=AddrSpace.GLOBAL): assert all_int(self.shape), "no placeholder-like on symbolic shape" - return UOp.placeholder(self.max_shard_shape, self.dtype, slot, addrspace, is_ptr=is_ptr) + return UOp.placeholder(self.max_shard_shape, self.dtype, slot, addrspace) # set is store+end+after def set(self:UOp, val:UOp|ConstType, end:UOp|tuple[UOp, ...]|list[UOp]=()) -> UOp: