work

the two remaining slow ones

.coveragerc

@@ -0,0 +1,3 @@
+[run]
+source = tinygrad
+branch = True

.github/workflows/benchmark.yml

@@ -199,7 +199,7 @@ jobs:
     - name: Test speed vs torch
       run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
     - name: Test speed vs theoretical
-      run: NV=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      run: NV=1 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
     - name: Test benchmark allreduce
       run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
     - name: Test tensor cores
@@ -409,7 +409,7 @@ jobs:
     #    python3 -c "import torch; print(torch.__version__)"
     #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
     - name: Test speed vs theoretical
-      run: AMD=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      run: AMD=1 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
     - name: Test tensor cores
       run: |
         AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
@@ -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=350 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=390 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
@@ -630,17 +630,17 @@ jobs:
     - name: openpilot compile3 0.9.9 dmonitoring
       run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
     - name: openpilot compile3 0.10.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=5 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
     - name: openpilot compile3 0.10.0 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=13 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
     - name: openpilot compile3 0.10.1 driving_vision
       # TODO: ASSERT_MIN_STEP_TIME=17
-      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=25 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
+      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=21 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
     - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=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_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 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_policy.onnx
     - name: openpilot compile3 0.10.1 dmonitoring
       # TODO: ASSERT_MIN_STEP_TIME=10
-      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=13 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/dmonitoring_model.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 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/dmonitoring_model.onnx
     - name: benchmark MobileNetV2 on DSP
       run: |
         # generate quantized weights

.github/workflows/mlperf.yml

@@ -27,4 +27,4 @@ jobs:
       run: |
         rm "~/.cache/tinygrad/cache_mlperf.db" || true
         BENCHMARK_LOG=mlpert_train_resnet LOGMLPERF=0 CACHEDB="~/.cache/tinygrad/cache_mlperf.db" examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
-        rm "~/.cache/tinygrad/cache_mlperf.db"
+        rm "~/.cache/tinygrad/cache_mlperf.db"

.github/workflows/test.yml

@@ -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
@@ -289,8 +290,8 @@ jobs:
         python extra/optimization/extract_dataset.py
         gzip -c /tmp/sops > extra/datasets/sops.gz
         #DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
-    - name: Repo line count < 18000 lines
-      run: MAX_LINE_COUNT=18000 python sz.py
+    - name: Repo line count < 18500 lines
+      run: MAX_LINE_COUNT=18500 python sz.py
 
   spec:
     strategy:

.gitignore

@@ -63,3 +63,5 @@ profile_stats
 *.log
 target
 .mypy_cache
+mutants
+.mutmut-cache

.pre-commit-config.yaml

@@ -28,7 +28,7 @@ repos:
         pass_filenames: false
       - id: tests
         name: subset of tests
-        entry: env OMP_NUM_THREADS=1 PYTHONPATH="." python3 -m pytest -n=8 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
+        entry: env OMP_NUM_THREADS=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
         language: system
         always_run: true
         pass_filenames: false

autogen_stubs.sh

@@ -31,7 +31,9 @@ $(for p in "$@"; do echo "  $p,"; done)
 ]
 def _try_dlopen_$name():
   library = ctypes.util.find_library("$name")
-  if library: return ctypes.CDLL(library)
+  if library:
+    try: return ctypes.CDLL(library)
+    except OSError: pass
   for candidate in PATHS_TO_TRY:
     try: return ctypes.CDLL(candidate)
     except OSError: pass
@@ -186,6 +188,7 @@ nv_status_codes = {}
     extra/nv_gpu_driver/g_rpc-message-header.h \
     extra/nv_gpu_driver/gsp_static_config.h \
     extra/nv_gpu_driver/vbios.h \
+    extra/nv_gpu_driver/pci_exp_table.h \
     --clang-args="-DRPC_MESSAGE_STRUCTURES -DRPC_STRUCTURES -include $NVKERN_SRC/src/common/sdk/nvidia/inc/nvtypes.h -I$NVKERN_SRC/src/nvidia/generated -I$NVKERN_SRC/src/common/inc -I$NVKERN_SRC/src/nvidia/inc -I$NVKERN_SRC/src/nvidia/interface/ -I$NVKERN_SRC/src/nvidia/inc/kernel -I$NVKERN_SRC/src/nvidia/inc/libraries -I$NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc -I$NVKERN_SRC/kernel-open/nvidia-uvm -I$NVKERN_SRC/kernel-open/common/inc -I$NVKERN_SRC/src/common/sdk/nvidia/inc -I$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include -I$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl" \
     -o $BASE/nv/nv.py
 
@@ -432,11 +435,13 @@ generate_sqtt() {
     $ROCPROF_SRC/include/rocprof_trace_decoder.h \
     $ROCPROF_SRC/include/trace_decoder_instrument.h \
     $ROCPROF_SRC/include/trace_decoder_types.h \
-    -o extra/sqtt/rocprof/rocprof.py
-  fixup extra/sqtt/rocprof/rocprof.py
-  sed -i '1s/^/# pylint: skip-file\n/' extra/sqtt/rocprof/rocprof.py
-  sed -i "s/import ctypes/import ctypes, ctypes.util/g" extra/sqtt/rocprof/rocprof.py
-  sed -i "s|FunctionFactoryStub()|ctypes.CDLL(ctypes.util.find_library('rocprof-trace-decoder'))|g" extra/sqtt/rocprof/rocprof.py
+    -o $BASE/rocprof.py
+  fixup $BASE/rocprof.py
+  sed -i '1s/^/# pylint: skip-file\n/' $BASE/rocprof.py
+  sed -i "s/import ctypes/import ctypes, ctypes.util/g" $BASE/rocprof.py
+  patch_dlopen $BASE/rocprof.py rocprof-trace-decoder "'/usr/local/lib/librocprof-trace-decoder.so'" "'/usr/local/lib/librocprof-trace-decoder.dylib'"
+  sed -i "s/def _try_dlopen_rocprof-trace-decoder():/def _try_dlopen_rocprof_trace_decoder():/g" $BASE/rocprof.py
+  sed -i "s|FunctionFactoryStub()|_try_dlopen_rocprof_trace_decoder()|g" $BASE/rocprof.py
 }
 
 generate_webgpu() {
@@ -545,7 +550,6 @@ elif [ "$1" == "kfd" ]; then generate_kfd
 elif [ "$1" == "nv" ]; then generate_nv
 elif [ "$1" == "amd" ]; then generate_amd
 elif [ "$1" == "am" ]; then generate_am
-elif [ "$1" == "nvdrv" ]; then generate_nvdrv
 elif [ "$1" == "sqtt" ]; then generate_sqtt
 elif [ "$1" == "qcom" ]; then generate_qcom
 elif [ "$1" == "io_uring" ]; then generate_io_uring

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh

@@ -15,7 +15,7 @@ export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"
 
 # pip install -e ".[mlperf]"
-export LOGMLPERF=1
+export LOGMLPERF=${LOGMLPERF:-1}
 
 export SEED=$RANDOM
 DATETIME=$(date "+%m%d%H%M")

examples/openpilot/compile3.py

@@ -4,8 +4,6 @@ import numpy as np
 from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device, dtypes
 from tinygrad.helpers import DEBUG, getenv
 from tinygrad.engine.realize import CompiledRunner
-
-import onnx
 from tinygrad.nn.onnx import OnnxRunner
 
 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
@@ -40,7 +38,7 @@ def compile(onnx_file):
   np.testing.assert_equal(test_val, ret, "JIT run failed")
   print("jit run validated")
 
-  # checks from compile2
+  # check gated read_image usage
   kernel_count = 0
   read_image_count = 0
   gated_read_image_count = 0
@@ -96,6 +94,7 @@ def test_vs_compile(run, inputs, test_val=None):
   return val
 
 def test_vs_onnx(new_inputs, test_val, onnx_file, tol):
+  import onnx
   import onnxruntime as ort
   
   onnx_inputs = {k:v.numpy() for k,v in new_inputs.items()}

examples/whisper.py

@@ -3,7 +3,7 @@
 import sys, base64, multiprocessing, itertools, collections
 from typing import Optional, Union, Literal, List
 
-from tinygrad import Tensor, TinyJit, Variable, nn
+from tinygrad import Tensor, TinyJit, Variable, nn, dtypes
 from tinygrad.nn.state import torch_load, load_state_dict
 from tinygrad.helpers import getenv, fetch
 
@@ -244,15 +244,16 @@ def transcribe_waveform(model: Whisper, enc, waveforms, truncate=False):
 
   log_spec = prep_audio(waveforms, model.batch_size, truncate)
   nsample = model.decoder.max_tokens_to_sample
+  nctx = model.decoder.max_self_attn_cache_len
 
   def inferloop(ctx: Union[np.ndarray, List[np.ndarray]], encoded_audio):
     pos, next_tokens = 0, ctx
-    for i in range((nsample-len(start_tokens))*2):
-      next_tokens = model.decoder(Tensor(next_tokens), pos, encoded_audio)[:, -1].argmax(axis=-1).numpy().astype(np.int32).reshape(-1, 1)
+    for i in range(nsample):
+      next_tokens = model.decoder(Tensor(next_tokens, dtype=dtypes.int32), pos, encoded_audio)[:, -1].argmax(axis=-1).numpy().astype(np.int32).reshape(-1, 1)
       next_tokens[ctx[:, -1] == eot] = eot
       ctx = np.concatenate((ctx, next_tokens), axis=1)
       pos = ctx.shape[-1] - 1
-      if (next_tokens == eot).all(): break
+      if (next_tokens == eot).all() or pos == nctx: break
     return ctx
 
   def gettexttoks(line): return [tok for tok in line if tok < eot or tok > enc._special_tokens["<|notimestamps|>"]][-nsample+len(start_tokens):]

extra/gemm/amd_uop_matmul.py

@@ -1,10 +1,11 @@
 from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
-from tinygrad.uop.ops import UOp, KernelInfo
+from tinygrad.uop.ops import UOp, KernelInfo, sint, AxisType
 from tinygrad.engine.realize import ExecItem, get_runner
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import getenv
 
 N = 4096
+M = K = N
 run_count = 5
 
 # ---------------------------
@@ -42,26 +43,17 @@ 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"
 
+def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.LOOP): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
+def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=False):
+  assert dest.shape == src.shape
+  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
+  copy = dest[*rngs].store(src[*rngs]).end(*rngs)
+  return dest.after(copy) if set else copy
+
 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")
-
-  waveIdx = (tid // WARP_SIZE) % WAVES_IN_BLOCK_X
-  waveIdy = (tid // WARP_SIZE) // WAVES_IN_BLOCK_X
-  assert waveIdy.vmax+1 == WAVES_IN_BLOCK_Y
-
-  idxInWave = (tid % WARP_SIZE) % LANES_PER_WAVE_X
-  idyInWave = (tid % WARP_SIZE) // LANES_PER_WAVE_X
-  assert idyInWave.vmax+1 == LANES_PER_WAVE_Y
-
   # ---------------------------
   # block indices & placeholders
   # ---------------------------
@@ -72,66 +64,66 @@ def hand_spec_kernel3():
   b = UOp.placeholder((N, N), dtypes.float, slot=2)
   c = UOp.placeholder((N, N), dtypes.float, slot=0)
 
-  BM_As_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
-  As = UOp.placeholder((BLOCK_K, BM_As_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
-  Bs = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
+  # index the output with the globals
+  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[blockIdx_y, :, blockIdx_x, :]
 
-  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)
+  # open the main reduction range
+  k_tile_range = UOp.range(N // BLOCK_K, 0, AxisType.REDUCE)
+  a = a.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_K, BLOCK_K)[blockIdx_y, :, k_tile_range, :]
+  b = b.reshape(N // BLOCK_K, BLOCK_K, N // BLOCK_N, BLOCK_N)[k_tile_range, :, blockIdx_x, :]
 
-  i = UOp.range(c_regs.size, 16)
-  c_regs = c_regs[i].set(0.0, end=i)
-
-  k_tile_range = UOp.range(N // BLOCK_K, 0)
+  # globals are no longer used, they are already in the indexes
+  del blockIdx_y, blockIdx_x
 
   # ---------------------------
   # 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)
+  tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
 
-  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)
+  # A: read BM x BK tiles (permute on store into locals)
+  BM_As_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
+  As = UOp.placeholder((BLOCK_K, BM_As_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
+  As_store = copy(As.permute((1,0)).reshape(-1, THREADS_PER_BLOCK)[:, tid], a.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=100)
+
+  # B: read BK x BN tiles
+  Bs = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
+  Bs_store = copy(Bs.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
 
   # TODO: can we automate barrier?
   barrier = UOp.barrier(As_store, Bs_store)
-  Bs = Bs.after(barrier)
-  As = As.after(barrier)
+  As, Bs = As.after(barrier), Bs.after(barrier)
 
   # open inner k range
-  k = UOp.range(BLOCK_K, 3)
+  k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
 
   # ---------------------------
   # 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))
+  waveIdx = (tid // WARP_SIZE) % WAVES_IN_BLOCK_X
+  waveIdy = (tid // WARP_SIZE) // WAVES_IN_BLOCK_X
+  assert waveIdy.vmax+1 == WAVES_IN_BLOCK_Y
 
-  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))
+  laneIdx = (tid % WARP_SIZE) % LANES_PER_WAVE_X
+  laneIdy = (tid % WARP_SIZE) // LANES_PER_WAVE_X
+  assert laneIdy.vmax+1 == LANES_PER_WAVE_Y
+
+  A_col = UOp.placeholder((ITERS_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  A_col = copy(A_col, As[k, :].reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM)[waveIdy, :, laneIdy, :], 300, set=True, upcast=True)
+
+  B_row = UOp.placeholder((ITERS_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+  B_row = copy(B_row, Bs[k, :].reshape(WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)[waveIdx, :, laneIdx, :], 400, set=True, upcast=True)
 
   # ---------------------------
   # FMA: c_regs += A_col * B_row
   # ---------------------------
-  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)
+  c_regs = UOp.placeholder((ITERS_PER_WAVE_M, TM, ITERS_PER_WAVE_N, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  i = UOp.range(c_regs.size, 16)
+  c_regs = c_regs.after(c_regs.flatten()[i].store(0.0).end(i))
+
+  # TODO: why don't these work as upcast?
+  # why if the ranges merge is it slow?!? (if you change the order on end, they will merge. big slowdown on METAL)
+  iterWaveM, yt, iterWaveN, xt = rngs = rngs_for_shape(c_regs.shape, 500)
+  sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iterWaveM, yt] * B_row[iterWaveN, xt]).end(iterWaveM, iterWaveN, yt, xt)
 
   # Close k, sync, and close K tiles
   sink = sink.end(k).barrier().end(k_tile_range)
@@ -139,15 +131,11 @@ def hand_spec_kernel3():
   # ---------------------------
   # REG -> GLOBAL (epilogue)
   # ---------------------------
-  c = c.reshape(N//BLOCK_M, WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM,
-                N//BLOCK_N, WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)
-  iterWaveM = UOp.range(ITERS_PER_WAVE_M, 1000)
-  yt = UOp.range(TM, 1001)
-  iterWaveN = UOp.range(ITERS_PER_WAVE_N, 1002)
-  xt = UOp.range(TN, 1003)
-  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)
+  c = c.reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM,
+                WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)
+  c = c[waveIdy, :, laneIdy, :,
+        waveIdx, :, laneIdx, :]
+  sink = copy(c, c_regs.after(sink), rng=600)
 
   return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()
 

extra/gemm/max_matmul.py

@@ -1,17 +1,11 @@
 import numpy as np, os
 from tinygrad.helpers import getenv, flat_mv
 from tinygrad import dtypes
-from typing import Optional, List, Tuple, cast, Dict, Final, DefaultDict, Self
 from tinygrad.engine.realize import get_program
 
 # for copied uops
-from tinygrad.codegen.opt.kernel import Kernel, KernelOptError
-from tinygrad.uop.ops import UOp, Ops, BinaryOps, UnaryOps, TernaryOps, KernelInfo
-from tinygrad.codegen.opt.search import Opt, OptOps
-from tinygrad import Device, dtypes, Tensor
-from tinygrad.dtype import PtrDType, DType, DTYPES_DICT
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
+from tinygrad import dtypes
+from tinygrad.dtype import DTYPES_DICT
 
 script_dir = os.path.dirname(os.path.abspath(__file__))
 
@@ -53,12 +47,6 @@ def randoms():
     nc = nc.astype(np.bfloat16 if DTYPE_IN == dtypes.bfloat16 else np.float16)
   return na, nb, nc
 
-def ast_to_cuda_prog(compiler, ast, opts):
-  k = Kernel(ast)
-  k.apply_opts(opts)
-  p = get_program(k.ast, k.opts, k.applied_opts)
-  return CUDAProgram(device, p.function_name, compiler.compile(p.src))
-
 if __name__ == "__main__":
   print(f"gemm variation: {GEMM_VARIATION=} {M=} {N=} {K=} {DTYPE_IN=} {DTYPE_OUT=} {DTYPE_ACC=}")
   prog, global_size, local_size = None, None, None
@@ -189,11 +177,11 @@ if __name__ == "__main__":
 
     tms = []
     na, nb, nc = randoms()
-    cudaalloc.copyin(a, bytearray(na))
-    cudaalloc.copyin(b, bytearray(nb))
+    cudaalloc._copyin(a, memoryview(bytearray(na)))
+    cudaalloc._copyin(b, memoryview(bytearray(nb)))
     for i in range(CNT):
       tms.append(prog(*args, **kwargs))
-    cudaalloc.copyout(flat_mv(nc.data), c)
+    cudaalloc._copyout(flat_mv(nc.data), c)
     comp = na.astype(np.float32) @ nb.astype(np.float32)
     result = nc.reshape(M, N).astype(np.float32)
 

extra/gemm/mi350x_uop_matmul.py

@@ -37,11 +37,11 @@ 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)
+#  -- DONE: add working slice a[gx, :, i] -> shape of the : (aka (16,16,32) becomes (16,))
+#  -- DONE(ish): 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
+#  -- DONE: 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)
 

extra/gemm/simple_matmul.py

@@ -17,7 +17,7 @@ M = getenv("M", N)
 K = getenv("K", N)
 CNT = getenv("CNT", 10)
 
-atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1)}.get(dtype_in, (1e-4, 3e-2))
+atol, rtol = {dtypes.half:{1e-3, 1e-2}, dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1)}.get(dtype_in, (1e-4, 3e-2))
 ATOL, RTOL = getenv("ATOL", atol), getenv("RTOL", rtol)
 
 INT_LOW = getenv("INT_LOW", 0)

extra/models/llama.py

@@ -89,6 +89,20 @@ class Attention:
     keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
     xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
     attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2)
+    if getenv("STUB_ATTENTION"):
+      # TODO: do we need mask?
+      from tinygrad.uop.ops import UOp, KernelInfo
+      def fa_custom_forward(attn:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
+        return UOp.sink(arg=KernelInfo(name="fa_custom_forward"))
+      def fa_custom_backward(out_q:UOp, out_k:UOp, out_v:UOp, grad:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
+        return UOp.sink(arg=KernelInfo(name="fa_custom_backward"))
+      def fa_backward(grad:UOp, kernel:UOp) -> tuple[None, UOp, UOp, UOp]:
+        grad_q = Tensor.empty_like(q:=Tensor(kernel.src[1]))
+        grad_k = Tensor.empty_like(k:=Tensor(kernel.src[2]))
+        grad_v = Tensor.empty_like(v:=Tensor(kernel.src[3]))
+        ck = Tensor.custom_kernel(grad_q, grad_k, grad_v, Tensor(grad), q, k, v, fxn=fa_custom_backward)[:3]
+        return (None, ck[0].uop, ck[1].uop, ck[2].uop)
+      attn = Tensor.empty_like(attn).custom_kernel(xq, keys, values, fxn=fa_custom_forward, grad_fxn=fa_backward)[0]
     attn = attn.reshape(bsz, seqlen, -1)
     return self.wo(attn)
 

extra/nv_gpu_driver/pci_exp_table.h

@@ -0,0 +1,134 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#ifndef PCIEXPTBL_H
+#define PCIEXPTBL_H
+
+#define NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_BASE   0x00
+#define NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_EXT    0xE0
+
+//
+// The VBIOS object comes from walking the PCI expansion code block
+// The following structure holds the expansion code format.
+//
+#define PCI_EXP_ROM_SIGNATURE     0xaa55
+#define PCI_EXP_ROM_SIGNATURE_NV  0x4e56 // "VN" in word format
+#define PCI_EXP_ROM_SIGNATURE_NV2 0xbb77
+#define IS_VALID_PCI_ROM_SIG(sig) ((sig == PCI_EXP_ROM_SIGNATURE) ||    \
+                                   (sig == PCI_EXP_ROM_SIGNATURE_NV) || \
+                                   (sig == PCI_EXP_ROM_SIGNATURE_NV2))
+
+#define OFFSETOF_PCI_EXP_ROM_SIG                    0x0
+#define OFFSETOF_PCI_EXP_ROM_NBSI_DATA_OFFSET       0x16
+#define OFFSETOF_PCI_EXP_ROM_PCI_DATA_STRUCT_PTR    0x18
+
+#pragma pack(1)
+typedef struct _PCI_EXP_ROM_STANDARD
+{
+    NvU16       sig;                //  00h: ROM Signature 0xaa55
+    NvU8        reserved [0x16];    //  02h: Reserved (processor architecture unique data)
+    NvU16       pciDataStrucPtr;    //  18h: Pointer to PCI Data Structure
+    NvU32       sizeOfBlock;        //  1Ah: <NBSI-specific appendage>
+} PCI_EXP_ROM_STANDARD, *PPCI_EXP_ROM_STANDARD;
+#pragma pack()
+
+#pragma pack(1)
+typedef struct _PCI_EXP_ROM_NBSI
+{
+    NvU16       sig;                //  00h: ROM Signature 0xaa55
+    NvU8        reserved [0x14];    //  02h: Reserved (processor architecture unique data)
+    NvU16       nbsiDataOffset;     //  16h: Offset from header to NBSI image
+    NvU16       pciDataStrucPtr;    //  18h: Pointer to PCI Data Structure
+    NvU32       sizeOfBlock;        //  1Ah: <NBSI-specific appendage>
+} PCI_EXP_ROM_NBSI, *PPCI_EXP_ROM_NBSI;
+#pragma pack()
+
+typedef union _PCI_EXP_ROM {
+    PCI_EXP_ROM_STANDARD standard;
+    PCI_EXP_ROM_NBSI nbsi;
+} PCI_EXP_ROM, *PPCI_EXP_ROM;
+
+#define PCI_DATA_STRUCT_SIGNATURE     0x52494350 // "PCIR" in dword format
+#define PCI_DATA_STRUCT_SIGNATURE_NV  0x5344504E // "NPDS" in dword format
+#define PCI_DATA_STRUCT_SIGNATURE_NV2 0x53494752 // "RGIS" in dword format
+#define IS_VALID_PCI_DATA_SIG(sig) ((sig == PCI_DATA_STRUCT_SIGNATURE) ||    \
+                                    (sig == PCI_DATA_STRUCT_SIGNATURE_NV) || \
+                                    (sig == PCI_DATA_STRUCT_SIGNATURE_NV2))
+
+#define PCI_LAST_IMAGE NVBIT(7)
+#define PCI_ROM_IMAGE_BLOCK_SIZE 512U
+
+#define OFFSETOF_PCI_DATA_STRUCT_SIG        0x0
+#define OFFSETOF_PCI_DATA_STRUCT_VENDOR_ID  0x4
+#define OFFSETOF_PCI_DATA_STRUCT_LEN        0xa
+#define OFFSETOF_PCI_DATA_STRUCT_CLASS_CODE 0xd
+#define OFFSETOF_PCI_DATA_STRUCT_CODE_TYPE  0x14
+#define OFFSETOF_PCI_DATA_STRUCT_IMAGE_LEN  0x10
+#define OFFSETOF_PCI_DATA_STRUCT_LAST_IMAGE 0x15
+
+#pragma pack(1)
+typedef struct _PCI_DATA_STRUCT
+{
+    NvU32       sig;                //  00h: Signature, the string "PCIR" or NVIDIA's alternate "NPDS"
+    NvU16       vendorID;           //  04h: Vendor Identification
+    NvU16       deviceID;           //  06h: Device Identification
+    NvU16       deviceListPtr;      //  08h: Device List Pointer
+    NvU16       pciDataStructLen;   //  0Ah: PCI Data Structure Length
+    NvU8        pciDataStructRev;   //  0Ch: PCI Data Structure Revision
+    NvU8        classCode[3];       //  0Dh: Class Code
+    NvU16       imageLen;           //  10h: Image Length (units of 512 bytes)
+    NvU16       vendorRomRev;       //  12h: Revision Level of the Vendor's ROM
+    NvU8        codeType;           //  14h: holds NBSI_OBJ_CODE_TYPE (0x70) and others
+    NvU8        lastImage;          //  15h: Last Image Indicator: bit7=1 is lastImage
+    NvU16       maxRunTimeImageLen; //  16h: Maximum Run-time Image Length (units of 512 bytes)
+} PCI_DATA_STRUCT, *PPCI_DATA_STRUCT;
+#pragma pack()
+
+#define NV_PCI_DATA_EXT_SIG 0x4544504E // "NPDE" in dword format
+#define NV_PCI_DATA_EXT_REV_10 0x100      // 1.0
+#define NV_PCI_DATA_EXT_REV_11 0x101      // 1.1
+
+#define OFFSETOF_PCI_DATA_EXT_STRUCT_SIG            0x0
+#define OFFSETOF_PCI_DATA_EXT_STRUCT_LEN            0x6
+#define OFFSETOF_PCI_DATA_EXT_STRUCT_REV            0x4
+#define OFFSETOF_PCI_DATA_EXT_STRUCT_SUBIMAGE_LEN   0x8
+#define OFFSETOF_PCI_DATA_EXT_STRUCT_LAST_IMAGE     0xa
+#define OFFSETOF_PCI_DATA_EXT_STRUCT_FLAGS          0xb
+
+#define PCI_DATA_EXT_STRUCT_FLAGS_CHECKSUM_DISABLED 0x04
+
+#pragma pack(1)
+typedef struct _NV_PCI_DATA_EXT_STRUCT
+{
+    NvU32   signature;          //  00h: Signature, the string "NPDE"
+    NvU16   nvPciDataExtRev;    //  04h: NVIDIA PCI Data Extension Revision
+    NvU16   nvPciDataExtLen;    //  06h: NVIDIA PCI Data Extension Length
+    NvU16   subimageLen;        //  08h: Sub-image Length
+    NvU8    privLastImage;      //  0Ah: Private Last Image Indicator
+    NvU8    flags;              //  0Bh: Private images enabled if bit0=1
+} NV_PCI_DATA_EXT_STRUCT, *PNV_PCI_DATA_EXT_STRUCT;
+#pragma pack()
+
+#endif // PCIEXPTBL_H
+
+

extra/sqtt/disasm.py

@@ -1,68 +0,0 @@
-import ctypes
-from dataclasses import dataclass
-import tinygrad.runtime.autogen.comgr as comgr
-from tinygrad.runtime.support.compiler_amd import check
-
-@dataclass
-class InstrCtx:
-  pc:int=0
-  inst:str=""
-
-@comgr.amd_comgr_create_disassembly_info.argtypes[2]
-def instr_cb(text, user_data):
-  c = ctypes.cast(user_data, ctypes.POINTER(ctypes.py_object)).contents.value
-  c.inst = ctypes.string_at(text).decode("utf-8","replace").strip()
-  return comgr.AMD_COMGR_STATUS_SUCCESS
-
-# nop callback
-@comgr.amd_comgr_create_disassembly_info.argtypes[3]
-def addr_cb(*args): return comgr.AMD_COMGR_STATUS_SUCCESS
-
-def comgr_get_address_table(lib:bytes) -> dict[int, tuple[str, int]]:
-  check(comgr.amd_comgr_create_data(comgr.AMD_COMGR_DATA_KIND_EXECUTABLE, ctypes.byref(data_src:=comgr.amd_comgr_data_t())))
-  lib_buf = ctypes.create_string_buffer(lib, len(lib))
-  check(comgr.amd_comgr_set_data(data_src, len(lib), lib_buf))
-  check(comgr.amd_comgr_get_data_isa_name(data_src, isa_sz:=ctypes.c_size_t(128), isa:=(ctypes.c_char*isa_sz.value)()))
-
-  @comgr.amd_comgr_create_disassembly_info.argtypes[1]
-  def memory_cb(from_addr, to, size, _):
-    base, buf_len = ctypes.addressof(lib_buf), len(lib_buf)
-    start = int(from_addr) - base
-    if start < 0 or start >= buf_len: return 0
-    ctypes.memmove(to, base + start, n:=min(int(size), buf_len - start))
-    return n
-
-  info_src = comgr.amd_comgr_disassembly_info_t()
-  check(comgr.amd_comgr_create_disassembly_info(ctypes.cast(isa, ctypes.POINTER(ctypes.c_char)), memory_cb, instr_cb, addr_cb, info_src))
-
-  @comgr.amd_comgr_iterate_symbols.argtypes[1]
-  def sym_callback(sym, udata):
-    check(comgr.amd_comgr_symbol_get_info(sym, comgr.AMD_COMGR_SYMBOL_INFO_TYPE, ctypes.byref(sym_type:=ctypes.c_int())))
-    if sym_type.value != comgr.AMD_COMGR_SYMBOL_TYPE_FUNC: return comgr.AMD_COMGR_STATUS_SUCCESS
-    check(comgr.amd_comgr_symbol_get_info(sym, comgr.AMD_COMGR_SYMBOL_INFO_VALUE, ctypes.byref(vaddr:=ctypes.c_uint64())))
-    check(comgr.amd_comgr_symbol_get_info(sym, comgr.AMD_COMGR_SYMBOL_INFO_SIZE, ctypes.byref(size:=ctypes.c_uint64())))
-    check(comgr.amd_comgr_map_elf_virtual_address_to_code_object_offset(data_src, vaddr.value, ctypes.byref(offset:=ctypes.c_uint64()),
-                                                                        ctypes.byref(ctypes.c_uint64()), ctypes.byref(nobits:=ctypes.c_bool())))
-    check(nobits.value)
-    base = ctypes.addressof(lib_buf)
-    pc = base + offset.value
-    end = pc + size.value
-    addr_table = ctypes.cast(udata, ctypes.POINTER(ctypes.py_object)).contents.value
-    instr_ref = ctypes.py_object(ctx:=InstrCtx())
-    instr_ptr = ctypes.cast(ctypes.pointer(instr_ref), ctypes.c_void_p)
-    while pc < end:
-      size_read = ctypes.c_uint64(0)
-      ctx.pc = pc
-      st = comgr.amd_comgr_disassemble_instruction(info_src, ctypes.c_uint64(pc), instr_ptr, ctypes.byref(size_read))
-      if st == comgr.AMD_COMGR_STATUS_SUCCESS and size_read.value:
-        rel = (pc - base) - offset.value
-        addr_table[vaddr.value + rel] = (ctx.inst, int(size_read.value))
-        pc += size_read.value
-      else: # don't inf loop if comgr fails
-        b = ctypes.c_ubyte.from_buffer(lib_buf, pc - base).value
-        addr_table[vaddr.value + (pc - base - offset.value)] = (f"DISASSEMBLER ISSUE 0x{b:02x}", 1)
-        pc += 1
-    return comgr.AMD_COMGR_STATUS_SUCCESS
-  addr_table:dict[int, tuple[str, int]] = {}
-  check(comgr.amd_comgr_iterate_symbols(data_src, sym_callback, ctypes.cast(ctypes.pointer(ctypes.py_object(addr_table)), ctypes.c_void_p)))
-  return addr_table

extra/sqtt/install_sqtt_decoder.py

@@ -13,6 +13,6 @@ if __name__ == "__main__":
     os.chmod(fp, 0o755)
     os.system(f"sudo {fp} --prefix={fp.parent} --include-subdir")
   else:
-    lib = fetch("https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so", name="librocprof-trace-decoder.so")
+    lib = fetch("https://github.com/ROCm/rocprof-trace-decoder/raw/43bf0fef74a83c3c25badfc5a09c0bd39ed8c6f9/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so", name="librocprof-trace-decoder.so")
   shutil.copy2(lib, DEST)
   print(f"Installed {lib.name} to", DEST)

extra/sqtt/rgptool.py

@@ -4,7 +4,7 @@ import argparse, ctypes, struct, hashlib, pickle, code, typing, functools
 import tinygrad.runtime.autogen.sqtt as sqtt
 from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent
-from tinygrad.helpers import round_up, flatten, all_same
+from tinygrad.helpers import round_up, flatten, all_same, temp
 from dataclasses import dataclass
 
 CHUNK_CLASSES = {
@@ -154,8 +154,22 @@ class RGP:
       if device not in device_events: raise RuntimeError(f"Device {device} not found in profile, devices in profile: {', '.join(device_events.keys())} ")
       device_event = device_events[device]
     sqtt_events = [x for x in profile if isinstance(x, ProfileSQTTEvent) and x.device == device_event.device]
+    device_props = device_event.props
+    # merge events per SE
+    merged_sqtt_events:dict[int, ProfileSQTTEvent] = {}
+    for ev in sqtt_events:
+      if ev.se not in merged_sqtt_events: merged_sqtt_events[ev.se] = ev
+      else:
+        merged_sqtt_events[ev.se] = ProfileSQTTEvent(
+          device=ev.device,
+          kern=ev.kern,
+          se=ev.se,
+          itrace=merged_sqtt_events[ev.se].itrace or ev.itrace,
+          blob=merged_sqtt_events[ev.se].blob + ev.blob,
+        )
+    sqtt_events = list(merged_sqtt_events.values())
+
     if len(sqtt_events) == 0: raise RuntimeError(f"Device {device_event.device} doesn't contain SQTT data")
-    device_props = sqtt_events[0].props
     gfx_ver = device_props['gfx_target_version'] // 10000
     gfx_iplvl = getattr(sqtt, f"SQTT_GFXIP_LEVEL_GFXIP_{device_props['gfx_target_version']//10000}_{(device_props['gfx_target_version']//100)%100}",
                         getattr(sqtt, f"SQTT_GFXIP_LEVEL_GFXIP_{device_props['gfx_target_version']//10000}", None))
@@ -196,7 +210,7 @@ class RGP:
         flags=0,
         trace_shader_core_clock=0x93f05080,
         trace_memory_clock=0x4a723a40,
-        device_id={110000: 0x744c, 110003: 0x7480, 120001: 0x7550}[device_props['gfx_target_version']],
+        device_id={110000: 0x744c, 110003: 0x7480, 120001: 0x7550, 120000: 0x7550}[device_props['gfx_target_version']],
         device_revision_id=0xc8,
         vgprs_per_simd=1536,
         sgprs_per_simd=128*16,
@@ -310,7 +324,7 @@ class RGP:
 if __name__ == '__main__':
   parser = argparse.ArgumentParser(prog='rgptool', description='A tool to create (from pickled tinygrad profile), inspect and modify Radeon GPU Profiler files')
   parser.add_argument('command')
-  parser.add_argument('input')
+  parser.add_argument('input', nargs='?', default=temp("profile.pkl", append_user=True))
   parser.add_argument('-d', '--device')
   parser.add_argument('-o', '--output')
   args = parser.parse_args()
@@ -332,3 +346,4 @@ if __name__ == '__main__':
 
   if args.output is not None:
     with open(args.output, 'wb+') as fd: fd.write(rgp.to_bytes())
+    print(f"Saved to {args.output}")

extra/sqtt/roc.py

@@ -1,9 +1,32 @@
 import ctypes, pathlib, argparse, pickle, re, functools, dataclasses, itertools
-from extra.sqtt.rocprof import rocprof
-from extra.sqtt.disasm import comgr_get_address_table
-from tinygrad.helpers import temp, DEBUG
-from tinygrad.device import ProfileEvent, ProfileProgramEvent
+from tinygrad.helpers import temp, unwrap, DEBUG
+from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent, ProfilePMCEvent
+from tinygrad.runtime.autogen import llvm, rocprof
+from tinygrad.runtime.support.elf import elf_loader
+
+# to pass NULL to callbacks
+llvm.LLVMCreateDisasmCPUFeatures.argtypes = tuple(llvm.LLVMCreateDisasmCPUFeatures.argtypes[:5]) + (ctypes.c_void_p, ctypes.c_void_p)
+def llvm_disasm(arch:str, lib:bytes) -> dict[int, tuple[str, int]]:
+  llvm.LLVMInitializeAMDGPUTargetInfo()
+  llvm.LLVMInitializeAMDGPUTargetMC()
+  llvm.LLVMInitializeAMDGPUAsmParser()
+  llvm.LLVMInitializeAMDGPUDisassembler()
+  ctx = llvm.LLVMCreateDisasmCPUFeatures("amdgcn-amd-amdhsa".encode(), arch.encode(), "".encode(), None, 0, None, None)
+
+  image, sections, relocs = elf_loader(lib)
+  text = next((sh.header for sh in sections if sh.name == ".text"), None)
+  off, sz = unwrap(text).sh_addr, unwrap(text).sh_size
+
+  addr_table:dict[int, tuple[str, int]] = {}
+  out = ctypes.create_string_buffer(128)
+  cur_off = off
+  while cur_off < sz + off:
+    view = (ctypes.c_ubyte * ((sz + off) - cur_off)).from_buffer_copy(memoryview(image)[cur_off:])
+    instr_sz = llvm.LLVMDisasmInstruction(ctx, view, ctypes.c_uint64(len(view)), ctypes.c_uint64(0), out, ctypes.c_size_t(128))
+    addr_table[cur_off] = (out.value.decode("utf-8", "replace").strip(), instr_sz)
+    cur_off += instr_sz
+  return addr_table
 
 @dataclasses.dataclass
 class InstInfo:
@@ -17,53 +40,77 @@ class InstInfo:
   def on_ev(self, ev):
     self.hit, self.lat, self.stall = self.hit + 1, self.lat + ev.duration, self.stall + ev.stall
 
+@dataclasses.dataclass(frozen=True)
+class InstExec:
+  typ:str
+  inst:str
+  stall:int
+  dur:int
+  time:int
+
+@dataclasses.dataclass(frozen=True)
+class PrgExec:
+  name:str
+  wave:int
+  cu:int
+  simd:int
+  def __str__(self): return f"{self.name},{self.wave},{self.cu},{self.simd}"
+
+@dataclasses.dataclass(frozen=True)
+class WaveExec:
+  wave_id:int
+  cu:int
+  simd:int
+  insts:list[InstExec]
+
 class _ROCParseCtx:
-  def __init__(self, sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
-    self.sqtt_evs, self.prog_evs = iter(sqtt_evs), prog_evs
-    self.wave_events, self.disasms, self.addr2prg = {}, {}, {}
+  def __init__(self, dev_evs:dict[str, ProfileDeviceEvent], sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
+    self.dev_evs, self.sqtt_evs, self.prog_evs = dev_evs, iter(sqtt_evs), prog_evs
+    self.wave_events:dict[PrgExec, dict[int, InstInfo]] = {}
+    self.disasms:dict[tuple[str, int], tuple[str, int]] = {}
+    self.inst_execs:dict[str, list[WaveExec]] = {}
 
     for prog in prog_evs:
-      for addr, info in comgr_get_address_table(prog.lib).items():
-        self.disasms[prog.base + addr] = info
-        self.addr2prg[prog.base + addr] = prog
+      arch = "gfx%d%x%x" % ((trgt:=unwrap(dev_evs[prog.device].props)['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
+      for addr, info in llvm_disasm(arch, unwrap(prog.lib)).items():
+        self.disasms[(prog.name, unwrap(prog.base) + addr)] = info
 
   def next_sqtt(self):
     x = next(self.sqtt_evs, None)
+    self.active_kern = x.kern if x is not None else None
     self.active_se = x.se if x is not None else None
     return x
 
-  def find_program(self, addr): return self.addr2prg[addr]
-
   def on_occupancy_ev(self, ev):
-    if DEBUG >= 4: print("OCC", ev.time, self.active_se, ev.cu, ev.simd, ev.wave_id, ev.start)
+    if DEBUG >= 5: print("OCC", ev.time, self.active_se, ev.cu, ev.simd, ev.wave_id, ev.start)
 
   def on_wave_ev(self, ev):
-    if DEBUG >= 4: print("WAVE", ev.wave_id, self.active_se, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time)
+    if DEBUG >= 5: print("WAVE", ev.wave_id, self.active_se, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time)
 
-    asm = {}
+    asm:dict[int, InstInfo] = {}
+    inst_execs:list[InstExec] = []
     for j in range(ev.instructions_size):
       inst_ev = ev.instructions_array[j]
       inst_typ = rocprof.rocprofiler_thread_trace_decoder_inst_category_t__enumvalues[inst_ev.category]
-      asm.setdefault(inst_ev.pc.address, InstInfo(typ=inst_typ, inst=self.disasms[inst_ev.pc.address][0]))
+      inst_disasm = self.disasms[(unwrap(self.active_kern), unwrap(inst_ev.pc.address))][0]
+      asm.setdefault(inst_ev.pc.address, InstInfo(typ=inst_typ, inst=inst_disasm))
       asm[inst_ev.pc.address].on_ev(inst_ev)
+      inst_execs.append(InstExec(inst_typ, inst_disasm, inst_ev.stall, inst_ev.duration, inst_ev.time))
 
-    self.wave_events[(self.find_program(ev.instructions_array[0].pc.address).name, ev.wave_id, ev.cu, ev.simd)] = asm
+    if ev.instructions_size > 0:
+      self.wave_events[key:=PrgExec(unwrap(self.active_kern), ev.wave_id, ev.cu, ev.simd)] = asm
+      self.inst_execs.setdefault(key.name, []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, inst_execs))
 
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser()
-  parser.add_argument('--profile', type=pathlib.Path, help='Path to profile', default=pathlib.Path(temp("profile.pkl", append_user=True)))
-  args = parser.parse_args()
-
-  with args.profile.open("rb") as f: profile = pickle.load(f)
+def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
+  dev_events:dict[str, ProfileDeviceEvent] = {}
   sqtt_events:list[ProfileSQTTEvent] = []
-  pmc_events:list[ProfilePMCEvent] = []
   prog_events:list[ProfileProgramEvent] = []
   for e in profile:
+    if isinstance(e, ProfileDeviceEvent): dev_events[e.device] = e
     if isinstance(e, ProfileSQTTEvent): sqtt_events.append(e)
-    if isinstance(e, ProfilePMCEvent): pmc_events.append(e)
     if isinstance(e, ProfileProgramEvent) and e.device.startswith("AMD"): prog_events.append(e)
 
-  ROCParseCtx = _ROCParseCtx(sqtt_events, prog_events)
+  ROCParseCtx = _ROCParseCtx(dev_events, sqtt_events, prog_events)
 
   @rocprof.rocprof_trace_decoder_se_data_callback_t
   def copy_cb(buf, buf_size, data_ptr):
@@ -80,12 +127,12 @@ if __name__ == "__main__":
       case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE:
         for ev in (rocprof.rocprofiler_thread_trace_decoder_wave_t * n).from_address(events_ptr): ROCParseCtx.on_wave_ev(ev)
       case _:
-        if DEBUG >= 2: print(rocprof.rocprofiler_thread_trace_decoder_record_type_t__enumvalues[record_type], events_ptr, n)
+        if DEBUG >= 5: print(rocprof.rocprofiler_thread_trace_decoder_record_type_t__enumvalues[record_type], events_ptr, n)
     return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
 
   @rocprof.rocprof_trace_decoder_isa_callback_t
   def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, data_ptr):
-    instr, mem_size_ptr[0] = ROCParseCtx.disasms[pc.address]
+    instr, mem_size_ptr[0] = ROCParseCtx.disasms[(unwrap(ROCParseCtx.active_kern), pc.address)]
 
     # this is the number of bytes to next instruction, set to 0 for end_pgm
     if instr == "s_endpgm": mem_size_ptr[0] = 0
@@ -100,10 +147,20 @@ if __name__ == "__main__":
 
   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)
+  except AttributeError as e: raise RuntimeError("Failed to find rocprof-trace-decoder. Run ./extra/sqtt/install_sqtt_decoder.py to install") from e
+  return ROCParseCtx
 
-  for ev in pmc_events:
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--profile', type=pathlib.Path, help='Path to profile', default=pathlib.Path(temp("profile.pkl", append_user=True)))
+  args = parser.parse_args()
+
+  with args.profile.open("rb") as f: profile = pickle.load(f)
+  rctx = decode(profile)
+  print('SQTT:', rctx.wave_events.keys())
+
+  for ev in profile:
+    if not isinstance(ev, ProfilePMCEvent): continue
     print(f"PMC Event: dev={ev.device} kern={ev.kern}")
     ptr = 0
     for s in ev.sched:

extra/sqtt/test_timing.py

@@ -0,0 +1,105 @@
+import os
+os.environ["PYTHONPATH"] = "."
+os.environ["SQTT"] = "1"
+os.environ["AMD"] = "1"
+os.environ["VIZ"] = "1"
+os.environ["AMD_LLVM"] = "0"
+
+import unittest
+import sys, contextlib
+from tinygrad import Tensor
+from tinygrad.dtype import dtypes
+from tinygrad.renderer import ProgramSpec
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.engine.realize import CompiledRunner
+from tinygrad.device import Device, ProfileDeviceEvent
+
+from extra.sqtt.roc import decode, InstExec, PrgExec
+
+dev = Device["AMD"]
+
+def custom(arg:str, s:UOp|None=None) -> UOp: return UOp(Ops.CUSTOM, src=(s,) if s is not None else (), arg=arg)
+
+def asm_kernel(instrs:list[str], l:int=1, g:int=1) -> Tensor:
+  name = sys._getframe(1).f_code.co_name
+  def fxn(_):
+    L = UOp.special(l, "lidx0")
+    G = UOp.special(g, "gidx0")
+    op = custom("asm volatile (")
+    for inst in instrs: op = custom(f'  "{inst}\\n\\t"', op)
+    op = custom(");", op)
+    return UOp.sink(op, L, G, arg=KernelInfo(name=name))
+  k = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
+  return k
+
+@contextlib.contextmanager
+def save_sqtt():
+  # clear the old traces
+  dev.profile_events.clear()
+  sqtt:dict[PrgExec, list[InstExec]] = {}
+  yield sqtt
+  # decode sqtt
+  rctx = decode(dev.profile_events+[ProfileDeviceEvent("AMD", props=dev.device_props())])
+  assert len(rctx.inst_execs) > 0, "empty sqtt output"
+  sqtt.update(rctx.inst_execs)
+
+class TestTiming(unittest.TestCase):
+  def test_v_add(self):
+    with save_sqtt() as sqtt:
+      asm_kernel([f"v_add_f32 v{10+i} v{10+i+1} {10+i}" for i in range(3)]).realize()
+    wave = list(sqtt.values())[0][:-1]
+    assert all(s.dur == 1 for s in wave)
+    assert all(s.stall == 0 for s in wave)
+
+  def test_chain_v_add_1l(self):
+    with save_sqtt() as sqtt:
+      asm_kernel([
+        "v_add_f32_e32 v1 v0 v0",
+        "v_add_f32_e32 v2 v1 v1",
+      ]).realize()
+    wave = list(sqtt.values())[0][:-1]
+    assert all(s.dur == 1 for s in wave)
+    assert all(s.stall == 0 for s in wave)
+
+  def test_multi_cycle_inst(self):
+    with save_sqtt() as sqtt:
+      asm_kernel([
+        "v_mov_b32_e32 v4 0x3f800000",
+        "v_rcp_f32_e32 v5 v4",
+        "v_mul_f32_e32 v6 v5 v4",
+      ]).realize()
+    w = list(sqtt.values())[0]
+    rcp, mul = w[1], w[2]
+    self.assertGreater(rcp.dur, 1) # 4 cycles on gfx11
+    self.assertEqual(mul.dur, 1)
+    # mul depends on v5, how can it run before rcp is done?
+    self.assertGreaterEqual(mul.time, rcp.time+rcp.dur)
+
+  def test_wmma(self):
+    with save_sqtt() as sqtt:
+      asm_kernel([
+        "v_wmma_f32_16x16x16_f16 v[16:23], v[0:7], v[8:15], v[16:23]",
+        "v_add_f32_e32 v0 v16 v0",
+      ], l=32*4).realize()
+    assert len(sqtt) == 2, f"expected two waves, got {len(sqtt)} {list(sqtt.keys())}"
+    wmma = list(sqtt.values())[0][0]
+    self.assertGreater(wmma.dur, 1) # rgp says 32 clocks
+
+  def test_sleep(self):
+    n = 1
+    def sleep_kernel(data0):
+      assert data0.dtype.base == dtypes.ulong
+      op = custom("unsigned long long t0 = __builtin_readcyclecounter();")
+      op = custom(f"__builtin_amdgcn_s_sleep({n});", op)
+      op = custom(f"unsigned long long t1 = __builtin_readcyclecounter();", op)
+      op = custom(f"data0_{data0.size}[0] = t1 - t0;", op)
+      return UOp.sink(data0, op, arg=KernelInfo(name=f"sleep_{n}"))
+    diff_hw_reg = Tensor.empty(1, dtype=dtypes.ulong)
+    diff_hw_reg = Tensor.custom_kernel(diff_hw_reg, fxn=sleep_kernel)[0]
+    with save_sqtt() as sqtt:
+      diff_hw_reg.realize()
+    diff_sqtt = list(sqtt.values())[0][2]
+    self.assertEqual(diff_sqtt.dur, diff_hw_reg.item()-1) # 1 cycle for reading the counter register
+
+if __name__ == "__main__":
+  unittest.main()

extra/thunder/cuda/fa.cu

@@ -10,7 +10,7 @@ constexpr int ATTN_N = 1024;
 constexpr int ATTN_H = 16;
 constexpr int ATTN_D = 64;
 
-template<int D> constexpr size_t ROWS = 16*(128/D); // height of each worker tile (rows)
+template<int D> constexpr size_t ROWS = 16*(64/D); // height of each worker tile (rows)
 template<int D, typename T=bf16, typename L=row_l> using qkvo_tile = rt<T, ROWS<D>, D, L>;
 template<int D, typename T=float> using attn_tile = rt<T, ROWS<D>, ROWS<D>>;
 template<int D> using shared_tile = st_bf<ROWS<D>, D>;

extra/thunder/cuda/fa.py

@@ -23,7 +23,7 @@ if __name__ == "__main__":
   Tensor.realize(q, k, v, out)
 
   NUM_WORKERS = 4
-  ROWS = 16 * (128 // D)
+  ROWS = 16 * (64 // D)
 
   gsz = (N // (ROWS*NUM_WORKERS), H, B)
   for _ in range(5):

extra/thunder/tiny/tk/__init__.py

@@ -0,0 +1 @@
+WARP_THREADS = 32

extra/thunder/tiny/tk/group.py

@@ -0,0 +1,272 @@
+import math, functools
+from typing import cast, Callable
+from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
+from tinygrad.uop.ops import AxisType, UOp, KernelInfo, Ops
+from tinygrad.engine.realize import ExecItem, get_runner
+from tinygrad.dtype import AddrSpace, PtrDType
+from tinygrad.helpers import getenv, prod
+
+from extra.thunder.tiny.tk import WARP_THREADS
+from extra.thunder.tiny.tk.tiles import TILE_ROW_DIM, TILE_COL_DIM, RT_BASE_TILE_NEPT, slots
+
+class Group:
+  def __init__(self, warps:int, ker):
+    self.warps = warps
+    self.group_threads = warps * WARP_THREADS
+    self.threadIdx_x = ker.threadIdx_x
+    self.ker = ker
+
+  # helpers
+  @property
+  def laneid(self): return self.threadIdx_x % self.group_threads
+  @property
+  def warpid(self): return self.laneid // WARP_THREADS
+  @property
+  def groupid(self): return self.threadIdx_x // self.group_threads
+
+  # ops that only work on a single warp
+
+  clear_rid = 1000
+  def clear(self, reg:UOp, value:float=0):
+    assert self.warps == 1
+
+    i = UOp.range(reg.size, Group.clear_rid)
+    Group.clear_rid += 1
+    return reg.reshape((reg.size,))[i].set(value, end=i).after(reg).reshape(reg.shape)
+
+  def zero(self, reg:UOp): return self.clear(reg, 0)
+  def neg_inf(self, reg:UOp): return self.clear(reg, -math.inf)
+
+  copy_rid = 300
+  def copy(self, dst:UOp, src:UOp):
+    assert self.warps == 1
+
+    assert dst.shape == src.shape
+    assert cast(PtrDType, dst.dtype).addrspace == AddrSpace.REG
+    assert cast(PtrDType, src.dtype).addrspace == AddrSpace.REG
+
+    rngs_for_shape = tuple(UOp.range(dim, Group.copy_rid + i) for i, dim in enumerate(dst.shape))
+    Group.copy_rid += len(dst.shape)
+
+    dst_store = dst[*rngs_for_shape].store(src[*rngs_for_shape].cast(dst.dtype.base)).end(*rngs_for_shape)
+
+    self.ker.push_store(dst_store, dst)
+    return dst.after(dst_store).reshape(dst.shape)
+
+  mma_rid = 600
+  def mma_AB(self, c:UOp, a:UOp, b:UOp, after=True):
+    assert self.warps == 1
+
+    mma_i_height = UOp.range(c.shape[-3], Group.mma_rid)
+    mma_i_width = UOp.range(c.shape[-2], Group.mma_rid+1)
+    mma_i_inner = UOp.range(a.shape[-2], Group.mma_rid+2, AxisType.REDUCE)
+    Group.mma_rid += 3
+
+    wmma_arg = ("WMMA_8_16_16_bfloat16_float", (8, 16, 16), dtypes.bfloat16, dtypes.float, "CUDA", 32, (((4, 2), (3, 2), (8, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
+
+    a_in = UOp.vectorize(*[a[mma_i_height, mma_i_inner, i] for i in range(8)])
+    b_in1 = UOp.vectorize(*([b[mma_i_inner, mma_i_width, i] for i in range(2)] + [b[mma_i_inner, mma_i_width, 4+i] for i in range(2)]))
+    c_out1 = UOp.vectorize(*[c[mma_i_height, mma_i_width, i] for i in range(4)])
+    b_in2 = UOp.vectorize(*([b[mma_i_inner, mma_i_width, 2+i] for i in range(2)] + [b[mma_i_inner, mma_i_width, 6+i] for i in range(2)]))
+    c_out2 = UOp.vectorize(*[c[mma_i_height, mma_i_width, 4+i] for i in range(4)])
+
+    out1 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in1, c_out1), arg=wmma_arg)
+    out2 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in2, c_out2), arg=wmma_arg)
+    c_i = [c[mma_i_height, mma_i_width, i].store(out1.gep(i)) for i in range(4)] + [c[mma_i_height, mma_i_width, 4+i].store(out2.gep(i)) for i in range(4)]
+    c_store = UOp.group(*c_i).end(mma_i_height, mma_i_width, mma_i_inner)
+
+    self.ker.push_store(c_store, c)
+    return c.after(c_store).reshape(c.shape) if after else c_store
+
+  def mma_ABt(self, c:UOp, a:UOp, b:UOp, after=True):
+    assert self.warps == 1
+
+    mma_i_height = UOp.range(c.shape[-3], Group.mma_rid)
+    mma_i_width = UOp.range(c.shape[-2], Group.mma_rid+1)
+    mma_i_inner = UOp.range(a.shape[-2], Group.mma_rid+2, AxisType.REDUCE)
+    Group.mma_rid += 3
+
+    wmma_arg = ("WMMA_8_16_16_bfloat16_float", (8, 16, 16), dtypes.bfloat16, dtypes.float, "CUDA", 32, (((4, 2), (3, 2), (8, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
+
+    a_in = UOp.vectorize(*[a[mma_i_height, mma_i_inner, i] for i in range(8)])
+    b_in1 = UOp.vectorize(*([b[mma_i_width, mma_i_inner, i] for i in range(2)] + [b[mma_i_width, mma_i_inner, 4+i] for i in range(2)]))
+    c_out1 = UOp.vectorize(*[c[mma_i_height, mma_i_width, i] for i in range(4)])
+    b_in2 = UOp.vectorize(*([b[mma_i_width, mma_i_inner, 2+i] for i in range(2)] + [b[mma_i_width, mma_i_inner, 6+i] for i in range(2)]))
+    c_out2 = UOp.vectorize(*[c[mma_i_height, mma_i_width, 4+i] for i in range(4)])
+
+    out1 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in1, c_out1), arg=wmma_arg)
+    out2 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in2, c_out2), arg=wmma_arg)
+    c_i = [c[mma_i_height, mma_i_width, i].store(out1.gep(i)) for i in range(4)] + [c[mma_i_height, mma_i_width, 4+i].store(out2.gep(i)) for i in range(4)]
+    c_store = UOp.group(*c_i).end(mma_i_height, mma_i_width, mma_i_inner)
+
+    self.ker.push_store(c_store, c)
+    return c.after(c_store).reshape(c.shape) if after else c_store
+
+  map_rid = 400
+  def map(self, a:UOp, op:Callable[[UOp], UOp]|Callable[[UOp, tuple], UOp]):
+    assert self.warps == 1
+
+    rngs_for_shape = tuple(UOp.range(dim, Group.map_rid + i) for i, dim in enumerate(a.shape))
+    Group.map_rid += len(a.shape)
+
+    if op.__code__.co_argcount == 1:
+      to_store = op(a[*rngs_for_shape])
+    else:
+      to_store = op(a[*rngs_for_shape], rngs_for_shape)
+
+    a_store = a[*rngs_for_shape].store(to_store).end(*rngs_for_shape)
+
+    self.ker.push_store(a_store, a)
+    return a.after(a_store).reshape(a.shape)
+
+  def row_reduce(self, vec:UOp, src:UOp, op:Callable[[UOp, UOp], UOp]):
+    assert self.warps == 1
+
+    red_local = UOp.placeholder((self.group_threads, 2), src.dtype.base, addrspace=AddrSpace.LOCAL, slot=slots.shared_slot)
+    slots.shared_slot += 1
+
+    for height in self.ker.range(src.shape[-3], track=False):
+      for i_outer in self.ker.range(2, track=False):
+        for width in self.ker.range(src.shape[-2], AxisType.REDUCE, track=False):
+          for i_inner in self.ker.range(4, AxisType.REDUCE, track=False):
+            elem_index = i_inner + 2 * (i_inner // 2) + i_outer * 2
+            vec_store = vec[height, 0, i_outer].store(op(vec[height, 0, i_outer], src[height, width, elem_index])).end(width, i_inner, i_outer)
+            vec = vec.after(vec_store).reshape(vec.shape)
+
+      # store to shared memory
+      for i_outer in self.ker.range(2, track=False):
+        red_local_store = red_local[self.laneid, i_outer].store(vec[height, 0, i_outer]).end(i_outer)
+        red_local = red_local.after(red_local_store).reshape(red_local.shape)
+
+      # reduce from shared memory
+      for i_outer in self.ker.range(2, track=False):
+        for i_inner in self.ker.range(3, AxisType.REDUCE, track=False):
+          offset = (self.laneid // 4) * 4 + ((self.laneid + 1 + i_inner) % 4)
+          vec_store = vec[height, 0, i_outer].store(op(vec[height, 0, i_outer], red_local[offset, i_outer])).end(i_inner, i_outer)
+
+    self.ker.push_store(vec_store, vec)
+    return vec.after(vec_store).reshape(vec.shape)
+
+  # ops that can work across multiple warps
+
+  LOAD_INNER = 8
+  load_rid = 100
+  def load(self, dst:UOp, src:UOp, dst_idxs:tuple[UOp|int,...]=(), idxs:tuple[UOp|int,...]=(), axis:int=0, transpose:bool=False):
+    assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
+    dst_dtype, src_dtype = cast(PtrDType, dst.dtype), cast(PtrDType, src.dtype)
+    if dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.LOCAL:
+      srcf = src.flatten(-2)
+
+      load_i_height = UOp.range(dst.shape[-3], Group.load_rid)
+      load_i_width = UOp.range(dst.shape[-2], Group.load_rid+1)
+      load_i_inner = UOp.range(RT_BASE_TILE_NEPT, Group.load_rid+2)
+      Group.load_rid += 3
+
+      if self.warps % 4 == 0: local_warpid = (self.warpid // 4) + (self.warpid % 4) * (self.warps // 4)
+      else: local_warpid = self.warpid
+      warp_laneid = self.threadIdx_x % WARP_THREADS
+
+      if not transpose:
+        row = (local_warpid * dst.shape[-3] + load_i_height) * TILE_ROW_DIM + (warp_laneid // 4)
+        col = load_i_width * TILE_COL_DIM + 2 * (warp_laneid % 4)
+
+        row_offset = ((load_i_inner % 4) // 2) * 8
+        col_offset = (load_i_inner % 2) + (load_i_inner // 4) * 8
+      else:
+        row = (local_warpid * dst.shape[-3] + load_i_height) * TILE_ROW_DIM + 2 * (warp_laneid % 4)
+        col = load_i_width * TILE_COL_DIM + (warp_laneid // 4)
+
+        row_offset = (load_i_inner % 2) + (load_i_inner // 4) * 8
+        col_offset = ((load_i_inner % 4) // 2) * 8
+
+      src_i_last = (row + row_offset) * src.shape[-1] + col + col_offset
+
+      dst_store = dst[*dst_idxs, load_i_height, load_i_width, load_i_inner].store(srcf[*idxs[:-2], src_i_last])
+      dst_store = dst_store.end(load_i_height, load_i_width, load_i_inner)
+    elif dst_dtype.addrspace == AddrSpace.LOCAL and src_dtype.addrspace == AddrSpace.GLOBAL:
+      dstf = dst.flatten(-2)
+
+      srcf = src.flatten()
+      row_stride = prod(src.shape[axis+1:])
+
+      idxs = tuple(idx * dst.shape[-2] if i == axis else idx for i, idx in enumerate(idxs))
+      idxs = tuple(idx * dst.shape[-1] if i == 3 else idx for i, idx in enumerate(idxs))
+      src_i = ((idxs[0] * src.shape[-3] + idxs[1]) * src.shape[-2] + idxs[2]) * src.shape[-1] + idxs[3]
+
+      memcpy_per_row = dst.shape[-1] // Group.LOAD_INNER
+      total_calls = prod(dst.shape[-2:]) // (self.group_threads * Group.LOAD_INNER)
+
+      load_i_outer = UOp.range(total_calls, Group.load_rid)
+      load_i_inner = UOp.range(Group.LOAD_INNER, Group.load_rid+1)
+      Group.load_rid += 2
+
+      load_idx = load_i_outer * self.group_threads + self.laneid
+      row = load_idx // memcpy_per_row
+      col = (load_idx * Group.LOAD_INNER) % dst.shape[-1]
+
+      dst_i = row * dst.shape[-1] + col + load_i_inner
+      src_i += row * row_stride + col + load_i_inner
+
+      dst_store = dstf[*dst_idxs, dst_i].store(srcf[src_i]).end(load_i_outer, load_i_inner)
+    else:
+      raise NotImplementedError(f"load from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented")
+
+    return dst.after(dst_store.barrier()).reshape(dst.shape)
+
+  STORE_INNER = 8
+  store_rid = 200
+  def store(self, dst:UOp, src:UOp, idxs:tuple[UOp|int,...]=(), src_idxs:tuple[UOp|int,...]=(), axis=0, after=True):
+    assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
+    dst_dtype, src_dtype = cast(PtrDType, dst.dtype), cast(PtrDType, src.dtype)
+    if src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.LOCAL:
+      dstf = dst.flatten(-2)
+
+      store_i_height = UOp.range(src.shape[-3], Group.store_rid)
+      store_i_width = UOp.range(src.shape[-2], Group.store_rid+1)
+      store_i_inner = UOp.range(RT_BASE_TILE_NEPT, Group.store_rid+2)
+      Group.store_rid += 3
+
+      if self.warps % 4 == 0: local_warpid = (self.warpid // 4) + (self.warpid % 4) * (self.warps // 4)
+      else: local_warpid = self.warpid
+      warp_laneid = self.threadIdx_x % WARP_THREADS
+
+      row = (local_warpid * src.shape[-3] + store_i_height) * TILE_ROW_DIM + (warp_laneid // 4)
+      col = store_i_width * TILE_COL_DIM + 2 * (warp_laneid % 4)
+
+      row_offset = ((store_i_inner % 4) // 2) * 8
+      col_offset = (store_i_inner % 2) + (store_i_inner // 4) * 8
+
+      dst_i_last = (row + row_offset) * dst.shape[-1] + col + col_offset
+
+      dst_store = dstf[*idxs[:-2], dst_i_last].store(src[*src_idxs, store_i_height, store_i_width, store_i_inner])
+      dst_store = dst_store.end(store_i_height, store_i_width, store_i_inner)
+    elif src_dtype.addrspace == AddrSpace.LOCAL and dst_dtype.addrspace == AddrSpace.GLOBAL:
+      dstf = dst.flatten()
+      row_stride = prod(dst.shape[axis+1:])
+
+      idxs = tuple(idx * src.shape[-2] if i == axis else idx for i, idx in enumerate(idxs))
+      idxs = tuple(idx * src.shape[-1] if i == 3 else idx for i, idx in enumerate(idxs))
+      dst_i = ((idxs[0] * dst.shape[-3] + idxs[1]) * dst.shape[-2] + idxs[2]) * dst.shape[-1] + idxs[3]
+
+      srcf = src.flatten(-2)
+
+      memcpy_per_row = src.shape[-1] // Group.STORE_INNER
+      total_calls = prod(src.shape[-2:]) // (self.group_threads * Group.STORE_INNER)
+
+      store_i_outer = UOp.range(total_calls, Group.store_rid)
+      store_i_inner = UOp.range(Group.STORE_INNER, Group.store_rid+1)
+      Group.store_rid += 2
+
+      load_idx = store_i_outer * self.group_threads + self.laneid
+      row = load_idx // memcpy_per_row
+      col = (load_idx * Group.STORE_INNER) % src.shape[-1]
+
+      src_i = row * src.shape[-1] + col + store_i_inner
+      dst_i += row * row_stride + col + store_i_inner
+
+      dst_store = dstf[dst_i].store(srcf[*src_idxs, src_i]).end(store_i_outer, store_i_inner)
+    else:
+      raise NotImplementedError(f"store from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented")
+
+    self.ker.push_store(dst_store, dst)
+    return dst.after(dst_store.barrier()).reshape(dst.shape) if after else dst_store

extra/thunder/tiny/tk/kernel.py

@@ -0,0 +1,57 @@
+from contextlib import AbstractContextManager
+from tinygrad.uop.ops import UOp, KernelInfo, AxisType
+from extra.thunder.tiny.tk import WARP_THREADS
+from extra.thunder.tiny.tk.group import Group
+
+class _tk_range:
+  user_rid = 0
+  def __init__(self, end:int, axis_type:AxisType): self.end, self.axis_type, self.done = end, axis_type, False
+  def __iter__(self): return self
+  def __next__(self):
+    if not self.done:
+      self.done = True
+      _tk_range.user_rid += 1
+      self._rng = UOp.range(self.end, _tk_range.user_rid-1, axis_type=self.axis_type)
+      return self._rng
+    raise StopIteration
+
+class Kernel(AbstractContextManager):
+  def __init__(self, grid_size:tuple[int, int, int], block_size:int):
+    self.blockIdx_x = UOp.special(grid_size[0], "gidx0")
+    self.blockIdx_y = UOp.special(grid_size[1], "gidx1")
+    self.blockIdx_z = UOp.special(grid_size[2], "gidx2")
+    self.threadIdx_x = UOp.special(block_size, "lidx0")
+
+    self.range_stack = []
+    self.store_stack = []
+
+  @property
+  def warpid(self): return self.threadIdx_x // WARP_THREADS
+
+  def __enter__(self): return self
+  def __exit__(self, exc_type, exc_value, traceback): pass
+
+  def group(self, size:int): return Group(size, self)
+  @property
+  def warp(self): return self.group(1)
+  @property
+  def warpgroup(self): return self.group(4)
+
+  def range(self, end:int, axis_type:AxisType=AxisType.LOOP, track:bool=True):
+    rng = _tk_range(end, axis_type)
+    if track: self.range_stack.append(rng)
+    return rng
+
+  def push_store(self, store:UOp, uop:UOp): self.store_stack.append((store, uop))
+
+  def finish(self):
+    # end all ranges
+    rngs = []
+    while self.range_stack: rngs.append(self.range_stack.pop(0)._rng)
+
+    return self.store_stack.pop()[0].end(*rngs).sink(arg=KernelInfo(opts_to_apply=())).simplify()
+
+  def endrange(self):
+    last_store = self.store_stack.pop()
+    last_range = self.range_stack.pop()
+    return last_store[1].after(last_store[0].barrier().end(last_range._rng)).reshape(last_store[1].shape)

extra/thunder/tiny/tk/tiles.py

@@ -0,0 +1,52 @@
+import math
+from typing import cast, Callable
+from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
+from tinygrad.uop.ops import AxisType, UOp, KernelInfo, Ops
+from tinygrad.engine.realize import ExecItem, get_runner
+from tinygrad.dtype import AddrSpace, PtrDType
+from tinygrad.helpers import getenv, prod
+
+from extra.thunder.tiny.tk import WARP_THREADS
+
+class _Slots:
+  def __init__(self):
+    self.global_slot = 0
+    self.shared_slot = 0
+    self.register_slot = 0
+slots = _Slots()
+
+def gl(shape, dtype):
+  slots.global_slot += 1
+  return UOp.placeholder(shape, dtype, slot=slots.global_slot-1)
+
+shared_slot = 0
+def st(shape, dtype):
+  slots.shared_slot += 1
+  return UOp.placeholder(shape, dtype, addrspace=AddrSpace.LOCAL, slot=slots.shared_slot-1)
+
+TILE_ROW_DIM, TILE_COL_DIM = 16, 16
+RT_BASE_TILE_NE = TILE_ROW_DIM * TILE_COL_DIM
+RT_BASE_TILE_NEPT = RT_BASE_TILE_NE // WARP_THREADS
+register_slot = 0
+def rt(shape, dtype):
+  assert len(shape) == 2
+
+  height = shape[0] // TILE_ROW_DIM
+  width = shape[1] // TILE_COL_DIM
+
+  slots.register_slot += 1
+  return UOp.placeholder((height, width, RT_BASE_TILE_NEPT), dtype, addrspace=AddrSpace.REG, slot=slots.register_slot-1)
+
+def rv(length, dtype, layout="naive"):
+  tiles = length // TILE_ROW_DIM
+  match layout:
+    case "naive":
+      inner_dim = 1
+      outer_dim = (tiles + 1) // 2
+    case "ortho":
+      inner_dim = 1
+      outer_dim = tiles
+    case _: raise NotImplementedError(f"rv layout {layout} not implemented")
+
+  slots.register_slot += 1
+  return UOp.placeholder((outer_dim, inner_dim, 2), dtype, addrspace=AddrSpace.REG, slot=slots.register_slot-1)

extra/usbgpu/tbgpu/installer/Shared/TinyGPUViewModel.swift

@@ -119,14 +119,7 @@ extension TinyGPUViewModel: OSSystemExtensionRequestDelegate {
 
 		os_log("sysex actionForReplacingExtension: %@ %@", existing, ext)
 
-		// Add appropriate logic here to determine whether to replace the extension
-		// with the new extension. Common things to check for include
-		// testing whether the new extension's version number is newer than
-		// the current version number, or whether the bundleIdentifier is different.
-		// For simplicity, this sample always replaces the current extension
-		// with the new one.
 		replacementAction = .replace
-
 		self.state = .activating
 		return replacementAction
 	}

extra/usbgpu/tbgpu/installer/TinyGPUDriverExtension/TinyGPUDriverUserClient.cpp

@@ -7,30 +7,48 @@
 struct TinyGPUDriverUserClient_IVars
 {
 	OSSharedPtr<TinyGPUDriver> provider = nullptr;
+
+	TinyGPUCreateDMAResp *dmas = nullptr;
+	size_t dmaCount = 0;
+	size_t dmaCap = 0;
+
+	int ensureDMACap(size_t need)
+	{
+		// not thread-safe
+		if (need <= dmaCap) return 0;
+
+		size_t newCap = dmaCap ? dmaCap * 2 : 16;
+		while (newCap < need) newCap *= 2;
+
+		auto *newArr = IONewZero(TinyGPUCreateDMAResp, newCap);
+		if (!newArr) return -kIOReturnNoMemory;
+
+		if (dmas && dmaCount) {
+			memcpy(newArr, dmas, dmaCount * sizeof(TinyGPUCreateDMAResp));
+		}
+
+		IOSafeDeleteNULL(dmas, TinyGPUCreateDMAResp, dmaCap);
+		dmas = newArr;
+		dmaCap = newCap;
+		return 0;
+	}
 };
 
 bool TinyGPUDriverUserClient::init()
 {
-	auto theAnswer = super::init();
-	if (!theAnswer) {
-		return false;
-	}
+	auto ok = super::init();
+	if (!ok) return false;
 
 	ivars = IONewZero(TinyGPUDriverUserClient_IVars, 1);
-	if (ivars == nullptr) {
-		return false;
-	}
-
+	if (!ivars) return false;
 	return true;
 }
 
 void TinyGPUDriverUserClient::free()
 {
-	if (ivars != nullptr) {
-		ivars->provider.reset();
+	if (ivars) {
+		IOSafeDeleteNULL(ivars, TinyGPUDriverUserClient_IVars, 1);
 	}
-
-	IOSafeDeleteNULL(ivars, TinyGPUDriverUserClient_IVars, 1);
 	super::free();
 }
 
@@ -59,6 +77,22 @@ error:
 
 kern_return_t TinyGPUDriverUserClient::Stop_Impl(IOService* in_provider)
 {
+	// release all DMA allocations for this client
+	if (ivars) {
+		for (size_t i = 0; i < ivars->dmaCount; i++) {
+			auto &d = ivars->dmas[i];
+			if (d.dmaCmd) {
+				d.dmaCmd->CompleteDMA(kIODMACommandCompleteDMANoOptions);
+				d.dmaCmd->release();
+				d.dmaCmd = nullptr;
+			}
+		}
+		ivars->dmaCount = 0;
+		IOSafeDeleteNULL(ivars->dmas, TinyGPUCreateDMAResp, ivars->dmaCap);
+		ivars->dmas = nullptr;
+		ivars->provider.reset();
+	}
+
 	return Stop(in_provider, SUPERDISPATCH);
 }
 
@@ -102,26 +136,26 @@ kern_return_t TinyGPUDriverUserClient::ExternalMethod(uint64_t selector, IOUserC
 
 kern_return_t IMPL(TinyGPUDriverUserClient, CopyClientMemoryForType)
 {
-	if (!memory) {
-		return kIOReturnBadArgument;
-	}
-
-	if (ivars->provider.get() == nullptr) {
-		return kIOReturnNotAttached;
-	}
+	if (!memory) return kIOReturnBadArgument;
+	if (!ivars->provider.get()) return kIOReturnNotAttached;
 
+	// bar handling, type is bar num
 	if (type < 6) {
 		uint32_t bar = (uint32_t)type;
 		return ivars->provider->MapBar(bar, memory);
 	}
 
-	// dma page buffer
-	TinyGPUCreateDMAResp buf;
-	kern_return_t err = ivars->provider->CreateDMA(type, &buf);
-	if (err) {
-		return err;
+	// dma handling, type is size
+	if (ivars->ensureDMACap(ivars->dmaCount + 1)) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: cannot grow dma array");
+		return kIOReturnNoMemory;
 	}
 
+	TinyGPUCreateDMAResp buf{};
+	kern_return_t err = ivars->provider->CreateDMA(type, &buf);
+	if (err) return err;
+
+	ivars->dmas[ivars->dmaCount++] = buf;
 	*memory = buf.sharedBuf;
 	return 0;
 }

extra/weekly_commits_table.py

@@ -1,7 +1,7 @@
 # extra/weekly_commits_table.py
 import os, subprocess, datetime as dt
 
-NAMES = ["chenyu","George Hotz","nimlgen","qazal","Sieds Lykles","wozeparrot"]
+NAMES = ["chenyu","George Hotz","nimlgen","qazal","wozeparrot"]
 REPO  = os.environ.get("REPO_PATH",".")
 today = dt.date.today()
 days  = [(today - dt.timedelta(i)).strftime("%Y-%m-%d") for i in range(6,-1,-1)]
@@ -40,4 +40,4 @@ for d in days:
 print("** Commits by day (last 7) **")
 print("```")
 print("\n".join([header, rule] + rows))
-print("```")
+print("```")

pytest.ini

@@ -1,5 +1,8 @@
 [pytest]
-norecursedirs = extra
+norecursedirs =
+    extra
+    .hypothesis
+    .git
 timeout = 300
 timeout_method = thread
 timeout_func_only = true

ruff.toml

@@ -1,6 +1,6 @@
 indent-width = 2
 preview = true
-target-version = "py310"
+target-version = "py311"
 
 lint.select = [
   "F",  # Pyflakes

setup.cfg

@@ -0,0 +1,21 @@
+[mutmut]
+paths_to_mutate=tinygrad
+do_not_mutate=
+  tinygrad/apps/*
+  tinygrad/codegen/*
+  tinygrad/engine/*
+  tinygrad/nn/*
+  tinygrad/renderer/*
+  tinygrad/runtime/*
+  tinygrad/schedule/*
+  tinygrad/uop/*
+  tinygrad/viz/*
+  tinygrad/device.py
+  tinygrad/dtype.py
+  tinygrad/gradient.py
+  tinygrad/helpers.py
+  tinygrad/tensor.py
+tests_dir=
+  test/test_tiny.py
+  test/test_ops.py
+debug=true

setup.py

@@ -32,6 +32,7 @@ setup(name='tinygrad',
         'tinygrad.codegen.opt',
         'tinygrad.codegen.late',
         'tinygrad.engine',
+        'tinygrad.mixin',
         'tinygrad.nn',
         'tinygrad.renderer',
         'tinygrad.runtime',

test/external/external_benchmark_op_cat.py

@@ -0,0 +1,163 @@
+# ruff: noqa: E501 E712
+from tinygrad import dtypes, Device
+from tinygrad.uop.ops import UOp, AxisType, Ops, KernelInfo
+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, Invalid
+
+c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1576), (), 0)
+c2 = UOp.range(1576, 20, AxisType.LOOP)
+c5 = c2<55
+c6 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 16, 4)), (), 1)
+c8 = UOp.range(16, 0, AxisType.REDUCE)
+c11 = UOp.range(4, 1, AxisType.REDUCE)
+c14 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((14, 64, 4)), (), 2)
+c25 = c5.where((c2%4*4+c11+c8*16+c2//4*256), UOp.const(dtypes.index, Invalid))
+c27 = c6.index((c8*4+c11))*c14.index(c25)
+c29 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(55), (), 3)
+c30 = c5.where(c2, UOp.const(dtypes.index, Invalid))
+c34 = c5.where((c27.reduce(c8, c11, arg=Ops.ADD)+c29.index(c30)), UOp.const(dtypes.float, 0.0))
+c38 = c2<87
+c39 = (c5!=True)&c38
+c40 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 8, 4)), (), 4)
+c42 = UOp.range(8, 2, AxisType.REDUCE)
+c44 = UOp.range(4, 3, AxisType.REDUCE)
+c47 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((8, 32, 4)), (), 5)
+c49 = c2+1
+c51 = c49%4*4
+c57 = c49//4*128
+c61 = c39.where((c51+c44+c42*16+c57+-1792), UOp.const(dtypes.index, Invalid))
+c63 = c40.index((c42*4+c44))*c47.index(c61)
+c65 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(32), (), 6)
+c68 = c39.where((c2+-55), UOp.const(dtypes.index, Invalid))
+c71 = c39.where((c63.reduce(c42, c44, arg=Ops.ADD)+c65.index(c68)), UOp.const(dtypes.float, 0.0))
+c75 = c2<99
+c76 = (c38!=True)&c75
+c77 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 8, 4)), (), 7)
+c78 = UOp.range(8, 4, AxisType.REDUCE)
+c80 = UOp.range(4, 5, AxisType.REDUCE)
+c83 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((3, 32, 4)), (), 8)
+c90 = c76.where((c51+c80+c78*16+c57+-2816), UOp.const(dtypes.index, Invalid))
+c92 = c77.index((c78*4+c80))*c83.index(c90)
+c94 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(12), (), 9)
+c97 = c76.where((c2+-87), UOp.const(dtypes.index, Invalid))
+c100 = c76.where((c92.reduce(c78, c80, arg=Ops.ADD)+c94.index(c97)), UOp.const(dtypes.float, 0.0))
+c104 = c2<105
+c105 = (c75!=True)&c104
+c106 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 8, 4)), (), 10)
+c107 = UOp.range(8, 6, AxisType.REDUCE)
+c109 = UOp.range(4, 7, AxisType.REDUCE)
+c112 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((2, 32, 4)), (), 11)
+c119 = c105.where((c51+c109+c107*16+c57+-3200), UOp.const(dtypes.index, Invalid))
+c121 = c106.index((c107*4+c109))*c112.index(c119)
+c123 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(6), (), 12)
+c126 = c105.where((c2+-99), UOp.const(dtypes.index, Invalid))
+c129 = c105.where((c121.reduce(c107, c109, arg=Ops.ADD)+c123.index(c126)), UOp.const(dtypes.float, 0.0))
+c133 = c2<117
+c134 = (c104!=True)&c133
+c135 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 8, 4)), (), 13)
+c136 = UOp.range(8, 8, AxisType.REDUCE)
+c138 = UOp.range(4, 9, AxisType.REDUCE)
+c141 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((3, 32, 4)), (), 14)
+c143 = c2+3
+c145 = c143%4*4
+c149 = c143//4
+c150 = c149*128
+c154 = c134.where((c145+c138+c136*16+c150+-3456), UOp.const(dtypes.index, Invalid))
+c156 = c135.index((c136*4+c138))*c141.index(c154)
+c158 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(12), (), 15)
+c161 = c134.where((c2+-105), UOp.const(dtypes.index, Invalid))
+c164 = c134.where((c156.reduce(c136, c138, arg=Ops.ADD)+c158.index(c161)), UOp.const(dtypes.float, 0.0))
+c168 = c2<645
+c169 = (c133!=True)&c168
+c170 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 16, 4)), (), 16)
+c171 = UOp.range(16, 10, AxisType.REDUCE)
+c173 = UOp.range(4, 11, AxisType.REDUCE)
+c176 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((132, 64, 4)), (), 17)
+c180 = c149*256
+c184 = c169.where((c145+c173+c171*16+c180+-7680), UOp.const(dtypes.index, Invalid))
+c186 = c170.index((c171*4+c173))*c176.index(c184)
+c188 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(528), (), 18)
+c191 = c169.where((c2+-117), UOp.const(dtypes.index, Invalid))
+c194 = c169.where((c186.reduce(c171, c173, arg=Ops.ADD)+c188.index(c191)), UOp.const(dtypes.float, 0.0))
+c198 = c2<653
+c199 = (c168!=True)&c198
+c200 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 4, 4)), (), 19)
+c201 = UOp.range(4, 12, AxisType.REDUCE)
+c203 = UOp.range(4, 13, AxisType.REDUCE)
+c206 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((2, 16, 4)), (), 20)
+c215 = c199.where((c145+c203+c201*16+c149*64+-10368), UOp.const(dtypes.index, Invalid))
+c217 = c200.index((c201*4+c203))*c206.index(c215)
+c219 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(8), (), 21)
+c222 = c199.where((c2+-645), UOp.const(dtypes.index, Invalid))
+c225 = c199.where((c217.reduce(c201, c203, arg=Ops.ADD)+c219.index(c222)), UOp.const(dtypes.float, 0.0))
+c229 = c2<917
+c230 = (c198!=True)&c229
+c231 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 8, 4)), (), 22)
+c232 = UOp.range(8, 14, AxisType.REDUCE)
+c234 = UOp.range(4, 15, AxisType.REDUCE)
+c237 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((66, 32, 4)), (), 23)
+c244 = c230.where((c145+c234+c232*16+c150+-20992), UOp.const(dtypes.index, Invalid))
+c246 = c231.index((c232*4+c234))*c237.index(c244)
+c248 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(264), (), 24)
+c251 = c230.where((c2+-653), UOp.const(dtypes.index, Invalid))
+c254 = c230.where((c246.reduce(c232, c234, arg=Ops.ADD)+c248.index(c251)), UOp.const(dtypes.float, 0.0))
+c258 = c2<1061
+c259 = (c229!=True)&c258
+c260 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 16, 4)), (), 25)
+c261 = UOp.range(16, 16, AxisType.REDUCE)
+c263 = UOp.range(4, 17, AxisType.REDUCE)
+c266 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((36, 64, 4)), (), 26)
+c273 = c259.where((c145+c263+c261*16+c180+-58880), UOp.const(dtypes.index, Invalid))
+c275 = c260.index((c261*4+c263))*c266.index(c273)
+c277 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(144), (), 27)
+c280 = c259.where((c2+-917), UOp.const(dtypes.index, Invalid))
+c283 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(144), (), 28)
+c286 = c259.where(((c275.reduce(c261, c263, arg=Ops.ADD)+c277.index(c280))*c283.index(c280)), UOp.const(dtypes.float, 0.0))
+c290 = c2<1064
+c291 = (c258!=True)&c290
+c292 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 4, 4)), (), 29)
+c293 = UOp.range(4, 18, AxisType.REDUCE)
+c295 = UOp.range(4, 19, AxisType.REDUCE)
+c298 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 16, 4)), (), 30)
+c305 = c291.where((c2*4+c295+c293*16+-4244), UOp.const(dtypes.index, Invalid))
+c307 = c292.index((c293*4+c295))*c298.index(c305)
+c309 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(3), (), 31)
+c312 = c291.where((c2+-1061), UOp.const(dtypes.index, Invalid))
+c315 = c291.where((c307.reduce(c293, c295, arg=Ops.ADD)+c309.index(c312)), UOp.const(dtypes.float, 0.0))
+c317 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 128, 4)), (), 32)
+c321 = (c290!=True).where((c2+-1064), UOp.const(dtypes.index, Invalid))
+c323 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1), (), 33)
+c328 = c290.where(UOp.const(dtypes.float, 0.0), (c317.index(c321)*c323.index(UOp.const(dtypes.index, 0)).reciprocal()))
+c329 = c34+c71+c100+c129+c164+c194+c225+c254+c286+c315+c328
+c331 = c0.index(c2, ptr=True).store(c329).end(c2)
+ast = c331.sink(arg=KernelInfo(name="cat", opts_to_apply=None))
+
+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)
+
+ps = ProgramSpec("cat", src, Device.DEFAULT, ast, uops)
+# print(ps.src)
+# print(ps.applied_opts)
+# NOTE: this is faster with no GROUP and with NOLOCALS
+# (Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UNROLL, axis=19, arg=4), Opt(op=OptOps.UNROLL, axis=17, arg=4), Opt(op=OptOps.UNROLL, axis=15, arg=4), Opt(op=OptOps.UNROLL, axis=13, arg=4), Opt(op=OptOps.UNROLL, axis=11, arg=4), Opt(op=OptOps.UNROLL, axis=9, arg=4), Opt(op=OptOps.UNROLL, axis=7, arg=4), Opt(op=OptOps.UNROLL, axis=5, arg=4), Opt(op=OptOps.UNROLL, axis=3, arg=4), Opt(op=OptOps.UNROLL, axis=1, 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/external/external_benchmark_op_conv.py

@@ -1,237 +1,65 @@
-# ruff: noqa: E501
+# ruff: noqa: E501 E712
 from tinygrad import dtypes, Device
-from tinygrad.uop.ops import UOp, AxisType, Ops
+from tinygrad.uop.ops import UOp, AxisType, Ops, KernelInfo
 from tinygrad.codegen import full_rewrite
+# from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.renderer import ProgramSpec
 from tinygrad.engine.realize import CompiledRunner
-from tinygrad.helpers import dedup
+from tinygrad.helpers import dedup, getenv
 from tinygrad.device import Buffer
-from tinygrad.dtype import ImageDType
+from tinygrad.dtype import ImageDType, Invalid
 
-# 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
+# PYTHONPATH="." 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
 
-""" 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']
-"""
+def vision_conv_143():
+  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 1024, 4)), (), 0)
+  c2 = UOp.range(32, 3, AxisType.LOOP)
+  c5 = UOp.range(128, 4, AxisType.LOOP)
+  c8 = UOp.range(16, 2, AxisType.LOOP)
+  c16 = UOp.range(7, 0, AxisType.REDUCE)
+  c17 = c8*2+c16
+  c24 = ((c17<3)!=True)&(c17<35)
+  c26 = UOp.range(7, 1, AxisType.REDUCE)
+  c27 = c2*2+c26
+  c32 = ((c27<3)!=True)&(c27<67)
+  c34 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 1)
+  c38 = c5//2
+  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.index, Invalid))
+  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
+  c49 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((64, 49, 4)), (), 2)
+  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
+  c63 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(128), (), 3)
+  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
+  c67 = c0.index((c2*128+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)
 
-""" 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__']
-"""
+  opts = None
+  return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
 
-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()
+def vision_conv_153():
+  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((8, 1024, 4)), (), 0)
+  c2 = UOp.range(16, 3, AxisType.LOOP)
+  c5 = UOp.range(256, 4, AxisType.LOOP)
+  c8 = UOp.range(8, 2, AxisType.LOOP)
+  c16 = UOp.range(7, 0, AxisType.REDUCE)
+  c17 = c8*2+c16
+  c24 = ((c17<3)!=True)&(c17<19)
+  c26 = UOp.range(7, 1, AxisType.REDUCE)
+  c27 = c2*2+c26
+  c32 = ((c27<3)!=True)&(c27<35)
+  c34 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 1024, 4)), (), 1)
+  c38 = c5//2
+  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.index, Invalid))
+  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
+  c49 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((128, 49, 4)), (), 2)
+  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
+  c63 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(256), (), 3)
+  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
+  c67 = c0.index((c2*256+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)
+
+  opts = None
+  return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
+
+ast = vision_conv_143() if getenv("NUM", 143) == 143 else vision_conv_153()
 
 compiler = Device.default.compiler
 renderer = Device.default.renderer
@@ -240,20 +68,13 @@ 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))
+ps = ProgramSpec("conv", src, Device.DEFAULT, ast, uops)
 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])
-
+# 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)

test/external/external_fuzz_am_interrupts.py

@@ -1,39 +0,0 @@
-import subprocess
-import random
-import time
-from concurrent.futures import ThreadPoolExecutor, as_completed
-
-def run_test(i, full_run=False):
-  print(f"\rRunning iteration {i}...", end=" ", flush=True)
-
-  p = subprocess.Popen(['python3', 'test/test_tiny.py', 'TestTiny.test_plus'],  stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-
-  if not full_run:
-    time.sleep(random.uniform(0, 1200) / 1000)
-    p.kill()
-    _, stderr = p.communicate()
-  else:
-    _, stderr = p.communicate()
-
-  if full_run:
-    stderr_text = stderr.decode()
-    print(stderr_text)
-    assert "Ran 1 test in" in stderr_text and "OK" in stderr_text
-
-max_workers = 4
-with ThreadPoolExecutor(max_workers=max_workers) as executor:
-  futures = []
-  for i in range(1000000):
-    if i % 100 == 0:
-      for future in as_completed(futures):
-        try: future.result()
-        except Exception as e:
-          print(f"\nError in iteration: {e}")
-      futures = []
-
-      run_test(i, True)
-    else:
-      future = executor.submit(run_test, i, False)
-      futures.append(future)
-
-    if len(futures) > max_workers * 2: futures = [f for f in futures if not f.done()]

test/external/external_fuzz_hcq_mp.py

@@ -0,0 +1,44 @@
+import subprocess
+import random
+import time
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from tinygrad.helpers import getenv
+
+# checks that HCQ drivers can be killed during operation without causing issues
+
+def run_test(i, full_run=False, force_ok=False):
+  print(f"\rRunning iteration {i}...", end=" ", flush=True)
+
+  p = subprocess.Popen(["python3", "test/test_tiny.py", "TestTiny.test_plus"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+  if not full_run:
+    time.sleep(random.uniform(0, 1200) / 1000.0)
+    p.kill()
+    _, stderr = p.communicate()
+  else:
+    _, stderr = p.communicate()
+    stderr_text = stderr.decode()
+    assert ("Ran 1 test in" in stderr_text and "OK" in stderr_text) or (not force_ok and "Failed to take lock file" in stderr_text), stderr_text
+
+if __name__ == "__main__":
+  max_workers = getenv("MAX_WORKERS", 4)
+  with ProcessPoolExecutor(max_workers=max_workers) as executor:
+    futures = []
+    for i in range(1000000):
+      if i % 100 == 0:
+        # wait for everything we launched so far
+        for f in as_completed(futures):
+          try:
+            f.result()
+          except Exception as e:
+            print(f"\nError in iteration: {e}")
+        futures = []
+
+        # do a full run in the main proc
+        run_test(i, True, force_ok=True)
+      else:
+        futures.append(executor.submit(run_test, i, bool(getenv("FULL_RUN", 0))))
+
+      # keep list small
+      if len(futures) > max_workers * 2:
+        futures = [f for f in futures if not f.done()]

test/external/external_test_hcq_open.py

@@ -0,0 +1,20 @@
+import os
+if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
+
+import unittest, time
+from tinygrad import Device
+
+class TestOpen(unittest.TestCase):
+  def generate_test_open(n):
+    def test(self):
+      dev = Device[Device.DEFAULT]
+      for i in range(10):
+        dev.allocator.alloc(10 << 20)
+        time.sleep(0.5)
+    test.__name__ = f'test_open_{n}'
+    return test
+
+  for i in range(64): locals()[f'test_open_{i}'] = generate_test_open(i)
+
+if __name__ == '__main__':
+  unittest.main()

test/external/external_test_tk.py

@@ -0,0 +1,345 @@
+import unittest
+
+from tinygrad import Tensor, Device, dtypes, Context
+from tinygrad.engine.realize import ExecItem, get_runner
+
+from extra.thunder.tiny.tk import WARP_THREADS
+from extra.thunder.tiny.tk.kernel import Kernel
+from extra.thunder.tiny.tk.tiles import gl, st, rt, rv
+
+class TestTK(unittest.TestCase):
+  @unittest.skip("store from float rt is wrong")
+  def test_simple_matmul(self):
+    N = 32
+    BLOCK_SIZE = 16
+    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      c = gl((1, 1, N, N), dtypes.float32)
+      a = gl((1, 1, N, N), dtypes.bfloat16)
+      b = gl((1, 1, N, N), dtypes.bfloat16)
+
+      a_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      b_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      c_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      b_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      c_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      col, row = ker.blockIdx_x, ker.blockIdx_y
+
+      c_reg = warp.zero(c_reg)
+      for tile in ker.range(N // BLOCK_SIZE):
+        a_smem = warp.load(a_smem, a, (), (0, 0, row, tile), axis=2)
+        b_smem = warp.load(b_smem, b, (), (0, 0, tile, col), axis=2)
+
+        a_reg = warp.load(a_reg, a_smem)
+        b_reg = warp.load(b_reg, b_smem, transpose=True)
+
+        c_reg = warp.mma_AB(c_reg, a_reg, b_reg)
+      c_reg = ker.endrange()
+
+      c_smem = warp.store(c_smem, c_reg)
+      c = warp.store(c, c_smem, (0, 0, row, col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
+      b = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
+      c = Tensor.empty(1, 1, N, N, dtype="float32")
+      Tensor.realize(a, b, c)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (c, a, b)])
+    for _ in range(5): ei.run(wait=True)
+    c = c.float()
+
+    ref = a.matmul(b, dtype=dtypes.float32).float()
+
+    assert ref.allclose(c)
+
+  @unittest.skip("store from float rt is wrong")
+  def test_simple_matmul_transposed(self):
+    N = 32
+    BLOCK_SIZE = 16
+    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      c = gl((1, 1, N, N), dtypes.float32)
+      a = gl((1, 1, N, N), dtypes.bfloat16)
+      b = gl((1, 1, N, N), dtypes.bfloat16)
+
+      a_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      b_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      c_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      b_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      c_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      col, row = ker.blockIdx_x, ker.blockIdx_y
+
+      c_reg = warp.zero(c_reg)
+      for tile in ker.range(N // BLOCK_SIZE):
+        a_smem = warp.load(a_smem, a, (), (0, 0, row, tile), axis=2)
+        b_smem = warp.load(b_smem, b, (), (0, 0, col, tile), axis=2)
+
+        a_reg = warp.load(a_reg, a_smem)
+        b_reg = warp.load(b_reg, b_smem)
+
+        c_reg = warp.mma_ABt(c_reg, a_reg, b_reg)
+      c_reg = ker.endrange()
+
+      c_smem = warp.store(c_smem, c_reg)
+      c = warp.store(c, c_smem, (0, 0, row, col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
+      b = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
+      c = Tensor.empty(1, 1, N, N, dtype="float32")
+      Tensor.realize(a, b, c)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (c, a, b)])
+    for _ in range(5): ei.run(wait=True)
+    c = c.float()
+
+    ref = a.matmul(b.transpose(2, 3), dtype=dtypes.float32).float()
+
+    assert ref.allclose(c)
+
+  def test_load_store(self):
+    N = 32
+    BLOCK_SIZE = 16
+    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = gl((1, 1, N, N), dtypes.float32)
+      a = gl((1, 1, N, N), dtypes.float32)
+
+      a_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      col, row = ker.blockIdx_x, ker.blockIdx_y
+
+      a_smem = warp.load(a_smem, a, (), (0, 0, row, col), axis=2)
+      a_reg = warp.load(a_reg, a_smem)
+      b_reg = warp.copy(b_reg, a_reg)
+      b_smem = warp.store(b_smem, b_reg)
+      b = warp.store(b, b_smem, (0, 0, row, col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
+      b = Tensor.empty(1, 1, N, N, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float()
+
+    assert ref.allclose(b)
+
+  def test_max(self):
+    N = 16
+    BLOCK_SIZE = 16
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = gl((1, 1, N, N), dtypes.float32)
+      a = gl((1, 1, N, N), dtypes.float32)
+
+      a_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      max_reg = rv(BLOCK_SIZE, dtypes.float32, "ortho")
+
+      max_reg = warp.neg_inf(max_reg)
+
+      for tile_row in ker.range(N // BLOCK_SIZE):
+        for tile_col in ker.range(N // BLOCK_SIZE):
+          a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
+          a_reg = warp.load(a_reg, a_smem)
+          max_reg = warp.row_reduce(max_reg, a_reg, lambda a, b: a.maximum(b))
+        sum_reg = ker.endrange()
+
+        b_reg = warp.zero(b_reg).after(tile_row)
+        b_reg = warp.map(b_reg, lambda _, idx: sum_reg[idx[0], 0, (idx[2]%4)//2])
+        b_smem = warp.store(b_smem, b_reg)
+
+        for tile_col in ker.range(N // BLOCK_SIZE):
+          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
+      b = Tensor.empty(1, 1, N, N, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float().max(axis=3, keepdim=True).expand(a.shape)
+
+    assert ref.allclose(b)
+
+  def test_max_nonsquare(self):
+    N, M = 16, 64
+    BLOCK_N, BLOCK_M = 16, 64
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = gl((1, 1, N, M), dtypes.float32)
+      a = gl((1, 1, N, M), dtypes.float32)
+
+      a_smem = st((BLOCK_N, BLOCK_M), dtypes.float32)
+      b_smem = st((BLOCK_N, BLOCK_M), dtypes.float32)
+
+      a_reg = rt((BLOCK_N, BLOCK_M), dtypes.float32)
+      b_reg = rt((BLOCK_N, BLOCK_M), dtypes.float32)
+
+      max_reg = rv(BLOCK_N, dtypes.float32, "ortho")
+
+      max_reg = warp.zero(max_reg)
+
+      for tile_row in ker.range(N // BLOCK_N):
+        for tile_col in ker.range(M // BLOCK_M):
+          a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
+          a_reg = warp.load(a_reg, a_smem)
+          sum_reg = warp.row_reduce(max_reg, a_reg, lambda a, b: a.maximum(b))
+        sum_reg = ker.endrange()
+
+        b_reg = warp.zero(b_reg).after(tile_row)
+        b_reg = warp.map(b_reg, lambda _, idx: sum_reg[idx[0], 0, (idx[2]%4)//2])
+        b_smem = warp.store(b_smem, b_reg)
+
+        for tile_col in ker.range(M // BLOCK_M):
+          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, M, dtype="float32").contiguous()
+      b = Tensor.empty(1, 1, N, M, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float().max(axis=3, keepdim=True).expand(a.shape)
+
+    assert ref.allclose(b)
+
+  def test_sum(self):
+    N = 16
+    BLOCK_SIZE = 16
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = gl((1, 1, N, N), dtypes.float32)
+      a = gl((1, 1, N, N), dtypes.float32)
+
+      a_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      sum_reg = rv(BLOCK_SIZE, dtypes.float32, "ortho")
+
+      for tile_row in ker.range(N // BLOCK_SIZE):
+        sum_reg = warp.zero(sum_reg).after(tile_row)
+
+        for tile_col in ker.range(N // BLOCK_SIZE):
+          a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
+          a_reg = warp.load(a_reg, a_smem)
+          sum_reg = warp.row_reduce(sum_reg, a_reg, lambda a, b: a + b)
+        sum_reg = ker.endrange()
+
+        b_reg = warp.zero(b_reg).after(tile_row)
+        b_reg = warp.map(b_reg, lambda _, idx: sum_reg[idx[0], 0, (idx[2]%4)//2])
+        b_smem = warp.store(b_smem, b_reg)
+
+        for tile_col in ker.range(N // BLOCK_SIZE):
+          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
+      a = Tensor.arange(1 * 1 * N * N).reshape(1, 1, N, N).cast(dtypes.float32).contiguous()
+      b = Tensor.empty(1, 1, N, N, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float().sum(axis=3, keepdim=True).expand(a.shape)
+
+    assert ref.allclose(b)
+
+  def test_sum_nonsquare(self):
+    N, M = 16, 64
+    BLOCK_N, BLOCK_M = 16, 64
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = gl((1, 1, N, M), dtypes.float32)
+      a = gl((1, 1, N, M), dtypes.float32)
+
+      a_smem = st((BLOCK_N, BLOCK_M), dtypes.float32)
+      b_smem = st((BLOCK_N, BLOCK_M), dtypes.float32)
+
+      a_reg = rt((BLOCK_N, BLOCK_M), dtypes.float32)
+      b_reg = rt((BLOCK_N, BLOCK_M), dtypes.float32)
+
+      sum_reg = rv(BLOCK_N, dtypes.float32, "ortho")
+
+      sum_reg = warp.zero(sum_reg)
+
+      for tile_row in ker.range(N // BLOCK_N):
+        for tile_col in ker.range(M // BLOCK_M):
+          a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
+          a_reg = warp.load(a_reg, a_smem)
+          sum_reg = warp.row_reduce(sum_reg, a_reg, lambda a, b: a + b)
+        sum_reg = ker.endrange()
+
+        b_reg = warp.zero(b_reg).after(tile_row)
+        b_reg = warp.map(b_reg, lambda _, idx: sum_reg[idx[0], 0, (idx[2]%4)//2])
+        b_smem = warp.store(b_smem, b_reg)
+
+        for tile_col in ker.range(M // BLOCK_M):
+          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, M, dtype="float32").contiguous()
+      b = Tensor.empty(1, 1, N, M, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float().sum(axis=3, keepdim=True).expand(a.shape)
+
+    assert ref.allclose(b)
+
+if __name__ == "__main__":
+  unittest.main()

test/external/speed_v_theoretical.py

@@ -85,12 +85,11 @@ class TestKernelSpeed(unittest.TestCase):
     gbs = mems / tm / 1e9
     self._compare(tm, tflops, gbs, nv_tflops, nv_gbs, amd_tflops, amd_gbs)
 
-  # NOTE: tiny7 was slower than tiny12
   # TODO: why are convs so slow?!?
   def test_conv_3x3_256_32_32_256_256(self): self._test_conv_3x3(256, 32, 32, 256, 256, nv_tflops=27, amd_tflops=14)
 
   # theoretical is nv_tflops=165, amd_tflops=123
-  def test_gemm_4096(self): self._test_matmul(4096, nv_tflops=115, amd_tflops=65)
+  def test_gemm_4096(self): self._test_matmul(4096, nv_tflops=110, amd_tflops=65)
   def test_gemm_8192(self): self._test_matmul(8192, nv_tflops=115, amd_tflops=60)
 
   # theoretical is nv_gbs=1008, amd_gbs=960

test/test_custom_kernel.py

@@ -1,5 +1,6 @@
 import unittest
 from tinygrad import Tensor, UOp, Context
+from tinygrad.dtype import AddrSpace
 from tinygrad.uop.ops import KernelInfo, AxisType
 
 # **** kernels ****
@@ -8,16 +9,24 @@ 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_eye_kernel(C:UOp) -> UOp:
+  i = UOp.range(C.shape[0], 0)
+  j = UOp.range(C.shape[1], 1)
+  return C[i, j].store((i.eq(j)).cast(C.dtype.base)).end(i, j).sink(arg=KernelInfo(name=f"custom_eye_{C.size}"))
+
 def custom_add_one_kernel(B:UOp, A:UOp) -> UOp:
+  A,B = A.flatten(), B.flatten()
   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:
+  C,A,B = C.flatten(), A.flatten(), B.flatten()
   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:
+  C,D,A,B = C.flatten(), D.flatten(), A.flatten(), B.flatten()
   assert C.size == D.size
   i = UOp.range(C.size, 0)
   store_c = C[i].store(A[i]+B[i])
@@ -39,13 +48,43 @@ def custom_sum(B:UOp, A:UOp) -> UOp:
   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)])
+  i = UOp.range(dest.shape[0], 0)
+  j = UOp.range(dest.shape[1], 1, AxisType.UPCAST)
+  vec = src[i, j].contract(j)
+  store = UOp.group(*[dest[i, 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=()))
 
+def slice_sum_kernel(dest:UOp, src:UOp):
+  G = UOp.range(src.shape[0], 0)
+  slice_src = src[G, :]
+  reg = UOp.placeholder((1,), dest.dtype.base, 0, addrspace=AddrSpace.REG)
+  reg = reg.after(G)[0].set(0)
+  R = UOp.range(src.shape[1], 1, AxisType.REDUCE)
+  reg = reg[0].set(reg.after(R)[0] + slice_src[R], end=R)
+  ast = dest[G].set(reg[0], end=G)
+  return ast.sink(arg=KernelInfo(name=f"slice_sum_{src.shape[0]}_{src.shape[1]}", opts_to_apply=()))
+
+def simple_qkv_kernel(O:UOp, Q:UOp, K:UOp, V:UOp) -> UOp:
+  # attention without softmax
+  N, d = Q.shape[0], Q.shape[1]
+
+  i = UOp.range(N, 0)  # output row
+  d_out = UOp.range(d, 1)  # output column
+  j = UOp.range(N, 2, axis_type=AxisType.REDUCE)
+
+  k_inner = UOp.range(d, 3, axis_type=AxisType.REDUCE)
+  qk_acc = UOp.placeholder((1,), Q.dtype.base, 0, addrspace=AddrSpace.REG)
+  qk_acc = qk_acc.after(i, j)[0].set(0.0)
+  qk_acc = qk_acc[0].set(qk_acc.after(k_inner)[0] + Q[i, k_inner] * K[j, k_inner], end=k_inner)
+  qk_score = qk_acc[0] / (d ** 0.5)
+
+  out_acc = UOp.placeholder((1,), Q.dtype.base, 1, addrspace=AddrSpace.REG)
+  out_acc = out_acc.after(i, d_out)[0].set(0.0)
+  out_acc = out_acc[0].set(out_acc.after(j)[0] + qk_score * V[j, d_out], end=j)
+
+  store = O[i, d_out].store(out_acc[0])
+  return store.end(d_out).end(i).sink(arg=KernelInfo(name=f"simple_qkv_{N}_{d}", opts_to_apply=()))
+
 # **** backward callbacks ****
 
 def backward_gemm(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
@@ -91,6 +130,12 @@ class TestCustomKernel(unittest.TestCase):
     tst = tst.custom_kernel(fxn=custom_arange_kernel)[0]
     self.assertTrue((ref == tst).all().item())
 
+  def test_eye(self):
+    ref = Tensor.eye(1024).contiguous().realize()
+    tst = Tensor.empty_like(ref)
+    tst = tst.custom_kernel(fxn=custom_eye_kernel)[0]
+    self.assertTrue((ref == tst).all().item())
+
   def test_flip_contract(self):
     a = Tensor.randn(10,4)
     b = Tensor.empty_like(a)
@@ -111,6 +156,12 @@ class TestCustomKernel(unittest.TestCase):
     b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
     self.assertEqual(b.item(), 15)
 
+  def test_slice_sum(self):
+    A = Tensor.randn(16, 16).contiguous()
+    B = Tensor.empty(16)
+    B = Tensor.custom_kernel(B, A, fxn=slice_sum_kernel)[0]
+    self.assertTrue(B.allclose(A.sum(1)))
+
   def test_gemm(self):
     N = 16
     a = Tensor.randn(N, N)
@@ -152,5 +203,19 @@ class TestCustomKernel(unittest.TestCase):
     err = (grad_b - real_grad_b).square().max()
     self.assertLess(err.item(), 1e-6)
 
+  def test_simple_qkv(self):
+    N, d = 8, 4
+    Q = Tensor.randn(N, d)
+    K = Tensor.randn(N, d)
+    V = Tensor.randn(N, d)
+    O = Tensor.empty(N, d)
+
+    O_custom = Tensor.custom_kernel(O, Q, K, V, fxn=lambda o,q,k,v: simple_qkv_kernel(o,q,k,v))[0]
+    O_ref = ((Q @ K.T) / (d ** 0.5)) @ V
+
+    Tensor.realize(O_custom, O_ref)
+    err = (O_custom - O_ref).square().max()
+    self.assertLess(err.item(), 1e-6)
+
 if __name__ == '__main__':
   unittest.main()

test/test_linearizer.py

@@ -4,7 +4,7 @@ from dataclasses import replace
 
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.codegen.gpudims import get_grouped_dims
-from tinygrad.uop.ops import UOp, Ops, GroupOp
+from tinygrad.uop.ops import UOp, Ops, GroupOp, AxisType, PatternMatcher, graph_rewrite, UPat
 from tinygrad.device import Device, Buffer, is_dtype_supported
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.engine.realize import run_schedule, lower_schedule, CompiledRunner, get_program
@@ -38,6 +38,22 @@ class TestLinearizer(unittest.TestCase):
     np.testing.assert_equal(a.numpy(), ta)
     np.testing.assert_equal(b.numpy(), tb)
 
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx")
+  def test_late_bias_load(self):
+    img = Tensor.empty(1, 3, 16, 16)
+    w = Tensor.empty(16, 3, 3, 3)
+    b = Tensor.empty(16)
+    out = img.conv2d(w, b)
+    ast = helper_linearizer_opt(out)
+    uops = get_program(ast, opts=[]).uops
+    # slice at the last loop end
+    uslice = [i for i,u in enumerate(uops) if u.op == Ops.END][-1]
+    # only valid test if outermost range is the reduce
+    if uops[uslice].src[-1].arg[-1] == AxisType.REDUCE:
+      load_types = [u.src[0].dtype for u in uops[uslice+1:] if u.op == Ops.LOAD]
+      # assert that there is a global load after the reduce ends
+      assert any(dt.addrspace == AddrSpace.GLOBAL for dt in load_types)
+
   def _test_no_nested_ranges(self, lins, skip=None):
     for l in lins:
       range_in_acc = flatten([[x for x in u.src if x.op is Ops.RANGE] for u in l.uops if u.op is Ops.DEFINE_REG])
@@ -262,6 +278,8 @@ class TestLinearizer(unittest.TestCase):
     _assert_grouped_dims("gidx", (65536,), (16,16,256), False, [16,16,256], False)
     #                              2             ->             3
     _assert_grouped_dims("gidx", (128,128), (16,16,256), False, [16,16,64], False)
+    #                              2             ->             2
+    _assert_grouped_dims("gidx", (65536,2), (65535,65535,65535), False, [32768,4], False)
     # test when the only divisor is the square root of dim
     _assert_grouped_dims("gidx", (121,), (12,12,12), False, [11,11], False)
 
@@ -286,6 +304,27 @@ class TestLinearizer(unittest.TestCase):
     with self.assertRaises(RuntimeError):
       get_grouped_dims("gidx", (2,3,4,5,6), (16,16,16))
 
+    # TODO: In the above cases we only test if the shape after reshape is correct, never the indices.
+    # We should check if the returned indices are correct, for all cases.
+    # (65536, 2) -> (32768, 4)
+    dims, expected_limited_dims = (65536,2), (32768, 4)
+    idxs = get_grouped_dims("gidx", dims, (65535,65535,65535))
+    def match_div(): raise RuntimeError("match_div")
+    def match_mod(): raise RuntimeError("match_mod")
+    flat_idx_pattern = UPat(Ops.SPECIAL, arg='gidx0')*expected_limited_dims[1]+UPat(Ops.SPECIAL, arg='gidx1')
+    pm = PatternMatcher([
+      (flat_idx_pattern//dims[1], match_div),
+      (flat_idx_pattern%dims[1], match_mod)
+    ])
+
+    with self.assertRaises(RuntimeError) as error:
+      graph_rewrite(idxs[0], pm)
+    self.assertIn("match_div", str(error.exception))
+
+    with self.assertRaises(RuntimeError) as error:
+      graph_rewrite(idxs[1], pm)
+    self.assertIn("match_mod", str(error.exception))
+
     # # variable too large
     # with self.assertRaises(AssertionError):
     #   get_grouped_dims("gidx", (Variable("start_pos",0,16),3,4), (16,16,16), False,)
@@ -432,6 +471,8 @@ def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
   # now all input buffers in s[-1] should be realized
   # create fresh buffers for the outputs
   bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(s[-1].ast.src) else x for i,x in enumerate(s[-1].bufs)]
+  # ensure buffers are allocated
+  for b in bufs: b.ensure_allocated()
   return s[-1].ast, bufs
 
 def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs):

test/test_multitensor.py

@@ -596,6 +596,12 @@ class TestMultiTensor(unittest.TestCase):
     # ast are the same on devices
     self.assertEqual(len(set(asts)), 1)
 
+  def test_flip(self):
+    rng = Tensor.rand((10, 10, 10))
+    t0 = rng.shard(devices_2, axis=1)
+    out = t0.flip(0) + 1
+    self.assertTrue((rng.flip(0)+1).allclose(out.to(rng.device)))
+
   def test_reshape_on_axis(self):
     t0 = Tensor.rand((26, 15, 7)).shard(devices_3, axis=1)
 

test/test_ops.py

@@ -1551,8 +1551,10 @@ class TestOps(unittest.TestCase):
                                 lambda x: Tensor.stack(*x.std_mean(axis=(1,2))))
 
   def test_std_mean_loaded_nan(self):
-    helper_test_op([(1,0,3,0,5)], lambda x: torch.stack(torch.std_mean(x, axis=(1,3))),
-                                  lambda x: Tensor.stack(*x.std_mean(axis=(1,3))))
+    with warnings.catch_warnings():
+      warnings.filterwarnings("ignore", message="std_mean\\(\\): degrees of freedom is <= 0")
+      helper_test_op([(1,0,3,0,5)], lambda x: torch.stack(torch.std_mean(x, axis=(1,3))),
+                                    lambda x: Tensor.stack(*x.std_mean(axis=(1,3))))
   def test_softmax(self):
     helper_test_op([(45,65)], torch.nn.Softmax(dim=1), Tensor.softmax, atol=1e-7, grad_atol=1e-7)
     helper_test_op([(45)], torch.nn.Softmax(dim=0), Tensor.softmax, atol=1e-7, grad_atol=1e-7)
@@ -2820,13 +2822,13 @@ class TestOps(unittest.TestCase):
   @slow_test
   def test_slice_fancy_indexing_list_indices(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[[0]]], lambda x: x[[[0]]])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[0],b,c,d,:], lambda x: x[[0],j,k,o,:])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[((0,),)])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[(0,),b,c,d,:], lambda x: x[(0,),j,k,o,:])
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[[[[0]]],b,c,d,[[1]]], lambda x: x[[[[0]]],j,k,o,[[1]]])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[1,0,-1],b,c,d,:], lambda x: x[[1,0,-1],j,k,o,:])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,b,c,[1,2,3],...], lambda x: x[i,j,k,[1,2,3],...])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[(1,0,-1),b,c,d,:], lambda x: x[(1,0,-1),j,k,o,:])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,b,c,(1,2,3),...], lambda x: x[i,j,k,(1,2,3),...])
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,b,c,[[1],[2],[3]],...], lambda x: x[i,j,k,[[1],[2],[3]],...])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,[2,1,0],c,[-2,1,0],e], lambda x: x[i,[2,1,0],k,[-2,1,0],p])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,(2,1,0),c,(-2,1,0),e], lambda x: x[i,(2,1,0),k,(-2,1,0),p])
 
   @slow_test
   def test_slice_fancy_indexing_tuple_indices(self):
@@ -2841,11 +2843,10 @@ class TestOps(unittest.TestCase):
   @slow_test
   def test_slice_fancy_indexing_list_with_tensors(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[a]], lambda x: x[[i]])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[a,1]], lambda x: x[[i,1]])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[a,[1,1]]], lambda x: x[[i,[1,1]]])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[a,(1,1)]], lambda x: x[[i,(1,1)]])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[a,b,c,d,e]], lambda x: x[[i,j,k,o,p]])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[(a,)], lambda x: x[(i,)])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[(a,1)], lambda x: x[(i,1)])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[(a,(1,1))], lambda x: x[(i,(1,1))])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[(a,b,c,d,e)], lambda x: x[(i,j,k,o,p)])
 
   def test_slice_fancy_indexing_errors(self):
     a = Tensor.ones(10,11,12)

test/test_outerworld_transformer.py

@@ -0,0 +1,18 @@
+from tinygrad import Tensor, UOp
+from tinygrad.uop.ops import Ops, AxisType
+import unittest
+# this test is only focused on transformers and using range for the layers
+
+class TestOuterworldTransformer(unittest.TestCase):
+  def test_three_mats(self):
+    w = Tensor.empty(3, 1024, 1024)
+    inp = Tensor.empty(1, 1024)
+    i = UOp.range(3, -1, AxisType.OUTER)
+    inp_after = Tensor(inp.uop.after(i))
+    inp_gemm = inp_after@w[i]
+    inp = inp.uop.after(inp.uop.store(inp_gemm.uop).end(i)).contiguous()
+    inp = Tensor(inp)
+    inp.realize()
+
+if __name__ == "__main__":
+  unittest.main()

test/test_rangeify.py

@@ -300,6 +300,12 @@ class TestRangeify(unittest.TestCase):
     w2 = Tensor.empty(12, 8, 3, 3)
     x.conv2d(w1).conv2d(w2).realize()
 
+  def test_resnet_conv2d(self):
+    x = Tensor.empty(1, 8, 32, 32)
+    w1 = Tensor.empty(8, 8, 3, 3)
+    w2 = Tensor.empty(8, 8, 1, 1)
+    x.conv2d(w1).conv2d(w2).realize()
+
   def test_xception_conv2d(self):
     # NOTE: this fusion is bad, it's recomputing the inner many times
     x = Tensor.empty(1, 4, 32, 32)

test/test_schedule.py

@@ -1042,13 +1042,12 @@ class TestSchedule(unittest.TestCase):
       compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
       np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)
 
-    with Context(FUSE_ATTENTION=1):
-      out = Tensor.scaled_dot_product_attention(q,k,v)
-      run_schedule(check_schedule(out, 4)) # TODO: should be 1?
-      if getenv("CHECK", 1):
-        import torch
-        compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
-        np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)
+    out = Tensor.scaled_dot_product_attention(q,k,v)
+    run_schedule(check_schedule(out, 4)) # TODO: should be 1?
+    if getenv("CHECK", 1):
+      import torch
+      compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
+      np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)
 
   def test_ugly_reduceop_pairing(self):
     Tensor.manual_seed(0)
@@ -1574,6 +1573,13 @@ class TestSchedule(unittest.TestCase):
   def test_conv2d(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4)
   def test_conv2d_fused(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4)
 
+  def test_resnet_conv2d(self):
+    x = Tensor.empty(1, 8, 32, 32)
+    w1 = Tensor.empty(8, 8, 3, 3)
+    w2 = Tensor.empty(8, 8, 1, 1)
+    out = x.conv2d(w1).conv2d(w2)
+    check_schedule(out, 2)
+
   @unittest.skipUnless(is_dtype_supported(dtypes.half) and is_dtype_supported(dtypes.ulong), "need half and ulong")
   def test_conv2d_half(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4, dtype=dtypes.half)
   @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")

test/test_softmax_fusion.py

@@ -32,7 +32,7 @@ def run_one_schedule_item(out): lower_schedule_item(get_single_element(out.sched
 class TestFuse(unittest.TestCase):
   def _test_fuse(self, fxn, *args, atol=1e-6, allow_multiple=False, **kwargs):
     GlobalCounters.reset()
-    out_single = fxn(*args, **kwargs).fuse()
+    out_single = fxn(*args, **kwargs)
     if not allow_multiple: run_one_schedule_item(out_single)
     np_single = out_single.numpy()
     GlobalCounters.reset()
@@ -100,7 +100,7 @@ class TestFuse(unittest.TestCase):
     q = (x @ wq).contiguous()
     k = (x @ wk).contiguous()
     v = (x @ wv).contiguous()
-    attn = q.scaled_dot_product_attention(k, v).fuse()
+    attn = q.scaled_dot_product_attention(k, v)
     s = attn.schedule()
     self.assertEqual(len(s), 4) # 3 matmul and 1 attention
 
@@ -121,7 +121,7 @@ class TestFuse(unittest.TestCase):
   def test_mismatch_reduce(self):
     a = Tensor.ones(16, 10).contiguous().realize()
     b = Tensor.ones(16, 20).contiguous().realize()
-    c = (a.sum(axis=1) + b.sum(axis=1)).fuse()
+    c = (a.sum(axis=1) + b.sum(axis=1))
     self.assertListEqual(c.tolist(), [30]*16)
 
   @unittest.skipUnless(Device.DEFAULT == "METAL", "METAL TC")
@@ -129,7 +129,7 @@ class TestFuse(unittest.TestCase):
     A = Tensor.randn(8, 8).realize()
     B = Tensor.randn(8, 8).realize()
     C = Tensor.ones(1, 8, 8).pad(((1,1), None, None),).sum(0)
-    out = (C + (A @ B)).fuse()
+    out = (C + (A @ B))
     out.realize()
 
 class TestSoftmaxFusion(unittest.TestCase):
@@ -180,7 +180,7 @@ class TestSoftmaxFusion(unittest.TestCase):
 
     print("*** auto single kernel softmax ***")
     with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
-      out = self.test.contiguous().softmax(-1).fuse()
+      out = self.test.contiguous().softmax(-1)
       run_one_schedule_item(out)
 
     np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)

test/test_uops.py

@@ -517,7 +517,7 @@ class TestUOpStr(unittest.TestCase):
 
 class TestUPatHelpers(unittest.TestCase):
   def test_location(self):
-    self.assertEqual(sym.patterns[-1][0].location[0].replace("\\", "/").split("/")[-1], "symbolic.py")
+    self.assertEqual(sym.patterns[-1][0].location[0].replace("\\", "/").split("/")[-1], "math.py")
     self.assertEqual(shared_spec.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "spec.py")
     test_upat = UPat(Ops.CONST, dtypes.bool)
     self.assertEqual(test_upat.location[0].split("/")[-1], __file__.replace("\\", "/").split("/")[-1])

test/unit/test_device.py

@@ -30,7 +30,7 @@ class TestDevice(unittest.TestCase):
 
   @unittest.skipIf(WIN and CI, "skipping windows test") # TODO: subproccess causes memory violation?
   def test_env_overwrite_default_compiler(self):
-    expect_failure = "\ntry: assert Device[Device.DEFAULT].compiler is None;\nexcept RuntimeError: pass"
+    expect_failure = "\ntry: assert Device[Device.DEFAULT].compiler is None;\nexcept Exception: pass"
 
     if Device.DEFAULT == "CPU":
       from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangJITCompiler

test/unit/test_indexing.py

@@ -894,7 +894,7 @@ class TestNumpy(unittest.TestCase):
     a = Tensor([[1, 2, 3],
                 [4, 5, 6],
                 [7, 8, 9]])
-    self.assertIsNot(a[...], a)
+    self.assertIs(a[...], a)
     numpy_testing_assert_equal_helper(a[...], a)
     # `a[...]` was `a` in numpy <1.9.
     #numpy_testing_assert_equal_helper(data_ptr(a[...]), data_ptr(a))
@@ -1037,9 +1037,9 @@ class TestNumpy(unittest.TestCase):
     # Before `...` would return a itself.
     a = Tensor([5])
 
-    self.assertIsNot(a, a[()])
-    self.assertIsNot(a, a[...])
-    self.assertIsNot(a, a[:])
+    self.assertIs(a, a[()])
+    self.assertIs(a, a[...])
+    self.assertIs(a, a[:])
 
   def test_broaderrors_indexing(self):
     a = Tensor.zeros(5, 5)

test/unit/test_uop_symbolic.py

@@ -643,6 +643,10 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(lidx+(gidx//4)*8+2*(gidx%4), 0, 372, "(lidx+(gidx*2))")
     self.helper_test_variable(lidx+2*(gidx%4)+(gidx//4)*8, 0, 372, "(lidx+(gidx*2))")
 
+  def test_div_mod_recombine_partial(self):
+    gidx = Variable("gidx", 0, 15)
+    self.helper_test_variable((gidx//2)%4+(gidx//8)*4, 0, 7, "gidx//2")
+
   def test_div_mod_recombine_folded_mod(self):
     a = Variable("a", 0, 2)
     b = Variable("b", 0, 100)

tinygrad/codegen/__init__.py

@@ -4,7 +4,7 @@ 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
 from tinygrad.renderer import Renderer
-from tinygrad.dtype import dtypes
+from tinygrad.dtype import dtypes, PtrDType
 from tinygrad.helpers import panic
 
 # import all pattern matchers here
@@ -19,13 +19,19 @@ from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_s
 from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 
+pm_syntactic_sugar = PatternMatcher([
+  # INDEX on ptr INDEX concats them
+  (UPat(Ops.INDEX, name="i1").f(Ops.INDEX, name="i2", allow_any_len=True),
+   lambda i1,i2: i2.replace(src=i1.src+i2.src[1:]) if isinstance(i1.dtype, PtrDType) and not isinstance(i2.dtype, PtrDType) else None),
+])
+
 def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -> UOp:
   if ren is None: ren = Renderer()
 
   if SPEC: type_verify(sink, kernel_spec)
 
   # preprocess
-  sink = graph_rewrite(sink, pm_mops, name="early movement ops")
+  sink = graph_rewrite(sink, pm_mops+pm_syntactic_sugar, name="early movement ops", bottom_up=True)
 
   # first we optimize
   if optimize:
@@ -130,7 +136,7 @@ def full_rewrite(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
   """
 
   full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
-  assert len(full_sink.ranges) == 0, "all ranges must end by the sink"
+  assert len(full_sink.ranges) == 0, f"all ranges must end by the sink, {full_sink.ranges}"
   lst = line_rewrite(linearize(full_sink), pm_linearize_cleanups)
   if SPEC: type_verify(lst, program_spec)
   return lst

tinygrad/codegen/gpudims.py

@@ -26,28 +26,35 @@ def _split_dims(dims, max_sizes):
   return tuple(_dims[:2] if _dims[2] == 1 else _dims[0] if _dims[1:3] == [1,1] else _dims)
 
 def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|None, reverse=False) -> list[UOp]:
-  if reverse: dims = dims[::-1]
-  # try to group first: (a, b, c, d) -> (ab, c, d)
-  limited = (grouped if (grouped := _group_dims(dims, max_sizes)) else dims) if max_sizes is not None else dims
-  # check if grouping failed
-  if max_sizes is not None and len(limited) > len(max_sizes): raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
-  # try to split up dims: (a,) -> (b, c)
-  if limited == dims: limited = _split_dims(dims, max_sizes) if max_sizes is not None else dims
-  ret = raw_idxs = [UOp(Ops.SPECIAL, dtypes.index, (sint_to_uop(s),), (f"{prefix}{i}")) for i,s in enumerate(limited)]
+  if reverse: return get_grouped_dims(prefix, dims[::-1], max_sizes)[::-1]
+  if max_sizes is None: limited = dims
+  else:
+    # try to group first: (a, b, c, d) -> (ab, c, d)
+    limited = grouped if (grouped := _group_dims(dims, max_sizes)) else dims
+    # check if grouping failed
+    if len(limited) > len(max_sizes): raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
+    # try to split up dims: (a,) -> (b, c)
+    if limited == dims: limited = _split_dims(dims, max_sizes)
+  raw_idxs = [UOp(Ops.SPECIAL, dtypes.index, (sint_to_uop(s),), (f"{prefix}{i}")) for i,s in enumerate(limited)]
   if len(limited) < len(dims):
     ret = []
-    if (contraction:=get_contraction(dims, limited)) is None: raise AssertionError(f"get_contraction should not be None {dims=} {limited=}")
+    if (contraction:=get_contraction(dims, limited)) is None: raise RuntimeError(f"get_contraction should not be None {dims=} {limited=}")
     for idx, contraction_group in zip(raw_idxs, contraction):
       for c in contraction_group[:-1]:
         ret.append(idx % dims[c])
         idx //= dims[c]
       ret.append(idx)
-  elif len(limited) > len(dims):
-    a, b = len(limited), len(dims)
-    if a == 2 and b == 1: ret = [raw_idxs[0] * limited[1] + raw_idxs[1]]
-    if a == 3 and b == 1: ret = [raw_idxs[0] * (limited[1] * limited[2]) + raw_idxs[1] * limited[2] + raw_idxs[2]]
-    if a == 3 and b == 2: ret = [raw_idxs[0] * limited[1] + raw_idxs[1], raw_idxs[2]]
-  return ret[::-1] if reverse else ret
+    return ret
+  elif (a:=len(limited)) > (b:=len(dims)):
+    if a == 2 and b == 1: return [raw_idxs[0] * limited[1] + raw_idxs[1]]
+    if a == 3 and b == 1: return [(raw_idxs[0] * limited[1] + raw_idxs[1]) * limited[2] + raw_idxs[2]]
+    if a == 3 and b == 2: return [raw_idxs[0] * limited[1] + raw_idxs[1], raw_idxs[2]]
+  elif limited != dims:
+    # Convert to 1D
+    flat = raw_idxs[0]*limited[1]+raw_idxs[1] if len(dims) == 2 else raw_idxs[0]*(limited[1]*limited[2])+raw_idxs[1]*limited[2]+raw_idxs[2]
+    # Get back original indices from 1D
+    return [flat//dims[1], flat%dims[1]] if len(dims) == 2 else [flat//(dims[2]*dims[1]), (flat//dims[2])%dims[1], flat%dims[2]]
+  return raw_idxs
 
 def add_gpudims(ctx:Renderer, s:UOp):
   if s.arg is None: return None

tinygrad/codegen/late/devectorizer.py

@@ -4,7 +4,7 @@ from collections import defaultdict
 from dataclasses import dataclass
 from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace, Invalid, PtrDType
 from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, graph_rewrite, GroupOp, identity_element
-from tinygrad.uop.symbolic import uop_given_valid, parse_valid, sym, symbolic_flat, invalid_gate
+from tinygrad.uop.symbolic import uop_given_valid, parse_valid, sym, symbolic, invalid_gate
 from tinygrad.helpers import getenv, flatten, AMX, prod
 from tinygrad.renderer import Renderer
 
@@ -61,7 +61,7 @@ def expand_index(buf:UOp, vec:UOp):
   if getenv("UNSAFE_DISABLE_MASK", 0): vec = vec.get_idx()
   # generate the individual indexes
   midx = graph_rewrite(UOp.sink(*[buf.index(vec.gep(i), ptr=True) for i in range(vec.dtype.count)]),
-                       symbolic_flat+load_store_indexing, name=f"index_buf_{buf.arg}")
+                       symbolic+load_store_indexing, name=f"index_buf_{buf.arg}")
   # extract all the relevant offsets
   offsets_rootsrc: defaultdict[Any, dict[int, list[int]]] = defaultdict(dict)
   for i in range(vec.dtype.count):
@@ -299,7 +299,7 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp):
     ended_ranges = flatten([x.ended_ranges for x in topo if x.op is Ops.END])
     input_ranges = tuple([x for x in topo if x.op is Ops.RANGE and x not in reduce_range and x not in ended_ranges])
     identity = red.const(red.dtype, identity_element(red.arg, red.dtype.scalar()))
-    acc = UOp(Ops.DEFINE_REG, red.dtype.ptr(size=1, addrspace=AddrSpace.REG), arg=(ctx.acc_num,))
+    acc = UOp(Ops.DEFINE_REG, red.dtype.ptr(size=1, addrspace=AddrSpace.REG), arg=ctx.acc_num)
     acc_init = acc.after(*input_ranges).index(UOp.const(dtypes.int, 0)).store(identity) if len(input_ranges) else \
                acc.index(UOp.const(dtypes.int, 0)).store(identity)
     lst = [acc.after(acc_init, *reduce_range).index(UOp.const(dtypes.int, 0))] + lst  # put acc as the first element

tinygrad/codegen/late/linearizer.py

@@ -1,42 +1,59 @@
 import heapq
+from typing import Any
 from collections import defaultdict
-from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat
+from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat, multirange_str
+from tinygrad.helpers import prod, getenv, TUPLE_ORDER
 
-def linearize(u:UOp) -> list[UOp]:
+def linearize(sink:UOp) -> list[UOp]:
   # this is a toposort with priority
-  lst = list(u.toposort())
+  lst = list(sink.toposort())
   consumers: defaultdict[UOp, list[UOp]] = defaultdict(list)
   in_degree:dict[UOp, int] = {}
-  priorities:dict[UOp, int] = {}
+  out_degree:dict[UOp, int] = {}
+  priorities:dict[UOp, tuple[int, int, Any]] = {}
 
   # get consumers and assign priorities
   # NOTE: this requires the lst be locally toposorted
   for u in reversed(lst):
     for s in u.src: consumers[s].append(u)
     in_degree[u] = len(u.src)
-    # put loads in the beginning of the block and prevent priority inversion. hack for BARRIER grouping too
-    priority = [0] + [priorities[x] for x in consumers[u]]
-    if u.op is Ops.LOAD: priority.append(-1000)
-    if u.op is Ops.BARRIER: priority.append(-1500)
-    # ranges are scheduled as late as possible so anything that can be outside is
-    # if u.op is Ops.RANGE: priority = [2000]
-    if u.op is Ops.END: priority = [-1000]
-    # move defines and consts to the top
-    if u.op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.CONST}: priority.append(-2000)
-    priorities[u] = min(priority)
+    out_degree[u] = len(consumers[u])
+
+    # we place UOps with higher run_counts later
+    run_count = prod([int(r.vmax)+1 for r in u.ranges])
+
+    # simple priority override. this is all bottom up now, smaller numbers will be closer to the top
+    extra = None
+    match u.op:
+      # the order and placement of these defines is important
+      case Ops.DEFINE_GLOBAL: priority, extra = -20, u.arg
+      case Ops.DEFINE_VAR: priority, extra = -19, u.arg
+      case Ops.DEFINE_LOCAL: priority = -18
+      case Ops.DEFINE_REG: priority = -17
+      case Ops.CONST: priority = -10  # early consts
+      case Ops.LOAD: priority = -1    # place loads early
+      case Ops.STORE: priority = 1    # place stores late
+      case Ops.RANGE: priority = 5    # placing RANGE is good
+      case Ops.END: priority = -5     # placing END is bad
+      case _: priority = 0            # everything else has priority 0
+    priorities[u] = (run_count, priority, extra)
 
   # number the uops in "ideal" order
-  nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: (priorities[x],)+x.tuplize))}
+  nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: priorities[x]+(x.tuplize if TUPLE_ORDER else ())))}
 
-  # then force then to be toposorted in as close to the ideal order as possible
-  heapq.heapify(heap:=[(nkey[u],u) for u in lst if in_degree[u] == 0])
+  # then force them to be toposorted in as close to the ideal order as possible
+  heap = [(-nkey[sink], sink)]
   newlst = []
   while heap:
     newlst.append(u:=heapq.heappop(heap)[1])
-    for v in consumers[u]:
-      in_degree[v] -= 1
-      if in_degree[v] == 0: heapq.heappush(heap, (nkey[v],v))
-  assert len(newlst) == len(lst), f"len mismatch {len(newlst)} != {len(lst)}"
+    for v in u.src:
+      out_degree[v] -= 1
+      if out_degree[v] == 0: heapq.heappush(heap, (-nkey[v],v))
+  newlst = newlst[::-1]
+
+  if getenv("DEBUG_LINEARIZE"):
+    for i,u in enumerate(newlst):
+      print(f"{i:4d} {str(u.op):20s} {multirange_str(u.ranges, color=True, pad=10)} {priorities[u]}")
   return newlst
 
 class CFGContext:

tinygrad/codegen/opt/heuristic.py

@@ -81,7 +81,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
             return k
 
   # are we grouping? (requires local shape support)
-  if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= 2048, False):
+  if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= (128 if NOLOCALS else 2048), False):
     for sz in [16]:
       try:
         k.apply_opt(Opt(OptOps.GROUPTOP, 0, sz))
@@ -107,7 +107,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
   # potentially do more upcasts of non reduce axes based on a heuristic
   is_dsp = k.ren is not None and k.ren.device == "DSP"
   upcasted_axis: set[int] = set()
-  while resolve(prod(k.output_shape[i] for i in k.upcastable_dims) >= 1024):
+  while resolve(prod(k.output_shape[i] for i in k.upcastable_dims) >= 1024) and (k.upcast_size() < 32):
     xb_choices = []
     # consider all upcastable axes with 3 or 4 upcast (128 on the DSP)
     for axis, upcast_amount in itertools.product(k.upcastable_dims, ([128] if not len(upcasted_axis) else []) if is_dsp else [3,4]):
@@ -135,8 +135,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
   # if last reduce dim is small(ish), loop unroll the reduce
   # NOTE: this can fail on multireduce with mismatching dimensions, this is okay
   try:
-    upcast_size = prod(k.full_shape[a] for a in k.axes_of(AxisType.UPCAST, AxisType.UNROLL))
-    if k.unrollable_dims and (upcast_size <= 4 or not k.axes_of(AxisType.UNROLL)) and (upcast_size < 64):
+    if k.unrollable_dims and (k.upcast_size() <= 4 or not k.axes_of(AxisType.UNROLL)) and (k.upcast_size() < 64):
       if (s:=k.full_shape[k.unrollable_dims[-1]]) <= 32:
         k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, 0))
         # if it's small, upcast a second reduce dimension too

tinygrad/codegen/opt/postrange.py

@@ -2,7 +2,8 @@ from __future__ import annotations
 import math, itertools
 from collections import defaultdict
 from typing import cast, Final
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp, axis_letters, axis_colors
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp
+from tinygrad.uop.ops import axis_letters, axis_colors, axis_to_pos
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes, ImageDType
 from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten
@@ -12,10 +13,6 @@ from tinygrad.renderer import Renderer
 
 remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
 
-# NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
-axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
-               AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
-
 class Scheduler:
   def __init__(self, ast:UOp, ren:Renderer):
     self.ast, self.ren = ast, ren
@@ -105,6 +102,8 @@ class Scheduler:
   def ranges_of(self, *axis_type:AxisType) -> list[UOp]: return [r for r in self.rngs if r.arg[-1] in axis_type]
   def axes_of(self, *axis_type:AxisType) -> list[int]: return [i for i,t in enumerate(self.axis_types) if t in axis_type]
 
+  def upcast_size(self) -> int: return prod(self.full_shape[a] for a in self.axes_of(AxisType.UPCAST, AxisType.UNROLL))
+
   # copied from kernel.py
   @property
   def upcastable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP) \
@@ -217,8 +216,7 @@ class Scheduler:
     return ret
 
   def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> None|list[UOp]:
-    reduceops = [x for x in self.ast.toposort() if x.op is Ops.REDUCE]
-    if not len(reduceops): raise KernelOptError("no reduce ops for TensorCore")
+    if not (reduceops := self.reduceops): raise KernelOptError("no reduce ops for TensorCore")
     reduceop = reduceops[0]
     if use_tensor_cores and reduceop is not None and reduceop.arg is Ops.ADD:
       mul = reduceop.src[0] if reduceop.src[0].op is not Ops.CAST else reduceop.src[0].src[0]
@@ -312,9 +310,10 @@ class Scheduler:
 
   # helpers for hand_coded_optimizations
   @property
+  def reduceops(self) -> list[UOp]: return [x for x in self.ast.backward_slice if x.op is Ops.REDUCE]
+  @property
   def reduceop(self) -> UOp|None:
-    red = [x for x in self.ast.backward_slice if x.op is Ops.REDUCE]
-    if not len(red): return None
+    if not (red := self.reduceops): return None
     return UOp(Ops.REDUCE_AXIS, red[0].dtype, red[0].src, (red[0].arg, ()))
   @property
   def bufs(self) -> list[UOp]: return [x for x in self.ast.toposort() if x.op is Ops.INDEX][::-1]

tinygrad/codegen/simplify.py

@@ -1,6 +1,6 @@
 import itertools
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, ImageDType
-from tinygrad.uop.symbolic import symbolic_flat
+from tinygrad.uop.symbolic import symbolic
 from tinygrad.helpers import partition, dedup
 from tinygrad.dtype import dtypes
 
@@ -28,13 +28,12 @@ def simplify_merge_adjacent(u:UOp) -> UOp|None:
         s0, s1 = r0.src[0], r1.src[0]
         # do the merge
         new_range = r0.replace(src=(s0*s1,))
-        nidx = graph_rewrite(u, _substitute+symbolic_flat+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
+        nidx = graph_rewrite(u, _substitute+symbolic+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
                              name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")
 
         # check if it simplifies
         if count_divmod(nidx) <= count_divmod(u):
           u = nidx
-          continue
   return u
 
 pm_simplify_ranges = PatternMatcher([
@@ -109,7 +108,7 @@ pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
     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
+])+symbolic
 
 pm_reduce_load_collapse = pm_reduce_collapse + PatternMatcher([
   # lift x+y out of reduce on ne

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, AMD_LLVM
+from tinygrad.helpers import unwrap_class_type, suppress_finalizing, AMD_LLVM, select_first_inited
 from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
 from tinygrad.renderer import Renderer
 
@@ -54,7 +54,7 @@ atexit.register(lambda: [Device[dn].finalize() for dn in Device._opened_devices]
 
 @dataclass(frozen=True)
 class ProfileDeviceEvent(ProfileEvent):
-  device:str; comp_tdiff:decimal.Decimal=decimal.Decimal(0); copy_tdiff:decimal.Decimal=decimal.Decimal(0) # noqa: E702
+  device:str; comp_tdiff:decimal.Decimal=decimal.Decimal(0); copy_tdiff:decimal.Decimal=decimal.Decimal(0); props:dict[str,Any]|None=None # noqa: E702
 
 @dataclass(frozen=True)
 class ProfileProgramEvent(ProfileEvent): device:str; name:str; lib:bytes|None; base:int|None # noqa: E702
@@ -291,8 +291,8 @@ class Compiled:
     if len(enable_comps) > 1: raise RuntimeError(f"{self.device}: multiple compilers set in env {enable_comps}")
     for _, comp_pair in disable_comps: self.compilers.remove(comp_pair)
 
-    try: self.renderer, self.compiler = next(self._get_available_compilers([list(enable_comps)[0][1]] if len(enable_comps) == 1 else self.compilers))
-    except StopIteration as exc: raise RuntimeError(f"no usable compilers for {self.device}") from exc
+    self.renderer, self.compiler = select_first_inited([list(enable_comps)[0][1]] if len(enable_comps) == 1 else self.compilers,
+                                                       f"No compiler for {self.device} is available")
 
     if DEBUG >= 1: print(f"{self.device}: using {self.compiler.__class__.__name__}")
 
@@ -300,10 +300,6 @@ class Compiled:
     compiler_name = f"{unwrap_class_type(c).__name__.upper().removesuffix('COMPILER').removeprefix(devname:=self.device.split(':')[0].upper())}"
     return f"{devname}_{compiler_name if len(compiler_name) > 0 else unwrap_class_type(c).__name__.upper()}"
 
-  def _get_available_compilers(self, compilers) -> Iterator[tuple[Renderer, Compiler]]:
-    for renderer, compiler in compilers:
-      with contextlib.suppress(Exception): yield renderer(), compiler()
-
   def synchronize(self):
     """
     Synchronize all pending operations on the device.
@@ -343,6 +339,7 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
   # PYTHON supports half memoryview in 3.12+ https://github.com/python/cpython/issues/90751
   if dtype == dtypes.half:
     if device == "CL": return not CI and not OSX
+    if device == "QCOM": return False # QCOM compiler is flaky with half
     if device in ["CUDA", "NV"]: return not CI
     if device == "CPU" and CPU_LLVM: return OSX
     if device == "PYTHON": return sys.version_info >= (3, 12)

tinygrad/gradient.py

@@ -29,14 +29,14 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.MUL, name="ret"), lambda ctx, ret: (ret.src[1]*ctx, ret.src[0]*ctx)),
   (UPat(Ops.WHERE, name="ret"), lambda ctx, ret: (None, ret.src[0].where(ctx, ctx.const_like(0)), ret.src[0].where(ctx.const_like(0), ctx))),
   (UPat(Ops.REDUCE_AXIS, name="ret"), reduce_gradient),
-  (UPat((Ops.CONTIGUOUS, Ops.FUSE)), lambda ctx: (ctx,)),
+  (UPat(Ops.CONTIGUOUS), lambda ctx: (ctx,)),
   (UPat(Ops.CONTIGUOUS_BACKWARD), lambda ctx: (ctx.contiguous(),)),
   (UPat(Ops.RESHAPE, name="ret"), lambda ctx, ret: (ctx.reshape(ret.src[0].shape), None)),
   (UPat(Ops.EXPAND, name="ret"), lambda ctx, ret: (ctx.r(Ops.ADD,tuple(i for i,(s,n) in enumerate(zip(ret.src[0].shape, ret.shape)) if s!=n)), None)),
   (UPat(Ops.PAD, name="ret"), lambda ctx, ret: (ctx.shrink(tuple([(p[0], s+p[0]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
   (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[0], s-p[1]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
   (UPat(Ops.PERMUTE, name="ret"), lambda ctx, ret: (ctx.permute(argsort(ret.marg)),)),
-  (UPat(Ops.FLIP, name="ret"), lambda ctx, ret: (ctx.flip(ret.marg),)),
+  (UPat(Ops.FLIP, name="ret"), lambda ctx, ret: (ctx.flip([i for i,x in enumerate(ret.marg) if x]),)),
   (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)),

tinygrad/helpers.py

@@ -44,8 +44,7 @@ def fully_flatten(l):
     return flattened
   return [l]
 def fromimport(mod, frm): return getattr(__import__(mod, fromlist=[frm]), frm)
-def _is_balanced(s:str) -> bool:
-  return (acc:=list(itertools.accumulate([(1 if ch=='(' else -1 if ch==')' else 0) for ch in s])))[-1]==0 and all(x>=0 for x in acc)
+def _is_balanced(s:str) -> bool: return (d := 0, all((d := d + (c == '(') - (c == ')')) >= 0 for c in s))[1] and d == 0
 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
@@ -115,6 +114,13 @@ def suppress_finalizing(func):
       if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing
   return wrapper
 
+def select_first_inited(candidates:Sequence[Callable[...,T]|Sequence[Callable[...,T]]], err_msg: str) -> tuple[T,...]|T:
+  excs = []
+  for typ in candidates:
+    try: return tuple([cast(Callable, t)() for t in typ]) if isinstance(typ, Sequence) else cast(Callable, typ)()
+    except Exception as e: excs.append(e)
+  raise ExceptionGroup(err_msg, excs)
+
 def unwrap_class_type(cls_t): return cls_t.func if isinstance(cls_t, functools.partial) else cls_t
 
 def pluralize(st:str, cnt:int): return f"{cnt} {st}"+('' if cnt == 1 else 's')
@@ -171,7 +177,6 @@ DISABLE_COMPILER_CACHE = ContextVar("DISABLE_COMPILER_CACHE", 0)
 VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
 CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0)
-FUSE_ATTENTION = ContextVar("FUSE_ATTENTION", 0)
 EMULATE = ContextVar("EMULATE", "")
 CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
 CPU_LLVM, CPU_LVP, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0), ContextVar("AMD_LLVM", 1)
@@ -181,6 +186,8 @@ SPEC = ContextVar("SPEC", 1)
 IGNORE_OOB = ContextVar("IGNORE_OOB", 1)
 PCONTIG = ContextVar("PCONTIG", 0)  # partial contiguous in rangeify
 DEBUG_RANGEIFY = ContextVar("DEBUG_RANGEIFY", 0)
+# set to 1, this uses tuplize in the linearizer sort order
+TUPLE_ORDER = ContextVar("TUPLE_ORDER", 1)
 
 @dataclass(frozen=True)
 class Metadata:

tinygrad/mixin/__init__.py

@@ -0,0 +1,4 @@
+from tinygrad.mixin.math import MathMixin
+from tinygrad.mixin.movement import MovementMixin
+
+class OpMixin(MathMixin, MovementMixin): pass

tinygrad/mixin/math.py

@@ -0,0 +1,172 @@
+from typing import Self
+from tinygrad.uop import Ops
+from tinygrad.dtype import dtypes, ConstType
+
+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)

tinygrad/mixin/movement.py

@@ -0,0 +1,328 @@
+# mixins add syntactic sugar to Tensor and UOp
+import functools
+from typing import TypeAlias, TYPE_CHECKING, Self
+from tinygrad.uop import Ops
+from tinygrad.helpers import prod, argfix, flatten, dedup
+if TYPE_CHECKING: from tinygrad.uop.ops import UOp
+sint: TypeAlias = "UOp | int"
+
+def _align_left(*shapes:tuple[sint, ...]) -> tuple[tuple[sint, ...], ...]:
+  # unsqueeze left to make every shape same length
+  max_dim = max(len(shape) for shape in shapes)
+  return tuple((1,) * (max_dim - len(shape)) + shape for shape in shapes)
+
+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!
+  @property
+  def ndim(self) -> int:
+    """
+    Returns the number of dimensions in the tensor.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor([[1, 2], [3, 4]])
+    print(t.ndim)
+    ```
+    """
+    return len(self.shape)
+
+  def numel(self) -> sint:
+    """
+    Returns the total number of elements in the tensor.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
+    print(t.numel())
+    ```
+    """
+    return prod(self.shape)
+
+  def _resolve_dim(self, dim:int, *, extra:bool=False) -> int:
+    total = self.ndim + int(extra)
+    if not -max(1, total) <= dim <= max(1, total)-1: raise IndexError(f"{dim=} out of range {[-max(1, total), max(1, total)-1]}")
+    return dim + total if dim < 0 else dim
+
+  def _broadcast_to(self, new_shape:tuple[sint, ...]) -> Self:
+    if self.shape == new_shape: return self
+    if self.ndim > len(new_shape): raise ValueError(f"cannot broadcast tensor to fewer dimensions. shape={self.shape} to {new_shape=}")
+    # first unsqueeze left with 1s https://data-apis.org/array-api/latest/API_specification/broadcasting.html
+    shape, _ = _align_left(self.shape, new_shape)
+    # for each dimension, check either dim is 1, or it does not change
+    if not all(s == ns or s == 1 for s,ns in zip(shape, new_shape)):
+      raise ValueError(f"cannot broadcast {self.shape} to {new_shape=}")
+    reshaped = self.reshape(shape)
+    ret = reshaped._mop(Ops.EXPAND, arg=new_shape)
+    return reshaped if ret.shape == reshaped.shape else ret
+
+  def expand(self, shape, *args) -> Self:
+    """
+    Returns a tensor that is expanded to the shape that is specified.
+    Expand can also increase the number of dimensions that a tensor has.
+
+    Passing a `-1` or `None` to a dimension means that its size will not be changed.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor([1, 2, 3])
+    print(t.expand(4, -1).numpy())
+    ```
+    """
+    new_shape = tuple(from_ if to == -1 or to is None else to for from_, to in zip(*(_align_left(self.shape, argfix(shape, *args)))))
+    return self._broadcast_to(new_shape)
+
+  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})")
+    ret = self._mop(Ops.RESHAPE, arg=new_shape)
+    return self if ret.shape == self.shape else ret
+
+  def shrink(self, arg:tuple[tuple[sint, sint]|None, ...]) -> Self:
+    """
+    Returns a tensor that shrinks the each axis based on input arg.
+    `arg` must have the same length as `self.ndim`.
+    For each axis, it can be `None`, which means no shrink, or a tuple `(start, end)` that works the same as Python slice.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor.arange(9).reshape(3, 3)
+    print(t.numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(t.shrink(((None, (1, 3)))).numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(t.shrink((((0, 2), (0, 2)))).numpy())
+    ```
+    """
+    if self.ndim != len(arg): raise ValueError(f"{self.ndim=} != {len(arg)=}")
+    ret = self._mop(Ops.SHRINK, arg=[x if x is not None else (0,s) for x,s in zip(arg, self.shape)])
+    return self if ret.shape == self.shape else ret
+
+  def permute(self, order, *args) -> Self:
+    """
+    Returns a tensor that is a permutation of the original tensor.
+    The new tensor has the same data as the original tensor but with the dimensions permuted according to the order specified.
+    `order` can be passed as a tuple or as separate arguments.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor.empty(2, 3, 5)
+    print(t.shape)
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(t.permute(2, 0, 1).shape)
+    ```
+    """
+    order_arg = tuple(self._resolve_dim(x) for x in argfix(order, *args))
+    if sorted(order_arg) != list(range(self.ndim)): raise RuntimeError(f"order is not a valid permutation, getting {order_arg}")
+    return self._mop(Ops.PERMUTE, arg=order_arg) if order_arg != tuple(range(self.ndim)) else self
+
+  def flip(self, axis, *args) -> Self:
+    """
+    Returns a tensor that reverses the order of the original tensor along given `axis`.
+    `axis` can be passed as a tuple or as separate arguments.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor.arange(6).reshape(2, 3)
+    print(t.numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(t.flip(0).numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(t.flip((0, 1)).numpy())
+    ```
+    """
+    axis_arg = tuple(self._resolve_dim(x) for x in argfix(axis, *args))
+    assert all(not isinstance(x, bool) and x >= 0 and x < self.ndim for x in axis_arg), f"flip args must be axis ints {axis_arg}"
+    if len(axis_arg) != len(dedup(axis_arg)): raise RuntimeError(f"dim can appear at most once, getting {axis_arg}")
+    flip_arg = tuple([i in axis_arg for i in range(len(self.shape))])
+    return self._mop(Ops.FLIP, arg=flip_arg) if any(flip_arg) else self
+
+  # **** high level ****
+
+  def shrink_to(self, shape, *args) -> Self:
+    return self.shrink(tuple([None if ns is None else (0, ns) for ns in argfix(shape, *args)]))
+
+  def view(self, shape, *args) -> Self:
+    """`.view` is an alias for `.reshape`."""
+    return self.reshape(shape, *args)
+
+  def squeeze(self, dim:int|None=None) -> Self:
+    """
+    Returns a tensor with specified dimensions of input of size 1 removed.
+    If `dim` is not specified, all dimensions with size 1 are removed.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor.zeros(2, 1, 2, 1, 2)
+    print(t.squeeze().shape)
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(t.squeeze(0).shape)
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(t.squeeze(1).shape)
+    ```
+    """
+    if dim is None: return self.reshape(tuple(dim for dim in self.shape if dim != 1))
+    dim = self._resolve_dim(dim)
+    return self if not self.ndim or self.shape[dim] != 1 else self.reshape(self.shape[:dim] + self.shape[dim+1:])
+
+  def unsqueeze(self, dim:int) -> Self:
+    """
+    Returns a tensor with a new dimension of size 1 inserted at the specified `dim`.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor([1, 2, 3, 4])
+    print(t.unsqueeze(0).numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(t.unsqueeze(1).numpy())
+    ```
+    """
+    dim = self._resolve_dim(dim, extra=True)
+    return self.reshape(self.shape[:dim] + (1,) + self.shape[dim:])
+
+  @property
+  def T(self) -> Self:
+    """`.T` is an alias for `.transpose()`."""
+    return self.transpose()
+
+  def transpose(self, dim0=1, dim1=0) -> Self:
+    """
+    Returns a tensor that is a transposed version of the original tensor.
+    The given dimensions `dim0` and `dim1` are swapped.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor.arange(6).reshape(2, 3)
+    print(t.numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(t.transpose(0, 1).numpy())
+    ```
+    """
+    order = list(range(self.ndim))
+    order[dim0], order[dim1] = order[dim1], order[dim0]
+    return self.permute(order)
+
+  def flatten(self, start_dim=0, end_dim=-1) -> Self:
+    """
+    Flattens the tensor by reshaping it into a one-dimensional tensor.
+    If `start_dim` or `end_dim` are passed, only dimensions starting with `start_dim` and ending with `end_dim` are flattened.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor.arange(8).reshape(2, 2, 2)
+    print(t.flatten().numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(t.flatten(start_dim=1).numpy())
+    ```
+    """
+    start_dim, end_dim = self._resolve_dim(start_dim), self._resolve_dim(end_dim)
+    return self.reshape(self.shape[:start_dim] + (prod(self.shape[start_dim:end_dim+1]), ) + self.shape[end_dim+1:])
+
+  def unflatten(self, dim:int, sizes:tuple[int,...]) -> Self:
+    """
+    Unflattens dimension `dim` of the tensor into multiple dimensions specified by `sizes`. `Tensor.flatten()` is the inverse of this function.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor.ones(3, 4, 1).unflatten(1, (2, 2)).shape)
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor.ones(3, 4, 1).unflatten(1, (-1, 2)).shape)
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor.ones(5, 12, 3).unflatten(-2, (2, 2, 3, 1, 1)).shape)
+    ```
+    """
+    dim = self._resolve_dim(dim)
+    return self.reshape(self.shape[:dim] + sizes + self.shape[dim+1:])
+
+  def rearrange(self, formula:str, **sizes) -> Self:
+    """
+    Rearranges input according to formula
+
+    See: https://einops.rocks/api/rearrange/
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    x = Tensor([[1, 2], [3, 4]])
+    print(Tensor.rearrange(x, "batch channel -> (batch channel)").numpy())
+    ```
+    """
+    def parse_formula(formula: str):
+      tokens = f" {formula} ".replace("…", "...").replace("(", " ( ").replace(")", " ) ").replace(" ", "  ").replace(" 1 ", " ( ) ").split()
+      lparens, rparens = map(lambda x: [i for i, ch in enumerate(tokens) if ch == x], ("(", ")"))
+      pairs = list(zip(lparens, rparens))
+      assert len(lparens) == len(rparens) and sorted(flatten(pairs)) == flatten(pairs), "bracket mismatch"
+      return [name for name in tokens if name not in ("(", ")")], [(s - 2*i, e - 1 - 2*i) for i, (s, e) in enumerate(pairs)]
+
+    assert formula.count("->") == 1, 'need exactly one "->" in formula'
+
+    (lhs, unflatten_dims), (rhs, flatten_dims) = map(parse_formula, formula.split("->"))
+
+    for name in sizes: assert name in lhs, f"axis {name} is not used in transform"
+    assert sorted(lhs) == sorted(rhs) and len(lhs) == len(set(lhs)), f"name mismatch in {formula}"
+    for name in flatten((lhs, rhs)): assert name == "..." or (name.isidentifier() and "_" not in (name[0], name[-1])), f"invalid axis name {name}"
+    assert "..." not in flatten([lhs[s:e] for s, e in unflatten_dims]), f"cannot have collapsed ellipsis (...) in lhs of {formula}"
+    assert lhs.count("...") <= 1, f"too many ellipses in {formula}"
+
+    # resolve ellipsis
+    if "..." in lhs: ell_len = len(self.shape) - len(lhs) + 1 + sum(e - s - 1 for s, e in unflatten_dims)
+    lhs, rhs = map(lambda l: l[:(i:=l.index("..."))] + [f"...{j}" for j in range(ell_len)] + l[i + 1:] if "..." in l else l, (lhs, rhs))
+    unflatten_dims = [(s + (ell_len - 1 if "...0" in lhs[:s] else 0), e + (ell_len - 1 if "...0" in lhs[:e] else 0)) for s, e in unflatten_dims]
+    flatten_dims = [(s + (ell_len - 1 if "...0" in rhs[:s] else 0), e + (ell_len - 1 if "...0" in rhs[:e] else 0)) for s, e in flatten_dims]
+
+    # apply movement ops in order unflatten -> permute -> flatten/unsqueeze
+    t = functools.reduce(lambda x, dims: x.unflatten(dims[0], tuple(sizes.get(lhs[d], -1) for d in range(*dims))), unflatten_dims, self)
+    for i, name in enumerate(lhs): assert (name not in sizes) or sizes[name] == t.shape[i], f"size provided for dimension {name} incorrect"
+    t = t.permute([lhs.index(name) for name in rhs])
+    return functools.reduce(lambda x, dims: x.flatten(dims[0], dims[1] - 1) if dims[0]<dims[1] else x.unsqueeze(dims[0]), reversed(flatten_dims), t)
+
+  # *** movement ops with expand ***
+
+  def repeat_interleave(self, repeats:int, dim:int|None=None) -> Self:
+    """
+    Repeats elements of a tensor.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor([1, 2, 3])
+    print(t.repeat_interleave(2).numpy())
+    ```
+    """
+    x, dim = (self.flatten(), 0) if dim is None else (self, self._resolve_dim(dim))
+    shp = x.shape
+    return x.reshape(*shp[:dim+1], 1, *shp[dim+1:]).expand(*shp[:dim+1], repeats, *shp[dim+1:]).reshape(*shp[:dim], shp[dim]*repeats, *shp[dim+1:])
+
+  def repeat(self, repeats, *args) -> Self:
+    """
+    Repeats tensor number of times along each dimension specified by `repeats`.
+    `repeats` can be passed as a tuple or as separate arguments.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor([1, 2, 3])
+    print(t.repeat(4, 2).numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(t.repeat(4, 2, 1).shape)
+    ```
+    """
+    repeats = argfix(repeats, *args)
+    base_shape = _align_left(self.shape, repeats)[0]
+    unsqueezed_shape = flatten([[s] if r == 1 else [1, s] for r,s in zip(repeats, base_shape)])
+    expanded_shape = flatten([[s] if r == 1 else [r, s] for r,s in zip(repeats, base_shape)])
+    final_shape = [r*s for r,s in zip(repeats, base_shape)]
+    return self.reshape(unsqueezed_shape).expand(expanded_shape).reshape(final_shape)

tinygrad/nn/__init__.py

@@ -323,7 +323,7 @@ class Embedding:
     if not dtypes.is_int(idx.dtype): raise TypeError(f"Expected integer dtype for index in embedding, got {idx.dtype}")
     big_shp = idx.shape+(self.vocab_sz, self.embed_sz)
     arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1)).expand(big_shp), self.weight.expand(big_shp)
-    return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
+    return (arange == idx).where(vals, 0).sum(-2, dtype=vals.dtype)
 
 class LSTMCell:
   """

tinygrad/renderer/nir.py

@@ -148,7 +148,7 @@ class NIRRenderer(Renderer):
     (UPat(Ops.CAST, name="x"), lambda ctx,x: ncast(ctx.b, ctx.r[x.src[0]], x.src[0].dtype, x.dtype)),
     (UPat(Ops.BITCAST, src=(UPat.var("a"),), allow_any_len=True), lambda ctx,a: ctx.r[a]),
     (UPat(Ops.GEP, src=(UPat.var("a"),), name="x"), lambda ctx,x,a: nchannel(ctx.b, ctx.r[a], x.arg[0])),
-    (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x:mesa.nir_local_variable_create(ctx.b.impl, glsl_type(x.dtype), f"acc{x.arg[0]}".encode()).contents),
+    (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x:mesa.nir_local_variable_create(ctx.b.impl, glsl_type(x.dtype), f"acc{x.arg}".encode()).contents),
     (UPat(Ops.BARRIER), lambda ctx: nbarrier(ctx.b)),
     (UPat(Ops.IF, name="x"), lambda ctx,x: mesa.nir_push_if(ctx.b, ctx.r[x.src[0]])),
     (UPat(Ops.ENDIF, name="x"), lambda ctx,x: (lambda _: mesa.nir_def())(mesa.nir_pop_if(ctx.b, ctx.r[x.src[0]])))

tinygrad/runtime/autogen/comgr.py

@@ -15,7 +15,9 @@ PATHS_TO_TRY = [
 ]
 def _try_dlopen_amd_comgr():
   library = ctypes.util.find_library("amd_comgr")
-  if library: return ctypes.CDLL(library)
+  if library:
+    try: return ctypes.CDLL(library)
+    except OSError: pass
   for candidate in PATHS_TO_TRY:
     try: return ctypes.CDLL(candidate)
     except OSError: pass

tinygrad/runtime/autogen/mesa.py

@@ -15,7 +15,9 @@ PATHS_TO_TRY = [
 ]
 def _try_dlopen_tinymesa_cpu():
   library = ctypes.util.find_library("tinymesa_cpu")
-  if library: return ctypes.CDLL(library)
+  if library:
+    try: return ctypes.CDLL(library)
+    except OSError: pass
   for candidate in PATHS_TO_TRY:
     try: return ctypes.CDLL(candidate)
     except OSError: pass

tinygrad/runtime/autogen/nv/nv.py

@@ -7357,6 +7357,113 @@ class struct_c__SA_FWSECLIC_FRTS_CMD(Structure):
      ]
 
 FWSECLIC_FRTS_CMD = struct_c__SA_FWSECLIC_FRTS_CMD
+PCIEXPTBL_H = True # macro
+NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_BASE = 0x00 # macro
+NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_EXT = 0xE0 # macro
+PCI_EXP_ROM_SIGNATURE = 0xaa55 # macro
+PCI_EXP_ROM_SIGNATURE_NV = 0x4e56 # macro
+PCI_EXP_ROM_SIGNATURE_NV2 = 0xbb77 # macro
+def IS_VALID_PCI_ROM_SIG(sig):  # macro
+   return ((sig==0xaa55) or (sig==0x4e56) or (sig==0xbb77))
+OFFSETOF_PCI_EXP_ROM_SIG = 0x0 # macro
+OFFSETOF_PCI_EXP_ROM_NBSI_DATA_OFFSET = 0x16 # macro
+OFFSETOF_PCI_EXP_ROM_PCI_DATA_STRUCT_PTR = 0x18 # macro
+PCI_DATA_STRUCT_SIGNATURE = 0x52494350 # macro
+PCI_DATA_STRUCT_SIGNATURE_NV = 0x5344504E # macro
+PCI_DATA_STRUCT_SIGNATURE_NV2 = 0x53494752 # macro
+def IS_VALID_PCI_DATA_SIG(sig):  # macro
+   return ((sig==0x52494350) or (sig==0x5344504E) or (sig==0x53494752))
+# PCI_LAST_IMAGE = NVBIT ( 7 ) # macro
+PCI_ROM_IMAGE_BLOCK_SIZE = 512 # macro
+OFFSETOF_PCI_DATA_STRUCT_SIG = 0x0 # macro
+OFFSETOF_PCI_DATA_STRUCT_VENDOR_ID = 0x4 # macro
+OFFSETOF_PCI_DATA_STRUCT_LEN = 0xa # macro
+OFFSETOF_PCI_DATA_STRUCT_CLASS_CODE = 0xd # macro
+OFFSETOF_PCI_DATA_STRUCT_CODE_TYPE = 0x14 # macro
+OFFSETOF_PCI_DATA_STRUCT_IMAGE_LEN = 0x10 # macro
+OFFSETOF_PCI_DATA_STRUCT_LAST_IMAGE = 0x15 # macro
+NV_PCI_DATA_EXT_SIG = 0x4544504E # macro
+NV_PCI_DATA_EXT_REV_10 = 0x100 # macro
+NV_PCI_DATA_EXT_REV_11 = 0x101 # macro
+OFFSETOF_PCI_DATA_EXT_STRUCT_SIG = 0x0 # macro
+OFFSETOF_PCI_DATA_EXT_STRUCT_LEN = 0x6 # macro
+OFFSETOF_PCI_DATA_EXT_STRUCT_REV = 0x4 # macro
+OFFSETOF_PCI_DATA_EXT_STRUCT_SUBIMAGE_LEN = 0x8 # macro
+OFFSETOF_PCI_DATA_EXT_STRUCT_LAST_IMAGE = 0xa # macro
+OFFSETOF_PCI_DATA_EXT_STRUCT_FLAGS = 0xb # macro
+PCI_DATA_EXT_STRUCT_FLAGS_CHECKSUM_DISABLED = 0x04 # macro
+class struct__PCI_EXP_ROM_STANDARD(Structure):
+    pass
+
+struct__PCI_EXP_ROM_STANDARD._pack_ = 1 # source:False
+struct__PCI_EXP_ROM_STANDARD._fields_ = [
+    ('sig', ctypes.c_uint16),
+    ('reserved', ctypes.c_ubyte * 22),
+    ('pciDataStrucPtr', ctypes.c_uint16),
+    ('sizeOfBlock', ctypes.c_uint32),
+]
+
+PCI_EXP_ROM_STANDARD = struct__PCI_EXP_ROM_STANDARD
+PPCI_EXP_ROM_STANDARD = ctypes.POINTER(struct__PCI_EXP_ROM_STANDARD)
+class struct__PCI_EXP_ROM_NBSI(Structure):
+    pass
+
+struct__PCI_EXP_ROM_NBSI._pack_ = 1 # source:False
+struct__PCI_EXP_ROM_NBSI._fields_ = [
+    ('sig', ctypes.c_uint16),
+    ('reserved', ctypes.c_ubyte * 20),
+    ('nbsiDataOffset', ctypes.c_uint16),
+    ('pciDataStrucPtr', ctypes.c_uint16),
+    ('sizeOfBlock', ctypes.c_uint32),
+]
+
+PCI_EXP_ROM_NBSI = struct__PCI_EXP_ROM_NBSI
+PPCI_EXP_ROM_NBSI = ctypes.POINTER(struct__PCI_EXP_ROM_NBSI)
+class union__PCI_EXP_ROM(Union):
+    _pack_ = 1 # source:False
+    _fields_ = [
+    ('standard', PCI_EXP_ROM_STANDARD),
+    ('nbsi', PCI_EXP_ROM_NBSI),
+     ]
+
+PCI_EXP_ROM = union__PCI_EXP_ROM
+PPCI_EXP_ROM = ctypes.POINTER(union__PCI_EXP_ROM)
+class struct__PCI_DATA_STRUCT(Structure):
+    pass
+
+struct__PCI_DATA_STRUCT._pack_ = 1 # source:False
+struct__PCI_DATA_STRUCT._fields_ = [
+    ('sig', ctypes.c_uint32),
+    ('vendorID', ctypes.c_uint16),
+    ('deviceID', ctypes.c_uint16),
+    ('deviceListPtr', ctypes.c_uint16),
+    ('pciDataStructLen', ctypes.c_uint16),
+    ('pciDataStructRev', ctypes.c_ubyte),
+    ('classCode', ctypes.c_ubyte * 3),
+    ('imageLen', ctypes.c_uint16),
+    ('vendorRomRev', ctypes.c_uint16),
+    ('codeType', ctypes.c_ubyte),
+    ('lastImage', ctypes.c_ubyte),
+    ('maxRunTimeImageLen', ctypes.c_uint16),
+]
+
+PCI_DATA_STRUCT = struct__PCI_DATA_STRUCT
+PPCI_DATA_STRUCT = ctypes.POINTER(struct__PCI_DATA_STRUCT)
+class struct__NV_PCI_DATA_EXT_STRUCT(Structure):
+    pass
+
+struct__NV_PCI_DATA_EXT_STRUCT._pack_ = 1 # source:False
+struct__NV_PCI_DATA_EXT_STRUCT._fields_ = [
+    ('signature', ctypes.c_uint32),
+    ('nvPciDataExtRev', ctypes.c_uint16),
+    ('nvPciDataExtLen', ctypes.c_uint16),
+    ('subimageLen', ctypes.c_uint16),
+    ('privLastImage', ctypes.c_ubyte),
+    ('flags', ctypes.c_ubyte),
+]
+
+NV_PCI_DATA_EXT_STRUCT = struct__NV_PCI_DATA_EXT_STRUCT
+PNV_PCI_DATA_EXT_STRUCT = ctypes.POINTER(struct__NV_PCI_DATA_EXT_STRUCT)
 __all__ = \
     ['ACPI_DATA', 'ACPI_DSM_CACHE', 'ACPI_DSM_FUNCTION_COUNT',
     'ACPI_DSM_FUNCTION_CURRENT', 'ACPI_DSM_FUNCTION_GPS',
@@ -7480,12 +7587,16 @@ __all__ = \
     'NVDM_TYPE_UEFI_XTL_DEBUG_INTR',
     'NVGPU_ENGINE_CAPS_MASK_ARRAY_MAX', 'NVGPU_ENGINE_CAPS_MASK_BITS',
     'NV_ACPI_GENERIC_FUNC_COUNT',
+    'NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_BASE',
+    'NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_EXT',
     'NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_FLAGS_VERSION_AVAILABLE',
     'NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_FLAGS_VERSION_UNAVAILABLE',
     'NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V1',
     'NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V2',
     'NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V3',
     'NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V4',
+    'NV_PCI_DATA_EXT_REV_10', 'NV_PCI_DATA_EXT_REV_11',
+    'NV_PCI_DATA_EXT_SIG', 'NV_PCI_DATA_EXT_STRUCT',
     'NV_RPC_UPDATE_PDE_BAR_1', 'NV_RPC_UPDATE_PDE_BAR_2',
     'NV_RPC_UPDATE_PDE_BAR_INVALID', 'NV_RPC_UPDATE_PDE_BAR_TYPE',
     'NV_RPC_UPDATE_PDE_BAR_TYPE__enumvalues',
@@ -7775,8 +7886,31 @@ __all__ = \
     'NV_VGPU_PTE_64_INDEX_SHIFT', 'NV_VGPU_PTE_64_PAGE_SIZE',
     'NV_VGPU_PTE_64_SIZE', 'NV_VGPU_PTE_INDEX_MASK',
     'NV_VGPU_PTE_INDEX_SHIFT', 'NV_VGPU_PTE_PAGE_SIZE',
-    'NV_VGPU_PTE_SIZE', 'PACKED_REGISTRY_ENTRY',
-    'PACKED_REGISTRY_TABLE', 'REGISTRY_TABLE_ENTRY_TYPE_BINARY',
+    'NV_VGPU_PTE_SIZE', 'OFFSETOF_PCI_DATA_EXT_STRUCT_FLAGS',
+    'OFFSETOF_PCI_DATA_EXT_STRUCT_LAST_IMAGE',
+    'OFFSETOF_PCI_DATA_EXT_STRUCT_LEN',
+    'OFFSETOF_PCI_DATA_EXT_STRUCT_REV',
+    'OFFSETOF_PCI_DATA_EXT_STRUCT_SIG',
+    'OFFSETOF_PCI_DATA_EXT_STRUCT_SUBIMAGE_LEN',
+    'OFFSETOF_PCI_DATA_STRUCT_CLASS_CODE',
+    'OFFSETOF_PCI_DATA_STRUCT_CODE_TYPE',
+    'OFFSETOF_PCI_DATA_STRUCT_IMAGE_LEN',
+    'OFFSETOF_PCI_DATA_STRUCT_LAST_IMAGE',
+    'OFFSETOF_PCI_DATA_STRUCT_LEN', 'OFFSETOF_PCI_DATA_STRUCT_SIG',
+    'OFFSETOF_PCI_DATA_STRUCT_VENDOR_ID',
+    'OFFSETOF_PCI_EXP_ROM_NBSI_DATA_OFFSET',
+    'OFFSETOF_PCI_EXP_ROM_PCI_DATA_STRUCT_PTR',
+    'OFFSETOF_PCI_EXP_ROM_SIG', 'PACKED_REGISTRY_ENTRY',
+    'PACKED_REGISTRY_TABLE', 'PCIEXPTBL_H',
+    'PCI_DATA_EXT_STRUCT_FLAGS_CHECKSUM_DISABLED', 'PCI_DATA_STRUCT',
+    'PCI_DATA_STRUCT_SIGNATURE', 'PCI_DATA_STRUCT_SIGNATURE_NV',
+    'PCI_DATA_STRUCT_SIGNATURE_NV2', 'PCI_EXP_ROM',
+    'PCI_EXP_ROM_NBSI', 'PCI_EXP_ROM_SIGNATURE',
+    'PCI_EXP_ROM_SIGNATURE_NV', 'PCI_EXP_ROM_SIGNATURE_NV2',
+    'PCI_EXP_ROM_STANDARD', 'PCI_ROM_IMAGE_BLOCK_SIZE',
+    'PNV_PCI_DATA_EXT_STRUCT', 'PPCI_DATA_STRUCT', 'PPCI_EXP_ROM',
+    'PPCI_EXP_ROM_NBSI', 'PPCI_EXP_ROM_STANDARD',
+    'REGISTRY_TABLE_ENTRY_TYPE_BINARY',
     'REGISTRY_TABLE_ENTRY_TYPE_DWORD',
     'REGISTRY_TABLE_ENTRY_TYPE_STRING',
     'REGISTRY_TABLE_ENTRY_TYPE_UNKNOWN', 'RM_ENGINE_TYPE',
@@ -8343,6 +8477,8 @@ __all__ = \
     'struct_UpdateBarPde_v15_00',
     'struct_VIRTUAL_DISPLAY_GET_MAX_RESOLUTION_PARAMS',
     'struct_VIRTUAL_DISPLAY_GET_NUM_HEADS_PARAMS',
+    'struct__NV_PCI_DATA_EXT_STRUCT', 'struct__PCI_DATA_STRUCT',
+    'struct__PCI_EXP_ROM_NBSI', 'struct__PCI_EXP_ROM_STANDARD',
     'struct_alloc_object_FERMI_CONTEXT_SHARE_A_v04_00',
     'struct_alloc_object_FERMI_VASPACE_A_v03_00',
     'struct_alloc_object_GF100_DISP_SW_v03_00',
@@ -8595,7 +8731,7 @@ __all__ = \
     'union_NV2080_CTRL_FB_FS_INFO_QUERY_DATA_v26_04',
     'union_NV2080_CTRL_GRMGR_GR_FS_INFO_QUERY_DATA_v1A_1D',
     'union_NV2080_CTRL_INTERNAL_PFM_REQ_HNDLR_STATE_SYNC_DATA_type_v21_04',
-    'union_alloc_object_params_v25_08',
+    'union__PCI_EXP_ROM', 'union_alloc_object_params_v25_08',
     'union_alloc_object_params_v26_00',
     'union_alloc_object_params_v27_00',
     'union_alloc_object_params_v29_06', 'union_c__SA_GspFwWprMeta_0',

tinygrad/runtime/autogen/rocprof.py

@@ -8,11 +8,22 @@
 # LONGDOUBLE_SIZE is: 16
 #
 import ctypes, ctypes.util
+PATHS_TO_TRY = [
+  '/usr/local/lib/librocprof-trace-decoder.so',
+  '/usr/local/lib/librocprof-trace-decoder.dylib',
+]
+def _try_dlopen_rocprof_trace_decoder():
+  library = ctypes.util.find_library("rocprof-trace-decoder")
+  if library:
+    try: return ctypes.CDLL(library)
+    except OSError: pass
+  for candidate in PATHS_TO_TRY:
+    try: return ctypes.CDLL(candidate)
+    except OSError: pass
+  return None
 
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}
@@ -157,7 +168,7 @@ class FunctionFactoryStub:
 # You can either re-run clan2py with -l /path/to/library.so
 # Or manually fix this by comment the ctypes.CDLL loading
 _libraries = {}
-_libraries['FIXME_STUB'] = ctypes.CDLL(ctypes.util.find_library('rocprof-trace-decoder')) #  ctypes.CDLL('FIXME_STUB')
+_libraries['FIXME_STUB'] = _try_dlopen_rocprof_trace_decoder() #  ctypes.CDLL('FIXME_STUB')
 
 
 

tinygrad/runtime/ops_amd.py

@@ -7,7 +7,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, H
 from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator, hcq_filter_visible_devices
 from tinygrad.uop.ops import sint
 from tinygrad.device import Compiled, DMAFdRef, BufferSpec, CompilerPairT
-from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, suppress_finalizing, lo32, hi32, colored, prod
+from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, suppress_finalizing, lo32, hi32, colored, prod, ContextVar
 from tinygrad.renderer.cstyle import AMDRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt
@@ -19,7 +19,7 @@ from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_so
 from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationMeta, PCIDevice, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
 if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
 
-SQTT, PMC = getenv("SQTT", 0), getenv("PMC", 0)
+SQTT, SQTT_ITRACE_SE_MASK, PMC = ContextVar("SQTT", 0), ContextVar("SQTT_ITRACE_SE_MASK", 0b11), ContextVar("PMC", 0)
 EVENT_INDEX_PARTIAL_FLUSH = 4 # based on a comment in nvd.h
 WAIT_REG_MEM_FUNCTION_EQ  = 3 # ==
 WAIT_REG_MEM_FUNCTION_NEQ = 4 # !=
@@ -28,10 +28,10 @@ AQL_HDR = (1 << hsa.HSA_PACKET_HEADER_BARRIER) | (hsa.HSA_FENCE_SCOPE_SYSTEM <<
         | (hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCRELEASE_FENCE_SCOPE)
 
 @dataclass(frozen=True)
-class ProfileSQTTEvent(ProfileEvent): device:str; se:int; props:dict; blob:bytes; itrace:bool # noqa: E702
+class ProfileSQTTEvent(ProfileEvent): device:str; kern:str; se:int; blob:bytes; itrace:bool # noqa: E702
 
 @dataclass(frozen=True)
-class PMCSample: name:str; block:str; xcc:int; inst:int; se:int; sa:int; wgp:int; off:int; size:int; reg:str # noqa: E702
+class PMCSample: name:str; block:str; xcc:int; inst:int; se:int; sa:int; wgp:int; off:int; size:int; regsample:str # noqa: E702
 
 @dataclass(frozen=True)
 class ProfilePMCEvent(ProfileEvent): device:str; kern:str; sched:list[PMCSample]; blob:bytes # noqa: E702
@@ -72,14 +72,10 @@ class AMDComputeQueue(HWQueue):
     if self.dev.xccs > 1:
       self._q[prev_len-1] |= (len(self._q) - prev_len)
 
-  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_inst(self, n):
-    self.wreg(self.gc.regGRBM_GFX_INDEX, **{f'{f}_broadcast_writes': 1 for f in ['se', 'sh' if self.dev.target[0] == 9 else 'sa']}, instance_index=n)
-  def set_grbm_se_sh(self, se, sh):
-    self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, **{f'{"sh" if self.dev.target[0] == 9 else "sa"}_index':sh}, instance_broadcast_writes=1)
-  def set_grbm_se_sh_wgp(self, se, sh, wgp): self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, sa_index=sh, instance_index=wgp << 2)
-  def set_grbm_se(self, se): self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, sh_broadcast_writes=1, instance_broadcast_writes=1)
+  def set_grbm(self, instance=None, se=None, sh=None, wgp=None):
+    instance_val = (wgp << 2 | (instance or 0)) if wgp is not None else instance
+    self.wreg(self.gc.regGRBM_GFX_INDEX, **{(f'{key}_broadcast_writes' if val is None else f'{key}_index'): (1 if val is None else val)
+      for key, val in [('instance', instance_val), ('se', se), ('sh' if self.dev.target[0] == 9 else 'sa', sh)]})
 
   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)) \
@@ -140,7 +136,7 @@ class AMDComputeQueue(HWQueue):
   ### PMC ###
 
   def pmc_reset_counters(self, en=True):
-    self.set_grbm_broadcast()
+    self.set_grbm()
     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
@@ -158,18 +154,20 @@ class AMDComputeQueue(HWQueue):
         "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)')
+      # gfx11+ and later require even-numbered SQ *_SELECT registers
+      regsample = f'reg{block}_PERFCOUNTER{(pcid:=next(block2pid[block]))}'
+      if (regsel:=getattr(self.gc, (f'reg{block}_PERFCOUNTER{(pcid*2) if self.dev.target[0]>=11 and block=="SQ" else pcid}_SELECT'), None)) is None:
+        raise RuntimeError(f'{block} is out of perfcounter registers: ({regsample} 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))
+      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, regsample))
 
     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.set_grbm()
     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:
@@ -178,12 +176,10 @@ class AMDComputeQueue(HWQueue):
       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)
+            self.set_grbm(**({'instance':inst} if s.inst > 1 else ({'se':se_idx}|({'sh':sa_idx, 'wgp':wgp_idx} if self.dev.target[0] != 9 else {}))))
 
             # 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)
+            lo, hi = getattr(self.gc, f'{s.regsample}_LO'), getattr(self.gc, f'{s.regsample}_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))
 
@@ -201,11 +197,11 @@ class AMDComputeQueue(HWQueue):
       _0=sqtt.union_rgp_sqtt_marker_event_0(_0=sqtt.struct_rgp_sqtt_marker_event_0_0(has_thread_dims=1)),
       _2=sqtt.union_rgp_sqtt_marker_event_2(cmd_id=next(prg.dev.sqtt_next_cmd_id))), *global_size)
 
+    se_cap = max(prod([x if isinstance(x, int) else 1 for x in global_size]) // 4, 1) // 32
     for xcc in range(self.dev.xccs):
       with self.pred_exec(xcc_mask=1 << xcc):
         for i in range(8 if prg.dev.target >= (11,0,0) else 4):
-          self.wreg(getattr(self.gc, f'regCOMPUTE_STATIC_THREAD_MGMT_SE{i}'),
-            ((prg.dev.sqtt_itrace_se_mask >> ((self.dev.se_cnt // self.dev.xccs) * xcc + i)) & 0b1) if SQTT >= 2 else 0xffffffff)
+          self.wreg(getattr(self.gc, f'regCOMPUTE_STATIC_THREAD_MGMT_SE{i}'), min(0xffffffff, (1 << (se_cap + (1 if i == 0 else 0))) - 1))
 
   def sqtt_userdata(self, data, *extra_dwords):
     data_ints = [x[0] for x in struct.iter_unpack('<I', bytes(data))] + list(extra_dwords)
@@ -217,18 +213,18 @@ class AMDComputeQueue(HWQueue):
     self.wreg(self.gc.regSQ_THREAD_TRACE_CTRL, draw_event_en=1, spi_stall_en=1, sq_stall_en=1, reg_at_hwm=2, hiwater=1, util_timer=1,
       mode=int(tracing), **trace_ctrl)
 
-  def sqtt_start(self, buf0s:list[HCQBuffer], se_mask:int):
+  def sqtt_start(self, buf0s:list[HCQBuffer]):
     self.memory_barrier()
     if self.dev.target[0] == 9:
-      self.set_grbm_broadcast()
+      self.set_grbm()
       self.wreg(self.gc.regSQ_THREAD_TRACE_MASK, simd_en=0xf, cu_sel=0, sq_stall_en=1, spi_stall_en=1, reg_stall_en=1, vm_id_mask=0)
       for se in range(len(buf0s)):
         mask = (__SQTT_MISC:=1<<0) | (__SQTT_TIME:=1<<1) | (__SQTT_REG:=1<<2) | (__SQTT_WAVE_START:=1<<3) | (__SQTT_WAVE_END:=1<<6) \
              | (__SQTT_USERDATA:=1<<12) | (__SQTT_REG_CS:=1<<5) | (__SQTT_REG_CS_PRIV:=1<<15)
-        if (se_mask >> se) & 0b1: mask |= (__SQTTINST:=1<<10) | (__SQTT_INST_PC:=1<<11) | (__SQTT_ISSUE:=1<<13)
+        if (SQTT_ITRACE_SE_MASK.value >> se) & 0b1: mask |= (__SQTTINST:=1<<10) | (__SQTT_INST_PC:=1<<11) | (__SQTT_ISSUE:=1<<13)
 
         with self.pred_exec(xcc_mask=1<<(se // (ses_per_xcc:=(self.dev.se_cnt // self.dev.xccs)))):
-          self.set_grbm_se_sh(se % ses_per_xcc, 0)
+          self.set_grbm(se=se % ses_per_xcc, sh=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))
@@ -240,7 +236,7 @@ class AMDComputeQueue(HWQueue):
       self.spi_config(tracing=True)
       # One buffer for one SE, mesa does it with a single buffer and ac_sqtt_get_data_offset, but this is simpler and should work just as well
       for se in range(len(buf0s)):
-        self.set_grbm_se_sh(se, 0)
+        self.set_grbm(se=se, sh=0)
 
         buf0_lo, buf0_hi = data64_le(buf0s[se].va_addr >> 12)
         if self.dev.target >= (12,0,0):
@@ -263,7 +259,7 @@ class AMDComputeQueue(HWQueue):
         token_exclude = (1 << self.soc.SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT) if self.dev.target < (12,0,0) else 0
 
         # disable instr tracing
-        if not (se_mask >> se) & 0b1:
+        if not (SQTT_ITRACE_SE_MASK.value >> se) & 0b1:
           # gfx12 doesn't have enums with all fields, so it's hardcoded, but it's the same as gfx11.
           token_exclude |= (1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT | \
                             1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT | \
@@ -273,15 +269,15 @@ class AMDComputeQueue(HWQueue):
                   **({} if self.dev.target < (12,0,0) else {'exclude_barrier_wait': 1}))
         self.sqtt_config(tracing=True)
 
-    self.set_grbm_broadcast()
+    self.set_grbm()
     if self.dev.target[0] > 9: self.wreg(self.gc.regCOMPUTE_THREAD_TRACE_ENABLE, 1)
     self.memory_barrier()
     return self
 
   # Magic values from src/amd/common/ac_sqtt.c:ac_sqtt_emit_stop and src/amd/common/ac_sqtt.c:ac_sqtt_emit_wait
-  def sqtt_stop(self, ses:int, wptrs:HCQBuffer):
+  def sqtt_stop(self, wptrs:HCQBuffer):
     self.memory_barrier()
-    self.set_grbm_broadcast()
+    self.set_grbm()
 
     # Start shutting everything down
     if self.dev.target[0] == 9: self.wreg(self.gc.regSQ_THREAD_TRACE_MODE, mask_cs=1, autoflush_en=1, mode=0)
@@ -290,20 +286,21 @@ class AMDComputeQueue(HWQueue):
       self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.THREAD_TRACE_FINISH) | self.pm4.EVENT_INDEX(0))
 
     # For each SE wait for finish to complete and copy regSQ_THREAD_TRACE_WPTR to know where in the buffer trace data ends
-    for se in range(ses):
-      self.set_grbm_se_sh(se, 0)
+    for se in range(self.dev.se_cnt):
+      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, sh=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):
-        self.wait_reg_mem(reg=status_reg, mask=self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('finish_pending'), op=WAIT_REG_MEM_FUNCTION_EQ, value=0)
-        self.sqtt_config(tracing=False)
-      self.wait_reg_mem(reg=status_reg, mask=self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('busy'), op=WAIT_REG_MEM_FUNCTION_EQ, value=0)
-      self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.CS_PARTIAL_FLUSH) | self.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))
+        regstatus = 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):
+          self.wait_reg_mem(reg=regstatus, mask=self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('finish_pending'), op=WAIT_REG_MEM_FUNCTION_EQ, value=0)
+          self.sqtt_config(tracing=False)
+        self.wait_reg_mem(reg=regstatus, mask=self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('busy'), op=WAIT_REG_MEM_FUNCTION_EQ, value=0)
+        self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.CS_PARTIAL_FLUSH) | self.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))
 
-      # Copy WPTR to memory (src_sel = perf, dst_sel = tc_l2, wr_confirm = True)
-      self.pkt3(self.pm4.PACKET3_COPY_DATA, 1 << 20 | 2 << 8 | 4, self.gc.regSQ_THREAD_TRACE_WPTR.addr[0], 0, *data64_le(wptrs.va_addr+(se*4)))
+        # Copy WPTR to memory (src_sel = perf, dst_sel = tc_l2, wr_confirm = True)
+        self.pkt3(self.pm4.PACKET3_COPY_DATA, 1 << 20 | 2 << 8 | 4, self.gc.regSQ_THREAD_TRACE_WPTR.addr[0], 0, *data64_le(wptrs.va_addr+(se*4)))
 
-    self.set_grbm_broadcast()
+    self.set_grbm()
     if self.dev.target[0] > 9: self.spi_config(tracing=False)
     self.memory_barrier()
     return self
@@ -584,12 +581,31 @@ class AMDProgram(HCQProgram):
     weakref.finalize(self, self._fini, self.dev, self.lib_gpu, buf_spec)
 
   def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
+    if self.dev.sqtt_enabled: cast(AMDComputeQueue, self.dev.hw_compute_queue_t()).sqtt_start(self.dev.sqtt_buffers).submit(self.dev)
     res = super().__call__(*bufs, global_size=global_size, local_size=local_size, vals=vals, wait=wait)
     if self.dev.pmc_enabled:
       cast(AMDComputeQueue, self.dev.hw_compute_queue_t()).pmc_read(self.dev.pmc_buffer, self.dev.pmc_sched) \
                                                           .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
       self.dev.allocator._copyout(pmc_buf:=memoryview(bytearray(self.dev.pmc_buffer.size)), self.dev.pmc_buffer)
       Compiled.profile_events += [ProfilePMCEvent(self.dev.device, self.name, self.dev.pmc_sched, bytes(pmc_buf))]
+    if self.dev.sqtt_enabled:
+      cast(AMDComputeQueue, self.dev.hw_compute_queue_t()).sqtt_stop(self.dev.sqtt_wptrs) \
+                                                          .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
+      self.dev.synchronize()
+
+      for se, buf in enumerate(self.dev.sqtt_buffers):
+        wptr = ((self.dev.sqtt_wptrs.cpu_view().view(fmt='I')[se]&0x1FFFFFFF)-(((buf.va_addr//32)&0x1FFFFFFF) if self.dev.target[0] == 11 else 0))*32
+
+        if DEBUG >= 5: print(f'\t{self.dev.device}: SE {se} blob size {wptr:#x}')
+        assert wptr >= 0 and wptr <= buf.size, f"{wptr} > {buf.size}, should never happen"
+
+        # When sqtt buffer overflows, wptr stops at the last dword
+        if wptr >= buf.size - 32:
+          print(colored(f"{self.dev.device}: Warning: SQTT buffer is full (SE {se})! Increase SQTT buffer with SQTT_BUFFER_SIZE=X (in MB)", "yellow"))
+
+        self.dev.allocator._copyout(sqtt_mv:=memoryview(bytearray(wptr)), buf)
+        resbuf = (struct.pack('<Q', 0x11 | (4 << 13) | (0xf << 16) | (se << 24)) + bytes(sqtt_mv)) if self.dev.target[0] == 9 else bytes(sqtt_mv)
+        Compiled.profile_events += [ProfileSQTTEvent(self.dev.device, self.name, se, resbuf, bool((SQTT_ITRACE_SE_MASK.value >> se) & 1))]
     return res
 
 class AMDAllocator(HCQAllocator['AMDDevice']):
@@ -815,7 +831,7 @@ class PCIIface(PCIIfaceBase):
 
 class USBIface(PCIIface):
   def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
-    self.dev, self.pci_dev = dev, USBPCIDevice(f"usb:{dev_id}", bars=[0, 2, 5])
+    self.dev, self.pci_dev = dev, USBPCIDevice(dev.__class__.__name__[:2], f"usb:{dev_id}", bars=[0, 2, 5])
     self._setup_adev(self.pci_dev, dma_regions=[(0x200000, self.pci_dev.dma_view(0xf000, 0x80000))])
     self.pci_dev.usb._pci_cacheable += [(self.pci_dev.bar_info[2].addr, self.pci_dev.bar_info[2].size)] # doorbell region is cacheable
 
@@ -932,9 +948,8 @@ class AMDDevice(HCQCompiled):
 
       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
+      self.sqtt_wptrs = self.allocator.alloc(round_up(self.se_cnt * 4, 0x1000), BufferSpec(cpu_access=True, nolru=True))
       self.sqtt_next_cmd_id = itertools.count(0)
-      cast(AMDComputeQueue, self.hw_compute_queue_t()).sqtt_start(self.sqtt_buffers, self.sqtt_itrace_se_mask).submit(self)
 
   def create_queue(self, queue_type, ring_size, ctx_save_restore_size=0, eop_buffer_size=0, ctl_stack_size=0, debug_memory_size=0):
     ring = self.iface.alloc(ring_size, uncached=True, cpu_access=True)
@@ -984,21 +999,4 @@ class AMDDevice(HCQCompiled):
 
   def on_device_hang(self): self.iface.on_device_hang()
 
-  def _at_profile_finalize(self):
-    if self.sqtt_enabled:
-      wptrs_buf = self.allocator.alloc(round_up(len(self.sqtt_buffers), 0x1000), BufferSpec(cpu_access=True, nolru=True))
-      cast(AMDComputeQueue, self.hw_compute_queue_t()).sqtt_stop(len(self.sqtt_buffers), wptrs_buf) \
-                                                      .signal(self.timeline_signal, self.next_timeline()).submit(self)
-      self.synchronize()
-      if DEBUG >= 2: print(f'{self.device}: Saving SQTT in profile...')
-      for i,buf0 in enumerate(self.sqtt_buffers):
-        wptr = ((wptrs_buf.cpu_view().view(fmt='I')[i] & 0x1FFFFFFF) - (((buf0.va_addr//32) & 0x1FFFFFFF) if self.target[0] == 11 else 0)) * 32
-        if DEBUG >= 2: print(f'\t{self.device}: SE {i} blob size {wptr:#x}')
-        assert wptr >= 0 and wptr <= buf0.size, f"{wptr} > {buf0.size}, should never happen"
-        # When sqtt buffer overflows, wptr stops at the last dword
-        if wptr >= buf0.size - 32:
-          print(colored(f"{self.device}: Warning: SQTT buffer is full (SE {i})! Increase SQTT buffer with SQTT_BUFFER_SIZE=X (in MB)", "yellow"))
-        self.allocator._copyout(sqtt_buf:=memoryview(bytearray(wptr)), buf0)
-        if self.target[0] == 9: sqtt_buf = memoryview(struct.pack('<Q', 0x11 | (4 << 13) | (0xf << 16) | (i << 24)) + sqtt_buf)
-        Compiled.profile_events += [ProfileSQTTEvent(self.device, i, self.iface.props, bytes(sqtt_buf), bool((self.sqtt_itrace_se_mask >> i) & 0b1))]
-    super()._at_profile_finalize()
+  def device_props(self): return self.iface.props

tinygrad/runtime/ops_nv.py

@@ -389,12 +389,12 @@ class NVKIface:
 
   def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, map_flags=0, cpu_addr=None, **kwargs) -> HCQBuffer:
     # Uncached memory is "system". Use huge pages only for gpu memory.
-    page_size = (4 << (12 if OSX else 10)) if uncached or host else ((2 << 20) if size >= (8 << 20) else (4 << (12 if OSX else 10)))
+    page_size = mmap.PAGESIZE if uncached or host else ((2 << 20) if size >= (8 << 20) else (mmap.PAGESIZE if MOCKGPU else 4 << 10))
     size = round_up(size, page_size)
-    va_addr = self._alloc_gpu_vaddr(size, alignment=page_size, force_low=cpu_access)
+    va_addr = self._alloc_gpu_vaddr(size, alignment=page_size, force_low=cpu_access) if (alloced:=cpu_addr is None) else cpu_addr
 
     if host:
-      va_addr = cpu_addr or FileIOInterface.anon_mmap(va_addr, size, mmap.PROT_READ|mmap.PROT_WRITE, MAP_FIXED|mmap.MAP_SHARED|mmap.MAP_ANONYMOUS, 0)
+      if alloced: va_addr = FileIOInterface.anon_mmap(va_addr, size, mmap.PROT_READ|mmap.PROT_WRITE, MAP_FIXED|mmap.MAP_SHARED|mmap.MAP_ANONYMOUS, 0)
 
       flags = (nv_gpu.NVOS02_FLAGS_PHYSICALITY_NONCONTIGUOUS << 4) | (nv_gpu.NVOS02_FLAGS_COHERENCY_CACHED << 12) \
             | (nv_gpu.NVOS02_FLAGS_MAPPING_NO_MAP << 30)
@@ -471,7 +471,7 @@ class PCIIface(PCIIfaceBase):
   def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, **kwargs) -> HCQBuffer:
     # Force use of huge pages for large allocations. NVDev will attempt to use huge pages in any case,
     # but if the size is not aligned, the tail will be allocated with 4KB pages, increasing TLB pressure.
-    page_size = (2 << 20) if size >= (8 << 20) and not uncached and not host else (4 << 10)
+    page_size = mmap.PAGESIZE if uncached or host else ((2 << 20) if size >= (8 << 20) else (4 << 10))
     return super().alloc(round_up(size, page_size), host=host, uncached=uncached, cpu_access=cpu_access, contiguous=contiguous, **kwargs)
 
   def setup_usermode(self): return 0xce000000, self.pci_dev.map_bar(bar=0, fmt='I', off=0xbb0000, size=0x10000)

tinygrad/runtime/ops_qcom.py

@@ -45,6 +45,10 @@ class QCOMSignal(HCQSignal):
       kgsl.IOCTL_KGSL_DEVICE_WAITTIMESTAMP_CTXTID(self.owner.fd, context_id=self.owner.ctx, timestamp=self.owner.last_cmd, timeout=0xffffffff)
 
 class QCOMComputeQueue(HWQueue):
+  def __init__(self, dev:QCOMDevice):
+    self.dev = dev
+    super().__init__()
+
   def __del__(self):
     if self.binded_device is not None: self.binded_device.allocator.free(self.hw_page, self.hw_page.size, BufferSpec(cpu_access=True, nolru=True))
 
@@ -54,7 +58,7 @@ class QCOMComputeQueue(HWQueue):
 
   def _cache_flush(self, write_back=True, invalidate=False, sync=True, memsync=False):
     # TODO: 7xx support.
-    if write_back: self.cmd(adreno.CP_EVENT_WRITE, adreno.CACHE_FLUSH_TS, *data64_le(QCOMDevice.dummy_addr), 0) # dirty cache write-back.
+    if write_back: self.cmd(adreno.CP_EVENT_WRITE, adreno.CACHE_FLUSH_TS, *data64_le(self.dev.dummy_addr), 0) # dirty cache write-back.
     if invalidate: self.cmd(adreno.CP_EVENT_WRITE, adreno.CACHE_INVALIDATE) # invalidate cache lines (following reads from RAM).
     if memsync: self.cmd(adreno.CP_WAIT_MEM_WRITES)
     if sync: self.cmd(adreno.CP_WAIT_FOR_IDLE)
@@ -65,7 +69,7 @@ class QCOMComputeQueue(HWQueue):
 
   def signal(self, signal:QCOMSignal, value=0, ts=False):
     self.cmd(adreno.CP_WAIT_FOR_IDLE)
-    if QCOMDevice.gpu_id < 700:
+    if self.dev.gpu_id[:2] < (7, 3):
       self.cmd(adreno.CP_EVENT_WRITE, qreg.cp_event_write_0(event=adreno.CACHE_FLUSH_TS, timestamp=ts),
                *data64_le(signal.timestamp_addr if ts else signal.value_addr), qreg.cp_event_write_3(value & 0xFFFFFFFF))
       self._cache_flush(write_back=True, invalidate=False, sync=False, memsync=False)
@@ -314,15 +318,12 @@ class QCOMAllocator(HCQAllocatorBase):
     self.dev._gpu_free(opaque)
 
 class QCOMDevice(HCQCompiled):
-  gpu_id: int = 0
-  dummy_addr: int = 0
-
   def __init__(self, device:str=""):
     self.fd = FileIOInterface('/dev/kgsl-3d0', os.O_RDWR)
-    QCOMDevice.dummy_addr = cast(int, self._gpu_alloc(0x1000).va_addr)
+    self.dummy_addr = cast(int, self._gpu_alloc(0x1000).va_addr)
 
     flags = kgsl.KGSL_CONTEXT_PREAMBLE | kgsl.KGSL_CONTEXT_PWR_CONSTRAINT | kgsl.KGSL_CONTEXT_NO_FAULT_TOLERANCE | kgsl.KGSL_CONTEXT_NO_GMEM_ALLOC \
-              | kgsl.KGSL_CONTEXT_PRIORITY(8) | kgsl.KGSL_CONTEXT_PREEMPT_STYLE(kgsl.KGSL_CONTEXT_PREEMPT_STYLE_FINEGRAIN)
+              | kgsl.KGSL_CONTEXT_PRIORITY(getenv("QCOM_PRIORITY", 8)) | kgsl.KGSL_CONTEXT_PREEMPT_STYLE(kgsl.KGSL_CONTEXT_PREEMPT_STYLE_FINEGRAIN)
     self.ctx = kgsl.IOCTL_KGSL_DRAWCTXT_CREATE(self.fd, flags=flags).drawctxt_id
 
     self.cmd_buf = self._gpu_alloc(16 << 20)
@@ -339,11 +340,14 @@ class QCOMDevice(HCQCompiled):
     # Load info about qcom device
     info = kgsl.struct_kgsl_devinfo()
     kgsl.IOCTL_KGSL_DEVICE_GETPROPERTY(self.fd, type=kgsl.KGSL_PROP_DEVICE_INFO, value=ctypes.addressof(info), sizebytes=ctypes.sizeof(info))
-    QCOMDevice.gpu_id = ((info.chip_id >> 24) & 0xFF) * 100 + ((info.chip_id >> 16) & 0xFF) * 10 + ((info.chip_id >>  8) & 0xFF)
-    if QCOMDevice.gpu_id >= 700: raise RuntimeError(f"Unsupported GPU: {QCOMDevice.gpu_id}")
+    self.gpu_id = (info.chip_id >> 24, (info.chip_id >> 16) & 0xFF, (info.chip_id >> 8) & 0xFF)
+
+    # a7xx start with 730x or 'Cxxx', a8xx starts 'Exxx'
+    if self.gpu_id[:2] >= (7, 3): raise RuntimeError(f"Unsupported GPU: chip_id={info.chip_id:#x}")
 
     compilers = [(QCOMRenderer, functools.partial(QCOMCompiler, device))]
-    super().__init__(device, QCOMAllocator(self), compilers, functools.partial(QCOMProgram, self), QCOMSignal, QCOMComputeQueue, None)
+    super().__init__(device, QCOMAllocator(self), compilers, functools.partial(QCOMProgram, self), QCOMSignal,
+                     functools.partial(QCOMComputeQueue, self), None)
 
   def _gpu_alloc(self, size:int, flags:int=0, uncached=False, fill_zeroes=False) -> HCQBuffer:
     flags |= kgsl.KGSL_MEMALIGN(alignment_hint:=12) | kgsl.KGSL_MEMFLAGS_USE_CPU_MAP

tinygrad/runtime/support/am/amdev.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
-import ctypes, collections, dataclasses, functools, os, hashlib, array
+import ctypes, collections, dataclasses, functools, hashlib, array
 from tinygrad.helpers import mv_address, getenv, DEBUG, fetch
 from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support.amd import AMDReg, import_module, import_asic_regs
 from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager
-from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
+from tinygrad.runtime.support.system import PCIDevice, PCIDevImplBase
 from tinygrad.runtime.support.am.ip import AM_SOC, AM_GMC, AM_IH, AM_PSP, AM_SMU, AM_GFX, AM_SDMA
 
 AM_DEBUG = getenv("AM_DEBUG", 0)
@@ -122,8 +122,6 @@ class AMDev(PCIDevImplBase):
     self.pci_dev, self.devfmt, self.dma_regions = pci_dev, pci_dev.pcibus, dma_regions
     self.vram, self.doorbell64, self.mmio = self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')
 
-    self.lock_fd = System.flock_acquire(f"am_{self.devfmt}.lock")
-
     self._run_discovery()
     self._build_regs()
 
@@ -190,7 +188,6 @@ class AMDev(PCIDevImplBase):
     for ip in [self.sdma, self.gfx]: ip.fini_hw()
     self.smu.set_clocks(level=0)
     self.ih.interrupt_handler()
-    os.close(self.lock_fd)
 
   def paddr2mc(self, paddr:int) -> int: return self.gmc.mc_base + paddr
 

tinygrad/runtime/support/compiler_amd.py

@@ -13,8 +13,9 @@ from tinygrad.runtime.support.compiler_cpu import LLVMCompiler
 from tinygrad.helpers import OSX, to_char_p_p
 
 def amdgpu_disassemble(lib:bytes):
-  asm = subprocess.check_output(["llvm-objdump" if OSX else "/opt/rocm/llvm/bin/llvm-objdump", '-d', '-'], input=lib)
-  print('\n'.join([x for x in asm.decode('utf-8').split("\n") if 's_code_end' not in x]))
+  asm = subprocess.check_output(["llvm-objdump" if OSX else "/opt/rocm/llvm/bin/llvm-objdump", '-d', '-'], input=lib).decode("utf-8").splitlines()
+  while asm and ("s_nop 0" in asm[-1] or "s_code_end" in asm[-1]): asm.pop()
+  print("\n".join(asm))
 
 def check(status):
   if status != 0:

tinygrad/runtime/support/hcq.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 from typing import cast, Callable, Type, TypeVar, Generic, Any, Sequence
-import contextlib, decimal, statistics, time, ctypes, array, os, struct, traceback, collections
+import contextlib, decimal, statistics, time, ctypes, array, os, struct, collections, functools
 try: import fcntl # windows misses that
 except ImportError: fcntl = None #type:ignore[assignment]
-from tinygrad.helpers import PROFILE, getenv, to_mv, ProfileRangeEvent
+from tinygrad.helpers import PROFILE, getenv, to_mv, ProfileRangeEvent, select_first_inited
 from tinygrad.device import BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent, CompilerPairT
 from tinygrad.uop.ops import sym_infer, sint, UOp
 from tinygrad.runtime.autogen import libc
@@ -409,6 +409,8 @@ class HCQCompiled(Compiled, Generic[SignalType]):
       for dev in HCQCompiled.peer_groups[pg]: cast(HCQAllocator, dev.allocator).map(alc)
     return self.signal_t(base_buf=HCQCompiled.signal_pool[pg].pop(), owner=self, **kwargs)
 
+  def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.
+
   def _at_profile_finalize(self):
     self.synchronize() # Expect device to be synchronizes
 
@@ -422,7 +424,7 @@ class HCQCompiled(Compiled, Generic[SignalType]):
     gpu2cpu_compute_time_diff = statistics.median([_sync(self, self.hw_compute_queue_t) for _ in range(40)])
     if self.hw_copy_queue_t is None: gpu2cpu_copy_time_diff = decimal.Decimal(0)
     else: gpu2cpu_copy_time_diff = statistics.median([_sync(self, self.hw_copy_queue_t) for _ in range(40)])
-    Compiled.profile_events += [ProfileDeviceEvent(self.device, gpu2cpu_compute_time_diff, gpu2cpu_copy_time_diff)]
+    Compiled.profile_events += [ProfileDeviceEvent(self.device, gpu2cpu_compute_time_diff, gpu2cpu_copy_time_diff, props=self.device_props())]
 
   def _wrap_timeline_signal(self):
     self.timeline_signal, self._shadow_timeline_signal, self.timeline_value = self._shadow_timeline_signal, self.timeline_signal, 1
@@ -435,19 +437,10 @@ class HCQCompiled(Compiled, Generic[SignalType]):
     except MemoryError: buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options), False
     return buf, realloced
 
-  def _make_no_iface_error(self, errs:str, err_short:str) -> RuntimeError:
-    # Keep it in a separate function to avoid creating a traceback <-> locals ref cycle
-    e = RuntimeError(f"No interface for {type(self).__name__[:-6]}:{self.device_id} is available")
-    if hasattr(e, "add_note"): e.add_note(errs + err_short)
-    return e
-
   def _select_iface(self, *ifaces:Type):
-    errs, err_short = "", ""
     if val:=getenv(f'{type(self).__name__[:-6].upper()}_IFACE', ""): ifaces = tuple(x for x in ifaces if x.__name__.startswith(val.upper()))
-    for iface_t in ifaces:
-      try: return iface_t(self, self.device_id)
-      except Exception as e: errs, err_short = errs + f"\n{iface_t.__name__}: {traceback.format_exc()}", err_short + f"\n{iface_t.__name__}: {e}."
-    raise self._make_no_iface_error(errs, err_short)
+    return select_first_inited([functools.partial(cast(Callable, iface), self, self.device_id) for iface in ifaces],
+                               f"No interface for {type(self).__name__[:-6]}:{self.device_id} is available")
 
   def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"
 

tinygrad/runtime/support/nv/ip.py

@@ -89,10 +89,18 @@ class NV_FLCN(NV_IP):
     self.prep_booter()
 
   def prep_ucode(self):
-    expansion_rom_off, bit_addr = {"GA": 0x16600, "AD": 0x14e00}[self.nvdev.chip_name[:2]], 0x1b0
-    vbios_bytes = bytes(array.array('I', self.nvdev.mmio[0x00300000//4:(0x00300000+0x98e00)//4]))
+    vbios_bytes, vbios_off = memoryview(bytes(array.array('I', self.nvdev.mmio[0x00300000//4:(0x00300000+0x100000)//4]))), 0
+    while True:
+      pci_blck = vbios_bytes[vbios_off + nv.OFFSETOF_PCI_EXP_ROM_PCI_DATA_STRUCT_PTR:].cast('H')[0]
+      imglen = vbios_bytes[vbios_off + pci_blck + nv.OFFSETOF_PCI_DATA_STRUCT_IMAGE_LEN:].cast('H')[0] * nv.PCI_ROM_IMAGE_BLOCK_SIZE
+      match vbios_bytes[vbios_off + pci_blck + nv.OFFSETOF_PCI_DATA_STRUCT_CODE_TYPE]:
+        case nv.NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_BASE: block_size = imglen
+        case nv.NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_EXT:
+          expansion_rom_off = vbios_off - block_size
+          break
+      vbios_off += imglen
 
-    bit_header = nv.BIT_HEADER_V1_00.from_buffer_copy(vbios_bytes[bit_addr:bit_addr + ctypes.sizeof(nv.BIT_HEADER_V1_00)])
+    bit_header = nv.BIT_HEADER_V1_00.from_buffer_copy(vbios_bytes[(bit_addr:=0x1b0):bit_addr + ctypes.sizeof(nv.BIT_HEADER_V1_00)])
     assert bit_header.Signature == 0x00544942, f"Invalid BIT header signature {hex(bit_header.Signature)}"
 
     for i in range(bit_header.TokenEntries):

tinygrad/runtime/support/nv/nvdev.py

@@ -73,8 +73,6 @@ class NVDev(PCIDevImplBase):
   def __init__(self, pci_dev:PCIDevice):
     self.pci_dev, self.devfmt, self.mmio = pci_dev, pci_dev.pcibus, pci_dev.map_bar(0, fmt='I')
 
-    self.lock_fd = System.flock_acquire(f"nv_{self.devfmt}.lock")
-
     self.smi_dev, self.is_booting = False, True
     self._early_init()
 

tinygrad/runtime/support/system.py

@@ -165,7 +165,8 @@ class _System:
 System = _System()
 
 class PCIDevice:
-  def __init__(self, pcibus:str, bars:list[int], resize_bars:list[int]|None=None):
+  def __init__(self, devpref:str, pcibus:str, bars:list[int], resize_bars:list[int]|None=None):
+    self.lock_fd = System.flock_acquire(f"{devpref.lower()}_{pcibus.lower()}.lock")
     self.pcibus, self.irq_poller = pcibus, None
 
     if FileIOInterface.exists(f"/sys/bus/pci/devices/{self.pcibus}/driver"):
@@ -215,7 +216,8 @@ class PCIDevice:
   def reset(self): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{self.pcibus}/reset'")
 
 class APLPCIDevice(PCIDevice):
-  def __init__(self, pcibus:str, bars:list[int], resize_bars:list[int]|None=None):
+  def __init__(self, devpref:str, pcibus:str, bars:list[int], resize_bars:list[int]|None=None):
+    self.lock_fd = System.flock_acquire(f"{devpref.lower()}_{pcibus.lower()}.lock")
     self.pcibus, self.bars = pcibus, {b: System.iokit_pci_memmap(b) for b in bars}
     self.bar_info = {b:PCIBarInfo(0, self.bars[b].nbytes-1 if b in self.bars else 0) for b in range(6)} # NOTE: fake bar info for nv.
   def map_bar(self, bar:int, off:int=0, addr:int=0, size:int|None=None, fmt='B') -> MMIOInterface: return self.bars[bar].view(off, size, fmt)
@@ -224,7 +226,8 @@ class APLPCIDevice(PCIDevice):
   def reset(self): System.iokit_pci_rpc(__TinyGPURPCReset:=2)
 
 class USBPCIDevice(PCIDevice):
-  def __init__(self, pcibus:str, bars:list[int], resize_bars:list[int]|None=None):
+  def __init__(self, devpref:str, pcibus:str, bars:list[int], resize_bars:list[int]|None=None):
+    self.lock_fd = System.flock_acquire(f"{devpref.lower()}_{pcibus.lower()}.lock")
     self.usb = ASM24Controller()
     self.pcibus, self.bar_info = pcibus, System.pci_setup_usb_bars(self.usb, gpu_bus=4, mem_base=0x10000000, pref_mem_base=(32 << 30))
   def map_bar(self, bar, off=0, addr=0, size=None, fmt='B'):
@@ -247,12 +250,12 @@ class LNXPCIIfaceBase:
 
       # Acquire va range to avoid collisions.
       FileIOInterface.anon_mmap(va_start, va_size, 0, mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS | MAP_NORESERVE | MAP_FIXED, 0)
-    self.pci_dev, self.dev, self.vram_bar = PCIDevice(cls.gpus[dev_id], bars=bars, resize_bars=[vram_bar]), dev, vram_bar
+    self.pci_dev, self.dev, self.vram_bar = PCIDevice(dev.__class__.__name__[:2], cls.gpus[dev_id], bars=bars, resize_bars=[vram_bar]), dev, vram_bar
     self.p2p_base_addr = self.pci_dev.bar_info[vram_bar].addr
 
   def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
     # NOTE: logic on macos is different, since bar is small
-    should_use_sysmem = host or (((uncached or cpu_access) if OSX else (uncached and cpu_access)) and not force_devmem)
+    should_use_sysmem = host or ((cpu_access if OSX else (uncached and cpu_access)) and not force_devmem)
     if should_use_sysmem:
       vaddr = self.dev_impl.mm.alloc_vaddr(size:=round_up(size, mmap.PAGESIZE), align=mmap.PAGESIZE)
       memview, paddrs = System.alloc_sysmem(size, vaddr=vaddr, contiguous=contiguous)
@@ -281,7 +284,8 @@ class LNXPCIIfaceBase:
 
 class APLPCIIfaceBase(LNXPCIIfaceBase):
   def __init__(self, dev, dev_id, vendor, devices, bars, vram_bar, va_start, va_size):
-    self.pci_dev, self.dev, self.vram_bar = APLPCIDevice(pcibus=f'usb4:{dev_id}', bars=bars), dev, vram_bar
+    self.pci_dev, self.dev, self.vram_bar = APLPCIDevice(dev.__class__.__name__[:2], pcibus=f'usb4:{dev_id}', bars=bars), dev, vram_bar
+    assert (read_vendor:=self.pci_dev.read_config(0x00, 2)) == vendor, f"Vendor ID mismatch: expected {vendor:#x}, got {read_vendor:#x}"
   def map(self, b:HCQBuffer): raise RuntimeError(f"map failed: {b.owner} -> {self.dev}")
 
 PCIIfaceBase:type = APLPCIIfaceBase if OSX else LNXPCIIfaceBase

tinygrad/schedule/indexing.py

@@ -16,6 +16,9 @@ def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None:
   for s in rb.src:
     if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None
 
+def realize_store(ctx:dict[UOp, None], a:UOp) -> None:
+  if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
+
 def realize_assign(ctx:dict[UOp, None], a:UOp) -> None:
   if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
   # if it's a kernel, we don't realize it
@@ -30,6 +33,8 @@ pm_generate_realize_map = PatternMatcher([
   (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
   # realize ASSIGN and input to assign (might be optimized out)
   (UPat(Ops.ASSIGN, name="a"), realize_assign),
+  # realize STORE
+  (UPat(Ops.STORE, name="a"), realize_store),
 ])
 
 @dataclass(frozen=True)
@@ -50,13 +55,14 @@ class IndexingContext:
     # if a range has a 1 src, it's the same as UOp.const(dtypes.index, 0)
     return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(dtypes.index, 0)
 
+ops_allowed_after = (Ops.KERNEL, Ops.RANGE)
 def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
   if x.op in {Ops.BUFFERIZE, Ops.INDEX}: return None
-  if x.op is Ops.AFTER and x.src[1].op is Ops.KERNEL: return None
+  if x.op is Ops.AFTER and x.src[1].op in ops_allowed_after: return None
   new_srcs = []
   for s in x.src:
     new_src = s
-    if s.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT} or (s.op is Ops.AFTER and s.src[1].op is Ops.KERNEL):
+    if s.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT} or (s.op is Ops.AFTER and s.src[1].op in ops_allowed_after):
       if x in ctx.range_map: new_src = new_src.index(*ctx.range_map[x][0])
     elif s in ctx.realize_map:
       realized_ranges = ctx.realize_map[s]
@@ -176,7 +182,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
       # mark all ranges as ended
       assert rctx.realize_map[x] is None
       rctx.realize_map[x] = list(range(len(x.shape)))
-    elif x.op in {Ops.MSTACK, Ops.MSELECT}:
+    elif x.op in {Ops.MSTACK, Ops.MSELECT, Ops.END}:
       # treat MSTACK/MSELECT like SINK
       continue
     elif len(consumer_rngs) == 0:
@@ -239,7 +245,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     # if the EXPAND is used to inject a range, we don't mark it as ending_ranges. otherwise we do.
     # NOTE: this doesn't actually always end a range, but this is why convs are realized, so for now we need it
     if x.op is Ops.EXPAND and all(isinstance(y, int) or y.op is not Ops.RANGE for y in x.shape):
-      ending_ranges[x] = list(UOp.sink(*[ro for ri, ro in zip(rngs, out_rngs) if ri is not ro]).ranges.keys())
+      ending_ranges[x] += list(UOp.sink(*[ro for ri, ro in zip(rngs, out_rngs) if ri is not ro]).ranges.keys())
 
     # REDUCE_AXIS creates ranges for the axes it is reducing
     if x.op is Ops.REDUCE_AXIS:
@@ -247,15 +253,23 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
 
     if debug:
       realized_ranges = rctx.realize_map.get(x, None)
-      disp = []
-      for i, (ri, ro) in enumerate(zip([r.render() for r in rngs], [r.render() for r in out_rngs])):
-        rng = f"{ri}" if ri == ro else f"{ri} -> {ro}"
-        if realized_ranges is not None and i in realized_ranges: rng = colored(rng, "yellow")
-        disp.append("["+rng+"]")
-      print("***" if x in rctx.realize_map else "   ", len(consumer_map[x]), f"{str(x.op):20s}", ''.join(disp))
+      if x.op is Ops.RESHAPE or len(rngs) != len(out_rngs):
+        disp = render_ranges(rngs, realized=realized_ranges) + " -> " + render_ranges(out_rngs, realized=realized_ranges)
+      else:
+        disp = render_ranges(rngs, out_rngs, realized=realized_ranges)
+      print("***" if x in rctx.realize_map else "   ",
+            f"{len(consumer_map[x]):2d} {str(x.op):20s} {str(x.shape):35s} {len(ending_ranges[x]):2d}", disp)
 
     # assign to the range map. rngs are the input ranges, out_rngs are the output ranges, from the x op.
     rctx.range_map[x] = (rngs, out_rngs)
 
   tsink = graph_rewrite(tsink, pm_apply_rangeify, ctx=rctx, bottom_up=True, name="apply rangeify")
   return tsink, rctx
+
+def render_ranges(*rngs_list, realized) -> str:
+  disp = []
+  for i, rs in enumerate(zip(*[[r.render() for r in rngs] for rngs in rngs_list])):
+    rng = rs[0] if all_same(rs) else " -> ".join(rs)
+    if realized is not None and i in realized: rng = colored(rng, "yellow")
+    disp.append("["+rng+"]")
+  return ''.join(disp)

tinygrad/schedule/multi.py

@@ -68,7 +68,7 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
   # allgather
   copied_chunks = []
   for i,c in enumerate(reduced_chunks):
-    this_chunk = [None] * len(buf.device)
+    this_chunk: list[UOp|None] = [None] * len(buf.device)
     this_chunk[(i+len(buf.device)-1)%n_lbs] = c
     for step in range(n_lbs-1):
       dest = (i+step)%n_lbs
@@ -186,7 +186,7 @@ def shrink_multi(root:UOp, multi:UOp):
 
 def flip_multi(root:UOp, multi:UOp):
   assert multi.axis is None or not root.marg[multi.axis], "flipping not supported on sharded axis"
-  return multi.src[0].flip(root.marg).multi(multi.axis)
+  return multi.src[0].flip([i for i,x in enumerate(root.marg) if x]).multi(multi.axis)
 
 # from multiple devices -> one
 def copy_multi(multi:UOp, device:UOp):
@@ -214,7 +214,7 @@ multi_pm = PatternMatcher([
   (UPat(Ops.COPY, src=(UPat(Ops.MULTI, name="multi"), UPat(Ops.DEVICE, name="device"))), copy_multi),
   (UPat(Ops.ALLREDUCE, src=(UPat(Ops.MULTI, name="multi"), UPat(Ops.DEVICE, name="device")), name="red"),
     lambda multi,device,red: multi.src[0].allreduce(red.arg, device).multi(axis=multi.axis)),
-  (UPat((Ops.CAST, Ops.BITCAST, Ops.CONTIGUOUS, Ops.DETACH, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE),
+  (UPat((Ops.CAST, Ops.BITCAST, Ops.CONTIGUOUS, Ops.DETACH, Ops.CONTIGUOUS_BACKWARD),
         src=(UPat(Ops.MULTI, name="multi"), ), name="root"), passthrough_multi),
 ])+replace_allreduce
 

tinygrad/schedule/rangeify.py

@@ -3,7 +3,7 @@ import itertools
 from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, ssimplify, KernelInfo
 from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType, BottomUpGate, Kernel, _remove_all_tags
-from tinygrad.uop.symbolic import symbolic_flat
+from tinygrad.uop.symbolic import symbolic
 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
@@ -17,7 +17,7 @@ sys.setrecursionlimit(10000)
 # movement op on INDEX as a PatternMatcher
 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
+   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)),
   # 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)),
@@ -65,7 +65,7 @@ mop_cleanup = PatternMatcher([
 
 earliest_rewrites = mop_cleanup+PatternMatcher([
   # just removing it works...
-  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE), name="x"), lambda x: x.src[0]),
+  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
 
   # remove CONTIGUOUS if the BUFFER is already contiguous
   (UPat(Ops.BUFFER).f(Ops.RESHAPE, allow_any_len=True, name="r").f(Ops.CONTIGUOUS, name="c"), lambda r,c: r.replace(tag=c.tag)),
@@ -396,6 +396,7 @@ def handle_after(ctx:LocalAddBufferContext, after:UOp):
   return buf
 
 def renumber_range(ctx:LocalAddBufferContext, r:UOp):
+  if r.arg[-1] == AxisType.OUTER: return None
   if r.tag != (): return None
   ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=None)
   ctx.range += 1
@@ -469,7 +470,10 @@ pm_add_range_tags = PatternMatcher([
 ])
 
 def split_store(ctx:list[UOp], x:UOp) -> UOp|None:
-  if len(x.ranges): return None
+  if len([r for r in x.ranges if r.arg[-1] != AxisType.OUTER]): return None
+
+  # ends of outer range don't go in kernels
+  if x.op is Ops.END and x.src[1].op is Ops.RANGE and x.src[1].arg[-1] == AxisType.OUTER: return None
 
   # local kernel rewrite
   lctx = LocalAddBufferContext()
@@ -536,8 +540,9 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   # convert movement ops to ranges
   tsink, rctx = run_rangeify(tsink, DEBUG_RANGEIFY)
 
-  tsink = graph_rewrite(tsink, symbolic_flat+pm_reduce_simplify+pm_const_buffer_folding, name="symbolic+reduce_collapse")  # this does const folding
+  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding, name="symbolic+reduce_collapse")  # this does const folding
   tsink = graph_rewrite(tsink, pm_remove_bufferize, bottom_up=True, name="remove bufferize with cost function")
+  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding, name="symbolic+reduce_collapse pt 2")
   tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
 
   # rebuild the sink with all the BUFFERIZEs with tags, this is what's ending up in the tensor graph

tinygrad/tensor.py

@@ -5,11 +5,12 @@ from contextlib import ContextDecorator
 from typing import Callable, ClassVar, Sequence, cast, get_args, Literal, SupportsIndex, ParamSpec, TypeVar, Generic
 from tinygrad.dtype import DType, DTypeLike, dtypes, ImageDType, ConstType, least_upper_float, least_upper_dtype, sum_acc_dtype, to_dtype, truncate
 from tinygrad.dtype import _from_np_dtype, _to_np_dtype
-from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten, dedup
-from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, DEBUG, is_numpy_ndarray, FUSE_ATTENTION, SPEC
+from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten
+from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, DEBUG, is_numpy_ndarray, SPEC
 from tinygrad.helpers import suppress_finalizing
 from tinygrad.gradient import compute_gradient
-from tinygrad.uop.mathtraits import MathTrait
+from tinygrad.mixin import OpMixin
+from tinygrad.mixin.movement import _align_left
 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
@@ -79,10 +80,6 @@ def _apply_winograd_matrix(mat, t:Tensor, dims:int) -> Tensor:
   assert isinstance(ret, Tensor), "sum didn't return a Tensor"
   return ret
 
-def _align_left(*shapes:tuple[sint, ...]) -> tuple[tuple[sint, ...], ...]:
-  # unsqueeze left to make every shape same length
-  max_dim = max(len(shape) for shape in shapes)
-  return tuple((1,) * (max_dim - len(shape)) + shape for shape in shapes)
 def _broadcast_shape(*shapes:tuple[sint, ...]) -> tuple[sint, ...]:
   return tuple(0 if 0 in nth_dim_sizes else smax(nth_dim_sizes) for nth_dim_sizes in zip(*_align_left(*shapes)))
 
@@ -100,7 +97,7 @@ def _flat_to_grouped(padding:Sequence[sint]) -> tuple[tuple[sint, sint], ...]: r
 
 ReductionStr = Literal["mean", "sum", "none"]
 
-class Tensor(MathTrait):
+class Tensor(OpMixin):
   """
   A `Tensor` is a multi-dimensional matrix containing elements of a single data type.
 
@@ -1040,80 +1037,6 @@ class Tensor(MathTrait):
 
   def _mop(self, op:Ops, arg) -> Tensor: return self._apply_uop(UOp._mop, extra_args=(op,), arg=arg)
 
-  def expand(self, shape, *args) -> Tensor:
-    """
-    Returns a tensor that is expanded to the shape that is specified.
-    Expand can also increase the number of dimensions that a tensor has.
-
-    Passing a `-1` or `None` to a dimension means that its size will not be changed.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor([1, 2, 3])
-    print(t.expand(4, -1).numpy())
-    ```
-    """
-    new_shape = tuple(from_ if to == -1 or to is None else to for from_, to in zip(*(_align_left(self.shape, argfix(shape, *args)))))
-    return self._broadcast_to(new_shape)
-
-  def permute(self, order, *args) -> Tensor:
-    """
-    Returns a tensor that is a permutation of the original tensor.
-    The new tensor has the same data as the original tensor but with the dimensions permuted according to the order specified.
-    `order` can be passed as a tuple or as separate arguments.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor.empty(2, 3, 5)
-    print(t.shape)
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.permute(2, 0, 1).shape)
-    ```
-    """
-    order_arg = tuple(self._resolve_dim(x) for x in argfix(order, *args))
-    if sorted(order_arg) != list(range(self.ndim)): raise RuntimeError(f"order is not a valid permutation, getting {order_arg}")
-    return self._apply_uop(UOp.permute, arg=order_arg)
-
-  def flip(self, axis, *args) -> Tensor:
-    """
-    Returns a tensor that reverses the order of the original tensor along given `axis`.
-    `axis` can be passed as a tuple or as separate arguments.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor.arange(6).reshape(2, 3)
-    print(t.numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.flip(0).numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.flip((0, 1)).numpy())
-    ```
-    """
-    axis_arg = tuple(self._resolve_dim(x) for x in argfix(axis, *args))
-    if len(axis_arg) != len(dedup(axis_arg)): raise RuntimeError(f"dim can appear at most once, getting {axis_arg}")
-    return self._apply_uop(UOp.flip, arg=tuple([i in axis_arg for i in range(len(self.shape))]))
-
-  def shrink(self, arg:tuple[tuple[sint, sint]|None, ...]) -> Tensor:
-    """
-    Returns a tensor that shrinks the each axis based on input arg.
-    `arg` must have the same length as `self.ndim`.
-    For each axis, it can be `None`, which means no shrink, or a tuple `(start, end)` that works the same as Python slice.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor.arange(9).reshape(3, 3)
-    print(t.numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.shrink(((None, (1, 3)))).numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.shrink((((0, 2), (0, 2)))).numpy())
-    ```
-    """
-    if self.ndim != len(arg): raise ValueError(f"{self.ndim=} != {len(arg)=}")
-    if (shrink_arg:=[x if x is not None else (0,s) for x,s in zip(arg, self.shape)]) == [(0,s) for s in self.shape]: return self
-    return self._apply_uop(UOp.shrink, arg=tuple(shrink_arg))
-
   def pad(self, padding:Sequence[sint]|Sequence[tuple[sint, sint]|None], mode:str="constant", value:float=0.0) -> Tensor:
     """
     Returns a tensor with padding applied based on the input `padding`.
@@ -1181,8 +1104,6 @@ class Tensor(MathTrait):
   def pad_to(self, shape, *args):
     if len(new_shape := argfix(shape, *args)) != self.ndim: raise ValueError(f"dim mismatch, cannot pad {self.shape} to {new_shape}")
     return self.pad(tuple([None if ns is None else (0, ns-s) for s,ns in zip(self.shape, new_shape)]))
-  def shrink_to(self, shape, *args):
-    return self.shrink(tuple([None if ns is None else (0, ns) for ns in argfix(shape, *args)]))
 
   # ***** movement high level ops *****
 
@@ -1402,44 +1323,6 @@ class Tensor(MathTrait):
     # checks for shapes and number of dimensions delegated to cat
     return Tensor.cat(*[t.unsqueeze(dim) for t in argfix(self, *args)], dim=dim)
 
-  def repeat_interleave(self, repeats:int, dim:int|None=None) -> Tensor:
-    """
-    Repeats elements of a tensor.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor([1, 2, 3])
-    print(t.repeat_interleave(2).numpy())
-    ```
-    """
-    x, dim = (self.flatten(), 0) if dim is None else (self, self._resolve_dim(dim))
-    shp = x.shape
-    return x.reshape(*shp[:dim+1], 1, *shp[dim+1:]).expand(*shp[:dim+1], repeats, *shp[dim+1:]).reshape(*shp[:dim], shp[dim]*repeats, *shp[dim+1:])
-
-  def repeat(self, repeats, *args) -> Tensor:
-    """
-    Repeats tensor number of times along each dimension specified by `repeats`.
-    `repeats` can be passed as a tuple or as separate arguments.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor([1, 2, 3])
-    print(t.repeat(4, 2).numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.repeat(4, 2, 1).shape)
-    ```
-    """
-    repeats = argfix(repeats, *args)
-    base_shape = _align_left(self.shape, repeats)[0]
-    unsqueezed_shape = flatten([[1, s] for s in base_shape])
-    expanded_shape = flatten([[r, s] for r,s in zip(repeats, base_shape)])
-    final_shape = [r*s for r,s in zip(repeats, base_shape)]
-    return self.reshape(unsqueezed_shape).expand(expanded_shape).reshape(final_shape)
-
-  def _resolve_dim(self, dim:int, *, extra:bool=False) -> int:
-    total = self.ndim + int(extra)
-    if not -max(1, total) <= dim <= max(1, total)-1: raise IndexError(f"{dim=} out of range {[-max(1, total), max(1, total)-1]}")
-    return dim + total if dim < 0 else dim
-
   def split(self, sizes:int|Sequence[int], dim:int=0) -> tuple[Tensor, ...]:
     """
     Splits the tensor into chunks along the dimension specified by `dim`.
@@ -1540,96 +1423,6 @@ class Tensor(MathTrait):
     output_shape = _broadcast_shape(*(t.shape for t in tensors))
     return tuple(t._broadcast_to(output_shape) for t in tensors)
 
-  def squeeze(self, dim:int|None=None) -> Tensor:
-    """
-    Returns a tensor with specified dimensions of input of size 1 removed.
-    If `dim` is not specified, all dimensions with size 1 are removed.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor.zeros(2, 1, 2, 1, 2)
-    print(t.squeeze().shape)
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.squeeze(0).shape)
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.squeeze(1).shape)
-    ```
-    """
-    if dim is None: return self.reshape(tuple(dim for dim in self.shape if dim != 1))
-    dim = self._resolve_dim(dim)
-    return self if not self.ndim or self.shape[dim] != 1 else self.reshape(self.shape[:dim] + self.shape[dim+1:])
-
-  def unsqueeze(self, dim:int) -> Tensor:
-    """
-    Returns a tensor with a new dimension of size 1 inserted at the specified `dim`.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor([1, 2, 3, 4])
-    print(t.unsqueeze(0).numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.unsqueeze(1).numpy())
-    ```
-    """
-    dim = self._resolve_dim(dim, extra=True)
-    return self.reshape(self.shape[:dim] + (1,) + self.shape[dim:])
-
-  @property
-  def T(self) -> Tensor:
-    """`.T` is an alias for `.transpose()`."""
-    return self.transpose()
-
-  def transpose(self, dim0=1, dim1=0) -> Tensor:
-    """
-    Returns a tensor that is a transposed version of the original tensor.
-    The given dimensions `dim0` and `dim1` are swapped.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor.arange(6).reshape(2, 3)
-    print(t.numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.transpose(0, 1).numpy())
-    ```
-    """
-    order = list(range(self.ndim))
-    order[dim0], order[dim1] = order[dim1], order[dim0]
-    return self.permute(order)
-
-  def flatten(self, start_dim=0, end_dim=-1) -> Tensor:
-    """
-    Flattens the tensor by reshaping it into a one-dimensional tensor.
-    If `start_dim` or `end_dim` are passed, only dimensions starting with `start_dim` and ending with `end_dim` are flattened.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor.arange(8).reshape(2, 2, 2)
-    print(t.flatten().numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.flatten(start_dim=1).numpy())
-    ```
-    """
-    start_dim, end_dim = self._resolve_dim(start_dim), self._resolve_dim(end_dim)
-    return self.reshape(self.shape[:start_dim] + (prod(self.shape[start_dim:end_dim+1]), ) + self.shape[end_dim+1:])
-
-  def unflatten(self, dim:int, sizes:tuple[int,...]) -> Tensor:
-    """
-    Unflattens dimension `dim` of the tensor into multiple dimensions specified by `sizes`. `Tensor.flatten()` is the inverse of this function.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor.ones(3, 4, 1).unflatten(1, (2, 2)).shape)
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor.ones(3, 4, 1).unflatten(1, (-1, 2)).shape)
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor.ones(5, 12, 3).unflatten(-2, (2, 2, 3, 1, 1)).shape)
-    ```
-    """
-    dim = self._resolve_dim(dim)
-    return self.reshape(self.shape[:dim] + sizes + self.shape[dim+1:])
-
   def diag(self) -> Tensor:
     """
     Returns a 2-D square tensor with the elements of input as the main diagonal.
@@ -1675,46 +1468,6 @@ class Tensor(MathTrait):
     for dim, shift in zip(dims, shifts): slices[dim] = slice(delta:=self.shape[dim]-shift%self.shape[dim], delta+self.shape[dim])
     return self.repeat(*tuple(2 if i in dims else 1 for i in range(self.ndim)))[slices]
 
-  def rearrange(self, formula:str, **sizes) -> Tensor:
-    """
-    Rearranges input according to formula
-
-    See: https://einops.rocks/api/rearrange/
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    x = Tensor([[1, 2], [3, 4]])
-    print(Tensor.rearrange(x, "batch channel -> (batch channel)").numpy())
-    ```
-    """
-    def parse_formula(formula: str):
-      tokens = f" {formula} ".replace("…", "...").replace("(", " ( ").replace(")", " ) ").replace(" ", "  ").replace(" 1 ", " ( ) ").split()
-      lparens, rparens = map(lambda x: [i for i, ch in enumerate(tokens) if ch == x], ("(", ")"))
-      pairs = list(zip(lparens, rparens))
-      assert len(lparens) == len(rparens) and sorted(flatten(pairs)) == flatten(pairs), "bracket mismatch"
-      return [name for name in tokens if name not in ("(", ")")], [(s - 2*i, e - 1 - 2*i) for i, (s, e) in enumerate(pairs)]
-
-    assert formula.count("->") == 1, 'need exactly one "->" in formula'
-
-    (lhs, unflatten_dims), (rhs, flatten_dims) = map(parse_formula, formula.split("->"))
-
-    for name in sizes: assert name in lhs, f"axis {name} is not used in transform"
-    assert sorted(lhs) == sorted(rhs) and len(lhs) == len(set(lhs)), f"name mismatch in {formula}"
-    for name in flatten((lhs, rhs)): assert name == "..." or (name.isidentifier() and "_" not in (name[0], name[-1])), f"invalid axis name {name}"
-    assert "..." not in flatten([lhs[s:e] for s, e in unflatten_dims]), f"cannot have collapsed ellipsis (...) in lhs of {formula}"
-    assert lhs.count("...") <= 1, f"too many ellipses in {formula}"
-
-    # resolve ellipsis
-    if "..." in lhs: ell_len = len(self.shape) - len(lhs) + 1 + sum(e - s - 1 for s, e in unflatten_dims)
-    lhs, rhs = map(lambda l: l[:(i:=l.index("..."))] + [f"...{j}" for j in range(ell_len)] + l[i + 1:] if "..." in l else l, (lhs, rhs))
-    unflatten_dims = [(s + (ell_len - 1 if "...0" in lhs[:s] else 0), e + (ell_len - 1 if "...0" in lhs[:e] else 0)) for s, e in unflatten_dims]
-    flatten_dims = [(s + (ell_len - 1 if "...0" in rhs[:s] else 0), e + (ell_len - 1 if "...0" in rhs[:e] else 0)) for s, e in flatten_dims]
-
-    # apply movement ops in order unflatten -> permute -> flatten/unsqueeze
-    t = functools.reduce(lambda x, dims: x.unflatten(dims[0], tuple(sizes.get(lhs[d], -1) for d in range(*dims))), unflatten_dims, self)
-    for i, name in enumerate(lhs): assert (name not in sizes) or sizes[name] == t.shape[i], f"size provided for dimension {name} incorrect"
-    t = t.permute([lhs.index(name) for name in rhs])
-    return functools.reduce(lambda x, dims: x.flatten(dims[0], dims[1] - 1) if dims[0]<dims[1] else x.unsqueeze(dims[0]), reversed(flatten_dims), t)
-
   def masked_select(self, mask):
     """
     Selects elements from `self` based on the boolean `mask`.
@@ -1927,6 +1680,12 @@ class Tensor(MathTrait):
     is_nan_close = (self.isnan() & other.isnan()) & equal_nan
     return is_finite_close | is_infinite_close | is_nan_close
 
+  def allclose(self, other:Tensor, rtol:float=1e-05, atol:float=1e-08, equal_nan=False) -> bool:
+    """
+    Check if all self and other are close. Return True or False.
+    """
+    return bool(self.isclose(other, rtol=rtol, atol=atol, equal_nan=equal_nan).all().item())
+
   def mean(self, axis:int|Sequence[int]|None=None, keepdim=False) -> Tensor:
     """
     Returns the mean value of the tensor along the specified axis or axes.
@@ -2131,7 +1890,7 @@ class Tensor(MathTrait):
     e = m.exp()
     return m, e, e.sum(axis=axis, keepdim=True)
 
-  def softmax(self, axis=-1, dtype:DTypeLike|None=None, _single_kernel=getenv("SINGLE_KERNEL_SOFTMAX")) -> Tensor:
+  def softmax(self, axis=-1, dtype:DTypeLike|None=None) -> Tensor:
     """
     Applies the softmax function to the tensor along the specified axis.
 
@@ -2151,9 +1910,6 @@ class Tensor(MathTrait):
     print(t.softmax(axis=0).numpy())
     ```
     """
-    if _single_kernel:
-      _, e, ss = self.contiguous()._softmax(axis, dtype)
-      return e.div(ss).fuse()
     _, e, ss = self._softmax(axis, dtype)
     return e.div(ss)
 
@@ -2344,23 +2100,17 @@ class Tensor(MathTrait):
     noop, i_ = [None] * (self.ndim-len(k_)), self.shape[-len(k_):]
     assert all(resolve(d*(k-1)+1 <= i) for k,d,i in zip(k_,d_,i_)), "kernel size cannot be greater than actual input size"
     o_ = [ceildiv(i-d*(k-1), s) for i,d,k,s in zip(i_,d_,k_,s_)]
-    if any(resolve(k > s) for k,s in zip(k_,s_)) or any(d != 1 for d in d_):
-      # input size scaling factor to make sure shrink for stride is possible
-      f_ = [1 + int(resolve(o*s > (i - d*(k-1)))) for o,s,i,d,k in zip(o_,s_,i_,d_,k_)]
-      # # repeats such that we don't need padding
-      x = self.repeat([1]*len(noop) + [ceildiv(k*(i*f+d),i) for k,i,d,f in zip(k_,i_,d_,f_)])
-      # handle dilation
-      x = x.shrink(tuple(noop + [(0,k*(i*f+d)) for k,i,d,f in zip(k_,i_,d_,f_)])).reshape(noop + flatten((k,(i*f+d)) for k,i,d,f in zip(k_,i_,d_,f_)))
-      # handle stride
-      x = x.shrink(tuple(noop + flatten(((0,k), (0,o*s)) for k,o,s in zip(k_,o_,s_)))).reshape(noop + flatten((k,o,s) for k,o,s in zip(k_,o_,s_)))
-      x = x.shrink(tuple(noop + flatten(((0,k), (0,o), (0,1)) for k,o in zip(k_,o_)))).reshape(noop + flatten((k,o) for k,o in zip(k_,o_)))
-      # permute to move reduce to the end
-      return x.permute(*range(len(noop)), *[len(noop)+i*2+1 for i in range(len(i_))], *[len(noop)+i*2 for i in range(len(i_))])
-    # TODO: once the shapetracker can optimize well, remove this alternative implementation
-    x = self.pad(tuple(noop + [(0, max(0,o*s-i)) for i,o,s in zip(i_,o_,s_)])).shrink(tuple(noop + [(0,o*s) for o,s in zip(o_,s_)]))
-    x = x.reshape(noop + flatten(((o,s) for o,s in zip(o_,s_))))
-    x = x.shrink(tuple(noop + flatten(((0,o), (0,k)) for o,k in zip(o_,k_))))
-    return x.permute(*range(len(noop)), *[len(noop)+i*2 for i in range(len(i_))], *[len(noop)+i*2+1 for i in range(len(i_))])
+    # input size scaling factor to make sure shrink for stride is possible
+    f_ = [smax(1, ceildiv(o*s - d, i)) for o,s,i,d in zip(o_,s_,i_,d_)]
+    # repeats such that we don't need padding
+    x = self.repeat([1]*len(noop) + [ceildiv(k*(i*f+d),i) for k,i,d,f in zip(k_,i_,d_,f_)])
+    # handle dilation
+    x = x.shrink_to(noop + [k*(i*f+d) for k,i,d,f in zip(k_,i_,d_,f_)]).reshape(noop + flatten((k,(i*f+d)) for k,i,d,f in zip(k_,i_,d_,f_)))
+    # handle stride
+    x = x.shrink_to(noop + flatten((k,o*s) for k,o,s in zip(k_,o_,s_))).reshape(noop + flatten((k,o,s) for k,o,s in zip(k_,o_,s_)))
+    x = x.shrink_to(noop + flatten((k,o,1) for k,o in zip(k_,o_))).reshape(noop + flatten((k,o) for k,o in zip(k_,o_)))
+    # permute to move reduce to the end
+    return x.permute(*range(len(noop)), *[len(noop)+i*2+1 for i in range(len(i_))], *[len(noop)+i*2 for i in range(len(i_))])
 
   def _resolve_pool_pads(self, padding:int|Sequence[int], dims:int) -> Sequence[int]:
     if not isinstance(padding, int) and not (len(padding) == 2*dims or len(padding) == dims):
@@ -3014,15 +2764,6 @@ class Tensor(MathTrait):
     """
     return self._apply_uop(UOp.contiguous, extra_args=args, **kwargs)
 
-  def fuse(self) -> Tensor:
-    """
-    Makes this a single kernel back to Ops.CONTIGUOUS on the inputs.
-
-    Useful for single kernel softmax and flash attention.
-    Careful, this can break codegen or make kernels really slow.
-    """
-    return self._apply_uop(UOp.fuse)
-
   def contiguous_backward(self) -> Tensor:
     """
     Inserts a contiguous operation in the backward pass.
@@ -3607,18 +3348,8 @@ class Tensor(MathTrait):
     return self / (1 + self.abs())
 
   # ***** broadcasted elementwise ops *****
-  def _broadcast_to(self, new_shape:tuple[sint, ...]) -> Tensor:
-    if self.shape == new_shape: return self
-    if self.ndim > len(new_shape): raise ValueError(f"cannot broadcast tensor to fewer dimensions. shape={self.shape} to {new_shape=}")
-    # first unsqueeze left with 1s https://data-apis.org/array-api/latest/API_specification/broadcasting.html
-    shape, _ = _align_left(self.shape, new_shape)
-    # for each dimension, check either dim is 1, or it does not change
-    if not all(resolve(s == ns) or resolve(s == 1) for s,ns in zip(shape, new_shape)):
-      raise ValueError(f"cannot broadcast {self.shape} to {new_shape=}")
-    # NOTE: this cast is no-op in forward and uses sum_acc_dtype in the backward sum
-    return self.reshape(shape).cast(sum_acc_dtype(self.dtype))._apply_uop(UOp.expand, arg=new_shape).cast(self.dtype)
 
-  def _broadcasted(self, y:Tensor|ConstType|UOp, reverse:bool=False, match_dtype:bool=True) -> tuple[Tensor, Tensor]:
+  def _broadcasted(self, y:Tensor|ConstType|UOp, reverse:bool=False, match_dtype:bool=True, backward_cast:bool=True) -> tuple[Tensor, Tensor]:
     x: Tensor = self
     if not isinstance(y, Tensor):
       # make y a Tensor
@@ -3634,8 +3365,13 @@ class Tensor(MathTrait):
 
     if reverse: x, y = y, x
 
+    # compute the output shape
+    out_shape = _broadcast_shape(x.shape, y.shape)
+
     # broadcast
-    return x._broadcast_to(out_shape:=_broadcast_shape(x.shape, y.shape)), y._broadcast_to(out_shape)
+    # NOTE: the backward cast is no-op in forward and uses sum_acc_dtype in the backward sum
+    return x.cast(sum_acc_dtype(x.dtype) if backward_cast else x.dtype)._broadcast_to(out_shape).cast(x.dtype), \
+           y.cast(sum_acc_dtype(y.dtype) if backward_cast else y.dtype)._broadcast_to(out_shape).cast(y.dtype)
 
   def sub(self, x:Tensor|ConstType, reverse=False) -> Tensor:
     """
@@ -4002,9 +3738,7 @@ class Tensor(MathTrait):
       key = key.repeat_interleave(self.shape[-3] // key.shape[-3], dim=-3)
       value = value.repeat_interleave(self.shape[-3] // value.shape[-3], dim=-3)
 
-    if FUSE_ATTENTION: q, key, value = self.contiguous(), key.contiguous(), value.contiguous()
-    else: q = self
-
+    q = self
     qk = q.matmul(key.transpose(-2,-1), dtype=least_upper_dtype(q.dtype, key.dtype, dtypes.float32)) / math.sqrt(q.shape[-1])
     # handle attention mask
     if is_causal:
@@ -4013,8 +3747,7 @@ class Tensor(MathTrait):
     if attn_mask is not None:
       if attn_mask.dtype == dtypes.bool: attn_mask = attn_mask.where(0, -float("inf"))
       qk = qk + attn_mask
-    attn = qk.cast(self.dtype).softmax(-1).dropout(dropout_p) @ value
-    return attn.fuse() if FUSE_ATTENTION else attn
+    return qk.cast(self.dtype).softmax(-1).dropout(dropout_p) @ value
 
   def _do_reduction(self, reduction:ReductionStr="mean") -> Tensor:
     if reduction not in get_args(ReductionStr): raise ValueError(f"{reduction=} must be one of {get_args(ReductionStr)}")
@@ -4209,29 +3942,6 @@ class Tensor(MathTrait):
 
   # ***** Tensor Properties *****
 
-  @property
-  def ndim(self) -> int:
-    """
-    Returns the number of dimensions in the tensor.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor([[1, 2], [3, 4]])
-    print(t.ndim)
-    ```
-    """
-    return len(self.shape)
-
-  def numel(self) -> sint:
-    """
-    Returns the total number of elements in the tensor.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
-    print(t.numel())
-    ```
-    """
-    return prod(self.shape)
-
   def element_size(self) -> int:
     """
     Returns the size in bytes of an individual element in the tensor.
@@ -4450,18 +4160,18 @@ class Tensor(MathTrait):
     x, w = x.contiguous(), w.contiguous()
 
     # expand out
-    rcin_hi, rcin_lo = cin//4 if cin >= 4 else 1, 4 if cin >= 4 else 1
-    cout_expand = [groups//4 if cin == 1 else groups, 4 if cin == 1 else 1, rcout//4 if rcout >= 4 else 1, 4 if rcout >= 4 else 1]
+    rcin_hi, rcin_lo = (cin//4, 4) if cin >= 4 else (1, 1)
+    group_shape, rcout_expand = (groups//4, 4) if cin == 1 else (groups, 1), (rcout//4, 4) if rcout >= 4 else (1, 1)
     x = x.reshape(bs, iy, ix, groups, rcin_hi, rcin_lo)
     if cin_last: w = w.reshape(cout//4, H, rcin_hi, W, 4, rcin_lo)
     else: w = w.reshape(cout//4, H, rcin_hi, W, rcin_lo, 4).permute(0,1,2,3,5,4)
 
     # prepare input
     x = x.permute(0,3,4,5,1,2).pad(self._resolve_pool_pads(padding,2))._pool((H,W), stride, dilation)# -> (bs, groups, rcin_hi, rcin_lo, oy, ox, H, W)
-    x = x.permute(0,4,5,1,2,3,6,7).reshape(bs, (oy := x.shape[4]), (ox := x.shape[5]), *cout_expand[0:2], 1, 1, rcin_hi, rcin_lo, H, W)
+    x = x.permute(0,4,5,1,2,3,6,7).reshape(bs, (oy := x.shape[4]), (ox := x.shape[5]), *group_shape, 1, 1, rcin_hi, rcin_lo, H, W)
 
     # prepare weights
-    w = w.permute(0,4,2,5,1,3).reshape((1, 1, 1, *cout_expand, rcin_hi, rcin_lo, H, W))
+    w = w.permute(0,4,2,5,1,3).reshape((1, 1, 1, *group_shape, *rcout_expand, rcin_hi, rcin_lo, H, W))
 
     # the conv!
     ret = (x*w).cast(base_image_type((bs*oy, ox*cout//4, 4)) if IMAGE >= 2 else dtypes.float32).sum((-4, -3, -2, -1), dtype=dtype)

tinygrad/uop/__init__.py

@@ -1,3 +1,5 @@
+# flake8: noqa: E702
+# allow semicolons to put multiple ops on one line
 from enum import auto, IntEnum, Enum
 
 # wrapper around IntEnum that preserves Enum.__str__ and makes auto() unique across all FastEnum subclasses
@@ -9,9 +11,21 @@ class FastEnum(IntEnum):
 
 # the order of these Ops controls the order of the toposort
 class Ops(FastEnum):
+  # ** 1 -- defines/special **
+
+  # TODO: unify these ops into the levels of the memory hierarchy
+  DEFINE_GLOBAL = auto(); DEFINE_LOCAL = auto(); DEFINE_REG = auto()
+
+  # this is for symbolic shapes
+  DEFINE_VAR = auto(); BIND = auto()
+
+  # this is a RANGE for GPU dimensions, similar to symbolic shapes but not exactly
+  SPECIAL = auto()
+
+  # ** 2 -- non op uops **
+
   # uops that aren't rendered
-  NOOP = auto(); SINK = auto(); UNIQUE = auto(); DEVICE = auto(); KERNEL = auto(); PRECAST = auto(); REWRITE_ERROR = auto()  # noqa: E702
-  SENTINEL = auto()
+  NOOP = auto(); SINK = auto(); PRECAST = auto()
 
   # AFTER passes src[0] through and promises in the toposort that any consumers of the AFTER run after src[1:]
   AFTER = auto()
@@ -19,64 +33,70 @@ class Ops(FastEnum):
   # GROUP is a NOOP that just merges things together
   GROUP = auto()
 
-  # buffer ops
-  COPY = auto(); BUFFER = auto(); BUFFER_VIEW = auto(); MSELECT = auto(); MSTACK = auto() # noqa: E702
+  # vector creation / item selection
+  GEP = auto(); VECTORIZE = auto()
 
-  # create buffer
-  BUFFERIZE = auto()
-
-  # ops that adjust the behavior of the scheduler
-  CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto(); FUSE = auto() # noqa: E702
-
-  # movement ops! these only exist in the tensor graph
-  RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); FLIP = auto() # noqa: E702
-  MULTI = auto()  # MULTI is really a movement op
-
-  # TODO: unify these ops into the levels of the memory hierarchy. depends on ASSIGN is STORE
-  DEFINE_GLOBAL = auto(); DEFINE_LOCAL = auto(); DEFINE_REG = auto() # noqa: E702
-
-  # this is for symbolic shapes
-  DEFINE_VAR = auto(); BIND = auto() # noqa: E702
-
-  # this is a RANGE for GPU dimensions, similar to symbolic shapes but not exactly
-  SPECIAL = auto()
-
-  # reduce
-  REDUCE_AXIS = auto(); REDUCE = auto(); ALLREDUCE = auto() # noqa: E702
-
-  # optimization helper ops
-  UNROLL = auto(); CONTRACT = auto(); GEP = auto(); VECTORIZE = auto(); CAT = auto(); PTRCAT = auto() # noqa: E702
-
-  # UnaryOps
-  CAST = auto(); BITCAST = auto(); EXP2 = auto(); LOG2 = auto(); SIN = auto(); SQRT = auto(); RECIPROCAL = auto(); NEG = auto(); TRUNC = auto() # noqa: E702
-
-  # load/store before math
-  LOAD = auto(); STORE = auto() # noqa: E702
-  ASSIGN = auto()  # TODO: ASSIGN is STORE, remove ASSIGN
-
-  # tensor core math op, not elementwise
-  WMMA = auto()
+  # ** 3 -- load/store **
 
   # INDEX is a BinaryOp similar to ADD, but it operates on pointers
   INDEX = auto()
 
+  # load/store before math
+  LOAD = auto(); STORE = auto()
+
+  # ** 4 -- math **
+
+  # tensor core math op, not elementwise
+  WMMA = auto()
+
+  # UnaryOps
+  CAST = auto(); BITCAST = auto(); EXP2 = auto(); LOG2 = auto(); SIN = auto()
+  SQRT = auto(); RECIPROCAL = auto(); NEG = auto(); TRUNC = auto()
+
   # BinaryOps
-  ADD = auto(); MUL = auto(); SHL = auto(); SHR = auto(); IDIV = auto(); MAX = auto(); MOD = auto() # noqa: E702
-  CMPLT = auto(); CMPNE = auto(); CMPEQ = auto() # noqa: E702
-  XOR = auto(); OR = auto(); AND = auto() # noqa: E702
-  THREEFRY = auto(); SUB = auto(); FDIV = auto(); POW = auto() # noqa: E702
+  ADD = auto(); MUL = auto(); SHL = auto(); SHR = auto(); IDIV = auto(); MAX = auto(); MOD = auto()
+  CMPLT = auto(); CMPNE = auto(); CMPEQ = auto()
+  XOR = auto(); OR = auto(); AND = auto()
+  THREEFRY = auto(); SUB = auto(); FDIV = auto(); POW = auto()
 
   # TernaryOps
-  WHERE = auto(); MULACC = auto() # noqa: E702
+  WHERE = auto(); MULACC = auto()
+
+  # ** 5 -- control flow / consts / custom **
 
   # control flow ops
-  BARRIER = auto(); RANGE = auto(); IF = auto(); END = auto(); ENDIF = auto() # noqa: E702
+  BARRIER = auto(); RANGE = auto(); IF = auto(); END = auto(); ENDIF = auto()
 
   # consts. VCONST is a vectorized const
-  VCONST = auto(); CONST = auto() # noqa: E702
+  VCONST = auto(); CONST = auto()
 
   # CUSTOM/CUSTOMI are used to output strings into codegen. the I makes the string inline
-  CUSTOM = auto(); CUSTOMI = auto() # noqa: E702
+  CUSTOM = auto(); CUSTOMI = auto()
+
+  # ** 6 -- ops that don't exist in programs **
+
+  # tensor graph ops
+  UNIQUE = auto(); DEVICE = auto(); KERNEL = auto()
+  ASSIGN = auto()
+
+  # buffer ops
+  BUFFERIZE = auto(); COPY = auto(); BUFFER = auto(); BUFFER_VIEW = auto(); MSELECT = auto(); MSTACK = auto()
+
+  # ops that adjust the behavior of the scheduler
+  CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto()
+
+  # movement ops! these only exist in the tensor graph
+  RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); FLIP = auto()
+  MULTI = auto()  # MULTI is really a movement op
+
+  # reduce
+  REDUCE_AXIS = auto(); REDUCE = auto(); ALLREDUCE = auto()
+
+  # errors/placeholders
+  REWRITE_ERROR = auto(); SENTINEL = auto()
+
+  # expander ops
+  UNROLL = auto(); CONTRACT = auto(); CAT = auto(); PTRCAT = auto()
 
 class GroupOp:
   Unary = {Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.SQRT, Ops.RECIPROCAL, Ops.NEG, Ops.TRUNC}

tinygrad/uop/decompositions.py

@@ -282,7 +282,7 @@ def magicgu(vmax:int, d:int) -> tuple[int,int]:
 
 def fast_idiv(device: str, x: UOp, d: int, dont_cast=False) -> UOp|None:
   # If d is a power of two this is not valid for signed ints!
-  is_unsigned = True if x.vmin>=0 or x.dtype in dtypes.uints else False
+  is_unsigned = x.vmin>=0 or x.dtype in dtypes.uints
   assert d>0, "Sign should have been taken out of divisor"
   vmin,vmax = max(x.vmin, x.dtype.min), min(x.vmax, x.dtype.max)
   m,s = magicgu(max(vmax, abs(vmin)), d)
@@ -293,7 +293,7 @@ def fast_idiv(device: str, x: UOp, d: int, dont_cast=False) -> UOp|None:
     if (ret:=fast_idiv(device, x//largest_factor_of_two_in_d, d//largest_factor_of_two_in_d, dont_cast=True)) is not None: return ret
   if dont_cast: return None
   # promo_lattice needs to return an unsigned type if the type is unsigned
-  if dtypes.is_int(next_dtype := promo_lattice[x.dtype.scalar()][-1]) and is_dtype_supported(next_dtype, None if device=='' else device):
+  if dtypes.is_int(next_dtype := promo_lattice[x.dtype.scalar()][-1]) and is_dtype_supported(next_dtype, device):
     if m*vmin >= dtypes.min(next_dtype) and m*vmax <= dtypes.max(next_dtype):
       return ((x.cast(next_dtype)*m) >> s).cast(x.dtype) if is_unsigned else ((x.cast(next_dtype)*m) >> s).cast(x.dtype) + (x<0).where(x.ufix(1), 0)
   return None
@@ -318,7 +318,7 @@ def threefry2x32(x: UOp, key: UOp):
 
 powers_of_two = {2**i:i for i in range(64)}
 @functools.cache
-def get_late_rewrite_patterns(ops:tuple[Ops, ...], force_transcendental=False):
+def get_late_rewrite_patterns(ops:tuple[Ops, ...], force_transcendental):
   pat: list[tuple[UPat, Callable]] = []
   for op,f in ((Ops.EXP2, xexp2), (Ops.LOG2, xlog2), (Ops.SIN, xsin)):
     if op not in ops or force_transcendental:

tinygrad/uop/mathtraits.py

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

tinygrad/uop/ops.py

@@ -4,22 +4,26 @@ 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.mixin import OpMixin
 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
+from tinygrad.helpers import strip_parens, colored, ansilen
 if TYPE_CHECKING:
   from tinygrad.device import Buffer, MultiBuffer
 
 class AxisType(Enum):
   def __repr__(self): return str(self)
   GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto(); UNROLL = auto() # noqa: E702
-  THREAD = auto()
+  THREAD = auto(); OUTER = auto() # noqa: E702
 axis_letters = {AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u",
-                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
+                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r", AxisType.OUTER: ("O")}
 axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE",
-               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}
+               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta", AxisType.OUTER: "green"}
+
+# NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
+axis_to_pos = {AxisType.OUTER: -2, AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2,
+               AxisType.UPCAST: 3, AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
 
 range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3, Ops.END: 1}
 
@@ -48,6 +52,11 @@ def range_str(u:UOp, color=False) -> str:
   ret = '_'.join([str(x) if x >= 0 else "m"+str(-x) for x in u.arg[0:-1]])
   return colored(ret, axis_colors[u.arg[-1]]) if color else ret
 
+def multirange_str(rngs:Iterable[UOp], color=False, pad=None) -> str:
+  ret = ','.join([range_str(x, color=color) for x in sorted(rngs, key=lambda x: x.arg)])
+  if pad is not None: ret += " " * (pad-ansilen(ret))
+  return ret
+
 def consumer_map_from_toposort(lst:Iterable[UOp]):
   ret: dict[UOp, dict[UOp, None]] = {}
   for u in lst:
@@ -104,7 +113,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(OpMixin, metaclass=UOpMetaClass):
   op:Ops
   dtype:DType = dtypes.void
   src:tuple[UOp, ...] = tuple()
@@ -187,10 +196,18 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   def _shape(self) -> tuple[sint, ...]|None:
     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 | \
+      case Ops.UNIQUE | Ops.DEVICE | Ops.RANGE | Ops.LOAD | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
            Ops.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.PRECAST | Ops.CONTRACT:
         return None
 
+      case Ops.INDEX:
+        # non pointer index doesn't have a shape
+        if not isinstance(self.dtype, PtrDType): return None
+        # fully indexed doesn't have a shape. TODO: remove this
+        if self.src[0]._shape is None or len(self.src[1:]) == len(self.src[0].shape): return None
+        # pointer index
+        return self.src[0].shape[len(self.src[1:]):]
+
       # some ops init the shape
       case Ops.CONST | Ops.DEFINE_VAR | Ops.BIND: return () if self._device is not None else None
       case Ops.BUFFER: return (self.arg,)
@@ -199,7 +216,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
       case Ops.DEFINE_GLOBAL | Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return (self.ptrdtype.size,)
 
       # passthrough ops
-      case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.FUSE | Ops.AFTER | Ops.END:
+      case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.END:
         return self.src[0]._shape
 
       # ops with custom handling
@@ -287,14 +304,20 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   def _ranges(self) -> dict[UOp, None]:
     ret: dict[UOp, None] = {}
     for s in self.src: ret.update(s.ranges)
-    if (er:=self.ended_ranges):
-      for s in UOp.sink(*er).ranges:
-        if s in ret: del ret[s]
+    for er in self.ended_ranges:
+      if er.op is Ops.RANGE:
+        # if it's a single RANGE, we don't flow through it.
+        if er in ret: del ret[er]
+      else:
+        # if it's not a RANGE, we include all ranges in srcs.
+        # technically we shouldn't flow through these ranges either, but this is pre pm_add_control_flow so it's the same.
+        for s in er.ranges:
+          if s in ret: del ret[s]
     return ret
 
   @property
   def ranges(self) -> dict[UOp, None]:
-    if self.op is Ops.RANGE: return {self:None}
+    if self.op is Ops.RANGE: return {self:None} | self._ranges
     return self._ranges
 
   # *** uop evaluation ***
@@ -344,7 +367,13 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   def index(self, *srcs:UOp|None, ptr=False, **kwargs):
     return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype if ptr else self.dtype.base), (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
   def __getitem__(self, idx):
-    return self.index(*[UOp.const(dtypes.index, x) if isinstance(x, int) else x for x in argfix(idx)])
+    idx = argfix(idx)
+    assert len(idx) == len(self.shape), f"__getitem__ shape mismatch, indexing {self.shape} with {len(idx)} args"
+    if len(slice_idx:=[i for i,x in enumerate(idx) if isinstance(x, slice)]):
+      perm = self.permute(tuple([i for i in range(self.ndim) if i not in slice_idx] + slice_idx))
+      return perm.index(*[UOp.const(dtypes.index, x) if isinstance(x, int) else x for x in idx if not isinstance(x, slice)], ptr=True)
+    else:
+      return self.index(*[UOp.const(dtypes.index, x) if isinstance(x, int) else x for x in 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)
@@ -368,7 +397,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
       i = (i,)
     return UOp(Ops.GEP, self.dtype.scalar().vec(len(i)) if len(i) > 1 else self.dtype.scalar(), (self,), i)
   def load(self, *src:UOp, **kwargs): return UOp(Ops.LOAD, dtype=kwargs.pop("dtype", self.dtype.base), src=(self,)+src, **kwargs)
-  def store(self, *src:UOp, **kwargs): return UOp(Ops.STORE, kwargs.pop("dtype", dtypes.void), (self,)+src, **kwargs)
+  def store(self, src:UOp|ConstType, **kwargs):
+    return UOp(Ops.STORE, kwargs.pop("dtype", dtypes.void), (self, UOp.const(self.dtype, src) if not isinstance(src, UOp) else src), **kwargs)
   def end(self, *src:UOp):
     if len(src) == 0: return self
     return UOp(Ops.END, src=(self,)+src)
@@ -425,7 +455,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     return UOp(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs)
   def contiguous_backward(self): return self.alu(Ops.CONTIGUOUS_BACKWARD)
   def bufferize(self, *args, **kwargs): return UOp(Ops.BUFFERIZE, dtype=self.dtype, src=(self,)+args, **kwargs)
-  def fuse(self): return self.alu(Ops.FUSE)
   def allreduce(self, op, device:str|tuple[str, ...]|UOp):
     assert isinstance(self.device, tuple), f"allreduce must be on tuple {self.device} isn't"
     return UOp(Ops.ALLREDUCE, self.dtype, (self, UOp(Ops.DEVICE, arg=device) if not isinstance(device, UOp) else device), op)
@@ -534,13 +563,13 @@ 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 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 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)
 
   # in these two, we have custom logic to check if they are a no-op
-  def permute(self, arg:tuple[int, ...]): return self._mop(Ops.PERMUTE, arg, same_shape_noop=False) if arg != tuple(range(len(self.shape))) else self
-  def flip(self, arg:tuple[bool, ...]): return self._mop(Ops.FLIP, arg, same_shape_noop=False) if any(arg) and len(arg) == len(self.shape) else self
+  #def permute(self, arg:tuple[int, ...]): return self._mop(Ops.PERMUTE, arg, same_shape_noop=False) if arg != tuple(range(len(self.shape))) else self
+  #def flip(self, arg:tuple[bool, ...]): return self._mop(Ops.FLIP, arg, same_shape_noop=False) if any(arg) and len(arg) == len(self.shape) else self
 
   # *** uop UNIQUE ***
 
@@ -764,8 +793,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   # *** uop high level syntactic sugar ***
 
-  def shrink_to(self, arg:tuple[sint, ...]): return self.shrink(tuple([(0,x) for x in arg]))
-
   @staticmethod
   def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL):
     lookup = {AddrSpace.GLOBAL: Ops.DEFINE_GLOBAL, AddrSpace.LOCAL: Ops.DEFINE_LOCAL, AddrSpace.REG: Ops.DEFINE_REG}
@@ -777,8 +804,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     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 set(self:UOp, val:UOp|ConstType, end:UOp|tuple[UOp, ...]|list[UOp]=()) -> UOp:
+    return self.src[0].after(self.store(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)]
@@ -831,9 +858,10 @@ def exec_alu(op:Ops, dtype:DType, operands, truncate_output=True):
 # ***** uop helpers *****
 
 def print_uops(uops:list[UOp]):
+  uops_index = {u:i for i,u in enumerate(uops)}
   for i,u in enumerate(uops):
-    formatted_srcs = [(uops.index(x) if x.op is not Ops.CONST else f"{x.arg}") if x in uops else "--" for x in u.src]
-    print(f"{i:4d} {str(u.op):20s}: {str(u.dtype):40s} " f"{str(formatted_srcs):32s} {u.arg}")
+    formatted_srcs = [(uops_index[x] if x.op is not Ops.CONST else f"{x.arg}") if x in uops else "--" for x in u.src]
+    print(f"{i:4d} {str(u.op):20s}: {multirange_str(u.ranges, color=True, pad=10)} {str(u.dtype):40s} " f"{str(formatted_srcs):32s} {u.arg}")
 
 # ***** pattern matcher *****
 
@@ -853,7 +881,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(OpMixin):
   __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,
@@ -916,7 +944,6 @@ class UPat(MathTrait):
   def store(self, *src:UPat, **kwargs): return UPat(Ops.STORE, self.dtype, (self,)+src, **kwargs)
   def assign(self, x:UPat, **kwargs): return UPat(Ops.ASSIGN, self.dtype, (self,x), **kwargs)
   def reduce(self, *src:UPat, **kwargs): return UPat(Ops.REDUCE, self.dtype, src=(self,)+src, **kwargs)
-  def fuse(self): return self.alu(Ops.FUSE)
   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)
@@ -1309,7 +1336,8 @@ renderer_infer = PatternMatcher([
 
 def srcs(ctx, src): return f"({ctx[src[0]]},)" if len(src) == 1 else f"({', '.join([ctx[x] for x in src])})"
 def render_marg(ctx,x:UOp):
-  if x.op in {Ops.PERMUTE, Ops.FLIP}: return str(x.marg)
+  if x.op is Ops.PERMUTE: return str(x.marg)
+  if x.op is Ops.FLIP: return str(tuple([i for i,x in enumerate(x.marg) if x]))
   pieces = []
   if x.op in {Ops.RESHAPE, Ops.EXPAND}:
     pieces = [f"{ctx[a] if isinstance(a, UOp) else str(a)}" for a in x.marg]

tinygrad/uop/spec.py

@@ -89,7 +89,7 @@ _tensor_spec = PatternMatcher([
 
   # DETACH and CONTIGUOUS change how we interpret the source UOp
   # CONTIGUOUS ensures the source UOp realizes
-  (UPat((Ops.DETACH, Ops.CONTIGUOUS, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE), name="root", src=(UPat.var("x"),), arg=None),
+  (UPat((Ops.DETACH, Ops.CONTIGUOUS, Ops.CONTIGUOUS_BACKWARD), name="root", src=(UPat.var("x"),), arg=None),
    lambda root,x: root.dtype == x.dtype),
 
   # CONTIGUOUS with a range
@@ -122,7 +122,7 @@ shared_codegen_spec = PatternMatcher([
   # DEFINEs
   (UPat(Ops.DEFINE_GLOBAL, name="x"), lambda x: isinstance(x.dtype, (PtrDType, ImageDType)) and x.dtype.addrspace == AddrSpace.GLOBAL),
   (UPat(Ops.DEFINE_LOCAL, name="x"), lambda x: isinstance(x.dtype, PtrDType) and x.dtype.addrspace == AddrSpace.LOCAL),
-  (UPat(Ops.DEFINE_REG, src=()), lambda: True),
+  (UPat(Ops.DEFINE_REG, src=(), name="x"), lambda x: isinstance(x.arg, int)),
 
   # allow AFTER on buffers, GROUP anywhere
   (UPat(Ops.AFTER, src=(UPat(GroupOp.Defines|{Ops.AFTER}),), allow_any_len=True), lambda: True),
@@ -134,10 +134,6 @@ shared_codegen_spec = PatternMatcher([
   # WMMA has a <a, b, acc>
   (UPat(Ops.WMMA, src=(UPat(), UPat(), UPat()), name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) == 8),
 
-  # UNROLL/CONTRACT is used here for WMMA
-  (UPat(Ops.CONTRACT, name="x"), lambda x: x.dtype.count == prod(y[1] for y in x.arg)),
-  (UPat(Ops.UNROLL, name="x"), lambda x: x.src[0].dtype.count == prod(y[1] for y in x.arg)),
-
   # VECTORIZE/GEP
   (UPat(Ops.VECTORIZE, name="x"), lambda x: len(x.src)>1 and len(x.src) == x.dtype.vcount and all(x.dtype == y.dtype.vec(len(x.src)) for y in x.src)),
   (UPat(Ops.GEP, src=(UPat.var("src"),), name="gep"), lambda gep,src: gep.dtype == src.dtype.scalar()),
@@ -166,6 +162,10 @@ kernel_spec = PatternMatcher([
   # index is allowed here
   (UPat(GroupOp.Elementwise|{Ops.CONST, Ops.RANGE, Ops.DEFINE_VAR}, dtype=dtypes.index), lambda: True),
 
+  # UNROLL/CONTRACT is used here for WMMA
+  (UPat(Ops.CONTRACT, name="x"), lambda x: x.dtype.count == prod(y[1] for y in x.arg)),
+  (UPat(Ops.UNROLL, name="x"), lambda x: x.src[0].dtype.count == prod(y[1] for y in x.arg)),
+
   # END can end multiple axes here
   (UPat(Ops.END, src=(UPat(), UPat()), allow_any_len=True, dtype=dtypes.void), lambda: True),
 

tinygrad/uop/symbolic.py

@@ -48,8 +48,10 @@ symbolic_simple = propagate_invalid + PatternMatcher([
   (UPat.var("x") // 1, lambda x: x),   # x//1 -> x
   (UPat.var("x") // -1, lambda x: -x), # x//-1 -> -x
   ((UPat.var() % UPat.var("y")).named("base") % UPat.var("y"), lambda base,y: base),  # (x%y)%y = -> x%y (rewritten with base for speed)
-  # 4 variations of (x%c)+(x//c)*c = x   TODO: add sorting to remove some variations
+  # variations of (x%c)+(x//c)*c = x   TODO: add sorting to remove some variations
   (UPat.var("x")%UPat.cvar("c")+(UPat.var("x")//UPat.cvar("c"))*UPat.cvar("c"), lambda x,c: x), # (x%c)+(x//c)*c = x
+  ((UPat.var("x")//UPat.cvar("a"))%UPat.cvar("c")+(UPat.var("x")//UPat.cvar("b"))*UPat.cvar("c"),
+    lambda x,a,b,c: x//a if a.arg*c.arg==b.arg else None), # ((x//a)%c)+(x//a*c)*c = x//a. Note if a = 1 it degenerates to the one above
   ((UPat.var("x")//UPat.cvar("c1"))*UPat.cvar("c3")+UPat.var("x")%UPat.cvar("c1")*UPat.cvar("c2"),
     lambda x,c1,c2,c3: x*c2 if c1.arg*c2.arg==c3.arg else None), # (x%c1)*c2+(x//c1)*c3 = x*c2 if c1*c2==c3
   ((UPat.var("y")+(UPat.var("x")//UPat.cvar("c"))*UPat.cvar("c"))+UPat.var("x")%UPat.cvar("c"), lambda y,x,c: y+x),
@@ -382,13 +384,6 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   (UPat(Ops.VECTORIZE, src=UPat(Ops.CONST), name="vec"), lambda vec: UOp.const(vec.dtype, tuple(x.arg for x in vec.src))),
 ])+gep_pushing
 
-symbolic_flat = symbolic+PatternMatcher([
-  # ** combine terms (opinionated) **
-  (-1 * (UPat.var("x") + UPat.var("y")), lambda x,y: (-x)+(-y)),  # -(x+y) -> -x + -y
-  # (x+y)*c -> x*c+y*c. only for int, float has inf*0=nan issue
-  ((UPat.var("x", dtypes.index) + UPat.var("y")) * UPat.cvar("c"), lambda x,y,c: x*c+y*c),
-])
-
 # ******** we take a small aside to "simplify_valid" to rewrite valids ********
 
 def parse_valid(valid:UOp) -> tuple[UOp, bool, int]|None:
@@ -503,7 +498,7 @@ pm_simplify_valid = PatternMatcher([
 
 # this is symbolic 2.0
 REMOVE_FROM_SINK_LIKE = {Ops.UNROLL, Ops.NOOP, Ops.VECTORIZE, Ops.SINK}
-sym = symbolic_flat+pm_simplify_valid+PatternMatcher([
+sym = symbolic+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]),
   (UPat(Ops.LOAD, src=(UPat.cvar('c'))), lambda c: c),
@@ -553,4 +548,8 @@ sym = symbolic_flat+pm_simplify_valid+PatternMatcher([
       if any(x.op in REMOVE_FROM_SINK_LIKE for x in root.src) else None),
   # remove END with empty NOOP
   (UPat(Ops.END, src=(UPat(Ops.NOOP, src=(), name="noop"),), allow_any_len=True), lambda noop:noop),
+  # ** combine terms (opinionated) **
+  (-1 * (UPat.var("x") + UPat.var("y")), lambda x,y: (-x)+(-y)),  # -(x+y) -> -x + -y
+  # (x+y)*c -> x*c+y*c. only for int, float has inf*0=nan issue
+  ((UPat.var("x", dtypes.index) + UPat.var("y")) * UPat.cvar("c"), lambda x,y,c: x*c+y*c),
 ])

tinygrad/viz/index.html

@@ -284,12 +284,15 @@
     height: fit-content;
   }
   .raw-text {
-    padding: 0 8px;
+    padding-left: 15px;
     width: 100%;
     height: 100%;
     max-height: 100vh;
     overflow-x: auto;
   }
+  .raw-text > pre {
+    display: inline-block;
+  }
   .raw-text code {
     max-height: none !important;
   }

tinygrad/viz/js/index.js

@@ -488,12 +488,15 @@ async function renderProfiler() {
     }
   }
 
-  canvas.addEventListener("click", e => {
+  const clickShape = (e) => {
     e.preventDefault();
     const foundRect = findRectAtPosition(e.clientX, e.clientY);
-    if (foundRect?.step != null && foundRect?.key == null) { return switchCtx(foundRect.ctx, foundRect.step); }
+    if (foundRect?.step != null && (foundRect?.key == null || e.type == "dblclick")) { return switchCtx(foundRect.ctx, foundRect.step); }
     if (foundRect?.key != focusedShape) { focusShape(foundRect); }
-  });
+  }
+  canvas.addEventListener("click", clickShape);
+
+  canvas.addEventListener("dblclick", clickShape);
 
   canvas.addEventListener("mousemove", e => {
     const foundRect = findRectAtPosition(e.clientX, e.clientY);
@@ -538,7 +541,9 @@ document.getElementById("zoom-to-fit-btn").addEventListener("click", () => {
 
 function codeBlock(st, language, { loc, wrap }={}) {
   const code = document.createElement("code");
-  code.innerHTML = hljs.highlight(st, { language }).value;
+  // plaintext renders like a terminal print, otherwise render with syntax highlighting
+  if (language === "txt") code.appendChild(colored(st));
+  else code.innerHTML = hljs.highlight(st, { language }).value;
   code.className = "hljs";
   const ret = document.createElement("pre");
   if (wrap) ret.className = "wrap";
@@ -648,7 +653,7 @@ async function main() {
         u.li = list.appendChild(document.createElement("ul"));
         u.li.id = `step-${i}-${j}`;
         const p = u.li.appendChild(document.createElement("p"));
-        p.innerText = `${u.name}`+(u.match_count ? ` - ${u.match_count}` : '');
+        p.appendChild(colored(`${u.name}`+(u.match_count ? ` - ${u.match_count}` : '')));
         p.onclick = (e) => {
           e.stopPropagation();
           const subrewrites = getSubrewrites(e.currentTarget.parentElement);
@@ -659,7 +664,7 @@ async function main() {
       }
       for (const l of ul.querySelectorAll("ul > ul > p")) {
         const subrewrites = getSubrewrites(l.parentElement);
-        if (subrewrites.length > 0) { l.innerText += ` (${subrewrites.length})`; l.parentElement.classList.add("has-children"); }
+        if (subrewrites.length > 0) { l.appendChild(d3.create("span").text(` (${subrewrites.length})`).node()); l.parentElement.classList.add("has-children"); }
       }
     }
     return setState({ currentCtx:-1 });
@@ -714,10 +719,10 @@ async function main() {
         }
       }
       metadata.appendChild(tabulate(ret.summary.map(s => {
-        const div = d3.create("div").style("background", cycleColors(colorScheme.CATEGORICAL, s.idx)).style("width", "24px").style("height", "100%");
-        return [s.label.trim(), div.node()];
+        const div = d3.create("div").style("background", cycleColors(colorScheme.CATEGORICAL, s.idx)).style("width", "100%").style("height", "100%");
+        return [s.label.trim(), div.text(s.value.toLocaleString()).node()];
       })).node());
-    } else root.appendChild(codeBlock(ret.src, ret.lang));
+    } else root.appendChild(codeBlock(ret.src, ret.lang || "txt"));
     return document.querySelector("#custom").replaceChildren(root);
   }
   // ** UOp view (default)

tinygrad/viz/js/worker.js

@@ -29,10 +29,10 @@ onmessage = (e) => {
       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);
+  // remove additions overlay if it's empty
+  if (!g.node("addition")?.width) g.removeNode("addition");
   postMessage(dagre.graphlib.json.write(g));
   self.close();
 }

tinygrad/viz/serve.py

@@ -1,14 +1,14 @@
 #!/usr/bin/env python3
 import multiprocessing, pickle, difflib, os, threading, json, time, sys, webbrowser, socket, argparse, socketserver, functools, codecs, io, struct
-import subprocess, ctypes, pathlib
-from contextlib import redirect_stdout
+import subprocess, ctypes, pathlib, traceback
+from contextlib import redirect_stdout, redirect_stderr
 from decimal import Decimal
 from http.server import BaseHTTPRequestHandler
 from urllib.parse import parse_qs, urlparse
 from typing import Any, TypedDict, TypeVar, Generator, Callable
 from tinygrad.helpers import colored, getenv, tqdm, unwrap, word_wrap, TRACEMETA, ProfileEvent, ProfileRangeEvent, TracingKey, ProfilePointEvent, temp
 from tinygrad.uop.ops import TrackedGraphRewrite, RewriteTrace, UOp, Ops, printable, GroupOp, srender, sint, sym_infer, range_str, pyrender
-from tinygrad.uop.ops import print_uops, range_start
+from tinygrad.uop.ops import print_uops, range_start, multirange_str
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, Device
 from tinygrad.renderer import ProgramSpec
 from tinygrad.dtype import dtypes
@@ -18,7 +18,7 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
                Ops.RANGE: "#c8a0e0", Ops.ASSIGN: "#909090", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff",
                Ops.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55",
                **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
-               Ops.BUFFER_VIEW: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.FUSE: "#FFa500",
+               Ops.BUFFER_VIEW: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0",
                Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",
                Ops.BUFFERIZE: "#FF991C", Ops.REWRITE_ERROR: "#ff2e2e", Ops.AFTER: "#8A7866", Ops.END: "#524C46"}
 
@@ -78,11 +78,14 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
         label += f"\n{x.op.name}{idx} {arg}" + (f" {x.src[0].op}" if len(x.src) else "")
     try:
       if len(rngs:=u.ranges):
-        label += f"\n({','.join([range_str(x, color=True) for x in sorted(rngs, key=lambda x: x.arg[0:-1])])})"
+        label += f"\n({multirange_str(rngs, color=True)})"
       if u.op not in {Ops.BUFFER, Ops.KERNEL, Ops.ASSIGN, Ops.COPY, Ops.SINK, *GroupOp.Buffer} and u._shape is not None:
         label += f"\n{shape_to_str(u.shape)}"
       if u.op in {Ops.INDEX, Ops.BUFFERIZE}:
         label += f"\n{u.render()}"
+        ranges: list[UOp] = []
+        for us in u.src[1:]: ranges += [s for s in us.toposort() if s.op in {Ops.RANGE, Ops.SPECIAL}]
+        if ranges: label += "\n"+' '.join([f"{s.render()}={s.vmax+1}" for s in ranges])
       if u.op in {Ops.END, Ops.REDUCE} and len(trngs:=list(UOp.sink(*u.src[range_start[u.op]:]).ranges)):
         label += "\n"+' '.join([f"{range_str(s, color=True)}({s.vmax+1})" for s in trngs])
     except Exception:
@@ -193,10 +196,40 @@ def mem_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:int,
   peaks.append(peak)
   return struct.pack("<BIQ", 1, len(events), peak)+b"".join(events) if events else None
 
+def load_sqtt(profile:list[ProfileEvent]) -> None:
+  from tinygrad.runtime.ops_amd import ProfileSQTTEvent
+  if not (sqtt_events:=[e for e in profile if isinstance(e, ProfileSQTTEvent)]): return None
+  def err(name:str, msg:str|None=None) -> None:
+    step = {"name":name, "data":{"src":msg or traceback.format_exc()}, "depth":0, "query":f"/render?ctx={len(ctxs)}&step=0&fmt=counters"}
+    return ctxs.append({"name":"Counters", "steps":[step]})
+  try: from extra.sqtt.roc import decode
+  except Exception: return err("DECODER IMPORT ISSUE")
+  try: rctx = decode(profile)
+  except Exception: return err("DECODER ERROR")
+  if not rctx.inst_execs: return err("EMPTY SQTT OUTPUT", f"{len(sqtt_events)} SQTT events recorded, none got decoded")
+  steps:list[dict] = []
+  for name,waves in rctx.inst_execs.items():
+    if (r:=ref_map.get(name)): name = ctxs[r]["name"]
+    steps.append({"name":name, "depth":0, "query":f"/render?ctx={len(ctxs)}&step={len(steps)}&fmt=counters",
+                  "data":{"src":trace.keys[r].ret.src if r else name, "lang":"cpp"}})
+    for w in waves:
+      rows = [(e.inst, e.time, e.time-(w.insts[i-1].time if i else 0), e.dur, e.stall, str(e.typ).split("_")[-1]) for i,e in enumerate(w.insts)]
+      summary = [{"label":"Total Cycles", "value":w.insts[-1].time-w.insts[0].time if w.insts else 0}, {"label":"CU", "value":w.cu},
+                 {"label":"SIMD", "value":w.simd}]
+      steps.append({"name":f"Wave {w.wave_id}", "depth":1, "query":f"/render?ctx={len(ctxs)}&step={len(steps)}&fmt=counters",
+                    "data":{"rows":rows, "cols":["Instruction", "Clk", "Wait", "Duration", "Stall", "Type"], "summary":summary}})
+  ctxs.append({"name":"Counters", "steps":steps})
+
 def get_profile(profile:list[ProfileEvent]) -> bytes|None:
   # start by getting the time diffs
   for ev in profile:
     if isinstance(ev,ProfileDeviceEvent): device_ts_diffs[ev.device] = (ev.comp_tdiff, ev.copy_tdiff if ev.copy_tdiff is not None else ev.comp_tdiff)
+  # load device specific counters
+  device_decoders:dict[str, Callable[[list[ProfileEvent]], None]] = {}
+  for device in device_ts_diffs:
+    d = device.split(":")[0]
+    if d == "AMD": device_decoders[d] = load_sqtt
+  for fxn in device_decoders.values(): fxn(profile)
   # map events per device
   dev_events:dict[str, list[tuple[int, int, float, DevEvent]]] = {}
   markers:list[ProfilePointEvent] = []
@@ -244,16 +277,20 @@ def get_llvm_mca(asm:str, mtriple:str, mcpu:str) -> dict:
   for i,usage in instr_usage.items(): rows[i].append([[k, v, (v/max_usage)*100] for k,v in usage.items()])
   return {"rows":rows, "cols":["Opcode", "Latency", {"title":"HW Resources", "labels":resource_labels}], "summary":summary}
 
-def get_stdout(f:Callable) -> str:
-  with redirect_stdout(buf:=io.StringIO()): f()
+def get_stdout(f: Callable) -> str:
+  buf = io.StringIO()
+  try:
+    with redirect_stdout(buf), redirect_stderr(buf): f()
+  except Exception: traceback.print_exc(file=buf)
   return buf.getvalue()
 
-def get_render(i:int, fmt:str) -> dict|None:
+def get_render(i:int, j:int, fmt:str) -> dict|None:
+  if fmt == "counters": return ctxs[i]["steps"][j]["data"]
   if not isinstance(prg:=trace.keys[i].ret, ProgramSpec): return None
-  if fmt == "uops": return {"src":get_stdout(lambda: print_uops(prg.uops or [])), "lang":"python"}
+  if fmt == "uops": return {"src":get_stdout(lambda: print_uops(prg.uops or [])), "lang":"txt"}
   if fmt == "src": return {"src":prg.src, "lang":"cpp"}
-  lib = (compiler:=Device[prg.device].compiler).compile(prg.src)
-  disasm_str = get_stdout(lambda: compiler.disassemble(lib))
+  compiler = Device[prg.device].compiler
+  disasm_str = get_stdout(lambda: compiler.disassemble(compiler.compile(prg.src)))
   from tinygrad.runtime.support.compiler_cpu import llvm, LLVMCompiler
   if isinstance(compiler, LLVMCompiler):
     mtriple = ctypes.string_at(llvm.LLVMGetTargetMachineTriple(tm:=compiler.target_machine)).decode()
@@ -264,7 +301,7 @@ def get_render(i:int, fmt:str) -> dict|None:
 
 # ** HTTP server
 
-def get_int(query:dict[str, list[str]], k:str) -> int: return int(query[k][0])
+def get_int(query:dict[str, list[str]], k:str) -> int: return int(query.get(k,["0"])[0])
 
 class Handler(BaseHTTPRequestHandler):
   def do_GET(self):
@@ -279,7 +316,9 @@ class Handler(BaseHTTPRequestHandler):
         if url.path.endswith(".css"): content_type = "text/css"
       except FileNotFoundError: status_code = 404
     elif (query:=parse_qs(url.query)):
-      if url.path == "/render": ret, content_type = json.dumps(get_render(get_int(query, "ctx"), query["fmt"][0])).encode(), "application/json"
+      if url.path == "/render":
+        render_src = get_render(get_int(query, "ctx"), get_int(query, "step"), query["fmt"][0])
+        ret, content_type = json.dumps(render_src).encode(), "application/json"
       else:
         try: return self.stream_json(get_full_rewrite(trace.rewrites[i:=get_int(query, "ctx")][get_int(query, "idx")], i))
         except (KeyError, IndexError): status_code = 404