forked from tinygrad/tinygrad
dtype.vec shapes (#16287)
* dtype.vec shapes * something * Closer * more passes * shape is in spec * fix reduce * image dtype shape correct * lil * use reshape on image * need BUFFER there * remove that test * fix ptx + x86 * fix nir * x86 fix maybe * x86 fixups * x86 fix * don't check that for NOOP
This commit is contained in:
George Hotz
@@ -507,14 +507,6 @@ class TestImageSimplification(unittest.TestCase):
|
||||
self.check(load, "(((lidx1<1)!=True)&(((lidx0+r0)<3)!=True)&((lidx0+r0)<11))",
|
||||
"(lidx2+gidx0*4+lidx1*256+(lidx0*1024+r0*1024)+-3264)", "0")
|
||||
|
||||
class TestUnfoldableImage(unittest.TestCase):
|
||||
def test_unfoldable_becomes_buffer(self):
|
||||
with Context(SPEC=0):
|
||||
lidx = Special("lidx", 2)
|
||||
load = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.PARAM, dtypes.imagef((10, 10, 4)), arg=0).index(lidx, ptr=True), UOp.const(dtypes.float, 0)))
|
||||
res = full_rewrite(load.sink()).src[0]
|
||||
self.assertEqual(res.src[0].src[0].dtype, dtypes.float.ptr(400))
|
||||
|
||||
class TestDropTrueGate(unittest.TestCase):
|
||||
def test_drop_true_gate_on_index(self):
|
||||
# test that INDEX with a constant True valid gets simplified to drop the valid
|
||||
|
||||
@@ -757,6 +757,24 @@ class TestReduceCollapse(unittest.TestCase):
|
||||
# Should become add of two separate reduces
|
||||
self.assertEqual(result.op, Ops.ADD)
|
||||
|
||||
class TestMovementOps(unittest.TestCase):
|
||||
def test_pm_mops_partial_reshape_index_removes_reshape(self):
|
||||
from tinygrad.schedule.rangeify import pm_mops
|
||||
src = UOp.param(0, dtypes.float, shape=(32, 4))
|
||||
r0, r1 = UOp.range(4, 0), UOp.range(8, 1)
|
||||
result = graph_rewrite(src.reshape((4, 8, 4)).index(r0, r1), pm_mops, name="test")
|
||||
self.assertEqual(result.op, Ops.INDEX)
|
||||
self.assertIs(result.src[0], src)
|
||||
self.assertEqual(result.shape, (4,))
|
||||
self.assertNotIn(Ops.RESHAPE, [u.op for u in result.toposort()])
|
||||
|
||||
def test_pm_mops_partial_reshape_index_suffix_mismatch_does_nothing(self):
|
||||
from tinygrad.schedule.rangeify import pm_mops
|
||||
src = UOp.param(0, dtypes.float, shape=(2, 6))
|
||||
result = graph_rewrite(src.reshape((2, 3, 2)).index(UOp.range(2, 0)), pm_mops, name="test")
|
||||
self.assertEqual(result.op, Ops.INDEX)
|
||||
self.assertEqual(result.src[0].op, Ops.RESHAPE)
|
||||
|
||||
class TestLoadStoreFolding(unittest.TestCase):
|
||||
def test_gated_load_gep_preserves_alt(self):
|
||||
"""Test that LOAD(GEP, alt) preserves alt value after rewrite"""
|
||||
|
||||
@@ -72,7 +72,8 @@ def expand_index(ctx, buf:UOp, vec:UOp):
|
||||
elif dropped == best_drop: cands.append((ch, cw, cidx))
|
||||
# and tiebreak with indexing complexity (ie. number of nodes)
|
||||
h, w, _ = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].gep(1).simplify().backward_slice))
|
||||
buf = buf.replace(dtype=(dtypes.imageh if dt.itemsize == 2 else dtypes.imagef)((h, w, 4)))
|
||||
assert buf.op is Ops.RESHAPE
|
||||
buf = buf.src[0].replace(dtype=(dtypes.imageh if dt.itemsize == 2 else dtypes.imagef)((h, w, 4))).flatten()
|
||||
if getenv("UNSAFE_DISABLE_MASK", 0): vec = vec.get_idx()
|
||||
# generate the individual indexes
|
||||
return UOp(Ops.STACK, buf.dtype, tuple(buf.index(vec.gep(i), ptr=True) for i in range(vec.dtype.count)))
|
||||
@@ -133,7 +134,7 @@ def gep_on_store(gep:UOp, st:UOp):
|
||||
return gep.src[0].store(st.gep(new_arg))
|
||||
|
||||
load_store_folding = PatternMatcher([
|
||||
(UPat(Ops.INDEX, src=(UPat(Ops.STACK, src=UPat(GroupOp.Defines).or_after(name="buf")), UPat.var("vec"))), expand_index),
|
||||
(UPat(Ops.INDEX, src=(UPat(Ops.STACK, src=UPat(name="buf")), UPat.var("vec"))), expand_index),
|
||||
(UPat(Ops.STACK, src=UPat(Ops.INDEX), name="midx"), fold_expanded_index),
|
||||
# GEP after LOAD
|
||||
(UPat(Ops.LOAD, src=(UPat(Ops.GEP, name="gep"),), name="ld", allow_any_len=True),
|
||||
@@ -198,7 +199,8 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
|
||||
def get_image_idx(idx:UOp, width:int):
|
||||
x, valid = idx.src[1].get_idx(), idx.src[1].get_valid()
|
||||
idx_x, idx_y = (x // 4) % width, x // (4*width)
|
||||
return idx.replace(src=(idx.src[0], idx_y.valid(valid), idx_x.valid(valid)))
|
||||
assert idx.src[0].op is Ops.RESHAPE, "image idx must be on reshape"
|
||||
return idx.replace(src=(idx.src[0].src[0], idx_y.valid(valid), idx_x.valid(valid)))
|
||||
|
||||
def image_fixup(ls:UOp):
|
||||
# normal image load or store, with the CAST from expand_index
|
||||
@@ -209,7 +211,8 @@ def image_fixup(ls:UOp):
|
||||
# this is an unprocessed image without a cast, we should just make it a buffer
|
||||
if isinstance(dt, ImageDType) and len(ls.src[0].src) == 2:
|
||||
off = ls.src[0].src[1]
|
||||
idx = ls.src[0].src[0].replace(dtype=(new_dt:=dtypes.half if dt.itemsize == 2 else dtypes.float).ptr(dt.size)).index(off)
|
||||
assert ls.src[0].src[0].op is Ops.RESHAPE, "image idx must be on reshape"
|
||||
idx = ls.src[0].src[0].src[0].replace(dtype=(new_dt:=dtypes.half if dt.itemsize == 2 else dtypes.float).ptr(dt.size)).index(off)
|
||||
return ls.replace(src=(idx,), dtype=new_dt).cast(dtypes.float) if ls.op is Ops.LOAD else ls.replace(src=(idx, ls.src[1].cast(new_dt)))
|
||||
|
||||
correct_load_store = PatternMatcher([
|
||||
@@ -372,7 +375,7 @@ pm_imageh_store = PatternMatcher([
|
||||
def make_image(ctx, ls, buf, off):
|
||||
if (vcount:=buf.dtype.vcount) != 1: buf = buf.src[0]
|
||||
if buf.op == Ops.PARAM and not isinstance(dt:=buf.dtype, ImageDType) and (dims:=ImageDType.valid_dims(dt, ctx)):
|
||||
buf = buf.replace(dtype=(dtypes.imageh if dt.base == dtypes.half else dtypes.imagef)((*dims[0], 4)))
|
||||
buf = buf.replace(dtype=(dtypes.imageh if dt.base == dtypes.half else dtypes.imagef)((*dims[0], 4))).flatten()
|
||||
if vcount != 1: buf = UOp.vectorize(*([buf] * vcount))
|
||||
if ls.op is Ops.LOAD: return ls.replace(src=(buf.index(off, ptr=True),), dtype=dtypes.float.vec(ls.dtype.vcount)).cast(dt.base)
|
||||
return buf.index(off, ptr=True).store(pm_imageh_store.rewrite(ls.src[1]) if dt.base == dtypes.half else ls.src[1])
|
||||
|
||||
@@ -268,6 +268,8 @@ ALLOW_TF32 = ContextVar("ALLOW_TF32", 0)
|
||||
SCACHE = ContextVar("SCACHE", 1)
|
||||
# allow use of atomics for embedding backward
|
||||
USE_ATOMICS = ContextVar("USE_ATOMICS", 0)
|
||||
# don't allow broadcast
|
||||
DISALLOW_BROADCAST = ContextVar("DISALLOW_BROADCAST", 1)
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Metadata:
|
||||
|
||||
@@ -139,8 +139,6 @@ extra_matcher = PatternMatcher([
|
||||
(UPat.var('x', dtypes.bool).ne(UPat.var('y')), lambda x,y: x^y),
|
||||
(UPat.var('x', dtypes.bool).alu(Ops.CMPEQ, UPat.var('y')), lambda x,y: (x^y)^True),
|
||||
(UPat.var('x', dtypes.bool)<UPat.var('y'), lambda x,y: (x^True)&y),
|
||||
# cast to pointer is a noop
|
||||
(UPat.var("y").cast(name="x"), lambda y,x: y if isinstance(x.dtype, PtrDType) or y.dtype == dtypes.void else None),
|
||||
# can't cast from float16 to ints/float64 directly and vice versa
|
||||
(UPat.var("y", dtypes.float16).cast((dtypes.float64,)+dtypes.ints, name="x"), lambda y,x: y.cast(dtypes.float32).cast(x.dtype)),
|
||||
(UPat.var("y", (dtypes.float64,)+dtypes.ints).cast(dtypes.float16, name="x"), lambda y,x: y.cast(dtypes.float32).cast(x.dtype)),
|
||||
@@ -175,6 +173,17 @@ extra_matcher = PatternMatcher([
|
||||
|
||||
# ***** X86 pre instruction selection *****
|
||||
|
||||
def gated_load(ctx, base:UOp, idx:UOp, cast:UOp, alt:UOp, gate:UOp, x:UOp):
|
||||
local = UOp(Ops.DEFINE_LOCAL, base.dtype.base.ptr(x.dtype.count, AddrSpace.LOCAL), arg=next(ctx))
|
||||
local_idx = local.index(UOp.const(dtypes.int32, 0), ptr=True)
|
||||
ptr = gate.where(base.index(idx, ptr=True), local_idx).after((local_idx if x.dtype.count == 1 else local).store(alt))
|
||||
return ptr.cast(cast.dtype).load(dtype=x.dtype)
|
||||
|
||||
def gated_store(base:UOp, idx:UOp, cast:UOp, gate:UOp, val:UOp):
|
||||
local = UOp(Ops.DEFINE_LOCAL, base.dtype.base.ptr(val.dtype.count, AddrSpace.LOCAL), arg=-1)
|
||||
ptr = gate.where(base.index(idx, ptr=True), local.index(UOp.const(dtypes.int32, 0), ptr=True))
|
||||
return ptr.cast(cast.dtype).store(val)
|
||||
|
||||
# these must be done in a separate matcher because they violate the spec
|
||||
pre_isel_matcher = PatternMatcher([
|
||||
# zero extending scalar 32bit int is a noop
|
||||
@@ -194,12 +203,8 @@ pre_isel_matcher = PatternMatcher([
|
||||
(UPat(Ops.STACK, src=(UPat.var("y"),), allow_any_len=True, name="x"),
|
||||
lambda y,x: UOp(Ops.NOOP, x.dtype, y.src) if all(s.op is Ops.GEP and s.src == y.src and s.arg[0] == i for i,s in enumerate(x.src)) else None),
|
||||
# gated load/store become a conditional move on the index, the load/store are unconditional
|
||||
(UPat.var("base").index(UPat.var("idx")).load(UPat.var("alt"), UPat.var("gate"), name="x"), lambda ctx,base,idx,gate,alt,x:
|
||||
gate.where(base.index(idx, ptr=True), (l:=UOp(Ops.DEFINE_LOCAL, base.dtype.base.ptr(x.dtype.count, AddrSpace.LOCAL), arg=next(ctx))
|
||||
.index(UOp.const(dtypes.int32, 0), ptr=True)).after(l.store(alt))).load(dtype=x.dtype)),
|
||||
(UPat.var("base").index(UPat.var("idx")).store(UPat.var("val"), UPat.var("gate")), lambda base,idx,gate,val:
|
||||
gate.where(base.index(idx, ptr=True), UOp(Ops.DEFINE_LOCAL, base.dtype.base.ptr(val.dtype.count, AddrSpace.LOCAL), arg=-1)
|
||||
.index(UOp.const(dtypes.int32, 0), ptr=True)).store(val)),
|
||||
(UPat.var("base").index(UPat.var("idx")).or_casted(name="cast").load(UPat.var("alt"), UPat.var("gate"), name="x"), gated_load),
|
||||
(UPat.var("base").index(UPat.var("idx")).or_casted(name="cast").store(UPat.var("val"), UPat.var("gate")), gated_store),
|
||||
# TODO: remove this once we allow all flag producing ops in cmove
|
||||
# if gate in scalar int cmove is not a comparison need to add one to set the flag
|
||||
(UPat.var("m", dtypes.bool).where(UPat.var("a"), UPat.var("b")),
|
||||
@@ -289,7 +294,7 @@ def vpins(x:UOp) -> UOp:
|
||||
# inserts scalar int in xmm0 into all lanes of xmm1
|
||||
def vpbroadcast(ctx:IselContext, x:UOp, y:UOp) -> UOp:
|
||||
n = x.ins({1: X86Ops.VPBROADCASTB, 2: X86Ops.VPBROADCASTW, 4: X86Ops.VPBROADCASTD, 8: X86Ops.VPBROADCASTQ}[y.dtype.itemsize], src=(y,))
|
||||
if y.op is Ops.LOAD and is_foldable(ctx, n, y): return n
|
||||
if y.op is Ops.LOAD and len(y.src) == 1 and is_foldable(ctx, n, y): return n
|
||||
# if there isn't a load we can fold we need to move y from gpr to xmm
|
||||
# this is hacky but required because int.vec(1) isn't supported
|
||||
y = y if y.dtype.itemsize > 1 else y.cast(dtypes.int16)
|
||||
@@ -359,6 +364,8 @@ dt_128bit = tuple(dt.vec(l) for dt in dts for l in [16,8,4,2,1] if l*dt.itemsize
|
||||
|
||||
isel_matcher = PatternMatcher([
|
||||
# **** Op -> Op ****
|
||||
# cast to pointer is a noop
|
||||
(UPat.var("y").cast(name="x"), lambda y,x: y if isinstance(x.dtype, PtrDType) or y.dtype == dtypes.void else None),
|
||||
# float gep(0) is a noop as it just moves the 0th element from one xmm register to another
|
||||
# this is done here to not interfere with shuffles
|
||||
(UPat(dtype=dtypes.floats).gep(0, name="x"), lambda x: x.replace(op=Ops.NOOP, arg=None)),
|
||||
@@ -551,12 +558,12 @@ isel_matcher = PatternMatcher([
|
||||
(UPat(Ops.COPY, dt_64bit, name="x"), lambda x: x.ins(X86Ops.VMOVSD)),
|
||||
(UPat(Ops.COPY, dt_32bit+dt_16bit, name="x"), lambda x: x.ins(X86Ops.VMOVSS)),
|
||||
(UPat(Ops.COPY, dtypes.ints+(dtypes.bool,), name="x"), lambda x: x.ins(X86Ops.MOV)),
|
||||
(UPat(Ops.LOAD, dt_128bit, name="x"), lambda x: x.ins(X86Ops.VMOVUPS, src=fold_address(x.src[0]))),
|
||||
(UPat(Ops.LOAD, dt_64bit, name="x"), lambda x: x.ins(X86Ops.VMOVSD, src=fold_address(x.src[0]))),
|
||||
(UPat(Ops.LOAD, dt_32bit, name="x"), lambda x: x.ins(X86Ops.VMOVSS, src=fold_address(x.src[0]))),
|
||||
(UPat(Ops.LOAD, dt_16bit, name="x"), lambda x:
|
||||
x.ins(X86Ops.VPINSRW, src=(def_reg(x.dtype, x.tag),) + fold_address(x.src[0]) + (imm(dtypes.uint8, 0),))),
|
||||
(UPat(Ops.LOAD, dtypes.ints+(dtypes.bool,), name="x"), lambda x: x.ins(X86Ops.MOV, src=fold_address(x.src[0]))),
|
||||
(UPat(Ops.LOAD, dt_128bit, src=(UPat(name="a"),), name="x"), lambda x,a: x.ins(X86Ops.VMOVUPS, src=fold_address(a))),
|
||||
(UPat(Ops.LOAD, dt_64bit, src=(UPat(name="a"),), name="x"), lambda x,a: x.ins(X86Ops.VMOVSD, src=fold_address(a))),
|
||||
(UPat(Ops.LOAD, dt_32bit, src=(UPat(name="a"),), name="x"), lambda x,a: x.ins(X86Ops.VMOVSS, src=fold_address(a))),
|
||||
(UPat(Ops.LOAD, dt_16bit, src=(UPat(name="a"),), name="x"), lambda x,a:
|
||||
x.ins(X86Ops.VPINSRW, src=(def_reg(x.dtype, x.tag),) + fold_address(a) + (imm(dtypes.uint8, 0),))),
|
||||
(UPat(Ops.LOAD, dtypes.ints+(dtypes.bool,), src=(UPat(name="a"),), name="x"), lambda x,a: x.ins(X86Ops.MOV, src=fold_address(a))),
|
||||
(UPat.var("a").store(UPat.var("b", dt_128bit), name="x"), lambda a,b,x: x.ins(X86Ops.VMOVUPSm, src=fold_address(a) + (b,))),
|
||||
(UPat.var("a").store(UPat.var("b", dt_64bit), name="x"), lambda a,b,x: x.ins(X86Ops.VMOVSDm, src=fold_address(a) + (b,))),
|
||||
(UPat.var("a").store(UPat.var("b", dt_32bit), name="x"), lambda a,b,x: x.ins(X86Ops.VMOVSSm, src=fold_address(a) + (b,))),
|
||||
@@ -565,12 +572,12 @@ isel_matcher = PatternMatcher([
|
||||
x.ins(X86Ops.MOVm, src=fold_address(a) + (b,)) if (i:=to_imm(b)) is None else x.ins(X86Ops.MOVi, src=fold_address(a) + (i,))),
|
||||
# **** X86Op -> X86Op ****
|
||||
# fold loads into X86Ops that allow it, if beneficial
|
||||
(UPat(Ops.INS, src=(UPat(Ops.LOAD, name="y"),), allow_any_len=True, name="x"), lambda ctx,y,x:
|
||||
x.replace(src=fold_address(y.src[0]) + x.src[1:]) if x.arg in X86GroupOp.ReadMem1st and is_foldable(ctx, x, y) else None),
|
||||
(UPat(Ops.INS, src=(UPat(), UPat(Ops.LOAD, name="y")), allow_any_len=True, name="x"), lambda ctx,y,x:
|
||||
x.replace(src=x.src[:1] + fold_address(y.src[0]) + x.src[2:]) if x.arg in X86GroupOp.ReadMem2nd and is_foldable(ctx, x, y) else None),
|
||||
(UPat(Ops.INS, src=(UPat(), UPat(), UPat(Ops.LOAD, name="y")), allow_any_len=True, name="x"), lambda ctx,y,x:
|
||||
x.replace(src=x.src[:2] + fold_address(y.src[0]) + x.src[3:]) if x.arg in X86GroupOp.ReadMem3rd and is_foldable(ctx, x, y) else None),
|
||||
(UPat(Ops.INS, src=(UPat(Ops.LOAD, src=(UPat(name="a"),), name="y"),), allow_any_len=True, name="x"), lambda ctx,y,a,x:
|
||||
x.replace(src=fold_address(a) + x.src[1:]) if x.arg in X86GroupOp.ReadMem1st and is_foldable(ctx, x, y) else None),
|
||||
(UPat(Ops.INS, src=(UPat(), UPat(Ops.LOAD, src=(UPat(name="a"),), name="y")), allow_any_len=True, name="x"), lambda ctx,y,a,x:
|
||||
x.replace(src=x.src[:1] + fold_address(a) + x.src[2:]) if x.arg in X86GroupOp.ReadMem2nd and is_foldable(ctx, x, y) else None),
|
||||
(UPat(Ops.INS, src=(UPat(), UPat(), UPat(Ops.LOAD, src=(UPat(name="a"),), name="y")), allow_any_len=True, name="x"), lambda ctx,y,a,x:
|
||||
x.replace(src=x.src[:2] + fold_address(a) + x.src[3:]) if x.arg in X86GroupOp.ReadMem3rd and is_foldable(ctx, x, y) else None),
|
||||
# allocate virtual registers
|
||||
(UPat((Ops.INS, Ops.DEFINE_REG, Ops.DEFINE_LOCAL), name="x"), alloc_vregs),
|
||||
])
|
||||
|
||||
@@ -138,7 +138,6 @@ class NIRRenderer(Renderer):
|
||||
# load/store use pointer arithmetic, and the cast does nothing. NOTE: this doesn't apply to image indexing cause it's 1-D
|
||||
(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off")), name="x"), lambda x,buf,off: x.replace(
|
||||
src=(buf,off.cast(dtypes.long))) if buf.dtype.addrspace != AddrSpace.REG and off.op not in (Ops.CAST, Ops.STACK) else None),
|
||||
(UPat(Ops.CAST, name="x"), lambda x: x.src[0] if isinstance(x.dtype, PtrDType) or x.src[0].dtype == dtypes.void else None),
|
||||
# images need index to be int for nir
|
||||
(UPat.var("buf").index(UPat.var("idx_y"), UPat.var("idx_x")),
|
||||
lambda buf,idx_y,idx_x: buf.index(idx_y.cast(dtypes.int), idx_x.cast(dtypes.int))),
|
||||
@@ -149,12 +148,12 @@ class NIRRenderer(Renderer):
|
||||
(UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 8)),
|
||||
(UPat(Ops.DEFINE_VAR, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 4)),
|
||||
(UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, {'g':ngid, 'l':nlid, 'i': nid}[x.arg[0]](ctx.b), int(x.arg[-1]))),
|
||||
(UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"),UPat.var("off"))), UPat.var("val")), allow_any_len=True),
|
||||
(UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"),UPat.var("off"))).or_casted(), UPat.var("val")), allow_any_len=True),
|
||||
lambda ctx,buf,off,val: nstore(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), ctx.r[val], val.dtype)),
|
||||
(UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))), UPat.var("alt"), UPat.var("gate")), name="x"),
|
||||
(UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))).or_casted(), UPat.var("alt"), UPat.var("gate")), name="x"),
|
||||
lambda ctx,x,buf,off,alt,gate: if_phi(ctx.b, ctx.r[gate],
|
||||
lambda: nload(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype, ctx.r[gate]), x.dtype), lambda: ctx.r[alt])),
|
||||
(UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))),), name="x"),
|
||||
(UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))).or_casted(),), name="x"),
|
||||
lambda ctx,x,buf,off: nload(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), x.dtype)),
|
||||
(UPat(Ops.STACK, name="x"), lambda ctx,x: nalu(ctx.b, f"vec{x.dtype.count}", *[ctx.r[src] for src in x.src])),
|
||||
(UPat(GroupOp.ALU, name="x"), lambda ctx,x: nalu(ctx.b, aop[x.src[0].dtype.scalar()][x.op], *[ctx.r[src] for src in x.src])),
|
||||
@@ -191,6 +190,7 @@ class NIRRenderer(Renderer):
|
||||
|
||||
for u in uops:
|
||||
if u.op in {Ops.NOOP, Ops.GROUP, Ops.INDEX}: pass
|
||||
elif u.op is Ops.CAST and isinstance(u.dtype, PtrDType): pass
|
||||
elif u.op is Ops.AFTER:
|
||||
self.r[u] = self.r[u.src[0]]
|
||||
elif u.op == Ops.SINK:
|
||||
|
||||
@@ -54,8 +54,6 @@ ptx_matcher = PatternMatcher([
|
||||
(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx")), name="op"), lambda buf,idx,op:
|
||||
UOp(Ops.INDEX, dtype=dtypes.int64, src=(buf, buf.cast(dtypes.int64)+idx.cast(dtypes.int64)*buf.dtype.itemsize)+op.src[2:]) \
|
||||
if op.dtype != dtypes.int64 and buf.dtype.addrspace != AddrSpace.REG else None),
|
||||
# load/store use pointer arithmetic, and the cast does nothing
|
||||
(UPat(Ops.CAST, name="x"), lambda x: x.src[0] if isinstance(x.dtype, PtrDType) or x.src[0].dtype == dtypes.void else None),
|
||||
# ptx shr and shl instructions require y to be uint
|
||||
(UPat.var("x") << UPat.var("y"), lambda x,y: UOp(Ops.SHL, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
|
||||
(UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
|
||||
@@ -102,18 +100,18 @@ string_rewrite = PatternMatcher([
|
||||
(UPat(Ops.CAST, name="x", src=(UPat.var("a"),)),
|
||||
lambda ctx, x, a: f"cvt{modifier(x.dtype, a.dtype)}.{ctx.cast_types[x.dtype]}.{ctx.cast_types[a.dtype]} {ctx.r[x]}, {ctx.r[a]};"),
|
||||
# store / gated load / load
|
||||
(UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))), UPat.var("var")), allow_any_len=True),
|
||||
(UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))).or_casted(), UPat.var("var")), allow_any_len=True),
|
||||
lambda ctx, loc, var, buf: f"st.{mem_type(buf)}" + \
|
||||
f"{f'.v{cnt}' if ((cnt:=var.dtype.count)>1) else ''}.{ctx.mem_types[var.dtype.scalar()]} " + \
|
||||
f"[{ctx.r[loc]}+0], {('{' + ', '.join(ctx.r[var]) + '}') if var.dtype.count > 1 else ctx.r[var]};"),
|
||||
(UPat(Ops.LOAD, name="x", src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))), UPat.var("alt"), UPat.var("gate")), allow_any_len=True),
|
||||
(UPat(Ops.LOAD, name="x", src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))).or_casted(), UPat.var("alt"), UPat.var("gate"))),
|
||||
lambda ctx, x, loc, alt, gate, buf: flatten([
|
||||
[f"mov.{ctx.mem_types[x.dtype.scalar()]} {v}, {render_val(0, x.dtype.scalar())};" for v in ctx.r[x]],
|
||||
[f"@{ctx.r[gate]} ld.{mem_type(buf)}.v{x.dtype.count}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];"]
|
||||
]) if alt.dtype.count > 1 else [
|
||||
f"@{ctx.r[gate]} ld.{mem_type(buf)}.{ctx.mem_types[x.dtype.scalar()]} {ctx.r[x]}, [{ctx.r[loc]}+0];",
|
||||
f"@!{ctx.r[gate]} mov.b{ctx.types[x.dtype.scalar()][1:]} {ctx.r[x]}, {ctx.r[alt]};"]),
|
||||
(UPat(Ops.LOAD, name="x", src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))),), allow_any_len=True),
|
||||
(UPat(Ops.LOAD, name="x", src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))).or_casted(),)),
|
||||
lambda ctx, x, loc, buf: f"ld.{mem_type(buf)}.v{x.dtype.count}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];" \
|
||||
if x.dtype.count > 1 else f"ld.{mem_type(buf)}.{ctx.mem_types[x.dtype]} {ctx.r[x]}, [{ctx.r[loc]}+0];"),
|
||||
# simple
|
||||
@@ -187,6 +185,7 @@ class PTXRenderer(Renderer):
|
||||
name = "test"
|
||||
for u in uops:
|
||||
if u.op in {Ops.NOOP, Ops.GROUP}: continue
|
||||
if u.op is Ops.CAST and isinstance(u.dtype, PtrDType): continue
|
||||
if u.op is Ops.AFTER:
|
||||
self.r[u] = self.r[u.src[0]]
|
||||
continue
|
||||
@@ -219,7 +218,6 @@ class PTXRenderer(Renderer):
|
||||
if u.op is Ops.SPECIAL: r[u] = "%" + u.arg
|
||||
elif u.op is Ops.DEFINE_VAR: bufs.append((u.expr, u.dtype))
|
||||
elif u.op is Ops.LOAD:
|
||||
assert u.src[0].dtype == dtypes.int64, "load isn't int64"
|
||||
r[u] = [ssa('val', dtype=self.types[u.dtype.scalar()]) for _ in range(u.dtype.count)] if u.dtype.count > 1 else ssa('val', u)
|
||||
elif u.op is Ops.PARAM: bufs.append((f"data{u.arg}", u.dtype))
|
||||
elif u.op is Ops.WMMA:
|
||||
|
||||
@@ -58,7 +58,7 @@ def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
|
||||
new_srcs = []
|
||||
for s in x.src:
|
||||
new_src = s
|
||||
if s.op in {Ops.PARAM, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
|
||||
if s.op in {Ops.PARAM, Ops.BUFFER, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
|
||||
if x in ctx.range_map: new_src = new_src.index(*ctx.range_map[x][0])
|
||||
elif s in ctx.realize_map:
|
||||
realized_ranges = ctx.realize_map[s]
|
||||
|
||||
@@ -64,10 +64,19 @@ pm_fold_moved_after = PatternMatcher([
|
||||
])
|
||||
|
||||
# movement op on INDEX as a PatternMatcher
|
||||
def _mop_index(r:UOp, idx:UOp):
|
||||
idxs = idx.src[1:]
|
||||
if len(idxs) == len(r.shape):
|
||||
return r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idxs), dtype=idx.dtype, arg=idx.arg)
|
||||
if r.op is Ops.RESHAPE:
|
||||
src_prefix = len(r.src[0].shape) - len(r.shape[len(idxs):])
|
||||
if src_prefix >= 0 and r.src[0].shape[src_prefix:] == r.shape[len(idxs):]:
|
||||
if src_prefix == 0: return r.src[0] if r.src[0].dtype == idx.dtype else None
|
||||
ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape[:src_prefix], r.shape[:len(idxs)], idxs), dtype=idx.dtype, arg=idx.arg)
|
||||
return ret if ret.shape == idx.shape else None
|
||||
|
||||
pm_mops = PatternMatcher([
|
||||
(UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
|
||||
lambda r,idx: r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idx.src[1:]), dtype=idx.dtype, arg=idx.arg)
|
||||
if len(idx.src[1:]) == len(r.shape) else None),
|
||||
(UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), _mop_index),
|
||||
# move movement ops and INDEX after AFTER (but not when AFTER has a raw STORE with shaped children — from replace_contig_with_store_after)
|
||||
(UPat(GroupOp.Movement|{Ops.INDEX}, name="r").after(name="a", allow_any_len=True),
|
||||
lambda r,a: UOp(r.op, r.dtype, (a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], r.arg)),
|
||||
@@ -381,7 +390,7 @@ pm_limit_bufs = PatternMatcher([(UPat(set.union(GroupOp.Binary, GroupOp.Ternary)
|
||||
# NOTE: this has been fixed up a bit
|
||||
|
||||
def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True):
|
||||
size = prod(x.shape)
|
||||
size = prod(x.shape) // x.dtype.count
|
||||
rngs = sorted(idx.ranges, key=lambda x: x.arg)
|
||||
assert size > 0 and isinstance(size, int), f"no zero sized or symbolic sized buffers {size}"
|
||||
|
||||
|
||||
+31
-11
@@ -9,7 +9,7 @@ from tinygrad.dtype import ConstFloat, PyConst, storage_fmt_for_dtype, to_storag
|
||||
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
|
||||
from tinygrad.helpers import PROFILE, dedup, cdiv, cmod, floordiv, floormod, diskcache_put, to_function_name, cpu_profile, TracingKey
|
||||
from tinygrad.helpers import VIZ, SPEC, CAPTURE_PROCESS_REPLAY
|
||||
from tinygrad.helpers import VIZ, SPEC, CAPTURE_PROCESS_REPLAY, DISALLOW_BROADCAST
|
||||
from tinygrad.helpers import colored, ansilen, printable
|
||||
if TYPE_CHECKING:
|
||||
from tinygrad.renderer import Estimates
|
||||
@@ -213,7 +213,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
|
||||
match self.op:
|
||||
# late ops don't have shape
|
||||
case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
|
||||
Ops.CONTRACT | Ops.SINK | Ops.END | Ops.REWRITE_ERROR | Ops.PTRCAT | Ops.ENDIF | \
|
||||
Ops.SINK | Ops.END | Ops.REWRITE_ERROR | Ops.PTRCAT | Ops.ENDIF | \
|
||||
Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.INS | Ops.TUPLE | Ops.CALL | Ops.FUNCTION:
|
||||
return None
|
||||
|
||||
@@ -237,6 +237,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
|
||||
return inner_shape
|
||||
|
||||
case Ops.CAST:
|
||||
# if it has a vec dtype, set the shape
|
||||
if self.dtype.count > 1: return (self.dtype.count,)
|
||||
# when PTX casts from ptr to non ptr, remove the shape of the buffer
|
||||
if isinstance(self.src[0].dtype, PtrDType) and not isinstance(self.src[0].dtype, ImageDType) and not isinstance(self.dtype, PtrDType):
|
||||
return ()
|
||||
@@ -246,12 +248,22 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
|
||||
for s in self.src[1:]: shp.extend(list(s.shape))
|
||||
return tuple(shp) + self.src[0].shape[len(self.src[1:]):]
|
||||
|
||||
# TODO: these should have the shape of the dtype.count
|
||||
case Ops.CONST | Ops.DEFINE_VAR: return ()
|
||||
case Ops.GEP | Ops.STACK | Ops.VCAT | Ops.GETADDR: return ()
|
||||
case Ops.GEP:
|
||||
return (len(self.arg),) if len(self.arg) > 1 else ()
|
||||
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
|
||||
# TODO: contract and unroll should be deleted
|
||||
case Ops.CONST | Ops.DEFINE_VAR | Ops.CONTRACT | Ops.UNROLL | Ops.VCAT:
|
||||
return (self.dtype.count,) if self.dtype.count > 1 else ()
|
||||
|
||||
# some ops init the shape
|
||||
case Ops.BIND | Ops.RANGE | Ops.SPECIAL | Ops.UNROLL: return ()
|
||||
case Ops.GETADDR: return ()
|
||||
case Ops.BIND | Ops.RANGE | Ops.SPECIAL: return ()
|
||||
case Ops.BINARY: return (len(self.arg),)
|
||||
case Ops.BUFFER: return (self.arg,)
|
||||
case Ops.BUFFER_VIEW:
|
||||
@@ -259,10 +271,15 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
|
||||
if self.src[0].op is Ops.INDEX: return ()
|
||||
return (self.arg[0],)
|
||||
case Ops.CUSTOM_FUNCTION: return None
|
||||
case Ops.STAGE: return tuple([int(r.vmax+1) for r in self.src[1:]])
|
||||
case Ops.DEFINE_LOCAL: return (self.ptrdtype.size,)
|
||||
case Ops.DEFINE_REG: return (self.ptrdtype.size,) if isinstance(self.dtype, PtrDType) else ()
|
||||
case Ops.STAGE:
|
||||
# STAGE adds the existing shape to the front, opposite of INDEX
|
||||
return tuple([int(r.vmax+1) for r in self.src[1:]])+self.src[0].shape
|
||||
case Ops.DEFINE_LOCAL | Ops.DEFINE_REG:
|
||||
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.PARAM:
|
||||
if isinstance(self.dtype, ImageDType): return self.dtype.shape
|
||||
if isinstance(self.dtype, PtrDType): return (self.ptrdtype.size,)
|
||||
# NOTE: copied from marg
|
||||
if len(self.src) >= 1: return tuple(self.src[0].sgep(i) for i in range(self.src[0].dtype.count))
|
||||
@@ -277,7 +294,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
|
||||
return self.src[0]._shape
|
||||
# REDUCE with empty axis is passthrough (lowered form)
|
||||
case Ops.REDUCE if len(self.arg[1]) == 0:
|
||||
return self.src[0]._shape
|
||||
# these can mismatch if there's a horizonal reduce
|
||||
return (self.dtype.count,) if self.dtype.count > 1 else ()
|
||||
|
||||
# TODO: disallow shape changing bitcast
|
||||
case Ops.BITCAST:
|
||||
@@ -334,6 +352,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
|
||||
if self.op in GroupOp.Broadcastable:
|
||||
input_shapes = [x._shape for x in self.src]
|
||||
assert len(self.src) > 0 and all(x is not None for x in input_shapes), f"None input shape not supported for {self.op}"
|
||||
if DISALLOW_BROADCAST and not all_same(input_shapes):
|
||||
raise RuntimeError(f"shape mismatch at {self.op}: {input_shapes} {[x.op for x in self.src]}")
|
||||
# broadcasting lives in _shape property now
|
||||
return _broadcast_shape(*input_shapes)
|
||||
|
||||
@@ -513,7 +533,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
|
||||
ret = UOp.vectorize(*stk)
|
||||
else:
|
||||
ret = UOp(Ops.CONST, dtype, arg=dtype.const(b), src=(UOp(Ops.DEVICE, arg=device),) if device is not None else ())
|
||||
return ret.reshape((1,)*len(shape)).expand(shape) if shape is not None and ret.shape != shape else ret
|
||||
return ret.reshape((1,)*len(shape)).expand(shape) if shape is not None and shape != () and ret.shape != shape else ret
|
||||
@staticmethod
|
||||
def unique_const(fill_value:ConstType, dtype:DTypeLike|None=None, device:str|tuple[str, ...]|None=None, # type: ignore[override]
|
||||
shape:tuple[sint, ...]|None=None, unique=True):
|
||||
|
||||
@@ -201,9 +201,16 @@ spec_program = PatternMatcher([
|
||||
# weakint is not allowed in programs
|
||||
(UPat(GroupOp.All, dtypes.weakint), lambda: False),
|
||||
|
||||
# movement ops are not allowed in programs
|
||||
(UPat(GroupOp.Movement), lambda: False),
|
||||
|
||||
# Invalid is not allowed in program
|
||||
(UPat(Ops.CONST, arg=Invalid), lambda: False),
|
||||
|
||||
# shape of uop must match dtype.count in program
|
||||
(UPat(GroupOp.All-{Ops.INS, Ops.NOOP}, name="x"),
|
||||
lambda x: False if x.dtype.count > 1 and (x.dtype.count,) != x.shape else None),
|
||||
|
||||
# STACK/GEP in program. TODO: this should match Tensor
|
||||
(UPat(Ops.STACK, name="x"), lambda x: len(x.src)>1 and len(x.src) == x.dtype.vcount and all(x.dtype == y.dtype.vec(len(x.src)) for y in x.src)),
|
||||
(UPat(Ops.GEP, src=(UPat.var("src"),), name="gep"), lambda gep,src: gep.dtype == src.dtype.scalar()),
|
||||
|
||||
@@ -119,8 +119,6 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]:
|
||||
if u.op in {Ops.DEVICE, Ops.CONST, Ops.UNIQUE, Ops.LUNIQUE} and u is not x: excluded.add(u)
|
||||
if u.op is Ops.CONST and len(u.src) and u.src[0].op in {Ops.UNIQUE, Ops.LUNIQUE}: excluded.remove(u)
|
||||
if u.op is Ops.STACK and len(u.src) == 0: excluded.add(u)
|
||||
# exclude RESHAPE/EXPAND that only serve to broadcast a CONST
|
||||
if u.op in {Ops.RESHAPE, Ops.EXPAND} and len(u.src) >= 1 and u.src[0] in excluded and u is not x: excluded.add(u)
|
||||
for u in toposort:
|
||||
if u in excluded: continue
|
||||
argst = codecs.decode(str(u.arg), "unicode_escape")
|
||||
|
||||
Reference in New Issue
Block a user