From 1e2becfeaef6abed7f8d5b24a9403d02f4d0b90e Mon Sep 17 00:00:00 2001 From: George Hotz Date: Wed, 2 Apr 2025 09:07:27 +0800 Subject: [PATCH] fix pad --- extra/replay_pkl.py | 2 +- tinygrad/codegen/devectorizer.py | 9 +++++++-- tinygrad/codegen/kernel.py | 2 +- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/extra/replay_pkl.py b/extra/replay_pkl.py index 4a88cb3214..619aa05eae 100644 --- a/extra/replay_pkl.py +++ b/extra/replay_pkl.py @@ -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() diff --git a/tinygrad/codegen/devectorizer.py b/tinygrad/codegen/devectorizer.py index 7c0f2b0dd0..84e59159a4 100644 --- a/tinygrad/codegen/devectorizer.py +++ b/tinygrad/codegen/devectorizer.py @@ -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): diff --git a/tinygrad/codegen/kernel.py b/tinygrad/codegen/kernel.py index c5e2e26bc1..f34e426edf 100644 --- a/tinygrad/codegen/kernel.py +++ b/tinygrad/codegen/kernel.py @@ -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))