From 48221d9024814ebd83a67b193f2dc8e725f4668b Mon Sep 17 00:00:00 2001 From: George Hotz Date: Mon, 31 Mar 2025 11:25:12 +0800 Subject: [PATCH] 2 global dim --- tinygrad/codegen/devectorizer.py | 5 +++++ tinygrad/codegen/kernel.py | 7 +++++-- tinygrad/spec.py | 3 ++- 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/tinygrad/codegen/devectorizer.py b/tinygrad/codegen/devectorizer.py index 70b0da8141..f50165e23d 100644 --- a/tinygrad/codegen/devectorizer.py +++ b/tinygrad/codegen/devectorizer.py @@ -99,6 +99,11 @@ load_store_folding = PatternMatcher([ def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None: if (idx:=uop_given_valid(valid, start_idx)) is None: return buf.const_like(0) + if getenv("DEBUG_SIMPLIFY"): + print("****") + print("idx in: ", start_idx.render()) + print("valid: ", valid.render()) + print("simp: ", idx.render()) if not isinstance(buf.dtype, ImageDType): return None if idx is start_idx else buf.index(idx, valid) # wait for it to be image indexed before running simplification diff --git a/tinygrad/codegen/kernel.py b/tinygrad/codegen/kernel.py index 22ea967f15..9ba28cf893 100644 --- a/tinygrad/codegen/kernel.py +++ b/tinygrad/codegen/kernel.py @@ -482,8 +482,11 @@ class Kernel: k.apply_opt(Opt(OptOps.UPCAST, 0, 128)) # make all non first dimensions local - # TODO: fix padding - for i in range(1, k.first_reduce): k.apply_opt(Opt(OptOps.LOCAL, 1, 0)) + if k.full_shape[0]%2 == 1: k.apply_opt(Opt(OptOps.PADTO, 0, 2)) + k.apply_opt(Opt(OptOps.LOCAL, 0, k.full_shape[0]//2)) + for i in range(1, k.first_reduce-1): k.apply_opt(Opt(OptOps.LOCAL, 1, 0)) + + #for i in range(1, k.first_reduce): k.apply_opt(Opt(OptOps.LOCAL, 1, 0)) return self diff --git a/tinygrad/spec.py b/tinygrad/spec.py index 5365498857..670e38e0bd 100644 --- a/tinygrad/spec.py +++ b/tinygrad/spec.py @@ -1,7 +1,7 @@ from typing import cast from tinygrad.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType -from tinygrad.helpers import all_same, dedup, prod +from tinygrad.helpers import all_same, dedup, prod, getenv buffer_spec = PatternMatcher([ (UPat(Ops.UNIQUE, dtypes.void, ()), lambda: True), @@ -46,6 +46,7 @@ tensor_uop_spec = buffer_spec+PatternMatcher([ # ***** uop type spec ***** def validate_index(idx:UOp, mask:UOp|None=None): + if getenv("IGNORE_OOB"): return True # this checks for out of bounds access. it is not complete but should catch some issues if mask is None and not isinstance(idx.dtype, ImageDType): # WEBGPU has a BITCAST in the index. TODO: fix