fix more tests

.coveragerc

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

.github/workflows/autogen.yml

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

.github/workflows/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 DISABLE_COMPILER_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 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 DISABLE_COMPILER_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 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=390 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=330 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=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
+      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
     - name: openpilot compile3 0.10.0 dmonitoring
-      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
+      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
     - 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=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
+      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
     - name: openpilot compile3 0.10.1 driving_policy
-      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
+      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
     - 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=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
+      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
     - 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

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

.gitignore

@@ -63,5 +63,3 @@ 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=6 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=8 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,9 +31,7 @@ $(for p in "$@"; do echo "  $p,"; done)
 ]
 def _try_dlopen_$name():
   library = ctypes.util.find_library("$name")
-  if library:
-    try: return ctypes.CDLL(library)
-    except OSError: pass
+  if library: return ctypes.CDLL(library)
   for candidate in PATHS_TO_TRY:
     try: return ctypes.CDLL(candidate)
     except OSError: pass
@@ -188,7 +186,6 @@ 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
 
@@ -435,13 +432,11 @@ generate_sqtt() {
     $ROCPROF_SRC/include/rocprof_trace_decoder.h \
     $ROCPROF_SRC/include/trace_decoder_instrument.h \
     $ROCPROF_SRC/include/trace_decoder_types.h \
-    -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
+    -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
 }
 
 generate_webgpu() {
@@ -536,7 +531,7 @@ generate_mesa() {
   sed -i "s/('fp_fast_math', ctypes.c_bool, 9)/('fp_fast_math', ctypes.c_uint32, 9)/" $BASE/mesa.py
   sed -i "s/('\(\w\+\)', pipe_shader_type, 8)/('\1', ctypes.c_ubyte)/" $BASE/mesa.py
   sed -i "s/\([0-9]\+\)()/\1/" $BASE/mesa.py
-  sed -i '/struct_nir_builder._pack_ = 1 # source:False/d' "$BASE/mesa.py"
+  sed -i "s/\(struct_nir_builder._pack_\) = 1/\1 = 0/" $BASE/mesa.py
   python3 -c "import tinygrad.runtime.autogen.mesa"
 }
 
@@ -550,6 +545,7 @@ 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=${LOGMLPERF:-1}
+export LOGMLPERF=1
 
 export SEED=$RANDOM
 DATETIME=$(date "+%m%d%H%M")

examples/openpilot/compile3.py

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

extra/gemm/max_matmul.py

@@ -1,11 +1,17 @@
 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 import dtypes
-from tinygrad.dtype import DTYPES_DICT
+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
 
 script_dir = os.path.dirname(os.path.abspath(__file__))
 
@@ -47,6 +53,12 @@ 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
@@ -177,11 +189,11 @@ if __name__ == "__main__":
 
     tms = []
     na, nb, nc = randoms()
-    cudaalloc._copyin(a, memoryview(bytearray(na)))
-    cudaalloc._copyin(b, memoryview(bytearray(nb)))
+    cudaalloc.copyin(a, bytearray(na))
+    cudaalloc.copyin(b, 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/metal_uop_matmul.py

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

extra/gemm/mi350x_uop_matmul.py

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

extra/gemm/simple_matmul.py

@@ -17,7 +17,7 @@ M = getenv("M", N)
 K = getenv("K", N)
 CNT = getenv("CNT", 10)
 
-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 = {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/gemm/torch_gemm.py

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

extra/mmapeak/mmapeak.py

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

extra/models/llama.py

@@ -89,20 +89,6 @@ 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

@@ -1,134 +0,0 @@
-/*
- * 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/qcom_gpu_driver/msm_kgsl.py

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

extra/sqtt/README.md

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

extra/sqtt/disasm.py

@@ -0,0 +1,68 @@
+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/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, temp
+from tinygrad.helpers import round_up, flatten, all_same
 from dataclasses import dataclass
 
 CHUNK_CLASSES = {
@@ -154,22 +154,8 @@ 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))
@@ -210,7 +196,7 @@ class RGP:
         flags=0,
         trace_shader_core_clock=0x93f05080,
         trace_memory_clock=0x4a723a40,
-        device_id={110000: 0x744c, 110003: 0x7480, 120001: 0x7550, 120000: 0x7550}[device_props['gfx_target_version']],
+        device_id={110000: 0x744c, 110003: 0x7480, 120001: 0x7550}[device_props['gfx_target_version']],
         device_revision_id=0xc8,
         vgprs_per_simd=1536,
         sgprs_per_simd=128*16,
@@ -324,7 +310,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', nargs='?', default=temp("profile.pkl", append_user=True))
+  parser.add_argument('input')
   parser.add_argument('-d', '--device')
   parser.add_argument('-o', '--output')
   args = parser.parse_args()
@@ -346,4 +332,3 @@ 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,32 +1,9 @@
 import ctypes, pathlib, argparse, pickle, re, functools, dataclasses, itertools
-from tinygrad.helpers import temp, unwrap, DEBUG
-from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
+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.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:
@@ -40,77 +17,53 @@ 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, 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]] = {}
+  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 = {}, {}, {}
 
     for prog in prog_evs:
-      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
+      for addr, info in comgr_get_address_table(prog.lib).items():
+        self.disasms[prog.base + addr] = info
+        self.addr2prg[prog.base + addr] = prog
 
   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 >= 5: print("OCC", ev.time, self.active_se, ev.cu, ev.simd, ev.wave_id, ev.start)
+    if DEBUG >= 4: print("OCC", ev.time, self.active_se, ev.cu, ev.simd, ev.wave_id, ev.start)
 
   def on_wave_ev(self, ev):
-    if DEBUG >= 5: print("WAVE", ev.wave_id, self.active_se, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time)
+    if DEBUG >= 4: print("WAVE", ev.wave_id, self.active_se, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time)
 
-    asm:dict[int, InstInfo] = {}
-    inst_execs:list[InstExec] = []
+    asm = {}
     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]
-      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.setdefault(inst_ev.pc.address, InstInfo(typ=inst_typ, inst=self.disasms[inst_ev.pc.address][0]))
       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))
 
-    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))
+    self.wave_events[(self.find_program(ev.instructions_array[0].pc.address).name, ev.wave_id, ev.cu, ev.simd)] = asm
 
-def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
-  dev_events:dict[str, ProfileDeviceEvent] = {}
+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)
   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(dev_events, sqtt_events, prog_events)
+  ROCParseCtx = _ROCParseCtx(sqtt_events, prog_events)
 
   @rocprof.rocprof_trace_decoder_se_data_callback_t
   def copy_cb(buf, buf_size, data_ptr):
@@ -127,12 +80,12 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
       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 >= 5: print(rocprof.rocprofiler_thread_trace_decoder_record_type_t__enumvalues[record_type], events_ptr, n)
+        if DEBUG >= 2: 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[(unwrap(ROCParseCtx.active_kern), pc.address)]
+    instr, mem_size_ptr[0] = ROCParseCtx.disasms[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
@@ -145,27 +98,15 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
 
     return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
 
-  try:
-    rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
-  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
+  rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
+  print('SQTT:', ROCParseCtx.wave_events.keys())
 
-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
+  for ev in pmc_events:
     print(f"PMC Event: dev={ev.device} kern={ev.kern}")
     ptr = 0
     for s in ev.sched:
       view = memoryview(ev.blob).cast('Q')
       print(f"\t{s.name}")
-      for xcc, inst, se_idx, sa_idx, wgp_idx in itertools.product(range(s.xcc), range(s.inst), range(s.se), range(s.sa), range(s.wgp)):
-        print(f"\t\tXCC {xcc} Inst {inst} SE {se_idx} SA {sa_idx} WGP {wgp_idx}: {view[ptr]:#x}")
+      for inst, se_idx, sa_idx, wgp_idx in itertools.product(range(s.inst), range(s.se), range(s.sa), range(s.wgp)):
+        print(f"\t\tInst {inst} SE {se_idx} SA {sa_idx} WGP {wgp_idx}: {view[ptr]}")
         ptr += 1

extra/sqtt/rocprof/install.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/43bf0fef74a83c3c25badfc5a09c0bd39ed8c6f9/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/5420409ad0963b2d76450add067b9058493ccbd0/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/rocprof/rocprof.py

@@ -8,19 +8,6 @@
 # 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:
@@ -168,7 +155,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'] = _try_dlopen_rocprof_trace_decoder() #  ctypes.CDLL('FIXME_STUB')
+_libraries['FIXME_STUB'] = ctypes.CDLL(ctypes.util.find_library('rocprof-trace-decoder')) #  ctypes.CDLL('FIXME_STUB')
 
 
 

extra/sqtt/test_timing.py

@@ -1,105 +0,0 @@
-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

@@ -2,7 +2,7 @@
 
 using namespace kittens;
 
-constexpr int NUM_WORKERS = 4;
+constexpr int NUM_WORKERS = 2;
 constexpr int PIPE_STAGES = 3;
 
 constexpr int ATTN_B = 16;
@@ -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*(64/D); // height of each worker tile (rows)
+template<int D> constexpr size_t ROWS = 16*(128/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

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

extra/thunder/cuda/matmul.cu

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

extra/thunder/cuda/matmul.py

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

extra/thunder/cuda/matmul2.cu

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

extra/thunder/tiny/tk/__init__.py

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

extra/thunder/tiny/tk/group.py

@@ -1,272 +0,0 @@
-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

@@ -1,57 +0,0 @@
-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

@@ -1,52 +0,0 @@
-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/tinyfs/fetch_file.py

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

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

@@ -119,7 +119,14 @@ 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,48 +7,30 @@
 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 ok = super::init();
-	if (!ok) return false;
+	auto theAnswer = super::init();
+	if (!theAnswer) {
+		return false;
+	}
 
 	ivars = IONewZero(TinyGPUDriverUserClient_IVars, 1);
-	if (!ivars) return false;
+	if (ivars == nullptr) {
+		return false;
+	}
+
 	return true;
 }
 
 void TinyGPUDriverUserClient::free()
 {
-	if (ivars) {
-		IOSafeDeleteNULL(ivars, TinyGPUDriverUserClient_IVars, 1);
+	if (ivars != nullptr) {
+		ivars->provider.reset();
 	}
+
+	IOSafeDeleteNULL(ivars, TinyGPUDriverUserClient_IVars, 1);
 	super::free();
 }
 
@@ -77,22 +59,6 @@ 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);
 }
 
@@ -136,26 +102,26 @@ kern_return_t TinyGPUDriverUserClient::ExternalMethod(uint64_t selector, IOUserC
 
 kern_return_t IMPL(TinyGPUDriverUserClient, CopyClientMemoryForType)
 {
-	if (!memory) return kIOReturnBadArgument;
-	if (!ivars->provider.get()) return kIOReturnNotAttached;
+	if (!memory) {
+		return kIOReturnBadArgument;
+	}
+
+	if (ivars->provider.get() == nullptr) {
+		return kIOReturnNotAttached;
+	}
 
-	// bar handling, type is bar num
 	if (type < 6) {
 		uint32_t bar = (uint32_t)type;
 		return ivars->provider->MapBar(bar, memory);
 	}
 
-	// dma handling, type is size
-	if (ivars->ensureDMACap(ivars->dmaCount + 1)) {
-		os_log(OS_LOG_DEFAULT, "tinygpu: cannot grow dma array");
-		return kIOReturnNoMemory;
+	// dma page buffer
+	TinyGPUCreateDMAResp buf;
+	kern_return_t err = ivars->provider->CreateDMA(type, &buf);
+	if (err) {
+		return err;
 	}
 
-	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","wozeparrot"]
+NAMES = ["chenyu","George Hotz","nimlgen","qazal","Sieds Lykles","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,8 +1,5 @@
 [pytest]
-norecursedirs =
-    extra
-    .hypothesis
-    .git
+norecursedirs = extra
 timeout = 300
 timeout_method = thread
 timeout_func_only = true

ruff.toml

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

setup.cfg

@@ -1,21 +0,0 @@
-[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,7 +32,6 @@ setup(name='tinygrad',
         'tinygrad.codegen.opt',
         'tinygrad.codegen.late',
         'tinygrad.engine',
-        'tinygrad.mixin',
         'tinygrad.nn',
         'tinygrad.renderer',
         'tinygrad.runtime',
@@ -53,7 +52,7 @@ setup(name='tinygrad',
         "License :: OSI Approved :: MIT License"
       ],
       install_requires=[],
-      python_requires='>=3.11',
+      python_requires='>=3.10',
       extras_require={
         'arm': ["unicorn"],
         'triton': ["triton-nightly>=2.1.0.dev20231014192330"],

test/external/external_benchmark_op_cat.py

@@ -1,163 +0,0 @@
-# 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,81 +0,0 @@
-# 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.codegen.opt import Opt, OptOps
-from tinygrad.renderer import ProgramSpec
-from tinygrad.engine.realize import CompiledRunner
-from tinygrad.helpers import dedup, getenv
-from tinygrad.device import Buffer
-from tinygrad.dtype import ImageDType, Invalid
-
-# 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
-
-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)
-
-  opts = None
-  return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
-
-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
-allocator = Device.default.allocator
-
-uops = full_rewrite(ast, renderer)
-src = renderer.render(uops)
-
-lib = compiler.compile(src)
-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])
-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_fuzz_am_interrupts.py

@@ -0,0 +1,39 @@
+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

@@ -1,44 +0,0 @@
-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

@@ -1,20 +0,0 @@
-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

@@ -1,345 +0,0 @@
-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,11 +85,12 @@ 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=110, amd_tflops=65)
+  def test_gemm_4096(self): self._test_matmul(4096, nv_tflops=115, 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,221 +0,0 @@
-import unittest
-from tinygrad import Tensor, UOp, Context
-from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import KernelInfo, AxisType
-
-# **** kernels ****
-
-def custom_arange_kernel(C:UOp) -> UOp:
-  i = UOp.range(C.size, 0)
-  return C[i].store(i.cast(C.dtype.base)).end(i).sink(arg=KernelInfo(name=f"custom_arange_{C.size}"))
-
-def custom_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])
-  store_d = D[i].store(A[i]*B[i])
-  return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name=f"custom_addmul_kernel_{C.size}")).simplify()
-
-def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
-  assert A.shape[1] == B.shape[0]
-  i, j, k = UOp.range(C.shape[0], 0), UOp.range(C.shape[1], 1), UOp.range(A.shape[1], 2, axis_type=AxisType.REDUCE)
-  C = C[i, j].set(0.0)
-  C = C[i, j].set(C.after(k)[i, j] + A[i, k] * B[k, j], end=k)
-  prog = C.end(i, j)
-  return prog.sink(arg=KernelInfo(name=f"custom_gemm_{C.shape[0]}_{C.shape[1]}_{A.shape[1]}", opts_to_apply=()))
-
-def custom_sum(B:UOp, A:UOp) -> UOp:
-  i = UOp.range(A.shape[0], 0, axis_type=AxisType.REDUCE)
-  B = B[0].set(0.0)
-  B = B[0].set(B.after(i)[0] + A[i], end=i)
-  return B.sink(arg=KernelInfo(name=f"custom_sum_{A.shape[0]}", opts_to_apply=()))
-
-def flip_contract_kernel(dest:UOp, src:UOp):
-  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]:
-  out, a, b = kernel.src
-  grad_a = (Tensor(gradient) @ Tensor(b).T).uop
-  grad_b = (Tensor(a).T @ Tensor(gradient)).uop
-  return (None, grad_a, grad_b)
-
-def backward_gemm_custom(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
-  out, a, b = kernel.src
-  grad_a = Tensor.empty_like(Tensor(a)).custom_kernel(Tensor(gradient), Tensor(b).T, fxn=custom_gemm)[0].uop
-  grad_b = Tensor.empty_like(Tensor(b)).custom_kernel(Tensor(a).T, Tensor(gradient), fxn=custom_gemm)[0].uop
-  return (None, grad_a, grad_b)
-
-# **** tests ****
-
-class TestCustomKernel(unittest.TestCase):
-  def test_simple(self):
-    a = Tensor.ones(16, 16).contiguous()
-    b = Tensor.ones(16, 16).contiguous()
-    c = Tensor.empty(16, 16)
-
-    c = Tensor.custom_kernel(c,a,b, fxn=custom_elementwise_add_kernel)[0]
-
-    out = c.flatten().tolist()
-    assert all(x == 2 for x in out), "all 2"
-
-  def test_multioutput(self):
-    a = Tensor.full((16, 16), 3.).contiguous()
-    b = Tensor.full((16, 16), 3.).contiguous()
-    c = Tensor.empty(16, 16)
-    d = Tensor.empty(16, 16)
-
-    c,d = Tensor.custom_kernel(c,d,a,b, fxn=custom_elementwise_addmul_kernel)[:2]
-    Tensor.realize(c,d)
-
-    assert all(x == 6 for x in c.flatten().tolist()), "all 6"
-    assert all(x == 9 for x in d.flatten().tolist()), "all 9"
-
-  def test_arange(self):
-    ref = Tensor.arange(100)
-    tst = Tensor.empty_like(ref)
-    tst = tst.custom_kernel(fxn=custom_arange_kernel)[0]
-    self.assertTrue((ref == tst).all().item())
-
-  def test_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)
-    b = b.custom_kernel(a, fxn=flip_contract_kernel)[0]
-    self.assertTrue((a.flip(1) == b).all().item())
-
-  def test_noncontig(self):
-    a = Tensor.ones(16, 16).contiguous()
-    tst = Tensor.empty_like(a)
-    b = a+1
-    b_p1 = Tensor.custom_kernel(tst, b, fxn=custom_add_one_kernel)[0]
-    self.assertTrue((b_p1 == 3).all().item())
-
-  def test_sum(self):
-    # TODO: this only works for float, and silently fails with int
-    a = Tensor([1.0, 2, 3, 4, 5])
-    tst = Tensor.empty(1)
-    b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
-    self.assertEqual(b.item(), 15)
-
-  def test_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)
-    b = Tensor.randn(N, N)
-    c = Tensor.empty(N, N)
-
-    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
-    err = (tst - (a@b)).square().max()
-    self.assertLess(err.item(), 1e-6)
-
-  def test_gemm_backward_custom(self): self.test_gemm_backward(True)
-  # NOTE: grad_fxn doesn't work with pyrender
-  @Context(SPEC=1)
-  def test_gemm_backward(self, custom_backward_gemm=False):
-    N = 4
-    a_rand = Tensor.randn(N, 8)
-    b_rand = Tensor.randn(8, N)
-    Tensor.realize(a_rand, b_rand)
-
-    a, b = Tensor(a_rand.numpy(), requires_grad=True), Tensor(b_rand.numpy(), requires_grad=True)
-    c = Tensor.empty(N, N)
-    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm, grad_fxn=backward_gemm_custom if custom_backward_gemm else backward_gemm)[0]
-    tst.sum().backward()
-    grad_a, grad_b = a.grad, b.grad
-    Tensor.realize(tst, grad_a, grad_b)
-
-    a, b = Tensor(a_rand.numpy(), requires_grad=True), Tensor(b_rand.numpy(), requires_grad=True)
-    ref = (a@b)
-    ref.sum().backward()
-    real_grad_a, real_grad_b = a.grad, b.grad
-    Tensor.realize(ref, real_grad_a, real_grad_b)
-
-    err = (tst - ref).square().max()
-    self.assertLess(err.item(), 1e-6)
-
-    err = (grad_a - real_grad_a).square().max()
-    self.assertLess(err.item(), 1e-6)
-
-    err = (grad_b - real_grad_b).square().max()
-    self.assertLess(err.item(), 1e-6)
-
-  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, AxisType, PatternMatcher, graph_rewrite, UPat
+from tinygrad.uop.ops import UOp, Ops, GroupOp
 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,22 +38,6 @@ 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])
@@ -278,8 +262,6 @@ 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)
 
@@ -304,27 +286,6 @@ 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,)
@@ -471,8 +432,6 @@ 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,12 +596,6 @@ 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,10 +1551,8 @@ class TestOps(unittest.TestCase):
                                 lambda x: Tensor.stack(*x.std_mean(axis=(1,2))))
 
   def test_std_mean_loaded_nan(self):
-    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))))
+    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)
@@ -2822,13 +2820,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,),)])
-    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]]], 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):
@@ -2843,10 +2841,11 @@ 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,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,(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

@@ -1,18 +0,0 @@
-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,12 +300,6 @@ 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

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

test/test_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)
+    out_single = fxn(*args, **kwargs).fuse()
     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)
+    attn = q.scaled_dot_product_attention(k, v).fuse()
     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))
+    c = (a.sum(axis=1) + b.sum(axis=1)).fuse()
     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))
+    out = (C + (A @ B)).fuse()
     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)
+      out = self.test.contiguous().softmax(-1).fuse()
       run_one_schedule_item(out)
 
     np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)

test/test_uops.py

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

test/unit/test_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 Exception: pass"
+    expect_failure = "\ntry: assert Device[Device.DEFAULT].compiler is None;\nexcept RuntimeError: pass"
 
     if Device.DEFAULT == "CPU":
       from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangJITCompiler

test/unit/test_helpers.py

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

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.assertIs(a[...], a)
+    self.assertIsNot(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.assertIs(a, a[()])
-    self.assertIs(a, a[...])
-    self.assertIs(a, a[:])
+    self.assertIsNot(a, a[()])
+    self.assertIsNot(a, a[...])
+    self.assertIsNot(a, a[:])
 
   def test_broaderrors_indexing(self):
     a = Tensor.zeros(5, 5)

test/unit/test_uop_symbolic.py

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

tinygrad/codegen/__init__.py

@@ -1,10 +1,9 @@
 from typing import cast
-import itertools
 from tinygrad.helpers import DEVECTORIZE, TRANSCENDENTAL, SPEC
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat
 from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
 from tinygrad.renderer import Renderer
-from tinygrad.dtype import dtypes, PtrDType
+from tinygrad.dtype import dtypes
 from tinygrad.helpers import panic
 
 # import all pattern matchers here
@@ -16,23 +15,14 @@ from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_in
   ReduceContext, correct_load_store, pm_render, pm_add_loads
 from tinygrad.codegen.opt.postrange import apply_opts
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse, pm_split_store
-from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops
+from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen
 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+pm_syntactic_sugar, name="early movement ops", bottom_up=True)
-
   # first we optimize
   if optimize:
     # collapse loads reduce (indexing by a tensor)
@@ -60,7 +50,7 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
   sink = graph_rewrite(sink, sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander")
 
   # add locals
-  sink = graph_rewrite(sink, pm_add_buffers_local+rangeify_codegen, ctx=itertools.count(0), name="add local buffers")
+  sink = graph_rewrite(sink, pm_add_buffers_local+rangeify_codegen, name="add local buffers")
 
   # ** devectorizer (full_graph_rewrite) **
   # remove reduce
@@ -136,7 +126,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, f"all ranges must end by the sink, {full_sink.ranges}"
+  assert len(full_sink.ranges) == 0, "all ranges must end by the sink"
   lst = line_rewrite(linearize(full_sink), pm_linearize_cleanups)
   if SPEC: type_verify(lst, program_spec)
   return lst

tinygrad/codegen/gpudims.py

@@ -26,35 +26,28 @@ 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: 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 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 len(limited) < len(dims):
     ret = []
-    if (contraction:=get_contraction(dims, limited)) is None: raise RuntimeError(f"get_contraction should not be None {dims=} {limited=}")
+    if (contraction:=get_contraction(dims, limited)) is None: raise AssertionError(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)
-    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
+  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
 
 def add_gpudims(ctx:Renderer, s:UOp):
   if s.arg is None: return None
@@ -87,7 +80,7 @@ def add_gpudims(ctx:Renderer, s:UOp):
   subs = {}
   for r in s_topo:
     # look for local INDEXes that are not used in the GLOBAL store, then add them as an INVALID
-    if r.op is Ops.STORE and r.buf_target().ptrdtype.addrspace == AddrSpace.GLOBAL:
+    if r.op is Ops.STORE and r.src[0].src[0].ptrdtype.addrspace == AddrSpace.GLOBAL:
       idx = r.src[0]
       missing_locals = [all_ranges[rng] for rng in local_dims if all_ranges[rng] not in idx.ranges]
       if len(missing_locals):

tinygrad/codegen/late/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, invalid_gate
+from tinygrad.uop.symbolic import uop_given_valid, parse_valid, sym, symbolic_flat, 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+load_store_indexing, name=f"index_buf_{buf.arg}")
+                       symbolic_flat+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/expander.py

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

tinygrad/codegen/late/linearizer.py

@@ -1,59 +1,42 @@
 import heapq
-from typing import Any
 from collections import defaultdict
-from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat, multirange_str
-from tinygrad.helpers import prod, getenv, TUPLE_ORDER
+from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat
 
-def linearize(sink:UOp) -> list[UOp]:
+def linearize(u:UOp) -> list[UOp]:
   # this is a toposort with priority
-  lst = list(sink.toposort())
+  lst = list(u.toposort())
   consumers: defaultdict[UOp, list[UOp]] = defaultdict(list)
   in_degree:dict[UOp, int] = {}
-  out_degree:dict[UOp, int] = {}
-  priorities:dict[UOp, tuple[int, int, Any]] = {}
+  priorities:dict[UOp, int] = {}
 
   # 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)
-    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)
+    # 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)
 
   # number the uops in "ideal" order
-  nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: priorities[x]+(x.tuplize if TUPLE_ORDER else ())))}
+  nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: (priorities[x],)+x.tuplize))}
 
-  # then force them to be toposorted in as close to the ideal order as possible
-  heap = [(-nkey[sink], sink)]
+  # 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])
   newlst = []
   while heap:
     newlst.append(u:=heapq.heappop(heap)[1])
-    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]}")
+    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)}"
   return newlst
 
 class CFGContext:
@@ -81,10 +64,7 @@ class CFGContext:
       # ranges that have dependencies on other siblings need to be scheduled after them
       order = sorted(v, key=lambda x: len([u for u in v if u in deps[x]]))
       zipped = zip(order, order[1:]) if k.op is Ops.SINK else zip([k.src[1]] + order, order)
-      for x,y in zipped:
-        # TODO: this can happen! it causes infinite loop in shufflenet
-        assert y.src[1] not in x.backward_slice_with_self
-        self.edges[y.src[1]] = x
+      for x,y in zipped: self.edges[y.src[1]] = x
 
 pm_add_control_flow = PatternMatcher([
   (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(src=x.src+(y,)) if (y:=ctx.edges.get(x)) is not None else None),
@@ -92,7 +72,7 @@ pm_add_control_flow = PatternMatcher([
 
 def do_split_ends(e:UOp):
   ret = e.src[0]
-  for r in sorted(UOp.sink(*e.src[1:]).ranges, key=lambda x: x.arg, reverse=True): ret = ret.end(r)
+  for r in list(UOp.sink(*e.src[1:]).ranges)[::-1]: ret = ret.end(r)
   return ret
 
 pm_split_ends = PatternMatcher([

tinygrad/codegen/opt/heuristic.py

@@ -73,15 +73,13 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
           if first_reduce_rng.src[0].divides(MV_THREADS_PER_ROW) is not None and k.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0:
             if DEBUG >= 3:
               print(f"MATVEC: {k.full_shape=} {first_reduce_rng.render()} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}")
-            try:
-              if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW))
-            except KernelOptError: pass
+            if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW))
             if MV_BLOCKSIZE > 1: k.apply_opt(Opt(OptOps.LOCAL, global_idx, MV_BLOCKSIZE))
             if MV_ROWS_PER_THREAD > 1: k.apply_opt(Opt(OptOps.UPCAST, global_idx, MV_ROWS_PER_THREAD))
             return k
 
   # are we grouping? (requires local shape support)
-  if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= (128 if NOLOCALS else 2048), False):
+  if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= 2048, False):
     for sz in [16]:
       try:
         k.apply_opt(Opt(OptOps.GROUPTOP, 0, sz))
@@ -107,7 +105,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) and (k.upcast_size() < 32):
+  while resolve(prod(k.output_shape[i] for i in k.upcastable_dims) >= 1024):
     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,7 +133,8 @@ 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:
-    if k.unrollable_dims and (k.upcast_size() <= 4 or not k.axes_of(AxisType.UNROLL)) and (k.upcast_size() < 64):
+    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 (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,8 +2,7 @@ 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
-from tinygrad.uop.ops import axis_letters, axis_colors, axis_to_pos
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp, axis_letters, axis_colors
 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
@@ -13,6 +12,10 @@ 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
@@ -61,9 +64,9 @@ class Scheduler:
     return self.ast.replace(arg=KernelInfo(name=name, applied_opts=tuple(self.applied_opts), dont_use_locals=self.dont_use_locals), tag=1)
 
   def _output_rngs(self) -> list[UOp]:
-    return flatten([[r for r in UOp.sink(*s.src[1:]).ranges if r.arg[-1] != AxisType.REDUCE] for s in self.ast.src if s.op is Ops.END])
+    return flatten([list(UOp.sink(*s.src[1:]).ranges) for s in self.ast.src if s.op is Ops.END])
   def _globalizable_rngs(self) -> list[UOp]:
-    ret = [r for r in self._output_rngs() if r.arg[-1] == AxisType.LOOP]
+    ret = self._output_rngs()
     # exclude any output ranges from global that don't appear in all BUFFERIZE
     for x in self.ast.toposort():
       if x.op is Ops.BUFFERIZE:
@@ -102,8 +105,6 @@ 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) \
@@ -216,7 +217,8 @@ class Scheduler:
     return ret
 
   def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> None|list[UOp]:
-    if not (reduceops := self.reduceops): raise KernelOptError("no reduce ops for TensorCore")
+    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")
     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]
@@ -310,10 +312,9 @@ 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:
-    if not (red := self.reduceops): return None
+    red = [x for x in self.ast.backward_slice if x.op is Ops.REDUCE]
+    if not len(red): 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/opt/search.py

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

tinygrad/codegen/opt/tc.py

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

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

tinygrad/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, select_first_inited
+from tinygrad.helpers import unwrap_class_type, suppress_finalizing
 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); props:dict[str,Any]|None=None # noqa: E702
+  device:str; comp_tdiff:decimal.Decimal=decimal.Decimal(0); copy_tdiff:decimal.Decimal=decimal.Decimal(0) # 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)
 
-    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")
+    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
 
     if DEBUG >= 1: print(f"{self.device}: using {self.compiler.__class__.__name__}")
 
@@ -300,6 +300,10 @@ 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.
@@ -329,7 +333,6 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
     return device in {"AMD", "PYTHON", "NULL"}
   if dtype in dtypes.fp8s:
     if device in {"CUDA", "NV"}: return not CI and not getenv(f"{device}_PTX") and not getenv("NV_NAK")
-    if device == "AMD": return not CI and not AMD_LLVM and getattr(Device["AMD"], "target") in {(9,4,2), (9,5,0)}
     return device in {"PYTHON", "NULL"}
   if device == "WEBGPU": return dtype in [dtypes.bool, dtypes.char, dtypes.uchar, dtypes.short,
                                           dtypes.ushort, dtypes.float, dtypes.int32, dtypes.uint32, dtypes.half]
@@ -339,7 +342,6 @@ 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/engine/realize.py

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

tinygrad/gradient.py

@@ -29,18 +29,15 @@ 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), lambda ctx: (ctx,)),
+  (UPat((Ops.CONTIGUOUS, Ops.FUSE)), 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([i for i,x in enumerate(ret.marg) if x]),)),
+  (UPat(Ops.FLIP, name="ret"), lambda ctx, ret: (ctx.flip(ret.marg),)),
   (UPat(Ops.MULTI, name="ret"), lambda ctx, ret: ctx.shard(ret.device, ret.axis).src),
-  # NOTE: this is only correct when the KERNEL has a single output
-  (UPat(Ops.AFTER), lambda ctx: (ctx, ctx)),
-  (UPat(Ops.KERNEL, name="k"), lambda ctx, k: k.arg.grad_fxn(ctx, k)),
   # there's no gradient for bitcast
   (UPat(Ops.BITCAST), lambda: (None,)),
 ])

tinygrad/helpers.py

@@ -44,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 (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 strip_parens(fst:str): return fst[1:-1] if fst[0] == '(' and fst[-1] == ')' and fst[1:-1].find('(') <= fst[1:-1].find(')') else fst
 def ceildiv(num, amt): return int(ret) if isinstance((ret:=-(num//-amt)), float) else ret
 def round_up(num:int, amt:int) -> int: return (num+amt-1)//amt * amt
 def round_down(num:int, amt:int) -> int: return -round_up(-num, amt)
@@ -114,13 +113,6 @@ 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')
@@ -177,6 +169,7 @@ 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)
@@ -186,8 +179,6 @@ 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

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

tinygrad/mixin/movement.py

@@ -1,328 +0,0 @@
-# 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).where(vals, 0).sum(-2, dtype=vals.dtype)
+    return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
 
 class LSTMCell:
   """

tinygrad/renderer/__init__.py

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

tinygrad/renderer/cstyle.py

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

tinygrad/renderer/llvmir.py

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

tinygrad/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}".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[0]}".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/am/am.py

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

tinygrad/runtime/autogen/am/pm4_nv.py

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

tinygrad/runtime/autogen/am/pm4_soc15.py

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

tinygrad/runtime/autogen/am/sdma_4_0_0.py

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

tinygrad/runtime/autogen/am/sdma_4_4_2.py

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

tinygrad/runtime/autogen/am/sdma_5_0_0.py

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

tinygrad/runtime/autogen/am/sdma_6_0_0.py

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

tinygrad/runtime/autogen/am/smu_v13_0_0.py

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

tinygrad/runtime/autogen/am/smu_v14_0_2.py

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

tinygrad/runtime/autogen/amd_gpu.py

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

tinygrad/runtime/autogen/comgr.py

@@ -15,9 +15,7 @@ PATHS_TO_TRY = [
 ]
 def _try_dlopen_amd_comgr():
   library = ctypes.util.find_library("amd_comgr")
-  if library:
-    try: return ctypes.CDLL(library)
-    except OSError: pass
+  if library: return ctypes.CDLL(library)
   for candidate in PATHS_TO_TRY:
     try: return ctypes.CDLL(candidate)
     except OSError: pass
@@ -54,8 +52,6 @@ else:
     c_long_double_t = ctypes.c_ubyte*16
 
 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
     @classmethod
     def as_dict(cls, self):
         result = {}

tinygrad/runtime/autogen/cuda.py

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

tinygrad/runtime/autogen/hip.py

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

tinygrad/runtime/autogen/hsa.py

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