forked from tinygrad/tinygrad
remove placeholder from is_ptr (#16893)
* remove placeholder from is_ptr * simpler * needed * remove rewriter * fix tests
This commit is contained in:
George Hotz
@@ -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)
|
||||
|
||||
@@ -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")
|
||||
|
||||
|
||||
@@ -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())
|
||||
|
||||
|
||||
+7
-9
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user