store

this can happen later
broken
2025-07-22 19:28:39 -07:00 · 2025-07-22 19:12:37 -07:00 · 2025-07-22 19:05:11 -07:00 · 2025-07-22 19:03:06 -07:00 · 2025-07-22 18:54:14 -07:00 · 2025-07-22 18:52:26 -07:00
90 changed files with 871 additions and 9745 deletions
@@ -112,16 +112,7 @@ runs:
        fi

    # ******************* apt *******************
-    - name: Setup apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
-      shell: bash
-      run: |
-        sudo chown -R $USER:$USER /var/cache/apt/archives

-        echo 'Acquire::GzipIndexes "true";' | sudo tee /etc/apt/apt.conf.d/gzip
-        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
-        echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' | sudo tee -a /etc/apt/apt.conf.d/99keep-debs
-    
    - name: Add OpenCL Repo
      if: inputs.opencl == 'true' && runner.os == 'Linux'
      shell: bash
@@ -144,11 +135,14 @@ runs:
        wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
        echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-20 main" | sudo tee /etc/apt/sources.list.d/llvm.list

-    - name: Compute Package List + Hash
+    - name: apt-get update + install
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
-      id: apt-pkgs
      shell: bash
      run: |
+        echo 'Acquire::GzipIndexes "true";' | sudo tee /etc/apt/apt.conf.d/gzip
+        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
+        sudo apt -qq update || true
+
        pkgs=""
        # **** OpenCL ****
        if [[ "${{ inputs.opencl }}" == "true" ]]; then
@@ -159,7 +153,7 @@ runs:
        fi
        # **** AMD ****
        if [[ "${{ inputs.amd }}" == "true" ]]; then
-          pkgs+=" hsa-rocr comgr hsa-rocr-dev liburing-dev libibverbs-dev libc6-dev"
+          pkgs+=" hsa-rocr comgr hsa-rocr-dev liburing-dev libc6-dev"
        fi
        # **** CUDA ****
        if [[ "${{ inputs.cuda }}" == "true" ]]; then
@@ -174,31 +168,14 @@ runs:
        if [[ "${{ inputs.llvm }}" == "true" ]]; then
          pkgs+=" libllvm20 clang-20 lld-20"
        fi
-        
-        echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
-        echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"

-    - name: Cache apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
-      uses: actions/cache@v4
-      with:
-        path: /var/cache/apt/archives/
-        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}
-
-    - name: Run apt Update + Install
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
-      shell: bash
-      run: |
-        sudo apt -qq update || true
-        
        # ******** do install ********
-        if [[ -n "${{ steps.apt-pkgs.outputs.pkgs }}" ]]; then
-          sudo apt-get -y --allow-unauthenticated --no-install-recommends install ${{ steps.apt-pkgs.outputs.pkgs }}
+        if [[ -n "$pkgs" ]]; then
+          sudo apt-get -y --allow-unauthenticated --no-install-recommends install $pkgs
        fi
-        
-        sudo chown -R $USER:$USER /var/cache/apt/archives/
-    
+
    # **** AMD ****
+
    - name: Setup AMD (Linux)
      if: inputs.amd == 'true' && runner.os == 'Linux'
      shell: bash
@@ -251,7 +228,7 @@ runs:
      shell: bash
      run: |
        cd ${{ github.workspace }}/gpuocelot/ocelot/build
-        sudo cp libgpuocelot.${{ runner.os == 'macOS' && 'dylib' || 'so' }} /usr/${{ runner.os == 'macOS' && 'local/' || '' }}lib/
+        sudo cp libgpuocelot.${{ runner.os == 'macOS' && 'dylib' || 'so' }} /usr/${{ runner.os == 'macOS' && 'local/' || ''}}lib/

    # **** WebGPU ****

@@ -617,10 +617,6 @@ jobs:
      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/e8bea2c78ffa92685ece511e9b554122aaf1a79d/selfdrive/modeld/models/supercombo.onnx
    - name: openpilot dmonitoring compile3 0.9.7
      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/dmonitoring_model.onnx
-    - name: openpilot compile3 Space Lab policy + vision
-      run: |
-        PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
-        PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
    - name: benchmark MobileNetV2 on DSP
      run: |
        # generate quantized weights
@@ -641,127 +637,3 @@ jobs:
          openpilot_0_9_7.txt
          openpilot_image_0_9_4.txt
          openpilot_image_0_9_7.txt
-
-  testreddriverbenchmark:
-    name: AM Benchmark
-    runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 15
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Remove amd modules
-      run: ./extra/hcq/hcq_smi.py amd rmmod
-    - name: Kill stale pids
-      run: ./extra/hcq/hcq_smi.py amd kill_pids
-    - name: Symlink models and datasets
-      run: |
-        mkdir -p weights
-        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
-        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
-        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
-        ln -s /raid/weights/mixtral-8x7b-32kseqlen weights/mixtral-8x7b-32kseqlen
-        ln -s /raid/weights/LLaMA-2 weights/LLaMA-2
-        mkdir -p extra/datasets
-        ln -s /raid/datasets/imagenet extra/datasets/imagenet
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: reset process replay
-      run: test/external/process_replay/reset.py
-    - name: Test driver cold start time
-      run: time DEBUG=3 AMD=1 AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
-    - name: Test driver warm start time
-      run: time DEBUG=3 AMD=1 python3 test/test_tiny.py TestTiny.test_plus
-    # Fails on 9070
-    # - name: Test tensor cores
-    #   run: |
-    #     AMD=1 AMD_LLVM=0 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
-    #     AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
-    #     AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-    - name: Run Tensor Core GEMM (AMD)
-      run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee am_matmul_amd.txt
-    - name: Test AMD=1
-      run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
-    - name: Test DISK copy time
-      run: AMD=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
-    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
-    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
-    - name: Run 10 MLPerf Bert training steps (1 gpu)
-      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee am_train_bert_one_gpu.txt
-    - uses: actions/upload-artifact@v4
-      with:
-        name: Speed (AM Driver)
-        path: |
-          am_matmul_amd.txt
-          am_train_cifar_one_gpu.txt
-          am_train_resnet_one_gpu.txt
-          am_train_bert_one_gpu.txt
-    - name: Run process replay tests
-      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
-
-  testgreendriverbenchmark:
-    name: NV Benchmark
-    runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 15
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Remove nv modules
-      run: ./extra/hcq/hcq_smi.py nv rmmod
-    - name: Kill stale pids
-      run: ./extra/hcq/hcq_smi.py nv kill_pids
-    - name: Symlink models and datasets
-      run: |
-        mkdir -p weights
-        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
-        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
-        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
-        ln -s /raid/weights/mixtral-8x7b-32kseqlen weights/mixtral-8x7b-32kseqlen
-        ln -s /raid/weights/LLaMA-2 weights/LLaMA-2
-        mkdir -p extra/datasets
-        ln -s /raid/datasets/imagenet extra/datasets/imagenet
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: reset process replay
-      run: test/external/process_replay/reset.py
-    - name: Test driver start time
-      run: time DEBUG=3 NV=1 python3 test/test_tiny.py TestTiny.test_plus
-    - name: Test tensor cores
-      run: NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
-    - name: Test DISK copy time
-      run: NV=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
-    - name: Test LLAMA-3
-      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
-    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
-    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
-    - name: Run 10 MLPerf Bert training steps (1 gpu)
-      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt
-    - uses: actions/upload-artifact@v4
-      with:
-        name: Speed (NV Driver)
-        path: |
-          nv_llama3_beam.txt
-          nv_train_cifar_one_gpu.txt
-          nv_train_resnet_one_gpu.txt
-          nv_train_bert_one_gpu.txt
-    - name: Run process replay tests
-      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
@@ -12,7 +12,7 @@ jobs:
  run_script_job:
    runs-on: [self-hosted, Linux, tinybox]
    if: github.repository_owner == 'tinygrad'
-    timeout-minutes: 360
+    timeout-minutes: 240

    steps:
    - name: Checkout Code
@@ -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"
@@ -132,13 +132,10 @@ jobs:
      run: |
        cp tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
        cp tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
-        cp tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
        ./autogen_stubs.sh libc
        ./autogen_stubs.sh io_uring
-        ./autogen_stubs.sh ib
        diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
        diff /tmp/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
-        diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
    - name: Verify WebGPU autogen
      run: |
        cp tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
@@ -338,7 +335,6 @@ jobs:
        python -m mypy --strict-equality --lineprecision-report .
        cat lineprecision.txt
        python -m mypy --strict-equality extra/onnx_parser.py
-        python -m mypy --strict-equality extra/onnx.py

  unittest:
    name: Unit Tests
@@ -512,6 +508,10 @@ jobs:
        run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_ops.py
      - name: Test Quantize ONNX
        run: CPU=1 PYTHONPATH=. python3 test/test_quantize_onnx.py
+      - name: Run REMOTE=1 Test (without process replay)
+        run: |
+          # TODO: re enable process replay, currently remote schedule opens devices
+          CAPTURE_PROCESS_REPLAY=0 REMOTEDEV=CPU REMOTE=1 python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -535,6 +535,10 @@ jobs:
          opencl: 'true'
      - name: Test ONNX (GPU)
        run: GPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+      - name: Run REMOTE=1 Test
+        run: |
+          REMOTEDEV=GPU REMOTE=1 python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py
+          REMOTEDEV=GPU IMAGE=2 REMOTE=1 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
      - name: Test Optimization Helpers
        run: PYTHONPATH="." DEBUG=1 python3 extra/optimization/test_helpers.py
      #- name: Test Action Space
@@ -895,11 +899,12 @@ jobs:
          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py

  amdremote:
-    name: Linux (remote)
+    name: Linux (remote amd)
    runs-on: ubuntu-22.04
    timeout-minutes: 20
    env:
      REMOTE: 1
+      HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
      PYTHONPATH: ${{ github.workspace }}
    steps:
      - name: Checkout Code
@@ -907,58 +912,38 @@ jobs:
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
-          key: linux-remote
+          key: linux-remote-amd
          deps: testing_minimal
          amd: 'true'
          llvm: 'true'
-          opencl: 'true'
      - name: Start remote server
        run: |
          start_server() {
            systemd-run --user \
              --unit="$1" \
-              --setenv=REMOTEDEV="$2" \
+              --setenv=REMOTEDEV=AMD \
              --setenv=MOCKGPU=1 \
              --setenv=PYTHONPATH=. \
-              --setenv=PORT="$3" \
+              --setenv=PORT="$2" \
              --working-directory="$(pwd)" \
              python tinygrad/runtime/ops_remote.py
          }

-          start_server "remote-server-amd-1" "AMD" 6667
-          start_server "remote-server-amd-2" "AMD" 6668
-          start_server "remote-server-gpu" "GPU" 7667
-          start_server "remote-server-cpu" "CPU" 8667
+          start_server "remote-server-1" 6667
+          start_server "remote-server-2" 6668
      - name: Check Device.DEFAULT and print some source
-        env:
-          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
        run: |
          python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
          python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'AMD', Device.default.properties.real_device"
          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
-      - name: Run REMOTE=1 Test (AMD)
-        env:
-          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
+      - name: Run REMOTE=1 Test
        run: |
          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py
-      - name: Run REMOTE=1 Test (GPU)
-        env:
-          HOST: 127.0.0.1:7667*6
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py
-          IMAGE=2 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
-      - name: Run REMOTE=1 Test (CPU)
-        env:
-          HOST: 127.0.0.1:8667*6
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py
      - name: Show remote server logs
        if: always()
        run: |
-          journalctl --user -u remote-server-amd-1 --no-pager
-          journalctl --user -u remote-server-amd-2 --no-pager
-          journalctl --user -u remote-server-gpu --no-pager
-          journalctl --user -u remote-server-cpu --no-pager
+          journalctl --user -u remote-server-1 --no-pager
+          journalctl --user -u remote-server-2 --no-pager

  osxtests:
    strategy:
@@ -240,21 +240,6 @@ generate_io_uring() {
  fixup $BASE/io_uring.py
 }

-generate_ib() {
-  clang2py -k cdefstum \
-    /usr/include/infiniband/verbs.h \
-    /usr/include/infiniband/verbs_api.h \
-    /usr/include/infiniband/ib_user_ioctl_verbs.h \
-    /usr/include/rdma/ib_user_verbs.h \
-    -o $BASE/ib.py
-
-  sed -i "s\import ctypes\import ctypes, ctypes.util\g" "$BASE/ib.py"
-  sed -i "s\FIXME_STUB\libibverbs\g" "$BASE/ib.py"
-  sed -i "s\FunctionFactoryStub()\ctypes.CDLL(ctypes.util.find_library('ibverbs'), use_errno=True)\g" "$BASE/ib.py"
-
-  fixup $BASE/ib.py
-}
-
 generate_libc() {
  clang2py -k cdefstum \
    $(dpkg -L libc6-dev | grep sys/mman.h) \
@@ -480,7 +465,6 @@ 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
-elif [ "$1" == "ib" ]; then generate_ib
 elif [ "$1" == "libc" ]; then generate_libc
 elif [ "$1" == "llvm" ]; then generate_llvm
 elif [ "$1" == "kgsl" ]; then generate_kgsl
@@ -1,6 +1,4 @@
-import functools
-import hashlib
-import os, random, pickle, queue, struct, math
+import os, random, pickle, queue
 from typing import List
 from pathlib import Path
 from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu_count
@@ -8,7 +6,6 @@ from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu
 import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX
-from tinygrad.nn.state import TensorIO

 ### ResNet

@@ -513,202 +510,6 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
      # happens with BENCHMARK set
      pass

-# llama3
-
-class BinIdxDataset:
-  def __init__(self, base_path:Path):
-    self.idx_t = Tensor(base_path.with_name(f"{base_path.name}.idx"))
-    self.idx = TensorIO(self.idx_t)
-
-    # parse idx file
-    magic = self.idx.read(9)
-    assert magic == b"MMIDIDX\x00\x00", "invalid index file format"
-    version, = struct.unpack("<Q", self.idx.read(8))
-    assert version == 1, "unsupported index version"
-    dtype_code, = struct.unpack("<B", self.idx.read(1))
-    self.dtype = {1:dtypes.uint8, 2:dtypes.int8, 3:dtypes.int16, 4:dtypes.int32, 5:dtypes.int64, 6:dtypes.float64, 7:dtypes.double, 8:dtypes.uint16}[dtype_code]
-    self.count, = struct.unpack("<Q", self.idx.read(8))
-    doc_count, = struct.unpack("<Q", self.idx.read(8))
-
-    start = self.idx.tell()
-    end = start + self.count * dtypes.int32.itemsize
-    self.sizes = self.idx_t[start:end].bitcast(dtypes.int32)
-
-    start = end
-    end = start + self.count * dtypes.int64.itemsize
-    self.pointers = self.idx_t[start:end].bitcast(dtypes.int64)
-
-    start = end
-    end = start + doc_count * dtypes.int64.itemsize
-    self.doc_idx = self.idx_t[start:end].bitcast(dtypes.int64)
-
-    # bin file
-    self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin"))
-
-  def _index(self, idx) -> tuple[int, int]:
-    return self.pointers[idx].item(), self.sizes[idx].item()
-
-  def get(self, idx, offset:int=0, length:int|None=None):
-    ptr, size = self._index(idx)
-    if length is None: length = size - offset
-    ptr += offset * self.dtype.itemsize
-    return self.bin_t[ptr:ptr+length*self.dtype.itemsize].bitcast(self.dtype).to(None)
-
-# https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/datasets.html
-class GPTDataset:
-  def __init__(self, base_path:Path, samples:int, seqlen:int, seed:int, shuffle:bool):
-    self.samples, self.seqlen = samples, seqlen
-    self.shuffle = shuffle
-    self.rng = np.random.RandomState(seed)
-
-    self.indexed_dataset = BinIdxDataset(base_path)
-
-    # check for cache
-    cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
-    cache_path = base_path.with_name(f"{base_path.name}.{cache_hash}.index_cache")
-    if cache_path.exists():
-      with open(cache_path, "rb") as f:
-        self.doc_idx, self.sample_idx, self.shuffle_idx = pickle.load(f)
-    else:
-      self.doc_idx = self._build_doc_idx()
-      self.sample_idx = self._build_sample_idx()
-      self.shuffle_idx = self._build_shuffle_idx()
-      # save cache
-      with open(cache_path, "wb") as f:
-        pickle.dump((self.doc_idx, self.sample_idx, self.shuffle_idx), f)
-
-  def __getitem__(self, idx):
-    if idx is None:
-      text = self._get(0)
-    else:
-      text = self._get(idx)
-
-    return text
-
-  def _get(self, idx):
-    idx = self.shuffle_idx[idx]
-
-    doc_idx_beg, doc_idx_beg_offset = self.sample_idx[idx]
-    doc_idx_end, doc_idx_end_offset = self.sample_idx[idx + 1]
-
-    doc_ids, sample_parts = [], []
-
-    if doc_idx_beg == doc_idx_end:
-      doc_ids.append(self.doc_idx[doc_idx_beg])
-
-      sample_parts.append(
-          self.indexed_dataset.get(
-            int(self.doc_idx[doc_idx_beg]), offset=int(doc_idx_beg_offset), length=int(doc_idx_end_offset - doc_idx_beg_offset + 1)))
-    else:
-      for i in range(doc_idx_beg, doc_idx_end + 1):
-        doc_ids.append(self.doc_idx[i])
-
-        offset = 0 if i > doc_idx_beg else doc_idx_beg_offset
-        length = None if i < doc_idx_end else int(doc_idx_end_offset + 1)
-        sample_parts.append(self.indexed_dataset.get(int(self.doc_idx[i]), offset=int(offset), length=length))
-
-    # concat all parts
-    text = Tensor.cat(*sample_parts)
-
-    return text
-
-  @functools.cached_property
-  def tokens_per_epoch(self) -> int:
-    return sum(self.indexed_dataset.sizes.tolist())
-
-  @functools.cached_property
-  def num_epochs(self) -> int:
-    # we need enough epochs to cover the requested amount of tokens
-    num_epochs = 1
-    num_tokens = self.tokens_per_epoch
-    while num_tokens < self.samples * self.seqlen:
-      num_epochs += 1
-      num_tokens += self.tokens_per_epoch
-    return num_epochs
-
-  # https://github.com/NVIDIA/Megatron-LM/blob/94bd476bd840c2fd4c3ebfc7448c2af220f4832b/megatron/core/datasets/gpt_dataset.py#L558
-  def _build_doc_idx(self):
-    doc_idx = np.mgrid[:self.num_epochs, :self.indexed_dataset.count][1]
-    doc_idx = doc_idx.reshape(-1)
-    doc_idx = doc_idx.astype(np.int32)
-    if self.shuffle: self.rng.shuffle(doc_idx)
-    return doc_idx
-
-  def _build_sample_idx(self):
-    sample_idx = np.empty((self.samples + 1, 2), dtype=np.int32)
-
-    sample_idx_idx, doc_idx_idx, doc_offset = 0, 0, 0
-    sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
-    sample_idx_idx += 1
-
-    for _ in tqdm(range(1, self.samples + 1)):
-      remaining_seqlen = self.seqlen + 1
-      while remaining_seqlen > 0:
-        doc_idx = int(self.doc_idx[doc_idx_idx])
-        doc_len = self.indexed_dataset.sizes[doc_idx].item() - doc_offset
-        remaining_seqlen -= doc_len
-        if remaining_seqlen <= 0:
-          doc_offset += remaining_seqlen + doc_len - 1
-          remaining_seqlen = 0
-        else:
-          if doc_idx_idx == len(self.doc_idx) - 1:
-            assert sample_idx_idx == self.samples
-            doc_idx = int(self.doc_idx[doc_idx_idx])
-            doc_offset = self.indexed_dataset.sizes[doc_idx].item() - 1
-            break
-          doc_idx_idx += 1
-          doc_offset = 0
-
-      sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
-      sample_idx_idx += 1
-
-    return sample_idx
-
-  def _build_shuffle_idx(self):
-    shuffle_idx = np.arange(self.samples, dtype=np.int32)
-    if self.shuffle: self.rng.shuffle(shuffle_idx)
-    return shuffle_idx
-
-class BlendedGPTDataset:
-  def __init__(self, paths:list[Path], weights:list[float], samples:int, seqlen:int, seed:int, shuffle:bool):
-    self.seed = seed
-
-    # normalize weights
-    total_weight = sum(weights)
-    self.weights = [w / total_weight for w in weights]
-
-    self.samples = samples
-    surplus = 0.005
-    samples_per_blend = [math.ceil(math.ceil(self.samples * w) * (1 + surplus)) for w in self.weights]
-
-    self.datasets = [GPTDataset(path, samples_per_blend[i], seqlen, seed + i, shuffle) for i,path in enumerate(paths)]
-
-  def get(self, idx:int):
-    tokens = self.datasets[0][idx]
-    return tokens
-
-def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
-  if val:
-    dataset = BlendedGPTDataset([
-      base_dir / "validation" / "c4-validationn-91205-samples.en_text_document",
-    ], [
-      1.0
-    ], samples, seqlen, seed, False)
-  else:
-    dataset = BlendedGPTDataset([
-      base_dir / "c4-train.en_6_text_document",
-      base_dir / "c4-train.en_7_text_document",
-    ], [
-      1.0, 1.0
-    ], samples, seqlen, seed, True)
-
-  for b in range(math.ceil(samples / bs)):
-    batch = []
-    for i in range(bs):
-      tokens = dataset.get(b * bs + i)
-      batch.append(tokens)
-    yield Tensor.stack(batch, dim=0)
-
 if __name__ == "__main__":
  def load_unet3d(val):
    assert not val, "validation set is not supported due to different sizes on inputs"
@@ -737,18 +538,6 @@ if __name__ == "__main__":
      for x in batch_load_retinanet(dataset, val, base_dir):
        pbar.update(x[0].shape[0])

-  def load_llama3(val):
-    bs = 24
-    samples = 5760 if val else 1_200_000
-    seqlen = 512
-
-    max_, min_ = 0, math.inf
-    for tokens in tqdm(batch_load_llama3(bs, samples, seqlen, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=5760, val=bool(val)), total=samples//bs):
-      max_ = max(max_, tokens.shape[1])
-      min_ = min(min_, tokens.shape[1])
-    print(f"max seq length: {max_}")
-    print(f"min seq length: {min_}")
-
  load_fn_name = f"load_{getenv('MODEL', 'resnet')}"
  if load_fn_name in globals():
    globals()[load_fn_name](getenv("VAL", 1))
@@ -1,4 +1,4 @@
-import time, math
+import time
 start = time.perf_counter()
 from pathlib import Path
 import numpy as np
@@ -241,34 +241,6 @@ def eval_mrcnn():
  evaluate_predictions_on_coco(bbox_output, iou_type='bbox')
  evaluate_predictions_on_coco(mask_output, iou_type='segm')

-def eval_llama3():
-  from extra.models.llama import Transformer
-  from examples.llama3 import MODEL_PARAMS
-  from tinygrad.helpers import tqdm
-
-  bs = 4
-  sequence_length = 512
-
-  model = Transformer(**(MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}), max_context=sequence_length, jit=False, disable_kv_cache=True)
-
-  @TinyJit
-  def eval_step(model, tokens):
-    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    return loss.flatten()
-
-  from examples.mlperf.dataloader import batch_load_llama3
-  iter = batch_load_llama3(bs, 5760, sequence_length, Path(getenv("BASEDIR", "/raid/datasets/c4/")), True)
-
-  losses = []
-  for tokens in tqdm(iter, total=5760//bs):
-    GlobalCounters.reset()
-    losses += eval_step(model, tokens).tolist()
-    tqdm.write(f"loss: {np.mean(losses)}")
-
-  log_perplexity = Tensor(losses).mean()
-  print(f"Log Perplexity: {log_perplexity.item()}")
-
 if __name__ == "__main__":
  # inference only
  Tensor.training = False
@@ -1290,16 +1290,9 @@ def train_llama3():
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup

  config = {}
-  BS                 = config["BS"]                     = getenv("BS", 16)
+  BS                 = config["BS"]                     = getenv("BS", 4)
  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
-  SEED               = config["SEED"]                   = getenv("SEED", 5760)
-  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
-  TRAIN_ON_VAL       = config["TRAIN_ON_VAL"]           = getenv("TRAIN_ON_VAL", 0)
-  SAMPLES            = config["SAMPLES"]                = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000)
-
-  # LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 FUSE_ARANGE=1 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
-  # trains to 7

  opt_adamw_beta_1 = 0.9
  opt_adamw_beta_2 = 0.95
@@ -1307,6 +1300,7 @@ def train_llama3():
  opt_adamw_weight_decay = 0.1

  opt_gradient_clip_norm = 1.0
+  sequence_length = 8192
  opt_learning_rate_warmup_steps = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
  opt_learning_rate_decay_steps = getenv("DECAY_STEPS", math.ceil(1_200_000 * 1152 / GBS) - opt_learning_rate_warmup_steps)
  opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152)  # NOTE: cannot change for benchmark
@@ -1314,9 +1308,7 @@ def train_llama3():

  # TODO: confirm weights are in bf16
  # vocab_size from the mixtral tokenizer
-  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}
-  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
-  model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
+  model = Transformer(**(MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}), max_context=sequence_length, jit=False, disable_kv_cache=True)

  optim = AdamW(get_parameters(model), lr=0.0,
                b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
@@ -1324,10 +1316,10 @@ def train_llama3():

  @TinyJit
  @Tensor.train()
-  def train_step(model, tokens):
+  def train_step(model, x, y):
    optim.zero_grad()
-    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    logits:Tensor = model(x, start_pos=0, temperature=math.nan)
+    loss = logits.cross_entropy(y)
    loss.backward()

    # L2 norm grad clip
@@ -1348,31 +1340,19 @@ def train_llama3():
    loss.realize(lr)
    return loss, lr

-  if getenv("FAKEDATA", 0):
-    def fake_data():
-      for _ in range(SAMPLES // GBS):
-        yield Tensor.randint(GBS, SEQLEN + 1, low=0, high=32000, dtype=dtypes.int32, device=Device.DEFAULT)
-    iter = fake_data()
-  else:
-    from examples.mlperf.dataloader import batch_load_llama3
-    iter = batch_load_llama3(GBS, SAMPLES, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=bool(TRAIN_ON_VAL))
+  # overfitting this example should give cross_entropy log(BS)
+  fake_input = Tensor([list(range(getenv("SEQLEN", 10)))], dtype="int16").expand(BS, -1)
+  fake_label = Tensor(list(range(BS)), dtype="int16")

-  i = 0
-  for tokens in tqdm(iter, total=SAMPLES//BS):
+  for _ in range(100):
    GlobalCounters.reset()
-    loss, lr = train_step(model, tokens)
-    # above as tqdm.write f-string
-    tqdm.write(f"{loss.item():.4f} loss, {lr.item():.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used")
-    if (fname:=getenv("LOSS_FILE", "")):
-      with open(fname, "a") as f:
-        f.write(f"{i} {loss.item():.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")
-
-    if getenv("CKPT") and (i % 200 == 0 or i == 10):
-      tqdm.write("saving checkpoint")
-      if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
-      fn = f"{ckpt_dir}/{i}.safe"
-      safe_save(get_state_dict(model), fn)
-    i += 1
+    loss, lr = train_step(model, fake_input, fake_label)
+    # BS=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=8B WARMUP_STEPS=2 DECAY_STEPS=300 PYTHONPATH=. AMD=1 MODEL=llama3 python3 examples/mlperf/model_train.py
+    # uses 43% ~= 83GB
+    # 8B bf16 = 16GB. model + grad + optim m and v = 64GB
+    # TODO: this OOM
+    # BS=1 SEQLEN=4000 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=8B WARMUP_STEPS=2 DECAY_STEPS=300 PYTHONPATH=. AMD=1 MODEL=llama3 python3 examples/mlperf/model_train.py
+    print(loss.item(), lr.item(), f"{GlobalCounters.global_mem//10**9=}")

 if __name__ == "__main__":
  multiprocessing.set_start_method('spawn')
@@ -19,9 +19,6 @@ if __name__ == "__main__":
  elif getenv("ASM") == -1:
    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel3_registers.cpp").read_text()
    prgfast = replace(prg, name="kernel3_registers", src=src, global_size=[N//128, N//128, 1], local_size=[256, 1, 1])
-  elif getenv("ASM") == -2:
-    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel4_gmem_df.cpp").read_text()
-    prgfast = replace(prg, name="kernel4_gmem_db", src=src, global_size=[N//128, N//128, 1], local_size=[256, 1, 1])
  else:
    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel5_lds_optim.cpp").read_text()
    prgfast = replace(prg, name="kernel5_lds_optim", src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1])
@@ -32,7 +29,7 @@ if __name__ == "__main__":
  c = Tensor.zeros(N, N).contiguous().realize()

  GlobalCounters.reset()
-  with Context(DEBUG=2):
+  with Context(DEBUG=2, BEAM=4):
    for _ in range(run_count): tc = (a@b).realize()

  GlobalCounters.reset()
@@ -10,8 +10,7 @@ __attribute__((device)) inline void __syncthreads() {
 }

 #define BLOCK_SIZE 256
-extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, BLOCK_SIZE)))
-kernel3_registers(float *a, float *b, float *c)
+extern "C" __attribute__((global)) void kernel3_registers(float *a, float *b, float *c)
 {
  constexpr int N = 4096;
  constexpr float alpha = 1.0;
@@ -81,8 +80,6 @@ kernel3_registers(float *a, float *b, float *c)

  // Iteration over BK blocks.
  for (int kId = 0; kId < N; kId += BK) {
-    __syncthreads();
-
    // We populate the Shared Memory with Ks row and columns
    for (int i = 0; i < nbReadsB; i++) {
      int index_x = BN * blockIdx.x + rBIdx;
@@ -126,6 +123,7 @@ kernel3_registers(float *a, float *b, float *c)
        }
      }
    }
+    __syncthreads();
  }

  for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
@@ -1,172 +0,0 @@
-typedef long unsigned int size_t;
-extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
-extern "C" __attribute__((device, const)) size_t __ockl_get_group_id(unsigned int);
-struct Dim3 { size_t x, y, z; };
-#define __shared__ __attribute__((shared, aligned(16)))
-__attribute__((device)) inline void __syncthreads() {
-  __builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");
-  __builtin_amdgcn_s_barrier();
-  __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");
-}
-
-#define BLOCK_SIZE 256
-extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, BLOCK_SIZE)))
-kernel4_gmem_db(float *a, float *b, float *c)
-{
-  constexpr int N = 4096;
-  constexpr float alpha = 1.0;
-  constexpr float beta = 0.0;
-
-  const Dim3 blockIdx{ __ockl_get_group_id(0), __ockl_get_group_id(1), __ockl_get_group_id(2) };
-  const Dim3 threadIdx{ __ockl_get_local_id(0), __ockl_get_local_id(1), __ockl_get_local_id(2) };
-
-  // Block Tile size
-  constexpr int BN = 128;
-  constexpr int BM = 128;
-  // Number of Row or column we read per batch
-  constexpr int BK = 8;
-
-  // Thread Tile size
-  constexpr int TN = 4;
-  constexpr int TM = 4;
-
-  constexpr int nbWaves = BLOCK_SIZE / 32;
-  // Wave Tile size
-  constexpr int WN = 64;
-  constexpr int WM = BN * BM / nbWaves / WN;
-
-  // Number of wave on X & Y axis in the Block tile
-  constexpr int nbWaveX = BN / WN;
-  constexpr int nbWaveY = BM / WM;
-
-  const int waveIndex = threadIdx.x / 32;
-  const int waveIdx = waveIndex % nbWaveX;
-  const int waveIdy = waveIndex / nbWaveX;
-  const int indexInWave = threadIdx.x % 32;
-
-  // A wave is a block of 8x4 of the output matrix
-  constexpr int nbThreadXPerWave = 8;
-  constexpr int nbThreadYPerWave = 4;
-
-  // Thread coordinates in Wave
-  const int idxInWave = indexInWave % nbThreadXPerWave;
-  const int idyInWave = indexInWave / nbThreadXPerWave;
-
-  constexpr int nbIterWaveN = WN / (nbThreadXPerWave * TN);
-  constexpr int nbIterWaveM = WM / (nbThreadYPerWave * TM);
-
-  // Wave Sub-tile size
-  constexpr int SUBWN = WN / nbIterWaveN;
-  constexpr int SUBWM = WM / nbIterWaveM;
-
-  // Thread mapping to read BKxBN block from A
-  int rAIdx = threadIdx.x % BK;
-  int rAIdy = threadIdx.x / BK;
-  // Thread mapping to read BNxBK block from B
-  int rBIdx = threadIdx.x % BN;
-  int rBIdy = threadIdx.x / BN;
-
-  constexpr int strideReadB = BLOCK_SIZE / BN;
-  constexpr int strideReadA = BLOCK_SIZE / BK;
-  constexpr int nbReadsB = BN * BK / BLOCK_SIZE;
-  constexpr int nbReadsA = BM * BK / BLOCK_SIZE;
-
-  float A_col[nbIterWaveM * TM];
-  float B_row[nbIterWaveN * TN];
-
-  __shared__ float As[BK][BM];
-  __shared__ float Bs[BK][BN];
-
-  float c_regs[TM * nbIterWaveM * TN * nbIterWaveN] = {0.0f};
-
-  for (int i = 0; i < nbReadsB; i++) {
-    int index_x = BN * blockIdx.x + rBIdx;
-    int index_y = rBIdy + i * strideReadB;
-    Bs[index_y % BK][index_x % BN] = b[N * index_y + index_x];
-  }
-
-  for (int i = 0; i < nbReadsA; i++) {
-    int index_x = rAIdx;
-    int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
-    As[(index_x % BK)][(index_y % BM)] = a[N * index_y + index_x];
-  }
-
-  __syncthreads();
-  // Iteration over BK blocks.
-  for (int kId = 0; kId < N; kId += BK) {
-    float regA[nbReadsA];
-    float regB[nbReadsB];
-    if (kId < N - BK) {
-      // We populate the Shared Memory with Ks row and columns
-      for (int i = 0; i < nbReadsB; i++) {
-        int index_x = BN * blockIdx.x + rBIdx;
-        int index_y = rBIdy + i * strideReadB + kId + BK;
-        regB[i] = b[N * index_y + index_x];
-      }
-
-      for (int i = 0; i < nbReadsA; i++) {
-        int index_x = rAIdx + kId + BK;
-        int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
-        regA[i] = a[N * index_y + index_x];
-      }
-    }
-
-    for (int k = 0; k < BK; k++) {
-      // we cache A & B for the entire Wave tile
-      for (int iterWave = 0; iterWave < nbIterWaveN; iterWave++) {
-        for (int i = 0; i < TN; i++) {
-          int index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i;
-          B_row[iterWave * TN + i] = Bs[k][index];
-        }
-      }
-
-      for (int iterWave = 0; iterWave < nbIterWaveM; iterWave++) {
-        for (int i = 0; i < TM; i++) {
-          int index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i;
-          A_col[iterWave * TM + i] = As[k][index];
-        }
-      }
-
-      // we accumulate to C_regs
-      for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
-        for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
-          for (int yt = 0; yt < TM; yt++) {
-            for (int xt = 0; xt < TN; xt++) {
-              const int x = iterWaveN * TN + xt;
-              const int y = iterWaveM * TM + yt;
-              c_regs[y * TN * nbIterWaveN + x] += A_col[y] * B_row[x];
-            }
-          }
-        }
-      }
-    }
-    __syncthreads();
-    if (kId < N - BK) {
-      for (int i = 0; i < nbReadsB; i++) {
-        int index_x = BN * blockIdx.x + rBIdx;
-        int index_y = rBIdy + i * strideReadB + kId + BK;
-        Bs[index_y % BK][index_x % BN] = regB[i]; // row
-      }
-
-      for (int i = 0; i < nbReadsA; i++) {
-        int index_x = rAIdx + kId + BK;
-        int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
-        As[(index_x % BK)][(index_y % BM)] = regA[i];
-      }
-      __syncthreads();
-    }
-  }
-
-  for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
-    for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
-      int xOut = blockIdx.x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave;
-      int yOut = blockIdx.y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave;
-      for (int yt = 0; yt < TM; yt++) {
-        for (int xt = 0; xt < TN; xt++) {
-          int indexC = N * (yOut + yt) + xOut + xt;
-          c[indexC] = beta * c[indexC] + alpha * c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)];
-        }
-      }
-    }
-  }
-}
@@ -26,7 +26,7 @@ kernel5_lds_optim(float *a, float *b, float *c)
  // Number of Row or column we read per batch
  constexpr int BK = 8;

-  // Thread Tile size
+  // Thread Tile size . 4x4
  constexpr int TN = 4;
  constexpr int TM = 4;

@@ -1,26 +1,17 @@
 from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, graph_rewrite, AxisType, PatternMatcher, UPat
+from tinygrad.helpers import prod, unwrap
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.opt.kernel import AxisType
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
+from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops, UOp, GroupOp
+from tinygrad.shape.shapetracker import ShapeTracker, strides_for_shape
+from tinygrad.schedule.kernelize import merge_views
+from tinygrad.shape.view import View
 from tinygrad.dtype import AddrSpace
-from tinygrad.schedule.kernelize import merge_views, view_left
-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.opt.kernel import axis_colors
-
-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
 run_count = 5

-BN = 128
-BM = 128
-BK = 8
-
-TN = 4
-TM = 4
-
-# 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):
@@ -31,305 +22,147 @@ def swizzle_reduceop(src:UOp, r:UOp, view:UOp):
  assert permute == tuple(range(len(permute))), f"reduce axis must already be in order, {permute} isn't"

  # append the reduce shape to each of the views
-  prshape = prod(rshape:=src.shape[-len(r.axis_arg):])
+  reduce_count = len(r.axis_arg)
+  prshape = prod(rshape:=src.shape[-reduce_count:])
  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]
+  nv = [View.create(v.shape[:-reduce_count]+rshape, tuple(x*prshape for x in v.strides[:-reduce_count])+rstrides, v.offset*prshape,
+                    v.mask[:-reduce_count]+tuple((0,s) for s in rshape) if v.mask is not None else None) for v in unwrap(view.st).views]

  # 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)))))
+             (r.arg[0], tuple(range(len(view.shape)-reduce_count, len(view.shape)))))

-pm = PatternMatcher([
+early_view_left = merge_views+PatternMatcher([
+  # view before elementwise and buffer ops
+  (UPat(Ops.VIEW, src=(UPat({*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.BIND, Ops.VALID, Ops.STORE, Ops.LOAD}, name="e"),), name="view"),
+   lambda e,view: e.replace(src=tuple(s.view(view.st) for s in e.src)) if e.tag is None else None),
+  # push a non contiguous ShapeTracker through reduceop
  (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), swizzle_reduceop),
 ])

-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(dtypes.int, N//BM, 0), 2:UOp.range(dtypes.int, 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))
-
-def hl_spec_kernel3():
+def hand_spec():
+  # Block Tile size  . 128x128
+  # Thread Tile size . 4x4
+  # Wave Tile size   . 128x32
+  # A wave is        . 8x4
+  # ────── problem size and tiling params (mirror the C kernel) ───────────────────
+  BK = 8             # depth of K-tile
+  BN = BM = 128      # block-tile (output) sizes
+  # the real thread is 16x8 = 128 regs
+  TM = 4
  nbIterWaveM = 2
-  nbIterWaveN = 2
+  TN = 4
+  nbIterWaveN = 4

-  # 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,)))
+  # ────── shared-memory tile sizes (unchanged) ───────────────────────────────────
+  LDS_A_SZ = BK * BM          # 1024 floats
+  LDS_B_SZ = BK * BN          # 1024 floats

-  # 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)
+  bC = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0)   # output C
+  bA = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1)   # input A
+  bB = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2)   # input B

-  #                     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[:]
+  # TODO: this should not be a string, just a number
+  lAs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(LDS_A_SZ, addrspace=AddrSpace.LOCAL), arg="As")
+  lBs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(LDS_B_SZ, addrspace=AddrSpace.LOCAL), arg="Bs")

-  # 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]
+  s0 = ShapeTracker.from_shape((N, N, N), (N, 0, 1))
+  s1 = ShapeTracker.from_shape((N, N, N), (0, 1, N))
+  s2 = ShapeTracker.from_shape((N, N, 1), (N, 1, 0))

-  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]
+  ls0 = ShapeTracker.from_shape((BM, BK))
+  ls1 = ShapeTracker.from_shape((BN, BK))

-  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)
+  buf_at = [AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.LOCAL, AxisType.LOCAL, AxisType.LOCAL, AxisType.UPCAST, AxisType.UPCAST]
+  buf_bt = [AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.LOCAL, AxisType.LOCAL, AxisType.LOCAL, AxisType.UPCAST, AxisType.UPCAST]
+  axis_types = buf_at + buf_bt + [AxisType.REDUCE, AxisType.UNROLL, AxisType.UNROLL, AxisType.UNROLL]

-  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)
+  # 128 x 128 x 8
+  full_shape = (N//BM, 2, 2, 2, 2, 2, 2, 2, N//BN, 2, 2, 2, 2, 2, 2, 2, N//BK, 2, 2, 2)

-  #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))
+  s0 = s0.reshape(full_shape)
+  s1 = s1.reshape(full_shape)
+  s2 = s2.reshape(full_shape[:-4] + (1,)*4)

-  axis_types = (
-    AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.UPCAST,
-    AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.UPCAST,
-    AxisType.REDUCE, AxisType.REDUCE)
+  ls0 = ls0.reshape((1, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2)).expand(s0.shape)
+  ls1 = ls1.reshape((1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 1, 2, 2, 2)).expand(s1.shape)
+  assert ls0.real_size() == LDS_A_SZ
+  assert ls1.real_size() == LDS_B_SZ

-  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
+  # BK is a loop of 8
+  #  each loop reads 8 in A, 16 in B

-def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
-  BLOCK_SIZE = 128 if kernel5 else 256
+  print(ls0)
+  print(ls1)

-  nbWaves = BLOCK_SIZE // 32
-  WN = 128 if kernel5 else 64
-  WM = BN * BM // nbWaves // WN
+  permaxis = []
+  for axis_order in [AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP, AxisType.UPCAST, AxisType.GROUP_REDUCE, AxisType.REDUCE, AxisType.UNROLL]:
+    permaxis += [i for i,a in enumerate(axis_types) if a == axis_order]
+  axis_types = [axis_types[x] for x in permaxis]
+  s0, s1, s2, ls0, ls1 = [x.permute(tuple(permaxis)) for x in [s0, s1, s2, ls0, ls1]]
+  print(axis_types)

-  nbWaveX = BN // WN
-  nbWaveY = BM // WM
+  lw0, lr0 = ls0, ls0
+  lw1, lr1 = ls1, ls1

-  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
+  # first round of permutes

-  nbThreadXPerWave = 8
-  nbThreadYPerWave = 4
+  permaxis = (0, 1, 19, 18, 17, 12, 11, 10, 5, 4, 3, 2, 6, 7, 8, 9, 16, 13, 14, 15)
+  s0 = s0.permute(permaxis)
+  lw0 = lw0.permute(permaxis)

-  idxInWave = indexInWave % nbThreadXPerWave
-  idyInWave = indexInWave // nbThreadXPerWave
+  permaxis = (0, 1, 15, 14, 9, 8, 7, 6, 13, 19, 18, 17, 5, 4, 3, 2, 16, 12, 11, 10)
+  s1 = s1.permute(permaxis)
+  lw1 = lw1.permute(permaxis)

-  nbIterWaveN = WN // (nbThreadXPerWave * TN)
-  nbIterWaveM = WM // (nbThreadYPerWave * TM)
+  # second round of permutes
+  #permaxis = (0, 1, 12, 11, 5, 4, 3, 2, 10, 6, 7, 8, 9, 13, 14, 15, 16, 17, 18, 19)
+  #lw0 = lw0.permute(permaxis)
+  #lr0 = lr0.permute(permaxis)

-  SUBWN = WN // nbIterWaveN
-  SUBWM = WM // nbIterWaveM
+  from tinygrad.opt.kernel import axis_colors, colored
+  print('_'.join([colored(f"{s}({st})", axis_colors[x]) for s,st,x in zip(s0.shape, s0.views[0].strides, axis_types)]))
+  print('_'.join([colored(f"{s}({st})", axis_colors[x]) for s,st,x in zip(s1.shape, s1.views[0].strides, axis_types)]))
+  print('_'.join([colored(f"{s}({st})", axis_colors[x]) for s,st,x in zip(s2.shape, s2.views[0].strides, axis_types)]))
+  print("lw")
+  print('_'.join([colored(f"{s}({st})", axis_colors[x]) for s,st,x in zip(lw0.shape, lw0.views[0].strides, axis_types)]))
+  print('_'.join([colored(f"{s}({st})", axis_colors[x]) for s,st,x in zip(lw1.shape, lw1.views[0].strides, axis_types)]))
+  print("lr")
+  print('_'.join([colored(f"{s}({st})", axis_colors[x]) for s,st,x in zip(lr0.shape, lr0.views[0].strides, axis_types)]))
+  print('_'.join([colored(f"{s}({st})", axis_colors[x]) for s,st,x in zip(lr1.shape, lr1.views[0].strides, axis_types)]))

-  # 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
+  # loads and stores
+  bs0 = bA.view(s0).load()
+  bs1 = bB.view(s1).load()
+  bs0 = lAs.view(lr0).load(lAs.view(lw0).store(bs0))
+  bs1 = lBs.view(lr1).load(lBs.view(lw1).store(bs1))

-  strideReadB = BLOCK_SIZE // BN
-  strideReadA = BLOCK_SIZE // BK
-  nbReadsB = BN * BK // BLOCK_SIZE
-  nbReadsA = BM * BK // BLOCK_SIZE
+  mat = (bs0 * bs1).r(Ops.ADD, tuple([i for i,a in enumerate(axis_types) if a in (AxisType.REDUCE, AxisType.UNROLL)]), permute=False)
+  st = bC.view(s2).store(mat)

-  blockIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx0", N//BN))
-  blockIdx_y = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx1", N//BM))
+  ast = st.sink(arg=KernelInfo(axis_types=tuple(axis_types), name="tinygemm"))
+  ast = graph_rewrite(ast, merge_views)
+  prg = get_program(ast, Device.default.renderer)
+  print(prg.src)
+  return prg

-  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(dtypes.int, 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(dtypes.int, 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(dtypes.int, 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(dtypes.int, 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(dtypes.int, 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(dtypes.int, 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(dtypes.int, BK, first_range+0)
-
-      # load from locals into registers
-      iterWave = UOp.range(dtypes.int, nbIterWaveN, first_range+1)
-      i = UOp.range(dtypes.int, 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(dtypes.int, nbIterWaveM, first_range+3)
-      i = UOp.range(dtypes.int, 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(dtypes.int, nbIterWaveM, first_range+5)
-      yt = UOp.range(dtypes.int, TM, first_range+6)
-      iterWaveN = UOp.range(dtypes.int, nbIterWaveN, first_range+7)
-      xt = UOp.range(dtypes.int, 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(dtypes.int, 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(dtypes.int, 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(dtypes.int, N//BK, 0)
-    kId = kId_range*BK
-
-    # load from globals into locals
-    i = UOp.range(dtypes.int, 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(dtypes.int, 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(dtypes.int, BK, 3)
-
-    # load from locals into registers
-    iterWave = UOp.range(dtypes.int, nbIterWaveN, 4)
-    i = UOp.range(dtypes.int, 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(dtypes.int, nbIterWaveM, 6)
-    i = UOp.range(dtypes.int, 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(dtypes.int, nbIterWaveM, 8)
-    yt = UOp.range(dtypes.int, TM, 9)
-    iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 10)
-    xt = UOp.range(dtypes.int, 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(dtypes.int, nbIterWaveM, 1000)
-  yt = UOp.range(dtypes.int, TM, 1001)
-  iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 1002)
-  xt = UOp.range(dtypes.int, 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__":
-  HL = getenv("HL")
-  if HL == 2: hprg = top_spec_kernel3()
-  elif HL == 1: hprg = hl_spec_kernel3()
-  else: hprg = hand_spec_kernel3()
-  prg = get_program(hprg, Device.default.renderer)
-  print(prg.src)
-  if getenv("SRC"): exit(0)
-  hrunner = CompiledRunner(prg)
+  hprg = hand_spec()
+  hrunner = CompiledRunner(hprg)

  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):
+  with Context(DEBUG=2, BEAM=4):
    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)
+  ei = ExecItem(hrunner, [hc.uop.buffer, a.uop.buffer, b.uop.buffer])
  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!")
+  assert err < 1e-06
@@ -1,122 +0,0 @@
-#!/usr/bin/env python3
-
-from tinygrad.runtime.support.system import System
-import argparse, glob, os, re, time, subprocess, sys
-
-def scan_devs_based_on_lock(prefix:str, args) -> list[str]:
-  target_dev = args.pci_bus if 'pci_bus' in args.__dir__() else ""
-
-  devs = []
-  for dev in glob.glob(f'/tmp/{prefix}_*.lock'):
-    dev_id = dev[8:-5]
-    if os.path.exists(f"/sys/bus/pci/devices/{dev_id}") and dev_id.startswith(target_dev): devs.append(dev_id)
-  return devs
-
-def _do_reset_device(pci_bus): System.pci_reset(pci_bus)
-def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")
-
-def cmd_remove_module(args):
-  modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia"] if args.backend == "nv" else ["amdgpu"]
-  to_unload = [m for m in modules if _is_module_loaded(m)]
-  if not to_unload: print("Kernel modules are not loaded")
-  else:
-    print("Removing kernel modules:", ", ".join(to_unload))
-    try: subprocess.run(["sudo", "modprobe", "-r", *to_unload], check=True)
-    except subprocess.CalledProcessError as e:
-      print("Failed to unload all modules — they may be in use.", file=sys.stderr)
-      sys.exit(e.returncode)
-
-def cmd_insert_module(args):
-  cmd_remove_module(args)
-  cmd_reset_devices(args)
-
-  module = "nvidia" if args.backend == "nv" else "amdgpu"
-  if _is_module_loaded(module):
-    print(f"{module} kernel module already loaded")
-    return
-
-  print(f"Inserting kernel module: {module}")
-  if args.backend == "nv":
-    subprocess.run(["nvidia-smi"], check=True)
-  elif args.backend == "amd":
-    subprocess.run(["sudo", "modprobe", "amdgpu"], check=True)
-
-def cmd_reset_devices(args):
-  devs = scan_devs_based_on_lock({"amd":"am", "nv":"nv"}[args.backend], args)
-
-  for dev in devs:
-    print(f"Resetting device {dev}")
-    if args.backend != "amd": _do_reset_device(dev)
-    time.sleep(0.2)
-
-def cmd_show_pids(args):
-  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
-
-  for dev in devs:
-    try:
-      pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
-      print(f"{dev}: {pid}")
-    except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
-
-def cmd_kill_pids(args):
-  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
-
-  for dev in devs:
-    try:
-      pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
-      print(f"{dev}: {pid}")
-    except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
-
-def cmd_kill_pids(args):
-  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
-
-  for dev in devs:
-    for i in range(128):
-      if i > 0: time.sleep(0.2)
-
-      try:
-        try: pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
-        except subprocess.CalledProcessError: break
-
-        print(f"Killing process {pid} (which uses {dev})")
-        subprocess.run(['sudo', 'kill', '-9', pid], check=True)
-      except subprocess.CalledProcessError as e:
-        print(f"Failed to kill process for device {dev}: {e}", file=sys.stderr)
-
-def add_common_commands(parent_subparsers):
-  p_insmod = parent_subparsers.add_parser("insmod", help="Insert a kernel module")
-  p_insmod.set_defaults(func=cmd_insert_module)
-
-  p_rmmod = parent_subparsers.add_parser("rmmod", help="Remove a kernel module")
-  p_rmmod.set_defaults(func=cmd_remove_module)
-
-  p_reset = parent_subparsers.add_parser("reset", help="Reset a device")
-  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device to reset")
-  p_reset.set_defaults(func=cmd_reset_devices)
-
-  p_reset = parent_subparsers.add_parser("pids", help="Show pids of processes using the device")
-  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
-  p_reset.set_defaults(func=cmd_show_pids)
-
-  p_reset = parent_subparsers.add_parser("kill_pids", help="Kill pids of processes using the device")
-  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
-  p_reset.set_defaults(func=cmd_kill_pids)
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser()
-  backend_subparsers = parser.add_subparsers(dest="backend", required=True, metavar="{nv,amd}", help="Hardware backend to target")
-
-  nv_parser = backend_subparsers.add_parser("nv", help="NVIDIA GPUs")
-  nv_commands = nv_parser.add_subparsers(dest="command", required=True)
-  add_common_commands(nv_commands)
-
-  amd_parser = backend_subparsers.add_parser("amd", help="AMD GPUs")
-  amd_commands = amd_parser.add_subparsers(dest="command", required=True)
-  add_common_commands(amd_commands)
-
-  args = parser.parse_args()
-  if args.command is None:
-    parser.print_help(sys.stderr)
-    sys.exit(1)
-
-  args.func(args)
@@ -99,9 +99,7 @@ class FeedForward:
    self.w3 = linear(dim, hidden_dim, bias=False) # the gate in Gated Linear Unit

  def __call__(self, x:Tensor) -> Tensor:
-    w1 = self.w1(x).silu()
-    w3 = self.w3(x.contiguous_backward())  # this fixes a strange fusion that makes tensor cores miss
-    return self.w2(w1 * w3)
+    return self.w2(self.w1(x).silu() * self.w3(x)) # SwiGLU [arxiv/2002.05202, eq (5)]

 class TransformerBlock:
  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int, norm_eps:float, max_context:int, linear=nn.Linear,
@@ -113,7 +111,7 @@ class TransformerBlock:

  def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]):
    h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)
-    return (h + self.feed_forward(self.ffn_norm(h))).contiguous().contiguous_backward()
+    return (h + self.feed_forward(self.ffn_norm(h))).contiguous()

 # standard openai sampling
 def sample(logits: Tensor, temp: float, k: int, p: float, af: float, ap: float):
@@ -187,10 +185,10 @@ class Transformer:

    mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1) if seqlen > 1 else None
    for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask)
-    logits = self.output(self.norm(h)).float()
+    logits = self.output(self.norm(h)).float()[:, -1, :]
    if math.isnan(temperature): return logits

-    return sample(logits[:, -1, :].flatten(), temperature, top_k, top_p, alpha_f, alpha_p)
+    return sample(logits.flatten(), temperature, top_k, top_p, alpha_f, alpha_p)

  def __call__(self, tokens:Tensor, start_pos:int, temperature:float=0.0, top_k:int=0, top_p:float=0.8, alpha_f:float=0.0, alpha_p:float=0.0):
    # TODO: better way to handle the first call v.s. the rest?
@@ -0,0 +1,2 @@
+GPU="$1"
+echo 1 | sudo tee /sys/bus/pci/devices/$GPU/reset 2>/dev/null
@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+
+from tinygrad.runtime.support.system import System
+import argparse, glob, os, re, time, subprocess, sys
+
+def scan_devs_based_on_lock(prefix:str) -> list[str]:
+  devs = []
+  for dev in glob.glob(f'/tmp/{prefix}_*.lock'):
+    dev_id = dev[8:-5]
+    if os.path.exists(f"/sys/bus/pci/devices/{dev_id}"): devs.append(dev_id)
+  return devs
+
+def _do_reset_device(pci_bus): System.pci_reset(pci_bus)
+def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")
+
+def cmd_remove_module(args):
+  to_unload = [m for m in ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia"] if _is_module_loaded(m)]
+  if not to_unload:
+    print("NVIDIA kernel modules are not loaded")
+  else:
+    print("Removing NVIDIA kernel modules:", ", ".join(to_unload))
+    try: subprocess.run(["sudo", "modprobe", "-r", *to_unload], check=True)
+    except subprocess.CalledProcessError as e:
+      print("Failed to unload all modules — they may be in use.", file=sys.stderr)
+      sys.exit(e.returncode)
+
+def cmd_insert_module(args):
+  cmd_remove_module(args)
+  cmd_reset_devices(args)
+
+  if not os.path.exists("/sys/module/nvidia"):
+    print("Inserting nvidia kernel module")
+    subprocess.run(["nvidia-smi"], check=True)
+  else: print("Nvidia kernel module already loaded")
+
+def cmd_reset_devices(args):
+  devs = scan_devs_based_on_lock("nv")
+  dev_to_reset = args.pci_bus if 'pci_bus' in args.__dir__() else ""
+
+  for dev in devs:
+    if dev.startswith(dev_to_reset):
+      print(f"Resetting device {dev}")
+      _do_reset_device(dev)
+      time.sleep(0.2)
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  subparsers = parser.add_subparsers(required=True, dest="cmd")
+
+  parser_insmod = subparsers.add_parser('insmod', help='Insert a nvidia kernel module')
+  parser_insmod.set_defaults(func=cmd_insert_module)
+
+  parser_rmmod = subparsers.add_parser('rmmod', help='Remove a nvidia kernel module')
+  parser_rmmod.set_defaults(func=cmd_remove_module)
+
+  parser_reset = subparsers.add_parser('reset', help='Reset a nvidia device')
+  parser_reset.add_argument('--pci_bus', type=str, default="", help='PCI bus ID of the device to reset')
+  parser_reset.set_defaults(func=cmd_reset_devices)
+
+  args = parser.parse_args()
+  if args.cmd is None:
+    parser.print_help(sys.stderr)
+    sys.exit(1)
+
+  args.func(args)
@@ -1,4 +1,3 @@
-# mypy: disable-error-code="misc, list-item, assignment, operator, index, arg-type"
 from types import SimpleNamespace
 from typing import Any, Sequence, cast, Literal, Callable, get_args, NamedTuple
 import dataclasses, functools, io, math, types, warnings, pathlib, sys, enum
@@ -84,7 +83,7 @@ class OnnxValue:
  is_optional: bool
  is_sequence: bool

-class Domain(enum.Enum):
+class Domain(enum.StrEnum):
  ONNX = "ai.onnx"
  ONNX_ML = "ai.onnx.ml"
  AI_ONNX_TRAINING = "ai.onnx.training"
@@ -798,11 +797,12 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT

  def Gather(x:Tensor, indices:Tensor, axis:int=0):
    if indices.numel() < 9: # NOTE lessor kernels for smaller indices but kernel number increases depending on size of indices
-      ret_shape = x.shape[:axis] + indices.shape + x.shape[axis+1:]
+      x_sh = list(x.shape)
+      ret_shape = x_sh[:axis] + list(indices.shape) + x_sh[axis+1:]
      if indices.ndim > 1: indices = indices.flatten()
-      index_consts = [_cached_to_python_const(indices)] if indices.shape == () else _cached_to_python_const(indices)
-      index_consts = [x.shape[axis]+i if i<0 else i for i in index_consts]
-      args = [[(0,x) if j != axis else (i,i+1) for j, x in enumerate(x.shape)] for i in index_consts]
+      indices = [_cached_to_python_const(indices)] if indices.shape == () else _cached_to_python_const(indices)
+      indices = [x_sh[axis]+x if x<0 else x for x in indices]
+      args = [[(0,x) if j != axis else (i,i+1) for j, x in enumerate(x_sh)] for i in indices] # type: ignore
      return x.shrink(arg=tuple(args[0])).cat(*[x.shrink(arg=tuple(arg)) for arg in args[1:]], dim=axis).reshape(ret_shape)
    # NOTE faster gather, fixed number of kernels, but exceeds limited kernels for openpilot
    return x[tuple([slice(None) if i != axis else indices for i in range(x.ndim)])]
@@ -128,12 +128,6 @@ def _linalg_eigh(self, UPLO: str = 'U'):
  w, v = torch.linalg.eigh(self.cpu(), UPLO=UPLO)
  return w.tiny(), v.tiny()

-@torch.library.impl("aten::_linalg_det", "privateuseone")
-# TODO: move to tinygrad
-def _linalg_det(self: torch.Tensor):
-  result = aten._linalg_det(self.cpu())
-  return result[0].tiny(), result[1].tiny(), result[2].tiny()
-
 def upsample_backward(grad_out, output_size, input_size, *args, f=None): return f(grad_out.cpu(), output_size, input_size, *args).tiny()

 for i in [
@@ -198,11 +198,6 @@ class TestTorchBackend(unittest.TestCase):
    recon = (v @ torch.diag(w) @ v.T).cpu().numpy()
    np.testing.assert_allclose(recon, a.cpu().numpy(), atol=1e-6)

-  def test_linalg_det(self):
-    a = torch.diag(torch.tensor([1,2,3,4,5], dtype = torch.float32, device=device))
-    b = torch.linalg.det(a)
-    np.testing.assert_equal(b.cpu().numpy(), 120.0)
-
  def test_scalar_assign(self):
    a = torch.tensor([1, 2, 3], device=device)
    a[1] = 4
@@ -1,8 +1,7 @@
 import pathlib
 from tinygrad import Tensor, Device, Context
-from tinygrad.helpers import getenv

 if __name__ == "__main__":
  with Context(DEBUG=2):
-    disk_llama = Tensor(pathlib.Path(getenv("TESTFILE", "/raid/weights/LLaMA-3/8B/consolidated.00.pth")))
+    disk_llama = Tensor(pathlib.Path("/raid/weights/LLaMA-3/8B/consolidated.00.pth"))
    device_llama = disk_llama.to(Device.DEFAULT).realize()
@@ -21,12 +21,12 @@ class FakeAM:
  def __init__(self):
    self.is_booting, self.smi_dev = True, False
    self.pcidev = FakePCIDev()
-    self.vram_size = (512 << 20)
+    self.vram_size = (4 << 30)
    self.vram_mv = memoryview(bytearray(self.vram_size))
    self.vram = MMIOInterface(mv_address(self.vram_mv), self.vram_mv.nbytes)
    self.gmc = FakeGMC(self)
    self.mm = AMMemoryManager(self, self.vram_size, boot_size=(32 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39], va_bits=48,
-      first_lv=am.AMDGPU_VM_PDB2, va_base=AMMemoryManager.va_allocator.base,
+      first_lv=am.AMDGPU_VM_PDB1, va_base=AMMemoryManager.va_allocator.base,
      palloc_ranges=[(1 << (i + 12), 0x1000) for i in range(9 * (3 - am.AMDGPU_VM_PDB2), -1, -1)])
    self.is_booting = False
    self.ip_ver = {am.GC_HWIP: (11, 0, 0)}
@@ -56,8 +56,6 @@ def helper_read_entry_components(entry_val):
          "read": (entry_val >> 5) & 0x1, "write": (entry_val >> 6) & 0x1, "exec": (entry_val >> 4) & 0x1,
          "mtype": (entry_val >> 48) & 0x7, "T": (entry_val >> 51) & 0x1, "L": (entry_val >> 55) & 0x1, "F": (entry_val >> 56) & 0x1}

-def helper_va(va:int): return va + AMMemoryManager.va_allocator.base
-
 class TestAMPageTable(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
@@ -68,9 +66,10 @@ class TestAMPageTable(unittest.TestCase):

    for va,sz in [(0x10000, 0x3000), (0x11000, 0x300000), (0x10000, 0x2000), (0x11000, 0x5000),
                  (0x2000000, 0x2000), (0x4000000, 0x4000000), (0x38000, 0x303000), (0x8000, 0x1000)]:
-      mm.map_range(vaddr=helper_va(va), size=sz, paddrs=[(va, sz)])
+      exteranl_va = va + AMMemoryManager.va_allocator.base
+      mm.map_range(vaddr=exteranl_va, size=sz, paddrs=[(va, sz)])

-      ctx = PageTableTraverseContext(self.d[0], mm.root_page_table, helper_va(va))
+      ctx = PageTableTraverseContext(self.d[0], mm.root_page_table, exteranl_va)
      results = list(ctx.next(sz))

      total_covered = 0
@@ -87,7 +86,7 @@ class TestAMPageTable(unittest.TestCase):
          assert pte['paddr'] == va + _offset + i * _pte_covers, f"Expected paddr {pte['paddr']:#x} to be {va + _offset + i * _pte_covers:#x}"
          assert pte['valid'] == 1

-      mm.unmap_range(helper_va(va), sz)
+      mm.unmap_range(va, sz)

      for tup in results:
        _offset, _pt, _pte_idx, _n_ptes, _pte_covers = tup
@@ -100,16 +99,18 @@ class TestAMPageTable(unittest.TestCase):
    mm0 = self.d[0].mm

    for (va1,sz1),(va2,sz2) in [((0x10000, (0x1000)), (0x11000, (2 << 20)))]:
-      mm0.map_range(vaddr=helper_va(va1), size=sz1, paddrs=[(va1, sz1)])
-      mm0.map_range(vaddr=helper_va(va2), size=sz2, paddrs=[(va2, sz2)])
-      mm0.unmap_range(helper_va(va2), sz2)
-      mm0.unmap_range(helper_va(va1), sz1)
+      exteranl_va1 = va1 + AMMemoryManager.va_allocator.base
+      exteranl_va2 = va2 + AMMemoryManager.va_allocator.base
+      mm0.map_range(vaddr=exteranl_va1, size=sz1, paddrs=[(va1, sz1)])
+      mm0.map_range(vaddr=exteranl_va2, size=sz2, paddrs=[(va2, sz2)])
+      mm0.unmap_range(va2, sz2)
+      mm0.unmap_range(va1, sz1)

  def test_double_map(self):
    mm0 = self.d[0].mm

    for va,sz in [(0x10000, 0x3000), (0x1000000, 0x1000000), (0x12000, 0x4000)]:
-      exteranl_va = helper_va(va)
+      exteranl_va = va + AMMemoryManager.va_allocator.base
      mm0.map_range(vaddr=exteranl_va, size=sz, paddrs=[(va, sz)])

      with self.assertRaises(AssertionError):
@@ -143,36 +144,36 @@ class TestAMPageTable(unittest.TestCase):
    mm0 = self.d[0].mm

    with self.assertRaises(AssertionError):
-      mm0.unmap_range(helper_va(0x10000), 0x3000)
+      mm0.unmap_range(0x10000, 0x3000)

-    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)])
-    mm0.unmap_range(helper_va(0x10000), 0x3000)
+    mm0.map_range(0x10000, 0x3000, paddrs=[(0x10000, 0x3000)])
+    mm0.unmap_range(0x10000, 0x3000)

    with self.assertRaises(AssertionError):
-      mm0.unmap_range(helper_va(0x10000), 0x3000)
+      mm0.unmap_range(0x10000, 0x3000)

-    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)])
-    mm0.unmap_range(helper_va(0x10000), 0x3000)
+    mm0.map_range(0x10000, 0x3000, paddrs=[(0x10000, 0x3000)])
+    mm0.unmap_range(0x10000, 0x3000)

    with self.assertRaises(AssertionError):
-      mm0.unmap_range(helper_va(0x10000), 0x3000)
+      mm0.unmap_range(0x10000, 0x3000)

  def test_free_pt(self):
    mm0 = self.d[0].mm

    # offset from start
    for off in [0, 0x3000, 0x10000]:
-      mm0.map_range(helper_va(0x1000000) + off, (2 << 20)  - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000))
-      mm0.unmap_range(helper_va(0x1000000) + off, (2 << 20) - off)
-      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)])
-      mm0.unmap_range(helper_va(0x1000000), 2 << 20)
+      mm0.map_range(0x1000000 + off, (2 << 20)  - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000))
+      mm0.unmap_range(0x1000000 + off, (2 << 20) - off)
+      mm0.map_range(0x1000000, 2 << 20, paddrs=[(0x10000, 2 << 20)])
+      mm0.unmap_range(0x1000000, 2 << 20)

    # offset from end
    for off in [0x1000, 0x20000]:
-      mm0.map_range(helper_va(0x1000000), (2 << 20) - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000))
-      mm0.unmap_range(helper_va(0x1000000), (2 << 20) - off)
-      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)])
-      mm0.unmap_range(helper_va(0x1000000), 2 << 20)
+      mm0.map_range(0x1000000, (2 << 20) - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000))
+      mm0.unmap_range(0x1000000, (2 << 20) - off)
+      mm0.map_range(0x1000000, 2 << 20, paddrs=[(0x10000, 2 << 20)])
+      mm0.unmap_range(0x1000000, 2 << 20)

  def test_frag_size(self):
    mm0 = self.d[0].mm
@@ -184,11 +184,6 @@ backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_mapping_cpu') # bad d
 backend_test.exclude('test_scatternd_min_cpu') # min not yet supported
 backend_test.exclude('test_scatternd_max_cpu') # max not yet supported

-# regression from removing StrEnum in Domain
-backend_test.exclude('test_adam_cpu')
-backend_test.exclude('test_gradient_of_add_and_mul_cpu')
-backend_test.exclude('test_gradient_of_add_cpu')
-
 if Device.DEFAULT in ['GPU', 'METAL']:
  backend_test.exclude('test_resize_upsample_sizes_nearest_axes_2_3_cpu')
  backend_test.exclude('test_resize_upsample_sizes_nearest_axes_3_2_cpu')
@@ -251,7 +251,6 @@ class TestTrainingOnnxOps(TestOnnxOps):
        outputs = ["X_out", "V_out"]
        self._validate_training("Momentum", onnx_fxn, inputs, attributes, outputs)

-  @unittest.expectedFailure  # TODO: regression from removing StrEnum in Domain
  def test_adam_t_greater_than_zero(self):
    from onnx.backend.test.case.node.adam import apply_adam
    for t in [1, 3, 100]:
@@ -2,7 +2,7 @@ import random
 from z3 import Int, Solver, sat
 from tinygrad import dtypes, Device
 from tinygrad.uop.ops import UOp, Ops, UPat, graph_rewrite, PatternMatcher
-from tinygrad.codegen.optional import fast_idiv
+from tinygrad.codegen.devectorizer import fast_idiv
 random.seed(42)

 z3_renderer = PatternMatcher([
@@ -1,9 +1,10 @@
-import unittest, numpy as np
-from tinygrad import Tensor, Device, TinyJit
+import unittest
+from tinygrad import Tensor
+from tinygrad import Device
 from tinygrad.helpers import Timing, CI, OSX
 import multiprocessing.shared_memory as shared_memory

-N = 256 if CI else 4096
+N = 4096
 class TestCopySpeed(unittest.TestCase):
  @classmethod
  def setUpClass(cls): Device[Device.DEFAULT].synchronize()
@@ -48,32 +49,6 @@ class TestCopySpeed(unittest.TestCase):
      with Timing("sync:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
        t.to('CPU').realize()

-  def testCopyDefaulttoCPUJit(self):
-    if Device.DEFAULT == "CPU": return unittest.skip("CPU to CPU copy is a no-op")
-
-    @TinyJit
-    def _do_copy(t): return t.to('CPU').realize()
-
-    t = Tensor.randn(N, N, 4).contiguous().realize()
-    for _ in range(5):
-      with Timing("sync:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
-        x = _do_copy(t)
-        Device[Device.DEFAULT].synchronize()
-      np.testing.assert_equal(t.numpy(), x.numpy())
-
-  def testCopytoCPUtoDefaultJit(self):
-    if Device.DEFAULT == "CPU": return unittest.skip("CPU to CPU copy is a no-op")
-
-    @TinyJit
-    def _do_copy(x): return t.to(Device.DEFAULT).realize()
-
-    for _ in range(5):
-      t = Tensor.randn(N, N, 4, device="CPU").contiguous().realize()
-      with Timing("sync:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
-        x = _do_copy(t)
-        Device[Device.DEFAULT].synchronize()
-      np.testing.assert_equal(t.numpy(), x.numpy())
-
  @unittest.skipIf(CI, "CI doesn't have 6 GPUs")
  @unittest.skipIf(Device.DEFAULT != "GPU", "only test this on GPU")
  def testCopyCPUto6GPUs(self):
@@ -1,32 +0,0 @@
-import unittest
-from tinygrad import dtypes, Device, Tensor, Context
-from tinygrad.dtype import AddrSpace
-from tinygrad.helpers import getenv
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.engine.realize import get_program, ExecItem, CompiledRunner
-
-class TestDefineReg(unittest.TestCase):
-  def test_simple(self, at=AxisType.UPCAST):
-    N = 16
-    bout = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0).view(ShapeTracker.from_shape((N,N)))
-    a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1).view(ShapeTracker.from_shape((N,N)))
-    a_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(N, AddrSpace.REG), arg=0).view(ShapeTracker.from_shape((N,N), (0,1)))
-
-    out = a_col.load(a_col.store(a.load()))
-    sink = bout.store(out).sink(arg=KernelInfo(name="regcopy", axis_types=(AxisType.LOOP, at)))
-    prg = get_program(sink, Device.default.renderer)
-
-    with Context(DEBUG=0):
-      a = Tensor.randn(N, N).realize()
-      b = Tensor.empty(N, N).realize()
-    hrunner = CompiledRunner(prg)
-    ExecItem(hrunner, [b.uop.buffer, a.uop.buffer]).run(wait=True)
-    with Context(DEBUG=0):
-      self.assertEqual((b-a).mean().item(), 0.0)
-
-  @unittest.skipIf(getenv("PTX"), "ptx needs regs to be unrolled")
-  def test_simple_loop(self): self.test_simple(AxisType.LOOP)
-
-if __name__ == '__main__':
-  unittest.main()
@@ -1,6 +1,6 @@
 import unittest, ctypes, struct, os, random, numpy as np
 from tinygrad import Device, Tensor, dtypes
-from tinygrad.helpers import getenv, CI, mv_address, DEBUG
+from tinygrad.helpers import getenv, CI, mv_address
 from tinygrad.device import Buffer, BufferSpec
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
 from tinygrad.runtime.autogen import libc
@@ -513,31 +513,6 @@ class TestHCQ(unittest.TestCase):

      assert buf2.as_buffer()[0] == i

-  def test_map_cpu_buffer_to_device(self):
-    if Device[Device.DEFAULT].hw_copy_queue_t is None: self.skipTest("skip device without copy queue")
-
-    sz = 0x2000
-    cpu_buffer = Buffer("CPU", sz, dtypes.uint8, options=BufferSpec(cpu_access=True)).ensure_allocated()
-    cpu_buffer._buf.cpu_view().view(fmt='B')[:] = bytes([x & 0xff for x in range(sz)])
-
-    for devid in range(6):
-      if DEBUG >= 2: print(f"Testing map to device {Device.DEFAULT}:{devid}")
-
-      try: d = Device[f"{Device.DEFAULT}:{devid}"]
-      except Exception: break
-
-      local_buf = Buffer(f"{Device.DEFAULT}:{devid}", sz, dtypes.uint8, options=BufferSpec(cpu_access=True)).ensure_allocated()
-
-      d.allocator.map(cpu_buffer._buf)
-
-      d.hw_copy_queue_t().wait(d.timeline_signal, d.timeline_value - 1) \
-                         .copy(local_buf._buf.va_addr, cpu_buffer._buf.va_addr, sz) \
-                         .signal(d.timeline_signal, d.timeline_value).submit(d)
-      d.timeline_signal.wait(d.timeline_value)
-      d.timeline_value += 1
-
-      np.testing.assert_equal(cpu_buffer.numpy(), local_buf.numpy(), "failed")
-
  @unittest.skipUnless(MOCKGPU, "Emulate this on MOCKGPU to check the path in CI")
  def test_on_device_hang(self):
    if not hasattr(self.d0, 'on_device_hang'): self.skipTest("device does not have on_device_hang")
@@ -114,6 +114,27 @@ class TestLinearizer(unittest.TestCase):
        if skip and i in skip: continue
        assert ranges[i-1] != u, f"multireduce nested the ranges! {ranges[i-1], {u}}"

+  @unittest.expectedFailure
+  def test_const_alu_indexing(self):
+    st = ShapeTracker.from_shape((4,)).to_uop()
+    load = UOp.load(UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=1, src=()), st, dtype=dtypes.float)
+    op = load+UOp.const(dtypes.float, 1.0)*UOp.const(dtypes.float, -1)
+    store = UOp.store(UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=0, src=()), st, op)
+    Tensor.manual_seed(0)
+    x = Tensor.randn(4,).realize()
+    helper_linearizer_ast(store.sink(), [x], wanna_output=[x.numpy()+1*-1], opts=[])
+
+  # shapeless CONST in AST is not supported
+  @unittest.expectedFailure
+  def test_const_alu_indexing_one_const_fine(self):
+    st = ShapeTracker.from_shape((4,)).to_uop()
+    load = UOp.load(UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=1, src=()), st, dtype=dtypes.float)
+    op = load+UOp.const(dtypes.float, 1.0)
+    store = UOp.store(UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=0, src=()), st, op)
+    Tensor.manual_seed(0)
+    x = Tensor.randn(4,).realize()
+    helper_linearizer_ast(store.sink(), [x], wanna_output=[x.numpy()+1], opts=[])
+
  @unittest.skipIf(CI and Device.DEFAULT in {"PTX", "AMD", "NV"}, "very slow")
  def test_indexing_multireduce(self):
    dataset = Tensor.rand(16384, 256).realize()
@@ -246,7 +267,7 @@ class TestLinearizer(unittest.TestCase):
    stores = [u for u in uops if u.op is Ops.STORE]
    assert len(accs) == 0  # it's removed now
    assert len(stores) == 1
-    assert stores[0].src[1].dtype == dtypes.float.vec(4)
+    assert stores[0].src[-1].dtype == dtypes.float.vec(4)

  # NOTE: can reenable, it does work. it just makes BEAM slow
  @unittest.expectedFailure
@@ -270,13 +291,13 @@ class TestLinearizer(unittest.TestCase):
    realized_ast = realized_ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
    program = get_program(realized_ast, Device[Device.DEFAULT].renderer)

-    stores = [u for u in program.uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
+    stores = [u for u in program.uops if u.op is Ops.STORE and u.dtype.addrspace != AddrSpace.REG]

    # the first store is to lds and can be upcasted
-    assert stores[0].src[1].dtype == dtypes.float.vec(4)
+    assert stores[0].src[-1].dtype == dtypes.float.vec(4)
    assert any(x.op is Ops.DEFINE_LOCAL for x in stores[0].toposort())
    # the second store is to gds with no upcasts
-    assert stores[1].src[1].dtype == dtypes.float
+    assert stores[1].src[-1].dtype == dtypes.float
    assert any(x.op is Ops.DEFINE_GLOBAL for x in stores[1].toposort())

  def test_zero_fold(self):
@@ -612,7 +633,6 @@ class TestLinearizer(unittest.TestCase):
    helper(Tensor.arange(255), max_ops=2)

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
-  @unittest.skipIf(getenv("PTX"), "broken on ptx for some reason")
  def test_grouped_store_phis(self):
    """
    float4 acc0 = float4(0.0,0.0,0.0,0.0);
@@ -628,7 +648,7 @@ class TestLinearizer(unittest.TestCase):
    k = helper_linearizer_opt(out)[-1]
    uops = get_program(k.get_optimized_ast(), k.opts).uops
    # check that the float4 cast collapses
-    store_vals = [u.src[1] for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
+    store_vals = [u.src[-1] for u in uops if u.op is Ops.STORE and u.dtype.addrspace != AddrSpace.REG]
    for val in store_vals:
      assert val.dtype == dtypes.float.vec(4) # and val.op is not Ops.VECTORIZE

@@ -651,7 +671,7 @@ class TestLinearizer(unittest.TestCase):
    x = Tensor.randn((4,3,6,6)).realize()
    out = x.flip((0,1)).contiguous()
    k = helper_linearizer_opt(out)[-1]
-    store_val = [u.src[1] for u in get_program(k.get_optimized_ast(), k.opts).uops if u.op is Ops.STORE][0]
+    store_val = [u.src[-1] for u in get_program(k.get_optimized_ast(), k.opts).uops if u.op is Ops.STORE][0]
    assert store_val.dtype == dtypes.float.vec(4) and store_val.op is not Ops.VECTORIZE

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@@ -670,7 +690,7 @@ class TestLinearizer(unittest.TestCase):
    barrier = [u for u in uops if u.op is Ops.BARRIER][0]
    # check that the float4 cast collapses for all stores
    for store in local_stores+global_stores:
-      assert store.src[1].dtype.count > 1 # and store.src[2].op is not Ops.VECTORIZE
+      assert store.src[-1].dtype.count > 1 # and store.src[2].op is not Ops.VECTORIZE
    # # check the children's vins
    # TODO: src ALU are not the same, should it?
    # assert barrier.src == tuple(local_stores)
@@ -679,20 +699,19 @@ class TestLinearizer(unittest.TestCase):
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
-  @unittest.skipIf(getenv("PTX"), "broken on ptx for some reason")
  def test_grouped_store_local_only(self):
    x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
    r = (x@y).relu()
    k = helper_linearizer_opt(r)[-1]
    uops = get_program(k.get_optimized_ast(), k.opts).uops
-    stores = [u for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
+    stores = [u for u in uops if u.op is Ops.STORE and u.dtype.addrspace != AddrSpace.REG]

    # the float4 value stores directly in lds and we skip upcast
-    self.assertEqual(stores[0].src[1].dtype, dtypes.float.vec(4))
+    self.assertEqual(stores[0].src[-1].dtype, dtypes.float.vec(4))
    #assert stores[0].src[-1].op is not Ops.VECTORIZE

    # the global store doesn't change
-    assert stores[1].src[1].dtype == dtypes.float
+    assert stores[1].src[-1].dtype == dtypes.float

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
@@ -711,7 +730,7 @@ class TestLinearizer(unittest.TestCase):
    ]
    k = helper_linearizer_ast(ast, [Tensor.randn(240*40).realize()], opts=[opt])[-1]
    out = [u for u in get_program(k.get_optimized_ast(), k.opts).uops if u.op is Ops.STORE][0]
-    assert out.src[1].op is Ops.VECTORIZE and out.src[1].dtype == dtypes.float.vec(4)
+    assert out.src[-1].op is Ops.VECTORIZE and out.src[-1].dtype == dtypes.float.vec(4)

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
@@ -729,18 +748,18 @@ class TestLinearizer(unittest.TestCase):
            Opt(op=OptOps.UPCAST, axis=1, arg=0), Opt(op=OptOps.UPCAST, axis=0, arg=2)]
    k = helper_linearizer_ast(ast, [Tensor.randn(8*32).realize()], opts=[opt])[-1]
    out = [u for u in get_program(k.get_optimized_ast(), k.opts).uops if u.op is Ops.STORE][0]
-    assert out.src[1].op is Ops.VECTORIZE and out.src[1].dtype.count != 1
+    assert out.src[-1].op is Ops.VECTORIZE and out.src[-1].dtype.count != 1

@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "need backends that support float4")
 class TestFloat4(unittest.TestCase):
  @staticmethod
  def count_float4(uops: list[UOp], n=4):
    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.float.vec(n)]),
-            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.float.vec(n)]))
+            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[-1].dtype == dtypes.float.vec(n)]))
  @staticmethod
  def count_half4(uops: list[UOp]):
    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.half.vec(4)]),
-            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.half.vec(4)]))
+            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[-1].dtype == dtypes.half.vec(4)]))

  def test_float4_basic(self):
    a = Tensor.empty(2, 8).realize()
@@ -82,7 +82,6 @@ class TestLinearizerDumb(unittest.TestCase):
    assert prg.uops is not None and not any(uop.op is Ops.MAX for uop in prg.uops), "leftover MAX"

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "need local")
-  @unittest.skip("not applicable")
  def test_expander_new_srcs(self):
    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
      UOp(Ops.STORE, dtypes.void, arg=None, src=(
@@ -1,6 +1,7 @@
 # ruff: noqa: E501
 import unittest
-from tinygrad import dtypes
+from tinygrad import dtypes, Device
+from tinygrad.helpers import CI
 from tinygrad.opt.kernel import Kernel
 from tinygrad.opt.search import Opt, OptOps, bufs_from_lin
 from extra.optimization.helpers import time_linearizer
@@ -161,5 +162,33 @@ class TestLinearizerOverflow(unittest.TestCase):
    opts = [Opt(op=OptOps.UPCAST, axis=3, arg=4), Opt(op=OptOps.LOCAL, axis=3, arg=16), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=2, arg=8), Opt(op=OptOps.UPCAST, axis=1, arg=2), Opt(op=OptOps.UPCAST, axis=2, arg=4)]
    _test_overflow(ast, opts)

+@unittest.skipIf(Device.DEFAULT not in {"GPU", "HSA", "CUDA", "METAL"}, "only backends with locals")
+@unittest.skipIf(CI, "slow")
+class TestLinearizerOverflowAlt(unittest.TestCase):
+  def test_overflow_1(self):
+    BS = 2
+    g0, g1, g2 = [UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=i) for i in range(3)]
+    in_st_1 = ShapeTracker(views=(View(shape=(1, BS, 1, 3, 8, 230, 8, 230), strides=(0, 150528, 0, 50176, 0, 224, 0, 1), offset=-675, mask=((0, 1), (0, BS), (0, 1), (0, 3), (0, 8), (3, 227), (0, 8), (3, 227)), contiguous=False),
+                                  View(shape=(BS, 1, 64, 112, 112, 3, 7, 7), strides=(10156800, 0, 0, 3680, 2, 3385600, 425040, 231), offset=0, mask=None, contiguous=False))).to_uop()
+    in_st_2 = ShapeTracker(views=(View(shape=(BS, 1, 64, 112, 112, 3, 7, 7), strides=(0, 0, 147, 0, 0, 49, 7, 1), offset=0, mask=None, contiguous=False),)).to_uop()
+    ot_st = ShapeTracker(views=(View(shape=(BS, 1, 64, 112, 112, 1, 1, 1), strides=(802816, 0, 12544, 112, 1, 0, 0, 0), offset=0, mask=None, contiguous=True),)).to_uop()
+    prod = UOp(Ops.LOAD, dtypes.float, (g1.view(in_st_1.arg),)) * UOp(Ops.LOAD, dtypes.float, (g2.view(in_st_2.arg),))
+    store = UOp(Ops.STORE, src=(g0.view(ot_st.arg), UOp(Ops.REDUCE_AXIS, dtypes.float, (prod,), (Ops.ADD, (7, 6, 5)))))
+    ast = UOp(Ops.SINK, src=(store,))
+    opts = [Opt(op=OptOps.LOCAL, axis=3, arg=16), Opt(op=OptOps.LOCAL, axis=2, arg=2), Opt(op=OptOps.UPCAST, axis=0, arg=2)]
+    _test_overflow(ast, opts)
+  def test_overflow_2(self):
+    BS = 2
+    g0, g1, g2 = [UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=i) for i in range(3)]
+    in_st_1 = ShapeTracker(views=(View(shape=(1, BS, 1, 3, 8, 230, 8, 230), strides=(0, 150528, 0, 50176, 0, 224, 0, 1), offset=-675, mask=((0, 1), (0, BS), (0, 1), (0, 3), (0, 8), (3, 227), (0, 8), (3, 227)), contiguous=False),
+                                  View(shape=(BS, 1, 64, 112, 112, 3, 7, 7), strides=(10156800, 0, 0, 3680, 2, 3385600, 425040, 231), offset=0, mask=None, contiguous=False))).to_uop()
+    in_st_2 = ShapeTracker(views=(View(shape=(BS, 1, 64, 112, 112, 3, 7, 7), strides=(0, 0, 147, 0, 0, 49, 7, 1), offset=0, mask=None, contiguous=False),)).to_uop()
+    ot_st = ShapeTracker(views=(View(shape=(BS, 1, 64, 112, 112, 1, 1, 1), strides=(802816, 0, 12544, 112, 1, 0, 0, 0), offset=0, mask=None, contiguous=True),)).to_uop()
+    prod = UOp(Ops.LOAD, dtypes.float, (g1.view(in_st_1.arg),)) * UOp(Ops.LOAD, dtypes.float, (g2.view(in_st_2.arg),))
+    store = UOp(Ops.STORE, src=(g0.view(ot_st.arg), UOp(Ops.REDUCE_AXIS, dtypes.float, (prod,), (Ops.ADD, (7, 6, 5)))))
+    ast = UOp(Ops.SINK, src=(store,))
+    opts = [Opt(op=OptOps.LOCAL, axis=3, arg=16), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=2, arg=16), Opt(op=OptOps.UPCAST, axis=4, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2), Opt(op=OptOps.UPCAST, axis=5, arg=2)]
+    _test_overflow(ast, opts)
+
 if __name__ == '__main__':
  unittest.main()
@@ -1126,7 +1126,6 @@ class TestMultiRamUsage(unittest.TestCase):
    # NOTE: the first one on the DEFAULT device should be freed
    self.assertUsed(self.N*self.N*4*2)

-  @unittest.skip("flaky")
  def test_zeros_shard(self, devices=(d1, d2)):
    _ = Tensor.zeros(self.N, self.N).contiguous().shard(devices, axis=0).realize()
    self.assertUsed(self.N*self.N*4) # sharding should not increase total ram usage
@@ -34,29 +34,25 @@ class TestBEAM(unittest.TestCase):
    capturing.clear()
    self.assertNotEqual(k_beam_0[-1].prg.p.src, k_beam_1[-1].prg.p.src)

-  def test_get_kernel_actions_dedup(self):
+  def test_get_kernel_actions(self):
    from test.test_linearizer import helper_realized_ast
-    from tinygrad.opt.search import get_kernel_actions
-    a = Tensor.empty(4, 3)
-    b = Tensor.empty(3)
+    a = Tensor.rand(4, 3)
+    b = Tensor.rand(3)
    realized_ast, _ = helper_realized_ast(a @ b)
-    candidates = [
-      Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UPCAST, axis=0, arg=4),
-      Opt(op=OptOps.LOCAL, axis=0, arg=0), Opt(op=OptOps.LOCAL, axis=0, arg=4),
-      Opt(op=OptOps.UNROLL, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=3),
-      Opt(op=OptOps.GROUP, axis=0, arg=0), Opt(op=OptOps.GROUP, axis=0, arg=3),
-      Opt(op=OptOps.GROUPTOP, axis=0, arg=0), Opt(op=OptOps.GROUPTOP, axis=0, arg=3),
-    ]
-    lins = get_kernel_actions(Kernel(realized_ast), include_0=False, candidates=candidates).values()
+    from tinygrad.opt.search import get_kernel_actions
+    lins = get_kernel_actions(Kernel(realized_ast), False).values()

    # ensure amt=0 are not duplicated
-    assert all(len(x.applied_opts) == 1 for x in lins)
-    kernel_actions = [x.applied_opts[0] for x in lins]
-    assert Opt(OptOps.UPCAST, axis=0, arg=4) not in kernel_actions, "did not de-dup UPCAST"
-    assert Opt(OptOps.LOCAL, axis=0, arg=4) not in kernel_actions, "did not de-dup LOCAL"
-    assert Opt(OptOps.UNROLL, axis=0, arg=3) not in kernel_actions, "did not de-dup UNROLL"
-    assert Opt(OptOps.GROUP, axis=0, arg=3) not in kernel_actions, "did not de-dup GROUP"
-    assert Opt(OptOps.GROUPTOP, axis=0, arg=3) not in kernel_actions, "did not de-dup GROUPTOP"
+    if Opt(OptOps.UPCAST, 0, 0) in actions:
+      assert len([x for x in lins if x.applied_opts[0] == Opt(OptOps.UPCAST, axis=0, arg=4)]) == 0, "did not de-dup UPCAST"
+    if Opt(OptOps.LOCAL, 0, 0) in actions:
+      assert len([x for x in lins if x.applied_opts[0] == Opt(OptOps.LOCAL, axis=0, arg=4)]) == 0, "did not de-dup LOCAL"
+    if Opt(OptOps.UNROLL, 0, 0) in actions:
+      assert len([x for x in lins if x.applied_opts[0] == Opt(OptOps.UNROLL, axis=0, arg=3)]) == 0, "did not de-dup UNROLL"
+    if Opt(OptOps.GROUP, 0, 0) in actions:
+      assert len([x for x in lins if x.applied_opts[0] == Opt(OptOps.GROUP, axis=0, arg=3)]) == 0, "did not de-dup GROUP"
+    if Opt(OptOps.GROUPTOP, 0, 0) in actions:
+      assert len([x for x in lins if x.applied_opts[0] == Opt(OptOps.GROUPTOP, axis=0, arg=3)]) == 0, "did not de-dup GROUPTOP"

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
  def test_search_over_shape(self):
@@ -158,13 +158,6 @@ class TestSetitem(unittest.TestCase):
    t[:-1] = t[1:]
    self.assertEqual(t.tolist(), [[2.0], [1.0], [1.0]])

-  def test_setitem_big(self):
-    idx_size, val = 256, 4
-    t = Tensor.arange(0, idx_size+1)
-    idx = Tensor.arange(0, idx_size)
-    t[idx] = val
-    self.assertEqual(t.tolist(), [val]*idx_size+[idx_size])
-
 class TestWithGrad(unittest.TestCase):
  def test_no_requires_grad_works(self):
    z = Tensor.rand(8, 8)
@@ -86,6 +86,7 @@ class TestFuse(unittest.TestCase):
      return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
    self._test_fuse(embedding, a, atol=1e-5)

+  @unittest.skip("still broken")
  def test_flash_attention(self):
    BS = 4
    HEADS = 2
@@ -97,7 +98,7 @@ class TestFuse(unittest.TestCase):
      v = Tensor.randn(BS, HEADS, MATDIM, EMB).realize()
    # TODO: OPT is breaking things. NOOPT isn't linearizing
    with Context(NOOPT=1):
-      self._test_fuse(Tensor.scaled_dot_product_attention, q, k, v, atol=1e-5)
+      self._test_fuse(Tensor.scaled_dot_product_attention, q, k, v)

 class TestSoftmaxFusion(unittest.TestCase):
  @classmethod
@@ -512,24 +512,6 @@ class TestTinygrad(unittest.TestCase):
    subprocess.run([f'NPY=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
                    shell=True, check=True)

-    if Device.DEFAULT != "CPU":
-      # setting multiple devices fail
-      with self.assertRaises(subprocess.CalledProcessError):
-        subprocess.run([f'{Device.DEFAULT}=1 CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                        shell=True, check=True)
-
-      # setting device via DEV
-      subprocess.run([f'DEV={Device.DEFAULT.capitalize()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                      shell=True, check=True)
-      subprocess.run([f'DEV={Device.DEFAULT.lower()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                      shell=True, check=True)
-      subprocess.run([f'DEV={Device.DEFAULT.upper()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                      shell=True, check=True)
-
-      with self.assertRaises(subprocess.CalledProcessError):
-        subprocess.run([f'DEV={Device.DEFAULT} CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                        shell=True, check=True)
-
  def test_no_attributeerror_after_apply_uop_exception(self):
    try:
      Tensor.arange(4).reshape(3,2)
@@ -33,35 +33,6 @@ class TestTiny(unittest.TestCase):
    self.assertListEqual((out:=a@b).flatten().tolist(), [1.0]*(N*N))
    if IMAGE < 2: self.assertEqual(out.dtype, out_dtype)

-  def test_double_gemm(self, N=64, BS=1):
-    a = Tensor.ones(BS,N,N).contiguous().realize()
-    b = Tensor.eye(N).contiguous().realize()
-    c = Tensor.eye(N).contiguous().realize()
-    d = Tensor.eye(N).contiguous().realize()
-    e = Tensor.eye(N).contiguous().realize()
-    f = Tensor.eye(N).contiguous().realize()
-    g = Tensor.eye(N).contiguous().realize()
-    out = ((a@b@c@d).contiguous()@e@f@g).contiguous().realize()
-    self.assertListEqual(out.flatten().tolist(), [1.0]*(BS*N*N))
-
-  def test_double_gemm_bs(self, N=64, BS=1): self.test_double_gemm(BS=4)
-
-  def test_conv2d(self):
-    N = 64
-    a = Tensor.ones(1,4,N,N).contiguous().realize()
-    w1 = Tensor.ones(16,4,3,3).contiguous().realize()
-    out = a.conv2d(w1).contiguous().realize()
-
-  def test_double_conv2d(self):
-    N = 64
-    a = Tensor.ones(1,4,N,N).contiguous().realize()
-    w1 = Tensor.ones(4,4,3,3).contiguous().realize()
-    w2 = Tensor.ones(4,4,3,3).contiguous().realize()
-    w3 = Tensor.ones(4,4,3,3).contiguous().realize()
-    w4 = Tensor.ones(4,4,3,3).contiguous().realize()
-    w5 = Tensor.ones(4,4,3,3).contiguous().realize()
-    out = a.conv2d(w1).conv2d(w2).conv2d(w3).conv2d(w4).conv2d(w5).contiguous().realize()
-
  # *** randomness ***

  def test_random(self):
@@ -317,7 +317,6 @@ class TestUOpGraph(unittest.TestCase):
      for uop, const in zip(uops, consts):
        self.assertEqual(uop, const)

-  @unittest.skip("no longer testable standalone")
  def test_wmma_vectorize_fold(self):
    for i in [2, 4, 8]:
      vec = UOp(Ops.VECTORIZE, dtypes.half.vec(i), tuple(UOp.const(dtypes.half, 0.0) for _ in range(i)))
@@ -2,21 +2,6 @@ from typing_extensions import Callable
 import hashlib, random, unittest
 from tinygrad import Tensor, Device, getenv, dtypes
 from tinygrad.device import is_dtype_supported
-from tinygrad.helpers import CI
-
-@unittest.skipUnless(is_dtype_supported(dtypes.uint8) and is_dtype_supported(dtypes.uint64), "Device must support uint8 and uint64")
-@unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "NV", "crashes in NV CI")
-class TestHashing(unittest.TestCase):
-  def _python_hash_1mb(self, data:bytes):
-    chunks = [data[i:i+4096] for i in range(0, len(data), 4096)]
-    chunk_hashes = [hashlib.shake_128(chunk).digest(16) for chunk in chunks]
-    return hashlib.shake_128(b''.join(chunk_hashes)).digest(16)
-
-  @unittest.skipIf(CI, "very slow")
-  def test_abc(self):
-    expected = self._python_hash_1mb(b"abc" + b"\x00" * (2**20 - 3))
-    out = Tensor(b"abc").hash()
-    self.assertEqual(bytes(out.data()), expected)

@unittest.skipUnless(is_dtype_supported(dtypes.uint8) and is_dtype_supported(dtypes.uint64), "Device must support uint8 and uint64")
@unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "NV", "crashes in NV CI")
@@ -48,25 +33,19 @@ class TestKeccak(unittest.TestCase):
      self.assertEqual(ha_ref, Tensor(a).keccak(name).data())
      self.assertEqual(hb_ref, hb)

-  def test_referenced(self):
+  def test_abc(self):
    # https://www.di-mgt.com.au/sha_testvectors.html
-    self.assertEqual(bytes(Tensor(b"abc").keccak().tolist()),
-                     bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
-    self.assertEqual(bytes(Tensor(b"").keccak().tolist()),
-                     bytearray.fromhex("a7ffc6f8bf1ed766 51c14756a061d662 f580ff4de43b49fa 82d80a4b80f8434a"))
-    t = Tensor(b"abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu").keccak()
-    self.assertEqual(bytes(t.tolist()),
-                     bytearray.fromhex("916f6061fe879741 ca6469b43971dfdb 28b1a32dc36cb325 4e812be27aad1d18"))
-    # TODO: this does not run or very slow
-    # self.assertEqual(bytes(Tensor(b"a" * 1000000).keccak().tolist()),
-    #                  bytearray.fromhex("5c8875ae474a3634 ba4fd55ec85bffd6 61f32aca75c6d699 d0cdcb6c115891c1"))
+    out = Tensor(b"abc").keccak()
+    self.assertEqual(bytes(out.tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))

  def test_long(self):
    data = b"\x00" * 4
    self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))

-    data = b"\x00" * (1000 if CI else 4096)
-    self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))
+    data = b"\x00" * 4096
+    with self.assertRaises(RecursionError):
+      # TODO: fix
+      self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))

 if __name__ == "__main__":
  unittest.main()
@@ -827,6 +827,39 @@ class TestShapeTrackerSize(unittest.TestCase):
    st = ShapeTracker.from_shape((10,10)).pad(((2,4), (3,1))).flip((True, True))
    self.assertEqual(st.real_size(), 100)

+class TestConsecutive(unittest.TestCase):
+  @classmethod
+  def setUpClass(self):
+    from tinygrad.tensor import Tensor  # easier test setup
+    self.t = Tensor([[1, 2, 3, 4], [5, 6, 7, 8]])
+    self.const = Tensor(2)
+    self.ones = Tensor.ones(2, 4)
+
+  def test_unmodified(self):
+    assert self.t.uop.st.consecutive
+    assert self.t.reshape(4, 2).uop.st.consecutive
+    assert self.t.reshape(1, 8).uop.st.consecutive
+
+  def test_sliced(self):
+    assert self.t[0].uop.st.consecutive
+    assert self.t[0, 1:2].uop.st.consecutive
+    assert self.t[1].uop.st.consecutive
+    assert not self.t[:, 0].uop.st.consecutive
+    assert not self.t[:, 1].uop.st.consecutive
+
+  def test_padded(self):
+    assert not self.t.pad(((1, 1), None)).uop.st.consecutive
+    assert not self.t.pad((None, (1, 1))).uop.st.consecutive
+
+  def test_const(self):
+    assert self.const.uop.st.consecutive
+
+  def test_ones(self):
+    assert not self.ones.uop.st.consecutive
+    assert not self.ones[0, :].uop.st.consecutive
+    # consecutive if sliced into size 1
+    assert self.ones[0, 0].uop.st.consecutive
+
 class TestRender(unittest.TestCase):
  def test_render(self):
    st = ShapeTracker.from_shape((2, 3))
@@ -14,7 +14,7 @@ def render(self) -> tuple[str, ConstType, ConstType]:
  # NOTE: we need STORE so the ALU op has children
  glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
  uops = full_rewrite(UOp(Ops.STORE, dtypes.void, (glbl.index(UOp.const(dtypes.int, 0)), self)).sink())
-  rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1]
+  rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[-1]
  return rewritten_uop.render(simplify=False), rewritten_uop.vmin, rewritten_uop.vmax

 def uconst(val): return UOp.const(dtypes.int, val)
@@ -642,7 +642,7 @@ class TestSymbolic(unittest.TestCase):
    # TODO: copied from render, render does not support cast
    glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
    uops = full_rewrite(UOp(Ops.STORE, dtypes.void, (glbl.index(UOp.const(dtypes.int, 0)), expr)).sink())
-    rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1]
+    rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[-1]

    self.assertEqual(rewritten_uop, cond.where(a.cast(dtypes.half), b.cast(dtypes.half)))

@@ -138,16 +138,6 @@ class TestViz(BaseTestViz):
    nop = UOp(Ops.NOOP, arg="infinite loop in fixed_point_rewrite")
    self.assertEqual(graphs[2], uop_to_json(nop)[id(nop)])

-  def test_const_node_visibility(self):
-    a = UOp.variable("a", 0, 10)
-    z = UOp.const(dtypes.int, 0)
-    alu = a*z
-    exec_rewrite(alu, [sym])
-    graphs = [x["graph"] for x in get_details(tracked_ctxs[0][0])]
-    # embed const in the parent node when possible
-    self.assertEqual(list(graphs[0]), [id(a), id(alu)])
-    self.assertEqual(list(graphs[1]), [id(z)])
-
 # VIZ displays nested graph_rewrites in a tree view

 def leaf_rewrite(x:UOp): return x.rtag(1) if x.tag is None else None
@@ -12,9 +12,8 @@ from tinygrad.codegen.quantize import pm_quant
 from tinygrad.codegen.gpudims import pm_add_gpudims
 from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing
 from tinygrad.codegen.expander import migrate_indexing, expander
-from tinygrad.codegen.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
-  ReduceContext, correct_load_store, pm_render
-from tinygrad.codegen.optional import get_late_rewrite_patterns
+from tinygrad.codegen.devectorizer import load_store_folding, load_store_indexing, devectorize, \
+  pm_reduce, ReduceContext, correct_load_store, pm_render, get_late_rewrite_patterns
 from tinygrad.codegen.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext

@dataclass
@@ -1,11 +1,13 @@
-from typing import Any, cast
+from typing import Any, Callable, cast
 import functools, operator, itertools
 from collections import defaultdict
 from dataclasses import dataclass
-from tinygrad.dtype import dtypes, ImageDType, PtrDType, DType, AddrSpace
+from tinygrad.device import is_dtype_supported
+from tinygrad.dtype import dtypes, ImageDType, PtrDType, promo_lattice, DType, AddrSpace
 from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, graph_rewrite, GroupOp, identity_element
 from tinygrad.uop.symbolic import split_uop, uop_given_valid, parse_valid, simplify_valid, sym, symbolic_flat
-from tinygrad.helpers import getenv, flatten, AMX, prod, partition
+from tinygrad.helpers import getenv, flatten, AMX, prod, partition, all_same
+from tinygrad.uop.transcendental import xexp2, xlog2, xsin, xpow, TRANSCENDENTAL_SUPPORTED_DTYPES
 from tinygrad.renderer import Renderer

 # ***** image load valid simplification *****
@@ -45,10 +47,10 @@ def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
  new_valid = functools.reduce(operator.and_, ss) if (ss:=[s for s in split_uop(valid, Ops.AND) if s not in drop_stmt]) else None
  return buf.index(idx, new_valid)

-def delete_redundant_gates(store:UOp, buf:UOp, idx:UOp, val:UOp, store_gate:UOp, cast:UOp|None=None) -> UOp|None:
+def delete_redundant_gates(buf:UOp, idx:UOp, val:UOp, store_gate:UOp, cast:UOp|None=None) -> UOp|None:
  if store_gate not in [gate.src[0] for gate in val.toposort() if gate.op is Ops.IF]: return None
  # remove the gate from the index
-  return UOp.store(buf.index(idx).cast(cast.dtype) if cast is not None else buf.index(idx), val, *store.src[2:])
+  return UOp.store(buf.index(idx).cast(cast.dtype) if cast is not None else buf.index(idx), val)

 load_store_indexing = PatternMatcher([
  # simplify valid
@@ -59,7 +61,7 @@ load_store_indexing = PatternMatcher([
  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("start_idx"), UPat(Ops.CONST, arg=True))), lambda buf,start_idx: buf.index(start_idx)),
  # delete_redundant_gates (after expand)
  (UPat(Ops.STORE, src=(UPat.any(stidx:=UPat.var("buf").index(UPat.var("idx"), UPat.var("store_gate")), stidx.cast().named("cast")),
-                                  UPat.var("val")), name="store", allow_any_len=True), delete_redundant_gates),
+                                  UPat.var("val"))), delete_redundant_gates),
 ])

 # ***** load/store grouping *****
@@ -111,7 +113,11 @@ def cat_after_store(cat:UOp, data:UOp, sto:UOp):
  for s in cat.src:
    ret.append(s.store(data.gep(tuple(range(offset, offset+s.dtype.count))), *sto.src[2:]))
    offset += s.dtype.count
-  return UOp(Ops.NOOP, src=tuple(ret))
+  # dtype CAT
+  dtypes: list[PtrDType] = [x.dtype for x in ret if isinstance(x.dtype, PtrDType)]
+  assert len(dtypes) == len(ret) and all_same([(x.size, x.addrspace) for x in dtypes])
+  out_dtype = dtypes[0].base.scalar().vec(sum([x.count for x in dtypes])).ptr(dtypes[0].size, dtypes[0].addrspace)
+  return UOp(Ops.PTRCAT, dtype=out_dtype, src=tuple(ret))

 def gep_on_store(gep:UOp, st:UOp, sto:UOp):
  # NOTE: we need to invert the gep here, but it may be an expanding gep
@@ -122,8 +128,8 @@ def gep_on_store(gep:UOp, st:UOp, sto:UOp):
  return gep.src[0].store(st.gep(new_arg), *sto.src[2:])

 load_store_folding = PatternMatcher([
-  (UPat(Ops.INDEX, src=(UPat(Ops.VECTORIZE, src=UPat(GroupOp.Defines, name="buf")), UPat.var("vec"))), expand_index),
-  (UPat(Ops.INDEX, src=(UPat(Ops.VECTORIZE, src=UPat(GroupOp.Defines, name="buf")), UPat.var("vec"),
+  (UPat(Ops.INDEX, src=(UPat(Ops.VECTORIZE, src=UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL), name="buf")), UPat.var("vec"))), expand_index),
+  (UPat(Ops.INDEX, src=(UPat(Ops.VECTORIZE, src=UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL), name="buf")), UPat.var("vec"),
                        UPat.var("mask"))), expand_index),
  # GEP after LOAD
  (UPat(Ops.LOAD, src=(UPat(Ops.GEP, name="gep"),), name="ld", allow_any_len=True),
@@ -137,6 +143,55 @@ load_store_folding = PatternMatcher([
  (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data")), allow_any_len=True, name="sto"), cat_after_store),
 ])

+# ***** optional patterns *****
+
+@functools.lru_cache(None)
+def magicgu(vmax:int, d:int) -> tuple[int,int]:
+  # calculate m,s such that x//d == (x*m) >> s for all 0 <= x <= vmax, d>0; adapted from Hacker's Delight, Chapter 10
+  nc = (vmax+1)//(d) * d - 1
+  nbits = vmax.bit_length()
+  for s in range(0, 2*nbits + 1):
+    if 2**s > nc*(d - 1 - (2**s - 1) % d):
+      m = (2**s + d - 1 - (2**s - 1) % d)//d
+      return m, s
+  assert False
+
+def fast_idiv(ctx: Renderer|None, x: UOp, d: int) -> UOp|None:
+  # idiv is truncated division, but arithmetic shift is floored division, so can only do non-negative numbers!
+  if x.vmin<0: return None
+  sign = 1 if d > 0 else -1
+  m,s = magicgu(vmax := min(x.vmax, dtypes.max(x.dtype)), abs(d))
+  if m * vmax <= dtypes.max(x.dtype): return sign * ((x*m) >> s)
+  # promo_lattice needs to return an unsigned type
+  if ctx is not None and dtypes.is_int(next_dtype := promo_lattice[x.dtype][-1]) and is_dtype_supported(next_dtype, ctx.device):
+    if m * vmax <= dtypes.max(next_dtype): return sign * ((x.cast(next_dtype)*m) >> s).cast(x.dtype)
+  return None
+
+powers_of_two = {2**i:i for i in range(64)}
+@functools.cache
+def get_late_rewrite_patterns(ops, force_transcendental=False):
+  pat: list[tuple[UPat, Callable]] = [(UPat(op, dtype=TRANSCENDENTAL_SUPPORTED_DTYPES, src=(UPat.var("d"),)), f) for op,f in \
+           ((Ops.EXP2, xexp2), (Ops.LOG2, xlog2), (Ops.SIN, xsin)) if op not in ops or force_transcendental]
+  # rewrite SQRT to xpow 0.5
+  if Ops.SQRT not in ops: pat.append((UPat(Ops.SQRT, src=UPat.var("d")), lambda d: xpow(d, d.const_like(0.5))))
+  # rewrite MOD to AND (which should always be supported, but not for generic in tests): x % (2**y) -> x & (2**y-1)
+  if Ops.AND in ops: pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.arg-1) if c.arg in powers_of_two else None)]
+  # rewrite MUL/IDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y)
+  if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.arg, 0)) else None)]
+  if Ops.SHR in ops:
+    # no reason to check x<0 for uints
+    pat += [(UPat.var("x", dtypes.uints)//UPat.cvar("c"), lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]
+    pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("c"), lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(
+      c-1, 0)) >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]  # (x+(x<0).where(c-1, 0)) >> v
+    if not getenv("DISABLE_FAST_IDIV"):
+      pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("d"), lambda ctx, x, d: fast_idiv(ctx, x, d.arg))]
+      pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("d"), lambda ctx, x, d: x - d*f if (f:=fast_idiv(ctx, x, d.arg)) is not None else None)]
+  if Ops.NEG in ops:
+    pat += [(UPat.var('x')*-1, lambda x: x.alu(Ops.NEG))]
+    if Ops.SUB in ops: pat += [(UPat.var('x')+UPat.var('y').alu(Ops.NEG), lambda x,y: x.alu(Ops.SUB, y))]
+  if Ops.MULACC in ops: pat += [(UPat.var('a')*UPat.var('b')+UPat.var('c'), lambda a,b,c: a.alu(Ops.MULACC, b, c))]
+  return PatternMatcher(pat)
+
 # *** correct load/store ***

 def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
@@ -154,8 +209,6 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
    must_divide = False
  elif buf.dtype.base != dtypes.float and buf.dtype.base != dtypes.half and not isinstance(buf.dtype, ImageDType):
    pass
-  elif cast(PtrDType, buf.dtype).addrspace == AddrSpace.REG:
-    pass
  elif isinstance(buf.dtype, ImageDType):
    lengths = [4]
  elif ctx is not None and ctx.supports_float4:
@@ -182,8 +235,7 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
      break

  # if it wasn't split, we return None. otherwise we CAT them
-  if len(ret) <= 1: return None
-  return UOp(Ops.CAT, ls.dtype, tuple(ret)) if ls.op is Ops.LOAD else UOp(Ops.NOOP, src=tuple(ret))
+  return UOp(Ops.CAT, ls.dtype, tuple(ret)) if len(ret) > 1 else None

 def image_fixup(ls:UOp):
  # normal image load or store, with the CAST from expand_index
@@ -235,8 +287,9 @@ def no_vectorized_alu(alu:UOp):
 def no_vectorized_acc(acc:UOp, c:UOp):
  if acc.dtype.count == 1: return None
  assert c.arg == 0, "this only supports index 0"
-  new_acc = acc.replace(dtype=acc.dtype.base.scalar().ptr(acc.dtype.count, cast(PtrDType, acc.dtype).addrspace))
-  return UOp(Ops.PTRCAT, acc.dtype, tuple([new_acc.index(UOp.const(dtypes.int, i)) for i in range(acc.dtype.count)]))
+  alus = tuple(UOp(acc.op, acc.dtype.base.scalar().ptr(1, cast(PtrDType, acc.dtype).addrspace),
+    tuple(s.gep(i) if j == 0 else s for j,s in enumerate(acc.src)), acc.arg+(i,)).index(UOp.const(dtypes.int, 0)) for i in range(acc.dtype.count))
+  return UOp(Ops.PTRCAT, acc.dtype, alus)

 devectorize = PatternMatcher([
  # no ALU on vectorized dtypes
@@ -257,9 +310,9 @@ pm_render = PatternMatcher([
  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat())).or_casted(),), allow_any_len=True, name="x"),
   lambda x: x.replace(src=(x.src[0], x.const_like(0))+x.src[1:]) if len(x.src) == 1 or x.src[1].op is Ops.CUSTOM else None),
  # gate any stores that aren't gated with ifs
-  (UPat(Ops.STORE, src=(UPat(src=(UPat(), UPat(), UPat(dtype=dtypes.bool)), name="idx").or_casted(), UPat()), name="store", allow_any_len=True),
-    lambda store,idx: UOp(Ops.STORE, dtype=store.dtype, src=store.src[:2]+(UOp(Ops.IF, src=(idx.src[2],)),)+store.src[2:]) if \
-      len(store.src) <= 2 or store.src[2].op != Ops.IF else None),
+  (UPat(Ops.STORE, src=(UPat(src=(UPat.var("buf"), UPat.var("idx"), UPat(name="gate", dtype=dtypes.bool))).or_casted(), UPat.var("val")),
+        name="store", allow_any_len=True),
+        lambda gate,store,buf,idx,val: UOp(Ops.STORE, dtype=store.dtype, src=(buf.index(idx), val, UOp(Ops.IF, src=(gate,)),)+store.src[2:])),
 ])

 # *** Ops.REDUCE -> Ops.DEFINE_ACC ***
@@ -282,16 +335,12 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp):
  assert all(x.dtype == red.dtype for x in lst), f"horizontal reduction mismatch {lst[0].dtype} != {red.dtype}"
  # if we have a range
  if len(reduce_range) != 0:
-    topo = inp.toposort()
-    stored_ranges = flatten([x.src[2:] for x in topo if x.op is Ops.STORE])
-    input_ranges = tuple([x for x in topo if x.op is Ops.RANGE and x not in reduce_range and x not in stored_ranges])
-    identity = red.const_like(identity_element(red.arg, red.dtype.scalar()))
-    acc = UOp(Ops.DEFINE_REG, red.dtype.ptr(size=1, addrspace=AddrSpace.REG), arg=(ctx.acc_num,)).index(UOp.const(dtypes.int, 0))
-    do_store = acc.store(identity, UOp(Ops.NOOP, src=input_ranges)) if len(input_ranges) else acc.store(identity)
-    lst = [acc.load(do_store, *reduce_range)] + lst  # put acc as the first element
+    acc = UOp(Ops.DEFINE_REG, red.dtype.ptr(size=1, addrspace=AddrSpace.REG),
+              (red.const_like(identity_element(red.arg, red.dtype.scalar())),) + tuple(reduce_range), (ctx.acc_num,)).index(UOp.const(dtypes.int, 0))
+    lst = [acc.load()] + lst  # put acc as the first element
    ctx.acc_num += 1
  ret = functools.reduce(lambda x,y: x.alu(red.arg, y), lst)
-  return acc.load(acc.store(ret, *reduce_range)) if len(reduce_range) != 0 else ret
+  return acc.store(ret, *reduce_range).load() if len(reduce_range) != 0 else ret

 def no_vectorized_reduce(inp:UOp, red:UOp):
  if inp.dtype != red.dtype:
@@ -86,6 +86,9 @@ expander = PatternMatcher([
  (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.GEP, Ops.WMMA, Ops.LOAD, Ops.STORE, Ops.INDEX,
         Ops.VECTORIZE, Ops.IF, Ops.REDUCE), name="root", custom_early_reject=set([Ops.UNROLL])), do_expand),
  (UPat(Ops.CONTRACT, name="con"), do_contract),
+  # vectorize DEFINE_ACC
+  (UPat(Ops.VECTORIZE, src=UPat(Ops.DEFINE_REG, name="acc"), name="v"),
+    lambda acc,v: acc.replace(dtype=v.dtype, src=(acc.src[0].broadcast(v.dtype.count),)+acc.src[1:])),
  # BARRIERs aren't actually expanded
  (UPat(Ops.BARRIER, src=(UPat(Ops.UNROLL, name="ex"),)),
   lambda ex: UOp(Ops.UNROLL, src=(UOp(Ops.BARRIER, src=ex.src),)*len(ex.src), arg=ex.arg)),
@@ -1,5 +1,5 @@
 import math
-from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType
+from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify
 from tinygrad.helpers import all_int
 from tinygrad.dtype import dtypes
 from tinygrad.shape.view import get_contraction
@@ -53,36 +53,20 @@ def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|No
 def add_gpudims(ctx:Renderer, s:UOp):
  if s.arg is None: return None
  ki: KernelInfo = s.arg
-  global_dims = [i for i,x in enumerate(ki.axis_types) if x is AxisType.GLOBAL]
-  local_dims = [i for i,x in enumerate(ki.axis_types) if x in (AxisType.LOCAL, AxisType.GROUP_REDUCE)]
-  if not global_dims and not local_dims: return None
+  if not ki.global_dims and not ki.local_dims: return None
  s_topo = list(s.toposort())
  if any(x.op is Ops.SPECIAL for x in s_topo): return None
-
-  # get global and local shape
-  all_ranges = {x.arg%1000:x for x in s_topo if x.op is Ops.RANGE}
-  ranges = [all_ranges[r] for r in global_dims+local_dims if r in all_ranges]
-  global_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg%1000 in global_dims])
-  local_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg%1000 in local_dims])
-
-  # get the idxs
+  ranges = sorted([x for x in s_topo if x.op is Ops.RANGE and x.arg in (ki.global_dims+ki.local_dims)], key=lambda x: x.arg)
+  if not len(ranges): return None
+  global_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg in ki.global_dims])
+  local_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg in ki.local_dims])
  if ki.dont_use_locals:
-    assert not local_dims, "can't use locals if there's no local dims"
+    assert not ki.local_dims, "can't use locals if there's no local dims"
    idxs = get_grouped_dims("idx", global_shape, ctx.global_max, reverse=True)
  else:
    # define indexes for GPU-like execution
    idxs = get_grouped_dims("gidx", global_shape, ctx.global_max, reverse=True) + get_grouped_dims("lidx", local_shape, ctx.local_max)
-
-  # apply to multiple ranges
-  subs = {}
-  for r in s_topo:
-    if r.op is not Ops.RANGE: continue
-    try:
-      ii = (global_dims+local_dims).index(r.arg%1000)
-      if r.arg < 2000 and ki.axis_types[r.arg%1000] == AxisType.GROUP_REDUCE: continue
-      subs[r] = idxs[ii]
-    except ValueError: continue
-  return s.substitute(subs)
+  return s.substitute(dict(zip(ranges, idxs)))

 pm_add_gpudims = PatternMatcher([
  (UPat(Ops.SINK, name="s"), add_gpudims),
@@ -3,7 +3,7 @@ import heapq
 from collections import defaultdict
 from dataclasses import dataclass, replace
 from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp
-from tinygrad.helpers import dedup, all_same, flatten, getenv
+from tinygrad.helpers import dedup, partition, all_same, flatten, getenv

 # NOTE: any toposort should be valid here, unlike last time this isn't required, it's just for speed
 def block_reorder(lst:list[UOp]) -> list[UOp]:
@@ -86,18 +86,27 @@ class BlockContext:
        ctx.child_count[s] += 1
        this_block_ctx += ctx.last_ctx(s)

-      # save the block ctx. SINK never has anything
-      ctx.block_ctxs[u] = _sort_ctx(this_block_ctx) if u.op is not Ops.SINK else ()
+      # save the block ctx
+      ctx.block_ctxs[u] = _sort_ctx(this_block_ctx)

      # RANGE/IF add to the next ctx
-      # STORE/REDUCE_AXIS subtract from the next ctx
+      # STORE/ASSIGN subtract from the next ctx
      if u.op in {Ops.RANGE, Ops.IF}: ctx.child_ctxs[u] = _sort_ctx(ctx.block_ctxs[u] + (u,))
-      elif u.op in {Ops.STORE, Ops.REDUCE_AXIS, Ops.CONTIGUOUS}: ctx.child_ctxs[u] = tuple([y for y in ctx.block_ctxs[u] if y not in u.src])
+      elif u.op is Ops.STORE:
+        if len(definereg:=[x for x in u.src[0].toposort() if x.op is Ops.DEFINE_REG]):
+          # old assign logic
+          ctx.child_ctxs[u] = tuple([y for y in ctx.last_ctx(u.src[1]) if y not in definereg[0].src[1:]])
+        elif any(x.op is Ops.DEFINE_LOCAL for x in u.src[0].toposort()):
+          # deal with non-reduce locals. probably wrong
+          idx_context, store_context = ctx.last_ctx(u.src[0]), ctx.last_ctx(u.src[1])
+          ctx.child_ctxs[u] = tuple([y for y in store_context if y not in idx_context and y.op is Ops.RANGE])
+        else: ctx.child_ctxs[u] = ()
+      #elif u.op is Ops.STORE: ctx.child_ctxs[u] = tuple([x for x in ctx.block_ctxs[u] if x not in u.src])
    return ctx

 # ***** make blocks *****

-DONT_PLACE_IN_BLOCK = {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.CONST}
+DONT_PLACE_IN_BLOCK = {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.CONST}

 def add_blockends(base_block:UOp, new_ctx:tuple[UOp, ...], current_ctx:tuple[UOp, ...], cnt:int=1) -> UOp:
  ends_to_add = [z for z in new_ctx if z not in current_ctx]
@@ -207,15 +216,12 @@ def remove_blockend(x:UOp):
    assert all_same(parent_blocks), f"should never have two parent blocks (has {len(parent_blocks)})"
    parent_block = parent_blocks[0]
    assert len(parent_blocks) == parent_block.arg.cnt
-    # NOTE: DEFINE_ACC doesn't have to be handled in any special way
-    late_ops = list(x.arg.lst)
+    # range needs DEFINE_ACC to be before the range (never in DEFINE_ACC for if)
+    early_ops, late_ops = partition(x.arg.lst, lambda y: y.op is Ops.DEFINE_REG and x.arg.end in y.src)
    # NOTE: we have to add a barrier at the start if barrier is used in the range
    if x.op is Ops.BLOCKEND and any(y.op is Ops.BARRIER for y in late_ops) and late_ops[-1].op is Ops.ENDRANGE:
      late_ops = [UOp(Ops.BARRIER)] + late_ops
-    # peephole opt, remove any BARRIERs next to each other
-    for i in range(len(late_ops)-1):
-      if late_ops[i].op is Ops.BARRIER and late_ops[i+1].op is Ops.BARRIER: late_ops[i+1] = UOp(Ops.NOOP)
-    arg = BasicBlock(parent_block.arg.lst+tuple(late_ops), tuple([y for y in x.arg.ctx if y is not x.arg.end]), cnt=x.arg.cnt)
+    arg = BasicBlock(tuple(early_ops)+parent_block.arg.lst+tuple(late_ops), tuple([y for y in x.arg.ctx if y is not x.arg.end]), cnt=x.arg.cnt)
    return UOp(Ops.BLOCK, src=tuple(y for y in x.src if y is not parent_block)+parent_block.src, arg=arg)

 block_merge = PatternMatcher([
@@ -1,94 +1,68 @@
 # the job of the lowerer is to do indexing
-import functools, operator
-from typing import cast
 from dataclasses import dataclass
-from tinygrad.dtype import dtypes, AddrSpace, PtrDType
-from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint_to_uop, AxisType, graph_rewrite
+from typing import cast
+from tinygrad.dtype import dtypes, PtrDType, AddrSpace
+from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint_to_uop, AxisType
 from tinygrad.helpers import prod, partition, flatten

 # ***** indexing *****

@dataclass
 class IndexContext:
-  axis_types: tuple[AxisType, ...]
  idxs: list[UOp]
-  start: int = 0
-
-def shape_to_idx(s, axis_types, start=0):
-  # indexes
-  idxs = []
-  for i, (s, at) in enumerate(zip(s, axis_types)):
-    if at in (AxisType.UPCAST, AxisType.UNROLL):
-      assert isinstance(s, int), "needs to be int to upcast/unroll"
-      idxs.append(UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(s), tuple(range(s))),), ((i,s),), tag=1))
-    else:
-      # all others are RANGES
-      idxs.append(UOp(Ops.RANGE, dtypes.int, (sint_to_uop(s),), start+i))
-  return idxs
+  ridxs: list[UOp]

 def get_index(ast:UOp) -> IndexContext:
  axis_types = ast.arg.axis_types if isinstance(ast.arg, KernelInfo) else ()
  if len(ast.full_shape) != len(axis_types): axis_types = (AxisType.LOOP,)*len(ast.full_shape)
-  return IndexContext(axis_types, [], 0)
+
+  # indexes
+  idxs = []
+  for i, (s, at) in enumerate(zip(ast.full_shape, axis_types)):
+    if at in (AxisType.UPCAST, AxisType.UNROLL):
+      assert isinstance(s, int), "needs to be int to upcast/unroll"
+      idxs.append(UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(s), tuple(range(s))),), ((i,s),)))
+    else:
+      # all others are RANGES
+      idxs.append(UOp(Ops.RANGE, dtypes.int, (sint_to_uop(s),), i))
+
+  # late indexes (group for reduce)
+  ridxs = idxs[:]
+  for i, (s, at) in enumerate(zip(ast.full_shape, axis_types)):
+    if at == AxisType.GROUP_REDUCE:
+      ridxs[i] = UOp(Ops.RANGE, dtypes.int, (sint_to_uop(s),), 1000+i)
+
+  return IndexContext(idxs, ridxs)

 # ***** lowering (given index) *****

-def subblock(ctx: IndexContext, full_new_idx: list[UOp], src: UOp):
-  lc = IndexContext(ctx.axis_types, full_new_idx, ctx.start+1000)
-  ctx.start = lc.start
-  return graph_rewrite(src, pm_lowerer, lc, name="subblock", bottom_up=True)
-
 def lower_reduce_axis(ctx: IndexContext, x: UOp):
-  new_idxs = shape_to_idx(x.src[0].shape, ctx.axis_types, ctx.start)
-  full_new_idx = list(ctx.idxs)
-  for a in x.axis_arg: full_new_idx[a] = new_idxs[a]
-
-  ret = subblock(ctx, full_new_idx, x.src[0])
-
  # NOTE: always using ridxs is fine here
-  reduce_range, reduce_expand = partition([full_new_idx[i] for i in x.axis_arg], lambda y: y.op is Ops.RANGE)
+  reduce_range, reduce_expand = partition([ctx.ridxs[i] for i in x.axis_arg], lambda y: y.op is Ops.RANGE)
  assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand} for {x.axis_arg}"
+  ret = x.src[0]
  if len(contract_axis:=flatten(x.arg for x in reduce_expand)):
-    ret = UOp(Ops.CONTRACT, x.dtype.vec(prod(x[1] for x in contract_axis)), (ret,), tuple(contract_axis), tag=1)
+    ret = UOp(Ops.CONTRACT, x.dtype.vec(prod(x[1] for x in contract_axis)), (ret,), tuple(contract_axis))
  # REDUCE supports both "horizontal" reduction and range reduction. the horizontal elements are taken in the nearest group
  return UOp(Ops.REDUCE, x.dtype, (ret,)+tuple(reduce_range), x.arg[0])

+def lower_load(ctx: IndexContext, x: UOp, buf: UOp):
+  idx, valid = x.st_arg.to_indexed_uops(ctx.ridxs if buf.op is Ops.DEFINE_LOCAL else ctx.idxs)
+  barrier = (UOp(Ops.BARRIER, dtypes.void, (x.src[1],)),) if buf.op is Ops.DEFINE_LOCAL else ()
+  return UOp(Ops.LOAD, x.dtype, (buf.index(idx, valid),) + barrier)
+
 def lower_store(ctx: IndexContext, x: UOp, buf: UOp):
-  # TODO: reenable after REDUCE_AXIS is fixed
-  #assert x.src[1].shape == x.src[0].shape, f"shape mismatch on store {x.src[1].shape} != {x.src[0].shape}"
+  idx, valid = x.st_arg.to_indexed_uops(ctx.idxs)
+  if cast(PtrDType, buf.dtype).addrspace == AddrSpace.GLOBAL:
+    # NOTE: only store the local reduceop in the threads that are actually doing the reduce
+    for oidx, ridx in zip(ctx.idxs, ctx.ridxs):
+      if oidx is not ridx: valid = valid * oidx.eq(0)
+  return buf.index(idx, valid).store(x.src[1], *[x for x in UOp.sink(idx, valid).toposort() if x.op is Ops.RANGE])

-  new_idxs = shape_to_idx(x.src[0].shape, ctx.axis_types, ctx.start)
-  idx, valid = x.st_arg.to_indexed_uops(new_idxs)
-  used_idxs = [x for x in UOp.sink(idx, valid).toposort() if x in new_idxs]
-  real_new_idxs = []
-  for i in range(len(x.src[0].shape)):
-    if new_idxs[i] in used_idxs or len(ctx.idxs) <= i: real_new_idxs.append(new_idxs[i])
-    else: real_new_idxs.append(ctx.idxs[i])
-
-  stored = subblock(ctx, real_new_idxs, x.src[1])
-  used_ranges = [x for x in used_idxs if x.op is Ops.RANGE]
-  ret = buf.index(idx, valid).store(stored, *used_ranges)
-
-  # insert BARRIER if we are ending a LOCAL, IF if we are ending a GROUP_REDUCE
-  if cast(PtrDType, buf.dtype).addrspace == AddrSpace.LOCAL and \
-      any(ctx.axis_types[x.arg%1000] in {AxisType.GROUP_REDUCE, AxisType.LOCAL} for x in used_ranges):
-    ret = ret.barrier()
-    range_gates = [x.eq(0) for x in used_ranges if ctx.axis_types[x.arg%1000] == AxisType.GROUP_REDUCE]
-    if len(range_gates): ret = UOp(Ops.IF, src=(functools.reduce(operator.and_, range_gates), ret))
-  return ret
-
-def fixup_wmma(ctx:IndexContext, x:UOp):
-  if x.tag is not None: return None
-  new_idxs = shape_to_idx(x.src[0].shape, ctx.axis_types, ctx.start)
-  full_new_idx = list(ctx.idxs)
-  for a in x.arg[-1]: full_new_idx[a] = new_idxs[a]
-
-  srcs = subblock(ctx, full_new_idx, UOp.sink(*x.src)).src
-
-  # NOTE: this assumes these are expanded. which now shouldn't change anything
-  new_x_arg_m2 = tuple([tuple([(full_new_idx[a].arg[0][0], sz) for a,sz in v]) for v in x.arg[-2]])
-  new_x_arg_m1 = tuple([full_new_idx[a].arg[0][0] for a in x.arg[-1]])
-  return x.replace(src=srcs, arg=x.arg[:-2]+(new_x_arg_m2, new_x_arg_m1), tag=1)
+def lower_const(ctx:IndexContext, view:UOp, c:UOp):
+  if all(x.mask is None for x in view.arg.views): return c
+  _, valid = view.arg.to_indexed_uops(ctx.idxs)
+  return valid.where(c, c.const_like(0))

 pm_lowerer = PatternMatcher([
  # TODO: remove these hacks
@@ -97,18 +71,10 @@ pm_lowerer = PatternMatcher([
  # hack for old style VALID (now it's just VIEW(CONST))
  (UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c"), UPat(Ops.CONST, arg=0)), lambda c,v: c.replace(src=()).view(v.arg)),

-  # consts and loads
-  (UPat(Ops.VIEW, src=(UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"),), name="view"),
-   lambda ctx,view,c: c if all(x.mask is None for x in view.arg.views) else view.arg.to_indexed_uops(ctx.idxs)[1].where(c, c.const_like(0))),
-  (UPat(Ops.LOAD, src=(UPat.var("buf").view(),), allow_any_len=True, name="x"),
-   lambda ctx,buf,x: UOp(Ops.LOAD, x.dtype, (buf.index(*x.st_arg.to_indexed_uops(ctx.idxs)),)+x.src[1:])),
-
  # reduce/view_const
  (UPat(Ops.REDUCE_AXIS, name="x"), lower_reduce_axis),
+  (UPat(Ops.VIEW, src=(UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"),), name="view"), lower_const),
+  # rewrite LOAD/STORE VIEW to LOAD/STORE with indexed
+  (UPat(Ops.LOAD, src=(UPat.var("buf").view(),), allow_any_len=True, name="x"), lower_load),
  (UPat(Ops.STORE, src=(UPat.var("buf").view(),), allow_any_len=True, name="x"), lower_store),
-  (UPat(Ops.WMMA, name="x"), fixup_wmma),
-
-  # axis fixups for WMMA
-  (UPat((Ops.CONTRACT, Ops.UNROLL), name="x"),
-   lambda ctx,x: x.replace(tag=1, arg=tuple([(ctx.idxs[a].arg[0][0], sz) for a,sz in x.arg])) if x.tag is None else None),
 ])
@@ -1,58 +0,0 @@
-# should this merge with transcendental?
-from typing import Callable
-import functools
-from tinygrad.device import is_dtype_supported
-from tinygrad.dtype import dtypes, promo_lattice
-from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher
-from tinygrad.helpers import getenv
-from tinygrad.uop.transcendental import xexp2, xlog2, xsin, xpow, TRANSCENDENTAL_SUPPORTED_DTYPES
-from tinygrad.renderer import Renderer
-
-# ***** optional patterns *****
-
-@functools.lru_cache(None)
-def magicgu(vmax:int, d:int) -> tuple[int,int]:
-  # calculate m,s such that x//d == (x*m) >> s for all 0 <= x <= vmax, d>0; adapted from Hacker's Delight, Chapter 10
-  nc = (vmax+1)//(d) * d - 1
-  nbits = vmax.bit_length()
-  for s in range(0, 2*nbits + 1):
-    if 2**s > nc*(d - 1 - (2**s - 1) % d):
-      m = (2**s + d - 1 - (2**s - 1) % d)//d
-      return m, s
-  assert False
-
-def fast_idiv(ctx: Renderer|None, x: UOp, d: int) -> UOp|None:
-  # idiv is truncated division, but arithmetic shift is floored division, so can only do non-negative numbers!
-  if x.vmin<0: return None
-  sign = 1 if d > 0 else -1
-  m,s = magicgu(vmax := min(x.vmax, dtypes.max(x.dtype)), abs(d))
-  if m * vmax <= dtypes.max(x.dtype): return sign * ((x*m) >> s)
-  # promo_lattice needs to return an unsigned type
-  if ctx is not None and dtypes.is_int(next_dtype := promo_lattice[x.dtype][-1]) and is_dtype_supported(next_dtype, ctx.device):
-    if m * vmax <= dtypes.max(next_dtype): return sign * ((x.cast(next_dtype)*m) >> s).cast(x.dtype)
-  return None
-
-powers_of_two = {2**i:i for i in range(64)}
-@functools.cache
-def get_late_rewrite_patterns(ops, force_transcendental=False):
-  pat: list[tuple[UPat, Callable]] = [(UPat(op, dtype=TRANSCENDENTAL_SUPPORTED_DTYPES, src=(UPat.var("d"),)), f) for op,f in \
-           ((Ops.EXP2, xexp2), (Ops.LOG2, xlog2), (Ops.SIN, xsin)) if op not in ops or force_transcendental]
-  # rewrite SQRT to xpow 0.5
-  if Ops.SQRT not in ops: pat.append((UPat(Ops.SQRT, src=UPat.var("d")), lambda d: xpow(d, d.const_like(0.5))))
-  # rewrite MOD to AND (which should always be supported, but not for generic in tests): x % (2**y) -> x & (2**y-1)
-  if Ops.AND in ops: pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.arg-1) if c.arg in powers_of_two else None)]
-  # rewrite MUL/IDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y)
-  if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.arg, 0)) else None)]
-  if Ops.SHR in ops:
-    # no reason to check x<0 for uints
-    pat += [(UPat.var("x", dtypes.uints)//UPat.cvar("c"), lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]
-    pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("c"), lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(
-      c-1, 0)) >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]  # (x+(x<0).where(c-1, 0)) >> v
-    if not getenv("DISABLE_FAST_IDIV"):
-      pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("d"), lambda ctx, x, d: fast_idiv(ctx, x, d.arg))]
-      pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("d"), lambda ctx, x, d: x - d*f if (f:=fast_idiv(ctx, x, d.arg)) is not None else None)]
-  if Ops.NEG in ops:
-    pat += [(UPat.var('x')*-1, lambda x: x.alu(Ops.NEG))]
-    if Ops.SUB in ops: pat += [(UPat.var('x')+UPat.var('y').alu(Ops.NEG), lambda x,y: x.alu(Ops.SUB, y))]
-  if Ops.MULACC in ops: pat += [(UPat.var('a')*UPat.var('b')+UPat.var('c'), lambda a,b,c: a.alu(Ops.MULACC, b, c))]
-  return PatternMatcher(pat)
@@ -4,7 +4,7 @@ from collections import defaultdict
 from typing import Any, Generic, TypeVar, Iterator
 import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal, time
 from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, \
-                             colored, Context, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, cpu_events, ProfileEvent, dedup
+                             colored, Context, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, cpu_events, ProfileEvent
 from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
 from tinygrad.renderer import Renderer

@@ -37,8 +37,7 @@ class _Device:
      with contextlib.suppress(Exception): yield self[device].device
  @functools.cached_property
  def DEFAULT(self) -> str:
-    dev = [dev] if (dev:=getenv("DEV", "").upper()) else []
-    from_env = dedup(dev + [d for d in self._devices if d not in ["DISK", "NPY"] and getenv(d) == 1])
+    from_env = [d for d in self._devices if d not in ["DISK", "NPY"] and getenv(d) == 1]
    assert len(from_env) < 2, f"multiple devices set in env: {from_env}"
    if len(from_env) == 1: return from_env[0]
    try:
@@ -336,7 +335,7 @@ if PROFILE:

    if not getenv("SQTT", 0):
      from tinygrad.uop.ops import launch_viz
-      launch_viz(PROFILE, fn)
+      launch_viz("PROFILE", fn)

 if __name__ == "__main__":
  for device in ALL_DEVICES:
@@ -42,13 +42,17 @@ def apply_graph_to_jit(jit_cache: list[ExecItem], input_rawbuffers: list[Buffer]

  for ji in jit_cache:
    match ji.prg:
-      case CompiledRunner(): ji_graph_dev = ji.prg.dev
-      case BufferXfer(): ji_graph_dev = Device[unwrap(ji.bufs[0]).device]
-      case BufferCopy(): ji_graph_dev = next((Device[unwrap(b).device] for b in ji.bufs if unwrap(b).device not in {"CPU", "LLVM"}), None)
+      case CompiledRunner():
+        ji_graph_dev = ji.prg.dev
+        # All GraphRunners can graph CompiledRunners
+        can_be_graphed = ji_graph_dev.graph is not None
+      case BufferXfer():
+        ji_graph_dev = Device[unwrap(ji.bufs[0]).device]
+        # All *Multi*GraphRunner support graphing BufferXfers
+        can_be_graphed = ji_graph_dev.graph is not None and issubclass(graph_class(ji_graph_dev), MultiGraphRunner)
      case ViewOp(): continue # ViewOps are just ignored
-      case _: ji_graph_dev = None # Everything else is not graphed and flushes existing graph if it's being constructed
+      case _: can_be_graphed = False # Everything else is not graphed and flushes existing graph if it's being constructed

-    can_be_graphed = ji_graph_dev is not None and ji_graph_dev.graph is not None and graph_class(ji_graph_dev).supports_exec_item(ji_graph_dev, ji)
    is_multigraph = can_be_graphed and issubclass(graph_class(ji_graph_dev), MultiGraphRunner)
    can_share_graph = can_be_graphed and (type(ji_graph_dev) is type(current_device) if is_multigraph else ji_graph_dev == current_device)
    can_extend_graph_batch = can_share_graph and (max_batch_size == 0 or len(current_batch) < max_batch_size)
@@ -126,13 +130,8 @@ class GraphRunner(Runner):

    return list({id(x):x for x in wait_nodes}.values())

-  @staticmethod
-  def supports_exec_item(dev, ei:ExecItem) -> bool: return isinstance(ei.prg, CompiledRunner)
-
 # a marker for your graph supporting multiple devices of the same type
-class MultiGraphRunner(GraphRunner):
-  @staticmethod
-  def supports_exec_item(dev, ei:ExecItem) -> bool: return isinstance(ei.prg, (CompiledRunner, BufferXfer))
+class MultiGraphRunner(GraphRunner): pass

 def get_out_buffers_for_ei(ei:ExecItem) -> list[Buffer]:
  if isinstance(ei.prg, CompiledRunner): return [cast(Buffer, ei.bufs[out]) for out in ei.prg.p.outs if out not in ei.prg.p.ins]
@@ -13,7 +13,7 @@ from tinygrad.uop.spec import type_verify

 # **************** Program Creation ****************

-@track_rewrites(name=lambda _ast,_renderer,ret: TracingKey(ret.name, (ret.function_name, ret.ast), ret.src, ret=ret))
+@track_rewrites(name=lambda _ast,_renderer,ret: TracingKey(ret.name, (ret.function_name, ret.ast), ret.src))
 def get_program(ast:UOp, renderer:Renderer) -> ProgramSpec:
  """
  Transform an AST into a ProgramSpec. May trigger BEAM search.
@@ -156,7 +156,7 @@ class ExecItem:
        lds_est = sym_infer(self.prg.estimates.lds, var_vals)
        mem_est = min(mem_est, lds_est)   # there can't be more memory accessed than loads/stores. remove this when symbolic is fixed
        ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
-        print(f"{colored(f'*** {self.prg.device[:7]:7s} {GlobalCounters.kernel_count:4d}', 'magenta' if jit else ('green' if self.prg.first_run else None))} {self.prg.display_name+' '*(44-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:5.2f} GB " +  # noqa: E501
+        print(f"{colored(f'*** {self.prg.device[:7]:7s} {GlobalCounters.kernel_count:4d}', 'magenta' if jit else ('green' if self.prg.first_run else None))} {self.prg.display_name+' '*(41-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:5.2f} GB " +  # noqa: E501
              (str() if et is None else f"tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({op_est/((et or 1e-20)*1e9):9.2f} GFLOPS {mem_est/((et or 1e-20)*1e9):6.1f}|{lds_est/((et or 1e-20)*1e9):<7.1f} GB/s)" +  # noqa: E501
               f" {[repr(m) if TRACEMETA >= 2 else str(m) for m in self.metadata] if self.metadata else ''}"))
      self.prg.first_run = False
@@ -81,13 +81,6 @@ def word_wrap(x, wrap=80):
  while len(ansistrip(x[:i])) < wrap and i < len(x): i += 1
  return x[:i] + "\n" + word_wrap(x[i:], wrap)

-def suppress_finalizing(func):
-  def wrapper(*args, **kwargs):
-    try: return func(*args, **kwargs)
-    except (AttributeError, TypeError, ImportError):
-      if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing
-  return wrapper
-
 def pluralize(st:str, cnt:int): return f"{cnt} {st}"+('' if cnt == 1 else 's')

 class LazySeq(Generic[T]): # NOTE: Mapping requires __iter__ and __len__, Sequence requires supporting __len__ and slicing in __getitem__
@@ -198,7 +191,6 @@ class TracingKey:
  keys:tuple[str, ...]=()                # optional keys to search for related traces
  fmt:str|None=None                      # optional detailed formatting
  cat:str|None=None                      # optional category to color this by
-  ret:Any=None

 class ProfileEvent: pass

@@ -5,7 +5,7 @@ from collections import defaultdict
 from typing import cast, Final, Callable, Sequence
 from enum import Enum, auto

-from tinygrad.uop.ops import GroupOp, KernelInfo, UOp, Ops, can_pad, resolve, Variable, sint, graph_rewrite, AxisType
+from tinygrad.uop.ops import GroupOp, KernelInfo, UOp, Ops, can_pad, resolve, Variable, sint, graph_rewrite, smax, AxisType
 from tinygrad.uop.spec import type_verify, ast_spec
 from tinygrad.device import Device
 from tinygrad.opt.tc import TensorCore
@@ -73,8 +73,7 @@ class Kernel:
      self.sts.append(unwrap(x.src[0].st))

    # add a shapetracker to the end to track the full shape, with 0 strides so it can merge
-    full_shape = ast.full_shape
-    self.sts.append(ShapeTracker.from_shape(full_shape, (0,)*len(full_shape)))
+    self.sts.append(ShapeTracker.from_shape(tuple([smax(*s) for s in zip(*[x.shape for x in self.sts])]), (0,)*len(self.sts[0].shape)))

    # parameters for optimization
    self.tensor_core: TensorCore|None = None
@@ -448,7 +447,9 @@ class Kernel:
      ret = op.replace(src=tuple(fixup_ast(x) for x in op.src)) # noqa: F821
      if op.op in GroupOp.Buffer and op in self.bufs:
        st = self.sts[self.bufs.index(op)]
-        # replace the VIEW source
+        # NOTE: if CONST got masked after applying opts, we create a new VALID
+        if op.op is Ops.CONST and any(v.mask is not None for v in st.views): return op.view(st).valid()
+        # otherwise we just replace the VIEW source
        return ret.replace(src=(ret.src[0].replace(arg=st),)+ret.src[1:])
      if op.op is Ops.SINK:
        # NOTE: should group_for_reduces be added to the local_dims?
@@ -462,7 +463,8 @@ class Kernel:
        if (tc := self.tensor_core) and self.use_tensor_cores == 1:
          # get reduce/upcast axes for the tensor cores
          tc_reduce_axes = self.shape_str_to_axis([f"r{i}" for i in range(len(tc.get_reduce_axes()))])
-          base_upcast_axes = tuple([(s,2) for s in self.shape_str_to_axis(tc.base_upcast_axes())])
+          base_upcast_axes = tuple([(s,2) for s in self.shape_str_to_axis([f"r{i}" for i in range(len(tc.get_reduce_axes()))] + \
+                                                                          [f"u{i}" for i in range(len(tc.get_upcast_axes()))])])[::-1]
          tc_upcast_axes = tuple([base_upcast_axes[:int(math.log2(tc.elements_per_thread[i]))] for i in range(3)])

          # permute the srcs
@@ -2,7 +2,7 @@ from typing import cast, Callable
 import itertools, functools, random, math, time, multiprocessing, traceback, signal, atexit
 from collections import defaultdict
 from dataclasses import replace
-from tinygrad.uop.ops import UOp, Ops, Variable, sym_infer, AxisType
+from tinygrad.uop.ops import UOp, Ops, Variable, sym_infer
 from tinygrad.device import Device, Buffer, Compiler
 from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str
 from tinygrad.helpers import IGNORE_BEAM_CACHE, TC_SEARCH_OVER_SHAPE
@@ -83,7 +83,7 @@ def _try_compile_linearized_w_idx(x:tuple[int,Kernel], compiler:Compiler) -> tup

 # workers should not open devices and should ignore ctrl c and should not launch VIZ
 def _init_worker():
-  Context(ALLOW_DEVICE_USAGE=0, VIZ=0, TRACK_MATCH_STATS=0).__enter__()
+  Context(ALLOW_DEVICE_USAGE=0, VIZ=0).__enter__()
  signal.signal(signal.SIGINT, signal.SIG_IGN)

 def _ensure_buffer_alloc(bufs:list[Buffer]) -> list[Buffer]: return [buf.ensure_allocated() if buf is not None else buf for buf in bufs]
@@ -108,9 +108,9 @@ def bufs_from_lin(lin:Kernel, allocate:bool=True) -> list[Buffer]:
  return cast(list[Buffer], rawbufs)

 # get dictionary of all possible actions
-def get_kernel_actions(lin:Kernel, include_0=True, candidates:list[Opt]|None=None) -> dict[int, Kernel]:
+def get_kernel_actions(lin:Kernel, include_0=True) -> dict[int, Kernel]:
  acted_lins, max_up, max_lcl = {0:lin} if include_0 else {}, getenv("BEAM_UPCAST_MAX", 256), getenv("BEAM_LOCAL_MAX", 1024)
-  kernel_actions = (actions if candidates is None else candidates).copy()
+  kernel_actions = actions.copy()

  if TC_SEARCH_OVER_SHAPE and len(lin.applied_opts) == 0: # tensor core opts must be first
    for i, action in enumerate(kernel_actions):
@@ -123,14 +123,14 @@ def get_kernel_actions(lin:Kernel, include_0=True, candidates:list[Opt]|None=Non
    if a.axis is not None and a.op is not OptOps.TC:
      try: ax = lin.real_axis(a.op, a.axis)
      except KernelOptError: continue
-      if (ax >= lin.shape_len) or (lin.full_shape[ax] == a.arg and Opt(a.op, a.axis, 0) in kernel_actions): continue
+      if (ax >= lin.shape_len) or (lin.full_shape[ax] == a.arg and Opt(a.op, ax, 0) in kernel_actions): continue
    lin2 = lin.copy()
    try:
      lin2.apply_opt(a)
      up, lcl, tc_up = 1, 1, prod(tc.dims)//tc.threads if (tc:=lin2.tensor_core) else 1
-      for s,c in zip(lin2.full_shape, lin2.axis_types):
-        if c in (AxisType.UPCAST, AxisType.UNROLL): up *= s
-        elif c in (AxisType.LOCAL, AxisType.GROUP_REDUCE): lcl *= s
+      for s,c in zip(lin2.full_shape, lin2.colors()):
+        if c in {"magenta", "yellow"}: up *= s
+        elif c in {"cyan", "green", "white"}: lcl *= s
      if up//tc_up > max_up or lcl > max_lcl:
        if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too many upcast/local. {up//tc_up=}, {max_up=}, {lcl=}, {max_lcl=}")
        continue
@@ -25,9 +25,6 @@ class TensorCore: # D = A * B + C, A is (M x K), B is (K x N), C and D are (M x
  def get_reduce_axes(self): return [(i, 2) for i in range(int(math.log2(self.dims[2])))]
  def get_upcast_axes(self): return [opt for opt in self.opts if opt[0] == "u"]
  def get_local_axes(self): return [opt for opt in self.opts if opt[0] == "l"]
-  def base_upcast_axes(self):
-    # this is defined in the swizzle. first we use the upcast axes, then the reduce
-    return ([f"r{i}" for i in range(len(self.get_reduce_axes()))] + [f"u{i}" for i in range(len(self.get_upcast_axes()))])[::-1]
  def __str__(self): return "_".join(["WMMA"] + list(map(str, self.dims)) + [self.dtype_in.name, self.dtype_out.name])
  def __post_init__(self):
    # all axes have size 2, <local> <reduce> <upcast> is the order
@@ -37,30 +34,12 @@ class TensorCore: # D = A * B + C, A is (M x K), B is (K x N), C and D are (M x
    assert 2**local_axes == self.threads, f"{self.threads} threads construct the warp but found {2**local_axes} in {self.opts}"
    assert 2**upcast_axes == self.elements_per_thread[2], \
      f"{self.elements_per_thread[2]} elements from C are processed per thread but found {2**upcast_axes} in {self.opts}"
-    # check dims match opts
-    assert self.dims[0] == 2**len(gd:=[x for x in self.opts if x[1] == '0']), f"opts wrong on dims[0], {self.dims[0]} vs {gd}"
-    assert self.dims[1] == 2**len(gd:=[x for x in self.opts if x[1] == '1']), f"opts wrong on dims[1], {self.dims[1]} vs {gd}"
-    # NOTE: the K opts is implictly set by the dim
    # check swizzle
    assert len(self.swizzle[0]) == 3 and len(self.swizzle[1]) == 3, "swizzle has wrong part count"
    assert len(self.swizzle[0][0]) == len(self.swizzle[1][0]) == local_axes, "local swizzle size is wrong"
    assert len(self.swizzle[0][1]) == len(self.swizzle[1][1]) == upcast_axes, "upcast swizzle size is wrong"
    assert len(self.swizzle[0][2]) == len(self.swizzle[1][2]) == reduce_axes, "reduce swizzle size is wrong"
    assert all(len(s) == local_axes+upcast_axes+reduce_axes for s in self._remaps()), "remaps are the wrong size"
-    # check elements_per_thread
-    un, ln = 0, 0
-    zero_stride_0 = []
-    zero_stride_1 = []
-    for o in self.opts:
-      if o[1] == '0': zero_stride_0.append(o[0] + str(un if o[0] == 'u' else ln))
-      if o[1] == '1': zero_stride_1.append(o[0] + str(un if o[0] == 'u' else ln))
-      if o[0] == 'u': un += 1
-      if o[0] == 'l': ln += 1
-    # NOTE: all the zero_stride dims can be placed in any order in the swizzle
-    upcasted_0 = [x for x in (self.swizzle[0][1] + self.swizzle[0][2]) if x not in zero_stride_0 and x[0] != 'l']
-    upcasted_1 = [x for x in (self.swizzle[1][1] + self.swizzle[1][2]) if x not in zero_stride_1 and x[0] != 'l']
-    assert 2**len(upcasted_0) == self.elements_per_thread[0], f"mismatch in elements_per_thread[0], {upcasted_0} vs {self.elements_per_thread[0]}"
-    assert 2**len(upcasted_1) == self.elements_per_thread[1], f"mismatch in elements_per_thread[1], {upcasted_1} vs {self.elements_per_thread[1]}"

 # ***** NVIDIA *****

@@ -86,14 +65,12 @@ cuda_sm75: list[TensorCore] = cuda_8168_f16

 # https://gpuopen.com/learn/wmma_on_rdna3/
 amd_rdna3 = [TensorCore(dims=(16,16,16), threads=32, elements_per_thread=(16,16,8), dtype_in=di, dtype_out=do,
-  opts=("l0","l0","l0","l0","l1","u1","u1","u1"),
-  swizzle=((('l4', 'u0', 'u1', 'u2', 'l0'), ('r1', 'r2', 'r3'), ('l1', 'l2', 'l3', 'r0')),
-           (('l0', 'l1', 'l2', 'l3', 'l4'), ('r1', 'r2', 'r3'), ('u0', 'u1', 'u2', 'r0'))))
+  opts=("l0","l0","l0","l0","l1","u1","u1","u1"), swizzle=((('l4', 'u0', 'u1', 'u2', 'l0'), ('r1', 'r2', 'r3'), ('l1', 'l2', 'l3', 'r0')),
+                                                           (('l0', 'l1', 'l2', 'l3', 'l4'), ('r1', 'r2', 'r3'), ('u0', 'u1', 'u2', 'r0'))))
  for di,do in [(dtypes.half,dtypes.float),(dtypes.half,dtypes.half),(dtypes.bfloat16,dtypes.float)]]
 amd_rdna4 = [TensorCore(dims=(16,16,16), threads=32, elements_per_thread=(8,8,8), dtype_in=di, dtype_out=do,
-  opts=("l0","l0","l0","l0","u1","u1","u1","l1"),
-  swizzle=((('u0', 'u1', 'u2', 'l4', 'r2'), ('r0', 'r1', 'r3'), ('l0', 'l1', 'l2', 'l3')),
-           (('l0', 'l1', 'l2', 'l3', 'r2'), ('r0', 'r1', 'r3'), ('l4', 'u0', 'u1', 'u2'))))
+  opts=("l0","l0","l0","l0","u1","u1","u1","l1"), swizzle=((('u0', 'u1', 'u2', 'l4', 'r2'), ('r0', 'r1', 'r3'), ('l0', 'l1', 'l2', 'l3')),
+                                                           (('l0', 'l1', 'l2', 'l3', 'r2'), ('r0', 'r1', 'r3'), ('l4', 'u0', 'u1', 'u2'))))
  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
@@ -9,7 +9,7 @@ from tinygrad.renderer import Renderer
 from tinygrad.codegen.devectorizer import no_vectorized_alu

 base_rewrite = PatternMatcher([
-  (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.dtype.size}];"),
+  (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.dtype.size}] = {{{ctx[x.src[0]]}}};"),
  (UPat(Ops.IF, name="x"), lambda ctx,x: f"if ({ctx[x.src[0]]}) {{"),
  (UPat((Ops.ENDIF, Ops.ENDRANGE)), lambda ctx: "}"),
  (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{x.arg[0]}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]})"),
@@ -25,7 +25,7 @@ base_rewrite = PatternMatcher([
  (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"(*(({ctx.buffer_prefix}{ctx.render_dtype(x.dtype)}*)&{ctx[x.src[0]]}))"),
  (UPat(Ops.DEFINE_LOCAL, name="x"), lambda ctx,x: f"{ctx.smem_align}{ctx.smem_prefix}{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.dtype.size}];"),
  (UPat(Ops.BARRIER), lambda ctx: ctx.barrier),
-  (UPat(Ops.PRECAST, name="x"), lambda ctx,x: ctx[x.src[0]]),
+  (UPat(Ops.NOOP, name="x"), lambda ctx,x: ctx[x.src[0]]),
  (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"{ctx.code_for_workitem[x.arg[0][0]](x.arg[0][-1])}; /* {x.arg[1]} */"),
  # const
  (UPat(Ops.CONST, arg=math.inf, name="x"), lambda ctx, x: f"({ctx.render_cast(x.dtype, ctx.infinity)})"),
@@ -47,7 +47,7 @@ base_rewrite = PatternMatcher([
   lambda ctx,buf,idx: f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})"),
  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat.var("gate"))).or_casted("bidx"), UPat.var("var")), allow_any_len=True),
   lambda ctx,bidx,var,gate: f"({ctx[gate]}?*{ctx[bidx]}:{ctx[var]})"),
-  (UPat(Ops.LOAD, src=(UPat.var('bidx'),), allow_any_len=True), lambda ctx,bidx: f"(*{ctx[bidx]})"),
+  (UPat(Ops.LOAD, src=(UPat.var('bidx'),), allow_any_len=True), lambda ctx,bidx: f"*{ctx[bidx]}"),
  (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var")), allow_any_len=True), lambda ctx,bidx,var: f"*{ctx[bidx]} = {ctx[var]};"),
  # alu/gep
  # TODO: look for left-associative
@@ -60,9 +60,9 @@ base_rewrite = PatternMatcher([
 ])

 extra_pm = PatternMatcher([
-  # insert a PRECAST before BITCAST to force it to be rendered. not needed on all backends?
-  (UPat(Ops.BITCAST, name="x"), lambda x: UOp(Ops.BITCAST, x.dtype, (UOp(Ops.PRECAST, x.src[0].dtype, x.src),))
-   if x.src[0].op not in {Ops.PRECAST, Ops.LOAD, Ops.CUSTOM} else None),
+  # insert a NOOP before BITCAST to force it to be rendered. not needed on all backends?
+  (UPat(Ops.BITCAST, name="x"),
+   lambda x: UOp(Ops.BITCAST, x.dtype, (UOp(Ops.NOOP, x.src[0].dtype, x.src),)) if x.src[0].op not in {Ops.NOOP, Ops.LOAD, Ops.CUSTOM} else None),
  # rewrite MAX to CMPLT + WHERE (max function is annoying on many cstyle backends)
  (UPat(Ops.MAX, name="m"), lambda m: (m.src[0] < m.src[1]).where(m.src[1], m.src[0])),
  # devectorize any bools
@@ -75,10 +75,6 @@ extra_pm = PatternMatcher([

 def uops_to_dtypes(uops:list[UOp]) -> list[DType]: return dedup(u.dtype for u in uops if not isinstance(u.dtype, (ImageDType, PtrDType)))

-# (name, dims, dtype_in, dtype_out, device, threads, upcast_axes, reduce_axes)
-def wmma_args(uops:list[UOp]):
-  return dedup((uop.arg[0], uop.arg[1], uop.src[0].dtype.scalar(), uop.dtype.scalar(), *(uop.arg[4:8])) for uop in uops if uop.op is Ops.WMMA)
-
 class CStyleLanguage(Renderer):
  kernel_typedef: str = "void"
  buffer_prefix: str = ""
@@ -122,9 +118,7 @@ class CStyleLanguage(Renderer):
  def render_dtype(self, dt:DType, mutable=True) -> str:
    if isinstance(dt, ImageDType): return f"{'write_only' if mutable else 'read_only'} image2d_t"
    if isinstance(dt, PtrDType):
-      prefix = ""
-      if dt.addrspace == AddrSpace.LOCAL and self.smem_prefix_for_cast: prefix = self.smem_prefix
-      if dt.addrspace == AddrSpace.GLOBAL: prefix = self.buffer_prefix
+      prefix = self.smem_prefix if dt.addrspace == AddrSpace.LOCAL and self.smem_prefix_for_cast else self.buffer_prefix
      return prefix + self.render_dtype(dt.base) + "*"
    if dt.count > 1: return self.type_map.get(scalar:=dt.scalar(), scalar.name).replace(" ", "_") + str(dt.count)
    return self.type_map.get(scalar:=dt.scalar(), scalar.name)
@@ -141,9 +135,6 @@ class CStyleLanguage(Renderer):
    c: defaultdict[str, int] = defaultdict(int)
    name = "test"
    for u in uops:
-      if u.op is Ops.NOOP:
-        if len(u.src): r[u] = r[u.src[0]]
-        continue
      if u.op is Ops.SINK:
        if u.arg is not None: name = u.arg.function_name
        continue
@@ -163,7 +154,7 @@ class CStyleLanguage(Renderer):
      elif u.op is Ops.RANGE: r[u] = f"ridx{u.arg}"
      else:
        prefix = {Ops.WMMA: "wmma", Ops.DEFINE_LOCAL: "temp", Ops.CONST: "const",
-                  Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.GEP: "gep", Ops.VECTORIZE: "cast", Ops.PRECAST: "precast",
+                  Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.GEP: "gep", Ops.VECTORIZE: "cast", Ops.NOOP: "precast",
                  Ops.INDEX: "bidx", Ops.DEFINE_REG: "acc", Ops.LOAD: "val"}.get(u.op, "alu")
        r[u] = f"{prefix}{c[prefix]}"

@@ -172,13 +163,13 @@ class CStyleLanguage(Renderer):

      if u.op in {Ops.ENDIF, Ops.ENDRANGE}: depth -= 1
      if (u.op is not Ops.CAST or u.dtype.vcount == 1) and (u.op in {Ops.CONST, Ops.GEP, Ops.INDEX, Ops.CUSTOMI} or \
-        (u.op is Ops.LOAD and cast(PtrDType, u.src[0].dtype).addrspace == AddrSpace.REG) or \
-        (u.op is Ops.CAST and isinstance(u.dtype, PtrDType)) or \
        (u.op in {Ops.VECTORIZE, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))):
        r[u] = l
      else:
-        if u.op in {Ops.RANGE, Ops.DEFINE_LOCAL, Ops.STORE, Ops.DEFINE_REG} or u.dtype == dtypes.void: pass
-        else: l = f"{self.render_dtype(u.dtype)} {r[u]} = {l}" + (";" if u.op is not Ops.SPECIAL else "")
+        if u.op in {Ops.RANGE, Ops.DEFINE_LOCAL, Ops.STORE, Ops.DEFINE_REG} or u.dtype == dtypes.void:
+          if u.op is Ops.STORE: r[u] = r[u.src[0]]
+        else:
+          l = f"{self.render_dtype(u.dtype)} {r[u]} = {l}" + (";" if u.op is not Ops.SPECIAL else "")
        kernel.append("  "*depth + l)
        if prefix: c[prefix] += 1  # if it was used, increment
      if u.op in {Ops.IF, Ops.RANGE}: depth += 1
@@ -218,7 +209,7 @@ class ClangRenderer(CStyleLanguage):
  def _render_defines(self, uops) -> list[str]:
    prefix = [self.render_vector_prefix(dt) for dt in uops_to_dtypes(uops) if dt.count > 1]
    # https://github.com/corsix/amx
-    for name, (N, M, _), dtype_in, _, _, _, _, _ in wmma_args(uops):
+    for name, (N, M, _), dtype_in, _, _, _, _, _ in dedup([uop.arg for uop in uops if uop.op is Ops.WMMA]):
      prefix += [
        '#define AMX_SET(imm5) __asm("nop\\nnop\\nnop\\n.word (0x201000+(%0<<5)+%1)" : : "i"(17), "i"(imm5) : "memory")',
        '#define AMX(op, gpr, btf) __asm(".word (0x201000+(%0 << 5)+0%1-((0%1>>4)*6))" : : "i"(op), "r"((unsigned long long)(gpr)+(btf)) : "memory")',
@@ -278,9 +269,9 @@ class IntelRenderer(OpenCLRenderer):

  def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str:
    prefix = []
-    for name, _, dtype_in, dtype_out, _, _, _, _ in wmma_args(uops):
-      dt_in = ("ushort", "bf16") if dtype_in == dtypes.bfloat16 else (dtype_in.name, "f16")
-      prefix.append(f"""{dtype_out.name}8 __{name}({dt_in[0]}16 a, {dt_in[0]}16 b, {dtype_out.name}8 c) {{
+    for arg in dedup([uop.arg for uop in uops if uop.op is Ops.WMMA]):
+      dt_in = ("ushort", "bf16") if arg[2] == dtypes.bfloat16 else (arg[2].name, "f16")
+      prefix.append(f"""{arg[3].name}8 __{arg[0]}({dt_in[0]}16 a, {dt_in[0]}16 b, {arg[3].name}8 c) {{
    return intel_sub_group_{dt_in[1]}_{dt_in[1]}_matrix_mad_k16(as_int8(a), as_int8(b), c);\n}}""")
    return super().render_kernel(function_name, kernel, bufs, uops, prefix or None)

@@ -316,13 +307,13 @@ class MetalRenderer(CStyleLanguage):
  ]) + base_rewrite

  def render_kernel(self, function_name, kernel, bufs, uops, prefix=None):
-    prefix = ["#include <metal_stdlib>","using namespace metal;"]
-    for name, _, dtype_in, dtype_out, _, _, _, _ in wmma_args(uops): prefix.append(
-  f"""{(dstr_out:=self.render_dtype(dtype_out.vec(2)))} __{name}({(dstr_in:=self.render_dtype(dtype_in.vec(2)))} a, {dstr_in} b, {dstr_out} c){{
-  simdgroup_{self.render_dtype(dtype_in)}8x8 mat_a, mat_b; simdgroup_{self.render_dtype(dtype_out)}8x8 mat_c;
+    prefix, wmma_args = ["#include <metal_stdlib>","using namespace metal;"], set([uop.arg for uop in uops if uop.op is Ops.WMMA])
+    for arg in wmma_args: prefix.append(
+  f"""{(dtype_out:=self.render_dtype(arg[3].vec(2)))} __{arg[0]}({(dtype_in:=self.render_dtype(arg[2].vec(2)))} a, {dtype_in} b, {dtype_out} c){{
+  simdgroup_{self.render_dtype(arg[2])}8x8 mat_a, mat_b; simdgroup_{self.render_dtype(arg[3])}8x8 mat_c;
  mat_a.thread_elements()[0] = a[0]; mat_b.thread_elements()[0] = b[0]; mat_c.thread_elements()[0] = c[0];
  mat_a.thread_elements()[1] = a[1]; mat_b.thread_elements()[1] = b[1]; mat_c.thread_elements()[1] = c[1];
-  simdgroup_multiply_accumulate(mat_c, mat_a, mat_b, mat_c);\n  return {dstr_out}(mat_c.thread_elements()[0], mat_c.thread_elements()[1]);\n}}""")
+  simdgroup_multiply_accumulate(mat_c, mat_a, mat_b, mat_c);\n  return {dtype_out}(mat_c.thread_elements()[0], mat_c.thread_elements()[1]);\n}}""")
    return super().render_kernel(function_name, kernel, bufs, uops, prefix)

 _nms = "xyzwabcdefghijkl"
@@ -371,7 +362,7 @@ class CUDARenderer(CStyleLanguage):

    dt_map_in = { dtypes.float: "tf32", dtypes.half: "f16", dtypes.bfloat16: "bf16" }
    dt_map_out = { dtypes.float: "f32", dtypes.half: "f16" }
-    for name, (N, M, K), dtype_in, dtype_out, _, _, upcast_axes, _ in wmma_args(uops):
+    for name, (N, M, K), dtype_in, dtype_out, _, _, upcast_axes, _ in dedup([uop.arg for uop in uops if uop.op is Ops.WMMA]):
      upcast_sizes = [prod(size for _, size in upcast) for upcast in upcast_axes]
      wmma_dtypes = [self.render_dtype(dtype.vec(size)) for dtype, size in zip([dtype_in, dtype_in, dtype_out], upcast_sizes)]
      n_operands = [size*dtype.itemsize//4 for dtype, size in zip([dtype_in, dtype_in, dtype_out], upcast_sizes)] # 4 => CUDA reg size in bytes
@@ -466,15 +457,15 @@ class AMDRenderer(CStyleLanguage):
    if any(dt.scalar() == dtypes.bfloat16 for dt in used_dtypes): prefix.append("typedef unsigned short hip_bfloat16;")
    prefix += [self.render_vector_prefix(dt) for dt in used_dtypes if dt.count > 1]

-    for name, _, dtype_in, dtype_out, _, _, _, _ in wmma_args(uops): # TODO: handle TCs f32_bf16 and bf16_bf16 w/ wrapper
+    for arg in dedup([uop.arg for uop in uops if uop.op is Ops.WMMA]): # 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'}")
+        prefix.append(f"#define __{arg[0]} __builtin_amdgcn_mfma_f32_16x16x16{'f16' if arg[2] == 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")
-      elif dtype_out == dtypes.float:
-        prefix.append(f"#define __{name} __builtin_amdgcn_wmma_f32_16x16x16_{'f16' if dtype_in == dtypes.half else 'bf16'}_w32")
-      else: prefix.append(f"static inline __attribute__((device)) half8 __{name}"+"""(half16 a, half16 b, half8 c) {
+        prefix.append(f"#define __{arg[0]} __builtin_amdgcn_wmma_{type_map[arg[3]]}_16x16x16_{type_map[arg[2]]}_w32_gfx12")
+      elif arg[3] == dtypes.float:
+        prefix.append(f"#define __{arg[0]} __builtin_amdgcn_wmma_f32_16x16x16_{'f16' if arg[2] == dtypes.half else 'bf16'}_w32")
+      else: prefix.append(f"static inline __attribute__((device)) half8 __{arg[0]}"+"""(half16 a, half16 b, half8 c) {
  half16 c_frag = {}; half8 d; for (int n = 0; n < 8; n++) { c_frag[n*2] = c[n]; }
  c_frag = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(a, b, c_frag, false);
  for (int n = 0; n < 8; n++) { d[n] = c_frag[n*2]; } return d;\n}""")
@@ -48,9 +48,8 @@ def render_wmma_amx(ctx, wmma: UOp) -> str:
 def render_wmma_amd(ctx, wmma: UOp, arch: str) -> str:
  dt_map = {dtypes.half: "f16", dtypes.float: "f32", dtypes.bfloat16: "bf16", dtypes.ushort: "bf16"}
  # https://github.com/llvm/llvm-project/blob/main/clang/test/CodeGenOpenCL/builtins-amdgcn-mfma.cl
-  if arch.split(":")[0] in {"gfx942", "gfx950"}:
-    return f"  {ctx[wmma]} = call {ldt(wmma.dtype)} @llvm.amdgcn.mfma.{dt_map[wmma.src[-1].dtype.scalar()]}" + \
-           f".16x16x16{dt_map[wmma.src[0].dtype.scalar()]}(" + ", ".join([f"{ldt(w.dtype)} {ctx[w]}" for w in wmma.src]) + ", i32 0, i32 0, i32 0)"
+  if arch.split(":")[0] == "gfx942": return f"  {ctx[wmma]} = call {ldt(wmma.dtype)} @llvm.amdgcn.mfma.{dt_map[wmma.src[-1].dtype.scalar()]}" + \
+    f".16x16x16{dt_map[wmma.src[0].dtype.scalar()]}(" + ", ".join([f"{ldt(w.dtype)} {ctx[w]}" for w in wmma.src]) + ", i32 0, i32 0, i32 0)"
  # https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AMDGPU/GlobalISel/llvm.amdgcn.wmma_32.ll
  # example: %wmma0 = call <8 x float> @llvm.amdgcn.wmma.f32.16x16x16.f16(<16 x half> %v99,<16 x half> %v100,<8 x float> %v101)
  return f"  {ctx[wmma]} = call {ldt(wmma.dtype)} @llvm.amdgcn.wmma.{dt_map[wmma.src[-1].dtype.scalar()]}.16x16x16." + \
@@ -161,7 +160,6 @@ class LLVMRenderer(Renderer):

    name = "test"
    for u in uops:
-      if u.op is Ops.NOOP: continue
      if u.op is Ops.SINK:
        if u.arg is not None: name = u.arg.function_name
        continue
@@ -172,7 +170,13 @@ class LLVMRenderer(Renderer):
        r[u] = f"%{'local' if u.op is Ops.DEFINE_LOCAL else 'reg'}_{str(u.arg).replace('(', '').replace(')', '').replace(',', '_').replace(' ', '')}"
        assert isinstance(u.dtype, PtrDType)
        if self.device == "LLVM" or u.op is Ops.DEFINE_REG:
-          kernel.append(f"  {r[u]} = alloca [{u.dtype.size} x {ldt(u.dtype.base)}]")
+          # put alloca in the beginning of the function always
+          kernel = [f"  {r[u]} = alloca [{u.dtype.size} x {ldt(u.dtype.base)}]"] + kernel
+          if u.op is Ops.DEFINE_REG:
+            # store the const here. TODO: this should be INDEX and STORE and shouldn't be handcoded here
+            for i in range(u.dtype.size):
+              kernel.append(f"  {r[u]}_idx_{i} = getelementptr inbounds {ldt(u.dtype.base)}, {ldt(u.dtype)} {r[u]}, i32 {i}")
+              kernel.append(f"  store {ldt(u.src[0].dtype)} {r[u.src[0]]}, {ldt(u.dtype)} {r[u]}_idx_{i}")
        else:
          local_args.append(f"@{r[u][1:]} = internal unnamed_addr addrspace(3) global [{u.dtype.size} x {ldt(u.dtype)}] undef, align 16")
          kernel.append(f"  {r[u]} = addrspacecast [{u.dtype.size} x {ldt(u.dtype)}] addrspace(3)* @{r[u][1:]} to [{u.dtype.size} x {ldt(u.dtype)}]*")
@@ -189,6 +193,9 @@ class LLVMRenderer(Renderer):
        if (l:=self.string_rewrite.rewrite(u, ctx=r)) is None:
          raise RuntimeError(f"failed to render {u.op} with {u.dtype} srcs {[x.dtype for x in u.src]}")
        kernel.append(cast(str, l))
+
+        # stores pass the first arg through
+        if u.op is Ops.STORE: r[u] = r[u.src[0]]
    return tuple(local_args), self._render_fn(name, args, kernel, prefix)

 barrier = 'fence syncscope("workgroup") release\ntail call void @llvm.amdgcn.s.barrier()\nfence syncscope("workgroup") acquire\n'
@@ -54,7 +54,7 @@ ptx_matcher = PatternMatcher([
  # move mask from INDEX to the load/store to enable pointer arithmetic
  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"), UPat.var("gate"))), UPat.var("alt"))),
   lambda buf,idx,gate,alt: UOp(Ops.LOAD, alt.dtype, (buf.index(idx), alt, gate))),
-  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"), UPat())), UPat.var("val"), UPat.var("gate")), allow_any_len=True),
+  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"), UPat())), UPat.var("val"), UPat.var("gate"))),
   lambda buf,idx,val,gate: UOp.store(buf.index(idx), val, gate)),
  # ptx shr and shl instructions require y to be uint
  (UPat.var("x") << UPat.var("y"), lambda x,y: UOp(Ops.SHL, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
@@ -110,7 +110,10 @@ string_rewrite = PatternMatcher([
  (UPat(Ops.LOAD, name="x", src=(UPat.var('loc'),), allow_any_len=True),
   lambda ctx, x, loc: f"ld.{mem_type(x)}.v{x.dtype.count}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];" \
     if x.dtype.count > 1 else f"ld.{mem_type(x)}.{ctx.mem_types[x.dtype]} {ctx.r[x]}, [{ctx.r[loc]}+0];"),
-  (UPat(Ops.DEFINE_REG, src=()), lambda ctx: []),
+  (UPat(Ops.DEFINE_REG, name="x", src=(UPat.cvar("pred", dtype=dtypes.bool),), allow_any_len=True), lambda ctx, x, pred: [
+    f"setp.ne.s16 {ctx.r[pred]}, {render_val(pred.arg, pred.dtype)}, 0;", f"mov.pred {ctx.r[x]}, {ctx.r[pred]};"]),
+  (UPat(Ops.DEFINE_REG, name="x", src=(UPat.cvar("pred"),), allow_any_len=True),
+   lambda ctx, x, pred: f"mov.b{ctx.types[x.dtype.base][1:]} {ctx.r[x]}, {render_val(pred.arg, x.dtype.base)};"),
  (UPat(Ops.RANGE, name="x"), lambda ctx, x: [f"mov.u32 {ctx.r[x]}, 0;", "LOOP_" + f"{ctx.r[x][1:]}:"]),
  (UPat(Ops.ENDRANGE, name="x", src=(UPat.var("src0"),)), lambda ctx, x, src0: [
    ctx.code_for_op[Ops.ADD](ctx.r[src0], ctx.r[src0], "1", dtypes.int, ctx.types[dtypes.int]),
@@ -173,7 +176,6 @@ class PTXRenderer(Renderer):

    name = "test"
    for u in uops:
-      if u.op is Ops.NOOP: continue
      if u.op is Ops.SINK:
        if u.arg is not None: name = u.arg.function_name
        continue
@@ -186,18 +188,11 @@ class PTXRenderer(Renderer):
      if u.op in {Ops.CAST, Ops.BITCAST} and (u.src[0].dtype == u.dtype or isinstance(u.src[0].dtype, PtrDType)):
        r[u] = r[u.src[0]]
        continue
-      if u.op is Ops.DEFINE_REG:
-        r[u] = [ssa("reg", u, self.types[u.dtype.base.scalar()]) for _ in range(cast(PtrDType, u.dtype).size)]
-        continue
      if u.op in {Ops.INDEX, Ops.LOAD, Ops.STORE} and isinstance(u.src[0].dtype, PtrDType) and u.src[0].dtype.addrspace == AddrSpace.REG:
-        if u.op is Ops.INDEX:
-          assert u.src[1].op == Ops.CONST, f"index on REG in ptx only supported on CONST, not {u.src[1].op}"
-          r[u] = r[u.src[0]][u.src[1].arg]
-        else:
-          r[u] = r[u.src[0]]
-          if u.op is Ops.STORE:
-            typ = "pred" if u.src[1].dtype == dtypes.bool else ("b"+self.types[u.src[1].dtype][1:])
-            kernel.append(f"mov.{typ} {self.r[u.src[0]]}, {self.r[u.src[1]]};")
+        r[u] = r[u.src[0]]
+        if u.op is Ops.STORE:
+          typ = "pred" if u.src[1].dtype == dtypes.bool else ("b"+self.types[u.src[1].dtype][1:])
+          kernel.append(f"mov.{typ} {self.r[u.src[0]]}, {self.r[u.src[1]]};")
        continue
      if u.op is Ops.SPECIAL: r[u] = "%" + u.arg[0]
      elif u.op is Ops.DEFINE_VAR: bufs.append((u.arg[0], u.dtype))
@@ -207,12 +202,12 @@ class PTXRenderer(Renderer):
      elif u.op is Ops.DEFINE_GLOBAL: bufs.append((f"data{u.arg}", u.dtype))
      elif u.op is Ops.WMMA:
        # registers for packing/unpacking input and acc
-        self.wmma_r = [[ssa("wmma_in", dtype="b32") for _ in range(0, len(r[u.src[0]]), 4 // u.src[0].dtype.scalar().itemsize)],
-                       [ssa("wmma_in", dtype="b32") for _ in range(0, len(r[u.src[1]]), 4 // u.src[0].dtype.scalar().itemsize)],
-                       [ssa("wmma_acc", dtype="b32") for _ in range(0, len(r[u.src[2]]), 4 // u.dtype.scalar().itemsize)]]
+        self.wmma_r = [[ssa("wmma_in", dtype="b32") for _ in range(0, len(r[u.src[0]]), 4 // u.arg[2].itemsize)],
+                       [ssa("wmma_in", dtype="b32") for _ in range(0, len(r[u.src[1]]), 4 // u.arg[2].itemsize)],
+                       [ssa("wmma_acc", dtype="b32") for _ in range(0, len(r[u.src[2]]), 4 // u.arg[3].itemsize)]]
        r[u] = [ssa("wmma", dtype=self.types[u.dtype.scalar()]) for _ in range(u.dtype.count)]
      prefix, dtype = {Ops.CAST: ("cast", None), Ops.BITCAST: ("cast", None), Ops.ENDRANGE: ("pred", "pred"), Ops.RANGE: ("ridx", None),
-        Ops.DEFINE_VAR: ("dat", None), Ops.CONST: ("const", None), Ops.DEFINE_LOCAL:("local",self.types[dtypes.ulong]),
+        Ops.DEFINE_REG: ("acc", None), Ops.DEFINE_VAR: ("dat", None), Ops.CONST: ("const", None), Ops.DEFINE_LOCAL:("local",self.types[dtypes.ulong]),
        Ops.DEFINE_GLOBAL: ("dat", self.types[dtypes.ulong]), **{op: ("alu", None) for op in GroupOp.ALU}}.get(u.op, (None, None))
      if prefix: r[u] = ssa(prefix, u, dtype)

@@ -40,6 +40,11 @@ wgsl_matcher = PatternMatcher([
  (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
  ]) + extra_pm

+def webgpu_define_reg(ctx, x):
+  ret = [f"var {ctx[x]}: array<{ctx.buf_map(x.dtype)},{x.dtype.size//(4//x.dtype.itemsize) if is_packed(x.dtype) else x.dtype.size}>;"]
+  for i in range(x.dtype.size): ret.append(f"{ctx[x]}[{i}] = {ctx[x.src[0]]};")
+  return ' '.join(ret)
+
 class WGSLRenderer(CStyleLanguage):
  device = "WEBGPU"
  global_max = (65535, 65535, 65535)
@@ -59,8 +64,7 @@ class WGSLRenderer(CStyleLanguage):
     lambda x: f"bitcast<u32>({x.arg})" if x.arg < 0 else f"{x.arg&0xFFFFFFFF}u"),
    (UPat(Ops.DEFINE_LOCAL, name="x"), lambda ctx,x:
      f"var<workgroup> {ctx[x]}: array<{ctx.buf_map(x.dtype.base)},{x.dtype.size//(4//x.dtype.itemsize) if is_packed(x.dtype) else x.dtype.size}>;"),
-    (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x:
-     f"var {ctx[x]}: array<{ctx.buf_map(x.dtype)},{x.dtype.size//(4//x.dtype.itemsize) if is_packed(x.dtype) else x.dtype.size}>;"),
+    (UPat(Ops.DEFINE_REG, name="x"), webgpu_define_reg),
    (UPat(Ops.BITCAST, dtype=dtypes.half, name="x", src=(UPat(dtype=(dtypes.short, dtypes.ushort, dtypes.uint32),),)),
     lambda ctx,x: f"bitcast<vec2<f16>>({ctx[x.src[0]]})[0]"),
    (UPat(Ops.BITCAST, dtype=(dtypes.char, dtypes.uchar), name="x"), lambda ctx,x: f"bitcast<{ctx.type_map[x.dtype]}>({ctx[x.src[0]]}&0xFF)"),
@@ -1,11 +1,11 @@
 import collections, time
 from typing import Any, cast
-from tinygrad.helpers import round_up, PROFILE, merge_dicts, getenv
+from tinygrad.helpers import round_up, PROFILE, merge_dicts
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQSignal, HCQBuffer, HWQueue, HCQArgsState, BumpAllocator
 from tinygrad.device import Buffer, BufferSpec, Compiled, Device, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Variable
-from tinygrad.engine.realize import ExecItem, BufferXfer, CompiledRunner, BufferCopy
+from tinygrad.engine.realize import ExecItem, BufferXfer, CompiledRunner
 from tinygrad.engine.jit import MultiGraphRunner

 class HCQGraph(MultiGraphRunner):
@@ -13,9 +13,6 @@ class HCQGraph(MultiGraphRunner):
    super().__init__(jit_cache, input_rawbuffers, var_vals)
    self.devices = list(set(cast(HCQCompiled, d) for ji in jit_cache for d in [Device[cast(Buffer, x).device] for x in ji.bufs]))

-    # CPU Device is always last
-    self.devices = sorted(self.devices, key=lambda x: 1 if x._is_cpu() else 0)
-
    # Replace input buffers with variables.
    self.hcq_bufs = [[cast(Buffer, x)._buf for x in ji.bufs] for ji in jit_cache]
    self.input_replace_to_var: dict[tuple[int, int], Variable] = {}
@@ -51,8 +48,7 @@ class HCQGraph(MultiGraphRunner):
    self.comp_queues: dict[HCQCompiled, HWQueue] = {dev: dev.hw_compute_queue_t() for dev in self.devices}
    self.copy_queues: dict[HCQCompiled, HWQueue] = {} # lazy allocation

-    self.signals: dict[Any, HCQSignal] = {**{dev: dev.new_signal(value=0) for dev in self.devices if dev.device != "CPU"},
-      **{"KICK": self.devices[0].new_signal(value=0)}, **{dev: self.devices[0].new_signal(value=0) for dev in self.devices if dev.device == "CPU"}}
+    self.signals: dict[Any, HCQSignal] = {**{dev: dev.new_signal(value=0) for dev in self.devices}, **{"KICK": self.devices[0].new_signal(value=0)}}
    self.kickoff_value: int = 0
    self.kickoff_var = UOp.variable("kickoff_var", 0, 0xffffffff, dtype=dtypes.uint32)

@@ -68,15 +64,10 @@ class HCQGraph(MultiGraphRunner):

    for dev, queue in self.comp_queues.items(): dev_access[queue].add(dev)

-    self.input_replace_map: dict[HCQCompiled, set[int]] = collections.defaultdict(set)
    self.fixedvars: dict[HCQCompiled, dict[Variable, int]] = {}

    for j,ji in enumerate(jit_cache):
-      if is_exec_prg:=isinstance(ji.prg, CompiledRunner): enqueue_dev: HCQCompiled = ji.prg.dev
-      else:
-        # For copy ops prioritize enqeueuing on the dest device, so reverse the buffers.
-        for b in cast(list[Buffer], ji.bufs[::-1]):
-          if (enqueue_dev:=cast(HCQCompiled, Device[b.device])).hw_copy_queue_t is not None: break
+      enqueue_dev: HCQCompiled = ji.prg.dev if (is_exec_prg:=isinstance(ji.prg, CompiledRunner)) else Device[ji.bufs[1].device] #type:ignore

      # set any fixedvars on the device
      self.fixedvars[enqueue_dev] = merge_dicts([self.fixedvars.get(enqueue_dev, {}), ji.fixedvars])
@@ -157,11 +148,10 @@ class HCQGraph(MultiGraphRunner):
      # Encode main commands based on ji type.
      if isinstance(ji.prg, CompiledRunner):
        enqueue_queue.exec(ji.prg._prg, self.ji_args[j], tuple(ji.prg.p.global_size or (1,1,1)), tuple(ji.prg.p.local_size or (1,1,1)))
-      elif isinstance(ji.prg, (BufferXfer, BufferCopy)):
+      elif isinstance(ji.prg, BufferXfer):
        dest, src = [cast(Buffer, x) for x in ji.bufs[0:2]]
-        for bufid, src in enumerate(cast(list[Buffer], ji.bufs)):
-          if (inprep_idx:=self.input_replace.get((j, bufid))) is not None: self.input_replace_map[enqueue_dev].add(inprep_idx)
-          else: cast(HCQAllocator, enqueue_dev.allocator).map(self.hcq_bufs[j][bufid])
+        cast(HCQAllocator, Device[src.device].allocator).map(dest._buf)
+
        enqueue_queue.copy(self.hcq_bufs[j][0].va_addr, self.hcq_bufs[j][1].va_addr, dest.nbytes)
        self.copy_to_devs[cast(HCQCompiled, Device[dest.device])].add(cast(HCQCompiled, Device[src.device]))

@@ -187,9 +177,6 @@ class HCQGraph(MultiGraphRunner):
    for sig in self.queue_signals_to_reset: sig.value = 0
    self.signals['KICK'].value = self.kickoff_value

-    for dev in self.devices:
-      for idx_to_map in self.input_replace_map[dev]: cast(HCQAllocator, dev.allocator).map(input_rawbuffers[idx_to_map]._buf)
-
    if PROFILE and self.kickoff_value > 1: self.collect_timestamps()

    hcq_var_vals = {self.kickoff_var: self.kickoff_value, **var_vals,
@@ -223,9 +210,3 @@ class HCQGraph(MultiGraphRunner):
    if PROFILE and self.kickoff_value >= 1: self.collect_timestamps()

    for fdev, buf in self.kernargs_bufs.items(): fdev.allocator._free(buf, BufferSpec(cpu_access=True))
-
-  @staticmethod
-  def supports_exec_item(dev, ei:ExecItem) -> bool:
-    # MOCKGPU is not supported, since it can't execute commands in parallel
-    copy = (isinstance(ei.prg, BufferCopy) and cast(HCQCompiled, dev).hw_copy_queue_t is not None) and not getenv("MOCKGPU")
-    return all(issubclass(type(Device[b.device]), HCQCompiled) for b in ei.bufs if b) and (isinstance(ei.prg, (CompiledRunner, BufferXfer)) or copy)
@@ -7,7 +7,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, H
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.uop.ops import sint
 from tinygrad.device import Compiled, DMAFdRef, BufferSpec
-from tinygrad.helpers import getenv, to_mv, round_up, data64_le, all_same, flatten, DEBUG, AMD_LLVM, PROFILE, ProfileEvent, suppress_finalizing
+from tinygrad.helpers import getenv, to_mv, round_up, data64_le, all_same, flatten, DEBUG, AMD_LLVM, PROFILE, ProfileEvent
 from tinygrad.renderer.cstyle import AMDRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt
@@ -473,12 +473,13 @@ class AMDAllocator(HCQAllocator['AMDDevice']):
  def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
    return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access)

-  @suppress_finalizing
  def _free(self, opaque, options:BufferSpec):
-    self.dev.synchronize()
-    self.dev.iface.free(opaque)
+    try:
+      self.dev.synchronize()
+      self.dev.iface.free(opaque)
+    except AttributeError: pass

-  def _map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+  def _map(self, buf:HCQBuffer): self.dev.iface.map(buf._base if buf._base is not None else buf)

@dataclass(frozen=True)
 class ProfileSQTTEvent(ProfileEvent): device:str; se:int; blob:bytes; itrace:bool # noqa: E702
@@ -562,7 +563,7 @@ class KFDIface:
    self.mem_fault_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_MEMORY)
    self.hw_fault_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_HW_EXCEPTION)

-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, cpu_addr=None) -> HCQBuffer:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False) -> HCQBuffer:
    flags = kfd.KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_EXECUTABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE

    if uncached: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED | kfd.KFD_IOC_ALLOC_MEM_FLAGS_GTT
@@ -571,7 +572,7 @@ class KFDIface:
    if cpu_access or host: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_PUBLIC

    if flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR:
-      buf = addr = cpu_addr or FileIOInterface.anon_mmap(0, size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | mmap.MAP_ANONYMOUS, 0)
+      buf = addr = FileIOInterface.anon_mmap(0, size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | mmap.MAP_ANONYMOUS, 0)
    else: buf, addr = 0, FileIOInterface.anon_mmap(0, size, 0, mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS | MAP_NORESERVE, 0)

    try: mem = kfd.AMDKFD_IOC_ALLOC_MEMORY_OF_GPU(self.kfd, va_addr=addr, size=size, base=addr, length=size, gpu_id=self.gpu_id,
@@ -592,7 +593,7 @@ class KFDIface:

  def free(self, mem):
    if len(mem.mapped_devs) > 0:
-      gpus = (ctypes.c_int32 * len(mem.mapped_devs))(*[x.iface.gpu_id for x in mem.mapped_devs])
+      gpus = (ctypes.c_int32 * len(mem.mapped_devs))(*[x.gpu_id for x in mem.mapped_devs])
      stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=len(gpus))
      assert stm.n_success == len(gpus)
    if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
@@ -605,8 +606,6 @@ class KFDIface:
    return dmaref

  def map(self, mem):
-    if mem.owner is not None and mem.owner._is_cpu(): return self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
-
    c_gpus = (ctypes.c_int32 * 1)(self.gpu_id)
    stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(c_gpus), n_devices=1)
    assert stm.n_success == 1
@@ -747,7 +746,7 @@ class AMDDevice(HCQCompiled):
                       (min((self.max_cu_id+1)*40, self.iface.props['array_count'] // self.iface.props['simd_arrays_per_engine'] * 512) - 1)
    self.xccs = self.iface.props.get('num_xcc', 1) if getenv("XCCS", 1) else 1
    # this is what llvm refers to as "architected flat scratch"
-    self.has_scratch_base_registers = self.target >= (11,0,0) or self.target in {(9,4,2), (9,5,0)}
+    self.has_scratch_base_registers = self.target >= (11,0,0) or self.target in {(9,4,2),(9,5)}

    # https://gitlab.freedesktop.org/agd5f/linux/-/blob/a1fc9f584c4aaf8bc1ebfa459fc57a3f26a290d8/drivers/gpu/drm/amd/amdkfd/kfd_queue.c#L391
    sgrp_size_per_cu, lds_size_per_cu, hwreg_size_per_cu = 0x4000, 0x10000, 0x1000
@@ -1,6 +1,6 @@
 from __future__ import annotations
-import platform, subprocess, sys, ctypes, functools, time, mmap
-from tinygrad.helpers import capstone_flatdump, getenv, from_mv, to_mv, OSX, mv_address, wait_cond, cpu_profile
+import platform, subprocess, sys, ctypes, functools, time
+from tinygrad.helpers import capstone_flatdump, getenv, from_mv, to_mv, OSX, mv_address, round_up, wait_cond
 from tinygrad.device import Compiler, BufferSpec, DMACPURef
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocatorBase, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface
 from tinygrad.runtime.support.elf import jit_loader
@@ -62,9 +62,10 @@ class CPUProgram(HCQProgram):
      ctypes.windll.kernel32.FlushInstructionCache(ctypes.c_void_p(proc), ctypes.c_void_p(self.mem), ctypes.c_size_t(len(lib)))
      self.fxn = ctypes.CFUNCTYPE(None)(self.mem)
    else:
+      from mmap import mmap, PROT_READ, PROT_WRITE, PROT_EXEC, MAP_ANON, MAP_PRIVATE
      # On apple silicon with SPRR enabled (it always is in macos) RWX pages are unrepresentable: https://blog.svenpeter.dev/posts/m1_sprr_gxf/
      # MAP_JIT allows us to easily flip pages from RW- to R-X and vice versa. It is a noop on intel cpus. (man pthread_jit_write_protect_np)
-      self.mem = mmap.mmap(-1, len(lib), mmap.MAP_ANON|mmap.MAP_PRIVATE|(MAP_JIT if OSX else 0), mmap.PROT_READ|mmap.PROT_WRITE|mmap.PROT_EXEC)
+      self.mem = mmap(-1, len(lib), MAP_ANON | MAP_PRIVATE | (MAP_JIT if OSX else 0), PROT_READ | PROT_WRITE | PROT_EXEC)

      if OSX: CPUProgram.rt_lib.pthread_jit_write_protect_np(False)
      self.mem.write(lib)
@@ -81,21 +82,19 @@ class CPUProgram(HCQProgram):
    super().__init__(HCQArgsState, dev, name, kernargs_alloc_size=0)

  def __del__(self):
-    if getattr(sys, 'is_finalizing', lambda: True)(): return
    if sys.platform == 'win32': ctypes.windll.kernel32.VirtualFree(ctypes.c_void_p(self.mem), ctypes.c_size_t(0), 0x8000) #0x8000 - MEM_RELEASE

 class CPUAllocator(HCQAllocatorBase):
  def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
-    if options.external_ptr: addr, buf = options.external_ptr, None
-    elif sys.platform == "win32": addr = mv_address(buf:=mmap.mmap(-1, size, access=mmap.ACCESS_WRITE))
-    else: addr = mv_address(buf:=mmap.mmap(-1, size, mmap.MAP_ANON | mmap.MAP_PRIVATE, mmap.PROT_READ | mmap.PROT_WRITE))
-    return HCQBuffer(va:=addr, sz:=size, meta=buf, view=MMIOInterface(va, sz, fmt='B'), owner=self.dev)
+    if options.external_ptr: buf = (ctypes.c_uint8 * size).from_address(options.external_ptr)
+    else:
+      offset = round_up(ctypes.addressof(tmpbuf:=(ctypes.c_uint8 * (size + 0x1000))()), 0x1000) - ctypes.addressof(tmpbuf)
+      buf = (ctypes.c_uint8 * size).from_buffer(tmpbuf, offset)
+    return HCQBuffer(va:=ctypes.addressof(buf), sz:=ctypes.sizeof(buf), meta=buf, view=MMIOInterface(va, sz, fmt='B'), owner=self.dev)
  def _as_buffer(self, src) -> memoryview: return to_mv(src.va_addr, src.size)
  def _as_dmaref(self, buf): return DMACPURef(buf.va_addr, buf.size)
-  def _copyin(self, dest, src:memoryview):
-    with cpu_profile('TINY -> CPU', self.dev.device, is_copy=True): ctypes.memmove(dest.va_addr, from_mv(src), len(src))
-  def _copyout(self, dest:memoryview, src):
-    with cpu_profile('CPU -> TINY', self.dev.device, is_copy=True): ctypes.memmove(from_mv(dest), src.va_addr, len(dest))
+  def _copyin(self, dest, src:memoryview): ctypes.memmove(dest.va_addr, from_mv(src), len(src))
+  def _copyout(self, dest:memoryview, src): ctypes.memmove(from_mv(dest), src.va_addr, len(dest))
  def _map(self, buf:HCQBuffer):
    if buf.view is None or not isinstance(buf.view, MMIOInterface): raise RuntimeError("Cannot map buffer without view to cpu")

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import ctypes, ctypes.util, functools
-from tinygrad.helpers import DEBUG, getenv, mv_address, init_c_var, init_c_struct_t, suppress_finalizing
+from tinygrad.helpers import DEBUG, getenv, mv_address, init_c_var, init_c_struct_t
 from tinygrad.device import Compiled, BufferSpec, LRUAllocator
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.renderer.ptx import PTXRenderer
@@ -45,8 +45,9 @@ class CUDAProgram:
    self.prg = prg
    if self.smem > 0: check(cuda.cuFuncSetAttribute(self.prg, cuda.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, self.smem))

-  @suppress_finalizing
-  def __del__(self): check(cuda.cuModuleUnload(self.module))
+  def __del__(self):
+    try: check(cuda.cuModuleUnload(self.module))
+    except AttributeError: pass

  def __call__(self, *args, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
    check(cuda.cuCtxSetCurrent(self.dev.context))
@@ -2,7 +2,7 @@ from __future__ import annotations
 from typing import cast
 import ctypes, functools, hashlib
 from tinygrad.runtime.autogen import opencl as cl
-from tinygrad.helpers import init_c_var, to_char_p_p, from_mv, OSX, DEBUG, getenv, mv_address, suppress_finalizing
+from tinygrad.helpers import init_c_var, to_char_p_p, from_mv, OSX, DEBUG, getenv, mv_address
 from tinygrad.renderer.cstyle import OpenCLRenderer, IntelRenderer
 from tinygrad.device import BufferSpec, LRUAllocator, Compiled, Compiler, CompileError

@@ -69,8 +69,9 @@ class CLAllocator(LRUAllocator['CLDevice']):
                                        cl.cl_image_format(cl.CL_RGBA, {2: cl.CL_HALF_FLOAT, 4: cl.CL_FLOAT}[options.image.itemsize]),
                                        options.image.shape[1], options.image.shape[0], 0, None, status := ctypes.c_int32()), status), options)
    return (checked(cl.clCreateBuffer(self.dev.context, cl.CL_MEM_READ_WRITE, size, None, status := ctypes.c_int32()), status), options)
-  @suppress_finalizing
-  def _free(self, opaque:tuple[ctypes._CData, BufferSpec], options:BufferSpec): check(cl.clReleaseMemObject(opaque[0]))
+  def _free(self, opaque:tuple[ctypes._CData, BufferSpec], options:BufferSpec):
+    try: check(cl.clReleaseMemObject(opaque[0]))
+    except AttributeError: pass
  def _copyin(self, dest:tuple[ctypes._CData, BufferSpec], src:memoryview):
    if dest[1].image is not None:
      check(cl.clEnqueueWriteImage(self.dev.queue, dest[0], False, (ctypes.c_size_t * 3)(0,0,0),
@@ -1,6 +1,6 @@
-import subprocess, pathlib, struct, ctypes, tempfile, functools, contextlib, decimal, platform
+import os, pathlib, struct, ctypes, tempfile, functools, contextlib, decimal, platform
 from typing import Any, cast
-from tinygrad.helpers import prod, to_mv, getenv, round_up, cache_dir, T, init_c_struct_t, PROFILE, ProfileRangeEvent, cpu_profile, unwrap
+from tinygrad.helpers import prod, to_mv, getenv, round_up, cache_dir, T, init_c_struct_t, PROFILE, ProfileRangeEvent, cpu_profile
 from tinygrad.device import Compiled, Compiler, CompileError, LRUAllocator, ProfileDeviceEvent
 from tinygrad.renderer.cstyle import MetalRenderer

@@ -144,10 +144,7 @@ class MetalCompiler(Compiler):
    with tempfile.NamedTemporaryFile(delete=True) as shader:
      shader.write(lib)
      shader.flush()
-      proc = subprocess.Popen(f"cd {pathlib.Path(__file__).parents[2]}/extra/disassemblers/applegpu && python3 compiler_explorer.py {shader.name}",
-                              stdout=subprocess.PIPE, shell=True, text=True, bufsize=1)
-      for line in unwrap(proc.stdout): print(line, end="")
-      ret = proc.wait()
+      ret = os.system(f"cd {pathlib.Path(__file__).parents[2]}/extra/disassemblers/applegpu && python3 compiler_explorer.py {shader.name}")
      if ret: print("Disassembler Error: Make sure you have https://github.com/dougallj/applegpu cloned to tinygrad/extra/disassemblers/applegpu")

 class MetalProgram:
@@ -226,6 +223,6 @@ class MetalAllocator(LRUAllocator[MetalDevice]):
  def _as_buffer(self, src:MetalBuffer) -> memoryview:
    self.dev.synchronize()
    return to_mv(cast(int, msg("contents", objc_id)(src.buf).value), src.size + src.offset)[src.offset:]
-  def _copyin(self, dest:MetalBuffer, src:memoryview): self._cp_mv(self._as_buffer(dest), src, "TINY -> METAL")
-  def _copyout(self, dest:memoryview, src:MetalBuffer): self._cp_mv(dest, self._as_buffer(src), "METAL -> TINY")
+  def _copyin(self, dest:MetalBuffer, src:memoryview): self._cp_mv(self._as_buffer(dest), src, "CPU -> METAL")
+  def _copyout(self, dest:memoryview, src:MetalBuffer): self._cp_mv(dest, self._as_buffer(src), "METAL -> CPU")
  def _offset(self, buf:MetalBuffer, size:int, offset:int): return MetalBuffer(buf.buf, size, offset)
@@ -7,7 +7,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, H
 from tinygrad.runtime.support.hcq import MMIOInterface, FileIOInterface, MOCKGPU
 from tinygrad.uop.ops import sint
 from tinygrad.device import BufferSpec
-from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, to_mv, hi32, lo32, suppress_finalizing
+from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, to_mv, hi32, lo32
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import NVRenderer
 from tinygrad.runtime.support.compiler_cuda import CUDACompiler, PTXCompiler, PTX, NVPTXCompiler, NVCompiler
@@ -276,12 +276,13 @@ class NVAllocator(HCQAllocator['NVDevice']):
  def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
    return self.dev.iface.alloc(size, cpu_access=options.cpu_access, host=options.host)

-  @suppress_finalizing
  def _free(self, opaque:HCQBuffer, options:BufferSpec):
-    self.dev.synchronize()
-    self.dev.iface.free(opaque)
+    try:
+      self.dev.synchronize()
+      self.dev.iface.free(opaque)
+    except AttributeError: pass

-  def _map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+  def _map(self, buf:HCQBuffer): self.dev.iface.map(buf._base if buf._base is not None else buf)

@dataclass
 class GPFifo:
@@ -381,14 +382,14 @@ class NVKIface:
    if made.params.status != 0: raise RuntimeError(f"_gpu_map_to_cpu returned {get_error_str(made.params.status)}")
    return fd_dev.mmap(target, size, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED | (MAP_FIXED if target is not None else 0), 0)

-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, map_flags=0, cpu_addr=None) -> HCQBuffer:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, map_flags=0) -> HCQBuffer:
    # Uncached memory is "system". Use huge pages only for gpu memory.
    page_size = (4 << (12 if OSX else 10)) if uncached or host else ((2 << 20) if size >= (8 << 20) else (4 << (12 if OSX else 10)))
    size = round_up(size, page_size)
    va_addr = self._alloc_gpu_vaddr(size, alignment=page_size, force_low=cpu_access)

    if host:
-      va_addr = cpu_addr or FileIOInterface.anon_mmap(va_addr, size, mmap.PROT_READ|mmap.PROT_WRITE, MAP_FIXED|mmap.MAP_SHARED|mmap.MAP_ANONYMOUS, 0)
+      va_addr = FileIOInterface.anon_mmap(va_addr, size, mmap.PROT_READ | mmap.PROT_WRITE, MAP_FIXED | mmap.MAP_SHARED | mmap.MAP_ANONYMOUS, 0)

      flags = (nv_gpu.NVOS02_FLAGS_PHYSICALITY_NONCONTIGUOUS << 4) | (nv_gpu.NVOS02_FLAGS_COHERENCY_CACHED << 12) \
            | (nv_gpu.NVOS02_FLAGS_MAPPING_NO_MAP << 30)
@@ -437,11 +438,7 @@ class NVKIface:
      hClient=self.root, hMemory=mem_handle, gpuAttributesCount=1, perGpuAttributes=attrs, mapped_gpu_ids=[self.gpu_uuid],
      has_cpu_mapping=has_cpu_mapping), view=MMIOInterface(va_base, size, fmt='B') if has_cpu_mapping else None, owner=self.dev)

-  def map(self, mem:HCQBuffer):
-    if mem.owner is not None and mem.owner._is_cpu():
-      if not any(x.device.startswith("NV") for x in mem.mapped_devs): return self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
-      mem = mem.mappings[next(x for x in mem.mapped_devs if x.device.startswith("NV"))]
-    self._gpu_uvm_map(mem.va_addr, mem.size, mem.meta.hMemory, create_range=False)
+  def map(self, mem:HCQBuffer): self._gpu_uvm_map(mem.va_addr, mem.size, mem.meta.hMemory, create_range=False)

  def _alloc_gpu_vaddr(self, size, alignment=(4 << 10), force_low=False):
    return NVKIface.low_uvm_vaddr_allocator.alloc(size, alignment) if force_low else NVKIface.uvm_vaddr_allocator.alloc(size, alignment)
@@ -40,7 +40,8 @@ class PythonProgram:
      loop_ends: dict[int, int] = {}
      while i < len(self.uops):
        uop, dtype, idp, arg = self.uops[i]
-        void_ops = {Ops.ENDRANGE, Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK, Ops.NOOP, Ops.STORE}
+        void_ops = {Ops.ENDRANGE, Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK}
+        if uop is Ops.DEFINE_REG: idp = [idp[0]]
        inp = [ul[v] for v in idp if self.uops[v][0] not in void_ops]
        dtp = [dl[v] for v in idp if self.uops[v][0] not in void_ops]
        if getenv("TRACE"): print(i, uop, dtype, arg, inp, dtp)
@@ -48,16 +49,18 @@ class PythonProgram:
          loop_ends[idp[0]] = i
          i = idp[0]
          continue
-        if uop in (Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK, Ops.NOOP):
+        if uop in (Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK):
          # in the python emulator, the warp is always in sync
          i += 1
          continue
        assert dtype is not None, f"{uop} is missing a dtype"
        dl[i] = dtype
        if uop is Ops.STORE:
+          #assert len(inp) == 2, "expected store is ([(memory, offset, gate)], [value])"
          for j,val in enumerate(inp[1] if dtp[1].count > 1 else [inp[1]]):
            for (m,o,g),v in zip(inp[0], val):
              if g: _store(m, o+j, v)
+          ul[i] = inp[0]
          i += 1
          continue
        if uop in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}:
@@ -66,6 +69,8 @@ class PythonProgram:
          if uop is Ops.DEFINE_REG:
            # REGs are per thread
            ul[i] = [memoryview(bytearray(dtype.size*dtype.itemsize)).cast(dtype.fmt) for _ in range(warp_size)]
+            for buf, val in zip(ul[i], inp[0]):
+              for x in range(dtype.size): buf[x] = val
          else:
            buf = memoryview(bytearray(dtype.size*dtype.itemsize)) if uop is not Ops.DEFINE_GLOBAL else pbufs.pop(0)
            ul[i] = [buf.cast(dtype.fmt)] * warp_size
@@ -123,27 +128,24 @@ class PythonProgram:
                  out[elem_idx][goff+lane_id] += sum(a_elem(inp[0], _k, c_j, goff) * b_elem(inp[1], c_i, _k, goff) for _k in range(K))
            return out

-          first_src_dtype = self.uops[idp[0]][1]
-          assert isinstance(first_src_dtype, DType) # mypy
-          dims, dtype_in, device, threads = arg[1], first_src_dtype.scalar(), arg[4], arg[5]
          # TODO: refactor these to a shared TensorCoreLayout in kernel.py
-          if device == "METAL":
+          if arg[4] == "METAL":
            # A (2 elements on 32 threads): row major
            def a_b_elem(x, i, j, goff): return x[(i%2)][goff+(i//2)%2+(j%4)*2+(i//4)*8+(j//4)*16]
            # (i, j), C, D (2 elements on 32 threads): row major same as A/B
            def c_map(lane, elem): return (elem + ((lane%2)*2) + ((lane//8)%2)*4, ((lane//2)%4) + (lane//16)*4)
            ul[i] = wmma_helper(32, 8, 2, 2, 2, a_b_elem, a_b_elem, c_map)
-          elif device == "AMD" and threads == 64:
+          elif arg[4] == "AMD" and arg[5] == 64:
            def a_elem(x, k, row, goff): return x[k%4][goff + (k//4)*16 + row]
            def b_elem(x, col, k, goff): return a_elem(x, k, col, goff) # pylint: disable=arguments-out-of-order
            def c_map(lane, elem): return (lane%16, (lane//16)*4 + elem)
            ul[i] = wmma_helper(64, 16, 4, 4, 4, a_elem, b_elem, c_map)
-          elif device == "AMD" and len(inp[0]) == 8: # RDNA4
+          elif arg[4] == "AMD" and len(inp[0]) == 8: # RDNA4
            def a_elem(x, k, row, goff): return x[k - [0, 4, 4, 8][k//4]][goff + row + [0, 16, 0, 16][k//4]]
            def b_elem(x, col, k, goff): return a_elem(x, k, col, goff)
            def c_map(lane, elem): return (lane%16, (lane//16)*8 + elem)
            ul[i] = wmma_helper(32, 16, 8, 8, 8, a_elem, b_elem, c_map)
-          elif device == "AMD":
+          elif arg[4] == "AMD":
            # A (16 elements on 32 threads): col major, lane 16-32 == lane 0-15
            def a_elem(x, k, row, goff):
              assert x[k][goff+row] == x[k][goff+row+16], "warp elements not duplicated properly across lanes"
@@ -152,27 +154,27 @@ class PythonProgram:
            def b_elem(x, col, k, goff): return a_elem(x, k, col, goff)  # pylint: disable=arguments-out-of-order
            def c_map(lane, elem): return (lane%16, lane//16+elem*2) # (i, j), C, D (8 elements on 32 threads): row major
            ul[i] = wmma_helper(32, 16, 16, 16, 8, a_elem, b_elem, c_map)
-          elif device == "CUDA":
+          elif arg[4] == "CUDA":
            # (col, row) given (lane, elem) for C & D (4 elements on 32 threads); shared by all tc shapes with M=16 N=8
            def c_map(lane, elem): return (elem%2 + (lane%4)*2, lane//4 + (elem//2)*8)

-            if dims == (8,16,16):
+            if arg[1] == (8,16,16):
              def a_elem(x, k, row, goff): return x[k%2 + (row//8)*2 + (k//8)*4][goff + (k//2)%4 + (row%8)*4]
              def b_elem(x, col, k, goff): return x[k%2 + (k//8)*2][goff + (k//2)%4 + col*4]
              ul[i] = wmma_helper(32, 16, 8, 4, 4, a_elem, b_elem, c_map)

-            elif dims == (8,16,8) and dtype_in == dtypes.half:
+            elif arg[1] == (8,16,8) and arg[2] == dtypes.half:
              def a_elem(x, k, row, goff): return x[k%2 + (row//8)*2][goff + k//2 + (row%8)*4]
              def b_elem(x, col, k, goff): return x[k%2][goff + k//2 + col*4]
              ul[i] = wmma_helper(32, 8, 4, 2, 4, a_elem, b_elem, c_map)

-            elif dims == (8,16,8) and dtype_in == dtypes.float:
+            elif arg[1] == (8,16,8) and arg[2] == dtypes.float:
              def a_elem(x, k, row, goff): return x[(k//4)*2 + row//8][goff + k%4 + (row%8)*4]
              def b_elem(x, col, k, goff): return x[k//4][goff + k%4 + col*4]
              ul[i] = wmma_helper(32, 8, 4, 2, 4, a_elem, b_elem, c_map)

            else: raise NotImplementedError(f"unimplemented tensor core {arg}")
-          elif device == "INTEL":
+          elif arg[4] == "INTEL":
            # A (16 elements on 8 threads)
            def a_elem(x, k, row, goff): return x[k%2+row*2][goff+k//2]
            # B (16 elements on 8 threads)
@@ -180,7 +182,7 @@ class PythonProgram:
            # C, D (8 elements on 8 threads)
            def c_map(lane, elem): return (lane, elem)
            ul[i] = wmma_helper(8, 16, 16, 16, 8, a_elem, b_elem, c_map)
-          elif device == "CPU":
+          elif arg[4] == "CPU":
            def elem(x, col, row, _): return x[col+row][0] # k is always 0
            def c_map(_, elem): return (elem%16, elem//16)
            ul[i] = wmma_helper(1, 1, 16, 16, 256, elem, elem, c_map)
@@ -16,7 +16,6 @@ from tinygrad.helpers import getenv, DEBUG, fromimport, unwrap, LazySeq, Timing
 from tinygrad.engine.jit import GraphRunner, MultiGraphRunner, ExecItem, graph_class
 from tinygrad.engine.realize import CompiledRunner, BufferXfer
 from tinygrad.device import Compiled, Buffer, Allocator, Compiler, Device, BufferSpec
-from tinygrad.runtime.support.ib import IBCtx, IBConn, SGE

 # ***** API *****

@@ -36,7 +35,6 @@ class RemoteProperties:
  offset_supported: bool
  graph_supported: bool
  graph_supports_multi: bool
-  ib_gid: bytes|None

@dataclass(frozen=True)
 class GetProperties(RemoteRequest): pass
@@ -47,18 +45,12 @@ class Event(RemoteRequest): event_session: SessionKey; event: int # noqa: E702
@dataclass(frozen=True)
 class Wait(RemoteRequest): event: int

-@dataclass(frozen=True)
-class IBConnect(RemoteRequest): host: str; gid: bytes; qp_num: int # noqa: E702
-
@dataclass(frozen=True)
 class BufferAlloc(RemoteRequest): buffer_num: int; size: int; options: BufferSpec # noqa: E702

@dataclass(frozen=True)
 class BufferOffset(RemoteRequest): buffer_num: int; size: int; offset: int; sbuffer_num: int # noqa: E702

-@dataclass(frozen=True)
-class BufferIOVAS(RemoteRequest): buffer_nums: list[tuple[SessionKey, int]] # noqa: E702
-
@dataclass(frozen=True)
 class BufferFree(RemoteRequest): buffer_num: int # noqa: E702

@@ -119,9 +111,9 @@ class GraphExec(RemoteRequest):
  wait: bool

 # for safe deserialization
-eval_globals = {x.__name__:x for x in [SessionKey, SessionFree, RemoteProperties, GetProperties, Event, Wait, BufferAlloc, BufferOffset, BufferIOVAS,
-                                       BufferFree, CopyIn, CopyOut, Transfer, BatchTransfer, IBConnect, ProgramAlloc, ProgramFree, ProgramExec,
-                                       GraphComputeItem, GraphAlloc, GraphFree, GraphExec, BufferSpec, UOp, Ops, dtypes]}
+eval_globals = {x.__name__:x for x in [SessionKey, SessionFree, RemoteProperties, GetProperties, Event, Wait, BufferAlloc, BufferOffset, BufferFree,
+                                       CopyIn, CopyOut, Transfer, BatchTransfer, ProgramAlloc, ProgramFree, ProgramExec, GraphComputeItem, GraphAlloc,
+                                       GraphFree, GraphExec, BufferSpec, UOp, Ops, dtypes]}
 attribute_whitelist: dict[Any, set[str]] = {dtypes: {*DTYPES_DICT.keys(), 'imagef', 'imageh'}, Ops: {x.name for x in Ops}}
 eval_fxns = {ast.Constant: lambda x: x.value, ast.Tuple: lambda x: tuple(map(safe_eval, x.elts)), ast.List: lambda x: list(map(safe_eval, x.elts)),
  ast.Dict: lambda x: {safe_eval(k):safe_eval(v) for k,v in zip(x.keys, x.values)},
@@ -168,12 +160,6 @@ class RemoteHandler:
    self.base_device = base_device
    self.sessions: defaultdict[SessionKey, RemoteSession] = defaultdict(RemoteSession)

-    try: self.ib_ctx: IBCtx|None = IBCtx(getenv("IB_DEV", 0))
-    except (IndexError, AttributeError): self.ib_ctx = None
-    self.ib_lock = asyncio.Lock()
-    self.ib_conns: dict[str, IBConn|None] = {}
-    self.iova_cache: dict[tuple[SessionKey, int], tuple[int, int, int]] = {}
-
  async def __call__(self, reader:asyncio.StreamReader, writer:asyncio.StreamWriter):
    while (req_hdr:=(await reader.readline()).decode().strip()):
      req_method, req_path, _ = req_hdr.split(' ')
@@ -185,30 +171,6 @@ class RemoteHandler:
      res_status, res_body = await self.handle(req_method, req_path, req_body)
      writer.write(f"HTTP/1.1 {res_status.value} {res_status.phrase}\r\nContent-Length: {len(res_body)}\r\n\r\n".encode() + res_body)

-  async def ib_connect(self, ssession:SessionKey, dsession:SessionKey) -> IBConn|None:
-    if self.ib_ctx is None: return None
-    await self.ib_lock.acquire()
-    conn = RemoteConnection(dsession.host)
-    if dsession.host not in self.ib_conns:
-      props = safe_eval(ast.parse(conn.q(GetProperties(session=dsession), wait=True), mode="eval").body)
-      if props.ib_gid is not None:
-        self.ib_conns[dsession.host] = ib_conn = IBConn(self.ib_ctx)
-        ibxc_ret = conn.q(IBConnect(ssession.host, ib_conn.gid, ib_conn.qp_num, session=dsession), wait=True)
-        ib_conn.connect(*struct.unpack('<16sQ', ibxc_ret))
-      else:
-        self.ib_conns[dsession.host] = None
-    self.ib_lock.release()
-    return self.ib_conns[dsession.host]
-
-  async def get_iovas(self, bufs:list[tuple[SessionKey, int]]) -> list[tuple[int, int, int]]:
-    await self.ib_lock.acquire()
-    if (rbufs:=[buf for buf in bufs if buf not in self.iova_cache]):
-      conn = RemoteConnection(rbufs[0][0].host)
-      resp = await conn.aq(BufferIOVAS(rbufs, session=rbufs[0][0]), wait=True)
-      self.iova_cache.update({rbuf: struct.unpack('<QQQ', resp[i*24:(i+1)*24]) for i,rbuf in enumerate(rbufs)})
-    self.ib_lock.release()
-    return [self.iova_cache[buf] for buf in bufs]
-
  async def handle(self, method:str, path:str, body:bytes) -> tuple[http.HTTPStatus, bytes]:
    status, ret = http.HTTPStatus.OK, b""
    if path == "/batch" and method == "POST":
@@ -226,7 +188,6 @@ class RemoteHandler:
            rp = RemoteProperties(
              real_device=dev.device, renderer=(cls.__module__, cls.__name__, args), offset_supported=hasattr(dev.allocator, '_offset'),
              graph_supported=graph_cls is not None, graph_supports_multi=graph_cls is not None and issubclass(graph_cls, MultiGraphRunner),
-              ib_gid=bytes(self.ib_ctx.gid_attr.raw) if self.ib_ctx is not None else None,
            )
            ret = repr(rp).encode()
          case Event():
@@ -239,19 +200,9 @@ class RemoteHandler:
          case Wait():
            assert await session.events[c.event].wait()
            del session.events[c.event] # do not leak memory
-          case IBConnect():
-            self.ib_conns[c.host] = ibc = IBConn(unwrap(self.ib_ctx))
-            ibc.connect(c.gid, c.qp_num)
-            ret = struct.pack('<16sQ', ibc.gid, ibc.qp_num)
          case BufferAlloc():
            assert c.buffer_num not in session.buffers, f"buffer {c.buffer_num} already allocated"
            session.buffers[c.buffer_num] = Buffer(dev.device, c.size, dtypes.uint8, options=c.options, preallocate=True)
-          case BufferIOVAS():
-            rets = []
-            for buffer_session,buffer_num in c.buffer_nums:
-              iova, mr = unwrap(self.ib_ctx).reg(buf:=self.sessions[buffer_session].buffers[buffer_num])
-              rets.append(struct.pack("<QQQ", iova, mr.contents.rkey, buf.nbytes))
-            ret = b"".join(rets)
          case BufferOffset():
            assert c.buffer_num not in session.buffers, f"buffer {c.buffer_num} already exists"
            session.buffers[c.buffer_num] = session.buffers[c.sbuffer_num].view(c.size, dtypes.uint8, c.offset).allocate()
@@ -269,29 +220,16 @@ class RemoteHandler:
                sbuf.copyout(data:=memoryview(bytearray(sbuf.nbytes)))
                dbuf.copyin(data)
            else:
-              conn, ib_conn = RemoteConnection(c.dsession.host), await self.ib_connect(unwrap(c.session), c.dsession)
+              conn = RemoteConnection(c.dsession.host)
              sbuf = session.buffers[c.buffer_num]
-              if ib_conn is not None:
-                src_iova, src_mr = unwrap(self.ib_ctx).reg(sbuf)
-                dst_iova, dst_key, dst_size = (await self.get_iovas([(c.dsession, c.dbuffer_num)]))[0]
-                assert sbuf.nbytes == dst_size, f"{sbuf.nbytes} != {dst_size}"
-                for d in Device._opened_devices: Device[d].synchronize()
-                ib_conn.rdma_write([SGE(dst_iova, dst_key, src_iova, src_mr.contents.lkey, dst_size)])
-              else:
-                sbuf.copyout(data:=memoryview(bytearray(sbuf.nbytes)))
-                await conn.aq(CopyIn(c.dbuffer_num, conn.req.h(data), session=c.dsession), wait=True)
+              sbuf.copyout(data:=memoryview(bytearray(sbuf.nbytes)))
+              conn.q(CopyIn(c.dbuffer_num, conn.req.h(data), session=c.dsession), wait=True)
          case BatchTransfer():
-            conn, ib_conn = RemoteConnection(c.dbuffer_nums[0][0].host), await self.ib_connect(c.sbuffer_nums[0][0], c.dbuffer_nums[0][0])
-            if ib_conn is not None:
-              sbufs = [unwrap(self.ib_ctx).reg(self.sessions[s].buffers[bi]) for s,bi in c.sbuffer_nums]
-              dbufs = await self.get_iovas(c.dbuffer_nums)
-              for d in Device._opened_devices: Device[d].synchronize()
-              ib_conn.rdma_write([SGE(di, dk, si, sm.contents.lkey, ds) for (di,dk,ds),(si,sm) in zip(dbufs, sbufs)])
-            else:
-              for (sbuf_session,sbuf_num),(dbuf_session,dbuf_num) in zip(c.sbuffer_nums, c.dbuffer_nums):
-                sbuf = self.sessions[sbuf_session].buffers[sbuf_num]
-                sbuf.copyout(data:=memoryview(bytearray(sbuf.nbytes)))
-                await conn.aq(CopyIn(dbuf_num, conn.req.h(data), session=dbuf_session), wait=True)
+            conn = RemoteConnection(c.dbuffer_nums[0][0].host)
+            for (sbuf_session,sbuf_num),(dbuf_session,dbuf_num) in zip(c.sbuffer_nums, c.dbuffer_nums):
+              sbuf = self.sessions[sbuf_session].buffers[sbuf_num]
+              sbuf.copyout(data:=memoryview(bytearray(sbuf.nbytes)))
+              await conn.aq(CopyIn(dbuf_num, conn.req.h(data), session=dbuf_session), wait=True)
          case ProgramAlloc():
            lib = dev.compiler.compile_cached(req._h[c.datahash].decode())
            session.programs[(c.name, c.datahash)] = dev.runtime(c.name, lib)
@@ -1,7 +1,7 @@
 import functools, struct
 from tinygrad.device import  Compiled, Allocator, Compiler, BufferSpec
 from tinygrad.renderer.wgsl import WGSLRenderer
-from tinygrad.helpers import round_up, suppress_finalizing
+from tinygrad.helpers import round_up
 from tinygrad.runtime.autogen import webgpu
 from typing import List, Any, TypeAlias
 import ctypes
@@ -188,8 +188,9 @@ class WebGpuAllocator(Allocator['WGPUDevPtr']):
  def _copyout(self, dest:memoryview, src:WGPUBufPtr):
    buffer_data = read_buffer(self.dev, src)
    dest[:] = buffer_data[:dest.nbytes] if webgpu.wgpuBufferGetSize(src)  > dest.nbytes else buffer_data
-  @suppress_finalizing
-  def _free(self, opaque:WGPUBufPtr, options:BufferSpec): webgpu.wgpuBufferDestroy(opaque)
+  def _free(self, opaque:WGPUBufPtr, options:BufferSpec):
+    try: webgpu.wgpuBufferDestroy(opaque)
+    except AttributeError: pass

 class WebGpuDevice(Compiled):
  def __init__(self, device:str):
@@ -169,12 +169,12 @@ class AM_SMU(AM_IP):
      self._send_msg(self.smu_mod.PPSMC_MSG_SetSoftMinByFreq, clck << 16 | (vals[level]))
      self._send_msg(self.smu_mod.PPSMC_MSG_SetSoftMaxByFreq, clck << 16 | (vals[level]))

-  def _smu_cmn_send_msg(self, msg:int, param=0, debug=False):
+  def _smu_cmn_send_msg(self, msg, param=0, debug=False):
    (self.adev.mmMP1_SMN_C2PMSG_90 if not debug else self.adev.mmMP1_SMN_C2PMSG_54).write(0) # resp reg
    (self.adev.mmMP1_SMN_C2PMSG_82 if not debug else self.adev.mmMP1_SMN_C2PMSG_53).write(param)
    (self.adev.mmMP1_SMN_C2PMSG_66 if not debug else self.adev.mmMP1_SMN_C2PMSG_75).write(msg)

-  def _send_msg(self, msg:int, param:int, read_back_arg=False, timeout=10000, debug=False): # default timeout is 10 seconds
+  def _send_msg(self, msg, param, read_back_arg=False, timeout=10000, debug=False): # 10s
    self._smu_cmn_send_msg(msg, param, debug=debug)
    wait_cond(lambda: (self.adev.mmMP1_SMN_C2PMSG_90 if not debug else self.adev.mmMP1_SMN_C2PMSG_54).read(), value=1, timeout_ms=timeout,
      msg=f"SMU msg {msg:#x} timeout")
@@ -414,12 +414,12 @@ class AM_PSP(AM_IP):

  def _wait_for_bootloader(self): wait_cond(lambda: self.adev.reg(f"{self.reg_pref}_35").read() & 0x80000000, value=0x80000000, msg="BL not ready")

-  def _prep_msg1(self, data:memoryview):
+  def _prep_msg1(self, data):
    assert len(data) <= self.msg1_view.nbytes, f"msg1 buffer is too small {len(data):#x} > {self.msg1_view.nbytes:#x}"
    self.msg1_view[:len(data)+4] = bytes(data) + b'\x00' * 4
    self.adev.gmc.flush_hdp()

-  def _bootloader_load_component(self, fw:int, compid:int):
+  def _bootloader_load_component(self, fw, compid):
    if fw not in self.adev.fw.sos_fw: return 0

    self._wait_for_bootloader()
@@ -458,7 +458,7 @@ class AM_PSP(AM_IP):

    wait_cond(lambda: self.adev.reg(f"{self.reg_pref}_64").read() & 0x8000FFFF, value=0x80000000, msg="sOS ring not created")

-  def _ring_submit(self, cmd:am.struct_psp_gfx_cmd_resp) -> am.struct_psp_gfx_cmd_resp:
+  def _ring_submit(self, cmd):
    msg = am.struct_psp_gfx_rb_frame(fence_value=(prev_wptr:=self.adev.reg(f"{self.reg_pref}_67").read()),
      cmd_buf_addr_lo=lo32(self.adev.paddr2mc(self.cmd_paddr)), cmd_buf_addr_hi=hi32(self.adev.paddr2mc(self.cmd_paddr)),
      fence_addr_lo=lo32(self.adev.paddr2mc(self.fence_paddr)), fence_addr_hi=hi32(self.adev.paddr2mc(self.fence_paddr)))
@@ -477,7 +477,7 @@ class AM_PSP(AM_IP):

    return resp

-  def _load_ip_fw_cmd(self, fw_types:list[int], fw_bytes:memoryview):
+  def _load_ip_fw_cmd(self, fw_types, fw_bytes):
    self._prep_msg1(fw_bytes)
    for fw_type in fw_types:
      if DEBUG >= 2: print(f"am {self.adev.devfmt}: loading fw: {am.psp_gfx_fw_type__enumvalues[fw_type]}")
@@ -487,7 +487,7 @@ class AM_PSP(AM_IP):
      cmd.cmd.cmd_load_ip_fw.fw_type = fw_type
      self._ring_submit(cmd)

-  def _tmr_load_cmd(self) -> am.struct_psp_gfx_cmd_resp:
+  def _tmr_load_cmd(self):
    cmd = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_SETUP_TMR)
    cmd.cmd.cmd_setup_tmr.buf_phy_addr_hi, cmd.cmd.cmd_setup_tmr.buf_phy_addr_lo = data64(self.adev.paddr2mc(self.tmr_paddr))
    cmd.cmd.cmd_setup_tmr.system_phy_addr_hi, cmd.cmd.cmd_setup_tmr.system_phy_addr_lo = data64(self.tmr_paddr)
@@ -495,7 +495,7 @@ class AM_PSP(AM_IP):
    cmd.cmd.cmd_setup_tmr.buf_size = self.tmr_size
    return self._ring_submit(cmd)

-  def _load_toc_cmd(self, toc_size:int) -> am.struct_psp_gfx_cmd_resp:
+  def _load_toc_cmd(self, toc_size):
    cmd = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_TOC)
    cmd.cmd.cmd_load_toc.toc_phy_addr_hi, cmd.cmd.cmd_load_toc.toc_phy_addr_lo = data64(self.msg1_addr)
    cmd.cmd.cmd_load_toc.toc_size = toc_size
@@ -1,8 +1,6 @@
 from __future__ import annotations
 from typing import cast, Callable, Type, TypeVar, Generic, Any
 import contextlib, decimal, statistics, time, ctypes, array, os, struct, traceback, collections
-try: import fcntl # windows misses that
-except ImportError: fcntl = None #type:ignore[assignment]
 from tinygrad.helpers import PROFILE, getenv, to_mv, round_up, ProfileRangeEvent
 from tinygrad.renderer import Renderer
 from tinygrad.device import BufferSpec, Compiler, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent
@@ -27,7 +25,9 @@ class FileIOInterface:
    self.fd:int = fd or os.open(path, flags)
  def __del__(self):
    if hasattr(self, 'fd'): os.close(self.fd)
-  def ioctl(self, request, arg): return fcntl.ioctl(self.fd, request, arg)
+  def ioctl(self, request, arg):
+    import fcntl # to support windows
+    return fcntl.ioctl(self.fd, request, arg)
  def mmap(self, start, sz, prot, flags, offset):
    x = libc.mmap(start, sz, prot, flags, self.fd, offset)
    if x == 0xffffffffffffffff: raise OSError(f"Failed to mmap {sz} bytes at {hex(start)}: {os.strerror(ctypes.get_errno())}")
@@ -406,8 +406,6 @@ class HCQCompiled(Compiled, Generic[SignalType]):
    return self.signal_t(base_buf=HCQCompiled.signal_pool[pg].pop(), owner=self, **kwargs)

  def _at_profile_finalize(self):
-    self.synchronize() # Expect device to be synchronizes
-
    def _sync(d:HCQCompiled, q_t:Callable[[], HWQueue]):
      q_t().timestamp(d.timeline_signal).signal(d.timeline_signal, d.next_timeline()).submit(d)
      st = time.perf_counter_ns()
@@ -439,8 +437,6 @@ class HCQCompiled(Compiled, Generic[SignalType]):
      except Exception: errs += f"\n{iface_t.__name__}: {traceback.format_exc()}"
    raise RuntimeError(f"Cannot find a usable interface for {type(self).__name__[:-6]}:{self.device_id}:\n{errs}")

-  def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] in ("CPU", "LLVM")
-
  def finalize(self):
    try: self.synchronize() # Try to finalize device in any case.
    except RuntimeError as e: print(f"{self.device} synchronization failed before finalizing: {e}")
@@ -452,8 +448,7 @@ class HCQBuffer:
  def __init__(self, va_addr:sint, size:int, texture_info:Any=None, meta:Any=None, _base:HCQBuffer|None=None, view:MMIOInterface|None=None,
               owner:HCQCompiled|None=None):
    self.va_addr, self.size, self.texture_info, self.meta, self._base, self.view = va_addr, size, texture_info, meta, _base, view
-    self._devs, self.owner = ([owner] if owner is not None else []), owner
-    self._mappings:dict[HCQCompiled, HCQBuffer] = {} # mapping to the other devices
+    self.devs, self.owner = ([owner] if owner is not None else []), owner

  def offset(self, offset:int=0, size:int|None=None) -> HCQBuffer:
    return HCQBuffer(self.va_addr+offset, size or (self.size - offset), owner=self.owner, texture_info=self.texture_info, meta=self.meta,
@@ -464,10 +459,7 @@ class HCQBuffer:
    return self.view

  @property
-  def mappings(self): return self._mappings if self._base is None else self._base._mappings
-
-  @property
-  def mapped_devs(self): return self._devs if self._base is None else self._base._devs
+  def mapped_devs(self): return self.devs if self._base is None else self._base.devs

 class HCQAllocatorBase(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
  """
@@ -485,10 +477,7 @@ class HCQAllocatorBase(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
    if self.dev in buf.mapped_devs: return
    if buf.owner is None: raise RuntimeError(f"map failed: buffer {buf.va_addr} has no owner, it's a virtual buffer")
    if not hasattr(self, '_map'): raise NotImplementedError("map failed: no method implemented")
-
-    # Since it's unified memory space, any buffer mapping is valid for all devices after successful map.
-    # Devices can save mappings and internal metadata as a new buffer.
-    if (mb:=self._map(buf)) is not None: buf.mappings[self.dev] = mb
+    self._map(buf)
    buf.mapped_devs.append(self.dev)

  def _offset(self, buf, size:int, offset:int) -> HCQBuffer: return buf.offset(offset=offset, size=size)
@@ -496,7 +485,7 @@ class HCQAllocatorBase(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
 class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
  def _copyin(self, dest:HCQBuffer, src:memoryview):
    assert self.dev.hw_copy_queue_t is not None
-    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"TINY -> {self.dev.device}", enabled=PROFILE):
+    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"CPU -> {self.dev.device}", enabled=PROFILE):
      for i in range(0, src.nbytes, self.b[0].size):
        self.b_next = (self.b_next + 1) % len(self.b)
        self.dev.timeline_signal.wait(self.b_timeline[self.b_next])
@@ -528,7 +517,7 @@ class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
    self.dev.synchronize()

    assert self.dev.hw_copy_queue_t is not None
-    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"{self.dev.device} -> TINY", enabled=PROFILE):
+    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"{self.dev.device} -> CPU", enabled=PROFILE):
      for i in range(0, dest.nbytes, cp_size:=(self.max_copyout_size or self.b[0].size)):
        self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
                                  .copy(self.b[0].va_addr, src.va_addr+i, lsize:=min(cp_size, dest.nbytes-i)) \
@@ -1,172 +0,0 @@
-from __future__ import annotations
-import resource, ctypes, weakref, functools, itertools, tinygrad.runtime.autogen.ib as ib
-from typing import Iterator
-from dataclasses import dataclass
-from weakref import WeakKeyDictionary
-from tinygrad.device import Buffer, DMACPURef, DMAFdRef
-from tinygrad.helpers import getenv, round_up, DEBUG
-
-DEFAULT_PORT, DEFAULT_GID = getenv("DEFAULT_PORT", 1), getenv("DEFAULT_GID", 3) # DEFAULT_GID=0 for RXE
-IOVA_ALIGN = resource.getpagesize()
-
-def checkz(x, ret=None):
-  assert x == 0, f'{x} != 0 (errno {ctypes.get_errno()})'
-  return ret
-
-@dataclass(frozen=True)
-class SGE:
-  dst_iova: int
-  dst_key: int
-  src_iova: int
-  src_key: int
-  size: int
-
-class IBCtx:
-  def __init__(self, idx:int):
-    # Open the device (aka Host Channel Adapter in ib-speak)
-    devs = ib.ibv_get_device_list(ctypes.byref(ndevs:=ctypes.c_int32()))
-    if idx >= ndevs.value: raise IndexError(f"{idx} > {ndevs.value}")
-    self.ctx = ib.ibv_open_device(devs[idx])
-    ib.ibv_free_device_list(devs)
-
-    # HACK: remove this (and all usage of `ctx.contents.ops`) when clang2py can deal with `static inline` wrapper-functions
-    self.vctx = ctypes.cast(ctypes.addressof(self.ctx.contents) - ib.struct_verbs_context.context.offset, ctypes.POINTER(ib.struct_verbs_context))
-
-    # Get attributes. Something like port_attr.max_msg_sz sound like it might requre taking the min of host's and remote's attributes if they differ
-    self.device_attr = checkz(ib.ibv_query_device(self.ctx, ctypes.byref(da:=ib.struct_ibv_device_attr())), da)
-    self.port_attr = checkz(self.vctx.contents.query_port(self.ctx, DEFAULT_PORT, ctypes.byref(pa:=ib.struct_ibv_port_attr()), ctypes.sizeof(pa)), pa)
-    self.gid_attr = checkz(ib.ibv_query_gid(self.ctx, DEFAULT_PORT, DEFAULT_GID, ctypes.byref(ga:=ib.union_ibv_gid())), ga)
-
-    # Allocate protection domain
-    self.pd = ib.ibv_alloc_pd(self.ctx)
-    self.next_iova: int = IOVA_ALIGN # don't start at zero (nullptr)
-
-    # weakref(buf) => (iova, mr, mr_dealloc). mr_dealloc is kept here to avoid double freeing mrs that are deallocated in __del__
-    self.mrs: WeakKeyDictionary[Buffer, tuple[int, ctypes._Pointer[ib.struct_ibv_mr], weakref.finalize]] = WeakKeyDictionary()
-
-    # Default soft fd limit is 1024, which is not enough, set soft to hard (maximum allowed by the os)
-    IBCtx.rlimit_fix()
-
-  def __del__(self):
-    # must deallocate all mrs in protection domain before deallocating the protection domain
-    if hasattr(self, "mrs"): [fin() for _,_,fin in self.mrs.values()]
-    if hasattr(self, "pd"): ib.ibv_dealloc_pd(self.pd)
-    if hasattr(self, "ctx"): ib.ibv_close_device(self.ctx)
-
-  @functools.cache # run once
-  @staticmethod
-  def rlimit_fix():
-    soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
-    resource.setrlimit(resource.RLIMIT_NOFILE, (hard, hard))
-    if DEBUG>=2: print(f"IB: Increased fd limit from {soft} to {hard}")
-
-  def alloc_iova(self, size:int, required_offset:int):
-    iova = round_up(self.next_iova - required_offset, IOVA_ALIGN) + required_offset
-    self.next_iova = iova + size
-    return iova
-
-  def reg(self, buf:Buffer) -> tuple[int, ctypes._Pointer[ib.struct_ibv_mr]]:
-    buf = buf.base
-    if buf not in self.mrs:
-      if buf.nbytes > self.device_attr.max_mr_size: raise RuntimeError(f"Buffer too big: {buf.nbytes:#x} > {self.device_attr.max_mr_size:#x}")
-      if len(self.mrs) >= self.device_attr.max_mr: raise RuntimeError(f"Out of memory region cap: {len(self.mrs)} >= {self.device_attr.max_mr}")
-      # Local read is implied (but still have to create the memory region, except for short sends/writes with IBV_SEND_INLINE that are inlined by cpu)
-      mr_flags = ib.IBV_ACCESS_LOCAL_WRITE | ib.IBV_ACCESS_REMOTE_READ | ib.IBV_ACCESS_REMOTE_WRITE
-      match (dmaref:=buf.as_dmaref()):
-        case DMACPURef():
-          iova = self.alloc_iova(dmaref.size, dmaref.addr % IOVA_ALIGN)
-          mr = ib.ibv_reg_mr_iova2(self.pd, ctypes.c_void_p(dmaref.addr), dmaref.size, iova, mr_flags)
-        case DMAFdRef():
-          iova = self.alloc_iova(dmaref.size, dmaref.offset % IOVA_ALIGN)
-          mr = ib.ibv_reg_dmabuf_mr(self.pd, dmaref.offset, dmaref.size, iova, dmaref.fd, mr_flags)
-        case _: raise RuntimeError(f"Unknown type of dma ref: {dmaref}")
-      if not mr: raise RuntimeError(f"Couldn't register memory region for {buf} {dmaref} (errno={ctypes.get_errno()})")
-      self.mrs[buf] = (iova, mr, weakref.finalize(buf, ib.ibv_dereg_mr, mr))
-    return self.mrs[buf][0:2]
-
-class IBConn:
-  def __init__(self, ctx:IBCtx):
-    self.ctx = ctx
-
-    # Create Completion Channel. It is a file descriptor that kernel sends notifications through, not a thing in infiniband spec, just linux-ism
-    self.comp_channel = ib.ibv_create_comp_channel(self.ctx.ctx)
-    # Create Completion Queue. When a Work Request with signaled flag is completed a Completion Queue Entry is pushed onto this queue
-    self.cq = ib.ibv_create_cq(self.ctx.ctx, _capacity:=256, _cq_context:=None, self.comp_channel, _comp_vector:=0)
-    self.pending_wrids: set[int] = set()
-    self.wrid_num: Iterator[int] = itertools.count(0) # wc_id is uint64, this will never overflow
-
-    # Create Queue Pair. It's the closest thing to a socket in infiniband with QP num being the closest thing to a port, except it's allocated by hca
-    qp_init_attrs_cap = ib.struct_ibv_qp_cap(max_send_wr=1024, max_recv_wr=64, max_send_sge=8, max_recv_sge=8, max_inline_data=64)
-    qp_init_attrs = ib.struct_ibv_qp_init_attr(send_cq=self.cq, recv_cq=self.cq, cap=qp_init_attrs_cap, qp_type=ib.IBV_QPT_RC) # Reliable Connection
-    self.qp = ib.ibv_create_qp(self.ctx.pd, ctypes.byref(qp_init_attrs))
-    self.qp_cap = qp_init_attrs.cap
-
-    # The most important thing about QPs is their state, when a new QP is created it's in the RESET state, before it can be properly used it has to go
-    # through Init, Ready To Receive, Ready To Send. A good docs on QP state machine: https://www.rdmamojo.com/2012/05/05/qp-state-machine/
-
-    # INIT
-    qp_access_flags = ib.IBV_ACCESS_REMOTE_WRITE | ib.IBV_ACCESS_REMOTE_READ
-    qpa = ib.struct_ibv_qp_attr(qp_state=ib.IBV_QPS_INIT, port_num=DEFAULT_PORT, qp_access_flags=qp_access_flags)
-    checkz(ib.ibv_modify_qp(self.qp, qpa, ib.IBV_QP_STATE | ib.IBV_QP_PORT | ib.IBV_QP_ACCESS_FLAGS | ib.IBV_QP_PKEY_INDEX))
-
-    self.gid, self.qp_num = bytes(self.ctx.gid_attr.raw), self.qp.contents.qp_num
-
-  # Exchange GID and QP num with remote. At least in RoCEv2 gid can be guessed from remote's ip, QP num can't.
-
-  def connect(self, remote_gid:bytes, remote_qp_num:int):
-    # RTR
-    qp_ah_attr_grh = ib.struct_ibv_global_route(hop_limit=1, dgid=ib.union_ibv_gid(raw=(ctypes.c_ubyte * 16)(*remote_gid)), sgid_index=DEFAULT_GID)
-    qp_ah_attr = ib.struct_ibv_ah_attr(is_global=1, port_num=DEFAULT_PORT, grh=qp_ah_attr_grh)
-    qpa = ib.struct_ibv_qp_attr(qp_state=ib.IBV_QPS_RTR, path_mtu=ib.IBV_MTU_4096, dest_qp_num=remote_qp_num, rq_psn=0, max_dest_rd_atomic=1,
-                                min_rnr_timer=12, ah_attr=qp_ah_attr)
-    checkz(ib.ibv_modify_qp(self.qp, qpa, ib.IBV_QP_STATE | ib.IBV_QP_PATH_MTU | ib.IBV_QP_DEST_QPN | ib.IBV_QP_RQ_PSN | \
-                                          ib.IBV_QP_MAX_DEST_RD_ATOMIC | ib.IBV_QP_MIN_RNR_TIMER | ib.IBV_QP_AV))
-
-    # RTS
-    qpa = ib.struct_ibv_qp_attr(qp_state=ib.IBV_QPS_RTS, timeout=14, retry_cnt=7, rnr_retry=7, sq_psn=0, max_rd_atomic=1)
-    checkz(ib.ibv_modify_qp(self.qp, qpa, ib.IBV_QP_STATE | ib.IBV_QP_TIMEOUT | ib.IBV_QP_RETRY_CNT | ib.IBV_QP_RNR_RETRY | ib.IBV_QP_SQ_PSN | \
-                                          ib.IBV_QP_MAX_QP_RD_ATOMIC))
-
-  def __del__(self):
-    self.wait_cq() # need to wait for **everything** to complete before it's safe to dealloc queues and stuff
-    ib.ibv_destroy_qp(self.qp)
-    ib.ibv_destroy_cq(self.cq)
-    ib.ibv_destroy_comp_channel(self.comp_channel)
-
-  def next_wrid(self):
-    self.pending_wrids.add(wrid:=next(self.wrid_num))
-    return wrid
-
-  def wait_cq(self, wr_id: int|None=None):
-    while (wr_id in self.pending_wrids) if wr_id is not None else self.pending_wrids:
-      if self.ctx.ctx.contents.ops.poll_cq(self.cq, _num_entries:=1, ctypes.byref(wc:=ib.struct_ibv_wc())):
-        if wc.status != ib.IBV_WC_SUCCESS:
-          raise RuntimeError(f'Work Request completed with error: wr_id={wc.wr_id} status={ib.ibv_wc_status__enumvalues.get(wc.status, wc.status)}')
-        self.pending_wrids.remove(wc.wr_id)
-
-  def rdma_write(self, sgl:list[SGE]):
-    swr: ctypes._Pointer[ib.struct_ibv_send_wr]|None = None
-    swr_cnt, wr_id = 0, self.next_wrid()
-    def _post():
-      nonlocal swr, swr_cnt, wr_id
-      if swr is not None:
-        # The swr can be freed when this returns, the memory that sge points to can be unmapped after work completion is retrieved from cq
-        checkz(self.ctx.ctx.contents.ops.post_send(self.qp, swr, ctypes.byref(_bad_wr:=ctypes.POINTER(ib.struct_ibv_send_wr)())))
-        # TODO: async
-        self.wait_cq(wr_id)
-        swr, swr_cnt, wr_id = None, 0, self.next_wrid()
-    # Everything is in reverse for elegant chaining
-    for sg in reversed(sgl):
-      # Message size limit (max 2GB per ib spec, 1GB on tinybox mellanoxes) applies to both scatter-gather entries and entire wrs
-      for off in reversed(range(0, sg.size, self.ctx.port_attr.max_msg_sz)):
-        # Scatter-Gather Entry for local memory
-        sge = ctypes.pointer(ib.struct_ibv_sge(addr=sg.src_iova+off, length=min(sg.size-off, self.ctx.port_attr.max_msg_sz), lkey=sg.src_key))
-        # RDMA struct for remote memory
-        wr = ib.union_ibv_send_wr_wr(rdma=ib.struct_ibv_send_wr_1_rdma(remote_addr=sg.dst_iova+off, rkey=sg.dst_key))
-        # Signal (with chosen work request id) if it's the last wr (first in the loop since it's reversed)
-        wid, flags = (wr_id, ib.IBV_SEND_SIGNALED) if swr is None else (0, 0)
-        # Create Send Request
-        swr = ctypes.pointer(ib.struct_ibv_send_wr(opcode=ib.IBV_WR_RDMA_WRITE, sg_list=sge, num_sge=1, wr=wr, wr_id=wid, send_flags=flags, next=swr))
-        # Flush if queue is being overrun
-        if (swr_cnt:=swr_cnt + 1) >= self.qp_cap.max_send_wr: _post()
-    _post()
@@ -1,6 +1,6 @@
 from __future__ import annotations
 import ctypes, time, array, struct, itertools, dataclasses
-from typing import cast, Any
+from typing import cast
 from tinygrad.runtime.autogen.nv import nv
 from tinygrad.helpers import to_mv, lo32, hi32, DEBUG, round_up, round_down, mv_address, fetch, wait_cond
 from tinygrad.runtime.support.system import System
@@ -8,7 +8,7 @@ from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.runtime.autogen import nv_gpu

@dataclasses.dataclass(frozen=True)
-class GRBufDesc: size:int; virt:bool; phys:bool; local:bool=False # noqa: E702
+class GRBufDesc: size:int; v:int; p:int; lc:int=0 # noqa: E702

 class NV_IP:
  def __init__(self, nvdev): self.nvdev = nvdev
@@ -26,13 +26,13 @@ class NVRpcQueue:
    self.gsp, self.va, self.queue_va, self.seq = gsp, va, va + self.tx.entryOff, 0
    self.queue_mv = to_mv(self.queue_va, self.tx.msgSize * self.tx.msgCount)

-  def _checksum(self, data:bytes):
+  def _checksum(self, data):
    if (pad_len:=(-len(data)) % 8): data += b'\x00' * pad_len
    checksum = 0
    for offset in range(0, len(data), 8): checksum ^= struct.unpack_from('Q', data, offset)[0]
    return hi32(checksum) ^ lo32(checksum)

-  def send_rpc(self, func:int, msg:bytes, wait=False):
+  def send_rpc(self, func, msg, wait=False):
    header = nv.rpc_message_header_v(signature=nv.NV_VGPU_MSG_SIGNATURE_VALID, rpc_result=nv.NV_VGPU_MSG_RESULT_RPC_PENDING,
      rpc_result_private=nv.NV_VGPU_MSG_RESULT_RPC_PENDING, header_version=(3<<24), function=func, length=len(msg) + 0x20)

@@ -49,7 +49,7 @@ class NVRpcQueue:
    self.seq += 1
    self.gsp.nvdev.NV_PGSP_QUEUE_HEAD[0].write(0x0)

-  def wait_resp(self, cmd:int) -> memoryview:
+  def wait_resp(self, cmd) -> memoryview:
    while True:
      System.memory_barrier()
      if self.rx.readPtr == self.tx.writePtr: continue
@@ -177,7 +177,7 @@ class NV_FLCN(NV_IP):
    self.nvdev.NV_PFALCON_FALCON_OS.with_base(self.falcon).write(0x0)
    assert self.nvdev.NV_PRISCV_RISCV_CPUCTL.with_base(self.falcon).read_bitfields()['active_stat'] == 1, "GSP Core is not active"

-  def execute_dma(self, base:int, cmd:int, dest:int, mem_off:int, sysmem:int, size:int):
+  def execute_dma(self, base, cmd, dest, mem_off, sysmem, size):
    wait_cond(lambda: self.nvdev.NV_PFALCON_FALCON_DMATRFCMD.with_base(base).read_bitfields()['full'], value=0, msg="DMA does not progress")

    self.nvdev.NV_PFALCON_FALCON_DMATRFBASE.with_base(base).write(lo32(sysmem >> 8))
@@ -194,7 +194,7 @@ class NV_FLCN(NV_IP):

    wait_cond(lambda: self.nvdev.NV_PFALCON_FALCON_DMATRFCMD.with_base(base).read_bitfields()['idle'], msg="DMA does not complete")

-  def start_cpu(self, base:int):
+  def start_cpu(self, base):
    if self.nvdev.NV_PFALCON_FALCON_CPUCTL.with_base(base).read_bitfields()['alias_en'] == 1:
      self.nvdev.wreg(base + self.nvdev.NV_PFALCON_FALCON_CPUCTL_ALIAS, 0x2)
    else: self.nvdev.NV_PFALCON_FALCON_CPUCTL.with_base(base).write(startcpu=1)
@@ -232,11 +232,11 @@ class NV_FLCN(NV_IP):
    if mailbox is not None:
      return self.nvdev.NV_PFALCON_FALCON_MAILBOX0.with_base(base).read(), self.nvdev.NV_PFALCON_FALCON_MAILBOX1.with_base(base).read()

-  def disable_ctx_req(self, base:int):
+  def disable_ctx_req(self, base):
    self.nvdev.NV_PFALCON_FBIF_CTL.with_base(base).update(allow_phys_no_ctx=1)
    self.nvdev.NV_PFALCON_FALCON_DMACTL.with_base(base).write(0x0)

-  def reset(self, base:int, riscv=False):
+  def reset(self, base, riscv=False):
    engine_reg = self.nvdev.NV_PGSP_FALCON_ENGINE if base == self.falcon else self.nvdev.NV_PSEC_FALCON_ENGINE
    engine_reg.write(reset=1)
    time.sleep(0.1)
@@ -408,10 +408,10 @@ class NV_GSP(NV_IP):
      assert self.nvdev.flcn.frts_offset == m.frtsOffset, f"FRTS mismatch: {self.nvdev.flcn.frts_offset} != {m.frtsOffset}"
    self.wpr_meta, self.wpr_meta_sysmem = self.nvdev._alloc_boot_struct(m)

-  def promote_ctx(self, client:int, subdevice:int, obj:int, ctxbufs:dict[int, GRBufDesc], bufs=None, virt=None, phys=None):
+  def promote_ctx(self, client, subdevice, obj, ctxbufs, bufs=None, virt=None, phys=None):
    res, prom = {}, nv_gpu.NV2080_CTRL_GPU_PROMOTE_CTX_PARAMS(entryCount=len(ctxbufs), engineType=0x1, hChanClient=client, hObject=obj)
    for i,(buf,desc) in enumerate(ctxbufs.items()):
-      use_v, use_p = (desc.virt if virt is None else virt), (desc.phys if phys is None else phys)
+      use_v, use_p = (desc.v if virt is None else virt), (desc.p if phys is None else phys)
      x = (bufs or {}).get(buf, self.nvdev.mm.valloc(desc.size, contiguous=True)) # allocate buffers
      prom.promoteEntry[i] = nv_gpu.NV2080_CTRL_GPU_PROMOTE_CTX_BUFFER_ENTRY(bufferId=buf, gpuVirtAddr=x.va_addr if use_v else 0, bInitialize=use_p,
        gpuPhysAddr=x.paddrs[0][0] if use_p else 0, size=desc.size if use_p else 0, physAttr=0x4 if use_p else 0, bNonmapped=(use_p and not use_v))
@@ -449,11 +449,10 @@ class NV_GSP(NV_IP):
    gr_size = _ctx_info(nv_gpu.NV0080_CTRL_FIFO_GET_ENGINE_CONTEXT_PROPERTIES_ENGINE_ID_GRAPHICS, add=0x40000)
    patch_size = _ctx_info(nv_gpu.NV0080_CTRL_FIFO_GET_ENGINE_CONTEXT_PROPERTIES_ENGINE_ID_GRAPHICS_PATCH)
    cfgs_sizes = {x: _ctx_info(x + 14, align=(2 << 20) if x == 5 else None) for x in range(3, 11)} # indices 3–10 are mapped to 17–24
-    self.grctx_bufs = {0: GRBufDesc(gr_size, phys=True, virt=True), 1: GRBufDesc(patch_size, phys=True, virt=True, local=True),
-      2: GRBufDesc(patch_size, phys=True, virt=True), **{x: GRBufDesc(cfgs_sizes[x], phys=False, virt=True) for x in range(3, 7)},
-      9: GRBufDesc(cfgs_sizes[9], phys=True, virt=True), 10: GRBufDesc(cfgs_sizes[10], phys=True, virt=False),
-      11: GRBufDesc(cfgs_sizes[10], phys=True, virt=True)} # NOTE: 11 reuses cfgs_sizes[10]
-    self.promote_ctx(self.priv_root, subdev, ch_gpfifo, {k:v for k, v in self.grctx_bufs.items() if not v.local})
+    self.grctx_bufs = {0: GRBufDesc(gr_size, p=1, v=1), 1: GRBufDesc(patch_size, p=1, v=1, lc=1), 2: GRBufDesc(patch_size, p=1, v=1),
+      **{x: GRBufDesc(cfgs_sizes[x], p=0, v=1) for x in range(3, 7)}, 9: GRBufDesc(cfgs_sizes[9], p=1, v=1),
+      10: GRBufDesc(cfgs_sizes[10], p=1, v=0), 11: GRBufDesc(cfgs_sizes[10], p=1, v=1)} # NOTE: 11 reuses cfgs_sizes[10]
+    self.promote_ctx(self.priv_root, subdev, ch_gpfifo, {k:v for k, v in self.grctx_bufs.items() if v.lc == 0})

    self.rpc_rm_alloc(hParent=ch_gpfifo, hClass=self.compute_class, params=None)
    self.rpc_rm_alloc(hParent=ch_gpfifo, hClass=self.dma_class, params=None)
@@ -474,7 +473,7 @@ class NV_GSP(NV_IP):

  ### RPCs

-  def rpc_rm_alloc(self, hParent:int, hClass:int, params:Any, client=None) -> int:
+  def rpc_rm_alloc(self, hParent, hClass, params, client=None) -> int:
    if hClass == self.gpfifo_class:
      ramfc_alloc = self.nvdev.mm.valloc(0x1000, contiguous=True)
      params.ramfcMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=ramfc_alloc.paddrs[0][0], size=0x200, addressSpace=2, cacheAttrib=0)
@@ -500,7 +499,7 @@ class NV_GSP(NV_IP):
      self.promote_ctx(client, self.subdevice, hParent, {k:v for k,v in self.grctx_bufs.items() if k in [0, 1, 2]}, phys_gr_ctx, phys=False)
    return obj if hClass != nv_gpu.NV1_ROOT else client

-  def rpc_rm_control(self, hObject:int, cmd:int, params:Any, client=None):
+  def rpc_rm_control(self, hObject, cmd, params, client=None):
    control_args = nv.rpc_gsp_rm_control_v(hClient=(client:=client or self.priv_root), hObject=hObject, cmd=cmd, flags=0x0,
      paramsSize=ctypes.sizeof(params) if params is not None else 0x0)
    self.cmd_q.send_rpc(nv.NV_VGPU_MSG_FUNCTION_GSP_RM_CONTROL, bytes(control_args) + (bytes(params) if params is not None else b''))
@@ -512,7 +511,7 @@ class NV_GSP(NV_IP):
      cast(nv_gpu.NVC36F_CTRL_CMD_GPFIFO_GET_WORK_SUBMIT_TOKEN_PARAMS, st).workSubmitToken |= (1 << 30)
    return st

-  def rpc_set_page_directory(self, device:int, hVASpace:int, pdir_paddr:int, client=None, pasid=0xffffffff):
+  def rpc_set_page_directory(self, device, hVASpace, pdir_paddr, client=None, pasid=0xffffffff):
    params = nv.struct_NV0080_CTRL_DMA_SET_PAGE_DIRECTORY_PARAMS_v1E_05(physAddress=pdir_paddr,
      numEntries=self.nvdev.mm.pte_cnt[0], flags=0x8, hVASpace=hVASpace, pasid=pasid, subDeviceId=1, chId=0) # flags field is all channels.
    alloc_args = nv.rpc_set_page_directory_v(hClient=client or self.priv_root, hDevice=device, pasid=pasid, params=params)
@@ -545,7 +544,7 @@ class NV_GSP(NV_IP):
    header = nv.PACKED_REGISTRY_TABLE(size=hdr_size + len(entries_bytes) + len(data_bytes), numEntries=len(table))
    self.cmd_q.send_rpc(nv.NV_VGPU_MSG_FUNCTION_SET_REGISTRY, bytes(header) + entries_bytes + data_bytes)

-  def run_cpu_seq(self, seq_buf:memoryview):
+  def run_cpu_seq(self, seq_buf):
    hdr = nv.rpc_run_cpu_sequencer_v17_00.from_address(mv_address(seq_buf))
    cmd_iter = iter(seq_buf[ctypes.sizeof(nv.rpc_run_cpu_sequencer_v17_00):].cast('I')[:hdr.cmdIndex])

@@ -71,7 +71,7 @@ class NVMemoryManager(MemoryManager):
  def on_range_mapped(self): self.dev.NV_VIRTUAL_FUNCTION_PRIV_MMU_INVALIDATE.write((1 << 0) | (1 << 1) | (1 << 6) | (1 << 31))

 class NVDev(PCIDevImplBase):
-  def __init__(self, devfmt:str, mmio:MMIOInterface, vram:MMIOInterface, venid:int, subvenid:int, rev:int, bars:dict):
+  def __init__(self, devfmt, mmio:MMIOInterface, vram:MMIOInterface, venid:int, subvenid:int, rev:int, bars:dict):
    self.devfmt, self.mmio, self.vram, self.venid, self.subvenid, self.rev, self.bars = devfmt, mmio, vram, venid, subvenid, rev, bars
    self.lock_fd = System.flock_acquire(f"nv_{self.devfmt}.lock")

@@ -101,10 +101,10 @@ class NVDev(PCIDevImplBase):
    for ip in [self.gsp, self.flcn]: ip.fini_hw()

  def reg(self, reg:str) -> NVReg: return self.__dict__[reg]
-  def wreg(self, addr:int, value:int):
+  def wreg(self, addr, value):
    self.mmio[addr // 4] = value
    if NV_DEBUG >= 4: print(f"wreg: {hex(addr)} = {hex(value)}")
-  def rreg(self, addr:int) -> int: return self.mmio[addr // 4]
+  def rreg(self, addr): return self.mmio[addr // 4]

  def _early_init(self):
    self.reg_names:set[str] = set()
@@ -134,12 +134,12 @@ class NVDev(PCIDevImplBase):

    self.vram_size = self.reg("NV_PGC6_AON_SECURE_SCRATCH_GROUP_42").read() << 20

-  def _alloc_boot_struct(self, struct:ctypes.Structure) -> tuple[ctypes.Structure, int]:
+  def _alloc_boot_struct(self, struct):
    va, paddrs = System.alloc_sysmem(sz:=ctypes.sizeof(type(struct)), contiguous=True)
    to_mv(va, sz)[:] = bytes(struct)
    return type(struct).from_address(va), paddrs[0]

-  def _download(self, file:str) -> str:
+  def _download(self, file) -> str:
    url = f"https://raw.githubusercontent.com/NVIDIA/open-gpu-kernel-modules/8ec351aeb96a93a4bb69ccc12a542bf8a8df2b6f/{file}"
    return fetch(url, subdir="defines").read_text()

@@ -12,20 +12,16 @@ class _System:

  def memory_barrier(self): lib.atomic_thread_fence(__ATOMIC_SEQ_CST:=5) if (lib:=self.atomic_lib()) is not None else None

-  def lock_memory(self, addr:int, size:int):
-    if libc.mlock(ctypes.c_void_p(addr), size): raise RuntimeError(f"Failed to lock memory at {addr:#x} with size {size:#x}")
-
-  def system_paddrs(self, vaddr:int, size:int) -> list[int]:
-    self.pagemap().seek(vaddr // mmap.PAGESIZE * 8)
-    return [(x & ((1<<55) - 1)) * mmap.PAGESIZE for x in array.array('Q', self.pagemap().read(size//mmap.PAGESIZE*8, binary=True))]
-
  def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False, data:bytes|None=None) -> tuple[int, list[int]]:
    assert not contiguous or size <= (2 << 20), "Contiguous allocation is only supported for sizes up to 2MB"
    flags = (libc.MAP_HUGETLB if contiguous and (size:=round_up(size, mmap.PAGESIZE)) > 0x1000 else 0) | (MAP_FIXED if vaddr else 0)
    va = FileIOInterface.anon_mmap(vaddr, size, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED|mmap.MAP_ANONYMOUS|MAP_POPULATE|MAP_LOCKED|flags, 0)

    if data is not None: to_mv(va, len(data))[:] = data
-    return va, self.system_paddrs(va, size)
+
+    # Read pagemap to get the physical address of each page. The pages are locked.
+    self.pagemap().seek(va // mmap.PAGESIZE * 8)
+    return va, [(x & ((1<<55) - 1)) * mmap.PAGESIZE for x in array.array('Q', self.pagemap().read(size//mmap.PAGESIZE*8, binary=True))]

  def pci_reset(self, gpu): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{gpu}/reset'")
  def pci_scan_bus(self, target_vendor:int, target_devices:list[int]) -> list[str]:
@@ -159,12 +155,6 @@ class PCIIfaceBase:
    if b.owner == self.dev and b.meta.has_cpu_mapping: FileIOInterface.munmap(b.va_addr, b.size)

  def map(self, b:HCQBuffer):
-    if b.owner is not None and b.owner._is_cpu():
-      System.lock_memory(cast(int, b.va_addr), b.size)
-      paddrs, snooped, uncached = [(x, 0x1000) for x in System.system_paddrs(cast(int, b.va_addr), round_up(b.size, 0x1000))], True, False
-    elif (ifa:=getattr(b.owner, "iface", None)) is not None and isinstance(ifa, PCIIfaceBase):
-      paddrs = [(paddr if b.meta.mapping.system else (paddr + ifa.p2p_base_addr), size) for paddr,size in b.meta.mapping.paddrs]
-      snooped, uncached = b.meta.mapping.snooped, b.meta.mapping.uncached
-    else: raise RuntimeError(f"map failed: {b.owner} -> {self.dev}")
-
-    self.dev_impl.mm.map_range(cast(int, b.va_addr), round_up(b.size, 0x1000), paddrs, system=True, snooped=snooped, uncached=uncached)
+    if (ifa:=getattr(b.owner, "iface", None)) is None or not isinstance(ifa, PCIIfaceBase): raise RuntimeError(f"map failed: {b.owner} -> {self.dev}")
+    paddrs = [(paddr if b.meta.mapping.system else (paddr + ifa.p2p_base_addr), size) for paddr,size in b.meta.mapping.paddrs]
+    self.dev_impl.mm.map_range(cast(int, b.va_addr), b.size, paddrs, system=True, snooped=b.meta.mapping.snooped, uncached=b.meta.mapping.uncached)
@@ -3,8 +3,8 @@ from tinygrad.uop.ops import UOp, Ops, GroupOp, PatternMatcher, UPat, graph_rewr
 from tinygrad.uop.ops import track_rewrites, _substitute
 from tinygrad.uop.spec import type_verify, tensor_uop_spec
 from tinygrad.uop.symbolic import symbolic_simple
-from tinygrad.helpers import Metadata, all_int, all_same, colored, prod, dedup, unwrap, getenv, pluralize, FUSE_ARANGE, DEBUG, SPLIT_REDUCEOP, argsort
-from tinygrad.dtype import ImageDType, dtypes
+from tinygrad.helpers import Metadata, all_int, all_same, colored, prod, dedup, unwrap, getenv, pluralize, FUSE_ARANGE, DEBUG, SPLIT_REDUCEOP
+from tinygrad.dtype import ImageDType
 from tinygrad.schedule.multi import multi_pm
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View, strides_for_shape, get_contraction_with_reduce
@@ -188,7 +188,7 @@ def reduce_push_add_ones(src:UOp, r:UOp, view:UOp):

 view_left = merge_views+PatternMatcher([
  # view before elementwise and buffer ops
-  (UPat(Ops.VIEW, src=(UPat({*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.BIND, Ops.LOAD, Ops.STORE, Ops.VALID, Ops.SINK}, name="e"),), name="view"),
+  (UPat(Ops.VIEW, src=(UPat({*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.BIND, Ops.LOAD, Ops.STORE, Ops.VALID}, name="e"),), name="view"),
   lambda e,view: e.replace(src=tuple(s.view(view.st) for s in e.src))),
  # if there's ones added after reduce, put this before the reduce
  (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), reduce_push_add_ones),
@@ -258,8 +258,7 @@ add_buffer_ops = PatternMatcher([
  # passthrough ASSIGN
  (UPat(Ops.ASSIGN, name="x"), lambda x: x.src[1]),
  # VALID
-  (UPat(Ops.VIEW, src=(UPat.cvar(),), name="self"),
-   lambda self: UOp.where(UOp(Ops.VALID, dtypes.bool, (UOp(Ops.VIEW, arg=self.st),)), self.const_like(self.base.arg), 0)),
+  (UPat(Ops.VIEW, src=(UPat.cvar(),), name="self"), UOp.valid),
 ])

 def check_load_st(glbl:UOp, view:UOp):
@@ -417,121 +416,6 @@ finalize_contiguous = PatternMatcher([

 remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])

-def contiguous_create_ranges(ctx:list[int], x:UOp):
-  if len(x.src) != 1: return None
-  ranges = []
-  for s in x.shape:
-    if resolve(s!=1):
-      ranges.append(UOp.range(dtypes.int, s, ctx[0]))
-      ctx[0] += 1
-    else:
-      ranges.append(UOp.const(dtypes.int, 0))
-  mm = UOp(Ops.MAP, dtype=x.src[0].dtype, src=(x.src[0],)+tuple(ranges))
-  buf = UOp.new_buffer(x.device, prod(x.shape), x.dtype).reshape(x.shape)
-  mm2 = UOp(Ops.MAP, dtype=x.src[0].dtype, src=(buf,)+tuple(ranges))
-  return UOp(Ops.STORE, src=(mm2, mm)+tuple(ranges))
-  #return x.replace(src=(mm,)+tuple(ranges))
-
-def map_reshape(x:UOp):
-  # don't push on the final buffer reshape for readable graph
-  #if x.src[0].src[0].op is Ops.BUFFER: return None
-  acc = 1
-  to_sum = []
-  for s,src in list(zip(x.shape, x.src[1:]))[::-1]:
-    to_sum.append(acc*src)
-    acc *= s
-  mish = sum(to_sum)
-  ret = []
-  for s in x.src[0].src[0].shape[::-1]:
-    if resolve(s!=1):
-      ret.append(mish % s)
-      mish //= s
-    else:
-      ret.append(UOp.const(dtypes.int, 0))
-  ret = UOp.sink(*ret).simplify().src[::-1] if len(ret) else ()
-  #return UOp(Ops.MAP, dtype=x.src[0].src[0].dtype, src=(x.src[0].src[0],)+ret)
-  # generic
-  return x.src[0].replace(src=tuple([UOp(Ops.MAP, dtype=s.dtype, src=(s,)+ret) for s in x.src[0].src]))
-
-def map_reduce(ctx:list[int], x:UOp):
-  rngs = list(x.src[1:])
-  r = x.src[0]
-  new_ranges = []
-  for i,s in enumerate(r.src[0].shape):
-    if i in r.arg[1]:
-      assert rngs[i].op == Ops.CONST
-      rngs[i] = UOp.range(dtypes.int, s, ctx[0])
-      new_ranges.append(rngs[i])
-      ctx[0] += 1
-  mm = UOp(Ops.MAP, r.src[0].dtype, src=(r.src[0],)+tuple(rngs))
-  return UOp(Ops.REDUCE_AXIS, r.dtype, src=(mm,)+tuple(new_ranges), arg=r.arg)
-
-def map_permute(x:UOp):
-  ret = x.src[1:]
-  # argsort or not?
-  perm = argsort(x.src[0].arg)
-  print(perm, x.src[0].arg)
-  ret = tuple([ret[p] for p in perm])
-  print(x.src[0].src[0].shape, ret)
-  #return UOp(Ops.MAP, dtype=x.src[0].src[0].dtype, src=(x.src[0].src[0],)+ret)
-  return x.src[0].replace(src=tuple([UOp(Ops.MAP, dtype=s.dtype, src=(s,)+ret) for s in x.src[0].src]))
-
-def map_expand(x:UOp):
-  r = x.src[0]
-  inp_shape, exp_shape = x.src[0].src[0].shape, x.src[0].shape
-  ret = list(x.src[1:])
-  exp_ranges = []
-  for i,(x,y) in enumerate(zip(inp_shape, exp_shape)):
-    if x != y:
-      exp_ranges.append(ret[i])
-      ret[i] = UOp.const(dtypes.int, 0)
-  mm = UOp(Ops.MAP, r.dtype, src=(r.src[0],)+tuple(ret))
-  return UOp(Ops.EXPAND, r.dtype, src=(mm,)+tuple(exp_ranges), arg=r.arg)
-
-def map_shrink(ctx:list[int], x:UOp):
-  r = x.src[0]
-  ret = list(x.src[1:])
-  for i,(s,(ss,se)) in enumerate(zip(r.src[0].shape, r.arg)):
-    assert ss == 0, "add to range?"
-    if se-ss != s and False:
-      new_ret_i = [ret[i]]
-      if ss != 0:
-        new_ret_i = [UOp.range(dtypes.int, ss, ctx[0])] + new_ret_i
-        ctx[0] += 1
-      if se != s:
-        new_ret_i = new_ret_i + [UOp.range(dtypes.int, s-se, ctx[0])]
-        ctx[0] += 1
-      ret[i] = UOp(Ops.CATRANGE, src=tuple(new_ret_i))
-  mm = UOp(Ops.MAP, r.dtype, src=(r.src[0],)+tuple(ret))
-  #return mm
-  # TODO: put the ranges on the shrink?
-  return UOp(Ops.SHRINK, r.dtype, src=(mm,), arg=r.arg)
-
-index_pushing = PatternMatcher([
-  (UPat(Ops.CONTIGUOUS, name="x"), contiguous_create_ranges),
-  (UPat(Ops.MAP, src=(UPat(Ops.RESHAPE),), allow_any_len=True, name="x"), map_reshape),
-  (UPat(Ops.MAP, src=(UPat(Ops.PERMUTE),), allow_any_len=True, name="x"), map_permute),
-  (UPat(Ops.MAP, src=(UPat(Ops.EXPAND),), allow_any_len=True, name="x"), map_expand),
-  (UPat(Ops.MAP, src=(UPat(Ops.SHRINK),), allow_any_len=True, name="x"), map_shrink),
-  (UPat(Ops.MAP, src=(UPat(Ops.REDUCE_AXIS),), allow_any_len=True, name="x"), map_reduce),
-  # move MAP through elementwise ALU
-  (UPat(Ops.MAP, src=(UPat(GroupOp.Elementwise),), allow_any_len=True, name="x"),
-   lambda x: x.src[0].replace(src=tuple([UOp(Ops.MAP, dtype=s.dtype, src=(s,)+x.src[1:]) for s in x.src[0].src]))),
-  # MAP on STORE is NOOP
-  (UPat(Ops.MAP, src=(UPat(Ops.STORE),), allow_any_len=True, name="x"), lambda x: x.src[0]),
-])
-
-fix_buffers = PatternMatcher([
-  (UPat(Ops.BUFFER, name="x"), lambda x: UOp(Ops.DEFINE_GLOBAL, dtype=x.dtype.ptr(x.arg), arg=x.src[0].arg)),
-  (UPat(Ops.MAP, name="x"), lambda x: x.replace(op=Ops.INDEX, dtype=x.src[0].dtype).load()),
-  (UPat(Ops.STORE, src=(UPat(Ops.LOAD),), name="x", allow_any_len=True), lambda x: x.replace(src=(x.src[0].src[0],)+x.src[1:])),
-  (UPat((Ops.RESHAPE, Ops.SHRINK, Ops.PERMUTE), name="x"), lambda x: x.src[0]),
-  # do EXPANDs need to track the ranges they end?
-  (UPat(Ops.EXPAND, name="x"), lambda x: x.src[0]),
-  #(UPat(Ops.EXPAND, name="x"), lambda x: x.replace(arg=None, op=Ops.NOOP)),
-  (UPat(Ops.REDUCE_AXIS, name="x"), lambda x: x.replace(op=Ops.REDUCE, arg=x.arg[0])),
-])
-
@track_rewrites(name=lambda sink,ret: f"Schedule {pluralize('Kernel',len([u for u in ret[sink].toposort() if u.op is Ops.KERNEL]))}")
 def get_kernelize_map(sink:UOp) -> dict[UOp, UOp]:
  """
@@ -543,20 +427,6 @@ def get_kernelize_map(sink:UOp) -> dict[UOp, UOp]:
  Returns:
    Map transforming each UOp in the sink to the Ops.KERNEL graph.
  """
-  pushed = graph_rewrite(sink, index_pushing, ctx=[0], name="index pushing", bottom_up=True)
-  pushed = graph_rewrite(pushed, fix_buffers, name="fix buffers")
-  from tinygrad.codegen.devectorizer import pm_reduce, ReduceContext
-  pushed = graph_rewrite(pushed, pm_reduce, ctx=ReduceContext(), name="remove reduce")
-  from tinygrad.codegen.linearize import block_create, BlockContext, pm_blockend_merge, block_merge, pm_finalize
-  pushed = graph_rewrite(pushed, block_create, ctx=BlockContext.from_sink(pushed), name="block create", bottom_up=True)
-  pushed = graph_rewrite(pushed, pm_blockend_merge, name="blockend merge")
-  pushed = graph_rewrite(pushed, block_merge, name="block merge")
-  pushed = graph_rewrite(pushed, pm_finalize, name="finalize")
-  from tinygrad.device import Device
-  try:
-    print(Device['CPU'].renderer.render(pushed.arg.lst))
-  except Exception as e:
-    print("render fail", e)

  # multi + merge_views + simplify
  tensor_map = graph_rewrite_map(sink, multi_pm+do_fuse+merge_views+sym+replace_contiguous, ctx={}, name="merge_views")
@@ -4,7 +4,7 @@ from dataclasses import dataclass
 import functools
 from typing import Callable
 from tinygrad.helpers import merge_dicts, getenv
-from tinygrad.shape.view import View, unravel
+from tinygrad.shape.view import View, strides_for_shape, unravel
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, Variable, sint, sint_to_uop, Context, PatternMatcher, UPat, GroupOp
 from tinygrad.uop.symbolic import split_uop, symbolic_flat, uop_given_valid, simplify_valid
@@ -75,6 +75,9 @@ class ShapeTracker:
  @property
  def contiguous(self) -> bool: return len(self.views) == 1 and self.views[0].contiguous

+  @property
+  def consecutive(self) -> bool: return len(self.views) == 1 and (v:=self.views[0]).mask is None and v.strides == strides_for_shape(v.shape)
+
  @property
  def shape(self) -> tuple[sint, ...]: return self.views[-1].shape

@@ -83,6 +86,7 @@ class ShapeTracker:

  def reduce(self, axis:tuple[int, ...]) -> tuple[sint, ...]: return tuple(1 if i in axis else s for i,s in enumerate(self.shape))

+  def to_uop(self) -> UOp: return UOp(Ops.VIEW, dtypes.void, (), self)
  def to_indexed_uops(self, _idxs:list[UOp]|tuple[UOp, ...]|None=None) -> tuple[UOp, UOp]:
    return views_to_indexed_uops(self.views, tuple(_idxs) if _idxs is not None else None)

@@ -282,7 +282,7 @@ class Tensor(MathTrait):
    # TODO: this is a hack for writing to DISK. remove with working assign
    if isinstance(self.device, str) and self.device.startswith("DISK"):
      if x.__class__ is not Tensor: x = Tensor(x, device="CPU", dtype=self.dtype)
-      self._buffer().copyin(x._data())
+      cast(Buffer, self.contiguous().realize().uop.base.buffer).ensure_allocated().copyin(x._data())
      return self
    if x.__class__ is not Tensor: x = Tensor(x, device=self.device, dtype=self.dtype)
    if self.uop is x.uop: return self  # a self assign is a NOOP
@@ -299,10 +299,7 @@ class Tensor(MathTrait):
    """
    return Tensor(self.uop.detach(), device=self.device, requires_grad=False)

-  def _buffer(self) -> Buffer:
-    x = self.cast(self.dtype.base).contiguous()
-    if isinstance(self.device, tuple): x = x.to("CPU")
-    return cast(Buffer, x.realize().uop.base.buffer).ensure_allocated()
+  def _buffer(self) -> Buffer: return cast(Buffer, self.cast(self.dtype.base).contiguous().to("CPU").realize().uop.base.buffer)
  def _data(self) -> memoryview: return self._buffer().as_buffer()

  def data(self) -> memoryview:
@@ -1979,14 +1976,12 @@ class Tensor(MathTrait):

    # https://keccak.team/keccak_specs_summary.html

-    def ctensor(l: Sequence[ConstType], dtype: DType = dtypes.uint64):
-      # TODO: contiguous is here for compile speed
-      return Tensor.stack(*(Tensor(v, dtype=dtype, device=self.device) for v in l)).contiguous()
+    def ctensor(l: Sequence[ConstType], dtype: DType = dtypes.uint64): return Tensor.stack(*(Tensor(v, dtype=dtype, device=self.device) for v in l))
    rot_offsets = [44, 43, 21, 14, 28, 20, 3, 45, 61, 1, 6, 25, 8, 18, 27, 36, 10, 15, 56, 62, 55, 39, 41, 2]
    rot_offsets_v0, rot_offsets_v1 =  ctensor([0] + [1 << v for v in rot_offsets]), ctensor([1] + [1 << (64 - v) for v in rot_offsets])

    # calculated from π step
-    reorder_indexes = ctensor([0,6,12,18,24,3,9,10,16,22,1,7,13,19,20,4,5,11,17,23,2,8,14,15,21], dtype=dtypes.int32)
+    reorder_indexes = ctensor([0,6,12,18,24,3,9,10,16,22,1,7,13,19,20,4,5,11,17,23,2,8,14,15,21])
    rnd_const_masks = [ctensor([v]).pad((0, 24)) for v in (1, 0x8082, 0x800000000000808a, 0x8000000080008000, 0x808b, 0x80000001, 0x8000000080008081,
    0x8000000000008009, 0x8a, 0x88, 0x80008009, 0x8000000a, 0x8000808b, 0x800000000000008b, 0x8000000000008089, 0x8000000000008003,
    0x8000000000008002, 0x8000000000000080, 0x800a, 0x800000008000000a, 0x8000000080008081, 0x8000000000008080, 0x80000001, 0x8000000080008008)]
@@ -1997,9 +1992,9 @@ class Tensor(MathTrait):
    data = data.pad((None, (0, data_pad))).reshape(bs := data.shape[0], -1, rate).pad((None, None, (0, 200 - rate)))

    # create pad mask
-    lbe = prod(data.shape[1:]) + rate - data_pad - 200
-    if data_pad == 1: mb = [(lbe, 0), (1, dsbyte ^ 0x80), (200 - rate, 0)]
-    else: mb = [(lbe, 0), (1, dsbyte), (data_pad - 2, 0), (1, 0x80), (200 - rate, 0)]
+    lbe = (blen := prod(data.shape[1:])) + rate - data_pad - 200
+    if data_pad == 1: mb = [(lbe, 0), (1, dsbyte ^ 0x80), (blen - lbe - 1, 0)]
+    else: mb = [(lbe, 0), (1, dsbyte), (blen + rate - lbe - 202, 0), (1, 0x80), (200 - rate, 0)]
    pad_mask = Tensor.cat(*(Tensor(v, dtype=dtypes.uint8, device=data.device).expand(l) for l, v in mb if l > 0)).unsqueeze(0)

    data = (data.flatten(1) ^ pad_mask).reshape(*data.shape[:2], 200).bitcast(dtypes.uint64)
@@ -2018,42 +2013,8 @@ class Tensor(MathTrait):
        # χ and ι step
        state = state.bitwise_xor(~state.roll(shifts=-1, dims=2) & state.roll(shifts=-2, dims=2))
        state = state.flatten(1) ^ rnd_const_masks[i]
-      # NOTE: kernelize here to prevent internal stack from growing propotional to data size
-      state = state.kernelize()
    return state.bitcast(dtypes.uint8)[:,:(obytes:=(200 - rate) // 2)].reshape(*self.shape[:-1], obytes)

-  def _hash_1mb(self) -> Tensor:
-    assert self.dtype == dtypes.uint8, "only support uint8 tensors for hashing"
-    assert self.ndim == 2, "only support batched 1d tensors"
-    assert self.shape[1] == 1024 * 1024, "only support messages of 1mb"
-
-    blocks = self.shape[0] * self.shape[1] // 4096
-    data = self.reshape(blocks, 4096)
-    block_hashes = data.keccak("shake_128").reshape(self.shape[0], 4096)
-    return block_hashes.keccak("shake_128").reshape(self.shape[0], 16)
-
-  def hash(self) -> Tensor:
-    """
-    Calculates a 16-byte hash of the tensor.
-    ```python exec="false source="above" session="tensor" result="python"
-    t = Tensor(b"Hello World!").hash()
-    print(t.data().hex())
-    ```
-    """
-
-    data = self.flatten().bitcast(dtypes.uint8)
-    if (tsize := data.shape[0]) % 2**20 != 0: data = data.pad((0, 2**20 - tsize % 2**20))
-    base_chunks = ceildiv(data.shape[0], 2**20)
-    tree_depth = math.ceil(math.log(base_chunks, 65536)) if base_chunks > 1 else 0
-
-    level_chunks = base_chunks
-    for _ in range(tree_depth + 1):
-      data = data.reshape(level_chunks, 2**20)._hash_1mb().flatten()
-      if (tsize := data.shape[0]) % 2**20 != 0: data = data.pad((0, 2**20 - tsize % 2**20))
-      level_chunks = ceildiv(data.shape[0], 2**20)
-
-    return data[:16]
-
  def _softmax(self, axis, dtype:DTypeLike|None=None) -> tuple[Tensor, Tensor, Tensor]:
    m = self - self.max(axis=axis, keepdim=True).detach()
    if dtype is not None: m = m.cast(dtype)
@@ -9,8 +9,7 @@ class FastEnum(IntEnum):
 # the order of these Ops controls the order of the toposort
 class Ops(FastEnum):
  # uops that aren't rendered
-  NOOP = auto(); SINK = auto(); UNIQUE = auto(); DEVICE = auto(); KERNEL = auto(); PRECAST = auto()  # noqa: E702
-  MAP = auto(); CATRANGE = auto()
+  NOOP = auto(); SINK = auto(); UNIQUE = auto(); DEVICE = auto(); KERNEL = auto()  # noqa: E702

  # buffer ops
  COPY = auto(); BUFFER = auto(); BUFFER_VIEW = auto(); MSELECT = auto(); MSTACK = auto() # noqa: E702
@@ -83,9 +82,6 @@ class GroupOp:
            Ops.XOR, Ops.SHL, Ops.SHR, Ops.OR, Ops.AND, Ops.THREEFRY, Ops.SUB, Ops.FDIV, Ops.POW}
  Ternary = {Ops.WHERE, Ops.MULACC}
  ALU = set.union(Unary, Binary, Ternary)
-  Elementwise = set.union(ALU, {Ops.CAST, Ops.BITCAST})
-
-  Defines = {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}

  Irreducible = {Ops.CONST, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.RANGE}
  Movement = {Ops.RESHAPE, Ops.EXPAND, Ops.PERMUTE, Ops.PAD, Ops.SHRINK, Ops.FLIP}
@@ -4,7 +4,7 @@ from tinygrad.dtype import dtypes

 class MathTrait:
  # required to implement
-  def alu(self:T, op:Ops, *src) -> T: raise NotImplementedError
+  def alu(self:T, arg:Ops, *src) -> T: raise NotImplementedError
  def const_like(self:T, b) -> T: raise NotImplementedError

  # great functions you get!
@@ -136,7 +136,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):

  @functools.cached_property
  def st(self) -> ShapeTracker|None:
-    if self.op in GroupOp.Block or self.op is Ops.INDEX: return None
    from tinygrad.shape.shapetracker import ShapeTracker
    # VIEW and MovementOps define a new ShapeTracker from the arg
    if self.op is Ops.VIEW: return self.arg
@@ -144,13 +143,12 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    # CONST with a DEVICE has a shape of ()
    if self.op is Ops.CONST and len(self.src) and self.src[0].op is Ops.DEVICE: return ShapeTracker.from_shape(())
    # BufferOps and ASSIGN flow ShapeTracker from a direct edge
-    if self.op in {Ops.STORE, Ops.ASSIGN, Ops.LOAD}: return self.src[0].st
    if self.op in GroupOp.Buffer: return views[0] if (views:=[x.st for x in self.src if x.op is Ops.VIEW]) else None
+    if self.op is Ops.ASSIGN: return self.src[0].st

    # BUFFER/BUFFER_VIEW and KERNEL only have a size
    if self.op in {Ops.BUFFER, Ops.BUFFER_VIEW}: return ShapeTracker.from_shape((self.size,))
    if self.op is Ops.KERNEL: return ShapeTracker.from_shape((self.arg.ast.size,))
-    #if self.op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}: return ShapeTracker.from_shape((self.dtype.size,))

    # otherwise we get the shape from sources
    if not (src_sts := [x.st for x in self.src if x.st is not None]): return None
@@ -169,7 +167,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    if self.op is Ops.VIEW: return self.shape
    # NOTE: if a parent doesn't have st its full_shape is empty
    parent_shapes = [x.full_shape for x in self.src]
-    return tuple(smax(x) for x in itertools.zip_longest(*parent_shapes, fillvalue=1))
+    return tuple(smax(x) for x in zip(*[x for x in parent_shapes if x != ()]))
  @property
  def shape(self) -> tuple[sint, ...]: return unwrap(self.st).shape
  @property
@@ -213,7 +211,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
  def sink(self, *srcs:UOp|None, **kwargs): return UOp(Ops.SINK, dtypes.void, (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
  def detach(self): return UOp(Ops.DETACH, self.dtype, (self,))
  def index(self, idx:UOp, valid:UOp|None=None): return UOp(Ops.INDEX, self.dtype, (self,idx,valid) if valid is not None else (self,idx))
-  def __getitem__(self, idx): return self.index(idx)
  def const_like(self, b:ConstLike):
    # constants can optionally have a DEVICE source
    return UOp.const(self.dtype, b, device=self._device, shape=self.shape if self.st is not None else None)
@@ -236,13 +233,12 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
      i = (i,)
    return UOp(Ops.GEP, self.dtype.scalar().vec(len(i)) if len(i) > 1 else self.dtype.scalar(), (self,), i)
  def load(self, *src:UOp, **kwargs): return UOp(Ops.LOAD, dtype=kwargs.pop("dtype", self.dtype.base), src=(self,)+src, **kwargs)
-  def store(self, *src:UOp, **kwargs): return UOp(Ops.STORE, dtypes.void, (self,)+src, **kwargs)
+  def store(self, *src:UOp, **kwargs): return UOp(Ops.STORE, self.dtype, (self,)+src, **kwargs)
  def assign(self, x:UOp): return UOp(Ops.ASSIGN, self.dtype, (self, x))
-  def barrier(self, *src:UOp): return UOp(Ops.BARRIER, src=(self,)+src)
-  def alu(self, op, *src:UOp, **kwargs):
+  def alu(self, arg, *src:UOp):
    out_dtype = (self, *src)[-1].dtype
-    if op in {Ops.CMPLT, Ops.CMPNE}: out_dtype = dtypes.bool.vec(out_dtype.count) if out_dtype.count > 1 else dtypes.bool
-    return UOp(op, out_dtype, (self,)+src, **kwargs)
+    if arg in {Ops.CMPLT, Ops.CMPNE}: out_dtype = dtypes.bool.vec(out_dtype.count) if out_dtype.count > 1 else dtypes.bool
+    return UOp(arg, out_dtype, (self,)+src)
  @staticmethod
  def const(dtype:DType, b:ConstLike, device:str|tuple[str, ...]|None=None, shape:tuple[sint, ...]|None=None):
    if isinstance(b, UOp): return b.unbind()[0] if b.op is Ops.BIND else b
@@ -250,21 +246,25 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    ret = UOp(Ops.VCONST if isinstance(b, tuple) else Ops.CONST, dtype, arg=dtypes.as_const(b, dtype))
    if shape is not None:
      from tinygrad.shape.shapetracker import ShapeTracker
-      ret = ret.replace(src=(UOp(Ops.VIEW, dtypes.void, (), ShapeTracker.from_shape(shape, (0,)*len(shape))),))
+      ret = ret.replace(src=(ShapeTracker.from_shape(()).reshape((1,)*len(shape)).expand(shape).to_uop(),))
    if device is not None:
      ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device).view(unwrap(ret.st)),))
    return ret
+  def valid(self): return UOp.where(UOp(Ops.VALID, dtypes.bool, (UOp(Ops.VIEW, arg=self.st),)), self.const_like(self.base.arg), 0)
  @staticmethod
  def range(dtype:DType, end:sint, idx:int): return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end),), arg=idx)
-  def r(self, op:Ops, axis:tuple[int, ...]):
+  def r(self, op:Ops, axis:tuple[int, ...], permute=True):
    axis = tuple(sorted([x for x in axis if resolve(self.shape[x] != 1)]))
    if len(axis) == 0: return self
    # move any non reduce axis before the first reduce axis
    move_early, rest = partition(range(axis[0], len(self.shape)), lambda i: i not in axis and resolve(self.shape[i] != 1))
-    permaxis = tuple(range(axis[0])) + tuple(move_early) + tuple(rest)
-    ret = self.permute(permaxis)
-    new_axis = tuple([x for x in range(axis[0]+len(move_early), len(self.shape)) if resolve(ret.shape[x] != 1)])
-    assert len(axis) == len(new_axis)
+    if move_early and permute:
+      permaxis = tuple(range(axis[0])) + tuple(move_early) + tuple(rest)
+      ret = self.permute(permaxis)
+      new_axis = tuple([x for x in range(axis[0]+len(move_early), len(self.shape)) if resolve(ret.shape[x] != 1)])
+      assert len(axis) == len(new_axis)
+    else:
+      ret, new_axis = self, axis
    ret = UOp(Ops.REDUCE_AXIS, self.dtype, (ret,), (op, new_axis))
    return ret.reshape(tuple([x if i not in axis else 1 for i,x in enumerate(self.shape)]))
  def reduce(self, *src:UOp, **kwargs): return UOp(Ops.REDUCE, kwargs.pop('dtype', self.dtype), src=(self,)+src, **kwargs)
@@ -336,7 +336,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    return self
  def view(self, new_st:ShapeTracker) -> UOp: return UOp(Ops.VIEW, self.dtype, (self,), new_st)

-  def _mop(self, op:Ops, arg) -> UOp:
+  def _mop(self, op:Ops, arg):
    ret = UOp(op, self.dtype, (self,), arg)
    if self.st == ret.st: return self  # ignore NOOPs, also check ret.st
    return ret
@@ -369,7 +369,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
      return self.src[0].device[self.arg]
    if self.op is Ops.MSTACK: return tuple(cast(str, x.device) for x in self.src)
    if self.op in {Ops.COPY, Ops.BUFFER, Ops.ALLREDUCE}: return self.src[1].device
-    return next((x._device for x in self.src if x._device is not None), None)
+    return dsrcs[0]._device if len(dsrcs:=[x for x in self.src if x._device is not None]) != 0 else None
  @property
  def buf_uop(self) -> UOp:
    if self.op is Ops.BUFFER: return self
@@ -537,6 +537,10 @@ class KernelInfo:
  opts_to_apply: tuple|None = None
  @property
  def function_name(self): return to_function_name(self.name)
+  @property
+  def global_dims(self) -> list[int]: return [i for i,x in enumerate(self.axis_types) if x is AxisType.GLOBAL]
+  @property
+  def local_dims(self) -> list[int]: return [i for i,x in enumerate(self.axis_types) if x in (AxisType.LOCAL, AxisType.GROUP_REDUCE)]

 # ******** ops in python ********

@@ -851,8 +855,8 @@ if TRACK_MATCH_STATS or PROFILE:
      with open(fn:=temp("rewrites.pkl", append_user=True), "wb") as f:
        print(f"rewrote {len(tracked_ctxs)} graphs and matched {sum(len(r.matches) for x in tracked_ctxs for r in x)} times, saved to {fn}")
        pickle.dump((tracked_keys, tracked_ctxs, uop_fields), f)
-    if VIZ: launch_viz(VIZ, temp("rewrites.pkl", append_user=True))
-    if getenv("PRINT_MATCH_STATS", TRACK_MATCH_STATS.value):
+    if VIZ: launch_viz("VIZ", temp("rewrites.pkl", append_user=True))
+    if getenv("PRINT_MATCH_STATS", 1):
      ret = [0,0,0.0,0.0]
      for k,v in sorted(list(match_stats.items()), key=lambda x: x[1][2]+x[1][3]):
        loc_str = f"{k.location[0].split('/')[-1]}:{k.location[1]}"
@@ -861,10 +865,9 @@ if TRACK_MATCH_STATS or PROFILE:
      print(f"{ret[0]:6d} / {ret[1]:7d} -- {ret[3]*1000.:9.2f} / {(ret[2]+ret[3])*1000.:9.2f} ms -- TOTAL")
      print(f"{len(match_stats)} rules, {sum(v[0] > 0 for v in match_stats.values())} matched once")

-  def launch_viz(var:ContextVar, data:str):
-    os.environ[(env_str:=var.key)] = "0"
+  def launch_viz(env_str:str, data:str):
+    os.environ[env_str] = "0"
    os.environ[f"{env_str}_DATA"] = data
-    os.environ[f"{env_str}_VALUE"] = str(var.value)
    if not int(os.getenv("VIZ", "0")) and not int(os.getenv("PROFILE", "0")):
      args = ['--kernels', getenv("VIZ_DATA", "")] if getenv("VIZ_DATA", "") else []
      args += ['--profile', getenv("PROFILE_DATA", "")] if getenv("PROFILE_DATA", "") else []
@@ -2,7 +2,6 @@ from typing import cast, Callable
 from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, python_alu, graph_rewrite, resolve
 from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace
 from tinygrad.helpers import all_same, prod, DEBUG, ContextVar, Context
-from tinygrad.shape.shapetracker import ShapeTracker
 try:
  import z3

@@ -131,43 +130,46 @@ index_pat = UPat(Ops.INDEX, name="idx").or_casted()
 spec = PatternMatcher([
  (UPat(Ops.DEFINE_GLOBAL, name="x"), lambda x: isinstance(x.dtype, (PtrDType, ImageDType)) and x.dtype.addrspace == AddrSpace.GLOBAL),
  (UPat(Ops.DEFINE_LOCAL, name="x"), lambda x: isinstance(x.dtype, PtrDType) and x.dtype.addrspace == AddrSpace.LOCAL),
-  (UPat(Ops.DEFINE_REG, src=()), lambda: True),
+  (UPat(Ops.DEFINE_REG, src=(UPat.var("c"),), name="x", allow_any_len=True),
+   lambda x,c: all(y.op is Ops.RANGE for y in x.src[1:]) and c.dtype.base == x.dtype.base),
  (UPat(Ops.DEFINE_VAR, name="x"), lambda x: isinstance(x.arg[1], int) and isinstance(x.arg[2], int)),

  (UPat(Ops.RANGE, src=(UPat.var("x"),), name="rng"), lambda rng,x: rng.dtype == x.dtype and isinstance(rng.arg, int)),
  (UPat(Ops.SPECIAL, src=()), lambda: True),

-  (UPat(Ops.VIEW, dtypes.void, src=(), name="x"), lambda x: isinstance(x.arg, ShapeTracker)),
-  (UPat(Ops.VIEW, src=(UPat.var("src"),), name="x"),
-   lambda x,src: isinstance(x.arg, ShapeTracker) and src.op is not Ops.STORE and x.dtype.base == src.dtype.base),
+  # TODO: confirm the args of both of these are shapetrackers
+  (UPat(Ops.VIEW, dtypes.void, src=()), lambda: True),
+  (UPat(Ops.VIEW, src=(UPat.var("src"),), name="x"), lambda x,src: src.op is not Ops.STORE and x.dtype.base == src.dtype.base),

  (UPat(Ops.VALID, dtypes.bool, (UPat(Ops.VIEW),)), lambda: True),
  (UPat(Ops.CONST, name="x"), lambda x: type(x.arg) is type(dtypes.as_const(x.arg, x.dtype))),

  # early LOAD has a <bufview, store?>
-  (UPat(Ops.LOAD, src=(UPat(Ops.VIEW, src=(UPat(GroupOp.Defines),)),)), lambda: True),
-  (UPat(Ops.LOAD, src=(UPat(Ops.VIEW, src=(UPat(GroupOp.Defines),)), UPat(Ops.STORE))), lambda: True),
+  (UPat(Ops.LOAD, src=(UPat(Ops.VIEW, src=(UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL)),)),)), lambda: True),
+  (UPat(Ops.LOAD, src=(UPat(Ops.VIEW, src=(UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL)),)), UPat(Ops.STORE))), lambda: True),

  # early STORE has a <bufview, val>
-  (UPat(Ops.STORE, src=(UPat(Ops.VIEW, src=(UPat(GroupOp.Defines),)), UPat())), lambda: True),
+  (UPat(Ops.STORE, src=(UPat(Ops.VIEW, src=(UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL)),)), UPat())), lambda: True),

  # **** new style load/store ****

  # INDEX is used in new style load/store
  # INDEX takes a <buf, alu, gate?>
-  (UPat(Ops.INDEX, src=(UPat(GroupOp.Defines), UPat())), lambda: True),
-  (UPat(Ops.INDEX, src=(UPat(GroupOp.Defines), UPat(), UPat(dtype=dtypes.bool))), lambda: True),
+  (UPat(Ops.INDEX, src=(UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG)), UPat())), lambda: True),
+  (UPat(Ops.INDEX, src=(UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG)), UPat(), UPat(dtype=dtypes.bool))), lambda: True),

  # LOAD on STORE
-  (UPat(Ops.LOAD, src=(UPat(Ops.STORE),), allow_any_len=True), lambda: True),
+  (UPat(Ops.LOAD, src=(UPat(Ops.STORE),)), lambda: True),

  # LOAD takes a <bufidx, alt?, barrier?>
-  (UPat(Ops.LOAD, src=(index_pat, UPat(Ops.IF, name="cond")), allow_any_len=True), lambda idx,cond: validate_index(idx,cond.src[0])),
-  (UPat(Ops.LOAD, src=(index_pat,), allow_any_len=True), validate_index),
+  (UPat(Ops.LOAD, src=(index_pat,)), validate_index),
+  (UPat(Ops.LOAD, src=(index_pat, UPat(Ops.BARRIER))), validate_index),
+  (UPat(Ops.LOAD, src=(index_pat, UPat(Ops.IF, name="cond"))), lambda idx,cond: validate_index(idx,cond.src[0])),
+  (UPat(Ops.LOAD, src=(index_pat, UPat.var("alt")), name="ld"), lambda ld,alt,idx: ld.dtype == alt.dtype and validate_index(idx)),

  # STORE takes a <bufidx, val, gate?>
-  (UPat(Ops.STORE, dtype=dtypes.void, src=(index_pat, UPat(name="val"), UPat(Ops.IF, name="gate")), allow_any_len=True), validate_store),
-  (UPat(Ops.STORE, dtype=dtypes.void, src=(index_pat, UPat(name="val")), allow_any_len=True), validate_store),
+  (UPat(Ops.STORE, src=(index_pat, UPat(name="val"), UPat(Ops.IF, name="gate")), allow_any_len=True), validate_store),
+  (UPat(Ops.STORE, src=(index_pat, UPat(name="val")), allow_any_len=True), validate_store),

  # most ALUs have all matching dtypes, except CMPLT, CMPNE, and WHERE
  (UPat(Ops.WHERE, name="w", src=(UPat(dtype=dtypes.bool), UPat.var("x"), UPat.var("y"))), lambda w,x,y: w.dtype == x.dtype == y.dtype),
@@ -199,7 +201,7 @@ spec = PatternMatcher([
  # NOTE: for testing, we let sinks be anything
  #(UPat(Ops.SINK, src=UPat(Ops.STORE)), lambda: True),
  (UPat(Ops.SINK, dtypes.void), lambda: True),
-  (UPat((Ops.NOOP, Ops.CUSTOMI, Ops.CUSTOM, Ops.PRECAST)), lambda: True),
+  (UPat((Ops.NOOP, Ops.CUSTOMI, Ops.CUSTOM)), lambda: True),

  # PTX LOAD/STORE
  (UPat((Ops.LOAD, Ops.STORE), src=(UPat(dtype=dtypes.int64),), allow_any_len=True), lambda: True),
@@ -209,7 +211,7 @@ spec = PatternMatcher([

 def verify_sink_dims(sink:UOp):
  if not all_same([s.shape for s in sink.src]): return False
-  for dims in zip(*[x.shape for x in sink.toposort() if x.op is Ops.VIEW]):
+  for dims in zip(*[x.shape for x in sink.toposort() if x.st is not None]):
    if len(n_dims:={s for s in dims if resolve(s!=1)}) > 1:
      print(f"# INVALID KERNEL DIMS: can only have 1 or n in each dimension: {n_dims}")
      return False
@@ -405,8 +405,8 @@ def reduce_mul_chain(r:UOp):
  return r.replace(src=(prod(inside) if len(inside) else r.src[0].const_like(1),)+r.src[1:])*prod(outside)

 # this is symbolic 2.0
-REMOVE_FROM_SINK = {Ops.SINK, Ops.UNROLL, Ops.PTRCAT, Ops.CAT, Ops.NOOP}
-REMOVE_FROM_BARRIER = {Ops.VECTORIZE, Ops.SINK, Ops.CAT, Ops.PTRCAT, Ops.NOOP}
+REMOVE_FROM_SINK = {Ops.SINK, Ops.UNROLL, Ops.PTRCAT, Ops.CAT}
+REMOVE_FROM_BARRIER = {Ops.VECTORIZE, Ops.SINK, Ops.CAT, Ops.PTRCAT}
 sym = symbolic_flat+PatternMatcher([
  # LOAD/STORE -> NOOP
  (UPat.var('x').store(UPat.var('x').load(), allow_any_len=True), lambda x: None if x.dtype.addrspace != AddrSpace.REG else x.src[0].src[0]),
@@ -437,6 +437,7 @@ sym = symbolic_flat+PatternMatcher([
  ((UPat.var('x', dtypes.uint64)&(UPat.var('y').where(UPat.const(dtypes.uint64, 0xFFFFFFFF), UPat.const(dtypes.uint64, 0)))).cast(dtypes.uint32),
   lambda x,y: y.where(x.cast(dtypes.uint32), UOp.const(dtypes.uint32, 0))),
  # ** self folding **
+  (UPat(Ops.DEFINE_REG, src=(UPat.var("x"),)), lambda x: x),            # a DEFINE_ACC without ranges is a CONST
  # x!=0 -> (bool)x
  (UPat.var("x")!=0, lambda x: x.cast(dtypes.bool.vec(x.dtype.count))),
  # ** where **
@@ -450,13 +451,15 @@ sym = symbolic_flat+PatternMatcher([
  (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx"), UPat.const(dtypes.bool, True))), lambda b, idx: b.index(idx)),
  # ** load/store folding **
  (UPat.store(UPat(Ops.INDEX, name="index"), UPat.load(UPat(Ops.INDEX, name="index"))), lambda index: UOp(Ops.NOOP)),
-  (UPat.store(UPat(Ops.INDEX, name="index"), UPat.var("gate").where(UPat.var("alt"),
-                                                                    UPat.load(UPat(Ops.INDEX, name="index"))), allow_any_len=True, name="store"),
-   lambda index, gate, alt, store: UOp.store(index.src[0].index(index.src[1], gate), alt, *store.src[2:])),
+  (UPat.store(UPat(Ops.INDEX, name="index"), UPat.var("gate").where(UPat.var("alt"), UPat.load(UPat(Ops.INDEX, name="index")))),
+   lambda index, gate, alt: UOp.store(index.src[0].index(index.src[1], gate), alt)),
  # fold gated LOAD/STORE
  (UPat().index(UPat(), UPat.const(dtypes.bool, True)).named("idx"), lambda idx: idx.replace(src=idx.src[0:2])), # remove True
  (UPat((Ops.LOAD, Ops.STORE), src=(UPat().index(UPat(), UPat.const(dtypes.bool, False)).or_casted(),), allow_any_len=True, name="x"),
    lambda x: UOp(Ops.NOOP) if x.op is Ops.STORE else x.const_like(0)), # NULL pointer store does nothing. NULL pointer load produces 0
+  # remove NOOPs from SINK
+  (UPat(Ops.SINK, name="root"),
+    lambda root: UOp(Ops.SINK, root.dtype, a, root.arg) if len(a:=tuple(x for x in root.src if x.op is not Ops.NOOP)) != len(root.src) else None),
  # remove VECTORIZE from SINK/BARRIER. TODO: SINK/BARRIER are really the same thing at GLOBAL/LOCAL levels
  (UPat(Ops.BARRIER, name="root"),
    lambda root: UOp(Ops.BARRIER, root.dtype, tuple(flatten(x.src if x.op in REMOVE_FROM_BARRIER else (x,) for x in root.src)), root.arg)
@@ -10,6 +10,5 @@ fetch "dagrejs.github.io/project/dagre/latest/dagre.min.js"
 fetch "cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/styles/default.min.css"
 fetch "cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/highlight.min.js"
 fetch "cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/python.min.js"
-fetch "cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/x86asm.min.js"
 fetch "cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/cpp.min.js"
 fetch "unpkg.com/@highlightjs/[email protected]/styles/tokyo-night-dark.min.css"
@@ -10,7 +10,6 @@
  <script src="assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/highlight.min.js"></script>
  <script src="assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/python.min.js"></script>
  <script src="assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/cpp.min.js"></script>
-  <script src="assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/x86asm.min.js"></script>
  <link rel="stylesheet" href="assets/unpkg.com/@highlightjs/[email protected]/styles/tokyo-night-dark.min.css" />
  <style>
  * {
@@ -104,17 +103,6 @@
  .metadata > * + *, .rewrite-container > * + *, .ctx-list > * + * {
    margin-top: 12px;
  }
-  .stats-list > * + * {
-    margin-top: 8px;
-  }
-  .stats-list > p > * + * {
-    margin-top: 12px;
-  }
-  .stats-list {
-    width: 100%;
-    max-height: 240px;
-    overflow: auto;
-  }
  .ctx-list > ul > * + * {
    margin-top: 4px;
  }
@@ -223,87 +211,19 @@
    pointer-events: none;
    display: none;
    font-size: 10px;
-    white-space: pre;
  }
  #device-list > div {
    min-height: 32px;
    max-width: 100px;
-    overflow-x: auto;
-    overflow-y: hidden;
+    overflow: hidden;
+    text-overflow: ellipsis;
    white-space: nowrap;
    display: flex;
+    cursor: pointer;
  }
  #device-list > div:hover {
    background-color: rgba(20, 23, 35, 0.3);
  }
-  .raw-text {
-    padding: 0 8px;
-    width: 100%;
-    height: 100%;
-    max-height: 100vh;
-    overflow-x: auto;
-  }
-  .raw-text code {
-    max-height: none !important;
-  }
-  table {
-    width: 100%;
-    border-collapse: separate;
-    border-spacing: 0;
-    background-color: #1a1b26;
-    color: #f0f0f5;
-    font-size: 0.95em;
-  }
-  table td {
-    border-bottom: 1px solid #2c2f40;
-    vertical-align: top;
-  }
-  table tr:last-child > td {
-     border-bottom: none;
-  }
-  tr.main-row:hover {
-    background-color: #2a2d3a;
-  }
-  tr.sub-row {
-    max-width: 150px;
-  }
-  tr.main-row > td, tr.sub-row > td {
-    padding: 8px 12px;
-  }
-  tr.code-row > td:first-child {
-    font-family: monospace;
-  }
-  td.pct-row > div {
-    height: 12px;
-    width: 100%;
-    display: flex;
-  }
-  td.pct-row > div > div {
-    height: 100%;
-  }
-  thead {
-    position: sticky;
-    top: 0;
-    z-index: 10;
-    background-color: #20222e;
-  }
-  thead th {
-    text-align: left;
-    padding: 10px 12px;
-    font-weight: 600;
-    border-bottom: 1px solid #3a3d52;
-    font-size: 0.95em;
-    letter-spacing: 0.03em;
-  }
-  .legend {
-    display: flex;
-    align-items: center;
-  }
-  .legend > div {
-    width: 0.95em;
-    height: 0.95em;
-    margin-right: 4px;
-  }
  </style>
 </head>
 <body>
@@ -318,10 +238,10 @@
        </svg>
      </button>
    </div>
-    <div class="progress-message"></div>
    <div class="container ctx-list-parent"><div class="ctx-list"></div></div>
    <div class="view profiler"></div>
    <div class="view graph">
+      <div class="progress-message">Rendering new layout...</div>
      <svg id="graph-svg" preserveAspectRatio="xMidYMid meet">
        <g id="render">
          <g id="edges"></g>
@@ -25,10 +25,6 @@ function intersectRect(r1, r2) {
 let [workerUrl, worker, timeout] = [null, null, null];
 async function renderDag(graph, additions, recenter=false) {
  // start calculating the new layout (non-blocking)
-  const progressMessage = document.querySelector(".progress-message");
-  progressMessage.innerText = "Rendering new graph...";
-  if (timeout != null) clearTimeout(timeout);
-  timeout = setTimeout(() => {progressMessage.style.display = "block"}, 2000);
  if (worker == null) {
    const resp = await Promise.all(["/assets/dagrejs.github.io/project/dagre/latest/dagre.min.js","/js/worker.js"].map(u => fetch(u)));
    workerUrl = URL.createObjectURL(new Blob([(await Promise.all(resp.map((r) => r.text()))).join("\n")], { type: "application/javascript" }));
@@ -37,6 +33,9 @@ async function renderDag(graph, additions, recenter=false) {
    worker.terminate();
    worker = new Worker(workerUrl);
  }
+  if (timeout != null) clearTimeout(timeout);
+  const progressMessage = document.querySelector(".progress-message");
+  timeout = setTimeout(() => {progressMessage.style.display = "block"}, 2000);
  worker.postMessage({graph, additions, ctxs});
  worker.onmessage = (e) => {
    displayGraph("graph");
@@ -109,11 +108,11 @@ function formatTime(ts, dur=ts) {
 }
 const formatUnit = (d, unit="") => d3.format(".3~s")(d)+unit;

-const colorScheme = {TINY:["#1b5745", "#354f52", "#354f52", "#46acc2", "#1d2e62"],
-  DEFAULT:["#2b2e39", "#2c2f3a", "#31343f", "#323544", "#2d303a", "#2e313c", "#343746", "#353847", "#3c4050", "#404459", "#444862", "#4a4e65"],
-  BUFFER:["#3A57B7","#5066C1","#6277CD","#7488D8","#8A9BE3","#A3B4F2"],
-  CATEGORICAL:["#ff8080", "#F4A261", "#C8F9D4", "#8D99AE", "#F4A261", "#ffffa2", "#ffffc0", "#87CEEB"],}
-const cycleColors = (lst, i) => lst[i%lst.length];
+const devColors = {"TINY":["rgb(27 87 69)", "rgb(53 79 82)", "rgb(53 79 82)", "rgb(70 172 194)", "rgb(29, 46, 98)"],
+                   "DEFAULT":["rgb(29,31,42)","rgb(42,45,61)","rgb(55,59,79)","rgb(68,72,98)","rgb(18,19,26)","rgb(47,50,68)","rgb(59,63,84)","rgb(74,78,101)","rgb(24,26,35)","rgb(35,37,50)","rgb(49,53,72)","rgb(64,68,89)"],}
+const bufColors = ["#3A57B7","#5066C1","#6277CD","#7488D8","#8A9BE3","#A3B4F2"];
+
+const lighten = (rgb, depth, step=0.08) => rgb.replace(/\d+/g, n => Math.round(parseInt(n)+(255-parseInt(n)) * Math.min(1, depth*step)));

 var profileRet, focusedDevice, canvasZoom, zoomLevel = d3.zoomIdentity;
 async function renderProfiler() {
@@ -153,18 +152,18 @@ async function renderProfiler() {
    for (const e of timeline.shapes) {
      if (e.depth === 0) colorKey = e.cat ?? e.name;
      if (!colorMap.has(colorKey)) {
-        const colors = colorScheme[k] ?? colorScheme.DEFAULT;
+        const colors = devColors[k] ?? devColors.DEFAULT;
        colorMap.set(colorKey, colors[colorMap.size%colors.length]);
      }
-      const fillColor = d3.color(colorMap.get(colorKey)).brighter(e.depth).toString();
+      const fillColor = lighten(colorMap.get(colorKey), e.depth);
      const label = parseColors(e.name).map(({ color, st }) => ({ color, st, width:ctx.measureText(st).width }));
      if (e.ref != null) ref = {ctx:e.ref, step:0};
      else if (ref != null) {
        const start = ref.step>0 ? ref.step+1 : 0;
        const stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name);
-        ref = stepIdx === -1 ? null : {ctx:ref.ctx, step:stepIdx};
+        if (stepIdx !== -1) ref = {ctx:ref.ctx, step:stepIdx};
      }
-      const arg = { tooltipText:formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...ref };
+      const arg = { tooltipText:formatTime(e.dur), ...ref };
      // offset y by depth
      data.shapes.push({x:e.st-st, y:offsetY+levelHeight*e.depth, width:e.dur, height:levelHeight, arg, label, fillColor });
    }
@@ -172,7 +171,6 @@ async function renderProfiler() {
    const startY = offsetY+(levelHeight*timeline.maxDepth)+padding/2;
    let area = mem.shapes.length === 0 ? 0 : areaScale(mem.peak);
    if (area === 0) div.style.pointerEvents = "none";
-    else div.style.cursor = "pointer";
    if (k === focusedDevice) {
      // expand memory graph for the focused device
      area = maxArea*4;
@@ -184,7 +182,7 @@ async function renderProfiler() {
      const y0 = e.y.map(yscale);
      const y1 = e.y.map(y => yscale(y+e.arg.nbytes));
      const arg = { tooltipText:`${e.arg.dtype} len:${formatUnit(e.arg.sz)}\n${formatUnit(e.arg.nbytes, "B")}` };
-      data.shapes.push({ x, y0, y1, arg, fillColor:cycleColors(colorScheme.BUFFER, i) });
+      data.shapes.push({ x, y0, y1, arg, fillColor:bufColors[i%bufColors.length] });
    }
    // lastly, adjust device rect by number of levels
    div.style.height = `${Math.max(levelHeight*timeline.maxDepth, baseHeight)+area+padding}px`;
@@ -230,7 +228,6 @@ async function renderProfiler() {
      ctx.fillRect(x, e.y, width, e.height);
      rectLst.push({ y0:e.y, y1:e.y+e.height, x0:x, x1:x+width, arg:e.arg });
      // add label
-      if (e.label == null) continue;
      ctx.textAlign = "left";
      ctx.textBaseline = "middle";
      let [labelX, labelWidth] = [x+2, 0];
@@ -362,7 +359,7 @@ document.getElementById("zoom-to-fit-btn").addEventListener("click", () => {

 // **** main VIZ interfacae

-function codeBlock(st, language, { loc, wrap }={}) {
+function codeBlock(st, language, { loc, wrap }) {
  const code = document.createElement("code");
  code.innerHTML = hljs.highlight(st, { language }).value;
  code.className = "hljs";
@@ -377,19 +374,6 @@ function codeBlock(st, language, { loc, wrap }={}) {
  return ret;
 }

-function appendTd(tr, value, unit=null) {
-  const fmt = (typeof value === "number" && !Number.isInteger(value)) ? value.toFixed(2) : value;
-  tr.appendChild(document.createElement("td")).innerText = unit == "us" ? formatTime(value) : fmt+(unit ?? "");
-}
-
-function appendRow(table, name, value, unit=null, cls="main-row") {
-  const tr = table.appendChild(document.createElement("tr"));
-  tr.className = cls;
-  tr.appendChild(document.createElement("td")).innerText = name;
-  appendTd(tr, value, unit);
-  return tr;
-}
-
 function setActive(e) {
  if (e == null) return;
  e.classList.add("active");
@@ -469,7 +453,7 @@ async function main() {
      for (const [j,u] of steps.entries()) {
        const inner = ul.appendChild(document.createElement("ul"));
        inner.id = `step-${i}-${j}`;
-        inner.innerText = `${u.name ?? u.loc[0].replaceAll("\\", "/").split("/").pop()+':'+u.loc[1]}`+(u.match_count ? ` - ${u.match_count}` : '');
+        inner.innerText = `${u.name ?? u.loc[0].replaceAll("\\", "/").split("/").pop()+':'+u.loc[1]} - ${u.match_count}`;
        inner.style.marginLeft = `${8*u.depth}px`;
        inner.onclick = (e) => {
          e.stopPropagation();
@@ -483,69 +467,23 @@ async function main() {
  const { currentCtx, currentStep, currentRewrite, expandSteps } = state;
  if (currentCtx == -1) return;
  const ctx = ctxs[currentCtx];
-  const step = ctx.steps[currentStep];
-  const ckey = step?.query;
+  const ckey = `ctx=${currentCtx-1}&idx=${currentStep}`;
  // close any pending event sources
  let activeSrc = null;
  for (const e of evtSources) {
-    const url = new URL(e.url);
-    if (url.pathname+url.search !== ckey) e.close();
+    if (e.url.split("?")[1] !== ckey) e.close();
    else if (e.readyState === EventSource.OPEN) activeSrc = e;
  }
  if (ctx.name === "Profiler") return renderProfiler();
  if (ckey in cache) {
    ret = cache[ckey];
  }
-  // ** Disassembly view
-  if (ckey.startsWith("/disasm")) {
-    if (!(ckey in cache)) cache[ckey] = ret = await (await fetch(ckey)).json();
-    displayGraph("profiler");
-    const root = document.createElement("div");
-    root.className = "raw-text";
-    const metadata = document.querySelector(".metadata");
-    metadata.innerHTML = "";
-    // detailed assembly view
-    if (ret.cols != null) {
-      const asm = root.appendChild(document.createElement("table"));
-      const thead = asm.appendChild(document.createElement("thead"));
-      const usage = {};
-      for (const c of ret.cols) thead.appendChild(document.createElement("th")).innerText = c;
-      for (const r of ret.rows) {
-        const tr = asm.appendChild(document.createElement("tr"));
-        tr.className = "main-row code-row";
-        for (const d of Object.values(r.data)) appendTd(tr, d);
-        const segmentsTd = tr.appendChild(document.createElement("td"));
-        segmentsTd.className = "pct-row";
-        const usageBar = segmentsTd.appendChild(document.createElement("div"));
-        for (const [k, {width, value}] of Object.entries(r.segs)) {
-          const seg = usageBar.appendChild(document.createElement("div"));
-          seg.style.width = width+"%";
-          seg.title = `${ret.segments[k]} ${value}`;
-          seg.style.background = cycleColors(colorScheme.CATEGORICAL, parseInt(k));
-          if (!(k in usage)) usage[k] = 0;
-          usage[k] += value;
-        }
-      }
-      const summary = metadata.appendChild(document.createElement("table"));
-      for (const [i,s] of ret.segments.entries()) {
-        const tr = summary.appendChild(document.createElement("tr"));
-        tr.className = "main-row";
-        const td = tr.appendChild(document.createElement("td"));
-        const div = td.appendChild(document.createElement("div"));
-        div.className = "legend";
-        div.appendChild(document.createElement("div")).style.background = cycleColors(colorScheme.CATEGORICAL, i);
-        div.appendChild(document.createElement("p")).textContent = s;
-        appendTd(tr, usage[i] ?? 0);
-      }
-    } else root.appendChild(codeBlock(ret.src, "x86asm"));
-    return document.querySelector(".profiler").replaceChildren(root);
-  }
-  // ** UOp view (default)
  // if we don't have a complete cache yet we start streaming rewrites in this step
+  const step = ctx.steps[currentStep];
  if (!(ckey in cache) || (cache[ckey].length !== step.match_count+1 && activeSrc == null)) {
    ret = [];
    cache[ckey] = ret;
-    const eventSource = new EventSource(ckey);
+    const eventSource = new EventSource(`/ctxs?${ckey}`);
    evtSources.push(eventSource);
    eventSource.onmessage = (e) => {
      if (e.data === "END") return eventSource.close();
@@ -564,28 +502,6 @@ async function main() {
  const metadata = document.querySelector(".metadata");
  const [code, lang] = ctx.fmt != null ? [ctx.fmt, "cpp"] : [ret[currentRewrite].uop, "python"];
  metadata.replaceChildren(codeBlock(step.code_line, "python", { loc:step.loc, wrap:true }), codeBlock(code, lang, { wrap:false }));
-  if (ctx.runtime_stats != null) {
-    const div = metadata.appendChild(document.createElement("div"));
-    div.className = "stats-list";
-    for (const [i, s] of ctx.runtime_stats.entries()) {
-      const p = div.appendChild(document.createElement("p"));
-      if (ctx.runtime_stats.length > 1) p.innerText = `Run ${i+1}/${ctx.runtime_stats.length}`;
-      const table = div.appendChild(document.createElement("table"));
-      const tbody = table.appendChild(document.createElement("tbody"));
-      for (const { name, value, unit, subunits } of s.data) {
-          const mainRow = appendRow(tbody, name, value, unit, "main-row");
-          if (!subunits?.length) continue;
-          const subunitRow = tbody.appendChild(document.createElement("tr"));
-          subunitRow.style.display = "none";
-          mainRow.onclick = () => subunitRow.style.display = subunitRow.style.display === "none" ? "table-row" : "none";
-          mainRow.style.cursor = "pointer";
-          const td = subunitRow.appendChild(document.createElement("td"));
-          td.colSpan = 2;
-          const table = td.appendChild(document.createElement("table"));
-          for (const u of subunits) appendRow(table, u.name, u.value, unit, "sub-row");
-      }
-    }
-  }
  // ** rewrite steps
  if (step.match_count >= 1) {
    const rewriteList = metadata.appendChild(document.createElement("div"));
@@ -1,15 +1,12 @@
 #!/usr/bin/env python3
-import multiprocessing, pickle, difflib, os, threading, json, time, sys, webbrowser, socket, argparse, socketserver, functools, codecs, io
-import subprocess, ctypes
-from contextlib import redirect_stdout
+import multiprocessing, pickle, difflib, os, threading, json, time, sys, webbrowser, socket, argparse, socketserver, functools, codecs
 from decimal import Decimal
 from http.server import BaseHTTPRequestHandler
 from urllib.parse import parse_qs, urlparse
 from typing import Any, TypedDict, Generator
 from tinygrad.helpers import colored, getenv, tqdm, unwrap, word_wrap, TRACEMETA, ProfileEvent, ProfileRangeEvent, TracingKey
 from tinygrad.uop.ops import TrackedGraphRewrite, UOp, Ops, printable, GroupOp, srender, sint
-from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, ProfilePointEvent, Device
-from tinygrad.renderer import ProgramSpec
+from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, ProfilePointEvent
 from tinygrad.dtype import dtypes

 uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
@@ -28,12 +25,8 @@ ref_map:dict[Any, int] = {}
 def get_metadata(keys:list[TracingKey], contexts:list[list[TrackedGraphRewrite]]) -> list[dict]:
  ret = []
  for i,(k,v) in enumerate(zip(keys, contexts)):
-    steps = [{"name":s.name, "loc":s.loc, "depth":s.depth, "match_count":len(s.matches), "code_line":printable(s.loc),
-              "query":f"/ctxs?ctx={i}&idx={j}"} for j,s in enumerate(v)]
-    if isinstance(k.ret, ProgramSpec): steps.append({"name":"View Disassembly", "query":f"/disasm?ctx={i}"})
-    ret.append(r:={"name":k.display_name, "fmt":k.fmt, "steps":steps})
-    # use the first key to get runtime profiling data about this context
-    if getenv("PROFILE_VALUE") >= 2 and k.keys: r["runtime_stats"] = get_runtime_stats(k.keys[0])
+    steps = [{"name":s.name, "loc":s.loc, "depth":s.depth, "match_count":len(s.matches), "code_line":printable(s.loc)} for s in v]
+    ret.append({"name":k.display_name, "fmt":k.fmt, "steps":steps})
    for key in k.keys: ref_map[key] = i
  return ret

@@ -55,7 +48,7 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
  excluded: set[UOp] = set()
  for u in (toposort:=x.toposort()):
    # always exclude DEVICE/CONST/UNIQUE
-    if u.op in {Ops.DEVICE, Ops.CONST, Ops.UNIQUE} and u is not x: excluded.add(u)
+    if u.op in {Ops.DEVICE, Ops.CONST, Ops.UNIQUE}: excluded.add(u)
    # only exclude CONST VIEW source if it has no other children in the graph
    if u.op is Ops.CONST and len(u.src) != 0 and all(cr.op is Ops.CONST for c in u.src[0].children if (cr:=c()) is not None and cr in toposort):
      excluded.update(u.src)
@@ -126,15 +119,12 @@ def timeline_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict:
    depth = next((i for i,level_et in enumerate(levels) if st>=level_et), len(levels))
    if depth < len(levels): levels[depth] = et
    else: levels.append(et)
-    name, cat, info = e.name, None, None
-    if (ref:=ref_map.get(name)) is not None:
-      name = ctxs[ref]["name"]
-      if isinstance(p:=contexts[0][ref].ret, ProgramSpec):
-        info = f"{p.estimates.ops/(t:=dur*1e3):.2f} GFLOPS {p.estimates.mem/t:4.1f}|{p.estimates.lds/t:.1f} GB/s"
+    name, cat = e.name, None
+    if (ref:=ref_map.get(name)) is not None: name = ctxs[ref]["name"]
    elif isinstance(e.name, TracingKey):
      name, cat = e.name.display_name, e.name.cat
      ref = next((v for k in e.name.keys if (v:=ref_map.get(k)) is not None), None)
-    shapes.append({"name":name, "ref":ref, "st":st, "dur":dur, "depth":depth, "cat":cat, "info":info})
+    shapes.append({"name":name, "ref":ref, "st":st, "dur":dur, "depth":depth, "cat":cat})
  return {"shapes":shapes, "maxDepth":len(levels)}

 def mem_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict:
@@ -182,39 +172,6 @@ def get_profile(profile:list[ProfileEvent]):
  dev_layout = {k:{"timeline":timeline_layout(v), "mem":mem_layout(v)} for k,v in dev_events.items()}
  return json.dumps({"layout":dev_layout, "st":min_ts, "et":max_ts}).encode("utf-8")

-def get_runtime_stats(key) -> list[dict]:
-  ret:list[dict] = []
-  for e in profile:
-    if isinstance(e, ProfileRangeEvent) and e.en is not None and e.name == key:
-      ret.append({"device":e.device, "data":[{"name":"Duration", "value":float(e.en-e.st), "unit":"us"}]})
-  return ret
-
-def get_disassembly(ctx:list[str]):
-  if not isinstance(prg:=contexts[0][int(ctx[0])].ret, ProgramSpec): return
-  lib = (compiler:=Device[prg.device].compiler).compile(prg.src)
-  with redirect_stdout(buf:=io.StringIO()): compiler.disassemble(lib)
-  disasm_str = buf.getvalue()
-  from tinygrad.runtime.ops_llvm import llvm, LLVMCompiler
-  if isinstance(compiler, LLVMCompiler):
-    mtriple = ctypes.string_at(llvm.LLVMGetTargetMachineTriple(tm:=compiler.target_machine)).decode()
-    mcpu = ctypes.string_at(llvm.LLVMGetTargetMachineCPU(tm)).decode()
-    # NOTE: llvm-objdump may contain headers, skip if llvm-mca can't parse those lines
-    data = json.loads(subprocess.check_output(["llvm-mca", f"-mtriple={mtriple}", f"-mcpu={mcpu}", "-skip-unsupported-instructions=parse-failure",
-                                               "--json", "-"], input=disasm_str.encode()))
-    cr = data["CodeRegions"][0]
-    instrs:list = [{"data":[rep], "segs":{}} for rep in cr["Instructions"]]
-    for i,info in enumerate(cr["InstructionInfoView"]["InstructionList"]): instrs[i]["data"].append(info["Latency"])
-    for d in cr["ResourcePressureView"]["ResourcePressureInfo"]:
-      i, r = d["InstructionIndex"], d["ResourceIndex"]
-      if i>len(instrs)-1: continue
-      instrs[i]["segs"][r] = instrs[i]["segs"].get(r, 0)+d["ResourceUsage"]
-    # rescale segment width to 0-100
-    if instrs:
-      hi = max([sum(ins["segs"].values()) for ins in instrs])
-      for n in instrs: n["segs"] = {k:{"width":v/hi*100, "value":v} for k,v in n["segs"].items()}
-    return json.dumps({"rows":instrs, "cols":["Opcode", "Latency", "HW Resources"], "segments":data["TargetInfo"]["Resources"]}).encode()
-  return json.dumps({"src":disasm_str}).encode()
-
 # ** HTTP server

 class Handler(BaseHTTPRequestHandler):
@@ -229,10 +186,9 @@ class Handler(BaseHTTPRequestHandler):
        if url.path.endswith(".js"): content_type = "application/javascript"
        if url.path.endswith(".css"): content_type = "text/css"
      except FileNotFoundError: status_code = 404
-    elif (query:=parse_qs(url.query)):
-      if url.path == "/disasm": ret, content_type = get_disassembly(**query), "application/json"
-      else: return self.stream_json(get_details(contexts[1][int(query["ctx"][0])][int(query["idx"][0])]))
-    elif url.path == "/ctxs": ret, content_type = json.dumps(ctxs).encode(), "application/json"
+    elif url.path == "/ctxs":
+      if "ctx" in (q:=parse_qs(url.query)): return self.stream_json(get_details(contexts[1][int(q["ctx"][0])][int(q["idx"][0])]))
+      ret, content_type = json.dumps(ctxs).encode(), "application/json"
    elif url.path == "/get_profile" and profile_ret is not None: ret, content_type = profile_ret, "application/json"
    else: status_code = 404
Author	SHA1	Message	Date
geohot	50be175f88	store	2025-07-22 19:28:39 -07:00
geohot	33d80db64c	this can happen later	2025-07-22 19:12:37 -07:00
geohot	28b6cff043	broken	2025-07-22 19:05:11 -07:00
geohot	23ee2769bd	gate store	2025-07-22 19:03:06 -07:00
geohot	5055668d1e	oops, forgot that	2025-07-22 18:54:14 -07:00
geohot	8305a5804c	just the range thing	2025-07-22 18:52:26 -07:00
George HotzandGitHub	911bb4ff44	Merge branch 'master' into endrange	2025-07-22 18:43:22 -07:00
geohot	942d69e139	store is endrange	2025-07-22 15:30:10 -07:00
geohot	343921f873	Revert "end the ranges in the stores" This reverts commit `c88ad9d24f`.	2025-07-22 15:12:21 -07:00
geohot	c88ad9d24f	end the ranges in the stores	2025-07-22 15:04:10 -07:00
geohot	60dcc9f4df	insert endrange	2025-07-22 14:38:49 -07:00