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

fix pad

Browse Source
This commit is contained in:
geohot
2025-04-02 09:07:27 +08:00
parent e18cdbcbe2
commit 1e2becfeae
3 changed files with 9 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/replay_pkl.py
+1 -1
View File
@@ -53,7 +53,7 @@ if __name__ == "__main__":
with Context(NOOPT=1):
lower_schedule_item(ScheduleItem(p.ast, ei.bufs)).run()
correct = ei.bufs[0].numpy()
ei.bufs[0].copyin(memoryview(bytearray(b'\x00'*ei.bufs[0].size)))
ei.bufs[0].copyin(memoryview(bytearray(b'\x00'*ei.bufs[0].nbytes)))
GlobalCounters.kernel_count -= 1
#if knum != 1 and not getenv("NOOPT"): k.hand_coded_optimizations()
tinygrad/codegen/devectorizer.py
+7 -2
View File
@@ -15,8 +15,13 @@ def expand_index(buf:UOp, vec:UOp, mask:UOp|None=None):
vectorize_mask = getenv("VECTORIZE_MASK", 0) and buf.arg == 0 and mask is not None
# generate the individual indexes
midx = graph_rewrite(UOp.sink(*[buf.index(vec.gep(i), mask.gep(i) if mask is not None else None) for i in range(vec.dtype.count)]),
symbolic_flat+commutative+load_store_indexing, name=f"index_buf_{buf.arg}")
if vectorize_mask:
# no load_store_indexing if we are doing vectorized mask
midx = graph_rewrite(UOp.sink(*[buf.index(vec.gep(i), mask.gep(i) if mask is not None else None) for i in range(vec.dtype.count)]),
symbolic_flat+commutative, name=f"index_buf_{buf.arg}")
else:
midx = graph_rewrite(UOp.sink(*[buf.index(vec.gep(i), mask.gep(i) if mask is not None else None) for i in range(vec.dtype.count)]),
symbolic_flat+commutative+load_store_indexing, name=f"index_buf_{buf.arg}")
# extract all the relevant offsets
offsets_rootsrc: defaultdict[Any, dict[int, list[int]]] = defaultdict(dict)
for i in range(vec.dtype.count):
tinygrad/codegen/kernel.py
+1 -1
View File
@@ -443,7 +443,7 @@ class Kernel:
if k.full_shape[-3:] == (32,3,3):
# 3x3 dwconv
# kernel 49 is broken
if k.full_shape[-4]%4 != 0 and k.full_shape[-4] != 7: k.apply_opt(Opt(OptOps.PADTO, len(k.full_shape)-4, 4))
if k.full_shape[-4]%4 != 0: k.apply_opt(Opt(OptOps.PADTO, len(k.full_shape)-4, 4))
k.apply_opt(Opt(OptOps.UNROLL, 0, 0))
k.apply_opt(Opt(OptOps.UNROLL, 0, 0))
k.apply_opt(Opt(OptOps.UPCAST, len(k.full_shape)-3, 32))