early lower cat

.github/actions/setup-tinygrad/action.yml

@@ -45,10 +45,6 @@ inputs:
     description: "Install mesa"
     required: false
     default: 'false'
-  tinydreno:
-    description: "Install tinydreno"
-    required: false
-    default: 'false'
 runs:
   using: "composite"
   steps:
@@ -330,9 +326,3 @@ runs:
       if: inputs.mesa == 'true' && runner.os == 'macOS'
       shell: bash
       run: brew install sirhcm/tinymesa/tinymesa_cpu
-
-    # *** tinydreno ***
-    - name: Install tinydreno (linux)
-      if: inputs.tinydreno == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: sudo curl -fL https://github.com/sirhcm/tinydreno/raw/refs/heads/master/libllvm-qcom.so -o /usr/lib/libllvm-qcom.so

.github/workflows/test.yml

@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '18'
+  CACHE_VERSION: '17'
   CAPTURE_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
   PYTHONPATH: ${{ github.workspace }}
@@ -1011,26 +1011,3 @@ jobs:
           python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
           DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
           python -m pytest -n=auto test/backend/test_ops.py --durations=20
-  qcomclcompiletests:
-    name: Compile-only (QCOM CL)
-    runs-on: ubuntu-24.04-arm
-    timeout-minutes: 15
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: compile-qcomcl
-          deps: testing_unit
-          tinydreno: 'true'
-          python-version: '3.12'
-      - name: Set env
-        shell: bash
-        run: printf "NULL=1\nNULL_ALLOW_COPYOUT=1\nNULL_QCOMCL=1" >> $GITHUB_ENV
-      - name: Run test_ops
-        shell: bash
-        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
-          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20

examples/mlperf/model_train.py

@@ -1416,9 +1416,8 @@ def train_llama3():
     loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
     loss.backward()
     assert all(p.grad is g for p,g in zip(optim.params, grads))
-    loss_cpu = loss.flatten().float().to("CPU")
-    Tensor.realize(loss_cpu, *grads)
-    return loss_cpu
+    Tensor.realize(loss, *grads)
+    return loss.flatten().float().to("CPU")
 
   @TinyJit
   def optim_step():
@@ -1428,11 +1427,10 @@ def train_llama3():
     for g in grads:
       g.assign(g.zeros_like()).realize()
 
-    lr_cpu = optim.lr.float().to("CPU")
-    grad_norm_cpu = grad_norm.float().to("CPU")
-    Tensor.realize(lr_cpu, grad_norm_cpu, *grads)
+    lr = optim.lr
+    Tensor.realize(lr, *grads)
 
-    return lr_cpu, grad_norm_cpu
+    return lr.float().to("CPU"), grad_norm.float().to("CPU")
 
   @TinyJit
   @Tensor.train(False)

extra/callrange/amd_call_matmul.py

@@ -1,139 +0,0 @@
-from typing import Callable
-from tinygrad import UOp, dtypes, Device, Tensor, getenv, function
-from tinygrad.uop.ops import AxisType, AddrSpace
-
-def simple_function(fxn:Callable[..., UOp]) -> Callable[..., UOp]:
-  def wrapper(*args:UOp) -> UOp:
-    params:list[UOp] = [x.param_like(i) for i,x in enumerate(args)]
-    return fxn(*params).call(*args)
-  return wrapper
-
-THREADS_PER_BLOCK = 128
-WARP_SIZE = 32
-
-# Register tile sizes (per-thread accumulator tile of C)
-TN = 4     # columns per thread
-TM = 4     # rows per thread
-
-WAVE_TILE_N = 128
-WAVE_TILE_M = 32
-
-LANES_PER_WAVE_X = 8
-LANES_PER_WAVE_Y = 4
-ITERS_PER_WAVE_N = 4 #WAVE_TILE_N // (LANES_PER_WAVE_X * TN)
-ITERS_PER_WAVE_M = 2 #WAVE_TILE_M // (LANES_PER_WAVE_Y * TM)
-
-WAVES_IN_BLOCK_Y = 4
-WAVES_IN_BLOCK_X = 1
-
-
-N = getenv("N", 4096)
-M = K = N
-
-# Threadblock tile sizes (block-level tile of C that a block computes)
-BLOCK_N = 128   # columns of C (N-dim) per block
-BLOCK_M = 128   # rows of C (M-dim) per block
-BLOCK_K = 8     # K-slice per block iteration
-
-@simple_function
-def slice_matmul(c_regs, a_local, b_local):
-  # 2x
-  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)
-
-
-  pass
-
-@simple_function
-def compute_local(c:UOp, a_local:UOp, b_local:UOp) -> UOp:
-  # this is the LID level on the GPU, here we can define regs
-  tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
-  waveIdx = (tid // WARP_SIZE) % WAVES_IN_BLOCK_X
-  waveIdy = (tid // WARP_SIZE) // WAVES_IN_BLOCK_X
-  assert waveIdy.vmax+1 == WAVES_IN_BLOCK_Y
-
-  laneIdx = (tid % WARP_SIZE) % LANES_PER_WAVE_X
-  laneIdy = (tid % WARP_SIZE) // LANES_PER_WAVE_X
-  assert laneIdy.vmax+1 == LANES_PER_WAVE_Y
-
-  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)
-
-  # do the math
-  A_col = A_col.assign(a_local[k_tile].reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM)[waveIdy, :, laneIdy, :].flatten())
-  B_row = B_row.assign(b_local[k_tile].reshape(WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)[waveIdx, :, laneIdx, :].flatten())
-  c_regs += A_col.reshape(-1, 1) * B_row.reshape(1, -1)  #
-  c_regs
-
-
-
-@simple_function
-def load_local(a_local, b_local, a_global, b_global):
-  # NOTE: it ends this range, so there's a BARRIER
-  tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
-  return UOp.group(
-    a_local[:, tid].store(a_global[tid, :]),
-    b_local[:, tid].store(b_global[:, tid]))
-
-@simple_function
-def reg_matmul(c_regs, a_local, b_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)
-
-
-
-@simple_function
-def local_matmul(c:UOp, a:UOp, b:UOp, a_local:UOp, b_local:UOp):
-  tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
-  waveIdx = (tid // WARP_SIZE) % WAVES_IN_BLOCK_X
-  waveIdy = (tid // WARP_SIZE) // WAVES_IN_BLOCK_X
-  laneIdx = (tid % WARP_SIZE) % LANES_PER_WAVE_X
-  laneIdy = (tid % WARP_SIZE) // LANES_PER_WAVE_X
-
-  # this is the LID level on the GPU, this (and below) is where we define REGs
-  c_regs = UOp.placeholder((ITERS_PER_WAVE_M*TM, ITERS_PER_WAVE_N*TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
-
-  # 128x128, Kx128, Kx128
-  k_tile = UOp.range(N // BLOCK_K, 0, AxisType.REDUCE)*BLOCK_K
-  fxn = reg_matmul(c_regs.assign(0),
-                   a_local[:, tid].assign(a[k_tile:k_tile+BLOCK_K, tid]),
-                   b_local[:, tid].assign(b[k_tile:k_tile+BLOCK_K, tid]))
-
-  # do math
-  c = c.reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM,
-                WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)
-  return c[waveIdy, :, laneIdy, :, waveIdx, :, laneIdx, :].store(c_regs.after(fxn))
-
-@simple_function
-def global_matmul(c:UOp, a:UOp, b:UOp):
-  # this is the GID level on the GPU, this is where we define LOCAL buffers shared across lids
-  gx = UOp.range(N//BLOCK_N, 0, AxisType.GLOBAL) * BLOCK_N
-  gy = UOp.range(M//BLOCK_M, 1, AxisType.GLOBAL) * BLOCK_M
-  a_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
-  b_local = UOp.placeholder((BLOCK_K, BLOCK_M), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
-  return local_matmul(c[gx:gx+BLOCK_N, gy:gy+BLOCK_M], a.permute(1,0)[:, gx:gx+BLOCK_N], b[:, gy:gy+BLOCK_M], a_local, b_local)
-
-  #ll = load_local(a_local, b_local, a.permute(1,0)[:, gx:gx+BLOCK_N], b[:, gy:gy+BLOCK_M])
-  #return compute_local(c[gx:gx+BLOCK_N, gy:gy+BLOCK_M], a_local.after(ll), b_local.after(ll))
-
-if __name__ == "__main__":
-  # this is the outer lvel on the GPU, this is where we define GLOBAL buffers
-  C = Tensor.empty(N, M)
-  A = Tensor.randn(N, K)
-  B = Tensor.randn(K, M)
-  c_out = C.call(A, B, fxn=global_matmul).numpy()
-
-  #C = UOp.new_buffer(Device.DEFAULT, N*M, dtypes.float).reshape(N,M)
-  #A = UOp.new_buffer(Device.DEFAULT, N*K, dtypes.float).reshape(N,K)
-  #B = UOp.new_buffer(Device.DEFAULT, K*M, dtypes.float).reshape(K,M)
-  #global_matmul(C, A, B).realize()
-
-  # input matmuls
-  #c = UOp.param(0, dtypes.float, (N, M))
-  #a = UOp.param(1, dtypes.float, (N, K))
-  #b = UOp.param(2, dtypes.float, (K, M))
-
-
-  #ba = a.rearrange("(n bn) (k bk) -> n k bn bk", bn=BLOCK_N, bk=BLOCK_K)[gx, k_tile_range]
-  #bb = b.rearrange("(k bk) (m bm) -> k m bk bm", bk=BLOCK_K, bm=BLOCK_M)[k_tile_range, gy]
-  #bc = c.rearrange("(n bn) (m bm) -> n m bn bm", bn=BLOCK_N, bm=BLOCK_M)[gx, gy]

extra/callrange/test.py

@@ -1,85 +0,0 @@
-from tinygrad import UOp, dtypes, Device, Tensor
-
-if __name__ == "__main__":
-  B0 = UOp.new_buffer(Device.DEFAULT, 100, dtypes.float).reshape(10,10)
-  B1 = UOp.new_buffer(Device.DEFAULT, 100, dtypes.float).reshape(10,10)
-
-
-  b0 = UOp.param(0, dtypes.float, (10,10))
-  b1 = UOp.param(1, dtypes.float, (10,10))
-  r0 = UOp.range(10, axis_id=0)
-  r1 = UOp.range(10, axis_id=1)
-
-  fxn = (b0[r0, r1] + b1[r0, r1]).call(B0, B1)
-  t = Tensor(fxn)
-  t.realize()
-
-# gemm (N,N)
-
-# (N//k, k, N//k, k)
-
-
-# what if call just implicitly ends all ranges and you don't need to connect them?
-# you do have to connect them, and it does end the ranges
-
-# if assign (store+after) is on call, we move the store into the call (indexed with the ranges) and replace the assign with an after
-
-
-def gemm(A, B):
-  N = 4096
-  k = 128
-
-  ia = UOp.param(0, dtypes.float, (k, k)).reshape(k, 1, k)
-  ib = UOp.param(1, dtypes.float, (k, k)).reshape(1, k, k)
-  gemm_fxn = (ia * ib).sum(2) # <-- rangeify this
-
-  a = UOp.param(0, dtypes.float, (N, N))
-  b = UOp.param(1, dtypes.float, (N, N))
-  r0 = UOp.range(N//k, 0)
-  r1 = UOp.range(N//k, 1)
-  local_fxn = gemm_fxn.call(a.reshape(N//k, k, N//k, k)[r0, :, r1, :], b.reshape(N//k, k, N//k, k)[r0, :, r1, :], r0, r1).permute(0,2,1,3).reshape(N,N)
-
-  fxn = local_fxn.call(A,B)
-
-
-
-  return
-
-
-
-
-
-  a = UOp.param(0, dtypes.float, (N//k, k, N//k, k))
-  b = UOp.param(1, dtypes.float, (N//k, k, N//k, k))
-
-
-
-  # inner kxk GEMM (are WMMAs calls?)
-  ia = UOp.param(0, dtypes.float, (k,k)).reshape(k, 1, k)
-  ib = UOp.param(1, dtypes.float, (k,k)).reshape(1, k, k)
-  r0 = UOp.range(N//k, 0)
-  r1 = UOp.range(N//k, 1)
-  fxn = (ia * ib).sum(2).call(a[:, r0, :, r1], b[:, r0, :, r1])   # this call ends these ranges implicitly
-  assert fxn.shape == (N//k, N//k, k, k)
-
-
-  #.call(A, B, UOp.range(N//k), UOp.range(N//k))
-
-  #r0 = UOp.param(2, dtypes.index, (), vmin_vmax=(0, N//k-1))
-  #r1 = UOp.param(3, dtypes.index, (), vmin_vmax=(0, N//k-1))
-
-
-
-
-
-
-
-
-
-
-# Q = [batch, seq_len, heads, dim]
-# K = [batch, seq_len, head_kv, dim]
-# V = [batch, seq_len, head_kv, dim]
-
-
-

extra/gemm/amd_uop_matmul.py

@@ -6,9 +6,8 @@ from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import getenv
 
 N = getenv("N", 4096)
-M = getenv("M", N)
-K = getenv("K", N)
-NUM_RUNS = getenv("CNT", 5)
+M = K = N
+run_count = getenv("CNT", 5)
 
 # ---------------------------
 # launch/config constants
@@ -20,9 +19,6 @@ WARP_SIZE = 32
 BLOCK_N = 128   # columns of C (N-dim) per block
 BLOCK_M = 128   # rows of C (M-dim) per block
 BLOCK_K = 8     # K-slice per block iteration
-assert N % BLOCK_N == 0, f"N ({N}) must be a multiple of BLOCK_N ({BLOCK_N})"
-assert M % BLOCK_M == 0, f"M ({M}) must be a multiple of BLOCK_M ({BLOCK_M})"
-assert K % BLOCK_K == 0, f"K ({K}) must be a multiple of BLOCK_K ({BLOCK_K})"
 
 # Register tile sizes (per-thread accumulator tile of C)
 TN = 4     # columns per thread
@@ -40,16 +36,16 @@ WAVE_TILE_N = 128 if is_kernel5 else 64
 WAVE_TILE_M = BLOCK_N * BLOCK_M // WARPS_PER_BLOCK // WAVE_TILE_N
 assert BLOCK_N % WAVE_TILE_N == 0, "BN must be a multiple of WN"
 assert BLOCK_M % WAVE_TILE_M == 0, "BM must be a multiple of WM"
-WAVES_PER_BLOCK_N = BLOCK_N // WAVE_TILE_N
-WAVES_PER_BLOCK_M = BLOCK_M // WAVE_TILE_M
-assert WAVES_PER_BLOCK_N * WAVES_PER_BLOCK_M == WARPS_PER_BLOCK, "wave grid must match warps/block"
+WAVES_IN_BLOCK_X = BLOCK_N // WAVE_TILE_N
+WAVES_IN_BLOCK_Y = BLOCK_M // WAVE_TILE_M
+assert WAVES_IN_BLOCK_X * WAVES_IN_BLOCK_Y == WARPS_PER_BLOCK, "wave grid must match warps/block"
 
-LANES_PER_WAVE_N = 8
-LANES_PER_WAVE_M = 4
-REG_TILES_PER_WAVE_N = WAVE_TILE_N // (LANES_PER_WAVE_N * TN)
-REG_TILES_PER_WAVE_M = WAVE_TILE_M // (LANES_PER_WAVE_M * TM)
-assert WAVE_TILE_N % (LANES_PER_WAVE_N * TN) == 0, "WAVE_TILE_N must be divisible by LANES_PER_WAVE_N*TN"
-assert WAVE_TILE_M % (LANES_PER_WAVE_M * TM) == 0, "WAVE_TILE_M must be divisible by LANES_PER_WAVE_M*TM"
+LANES_PER_WAVE_X = 8
+LANES_PER_WAVE_Y = 4
+ITERS_PER_WAVE_N = WAVE_TILE_N // (LANES_PER_WAVE_X * TN)
+ITERS_PER_WAVE_M = WAVE_TILE_M // (LANES_PER_WAVE_Y * TM)
+assert WAVE_TILE_N % (LANES_PER_WAVE_X * TN) == 0, "WAVE_TILE_N must be divisible by LANES_PER_WAVE_X*TN"
+assert WAVE_TILE_M % (LANES_PER_WAVE_Y * TM) == 0, "WAVE_TILE_M must be divisible by LANES_PER_WAVE_Y*TM"
 
 def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.LOOP): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
 def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=False):
@@ -62,41 +58,41 @@ def hand_spec_kernel3():
   # ---------------------------
   # block indices & placeholders
   # ---------------------------
-  block_id_n = UOp.special(N // BLOCK_N, "gidx0")
-  block_id_m = UOp.special(M // BLOCK_M, "gidx1")
+  blockIdx_x = UOp.special(N // BLOCK_N, "gidx0")
+  blockIdx_y = UOp.special(N // BLOCK_M, "gidx1")
 
-  a = UOp.placeholder((M, K), dtypes.float, slot=1)
-  b = UOp.placeholder((K, N), dtypes.float, slot=2)
-  c = UOp.placeholder((M, N), dtypes.float, slot=0)
+  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)
 
   # index the output with the globals
-  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[block_id_m, :, block_id_n, :]
+  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[blockIdx_y, :, blockIdx_x, :]
 
   # open the main reduction range
-  k_tile_range = UOp.range(K // BLOCK_K, 0, AxisType.REDUCE)
-  a = a.reshape(M // BLOCK_M, BLOCK_M, K // BLOCK_K, BLOCK_K)[block_id_m, :, k_tile_range, :]
-  b = b.reshape(K // BLOCK_K, BLOCK_K, N // BLOCK_N, BLOCK_N)[k_tile_range, :, block_id_n, :]
+  k_tile_range = UOp.range(N // BLOCK_K, 0, AxisType.REDUCE)
+  a = a.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_K, BLOCK_K)[blockIdx_y, :, k_tile_range, :]
+  b = b.reshape(N // BLOCK_K, BLOCK_K, N // BLOCK_N, BLOCK_N)[k_tile_range, :, blockIdx_x, :]
 
   # globals are no longer used, they are already in the indexes
-  del block_id_m, block_id_n
+  del blockIdx_y, blockIdx_x
 
   # ---------------------------
-  # GLOBAL -> LOCAL (A_local, B_local)
+  # GLOBAL -> LOCAL (As, Bs)
   # ---------------------------
   tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
 
   # A: read BM x BK tiles (permute on store into locals)
-  BM_A_local_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
-  A_local = UOp.placeholder((BLOCK_K, BM_A_local_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
-  A_local_store = copy(A_local.permute((1,0)).reshape(-1, THREADS_PER_BLOCK)[:, tid], a.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=100)
+  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))
+  As_store = copy(As.permute((1,0)).reshape(-1, THREADS_PER_BLOCK)[:, tid], a.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=100)
 
   # B: read BK x BN tiles
-  B_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
-  B_local_store = copy(B_local.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
+  Bs = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
+  Bs_store = copy(Bs.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
 
   # TODO: can we automate barrier?
-  barrier = UOp.barrier(A_local_store, B_local_store)
-  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
+  barrier = UOp.barrier(As_store, Bs_store)
+  As, Bs = As.after(barrier), Bs.after(barrier)
 
   # open inner k range
   k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
@@ -104,33 +100,31 @@ def hand_spec_kernel3():
   # ---------------------------
   # LOCAL -> REG (per-wave tiles)
   # ---------------------------
-  waveIdx = (tid // WARP_SIZE) % WAVES_PER_BLOCK_N
-  waveIdy = (tid // WARP_SIZE) // WAVES_PER_BLOCK_N
-  assert waveIdy.vmax+1 == WAVES_PER_BLOCK_M
+  waveIdx = (tid // WARP_SIZE) % WAVES_IN_BLOCK_X
+  waveIdy = (tid // WARP_SIZE) // WAVES_IN_BLOCK_X
+  assert waveIdy.vmax+1 == WAVES_IN_BLOCK_Y
 
-  laneIdx = (tid % WARP_SIZE) % LANES_PER_WAVE_N
-  laneIdy = (tid % WARP_SIZE) // LANES_PER_WAVE_N
-  assert laneIdy.vmax+1 == LANES_PER_WAVE_M
+  laneIdx = (tid % WARP_SIZE) % LANES_PER_WAVE_X
+  laneIdy = (tid % WARP_SIZE) // LANES_PER_WAVE_X
+  assert laneIdy.vmax+1 == LANES_PER_WAVE_Y
 
-  A_col = UOp.placeholder((REG_TILES_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
-  A_local_slice = A_local[k, :].reshape(WAVES_PER_BLOCK_M, REG_TILES_PER_WAVE_M, LANES_PER_WAVE_M, TM)[waveIdy, :, laneIdy, :]
-  A_col = copy(A_col, A_local_slice , 300, set=True, upcast=True)
+  A_col = UOp.placeholder((ITERS_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  A_col = copy(A_col, As[k, :].reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM)[waveIdy, :, laneIdy, :], 300, set=True, upcast=True)
 
-  B_row = UOp.placeholder((REG_TILES_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
-  B_local_slice = B_local[k, :].reshape(WAVES_PER_BLOCK_N, REG_TILES_PER_WAVE_N, LANES_PER_WAVE_N, TN)[waveIdx, :, laneIdx, :]
-  B_row = copy(B_row, B_local_slice, 400, set=True, upcast=True)
+  B_row = UOp.placeholder((ITERS_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+  B_row = copy(B_row, Bs[k, :].reshape(WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)[waveIdx, :, laneIdx, :], 400, set=True, upcast=True)
 
   # ---------------------------
   # FMA: c_regs += A_col * B_row
   # ---------------------------
-  c_regs = UOp.placeholder((REG_TILES_PER_WAVE_M, TM, REG_TILES_PER_WAVE_N, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  c_regs = UOp.placeholder((ITERS_PER_WAVE_M, TM, ITERS_PER_WAVE_N, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
   i = UOp.range(c_regs.size, 16)
   c_regs = c_regs.after(c_regs.flatten()[i].store(0.0).end(i))
 
   # TODO: why don't these work as upcast?
   # why if the ranges merge is it slow?!? (if you change the order on end, they will merge. big slowdown on METAL)
-  iter_m, t_m, iter_n, t_n = rngs = rngs_for_shape(c_regs.shape, 500)
-  sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iter_m, t_m] * B_row[iter_n, t_n]).end(iter_m, iter_n, t_m, t_n)
+  iterWaveM, yt, iterWaveN, xt = rngs = rngs_for_shape(c_regs.shape, 500)
+  sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iterWaveM, yt] * B_row[iterWaveN, xt]).end(iterWaveM, iterWaveN, yt, xt)
 
   # Close k, sync, and close K tiles
   sink = sink.end(k).barrier().end(k_tile_range)
@@ -138,28 +132,28 @@ def hand_spec_kernel3():
   # ---------------------------
   # REG -> GLOBAL (epilogue)
   # ---------------------------
-  c = c.reshape(WAVES_PER_BLOCK_M, REG_TILES_PER_WAVE_M, LANES_PER_WAVE_M, TM,
-                WAVES_PER_BLOCK_N, REG_TILES_PER_WAVE_N, LANES_PER_WAVE_N, TN)
+  c = c.reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM,
+                WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)
   c = c[waveIdy, :, laneIdy, :,
         waveIdx, :, laneIdx, :]
   sink = copy(c, c_regs.after(sink), rng=600)
 
   return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()
 
-def test_matmul(sink:UOp, dtype=dtypes.float32, M=M, N=N, K=K):
+def test_matmul(sink:UOp, dtype=dtypes.float32, N=N):
   rng = np.random.default_rng()
-  a = Tensor(rng.random((M, K), dtype=np.float32)-0.5, dtype=dtype)
-  b = Tensor(rng.random((K, N), dtype=np.float32)-0.5, dtype=dtype)
-  hc = Tensor.empty(M, N, dtype=dtype)
+  a = Tensor(rng.random((N, N), dtype=np.float32)-0.5, dtype=dtype)
+  b = Tensor(rng.random((N, N), dtype=np.float32)-0.5, dtype=dtype)
+  hc = Tensor.empty(N, N, dtype=dtype)
   Tensor.realize(a, b, hc)
 
   ei = ExecItem(sink, [t.uop.buffer for t in [hc, a, b]], prg=get_runner(Device.DEFAULT, sink))
 
   ets = []
   with Context(DEBUG=2):
-    for _ in range(NUM_RUNS):
+    for _ in range(run_count):
       ets.append(ei.run(wait=True))
-  print(f"REAL TFLOPS {M * N * K * 2 / min(ets) * 1e-12:.2f}")
+  print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")
 
   if getenv("VERIFY", 1):
     GlobalCounters.reset()

extra/usbgpu/tbgpu/installer/notary_tool.sh

@@ -3,7 +3,3 @@ set -e
 
 ditto -c -k --keepParent ./build/Release/TinyGPU.app ./build/Release/TinyGPU.zip
 xcrun notarytool submit ./build/Release/TinyGPU.zip --keychain-profile "hgwJFhdheiIEy82nDN" --wait
-
-rm ./build/Release/TinyGPU.zip
-xcrun stapler staple ./build/Release/TinyGPU.app
-ditto -c -k --keepParent ./build/Release/TinyGPU.app ./build/Release/TinyGPU.zip

pyproject.toml

@@ -74,7 +74,7 @@ testing_minimal = [
   "hypothesis>=6.148.9",
   "z3-solver<4.15.4",  # 4.15.4 has a segfault when creating many z3.Context()
 ]
-testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai", "gguf>=0.18"]
+testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai", "gguf"]
 testing = [
   "tinygrad[testing_unit]",
   "pillow",

test/amd/test_sqtt_examples.py

@@ -9,8 +9,8 @@ from tinygrad.renderer.amd import decode_inst
 from tinygrad.runtime.autogen.amd.rdna3.ins import SOPP
 from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp
 from tinygrad.renderer.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_RDNA4, WAVEEND, INST, INST_RDNA4, VALUINST,
-                                     IMMEDIATE, IMMEDIATE_MASK, PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4, PACKET_TYPES_CDNA, CDNA_WAVESTART,
-                                     InstOp, InstOpRDNA4, print_packets, CDNA_WAVEEND)
+                                     IMMEDIATE, IMMEDIATE_MASK, PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4,
+                                     InstOp, InstOpRDNA4, print_packets)
 from test.amd.helpers import TARGET_TO_ARCH
 
 import tinygrad
@@ -125,7 +125,7 @@ class SQTTExamplesTestBase(unittest.TestCase):
           self.assertIsInstance(packets[0], LAYOUT_HEADER, f"first packet should be LAYOUT_HEADER in {name}")
 
   def test_packet_types_valid(self):
-    all_classes = set(PACKET_TYPES_RDNA3.values()) | set(PACKET_TYPES_RDNA4.values()) | set(PACKET_TYPES_CDNA.values())
+    all_classes = set(PACKET_TYPES_RDNA3.values()) | set(PACKET_TYPES_RDNA4.values())
     for name, (events, *_) in self.examples.items():
       for i, event in enumerate(events):
         with self.subTest(example=name, event=i):
@@ -138,8 +138,8 @@ class SQTTExamplesTestBase(unittest.TestCase):
       if "empty" in name: continue
       with self.subTest(example=name):
         all_packets = [p for e in events for p in decode(e.blob)]
-        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVESTART, WAVESTART_RDNA4, CDNA_WAVESTART))]), 0, f"no WAVESTART in {name}")
-        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVEEND, CDNA_WAVEEND))]), 0, f"no WAVEEND in {name}")
+        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVESTART, WAVESTART_RDNA4))]), 0, f"no WAVESTART in {name}")
+        self.assertGreater(len([p for p in all_packets if isinstance(p, WAVEEND)]), 0, f"no WAVEEND in {name}")
 
   def test_time_monotonic(self):
     for name, (events, *_) in self.examples.items():
@@ -183,8 +183,8 @@ class SQTTExamplesTestBase(unittest.TestCase):
         for event in events:
           wave_starts: dict[tuple[int, int, int], int] = {}
           for p in decode(event.blob):
-            if isinstance(p, (WAVESTART, CDNA_WAVESTART, WAVESTART_RDNA4)): wave_starts[(p.wave, p.simd, p.cu)] = p._time
-            elif isinstance(p, (WAVEEND, CDNA_WAVEEND)) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
+            if isinstance(p, (WAVESTART, WAVESTART_RDNA4)): wave_starts[(p.wave, p.simd, p.cu)] = p._time
+            elif isinstance(p, WAVEEND) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
               our_waves.append((wave_starts[key], p._time))
         self.assertEqual(sorted(our_waves), sorted(roc_waves), f"wave times mismatch in {name}")
 
@@ -221,12 +221,8 @@ class TestSQTTExamplesRDNA3(SQTTExamplesTestBase):
   }
 
 class TestSQTTExamplesRDNA4(SQTTExamplesTestBase): target = "gfx1200"
-
-class TestSQTTExamplesCDNA(SQTTExamplesTestBase):
-  target = "gfx950"
-  def test_decode_all_examples(self): self.skipTest("TODO: correct deltas in the timestamp packet types, first packet is REGCS_CDNA")
-  def test_gemm_has_instructions(self): self.skipTest("TODO: decode CDNA inst packets")
-  def test_rocprof_wave_times_match(self): self.skipTest("TODO: requires timestamp patching")
+@unittest.skip("TODO: fix CDNA")
+class TestSQTTExamplesCDNA(SQTTExamplesTestBase): target = "gfx950"
 
 if __name__ == "__main__":
   unittest.main()

test/amd/test_sqttmap.py

@@ -2,10 +2,9 @@
 import unittest, pickle
 from typing import Iterator
 from pathlib import Path
-from tinygrad.helpers import DEBUG, OSX, getenv, temp
+from tinygrad.helpers import DEBUG, OSX
 from tinygrad.renderer.amd.sqtt import print_packets, map_insts
 from tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm
-from tinygrad.viz.serve import sqtt_timeline
 from test.amd.disasm import disasm
 
 import tinygrad
@@ -54,7 +53,7 @@ class TestSQTTMapBase(unittest.TestCase):
   def setUpClass(cls):
     if cls is TestSQTTMapBase: raise unittest.SkipTest("base class")
     cls.examples = {}
-    for pkl_path in ([Path(temp("profile.pkl", append_user=True))] if getenv("LOAD_PROFILE") else sorted((EXAMPLES_DIR/cls.target).glob("*.pkl"))):
+    for pkl_path in sorted((EXAMPLES_DIR/cls.target).glob("*.pkl")):
       with open(pkl_path, "rb") as f:
         data = pickle.load(f)
       sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
@@ -73,22 +72,6 @@ class TestSQTTMapBase(unittest.TestCase):
           passed_insts, n_waves, n_units = rocprof_inst_traces_match(event, kern_events[event.kern], target, pass_rocprof_err)
           if n_waves: print(f"{name}: passed for {passed_insts} instructions across {n_waves} waves scheduled on {n_units} wave units")
 
-  def test_sqtt_timeline(self):
-    for name, (events, kern_events, target) in self.examples.items():
-      for event in events:
-        if (p:=kern_events.get(event.kern)) is None: continue
-        with self.subTest(example=name, kern=event.kern):
-          events = [e for e in sqtt_timeline(event.blob, p.lib, target) if type(e).__name__ == "ProfileRangeEvent"]
-          insts, execs = 0, 0
-          for e in events:
-            if "EXEC" in e.device:
-              if "ALT" not in e.name.display_name: execs += 1
-            elif "WAVE" in e.device:
-              # sopk/immediates don't get ALU/MEM EXEC
-              if e.name.display_name not in {"IMMEDIATE", "IMMEDIATE_MASK", "JUMP", "JUMP_NO", "MESSAGE"}: insts += 1
-            else: raise Exception(f"timeline row must be INST or EXEC, got {e.device}")
-          self.assertEqual(execs, insts)
-
 class TestSQTTMapRDNA3(TestSQTTMapBase): target = "gfx1100"
 
 class TestSQTTMapRDNA4(TestSQTTMapBase): target = "gfx1200"

test/backend/test_multitensor.py

@@ -199,6 +199,13 @@ class TestMultiTensor(unittest.TestCase):
       run_schedule(sched)
       np.testing.assert_equal(xt.numpy(), X_np[i*2:i*2+2])
 
+  def test_cat_on_non_shard_axis(self):
+    # cat must be lowered to PAD/ADD before multi_pm runs, otherwise MULTI nodes are not handled
+    X = Tensor.arange(8).reshape(4, 2).realize().shard_(devices_2, 0)
+    Y = Tensor.arange(8, 16).reshape(4, 2).realize().shard_(devices_2, 0)
+    Z = X.cat(Y, dim=1)
+    np.testing.assert_equal(Z.numpy(), np.concatenate([np.arange(8).reshape(4, 2), np.arange(8, 16).reshape(4, 2)], axis=1))
+
   @given(strat.sampled_from((devices_2, devices_3)),
          strat.sampled_from((Ops.ADD, Ops.MUL, Ops.MAX)),
          strat.sampled_from((None, 0, 1)), strat.sampled_from((None, 0, 1)))

test/backend/test_ops.py

@@ -6,7 +6,6 @@ from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, CPU_LLVM, AMD_LL
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
-from tinygrad.renderer.cstyle import QCOMCLRenderer
 from tinygrad.renderer.nir import NIRRenderer
 
 TINY_BACKEND = getenv("TINY_BACKEND")
@@ -437,7 +436,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,35), (45,35), (45,35)], lambda x,y,z: x.lerp(y,z))
     helper_test_op(None, lambda x,y,z: x.lerp(y,z), vals=[[1.,2.,3.], [4.,5.,6.], 0.5])
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
+  @unittest.skipIf(Device.DEFAULT == "QCOM", "OpenCL fails to compile this (both on GPU(qcom)/QCOM backends)")
   def test_tril(self):
     helper_test_op([(3,3)], lambda x: x.tril())
     helper_test_op([(3,3)], lambda x: x.tril(1))
@@ -455,7 +454,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(5,3,3)], lambda x: x.tril(1))
     helper_test_op(None, lambda x: x.tril(), vals=[[[True] * 3] * 3], forward_only=True)
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
+  @unittest.skipIf(Device.DEFAULT == "QCOM", "OpenCL fails to compile this (both on GPU(qcom)/QCOM backends)")
   def test_triu(self):
     helper_test_op([(3,3)], lambda x: x.triu())
     helper_test_op([(3,3)], lambda x: x.triu(1))
@@ -766,7 +765,6 @@ class TestOps(unittest.TestCase):
 
     self.helper_test_exception([(4), (4)], lambda x,y: x.bitwise_xor(y), expected=RuntimeError)
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_and(self):
     data = [[1,-8,1],[32,1,6]]
     tor = torch.tensor(data, dtype=torch.int)
@@ -784,7 +782,6 @@ class TestOps(unittest.TestCase):
 
     self.helper_test_exception([(4), (4)], lambda x,y: x.bitwise_and(y), expected=RuntimeError)
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_or(self):
     data = [[1,-8,1],[32,1,6]]
     tor = torch.tensor(data, dtype=torch.int)
@@ -1173,7 +1170,6 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x: x.type(torch.int32).argmax().type(torch.int32), lambda x: x.argmax(), forward_only=True, vals=[[False, True]])
     helper_test_op(None, lambda x: x.type(torch.int32).argmax().type(torch.int32), lambda x: x.argmax(), forward_only=True, vals=[[True, False]])
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_argmin(self):
     # check if it returns the first index for multiple occurrences
     helper_test_op(None, lambda x: x.argmin().type(torch.int32), lambda x: x.argmin(), forward_only=True, vals=[[2, 2]])
@@ -1479,7 +1475,6 @@ class TestOps(unittest.TestCase):
   def test_prod_dtype_arg(self):
     with self.assertRaises(AttributeError): Tensor([1.0, 2.0]).prod(dtype="")
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_min(self):
     helper_test_op([(3,3)], lambda x: x.min())
     helper_test_op([(45,3)], lambda x: x.min())
@@ -1508,6 +1503,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(3,3)], lambda x: torch.full_like(x, 2).prod(), lambda x: (x.full_like(2)).prod(), forward_only=True)
     helper_test_op([(3,3)], lambda x: torch.full_like(x, 2).max(), lambda x: (x.full_like(2)).max(), forward_only=True)
 
+  @unittest.skipIf(Device.DEFAULT == "QCOM", "OpenCL fails to compile this (both on GPU(qcom)/QCOM backends)")
   def test_any(self):
     helper_test_op([(3,4,5,6)], lambda x: x.any(), forward_only=True)
     helper_test_op(None, lambda x: x.any(), vals=[[True, True]], forward_only=True)
@@ -1519,7 +1515,7 @@ class TestOps(unittest.TestCase):
   def test_any_zero_axis(self):
     helper_test_op([(1,0,3,0,5)], lambda x: x.any(axis=(1,3)), forward_only=True)
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
+  @unittest.skipIf(Device.DEFAULT == "QCOM", "OpenCL fails to compile this (both on GPU(qcom)/QCOM backends)")
   def test_all(self):
     helper_test_op([(3,4,5,6)], lambda x: x.all(), forward_only=True)
     helper_test_op(None, lambda x: x.all(), vals=[[True, True]], forward_only=True)
@@ -2893,7 +2889,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[...,c,:,e], lambda x: x[...,k,:,p])
 
   @slow_test
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_slice_fancy_indexing_dim_collapse_int(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     # dim collapse from int
@@ -2904,7 +2899,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[1,:,3:11:2,d,0:2], lambda x: x[1,:,3:11:2,o,0:2])
 
   @slow_test
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_slice_fancy_indexing_dim_inject_none(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     # dim injection from None
@@ -2939,7 +2933,6 @@ class TestOps(unittest.TestCase):
                             lambda x: x[Tensor([[0,1,-1],[-1,-2,0]]), Tensor([2,1,-1])])
 
   @slow_test
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_slice_fancy_indexing_list_indices(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[((0,),)])
@@ -2951,7 +2944,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,(2,1,0),c,(-2,1,0),e], lambda x: x[i,(2,1,0),k,(-2,1,0),p])
 
   @slow_test
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_slice_fancy_indexing_tuple_indices(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[(((0,),),)], lambda x: x[(((0,),),)])
@@ -3293,6 +3285,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(20,)], lambda x: (x>0.5).nonzero().int(), lambda x: (x>0.5).nonzero(), forward_only=True)
     helper_test_op([(10, 5, 3)], lambda x: (x>0.5).nonzero().int(), lambda x: (x>0.5).nonzero(), forward_only=True)
 
+  @unittest.skipIf(Device.DEFAULT == "QCOM", "OpenCL fails to compile this (both on GPU(qcom)/QCOM backends)")
   def test_cast(self):
     helper_test_op([(3, 3)], lambda x: x.float())
     helper_test_op(None, lambda x: x.float(), vals=[[0, 1, 2, 3]], forward_only=True)

test/backend/test_schedule.py

@@ -811,26 +811,6 @@ class TestSchedule(unittest.TestCase):
       self.assertEqual(cnt1, 5)
       self.assertEqual(cnt2, 5)
 
-  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
-  def test_image_f16_residual_fusion(self):
-    with Context(FLOAT16=1, OPENPILOT_HACKS=1):
-      def cnt():
-        inp = Tensor.empty((512,), dtype='float')
-        b1, b2 = Tensor.empty((512, 1024), dtype='float'), Tensor.empty((1024, 512), dtype='float')
-        c1, c2 = Tensor.empty((1024,), dtype='float'), Tensor.empty((512,), dtype='float')
-        rb = (((((inp @ b1) + c1).relu() @ b2) + c2).relu() + inp).relu()
-        b16, c16 = Tensor.empty((512, 16), dtype='float'), Tensor.empty((16,), dtype='float')
-        b32, c32 = Tensor.empty((512, 32), dtype='float'), Tensor.empty((32,), dtype='float')
-        sched = Tensor.schedule((rb @ b16 + c16).relu(), (rb @ b32 + c32).relu())
-        for si in sched: si.lower()
-        return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
-
-      with Context(IMAGE=1): cnt1 = cnt()
-      with Context(IMAGE=2): cnt2 = cnt()
-
-      self.assertEqual(cnt1, 9)
-      self.assertEqual(cnt2, 9)
-
   @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
   @unittest.expectedFailure
   def test_image_conv_fusion(self):

test/mockgpu/amd/amdgpu.py

@@ -271,7 +271,6 @@ class SDMAExecutor(AMDQueue):
       elif op == amd_gpu.SDMA_OP_GCR: self._execute_gcr()
       elif op == amd_gpu.SDMA_OP_COPY: self._execute_copy()
       elif op == amd_gpu.SDMA_OP_TIMESTAMP: self._execute_timestamp()
-      elif op == 32: self.rptr[0] += 4  # SDMA_OP_DUMMY_TRAP: pipeline flush, no interrupt
       else: raise RuntimeError(f"Unknown SDMA op {op}")
     return self.rptr[0] - prev_rptr
 

test/null/test_memory_planner.py

@@ -11,8 +11,8 @@ def b(i, base=None, offset=0, pin=False, size=16):
   if pin: global_map[i].ref(1)
   return global_map[i]
 
-def check_assign(buffers:list[list[Buffer]|tuple[Buffer, ...]], copies:list[tuple[Buffer, Buffer]]|None=None):
-  assigned = _internal_memory_planner(buffers, copies=copies)
+def check_assign(buffers:list[list[Buffer]|tuple[Buffer, ...]]):
+  assigned = _internal_memory_planner(buffers, noopt_buffers=None)
 
   taken_parts = set()
   first_appearance, last_appearance = {}, {}
@@ -134,75 +134,5 @@ class TestMemoryPlanner(unittest.TestCase):
     ]
     check_assign(bs)
 
-  def test_copy_bufs_separate_from_compute(self):
-    bs = [
-      [b(0), b(1)],
-      [b(1), b(2)],
-      [b(3), b(2)],
-    ]
-    assigned = _internal_memory_planner(bs, copies=[(b(1), b(0))])
-    r1, r2 = assigned.get(b(1), b(1)), assigned.get(b(2), b(2))
-    assert r1.base != r2.base
-
-  def test_copy_bufs_reuse_among_copies(self):
-    bs = [
-      [b(0), b(1)],
-      [b(2), b(1)],
-      [b(3), b(2)],
-    ]
-    assigned = _internal_memory_planner(bs, copies=[(b(1), b(0)), (b(2), b(1))])
-    r1, r2 = assigned.get(b(1), b(1)), assigned.get(b(2), b(2))
-    assert r1.base == r2.base
-
-  def test_compute_bufs_reuse_among_compute(self):
-    bs = [
-      [b(0), b(1)],
-      [b(2), b(1)],
-      [b(3), b(2)],
-      [b(4), b(3)],
-    ]
-    assigned = _internal_memory_planner(bs, copies=[(b(1), b(0))])
-    r2, r3 = assigned.get(b(2), b(2)), assigned.get(b(3), b(3))
-    assert r2.base == r3.base
-
-  def test_copy_and_compute_no_cross_reuse(self):
-    bs = [
-      [b(0), b(1)],
-      [b(2), b(1)],
-      [b(3), b(2)],
-    ]
-    assigned = _internal_memory_planner(bs, copies=[(b(2), b(1))])
-    r0, r2 = assigned.get(b(0), b(0)), assigned.get(b(2), b(2))
-    assert r0.base != r2.base
-
-  def test_multiple_copy_bufs_with_offsets(self):
-    bs = [
-      [b(0, pin=True), b(1), b(2)],
-      [b(3, base=0, offset=1, size=8), b(1), b(2)],
-      [b(4), b(3)],
-      [b(5), b(4)],
-    ]
-    check_assign(bs, copies=[(b(1), b(0)), (b(2), b(0))])
-
-  def test_copy_bufs_pinned_mixed(self):
-    bs = [
-      [b(0, pin=True), b(1), b(2)],
-      [b(1), b(3), b(2)],
-      [b(4), b(3)],
-      [b(5), b(4), b(0)],
-    ]
-    check_assign(bs, copies=[(b(1), b(0)), (b(3), b(1))])
-
-  def test_deferred_copy_frees_chain(self):
-    bs = []
-    copies = []
-    for i in range(6):
-      copy_buf, compute_buf = b(i * 2 + 1), b(i * 2 + 2)
-      bs.append([copy_buf, b(0, pin=True)])
-      bs.append([compute_buf, copy_buf])
-      copies.append((copy_buf, b(0, pin=True)))
-    bs.append([b(100, pin=True)])
-    check_assign(bs, copies=copies)
-
 if __name__ == "__main__":
   unittest.main()

test/null/test_multitensor.py

@@ -108,6 +108,7 @@ class TestMultiRamUsage(unittest.TestCase):
   def test_matmul_half(self): self._test_matmul_half(dev_count=2)
   def test_matmul_half_alt(self): self._test_matmul_half(dev_count=4)
 
+  @unittest.expectedFailure
   def test_multi_layer_allreduce(self):
     N = 32
     devices_2 = ("NULL:1", "NULL:2")

test/null/test_schedule.py

@@ -1189,11 +1189,5 @@ class TestBufferView(unittest.TestCase):
     b = a.shrink(((200, 800),)).shrink(((0, 300),)).reshape((30, 10)).shrink(((20, 25), (0, 10))).contiguous()
     run_schedule(check_schedule(b, 0))
 
-class TestInvalidTensor(unittest.TestCase):
-  def test_full_invalid_is_zero_kernels(self):
-    from tinygrad.dtype import Invalid
-    t = Tensor.full((4,), Invalid, dtype=dtypes.float)
-    check_schedule(t, 0)
-
 if __name__ == '__main__':
   unittest.main(verbosity=2)

test/null/test_simplify_valid_idx.py

@@ -351,7 +351,7 @@ class TestImageSimplification(unittest.TestCase):
 
     self.check(load,
                "((((idx2*2)+r0)<11)&((((idx1*8)+r1)<3)!=True))",
-               "(idx0+(idx1*512+r1*64)+-192)",
+               "(((idx0+((idx1*512)+(r1*64)))+832)%1024)",
                "((((idx2*2)+r0)+(((idx1+((r1+5)//8))+1)//2))+-4)")
 
   def test_simplify1(self):
@@ -388,17 +388,18 @@ class TestImageSimplification(unittest.TestCase):
     alu8 = (idx0//8%32//4)
     alu9 = idx0<256
 
+    # TODO: can this be simplified further?
     load = get_load_image_uop(shape, alu9, (((alu8+(alu2*8))%64),(alu2//8)))
-    self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32+8)", "(idx0//2%4)")
+    self.check(load, "(idx0<256)", "(((((idx0%8)*32)+(idx0//32))+8)%64)", "(idx0//2%4)")
 
     load = get_load_image_uop(shape, alu9, (((alu8+(alu3*8))%64),(alu3//8)))
-    self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32+16)", "(idx0//2%4)")
+    self.check(load, "(idx0<256)", "(((((idx0%8)*32)+(idx0//32))+16)%64)", "(idx0//2%4)")
 
     load = get_load_image_uop(shape, alu9, (((alu8+(alu4*8))%64),(alu4//8)))
-    self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32+24)", "(idx0//2%4)")
+    self.check(load, "(idx0<256)", "(((((idx0%8)*32)+(idx0//32))+24)%64)", "(idx0//2%4)")
 
     load = get_load_image_uop(shape, alu9, (((alu8+(alu5*8))%64),(alu5//8)))
-    self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32)", "(idx0//2%4)")
+    self.check(load, "(idx0<256)", "((((idx0%8)*32)+(idx0//32))%64)", "(idx0//2%4)")
 
   def test_simplify5(self):
     # openpilot 0.9.7, chunk replacement to simplify
@@ -413,7 +414,7 @@ class TestImageSimplification(unittest.TestCase):
     valid = alu3<640
 
     load = get_load_image_uop(shape, valid, idx)
-    self.check(load, None, "((idx0+((idx1//3)*16))+128)", "((idx1%3)*4)")
+    self.check(load, "(((idx0+(idx1*64))%192)<160)", "((idx0+((idx1//3)*16))+128)", "((idx1%3)*4)")
 
   def test_simplify6(self):
     # from openpilot

test/null/test_uop_symbolic.py

@@ -286,11 +286,6 @@ class TestSymbolic(unittest.TestCase):
                               "(((z+(x*-1))+(y*-1))+7)")
     self.helper_test_variable((10+12*Variable("x",0,2)+Variable("y", 0, 4)%3)%13, 8, 12, "(((x*-1)+(y%3))+10)")
 
-  def test_mod_congruence_tied_remainder(self):
-    # when f%c == c/2, both r and r-c have equal abs — try both signs
-    self.helper_test_variable((3+2*Variable("x",0,1)+3*Variable("y",0,1))%4, 0, 3, "((x*-2)+(y*-1)+3)")
-    self.helper_test_variable((3+6*Variable("x",0,1)+7*Variable("y",0,1))%4, 0, 3, "((x*-2)+(y*-1)+3)")
-
   def test_div_congruence(self):
     self.helper_test_variable((3+3*Variable("a",0,3))//4, 0, 3, "a")
     self.helper_test_variable((18+17*Variable("a",0,2)+17)//18, 1, 3, "(a+1)")
@@ -708,24 +703,6 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((31*b+1)//18, 0, 172, "(((b*13)+1)//18+b)")
     self.helper_test_variable((19*b+3)//7, 0, 271, "(((b*5)+3)//7+(b*2))")
 
-  def test_gcd_with_remainder(self):
-    # gcd_with_remainder: factor GCD out of non-constant terms and denominator
-    a = Variable("a", 0, 2)
-    self.helper_test_variable((a*4)//6, 0, 1, "(a*2//3)")
-    self.helper_test_variable((a*4+1)//6, 0, 1, "(a*2//3)")
-    self.helper_test_variable((a*4+2)//6, 0, 1, "((a*2+1)//3)")
-    self.helper_test_variable((a*4+3)//6, 0, 1, "((a*2+1)//3)")
-    self.helper_test_variable((a*4)%6, 0, 4, "(a*2%3*2)")
-    self.helper_test_variable((a*4+1)%6, 1, 5, "(a*2%3*2+1)")
-    self.helper_test_variable((a*4+2)%6, 0, 4, "((a*2+1)%3*2)")
-    self.helper_test_variable((a*4+3)%6, 1, 5, "((a*2+1)%3*2+1)")
-
-  def test_div_by_factor_tie_break(self):
-    a = Variable("a", 0, 1)
-    b = Variable("b", 0, 1)
-    with Context(CORRECT_DIVMOD_FOLDING=1):
-      self.helper_test_variable((a*2+b*3+2)//6, 0, 1, "((a+b+1)//3)")
-
   def test_div_mod_recombine_large_coeff(self):
     # recombine must work even when coeff > divisor: both mod and div reduce the coeff the same way
     b = Variable("b", 0, 100)
@@ -733,79 +710,6 @@ class TestSymbolic(unittest.TestCase):
     a = Variable("a", 0, 10)
     self.helper_test_variable((25*a+3)%10 + ((25*a+3)//10)*10, 3, 253, "((a*25)+3)")
 
-  def test_mod_nest_by_factor(self):
-    # (a*f+b) % (f*k) = (a%k)*f + b when 0<=b<f — mirrors nest_div_by_factor for MOD
-    gidx0 = Variable("gidx0", 0, 15)
-    lidx0 = Variable("lidx0", 0, 3)
-    # f=4, k=2, c=8: (gidx0*4+lidx0)%8 = (gidx0%2)*4 + lidx0
-    self.helper_test_variable((gidx0*4+lidx0)%8, 0, 7, "(lidx0+gidx0%2*4)")
-    # f=2, k=4: (gidx0*2+lidx0)%8 where lidx0 in [0,1]
-    lidx1 = Variable("lidx1", 0, 1)
-    self.helper_test_variable((gidx0*2+lidx1)%8, 0, 7, "(lidx1+gidx0%4*2)")
-    # f=3, k=3: (a*3+b)%9 where b in [0,2]
-    a = Variable("a", 0, 10)
-    b = Variable("b", 0, 2)
-    self.helper_test_variable((a*3+b)%9, 0, 8, "(b+a%3*3)")
-
-  def test_mod_nest_by_factor_with_const(self):
-    # nest_by_factor MOD with non-zero constant offset: (a*f+b+const) % (f*k) = (a%k)*f + b + const when 0<=b+const<f
-    a = Variable("a", 0, 7)
-    b = Variable("b", 0, 1)
-    # f=4, k=2, const=2: (a*4+b+2)%8 = (a%2)*4 + b + 2
-    self.helper_test_variable((a*4+b+2)%8, 2, 7, "(b+a%2*4+2)")
-    # f=6, k=2, const=3: (a*6+b+3)%12 = (a%2)*6 + b + 3
-    b2 = Variable("b", 0, 2)
-    self.helper_test_variable((a*6+b2+3)%12, 3, 11, "(b+a%2*6+3)")
-    # f=3, k=2, const=1: (a*3+b+1)%6 = (a%2)*3 + b + 1
-    self.helper_test_variable((a*3+b+1)%6, 1, 5, "(b+a%2*3+1)")
-
-  def test_div_nest_by_factor_with_const(self):
-    # nest_by_factor IDIV: (160*a + 5*b + 4*c + K) // 60 should pick div=5 (clean) over div=4 (dirty)
-    a = Variable("a", 0, 2)
-    b = Variable("b", 0, 31)
-    c = Variable("c", 0, 1)
-    self.helper_test_variable((160*a + 5*b + 4*c) // 60, 0, 7, "(a*2+(b+a*8)//12)")
-    self.helper_test_variable((160*a + 5*b + 4*c + 1) // 60, 0, 8, "(a*2+(b+c+a*8)//12)")
-    self.helper_test_variable((160*a + 5*b + 4*c + 2) // 60, 0, 8, "(a*2+(b+c+a*8)//12)")
-    self.helper_test_variable((160*a + 5*b + 4*c + 3) // 60, 0, 8, "(a*2+(b+c+a*8)//12)")
-    self.helper_test_variable((160*a + 5*b + 4*c + 59) // 60, 0, 8, "(a*2+(b+c+a*8+11)//12)")
-
-  def test_div_mod_recombine_after_nesting(self):
-    # when nest_div_by_factor simplifies the div, the mod must also nest so recombine can fire
-    gidx0 = Variable("gidx0", 0, 15)
-    lidx0 = Variable("lidx0", 0, 3)
-    x = gidx0*4+lidx0
-    # div nests: x//8 -> gidx0//2, mod nests: x%8 -> (gidx0%2)*4+lidx0, then recombine gives x back
-    self.helper_test_variable((x//8)*8 + x%8, 0, 63, "(lidx0+gidx0*4)")
-    # with a scaling factor: recombine gives x*2
-    self.helper_test_variable((x//8)*16 + (x%8)*2, 0, 126, "(gidx0*8+lidx0*2)")
-    # two variables with different factors
-    a = Variable("a", 0, 7)
-    b = Variable("b", 0, 1)
-    y = a*6+b
-    # div nests: y//12 -> a//2, mod nests: y%12 -> (a%2)*6+b, recombine
-    self.helper_test_variable((y//12)*12 + y%12, 0, 43, "(b+a*6)")
-
-  def test_div_mod_recombine_in_additive_sum(self):
-    x = Variable("x", 0, 31)
-    y = Variable("y", 0, 5)
-    # recombine should work inside larger additive sums, not just in the two special y+... tree shapes
-    self.helper_test_variable((x//8)*4 + y + (x//2)%4, 0, 20, "(y+x//2)")
-    self.helper_test_variable(y + (x//8)*4 + (x//2)%4, 0, 20, "(y+x//2)")
-
-  def test_div_mod_recompose_low_order_remainder(self):
-    x = Variable("x", 0, 127)
-    self.helper_test_variable((x//2)%4*2 + x%2, 0, 7, "(x%8)")
-
-  def test_reshape_index_roundtrip(self):
-    # simulate reshape index decompose then recompose — the core pattern this enables
-    # (8,8) decomposed for (16,4): combined=r0*8+r1, div and mod by 4
-    r0 = Variable("r0", 0, 7)
-    r1 = Variable("r1", 0, 7)
-    combined = r0*8+r1
-    src_idx = (combined//4)*4 + combined%4
-    self.helper_test_variable(src_idx, 0, 63, "(r1+r0*8)")
-
   def test_gated_load(self):
     idx = Variable("idx", 0, 24)
     self.helper_test_variable(idx//4, 0, 6, "(idx//4)")
@@ -952,12 +856,6 @@ class TestSymbolic(unittest.TestCase):
     self.assertIn((a.cast(dtypes.long)+b.cast(dtypes.long)).render(), "(long)((a+b))")
     self.assertIn((a.cast(dtypes.long)*b.cast(dtypes.long)).render(), "(long)((a*b))")
 
-  def test_nested_mod_negative_range(self):
-    # (x%(k*c))%c = x%c holds for cmod regardless of signs since sign(x%(k*c)) = sign(x)
-    x = Variable("x", 0, 1575)
-    self.helper_test_variable(((x + (-1064)) % 512) % 4, -3, 3, "((x+-1064)%4)")
-    self.helper_test_variable(((x + (-1064)) % 512) % 128, -127, 127, "((x+-1064)%128)")
-
 class TestSymbolicNumeric(unittest.TestCase):
   def helper_test_numeric(self, f):
     MIN, MAX = 0, 10

test/null/test_viz.py

@@ -411,7 +411,7 @@ class TestVizProfiler(BaseTestViz):
     j = load_profile(prof)
     event = j['layout']['NV:SDMA:0']['events'][0]
     gbs = sz/(dur*1e-6)*1e-9
-    self.assertEqual(event['fmt'], f"{gbs:.0f} GB/s")
+    self.assertTrue(event['fmt'].startswith(f"{gbs:.0f} GB/s"))
 
   def test_graph(self):
     prof = [ProfileDeviceEvent(device='NV', tdiff=decimal.Decimal(-1000)),
@@ -453,7 +453,9 @@ class TestVizProfiler(BaseTestViz):
     j = load_profile(prof)
     sdma_events = j['layout']['NV:1:SDMA:0']['events']
     gbs = sz/(dur*1e-6)*1e-9
-    self.assertEqual(sdma_events[0]['fmt'], f"{gbs:.0f} GB/s")
+    timing_txt, trace_txt = sdma_events[0]['fmt'].split("\nTB:")
+    self.assertEqual(timing_txt, f"{gbs:.0f} GB/s")
+    self.assertEqual(json.loads(trace_txt)[0][0], __file__)
 
   def test_block_ordering(self):
     prof = [ProfileDeviceEvent(device='NV', tdiff=decimal.Decimal(-1000)),
@@ -509,15 +511,9 @@ class TestVizProfiler(BaseTestViz):
 
   def test_calltrace(self):
     def fxn(): return Tensor.empty(10).mul(2).realize()
-    with cpu_profile(TracingKey("test_fxn"), "CUSTOM"):
-      fxn()
-    codegen_trace = get_viz_list()[0]["steps"][0]["trace"]
-    assert any(fxn.__code__.co_filename == f and fxn.__code__.co_firstlineno == l for f,l,*_ in codegen_trace), str(codegen_trace)
-    profile_ret = load_profile(cpu_events)
-    e = profile_ret["layout"]["CUSTOM"]["events"][0]
-    self.assertEqual(e["name"], "test_fxn")
-    runtime_trace = json.loads(e["fmt"].replace("TB:", ""))
-    assert any(fxn.__code__.co_filename == f and fxn.__code__.co_firstlineno+1 == l for f,l,*_ in runtime_trace), str(runtime_trace)
+    fxn()
+    trace = get_viz_list()[0]["steps"][0]["trace"]
+    assert any(fxn.__code__.co_filename == f and fxn.__code__.co_firstlineno == l for f,l,*_ in trace), str(trace)
 
   # can pack up to 1hr 11 min of trace events
   def test_trace_duration(self):

test/unit/test_assign.py

@@ -911,36 +911,5 @@ class TestAssignToUnrealizedView(unittest.TestCase):
       # TODO: broken now, silently dropped
       self.assertEqual(c.tolist(), [[5,5],[5,5]])
 
-class TestPartialAssignToSharedBuffer(unittest.TestCase):
-  def test_five_slices(self):
-    big = Tensor.zeros(50).contiguous().realize()
-    views = [big[i*10:(i+1)*10].reshape(2, 5) for i in range(5)]
-    for v in views: v.assign(v + 1)
-    Tensor.realize(*views)
-    for v in views:
-      np.testing.assert_allclose(v.numpy(), np.ones((2, 5)))
-
-  def test_many_slices(self):
-    n_params = 10
-    big = Tensor.zeros(n_params * 12).contiguous().realize()
-    grads = [big[i*12:(i+1)*12].reshape(3, 4) for i in range(n_params)]
-    for g in grads: g.assign(g + 1)
-    Tensor.realize(*grads)
-    for g in grads:
-      np.testing.assert_allclose(g.numpy(), np.ones((3, 4)))
-
-  def test_mixed_shapes(self):
-    big = Tensor.zeros(100).contiguous().realize()
-    shapes = [(3, 4), (4, 6), (6, 4), (2, 5), (4, 3)]
-    pos, views = 0, []
-    for s in shapes:
-      n = s[0] * s[1]
-      views.append(big[pos:pos+n].reshape(*s))
-      pos += n
-    for v in views: v.assign(v + 1)
-    Tensor.realize(*views)
-    for v, s in zip(views, shapes):
-      np.testing.assert_allclose(v.numpy(), np.ones(s))
-
 if __name__ == "__main__":
   unittest.main()

test/unit/test_gguf.py

@@ -26,7 +26,7 @@ class TestGGUF(unittest.TestCase):
     # codes 0-7 = [0, 1, 2, 3, 4, 6, 8, 12], codes 8-15 are their negatives
     block = np.array([0x80] + list(range(16)), dtype=np.uint8)  # E=128, nibbles 0-15 in low, zeros in high
     expected = np.array([0., 1., 2., 3., 4., 6., 8., 12., -0., -1., -2., -3., -4., -6., -8., -12.] + [0.]*16, dtype=np.float32)
-    np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 32, GGMLQuantizationType.MXFP4.value).numpy().flatten(), expected)
+    np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 32, 39).numpy().flatten(), expected)
 
   def test_dequantization_q4_0(self): self._test_dequantization(GGMLQuantizationType.Q4_0)
   def test_dequantization_q4_1(self): self._test_dequantization(GGMLQuantizationType.Q4_1)
@@ -34,8 +34,9 @@ class TestGGUF(unittest.TestCase):
   def test_dequantization_q4_k(self): self._test_dequantization(GGMLQuantizationType.Q4_K)
   def test_dequantization_q5_k(self): self._test_dequantization(GGMLQuantizationType.Q5_K)
   def test_dequantization_q6_k(self): self._test_dequantization(GGMLQuantizationType.Q6_K)
-  def test_dequantization_mxfp4(self): self._test_dequantization(GGMLQuantizationType.MXFP4)
-  def test_dequantization_mxfp4_old(self):
+  def test_dequantization_mxfp4(self):
+    MXFP4 = 39
+
     def encode(nibbles, E):
       packed = [(low & 0xF) | ((high & 0xF) << 4) for low, high in zip(nibbles[:16], nibbles[16:])]
       return np.array([E] + packed, dtype=np.uint8)
@@ -56,10 +57,12 @@ class TestGGUF(unittest.TestCase):
       blocks.append(encode(codes, E))
       expected.extend(decode(c, E) for c in codes)
     tensor = Tensor(np.concatenate(blocks))
-    out = ggml_data_to_tensor(tensor, len(expected), GGMLQuantizationType.MXFP4.value)
-    np.testing.assert_equal(out.numpy(), expected)
+    out = ggml_data_to_tensor(tensor, len(expected), MXFP4)
+    # TODO: should this be exact equal? somehow failed on CI
+    np.testing.assert_allclose(out.numpy(), expected, atol=0.0, rtol=1e-6)
 
   def test_dequantization_mxfp4_block(self):
+    MXFP4 = 39
     # https://gist.github.com/Ananta-Ranganathan/3317b6ed51a3b033e9c2564fafb4e043
     # used the above script to download the first block of blk.0.attn_k_b.weight from
     # https://huggingface.co/unsloth/GLM-4.7-Flash-GGUF/blob/main/GLM-4.7-Flash-MXFP4_MOE.gguf
@@ -74,8 +77,9 @@ class TestGGUF(unittest.TestCase):
                         0.03125000,  0.00000000, 0.06250000,  0.01562500,
                         -0.06250000,  0.00000000, 0.00000000, -0.01562500,
                         0.04687500,  0.00000000, 0.00000000,  0.01562500], dtype=np.float32)
-    out = ggml_data_to_tensor(Tensor(block), 32, GGMLQuantizationType.MXFP4.value)
-    np.testing.assert_equal(out.numpy(), expected)
+    out = ggml_data_to_tensor(Tensor(block), 32, MXFP4)
+    # TODO: similar to previous test fails on Mac CI with assert_equal for unclear reason
+    np.testing.assert_allclose(out.numpy(), expected, atol=0.0, rtol=1e-6)
 
   def test_expected_failure_unknown_type(self):
     with self.assertRaises(ValueError):
@@ -129,7 +133,6 @@ class TestGGUFGEMV(unittest.TestCase):
     if qtype == GGMLQuantizationType.Q8_0: q_data[:, :2] = scales[:, :2]                    # d at offset 0
     elif qtype in (GGMLQuantizationType.Q4_K, GGMLQuantizationType.Q5_K): q_data[:, :4] = scales[:, :4] # d, dmin at offset 0
     elif qtype == GGMLQuantizationType.Q6_K: q_data[:, -2:] = scales[:, :2]                 # d at end
-    elif qtype == GGMLQuantizationType.MXFP4: q_data[:, 0] = rng.integers(120, 136, size=n_blocks, dtype=np.uint8) # constrain byte0
     q_data = q_data.flatten()
     ref = dequantize(q_data, qtype).reshape(rows, cols)
 
@@ -155,7 +158,6 @@ class TestGGUFGEMV(unittest.TestCase):
   def test_gguf_gemv_q4_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q4_K)
   def test_gguf_gemv_q5_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_K)
   def test_gguf_gemv_q6_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q6_K)
-  def test_gguf_gemv_mxfp4(self): self._test_gguf_gemv(GGMLQuantizationType.MXFP4)
 
 if __name__ == '__main__':
   unittest.main()

test/unit/test_invalid_tensor.py

@@ -1,141 +0,0 @@
-import math, unittest
-from unittest.mock import patch
-from tinygrad import Tensor
-from tinygrad.device import Device
-from tinygrad.dtype import Invalid, dtypes
-from tinygrad.helpers import unwrap_class_type
-
-class TestInvalidTensor(unittest.TestCase):
-  def _invalid_test_helper(self, out, expected):
-    before = None
-    original_call = (runtime_cls:=unwrap_class_type(Device[Device.DEFAULT].runtime)).__call__
-
-    def patched_call(self_prg, *bufs, **kwargs):
-      nonlocal before
-      before = Device[Device.DEFAULT].allocator._as_buffer(bufs[0]).cast(out.dtype.fmt).tolist()
-      return original_call(self_prg, *bufs, **kwargs)
-
-    with patch.object(runtime_cls, '__call__', patched_call): ret = out.tolist()
-
-    for i,v in enumerate(expected):
-      if v is None: assert before[i] == ret[i] or (math.isnan(before[i]) and math.isnan(ret[i]))
-      else: assert ret[i] == v
-
-    return before, ret
-
-  def test_where_x_invalid(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid)
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_where_invalid_x(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Invalid, Tensor([1.0, 2.0, 3.0, 4.0]))
-    self._invalid_test_helper(out, [None, None, 3.0, 4.0])
-
-  def test_where_invalid_2d(self):
-    mask = Tensor.arange(6).reshape(2, 3) < 3
-    vals = Tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
-    out = mask.where(vals, Invalid)
-    before, ret = self._invalid_test_helper(out, [])
-    assert ret[0] == [1.0, 2.0, 3.0]
-    assert before[3] == ret[1][0] or (math.isnan(before[3]) and math.isnan(ret[1][0]))
-    assert before[4] == ret[1][1] or (math.isnan(before[4]) and math.isnan(ret[1][1]))
-    assert before[5] == ret[1][2] or (math.isnan(before[5]) and math.isnan(ret[1][2]))
-
-  def test_where_invalid_int(self):
-    mask = Tensor.arange(3) < 2
-    out = mask.where(Tensor([10, 20, 30]), Invalid)
-    self._invalid_test_helper(out, [10, 20, None])
-
-  def test_where_invalid_add(self):
-    mask = Tensor.arange(3) < 2
-    mixed = mask.where(Tensor([10.0, 20.0, 30.0]), Invalid)
-    out = mixed + Tensor([1.0, 2.0, 3.0])
-    self._invalid_test_helper(out, [11.0, 22.0, None])
-
-  def test_where_invalid_add_left(self):
-    mask = Tensor.arange(3) < 2
-    mixed = mask.where(Tensor([10.0, 20.0, 30.0]), Invalid)
-    out = Tensor([1.0, 2.0, 3.0]) + mixed
-    self._invalid_test_helper(out, [11.0, 22.0, None])
-
-  def test_where_always_true(self):
-    mask = Tensor.arange(3) < 10
-    out = mask.where(Tensor([10.0, 20.0, 30.0]), Invalid)
-    self._invalid_test_helper(out, [10.0, 20.0, 30.0])
-
-  def test_where_cast(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid).cast(dtypes.int)
-    self._invalid_test_helper(out, [1, 2, None, None])
-
-  def test_where_compare(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid) > 1
-    self._invalid_test_helper(out, [False, True, None, None])
-
-  def test_where_unary(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 4.0, 9.0, 16.0]), Invalid).sqrt()
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_where_where(self):
-    mask1 = Tensor.arange(4) < 2
-    mask2 = Tensor.arange(4) > 0
-    out = mask2.where(mask1.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid), Invalid)
-    self._invalid_test_helper(out, [None, 2.0, None, None])
-
-  def test_where_reduce_always_true(self):
-    mask = Tensor.arange(4) < 9
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid).sum()
-    before, ret = self._invalid_test_helper(out, [])
-    assert ret == 10.0
-
-  def test_invalid_unary(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.float).sqrt())
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_invalid_binary(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.float) + 2)
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_invalid_binary_left(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), 2 + Tensor.full((4,), Invalid, dtype=dtypes.float))
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_invalid_reshape(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid).reshape(2,2)
-    before, ret = self._invalid_test_helper(out, [])
-    assert ret[0] == [1.0, 2.0]
-    assert ret[1][0] == before[2] or (math.isnan(ret[1][0]) and math.isnan(before[2]))
-    assert ret[1][1] == before[3] or (math.isnan(ret[1][1]) and math.isnan(before[3]))
-
-  def test_invalid_cast(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int).cast(dtypes.float))
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_invalid_bitcast(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int).bitcast(dtypes.float))
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_where_bitcast(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int)).bitcast(dtypes.int)
-    self._invalid_test_helper(out, [0x3f800000, 0x40000000, None, None])
-
-  # tensor indexing uses reduce, so the entire result becomes invalid
-  @unittest.expectedFailure
-  def test_tensor_index(self):
-    idx = (Tensor.arange(4) < 2).where(Tensor([0, 1, 2, 3]), Invalid)
-    out = Tensor([1.0, 2.0, 3.0, 4.0])[idx]
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-if __name__ == '__main__':
-  unittest.main()

test/unit/test_llm_server.py

@@ -83,30 +83,5 @@ class TestTransformerGenerate(unittest.TestCase):
     self.assertEqual(cache_size_after_warmup, len(schedule_cache),
       f"third prompt added {len(schedule_cache) - cache_size_after_warmup} new schedule cache entries (expected 0)")
 
-  def test_chunked_prefill(self):
-    """When prompt > chunk_size, all chunks should be prefill"""
-    from tinygrad.apps.llm import Transformer
-    from tinygrad.uop.ops import resolve
-    model = Transformer(num_blocks=1, dim=64, hidden_dim=128, n_heads=2, n_kv_heads=2,
-                        norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, max_context=64)
-
-    def get_prefill_flags(tokens, chunk_size):
-      is_prefill = []
-      def mock_call(self, tokens, start_pos):
-        is_prefill.append(resolve(tokens.shape[1] != 1))
-        return Tensor([[42]])
-      with patch.object(Transformer, '__call__', mock_call):
-        gen = model.generate(tokens, chunk_size=chunk_size)
-        for _ in range(3): next(gen)
-      model._cached_tokens = []
-      return is_prefill
-
-    # 8 tokens, chunk_size=4 -> 2 prefill chunks
-    self.assertEqual(get_prefill_flags(list(range(8)), 4), [True, True, False, False])
-    # 9 tokens, chunk_size=4 -> 3 prefill chunks (4+4+1)
-    self.assertEqual(get_prefill_flags(list(range(9)), 4), [True, True, True, False, False])
-    # 4 tokens, chunk_size=4 -> 1 prefill chunk
-    self.assertEqual(get_prefill_flags(list(range(4)), 4), [True, False, False])
-
 if __name__ == '__main__':
   unittest.main()

tinygrad/apps/llm.py

@@ -239,10 +239,10 @@ class Transformer:
     t = Tensor(tokens + [0] * (self.max_context - len(tokens)), dtype="int32").reshape(1, self.max_context)
     # recompute start_pos from what's currently valid in the kv cache
     start_pos = self.get_start_pos(tokens)
-    out, prompt_len = None, len(tokens)
+    out = None
     while len(tokens) < self.max_context:
       sp, nt = v_start_pos.bind(start_pos), v_toks.bind(min(chunk_size, len(tokens) - start_pos))
-      out = self(t[:, sp:sp+nt] if start_pos < prompt_len or out is None else out, sp).realize()
+      out = self(t[:, sp:sp+nt] if out is None else out, sp).realize()
       start_pos += nt.val
       # chunked prefill: keep processing until all prompt tokens are consumed
       if start_pos < len(tokens): continue

tinygrad/codegen/late/devectorizer.py

@@ -192,7 +192,7 @@ def _do_image_fixup(dt:ImageDType, idx:UOp) -> tuple[UOp, UOp, int, int]:
                 # maximize number of valids removed
                (len(_drop_valid_stmts(valid, idx:=uop_given_valid(valid, UOp.vectorize((x//4)%hw[1], x//(4*hw[1]))), *hw)),
                 # and minimize idx complexity (number of nodes)
-                -len(idx.backward_slice)))
+                -len(idx.simplify().backward_slice)))
     buf = buf.replace(dtype=(dtypes.imageh if dt.itemsize == 2 else dtypes.imagef)((h, w, 4), w * 4 * dt.itemsize))
   oidx = UOp(Ops.VECTORIZE, dtypes.index.vec(2), ((x // 4) % w, (x // (4*w))))
   return x, idx.replace(src=(buf, oidx.valid(valid))), w, h

tinygrad/codegen/simplify.py

@@ -16,7 +16,7 @@ pm_flatten_range = PatternMatcher([
   (UPat((Ops.REDUCE, Ops.STORE, Ops.END), name="r"), flatten_range),
 ])
 
-def count_divmod(x:UOp) -> int: return sum(u.op in {Ops.IDIV, Ops.MOD} for u in x.backward_slice)
+def count_divmod(x:UOp) -> int: return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
 def simplify_merge_adjacent(u:UOp) -> UOp|None:
   reduce_ranges = [x.ranges for x in u.backward_slice_with_self if x.op is Ops.REDUCE]
   # on END we only want to merge adjacent ranges, on REDUCE we want to try all combinations

tinygrad/device.py

@@ -6,7 +6,7 @@ import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re
 from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup, ContextVar
 from tinygrad.helpers import unwrap_class_type, suppress_finalizing, select_first_inited, VIZ, CPU_LLVM, CPU_LVP, NV_PTX, CUDA_PTX, NV_NAK
-from tinygrad.helpers import EMULATED_DTYPES, NULL_IR3, NULL_QCOMCL, TracingKey
+from tinygrad.helpers import EMULATED_DTYPES, TracingKey
 from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
 if TYPE_CHECKING: from tinygrad.renderer import Renderer
 
@@ -371,7 +371,7 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
     if device in ["CUDA", "NV"]: return not CI
     if device == "CPU" and CPU_LLVM: return OSX
     if device == "PYTHON": return sys.version_info >= (3, 12)
-  if dtype == dtypes.float64: return (device not in {"METAL", "QCOM"} and not (OSX and device == "CL") and not NULL_IR3 and not NULL_QCOMCL
+  if dtype == dtypes.float64: return (device not in {"METAL", "QCOM"} and not (OSX and device == "CL") and not getenv("NULL_IR3")
                                       and dtypes.long not in EMULATED_DTYPES.tolist(dtypes))
   return True
 

tinygrad/dtype.py

@@ -30,7 +30,6 @@ class InvalidType:
   def __hash__(self): return id(self)
   def __repr__(self): return "Invalid"
   def __reduce__(self): return (InvalidType, ())  # unpickle returns the singleton
-  def __format__(self, spec): return "Invalid"
 
 Invalid = InvalidType()
 

tinygrad/engine/jit.py

@@ -114,9 +114,9 @@ class GraphRunner(Runner):
           assert ji.prg.p.local_size is not None
           self.launch_dims_base[j] = (tuple(ji.prg.p.global_size), tuple(ji.prg.p.local_size))
 
-    # used in MultiGraphRunner. tracks (offset, end, dep) ranges per base buffer id to handle suballocated buffers correctly.
-    self.w_dependency_map: dict[int, list[tuple[int, int, Any]]] = collections.defaultdict(list)
-    self.r_dependency_map: dict[int, list[tuple[int, int, Any]]] = collections.defaultdict(list)
+    # used in MultiGraphRunner. the ints are id() of _bufs
+    self.w_dependency_map: dict[int, Any] = {}
+    self.r_dependency_map: dict[int, list[Any]] = collections.defaultdict(list)
 
     assert jit_cache[0].prg is not None
     super().__init__(colored(f"<batched {len(jit_cache)}>", "cyan"), jit_cache[0].prg.device.split(":")[0], estimates.simplify())
@@ -132,22 +132,19 @@ class GraphRunner(Runner):
       yield j, (dims[gl] if gl is not None else self.launch_dims_base[j][0]), (dims[lc] if lc is not None else self.launch_dims_base[j][1])
 
   def _access_resources(self, bufs:list[Buffer], write:list[int], new_dependency:Any):
+    # To synchronize access to resources, we monitor the necessary prerequisites for accessing each resource,
+    # whether for write or read operations. A resource can be accessed by either a single writer or multiple readers.
     wait_nodes = []
+
     for i,buf in enumerate(bufs):
-      key, s, e = id(buf.base._buf), buf.offset, buf.offset + buf.nbytes
-      wait_nodes += [dep for st,en,dep in self.w_dependency_map[key] if st < e and s < en]
-      if i in write: wait_nodes += [dep for st,en,dep in self.r_dependency_map[key] if st < e and s < en]
-    for i,buf in enumerate(bufs):
-      key, s, e = id(buf.base._buf), buf.offset, buf.offset + buf.nbytes
+      if id(buf.base._buf) in self.w_dependency_map: wait_nodes.append(self.w_dependency_map[id(buf.base._buf)])
       if i in write:
-        for dmap in [self.w_dependency_map, self.r_dependency_map]:
-          kept = []
-          for st,en,dep in dmap[key]:
-            if st < min(s, en): kept.append((st, min(s, en), dep))
-            if max(e, st) < en: kept.append((max(e, st), en, dep))
-          dmap[key] = kept
-        self.w_dependency_map[key].append((s, e, new_dependency))
-      else: self.r_dependency_map[key].append((s, e, new_dependency))
+        if id(buf.base._buf) in self.r_dependency_map: wait_nodes.extend(self.r_dependency_map.pop(id(buf.base._buf)))
+
+    for i,buf in enumerate(bufs):
+      if i in write: self.w_dependency_map[id(buf.base._buf)] = new_dependency
+      else: self.r_dependency_map[id(buf.base._buf)].append(new_dependency)
+
     return list({id(x):x for x in wait_nodes}.values())
 
   @staticmethod
@@ -360,8 +357,9 @@ class TinyJit(Generic[ReturnType]):
         jit_cache = pruned
 
       # memory planning (optional)
-      copies = [(cast(Buffer,ji.bufs[0]),cast(Buffer,ji.bufs[1])) for ji in jit_cache if isinstance(ji.prg, (BufferXfer, BufferCopy, EncDec))]
-      assigned = _internal_memory_planner([cast(list[Buffer], item.bufs) for item in jit_cache], copies, debug_prefix="JIT ")
+      # Exclude buffers involved in transfer ops to preserve parallelism.
+      noopt_buffers = {b for ji in jit_cache if isinstance(ji.prg, (BufferXfer, BufferCopy, EncDec)) for b in ji.bufs}
+      assigned = _internal_memory_planner([cast(list[Buffer], item.bufs) for item in jit_cache], noopt_buffers, debug_prefix="JIT ")
       jit_cache = [replace(item, bufs=[assigned.get(b,b).ensure_allocated() for b in item.bufs if b is not None]) for item in jit_cache]
 
       input_replace = get_input_replace(jit_cache, input_buffers)

tinygrad/engine/memory.py

@@ -7,50 +7,43 @@ from tinygrad.uop.ops import Ops
 from tinygrad.dtype import dtypes, ImageDType
 from tinygrad.runtime.support.memory import TLSFAllocator
 
-LaneKey = tuple[str, int]
-
 # **************** memory planning ****************
 
-def _internal_memory_planner(buffers:list[list[Buffer]], copies:list[tuple[Buffer, Buffer]]|None=None,
-                             ignore_checks=False, debug_prefix="") -> dict[Buffer, Buffer]:
+def _internal_memory_planner(buffers:list[list[Buffer]], noopt_buffers=None, ignore_checks=False, debug_prefix="") -> dict[Buffer, Buffer]:
   if NO_MEMORY_PLANNER: return {}
   first_appearance, last_appearance, buf_to_opt = {}, {}, set()
   for i,u in enumerate(buffers):
     for buf in u:
-      if not ignore_checks and (buf.is_allocated() or buf.base.is_allocated() or buf.uop_refcount > 0): continue
+      should_skip = buf.is_allocated() or buf.base.is_allocated() or buf.uop_refcount > 0 or (noopt_buffers is not None and buf.base in noopt_buffers)
+      if not ignore_checks and should_skip: continue
       if buf.base not in first_appearance: first_appearance[buf.base] = i
       last_appearance[buf.base] = i
       buf_to_opt.add(buf)
 
-  # Separate copy and compute buffers into different lanes and defer cross-queue frees to avoid introducing dependencies (copy->compute->copy)
-  copy_dsts, copy_srcs = ({dst.base for dst,_ in copies}, {src.base for _,src in copies}) if copies else (set(), set())
-  def _key(buf) -> LaneKey: return (buf.device, 1 if buf in copy_dsts or buf in copy_srcs else 0)
-  buf_hold = {buf: last_appearance[buf] - first_appearance[buf] + 1 for buf in first_appearance if buf in copy_dsts or buf in copy_srcs}
-
   # Sort buffer operations in timeline order. Two events: buffer is allocated or buffer is freed.
   buffer_requests = sorted([((first_appearance[buf], True), buf) for buf in first_appearance.keys()] + \
-    [((last_appearance[buf] + 1 + buf_hold.get(buf, 0), False), buf) for buf in first_appearance.keys()], key=lambda x: x[0])
-  total_memory = sum(round_up(buf.nbytes, BLK:=0x1000) for buf in first_appearance.keys()) * 2 # *2 for fragmentation (which is about 15%)
+                           [((last_appearance[buf] + 1, False), buf) for buf in first_appearance.keys()], key=lambda x: x[0])
+  total_memory = sum(round_up(buf.nbytes, min_block_size:=0x1000) for buf in first_appearance.keys()) * 2 # *2 for fragmentation (which is about 15%)
 
   # Try to suballocate from a shared buffer managed by global_planner using TLSFAllocator.
   # Also track buffer replacements for buffers that do not support suballocation.
   buffer_replace:dict[Buffer, tuple[Buffer|None, int|None]] = {}
   reuse_buffers:dict[tuple, list[Buffer]] = defaultdict(list)
-  global_planner:dict[LaneKey, tuple[int, TLSFAllocator]] = defaultdict(lambda: (0, TLSFAllocator(total_memory, block_size=BLK, lv2_cnt=32)))
+  global_planner:dict[str, tuple[int, TLSFAllocator]] = defaultdict(lambda: (0, TLSFAllocator(total_memory, block_size=min_block_size, lv2_cnt=32)))
   for (_, is_open_ev), buf in buffer_requests:
     # Check if suballocation is possible for the given buffer and device.
     if hasattr(Device[buf.device].allocator, "_offset") and not isinstance(buf.dtype, ImageDType):
-      if is_open_ev: buffer_replace[buf] = (None, global_planner[_key(buf)][1].alloc(round_up(buf.nbytes, BLK)))
-      else: global_planner[_key(buf)][1].free(cast(int, buffer_replace[buf][1]))
-      global_planner[_key(buf)] = (max(global_planner[_key(buf)][0], buffer_replace[buf][1] + buf.nbytes), global_planner[_key(buf)][1])
+      if is_open_ev: buffer_replace[buf] = (None, global_planner[buf.device][1].alloc(round_up(buf.nbytes, 0x1000)))
+      else: global_planner[buf.device][1].free(cast(int, buffer_replace[buf][1]))
+      global_planner[buf.device] = (max(global_planner[buf.device][0], buffer_replace[buf][1] + buf.nbytes), global_planner[buf.device][1])
     else:
-      key = (_key(buf), buf.dtype, buf.options, buf.nbytes)
+      key = (buf.device, buf.dtype, buf.options, buf.nbytes)
       if is_open_ev: buffer_replace[buf] = (reuse_buffers[key].pop(), None) if key in reuse_buffers and len(reuse_buffers[key]) > 0 else (buf, None)
       else: reuse_buffers[key].append(cast(Buffer, buffer_replace[buf][0]))
 
   # Allocate global buffers based on the memory planner.
-  global_buffers = {key: Buffer(key[0], round_up(sz, BLK), dtypes.int8) for key, (sz, _) in global_planner.items()}
-  buffer_resolve:dict[Buffer, tuple[Buffer, int|None]] = {buf: (base or global_buffers[_key(buf)], off) for buf,(base,off) in buffer_replace.items()}
+  global_buffers = {dev: Buffer(dev, round_up(sz, 0x1000), dtypes.int8) for dev, (sz, _) in global_planner.items()}
+  buffer_resolve:dict[Buffer, tuple[Buffer, int|None]] = {buf: (base or global_buffers[buf.device], off) for buf,(base,off) in buffer_replace.items()}
 
   # Assign buffers. First, assign full buffers (not sub-buffers).
   assigned:dict[Buffer, Buffer] = {}
@@ -73,5 +66,5 @@ def _internal_memory_planner(buffers:list[list[Buffer]], copies:list[tuple[Buffe
 def memory_planner(schedule:list[ExecItem]) -> list[ExecItem]:
   # Exclude buffers involved in load ops (e.g transfers) to preserve parallelism in graphs.
   assigned = _internal_memory_planner([[b for b in si.bufs if b is not None] for si in schedule],
-                                      copies=[(cast(Buffer,si.bufs[0]),cast(Buffer,si.bufs[1])) for si in schedule if si.ast.op is Ops.COPY])
+                                      noopt_buffers={b for si in schedule if si.ast.op is not Ops.SINK for b in si.bufs if b is not None})
   return [ExecItem(si.ast, [assigned.get(x, x) if x is not None else None for x in si.bufs], si.metadata, si.fixedvars) for si in schedule]

tinygrad/engine/realize.py

@@ -93,8 +93,8 @@ class BufferXfer(BufferCopy):
   def copy(self, dest, src): dest.allocator._transfer(dest._buf, src._buf, dest.nbytes, src_dev=src.allocator.dev, dest_dev=dest.allocator.dev)
 
 class EncDec(Runner):
-  def __init__(self, cf:UOp, total_sz:int, device:str):
-    self.shape, self.pos_var = tuple(s.arg for s in cf.src if s.op is Ops.CONST), cf.variables()[0].expr
+  def __init__(self, encdec:UOp, total_sz:int, device:str):
+    self.shape, self.pos_var = encdec.arg[0], encdec.variables()[0].expr
     name = f"enc/dec {total_sz/1e6:7.2f}M, HEVC" if total_sz >= 1e6 else f"enc/dec {total_sz:8d}, HEVC"
     super().__init__(colored(name, "yellow"), device, Estimates(lds=total_sz, mem=total_sz))
   def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int], wait=False):
@@ -130,7 +130,7 @@ si_lowerer = PatternMatcher([
   (UPat(Ops.COPY, name="copy"), lambda ctx,copy: (BufferXfer(ctx[0].nbytes, ctx[0].device, ctx[1].device) \
       if hasattr(alc:=Device[ctx[0].device].allocator, '_transfer') and alc.supports_transfer and all_same([x.device.split(":")[0] for x in ctx]) \
       else BufferCopy(ctx[0].nbytes, ctx[0].device, ctx[1].device))),
-  (UPat(Ops.CUSTOM_FUNCTION, arg="encdec", name="cf"), lambda ctx,cf: EncDec(cf, ctx[0].nbytes, ctx[0].device)),
+  (UPat(Ops.ENCDEC, name="encdec"), lambda ctx,encdec: EncDec(encdec, ctx[0].nbytes, ctx[1].device)),
 ])
 
 @dataclass

tinygrad/engine/schedule.py

@@ -71,7 +71,7 @@ def linear_to_schedule(linear:UOp) -> list[ExecItem]:
       base = buf_uops[1].buffer
       assert isinstance(base, Buffer), "base can't be MultiBuffer"
       buffers[buf_uops[0]] = base.view(buf_uops[0].arg, ast.dtype, ast.arg[1]*base.dtype.itemsize)
-    ubufs = [b.buffer for b in buf_uops if b.op is not Ops.BIND]
+    ubufs = [b.buffer for b in buf_uops]
     metadata = si.arg.metadata
     if any(isinstance(x, MultiBuffer) for x in ubufs):
       assert all(isinstance(x, MultiBuffer) for x in ubufs), "kernel must all be multibuffer"

tinygrad/function.py

@@ -36,7 +36,7 @@ class _function(Generic[ReturnType]):
     call_uops: list[UOp] = dedup(input_uops)
 
     # disable realize/schedule while this is running
-    # run it and do surgery later. TODO: why am i not calling it with the params?
+    # run it and do surgery later
     with Context(ALLOW_DEVICE_USAGE=getenv("DEVICE_IN_FUNCTION_BUG", 0)):
       ret = self.fxn(*args, **kwargs)
     assert isinstance(ret, Tensor), "only supports one tensor return for now"

tinygrad/gradient.py

@@ -1,8 +1,14 @@
 from typing import cast
-import math, dataclasses
+import math, itertools, dataclasses
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata
 from tinygrad.helpers import argsort
 
+def cat_gradient(ctx:UOp, ret:UOp) -> tuple[UOp, ...]:
+  axis = ret.arg
+  dim_acc = list(itertools.accumulate([s.shape[axis] for s in ret.src], initial=0))
+  return tuple(ctx.shrink(tuple([(dim_acc[i], dim_acc[i+1]) if j==axis else (0, ctx.shape[j])
+    for j in range(len(ctx.shape))])) for i in range(len(ret.src)))
+
 def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
   def broadcast_to_input(x): return x.reshape(x.shape+(1,)*(len(ret.src[0].shape)-len(x.shape))).expand(ret.src[0].shape)
   if op == Ops.ADD: return (broadcast_to_input(ctx),)
@@ -54,6 +60,7 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[0], s-p[1]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
   (UPat(Ops.PERMUTE, name="ret"), lambda ctx, ret: (ctx.permute(argsort(ret.marg)),)),
   (UPat(Ops.FLIP, name="ret"), lambda ctx, ret: (ctx.flip([i for i,x in enumerate(ret.marg) if x]),)),
+  (UPat(Ops.CAT, name="ret"), lambda ctx, ret: cat_gradient(ctx, ret)),
   (UPat(Ops.COPY, name="ret"), lambda ctx, ret: (ctx.copy_to_device(ret.src[0].device), None)),
   (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

tinygrad/helpers.py

@@ -195,8 +195,7 @@ CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasat
 CPU_CC, CPU_LLVM, CPU_LVP = ContextVar("CPU_CC", ""), ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0)
 NV_CC, NV_PTX, NV_NAK, NV_NVCC = ContextVar("NV_CC", ""), ContextVar("NV_PTX", 0), ContextVar("NV_NAK", 0), ContextVar("NV_NVCC", 0)
 CUDA_CC, CUDA_PTX, CUDA_NVCC = ContextVar("CUDA_CC", ""), ContextVar("CUDA_PTX", 0), ContextVar("CUDA_NVCC", 0)
-NULL_QCOMCL, NULL_IR3, NULL_NAK = ContextVar("NULL_QCOMCL", 0), ContextVar("NULL_IR3", 0), ContextVar("NULL_NAK", 0)
-NULL_ALLOW_COPYOUT = ContextVar("NULL_ALLOW_COPYOUT", 0)
+NULL_IR3, NULL_NAK, NULL_ALLOW_COPYOUT = ContextVar("NULL_IR3", 0), ContextVar("NULL_NAK", 0), ContextVar("NULL_ALLOW_COPYOUT", 0)
 AMD_CC, AMD_LLVM, AMD_HIPCC  = ContextVar("AMD_CC", ""), ContextVar("AMD_LLVM", 0), ContextVar("AMD_HIPCC", 0)
 QCOM_CC, QCOM_IR3 = ContextVar("QCOM_CC", ""), ContextVar("QCOM_IR3", 0)
 # VIZ implies PROFILE, but you can run PROFILE without VIZ

tinygrad/nn/onnx.py

@@ -153,19 +153,21 @@ class OnnxPBParser:
 
   def _parse_ModelProto(self) -> dict:
     """Entry point for parsing the ONNX model."""
-    graph: dict|None = None
-    opset_imports: list[OpSetId] = []
+    obj: dict[str, Any] = {"opset_import": []}
     for fid, wire_type in self._parse_message(self.reader.len):
       match fid:
-        case 7: graph = self._parse_GraphProto()
-        case 8: opset_imports.append(self._parse_OperatorSetIdProto())
+        case 4: obj["domain"] = self.reader.read_string()
+        case 5: obj["model_version"] = self.reader.read_int64()
+        case 7: obj["graph"] = self._parse_GraphProto()
+        case 8: obj["opset_import"].append(self._parse_OperatorSetIdProto())
         case _: self.reader.skip_field(wire_type)
-    assert graph is not None
+
     # update opset version
-    versions = {opset.domain: opset.version for opset in opset_imports}
-    graph["node"] = [OnnxNode(n.op, OpSetId(n.opset_id.domain, versions.get(n.opset_id.domain, 1)), n.inputs, n.outputs, n.opts)
-                     for n in graph["node"]]
-    return graph
+    opset_imports = {Domain.from_onnx(x.get('domain')):x.get('version', 1) for x in obj["opset_import"]}
+    for n in obj["graph"]["node"]:
+      n_ = n["parsed_node"]
+      n["parsed_node"] = OnnxNode(n_.op, OpSetId(n_.opset_id.domain, opset_imports.get(n_.opset_id.domain, 1)), n_.inputs, n_.outputs, n_.opts)
+    return obj
 
   def _parse_GraphProto(self) -> dict:
     obj: dict[str, Any] = {"node": [], "initializer": [], "input": [], "output": []}
@@ -179,23 +181,26 @@ class OnnxPBParser:
         case _: self.reader.skip_field(wire_type)
     return obj
 
-  def _parse_NodeProto(self) -> OnnxNode:
-    inputs: list[str] = []
-    outputs: list[str] = []
-    attributes: list[tuple[str, Any]] = []
-    domain: str|None = None
-    op_type = ""
+  def _parse_NodeProto(self) -> dict:
+    obj: dict[str, Any] = {"input": [], "output": [], "attribute": [], "domain": None}
     for fid, wire_type in self._parse_message(self._decode_end_pos()):
       match fid:
-        case 1: inputs.append(self.reader.read_string())
-        case 2: outputs.append(self.reader.read_string())
-        case 4: op_type = self.reader.read_string()
-        case 5: attributes.append(self._parse_AttributeProto())
-        case 7: domain = self.reader.read_string()
+        case 1: obj["input"].append(self.reader.read_string())
+        case 2: obj["output"].append(self.reader.read_string())
+        case 3: obj["name"] = self.reader.read_string()
+        case 4: obj["op_type"] = self.reader.read_string()
+        case 5: obj["attribute"].append(self._parse_AttributeProto())
+        case 6: obj["doc_string"] = self.reader.read_string()
+        case 7: obj["domain"] = self.reader.read_string()
         case _: self.reader.skip_field(wire_type)
-    return OnnxNode(op_type, OpSetId(Domain.from_onnx(domain), 1), tuple(inputs), tuple(outputs), dict(attributes))
 
-  def _parse_TensorProto(self) -> tuple[str, Tensor]:
+    # parse node
+    attributes = {attr_dict["name"]: attr_dict[AttributeType(attr_dict["type"]).to_field_name()] for attr_dict in obj["attribute"]}
+    opset_id = OpSetId(Domain.from_onnx(obj.get('domain')), 1)  # default version, to be updated later in _parse_ModelProto
+    obj["parsed_node"] = OnnxNode(obj["op_type"], opset_id, tuple(obj["input"]), tuple(obj["output"]), attributes)
+    return obj
+
+  def _parse_TensorProto(self) -> dict:
     obj: dict[str, Any] = {"dims": []}
     for fid, wire_type in self._parse_message(self._decode_end_pos()):
       match fid:
@@ -215,16 +220,18 @@ class OnnxPBParser:
     # load external data
     if self.load_external_data and obj.get("data_location", 0) == 1:
       if "external_data" not in obj: raise ValueError("no external_data")
-      ext = dict(obj["external_data"])
-      if "location" not in ext: raise ValueError("no location in external_data")
-      offset = int(ext.get("offset", "0"))
-      length = int(ext["length"]) if "length" in ext else None
+      location, length, offset = None, None, 0
+      for kv in obj["external_data"]:
+        if kv["key"] == "location": location = kv["value"]
+        elif kv["key"] == "offset": offset = int(kv["value"])
+        elif kv["key"] == "length": length = int(kv["value"])
+      if location is None: raise ValueError("no location in external_data")
 
       if self.file_path is None:
         if isinstance(self.tensor.device, str) and self.tensor.device.startswith("DISK:"):
           self.file_path = pathlib.Path(self.tensor.device[5:])
         else: raise ValueError("onnx external_data needs the origin file path, try passing onnx file path to onnx_load")
-      ext_path = self.file_path.parent.joinpath(ext["location"])
+      ext_path = self.file_path.parent.joinpath(location)
       if not ext_path.exists(): raise FileNotFoundError(f"external location not exists: {ext_path}")
 
       ext_tensor = Tensor(ext_path)
@@ -234,20 +241,23 @@ class OnnxPBParser:
     # parse tensor
     to_dtype = dtype_fallback(true_dtype := OnnxDataType(obj['data_type']).to_dtype(), "buffer parse")
     shape = tuple(obj['dims'])
-    data_fields = [f for f in ('float_data','int32_data','int64_data','double_data','uint64_data','raw_data') if f in obj]
-    data = obj[get_single_element(data_fields)]
-    name = obj.get("name", "")
-    if not isinstance(data, Tensor): return name, Tensor(data, dtype=to_dtype).reshape(shape)
-    assert data.dtype == dtypes.uint8, data
+    present_fields = [field for field in ['float_data', 'int32_data', 'int64_data', 'double_data', 'uint64_data', 'raw_data'] if field in obj]
+    assert len(present_fields) == 1, f"only 1 data field is allowed from {obj=}"
+    data = obj[present_fields[0]]
+    if not isinstance(data, Tensor):
+      obj["parsed_tensor"] = Tensor(data, dtype=to_dtype).reshape(shape)
+      return obj
+    assert isinstance(data, Tensor) and data.dtype == dtypes.uint8, data
     data = data.bitcast(true_dtype).reshape(shape)
     data = data.to(Device.DEFAULT) if true_dtype is to_dtype else data.to("cpu").cast(to_dtype).to(Device.DEFAULT)
     # const folding
     if shape == ():
       if data.dtype == dtypes.float16 and sys.version_info < (3, 12): data = data.cast(dtypes.float32)
       data = Tensor(data.item(), dtype=to_dtype).reshape(shape)
-    return name, data
+    obj["parsed_tensor"] = data
+    return obj
 
-  def _parse_AttributeProto(self) -> tuple[str, Any]:
+  def _parse_AttributeProto(self) -> dict:
     obj: dict[str, Any] = {"floats": [], "ints": [], "strings": []}
     for fid, wire_type in self._parse_message(self._decode_end_pos()):
       match fid:
@@ -255,7 +265,7 @@ class OnnxPBParser:
         case 2: obj["f"] = self.reader.read_float()
         case 3: obj["i"] = self.reader.read_int64()
         case 4: obj["s"] = self.reader.read_bytes().data().tobytes().decode("utf8")
-        case 5: obj["t"] = self._parse_TensorProto()[1]
+        case 5: obj["t"] = self._parse_TensorProto()['parsed_tensor']
         case 6: obj["g"] = OnnxRunner._from_subgraph(self._parse_GraphProto())
         case 7: obj["floats"].append(self.reader.read_float())
         case 8: obj["ints"].append(self.reader.read_int64())
@@ -263,22 +273,26 @@ class OnnxPBParser:
         case 20: obj["type"] = self.reader.read_int64()
         case _: self.reader.skip_field(wire_type)
     obj["floats"], obj["ints"], obj["strings"] = tuple(obj["floats"]), tuple(obj["ints"]), tuple(obj["strings"])
-    return obj["name"], obj[AttributeType(obj["type"]).to_field_name()]
+    return obj
 
-  def _parse_ValueInfoProto(self) -> tuple[str, OnnxValue|None]:
-    name, type_obj = "", None
+  def _parse_ValueInfoProto(self) -> dict:
+    obj: dict[str, Any] = {}
     for fid, wire_type in self._parse_message(self._decode_end_pos()):
       match fid:
-        case 1: name = self.reader.read_string()
-        case 2: type_obj = self._parse_TypeProto()
+        case 1: obj["name"] = self.reader.read_string()
+        case 2: obj["type"] = self._parse_TypeProto()
         case _: self.reader.skip_field(wire_type)
-    if type_obj is None: return name, None
+
+    # parse type
+    if "type" not in obj: return {**obj, "parsed_type": None}
+    type_obj = obj["type"]
     if is_optional := "optional_type" in type_obj: type_obj = type_obj["optional_type"]["elem_type"]
     if is_sequence := "sequence_type" in type_obj: type_obj = type_obj["sequence_type"]["elem_type"]
     assert "tensor_type" in type_obj, type_obj
     shape_dims = type_obj['tensor_type'].get('shape', {}).get('dim', [])
-    return name, OnnxValue(tuple(d.get('dim_param') or d.get('dim_value') for d in shape_dims),
-                           OnnxDataType(type_obj['tensor_type']['elem_type']).to_dtype(), is_optional, is_sequence)
+    obj['parsed_type'] = OnnxValue(tuple(d.get('dim_param') or d.get('dim_value') for d in shape_dims),
+                                   OnnxDataType(type_obj['tensor_type']['elem_type']).to_dtype(), is_optional, is_sequence)
+    return obj
 
   def _parse_TypeProto(self) -> dict:
     obj: dict[str, Any] = {}
@@ -324,24 +338,23 @@ class OnnxPBParser:
         case _: self.reader.skip_field(wire_type)
     return obj
 
-  def _parse_StringStringEntryProto(self) -> tuple[str, str]:
-    key, value = "", ""
+  def _parse_StringStringEntryProto(self) -> dict:
+    obj: dict[str, Any] = {}
     for fid, wire_type in self._parse_message(self._decode_end_pos()):
       match fid:
-        case 1: key = self.reader.read_string()
-        case 2: value = self.reader.read_string()
+        case 1: obj["key"] = self.reader.read_string()
+        case 2: obj["value"] = self.reader.read_string()
         case _: self.reader.skip_field(wire_type)
-    return key, value
+    return obj
 
-  def _parse_OperatorSetIdProto(self) -> OpSetId:
-    domain: str|None = None
-    version = 1
+  def _parse_OperatorSetIdProto(self) -> dict:
+    obj: dict[str, Any] = {}
     for fid, wire_type in self._parse_message(self._decode_end_pos()):
       match fid:
-        case 1: domain = self.reader.read_string()
-        case 2: version = self.reader.read_int64()
+        case 1: obj["domain"] = self.reader.read_string()
+        case 2: obj["version"] = self.reader.read_int64()
         case _: self.reader.skip_field(wire_type)
-    return OpSetId(Domain.from_onnx(domain), version)
+    return obj
 
 # ***** python const *****
 required_input_python_consts: dict[str, tuple[int, ...]] = {
@@ -367,15 +380,16 @@ class OnnxRunner:
     model_path: The ONNX model, provided as a file path (a string or Path object) or a Tensor.
   """
   def __init__(self, model_path: Tensor | str | pathlib.Path):
-    self._init_from_graph(OnnxPBParser(model_path, load_external_data=True).parse())
+    model = OnnxPBParser(model_path, load_external_data=True).parse()
+    self._init_from_graph(model["graph"])
 
   def _init_from_graph(self, graph: dict, is_subgraph: bool = False):
-    self.is_training = any(n.opset_id.domain in {Domain.AI_ONNX_TRAINING, Domain.AI_ONNX_PREVIEW_TRAINING} for n in graph["node"])
+    self.is_training = any(n['parsed_node'].opset_id.domain in {Domain.AI_ONNX_TRAINING, Domain.AI_ONNX_PREVIEW_TRAINING} for n in graph["node"])
     self.graph_name = graph["name"] if is_subgraph else ""
-    self.graph_values: dict[str, Any] = {"": None, **dict(graph["initializer"])}
-    self.graph_inputs = {name: typ for name, typ in graph["input"] if name not in self.graph_values}
-    self.graph_outputs = tuple(name for name, _ in graph["output"])
-    self.graph_nodes = tuple(graph["node"])
+    self.graph_values = {"": None, **{i["name"]: i["parsed_tensor"] for i in graph["initializer"]}}
+    self.graph_inputs = {i["name"]: i["parsed_type"] for i in graph["input"] if i["name"] not in self.graph_values}
+    self.graph_outputs = tuple(o["name"] for o in graph["output"])
+    self.graph_nodes = tuple(n["parsed_node"] for n in graph["node"])
     # track names from initializers and Constant nodes for fast path optimizations
     self.const_names: set[str] = set(self.graph_values.keys()) | {o for n in self.graph_nodes if n.op == "Constant" for o in n.outputs}
 

tinygrad/nn/state.py

@@ -336,14 +336,15 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
       d = blocks[:,-2:].bitcast(dtypes.float16).cast(dtypes.float32).expand((-1, 256))
       return d * (xl.bitwise_or(xh).bitcast(dtypes.int8) - 32).flatten(-2) * scales
     if ggml_type == 39:
-      e = blocks[:, 0].cast(dtypes.uint32)
-      small_bits = Tensor([0x00200000, 0x00400000], dtype=dtypes.uint32, device=t.device)[e.clip(0, 1).cast(dtypes.int32)] # e = 0 or e = 1 case
-      d = (e < 2).where(small_bits, ((e - 1) * 0x00800000).cast(dtypes.uint32)).bitcast(dtypes.float32).unsqueeze(-1)
+      e_int = blocks[:, 0].cast(dtypes.int32)
+      d = ((e_int >= 2).cast(dtypes.float32) * (e_int.cast(dtypes.float32) - 128).exp2() +
+           (e_int == 1).cast(dtypes.float32) * 2.0**(-127) +
+           (e_int == 0).cast(dtypes.float32) * 2.0**(-128)).unsqueeze(-1)
       codes = q_to_uint8(blocks[:, 1:17], 4)
-      fp4_lut = Tensor([0.0, 1.0, 2.0, 3.0, 4.0, 6.0, 8.0, 12.0,
-                       -0.0,-1.0,-2.0,-3.0,-4.0,-6.0,-8.0,-12.0],
-                       dtype=dtypes.float32, device=t.device)
-      fp4_val = fp4_lut[codes]
+      sign = 1.0 - codes.rshift(3).cast(dtypes.float32) * 2.0
+      exp, mant = codes.rshift(1).bitwise_and(0x3).cast(dtypes.float32), codes.bitwise_and(0x1).cast(dtypes.float32)
+      fp4_val = sign * 2.0 * ((exp != 0).cast(dtypes.float32) * (1.0 + 0.5 * mant) * (exp - 1.0).exp2() +
+                              (exp == 0).cast(dtypes.float32) * 0.5 * mant)
       return (fp4_val * d).flatten(-2)[:n]
   raise ValueError(f"GGML type '{ggml_type}' is not supported!")
 

tinygrad/renderer/amd/sqtt.py

@@ -106,7 +106,7 @@ class InstOpRDNA4(Enum):
   SMEM = 0x1
   JUMP = 0x3
   JUMP_NO = 0x4
-  CALL = 0x5
+  JUMP_UNCOND = 0x5
   MESSAGE = 0x9
   VALU_TRANS = 0xb
   VALU_B2 = 0xd
@@ -187,16 +187,19 @@ class TS_DELTA_SHORT(PacketType):
 class TS_DELTA_OR_MARK(PacketType):
   encoding = bits[6:0] == 0b0000001
   delta = bits[47:12]
-  pl = bits[8:8]
-  rt = bits[9:9]
+  bit8 = bits[8:8]
+  bit9 = bits[9:9]
   @property
-  def is_marker(self) -> bool: return bool(self.rt and not self.pl)
+  def is_marker(self) -> bool: return bool(self.bit9 and not self.bit8)
 
-class TS_DELTA_OR_MARK_RDNA4(TS_DELTA_OR_MARK):
+class TS_DELTA_OR_MARK_RDNA4(PacketType):  # Layout 4: 48->64 bits
+  encoding = bits[6:0] == 0b0000001
   delta = bits[63:12]
-  rt = bits[7:7]
-  pl = bits[8:8]
-  tl = bits[9:9]
+  bit7 = bits[7:7]
+  bit8 = bits[8:8]
+  bit9 = bits[9:9]
+  @property
+  def is_marker(self) -> bool: return bool((self.bit9 and not self.bit8) or self.bit7)
 
 class TS_DELTA_S5_W2(PacketType):
   encoding = bits[4:0] == 0b11100
@@ -402,17 +405,16 @@ class CDNA_PKT_2(PacketType):
   unk_padding = bits[63:8]
 
 class CDNA_WAVESTART(PacketType):
-  """type 3: 32-bit wave start (Wave/group_id)"""
+  """pkt_fmt=3: 32-bit WAVESTART packet (case 0x8)"""
   encoding = bits[3:0] == 3
-  sh = bits[5:5]
-  cu = bits[9:6]
-  wave = bits[13:10]
-  simd = bits[15:14]
-  pipe = bits[17:16]
-  me = bits[19:18]
-  _gap = bits[21:20]
-  count = bits[28:22]
-  _padding = bits[31:29]
+  unk_0 = bits[5:5]       # (data >> 5) & 1
+  unk_1 = bits[9:6]       # (data >> 6) & 0xf
+  wave = bits[13:10]      # (data >> 10) & 0xf
+  simd = bits[15:14]      # (data >> 0xe) & 3
+  cu = bits[17:16]        # (data >> 0x10) & 3
+  unk_5 = bits[19:18]     # (data >> 0x12) & 3
+  unk_6 = bits[28:22]     # (data >> 0x16) & 0x7f
+  unk_padding = bits[31:29]
 
 class CDNA_PKT_4(PacketType):
   """pkt_fmt=4: 16-bit packet (case 0xc, same as 0x8/0x14)"""
@@ -422,21 +424,21 @@ class CDNA_PKT_4(PacketType):
   unk_2 = bits[13:10]     # (data_word >> 10) & 0xf
   unk_3 = bits[15:14]     # (data_word >> 0xe)
 
-class REGCS_CDNA(PacketType):
-  """type 5: 48-bit register CS write (RegCs)"""
+class CDNA_PKT_5(PacketType):
+  """pkt_fmt=5: 48-bit packet (case 0x10)"""
   encoding = bits[3:0] == 5
-  pipe = bits[6:5]
-  _me_raw = bits[8:7]
-  regaddr = bits[15:9]
-  regdata = bits[47:16]
+  unk_0 = bits[6:5]       # (data >> 5) & 3
+  unk_1 = bits[7:7]       # (data >> 7) + 1 & 1
+  unk_2 = bits[15:9]      # (data >> 9) & 0x7f
+  unk_padding = bits[47:16]
 
 class CDNA_WAVEEND(PacketType):
-  """type 6: 16-bit wave end (group_id)"""
+  """pkt_fmt=6: 16-bit WAVEEND packet (case 0x14, same as 0x8/0xc)"""
   encoding = bits[3:0] == 6
-  sh = bits[5:5]
-  cu = bits[9:6]
-  wave = bits[13:10]
-  simd = bits[15:14]
+  unk_0 = bits[5:5]       # (data_word >> 5) & 1
+  unk_1 = bits[9:6]       # (data_word >> 6) & 0xf
+  wave = bits[13:10]      # (data_word >> 10) & 0xf
+  simd = bits[15:14]      # (data_word >> 0xe)
 
 class CDNA_EXEC(PacketType):
   """pkt_fmt=10: 16-bit EXEC packet (case 0x24)"""
@@ -509,7 +511,7 @@ class CDNA_PKT_15(PacketType):
   unk_padding = bits[47:16]
 
 PACKET_TYPES_CDNA: dict[int, type[PacketType]] = {
-  0: CDNA_DELTA, 1: CDNA_TIMESTAMP, 2: CDNA_PKT_2, 3: CDNA_WAVESTART, 4: CDNA_PKT_4, 5: REGCS_CDNA, 6: CDNA_WAVEEND,
+  0: CDNA_DELTA, 1: CDNA_TIMESTAMP, 2: CDNA_PKT_2, 3: CDNA_WAVESTART, 4: CDNA_PKT_4, 5: CDNA_PKT_5, 6: CDNA_WAVEEND,
   7: CDNA_PKT_7, 8: CDNA_PKT_8, 9: CDNA_PKT_9, 10: CDNA_EXEC, 11: CDNA_PKT_11, 12: CDNA_PKT_12,
   13: CDNA_INST, 14: CDNA_PKT_14, 15: CDNA_PKT_15,
 }

tinygrad/renderer/cstyle.py

@@ -566,11 +566,4 @@ class AMDHIPCCRenderer(AMDHIPRenderer):
     super().__init__(arch)
     self.compiler = HIPCCCompiler(arch)
 
-class QCOMCLRenderer(OpenCLRenderer):
-  device = "QCOM"
-
-  def __init__(self, chip_id):
-    from tinygrad.runtime.support.compiler_qcom import QCOMCompiler
-    self.chip_id, self.compiler = chip_id, QCOMCompiler(chip_id)
-
-  def __reduce__(self): return self.__class__, (self.chip_id,)
+class QCOMRenderer(OpenCLRenderer): device = "QCOM"

tinygrad/runtime/ops_amd.py

@@ -485,7 +485,6 @@ class AMDCopyQueue(HWQueue):
     if (dev:=signal.owner) is not None and signal.is_timeline and not dev.is_am():
       self.q(self.sdma.SDMA_OP_FENCE | fence_flags, *data64_le(dev.queue_event_mailbox_ptr), dev.queue_event.event_id)
       self.q(self.sdma.SDMA_OP_TRAP, self.sdma.SDMA_PKT_TRAP_INT_CONTEXT_INT_CONTEXT(dev.queue_event.event_id))
-    elif dev is not None and dev.is_am(): self.q(self.sdma.SDMA_OP_TRAP, 0)
 
     return self
 
@@ -867,12 +866,10 @@ class PCIIface(PCIIfaceBase):
     return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbell=self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q'), put_value=0,
       read_ptr=gart.cpu_view().view(offset=rptr, size=8, fmt='Q'), write_ptr=gart.cpu_view().view(offset=wptr, size=8, fmt='Q'), params=rcvr_params)
 
-  def _collect_interrupts(self, reset=False, drain_only=False):
+  def _collect_faults(self, reset=False):
     devs:list[AMDDevice] = [d for pg in HCQCompiled.peer_groups.values() for d in pg if isinstance(d, AMDDevice) and d.is_am()]
     for d in devs:
-      if drain_only: d.iface.dev_impl.ih.drain()
-      else: d.iface.dev_impl.ih.interrupt_handler()
-
+      d.iface.dev_impl.ih.interrupt_handler()
       if reset and d.iface.dev_impl.recover(force=d.error_state is not None):
         d.compute_queue.put_value = d.compute_queue.read_ptr[0] = d.compute_queue.write_ptr[0] = 0
         d.iface.dev_impl.gfx.setup_ring(*d.compute_queue.params)
@@ -882,11 +879,11 @@ class PCIIface(PCIIfaceBase):
   def sleep(self, timeout):
     if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
       self.pci_dev.irq_fd.read(8 * events_cnt)
-    self._collect_interrupts()
+    self._collect_faults()
     if self.dev_impl.is_err_state: raise RuntimeError("Device is in error state")
 
   def on_device_hang(self):
-    self._collect_interrupts(reset=True)
+    self._collect_faults(reset=True)
     raise RuntimeError("Device hang detected")
 
   def device_fini(self): self.dev_impl.fini()
@@ -1079,10 +1076,6 @@ class AMDDevice(HCQCompiled):
     self.hw_compute_queue_t().memory_barrier().signal(self.timeline_signal, self.next_timeline()).submit(self)
     self.synchronize()
 
-  def synchronize(self, timeout:int|None=None):
-    super().synchronize(timeout)
-    if self.is_am() and self.error_state is None: self.iface._collect_interrupts(reset=False, drain_only=True)
-
   def on_device_hang(self): self.iface.on_device_hang()
 
   def device_props(self): return self.iface.props

tinygrad/runtime/ops_null.py

@@ -1,9 +1,9 @@
 import functools
 from tinygrad.device import Compiled, Allocator, CompilerSet
 from tinygrad.engine.jit import MultiGraphRunner
-from tinygrad.renderer.cstyle import Renderer, CStyleLanguage, AMDHIPRenderer, QCOMCLRenderer
+from tinygrad.renderer.cstyle import Renderer, CStyleLanguage, AMDHIPRenderer
 from tinygrad.uop.ops import Ops
-from tinygrad.helpers import cpu_profile, EMULATE, NULL_QCOMCL, NULL_IR3, NULL_NAK, NULL_ALLOW_COPYOUT
+from tinygrad.helpers import cpu_profile, EMULATE, NULL_IR3, NULL_NAK, NULL_ALLOW_COPYOUT
 from tinygrad.renderer.nir import IR3Renderer, NAKRenderer
 
 class NullRenderer(CStyleLanguage):
@@ -39,7 +39,6 @@ class NullDevice(Compiled):
       case "AMD_CDNA4": renderer = functools.partial(AMDHIPRenderer, "gfx950")
       case "": renderer = NullRenderer
       case _: raise RuntimeError(f"can't EMULATE device: {EMULATE.value}")
-    compilers = CompilerSet([(renderer, None), (functools.partial(QCOMCLRenderer, 0x6030001), NULL_QCOMCL), # adreno 630
-                             (functools.partial(IR3Renderer, 0x6030001), NULL_IR3), # adreno 630
+    compilers = CompilerSet([(renderer, None), (functools.partial(IR3Renderer, 0x6030001), NULL_IR3), # adreno 630
                              (functools.partial(NAKRenderer, "sm_120", 48), NULL_NAK)]) # 5090
     super().__init__(device, NullAllocator(self), compilers, functools.partial(NullProgram, device), NullGraph)

tinygrad/runtime/ops_nv.py

@@ -718,8 +718,6 @@ class NVDevice(HCQCompiled[NVSignal]):
       if coloc_size > self.vid_coloc_buf.size: self.vid_coloc_buf, _ = self._realloc(self.vid_coloc_buf, coloc_size, force=True)
       if filter_size > self.vid_filter_buf.size: self.vid_filter_buf, _ = self._realloc(self.vid_filter_buf, filter_size, force=True)
 
-  def hw_copy_queues(self): return super().hw_copy_queues() + ([("NVDEC:0", NVVideoQueue)] if hasattr(self, 'vid_gpfifo') else [])
-
   def invalidate_caches(self):
     if self.is_nvd(): self.iface.rm_control(self.subdevice, nv_gpu.NV2080_CTRL_CMD_INTERNAL_BUS_FLUSH_WITH_SYSMEMBAR, None)
     else:

tinygrad/runtime/ops_qcom.py

@@ -6,10 +6,11 @@ from tinygrad.device import BufferSpec, CompilerSet, Device
 from tinygrad.runtime.support.hcq import HCQBuffer, HWQueue, HCQProgram, HCQCompiled, HCQAllocatorBase, HCQSignal, HCQArgsState, BumpAllocator
 from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface
 from tinygrad.runtime.autogen import kgsl, mesa
-from tinygrad.renderer.cstyle import QCOMCLRenderer
+from tinygrad.runtime.ops_cl import CLDevice
+from tinygrad.renderer.cstyle import QCOMRenderer
 from tinygrad.renderer.nir import IR3Renderer
-from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, ceildiv, prod, cpu_profile, lo32, suppress_finalizing
-from tinygrad.helpers import next_power2, flatten, QCOM_IR3, QCOM_CC, PROFILE
+from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, ceildiv, prod, fromimport, cpu_profile, lo32, suppress_finalizing
+from tinygrad.helpers import next_power2, flatten, QCOM_IR3, QCOM_CC, PROFILE, DEBUG
 from tinygrad.dtype import ImageDType, dtypes
 from tinygrad.runtime.support.system import System
 if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl  # noqa: F401  # pylint: disable=unused-import
@@ -247,7 +248,9 @@ class QCOMProgram(HCQProgram):
       self.tex_off, self.ibo_off, self.samp_off = 2048, 2048 + 0x40 * self.tex_cnt, 2048 + 0x40 * (self.tex_cnt + self.ibo_cnt)
       self.fregs, self.hregs = v.info.max_reg + 1, v.info.max_half_reg + 1
       self.consts_info:list[tuple] = []
-    else: self._parse_lib(lib)
+    else:
+      self._parse_lib(lib:=self.dev.cl_dev.cl_compiler.compile_cached(lib.decode()))
+      if DEBUG >= 7: fromimport('tinygrad.runtime.support.compiler_mesa', 'disas_adreno')(lib[(ofs:=_read_lib(lib, 0xc0)):ofs+_read_lib(lib, 0x100)])
 
     self.lib_gpu: HCQBuffer = self.dev.allocator.alloc(self.image_size, buf_spec:=BufferSpec(cpu_access=True, nolru=True))
     to_mv(self.lib_gpu.va_addr, self.image_size)[:] = self.image
@@ -381,8 +384,8 @@ class QCOMDevice(HCQCompiled):
     if PROFILE and self.gpu_id[:2] < (7, 3):
       System.write_sysfs("/sys/class/kgsl/kgsl-3d0/idle_timer", value="4000000000", msg="Failed to disable suspend mode", expected="4294967276")
 
-    compilers = CompilerSet(ctrl_var=QCOM_CC, cset=[(functools.partial(QCOMCLRenderer, info.chip_id), None),
-                                                    (functools.partial(IR3Renderer, info.chip_id), QCOM_IR3)])
+    self.cl_dev = CLDevice(device)
+    compilers = CompilerSet(ctrl_var=QCOM_CC, cset=[(QCOMRenderer, None), (functools.partial(IR3Renderer, info.chip_id), QCOM_IR3)])
     super().__init__(device, QCOMAllocator(self), compilers, functools.partial(QCOMProgram, self), QCOMSignal,
                      functools.partial(QCOMComputeQueue, self), None)
 

tinygrad/runtime/support/am/ip.py

@@ -425,16 +425,6 @@ class AM_IH(AM_IP):
     if self.adev.ip_ver[am.NBIO_HWIP][:2] != (7,9):
       self.adev.soc.doorbell_enable(port=1, awid=0x0, awaddr_31_28_value=0x0, offset=am.AMDGPU_NAVI10_DOORBELL_IH*2, size=2)
 
-  def drain(self):
-    _, _, suf, _ = self.rings[0]
-    wptr = self.adev.reg(f"regIH_RB_WPTR{suf}").read_bitfields()
-    self.adev.regIH_RB_RPTR.write(wptr['offset'] % (self.ring_size // 4))
-
-    if wptr['rb_overflow']:
-      self.adev.reg(f"regIH_RB_WPTR{suf}").update(rb_overflow=0)
-      self.adev.reg(f"regIH_RB_CNTL{suf}").update(wptr_overflow_clear=1)
-      self.adev.reg(f"regIH_RB_CNTL{suf}").update(wptr_overflow_clear=0)
-
   def interrupt_handler(self):
     _, _, suf, _ = self.rings[0]
     wptr = self.adev.reg(f"regIH_RB_WPTR{suf}").read_bitfields()
@@ -442,15 +432,11 @@ class AM_IH(AM_IP):
 
     while rptr != wptr['offset']:
       entry = [self.ring_view[(rptr + i) % (self.ring_size // 4)] for i in range(8)]
-      rptr = (rptr + 8) % (self.ring_size // 4)
-
       client, src, ring_id, vmid, vmid_type, pasid, node = \
         [getattr(am, f'SOC15_{n}_FROM_IH_ENTRY')(entry) for n in ['CLIENT_ID', 'SOURCE_ID', 'RING_ID', 'VMID', 'VMID_TYPE', 'PASID', 'NODEID']]
       ctx = [getattr(am, f'SOC15_CONTEXT_ID{i}_FROM_IH_ENTRY')(entry) for i in range(4)]
 
       src_name = self.adev.soc.ih_srcs_names.get(client, {}).get(src, '')
-      if src_name in {"SDMA_TRAP", "CP_EOP_INTR"}: continue
-
       print(f"am {self.adev.devfmt}: IH ({rptr:#x}/{wptr['offset']:#x}) client={self.adev.soc.ih_clients.get(client)} src={src_name}({src}) "
             f"ring={ring_id} vmid={vmid}({vmid_type}) pasid={pasid} node={node} ctx=[{ctx[0]:#x}, {ctx[1]:#x}, {ctx[2]:#x}, {ctx[3]:#x}]")
 
@@ -468,7 +454,14 @@ class AM_IH(AM_IP):
         self.adev.is_err_state = True
       else: self.adev.is_err_state = True
 
-    self.drain()
+      rptr = (rptr + 8) % (self.ring_size // 4)
+
+    if wptr['rb_overflow']:
+      self.adev.reg(f"regIH_RB_WPTR{suf}").update(rb_overflow=0)
+      self.adev.reg(f"regIH_RB_CNTL{suf}").update(wptr_overflow_clear=1)
+      self.adev.reg(f"regIH_RB_CNTL{suf}").update(wptr_overflow_clear=0)
+
+    self.adev.regIH_RB_RPTR.write(wptr['offset'] % (self.ring_size // 4))
 
     bif_intr = self.adev.regBIF_BX0_BIF_DOORBELL_INT_CNTL.read_bitfields()
     athub_err, cntlr_err = bif_intr['ras_athub_err_event_interrupt_status'], bif_intr['ras_cntlr_interrupt_status']
@@ -499,8 +492,7 @@ class AM_SDMA(AM_IP):
                                                           inst=inst)
         self.adev.reg(f"regSDMA{pipe}_{self.sdma_name}_CNTL").update(halt=0, **{f"{'th1_' if self.sdma_name == 'F32' else ''}reset":0}, inst=inst)
 
-      self.adev.reg(f"regSDMA{pipe}_CNTL").update(trap_enable=1,
-        **({'utc_l1_enable':1} if self.adev.ip_ver[am.SDMA0_HWIP] <= (5,2,0) else {}), inst=inst)
+      self.adev.reg(f"regSDMA{pipe}_CNTL").update(**({'utc_l1_enable':1} if self.adev.ip_ver[am.SDMA0_HWIP] <= (5,2,0) else {}), inst=inst)
 
     if self.adev.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)}:
       for aid_id in range(4):

tinygrad/runtime/support/compiler_qcom.py

@@ -1,62 +0,0 @@
-import ctypes, struct
-from tinygrad.device import Compiler
-from tinygrad.helpers import DEBUG, system
-from tinygrad.runtime.support.c import DLL
-from tinygrad.runtime.support.compiler_mesa import disas_adreno
-
-# see https://github.com/sirhcm/tinydreno
-dll = DLL("llvm-qcom", ["llvm-qcom"])
-
-(create_llvm_instance:=dll.cl_compiler_create_llvm_instance).restype, create_llvm_instance.argtypes = ctypes.c_void_p, []
-
-(compile_source:=dll.cl_compiler_compile_source).restype = ctypes.c_void_p
-compile_source.argtypes = [ctypes.c_void_p, ctypes.c_uint64, ctypes.c_int, ctypes.c_char_p, ctypes.c_int, ctypes.c_uint64, ctypes.c_uint64,
-                           ctypes.c_char_p, ctypes.c_uint64, ctypes.c_uint64, ctypes.c_void_p]
-
-(link_program:=dll.cl_compiler_link_program).restype = ctypes.c_void_p
-link_program.argtypes = [ctypes.c_void_p, ctypes.c_uint64, ctypes.c_int, ctypes.c_char_p, ctypes.c_int, ctypes.c_void_p]
-
-(get_error_code:=dll.cl_compiler_get_error_code).restype, get_error_code.argtypes = ctypes.c_int, [ctypes.c_void_p]
-(get_build_log:=dll.cl_compiler_get_build_log).restype, get_build_log.argtypes = ctypes.c_char_p, [ctypes.c_void_p]
-
-(handle_create_binary:=dll.cl_compiler_handle_create_binary).restype = None
-handle_create_binary.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p), ctypes.POINTER(ctypes.c_size_t)]
-
-(free_handle:=dll.cl_compiler_free_handle).restype, free_handle.argtypes = None, [ctypes.c_void_p]
-(free_assembly:=dll.cl_compiler_free_assembly).restype, free_assembly.argtypes = None, [ctypes.c_void_p]
-(destroy_llvm_instance:=dll.cl_compiler_destroy_llvm_instance).restype, destroy_llvm_instance.argtypes = None, [ctypes.c_void_p]
-
-MODE_32BIT, MODE_64BIT, SRC_STR, SRC_BLOB = 0, 1, 0, 1
-
-def _read_lib(lib, off) -> int: return struct.unpack("I", lib[off:off+4])[0]
-
-class QCOMCompiler(Compiler):
-  def __init__(self, chip_id):
-    self.chip_id, self.llvm_inst = chip_id, create_llvm_instance()
-    super().__init__(f"compile_qcomcl_{chip_id}")
-
-  def __del__(self): destroy_llvm_instance(self.llvm_inst)
-
-  def __reduce__(self): return QCOMCompiler, (self.chip_id,)
-
-  def checked(self, handle):
-    if handle is None or get_error_code(handle) != 0:
-      destroy_llvm_instance(self.llvm_inst)
-      self.llvm_inst = create_llvm_instance()
-      raise RuntimeError("QCOM Compilation Error" + ("" if handle is None else f": {get_build_log(handle)}"))
-    return handle
-
-  def compile(self, src) -> bytes:
-    ch = self.checked(compile_source(self.llvm_inst, self.chip_id, MODE_64BIT, b"", 0, 0, 0, src.encode(), 0, SRC_STR, None))
-    if DEBUG >= 8:
-      handle_create_binary(ch, ctypes.byref(ptr:=ctypes.c_void_p()), ctypes.byref(sz:=ctypes.c_size_t()))
-      print(system("llvm-dis", input=ctypes.string_at(ptr, sz.value)[16:]))
-      free_assembly(ptr)
-    lh = self.checked(link_program(self.llvm_inst, self.chip_id, MODE_64BIT, None, 1, ctypes.pointer(ctypes.c_void_p(ch))))
-    handle_create_binary(lh, ctypes.byref(ptr:=ctypes.c_void_p()), ctypes.byref(sz:=ctypes.c_size_t()))
-    for h in [ch, lh]: free_handle(h)
-    ret = ctypes.string_at(ptr, sz.value)
-    free_assembly(ptr)
-    return ret
-
-  def disassemble(self, lib: bytes): disas_adreno(lib[(ofs:=_read_lib(lib, 0xc0)):ofs+_read_lib(lib, 0x100)], self.chip_id)

tinygrad/schedule/indexing.py

@@ -9,7 +9,7 @@ from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored
 
 ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
                      Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.PARAM,
-                     Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.CALL}
+                     Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.CALL, Ops.ENCDEC}
 
 def realize(ctx:dict[UOp, None], tr:UOp) -> None: ctx[tr] = None
 
@@ -18,8 +18,8 @@ def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None:
     if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None
 
 def realize_assign_src(ctx:dict[UOp, None], buf:UOp, x:UOp):
-  # don't realize COPY/BUFFER_VIEW when they are the direct source of ASSIGN — the ASSIGN target buffer is the output
-  if x.op in {Ops.COPY, Ops.BUFFER_VIEW} and x in ctx \
+  # don't realize COPY/BUFFER_VIEW/ENCDEC when they are the direct source of ASSIGN — the ASSIGN target buffer is the output
+  if x.op in {Ops.COPY, Ops.BUFFER_VIEW, Ops.ENCDEC} and x in ctx \
      and not buf.op_in_backward_slice_with_self(Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD):
     del ctx[x]
   # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce
@@ -29,9 +29,9 @@ pm_generate_realize_map = PatternMatcher([
   # always realize SINK src
   (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
   # always realize
-  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS, Ops.STORE, Ops.ASSIGN}, name="tr"), realize),
+  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS, Ops.STORE, Ops.ASSIGN, Ops.ENCDEC}, name="tr"), realize),
   # realize srcs of these
-  (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
+  (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK, Ops.ENCDEC), name="rb"), realize_srcs),
   # sometimes realize src of assign
   (UPat(Ops.ASSIGN, src=(UPat.var("buf"), UPat.var("x"))), realize_assign_src),
 ])

tinygrad/schedule/multi.py

@@ -37,7 +37,7 @@ replace_allreduce = PatternMatcher([
 _early_allreduce = PatternMatcher([
   (UPat(Ops.ALLREDUCE, src=(UPat.var("buf"), UPat()), name="red"), handle_allreduce),
 ])
-if not getenv("LATE_ALLREDUCE", 1): replace_allreduce = _early_allreduce + replace_allreduce
+if not getenv("LATE_ALLREDUCE", 0): replace_allreduce = _early_allreduce + replace_allreduce
 
 # ***** multi functions *****
 

tinygrad/schedule/rangeify.py

@@ -112,7 +112,19 @@ def resolve_call(c:UOp, allow_param_mismatch=True) -> UOp|None:
     if p.dtype != a.dtype: raise TypeError(f"arg {i} dtype mismatch: expected {p.dtype}, got {a.dtype}")
   return c.src[0].substitute(dict_map, walk=True)
 
-earliest_rewrites = mop_cleanup+PatternMatcher([
+def lower_cat(cat:UOp) -> UOp:
+  axis = cat.arg
+  dim_acc = list(itertools.accumulate([s.shape[axis] for s in cat.src], initial=0))
+  padded = [s.pad(tuple((dim_acc[i], dim_acc[-1]-dim_acc[i+1]) if j==axis else (0,0) for j in range(len(s.shape)))) for i,s in enumerate(cat.src)]
+  ret = padded[0]
+  for p in padded[1:]: ret = ret.alu(Ops.ADD, p)
+  return ret
+
+pm_lower_cat = PatternMatcher([
+  (UPat(Ops.CAT, name="cat"), lower_cat),
+])
+
+earliest_rewrites = mop_cleanup+pm_lower_cat+PatternMatcher([
   # early fixup const copy
   (UPat(Ops.COPY, src=(UPat.var("s"), UPat.var("d"))),
    lambda s,d: s.substitute({UOp(Ops.DEVICE, arg=s.device):d}) if s.base.op is Ops.CONST else None),
@@ -174,7 +186,7 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
 # *****************
 # 3.5 cleanups
 
-ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN, Ops.NOOP}
+ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN, Ops.ENCDEC, Ops.NOOP}
 
 # you don't know in the first pass if axes are going to die, this happens if there's an EXPAND to the left
 def cleanup_dead_axes(b:UOp):
@@ -359,15 +371,10 @@ def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True):
     assign_target, assign_src = assign.src[0], assign.src[1]
     assert assign_target.op is Ops.INDEX, f"{assign_target.op} is not index"
     while assign_src.op is Ops.NOOP: assign_src = assign_src.src[0]
-
-    store_target = assign_target
-    if assign.arg and assign_target.src[0].op is Ops.BUFFERIZE and assign_target.src[0].src[0].op is Ops.INDEX:
-      # BUFFERIZE(INDEX(...)); store through the underlying global index instead.
-      store_target = assign_target.src[0].src[0]
-
-    end_rngs = sorted(dedup(tuple(store_target.ranges) + tuple(rngs)), key=lambda x: x.arg)
-    ret = store_target.buf_uop.base
-    if assign_src is not store_target: ret = ret.after(store_target.replace(dtype=sdtype).store(assign_src).end(*end_rngs))
+    # skip self-assign from same-device copy, otherwise create the store
+    # in assign, this is the buffer size, not the bufferize size
+    if assign_src is assign_target: ret = assign_target.src[0]
+    else: ret = assign_target.src[0].after(assign_target.replace(dtype=sdtype).store(assign_src).end(*rngs))
     for op, marg in reversed(assign.arg or ()): ret = ret._mop(op, marg)
     return ret
 
@@ -523,11 +530,12 @@ def split_store(x:UOp) -> UOp|None:
   lctx = LocalAddBufferContext()
   ret = graph_rewrite(x, to_define_global+pm_flatten_range+rangeify_codegen, ctx=lctx, name="kernel split", bottom_up=True)
 
-  # SINK requires all buffers on the same device, but COPY/BUFFER_VIEW are cross-device or special hardware ops
+  # SINK requires all buffers on the same device, but COPY/BUFFER_VIEW/ENCDEC are cross-device or special hardware ops
   if ret.op is Ops.STORE: stored = ret.src[1]
   elif ret.op is Ops.END and ret.src[0].op is Ops.STORE: stored = ret.src[0].src[1]
   else: raise RuntimeError(f"unknown kernel type {ret.op}")
   if stored.op in {Ops.COPY, Ops.BUFFER_VIEW}: ret = stored.replace(src=stored.src + ret.ended_ranges)
+  elif stored.op is Ops.ENCDEC: ret = stored
   else: ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts))
 
   kernel = ret.call(*lctx.map.values(), *lctx.vars.keys())
@@ -541,7 +549,8 @@ split_kernels = PatternMatcher([
 
 @profile_matches
 def get_kernel_graph(sink:UOp) -> UOp:
-  tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
+  tsink = graph_rewrite(sink, pm_lower_cat, name="lower_cat")
+  tsink = graph_rewrite(tsink, multi_pm, name="multi_pm")
   if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_assign, ctx={}, name="fold moved assigns")
   tsink = graph_rewrite(tsink, pm_syntactic_sugar+pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")
 

tinygrad/tensor.py

@@ -5,7 +5,7 @@ from contextlib import ContextDecorator
 from typing import Any, Callable, ClassVar, Sequence, cast, get_args, Literal, SupportsIndex, ParamSpec, TypeVar, Generic, TYPE_CHECKING
 if TYPE_CHECKING: import numpy
 from tinygrad.dtype import DType, DTypeLike, dtypes, ImageDType, ConstType, least_upper_float, least_upper_dtype, sum_acc_dtype, to_dtype, truncate
-from tinygrad.dtype import _from_np_dtype, _to_np_dtype, PyConst, Invalid, InvalidType
+from tinygrad.dtype import _from_np_dtype, _to_np_dtype, PyConst
 from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten
 from tinygrad.helpers import IMAGE, FLOAT16, WINO, Metadata, TRACEMETA, ASM_GEMM, ceildiv, fetch, is_numpy_ndarray, TracingKey, cpu_profile
 from tinygrad.helpers import suppress_finalizing, disable_gc
@@ -113,7 +113,7 @@ class Tensor(OpMixin):
   __slots__ = "uop", "requires_grad", "grad"
   training: ClassVar[bool] = False
 
-  def __init__(self, data:ConstType|bytes|list|tuple|UOp|'numpy.ndarray'|pathlib.Path|None,
+  def __init__(self, data:PyConst|bytes|list|tuple|UOp|'numpy.ndarray'|pathlib.Path|None,
                device:str|tuple|list|None=None, dtype:DTypeLike|None=None, requires_grad:bool|None=None, _force_unique:bool=False):
     if device is None and isinstance(data, pathlib.Path): device = f"DISK:{data.resolve()}"  # keep it on the disk if device is None
     _dtype:DType|None = to_dtype(dtype) if dtype is not None else None
@@ -141,9 +141,6 @@ class Tensor(OpMixin):
       data = Tensor(0, device=_device, dtype=_dtype or dtypes.default_float, requires_grad=requires_grad).uop
     elif isinstance(data, get_args(PyConst)):
       data = (UOp.unique_const if _force_unique or requires_grad else UOp.const)(_dtype or dtypes.from_py(data), data, _device)
-    elif isinstance(data, InvalidType):
-      assert _dtype is not None
-      data = UOp.const(_dtype, data, _device)
     elif isinstance(data, bytes): data = _frompy(data, dtypes.uint8 if _dtype is None else _dtype)
     elif isinstance(data, (list, tuple)):
       if _dtype is None:
@@ -1387,10 +1384,8 @@ class Tensor(OpMixin):
     """
     dim = self._resolve_dim(dim)
     for arg in args: assert arg.ndim==self.ndim and all(ti==ai for i,(ti,ai) in enumerate(zip(self.shape, arg.shape)) if i!=dim)
-    tensors = [self, *args]
-    dim_cumsum = list(itertools.accumulate([t.shape[dim] for t in tensors], initial=0))
-    for i,t in enumerate(tensors): tensors[i] = t.pad([(dim_cumsum[i], dim_cumsum[-1]-dim_cumsum[i+1]) if j==dim else None for j in range(t.ndim)])
-    return functools.reduce(Tensor.add, tensors)
+    _ = [t.shape[dim] for t in [self, *args]]  # validate dim in bounds (catches scalar cat)
+    return self._apply_uop(lambda *uops, arg: UOp(Ops.CAT, uops[0].dtype, uops, arg), *args, arg=dim)
 
   def stack(self:Tensor, *args:Tensor, dim:int=0) -> Tensor:
     """
@@ -2947,8 +2942,7 @@ class Tensor(OpMixin):
     if not isinstance(y, Tensor):
       # make y a Tensor
       assert isinstance(y, (*get_args(ConstType), UOp)), f"{type(y)=}, {y=}"
-      if y is Invalid or isinstance(x.dtype, ImageDType) or dtypes.is_float(x.dtype) or (dtypes.is_int(x.dtype) and isinstance(y, int)):
-        y_dtype = x.dtype
+      if isinstance(x.dtype, ImageDType) or dtypes.is_float(x.dtype) or (dtypes.is_int(x.dtype) and isinstance(y, int)): y_dtype = x.dtype
       elif not isinstance(y, UOp): y_dtype = dtypes.from_py(y)
       if isinstance(y, UOp): y = Tensor.from_uop(y, device=x.device)
       else: y = Tensor(dtypes.as_const(y, y_dtype), x.device, y_dtype, requires_grad=False)
@@ -3178,10 +3172,8 @@ class Tensor(OpMixin):
     the reference frames (`ref_frames`).
     """
     ref_frames = [x.contiguous() for x in ref_frames or []]
-    assert frame_pos.op is Ops.BIND, "frame_pos must be a bound Variable"
-    srcs = (out:=Tensor.empty(*shape, device=self.device, dtype=self.dtype), self.contiguous(), state.contiguous(), *ref_frames)
-    fn = UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(frame_pos.src[0], *[UOp.const(dtypes.int, s) for s in shape]), arg="encdec")
-    return Tensor(out.uop.after(fn.call(*[s.uop for s in srcs], frame_pos)), device=self.device)
+    assert isinstance(frame_pos, Variable), "frame_pos must be a Variable"
+    return self.contiguous()._apply_uop(UOp.encdec, state.contiguous(), *ref_frames, extra_args=(frame_pos,), arg=(shape,))
 
   # ***** functional nn ops *****
 
@@ -3655,15 +3647,11 @@ class Tensor(OpMixin):
 
     # contiguous creates the image, and early realize static weights (TODO: test for the static weight)
     if IMAGE == 1:
-      # pad with Invalid
-      def _invalid_pad_to(t, shape):
-        if all(p is None or p == s for p,s in zip(shape, t.shape)): return t
-        return Tensor(True, device=t.device).expand(t.shape).pad_to(shape).where(t.pad_to(shape), Invalid)
       # hacks for pitch alignment
       assert isinstance(ix, int) and isinstance(H, int)
       ALIGN = 64 // dtsz
-      x = _invalid_pad_to(x, (None, None, round_up(ix, ALIGN // math.gcd(groups * cin, ALIGN)), None))
-      w = _invalid_pad_to(w, (None, round_up(H, ALIGN // math.gcd(W * cin * 4, ALIGN))) + (None,) * (w.ndim - 2))
+      x = x.pad_to(None, None, round_up(ix, ALIGN // math.gcd(groups * cin, ALIGN)), None)
+      w = w.pad_to((None, round_up(H, ALIGN // math.gcd(W * cin * 4, ALIGN))) + (None,) * (w.ndim - 2))
 
       if FLOAT16: x, w = x.cast(dtypes.half).contiguous().cast(dtypes.float), w.cast(dtypes.half).contiguous().cast(dtypes.float)
       else: x, w = x.contiguous(), w.contiguous()

tinygrad/uop/__init__.py

@@ -91,7 +91,7 @@ class Ops(FastEnum):
   CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto()
 
   # buffer ops
-  BUFFERIZE = auto(); COPY = auto(); BUFFER = auto(); BUFFER_VIEW = auto(); MSELECT = auto(); MSTACK = auto(); CUSTOM_FUNCTION = auto()
+  BUFFERIZE = auto(); COPY = auto(); BUFFER = auto(); BUFFER_VIEW = auto(); MSELECT = auto(); MSTACK = auto(); ENCDEC = auto()
 
   # the core 6 movement ops! these only exist in the tensor graph
   RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); FLIP = auto()
@@ -103,6 +103,9 @@ class Ops(FastEnum):
   # expander ops
   UNROLL = auto(); CONTRACT = auto(); VCAT = auto(); PTRCAT = auto()
 
+  # CAT is a movement op (placed here to preserve enum ordering of existing ops)
+  CAT = auto()
+
 class GroupOp:
   Unary = {Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.SQRT, Ops.RECIPROCAL, Ops.NEG, Ops.TRUNC}
   Binary = {Ops.ADD, Ops.MUL, Ops.IDIV, Ops.MAX, Ops.MOD, Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ,

tinygrad/uop/divandmod.py

@@ -1,4 +1,4 @@
-import functools, itertools, math
+import functools
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import cdiv, cmod, CORRECT_DIVMOD_FOLDING, unwrap
@@ -12,8 +12,8 @@ def fold_divmod_general(d: UOp, correct_divmod_folding: bool) -> UOp|None:
   x_min, x_max, y_min, y_max = x.vmin, x.vmax, y.vmin, y.vmax
   assert isinstance(x_min, int) and isinstance(x_max, int) and isinstance(y_min, int) and isinstance(y_max, int)
   if y_min==y_max==0: raise ZeroDivisionError(f"{'Division' if d.op is Ops.IDIV else 'Mod'} by zero trying to rewrite {x.alu(d.op, y)}")
-  if y_min*y_max > 0 and (qv:=cdiv(x_min,y_min)) == cdiv(x_min,y_max) == cdiv(x_max,y_min) == cdiv(x_max,y_max):
-    return x - qv*y if d.op is Ops.MOD else d.const_like(qv)
+  if y_min*y_max > 0 and (q:=cdiv(x_min,y_min)) == cdiv(x_min,y_max) == cdiv(x_max,y_min) == cdiv(x_max,y_max):
+    return x - q*y if d.op is Ops.MOD else d.const_like(q)
 
   # split uops for the rest of the processing
   x_peeled, const = x.pop_const()
@@ -22,11 +22,11 @@ def fold_divmod_general(d: UOp, correct_divmod_folding: bool) -> UOp|None:
   # ** Constant Denominator Rules **
   # these rules strictly require y to be a scalar constant > 0
   if y.op is Ops.CONST and (c := y.arg) > 0:
-    # nested_div_mod: (x%(k*c))//c -> (x//c)%k, and (x%(k*c))%c -> x%c
-    if x.op is Ops.MOD and (k := x.src[1].divides(c)) is not None:
-      return x.src[0] // y % k if d.op is Ops.IDIV else x.src[0] % y
+    # canonicalize_mod_div: (x%(d*k))//d -> (x//d)%k, puts nested div/mod in div-first canonical form for recombine
+    if d.op is Ops.IDIV and x.op is Ops.MOD and x.src[1].op is Ops.CONST and x.vmin >= 0 and x.src[1].arg % c == 0:
+      return x.src[0] // y % x.ufix(x.src[1].arg // c)
 
-    # remove_nested_mod in sum: (a%4 + b)%2 -> (a+b)%2, requires non-negative sums
+    # remove_nested_mod: remove nested mod in case the inner mod is a multiple of the outer mod, example: (a%4 + b)%2 -> (a+b)%2
     if d.op is Ops.MOD and x.vmin >= 0:
       new_xs, changed = [], False
       for u in uops_no_const:
@@ -42,45 +42,40 @@ def fold_divmod_general(d: UOp, correct_divmod_folding: bool) -> UOp|None:
 
     # fold_binary_numerator: fold if expression has one non-constant term that takes on two values
     if len(terms)==1 and (v:=terms[0]).vmax-v.vmin == 1:
-      y1 = (cmod if d.op is Ops.MOD else cdiv)(factors[0]*v.vmin+const, c)
-      y2 = (cmod if d.op is Ops.MOD else cdiv)(factors[0]*v.vmax+const, c)
+      y1 = cmod(factors[0]*v.vmin+const, c) if d.op is Ops.MOD else cdiv(factors[0]*v.vmin+const, c)
+      y2 = cmod(factors[0]*v.vmax+const, c) if d.op is Ops.MOD else cdiv(factors[0]*v.vmax+const, c)
       return (y2-y1)*(v-v.vmin) + y1
 
     # fold_divmod_congruence: fold if a is congruent to an expression whose range is between 0 and c
     if not (x.vmin<0 and correct_divmod_folding):
-      # when f%c == c//2, abs(r) == abs(r-c) is a tie, try both signs since either may fit in one period
-      rem_choices = [(r, r-c) if (r:=f%c)*2 == c else (min(r, r-c, key=abs),) for f in factors]
-      for rems in itertools.product(*rem_choices):
-        if (rem:=sum(r*v for r,v in zip(rems,terms))+const%c).vmin//c==rem.vmax//c:
-          if d.op is Ops.MOD: return rem - rem.vmin//c*c
-          return sum((f-r)//c * v for f,r,v in zip(factors,rems,terms)) + const//c + rem.vmin//c
+      rems = [min((r:=f%c), r-c, key=abs) for f in factors]
+      if (rem:=sum(r*v for r,v in zip(rems,terms))+const%c).vmin//c==rem.vmax//c:
+        if d.op is Ops.MOD: return rem - rem.vmin//c*c
+        return sum((f-r)//c * v for f,r,v in zip(factors,rems,terms)) + (const-const%c+rem.vmin//c*c)//c
 
     # gcd_with_remainder: factor out common gcd from numerator
-    if x.vmin >= 0 and (g:=math.gcd(*factors, c)) > 1:
-      new_x = unwrap(x_peeled.divides(g)).simplify() + (const//g)%(c//g)
-      if new_x.vmin >= 0:
-        if d.op is Ops.MOD: return new_x % (c//g) * g + const%g
-        return new_x // (c//g) + const//c
-
-    # nest_by_factor: x//c -> (x//f)//(c//f), x%c -> (x//f%(c//f))*f + b where b=x%f
+    # Note: this rule uses uops_no_const to exclude the additive constant from the GCD calculation
     if x.vmin >= 0:
+      gcd = UOp.gcd(*uops_no_const, y).simplify()
+      if gcd.op is Ops.CONST and gcd.arg > 1:
+        new_x = unwrap(x_peeled.divide_exact(gcd)).simplify() + (const%c)//gcd.arg
+        if new_x.vmin >= 0:
+          ret = new_x.alu(d.op, x.ufix(c//gcd.arg))
+          return ret*gcd + const%gcd.arg if d.op is Ops.MOD else ret+const//c
+
+    # nest_div_by_factor: try nesting the div with each candidate factor and pick the simplest result
+    if d.op is Ops.IDIV and x.vmin >= 0:
+      # NOTE: this is recursive!
       results = []
       for div in {abs(f) for u, f in zip(uops_no_const, factors) if u.op not in (Ops.CONST, Ops.VCONST) and 1 < abs(f) < c and (c%f)==0}:
         if (newxs := fold_divmod_general(x//div, correct_divmod_folding)) is not None and newxs.vmin >= 0:
-          if d.op is Ops.IDIV:
-            results.append((len(newxs.backward_slice), newxs // (c // div)))
-          else:
-            b_parts = [f%div*t for f, t in zip(factors, terms) if f%div]
-            if const % div: b_parts.append(x.const_like(const % div))
-            b = UOp.sum(*b_parts) if b_parts else x.const_like(0)
-            if 0 <= b.vmin and b.vmax < div:
-              results.append((len((r:=(newxs % x.ufix(c//div))*div + b).backward_slice), r))
-      if results: return min(results, key=lambda r: r[0])[1]
+          results.append((len(newxs.backward_slice), newxs // (c // div)))
+      if results: return min(results)[1]
 
   # ** Variable Denominator / Fallback Rules **
   # These rules apply to variables OR constants that failed the checks above.
   # Reconstruct all uops including const for these checks.
-  all_uops = list(x.split_uop(Ops.ADD))
+  all_uops = uops_no_const + ([x.const_like(const)] if const != 0 else [])
 
   # divide_by_gcd: x//y -> (x//gcd)//(y//gcd)
   gcd = UOp.gcd(*all_uops, y).simplify()

tinygrad/uop/ops.py

@@ -207,7 +207,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   def _shape(self) -> tuple[sint, ...]|None:
     match self.op:
       # late ops don't have shape
-      case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.LOAD | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
+      case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.RANGE | Ops.LOAD | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
            Ops.VECTORIZE | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.CONTRACT | Ops.SINK | \
            Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY | Ops.INS:
         return None
@@ -218,20 +218,18 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
           return None
 
       case Ops.INDEX:
-        # non pointer index
-        if not isinstance(self.dtype, PtrDType):
-          idxs = flatten([d.shape for d in self.src[1:]])
-          return tuple(idxs) + self.src[0].shape[len(self.src)-1:]
+        # non pointer index doesn't have a shape
+        if not isinstance(self.dtype, PtrDType): return None
         # fully indexed doesn't have a shape. TODO: remove this
         if self.src[0]._shape is None or len(self.src[1:]) == len(self.src[0].shape): return None
         # pointer index
         return self.src[0].shape[len(self.src[1:]):]
 
       # some ops init the shape
-      case Ops.CONST | Ops.VCONST | Ops.DEFINE_VAR | Ops.BIND | Ops.RANGE: return ()
+      case Ops.CONST | Ops.VCONST | Ops.DEFINE_VAR | Ops.BIND: return ()
       case Ops.BUFFER: return (self.arg,)
       case Ops.BUFFER_VIEW: return (self.arg[0],)
-      case Ops.CUSTOM_FUNCTION: return None
+      case Ops.ENCDEC: return self.arg[0]
       case Ops.BUFFERIZE: return tuple([int(r.vmax+1) for r in self.src[1:]])
       case Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return (self.ptrdtype.size,)
       case Ops.PARAM:
@@ -248,9 +246,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
         inner_shape = self.src[0]._shape
         if inner_shape is None: return None
         # substitute internal PARAMs in the shape with corresponding args
-        ret = tuple(graph_rewrite(s, _pm_resolve_params, self.src[1:], walk=True) if isinstance(s, UOp) else s for s in inner_shape)
-        prepend = tuple([x.vmax+1 for x in self.src[0].ranges])
-        return prepend+ret
+        return tuple(graph_rewrite(s, _pm_resolve_params, self.src[1:], walk=True) if isinstance(s, UOp) else s for s in inner_shape)
 
       # TODO: disallow shape changing bitcast
       case Ops.BITCAST:
@@ -267,6 +263,14 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       # MULTI marker (axis info in PARAM sources) has no shape
       case Ops.MULTI if len(self.src) == 0: return None
 
+      case Ops.CAT:
+        shapes = [s.shape for s in self.src]
+        axis = self.arg
+        for s in shapes[1:]:
+          if len(s) != len(shapes[0]) or not all(a==b for i,(a,b) in enumerate(zip(s, shapes[0])) if i!=axis):
+            raise ValueError(f"CAT shape mismatch: {shapes}")
+        return tuple(ssimplify(sum(s[i] for s in shapes)) if i==axis else shapes[0][i] for i in range(len(shapes[0])))
+
     # movement ops change the shape
     # NOTE: ssimplify is required because the shape needs to be canonical for broadcasting and same shape checking
     if self.op in GroupOp.Movement.union({Ops.MULTI, Ops.REDUCE_AXIS, Ops.WMMA}):
@@ -352,11 +356,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   @recursive_property
   def _ranges(self) -> dict[UOp, None]:
     ret: dict[UOp, None] = {}
-    if self.op is Ops.CALL:
-      # ranges do not flow through calls
-      for s in self.src[1:]: ret.update(s.ranges)
-    else:
-      for s in self.src: ret.update(s.ranges)
+    for s in self.src: ret.update(s.ranges)
     for er in self.ended_ranges:
       if er.op is Ops.RANGE:
         # if it's a single RANGE, we don't flow through it.
@@ -422,7 +422,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     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):
     idx = argfix(idx)
-    #assert len(idx) == len(self.shape), f"__getitem__ shape mismatch, indexing {self.shape} with {len(idx)} args"
+    assert len(idx) == len(self.shape), f"__getitem__ shape mismatch, indexing {self.shape} with {len(idx)} args"
     if len(slice_idx:=[i for i,x in enumerate(idx) if isinstance(x, slice)]):
       perm = self.permute(tuple([i for i in range(self.ndim) if i not in slice_idx] + slice_idx))
       return perm.index(*[UOp.const(dtypes.index, x) if isinstance(x, int) else x for x in idx if not isinstance(x, slice)], ptr=True)
@@ -575,6 +575,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   def mstack(self, *srcs: UOp) -> UOp: return UOp(Ops.MSTACK, self.dtype, (self,)+srcs)
   @property
   def metadata(self) -> tuple[Metadata, ...]|None: return all_metadata.get(self, None)
+  def encdec(self, *src, arg=None): return UOp(Ops.ENCDEC, self.dtype, src=(self,)+src, arg=arg)
 
   # *** uop movement ops ***
 
@@ -910,7 +911,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     return p
 
   def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=(), name:str|None=None, precompile:bool=False) -> UOp:
-    # ranges don't leak through calls, they end!
+    # TODO: reenable this after ENCDEC is fixed
     #assert len(self.ranges) == 0, f"ranges {self.ranges} are leaking out of the call in {self.pyrender()}"
     return UOp(Ops.CALL, self.dtype, (self,)+srcs, CallInfo(grad_fxn, metadata, name, precompile))
   def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]:
@@ -943,7 +944,7 @@ class CallInfo:
 def should_resolve_call(c:UOp) -> bool:
   # don't resolve real kernel calls, sink or program
   if c.src[0].op is Ops.SINK and isinstance(c.src[0].arg, KernelInfo): return False
-  if c.src[0].op in {Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.CUSTOM_FUNCTION}: return False
+  if c.src[0].op in {Ops.PROGRAM, Ops.LINEAR, Ops.COPY}: return False
   if c.arg.precompile: return False
   return True
 
@@ -1471,7 +1472,6 @@ renderer = PatternMatcher([
   (UPat(Ops.PARAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.NOOP, name="x"))), lambda x: x.arg),
   (UPat((Ops.SPECIAL), name="x"), lambda x: x.arg),
   (UPat(Ops.RANGE, name="x"), lambda x: f"r{range_str(x)}"),
-  (UPat(Ops.PARAM, name="x"), lambda x: f"p{x.arg}"),
   (UPat((Ops.CONST, Ops.VCONST), name="x"), lambda x: str(x.arg)),
   (UPat(Ops.UNROLL, name="x"), lambda ctx,x,u: f"UNROLL({ctx[x.src[0]]}, {u.arg})"),
   (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({str(x.dtype)[7:]})({ctx[x.src[0]]})"),
@@ -1521,7 +1521,7 @@ pm_pyrender_extra = PatternMatcher([
   (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"), lambda x,u,d:
     f"UOp.new_buffer({repr(d.arg)}, {x.size}, {x.dtype}, {u.arg})"),
   (UPat(Ops.COPY, src=(UPat(name="x"), UPat(Ops.DEVICE, name="d"))), lambda ctx,x,d: f"{ctx[x]}.copy_to_device({repr(d.arg)})"),
-  (UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda ctx,x: f"UOp(Ops.CUSTOM_FUNCTION, {x.dtype}, src={srcs(ctx, x.src)}, arg={x.arg!r})"),
+  (UPat(Ops.ENCDEC, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.encdec({''.join([str(ctx[s])+', ' for s in x.src[1:]])}arg={x.arg!r})"),
   (UPat(Ops.REDUCE_AXIS, name="r"), lambda ctx,r: f"{ctx[r.src[0]]}.r({r.arg[0]}, {r.arg[1]})"),
   # NOTE: range has srcs sometimes after control flow
   (UPat(Ops.RANGE, src=(UPat(Ops.CONST, name="c"),), allow_any_len=True, name="x"), lambda ctx,x,c:
@@ -1529,7 +1529,7 @@ pm_pyrender_extra = PatternMatcher([
       (f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+(', dtype='+str(x.dtype) if x.dtype is not dtypes.index else '')+")"),
   # TODO: index shouldn't mismatch dtype
   (UPat(Ops.INDEX, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda ctx,x:
-   f"{ctx[x.src[0]]}.index({ctx[x.src[1]]}, "+''.join([f"{ctx[xx]}, " for xx in x.src[2:]])+
+   f"{ctx[x.src[0]]}.index({ctx[x.src[1]]}, "+(f"{ctx[x.src[2]]}, " if len(x.src) > 2 else "")+
     (f"dtype={x.dtype})" if x.src[0].dtype != x.dtype else "ptr=True)") if x.src[0].dtype.base != x.dtype else None),
   # TODO: movement ops simplify stuff, this can break SPEC=2
   #(UPat(GroupOp.Movement, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({render_marg(ctx,x)})"),

tinygrad/uop/spec.py

@@ -69,6 +69,7 @@ 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)),
+  (UPat(Ops.CAT, name="mv"), lambda mv: isinstance(mv.arg, int) and len(mv.src) >= 1),
 
   # inputs to movement ops
   (UPat((Ops.VECTORIZE, Ops.VCONST), dtype=dtypes.index), lambda: True),
@@ -84,7 +85,7 @@ _tensor_spec = PatternMatcher([
   (UPat(Ops.LUNIQUE, dtypes.void, ()), lambda: True),
   (UPat(Ops.DEVICE, dtypes.void, (), name="d"), lambda d:
    isinstance(d.arg, str) or (isinstance(d.arg, tuple) and all(isinstance(s, str) for s in d.arg))),
-  (UPat(Ops.BUFFER, src=(UPat((Ops.LUNIQUE, Ops.UNIQUE, Ops.NOOP)), UPat(Ops.DEVICE)), name="buf"),
+  (UPat(Ops.BUFFER, src=(UPat((Ops.LUNIQUE, Ops.UNIQUE)), UPat(Ops.DEVICE)), name="buf"),
    lambda buf: isinstance(buf.arg, int) and isinstance(buf.dtype, (DType, ImageDType))),
 
   # BUFFER_VIEW on BUFFER is allowed if BUFFER is
@@ -120,10 +121,11 @@ _tensor_spec = PatternMatcher([
   (UPat(Ops.CONTIGUOUS, name="root", src=(UPat.var("x"),), allow_any_len=True, arg=None),
    lambda root,x: root.dtype == x.dtype and all(u.op is Ops.RANGE for u in root.src[1:])),
 
-  # COPY/ALLREDUCE/MULTI
+  # COPY/ALLREDUCE/MULTI/ENCDEC
   (UPat(Ops.COPY, name="copy", src=(UPat.var("x"), UPat(Ops.DEVICE)), arg=None), lambda copy,x: copy.dtype == x.dtype),
   (UPat(Ops.ALLREDUCE, name="red", src=(UPat.var("x"), UPat(Ops.DEVICE))), lambda red,x: red.dtype == x.dtype and isinstance(red.arg, Ops)),
   (UPat(Ops.MULTI, name="multi"), lambda multi: all(x.dtype == multi.dtype for x in multi.src) and isinstance(multi.arg, int)),
+  (UPat(Ops.ENCDEC, name="x"), lambda x: len(x.src) >= 2), # state + inbuffer
 
   # REDUCE_AXIS is the reduce in the tensor graph
   (UPat(Ops.REDUCE_AXIS, name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) >= 2 and x.arg[0] in {Ops.ADD, Ops.MUL, Ops.MAX}),
@@ -131,10 +133,9 @@ _tensor_spec = PatternMatcher([
   # AFTER if things were kernelized
   (UPat(Ops.AFTER, src=(UPat((Ops.BUFFER, Ops.AFTER)),), allow_any_len=True), lambda: True),
 
-  # allow CALL/PARAM/CUSTOM_FUNCTION
+  # allow CALL/PARAM
   (UPat(Ops.CALL, src=(UPat(name="f"),), name="c", allow_any_len=True), lambda c,f: c.dtype == f.dtype),
   (UPat(Ops.PARAM), lambda: True),
-  (UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda x: isinstance(x.arg, str)),
 
   # ** for custom kernels **
 
@@ -267,6 +268,8 @@ full_spec = PatternMatcher([
   # linearizer: outputs + intermediate KERNELs
   (UPat(Ops.CALL, dtype=dtypes.void), lambda: True),
 
+  # Invalid must have type Index
+  (UPat(Ops.CONST, arg=Invalid, name="x"), lambda x: x.dtype.scalar() == dtypes.index),
   # where on index in rhs position is fine
   (UPat(Ops.WHERE, dtype=dtypes.index, src=(UPat(dtype=dtypes.bool), UPat(), UPat(dtype=dtypes.index))), lambda: True),
   # allow index dtype on a restricted set of UOps

tinygrad/uop/symbolic.py

@@ -25,48 +25,16 @@ def fold_bitcast(root:UOp, c:UOp) -> UOp|None:
 invalid_pat = UPat(Ops.CONST, arg=Invalid, name="i")
 invalid_gate = UPat.var("cond").where(UPat.var("x"), invalid_pat)
 
-def fold_add_divmod_recombine(x:UOp) -> UOp|None:
-  terms = list(x.split_uop(Ops.ADD))
-  for i,u in enumerate(terms):
-    if u.op is Ops.MOD and u.src[1].op is Ops.CONST: base, div, mul = u.src[0], u.src[1].arg, 1
-    elif u.op is Ops.MUL and u.src[1].op is Ops.CONST and (m:=u.src[0]).op is Ops.MOD and m.src[1].op is Ops.CONST:
-      base, div, mul = m.src[0], m.src[1].arg, u.src[1].arg
-    else: continue
-    for j,v in enumerate(terms):
-      if i == j: continue
-      if v.op is not Ops.MUL or v.src[1].op is not Ops.CONST or v.src[1].arg != div*mul: continue
-      q, exact = v.src[0], False
-      # (base%div)*mul + (base//div)*(div*mul) -> base*mul
-      if q.op is Ops.IDIV and q.src[1].op is Ops.CONST and q.src[1].arg == div: exact = q.src[0] is base
-      # ((base//d)%div)*mul + (base//(d*div))*(div*mul) -> (base//d)*mul
-      if not exact and base.op is Ops.IDIV and base.src[1].op is Ops.CONST:
-        exact = q.op is Ops.IDIV and q.src[1].op is Ops.CONST and q.src[0] is base.src[0] and q.src[1].arg == base.src[1].arg*div
-      if exact: return functools.reduce(operator.add, (t for k,t in enumerate(terms) if k not in (i,j)), base*mul)
-      # ((base//div)%d)*div + base%div -> base%(div*d)
-      if mul == 1 and div > 0 and q.op is Ops.MOD and q.src[1].op is Ops.CONST and (d:=q.src[1].arg) > 0 and q.src[0].op is Ops.IDIV:
-        if q.src[0].src[0] is base and q.src[0].src[1].op is Ops.CONST and q.src[0].src[1].arg == div:
-          return functools.reduce(operator.add, (t for k,t in enumerate(terms) if k not in (i,j)), base % (div*d))
-  return None
-
 # this needs to be before symbolic so that 0*something_that_might_be_invalid doesnt become 0
 propagate_invalid = PatternMatcher([
   # propagate invalid, push it past children
-  (invalid_gate.cast(name="cast"), lambda i,x,cond,cast: x.cast(cast.dtype) if i.dtype is dtypes.index else None),
-  (UPat(GroupOp.Unary, src=(invalid_gate,), name="alu"), lambda cond,x,alu,i: cond.where(x.alu(alu.op), i)),
-  (UPat(GroupOp.Binary-GroupOp.Comparison, src=(invalid_gate, UPat.var("y")), name="alu"), lambda cond,x,y,alu,i: cond.where(x.alu(alu.op,y), i)),
-  (UPat(GroupOp.Binary-GroupOp.Comparison, src=(UPat.var("y"), invalid_gate), name="alu"), lambda cond,x,y,alu,i: cond.where(y.alu(alu.op,x), i)),
+  (invalid_gate.cast(name="cast"), lambda i,x,cond,cast: x.cast(cast.dtype)),
+  *((invalid_gate.alu(op, UPat.var("y")).named("alu"), lambda cond,x,y,alu,i: cond.where(x.alu(alu.op,y), i))
+    for op in GroupOp.Binary-GroupOp.Comparison),
   # TODO: when can this happen? and is it always safe to just drop invalid?
-  (UPat(GroupOp.Comparison, src=(invalid_gate, UPat.var("y")), name="alu"), lambda cond,x,y,alu,i:
-     x.alu(alu.op,y) if i.dtype is dtypes.index else cond.where(x.alu(alu.op,y), i.cast(dtypes.bool))),
-  (UPat(GroupOp.Comparison, src=(UPat.var("y"), invalid_gate), name="alu"), lambda cond,x,y,alu,i:
-     y.alu(alu.op,x) if i.dtype is dtypes.index else cond.where(y.alu(alu.op,x), i.cast(dtypes.bool))),
+  *((invalid_gate.alu(op, UPat.var("y")).named("alu"), lambda cond,x,y,alu,i: x.alu(alu.op,y)) for op in GroupOp.Comparison),
   # alu with invalid -> invalid
-  (UPat(GroupOp.Unary, src=(invalid_pat,)), lambda i: i),
-  (UPat(GroupOp.Binary-GroupOp.Comparison, src=[invalid_pat, UPat()]), lambda i: i),
-  # normalize where(cond, Invalid, val) -> where(~cond, val, Invalid)
-  (UPat.var("cond").where(invalid_pat, UPat.var("val")), lambda cond, i, val: cond.logical_not().where(val, i) if val.arg != Invalid else i),
-  (UPat(Ops.BITCAST, src=(invalid_pat,), name="bc"), lambda bc,i: i.cast(bc.dtype)),
-  (UPat(Ops.BITCAST, src=(invalid_gate,), name="bc"), lambda bc,cond,x,i: cond.where(x.bitcast(bc.dtype), i.bitcast(bc.dtype))),
+  *((invalid_pat.alu(op, UPat(dtype=dtypes.index)), lambda i: i) for op in GroupOp.Binary-GroupOp.Comparison),
 ])
 
 symbolic_simple = propagate_invalid + PatternMatcher([
@@ -78,8 +46,24 @@ symbolic_simple = propagate_invalid + PatternMatcher([
   (UPat.var("x") // 1, lambda x: x),   # x//1 -> x
   (UPat.var("x") // -1, lambda x: -x), # x//-1 -> -x
   ((UPat.var() % UPat.var("y")).named("base") % UPat.var("y"), lambda base,y: base),  # (x%y)%y = -> x%y (rewritten with base for speed)
-  # variations of (x%c)+(x//c)*c = x
-  (UPat(Ops.ADD, dtype=dtypes.index, name="x"), fold_add_divmod_recombine),
+  # variations of (x%c)+(x//c)*c = x   TODO: add sorting to remove some variations
+  (UPat.var("x")%UPat.cvar("c")+(UPat.var("x")//UPat.cvar("c"))*UPat.cvar("c"), lambda x,c: x), # (x%c)+(x//c)*c = x
+  ((UPat.var("x")//UPat.cvar("a"))%UPat.cvar("c")+(UPat.var("x")//UPat.cvar("b"))*UPat.cvar("c"),
+    lambda x,a,b,c: x//a if a.arg*c.arg==b.arg else None), # ((x//a)%c)+(x//a*c)*c = x//a. Note if a = 1 it degenerates to the one above
+  ((UPat.var("x")//UPat.cvar("a"))%UPat.cvar("c1")*UPat.cvar("c2")+(UPat.var("x")//UPat.cvar("b"))*UPat.cvar("c3"),
+    lambda x,a,b,c1,c2,c3: x//a*c2 if c1.arg>0 and a.arg*c1.arg==b.arg and c1.arg*c2.arg==c3.arg else None),
+  ((UPat.var("x")//UPat.cvar("c1"))*UPat.cvar("c3")+UPat.var("x")%UPat.cvar("c1")*UPat.cvar("c2"),
+    lambda x,c1,c2,c3: x*c2 if c1.arg*c2.arg==c3.arg else None), # (x%c1)*c2+(x//c1)*c3 = x*c2 if c1*c2==c3
+  ((UPat.var("y")+(UPat.var("x")//UPat.cvar("c"))*UPat.cvar("c"))+UPat.var("x")%UPat.cvar("c"), lambda y,x,c: y+x),
+  ((UPat.var("y")+UPat.var("x")%UPat.cvar("c"))+(UPat.var("x")//UPat.cvar("c"))*UPat.cvar("c"), lambda y,x,c: y+x),
+  ((UPat.var("y")+(UPat.var("x")//UPat.cvar("c1"))*UPat.cvar("c3"))+UPat.var("x")%UPat.cvar("c1")*UPat.cvar("c2"),
+    lambda y,x,c1,c2,c3: y+x*c2 if c1.arg*c2.arg==c3.arg else None),
+  ((UPat.var("y")+UPat.var("x")%UPat.cvar("c1")*UPat.cvar("c2"))+(UPat.var("x")//UPat.cvar("c1"))*UPat.cvar("c3"),
+    lambda y,x,c1,c2,c3: y+x*c2 if c1.arg*c2.arg==c3.arg else None),
+  ((UPat.var("y")+(UPat.var("x")//UPat.cvar("a"))%UPat.cvar("c1")*UPat.cvar("c2"))+(UPat.var("x")//UPat.cvar("b"))*UPat.cvar("c3"),
+    lambda y,x,a,b,c1,c2,c3: y+x//a*c2 if c1.arg>0 and a.arg*c1.arg==b.arg and c1.arg*c2.arg==c3.arg else None),
+  ((UPat.var("y")+(UPat.var("x")//UPat.cvar("b"))*UPat.cvar("c3"))+(UPat.var("x")//UPat.cvar("a"))%UPat.cvar("c1")*UPat.cvar("c2"),
+    lambda y,x,a,b,c1,c2,c3: y+x//a*c2 if c1.arg>0 and a.arg*c1.arg==b.arg and c1.arg*c2.arg==c3.arg else None),
   (UPat.var("x", dtype=dtypes.bool) & UPat.cvar("c", vec=False), lambda x,c: x if c.arg else c),
   (UPat.var("x", dtype=dtypes.bool) | UPat.cvar("c", vec=False), lambda x,c: c if c.arg else x),
   (UPat(GroupOp.Idempotent, src=(UPat.var("x"), UPat.var("x"))), lambda x: x),
@@ -138,12 +122,6 @@ symbolic_simple = propagate_invalid + PatternMatcher([
   (UPat.cvar("gate", vec=False).where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1),
   # a.where(b.where(c, d), d) -> (a & b).where(c, d)
   (UPat.var("a").where(UPat.var("b").where(UPat.var("c"), UPat.var("d")), UPat.var("d")), lambda a,b,c,d: (a&b).where(c,d)),
-  # where over invalid -> invalid
-  (invalid_gate.where(UPat.var("a"), UPat.var("b")), lambda a,b,cond,x,i: i.cast(a.dtype)),
-  (invalid_pat.where(UPat.var("a"), UPat.var("b")), lambda a,b,i: i.cast(a.dtype)),
-  # reduce with invalid -> invalid
-  (UPat(Ops.REDUCE, src=(invalid_gate,), allow_any_len=True, name="r"), lambda r,cond,x,i: i.cast(r.dtype)),
-  (UPat(Ops.REDUCE, src=(invalid_pat,), allow_any_len=True, name="r"), lambda r,i: i.cast(r.dtype)),
 ])
 
 # ******** phase 2 builds on phase 1, it includes the old "symbolic", rules that match deeper ********
@@ -364,8 +342,8 @@ def reduce_mul_chain(r:UOp) -> UOp|None:
   if r.dtype != r.src[0].dtype: return None
   inside, outside = [], []
   for m in r.src[0].split_uop(Ops.MUL):
-    m_parents = m.backward_slice
-    if m not in r.src[1:] and all(r not in m_parents for r in r.src[1:]) and (r.arg != Ops.MAX or m.vmin >= 0): outside.append(m)
+    m_parents = m.toposort()
+    if all(r not in m_parents for r in r.src[1:]) and (r.arg != Ops.MAX or m.vmin >= 0): outside.append(m)
     else: inside.append(m)
   if len(outside) == 0: return None
   return r.replace(src=(prod(inside) if len(inside) else r.src[0].const_like(1),)+r.src[1:])*prod(outside)
@@ -434,10 +412,6 @@ sym = symbolic+pm_simplify_valid+PatternMatcher([
   # fold gated LOAD/STORE
   (UPat((Ops.LOAD, Ops.STORE), src=(UPat().index(UPat.const(dtypes.index, Invalid)).or_casted(),), allow_any_len=True, name="x"),
     lambda x: UOp(Ops.NOOP) if x.op is Ops.STORE else x.const_like(0)), # invalid store does nothing. invalid load produces 0
-  (UPat(Ops.STORE, src=(UPat(), invalid_pat), allow_any_len=True), lambda i: UOp(Ops.NOOP)),
-  # store of where with invalid -> gated store
-  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, name="index"), UPat.var("cond").where(UPat.var("val"), invalid_pat)), allow_any_len=True, name="store"),
-   lambda index, cond, val, store, i: UOp.store(index.src[0].index(cond.where(index.src[1], UOp.invalid())), val, *store.src[2:])),
   ((UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()),  # 1/(x^c) -> (1/x)^c
   ((UPat.var("x") * UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()*x.reciprocal()),
   ((UPat.var("x") * UPat.cvar("c")).reciprocal(), lambda x,c: x.reciprocal()*c.reciprocal()), # 1/(x*c) -> (1/c)*(1/x)

tinygrad/viz/js/index.js

@@ -236,7 +236,7 @@ const drawLine = (ctx, x, y, opts) => {
 function tabulate(rows) {
   const root = d3.create("div").style("display", "grid").style("grid-template-columns", `${Math.max(...rows.map(x => x[0].length), 0)}ch 1fr`).style("gap", "0.2em").style("white-space", "nowrap");
   for (const [k,v] of rows) { root.append("div").text(k); root.append("div").node().append(v); }
-  return root.node();
+  return root;
 }
 
 var data, focusedDevice, focusedShape, formatTime, canvasZoom, zoomLevel = d3.zoomIdentity;
@@ -268,7 +268,7 @@ function setFocus(key) {
   const html = d3.select(".info").html("");
   if (eventType === EventTypes.EXEC) {
     const [n, _, ...rest] = e.arg.tooltipText.split("\n");
-    html.append(() => tabulate([["Name", colored(e.arg.label)], ["Duration", formatTime(e.width)], ["Start Time", formatTime(e.x)]]));
+    html.append(() => tabulate([["Name", d3.create("p").html(n).node()], ["Duration", formatTime(e.width)], ["Start Time", formatTime(e.x)]]).node());
     let group = html.append("div").classed("args", true);
     for (const r of rest) group.append("p").text(r);
     group = html.append("div").classed("args", true);
@@ -290,7 +290,7 @@ function setFocus(key) {
     const [dtype, sz, nbytes, dur] = e.arg.tooltipText.split("\n");
     const rows = [["DType", dtype], ["Len", sz], ["Size", nbytes], ["Lifetime", dur]];
     if (e.arg.users != null) rows.push(["Users", e.arg.users.length]);
-    html.append(() => tabulate(rows));
+    html.append(() => tabulate(rows).node());
     const kernels = html.append("div").classed("args", true);
     for (let u=0; u<e.arg.users?.length; u++) {
       const { repr, num, mode, shape } = e.arg.users[u];
@@ -302,8 +302,7 @@ function setFocus(key) {
   }
   // instructions list renderer
   let instList = document.getElementById("insts");
-  if (data.pcToShape.size == 0) return d3.select(instList?.parentElement).html("");
-  if (instList == null) {
+  if (data.pcToShape.size > 0 && instList == null) {
     let contents = "", i = 0;
     for (const [k, v] of data.pcToShape) {
       contents += `<div class="line" data-k="${k}"><span class="left" id="inst-${k}"><span class="n">${i++}</span><span class="wave">${v.wave}</span>
@@ -315,7 +314,7 @@ function setFocus(key) {
   }
   d3.select(instList).selectAll("span").classed("highlight", false);
   const instLine = document.getElementById(`inst-${key}`); instLine?.classList.add("highlight");
-  if (instLine != null) {
+  if (instLine != null && instList != null) {
     const r = rect(instLine), c = rect(instList);
     if (Math.max(c.top-r.bottom, r.top-c.bottom)>=-30) instLine.scrollIntoView({ block:"center" });
   }
@@ -323,10 +322,10 @@ function setFocus(key) {
 
 const EventTypes = { EXEC:0, BUF:1 };
 
-async function renderProfiler(path, opts) {
+async function renderProfiler(path, unit, opts) {
   displaySelection("#profiler");
   // support non realtime x axis units
-  formatTime = opts.unit === "ms" ? formatMicroseconds : formatCycles;
+  formatTime = unit === "realtime" ? formatMicroseconds : formatCycles;
   if (data?.path !== path) { data = {tracks:new Map(), axes:{}, path, first:null, pcToShape:new Map()}; focusedDevice = null; focusedShape = null; }
   setFocus(focusedShape);
   // layout once!
@@ -378,12 +377,16 @@ async function renderProfiler(path, opts) {
       for (let j=0; j<eventsLen; j++) {
         const e = {name:strings[u32()], ref:optional(u32()), key:optional(u32()), st:u32(), dur:f32(), info:strings[u32()] || null};
         // find a free level to put the event
-        let depth = levels.findIndex(levelEt => e.st >= levelEt);
-        const et = e.st+Math.trunc(e.dur);
-        if (depth === -1) {
-          depth = levels.length;
-          levels.push(et);
-        } else levels[depth] = et;
+        let depth = 0;
+        if (opts.levelKey != null) { depth = opts.levelKey(e); levels[depth] = 0; }
+        else {
+          depth = levels.findIndex(levelEt => e.st >= levelEt);
+          const et = e.st+Math.trunc(e.dur);
+          if (depth === -1) {
+            depth = levels.length;
+            levels.push(et);
+          } else levels[depth] = et;
+        }
         if (depth === 0 || opts.colorByName) colorKey = e.name.split(" ")[0];
         if (!colorMap.has(colorKey)) {
           const color = typeof colors === "function" ? colors(colorKey)
@@ -414,10 +417,11 @@ async function renderProfiler(path, opts) {
         }
         // tiny device events go straight to the rewrite rule
         const key = k.startsWith("TINY") ? null : `${k}-${j}`;
+        const labelHTML = label.map(l=>`<span style="color:${l.color}">${l.st}</span>`).join("");
         let info = e.info != null ? "\n"+e.info : "", trace = null
         if (info.startsWith("\nPC:")) { data.pcToShape.set(key, {wave:dnum, pc:parseInt(e.info.split(":")[1]), st:e.st}); info = ""; }
         if (info.startsWith("\nTB:")) { trace = info; info = ""; }
-        const arg = { tooltipText:" N:"+shapes.length+"\n"+formatTime(e.dur)+info, label, trace, bufs:[], key, ctx:shapeRef?.ctx, step:shapeRef?.step };
+        const arg = { tooltipText:labelHTML+" N:"+shapes.length+"\n"+formatTime(e.dur)+info, trace, bufs:[], key, ctx:shapeRef?.ctx, step:shapeRef?.step };
         if (e.key != null) shapeMap.set(e.key, key);
         // offset y by depth
         shapes.push({x:e.st, y:levelHeight*depth, width:e.dur, height:levelHeight, arg, label:opts.hideLabels ? null : label, fillColor });
@@ -593,7 +597,7 @@ async function renderProfiler(path, opts) {
       const labelX = x+ctx.lineWidth+2;
       if (labelX <= lastLabelEnd) continue;
 
-      const label = formatTime(tick, et-st <= 1e3);
+      const label = formatTime(tick, et-st <= 1e3 ? true : false);
       ctx.textBaseline = "top";
       ctx.fillText(label, labelX, tickSize);
       lastLabelEnd = labelX + ctx.measureText(label).width + 4;
@@ -670,12 +674,12 @@ async function renderProfiler(path, opts) {
 
   canvas.addEventListener("mousemove", e => {
     const foundRect = findRectAtPosition(e.clientX, e.clientY);
-    const tooltip = document.getElementById("tooltip");
     if (foundRect?.tooltipText != null) {
-      tooltip.replaceChildren(colored(foundRect.label||[]), document.createTextNode(foundRect.tooltipText));
+      const tooltip = document.getElementById("tooltip");
       tooltip.style.display = "block";
       tooltip.style.left = (e.pageX+10)+"px";
       tooltip.style.top = (e.pageY)+"px";
+      tooltip.innerHTML = foundRect.tooltipText;
     } else tooltip.style.display = "none";
   });
   canvas.addEventListener("mouseleave", () => document.getElementById("tooltip").style.display = "none");
@@ -866,6 +870,8 @@ async function main() {
   // ** center graph
   const { currentCtx, currentStep, currentRewrite, expandSteps } = state;
   if (currentCtx == -1) return;
+  // always have a new sidebar when view changes
+  metadata.innerHTML = "";
   const ctx = ctxs[currentCtx];
   const step = ctx.steps[currentStep];
   const ckey = step?.query;
@@ -876,7 +882,7 @@ async function main() {
     if (url.pathname+url.search !== ckey) e.close();
     else if (e.readyState === EventSource.OPEN) activeSrc = e;
   }
-  if (ctx.name === "Profiler") return renderProfiler("/get_profile", {unit:"ms", width:"132px"});
+  if (ctx.name === "Profiler") return renderProfiler("/get_profile", "realtime", { width:"132px" });
   if (workerUrl == null) await initWorker();
   if (ckey in cache) {
     ret = cache[ckey];
@@ -894,12 +900,15 @@ async function main() {
     }
     // timeline with cycles on the x axis
     if (ret instanceof ArrayBuffer) {
-      const pkts = step.name.includes("PKTS");
-      return renderProfiler(ckey, {unit:"clk", heightScale:0.5, hideLabels:true, colorByName:pkts});
+      opts = {heightScale:0.5, hideLabels:true, levelKey:step.name.includes("PKTS") ? (e) => parseInt(e.name.split(" ")[1].split(":")[1]) : null, colorByName:ckey.includes("pkts")};
+      return renderProfiler(ckey, "clk", opts);
     }
-    metadata.replaceChildren(...((ret.metadata ?? []).map((m) => {
-      return tabulate(m.map((e) => [e.label.trim(), typeof e.value === "string" ? e.value : formatUnit(e.value)]));
-    })));
+    ret.metadata?.forEach(m => {
+      if (Array.isArray(m)) return metadata.appendChild(tabulate(m.map(({ label, value }) => {
+        return [label.trim(), typeof value === "string" ? value : formatUnit(value)];
+      })).node());
+      metadata.appendChild(codeBlock(m.src)).classList.add("full-height")
+    });
     // graph render
     if (ret.data != null) {
       metadata.prepend(showGraph.label);

tinygrad/viz/serve.py

@@ -47,8 +47,8 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
                **{x:"#f2cb91" for x in {Ops.DEFINE_LOCAL, Ops.DEFINE_REG}}, 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.INS: "#eec4ff",
-               **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
-               Ops.BUFFER_VIEW: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",
+               **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80", Ops.CAT:"#C1FFD7",
+               Ops.BUFFER_VIEW: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.ENCDEC: "#bf71b6",
                Ops.CALL: "#00B7C8", Ops.PARAM: "#14686F", Ops.SOURCE: "#c0c0c0", Ops.LINEAR: "#7DF4FF", Ops.BINARY: "#404040",
                Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",
                Ops.BUFFERIZE: "#FF991C", Ops.REWRITE_ERROR: "#ff2e2e", Ops.AFTER: "#8A7866", Ops.END: "#524C46"}
@@ -71,9 +71,10 @@ def get_rewrites(t:RewriteTrace) -> list[dict]:
     steps = [create_step(s.name, ("/graph-rewrites", i, j), loc=s.loc, match_count=len(s.matches), code_line=printable(s.loc),
                          trace=k.tb if j==0 else None, depth=s.depth) for j,s in enumerate(v)]
     if (p:=get_prg_uop(i)) is not None:
-      steps.append(create_step("View UOp List", ("/uops", i, len(steps))))
-      steps.append(create_step("View Source", ("/code", i, len(steps)), p.src[3].arg))
-      steps.append(create_step("View Disassembly", ("/asm", i, len(steps)), (k.ret, p.src[4].arg)))
+      _, __, lin, src, binary = p.src
+      steps.append(create_step("View UOp List", ("/uops", i, len(steps)), lin.src))
+      steps.append(create_step("View Source", ("/code", i, len(steps)), src.arg))
+      steps.append(create_step("View Disassembly", ("/asm", i, len(steps)), (k.ret, binary.arg)))
     for key in k.keys: ref_map[key] = i
     ret.append({"name":k.display_name, "steps":steps})
   return ret
@@ -149,10 +150,9 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
   return graph
 
 @functools.cache
-def _reconstruct(a:int, depth:int|None=None):
+def _reconstruct(a:int):
   op, dtype, src, arg, *rest = trace.uop_fields[a]
-  if depth is not None and depth <= 0: return UOp(op, dtype, (), arg, *rest)
-  return UOp(op, dtype, tuple(_reconstruct(s, None if depth is None else depth-1) for s in src), arg, *rest)
+  return UOp(op, dtype, tuple(_reconstruct(s) for s in src), arg, *rest)
 
 def get_full_rewrite(ctx:TrackedGraphRewrite) -> Generator[GraphRewriteDetails, None, None]:
   next_sink = _reconstruct(ctx.sink)
@@ -169,7 +169,7 @@ def get_full_rewrite(ctx:TrackedGraphRewrite) -> Generator[GraphRewriteDetails,
 
 def get_prg_uop(i:int) -> UOp|None:
   s = next((s for s in trace.rewrites[i] if s.name == "View Program"), None)
-  return _reconstruct(s.sink, depth=1) if s is not None else None
+  return _reconstruct(s.sink) if s is not None else None
 
 # encoder helpers
 
@@ -227,7 +227,7 @@ def timeline_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:
       elif isinstance(e.name.ret, int):
         membw = e.name.ret / (dur * 1e-6)
         fmt.append(f"{membw*1e-9:.0f} GB/s" if membw < 1e13 else f"{membw*1e-12:.0f} TB/s")
-      elif e.name.tb: fmt.append("TB:"+json.dumps(e.name.tb))
+      if e.name.tb: fmt.append("TB:"+json.dumps(e.name.tb))
     events.append(struct.pack("<IIIIfI", enum_str(name, scache), option(ref), option(key), rel_ts(st,start_ts), dur, enum_str("\n".join(fmt),scache)))
   return struct.pack("<BI", 0, len(events))+b"".join(events) if events else None
 
@@ -338,17 +338,19 @@ def sqtt_timeline(data:bytes, lib:bytes, target:str) -> list[ProfileEvent]:
   from tinygrad.renderer.amd.sqtt import map_insts, InstructionInfo, PacketType, INST, InstOp, VALUINST, IMMEDIATE, IMMEDIATE_MASK, VMEMEXEC, ALUEXEC
   from tinygrad.renderer.amd.sqtt import INST_RDNA4, InstOpRDNA4
   ret:list[ProfileEvent] = []
-  row_ends:dict[str, Decimal] = {}
-  def add(name:str, p:PacketType, width=1, op:str|None=None, wave:int|None=None, info:InstructionInfo|None=None) -> None:
-    row = f"WAVE:{wave}" if (wave:=getattr(p, "wave", wave)) is not None else f"{p.__class__.__name__}:0 {name}"
-    ret.append(e:=ProfileRangeEvent(row, TracingKey(op or name, ret=f"PC:{info.pc}" if info else None), Decimal(p._time), Decimal(p._time+width)))
-    if (et:=row_ends.get(row)) is not None and e.st < et: raise RuntimeError(f"packet {p} overlaps another packet in {row}.")
-    row_ends[row] = unwrap(e.en)
+  rows:dict[str, None] = {}
+  trace:dict[str, set[int]] = {}
+  def add(name:str, p:PacketType, idx=0, width=1, op_name=None, wave=None, info:InstructionInfo|None=None) -> None:
+    if hasattr(p, "wave"): wave = p.wave
+    rows.setdefault(r:=(f"WAVE:{wave}" if wave is not None else f"{p.__class__.__name__}:0 {name}"))
+    key = TracingKey(f"{op_name if op_name is not None else name} OP:{idx}", ret=f"PC:{info.pc}" if info is not None else None)
+    ret.append(ProfileRangeEvent(r, key, Decimal(p._time), Decimal(p._time+width)))
   for p, info in map_insts(data, lib, target):
     if len(ret) > getenv("MAX_SQTT_PKTS", 50_000): break
     if isinstance(p, (INST, INST_RDNA4)):
-      name = p.op.name if isinstance(p.op, (InstOp, InstOpRDNA4)) else f"0x{p.op:02x}"
-      add(name, p, width=10 if "BARRIER" in name else 1, info=info)
+      op_name = p.op.name if isinstance(p.op, (InstOp, InstOpRDNA4)) else f"0x{p.op:02x}"
+      name, width = (op_name, 10 if "BARRIER" in op_name else 1)
+      add(name, p, width=width, idx=int("OTHER" in name), info=info)
     if isinstance(p, (VALUINST, IMMEDIATE)): add(p.__class__.__name__, p, info=info)
     if isinstance(p, IMMEDIATE_MASK): add("IMMEDIATE", p, wave=unwrap(info).wave, info=info)
     if isinstance(p, (VMEMEXEC, ALUEXEC)):
@@ -357,9 +359,11 @@ def sqtt_timeline(data:bytes, lib:bytes, target:str) -> list[ProfileEvent]:
         add("VALU", p)
         add("SALU", p)
       else:
-        add(name.replace("_ALT", ""), p, op=name)
+        add(name.replace("_ALT", ""), p, op_name=name)
+      if p._time in trace.setdefault(name, set()): raise AssertionError(f"packets overlap in shared resource! {name}")
+      trace[name].add(p._time)
   pc_map = {addr:str(inst) for addr,inst in amd_decode(lib, target).items()}
-  return [ProfilePointEvent(r, "JSON", "pcMap", pc_map, ts=Decimal(0)) for r in row_ends]+ret
+  return [ProfilePointEvent(r, "JSON", "pcMap", pc_map, ts=Decimal(0)) for r in rows]+ret
 
 # ** SQTT OCC only unpacks wave start, end time and SIMD location
 
@@ -563,7 +567,7 @@ def get_render(query:str) -> dict:
   i, j, fmt = get_int(qs:=parse_qs(url.query), "ctx"), get_int(qs, "step"), url.path.lstrip("/")
   data = ctxs[i]["steps"][j]["data"]
   if fmt == "graph-rewrites": return {"value":get_full_rewrite(trace.rewrites[i][j]), "content_type":"text/event-stream"}
-  if fmt == "uops": return {"src":get_stdout(lambda: print_uops(_reconstruct(trace.rewrites[i][j-1].sink).src[2].src)), "lang":"txt"}
+  if fmt == "uops": return {"src":get_stdout(lambda: print_uops(data)), "lang":"txt"}
   if fmt == "code": return {"src":data, "lang":"cpp"}
   if fmt == "asm":
     ret:dict = {}