faster devectorizer with one line

merge devectorize + indexing (#17354 )
* external benchmark schedule in 5 sec (codex slop) * prune * real? * delete
2026-08-01 20:00:55 +00:00 · 2026-08-01 11:40:15 -07:00
1 changed files with 2 additions and 6 deletions
@@ -125,8 +125,7 @@ def do_devectorize(b:UOp):
  # broadcasting needs to be already unpacked, Invalid matches any dtype and shape
  if not all(x.shape == b.shape or x.base.is_invalid for x in b.src): return None
  src = []
-  for idx in itertools.product(*[range(x) for x in b.shape]):
-    idx_c = [UOp.const(i) for i in idx]
+  for idx_c in itertools.product(*[[UOp.const(i) for i in range(x)] for x in b.shape]):
    src.append(b.replace(dtype=None, src=tuple(x.base if x.base.is_invalid else x.index(*idx_c) for x in b.src)))
  return UOp.stack(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src)

@@ -330,10 +329,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
  sink = graph_rewrite(sink, symbolic_simple+pm_expand_broadcast+pm_add_loads, name="*** expand broadcast / add loads")

  # devectorize
-  sink = graph_rewrite(sink, symbolic_simple+devectorizer2, ctx=ren, name="devectorize2")
-
-  # simplify indexing
-  sink = graph_rewrite(sink, indexing_simplify, name="simplify load/store indexing")
+  sink = graph_rewrite(sink, symbolic_simple+devectorizer2+indexing_simplify, ctx=ren, name="devectorize2")

  # some coalescing misses without this
  sink = graph_rewrite(sink, sym, name="early symbolic")
Author	SHA1	Message	Date
geohot	f50d40014c	faster devectorizer with one line	2026-08-01 20:00:55 +00:00
George HotzandGitHub	15c936db01	merge devectorize + indexing (#17354 ) * external benchmark schedule in 5 sec (codex slop) * prune * real? * delete	2026-08-01 11:40:15 -07:00