Merge branch 'master' into bitcast_spec

simpler, but not inverse
fix bitcast
2026-08-24 16:06:07 +00:00 · 2026-08-12 19:24:33 -07:00 · 2026-06-25 10:24:01 -07:00 · 2026-06-25 10:11:13 -07:00 · 2026-06-25 10:00:14 -07:00
157 changed files with 3689 additions and 3247 deletions
@@ -4,7 +4,7 @@ inputs:
  python-version:
    description: 'Python version to use'
    required: false
-    default: '3.14'
+    default: '' # if you don't set a version, the native python version will be used
  key:
    description: 'Key for the python cache'
    required: false
@@ -41,12 +41,12 @@ inputs:
    description: "Install LLVM?"
    required: false
    default: 'false'
-  qemu:
-    description: "Install qemu?"
+  tinydreno:
+    description: "Install tinydreno"
    required: false
    default: 'false'
-  ninja:
-    description: "Install ninja?"
+  qemu:
+    description: "Install qemu"
    required: false
    default: 'false'
 runs:
@@ -59,18 +59,18 @@ runs:
        echo "OMP_NUM_THREADS=1" >> "$GITHUB_ENV"
        # no buffers should be over 300MB in CI
        echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
-        if [[ "$RUNNER_OS" == "Linux" ]]; then
-          echo "VIRTUAL_ENV=/opt/venv/${{ inputs.python-version }}" >> "$GITHUB_ENV"
-          echo "UV_PYTHON_INSTALL_DIR=/opt/python" >> "$GITHUB_ENV"
-        else
-          echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
-        fi

    - name: Set up uv
      uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b
      with:
        enable-cache: 'false' # see below for manual caching

+    - name: Set up Python ${{ inputs.python-version }}
+      uses: actions/setup-python@v6
+      if: inputs.python-version != ''
+      with:
+        python-version: ${{ inputs.python-version }}
+
    # **** Caching packages ****

    - name: Cache Python packages (PR)
@@ -109,15 +109,15 @@ runs:
      if: inputs.deps != ''
      shell: bash
      run: |
-        uv venv --allow-existing --python ${{ inputs.python-version }} "$VIRTUAL_ENV"
+        uv venv .venv
        DEPS="${{ inputs.deps }}"
-        uv pip install --python "$VIRTUAL_ENV" -e ".[${DEPS// /,}]" ${{ inputs.pydeps }} --torch-backend cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
+        uv pip install --python .venv -e ".[${DEPS// /,}]" ${{ inputs.pydeps }} --torch-backend cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
    - name: Install dependencies in venv (without extra)
      if: inputs.deps == ''
      shell: bash
      run: |
-        uv venv --allow-existing --python ${{ inputs.python-version }} "$VIRTUAL_ENV"
-        uv pip install --python "$VIRTUAL_ENV" -e . ${{ inputs.pydeps }}
+        uv venv .venv
+        uv pip install --python .venv -e . ${{ inputs.pydeps }}
    - name: Prune uv cache
      if: github.event_name != 'pull_request'
      shell: bash
@@ -125,15 +125,16 @@ runs:
    - name: Configure venv
      shell: bash
      run: |
+        echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
        if [[ "$RUNNER_OS" == "Windows" ]]; then
-          echo "$VIRTUAL_ENV/Scripts" >> "$GITHUB_PATH"
+          echo "${{ github.workspace }}/.venv/Scripts" >> "$GITHUB_PATH"
        else
-          echo "$VIRTUAL_ENV/bin" >> "$GITHUB_PATH"
+          echo "${{ github.workspace }}/.venv/bin" >> "$GITHUB_PATH"
        fi

    # ******************* apt *******************
    - name: Setup apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
      shell: bash
      run: |
        sudo mkdir -p /var/cache/apt/archives
@@ -161,7 +162,7 @@ runs:
        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
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
      id: apt-pkgs
      shell: bash
      run: |
@@ -186,29 +187,25 @@ runs:
        if [[ "${{ inputs.qemu }}" == "true" ]]; then
          pkgs+=" qemu-user-static"
        fi
-        # **** ninja ****
-        if [[ "${{ inputs.ninja }}" == "true" ]]; then
-          pkgs+=" ninja-build"
-        fi

        echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
        echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"

    - name: Cache apt (PR)
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name == 'pull_request'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true') && github.event_name == 'pull_request'
      uses: actions/cache/restore@v5
      with:
        path: /var/cache/apt/archives/
        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
    - name: Cache apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name != 'pull_request'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true') && github.event_name != 'pull_request'
      uses: actions/cache@v5
      with:
        path: /var/cache/apt/archives/
        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}

    - name: Run apt Update + Install
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
      shell: bash
      run: |
        sudo apt -qq update || true
@@ -280,6 +277,12 @@ runs:
      shell: bash
      run: brew install llvm@20

+    # *** tinydreno ***
+    - name: Install tinydreno (linux)
+      if: inputs.tinydreno == 'true' && runner.os == 'Linux'
+      shell: bash
+      run: sudo curl -fL https://github.com/sirhcm/tinydreno/raw/refs/heads/master/libllvm-qcom.so -o /usr/lib/libllvm-qcom.so
+
    # *** OpenCL ***
    - name: Install rusticl
      if: inputs.opencl == 'true'
@@ -35,7 +35,7 @@ jobs:
        key: 'autogen'
        amd: 'true'
        llvm: 'true'
-        deps: 'autogen'
+        pydeps: 'pyyaml mako'
    - name: Install autogen support packages
      run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev liburing-dev
    - name: Regenerate autogen files
@@ -108,6 +108,10 @@ jobs:
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -125,6 +129,10 @@ jobs:
      # just metal for now
      if: ${{ matrix.dev == 'METAL' }}
      run: BENCHMARK_LOG=olmoe JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m olmoe --benchmark --warmup
+    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
+      # only run on machines with multiple gpus
+      if: ${{ matrix.dev != 'METAL' }}
+      run: BENCHMARK_LOG=llama3_beam_4gpu JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -174,6 +182,10 @@ jobs:
      # slow on metal
      if: ${{ matrix.dev != 'METAL' }}
      run: time BENCHMARK_LOG=cifar DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Run full CIFAR training steps w 6 GPUS
+      # only run on machines with multiple gpus
+      if: ${{ matrix.dev != 'METAL' }}
+      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -215,8 +227,15 @@ jobs:
        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: Run MLPerf resnet eval on training data
+      run: time BENCHMARK_LOG=resnet_eval MODEL=resnet python3 examples/mlperf/model_eval.py
    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
      run: BENCHMARK_LOG=resnet_10steps DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
+    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+      run: BENCHMARK_LOG=resnet_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
+    - name: Run 10 MLPerf Bert training steps (6 gpu)
+      # TODO: remove BERT_LAYERS once scheduler is fast
+      run: BENCHMARK_LOG=bert_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -266,60 +285,6 @@ jobs:
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

-  multigpubenchmark:
-    name: Multi-GPU Benchmarks (DEV=${{ matrix.dev }})
-    runs-on: [self-hosted, "${{ matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
-    strategy:
-      fail-fast: false
-      matrix:
-        dev: ['AMD', 'NV']
-    timeout-minutes: 60
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    env:
-      DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup (AMD)
-      if: ${{ matrix.dev == 'AMD' }}
-      run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
-    - name: Setup (NV)
-      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
-    - name: Symlink models and datasets
-      run: |
-        mkdir -p weights
-        mkdir -p extra/datasets
-        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
-        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: python3 test/external/process_replay/reset.py
-    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      run: BENCHMARK_LOG=llama3_beam_4gpu JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
-    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
-    - name: Run MLPerf resnet eval on training data
-      run: time BENCHMARK_LOG=resnet_eval MODEL=resnet python3 examples/mlperf/model_eval.py
-    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
-    - name: Run 10 MLPerf Bert training steps (6 gpu)
-      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
  tests:
    name: Tests (DEV=${{ matrix.dev }})
    runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
@@ -445,7 +410,7 @@ jobs:
    - name: UsbGPU tiny tests
      run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
    - name: UsbGPU copy speeds
-      run: sudo -E PYTHONDONTWRITEBYTECODE=1 SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
    #- name: UsbGPU openpilot test
    #  run: sudo -E PYTHONPATH=. GMMU=0 DEV=USB+AMD GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
    - name: UsbGPU (USB4/TB) install script
@@ -576,7 +541,7 @@ jobs:
    - name: openpilot run_pickle big_driving_supercombo
      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_run_pickle RUN_PICKLE=1 PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py - openpilot.pkl
    - name: Test copy speeds
-      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      run: SIZE=64e6 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds

  driverbenchmarks:
    name: PCI Driver Benchmark (DEV=${{ matrix.dev }})
@@ -598,8 +563,8 @@ jobs:
    - name: Setup
      run: |
        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} rmmod
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} kill_pids
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev == 'AMD' && 'amd' || 'nv' }} rmmod
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev == 'AMD' && 'amd' || 'nv' }} kill_pids
        mkdir -p extra/datasets
        ln -s /raid/datasets/imagenet extra/datasets/imagenet
    - name: setup staging db
@@ -634,9 +599,6 @@ jobs:
      run: |
        GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
        GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
-    - name: HEVC Decode Benchmark
-      if: ${{ matrix.dev == 'NV' }}
-      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
      if: ${{ matrix.dev == 'NV' }}
      run: BENCHMARK_LOG=resnet_10steps MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
@@ -8,7 +8,7 @@ permissions:
  contents: write
 jobs:
  deploy:
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v6
      - name: Configure Git Credentials
@@ -166,7 +166,7 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: windows-${{ matrix.dev }}-minimal
-          deps: testing_minimal
+          deps: testing_unit
          pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }}
      - name: Set env
        shell: bash
@@ -179,3 +179,35 @@ jobs:
      - name: Run test_tiny
        shell: bash
        run: python -m pytest -n=auto test/test_tiny.py --durations=20
+
+
+  qcomclcompiletests:
+    name: Compile-only (QCOM CL)
+    runs-on: ubuntu-24.04-arm
+    timeout-minutes: 15
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v6
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: compile-qcomcl
+          deps: testing_unit
+          tinydreno: 'true'
+      - name: Set env
+        shell: bash
+        run: printf "DEV=NULL:QCOMCL:a630\nNULL_ALLOW_COPYOUT=1" >> $GITHUB_ENV
+      - name: Run test_ops
+        shell: bash
+        run: |
+          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
+          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20
+      - name: Run test_ops (IMAGE)
+        shell: bash
+        env:
+          IMAGE: 1
+          DEV: "NULL:QCOMCL:a630,IMAGE_PITCH_ALIGNMENT=64"
+        run: |
+          DEBUG=4 python test/backend/test_ops.py TestOps.test_gemm | grep read_imagef
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20
@@ -10,7 +10,7 @@ on:

 jobs:
  deploy:
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-latest
    steps:
    - uses: actions/checkout@v6
    - name: Set up Python
@@ -10,7 +10,7 @@ concurrency:
 jobs:
  checkbranch:
    name: Check PR Branch status
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-latest
    outputs:
      branchstat: ${{ steps.brstat.outputs.stat}}
    steps:
@@ -44,7 +44,7 @@ jobs:
    permissions:
      contents: read
      pull-requests: write
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-latest
    needs: checkbranch
    if: needs.checkbranch.outputs.branchstat == 'false'
    steps:
@@ -87,7 +87,7 @@ jobs:
    name: Core Library Line Difference
    permissions:
      pull-requests: write
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-latest
    needs: checkbranch
    if: needs.checkbranch.outputs.branchstat == 'true'
    steps:
@@ -21,7 +21,7 @@ concurrency:
 jobs:
  docs:
    name: Docs
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: &linux ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 10
    env:
      CHECK_OOB: 0
@@ -31,7 +31,8 @@ jobs:
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        deps: "docs testing_minimal"
+        deps: docs
+        pydeps: "capstone torch"
    - name: Build wheel and show size
      run: |
        uv build --wheel
@@ -60,7 +61,7 @@ jobs:

  torchbackend:
    name: Torch Backend Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -72,7 +73,10 @@ jobs:
        deps: testing_unit
        pydeps: "pillow torchvision expecttest"
        llvm: 'true'
-        ninja: 'true'
+    - name: Install ninja
+      run: |
+        sudo apt update || true
+        sudo apt install -y --no-install-recommends ninja-build
    - name: Test ResNet-18
      run: DEBUG=2 python3 extra/torch_backend/example.py
    - name: Test one op in torch tests
@@ -82,26 +86,9 @@ jobs:
    - name: Custom tests
      run: DEV=CPU:LLVM GPUS=4 TINY_BACKEND=1 python3 -m pytest -nauto extra/torch_backend/test.py extra/torch_backend/test_inplace.py extra/torch_backend/test_multigpu.py extra/torch_backend/test_kernel_fusion.py --durations=20

-  torchbackendtrain:
-    name: Torch Backend Training
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_unit
-        llvm: 'true'
-        ninja: 'true'
-    - name: Test beautiful_mnist in torch with TINY_BACKEND
-      run: STEPS=20 DEV=CPU TARGET_EVAL_ACC_PCT=90.0 MAX_BUFFER_SIZE=0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
-
  bepython:
    name: Python Backend
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -139,7 +126,7 @@ jobs:

  linter:
    name: Linters
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 10

    steps:
@@ -170,7 +157,7 @@ jobs:

  nulltest:
    name: Null Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15

    steps:
@@ -204,7 +191,7 @@ jobs:

  unittest:
    name: Unit Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15

    steps:
@@ -241,7 +228,7 @@ jobs:
      matrix:
        group: [1, 2]
    name: SPEC=2 (${{ matrix.group }})
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -257,7 +244,7 @@ jobs:

  fuzzing:
    name: Fuzzing
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 10
    steps:
    - name: Checkout Code
@@ -273,7 +260,7 @@ jobs:

  testopenclimage:
    name: CL IMAGE Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -293,7 +280,7 @@ jobs:

  testopenpilot:
    name: openpilot Compile Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -322,7 +309,7 @@ jobs:

  testonnxcpu:
    name: ONNX (CPU) Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 20

    steps:
@@ -341,7 +328,7 @@ jobs:

  testoptim:
    name: Optimization Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 20
    steps:
      - name: Checkout Code
@@ -373,7 +360,7 @@ jobs:

  testllm:
    name: Test LLM
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    env:
      CHECK_OOB: 0
@@ -400,7 +387,7 @@ jobs:

  testmodels:
    name: Models
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -420,7 +407,7 @@ jobs:

  testdsp:
    name: Linux (DSP)
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -448,7 +435,7 @@ jobs:
          - 'WEBGPU'

    name: Linux (DEV=${{ matrix.dev }})
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 20
    steps:
    - name: Checkout Code
@@ -474,7 +461,7 @@ jobs:

  testamdasm:
    name: AMD ASM IDE
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 20
    env:
      DEV: MOCKKFD+AMD
@@ -504,7 +491,7 @@ jobs:
      - name: Run AMD renderer tests (AMD:LLVM)
        run: DEV=MOCKKFD+AMD:LLVM python -m pytest -n=auto test/amd/ --durations 20
      - name: Run SQTT profiling tests
-        run: VIZ=-2 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
+        run: PROFILE=1 SQTT=1 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
      - name: Run AMD emulated tests on NULL backend
        env:
          AMD: 0
@@ -520,7 +507,7 @@ jobs:

  hcq2:
    name: hcq2
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 5
    steps:
      - name: Checkout Code
@@ -542,7 +529,7 @@ jobs:

  testmockam:
    name: Linux (am)
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    env:
      DEV: MOCKPCI+AMD
@@ -578,7 +565,7 @@ jobs:
        arch: [gfx1100, gfx1201, gfx950]

    name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    env:
      DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
@@ -616,7 +603,7 @@ jobs:
        backend: [ptx, nv]

    name: Linux (${{ matrix.backend }})
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 20
    env:
      FORWARD_ONLY: 1
@@ -650,17 +637,10 @@ jobs:
    strategy:
      fail-fast: false
      matrix:
-        dev:
-          - 'NULL:IR3:a630'
-          - 'NULL:QCOMCL:a630'
-          - 'NULL:NAK:sm_120'
-    name: Compile-only (DEV=${{ matrix.dev }})
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+        backend: [ir3, nak]
+    name: Compile-only (${{ matrix.backend }})
+    runs-on: *linux
    timeout-minutes: 15
-    env:
-      NULL_ALLOW_COPYOUT: 1
-      DEV: ${{ matrix.dev }}${{ contains(matrix.dev, 'a630') && ',IMAGE_PITCH_ALIGNMENT=64' || '' }}
-      IMAGE: ${{ contains(matrix.dev, 'a630') && '1' || '0' }}
    steps:
      - name: Checkout Code
        uses: actions/checkout@v6
@@ -669,15 +649,21 @@ jobs:
        with:
          key: compile-${{ matrix.backend }}
          deps: "testing_unit mesa"
-          qemu: ${{ contains(matrix.dev, 'QCOMCL') }}
-      - name: Test IMAGE
+      - name: Set env
        shell: bash
-        if: contains(matrix.dev, 'a630')
-        run: DEBUG=7 python3 test/backend/test_ops.py TestOps.test_gemm | grep isam
+        run: printf "NULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'DEV=NULL:IR3:a630' || matrix.backend == 'nak' && 'DEV=NULL:NAK:sm_120' }}" >> $GITHUB_ENV
      - name: Run test_ops
        shell: bash
        run: |
          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
-          # QCOMCL compiles in qemu, too slow for parallel workers
-          ${{ contains(matrix.dev, 'QCOMCL') && 'PARALLEL=0' || '' }} python -m pytest -n=auto test/backend/test_ops.py --durations=20
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20
+      - name: Run test_ops (IMAGE)
+        if: matrix.backend == 'ir3'
+        shell: bash
+        env:
+          IMAGE: 1
+          DEV: "NULL:IR3:a630,IMAGE_PITCH_ALIGNMENT=64"
+        run: |
+          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_gemm | grep image_load
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20
@@ -69,4 +69,3 @@ mutants
 dagre/
 graphlib/
 uv.lock
-pi_session_window0.jsonl
@@ -140,7 +140,7 @@ Documentation along with a quick start guide can be found on the [docs website](
 ```python
 from tinygrad import Tensor

-x = Tensor.eye(3).clone()  # clone to make it a buffer
+x = Tensor.eye(3)
 y = Tensor([[2.0,0,-2.0]])
 z = y.matmul(x).sum()
 z.backward()
@@ -1742,8 +1742,8 @@ def train_gptoss():
  )

  for p in optim.params:
-    p.grad = p.zeros_like(dtype=dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype).contiguous()
-    if getattr(p, "_zero2", False): p.grad = optim.optimizers[0]._zero_shard(p.grad)
+    grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
+    p.grad = p.zeros_like(dtype=grad_dtype).contiguous()
  grads = [p.grad for p in optim.params]

  from extra.gemm.cdna_asm_gemm import _mx_block_scale
@@ -146,7 +146,6 @@ class GPTOSS:
      return w_q, w_e8.is_param_(False)
    if moe:
      qs = [_one(*shape[1:]) for _ in range(shape[0])]
-      for q in qs: q[0]._zero2 = True  # grad arrives sharded on the expert axis under ZeRO-2 (moe_gemm)
      return [q[0] for q in qs], [q[1] for q in qs]
    return _one(*shape)

@@ -183,14 +182,12 @@ class GPTOSS:
    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
    xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)  # (B,N,H,D)/(B,N,KV,D)

-    fa_saves = []
-    if getenv("HK_FLASH_ATTENTION"):
-      from extra.thunder.amd.fa import flash_attention
-      attn, _, l_vec = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks, window=self.sliding_window if sliding else 0)
-      attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
-      fa_saves = [xq, xk, xv, l_vec]
-    elif sliding:
+    if sliding:
      attn = self._sliding_attention(xq, xk, xv, sinks)
+    elif getenv("HK_FLASH_ATTENTION"):
+      from extra.thunder.amd.fa import flash_attention
+      attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks)
+      attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
    else:
      xqm = xq.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
      xkm, xvm = xk.permute(0, 2, 1, 3).unsqueeze(2), xv.permute(0, 2, 1, 3).unsqueeze(2)
@@ -202,7 +199,7 @@ class GPTOSS:
      attn = (w @ xvm).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, self.n_heads * self.head_dim)

    out = matmul_mx(attn, wo, wo_scale) + wo_bias
-    return out, [x_normed, rrms, attn] + fa_saves
+    return out, [x_normed, rrms, attn]

  def feed_forward(self, x:Tensor, *, ffn_norm:Tensor, gate:Tensor, gate_bias:Tensor,
                   w_gate_up:Tensor, w_gate_up_scale:Tensor, w_gate_up_bias:Tensor,
@@ -223,7 +220,6 @@ class GPTOSS:
      z = grouped_mx_gemm(_pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \
          + (onehot @ w_down_bias.float()).cast(dtypes.bfloat16)
      out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim)
-      return out, [x_normed, rrms, xg, h, y, z, r.weights, r.dest_row, r.off]
    else:
      thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
      weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
@@ -11,7 +11,6 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
-export GROUPED_MOE=${GROUPED_MOE:-1}
 export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
@@ -11,7 +11,6 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
-export GROUPED_MOE=${GROUPED_MOE:-1}
 export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
@@ -35,7 +35,7 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li
    return name

  for call in iter_kernel_calls(linear):
-    arg_uops = [b for b in call.src[1:] if not b.is_bound_var]
+    arg_uops = [b for b in call.src[1:] if b.op is not Ops.BIND]
    prg = to_program(call.src[0], Device[arg_uops[0].device].renderer)
    info = prg.arg
    functions[info.function_name] = prg.src[2].arg
@@ -122,8 +122,7 @@ def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp,
  groups_x, groups_y = UOp.special(ceildiv(N, tile_n), "gidx0"), UOp.special(ceildiv(M, tile_m), "gidx1")
  lds = UOp.placeholder((163840,), dtypes.uint8, 0, AddrSpace.LOCAL)
  sink = UOp.sink(C.base, A.base, B.base, scale_a.base, scale_b.base, *(x.base for x in extra), lds, threads, groups_x, groups_y,
-                  arg=KernelInfo(f"mxfp4_gemm_{M}_{N}_{K}",
-                                 estimates=Estimates(ops=2*M*N*K, mem=(M*half_k+N*half_k)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
+                  arg=KernelInfo(f"custom_mxfp4_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K)))
  insts = build_kernel(M, N, K, tile_m, tile_n)
  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts))))

@@ -1,32 +1,10 @@
 import functools, pathlib
 from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
-from tinygrad.helpers import getenv
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 from extra.gemm.cdna_asm_gemm import quantize_mxfp8, _mx_block_scale, _mx_block_scale_3d

-ZERO_OPTIM = getenv("ZERO_OPTIM", 0)
-
-def reduce_scatter_devaxis(out:Tensor, shard_axis:int=0) -> Tensor:
-  # out: sharded on the device axis, shape (ndev, *rest); return the device-axis sum left sharded on shard_axis.
-  u = out.uop
-  devs, rest = u.device, u.shape[1:]
-  assert rest[shard_axis] % len(devs) == 0, f"reduce_scatter needs even shards: {rest[shard_axis]} % {len(devs)}"
-  # reach the raw per-device buffer below the UNSHARD, keeping the AFTERs so reads stay ordered after the kernel writes
-  node, barriers = u, []
-  while node.op is not Ops.UNSHARD:
-    if node.op is Ops.AFTER: barriers += node.src[1:]
-    node = node.src[0]
-  mbuf = node.src[0].after(*barriers) if barriers else node.src[0]
-  sz = rest[shard_axis] // len(devs)
-  shards = []
-  for i in range(len(devs)):
-    bounds = tuple((0,s) if a != shard_axis else (i*sz,(i+1)*sz) for a,s in enumerate(rest))
-    contribs = [mbuf.mselect(j).reshape(rest).shrink(bounds).copy_to_device(devs[i]) for j in range(len(devs))]
-    shards.append(functools.reduce(lambda a,b: a.alu(Ops.ADD, b), contribs))
-  return Tensor(UOp.mstack(*shards).unshard(shard_axis, UOp.range(len(devs), -1, AxisType.DEVICE)), device=devs)
-
@functools.cache
 def custom_hk_grouped_mxfp8_gemm(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, *extra:UOp, dname:str, n_experts:int) -> UOp:
  M, K = A.shape
@@ -80,8 +58,7 @@ def grouped_mx_wgrad(g:Tensor, xg:Tensor, expert_off:Tensor, n_experts:int) -> T
  out = Tensor(inv.uop.unshard(0), device=g.device) if is_multi else inv
  out = Tensor.custom_kernel(out, gT, xT, g_si, x_si, expert_off,
                             fxn=functools.partial(custom_hk_grouped_mxfp8_wgrad, dname=dname, n_experts=n_experts))[0]
-  if is_multi and ZERO_OPTIM: out = reduce_scatter_devaxis(out, 0)
-  else: out = out.sum(0) if is_multi else out.squeeze(0)
+  out = out.sum(0) if is_multi else out.squeeze(0)
  return out.reshape(n_experts, N, K)

 def mx_pack_3d(e8:Tensor) -> Tensor:
@@ -79,7 +79,7 @@ if __name__ == "__main__":
  linear, var_vals = C.linear_with_vars()
  last_call = linear.src[-1]
  ast = last_call.src[0]
-  bufs = [s.buffer for s in last_call.src[1:] if not s.is_bound_var]
+  bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]

  src = compiled.asm["ptx"]
  # specify the shared memory here so we don't need to do it dynamically
@@ -97,13 +97,11 @@ 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", aliases=["NV"], help="NVIDIA GPUs")
-  nv_parser.set_defaults(backend="nv")
+  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", aliases=["AMD"], help="AMD GPUs")
-  amd_parser.set_defaults(backend="amd")
+  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)

@@ -3,7 +3,7 @@ from typing import cast, Any, Callable
 import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit
 assert sys.platform != 'win32'
 from dataclasses import dataclass
-from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, encode_kernargs_clike, make_cmdbuf
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_cmdbuf
 from tinygrad.runtime.support.hcq2 import make_binary_patch
 from tinygrad.uop.ops import sint, UOp
 from tinygrad.device import Compiled, BufferSpec, Buffer, Device
@@ -179,10 +179,11 @@ class SDMAOps(FastEnum): COPY = auto(); POLL_REGMEM = auto(); FENCE = auto(); TR

 def sdma_copy(ctx, call):
  sz = call.src[2].max_numel() * call.src[2].dtype.itemsize
-  hdr = ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR)
-  return call.ins(SDMAOps.COPY, src=tuple(x for off in range(0, sz, ctx.max_copy_size) for x in (
-    *(UOp.const(v, dtypes.uint32) for v in (hdr, ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0)),
-    *(a + UOp.const(off, dtypes.uint64) if off else a for a in (call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs))))))
+  src_addr, dst_addr = call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs)
+  return call.ins(SDMAOps.COPY, src=tuple(UOp.const(x, dtypes.uint32) for off in range(0, sz, ctx.max_copy_size) for x in (
+    ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR),
+    ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0,
+    *data64_le(src_addr+UOp.const(off, dtypes.uint64)), *data64_le(dst_addr+UOp.const(off, dtypes.uint64)))))

 def sdma_wait(ctx, ins, dst, val):
  op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
@@ -261,11 +262,10 @@ class AMDProgramData:
  private_segment_size:int; kernargs_segment_size:int; kernargs_alloc_size:int
  enable_dispatch_ptr:int; enable_private_segment_sgpr:int

-_amd_program_cache:dict[tuple[bytes, tuple[str, ...]], UOp] = {}
+_amd_program_cache:dict[tuple[bytes,str], tuple[AMDProgramData,bytes]] = {}
 def amd_build_program(prg:UOp) -> UOp:
  dev = Device[to_tuple(prg.device)[0]] # TODO: rm this
-  # key on the full device tuple: the same lib can be built for different device sets, each needs its own program buffer
-  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, to_tuple(prg.device)))) is None:
+  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, dev.device))) is None:
    image, sections, relocs = elf_loader(lib)
    rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata")
    for off, sym, typ, addent in relocs:
@@ -290,14 +290,14 @@ class AMDAllocator(HCQAllocator['AMDDevice']):
  def __init__(self, dev:AMDDevice):
    super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb())

-  def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
+  def _alloc(self, size:int, options:BufferSpec) -> HCQ2Buffer:
    return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_copy_queue)

  def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque)

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

-  def _do_unmap(self, buf:HCQBuffer): self.dev.iface.unmap(buf)
+  def _do_unmap(self, buf:HCQ2Buffer): self.dev.iface.unmap(buf)

@dataclass
 class AMDQueueDesc:
@@ -561,7 +561,9 @@ class AMDDevice(HCQ2Compiled):
  def is_usb(self) -> bool: return False

  def __init__(self, device:str=""):
-    self.iface = self._select_iface(device)
+    self.device_id = int(device.split(":")[1]) if ":" in device else 0
+
+    self.iface = self._select_iface()

    self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
    self.arch = "gfx%d%x%x" % self.target
@@ -16,9 +16,8 @@ apt-get install -y python3-pip python3-venv git tmux rclone clang

 ### 1.2 Install Python deps
 ```bash
-python3 -m pip install --break-system-packages --ignore-installed typing-extensions numpy tqdm wandb tiktoken sentencepiece
+python3 -m pip install --break-system-packages numpy tqdm wandb tiktoken sentencepiece
 ```
-Note: `--ignore-installed typing-extensions` is needed because the base image ships typing-extensions 4.10.0 without a RECORD file, so pip cannot uninstall it.

 ### 1.3 Install ROCm dev headers
 The base image has ROCm runtime but NOT the HIP dev headers. Need:
@@ -86,45 +85,36 @@ Files downloaded (~85GB total, ~6 minutes):
 - `c4-validation-91205-samples.en_text_document.idx` (1.8 MB)
 - `LICENSE.txt`, `NOTICE.txt`

-**Wait for rclone to fully complete before starting training.** Starting training while the dataset is still downloading will read a truncated .bin file, causing `ValueError: all input arrays must have the same shape` in the dataloader. The stale `.index_cache` and `.blend_cache` files must also be deleted if this happens:
-```bash
-rm -f /raid/datasets/c4-8b/*.index_cache /raid/datasets/c4-8b/*.blend_cache
-```
-
 ## Phase 4: wandb Login
 ```bash
 wandb login
 ```
 Enter API key from https://wandb.ai/authorize

-Alternatively, pass the key directly:
-```bash
-wandb login <API_KEY>
-```
-
 ## Phase 5: Run Training

-Run training in tmux so it survives SSH disconnects:
-```bash
-tmux new-session -d -s train 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=AMD:HIP ROCM_PATH=/opt/rocm WANDB=1 bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh 2>&1 | tee /root/train.log'
-```
-Attach with `tmux attach -t train`.
-
 ### 5.1 Smoke test (beam search, 2 layers, real data)
 Always run beam first to validate the pipeline:
 ```bash
-tmux new-session -d -s beam 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=AMD:HIP ROCM_PATH=/opt/rocm bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh 2>&1 | tee /root/beam.log'
+cd /root/tinygrad
+COMGR_PATH=/opt/rocm/lib/libamd_comgr.so \
+COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so \
+CC=/opt/rocm/core-7.14/lib/llvm/bin/clang \
+DEV=AMD:HIP \
+ROCM_PATH=/opt/rocm \
+  bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
 ```

-The beam test runs 10 training steps with 2 layers. Expected results:
- ~0.29s per step after warmup
- ~700K GFLOPS, ~7% MFU (low because only 2 layers)
- ~380 GB VRAM used
- Loss stable at ~12.55 with random init
-
 ### 5.2 Full training run
 ```bash
-tmux new-session -d -s train 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=AMD:HIP ROCM_PATH=/opt/rocm WANDB=1 bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh 2>&1 | tee /root/train.log'
+cd /root/tinygrad
+COMGR_PATH=/opt/rocm/lib/libamd_comgr.so \
+COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so \
+CC=/opt/rocm/core-7.14/lib/llvm/bin/clang \
+DEV=AMD:HIP \
+ROCM_PATH=/opt/rocm \
+WANDB=1 \
+  bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh
 ```

 ## Environment Variable Reference
@@ -190,18 +180,13 @@ $ lspci -nn | grep AMD
 ```
 CPU flags include `hypervisor`. `dmesg` shows `Hypervisor detected: KVM`.

-### Working path: amdgpu driver (KFDIface)
-The amdgpu driver loads on boot and binds to all 8 GPUs, creating `/dev/kfd` and 64 renderD nodes (`/dev/dri/renderD128` through `/dev/dri/renderD191`). tinygrad's `KFDIface` enumerates GPUs through `/sys/devices/virtual/kfd/kfd/topology/nodes` and uses `/dev/kfd` for ioctl. No PCI device ID patching is needed — the KFD path does not use `PCIIface` or `AMDev._run_discovery()`.
-
-This is the working configuration. No code changes to tinygrad are required.
-
-### PCIIface path (does not work on this VM)
-For reference, the `PCIIface` path was also explored but does not work in this KVM guest:
-
- `PCIIface` in `ops_amd.py` does not list device ID `0x75b0`. Adding it allows PCI detection but `AMDev._run_discovery()` fails because the VRAM BAR reads all `0xFF`.
- This was observed with the GPU unbound from any driver, after PCI reset, and with VFIO bound.
- VFIO binding (`vfio-pci` with `enable_unsafe_noiommu_mode=1`) succeeded but VRAM BAR still reads all `0xFF`.
- No IOMMU in guest — `dmesg` has no `AMD-Vi` entries, PCI devices have no `iommu_group` symlink.
+### PCI device ID
+`lspci -v` shows device ID `0x75b0` and subsystem ID `0x75a0`:
+```
+83:00.0 Processing accelerators: ... Device 75b0
+    Subsystem: ... Device 75a0
+```
+tinygrad's `PCIIface` in `ops_amd.py` and `hive_reset.py` did not list `0x75b0`, so the GPU was not found. Adding `0x75b0` to the device ID list in both files fixes the detection.

 ### amdgpu driver behavior
 On first boot, amdgpu loaded and bound to all 8 GPUs. On one boot it failed to initialize:
@@ -213,5 +198,24 @@ On first boot, amdgpu loaded and bound to all 8 GPUs. On one boot it failed to i
 ```
 On a subsequent boot, amdgpu initialized successfully (SMU initialized, VRAM ready). After unbinding all 8 GPUs from amdgpu, `rmmod amdgpu` wedged the module (stuck in "Unloading" state in `/proc/modules`), requiring a full VM reboot.

+### `/dev/kfd`
+`/dev/kfd` exists when amdgpu is loaded. Opening it returns `OSError: [Errno 22] Invalid argument`.
+
+### VRAM BAR reads all 0xFF
+After amdgpu initializes the GPU and is then unbound, reading the VRAM BAR (via `/sys/bus/pci/devices/0000:83:00.0/resource0`) returns all `0xFF` at all offsets — including the discovery table at `vram_size - 64KB`. tinygrad's `AMDev._run_discovery()` fails with `AssertionError: discovery signatures mismatch`.
+
+A PCI reset (`echo 1 > /sys/bus/pci/devices/0000:83:00.0/reset`) did not change the VRAM contents — still all `0xFF`.
+
+VRAM was also all `0xFF` when read via `/dev/mem` at the BAR physical address (`0xa0000000000`).
+
+### VFIO attempt
+Bound the GPU to `vfio-pci` with `enable_unsafe_noiommu_mode=1`. The GPU bound successfully and `/dev/vfio/noiommu-0` appeared. Running tinygrad with `VFIO=1` still failed with the same `discovery signatures mismatch` — VRAM BAR still reads all `0xFF`.
+
+### No IOMMU in guest
+`dmesg` has no `AMD-Vi` entries. PCI devices have no `iommu_group` symlink.
+
 ### No fan control
 No `fan*` or `pwm*` hwmon entries exist. Only `temp*`, `power*`, `freq*` are exposed. GPU temps read 56-63°C, power ~265W per GPU.
+
+### Current status: NOT WORKING
+tinygrad's `PCIIface` finds the GPU (after adding `0x75b0`) but `AMDev._run_discovery()` fails because the VRAM discovery table reads all `0xFF`. This was observed with the GPU unbound from any driver, after PCI reset, and with VFIO bound.
@@ -192,7 +192,7 @@ def unpack_insts(viz_data, i:int, j:int, data:dict) -> dict:
    prev_instr = max(prev_instr, e.time + e.dur)
  summary = [{"label":"Total Cycles", "value":w.end_time-w.begin_time}, {"label":"SE", "value":w.se}, {"label":"CU", "value":w.cu},
             {"label":"SIMD", "value":w.simd}, {"label":"Wave ID", "value":w.wave_id}, {"label":"Run number", "value":data["run_number"]}]
-  return {"rows":[tuple(v.values()) for v in rows.values()], "cols":columns, "metadata":[summary],"ref":viz_data.ref_map.get(data["prg"].profile_key)}
+  return {"rows":[tuple(v.values()) for v in rows.values()], "cols":columns, "metadata":[summary], "ref":viz_data.ref_map.get(data["prg"].name)}

 def print_data(data:dict) -> None:
  from tabulate import tabulate
@@ -19,33 +19,16 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None
@functools.cache
 def custom_fused_qkv_rope_forward(q:UOp, k:UOp, v:UOp, xqkv:UOp, freqs_cis:UOp,
                                  device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
-  group_size = H // H_KV
-  q, k, v = q.reshape(B, N, H, D), k.reshape(B, N, H_KV, D), v.reshape(B, N, H_KV, D)
-  xqkv = xqkv.reshape(B, N, H_KV, group_size + 2, D)
-  b, n = UOp.range(B, 0), UOp.range(N, 1)
-  pair = UOp.range(D // 2, 2)
-  even = pair * 2
-  c = freqs_cis[0, n, 0, pair, 0].cast(dtypes.float)
-  s = freqs_cis[0, n, 0, pair, 1].cast(dtypes.float)
-  ordered:UOp|None = None
-  for kvh in range(H_KV):
-    q_out, k_out, v_out = (x.after(ordered) if ordered is not None else x for x in (q, k, v))
-    x_in = xqkv.after(ordered) if ordered is not None else xqkv
-    stores:list[UOp] = []
-    for rep in range(group_size):
-      a = x_in[b, n, kvh, rep, even].cast(dtypes.float)
-      bb = x_in[b, n, kvh, rep, even + 1].cast(dtypes.float)
-      h = kvh * group_size + rep
-      stores += [q_out[b, n, h, even].store((a * c - bb * s).cast(q.dtype)), q_out[b, n, h, even + 1].store((a * s + bb * c).cast(q.dtype))]
-    a = x_in[b, n, kvh, group_size, even].cast(dtypes.float)
-    bb = x_in[b, n, kvh, group_size, even + 1].cast(dtypes.float)
-    stores += [k_out[b, n, kvh, even].store((a * c - bb * s).cast(k.dtype)),
-               k_out[b, n, kvh, even + 1].store((a * s + bb * c).cast(k.dtype)),
-               v_out[b, n, kvh, even].store(x_in[b, n, kvh, group_size + 1, even]),
-               v_out[b, n, kvh, even + 1].store(x_in[b, n, kvh, group_size + 1, even + 1])]
-    ordered = UOp.group(*stores)
-  assert ordered is not None
-  return ordered.end(pair, n, b).sink(arg=KernelInfo(name="fused_qkv_rope_forward"))
+  code = (pathlib.Path(__file__).parent / "fused_qkv_rope.cpp").read_text()
+  threads = 256
+  thread_idx = UOp.special(threads, "lidx0")
+  block_idx_x, block_idx_y = UOp.special(B, "gidx0"), UOp.special(N, "gidx1")
+  sink = UOp.sink(q.base, k.base, v.base, xqkv.base, freqs_cis.base, thread_idx, block_idx_x, block_idx_y,
+                  arg=KernelInfo(name="fused_qkv_rope_forward"))
+  compile_args = ["-std=c++20", "-ffast-math", f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}",
+                  f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DTHREADS_PER_BLOCK={threads}"]
+  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))

@functools.cache
 def custom_fused_qkv_rope_backward(dxqkv:UOp, dq:UOp, dk:UOp, dv:UOp, freqs_cis:UOp,
@@ -126,7 +109,8 @@ def fused_qkv_rope(xqkv:Tensor, freqs_cis:Tensor, n_heads:int, n_kv_heads:int, h
 def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
  return _sharded_empty(ref.shape, ref, axis)

-def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=0):
+@functools.cache
+def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink):
  def grad(dou:UOp, ker:UOp) -> tuple:
    do = Tensor(dou, device=dou.device)
    attn = Tensor(ker.src[1].after(ker), device=ker.src[1].device)
@@ -145,7 +129,7 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
    delta_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
    delta_vec, dq = Tensor.custom_kernel(delta_vec, dq, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]

-    dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, window=window))[:3]
+    dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:3]

    if D == 64:
      dq = dq.reshape(B, H, N//16, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2).permute(0, 1, 2, 8, 9, 10, 11, 3, 4, 6, 7, 5, 12).reshape(B, H, N, D).transpose(1, 2)
@@ -165,7 +149,7 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
  return grad

 # TODO: remove write_flat once scheduler can remove reshapes between custom_kernel. TestCustomKernel.test_simple_reshape
-def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None, window:int=0):
+def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None):
  assert attn_mask is None, "attn_mask not supported"
  assert is_causal, "only causal attention supported"

@@ -192,18 +176,18 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
  attn = _sharded_empty((B, N, H * D), xq, axis=shard_axis) if write_flat else _sharded_empty_like(xq, axis=shard_axis)
  l_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)

-  grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=window)
+  grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink)

  fwd_inputs = (attn, l_vec, xq, xk, xv) + ((sinks,) if has_sink else ())
-  attn, l_vec = Tensor.custom_kernel(*fwd_inputs, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, has_sink=has_sink, window=window), grad_fxn=grad)[:2]
+  attn, l_vec = Tensor.custom_kernel(*fwd_inputs, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, has_sink=has_sink), grad_fxn=grad)[:2]

  return attn, attn, l_vec

@functools.cache
-def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True, window:int=0):
+def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True):
  code = (pathlib.Path(__file__).parent / "fa_fwd_causal.cpp").read_text()
  compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math",
-                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}", f"-DWINDOW={window}"]
+                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}"]

  Q_BLOCK_SIZE = 32
  NUM_WARPS = 8
@@ -263,10 +247,10 @@ def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arc
             src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))

@functools.cache
-def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, window:int=0):
+def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
  code = (pathlib.Path(__file__).parent / "fa_bwd_causal.cpp").read_text()
  compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math",
-                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DWINDOW={window}"]
+                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}"]

  BLOCK_SIZE_KV = 256
  GROUP_SIZE = H // H_KV
@@ -269,9 +269,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p

    qo_tile<D, float> q_reg_fl;
    load<1, qo_tile<D, float>, _gl_QKVO>(q_reg_fl, g.Qg, {batch_idx, tile_idx, head_idx, 0});
-    #if !WINDOW
    mul(q_reg_fl, q_reg_fl, TEMPERATURE_SCALE);  // Use sqrtf for clarity
-    #endif
    copy(q_reg, q_reg_fl);
    transpose(q_reg_transposed, q_reg);

@@ -290,9 +288,6 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    zero(att_block[0]);
    transpose(k_reg_transposed, k_reg);
    mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
-    #if WINDOW
-    mul(att_block[0], att_block[0], TEMPERATURE_SCALE);
-    #endif
    __builtin_amdgcn_sched_barrier(0);
    if constexpr (causal) {
        const int kv_end_pos = (min_tile + 1) * KV_BLOCK_SIZE;
@@ -342,9 +337,6 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
        zero(att_block[1]);
        transpose(k_reg_transposed, k_reg);
        mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
-        #if WINDOW
-        mul(att_block[1], att_block[1], TEMPERATURE_SCALE);
-        #endif
 #if WINDOW
        // window masks interior tiles that causal skips
        mask_kv_tile(att_block[1], tile_idx, j - 2, neg_inf_v, lane);
@@ -409,9 +401,6 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
        zero(att_block[0]);
        transpose(k_reg_transposed, k_reg);
        mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
-        #if WINDOW
-        mul(att_block[0], att_block[0], TEMPERATURE_SCALE);
-        #endif
        //      Finish softmax for QK1
        exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]);
        mul(norm_vec, norm_vec, scale_vec);
@@ -480,9 +469,6 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    zero(att_block[1]);
    transpose(k_reg_transposed, k_reg);
    mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
-    #if WINDOW
-    mul(att_block[1], att_block[1], TEMPERATURE_SCALE);
-    #endif
    //      Finish softmax for QK2
    exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]);
    mul(norm_vec, norm_vec, scale_vec);
@@ -549,9 +535,6 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    zero(att_block[0]);
    transpose(k_reg_transposed, k_reg);
    mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
-    #if WINDOW
-    mul(att_block[0], att_block[0], TEMPERATURE_SCALE);
-    #endif
    //      Finish softmax for QK3
    exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]);
    mul(norm_vec, norm_vec, scale_vec);
@@ -614,9 +597,6 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    zero(att_block[1]);
    transpose(k_reg_transposed, k_reg);
    mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
-    #if WINDOW
-    mul(att_block[1], att_block[1], TEMPERATURE_SCALE);
-    #endif
    //      Finish softmax for QK4
    exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]);
    mul(norm_vec, norm_vec, scale_vec);
@@ -0,0 +1,69 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+
+#ifndef ATTN_B
+#define ATTN_B 2
+#endif
+#ifndef ATTN_N
+#define ATTN_N 8192
+#endif
+#ifndef ATTN_H
+#define ATTN_H 32
+#endif
+#ifndef ATTN_H_KV
+#define ATTN_H_KV 8
+#endif
+#ifndef ATTN_D
+#define ATTN_D 128
+#endif
+#ifndef THREADS_PER_BLOCK
+#define THREADS_PER_BLOCK 256
+#endif
+
+constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV;
+constexpr int HALF_D = ATTN_D / 2;
+constexpr int PACKED_D = (GROUP_SIZE + 2) * ATTN_D;
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_BLOCK) void
+fused_qkv_rope_forward(
+    __hip_bfloat16*       __restrict__ q,
+    __hip_bfloat16*       __restrict__ k,
+    __hip_bfloat16*       __restrict__ v,
+    const __hip_bfloat16* __restrict__ xqkv,
+    const __hip_bfloat16* __restrict__ freqs_cis) {
+  const int b = blockIdx.x;
+  const int n = blockIdx.y;
+  const int bn = b * ATTN_N + n;
+  const int packed_bn = bn * ATTN_H_KV * PACKED_D;
+  const int q_bn = bn * ATTN_H * ATTN_D;
+  const int kv_bn = bn * ATTN_H_KV * ATTN_D;
+
+  if (threadIdx.x < HALF_D) {
+    const int pair = threadIdx.x;
+    const int even = pair << 1;
+    const float c = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 0]);
+    const float s = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 1]);
+
+    for (int kvh = 0; kvh < ATTN_H_KV; kvh++) {
+      const int base = packed_bn + kvh * PACKED_D;
+
+      for (int rep = 0; rep < GROUP_SIZE; rep++) {
+        const int qbase = base + rep * ATTN_D;
+        const int h = kvh * GROUP_SIZE + rep;
+        const float a = static_cast<float>(xqkv[qbase + even]);
+        const float bb = static_cast<float>(xqkv[qbase + even + 1]);
+        const int out = q_bn + h * ATTN_D + even;
+        q[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
+        q[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
+      }
+
+      const float a = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even]);
+      const float bb = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even + 1]);
+      const int out = kv_bn + kvh * ATTN_D + even;
+      k[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
+      k[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
+      v[out] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even];
+      v[out + 1] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even + 1];
+    }
+  }
+}
@@ -209,7 +209,7 @@ class ST:
    return cls(uop, rows, cols, layout, base_shape, ker)

  def swizzle(self, row, col):
-    swizzled_offset = self.base_shape.swizzle(row, col, self._uop.dtype)
+    swizzled_offset = self.base_shape.swizzle(row, col, self._uop.dtype.scalar())

    row = swizzled_offset // self.base_shape.cols
    col = swizzled_offset % self.base_shape.cols
@@ -4,7 +4,7 @@
 # A006 Lambda argument `input` is shadowing a Python builtin
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.uop.ops import Ops, GroupOp
-from tinygrad.helpers import getenv, prod, strides_for_shape
+from tinygrad.helpers import getenv, prod, strides_for_shape, argfix
 import torch.lib
 TORCH_DEBUG = getenv("TORCH_DEBUG")
 import torch, pathlib, operator, functools, weakref
@@ -73,12 +73,6 @@ def wrap_view_op(fn):
    return wrap(ret)
  return _wrap

-# NOTE: list assignment raises IndexError on an out of range dim, and the index must be a tuple: a list of all ints is one advanced index
-def _index_dim(self, dim, idx):
-  idxs = [slice(None)] * self.ndim
-  idxs[dim] = idx
-  return self[tuple(idxs)]
-
 view_ops = {
  "aten.view": Tensor.reshape,
  "aten._unsafe_view": Tensor.reshape,  # when are views unsafe, and do we care?
@@ -88,13 +82,15 @@ view_ops = {
  "aten.transpose.int": Tensor.transpose,
  "aten.squeeze.dim": Tensor.squeeze,
  "aten.unsqueeze": Tensor.unsqueeze,
-  "aten.select.int": _index_dim,
+  "aten.select.int": lambda self, dim, idx: self[(slice(None),) * (dim%self.ndim) + (idx,)],
  "aten.permute": Tensor.permute,
  "aten.alias": lambda self: self,
  "aten.diagonal": Tensor.diagonal,
-  "aten.slice.Tensor": lambda self, dim=0, start=None, end=None, step=1: _index_dim(self, dim, slice(start, end, step)),
  }

+# torch 2.10 handles this natively
+if tuple(map(int, torch.__version__.split('.')[:2])) < (2, 10): view_ops.update({"aten.detach": Tensor.detach})
+
 for k,v in view_ops.items(): torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrap_view_op(v))

 def _get_view_ops(view): return getattr(view, "_view_ops", [])
@@ -103,21 +99,46 @@ def _apply_view_ops(target, ops):
  for fn, args, kwargs in ops: target = fn(target, *args, **kwargs)
  return target

-# a chain of reshapes is undone by reshaping the value back to the base
+# similar to https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/InferSize.h
+def _reshape_target_shape(shape:tuple[int, ...], args) -> tuple[int, ...]|None:
+  if not (req := argfix(*args)): return None
+  new_shape, infer_idx = [], -1
+  for i, s in enumerate(req):
+    if s is None: s = shape[i] if i < len(shape) else None
+    if not isinstance(s, int): return None
+    if s == -1:
+      if infer_idx != -1: return None
+      infer_idx = len(new_shape)
+    new_shape.append(s)
+  total = prod(shape)
+  if infer_idx != -1:
+    known = prod(x for x in new_shape if x != -1)
+    if known == 0:
+      if total != 0: return None
+      new_shape[infer_idx] = 0
+    else: new_shape[infer_idx] = total // known
+  return tuple(new_shape) if prod(new_shape) == total else None
+
+# TODO: can we get rid of this? only for test_flatten_reshape_add
 def _try_simple_reshape_view_write(base: Tensor, view: Tensor, val: Tensor) -> bool:
  if not (ops := _get_view_ops(view)): return False
-  if any(fn is not Tensor.reshape for fn, _, _ in ops): return False
-  base.assign(val.reshape(base.shape))
+  shapes = [base.shape]
+  for fn, args, _ in ops:
+    if fn is Tensor.reshape:
+      if not (next_shape := _reshape_target_shape(shapes[-1], args)): return False
+      shapes.append(next_shape)
+  if shapes[-1] != view.shape: return False
+  for s in reversed(shapes[:-1]): val = val.reshape(s)
+  base.assign(val)
  return True

 def _view_write(base: Tensor, view: Tensor, value: Tensor) -> None:
  val = value if value.dtype == base.dtype else value.cast(base.dtype)
+  if view.shape == base.shape: return base.assign(val)
  if _try_simple_reshape_view_write(base, view, val): return
  idx_base = Tensor.arange(base.numel(), dtype=dtypes.int32).reshape(base.shape)
  idx_view = _apply_view_ops(idx_base, _get_view_ops(view)).reshape(-1)
-  # clone, not contiguous: contiguous() on a base that already owns its buffer returns the base itself, and scattering
-  # into that is an in-place write to a buffer other tensors still hold, which setitem refuses
-  flat_base = base.reshape(base.numel()).clone()
+  flat_base = base.reshape(base.numel()).contiguous()
  flat_base[idx_view] = val.reshape(-1)
  base.assign(flat_base.reshape(base.shape))

@@ -145,6 +166,11 @@ def _index_put_impl_(self, indices, values, accumulate=False, unsafe=False):
 def index_put(self, indices, values, accumulate=False):
  return aten.index_put(self.cpu(), [z.cpu() if isinstance(z, torch.Tensor) else None for z in indices], values.clone().cpu(), accumulate).tiny()

+@torch.library.impl("aten::isin.Tensor_Tensor_out", "privateuseone")
+def isin_tensor_tensor_out(x, y, *, assume_unique=False, invert=False, out=None):
+  result = (unwrap(x).unsqueeze(-1) == unwrap(y).flatten()).any(-1)
+  return out.copy_(wrap(~result if invert else result))
+
@torch.library.impl("aten::randperm.generator_out", "privateuseone")
 def randperm_generator(n, generator=None, out=None):
  if generator is not None: raise NotImplementedError("tinygrad torch backend does not support torch.Generator for randperm")
@@ -205,6 +231,49 @@ def as_strided(tensor:torch.Tensor, size, stride, storage_offset=None):
 def _reshape_alias(tensor:torch.Tensor, size, stride):
  return _as_strided(tensor, size, stride)

+@torch.library.impl("aten::empty_strided", "privateuseone")
+def empty_strided(size, stride, dtype=None, layout=None, device=None, pin_memory=False):
+  if TORCH_DEBUG: print(f"empty_strided {size=} {stride=} {dtype=} {layout=} {device=} {pin_memory=}")
+  ret = Tensor.empty(*size, dtype=_from_torch_dtype(dtype or torch.get_default_dtype()), device=_from_torch_device(device))
+  # TODO: should return with requested strides
+  return wrap(ret)
+
+@torch.library.impl("aten::empty.memory_format", "privateuseone")
+def empty_memory_format(size, dtype=None, layout=None, device=None, pin_memory=False, memory_format=None):
+  if TORCH_DEBUG: print(f"empty.memory_format {size=} {dtype=} {layout=} {device=} {pin_memory=} {memory_format=}")
+  ret = Tensor.empty(*size, dtype=_from_torch_dtype(dtype or torch.get_default_dtype()), device=_from_torch_device(device))
+  return wrap(ret)
+
+@torch.library.impl("aten::max_pool2d_with_indices", "privateuseone")
+def max_pool2d_with_indices(self:torch.Tensor, kernel_size:tuple[int, ...], stride=None, padding=0, dilation=1, ceil_mode=False):
+  # TODO: supprt stride [] in tinygrad?
+  if stride is not None and len(stride) == 0: stride = None
+  ret, idx = unwrap(self).max_pool2d(kernel_size, stride, dilation, padding, ceil_mode, return_indices=True)
+  return (wrap(ret), wrap(idx.cast(dtypes.int64)))
+
+@torch.library.impl("aten::max_pool2d_with_indices_backward", "privateuseone")
+def max_pool2d_with_indices_backward(grad_out:torch.Tensor, self:torch.Tensor, kernel_size:tuple[int, ...], stride=None, padding=0, dilation=1, ceil_mode=False, indices=None):
+  return wrap(Tensor.max_unpool2d(unwrap(grad_out), unwrap(indices), output_size=unwrap(self).shape))
+
+@torch.library.impl("aten::max_unpool2d", "privateuseone")
+def max_unpool2d(self:torch.Tensor, indices:torch.Tensor, output_size):
+  return wrap(unwrap(self).max_unpool2d(unwrap(indices), output_size=output_size))
+
+@torch.library.impl("aten::arange", "privateuseone")
+def arange(end, dtype=None, device=None, pin_memory=None):
+  has_float = isinstance(end, float)
+  return wrap(Tensor.arange(0, end, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))
+
+@torch.library.impl("aten::arange.start", "privateuseone")
+def arange_start(start, end, dtype=None, device=None, pin_memory=None):
+  has_float = any(isinstance(x, float) for x in (start, end))
+  return wrap(Tensor.arange(start, end, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))
+
+@torch.library.impl("aten::arange.start_step", "privateuseone")
+def arange_start_step(start, end, step, dtype=None, device=None, pin_memory=None):
+  has_float = any(isinstance(x, float) for x in (start, end, step))
+  return wrap(Tensor.arange(start, end, step, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))
+
@torch.library.impl("aten::convolution_overrideable", "privateuseone")
 def convolution_overrideable(input, weight, bias, stride, padding, dilation, transposed, output_padding, groups):
  if TORCH_DEBUG >= 1:
@@ -225,27 +294,12 @@ def convolution_backward_overrideable(grad_out, input, weight, stride, padding,
  grads = out.gradient(*[t for t,m in zip([input, weight, bias], output_mask) if m], gradient=grad_out)
  return tuple([wrap(grads.pop(0)) if m else None for m in output_mask])

-# the functional scatters. without an impl aten falls back to a path that assumes a real storage: "self.has_storage() INTERNAL ASSERT FAILED"
-def _scatter_into(self, src, dim, index):
-  out = unwrap(self).clone()
-  slices = [slice(None)] * out.ndim
-  slices[dim] = index
-  out[slices] = unwrap(src).cast(out.dtype)  # torch casts src to self's dtype, tinygrad setitem demands they already match
-  return wrap(out)
-
-@torch.library.impl("aten::slice_scatter", "privateuseone")
-def slice_scatter(self, src, dim=0, start=None, end=None, step=1): return _scatter_into(self, src, dim, slice(start, end, step))
-
-@torch.library.impl("aten::select_scatter", "privateuseone")
-def select_scatter(self, src, dim, index): return _scatter_into(self, src, dim, index)
-
-@torch.library.impl("aten::diagonal_scatter", "privateuseone")
-def diagonal_scatter(self, src, offset=0, dim1=0, dim2=1):
-  # a diagonal is not one axis, so scatter through the flat indices it picks out
-  base, out = unwrap(self), unwrap(self).clone().reshape(-1)
-  idx = Tensor.arange(base.numel(), dtype=dtypes.int32).reshape(base.shape).diagonal(offset, dim1, dim2).reshape(-1)
-  out[idx] = unwrap(src).cast(base.dtype).reshape(-1)
-  return wrap(out.reshape(base.shape))
+@torch.library.impl("aten::slice.Tensor", "privateuseone")
+@wrap_view_op
+def slice_tensor(self, dim=0, start=None, end=None, step=1):
+  slices = [slice(None)] * self.ndim
+  slices[dim] = slice(start, end, step)
+  return self[slices]

@torch.library.impl("aten::slice_backward", "privateuseone")
 def slice_backward(grad_out, input_sizes, dim, start, end, step):
@@ -287,14 +341,19 @@ for dim in [1, 2, 3]:
    torch.library.impl(f"aten::{pad_type}_pad{dim}d", "privateuseone")(functools.partial(pad_forward, mode=mode))
    torch.library.impl(f"aten::{pad_type}_pad{dim}d_backward", "privateuseone")(functools.partial(pad_backward, mode=mode))

-# the schemas are all positional: (self, output_size, align_corners, *scales) for linear, (self, output_size, *scales) for nearest.
-def upsample(self, size, *args, mode=None):
-  return wrap(Tensor.interpolate(unwrap(self), size, mode=mode, align_corners=args[0] if mode == "linear" else False))
+def upsample(self, size, align_corners=False, mode=None): return wrap(Tensor.interpolate(unwrap(self), size, mode=mode, align_corners=align_corners))
 for i,pre in enumerate(["", "bi", "tri"]):
  torch.library.impl(f"aten::upsample_{pre}linear{i+1}d", "privateuseone")(functools.partial(upsample, mode="linear"))
  torch.library.impl(f"aten::upsample_nearest{i+1}d", "privateuseone")(functools.partial(upsample, mode="nearest"))
  torch.library.impl(f"aten::_upsample_nearest_exact{i+1}d", "privateuseone")(functools.partial(upsample, mode="nearest-exact"))

+@torch.library.impl("aten::scatter_add.out", "privateuseone")
+def scatter_add(self, dim, index, src, out):
+  self, index, src, out_unwrapped = unwrap(self), unwrap(index), unwrap(src), unwrap(out)
+  if self.shape == (): _apply_inplace(out_unwrapped, src)
+  else: _apply_inplace(out_unwrapped, Tensor.scatter_reduce(self, dim, index, src, reduce='sum'))
+  return out
+
 def _copy_between_devices(src, dest, cast_dtype, to_device, non_blocking=False):
  if src.is_tiny and dest.is_tiny:
    src_t, dest_t = unwrap(src), unwrap(dest)
@@ -345,11 +404,15 @@ def sort_values(input, dim=-1, descending=False, stable=True, values=None, indic
  _apply_inplace(unwrap(indices), out_indices.cast(dtypes.int64))
  return values, indices

+@torch.library.impl("aten::_linalg_svd", "privateuseone")
+def _linalg_svd(self, full_matrices=False):
+  U, S, Vh = unwrap(self).svd(full_matrices)
+  return wrap(U), wrap(S), wrap(Vh)
+
 # register some decompositions
 from torch._decomp import get_decompositions
 decomps = [
  aten.native_layer_norm_backward,
-  aten.native_group_norm_backward,
  aten.linalg_cross,
  aten.addmm,
  aten.addcmul,
@@ -384,20 +447,12 @@ decomps = [
  aten._softmax_backward_data, aten.embedding_dense_backward,
  aten.linalg_vector_norm,
  aten.binary_cross_entropy, aten.binary_cross_entropy_backward,
-  # the C++ mse/smooth_l1 kernels resize their out tensor, and a tiny tensor has no storage to resize
-  aten.mse_loss, aten.mse_loss_backward,
-  aten.smooth_l1_loss, aten.smooth_l1_loss_backward,
  aten.upsample_nearest2d.out,
-  # NOTE: only the "out" overload, the "vec" one is CompositeImplicitAutograd and overriding it loses the autograd kernel
-  aten.upsample_bicubic2d.out,
-  aten._adaptive_avg_pool2d,
  # activations
  aten.hardswish, aten.hardswish_backward,
  aten.hardtanh, aten.hardtanh_backward,
  aten.gelu, aten.gelu_backward,
-  # NOTE: no aten.logical_or here, its decomposition reaches aten.bitwise_or through a path that checks aliasing by
-  # reading storage, which a tiny tensor has none of. it gets a direct impl below instead
-  aten.logical_and, aten.logical_xor,
+  aten.logical_and,
  aten.randint,
  aten.eye,
  aten.hardsigmoid_backward,
@@ -440,7 +495,7 @@ simple_tensor_methods = [
  # reduce
  "all", "any", "argmax", "argmin", "cumsum", "cumprod",
  # complex
-  "linspace"]
+  "avg_pool2d", "linspace"]

 tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_methods}, **{
  "aten.add.out": lambda input,other,alpha=1: input+alpha*other,
@@ -485,8 +540,6 @@ tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_
  "aten.where.self_out": Tensor.where,
  "aten.prod.int_out": Tensor.prod,
  "aten.scatter.src_out": Tensor.scatter,
-  "aten.scatter_add.out": lambda self,dim,index,src: src if self.shape == () else Tensor.scatter_reduce(self, dim, index, src, reduce="sum"),
-  "aten.isin.Tensor_Tensor_out": lambda x,y,assume_unique=False,invert=False: (x.unsqueeze(-1)==y.flatten()).any(-1) != invert,
  # NOTE: axis=[] in torch means all, change tinygrad?
  "aten.sum.IntList_out": lambda self,axis,keepdim=False,dtype=None:
    self.sum(axis if axis is None or len(axis) else None, keepdim,
@@ -502,9 +555,10 @@ def wrap_out(f):
    assert out.shape == assigned.shape, f"shape mismatch: {assigned.shape} -> {out.shape}"
    assert out.device == assigned.device or out.device is None or assigned.device is None, f"device mismatch: {assigned.device} -> {out.device}"
    assert out.dtype == assigned.dtype, f"dtype mismatch: {assigned.dtype} -> {out.dtype}"
-    # writing out= is an in-place write like any other: through the base if it is a view, refreshing any derived views
-    _apply_inplace(out, assigned)
-    return out
+    # an out= that is a view has to be written through its base, and _apply_inplace gives a deviceless base its buffer first
+    if canonical_base(out) is not out: return _apply_inplace(out, assigned) or out
+    if out.device is None and assigned.device is not None: out.replace(out.empty_like(device=assigned.device))
+    return out.assign(assigned)
  return _wrap_out

 def _inplace_op(t, new_value):
@@ -512,14 +566,7 @@ def _inplace_op(t, new_value):
  else: _apply_inplace(t, new_value)
  return t

-# the three arange overloads are one function at different arity, and dtype/layout/device/pin_memory are keyword only in all of them
-def _arange(*args, dtype=None, **_):
-  return Tensor.arange(*args, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if any(isinstance(x, float) for x in args) else torch.int64)))
-
-def _empty(size, dtype=None, device=None, **_):
-  return Tensor.empty(*size, dtype=_from_torch_dtype(dtype or torch.get_default_dtype()), device=_from_torch_device(device))
-
-tiny_backend = {**tiny_backend_out, **{
+tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
  "aten.remainder.Scalar_Tensor": lambda x,y: x%y,
  "aten.floor_divide": lambda x,y: x//y,
  "aten.floor_divide_.Tensor": lambda x,y: x//y,
@@ -532,8 +579,8 @@ tiny_backend = {**tiny_backend_out, **{
  # inplace ops using replace for fusion
  "aten.zero_": lambda x: x.const_like(0),
  "aten.fill_.Scalar": lambda x, y: x.const_like(y),
-  "aten.add_.Tensor": lambda self, other, alpha=1: self + other * alpha,
-  "aten.add_.Scalar": lambda self, other, alpha=1: self + other * alpha,
+  "aten.add_.Tensor": lambda self, other, alpha=1.0: self + other * alpha,
+  "aten.add_.Scalar": lambda self, other, alpha=1.0: self + other * alpha,
  "aten.mul_.Tensor": lambda self, other: self * other,
  "aten.mul_.Scalar": lambda self, other: self * other,
  # relu doesn't have an out form?
@@ -566,9 +613,7 @@ tiny_backend = {**tiny_backend_out, **{
  # these don't work in out form, they have size 0
  "aten.abs": Tensor.abs,
  "aten.logical_not": Tensor.logical_not,
-  # compare against zero first: logical_* is bool-valued for any input dtype, while | is bitwise
-  "aten.logical_or": lambda x, y: (x != 0) | (y != 0),
-  "aten.logical_or_": lambda x, y: (x != 0) | (y != 0),
+  "aten.logical_or_": lambda x, y: x | y,
  "aten.multinomial": Tensor.multinomial,
  "aten.masked_fill_.Scalar": lambda self, mask, value: self.masked_fill(mask, value),
  "aten.masked_fill_.Tensor": lambda self, mask, value: self.masked_fill(mask, value),
@@ -577,7 +622,14 @@ tiny_backend = {**tiny_backend_out, **{
  "aten.masked_select": Tensor.masked_select,
  "aten.all": Tensor.all,
  "aten.sgn": Tensor.sign,
+  "aten.acos": Tensor.acos,
  "aten.any": Tensor.any,
+  "aten.bitwise_not": Tensor.bitwise_not,
+  "aten.argmax": Tensor.argmax,
+  "aten.argmin": Tensor.argmin,
+  "aten.asinh": Tensor.asinh,
+  "aten.mul": Tensor.mul,
+  "aten.atanh": Tensor.atanh,
  "aten.fill_.Tensor": lambda self, value: self.const_like(value.reshape(()).item()),
  "aten.flip": Tensor.flip,
  "aten.scatter_reduce.two": Tensor.scatter_reduce,
@@ -588,22 +640,10 @@ tiny_backend = {**tiny_backend_out, **{
  "aten.add.Tensor": lambda input,other,alpha=1: input+alpha*other,
  "aten.linspace": lambda start, stop, steps, dtype=None, **kwargs:
    Tensor.linspace(start, stop, steps, **({"dtype": _from_torch_dtype(dtype)} if dtype is not None else {})),
-  # the functional copy_. without an impl the fallback segfaults on a tensor with no storage
-  "aten.copy": lambda self,src,non_blocking=False: src.cast(self.dtype).to(self.device).expand(self.shape),
-  "aten.arange": lambda end, **kwargs: _arange(0, end, **kwargs),
-  "aten.arange.start": _arange,
-  "aten.arange.start_step": _arange,
-  # empty_strided takes the strides and drops them: we always allocate contiguous
-  "aten.empty_strided": lambda size, stride, **kwargs: _empty(size, **kwargs),
-  "aten.empty.memory_format": _empty,
-  # TODO: supprt stride [] in tinygrad?
-  "aten.max_pool2d_with_indices": lambda self,kernel_size,stride=None,padding=0,dilation=1,ceil_mode=False: ((r:=Tensor.max_pool2d(self, kernel_size, stride or None, dilation, padding, ceil_mode, return_indices=True))[0], r[1].cast(dtypes.int64)),
-  "aten.max_pool2d_with_indices_backward": lambda grad_out,self,kernel_size,stride=None,padding=0,dilation=1,ceil_mode=False,indices=None: Tensor.max_unpool2d(grad_out, indices, output_size=self.shape),
-  "aten.max_unpool2d": lambda self,indices,output_size: Tensor.max_unpool2d(self, indices, output_size=output_size),
-  "aten._linalg_svd": lambda self,full_matrices=False: Tensor.svd(self, full_matrices),
  "aten.topk": Tensor.topk,
  "aten.constant_pad_nd": lambda self, padding, value=0.0: self.pad(padding, mode="constant", value=value).contiguous(),
-  "aten.cumsum": lambda self, dim: self.cumsum(dim),
+  # TODO: input contiguous is needed to prevent CFGContext circular dependency assertion for shapes >512 (see test_cumsum_arange_large)
+  "aten.cumsum": lambda self, dim: self.contiguous().cumsum(dim),
  "aten.logsumexp": lambda self, axis, keepdim=False: self.logsumexp(axis[0], keepdim=keepdim),
  "aten.roll": Tensor.roll,
  "aten.logcumsumexp": Tensor.logcumsumexp,
@@ -612,7 +652,6 @@ tiny_backend = {**tiny_backend_out, **{
    self.ones_like(**{k: v for k, v in {"dtype": _from_torch_dtype(dtype) if dtype else None,
                                        "device": _from_torch_device(device) if device else None}.items() if v is not None}),
  "aten.max.dim": lambda self, dim, keepdim=False: (self.max(dim, keepdim), self.argmax(dim, keepdim).cast(dtype=dtypes.int64)),
-  "aten.min.dim": lambda self, dim, keepdim=False: (self.min(dim, keepdim), self.argmin(dim, keepdim).cast(dtype=dtypes.int64)),
  "aten.cummax": lambda self, dim: ((r := self.cummax(dim))[0], r[1].cast(dtypes.int64)),
  "aten.cummin": lambda self, dim: ((r := self.cummin(dim))[0], r[1].cast(dtypes.int64)),
  "aten.nonzero": Tensor.nonzero,
@@ -674,16 +713,15 @@ def wrap_inplace_view_op(f):
  return nf

 # the aten schema says how an op is called: an inplace view retargets the view, a writable first arg is inplace,
-# and a writable out arg gets wrap_out's dtype cast, shape assert, and view write-through
+# and a writable out arg must have come from tiny_backend_out so that wrap_out was applied
 for k,v in tiny_backend.items():
  name, _, overload = k.removeprefix("aten.").partition(".")
  op = getattr(getattr(aten, name), overload or "default")
  writes = [a.name for a in op._schema.arguments if a.alias_info is not None and a.alias_info.is_write]
  if torch.Tag.inplace_view in op.tags: fxn = wrap_inplace_view_op(v)
  elif writes == [op._schema.arguments[0].name] and op._schema.returns: fxn = wrap_inplace(v)
-  elif not writes: fxn = wrap_fxn(k, v)
-  elif writes == ["out"]: fxn = wrap_fxn(k, wrap_out(v))
-  else: raise RuntimeError(f"{k} writes {writes}: unhandled writable arg in schema")
+  elif not writes or (writes == ["out"] and k in tiny_backend_out): fxn = wrap_fxn(k, v)
+  else: raise RuntimeError(f"{k} writes {writes}: expected an inplace first arg, or an out arg with {k} in tiny_backend_out")
  torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(fxn)

@torch.library.impl("aten::equal", "privateuseone")
@@ -83,12 +83,6 @@ class TestTorchBackend(unittest.TestCase):
    torch.add(torch.ones(5, device=device), torch.ones(5, device=device), out=a)
    self.assertEqual(a.detach().storage_offset(), 3)

-  def test_out_refreshes_views_of_base(self):
-    a = torch.zeros(4, device=device)
-    v = a[2:]
-    torch.add(torch.ones(4, device=device), torch.ones(4, device=device), out=a)
-    np.testing.assert_equal(v.cpu().numpy(), [2., 2.])
-
  @unittest.expectedFailure  # TODO: storage offset assumes a contiguous source, use UOp.contiguous_view_offset
  def test_storage_offset_non_contiguous_source(self):
    a = torch.arange(12., device=device).reshape(3,4)
@@ -172,15 +166,6 @@ class TestTorchBackend(unittest.TestCase):
    expected = np.array([[1.5, 5.2, 9.0], [13.2, 17.1, 18.4]], dtype=np.float32)
    np.testing.assert_equal(y3.cpu().numpy(), expected)

-  def test_argmax_argmin(self):
-    a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
-    c = a.cpu()
-    for got, want in [(a.argmax(), c.argmax()), (a.argmin(0), c.argmin(0)), (a.argmax(1, keepdim=True), c.argmax(1, keepdim=True)),
-                      (torch.min(a, 1).indices, torch.min(c, 1).indices), (torch.max(a, 1).indices, torch.max(c, 1).indices),
-                      (torch.min(a, 1).values, torch.min(c, 1).values), (torch.min(a, 1, keepdim=True).indices, torch.min(c, 1, keepdim=True).indices)]:
-      self.assertEqual(got.dtype, want.dtype)  # torch's arg reduces are int64, tinygrad's are int32
-      np.testing.assert_equal(got.cpu().numpy(), want.numpy())
-
  def test_isfinite(self):
    a = torch.ones(4, device=device)
    np.testing.assert_equal(torch.isfinite(a).cpu().numpy(), [True, True, True, True])
@@ -388,22 +373,6 @@ class TestTorchBackend(unittest.TestCase):
    for bwd_eps in [1e-5, 0.3]:
      for got, want in zip(run(device, bwd_eps), run("cpu", bwd_eps)): np.testing.assert_allclose(got, want, atol=1e-4, rtol=1e-3)

-  def test_groupnorm_backward(self):
-    def run(dev):
-      x = torch.arange(24., device=dev).reshape(2, 4, 3).requires_grad_()
-      w = torch.linspace(0.5, 2.0, 4).to(dev).requires_grad_()
-      torch.nn.functional.group_norm(x, 2, w, torch.zeros(4, device=dev)).square().sum().backward()
-      return x.grad.cpu().numpy(), w.grad.cpu().numpy()
-    for got, want in zip(run(device), run("cpu")): np.testing.assert_allclose(got, want, atol=1e-4, rtol=1e-3)
-
-  def test_mse_smooth_l1_loss_backward(self):
-    def run(dev, loss):
-      x = torch.arange(4., device=dev).requires_grad_()
-      loss(x, torch.ones(4, device=dev)).backward()
-      return x.grad.cpu().numpy()
-    for loss in [torch.nn.functional.mse_loss, torch.nn.functional.smooth_l1_loss]:
-      np.testing.assert_allclose(run(device, loss), run("cpu", loss), atol=1e-6)
-
  def test_batchnorm_unsqueeze(self):
    bn = torch.nn.BatchNorm2d(4).to(device)
    x = torch.randn(8, 4, 3, 3, device=device)
@@ -547,15 +516,6 @@ class TestTorchBackend(unittest.TestCase):
    cpu_res = torch.arange(20, dtype=torch.float32)[::2][1:4].numpy()
    np.testing.assert_equal(torch_res, cpu_res)

-  def test_select_out_of_range_dim(self):
-    a = torch.arange(12, dtype=torch.int32, device=device).reshape(3, 4)
-    with self.assertRaises(IndexError): a.select(5, 0)
-
-  def test_select_collapses_the_only_dim(self):
-    a = torch.arange(3, dtype=torch.int32, device=device)
-    self.assertEqual(a.select(0, 1).shape, ())
-    np.testing.assert_equal(a.select(0, 1).cpu().numpy(), 1)
-
  def test_slice_negative_dim(self):
    a = torch.arange(13, dtype=torch.int32, device=device).repeat(8, 1)
    torch_chunks = a.chunk(3, -1)
@@ -836,86 +796,6 @@ class TestTorchBackend(unittest.TestCase):
    np.testing.assert_allclose(w_tiny.grad.cpu().numpy(), w_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
    np.testing.assert_allclose(b_tiny.grad.cpu().numpy(), b_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)

-  def test_write_through_detach_of_unrealized(self):
-    a = torch.empty(4, device=device)
-    a.detach().fill_(3)
-    np.testing.assert_equal(a.cpu().numpy(), [3, 3, 3, 3])
-
-  def test_square_transpose_inplace(self):
-    # a same-shape transpose is not a reshape: writing the transposed values straight back would scramble the base
-    a = torch.tensor([[0., 1., 2.], [3., 4., 5.], [6., 7., 8.]], device=device)
-    a.transpose(0, 1).add_(100)
-    np.testing.assert_equal(a.cpu().numpy(), [[100., 101., 102.], [103., 104., 105.], [106., 107., 108.]])
-
-  def test_interpolate(self):
-    a = torch.arange(4, dtype=torch.float32, device=device).reshape(1, 1, 2, 2)
-    nearest = torch.nn.functional.interpolate(a, scale_factor=2.0)
-    np.testing.assert_equal(nearest.cpu().numpy()[0, 0], [[0, 0, 1, 1], [0, 0, 1, 1], [2, 2, 3, 3], [2, 2, 3, 3]])
-    linear = torch.nn.functional.interpolate(a, size=(4, 4), mode="bilinear", align_corners=False)
-    ref = torch.nn.functional.interpolate(a.cpu(), size=(4, 4), mode="bilinear", align_corners=False)
-    np.testing.assert_allclose(linear.cpu().numpy(), ref.numpy(), rtol=1e-5)
-
-  def test_interpolate_bicubic_area(self):
-    a = torch.arange(32, dtype=torch.float32, device=device).reshape(1, 2, 4, 4)
-    for mode, scale in [("bicubic", 2.0), ("area", 0.5)]:
-      ref = torch.nn.functional.interpolate(a.cpu(), scale_factor=scale, mode=mode)
-      np.testing.assert_allclose(torch.nn.functional.interpolate(a, scale_factor=scale, mode=mode).cpu().numpy(), ref.numpy(), atol=1e-4)
-
-  @unittest.expectedFailure
-  def test_interpolate_bicubic_backward(self):
-    # the forward comes from a decomposition, but aten::upsample_bicubic2d_backward has none (nor does
-    # aten::_adaptive_avg_pool2d_backward, for area), so training through these modes needs a real kernel
-    x = torch.arange(32., dtype=torch.float32, device=device).reshape(1, 2, 4, 4).requires_grad_()
-    torch.nn.functional.interpolate(x, scale_factor=2.0, mode="bicubic").sum().backward()
-
-  @unittest.expectedFailure
-  def test_interpolate_inexact_scale(self):
-    # torch forwards the raw scale_factor, Tensor.interpolate recomputes it from output_size, and they disagree here
-    a = torch.arange(6, dtype=torch.float32, device=device).reshape(1, 1, 2, 3)
-    tiny = torch.nn.functional.interpolate(a, scale_factor=2.5, mode="bilinear")
-    ref = torch.nn.functional.interpolate(a.cpu(), scale_factor=2.5, mode="bilinear")
-    np.testing.assert_allclose(tiny.cpu().numpy(), ref.numpy(), rtol=1e-5)
-
-  def test_logical_or_xor(self):
-    a = torch.tensor([True, True, False, False], device=device)
-    b = torch.tensor([True, False, True, False], device=device)
-    np.testing.assert_equal(torch.logical_or(a, b).cpu().numpy(), [True, True, True, False])
-    np.testing.assert_equal(torch.logical_xor(a, b).cpu().numpy(), [False, True, True, False])
-    # bool-valued whatever the input dtype, so this is not | and ^
-    i, j = torch.tensor([2, 0, 5, 0], device=device), torch.tensor([0, 0, 1, 1], device=device)
-    np.testing.assert_equal(torch.logical_or(i, j).cpu().numpy(), [True, False, True, True])
-    np.testing.assert_equal(torch.logical_xor(i, j).cpu().numpy(), [True, False, False, True])
-
-  def test_slice_scatter(self):
-    # the scatters are functional: they return a new tensor and must leave the one they were given alone
-    a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
-    out = torch.slice_scatter(a, torch.ones(1, 4, device=device), 0, 0, 1)
-    np.testing.assert_equal(out.cpu().numpy(), [[1, 1, 1, 1], [4, 5, 6, 7], [8, 9, 10, 11]])
-    np.testing.assert_equal(a.cpu().numpy(), np.arange(12, dtype=np.float32).reshape(3, 4))
-
-  def test_slice_scatter_casts_src(self):
-    a = torch.zeros(3, 4, device=device)
-    out = torch.slice_scatter(a, torch.ones(1, 4, dtype=torch.int32, device=device), 0, 0, 1)
-    self.assertEqual(out.dtype, torch.float32)
-    np.testing.assert_equal(out.cpu().numpy()[0], np.ones(4, dtype=np.float32))
-
-  def test_select_scatter(self):
-    a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
-    out = torch.select_scatter(a, torch.ones(4, device=device), 0, 1)
-    np.testing.assert_equal(out.cpu().numpy(), [[0, 1, 2, 3], [1, 1, 1, 1], [8, 9, 10, 11]])
-
-  def test_diagonal_scatter(self):
-    a = torch.zeros(3, 3, device=device)
-    out = torch.diagonal_scatter(a, torch.arange(3, dtype=torch.float32, device=device))
-    np.testing.assert_equal(out.cpu().numpy(), np.diag([0., 1., 2.]))
-    np.testing.assert_equal(a.cpu().numpy(), np.zeros((3, 3), dtype=np.float32))
-
-  def test_copy_functional(self):
-    # without an impl this segfaults rather than fails: a regression here takes the whole run down
-    a = torch.arange(4, dtype=torch.float32, device=device)
-    out = torch.ops.aten.copy(a, torch.zeros(4, device=device))
-    np.testing.assert_equal(out.cpu().numpy(), [0., 0., 0., 0.])
-    np.testing.assert_equal(a.cpu().numpy(), [0., 1., 2., 3.])

 from tinygrad import Tensor
 class TestBackendHelpers(unittest.TestCase):
@@ -111,10 +111,6 @@ docs = [
  "numpy",
 ]
 mesa = ["tinymesa==25.2.7.2"]
-autogen = [
-  "pyyaml",
-  "mako",
-]


 [tool.mutmut]
@@ -67,26 +67,32 @@ class TestParseExpr(unittest.TestCase):

  def test_integer_literals(self):
    """Test parsing integer literals."""
-    self.assertIs(parse_expr('0', {}), UOp.const(0, dtypes.uint32))
-    self.assertIs(parse_expr('42', {}), UOp.const(42, dtypes.uint32))
-    self.assertIs(parse_expr('42U', {}), UOp.const(42, dtypes.uint32))
+    self.assertEqual(parse_expr('0', {}).val, 0)
+    self.assertEqual(parse_expr('42', {}).val, 42)
+    self.assertEqual(parse_expr('42U', {}).val, 42)

  def test_negative_integers(self):
    """Test parsing negative integer literals."""
-    self.assertIs(parse_expr('-1', {}), UOp.const(-1, dtypes.int))
+    result = parse_expr('-1', {})
+    self.assertEqual(result.val, -1)
+    self.assertEqual(result.dtype, dtypes.int)

  def test_float_literals(self):
    """Test parsing float literals."""
-    self.assertIs(parse_expr('1.0F', {}), UOp.const(1.0, dtypes.float32))
+    result = parse_expr('1.0F', {})
+    self.assertEqual(result.val, 1.0)
+    self.assertEqual(result.dtype, dtypes.float32)

  def test_hex_literals(self):
    """Test parsing hex literals."""
-    self.assertIs(parse_expr('0xFF', {}), UOp.const(255, dtypes.uint32))
+    result = parse_expr('0xFF', {})
+    self.assertEqual(result.val, 255)

  def test_variable_lookup(self):
    """Test variable lookup in parse_expr."""
    vrs = {'x': UOp.const(42, dtypes.uint32)}
-    self.assertIs(parse_expr('x', vrs), vrs['x'])
+    result = parse_expr('x', vrs)
+    self.assertEqual(result.val, 42)

  def test_binary_ops(self):
    """Test parsing binary operations."""
@@ -97,7 +103,9 @@ class TestParseExpr(unittest.TestCase):
    self.assertEqual(result.op, Ops.ADD)

    # Subtraction with constant folding
-    self.assertIs(parse_expr('10 - 5', {}), UOp.const(5, dtypes.uint32))
+    result = parse_expr('10 - 5', {})
+    self.assertEqual(result.op, Ops.CONST)
+    self.assertEqual(result.val, 5)

  def test_ternary(self):
    """Test parsing ternary expressions."""
@@ -134,8 +142,15 @@ class TestForLoopParsing(unittest.TestCase):
    S0 = UOp.const(0, dtypes.uint32)
    _vrs, assigns = parse_pcode(pcode, {'S0': S0})

-    # every cond folds (S0 is a const), leaving the default branch: -1 in the destination dtype
-    self.assertIs(assigns[0][1].simplify(), UOp.const(-1, dtypes.uint32))
+    # Check that the innermost value (default) is -1 (may be wrapped in CAST)
+    val = assigns[0][1]
+    # Traverse to innermost WHERE
+    while val.op == Ops.WHERE:
+      val = val.src[2]  # false branch
+    # Unwrap CAST if present
+    while val.op == Ops.CAST:
+      val = val.src[0]
+    self.assertEqual(val.val, -1)

  def test_ctz_parsing(self):
    """Test CTZ pcode parsing."""
@@ -247,8 +262,8 @@ class TestDSPcodePatterns(unittest.TestCase):
    _, assigns = parse_pcode(pcode, srcs)
    # Check addresses: 100 + 2*4 = 108, 100 + 5*4 = 120
    # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
-    self.assertIs(assigns[0][1][0].simplify(), UOp.const(108, dtypes.uint32))  # type: ignore[index]
-    self.assertIs(assigns[1][1][0].simplify(), UOp.const(120, dtypes.uint32))  # type: ignore[index]
+    self.assertEqual(assigns[0][1][0].simplify().val, 108)  # type: ignore[index]
+    self.assertEqual(assigns[1][1][0].simplify().val, 120)  # type: ignore[index]

  def test_ds_store_data_values(self):
    """Test DS_STORE_2ADDR_B32 uses correct data values."""
@@ -265,8 +280,8 @@ class TestDSPcodePatterns(unittest.TestCase):
    _, assigns = parse_pcode(pcode, srcs)
    # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
    # DATA[31:0] should preserve the value
-    self.assertIs(assigns[0][1][1].simplify(), UOp.const(0xAAAAAAAA, dtypes.uint32))  # type: ignore[index]
-    self.assertIs(assigns[1][1][1].simplify(), UOp.const(0xBBBBBBBB, dtypes.uint32))  # type: ignore[index]
+    self.assertEqual(assigns[0][1][1].simplify().val, 0xAAAAAAAA)  # type: ignore[index]
+    self.assertEqual(assigns[1][1][1].simplify().val, 0xBBBBBBBB)  # type: ignore[index]

 class TestConditionalParsing(unittest.TestCase):
  """Test conditional (if/elsif/else) pcode parsing."""
@@ -291,12 +306,12 @@ class TestConcatWidthParsing(unittest.TestCase):
  def test_permlanex16_altrow_concat(self):
    for row, expected in [(0, 1), (1, 0), (2, 3), (3, 2)]:
      parsed = parse_expr('{ row[1], ~row[0] }', {'row': UOp.const(row, dtypes.uint32)})
-      self.assertIs(parsed.simplify(), UOp.const(expected, dtypes.uint32))
+      self.assertEqual(parsed.simplify().val, expected)

  def test_permlane64_altlane_concat(self):
    for lane, expected in [(0, 32), (1, 33), (31, 63), (32, 0), (63, 31)]:
      parsed = parse_expr('{ ~lane[5], lane[4:0] }', {'lane': UOp.const(lane, dtypes.uint32)})
-      self.assertIs(parsed.simplify(), UOp.const(expected, dtypes.uint32))
+      self.assertEqual(parsed.simplify().val, expected)

  def test_permlane64_wave64_pcode_indices(self):
    vgpr = UOp.param(0, dtypes.uint32, (256,))
@@ -312,17 +327,19 @@ class TestConcatWidthParsing(unittest.TestCase):
      'S2': UOp.const(0, dtypes.uint32),
    }

-    def check_load_idx(v: UOp, expected: int):
+    def load_idx(v: UOp) -> int:
      simp = v.simplify()
      self.assertEqual(simp.op, Ops.LOAD)
      self.assertEqual(simp.src[0].op, Ops.INDEX)
-      self.assertIs(simp.src[0].src[1].simplify(), UOp.const(expected, dtypes.uint32))
+      idx = simp.src[0].src[1].simplify()
+      self.assertEqual(idx.op, Ops.CONST)
+      return idx.val

    _, assigns = parse_pcode(PCODE[VOP1Op.V_PERMLANE64_B32_E32], srcs)
    self.assertEqual(len(assigns), 64)
    for lane, (dst_idx, src_idx) in {0: (64, 32), 31: (95, 63), 32: (96, 0), 63: (127, 31)}.items():
-      self.assertIs(assigns[lane][1][0].simplify(), UOp.const(dst_idx, dtypes.uint32))  # type: ignore[index]
-      check_load_idx(assigns[lane][1][1], src_idx)  # type: ignore[index]
+      self.assertEqual(assigns[lane][1][0].simplify().val, dst_idx)  # type: ignore[index]
+      self.assertEqual(load_idx(assigns[lane][1][1]), src_idx)  # type: ignore[index]

 class TestAllPcode(unittest.TestCase):
  """Test that all pcode from all architectures can be parsed."""
@@ -1,28 +1,30 @@
 import unittest, contextlib
 from tinygrad import Device, Tensor, Context, TinyJit
-from tinygrad.device import Compiled, ProfileProgramEvent
+from tinygrad.device import Compiled, ProfileProgramEvent, ProfileDeviceEvent
 from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
 from tinygrad.viz.serve import load_amd_counters, VizData

@contextlib.contextmanager
 def save_sqtt():
-  Device[Device.DEFAULT].synchronize()
-  profile_start = len(Compiled.profile_events)
  data = VizData()
  yield data.ctxs
  Device[Device.DEFAULT].synchronize()
  Device[Device.DEFAULT]._at_profile_finalize()
-  load_amd_counters(data, [e for e in Compiled.profile_events[:profile_start] if isinstance(e, ProfileProgramEvent)] +
-                          Compiled.profile_events[profile_start:])
+  load_amd_counters(data, Compiled.profile_events)
  data.ctxs[:] = [r for r in data.ctxs if r["name"].startswith("SQTT")]

@unittest.skipUnless(Device.DEFAULT == "AMD", "only runs on AMD")
 class TestSQTTProfiler(unittest.TestCase):
+  # TODO: can we enable SQTT profiling in context?
  @classmethod
  def setUpClass(cls):
    if not Device[Device.DEFAULT].sqtt_enabled: raise unittest.SkipTest("device must be in SQTT profiling mode")

+  def setUp(self):
+    Device[Device.DEFAULT].synchronize()
+    Compiled.profile_events[:] = [e for e in Compiled.profile_events if isinstance(e, (ProfileProgramEvent, ProfileDeviceEvent))]
+
  def test_simple(self):
    t = Tensor.empty(1) + 1
    with save_sqtt() as sqtt:
@@ -1,5 +1,4 @@
 import unittest
-import functools
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.helpers import getenv, system, DEV
 from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm
@@ -10,7 +9,6 @@ from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX
 # Use DEV=NULL:HIP:gfx950 to also test the assembly
 def is_cdna4(): return Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950")

-@functools.cache
 def has_hipcc():
  try: system("hipcc --version")
  except Exception: return False
@@ -1,7 +1,7 @@
 import unittest, math
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.dtype import DTYPES_DICT
-from tinygrad.uop.ops import Ops, UOp, GroupOp
+from tinygrad.uop.ops import Ops, UOp
 from tinygrad.codegen.decomp.op import threefry2x32
 import numpy as np
 from test.helpers import not_support_multi_device
@@ -17,7 +17,7 @@ def _check_ast_count(desired_count:int, t:Tensor):
 class TestMovedConstFolding(unittest.TestCase):
  def test_contiguous_deviceless_const(self):
    t = Tensor(UOp.const(2.0, dtypes.float)).contiguous()
-    self.assertIs(t.uop, UOp.const(2.0, dtypes.float))
+    self.assertIs(t.uop.op, Ops.CONST)
    self.assertIsNone(t.uop.device)

  def test_add_shrunk_zero(self):
@@ -169,8 +169,8 @@ class TestMultiConstFolding(unittest.TestCase):
 class TestThreefryConstFolding(unittest.TestCase):
  def test_threefry(self):
    # THREEFRY(const,const) folds to a const once decomposed
-    x = threefry2x32(UOp.const(5, dtypes.uint64), UOp.const(10, dtypes.uint64)).simplify()
-    self.assertEqual([u.op for u in x.toposort() if u.op in GroupOp.ALU], [])
+    x = threefry2x32(UOp.const(5, dtypes.uint64), UOp.const(10, dtypes.uint64))
+    self.assertIs(x.simplify().op, Ops.CONST)

 class TestTautologicalCompare(unittest.TestCase):
  # without const folding, these would have triggered -Wtautological-compare in clang
@@ -188,6 +188,7 @@ class TestTautologicalCompare(unittest.TestCase):
    np.testing.assert_equal((Tensor(True) < Tensor(False)).numpy(), False)
    np.testing.assert_equal((Tensor(True) < Tensor(True)).numpy(), False)

+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support NaN comparison correctly")
  def test_a_eq_a(self):
    # self eq is always true for int or bool
    a = Tensor([1, 2, 3])
@@ -4,7 +4,7 @@ import numpy as np
 from tinygrad.dtype import AddrSpace, dtypes, Invalid
 from tinygrad.uop.ops import KernelInfo, AxisType, Ops
 from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import assert_kernel_count, KernelCountException
+from test.helpers import assert_kernel_count

 # **** kernels ****

@@ -474,7 +474,7 @@ class TestCustomKernelInput(unittest.TestCase):
    y.realize()
    kernel_count = GlobalCounters.kernel_count
    self.assertEqual(y.tolist(), x.add(1).tolist())
-    if kernel_count > max_kernels: raise KernelCountException(max_kernels, kernel_count)
+    self.assertLessEqual(kernel_count, max_kernels)
    # same test with @function, input is PARAM
    from tinygrad import function
    x0 = Tensor.arange(32).clone("CPU").realize()
@@ -487,7 +487,7 @@ class TestCustomKernelInput(unittest.TestCase):
    y = run(x0).realize()
    kernel_count = GlobalCounters.kernel_count
    self.assertEqual(y.tolist(), mop_fxn(x0).add(1).tolist())
-    if kernel_count > max_kernels: raise KernelCountException(max_kernels, kernel_count)
+    self.assertLessEqual(kernel_count, max_kernels)

  def test_reshape(self): self._test_mop(lambda x: x.reshape(16, 2), max_kernels=2)
  def test_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T, max_kernels=3)
@@ -340,9 +340,6 @@ class TestUint64DType(TestDType):
  DTYPE = dtypes.uint64
  def test_uint64_load(self):
    assert Tensor(2**64 - 1, dtype=dtypes.uint64).numpy() == 2**64 - 1
-  @unittest.skipIf(dtypes.double not in supported_dtypes, "needs float64")
-  def test_uint64_cast_double(self):
-    assert Tensor([2**32 + 1], dtype=dtypes.uint64).cast(dtypes.double).numpy() == 2**32 + 1

@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
 class TestEmulatedUInt64DType(TestUint64DType):
@@ -6,8 +6,6 @@ from tinygrad.tensor import _to_np_dtype
 from tinygrad.runtime.ops_python import from_storage_scalar
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
-from tinygrad.renderer.llvmir import CPULLVMRenderer
-from tinygrad.renderer.isa.x86 import X86Renderer
 from tinygrad.uop import Ops
 import numpy as np
 import pytest
@@ -66,8 +64,6 @@ ht.fp8e5m2fnuz = ht.uint8
 def universal_test(a, b, dtype, op):
  if not isinstance(op, tuple): op = (op, op)
  if op[0] == operator.mod and b == 0: return
-  # TODO: throws floating point exception
-  if isinstance(Device[Device.DEFAULT].renderer, (X86Renderer, CPULLVMRenderer)) and op[0] == operator.mod and a == dtype.min and b == -1: return
  # lt and max with nan is undefined in tinygrad
  if op[0] in (operator.lt, Tensor.maximum) and (math.isnan(a) or math.isnan(b)): return
  ta, tb = Tensor([a], dtype=dtype), Tensor([b], dtype=dtype)
@@ -7,7 +7,7 @@ from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
 from tinygrad.renderer.isa import IselContext

 # INDEX on a register value with a constant index extracts a single element (the old GEP)
-def lane(y:UOp, i:int) -> UOp: return y.index(UOp.cconst(i, dtypes.int), dtype=y.dtype)
+def lane(y:UOp, i:int) -> UOp: return y.index(UOp.const(i, dtypes.int), dtype=y.dtype.scalar())

@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86")
 class TestIselX86(unittest.TestCase):
@@ -46,10 +46,10 @@ class TestIselX86(unittest.TestCase):
  # complex address is [base + index*scale + displacement]
  def test_complex_address(self):
    a = UOp.variable("a", 0, 0, dtypes.int32)
-    load = UOp.param(0, dtypes.int32, (16,)).index(a + UOp.cconst(1, dtypes.int32)).load()
+    load = UOp.param(0, dtypes.int32, (16,)).index(a + 1).load()
    n = self.isel_rewrite(load)
    # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32
-    self.assertTrue(n.src[2].dtype is dtypes.int8 and n.src[2].src[0].op is Ops.CONST and n.src[2].src[0].val == 4)
+    self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].val == 4)

 if __name__ == "__main__":
  unittest.main()
@@ -360,7 +360,7 @@ class TestJitGraphSplit(unittest.TestCase):
    self.expect(f, inp, inp_cpu,
      graph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
      multigraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
-      hcqgraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()]) # cpu is hcq2 now, it does not join hcq graphs
+      hcqgraph=[self.ji_graph(4)])

  def test_jit_cpu_several(self):
    if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
@@ -377,9 +377,9 @@ class TestJitGraphSplit(unittest.TestCase):
    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
    inp_cpu = Tensor.randn(10, 10, device="CPU").realize()
    self.expect(f, inp, inp_cpu,
-      graph=[self.ji_graph(2), self.ji_comp(), self.ji_comp(), self.ji_comp()],
-      multigraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp(), self.ji_comp()],
-      hcqgraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp(), self.ji_comp()])
+      graph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
+      multigraph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
+      hcqgraph=[self.ji_graph(5)])

  def test_jit_multidev(self):
    if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
@@ -16,6 +16,8 @@ from test.helpers import replace_opts, check_schedule
 from test.backend.test_softmax_fusion import single_kernel_softmax
 MOCKGPU = DEV.interface.startswith("MOCK")

+from tinygrad.uop.render import print_uops # noqa: F401 # pylint: disable=unused-import
+
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, ISARenderer), "isa backends don't preserve the op spec when lowering")
 class TestLinearizer(unittest.TestCase):
  def test_arg_dedup(self):
@@ -28,7 +30,7 @@ class TestLinearizer(unittest.TestCase):
    c = ((a.shrink(((0, 2),)) - a.shrink(((2, 4),))) - (b.shrink(((0, 2),)) - b.shrink(((2, 4),))))
    linear = c.schedule_linear()
    run_linear(linear)
-    rawbufs = [s.buffer for s in linear.src[-1].src[1:] if not s.is_bound_var]
+    rawbufs = [s.buffer for s in linear.src[-1].src[1:] if s.op is not Ops.BIND]
    assert len(rawbufs) == 3 and set(rawbufs[1:]) == {a.uop.base.realized, b.uop.base.realized}
    np_c = (np_a[:2] - np_a[2:]) - (np_b[:2] - np_b[2:])
    np.testing.assert_allclose(np_c, c.numpy(), atol=1e-4, rtol=1e-4)
@@ -246,10 +248,11 @@ class TestLinearizer(unittest.TestCase):
    uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[1].src)
    begin_range = [i for i, x in enumerate(uops) if x.op is Ops.RANGE][-1]
    end_range = [i for i, x in enumerate(uops) if x.op is Ops.END][0]
+    for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype)
    for u in uops:
      if u.op is Ops.STORE and u.src[0].addrspace is AddrSpace.REG:
        if uops.index(u) < begin_range:
-          assert u.src[1].op not in GroupOp.ALU
+          assert u.src[1].op is Ops.CONST
        else:
          assert u.src[1].op in GroupOp.ALU
          assert begin_range < uops.index(u) < end_range
@@ -265,9 +268,9 @@ class TestLinearizer(unittest.TestCase):
    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
    idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
    idxs = sorted(idxs, key=lambda uop: uop.arg)
-    assert (idxs[0].arg, idxs[0].src[0].src[0].val) == ('gidx0', 6), idxs[0]
-    assert (idxs[1].arg, idxs[1].src[0].src[0].val) == ('gidx1', 5), idxs[1].arg
-    assert (idxs[2].arg, idxs[2].src[0].src[0].val) == ('gidx2', 4), idxs[2].arg
+    assert (idxs[0].arg, idxs[0].src[0].val) == ('gidx0', 6), idxs[0]
+    assert (idxs[1].arg, idxs[1].src[0].val) == ('gidx1', 5), idxs[1].arg
+    assert (idxs[2].arg, idxs[2].src[0].val) == ('gidx2', 4), idxs[2].arg

  def test_sum_collapse(self):
    t = Tensor([2]).reshape(1, 1).expand(256, 256).sum()
@@ -408,7 +411,7 @@ def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
  last_call = linear.src[-1]
  ast = last_call.src[0]
  assert ast.op is Ops.SINK, f"helper_realized_ast expects a SINK {last_call}"
-  last_bufs = [s.buffer for s in last_call.src[1:] if not s.is_bound_var]
+  last_bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]
  # now all input buffers in last_call should be realized
  # create fresh buffers for the outputs
  bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(ast.src) else x for i,x in enumerate(last_bufs)]
@@ -9,7 +9,7 @@ from extra.llama_kernels.swiglu import swiglu
 from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
 from extra.thunder.amd.fa import custom_fused_qkv_rope_backward, fused_qkv_rope
 from test.helpers import needs_second_gpu, assert_kernel_count
-from test.backend.test_asm_gemm import has_hipcc, is_cdna4
+from test.backend.test_asm_gemm import has_hipcc

 def run_fused_ce(bs:int, seqlen:int, vocab:int, label_smoothing:float=0.0) -> None:
  Tensor.manual_seed(0)
@@ -99,20 +99,22 @@ class TestLocalAmax(unittest.TestCase):
    assert_kernel_count(2)
    self.assertEqual(out.tolist(), [[0., 7., 14., 21.], [28., 35., 42., 49.], [120., 135., 150., 165.], [180., 195., 210., 225.]])

+@unittest.skipUnless(has_hipcc() and Device.DEFAULT == "AMD", "requires hipcc to compile and amd device to run")
 class TestFusedQKVRoPE(unittest.TestCase):
  SHAPE = (2, 8192, 32, 8, 128)

-  def setUp(self):
-    if dtypes.bfloat16 not in Device[Device.DEFAULT].renderer.supported_dtypes(): self.skipTest("test uses bf16 inputs")
-
  def rand_bf16(self, *shape:int) -> Tensor:
    return (Tensor.randn(*shape) * 0.1).cast(dtypes.bfloat16).contiguous().realize()

-  def test_forward(self):
+  def freqs_cis(self) -> Tensor:
+    _, N, _, _, D = self.SHAPE
+    return precompute_freqs_cis(D, N * 2).cast(dtypes.bfloat16).clone().realize()
+
+  def test_llama31_8b_forward(self):
    Tensor.manual_seed(0)
-    B, N, H, H_KV, D = 1, 32, 8, 2, 16
+    B, N, H, H_KV, D = self.SHAPE
    GROUP = H // H_KV
-    freqs_cis = (Tensor.randn(1, N * 2, 1, D // 2, 2) * 0.1).cast(dtypes.bfloat16).contiguous().realize()
+    freqs_cis = self.freqs_cis()

    x = self.rand_bf16(B, N, H_KV * (GROUP + 2) * D)
    q, k, v = fused_qkv_rope(x, freqs_cis, H, H_KV, D)
@@ -129,13 +131,12 @@ class TestFusedQKVRoPE(unittest.TestCase):
      self.assertTrue(k.allclose(k_ref, atol=2e-2, rtol=0).item(), "K forward mismatch")
      self.assertTrue(v.allclose(v_ref, atol=0, rtol=0).item(), "V forward mismatch")

-  @unittest.skipUnless(has_hipcc() and is_cdna4(), "backward kernel requires hipcc to compile")
-  def test_llama31_8b(self):
+  def test_llama31_8b_backward(self):
    Tensor.manual_seed(1)
    B, N, H, H_KV, D = self.SHAPE
    PARTIALS = 2
    GROUP = H // H_KV
-    freqs_cis = precompute_freqs_cis(D, N * 2).cast(dtypes.bfloat16).clone().realize()
+    freqs_cis = self.freqs_cis()
    dq = self.rand_bf16(B, N, H, D)
    dk_partial = self.rand_bf16(B * PARTIALS, N, H_KV, D)
    dv_partial = self.rand_bf16(B * PARTIALS, N, H_KV, D)
@@ -1,12 +1,12 @@
 import unittest, random
 from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable
-from tinygrad.uop.ops import Ops, UOp, AxisType, graph_rewrite
+from tinygrad.uop.ops import Ops, UOp, AxisType
 from tinygrad.helpers import getenv, prod, Context
 from tinygrad.nn.state import get_parameters
-from tinygrad.engine.realize import run_linear, compile_linear, lower_and_compile, pm_beam
+from tinygrad.engine.realize import run_linear, compile_linear
 import numpy as np
 from hypothesis import given, strategies as strat, settings
-from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count, KernelCountException
+from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count

 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
@@ -72,16 +72,16 @@ class TestMultiTensor(unittest.TestCase):
    X.shard_(devices_2, 0)
    out = (X + X)
    linear = compile_linear(out.schedule_linear())
-    uops = [call.src[0].src[0] for call in linear.src if call.src[0].op is Ops.PROGRAM]
+    names = [call.src[0].src[0].arg.name for call in linear.src if call.src[0].op is Ops.PROGRAM]
    run_linear(linear)
-    self.assertEqual(len(set(uops)), 1, "function was relinearized")
+    self.assertEqual(len(set(names)), 1, "function was relinearized")

  def test_shard_beam(self):
    cpu_2 = ("CPU:1", "CPU:2")
    src = Tensor.ones(16).shard(cpu_2, 0).realize()
-    lin = UOp(Ops.LINEAR, src=(src.to(cpu_2[::-1]).schedule_linear().src[0],))
-    with Context(BEAM=1, IGNORE_BEAM_CACHE=1): call = lower_and_compile(graph_rewrite(lin, pm_beam, ctx=1, walk=True)).src[0]
-    self.assertNotEqual(call.src[0].src[0].arg.applied_opts, ())
+    pad = src.to(cpu_2[::-1]).schedule_linear().src[0]
+    with Context(BEAM=1, IGNORE_BEAM_CACHE=1): prg = compile_linear(UOp(Ops.LINEAR, src=(pad,))).src[0].src[0]
+    self.assertNotEqual(prg.src[0].arg.applied_opts, ())

  def test_shard_same_device(self):
    X = Tensor.ones(256).contiguous().realize()
@@ -395,7 +395,7 @@ class TestMultiBufferView(unittest.TestCase):
    linear, var_vals = b_multi.linear_with_vars()
    if all(not d.startswith(("WEBGPU", "CL")) for d in b_multi.device):
      compiled = [call for call in linear.src if call.src[0].op is Ops.SINK]
-      if len(compiled) != 0: raise KernelCountException(0, len(compiled))
+      self.assertEqual(len(compiled), 0, f"expected zero compiled kernels, got {len(compiled)}")
    run_linear(linear, var_vals)
    np.testing.assert_equal(b_multi.numpy(), b_ref.numpy())

@@ -720,11 +720,10 @@ class TestOps(unittest.TestCase):
      return torch.autograd.grad(t ** c, t)[0].item()
    for x in [-math.inf, 0, 1, math.inf]:
      for c in [-1, 0, 0.3, 1, 2]:
-        torch_out = get_torch_gradient(x, c)
-        # the pow backward routes through exp2/log2, whose 0/inf behavior is undefined on WEBGPU
-        if Device.DEFAULT == "WEBGPU" and not math.isfinite(torch_out): continue
        tiny_out = get_tiny_gradient(x, c)
+        torch_out = get_torch_gradient(x, c)
        if math.isnan(tiny_out):
+          if Device.DEFAULT == "WEBGPU": continue # TODO: WEBGPU issue with nan
          assert math.isnan(torch_out)
        else:
          self.assertAlmostEqual(tiny_out, torch_out, msg=f"{x}, {c}")
@@ -750,7 +749,6 @@ class TestOps(unittest.TestCase):
  def test_exp2_log2_zero_times_negative(self):
    # gallivm's exp2/log2 have "undefined behavior with infs, 0s and nans", so exp2(log2(0)*y) returns 0 instead of inf
    helper_test_op(None, lambda x,y: (x.log2()*y).exp2(), lambda x,y: (x.log2()*y).exp2(), vals=[[0.0], [-0.7]], forward_only=True)
-  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "pow at 0 routes through exp2/log2, whose 0/inf behavior is undefined on WEBGPU")
  def test_pow_zero_const(self):
    helper_test_op(None, lambda x: x**0.3, vals=[[0.0]])
    helper_test_op(None, lambda x: x**0.0, vals=[[0.0]])
@@ -822,10 +820,6 @@ class TestOps(unittest.TestCase):
    helper_test_op([], lambda: tor0&tor1, lambda: ten0&ten1, forward_only=True)

    helper_test_op(None, lambda x: (1 < x) & (x < 2), forward_only=True, vals=[[1.2, 1.2, 1.2, 3.2]])
-    helper_test_op([(3000,)]*10, lambda *xs: (sum(xs[1:], xs[0]) > 5) & (xs[0] < 0.9), forward_only=True)
-
-    if not COMPILE_ONLY:
-      np.testing.assert_equal((Tensor(2**64-1, dtype=dtypes.uint64) & 0xFFFFFFFF).numpy(), 0xFFFFFFFF)

  def test_or(self):
    data = [[1,-8,1],[32,1,6]]
@@ -2168,10 +2162,6 @@ class TestOps(unittest.TestCase):
  def test_roll(self):
    helper_test_op([(2, 4)], lambda x: x.roll(1))
    helper_test_op([(2, 4)], lambda x: x.roll((1,)))
-    helper_test_op([(0,)], lambda x: x.roll(1, 0))
-    helper_test_op([(2, 0, 3)], lambda x: x.roll(1, 0))
-    helper_test_op([(2, 0, 3)], lambda x: x.roll(1, 1))
-    helper_test_op([(2, 0, 3)], lambda x: x.roll(1))
    self.helper_test_exception([(2, 4)], lambda x: x.roll((1, 2)), expected=RuntimeError)
    helper_test_op([(2, 4)], lambda x: x.roll(1, 0))
    helper_test_op([(2, 4)], lambda x: x.roll(-1, 0))
@@ -2,8 +2,7 @@ import unittest, pickle, types, tracemalloc
 import numpy as np
 from tinygrad import Tensor, Device, TinyJit, Variable, dtypes
 from tinygrad.helpers import GlobalCounters, ContextVar, Context, DEV
-from tinygrad.uop.ops import PatternMatcher, UPat, UOp, deconstruct_function
-from test.helpers import KernelCountException
+from tinygrad.uop.ops import PatternMatcher, UPat, UOp

 class TestPickle(unittest.TestCase):
  def test_pickle_code_object(self):
@@ -12,11 +11,6 @@ class TestPickle(unittest.TestCase):
    fxn = types.FunctionType(pickle.loads(code_str), globals())
    self.assertEqual(fxn(2), 4)

-  def test_deconstruct_function_nested_comprehension(self):
-    # pre PEP 709, each comprehension is its own code object, so dtypes here is referenced two code objects deep
-    def fxn(): return [[dtypes.int for _ in range(2)] for _ in range(2)]
-    self.assertEqual(types.FunctionType(*deconstruct_function(fxn))(), fxn())
-
  def test_pickle_pattern_matcher(self):
    pm = PatternMatcher([(UPat.cvar('x'), lambda x: x*2)])
    sink = UOp.const(2)
@@ -42,7 +36,7 @@ class TestPickle(unittest.TestCase):
    t2:Tensor = pickle.loads(st)
    np.testing.assert_equal(t_values, t2.numpy())
    # expect at most one COPY kernel
-    if GlobalCounters.kernel_count > 1: raise KernelCountException(1, GlobalCounters.kernel_count)
+    self.assertLessEqual(GlobalCounters.kernel_count, 1)

  def test_pickle_realized_tensor_alt(self):
    print("** init")
@@ -82,7 +82,7 @@ class TestQuantizeOnnxCPU(unittest.TestCase):
      linear = run_onnx({"input":inp})["output"].schedule_linear()
      prg = to_program(linear.src[-2].src[0], renderer=Device[Device.DEFAULT].renderer)
      daccs = [u for u in tuple(prg.src[1].src) if u.op is Ops.BUFFER and u.addrspace is AddrSpace.REG]
-      assert all(u.dtype is dtypes.int for u in daccs)
+      assert all(u.dtype.scalar() is dtypes.int for u in daccs)

@unittest.skipIf(Device.DEFAULT != "DSP", "only tests for DSP")
 class TestQuantizeOnnx(unittest.TestCase):
@@ -77,14 +77,6 @@ class TestCStyleFailures(unittest.TestCase):

@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "tests for wgsl renderer")
 class TestWGSLFailures(unittest.TestCase):
-  def test_folded_packed_store(self):
-    b = UOp.param(0, dtypes.char, (4,))
-    idx = b.index(UOp.const(0).cast(dtypes.int))
-    store = UOp.store(idx, UOp.load(idx, dtype=dtypes.uint32) & UOp.const(0xffffff00).cast(dtypes.uint32))
-    src = Device[Device.DEFAULT].renderer.render(UOp.sink(store, arg=KernelInfo()).toposort())
-    self.assertIn("atomicAnd(&data0_4[0],4294967040u);", src)
-    self.assertNotIn("atomicAdd", src)
-
  def test_multiply_infinity(self):
    # multiplying a positive constant by infinity should return infinity
    # WGSL pipelines do not handle this reliably, some of which return zero, unless infinity always comes from a read on a dynamic buffer
@@ -176,7 +176,7 @@ class TestLimitBufs(unittest.TestCase):

  def test_limit_bufs_linear_scaling(self):
    def sched_time(n):
-      with Context(TRACK_MATCH_STATS=0, DEBUG=0, PARALLEL=0):
+      with Context(TRACK_MATCH_STATS=0, DEBUG=0):
        bufs = [Tensor.ones(16).contiguous().realize() for _ in range(4)]
        root = bufs[0]
        for i in range(n): root = root + bufs[i % 4]
@@ -653,19 +653,9 @@ class TestZeroShapeTensor(unittest.TestCase):
    np.testing.assert_equal(Tensor([[1, 2]]).pad_to(2, 3).numpy(), [[1, 2, 0], [0, 0, 0]])
    np.testing.assert_equal(Tensor([[1, 2]]).pad_to(1, 3).numpy(), [[1, 2, 0]])
    np.testing.assert_equal(Tensor([[1, 2]]).pad_to(None, 3).numpy(), [[1, 2, 0]])
-    np.testing.assert_equal(Tensor([1, 2]).pad_to(4, value=2).numpy(), [1, 2, 2, 2])
-    np.testing.assert_equal(Tensor([[1, 2]]).pad_to(2, 3, value=-1).numpy(), [[1, 2, -1], [-1, -1, -1]])
-    np.testing.assert_equal(Tensor([1, 2]).pad_to(None, value=5).numpy(), [1, 2])  # no-op pad ignores the fill
    with self.assertRaises(ValueError): Tensor([1, 2]).pad_to(2, 3)
    with self.assertRaises(ValueError): Tensor([[1, 2]]).pad_to(3)

-  def test_max_shape(self):
-    from tinygrad import UOp
-    t = Tensor.empty(2, UOp.variable('v', 1, 32), 4)
-    self.assertEqual(t.max_shape, (2, 32, 4))
-    self.assertEqual(t.max_numel(), 2*32*4)
-    self.assertEqual(Tensor.empty(2, 3).max_shape, (2, 3))
-
  def test_shrink_into_zero(self):
    t = Tensor.rand(3, 4).realize()
    assert t.shrink((None, (2, 2))).realize().shape == (3, 0)
@@ -2,7 +2,7 @@ from typing import Optional, Any
 import unittest, math
 import numpy as np
 from tinygrad.tensor import Tensor, _to_np_dtype
-from tinygrad.helpers import Context, ceildiv
+from tinygrad.helpers import Context
 from tinygrad.dtype import dtypes, DType, AddrSpace, ConstFloat  # noqa: F401
 from tinygrad.device import Buffer, Device
 from tinygrad.uop.ops import Ops, UOp, KernelInfo, AxisType, buffers
@@ -57,35 +57,6 @@ def _test_uops_result(output_dtype, uops, res):
  run_uops([out], [buf])
  return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]

-@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage) and
-                     dtypes.uint64 in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires C-style pointer bitcast and 64-bit ints")
-class TestBitcastBufferView(unittest.TestCase):
-  @Context(SPEC=2)
-  def test_render(self):
-    buf = UOp.param(0, dtypes.uint32, (4,))
-    uops = to_uops_list([buf.shrink(((1, 3),)).bitcast(dtypes.uint64).index(0).store(1)], ren=Device[Device.DEFAULT].renderer)
-    idx = next(u for u in uops if u.op is Ops.INDEX and u.src[0].op is Ops.BITCAST)
-    self.assertEqual(idx.src[0].src[0].op, Ops.SHRINK)
-    Device[Device.DEFAULT].renderer.render(uops)
-
-  @Context(SPEC=2)
-  def test_load(self):
-    val = 0x1122334455667788
-    src, out = UOp.param(0, dtypes.uint32, (4,)), UOp.param(1, dtypes.uint64, (1,))
-    ibuf = Buffer(Device.DEFAULT, 4, dtypes.uint32, initial_value=np.array([0, 0x55667788, 0x11223344, 0], dtype=np.uint32).tobytes())
-    obuf = Buffer(Device.DEFAULT, 1, dtypes.uint64).allocate()
-    run_uops([out.index(0).store(src.shrink(((1, 3),)).bitcast(dtypes.uint64).index(0))], [ibuf, obuf])
-    self.assertEqual(np.frombuffer(obuf.as_memoryview(), dtype=np.uint64)[0], val)
-
-  @Context(SPEC=2)
-  def test_store(self):
-    val = 0x1122334455667788
-    dst = UOp.param(0, dtypes.uint32, (6,))
-    buf = Buffer(Device.DEFAULT, 6, dtypes.uint32, initial_value=bytes(24))
-    view = dst.shrink(((1, 5),)).bitcast(dtypes.uint64)  # two stores through one view: it must inline, not get a declared vector-pointer
-    run_uops([view.index(0).store(val ^ 0xff), view.index(1).store(val)], [buf])
-    self.assertEqual(np.frombuffer(buf.as_memoryview(), dtype=np.uint64, count=2, offset=4).tolist(), [val ^ 0xff, val])
-
 class TestUOps(unittest.TestCase):
  def _equal(self, v1, v2):
    assert isinstance(v2, (float, int, bool))
@@ -222,16 +193,15 @@ class TestLocalAccess(unittest.TestCase):
  @unittest.skipUnless(Device.DEFAULT == "WEBGPU", "Test local memory size for packed data types")
  def test_packed_smem_size(self):
    _dtypes = [dtypes.char, dtypes.uchar, dtypes.short, dtypes.ushort, dtypes.half]
-    # a partial word still needs a whole word, so sizes that don't fill one must round up
-    for size in (16, 5):
-      for dtype in _dtypes:
-        temp = UOp.placeholder((size,), dtype, slot=0, addrspace=AddrSpace.LOCAL)
-        uops = to_uops_list([temp], ren=Device[Device.DEFAULT].renderer)
-        out = Device[Device.DEFAULT].renderer.render(uops)
-        # half is supported in wgsl, so it doesn't have to be packed
-        corrected_size = ceildiv(size, 4//dtype.itemsize) if dtype != dtypes.half else size
-        # temp0: array<{Device[Device.DEFAULT].renderer.buf_map(dtype)},{corrected_size}>;
-        self.assertIn(f",{corrected_size}>;", out)
+    size = 16
+    for dtype in _dtypes:
+      temp = UOp.placeholder((size,), dtype, slot=0, addrspace=AddrSpace.LOCAL)
+      uops = to_uops_list([temp], ren=Device[Device.DEFAULT].renderer)
+      out = Device[Device.DEFAULT].renderer.render(uops)
+      # half is supported in wgsl, so it doesn't have to be packed
+      corrected_size = size//(4//dtype.itemsize) if dtype != dtypes.half else size
+      # temp0: array<{Device[Device.DEFAULT].renderer.buf_map(dtype)},{corrected_size}>;
+      self.assertIn(f",{corrected_size}>;", out)

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared memory")
  @unittest.skip("tinygrad doesn't support this behavior")
@@ -1,8 +1,6 @@
 import unittest, numpy as np
 from unittest.mock import patch
 from tinygrad import Device, Tensor
-from tinygrad.device import Buffer
-from tinygrad.dtype import dtypes
 from tinygrad.helpers import getenv
 from tinygrad.runtime.support.hcq2 import HCQ_DEVS, all_devices_in

@@ -12,22 +10,5 @@ class TestHCQ2(unittest.TestCase):
    with patch.object(Device[Device.DEFAULT], "has_copy_queue", False):
      np.testing.assert_equal(Tensor(np.arange(61, dtype=np.float32)).to(Device.DEFAULT).contiguous().realize().numpy(), np.arange(61))

-  @unittest.skipIf(Device.DEFAULT == "CPU", "staged copies need a non-CPU hcq2 device")
-  def test_staged_copy_slot_reuse(self):
-    # chunks of a staged copy rotate through the staging buffer slots, many rotations must stay bit-exact in both directions
-    import tinygrad.runtime.support.hcq2 as hcq2
-    buf = Buffer("CPU", 1 << 20, dtypes.uint8, preallocate=True)
-    data = np.random.default_rng(42).integers(0, 256, (5 << 20) + 123, dtype=np.uint8)
-    with patch.object(hcq2, "STAGING_SIZE", 1 << 20), patch.object(hcq2, "STAGING_SLOTS", 4), patch.object(hcq2, "_staging", lambda: buf):
-      np.testing.assert_equal(Tensor(data).to(Device.DEFAULT).realize().numpy(), data)
-
-  def test_overlapping_device_tuples(self):
-    # an op on a wide device tuple followed by an op on an overlapping smaller tuple used to MMU-fault the smaller one
-    d4, d2 = tuple(f"{Device.DEFAULT}:{i}" for i in range(4)), tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
-    ref = Tensor.arange(16).contiguous().realize()
-    Tensor(ref.uop.copy_to_device(d4)).realize()
-    out = Tensor.ones(8).shard(d2, axis=0).contiguous().realize()
-    np.testing.assert_equal(out.numpy(), np.ones(8))
-
 if __name__ == "__main__":
  unittest.main()
@@ -1,17 +0,0 @@
-from tinygrad import Device, Tensor, TinyJit, dtypes
-from tinygrad.helpers import Timing, Context
-
-GPUS, DEPTH, SZ = 8, 4, 128 * 2**20
-WARMUP, ITERS = 3, 5
-devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(GPUS))
-bufs = tuple(Tensor.empty(SZ, dtype=dtypes.uint8, device=dev).contiguous().realize() for _ in range(DEPTH) for dev in devs)
-
-@TinyJit
-def all_to_all(*srcs:Tensor): return Tensor.realize(*(src.to(dst) for i,src in enumerate(srcs) for j,dst in enumerate(devs) if i % GPUS != j))
-
-if __name__ == "__main__":
-  with Context(ALL2ALL=1, JIT_BATCH_SIZE=0):
-    for i in range(-WARMUP, ITERS):
-      with Timing("ALL2ALL ", lambda ns: f" {SZ*GPUS*(GPUS-1)*DEPTH/ns:.2f} GB/s", enabled=i>=0):
-        all_to_all(*bufs)
-        for dev in devs: Device[dev].synchronize()
@@ -1,5 +1,5 @@
-import unittest, time, itertools
-from tinygrad import Tensor, Context
+import unittest, time
+from tinygrad import Tensor

 class TestScheduleScaling(unittest.TestCase):
  """Test that .schedule() scales linearly with graph size (no O(n^2) behavior)."""
@@ -130,18 +130,5 @@ class TestScheduleScaling(unittest.TestCase):
      return parts[0].cat(*parts[1:])
    self._assert_linear(concat_chain)

-  @Context(DEV="NULL:HIP:gfx1100")
-  def test_custom_kernel_assign_scaling(self):
-    from tinygrad.uop.ops import UOp, Ops, KernelInfo
-    from tinygrad.runtime.autogen.amd.rdna3.ins import s_nop
-    count = itertools.count(0)
-    def custom_kernel_assign(n):
-      def custom_asm(out):
-        return UOp(Ops.PROGRAM, src=(UOp.sink(out, arg=KernelInfo(f"fxn_{next(count)}")),
-                                     UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=s_nop(i)) for i in range(n*8)))))
-      call = Tensor.custom_kernel(Tensor.empty(1), fxn=custom_asm)[0]
-      return Tensor.cat(*[Tensor.empty(1).assign(call+i) for i in range(n)])
-    self._assert_linear(custom_kernel_assign, n_small=50, n_large=500)
-
 if __name__ == '__main__':
  unittest.main(verbosity=2)
@@ -6,7 +6,7 @@ import numpy as np
 class TestDevCopySpeeds(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
-    cls.sz = getenv("SIZE", 2000000)
+    cls.sz = getenv("SIZE", 2e6)
    cls.dev = Device["AMD"]
    if not cls.dev.is_usb(): raise unittest.SkipTest("only test this on USB devices")

@@ -44,7 +44,7 @@ def realized_matmul():
  z = y.matmul(x)
  Tensor.realize(z)
 def realized_gradient():
-  x = Tensor.eye(3).clone()
+  x = Tensor.eye(3)
  y = Tensor([[2.0,0,-2.0]])
  z = y.matmul(x).sum()
  z.backward()
@@ -86,9 +86,7 @@ def assert_jit_cache_len(fxn, expected_len):
  if linear is None or not linear.src:
    if expected_len != 0: raise KernelCountException(expected_len, 0)
    return
-  if expected_len and all(call_is_hcq(call) for call in linear.src): # HCQ2: one batch submitter, or fence + reset + merged calls + finalizer
-    from tinygrad.runtime.support.hcq2 import HCQ_RUNTIME_DEV
-    expected_len = 1 if HCQ_RUNTIME_DEV.value == "CPU" else 4
+  if expected_len and all(call_is_hcq(call) for call in linear.src): expected_len = 3 # HCQ2: merged same-queue calls + finalizer + bumps
  if call_is_graph(linear.src[0]):
    if len(linear.src) != 1: raise KernelCountException(1, len(linear.src))
    inner = linear.src[0].src[0].src[0]  # LINEAR UOp inside CUSTOM_FUNCTION
@@ -1,20 +1,5 @@
 # Tokenizer-based expression parser for AMD pcode
-import ast, itertools, operator, re
 from typing import Any, Callable
-
-_BINOPS = {ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, ast.FloorDiv: operator.floordiv,
-           ast.Mod: operator.mod, ast.LShift: operator.lshift, ast.RShift: operator.rshift,
-           ast.BitAnd: operator.and_, ast.BitOr: operator.or_, ast.BitXor: operator.xor}
-def _const_int(expr: str) -> int:
-  """Evaluate a compile-time integer expression (integer literals and basic arithmetic only)."""
-  def ev(node: ast.AST) -> int:
-    if isinstance(node, ast.Expression): return ev(node.body)
-    if isinstance(node, ast.Constant) and isinstance(node.value, int): return node.value
-    if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.USub, ast.UAdd)):
-      return (-1 if isinstance(node.op, ast.USub) else 1) * ev(node.operand)
-    if isinstance(node, ast.BinOp) and type(node.op) in _BINOPS: return _BINOPS[type(node.op)](ev(node.left), ev(node.right))
-    raise ValueError(f"not a constant integer expression: {expr!r}")
-  return ev(ast.parse(expr.strip(), mode='eval'))
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.codegen.decomp.dtype import f2f
@@ -375,13 +360,22 @@ _FUNCS: dict[str, Callable[..., UOp]] = {
  'fp8_to_f32': _fp8_to_f32, 'bf8_to_f32': _bf8_to_f32, 'f32_to_fp8': _f32_to_fp8, 'f32_to_bf8': _f32_to_bf8,
  'f32_to_bf16': _f32_to_bf16, 'f32_to_bf16_SR': _f32_to_bf16_sr, 'f32_to_bf16_sr': _f32_to_bf16_sr,
 }
-# min/max family: min/max + 3-input (x3), IEEE num variants (f16/f32 only), and long names minimum/maximum (f16/f32 only)
-for is_max, name, full in [(False, 'min', 'minimum'), (True, 'max', 'maximum')]:
-  for dt, sfx, pre in [(dtypes.float32, 'f32', None), (dtypes.int, 'i32', None), (dtypes.uint32, 'u32', None),
-                       (dtypes.int16, 'i16', None), (dtypes.uint16, 'u16', None), (dtypes.half, 'f16', _f16_extract)]:
-    def mm(*a, im=is_max, d=dt, p=pre): return _minmax_reduce(im, d, *(a if p is None else [p(x) for x in a]))
-    extra = (f'v_{name}_num_{sfx}', f'v_{name}3_num_{sfx}', f'v_{full}_{sfx}', f'v_{full}3_{sfx}') if dt in (dtypes.float32, dtypes.half) else ()
-    for fn in (f'v_{name}_{sfx}', f'v_{name}3_{sfx}', *extra): _FUNCS[fn] = mm
+for is_max, name in [(False, 'min'), (True, 'max')]:
+  for dt, sfx in [(dtypes.float32, 'f32'), (dtypes.int, 'i32'), (dtypes.uint32, 'u32'), (dtypes.int16, 'i16'), (dtypes.uint16, 'u16')]:
+    _FUNCS[f'v_{name}_{sfx}'] = lambda *a, im=is_max, d=dt: _minmax_reduce(im, d, *a)
+    _FUNCS[f'v_{name}3_{sfx}'] = lambda *a, im=is_max, d=dt: _minmax_reduce(im, d, *a)
+# f16 min/max/min3/max3/med3
+for is_max, name in [(False, 'min'), (True, 'max')]:
+  _FUNCS[f'v_{name}_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}3_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}_num_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}_num_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
+  _FUNCS[f'v_{name}3_num_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}3_num_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
+  _FUNCS[f'v_{name}imum_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}imum_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
+  _FUNCS[f'v_{name}imum3_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}imum3_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)

 # ═══════════════════════════════════════════════════════════════════════════════
 # TOKENIZER/PARSER
@@ -896,8 +890,6 @@ class Parser:
      return result & _isnan(l).logical_not() & _isnan(r).logical_not()
    return result

-_break_var_ids = itertools.count()  # unique names for per-loop break-tracking variables
-
 def _match_bracket(toks: list[Token], start: int) -> tuple[int, list[Token]]:
  """Match brackets from start, return (end_idx, inner_tokens)."""
  j, depth = start + 1, 1
@@ -995,7 +987,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
        i += 1
      # Execute loop with break support
      has_break = any('break' in bl.lower() for bl in body_lines)
-      found_var = f'_found_{next(_break_var_ids)}' if has_break else None
+      found_var = f'_found_{id(body_lines)}' if has_break else None
      if found_var: env[found_var] = block_assigns[found_var] = _const(dtypes.bool, False)
      for loop_i in range(start_val, end_val + 1):
        subst_lines = [_subst_loop_var(bl, loop_var, loop_i) for bl in body_lines if not (has_break and bl.strip().lower() == 'break')]
@@ -1095,7 +1087,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
          j, slice_toks = _match_bracket(toks, j)
          slice_str = _tok_str(slice_toks)
          hi_str, lo_str = slice_str.split(':')
-          hi_val, lo_val = _const_int(hi_str), _const_int(lo_str)
+          hi_val, lo_val = int(eval(hi_str.strip())), int(eval(lo_str.strip()))
          if j < len(toks) and toks[j].type == 'DOT': j += 2  # skip .type suffix
          if j < len(toks) and toks[j].type == 'EQUALS': j += 1
          ln = parse_tokens(lane_toks, env, funcs)
@@ -1153,7 +1145,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
        hi_str = ' '.join(t.val for t in toks[bracket_start:colon_pos] if t.type != 'EOF')
        lo_str = ' '.join(t.val for t in toks[colon_pos+1:j] if t.type != 'EOF')
        try:
-          hi_val, lo_val = _const_int(hi_str), _const_int(lo_str)
+          hi_val, lo_val = int(eval(hi_str)), int(eval(lo_str))
          hi, lo = max(hi_val, lo_val), min(hi_val, lo_val)
          j += 1
          if j < len(toks) and toks[j].type == 'DOT': j += 2
@@ -1167,7 +1159,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
          block_assigns[var] = env[var] = _set_bits(old, _val_to_bits(val), hi - lo + 1, lo)
          i += 1
          continue
-        except (ValueError, SyntaxError): pass  # non-constant slice bounds - fall through to other statement forms
+        except Exception: pass
      elif toks[1].type == 'LBRACKET':  # bit index: var[expr] (only for var[...], not var.type[...])
        existing = block_assigns.get(var, env.get(var))
        if existing is not None and isinstance(existing, UOp) and \
@@ -1368,25 +1360,3 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic

 def parse_expr(expr: str, env: dict[str, VarVal], funcs: dict | None = None) -> UOp:
  return parse_tokens(tokenize(expr.strip().rstrip(';')), env, funcs)
-
-def parse_pcode(pcode: str, srcs: dict[str, UOp | int] | None = None) -> tuple[dict, list]:
-  env: dict = srcs.copy() if srcs else {}
-  assigns: list[tuple[str, UOp]] = []
-  raw_lines = [l.strip().rstrip(';') for l in pcode.split('\n') if l.strip() and not l.strip().startswith('//')]
-  # TODO: pcode.py should tokenize full pcode string instead of line-by-line, then this hack can be removed
-  lines: list[str] = []
-  for l in raw_lines:
-    if lines and re.search(r'(&&|\|\||[&|+\-*/^])\s*$', lines[-1]): lines[-1] = lines[-1] + ' ' + l
-    else: lines.append(l)
-  _, final, _ = parse_block(lines, 0, env, assigns=assigns)
-  sliced = set(d.split('[')[0] for d, _ in assigns if '[' in d)
-  for var, val in final.items():
-    if var in ['D0', 'S0', 'SCC', 'VCC', 'EXEC', 'PC', 'RETURN_DATA', 'VDATA'] and isinstance(val, UOp):
-      if var in sliced and not any(re.match(rf'{var}\.\w+\s*=', l) for l in lines): continue
-      for l in lines:
-        if (m := re.match(rf'{var}\.(\w+(?:\[\w+\])?)', l)):
-          assigns.append((f'{var}.{m.group(1)}', val))
-          break
-      else: assigns.append((var, val))
-  return env, assigns
-
@@ -1,100 +0,0 @@
-# SQTT trace encoder for the emulator (the decoder lives in tinygrad/renderer/amd/sqtt.py).
-# run_asm emits packets inline as instructions execute; finished traces end up in emu.sqtt_traces.
-from __future__ import annotations
-from tinygrad.renderer.amd.dsl import Inst
-from tinygrad.renderer.amd.sqtt import (_build_decode_tables, PACKET_TYPES_RDNA3, PacketType, InstOp,
-                                        LAYOUT_HEADER, WAVESTART, WAVEEND, INST, IMMEDIATE, VALUINST)
-
-_NIB_COUNTS = {cls: nc for _, (cls, nc, *_) in _build_decode_tables(PACKET_TYPES_RDNA3)[0].items()}
-
-def _emit_nibbles(nibbles: list[int], pkt_cls: type[PacketType], **kwargs):
-  raw = pkt_cls.encoding.default
-  for k, v in kwargs.items(): raw = pkt_cls.__dict__[k].set(raw, v)
-  nibbles.extend((raw >> (i * 4)) & 0xF for i in range(_NIB_COUNTS[pkt_cls]))
-
-def make_encoder():
-  """Build an SQTT trace encoder for the emulator. Returns (emit, finish, finalize)."""
-  from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp as SOPPOp3
-  from tinygrad.runtime.autogen.amd.rdna4.enum import SOPPOp as SOPPOp4
-  from tinygrad.runtime.autogen.amd.rdna3 import ins as ir3
-  from tinygrad.runtime.autogen.amd.rdna4 import ins as ir4
-  from tinygrad.runtime.autogen.amd.cdna import ins as irc
-  import re
-
-  def _kinds(*names: str) -> tuple[type[Inst], ...]:
-    return tuple(getattr(m, n) for m in (ir3, ir4, irc) for n in names if hasattr(m, n))
-  _SOPP, _SMEM, _DS = _kinds('SOPP'), _kinds('SMEM'), _kinds('DS')
-  _GLOBAL, _FLAT, _SCRATCH = _kinds('GLOBAL', 'VGLOBAL'), _kinds('FLAT', 'VFLAT'), _kinds('SCRATCH', 'VSCRATCH')
-  _VALU = _kinds('VOP1', 'VOP2', 'VOP3', 'VOP3P', 'VOP3PX2', 'VOPC', 'VOPD', 'VOP3SD', 'VOP3_SDST', 'VOP1_SDST')
-
-  # SOPP classification sets
-  _SOPP_SKIP = {SOPPOp3.S_ENDPGM.value, SOPPOp3.S_ENDPGM_SAVED.value, SOPPOp3.S_ENDPGM_ORDERED_PS_DONE.value, SOPPOp3.S_DELAY_ALU.value}
-  _SOPP_IMMEDIATE = {SOPPOp3.S_NOP.value, SOPPOp3.S_CLAUSE.value, SOPPOp3.S_WAITCNT.value, SOPPOp3.S_WAITCNT_DEPCTR.value,
-                     SOPPOp3.S_WAIT_IDLE.value, SOPPOp3.S_WAIT_EVENT.value, SOPPOp3.S_SLEEP.value, SOPPOp3.S_SET_INST_PREFETCH_DISTANCE.value}
-  for _op in (SOPPOp4.S_WAIT_ALU, SOPPOp4.S_WAIT_LOADCNT, SOPPOp4.S_WAIT_STORECNT, SOPPOp4.S_WAIT_SAMPLECNT,
-              SOPPOp4.S_WAIT_BVHCNT, SOPPOp4.S_WAIT_EXPCNT, SOPPOp4.S_WAIT_DSCNT, SOPPOp4.S_WAIT_KMCNT,
-              SOPPOp4.S_WAIT_LOADCNT_DSCNT, SOPPOp4.S_WAIT_STORECNT_DSCNT):
-    _SOPP_IMMEDIATE.add(_op.value)
-  _SOPP_BARRIER = {SOPPOp3.S_BARRIER.value}
-  if hasattr(SOPPOp4, 'S_BARRIER_WAIT'): _SOPP_BARRIER.add(SOPPOp4.S_BARRIER_WAIT.value)
-  if hasattr(SOPPOp4, 'S_BARRIER_LEAVE'): _SOPP_BARRIER.add(SOPPOp4.S_BARRIER_LEAVE.value)
-  _SOPP_BRANCH = {SOPPOp3.S_BRANCH.value, SOPPOp3.S_CBRANCH_SCC0.value, SOPPOp3.S_CBRANCH_SCC1.value,
-                  SOPPOp3.S_CBRANCH_VCCZ.value, SOPPOp3.S_CBRANCH_VCCNZ.value,
-                  SOPPOp3.S_CBRANCH_EXECZ.value, SOPPOp3.S_CBRANCH_EXECNZ.value}
-
-  # VALU sub-classification patterns
-  _VALUT_4_RE = re.compile(r'V_(EXP|LOG|RCP|RSQ|SQRT|SIN|COS|CEIL|FLOOR|TRUNC|RNDNE|FRACT|FREXP)_')
-  _VALUB_2_RE = re.compile(r'V_(LSHLREV|LSHRREV|ASHRREV)_(B|I)64')
-  _VALUB_4_RE = re.compile(r'V_MAD_(U|I)64')
-  _VALUB_16_RE = re.compile(r'V_\w+_F64')
-
-  def _valu_op(op_name: str) -> InstOp|None:
-    if 'CMPX' in op_name: return InstOp.VALU1_WR_EXEC
-    if _VALUB_2_RE.search(op_name): return InstOp.VALUB_2
-    if _VALUB_4_RE.search(op_name): return InstOp.VALUB_4
-    if _VALUB_16_RE.search(op_name): return InstOp.VALUB_16
-    if _VALUT_4_RE.search(op_name): return InstOp.VALUT_4
-    return None
-
-  def _mem_op(t: type[Inst], op_name: str) -> InstOp:
-    is_store = "STORE" in op_name
-    if issubclass(t, _DS): return InstOp.LDS_WR_2 if is_store else InstOp.LDS_RD
-    if issubclass(t, _GLOBAL): return InstOp.SGMEM_WR_2 if is_store else InstOp.SGMEM_RD_1
-    if issubclass(t, _FLAT) or issubclass(t, _SCRATCH): return InstOp.FLAT_WR_3 if is_store else InstOp.FLAT_RD_2
-    return InstOp.SALU
-
-  nibbles: list[int] = []
-  started: set[int] = set()
-  _emit_nibbles(nibbles, LAYOUT_HEADER, layout=3, sel_a=6)
-
-  def emit(wave_id: int, inst: Inst, branch_taken: bool|None):
-    """Emit an SQTT packet for one executed instruction."""
-    w = wave_id & 0x1F
-    if wave_id not in started:
-      _emit_nibbles(nibbles, WAVESTART, delta=1, simd=0, wgp=0, wave=w, id7=wave_id)
-      started.add(wave_id)
-    inst_type, inst_op, op_name = type(inst), inst.op.value if hasattr(inst, 'op') else 0, inst.op.name if hasattr(inst, 'op') else ""
-    if issubclass(inst_type, _SOPP):
-      if inst_op in _SOPP_SKIP: return
-      if inst_op in _SOPP_IMMEDIATE: _emit_nibbles(nibbles, IMMEDIATE, delta=1, wave=w)
-      elif inst_op in _SOPP_BARRIER: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.BARRIER)
-      elif inst_op in _SOPP_BRANCH: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.JUMP if branch_taken else InstOp.JUMP_NO)
-      else: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.SALU)
-    elif issubclass(inst_type, _VALU):
-      if (op := _valu_op(op_name)) is None: _emit_nibbles(nibbles, VALUINST, delta=1, wave=w)
-      else: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=op)
-    elif issubclass(inst_type, _SMEM): _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.SMEM_RD)
-    else: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=_mem_op(inst_type, op_name))
-
-  def finish(wave_id: int):
-    """Emit WAVEEND for a completed wave."""
-    if wave_id in started: _emit_nibbles(nibbles, WAVEEND, delta=1, simd=0, wgp=0, wave=wave_id & 0x1F)
-
-  def finalize() -> bytes:
-    """Pad and return the encoded SQTT blob."""
-    while len(nibbles) % 2 != 0: nibbles.append(0)
-    nibbles.extend([0] * 32)
-    while len(nibbles) % 64 != 0: nibbles.append(0)
-    return bytes(nibbles[i] | ((nibbles[i + 1] if i + 1 < len(nibbles) else 0) << 4) for i in range(0, len(nibbles), 2))
-
-  return emit, finish, finalize
@@ -160,7 +160,7 @@ class MockUSB3:
    elif request == 0xE5:
      self.state._xram_write_byte(value, index)
    elif request == 0xF2:
-      op = ("sram_read" if value & 0x8000 else "sram_write", 0xF000 + (index & 0xFF) * 0x4000, (value & 0x7FFF) * 512)
+      op = ("sram_read" if value & 0x8000 else "sram_write", 0xF000, (value & 0x7FFF) * 512)
      if value & 0x8000: self._bulk_read_op = op
      else: self._bulk_write_op = op
    elif request == 0xF0:
@@ -193,33 +193,19 @@ class MockUSB3:
    op, address, size = self._bulk_write_op
    assert len(data) == size
    if op == "sram_write":
-      ctrl, (host_addr, region_size) = next((ca, r) for ca, r in self.state._dma_regions.items() if ca <= address < ca + r[1])
-      ctypes.memmove(host_addr + (address - ctrl), data, min(len(data), region_size - (address - ctrl)))
-      self.state.driver._emulate_execute()  # landed data may un-stall a ring polling on it (e.g. copyin sentinels)
+      host_addr, region_size = self.state._dma_regions[address]
+      ctypes.memmove(host_addr, data, min(len(data), region_size))
    elif op == "pcie_write": self.state._pcie_write(address, data)
    else: raise RuntimeError(f"cannot bulk write for {op}")
    self._bulk_write_op = None

-  def bulk_write_async(self, payload:memoryview, timeout:int=10000) -> int:  # the mock completes transfers synchronously
-    self.bulk_write(bytes(payload), timeout)
-    return 0
-
-  def control_write_async(self, request:int, value:int=0, index:int=0, data:bytes=b"", timeout:int=1000) -> int:
-    self.control_write(request, value, index, data, timeout)
-    return 0
-
-  def control_read_async(self, request:int, length:int, value:int=0, index:int=0, timeout:int=1000) -> tuple[int, memoryview]:
-    return 0, self.control_read(request, length, value, index, timeout)
-
-  def bulk_wait(self, tag:int): pass
-
  def bulk_read(self, length:int, timeout:int=1000) -> memoryview:
    assert self._bulk_read_op is not None
    op, address, size = self._bulk_read_op
    assert length == size
    if op == "sram_read":
-      ctrl, (host_addr, region_size) = next((ca, r) for ca, r in self.state._dma_regions.items() if ca <= address < ca + r[1])
-      data = bytes((ctypes.c_ubyte * min(length, region_size - (address - ctrl))).from_address(host_addr + (address - ctrl)))
+      host_addr, region_size = self.state._dma_regions[address]
+      data = bytes((ctypes.c_ubyte * min(length, region_size)).from_address(host_addr))
    elif op == "pcie_read": data = self.state._pcie_read(address, length)
    else: raise RuntimeError(f"cannot bulk read for {op}")
    self._bulk_read_op = None
@@ -0,0 +1,45 @@
+import unittest
+from tinygrad import UOp, dtypes
+from tinygrad.uop.ops import shape_to_shape_arg, ParamArg, Ops, AddrSpace
+
+def placeholder(shape, dtype, slot):
+  return UOp(Ops.PARAM, dtype, (shape_to_shape_arg(shape),), arg=ParamArg(slot, AddrSpace.GLOBAL))
+
+class TestBitcastSpec(unittest.TestCase):
+  def test_bitcast_no_shape_change(self):
+    pl = placeholder((10,10), dtypes.int, 0)
+    out = pl.bitcast(dtypes.float)
+    self.assertEqual(out.shape, (10,10))
+
+  def test_bitcast_increase_shape(self):
+    pl = placeholder((10,10), dtypes.int, 0)
+    out = pl.bitcast(dtypes.short)
+    self.assertEqual(out.shape, (10,20))
+
+  def test_bitcast_decrease_shape(self):
+    pl = placeholder((10,10), dtypes.int, 0)
+    out = pl.bitcast(dtypes.long)
+    self.assertEqual(out.shape, (10,5))
+
+  def test_bitcast_remove_ones(self):
+    pl = placeholder((10,2), dtypes.int, 0)
+    out = pl.bitcast(dtypes.long)
+    self.assertEqual(out.shape, (10,1))
+
+  def test_bitcast_remove_ones_full(self):
+    pl = placeholder((2,), dtypes.int, 0)
+    out = pl.bitcast(dtypes.long)
+    self.assertEqual(out.shape, (1,))
+
+  def test_bitcast_add_ones_full(self):
+    pl = placeholder((), dtypes.long, 0)
+    out = pl.bitcast(dtypes.int)
+    self.assertEqual(out.shape, (2,))
+
+  def test_bitcast_add_ones_full_uchar(self):
+    pl = placeholder((), dtypes.long, 0)
+    out = pl.bitcast(dtypes.uchar)
+    self.assertEqual(out.shape, (8,))
+
+if __name__ == '__main__':
+  unittest.main()
@@ -1,10 +1,39 @@
 import unittest, itertools, math
-from tinygrad import dtypes, Context
-from tinygrad.dtype import DType, ConstType
+from tinygrad import Tensor, dtypes, Context
+from tinygrad.dtype import DType, ConstType, truncate
 from tinygrad.uop.ops import Ops, UOp
 from test.helpers import full_rewrite
 import numpy as np

+def _check_ast_count(desired_count:int, t:Tensor):
+  # NOTE: this has side effect because everything can be scheduled only once
+  linear = t.schedule_linear()
+  asts = [s for s in linear.src if s.src[0].op is Ops.SINK]
+  len(asts)
+  # NOT SUPPORTED ANYMORE
+  #assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
+
+class TestUnaryOpsConstFolding(unittest.TestCase):
+  def test_all_consts_ops(self):
+    _check_ast_count(0, Tensor.ones(4).exp())
+    _check_ast_count(0, Tensor.ones(4).sqrt())
+    _check_ast_count(0, Tensor.ones(4) + Tensor.ones(4))
+    _check_ast_count(0, Tensor.ones(4) / Tensor.ones(4))
+
+  def test_cast(self):
+    _check_ast_count(0, Tensor.ones(4).cast(dtypes.int16))
+    _check_ast_count(0, Tensor.full(4, fill_value=-1).cast(dtypes.uint16))
+
+  def test_neg_folding(self):
+    _check_ast_count(0, Tensor([1, 2, 3]).mul(-1).neg())
+    _check_ast_count(0, Tensor([1, 2, 3]).neg().mul(-1))
+    _check_ast_count(0, Tensor([1, 2, 3]).neg().neg())
+
+  def test_neg_realized_no_fold(self):
+    x = Tensor.randn(32, 32)
+    x = x.clip(0, 1).realize()
+    _check_ast_count(1, x.neg())
+
 class TestWeakConstFolding(unittest.TestCase):
  def test_weakint_math(self):
    out = (UOp.const(2**40) + UOp.const(2**40)).simplify()
@@ -22,18 +51,87 @@ class TestWeakConstFolding(unittest.TestCase):
  def test_invalid_poison(self):
    self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid)

-class TestBitcastConstFolding(unittest.TestCase):
-  def test_out_of_range_source_value(self):
-    for val, src_dt, dst_dt, bits in ((3000000000, dtypes.int32, dtypes.uint32, 3000000000),
-                                      (70000, dtypes.int16, dtypes.uint16, 4464),
-                                      (-5, dtypes.uint32, dtypes.int32, -5)):
-      self.assertEqual(UOp.const(val, src_dt).bitcast(dst_dt).simplify().val, bits)
+  def test_cast_commits_to_dtype_grid(self):
+    # committing a weak const to a stated width puts the value on that width's grid, same as storage packing and native compilers
+    v = 1/123008  # not representable in float16
+    out = UOp.const(v).cast(dtypes.half).simplify()
+    self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.half, truncate[dtypes.half](v)))
+    self.assertNotEqual(out.val, v)
+    # the grid commit preserves the sign of zero
+    self.assertEqual(math.copysign(1, UOp.const(-0.0).cast(dtypes.half).simplify().val), -1)
+    # observable at tensor level: the const-folded comparison agrees with the committed value
+    self.assertTrue((Tensor(-3.2).cast(dtypes.float32) <= truncate[dtypes.float32](-3.2)).item())

+class TestBinaryOpsConstFolding(unittest.TestCase):
+  def test_add_literal_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + 0)
+  def test_add_tensor_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(4))
+  def test_literal_zero_add(self):
+    _check_ast_count(0, 0 + Tensor([1.0, 2, 3, 4]))
+  def test_tensor_zero_add(self):
+    _check_ast_count(0, Tensor.zeros(4) + Tensor([1.0, 2, 3, 4]))
+
+  def test_sub_literal_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) - 0)
+  def test_sub_tensor_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) - Tensor.zeros(4))
+
+  def test_mul_literal_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * 0)
+  def test_mul_tensor_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.zeros(4))
+  def test_literal_zero_mul(self):
+    _check_ast_count(0, 0 * Tensor([1.0, 2, 3, 4]) * 0)
+  def test_tensor_zero_mul(self):
+    _check_ast_count(0, Tensor.zeros(4) * Tensor([1.0, 2, 3, 4]))
+
+  def test_mul_literal_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * 1)
+  def test_mul_tensor_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.ones(4))
+  def test_literal_one_mul(self):
+    _check_ast_count(0, 1 * Tensor([1.0, 2, 3, 4]))
+  def test_tensor_one_mul(self):
+    _check_ast_count(0, Tensor.ones(4) * Tensor([1.0, 2, 3, 4]))
+
+  def test_bool_tensor_mul_bool(self):
+    _check_ast_count(0, Tensor([True, False]) * True)
+    _check_ast_count(0, Tensor([True, False]) * False)
+  def test_bool_mul_bool_tensor(self):
+    _check_ast_count(0, True * Tensor([True, False]))
+    _check_ast_count(0, False * Tensor([True, False]))
+
+  def test_div_literal_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) / 1)
+  def test_div_tensor_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) / Tensor.ones(4))
+
+  def test_floordiv_literal_one(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) // 1)
+  def test_floordiv_tensor_one(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) // Tensor.ones(4, dtype=dtypes.int32))
+
+  def test_pow_literal_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** 0)
+  def test_pow_tensor_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** Tensor.zeros(4))
+
+  def test_pow_literal_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** 1)
+  def test_pow_tensor_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** Tensor.ones(4))
+  def test_literal_one_pow(self):
+    _check_ast_count(0, 1 ** Tensor([1.0, 2, 3, 4]))
+  def test_tensor_one_pow(self):
+    _check_ast_count(0, Tensor.ones(4) ** Tensor([1.0, 2, 3, 4]))
+
+class TestBitcastConstFolding(unittest.TestCase):
  def test_scalar_bitcast(self):
    def t(cases: dict[DType, ConstType]):
      for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()):
        if not math.isnan(from_v):
-          r = UOp.const(from_v, from_dt).bitcast(to_dt).simplify()
+          r = full_rewrite(UOp.const(from_v, from_dt).bitcast(to_dt).sink()).src[0]
          self.assertEqual(r.op, Ops.CONST, msg:=f"{from_dt} -> {to_dt} ({from_v} -> {to_v})")
          self.assertEqual(r.dtype, to_dt, msg)
          np.testing.assert_equal(r.val, to_v, msg)
@@ -57,9 +155,24 @@ class TestBitcastConstFolding(unittest.TestCase):

  def test_vec_bitcast(self):
    with Context(SPEC=0):
-      result = full_rewrite(UOp.const((-1, -2**31, 75), dtypes.int32).bitcast(dtypes.uint32).sink())
-      expected = full_rewrite(UOp.const((2**32-1, 2**31, 75), dtypes.uint32).sink())
-    self.assertEqual(result.src, expected.src)
+      srcs = full_rewrite(UOp.const((-1, -2**31, 75), dtypes.int32).bitcast(dtypes.uint32).sink()).src
+    self.assertTrue(all(r.op is Ops.CONST and r.dtype == dtypes.uint32 for r in srcs))
+    self.assertEqual(tuple(x.val for x in srcs), (2**32-1, 2**31, 75))
+
+# folds advance indexing into basic indexing
+class TestIndexingConstFolding(unittest.TestCase):
+  def test_scalar_index(self):
+    t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize()
+    _check_ast_count(1, t[:,:,Tensor(1),:])
+    _check_ast_count(1, t[:,:,Tensor(1)+2,:])
+    _check_ast_count(1, t[:,:,Tensor(1),Tensor(0)])
+
+  def test_const_tensor_index(self):
+    # TODO: these can be 0, implement const tensor folded indexing
+    t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize()
+    _check_ast_count(1, t[:,:,Tensor.ones(2,1,dtype=dtypes.int),:])
+    _check_ast_count(1, t[:,:,Tensor.ones(1,2,dtype=dtypes.int)+2,:])
+    _check_ast_count(1, t[:,:,Tensor.ones(1,1,dtype=dtypes.int),Tensor.zeros(2,1,2,dtype=dtypes.int)])

 if __name__ == '__main__':
  unittest.main()
@@ -51,6 +51,10 @@ class TestHelpers(unittest.TestCase):
    assert dtypes.is_float(dtypes.fp8e4m3)
    assert dtypes.is_float(dtypes.fp8e5m2)

+  @given(strat.sampled_from([d for d in DTYPES_DICT.values() if dtypes.is_float(d) or dtypes.is_int(d)]))
+  def test_scalar(self, dtype):
+    assert dtype.scalar() == dtype
+
  def test_from_py(self):
    assert dtypes.from_py(True) == dtypes.bool
    assert dtypes.from_py(Invalid) == dtypes.bool
@@ -106,8 +110,7 @@ class TestHelpers(unittest.TestCase):

  def test_float_to_bf16(self):
    max_bf16 = torch.finfo(torch.bfloat16).max
-    for a in [1, 1.1, 1234, 23456, -777.777, max_bf16, max_bf16 * 1.00001, -max_bf16, -max_bf16 * 1.00001,
-              max_bf16 * 2, -max_bf16 * 2, math.inf, -math.inf]:
+    for a in [1, 1.1, 1234, 23456, -777.777, max_bf16, max_bf16 * 1.00001, -max_bf16, -max_bf16 * 1.00001, math.inf, -math.inf]:
      self.assertEqual(float_to_bf16(a), torch.tensor([a], dtype=torch.bfloat16).item())
    self.assertTrue(math.isnan(float_to_bf16(math.nan)))

@@ -1,7 +1,6 @@
 import unittest, subprocess, platform
 from tinygrad.runtime.support.compiler_cpu import ClangCompiler
 from tinygrad.runtime.support.elf import elf_loader
-from tinygrad.runtime.support.c import DLL

 class TestElfLoader(unittest.TestCase):
  def test_load_clang_jit_strtab(self):
@@ -24,7 +23,7 @@ class TestElfLoader(unittest.TestCase):
      }
    '''
    with self.assertRaisesRegex(RuntimeError, 'evil_external_function'):
-      elf_loader(ClangCompiler([{'AMD64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine(), m), "native"]).compile(src))
+      ClangCompiler([{'AMD64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine(), m), "native"]).compile(src)
  def test_link(self):
    src = '''
      float powf(float, float); // from libm
@@ -33,7 +32,7 @@ class TestElfLoader(unittest.TestCase):
    args = ('-x', 'c', '-c', '-target', f'{platform.machine()}-none-unknown-elf', '-march=native', '-fPIC', '-O2', '-ffreestanding', '-nostdlib')
    obj = subprocess.check_output(('clang',) + args + ('-', '-o', '-'), input=src.encode())
    with self.assertRaisesRegex(RuntimeError, 'powf'): elf_loader(obj)
-    elf_loader(obj, link_libs=[DLL('m', 'm')])
+    elf_loader(obj, link_libs=['m'])

 if __name__ == '__main__':
  unittest.main()
@@ -1,7 +1,8 @@
 import unittest, math
 from tinygrad import dtypes
+from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import all_same, Context
-from tinygrad.uop.ops import GroupOp, UOp, Ops, PatternMatcher, TrackedPatternMatcher, UPat
+from tinygrad.uop.ops import GroupOp, UOp, Ops, exec_alu, PatternMatcher, TrackedPatternMatcher, UPat
 from test.helpers import full_rewrite
 from hypothesis import given, strategies as strat

@@ -10,14 +11,125 @@ from hypothesis import given, strategies as strat
 def apply_rewrite(expr):
  return full_rewrite(expr.sink()).src[0]

-def const_value(uop:UOp):
-  if uop.op is Ops.CAST: uop = uop.src[0]
-  assert uop.op is Ops.CONST
-  return uop.val
+@Context(SPEC=0)
+def apply_rewrite_values(expr):
+  srcs = full_rewrite(expr.sink()).src
+  if len(srcs) == 1:
+    if srcs[0].op is Ops.CONST: return (srcs[0].val,)
+    if srcs[0].op is Ops.STACK: return tuple(s.val for s in srcs[0].src)
+  return tuple(s.val for s in srcs)
+
+def evaluate_uop(uop, variables):
+  if uop.op == Ops.CONST:
+    return uop.val
+  elif uop.op == Ops.PARAM and uop.arg.addrspace is AddrSpace.ALU:
+    return variables[uop.expr]
+  elif uop.op in GroupOp.ALU:
+    src_values = [evaluate_uop(src, variables) for src in uop.src]
+    return exec_alu(uop.op, uop.dtype, src_values)
+  else:
+    raise NotImplementedError(f"Unsupported UOp {uop.op}")
+
+class TestArithmeticSimplifications(unittest.TestCase):
+  def test_full_graph_rewrite_division_by_zero(self):
+    optimized_div_uop = apply_rewrite(UOp.const(10.0) / UOp.const(0.0))
+    self.assertEqual(optimized_div_uop.op, Ops.CONST)
+    self.assertTrue(math.isinf(optimized_div_uop.val) or math.isnan(optimized_div_uop.val))
+
+  def test_full_graph_rewrite_redundant_operations(self):
+    optimized_uop = apply_rewrite((UOp.const(10.0) + UOp.const(0.0)) * UOp.const(1.0))
+    self.assertEqual(optimized_uop.op, Ops.CONST)
+    self.assertEqual(optimized_uop.val, 10.0)
+
+  def test_full_graph_rewrite_large_graph(self):
+    prev_uop = UOp.const(0)
+    for i in range(1, 101):
+      prev_uop += UOp.const(i)
+    optimized_uop = apply_rewrite(prev_uop)
+    self.assertEqual(optimized_uop.op, Ops.CONST)
+    self.assertEqual(optimized_uop.val, sum(range(1, 101)))
+
+  def test_full_graph_rewrite_division_by_one(self):
+    optimized_uop = apply_rewrite(UOp.const(42.0) / UOp.const(1.0))
+    self.assertEqual(optimized_uop.op, Ops.CONST)
+    self.assertEqual(optimized_uop.val, 42.0)
+
+  def test_full_graph_rewrite_modulo_by_one(self):
+    optimized_uop = apply_rewrite(UOp.const(42) % UOp.const(1))
+    self.assertEqual(optimized_uop.op, Ops.CONST)
+    self.assertEqual(optimized_uop.val, 0)
+
+
+class TestFoldingAndReduction(unittest.TestCase):
+  @unittest.skip("reduce is removed now")
+  def test_full_graph_rewrite_constant_reduction_folding(self):
+    const1 = UOp.const(5)
+    const2 = UOp.const(10)
+    const3 = UOp.const(20)
+    optimized_sink = apply_rewrite((const1 + const2 + const3).reduce(Ops.ADD))
+    expected_sum = 5 + 10 + 20
+    self.assertEqual(optimized_sink.val, expected_sum)
+
+  @unittest.skip("reduce is removed now")
+  def test_full_graph_rewrite_reduction_with_unused_range(self):
+    const1 = UOp.const(15)
+    const2 = UOp.const(25)
+    rng = UOp.range(10, idx=0)
+    optimized_sink = apply_rewrite((const1 + const2).reduce(Ops.ADD, rng))
+    expected_sum = 10 * (15 + 25)
+    self.assertEqual(optimized_sink.val, expected_sum)
+
+  @unittest.skip("currently failing")
+  def test_full_graph_rewrite_range_reduction(self):
+    simple_range = UOp.range(5, idx=0)
+    optimized_sink = apply_rewrite(simple_range.reduce(Ops.ADD, simple_range))
+    expected_sum = sum(range(5))
+    self.assertEqual(optimized_sink.val, expected_sum)
+
+  @unittest.skip("currently failing")
+  def test_full_graph_rewrite_simple_reduction_folding(self):
+    simple_range = UOp.range(4, idx=0)
+    add_uop = simple_range + UOp.const(1)
+    optimized_sink = apply_rewrite(add_uop.reduce(Ops.ADD, simple_range))
+    expected_sum = sum(i + 1 for i in range(4))
+    self.assertEqual(optimized_sink.val, expected_sum)
+
+  @unittest.skip("currently failing")
+  def test_full_graph_rewrite_nested_loop_collapse(self):
+    outer_range = UOp.range(8, 0)
+    inner_range = UOp.range(4, 1)
+    expr = (outer_range * 10) + inner_range
+    optimized_reduce_uop = apply_rewrite(expr.reduce(Ops.ADD, outer_range, inner_range))
+    self.assertEqual(optimized_reduce_uop.op, Ops.CONST)
+    self.assertEqual(optimized_reduce_uop.val, sum((i * 10) + j for i in range(8) for j in range(4)))
+

 class TestModuloAndDivisionFolding(unittest.TestCase):
+  def test_full_graph_rewrite_modulo_folding_with_define_var(self):
+    # index dtype because div-mod rules only work on index
+    x_var_uop = UOp.variable('x', 0, 100).cast(dtypes.weakint)
+    optimized_mod_uop = apply_rewrite(((x_var_uop * 4) + 2) % 4)
+    self.assertEqual(optimized_mod_uop.op, Ops.CONST)
+    self.assertEqual(optimized_mod_uop.val, 2)
+
+  def test_full_graph_rewrite_division_folding_with_define_var(self):
+    # index dtype because div-mod rules only work on index
+    n_var_uop = UOp.variable('n', 1, 1000).cast(dtypes.weakint)
+    optimized_div_uop = apply_rewrite((n_var_uop * 6) // 3)
+    self.assertEqual(optimized_div_uop.op, Ops.MUL)
+    self.assertEqual(optimized_div_uop.src[1].val, 2)
+
+  def test_full_graph_rewrite_complex_mod_div_folding(self):
+    # index dtype because div-mod rules only work on index
+    k_var_uop = UOp.variable('k', 0, 50).cast(dtypes.weakint)
+    optimized_div_uop = apply_rewrite(((k_var_uop * 12 + 8) % 6) // 2)
+    self.assertEqual(optimized_div_uop.op, Ops.CONST)
+    self.assertEqual(optimized_div_uop.val, 1)
+
  def test_graph_rewrite_div_folding_bug(self):
-    lhs = UOp.stack(*(UOp.special(32, 'lidx0'),)*4) + UOp.const((0, 256, 512, 768))
+    lhs = UOp(Ops.ADD, src=(
+      UOp(Ops.STACK, arg=None, src=(UOp(Ops.SPECIAL, src=(UOp.const(32),), arg='lidx0'),)*4),
+      UOp.const((0, 256, 512, 768))))
    rhs = UOp.const((2,)*4)
    unopt = lhs<rhs
    opt = apply_rewrite(unopt)
@@ -25,31 +137,74 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
    print(opt)
    if opt.op is Ops.STACK: self.assertFalse(all_same(opt.src))

+  def test_full_graph_rewrite_modulo_large_divisor(self):
+    # index dtype because div-mod rules only work on index
+    x_var_uop = UOp.variable('x', 1, 5)
+    self.assertIs(apply_rewrite(x_var_uop.cast(dtypes.weakint) % 10).render(simplify=False), x_var_uop.render(simplify=False))
+
+  def test_full_graph_rewrite_division_with_remainder(self):
+    x_var_uop = UOp.variable('x', 7, 9)
+    optimized_sink = apply_rewrite(x_var_uop // 2)
+    for x_value in range(7, 10):
+      self.assertEqual(x_value // 2, evaluate_uop(optimized_sink, {'x': x_value}))
+
+  def test_full_graph_rewrite_complex_mod_div_expression(self):
+    x_var_uop = UOp.variable('x', 1, 10)
+    optimized_sink = apply_rewrite(((x_var_uop * 5) % 3) // 2)
+    for x_value in range(1, 11):
+      original_result = ((x_value * 5) % 3) // 2
+      optimized_result = evaluate_uop(optimized_sink, {'x': x_value})
+      self.assertEqual(original_result, optimized_result)
+
+
 class TestEdgeCasesAndSpecialOperations(unittest.TestCase):
  def test_full_graph_rewrite_transcendental_edge_cases(self):
    optimized_sink = full_rewrite(UOp.const(-1.0).log2().sink(UOp.const(0.0).reciprocal()))
    optimized_log2_neg, optimized_recip_zero = optimized_sink.src
-    log2_neg, recip_zero = const_value(optimized_log2_neg), const_value(optimized_recip_zero)
-    self.assertTrue(math.isnan(log2_neg), f"Expected NaN for log2(-1.0), got {log2_neg}")
-    self.assertTrue(math.isinf(recip_zero) and recip_zero > 0, f"Expected +inf for reciprocal(0.0), got {recip_zero}")
+    self.assertTrue(math.isnan(optimized_log2_neg.val), f"Expected NaN for log2(-1.0), got {optimized_log2_neg.val}")
+    self.assertTrue(math.isinf(optimized_recip_zero.val) and optimized_recip_zero.val > 0,
+                    f"Expected +inf for reciprocal(0.0), got {optimized_recip_zero.val}")
+
+  @unittest.skip("broken")
+  def test_full_graph_rewrite_modulo_negative_dividend(self):
+    x_var_uop = UOp.variable('x', -5, -1)
+    optimized_sink = full_rewrite((x_var_uop % 3).sink())
+    for x_value in range(-5, 0):
+      self.assertEqual(x_value % 3, evaluate_uop(optimized_sink.src[0], {'x': x_value}))
+
+  @unittest.skip("broken")
+  def test_full_graph_rewrite_division_negative_divisor(self):
+    x_var_uop = UOp.variable('x', 1, 5)
+    optimized_sink = full_rewrite((x_var_uop // -2).sink())
+    for x_value in range(1, 6):
+      self.assertEqual(x_value // -2, evaluate_uop(optimized_sink.src[0], {'x': x_value}))

 class TestGEPAndVectorizeRewrite(unittest.TestCase):
  def test_gep_single_element_extraction(self):
    # GEP on a vector dtype to extract a single element
    base_vector = UOp.const((1.0, 2.0, 3.0, 4.0))
-    self.assertIs(apply_rewrite(base_vector.index(2)), apply_rewrite(base_vector.src[2]))
+    self.assertEqual(apply_rewrite(base_vector.index(2)).val, 3.0)

  def test_gep_tuple_extraction(self):
    # GEP on a vector dtype to extract multiple elements as a vector
    base_vector = UOp.const((1.0, 2.0, 3.0, 4.0))
-    self.assertIs(apply_rewrite(UOp.stack(*[base_vector.index(i) for i in (2, 3)])),
-                  apply_rewrite(UOp.stack(base_vector.src[2], base_vector.src[3])))
+    self.assertEqual(list(apply_rewrite_values(UOp.stack(*[base_vector.index(i) for i in (2, 3)]))), [3.0, 4.0])
+
+  def test_gep_on_const_stack(self):
+    # GEP on a const STACK to extract a single element
+    const_stack = UOp.const((1.0, 2.0, 3.0, 4.0))
+    self.assertEqual(apply_rewrite(const_stack.index(2)).val, 3.0)
+
+  def test_gep_tuple_on_const_stack(self):
+    # GEP on a const STACK using a tuple to extract multiple elements
+    const_stack = UOp.const((7.0, 8.0, 9.0, 10.0))
+    self.assertEqual(list(apply_rewrite_values(UOp.stack(*[const_stack.index(i) for i in (1, 3)]))), [8.0, 10.0])

  def test_vectorize_multiple_elements(self):
    # Vectorizing multiple elements using GEP
    base_vector = UOp.const((5.0, 10.0, 15.0, 20.0))
-    vectorized_uop = UOp.stack(*(base_vector.index(i) for i in range(4)))
-    self.assertIs(apply_rewrite(vectorized_uop), apply_rewrite(base_vector))
+    vectorized_uop = UOp(Ops.STACK, src=tuple(base_vector.index(i) for i in range(4)))
+    self.assertEqual(list(apply_rewrite_values(vectorized_uop)), [5.0, 10.0, 15.0, 20.0])


 import inspect
@@ -101,6 +256,16 @@ class TestSubstitute(unittest.TestCase):
    ret = substitute(ret, {a.sin():b})
    self.assertIs(ret, b.sin())

+  # broken due to infinite recursion
+  # NOTE: VIZ hangs and doesn't recover if you click this one
+  @unittest.skip("recursion error no longer raised")
+  def test_assert_inf_recurse(self):
+    a = UOp.variable('a', 0, 10)
+    n1 = a.sin()
+    ret = n1
+    with self.assertRaises(RecursionError):
+      ret = substitute(ret, {n1:n1.sqrt()})
+
  def test_sin_to_sqrt(self):
    a = UOp.variable('a', 0, 10, dtype=dtypes.float)
    n1 = a.sin()
@@ -1,7 +1,6 @@
 import unittest
 from tinygrad.helpers import GlobalCounters
 from tinygrad.nn.datasets import mnist
-from test.helpers import KernelCountException

 class TestDataset(unittest.TestCase):
  def test_dataset_is_realized(self):
@@ -9,7 +8,7 @@ class TestDataset(unittest.TestCase):
    X_train[0].contiguous().realize()
    GlobalCounters.reset()
    X_train[0].contiguous().realize()
-    if GlobalCounters.kernel_count > 1: raise KernelCountException(1, GlobalCounters.kernel_count)  # 0 if SLICE (zero-copy), 1 otherwise
+    self.assertLessEqual(GlobalCounters.kernel_count, 1)  # 0 if SLICE (zero-copy), 1 otherwise

 if __name__ == '__main__':
  unittest.main()
@@ -15,13 +15,17 @@ def simplify_valid_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move
 def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+indexing_simplify, name="simplify_image_idx")

 def get_gated_load_uop(valid:UOp, idx:UOp):
-  return UOp.param(0, dtypes.float, (1024,)).index(idx.valid(valid)).load()
+  return UOp(Ops.LOAD, src=(
+    UOp.param(0, dtypes.float, (1024,)).index(idx.valid(valid)),
+  ))

 def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
-  return UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)).load()
+  return UOp(Ops.LOAD, src=(
+    UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)),
+  ))

-def Special(expr, nmax): return UOp.special(nmax, expr)
-def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax, param=True)
+def Special(expr, nmax): return UOp(Ops.SPECIAL, src=(UOp.const(nmax),), arg=expr)
+def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax)
 def Range(n, nmax): return UOp.range(nmax, n)

 class TestValidIdxSimplification(unittest.TestCase):
@@ -508,7 +512,7 @@ class TestDropTrueGate(unittest.TestCase):
    buf = UOp.param(0, dtypes.int, (1,))
    idx = UOp.const(0)
    true_gate = UOp.const(True)
-    index_with_gate = buf.index(idx.valid(true_gate))
+    index_with_gate = UOp(Ops.INDEX, src=(buf, idx.valid(true_gate)))
    # apply the optimization
    result = graph_rewrite(index_with_gate, sym+indexing_simplify)
    # the True valid should be dropped (INDEX should only have 2 sources)
@@ -520,17 +524,13 @@ class TestRangeShrink(unittest.TestCase):
      result = full_rewrite(sink)
    return [u for u in result.toposort() if u.op is Ops.RANGE]

-  def assert_range_end(self, ranges:list[UOp], end:int):
-    self.assertEqual(len(ranges), 1)
-    with Context(NOOPT=1, SPEC=0): expected = full_rewrite(UOp.const(end, dtypes.int).sink()).src[0]
-    self.assertIs(ranges[0].src[0], expected)
-
  def test_range_shrink_single_guard(self):
    # range 0..203 guarded by r < 4 everywhere -> shrink to 0..3
    r = Range(0, 204)
    load = get_gated_load_uop(r < UOp.const(4), r)
    ranges = self.get_ranges(load.sink())
-    self.assert_range_end(ranges, 4)
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].val, 4)

  def test_range_shrink_picks_max_guard(self):
    # two loads guard the same range with r < 4 and r < 8 -> shrink to max(4, 8) = 8
@@ -538,22 +538,25 @@ class TestRangeShrink(unittest.TestCase):
    load1 = get_gated_load_uop(r < UOp.const(4), r)
    load2 = get_gated_load_uop(r < UOp.const(8), r)
    ranges = self.get_ranges(UOp.sink(load1, load2))
-    self.assert_range_end(ranges, 8)
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].val, 8)

  def test_range_no_shrink_guard_ge_max(self):
    # guard r < 300 with range max 204 -> no shrink (guard doesn't constrain)
    r = Range(0, 204)
    load = get_gated_load_uop(r < UOp.const(300), r)
    ranges = self.get_ranges(load.sink())
-    self.assert_range_end(ranges, 204)
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].val, 204)

  def test_range_no_shrink_when_unguarded_elsewhere(self):
    # one load guards r < 4, but another load uses r without a gate -> no shrink
    r = Range(0, 204)
    load1 = get_gated_load_uop(r < UOp.const(4), r)
-    load2 = UOp.param(1, dtypes.float, (204,)).index(r).load()
+    load2 = UOp(Ops.LOAD, src=(UOp.param(1, dtypes.float, (204,)).index(r),))
    ranges = self.get_ranges(UOp.sink(load1, load2))
-    self.assert_range_end(ranges, 204)
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].val, 204)

  def test_range_no_shrink_when_used_in_reduce(self):
    # range used in both a gated load AND directly in the reduce expression -> no shrink
@@ -561,7 +564,8 @@ class TestRangeShrink(unittest.TestCase):
    gated_load = get_gated_load_uop(r < UOp.const(4), r)
    red = (r.cast(dtypes.float) + gated_load).reduce(r, arg=Ops.ADD)
    ranges = self.get_ranges(red.sink())
-    self.assert_range_end(ranges, 204)
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].val, 204)

  def test_range_shrink_to_single_iteration(self):
    # guard r < 1 shrinks range to 1 -> single iteration, range eliminated entirely
@@ -576,7 +580,8 @@ class TestRangeShrink(unittest.TestCase):
    r = Range(0, 204)
    x = (r < 4).where(UOp.const(1.0), Invalid)
    ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink())
-    self.assert_range_end(ranges, 4)
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].val, 4)

  def test_range_shrink_store_where_invalid_flipped(self):
    # above, but flipped
@@ -584,7 +589,8 @@ class TestRangeShrink(unittest.TestCase):
    r = Range(0, 204)
    x = (r < 4).where(UOp.const(1.0), Invalid)
    ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink())
-    self.assert_range_end(ranges, 4)
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].val, 4)

 if __name__ == '__main__':
  unittest.main()
@@ -69,7 +69,7 @@ class TestIdxUpcast(unittest.TestCase):
    if not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)):
      assert idx.op is Ops.INDEX
      idx_val = idx.src[1]
-      self.assertFalse(idx_val.overflows(idx_val.dtype))
+      self.assertFalse(idx_val.overflows(idx_val.dtype.scalar()))

  # use expand to generate kernel that uses large idx
  def do_op_then_assert(self, dtype: DType, dim1, dim2, dim3):
@@ -171,11 +171,6 @@ class TestTensorConstLike(unittest.TestCase):
    t = Tensor.ones(8, 4).shard(("NULL:0", "NULL:1"), axis=0)
    with self.assertRaises(RuntimeError): t.full_like(5, device="NULL")

-class TestTensorShape(unittest.TestCase):
-  def test_float_shape_raises(self):
-    for dim in (2.0, 2.5):
-      with self.subTest(dim=dim), self.assertRaisesRegex(RuntimeError, "shape must be int"): Tensor.ones(dim)
-
 class TestTensorDevice(unittest.TestCase):
  def test_create_from_single_device_tuple(self):
    (Tensor([1.0], device=(Device.DEFAULT,)) + Tensor([2.0])).realize()
@@ -1,9 +1,10 @@
 import unittest, pytest
 from tinygrad import dtypes, Variable, Device
 from tinygrad.dtype import AddrSpace
+from tinygrad.helpers import DEBUG, Context
 from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType, broadcast_axes, KernelInfo
 from tinygrad.uop.symbolic import sym
-from test.helpers import full_rewrite, to_uops_list
+from test.helpers import to_uops_list
 from tinygrad.codegen import full_rewrite_to_sink

 simple_pm = PatternMatcher([
@@ -13,27 +14,43 @@ simple_pm = PatternMatcher([
  ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.val+c2.val)),
 ])

+def const_values(u:UOp):
+  if u.op is Ops.CONST: return (u.val,)
+  if u.op is Ops.STACK: return tuple(x.val for x in u.src)
+  raise AssertionError(f"expected const-like UOp, got {u.op}")
+
 class TestGraphRewriteConst(unittest.TestCase):
  def test_gep_const(self):
    v1 = UOp.const((0,1,2), dtypes.int)
    v2 = v1.index(1)
    ret = graph_rewrite(v2, sym)
-    self.assertIs(ret, UOp.const(1, dtypes.int))
+    self.assertEqual(ret.dtype, dtypes.int)
+    self.assertEqual(ret.val, 1)

  def test_add_const(self):
    v1 = UOp.const((0,1,2))
    v2 = UOp.const((5,6,7))
-    self.assertIs(graph_rewrite(v1+v2, sym), UOp.const((5,7,9)))
+    ret = graph_rewrite(v1+v2, sym)
+    self.assertEqual(ret.op, Ops.STACK)
+    self.assertEqual(const_values(ret), (5,7,9))
+
+  def test_add_const_lose_v(self):
+    v1 = UOp.const((0,1,2))
+    v2 = UOp.const((2,1,0))
+    ret = graph_rewrite(v1+v2, sym)
+    self.assertEqual(ret.op, Ops.STACK)
+    self.assertEqual(const_values(ret), (2,2,2))

 def xfail_broken_const_wraparound(fn):
  fn = pytest.mark.xfail(reason="const folding does not properly implement modular arithmetic")(fn)
  return unittest.expectedFailure(fn)
 class TestModularWraparound(unittest.TestCase):
  def _test(self, uop:UOp, expected:int):
-    result = uop.simplify()
-    self.assertEqual(result.op, Ops.CONST)
-    self.assertEqual(result.dtype, uop.dtype)
-    self.assertEqual(result.val, expected)
+    results = to_uops_list([uop])
+    self.assertEqual(len(results), 2)  # +1 for SINK
+    self.assertEqual(results[0].op, Ops.CONST)
+    self.assertEqual(results[0].dtype, uop.dtype)
+    self.assertEqual(results[0].val, expected)

  @xfail_broken_const_wraparound
  def test_cast(self):
@@ -140,7 +157,7 @@ class TestGraphRewrite(unittest.TestCase):
    self.assertEqual(nout.val, 3.0)

  def test_depth_2_fold(self):
-    v = UOp.variable("v", 0, 1, dtypes.float, param=True)
+    v = UOp.variable("v", 0, 1, dtypes.float)
    c1 = UOp.const(1.0)
    c2 = UOp.const(2.0)
    nout = graph_rewrite(v+c1+c2, simple_pm)
@@ -174,25 +191,63 @@ class TestGraphRewrite(unittest.TestCase):
      self.assertEqual(len([x for x in sink.toposort() if x.op is Ops.CONST]), 1)

 class TestUOpGraph(unittest.TestCase):
+  def test_add_constant_fold(self):
+    c1 = UOp.const(1.0, dtypes.float)
+    c2 = UOp.const(2.0, dtypes.float)
+    out = c1+c2
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 2)  # +1 for SINK
+    out = uops[-2]
+    self.assertEqual(out.op, Ops.CONST)
+    self.assertEqual(out.val, 3.0)
+
  def test_where_same_fold(self):
    v = UOp.variable('tmp', 0, 1)
    c0 = UOp.const(0)
    vc = v != c0
    c1 = UOp.const(1.0, dtypes.float)
    out = vc.where(c1, c1)
-    self.assertIs(out.simplify(), c1)
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 2)  # +1 for SINK
+    out = uops[-2]
+    self.assertEqual(out.op, Ops.CONST)
+    self.assertEqual(out.val, 1.0)

  def test_where_const_fold(self):
    bf = UOp.const(False)
    c1 = UOp.const(1.0, dtypes.float)
    c2 = UOp.const(2.0, dtypes.float)
    out = bf.where(c1, c2)
-    self.assertIs(out.simplify(), c2)
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 2)  # +1 for SINK
+    out = uops[-2]
+    self.assertEqual(out.op, Ops.CONST)
+    self.assertEqual(out.val, 2.0)

  def test_const_cast(self):
    bf = UOp.const(False)
    out = bf.cast(dtypes.int)
-    self.assertIs(full_rewrite(out.sink()).src[0], full_rewrite(UOp.const(0, dtypes.int).sink()).src[0])
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 2)  # +1 for SINK
+    out = uops[-2]
+    self.assertEqual(out.op, Ops.CONST)
+    self.assertEqual(out.val, 0)
+
+  def test_const_bitcast(self):
+    bf = UOp.const(1.0, dtypes.float)
+    out = bf.bitcast(dtypes.uint32)
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 2)  # +1 for SINK
+    out = uops[-2]
+    self.assertEqual(out.op, Ops.CONST)
+    self.assertEqual(out.val, 0x3F800000)
+
+  @unittest.expectedFailure
+  def test_const_shape_change_bitcast(self):
+    bf = UOp.const(0x3F).cast(dtypes.uint8)
+    out = bf.bitcast(dtypes.half)
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 2)  # +1 for SINK

  def test_devectorize_derives_lane_dtype(self):
    from tinygrad.codegen import do_devectorize
@@ -202,11 +257,66 @@ class TestUOpGraph(unittest.TestCase):
    invalid_lane_mul = next(u for u in out.src[0].toposort() if u.op is Ops.MUL)
    self.assertIs(invalid_lane_mul.dtype, dtypes.bool)

+  @unittest.skip("this test isn't valid uops")
+  def test_noop_vectorize_fold(self):
+    d0 = UOp.param(0, dtypes.float, (1,))
+    idx = UOp.const(0)
+    ld = d0.load(idx, dtype=dtypes.float)
+    vec = UOp(Ops.STACK, dtypes.float, (ld,))
+    x = vec.index(0)
+    alu = UOp(Ops.SQRT, src=(x, ))
+    out = UOp(Ops.STORE, src=(d0, idx, alu))
+    uops = to_uops_list([out])
+    self.assertEqual(len([x for x in uops if x.op is Ops.STACK]), 0)
+
+  @unittest.skip("this test isn't valid uops")
+  def test_gep_vec_fold(self):
+    d0 = UOp.param(0, dtypes.float, (1,))
+    d1 = UOp.param(1, dtypes.float, (1,))
+    d2 = UOp.param(2, dtypes.float, (1,))
+    idx = UOp.const(0)
+    def _test_vec(geps, count=4):
+      vec = UOp(Ops.STACK, dtypes.float, geps)
+      out = d0.index(idx).store(vec)
+      uops = to_uops_list([out])
+      if DEBUG >= 4:
+        from tinygrad import Device
+        print(Device[Device.DEFAULT].renderer.render(uops))
+      return uops[-2].src[-1]  # -2 to skip SINK
+
+    # possible
+    val = d1.index(idx).load(dtype=dtypes.float)
+    xyzw = tuple(val.index(i) for i in range(4))
+    self.assertIs(_test_vec(xyzw).op, Ops.LOAD)
+
+    # unaligned
+    val = d1.index(idx).load(dtype=dtypes.float)
+    wzyx = tuple(val.index(i) for i in reversed(range(4)))
+    self.assertIs(_test_vec(wzyx).op, Ops.STACK)
+
+    # different_size
+    val = d1.index(idx).load(dtype=dtypes.float)
+    xy = tuple(val.index(i) for i in range(2))
+    self.assertIs(_test_vec(xy+xy).op, Ops.STACK)
+    val = d1.index(idx).load(dtype=dtypes.float)
+    xy = tuple(val.index(i) for i in range(2))
+    self.assertIs(_test_vec(xy, count=2).op, Ops.STACK)
+
+    # different vals
+    val1 = d1.index(idx).load(dtype=dtypes.float)
+    val2 = d2.index(idx).load(dtype=dtypes.float)
+    xy1 = tuple(val1.index(i) for i in range(2))
+    xy2 = tuple(val2.index(i) for i in range(2))
+    self.assertIs(_test_vec(xy1+xy2).op, Ops.STACK)
+
  def test_gep_vec_const_fold(self):
    for vec_size in [2, 4, 8]:
      consts = [UOp.const(float(i), dtypes.float) for i in range(vec_size)]
-      vec = UOp.stack(*consts)
-      for i, const in enumerate(consts): self.assertIs(vec.index(i), const)
+      vec = UOp(Ops.STACK, src=tuple(consts))
+      with Context(SPEC=0):
+        uops = to_uops_list([vec.index(i) for i in range(vec_size)])
+        for uop, const in zip(uops, consts):
+          self.assertEqual(uop, const)

  def test_cast_alu_fold(self):
    d0 = UOp.param(0, dtypes.bool, (1,))
@@ -216,7 +326,7 @@ class TestUOpGraph(unittest.TestCase):
    alu = (ld<1).cast(dtypes.bool)
    out = d0.index(idx).store(alu)
    uops = to_uops_list([out])
-    self.assertEqual(len([x for x in uops if x.op is Ops.CAST and x.src[0].op is not Ops.CONST]), 0)
+    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)

  def test_double_cast_fold(self):
    d0 = UOp.param(0, dtypes.float, (1,))
@@ -226,15 +336,20 @@ class TestUOpGraph(unittest.TestCase):
    alu = ld.cast(dtypes.float).cast(dtypes.float)
    out = d0.index(idx).store(alu)
    uops = to_uops_list([out])
-    self.assertEqual(len([x for x in uops if x.op is Ops.CAST and x.src[0].op is not Ops.CONST]), 1)
+    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)

  def test_depth_2_const_fold(self):
-    v = UOp.variable("tmp", 0, 1, dtypes.int, param=True)
+    v = UOp.variable("tmp", 0, 1, dtypes.int)
    c2 = UOp.const(2, dtypes.int)
    c4 = UOp.const(4, dtypes.int)
    vc = v+c2
    out = vc+c4
-    self.assertIs(out.simplify(), (v+UOp.const(6, dtypes.int)).simplify())
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 5)  # +1 for SINK, +1 for the PARAM shape STACK
+    out = uops[-2]  # -2 to skip SINK
+    self.assertEqual(out.op, Ops.ADD)
+    self.assertEqual(out.src[1].op, Ops.CONST)
+    self.assertEqual(out.src[1].val, 6)

  def test_bitcast_to_same_dtype_fold(self):
    for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
@@ -245,8 +360,9 @@ class TestUOpGraph(unittest.TestCase):

  def test_sub_with_cast_folds(self):
    a = Variable("a", 0, 5)
-    out = a.cast(dtypes.int)+(-a).cast(dtypes.int)
-    self.assertIs(full_rewrite(out.sink()).src[0], full_rewrite(UOp.const(0, dtypes.int).sink()).src[0])
+    uops = to_uops_list([a.cast(dtypes.int)+(-a).cast(dtypes.int)])
+    assert uops[0] == UOp.const(0, dtypes.int)
+    assert uops[-1].op == Ops.SINK

  def test_where_on_gated_load_fold(self):
    ridx0 = UOp.range(100, 0)
@@ -255,10 +371,9 @@ class TestUOpGraph(unittest.TestCase):
    w = (ridx0<50).where(ld, 5)
    out = UOp.param(1, dtypes.long, (100,))
    uops = to_uops_list([out.index(ridx0).store(w)])
-    expected = full_rewrite(UOp.const(5, dtypes.long).sink()).src[0]
    for u in uops:
      assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: self.assertIs(u.src[1], expected)
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val==5

  def test_where_on_gated_load_folds_swapped_branches(self):
    ridx0 = UOp.range(100, 0)
@@ -266,10 +381,9 @@ class TestUOpGraph(unittest.TestCase):
    ld = d0.index(ridx0.valid((ridx0<50).logical_not()))
    w = (ridx0<50).where(5, ld)
    uops = to_uops_list([w])
-    expected = full_rewrite(UOp.const(5, dtypes.long).sink()).src[0]
    for u in uops:
      assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD: self.assertIs(u.src[1], expected)
+      if u.op is Ops.LOAD: assert u.src[1].val==5

  def test_where_on_gated_load_with_cast(self):
    ridx0 = UOp.range(100, 0)
@@ -279,10 +393,9 @@ class TestUOpGraph(unittest.TestCase):
    w = (ridx0<50).where(ld, 5.0)
    out = UOp.param(1, dtypes.float, (100,))
    uops = to_uops_list([out.index(ridx0).store(w)])
-    expected = full_rewrite(UOp.const(5, dtypes.int).sink()).src[0]
    for u in uops:
      assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: self.assertIs(u.src[1], expected)
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val == 5

  def test_where_on_casted_gated_load_extra_cond(self):
    ridx0 = UOp.range(100, 0)
@@ -312,10 +425,9 @@ class TestUOpGraph(unittest.TestCase):
    val = (ridx0<50).where(5, ld)
    st = idx.store(val).end(ridx0)
    uops = to_uops_list([st])
-    expected = full_rewrite(UOp.const(5, dtypes.long).sink()).src[0]
    for u in uops:
      assert u.op is not Ops.WHERE
-      if u.op is Ops.STORE: self.assertIs(u.src[1], expected)
+      if u.op is Ops.STORE: assert u.src[1].val==5

  def test_load_idx_becomes_int(self):
    # mnist indexing with split reduceop
@@ -389,6 +501,13 @@ class TestUOpGraph(unittest.TestCase):
    # only the second store happens
    self.assertEqual(len([u for u in uops if u.op is Ops.STORE]), 1)

+  @unittest.skip("this is a uop type error")
+  def test_asserts_bad_gate(self):
+    glbl0 = UOp.param(0, dtypes.int, (1,))
+    idx = UOp.const(0)
+    bad_gate = UOp.const(1)
+    with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, src=(glbl0, idx, UOp.const(42), bad_gate))])
+
  def test_after_end(self):
    r = UOp.range(10, 0)

@@ -456,7 +575,7 @@ class TestConstBufferize(unittest.TestCase):
    from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts
    c = UOp.const(42.0)
    r1 = UOp.range(3, 0)
-    bufferize_with_range = c.bufferize(r1, arg=BufferizeOpts(device="CPU"))
+    bufferize_with_range = UOp(Ops.STAGE, src=(c, r1), arg=BufferizeOpts(device="CPU"))
    self.assertEqual(len(bufferize_with_range.src), 2)  # const + 1 range

    result = graph_rewrite(bufferize_with_range, pm_const_buffer_folding, name='test')
@@ -471,7 +590,7 @@ class TestConstBufferize(unittest.TestCase):
    c = UOp.const(3.14)
    r1 = UOp.range(3, 0)
    r2 = UOp.range(4, 1)
-    bufferize_with_ranges = c.bufferize(r1, r2, arg=BufferizeOpts(device="CPU"))
+    bufferize_with_ranges = UOp(Ops.STAGE, src=(c, r1, r2), arg=BufferizeOpts(device="CPU"))
    self.assertEqual(len(bufferize_with_ranges.src), 3)  # const + 2 ranges

    result = graph_rewrite(bufferize_with_ranges, pm_const_buffer_folding, name='test')
@@ -3,10 +3,10 @@ import unittest, pickle, functools, math
 import z3

 from tinygrad.dtype import dtypes, ConstType, DType, Invalid
+from test.helpers import get_uops
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad.uop.spec import spec_shared, type_verify
-from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load
-from tinygrad.uop.weak import pm_cast_weak
+from tinygrad.uop.symbolic import sym, pm_fold_cast_const, commutative, pm_simplify_valid, pm_move_where_on_load
 from tinygrad.uop.validate import uops_to_z3

 def check_uop_against_string(self, v:UOp, s:str):
@@ -16,8 +16,7 @@ def check_uop_against_string(self, v:UOp, s:str):
  s_eval = graph_rewrite(s_eval, commutative, name="cannonicalize eval")
  self.assertIs(s_eval, v, f"eval did not match simplified: {s_eval} != {v.render()} for {s}")

-def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.weakint):
-  return UOp.variable(name, min_val, max_val, dtype, param=True)
+def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.weakint): return UOp.variable(name,min_val,max_val,dtype)
 def uconst(val): return UOp.const(val)
 def usum(ops): return functools.reduce(lambda x,y: x+y, ops)
 def uand(ops): return functools.reduce(lambda x,y: x*y, ops)
@@ -36,7 +35,7 @@ class TestSymbolic(unittest.TestCase):
    self.assertEqual(solver.check(expr1 != expr2), z3.unsat, "simplified expression not equal to original")

  def helper_test_variable(self, v, n, m, s, test_z3:bool=True):
-    v_simplified = graph_rewrite(v, sym+pm_cast_weak, name="simplify symbolic uop")
+    v_simplified = graph_rewrite(v, sym+pm_fold_cast_const, name="simplify symbolic uop")
    if test_z3: self.check_equal_z3(v, v_simplified)
    nmin, nmax = v_simplified.vmin, v_simplified.vmax
    check_uop_against_string(self, v_simplified, s)
@@ -149,13 +148,6 @@ class TestSymbolic(unittest.TestCase):
  def test_xor_0(self):
    self.helper_test_variable(Variable("a", 0, 8, dtypes.int) ^ 0, 0, 8, "a", test_z3=False)

-  def test_or_0(self):
-    self.helper_test_variable(Variable("a", 0, 8, dtypes.int) | 0, 0, 8, "a", test_z3=False)
-
-  def test_shift_0(self):
-    self.helper_test_variable(Variable("a", 0, 8, dtypes.int) << 0, 0, 8, "a")
-    self.helper_test_variable(Variable("a", 0, 8, dtypes.int) >> 0, 0, 8, "a")
-
  def test_xor_self_inverse(self):
    self.helper_test_variable((Variable("a", 0, 8, dtypes.int) ^ 5) ^ 5, 0, 8, "a", test_z3=False)

@@ -450,7 +442,7 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(uand([uconst(1), Variable("a", 0, 1)]), 0, 1, "a")

  def test_masked_shr_fold(self):
-    x = UOp.variable('x', 0, 255, dtype=dtypes.uint32, param=True)
+    x = UOp.variable('x', 0, 255, dtype=dtypes.uint32)
    self.helper_test_variable((x & -4) >> 2, 0, 63, "(x>>2)")

  def test_bool_or_not_tautology(self):
@@ -491,12 +483,12 @@ class TestSymbolic(unittest.TestCase):

  def test_div_drop_small_terms(self):
    # from openpilot, shouldnt simplify
-    gidx0 = UOp.variable("gidx0", 0, 10, param=True)
-    gidx1 = UOp.variable("gidx1", 0, 10, param=True)
-    lidx0 = UOp.variable("lidx0", 0, 1, param=True)
-    lidx1 = UOp.variable("lidx1", 0, 1, param=True)
-    ridx1005 = UOp.variable("ridx1005", 0, 2, param=True)
-    ridx1006 = UOp.variable("ridx1006", 0, 2, param=True)
+    gidx0 = UOp.variable("gidx0", 0, 10)
+    gidx1 = UOp.variable("gidx1", 0, 10)
+    lidx0 = UOp.variable("lidx0", 0, 1)
+    lidx1 = UOp.variable("lidx1", 0, 1)
+    ridx1005 = UOp.variable("ridx1005", 0, 2)
+    ridx1006 = UOp.variable("ridx1006", 0, 2)
    self.helper_test_variable((lidx1+((gidx1*18)+(ridx1005*18)+(lidx0*162))+(gidx0*2)+(ridx1006*2)+-40)//18, -3, 20,
      "(gidx1+ridx1005+lidx0*9+(gidx0+ridx1006+7)//9+-3)")

@@ -1005,7 +997,7 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(cond.ne(False), 0, 1, "(x<2)")

  def test_bitcast_chain(self):
-    a = UOp.variable("a", 0, 3, dtype=dtypes.int32, param=True)
+    a = UOp.variable("a", 0, 3, dtype=dtypes.int32)
    self.assertIs(graph_rewrite(a.bitcast(dtypes.float32).bitcast(a.dtype), sym), a)

  def test_negation_in_where(self):
@@ -1021,19 +1013,20 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(-a<-b, False, True, "(b<a)")

  def test_where_cast(self):
-    cond = Variable("s", 0, 3, dtypes.int) < 2
+    s = Variable("s", 0, 3, dtypes.int)
+    cond = s < 2
    a = Variable("a", 0, 3, dtypes.int)
-    self.assertIs(graph_rewrite(cond.where(a, a+1).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), (a+1).cast(dtypes.half)))
-    self.assertIs(graph_rewrite(cond.where(a, uconst(2)).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.const(2, dtypes.half)))
-    self.assertIs(graph_rewrite(cond.where(a, UOp.invalid()).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.invalid()))
+    b = Variable("b", 0, 3, dtypes.int)
+    expr = cond.where(a, b).cast(dtypes.half)

-  def test_where_const_gate_keeps_stated_width(self):
-    a = Variable("a", 0, 3, dtypes.half)
-    self.assertIs(graph_rewrite(UOp.const(True, dtypes.bool).where(uconst(0.0), a), sym), UOp.const(0.0, dtypes.half))
-    self.assertIs(graph_rewrite(UOp.const(True, dtypes.bool).where(uconst(0), Variable("i", 0, 3, dtypes.int)), sym), UOp.const(0, dtypes.int))
-    self.assertIs(graph_rewrite(UOp.const(False, dtypes.bool).where(uconst(0.0), a), sym), a)
-    self.assertIs(graph_rewrite(UOp.const(False, dtypes.bool).where(uconst(0.0), UOp.invalid()), sym), UOp.invalid())
-    self.assertIs(graph_rewrite(UOp.const(True, dtypes.bool).where(uconst(0.0), uconst(1)), sym), uconst(0.0))
+    # TODO: copied from render, render does not support cast
+    glbl = UOp.param(0, dtypes.int, (1,))
+    uops = get_uops(UOp(Ops.STORE, src=(glbl.index(UOp.const(0, dtypes.int)), expr)).sink())
+    rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1]
+
+    # the vars are now scalar PARAMs
+    pvar = {u.expr: u for u in rewritten_uop.toposort() if u.op is Ops.PARAM}
+    self.assertEqual(rewritten_uop, (pvar['s']<UOp.const(2, dtypes.int)).where(pvar['a'].cast(dtypes.half), pvar['b'].cast(dtypes.half)))

  def test_where_merge_branches(self):
    cond1 = Variable("s", 0, 10) < 6
@@ -1187,7 +1180,7 @@ class TestSymbolicVariables(unittest.TestCase):
    assert (a//4 + a//6).variables() == [a]

  def test_variable_min_eq_max_bind_folds(self):
-    b = UOp.variable("x", 1, 1).bind(1)
+    b = Variable("x", 1, 1).bind(1)
    s = b.simplify()
    self.assertEqual(s.op, Ops.CONST)
    self.assertEqual(s.val, 1)
@@ -1381,16 +1374,6 @@ class TestInvalidIndex(unittest.TestCase):
    c2 = UOp.const((1, Invalid, 1, 1))
    self.assertIs((c1+c2).simplify(), UOp.const((2, Invalid, Invalid, Invalid)))

-  def test_gated_load_keeps_index_valid(self):
-    # the load executes even on gated-off iterations: gated_given_valid must not erase its mask (PADTO OOB shape)
-    buf = UOp.param(0, dtypes.bool, (17,))
-    ridx = Variable("ridx", 0, 31)
-    cond = ridx < 17
-    load = buf.index(ridx.valid(cond))
-    out = graph_rewrite(cond.where(load.where(uconst(2), uconst(0)), UOp.invalid()), sym)
-    idx = next(u for u in out.toposort() if u.op is Ops.INDEX)
-    self.assertIs(idx.src[1].get_valid(), cond.simplify())
-
 class TestStoreLoadFolding(unittest.TestCase):
  """Tests for store(index, load(index)) -> NOOP rule. This rule matches patterns that EMERGE during simplification."""
  def test_store_load_folding(self):
@@ -5,7 +5,7 @@ from tinygrad.tensor import Tensor
 from tinygrad.helpers import Timing, Context, cdiv
 from tinygrad.dtype import dtypes, AddrSpace, ConstFloat, Invalid  # noqa: F401
 from tinygrad.device import Device
-from tinygrad.uop.ops import Ops, AxisType, ParamArg, PatternMatcher, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
+from tinygrad.uop.ops import Ops, ParamArg, PatternMatcher, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
 from tinygrad.uop.weak import pm_lower_index_dtype
 from tinygrad.uop.spec import spec_program, spec_shared, type_verify
 from tinygrad.uop.symbolic import sym, pm_remove_invalid
@@ -41,6 +41,11 @@ class TestDTypeFromUOp(unittest.TestCase):
    # an explicit (strong) const dtype is legal until the field is removed
    self.assertEqual(UOp.const(3, dtypes.int32).dtype, dtypes.int32)

+  def test_weak_dtype_rejected_by_program_spec(self):
+    for weak, concrete, value in ((dtypes.weakint, dtypes.int32, 1), (dtypes.weakfloat, dtypes.float32, 1.0)):
+      with self.assertRaises(RuntimeError): type_verify(UOp.const(value, weak).sink(), spec_program)
+      type_verify(UOp.const(value, concrete).sink(), spec_program)
+
  def test_invalid_stated_dtype(self):
    # UOp.const normalizes a stated dtype away (const_like/full pass their position's); the core constructor does not,
    # and the spec is what rejects a non-bool Invalid
@@ -129,7 +134,7 @@ class TestConstFloatEq(unittest.TestCase):
    self.assertFalse(Invalid != HoldsInvalid())

  def test_matchers_agree_on_nan(self):
-    n = UOp.const(math.nan)
+    n = UOp.const(math.nan, dtypes.float32)
    for compiled in (False, True):
      pm = PatternMatcher([(UPat(Ops.CONST, arg=math.nan), lambda: True)], compiled=compiled)
      self.assertTrue(pm.rewrite(n), f"{compiled=}")
@@ -301,9 +306,9 @@ class TestFastIdiv(unittest.TestCase):
      self.assertNotIn(Ops.CMOD, ops, f"For dtype={dt} FLOORMOD by pow2 left a MOD")
      self.assertNotIn(Ops.FLOORMOD, ops, f"For dtype={dt} FLOORMOD survived past late rewrite")

-  def test_floordiv_power_of_two(self):
-    # FLOORDIV by a power of two lowers to a shift, with no round toward zero correction (a shift is exactly floor division)
-    for dt in (dtypes.int32, dtypes.uint32, dtypes.int64, dtypes.uint64):
+  def test_floordiv_power_of_two_uint(self):
+    # uint FLOORDIV by a power of two lowers to a shift, leaving no IDIV/FLOORDIV in the kernel
+    for dt in (dtypes.uint32, dtypes.uint64):
      g = UOp.param(0, dt, (3,))
      c = UOp.const(2).cast(dt)
      a = UOp(Ops.FLOORDIV, dt, (g.index(c), c))
@@ -311,7 +316,6 @@ class TestFastIdiv(unittest.TestCase):
      ops = [x.op for x in uops]
      self.assertIn(Ops.SHR, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
      self.assertNotIn(Ops.CDIV, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
-      self.assertNotIn(Ops.CMOD, ops, f"For dtype={dt} FLOORDIV by pow2 kept the round toward zero correction")
      self.assertNotIn(Ops.FLOORDIV, ops, f"For dtype={dt} FLOORDIV survived past late rewrite")

  @Context(DISABLE_FAST_IDIV=0)
@@ -344,9 +348,10 @@ class TestFastIdiv(unittest.TestCase):
  def test_fast_idiv_remove_powers_of_two(self):
    ridx = UOp.range(2**20, 0)
    uops = to_uops_list([ridx//(7*64)], ren=Device[Device.DEFAULT].renderer)
+    ops = [x.op for x in uops]
    # this requires shifting out the powers of two before doing fast_idiv
    # (((ridx0>>6)*18725)>>17) instead of (int)((((long)(ridx0)*1198373)>>29))
-    self.assertNotIn(dtypes.long, [x.dtype for x in uops])
+    self.assertNotIn(Ops.CAST, ops)

  @unittest.expectedFailure
  def test_fast_idiv_overflow(self):
@@ -458,14 +463,6 @@ class TestUopsObject(unittest.TestCase):
    self.assertEqual(a.device, Device.DEFAULT)

 class TestUOpRender(unittest.TestCase):
-  def test_render_ssimplified_marg_outside_toposort(self):
-    r = UOp.range(UOp.const(16, dtypes.int), 2, AxisType.WEAK, dtype=dtypes.int)
-    offset = (r * 2) + (r * 2)
-    shrink = UOp(Ops.SHRINK, src=(UOp.param(0, dtypes.uint, (32,)), offset, UOp.const(2, dtypes.int)))
-    self.assertIsNot(shrink.src[1], shrink.marg[0][0])
-    self.assertEqual(shrink.render(simplify=False), "p0.shrink((((r2*4), 2),))")
-    self.assertEqual(UOp.range(1, 0, src=(shrink,), dtype=dtypes.int).render(simplify=False), "r0")
-
  def test_render_vectorize_empty(self):
    u = UOp(Ops.STACK, dtype=dtypes.void, src=())
    self.assertEqual(u.render(simplify=False), "{}")
@@ -1,12 +1,10 @@
 import unittest
-from tinygrad import dtypes
+from tinygrad import dtypes, Variable
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import Ops, UOp, AxisType
 from test.helpers import to_uops_list

-def Variable(name, nmin, nmax): return UOp.variable(name, nmin, nmax, param=True)
-
 class TestValidateOOB(unittest.TestCase):
  """Test z3 validation of index bounds for different ALU ops and patterns."""

@@ -1,5 +1,5 @@
-import unittest
-import decimal, sys, json, contextlib, tempfile, pickle, io, math, pathlib
+import unittest, decimal, sys, json, contextlib, tempfile, pickle, io, math
+from pathlib import Path
 from dataclasses import dataclass
 from typing import Generator

@@ -43,7 +43,7 @@ def save_viz():
  Buffer.profile_events.clear()
  cpu_events.clear()
  viz = VizTrace()
-  with Context(VIZ=-1, TRACK_MATCH_STATS=2, PROFILE=1, PARALLEL=0):
+  with Context(VIZ=-1, TRACK_MATCH_STATS=2, PROFILE=1):
    yield viz
  viz.set_data()

@@ -236,8 +236,8 @@ class TestViz(unittest.TestCase):
  def test_const_node_visibility(self):
    with save_viz() as viz:
      a = UOp.variable("a", 0, 10, dtype=dtypes.int)
-      z = UOp.const(0)
-      y = UOp.const(math.pi)
+      z = UOp.const(0, a.dtype)
+      y = UOp.const(math.pi, dtypes.float)
      alu = a*z
      ret = exec_rewrite(sink:=UOp.sink(alu, y), [sym])
    lst = viz.list_items()
@@ -249,7 +249,7 @@ class TestViz(unittest.TestCase):
    self.assertTrue(graphs[0][id(y)]["exclude"])
    self.assertFalse(graphs[0][id(alu)]["exclude"])
    self.assertEqual(graphs[0][id(y)]["label"].split("\n")[:2], ["CONST", "3.14159"])
-    self.assertEqual(list(graphs[1]), [id(u) for u in ret.toposort()])  # rewrite graph keys follow the rewritten sink's toposort
+    self.assertEqual(list(graphs[1]), [id(z), id(y), id(ret)])

  def test_const_reshape_expand_folded(self):
    # CONST->EXPAND should be folded into the ALU node, not shown as separate EXPAND nodes
@@ -305,10 +305,10 @@ class TestVizTree(unittest.TestCase):

  def test_tree_view(self):
    with save_viz() as viz:
-      a = UOp.variable("a",0,10,param=True)
-      b = UOp.variable("b",0,10,param=True)
-      c = UOp.variable("c",0,10,param=True)
-      d = UOp.variable("d",0,10,param=True)
+      a = UOp.variable("a",0,10)
+      b = UOp.variable("b",0,10)
+      c = UOp.variable("c",0,10)
+      d = UOp.variable("d",0,10)
      sink = UOp.sink(a+b, c+d)
      def tree_rewrite(): return graph_rewrite(sink, root, name="root")
      tree_rewrite()
@@ -516,22 +516,6 @@ class TestVizIntegration(unittest.TestCase):
    src_render = get_render(viz.data, steps[src_idx]["query"])["src"]
    self.assertEqual(src, src_render)

-  def test_profiler_duplicate_name(self):
-    kernel_name = "duplicate_name"
-    def one(A:UOp): return A[0].store(UOp.const(1.0, dtypes.float)).sink(arg=KernelInfo(kernel_name))
-    def zero(A:UOp): return A[0].store(UOp.const(0.0, dtypes.float)).sink(arg=KernelInfo(kernel_name))
-    with save_viz() as viz:
-      @TinyJit
-      def f(a:Tensor, b:Tensor): return Tensor.custom_kernel(a, fxn=one)[0], Tensor.custom_kernel(b, fxn=zero)[0]
-      a, b = Tensor.empty(4, device="NULL"), Tensor.empty(4, device="NULL")
-      # warmup
-      for _ in range(2): Tensor.realize(*f(a, b))
-      Tensor.realize(*f(a, b))
-    kernels = {i for i,c in enumerate(viz.list_items()) if c["name"] == kernel_name}
-    profile = decode_profile(unwrap(get_profile(viz.data, cpu_events)))
-    events = [e for e in profile["layout"]["NULL"]["events"] if e["name"] == kernel_name]
-    self.assertEqual({e["ref"] for e in events}, kernels)
-
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry
 from tinygrad.viz.serve import get_profile
 from tinygrad.viz.cli import decode_profile
@@ -835,6 +819,8 @@ from extra.gemm.amd_asm_matmul import Kernel

@needs_tracked_pm
 class TestCfg(unittest.TestCase):
+  def setUp(self): self.arch = "gfx1100"
+
  def get_cfg(self, name:str, k:Kernel):
    insts = k.finalize()
    def fxn(out:UOp) -> UOp:
@@ -843,7 +829,7 @@ class TestCfg(unittest.TestCase):
      sink = UOp.sink(out.base, lidx, gidx, arg=KernelInfo(name=name))
      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
    with save_viz() as viz:
-      with Context(DEV="NULL::gfx1100"):
+      with Context(DEV=f"NULL::{self.arch}"):
        out = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
        _ = do_to_program(out.schedule_linear().src[-1].src[0], Device[out.device].renderer)
    codegen_rewrites = next(s for s in viz.list_items() if s["name"] == name)
@@ -1025,8 +1011,8 @@ def run_cli(*cli_args) -> list[dict]:
@contextlib.contextmanager
 def write_files(viz) -> list[str]:
  with tempfile.TemporaryDirectory() as tmpdir:
-    (r:=pathlib.Path(tmpdir)/"rewrites.pkl").write_bytes(pickle.dumps(viz.data.trace))
-    (p:=pathlib.Path(tmpdir)/"profile.pkl").write_bytes(pickle.dumps(cpu_events))
+    (r:=Path(tmpdir)/"rewrites.pkl").write_bytes(pickle.dumps(viz.data.trace))
+    (p:=Path(tmpdir)/"profile.pkl").write_bytes(pickle.dumps(cpu_events))
    yield ["--rewrites-path", str(r), "--profile-path", str(p)]

 class TestCLI(unittest.TestCase):
@@ -10,12 +10,12 @@ from test.helpers import replace_opts
 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 and uop.shape == (4,)]),
-            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.float and uop.shape == (4,)]))
+    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype.scalar() == dtypes.float and uop.shape == (4,)]),
+            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype.scalar() == dtypes.float and uop.shape == (4,)]))
  @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 and uop.shape == (4,)]),
-            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.half and uop.shape == (4,)]))
+    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype.scalar() == dtypes.half and uop.shape == (4,)]),
+            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype.scalar() == dtypes.half and uop.shape == (4,)]))

  def test_float4_basic(self):
    a = Tensor.empty(2, 8).realize()
@@ -2,7 +2,6 @@ import unittest
 from tinygrad import Tensor, UOp, dtypes
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import Ops
-from test.helpers import KernelCountException

 class TestRingAllReduce(unittest.TestCase):
  def test_schedule_ring(self):
@@ -14,7 +13,7 @@ class TestRingAllReduce(unittest.TestCase):
      copies = [si for si in linear.src if si.src[0].op is Ops.COPY]
      pairs = [(c.src[1].buffer.device, c.src[2].buffer.device) for c in copies]
      # N*(N-1) scatter reduce, and N*(N-1) allgather
-      if len(pairs) != N*(N-1)*2: raise KernelCountException(N*(N-1)*2, len(pairs))
+      self.assertEqual(len(pairs), N*(N-1)*2)
      # copy topology forms a ring
      self.assertEqual(len(set(pairs)), N)

@@ -26,8 +25,8 @@ class TestRingAllReduce(unittest.TestCase):
      linear = t.sum(0).mul(2.0).contiguous().linear_with_vars()[0]
      copies = [si for si in linear.src if si.src[0].op is Ops.COPY]
      sinks = [si for si in linear.src if si.src[0].op is Ops.SINK]
-      if len(copies) != 24: raise KernelCountException(24, len(copies))
-      if len(sinks) != 26: raise KernelCountException(26, len(sinks))
+      self.assertEqual(len(copies), 24)
+      self.assertEqual(len(sinks), 26)

  @Context(RING=0, ALL2ALL=0)
  def test_schedule_naive(self):
@@ -40,8 +39,8 @@ class TestRingAllReduce(unittest.TestCase):
    sinks = [si for si in linear.src if si.src[0].op is Ops.SINK]
    pairs = [(c.src[1].buffer.device, c.src[2].buffer.device) for c in copies]

-    if len(pairs) != N*(N-1): raise KernelCountException(N*(N-1), len(pairs))
-    if len(sinks) != 2: raise KernelCountException(2, len(sinks))
+    self.assertEqual(len(pairs), N*(N-1))
+    self.assertEqual(len(sinks), 2)
    self.assertTrue(all(dst != src for dst, src in pairs))

  def test_symbolic_shape(self):
@@ -65,7 +64,7 @@ class TestAllreduceCast(unittest.TestCase):
    with Context(ALLREDUCE_CAST=allreduce_cast, RING=0, SCACHE=0):
      t = Tensor.empty(4, 4, dtype=dtype).shard(ds, axis=0)
      linear = t.sum(0).linear_with_vars()[0]
-      return {si.src[1].buffer.dtype for si in linear.src if si.src[0].op is Ops.COPY}
+      return {si.src[1].buffer.dtype.scalar() for si in linear.src if si.src[0].op is Ops.COPY}

  def test_allreduce_cast_bf16(self):
    # with ALLREDUCE_CAST, allreduce copies stay in bfloat16 instead of promoting to float32
@@ -540,10 +540,6 @@ class TestAssign(unittest.TestCase):
    c = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
    c[0:2].bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000], dtype=dtypes.uint32)).realize()
    np.testing.assert_allclose(c.numpy(), [4.0, 3.0, 3.0, 4.0])
-    # without .realize()
-    a = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
-    a.bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.uint32))
-    np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])

  def test_assign_bitcast_different_size(self):
    # assign to a shape-changing bitcast view (only works on DISK currently)
@@ -1,6 +1,6 @@
 import unittest
 import numpy as np
-from tinygrad import Tensor, dtypes, nn
+from tinygrad import Tensor, dtypes
 from tinygrad.llm.model import (
  GatedDeltaNetBlock, SSMConfig, TransformerBlock, TransformerConfig,
  apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk,
@@ -45,10 +45,10 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
    return Tensor.linspace(start, stop, int(np.prod(shape)), dtype=dtypes.float32).reshape(*shape)

  def _make_config(self, **kwargs):
-    return TransformerConfig(**({"num_blocks":1, "dim":32, "hidden_dim":64, "n_heads":1, "n_kv_heads":1,
-                                 "norm_eps":1e-5, "vocab_size":32, "head_dim":32, "rope_theta":10000.0,
-                                 "rope_dim":32, "v_head_dim":32, "max_context":4, "ssm_layers":(True,),
-                                 "ssm":SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32)} | kwargs))
+    return TransformerConfig(**({"num_blocks":1, "dim":4, "hidden_dim":8, "n_heads":1, "n_kv_heads":1,
+                                 "norm_eps":1e-5, "vocab_size":32, "head_dim":4, "rope_theta":10000.0,
+                                 "rope_dim":4, "v_head_dim":4, "max_context":4, "ssm_layers":(True,),
+                                 "ssm":SSMConfig(conv_kernel=2, state_size=2, group_count=1, time_step_rank=1, inner_size=2)} | kwargs))

  def _make_block(self, config:TransformerConfig) -> GatedDeltaNetBlock:
    block = GatedDeltaNetBlock(config, config.ssm)
@@ -79,10 +79,6 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
      recurrent_state = cache[:, conv_flat:].reshape(cache.shape[0], block.num_v_heads, block.head_v_dim, block.head_v_dim)
      return conv_state, recurrent_state

-  def _reset_state(self, block:GatedDeltaNetBlock):
-    Tensor.realize(block.conv_state.assign(block.conv_state.const_like(0)),
-                   block.recurrent_state.assign(block.recurrent_state.const_like(0)))
-
  def _linear_np(self, x:np.ndarray, weight:np.ndarray) -> np.ndarray:
    return x.astype(np.float32) @ weight.T.astype(np.float32)

@@ -90,7 +86,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
    x_float = x.astype(np.float32)
    return (x_float / np.sqrt((x_float * x_float).mean(axis=-1, keepdims=True) + eps)) * weight.astype(np.float32)

-  def _normalize_np(self, x:np.ndarray, eps:float=1e-6) -> np.ndarray:
+  def _normalize_np(self, x:np.ndarray, eps:float=1e-12) -> np.ndarray:
    return x / np.maximum(np.sqrt((x * x).sum(axis=-1, keepdims=True)), eps)

  def _softplus_np(self, x:np.ndarray) -> np.ndarray:
@@ -152,12 +148,6 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
    x = Tensor.linspace(-1.0, 1.0, 3 * config.dim, dtype=dtypes.float32).reshape(1, 3, config.dim)

    expected_outs, expected_conv, expected_recurrent = self._naive_attention(block, x)
-    out = self._run_attention(block, x, 0)
-    conv_state, recurrent_state = self._cache_views(block)
-    np.testing.assert_allclose(out, np.concatenate(expected_outs, axis=1), rtol=1e-3, atol=1e-3)
-    np.testing.assert_allclose(conv_state, expected_conv[-1], rtol=1e-3, atol=1e-3)
-    np.testing.assert_allclose(recurrent_state, expected_recurrent[-1], rtol=1e-3, atol=1e-3)
-    self._reset_state(block)

    for step in range(x.shape[1]):
      out = self._run_attention(block, x[:, step:step+1], step)
@@ -173,7 +163,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
    prompt = Tensor.linspace(0.75, -0.75, 2 * config.dim, dtype=dtypes.float32).reshape(1, 2, config.dim)

    for i in range(warmup.shape[1]): self._run_attention(block, warmup[:, i:i+1], i)
-    self._reset_state(block)
+    Tensor.realize(*block._state_reset_ops())
    expected_outs, expected_conv, expected_recurrent = self._naive_attention(block, prompt)

    for step in range(prompt.shape[1]):
@@ -187,64 +177,18 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
                                 err_msg=f"GatedDeltaNet reset recurrent cache mismatch at step {step}")

  def test_kda_channel_decay(self):
-    config = self._make_config(dim=4, hidden_dim=8, n_heads=2, head_dim=4, rope_dim=4, v_head_dim=4,
-      ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True))
-    block, x = GatedDeltaNetBlock(config, config.ssm), Tensor([[[1., 2., 0., 0.], [2., 1., 0., 0.]]])
+    config = self._make_config(n_heads=2, ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True))
+    block, x = GatedDeltaNetBlock(config, config.ssm), Tensor([[[1., 2., 0., 0.]]])
+    # f_b(f_a(x)) = [1, 2, 3, 4]
    block.ssm_f_a.weight = Tensor([[1., 0., 0., 0.], [0., 1., 0., 0.]])
    block.ssm_f_b.weight = Tensor([[1., 0.], [0., 1.], [1., 1.], [2., 1.]])
    block._init_state(x)
    initial_state = Tensor.arange(8, dtype=dtypes.float32).reshape(1, 2, 2, 2)
    block.recurrent_state.assign(initial_state).realize()
    block.ssm_a = Tensor([[-1.], [-1.]])
-    block._attention(x, x.shape[1]).realize()
-    alpha = np.exp(-self._softplus_np(np.array([[1, 2, 3, 4], [2, 1, 3, 5]])).reshape(2, 2, 2)).prod(0)
-    np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha[..., None], rtol=1e-5, atol=1e-5)
-
-  def test_kda_prefill_matches_decode(self):
-    config = self._make_config(ssm=SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32, kda=True))
-    block = GatedDeltaNetBlock(config, config.ssm)
-    for p in nn.state.get_parameters(block):
-      p.replace(self._tensor_linspace(-0.05, 0.05, p.shape) if len(p.shape) > 1 else self._tensor_linspace(0.05, 0.1, p.shape))
-    x = self._tensor_linspace(-0.5, 0.5, (1, 3, config.dim))
-    prefill = self._run_attention(block, x, 0)
-    prefill_conv, prefill_recurrent = self._cache_views(block)
-    self._reset_state(block)
-    decode = np.concatenate([self._run_attention(block, x[:, i:i+1], i) for i in range(3)], axis=1)
-    decode_conv, decode_recurrent = self._cache_views(block)
-    np.testing.assert_allclose(prefill, decode, rtol=1e-3, atol=1e-3)
-    np.testing.assert_allclose(prefill_conv, decode_conv, rtol=1e-3, atol=1e-3)
-    np.testing.assert_allclose(prefill_recurrent, decode_recurrent, rtol=1e-3, atol=1e-3)
-
-  def test_varied_chunk_sizes_match_decode(self):
-    for kda in (False, True):
-      ssm = SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32, kda=kda)
-      config = self._make_config(ssm=ssm)
-      if kda:
-        block = GatedDeltaNetBlock(config, config.ssm)
-        for p in nn.state.get_parameters(block):
-          p.replace(self._tensor_linspace(-0.05, 0.05, p.shape) if len(p.shape) > 1 else self._tensor_linspace(0.05, 0.1, p.shape))
-      else: block = self._make_block(config)
-      x = self._tensor_linspace(-0.5, 0.5, (1, 4, config.dim))
-      decode = np.concatenate([self._run_attention(block, x[:, i:i+1], i) for i in range(4)], axis=1)
-      decode_conv, decode_recurrent = self._cache_views(block)
-      for chunking in ([4], [2, 2], [1, 3], [3, 1], [2, 1, 1]):
-        self._reset_state(block)
-        outs, start = [], 0
-        for size in chunking:
-          outs.append(self._run_attention(block, x[:, start:start+size], start))
-          start += size
-        chunked_conv, chunked_recurrent = self._cache_views(block)
-        np.testing.assert_allclose(np.concatenate(outs, axis=1), decode, rtol=1e-3, atol=1e-3, err_msg=f"{kda=} {chunking=}")
-        np.testing.assert_allclose(chunked_conv, decode_conv, rtol=1e-3, atol=1e-3, err_msg=f"{kda=} {chunking=}")
-        np.testing.assert_allclose(chunked_recurrent, decode_recurrent, rtol=1e-3, atol=1e-3, err_msg=f"{kda=} {chunking=}")
-
-  def test_start_zero_resets_realized_state(self):
-    config, x = self._make_config(max_context=3), self._tensor_linspace(-1, 1, (1, 3, 32))
-    block = self._make_block(config)
-    self._run_attention(block, x, 0)
-    restarted = self._run_attention(block, x[:, :2], 0)
-    fresh = self._run_attention(self._make_block(config), x[:, :2], 0)
-    np.testing.assert_allclose(restarted, fresh, rtol=1e-3, atol=1e-3)
+    block._attention(x, 0).realize()
+    alpha = np.exp(-self._softplus_np(np.arange(1, 5)).reshape(1, 2, 1, 2))
+    np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha, rtol=1e-5, atol=1e-5)

 class TestPairwiseTopk(unittest.TestCase):
  def test_basic_topk(self):
@@ -224,21 +224,6 @@ class TestCallSchedule(unittest.TestCase):
    np.testing.assert_equal(x.numpy(), [2, 2, 2])
    np.testing.assert_equal(y.numpy(), [3, 3, 3])

-  def test_precompile_nested_scope_collision(self):
-    # a precompiled function body gets its own positional p{slot} params; they must not be renumbered when the call is
-    # scheduled inside an enclosing realize with a different slot ordering. the store must use this call's Variable
-    cache = Tensor.zeros(16)
-    @function(precompile=True, allow_implicit=True)
-    def store(x:Tensor, sp:UOp) -> Tensor:
-      # update a cache at a symbolic offset, like an attention KV cache update
-      return Tensor(cache.uop.after(cache[sp:sp+x.shape[0]].uop.store(x.uop)))[:sp+x.shape[0]].sum()
-    sp_v, nt_v = UOp.variable("sp", 0, 8), UOp.variable("nt", 1, 8)
-    t = Tensor.arange(16).float().realize()
-    sp, nt = sp_v.bind(0), nt_v.bind(8)
-    store(t[sp:sp+nt].clone().realize(), sp).realize()
-    np.testing.assert_equal(cache.numpy()[:8], t[:8].numpy())
-    np.testing.assert_equal(cache.numpy()[8:], np.zeros(8))
-
  def test_precompile_schedule_cache_hit(self):
    """two instances of the same @function should produce identical function body keys (schedule cache hit)"""
    @function(precompile=True)
@@ -3,12 +3,11 @@ import tempfile, unittest, math
 from tinygrad import Tensor, dtypes, TinyJit
 from tinygrad.helpers import Context
 from tinygrad.dtype import least_upper_float
-from tinygrad.uop.ops import UOp, Ops, GroupOp, dtype_from_uop, graph_rewrite
+from tinygrad.uop.ops import UOp, Ops, dtype_from_uop, graph_rewrite
 from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak
 from tinygrad.uop.symbolic import symbolic_simple
 from tinygrad.uop.spec import spec_shared, type_verify
 from tinygrad.engine.jit import JitError
-from test.helpers import full_rewrite


 class TestWeakPromotion(unittest.TestCase):
@@ -77,11 +76,6 @@ class TestWeakPromotion(unittest.TestCase):
      committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_lower_index_dtype, ctx={})
    self.assertEqual([u.dtype for u in committed.toposort() if u.op is Ops.ADD], [dtypes.float32])

-  def test_div_sub_operand_kept_weak(self):
-    a = Tensor.empty(4, dtype=dtypes.float32)
-    for t in (a / 1, a - 0):
-      self.assertEqual(t.uop.src[1].dtype, dtypes.weakfloat)
-
  def test_cast_weak_expression_commits_at_cast_floor(self):
    # the floor never narrows: a cast BELOW the default does not pull the compute width down with it
    with Context(DEFAULT_FLOAT=dtypes.float32):
@@ -94,13 +88,6 @@ class TestWeakPromotion(unittest.TestCase):
      out = Tensor(1.0, dtype=dtypes.float32, device="CPU") / denom
      self.assertAlmostEqual(out.item(), 1 / (70000 + 1e-5), places=10)

-  def test_stacked_weak_casts_convert_each_kind(self):
-    # each weak cast is a kind conversion: weakint truncates before weakfloat re-lifts (neither is only a marker)
-    x = Tensor([2.5, -3.7], dtype=dtypes.float32, device="CPU")
-    stacked = x.cast(dtypes.weakint).cast(dtypes.weakfloat)
-    self.assertIs(stacked.dtype, dtypes.weakfloat)
-    self.assertEqual(stacked.tolist(), [2.0, -3.0])
-
  def test_uop_scalar_const_lifts_kind(self):
    for dtype, value, out_dtype, const_dtype in ((dtypes.weakint, 1, dtypes.weakint, dtypes.weakint),
                                                 (dtypes.int32, 1, dtypes.int32, dtypes.weakint),
@@ -289,18 +276,5 @@ class TestSignedUint64Weakfloat(unittest.TestCase):
    self.assertAlmostEqual((i64 + u64).sin().item(), math.sin(2), places=5)  # Unary lowers before transcendental


-class TestNoRedundantWide(unittest.TestCase):
-  def wide_alu(self, t:Tensor) -> int:
-    return sum(sum(1 for u in full_rewrite(call.src[0]).toposort() if u.op in GroupOp.ALU and u.dtype in {dtypes.long, dtypes.ulong})
-               for call in t.schedule_linear().src if call.src[0].op is Ops.SINK)
-
-  def test_unbounded_long_stays_long(self):
-    self.assertGreater(self.wide_alu(Tensor.empty(16, dtype=dtypes.long)*3 + 1), 0)
-
-  def test_fancy_index_has_no_wide_alu(self):
-    j, o = Tensor([0, 1, 2]).reshape(3, 1), Tensor([0, 1]).reshape(1, 2)
-    self.assertEqual(self.wide_alu(Tensor.empty(8, 9, 10, 11, 12)[1, j, 2, o, 2]), 0)
-
-
 if __name__ == "__main__":
  unittest.main()
@@ -4,7 +4,7 @@ from tinygrad.function import function
 from tinygrad import Tensor, GlobalCounters, Device
 from tinygrad.dtype import Invalid
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, ProgramInfo
-from test.helpers import assert_kernel_count, KernelCountException
+from test.helpers import assert_kernel_count

 class TestFunction(unittest.TestCase):
  def test_simple(self):
@@ -516,7 +516,7 @@ class TestFunctionTuple(unittest.TestCase):
    Tensor.realize(a)
    c = f(a)

-    if count_kernels(c) != 1: raise KernelCountException(1, count_kernels(c))
+    self.assertEqual(count_kernels(c), 1)

    c.sum().backward()
    Tensor.realize(a.grad)
@@ -51,10 +51,6 @@ class TestTensorGradient(unittest.TestCase):
    with self.assertRaises(RuntimeError): x.sum().gradient(x)
    with self.assertRaises(RuntimeError): x.float().sum().gradient(x)

-  def test_const_target_raise(self):
-    t = Tensor(2.0)
-    with self.assertRaises(RuntimeError): (t * 2.0).gradient(t)
-
  def test_copy_to_device_gradient(self):
    t = Tensor([1.0, 2, 3]).realize()
    t.to("CPU:1").square().sum().backward()
@@ -104,7 +100,7 @@ class TestTensorGradient(unittest.TestCase):

  def test_implicit_broadcast_where_gradient(self):
    # WHERE with a bare ()-shape branch: the scalar's gradient counts the positions where it is selected
-    cond, x, w = Tensor([True, False, True]), Tensor([1.0, 2.0, 3.0]), Tensor(4.0, dtype=dtypes.float32)
+    cond, x, w = Tensor([True, False, True]), Tensor([1.0, 2.0, 3.0]), Tensor(4.0)
    dw = Tensor(cond.uop.alu(Ops.WHERE, x.uop, w.uop)).sum().gradient(w)[0]
    self.assertEqual(dw.shape, ())
    self.assertEqual(dw.item(), 1.0)
@@ -113,7 +109,7 @@ class TestTensorGradient(unittest.TestCase):

  def test_implicit_broadcast_alu_gradient(self):
    # MUL with a bare ()-shape src, no EXPAND in the graph
-    x, w = Tensor([1.0, 2.0, 3.0]), Tensor(2.0, dtype=dtypes.float32)
+    x, w = Tensor([1.0, 2.0, 3.0]), Tensor(2.0)
    m = x.uop.alu(Ops.MUL, w.uop)
    self.assertIs(m.src[1], w.uop)
    dw = Tensor(m).sum().gradient(w)[0]
@@ -122,7 +118,7 @@ class TestTensorGradient(unittest.TestCase):

  def test_implicit_broadcast_intermediate_accumulation(self):
    # s is used directly and through an implicit broadcast edge, each edge's gradient reduces to s's shape before they sum
-    x, p = Tensor([1.0, 2.0, 3.0]), Tensor(0.5, dtype=dtypes.float32)
+    x, p = Tensor([1.0, 2.0, 3.0]), Tensor(0.5)
    s = p.sin()
    z = Tensor(x.uop.alu(Ops.MUL, s.uop)).sum() + s
    dp = z.gradient(p)[0]
@@ -25,10 +25,10 @@ class TestHCQUnit(unittest.TestCase):
    cpu_call = UOp(Ops.PROGRAM, src=(UOp.sink(),)).call(UOp.new_buffer("CPU", 1, dtypes.float))
    gpu_devs = [d0]

-    # CPU uses HCQ2 and is no longer batched into legacy HCQ graphs.
+    # local MMIO: GPU works alone and with CPU in batch (cpu_support=True)
    assert HCQGraph.supports_uop(gpu_devs, gpu_call) is True
-    assert HCQGraph.supports_uop(gpu_devs, cpu_call) is False
-    assert HCQGraph.supports_uop(gpu_devs + [cpu_dev], gpu_call) is False
+    assert HCQGraph.supports_uop(gpu_devs, cpu_call) is True
+    assert HCQGraph.supports_uop(gpu_devs + [cpu_dev], gpu_call) is True

    # USB MMIO: GPU-only still works, but CPU batching must be rejected (cpu_support=False)
    orig_view = d0.timeline_signal.base_buf.view
@@ -2,7 +2,7 @@ import unittest
 import numpy as np
 from dataclasses import replace
 from tinygrad import Tensor
-from tinygrad.llm.model import ExpertGating, TransformerBlock, TransformerConfig
+from tinygrad.llm.model import TransformerBlock, TransformerConfig

 def _moe_config(dim=8, hidden=16, n_heads=2, num_experts=4, num_experts_per_tok=2):
  return TransformerConfig(
@@ -96,32 +96,5 @@ class TestMoEFeedForward(unittest.TestCase):
    expected = moe_expected + shared_expected
    np.testing.assert_allclose(out.numpy(), expected, rtol=1e-2)

-  def test_moe_feed_forward_gating_funcs(self):
-    dim, hidden, n_heads = 8, 16, 2
-    num_experts, k = 4, 2
-    logits = np.array([4.0, 3.0, 0.0, -1.0], dtype=np.float32)
-    def softmax(x):
-      probs = np.exp(x - x.max())
-      return probs / probs.sum()
-    for gating_func in ExpertGating:
-      for norm_topk_prob in (False, True):
-        block = TransformerBlock(replace(_moe_config(dim, hidden, n_heads, num_experts, k),
-                                         expert_gating_func=gating_func, norm_topk_prob=norm_topk_prob))
-        block.ffn_gate_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) for _ in range(num_experts)])
-        block.ffn_up_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) * (i + 1) for i in range(num_experts)])
-        block.ffn_down_exps.weight = Tensor.stack(*[Tensor.eye(dim, hidden) for _ in range(num_experts)])
-        block.ffn_gate_inp.weight = Tensor((logits / dim)[None, :].repeat(dim, 0).T)
-        out = block._feed_forward(Tensor.ones(1, 1, dim)).numpy()[0, 0, 0]
-
-        if gating_func == ExpertGating.SOFTMAX: selection_scores = softmax(logits)
-        elif gating_func == ExpertGating.SIGMOID: selection_scores = 1 / (1 + np.exp(-logits))
-        elif gating_func == ExpertGating.SOFTMAX_WEIGHT: selection_scores = logits
-        else: selection_scores = np.sqrt(np.logaddexp(0, logits))
-        sel = np.argsort(selection_scores)[-k:]
-        weights = softmax(logits[sel]) if gating_func == ExpertGating.SOFTMAX_WEIGHT else selection_scores[sel]
-        if norm_topk_prob: weights /= weights.sum()
-        expected = (weights * (sel + 1)).sum() / (1 + np.exp(-1))
-        np.testing.assert_allclose(out, expected, rtol=1e-3)
-
 if __name__ == '__main__':
  unittest.main()
@@ -1,8 +1,6 @@
 import unittest
-import numpy as np
 from unittest.mock import patch
 from tinygrad import Tensor, UOp
-from tinygrad.nn.state import get_state_dict
 from tinygrad.schedule import schedule_cache
 from tinygrad.llm.model import Transformer, TransformerConfig
 from tinygrad.llm.serve import StreamRouter
@@ -15,18 +13,12 @@ V_TOKS = UOp.variable("toks", 1, 32)  # 32 is the default chunk_size in generate
 class TestTransformerGenerate(unittest.TestCase):
  def test_warmup(self):
    model, calls = Transformer(TEST_CONFIG), []
-    def generate(tokens, **kwargs):
+    def generate(tokens):
      calls.append(tokens)
      yield from (1, 2)
    with patch.object(model, "generate", generate): model.warmup()
    self.assertEqual(calls, [[0], [0]])

-  def test_warmup_then_generate_with_default_chunk(self):
-    # warmup must not capture JIT graphs that generate()'s default chunk_size then rejects
-    model = Transformer(TEST_CONFIG)
-    model.warmup()
-    self.assertIsInstance(next(model.generate([5, 6, 7, 8])), int)
-
  def test_first_recurrent_generate_before_state_init(self):
    model = Transformer(TEST_CONFIG)
    model.has_recurrent_block = True
@@ -46,15 +38,6 @@ class TestTransformerGenerate(unittest.TestCase):
      next(model.generate([1, 2, 3, 4, 5, 42, 10]))
    self.assertEqual(calls, [((1, 1), V_START_POS.bind(5)), ((1, 1), V_START_POS.bind(6))])

-  def test_recurrent_divergent_prompt_restarts(self):
-    model, calls = Transformer(TEST_CONFIG), []
-    model.has_recurrent_block, model._cached_tokens = True, [1, 2, 9]
-    def mock_call(self, tokens, start_pos, temperature):
-      calls.append(start_pos)
-      return Tensor([[42]])
-    with patch.object(Transformer, '__call__', mock_call): next(model.generate([1, 2, 10, 11]))
-    self.assertEqual(calls[0], V_START_POS.bind(0))
-
  def test_template_starts_reasoning(self):
    router = StreamRouter(reasoning=True)
    self.assertEqual(list(router.route("reasoning</think>answer")),
@@ -154,22 +137,6 @@ class TestTransformerGenerate(unittest.TestCase):
    # 4 tokens, chunk_size=4 -> 1 prefill chunk
    self.assertEqual(get_prefill_flags(list(range(4)), 4), [True, False, False])

-  def test_chunked_prefill_kv_cache_matches_single_chunk(self):
-    config = TransformerConfig(num_blocks=1, dim=8, hidden_dim=16, n_heads=1, n_kv_heads=1, norm_eps=1e-5,
-      vocab_size=32, head_dim=4, rope_theta=1000000, rope_dim=4, qk_norm=4, v_head_dim=4, max_context=16)
-    def model():
-      m = Transformer(config)
-      rng = np.random.RandomState(1234)
-      for t in get_state_dict(m).values():
-        t.assign(Tensor(rng.uniform(-1, 1, t.shape).astype(np.float32))).realize()
-      return m
-    def prefill(m, chunk_size):
-      gen = m.generate(list(range(1, 9)), chunk_size=chunk_size, temperature=0.0)
-      next(gen)
-      return [b.cache_kv.numpy() for b in m.blk]
-    for g, r in zip(prefill(model(), 4), prefill(model(), 8)):
-      np.testing.assert_allclose(g[:, :, :, :8, :], r[:, :, :, :8, :], atol=1e-5)
-
  def test_kv_cache_resume_matches_fresh(self):
    model = Transformer(TEST_CONFIG)

@@ -211,12 +178,6 @@ class TestTransformerGenerate(unittest.TestCase):
    # with temperature=2.0, we should see at least 2 distinct outputs across 5 runs
    self.assertGreater(len(runs), 1, "high temperature should produce varied outputs")

-  def test_recurrent_temperature_high_produces_variety(self):
-    model = Transformer(TEST_CONFIG)
-    model.has_recurrent_block = True
-    outputs = {model.forward(Tensor([[1]]), 0, Tensor([2.0])).item() for _ in range(5)}
-    self.assertGreater(len(outputs), 1)
-
  def test_temperature_passed_to_forward(self):
    """Temperature from generate should be passed through to __call__."""
    model = Transformer(TEST_CONFIG)
@@ -390,12 +390,6 @@ class TestMultiTensor(unittest.TestCase):
    self.assertEqual(out.shape, (rows, 8))
    np.testing.assert_equal(out[:3].to(Device.DEFAULT).numpy(), np.ones((3, 8)))

-  def test_symbolic_broadcast_consumed(self):
-    rows = Variable("rows", 1, 4).bind(3)
-    out = (Tensor.ones(rows).to(devices_2) + 1).realize()
-    self.assertEqual(out.shape, (rows,))
-    np.testing.assert_equal(out[:3].to(Device.DEFAULT).numpy(), np.full(3, 2))
-
  def test_multitensor_jit_in_list(self):
    # test MULTI tensor inside a list container - exercises the container unpacking + MULTI unpacking
    @TinyJit
@@ -1,16 +1,11 @@
 import unittest
 import functools
-from tinygrad import Tensor, Variable, UOp, function
+from tinygrad import Tensor, Variable, UOp
 from tinygrad.uop.ops import KernelInfo
 from tinygrad.schedule import schedule_cache

-def custom_add_kernel(A:UOp, B:UOp, num:int=0) -> UOp:
-  return A[0].set(B[0] + num).sink(arg=KernelInfo(f"custom_add_{num}"))
-
-def custom_add_backward(grad_output:UOp, _) -> tuple[None, UOp]:
-  grad = Tensor.invalids(*grad_output.shape, dtype=grad_output.dtype, device=grad_output.device)
-  grad = Tensor.custom_kernel(grad, Tensor(grad_output, device=grad_output.device), fxn=functools.partial(custom_add_kernel, num=0))[0]
-  return None, grad.uop
+def custom_set0_kernel(A:UOp, num:int) -> UOp:
+  return A[0].set(num).sink(arg=KernelInfo(f"custom_set0_{num}"))

 class TestScheduleCache(unittest.TestCase):
  def test_bound_variable_reuses_cache(self):
@@ -30,27 +25,27 @@ class TestScheduleCache(unittest.TestCase):

  def test_custom_kernel(self):
    for i in range(4):
-      a, b = Tensor.empty(1), Tensor.ones(1)
-      a = Tensor.custom_kernel(a, b, fxn=functools.partial(custom_add_kernel, num=i))[0]
+      a = Tensor.empty(1)
+      a = Tensor.custom_kernel(a, fxn=functools.partial(custom_set0_kernel, num=i))[0]
      a.realize()
-      self.assertEqual(a.item(), i+1)
+      self.assertEqual(a.item(), i)

  def test_same_custom_function_reuses_cache(self):
    schedule_cache.clear()
-    fxn = functools.partial(custom_add_kernel, num=10)
+    fxn = functools.partial(custom_set0_kernel, num=10)

    # first run
-    a, x = Tensor.empty(1), Tensor.ones(1)
-    a = Tensor.custom_kernel(a, x, fxn=fxn)[0]
+    a = Tensor.empty(1)
+    a = Tensor.custom_kernel(a, fxn=fxn)[0]
    a.realize()
-    self.assertEqual(a.item(), 11)
+    self.assertEqual(a.item(), 10)
    cache_size_after_first = len(schedule_cache)

    # second run with same function should reuse cache
-    b, x = Tensor.empty(1), Tensor.ones(1)
-    b = Tensor.custom_kernel(b, x, fxn=fxn)[0]
+    b = Tensor.empty(1)
+    b = Tensor.custom_kernel(b, fxn=fxn)[0]
    b.realize()
-    self.assertEqual(b.item(), 11)
+    self.assertEqual(b.item(), 10)
    self.assertEqual(len(schedule_cache), cache_size_after_first)

  def test_simple(self):
@@ -70,29 +65,5 @@ class TestScheduleCache(unittest.TestCase):
      print(num)
    self.assertEqual(len(schedule_cache), start_len_schedule_cache)

-  def test_simple_precompile(self):
-    @function(precompile=True, precompile_backward=True)
-    def f(x:Tensor) -> Tensor:
-      out = Tensor.invalids(*x.shape, dtype=x.dtype, device=x.device)
-      out = Tensor.custom_kernel(out, x, fxn=functools.partial(custom_add_kernel, num=10), grad_fxn=custom_add_backward)[0]
-      return out + x
-
-    # warmup
-    x = Tensor.ones(1).realize()
-    out = f(x)
-    out.backward(x)
-    self.assertEqual(out.item(), 12)
-    self.assertEqual(x.grad.item(), 2)
-
-    # use the cache next time function is called
-    start_len_schedule_cache = len(schedule_cache)
-    for _ in range(3):
-      x = Tensor.ones(1).realize()
-      out = f(x)
-      out.backward(x)
-      self.assertEqual(out.item(), 12)
-      self.assertEqual(x.grad.item(), 2)
-    self.assertEqual(len(schedule_cache), start_len_schedule_cache)
-
 if __name__ == "__main__":
  unittest.main()
@@ -1,7 +1,7 @@
 from dataclasses import replace, dataclass
 import itertools, functools
 from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
-from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, NUM_CPU_THREADS, TC_SELECT, TC_OPT, TracingKey, Context, panic
+from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, TracingKey, Context, panic
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, Ops, UPat, rewrite_group, KernelInfo, ProgramInfo, GroupOp, AxisType
 from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak, pm_cast_weak
 from tinygrad.uop.render import pyrender
@@ -12,7 +12,7 @@ from tinygrad.dtype import dtypes, AddrSpace

 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
-from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid
+from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_fold_cast_const, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid
 from tinygrad.uop.movement import mop_cleanup
 from tinygrad.codegen.decomp.dtype import pm_dtype_decomps
 from tinygrad.codegen.decomp.op import get_late_rewrite_patterns, get_simplifying_rewrite_patterns
@@ -153,8 +153,8 @@ devectorizer2 = mop_cleanup+pm_mops+PatternMatcher([
  # unpack WMMA
  (UPat(Ops.WMMA, name="u"), do_stack_wmma),
  # stacked INDEX is many INDEX
-  (UPat(Ops.INDEX, src=(UPat((Ops.PARAM, Ops.BUFFER), name="b"), UPat(Ops.STACK, name="s")), name="x"),
-   lambda b,s,x: UOp.stack(*[x.replace(src=(b,u)) for u in s.src])),
+  (UPat(Ops.INDEX, src=(UPat((Ops.PARAM, Ops.BUFFER), name="b"), UPat(Ops.STACK, name="s"))),
+   lambda b,s: UOp.stack(*[b.index(u) for u in s.src])),
  # INDEX into RESHAPE moves the RESHAPE
  (UPat(Ops.INDEX, src=(UPat((Ops.PARAM, Ops.BUFFER), name="b"), UPat(Ops.RESHAPE, name="s"))),
   lambda b,s: b.index(s.src[0]).reshape(s.shape)),
@@ -233,11 +233,10 @@ pm_reduce_local = pm_wmma_add+PatternMatcher([
  (UPat(Ops.SINK, name="sink"), merge_reduce_ends),
 ])+pm_clean_up_group_sink

-def is_shape_changing_bitcast(u:UOp): return u.op is Ops.BITCAST and u.shape != u.src[0].shape
 def maybe_load(u:UOp): return u.load() if u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL, AddrSpace.REG) else u
 pm_add_loads = PatternMatcher([
-  (UPat(GroupOp.Elementwise|{Ops.REDUCE,Ops.WMMA,Ops.STACK}, name="x"),
-   lambda x: None if is_shape_changing_bitcast(x) else x.replace(src=tuple(map(maybe_load, x.src)))),
+  # BITCAST?
+  (UPat(GroupOp.Elementwise|{Ops.REDUCE,Ops.WMMA,Ops.STACK}, name="x"), lambda x: x.replace(src=tuple([maybe_load(u) for u in x.src]))),
  (UPat(Ops.STORE, name="x"), lambda x: x.replace(src=(x.src[0], maybe_load(x.src[1]))+x.src[2:])),
 ])

@@ -282,10 +281,6 @@ pm_implicit_barriers = PatternMatcher([
  (UPat(Ops.END, name="end"), add_war_barrier),
 ])

-pm_casted_consts = PatternMatcher([
-  (UPat(Ops.CONST, dtypes.all, name="c"), lambda c: UOp.cconst(c.val, c.dtype)),
-])
-
 def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
  if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
  if DEBUG >= 5: print(pyrender(ast))
@@ -306,7 +301,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
    sink = graph_rewrite(sink, pm_split_ranges+pm_flatten_range, ctx={}, name="split ranges")

    # symbolic (NOTE: this is a requirement for pm_simplify_ranges to be correct)
-    sink = graph_rewrite(sink, sym+pm_flatten_range, name="initial symbolic")
+    sink = graph_rewrite(sink, sym+pm_fold_cast_const+pm_flatten_range, name="initial symbolic")

    # optimize (schedule) the AST
    sink = graph_rewrite(sink, pm_flatten_range+pm_simplify_ranges, ctx={}, name="simplify ranges")
@@ -351,7 +346,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:

  # lower index dtype
  # NOTE: we need indexing_simplify to remove the cast to long using the Invalid
-  sink = graph_rewrite(sink, symbolic_simple+pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes")
+  sink = graph_rewrite(sink, symbolic_simple+pm_fold_cast_const+pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes")

  # final symbolic before decomp
  sink = graph_rewrite(sink, symbolic, name="final symbolic")
@@ -362,7 +357,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:

  # floordiv+mod / dtype decomp (early)
  supported_ops = tuple(ren.code_for_op.keys())
-  pm_decomp = symbolic_simple+get_simplifying_rewrite_patterns(supported_ops)
+  pm_decomp = symbolic_simple+pm_fold_cast_const+get_simplifying_rewrite_patterns(supported_ops)
  sink = graph_rewrite(sink, pm_decomp, name="early decompositions")

  # late decomps + move gates from unrenderable INVALID where
@@ -378,10 +373,6 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
  pm_final_rewrite = pm_commit_weak+pm_cast_weak+pm_decomp+extra_matcher+pm_split_ends
  sink = graph_rewrite(sink, pm_final_rewrite+pm_remove_invalid, ctx=ren, name="final rewrite")

-  # spell every literal as a casted const CAST(dt, CONST(value))
-  # TODO: remove once consts are always weak
-  sink = graph_rewrite(sink, pm_casted_consts, name="casted consts", walk=True)
-
  # add implicit barriers (stores/loads through LOCAL memory ordered by AFTER or across loop iterations need workgroup barriers)
  sink = graph_rewrite(sink, pm_implicit_barriers, name="add implicit barriers")

@@ -460,7 +451,7 @@ pm_to_program = PatternMatcher([
  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.LINEAR), UPat(Ops.SOURCE, name="source")), name="prg"), do_compile),
 ])

-@rewrite_group(name=lambda ast,renderer,ret,**_: TracingKey((k:=ret.src[0].arg).name,(k.function_name, ast, ret.key),ret=renderer), replay=True)
+@rewrite_group(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.src[0].arg.name,(ret.src[0].arg.function_name, ast), ret=renderer), replay=True)
@Context(ALLOW_DEVICE_USAGE=0)
 def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
  """
@@ -489,14 +480,9 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
  if VIZ: graph_rewrite(prg, PatternMatcher([]), name="View Program")
  return prg

-# config affects generated programs and cache keys; context also carries compile-only behavior to workers
-to_program_config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32,
-                     DEFAULT_FLOAT, DEFAULT_INT, NUM_CPU_THREADS, TC_SELECT, TC_OPT)
-to_program_context = (*to_program_config, SPEC, DEBUG)
-def to_program_key(ast:UOp, renderer:Renderer) -> tuple:
-  return (ast.key, type(renderer), renderer.target, *[x.value for x in to_program_config])
-
 to_program_cache: dict[tuple, UOp] = {}
 def to_program(ast:UOp, renderer:Renderer) -> UOp:
-  if (prg:=to_program_cache.get(key:=to_program_key(ast, renderer))) is None: to_program_cache[key] = prg = do_to_program(ast, renderer)
+  config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT)
+  key = (ast.key, type(renderer), renderer.target, *[x.value for x in config])
+  if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer)
  return prg
@@ -25,16 +25,15 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
  match op:
    case Ops.NEG: return l2i(Ops.SUB, dt, zero, zero, *uops)
    case Ops.CAST if dt in (dtypes.long, dtypes.ulong) and uops[0].dtype not in dtypes.floats:
-      # the high word is the sign extension, and unsigned and bool sources zero extend
+      # the high word is the sign extension; bool has no sign, test the already-cast low word instead (bool < 0 would promote to weakint)
      x, lo = uops[0], uops[0].cast(l2i_dt[dt])
-      if x.dtype is dtypes.bool or x.dtype in dtypes.uints: return lo, lo.const_like(0)
-      return lo, (x < x.const_like(0)).where(lo.const_like(-1), lo.const_like(0))
+      sign = lo if x.dtype is dtypes.bool else x
+      return lo, (sign < sign.const_like(0)).where(lo.const_like(-1), lo.const_like(0))
    case Ops.CAST if dt in (dtypes.long, dtypes.ulong):
      return (lo:=uops[0].cast(l2i_dt[dt])), (uops[0] / 2**32).cast(l2i_dt[dt]) - ((uops[0] < 0) & lo.ne(0))
    case Ops.CAST if dt in dtypes.floats:
      small = (a1.eq(0) & (a0 >= 0)) | (a1.eq(-1) & (a0 < 0))
-      cdt = dt if dt == dtypes.float64 else dtypes.float32
-      return small.where(a0.cast(dt), ((a1.cast(cdt) * (2**32)) + a0.bitcast(dtypes.uint).cast(cdt)).cast(dt))
+      return small.where(a0.cast(dt), ((a1.cast(dtypes.float32) * (2**32)) + a0.bitcast(dtypes.uint).cast(dtypes.float32)).cast(dt))
    case Ops.CAST: return a0.bitcast(dtypes.uint).cast(dt)
    case Ops.BITCAST: return a0.bitcast(dt), a1.bitcast(dt)
    case Ops.SHL:
@@ -75,11 +75,7 @@ powers_of_two: dict[int, int] = {2**i:i for i in range(64)}
@functools.cache
 def get_simplifying_rewrite_patterns(ops:tuple[Ops, ...]) -> PatternMatcher:
  # these are rewrites that make things simpler
-  pat: list[tuple[UPat, Callable]] = []
-  # FLOORDIV by 2**y -> x >> y (an arithmetic shift is exactly floor division for any sign); fires before floordiv_to_idiv
-  if Ops.SHR in ops: pat.append((UPat.var("x", dtypes.ints)//UPat.cvar("c"),
-    lambda x,c: x >> v if (v:=powers_of_two.get(c.val, 0)) else None))
-  pat.append((UPat.var("a")//UPat.var("b"), floordiv_to_idiv))
+  pat: list[tuple[UPat, Callable]] = [(UPat.var("a")//UPat.var("b"), floordiv_to_idiv)]
  # FLOORMOD by 2**y -> x & (2**y-1) (correct floor mod for any sign in two's complement); fires before floormod_to_mod
  if Ops.AND in ops: pat.append((UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.val-1) if c.val in powers_of_two else None))
  pat.append((UPat.var("a")%UPat.var("b"), floormod_to_mod))
@@ -132,6 +128,6 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa
    if Ops.SHL in ops: pat += [(UPat.var('x').alu(Ops.SHL, UPat.cvar('n'))+UPat.var('c'), lambda x,n,c: x.alu(Ops.MULACC, x.const_like(1<<n.val), c))]
  # some backends emit FDIV for RECIP, in that case: a*(1/b) -> a/b
  if Ops.FDIV in ops:
-    pat += [(UPat.var("x").reciprocal(), lambda x: UOp.const(1.0).alu(Ops.FDIV, x))]
-    pat += [(UPat.var("a") * UPat(Ops.FDIV, dtypes.floats, src=(UPat.const(1), UPat.var("b"))), lambda a,b: a.alu(Ops.FDIV, b))]
+    pat += [(UPat.var("x").reciprocal(), lambda x: x.const_like(1).alu(Ops.FDIV, x))]
+    pat += [(UPat.var("a", dtypes.floats) * UPat(Ops.FDIV, dtypes.floats, src=(UPat.const(1), UPat.var("b"))), lambda a,b: a.alu(Ops.FDIV, b))]
  return PatternMatcher(pat)
@@ -57,7 +57,7 @@ def add_gpudims(ctx:Renderer, s:UOp):

  # get the idxs
  ki: KernelInfo = s.arg
-  if ctx.has_threads: idxs = [UOp.variable("core_id", 0, int(global_shape[0])-1, dtypes.int, param=True).cast(dtypes.weakint)]
+  if ctx.has_threads: idxs = [UOp.variable("core_id", 0, int(global_shape[0])-1, dtypes.int).cast(dtypes.weakint)]
  elif ki.dont_use_locals:
    assert not local_dims, "can't use locals if there's no local dims"
    idxs = get_grouped_dims("idx", global_shape, ctx.global_max, reverse=True)
@@ -89,7 +89,7 @@ def add_gpudims(ctx:Renderer, s:UOp):

 pm_device_to_var = PatternMatcher([
  # the DEVICE axis is not a program axis, it's bound per device at launch. lower it to the _device_num variable (like SPECIAL for devices)
-  (UPat(Ops.RANGE, name="r"), lambda r: UOp.variable("_device_num", 0, r.vmax, dtype=r.dtype, param=True) if r.arg[-1] is AxisType.DEVICE else None),
+  (UPat(Ops.RANGE, name="r"), lambda r: UOp.variable("_device_num", 0, r.vmax, dtype=r.dtype) if r.arg[-1] is AxisType.DEVICE else None),
  # ENDs that closed a DEVICE range no longer close it
  (UPat(Ops.END, name="e"), lambda e: e.replace(src=(e.src[0],)+tuple(s for s in e.src[1:] if s.op is not Ops.PARAM))
   if any(s.op is Ops.PARAM and s.arg.name == '_device_num' for s in e.src[1:]) else None),
@@ -26,7 +26,7 @@ def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]:
    # check if idx is out of bound when X is on the wrong side of the bound: X in [c+1, vmax] or [vmin, c-1]
    lo, hi = (c + 1, X.vmax) if is_upper_bound else (X.vmin, c - 1)
    if lo <= hi:
-      fake = UOp.variable(f"fake{i}", lo, hi, X.dtype, param=True)
+      fake = UOp.variable(f"fake{i}", lo, hi, X.dtype)
      subs = [{X: fake}]
      # idx may not have X itself, so also substitute a term of X: v -> fake - (X - v)
      terms = list(X.split_uop(Ops.ADD))
@@ -51,8 +51,8 @@ def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|N
  if not drop_stmt and idx is start_idx: return None
  new_valid = UOp.uprod(*ss) if (ss:=[s for s in valid.split_uop(Ops.AND) if s not in drop_stmt]) else None
  idx_y, idx_x = idx.index(1), idx.index(0)
-  if new_valid is not None: return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid))
-  return buf.index(idx_y, idx_x)
+  if new_valid is not None: return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid), dtype=dtypes.float)
+  return buf.index(idx_y, idx_x, dtype=dtypes.float)

 indexing_simplify = PatternMatcher([
  # image load valid idx simplification
@@ -88,9 +88,9 @@ def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
  buf = buf.replace(src=(shape_to_shape_arg((h, w, 4)),))
  shapes[buf.arg.slot] = (h, w)
  if valid.op is not Ops.CONST or valid.val is not True:
-    return buf.index(cidx.src[1].valid(valid), cidx.src[0].valid(valid))
+    return buf.index(cidx.src[1].valid(valid), cidx.src[0].valid(valid), dtype=dtypes.float)
  else:
-    return buf.index(cidx.src[1], cidx.src[0])
+    return buf.index(cidx.src[1], cidx.src[0], dtype=dtypes.float)

 pm_simplify_add_image = PatternMatcher([
  (UPat(Ops.SHRINK, src=(UPat(Ops.PARAM, name="buf"), UPat(name="x"), UPat(arg=4))), transform_to_image),
@@ -3,17 +3,17 @@ from tinygrad.uop.ops import PatternMatcher, UPat, Ops
 from tinygrad.dtype import Invalid, dtypes

 def move_where_load(gate, l, a, w):
-  return l.replace(src=(l.src[0], l.vconst_like(0) if a.is_invalid else l.const_like(a.val) if a.op is Ops.CONST else
+  return l.replace(src=(l.src[0], l.vconst_like(0) if a.is_invalid else
    a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype), l.src[2])).cast(w.dtype)

 pm_move_gates_from_index = PatternMatcher([
  # for image idx (must be first)
  (UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx_y"), UPat(arg=Invalid)),
                         UPat.var("gate").where(UPat.var("idx_x"), UPat(arg=Invalid))).load(name="l"),
-   lambda buf,gate,idx_y,idx_x,l: buf.index(idx_y, idx_x).load(l.vconst_like(0), gate)),
+   lambda buf,gate,idx_y,idx_x,l: buf.index(idx_y, idx_x, dtype=dtypes.float).load(l.vconst_like(0), gate)),
  (UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx_y"), UPat(arg=Invalid)),
                         UPat.var("gate").where(UPat.var("idx_x"), UPat(arg=Invalid))).store(UPat.var("data")),
-   lambda buf,gate,idx_y,idx_x,data: buf.index(idx_y, idx_x).store(data, gate)),
+   lambda buf,gate,idx_y,idx_x,data: buf.index(idx_y, idx_x, dtype=dtypes.float).store(data, gate)),

  # here we create the alt value for load to be 0s and remove the where Invalid
  (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat(), UPat.var("gate").where(UPat.var("idx"), UPat(arg=Invalid)),), name="mop", allow_any_len=True) \
@@ -4,7 +4,7 @@ from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
 from tinygrad.renderer.isa import ISARenderer, Register, greg
 from tinygrad.dtype import dtypes

-PSEUDO_OPS = {Ops.CONST, Ops.CAST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK}
+PSEUDO_OPS = {Ops.CONST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK}

 class LinearScanRegallocContext:
  # returns the uop that defines the virtual register
@@ -52,7 +52,7 @@ class LinearScanRegallocContext:
        # the value of a BUFFER is its 64bit address, XMM registers need 16 bytes
        sz = 16 if v.cons[0].size == 16 else (8 if self.vdef(v).op is Ops.BUFFER else self.vdef(v).dtype.itemsize)
        offset = self.stack_size + (sz - self.stack_size % sz) % sz
-        self.spills[v] = UOp.cconst(offset, dtypes.int32)
+        self.spills[v] = UOp.const(offset, dtypes.int32)
        self.stack_size = offset + sz
      r = alloc(cons if cons is not None else v.cons, i)
      self.insert_before.setdefault(i, []).append((v, r))
@@ -84,7 +84,7 @@ class LinearScanRegallocContext:

      # allocate stack array
      if u.op is Ops.BUFFER:
-        self.locals[u] = UOp.cconst(self.stack_size, dtypes.int32)
+        self.locals[u] = UOp.const(self.stack_size, dtypes.int32)
        self.stack_size += u.max_numel() * u.dtype.itemsize

      # loop prologue, avoid loading inside the loop
@@ -125,7 +125,7 @@ def regalloc_rewrite(ctx:LinearScanRegallocContext, x:UOp):
  # alloc/dealloc stack
  if ctx.stack_size > 0:
    sp = ctx.ren.stack_pointer()
-    offset = UOp.cconst(ctx.stack_size, sp.dtype)
+    offset = UOp.const(ctx.stack_size, sp.dtype)
    if i == 0: before = [ctx.ren.isel_matcher.rewrite(UOp(Ops.SUB, src=(sp, offset), tag=sp.tag))] + before
    elif i == len(ctx.uops) - 2: before += [ctx.ren.isel_matcher.rewrite(UOp(Ops.ADD, src=(sp, offset), tag=sp.tag))]

@@ -1,11 +1,12 @@
 from __future__ import annotations
 import math, itertools
-from typing import cast
+from collections import defaultdict
+from typing import cast, Final
 from tinygrad.uop.ops import Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, remove_all_tags
 from tinygrad.uop.ops import axis_letters, axis_colors, axis_to_pos
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes, Invalid
-from tinygrad.helpers import colored, getenv, DEBUG, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten
+from tinygrad.helpers import colored, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten
 from tinygrad.helpers import ALLOW_TF32, count, Context
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError, check
 from tinygrad.codegen.simplify import pm_flatten_range
@@ -47,6 +48,7 @@ class Scheduler:
    if hasattr(self, 'tensor_core'): ret.tensor_core = self.tensor_core
    return ret

+  kernel_cnt: Final[defaultdict[str, int]] = defaultdict(int)
  def get_optimized_ast(self, name_override:str|None=None) -> UOp:
    if name_override is not None: name = name_override
    else:
@@ -54,6 +56,9 @@ class Scheduler:
      special_uops = sorted([x for x in self.ast.toposort() if x.op is Ops.SPECIAL], key=lambda x: x.arg)
      special_ops = [colored(str(x.vmax+1), "blue" if x.arg[0] == "g" else "cyan") for x in special_uops]
      name = k_type + colored('_', 'BLACK').join(['']+special_ops+[colored(x.src[0].render(), color) for x,color in zip(self.rngs, self.colors())])
+      Scheduler.kernel_cnt[(function_name := to_function_name(name))] += 1
+      num = f"n{Scheduler.kernel_cnt[function_name]-1}" if Scheduler.kernel_cnt[function_name] > 1 else ""
+      name += colored(num, 'BLACK')
    self.ast = graph_rewrite(self.ast, pm_flatten_range, name="flatten range")
    return self.ast.replace(arg=KernelInfo(name=name, applied_opts=tuple(self.applied_opts), dont_use_locals=self.dont_use_locals), tag=1)

@@ -1,13 +1,12 @@
-import math, time, traceback, signal
+import math, time, multiprocessing, traceback, signal, atexit
 from dataclasses import replace
 from tinygrad.uop.ops import sym_infer, AxisType, UOp, Ops
 from tinygrad.uop.render import pyrender
 from tinygrad.device import Device, Buffer
-from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, colored, time_to_str
+from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str
 from tinygrad.helpers import IGNORE_BEAM_CACHE
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.engine.realize import time_call
-from tinygrad.engine.worker import get_worker_pool, terminate_worker_pool
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt.postrange import Scheduler

@@ -43,9 +42,9 @@ def _time_program(prg:UOp, var_vals:dict[str, int], rawbufs:list[Buffer], early_
    global_size, factor = get_test_global_size(prg.arg.global_size, max_global_size, var_vals)
    prg = prg.replace(arg=replace(prg.arg, global_size=tuple(global_size)))
  call = prg.call(*[UOp.from_buffer(b) for b in rawbufs])
-  tms, timer = [], time_call(call, var_vals, timeout=timeout, clear_l2=clear_l2)
+  tms = []
  for _ in range(cnt):
-    try: tms.append(next(timer) * factor)
+    try: tms.append(time_call(call, var_vals, timeout=timeout, clear_l2=clear_l2) * factor)
    except AssertionError: return [math.inf] * cnt
    if early_stop is not None and early_stop < min(tms): break
  return tms
@@ -79,6 +78,11 @@ def _try_compile(x:tuple[int,Scheduler]) -> tuple[int, tuple[UOp, float]|None]:
    if hasattr(signal, "alarm"): signal.alarm(0)
  return x[0], ret

+# 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__()
+  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]

 # *** external API ***
@@ -107,8 +111,9 @@ def get_kernel_actions(s:Scheduler, include_0=True, max_up:int|None=None) -> dic
    except KernelOptError: pass
  return acted

-BEAM_DEBUG = getenv("BEAM_DEBUG")
+beam_pool, BEAM_DEBUG = None, getenv("BEAM_DEBUG")
 def beam_search(s:Scheduler, rawbufs:list[Buffer], var_vals:dict[str,int], amt:int, allow_test_size=True, disable_cache=IGNORE_BEAM_CACHE.value):
+  global beam_pool
  key = {"ast": s.ast.key, "amt": amt, "allow_test_size": allow_test_size, "device": s.ren.target.device, "suffix": s.ren.suffix}
  if not disable_cache and CACHELEVEL >= 1 and (val:=diskcache_get("beam_search", key)) is not None:
    ret = s.copy()
@@ -118,7 +123,11 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], var_vals:dict[str,int], amt:i
  beam: list[tuple[Scheduler, float]] = [(s, float("inf"))]
  seen_libs = set()

-  pool = get_worker_pool()
+  default_parallel = multiprocessing.cpu_count() if s.ren.target.device in {"CUDA", "AMD", "NV", "METAL", "HIP"} else 0
+  if beam_pool is None and (workers := getenv("PARALLEL", default_parallel)):
+    beam_pool = multiprocessing.get_context("spawn").Pool(workers, _init_worker, (), getenv("BEAM_MAX_TASKS_PER_CHILD", 16))
+    @atexit.register
+    def close_pool(): beam_pool.close()

  min_progress = getenv("BEAM_MIN_PROGRESS", 0.01)/1e6
  if BEAM_DEBUG:
@@ -134,7 +143,7 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], var_vals:dict[str,int], amt:i
      candidates: list[Scheduler] = flatten([get_kernel_actions(si, include_0=False).values() for si,_ in beam])
      timed: list[tuple[Scheduler, float]] = []
      least_compute_ops = math.inf
-      for i, proc in ((map if pool is None else pool.imap_unordered)(_try_compile, enumerate(candidates))):
+      for i, proc in ((map if beam_pool is None else beam_pool.imap_unordered)(_try_compile, enumerate(candidates))):
        if proc is None: continue
        prg, compile_et = proc
        if (lib:=prg.src[3].arg) in seen_libs: continue
@@ -170,7 +179,7 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], var_vals:dict[str,int], amt:i
        print(f"\r{time.perf_counter() - st:7.2f}s:", colored(time_to_str(beam[0][1], w=12), "green" if exiting else None),
              f"from {len(candidates):3d} -> {len(opts):3d} actions\033[K", beam[0][0].colored_shape())
  except KeyboardInterrupt as e:
-    terminate_worker_pool()
+    if beam_pool is not None: beam_pool.terminate()
    raise e

  if CACHELEVEL >= 1: diskcache_put("beam_search", key, beam[0][0].applied_opts)
@@ -1,7 +1,7 @@
 import itertools
 from typing import Callable
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, AxisType
-from tinygrad.uop.symbolic import symbolic, invalid_gate
+from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const, invalid_gate
 from tinygrad.helpers import partition
 from tinygrad.dtype import dtypes

@@ -32,7 +32,7 @@ def simplify_merge_adjacent(u:UOp) -> UOp|None:
        s0, s1 = r0.src[0], r1.src[0]
        # do the merge
        new_range = r0.replace(src=(s0*s1,))
-        nidx = graph_rewrite(u, _substitute+symbolic+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
+        nidx = graph_rewrite(u, _substitute+symbolic+pm_fold_cast_const+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
                             name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")

        # check if it simplifies
@@ -137,7 +137,7 @@ def reduce_collapse(red:UOp, u:UOp, pm:PatternMatcher=pm_reduce_collapse) -> UOp
    for u in included:
      for s in u.src:
        if s in included or s in replaces or s.op in {Ops.CONST, Ops.PARAM, Ops.BUFFER}: continue
-        replaces[s] = UOp.variable(f'in{len(replaces)}', s.vmin, s.vmax, s.dtype, param=True)
+        replaces[s] = UOp.variable(f'in{len(replaces)}', s.vmin, s.vmax, s.dtype)
    collapse_fxn = u.substitute(replaces).reduce(r, arg=Ops.ADD)
    sink = graph_rewrite(collapse_fxn, pm, name="reduce_collapse")
    if not no_range(sink): return None
@@ -1,8 +1,8 @@
 from __future__ import annotations
 from dataclasses import dataclass, replace
 from collections import defaultdict
-from typing import Any, Callable, Generic, TypeVar, Iterator, Generator, Self, TYPE_CHECKING
-import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal, subprocess, struct
+from typing import Any, Generic, TypeVar, Iterator, Generator, Self, TYPE_CHECKING
+import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal
 from tinygrad.helpers import LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing
 from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize, Target, unwrap, round_up
@@ -66,10 +66,10 @@ def canonicalize_device(device:str|tuple|list|None) -> str|tuple[str, ...]:
 class ProfileDeviceEvent(ProfileEvent): device:str; tdiff:decimal.Decimal=decimal.Decimal(0); props:dict[str,Any]|None=None # noqa: E702

@dataclass(frozen=True)
-class ProfileProgramEvent(ProfileEvent): device:str; name:str; lib:bytes|None; base:int|None; tag:int|None=None; profile_key:bytes|None=None # noqa: E702
+class ProfileProgramEvent(ProfileEvent): device:str; name:str; lib:bytes|None; base:int|None; tag:int|None=None # noqa: E702

@dataclass(frozen=True)
-class ProfileGraphEntry: device:str; name:str|TracingKey; st_id:int; en_id:int; profile_key:bytes|None=None # noqa: E702
+class ProfileGraphEntry: device:str; name:str|TracingKey; st_id:int; en_id:int # noqa: E702

@dataclass(frozen=True)
 class ProfileGraphEvent(ProfileEvent): ents:list[ProfileGraphEntry]; deps:list[list[int]]; sigs:list[decimal.Decimal] # noqa: E702
@@ -83,7 +83,6 @@ class BufferSpec:
  cpu_access: bool = False
  host: bool = False
  nolru: bool = False
-  zero: bool = False
  external_ptr: int|None = None

 class MultiBuffer:
@@ -104,7 +103,7 @@ class Buffer:
  def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:BufferSpec|None=None,
               initial_value:bytes|pickle.PickleBuffer|None=None, uop_refcount=0, base:Buffer|None=None, offset:int=0, preallocate=False):
    assert isinstance(dtype, DType)
-    self.device, self.size, self.dtype, self.options, self.offset, self.allocated_views = Device.canonicalize(device), size, dtype, options, offset, 0
+    self.device, self.size, self.dtype, self.options, self.offset, self.allocated_views = device, size, dtype, options, offset, 0
    self._bufs: dict[str, Any] = {}
    if base is None:
      assert offset == 0, "base buffers can't have offset"
@@ -117,7 +116,7 @@ class Buffer:
        if isinstance(initial_value, pickle.PickleBuffer): initial_value.release()
    else:
      assert base._base is None, "base can't have a base"
-      assert self.device == base.device, "base must have the same device"
+      assert device == base.device, "base must have the same device"
      self._base = base
    if preallocate: self.allocate()
  @property
@@ -134,7 +133,7 @@ class Buffer:
  # check if the underlying buffer is allocated, possibly from the base object
  def is_allocated(self) -> bool: return self.base.is_allocated() if self._base is not None else self.device in self._bufs
  def get_buf(self, device: str) -> Any:
-    if device not in self._bufs and (device:=Device.canonicalize(device)) not in self._bufs:
+    if device not in self._bufs:
      allocator = Device[device].allocator
      if device == self.device: self.ensure_allocated()
      elif self._base is not None: self._bufs[device] = allocator._offset(self._base.get_buf(device), self.nbytes, self.offset)
@@ -266,7 +265,7 @@ class LRUAllocator(Allocator, Generic[DeviceType]):
      for opaque in opaques: super().free(opaque, sz, options)
      opaques.clear()
  def free(self, opaque:Any, size:int, options:BufferSpec|None=None):
-    if LRU and (options is None or (not (options.nolru or options.zero) and options.external_ptr is None)): self.cache[(size, options)].append(opaque)
+    if LRU and (options is None or (not options.nolru and options.external_ptr is None)): self.cache[(size, options)].append(opaque)
    else: super().free(opaque, size, options)

 class DepsTracker:
@@ -311,14 +310,6 @@ class Compiler:
      if self.cachekey is not None: diskcache_put(self.cachekey, src, lib)
    return lib
  def disassemble(self, lib:bytes): pass
-  def server(self, cmd:str, arch:str, *args) -> subprocess.Popen:
-    argv = f"{cmd} {pathlib.Path(__file__).parent}/runtime/support/compileserver.py {type(self).__module__}:{type(self).__name__} {arch}"
-    return subprocess.Popen(argv.split() + [str(a) for a in args], stdout=subprocess.PIPE, stdin=subprocess.PIPE, bufsize=0)
-  def compile_server(self, src:str, proc:subprocess.Popen) -> bytes:
-    unwrap(proc.stdin).write(struct.pack("I", len(src.encode())) + src.encode())
-    if (lib:=unwrap(proc.stdout).read(struct.unpack("I", unwrap(proc.stdout).read(4))[0])): return lib
-    raise CompileError("Compilation Error")
-

@dataclass
 class TinyELF:
@@ -327,7 +318,6 @@ class TinyELF:
  target: Target
  # tuple of (name, slot, dtype, shape)
  signature: tuple[tuple[str|None, int, DType, tuple], ...]
-  profile_key: bytes|None = None

  @staticmethod
  def iter_sig(signature:tuple[tuple[str|None, int, DType, tuple], ...], offset:int=0) -> Generator[tuple[int, DType], None, None]:
@@ -341,18 +331,17 @@ class Program(Generic[DeviceType]):
               wait=False) -> float|None: pass

 class Compiled:
-  ifaces:list[Callable] = []
  profile_events:list[ProfileEvent] = [ProfileDeviceEvent("CPU")] # NOTE: CPU is the default device.

-  has_copy_queue:bool = True
-
  pm_lower:Any = None
  pm_bufferize:Any = None

+  has_copy_queue:bool = True
+
  def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime:type[Program[Self]]|None, graph=None, arch=None):
    from tinygrad.renderer import Renderer
    self.device, self.allocator, self.runtime_t, self.graph, self.renderers = device, allocator, runtime, graph, renderers or [Renderer]
-    self.device_id, self.arch = (int(idx) if ":" in device and (idx:=device.split(":")[1]).isdigit() else 0), arch
+    self.arch = arch
    self.cached_renderer:dict[Any, Renderer] = {}

  @property
@@ -375,21 +364,11 @@ class Compiled:
    return select_first_inited(select_by_name(self.renderers, self._renderer_name, t.renderer, f"{self.device} has no renderer {t.renderer!r}"),
                               f"No renderer for {self.device} is available", self.cached_renderer, t)

-  def _select_iface(self, device:str):
-    self.device_id = int(device.split(":")[1]) if ":" in device else 0
-    assert (v:=getenv(k:=f'{type(self).__name__[:-6].upper()}_IFACE', "")) == "",  \
-      f"{k}={v} is deprecated, use DEV={replace(DEV.target(type(self).__name__[:-6]), interface=v)} instead"
-    t = DEV.target(dev:=type(self).__name__[:-6])
-    filtered = select_by_name(self.ifaces, lambda i: i.__name__[:-5], t.interface, f"{dev} has no interface {t.interface!r}")
-    filtered = [i for i in filtered if t.interface.startswith("MOCK") or not i.__name__[:-5].startswith("MOCK")] # never fallback to mock ifaces
-    return select_first_inited([functools.partial(iface, self, self.device_id) for iface in filtered],
-                               f"No interface for {dev}:{self.device_id} is available")
-
  def count(self) -> int:
    """
    Returns the number of physical accelerators available to the runtime.
    """
-    return self.iface.count if hasattr(self, 'iface') else 1
+    return 1

  def synchronize(self):
    """
@@ -407,7 +386,7 @@ class Compiled:
    """
    Called at the end of process lifetime to allow the device to finalize.
    """
-    if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini()
+    # override this in your device implementation

 if PROFILE:
  @atexit.register
@@ -429,7 +408,7 @@ def enumerate_devices_str() -> Generator[str, None, None]:
    ren_results, iface_results = [], []
    try:
      d = Device[device]
-      for iface in [i for i in d.ifaces if not i.__name__.startswith("MOCK")]:
+      for iface in [i for i in getattr(d, 'ifaces', []) if not i.__name__.startswith("MOCK")]:
        try:
          name = iface.__name__[:-5]
          default_text, count = ("(default)", d.count()) if type(d.iface) is iface else (f"(DEV={name}+{device} to make default)", iface(d, 0).count) # type: ignore
@@ -66,6 +66,7 @@ class DType(metaclass=DTypeMetaClass):
  def __reduce__(self): return type(self), tuple(getattr(self, f.name) for f in fields(self))
  def __repr__(self): return f"dtypes.{INVERSE_DTYPES_DICT[self.name]}"
  def __lt__(self, o:DType): return (self.priority, self.bitsize, self.name, self.fmt) < (o.priority, o.bitsize, o.name, o.fmt)
+  def scalar(self) -> DType: return self
  @functools.cached_property
  def min(self):
    if dtypes.is_int(self): return 0 if dtypes.is_unsigned(self) else -2**(self.bitsize-1)
@@ -221,7 +222,7 @@ def float_to_fp16(x):

 def float_to_bf16(x):
  if not math.isfinite(x): return x
-  u = struct.unpack('I', struct.pack('f', truncate[dtypes.float](x)))[0]
+  u = struct.unpack('I', struct.pack('f', x))[0]
  u = (u + 0x7FFF + ((u >> 16) & 1)) & 0xFFFF0000
  return struct.unpack('f', struct.pack('I', u))[0]

@@ -44,7 +44,7 @@ def graph_split_rewrite(linear:UOp, max_batch_size:int=0) -> UOp:
    current_batch, current_batch_devs = [], []

  for si in linear.src:
-    devs = dedup([Device[x] for b in si.src[1:] if not b.is_bound_var for x in (b.device if isinstance(b.device, tuple) else (b.device,))])
+    devs = dedup([Device[x] for b in si.src[1:] if b.op is not Ops.BIND for x in (b.device if isinstance(b.device, tuple) else (b.device,))])
    graph_t = graph_class(devs[0]) if devs[0].graph is not None else None

    can_graph = graph_t is not None and graph_t.supports_uop(devs, si)
@@ -269,14 +269,10 @@ class _TinyJit(Generic[ReturnType]):
        big_linear, onetime_linear = prune_linear(big_linear, set(input_buf_uops))
        if DEBUG >= 1: print(f"pruned from {len(big_linear.src) + len(onetime_linear.src)} -> {len(big_linear.src)} kernels")
        run_linear(onetime_linear, var_vals)
-        del onetime_linear

      # hold all buffers reachable from live Tensors (e.g. lazy .grad created during capture), the memory planner can't suballocate those
      held_bufs = set(buffers) | {u for tref in list(all_tensors) if (t:=tref()) is not None for u in t.uop.toposort() if u.op is Ops.BUFFER}
      linear = jit_lower(big_linear, held_bufs, input_buf_uops)
-      # drop the pre-planning graph: it keeps the whole capture-time working set allocated (big_linear) or referenced (held_bufs).
-      # the planned linear only uses the arena/held buffers, so the intermediates must be freed before linking and first exec
-      del big_linear, held_bufs
      self.captured = CapturedJit(ret, linear, names, expected_input_info)
      ret = self.captured(input_buf_uops, var_vals)
    elif self.cnt >= 2:
@@ -1,23 +1,18 @@
 from __future__ import annotations
 from typing import cast, Iterator, Any, Sequence
-import random, itertools, math, weakref, array, decimal
+import time, random, itertools, math, contextlib, weakref, array
 from dataclasses import dataclass, replace, field
-from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansipad, all_int, prod, flatten, Context, getenv, to_tuple, tqdm
-from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us
-from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite, ProgramInfo
+from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple
+from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, wait_cond
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite
 from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
-from tinygrad.dtype import dtypes
-from tinygrad.renderer import Estimates, Renderer
-from tinygrad.codegen import to_program, to_program_cache, to_program_key, to_program_context
+from tinygrad.renderer import Estimates
+from tinygrad.codegen import to_program
 from tinygrad.codegen.opt.postrange import args_from_ast
-from tinygrad.engine.worker import get_worker_pool, terminate_worker_pool

 # **************** Helpers ****************

-def get_call_arg_uops(call:UOp) -> tuple[UOp, ...]: return tuple(s for s in call.src[1:] if not s.is_bound_var)
-def get_call_var_uops(call:UOp, prg:UOp) -> list[UOp]:
-  bound = {s.src[0].expr: s.src[1].src[1] for s in call.src[1:] if s.is_bound_var}
-  return [bound.get(v.expr, v) for v in prg.arg.vars]
+def get_call_arg_uops(call:UOp) -> tuple[UOp, ...]: return tuple(s for s in call.src[1:] if s.op is not Ops.BIND)

 def get_call_outs_ins(call:UOp) -> tuple[tuple[int, ...], tuple[int, ...]]:
  ast = call.src[0]
@@ -26,12 +21,6 @@ def get_call_outs_ins(call:UOp) -> tuple[tuple[int, ...], tuple[int, ...]]:
  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return (0,), tuple(range(1, len(get_call_arg_uops(call))))
  return (), ()

-def get_call_kernels(call:UOp) -> list[tuple[str, UOp]]:
-  if (ast:=call.src[0]).op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq": return [(d, k) for devs, k, _ in call.arg.aux.kernels for d in devs]
-  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return [(to_tuple(ast.device)[0], call)]
-  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "validate": return []
-  return [(d, call) for d in to_tuple(call.src[1].device)]
-
 def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|None=None) -> str:
  def _uop_sz_to_str(uop:UOp) -> str: return size_to_str(sym_infer(prod(uop.shape) * uop.dtype.itemsize, var_vals or {}))
  def _dev_str(buf:Buffer|UOp) -> str: return ', '.join(d[:7] for d in to_tuple(buf.device))
@@ -47,52 +36,49 @@ def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|N
 # **************** Stat ****************

 def estimate_uop(call:UOp) -> Estimates:
-  if (ast:=call.src[0]).op is Ops.PROGRAM: return ast.src[0].arg.estimates or Estimates()
+  ast = call.src[0]
+  if ast.op is Ops.PROGRAM: return ast.src[0].arg.estimates or Estimates()
  if ast.op is Ops.COPY or (ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec"):
-    return Estimates(lds=(nbytes:=prod(call.src[1].shape) * call.src[1].dtype.itemsize), mem=nbytes)
+    nbytes = prod(call.src[1].shape) * call.src[1].dtype.itemsize
+    return Estimates(lds=nbytes, mem=nbytes)
  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return get_graph_runtime(ast).estimates
  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq": return call.arg.aux.estimates
  return Estimates()

 first_run_cache:set[bytes] = set()
-def track_stats(ctx:ExecContext, call:UOp, st:decimal.Decimal, ets:list[float|None]):
-  if ctx.update_stats:
-    is_hcq = (ast:=call.src[0]).op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq"
-    estimates, n = estimate_uop(call), 1 if is_hcq else len(get_call_kernels(call))
-    GlobalCounters.kernel_count += len(call.arg.aux.kernels) if is_hcq else n
-    GlobalCounters.global_ops += n*sym_infer(estimates.ops, ctx.var_vals)
-    GlobalCounters.global_mem += n*sym_infer(estimates.mem, ctx.var_vals)
-    GlobalCounters.time_sum_s += sum(et for et in ets if et is not None)
-  if DEBUG < 2 and not PROFILE: return
+@contextlib.contextmanager
+def track_stats(ctx:ExecContext, call:UOp, device:str, bufs:list[Buffer], var_vals:dict[str, int]):
+  if PROFILE:
+    outputs, inputs = get_call_outs_ins(call)
+    cpu_events.append(ProfilePointEvent(device, "exec", len(cpu_events), {"var_vals": var_vals,
+      "bufs": [b.trace_num for b in bufs], "name": get_call_name(call, bufs, var_vals), "outputs": outputs, "inputs": inputs}))
+  et: list[float|None] = [None]
+  if DEBUG >= 2: st = time.perf_counter()
+  yield et
+  if not ctx.update_stats: return

-  kernels = get_call_kernels(call) # everything below is the per kernel display: exec events for the profiler and DEBUG=2 lines
-  args = resolve_params(call, ctx.input_uops) if kernels and kernels[0][1] is call else []
-  lanes = list(unwrap_multi(call, [args[g] for g in call.src[0].arg.globals] if call.src[0].op is Ops.PROGRAM else args)) if args else []
-  for i, (device, kcall) in enumerate(kernels):
-    et, bufs = ets[i] if i < len(ets) else None, lanes[i][0] if i < len(lanes) else []
-    if PROFILE: # backdate the event to the start of the call, the viz matches a device range with the exec event before it
-      outputs, inputs = get_call_outs_ins(kcall)
-      cpu_events.append(ProfilePointEvent(device, "exec", len(cpu_events), {"var_vals": ctx.var_vals,
-        "bufs": [b.trace_num for b in bufs], "name": get_call_name(kcall, bufs, ctx.var_vals), "outputs": outputs, "inputs": inputs}, ts=st))
-    if DEBUG < 2 or not ctx.update_stats: continue
-    if et is None:
-      Device[device].synchronize()
-      et, st = float(perf_counter_us() - st)*1e-6, perf_counter_us()
-      GlobalCounters.time_sum_s += et
+  if DEBUG >= 2 and et[0] is None:
+    Device[device].synchronize()
+    et[0] = time.perf_counter() - st

-    estimates = estimate_uop(kcall)
-    display_name = get_call_name(kcall, bufs, ctx.var_vals)
-    op_est, mem_est, lds_est = (sym_infer(x, ctx.var_vals) for x in (estimates.ops, estimates.mem, estimates.lds))
-    header_color = 'magenta' if ctx.jit else ('green' if kcall.src[0].key not in first_run_cache else None)
-    ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
-    flops, membw, ldsbw = op_est/(et or 1e-20), mem_est/(et or 1e-20), lds_est/(et or 1e-20)
+  estimates = estimate_uop(call)
+  GlobalCounters.kernel_count += 1
+  GlobalCounters.global_ops += (op_est:=sym_infer(estimates.ops, var_vals))
+  GlobalCounters.global_mem += (mem_est:=sym_infer(estimates.mem, var_vals))
+  if et[0] is not None: GlobalCounters.time_sum_s += et[0]
+  if DEBUG >= 2:
+    display_name = get_call_name(call, bufs, var_vals)
+    lds_est = sym_infer(estimates.lds, var_vals)
+    header_color = 'magenta' if ctx.jit else ('green' if call.src[0].key not in first_run_cache else None)
+    ptm = colored(time_to_str(et[0], w=9), "yellow" if et[0] > 0.01 else None) if et[0] is not None else ""
+    flops, membw, ldsbw = op_est/(et[0] or 1e-20), mem_est/(et[0] or 1e-20), lds_est/(et[0] or 1e-20)
    flops_str = f"{flops*1e-9:7.0f} GFLOPS" if flops < 1e14 else colored(f"{flops*1e-12:7.0f} TFLOPS", 'green')
    mem_str = f"{membw*1e-9:4.0f}|{ldsbw*1e-9:<6.0f} GB/s" if membw < 1e13 and ldsbw < 1e15 else \
      colored(f"{membw*1e-12:4.0f}|{ldsbw*1e-12:<6.0f} TB/s", 'green')
    print(f"{colored(f'*** {device[:7]:7s} {GlobalCounters.kernel_count:4d}', header_color)}"+
-      f" {ansipad(display_name, 46)} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:6.2f} GB"+
-      ("" if et is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})"))
-    first_run_cache.add(kcall.src[0].key)
+      f" {display_name+' '*(46-ansilen(display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:6.2f} GB"+
+      ("" if et[0] is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})"))
+    first_run_cache.add(call.src[0].key)

 local_size_cache: dict[bytes, tuple[int, ...]] = {}
 def optimize_local_size(call:UOp, prg:UOp) -> UOp|None:
@@ -165,31 +151,32 @@ def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], d
                   for x in call.src[0].toposort())
    for j, per_dev in enumerate(zip(*[cast(MultiBuffer, b).bufs for b in bufs])): yield list(per_dev), {"_device_num": j} if has_dnum else {}

-def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
+def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
  for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)):
    dest, src = bufs[0].ensure_allocated(), bufs[1].ensure_allocated()
-    if hasattr(dest.allocator,'_transfer') and dest.allocator.supports_transfer and dest.device.split(":")[0] == src.device.split(":")[0]:
-      dest.allocator._transfer(dest._buf, src._buf, dest.nbytes, src_dev=src.allocator.dev, dest_dev=dest.allocator.dev)
-    elif src.device.startswith("DISK") and getattr(src.allocator.dev, 'fd', None) is not None \
-         and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096 and dest.allocator.supports_copy_from_disk:
-      dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes)
-    elif hasattr(dest.allocator, '_as_buffer'): src.allocator._copyout(dest.as_memoryview(force_zero_copy=True), src._buf)
-    else: dest.allocator._copyin(dest._buf, src.as_memoryview(allow_zero_copy=True))
-  return []
+    with track_stats(ctx, call, dest.device, [dest, src], ctx.var_vals):
+      if hasattr(dest.allocator,'_transfer') and dest.allocator.supports_transfer and dest.device.split(":")[0] == src.device.split(":")[0]:
+        dest.allocator._transfer(dest._buf, src._buf, dest.nbytes, src_dev=src.allocator.dev, dest_dev=dest.allocator.dev)
+      elif src.device.startswith("DISK") and getattr(src.allocator.dev, 'fd', None) is not None \
+           and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096 and dest.allocator.supports_copy_from_disk:
+        dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes)
+      elif hasattr(dest.allocator, '_as_buffer'): src.allocator._copyout(dest.as_memoryview(force_zero_copy=True), src._buf)
+      else: dest.allocator._copyin(dest._buf, src.as_memoryview(allow_zero_copy=True))
+  return None

-def exec_kernel(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
-  ets:list[float|None] = []
-  resolved = resolve_params(call, ctx.input_uops)
-  for device, (bufs, device_vars) in zip(to_tuple(call.src[1].device), unwrap_multi(call, [resolved[i] for i in ast.arg.globals])):
+def exec_kernel(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
+  et = None
+  for device, (bufs, device_vars) in zip(to_tuple(call.src[1].device), unwrap_multi(call, resolve_params(call, ctx.input_uops))):
    var_vals = {**ctx.var_vals, **device_vars}
-    prg_bufs = [b.ensure_allocated() for b in bufs]
+    prg_bufs = [bufs[i].ensure_allocated() for i in ast.arg.globals]
    rt = get_runtime(device, ast, cache=ctx.cache)
    global_size, local_size = ast.arg.launch_dims(var_vals)
-    ets.append(rt(*[b.get_buf(device) for b in prg_bufs], global_size=global_size, local_size=local_size, vals=ast.arg.vals(var_vals),
-                  wait=ctx.wait, timeout=ctx.timeout))
-  return ets
+    with track_stats(ctx, call, device, prg_bufs, var_vals) as tm:
+      et = tm[0] = rt(*[b.get_buf(device) for b in prg_bufs], global_size=global_size, local_size=local_size, vals=ast.arg.vals(var_vals),
+                       wait=ctx.wait, timeout=ctx.timeout)
+  return et

-def exec_validate(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
+def exec_validate(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
  import numpy as np
  for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)):
    bufs, dev_bufs = bufs[:len(bufs)//2], bufs[len(bufs)//2:]
@@ -198,36 +185,45 @@ def exec_validate(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
    global_size, local_size = prg.arg.launch_dims(var_vals)
    cpu_rt(*[bufs[i].ensure_allocated()._buf for i in prg.arg.globals], global_size=global_size, local_size=local_size, vals=prg.arg.vals(var_vals))
    for i in prg.arg.outs: np.testing.assert_allclose(dev_bufs[i].ensure_allocated().numpy(), bufs[i].numpy(), rtol=1e-3, atol=1e-3)
-  return []
+  return None

-def exec_encdec(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
+def exec_encdec(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
  bufs = [cast(Buffer, b.buffer).ensure_allocated() for b in resolve_params(call, ctx.input_uops)]
  shape, pos_var = tuple(s.val for s in ast.src if s.op is Ops.CONST), ast.variables()[0].expr
-  bufs[0].allocator._encode_decode(bufs[0]._buf, bufs[1]._buf, bufs[2]._buf, [x._buf for x in bufs[3:]], shape, ctx.var_vals[pos_var])
-  return []
+  with track_stats(ctx, call, bufs[0].device, bufs, ctx.var_vals):
+    bufs[0].allocator._encode_decode(bufs[0]._buf, bufs[1]._buf, bufs[2]._buf, [x._buf for x in bufs[3:]], shape, ctx.var_vals[pos_var])
+  return None

-def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
-  return [get_graph_runtime(ast, ctx.input_uops)(ctx.input_uops, ctx.var_vals, wait=ctx.wait)]
+def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
+  rt = get_graph_runtime(ast, ctx.input_uops)
+  with track_stats(ctx, call, rt.device, [], ctx.var_vals) as t: t[0] = rt(ctx.input_uops, ctx.var_vals, wait=ctx.wait)
+  return t[0]

-def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
-  dev = cast(Any, Device[(info:= call.arg.aux).device[0]])
-  addrs = [(b.bufs[j] if isinstance(b:=_resolve(ctx.input_uops[k], ctx.input_uops).buffer, MultiBuffer) else b).get_buf(dev_name).va_addr
-           for devs, idxs in info.input_idxs for j, dev_name in enumerate(devs) for k in idxs]
-  dev.rt_buffer()._buf.cpu_view().view(offset=(base:=dev.rt_allocator.alloc(len(addrs) * 8)), fmt='Q')[:len(addrs)] = array.array('Q', addrs)
+def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
+  if (info:=call.arg.aux).inputs is not None:
+    bufs = [_resolve(ctx.input_uops[i], ctx.input_uops).buffer for i in call.arg.aux.input_idxs]
+    table = call.src[1+info.inputs].buffer
+    for j,dev in enumerate(call.arg.aux.device):
+      addrs = array.array('Q', [(b.bufs[j] if isinstance(b, MultiBuffer) else b).get_buf(dev).va_addr for b in bufs])
+      mv = (table.bufs[j] if isinstance(table, MultiBuffer) else table).ensure_allocated()._buf.cpu_view().view(fmt='Q')
+      wait_cond(lambda: mv[0], value=0, timeout_ms=ctx.timeout or getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000), msg=f"{dev} hang detected")
+      mv[:len(addrs)] = addrs

-  if info.inputs is not None:
-    tables = [UOp.from_buffer(dev.rt_buffer().view(len(idxs), dtypes.uint64, base + j*len(idxs)*8), HCQ_RUNTIME_DEV.value)
-              for devs, idxs in info.input_idxs for j in range(len(devs))]
-    call = call.substitute({call.src[1+info.inputs]: UOp.mstack(*tables)})
-  exec_kernel(replace(ctx, var_vals={**ctx.var_vals, "hcq_inputs_ptr": dev.rt_buffer()._buf.va_addr + base}), call, ast)
+  exec_kernel(replace(ctx, update_stats=DEBUG>=3), call, ast)

-  def _prof_tm(device:str, stat_call:UOp, prof:tuple[int, ...]) -> float|None:
-    (d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, stat_call.arg.name, prof[0], prof[1], stat_call.key)
-    if not ctx.wait: return None
-    d.synchronize(timeout=ctx.timeout)
-    st, en = (d.signal(x)._buf.cpu_view().view(fmt='Q')[0] for x in prof)
-    return float(en-st)/d.timestamp_divider/1e6
-  return [_prof_tm(device, k, prof) for devices, k, prof in info.kernels if prof for device in devices] if PROFILE or ctx.wait else []
+  tms = []
+  for devices,name,estimates,prof in info.kernels:
+    for device in devices:
+      tm = None
+      if prof:
+        (d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, name, *prof)
+        if ctx.wait:
+          d.synchronize(timeout=ctx.timeout)
+          st, en = (d.signal(x)._buf.cpu_view().view(fmt='Q')[0] for x in prof)
+          tms.append(tm:=float(en-st)/d.timestamp_divider/1e6)
+      stat_call = call.replace(arg=replace(call.arg, name=name, aux=replace(info, estimates=estimates, kernels=())))
+      with track_stats(ctx, stat_call, device, [], ctx.var_vals) as et: et[0] = tm
+  return max(tms) if tms else None

 # flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
 pm_flatten_linear = PatternMatcher([
@@ -248,44 +244,10 @@ pm_beam = PatternMatcher([
   lambda ctx,call,sink: call.replace(src=(sink.replace(arg=replace(sink.arg, beam=ctx)), *call.src[1:])) if sink.arg.beam == 0 else None),
 ])

-# **************** parallel lowering + compilation ****************
-
-def _compile_kernel(x:tuple[int, tuple[UOp, Renderer], dict]) -> tuple[int, UOp]:
-  with Context(**x[2]): return x[0], to_program(*x[1])
-
-def _get_call_to_compile(c:UOp) -> tuple[UOp, Renderer]|None:
-  ast = a0.src[0] if (a0:=c.src[0]).op is Ops.CUSTOM_FUNCTION and a0.arg == "hcq" else a0
-  # a PROGRAM with a ProgramInfo and a BINARY is already compiled
-  if ast.op is Ops.SINK or (ast.op is Ops.PROGRAM and not (isinstance(ast.arg, ProgramInfo) and ast.src[-1].op is Ops.BINARY)):
-    return ast, Device[c.device if isinstance(c.device, str) else c.device[0]].renderer
-  return None
-
-def lower_and_compile(linear:UOp) -> UOp:
-  # collect the kernels to lower and compile, deduped by their compile cache key
-  if not len(ar:={c: a for c in linear.toposort() if c.op is Ops.CALL and (a:=_get_call_to_compile(c)) is not None}): return linear
-
-  # lower and compile what's not cached, in parallel if there's a worker pool
-  keys = {c: to_program_key(*a) for c, a in ar.items()}
-  todo = list({keys[c]: a for c, a in ar.items() if keys[c] not in to_program_cache}.items())
-  if len(todo):
-    # kernels that beam search must compile in the parent, beam needs device access to time candidates
-
-    pool = None if len(todo) == 1 or any(getattr(c.src[0].arg, "beam", 0) for c in ar) else get_worker_pool()
-    ctx = {v.key: v.value for v in to_program_context}
-    tasks = ((i, ast_ren, ctx) for i, (_, ast_ren) in enumerate(todo))
-    try:
-      with tqdm(total=len(todo), desc="compiling", disable=DEBUG<1) as pbar:
-        for i, prg in (map if pool is None else pool.imap_unordered)(_compile_kernel, tasks):
-          pbar.set_description(f"compiling {ansipad(prg.src[0].arg.name, 40)}")
-          to_program_cache[todo[i][0]] = prg
-          pbar.update(1)
-    except KeyboardInterrupt:
-      if pool is not None: terminate_worker_pool()
-      raise
-
-  # swap the compiled PROGRAMs into the calls
-  return linear.substitute({c: c.replace(src=(c.src[0].substitute({a[0]: to_program_cache[keys[c]]}), *c.src[1:])) for c, a in ar.items()},
-                           name="precompile kernels")
+pm_compile = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.PROGRAM), name="ast"),), name="call", allow_any_len=True), lambda call,ast:
+    call.replace(src=(to_program(ast, Device[call.device if isinstance(call.device, str) else call.device[0]].renderer), *call.src[1:]))),
+])

 pm_optimize_local_size = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), optimize_local_size),
@@ -300,15 +262,14 @@ pm_exec = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="validate", name="ast"),), name="call", allow_any_len=True), exec_validate),
 ])

-if getenv("HCQ2"): from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link, HCQ_RUNTIME_DEV # noqa: E402 # down here, hcq2 imports realize
+if getenv("HCQ2"): from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link # noqa: E402 # down here, hcq2 imports the helpers above

 def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, profile:bool|None=None) -> UOp:
  if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
  if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
-  linear = lower_and_compile(linear)
-  linear = graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
+  linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True)
  if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
-  return linear
+  return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)

 def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache) if getenv("HCQ2") else linear

@@ -316,15 +277,14 @@ def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:Sequenc
  inputs = list(input_uops)
  if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU, input_uops=inputs))
  ctx = ExecContext(var_vals or {}, tuple(inputs), update_stats, jit, wait or DEBUG>=2)
-  for call in linear.src: track_stats(ctx, call, perf_counter_us(), pm_exec.rewrite(call, ctx))
+  for call in linear.src: pm_exec.rewrite(call, ctx)

-def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None, clear_l2:bool=False) -> Iterator[float]:
+def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None, clear_l2:bool=False) -> float:
+  if clear_l2:
+    if hasattr(dev:=Device[call.src[1].device], 'invalidate_caches'): dev.invalidate_caches()
+    else:
+      from tinygrad.tensor import Tensor
+      with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024, 1024).contiguous().realize(do_update_stats=False)
  ctx = ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False)
  linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True), cache=ctx.cache)
-  while True:
-    if clear_l2:
-      if hasattr(dev:=Device[call.src[1].device], 'invalidate_caches'): dev.invalidate_caches()
-      else:
-        from tinygrad.tensor import Tensor
-        with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024, 1024).contiguous().realize(do_update_stats=False)
-    yield max(et for c in linear.src for et in pm_exec.rewrite(c, ctx) or [0.0])
+  return max(pm_exec.rewrite(c, ctx) or 0.0 for c in linear.src)
@@ -1,49 +0,0 @@
-import multiprocessing, atexit, signal, sys, threading, contextlib
-from multiprocessing.context import SpawnContext, SpawnProcess
-from tinygrad.helpers import Context, getenv, PARALLEL
-
-# generic pool of worker processes for parallel compilation, shared by kernel lowering and BEAM search
-
-# 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__()
-  signal.signal(signal.SIGINT, signal.SIG_IGN)
-
-# spawn normally reimports the user's __main__ before _init_worker. This replays top-level code and can recursively create pools. There is no public
-# multiprocessing switch to skip that import, so hide the two attributes used to locate __main__ while each worker (including replacements) starts.
-_spawn_lock, _missing = threading.Lock(), object()
-@contextlib.contextmanager
-def _without_main():
-  main = sys.modules.get("__main__")
-  if main is None:
-    yield
-    return
-  with _spawn_lock:
-    saved = {name:getattr(main, name, _missing) for name in ("__file__", "__spec__")}
-    try:
-      for name in saved: setattr(main, name, None)
-      yield
-    finally:
-      for name,value in saved.items(): delattr(main, name) if value is _missing else setattr(main, name, value)
-
-class _WorkerProcess(SpawnProcess):
-  @staticmethod
-  def _Popen(process_obj):
-    with _without_main(): return SpawnProcess._Popen(process_obj)
-
-class _WorkerContext(SpawnContext): Process = _WorkerProcess
-
-worker_pool = None
-def get_worker_pool():
-  global worker_pool
-  if multiprocessing.current_process().daemon or PARALLEL == 0: return None
-  if worker_pool is None:
-    worker_pool = _WorkerContext().Pool(PARALLEL.value, _init_worker, (), getenv("BEAM_MAX_TASKS_PER_CHILD", 16))
-    @atexit.register
-    def close_pool(pool=worker_pool): pool.close()
-  return worker_pool
-
-def terminate_worker_pool():
-  global worker_pool
-  if worker_pool is not None: worker_pool.terminate()
-  worker_pool = None
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
George HotzandGitHub	26f2049d43	Merge branch 'master' into bitcast_spec	2026-08-12 19:24:33 -07:00
geohot	a8d13380f9	simpler, but not inverse	2026-06-25 10:24:01 -07:00
geohot	afb0463f10	fix bitcast	2026-06-25 10:11:13 -07:00
geohot	947d6a3c69	enforce a cast for bitcast	2026-06-25 10:00:14 -07:00