tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 20:56:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

hacks

Browse Source
This commit is contained in:
geohot
2025-03-27 10:51:41 +08:00
parent 38488ec3b0
commit f27f484621
2 changed files with 9 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/codegen/devectorizer.py
+7 -6
View File
@@ -46,19 +46,19 @@ def expand_index(buf:UOp, vec:UOp, mask:UOp|None=None):
offsets[i+3] = {}
grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])]
for grp in grouped_offsets:
if len(grp) == 32:
for jj in range(grp[0]+32, grp[0]+128):
grp.append(jj)
offsets[jj] = []
load_size = len(grp)
if load_size == 32 and buf.arg not in [0, 3]:
# make loads of 32 be 128, skip the bias and the output (fix this)
load_size = 128
# get the index offset for this element. using [0] is okay, because they are the same
lidx = midx.src[offsets[grp[0]][0]]
if len(grp) > 1: lidx = lidx.cast(ptrdtype.base.vec(len(grp)).ptr(size=ptrdtype.size, local=ptrdtype.local))
if load_size > 1: lidx = lidx.cast(ptrdtype.base.vec(load_size).ptr(size=ptrdtype.size, local=ptrdtype.local))
# set the idxs of the output
for i,g in enumerate(grp):
for oo in offsets[g]: idxs[oo] = global_offset+i
# add this lidx to the CAT
ret.append(lidx)
global_offset += len(grp)
global_offset += load_size
assert None not in idxs, f"some idxs are missing {idxs}"
# this base thing is for image, we want the CAT to be a normal pointer
post_cat = UOp(Ops.PTRCAT, ptrdtype.base.ptr(size=ptrdtype.size, local=ptrdtype.local).vec(vec.dtype.count), tuple(ret))
@@ -271,6 +271,7 @@ pm_render = PatternMatcher([
lambda c: UOp(Ops.VECTORIZE, c.dtype, (UOp.const(c.dtype.scalar(), c.arg),)*c.dtype.vcount) if c.dtype.vcount > 1 else None),
(UPat(Ops.VCONST, name='c'), lambda c: UOp(Ops.VECTORIZE, c.dtype, tuple(UOp.const(c.dtype.scalar(), x) for x in c.arg))),
(UPat(Ops.GEP, name='gep'), lambda gep: UOp(Ops.VECTORIZE, gep.dtype, tuple(gep.src[0].gep(x) for x in gep.arg)) if len(gep.arg) > 1 else None),
(UPat(Ops.GEP, name='gep'), lambda gep: gep.src[0] if gep.src[0].dtype.vcount == 1 and gep.arg == (0,) else None),
(UPat(Ops.VECTORIZE, src=(UPat(name='x'),)), lambda x: x),
# give any loads that are masked an alt value
(UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat())).or_casted(),), name="x"), lambda x: x.replace(src=x.src+(x.const_like(0),))),
tinygrad/ops.py
+2 -1
View File
@@ -383,7 +383,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
if self.op is Ops.VCONST: return UOp.const(self.dtype.scalar(), self.arg[i])
if self.op is Ops.CONST: return UOp.const(self.dtype.scalar(), self.arg)
i = (i,)
if (self.dtype.vcount == len(i) and i == tuple(range(len(i)))) or self.dtype == dtypes.void: return self
#if self.dtype.vcount == 1 and i == (0,): return self
if (self.dtype.count == len(i) and i == tuple(range(len(i)))) or self.dtype == dtypes.void: return self
return UOp(Ops.GEP, self.dtype.scalar().vec(len(i)) if len(i) > 1 else self.dtype.scalar(), (self,), i)
def load(self, *src:UOp, **kwargs): return UOp(Ops.LOAD, src=(self,)+src, **kwargs)
def store(self, *src:UOp, **kwargs): return UOp(Ops.STORE, dtypes.void, (self,)+src, **kwargs)