colors

.github/workflows/autogen.yml

@@ -2,7 +2,7 @@ name: Autogen
 env:
   # increment this when downloads substantially change to avoid the internet
   DOWNLOAD_CACHE_VERSION: '12'
-  PYTHON_CACHE_VERSION: '4'
+  PYTHON_CACHE_VERSION: '3'
   APT_CACHE_VERSION: '1'
   BUILD_CACHE_VERSION: '1'
   CAPTURE_PROCESS_REPLAY: 1

.github/workflows/test.yml

@@ -2,7 +2,7 @@ name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
   DOWNLOAD_CACHE_VERSION: '12'
-  PYTHON_CACHE_VERSION: '4'
+  PYTHON_CACHE_VERSION: '3'
   APT_CACHE_VERSION: '1'
   BUILD_CACHE_VERSION: '1'
   CAPTURE_PROCESS_REPLAY: 1
@@ -230,7 +230,7 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: linting-only
-        python-version: '3.11'
+        python-version: '3.10'
         deps: linting
     - name: Lint bad-indentation and trailing-whitespace with pylint
       run: python -m pylint --disable=all -e W0311 -e C0303 --jobs=0 --indent-string='  ' --recursive=y .

autogen_stubs.sh

@@ -531,7 +531,7 @@ generate_mesa() {
   sed -i "s/('fp_fast_math', ctypes.c_bool, 9)/('fp_fast_math', ctypes.c_uint32, 9)/" $BASE/mesa.py
   sed -i "s/('\(\w\+\)', pipe_shader_type, 8)/('\1', ctypes.c_ubyte)/" $BASE/mesa.py
   sed -i "s/\([0-9]\+\)()/\1/" $BASE/mesa.py
-  sed -i '/struct_nir_builder._pack_ = 1 # source:False/d' "$BASE/mesa.py"
+  sed -i "s/\(struct_nir_builder._pack_\) = 1/\1 = 0/" $BASE/mesa.py
   python3 -c "import tinygrad.runtime.autogen.mesa"
 }
 

extra/gemm/amd_uop_matmul.py

@@ -68,17 +68,17 @@ def hand_spec_kernel3():
   blockIdx_x = UOp.special(N // BLOCK_N, "gidx0")
   blockIdx_y = UOp.special(N // BLOCK_M, "gidx1")
 
-  a = UOp.placeholder((N, N), dtypes.float, slot=1)
-  b = UOp.placeholder((N, N), dtypes.float, slot=2)
-  c = UOp.placeholder((N, N), dtypes.float, slot=0)
+  a = UOp.placeholder(dtypes.float, (N, N), slot=1)
+  b = UOp.placeholder(dtypes.float, (N, N), slot=2)
+  c = UOp.placeholder(dtypes.float, (N, N), slot=0)
 
   BM_As_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
-  As = UOp.placeholder((BLOCK_K, BM_As_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
-  Bs = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
+  As = UOp.placeholder(dtypes.float, (BLOCK_K, BM_As_stride), slot=0, addrspace=AddrSpace.LOCAL)
+  Bs = UOp.placeholder(dtypes.float, (BLOCK_K, BLOCK_N), slot=1, addrspace=AddrSpace.LOCAL)
 
-  A_col = UOp.placeholder((ITERS_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
-  B_row = UOp.placeholder((ITERS_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
-  c_regs = UOp.placeholder((ITERS_PER_WAVE_M, TM, ITERS_PER_WAVE_N, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  A_col = UOp.placeholder(dtypes.float, (ITERS_PER_WAVE_M, TM), slot=0, addrspace=AddrSpace.REG)
+  B_row = UOp.placeholder(dtypes.float, (ITERS_PER_WAVE_N, TN), slot=1, addrspace=AddrSpace.REG)
+  c_regs = UOp.placeholder(dtypes.float, (ITERS_PER_WAVE_M, TM, ITERS_PER_WAVE_N, TN), slot=2, addrspace=AddrSpace.REG)
 
   i = UOp.range(c_regs.size, 16)
   c_regs = c_regs[i].set(0.0, end=i)
@@ -88,15 +88,15 @@ def hand_spec_kernel3():
   # ---------------------------
   # GLOBAL -> LOCAL (As, Bs)
   # ---------------------------
-  b = b.reshape(N // BLOCK_K, BLOCK_K,
-                N // BLOCK_N, BLOCK_N)
+  b = b.reshape((N // BLOCK_K, BLOCK_K,
+                 N // BLOCK_N, BLOCK_N))
   i = UOp.range(BLOCK_N * BLOCK_K // THREADS_PER_BLOCK, 1)
   index_x = tid % BLOCK_N
   index_y = (tid // BLOCK_N) + (THREADS_PER_BLOCK // BLOCK_N) * i
   Bs_store = Bs[index_y, index_x].store(b[k_tile_range, index_y, blockIdx_x, index_x]).end(i)
 
-  a = a.reshape(N // BLOCK_M, BLOCK_M,
-                N // BLOCK_K, BLOCK_K)
+  a = a.reshape((N // BLOCK_M, BLOCK_M,
+                 N // BLOCK_K, BLOCK_K))
   i = UOp.range(BLOCK_M * BLOCK_K // THREADS_PER_BLOCK, 2)
   index_x = tid % BLOCK_K
   index_y = (tid // BLOCK_K) + (THREADS_PER_BLOCK // BLOCK_K) * i
@@ -113,15 +113,15 @@ def hand_spec_kernel3():
   # ---------------------------
   # LOCAL -> REG (per-wave tiles)
   # ---------------------------
-  Bs_view = Bs.reshape(BLOCK_K, WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)
   iterWaveN = UOp.range(ITERS_PER_WAVE_N, 4)
   i = UOp.range(TN, 5)
-  B_row = B_row[iterWaveN, i].set(Bs_view[k, waveIdx, iterWaveN, idxInWave, i], end=(iterWaveN, i))
+  index = waveIdx * WAVE_TILE_N + iterWaveN * N_PER_ITER + idxInWave * TN + i
+  B_row = B_row[iterWaveN, i].set(Bs[k, index], end=(iterWaveN, i))
 
-  As_view = As.reshape(BLOCK_K, WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM)
   iterWaveM = UOp.range(ITERS_PER_WAVE_M, 6)
   i = UOp.range(TM, 7)
-  A_col = A_col[iterWaveM, i].set(As_view[k, waveIdy, iterWaveM, idyInWave, i], end=(iterWaveM, i))
+  index = waveIdy * WAVE_TILE_M + iterWaveM * M_PER_ITER + idyInWave * TM + i
+  A_col = A_col[iterWaveM, i].set(As[k, index], end=(iterWaveM, i))
 
   # ---------------------------
   # FMA: c_regs += A_col * B_row
@@ -139,8 +139,8 @@ def hand_spec_kernel3():
   # ---------------------------
   # REG -> GLOBAL (epilogue)
   # ---------------------------
-  c = c.reshape(N//BLOCK_M, WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM,
-                N//BLOCK_N, WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)
+  c = c.reshape((N//BLOCK_M, WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM,
+                 N//BLOCK_N, WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN))
   iterWaveM = UOp.range(ITERS_PER_WAVE_M, 1000)
   yt = UOp.range(TM, 1001)
   iterWaveN = UOp.range(ITERS_PER_WAVE_N, 1002)
@@ -149,15 +149,17 @@ def hand_spec_kernel3():
   sink = c_glbl_idx.store(c_regs.after(sink)[iterWaveM, yt, iterWaveN, xt])
   sink = sink.end(iterWaveM, iterWaveN, yt, xt)
 
-  return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()
+  return sink.sink(arg=KernelInfo(opts_to_apply=()))
 
-def test_matmul(sink:UOp, N=N):
+
+if __name__ == "__main__":
   with Context(DEBUG=0):
     a = Tensor.randn(N, N)
     b = Tensor.randn(N, N)
     hc = Tensor.empty(N, N)
     Tensor.realize(a, b, hc)
 
+  sink = hand_spec_kernel3()
   ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in [hc, a, b]])
 
   GlobalCounters.reset()
@@ -175,6 +177,3 @@ def test_matmul(sink:UOp, N=N):
   print(f"mean squared error {err}")
   if err > 1e-06:
     raise RuntimeError("matmul is wrong!")
-
-if __name__ == "__main__":
-  test_matmul(hand_spec_kernel3(), N=N)

extra/gemm/metal_uop_matmul.py

@@ -1,42 +0,0 @@
-from tinygrad import UOp, dtypes
-from tinygrad.uop.ops import AxisType, Ops, KernelInfo, AddrSpace
-from extra.gemm.amd_uop_matmul import test_matmul
-
-N = 2048
-
-# metal has an 8x8 tensor core. this is the indexing
-def mat_idx(buf, g0, g1, warp, u):
-  l = [(warp//2**i)%2 for i in range(5)]
-  return buf[g0, l[4]*4 + l[2]*2 + l[1], g1, l[3]*4 + l[0]*2 + u]
-
-def hand_spec_tc_cores():
-  gx = UOp.special(N // 8, "gidx0")
-  gy = UOp.special(N // 8, "gidx1")
-  warp = UOp.special(32, "lidx0")
-
-  c = UOp.placeholder((N, N), dtypes.float, slot=0).reshape((N//8, 8, N//8, 8))
-  a = UOp.placeholder((N, N), dtypes.float, slot=1).reshape((N//8, 8, N//8, 8))
-  b = UOp.placeholder((N, N), dtypes.float, slot=2).reshape((N//8, 8, N//8, 8))
-
-  gk = UOp.range(N // 8, 0, AxisType.REDUCE)
-
-  a_tc = UOp.vectorize(*[mat_idx(a, gx, gk, warp, i) for i in range(2)])
-  b_tc = UOp.vectorize(*[mat_idx(b, gk, gy, warp, i) for i in range(2)])
-
-  acc = UOp.placeholder((2,), dtypes.float, slot=0, addrspace=AddrSpace.REG)
-  acc = acc[0].set(0.0)
-  acc = acc[1].set(0.0)
-
-  # TODO: make this simple
-  wmma_arg = ('WMMA_8_8_8_float_float', (8, 8, 8), dtypes.float, dtypes.float, 'METAL', 32, (((3, 2),), ((3, 2),), ((3, 2),)), ())
-
-  acc_load = UOp.vectorize(acc.after(gk)[0], acc.after(gk)[1])
-  out = UOp(Ops.WMMA, dtypes.float.vec(2), (a_tc, b_tc, acc_load), arg=wmma_arg)
-
-  end_loop = UOp.group(*[acc[i].store(out.gep(i)) for i in range(2)]).end(gk)
-
-  sink = UOp.group(*[mat_idx(c.after(end_loop), gx, gy, warp, i).store(acc[i]) for i in range(2)])
-  return sink.sink(arg=KernelInfo(name="custom_metal_matmul", opts_to_apply=())).simplify()
-
-if __name__ == "__main__":
-  test_matmul(hand_spec_tc_cores(), N=N)

extra/gemm/mi350x_uop_matmul.py

@@ -1,229 +0,0 @@
-import os
-import numpy as np
-np.set_printoptions(linewidth=1000000)
-os.environ["AMD_LLVM"] = "0"
-
-from tinygrad import Tensor, Context, dtypes, UOp, GlobalCounters
-from tinygrad.helpers import DEBUG, getenv
-from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import AxisType, KernelInfo, Ops
-
-WARP_SIZE = 64
-
-# Reg tile sizes (tensor cores)
-TC_M = 16
-TC_N = 16
-TC_K = 32
-
-# 1024 matrix cores
-# 16 cycle mfma
-# 2.2 GHz
-# 16x16x32x2 FLOPS/mma = 16384
-# 2.2*1e9*16384*1024/16*1e-12 TFLOPS = 2306 TFLOPS
-
-#N,M,K = 256,256,64
-N,M,K = 4096,4096,4096
-
-# Threadblock tile sizes (block-level tile of C that a block computes)
-#BLOCK_M = 128   # rows of C (M-dim) per block
-#BLOCK_N = 128   # columns of C (N-dim) per block
-#BLOCK_K = 128   # K-slice per block iteration
-
-BLOCK_M = 64
-BLOCK_N = 64
-BLOCK_K = 128
-
-WARPGROUP_SIZE = 1
-BLOCK_M = BLOCK_M * WARPGROUP_SIZE
-
-# TODO: improve the syntax of this. better syntax, faster iteration
-#  -- add working slice a[gx, :, i] -> shape of the : (aka (16,16,32) becomes (16,))
-#  -- add argfix to movement (traits shared with Tensor)
-#  -- fix WMMA to not require all the junk
-#  -- improve syntax for vectorized loads/stores (both with DEVECTORIZE and without)
-#  -- be able to use CONTRACT on a range
-#  -- fix upcasted RANGE on an already vectorized buffer
-#  -- improve "all ranges not ended error" / fix the bug with after on ended ranges (if you are after end of range, range is closed)
-
-CUS_PER_GPU = 256
-assert ((M//BLOCK_M) * (N//BLOCK_N)) >= CUS_PER_GPU, "not enough globals"
-
-def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
-  # A = (M x K)
-  # B = (K x N)
-  # C = (M x N)
-
-  # check it's proper matmul
-  assert C.shape[0] == A.shape[0]
-  assert C.shape[1] == B.shape[1]
-  assert A.shape[1] == B.shape[0]
-
-  gx, gy = UOp.special(M//BLOCK_M, "gidx0"), UOp.special(N//BLOCK_N, "gidx1")
-  warp = UOp.special(WARP_SIZE, "lidx0")
-  warpgroup = UOp.special(WARPGROUP_SIZE, "lidx1")
-
-  # generic copy logic (not good)
-  def generic_copy(glbl, gargs, lcl, rng):
-    # Fully coalesced 128-bit loads/stores.
-    INNER_SIZE = 8
-    cp_i = UOp.range(lcl.size//(WARPGROUP_SIZE*WARP_SIZE*INNER_SIZE), rng)
-    cp_inner = UOp.range(INNER_SIZE, rng+1, AxisType.UPCAST)
-    idx_i = cp_i*WARPGROUP_SIZE*WARP_SIZE*INNER_SIZE + warpgroup*WARP_SIZE*INNER_SIZE + warp*INNER_SIZE + cp_inner
-    return lcl[idx_i].store(glbl[*gargs, idx_i]).end(cp_i, cp_inner)
-
-  # split out the globals into blocks
-  C = C.reshape((M//BLOCK_M, BLOCK_M, N//BLOCK_N, BLOCK_N))
-  A = A.reshape((M//BLOCK_M, BLOCK_M, K//BLOCK_K, BLOCK_K))
-  B = B.reshape((K//BLOCK_K, BLOCK_K, N//BLOCK_N, BLOCK_N))
-
-  # this is the big accumulator
-  acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float.vec(4), 0, AddrSpace.REG)
-  assert acc.size*WARP_SIZE*WARPGROUP_SIZE*4 == BLOCK_M*BLOCK_N
-  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const(dtypes.float.vec(4), 0.0), end=init_l)
-
-  # create locals (note A is permuted, and the stride is changed to avoid bank conflicts)
-  def make_locals(slot) -> tuple[UOp, UOp]:
-    BM_As_stride = (BLOCK_M + 1)
-    BN_Bs_stride = (BLOCK_N + 0)
-    INNER_SLICE = 8
-    As = UOp.placeholder((BLOCK_K//INNER_SLICE, BM_As_stride, INNER_SLICE), dtypes.half, slot=slot, addrspace=AddrSpace.LOCAL)
-    INNER_SLICE = 1
-    Bs = UOp.placeholder((BLOCK_K//INNER_SLICE, BN_Bs_stride, INNER_SLICE), dtypes.half, slot=slot+1, addrspace=AddrSpace.LOCAL)
-    As = As.permute((0,2,1)).reshape((BLOCK_K, BM_As_stride)).shrink_to((BLOCK_K, BLOCK_M))
-    Bs = Bs.permute((0,2,1)).reshape((BLOCK_K, BN_Bs_stride)).shrink_to((BLOCK_K, BLOCK_N))
-    return As, Bs
-
-  # load from globals into locals (TODO: use the warpgroup)
-
-  def load_to_locals(l_K_outer_loop:UOp, Asl:UOp, Bsl:UOp, rng:int, barrier=True) -> tuple[UOp, UOp]:
-    if getenv("FAKE"):
-      return Asl[0].set(0), Bsl[0].set(0)
-    else:
-      pA = A.permute((0,2,1,3)).reshape((M//BLOCK_M, K//BLOCK_K, BLOCK_M*BLOCK_K))
-      pas = Asl.permute((1,0)).reshape((BLOCK_M*BLOCK_K,))
-      As_store = generic_copy(pA, (gx, l_K_outer_loop), pas, rng)
-
-      pB = B.permute((0,2,1,3)).reshape((K//BLOCK_K, N//BLOCK_N, BLOCK_K*BLOCK_N))
-      pbs = Bsl.reshape((BLOCK_K*BLOCK_N,))
-      Bs_store = generic_copy(pB, (l_K_outer_loop, gy), pbs, rng+2)
-
-      barrier = UOp.barrier(As_store, Bs_store) if barrier else UOp.group(As_store, Bs_store)
-      return Asl.after(barrier), Bsl.after(barrier)
-
-  def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...]=()) -> UOp:
-    K_inner_loop = UOp.range(BLOCK_K//TC_K, rng, AxisType.REDUCE)
-
-    # load from locals into registers
-    Ar = UOp.placeholder((BLOCK_M//TC_M//WARPGROUP_SIZE,), dtypes.half.vec(8), slot=1, addrspace=AddrSpace.REG)
-    Br = UOp.placeholder((BLOCK_N//TC_N,), dtypes.half.vec(8), slot=2, addrspace=AddrSpace.REG)
-
-    M_load_loop = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, rng+10)
-    Asl = Asl.reshape((BLOCK_K//TC_K, TC_K, BLOCK_M//TC_M//WARPGROUP_SIZE, WARPGROUP_SIZE, TC_M))
-    load_rng = UOp.range(8, rng+11, axis_type=AxisType.UPCAST)
-    A_in = Asl[K_inner_loop, (warp//16)*8+load_rng, M_load_loop, warpgroup, warp%16].contract(load_rng)
-    Ar = Ar[M_load_loop].set(A_in, end=M_load_loop)
-
-    N_load_loop = UOp.range(BLOCK_N//TC_N, rng+20)
-    Bsl = Bsl.reshape((BLOCK_K//TC_K, TC_K, BLOCK_N//TC_N, TC_N))
-    load_rng = UOp.range(8, rng+21, axis_type=AxisType.UPCAST)
-    B_in = Bsl[K_inner_loop, (warp//16)*8+load_rng, N_load_loop, warp%16].contract(load_rng)
-    Br = Br[N_load_loop].set(B_in, end=N_load_loop)
-
-    M_inner_loop = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, rng+30)
-    N_inner_loop = UOp.range(BLOCK_N//TC_N, rng+31)
-
-    # load values
-    acc_after = acc.after(*afters, M_inner_loop, N_inner_loop, K_inner_loop)
-    acc_load = acc_after[N_inner_loop, M_inner_loop]
-
-    # do WMMA
-    wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-    out = UOp(Ops.WMMA, dtypes.float.vec(4), (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)
-
-    # store back the acc
-    acc_store = acc[N_inner_loop, M_inner_loop].store(out)
-    return acc_store.end(M_inner_loop, N_inner_loop, K_inner_loop)
-
-  # **** START INNER LOOP *****
-  # inner loop -- locals -> regs
-
-  # no pipeline
-  if not getenv("PIPELINE"):
-    As, Bs = make_locals(slot=0)
-
-    K_outer_loop = UOp.range(K//BLOCK_K, 0, AxisType.REDUCE)
-    As, Bs = load_to_locals(K_outer_loop, As, Bs, 1000, barrier=True)
-    acc_store = compute_on_locals(acc, As, Bs, 1500, afters=(K_outer_loop,))
-    acc = acc.after(acc_store.barrier().end(K_outer_loop))
-  else:
-    # this doesn't work
-    As0, Bs0 = make_locals(slot=0)
-    As1, Bs1 = make_locals(slot=2)
-    As0, Bs0 = load_to_locals(0, As0, Bs0, 1000)
-
-    K_outer_loop = UOp.range((K//BLOCK_K-2)//2, 0, AxisType.REDUCE)
-    As1, Bs1 = load_to_locals(K_outer_loop+1, As1, Bs1, 2000, barrier=False)
-    acc_store = compute_on_locals(acc, As0, Bs0, 1500, afters=(K_outer_loop,))
-    As0, Bs0 = load_to_locals(K_outer_loop+2, As0, Bs0, 3000, barrier=False)
-    acc_store = compute_on_locals(acc, As1, Bs1, 2500, afters=(acc_store, As0, Bs0))
-    acc = acc.after(acc_store.barrier().end(K_outer_loop))
-
-    #acc_store = compute_on_locals(acc, As0, Bs0, 3500, afters=(acc_store.barrier().end(K_outer_loop)))
-    """
-    As1, Bs1 = load_to_locals(K//BLOCK_K-1, As1, Bs1, 4000)
-    acc_store = compute_on_locals(acc, As1, Bs1, 4500, afters=(acc_store))
-    """
-    #acc = acc.after(acc_store)
-
-  # **** END LOOPS *****
-
-  # store the acc into gmem
-  cp_i, cp_j = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, 10004), UOp.range(BLOCK_N//TC_N, 10005)
-  c_load = lambda i: C[gx, cp_i*TC_M*WARPGROUP_SIZE + warpgroup*TC_M + (warp//16)*4+i, gy, cp_j*TC_N + warp%16]
-  store = UOp.group(*[c_load(i).store(acc[cp_j, cp_i].gep(i)) for i in range(4)])
-  store = store.end(cp_i, cp_j)
-
-  return store.sink(arg=KernelInfo(name="custom_gemm", opts_to_apply=())).simplify()
-
-# simplest WMMA
-"""
-# init the acc
-acc = UOp.placeholder((4,), dtypes.float, 0, AddrSpace.REG)
-acc = acc[init_l:=UOp.range(4, 1)].set(0.0, end=init_l)
-
-# do the wmma
-acc_load = UOp.vectorize(*[acc.after(K_loop)[i] for i in range(4)])
-wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-out = UOp(Ops.WMMA, dtypes.float.vec(4), (A_in, B_in, acc_load), arg=wmma_arg)
-
-# store back the acc
-acc = acc.after(UOp.group(*[acc[i].store(out.gep(i)) for i in range(4)]).end(K_loop))
-
-# store the acc into gmem
-store = UOp.group(*[C[gx, (warp//16)*4+i, gy, warp%16].store(acc[i]) for i in range(4)])
-"""
-
-if __name__ == "__main__":
-  a = Tensor.randn(M, K, dtype=dtypes.half)
-  b = Tensor.randn(K, N, dtype=dtypes.half)
-
-  #a = Tensor.zeros(M, K, dtype=dtypes.half).contiguous()
-  #a[0,16] = 1
-  #b = Tensor.ones(K, N, dtype=dtypes.half).contiguous()
-
-  c = Tensor.empty(M, N, dtype=dtypes.float)
-  with Context(DEBUG=0): Tensor.realize(a,b)
-
-  ref = a.dot(b, dtype=dtypes.float)
-  ref.realize()
-
-  GlobalCounters.reset()
-  with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2):
-    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
-    tst.realize()
-  print(f"{(N*M*K*2 / GlobalCounters.time_sum_s)*1e-12:.2f} REAL TFLOPS")
-
-  with Context(DEBUG=0):
-    #print(ref.numpy())
-    #print(tst.numpy())
-    assert Tensor.isclose(ref, tst, atol=1e-2).all().item(), "matrix not close"

extra/mmapeak/mmapeak.py

@@ -84,14 +84,12 @@ if __name__=="__main__":
     NUM_WORKGROUPS = 256
     WAVE_SIZE = 64
     NUM_WAVES = 4
-    launchBenchmark("v_mfma_f32_16x16x16_f16", (3,0,1), accum=True)
     launchBenchmark("v_mfma_f32_16x16x16_bf16", (3,0,1), accum=True)
     FLOPS_PER_MATMUL = 16*16*32*2
-    launchBenchmark("v_mfma_f32_16x16x32_f16", (3,0,3), accum=True)
     launchBenchmark("v_mfma_f32_16x16x32_bf16", (3,0,3), accum=True)
     FLOPS_PER_MATMUL = 16*16*128*2
     launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,7), accum=True) # fp8
     launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,5), accum=True, extra=", cbsz:2 blgp:2") # fp6
     launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,3), accum=True, extra=", cbsz:4 blgp:4") # fp4
   else:
-    raise RuntimeError(f"arch {DEV.arch} not supported.")
+    raise RuntimeError(f"arch {DEV.arch} not supported.")

extra/qcom_gpu_driver/msm_kgsl.py

@@ -10,8 +10,6 @@ import ctypes
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

extra/sqtt/roc.py

@@ -109,6 +109,6 @@ if __name__ == "__main__":
     for s in ev.sched:
       view = memoryview(ev.blob).cast('Q')
       print(f"\t{s.name}")
-      for xcc, inst, se_idx, sa_idx, wgp_idx in itertools.product(range(s.xcc), range(s.inst), range(s.se), range(s.sa), range(s.wgp)):
-        print(f"\t\tXCC {xcc} Inst {inst} SE {se_idx} SA {sa_idx} WGP {wgp_idx}: {view[ptr]:#x}")
+      for inst, se_idx, sa_idx, wgp_idx in itertools.product(range(s.inst), range(s.se), range(s.sa), range(s.wgp)):
+        print(f"\t\tInst {inst} SE {se_idx} SA {sa_idx} WGP {wgp_idx}: {view[ptr]:#x}")
         ptr += 1

extra/sqtt/rocprof/rocprof.py

@@ -11,8 +11,6 @@ import ctypes, ctypes.util
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

extra/thunder/cuda/fa.cu

@@ -2,7 +2,7 @@
 
 using namespace kittens;
 
-constexpr int NUM_WORKERS = 4;
+constexpr int NUM_WORKERS = 2;
 constexpr int PIPE_STAGES = 3;
 
 constexpr int ATTN_B = 16;

extra/thunder/cuda/fa.py

@@ -13,7 +13,7 @@ if __name__ == "__main__":
   print(pretty_ptx(lib.decode()))
 
   prg = device.runtime(kernel_name, lib)
-  prg.smem = 16384 * 3
+  prg.smem = 16384 * 2
 
   B, N, H, D = 16, 1024, 16, 64
   q = Tensor.randn(B, N, H, D, device='CUDA', dtype="bfloat16")
@@ -22,7 +22,7 @@ if __name__ == "__main__":
   out = Tensor.empty(B, N, H, D, device='CUDA', dtype="bfloat16")
   Tensor.realize(q, k, v, out)
 
-  NUM_WORKERS = 4
+  NUM_WORKERS = 2
   ROWS = 16 * (128 // D)
 
   gsz = (N // (ROWS*NUM_WORKERS), H, B)

extra/thunder/cuda/matmul.cu

@@ -5,11 +5,11 @@ using namespace kittens;
 constexpr int g_N = 8192;
 constexpr int BLOCK_SIZE = 32;
 #define NUM_WORKERS (1)
+#define NUM_THREADS (NUM_WORKERS*kittens::WARP_THREADS)
 
 using sub_tile = st_bf<BLOCK_SIZE,BLOCK_SIZE>;
-using tile_gl =  gl<bf16, 1, 1, g_N, g_N>;
+using tile_gl =  gl<bf16,  1, 1, g_N, g_N, sub_tile>;
 
-__launch_bounds__(NUM_WORKERS*WARP_THREADS, 1)
 __global__ void kernel(bf16 *c_ptr, bf16 *a_ptr, bf16 *b_ptr) {
   tile_gl g_C{c_ptr, nullptr, nullptr, nullptr, nullptr};
   tile_gl g_A{a_ptr, nullptr, nullptr, nullptr, nullptr};

extra/thunder/cuda/matmul.py

@@ -1,14 +1,10 @@
 import pathlib
 from tinygrad import Device, Tensor
-from tinygrad.helpers import Context, getenv
+from tinygrad.helpers import Context
 from tinygrad.runtime.support.compiler_cuda import pretty_ptx, NVCCCompiler
 
 if __name__ == "__main__":
-  if getenv("MATMUL2"):
-    code = (pathlib.Path(__file__).parent / "matmul2.cu").read_text()
-  else:
-    code = (pathlib.Path(__file__).parent / "matmul.cu").read_text()
-
+  code = (pathlib.Path(__file__).parent / "matmul.cu").read_text()
   device = Device["CUDA"]
   kitten_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "--expt-relaxed-constexpr"]
   lib = NVCCCompiler(device.compiler.arch, kitten_args).compile(code)
@@ -17,10 +13,7 @@ if __name__ == "__main__":
   print(pretty_ptx(lib.decode()))
 
   prg = device.runtime(kernel_name, lib)
-  if getenv("MATMUL2"):
-    prg.smem = 16384 * 2
-  else:
-    prg.smem = 10000
+  prg.smem = 10000
 
   N = 8192
   a = Tensor.randn(N, N, device='CUDA', dtype="bfloat16")
@@ -28,25 +21,14 @@ if __name__ == "__main__":
   c = Tensor.empty(N, N, device='CUDA', dtype="bfloat16")
   Tensor.realize(a, b, c)
 
-  WARP_THREADS = 32
-  if getenv("MATMUL2"):
-    SUPER_N = 2
-    SUPER_M = 2
-    NUM_WORKERS = SUPER_N * SUPER_M
-    BLOCK_SIZE = 32
-    gsz = (N // (BLOCK_SIZE * SUPER_N), N // (BLOCK_SIZE * SUPER_M), 1)
-  else:
-    NUM_WORKERS = 1
-    BLOCK_SIZE = 32
-    gsz = (N // (BLOCK_SIZE), N // (BLOCK_SIZE), 1)
+  BLOCK_SIZE = 32
 
+  gsz = (N // BLOCK_SIZE, N // BLOCK_SIZE, 1)
   for _ in range(5):
     et = prg(c.uop.buffer.ensure_allocated()._buf, a.uop.buffer._buf, b.uop.buffer._buf,
-             global_size=gsz, local_size=(NUM_WORKERS*WARP_THREADS,1,1), wait=True)
+             global_size=gsz, local_size=(32,1,1), wait=True)
     print(f"{N*N*N*2/(et*1e9):2f} GFLOPS")
 
-  # print(c.tolist())
-
   for _ in range(5):
     with Context(DEBUG=2):
       ref = (a@b).realize()

extra/thunder/cuda/matmul2.cu

@@ -1,105 +0,0 @@
-#include "kittens.cuh"
-using namespace kittens;
-
-constexpr int g_N = 8192;
-
-constexpr int SUPER_N = 2;
-constexpr int SUPER_M = 2;
-constexpr int NUM_WORKERS = SUPER_N * SUPER_M;
-constexpr int LOAD_TASKS = SUPER_N + SUPER_M;
-
-constexpr int WORKER_M = 32;
-constexpr int WORKER_N = 32;
-
-constexpr int BLOCK_K = 32;
-constexpr int BLOCK_M = WORKER_M * SUPER_M;
-constexpr int BLOCK_N = WORKER_N * SUPER_N;
-
-constexpr int PIPE_STAGES = 2;
-
-using reg_tile_A = rt_bf<WORKER_M, BLOCK_K>;
-using reg_tile_B_col = rt_bf<BLOCK_K, WORKER_N, ducks::rt_layout::col>;
-using reg_tile_C = rt_fl<WORKER_M, WORKER_N>;
-
-using shared_tile_A = st_bf<WORKER_M, BLOCK_K>;
-using shared_tile_B = st_bf<BLOCK_K, WORKER_N>;
-using shared_tile_C = st_bf<WORKER_M, WORKER_N>;
-
-using gl_tile_A = gl<bf16, 1, 1, g_N, g_N, shared_tile_A>;
-using gl_tile_B = gl<bf16, 1, 1, g_N, g_N, shared_tile_B>;
-using gl_tile_C = gl<bf16, 1, 1, g_N, g_N, shared_tile_C>;
-
-__launch_bounds__(NUM_WORKERS *WARP_THREADS, 1) __global__
-    void kernel(bf16 *c_ptr, bf16 *a_ptr, bf16 *b_ptr) {
-  gl_tile_C g_C{c_ptr, nullptr, nullptr, nullptr, nullptr};
-  gl_tile_A g_A{a_ptr, nullptr, nullptr, nullptr, nullptr};
-  gl_tile_B g_B{b_ptr, nullptr, nullptr, nullptr, nullptr};
-
-  extern __shared__ alignment_dummy __shm[];
-  shared_allocator al((int *)&__shm[0]);
-
-  shared_tile_A(&As)[SUPER_M][PIPE_STAGES] =
-      al.allocate<shared_tile_A, SUPER_M, PIPE_STAGES>();
-  shared_tile_B(&Bs)[SUPER_N][PIPE_STAGES] =
-      al.allocate<shared_tile_B, SUPER_N, PIPE_STAGES>();
-
-  reg_tile_A A_reg;
-  reg_tile_B_col B_reg_col;
-  reg_tile_C C_accum;
-
-  int warpid = kittens::warpid();
-  int warp_m = warpid % SUPER_M;
-  int warp_n = warpid / SUPER_M;
-
-  int load_group_id = warpgroup::groupid();
-
-  int block_row = blockIdx.y * SUPER_M;
-  int block_col = blockIdx.x * SUPER_N;
-
-  warp::zero(C_accum);
-  int num_tiles = (g_N + BLOCK_K - 1) / BLOCK_K;
-
-  for (int load_tile = 0; load_tile < (PIPE_STAGES - 1); load_tile++) {
-    if (load_tile < num_tiles) {
-      int load_smem_idx = load_tile % PIPE_STAGES;
-      for (int task_id = warpid; task_id < LOAD_TASKS; task_id += NUM_WORKERS) {
-        if (task_id < SUPER_M) {
-          warp::load_async(As[task_id][load_smem_idx], g_A, {0, 0, block_row + task_id, load_tile});
-        } else {
-          int n_index = task_id - SUPER_M;
-          warp::load_async(Bs[n_index][load_smem_idx], g_B, {0, 0, load_tile, block_col + n_index});
-        }
-      }
-    }
-  }
-
-  for (int tile = 0; tile < num_tiles; tile++) {
-    int compute_smem_idx = tile % PIPE_STAGES;
-
-    int load_tile = tile + PIPE_STAGES - 1;
-    int load_smem_idx = load_tile % PIPE_STAGES;
-
-    if (load_tile < num_tiles) {
-      for (int task_id = warpid; task_id < LOAD_TASKS; task_id += NUM_WORKERS) {
-        if (task_id < SUPER_M) {
-          warp::load_async(As[task_id][load_smem_idx], g_A,
-                           {0, 0, block_row + task_id, load_tile});
-        } else {
-          int n_index = task_id - SUPER_M;
-          warp::load_async(Bs[n_index][load_smem_idx], g_B,
-                           {0, 0, load_tile, block_col + n_index});
-        }
-      }
-      load_async_wait<1>();
-    } else
-      load_async_wait();
-    __syncthreads();
-
-    warp::load(A_reg, As[warp_m][compute_smem_idx]);
-    warp::load(B_reg_col, Bs[warp_n][compute_smem_idx]);
-
-    warp::mma_AB(C_accum, A_reg, B_reg_col, C_accum);
-    __syncthreads();
-  }
-  warp::store(g_C, C_accum, {0, 0, block_row + warp_m, block_col + warp_n});
-}

test/test_custom_kernel.py

@@ -1,156 +0,0 @@
-import unittest
-from tinygrad import Tensor, UOp, Context
-from tinygrad.uop.ops import KernelInfo, AxisType
-
-# **** kernels ****
-
-def custom_arange_kernel(C:UOp) -> UOp:
-  i = UOp.range(C.size, 0)
-  return C[i].store(i.cast(C.dtype.base)).end(i).sink(arg=KernelInfo(name=f"custom_arange_{C.size}"))
-
-def custom_add_one_kernel(B:UOp, A:UOp) -> UOp:
-  assert B.size == A.size
-  i = UOp.range(A.size, 0)
-  return B[i].store(A[i] + 1).end(i).sink(arg=KernelInfo(name=f"add_one_{A.size}"))
-
-def custom_elementwise_add_kernel(C:UOp, A:UOp, B:UOp) -> UOp:
-  i = UOp.range(C.size, 0)
-  return C[i].store(A[i]+B[i]).end(i).sink(arg=KernelInfo(name=f"custom_add_kernel_{C.size}")).simplify()
-
-def custom_elementwise_addmul_kernel(C:UOp, D:UOp, A:UOp, B:UOp) -> UOp:
-  assert C.size == D.size
-  i = UOp.range(C.size, 0)
-  store_c = C[i].store(A[i]+B[i])
-  store_d = D[i].store(A[i]*B[i])
-  return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name=f"custom_addmul_kernel_{C.size}")).simplify()
-
-def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
-  assert A.shape[1] == B.shape[0]
-  i, j, k = UOp.range(C.shape[0], 0), UOp.range(C.shape[1], 1), UOp.range(A.shape[1], 2, axis_type=AxisType.REDUCE)
-  C = C[i, j].set(0.0)
-  C = C[i, j].set(C.after(k)[i, j] + A[i, k] * B[k, j], end=k)
-  prog = C.end(i, j)
-  return prog.sink(arg=KernelInfo(name=f"custom_gemm_{C.shape[0]}_{C.shape[1]}_{A.shape[1]}", opts_to_apply=()))
-
-def custom_sum(B:UOp, A:UOp) -> UOp:
-  i = UOp.range(A.shape[0], 0, axis_type=AxisType.REDUCE)
-  B = B[0].set(0.0)
-  B = B[0].set(B.after(i)[0] + A[i], end=i)
-  return B.sink(arg=KernelInfo(name=f"custom_sum_{A.shape[0]}", opts_to_apply=()))
-
-def flip_contract_kernel(dest:UOp, src:UOp):
-  assert dest.size%4 == 0
-  i = UOp.range(dest.size//4, 0)
-  j = UOp.range(4, 1, AxisType.UPCAST)
-  vec = src[i*4+j].contract(j)
-  store = UOp.group(*[dest[i*4+k].store(vec.gep(3-k)) for k in range(4)])
-  return store.end(i).sink(arg=KernelInfo(name=f"flip_contract_{dest.size}", opts_to_apply=()))
-
-# **** backward callbacks ****
-
-def backward_gemm(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
-  out, a, b = kernel.src
-  grad_a = (Tensor(gradient) @ Tensor(b).T).uop
-  grad_b = (Tensor(a).T @ Tensor(gradient)).uop
-  return (None, grad_a, grad_b)
-
-def backward_gemm_custom(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
-  out, a, b = kernel.src
-  grad_a = Tensor.empty_like(Tensor(a)).custom_kernel(Tensor(gradient), Tensor(b).T, fxn=custom_gemm)[0].uop
-  grad_b = Tensor.empty_like(Tensor(b)).custom_kernel(Tensor(a).T, Tensor(gradient), fxn=custom_gemm)[0].uop
-  return (None, grad_a, grad_b)
-
-# **** tests ****
-
-class TestCustomKernel(unittest.TestCase):
-  def test_simple(self):
-    a = Tensor.ones(16, 16).contiguous()
-    b = Tensor.ones(16, 16).contiguous()
-    c = Tensor.empty(16, 16)
-
-    c = Tensor.custom_kernel(c,a,b, fxn=custom_elementwise_add_kernel)[0]
-
-    out = c.flatten().tolist()
-    assert all(x == 2 for x in out), "all 2"
-
-  def test_multioutput(self):
-    a = Tensor.full((16, 16), 3.).contiguous()
-    b = Tensor.full((16, 16), 3.).contiguous()
-    c = Tensor.empty(16, 16)
-    d = Tensor.empty(16, 16)
-
-    c,d = Tensor.custom_kernel(c,d,a,b, fxn=custom_elementwise_addmul_kernel)[:2]
-    Tensor.realize(c,d)
-
-    assert all(x == 6 for x in c.flatten().tolist()), "all 6"
-    assert all(x == 9 for x in d.flatten().tolist()), "all 9"
-
-  def test_arange(self):
-    ref = Tensor.arange(100)
-    tst = Tensor.empty_like(ref)
-    tst = tst.custom_kernel(fxn=custom_arange_kernel)[0]
-    self.assertTrue((ref == tst).all().item())
-
-  def test_flip_contract(self):
-    a = Tensor.randn(10,4)
-    b = Tensor.empty_like(a)
-    b = b.custom_kernel(a, fxn=flip_contract_kernel)[0]
-    self.assertTrue((a.flip(1) == b).all().item())
-
-  def test_noncontig(self):
-    a = Tensor.ones(16, 16).contiguous()
-    tst = Tensor.empty_like(a)
-    b = a+1
-    b_p1 = Tensor.custom_kernel(tst, b, fxn=custom_add_one_kernel)[0]
-    self.assertTrue((b_p1 == 3).all().item())
-
-  def test_sum(self):
-    # TODO: this only works for float, and silently fails with int
-    a = Tensor([1.0, 2, 3, 4, 5])
-    tst = Tensor.empty(1)
-    b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
-    self.assertEqual(b.item(), 15)
-
-  def test_gemm(self):
-    N = 16
-    a = Tensor.randn(N, N)
-    b = Tensor.randn(N, N)
-    c = Tensor.empty(N, N)
-
-    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
-    err = (tst - (a@b)).square().max()
-    self.assertLess(err.item(), 1e-6)
-
-  def test_gemm_backward_custom(self): self.test_gemm_backward(True)
-  # NOTE: grad_fxn doesn't work with pyrender
-  @Context(SPEC=1)
-  def test_gemm_backward(self, custom_backward_gemm=False):
-    N = 4
-    a_rand = Tensor.randn(N, 8)
-    b_rand = Tensor.randn(8, N)
-    Tensor.realize(a_rand, b_rand)
-
-    a, b = Tensor(a_rand.numpy(), requires_grad=True), Tensor(b_rand.numpy(), requires_grad=True)
-    c = Tensor.empty(N, N)
-    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm, grad_fxn=backward_gemm_custom if custom_backward_gemm else backward_gemm)[0]
-    tst.sum().backward()
-    grad_a, grad_b = a.grad, b.grad
-    Tensor.realize(tst, grad_a, grad_b)
-
-    a, b = Tensor(a_rand.numpy(), requires_grad=True), Tensor(b_rand.numpy(), requires_grad=True)
-    ref = (a@b)
-    ref.sum().backward()
-    real_grad_a, real_grad_b = a.grad, b.grad
-    Tensor.realize(ref, real_grad_a, real_grad_b)
-
-    err = (tst - ref).square().max()
-    self.assertLess(err.item(), 1e-6)
-
-    err = (grad_a - real_grad_a).square().max()
-    self.assertLess(err.item(), 1e-6)
-
-    err = (grad_b - real_grad_b).square().max()
-    self.assertLess(err.item(), 1e-6)
-
-if __name__ == '__main__':
-  unittest.main()

test/test_schedule.py

@@ -711,7 +711,7 @@ class TestSchedule(unittest.TestCase):
     self.assertEqual(b.buffer.numpy(), [12])
 
   # unlike schedule, kernelize can be called multiple times on a Tensor
-  def test_double_kernelize(self):
+  def test_double_kerenlize(self):
     a = Tensor.empty(10)
     b = Tensor.empty(10)
     c = (a+b)
@@ -1503,18 +1503,6 @@ class TestSchedule(unittest.TestCase):
     run_schedule(sched)
     np.testing.assert_allclose(dx.numpy(), [[[[0.,3.,9.],[0,1.,3.],[0.,0.,0.]]]*3]*3)
 
-  def test_fuse_arange_avg_pool2d_ceil_mode(self):
-    x = Tensor.avg_pool2d(Tensor.empty(1,1,6,6), kernel_size=(3,3), padding=1, stride=3, ceil_mode=True)
-    sched = check_schedule(x, 1)
-    self.assertEqual(len([x for x in sched[0].ast.backward_slice_with_self if x.op is Ops.REDUCE]), 1)
-
-  def test_fuse_arange_pad_circular_mode_bw(self):
-    x = Tensor.empty(1,1,5,5,5)
-    out = x.pad((1,2,3,5,1,2), mode="circular")
-    g = out.sum().gradient(x)[0]
-    sched = check_schedule(g, 1)
-    self.assertEqual(len([x for x in sched[0].ast.backward_slice_with_self if x.op is Ops.REDUCE]), 0)
-
   # TODO like openpilot with imagef
   @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
   def test_base_change_expand_expand(self):
@@ -2279,7 +2267,7 @@ class TestContiguous(unittest.TestCase):
   def test_double_contiguous_realizes_once(self):
     a = Tensor.empty(4, 1)
     b = a.expand((4, 4)).contiguous().contiguous()
-    check_schedule(b, 1)
+    check_schedule(b, 2) # TODO: should be 1?
 
   def test_view_does_not_realize(self):
     a = Tensor.empty(4)

test/test_uops.py

@@ -572,7 +572,7 @@ class TestUOpPrograms(unittest.TestCase):
   def test_simple(self):
     out = Tensor.empty(10,10,dtype=dtypes.int)
 
-    ptr = UOp.placeholder(out.shape, out.dtype, slot=0)
+    ptr = UOp.placeholder(out.dtype, out.shape, slot=0)
     i, j = UOp.range(10, axis_id=0), UOp.range(10, axis_id=1)
     prog = ptr[i,j].set(42).end(i,j)
     self._run(prog.sink(), out)
@@ -592,9 +592,9 @@ class TestUOpPrograms(unittest.TestCase):
     DT = dtypes.float32
 
     # Placeholders (bind slots explicitly)
-    A = UOp.placeholder((M, K), DT, slot=0)
-    B = UOp.placeholder((K, N), DT, slot=1)
-    C = UOp.placeholder((M, N), DT, slot=2)
+    A = UOp.placeholder(DT, (M, K), slot=0)
+    B = UOp.placeholder(DT, (K, N), slot=1)
+    C = UOp.placeholder(DT, (M, N), slot=2)
 
     # Axes: i,j are spatial; k is a reduction axis over the shared dim K
     i = UOp.range(M, axis_id=0)                             # rows of A/C

test/unit/test_helpers.py

@@ -99,11 +99,6 @@ class TestStripParens(unittest.TestCase):
   def test_simple(self): self.assertEqual("1+2", strip_parens("(1+2)"))
   def test_nested(self): self.assertEqual("1+(2+3)", strip_parens("(1+(2+3))"))
   def test_casted_no_strip(self): self.assertEqual("(int)(1+2)", strip_parens("(int)(1+2)"))
-  def test_unmatched_parens(self): self.assertEqual("((c35+c39>>23&255)+-127).cast(dtypes.float)",
-    strip_parens("((c35+c39>>23&255)+-127).cast(dtypes.float)"))
-  def test_single_paren_left(self): self.assertEqual("(abc", strip_parens("(abc"))
-  def test_single_paren_right(self): self.assertEqual("abc)", strip_parens("abc)"))
-  def test_parens_at_different_depths(self): self.assertEqual("(a+(b))*(c)", strip_parens("(a+(b))*(c)"))
 
 class TestProd(unittest.TestCase):
   def test_empty(self): self.assertEqual(1, prod(tuple()))

tinygrad/codegen/__init__.py

@@ -19,13 +19,18 @@ from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_s
 from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 
+pm_preprocess = PatternMatcher([
+  (UPat(Ops.RESHAPE, name="r").after(name="a", allow_any_len=True), lambda r,a: a.replace(src=(r.src[0],)+a.src[1:]).reshape(r.shape)),
+  (UPat(Ops.RESHAPE, name="r").end(name="a", allow_any_len=True), lambda r,a: a.replace(src=(r.src[0],)+a.src[1:])),
+])
+
 def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -> UOp:
   if ren is None: ren = Renderer()
 
   if SPEC: type_verify(sink, kernel_spec)
 
   # preprocess
-  sink = graph_rewrite(sink, pm_mops, name="early movement ops")
+  sink = graph_rewrite(sink, pm_preprocess+pm_mops, name="early movement ops")
 
   # first we optimize
   if optimize:

tinygrad/codegen/late/expander.py

@@ -75,18 +75,10 @@ def do_contract(con:UOp):
     idxs += [_expand_arg_to_idx(ex.arg, {**rpk, **lrpk}) for lrpk in _choices_from_args(con.arg)]
   return UOp(Ops.UNROLL, con.dtype, (ex.src[0].gep(tuple(idxs)),), new_ex_args)
 
-def end_unrolls(u:UOp):
-  unrolls, src = partition(u.src[1:], lambda x: x.op is Ops.UNROLL)
-  if not len(unrolls): return None
-  ret = UOp(Ops.CONTRACT, dtypes.void, (u.src[0],), sum([x.arg for x in unrolls], start=()))
-  return u.replace(src=(ret,)+tuple(src))
-
 expander = PatternMatcher([
   # push broadcast through AFTER
   (UPat.var("x").broadcast(name="b").after(name="a", allow_any_len=True), lambda x,b,a: x.after(*a.src[1:]).broadcast(len(b.src))),
   (UPat.var("x").broadcast(name="b").end(name="a", allow_any_len=True), lambda x,b,a: x.end(*a.src[1:]).broadcast(len(b.src))),
-  # END on UNROLL ends the UNROLL
-  (UPat(Ops.END, name="u"), end_unrolls),
   # BUFFERIZE puts UNROLLs for ranges as contract
   (UPat(Ops.BUFFERIZE, src=(UPat(Ops.UNROLL), UPat(Ops.UNROLL)), name="x"),
     lambda x: x.replace(src=tuple(UOp(Ops.CONTRACT, dtype=s.dtype.vec(x.src[1].src[0].dtype.count), src=(s,), arg=x.src[1].arg) for s in x.src))),

tinygrad/codegen/opt/heuristic.py

@@ -73,9 +73,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
           if first_reduce_rng.src[0].divides(MV_THREADS_PER_ROW) is not None and k.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0:
             if DEBUG >= 3:
               print(f"MATVEC: {k.full_shape=} {first_reduce_rng.render()} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}")
-            try:
-              if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW))
-            except KernelOptError: pass
+            if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW))
             if MV_BLOCKSIZE > 1: k.apply_opt(Opt(OptOps.LOCAL, global_idx, MV_BLOCKSIZE))
             if MV_ROWS_PER_THREAD > 1: k.apply_opt(Opt(OptOps.UPCAST, global_idx, MV_ROWS_PER_THREAD))
             return k

tinygrad/codegen/opt/postrange.py

@@ -66,7 +66,7 @@ class Scheduler:
   def _output_rngs(self) -> list[UOp]:
     return flatten([[r for r in UOp.sink(*s.src[1:]).ranges if r.arg[-1] != AxisType.REDUCE] for s in self.ast.src if s.op is Ops.END])
   def _globalizable_rngs(self) -> list[UOp]:
-    ret = [r for r in self._output_rngs() if r.arg[-1] == AxisType.LOOP]
+    ret = self._output_rngs()
     # exclude any output ranges from global that don't appear in all BUFFERIZE
     for x in self.ast.toposort():
       if x.op is Ops.BUFFERIZE:

tinygrad/codegen/simplify.py

@@ -91,59 +91,47 @@ pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
   # lift x*y out of reduce
   ((UPat.var("x")*UPat.var("y")) < UPat.var("c"), lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and y.vmin > 0 else None),
   # fold the range
-  # bound from below
-  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(0, UPat.var("val")).reduce(UPat.var("r"), arg=Ops.ADD),
-   lambda r,cut,val: (r.src[0]-cut).maximum(0).minimum(r.src[0]).cast(val.dtype) * val if no_range(val) else None),
-  # bound from two sides
-  (((UPat.var("r")<UPat.var("lower")).logical_not()&(UPat(Ops.RANGE, name="r")<UPat.var("upper"))).where(UPat.var("val"), 0).reduce(UPat.var("r"),
-    arg=Ops.ADD), lambda r,lower,upper,val:
-      (upper.minimum(r.src[0])-lower.maximum(0)).maximum(0).minimum(r.src[0]).cast(val.dtype) * val if no_range(val) else None),
-  # bound from above
-  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(UPat.var("val"), 0).reduce(UPat.var("r"), arg=Ops.ADD),
-   lambda r,cut,val: cut.maximum(0).minimum(r.src[0]).cast(val.dtype) * val if no_range(val) else None),
+  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(0, UPat.cvar("val")).reduce(UPat.var("r"), arg=Ops.ADD),
+   lambda r,cut,val: (r.src[0]-cut).maximum(0).minimum(r.src[0]).cast(val.dtype) * val),
+  (((UPat.var("r")<UPat.var("lower")).logical_not()&(UPat(Ops.RANGE, name="r")<UPat.var("upper"))).where(UPat.cvar("val"), 0).reduce(UPat.var("r"),
+    arg=Ops.ADD), lambda r,lower,upper,val: (upper.minimum(r.src[0])-lower.maximum(0)).maximum(0).minimum(r.src[0]).cast(val.dtype) * val),
+  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(UPat.cvar("val"), 0).reduce(UPat.var("r"), arg=Ops.ADD),
+   lambda r,cut,val: cut.maximum(0).minimum(r.src[0]).cast(val.dtype) * val),
   # REDUCE on ADD
   ((UPat.var("x")+UPat.var("y")).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
    lambda x,y,r: x.reduce(*r.src[1:], arg=Ops.ADD) + y.reduce(*r.src[1:],arg=Ops.ADD)),
-  # AND on WHERE
-  ((UPat(Ops.DEFINE_VAR, name="x") & UPat.var("y")).where(UPat.var("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
-    lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x.cast(c.dtype)),
-  # MUL casted bool
-  ((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast()), lambda x,gate: gate.where(x, 0)),
 ])+symbolic_flat
 
-pm_reduce_load_collapse = pm_reduce_collapse + PatternMatcher([
+pm_reduce_load_collapse = PatternMatcher([
+  # MUL casted bool
+  ((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast()), lambda x,gate: gate.where(x, 0)),
   # lift x+y out of reduce on ne
   ((UPat.var("x")+UPat.var("y")).or_casted() != UPat.var("c"), lambda x,y,c: (x != (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
   # reduce on gated load becomes can substitute the range and remove the reduce
   ((UPat.var("idx")!=(UPat(Ops.RANGE, name="r").or_casted())).where(0, UPat.var("expr")).reduce(UPat.var("r"), arg=Ops.ADD),
    lambda r,idx,expr: (v:=(idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])).where(expr.substitute({r:idx.cast(r.dtype).valid(v)}),0)),
-])
+])+symbolic_flat
 
-def reduce_collapse(red:UOp, u:UOp, pm=pm_reduce_collapse):
-  for r in red.src[1:]:
-    included = u.toposort(gate=lambda x: r in x.ranges)
-    if any(x.op in {Ops.STORE, Ops.REDUCE} for x in included): return None
-    replaces: dict[UOp, UOp] = {}
-    for u in included:
-      for s in u.src:
-        if s in included or s in replaces or s.op in {Ops.CONST, Ops.VCONST, Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}: continue
-        replaces[s] = UOp(Ops.DEFINE_VAR, dtype=s.dtype, arg=(f'in{len(replaces)}', s.vmin, s.vmax))
-    collapse_fxn = u.substitute(replaces).reduce(r, arg=Ops.ADD)
-    sink = graph_rewrite(collapse_fxn, pm, name="reduce_collapse")
-    if not no_range(sink): return None
-    u = sink.substitute({v:k for k,v in replaces.items()})
-  return u
+def reduce_collapse(red:UOp, pm=pm_reduce_collapse):
+  included = red.src[0].toposort(gate=lambda x: any(y in x.ranges for y in red.src[1:]))
+  if any(x.op in {Ops.STORE, Ops.REDUCE} for x in included): return None
+  replaces: dict[UOp, UOp] = {}
+  for u in included:
+    for s in u.src:
+      if s in included or s in replaces or s.op in {Ops.CONST, Ops.VCONST, Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}: continue
+      replaces[s] = UOp(Ops.DEFINE_VAR, dtype=s.dtype, arg=(f'in{len(replaces)}', s.vmin, s.vmax))
+  collapse_fxn = red.substitute(replaces)
+  sink = graph_rewrite(collapse_fxn, pm, name="reduce_collapse")
+  return sink.substitute({v:k for k,v in replaces.items()}) if no_range(sink) else None
 
-def reduce_load_collapse(red:UOp, u:UOp): return reduce_collapse(red, u, pm=pm_reduce_load_collapse)
+def reduce_load_collapse(red:UOp): return reduce_collapse(red, pm=pm_reduce_load_collapse)
 
-# remove REDUCE without loads (generic arange opt / indexing).
-pm_reduce_simplify = pm_reduce_unparented + PatternMatcher([
-  (UPat(Ops.REDUCE, src=(UPat.var("u"),), allow_any_len=True, arg=Ops.ADD, name="red"), reduce_collapse),
-])
+# remove REDUCE without loads (generic arange opt / indexing). TODO: support multi range
+pm_reduce_simplify = pm_reduce_unparented + PatternMatcher([(UPat(Ops.REDUCE, src=(UPat(), UPat()), name="red"), reduce_collapse),])
 # remove REDUCE on load, comes from indexing a tensor with another tensor
 def no_load(u:UOp) -> bool: return not any(x.op is Ops.INDEX for x in u.backward_slice_with_self)
 pm_load_collapse = PatternMatcher([
-  (UPat(Ops.REDUCE, src=(UPat.var("u"), UPat()), name="red"), reduce_load_collapse),
+  (UPat(Ops.REDUCE, src=(UPat(), UPat()), name="red"), reduce_load_collapse),
   # we want to make sure we dont do math on a loaded index since that can cause overflow, this undoes the rule in pm_reduce_load_collapse
   ((UPat.var("x", dtypes.index)+UPat.var("y"))<UPat.var("c"), lambda x,y,c: x < c-y if no_load(y) and no_load(c) and not no_load(x) else None),
 ])

tinygrad/gradient.py

@@ -38,9 +38,6 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.PERMUTE, name="ret"), lambda ctx, ret: (ctx.permute(argsort(ret.marg)),)),
   (UPat(Ops.FLIP, name="ret"), lambda ctx, ret: (ctx.flip(ret.marg),)),
   (UPat(Ops.MULTI, name="ret"), lambda ctx, ret: ctx.shard(ret.device, ret.axis).src),
-  # NOTE: this is only correct when the KERNEL has a single output
-  (UPat(Ops.AFTER), lambda ctx: (ctx, ctx)),
-  (UPat(Ops.KERNEL, name="k"), lambda ctx, k: k.arg.grad_fxn(ctx, k)),
   # there's no gradient for bitcast
   (UPat(Ops.BITCAST), lambda: (None,)),
 ])

tinygrad/helpers.py

@@ -44,9 +44,7 @@ def fully_flatten(l):
     return flattened
   return [l]
 def fromimport(mod, frm): return getattr(__import__(mod, fromlist=[frm]), frm)
-def _is_balanced(s:str) -> bool:
-  return (acc:=list(itertools.accumulate([(1 if ch=='(' else -1 if ch==')' else 0) for ch in s])))[-1]==0 and all(x>=0 for x in acc)
-def strip_parens(fst:str) -> str: return fst[1:-1] if fst and fst[0]=='(' and fst[-1] == ')' and _is_balanced(fst[1:-1]) else fst
+def strip_parens(fst:str): return fst[1:-1] if fst[0] == '(' and fst[-1] == ')' and fst[1:-1].find('(') <= fst[1:-1].find(')') else fst
 def ceildiv(num, amt): return int(ret) if isinstance((ret:=-(num//-amt)), float) else ret
 def round_up(num:int, amt:int) -> int: return (num+amt-1)//amt * amt
 def round_down(num:int, amt:int) -> int: return -round_up(-num, amt)

tinygrad/runtime/autogen/am/am.py

@@ -10,8 +10,6 @@ import ctypes
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/am/pm4_nv.py

@@ -10,8 +10,6 @@ import ctypes
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/am/pm4_soc15.py

@@ -10,8 +10,6 @@ import ctypes
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/am/sdma_4_0_0.py

@@ -10,8 +10,6 @@ import ctypes
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/am/sdma_4_4_2.py

@@ -10,8 +10,6 @@ import ctypes
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/am/sdma_5_0_0.py

@@ -10,8 +10,6 @@ import ctypes
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/am/sdma_6_0_0.py

@@ -10,8 +10,6 @@ import ctypes
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/am/smu_v13_0_0.py

@@ -10,8 +10,6 @@ import ctypes
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/am/smu_v14_0_2.py

@@ -10,8 +10,6 @@ import ctypes
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/amd_gpu.py

@@ -10,8 +10,6 @@ import ctypes, os
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/comgr.py

@@ -52,8 +52,6 @@ else:
     c_long_double_t = ctypes.c_ubyte*16
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/cuda.py

@@ -10,8 +10,6 @@ import ctypes, ctypes.util
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/hip.py

@@ -10,8 +10,6 @@ import ctypes, os
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/hsa.py

@@ -31,8 +31,6 @@ def char_pointer_cast(string, encoding='utf-8'):
 _libraries = {}
 _libraries['libhsa-runtime64.so'] = ctypes.CDLL(os.getenv('ROCM_PATH')+'/lib/libhsa-runtime64.so' if os.getenv('ROCM_PATH') else ctypes.util.find_library('hsa-runtime64'))
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/ib.py

@@ -10,8 +10,6 @@ import ctypes, ctypes.util
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/io_uring.py

@@ -10,8 +10,6 @@ import ctypes
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/kfd.py

@@ -24,8 +24,6 @@ def _IOR(base, nr, type): return functools.partial(_do_ioctl, 2, ord(base) if is
 def _IOWR(base, nr, type): return functools.partial(_do_ioctl, 3, ord(base) if isinstance(base, str) else base, nr, type)
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/kgsl.py

@@ -10,8 +10,6 @@ import ctypes, os
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/libc.py

@@ -48,8 +48,6 @@ def char_pointer_cast(string, encoding='utf-8'):
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/libusb.py

@@ -21,8 +21,6 @@ class FunctionFactoryStub:
 _libraries = {}
 _libraries['libusb'] = None if (lib_path:=os.getenv('LIBUSB_PATH', ctypes.util.find_library('usb-1.0'))) is None else ctypes.CDLL(lib_path) #  ctypes.CDLL('libusb')
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/llvm.py

@@ -10,8 +10,6 @@ import ctypes, tinygrad.runtime.support.llvm as llvm_support
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/mesa.py

@@ -23,8 +23,6 @@ def _try_dlopen_tinymesa_cpu():
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}
@@ -10256,6 +10254,7 @@ nir_instr_writemask_filter_cb = ctypes.CFUNCTYPE(ctypes.c_bool, ctypes.POINTER(s
 class struct_nir_builder(Structure):
     pass
 
+struct_nir_builder._pack_ = 0 # source:False
 struct_nir_builder._fields_ = [
     ('cursor', nir_cursor),
     ('exact', ctypes.c_bool),

tinygrad/runtime/autogen/nv/nv.py

@@ -10,8 +10,6 @@ import ctypes
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/nv_gpu.py

@@ -10,8 +10,6 @@ import ctypes, os
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/nvrtc.py

@@ -31,8 +31,6 @@ def char_pointer_cast(string, encoding='utf-8'):
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/opencl.py

@@ -10,8 +10,6 @@ import ctypes, ctypes.util
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/qcom_dsp.py

@@ -10,8 +10,6 @@ import ctypes
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/sqtt.py

@@ -10,8 +10,6 @@ import ctypes, os
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/vfio.py

@@ -26,8 +26,6 @@ def _IOR(base, nr, type): return functools.partial(_do_ioctl, 2, ord(base) if is
 def _IOWR(base, nr, type): return functools.partial(_do_ioctl, 3, ord(base) if isinstance(base, str) else base, nr, type)
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/webgpu.py

@@ -10,8 +10,6 @@ import ctypes, tinygrad.runtime.support.webgpu as webgpu_support
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/ops_amd.py

@@ -31,7 +31,7 @@ AQL_HDR = (1 << hsa.HSA_PACKET_HEADER_BARRIER) | (hsa.HSA_FENCE_SCOPE_SYSTEM <<
 class ProfileSQTTEvent(ProfileEvent): device:str; se:int; props:dict; blob:bytes; itrace:bool # noqa: E702
 
 @dataclass(frozen=True)
-class PMCSample: name:str; block:str; xcc:int; inst:int; se:int; sa:int; wgp:int; off:int; size:int; reg:str # noqa: E702
+class PMCSample: name:str; block:str; inst:int; se:int; sa:int; wgp:int; off:int; size:int; reg:str # noqa: E702
 
 @dataclass(frozen=True)
 class ProfilePMCEvent(ProfileEvent): device:str; kern:str; sched:list[PMCSample]; blob:bytes # noqa: E702
@@ -74,12 +74,10 @@ class AMDComputeQueue(HWQueue):
 
   def set_grbm_broadcast(self):
     self.wreg(self.gc.regGRBM_GFX_INDEX, **{f'{f}_broadcast_writes': 1 for f in ['se', 'sh' if self.dev.target[0] == 9 else 'sa', 'instance']})
+  def set_grbm_se(self, se): self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, instance_broadcast_writes=1)
   def set_grbm_inst(self, n):
     self.wreg(self.gc.regGRBM_GFX_INDEX, **{f'{f}_broadcast_writes': 1 for f in ['se', 'sh' if self.dev.target[0] == 9 else 'sa']}, instance_index=n)
-  def set_grbm_se_sh(self, se, sh):
-    self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, **{f'{"sh" if self.dev.target[0] == 9 else "sa"}_index':sh}, instance_broadcast_writes=1)
-  def set_grbm_se_sh_wgp(self, se, sh, wgp): self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, sa_index=sh, instance_index=wgp << 2)
-  def set_grbm_se(self, se): self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, sh_broadcast_writes=1, instance_broadcast_writes=1)
+  def set_grbm_se_sh_wgp(self, se, sa, wgp): self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, sa_index=sa, instance_index=wgp << 2)
 
   def wait_reg_mem(self, value, mask=0xffffffff, mem=None, reg=None, reg_done=0, op=WAIT_REG_MEM_FUNCTION_GEQ):
     wrm_info_dw = self.pm4.WAIT_REG_MEM_MEM_SPACE(int(mem is not None)) | self.pm4.WAIT_REG_MEM_OPERATION(int(mem is None and reg_done > 0)) \
@@ -147,24 +145,18 @@ class AMDComputeQueue(HWQueue):
 
   def pmc_start(self, counters):
     self.pmc_reset_counters(en=False)
-    self.wreg(self.gc.regSQ_PERFCOUNTER_CTRL, cs_en=1, ps_en=1, gs_en=1, hs_en=1, **({'vmid_mask':0xffff} if (gfx9:=self.dev.target[0] == 9) else {}))
-    if self.dev.target[0] >= 11: self.wreg(self.gc.regSQ_PERFCOUNTER_CTRL2, force_en=1, vmid_en=0xffff)
+    self.wreg(self.gc.regSQ_PERFCOUNTER_CTRL, cs_en=1, ps_en=1, gs_en=1, hs_en=1)
+    self.wreg(self.gc.regSQ_PERFCOUNTER_CTRL2, force_en=1, vmid_en=0xffff)
 
-    end_off = 0
+    out_off = 0
     block2pid:dict[str, itertools.count] = collections.defaultdict(lambda: itertools.count())
     for name,block,idx in counters:
-      # sq block on gfx11+ goes down to wgps
-      inst_cnt, se_cnt, sa_cnt, wgp_cnt = {"GRBM": (1, 1, 1, 1), "GL2C": (32, 1, 1, 1), "TCC": (16, 1, 1, 1),
-        "SQ": (1, self.dev.se_cnt // self.dev.xccs) + ((1, 1) if gfx9 else (2, self.dev.iface.props['cu_per_simd_array'] // 2))}[block]
-      end_off += (rec_size:=prod((self.dev.xccs, inst_cnt, se_cnt, sa_cnt, wgp_cnt)) * 8)
+      inst_cnt, se_cnt, sa_cnt, wgp_cnt = {"GRBM": (1, 1, 1, 1), "GL2C": (32, 1, 1, 1),
+                                           "SQ": (1, self.dev.se_cnt, 2, self.dev.iface.props['cu_per_simd_array'] // 2)}[block]
+      reg, out_off = f'reg{block}_PERFCOUNTER{next(block2pid[block])}', out_off + (rec_size:=prod((inst_cnt, se_cnt, sa_cnt, wgp_cnt)) * 8)
+      self.wreg(getattr(self.gc, f'{reg}_SELECT'), idx)
+      self.dev.pmc_sched.append(PMCSample(name, block, inst_cnt, se_cnt, sa_cnt, wgp_cnt, out_off-rec_size, rec_size, reg))
 
-      if (regsel:=getattr(self.gc, (reg:=f'reg{block}_PERFCOUNTER{next(block2pid[block])}') + '_SELECT', None)) is None:
-        raise RuntimeError(f'{block} is out of perfcounter registers: ({reg} is not found)')
-
-      self.wreg(regsel, perf_sel=idx, **({'simd_mask':0xf, 'sqc_bank_mask':0xf, 'sqc_client_mask':0xf} if gfx9 and block == "SQ" else {}))
-      self.dev.pmc_sched.append(PMCSample(name, block, self.dev.xccs, inst_cnt, se_cnt, sa_cnt, wgp_cnt, end_off-rec_size, rec_size, reg))
-
-    if gfx9: self.wreg(self.gc.regSQ_PERFCOUNTER_MASK, sh0_mask=0xffff, sh1_mask=0xffff)
     self.wreg(self.gc.regCOMPUTE_PERFCOUNT_ENABLE, 1)
     return self.pmc_reset_counters(en=True)
 
@@ -175,17 +167,14 @@ class AMDComputeQueue(HWQueue):
     for s in sched:
       offset = itertools.count(s.off, step=8)
 
-      for xcc in range(s.xcc):
-        with self.pred_exec(xcc_mask=1 << xcc):
-          for inst, se_idx, sa_idx, wgp_idx in itertools.product(range(s.inst), range(s.se), range(s.sa), range(s.wgp)):
-            if s.inst > 1: self.set_grbm_inst(inst)
-            elif self.dev.target[0] == 9: self.set_grbm_se(se_idx)
-            else: self.set_grbm_se_sh_wgp(se_idx, sa_idx, wgp_idx)
+      for inst, se_idx, sa_idx, wgp_idx in itertools.product(range(s.inst), range(s.se), range(s.sa), range(s.wgp)):
+        if s.inst > 1: self.set_grbm_inst(inst)
+        else: self.set_grbm_se_sh_wgp(se_idx, sa_idx, wgp_idx)
 
-            # Copy counter to memory (src_sel = perf, dst_sel = tc_l2)
-            lo, hi = getattr(self.gc, f'{s.reg}_LO'), getattr(self.gc, f'{s.reg}_HI', None)
-            self.pkt3(self.pm4.PACKET3_COPY_DATA, (2 << 8) | 4, lo.addr[0], 0, *data64_le(buf.va_addr+(loff:=next(offset))))
-            if hi is not None: self.pkt3(self.pm4.PACKET3_COPY_DATA, (2 << 8) | 4, hi.addr[0], 0, *data64_le(buf.va_addr+loff+4))
+        # Copy counter to memory (src_sel = perf, dst_sel = tc_l2)
+        lo, hi = getattr(self.gc, f'{s.reg}_LO'), getattr(self.gc, f'{s.reg}_HI', None)
+        self.pkt3(self.pm4.PACKET3_COPY_DATA, 2 << 8 | 4, lo.addr[0], 0, *data64_le(buf.va_addr+(loff:=next(offset))))
+        if hi is not None: self.pkt3(self.pm4.PACKET3_COPY_DATA, 2 << 8 | 4, hi.addr[0], 0, *data64_le(buf.va_addr+loff+4))
 
     return self.pmc_reset_counters(en=True)
 
@@ -228,7 +217,7 @@ class AMDComputeQueue(HWQueue):
         if (se_mask >> se) & 0b1: mask |= (__SQTTINST:=1<<10) | (__SQTT_INST_PC:=1<<11) | (__SQTT_ISSUE:=1<<13)
 
         with self.pred_exec(xcc_mask=1<<(se // (ses_per_xcc:=(self.dev.se_cnt // self.dev.xccs)))):
-          self.set_grbm_se_sh(se % ses_per_xcc, 0)
+          self.set_grbm_se(se % ses_per_xcc)
           self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK, reg_mask=0xf, token_mask=mask)
           self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK2, inst_mask=0xffffffff)
           self.wreg(self.gc.regSQ_THREAD_TRACE_BASE, addr=lo32(buf0s[se].va_addr >> 12))
@@ -240,7 +229,7 @@ class AMDComputeQueue(HWQueue):
       self.spi_config(tracing=True)
       # One buffer for one SE, mesa does it with a single buffer and ac_sqtt_get_data_offset, but this is simpler and should work just as well
       for se in range(len(buf0s)):
-        self.set_grbm_se_sh(se, 0)
+        self.set_grbm_se(se)
 
         buf0_lo, buf0_hi = data64_le(buf0s[se].va_addr >> 12)
         if self.dev.target >= (12,0,0):
@@ -291,7 +280,7 @@ class AMDComputeQueue(HWQueue):
 
     # For each SE wait for finish to complete and copy regSQ_THREAD_TRACE_WPTR to know where in the buffer trace data ends
     for se in range(ses):
-      self.set_grbm_se_sh(se, 0)
+      self.set_grbm_se(se)
 
       status_reg = self.gc.regSQ_THREAD_TRACE_STATUS.addr[0] - (self.pm4.PACKET3_SET_UCONFIG_REG_START if self.dev.target[0] == 9 else 0)
       if self.dev.target >= (10, 0, 0):
@@ -762,8 +751,7 @@ class KFDIface:
 
     raise RuntimeError("\n".join(report))
 
-  def is_in_profile_mode(self):
-    return self.dev.target[0] == 9 or FileIOInterface(f'{self.dev_sysfs_path}/power_dpm_force_performance_level').read()[:16] == 'profile_standard'
+  def is_in_profile_mode(self): return FileIOInterface(f'{self.dev_sysfs_path}/power_dpm_force_performance_level').read()[:16] == 'profile_standard'
 
 class PCIIface(PCIIfaceBase):
   gpus:ClassVar[list[str]] = []
@@ -909,18 +897,17 @@ class AMDDevice(HCQCompiled):
 
     self.pmc_enabled = PROFILE and PMC > 0
     if self.pmc_enabled:
-      if self.target[0] not in {9, 11, 12}: raise RuntimeError(f'PMC are not supported on gc:{self.target}')
+      if self.target[0] not in {11, 12}: raise RuntimeError(f'PMC are not supported on gc:{self.target}')
       if not self.iface.is_in_profile_mode(): raise RuntimeError("PMC requires stable power state: run `amd-smi set -l stable_std` for KFD iface")
 
       self.pmc_sched:list[PMCSample] = []
       self.pmc_counters = import_pmc(self.target)
 
       # validate counters
-      pmc_default = "TCC_HIT,TCC_MISS,SQ_LDS_BANK_CONFLICT" if self.target[0] == 9 else "GL2C_HIT,GL2C_MISS,SQC_LDS_IDX_ACTIVE,SQC_LDS_BANK_CONFLICT"
-      for k in (PMC_COUNTERS:=getenv("PMC_COUNTERS", pmc_default).split(",")):
+      for k in (PMC_COUNTERS:=getenv("PMC_COUNTERS", "GL2C_HIT,GL2C_MISS,SQC_LDS_IDX_ACTIVE,SQC_LDS_BANK_CONFLICT").split(",")):
         if k not in self.pmc_counters: raise RuntimeError(f"PMC counter {k} is not supported. Available: {','.join(self.pmc_counters.keys())}")
 
-      cast(AMDComputeQueue, self.hw_compute_queue_t()).pmc_start([(k, *self.pmc_counters[k]) for k in PMC_COUNTERS]).submit(self)
+      cast(AMDComputeQueue, self.hw_compute_queue_t()).pmc_start([self.pmc_counters[k] for k in PMC_COUNTERS]).submit(self)
       self.pmc_buffer = self.allocator.alloc(self.pmc_sched[-1].off + self.pmc_sched[-1].size, BufferSpec(nolru=True, uncached=True))
       self.allocator._copyin(self.pmc_buffer, memoryview(bytearray(self.pmc_buffer.size))) # zero pmc buffers, some counters have only lo part.
 

tinygrad/runtime/support/amd.py

@@ -63,16 +63,10 @@ def import_soc(ip):
 
 def import_ip_offsets(ip): return type("IPOFF", (object,), import_header(f"include/{('sienna_cichlid' if ip[0] > 9 else 'vega20')}_ip_offset.h"))
 
-def import_pmc(ip) -> dict[str, tuple[str, int]]:
-  res:dict[str, tuple[str, int]] = {}
-  arch = f"gfx{ip[0]}{ip[1]:x}{ip[2]:x}"
-
-  for sec in header_download("rocprofiler-compute/src/rocprof_compute_soc/profile_configs/counter_defs.yaml", url=ROCM_URL).split('- name: ')[1:]:
-    for arch_spec in sec.split('- architectures:')[1:]:
-      if arch in arch_spec and (block:=re.search(r'block:\s*([A-Za-z0-9_]+)', arch_spec)) and (ev:=re.search(r'event:\s*(\d+)', arch_spec)):
-        res[sec.splitlines()[0].strip()] = (block.group(1), int(ev.group(1)))
-
-  return res
+def import_pmc(ip) -> dict[str, tuple[str, str, int]]:
+  ver = min(ip[0], 11) # 12 is same as 11
+  m = re.search(rf'<gfx{ver}>(.*?)</gfx{ver}>', header_download("rocprofiler/src/core/counters/basic/gfx_metrics.xml", url=ROCM_URL), re.S)
+  return {n:(n,b,int(e)) for n,b,e in re.findall(r'<metric name="([A-Za-z0-9_]+)" block="([A-Za-z0-9_]+)" event="([0-9]+)"', m.group(1))} if m else {}
 
 def import_asic_regs(prefix:str, version:tuple[int, ...], cls=AMDReg) -> dict[str, AMDReg]:
   def _split_name(name): return name[:(pos:=next((i for i,c in enumerate(name) if c.isupper()), len(name)))], name[pos:]

tinygrad/schedule/indexing.py

@@ -51,7 +51,7 @@ class IndexingContext:
     return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(dtypes.index, 0)
 
 def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
-  if x.op in {Ops.BUFFERIZE, Ops.INDEX}: return None
+  if x.op in {Ops.BUFFERIZE, Ops.INDEX, Ops.KERNEL}: return None
   if x.op is Ops.AFTER and x.src[1].op is Ops.KERNEL: return None
   new_srcs = []
   for s in x.src:
@@ -155,10 +155,6 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
   ending_ranges: dict[UOp, list[UOp]] = {}
   for x in tsink_reverse_toposort:
     if x.op in {Ops.DEVICE, Ops.UNIQUE}: continue
-
-    # no ranges on kernels, they are internal
-    if x.op is Ops.KERNEL: continue
-
     if x.dtype.scalar() == dtypes.index: continue  # TODO: why do I need this?
     ending_ranges[x] = sum([ending_ranges.get(u, []) for u in consumer_map[x]], [])
 

tinygrad/schedule/rangeify.py

@@ -2,9 +2,9 @@ from dataclasses import dataclass, field
 import itertools
 from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, ssimplify, KernelInfo
-from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType, BottomUpGate, Kernel, _remove_all_tags
+from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType, BottomUpGate
 from tinygrad.uop.symbolic import symbolic_flat
-from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY
+from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, Metadata, DEBUG_RANGEIFY
 from tinygrad.helpers import PCONTIG, partition, get_single_element, unwrap
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
@@ -18,10 +18,6 @@ sys.setrecursionlimit(10000)
 pm_mops = PatternMatcher([
   (UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
    lambda r,idx: r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idx.src[1:]), dtype=idx.dtype, arg=idx.arg)),  # type: ignore
-  # move movement ops after AFTER
-  (UPat(GroupOp.Movement, name="r").after(name="a", allow_any_len=True),
-   lambda r,a: UOp(r.op, r.dtype, (a.replace(src=(r.src[0],)+a.src[1:], tag=None),)+r.src[1:], r.arg, tag=a.tag)),
-  (UPat(GroupOp.Movement, name="r").end(name="a", allow_any_len=True), lambda r,a: a.replace(src=(r.src[0],)+a.src[1:])),
 ])
 
 # *****************
@@ -355,9 +351,6 @@ pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
   # move RESHAPEs through MSELECT/MSTACK
   (UPat((Ops.MSELECT, Ops.MSTACK), src=UPat(Ops.RESHAPE), name="m"),
    lambda m: m.replace(src=tuple([x.src[0].base for x in m.src]), tag=None).reshape(m.shape).rtag(m.tag)),
-
-  # remove any RESHAPEs on KERNEL
-  (UPat(Ops.KERNEL, name="k"), lambda k: k.replace(src=tuple(x.src[0] if x.op is Ops.RESHAPE else x for x in k.src))),
 ])
 
 pm_add_buffers_local = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
@@ -461,13 +454,19 @@ def remove_metadata_tags(ctx:LocalAddBufferContext, x:UOp):
   return x.replace(tag=None)
 
 pm_remove_tags = PatternMatcher([
+  # remove all the tags
   (UPat(GroupOp.All, name="x"), remove_metadata_tags),
 ])
 
 pm_add_range_tags = PatternMatcher([
-  (UPat(Ops.RANGE, name="x"), lambda x: x.rtag(())),
+  (UPat(Ops.RANGE, name="x"), lambda x: x.rtag(()))
 ])
 
+@dataclass(frozen=True)
+class Kernel:
+  ast: UOp
+  metadata: tuple[Metadata, ...] = ()
+
 def split_store(ctx:list[UOp], x:UOp) -> UOp|None:
   if len(x.ranges): return None
 
@@ -504,7 +503,7 @@ def tag_uop(ctx:list[UOp], x:UOp):
   return x.replace(tag=(len(ctx)-1,))
 add_tags = PatternMatcher([
   # don't tag BUFFERs, they are global
-  (UPat(GroupOp.All-{Ops.BUFFER, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.DEFINE_VAR, Ops.BIND, Ops.KERNEL,
+  (UPat(GroupOp.All-{Ops.BUFFER, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.DEFINE_VAR, Ops.BIND,
                      Ops.MSTACK, Ops.MSELECT, Ops.RANGE}.union(GroupOp.Movement), name="x"), tag_uop),
   (UPat({Ops.MSTACK, Ops.MSELECT}, name="x"), lambda ctx,x: None if all(s.op is Ops.BUFFER for s in x.src) else tag_uop(ctx, x)),
 ])
@@ -531,7 +530,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   uop_list: list[UOp] = []
   tsink = graph_rewrite(sink, add_tags, ctx=uop_list, bottom_up=True, name="number the uops")
 
-  tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites+replace_contiguous, ctx={}, name="earliest rewrites")
+  tsink = graph_rewrite(tsink, earliest_rewrites+replace_contiguous, ctx={}, name="earliest rewrites")
 
   # convert movement ops to ranges
   tsink, rctx = run_rangeify(tsink, DEBUG_RANGEIFY)
@@ -543,7 +542,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   # rebuild the sink with all the BUFFERIZEs with tags, this is what's ending up in the tensor graph
   # MSTACK stacks multiple BUFFERIZEs in one tagged tensor
   # if it's not tagged by here, it's out
-  tsink = UOp.sink(*[x for x in tsink.backward_slice if x.base.op in {Ops.BUFFERIZE, Ops.MSTACK, Ops.CONST, Ops.BUFFER, Ops.AFTER} and \
+  tsink = UOp.sink(*[x for x in tsink.backward_slice if x.base.op in {Ops.BUFFERIZE, Ops.MSTACK, Ops.CONST, Ops.BUFFER} and \
                      x.tag is not None and len(x.tag)])
 
   if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Tagged Rangeify")
@@ -568,14 +567,10 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
 
   if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
 
-  # TODO: we can probably get this earlier
-  sink_tags = [s.tag for s in tsink.src]
-  tsink = graph_rewrite(tsink, _remove_all_tags, name="remove all tags")
-
   becomes_map: dict[UOp, UOp] = {}
-  for tag, s in zip(sink_tags, tsink.src):
-    assert tag is not None
-    for a in tag:
+  for s in tsink.src:
+    assert s.tag is not None
+    for a in s.tag:
       if a is None: continue
-      becomes_map[uop_list[int(a)]] = s
+      becomes_map[uop_list[int(a)]] = s.replace(tag=None)
   return becomes_map

tinygrad/tensor.py

@@ -10,7 +10,7 @@ from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, p
 from tinygrad.helpers import suppress_finalizing
 from tinygrad.gradient import compute_gradient
 from tinygrad.uop.mathtraits import MathTrait
-from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, identity_element, all_metadata, _index_to_concrete_int, sint_to_uop
+from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, identity_element, all_metadata, _index_to_concrete_int, sint_to_uop, srender
 from tinygrad.uop.spec import type_verify, tensor_spec
 from tinygrad.device import Device, Buffer
 from tinygrad.engine.realize import run_schedule
@@ -239,14 +239,6 @@ class Tensor(MathTrait):
     _apply_map_to_tensors(becomes_map, name="Apply Kernelize Map")
     return self
 
-  def custom_kernel(self, *lst:Tensor, fxn:Callable, grad_fxn:Callable|None=None) -> list[Tensor]:
-    """
-    Call into a custom kernel written in UOps. Returns the Tensors after the Kernel has been applied.
-
-    This API is alpha and may change.
-    """
-    return [Tensor(u) for u in UOp.custom_kernel(*[t.uop for t in (self,)+lst], fxn=fxn, grad_fxn=grad_fxn)]
-
   def schedule_with_vars(self, *lst:Tensor) -> tuple[list[ScheduleItem], dict[str, int]]:
     """
     Creates the schedule needed to realize these Tensor(s), with Variables.
@@ -1038,7 +1030,28 @@ class Tensor(MathTrait):
 
   # ***** movement low level ops *****
 
-  def _mop(self, op:Ops, arg) -> Tensor: return self._apply_uop(UOp._mop, extra_args=(op,), arg=arg)
+  def view(self, shape:tuple[sint, ...], *args) -> Tensor:
+    """`.view` is an alias for `.reshape`."""
+    return self.reshape(shape, *args)
+
+  def reshape(self, shape, *args) -> Tensor:
+    """
+    Returns a tensor with the same data as the original tensor but with a different shape.
+    `shape` can be passed as a tuple or as separate arguments.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor.arange(6)
+    print(t.reshape(2, 3).numpy())
+    ```
+    """
+    # resolve None and args
+    new_shape = tuple([s if s is not None else self.shape[i] for i,s in enumerate(argfix(shape, *args))])
+    # resolve -1
+    if (c := new_shape.count(-1)) > 1: raise RuntimeError(f"only one dimension can be inferred using -1, getting {new_shape}")
+    if c: new_shape = tuple([-prod(self.shape) // prod(new_shape) if s == -1 else s for s in new_shape])
+    if resolve(prod(self.shape) != prod(new_shape), True):
+      raise ValueError(f"size mismatch, can't reshape ({', '.join(srender(d) for d in self.shape)}) -> ({', '.join(srender(d) for d in new_shape)})")
+    return self._apply_uop(UOp.reshape, arg=new_shape) if new_shape != self.shape else self
 
   def expand(self, shape, *args) -> Tensor:
     """

tinygrad/uop/mathtraits.py

@@ -1,10 +1,6 @@
-from typing import TypeVar, TypeAlias, TYPE_CHECKING
+from typing import TypeVar
 from tinygrad.uop import Ops
 from tinygrad.dtype import dtypes, ConstType
-from tinygrad.helpers import prod, argfix
-if TYPE_CHECKING:
-  from tinygrad.uop.ops import UOp
-  sint:TypeAlias = UOp|int
 
 TMT = TypeVar("TMT", bound="MathTrait")
 class MathTrait:
@@ -175,32 +171,3 @@ class MathTrait:
   def exp2(self): return self.alu(Ops.EXP2)
   def pow(self:TMT, x:TMT|ConstType): return self.alu(Ops.POW, self.ufix(x))
   def __pow__(self:TMT, x:TMT|ConstType): return self.pow(x)
-
-  # **** movement ops ****
-
-  # required to implement
-  def _mop(self:TMT, op:Ops, arg) -> TMT: raise NotImplementedError
-  @property
-  def shape(self) -> tuple["sint", ...]: raise NotImplementedError
-
-  def view(self:TMT, shape, *args) -> TMT:
-    """`.view` is an alias for `.reshape`."""
-    return self.reshape(shape, *args)
-
-  def reshape(self:TMT, shape, *args) -> TMT:
-    """
-    Returns a tensor with the same data as the original tensor but with a different shape.
-    `shape` can be passed as a tuple or as separate arguments.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor.arange(6)
-    print(t.reshape(2, 3).numpy())
-    ```
-    """
-    # resolve None and args
-    new_shape = tuple([s if s is not None else self.shape[i] for i,s in enumerate(argfix(shape, *args))])
-    # resolve -1
-    if (c := new_shape.count(-1)) > 1: raise RuntimeError(f"only one dimension can be inferred using -1, getting {new_shape}")
-    if c: new_shape = tuple([-prod(self.shape) // prod(new_shape) if s == -1 else s for s in new_shape])
-    if prod(self.shape) != prod(new_shape): raise ValueError(f"size mismatch, can't reshape ({self.shape}) -> ({new_shape})")
-    return self._mop(Ops.RESHAPE, arg=new_shape) if new_shape != self.shape else self

tinygrad/uop/ops.py

@@ -188,7 +188,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     match self.op:
       # late ops don't have shape
       case Ops.UNIQUE | Ops.DEVICE | Ops.RANGE | Ops.INDEX | Ops.LOAD | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
-           Ops.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.PRECAST | Ops.CONTRACT:
+           Ops.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.PRECAST:
         return None
 
       # some ops init the shape
@@ -338,13 +338,10 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   def group(*srcs:UOp|None):  # pylint: disable=no-self-argument
     if len(srcs) == 1 and isinstance(srcs[0], UOp): return srcs[0]
     return UOp(Ops.GROUP, dtypes.void, tuple([x for x in srcs if x is not None]))
-  def vectorize(self, *srcs, **kwargs):
-    return UOp(Ops.VECTORIZE, self.dtype.vec(len(srcs)+1), (self,)+srcs, **kwargs)
   def detach(self): return UOp(Ops.DETACH, self.dtype, (self,))
   def index(self, *srcs:UOp|None, ptr=False, **kwargs):
     return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype if ptr else self.dtype.base), (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
-  def __getitem__(self, idx):
-    return self.index(*[UOp.const(dtypes.index, x) if isinstance(x, int) else x for x in argfix(idx)])
+  def __getitem__(self, idx): return self.index(*argfix(idx))
   def const_like(self, b:ConstLike):
     # constants can optionally have a DEVICE source
     return UOp.const(self.dtype, b, device=self._device, shape=self._shape)
@@ -375,9 +372,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   def after(self, *src:UOp, **kwargs): return UOp(Ops.AFTER, self.dtype, (self,)+src, **kwargs)
   def assign(self, x:UOp): return UOp(Ops.ASSIGN, self.dtype, (self, x))
   def barrier(self, *src:UOp): return UOp(Ops.BARRIER, src=(self,)+src)
-  def contract(self, *rngs:UOp):
-    assert all(x.arg[-1] == AxisType.UPCAST for x in rngs), "all contract ranges must be upcast"
-    return UOp(Ops.CONTRACT, dtype=self.dtype.vec(prod([x.vmax+1 for x in rngs])), src=(self,), arg=tuple((x.arg[0], x.vmax+1) for x in rngs))
   def alu(self, op, *src:UOp, **kwargs):
     out_dtype = (self, *src)[-1].dtype
     if op in {Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ}: out_dtype = dtypes.bool.vec(out_dtype.count) if out_dtype.count > 1 else dtypes.bool
@@ -420,7 +414,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     return self.op is Ops.BUFFER
 
   def contiguous(self, *args, **kwargs):
-    if self.op is Ops.CONTIGUOUS: return self
     if self.is_contiguous(): return self
     return UOp(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs)
   def contiguous_backward(self): return self.alu(Ops.CONTIGUOUS_BACKWARD)
@@ -533,7 +526,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   # in these four, if the shape doesn't change we can return self
   def forced_reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, same_shape_noop=False)
-  #def reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, same_shape_noop=True)
+  def reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, same_shape_noop=True)
   def expand(self, arg:tuple[sint, ...]): return self._mop(Ops.EXPAND, arg, same_shape_noop=True)
   def shrink(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.SHRINK, arg, same_shape_noop=True)
   def pad(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.PAD, arg, same_shape_noop=True)
@@ -764,28 +757,20 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   # *** uop high level syntactic sugar ***
 
-  def shrink_to(self, arg:tuple[sint, ...]): return self.shrink(tuple([(0,x) for x in arg]))
-
   @staticmethod
-  def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL):
+  def placeholder(dtype:DType, shape:tuple[int, ...], slot:int, addrspace=AddrSpace.GLOBAL):
     lookup = {AddrSpace.GLOBAL: Ops.DEFINE_GLOBAL, AddrSpace.LOCAL: Ops.DEFINE_LOCAL, AddrSpace.REG: Ops.DEFINE_REG}
     ret = UOp(lookup[addrspace], dtype.ptr(prod(shape), addrspace), arg=slot)
     if len(shape) > 1: ret = ret.reshape(shape)
     return ret
   def placeholder_like(self, slot:int):
     assert all_int(self.shape), "no placeholder-like on symbolic shape"
-    return UOp.placeholder(self.shape, self.dtype, slot)
+    return UOp.placeholder(self.dtype, self.shape, slot)
 
   # set is store+end+after
   def set(self:UOp, val:UOp|ConstType, end:UOp|tuple[UOp, ...]=()) -> UOp:
     return self.src[0].after(self.store(UOp.const(self.dtype, val) if not isinstance(val, UOp) else val).end(*argfix(end)))
 
-  def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]:
-    placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(srcs)]
-    contig_srcs = tuple(x.contiguous() for x in srcs)
-    kernel = UOp(Ops.KERNEL, src=tuple(x.base for x in contig_srcs), arg=Kernel(fxn(*placeholders), grad_fxn=grad_fxn))
-    return [s.after(kernel) for s in contig_srcs]
-
 @dataclass(frozen=True)
 class KernelInfo:
   name: str = "test"            # name of the kernel
@@ -796,12 +781,6 @@ class KernelInfo:
   @property
   def function_name(self): return to_function_name(self.name)
 
-@dataclass(frozen=True)
-class Kernel:
-  ast: UOp
-  metadata: tuple[Metadata, ...] = ()
-  grad_fxn: Callable|None = None
-
 # ******** ops in python ********
 
 def safe_exp2(x):
@@ -1125,7 +1104,7 @@ if TRACK_MATCH_STATS or PROFILE:
   def launch_viz(env_str:str, data:str):
     os.environ[env_str] = "0"
     os.environ[f"{env_str}_DATA"] = data
-    if not int(os.getenv("VIZ", "0")) and not int(os.getenv("PROFILE", "0")) and not CI:
+    if not int(os.getenv("VIZ", "0")) and not int(os.getenv("PROFILE", "0")) and not int(os.getenv("SQTT", "0")) and not CI:
       args = ['--kernels', getenv("VIZ_DATA", "")] if getenv("VIZ_DATA", "") else []
       args += ['--profile', getenv("PROFILE_DATA", "")] if getenv("PROFILE_DATA", "") else []
       viz_path = pathlib.Path(__file__).resolve().parent.parent / "viz" / "serve.py"
@@ -1262,7 +1241,6 @@ pm_lower_index_dtype = PatternMatcher([
 def _index_to_concrete_int(u:UOp): return graph_rewrite(u.sink(), pm_lower_index_dtype).src[0]
 
 _substitute = PatternMatcher([(UPat(tuple(Ops), name="x"), lambda ctx,x: ctx.get(x,None))])
-_remove_all_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
 
 def do_unbind(ctx:dict[Variable, int], x:UOp):
   v,i = x.unbind()
@@ -1274,7 +1252,7 @@ pm_unbind = PatternMatcher([(UPat(Ops.BIND, name="x"), do_unbind)])
 syms = { Ops.ADD: "+", Ops.SUB: "-", Ops.IDIV: "//", Ops.MOD: "%", Ops.SHL: "<<", Ops.SHR: ">>",
          Ops.MUL: "*", Ops.CMPLT: "<", Ops.CMPNE: "!=", Ops.AND: "&", Ops.OR: "|", Ops.XOR: "^"}
 # comparison operators are not in here because they are chained in python, not left-associative
-precedence = {Ops.MUL:1, Ops.IDIV:1, Ops.MOD:1, Ops.ADD:2, Ops.SUB:2, Ops.SHL:3, Ops.SHR:3, Ops.AND:4, Ops.XOR:5, Ops.OR:6}
+precedence = {Ops.NEG:0, Ops.MUL:1, Ops.IDIV:1, Ops.MOD:1, Ops.ADD:2, Ops.SUB:2, Ops.SHL:3, Ops.SHR:3, Ops.AND:4, Ops.XOR:5, Ops.OR:6}
 def strip_binary_parens(x:UOp, left:str, right:str, code_for_op) -> str:
   if x.op not in precedence: return code_for_op(left, right)
   return code_for_op(strip_parens(left) if precedence.get(x.src[0].op,99)<=precedence[x.op] else left, strip_parens(right) if
@@ -1345,12 +1323,10 @@ pm_pyrender_extra = PatternMatcher([
   (UPat(GroupOp.Movement, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({render_marg(ctx,x)})"),
   # NOTE: CMPNE doesn't work cause there's no __rne__
   (UPat(set(syms.keys())-{Ops.SUB, Ops.CMPNE}, src=(UPat(Ops.CONST, name="y"), UPat(name="z")), name="x"),
-   lambda ctx,x,y,z: strip_binary_parens(x, str(y.arg), ctx[z], lambda a,b: f"({a}{syms[x.op]}{b})")),
+   lambda ctx,x,y,z: f"({y.arg}{syms[x.op]}{ctx[z]})"),
   # NOTE: sub doesn't work cause it's written as add/mul
-  (UPat(set(syms.keys())-{Ops.SUB}, src=(UPat(name="y"), UPat(Ops.CONST, name="z")), name="x"), lambda ctx,x,y,z:
-    strip_binary_parens(x, ctx[y], str(z.arg), lambda a,b: f"({a}{syms[x.op]}{b})")),
-  (UPat(set(syms.keys())-{Ops.SUB}, name="x"), lambda ctx,x:
-    strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
+  (UPat(set(syms.keys())-{Ops.SUB}, src=(UPat(name="y"), UPat(Ops.CONST, name="z")), name="x"), lambda ctx,x,y,z: f"({ctx[y]}{syms[x.op]}{z.arg})"),
+  (UPat(set(syms.keys())-{Ops.SUB}, name="x"), lambda ctx,x: f"({ctx[x.src[0]]}{syms[x.op]}{ctx[x.src[1]]})"),
   (UPat(sugar, src=(), name="x"), lambda x: f"UOp.{x.op.name.lower()}("+', '.join(([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
   (UPat(sugar, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}("+', '.join([ctx[y] for y in x.src[1:]] + \
     ([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
@@ -1396,7 +1372,7 @@ def pyrender(ast:UOp) -> str:
     else:
       r[u] = f"c{i}" if u is not lst[-1] else "ast"
       ret[r[u]] = ren
-  return ''.join([v[1] for v in kernels.values()]) + '\n'.join([f"{k} = {strip_parens(v)}" for k,v in ret.items()])
+  return ''.join([v[1] for v in kernels.values()]) + '\n'.join([f"{k} = {v}" for k,v in ret.items()])
 
 # *** what was symbolic.py ***
 

tinygrad/uop/spec.py

@@ -1,6 +1,6 @@
 import math
 from typing import cast, Any
-from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, AxisType, KernelInfo, pyrender, Kernel
+from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, AxisType, KernelInfo, pyrender
 from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid
 from tinygrad.helpers import DEBUG, Context, prod, SPEC, Metadata
 from tinygrad.uop.validate import validate_index
@@ -44,20 +44,7 @@ shared_spec = PatternMatcher([
 
 # ***** UOp spec in the Tensor graph *****
 
-movement_ops = PatternMatcher([
-  (UPat((Ops.RESHAPE, Ops.EXPAND), name="mv", src=(UPat.var("x"), UPat(dtype=dtypes.index))), lambda mv,x: True),
-  (UPat((Ops.PAD, Ops.SHRINK), name="mv", src=(UPat.var("x"), UPat(dtype=dtypes.index), UPat(dtype=dtypes.index))), lambda mv,x: True),
-  (UPat((Ops.PERMUTE, Ops.FLIP), name="mv", src=(UPat.var("x"),)), lambda mv,x: isinstance(mv.arg, tuple)),
-
-  # inputs to movement ops
-  (UPat((Ops.VECTORIZE, Ops.VCONST), dtype=dtypes.index), lambda: True),
-  (UPat({Ops.ADD, Ops.MUL, Ops.IDIV}, dtype=dtypes.index), lambda: True),
-
-  # AFTER on Movement Op
-  (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement),), allow_any_len=True), lambda: True),
-])
-
-_tensor_spec = PatternMatcher([
+tensor_spec = PatternMatcher([
   # buffer spec
   (UPat(Ops.UNIQUE, dtypes.void, ()), lambda: True),
   (UPat(Ops.DEVICE, dtypes.void, (), name="d"), lambda d:
@@ -69,7 +56,7 @@ _tensor_spec = PatternMatcher([
   (UPat(Ops.BUFFER_VIEW, src=(UPat(Ops.MSTACK, src=UPat(Ops.BUFFER)),)), lambda: True),
 
   # KERNEL can attach to an AFTER to describe the compute required to realize a BUFFER
-  (UPat(Ops.KERNEL, src=UPat((Ops.BUFFER, Ops.BUFFER_VIEW, Ops.AFTER, Ops.MSELECT, Ops.MSTACK, Ops.BIND, Ops.CONTIGUOUS))), lambda: True),
+  (UPat(Ops.KERNEL, src=UPat((Ops.BUFFER, Ops.BUFFER_VIEW, Ops.AFTER, Ops.MSELECT, Ops.MSTACK, Ops.BIND))), lambda: True),
 
   # ASSIGN has a target and a value. It can also optionally depend on other assigns
   (UPat(Ops.ASSIGN, name="x"), lambda x: len(x.src) >= 2 and all(s.op is Ops.ASSIGN for s in x.src[2:])),
@@ -80,6 +67,14 @@ _tensor_spec = PatternMatcher([
   # MSTACK combines buffers into multi
   (UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(x.device, str) for x in x.src)),
 
+  (UPat((Ops.RESHAPE, Ops.EXPAND), name="mv", src=(UPat.var("x"), UPat(dtype=dtypes.index))), lambda mv,x: True),
+  (UPat((Ops.PAD, Ops.SHRINK), name="mv", src=(UPat.var("x"), UPat(dtype=dtypes.index), UPat(dtype=dtypes.index))), lambda mv,x: True),
+  (UPat((Ops.PERMUTE, Ops.FLIP), name="mv", src=(UPat.var("x"),)), lambda mv,x: isinstance(mv.arg, tuple)),
+
+  # inputs to movement ops
+  (UPat((Ops.VECTORIZE, Ops.VCONST), dtype=dtypes.index), lambda: True),
+  (UPat({Ops.ADD, Ops.MUL, Ops.IDIV}, dtype=dtypes.index), lambda: True),
+
   # Tensor variable bindings
   (UPat(Ops.BIND, (dtypes.int,dtypes.index,), (UPat(Ops.DEFINE_VAR), UPat.cvar(dtype=(dtypes.int,dtypes.index,))), arg=None), lambda: True),
 
@@ -109,12 +104,7 @@ _tensor_spec = PatternMatcher([
 
   # AFTER if things were kernelized
   (UPat(Ops.AFTER, src=(UPat((Ops.BUFFER, Ops.AFTER)),), allow_any_len=True), lambda: True),
-])+movement_ops+shared_spec
-
-tensor_spec = PatternMatcher([
-  # no tags allowed in tensor graph
-  (UPat(GroupOp.All, name="x"), lambda x: None if x.tag is None else False),
-])+_tensor_spec
+])+shared_spec
 
 # ***** UOp spec in codegen shared between kernel and program *****
 
@@ -163,6 +153,11 @@ shared_codegen_spec = PatternMatcher([
 # ***** UOp spec in kernel graph *****
 
 kernel_spec = PatternMatcher([
+  # RESHAPE (but only RESHAPE) is allowed here
+  (UPat(Ops.RESHAPE, name="mv", src=(UPat.var("x"), UPat(dtype=dtypes.index))), lambda mv,x: True),
+  (UPat(Ops.AFTER, src=(UPat(Ops.RESHAPE),), allow_any_len=True), lambda: True),
+  (UPat(Ops.VCONST, dtype=dtypes.index), lambda: True),
+
   # index is allowed here
   (UPat(GroupOp.Elementwise|{Ops.CONST, Ops.RANGE, Ops.DEFINE_VAR}, dtype=dtypes.index), lambda: True),
 
@@ -174,7 +169,7 @@ kernel_spec = PatternMatcher([
 
   # reduce must be on ranges
   (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:])),
-])+movement_ops+shared_codegen_spec+shared_spec
+])+shared_codegen_spec+shared_spec
 
 # ***** UOp spec in linearized programs *****
 
@@ -251,7 +246,7 @@ full_spec = PatternMatcher([
   (UPat(Ops.DEFINE_VAR, dtype=dtypes.floats), lambda: True),
   # allow any AFTER
   (UPat(Ops.AFTER, src=(UPat(),), allow_any_len=True), lambda: True),
-])+_tensor_spec+kernel_spec+program_spec+shared_spec
+])+tensor_spec+kernel_spec+program_spec+shared_spec
 
 # ***** uop helpers *****
 
@@ -267,7 +262,7 @@ def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher):
 
 # late imports to avoid circular import
 from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.schedule.rangeify import BufferizeOpts
+from tinygrad.schedule.rangeify import BufferizeOpts, Kernel
 glbls:dict[str, Any] = {"inf": math.inf, "nan": math.nan, "KernelInfo": KernelInfo, "Kernel": Kernel, "Metadata": Metadata,
                         "UOp": UOp, "dtypes": dtypes, "Ops": Ops, "AxisType": AxisType, "Invalid": Invalid,
                         "Opt": Opt, "OptOps": OptOps, "BufferizeOpts": BufferizeOpts, "AddrSpace": AddrSpace}

tinygrad/uop/symbolic.py

@@ -514,8 +514,9 @@ sym = symbolic_flat+pm_simplify_valid+PatternMatcher([
    lambda x,y,alu: UOp(Ops.VECTORIZE, alu.dtype, (UOp(alu.op, alu.dtype.scalar(), (x,y)),)*alu.dtype.count)),
   # VECTORIZE of a single element is just that element
   (UPat(Ops.VECTORIZE, src=(UPat(name='x'),)), lambda x: x),
-  # VECTORIZE void is GROUP
-  (UPat(Ops.VECTORIZE, dtype=dtypes.void, name='x'), lambda x: UOp.group(*x.src)),
+  # VECTORIZE void is SINK
+  (UPat(Ops.VECTORIZE, dtype=dtypes.void, src=UPat(Ops.BARRIER, name='b')), lambda b: b),
+  (UPat(Ops.VECTORIZE, dtype=dtypes.void, name='x'), lambda x: UOp(Ops.SINK, dtypes.void, x.src)),
   # tensor core with a 0 input is acc
   (UPat(Ops.WMMA, src=(UPat.const(None, 0.0), UPat.var(), UPat.var("acc"))), lambda acc: acc),
   (UPat(Ops.WMMA, src=(UPat.var(), UPat.const(None, 0.0), UPat.var("acc"))), lambda acc: acc),

tinygrad/viz/index.html

@@ -77,12 +77,6 @@
     display: inline-flex;
     align-items: center;
     gap: 4px;
-    line-height: 1;
-    user-select: none;
-    cursor: pointer;
-  }
-  input {
-    outline: none;
   }
   #graph svg {
     width: 100%;

tinygrad/viz/js/index.js

@@ -109,12 +109,12 @@ async function initWorker() {
   workerUrl = URL.createObjectURL(new Blob([(await Promise.all(resp.map((r) => r.text()))).join("\n")], { type: "application/javascript" }));
 }
 
-function renderDag(graph, additions, recenter, layoutOpts) {
+function renderDag(graph, additions, recenter) {
   // start calculating the new layout (non-blocking)
   updateProgress({ start:true });
   if (worker != null) worker.terminate();
   worker = new Worker(workerUrl);
-  worker.postMessage({graph, additions, opts:layoutOpts });
+  worker.postMessage({graph, additions});
   worker.onmessage = (e) => {
     displaySelection("#graph");
     updateProgress({ start:false });
@@ -623,10 +623,6 @@ window.addEventListener("popstate", (e) => {
   if (e.state != null) setState(e.state);
 });
 
-const toggleLabel = d3.create("label").text("Show indexing (r)").node();
-const toggle = d3.create("input").attr("type", "checkbox").attr("id", "show-indexing").property("checked", true).node();
-toggleLabel.prepend(toggle);
-
 async function main() {
   // ** left sidebar context list
   if (ctxs == null) {
@@ -739,13 +735,10 @@ async function main() {
     };
   }
   if (ret.length === 0) return;
-  // ** center UOp graph
-  const render = (opts) => renderDag(ret[currentRewrite].graph, ret[currentRewrite].changed_nodes ?? [], currentRewrite === 0, opts);
-  render({ showIndexing:toggle.checked });
-  toggle.onchange = (e) => render({ showIndexing:e.target.checked });
+  renderDag(ret[currentRewrite].graph, ret[currentRewrite].changed_nodes ?? [], currentRewrite === 0);
   // ** right sidebar code blocks
   const codeElement = codeBlock(ret[currentRewrite].uop, "python", { wrap:false });
-  metadata.replaceChildren(toggleLabel, codeBlock(step.code_line, "python", { loc:step.loc, wrap:true }), codeElement);
+  metadata.replaceChildren(codeBlock(step.code_line, "python", { loc:step.loc, wrap:true }), codeElement);
   // ** rewrite steps
   if (step.match_count >= 1) {
     const rewriteList = metadata.appendChild(document.createElement("div"));
@@ -862,10 +855,6 @@ document.addEventListener("keydown", (event) => {
     event.preventDefault()
     document.getElementById("zoom-to-fit-btn").click();
   }
-  // r key toggles indexing
-  if (event.key === "r") {
-    toggle.click();
-  }
 });
 
 main()

tinygrad/viz/js/worker.js

@@ -5,7 +5,7 @@ const ctx = canvas.getContext("2d");
 ctx.font = `350 ${LINE_HEIGHT}px sans-serif`;
 
 onmessage = (e) => {
-  const { graph, additions, opts } = e.data;
+  const { graph, additions } = e.data;
   const g = new dagre.graphlib.Graph({ compound: true });
   g.setGraph({ rankdir: "LR" }).setDefaultEdgeLabel(function() { return {}; });
   if (additions.length !== 0) g.setNode("addition", {label:"", labelWidth:0, labelHeight:0, className:"overlay"});
@@ -23,15 +23,6 @@ onmessage = (e) => {
     for (const [port, s] of src) g.setEdge(s, k, { label: edgeCounts[s] > 1 ? {type:"tag", text:edgeCounts[s]} : {type:"port", text:port}});
     if (additions.includes(parseInt(k))) g.setParent(k, "addition");
   }
-  // optionally hide nodes from the layuot
-  if (!opts.showIndexing) {
-    for (const n of g.nodes()) {
-      const node = g.node(n);
-      if (node.label.includes("dtypes.index")) g.removeNode(n);
-    }
-    // After all layout changes are complete, remove the overlay node if it's empty
-    if (!g.node("addition")?.width) g.removeNode("addition");
-  }
   dagre.layout(g);
   postMessage(dagre.graphlib.json.write(g));
   self.close();

tinygrad/viz/serve.py

@@ -14,7 +14,7 @@ from tinygrad.renderer import ProgramSpec
 from tinygrad.dtype import dtypes
 
 uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
-               Ops.DEFINE_GLOBAL:"#cb9037", **{x:"#f2cb91" for x in {Ops.DEFINE_LOCAL, Ops.DEFINE_REG}}, Ops.REDUCE_AXIS: "#FF6B6B",
+               **{x:"#f2cb91" for x in GroupOp.Defines}, Ops.REDUCE_AXIS: "#FF6B6B",
                Ops.RANGE: "#c8a0e0", Ops.ASSIGN: "#909090", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff",
                Ops.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55",
                **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
@@ -55,7 +55,7 @@ def pystr(u:UOp, i:int) -> str:
   try: return pyrender(u)
   except Exception: return str(u)
 
-def uop_to_json(x:UOp) -> dict[int, dict]:
+def uop_to_json(x:UOp, ignore_indexing=False) -> dict[int, dict]:
   assert isinstance(x, UOp)
   graph: dict[int, dict] = {}
   excluded: set[UOp] = set()
@@ -63,6 +63,7 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
     # always exclude DEVICE/CONST/UNIQUE
     if u.op in {Ops.DEVICE, Ops.CONST, Ops.UNIQUE} and u is not x: excluded.add(u)
     if u.op is Ops.VCONST and u.dtype.scalar() == dtypes.index and u is not x: excluded.add(u)
+    if u.dtype.scalar() is dtypes.index and ignore_indexing: excluded.update(u.backward_slice_with_self)
   for u in toposort:
     if u in excluded: continue
     argst = codecs.decode(str(u.arg), "unicode_escape")
@@ -103,14 +104,16 @@ def _reconstruct(a:int):
 def get_full_rewrite(ctx:TrackedGraphRewrite, i:int=0) -> Generator[GraphRewriteDetails, None, None]:
   next_sink = _reconstruct(ctx.sink)
   # in the schedule graph we don't show indexing ops (unless it's in a kernel AST or rewriting dtypes.index sink)
-  yield {"graph":uop_to_json(next_sink), "uop":pystr(next_sink,i), "changed_nodes":None, "diff":None, "upat":None}
+  ignore_indexing = trace.keys[i].display_name.startswith("Schedule") and not (ctx.name in {"kernel split"} or \
+      any(s.dtype is dtypes.index for s in next_sink.src+(next_sink,)))
+  yield {"graph":uop_to_json(next_sink, ignore_indexing), "uop":pystr(next_sink,i), "changed_nodes":None, "diff":None, "upat":None}
   replaces: dict[UOp, UOp] = {}
   for u0_num,u1_num,upat_loc,dur in tqdm(ctx.matches):
     replaces[u0:=_reconstruct(u0_num)] = u1 = _reconstruct(u1_num)
     try: new_sink = next_sink.substitute(replaces)
     except RuntimeError as e: new_sink = UOp(Ops.NOOP, arg=str(e))
     match_repr = f"# {dur*1e6:.2f} us\n"+printable(upat_loc)
-    yield {"graph":(sink_json:=uop_to_json(new_sink)), "uop":pystr(new_sink,i),
+    yield {"graph":(sink_json:=uop_to_json(new_sink, ignore_indexing)), "uop":pystr(new_sink,i),
            "changed_nodes":[id(x) for x in u1.toposort() if id(x) in sink_json],
            "diff":list(difflib.unified_diff(pystr(u0,i).splitlines(),pystr(u1,i).splitlines())), "upat":(upat_loc, match_repr)}
     if not ctx.bottom_up: next_sink = new_sink