2 global dim

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

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
 

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