tune long AMD GEMMs and reject irrelevant store hazards

tinygrad/codegen/opt/gemm.py

@@ -116,9 +116,11 @@ def _gemm_block_m(g:GemmMatch) -> int:
     64 if g.old is not None or _gemm_block_n(g) == 64 else BLOCK_M
 def _gemm_block_n(g:GemmMatch) -> int:
   if g.n == 576 and g.old is None and not g.a_kxm and g.b_kxn: return 192
+  if g.old is not None and g.a_kxm and g.b_kxn and g.m == 256 and g.k >= 65536: return 64
   return 64 if g.n == 64 and g.old is None and not g.a_kxm and not g.b_kxn else BLOCK_N
 def _gemm_block_k(g:GemmMatch) -> int:
-  if g.old is not None and g.m < 512 and g.k % 64 == 0: return 64
+  if g.old is not None and g.a_kxm and g.b_kxn and g.m == 256 and g.k >= 65536: return BLOCK_K
+  if g.old is not None and (g.m < 512 or (g.a_kxm and g.b_kxn and g.k >= 65536)) and g.k % 64 == 0: return 64
   return 96 if _gemm_block_n(g) == 64 and g.k == 288 else BLOCK_K
 def _batched_block_k(g:BatchedGemmMatch) -> int:
   return 128 if g.m == 64 and g.batch >= 96 and g.k % 128 == 0 else BLOCK_K
@@ -128,7 +130,7 @@ def _batched_threads(g:BatchedGemmMatch) -> int:
   return _batched_block_n(g)
 def _gemm_threads(g:GemmMatch) -> int:
   if _gemm_block_n(g) == 192: return 192
-  return 64 if g.old is not None and g.a_kxm and g.b_kxn and g.m >= 512 and g.k >= 65536 else THREADS
+  return 64 if g.old is not None and g.a_kxm and g.b_kxn and g.m >= 256 and g.k >= 65536 else THREADS
 def _render_gemm(g:GemmMatch, name:str) -> str:
   bm, bn_size, bk, threads = _gemm_block_m(g), _gemm_block_n(g), _gemm_block_k(g), _gemm_threads(g)
   waves_m, waves_n = (2, 3) if threads == 192 else (1, 2) if threads == 64 else \

tinygrad/schedule/rangeify.py

@@ -60,9 +60,10 @@ pm_mops = PatternMatcher([
 # 0. do some cleanup rewrites, mostly copied from the old stuff
 
 def fix_store_hazard(target:UOp, src:UOp):
+  base = target.base
+  if src is not base and base not in src.backward_slice: return None
   # PERMUTE and FLIP reorder indices, SHRINK can have overlapping regions when dest is also shrunk
   unsafe = {Ops.PERMUTE, Ops.FLIP} | ({Ops.SHRINK} if target.op_in_backward_slice_with_self(Ops.SHRINK) else set())
-  base = target.base
   reaches_base: dict[UOp, bool] = {}
   for s in src.toposort(gate=lambda s: s.op is not Ops.CONTIGUOUS):
     reaches_base[s] = s is base or any(reaches_base.get(c) for c in s.src)