fix later

* move reshape to MathTraits

* confirm it works in amd_uop_matmul

.github/workflows/autogen.yml

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

.github/workflows/benchmark.yml

@@ -527,7 +527,7 @@ jobs:
     - name: Run 10 CIFAR training steps
       run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=330 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
     - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=330 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=350 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
     # - name: Run 10 CIFAR training steps w BF16
     #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
     # TODO: too slow

.github/workflows/test.yml

@@ -2,7 +2,7 @@ name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
   DOWNLOAD_CACHE_VERSION: '12'
-  PYTHON_CACHE_VERSION: '3'
+  PYTHON_CACHE_VERSION: '4'
   APT_CACHE_VERSION: '1'
   BUILD_CACHE_VERSION: '1'
   CAPTURE_PROCESS_REPLAY: 1
@@ -230,7 +230,7 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: linting-only
-        python-version: '3.10'
+        python-version: '3.11'
         deps: linting
     - name: Lint bad-indentation and trailing-whitespace with pylint
       run: python -m pylint --disable=all -e W0311 -e C0303 --jobs=0 --indent-string='  ' --recursive=y .
@@ -243,8 +243,9 @@ jobs:
       run: |
         python -m mypy --strict-equality --lineprecision-report .
         cat lineprecision.txt
-    - name: Run TYPED=1
-      run: TYPED=1 python -c "import tinygrad"
+    # broken because of UPatAny
+    #- name: Run TYPED=1
+    #  run: TYPED=1 python -c "import tinygrad"
 
   unittest:
     name: Unit Tests

autogen_stubs.sh

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

extra/gemm/amd_uop_matmul.py

@@ -1,353 +1,180 @@
 from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, graph_rewrite, AxisType, PatternMatcher, UPat
-from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
+from tinygrad.uop.ops import UOp, KernelInfo
+from tinygrad.engine.realize import ExecItem, get_runner
 from tinygrad.dtype import AddrSpace
-from tinygrad.helpers import getenv, colored, prod, unwrap
-from tinygrad.shape.shapetracker import ShapeTracker, View
-from tinygrad.shape.view import strides_for_shape
-from tinygrad.codegen.opt.kernel import axis_colors, Opt, OptOps
-from tinygrad.codegen.opt.swizzler import merge_views, view_left
-
-def to_colored(full_shape, axis_types): return '_'.join([colored(str(s), axis_colors[at]) for s,at in zip(full_shape, axis_types)])
+from tinygrad.helpers import getenv
 
 N = 4096
 run_count = 5
 
-BN = 128
-BM = 128
-BK = 8
+# ---------------------------
+# launch/config constants
+# ---------------------------
 
-TN = 4
-TM = 4
+WARP_SIZE = 32
 
-# NOTE: this is from testgrad
-# change reduceop axes and input ShapeTrackers, view gets replaced with a reshape.
-# src->r->view  -->   src->view->r
-def swizzle_reduceop(src:UOp, r:UOp, view:UOp):
-  if r.tag is not None: return None
-  # confirm the input is in order
-  # TODO: replace this with a UOp that allows for nothing else then remove this
-  permute = tuple(i for i in range(len(src.shape)) if i not in r.axis_arg)+r.axis_arg
-  assert permute == tuple(range(len(permute))), f"reduce axis must already be in order, {permute} isn't"
+# 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
 
-  # append the reduce shape to each of the views
-  prshape = prod(rshape:=src.shape[-len(r.axis_arg):])
-  rstrides = strides_for_shape(rshape)
-  nv = [View.create(v.shape+rshape, tuple(x*prshape for x in v.strides)+rstrides, v.offset*prshape,
-                    v.mask+tuple((0,s) for s in rshape) if v.mask is not None else None) for v in unwrap(view.st).views]
+# Register tile sizes (per-thread accumulator tile of C)
+TN = 4     # columns per thread
+TM = 4     # rows per thread
 
-  # no reshape required with shrinking REDUCE_AXIS
-  return UOp(Ops.REDUCE_AXIS, r.dtype, (src.view(ShapeTracker(tuple(nv))),),
-             (r.arg[0], tuple(range(len(view.shape), len(view.shape) + len(r.axis_arg)))))
+is_kernel5 = getenv("K5", 0)
+THREADS_PER_BLOCK = 128 if is_kernel5 else 256
+assert THREADS_PER_BLOCK % BLOCK_N == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_N"
+assert THREADS_PER_BLOCK % BLOCK_K == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_K"
+assert (BLOCK_N * BLOCK_K) % THREADS_PER_BLOCK == 0
+assert (BLOCK_M * BLOCK_K) % THREADS_PER_BLOCK == 0
 
-pm = PatternMatcher([
-  (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), swizzle_reduceop),
-])
+WARPS_PER_BLOCK = THREADS_PER_BLOCK // WARP_SIZE
+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_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"
 
-def rangeify_kernel3():
-  a = Tensor.empty(N,N)
-  b = Tensor.empty(N,N)
-  c = a@b
-  #c = c.reshape((32,2,16,4,32,2,16,4)).contiguous()
-  sink = c.schedule()[-1].ast
-  #print(sink)
+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)
+N_PER_ITER = WAVE_TILE_N // ITERS_PER_WAVE_N
+M_PER_ITER = WAVE_TILE_M // ITERS_PER_WAVE_M
+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"
 
-  opts  = [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.UPCAST, 0, 2)]
-  opts += [Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 1, 16), Opt(OptOps.UPCAST, 1, 2)]
-  opts += [Opt(OptOps.UNROLL, 0, 8)]
+def hand_spec_kernel3():
+  # ---------------------------
+  # per-thread read mapping
+  # ---------------------------
+  # A: read BK x BN tiles; B: read BN x BK tiles
+  tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
 
-  return sink.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
+  waveIdx = (tid // WARP_SIZE) % WAVES_IN_BLOCK_X
+  waveIdy = (tid // WARP_SIZE) // WAVES_IN_BLOCK_X
+  assert waveIdy.vmax+1 == WAVES_IN_BLOCK_Y
 
-def top_spec_kernel3():
-  a = Tensor.empty(N,N)
-  b = Tensor.empty(N,N)
-  c = a@b
-  sink = c.schedule()[-1].ast
-  L = 16
-  sink = sink.reshape((N//L, L, N//L, L)) #.lift({0:UOp.range(N//BM, 0), 2:UOp.range(N//BN, 1)})
-  sink = graph_rewrite(sink, view_left+pm)
-  axis_types = (AxisType.GLOBAL, AxisType.LOCAL, AxisType.GLOBAL, AxisType.LOCAL, AxisType.REDUCE)
-  return sink.replace(arg=KernelInfo(name="top_"+to_colored(sink.full_shape, axis_types), axis_types=axis_types))
+  idxInWave = (tid % WARP_SIZE) % LANES_PER_WAVE_X
+  idyInWave = (tid % WARP_SIZE) // LANES_PER_WAVE_X
+  assert idyInWave.vmax+1 == LANES_PER_WAVE_Y
 
-def hl_spec_kernel3():
-  nbIterWaveM = 2
-  nbIterWaveN = 2
+  # ---------------------------
+  # block indices & placeholders
+  # ---------------------------
+  blockIdx_x = UOp.special(N // BLOCK_N, "gidx0")
+  blockIdx_y = UOp.special(N // BLOCK_M, "gidx1")
 
-  # define buffers
-  # TODO: remove these views once the defines have a shape
-  a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1).view(ShapeTracker.from_shape((N,N)))
-  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2).view(ShapeTracker.from_shape((N,N))).permute((1,0))
-  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0).view(ShapeTracker.from_shape((N,N)))
-  As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BM, AddrSpace.LOCAL), arg=0).view(ShapeTracker.from_shape((BK, BM))).permute((1,0))
-  Bs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BN, AddrSpace.LOCAL), arg=1).view(ShapeTracker.from_shape((BK, BN))).permute((1,0))
-  A_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveM * TM, AddrSpace.REG), arg=0).view(ShapeTracker.from_shape((nbIterWaveM * TM,)))
-  B_row = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveN * TN, AddrSpace.REG), arg=1).view(ShapeTracker.from_shape((nbIterWaveN * TN,)))
+  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)
 
-  # shape buffers. TODO: permutes
-  full_shape = (N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, N//BK, BK)
-  a = a.reshape((N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, 1, 1, 1, 1, N//BK, BK)).expand(full_shape)
-  b = b.reshape((1, 1, 1, 1, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, N//BK, BK)).expand(full_shape)
-  c = c.reshape((N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, 1, 1))
-  As = As.reshape((1, nbIterWaveM, BM//(nbIterWaveM * TM), TM, 1, 1, 1, 1, 1, BK)).expand(full_shape)
-  Bs = Bs.reshape((1, 1, 1, 1, 1, nbIterWaveN, BN//(nbIterWaveN * TN), TN, 1, BK)).expand(full_shape)
-  A_col = A_col.reshape((1, nbIterWaveM, 1, TM, 1, 1, 1, 1, 1, 1)).expand(full_shape)
-  B_row = B_row.reshape((1, 1, 1, 1, 1, nbIterWaveN, 1, TN, 1, 1)).expand(full_shape)
+  BM_As_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
+  As = UOp.placeholder((BLOCK_K, BM_As_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
+  Bs = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
 
-  #                     U1   L2 L3 L4 L5   U6 U7      U9   L10 L11 L12 L13   U14 U15      U17  U18  U19
-  expanded_shape = (32, 2,   2, 2, 2, 2,   2, 2,  32, 2,   2,  2,  2,  2,    2,  2,  512, 2,   2,   2)
-  assert len(expanded_shape) == 20
-  permute_a = list(range(len(expanded_shape)))
-  permute_b = permute_a[:]
+  A_col = UOp.placeholder((ITERS_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  B_row = UOp.placeholder((ITERS_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+  c_regs = UOp.placeholder((ITERS_PER_WAVE_M, TM, ITERS_PER_WAVE_N, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
 
-  # this makes all the global loads match
-  # this can also be more simply done by rebinding the RANGEs
-  # but sadly, rebinding the RANGEs doesn't work to change the order of the local axes
-  permute_a[17:20] = [11,12,13]
-  permute_a[11:14] = [17,18,19]
-  permute_a[7], permute_a[10] = permute_a[10], permute_a[7]
-  permute_a[2:7] = [3,4,5,6,2]
+  i = UOp.range(c_regs.size, 16)
+  c_regs = c_regs[i].set(0.0, end=i)
 
-  permute_b[2:16] = [19,9,10,11,17,18,8,2,12,13,14,15,3,4]
-  permute_b[17:20] = [5,6,7]
+  k_tile_range = UOp.range(N // BLOCK_K, 0)
 
-  a_permute   = a.reshape(expanded_shape).permute(tuple(permute_a)).reshape(full_shape)
-  As_permute = As.reshape(expanded_shape).permute(tuple(permute_a)).reshape(full_shape)
+  # ---------------------------
+  # GLOBAL -> LOCAL (As, Bs)
+  # ---------------------------
+  b = b.reshape(N // BLOCK_K, BLOCK_K,
+                N // BLOCK_N, BLOCK_N)
+  i = UOp.range(BLOCK_N * BLOCK_K // THREADS_PER_BLOCK, 1)
+  index_x = tid % BLOCK_N
+  index_y = (tid // BLOCK_N) + (THREADS_PER_BLOCK // BLOCK_N) * i
+  Bs_store = Bs[index_y, index_x].store(b[k_tile_range, index_y, blockIdx_x, index_x]).end(i)
 
-  b_permute   = b.reshape(expanded_shape).permute(tuple(permute_b)).reshape(full_shape)
-  Bs_permute = Bs.reshape(expanded_shape).permute(tuple(permute_b)).reshape(full_shape)
+  a = a.reshape(N // BLOCK_M, BLOCK_M,
+                N // BLOCK_K, BLOCK_K)
+  i = UOp.range(BLOCK_M * BLOCK_K // THREADS_PER_BLOCK, 2)
+  index_x = tid % BLOCK_K
+  index_y = (tid // BLOCK_K) + (THREADS_PER_BLOCK // BLOCK_K) * i
+  As_store = As[index_x, index_y].store(a[blockIdx_y, index_y, k_tile_range, index_x]).end(i)
 
-  #out = (a.load() * b.load()).r(Ops.ADD, (8, 9))
-  out = (As.load(As_permute.store(a_permute.load())) * Bs.load(Bs_permute.store(b_permute.load()))).r(Ops.ADD, (8, 9))
-  #out = (A_col.load(A_col.store(As.load(As.store(a.load())))) * B_row.load(B_row.store(Bs.load(Bs.store(b.load()))))).r(Ops.ADD, (8, 9))
+  # TODO: can we automate barrier?
+  barrier = UOp.barrier(As_store, Bs_store)
+  Bs = Bs.after(barrier)
+  As = As.after(barrier)
 
-  axis_types = (
-    AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.UPCAST,
-    AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.UPCAST,
-    AxisType.REDUCE, AxisType.REDUCE)
+  # open inner k range
+  k = UOp.range(BLOCK_K, 3)
 
-  sink = c.store(out).sink(arg=KernelInfo(name="tg_"+to_colored(full_shape, axis_types), axis_types=axis_types))
-  sink = graph_rewrite(sink, merge_views)
-  return sink
+  # ---------------------------
+  # LOCAL -> REG (per-wave tiles)
+  # ---------------------------
+  Bs_view = Bs.reshape(BLOCK_K, WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)
+  iterWaveN = UOp.range(ITERS_PER_WAVE_N, 4)
+  i = UOp.range(TN, 5)
+  B_row = B_row[iterWaveN, i].set(Bs_view[k, waveIdx, iterWaveN, idxInWave, i], end=(iterWaveN, i))
 
-def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
-  BLOCK_SIZE = 128 if kernel5 else 256
+  As_view = As.reshape(BLOCK_K, WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM)
+  iterWaveM = UOp.range(ITERS_PER_WAVE_M, 6)
+  i = UOp.range(TM, 7)
+  A_col = A_col[iterWaveM, i].set(As_view[k, waveIdy, iterWaveM, idyInWave, i], end=(iterWaveM, i))
 
-  nbWaves = BLOCK_SIZE // 32
-  WN = 128 if kernel5 else 64
-  WM = BN * BM // nbWaves // WN
+  # ---------------------------
+  # FMA: c_regs += A_col * B_row
+  # ---------------------------
+  iterWaveM = UOp.range(ITERS_PER_WAVE_M, 8)
+  yt = UOp.range(TM, 9)
+  iterWaveN = UOp.range(ITERS_PER_WAVE_N, 10)
+  xt = UOp.range(TN, 12)
+  c_idx = c_regs.after(k, k_tile_range)[iterWaveM, yt, iterWaveN, xt]
+  sink = c_idx.store(c_idx + A_col[iterWaveM, yt] * B_row[iterWaveN, xt]).end(iterWaveM, iterWaveN, yt, xt)
 
-  nbWaveX = BN // WN
-  nbWaveY = BM // WM
+  # Close k, sync, and close K tiles
+  sink = sink.end(k).barrier().end(k_tile_range)
 
-  threadIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("lidx0", BLOCK_SIZE))
-  waveIndex = threadIdx_x // 32
-  waveIdx = waveIndex % nbWaveX
-  waveIdy = waveIndex // nbWaveX
-  indexInWave = threadIdx_x % 32
-
-  nbThreadXPerWave = 8
-  nbThreadYPerWave = 4
-
-  idxInWave = indexInWave % nbThreadXPerWave
-  idyInWave = indexInWave // nbThreadXPerWave
-
-  nbIterWaveN = WN // (nbThreadXPerWave * TN)
-  nbIterWaveM = WM // (nbThreadYPerWave * TM)
-
-  SUBWN = WN // nbIterWaveN
-  SUBWM = WM // nbIterWaveM
-
-  # Thread mapping to read BKxBN block from A
-  rAIdx = threadIdx_x % BK
-  rAIdy = threadIdx_x // BK
-  # Thread mapping to read BNxBK block from B
-  rBIdx = threadIdx_x % BN
-  rBIdy = threadIdx_x // BN
-
-  strideReadB = BLOCK_SIZE // BN
-  strideReadA = BLOCK_SIZE // BK
-  nbReadsB = BN * BK // BLOCK_SIZE
-  nbReadsA = BM * BK // BLOCK_SIZE
-
-  blockIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx0", N//BN))
-  blockIdx_y = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx1", N//BM))
-
-  a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1)
-  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2)
-  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0)
-
-  A_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveM * TM, AddrSpace.REG), arg=0)
-  B_row = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveN * TN, AddrSpace.REG), arg=1)
-
-  BM_As_stride = (BM+4) if kernel5 else BM
-  As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BM_As_stride, AddrSpace.LOCAL), arg=0)
-  Bs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BN, AddrSpace.LOCAL), arg=1)
-
-  c_regs = UOp(Ops.DEFINE_REG, dtypes.float.ptr(TM * nbIterWaveM * TN * nbIterWaveN), arg=2)
-
-  i = UOp.range(c_regs.dtype.size, 16)
-  init_store = c_regs[i].store(UOp.const(dtypes.float, 0.0), i)
-
-  if kernel4:
-    regA = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbReadsA, AddrSpace.REG), arg=3)
-    regB = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbReadsB, AddrSpace.REG), arg=4)
-
-    # initial load from globals into locals (0)
-    kId = 0
-
-    # load from globals into locals
-    i = UOp.range(nbReadsB, 0)
-    index_x = BN * blockIdx_x + rBIdx
-    index_y = rBIdy + i * strideReadB + kId
-    Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(b[N * index_y + index_x].load(), i)
-
-    i = UOp.range(nbReadsA, 1)
-    index_x = rAIdx + kId
-    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
-    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)
-
-    # iterate over the middle chunk
-    kId_range = UOp.range(N//BK-1, 2)
-    kId = kId_range*BK
-
-    barrier = UOp.barrier(As_store, Bs_store)
-
-    # load from globals into registers (next round)
-    i = UOp.range(nbReadsB, 3)
-    index_x = BN * blockIdx_x + rBIdx
-    index_y = rBIdy + i * strideReadB + kId + BK
-    regB_store = regB[i].store(b[N * index_y + index_x].load(), i)
-
-    i = UOp.range(nbReadsA, 4)
-    index_x = rAIdx + kId + BK
-    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
-    regA_store = regA[i].store(a[N * index_y + index_x].load(), i)
-
-    def inner_loop(first_range, inp_dep=()):
-      # inner unroll
-      k = UOp.range(BK, first_range+0)
-
-      # load from locals into registers
-      iterWave = UOp.range(nbIterWaveN, first_range+1)
-      i = UOp.range(TN, first_range+2)
-      index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
-      B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(*inp_dep), iterWave, i)
-
-      iterWave = UOp.range(nbIterWaveM, first_range+3)
-      i = UOp.range(TM, first_range+4)
-      index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
-      A_col_store = A_col[iterWave*TM + i].store(As[k*BM_As_stride + index].load(*inp_dep), iterWave, i)
-
-      # do the GEMM math
-      iterWaveM = UOp.range(nbIterWaveM, first_range+5)
-      yt = UOp.range(TM, first_range+6)
-      iterWaveN = UOp.range(nbIterWaveN, first_range+7)
-      xt = UOp.range(TN, first_range+8)
-      x = iterWaveN * TN + xt
-      y = iterWaveM * TM + yt
-      c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
-      # sketchy, this should end the kId_range but it doesn't
-      sink = c_regs_idx.store(c_regs_idx.load(init_store) + A_col[y].load(A_col_store) * B_row[x].load(B_row_store),
-                              iterWaveM, iterWaveN, yt, xt, k)
-      return sink
-
-    # TODO: kId_range should endrange after a barrier
-    sink = inner_loop(5, (barrier, regB_store, regA_store)).barrier()
-
-    # load from registers into locals
-    i = UOp.range(nbReadsB, 14)
-    index_x = BN * blockIdx_x + rBIdx
-    index_y = rBIdy + i * strideReadB + kId + BK
-    Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(regB[i].load(sink), i, kId_range)
-
-    i = UOp.range(nbReadsA, 15)
-    index_x = rAIdx + kId + BK
-    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
-    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(regA[i].load(sink), i, kId_range)
-
-    # final iteration without the copy
-    sink = inner_loop(16, (UOp.barrier(Bs_store, As_store),))
-  else:
-    kId_range = UOp.range(N//BK, 0)
-    kId = kId_range*BK
-
-    # load from globals into locals
-    i = UOp.range(nbReadsB, 1)
-    index_x = BN * blockIdx_x + rBIdx
-    index_y = rBIdy + i * strideReadB + kId
-    Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(b[N * index_y + index_x].load(), i)
-
-    i = UOp.range(nbReadsA, 2)
-    index_x = rAIdx + kId
-    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
-    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)
-
-    barrier = UOp.barrier(As_store, Bs_store)
-
-    k = UOp.range(BK, 3)
-
-    # load from locals into registers
-    iterWave = UOp.range(nbIterWaveN, 4)
-    i = UOp.range(TN, 5)
-    index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
-    B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(barrier), iterWave, i)
-
-    iterWave = UOp.range(nbIterWaveM, 6)
-    i = UOp.range(TM, 7)
-    index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
-    A_col_store = A_col[iterWave*TM + i].store(As[k*BM_As_stride + index].load(barrier), iterWave, i)
-
-    # do the GEMM math
-    iterWaveM = UOp.range(nbIterWaveM, 8)
-    yt = UOp.range(TM, 9)
-    iterWaveN = UOp.range(nbIterWaveN, 10)
-    xt = UOp.range(TN, 12)
-    x = iterWaveN * TN + xt
-    y = iterWaveM * TM + yt
-    c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
-    sink = c_regs_idx.store(c_regs_idx.load(init_store) + A_col[y].load(A_col_store) * B_row[x].load(B_row_store),
-                            iterWaveM, iterWaveN, yt, xt, k, kId_range)
-
-  # store c_regs into c
-  iterWaveM = UOp.range(nbIterWaveM, 1000)
+  # ---------------------------
+  # REG -> GLOBAL (epilogue)
+  # ---------------------------
+  c = c.reshape(N//BLOCK_M, WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM,
+                N//BLOCK_N, WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)
+  iterWaveM = UOp.range(ITERS_PER_WAVE_M, 1000)
   yt = UOp.range(TM, 1001)
-  iterWaveN = UOp.range(nbIterWaveN, 1002)
+  iterWaveN = UOp.range(ITERS_PER_WAVE_N, 1002)
   xt = UOp.range(TN, 1003)
-  xOut = blockIdx_x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave
-  yOut = blockIdx_y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave
-  indexC = N * (yOut + yt) + xOut + xt
-  sink = c[indexC].store(c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)].load(sink),
-                         iterWaveM, iterWaveN, yt, xt)
+  c_glbl_idx = c[blockIdx_y, waveIdy, iterWaveM, idyInWave, yt, blockIdx_x, waveIdx, iterWaveN, idxInWave, xt]
+  sink = c_glbl_idx.store(c_regs.after(sink)[iterWaveM, yt, iterWaveN, xt])
+  sink = sink.end(iterWaveM, iterWaveN, yt, xt)
 
-  return sink.sink(arg=KernelInfo(name="tinygemm"))
+  return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()
+
+def test_matmul(sink:UOp, N=N):
+  with Context(DEBUG=0):
+    a = Tensor.randn(N, N)
+    b = Tensor.randn(N, N)
+    hc = Tensor.empty(N, N)
+    Tensor.realize(a, b, hc)
+
+  ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in [hc, a, b]])
+
+  GlobalCounters.reset()
+  ets = []
+  with Context(DEBUG=2):
+    for _ in range(run_count):
+      ets.append(ei.run(wait=True))
+  print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")
+
+  GlobalCounters.reset()
+  with Context(DEBUG=2):
+    tc = (a @ b).realize()
+  with Context(DEBUG=0):
+    err = (hc - tc).square().mean().item()
+  print(f"mean squared error {err}")
+  if err > 1e-06:
+    raise RuntimeError("matmul is wrong!")
 
 if __name__ == "__main__":
-  HL = getenv("HL")
-  if HL == 3: hprg = rangeify_kernel3()
-  elif HL == 2: hprg = top_spec_kernel3()
-  elif HL == 1: hprg = hl_spec_kernel3()
-  else: hprg = hand_spec_kernel3()
-  if HL == 3:
-    prg = get_program(hprg, Device.default.renderer)
-  else:
-    prg = get_program(hprg, Device.default.renderer)
-  print(prg.src)
-  if getenv("SRC"): exit(0)
-  hrunner = CompiledRunner(prg)
-
-  a = Tensor.randn(N, N).realize()
-  b = Tensor.randn(N, N).realize()
-  hc = Tensor.zeros(N, N).contiguous().realize()
-
-  GlobalCounters.reset()
-  with Context(DEBUG=2):
-    for _ in range(run_count): tc = (a@b).realize()
-
-  GlobalCounters.reset()
-  buffers = [hc.uop.buffer, a.uop.buffer, b.uop.buffer]
-  ei = ExecItem(hrunner, buffers)
-  with Context(DEBUG=2):
-    for _ in range(run_count): ei.run(wait=True)
-  err = (hc-tc).square().mean().item()
-  print(f"hrunner {err}")
-  if err > 1e-06: raise RuntimeError("matmul is wrong!")
+  test_matmul(hand_spec_kernel3(), N=N)

extra/gemm/metal_uop_matmul.py

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

extra/gemm/mi350x_uop_matmul.py

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

extra/gemm/torch_gemm.py

@@ -9,6 +9,7 @@ torch.set_num_threads(1)
 from tinygrad.helpers import getenv
 CUDA = getenv("CUDA", 1)
 MPS = getenv("MPS", 0)
+if getenv("FP16_ACC"): torch.backends.cuda.matmul.allow_fp16_accumulation = True
 
 for dtype in [torch.float32, torch.float16, torch.bfloat16]:
   for N in [256, 512, 1024, 2048, 4096]:

extra/mmapeak/mmapeak.py

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

extra/qcom_gpu_driver/msm_kgsl.py

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

extra/sqtt/README.md

@@ -2,8 +2,6 @@
 
 ## Getting SQ Thread Trace
 
-Only supported on 7900XTX, requires either AM (`rmmod amdgpu`) or disabling power gating on AMD (`ppfeaturemask=0xffff3fff`, don't forget to rebuild initramfs)
-
 SQTT is implemented on top of normal tinygrad profiling, `VIZ=1 SQTT=1` to get profile pickle with sqtt data embedded in it.
 
 `SQTT_BUFFER_SIZE=X` to change size of SQTT buffer (per shader engine, 6 SEs on 7900xtx) in megabytes, default 256.

extra/sqtt/roc.py

@@ -98,8 +98,10 @@ if __name__ == "__main__":
 
     return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
 
-  rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
-  print('SQTT:', ROCParseCtx.wave_events.keys())
+  try:
+    rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
+    print('SQTT:', ROCParseCtx.wave_events.keys())
+  except Exception as e: print("Error in sqtt decoder:", e)
 
   for ev in pmc_events:
     print(f"PMC Event: dev={ev.device} kern={ev.kern}")
@@ -107,6 +109,6 @@ if __name__ == "__main__":
     for s in ev.sched:
       view = memoryview(ev.blob).cast('Q')
       print(f"\t{s.name}")
-      for inst, se_idx, sa_idx, wgp_idx in itertools.product(range(s.inst), range(s.se), range(s.sa), range(s.wgp)):
-        print(f"\t\tInst {inst} SE {se_idx} SA {sa_idx} WGP {wgp_idx}: {view[ptr]}")
+      for xcc, inst, se_idx, sa_idx, wgp_idx in itertools.product(range(s.xcc), range(s.inst), range(s.se), range(s.sa), range(s.wgp)):
+        print(f"\t\tXCC {xcc} Inst {inst} SE {se_idx} SA {sa_idx} WGP {wgp_idx}: {view[ptr]:#x}")
         ptr += 1

extra/sqtt/rocprof/rocprof.py

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

extra/thunder/cuda/fa.cu

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

extra/thunder/cuda/fa.py

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

extra/thunder/cuda/matmul.cu

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

extra/thunder/cuda/matmul.py

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

extra/thunder/cuda/matmul2.cu

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

extra/tinyfs/fetch_file.py

@@ -3,9 +3,9 @@ import argparse
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
-  parser.add_argument("hash", type=str, required=True, help="file hash to fetch")
-  parser.add_argument("len", type=int, required=True, help="file length to fetch")
-  parser.add_argument("dest", type=str, required=True, help="destination path to save the file")
+  parser.add_argument("--hash", type=str, required=True, help="file hash to fetch")
+  parser.add_argument("--len", type=int, required=True, help="file length to fetch")
+  parser.add_argument("--dest", type=str, required=True, help="destination path to save the file")
   args = parser.parse_args()
 
   Tensor(bytes.fromhex(args.hash), device="CPU").load(args.len).to(f"disk:{args.dest}").realize()

setup.py

@@ -52,7 +52,7 @@ setup(name='tinygrad',
         "License :: OSI Approved :: MIT License"
       ],
       install_requires=[],
-      python_requires='>=3.10',
+      python_requires='>=3.11',
       extras_require={
         'arm': ["unicorn"],
         'triton': ["triton-nightly>=2.1.0.dev20231014192330"],

test/external/external_benchmark_op_conv.py

@@ -0,0 +1,260 @@
+# ruff: noqa: E501
+from tinygrad import dtypes, Device
+from tinygrad.uop.ops import UOp, AxisType, Ops
+from tinygrad.codegen import full_rewrite
+from tinygrad.renderer import ProgramSpec
+from tinygrad.engine.realize import CompiledRunner
+from tinygrad.helpers import dedup
+from tinygrad.device import Buffer
+from tinygrad.dtype import ImageDType
+
+# PYTHONPATH="." DEBUG=5 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+# kernel 672
+# faster on d59d4cd, 50% slower with the new linearizer
+
+""" d59d4cd
+c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), arg=0, src=())
+c1 = UOp.range(UOp.const(dtypes.index, 64), 3, AxisType.LOOP)
+c2 = UOp.range(UOp.const(dtypes.index, 64), 4, AxisType.LOOP)
+c3 = UOp.range(UOp.const(dtypes.index, 32), 2, AxisType.LOOP)
+c4 = (((c1*UOp.const(dtypes.index, 64))+c2)+(c3*UOp.const(dtypes.index, 4096)))
+c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), arg=1, src=())
+c6 = c5.index(c4).load()
+c7 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 3072, 4)), arg=2, src=())
+c8 = UOp.range(UOp.const(dtypes.index, 48), 0, AxisType.REDUCE)
+c9 = UOp.range(UOp.const(dtypes.index, 4), 1, AxisType.REDUCE)
+c10 = c7.index(((((c8*UOp.const(dtypes.index, 4))+c9)+(c1*UOp.const(dtypes.index, 192)))+(c3*UOp.const(dtypes.index, 12288)))).load()
+c11 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 192, 4)), arg=3, src=())
+c12 = c11.index(((((c9*UOp.const(dtypes.index, 4))+(c2%UOp.const(dtypes.index, 4)))+(c8*UOp.const(dtypes.index, 16)))+((c2//UOp.const(dtypes.index, 4))*UOp.const(dtypes.index, 768)))).load()
+c13 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(64), arg=4, src=())
+c14 = c13.index(c2).load()
+c15 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(64), arg=5, src=())
+c16 = c15.index(c2).load()
+c17 = (c6+(((c10*c12.cast(dtypes.float)).cast(dtypes.float).reduce(c8, c9, arg=Ops.ADD)+c14.cast(dtypes.float))*c16.cast(dtypes.float)))
+c18 = c0.index(c4).store(c17, c3, c1, c2)
+ast = c18.sink()
+more upcast axis : [(3, 320, 0, 4)]
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+__kernel void r_512_16_4_4_48_4(write_only image2d_t data0_131072, read_only image2d_t data1_131072, read_only image2d_t data2_393216, read_only image2d_t data3_12288, __global half* data4_64, __global half* data5_64) {
+const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;
+  float acc0[16];
+  int idx0 = get_global_id(0); /* 16 */
+  int idx1 = get_global_id(1); /* 512 */
+  int alu0 = (idx1>>4);
+  *(acc0+0) = 0.0f;
+  *(acc0+1) = 0.0f;
+  *(acc0+2) = 0.0f;
+  *(acc0+3) = 0.0f;
+  *(acc0+4) = 0.0f;
+  *(acc0+5) = 0.0f;
+  *(acc0+6) = 0.0f;
+  *(acc0+7) = 0.0f;
+  *(acc0+8) = 0.0f;
+  *(acc0+9) = 0.0f;
+  *(acc0+10) = 0.0f;
+  *(acc0+11) = 0.0f;
+  *(acc0+12) = 0.0f;
+  *(acc0+13) = 0.0f;
+  *(acc0+14) = 0.0f;
+  *(acc0+15) = 0.0f;
+  for (int Ridx0 = 0; Ridx0 < 48; Ridx0++) {
+    int alu17 = ((idx1*192)+Ridx0);
+    int alu18 = (alu17+48);
+    int alu19 = (alu17+96);
+    int alu20 = (alu17+144);
+    int alu21 = (Ridx0<<2);
+    float4 val0 = read_imagef(data3_12288, smp, (int2)(alu21,idx0));
+    float4 val1 = read_imagef(data3_12288, smp, (int2)((alu21+1),idx0));
+    float4 val2 = read_imagef(data3_12288, smp, (int2)((alu21+2),idx0));
+    float4 val3 = read_imagef(data3_12288, smp, (int2)((alu21+3),idx0));
+    float4 val4 = read_imagef(data2_393216, smp, (int2)((alu18-(3072*(((alu18>>10)*43)>>7))),alu0));
+    float4 val5 = read_imagef(data2_393216, smp, (int2)((alu19-(3072*(((alu19>>10)*43)>>7))),alu0));
+    float4 val6 = read_imagef(data2_393216, smp, (int2)((alu20-(3072*(((alu20>>10)*43)>>7))),alu0));
+    float4 val7 = read_imagef(data2_393216, smp, (int2)((alu17-(3072*(((alu17>>10)*43)>>7))),alu0));
+    *(acc0+1) = ((*(acc0+1))+(val4.x*val0.x)+(val4.y*val1.x)+(val4.z*val2.x)+(val4.w*val3.x));
+    *(acc0+5) = ((*(acc0+5))+(val4.x*val0.y)+(val4.y*val1.y)+(val4.z*val2.y)+(val4.w*val3.y));
+    *(acc0+9) = ((*(acc0+9))+(val4.x*val0.z)+(val4.y*val1.z)+(val4.z*val2.z)+(val4.w*val3.z));
+    *(acc0+13) = ((*(acc0+13))+(val4.x*val0.w)+(val4.y*val1.w)+(val4.z*val2.w)+(val4.w*val3.w));
+    *(acc0+2) = ((*(acc0+2))+(val5.x*val0.x)+(val5.y*val1.x)+(val5.z*val2.x)+(val5.w*val3.x));
+    *(acc0+6) = ((*(acc0+6))+(val5.x*val0.y)+(val5.y*val1.y)+(val5.z*val2.y)+(val5.w*val3.y));
+    *(acc0+10) = ((*(acc0+10))+(val5.x*val0.z)+(val5.y*val1.z)+(val5.z*val2.z)+(val5.w*val3.z));
+    *(acc0+14) = ((*(acc0+14))+(val5.x*val0.w)+(val5.y*val1.w)+(val5.z*val2.w)+(val5.w*val3.w));
+    *(acc0+3) = ((*(acc0+3))+(val6.x*val0.x)+(val6.y*val1.x)+(val6.z*val2.x)+(val6.w*val3.x));
+    *(acc0+7) = ((*(acc0+7))+(val6.x*val0.y)+(val6.y*val1.y)+(val6.z*val2.y)+(val6.w*val3.y));
+    *(acc0+11) = ((*(acc0+11))+(val6.x*val0.z)+(val6.y*val1.z)+(val6.z*val2.z)+(val6.w*val3.z));
+    *(acc0+15) = ((*(acc0+15))+(val6.x*val0.w)+(val6.y*val1.w)+(val6.z*val2.w)+(val6.w*val3.w));
+    *(acc0+0) = ((*(acc0+0))+(val7.x*val0.x)+(val7.y*val1.x)+(val7.z*val2.x)+(val7.w*val3.x));
+    *(acc0+4) = ((*(acc0+4))+(val7.x*val0.y)+(val7.y*val1.y)+(val7.z*val2.y)+(val7.w*val3.y));
+    *(acc0+8) = ((*(acc0+8))+(val7.x*val0.z)+(val7.y*val1.z)+(val7.z*val2.z)+(val7.w*val3.z));
+    *(acc0+12) = ((*(acc0+12))+(val7.x*val0.w)+(val7.y*val1.w)+(val7.z*val2.w)+(val7.w*val3.w));
+  }
+  int alu39 = (idx0<<2);
+  half4 val8 = (*((__global half4*)((data4_64+alu39))));
+  half4 val9 = (*((__global half4*)((data5_64+alu39))));
+  int alu40 = (idx0+(idx1<<6));
+  int2 cast0 = (int2)((alu40&1023),alu0);
+  float4 val10 = read_imagef(data1_131072, smp, cast0);
+  int2 cast1 = (int2)(((alu40+16)&1023),alu0);
+  float4 val11 = read_imagef(data1_131072, smp, cast1);
+  int2 cast2 = (int2)(((alu40+32)&1023),alu0);
+  float4 val12 = read_imagef(data1_131072, smp, cast2);
+  int2 cast3 = (int2)(((alu40+48)&1023),alu0);
+  float4 val13 = read_imagef(data1_131072, smp, cast3);
+  float cast4 = ((float)(val8.x));
+  float cast5 = ((float)(val9.x));
+  float cast6 = ((float)(val8.y));
+  float cast7 = ((float)(val9.y));
+  float cast8 = ((float)(val8.z));
+  float cast9 = ((float)(val9.z));
+  float cast10 = ((float)(val8.w));
+  float cast11 = ((float)(val9.w));
+  write_imagef(data0_131072, cast0, (float4)((val10.x+(((*(acc0+0))+cast4)*cast5)),(val10.y+(((*(acc0+4))+cast6)*cast7)),(val10.z+(((*(acc0+8))+cast8)*cast9)),(val10.w+(((*(acc0+12))+cast10)*cast11))));
+  write_imagef(data0_131072, cast1, (float4)((val11.x+(((*(acc0+1))+cast4)*cast5)),(val11.y+(((*(acc0+5))+cast6)*cast7)),(val11.z+(((*(acc0+9))+cast8)*cast9)),(val11.w+(((*(acc0+13))+cast10)*cast11))));
+  write_imagef(data0_131072, cast2, (float4)((val12.x+(((*(acc0+2))+cast4)*cast5)),(val12.y+(((*(acc0+6))+cast6)*cast7)),(val12.z+(((*(acc0+10))+cast8)*cast9)),(val12.w+(((*(acc0+14))+cast10)*cast11))));
+  write_imagef(data0_131072, cast3, (float4)((val13.x+(((*(acc0+3))+cast4)*cast5)),(val13.y+(((*(acc0+7))+cast6)*cast7)),(val13.z+(((*(acc0+11))+cast8)*cast9)),(val13.w+(((*(acc0+15))+cast10)*cast11))));
+}
+*** QCOM     672 r_512_16_4_4_48_4                              arg  6 mem   0.10 GB tm    322.55us/    77.83ms (    157 GFLOPS    4|160    GB/s) ['mul', '__add__', 'conv2d']
+"""
+
+""" master 99e76f33a0f4ec84c79c1271dbc955fe6b5a7778
+c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 0)
+c2 = UOp.range(64, 3, AxisType.LOOP)
+c4 = UOp.range(64, 4, AxisType.LOOP)
+c7 = UOp.range(32, 2, AxisType.LOOP)
+c10 = (((c2*64)+c4)+(c7*4096))
+c12 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 1)
+c14 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 3072, 4)), (), 2)
+c16 = UOp.range(48, 0, AxisType.REDUCE)
+c19 = UOp.range(4, 1, AxisType.REDUCE)
+c28 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 192, 4)), (), 3)
+c40 = (c14.index(((((c16*4)+c19)+(c2*192))+(c7*12288)))*c28.index(((((c19*4)+(c4%4))+(c16*16))+((c4//4)*768))))
+c42 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(64), (), 4)
+c46 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(64), (), 5)
+c50 = (c12.index(c10)+((c40.reduce(c16, c19, arg=Ops.ADD)+c42.index(c4).cast(dtypes.float))*c46.index(c4).cast(dtypes.float)))
+c52 = c0.index(c10, ptr=True).store(c50).end(c7, c2, c4)
+ast = c52.sink()
+more upcast axis : [(3, 320, 0, 4)]
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+__kernel void r_512_16_4_4_48_4(write_only image2d_t data0_131072, read_only image2d_t data1_131072, read_only image2d_t data2_393216, read_only image2d_t data3_12288, __global half* data4_64, __global half* data5_64) {
+const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;
+  float acc0[16];
+  int idx0 = get_global_id(0); /* 16 */
+  int idx1 = get_global_id(1); /* 512 */
+  *(acc0+0) = 0.0f;
+  *(acc0+1) = 0.0f;
+  *(acc0+2) = 0.0f;
+  *(acc0+3) = 0.0f;
+  *(acc0+4) = 0.0f;
+  *(acc0+5) = 0.0f;
+  *(acc0+6) = 0.0f;
+  *(acc0+7) = 0.0f;
+  *(acc0+8) = 0.0f;
+  *(acc0+9) = 0.0f;
+  *(acc0+10) = 0.0f;
+  *(acc0+11) = 0.0f;
+  *(acc0+12) = 0.0f;
+  *(acc0+13) = 0.0f;
+  *(acc0+14) = 0.0f;
+  *(acc0+15) = 0.0f;
+  int alu16 = (idx0<<2);
+  half4 val0 = (*((__global half4*)((data4_64+alu16))));
+  half4 val1 = (*((__global half4*)((data5_64+alu16))));
+  int alu17 = (idx0+(idx1<<6));
+  int alu18 = (idx1>>4);
+  int2 cast0 = (int2)((alu17&1023),alu18);
+  float4 val2 = read_imagef(data1_131072, smp, cast0);
+  int2 cast1 = (int2)(((alu17+16)&1023),alu18);
+  float4 val3 = read_imagef(data1_131072, smp, cast1);
+  int2 cast2 = (int2)(((alu17+32)&1023),alu18);
+  float4 val4 = read_imagef(data1_131072, smp, cast2);
+  int2 cast3 = (int2)(((alu17+48)&1023),alu18);
+  float4 val5 = read_imagef(data1_131072, smp, cast3);
+  for (int Ridx0 = 0; Ridx0 < 48; Ridx0++) {
+    int alu19 = ((idx1*192)+Ridx0);
+    int alu20 = (alu19+48);
+    int alu21 = (alu19+96);
+    int alu22 = (alu19+144);
+    int alu23 = (Ridx0<<2);
+    float4 val6 = read_imagef(data3_12288, smp, (int2)(alu23,idx0));
+    float4 val7 = read_imagef(data3_12288, smp, (int2)((alu23+1),idx0));
+    float4 val8 = read_imagef(data3_12288, smp, (int2)((alu23+2),idx0));
+    float4 val9 = read_imagef(data3_12288, smp, (int2)((alu23+3),idx0));
+    float4 val10 = read_imagef(data2_393216, smp, (int2)((alu20-(3072*(((alu20>>10)*43)>>7))),alu18));
+    *(acc0+1) = ((*(acc0+1))+(val10.x*val6.x)+(val10.y*val7.x)+(val10.z*val8.x)+(val10.w*val9.x));
+    *(acc0+5) = ((*(acc0+5))+(val10.x*val6.y)+(val10.y*val7.y)+(val10.z*val8.y)+(val10.w*val9.y));
+    *(acc0+9) = ((*(acc0+9))+(val10.x*val6.z)+(val10.y*val7.z)+(val10.z*val8.z)+(val10.w*val9.z));
+    *(acc0+13) = ((*(acc0+13))+(val10.x*val6.w)+(val10.y*val7.w)+(val10.z*val8.w)+(val10.w*val9.w));
+    float4 val11 = read_imagef(data2_393216, smp, (int2)((alu21-(3072*(((alu21>>10)*43)>>7))),alu18));
+    *(acc0+2) = ((*(acc0+2))+(val11.x*val6.x)+(val11.y*val7.x)+(val11.z*val8.x)+(val11.w*val9.x));
+    *(acc0+6) = ((*(acc0+6))+(val11.x*val6.y)+(val11.y*val7.y)+(val11.z*val8.y)+(val11.w*val9.y));
+    *(acc0+10) = ((*(acc0+10))+(val11.x*val6.z)+(val11.y*val7.z)+(val11.z*val8.z)+(val11.w*val9.z));
+    *(acc0+14) = ((*(acc0+14))+(val11.x*val6.w)+(val11.y*val7.w)+(val11.z*val8.w)+(val11.w*val9.w));
+    float4 val12 = read_imagef(data2_393216, smp, (int2)((alu22-(3072*(((alu22>>10)*43)>>7))),alu18));
+    *(acc0+3) = ((*(acc0+3))+(val12.x*val6.x)+(val12.y*val7.x)+(val12.z*val8.x)+(val12.w*val9.x));
+    *(acc0+7) = ((*(acc0+7))+(val12.x*val6.y)+(val12.y*val7.y)+(val12.z*val8.y)+(val12.w*val9.y));
+    *(acc0+11) = ((*(acc0+11))+(val12.x*val6.z)+(val12.y*val7.z)+(val12.z*val8.z)+(val12.w*val9.z));
+    *(acc0+15) = ((*(acc0+15))+(val12.x*val6.w)+(val12.y*val7.w)+(val12.z*val8.w)+(val12.w*val9.w));
+    float4 val13 = read_imagef(data2_393216, smp, (int2)((alu19-(3072*(((alu19>>10)*43)>>7))),alu18));
+    *(acc0+0) = ((*(acc0+0))+(val13.x*val6.x)+(val13.y*val7.x)+(val13.z*val8.x)+(val13.w*val9.x));
+    *(acc0+4) = ((*(acc0+4))+(val13.x*val6.y)+(val13.y*val7.y)+(val13.z*val8.y)+(val13.w*val9.y));
+    *(acc0+8) = ((*(acc0+8))+(val13.x*val6.z)+(val13.y*val7.z)+(val13.z*val8.z)+(val13.w*val9.z));
+    *(acc0+12) = ((*(acc0+12))+(val13.x*val6.w)+(val13.y*val7.w)+(val13.z*val8.w)+(val13.w*val9.w));
+  }
+  float cast4 = ((float)(val0.x));
+  float cast5 = ((float)(val1.x));
+  float cast6 = ((float)(val0.y));
+  float cast7 = ((float)(val1.y));
+  float cast8 = ((float)(val0.z));
+  float cast9 = ((float)(val1.z));
+  float cast10 = ((float)(val0.w));
+  float cast11 = ((float)(val1.w));
+  write_imagef(data0_131072, cast0, (float4)((val2.x+(((*(acc0+0))+cast4)*cast5)),(val2.y+(((*(acc0+4))+cast6)*cast7)),(val2.z+(((*(acc0+8))+cast8)*cast9)),(val2.w+(((*(acc0+12))+cast10)*cast11))));
+  write_imagef(data0_131072, cast1, (float4)((val3.x+(((*(acc0+1))+cast4)*cast5)),(val3.y+(((*(acc0+5))+cast6)*cast7)),(val3.z+(((*(acc0+9))+cast8)*cast9)),(val3.w+(((*(acc0+13))+cast10)*cast11))));
+  write_imagef(data0_131072, cast2, (float4)((val4.x+(((*(acc0+2))+cast4)*cast5)),(val4.y+(((*(acc0+6))+cast6)*cast7)),(val4.z+(((*(acc0+10))+cast8)*cast9)),(val4.w+(((*(acc0+14))+cast10)*cast11))));
+  write_imagef(data0_131072, cast3, (float4)((val5.x+(((*(acc0+3))+cast4)*cast5)),(val5.y+(((*(acc0+7))+cast6)*cast7)),(val5.z+(((*(acc0+11))+cast8)*cast9)),(val5.w+(((*(acc0+15))+cast10)*cast11))));
+}
+*** QCOM     672 r_512_16_4_4_48_4                              arg  6 mem   0.10 GB tm    527.97us/    78.94ms (     96 GFLOPS    3|98     GB/s) ['conv2d', 'mul', '__add__']
+"""
+
+c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 0)
+c2 = UOp.range(64, 3, AxisType.LOOP)
+c4 = UOp.range(64, 4, AxisType.LOOP)
+c7 = UOp.range(32, 2, AxisType.LOOP)
+c10 = (((c2*64)+c4)+(c7*4096))
+c12 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 1)
+c14 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 3072, 4)), (), 2)
+c16 = UOp.range(48, 0, AxisType.REDUCE)
+c19 = UOp.range(4, 1, AxisType.REDUCE)
+c28 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 192, 4)), (), 3)
+c40 = (c14.index(((((c16*4)+c19)+(c2*192))+(c7*12288)))*c28.index(((((c19*4)+(c4%4))+(c16*16))+((c4//4)*768))))
+c42 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(64), (), 4)
+c46 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(64), (), 5)
+c50 = (c12.index(c10)+((c40.reduce(c16, c19, arg=Ops.ADD)+c42.index(c4).cast(dtypes.float))*c46.index(c4).cast(dtypes.float)))
+c52 = c0.index(c10, ptr=True).store(c50).end(c7, c2, c4)
+ast = c52.sink()
+
+compiler = Device.default.compiler
+renderer = Device.default.renderer
+allocator = Device.default.allocator
+
+uops = full_rewrite(ast, renderer)
+src = renderer.render(uops)
+
+# NOLOCALS=1 IMAGE=2 DEV=CL
+lib = compiler.compile(src)
+# r_64_8_16_4_4_48_4
+# NOLOCALS: r_512_16_4_4_48_4
+ps = ProgramSpec("r_512_16_4_4_48_4", src, Device.DEFAULT, ast, uops)
+print(ps.src)
+print(ps.applied_opts)
+# (Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None))
+cr = CompiledRunner(ps, precompiled=lib)
+
+gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.DEFINE_GLOBAL]), key=lambda u: u.arg)
+print(len(gs))
+print([g.dtype for g in gs])
+
+bufs = [Buffer(ps.device, g.size, g.dtype if isinstance(g.dtype, ImageDType) else g.dtype._base).ensure_allocated() for g in gs]
+
+t = cr(bufs, wait=True)
+print(f"{t*1e6:.2f} us")

test/test_custom_kernel.py

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

test/test_schedule.py

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

test/test_uops.py

@@ -2,13 +2,13 @@ from typing import Optional, Any
 import unittest, math
 import numpy as np
 from tinygrad.tensor import Tensor, _to_np_dtype
-from tinygrad.helpers import CI, DEBUG, getenv, Timing
+from tinygrad.helpers import CI, DEBUG, getenv, Timing, Context
 from tinygrad.dtype import dtypes, DType, AddrSpace
 from tinygrad.device import Buffer, Device
-from tinygrad.uop.ops import Ops, UOp, UPat, KernelInfo, exec_alu # noqa F401
+from tinygrad.uop.ops import Ops, UOp, UPat, KernelInfo, exec_alu, AxisType
 from tinygrad.uop.spec import shared_spec
 from tinygrad.renderer import ProgramSpec
-from tinygrad.engine.realize import CompiledRunner, get_program
+from tinygrad.engine.realize import CompiledRunner, get_program, get_runner, ExecItem
 from tinygrad.codegen import full_rewrite
 from tinygrad.uop.symbolic import sym
 from tinygrad.device import is_dtype_supported
@@ -39,7 +39,7 @@ def _test_single_value(vals, op, dts):
   output_dtype = dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else dts[-1]
   buf_store = uop(uops, Ops.DEFINE_GLOBAL, output_dtype.ptr(), (), 0)
   buf_loads = [uop(uops, Ops.DEFINE_GLOBAL, dtype.ptr(), (), i+1) for i,dtype in enumerate(dts)]
-  loads = (uop(uops, Ops.LOAD, dtype, [buf_loads[i].index(uop(uops, Ops.CONST, dtypes.int32, (), 0), ptr=True)]) for i, dtype in enumerate(dts))
+  loads = (buf_loads[i].index(uop(uops, Ops.CONST, dtypes.int32, (), 0)) for i, dtype in enumerate(dts))
   alu = uop(uops, op, output_dtype, loads)
   out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0), ptr=True), alu))
   buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
@@ -56,7 +56,7 @@ def _test_single_value_const(vals, op, dts):
   buf_store = uop(uops, Ops.DEFINE_GLOBAL, output_dtype.ptr(), (), 0)
   loads = (uop(uops, Ops.CONST, dtype, [], a) for a,dtype in zip(vals, dts))
   alu = uop(uops, op, output_dtype, loads)
-  out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), alu))
+  out = buf_store[UOp.const(dtypes.int32, 0)].store(alu)
   buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
   prg = _uops_to_prg([out])
   prg.exec([buf])
@@ -565,5 +565,73 @@ class TestZeroRange(unittest.TestCase):
       out = Tensor.ones(10, dtype=dtypes.int).contiguous().shrink(((0,v),)).sum()
       self.assertEqual(out.item(), i)
 
+class TestUOpPrograms(unittest.TestCase):
+  def _run(self, prog:UOp, *tensors:Tensor):
+    ExecItem(get_runner(Device.DEFAULT, prog), [t.uop.buffer for t in tensors]).run(wait=True)
+
+  def test_simple(self):
+    out = Tensor.empty(10,10,dtype=dtypes.int)
+
+    ptr = UOp.placeholder(out.shape, out.dtype, slot=0)
+    i, j = UOp.range(10, axis_id=0), UOp.range(10, axis_id=1)
+    prog = ptr[i,j].set(42).end(i,j)
+    self._run(prog.sink(), out)
+
+    with Context(DEBUG=0): self.assertTrue((out == 42).all().item())
+
+  def test_matmul(self):
+    a = Tensor.randn(10,10)
+    b = Tensor.randn(10,10)
+    c = Tensor.empty(10,10)
+    ref = (a@b)
+    with Context(DEBUG=0): Tensor.realize(a, b, c, ref)
+
+    # C[i,j] = sum_k A[i,k] * B[k,j]
+    # Shapes: A[M,K], B[K,N], C[M,N]
+    M = N = K = 10
+    DT = dtypes.float32
+
+    # Placeholders (bind slots explicitly)
+    A = UOp.placeholder((M, K), DT, slot=0)
+    B = UOp.placeholder((K, N), DT, slot=1)
+    C = UOp.placeholder((M, N), DT, slot=2)
+
+    # Axes: i,j are spatial; k is a reduction axis over the shared dim K
+    i = UOp.range(M, axis_id=0)                             # rows of A/C
+    j = UOp.range(N, axis_id=1)                             # cols of B/C
+    k = UOp.range(K, axis_id=2, axis_type=AxisType.REDUCE)  # reduction over K
+
+    # Zero-init: write a scalar 0 to each (i,j).
+    C = C[i, j].set(0.0)
+
+    # Accumulate: C_after(k) enforces the dependency along the reduction axis
+    C = C[i, j].set(C.after(k)[i, j] + A[i, k] * B[k, j])
+
+    # Finalize the loop nest / schedule in (i, j, k) order
+    prog = C.end(i, j, k)
+
+    # run program
+    # TODO: make this work with opts_to_apply
+    self._run(prog.sink(arg=KernelInfo(opts_to_apply=())), a, b, c)
+
+    with Context(DEBUG=0): self.assertLessEqual((c-ref).square().mean().item(), 1e-6)
+
+  def test_matmul_relu(self):
+    a, b, c = Tensor.randn(10,10), Tensor.randn(10,10), Tensor.empty(10,10)
+    ref = (a@b).relu()
+    with Context(DEBUG=0): Tensor.realize(a, b, c, ref)
+
+    A, B, C = a.uop.placeholder_like(0), b.uop.placeholder_like(1), c.uop.placeholder_like(2)
+    i, j, k = UOp.range(10, 0), UOp.range(10, 1), UOp.range(10, 2, axis_type=AxisType.REDUCE)
+
+    C = C[i, j].set(0.0)
+    C = C[i, j].set(C.after(k)[i, j] + A[i, k] * B[k, j], end=k)
+    C = C[i, j].set(C[i, j].maximum(0.0))
+
+    prog = C.end(i, j)
+
+    self._run(prog.sink(arg=KernelInfo(opts_to_apply=())), a, b, c)
+    with Context(DEBUG=0): self.assertLessEqual((c-ref).square().mean().item(), 1e-6)
+
 if __name__ == '__main__':
   unittest.main(verbosity=2)

test/unit/test_helpers.py

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

test/unit/test_uop_symbolic.py

@@ -1019,10 +1019,7 @@ class TestSymbolicRealWorld(unittest.TestCase):
     #print(idx.render())
     # NOTE: this used to have 13,151,129,600 in the output which is out of int32 range.
     self.assertIn(idx.render(),
-      ("((((((((((lidx5+1)//16)*802816)+(((lidx5+1)%16)*49))+(gidx0*3211264))+(gidx1*784))+(gidx2*8))+(lidx4*100352))+lidx3)+2207744)",
-       '((lidx3+((((((((lidx5+1)//16)*802816)+(((lidx5+1)%16)*49))+(gidx0*3211264))+(gidx1*784))+(gidx2*8))+(lidx4*100352)))+2207744)',
-       '((lidx3+((lidx4*100352)+((gidx2*8)+((gidx1*784)+((gidx0*3211264)+((((lidx5+1)//16)*802816)+(((lidx5+1)%16)*49)))))))+2207744)',
-      ))
+      ("(lidx3+((lidx5+1)//16*802816+(lidx5+1)%16*49+gidx0*3211264+gidx1*784+gidx2*8+lidx4*100352)+2207744)",))
 
 class TestBounds(unittest.TestCase):
   def test_unrolled_arange(self):

tinygrad/codegen/__init__.py

@@ -1,4 +1,5 @@
 from typing import cast
+import itertools
 from tinygrad.helpers import DEVECTORIZE, TRANSCENDENTAL, SPEC
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat
 from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
@@ -15,7 +16,7 @@ from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_in
   ReduceContext, correct_load_store, pm_render, pm_add_loads
 from tinygrad.codegen.opt.postrange import apply_opts
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse, pm_split_store
-from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen
+from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 
 def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -> UOp:
@@ -23,6 +24,9 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
 
   if SPEC: type_verify(sink, kernel_spec)
 
+  # preprocess
+  sink = graph_rewrite(sink, pm_mops, name="early movement ops")
+
   # first we optimize
   if optimize:
     # collapse loads reduce (indexing by a tensor)
@@ -50,7 +54,7 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
   sink = graph_rewrite(sink, sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander")
 
   # add locals
-  sink = graph_rewrite(sink, pm_add_buffers_local+rangeify_codegen, name="add local buffers")
+  sink = graph_rewrite(sink, pm_add_buffers_local+rangeify_codegen, ctx=itertools.count(0), name="add local buffers")
 
   # ** devectorizer (full_graph_rewrite) **
   # remove reduce

tinygrad/codegen/gpudims.py

@@ -80,7 +80,7 @@ def add_gpudims(ctx:Renderer, s:UOp):
   subs = {}
   for r in s_topo:
     # look for local INDEXes that are not used in the GLOBAL store, then add them as an INVALID
-    if r.op is Ops.STORE and r.src[0].src[0].ptrdtype.addrspace == AddrSpace.GLOBAL:
+    if r.op is Ops.STORE and r.buf_target().ptrdtype.addrspace == AddrSpace.GLOBAL:
       idx = r.src[0]
       missing_locals = [all_ranges[rng] for rng in local_dims if all_ranges[rng] not in idx.ranges]
       if len(missing_locals):

tinygrad/codegen/late/expander.py

@@ -75,7 +75,18 @@ def do_contract(con:UOp):
     idxs += [_expand_arg_to_idx(ex.arg, {**rpk, **lrpk}) for lrpk in _choices_from_args(con.arg)]
   return UOp(Ops.UNROLL, con.dtype, (ex.src[0].gep(tuple(idxs)),), new_ex_args)
 
+def end_unrolls(u:UOp):
+  unrolls, src = partition(u.src[1:], lambda x: x.op is Ops.UNROLL)
+  if not len(unrolls): return None
+  ret = UOp(Ops.CONTRACT, dtypes.void, (u.src[0],), sum([x.arg for x in unrolls], start=()))
+  return u.replace(src=(ret,)+tuple(src))
+
 expander = PatternMatcher([
+  # push broadcast through AFTER
+  (UPat.var("x").broadcast(name="b").after(name="a", allow_any_len=True), lambda x,b,a: x.after(*a.src[1:]).broadcast(len(b.src))),
+  (UPat.var("x").broadcast(name="b").end(name="a", allow_any_len=True), lambda x,b,a: x.end(*a.src[1:]).broadcast(len(b.src))),
+  # END on UNROLL ends the UNROLL
+  (UPat(Ops.END, name="u"), end_unrolls),
   # BUFFERIZE puts UNROLLs for ranges as contract
   (UPat(Ops.BUFFERIZE, src=(UPat(Ops.UNROLL), UPat(Ops.UNROLL)), name="x"),
     lambda x: x.replace(src=tuple(UOp(Ops.CONTRACT, dtype=s.dtype.vec(x.src[1].src[0].dtype.count), src=(s,), arg=x.src[1].arg) for s in x.src))),
@@ -84,7 +95,7 @@ expander = PatternMatcher([
    lambda outer, inner: UOp(Ops.UNROLL, outer.dtype, (inner.src[0],), inner.arg+outer.arg)),
   # do expansion
   (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.GEP, Ops.WMMA, Ops.LOAD, Ops.STORE, Ops.INDEX, Ops.BUFFERIZE,
-         Ops.VECTORIZE, Ops.REDUCE, Ops.END), name="root", custom_early_reject=set([Ops.UNROLL])), do_expand),
+         Ops.VECTORIZE, Ops.REDUCE, Ops.END, Ops.AFTER), name="root", custom_early_reject=set([Ops.UNROLL])), do_expand),
   (UPat(Ops.CONTRACT, name="con"), do_contract),
   # BARRIERs aren't actually expanded
   (UPat(Ops.BARRIER, src=(UPat(Ops.UNROLL, name="ex"),)),

tinygrad/codegen/late/linearizer.py

@@ -64,7 +64,10 @@ class CFGContext:
       # ranges that have dependencies on other siblings need to be scheduled after them
       order = sorted(v, key=lambda x: len([u for u in v if u in deps[x]]))
       zipped = zip(order, order[1:]) if k.op is Ops.SINK else zip([k.src[1]] + order, order)
-      for x,y in zipped: self.edges[y.src[1]] = x
+      for x,y in zipped:
+        # TODO: this can happen! it causes infinite loop in shufflenet
+        assert y.src[1] not in x.backward_slice_with_self
+        self.edges[y.src[1]] = x
 
 pm_add_control_flow = PatternMatcher([
   (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(src=x.src+(y,)) if (y:=ctx.edges.get(x)) is not None else None),
@@ -72,7 +75,7 @@ pm_add_control_flow = PatternMatcher([
 
 def do_split_ends(e:UOp):
   ret = e.src[0]
-  for r in list(UOp.sink(*e.src[1:]).ranges)[::-1]: ret = ret.end(r)
+  for r in sorted(UOp.sink(*e.src[1:]).ranges, key=lambda x: x.arg, reverse=True): ret = ret.end(r)
   return ret
 
 pm_split_ends = PatternMatcher([

tinygrad/codegen/opt/heuristic.py

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

tinygrad/codegen/opt/postrange.py

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

tinygrad/codegen/opt/search.py

@@ -1,9 +1,8 @@
-from typing import cast
 import functools, math, time, multiprocessing, traceback, signal, atexit
 from dataclasses import replace
 from tinygrad.uop.ops import sym_infer, AxisType, pyrender
 from tinygrad.device import Device, Buffer, Compiler
-from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str
+from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str, unwrap
 from tinygrad.helpers import IGNORE_BEAM_CACHE
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.tensor import Tensor
@@ -50,7 +49,7 @@ def _time_program(p:ProgramSpec, lib:bytes, var_vals:dict[str, int], rawbufs:lis
       if hasattr(dev:=Device[p.device], 'invalidate_caches'): dev.invalidate_caches()
       else:
         with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024,1024).contiguous().realize(do_update_stats=False)
-    tms.append(cast(float, car(input_bufs, var_vals, wait=True))*factor)
+    tms.append(unwrap(car(input_bufs, var_vals, wait=True))*factor)
     if early_stop is not None and early_stop < min(tms): break
   return tms
 
@@ -168,7 +167,7 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True
           raise
         timed.append((candidates[i], min(tms)))
         if BEAM_DEBUG > 1:
-          print(f"{time.perf_counter() - st:7.2f}s: {i:5d} {len(cast(list, p.uops)):5d} uops",
+          print(f"{time.perf_counter() - st:7.2f}s: {i:5d} {len(unwrap(p.uops)):5d} uops",
                 f"{time_to_str(compile_et, w=12)} compile/{time_to_str(timed[-1][1], w=12)} run",
                 f"      {len(timed):4d}/{len(candidates):4d}         {timed[-1][0].colored_shape()}")
         elif DEBUG >= 2:

tinygrad/codegen/opt/tc.py

@@ -111,7 +111,7 @@ amd_rdna4 = [TensorCore(dims=(16,16,16), threads=32, elements_per_thread=(8,8,8)
   for di,do in [(dtypes.half,dtypes.float),(dtypes.half,dtypes.half),(dtypes.bfloat16,dtypes.float),(dtypes.bfloat16,dtypes.bfloat16)]]
 
 # https://gpuopen.com/learn/amd-lab-notes/amd-lab-notes-matrix-cores-readme
-amd_cdna = [TensorCore(dims=(16,16,16), threads=64, elements_per_thread=(4,4,4), dtype_in=di, dtype_out=do,
+amd_cdna_161616 = [TensorCore(dims=(16,16,16), threads=64, elements_per_thread=(4,4,4), dtype_in=di, dtype_out=do,
   opts=("l0","l0","l0","l0","u1","u1","l1","l1"),
   swizzle=((('u0', 'u1', 'l4', 'l5', 'r2', 'r3'), ('r0', 'r1'), ('l0', 'l1', 'l2', 'l3')),
            (('l0', 'l1', 'l2', 'l3', 'r2', 'r3'), ('r0', 'r1'), ('l4', 'l5', 'u0', 'u1'))))
@@ -119,11 +119,13 @@ amd_cdna = [TensorCore(dims=(16,16,16), threads=64, elements_per_thread=(4,4,4),
 
 amd_cdna_161632 = [TensorCore(dims=(16,16,32), threads=64, elements_per_thread=(8,8,4), dtype_in=di, dtype_out=do,
   opts=("l0","l0","l0","l0","u1","u1","l1","l1"),
-  swizzle=((('u0','u1','l4','l5','r3','r4'), ('r0','r1'), ('l0','l1','l2','l3','r2')),
-           (('l0','l1','l2','l3','r3','r4'), ('r0','r1'), ('l4','l5','u0','u1','r2'))))
-  for di,do in [(dtypes.half,dtypes.float),(dtypes.bfloat16,dtypes.float)]]
+  swizzle=((('u0', 'u1', 'l4', 'l5', 'r3', 'r4'), ('r0', 'r1'), ('l0', 'l1', 'l2', 'l3', 'r2')),
+           (('l0', 'l1', 'l2', 'l3', 'r3', 'r4'), ('r0', 'r1'), ('l4', 'l5', 'u0', 'u1', 'r2'))))
+  for di,do in [(dtypes.fp8e5m2,dtypes.float),(dtypes.fp8e4m3,dtypes.float),(dtypes.half,dtypes.float),(dtypes.bfloat16,dtypes.float)]]
 
-amd_cdna4 = amd_cdna_161632 + amd_cdna
+amd_cdna3 = amd_cdna_161632[:2] + amd_cdna_161616
+
+amd_cdna4 = amd_cdna_161632 + amd_cdna_161616
 
 # ***** Apple Metal *****
 

tinygrad/codegen/simplify.py

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

tinygrad/device.py

@@ -5,7 +5,7 @@ from typing import Any, Generic, TypeVar, Iterator, Sequence, cast, Generator
 import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
 from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored, CPU_LLVM
 from tinygrad.helpers import Context, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup
-from tinygrad.helpers import unwrap_class_type, suppress_finalizing
+from tinygrad.helpers import unwrap_class_type, suppress_finalizing, AMD_LLVM
 from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
 from tinygrad.renderer import Renderer
 
@@ -333,6 +333,7 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
     return device in {"AMD", "PYTHON", "NULL"}
   if dtype in dtypes.fp8s:
     if device in {"CUDA", "NV"}: return not CI and not getenv(f"{device}_PTX") and not getenv("NV_NAK")
+    if device == "AMD": return not CI and not AMD_LLVM and getattr(Device["AMD"], "target") in {(9,4,2), (9,5,0)}
     return device in {"PYTHON", "NULL"}
   if device == "WEBGPU": return dtype in [dtypes.bool, dtypes.char, dtypes.uchar, dtypes.short,
                                           dtypes.ushort, dtypes.float, dtypes.int32, dtypes.uint32, dtypes.half]

tinygrad/engine/realize.py

@@ -3,6 +3,7 @@ import time, pprint, random, itertools, math
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA, TracingKey
 from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, cpu_profile, PROFILE, ProfilePointEvent, cpu_events, prod, Context
+from tinygrad.helpers import unwrap
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, graph_rewrite, print_uops, track_rewrites, KernelInfo, pyrender
 from tinygrad.device import Device, Buffer
 from tinygrad.renderer import Renderer, ProgramSpec, Estimates
@@ -78,6 +79,7 @@ def optimize_local_size(_prg:Callable, global_size:list[int], rawbufs:list[Buffe
 
 class CompiledRunner(Runner):
   def __init__(self, p:ProgramSpec, precompiled:bytes|None=None, prg=None):
+    if DEBUG >= 3: print(p.applied_opts)
     if DEBUG >= 4: print(p.src)
     self.p:ProgramSpec = p
     if precompiled is not None: self.lib = precompiled
@@ -90,7 +92,8 @@ class CompiledRunner(Runner):
 
   def __reduce__(self): return self.__class__, (self.p, self.lib)
 
-  def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int], wait=False) -> float|None:
+  def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int]|None=None, wait=False) -> float|None:
+    if var_vals is None: var_vals = {}
     has_local = Device[self.p.device].renderer.has_local
     global_size, local_size = self.p.launch_dims(var_vals)
     if has_local and global_size is not None and local_size is None and all_int(self.p.global_size): # type: ignore[arg-type]
@@ -164,7 +167,7 @@ class ExecItem:
   fixedvars: dict[str, int] = field(default_factory=dict)
   def run(self, _var_vals:dict[str, int]|None=None, wait=False, jit=False, do_update_stats=True) -> float|None:
     var_vals = self.fixedvars if _var_vals is None else (_var_vals|self.fixedvars)
-    bufs = [cast(Buffer, x) for x in self.bufs] if jit else [cast(Buffer, x).ensure_allocated() for x in self.bufs]
+    bufs = [unwrap(x) for x in self.bufs] if jit else [unwrap(x).ensure_allocated() for x in self.bufs]
     if PROFILE:
       payload = {"metadata":self.metadata, "var_vals":var_vals, "bufs":[b.trace_num for b in bufs], "name":self.prg.display_name}
       payload["outputs"], payload["inputs"] = (self.prg.p.outs, self.prg.p.ins) if isinstance(self.prg, CompiledRunner) else ([0], [1])

tinygrad/gradient.py

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

tinygrad/helpers.py

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

tinygrad/renderer/__init__.py

@@ -1,13 +1,12 @@
 from __future__ import annotations
-from typing import Callable, cast, TYPE_CHECKING
+from typing import Callable, cast
 import functools
 from dataclasses import dataclass, field
 from tinygrad.helpers import to_function_name, dedup, prod
 from tinygrad.uop.ops import Ops, UOp, sym_infer, sint, Variable, ssimplify, GroupOp, PatternMatcher
 from tinygrad.dtype import AddrSpace, PtrDType
-if TYPE_CHECKING:
-  from tinygrad.codegen.opt.tc import TensorCore
-  from tinygrad.codegen.opt import Opt
+from tinygrad.codegen.opt.tc import TensorCore
+from tinygrad.codegen.opt import Opt
 
 @dataclass(frozen=True)
 class Estimates:
@@ -30,7 +29,7 @@ class Estimates:
     if ignore_indexing:
       def range_gate(x): return x.op is not Ops.RANGE
       for u in uops:
-        if u.op in {Ops.LOAD, Ops.STORE} and (not isinstance(u.src[0].dtype, PtrDType) or u.src[0].dtype.addrspace != AddrSpace.REG):
+        if u.op in {Ops.LOAD, Ops.STORE}:
           # if u.src[0] is INDEX, we have to include the buffer since it might be an AFTER
           dont_count = dont_count.union((UOp.sink(*u.src[0].src[1:]) if u.src[0].op is Ops.INDEX else u.src[0]).toposort(range_gate))
           # TODO: is this correct? this all needs to be cleaned up
@@ -81,16 +80,15 @@ class ProgramSpec:
       for u in self.uops:
         if u.op is Ops.DEFINE_VAR: self.vars.append(u)
         if u.op is Ops.DEFINE_GLOBAL: self.globals.append(u.arg)
-        if u.op is Ops.STORE and (u.src[0].op is Ops.INDEX or (u.src[0].op is Ops.CAST and u.src[0].src[0].op is Ops.INDEX)):
-          idx = u.src[0] if u.src[0].op is Ops.INDEX else u.src[0].src[0]
-          if (buf:=idx.src[0]).op is Ops.DEFINE_GLOBAL: self.outs.append(buf.arg)
-        if u.op is Ops.LOAD and (u.src[0].op is Ops.INDEX or (u.src[0].op is Ops.CAST and u.src[0].src[0].op is Ops.INDEX)):
-          idx = u.src[0] if u.src[0].op is Ops.INDEX else u.src[0].src[0]
-          if (buf:=idx.src[0]).op is Ops.DEFINE_GLOBAL: self.ins.append(buf.arg)
+        if u.op in (Ops.STORE, Ops.LOAD):
+          if (idx:=u.src[0]).op is Ops.INDEX or (u.src[0].op is Ops.CAST and (idx:=u.src[0].src[0]).op is Ops.INDEX):
+            if (buf:=idx.src[0]).op is Ops.DEFINE_GLOBAL: (self.outs if u.op is Ops.STORE else self.ins).append(buf.arg)
+          # TODO: can else happen?
         if u.op is Ops.SPECIAL:
           # NOTE: you have to set local_size and global_size to the base [1,1,1] outside this
           if u.arg[0] == 'i': self.local_size = None
           special_size = self.local_size if u.arg[0] == 'l' else self.global_size
+          # TODO: this cast is wrong, u.src[0].ssimplify() can be sint
           if special_size is not None: special_size[int(u.arg[-1])] = cast(int, u.src[0].ssimplify())
       self.vars = sorted(self.vars, key=lambda v: v.arg)
       self.outs = sorted(dedup(self.outs))

tinygrad/renderer/cstyle.py

@@ -71,11 +71,26 @@ extra_pm = PatternMatcher([
   (UPat(Ops.WHERE, name="alu"), no_vectorized_alu),
 ])
 
+def create_non_native_float_pats(dts:tuple[DType, ...], casting:bool=True):
+  patterns = PatternMatcher([
+    (UPat(Ops.WHERE, src=(UPat.var("b"), UPat.var("x", dtype=dts), UPat.var("y", dtype=dts))),
+     lambda b,x,y: UOp(Ops.WHERE, dtype=dtypes.float, src=(b,x.cast(dtypes.float),y.cast(dtypes.float))).cast(x.dtype)),
+    (UPat(GroupOp.ALU, dtype=dts, name="x"),
+     lambda x: UOp(x.op, dtypes.float, tuple(vv.cast(dtypes.float) for vv in x.src), x.arg).cast(x.dtype)),
+    (UPat(GroupOp.ALU, dtypes.bool, name="alu", src=(UPat.var("x", dtype=dts), UPat.var("y", dtype=dts))),
+     lambda alu,x,y: UOp(alu.op, dtypes.bool, (x.cast(dtypes.float), y.cast(dtypes.float)), alu.arg))])
+  if casting:
+    # add float intermediate casting
+    patterns += PatternMatcher([
+      (UPat(Ops.CAST, dts, (UPat.var("x"),), name="y"), lambda x,y: x.cast(dtypes.float).cast(y.dtype) if x.dtype!=dtypes.float else None),
+      (UPat(Ops.CAST, name="x", src=(UPat.var("y", dts),)), lambda x,y: y.cast(dtypes.float).cast(x.dtype) if x.dtype!=dtypes.float else None)])
+  return patterns
+
 def uops_to_dtypes(uops:list[UOp]) -> list[DType]: return dedup(u.dtype for u in uops if not isinstance(u.dtype, (ImageDType, PtrDType)))
 
 # (name, dims, dtype_in, dtype_out, device, threads, upcast_axes, reduce_axes)
 def wmma_args(uops:list[UOp]):
-  return dedup((uop.arg[0], uop.arg[1], uop.src[0].dtype.scalar(), uop.dtype.scalar(), *(uop.arg[4:8])) for uop in uops if uop.op is Ops.WMMA)
+  return dedup((uop.arg[0], uop.arg[1], uop.arg[2], uop.dtype.scalar(), *(uop.arg[4:8])) for uop in uops if uop.op is Ops.WMMA)
 
 class CStyleLanguage(Renderer):
   kernel_typedef: str = "void"
@@ -367,12 +382,8 @@ class CUDARenderer(CStyleLanguage):
     Ops.SQRT: lambda x,dtype: f"hsqrt({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"sqrt({x})",
     Ops.RECIPROCAL: lambda x,dtype: f"hrcp({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"(1/{x})" }
   type_map = {dtypes.bfloat16: "nv_bfloat16", dtypes.fp8e4m3: "__nv_fp8_e4m3", dtypes.fp8e5m2: "__nv_fp8_e5m2"}
-  extra_matcher = PatternMatcher([
+  extra_matcher = create_non_native_float_pats(dtypes.fp8s, casting=False) + PatternMatcher([
     (UPat(Ops.CAST, dtypes.fp8s, UPat.var("x", dtypes.fp8s), name='y'), lambda x,y: x.cast(dtypes.float).cast(y.dtype) if x.dtype!=y.dtype else None),
-    (UPat(GroupOp.ALU, dtype=dtypes.fp8s, name="x"),
-     lambda x: UOp(x.op, dtypes.float, tuple(vv.cast(dtypes.float) for vv in x.src), x.arg).cast(x.dtype)),
-    (UPat(GroupOp.ALU, dtypes.bool, name="alu", src=(UPat.var("x", dtype=dtypes.fp8s), UPat.var("y", dtype=dtypes.fp8s))),
-     lambda alu,x,y: UOp(alu.op, dtypes.bool, (x.cast(dtypes.float), y.cast(dtypes.float)), alu.arg)),
   ]) + extra_pm
   def render_vector_prefix(self, dt:DType) -> str:
     vec, scal = self.render_dtype(dt), self.render_dtype(dt.scalar()),
@@ -423,13 +434,20 @@ class AMDRenderer(CStyleLanguage):
 
   @staticmethod
   def get_tensor_cores(arch):
-    return {"gfx942": tc.amd_cdna, "gfx950": tc.amd_cdna4, "gfx1200": tc.amd_rdna4, "gfx1201": tc.amd_rdna4}.get(arch.split(":")[0], tc.amd_rdna3)
+    return {"gfx942": tc.amd_cdna3, "gfx950": tc.amd_cdna4, "gfx1200": tc.amd_rdna4, "gfx1201": tc.amd_rdna4}.get(arch.split(":")[0], tc.amd_rdna3)
+  @staticmethod
+  def is_cdna(arch): return arch.split(":")[0] in {"gfx942", "gfx950"}
   def __init__(self, arch:str): # gfx942 => MI300, gfx1100 => RX 7900, gfx1201 => RX 9700
     self.arch = arch
     self.tensor_cores = self.get_tensor_cores(arch)
-    if self.tensor_cores == tc.amd_cdna:
+    if self.is_cdna(self.arch):
       self.string_rewrite = PatternMatcher([
-        (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{x.arg[0]}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}, 0, 0, 0)")]) + base_rewrite
+        (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{x.arg[0]}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}, 0, 0, 0)"),
+        (UPat(Ops.CAST, dtypes.fp8s, (UPat.var("y", dtypes.float),), name="x",),
+          lambda ctx,x, y: f"f32_to_fp8({ctx[x.src[0]]}, {'1' if x.dtype == dtypes.fp8e5m2 else '0'})"),
+        (UPat(Ops.CAST, dtypes.float, (UPat.var("y", dtypes.fp8s),), name="x",),
+          lambda ctx,x, y: f"__builtin_amdgcn_cvt_f32_{'bf8' if y.dtype == dtypes.fp8e5m2 else 'fp8'}((unsigned int){ctx[x.src[0]]}, 0)"),
+      ]) + base_rewrite
   def __reduce__(self): return self.__class__, (self.arch,)
 
   # language options
@@ -455,20 +473,11 @@ class AMDRenderer(CStyleLanguage):
   barrier = '__builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");' + '__builtin_amdgcn_s_barrier();' + \
             '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");'
   float4 = "make_float4"
-  type_map = {dtypes.bfloat16: "hip_bfloat16"}
-  extra_matcher = PatternMatcher([
-    # cast bfloat16 alus to float
-    (UPat(Ops.WHERE, src=(UPat.var("b"), UPat.var("x", dtype=dtypes.bfloat16), UPat.var("y", dtype=dtypes.bfloat16))),
-      lambda b,x,y: UOp(Ops.WHERE, dtype=dtypes.float, src=(b,x.cast(dtypes.float),y.cast(dtypes.float))).cast(dtypes.bfloat16)),
-    (UPat(GroupOp.ALU, dtype=dtypes.bfloat16, name="x"),
-      lambda x: UOp(x.op, dtypes.float, tuple(vv.cast(dtypes.float) for vv in x.src), x.arg).cast(dtypes.bfloat16)),
-    (UPat(GroupOp.ALU, dtypes.bool, name="alu", src=(UPat.var("x", dtype=dtypes.bfloat16), UPat.var("y", dtype=dtypes.bfloat16))),
-      lambda alu,x,y: UOp(alu.op, dtypes.bool, (x.cast(dtypes.float), y.cast(dtypes.float)), alu.arg)),
-    # add float intermediate casting for bfloat16
-    (UPat(Ops.CAST, name="x", src=(UPat.var("y", dtypes.bfloat16),)),
-      lambda x,y: y.cast(dtypes.float).cast(x.dtype) if x.dtype!=dtypes.float else None),
-    (UPat(Ops.CAST, dtypes.bfloat16, (UPat.var("x"),)),
-      lambda x: x.cast(dtypes.float).cast(dtypes.bfloat16) if x.dtype!=dtypes.float else None),
+  type_map = {dtypes.bfloat16: "hip_bfloat16", dtypes.fp8e4m3: "hip_fp8", dtypes.fp8e5m2: "hip_bf8"}
+  extra_matcher = create_non_native_float_pats((dtypes.bfloat16, *dtypes.fp8s)) + PatternMatcher([
+    (UPat(Ops.WMMA, name="x", dtype=dtypes.float.vec(4)),
+      lambda x: UOp(Ops.WMMA, x.dtype, (x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64),
+        x.src[2]), (*x.arg,)) if x.src[0].dtype in (dtypes.fp8e4m3.vec(8), dtypes.fp8e5m2.vec(8)) else None),
     # bfloat16 casting
     (UPat.cvar('x', dtypes.bfloat16), lambda x: cast_float_to_bf16(UOp.const(dtypes.float, x.arg))),
     (UPat(Ops.CAST, dtypes.float, (UPat.var("x", dtypes.bfloat16),)),
@@ -482,14 +491,21 @@ class AMDRenderer(CStyleLanguage):
 
   def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str:
     prefix = ["#define INFINITY (__builtin_inff())","#define NAN (__builtin_nanf(\"\"))","typedef long unsigned int size_t;","#define half _Float16"]
-    type_map = { dtypes.bfloat16: "bf16", dtypes.float: "f32", dtypes.half: "f16" }
+    type_map = { dtypes.bfloat16: "bf16", dtypes.float: "f32", dtypes.half: "f16", dtypes.fp8e4m3: "_fp8_fp8", dtypes.fp8e5m2: "_bf8_bf8" }
     used_dtypes = uops_to_dtypes(uops)
     if any(dt.scalar() == dtypes.bfloat16 for dt in used_dtypes): prefix.append("typedef unsigned short hip_bfloat16;")
+    if any(dt.scalar() in dtypes.fp8s for dt in used_dtypes):
+      prefix += ["typedef unsigned char hip_bf8;", "typedef unsigned char hip_fp8;"]
+      prefix.append("""static inline __attribute__((device)) unsigned char f32_to_fp8(float v, int is_bf8) {
+  v = (((*(unsigned*)&v)&0x7F800000)!=0x7F800000)?__builtin_amdgcn_fmed3f(v,is_bf8?57344.0f:448.0f,is_bf8?-57344.0f:-448.0f) : v;
+  return (unsigned char)(is_bf8?__builtin_amdgcn_cvt_pk_bf8_f32(v,v,0,false):__builtin_amdgcn_cvt_pk_fp8_f32(v,v,0,false));\n}""")
     prefix += [self.render_vector_prefix(dt) for dt in used_dtypes if dt.count > 1]
 
-    for name, _, dtype_in, dtype_out, _, _, _, _ in wmma_args(uops): # TODO: handle TCs f32_bf16 and bf16_bf16 w/ wrapper
-      if self.tensor_cores == tc.amd_cdna:
-        prefix.append(f"#define __{name} __builtin_amdgcn_mfma_f32_16x16x16{'f16' if dtype_in == dtypes.half else 'bf16_1k'}")
+    for name, (N, M, K), dtype_in, dtype_out, _, _, _, _ in wmma_args(uops): # TODO: handle TCs f32_bf16 and bf16_bf16 w/ wrapper
+      if self.is_cdna(self.arch):
+        if (N, M, K) == (16, 16, 16): type_map[dtypes.bfloat16] = 'bf16_1k'
+        elif (N, M, K) == (16, 16, 32): type_map = {**type_map, dtypes.bfloat16: "_bf16", dtypes.half: "_f16"}
+        prefix.append(f"#define __{name} __builtin_amdgcn_mfma_f32_{N}x{M}x{K}{type_map[dtype_in]}")
       # #define __WMMA_16_16_16_half_half __builtin_amdgcn_wmma_f16_16x16x16_f16_w32_gfx12
       elif self.tensor_cores == tc.amd_rdna4:
         prefix.append(f"#define __{name} __builtin_amdgcn_wmma_{type_map[dtype_out]}_16x16x16_{type_map[dtype_in]}_w32_gfx12")

tinygrad/renderer/llvmir.py

@@ -246,7 +246,7 @@ class AMDLLVMRenderer(LLVMRenderer):
   def __init__(self, arch:str):
     self.arch = arch
     self.tensor_cores = AMDRenderer.get_tensor_cores(arch)
-    self.is_cdna = arch.split(":")[0] in {"gfx942", "gfx950"}
+    self.is_cdna = AMDRenderer.is_cdna(arch)
     self.string_rewrite += PatternMatcher([(UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, cdna=self.is_cdna: render_wmma_amd(ctx, wmma, cdna))])
     if self.is_cdna:
       self.extra_matcher += PatternMatcher([

tinygrad/runtime/autogen/am/am.py

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

tinygrad/runtime/autogen/am/pm4_nv.py

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

tinygrad/runtime/autogen/am/pm4_soc15.py

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

tinygrad/runtime/autogen/am/sdma_4_0_0.py

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

tinygrad/runtime/autogen/am/sdma_4_4_2.py

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

tinygrad/runtime/autogen/am/sdma_5_0_0.py

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

tinygrad/runtime/autogen/am/sdma_6_0_0.py

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

tinygrad/runtime/autogen/am/smu_v13_0_0.py

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

tinygrad/runtime/autogen/am/smu_v14_0_2.py

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

tinygrad/runtime/autogen/amd_gpu.py

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

tinygrad/runtime/autogen/comgr.py

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

tinygrad/runtime/autogen/cuda.py

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

tinygrad/runtime/autogen/hip.py

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

tinygrad/runtime/autogen/hsa.py

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

tinygrad/runtime/autogen/ib.py

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

tinygrad/runtime/autogen/io_uring.py

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

tinygrad/runtime/autogen/kfd.py

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

tinygrad/runtime/autogen/kgsl.py

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

tinygrad/runtime/autogen/libc.py

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

tinygrad/runtime/autogen/libusb.py

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

tinygrad/runtime/autogen/llvm.py

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

tinygrad/runtime/autogen/mesa.py

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

tinygrad/runtime/autogen/nv/nv.py

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

tinygrad/runtime/autogen/nv_gpu.py

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

tinygrad/runtime/autogen/nvrtc.py

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

tinygrad/runtime/autogen/opencl.py

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

tinygrad/runtime/autogen/qcom_dsp.py

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

tinygrad/runtime/autogen/sqtt.py

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

tinygrad/runtime/autogen/vfio.py

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

tinygrad/runtime/autogen/webgpu.py

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

tinygrad/runtime/ops_amd.py

@@ -31,7 +31,7 @@ AQL_HDR = (1 << hsa.HSA_PACKET_HEADER_BARRIER) | (hsa.HSA_FENCE_SCOPE_SYSTEM <<
 class ProfileSQTTEvent(ProfileEvent): device:str; se:int; props:dict; blob:bytes; itrace:bool # noqa: E702
 
 @dataclass(frozen=True)
-class PMCSample: name:str; block:str; inst:int; se:int; sa:int; wgp:int; off:int; size:int; reg:str # noqa: E702
+class PMCSample: name:str; block:str; xcc:int; inst:int; se:int; sa:int; wgp:int; off:int; size:int; reg:str # noqa: E702
 
 @dataclass(frozen=True)
 class ProfilePMCEvent(ProfileEvent): device:str; kern:str; sched:list[PMCSample]; blob:bytes # noqa: E702
@@ -74,10 +74,12 @@ class AMDComputeQueue(HWQueue):
 
   def set_grbm_broadcast(self):
     self.wreg(self.gc.regGRBM_GFX_INDEX, **{f'{f}_broadcast_writes': 1 for f in ['se', 'sh' if self.dev.target[0] == 9 else 'sa', 'instance']})
-  def set_grbm_se(self, se): self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, instance_broadcast_writes=1)
   def set_grbm_inst(self, n):
     self.wreg(self.gc.regGRBM_GFX_INDEX, **{f'{f}_broadcast_writes': 1 for f in ['se', 'sh' if self.dev.target[0] == 9 else 'sa']}, instance_index=n)
-  def set_grbm_se_sh_wgp(self, se, sa, wgp): self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, sa_index=sa, instance_index=wgp << 2)
+  def set_grbm_se_sh(self, se, sh):
+    self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, **{f'{"sh" if self.dev.target[0] == 9 else "sa"}_index':sh}, instance_broadcast_writes=1)
+  def set_grbm_se_sh_wgp(self, se, sh, wgp): self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, sa_index=sh, instance_index=wgp << 2)
+  def set_grbm_se(self, se): self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, sh_broadcast_writes=1, instance_broadcast_writes=1)
 
   def wait_reg_mem(self, value, mask=0xffffffff, mem=None, reg=None, reg_done=0, op=WAIT_REG_MEM_FUNCTION_GEQ):
     wrm_info_dw = self.pm4.WAIT_REG_MEM_MEM_SPACE(int(mem is not None)) | self.pm4.WAIT_REG_MEM_OPERATION(int(mem is None and reg_done > 0)) \
@@ -139,41 +141,51 @@ class AMDComputeQueue(HWQueue):
 
   def pmc_reset_counters(self, en=True):
     self.set_grbm_broadcast()
-    self.wreg(self.gc.regCP_PERFMON_CNTL, perfmon_state=0)
-    if en: self.wreg(self.gc.regCP_PERFMON_CNTL, perfmon_state=1)
+    self.wreg(self.gc.regCP_PERFMON_CNTL if self.dev.target[0] <= 11 else self.gc.regCP_PERFMON_CNTL_1, perfmon_state=0)
+    if en: self.wreg(self.gc.regCP_PERFMON_CNTL if self.dev.target[0] <= 11 else self.gc.regCP_PERFMON_CNTL_1, perfmon_state=1)
     return self
 
   def pmc_start(self, counters):
     self.pmc_reset_counters(en=False)
-    self.wreg(self.gc.regSQ_PERFCOUNTER_CTRL, cs_en=1, ps_en=1, gs_en=1, hs_en=1)
-    self.wreg(self.gc.regSQ_PERFCOUNTER_CTRL2, force_en=1, vmid_en=0xffff)
+    self.wreg(self.gc.regSQ_PERFCOUNTER_CTRL, cs_en=1, ps_en=1, gs_en=1, hs_en=1, **({'vmid_mask':0xffff} if (gfx9:=self.dev.target[0] == 9) else {}))
+    if self.dev.target[0] >= 11: self.wreg(self.gc.regSQ_PERFCOUNTER_CTRL2, force_en=1, vmid_en=0xffff)
 
-    out_off = 0
+    end_off = 0
     block2pid:dict[str, itertools.count] = collections.defaultdict(lambda: itertools.count())
     for name,block,idx in counters:
-      inst_cnt, se_cnt, sa_cnt, wgp_cnt = (32, 1, 1, 1) if block != "SQ" else (1, self.dev.se_cnt, 2, self.dev.iface.props['cu_per_simd_array'] // 2)
-      reg, out_off = f'reg{block}_PERFCOUNTER{next(block2pid[block])}', out_off + (rec_size:=prod((inst_cnt, se_cnt, sa_cnt, wgp_cnt)) * 8)
-      self.wreg(getattr(self.gc, f'{reg}_SELECT'), idx)
-      self.dev.pmc_sched.append(PMCSample(name, block, inst_cnt, se_cnt, sa_cnt, wgp_cnt, out_off-rec_size, rec_size, reg))
+      # sq block on gfx11+ goes down to wgps
+      inst_cnt, se_cnt, sa_cnt, wgp_cnt = {"GRBM": (1, 1, 1, 1), "GL2C": (32, 1, 1, 1), "TCC": (16, 1, 1, 1),
+        "SQ": (1, self.dev.se_cnt // self.dev.xccs) + ((1, 1) if gfx9 else (2, self.dev.iface.props['cu_per_simd_array'] // 2))}[block]
+      end_off += (rec_size:=prod((self.dev.xccs, inst_cnt, se_cnt, sa_cnt, wgp_cnt)) * 8)
 
+      if (regsel:=getattr(self.gc, (reg:=f'reg{block}_PERFCOUNTER{next(block2pid[block])}') + '_SELECT', None)) is None:
+        raise RuntimeError(f'{block} is out of perfcounter registers: ({reg} is not found)')
+
+      self.wreg(regsel, perf_sel=idx, **({'simd_mask':0xf, 'sqc_bank_mask':0xf, 'sqc_client_mask':0xf} if gfx9 and block == "SQ" else {}))
+      self.dev.pmc_sched.append(PMCSample(name, block, self.dev.xccs, inst_cnt, se_cnt, sa_cnt, wgp_cnt, end_off-rec_size, rec_size, reg))
+
+    if gfx9: self.wreg(self.gc.regSQ_PERFCOUNTER_MASK, sh0_mask=0xffff, sh1_mask=0xffff)
     self.wreg(self.gc.regCOMPUTE_PERFCOUNT_ENABLE, 1)
     return self.pmc_reset_counters(en=True)
 
   def pmc_read(self, buf, sched):
     self.set_grbm_broadcast()
-    self.wreg(self.gc.regCP_PERFMON_CNTL, perfmon_state=1, perfmon_sample_enable=1) # read counters
+    self.wreg(self.gc.regCP_PERFMON_CNTL if self.dev.target[0] <= 11 else self.gc.regCP_PERFMON_CNTL_1, perfmon_state=1, perfmon_sample_enable=1)
 
     for s in sched:
       offset = itertools.count(s.off, step=8)
 
-      for inst, se_idx, sa_idx, wgp_idx in itertools.product(range(s.inst), range(s.se), range(s.sa), range(s.wgp)):
-        if s.inst > 1: self.set_grbm_inst(inst)
-        else: self.set_grbm_se_sh_wgp(se_idx, sa_idx, wgp_idx)
+      for xcc in range(s.xcc):
+        with self.pred_exec(xcc_mask=1 << xcc):
+          for inst, se_idx, sa_idx, wgp_idx in itertools.product(range(s.inst), range(s.se), range(s.sa), range(s.wgp)):
+            if s.inst > 1: self.set_grbm_inst(inst)
+            elif self.dev.target[0] == 9: self.set_grbm_se(se_idx)
+            else: self.set_grbm_se_sh_wgp(se_idx, sa_idx, wgp_idx)
 
-        # Copy counter to memory (src_sel = perf, dst_sel = tc_l2)
-        lo, hi = getattr(self.gc, f'{s.reg}_LO'), getattr(self.gc, f'{s.reg}_HI', None)
-        self.pkt3(self.pm4.PACKET3_COPY_DATA, 2 << 8 | 4, lo.addr[0], 0, *data64_le(buf.va_addr+(loff:=next(offset))))
-        if hi is not None: self.pkt3(self.pm4.PACKET3_COPY_DATA, 2 << 8 | 4, hi.addr[0], 0, *data64_le(buf.va_addr+loff+4))
+            # Copy counter to memory (src_sel = perf, dst_sel = tc_l2)
+            lo, hi = getattr(self.gc, f'{s.reg}_LO'), getattr(self.gc, f'{s.reg}_HI', None)
+            self.pkt3(self.pm4.PACKET3_COPY_DATA, (2 << 8) | 4, lo.addr[0], 0, *data64_le(buf.va_addr+(loff:=next(offset))))
+            if hi is not None: self.pkt3(self.pm4.PACKET3_COPY_DATA, (2 << 8) | 4, hi.addr[0], 0, *data64_le(buf.va_addr+loff+4))
 
     return self.pmc_reset_counters(en=True)
 
@@ -216,7 +228,7 @@ class AMDComputeQueue(HWQueue):
         if (se_mask >> se) & 0b1: mask |= (__SQTTINST:=1<<10) | (__SQTT_INST_PC:=1<<11) | (__SQTT_ISSUE:=1<<13)
 
         with self.pred_exec(xcc_mask=1<<(se // (ses_per_xcc:=(self.dev.se_cnt // self.dev.xccs)))):
-          self.set_grbm_se(se % ses_per_xcc)
+          self.set_grbm_se_sh(se % ses_per_xcc, 0)
           self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK, reg_mask=0xf, token_mask=mask)
           self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK2, inst_mask=0xffffffff)
           self.wreg(self.gc.regSQ_THREAD_TRACE_BASE, addr=lo32(buf0s[se].va_addr >> 12))
@@ -228,7 +240,7 @@ class AMDComputeQueue(HWQueue):
       self.spi_config(tracing=True)
       # One buffer for one SE, mesa does it with a single buffer and ac_sqtt_get_data_offset, but this is simpler and should work just as well
       for se in range(len(buf0s)):
-        self.set_grbm_se(se)
+        self.set_grbm_se_sh(se, 0)
 
         buf0_lo, buf0_hi = data64_le(buf0s[se].va_addr >> 12)
         if self.dev.target >= (12,0,0):
@@ -279,7 +291,7 @@ class AMDComputeQueue(HWQueue):
 
     # For each SE wait for finish to complete and copy regSQ_THREAD_TRACE_WPTR to know where in the buffer trace data ends
     for se in range(ses):
-      self.set_grbm_se(se)
+      self.set_grbm_se_sh(se, 0)
 
       status_reg = self.gc.regSQ_THREAD_TRACE_STATUS.addr[0] - (self.pm4.PACKET3_SET_UCONFIG_REG_START if self.dev.target[0] == 9 else 0)
       if self.dev.target >= (10, 0, 0):
@@ -750,7 +762,8 @@ class KFDIface:
 
     raise RuntimeError("\n".join(report))
 
-  def is_in_profile_mode(self): return FileIOInterface(f'{self.dev_sysfs_path}/power_dpm_force_performance_level').read()[:16] == 'profile_standard'
+  def is_in_profile_mode(self):
+    return self.dev.target[0] == 9 or FileIOInterface(f'{self.dev_sysfs_path}/power_dpm_force_performance_level').read()[:16] == 'profile_standard'
 
 class PCIIface(PCIIfaceBase):
   gpus:ClassVar[list[str]] = []
@@ -761,7 +774,7 @@ class PCIIface(PCIIfaceBase):
     self._setup_adev(self.pci_dev)
     self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
 
-  def is_in_profile_mode(self): return False
+  def is_in_profile_mode(self): return True
 
   def _setup_adev(self, pci_dev:PCIDevice, dma_regions:list[tuple[int, MMIOInterface]]|None=None):
     self.dev_impl:AMDev = AMDev(pci_dev, dma_regions)
@@ -896,27 +909,27 @@ class AMDDevice(HCQCompiled):
 
     self.pmc_enabled = PROFILE and PMC > 0
     if self.pmc_enabled:
-      if self.target[0] not in {11}: raise RuntimeError(f'PMC are not supported on gc:{self.target}')
-      if not self.iface.is_in_profile_mode(): raise RuntimeError("PMC requires stable power state: AMD_IFACE=KFD and `amd-smi set -l stable_std`")
+      if self.target[0] not in {9, 11, 12}: raise RuntimeError(f'PMC are not supported on gc:{self.target}')
+      if not self.iface.is_in_profile_mode(): raise RuntimeError("PMC requires stable power state: run `amd-smi set -l stable_std` for KFD iface")
 
       self.pmc_sched:list[PMCSample] = []
       self.pmc_counters = import_pmc(self.target)
 
       # validate counters
-      for k in (PMC_COUNTERS:=getenv("PMC_COUNTERS", "GL2C_HIT,GL2C_MISS,SQC_LDS_IDX_ACTIVE,SQC_LDS_BANK_CONFLICT").split(",")):
+      pmc_default = "TCC_HIT,TCC_MISS,SQ_LDS_BANK_CONFLICT" if self.target[0] == 9 else "GL2C_HIT,GL2C_MISS,SQC_LDS_IDX_ACTIVE,SQC_LDS_BANK_CONFLICT"
+      for k in (PMC_COUNTERS:=getenv("PMC_COUNTERS", pmc_default).split(",")):
         if k not in self.pmc_counters: raise RuntimeError(f"PMC counter {k} is not supported. Available: {','.join(self.pmc_counters.keys())}")
 
-      cast(AMDComputeQueue, self.hw_compute_queue_t()).pmc_start([self.pmc_counters[k] for k in PMC_COUNTERS]).submit(self)
+      cast(AMDComputeQueue, self.hw_compute_queue_t()).pmc_start([(k, *self.pmc_counters[k]) for k in PMC_COUNTERS]).submit(self)
       self.pmc_buffer = self.allocator.alloc(self.pmc_sched[-1].off + self.pmc_sched[-1].size, BufferSpec(nolru=True, uncached=True))
+      self.allocator._copyin(self.pmc_buffer, memoryview(bytearray(self.pmc_buffer.size))) # zero pmc buffers, some counters have only lo part.
 
     # SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them)
     self.sqtt_enabled = PROFILE and SQTT > 0
     if self.sqtt_enabled:
       if self.target[0] not in {9, 11, 12}: raise RuntimeError(f'SQ Thread Tracing is not supported on gc:{self.target}')
-      if not self.is_am() and (ppfeaturemask:=int(FileIOInterface('/sys/module/amdgpu/parameters/ppfeaturemask', os.O_RDONLY).read(), 16))&0x8000:
-        raise RuntimeError("SQTT can't be enabled because of hardware bug, to workaround either use AMD_IFACE=PCI or add "
-                           f"ppfeaturemask={(ppfeaturemask&~0x8000):#x} (current {ppfeaturemask=:#x} & ~PP_GFXOFF_MASK) to amdgpu module parameters\n"
-                           "For more information read https://github.com/tinygrad/tinygrad/blob/master/extra/sqtt/README.md")
+      if not self.iface.is_in_profile_mode(): raise RuntimeError("SQTT requires stable power state: run `amd-smi set -l stable_std` for KFD iface")
+
       SQTT_BUFFER_SIZE = getenv("SQTT_BUFFER_SIZE", 256) # in mb, per shader engine
       self.sqtt_buffers = [self.allocator.alloc(SQTT_BUFFER_SIZE << 20, BufferSpec(nolru=True, uncached=True)) for _ in range(self.se_cnt)]
       self.sqtt_itrace_se_mask = getenv("SQTT_ITRACE_SE_MASK", -1 if SQTT >= 2 else (1 << 1)) # se bitmask: -1 enable all, 0 disable all

tinygrad/runtime/ops_null.py

@@ -1,5 +1,4 @@
 import functools
-from typing import cast
 from tinygrad.device import Compiled, Compiler, Allocator
 from tinygrad.engine.jit import MultiGraphRunner
 from tinygrad.renderer.cstyle import Renderer, CStyleLanguage
@@ -33,7 +32,7 @@ class NullGraph(MultiGraphRunner):
 class NullDevice(Compiled):
   def __init__(self, device:str):
     renderer:functools.partial|type[Renderer]
-    match cast(str, EMULATE.value):
+    match str(EMULATE.value):
       case "AMD": renderer = functools.partial(AMDLLVMRenderer, "gfx1100")
       case "AMD_RDNA4": renderer = functools.partial(AMDLLVMRenderer, "gfx1201")
       case "": renderer = NullRenderer

tinygrad/runtime/ops_tinyfs.py

@@ -13,6 +13,7 @@ class TinyFSDevice(Compiled):
 
     self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
     self.sock.connect((TINYFS_ENDPOINT.rsplit(":", 1)[0], int(TINYFS_ENDPOINT.rsplit(":", 1)[1])))
+    self.sock.settimeout(10)
     self.sfile = self.sock.makefile("rwb")
 
     # fetch node info
@@ -112,9 +113,9 @@ class TinyFSAllocator(Allocator[TinyFSDevice]):
 
         writer.write(f"CHUNK_OUT {size}\r\n".encode())
         writer.write(src.hash_buf[i*16:(i+1)*16])
-        await writer.drain()
+        await asyncio.wait_for(writer.drain(), timeout=10)
 
-        chunk = await reader.readexactly(size)
+        chunk = await asyncio.wait_for(reader.readexactly(size), timeout=10)
 
         view = dest[ptr:ptr+len(chunk)]
         view[:] = chunk

tinygrad/runtime/support/am/ip.py

@@ -272,7 +272,7 @@ class AM_GFX(AM_IP):
     self.adev.regSDMA0_RLC_CGCG_CTRL.update(cgcg_int_enable=1)
     self.adev.regSDMA1_RLC_CGCG_CTRL.update(cgcg_int_enable=1)
 
-    self.adev.regRLC_CGTT_MGCG_OVERRIDE.update(perfmon_clock_state=0, gfxip_fgcg_override=0, gfxip_repeater_fgcg_override=0,
+    self.adev.regRLC_CGTT_MGCG_OVERRIDE.update(perfmon_clock_state=1, gfxip_fgcg_override=0, gfxip_repeater_fgcg_override=0,
       grbm_cgtt_sclk_override=0, rlc_cgtt_sclk_override=0, gfxip_mgcg_override=0, gfxip_cgls_override=0, gfxip_cgcg_override=0)
 
     self.adev.regRLC_SAFE_MODE.write(message=0, cmd=1)

tinygrad/runtime/support/amd.py

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

tinygrad/schedule/indexing.py

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

tinygrad/schedule/rangeify.py

@@ -1,11 +1,11 @@
-from typing import cast
 from dataclasses import dataclass, field
+import itertools
 from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, ssimplify, KernelInfo
-from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType, BottomUpGate
+from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType, BottomUpGate, Kernel, _remove_all_tags
 from tinygrad.uop.symbolic import symbolic_flat
-from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, Metadata, DEBUG_RANGEIFY
-from tinygrad.helpers import PCONTIG, partition, get_single_element
+from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY
+from tinygrad.helpers import PCONTIG, partition, get_single_element, unwrap
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
 from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, ALWAYS_CONTIGUOUS, IndexingContext, apply_movement_op
@@ -18,6 +18,10 @@ sys.setrecursionlimit(10000)
 pm_mops = PatternMatcher([
   (UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
    lambda r,idx: r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idx.src[1:]), dtype=idx.dtype, arg=idx.arg)),  # type: ignore
+  # move movement ops after AFTER
+  (UPat(GroupOp.Movement, name="r").after(name="a", allow_any_len=True),
+   lambda r,a: UOp(r.op, r.dtype, (a.replace(src=(r.src[0],)+a.src[1:], tag=None),)+r.src[1:], r.arg, tag=a.tag)),
+  (UPat(GroupOp.Movement, name="r").end(name="a", allow_any_len=True), lambda r,a: a.replace(src=(r.src[0],)+a.src[1:])),
 ])
 
 # *****************
@@ -299,7 +303,7 @@ pm_limit_bufs = PatternMatcher([(UPat(set.union(GroupOp.Binary, GroupOp.Ternary)
 # BUFFERIZE returns the BUFFER ready for INDEXing (doing this will make splitting a lot easier)
 # NOTE: this has been fixed up a bit
 
-def bufferize_to_store(x:UOp, idx:UOp, allow_locals=True):
+def bufferize_to_store(ctx:itertools.count|None, x:UOp, idx:UOp, allow_locals=True):
   #assert isinstance(x.tag, Flat), "bufferize must be flat"
   size = prod(x.shape)
   rngs = sorted(idx.ranges, key=lambda x: x.arg)
@@ -329,9 +333,7 @@ def bufferize_to_store(x:UOp, idx:UOp, allow_locals=True):
 
   if allow_locals:
     # handle locals
-    tag = x.arg.device
-    if tag is None: tag = UOp.unique().arg # TODO: hack
-    buf = UOp(Ops.DEFINE_LOCAL, sdtype, arg=tag)
+    buf = UOp(Ops.DEFINE_LOCAL, sdtype, arg=next(unwrap(ctx)))
     do_store = buf.broadcast(x.src[1].dtype.count).index(idx, dtype=sdtype).store(x.src[0]).end(*rngs)
     return buf.after(do_store.barrier())
 
@@ -348,11 +350,14 @@ def flatten_bufferize(x:UOp):
 pm_flatten_bufferize = PatternMatcher([(UPat(Ops.BUFFERIZE, name="x"), flatten_bufferize)])
 
 pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
-  (UPat(Ops.BUFFERIZE, src=(UPat(), UPat(name="idx")), name="x"), lambda x, idx: bufferize_to_store(x, idx, allow_locals=False)),
+  (UPat(Ops.BUFFERIZE, src=(UPat(), UPat(name="idx")), name="x"), lambda x, idx: bufferize_to_store(None, x, idx, allow_locals=False)),
 
   # move RESHAPEs through MSELECT/MSTACK
   (UPat((Ops.MSELECT, Ops.MSTACK), src=UPat(Ops.RESHAPE), name="m"),
    lambda m: m.replace(src=tuple([x.src[0].base for x in m.src]), tag=None).reshape(m.shape).rtag(m.tag)),
+
+  # remove any RESHAPEs on KERNEL
+  (UPat(Ops.KERNEL, name="k"), lambda k: k.replace(src=tuple(x.src[0] if x.op is Ops.RESHAPE else x for x in k.src))),
 ])
 
 pm_add_buffers_local = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
@@ -456,19 +461,13 @@ def remove_metadata_tags(ctx:LocalAddBufferContext, x:UOp):
   return x.replace(tag=None)
 
 pm_remove_tags = PatternMatcher([
-  # remove all the tags
   (UPat(GroupOp.All, name="x"), remove_metadata_tags),
 ])
 
 pm_add_range_tags = PatternMatcher([
-  (UPat(Ops.RANGE, name="x"), lambda x: x.rtag(()))
+  (UPat(Ops.RANGE, name="x"), lambda x: x.rtag(())),
 ])
 
-@dataclass(frozen=True)
-class Kernel:
-  ast: UOp
-  metadata: tuple[Metadata, ...] = ()
-
 def split_store(ctx:list[UOp], x:UOp) -> UOp|None:
   if len(x.ranges): return None
 
@@ -505,7 +504,7 @@ def tag_uop(ctx:list[UOp], x:UOp):
   return x.replace(tag=(len(ctx)-1,))
 add_tags = PatternMatcher([
   # don't tag BUFFERs, they are global
-  (UPat(GroupOp.All-{Ops.BUFFER, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.DEFINE_VAR, Ops.BIND,
+  (UPat(GroupOp.All-{Ops.BUFFER, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.DEFINE_VAR, Ops.BIND, Ops.KERNEL,
                      Ops.MSTACK, Ops.MSELECT, Ops.RANGE}.union(GroupOp.Movement), name="x"), tag_uop),
   (UPat({Ops.MSTACK, Ops.MSELECT}, name="x"), lambda ctx,x: None if all(s.op is Ops.BUFFER for s in x.src) else tag_uop(ctx, x)),
 ])
@@ -532,7 +531,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   uop_list: list[UOp] = []
   tsink = graph_rewrite(sink, add_tags, ctx=uop_list, bottom_up=True, name="number the uops")
 
-  tsink = graph_rewrite(tsink, earliest_rewrites+replace_contiguous, ctx={}, name="earliest rewrites")
+  tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites+replace_contiguous, ctx={}, name="earliest rewrites")
 
   # convert movement ops to ranges
   tsink, rctx = run_rangeify(tsink, DEBUG_RANGEIFY)
@@ -544,7 +543,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   # rebuild the sink with all the BUFFERIZEs with tags, this is what's ending up in the tensor graph
   # MSTACK stacks multiple BUFFERIZEs in one tagged tensor
   # if it's not tagged by here, it's out
-  tsink = UOp.sink(*[x for x in tsink.backward_slice if x.base.op in {Ops.BUFFERIZE, Ops.MSTACK, Ops.CONST, Ops.BUFFER} and \
+  tsink = UOp.sink(*[x for x in tsink.backward_slice if x.base.op in {Ops.BUFFERIZE, Ops.MSTACK, Ops.CONST, Ops.BUFFER, Ops.AFTER} and \
                      x.tag is not None and len(x.tag)])
 
   if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Tagged Rangeify")
@@ -569,10 +568,14 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
 
   if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
 
+  # TODO: we can probably get this earlier
+  sink_tags = [s.tag for s in tsink.src]
+  tsink = graph_rewrite(tsink, _remove_all_tags, name="remove all tags")
+
   becomes_map: dict[UOp, UOp] = {}
-  for s in tsink.src:
-    assert s.tag is not None
-    for a in s.tag:
+  for tag, s in zip(sink_tags, tsink.src):
+    assert tag is not None
+    for a in tag:
       if a is None: continue
-      becomes_map[uop_list[cast(int, a)]] = s.replace(tag=None)
+      becomes_map[uop_list[int(a)]] = s
   return becomes_map

tinygrad/tensor.py

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

tinygrad/uop/mathtraits.py

@@ -1,173 +0,0 @@
-from typing import TypeVar
-from tinygrad.uop import Ops
-from tinygrad.dtype import dtypes, ConstType
-
-TMT = TypeVar("TMT", bound="MathTrait")
-class MathTrait:
-  # required to implement
-  def alu(self:TMT, op:Ops, *src:TMT) -> TMT: raise NotImplementedError
-  def const_like(self:TMT, b:ConstType) -> TMT: raise NotImplementedError
-
-  # great functions you get!
-  def ufix(self:TMT, x:TMT|ConstType) -> TMT: return self.const_like(x) if not isinstance(x, MathTrait) else x
-  def _binop(self:TMT, op:Ops, x:TMT|ConstType, reverse:bool) -> TMT:
-    return self.ufix(x).alu(op, self) if reverse else self.alu(op, self.ufix(x))
-  def logical_not(self): return self.ne(True)
-  def neg(self):
-    if (dtype:=getattr(self, 'dtype')) is None: raise TypeError(f"MathTraits __neg__ requires a dtype, {self=}")
-    return self.logical_not() if dtype.scalar() == dtypes.bool else self*(-1)
-  def _check_dtype(self):
-    if (dtype:=getattr(self, 'dtype')) is not None:
-      if isinstance(dtype, tuple): dtype = dtype[0]
-      if not (dtypes.is_bool(dtype) or dtypes.is_int(dtype)): raise RuntimeError(f"{dtype} is not supported")
-  def add(self:TMT, x:TMT|ConstType, reverse:bool=False):
-    """
-    Adds `self` and `x`.
-    Equivalent to `self + x`.
-    Supports broadcasting to a common shape, type promotion, and integer, float, boolean inputs.
-    ```python exec="true" source="above" session="tensor" result="python"
-    Tensor.manual_seed(42)
-    t = Tensor.randn(4)
-    print(t.numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.add(20).numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.add(Tensor([[2.0], [3.5]])).numpy())
-    ```
-    """
-    return self._binop(Ops.ADD, x, reverse)
-  def mul(self:TMT, x:TMT|ConstType, reverse:bool=False):
-    """
-    Multiplies `self` and `x`.
-    Equivalent to `self * x`.
-    Supports broadcasting to a common shape, type promotion, and integer, float, boolean inputs.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    Tensor.manual_seed(42)
-    t = Tensor.randn(4)
-    print(t.numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.mul(3).numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.mul(Tensor([[-1.0], [2.0]])).numpy())
-    ```
-    """
-    return self._binop(Ops.MUL, x, reverse)
-  def bitwise_and(self:TMT, x:TMT|ConstType, reverse:bool=False):
-    """
-    Computes the bitwise AND of `self` and `x`.
-    Equivalent to `self & x`.
-    Supports broadcasting to a common shape, type promotion, and integer, boolean inputs.
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor([2, 5, 255]).bitwise_and(Tensor([3, 14, 16])).numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor([True, True, False, False]).bitwise_and(Tensor([True, False, True, False])).numpy())
-    ```
-    """
-    self._check_dtype()
-    return self._binop(Ops.AND, x, reverse)
-  def bitwise_or(self:TMT, x:TMT|ConstType, reverse:bool=False):
-    """
-    Computes the bitwise OR of `self` and `x`.
-    Equivalent to `self | x`.
-    Supports broadcasting to a common shape, type promotion, and integer, boolean inputs.
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor([2, 5, 255]).bitwise_or(Tensor([4, 4, 4])).numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor([True, True, False, False]).bitwise_or(Tensor([True, False, True, False])).numpy())
-    ```
-    """
-    self._check_dtype()
-    return self._binop(Ops.OR, x, reverse)
-  def bitwise_xor(self:TMT, x:TMT|ConstType, reverse:bool=False):
-    """
-    Computes bitwise xor of `self` and `x`.
-    Equivalent to `self ^ x`.
-    Supports broadcasting to a common shape, type promotion, and integer, boolean inputs.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor([-1, -2, 3]).bitwise_xor(Tensor([1, 0, 3])).numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor([True, True, False, False]).bitwise_xor(Tensor([True, False, True, False])).numpy())
-    ```
-    """
-    self._check_dtype()
-    return self._binop(Ops.XOR, x, reverse)
-  def idiv(self:TMT, x:TMT|ConstType, reverse:bool=False):
-    """
-    Divides `self` by `x`.
-    Equivalent to `self // x`.
-    Supports broadcasting to a common shape, type promotion, and integer inputs.
-    `idiv` performs integer division (truncate towards zero).
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor([-4, 7, 5, 4, -7, 8]).idiv(Tensor([2, -3, 8, -2, 3, 5])).numpy())
-    ```
-    """
-    return self._binop(Ops.IDIV, x, reverse)
-  def mod(self:TMT, x:TMT|ConstType, reverse:bool=False): return self._binop(Ops.MOD, x, reverse)
-  def sub(self:TMT, x:TMT|ConstType, reverse:bool=False): return self.ufix(x).alu(Ops.ADD, -self) if reverse else self.alu(Ops.ADD, self.ufix(-x))
-  def div(self:TMT, x:TMT|ConstType, reverse:bool=False):
-    return (self.ufix(x)*self.alu(Ops.RECIPROCAL)) if reverse else (self*self.ufix(x).alu(Ops.RECIPROCAL))
-
-  def __neg__(self): return self.neg()
-
-  def __add__(self:TMT, x:TMT|ConstType): return self.add(x)
-  def __sub__(self:TMT, x:TMT|ConstType): return self.sub(x)
-  def __mul__(self:TMT, x:TMT|ConstType): return self.mul(x)
-  def __truediv__(self:TMT, x:TMT|ConstType): return self.div(x)
-  def __floordiv__(self:TMT, x:TMT|ConstType): return self.idiv(x)  # TODO: idiv is trunc div, not floordiv
-  def __mod__(self:TMT, x:TMT|ConstType): return self.mod(x)
-  def __and__(self:TMT, x:TMT|ConstType): return self.bitwise_and(x)
-  def __or__(self:TMT, x:TMT|ConstType): return self.bitwise_or(x)
-  def __xor__(self:TMT, x:TMT|ConstType): return self.bitwise_xor(x)
-
-  def __radd__(self:TMT, x:TMT|ConstType): return self.add(x, True)
-  def __rsub__(self:TMT, x:TMT|ConstType): return self.sub(x, True)
-  def __rmul__(self:TMT, x:TMT|ConstType): return self.mul(x, True)
-  def __rtruediv__(self:TMT, x:TMT|ConstType): return self.div(x, True)
-  def __rfloordiv__(self:TMT, x:TMT|ConstType): return self.idiv(x, True)
-  def __rand__(self:TMT, x:TMT|ConstType): return self.bitwise_and(x, True)
-  def __ror__(self:TMT, x:TMT|ConstType): return self.bitwise_or(x, True)
-  def __rxor__(self:TMT, x:TMT|ConstType): return self.bitwise_xor(x, True)
-  def __rmod__(self:TMT, x:TMT|ConstType): return self.mod(x, True)
-
-  def __lt__(self:TMT, x:TMT|ConstType): return self.alu(Ops.CMPLT, self.ufix(x))
-  def __gt__(self:TMT, x:TMT|ConstType): return self.ufix(x).alu(Ops.CMPLT, self)
-  def __ge__(self:TMT, x:TMT|ConstType): return (self < x).logical_not()
-  def __le__(self:TMT, x:TMT|ConstType): return (self > x).logical_not()
-
-  def ne(self:TMT, x:TMT|ConstType): return self.alu(Ops.CMPNE, self.ufix(x))
-  def eq(self:TMT, x:TMT|ConstType): return self.ne(x).logical_not()
-  def __ne__(self:TMT, x:TMT|ConstType): return self.ne(x)  # type: ignore[override]
-  # NOTE: __eq__ isn't overridden, and means the same thing as is by default
-
-  def lshift(self:TMT, x:TMT|int, reverse:bool=False): return self._binop(Ops.SHL, x, reverse)
-  def rshift(self:TMT, x:TMT|int, reverse:bool=False): return self._binop(Ops.SHR, x, reverse)
-  def __lshift__(self:TMT, x:TMT|int): return self.lshift(x)
-  def __rshift__(self:TMT, x:TMT|int): return self.rshift(x)
-  def __rlshift__(self:TMT, x:TMT|int): return self.lshift(x, True)
-  def __rrshift__(self:TMT, x:TMT|int): return self.rshift(x, True)
-
-  def maximum(self:TMT, x:TMT|ConstType): return self.alu(Ops.MAX, self.ufix(x))
-  def minimum(self:TMT, x:TMT|ConstType): return -(-self).maximum(-x)
-  def where(self:TMT, x:TMT|ConstType, y:TMT|ConstType):
-    if isinstance(x, type(self)): return self.alu(Ops.WHERE, x, x.ufix(y))
-    if isinstance(y, type(self)): return self.alu(Ops.WHERE, y.ufix(x), y)
-    raise RuntimeError("where needs at least one UOp arg")
-  def threefry(self:TMT, seed:TMT): return self.alu(Ops.THREEFRY, seed)
-  def reciprocal(self): return self.alu(Ops.RECIPROCAL)
-  def trunc(self): return self.alu(Ops.TRUNC)
-  def sqrt(self): return self.alu(Ops.SQRT)
-  def sin(self): return self.alu(Ops.SIN)
-  def log2(self): return self.alu(Ops.LOG2)
-  def exp2(self): return self.alu(Ops.EXP2)
-  def pow(self:TMT, x:TMT|ConstType): return self.alu(Ops.POW, self.ufix(x))
-  def __pow__(self:TMT, x:TMT|ConstType): return self.pow(x)

tinygrad/uop/mixins.py

@@ -0,0 +1,206 @@
+# mixins add syntactic sugar to Tensor and UOp
+from typing import TypeAlias, TYPE_CHECKING, Self
+from tinygrad.uop import Ops
+from tinygrad.dtype import dtypes, ConstType
+from tinygrad.helpers import prod, argfix
+if TYPE_CHECKING:
+  from tinygrad.uop.ops import UOp
+  sint:TypeAlias = UOp|int
+
+class MathMixin:
+  # required to implement
+  def alu(self, op:Ops, *src:Self) -> Self: raise NotImplementedError
+  def const_like(self, b:ConstType) -> Self: raise NotImplementedError
+
+  # great functions you get!
+  def ufix(self, x:Self|ConstType) -> Self: return self.const_like(x) if not isinstance(x, MathMixin) else x
+  def _binop(self, op:Ops, x:Self|ConstType, reverse:bool) -> Self:
+    return self.ufix(x).alu(op, self) if reverse else self.alu(op, self.ufix(x))
+  def logical_not(self): return self.ne(True)
+  def neg(self):
+    if (dtype:=getattr(self, 'dtype')) is None: raise TypeError(f"MathTraits __neg__ requires a dtype, {self=}")
+    return self.logical_not() if dtype.scalar() == dtypes.bool else self*(-1)
+  def _check_dtype(self):
+    if (dtype:=getattr(self, 'dtype')) is not None:
+      if isinstance(dtype, tuple): dtype = dtype[0]
+      if not (dtypes.is_bool(dtype) or dtypes.is_int(dtype)): raise RuntimeError(f"{dtype} is not supported")
+  def add(self, x:Self|ConstType, reverse:bool=False):
+    """
+    Adds `self` and `x`.
+    Equivalent to `self + x`.
+    Supports broadcasting to a common shape, type promotion, and integer, float, boolean inputs.
+    ```python exec="true" source="above" session="tensor" result="python"
+    Tensor.manual_seed(42)
+    t = Tensor.randn(4)
+    print(t.numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(t.add(20).numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(t.add(Tensor([[2.0], [3.5]])).numpy())
+    ```
+    """
+    return self._binop(Ops.ADD, x, reverse)
+  def mul(self, x:Self|ConstType, reverse:bool=False):
+    """
+    Multiplies `self` and `x`.
+    Equivalent to `self * x`.
+    Supports broadcasting to a common shape, type promotion, and integer, float, boolean inputs.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    Tensor.manual_seed(42)
+    t = Tensor.randn(4)
+    print(t.numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(t.mul(3).numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(t.mul(Tensor([[-1.0], [2.0]])).numpy())
+    ```
+    """
+    return self._binop(Ops.MUL, x, reverse)
+  def bitwise_and(self, x:Self|ConstType, reverse:bool=False):
+    """
+    Computes the bitwise AND of `self` and `x`.
+    Equivalent to `self & x`.
+    Supports broadcasting to a common shape, type promotion, and integer, boolean inputs.
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([2, 5, 255]).bitwise_and(Tensor([3, 14, 16])).numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([True, True, False, False]).bitwise_and(Tensor([True, False, True, False])).numpy())
+    ```
+    """
+    self._check_dtype()
+    return self._binop(Ops.AND, x, reverse)
+  def bitwise_or(self, x:Self|ConstType, reverse:bool=False):
+    """
+    Computes the bitwise OR of `self` and `x`.
+    Equivalent to `self | x`.
+    Supports broadcasting to a common shape, type promotion, and integer, boolean inputs.
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([2, 5, 255]).bitwise_or(Tensor([4, 4, 4])).numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([True, True, False, False]).bitwise_or(Tensor([True, False, True, False])).numpy())
+    ```
+    """
+    self._check_dtype()
+    return self._binop(Ops.OR, x, reverse)
+  def bitwise_xor(self, x:Self|ConstType, reverse:bool=False):
+    """
+    Computes bitwise xor of `self` and `x`.
+    Equivalent to `self ^ x`.
+    Supports broadcasting to a common shape, type promotion, and integer, boolean inputs.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-1, -2, 3]).bitwise_xor(Tensor([1, 0, 3])).numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([True, True, False, False]).bitwise_xor(Tensor([True, False, True, False])).numpy())
+    ```
+    """
+    self._check_dtype()
+    return self._binop(Ops.XOR, x, reverse)
+  def idiv(self, x:Self|ConstType, reverse:bool=False):
+    """
+    Divides `self` by `x`.
+    Equivalent to `self // x`.
+    Supports broadcasting to a common shape, type promotion, and integer inputs.
+    `idiv` performs integer division (truncate towards zero).
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-4, 7, 5, 4, -7, 8]).idiv(Tensor([2, -3, 8, -2, 3, 5])).numpy())
+    ```
+    """
+    return self._binop(Ops.IDIV, x, reverse)
+  def mod(self, x:Self|ConstType, reverse:bool=False): return self._binop(Ops.MOD, x, reverse)
+  def sub(self, x:Self|ConstType, reverse:bool=False): return self.ufix(x).alu(Ops.ADD, -self) if reverse else self.alu(Ops.ADD, self.ufix(-x))
+  def div(self, x:Self|ConstType, reverse:bool=False):
+    return (self.ufix(x)*self.alu(Ops.RECIPROCAL)) if reverse else (self*self.ufix(x).alu(Ops.RECIPROCAL))
+
+  def __neg__(self): return self.neg()
+
+  def __add__(self, x:Self|ConstType): return self.add(x)
+  def __sub__(self, x:Self|ConstType): return self.sub(x)
+  def __mul__(self, x:Self|ConstType): return self.mul(x)
+  def __truediv__(self, x:Self|ConstType): return self.div(x)
+  def __floordiv__(self, x:Self|ConstType): return self.idiv(x)  # TODO: idiv is trunc div, not floordiv
+  def __mod__(self, x:Self|ConstType): return self.mod(x)
+  def __and__(self, x:Self|ConstType): return self.bitwise_and(x)
+  def __or__(self, x:Self|ConstType): return self.bitwise_or(x)
+  def __xor__(self, x:Self|ConstType): return self.bitwise_xor(x)
+
+  def __radd__(self, x:Self|ConstType): return self.add(x, True)
+  def __rsub__(self, x:Self|ConstType): return self.sub(x, True)
+  def __rmul__(self, x:Self|ConstType): return self.mul(x, True)
+  def __rtruediv__(self, x:Self|ConstType): return self.div(x, True)
+  def __rfloordiv__(self, x:Self|ConstType): return self.idiv(x, True)
+  def __rand__(self, x:Self|ConstType): return self.bitwise_and(x, True)
+  def __ror__(self, x:Self|ConstType): return self.bitwise_or(x, True)
+  def __rxor__(self, x:Self|ConstType): return self.bitwise_xor(x, True)
+  def __rmod__(self, x:Self|ConstType): return self.mod(x, True)
+
+  def __lt__(self, x:Self|ConstType): return self.alu(Ops.CMPLT, self.ufix(x))
+  def __gt__(self, x:Self|ConstType): return self.ufix(x).alu(Ops.CMPLT, self)
+  def __ge__(self, x:Self|ConstType): return (self < x).logical_not()
+  def __le__(self, x:Self|ConstType): return (self > x).logical_not()
+
+  def ne(self, x:Self|ConstType): return self.alu(Ops.CMPNE, self.ufix(x))
+  def eq(self, x:Self|ConstType): return self.ne(x).logical_not()
+  def __ne__(self, x:Self|ConstType): return self.ne(x)  # type: ignore[override]
+  # NOTE: __eq__ isn't overridden, and means the same thing as is by default
+
+  def lshift(self, x:Self|int, reverse:bool=False): return self._binop(Ops.SHL, x, reverse)
+  def rshift(self, x:Self|int, reverse:bool=False): return self._binop(Ops.SHR, x, reverse)
+  def __lshift__(self, x:Self|int): return self.lshift(x)
+  def __rshift__(self, x:Self|int): return self.rshift(x)
+  def __rlshift__(self, x:Self|int): return self.lshift(x, True)
+  def __rrshift__(self, x:Self|int): return self.rshift(x, True)
+
+  def maximum(self, x:Self|ConstType): return self.alu(Ops.MAX, self.ufix(x))
+  def minimum(self, x:Self|ConstType): return -(-self).maximum(-x)
+  def where(self, x:Self|ConstType, y:Self|ConstType):
+    if isinstance(x, type(self)): return self.alu(Ops.WHERE, x, x.ufix(y))
+    if isinstance(y, type(self)): return self.alu(Ops.WHERE, y.ufix(x), y)
+    raise RuntimeError("where needs at least one UOp arg")
+  def threefry(self, seed:Self): return self.alu(Ops.THREEFRY, seed)
+  def reciprocal(self): return self.alu(Ops.RECIPROCAL)
+  def trunc(self): return self.alu(Ops.TRUNC)
+  def sqrt(self): return self.alu(Ops.SQRT)
+  def sin(self): return self.alu(Ops.SIN)
+  def log2(self): return self.alu(Ops.LOG2)
+  def exp2(self): return self.alu(Ops.EXP2)
+  def pow(self, x:Self|ConstType): return self.alu(Ops.POW, self.ufix(x))
+  def __pow__(self, x:Self|ConstType): return self.pow(x)
+
+class MovementMixin:
+  # required to implement
+  def _mop(self, op:Ops, arg) -> Self: raise NotImplementedError
+  @property
+  def shape(self) -> tuple["sint", ...]: raise NotImplementedError
+
+  # great functions you get!
+  def view(self, shape, *args) -> Self:
+    """`.view` is an alias for `.reshape`."""
+    return self.reshape(shape, *args)
+
+  def reshape(self, shape, *args) -> Self:
+    """
+    Returns a tensor with the same data as the original tensor but with a different shape.
+    `shape` can be passed as a tuple or as separate arguments.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor.arange(6)
+    print(t.reshape(2, 3).numpy())
+    ```
+    """
+    # resolve None and args
+    new_shape = tuple([s if s is not None else self.shape[i] for i,s in enumerate(argfix(shape, *args))])
+    # resolve -1
+    if (c := new_shape.count(-1)) > 1: raise RuntimeError(f"only one dimension can be inferred using -1, getting {new_shape}")
+    if c: new_shape = tuple([-prod(self.shape) // prod(new_shape) if s == -1 else s for s in new_shape])
+    if prod(self.shape) != prod(new_shape): raise ValueError(f"size mismatch, can't reshape ({self.shape}) -> ({new_shape})")
+    return self._mop(Ops.RESHAPE, arg=new_shape) if new_shape != self.shape else self

tinygrad/uop/ops.py

@@ -4,8 +4,8 @@ import sys, time, functools, itertools, math, operator, hashlib, os, types, pick
 from dataclasses import dataclass
 from enum import Enum, auto
 from tinygrad.uop import Ops, GroupOp
-from tinygrad.uop.mathtraits import MathTrait
-from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, truncate, PtrDType, least_upper_dtype, Invalid, InvalidType
+from tinygrad.uop.mixins import MathMixin, MovementMixin
+from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, truncate, PtrDType, least_upper_dtype, Invalid, InvalidType, AddrSpace
 from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA
 from tinygrad.helpers import PICKLE_BUFFERS, PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, VIZ, SPEC, CI
 from tinygrad.helpers import strip_parens, colored
@@ -104,7 +104,7 @@ class recursive_property(property):
 
 # NOTE: this should be frozen, but frozen is slower
 @dataclass(eq=False, slots=True)
-class UOp(MathTrait, metaclass=UOpMetaClass):
+class UOp(MathMixin, MovementMixin, metaclass=UOpMetaClass):
   op:Ops
   dtype:DType = dtypes.void
   src:tuple[UOp, ...] = tuple()
@@ -188,7 +188,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     match self.op:
       # late ops don't have shape
       case Ops.UNIQUE | Ops.DEVICE | Ops.RANGE | Ops.INDEX | Ops.LOAD | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
-           Ops.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.PRECAST:
+           Ops.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.PRECAST | Ops.CONTRACT:
         return None
 
       # some ops init the shape
@@ -338,10 +338,13 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   def group(*srcs:UOp|None):  # pylint: disable=no-self-argument
     if len(srcs) == 1 and isinstance(srcs[0], UOp): return srcs[0]
     return UOp(Ops.GROUP, dtypes.void, tuple([x for x in srcs if x is not None]))
+  def vectorize(self, *srcs, **kwargs):
+    return UOp(Ops.VECTORIZE, self.dtype.vec(len(srcs)+1), (self,)+srcs, **kwargs)
   def detach(self): return UOp(Ops.DETACH, self.dtype, (self,))
   def index(self, *srcs:UOp|None, ptr=False, **kwargs):
     return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype if ptr else self.dtype.base), (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
-  def __getitem__(self, *idx): return self.index(*idx)
+  def __getitem__(self, idx):
+    return self.index(*[UOp.const(dtypes.index, x) if isinstance(x, int) else x for x in argfix(idx)])
   def const_like(self, b:ConstLike):
     # constants can optionally have a DEVICE source
     return UOp.const(self.dtype, b, device=self._device, shape=self._shape)
@@ -372,6 +375,9 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   def after(self, *src:UOp, **kwargs): return UOp(Ops.AFTER, self.dtype, (self,)+src, **kwargs)
   def assign(self, x:UOp): return UOp(Ops.ASSIGN, self.dtype, (self, x))
   def barrier(self, *src:UOp): return UOp(Ops.BARRIER, src=(self,)+src)
+  def contract(self, *rngs:UOp):
+    assert all(x.arg[-1] == AxisType.UPCAST for x in rngs), "all contract ranges must be upcast"
+    return UOp(Ops.CONTRACT, dtype=self.dtype.vec(prod([x.vmax+1 for x in rngs])), src=(self,), arg=tuple((x.arg[0], x.vmax+1) for x in rngs))
   def alu(self, op, *src:UOp, **kwargs):
     out_dtype = (self, *src)[-1].dtype
     if op in {Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ}: out_dtype = dtypes.bool.vec(out_dtype.count) if out_dtype.count > 1 else dtypes.bool
@@ -390,10 +396,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     if shape is not None: ret = ret.reshape((1,)*len(shape)).expand(shape)
     return ret
   @staticmethod
-  def range(end:sint, *arg, dtype=dtypes.index, src=(), **kwargs):
-    if len(arg) == 0: raise RuntimeError("range needs an arg")
-    if len(arg) == 1: arg = arg+(AxisType.LOOP,)
-    return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end, dtype),)+src, arg=arg, **kwargs)
+  def range(end:sint, axis_id, axis_type=AxisType.LOOP, *arg, dtype=dtypes.index, src=(), **kwargs):
+    return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end, dtype),)+src, arg=(axis_id, axis_type)+arg, **kwargs)
   @staticmethod
   def special(end:sint, name:str, dtype=dtypes.index): return UOp(Ops.SPECIAL, dtype=dtype, src=(sint_to_uop(end, dtype),), arg=name)
   def r(self, op:Ops, axis:tuple[int, ...]):
@@ -416,6 +420,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     return self.op is Ops.BUFFER
 
   def contiguous(self, *args, **kwargs):
+    if self.op is Ops.CONTIGUOUS: return self
     if self.is_contiguous(): return self
     return UOp(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs)
   def contiguous_backward(self): return self.alu(Ops.CONTIGUOUS_BACKWARD)
@@ -528,7 +533,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   # in these four, if the shape doesn't change we can return self
   def forced_reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, same_shape_noop=False)
-  def reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, same_shape_noop=True)
+  #def reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, same_shape_noop=True)
   def expand(self, arg:tuple[sint, ...]): return self._mop(Ops.EXPAND, arg, same_shape_noop=True)
   def shrink(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.SHRINK, arg, same_shape_noop=True)
   def pad(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.PAD, arg, same_shape_noop=True)
@@ -580,6 +585,16 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     while len(s.src) and s.op not in {Ops.BUFFER, Ops.BUFFERIZE, Ops.MSTACK}: s = s.src[0]
     return s
 
+  def buf_target(self) -> UOp:
+    # the buffer that's being loaded from or store to
+    match self.op:
+      case Ops.DEFINE_GLOBAL | Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return self
+      case Ops.AFTER | Ops.INDEX | Ops.STORE | Ops.LOAD: return self.src[0].buf_target()
+      case Ops.VECTORIZE:
+        assert all_same(self.src)
+        return self.src[0].buf_target()
+      case _: raise RuntimeError(f"buf_target called on non load/index/store {self.op}")
+
   @property
   def buffer(self) -> Buffer|MultiBuffer:
     from tinygrad.device import Buffer, MultiBuffer
@@ -739,12 +754,38 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     return fxn(**{k:v for k,v in var_vals.items() if k in varnames})
 
   def render(self, simplify=True, pm:PatternMatcher|None=None) -> str:
-    with Context(TRACK_MATCH_STATS=0, SPEC=0):
-      ret = graph_rewrite(self.simplify() if simplify else self, renderer if pm is None else pm)
-    return ret.arg if ret.op is Ops.NOOP else str(ret)
+    ctx: dict[UOp, str] = {}
+    pm = renderer if pm is None else pm
+    for u in (s:=self.simplify() if simplify else self).toposort():
+      ctx[u] = cast(str, pm.rewrite(u, ctx=ctx))
+    return ctx[s]
 
   def pyrender(self): return pyrender(self)
 
+  # *** uop high level syntactic sugar ***
+
+  def shrink_to(self, arg:tuple[sint, ...]): return self.shrink(tuple([(0,x) for x in arg]))
+
+  @staticmethod
+  def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL):
+    lookup = {AddrSpace.GLOBAL: Ops.DEFINE_GLOBAL, AddrSpace.LOCAL: Ops.DEFINE_LOCAL, AddrSpace.REG: Ops.DEFINE_REG}
+    ret = UOp(lookup[addrspace], dtype.ptr(prod(shape), addrspace), arg=slot)
+    if len(shape) > 1: ret = ret.reshape(shape)
+    return ret
+  def placeholder_like(self, slot:int):
+    assert all_int(self.shape), "no placeholder-like on symbolic shape"
+    return UOp.placeholder(self.shape, self.dtype, slot)
+
+  # set is store+end+after
+  def set(self:UOp, val:UOp|ConstType, end:UOp|tuple[UOp, ...]=()) -> UOp:
+    return self.src[0].after(self.store(UOp.const(self.dtype, val) if not isinstance(val, UOp) else val).end(*argfix(end)))
+
+  def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]:
+    placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(srcs)]
+    contig_srcs = tuple(x.contiguous() for x in srcs)
+    kernel = UOp(Ops.KERNEL, src=tuple(x.base for x in contig_srcs), arg=Kernel(fxn(*placeholders), grad_fxn=grad_fxn))
+    return [s.after(kernel) for s in contig_srcs]
+
 @dataclass(frozen=True)
 class KernelInfo:
   name: str = "test"            # name of the kernel
@@ -755,6 +796,12 @@ class KernelInfo:
   @property
   def function_name(self): return to_function_name(self.name)
 
+@dataclass(frozen=True)
+class Kernel:
+  ast: UOp
+  metadata: tuple[Metadata, ...] = ()
+  grad_fxn: Callable|None = None
+
 # ******** ops in python ********
 
 def safe_exp2(x):
@@ -806,7 +853,7 @@ def printable(loc:tuple[str, int]) -> str:
   try: return lines(loc[0])[loc[1]-1].strip()
   except FileNotFoundError: return "<missing>"
 
-class UPat(MathTrait):
+class UPat(MathMixin, MovementMixin):
   __slots__ = ("op", "dtype", "arg", "name", "src")
   def __init__(self, op:Ops|tuple[Ops, ...]|set[Ops]|None=None, dtype:DType|tuple[DType, ...]|None=None,
                src:tuple[UPat, ...]|list[UPat]|UPat|None=None, arg:Any=None,
@@ -873,6 +920,7 @@ class UPat(MathTrait):
   def broadcast(self, **kwargs): return UPat(Ops.VECTORIZE, self.dtype, src=self, **kwargs)
   def contiguous(self, *args, **kwargs): return UPat(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs)
   def after(self, *src:UPat, **kwargs): return UPat(Ops.AFTER, self.dtype, (self,)+src, **kwargs)
+  def end(self, *src:UPat, **kwargs): return UPat(Ops.END, self.dtype, (self,)+src, **kwargs)
 
   def const_like(self, b:ConstLike): return UPat.const(self.dtype, cast(ConstType, b))
   def alu(self, op:Ops, *src:UPat):
@@ -1077,7 +1125,7 @@ if TRACK_MATCH_STATS or PROFILE:
   def launch_viz(env_str:str, data:str):
     os.environ[env_str] = "0"
     os.environ[f"{env_str}_DATA"] = data
-    if not int(os.getenv("VIZ", "0")) and not int(os.getenv("PROFILE", "0")) and not int(os.getenv("SQTT", "0")) and not CI:
+    if not int(os.getenv("VIZ", "0")) and not int(os.getenv("PROFILE", "0")) and not CI:
       args = ['--kernels', getenv("VIZ_DATA", "")] if getenv("VIZ_DATA", "") else []
       args += ['--profile', getenv("PROFILE_DATA", "")] if getenv("PROFILE_DATA", "") else []
       viz_path = pathlib.Path(__file__).resolve().parent.parent / "viz" / "serve.py"
@@ -1214,6 +1262,7 @@ pm_lower_index_dtype = PatternMatcher([
 def _index_to_concrete_int(u:UOp): return graph_rewrite(u.sink(), pm_lower_index_dtype).src[0]
 
 _substitute = PatternMatcher([(UPat(tuple(Ops), name="x"), lambda ctx,x: ctx.get(x,None))])
+_remove_all_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
 
 def do_unbind(ctx:dict[Variable, int], x:UOp):
   v,i = x.unbind()
@@ -1224,31 +1273,37 @@ pm_unbind = PatternMatcher([(UPat(Ops.BIND, name="x"), do_unbind)])
 # for debug
 syms = { Ops.ADD: "+", Ops.SUB: "-", Ops.IDIV: "//", Ops.MOD: "%", Ops.SHL: "<<", Ops.SHR: ">>",
          Ops.MUL: "*", Ops.CMPLT: "<", Ops.CMPNE: "!=", Ops.AND: "&", Ops.OR: "|", Ops.XOR: "^"}
+# comparison operators are not in here because they are chained in python, not left-associative
+precedence = {Ops.MUL:1, Ops.IDIV:1, Ops.MOD:1, Ops.ADD:2, Ops.SUB:2, Ops.SHL:3, Ops.SHR:3, Ops.AND:4, Ops.XOR:5, Ops.OR:6}
+def strip_binary_parens(x:UOp, left:str, right:str, code_for_op) -> str:
+  if x.op not in precedence: return code_for_op(left, right)
+  return code_for_op(strip_parens(left) if precedence.get(x.src[0].op,99)<=precedence[x.op] else left, strip_parens(right) if
+    precedence.get(x.src[1].op,99)<precedence[x.op] else right)
+
 renderer = PatternMatcher([
-  (UPat((Ops.DEFINE_VAR,), name="x"), lambda x: UOp(Ops.NOOP, arg=x.arg[0])),
-  (UPat((Ops.SPECIAL), name="x"), lambda x: UOp(Ops.NOOP, arg=x.arg)),
-  (UPat(Ops.RANGE, name="x"), lambda x: UOp(Ops.NOOP, arg=f"r{range_str(x)}")),
-  (UPat((Ops.CONST, Ops.VCONST), name="x"), lambda x: UOp(Ops.NOOP, arg=str(x.arg))),
-  (UPat(Ops.UNROLL, name="x"), lambda x: UOp(Ops.NOOP, arg=f"UNROLL({x.src[0].arg}, {x.arg})")),
-  (UPat(Ops.CAST, name="x"), lambda x: UOp(Ops.NOOP, arg=f"({str(x.dtype)[7:]})({x.src[0].arg})")),
-  (UPat(Ops.BIND, src=UPat(Ops.NOOP), name="x"), lambda x: x.src[0]),
-  #(UPat(Ops.BIND, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}[={x.src[1].arg}]")),
-  (UPat(Ops.NEG, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"(-{x.src[0].arg})")),
-  (UPat(Ops.RECIPROCAL, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"(1/{x.src[0].arg})")),
-  (UPat(Ops.MAX, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"max({x.src[0].arg}, {x.src[1].arg})")),
-  (UPat(Ops.MULACC, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"({x.src[0].arg}*{x.src[1].arg}+{x.src[2].arg})")),
-  (UPat(Ops.WHERE, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"({x.src[1].arg} if {x.src[0].arg} else {x.src[2].arg})")),
-  (UPat(set(syms.keys()), src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"({x.src[0].arg}{syms[x.op]}{x.src[1].arg})")),
-  (UPat((Ops.INDEX, Ops.BUFFERIZE), name="x"), lambda x:
-   UOp(Ops.NOOP, arg=''.join([f"[{strip_parens(y.arg)}]" for y in x.src[1:]])) if all(y.op is Ops.NOOP for y in x.src[1:]) else None),
-  (UPat(Ops.VECTORIZE, src=UPat(Ops.NOOP), name="x"),
-   lambda x: UOp(Ops.NOOP, arg=f"{{{','.join([y.arg for y in x.src])}}}" if not all_same(x.src) else f"{{{x.src[0].arg}, ...}}")),
+  (UPat((Ops.DEFINE_VAR,), name="x"), lambda x: x.arg[0]),
+  (UPat((Ops.SPECIAL), name="x"), lambda x: x.arg),
+  (UPat(Ops.RANGE, name="x"), lambda x: f"r{range_str(x)}"),
+  (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]]})"),
+  (UPat(Ops.BIND, name="x"), lambda ctx,x: ctx[x.src[0]]),
+  (UPat(Ops.NEG, name="x"), lambda ctx,x: f"(-{ctx[x.src[0]]})"),
+  (UPat(Ops.RECIPROCAL, name="x"), lambda ctx,x: f"(1/{ctx[x.src[0]]})"),
+  (UPat(Ops.MAX, name="x"), lambda ctx,x: f"max({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
+  (UPat(Ops.MULACC, name="x"), lambda ctx,x: f"({ctx[x.src[0]]}*{ctx[x.src[1]]}+{ctx[x.src[2]]})"),
+  (UPat(Ops.WHERE, name="x"), lambda ctx,x: f"({ctx[x.src[1]]} if {ctx[x.src[0]]} else {ctx[x.src[2]]})"),
+  (UPat(set(syms.keys()), name="x"), lambda ctx,x: strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
+  (UPat((Ops.INDEX, Ops.BUFFERIZE), name="x"), lambda x, ctx: ''.join([f"[{strip_parens(ctx[y])}]" for y in x.src[1:]])),
+  (UPat(Ops.VECTORIZE, name="x"),
+   lambda ctx,x: f"{{{','.join([ctx[y] for y in x.src])}}}" if not all_same(x.src) else f"{{{ctx[x.src[0]]}, ...}}"),
+  (UPat(GroupOp.All, name="x"), lambda x: str(x)),
 ])
+
 renderer_infer = PatternMatcher([
-  (UPat(Ops.MOD, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"cmod({x.src[0].arg}, {x.src[1].arg})")),
-  (UPat(Ops.IDIV, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"cdiv({x.src[0].arg}, {x.src[1].arg})")),
-  *renderer.patterns
-])
+  (UPat(Ops.MOD, name="x"), lambda ctx,x: f"cmod({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
+  (UPat(Ops.IDIV, name="x"), lambda ctx,x: f"cdiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
+]) + renderer
 
 # *** pyrender ***
 
@@ -1290,10 +1345,12 @@ pm_pyrender_extra = PatternMatcher([
   (UPat(GroupOp.Movement, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({render_marg(ctx,x)})"),
   # NOTE: CMPNE doesn't work cause there's no __rne__
   (UPat(set(syms.keys())-{Ops.SUB, Ops.CMPNE}, src=(UPat(Ops.CONST, name="y"), UPat(name="z")), name="x"),
-   lambda ctx,x,y,z: f"({y.arg}{syms[x.op]}{ctx[z]})"),
+   lambda ctx,x,y,z: strip_binary_parens(x, str(y.arg), ctx[z], lambda a,b: f"({a}{syms[x.op]}{b})")),
   # NOTE: sub doesn't work cause it's written as add/mul
-  (UPat(set(syms.keys())-{Ops.SUB}, src=(UPat(name="y"), UPat(Ops.CONST, name="z")), name="x"), lambda ctx,x,y,z: f"({ctx[y]}{syms[x.op]}{z.arg})"),
-  (UPat(set(syms.keys())-{Ops.SUB}, name="x"), lambda ctx,x: f"({ctx[x.src[0]]}{syms[x.op]}{ctx[x.src[1]]})"),
+  (UPat(set(syms.keys())-{Ops.SUB}, src=(UPat(name="y"), UPat(Ops.CONST, name="z")), name="x"), lambda ctx,x,y,z:
+    strip_binary_parens(x, ctx[y], str(z.arg), lambda a,b: f"({a}{syms[x.op]}{b})")),
+  (UPat(set(syms.keys())-{Ops.SUB}, name="x"), lambda ctx,x:
+    strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
   (UPat(sugar, src=(), name="x"), lambda x: f"UOp.{x.op.name.lower()}("+', '.join(([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
   (UPat(sugar, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}("+', '.join([ctx[y] for y in x.src[1:]] + \
     ([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
@@ -1339,7 +1396,7 @@ def pyrender(ast:UOp) -> str:
     else:
       r[u] = f"c{i}" if u is not lst[-1] else "ast"
       ret[r[u]] = ren
-  return ''.join([v[1] for v in kernels.values()]) + '\n'.join([f"{k} = {v}" for k,v in ret.items()])
+  return ''.join([v[1] for v in kernels.values()]) + '\n'.join([f"{k} = {strip_parens(v)}" for k,v in ret.items()])
 
 # *** what was symbolic.py ***
 

tinygrad/uop/spec.py

@@ -1,6 +1,6 @@
 import math
 from typing import cast, Any
-from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, AxisType, KernelInfo, pyrender
+from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, AxisType, KernelInfo, pyrender, Kernel
 from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid
 from tinygrad.helpers import DEBUG, Context, prod, SPEC, Metadata
 from tinygrad.uop.validate import validate_index
@@ -44,7 +44,20 @@ shared_spec = PatternMatcher([
 
 # ***** UOp spec in the Tensor graph *****
 
-tensor_spec = PatternMatcher([
+movement_ops = PatternMatcher([
+  (UPat((Ops.RESHAPE, Ops.EXPAND), name="mv", src=(UPat.var("x"), UPat(dtype=dtypes.index))), lambda mv,x: True),
+  (UPat((Ops.PAD, Ops.SHRINK), name="mv", src=(UPat.var("x"), UPat(dtype=dtypes.index), UPat(dtype=dtypes.index))), lambda mv,x: True),
+  (UPat((Ops.PERMUTE, Ops.FLIP), name="mv", src=(UPat.var("x"),)), lambda mv,x: isinstance(mv.arg, tuple)),
+
+  # inputs to movement ops
+  (UPat((Ops.VECTORIZE, Ops.VCONST), dtype=dtypes.index), lambda: True),
+  (UPat({Ops.ADD, Ops.MUL, Ops.IDIV}, dtype=dtypes.index), lambda: True),
+
+  # AFTER on Movement Op
+  (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement),), allow_any_len=True), lambda: True),
+])
+
+_tensor_spec = PatternMatcher([
   # buffer spec
   (UPat(Ops.UNIQUE, dtypes.void, ()), lambda: True),
   (UPat(Ops.DEVICE, dtypes.void, (), name="d"), lambda d:
@@ -56,7 +69,7 @@ tensor_spec = PatternMatcher([
   (UPat(Ops.BUFFER_VIEW, src=(UPat(Ops.MSTACK, src=UPat(Ops.BUFFER)),)), lambda: True),
 
   # KERNEL can attach to an AFTER to describe the compute required to realize a BUFFER
-  (UPat(Ops.KERNEL, src=UPat((Ops.BUFFER, Ops.BUFFER_VIEW, Ops.AFTER, Ops.MSELECT, Ops.MSTACK, Ops.BIND))), lambda: True),
+  (UPat(Ops.KERNEL, src=UPat((Ops.BUFFER, Ops.BUFFER_VIEW, Ops.AFTER, Ops.MSELECT, Ops.MSTACK, Ops.BIND, Ops.CONTIGUOUS))), lambda: True),
 
   # ASSIGN has a target and a value. It can also optionally depend on other assigns
   (UPat(Ops.ASSIGN, name="x"), lambda x: len(x.src) >= 2 and all(s.op is Ops.ASSIGN for s in x.src[2:])),
@@ -67,14 +80,6 @@ tensor_spec = PatternMatcher([
   # MSTACK combines buffers into multi
   (UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(x.device, str) for x in x.src)),
 
-  (UPat((Ops.RESHAPE, Ops.EXPAND), name="mv", src=(UPat.var("x"), UPat(dtype=dtypes.index))), lambda mv,x: True),
-  (UPat((Ops.PAD, Ops.SHRINK), name="mv", src=(UPat.var("x"), UPat(dtype=dtypes.index), UPat(dtype=dtypes.index))), lambda mv,x: True),
-  (UPat((Ops.PERMUTE, Ops.FLIP), name="mv", src=(UPat.var("x"),)), lambda mv,x: isinstance(mv.arg, tuple)),
-
-  # inputs to movement ops
-  (UPat((Ops.VECTORIZE, Ops.VCONST), dtype=dtypes.index), lambda: True),
-  (UPat({Ops.ADD, Ops.MUL, Ops.IDIV}, dtype=dtypes.index), lambda: True),
-
   # Tensor variable bindings
   (UPat(Ops.BIND, (dtypes.int,dtypes.index,), (UPat(Ops.DEFINE_VAR), UPat.cvar(dtype=(dtypes.int,dtypes.index,))), arg=None), lambda: True),
 
@@ -104,7 +109,12 @@ tensor_spec = PatternMatcher([
 
   # AFTER if things were kernelized
   (UPat(Ops.AFTER, src=(UPat((Ops.BUFFER, Ops.AFTER)),), allow_any_len=True), lambda: True),
-])+shared_spec
+])+movement_ops+shared_spec
+
+tensor_spec = PatternMatcher([
+  # no tags allowed in tensor graph
+  (UPat(GroupOp.All, name="x"), lambda x: None if x.tag is None else False),
+])+_tensor_spec
 
 # ***** UOp spec in codegen shared between kernel and program *****
 
@@ -115,7 +125,7 @@ shared_codegen_spec = PatternMatcher([
   (UPat(Ops.DEFINE_REG, src=()), lambda: True),
 
   # allow AFTER on buffers, GROUP anywhere
-  (UPat(Ops.AFTER, src=(UPat(GroupOp.Defines),), allow_any_len=True), lambda: True),
+  (UPat(Ops.AFTER, src=(UPat(GroupOp.Defines|{Ops.AFTER}),), allow_any_len=True), lambda: True),
   (UPat(Ops.GROUP, dtypes.void), lambda: True),
 
   # RANGE/SPECIAL define loops, END closes them
@@ -141,20 +151,36 @@ shared_codegen_spec = PatternMatcher([
   (UPat((Ops.CUSTOMI, Ops.CUSTOM, Ops.PRECAST)), lambda: True),
 
   # INDEX
-  (UPat(GroupOp.Defines, name="buf").or_after().index(UPat.var("idx")), validate_index),
+  (UPat(GroupOp.Defines|{Ops.AFTER}, name="buf").index(UPat.var("idx")), validate_index),
 
   # SPECIAL
   (UPat(Ops.SPECIAL, src=(UPat.var("x", (dtypes.index, dtypes.int32)),), name="s"), lambda s,x: s.dtype == x.dtype and isinstance(s.arg, str)),
 
-  # BARRIER
-  (UPat(Ops.BARRIER, dtypes.void, src=(UPat(),)), lambda: True),
+  # BARRIER (on any length)
+  (UPat(Ops.BARRIER, dtypes.void), lambda: True),
 ])
 
+# ***** UOp spec in kernel graph *****
+
+kernel_spec = PatternMatcher([
+  # index is allowed here
+  (UPat(GroupOp.Elementwise|{Ops.CONST, Ops.RANGE, Ops.DEFINE_VAR}, dtype=dtypes.index), lambda: True),
+
+  # END can end multiple axes here
+  (UPat(Ops.END, src=(UPat(), UPat()), allow_any_len=True, dtype=dtypes.void), lambda: True),
+
+  # bufferize can be on anything
+  (UPat(Ops.BUFFERIZE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: True),
+
+  # reduce must be on ranges
+  (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:])),
+])+movement_ops+shared_codegen_spec+shared_spec
+
 # ***** UOp spec in linearized programs *****
 
 program_spec = PatternMatcher([
   # INDEX with a gate as third src
-  (UPat(Ops.INDEX, src=(UPat(GroupOp.Defines, name="buf").or_after(), UPat.var("idx"), UPat.var("gate", dtype=dtypes.bool))), validate_index),
+  (UPat(Ops.INDEX, src=(UPat(GroupOp.Defines|{Ops.AFTER}, name="buf"), UPat.var("idx"), UPat.var("gate", dtype=dtypes.bool))), validate_index),
 
   # LOAD (idx, alt_value), LOAD can have an alt value, but only if the index has a gate
   (UPat().index(UPat(), UPat(dtype=dtypes.bool)).or_casted().load(UPat()), lambda: True),
@@ -173,22 +199,6 @@ program_spec = PatternMatcher([
   (UPat(Ops.ENDIF, dtype=dtypes.void, src=(UPat(Ops.IF),)), lambda: True),
 ])+shared_codegen_spec+shared_spec
 
-# ***** UOp spec in kernel graph *****
-
-kernel_spec = PatternMatcher([
-  # index is allowed here
-  (UPat(GroupOp.Elementwise|{Ops.CONST, Ops.RANGE, Ops.DEFINE_VAR}, dtype=dtypes.index), lambda: True),
-
-  # END can end multiple axes here
-  (UPat(Ops.END, src=(UPat(), UPat()), allow_any_len=True, dtype=dtypes.void), lambda: True),
-
-  # bufferize can be on anything
-  (UPat(Ops.BUFFERIZE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: True),
-
-  # reduce must be on ranges
-  (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:])),
-])+shared_codegen_spec+shared_spec
-
 # *** this spec should match all UOps ever created ***
 
 full_spec = PatternMatcher([
@@ -241,7 +251,7 @@ full_spec = PatternMatcher([
   (UPat(Ops.DEFINE_VAR, dtype=dtypes.floats), lambda: True),
   # allow any AFTER
   (UPat(Ops.AFTER, src=(UPat(),), allow_any_len=True), lambda: True),
-])+tensor_spec+kernel_spec+program_spec+shared_spec
+])+_tensor_spec+kernel_spec+program_spec+shared_spec
 
 # ***** uop helpers *****
 
@@ -257,7 +267,7 @@ def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher):
 
 # late imports to avoid circular import
 from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.schedule.rangeify import BufferizeOpts, Kernel
+from tinygrad.schedule.rangeify import BufferizeOpts
 glbls:dict[str, Any] = {"inf": math.inf, "nan": math.nan, "KernelInfo": KernelInfo, "Kernel": Kernel, "Metadata": Metadata,
                         "UOp": UOp, "dtypes": dtypes, "Ops": Ops, "AxisType": AxisType, "Invalid": Invalid,
                         "Opt": Opt, "OptOps": OptOps, "BufferizeOpts": BufferizeOpts, "AddrSpace": AddrSpace}

tinygrad/uop/symbolic.py

@@ -452,7 +452,7 @@ def _valid_priority(v: UOp, valids:list[UOp]):
   return sum(-1 if (res:=parse_valid(v)) is not None and res[0] in other.toposort() else 0 for other in valids)
 
 def simplify_valid(valid:UOp) -> UOp|None:
-  if valid.op_in_backward_slice_with_self(Ops.LOAD): return None  # this should only be for indexing, skip if there's a LOAD
+  if valid.op_in_backward_slice_with_self(Ops.INDEX): return None  # this should only be for indexing, skip if there's a INDEX
   ret:list[UOp] = []
   something_changed = False
   valids = list(valid.split_uop(Ops.AND))
@@ -502,7 +502,7 @@ pm_simplify_valid = PatternMatcher([
 ])
 
 # this is symbolic 2.0
-REMOVE_FROM_SINK_LIKE = {Ops.UNROLL, Ops.NOOP}
+REMOVE_FROM_SINK_LIKE = {Ops.UNROLL, Ops.NOOP, Ops.VECTORIZE, Ops.SINK}
 sym = symbolic_flat+pm_simplify_valid+PatternMatcher([
   # LOAD/STORE -> NOOP
   (UPat.var('x').store(UPat.var('x').load(), allow_any_len=True), lambda x: None if x.dtype.addrspace != AddrSpace.REG else x.src[0].src[0]),
@@ -514,9 +514,8 @@ sym = symbolic_flat+pm_simplify_valid+PatternMatcher([
    lambda x,y,alu: UOp(Ops.VECTORIZE, alu.dtype, (UOp(alu.op, alu.dtype.scalar(), (x,y)),)*alu.dtype.count)),
   # VECTORIZE of a single element is just that element
   (UPat(Ops.VECTORIZE, src=(UPat(name='x'),)), lambda x: x),
-  # VECTORIZE void is SINK
-  (UPat(Ops.VECTORIZE, dtype=dtypes.void, src=UPat(Ops.BARRIER, name='b')), lambda b: b),
-  (UPat(Ops.VECTORIZE, dtype=dtypes.void, name='x'), lambda x: UOp(Ops.SINK, dtypes.void, x.src)),
+  # VECTORIZE void is GROUP
+  (UPat(Ops.VECTORIZE, dtype=dtypes.void, name='x'), lambda x: UOp.group(*x.src)),
   # tensor core with a 0 input is acc
   (UPat(Ops.WMMA, src=(UPat.const(None, 0.0), UPat.var(), UPat.var("acc"))), lambda acc: acc),
   (UPat(Ops.WMMA, src=(UPat.var(), UPat.const(None, 0.0), UPat.var("acc"))), lambda acc: acc),

tinygrad/viz/index.html

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

tinygrad/viz/js/index.js

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

tinygrad/viz/js/worker.js

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

tinygrad/viz/serve.py

@@ -14,9 +14,9 @@ from tinygrad.renderer import ProgramSpec
 from tinygrad.dtype import dtypes
 
 uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
-               Ops.DEFINE_GLOBAL: "#ffe0b0", Ops.DEFINE_LOCAL: "#ffe0d0", Ops.DEFINE_REG: "#f0ffe0", Ops.REDUCE_AXIS: "#FF6B6B",
+               Ops.DEFINE_GLOBAL:"#cb9037", **{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: "#e8ffa0", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55",
+               Ops.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55",
                **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
                Ops.BUFFER_VIEW: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.FUSE: "#FFa500",
                Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",
@@ -55,7 +55,7 @@ def pystr(u:UOp, i:int) -> str:
   try: return pyrender(u)
   except Exception: return str(u)
 
-def uop_to_json(x:UOp, ignore_indexing=False) -> dict[int, dict]:
+def uop_to_json(x:UOp) -> dict[int, dict]:
   assert isinstance(x, UOp)
   graph: dict[int, dict] = {}
   excluded: set[UOp] = set()
@@ -63,7 +63,6 @@ def uop_to_json(x:UOp, ignore_indexing=False) -> dict[int, dict]:
     # always exclude DEVICE/CONST/UNIQUE
     if u.op in {Ops.DEVICE, Ops.CONST, Ops.UNIQUE} and u is not x: excluded.add(u)
     if u.op is Ops.VCONST and u.dtype.scalar() == dtypes.index and u is not x: excluded.add(u)
-    if u.dtype.scalar() is dtypes.index and ignore_indexing: excluded.update(u.backward_slice_with_self)
   for u in toposort:
     if u in excluded: continue
     argst = codecs.decode(str(u.arg), "unicode_escape")
@@ -104,16 +103,14 @@ def _reconstruct(a:int):
 def get_full_rewrite(ctx:TrackedGraphRewrite, i:int=0) -> Generator[GraphRewriteDetails, None, None]:
   next_sink = _reconstruct(ctx.sink)
   # in the schedule graph we don't show indexing ops (unless it's in a kernel AST or rewriting dtypes.index sink)
-  ignore_indexing = trace.keys[i].display_name.startswith("Schedule") and not (ctx.name in {"kernel split"} or \
-      any(s.dtype is dtypes.index for s in next_sink.src+(next_sink,)))
-  yield {"graph":uop_to_json(next_sink, ignore_indexing), "uop":pystr(next_sink,i), "changed_nodes":None, "diff":None, "upat":None}
+  yield {"graph":uop_to_json(next_sink), "uop":pystr(next_sink,i), "changed_nodes":None, "diff":None, "upat":None}
   replaces: dict[UOp, UOp] = {}
   for u0_num,u1_num,upat_loc,dur in tqdm(ctx.matches):
     replaces[u0:=_reconstruct(u0_num)] = u1 = _reconstruct(u1_num)
     try: new_sink = next_sink.substitute(replaces)
     except RuntimeError as e: new_sink = UOp(Ops.NOOP, arg=str(e))
     match_repr = f"# {dur*1e6:.2f} us\n"+printable(upat_loc)
-    yield {"graph":(sink_json:=uop_to_json(new_sink, ignore_indexing)), "uop":pystr(new_sink,i),
+    yield {"graph":(sink_json:=uop_to_json(new_sink)), "uop":pystr(new_sink,i),
            "changed_nodes":[id(x) for x in u1.toposort() if id(x) in sink_json],
            "diff":list(difflib.unified_diff(pystr(u0,i).splitlines(),pystr(u1,i).splitlines())), "upat":(upat_loc, match_repr)}
     if not ctx.bottom_up: next_sink = new_sink
@@ -155,7 +152,7 @@ def timeline_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:
       name = ctxs[ref]["name"]
       if isinstance(p:=trace.keys[ref].ret, ProgramSpec) and (ei:=exec_points.get(p.name)) is not None:
         info = f"{sym_infer(p.estimates.ops, ei.arg['var_vals'])/(t:=dur*1e3):.2f} GFLOPS {sym_infer(p.estimates.mem, ei.arg['var_vals'])/t:4.1f}"+ \
-               f"|{sym_infer(p.estimates.lds,ei.arg['var_vals'])/t:.1f} GB/s\n{[str(m) for m in ei.arg['metadata']]}"
+               f"|{sym_infer(p.estimates.lds,ei.arg['var_vals'])/t:.1f} GB/s\n{[str(m) for m in (ei.arg['metadata'] or ())]}"
         key = ei.key
     elif isinstance(e.name, TracingKey):
       name = e.name.display_name