fix flaky

* remove pm_render

* move merge_reduce_ends

* ext devectorize

.github/actions/setup-tinygrad/action.yml

@@ -10,7 +10,7 @@ inputs:
     required: false
     default: ''  # if you don't set a key, it doesn't cache
   deps:
-    description: 'Extra dependency groups (comma separated)'
+    description: 'Extra dependency groups (space separated)'
     required: false
     default: ''
   pydeps:
@@ -41,10 +41,6 @@ inputs:
     description: "Install LLVM?"
     required: false
     default: 'false'
-  mesa:
-    description: "Install mesa (true, false, cpu)"
-    required: false
-    default: 'false'
   tinydreno:
     description: "Install tinydreno"
     required: false
@@ -114,7 +110,8 @@ runs:
       shell: bash
       run: |
         uv venv .venv
-        uv pip install --python .venv -e ".[${{ inputs.deps }}]" ${{ inputs.pydeps }} --torch-backend cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
+        DEPS="${{ inputs.deps }}"
+        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
@@ -146,11 +143,6 @@ runs:
         echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
         echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' | sudo tee -a /etc/apt/apt.conf.d/99keep-debs
 
-    - name: Add OpenCL Repo
-      if: inputs.opencl == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: echo "deb [ allow-insecure=yes ] https://apt.repos.intel.com/oneapi all main" | sudo tee /etc/apt/sources.list.d/oneAPI.list
-
     - name: Add AMD Repo (Linux)
       if: inputs.amd == 'true' && runner.os == 'Linux'
       shell: bash
@@ -176,10 +168,7 @@ runs:
         pkgs=""
         # **** OpenCL ****
         if [[ "${{ inputs.opencl }}" == "true" ]]; then
-          pkgs+=" opencl-headers \
-            intel-oneapi-runtime-openmp=2023.2.1-16 intel-oneapi-runtime-compilers-common=2023.2.1-16 intel-oneapi-runtime-compilers=2023.2.1-16 \
-            intel-oneapi-runtime-dpcpp-sycl-opencl-cpu=2023.2.1-16 intel-oneapi-runtime-tbb-common=2021.10.0-49541 \
-            intel-oneapi-runtime-tbb=2021.10.0-49541 intel-oneapi-runtime-opencl=2023.2.1-16"
+          pkgs+=" ocl-icd-opencl-dev"
         fi
         # **** AMD ****
         if [[ "${{ inputs.amd }}" == "true" ]]; then
@@ -286,18 +275,18 @@ runs:
       shell: bash
       run: brew install llvm@20
 
-    # **** mesa ****
-    - name: Install mesa (linux)
-      if: inputs.mesa != 'false' && runner.os == 'Linux'
-      shell: bash
-      run: sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa${{ inputs.mesa == 'cpu' && '_cpu' || '' }}-mesa-25.2.7-linux-amd64.so -o /usr/lib/libtinymesa${{ inputs.mesa == 'cpu' && '_cpu' || '' }}.so
-    - name: Install mesa (macOS)
-      if: inputs.mesa != 'false' && runner.os == 'macOS'
-      shell: bash
-      run: brew install sirhcm/tinymesa/tinymesa${{ inputs.mesa == 'cpu' && '_cpu' || '' }}
-
     # *** 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'
+      shell: bash
+      run: |
+        sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/rusticl-v1/libRusticlOpenCL.so.1.0.0 -o /usr/lib/libRusticlOpenCL.so
+        sudo mkdir -p /etc/OpenCL/vendors
+        echo "/usr/lib/libRusticlOpenCL.so" | sudo tee /etc/OpenCL/vendors/rusticl.icd
+        echo "RUSTICL_ENABLE=llvmpipe" >> "$GITHUB_ENV"

.github/workflows/benchmark.yml

@@ -81,8 +81,59 @@ jobs:
   #      source /tmp/tinygrad_pytest_ci/bin/activate
   #      pytest -nauto --durations=20
 
-  sharedbenchmarks:
-    name: Benchmark (DEV=${{ matrix.dev }})
+  llmbenchmark:
+    name: LLM (DEV=${{ matrix.dev }})
+    runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
+    strategy:
+      fail-fast: false
+      matrix:
+        dev: ['METAL', 'AMD', 'NV']
+    timeout-minutes: 60
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    env:
+      DEV: ${{ matrix.dev }}
+    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: 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: |
+        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 llama3.2
+      run: BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
+    - name: Run qwen3.5
+      # qwen3.5:35b-a3b doesn't fit on mac
+      if: ${{ matrix.dev != 'METAL' }}
+      run: BENCHMARK_LOG=qwen35_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.5:35b-a3b --benchmark --warmup
+    - name: Run olmoe
+      # 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
+
+  cifarbenchmark:
+    name: HLB-CIFAR10 (DEV=${{ matrix.dev }})
     runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
     strategy:
       fail-fast: false
@@ -111,6 +162,149 @@ jobs:
         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 10 CIFAR training steps
+      env:
+        ASSERT_MIN_STEP_TIME: ${{ matrix.dev == 'NV' && '130' || matrix.dev == 'AMD' && '200' || '3000' }}
+      run: BENCHMARK_LOG=cifar_10steps STEPS=10 python3 examples/hlb_cifar10.py
+    - name: Run 10 CIFAR training steps w HALF
+      env:
+        ASSERT_MIN_STEP_TIME: ${{ matrix.dev == 'NV' && '120' || matrix.dev == 'AMD' && '235' || '3000' }}
+      run: BENCHMARK_LOG=cifar_10steps_half STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
+    - name: Run full CIFAR training w 1 GPU
+      # 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
+
+  mlperfbenchmark:
+    name: MLPerf (${{ matrix.dev }})
+    runs-on: [self-hosted, Linux, "${{ 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 }}
+    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: Symlink models and datasets
+      run: |
+        mkdir -p extra/datasets
+        ln -s /raid/datasets/imagenet extra/datasets/imagenet
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: reset process replay
+      run: test/external/process_replay/reset.py
+    - name: 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
+
+  sdbenchmark:
+    name: Stable Diffusion (DEV=${{ matrix.dev }})
+    runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
+    strategy:
+      fail-fast: false
+      matrix:
+        dev: ['METAL', 'AMD', 'NV']
+    timeout-minutes: 60
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    env:
+      DEV: ${{ matrix.dev }}
+    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 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 Stable Diffusion
+      env:
+        ASSERT_MIN_STEP_TIME: ${{ matrix.dev == 'METAL' && '720' || matrix.dev == 'AMD' && '550' || '0' }}
+      run: BENCHMARK_LOG=stable_diffusion python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
+    - name: Run SDXL
+      if: ${{ matrix.dev != 'NV' }}
+      env:
+        ASSERT_MIN_STEP_TIME: ${{ matrix.dev == 'METAL' && '5000' || matrix.dev == 'AMD' && '3200' || '2000' }}
+      run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing
+    - 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' }}"]
+    strategy:
+      fail-fast: false
+      matrix:
+        dev: ['METAL', 'AMD', 'NV']
+    timeout-minutes: 60
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    env:
+      DEV: ${{ matrix.dev }}
+    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 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: Test tiny
+      run: |
+        DEBUG=2 python -m pytest -rA test/test_tiny.py
+        if [[ "${{ matrix.dev }}" == "NV" ]]; then
+          DEBUG=2 DEV=CUDA python -m pytest -rA test/test_tiny.py
+        fi
     - name: Test tensor cores
       run: |
         if [[ "${{ matrix.dev }}" == "METAL" ]]; then
@@ -150,109 +344,24 @@ jobs:
       env:
         HALF: ${{ matrix.dev == 'NV' && '1' || '0' }}
       run: CAPTURE_PROCESS_REPLAY=0 BIG=2 ${{ matrix.dev == 'METAL' && 'MPS=1' || 'TORCHCUDA=1' }} python3 test/speed/external_test_speed_v_torch.py
-    - name: Run Stable Diffusion
-      env:
-        ASSERT_MIN_STEP_TIME: ${{ matrix.dev == 'METAL' && '720' || matrix.dev == 'AMD' && '550' || '0' }}
-      run: BENCHMARK_LOG=stable_diffusion python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
-    - name: Run SDXL
-      if: ${{ matrix.dev != 'NV' }}
-      env:
-        ASSERT_MIN_STEP_TIME: ${{ matrix.dev == 'METAL' && '5000' || matrix.dev == 'AMD' && '3200' || '2000' }}
-      run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing
-    - name: Run llama3.2
-      run: BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
-    - name: Run qwen3.5
-      # qwen3.5:35b-a3b doesn't fit on mac
+    - name: Test speed vs theoretical
+      # no targets for METAL
       if: ${{ matrix.dev != 'METAL' }}
-      run: BENCHMARK_LOG=qwen35_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.5:35b-a3b --benchmark --warmup
-    - name: Run olmoe
-      # 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
+      run: IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
     - name: Train MNIST
       run: time TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
-    - name: Run 10 CIFAR training steps
-      env:
-        ASSERT_MIN_STEP_TIME: ${{ matrix.dev == 'NV' && '130' || matrix.dev == 'AMD' && '200' || '3000' }}
-      run: BENCHMARK_LOG=cifar_10steps STEPS=10 python3 examples/hlb_cifar10.py
-    - name: Run 10 CIFAR training steps w HALF
-      env:
-        ASSERT_MIN_STEP_TIME: ${{ matrix.dev == 'NV' && '120' || matrix.dev == 'AMD' && '230' || '3000' }}
-      run: BENCHMARK_LOG=cifar_10steps_half STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
-    - uses: actions/upload-artifact@v7
-      if: ${{ matrix.dev != 'AMD' }}
-      with:
-        name: Speed (${{ matrix.dev }})
-        path: |
-          onnx_inference_speed.csv
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-  tinyboxbenchmark:
-    name: Tinybox Benchmark (${{ matrix.dev }})
-    runs-on: [self-hosted, Linux, "${{ 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 }}
-    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: 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: test/external/process_replay/reset.py
-    - name: Test tiny
-      run: |
-        DEBUG=2 python -m pytest -rA test/test_tiny.py
-        if [[ "${{ matrix.dev }}" == "NV" ]]; then
-          DEBUG=2 DEV=CUDA python -m pytest -rA test/test_tiny.py
-        fi
-    - name: Test speed vs theoretical
-      run: IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
     - name: Test benchmark allreduce
       if: ${{ matrix.dev == 'NV' }}
       run: python test/external/external_benchmark_multitensor_allreduce.py
     - 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 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 w 1 GPU
-      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
-      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 (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
+    - uses: actions/upload-artifact@v7
+      if: ${{ matrix.dev != 'AMD' }}
+      with:
+        name: Speed (${{ matrix.dev }})
+        path: |
+          onnx_inference_speed.csv
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
@@ -469,8 +578,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: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.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

.github/workflows/test.yml

@@ -251,14 +251,9 @@ jobs:
       with:
         key: fuzzing-unit
         deps: testing_unit
-    - name: Fuzz Test symbolic
-      run: python test/external/fuzz_symbolic.py
-    - name: Fuzz Test symbolic (symbolic divisors)
-      run: python test/external/fuzz_symbolic_symbolic_div.py
-    - name: Fuzz Test fast idiv
-      run: python test/external/fuzz_fast_idiv.py
-    - name: Fuzz Test shape ops
-      run: python test/external/fuzz_shape_ops.py
+    - name: Fuzz Tests
+      run: |
+        parallel --tagstring '[{}]' 'python test/external/fuzz_{}.py' ::: symbolic symbolic_div fast_idiv shape_ops
 
   testopenclimage:
     name: CL IMAGE Tests
@@ -296,7 +291,8 @@ jobs:
           llvm: 'true'
       - name: Test openpilot model kernel count and gate usage
         run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1361 ALLOWED_GATED_READ_IMAGE=55 FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1361 ALLOWED_GATED_READ_IMAGE=55 FLOAT16=1 DEV="CL::IMAGE_PITCH_ALIGNMENT=64" IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          # IMAGE_PITCH_ALIGNMENT=64 matches adreno 630
       - name: Test openpilot CL compile fp32 (test correctness)
         run: |
           DEV=CL IMAGE=1 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
@@ -443,9 +439,8 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: linux-${{ matrix.dev }}
-        deps: testing_unit
+        deps: "testing_unit${{ contains(matrix.dev, 'LVP') && ' mesa' || '' }}"
         llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') || contains(matrix.dev, 'CLANG') }}
-        mesa: ${{ contains(matrix.dev, 'LVP') && 'cpu' || 'false' }}
         webgpu: ${{ matrix.dev == 'WEBGPU' }}
         opencl: ${{ matrix.dev == 'CL' }}
     - name: Set env
@@ -497,7 +492,7 @@ jobs:
           AMD: 0
         run: |
           PYTHONPATH=. DEV=NULL:HIP:gfx1100 python extra/mmapeak/mmapeak.py
-          PYTHONPATH=. DEV=NULL:HIP:gfx950 python3 -m pytest -n=auto test/testextra/test_tk.py test/backend/test_asm_gemm.py
+          PYTHONPATH=. DEV=NULL:HIP:gfx950 python3 -m pytest -n=auto test/testextra/test_tk.py
       - name: Run matmul on MOCKKFD
         run: |
           PYTHONPATH="." DEV=MOCKKFD+AMD N=256 python3 extra/gemm/amd_asm_matmul.py
@@ -700,9 +695,8 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: macos-${{ matrix.dev }}
-        deps: testing_unit
+        deps: "testing_unit${{ contains(matrix.dev, 'LVP') && ' mesa' || '' }}"
         llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') }}
-        mesa: ${{ contains(matrix.dev, 'LVP') && 'cpu' || 'false' }}
         webgpu: ${{ matrix.dev == 'WEBGPU' }}
     - name: Set env
       run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
@@ -768,8 +762,7 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: compile-${{ matrix.backend }}
-          deps: testing_unit
-          mesa: ${{ (matrix.backend == 'ir3' || matrix.backend == 'nak') && 'true' }}
+          deps: "testing_unit mesa"
       - name: Set env
         shell: bash
         run: printf "NULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'DEV=NULL:IR3:a630' || matrix.backend == 'nak' && 'DEV=NULL:NAK:sm_120' }}" >> $GITHUB_ENV

examples/anthropic_challenge.py

@@ -100,7 +100,7 @@ class VLIWRenderer(Renderer):
       assert u.dtype.count in (1,8), "dtype count must be 1 or 8"
 
       # dumb register allocator
-      if u.op not in {Ops.STORE, Ops.SINK, Ops.GEP}:
+      if u.op not in {Ops.STORE, Ops.SINK, Ops.INDEX}:
         r[u] = reg
         reg += u.dtype.count
 
@@ -110,9 +110,9 @@ class VLIWRenderer(Renderer):
           inst.append({"flow": [("halt",)]})
         case Ops.CONST:
           inst.append({"load": [("const", r[u], u.arg)]})
-        case Ops.GEP:
-          # a GEP is just an alias to a special register in the vector
-          r[u] = r[u.src[0]] + u.arg[0]
+        case Ops.INDEX:
+          # an INDEX is just an alias to a special register in the vector
+          r[u] = r[u.src[0]] + u.src[1].arg
         case Ops.STACK:
           if all(s == u.src[0] for s in u.src):
             # if all sources are the same, we can broadcast

examples/mlperf/models/gpt_oss.py

@@ -0,0 +1,285 @@
+import math, os, functools
+if __name__ == "__main__":
+  os.environ["DEFAULT_FLOAT"] = "bfloat16"
+  os.environ["OPTIM_DTYPE"] = "bfloat16"
+  if "DEV" not in os.environ: os.environ["DEV"] = "NULL::gfx950"
+  # CDNA
+  os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
+  os.environ["ALL2ALL"] = "1"
+  os.environ["USE_ATOMICS"] = "1"
+from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
+from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
+from tinygrad.uop.ops import Ops, UOp
+from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
+from extra.llama_kernels.rmsnorm import rmsnorm
+from extra.gemm.cdna_asm_gemm import _mx_block_scale
+
+FP8_DTYPE = dtypes.fp8e4m3
+FP8_MAX = 448.0
+INIT_STD = 0.008
+
+def quantize_mx(x:Tensor) -> tuple[Tensor, Tensor]:
+  *batch, K = x.shape
+  scale_K = K // 32
+  amax = x.detach().float().reshape(*batch, scale_K, 32).abs().max(axis=-1)
+  e8 = (amax.maximum(1e-38).log2().floor() + 127).clamp(0, 254).cast(dtypes.uint8)
+  qscale = (127.0 - e8.cast(dtypes.float32)).exp2().reshape(*batch, scale_K, 1).expand(*batch, scale_K, 32).reshape(*batch, K)
+  x_scaled = x.float() * qscale
+  x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach())  # STE
+  return x_clamped.cast(FP8_DTYPE), e8
+
+def _quant_dequant_fwd(x:Tensor) -> Tensor:
+  # x (2d bf16) -> bf16 value after an mxfp8 round-trip (1x32 block scaling on the last axis)
+  M, K = x.shape
+  scale_K = K // 32
+  amax = x.float().reshape(M, scale_K, 32).abs().max(axis=-1)
+  e8 = (amax.maximum(1e-38).log2().floor() + 127).clamp(0, 254).cast(dtypes.uint8)
+  qscale = (127.0 - e8.cast(dtypes.float32)).exp2().reshape(M, scale_K, 1).expand(M, scale_K, 32).reshape(M, K)
+  x_fp8 = (x.float() * qscale).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE).cast(dtypes.float32)
+  return (x_fp8 * _mx_block_scale(e8)).cast(dtypes.bfloat16)
+
+@functools.cache
+def _quant_dequant_fwd_fxn(x_p, device):
+  return _quant_dequant_fwd(Tensor(x_p, device=device))
+
+def _quant_dequant_bwd(grad:UOp, call:UOp) -> tuple:
+  return (Tensor(grad).cast(dtypes.bfloat16).uop,)
+
+def quant_dequant_mx(x:Tensor) -> Tensor:
+  fxn = _quant_dequant_fwd_fxn(x.as_param(0).uop, x.device)
+  return Tensor(UOp.maketuple(fxn.uop).call(x.uop, grad_fxn=_quant_dequant_bwd).gettuple(0))
+
+def _dequant_fwd(w_q:Tensor, w_scale:Tensor) -> Tensor:
+  return w_q.cast(dtypes.bfloat16) * _mx_block_scale(w_scale)
+
+@functools.cache
+def _dequant_fwd_fxn(wq_p, ws_p, device):
+  return _dequant_fwd(Tensor(wq_p, device=device), Tensor(ws_p, device=device))
+
+def _dequant_bwd(grad:UOp, call:UOp) -> tuple:
+  w_scale = Tensor(call.src[2])
+  return ((Tensor(grad).cast(dtypes.bfloat16) * _mx_block_scale(w_scale).cast(dtypes.bfloat16)).uop, None)
+
+def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
+  fxn = _dequant_fwd_fxn(w_q.as_param(0).uop, w_scale.as_param(1).uop, w_q.device)
+  call = UOp.maketuple(fxn.uop).call(w_q.uop, w_scale.uop, grad_fxn=_dequant_bwd)
+  return Tensor(call.gettuple(0))
+
+def matmul_mx(x:Tensor, w_q:Tensor, w_scale:Tensor) -> Tensor:
+  l_shape = x.shape[:-1]
+  x_phys = quant_dequant_mx(x.reshape(-1, x.shape[-1])).reshape(*l_shape, x.shape[-1])
+  w_phys = dequant_weight(w_q, w_scale)
+  return (x_phys @ w_phys.T).cast(dtypes.bfloat16)
+
+def swiglu(x:Tensor, limit:float=7.0, alpha:float=1.702) -> Tensor:
+  x_glu, x_linear = x[..., ::2], x[..., 1::2]
+  x_glu = x_glu.clamp(max_=limit)
+  x_linear = x_linear.clamp(-limit, limit)
+  return (x_glu * (alpha * x_glu).sigmoid()) * (x_linear + 1)
+
+class GPTOSS:
+  def __init__(self, dim:int, n_layers:int, n_heads:int, n_kv_heads:int, head_dim:int, n_experts:int, experts_per_tok:int,
+               intermediate_size:int, vocab_size:int, norm_eps:float=1e-5, rope_theta:int=150000, sliding_window:int=128,
+               swiglu_limit:float=7.0, max_context:int=8192):
+    self.dim, self.n_layers, self.n_heads, self.n_kv_heads, self.head_dim = dim, n_layers, n_heads, n_kv_heads, head_dim
+    self.n_rep = n_heads // n_kv_heads
+    self.n_experts, self.experts_per_tok, self.intermediate_size = n_experts, experts_per_tok, intermediate_size
+    self.vocab_size, self.norm_eps, self.sliding_window, self.swiglu_limit = vocab_size, norm_eps, sliding_window, swiglu_limit
+    self.sm_scale = 1.0 / math.sqrt(head_dim)
+
+    scaled_std = INIT_STD / math.sqrt(2 * n_layers)
+    q_dim, qkv_dim = n_heads * head_dim, head_dim * (n_heads + 2 * n_kv_heads)
+
+    # attn
+    self.wqkv, self.wqkv_scale = self._quant_weight(n_layers, qkv_dim, dim)
+    self.wqkv_bias = Tensor.zeros(n_layers, qkv_dim, dtype=dtypes.bfloat16).contiguous()
+    self.wo, self.wo_scale = self._quant_weight(n_layers, dim, q_dim, std=scaled_std)
+    self.wo_bias = Tensor.zeros(n_layers, dim, dtype=dtypes.bfloat16).contiguous()
+    self.sinks = Tensor.zeros(n_layers, n_heads, dtype=dtypes.bfloat16).contiguous()
+    self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
+
+    # moe ffn
+    self.ffn_norm = Tensor.ones(n_layers, dim).contiguous()
+    self.gate = Tensor.normal(n_layers, n_experts, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
+    self.gate_bias = Tensor.zeros(n_layers, n_experts, dtype=dtypes.bfloat16).contiguous()
+    self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim)
+    self.w_gate_up_bias = Tensor.zeros(n_layers, n_experts, intermediate_size * 2, dtype=dtypes.bfloat16).contiguous()
+    self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std)
+    self.w_down_bias = Tensor.zeros(n_layers, n_experts, dim, dtype=dtypes.bfloat16).contiguous()
+
+    # output
+    self.norm = nn.RMSNorm(dim, norm_eps)
+    self.tok_embeddings = nn.Embedding(vocab_size, dim)
+    self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
+    self.output = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
+    self.freqs_cis = precompute_freqs_cis(head_dim, max_context * 2, rope_theta).contiguous().is_param_(False)
+
+  def _quant_weight(self, *shape:int, std:float=INIT_STD):
+    w = Tensor.zeros(*shape) if getenv("ZEROS") else Tensor.normal(*shape, mean=0.0, std=std)
+    w_q, w_e8 = quantize_mx(w)
+    return w_q, w_e8.is_param_(False)
+
+  def _attn_mask(self, seqlen:int, sliding:bool, dtype) -> Tensor:
+    i, j = Tensor.arange(seqlen).reshape(seqlen, 1), Tensor.arange(seqlen).reshape(1, seqlen)
+    allowed = j <= i
+    if sliding: allowed = allowed & (i - j < self.sliding_window)
+    return allowed.where(0.0, -1e30).cast(dtype).contiguous()
+
+  def attention(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wqkv_scale:Tensor,
+                wqkv_bias:Tensor, wo:Tensor, wo_scale:Tensor, wo_bias:Tensor, sinks:Tensor):
+    bsz, seqlen, _ = x.shape
+    x_normed, rrms = rmsnorm(x, self.norm_eps)
+    qkv = matmul_mx(x_normed * attention_norm, wqkv, wqkv_scale) + wqkv_bias
+    qkv = qkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
+    xq = qkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
+    xk, xv = qkv[:, :, :, self.n_rep], qkv[:, :, :, self.n_rep + 1]
+    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
+
+    xq = xq.cast(dtypes.bfloat16).reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
+    xk = xk.cast(dtypes.bfloat16).permute(0, 2, 1, 3).unsqueeze(2)
+    xv = xv.cast(dtypes.bfloat16).permute(0, 2, 1, 3).unsqueeze(2)
+    scores = ((xq @ xk.transpose(-2, -1)).float() * self.sm_scale + mask).contiguous()
+    sink = sinks.reshape(1, self.n_kv_heads, self.n_rep, 1, 1).float()
+    m = scores.max(-1, keepdim=True).maximum(sink)
+    e = (scores - m).exp()
+    w = (e / (e.sum(-1, keepdim=True) + (sink - m).exp())).cast(dtypes.bfloat16).contiguous()
+    attn = (w @ xv).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]
+
+  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,
+                   w_down:Tensor, w_down_scale:Tensor, w_down_bias:Tensor):
+    x_normed, rrms = rmsnorm(x, self.norm_eps)
+    inp = x_normed * ffn_norm
+
+    logits = inp.float() @ gate.float().T + gate_bias.float()
+    thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
+    weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
+
+    out = None
+    for e in range(self.n_experts):
+      gate_up = (matmul_mx(inp.contiguous_backward(), w_gate_up[e], w_gate_up_scale[e]) + w_gate_up_bias[e]).contiguous()
+      y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), w_down[e], w_down_scale[e]) + w_down_bias[e]).contiguous()
+      contrib = (weights[..., e:e+1].cast(y.dtype) * y).contiguous()
+      out = contrib if out is None else out + contrib
+    return out, [x_normed, rrms]
+
+  @function(precompile=True, precompile_backward=True)
+  def run_layer(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
+    attn, attn_saves = self.attention(x, freqs_cis, mask, **attn_kwargs)
+    h = x + attn
+    ffn, ffn_saves = self.feed_forward(h, **ffn_kwargs)
+    h = h + ffn
+    if save: return (h, *attn_saves, *ffn_saves)
+    return (h,)
+
+  def shard(self, device:tuple[str, ...], mp:bool=False):
+    assert not mp, "MP not supported"
+    from tinygrad.nn.state import get_parameters
+    for v in get_parameters(self): v.shard_(device, axis=None)
+    Tensor.realize(*get_parameters(self))
+
+  def __call__(self, tokens:Tensor, save:bool=True):
+    h = self.tok_embeddings(tokens)
+    bsz, seqlen = tokens.shape
+    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :seqlen, :, :, :]
+    mask_full = self._attn_mask(seqlen, False, dtypes.float32)
+    mask_sliding = self._attn_mask(seqlen, True, dtypes.float32)
+    for i in range(self.n_layers):
+      attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wqkv_scale=self.wqkv_scale[i],
+                         wqkv_bias=self.wqkv_bias[i], wo=self.wo[i], wo_scale=self.wo_scale[i], wo_bias=self.wo_bias[i],
+                         sinks=self.sinks[i])
+      ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], gate=self.gate[i], gate_bias=self.gate_bias[i],
+                        w_gate_up=self.w_gate_up[i], w_gate_up_scale=self.w_gate_up_scale[i], w_gate_up_bias=self.w_gate_up_bias[i],
+                        w_down=self.w_down[i], w_down_scale=self.w_down_scale[i], w_down_bias=self.w_down_bias[i])
+      mask = mask_sliding if i % 2 == 0 else mask_full
+      h, *_ = self.run_layer(h, freqs_cis, mask, attn_kwargs, ffn_kwargs, save=save)
+
+    logits = self.norm(h) @ self.output.T
+    return logits
+
+def _get_pads(uop:UOp) -> list[UOp]:
+  if uop.op == Ops.ADD: return _get_pads(uop.src[0]) + _get_pads(uop.src[1])
+  return [uop]
+
+def apply_grad(grad_buf:Tensor, new_grad:UOp):
+  pads = _get_pads(new_grad)
+  if len(pads) <= 1:
+    new_grad = new_grad.cast(grad_buf.dtype)
+    grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
+    return
+  cur = grad_buf.uop
+  for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
+    if pad.op == Ops.PAD:
+      grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)])
+      buf_slice = cur.shrink(grad_shrink)
+      cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
+    else:
+      cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
+  grad_buf.uop = cur
+
+GPT_OSS_20B = dict(dim=2880, n_layers=24, n_heads=64, n_kv_heads=8, head_dim=64, n_experts=32, experts_per_tok=4,
+                   intermediate_size=2880, vocab_size=128256, norm_eps=1e-5, rope_theta=150000, sliding_window=128,
+                   swiglu_limit=7.0)
+
+if __name__ == "__main__":
+  config = {}
+  BS      = config["BS"]      = getenv("BS", 16)
+  SEQLEN  = config["SEQLEN"]  = getenv("SEQLEN", 8192)
+
+  model_params = GPT_OSS_20B
+  real_vocab_size = model_params["vocab_size"]
+  if (layers := getenv("LAYERS")) != 0: model_params["n_layers"] = layers
+
+  model = GPTOSS(**model_params, max_context=SEQLEN)
+
+  state = nn.state.get_state_dict(model)
+  print("tensor count:", len(state))
+
+  from tinygrad import Device
+  is_dp = (DP := getenv("DP", 1)) > 1
+  device_count = DP
+  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
+
+  if is_dp: model.shard(device)
+
+  # preallocate all the grad buffers and zero them out
+  grad_dtype = lambda x: dtypes.bfloat16 if x.dtype in dtypes.fp8s else x.dtype
+  grads = {x:x.zeros_like(dtype=grad_dtype(x)).contiguous() for x in state.values() if x.is_param}
+
+  # print model size
+  sz = 0
+  for k,v in state.items():
+    print(f"{colored(k, 'green' if v in grads else 'white'):30s} {str(v.shape):30s} {str(v.dtype):20s} {v.device}  {v.nbytes()/1e9:.2f} GB")
+    sz += v.nbytes()
+  print(f"total sz: {sz/1e9:.2f} GB")
+
+  with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=real_vocab_size, dtype=dtypes.int)
+  with Timing("realize weights/grads/data: "): Tensor.realize(*state.values(), *grads.values(), tokens)
+  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
+  if is_dp: tokens = tokens.shard(device, axis=0)
+
+  @TinyJit
+  def fwd_bwd(tokens:Tensor):
+    with Timing("python forward: "):
+      logits = model(tokens[:, :-1], save=True)
+      loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    with Timing("python backward: "):
+      for t,g in zip(grads, loss.gradient(*grads)):
+        apply_grad(grads[t], g.uop)
+    with Timing("run fwd_bwd: "): loss.realize(*grads.values())
+
+  @TinyJit
+  def optim_step():
+    for g in grads.values(): g.assign(g.zeros_like())
+    Tensor.realize(*grads.values())
+
+  for i in range(6):
+    GlobalCounters.reset()
+    profile_marker(f"step {i}")
+    with Timing(colored(f"*** step {i}: ", "red")):
+      fwd_bwd(tokens)
+      optim_step()
+  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -14,7 +14,6 @@ export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
-export USE_HK_BF16_GEMM=${USE_HK_BF16_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam_mp.sh

@@ -14,7 +14,6 @@ export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
-export USE_HK_BF16_GEMM=${USE_HK_BF16_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh

@@ -14,7 +14,6 @@ export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
-export USE_HK_BF16_GEMM=${USE_HK_BF16_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run_mp.sh

@@ -14,7 +14,6 @@ export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
-export USE_HK_BF16_GEMM=${USE_HK_BF16_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}

extra/gemm/metal_uop_matmul.py

@@ -33,7 +33,7 @@ def hand_spec_tc_cores():
   acc_load = UOp.vectorize(acc.after(gk)[0], acc.after(gk)[1])
   out = UOp(Ops.WMMA, dtypes.float.vec(2), (a_tc, b_tc, acc_load), arg=wmma_arg)
 
-  end_loop = UOp.group(*[acc[i].store(out.gep(i)) for i in range(2)]).end(gk)
+  end_loop = UOp.group(*[acc[i].store(out.index(i)) for i in range(2)]).end(gk)
 
   sink = UOp.group(*[mat_idx(c.after(end_loop), gx, gy, warp, i).store(acc[i]) for i in range(2)])
   return sink.sink(arg=KernelInfo(name="custom_metal_matmul", opts_to_apply=())).simplify()

extra/gemm/mi350x_uop_matmul.py

@@ -180,7 +180,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
   # store the acc into gmem
   cp_i, cp_j = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, 10004), UOp.range(BLOCK_N//TC_N, 10005)
   c_load = lambda i: C[gx, cp_i*TC_M*WARPGROUP_SIZE + warpgroup*TC_M + (warp//16)*4+i, gy, cp_j*TC_N + warp%16]
-  store = UOp.group(*[c_load(i).store(acc[cp_j, cp_i].gep(i)) for i in range(4)])
+  store = UOp.group(*[c_load(i).store(acc[cp_j, cp_i].index(i)) for i in range(4)])
   store = store.end(cp_i, cp_j)
 
   return store.sink(arg=KernelInfo(name="custom_gemm", opts_to_apply=())).simplify()
@@ -197,7 +197,7 @@ wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float,
 out = UOp(Ops.WMMA, dtypes.float.vec(4), (A_in, B_in, acc_load), arg=wmma_arg)
 
 # store back the acc
-acc = acc.after(UOp.group(*[acc[i].store(out.gep(i)) for i in range(4)]).end(K_loop))
+acc = acc.after(UOp.group(*[acc[i].store(out.index(i)) for i in range(4)]).end(K_loop))
 
 # store the acc into gmem
 store = UOp.group(*[C[gx, (warp//16)*4+i, gy, warp%16].store(acc[i]) for i in range(4)])

extra/hcq2/hcq2.py

@@ -2,12 +2,12 @@ from __future__ import annotations
 from typing import cast, Callable, TypeVar, Generic, Any
 import struct, functools, time, collections, importlib, itertools, weakref
 from dataclasses import replace, dataclass, field
-from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, DEBUG, dedup, flatten, pluralize
-from tinygrad.helpers import to_tuple, round_up
+from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, DEBUG, dedup, pluralize
+from tinygrad.helpers import to_tuple, round_up, partition
 from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer
 from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites, GroupOp
 from tinygrad.uop.symbolic import symbolic_simple, symbolic
-from tinygrad.dtype import dtypes, AddrSpace
+from tinygrad.dtype import dtypes, AddrSpace, truncate
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.renderer import Renderer, Estimates
 from tinygrad.engine.realize import to_program, get_call_arg_uops, get_call_name, get_call_outs_ins, estimate_uop, pm_flatten_linear
@@ -23,11 +23,12 @@ class HCQ2Compiled(Compiled):
 
     # default pm bufferize
     self.pm_bufferize = PatternMatcher([
-      (UPat(Ops.BUFFER, tag="timeline_signal"), lambda ctx: ctx.timeline_signal()),
-      (UPat(Ops.BUFFER, tag="timeline_value"), lambda ctx: ctx.timeline_value()),
-      (UPat(Ops.BUFFER, tag="sentinel_signal"), lambda ctx: ctx.timeline_signal("sentinel", (1 << 64) - 1)),
-      (UPat(Ops.BUFFER, name="b"), lambda ctx, b:
-        Buffer(ctx.device, b.max_numel(), b.dtype, options=BufferSpec(host=False, uncached=True, cpu_access=True, nolru=True))), # TODO: remove nolru
+      (UPat(Ops.PARAM, tag="timeline_signal"), lambda ctx: ctx.timeline_signal()),
+      (UPat(Ops.PARAM, tag="timeline_value"), lambda ctx: ctx.timeline_value()),
+      (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx.timeline_signal("sentinel", (1 << 64) - 1)),
+      (UPat(Ops.PARAM, name="b"), lambda ctx, b:
+        Buffer(ctx.device, b.max_numel(), b.dtype.base, options=BufferSpec(host=False, uncached=True, cpu_access=True, nolru=True))
+        if b.tag is not None else None), # TODO: remove nolru
     ])
 
     super().__init__(device, allocator, compilers, lambda *a, **kw: None, None, arch=arch)
@@ -44,10 +45,6 @@ class HCQ2Compiled(Compiled):
     buf.as_memoryview(force_zero_copy=True).cast('Q')[0] = init_value
     return buf
 
-  @functools.cached_property
-  def timestamps_buf(self) -> Buffer:
-    return Buffer(self.device, 0x1000, dtypes.uint8, options=BufferSpec(cpu_access=True), preallocate=True)
-
   def synchronize(self, timeout:int|None=None):
     if not hasattr(self, 'iface'): return
     sig = self.timeline_signal()._buf.cpu_view().mv.cast('Q')
@@ -132,35 +129,6 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
 
   # def _as_buffer(self, buf): return buf.cpu_view().mv
 
-def unwrap_after(uop):
-  while uop.op is Ops.AFTER: uop = uop.src[0]
-  return uop
-
-def make_getaddr(u, device=None):
-  if unwrap_after(u).op not in (Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM): return u
-  return UOp(Ops.GETADDR, dtypes.uint64, src=(u, UOp(Ops.DEVICE, arg=device or to_tuple(u.device)[0])))
-
-def make_ins(op, *srcs):
-  return UOp(Ops.INS, dtypes.void, tuple(UOp.const(dtypes.uint32, s) if isinstance(s, int) else s.cast(dtypes.uint32) for s in srcs), op)
-
-def make_patch(buf:UOp, off:sint, val:UOp, dtype=None) -> UOp:
-  dt = dtype or val.dtype
-  return UOp(Ops.SHRINK, buf.dtype.base, (buf, UOp.const(dtypes.int, off), UOp.const(dtypes.int, dt.itemsize))).bitcast(dt).store(val.cast(dt))
-
-def make_cmdbuf(lin, devs, tag):
-  blob, patches = b'', []
-  for s in (s for ins in lin.src for s in ins.src):
-    if s.op is not Ops.CONST: patches.append((len(blob), s))
-    blob += struct.pack(f'<{s.dtype.fmt}', s.arg if s.op is Ops.CONST else 0x0)
-  buf = UOp.new_buffer(devs, len(blob), dtypes.uint8).rtag(tag)
-  return buf.after(buf.store(UOp(Ops.BINARY, dtypes.void, src=(), arg=blob)), *[make_patch(buf, off, s) for off, s in patches])
-
-def make_mstack(uops): return uops[0] if len(uops) == 1 else UOp(Ops.MSTACK, uops[0].dtype, tuple(uops))
-
-def make_signal(devs, queue=None, sentinel=False):
-  return UOp.new_buffer(devs, 1, dtypes.uint64).rtag("sentinel_signal" if sentinel else (queue, "timeline_signal") if queue else "timeline_signal")
-def make_signal_value(devs, queue=None): return UOp.new_buffer(devs, 1, dtypes.uint64).rtag((queue, "timeline_value") if queue else "timeline_value")
-
 # *****************
 # 0. helpers
 
@@ -169,18 +137,65 @@ HCQ_P2P_DEVS = HCQ_DEVS | frozenset(("CPU",))
 
 def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
 
+def unwrap_after(uop):
+  while uop.op is Ops.AFTER: uop = uop.src[0]
+  return uop
+
+def make_getaddr(u, device=None):
+  if unwrap_after(u).op not in (Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM): return u
+  return UOp(Ops.GETADDR, dtypes.uint64, src=(u,), arg=device or to_tuple(u.device)[0])
+
+def make_ins(op, *srcs):
+  return UOp(Ops.INS, dtypes.void, tuple(UOp.const(dtypes.uint32, s) if isinstance(s, int) else s.cast(dtypes.uint32) for s in srcs), op)
+
+def make_placeholder(devs, size:int, dtype, name=None, unique=True) -> UOp:
+  return UOp.param(next(UOp.unique_num) if unique else 0, dtype.ptr(size), device=devs).rtag(name or "buf")
+
+def make_patch(buf:UOp, off:sint, val:UOp, dtype=None) -> UOp:
+  dt = dtype or val.dtype
+  return UOp(Ops.SHRINK, buf.dtype.base, (buf, UOp.const(dtypes.int, off), UOp.const(dtypes.int, dt.itemsize))).bitcast(dt).store(val.cast(dt))
+
+def make_cmdbuf(lin, devs):
+  blob, patches = b'', []
+  for s in (s for ins in lin.src for s in ins.src):
+    if s.op is not Ops.CONST: patches.append((len(blob), s))
+    blob += struct.pack(f'<{s.dtype.fmt}', s.arg if s.op is Ops.CONST else 0x0)
+  buf = make_placeholder(devs, len(blob), dtypes.uint8)
+
+  # pull patches to cmdbuf
+  afters = dedup(u for _, s in patches for u in s.toposort() if u.op is Ops.AFTER)
+  deps = tuple(d for p in afters for d in p.src[1:])
+  cmdbuf = buf.after(buf.store(UOp(Ops.BINARY, dtypes.void, src=(), arg=blob)), *[make_patch(buf, off, s) for off, s in patches], *deps)
+  return cmdbuf.substitute({p: p.src[0] for p in afters}) if afters else cmdbuf
+
+def make_mstack(uops): return uops[0] if len(uops) == 1 else UOp(Ops.MSTACK, uops[0].dtype, tuple(uops))
+
+def make_signal(devs, queue=None, sentinel=False):
+  return make_placeholder(devs, 1, dtypes.uint64, "sentinel_signal" if sentinel else (queue, "timeline_signal") if queue else "timeline_signal", unique=False)
+def make_signal_value(devs, queue=None):
+  return make_placeholder(devs, 1, dtypes.uint64, (queue, "timeline_value") if queue else "timeline_value", unique=False)
+
+def make_submit(*cmds, devs:str|tuple[str, ...], queue:str) -> UOp:
+  return UOp.custom_function("submit_cmdbuf", UOp(Ops.LINEAR, dtypes.void, src=tuple(cmds), arg=(to_tuple(devs), queue)))
+def get_submit(ast:UOp) -> UOp: return next(u for u in ast.toposort() if u.op is Ops.CUSTOM_FUNCTION and u.arg == "submit_cmdbuf")
+
 @dataclass(frozen=True)
 class HCQInfo:
-  name:str = ""
-  estimates:Estimates = Estimates()
-  outs:tuple[int, ...] = ()
-  devs:tuple[str, ...] = ()
+  name:str
+  estimates:Estimates
+  device:tuple[str, ...]
+  queue:str
 
-  params:tuple[int, ...] = ()
+  input_idxs:tuple[int, ...] = () # indexes into input_uops used by this call
   inputs:int|None = None
 
-  @staticmethod
-  def from_call(call:UOp) -> HCQInfo: return HCQInfo(get_call_name(call, get_call_arg_uops(call)), estimate_uop(call), get_call_outs_ins(call)[0])
+# *****************
+# 0.1. prep: replace buffers with params
+
+def replace_call_buffers(ctx:list[UOp], call:UOp) -> UOp|None:
+  ctx += [s for s in dedup(call.src[1:]) if s not in ctx and s.op not in (Ops.PARAM, Ops.BIND)]
+  return call.replace(src=call.src[:1] + tuple(s if s.op in (Ops.PARAM, Ops.BIND) else s.param_like(ctx.index(s)) for s in call.src[1:]))
+pm_replace_buffers = PatternMatcher([(UPat(Ops.CALL, name="call"), replace_call_buffers)])
 
 # *****************
 # 1.1. prep: staging copies
@@ -190,70 +205,24 @@ def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS) and not all_d
 def stage_copy(dst:UOp, src:UOp) -> UOp|None:
   if not (_need_staging(src, dst) or _need_staging(dst, src)): return None
 
-  stage = UOp.new_buffer("CPU", src.nbytes(), dtypes.uint8)
+  stage = UOp.new_buffer("CPU", src.max_numel() * src.dtype.base.itemsize, dtypes.uint8)
   return UOp(Ops.LINEAR, dtypes.void, (src.copy_to_device("CPU").call(stage, src), stage.copy_to_device(dst.device).call(dst, stage)))
 pm_insert_copy_staging = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy)])
 
 # *****************
-# 2.1. hcq lowering: programs/kernargs
+# 2.1. tag hcq calls
 
-@functools.cache
-def get_pm_prep_program(name:str) -> PatternMatcher|None:
-  try:
-    importlib.import_module(f'tinygrad.runtime.ops_{name.lower()}') # TODO: remove that
-    return importlib.import_module(f'extra.hcq2.ops_{name.lower()}2').pm_prep_program
-  except ImportError: return None
+def tag_hcq_call(ctx:itertools.count, call:UOp) -> UOp:
+  if (hcq_devs:=next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)) is None: return call
 
-def prep_program(call:UOp, prg:UOp) -> UOp|None:
-  dev = call.src[1].device
-  if (pm:=get_pm_prep_program(to_tuple(dev)[0].split(":")[0])) is None or (lowered:=pm.rewrite(prg)) is None: return None
-
-  data, image_bytes = lowered
-  buf = UOp.new_buffer(dev, len(image_bytes), dtypes.uint8).rtag("program")
-  blob = UOp(Ops.BINARY, dtypes.void, src=(), arg=image_bytes)
-  return prg.replace(src=(buf.after(buf.store(blob)),), arg=(data, prg.arg)).call(*call.src[1:], aux=HCQInfo.from_call(call))
-
-def prep_kernargs(call:UOp, prg:UOp) -> UOp:
-  (data, info), dev_uop = prg.arg, UOp(Ops.DEVICE, arg=call.src[1].device)
-  buf = UOp.new_buffer(dev_uop.arg, data.kernargs_alloc_size, dtypes.uint8).rtag("kernargs")
-  patches = [make_patch(buf, i*8, make_getaddr(call.src[1+gi], dev_uop.arg)) for i,gi in enumerate(info.globals)] \
-          + [make_patch(buf, len(info.globals)*8 + i*4, v, dtypes.uint32) for i,v in enumerate(info.vars)]
-  return call.replace(src=(prg.replace(src=prg.src + (buf.after(*patches),), arg=(data, info)),) + call.src[1:])
-
-pm_prep_runtime = PatternMatcher([
-  # bind generic PROGRAM device to the call's actual dev(s), then run device-specific lowering
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"),),
-    name="call", allow_any_len=True), prep_program),
-
-  # lower kernargs (PROGRAM.src[0] is now AFTER(BUFFER, COPY) — the lowered program image)
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.BUFFER).or_after(),), name="prg"),), name="call", allow_any_len=True), prep_kernargs),
-])
+  queue = "COMPUTE:0" if call.src[0].op is Ops.PROGRAM else "COPY:0"
+  info = HCQInfo(get_call_name(call, get_call_arg_uops(call)), estimate_uop(call), to_tuple(hcq_devs), queue)
+  return call.replace(arg=replace(call.arg, aux=info)).rtag(next(ctx))
+pm_tag_hcq_calls = PatternMatcher([(UPat(Ops.LINEAR, name="linear"),
+  lambda ctx, linear: linear.replace(src=tuple(tag_hcq_call(ctx, s) for s in linear.src)))])
 
 # *****************
-# 2.2. hcq lowering: ops to ir
-
-def make_submit(*cmds, devs:str|tuple[str, ...], queue:str) -> UOp:
-  devs:tuple[str, ...] = to_tuple(devs)
-  return UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(UOp(Ops.LINEAR, dtypes.void, src=tuple(cmds), arg=(devs, queue)),), arg="submit")
-
-def lower_program(call:UOp, prg:UOp) -> UOp:
-  return make_submit(prg, devs=call.src[1].device, queue="COMPUTE:0").sink().call(*call.src[1:], aux=call.arg.aux).rtag("hcq")
-
-def lower_copy(call:UOp, copy:UOp) -> UOp|None:
-  dst, src = call.src[1], call.src[2]
-  if (hcq_dev:=next((b.device for b in (dst, src) if b.device.split(":")[0] in HCQ_DEVS), None)) is None: return None
-
-  cp_op = UOp(Ops.COPY, dtypes.void, src=(dst, src), arg=src.buffer.nbytes)
-  return make_submit(cp_op, devs=hcq_dev, queue="COPY:0").sink().call(*call.src[1:], aux=HCQInfo.from_call(call)).rtag("hcq")
-
-pm_lower_ops = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.BUFFER).or_after(), UPat(Ops.BUFFER).or_after()), name="prg"),),
-    name="call", allow_any_len=True), lower_program),
-  (UPat(Ops.CALL, src=(UPat(Ops.COPY, name="copy"),), name="call", allow_any_len=True), lower_copy),
-])
-
-# *****************
-# 3.1. deps tracking
+# 2.2. deps tracking
 # device.timeline_signal/value are the per-device schedule epoch. Before a schedule queue accesses memory owned by device N for the first time,
 # it waits for device[N].timeline_signal >= device[N].timeline_value - 1. This orders the schedule after all prior schedules that touched device N.
 #
@@ -266,142 +235,95 @@ pm_lower_ops = PatternMatcher([
 #
 # C programs reserve and bump timeline values, then patch command buffers with the concrete wait/signal values.
 
-@dataclass
-class DepsCtx:
-  deps:DepsTracker = field(default_factory=DepsTracker)
-  opid:itertools.count = field(default_factory=lambda: itertools.count(0))
-  last_per_queue:weakref.WeakValueDictionary[tuple[Any, str], UOp] = field(default_factory=weakref.WeakValueDictionary)
-  params:dict[tuple[int, int], Buffer] = field(default_factory=dict)
+class HCQDepsTracker(DepsTracker):
+  @staticmethod
+  def _key(buf:Any) -> tuple[Any, int, int]:
+    return (buf.arg.slot, 0, buf.max_numel() * buf.dtype.base.itemsize) if isinstance(buf, UOp) else DepsTracker._key(buf)
 
-def get_dep_buf(ctx:DepsCtx, u:UOp, lane:int) -> Buffer:
-  # TODO: should this be a part of DepsTracker?
-  if u.op is Ops.PARAM: return ctx.params.setdefault((u.arg.slot, lane), Buffer("NULL", u.max_numel(), u.dtype.base))
-  if u.op is Ops.MSTACK: return get_dep_buf(ctx, u.src[lane], 0)
-  if u.op in (Ops.SLICE, Ops.MSELECT): return get_dep_buf(ctx, u.src[0], u.arg if u.op is Ops.MSELECT else lane)
-  return b.bufs[lane] if isinstance(b:=u.buffer, MultiBuffer) else b
+def make_deps(u:UOp, dep_lanes:list[tuple[UOp, int, int]], nlanes:int) -> UOp:
+  deps:dict[UOp, list[int|None]] = collections.defaultdict(lambda: [None]*nlanes)
+  for dep, dlane, lane in dep_lanes: deps[dep][lane] = dlane
+  return u.after(*deps, arg=tuple(tuple(v) for v in deps.values()))
 
-def schedule_inner_sync(ctx:DepsCtx, linear:UOp) -> UOp:
-  new_src = []
-  for call in linear.src:
-    if call.tag != "hcq":
-      new_src.append(call)
-      continue
+def sched_sync(ctx:DepsTracker, call:UOp) -> UOp|None:
+  if not isinstance(call.arg.aux, HCQInfo): return None
 
-    new_q = ctx.last_per_queue[q.arg] = (q:=get_submit(call.src[0]).src[0]).rtag(next(ctx.opid))
-    qdevs, refs = to_tuple(new_q.arg[0]), get_call_arg_uops(call)
+  refs = get_call_arg_uops(call)
+  outs, _ = get_call_outs_ins(call)
+  devices, queue = call.arg.aux.device, call.arg.aux.queue
 
-    # per-lane deps, tracked per (device, queue). skip self
-    dep_lanes:list[tuple[UOp, int]] = []
-    for lane, d in enumerate(qdevs):
-      for dep in ctx.deps.access_resources([get_dep_buf(ctx, b, lane) for b in refs], call.arg.aux.outs, new_q.replace(arg=(d, new_q.arg[1]))):
-        if dep.tag != new_q.tag: dep_lanes.append((dep, lane))
+  dep_lanes:list[tuple[UOp, int, int]] = []
+  for lane, d in enumerate(devices):
+    lane_refs = [b if b.op is Ops.PARAM else mb.bufs[lane] if isinstance(mb:=b.buffer, MultiBuffer) else mb for b in refs]
+    for dep, dlane in ctx.access_resources(lane_refs, outs, (call, lane)): dep_lanes.append((dep, dlane, lane))
 
-    # drop self-queue waits, queue self-orders
-    if qdevs[0].split(":")[0] in {"AMD", "QCOM"} or new_q.arg[1].startswith("COPY"):
-      dep_lanes = [(dep, lane) for dep, lane in dep_lanes if dep.arg != (qdevs[lane], new_q.arg[1])]
+  if devices[0].split(":")[0] in {"AMD", "QCOM"} or queue.startswith("COPY"):
+    dep_lanes = [(dep, dlane, lane) for dep, dlane, lane in dep_lanes if (dep.arg.aux.device[dlane], dep.arg.aux.queue) != (devices[lane], queue)]
 
-    # keep latest dep per lane, group lanes
-    latest = {(dep.arg, lane): dep for dep, lane in sorted(dep_lanes, key=lambda x: x[0].tag)}
-    deps:dict[UOp, tuple[int, ...]] = collections.defaultdict(tuple)
-    for (_, lane), dep in latest.items(): deps[dep] += (lane,)
-
-    if deps: new_q = new_q.after(*deps, arg=tuple(deps.values())).rtag("deps")
-    new_src.append(call.replace(src=(call.src[0].substitute({q:new_q}),)))
-  return linear.replace(src=tuple(new_src))
-pm_schedule_inner_sync = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), schedule_inner_sync)])
+  # keep latest dep per (dep device, queue, cur lane)
+  latest = {((dep.arg.aux.device[dlane], dep.arg.aux.queue), lane): (dep, dlane) for dep, dlane, lane in sorted(dep_lanes, key=lambda x: x[0].tag)}
+  return make_deps(call, [(dep, dlane, lane) for (_, lane), (dep, dlane) in latest.items()], len(devices))
+pm_sched_sync = PatternMatcher([(UPat(Ops.CALL, name="call"), sched_sync)])
 
 # *****************
-# 3.2. finalizer
+# 2.3. merge into queues
 
-def make_finalizer(queues:list[UOp], nbump:int) -> UOp:
-  devs = tuple(dedup([d for q in queues for d in to_tuple(q.arg[0])]))
-  zero = UOp.const(dtypes.int, 0)
-  tl = make_signal_value(devs)
-
-  # queue is inc with deps
-  submit = make_submit(make_signal(devs).store(tl.index(zero)), devs=devs, queue="COMPUTE:0")
-
-  # split each (multi-device) queue into per-device deps so each finalizer lane waits on the matching device's signal
-  lane_queues = [(q.replace(arg=(d, q.arg[1])), (devs.index(d),)) for q in queues for d in to_tuple(q.arg[0])]
-  submit = submit.replace(src=(submit.src[0].after(*(q for q, _ in lane_queues), arg=tuple(l for _, l in lane_queues)).rtag("deps"),))
-
-  upd = [(tl, 1)] + [(make_signal_value(devs, queue=qn), nbump) for qn in dedup([q.arg[1] for q in queues])]
-  patches = [s.after(submit).index(zero, dtype=s.dtype.ptr()).store(s.index(zero) + inc) for s, inc in upd]
-  return UOp.barrier(*patches).sink().call(aux=HCQInfo("hcq finalizer")).rtag("hcq")
-
-def add_finalizer(ctx:DepsCtx, linear:UOp) -> UOp:
-  parts:dict[str, list[UOp]] = collections.defaultdict(list)
-  for d, q in ctx.last_per_queue.items(): parts[to_tuple(d[0])[0].split(':')[0]].append(q)
-
-  nbump = next(ctx.opid)
-  return linear.replace(src=linear.src + tuple([make_finalizer(queues, nbump) for queues in parts.values()]))
-pm_add_finalizer = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), add_finalizer)])
-
-# *****************
-# 3.3. lower loads/stores
-
-def add_loads(ctx:set[int], deps:UOp) -> UOp:
-  cur_devs = to_tuple((cur:=deps.src[0]).arg[0])
-
-  waits = []
-  for lanes, dep in zip(deps.arg, deps.src[1:]):
-    dep_dev, queue = dep.arg # dep_dev is a single device (deps are recorded per-device)
-    ctx.add(dep.tag) # mark op to update signal.
-
-    # for lanes that need this dep, wait on the dep device's signal/value; other lanes get a passing sentinel
-    lanes = set(lanes)
-    sig = make_mstack([make_signal(dep_dev if j in lanes else d, queue=queue, sentinel=j not in lanes) for j, d in enumerate(cur_devs)])
-    val = make_mstack([make_signal_value(dep_dev if j in lanes else d, queue=queue) for j, d in enumerate(cur_devs)]).index(UOp.const(dtypes.int, 0))
-    waits.append(sig.wait(val + dep.tag))
-  return cur.replace(src=tuple(waits) + cur.src)
-pm_add_inner_loads = PatternMatcher([(UPat(Ops.AFTER, tag="deps", name="deps"), add_loads)])
-
-def add_stores(ctx:set[int], submit:UOp, q:UOp) -> UOp|None:
-  if q.tag not in ctx: return None
-  devs, queue = q.arg
-  src = q.src + (make_signal(devs, queue=queue).store(make_signal_value(devs, queue=queue).index(UOp.const(dtypes.int, 0)) + q.tag),)
-  return submit.replace(src=(q.replace(src=src, tag=None),))
-pm_add_inner_stores = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_stores)])
-
-# *****************
-# 4.1. merge queues
-
-def get_submit(ast:UOp) -> UOp: return next(u for u in ast.toposort() if u.op is Ops.CUSTOM_FUNCTION and u.arg == "submit")
-
-def merge_sink(sinks:list[UOp]) -> UOp:
-  if len(sinks) == 1: return sinks[0]
-  submits = [get_submit(sink) for sink in sinks]
-  queues = [submit.src[0] for submit in submits]
-  anchor = submits[-1].replace(src=(queues[-1].replace(src=tuple(x for q in queues for x in q.src)),))
-  for sink, submit in zip(sinks[:-1], submits[:-1]):
-    if sink.src[0] is not submit: anchor = sink.src[0].substitute({submit: anchor}, walk=True)
-  return sinks[-1].substitute({submits[-1]: anchor}, walk=True)
+def _merged_hcq_call(calls:list[UOp]):
+  info = replace(unwrap_after(calls[0]).arg.aux, estimates=sum((unwrap_after(c).arg.aux.estimates for c in calls), start=Estimates()))
+  cmdbuf = make_submit(*calls, devs=info.device, queue=info.queue)
+  return UOp.custom_function("hcq", cmdbuf.sink()).call(name="hcq", aux=info)
 
 def merge_queues(linear:UOp) -> UOp:
   new_src:list[UOp] = []
-  opened_qs:dict[tuple[tuple[str, ...], str], tuple[list[UOp], HCQInfo]] = {} # (devs, queue) -> (sinks, aux), kept in submit order
+  opened_qs:dict[tuple[tuple[str, ...], str], list[UOp]] = {} # (devs, queue) -> list of calls, kept in submit order
 
   for call in linear.src:
-    # finalizer cannot be merged, since it bumps inner signal (this introduces race when multidevs).
-    if call.tag != "hcq" or (call.tag == "hcq" and call.arg.aux.name == "hcq finalizer"):
-      new_src += [merge_sink((sa:=opened_qs.pop(k))[0]).call(aux=sa[1]).rtag("hcq") for k in list(opened_qs)] + [call]
+    if not isinstance(unwrap_after(call).arg.aux, HCQInfo):
+      new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in list(opened_qs)] + [call]
       continue
 
-    devs, queue = get_submit(new_sink:=call.src[0]).src[0].arg
-    new_rec = ([new_sink], call.arg.aux)
-    if (old:=opened_qs.pop((devs, queue), None)) is not None:
-      new_rec = (old[0] + [new_sink], replace(new_rec[1], name=f"{queue.lower()} submit", estimates=old[1].estimates + new_rec[1].estimates))
+    devices, queue = unwrap_after(call).arg.aux.device, unwrap_after(call).arg.aux.queue
+
+    if (old:=opened_qs.pop((devices, queue), None)) is not None: new_rec = old + [call]
     else:
       # no such queue opened: close every open submit on this queue that shares a device, so submit order is kept
-      closing = [k for k in opened_qs if k[1] == queue and set(k[0]) & set(devs)]
-      new_src += [merge_sink((sa:=opened_qs.pop(k))[0]).call(aux=sa[1]).rtag("hcq") for k in closing]
-    opened_qs[(devs, queue)] = new_rec
-
-  return linear.replace(src=tuple(new_src + [merge_sink(sinks).call(aux=aux).rtag("hcq") for sinks, aux in opened_qs.values()]))
+      closing = [k for k in opened_qs if k[1] == queue and set(k[0]) & set(devices)]
+      new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in closing]
+      new_rec = [call]
+    opened_qs[(devices, queue)] = new_rec
+  return linear.replace(src=tuple(new_src + [_merged_hcq_call(c) for c in opened_qs.values()]))
 pm_merge_queues = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), merge_queues)])
 
 # *****************
-# 4.2. global sync
+# 2.4. finalizer
+
+def add_finalizer(ctx:itertools.count, linear:UOp) -> UOp:
+  # collect by device type
+  parts:dict[str, list[UOp]] = collections.defaultdict(list)
+  for call in linear.src:
+    if (c:=unwrap_after(call)).src[0].op is not Ops.CUSTOM_FUNCTION or c.src[0].arg != "hcq": continue
+    parts[c.arg.aux.device[0].split(':')[0]].append(unwrap_after(get_submit(call).src[0].src[0]))
+
+  nbump = next(ctx)
+  finalizers = []
+  for calls in parts.values():
+    devs = tuple(dedup(d for call in calls for d in unwrap_after(call).arg.aux.device))
+    zero = UOp.const(dtypes.int, 0)
+    tl = make_signal_value(devs)
+
+    # split each (multi-device) call into per-device deps, then store the device timeline value into the device signal after them
+    dep_lanes = [(call, dlane, devs.index(d)) for call in calls for dlane, d in enumerate(unwrap_after(call).arg.aux.device)]
+    store = make_deps(make_signal(devs).store(tl.index(zero)), dep_lanes, len(devs))
+    submit = make_submit(store, devs=devs, queue="COMPUTE:0")
+
+    upd = [(tl, 1)] + [(make_signal_value(devs, queue=qn), nbump) for qn in dedup([unwrap_after(call).arg.aux.queue for call in calls])]
+    patches = [s.after(submit).index(zero, dtype=s.dtype).store(s.index(zero) + inc) for s, inc in upd]
+    finalizers.append(UOp.custom_function("hcq", UOp.barrier(*patches).sink()).call(aux=HCQInfo("hcq finalizer", Estimates(), devs, "COMPUTE:0")))
+  return linear.replace(src=linear.src + tuple(finalizers))
+pm_add_finalizer = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), add_finalizer)])
+
+# *****************
+# 2.4. global sync
 
 def add_global_sync(ctx:set[tuple[str, ...]], submit:UOp, q:UOp) -> UOp|None:
   if (devs:=q.arg[0]) in ctx: return None
@@ -410,62 +332,118 @@ def add_global_sync(ctx:set[tuple[str, ...]], submit:UOp, q:UOp) -> UOp|None:
   # some devices from a command buffer might be used for the first time this schedule, so we wait for their global timeline epoch.
   wait = make_signal(devs).wait(make_signal_value(devs).index(UOp.const(dtypes.int, 0)) - 1)
   return submit.replace(src=(q.replace(src=(UOp(Ops.BARRIER, dtypes.void), wait, *q.src)),))
-pm_add_global_sync = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_global_sync)])
+pm_add_global_sync = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_global_sync)])
 
 # *****************
-# 4.3. annotate exec devs
+# 3.1. lower loads/stores
 
-pm_annotate_devs = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"),
-  lambda call: call.replace(arg=replace(call.arg, aux=replace(call.arg.aux, devs=get_submit(call.src[0]).src[0].arg[0]))))])
+def add_loads(ctx:set[int], submit:UOp, q:UOp) -> UOp|None:
+  cur_devs = q.arg[0]
+  new_src:list[UOp] = []
+  for s in q.src:
+    if s.op is Ops.AFTER:
+      for lanes, dep in zip(s.arg, s.src[1:]):
+        devs, queue = dep.arg.aux.device, dep.arg.aux.queue
+        ctx.add(dep.tag) # mark op to update signal.
+
+        sig = make_mstack([make_signal(d if dl is None else devs[dl], queue=queue, sentinel=dl is None) for dl, d in zip(lanes, cur_devs)])
+        val = make_mstack([make_signal_value(d if dl is None else devs[dl], queue=queue) for dl, d in zip(lanes, cur_devs)]).index(UOp.const(dtypes.int, 0))
+        new_src.append(sig.wait(val + dep.tag))
+      s = s.src[0]
+    new_src.append(s)
+  return submit.replace(src=(q.replace(src=tuple(new_src)),))
+pm_add_inner_loads = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_loads)])
+
+def add_stores(ctx:set[int], submit:UOp, q:UOp) -> UOp|None:
+  devs, queue = q.arg
+  new_src:list[UOp] = []
+  for op in q.src:
+    new_src.append(op)
+    if (sigval:=unwrap_after(op).tag) in ctx:
+      new_src.append(make_signal(devs, queue=queue).store(make_signal_value(devs, queue=queue).index(UOp.const(dtypes.int, 0)) + sigval))
+  return submit.replace(src=(q.replace(src=tuple(new_src)),))
+pm_add_inner_stores = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_stores)])
 
 # *****************
-# 4.4. replace params with per-submit input address loads
+# 2.1. hcq lowering: programs
 
-def replace_params(call:UOp) -> UOp|None:
-  if not (params:={u:u.arg.slot for u in call.src[0].toposort() if u.op is Ops.PARAM and u.addrspace is AddrSpace.GLOBAL}): return None
-
-  # fill new info
-  hcqinfo = replace(call.arg.aux, params=tuple(sorted(set(params.values()))), inputs=len(get_call_arg_uops(call)))
-
-  inputs = UOp.new_buffer(get_submit(call.src[0]).src[0].arg[0], len(hcqinfo.params), dtypes.uint64).rtag("inputs")
-
-  slot2idx = {s:i for i,s in enumerate(hcqinfo.params)}
-  body = call.src[0].substitute({u:inputs.index(UOp.const(dtypes.int, slot2idx[s])).load() for u,s in params.items()})
-
-  return call.replace(src=(body, *call.src[1:], inputs), arg=replace(call.arg, aux=hcqinfo))
-pm_replace_params = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), replace_params)])
+def encode_kernargs_clike(call:UOp, prg:UOp, devs:str|tuple[str, ...]) -> UOp:
+  data, info = prg.arg
+  call_args = get_call_arg_uops(call)
+  buf = make_placeholder(devs, data.kernargs_alloc_size, dtypes.uint8, name="kernargs")
+  patches = [make_patch(buf, i*8, make_getaddr(call_args[gi], devs)) for i,gi in enumerate(info.globals)] \
+          + [make_patch(buf, len(info.globals)*8 + i*4, v, dtypes.uint32) for i,v in enumerate(info.vars)]
+  return buf.after(*patches)
 
 # *****************
-# 5.1. encode cmdbufs
-
-@functools.cache
-def get_pm_lower(name:str) -> PatternMatcher|None:
-  try:
-    importlib.import_module(f'tinygrad.runtime.ops_{name.lower()}') # TODO: remove that
-    return importlib.import_module(f'extra.hcq2.ops_{name.lower()}2').pm_lower
-  except ImportError: return None
+# 2.2. hcq lowering: ops to ir
 
 def encode_cmdbuf(submit:UOp, lin:UOp) -> UOp|None:
-  if (pm:=get_pm_lower(to_tuple(lin.arg[0])[0].split(":")[0])) is None: return None
-  return pm.rewrite(submit)
-pm_encode_cmdbufs = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit", src=(UPat(Ops.LINEAR, name="lin"),), name="submit"), encode_cmdbuf)])
+  if (pm:=Device.get_class(lin.arg[0][0]).pm_lower) is None: return None
+  return graph_rewrite(submit, pm, name=f"encode {lin.arg[0]}", enter_calls=True)
+pm_encode_cmdbufs = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="lin"),), name="submit"), encode_cmdbuf)])
 
 # *****************
-# 5.2. lift patches to the command buffer (root)
 
-def lift_patches_to_cmdbuf(cmdbuf:UOp) -> UOp|None:
-  if not (patches:=dedup(u for store in cmdbuf.src[1:] for u in store.toposort() if u.op is Ops.AFTER)): return None
-  deps = tuple(d for p in patches for d in p.src[1:])
-  return cmdbuf.replace(src=cmdbuf.src + deps).substitute({p: p.src[0] for p in patches})
-pm_lift_patches_to_cmdbuf = PatternMatcher([
-  (UPat(Ops.AFTER, src=(UPat(Ops.BUFFER, tag={"compute", "copy"}),), allow_any_len=True, name="cmdbuf"), lift_patches_to_cmdbuf),
+pm_early_simplify = PatternMatcher([
+  # getaddr(slice(base, off)) -> getaddr(base) + byte offset
+  (UPat(Ops.GETADDR, src=(UPat(Ops.SLICE, src=(UPat.var("base"), UPat.cvar("off"))),), name="g"),
+   lambda g, base, off: g.replace(src=(base,)) + UOp.const(dtypes.uint64, off.arg * base.dtype.itemsize)),
 ])
 
+# *****************
+
+def replace_params(call:UOp) -> UOp|None:
+  gaddrs = [u for u in call.src[0].toposort(enter_calls=False) if u.op is Ops.GETADDR and u.src[0].op is Ops.PARAM and u.src[0].tag is None]
+  if not gaddrs: return None
+
+  idxs:dict[int, int] = {}
+  for g in gaddrs: idxs.setdefault(g.src[0].arg.slot, len(idxs))
+
+  inputs = make_placeholder(call.arg.aux.device, len(idxs), dtypes.uint64, "inputs")
+  body = call.src[0].substitute({g: inputs.index(UOp.const(dtypes.int, idxs[g.src[0].arg.slot])).load() for g in gaddrs})
+  return call.replace(src=(body, *call.src[1:], inputs), arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=tuple(idxs))))
+pm_replace_params = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), replace_params)])
+
+# *****************
+
+def changes_per_submit(u:UOp) -> bool: return u.op in (Ops.LOAD, Ops.INDEX) or (u.op is Ops.PARAM and u.tag is None)
+def is_placeholder(b:UOp) -> bool: return (b.op is Ops.PARAM and b.tag is not None) or (b.op is Ops.MSTACK and all(is_placeholder(x) for x in b.src))
+def is_link_patch(s:UOp) -> bool: return is_placeholder(s.buf_uop) and not any(changes_per_submit(u) for u in s.backward_slice)
+
+def trim_link_patches(ctx:list[UOp], a:UOp) -> UOp|None:
+  links, kept = partition(a.src[1:], is_link_patch)
+  ctx += links
+  return a.src[0].after(*kept) if links else None
+pm_trim_link_patches = PatternMatcher([(UPat(Ops.AFTER, src=(UPat((Ops.PARAM, Ops.MSTACK)),), allow_any_len=True, name="a"), trim_link_patches)])
+
+def split_patches(call:UOp) -> UOp|None:
+  # trim link-time patches
+  body = graph_rewrite(call.src[0], pm_trim_link_patches, ctx=(lt_patches:=[]), name=f"trim link-time patches ({call.arg.aux.name})")
+
+  units:dict[UOp, None] = {}
+  def unit_gate(u:UOp) -> bool:
+    if (is_plc:=is_placeholder(u)): units[u] = None
+    return not is_plc
+  body.toposort(gate=unit_gate)
+
+  srcs = dedup(list(call.src[1:]) + list(units) + [s.buf_uop for s in lt_patches])
+  param_sub = {u: UOp.param(i, u.dtype, device=u.device) for i,u in enumerate(srcs)}
+  for b in dedup(s.buf_uop for s in lt_patches):
+    idx = param_sub[b].arg.slot
+    srcs[idx] = srcs[idx].after(*dedup(s for s in lt_patches if s.buf_uop is b))
+
+  param_sub |= {v: v.replace(arg=replace(v.arg, slot=-1)) for v in body.variables() if v.op is Ops.PARAM}
+
+  info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(srcs) if unwrap_after(u).tag == "inputs"), None))
+  return call.replace(src=(body.substitute(param_sub), *srcs), arg=replace(call.arg, aux=info))
+pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), split_patches)])
+
 # *****************
 # 5.3. pack placeholders buffers
 
 def pack_hcq_placeholders(call:UOp) -> UOp|None:
-  bufs = [b for b in call.src[0].toposort() if b.op is Ops.BUFFER and b.tag in (maxtags:={"scratch"}) | (sumtags:={"program", "kernargs"})]
+  bufs = [b for b in call.src[0].toposort() if b.op is Ops.PARAM and b.tag in (maxtags:={"scratch"}) | (sumtags:={"program", "kernargs"})]
 
   off_per_buf:dict[UOp, int] = {}
   size_per_tag:dict[str, int] = {}
@@ -481,15 +459,7 @@ def pack_hcq_placeholders(call:UOp) -> UOp|None:
   bases = {tag:UOp.new_buffer(b.device, size_per_tag[tag], b.dtype).rtag(tag) for tag,b in ref_bufs.items()}
   subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(dtypes.weakint, off_per_buf.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases}
   return call.replace(src=(call.src[0].substitute(subs, walk=True), *call.src[1:])) if subs else None
-pm_pack_placeholders = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), pack_hcq_placeholders)])
-
-# *****************
-# 5.4. capture buffers reachable from each hcq call as BIND, so we don't drop their refs
-
-def hold_call_buffers(call:UOp) -> UOp|None:
-  if not (bufs:=tuple(dedup(u for u in call.src[0].toposort() if u.op is Ops.BUFFER and u not in call.src))): return None
-  return call.replace(src=call.src + (UOp(Ops.BIND, dtypes.void, src=bufs),))
-pm_hold_call_buffers = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), hold_call_buffers)])
+pm_pack_placeholders = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), pack_hcq_placeholders)])
 
 # *****************
 # 6. bufferize placeholders: replace placeholders with real buffers.
@@ -498,7 +468,7 @@ def bufferize_buf(buf:UOp) -> UOp|None:
   if buf.tag is None: return None
   uops = tuple(UOp.from_buffer((dv:=Device[dev]).pm_bufferize.rewrite(buf, ctx=dv), "CPU") for dev in to_tuple(buf.device))
   return make_mstack(uops)
-pm_bufferize = PatternMatcher([(UPat(Ops.BUFFER, name="buf"), bufferize_buf)])
+pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, name="buf"), bufferize_buf)])
 
 # *****************
 # 7. resolve patches
@@ -512,20 +482,20 @@ def fold_blob_store(buf:UOp, blob:UOp) -> UOp:
 
 def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp:
   for b, v in zip((bs:=mb.bufs if isinstance((mb:=buf.buffer), MultiBuffer) else (mb,)), val.src if val.op is Ops.STACK else (val,)*len(bs)):
-    struct.pack_into(f'<{v.dtype.fmt}', b.ensure_allocated()._buf.cpu_view().mv.cast('B'), off.arg * buf.dtype.base.itemsize, v.arg)
+    struct.pack_into(f'<{v.dtype.fmt}', b.ensure_allocated()._buf.cpu_view().mv.cast('B'), off.arg * buf.dtype.base.itemsize, truncate[v.dtype](v.arg))
   return UOp(Ops.NOOP)
 
 def resolve_getaddr(buf:UOp, g:UOp) -> UOp:
   if buf.op not in (Ops.BUFFER, Ops.MSTACK, Ops.MSELECT): return buf
-  devs, b = to_tuple(g.src[1].arg), buf.buffer
+  devs, b = to_tuple(g.arg), buf.buffer
   bufs = tuple(cast(Buffer, x.buffer) for x in buf.src) if buf.op is Ops.MSTACK else tuple(b.bufs if isinstance(b, MultiBuffer) else (b,)*len(devs))
   assert len(bufs) == len(devs), f"can't resolve {len(bufs)} buffers on {len(devs)} devices"
   addrs = tuple(UOp.const(dtypes.uint64, x.get_buf(d).va_addr) for x, d in zip(bufs, devs))
   return addrs[0] if len(addrs) == 1 else UOp(Ops.STACK, dtypes.uint64.vec(len(addrs)), addrs)
 
-def resolve_getaddr_slice(bv:UOp, dev:UOp) -> UOp:
+def resolve_getaddr_slice(bv:UOp, g:UOp) -> UOp:
   itemsize = bv.src[0].dtype.itemsize if unwrap_after(bv.src[0]).op in (Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize
-  return UOp(Ops.GETADDR, dtypes.uint64, src=(bv.src[0], dev)) + UOp.const(dtypes.uint64, bv.src[1].arg * itemsize)
+  return UOp(Ops.GETADDR, dtypes.uint64, src=(bv.src[0],), arg=g.arg) + UOp.const(dtypes.uint64, bv.src[1].arg * itemsize)
 
 pm_resolve_patches = PatternMatcher([
   # multi
@@ -537,61 +507,54 @@ pm_resolve_patches = PatternMatcher([
     lambda shr, bv: shr.replace(src=(bv.src[0], shr.src[1] + bv.src[1].cast(shr.src[1].dtype), shr.src[2]))),
 
   # getaddr
-  (UPat(Ops.GETADDR, src=(UPat(Ops.SLICE, name="bv"), UPat(Ops.DEVICE, name="dev"))), resolve_getaddr_slice), # getaddr(slice(x)) -> offset+getaddr(x)
-  (UPat(Ops.GETADDR, src=(UPat(name="buf"), UPat(Ops.DEVICE)), name="g"), resolve_getaddr),
+  (UPat(Ops.GETADDR, src=(UPat(Ops.SLICE, name="bv"),), name="g"), resolve_getaddr_slice), # getaddr(slice(x)) -> offset+getaddr(x)
+  (UPat(Ops.GETADDR, src=(UPat(name="buf"),), name="g"), resolve_getaddr),
 
   # folders
   (UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf").store(UPat(Ops.BINARY, name="blob")), fold_blob_store),
   (UPat(Ops.SHRINK, src=(UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf"), UPat.cvar("off"), UPat(Ops.CONST))).bitcast()
     .store(UPat.any(UPat.cvar("val"), UPat(Ops.STACK, name="val"))), fold_const_store),
-]) + symbolic_simple
+])
 
 # *****************
 # 8. callify hcq programs
 
-def to_param(bufs:list[UOp], ref:UOp) -> UOp:
-  if ref not in bufs: bufs.append(ref)
-  return UOp.placeholder((ref.buffer.size,), ref.dtype, bufs.index(ref))
-pm_to_param = PatternMatcher([(UPat({Ops.MSELECT, Ops.MSTACK, Ops.BUFFER}, name="r"), lambda ctx, r: to_param(ctx, r))])
+pm_callify_hcq = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="hcq", src=(UPat(Ops.SINK),), name="cf"),
+  lambda cf: cf.replace(src=(to_program(cf.src[0].replace(arg=KernelInfo("hcq_submit"), tag=1), Device["CPU"].renderer),)))])
 
-def parametrize_host_buffers(call:UOp) -> UOp:
-  # preserve original order of args
-  body = graph_rewrite(call.src[0], pm_to_param, ctx=(bufs:=list(get_call_arg_uops(call))), bottom_up=True, name="parametrize host buffers")
+hcq_compile_cache:dict[bytes, tuple[UOp, tuple[UOp, ...]]] = {}
 
-  # move vars to new slots
-  var_slots = {nm:len(bufs)+i for i,nm in enumerate(sorted({v.expr for v in body.variables() if v.op is Ops.PARAM}))}
-  body = body.substitute({v:v.replace(arg=replace(v.arg, slot=var_slots[v.expr])) for v in body.variables() if v.op is Ops.PARAM})
+@track_rewrites(lambda linear,input_uops,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
+def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None) -> UOp:
+  if input_uops is not None: linear = graph_rewrite(linear, pm_replace_buffers, ctx=input_uops, walk=True, enter_calls=True, name="replace buffer")
 
-  return call.replace(src=(body, *bufs) + tuple(x for x in call.src[1:] if x.op is Ops.BIND))
-pm_parametrize_host_buffers = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), parametrize_host_buffers)])
+  if (final_linear:=(hcq_compile_cache.get(cache_key:=linear.key))) is None:
+    # schedule
+    linear = linear.substitute(back_map:={s.param_like(i): s for i,s in enumerate(input_uops)} if input_uops is not None else {}, walk=True)
+    linear = graph_rewrite(linear, pm_insert_copy_staging + pm_flatten_linear, name="insert copy staging")
+    linear = graph_rewrite(linear, pm_tag_hcq_calls, ctx=(enumerator:=itertools.count(0)), walk=True, name="tag hcq calls")
+    linear = graph_rewrite(linear, pm_sched_sync, ctx=HCQDepsTracker(), walk=True, name="schedule sync")
+    linear = linear.substitute({s: p for p, s in back_map.items()}, walk=True)
+    linear = graph_rewrite(linear, pm_merge_queues, walk=True, name="merge queues")
+    linear = graph_rewrite(linear, pm_add_finalizer, ctx=enumerator, walk=True, name="add finalizer")
+    linear = graph_rewrite(linear, pm_add_global_sync, ctx=set(), walk=True, name="add global sync", enter_calls=True)
 
-def callify_hcq(call:UOp) -> UOp:
-  prg = to_program(call.src[0].sink(arg=KernelInfo("hcq_submit"), tag=1), Device["CPU"].renderer)
-  return UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(prg,), arg="hcq").call(*call.src[1:], aux=call.arg.aux)
-pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), callify_hcq)])
+    # lowering to hcq ir
+    linear = graph_rewrite(linear, pm_add_inner_loads, ctx=(waited:=set()), walk=True, name="add loads", enter_calls=True)
+    linear = graph_rewrite(linear, pm_add_inner_stores, ctx=waited, walk=True, name="add stores", enter_calls=True)
+    linear = graph_rewrite(linear, pm_encode_cmdbufs, walk=True, name="encode cmdbufs", enter_calls=True)
 
-@track_rewrites(lambda _,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
-def hcq_compile(linear:UOp) -> UOp:
-  linear = graph_rewrite(linear, pm_insert_copy_staging + pm_flatten_linear, name="insert copy staging")
-  linear = graph_rewrite(linear, pm_prep_runtime, name="prepare runtime")
+    # pie
+    linear = graph_rewrite(linear, pm_early_simplify + symbolic, bottom_up=False, name="early simplify patches")
+    linear = graph_rewrite(linear, pm_replace_params, walk=True, name="replace params")
+    linear = graph_rewrite(linear, pm_split_patches, walk=True, name="split patches")
 
-  linear = graph_rewrite(linear, pm_lower_ops, name="lower ops into hcq ir")
-  linear = graph_rewrite(linear, pm_schedule_inner_sync, ctx=(deps_ctx:=DepsCtx()), walk=True, name="schedule inner sync")
-  linear = graph_rewrite(linear, pm_add_finalizer, ctx=deps_ctx, walk=True, name="add finalizer")
-  linear = graph_rewrite(linear, pm_add_inner_loads, ctx=(waited:=set()), walk=True, name="add loads", enter_calls=True)
-  linear = graph_rewrite(linear, pm_add_inner_stores, ctx=waited, walk=True, name="add stores", enter_calls=True)
-  linear = graph_rewrite(linear, pm_merge_queues, name="merge queues")
-  linear = graph_rewrite(linear, pm_add_global_sync, ctx=set(), walk=True, name="add global sync", enter_calls=True)
-  linear = graph_rewrite(linear, pm_annotate_devs, name="annotate devs")
-  linear = graph_rewrite(linear, pm_replace_params, name="replace params")
-  linear = graph_rewrite(linear, pm_encode_cmdbufs, walk=True, name="encode cmdbufs", enter_calls=True)
-  linear = graph_rewrite(linear, pm_lift_patches_to_cmdbuf, name="lift patches to cmdbuf", enter_calls=True)
-  linear = graph_rewrite(linear, pm_pack_placeholders, walk=True, name="pack placeholders")
-  return graph_rewrite(linear, pm_hold_call_buffers, walk=True, name="hold call buffers")
+    # and compile it
+    final_linear = hcq_compile_cache[cache_key] = graph_rewrite(linear, pm_callify_hcq, name="callify hcq", enter_calls=True)
+
+  return final_linear
 
 @track_rewrites(lambda _,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}")
 def hcq_link(linear:UOp) -> UOp:
-  linear = graph_rewrite(linear, pm_bufferize, bottom_up=True, walk=True, name="bufferize placeholders", enter_calls=True)
-  linear = graph_rewrite(linear, pm_resolve_patches, bottom_up=False, name="simplify patches", enter_calls=True)
-  linear = graph_rewrite(linear, pm_parametrize_host_buffers, walk=True, name="parametrize host buffers")
-  return graph_rewrite(linear, pm_callify_hcq, name="callify hcq")
+  linear = graph_rewrite(linear, pm_bufferize, bottom_up=True, walk=True, name="bufferize placeholders")
+  return graph_rewrite(linear, pm_resolve_patches + symbolic, bottom_up=False, name="simplify patches")

extra/hcq2/ops_amd2.py

@@ -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 extra.hcq2.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, make_getaddr, make_ins, make_cmdbuf
+from extra.hcq2.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_getaddr, make_ins, make_cmdbuf, make_placeholder
 from tinygrad.uop.ops import sint, UOp
 from tinygrad.device import Compiled, BufferSpec, Buffer, Device
 from tinygrad.dtype import dtypes
@@ -102,11 +102,12 @@ def pm4_timestamp(ctx, dst):
   return release_mem(ctx, make_getaddr(dst, ctx.devs), 0, ctx.pm4.data_sel__mec_release_mem__send_gpu_clock_counter,
                      ctx.pm4.int_sel__mec_release_mem__none)
 
-def pm4_program(ctx, prg):
+def pm4_program(ctx, call, prg):
   data, info = prg.arg
-  lib_gpu, args = prg.src
+  lib_gpu = prg.src[0]
+  args = encode_kernargs_clike(call, prg, ctx.devs)
   prog_addr = make_getaddr(lib_gpu, ctx.devs) + data.entry_point_offset
-  scratch_addr = make_getaddr(UOp.new_buffer(lib_gpu.device, data.private_segment_size, dtypes.uint8).rtag("scratch"), ctx.devs)
+  scratch_addr = make_getaddr(make_placeholder(ctx.devs, data.private_segment_size, dtypes.uint8, "scratch", unique=False), ctx.devs)
   args_addr = make_getaddr(args, ctx.devs)
 
   user_regs = []
@@ -134,9 +135,10 @@ def pm4_program(ctx, prg):
   return UOp(Ops.LINEAR, dtypes.void, tuple(ins))
 
 pm_pm4_opsel = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), pm4_program),
+
   (UPat(Ops.WAIT, src=(UPat(name="dst"), UPat(name="val"))), pm4_wait),
   (UPat(Ops.BARRIER), pm4_barrier),
-  (UPat(Ops.PROGRAM, name="prg"), pm4_program),
   (UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", src=(UPat(name="dst"),)), pm4_timestamp),
   (UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
 ])
@@ -146,7 +148,7 @@ def pm4_submit(cmdbuf, devs):
 
   # the compute queue's ring and its host-side ring/write/put pointers (placeholders, resolved in pm_bufferize)
   for d in devs: q = Device[d].compute_queue
-  ring, wptr, doorbell, put_ptr = (UOp.new_buffer(devs, b.size, b.dtype).rtag(("COMPUTE:0", name))
+  ring, wptr, doorbell, put_ptr = (make_placeholder(devs, b.size, b.dtype, ("COMPUTE:0", name), unique=False)
     for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
 
   # place the cmdbuf at the ring's write offset, wrapping the ring
@@ -156,29 +158,31 @@ def pm4_submit(cmdbuf, devs):
   ring_idx = ((put + i.cast(put.dtype)) % q.ring.size).cast(dtypes.int)
 
   # copy the cmdbuf into the ring and advance the put/write pointers
-  copy_to_ring = ring.index(ring_idx, dtype=ring.dtype.ptr()).store(
-    cmdbuf.index(i*4, dtype=cmdbuf.dtype.ptr()).cast(dtypes.uint32.ptr()).load()).end(i)
-  bump_put_ptr = put_ptr.index(zero, dtype=put_ptr.dtype.ptr()).store(next_put)
-  bump_wptr = wptr.index(zero, dtype=wptr.dtype.ptr()).store(next_put)
+  copy_to_ring = ring.index(ring_idx, ptr=True).store(
+    cmdbuf.index(i*4, ptr=True).cast(dtypes.uint32.ptr()).load()).end(i)
+  bump_put_ptr = put_ptr.index(zero, ptr=True).store(next_put)
+  bump_wptr = wptr.index(zero, ptr=True).store(next_put)
 
   # ring the doorbell once the copy and pointer bumps have landed
   flush = UOp.barrier(copy_to_ring, bump_put_ptr, bump_wptr)
-  return doorbell.after(flush).index(zero, dtype=doorbell.dtype.ptr()).store(next_put)
+  return doorbell.after(flush).index(zero, ptr=True).store(next_put)
 
 pm_pm4_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
-  lambda lin: pm4_submit(make_cmdbuf(lin, to_tuple(lin.arg[0]), "compute"), to_tuple(lin.arg[0])))])
+  lambda lin: pm4_submit(make_cmdbuf(lin, to_tuple(lin.arg[0])), to_tuple(lin.arg[0])))])
 
 # *****************
 # SDMA
 
 class SDMAOps(FastEnum): COPY = auto(); POLL_REGMEM = auto(); FENCE = auto(); TRAP = auto(); TIMESTAMP = auto()  # noqa: E702
 
-def sdma_copy(ctx, dst, src, copy):
+def sdma_copy(ctx, call):
+  dst, src = call.src[1], call.src[2]
+  sz = src.max_numel() * src.dtype.base.itemsize
   src_addr, dst_addr = make_getaddr(src, ctx.devs), make_getaddr(dst, ctx.devs)
   return UOp(Ops.LINEAR, dtypes.void, tuple([make_ins(SDMAOps.COPY,
      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(copy.arg - off, ctx.max_copy_size) - 1), 0,
-     *data64_le(src_addr + off), *data64_le(dst_addr + off)) for off in range(0, copy.arg, ctx.max_copy_size)]))
+     ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz - off, ctx.max_copy_size) - 1), 0,
+     *data64_le(src_addr + off), *data64_le(dst_addr + off)) for off in range(0, sz, ctx.max_copy_size)]))
 
 def sdma_wait(ctx, dst, val):
   op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
@@ -196,9 +200,10 @@ def sdma_timestamp(ctx, dst):
   return make_ins(SDMAOps.TIMESTAMP, op, *data64_le(make_getaddr(dst, ctx.devs)))
 
 pm_sdma_opsel = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),
+
   (UPat(Ops.BARRIER), lambda: UOp(Ops.NOOP, dtypes.void, ())),
   (UPat(Ops.WAIT, src=(UPat(name="dst"), UPat(name="val"))), sdma_wait),
-  (UPat(Ops.COPY, src=(UPat(name="dst"), UPat(name="src")), name="copy"), sdma_copy),
   (UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", src=(UPat(name="dst"),)), sdma_timestamp),
   (UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), sdma_store),
 ])
@@ -209,7 +214,7 @@ def sdma_submit(cmdbuf, devs):
 
   # the sdma queue's ring and its host-side ring/write/put pointers
   for d in devs: q = Device[d].sdma_queue(0)
-  ring, wptr, doorbell, put_ptr = (UOp.new_buffer(devs, b.size, b.dtype).rtag(("COPY:0", name))
+  ring, wptr, doorbell, put_ptr = (make_placeholder(devs, b.size, b.dtype, ("COPY:0", name), unique=False)
     for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
 
   # sdma needs the cmdbuf contiguous: if it won't fit before the ring end, restart at 0 and zero the tail
@@ -221,22 +226,33 @@ def sdma_submit(cmdbuf, devs):
 
   # zero the wrapped tail, then copy the cmdbuf into the ring
   zi = UOp.range(zero_amt_dw, 0, dtype=dtypes.int, src=(cmdbuf,))
-  zero_tail = ring.index(tail_off_dw + zi, dtype=ring.dtype.ptr()).store(UOp.const(dtypes.uint32, 0)).end(zi)
+  zero_tail = ring.index(tail_off_dw + zi, ptr=True).store(UOp.const(dtypes.uint32, 0)).end(zi)
   i = UOp.range(UOp.const(dtypes.int, size_dw), 0, dtype=dtypes.int, src=(cmdbuf,))
-  copy_to_ring = ring.index(start_dw + i, dtype=ring.dtype.ptr()).store(
-    cmdbuf.index(i*4, dtype=cmdbuf.dtype.ptr()).cast(dtypes.uint32.ptr()).load()).end(i)
+  copy_to_ring = ring.index(start_dw + i, ptr=True).store(
+    cmdbuf.index(i*4, ptr=True).cast(dtypes.uint32.ptr()).load()).end(i)
 
   # advance the put/write pointers past the zeroed tail and the cmdbuf
   next_put_b = put_b + ((zero_amt_dw + size_dw) * 4).cast(put_b.dtype)
-  bump_put_ptr = put_ptr.index(zero, dtype=put_ptr.dtype.ptr()).store(next_put_b)
-  bump_wptr = wptr.index(zero, dtype=wptr.dtype.ptr()).store(next_put_b)
+  bump_put_ptr = put_ptr.index(zero, ptr=True).store(next_put_b)
+  bump_wptr = wptr.index(zero, ptr=True).store(next_put_b)
 
   # ring the doorbell once the writes have landed
   flush = UOp.barrier(zero_tail, copy_to_ring, bump_put_ptr, bump_wptr)
-  return doorbell.after(flush).index(zero, dtype=doorbell.dtype.ptr()).store(next_put_b)
+  return doorbell.after(flush).index(zero, ptr=True).store(next_put_b)
 
 pm_sdma_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
-  lambda lin: sdma_submit(make_cmdbuf(lin, to_tuple(lin.arg[0]), "copy"), to_tuple(lin.arg[0])))])
+  lambda lin: sdma_submit(make_cmdbuf(lin, to_tuple(lin.arg[0])), to_tuple(lin.arg[0])))])
+
+@dataclass(frozen=True)
+class AMDEncodeCtx:  # encode-time constants for one queue: devs (every cmdbuf address resolves into these) + gfx version + packet/ip modules
+  devs: tuple[str, ...]; target: tuple[int, ...]; pm4: Any; sdma: Any; soc: Any  # noqa: E702
+  gc: AMDIP; nbio: AMDIP; xccs: int; max_copy_size: int; tmpring_size: Callable  # noqa: E702
+
+def encode_queue(q:UOp) -> UOp|None:
+  d = Device[(devs:=to_tuple(q.arg[0]))[0]]
+  ctx = AMDEncodeCtx(devs, d.target, d.pm4, d.sdma, d.soc, d.gc, d.nbio, d.xccs, d.max_copy_size, d.tmpring_size)
+  opsel, submit = (pm_pm4_opsel, pm_pm4_submit) if q.arg[1].startswith("COMPUTE") else (pm_sdma_opsel, pm_sdma_submit)
+  return submit.rewrite(graph_rewrite(q, opsel + pm_flatten_linear, walk=True, ctx=ctx, name=f"{q.arg[1]} opsel"))
 
 @dataclass(frozen=True)
 class AMDProgramData:
@@ -245,7 +261,6 @@ class AMDProgramData:
   enable_dispatch_ptr:int; enable_private_segment_sgpr:int
 
 _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
   if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, dev.device))) is None:
@@ -258,22 +273,17 @@ def amd_build_program(prg:UOp) -> UOp:
     if (lds:=((desc.group_segment_fixed_size+511)//512)&0x1FF) > (dev.iface.props['lds_size_in_kb']*1024)//512:
       raise RuntimeError("Too many resources requested: group_segment_size")
     edp = desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_DISPATCH_PTR
-    cached = _amd_program_cache[key] = (AMDProgramData(
-      entry_point_offset=rodata + desc.kernel_code_entry_byte_offset,
+
+    data = AMDProgramData(entry_point_offset=rodata + desc.kernel_code_entry_byte_offset,
       rsrc1=desc.compute_pgm_rsrc1 | ((1<<20) if dev.target[0]==11 else 0),  # priv=1 on gfx11 for cwsr
       rsrc2=desc.compute_pgm_rsrc2 | (lds<<15), rsrc3=desc.compute_pgm_rsrc3,
-      wave32=bool(desc.kernel_code_properties & 0x400),
-      private_segment_size=desc.private_segment_fixed_size,
-      kernargs_segment_size=desc.kernarg_size,
-      kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0),
-      enable_dispatch_ptr=edp,
-      enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER), bytes(image))
+      wave32=bool(desc.kernel_code_properties & 0x400), private_segment_size=desc.private_segment_fixed_size, kernargs_segment_size=desc.kernarg_size,
+      kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0), enable_dispatch_ptr=edp,
+      enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER)
+    buf = make_placeholder(prg.device, len(image), dtypes.uint8, "program")
+    cached = _amd_program_cache[key] = prg.replace(src=(buf.after(buf.store(UOp(Ops.BINARY, dtypes.void, src=(), arg=bytes(image)))),), arg=(data, prg.arg))
   return cached
 
-pm_prep_program = PatternMatcher([
-  (UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program),
-])
-
 class AMDAllocator(HCQAllocator['AMDDevice']):
   def __init__(self, dev:AMDDevice):
     super().__init__(dev, supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_queue and not dev.is_usb())
@@ -516,23 +526,15 @@ class PCIIface(PCIIfaceBase):
 
 def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {})
 
-@dataclass(frozen=True)
-class AMDEncodeCtx:  # encode-time constants for one queue: devs (every cmdbuf address resolves into these) + gfx version + packet/ip modules
-  devs: tuple[str, ...]; target: tuple[int, ...]; pm4: Any; sdma: Any; soc: Any  # noqa: E702
-  gc: AMDIP; nbio: AMDIP; xccs: int; max_copy_size: int; tmpring_size: Callable  # noqa: E702
-
-def encode_queue(q:UOp) -> UOp|None:
-  if not (isinstance(q.arg, tuple) and len(q.arg) == 2 and isinstance(q.arg[1], str) and q.arg[1].startswith(("COMPUTE", "COPY"))): return None
-  d = Device[(devs:=to_tuple(q.arg[0]))[0]]
-  ctx = AMDEncodeCtx(devs, d.target, d.pm4, d.sdma, d.soc, d.gc, d.nbio, d.xccs, d.max_copy_size, d.tmpring_size)
-  opsel, submit = (pm_pm4_opsel, pm_pm4_submit) if q.arg[1].startswith("COMPUTE") else (pm_sdma_opsel, pm_sdma_submit)
-  return submit.rewrite(graph_rewrite(q, opsel + pm_flatten_linear, walk=True, ctx=ctx, name=f"{q.arg[1]} opsel"))
-
-pm_lower = PatternMatcher([
-  (UPat(Ops.CUSTOM_FUNCTION, arg="submit", src=(UPat(Ops.LINEAR, name="q"),)), encode_queue),
-])
-
 class AMDDevice(HCQ2Compiled):
+  pm_lower = PatternMatcher([
+    # prep program
+    (UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program),
+
+    # encoding of cmdbuf
+    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),)), encode_queue),
+  ])
+
   timestamp_divider = 100.0  # AMD GPU clock: ticks/us
 
   ifaces = [KFDIface, PCIIface]
@@ -579,7 +581,7 @@ class AMDDevice(HCQ2Compiled):
 
     # Scratch setup
     self.max_private_segment_size = 0
-    self.pm_bufferize = PatternMatcher([(UPat(Ops.BUFFER, tag="scratch", name="b"), lambda ctx, b: ctx.scratch_buffer(b.max_numel()))]) + self.pm_bufferize
+    self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag="scratch", name="b"), lambda ctx, b: ctx.scratch_buffer(b.max_numel()))]) + self.pm_bufferize
 
     self.pmc_enabled:bool = PROFILE > 0 and PMC > 0
     if self.pmc_enabled:
@@ -627,10 +629,10 @@ class AMDDevice(HCQ2Compiled):
 
     qname = f"{'COPY' if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA else 'COMPUTE'}:{idx}"
     self.pm_bufferize = PatternMatcher([
-      (UPat(Ops.BUFFER, tag={(qname, name)}), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
+      (UPat(Ops.PARAM, tag={(qname, name)}), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
     ] + [
-      (UPat(Ops.BUFFER, tag={(qname, "timeline_signal")}), lambda ctx, q=qname: ctx.timeline_signal(q)),
-      (UPat(Ops.BUFFER, tag={(qname, "timeline_value")}), lambda ctx, q=qname: ctx.timeline_value(q)),
+      (UPat(Ops.PARAM, tag={(qname, "timeline_signal")}), lambda ctx, q=qname: ctx.timeline_signal(q)),
+      (UPat(Ops.PARAM, tag={(qname, "timeline_value")}), lambda ctx, q=qname: ctx.timeline_value(q)),
     ]) + self.pm_bufferize
 
     return queue

extra/thunder/tiny/tk/group.py

@@ -93,7 +93,7 @@ class Group:
           d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
 
           out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
-          c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
+          c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
           c_store = UOp.group(*c_i).end(height, width, inner)
 
     self.ker.push_store(c_store, c)
@@ -123,7 +123,7 @@ class Group:
           d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
 
           out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
-          c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
+          c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
           c_store = UOp.group(*c_i).end(height, width, inner)
 
     self.ker.push_store(c_store, c)
@@ -153,7 +153,7 @@ class Group:
           d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
 
           out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
-          c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
+          c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
           c_store = UOp.group(*c_i).end(height, width, inner)
 
     self.ker.push_store(c_store, c)
@@ -183,7 +183,7 @@ class Group:
           d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
 
           out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
-          c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
+          c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
           c_store = UOp.group(*c_i).end(height, width, inner)
 
     self.ker.push_store(c_store, c)

extra/tinyfs/fetch_file.py

@@ -1,4 +1,6 @@
 from tinygrad.tensor import Tensor
+from tinygrad.helpers import CHUNK_SIZE
+from tinygrad.nn.state import fs_load
 import argparse, math, hashlib
 
 def _python_hash_1mb(data:bytes|bytearray):
@@ -7,15 +9,15 @@ def _python_hash_1mb(data:bytes|bytearray):
   return hashlib.shake_128(b''.join(chunk_hashes)).digest(16)
 
 def hash_file(data: bytes|bytearray):
-  if len(data) % Tensor.CHUNK_SIZE != 0: data += bytes(Tensor.CHUNK_SIZE - len(data) % Tensor.CHUNK_SIZE)
-  base_chunks = math.ceil(len(data) / Tensor.CHUNK_SIZE)
-  tree_depth = math.ceil(math.log(base_chunks, Tensor.CHUNK_SIZE // 16))
+  if len(data) % CHUNK_SIZE != 0: data += bytes(CHUNK_SIZE - len(data) % CHUNK_SIZE)
+  base_chunks = math.ceil(len(data) / CHUNK_SIZE)
+  tree_depth = math.ceil(math.log(base_chunks, CHUNK_SIZE // 16))
 
   for _ in range(tree_depth + 1):
-    data_chunks = [data[i:i+Tensor.CHUNK_SIZE] for i in range(0, len(data), Tensor.CHUNK_SIZE)]
+    data_chunks = [data[i:i+CHUNK_SIZE] for i in range(0, len(data), CHUNK_SIZE)]
     data_chunk_hashes = [_python_hash_1mb(chunk) for chunk in data_chunks]
     data = b''.join(data_chunk_hashes)
-    if len(data) % Tensor.CHUNK_SIZE != 0: data += bytes(Tensor.CHUNK_SIZE - len(data) % Tensor.CHUNK_SIZE)
+    if len(data) % CHUNK_SIZE != 0: data += bytes(CHUNK_SIZE - len(data) % CHUNK_SIZE)
 
   return data[:16]
 
@@ -27,7 +29,7 @@ if __name__ == "__main__":
   parser.add_argument("--check", action="store_true", help="verify the file hash after fetching")
   args = parser.parse_args()
 
-  Tensor(bytes.fromhex(args.hash), device="CPU").fs_load(args.len).to(f"disk:{args.dest}").realize()
+  fs_load(Tensor(bytes.fromhex(args.hash), device="CPU"), args.len).to(f"disk:{args.dest}").realize()
 
   if args.check:
     with open(args.dest, "rb") as f:

extra/tinyfs/fetch_raid.py

@@ -3,6 +3,7 @@ from pathlib import Path
 
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import tqdm, getenv
+from tinygrad.nn.state import fs_load
 
 raid_root = Path(getenv("RAID_ROOT", "/raid"))
 
@@ -14,7 +15,7 @@ def fetch_file(item):
   path.parent.mkdir(parents=True, exist_ok=True)
 
   try:
-    pt = Tensor(bytes.fromhex(h), device="CPU").fs_load(size).to(f"disk:{path.as_posix()}").realize()
+    pt = fs_load(Tensor(bytes.fromhex(h), device="CPU"), size).to(f"disk:{path.as_posix()}").realize()
   except Exception as e:
     print(f"error fetching {path}, {h}, {size}: {e}")
     raise
@@ -22,7 +23,7 @@ def fetch_file(item):
   pt.uop.buffer.deallocate()
 
 def fetch_mapping(h, l):
-  mapping_tensor = Tensor(bytes.fromhex(h)).fs_load(l).realize()
+  mapping_tensor = fs_load(Tensor(bytes.fromhex(h)), l).realize()
   mapping = mapping_tensor.data().tobytes().decode()
   mapping = json.loads(mapping)
   mapped_files = mapping.items()

extra/tinyfs/upload_raid.py

@@ -3,12 +3,13 @@ import multiprocessing, json
 
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import tqdm
+from tinygrad.nn.state import fs_store
 
 raid_root = Path("/raid")
 
 def upload_file(path: Path):
   pt = Tensor(path).realize()
-  h = pt.fs_store().realize()
+  h = fs_store(pt).realize()
   pt.uop.realized.deallocate()
   return h.data().hex(), path, pt.nbytes()
 
@@ -26,6 +27,6 @@ if __name__ == "__main__":
 
   mapping = json.dumps(mapping).encode()
   mapping_tensor = Tensor(mapping, device="CPU")
-  h = mapping_tensor.fs_store().realize()
+  h = fs_store(mapping_tensor).realize()
 
   print(f"final hash: {h.data().hex()}, size: {len(mapping)}")

pyproject.toml

@@ -111,6 +111,7 @@ docs = [
   "black",
   "numpy",
 ]
+mesa = ["tinymesa==25.2.7.2"]
 
 
 [tool.mutmut]

spec/tinyspec.tex

@@ -227,7 +227,7 @@ Ternary & $(P, A, B)$
 \midrule
 \op{Barrier} & (deps\ldots) & --- & Synchronize threads within a workgroup. \\
 \op{Ins} & \ldots & \ldots & A single machine instruction (e.g.\ AMD ISA). \\
-\op{GetAddr} & (buf, dev) & --- & Lower buf to its address on device dev. \\
+\op{GetAddr} & (buf,) & dev & Lower buf to its address on device dev. \\
 \op{Special} & (bound,) & name & GPU thread/workgroup index (e.g.\ \texttt{gidx0}, \texttt{lidx1}). \\
 \op{If} & (gate,) & --- & Begin conditional execution block. \\
 \op{Endif} & (if,) & --- & End conditional execution block. \\

test/backend/test_arange.py

@@ -187,6 +187,26 @@ class TestIndexing(unittest.TestCase):
     for i in idx.flatten().numpy(): expected_grad[i] += 2
     np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)
 
+  @unittest.skipIf(Device.DEFAULT not in ("CPU", "AMD"), "atomics only on AMD/CPU")
+  @Context(USE_ATOMICS=1, SPEC=1)
+  def test_embedding_backward_padded_embed(self):
+    from tinygrad.renderer.cstyle import CStyleLanguage
+    if Device.DEFAULT == "CPU" and not isinstance(Device["CPU"].renderer, CStyleLanguage): self.skipTest("CPU needs Clang renderer")
+    vocab_size, embed_size = 1000, 300
+    bs, seqlen = 4, 256
+    idx = Tensor.randint(bs, seqlen, high=vocab_size)
+    emb = nn.Embedding(vocab_size, embed_size)
+    emb.weight = Tensor.ones(vocab_size, embed_size)
+    gt = Tensor.zeros(bs, seqlen, embed_size)
+    Tensor.realize(idx, emb.weight, gt)
+    loss = (emb(idx)-gt).square().sum()
+    loss.backward()
+    emb.weight.grad.realize()
+    # correctness check
+    expected_grad = np.zeros((vocab_size, embed_size), dtype=np.float32)
+    for i in idx.flatten().numpy(): expected_grad[i] += 2
+    np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)
+
   @needs_second_gpu
   @unittest.skipIf(Device.DEFAULT not in ("CPU", "AMD"), "atomics only on AMD/CPU")
   @Context(USE_ATOMICS=1, SPEC=1)

test/backend/test_asm_gemm.py

@@ -9,7 +9,7 @@ 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")
 
-def run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=None, gpus:int=1) -> None:
+def run_asm_gemm(a_shape, b_shape, dtype=dtypes.bfloat16, a_shard=None, b_shard=None, gpus:int=1) -> None:
   Tensor.manual_seed(0)
   input_dtype = dtypes.bfloat16 if dtype == FP8_DTYPE else dtype
   a_rand = Tensor.randn(a_shape, dtype=dtypes.float).sub(0.5).cast(input_dtype)
@@ -64,31 +64,31 @@ def run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=N
     assert a.grad.allclose(a_ref.grad, atol=grad_atol, rtol=grad_rtol).item(), "grad_a mismatch"
     assert b.grad.allclose(b_ref.grad, atol=grad_atol, rtol=grad_rtol).item(), "grad_b mismatch"
 
-def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=1) -> None:
+def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.bfloat16, gpus:int=1) -> None:
   run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=0, b_shard=None, gpus=gpus)
 
-def verify_asm_gemm_k_sharded(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=8) -> None:
+def verify_asm_gemm_k_sharded(M:int, N:int, K:int, dtype=dtypes.bfloat16, gpus:int=8) -> None:
   run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=1, b_shard=0, gpus=gpus)
 
-def verify_asm_gemm_n_sharded(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
+def verify_asm_gemm_n_sharded(batch:int, M:int, N:int, K:int, dtype=dtypes.bfloat16, gpus:int=2) -> None:
   run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=None, b_shard=1, gpus=gpus)
 
-def verify_asm_gemm_m_sharded(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
+def verify_asm_gemm_m_sharded(M:int, N:int, K:int, dtype=dtypes.bfloat16, gpus:int=2) -> None:
   run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=0, b_shard=None, gpus=gpus)
 
-def verify_asm_gemm_n_sharded_2d(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
+def verify_asm_gemm_n_sharded_2d(M:int, N:int, K:int, dtype=dtypes.bfloat16, gpus:int=2) -> None:
   run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=None, b_shard=1, gpus=gpus)
 
-def verify_asm_gemm_k_sharded_3d(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
+def verify_asm_gemm_k_sharded_3d(batch:int, M:int, N:int, K:int, dtype=dtypes.bfloat16, gpus:int=2) -> None:
   run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=2, b_shard=0, gpus=gpus)
 
 # 128x smaller than usual
 # uses the UOp GEMM, runs on non CDNA4 and CI
-@unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
+@unittest.skipUnless(dtypes.bfloat16 in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
 class TestGemm(unittest.TestCase):
   def setUp(self):
     if is_cdna4(): self.skipTest("shapes are too small for the assembly GEMM")
-  def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 32), N, N, dtype=dtypes.half)
+  def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 32), N, N, dtype=dtypes.bfloat16)
   def test_gemm(self): verify_asm_gemm(1, 64, 32, 112)
   def test_gemm_batched(self): verify_asm_gemm(2, 64, 32, 32)
   @needs_second_gpu
@@ -107,7 +107,7 @@ class TestGemm(unittest.TestCase):
 # uses the smallest size for the cdna assembly gemm
 class TestAsmGEMM(unittest.TestCase):
   def setUp(self):
-    if not is_cdna4():
+    if not is_cdna4() or not has_hipcc():
       self.skipTest("assembly gemm is only for cdna4")
 
   def test_tiny(self): verify_asm_gemm(1, 256, 256, 64)
@@ -145,7 +145,7 @@ class TestGemmLlama(unittest.TestCase):
   dtype = dtypes.bfloat16
 
   def setUp(self):
-    if not is_cdna4() or DEV.interface.startswith("MOCK"):
+    if not is_cdna4() or DEV.interface.startswith("MOCK") or not has_hipcc():
       self.skipTest("very slow on non mi350x")
 
   def test_empty(self): asm_gemm(Tensor.empty(N:=getenv("N", 4096), N, dtype=self.dtype), Tensor.empty(N, N, dtype=self.dtype)).realize()
@@ -380,23 +380,5 @@ class TestHkBf16AtbGemm(unittest.TestCase):
   @needs_second_gpu
   def test_m_sharded(self): run_atb_gemm(256, 512, 256, a_shard=2, b_shard=None, gpus=2)
 
-class TestMagicGu(unittest.TestCase):
-  def test_magicgu_matches_old(self):
-    from extra.gemm.cdna_asm_gemm import _magicgu_mulhi, TILE_M, TILE_N, TILE_K
-    old_iters_args = {64: (67108864, 0), 128: (33554432, 0), 224: (613566757, 2147483656)}
-    old_gemm_shapes = [
-      (8192, 4096, 4096), (8192, 14336, 4096), (8192, 4096, 14336),
-      (8192, 8192, 8192), (4096, 4096, 4096), (4096, 14336, 4096),
-      (4096, 14336, 8192), (4096, 4096, 14336), (14336, 4096, 8192),
-      (4096, 8192, 14336), (4096, 4096, 8192), (4096, 8192, 4096),
-    ]
-    for M, N, K in old_gemm_shapes:
-      iters = K // TILE_K
-      total = (M // TILE_M) * (N // TILE_N) * iters
-      for batch in [1, 2]:
-        magic, shift = _magicgu_mulhi(iters, total * batch)
-        old_magic, old_shift = old_iters_args[iters]
-        self.assertEqual((magic, shift), (old_magic, old_shift), f"mismatch for ({M},{N},{K}) batch={batch} iters={iters}")
-
 if __name__ == "__main__":
   unittest.main()

test/backend/test_const_folding.py

@@ -2,6 +2,7 @@ import unittest, math
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.dtype import DTYPES_DICT
 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
 
@@ -167,7 +168,8 @@ class TestMultiConstFolding(unittest.TestCase):
 
 class TestThreefryConstFolding(unittest.TestCase):
   def test_threefry(self):
-    x = UOp.const(dtypes.uint64, 5).threefry(UOp.const(dtypes.uint64, 10))
+    # THREEFRY(const,const) folds to a const once decomposed
+    x = threefry2x32(UOp.const(dtypes.uint64, 5), UOp.const(dtypes.uint64, 10))
     self.assertIs(x.simplify().op, Ops.CONST)
 
 class TestTautologicalCompare(unittest.TestCase):

test/backend/test_custom_kernel.py

@@ -51,7 +51,7 @@ def flip_contract_kernel(dest:UOp, src:UOp):
   i = UOp.range(dest.shape[0], 0)
   j = UOp.range(dest.shape[1], 1, AxisType.UPCAST)
   vec = src[i, j].contract(j)
-  store = UOp.group(*[dest[i, k].store(vec.gep(3-k)) for k in range(4)])
+  store = UOp.group(*[dest[i, k].store(vec.index(3-k)) for k in range(4)])
   return store.end(i, j).sink(arg=KernelInfo(name=f"flip_contract_{dest.numel()}", opts_to_apply=()))
 
 def slice_sum_kernel(dest:UOp, src:UOp):

test/backend/test_llama_kernels.py

@@ -89,7 +89,7 @@ class TestLocalAmax(unittest.TestCase):
     x = Tensor.arange(16).reshape(4, 4).cast(dtypes.float).clone(devices[0]).realize().shard(devices, axis=0).realize()
     GlobalCounters.reset()
     out = (x * local_abs_max(x)).clone().realize()
-    self.assertEqual(GlobalCounters.kernel_count, 4)
+    self.assertEqual(GlobalCounters.kernel_count, 2)
     self.assertEqual(out.tolist(), [[0., 7., 14., 21.], [28., 35., 42., 49.], [120., 135., 150., 165.], [180., 195., 210., 225.]])
 
 if __name__ == '__main__':

test/backend/test_ops.py

@@ -1449,6 +1449,7 @@ class TestOps(unittest.TestCase):
   def test_small_gemm_eye(self):
     helper_test_op(None, lambda x,y: x.matmul(y), lambda x,y: x@y, vals=[np.eye(8).astype(np.float32), np.eye(8).astype(np.float32)])
   @unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "not precise enough when emulating")
+  @unittest.skipIf(IMAGE>0, "image does math in float32")
   def test_gemm_fp16(self):
     helper_test_op([(64,64), (64,64)], lambda x,y: x.half().matmul(y.half()), atol=5e-3, rtol=5e-3, grad_atol=5e-3, grad_rtol=5e-3)
   def test_gemm(self):
@@ -1601,6 +1602,12 @@ class TestOps(unittest.TestCase):
         helper_test_op(None, lambda x,y: x.isclose(y), vals=[[a], [b]], forward_only=True)
         helper_test_op(None, lambda x,y: x.isclose(y, equal_nan=True), vals=[[a], [b]], forward_only=True)
 
+  def test_isclose_scalar(self):
+    # torch needs a tensor
+    helper_test_op([(3, 4, 5, 6)], lambda x: x.isclose(torch.tensor(1.0)), lambda x: x.isclose(1.0), forward_only=True)
+    helper_test_op(None, lambda x: x.isclose(torch.tensor(1.0)), lambda x: x.isclose(1.0),
+                   vals=[[1.0, 1.0 + 1e-7, 2.0, math.inf, -math.inf, math.nan]], forward_only=True)
+
   def test_mean(self):
     helper_test_op([(3,4,5,6)], lambda x: x.mean())
     helper_test_op([()], lambda x: x.mean())

test/backend/test_uops.py

@@ -239,6 +239,7 @@ class TestAssembly(unittest.TestCase):
     self.assertIn(Ops.SHL, ops)
     self.assertIn(Ops.MUL, ops)
 
+  @unittest.skip("this is a questionable microoptimization i won't enforce")
   def test_mulacc_unrolled(self):
     # test that     acc = acc + a0*b0 + a1*b1 + a2*b2 + a3*b3
     # is not        acc = acc + (a0*b0 + a1*b1 + a2*b2 + a3*b3)

test/external/external_uop_gc.py

@@ -3,7 +3,7 @@ from tinygrad import Tensor, UOp, Device, nn
 from tinygrad.schedule import schedule_cache
 from tinygrad.codegen import to_program, to_program_cache
 from tinygrad.schedule.indexing import apply_movement_op, _apply_reshape
-from tinygrad.codegen.decomp.divandmod import fold_divmod_general
+from tinygrad.uop.divandmod import fold_divmod_general
 from test.test_tiny import TestTiny
 
 def uops_allocated(): return sum([isinstance(x, UOp) for x in gc.get_objects()])

test/null/test_gpudims.py

@@ -94,6 +94,17 @@ class TestGroupedDims(unittest.TestCase):
     assert idxs[2].op is Ops.SPECIAL, f"expected SPECIAL for direct-mapped dim, got {idxs[2].op}"
     assert idxs[3].op is Ops.SPECIAL, f"expected SPECIAL for direct-mapped dim, got {idxs[3].op}"
 
+  def test_grouped_dims_high_rank(self):
+    # 4D collapsed onto 2 axes
+    self._check_grouped_dims("gidx", (4,4,4,4), (16,16), False, [16,16])
+    # 4D untouched
+    self._check_grouped_dims("gidx", (2,3,4,5), None, False, [2,3,4,5])
+    idxs = get_grouped_dims("gidx", (2,3,4,5), None, False)
+    assert all(u.op is Ops.SPECIAL for u in idxs), f"expected all-SPECIAL when untouched, got {[u.op for u in idxs]}"
+    # 5D and 6D collapsed onto 3 axes
+    self._check_grouped_dims("gidx", (2,2,2,2,2), (4,4,4), False, [4,4,2])
+    self._check_grouped_dims("gidx", (2,2,2,2,2,2), (8,8,8), False, [8,4,2])
+
   def test_global_prod_max(self):
     g, l = UOp.range(256, 0, AxisType.GLOBAL), UOp.range(256, 1, AxisType.LOCAL)
     sink = UOp.param(0, dtypes.float.ptr()).index(g + l).store(UOp.const(dtypes.float, 1.0)).end(g, l).sink(arg=KernelInfo())

test/null/test_graph_rewrite.py

@@ -183,33 +183,27 @@ 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(dtypes.float32.vec(4), (1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(apply_rewrite(base_vector.gep(2)).arg, 3.0)
+    self.assertEqual(apply_rewrite(base_vector.index(2)).arg, 3.0)
 
   def test_gep_tuple_extraction(self):
     # GEP on a vector dtype to extract multiple elements as a vector
     base_vector = UOp.const(dtypes.float32.vec(4), (1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(list(apply_rewrite_values(base_vector.gep((2, 3)))), [3.0, 4.0])
+    self.assertEqual(list(apply_rewrite_values(UOp.vectorize(*[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(dtypes.float32.vec(4), (1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(apply_rewrite(const_stack.gep(2)).arg, 3.0)
+    self.assertEqual(apply_rewrite(const_stack.index(2)).arg, 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(dtypes.float32.vec(4), (7.0, 8.0, 9.0, 10.0))
-    self.assertEqual(list(apply_rewrite_values(const_stack.gep((1, 3)))), [8.0, 10.0])
-
-  def test_gep_gep_simplification(self):
-    # Nested GEP simplification on a vector dtype
-    base_vector = UOp.const(dtypes.float32.vec(4), (10.0, 20.0, 30.0, 40.0))
-    gep_inner = base_vector.gep(1)  # Extract 2nd element (20.0)
-    self.assertEqual(apply_rewrite(gep_inner.gep(0)).arg, 20.0)
+    self.assertEqual(list(apply_rewrite_values(UOp.vectorize(*[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(dtypes.float32.vec(4), (5.0, 10.0, 15.0, 20.0))
-    vectorized_uop = UOp(Ops.STACK, dtypes.float32.vec(4), src=(base_vector.gep(0), base_vector.gep(1), base_vector.gep(2), base_vector.gep(3)))
+    vectorized_uop = UOp(Ops.STACK, dtypes.float32.vec(4), 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])
 
 

test/null/test_helpers.py

@@ -1,6 +1,6 @@
 import ctypes, gzip, unittest, timeit, pickle
 from tinygrad import Variable
-from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, get_contraction, count, all_same
+from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, count, all_same
 from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
 from tinygrad.helpers import ceildiv, ansistrip, get_shape
 from tinygrad.tensor import Tensor
@@ -273,75 +273,6 @@ class TestMemoryview(unittest.TestCase):
     mva_us = timeit.timeit(lambda: mv_address(x), number=iters) * 1e6 / iters
     print(f"from_mv vs mv_address: {fmv_us:8.3f} µs vs {mva_us:8.3f} µs")
 
-class TestGetContraction(unittest.TestCase):
-  def test_contraction(self):
-    r = get_contraction((1,2,3,4), (2,3,4))
-    self.assertEqual(r, [[0, 1], [2], [3]])
-
-    r = get_contraction((2,1,3,4), (2,3,4))
-    self.assertEqual(r, [[0], [1, 2], [3]])
-
-    r = get_contraction((1,2,3,1,4), (1,2,3,4))
-    self.assertEqual(r, [[], [0, 1], [2], [3, 4]])
-
-    r = get_contraction((1,2,3,1,4,1,1), (2,3,4))
-    self.assertEqual(r, [[0, 1], [2], [3, 4, 5, 6]])
-
-    r = get_contraction((1,2,3,4), (1,2,3*4))
-    self.assertEqual(r, [[], [0, 1], [2, 3]])
-
-    r = get_contraction((1,2,3,4), (2,1,3,4))
-    self.assertEqual(r, [[0, 1], [], [2], [3]])
-
-    r = get_contraction((1,2,3,4), (1,1,2*3*4,1))
-    self.assertEqual(r, [[], [], [0,1,2,3], []])
-
-    r = get_contraction((2,1,3,4), (1,2,3,4))
-    self.assertEqual(r, [[], [0], [1, 2], [3]])
-
-    r = get_contraction((1,2,3,4), (2*3*4,1,1,1))
-    self.assertEqual(r, [[0, 1, 2, 3], [], [], []])
-
-    r = get_contraction((4,4,4,4), (16,1,16))
-    self.assertEqual(r, [[0, 1], [], [2, 3]])
-
-    r = get_contraction((1,2,3,4,1,1,1), (2,3,4))
-    self.assertEqual(r, [[0, 1], [2], [3, 4, 5, 6]])
-
-    r = get_contraction((1,2,3,4), (1,2,3,4,1))
-    self.assertEqual(r, [[], [0, 1], [2], [3], []])
-
-    r = get_contraction((14,1,384,14,1,1,1,1), (1,14,384,14))
-    self.assertEqual(r, [[], [0], [1,2], [3,4,5,6,7]])
-
-    r = get_contraction((14,1,384,1,14,1,1,1,1), (1,14,384,14))
-    self.assertEqual(r, [[], [0], [1,2], [3,4,5,6,7,8]])
-
-    r = get_contraction((512, 512), (1, 1, 512, 1, 1, 1, 1, 512))
-    self.assertEqual(r, [[], [], [0], [], [], [], [], [1]])
-
-    r = get_contraction((1,2,3,4), (1,2,6,2))
-    self.assertEqual(r, None)
-
-  def test_contraction_ones(self):
-    r = get_contraction((1,), (1,1,1))
-    self.assertEqual(r, [[], [], [0]])
-
-    r = get_contraction((1,1), (1,1,1))
-    self.assertEqual(r, [[], [], [0, 1]])
-
-    r = get_contraction((1,1,1,1), (1,))
-    self.assertEqual(r, [[0,1,2,3]])
-
-    r = get_contraction((1,1,1,1), (1,1))
-    self.assertEqual(r, [[], [0,1,2,3]])
-
-    r = get_contraction((1,1,1,1), (1,1,1))
-    self.assertEqual(r, [[], [], [0,1,2,3]])
-
-    r = get_contraction((1,1,1,1), (1,1,1,1))
-    self.assertEqual(r, [[], [], [], [0,1,2,3]])
-
 class TestGetShape(unittest.TestCase):
   def test_get_shape(self):
     assert get_shape(2) == ()

test/null/test_pattern_matcher.py

@@ -50,7 +50,7 @@ class TestPatternMatcher(unittest.TestCase):
     def fxn(ctx, x):
       ctx.append(True)
       assert len(x.src) == 0
-      return x.replace(src=(UOp(Ops.DEVICE, arg="blah"),))
+      return x.replace(src=(UOp(Ops.NOOP),))
     matcher = PatternMatcher([(UPat(Ops.CONST, src=(), name="x"), fxn)])
     c1 = UOp.const(dtypes.float, 1.0)
     # second rewrite shouldn't match anything

test/null/test_schedule.py

@@ -266,7 +266,7 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.empty(big_enough).realize()
     with Context(SPLIT_REDUCEOP=1):
       out = (x - x.max(keepdim=True)).max()
-      check_schedule(out, 4)
+      check_schedule(out, 3)
 
   def test_example_matmul_contig(self):
     x = Tensor.eye(64).clone().realize()
@@ -355,8 +355,7 @@ class TestSchedule(unittest.TestCase):
     b = Tensor.empty((1, 16)).realize()
     out0 = a.sum() + 2
     out1 = a.sum() + b
-    # check_schedule([out0, out1], 2)
-    check_schedule([out0, out1], 3)
+    check_schedule([out0, out1], 2)
 
   def test_scaled_dot_product_attention_multireduce_fusion(self):
     q = Tensor.empty(32,8,16,8).realize()
@@ -546,8 +545,7 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.empty(3, 4, 5).abs().realize()
     b = Tensor.empty(3, 4, 5).abs().realize()
     out = (a.log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum()+b).abs().log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum().contiguous()
-    # check_schedule(out, 1)
-    check_schedule(out, 2)
+    check_schedule(out, 1)
 
   def test_shrink_pad_safe(self):
     a = Tensor.ones((3, )).contiguous().realize()

test/null/test_simplify_valid_idx.py

@@ -1,6 +1,6 @@
 import unittest, itertools
 
-from tinygrad.codegen.late.devectorizer import indexing_simplify
+from tinygrad.codegen.late.coalese import indexing_simplify
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite
 from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load
@@ -495,7 +495,7 @@ class TestImageSimplification(unittest.TestCase):
 class TestDropTrueGate(unittest.TestCase):
   def test_drop_true_gate_on_index(self):
     # test that INDEX with a constant True valid gets simplified to drop the valid
-    from tinygrad.codegen.late.devectorizer import indexing_simplify
+    from tinygrad.codegen.late.coalese import indexing_simplify
     from tinygrad.uop.ops import graph_rewrite
     from tinygrad.uop.symbolic import sym
     buf = UOp.param(0, dtypes.int.ptr())

test/null/test_tensor_uop_mixin.py

@@ -456,7 +456,6 @@ class TestTensorUOpSVD(unittest.TestCase):
   def test_svd_batched(self):   self._check(_t(2, 2, 2).float())
   def test_svd_nonfull(self):   self._check(_t(3, 2).float(), full_matrices=False)
 
-# UOp.empty / UOp.empty_like are the canonical buffer allocators; Tensor.empty / Tensor.empty_like just forward.
 class TestUOpEmpty(unittest.TestCase):
   def test_empty_dtype_string(self):
     self.assertEqual(UOp.empty((3, 4), dtype="float32").dtype, dtypes.float32)
@@ -475,11 +474,12 @@ class TestUOpEmpty(unittest.TestCase):
       self.assertTrue(u.has_buffer_identity())
 
   def test_empty_direct_singleton_tuple_device(self):
-    # regression: direct UOp.empty with a singleton-tuple device + axis must not trip .multi()'s tuple assert
-    u = UOp.empty((4,), dtype=dtypes.float32, device=("NULL:0",), axis=0)
+    u = UOp.empty((4,), dtype=dtypes.float32, device=("NULL:0",))
     self.assertEqual((u.shape, u.device, u.axis), ((4,), "NULL", None))
 
 class TestTensorUOpCreation(unittest.TestCase):
+  def test_empty(self):
+    self.assertIs(_strip_unique(Tensor.empty(2, 3).uop), _strip_unique(UOp.empty(2, 3)))
   def test_full(self):
     self.assertIs(_strip_unique(Tensor.full((2, 3), 42).uop), _strip_unique(UOp.full((2, 3), 42)))
   def test_full_kwargs(self):

test/null/test_tinyfs.py

@@ -1,24 +1,25 @@
 import unittest
 from tinygrad import Tensor
+from tinygrad.nn.state import fs_store, fs_load
 
 class TestLoadStore(unittest.TestCase):
   def test_load_shape(self):
-    t = Tensor(bytes(16)).fs_load(1024)
+    t = fs_load(Tensor(bytes(16)), 1024)
     assert t.shape == (1024,), t.shape
     t.schedule_linear()
 
   def test_store_shape(self):
-    t = Tensor.zeros(1024).fs_store()
+    t = fs_store(Tensor.zeros(1024))
     assert t.shape == (16,), t.shape
     t.schedule_linear()
 
   def test_load_large_shape(self):
-    t = Tensor(bytes(16)).fs_load(10_000_000)
+    t = fs_load(Tensor(bytes(16)), 10_000_000)
     assert t.shape == (10_000_000,), t.shape
     t.schedule_linear()
 
   def test_store_large_shape(self):
-    t = Tensor.zeros(10_000_000).fs_store()
+    t = fs_store(Tensor.zeros(10_000_000))
     assert t.shape == (16,), t.shape
     t.schedule_linear()
 

test/null/test_uop_graph.py

@@ -2,9 +2,8 @@ import unittest, pytest
 from tinygrad import dtypes, Variable
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import DEBUG, Context
-from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp, AxisType
+from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType
 from tinygrad.uop.symbolic import sym
-from tinygrad.codegen.late.expander import expander
 from test.helpers import to_uops_list
 
 simple_pm = PatternMatcher([
@@ -22,18 +21,11 @@ def const_values(u:UOp):
 class TestGraphRewriteConst(unittest.TestCase):
   def test_gep_const(self):
     v1 = UOp.const(dtypes.int.vec(3), (0,1,2))
-    v2 = v1.gep(1)
+    v2 = v1.index(1)
     ret = graph_rewrite(v2, sym)
     self.assertEqual(ret.dtype, dtypes.int)
     self.assertEqual(ret.arg, 1)
 
-  def test_gep_const_single(self):
-    v1 = UOp.const(dtypes.int.vec(3), 4)
-    v2 = v1.gep(1)
-    ret = graph_rewrite(v2, sym)
-    self.assertEqual(ret.dtype, dtypes.int)
-    self.assertEqual(ret.arg, 4)
-
   def test_add_const(self):
     v1 = UOp.const(dtypes.int, (0,1,2))
     v2 = UOp.const(dtypes.int, (5,6,7))
@@ -262,7 +254,7 @@ class TestUOpGraph(unittest.TestCase):
     idx = UOp.const(dtypes.int, 0)
     ld = d0.load(idx, dtype=dtypes.float.vec(2))
     vec = UOp(Ops.STACK, dtypes.float.vec(2), (ld,))
-    x = UOp(Ops.GEP, dtypes.float, (vec, ), arg=0)
+    x = vec.index(0)
     alu = UOp(Ops.SQRT, dtypes.float, (x, ))
     out = UOp(Ops.STORE, dtypes.void, (d0, idx, alu))
     uops = to_uops_list([out])
@@ -285,27 +277,27 @@ class TestUOpGraph(unittest.TestCase):
 
     # possible
     val = d1.index(idx).load(dtype=dtypes.float.vec(4))
-    xyzw = tuple(val.gep(i) for i in range(4))
+    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.vec(4))
-    wzyx = tuple(val.gep(i) for i in reversed(range(4)))
+    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.vec(2))
-    xy = tuple(val.gep(i) for i in range(2))
+    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.vec(4))
-    xy = tuple(val.gep(i) for i in range(2))
+    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.vec(2))
     val2 = d2.index(idx).load(dtype=dtypes.float.vec(2))
-    xy1 = tuple(val1.gep(i) for i in range(2))
-    xy2 = tuple(val2.gep(i) for i in range(2))
+    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):
@@ -313,7 +305,7 @@ class TestUOpGraph(unittest.TestCase):
       consts = [UOp.const(dtypes.float, float(i)) for i in range(vec_size)]
       vec = UOp(Ops.STACK, dtypes.float.vec(vec_size), tuple(consts))
       with Context(SPEC=0):
-        uops = to_uops_list([UOp(Ops.GEP, dtypes.float, (vec,), (i,)) for i in range(vec_size)])
+        uops = to_uops_list([vec.index(i) for i in range(vec_size)])
         for uop, const in zip(uops, consts):
           self.assertEqual(uop, const)
 
@@ -579,161 +571,6 @@ class TestUOpGraph(unittest.TestCase):
     a = c.after(e)
     self.assertNotIn(r, a.ranges)
 
-@track_rewrites()
-def expander_rewrite(sink): return graph_rewrite(sink, sym + expander)
-
-class TestExpander(unittest.TestCase):
-  def test_expand_add_broadcast(self):
-    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(x for x in range(4))),), ((1,4),))
-    sink = expander_rewrite(e1+3)
-    assert sink.op is Ops.UNROLL and len(const_values(sink.src[0])) == 4
-    self.assertTupleEqual(const_values(sink.src[0]), (3,4,5,6))
-
-  def test_contract_simple(self):
-    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(x for x in range(4))),), ((1,4),))
-    con = UOp(Ops.CONTRACT, dtypes.int.vec(4), (e1,), ((1,4),))
-    sink = expander_rewrite(con)
-    self.assertEqual(sink.op, Ops.STACK)
-    self.assertTupleEqual(const_values(sink), (0,1,2,3))
-
-  def test_contract_axis_1(self):
-    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(16), tuple(x for x in range(16))),), ((1,4),(2,4)))
-    con = UOp(Ops.CONTRACT, dtypes.int.vec(4), (e1,), ((1,4),))
-    sink = expander_rewrite(con)
-    vals = const_values(sink.src[0])
-    assert sink.op is Ops.UNROLL and len(vals) == 16 and sink.arg == ((2,4),)
-    assert sink.src[0].op is Ops.STACK
-    self.assertTupleEqual(vals[0:4], (0,4,8,12))
-    self.assertTupleEqual(vals[12:], (3,7,11,15))
-
-  def test_contract_axis_2(self):
-    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(16), tuple(x for x in range(16))),), ((1,4),(2,4)))
-    con = UOp(Ops.CONTRACT, dtypes.int.vec(4), (e1,), ((2,4),))
-    sink = expander_rewrite(con)
-    vals = const_values(sink.src[0])
-    assert sink.op is Ops.UNROLL and len(vals) == 16 and sink.arg == ((1,4),)
-    assert sink.src[0].op is Ops.STACK
-    self.assertTupleEqual(vals[0:4], (0,1,2,3))
-    self.assertTupleEqual(vals[12:], (12,13,14,15))
-
-  def test_contract_axis_2_big(self):
-    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(16), tuple(x for x in range(16))),), ((1,2),(2,2),(3,2),(4,2)))
-    con = UOp(Ops.CONTRACT, dtypes.int.vec(2), (e1,), ((2,2),))
-    sink = expander_rewrite(con)
-    assert sink.op is Ops.UNROLL and sink.arg == ((1, 2), (3, 2), (4, 2))
-    vals = const_values(sink.src[0])
-    self.assertTupleEqual(vals[0:2], (0,4))
-    self.assertTupleEqual(vals[12:14], (10,14))
-
-  def test_contract_multi_axis(self):
-    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(16), tuple(x for x in range(16))),), ((1,2),(2,2),(3,2),(4,2)))
-    sink = expander_rewrite(UOp(Ops.CONTRACT, dtypes.int.vec(4), (e1,), ((3, 2), (2, 2))))
-    assert sink.op is Ops.UNROLL and sink.arg == ((1, 2), (4, 2))
-    self.assertTupleEqual(const_values(sink.src[0])[0:4], (0, 4, 2, 6))
-    sink = expander_rewrite(UOp(Ops.CONTRACT, dtypes.int.vec(4), (e1,), ((2, 2), (3, 2))))
-    assert sink.op is Ops.UNROLL and sink.arg == ((1, 2), (4, 2))
-    self.assertTupleEqual(const_values(sink.src[0])[0:4], (0, 2, 4, 6))
-
-  def test_contract_mid(self):
-    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(8), tuple(x for x in range(8))),), ((1,2),(2,2),(3,2)))
-    con = UOp(Ops.CONTRACT, dtypes.int.vec(2), (e1,), ((2,2),))
-    sink = expander_rewrite(con)
-    assert sink.op is Ops.UNROLL and sink.arg == ((1,2),(3,2))
-    assert sink.src[0].op is Ops.STACK and len(const_values(sink.src[0])) == 8
-    self.assertTupleEqual(const_values(sink.src[0]), (0,2,1,3,4,6,5,7))
-
-  def test_contract_no_expand(self):
-    e1 = UOp.variable("i", 0, 10, dtype=dtypes.int)
-    con = UOp(Ops.CONTRACT, dtypes.int.vec(2), (e1,), ((2,2),))
-    sink = expander_rewrite(con)
-    assert sink.op is Ops.STACK and len(sink.src) == 2
-    assert sink.src[0] == sink.src[1]
-
-  def test_contract_half_expand(self):
-    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(x for x in range(4))),), ((1,4),))
-    con = UOp(Ops.CONTRACT, dtypes.int.vec(8), (e1,), ((1,4), (2,2)))
-    sink = expander_rewrite(con)
-    vals = const_values(sink)
-    assert sink.op is Ops.STACK and len(vals) == 8
-    assert vals[0] == vals[1]
-    assert vals[0] != vals[2]
-    assert vals[6] == vals[7]
-
-  def test_expand_same_axis(self):
-    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(x for x in range(4))),), ((1,4),))
-    e2 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(4*x for x in range(4))),), ((1,4),))
-    sink = expander_rewrite(e1+e2)
-    self.assertEqual(sink.op, Ops.UNROLL)
-    self.assertEqual(sink.src[0].op, Ops.STACK)
-    self.assertTupleEqual(const_values(sink.src[0]), (0,5,10,15))
-
-  def test_expand_different_axis(self, flip=False):
-    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(4*x for x in range(4))),), ((1,4),))
-    e2 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(x for x in range(4))),), ((2,4),))
-    sink = expander_rewrite((e2+e1) if flip else (e1+e2))
-    vals = const_values(sink.src[0])
-    assert sink.op is Ops.UNROLL and len(vals) == 16
-    assert sink.arg == ((1, 4), (2, 4))
-    self.assertTupleEqual(vals, (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15))
-
-  def test_expand_different_axis_flip(self): self.test_expand_different_axis(True)
-
-  @unittest.skip("no longer supported")
-  def test_reduce_known_axis(self):
-    e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((1,4),))
-    sink = (3*e1).reduce(e1, arg=Ops.ADD)
-    sink = expander_rewrite(sink)
-    assert sink.op is Ops.CONST
-    self.assertEqual(sink.arg, 3*(0+1+2+3))
-
-  @unittest.skip("no longer supported")
-  def test_reduce_const(self):
-    e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((1,4),))
-    sink = UOp.const(dtypes.int, 3).reduce(e1, arg=Ops.ADD)
-    sink = expander_rewrite(sink)
-    assert sink.op is Ops.CONST
-    self.assertEqual(sink.arg, 3*4)
-
-  @unittest.skip("no longer supported")
-  def test_double_expand(self):
-    e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((2,4),))
-    e2 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, 4+x) for x in range(4)), ((2,4),))
-    e = UOp(Ops.UNROLL, dtypes.int, (e1, e2), ((1,2),))
-    sink = expander_rewrite(e)
-    assert sink.op is Ops.UNROLL and len(sink.src) == 8
-    assert sink.arg == ((1, 2), (2, 4))
-    self.assertListEqual([x.arg for x in sink.src], [0,1,2,3,4,5,6,7])
-
-  @unittest.skip("no longer supported")
-  def test_double_expand_reverse(self):
-    e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((1,4),))
-    e2 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, 4+x) for x in range(4)), ((1,4),))
-    e = UOp(Ops.UNROLL, dtypes.int, (e1, e2), ((2,2),))
-    sink = expander_rewrite(e)
-    assert sink.op is Ops.UNROLL and len(sink.src) == 8
-    assert sink.arg == ((1, 4), (2, 2))
-    self.assertListEqual([x.arg for x in sink.src], [0, 4, 1, 5, 2, 6, 3, 7])
-
-  @unittest.skip("no longer supported")
-  def test_double_expand_middle(self):
-    e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((1,2),(3,2)))
-    e2 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, 4+x) for x in range(4)), ((1,2),(3,2)))
-    e = UOp(Ops.UNROLL, dtypes.int, (e1, e2), ((2,2),))
-    sink = expander_rewrite(e)
-    assert sink.op is Ops.UNROLL and len(sink.src) == 8
-    assert sink.arg == ((1, 2), (2, 2), (3, 2))
-    self.assertListEqual([x.arg for x in sink.src], [0, 1, 4, 5, 2, 3, 6, 7])
-
-  # does this need to work?
-  @unittest.expectedFailure
-  @unittest.skip
-  def test_reduce_different_axis(self):
-    e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((1,4),))
-    e2 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((2,4),))
-    sink = e1.reduce(e2, arg=Ops.ADD)
-    sink = expander_rewrite(sink)
-    print(sink)
-
 class TestReduceCollapse(unittest.TestCase):
   def test_multi_range_reduce_add(self):
     """Test that (x + y).reduce(r1, r2) distributes over multiple ranges"""

test/null/test_uop_symbolic.py

@@ -825,6 +825,26 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((x//10)*10 + x%10, 0, 119, "(a*10+(a+b//5)//2*10+(b+a*5)%10)")
     self.helper_test_variable((x//10)*2 + (x//5)%2, 0, 23, "(a*3+b//5)")
 
+  def test_div_mod_recombine_merged_quotient(self):
+    # recombine finds the quotient base//div even when stored merged as base0//(d0*div), including with an offset
+    x = Variable("x", 0, 199)
+    self.helper_test_variable(((x//3)%4)*2 + ((x//12)%5)*8, 0, 38, "x//3%20*2")  # nested-merged quotient
+    self.helper_test_variable(((x//3 + 1)%4) + ((x+3)//12)*4, 1, 67, "x//3+1")   # offset-merged quotient
+
+  def test_div_mod_recombine_negative_div(self):
+    # partial recombine only needs d>0, div can be negative: (x%div) + ((x//div)%d)*div -> x%(div*d)
+    x = Variable("x", 0, 199)
+    self.helper_test_variable(x%(-3) + ((x//(-3))%5)*(-3), -14, 0, "x%-15")
+
+  def test_div_mod_recombine_shifted_quotient(self):
+    # when vmin<0 blocks const reduction on the mod side, the quotient is stored const-shifted: (x-50)//3 -> (x+1)//3 - 17.
+    # recombine only needs a quotient of some b congruent to base mod div, so the shift folds into the result
+    x = Variable("x", 0, 100)
+    y = Variable("y", 0, 99)
+    self.helper_test_variable((x-50)%3 + ((x-50)//3)*3, -50, 50, "(x+-50)")                # shifted literal quotient
+    self.helper_test_variable((x-50)%3 + (((x-50)//3)%5)*3, 0, 14, "((x+-50)%15)")         # shift inside the partial's mod
+    self.helper_test_variable(((y-50)//5)%4 + ((y-50)//20)*4, -10, 9, "(y//5+-10)")        # merged and shifted
+
   def test_div_mod_recombine_in_additive_sum(self):
     x = Variable("x", 0, 31)
     y = Variable("y", 0, 5)

test/null/test_uop_vmin_vmax.py

@@ -319,7 +319,7 @@ class TestVminVmaxVConst(unittest.TestCase):
     d1 = UOp.param(1, dtypes.int.ptr())
     idx = UOp.const(dtypes.int, 0)
     val = UOp(Ops.LOAD, dtypes.int.vec(2), (d1.index(idx).cast(dtypes.int.vec(2).ptr()),))
-    uop = (val // 32).gep(0)
+    uop = (val // 32).index(0)
     self.assertEqual(uop.vmin, -67108864)
     self.assertEqual(uop.vmax, 67108863)
 

test/null/test_uops.py

@@ -318,10 +318,6 @@ class TestUOpStr(unittest.TestCase):
     vec = UOp(Ops.STACK, dtypes.int.vec(4), tuple(UOp.const(dtypes.int, x) for x in range(4)))
     assert str(eval(str(vec))) == str(vec)
 
-  def test_device_arg(self):
-    device = UOp(Ops.DEVICE, arg="CL")
-    assert str(eval(str(device))) == str(device)
-
   def test_reduceop_arg(self):
     sum_uop = Tensor.empty(32, 32).sum().uop
     assert str(eval(str(sum_uop))) == str(sum_uop)

test/null/test_uops_stats.py

@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor
-from tinygrad.helpers import GlobalCounters, DEV
+from tinygrad.helpers import GlobalCounters
 from tinygrad.engine.realize import compile_linear, estimate_uop
 from tinygrad.codegen import to_program
 from tinygrad.renderer import Estimates
@@ -90,7 +90,6 @@ class TestUOpsStatsMatmulHalf(unittest.TestCase):
     expected_ops = N ** 3 * 2
     self.assertEqual(expected_ops, GlobalCounters.global_ops)
 
-  @unittest.skipIf(DEV.arch=="INTEL", "intel gets 524288 != 524352")
   def test_bigger_matmul_half(self): self.test_simple_matmul_half(64)
 
   def test_batched_matmul_half(self, N=16):
@@ -166,19 +165,30 @@ class TestStatsOptimized(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
     cls.ast_gemm = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule_linear().src[-1].src[0]
+    cls.ast_gemm_half = (Tensor.empty(N, N, dtype=dtypes.half) @ Tensor.empty(N, N, dtype=dtypes.half)).schedule_linear().src[-1].src[0]
     cls.ast_reduce = (Tensor.empty(N*N).sum()).schedule_linear().src[-1].src[0]
 
-  def check_gemm(self, p:UOp, extra_flops=0):
+  def check_gemm(self, p:UOp, extra_flops=0, half=False):
     est = p.src[0].arg.estimates
     print(p.arg.name, est.ops, est.mem, est.lds)
     self.assertEqual(est.ops, 2*N*N*N + extra_flops)  # N**3 mulaccs
-    self.assertEqual(est.mem, 3*N*N*4) # 3 NxN mats with floats
+    self.assertEqual(est.mem, 3*N*N*(2 if half else 4)) # 3 NxN mats with floats
 
   def test_gemm(self):
     p = to_program(replace_opts(self.ast_gemm, []), renderer=Device[Device.DEFAULT].renderer)
     self.check_gemm(p)
     self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4 + 4*N*N)
 
+  @unittest.skip("fails locally on AMD")
+  def test_gemm_tc_unroll_half(self):
+    try:
+      p = to_program(replace_opts(self.ast_gemm_half, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)]),
+                      renderer=Device[Device.DEFAULT].renderer)
+    except KernelOptError:
+      raise unittest.SkipTest("no tensor cores")
+    print(p.src[2].arg)
+    self.check_gemm(p, half=True)
+
   def test_gemm_tc_unroll(self):
     try:
       p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)]),

test/unit/test_assign.py

@@ -103,6 +103,16 @@ class TestAssign(unittest.TestCase):
     out = x.item()
     assert out == 1, f"expected 1, got {out}"
 
+  def test_pending_assign_chain_preserves_intermediate_reads(self):
+    x = Tensor([0.0]).contiguous().realize()
+    y0 = x + 0
+    x.assign(x + 1)
+    y1 = x + 0
+    x.assign(x + 1)
+    y2 = x + 0
+    x.assign(x + 1)
+    assert [y0.item(), y1.item(), y2.item(), x.item()] == [0.0, 1.0, 2.0, 3.0]
+
   def test_assign_add_jit(self):
     @TinyJit
     def f(x):
@@ -306,6 +316,16 @@ class TestAssign(unittest.TestCase):
     t.assign(t + 100)
     np.testing.assert_equal(t.numpy(), [[100, 104, 108, 112], [101, 105, 109, 113], [102, 106, 110, 114], [103, 107, 111, 115]])
 
+  def test_assign_corealize_order_independent(self):
+    for order in [lambda x,y: Tensor.realize(x, y), lambda x,y: Tensor.realize(y, x)]:
+      x = Tensor([1.0]).realize()
+      y = x + 10
+      x.assign(x*2)
+      x.assign(x+3)
+      order(x, y)
+      self.assertEqual(y.tolist(), [11.0])
+      self.assertEqual(x.tolist(), [5.0])
+
   def test_assign_contiguous(self):
     b = Tensor.arange(16).reshape(4,4).clone().realize()
     a = (Tensor.arange(16).reshape(4,4).clone().realize() + 1)
@@ -853,6 +873,28 @@ class TestAssignOrdering(unittest.TestCase):
       b_np *= 0.9
     np.testing.assert_allclose(param.item(), p_np, atol=1e-5)
 
+  def test_war_reader_already_depends_on_write(self):
+    x = Tensor([1.0]).contiguous().realize()
+    y = Tensor([2.0]).contiguous().realize()
+    x_expr = x + 10
+    x.assign(x * 2)
+    y.assign(y + x)
+    z = y + x_expr
+    Tensor.realize(x, y, z)
+    # TODO: z should be 15: x_expr means 11 (x captured at build time), but the read is fused past the assign and
+    # sees the new bytes. once stale readers are scheduled before the overwrite, update this to 15
+    np.testing.assert_allclose([x.item(), y.item(), z.item()], [2.0, 4.0, 16.0])
+
+  def test_war_multi_read_then_assign(self):
+    devices = ("CPU:0", "CPU:1")
+    for realize_reader_first in (False, True):
+      buf = Tensor([1., 2., 3., 4.], device="CPU").contiguous().realize().shard(devices, 0).realize()
+      stale = buf.to("CPU")
+      buf.assign(Tensor.full(buf.shape, 10.0, device="CPU").shard(devices, 0).contiguous().realize())
+      Tensor.realize(stale, buf) if realize_reader_first else Tensor.realize(buf, stale)
+      np.testing.assert_equal(stale.numpy(), [1., 2., 3., 4.])
+      np.testing.assert_equal(buf.numpy(), [10., 10., 10., 10.])
+
   def test_multiple_slice_assigns_then_read(self):
     """Multiple non-overlapping slice assigns then read."""
     buf = Tensor.zeros(4).contiguous().realize()

test/unit/test_function.py

@@ -246,6 +246,24 @@ class TestFunction(unittest.TestCase):
     r0 = f(buf, x, v.bind(0)).numpy()
     np.testing.assert_equal(r0, [[1.,0.,0.,0.,0.,0.,0.,0.], [2.,0.,0.,0.,0.,0.,0.,0.]])
 
+  def test_single_after_store_precompile(self):
+    """precompiled AFTER(buf, STORE(view, data)) should return buf after the store."""
+    @function(precompile=True)
+    def f(buf:Tensor, x:Tensor, start_pos:int|UOp) -> Tensor:
+      slice_uop = buf[:, start_pos:start_pos+1].uop
+      assigned = Tensor(buf.uop.after(slice_uop.store(x.uop)))
+      return assigned
+
+    x = Tensor([[1.], [2.]]).realize()
+    v = UOp.variable("sp", 0, 7)
+    for sp in (0, 2):
+      with self.subTest(sp=sp):
+        buf = Tensor.zeros(2, 8).clone().realize()
+        expected = np.zeros((2, 8), dtype=np.float32)
+        expected[:, sp] = [1., 2.]
+        np.testing.assert_equal(f(buf, x, v.bind(sp)).numpy(), expected)
+        np.testing.assert_equal(buf.numpy(), expected)
+
   @unittest.expectedFailure
   def test_assign_slice(self):
     @function

test/unit/test_randomness.py

@@ -53,6 +53,12 @@ def equal_distribution(tiny_func, torch_func=None, numpy_func=None, shape=(40, 4
 def normal_test(func, shape=(20, 45), alpha=0.05): return equal_distribution(func, numpy_func=lambda x: np.random.randn(*x), shape=shape, alpha=alpha)
 
 class TestRandomness(unittest.TestCase):
+  def test_three_lazy_rands_realized_one_at_a_time_are_distinct(self):
+    Tensor.manual_seed(123)
+    r1, r2, r3 = [Tensor.rand(4) for _ in range(3)]
+    self.assertNotEqual(r1.tolist(), r2.tolist())
+    self.assertNotEqual(r2.tolist(), r3.tolist())
+
   def test_randn(self):
     self.assertEqual(Tensor.randn(3,3,dtype=dtypes.half).dtype, dtypes.half)
     self.assertTrue(normal_test(Tensor.randn))

test/unit/test_tinyfs.py

@@ -1,6 +1,8 @@
 import json, math, os, socketserver, threading, unittest
 import numpy as np
 from tinygrad import Tensor, dtypes
+from tinygrad.helpers import CHUNK_SIZE
+from tinygrad.nn.state import fs_store, fs_load
 from extra.tinyfs.fetch_file import hash_file, _python_hash_1mb
 
 _chunks: dict[bytes, bytes] = {}
@@ -14,8 +16,8 @@ class _Handler(socketserver.StreamRequestHandler):
       elif cmd.startswith("STORE_IN"):
         data = self.rfile.read(int(cmd.split()[1]))
         hashes = bytearray()
-        for i in range(math.ceil(len(data) / Tensor.CHUNK_SIZE)):
-          chunk = data[i*Tensor.CHUNK_SIZE:(i+1)*Tensor.CHUNK_SIZE].ljust(Tensor.CHUNK_SIZE, b'\0')
+        for i in range(math.ceil(len(data) / CHUNK_SIZE)):
+          chunk = data[i*CHUNK_SIZE:(i+1)*CHUNK_SIZE].ljust(CHUNK_SIZE, b'\0')
           h = _python_hash_1mb(chunk)
           _chunks[h] = chunk
           hashes.extend(h)
@@ -46,35 +48,35 @@ class TestTinyFS(unittest.TestCase):
     cls._server.server_close()
 
   def test_store(self):
-    h = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
+    h = fs_store(Tensor([1.0, 2.0, 3.0, 4.0])).realize()
     self.assertEqual(h.shape, (16,))
     self.assertEqual(h.dtype, dtypes.uint8)
 
   def test_store_deterministic(self):
-    a = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
-    b = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
+    a = fs_store(Tensor([1.0, 2.0, 3.0, 4.0])).realize()
+    b = fs_store(Tensor([1.0, 2.0, 3.0, 4.0])).realize()
     np.testing.assert_array_equal(a.numpy(), b.numpy())
 
   def test_store_different_data(self):
-    a = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
-    b = Tensor([5.0, 6.0, 7.0, 8.0]).fs_store().realize()
+    a = fs_store(Tensor([1.0, 2.0, 3.0, 4.0])).realize()
+    b = fs_store(Tensor([5.0, 6.0, 7.0, 8.0])).realize()
     self.assertNotEqual(a.tolist(), b.tolist())
 
   def test_roundtrip_uint8(self):
     arr = np.arange(256, dtype=np.uint8)
-    loaded = Tensor(arr).fs_store().realize().fs_load(len(arr)).to("CPU")
+    loaded = fs_load(fs_store(Tensor(arr)).realize(), len(arr)).to("CPU")
     np.testing.assert_array_equal(loaded.numpy(), arr)
 
   def test_roundtrip_multichunk_uint8(self):
-    arr = np.random.default_rng(42).integers(0, 256, size=Tensor.CHUNK_SIZE + 1024, dtype=np.uint8)
-    loaded = Tensor(arr).fs_store().realize().fs_load(len(arr)).to("CPU")
+    arr = np.random.default_rng(42).integers(0, 256, size=CHUNK_SIZE + 1024, dtype=np.uint8)
+    loaded = fs_load(fs_store(Tensor(arr)).realize(), len(arr)).to("CPU")
     np.testing.assert_array_equal(loaded.numpy(), arr)
 
   def test_hash_matches_python_impl(self):
     arr = np.arange(256, dtype=np.uint8)
-    h = Tensor(arr).fs_store().realize()
+    h = fs_store(Tensor(arr)).realize()
     # the hash from fs_store should match the pure-Python hash_file reference
-    padded = arr.tobytes().ljust(Tensor.CHUNK_SIZE, b'\0')
+    padded = arr.tobytes().ljust(CHUNK_SIZE, b'\0')
     self.assertEqual(h.data().tobytes(), hash_file(padded))
 
 if __name__ == "__main__":

tinygrad/callify.py

@@ -16,28 +16,24 @@ def tag_uop(ctx:AllocCtx, x:UOp):
   ctx.uop_list.append(x)
   return x.replace(tag=(len(ctx.uop_list)-1,))
 
+def disk_like(u:UOp): return isinstance(u.device, str) and u.device.startswith(("DISK", "TINYFS"))
+
 def disk_copy_is_buffer(ctx:AllocCtx, u:UOp):
   # copies to disk are replaced with the disk buffer
-  to_disk = isinstance(u.device, str) and u.device.startswith(("DISK", "TINYFS"))
-  if to_disk: ctx.buffer_map[u] = u.empty_like()
+  if disk_like(u) and u.tag is None:
+    ctx.buffer_map[u] = u.empty_like()
+    return u.rtag(())
   # all copies from disk/numpy are realized into a real buffer
-  from_creation = isinstance(u.src[0].device, str) and any(u.src[0].device.startswith(x) for x in ["NPY", "DISK", "PYTHON", "TINYFS"])
+  from_creation = isinstance(u.src[0].device, str) and u.src[0].device.startswith(("NPY", "DISK", "PYTHON", "TINYFS"))
   if from_creation: return tag_uop(ctx, u)
 
-def apply_after(ctx:AllocCtx, u:UOp):
-  base = u.src[0]
-  while base.op is Ops.AFTER: base = base.src[0]
-  ctx.buffer_map[u] = base
-
-# CONTIGUOUS and AFTER+STORE + parents are the only nodes that get updated
+# CONTIGUOUS and AFTER + parents are the only nodes that get updated
 add_tags = PatternMatcher([
   (UPat(Ops.COPY, name="u"), disk_copy_is_buffer),
   # no tag on copies that are assigned via STORE+AFTER — merge COPY tag into AFTER
   (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(name="dest"), UPat(Ops.COPY, name="c")))), name="a"),
    lambda a,c,dest: a.replace(src=(a.src[0], a.src[1].replace(src=(dest, c.rtag(())))), tag=a.tag+c.tag) if a.tag and c.tag else None),
-  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE)), name="x"), tag_uop),
-  (UPat(Ops.AFTER, name="u"), apply_after),
-  (UPat(Ops.CONTIGUOUS, name="x"), tag_uop),
+  (UPat((Ops.CONTIGUOUS, Ops.AFTER), name="x"), tag_uop),
   (UPat(GroupOp.All, name="x"), lambda ctx,x: tag_uop(ctx,x) if x in ctx.bases else None),
 ])
 
@@ -47,13 +43,13 @@ def replace_contig_with_store_after(u:UOp):
   # if size is 0, remove the contig
   if 0 in u.shape: return u.src[0]
   # no real contig for DISK/TINYFS tensors, they are left alone
-  if isinstance(u.device, str) and u.device.startswith(("DISK", "TINYFS")): return u.rtag(None)
+  if disk_like(u): return u.rtag(None)
   buf = u.empty_like()
   return buf.after(buf.store(u.src[0])).rtag(u.tag)
 
 def replace_store_after_with_contig(u:UOp, src:UOp):
   assigned_to = u
-  while assigned_to.op in {Ops.BITCAST, Ops.AFTER}: assigned_to = assigned_to.src[0].base
+  while assigned_to.op in {Ops.BITCAST, Ops.AFTER, Ops.MULTI}: assigned_to = assigned_to.src[0].base
   if assigned_to.op is not Ops.BUFFER: return src.contiguous(tag=u.tag)
 
 def _make_buffer_view(src:UOp) -> UOp|None:
@@ -70,7 +66,7 @@ def _make_buffer_view(src:UOp) -> UOp|None:
 def contiguous_mops_to_view(c:UOp, src:UOp):
   """CONTIGUOUS(MOPS(BUFFER)) → CONTIGUOUS(SLICE) when movement ops collapse to a contiguous range."""
   buf = src.base
-  if buf.op not in {Ops.BUFFER, Ops.SLICE}: return None
+  if buf.op not in {Ops.BUFFER, Ops.SLICE, Ops.MULTI}: return None
   if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.SLICE}: return None
 
   # no symbolic shape
@@ -78,14 +74,11 @@ def contiguous_mops_to_view(c:UOp, src:UOp):
 
   # check if view is supported
   from tinygrad.device import Device
-  if isinstance(c.device, str):
-    if not hasattr(Device[c.device].allocator, "_offset"): return None
-  elif not all(hasattr(Device[d].allocator, "_offset") for d in c.device): return None
+  devs = (c.device,) if isinstance(c.device, str) else c.device
+  if not all(hasattr(Device[d].allocator, "_offset") for d in devs): return None
 
-  x = src
-  while x.op in GroupOp.Movement: x = x.src[0]
   # NOTE: this contiguous is removed because this SLICE/RESHAPE has_buffer_identity
-  if x.op is not Ops.MULTI and (view := _make_buffer_view(src)) is not None:
+  if buf.op is not Ops.MULTI and (view := _make_buffer_view(src)) is not None:
     return view.contiguous(tag=c.tag)
 
   # for MULTI tensors, use multi_pm to resolve per-shard movement ops, then create SLICE on the resolved result
@@ -128,7 +121,7 @@ def transform_precompiled_call(c:UOp) -> UOp|None:
       subs[s] = placed
       items.append(s.after(*after_deps) if after_deps else s)
     else:
-      items.append(t.after(t.store(s), *after_deps))
+      items.append(t.after(t.store(s.after(*after_deps))))
   fxn = UOp.sink(*(x.substitute(subs) for x in items))
 
   # body switches from TUPLE to SINK, so the node becomes an opaque CALL (not FUNCTION)
@@ -205,9 +198,9 @@ pm_replace_buf = PatternMatcher([
 def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
   if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Tensor Graph")
   # uop list is a list in the original_sink graph and we can map to the tags later
-  # here we build buffer map
-  dont_realize = {Ops.CONST, Ops.BUFFER, Ops.BIND, Ops.AFTER}
-  ctx = AllocCtx(bases=set([x.multibase for x in big_sink.src if x.base.op not in dont_realize and x.base.addrspace is not AddrSpace.ALU]))
+  # same predicate as Tensor.realize
+  ctx = AllocCtx(bases={base for x in big_sink.src if (base:=x.base).device is not None and not base.has_buffer_identity()
+                        and base.op is not Ops.AFTER and base.addrspace is not AddrSpace.ALU})
 
   # this rewrite is "read-only", it adds simple things to buffer_map and may sink things on big_sink, bottom_up
   # this is the only one where we have to be careful to not break the tensor graph
@@ -216,8 +209,9 @@ def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
   # here we can break the tensor graph. this is the only place you need to maintain numbered tags
   big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, name="early transform tensor graph")
 
-  # here we construct the final buffer_map. this is everything that will go into the tensor map
+  # here we construct the final buffer_map: as-built nodes -> their final storage. values are never keys
   graph_rewrite(big_sink, pm_finalize_call, ctx=ctx, name="finalize call")
   ret = graph_rewrite(UOp.sink(*ctx.assigns), pm_replace_buf, ctx=ctx, bottom_up=True, name="replace bufs").call(*ctx.replacements)
+  assert not any(x in ctx.buffer_map for x in ctx.buffer_map.values())
   if VIZ: graph_rewrite(ret, PatternMatcher([]), name="View Call")
   return ret, ctx.buffer_map

tinygrad/codegen/__init__.py

@@ -1,32 +1,37 @@
-from dataclasses import replace
-import itertools
+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, TracingKey, Context, panic
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, GroupOp
+from tinygrad.uop.ops import AxisType
 from tinygrad.uop.render import pyrender
 from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
 from tinygrad.renderer import Renderer, Estimates
 from tinygrad.renderer.isa import ISARenderer, IselContext, PreRegAllocContext
-from tinygrad.dtype import dtypes, PtrDType, ImageDType
+from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
 
 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
-from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid
+from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid
 from tinygrad.codegen.decomp.dtype import pm_dtype_decomps
 from tinygrad.codegen.decomp.op import get_late_rewrite_patterns, get_simplifying_rewrite_patterns
 from tinygrad.codegen.decomp.transcendental import get_transcendental_patterns
-from tinygrad.codegen.late.expander import expander, pm_pre_expander, pm_group_for_reduce
-from tinygrad.codegen.late.devectorizer import load_store_folding, indexing_simplify, devectorize_buf_and_index, devectorize_alu, pm_reduce, \
-  ReduceContext, pm_render, pm_add_loads
+from tinygrad.codegen.late.coalese import indexing_simplify
 from tinygrad.codegen.opt.postrange import apply_opts
 from tinygrad.codegen.late.gater import pm_move_gates_from_index
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse
-from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops, pm_syntactic_sugar, pm_store_ranges
+from tinygrad.schedule.rangeify import pm_mops, pm_syntactic_sugar, pm_store_ranges, mop_cleanup
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite
 from tinygrad.codegen.late.coalese import memory_coalesing, pm_simplify_add_image
+from tinygrad.helpers import all_same, flatten, argsort, partition
+from tinygrad.uop.ops import _align_left, _broadcast_shape, identity_element
+from tinygrad.schedule.rangeify import BufferizeOpts
 
 pm_remove_vec_dtypes = PatternMatcher([
+  # CONST must be stacked CONST
+  (UPat(Ops.CONST, name='c'),
+   lambda c: UOp(Ops.STACK, c.dtype, (UOp.const(c.dtype.scalar(), c.arg),)*c.dtype.vcount) if c.dtype.vcount > 1 else None),
   # rewrite PARAM to non pointer
   (UPat((Ops.PARAM, Ops.BUFFER), name="buf"), lambda buf:
    buf.replace(dtype=buf.dtype.base, src=(UOp.const(dtypes.int, buf.ptrdtype.size),)) \
@@ -34,8 +39,6 @@ pm_remove_vec_dtypes = PatternMatcher([
   # remove all vec dtypes
   (UPat(GroupOp.All-{Ops.PARAM, Ops.BUFFER}, name="x"),
    lambda x: x.replace(dtype=x.dtype.base.scalar().base)),
-  # rewrite GEP to INDEX
-  (UPat(Ops.GEP, name="x"), lambda x: x.replace(op=Ops.INDEX, src=x.src+(UOp.const(dtypes.int, x.arg if len(x.arg) > 1 else x.arg[0]),), arg=None)),
 ])+pm_clean_up_group_sink
 
 def do_number_param(ctx:list[int], x:UOp):
@@ -51,6 +54,227 @@ pm_no_weakints = PatternMatcher([
   (UPat(GroupOp.All, dtype=dtypes.weakint, name="x"), lambda x: x.replace(dtype=dtypes.int))
 ])
 
+def build_range_map(sink:UOp) -> dict[int, int]:
+  ctx: dict[int, int] = {}
+  for x in sink.toposort():
+    if x.op is Ops.RANGE and x.arg[1] in {AxisType.UNROLL, AxisType.UPCAST}:
+      ctx[x.arg[0]] = len(ctx)
+  return ctx
+
+def expand_reduce(r:UOp):
+  range_srcs = []
+  new_axes = []
+  for u in r.src[1:]:
+    if u.op == Ops.RANGE:
+      range_srcs.append(u)
+    else:
+      for i,s in enumerate(u.shape):
+        if s > 1: new_axes.append(i)
+  if len(new_axes) == 0: return None
+  assert r.arg[1] == ()
+  # move to the front
+  out_shape = tuple([1 if i in new_axes else s for i,s in enumerate(r.src[0].shape)])
+  return r.src[0].reduce(*range_srcs, arg=(r.arg[0], tuple(new_axes))).reshape(out_shape)
+
+def do_contract(ctx:dict[int, int], u:UOp):
+  # the context is a mapping from range number (in contract) to axis number
+  permute_tail = [ctx[rn] for rn,_ in u.arg]
+  permute_head = [i for i in range(len(u.src[0].shape)) if i not in permute_tail]
+  out = u.src[0].permute(permute_head+permute_tail)
+  return out.reshape(*out.shape[:len(permute_head)], -1)
+
+def do_unroll(ctx:dict[int, int], u:UOp):
+  # this is the opposite of contract
+  permute_tail = [ctx[rn] for rn,_ in u.arg]
+  out = u.src[0].reshape(*u.src[0].shape[:-1], *[nm for _,nm in u.arg])
+  permute_head = [i for i in range(len(out.shape)) if i not in permute_tail]
+  return out.permute(argsort(permute_head+permute_tail))
+
+expander2 = PatternMatcher([
+  (UPat(Ops.REDUCE, name="r"), expand_reduce),
+  (UPat(Ops.RANGE, name="r"),
+   lambda ctx, r: UOp.const(r.dtype, tuple(range(r.vmax+1))) \
+    .reshape(tuple([r.vmax+1 if i == ctx[r.arg[0]] else 1 for i in range(len(ctx))])) if r.arg[0] in ctx else None),
+  (UPat(Ops.CONTRACT, name="u"), do_contract),
+  (UPat(Ops.UNROLL, name="u"), do_unroll),
+])+pm_flatten_range+mop_cleanup
+
+def broadcast_binary(x:UOp):
+  shapes = [u._shape for u in x.src]
+  if any(s is None for s in shapes) or all_same(shapes): return None
+  shaped_aligned = _align_left(*shapes)
+  broadcasted = _broadcast_shape(*shapes)
+  src_reshaped = [u.reshape(shp).expand(broadcasted) for u,shp in zip(x.src, shaped_aligned)]
+  return x.replace(src=tuple(src_reshaped))
+
+def broadcast_and_devec_wmma(b:UOp):
+  shapes = [u.shape[:-1] for u in b.src]
+  if all_same(shapes): return None
+  shaped_aligned = _align_left(*shapes)
+  broadcasted = _broadcast_shape(*shapes)
+  src_reshaped = [u.reshape(shp+(u.shape[-1],)).expand(broadcasted+(u.shape[-1],))
+                  for u,shp in zip(b.src, shaped_aligned)]
+  src = []
+  for idx in itertools.product(*[range(i) for i in b.shape[:-1]]):
+    idx_c = [UOp.const(dtypes.weakint, i) for i in idx]
+    src.append(b.replace(src=tuple([x.index(*idx_c) for x in src_reshaped])))
+  return UOp.vectorize(*src).reshape(b.shape)
+
+pm_wmma_add = PatternMatcher([
+  (UPat(Ops.WMMA, name="wmma") + UPat.var("add"),
+   lambda add, wmma: UOp(wmma.op, wmma.dtype, (wmma.src[0], wmma.src[1], wmma.src[2]+add), wmma.arg)),
+  # push permute/reshape to the other side of the add
+  (UPat(Ops.PERMUTE, src=(UPat(Ops.WMMA, name="wmma"),), name="permute") + UPat.var("add"),
+    lambda wmma,permute,add: (wmma + add.permute(argsort(permute.arg))).permute(permute.arg)),
+  (UPat(Ops.PERMUTE, src=(UPat(Ops.RESHAPE, src=(UPat(Ops.WMMA, name="wmma"), UPat()), name="reshape"),), name="permute") + UPat.var("add"),
+    lambda wmma,reshape,permute,add: (wmma + add.permute(argsort(permute.arg)).reshape(wmma.shape)).reshape(reshape.shape).permute(permute.arg)),
+])
+
+unbroadcast = pm_wmma_add+PatternMatcher([
+  (UPat(GroupOp.Binary|GroupOp.Ternary|{Ops.STORE}, name="x"), broadcast_binary),
+  (UPat(Ops.WMMA, name="b"), broadcast_and_devec_wmma),
+])
+
+def do_devectorize(b:UOp):
+  if b.shape == (): return None
+  # broadcasting needs to be already unpacked
+  if not all_same([x.shape for x in b.src]): return None
+  src = []
+  for idx in itertools.product(*[range(x) for x in b.shape]):
+    idx_c = [UOp.const(dtypes.weakint, i) for i in idx]
+    src.append(b.replace(src=tuple([x.index(*idx_c) for x in b.src])))
+  return UOp.vectorize(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src)
+
+def do_stack_wmma(u:UOp):
+  if all(x.op in (Ops.STACK, Ops.WMMA) for x in u.src): return None
+  assert len(u.shape) == 1
+  src = []
+  for b in u.src:
+    if b.op != Ops.STACK:
+      src.append(UOp._stack(*[b.index(UOp.const(dtypes.weakint, i)) for i in range(b.max_numel())]))
+    else:
+      src.append(b)
+  return u.replace(src=tuple(src))
+
+ew_devectorizer = PatternMatcher([
+  # unpack broadcasting
+  (UPat(GroupOp.Elementwise, name="b"), do_devectorize),
+])
+
+devectorizer2 = pm_mops+PatternMatcher([
+  # unpack broadcasting
+  (UPat(GroupOp.Elementwise|{Ops.LOAD,Ops.STORE}, name="b"), do_devectorize),
+  # const INDEX into STACK is src (this is symbolic)
+  (UPat(Ops.INDEX, src=(UPat(Ops.STACK, name="a"), UPat.cvar("i")), name="idx", allow_any_len=True),
+   lambda a,i,idx: a.src[i.arg].index(*idx.src[2:])),
+  # INDEX without src is nothing
+  (UPat(Ops.INDEX, src=(UPat.var('x'),)), lambda x: x),
+  # 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"))),
+   lambda b,s: UOp.vectorize(*[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)),
+  # RESHAPE a void is removed (hack for AFTER)
+  (UPat(Ops.RESHAPE, dtype=dtypes.void, name="x"), lambda x: x.src[0]),
+  # reshape of a single element shaped value to scalar is an index
+  (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(UOp.const(dtypes.weakint, 0)) if x.marg == () and x.src[0].shape == (1,) else None),
+  # RESHAPE+EXPAND -> STACK
+  (UPat(Ops.EXPAND, src=(UPat(Ops.RESHAPE, src=(UPat.var("x"), UPat())), UPat()), name="out"),
+   lambda x,out: UOp.vectorize(*([x]*out.max_numel())) if out.shape == (out.max_numel(),) else None),
+  # INDEX on INDEX is INDEX
+  (UPat(Ops.INDEX, src=(UPat(Ops.INDEX, name="idx1", allow_any_len=True),), allow_any_len=True, name="idx2"),
+   lambda idx1, idx2: idx1.src[0].index(*idx1.src[1:], *idx2.src[1:])),
+])
+
+def fix_group_for_reduce(x:UOp):
+  reduce_gfr, reduce_r = partition(x.src[1:], lambda u: u.op is Ops.RANGE and u.arg[1] == AxisType.GROUP_REDUCE)
+  if len(reduce_gfr) == 0: return None
+
+  # NOTE: if there's other locals here, we need them in the buffer too
+  upstream_locals = [u for u in x.toposort() if u.op is Ops.RANGE and u.arg[1] == AxisType.LOCAL]
+
+  # do only the non grouped reduces early
+  ret = x.replace(src=(x.src[0],)+tuple(reduce_r))
+  reduce_loop = [x.replace(arg=(x.arg[0]+100, AxisType.REDUCE)) for x in reduce_gfr]
+  buf = ret.bufferize(*upstream_locals, *reduce_gfr, arg=BufferizeOpts(reduce_gfr[0].arg[0], AddrSpace.LOCAL)).index(*upstream_locals, *reduce_loop)
+
+  # do the final reduce (if/barrier are added in gpudims step)
+  # NOTE: we remove all horizontal reduces here, they remain in the first reduce
+  return buf.reduce(*reduce_loop, arg=(x.arg[0], ()))
+
+pm_group_for_reduce = PatternMatcher([
+  # fix group for reduce
+  (UPat(Ops.REDUCE, name="x"), fix_group_for_reduce),
+])
+
+@dataclass
+class ReduceContext:
+  acc_num: int = 0
+
+def merge_reduce_ends(sink:UOp):
+  # merge ENDs that share the same range and nesting context (only those created by reduce_to_acc)
+  # ENDs at different nesting depths get cloned RANGEs so each RANGE maps to one END
+  range_to_ends: dict[tuple[UOp, ...], list[UOp]] = {}
+  for u in sink.backward_slice:
+    if u.op is Ops.END and u.tag == "mergeable": range_to_ends.setdefault(u.src[1:], []).append(u)
+  subs: dict[UOp, UOp] = {}
+  next_axis = max((u.arg[0] for u in sink.backward_slice if u.op is Ops.RANGE), default=-1) + 1
+  for r, ends in range_to_ends.items():
+    if len(ends) <= 1: continue
+    by_ctx: dict[frozenset[UOp], list[UOp]] = {}
+    for e in ends: by_ctx.setdefault(frozenset(e.ranges), []).append(e)
+    for i, group in enumerate(by_ctx.values()):
+      tr = r if i == 0 else tuple(rr.replace(arg=(next_axis + j, *rr.arg[1:])) for j, rr in enumerate(r))
+      if i > 0: next_axis += len(r)
+      mapped = [e.substitute(dict(zip(r, tr))) if i > 0 else e for e in group]
+      merged = mapped[0] if len(mapped) == 1 else UOp.group(*(e.src[0] for e in mapped)).end(*tr)
+      for e in group: subs[e] = merged
+  return sink.substitute(subs) if subs else None
+
+def reduce_ranges_to_acc(ctx:ReduceContext, r:UOp):
+  # TODO: remove this is_ptr when placeholder isn't ptr
+  acc = UOp.placeholder_like(r, ctx.acc_num, AddrSpace.REG, is_ptr=False)
+  ctx.acc_num += 1
+  topo = r.src[0].toposort()
+  ended_ranges = flatten([x.ended_ranges for x in topo if x.op is Ops.END])
+  input_ranges = tuple(x for x in topo if x.op is Ops.RANGE and x not in r.src[1:] and x not in ended_ranges)
+  acc_init = acc.after(*input_ranges).store(identity_element(r.arg[0], r.dtype.scalar()))
+  acc_initted = acc.after(acc_init, *r.src[1:])
+  inp = r.src[0].reduce(arg=r.arg) if r.arg[1] else r.src[0]
+  acc_out = acc_initted.store(acc_initted.alu(r.arg[0], inp)).end(*r.src[1:]).rtag("mergeable")
+  return acc.after(acc_out)
+
+def expand_horizontal_reduce(r:UOp):
+  permute = [i for i in range(len(r.src[0].shape)) if i in r.arg[1]] + [i for i in range(len(r.src[0].shape)) if i not in r.arg[1]]
+  inp = r.src[0].permute(permute)
+  vals = [inp.index(*idx) for idx in itertools.product(*[range(inp.max_shape[a]) for a in range(len(r.arg[1]))])]
+  return functools.reduce(lambda x,y: x.alu(r.arg[0], y), vals)
+
+pm_reduce_local = pm_wmma_add+PatternMatcher([
+  (UPat(Ops.REDUCE, src=(UPat(), UPat()), allow_any_len=True, name="r"), reduce_ranges_to_acc),
+  (UPat(Ops.REDUCE, src=(UPat(),), name="r"), expand_horizontal_reduce),
+  (UPat(Ops.SINK, name="sink"), merge_reduce_ends),
+])+pm_clean_up_group_sink
+
+def maybe_load(u:UOp): return u.load() if u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL, AddrSpace.REG) else u
+pm_move_regs = PatternMatcher([
+  # 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:])),
+])
+
+def add_local_buffer(ctx, x:UOp):
+  # TODO: remove this is_ptr when placeholder isn't ptr
+  buf = UOp.placeholder(x.max_shape, x.dtype, slot=next(ctx), addrspace=x.arg.addrspace, is_ptr=False)
+  return buf.after(buf.index(*x.src[1:]).store(x.src[0]).end(*x.src[1:]).barrier())
+
+pm_add_local_buffers = PatternMatcher([
+  (UPat(Ops.STAGE, name="x"), add_local_buffer),
+])+pm_mops
+
 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))
@@ -76,40 +300,45 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
     # do postrange optimization, BEAM or hand_coded_optimizations
     sink = apply_opts(sink, ren, beam=ast.arg.beam)
 
+  # this is new style (TODO: this should all be removed)
+  sink = graph_rewrite(sink, pm_remove_vec_dtypes, name="transform to new style")
+
   # ** expander (expand_rewrite) **
   sink = graph_rewrite(sink, sym+pm_move_where_on_load+pm_flatten_range, name="postopt symbolic")
 
   # expand
-  sink = graph_rewrite(sink, sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander")
+  sink = graph_rewrite(sink, expander2, ctx=build_range_map(sink), name="expander")
+  sink = graph_rewrite(sink, pm_group_for_reduce, name="group for reduce")
 
   # add locals
-  sink = graph_rewrite(sink, pm_add_buffers_local+rangeify_codegen, ctx=itertools.count(0), name="add local buffers")
+  sink = graph_rewrite(sink, pm_add_local_buffers, ctx=itertools.count(0), name="add local buffers")
 
   # ** devectorizer (full_graph_rewrite) **
   # remove reduce
-  sink = graph_rewrite(sink, pm_reduce+gep_pushing, ctx=ReduceContext(), name="remove_reduce")
+  sink = graph_rewrite(sink, mop_cleanup+pm_reduce_local, ctx=ReduceContext(), name="remove_reduce")
 
   # add gpu dims (late). this works after devectorize, but it's faster here
   sink = graph_rewrite(sink, pm_add_gpudims, ctx=ren, name="add gpudims")
 
   # **** optimizations are done, now we lower to actual code ****
 
+  sink = graph_rewrite(sink, symbolic_simple+unbroadcast, name="*** unbroadcast")
+
   # add loads and remove invalids
-  sink = graph_rewrite(sink, pm_add_loads+pm_remove_invalid, name="** add loads (code)")
+  sink = graph_rewrite(sink, pm_move_regs, name="** add loads")
 
   # devectorize
-  sink = graph_rewrite(sink, sym+devectorize_alu+devectorize_buf_and_index+load_store_folding, ctx=ren, name="devectorize")
-
-  # this is new style (TODO: this should all be removed)
-  sink = graph_rewrite(sink, pm_render, name="pm_render gep/stack")
-  sink = graph_rewrite(sink, pm_remove_vec_dtypes, name="transform to new style")
+  sink = graph_rewrite(sink, symbolic_simple+devectorizer2, ctx=ren, name="devectorize2")
 
   # simplify indexing
   sink = graph_rewrite(sink, indexing_simplify, name="simplify load/store indexing")
 
+  # some coalesing misses without this
+  sink = graph_rewrite(sink, sym, name="early symbolic")
+
   # do memory coalesing (late)
   sink = memory_coalesing(sink, ren)
-  sink = graph_rewrite(sink, pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True)
+  sink = graph_rewrite(sink, symbolic_simple+ew_devectorizer+pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True)
 
   # extra symbolic before decomp. crashes without this?
   sink = graph_rewrite(sink, sym, name="extra symbolic")
@@ -142,7 +371,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   # final rules for the renderer (without sym)
   extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([])
   pm_final_rewrite = pm_decomp+extra_matcher+pm_split_ends+pm_no_weakints
-  sink = graph_rewrite(sink, pm_final_rewrite, ctx=ren, name="final rewrite")
+  sink = graph_rewrite(sink, pm_final_rewrite+pm_remove_invalid, ctx=ren, name="final rewrite")
 
   # this was the linearizer
   sink = graph_rewrite(sink, pm_add_control_flow, ctx=CFGContext(sink), name="add control flow", bottom_up=True)

tinygrad/codegen/decomp/dtype.py

@@ -116,7 +116,7 @@ def f2f_load(x: UOp, fr:DType, to:DType) -> UOp:
 
 def f2f_store(st, idx, val, fr:DType, to:DType):
   if (n:=val.max_numel()) == 1: return st.replace(src=(idx, f2f(val.bitcast(f2f_dt[to]), to, fr)))
-  return UOp.group(*(st.replace(src=(reindex(idx, i, 1), f2f(val.gep(i).bitcast(f2f_dt[to]), to, fr))) for i in range(n)))
+  return UOp.group(*(st.replace(src=(reindex(idx, i, 1), f2f(val.index(i).bitcast(f2f_dt[to]), to, fr))) for i in range(n)))
 
 pm_long_decomp = PatternMatcher([
   (UPat(GroupOp.Defines, src=(UPat.var("sz"),), name="x"), lambda x,sz:

tinygrad/codegen/gpudims.py

@@ -1,6 +1,5 @@
 import math
 from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType
-from tinygrad.helpers import dedup, get_contraction
 from tinygrad.dtype import dtypes, AddrSpace, Invalid
 from tinygrad.renderer import Renderer
 
@@ -23,7 +22,7 @@ def _split_dims(dims, max_sizes):
       div = next((d for d in range(2, math.ceil(math.sqrt(_dims[i])) + 1) if (_dims[i] % d) == 0), 1)
       if div == 1: raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
       _dims[i], _dims[(i+1)%len(_dims)] = _dims[i]//div, _dims[(i+1)%len(_dims)]*div
-  return tuple(_dims[:2] if _dims[2] == 1 else _dims[0] if _dims[1:3] == [1,1] else _dims)
+  return tuple(_dims[:2] if _dims[2] == 1 else _dims)
 
 def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|None, reverse=False) -> list[UOp]:
   if reverse: return get_grouped_dims(prefix, dims[::-1], max_sizes)[::-1]
@@ -36,24 +35,8 @@ def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|No
     # try to split up dims: (a,) -> (b, c)
     if limited == dims: limited = _split_dims(dims, max_sizes)
   raw_idxs = [UOp.special(s, f"{prefix}{i}") for i,s in enumerate(limited)]
-  if len(limited) < len(dims):
-    ret = []
-    if (contraction:=get_contraction(dims, limited)) is None: raise RuntimeError(f"get_contraction should not be None {dims=} {limited=}")
-    for idx, contraction_group in zip(raw_idxs, contraction):
-      for c in contraction_group[:-1]:
-        ret.append(idx % dims[c])
-        idx //= dims[c]
-      ret.append(idx)
-    return ret
-  elif (a:=len(limited)) > (b:=len(dims)):
-    if a == 2 and b == 1: return [raw_idxs[0] * limited[1] + raw_idxs[1]]
-    if a == 3 and b == 1: return [(raw_idxs[0] * limited[1] + raw_idxs[1]) * limited[2] + raw_idxs[2]]
-  if limited != dims:
-    # Convert to 1D
-    flat = raw_idxs[0]*limited[1]+raw_idxs[1] if len(limited) == 2 else raw_idxs[0]*(limited[1]*limited[2])+raw_idxs[1]*limited[2]+raw_idxs[2]
-    # Get back original indices from 1D
-    return [flat//dims[1], flat%dims[1]] if len(dims) == 2 else [flat//(dims[2]*dims[1]), (flat//dims[2])%dims[1], flat%dims[2]]
-  return raw_idxs
+  flat = sum(idx * math.prod(limited[i+1:]) for i,idx in enumerate(raw_idxs))
+  return [ssimplify(flat // math.prod(dims[i+1:])) if i == 0 else ssimplify((flat // math.prod(dims[i+1:])) % dims[i]) for i in range(len(dims))]
 
 def add_gpudims(ctx:Renderer, s:UOp):
   if s.arg is None: return None
@@ -64,14 +47,13 @@ def add_gpudims(ctx:Renderer, s:UOp):
   all_ranges = {x.arg[0:-1]:x for x in s_topo if x.op is Ops.RANGE}
 
   # extract global/local dims
-  global_dims = sorted(dedup([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] in (AxisType.GLOBAL, AxisType.THREAD)]))
-  local_dims = sorted(dedup([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] in (AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE)]))
+  global_dims = sorted([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] in (AxisType.GLOBAL, AxisType.THREAD)])
+  local_dims = sorted([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] in (AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE)])
   if not global_dims and not local_dims: return None
 
   # get global and local shape
-  ranges = [all_ranges[r] for r in global_dims+local_dims if r in all_ranges]
-  global_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg[0:-1] in global_dims])
-  local_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg[0:-1] in local_dims])
+  global_shape = tuple(ssimplify(all_ranges[r].src[0]) for r in global_dims)
+  local_shape = tuple(ssimplify(all_ranges[r].src[0]) for r in local_dims)
 
   # get the idxs
   ki: KernelInfo = s.arg

tinygrad/codegen/late/coalese.py

@@ -1,28 +1,71 @@
 from typing import Any
-import itertools
+import itertools, functools
 from collections import defaultdict
-from tinygrad.dtype import dtypes, AddrSpace, Invalid, ImageDType
+from tinygrad.dtype import dtypes, AddrSpace, Invalid, ImageDType, DType
 from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp
-from tinygrad.helpers import getenv, IMAGE, all_same
+from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
+from tinygrad.helpers import getenv, IMAGE, OSX, ceildiv
 from tinygrad.renderer import Renderer
-from tinygrad.uop.symbolic import symbolic_simple
-from tinygrad.codegen.late.devectorizer import image_valid_dims, _drop_valid_stmts, uop_given_valid
 
-def do_devectorize(b:UOp):
-  if b.shape == (): return None
-  # broadcasting needs to be already unpacked
-  if not all_same([x.shape for x in b.src]): return None
-  src = []
-  for idx in itertools.product(*[range(x) for x in b.shape]):
-    idx_c = [UOp.const(dtypes.weakint, i) for i in idx]
-    src.append(b.replace(src=tuple([x.index(*idx_c) for x in b.src])))
-  return UOp._stack(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src)
+# ***** image load valid simplification *****
 
-devectorizer2 = PatternMatcher([
-  # unpack broadcasting
-  (UPat(GroupOp.Elementwise, name="b"), do_devectorize),
+@functools.cache
+def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]:
+  # can drop valid if idx is out of bound when valid is False
+  drop_stmt = []
+  for i,stmt in enumerate(valid.split_uop(Ops.AND)):
+    if (res:=parse_valid(stmt)) is None: continue
+    X, is_upper_bound, c = res
+
+    # for X0 + X1 + ... >= 1, check if it's out of bound when Xi = 0 for all i
+    if not is_upper_bound and c == 1 and all(u.op in GroupOp.Irreducible and u.vmin == 0 for u in X.split_uop(Ops.ADD)):
+      testidx = functools.reduce(lambda nowidx,u: nowidx.substitute({u:u.const_like(0)}), X.split_uop(Ops.ADD), idx)
+      if testidx.index(0).vmax < 0 or testidx.index(1).vmax < 0:
+        drop_stmt.append(stmt)
+        continue
+
+    # 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)
+      for coord,b in zip(idx.src, (width, height)):
+        rw = coord.substitute({X:fake}).simplify()
+        if rw.vmin >= b or rw.vmax < 0:
+          drop_stmt.append(stmt)
+          break
+  return drop_stmt
+
+def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
+  idx = uop_given_valid(valid, start_idx)
+  return None if idx is start_idx else buf.index(idx.valid(valid), ptr=True)
+
+def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|None:
+  if not isinstance(buf.dtype, ImageDType): return None
+  start_idx = idx_x._stack(idx_y)
+  idx = uop_given_valid(valid, start_idx)
+  drop_stmt = _drop_valid_stmts(valid, idx, buf.dtype.shape[0], buf.dtype.shape[1])
+
+  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)
+  return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid), ptr=True) if new_valid is not None else buf.index(idx_y, idx_x, ptr=True)
+
+indexing_simplify = PatternMatcher([
+  # image load valid idx simplification
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)),
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("valid").where(UPat.var("idx_y"), UPat(arg=Invalid)),
+                                         UPat.var("valid").where(UPat.var("idx_x"), UPat(arg=Invalid)))), simplify_valid_image_load),
 ])
 
+# get list of (height, width) that do not require pitch padding
+def image_valid_dims(base:DType, size:int, arch:str) -> list[tuple[int,int]]:
+  if (ALIGN:=next((int(p.split('=')[1]) for p in arch.split(',') if p.startswith("IMAGE_PITCH_ALIGNMENT=")), 0)) == 0: return []
+  MAXW, pxls = 16384, size // 4
+  if base not in (dtypes.half, dtypes.float) or size > 4*MAXW*MAXW: return []
+  # height=1 images just need to abide by alignment requirements in bytes, not pixels!
+  if size % (ALIGN * 4) != 0: return [] if (base.itemsize * size) % (64 if OSX else ALIGN) != 0 or pxls > MAXW else [(1, pxls)]
+  return [(pxls//ALIGN//k, ALIGN*k) for k in range(ceildiv(pxls//ALIGN, MAXW), min(pxls//ALIGN, MAXW//ALIGN)+1) if (pxls//ALIGN)%k == 0]
+
 def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
   shapes, ren = ctx
   if not IMAGE or ren.target.device not in {"QCOM", "CL", "PYTHON", "NULL"}: return None
@@ -31,14 +74,14 @@ def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
   # search for dims that drop the most valid statements
   best_drop, cands = -1, []
   for ch, cw in [shapes[buf.arg.slot]] if buf.arg.slot in shapes else image_valid_dims(buf.dtype, buf.max_numel(), ren.target.arch):
-    cidx = uop_given_valid(valid, UOp.vectorize((x//4)%cw, x//(4*cw)))
+    cidx = uop_given_valid(valid, ((x//4)%cw)._stack(x//(4*cw)))
     dropped = len(_drop_valid_stmts(valid, cidx, ch, cw))
     if dropped > best_drop: best_drop, cands = dropped, [(ch, cw, cidx)]
     elif dropped == best_drop: cands.append((ch, cw, cidx))
   # if no candidates, we don't rewrite
   if len(cands) == 0: return None
   # and tiebreak with indexing complexity (ie. number of nodes)
-  h, w, cidx = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].gep(1).simplify().backward_slice))
+  h, w, cidx = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].index(1).simplify().backward_slice))
   buf = buf.replace(dtype=(dtypes.imageh if buf.dtype.itemsize == 2 else dtypes.imagef)((h, w, 4)))
   shapes[buf.arg.slot] = (h, w)
   if valid.op is not Ops.CONST or valid.arg is not True:
@@ -53,7 +96,7 @@ pm_simplify_add_image = PatternMatcher([
   (UPat(Ops.INDEX, dtype=dtypes.float, name="x").load(dtype=dtypes.half), lambda x: x.load().cast(dtypes.half)),
   (UPat(Ops.INDEX, dtype=dtypes.float, name="x").store(UPat(name="d", dtype=dtypes.half)), lambda x,d: x.store(d.cast(dtypes.float))),
   (UPat.var("x", dtype=dtypes.float).cast(dtypes.half).cast(dtypes.float), lambda x: x),
-])+devectorizer2+symbolic_simple
+])
 
 def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp:
   if getenv("DMC"): return sink
@@ -64,7 +107,7 @@ def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp:
     # TODO: this should handle images too, it's just memory coalesing
     if u.op in {Ops.LOAD, Ops.STORE}:
       assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalesing does not support gated loads/stores"
-      assert u.src[0].op is Ops.INDEX
+      assert u.src[0].op is Ops.INDEX, f"memory coalesing should be on INDEX, not {u.src[0].op}"
       buf, idx_u = u.src[0].src
       if buf.addrspace == AddrSpace.REG: continue
       idx: Any = idx_u.src[1] if idx_u.op is Ops.WHERE and idx_u.src[2].arg is Invalid else idx_u

tinygrad/codegen/late/devectorizer.py

@@ -1,242 +0,0 @@
-import functools, itertools
-from dataclasses import dataclass
-from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace, Invalid, PtrDType
-from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, identity_element
-from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
-from tinygrad.helpers import getenv, flatten, prod, OSX, ceildiv
-
-# ***** image load valid simplification *****
-
-@functools.cache
-def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]:
-  # can drop valid if idx is out of bound when valid is False
-  drop_stmt = []
-  for i,stmt in enumerate(valid.split_uop(Ops.AND)):
-    if (res:=parse_valid(stmt)) is None: continue
-    X, is_upper_bound, c = res
-
-    # for X0 + X1 + ... >= 1, check if it's out of bound when Xi = 0 for all i
-    if not is_upper_bound and c == 1 and all(u.op in GroupOp.Irreducible and u.vmin == 0 for u in X.split_uop(Ops.ADD)):
-      testidx = functools.reduce(lambda nowidx,u: nowidx.substitute({u:u.const_like(0)}), X.split_uop(Ops.ADD), idx)
-      if testidx.gep(0).vmax < 0 or testidx.gep(1).vmax < 0:
-        drop_stmt.append(stmt)
-        continue
-
-    # 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)
-      for coord,b in zip(idx.src, (width, height)):
-        rw = coord.substitute({X:fake}).simplify()
-        if rw.vmin >= b or rw.vmax < 0:
-          drop_stmt.append(stmt)
-          break
-  return drop_stmt
-
-def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
-  idx = uop_given_valid(valid, start_idx)
-  return None if idx is start_idx else buf.index(idx.valid(valid), ptr=True)
-
-def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|None:
-  if not isinstance(buf.dtype, ImageDType): return None
-  start_idx = UOp.vectorize(idx_x, idx_y)
-  idx = uop_given_valid(valid, start_idx)
-  drop_stmt = _drop_valid_stmts(valid, idx, buf.dtype.shape[0], buf.dtype.shape[1])
-
-  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.gep(1), idx.gep(0)
-  return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid), ptr=True) if new_valid is not None else buf.index(idx_y, idx_x, ptr=True)
-
-indexing_simplify = PatternMatcher([
-  # image load valid idx simplification
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)),
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("valid").where(UPat.var("idx_y"), UPat(arg=Invalid)),
-                                         UPat.var("valid").where(UPat.var("idx_x"), UPat(arg=Invalid)))), simplify_valid_image_load),
-])
-
-# ***** load/store grouping *****
-
-# get list of (height, width) that do not require pitch padding
-def image_valid_dims(base:DType, size:int, arch:str) -> list[tuple[int,int]]:
-  if (ALIGN:=next((int(p.split('=')[1]) for p in arch.split(',') if p.startswith("IMAGE_PITCH_ALIGNMENT=")), 0)) == 0: return []
-  MAXW, pxls = 16384, size // 4
-  if base not in (dtypes.half, dtypes.float) or size > 4*MAXW*MAXW: return []
-  # height=1 images just need to abide by alignment requirements in bytes, not pixels!
-  if size % (ALIGN * 4) != 0: return [] if (base.itemsize * size) % (64 if OSX else ALIGN) != 0 or pxls > MAXW else [(1, pxls)]
-  return [(pxls//ALIGN//k, ALIGN*k) for k in range(ceildiv(pxls//ALIGN, MAXW), min(pxls//ALIGN, MAXW//ALIGN)+1) if (pxls//ALIGN)%k == 0]
-
-def expand_index(ctx, buf:UOp, vec:UOp):
-  if getenv("UNSAFE_DISABLE_MASK", 0): vec = vec.get_idx()
-  # generate the individual indexes
-  return UOp(Ops.STACK, buf.dtype, tuple(buf.index(vec.gep(i), ptr=True) for i in range(vec.dtype.count)))
-
-def load_stack(stack:UOp, ld:UOp):
-  offset, ret = 0, []
-  for x in stack.src:
-    src = [x]
-    for s in ld.src[1:]:
-      src.append(s.gep(tuple(range(offset, offset+x.dtype.count))) if s.dtype.vcount > 1 else s)
-    ret.append(ld.replace(dtype=x.dtype.base, src=tuple(src)))
-    offset += x.dtype.count
-  return UOp(Ops.STACK, stack.dtype.base.vec(len(stack.src)), tuple(ret))
-
-def store_stack(stack:UOp, data:UOp):
-  offset, ret = 0, []
-  for x in stack.src:
-    ret.append(x.store(data.gep(tuple(range(offset, offset+x.dtype.count)))))
-    offset += x.dtype.count
-  return UOp.group(*ret)
-
-load_store_folding = PatternMatcher([
-  (UPat(Ops.INDEX, src=(UPat(Ops.STACK, src=UPat(name="buf")), UPat.var("vec"))), expand_index),
-  # put STACK of indexes after LOAD/STORE
-  (UPat(Ops.LOAD, src=(UPat(Ops.STACK, src=UPat(Ops.INDEX), name="stack"),), name="ld", allow_any_len=True), load_stack),
-  (UPat(Ops.STORE, src=(UPat(Ops.STACK, src=UPat(Ops.INDEX), name="stack"), UPat(name="data"))), store_stack),
-])
-
-# *** uop expander ***
-
-# TODO: there's a lot shared with gep_through_wmma here
-def no_vectorized_wmma(wmma:UOp):
-  out_sz = prod(x[1] for x in wmma.arg[6][-1])
-  if wmma.dtype.count == out_sz: return None
-  tsrcs = []
-  for s,sz in zip(wmma.src, wmma.arg[6]):
-    ssz = prod(x[1] for x in sz)
-    tsrcs.append([s.gep(tuple(range(grp, grp+ssz))) for grp in range(0, s.dtype.count, ssz)])
-  wmmas = [UOp(Ops.WMMA, wmma.dtype.scalar().vec(out_sz), tsrc, wmma.arg) for tsrc in zip(*tsrcs)]
-  wmma_ex = flatten([[e.gep(i) for i in range(out_sz)] for e in wmmas])
-  return UOp(Ops.STACK, wmma.dtype, tuple(wmma_ex))
-
-def no_vectorized_alu(alu:UOp):
-  if alu.dtype.vcount == 1: return None
-  alus = tuple(UOp(alu.op, alu.dtype.scalar(), tuple(s.gep(i) for s in alu.src), alu.arg) for i in range(alu.dtype.vcount))
-  return UOp(Ops.STACK, alu.dtype, alus)
-
-def no_vectorized_buf(buf:UOp):
-  if not isinstance(buf.dtype, PtrDType): return None
-  if buf.addrspace not in (AddrSpace.LOCAL, AddrSpace.REG): return None
-  # TODO: this fails on regs
-  #assert buf.max_numel() == buf.ptrdtype.size
-  sz = buf.ptrdtype.size*buf.ptrdtype.count
-  return buf.replace(dtype=buf.ptrdtype.base.scalar().ptr(sz, buf.addrspace), src=(UOp.const(dtypes.int, sz),)).cast(buf.dtype)
-
-def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp, bcast:UOp|None=None):
-  if buf.addrspace not in (AddrSpace.LOCAL, AddrSpace.REG): return None
-  cnt = cast.dtype.count
-  if bcast is not None and bcast.op is Ops.GEP:
-    # GEP selects specific lanes; bcast.arg[k] is the offset for lane k, iterate groups × selected lanes
-    pairs = [(k, g + bcast.arg[k]) for g, k in itertools.product(range(cast.dtype.vcount), range(len(bcast.arg)))]
-  elif bcast is not None:
-    # BROADCAST: cross product of components × lanes
-    pairs = [(j, c) for c, j in itertools.product(range(cnt), range(bcast.dtype.vcount))]
-  else:
-    # simple scalar index: one lane, all components
-    pairs = [(0, c) for c in range(cnt)]
-  idx_lanes, offsets = (tuple(x) for x in zip(*pairs))
-  return buf.broadcast(len(pairs)).index(idx.gep(idx_lanes)*cnt + UOp.const(dtypes.weakint.vec(len(pairs)), offsets), ptr=True)
-
-devectorize_buf_and_index = PatternMatcher([
-  (UPat(Ops.BUFFER, name="buf"), no_vectorized_buf),
-  (UPat(Ops.BUFFER).or_after(name="buf").cast(name="cast").index(UPat.var("idx")), no_vectorized_index),
-  (UPat(Ops.BUFFER).or_after(name="buf").cast(name="cast").broadcast(name="bcast").index(UPat.var("idx")), no_vectorized_index),
-  (UPat(Ops.BUFFER).or_after(name="buf").cast(name="cast").gep(name="bcast").index(UPat.var("idx")), no_vectorized_index),
-])
-
-devectorize_alu = PatternMatcher([
-  # CAST after AFTER
-  (UPat(Ops.CAST, name="c").f(Ops.AFTER, allow_any_len=True, name="a"), lambda c,a: c.src[0].after(*a.src[1:]).cast(c.dtype)),
-  # no ALU on vectorized dtypes
-  (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name="alu"), no_vectorized_alu),
-  (UPat(Ops.WMMA, name="wmma"), no_vectorized_wmma),
-])
-
-pm_render = PatternMatcher([
-  # for rendering, we use explicit VECTORIZE
-  (UPat(Ops.CONST, name='c'),
-   lambda c: UOp(Ops.STACK, c.dtype, (UOp.const(c.dtype.scalar(), c.arg),)*c.dtype.vcount) if c.dtype.vcount > 1 else None),
-  (UPat(Ops.GEP, name='gep'), lambda gep: UOp(Ops.STACK, gep.dtype, tuple(gep.src[0].gep(x) for x in gep.arg)) if len(gep.arg) > 1 else None),
-  (UPat(Ops.GEP, name='gep'), lambda gep: gep.src[0] if gep.src[0].dtype.vcount == 1 and gep.arg == (0,) else None),
-  (UPat(Ops.STACK, src=(UPat(name='x'),)), lambda x: x),
-])
-
-# *** Ops.REDUCE -> Ops.DEFINE_ACC ***
-
-@dataclass
-class ReduceContext:
-  acc_num: int = 0
-
-def horizontal_reduce(inp:UOp, out_dtype:DType) -> list[UOp]:
-  # if this has a horizontal reduction component, do that first
-  if inp.dtype != out_dtype:
-    # NOTE: [0 1 2 3 4 5 6 7] -> [0+4, 1+5, 2+6, 3+7]
-    horizontal_amount = inp.dtype.count//out_dtype.count
-    return [inp.gep(tuple(range(i, inp.dtype.count, horizontal_amount))) for i in range(0, horizontal_amount)]
-  return [inp]
-
-def reduce_to_acc(ctx:ReduceContext, red:UOp):
-  inp, reduce_range = red.src[0], red.src[1:]
-  lst = horizontal_reduce(inp, red.dtype)
-  assert all(x.dtype == red.dtype for x in lst), f"horizontal reduction mismatch {lst[0].dtype} != {red.dtype}"
-  # if we have a range
-  if len(reduce_range) != 0:
-    topo = inp.toposort()
-    ended_ranges = flatten([x.ended_ranges for x in topo if x.op is Ops.END])
-    input_ranges = tuple([x for x in topo if x.op is Ops.RANGE and x not in reduce_range and x not in ended_ranges])
-    identity = red.const(red.dtype, identity_element(red.arg[0], red.dtype.scalar()))
-    acc = UOp.placeholder((1,), red.dtype, ctx.acc_num, AddrSpace.REG)
-    acc_init = acc.after(*input_ranges).index(UOp.const(dtypes.weakint, 0)).store(identity)
-    lst = [acc.after(acc_init, *reduce_range).index(UOp.const(dtypes.weakint, 0))] + lst  # put acc as the first element
-    ctx.acc_num += 1
-  ret = functools.reduce(lambda x,y: x.alu(red.arg[0], y), lst)
-  if len(reduce_range) == 0: return ret
-  end = acc.index(UOp.const(dtypes.weakint, 0)).store(ret).end(*reduce_range).rtag("mergeable")
-  return acc.after(end).index(UOp.const(dtypes.weakint, 0))
-
-def merge_reduce_ends(ctx:ReduceContext, sink:UOp):
-  # merge ENDs that share the same range and nesting context (only those created by reduce_to_acc)
-  # ENDs at different nesting depths get cloned RANGEs so each RANGE maps to one END
-  range_to_ends: dict[tuple[UOp, ...], list[UOp]] = {}
-  for u in sink.backward_slice:
-    if u.op is Ops.END and u.tag == "mergeable": range_to_ends.setdefault(u.src[1:], []).append(u)
-  subs: dict[UOp, UOp] = {}
-  next_axis = max((u.arg[0] for u in sink.backward_slice if u.op is Ops.RANGE), default=-1) + 1
-  for r, ends in range_to_ends.items():
-    if len(ends) <= 1: continue
-    by_ctx: dict[frozenset[UOp], list[UOp]] = {}
-    for e in ends: by_ctx.setdefault(frozenset(e.ranges), []).append(e)
-    for i, group in enumerate(by_ctx.values()):
-      tr = r if i == 0 else tuple(rr.replace(arg=(next_axis + j, *rr.arg[1:])) for j, rr in enumerate(r))
-      if i > 0: next_axis += len(r)
-      mapped = [e.substitute(dict(zip(r, tr))) if i > 0 else e for e in group]
-      merged = mapped[0] if len(mapped) == 1 else UOp.group(*(e.src[0] for e in mapped)).end(*tr)
-      for e in group: subs[e] = merged
-  return sink.substitute(subs) if subs else None
-
-pm_reduce = PatternMatcher([
-  # invalid -> identity element
-  (UPat(Ops.REDUCE, src=(invalid_gate,), allow_any_len=True, name="red"), lambda red,cond,x,i:
-    red.replace(src=(cond.where(x, identity_element(red.arg[0], x.dtype.scalar())),)+red.src[1:])),
-  # REDUCE -> DEFINE_ACC+ASSIGN, then merge ENDs with same range
-  (UPat(Ops.REDUCE, name="red"), reduce_to_acc),
-  (UPat(Ops.SINK, name="sink"), merge_reduce_ends),
-  # tensor core built in accumulate
-  (UPat(Ops.WMMA, name="wmma") + UPat.var("add"),
-    lambda add, wmma: UOp(wmma.op, wmma.dtype, (wmma.src[0], wmma.src[1], wmma.src[2]+add), wmma.arg)),
-])
-
-# add loads
-
-def add_load(idx:UOp):
-  if isinstance(idx.dtype, PtrDType): return None
-  assert isinstance(idx.src[0].dtype, PtrDType), f"param is not PtrDType {idx.src[0].dtype}"
-  return idx.replace(dtype=idx.src[0].dtype).load(dtype=idx.dtype.base)
-
-pm_add_loads = PatternMatcher([
-  # add loads to non ptr index
-  (UPat(Ops.INDEX, name="idx"), add_load),
-  # remove loads from stores
-  (UPat(Ops.STORE, src=(UPat(Ops.LOAD),), allow_any_len=True, name="s"), lambda s: s.replace(src=(s.src[0].src[0],)+s.src[1:])),
-  (UPat(Ops.LOAD, src=(UPat(Ops.LOAD),), allow_any_len=True, name="l"), lambda l: l.replace(src=(l.src[0].src[0],)+l.src[1:])),
-])

tinygrad/codegen/late/expander.py

@@ -1,160 +0,0 @@
-# this converts a lowerer program into a vectorized program
-import functools, itertools
-from tinygrad.dtype import dtypes, PtrDType, AddrSpace
-from tinygrad.helpers import dedup, flatten, all_same, prod, partition
-from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, AxisType, range_start
-from tinygrad.schedule.rangeify import BufferizeOpts
-
-def _expand_arg_to_idx(args:tuple[tuple[int, int], ...], rpk:dict[int, int]) -> int:
-  idx, mul = 0, 1
-  for axis,m in args[::-1]:
-    idx += rpk[axis] * mul
-    mul *= m
-  return idx
-
-def _choices_from_args(args:tuple[tuple[int, int], ...]) -> list[dict[int, int]]:
-  return [dict(x) for x in itertools.product(*[zip(itertools.repeat(axis), range(m)) for axis,m in args])]
-
-@functools.cache
-def _swizzle_args(cargs:tuple[tuple[int, int], ...], eargs:tuple[tuple[int, int], ...], exclude_args:tuple[int, ...]) -> list[int]:
-  return [_expand_arg_to_idx(eargs, {**rpk, **{x:0 for x in exclude_args}} if exclude_args else rpk) for rpk in _choices_from_args(cargs)]
-
-def do_expand(root:UOp):
-  expands = [x for x in root.src if x.op is Ops.UNROLL]
-  if len(expands) == 0: return None
-  # NOTE: we 0 out the reduce axis for WMMA. in theory they should all be the same, but is this always correct?
-  exclude_args = tuple(dedup(root.arg[-1] + tuple(y[0] for y in flatten(root.arg[-2])))) if root.op is Ops.WMMA else ()
-  if all_same(expands_args:=[x.arg for x in expands]) and len(exclude_args) == 0:
-    # if there's only one expand arg, it's okay to use it (optimization)
-    expand_args = expands[0].arg
-  else:
-    # otherwise, we sort them and GEP
-    expand_args = tuple(x for x in sorted(dedup(flatten(expands_args))) if x[0] not in exclude_args)
-  expand_sz = prod([x[1] for x in expand_args])
-  new_srcs = []
-  for i,src in enumerate(root.src):
-    if src.op is Ops.UNROLL:
-      if expand_args == src.arg:
-        # just remove the expand
-        new_srcs.append(src.src[0])
-      else:
-        lst = _swizzle_args(expand_args, src.arg, exclude_args)
-        # if the base dtype is > 1, put those at the end
-        if src.dtype.count > 1: lst = flatten([[i*src.dtype.count+j for j in range(src.dtype.count)] for i in lst])
-        new_srcs.append(src.src[0].gep(tuple(lst)))
-    else:
-      # non-UNROLL input
-      if root.op in range_start and i >= range_start[root.op]:
-        # for any range args of REDUCE/WMMA/END/etc., pass them through
-        new_srcs.append(src)
-      elif root.op is Ops.INDEX and i >= 1 and not isinstance(root.dtype, PtrDType):
-        new_srcs.append(src)
-      elif src.dtype.count > 1:
-        # put any input dtype > 1 grouped together
-        new_srcs.append(src.gep(tuple(i for _ in range(expand_sz) for i in range(src.dtype.count))))
-      else:
-        # repeat the arg
-        new_srcs.append(src.broadcast(expand_sz))
-
-  # for non-PtrDType INDEX on REG buffers, expand into individual scalar INDEXes instead of one vectorized INDEX
-  # this avoids creating a VECTORIZE of REG pointers which the devectorizer can't resolve
-  if root.op is Ops.INDEX and not isinstance(root.dtype, PtrDType) and \
-     isinstance(root.src[0].dtype, PtrDType) and root.src[0].dtype.addrspace == AddrSpace.REG:
-    idxs = []
-    for j in range(expand_sz):
-      idx_srcs = tuple(s.gep(j) if isinstance(s.dtype, PtrDType) or s.dtype.count > 1 else s for s in new_srcs)
-      idxs.append(UOp(Ops.INDEX, root.dtype, idx_srcs, root.arg))
-    return UOp(Ops.UNROLL, root.dtype, (UOp(Ops.STACK, root.dtype.vec(expand_sz), tuple(idxs)),), expand_args)
-
-  new_arg = root.arg
-  if root.op is Ops.GEP:
-    assert root.dtype.count == 1
-    # is this right?
-    new_arg = tuple(range(root.arg[0], new_srcs[0].dtype.count, new_srcs[0].dtype.count // expand_sz))
-  nsrc = UOp(root.op, root.dtype.scalar().vec(root.dtype.count*expand_sz), tuple(new_srcs), new_arg)
-  return UOp(Ops.UNROLL, root.dtype, (nsrc,), expand_args)
-
-def do_contract(con:UOp):
-  ex = con.src[0]
-  # CONTRACT without UNROLL repeats the element VECTORIZED
-  if ex.op is not Ops.UNROLL: return UOp(Ops.STACK, con.dtype, con.src*con.dtype.count)
-  # CONTRACT may remove several axes from UNROLL
-  assert con.dtype == dtypes.void or con.dtype.count == prod([x[1] for x in con.arg]), "dtype is wrong"
-  idxs = []
-  for rpk in _choices_from_args(new_ex_args:=tuple(x for x in ex.arg if x not in con.arg)):
-    idxs += [_expand_arg_to_idx(ex.arg, {**rpk, **lrpk}) for lrpk in _choices_from_args(con.arg)]
-  return UOp(Ops.UNROLL, con.dtype, (ex.src[0].gep(tuple(idxs)),), new_ex_args)
-
-def end_unrolls(u:UOp):
-  unrolls, src = partition(u.src[1:], lambda x: x.op is Ops.UNROLL)
-  if not len(unrolls): return None
-  ret = UOp(Ops.CONTRACT, dtypes.void, (u.src[0],), sum([x.arg for x in unrolls], start=()))
-  return u.replace(src=(ret,)+tuple(src))
-
-expander = PatternMatcher([
-  # push broadcast through AFTER/END
-  (UPat.var("x").broadcast(name="b").after(name="a", allow_any_len=True), lambda x,b,a: x.after(*a.src[1:]).broadcast(len(b.src))),
-  (UPat.var("x").broadcast(name="b").end(name="a", allow_any_len=True), lambda x,b,a: x.end(*a.src[1:]).broadcast(len(b.src))),
-  # END on UNROLL ends the UNROLL
-  (UPat(Ops.END, name="u"), end_unrolls),
-  # BUFFERIZE puts UNROLLs for ranges as contract
-  (UPat(Ops.STAGE, src=(UPat(Ops.UNROLL), UPat(Ops.UNROLL)), name="x"),
-    lambda x: x.replace(src=tuple(UOp(Ops.CONTRACT, dtype=s.dtype.vec(x.src[1].src[0].dtype.count), src=(s,), arg=x.src[1].arg) for s in x.src))),
-  # double expand
-  (UPat(Ops.UNROLL, name="outer", src=(UPat(Ops.UNROLL, name="inner"),)),
-   lambda outer, inner: UOp(Ops.UNROLL, outer.dtype, (inner.src[0],), inner.arg+outer.arg)),
-  # do expansion
-  (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.GEP, Ops.WMMA, Ops.LOAD, Ops.STORE, Ops.INDEX, Ops.STAGE,
-         Ops.STACK, Ops.REDUCE, Ops.END, Ops.AFTER), name="root", custom_early_reject=set([Ops.UNROLL])), do_expand),
-  (UPat(Ops.CONTRACT, name="con"), do_contract),
-  # empty UNROLL is NOOP
-  (UPat(Ops.UNROLL, src=(UPat.var('x'),), arg=()), lambda x: x),
-])
-
-# ****
-
-def fix_reduce_unroll(x:UOp):
-  reduce_range, reduce_expand = partition(x.src[1:], lambda y: y.op is Ops.RANGE)
-  if len(reduce_expand) == 0: return None
-  reduce_expand = [x for x in reduce_expand if x.op is not Ops.CONST]
-  assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand}"
-  ret = x.src[0]
-  if len(contract_axis:=flatten(x.arg for x in reduce_expand)):
-    ret = UOp(Ops.CONTRACT, x.dtype.vec(prod(x[1] for x in contract_axis)), (ret,), tuple(contract_axis), tag=1)
-  # REDUCE supports both "horizontal" reduction and range reduction. the horizontal elements are taken in the nearest group
-  return x.replace(src=(ret,)+tuple(reduce_range))
-
-def fix_store_unroll(x:UOp):
-  store_expand, store_range = partition(x.src[2:], lambda y: y.op is Ops.UNROLL)
-  if len(store_expand) == 0: return None
-  return UOp(Ops.CONTRACT, dtypes.void, (x.replace(src=x.src[:2]+tuple(store_range)),), tuple(flatten(x.arg for x in store_expand)), tag=1)
-
-def fix_group_for_reduce(x:UOp):
-  reduce_gfr, reduce_r = partition(x.src[1:], lambda u: u.op is Ops.RANGE and u.arg[1] == AxisType.GROUP_REDUCE)
-  if len(reduce_gfr) == 0: return None
-
-  # NOTE: if there's other locals here, we need them in the buffer too
-  upstream_locals = [u for u in x.toposort() if u.op is Ops.RANGE and u.arg[1] == AxisType.LOCAL]
-
-  # do only the non grouped reduces early
-  ret = x.replace(src=(x.src[0],)+tuple(reduce_r))
-  reduce_loop = [x.replace(arg=(x.arg[0]+100, AxisType.REDUCE)) for x in reduce_gfr]
-  buf = ret.bufferize(*upstream_locals, *reduce_gfr, arg=BufferizeOpts(reduce_gfr[0].arg[0], AddrSpace.LOCAL)).index(*upstream_locals, *reduce_loop)
-
-  # do the final reduce (if/barrier are added in gpudims step)
-  return buf.reduce(*reduce_loop, arg=x.arg)
-
-pm_pre_expander = PatternMatcher([
-  # rewrite UPCAST/UNROLL range to something to be expanded
-  (UPat(Ops.RANGE, name="r"),
-   lambda r: UOp(Ops.UNROLL, r.dtype, (UOp.const(r.dtype.vec(s:=r.vmax+1), tuple(range(s))),), ((r.arg[0],s),)) \
-    if r.arg[1] in {AxisType.UNROLL, AxisType.UPCAST} else None),
-  # fix REDUCEs with UNROLLs
-  (UPat(Ops.REDUCE, name="x"), fix_reduce_unroll),
-  (UPat(Ops.STORE, name="x"), fix_store_unroll),
-])
-
-pm_group_for_reduce = PatternMatcher([
-  # fix group for reduce
-  (UPat(Ops.REDUCE, name="x"), fix_group_for_reduce),
-])

tinygrad/codegen/opt/heuristic.py

@@ -3,7 +3,7 @@ from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.helpers import getenv, DEBUG, prod, NOLOCALS, TC_OPT, TC_SELECT, USE_TC, IMAGE
 from tinygrad.dtype import PtrDType
 from tinygrad.uop.ops import Ops, resolve, AxisType
-from tinygrad.codegen.late.devectorizer import image_valid_dims
+from tinygrad.codegen.late.coalese import image_valid_dims
 from tinygrad.codegen.opt.postrange import Scheduler
 
 def hand_coded_optimizations(k:Scheduler) -> Scheduler:

tinygrad/codegen/opt/postrange.py

@@ -290,6 +290,12 @@ class Scheduler:
             # axes to range number (was done in lowerer)
             tc_upcast_axes = tuple([tuple([(self.rngs[a].arg[0], sz) for a,sz in v]) for v in tc_upcast_axes])
             tc_reduce_axes = tuple([self.rngs[a].arg[0] for a in tc_reduce_axes])
+            def with_missing_tc_axes(arg):
+              ret = list(arg)
+              for rn,_ in tc_upcast_axes[0]+tc_upcast_axes[1]:
+                if rn not in [x[0] for x in ret]: ret.append((rn, 1))
+              return tuple(ret)
+            tc_upcast_axes = tuple(with_missing_tc_axes(v) for v in tc_upcast_axes)
 
             # construct the op
             # TODO: remove tc_upcast_axes from the arg

tinygrad/codegen/simplify.py

@@ -39,7 +39,7 @@ def simplify_merge_adjacent(u:UOp) -> UOp|None:
   return u
 
 def mark_gated(ctx, idx):
-  if idx.src[1].op is Ops.WHERE:
+  if len(idx.src) > 1 and idx.src[1].op is Ops.WHERE:
     x, cond = idx.src[1].get_idx(), idx.src[1].get_valid()
     # get all ranges r with guards "r < c" for some const c
     guards = {r:c for v in cond.split_uop(Ops.AND) if v.op is Ops.CMPLT and (r:=v.src[0]).op is Ops.RANGE and (c:=v.src[1]).op is Ops.CONST}

tinygrad/device.py

@@ -25,11 +25,13 @@ class _Device:
     assert ALLOW_DEVICE_USAGE or ix.split(":")[0] in ["DISK", "TINYFS", "NPY", "PYTHON"], f"usage of device {ix} disallowed"
     return self.__get_canonicalized_item(ix)
   @functools.cache  # this class is a singleton, pylint: disable=method-cache-max-size-none
-  def __get_canonicalized_item(self, ix:str) -> Compiled:
+  def get_class(self, ix:str):
     base = (__package__ or __name__).split('.')[0]  # tinygrad
     x = ix.split(":")[0].lower()
-    ret = [cls for cname, cls in inspect.getmembers(importlib.import_module(f'{base}.runtime.ops_{x}')) \
-           if (cname.lower() == x + "device")][0](ix)
+    return [cls for cname, cls in inspect.getmembers(importlib.import_module(f'{base}.runtime.ops_{x}')) if (cname.lower() == x + "device")][0]
+  @functools.cache  # this class is a singleton, pylint: disable=method-cache-max-size-none
+  def __get_canonicalized_item(self, ix:str) -> Compiled:
+    ret = self.get_class(ix)(ix)
     if DEBUG >= 1: print(f"opened device {ix} from pid:{os.getpid()}")
     self._opened_devices.add(ix)
     return ret

tinygrad/engine/jit.py

@@ -94,16 +94,16 @@ class DepsTracker:
     self.r_dependency_map: dict[int, list[tuple[int, int, Any]]] = collections.defaultdict(list)
 
   @staticmethod
-  def _buf_key(buf:Buffer) -> int: return id(buf.base)
+  def _key(buf:Any) -> tuple[Any, int, int]: return id(buf.base), buf.offset, buf.offset + buf.nbytes
 
-  def access_resources(self, bufs:list[Buffer], write:list[int], new_dependency:Any):
+  def access_resources(self, bufs:list[Any], write:list[int], new_dependency:Any):
     wait_nodes = []
     for i,buf in enumerate(bufs):
-      key, s, e = self._buf_key(buf), buf.offset, buf.offset + buf.nbytes
+      key, s, e = self._key(buf)
       wait_nodes += [dep for st,en,dep in self.w_dependency_map[key] if st < e and s < en]
       if i in write: wait_nodes += [dep for st,en,dep in self.r_dependency_map[key] if st < e and s < en]
     for i,buf in enumerate(bufs):
-      key, s, e = self._buf_key(buf), buf.offset, buf.offset + buf.nbytes
+      key, s, e = self._key(buf)
       if i in write:
         for dmap in [self.w_dependency_map, self.r_dependency_map]:
           kept = []

tinygrad/engine/realize.py

@@ -210,14 +210,15 @@ def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
 
 def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
   if call.arg.aux.inputs is not None:
-    for j,dev in enumerate(call.arg.aux.devs):
-      addrs = [(b.bufs[j] if isinstance(b:=ctx.input_uops[i].buffer, MultiBuffer) else b).get_buf(dev).va_addr for i in call.arg.aux.params]
+    bufs = [_resolve(ctx.input_uops[i], ctx.input_uops).buffer for i in call.arg.aux.input_idxs]
+    for j,dev in enumerate(call.arg.aux.device):
+      addrs = [(b.bufs[j] if isinstance(b, MultiBuffer) else b).get_buf(dev).va_addr for b in bufs]
       buf = b.bufs[j] if isinstance(b:=call.src[1+call.arg.aux.inputs].buffer, MultiBuffer) else b
       buf.ensure_allocated()._buf.cpu_view().view(fmt='Q')[:len(addrs)] = array.array('Q', addrs)
 
   pm_exec.rewrite(call.replace(src=(ast,) + call.src[1:]), replace(ctx, update_stats=False, wait=True))
 
-  for d in call.arg.aux.devs:
+  for d in call.arg.aux.device:
     with track_stats(ctx, call, d, [], ctx.var_vals):
       if ctx.wait: Device[d].synchronize()
   return None
@@ -260,13 +261,13 @@ pm_exec = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="validate", name="ast"),), name="call", allow_any_len=True), exec_validate),
 ])
 
-def compile_linear(linear:UOp, beam:int|None=None, validate=False) -> UOp:
+def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|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 = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True)
   if getenv("HCQ2"):
     from extra.hcq2.hcq2 import hcq_compile
-    linear = hcq_compile(linear)
+    linear = hcq_compile(linear, input_uops)
   return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
 
 def link_linear(linear:UOp) -> UOp:
@@ -275,9 +276,10 @@ def link_linear(linear:UOp) -> UOp:
     linear = hcq_link(linear)
   return linear
 
-def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:tuple[UOp, ...]=(), update_stats=True, jit=False, wait=False):
-  if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU))
-  ctx = ExecContext(var_vals or {}, input_uops, update_stats, jit, wait or DEBUG>=2)
+def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:Sequence[UOp]=(), update_stats=True, jit=False, wait=False):
+  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: 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) -> float:

tinygrad/helpers.py

@@ -122,13 +122,6 @@ def strides_for_shape(shape:tuple[T, ...]) -> tuple[T, ...]:
   strides = tuple(itertools.accumulate(reversed(shape[1:]), operator.mul, initial=1))[::-1]
   return canonicalize_strides(shape, strides)
 
-# returns the axes to create new_shape if new_shape can be created by combining axis from old_shape
-def get_contraction(old_shape:tuple[T, ...], new_shape:tuple[T, ...]) -> list[list[int]]|None: # T is sint
-  acc_old, acc_new = list(itertools.accumulate(old_shape, operator.mul)), list(itertools.accumulate(new_shape, operator.mul))
-  try: split = [0 if isinstance(acc, int) and acc == 1 else acc_old.index(acc)+1 for acc in acc_new]
-  except ValueError: return None
-  return [list(range(st,ed)) for st,ed in zip([0]+split[:-1], split[:-1]+[len(old_shape)])]
-
 def suppress_finalizing(func):
   def wrapper(*args, **kwargs):
     try: return func(*args, **kwargs)
@@ -239,6 +232,7 @@ class _DEV(ContextVar):
 DEV, DEBUG, BEAM, NOOPT = _DEV("DEV", ""), ContextVar("DEBUG", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0)
 IMAGE, FLOAT16, OPENPILOT_HACKS = ContextVar("IMAGE", 0), ContextVar("FLOAT16", 0), ContextVar("OPENPILOT_HACKS", 0)
 JIT, JIT_BATCH_SIZE = ContextVar("JIT", 2 if OSX and ARCH_X86 else 1), ContextVar("JIT_BATCH_SIZE", 32)
+CHUNK_SIZE = 2**20  # TinyFS content-addressed store: blob chunk + hash-tree node granularity
 WINO, CAPTURING, TRACEMETA, NO_COLOR = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1), ContextVar("NO_COLOR", 0)
 TRAINING = ContextVar("TRAINING", 0)
 USE_TC, TC_SELECT, TC_OPT = ContextVar("TC", 1), ContextVar("TC_SELECT", -1), ContextVar("TC_OPT", 0)
@@ -275,7 +269,7 @@ SCACHE = ContextVar("SCACHE", 1)
 # allow use of atomics for embedding backward
 USE_ATOMICS = ContextVar("USE_ATOMICS", 0)
 # don't allow broadcast
-DISALLOW_BROADCAST = ContextVar("DISALLOW_BROADCAST", 1)
+DISALLOW_BROADCAST = ContextVar("DISALLOW_BROADCAST", 0)
 
 @dataclass(frozen=True)
 class Metadata:

tinygrad/mixin/creation.py

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING, Callable, Self
 from tinygrad.dtype import ConstType, DTypeLike, Invalid, dtypes, to_dtype
-from tinygrad.helpers import argfix
+from tinygrad.helpers import argfix, prod
 from tinygrad.mixin.dtype import DTypeMixin
 from tinygrad.mixin.movement import MovementMixin
 
@@ -19,6 +19,26 @@ class CreationMixin(DTypeMixin, MovementMixin):
     if self._uop.axis is None: return self._wrap_uop(fxn(self.shape, None)._uop.shard(self.device, None))
     return self._wrap_uop(UOp.mstack(*[fxn(self._uop.shard_shape, d)._uop for d in self.device]).multi(self._uop.axis))
 
+  @classmethod
+  def empty(cls, *shape, device:str|tuple[str, ...]|None=None, dtype:DTypeLike|None=None) -> Self:
+    """
+    Creates an empty tensor with the given shape.
+
+    You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor.empty(2, 3)
+    print(t.shape)
+    ```
+    """
+    from tinygrad.uop.ops import UOp, to_max_shape
+    from tinygrad.device import canonicalize_device
+    dt = to_dtype(dtype) if dtype is not None else dtypes.default_float
+    new_shape = argfix(*shape)
+    max_shape = to_max_shape(new_shape)
+    u = UOp.new_buffer(canonicalize_device(device), prod(max_shape), dt).reshape(max_shape).shrink_to(new_shape)
+    return cls._wrap_uop(u)
+
   def empty_like(self, dtype: DTypeLike|None=None, device: str|tuple[str, ...]|None=None) -> Self:
     """
     Creates an empty tensor with the same shape as `self`.

tinygrad/mixin/dtype.py

@@ -9,7 +9,8 @@ class DTypeMixin:
   def dtype(self) -> DType: raise NotImplementedError
   @property
   def _uop(self) -> 'UOp': raise NotImplementedError
-  def _wrap_uop(self, u:'UOp') -> Self: raise NotImplementedError
+  @classmethod
+  def _wrap_uop(cls, u:'UOp') -> Self: raise NotImplementedError
 
   def cast(self, dtype:DTypeLike) -> Self:
     """

tinygrad/mixin/elementwise.py

@@ -546,7 +546,7 @@ class ElementwiseMixin(CreationMixin):
     """
     base, exponent = self._broadcasted(x, reverse=reverse)
     # TODO: int pow
-    if not base.is_floating_point() and not isinstance(x, ElementwiseMixin) and not (isinstance(x, int) and x >= 0):
+    if not base.is_floating_point() and isinstance(x, ConstType) and not (isinstance(x, int) and x >= 0):
       raise RuntimeError("base needs to be float")
     ret = base.alu(Ops.POW, exponent)
     # NOTE: pow(int, float) -> int
@@ -631,6 +631,7 @@ class ElementwiseMixin(CreationMixin):
     print(Tensor([float('nan')]).isclose(Tensor([float('nan')]), equal_nan=True).numpy())
     ```
     """
+    other = self.ufix(other)
     is_finite_close = self.isfinite() & other.isfinite() & ((self - other).abs() <= atol + rtol * other.abs())
     is_infinite_close = (self.isinf() | other.isinf()) & self.eq(other)
     is_nan_close = (self.isnan() & other.isnan()) & equal_nan
@@ -1071,7 +1072,7 @@ class ElementwiseMixin(CreationMixin):
     print(Tensor([1., 2., 3.]).lerp(Tensor([4., 5., 6.]), 0.5).numpy())
     ```
     """
-    if self.dtype == dtypes.uint8 and isinstance(weight, ElementwiseMixin):
+    if self.dtype == dtypes.uint8 and not isinstance(weight, ConstType):
       w_i = (weight * (1<<(W_PREC:=7)) + 0.5).cast(dtypes.int16)
       return (self+(((end - self).cast(dtypes.int8) * w_i + (1<<W_PREC-1)).cast(dtypes.uint16) >> W_PREC)).cast(dtypes.uint8)
     return self + (end - self) * weight

tinygrad/mixin/gradient.py

@@ -5,7 +5,14 @@ from tinygrad.helpers import argsort
 from tinygrad.dtype import sum_acc_dtype
 
 def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
-  def broadcast_to_input(x): return x.reshape(x.shape+(1,)*(len(ret.src[0].shape)-len(x.shape))).expand(ret.src[0].shape)
+  def broadcast_to_input(x):
+    shape, j = [], 0
+    for i in range(len(ret.src[0].shape)):
+      if i in ret.arg[1]: shape.append(1)
+      else:
+        shape.append(x.shape[j])
+        j += 1
+    return x.reshape(tuple(shape)).expand(ret.src[0].shape)
   if op == Ops.ADD: return (broadcast_to_input(ctx),)
   if op == Ops.MAX:
     assert ret.op is Ops.REDUCE, "only works on REDUCE"
@@ -69,7 +76,7 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.RESHAPE, name="ret"), lambda ctx, ret: (ctx.reshape(ret.src[0].shape), None)),
   (UPat(Ops.EXPAND, name="ret"), lambda ctx, ret:
     (ctx.cast(sum_acc_dtype(ctx.dtype))._rop(Ops.ADD, tuple(i for i,(s,n) in enumerate(zip(ret.src[0].shape, ret.shape)) if s!=n))
-     .cast(ctx.dtype), None)),
+     .reshape(ret.src[0].shape).cast(ctx.dtype), None)),
   (UPat(Ops.PAD, name="ret"), lambda ctx, ret: (ctx.shrink(tuple([(p[0], s+p[0]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
   (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[0], s-p[0]-p[1]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
   (UPat(Ops.PERMUTE, name="ret"), lambda ctx, ret: (ctx.permute(argsort(ret.marg)),)),

tinygrad/mixin/reduce.py

@@ -14,7 +14,7 @@ class ReduceMixin(DTypeMixin, MovementMixin):
     axis = tuple(self._resolve_dim(x) for x in (range(self.ndim) if axis is None else make_tuple(axis, 1)))
     if self.ndim == 0: axis = ()
     ret = self._rop(op, axis)
-    return ret if keepdim else ret.reshape(tuple(s for i,s in enumerate(self.shape) if i not in axis))
+    return ret.reshape(tuple(1 if i in axis else s for i,s in enumerate(self.shape))) if keepdim else ret
 
   def sum(self, axis:int|Sequence[int]|None=None, keepdim=False, dtype:DTypeLike|None=None) -> Self:
     """

tinygrad/nn/__init__.py

@@ -335,13 +335,14 @@ def _embedding_bwd(grad_emb:UOp, call:UOp) -> tuple:
 
     embed_size = grad_weight.shape[-1]
     BLOCK_J = min(256, embed_size)
-    assert embed_size % BLOCK_J == 0, f"embed_size {embed_size} must be divisible by {BLOCK_J}"
-
-    n_j_blocks = embed_size // BLOCK_J
+    n_j_blocks = (embed_size + BLOCK_J - 1) // BLOCK_J
     i = UOp.range(grad_emb_flat.shape[0], 0)         # batch_size * sequence_length -> GLOBAL
     j_inner = UOp.range(BLOCK_J, 2, AxisType.LOOP if device in ("CPU", "NULL") else AxisType.LOCAL)  # BLOCK_J threads per workgroup
     j_outer = UOp.range(n_j_blocks, 1)
     j = j_outer * BLOCK_J + j_inner
+    # mask padded embed
+    j_ok = j < embed_size
+    j_idx = j.clip(0, embed_size-1)
 
     if is_vocab_sharded:
       # each device owns [offset, offset+local_vocab_size) of the global vocabulary
@@ -349,16 +350,16 @@ def _embedding_bwd(grad_emb:UOp, call:UOp) -> tuple:
       offset = dnum * local_vocab_size
       global_token_id = idx_flat[i].cast(dtypes.weakint)
       local_token_id = (global_token_id - offset).clip(0, grad_weight.shape[0]-1)
-      in_range = (global_token_id >= offset) & (global_token_id < (offset + local_vocab_size))
-      grad_val = in_range.where(grad_emb_flat[i, j].load().cast(dtypes.float), 0.0)
+      in_range = (global_token_id >= offset) & (global_token_id < (offset + local_vocab_size)) & j_ok
+      grad_val = in_range.where(grad_emb_flat[i, j_idx].load().cast(dtypes.float), 0.0)
     else:
       local_token_id = idx_flat[i].clip(0, grad_weight.shape[0]-1).cast(dtypes.weakint)
-      grad_val = grad_emb_flat[i, j].load().cast(dtypes.float)
+      grad_val = j_ok.where(grad_emb_flat[i, j_idx].load().cast(dtypes.float), 0.0)
     # atomic scatter-add: grad_weight[token_id, j] += grad_emb_flat[i, j]
     if device in ("CPU", "NULL"): atomic_arg = "__atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED);"
     elif device == "AMD": atomic_arg = "__hip_atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);"
     else: raise NotImplementedError(f"no atomics for device {device}")
-    atomic = UOp(Ops.CUSTOM, dtypes.void, (grad_weight.index(local_token_id, j, ptr=True), grad_val), arg = atomic_arg)
+    atomic = UOp(Ops.CUSTOM, dtypes.void, (grad_weight.index(local_token_id, j_idx, ptr=True), grad_val), arg = atomic_arg)
     return atomic.end(i, j_outer, j_inner).sink(arg=KernelInfo(name="embedding_bwd", opts_to_apply=()))
 
   grad_weight_uop = grad_weight_uop.custom_kernel(grad_emb, idx, fxn=_embedding_bwd_kernel)[0]

tinygrad/nn/state.py

@@ -1,9 +1,9 @@
-import json, pathlib, zipfile, pickle, tarfile, struct, functools, io, zlib
+import json, math, pathlib, zipfile, pickle, tarfile, struct, functools, io, zlib
 from collections import OrderedDict
 from typing import Any, Callable, BinaryIO, Iterable, cast
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import prod, argsort, DEBUG, Timing, GlobalCounters, tqdm, round_up, T, strides_for_shape
+from tinygrad.helpers import prod, argsort, DEBUG, Timing, GlobalCounters, tqdm, round_up, T, strides_for_shape, CHUNK_SIZE
 
 class TensorIO(io.RawIOBase, BinaryIO):
   def __init__(self, t: Tensor):
@@ -82,6 +82,60 @@ def safe_save(tensors:dict[str, Tensor], fn:str, metadata:dict[str, Any]|None=No
   t[8:8+len(j)].assign(list(j.encode('utf-8')))
   for k,v in safe_load(t).items(): v.assign(tensors[k])
 
+# tinyfs
+
+def fs_store(t:Tensor) -> Tensor:
+  """
+  Store a tensor to storage.
+  """
+  # TODO: this should work locally as well
+  data = t.contiguous().flatten().bitcast(dtypes.uint8)
+
+  # pad to a multiple of 1mb
+  if (tsize := data.shape[0]) % CHUNK_SIZE != 0: data = data.pad((0, CHUNK_SIZE - tsize % CHUNK_SIZE))
+  size = data.shape[0]
+
+  base_chunks = math.ceil(size / CHUNK_SIZE)
+  tree_depth = math.ceil(math.log(base_chunks, CHUNK_SIZE // 16))
+
+  to_device = "CPU" if isinstance(t.device, str) and t.device.startswith("DISK") else t.device
+
+  level_chunks = base_chunks
+  for _ in range(tree_depth + 1):
+    data = data.to("tinyfs:store")[:level_chunks * 16].contiguous().to(to_device)
+    if (tsize := data.shape[0]) % CHUNK_SIZE != 0: data = data.pad((0, CHUNK_SIZE - tsize % CHUNK_SIZE))
+    level_chunks = math.ceil(data.shape[0] / CHUNK_SIZE)
+
+  return data[:16].contiguous()
+
+def fs_load(t:Tensor, size:int) -> Tensor:
+  """
+  Load a tensor from storage.
+
+  t should be a tensor of the hash to load
+  """
+  # TODO: this should work locally as well
+  assert t.dtype == dtypes.uint8, "hash is expected to be uint8"
+  h = t.contiguous().flatten()
+  assert h.shape[0] == 16, "expected hash"
+
+  base_chunks = math.ceil(size / CHUNK_SIZE)
+  tree_depth = math.ceil(math.log(base_chunks, CHUNK_SIZE // 16))
+  data, level_chunks = h, 0
+  for i in reversed(range(tree_depth + 1)):
+    data = data.to("tinyfs:load")
+
+    # if not last level, its still hashes
+    if i > 0 or tree_depth == 0:
+      level_chunks = max(1, math.ceil(base_chunks / (CHUNK_SIZE // 16)**(i-1)))
+      pad_amt = 16 * level_chunks
+    else: pad_amt = CHUNK_SIZE * level_chunks
+    if (tsize := data.shape[0]) < pad_amt: data = data.pad((0, pad_amt - tsize))
+    data = data[:pad_amt].contiguous()
+    if i != 0: data = data.to(t.device)
+
+  return data[:size]
+
 # state dict
 
 def get_state_dict(obj, prefix:str='', tensor_type=Tensor) -> dict[str, Tensor]:

tinygrad/renderer/amd/elf.py

@@ -39,7 +39,7 @@ def assemble_linear(prg:UOp, lin:UOp, arch:str) -> bytes:
   for u in sink.toposort():
     if u.op is Ops.PARAM and u.addrspace is AddrSpace.ALU: n_vars += 1
     elif u.op is Ops.PARAM: n_bufs += 1
-    elif u.op is Ops.BUFFER and u.addrspace is AddrSpace.LOCAL: lds_size += u.ptrdtype.size * u.ptrdtype.base.itemsize
+    elif u.op is Ops.BUFFER and u.addrspace is AddrSpace.LOCAL: lds_size += u.max_numel() * u.dtype.itemsize
     elif u.op is Ops.SPECIAL and u.arg.startswith("gidx"): gids.add(int(u.arg[-1]))
   code_bytes = b"".join(inst.to_bytes() for inst in insts)
   arch = next(v for k, v in _arch_map.items() if arch.startswith(k))

tinygrad/renderer/cstyle.py

@@ -34,8 +34,8 @@ base_rewrite = PatternMatcher([
   (UPat(Ops.CONST, arg=-math.inf, name="x"), lambda ctx, x: f"({ctx.render_cast(x, f'-{ctx.infinity}')})"),
   (UPat(Ops.CONST, dtype=dtypes.floats, name="x"), lambda ctx,x: f"({ctx.render_cast(x, ctx.nan)})" if math.isnan(x.arg) else None),
   (UPat(Ops.CONST, dtype=dtypes.float, name="x"), lambda ctx,x: f"{x.arg}f"),
-  (UPat(Ops.CONST, dtype=dtypes.int64, name="x"), lambda ctx,x: f"{x.arg}ll"),
-  (UPat(Ops.CONST, dtype=dtypes.uint64, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}ull"),
+  (UPat(Ops.CONST, dtype=dtypes.int64, name="x"), lambda ctx,x: f"{x.arg}l"),
+  (UPat(Ops.CONST, dtype=dtypes.uint64, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}ul"),
   (UPat(Ops.CONST, dtype=dtypes.uint32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}u"),
   (UPat(Ops.CONST, dtype=dtypes.bool, name="x"), lambda ctx,x: "1" if x.arg else "0"),
   # consts are rendered to larger type and casted

tinygrad/renderer/isa/x86.py

@@ -387,6 +387,8 @@ dt_128bit = tuple(dt.vec(l) for dt in dts for l in [16,8,4,2,1] if l*dt.itemsize
 
 isel_matcher = PatternMatcher([
   # **** Op -> Op ****
+  # materialize the structural width of a STACK into a vec dtype
+  (UPat(Ops.STACK, name="x"), lambda x: x.replace(dtype=x.dtype.scalar().vec(len(x.src))) if 1 < len(x.src) != x.dtype.count else None),
   # cast of void is a noop
   (UPat.var("y").cast(name="x"), lambda y,x: y if y.dtype == dtypes.void else None),
   # extracting the 0th float element is a noop as it just moves the 0th element from one xmm register to another

tinygrad/renderer/llvmir.py

@@ -218,9 +218,9 @@ class AMDLLVMRenderer(LLVMRenderer):
   ]) + base_rewrite
   extra_matcher = LLVMRenderer.extra_matcher + create_non_native_float_pats(dtypes.fp8s) + PatternMatcher([
     (UPat(Ops.CAST, dtype=dtypes.half.vec(16), src=UPat.var("y", dtypes.half.vec(8))),
-      lambda y: UOp(Ops.STACK, dtypes.half.vec(16), tuple(y.gep(i // 2) if i % 2 == 0 else UOp.const(dtypes.half, 0.0) for i in range(16)))),
+      lambda y: UOp(Ops.STACK, dtypes.half.vec(16), tuple(y.index(i // 2) if i % 2 == 0 else UOp.const(dtypes.half, 0.0) for i in range(16)))),
     (UPat(Ops.CAST, dtype=dtypes.half.vec(8), src=UPat.var("y", dtypes.half.vec(16))),
-      lambda y: UOp(Ops.STACK, dtypes.half.vec(8), tuple(y.gep(i * 2) for i in range(8)))),
+      lambda y: UOp(Ops.STACK, dtypes.half.vec(8), tuple(y.index(i * 2) for i in range(8)))),
     # amd llvm intrinsics llvm.log2/llvm.exp2 don't support double
     (UPat(Ops.LOG2, dtype=dtypes.double, src=(UPat.var("d"),)), xlog2),
     (UPat(Ops.EXP2, dtype=dtypes.double, src=(UPat.var("d"),)), xexp2),

tinygrad/renderer/nir.py

@@ -4,9 +4,9 @@ from tinygrad.helpers import DEBUG, OSX, unwrap, fromimport, Target
 from tinygrad.renderer import Renderer
 from tinygrad.renderer.cstyle import CUDARenderer, OpenCLRenderer
 from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str
-from tinygrad.runtime.autogen import mesa
+from tinygrad.runtime.autogen import mesa, libc
 from tinygrad.runtime.support.c import POINTER
-import base64, ctypes, ctypes.util, struct, functools, inspect, itertools
+import base64, ctypes, struct, functools, inspect, itertools
 
 def g(s:str): return getattr(mesa, s)
 def nsrc(d:mesa.nir_def) -> mesa.nir_src: return mesa.nir_src(ssa=ctypes.pointer(d))
@@ -222,8 +222,7 @@ class NIRRenderer(Renderer):
     self.postrender(uops)
 
     mesa.nir_validate_shader(self.b.shader, b"after render")
-    if DEBUG >= 4: mesa.nir_print_shader(self.b.shader, ctypes.POINTER(mesa.struct__IO_FILE).in_dll(ctypes.CDLL(ctypes.util.find_library('c')),
-                                                                                                    "__stdoutp" if OSX else "stdout"))
+    if DEBUG >= 4: mesa.nir_print_shader(self.b.shader, ctypes.POINTER(mesa.struct__IO_FILE).in_dll(libc.dll, "__stdoutp" if OSX else "stdout"))
     mesa.nir_serialize(blob:=mesa.struct_blob(), self.b.shader, False)
     ret = base64.b64encode(ctypes.string_at(blob.data, blob.size)).decode()
 

tinygrad/runtime/autogen/__init__.py

@@ -22,6 +22,7 @@ llvm_lib = (
 clang_lib = win_llvm.replace("LLVM-C", "libclang") + (mac_llvm + other_llvm).replace("LLVM", "clang")
 
 webgpu_lib = "os.path.join(sysconfig.get_paths()['purelib'], 'pydawn', 'lib', 'libwebgpu_dawn.dll') if WIN else 'webgpu_dawn'"
+tinymesa_path = "os.path.join(sysconfig.get_paths()['platlib'], 'tinymesa')"
 nv_lib_path = ("[f'/{pre}/cuda/targets/{tgt}/lib' for pre in ['opt', 'usr/local'] for tgt in "
                "[sysconfig.get_config_vars().get(\"MULTIARCH\", \"\").rsplit(\"-\", 1)[0], 'sbsa-linux']]")
 
@@ -51,8 +52,8 @@ def __getattr__(nm):
   match nm:
     case "libc":
       return load("libc", lambda: ([i for i in system("dpkg -L libc6-dev").split() if 'sys/mman.h' in i or 'bits/mman-shared.h' in i] +
-                                   ["/usr/include/string.h", "/usr/include/elf.h", "/usr/include/unistd.h", "/usr/include/asm-generic/mman-common.h"]),
-                  args=["-D__USE_GNU", "-D_GNU_SOURCE"], dll="'c'", errno=True, recsym=True)
+                                   ["/usr/include/string.h", "/usr/include/elf.h", "/usr/include/unistd.h", "/usr/include/stdio.h", "/usr/include/asm-generic/mman-common.h"]),
+                  args=["-D__USE_GNU", "-D_GNU_SOURCE"], dll="'c'", errno=True, recsym=True, rules=[(r'([a-z]+) = \1', '')]) # removes stdin = stdin
     case "avcodec": return load("avcodec", ["{}/libavcodec/hevc/hevc.h", "{}/libavcodec/cbs_h265.h"], srcs=ffmpeg_src)
     case "opencl": return load("opencl", ["{}/CL/cl.h"], dll="'OpenCL'", args=["-I{}"], srcs=opencl_src)
     case "cuda": return load("cuda", ["{}/include/cuda.h"], dll="'nvcuda' if WIN else 'cuda'", args=["-D__CUDA_API_VERSION_INTERNAL"], srcs=cudart_src, macros=False, prolog=["from tinygrad.helpers import WIN"])
@@ -153,10 +154,8 @@ def __getattr__(nm):
           *[f"python3 src/compiler/{s}_h.py > gen/{s.split('/')[-1]}.h" for s in ["nir/nir_opcodes", "nir/nir_builder_opcodes"]],
           *[f"python3 src/compiler/nir/nir_{s}_h.py --outdir gen" for s in ["intrinsics", "intrinsics_indices"]]]), cwd=path, shell=True, check=True),
   srcs="https://gitlab.freedesktop.org/mesa/mesa/-/archive/mesa-25.2.7/mesa-25.2.7.tar.gz",
-  dll="'tinymesa_cpu' if (_cpu:=DEV.renderer == 'LVP') else 'tinymesa', " \
-      'emsg="not available on this platform" if WIN or (OSX and (platform.machine() != "arm64" or (_mv:=platform.mac_ver()[0][:2]) not in {"14","15","26"})) or (platform.system() == "Linux" and platform.machine() not in {"x86_64", "aarch64"}) else ' \
-      'f"run `sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa{\'_cpu\'*_cpu}-mesa-25.2.7-{\'macos-\'+_mv if OSX else \'linux\'}-{\'amd64\' if ARCH_X86 else \'arm64\'}.{\'dylib\' if OSX else \'so\'} -o /usr/local/lib/libtinymesa{\'_cpu\'*_cpu}.{\'dylib\' if OSX else \'so\'}`"',
-  prolog=["from tinygrad.helpers import DEV, ARCH_X86, WIN, OSX", "import gzip, base64, platform"],
+  dll=f"'tinymesa_cpu' if DEV.renderer == 'LVP' else 'tinymesa', {tinymesa_path}, emsg='pip install tinymesa==25.2.7.2'",
+  prolog=["from tinygrad.helpers import DEV", "import gzip, base64, platform, sysconfig, os"],
   epilog=lambda path: [system(f"{root}/extra/mesa/lvp_nir_options.sh {path}")])
     case "libclang":
       return load("libclang",

tinygrad/runtime/autogen/libc.py

@@ -896,6 +896,304 @@ def getentropy(__buffer:ctypes.c_void_p, __length:size_t) -> int: ...
 def close_range(__fd:int, __max_fd:int, __flags:int) -> int: ...
 @dll.bind(ctypes.c_int32)
 def gettid() -> ctypes.c_int32: ...
+@c.record
+class struct___va_list_tag(c.Struct):
+  SIZE = 24
+  gp_offset: int
+  fp_offset: int
+  overflow_arg_area: ctypes.c_void_p
+  reg_save_area: ctypes.c_void_p
+struct___va_list_tag.register_fields([('gp_offset', ctypes.c_uint32, 0), ('fp_offset', ctypes.c_uint32, 4), ('overflow_arg_area', ctypes.c_void_p, 8), ('reg_save_area', ctypes.c_void_p, 16)])
+va_list: TypeAlias = c.Array[struct___va_list_tag, Literal[1]]
+@c.record
+class struct__G_fpos_t(c.Struct):
+  SIZE = 16
+  __pos: int
+  __state: __mbstate_t
+fpos_t: TypeAlias = struct__G_fpos_t
+@c.record
+class __mbstate_t(c.Struct):
+  SIZE = 8
+  __count: int
+  __value: __mbstate_t___value
+@c.record
+class __mbstate_t___value(c.Struct):
+  SIZE = 4
+  __wch: int
+  __wchb: c.Array[ctypes.c_char, Literal[4]]
+__mbstate_t___value.register_fields([('__wch', ctypes.c_uint32, 0), ('__wchb', c.Array[ctypes.c_char, Literal[4]], 0)])
+__mbstate_t.register_fields([('__count', ctypes.c_int32, 0), ('__value', __mbstate_t___value, 4)])
+struct__G_fpos_t.register_fields([('__pos', ctypes.c_int64, 0), ('__state', __mbstate_t, 8)])
+@c.record
+class struct__G_fpos64_t(c.Struct):
+  SIZE = 16
+  __pos: int
+  __state: __mbstate_t
+fpos64_t: TypeAlias = struct__G_fpos64_t
+struct__G_fpos64_t.register_fields([('__pos', ctypes.c_int64, 0), ('__state', __mbstate_t, 8)])
+@c.record
+class struct__IO_FILE(c.Struct):
+  SIZE = 216
+  _flags: int
+  _IO_read_ptr: c.POINTER[ctypes.c_char]
+  _IO_read_end: c.POINTER[ctypes.c_char]
+  _IO_read_base: c.POINTER[ctypes.c_char]
+  _IO_write_base: c.POINTER[ctypes.c_char]
+  _IO_write_ptr: c.POINTER[ctypes.c_char]
+  _IO_write_end: c.POINTER[ctypes.c_char]
+  _IO_buf_base: c.POINTER[ctypes.c_char]
+  _IO_buf_end: c.POINTER[ctypes.c_char]
+  _IO_save_base: c.POINTER[ctypes.c_char]
+  _IO_backup_base: c.POINTER[ctypes.c_char]
+  _IO_save_end: c.POINTER[ctypes.c_char]
+  _markers: c.POINTER[struct__IO_marker]
+  _chain: c.POINTER[struct__IO_FILE]
+  _fileno: int
+  _flags2: int
+  _old_offset: int
+  _cur_column: int
+  _vtable_offset: int
+  _shortbuf: c.Array[ctypes.c_char, Literal[1]]
+  _lock: ctypes.c_void_p
+  _offset: int
+  _codecvt: c.POINTER[struct__IO_codecvt]
+  _wide_data: c.POINTER[struct__IO_wide_data]
+  _freeres_list: c.POINTER[struct__IO_FILE]
+  _freeres_buf: ctypes.c_void_p
+  __pad5: int
+  _mode: int
+  _unused2: c.Array[ctypes.c_char, Literal[20]]
+FILE: TypeAlias = struct__IO_FILE
+class struct__IO_marker(c.Struct): pass
+_IO_lock_t: TypeAlias = None
+class struct__IO_codecvt(c.Struct): pass
+class struct__IO_wide_data(c.Struct): pass
+struct__IO_FILE.register_fields([('_flags', ctypes.c_int32, 0), ('_IO_read_ptr', c.POINTER[ctypes.c_char], 8), ('_IO_read_end', c.POINTER[ctypes.c_char], 16), ('_IO_read_base', c.POINTER[ctypes.c_char], 24), ('_IO_write_base', c.POINTER[ctypes.c_char], 32), ('_IO_write_ptr', c.POINTER[ctypes.c_char], 40), ('_IO_write_end', c.POINTER[ctypes.c_char], 48), ('_IO_buf_base', c.POINTER[ctypes.c_char], 56), ('_IO_buf_end', c.POINTER[ctypes.c_char], 64), ('_IO_save_base', c.POINTER[ctypes.c_char], 72), ('_IO_backup_base', c.POINTER[ctypes.c_char], 80), ('_IO_save_end', c.POINTER[ctypes.c_char], 88), ('_markers', c.POINTER[struct__IO_marker], 96), ('_chain', c.POINTER[struct__IO_FILE], 104), ('_fileno', ctypes.c_int32, 112), ('_flags2', ctypes.c_int32, 116), ('_old_offset', ctypes.c_int64, 120), ('_cur_column', ctypes.c_uint16, 128), ('_vtable_offset', ctypes.c_byte, 130), ('_shortbuf', c.Array[ctypes.c_char, Literal[1]], 131), ('_lock', c.POINTER[_IO_lock_t], 136), ('_offset', ctypes.c_int64, 144), ('_codecvt', c.POINTER[struct__IO_codecvt], 152), ('_wide_data', c.POINTER[struct__IO_wide_data], 160), ('_freeres_list', c.POINTER[struct__IO_FILE], 168), ('_freeres_buf', ctypes.c_void_p, 176), ('__pad5', size_t, 184), ('_mode', ctypes.c_int32, 192), ('_unused2', c.Array[ctypes.c_char, Literal[20]], 196)])
+try: stdin = c.POINTER[FILE].in_dll(dll, 'stdin') # type: ignore
+except (ValueError,AttributeError): pass
+try: stdout = c.POINTER[FILE].in_dll(dll, 'stdout') # type: ignore
+except (ValueError,AttributeError): pass
+try: stderr = c.POINTER[FILE].in_dll(dll, 'stderr') # type: ignore
+except (ValueError,AttributeError): pass
+@dll.bind(ctypes.c_int32, c.POINTER[ctypes.c_char])
+def remove(__filename:c.POINTER[ctypes.c_char]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[ctypes.c_char], c.POINTER[ctypes.c_char])
+def rename(__old:c.POINTER[ctypes.c_char], __new:c.POINTER[ctypes.c_char]) -> int: ...
+@dll.bind(ctypes.c_int32, ctypes.c_int32, c.POINTER[ctypes.c_char], ctypes.c_int32, c.POINTER[ctypes.c_char])
+def renameat(__oldfd:int, __old:c.POINTER[ctypes.c_char], __newfd:int, __new:c.POINTER[ctypes.c_char]) -> int: ...
+@dll.bind(ctypes.c_int32, ctypes.c_int32, c.POINTER[ctypes.c_char], ctypes.c_int32, c.POINTER[ctypes.c_char], ctypes.c_uint32)
+def renameat2(__oldfd:int, __old:c.POINTER[ctypes.c_char], __newfd:int, __new:c.POINTER[ctypes.c_char], __flags:int) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE])
+def fclose(__stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(c.POINTER[FILE])
+def tmpfile() -> c.POINTER[FILE]: ...
+@dll.bind(c.POINTER[FILE])
+def tmpfile64() -> c.POINTER[FILE]: ...
+@dll.bind(c.POINTER[ctypes.c_char], c.Array[ctypes.c_char, Literal[20]])
+def tmpnam(_0:c.Array[ctypes.c_char, Literal[20]]) -> c.POINTER[ctypes.c_char]: ...
+@dll.bind(c.POINTER[ctypes.c_char], c.Array[ctypes.c_char, Literal[20]])
+def tmpnam_r(__s:c.Array[ctypes.c_char, Literal[20]]) -> c.POINTER[ctypes.c_char]: ...
+@dll.bind(c.POINTER[ctypes.c_char], c.POINTER[ctypes.c_char], c.POINTER[ctypes.c_char])
+def tempnam(__dir:c.POINTER[ctypes.c_char], __pfx:c.POINTER[ctypes.c_char]) -> c.POINTER[ctypes.c_char]: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE])
+def fflush(__stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE])
+def fflush_unlocked(__stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_int32)
+def fcloseall() -> int: ...
+@dll.bind(c.POINTER[FILE], c.POINTER[ctypes.c_char], c.POINTER[ctypes.c_char])
+def fopen(__filename:c.POINTER[ctypes.c_char], __modes:c.POINTER[ctypes.c_char]) -> c.POINTER[FILE]: ...
+@dll.bind(c.POINTER[FILE], c.POINTER[ctypes.c_char], c.POINTER[ctypes.c_char], c.POINTER[FILE])
+def freopen(__filename:c.POINTER[ctypes.c_char], __modes:c.POINTER[ctypes.c_char], __stream:c.POINTER[FILE]) -> c.POINTER[FILE]: ...
+@dll.bind(c.POINTER[FILE], c.POINTER[ctypes.c_char], c.POINTER[ctypes.c_char])
+def fopen64(__filename:c.POINTER[ctypes.c_char], __modes:c.POINTER[ctypes.c_char]) -> c.POINTER[FILE]: ...
+@dll.bind(c.POINTER[FILE], c.POINTER[ctypes.c_char], c.POINTER[ctypes.c_char], c.POINTER[FILE])
+def freopen64(__filename:c.POINTER[ctypes.c_char], __modes:c.POINTER[ctypes.c_char], __stream:c.POINTER[FILE]) -> c.POINTER[FILE]: ...
+@dll.bind(c.POINTER[FILE], ctypes.c_int32, c.POINTER[ctypes.c_char])
+def fdopen(__fd:int, __modes:c.POINTER[ctypes.c_char]) -> c.POINTER[FILE]: ...
+@c.record
+class struct__IO_cookie_io_functions_t(c.Struct):
+  SIZE = 32
+  read: c.CFUNCTYPE[ctypes.c_int64, [ctypes.c_void_p, c.POINTER[ctypes.c_char], ctypes.c_uint64]]
+  write: c.CFUNCTYPE[ctypes.c_int64, [ctypes.c_void_p, c.POINTER[ctypes.c_char], ctypes.c_uint64]]
+  seek: c.CFUNCTYPE[ctypes.c_int32, [ctypes.c_void_p, c.POINTER[ctypes.c_int64], ctypes.c_int32]]
+  close: c.CFUNCTYPE[ctypes.c_int32, [ctypes.c_void_p]]
+cookie_io_functions_t: TypeAlias = struct__IO_cookie_io_functions_t
+cookie_read_function_t: TypeAlias = c.CFUNCTYPE[ctypes.c_int64, [ctypes.c_void_p, c.POINTER[ctypes.c_char], ctypes.c_uint64]]
+cookie_write_function_t: TypeAlias = c.CFUNCTYPE[ctypes.c_int64, [ctypes.c_void_p, c.POINTER[ctypes.c_char], ctypes.c_uint64]]
+cookie_seek_function_t: TypeAlias = c.CFUNCTYPE[ctypes.c_int32, [ctypes.c_void_p, c.POINTER[ctypes.c_int64], ctypes.c_int32]]
+cookie_close_function_t: TypeAlias = c.CFUNCTYPE[ctypes.c_int32, [ctypes.c_void_p]]
+struct__IO_cookie_io_functions_t.register_fields([('read', c.POINTER[cookie_read_function_t], 0), ('write', c.POINTER[cookie_write_function_t], 8), ('seek', c.POINTER[cookie_seek_function_t], 16), ('close', c.POINTER[cookie_close_function_t], 24)])
+@dll.bind(c.POINTER[FILE], ctypes.c_void_p, c.POINTER[ctypes.c_char], cookie_io_functions_t)
+def fopencookie(__magic_cookie:ctypes.c_void_p, __modes:c.POINTER[ctypes.c_char], __io_funcs:cookie_io_functions_t) -> c.POINTER[FILE]: ...
+@dll.bind(c.POINTER[FILE], ctypes.c_void_p, size_t, c.POINTER[ctypes.c_char])
+def fmemopen(__s:ctypes.c_void_p, __len:size_t, __modes:c.POINTER[ctypes.c_char]) -> c.POINTER[FILE]: ...
+@dll.bind(c.POINTER[FILE], c.POINTER[c.POINTER[ctypes.c_char]], c.POINTER[size_t])
+def open_memstream(__bufloc:c.POINTER[c.POINTER[ctypes.c_char]], __sizeloc:c.POINTER[size_t]) -> c.POINTER[FILE]: ...
+@dll.bind(None, c.POINTER[FILE], c.POINTER[ctypes.c_char])
+def setbuf(__stream:c.POINTER[FILE], __buf:c.POINTER[ctypes.c_char]) -> None: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE], c.POINTER[ctypes.c_char], ctypes.c_int32, size_t)
+def setvbuf(__stream:c.POINTER[FILE], __buf:c.POINTER[ctypes.c_char], __modes:int, __n:size_t) -> int: ...
+@dll.bind(None, c.POINTER[FILE], c.POINTER[ctypes.c_char], size_t)
+def setbuffer(__stream:c.POINTER[FILE], __buf:c.POINTER[ctypes.c_char], __size:size_t) -> None: ...
+@dll.bind(None, c.POINTER[FILE])
+def setlinebuf(__stream:c.POINTER[FILE]) -> None: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE], c.POINTER[ctypes.c_char])
+def fprintf(__stream:c.POINTER[FILE], __format:c.POINTER[ctypes.c_char]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[ctypes.c_char])
+def printf(__format:c.POINTER[ctypes.c_char]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[ctypes.c_char], c.POINTER[ctypes.c_char])
+def sprintf(__s:c.POINTER[ctypes.c_char], __format:c.POINTER[ctypes.c_char]) -> int: ...
+__gnuc_va_list: TypeAlias = c.Array[struct___va_list_tag, Literal[1]]
+@dll.bind(ctypes.c_int32, c.POINTER[FILE], c.POINTER[ctypes.c_char], c.Array[struct___va_list_tag, Literal[1]])
+def vfprintf(__s:c.POINTER[FILE], __format:c.POINTER[ctypes.c_char], __arg:c.Array[struct___va_list_tag, Literal[1]]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[ctypes.c_char], c.Array[struct___va_list_tag, Literal[1]])
+def vprintf(__format:c.POINTER[ctypes.c_char], __arg:c.Array[struct___va_list_tag, Literal[1]]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[ctypes.c_char], c.POINTER[ctypes.c_char], c.Array[struct___va_list_tag, Literal[1]])
+def vsprintf(__s:c.POINTER[ctypes.c_char], __format:c.POINTER[ctypes.c_char], __arg:c.Array[struct___va_list_tag, Literal[1]]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[ctypes.c_char], size_t, c.POINTER[ctypes.c_char])
+def snprintf(__s:c.POINTER[ctypes.c_char], __maxlen:size_t, __format:c.POINTER[ctypes.c_char]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[ctypes.c_char], size_t, c.POINTER[ctypes.c_char], c.Array[struct___va_list_tag, Literal[1]])
+def vsnprintf(__s:c.POINTER[ctypes.c_char], __maxlen:size_t, __format:c.POINTER[ctypes.c_char], __arg:c.Array[struct___va_list_tag, Literal[1]]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[c.POINTER[ctypes.c_char]], c.POINTER[ctypes.c_char], c.Array[struct___va_list_tag, Literal[1]])
+def vasprintf(__ptr:c.POINTER[c.POINTER[ctypes.c_char]], __f:c.POINTER[ctypes.c_char], __arg:c.Array[struct___va_list_tag, Literal[1]]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[c.POINTER[ctypes.c_char]], c.POINTER[ctypes.c_char])
+def __asprintf(__ptr:c.POINTER[c.POINTER[ctypes.c_char]], __fmt:c.POINTER[ctypes.c_char]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[c.POINTER[ctypes.c_char]], c.POINTER[ctypes.c_char])
+def asprintf(__ptr:c.POINTER[c.POINTER[ctypes.c_char]], __fmt:c.POINTER[ctypes.c_char]) -> int: ...
+@dll.bind(ctypes.c_int32, ctypes.c_int32, c.POINTER[ctypes.c_char], c.Array[struct___va_list_tag, Literal[1]])
+def vdprintf(__fd:int, __fmt:c.POINTER[ctypes.c_char], __arg:c.Array[struct___va_list_tag, Literal[1]]) -> int: ...
+@dll.bind(ctypes.c_int32, ctypes.c_int32, c.POINTER[ctypes.c_char])
+def dprintf(__fd:int, __fmt:c.POINTER[ctypes.c_char]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE], c.POINTER[ctypes.c_char])
+def fscanf(__stream:c.POINTER[FILE], __format:c.POINTER[ctypes.c_char]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[ctypes.c_char])
+def scanf(__format:c.POINTER[ctypes.c_char]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[ctypes.c_char], c.POINTER[ctypes.c_char])
+def sscanf(__s:c.POINTER[ctypes.c_char], __format:c.POINTER[ctypes.c_char]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE], c.POINTER[ctypes.c_char], c.Array[struct___va_list_tag, Literal[1]])
+def vfscanf(__s:c.POINTER[FILE], __format:c.POINTER[ctypes.c_char], __arg:c.Array[struct___va_list_tag, Literal[1]]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[ctypes.c_char], c.Array[struct___va_list_tag, Literal[1]])
+def vscanf(__format:c.POINTER[ctypes.c_char], __arg:c.Array[struct___va_list_tag, Literal[1]]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[ctypes.c_char], c.POINTER[ctypes.c_char], c.Array[struct___va_list_tag, Literal[1]])
+def vsscanf(__s:c.POINTER[ctypes.c_char], __format:c.POINTER[ctypes.c_char], __arg:c.Array[struct___va_list_tag, Literal[1]]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE])
+def fgetc(__stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE])
+def getc(__stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_int32)
+def getchar() -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE])
+def getc_unlocked(__stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_int32)
+def getchar_unlocked() -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE])
+def fgetc_unlocked(__stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_int32, ctypes.c_int32, c.POINTER[FILE])
+def fputc(__c:int, __stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_int32, ctypes.c_int32, c.POINTER[FILE])
+def putc(__c:int, __stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_int32, ctypes.c_int32)
+def putchar(__c:int) -> int: ...
+@dll.bind(ctypes.c_int32, ctypes.c_int32, c.POINTER[FILE])
+def fputc_unlocked(__c:int, __stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_int32, ctypes.c_int32, c.POINTER[FILE])
+def putc_unlocked(__c:int, __stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_int32, ctypes.c_int32)
+def putchar_unlocked(__c:int) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE])
+def getw(__stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_int32, ctypes.c_int32, c.POINTER[FILE])
+def putw(__w:int, __stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(c.POINTER[ctypes.c_char], c.POINTER[ctypes.c_char], ctypes.c_int32, c.POINTER[FILE])
+def fgets(__s:c.POINTER[ctypes.c_char], __n:int, __stream:c.POINTER[FILE]) -> c.POINTER[ctypes.c_char]: ...
+@dll.bind(c.POINTER[ctypes.c_char], c.POINTER[ctypes.c_char], ctypes.c_int32, c.POINTER[FILE])
+def fgets_unlocked(__s:c.POINTER[ctypes.c_char], __n:int, __stream:c.POINTER[FILE]) -> c.POINTER[ctypes.c_char]: ...
+@dll.bind(ctypes.c_int64, c.POINTER[c.POINTER[ctypes.c_char]], c.POINTER[size_t], ctypes.c_int32, c.POINTER[FILE])
+def __getdelim(__lineptr:c.POINTER[c.POINTER[ctypes.c_char]], __n:c.POINTER[size_t], __delimiter:int, __stream:c.POINTER[FILE]) -> ctypes.c_int64: ...
+@dll.bind(ctypes.c_int64, c.POINTER[c.POINTER[ctypes.c_char]], c.POINTER[size_t], ctypes.c_int32, c.POINTER[FILE])
+def getdelim(__lineptr:c.POINTER[c.POINTER[ctypes.c_char]], __n:c.POINTER[size_t], __delimiter:int, __stream:c.POINTER[FILE]) -> ctypes.c_int64: ...
+@dll.bind(ctypes.c_int64, c.POINTER[c.POINTER[ctypes.c_char]], c.POINTER[size_t], c.POINTER[FILE])
+def getline(__lineptr:c.POINTER[c.POINTER[ctypes.c_char]], __n:c.POINTER[size_t], __stream:c.POINTER[FILE]) -> ctypes.c_int64: ...
+@dll.bind(ctypes.c_int32, c.POINTER[ctypes.c_char], c.POINTER[FILE])
+def fputs(__s:c.POINTER[ctypes.c_char], __stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[ctypes.c_char])
+def puts(__s:c.POINTER[ctypes.c_char]) -> int: ...
+@dll.bind(ctypes.c_int32, ctypes.c_int32, c.POINTER[FILE])
+def ungetc(__c:int, __stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_uint64, ctypes.c_void_p, size_t, size_t, c.POINTER[FILE])
+def fread(__ptr:ctypes.c_void_p, __size:size_t, __n:size_t, __stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_uint64, ctypes.c_void_p, size_t, size_t, c.POINTER[FILE])
+def fwrite(__ptr:ctypes.c_void_p, __size:size_t, __n:size_t, __s:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[ctypes.c_char], c.POINTER[FILE])
+def fputs_unlocked(__s:c.POINTER[ctypes.c_char], __stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(size_t, ctypes.c_void_p, size_t, size_t, c.POINTER[FILE])
+def fread_unlocked(__ptr:ctypes.c_void_p, __size:size_t, __n:size_t, __stream:c.POINTER[FILE]) -> size_t: ...
+@dll.bind(size_t, ctypes.c_void_p, size_t, size_t, c.POINTER[FILE])
+def fwrite_unlocked(__ptr:ctypes.c_void_p, __size:size_t, __n:size_t, __stream:c.POINTER[FILE]) -> size_t: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE], ctypes.c_int64, ctypes.c_int32)
+def fseek(__stream:c.POINTER[FILE], __off:int, __whence:int) -> int: ...
+@dll.bind(ctypes.c_int64, c.POINTER[FILE])
+def ftell(__stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(None, c.POINTER[FILE])
+def rewind(__stream:c.POINTER[FILE]) -> None: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE], ctypes.c_int64, ctypes.c_int32)
+def fseeko(__stream:c.POINTER[FILE], __off:ctypes.c_int64, __whence:int) -> int: ...
+@dll.bind(ctypes.c_int64, c.POINTER[FILE])
+def ftello(__stream:c.POINTER[FILE]) -> ctypes.c_int64: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE], c.POINTER[fpos_t])
+def fgetpos(__stream:c.POINTER[FILE], __pos:c.POINTER[fpos_t]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE], c.POINTER[fpos_t])
+def fsetpos(__stream:c.POINTER[FILE], __pos:c.POINTER[fpos_t]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE], ctypes.c_int64, ctypes.c_int32)
+def fseeko64(__stream:c.POINTER[FILE], __off:ctypes.c_int64, __whence:int) -> int: ...
+@dll.bind(ctypes.c_int64, c.POINTER[FILE])
+def ftello64(__stream:c.POINTER[FILE]) -> ctypes.c_int64: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE], c.POINTER[fpos64_t])
+def fgetpos64(__stream:c.POINTER[FILE], __pos:c.POINTER[fpos64_t]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE], c.POINTER[fpos64_t])
+def fsetpos64(__stream:c.POINTER[FILE], __pos:c.POINTER[fpos64_t]) -> int: ...
+@dll.bind(None, c.POINTER[FILE])
+def clearerr(__stream:c.POINTER[FILE]) -> None: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE])
+def feof(__stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE])
+def ferror(__stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(None, c.POINTER[FILE])
+def clearerr_unlocked(__stream:c.POINTER[FILE]) -> None: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE])
+def feof_unlocked(__stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE])
+def ferror_unlocked(__stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(None, c.POINTER[ctypes.c_char])
+def perror(__s:c.POINTER[ctypes.c_char]) -> None: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE])
+def fileno(__stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE])
+def fileno_unlocked(__stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE])
+def pclose(__stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(c.POINTER[FILE], c.POINTER[ctypes.c_char], c.POINTER[ctypes.c_char])
+def popen(__command:c.POINTER[ctypes.c_char], __modes:c.POINTER[ctypes.c_char]) -> c.POINTER[FILE]: ...
+@dll.bind(c.POINTER[ctypes.c_char], c.POINTER[ctypes.c_char])
+def ctermid(__s:c.POINTER[ctypes.c_char]) -> c.POINTER[ctypes.c_char]: ...
+@dll.bind(c.POINTER[ctypes.c_char], c.POINTER[ctypes.c_char])
+def cuserid(__s:c.POINTER[ctypes.c_char]) -> c.POINTER[ctypes.c_char]: ...
+class struct_obstack(c.Struct): pass
+@dll.bind(ctypes.c_int32, c.POINTER[struct_obstack], c.POINTER[ctypes.c_char])
+def obstack_printf(__obstack:c.POINTER[struct_obstack], __format:c.POINTER[ctypes.c_char]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[struct_obstack], c.POINTER[ctypes.c_char], c.Array[struct___va_list_tag, Literal[1]])
+def obstack_vprintf(__obstack:c.POINTER[struct_obstack], __format:c.POINTER[ctypes.c_char], __args:c.Array[struct___va_list_tag, Literal[1]]) -> int: ...
+@dll.bind(None, c.POINTER[FILE])
+def flockfile(__stream:c.POINTER[FILE]) -> None: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE])
+def ftrylockfile(__stream:c.POINTER[FILE]) -> int: ...
+@dll.bind(None, c.POINTER[FILE])
+def funlockfile(__stream:c.POINTER[FILE]) -> None: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE])
+def __uflow(_0:c.POINTER[FILE]) -> int: ...
+@dll.bind(ctypes.c_int32, c.POINTER[FILE], ctypes.c_int32)
+def __overflow(_0:c.POINTER[FILE], _1:int) -> int: ...
 MREMAP_MAYMOVE = 1
 MREMAP_FIXED = 2
 MREMAP_DONTUNMAP = 4
@@ -4027,6 +4325,27 @@ F_ULOCK = 0
 F_LOCK = 1
 F_TLOCK = 2
 F_TEST = 3
+_STDIO_H = 1
+_IOFBF = 0
+_IOLBF = 1
+_IONBF = 2
+BUFSIZ = 8192
+EOF = (-1)
+SEEK_SET = 0
+SEEK_CUR = 1
+SEEK_END = 2
+SEEK_DATA = 3
+SEEK_HOLE = 4
+P_tmpdir = "/tmp"
+L_tmpnam = 20
+TMP_MAX = 238328
+L_ctermid = 9
+L_cuserid = 9
+FOPEN_MAX = 16
+_PRINTF_NAN_LEN_MAX = 4
+RENAME_NOREPLACE = (1 << 0)
+RENAME_EXCHANGE = (1 << 1)
+RENAME_WHITEOUT = (1 << 2)
 PROT_READ = 0x1
 PROT_WRITE = 0x2
 PROT_EXEC = 0x4

tinygrad/runtime/autogen/mesa.py

@@ -4,9 +4,9 @@ import ctypes
 from typing import Literal, TypeAlias
 from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR
 from tinygrad.runtime.support import c
-from tinygrad.helpers import DEV, ARCH_X86, WIN, OSX
-import gzip, base64, platform
-dll = c.DLL('mesa', 'tinymesa_cpu' if (_cpu:=DEV.renderer == 'LVP') else 'tinymesa', emsg="not available on this platform" if WIN or (OSX and (platform.machine() != "arm64" or (_mv:=platform.mac_ver()[0][:2]) not in {"14","15","26"})) or (platform.system() == "Linux" and platform.machine() not in {"x86_64", "aarch64"}) else f"run `sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa{'_cpu'*_cpu}-mesa-25.2.7-{'macos-'+_mv if OSX else 'linux'}-{'amd64' if ARCH_X86 else 'arm64'}.{'dylib' if OSX else 'so'} -o /usr/local/lib/libtinymesa{'_cpu'*_cpu}.{'dylib' if OSX else 'so'}`")
+from tinygrad.helpers import DEV
+import gzip, base64, platform, sysconfig, os
+dll = c.DLL('mesa', 'tinymesa_cpu' if DEV.renderer == 'LVP' else 'tinymesa', os.path.join(sysconfig.get_paths()['platlib'], 'tinymesa'), emsg='pip install tinymesa==25.2.7.2')
 class struct_u_printf_info(c.Struct): pass
 u_printf_info: TypeAlias = struct_u_printf_info
 uint32_t: TypeAlias = ctypes.c_uint32

tinygrad/runtime/ops_dsp.py

@@ -13,10 +13,8 @@ from tinygrad.uop.ops import PatternMatcher, UPat
 
 dsp_pm = PatternMatcher([
   (((UPat.var('x').maximum(0) ^ -1).maximum(-256) ^ -1).cast(dtypes.uchar.vec(128)),
-   lambda x: UOp(Ops.CUSTOM, dtypes.uchar.vec(128), src=tuple(x.gep(tuple(range(i, i+32))) for i in range(0, 128, 32)),
+   lambda x: UOp(Ops.CUSTOM, dtypes.uchar.vec(128), src=tuple(UOp.vectorize(*[x.index(j) for j in range(i, i+32)]) for i in range(0, 128, 32)),
      arg="__builtin_HEXAGON_V6_vpackhub_sat_128B(__builtin_HEXAGON_V6_vpackwh_sat_128B({3}, {2}), __builtin_HEXAGON_V6_vpackwh_sat_128B({1}, {0}))")),
-  (UPat(Ops.GEP, name="x"), lambda x: UOp(Ops.CUSTOM, x.dtype, x.src+x.src,
-    "__builtin_shufflevector({0}, {1}, "+','.join([str(y) for y in x.arg])+")") if len(x.arg) > 1 and x.src[0].dtype.count > 1 else None),
 ])
 
 dsp_pm_late = PatternMatcher([

tinygrad/runtime/ops_tinyfs.py

@@ -1,8 +1,7 @@
 import socket, json, asyncio, threading, math
 from contextlib import asynccontextmanager
 from tinygrad.device import Compiled, Allocator
-from tinygrad.helpers import DEBUG, getenv
-from tinygrad import Tensor
+from tinygrad.helpers import DEBUG, getenv, CHUNK_SIZE
 
 TINYFS_ENDPOINT = getenv("TINYFS_ENDPOINT", "localhost:6767")
 TINYFS_TIMEOUT = getenv("TINYFS_TIMEOUT", 60)
@@ -95,7 +94,7 @@ class TinyFSAllocator(Allocator[TinyFSDevice]):
       dest.copyout_queue = json.loads(locs)
       dest.hash_buf = src.tobytes()
     elif dest.device.op == "STORE":
-      expected_hashes = math.ceil(dest.size / Tensor.CHUNK_SIZE)
+      expected_hashes = math.ceil(dest.size / CHUNK_SIZE)
       dest.hash_buf = bytearray(expected_hashes * 16)
       self.dev.sfile.readinto(dest.hash_buf)
 
@@ -109,8 +108,8 @@ class TinyFSAllocator(Allocator[TinyFSDevice]):
   async def _copyout_async(self, dest:memoryview, src:TinyFSBuffer):
     async def _worker(i, loc):
       async with self.dev.connection(loc) as (reader, writer):
-        ptr = i * Tensor.CHUNK_SIZE
-        size = min(len(dest[ptr:ptr+Tensor.CHUNK_SIZE]), Tensor.CHUNK_SIZE)
+        ptr = i * CHUNK_SIZE
+        size = min(len(dest[ptr:ptr+CHUNK_SIZE]), CHUNK_SIZE)
 
         writer.write(f"CHUNK_OUT {size}\r\n".encode())
         writer.write(src.hash_buf[i*16:(i+1)*16])

tinygrad/runtime/support/compiler_mesa.py

@@ -1,7 +1,7 @@
 import base64, ctypes, pathlib, tempfile, hashlib
 from tinygrad.device import Compiler
 from tinygrad.helpers import cpu_objdump, system, data64
-from tinygrad.runtime.autogen import mesa, llvm
+from tinygrad.runtime.autogen import mesa, llvm, libc
 from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, expect, cerr
 
 # NB: compilers assume mesa's glsl type cache is managed externally with mesa.glsl_type_singleton_init_or_ref() and mesa.glsl_type_singleton_decref()
@@ -80,14 +80,16 @@ class NAKCompiler(Compiler):
     except Exception as e: print("Failed to generate SASS", str(e), "Make sure your PATH contains nvdisasm binary of compatible version.")
 
 def disas_adreno(lib:bytes, gpu_id=630):
-  with tempfile.TemporaryFile('w+', buffering=1) as tf:
+  with tempfile.TemporaryFile('w+') as tf:
+    mesa_fp = ctypes.cast(fp:=libc.fdopen(tf.fileno(), b"w"), ctypes.POINTER(mesa.struct__IO_FILE))
     @ctypes.CFUNCTYPE(None, ctypes.c_void_p, ctypes.c_uint32, ctypes.c_void_p)
     def hd(data, n, instr):
       fst, snd = data64(ctypes.cast(instr, ctypes.POINTER(ctypes.c_uint64)).contents.value)
-      print(f"{n:04} [{fst:08x}_{snd:08x}] ", end="", flush=True, file=tf)
+      # print with libc so that output interleaves properly
+      libc.fputs(f"{n:04} [{fst:08x}_{snd:08x}] ".encode(), fp)
 
-    ctypes.CDLL(None).setlinebuf(fp:=ctypes.cast(ctypes.CDLL(None).fdopen(tf.fileno(), b"w"), ctypes.POINTER(mesa.struct__IO_FILE)))
-    mesa.ir3_isa_disasm(lib, len(lib), fp, mesa.struct_isa_decode_options(gpu_id, True, 0, True, pre_instr_cb=hd))
+    mesa.ir3_isa_disasm(lib, len(lib), mesa_fp, mesa.struct_isa_decode_options(gpu_id, True, 0, True, pre_instr_cb=hd))
+    libc.fflush(fp)
     tf.seek(0)
     print(tf.read())
 

tinygrad/runtime/support/system.py

@@ -163,6 +163,7 @@ class PCIDevice:
 
     if FileIOInterface.exists(f"/sys/bus/pci/devices/{self.pcibus}/driver"):
       FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/driver/unbind", os.O_WRONLY).write(self.pcibus)
+    if FileIOInterface.exists(f"/sys/bus/pci/devices/{self.pcibus}/driver"): raise RuntimeError(f"Driver is bound to {pcibus}")
 
     if getenv("VFIO", 0) and (vfio_fd:=System.vfio) is not None:
       FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/driver_override", os.O_WRONLY).write("vfio-pci")

tinygrad/schedule/__init__.py

@@ -11,6 +11,14 @@ def _unwrap_src(s: UOp) -> UOp:
   while len(s.src) and s.op not in {Ops.AFTER, Ops.BUFFER, Ops.PARAM, Ops.MSELECT, Ops.MSTACK, Ops.BIND}: s = s.src[0]
   return s
 
+# a buffer state is AFTER | BUFFER | PARAM. MSELECT/MSTACK join per-device states, BIND is not a buffer dependency
+def _states(s: UOp) -> list[UOp]:
+  s = _unwrap_src(s)
+  if s.op in {Ops.MSELECT, Ops.MSTACK}: return [st for ss in s.src for st in _states(ss)]
+  if s.op is Ops.BIND: return []
+  assert s.op in {Ops.AFTER, Ops.BUFFER, Ops.PARAM}, f"input to kernel must resolve to a buffer state, not {s.op}"
+  return [s]
+
 def _split_after(after: UOp) -> tuple[tuple[UOp, ...], tuple[UOp, ...]]:
   kernels, remaining = partition(after.src[1:], lambda s: s.op in {Ops.CALL, Ops.END})
   deps, remaining = partition(remaining, lambda s: s.op is Ops.AFTER)
@@ -23,32 +31,35 @@ def create_schedule(sched_sink:UOp) -> UOp:
     # build kernel dependency graph: edges from producer kernel to consumer kernels
     children: dict[UOp, list[UOp]] = {}
     in_degree: dict[UOp, int] = {}
+    writes: dict[UOp, list[tuple[UOp, UOp, tuple[UOp, ...]]]] = {}  # buffer -> (AFTER, prior state, new kernels)
+    reads: list[tuple[UOp, UOp, UOp]] = []  # (reader AFTER, reader kernel, buffer state read)
     for u in sched_sink.toposort(gate_kernel_sink):
       if u.op is not Ops.AFTER: continue
       kernels, after_deps = _split_after(u)
+      prev_state = _unwrap_src(u.src[0])
+      prev_kernels = set(_split_after(prev_state)[0]) if prev_state.op is Ops.AFTER else set()
+      writes.setdefault(u.buf_uop, []).append((u, prev_state, tuple(k for k in kernels if k not in prev_kernels)))
       for k in kernels:
         in_degree.setdefault(k, 0)
         if k.op is Ops.END: assert k.src[0].op is Ops.CALL, f"END src[0] should be KERNEL, not {k.src[0].op}"
         kernel_deps = k.src[0].src[1:] if k.op is Ops.END else k.src[1:]
-        for s in kernel_deps + after_deps:
-          match (s := _unwrap_src(s)).op:
-            case Ops.AFTER:
-              for t in _split_after(s)[0]:
-                children.setdefault(t, []).append(k)
-                in_degree[k] += 1
-            case Ops.MSELECT | Ops.MSTACK:
-              for ss in s.src:
-                if ss.op is Ops.MSELECT: ss = ss.src[0]
-                ss = _unwrap_src(ss)
-                if ss.op not in {Ops.BUFFER, Ops.PARAM}:
-                  assert ss.op is Ops.AFTER, f"ss.op is not AFTER, it's {ss.op}"
-                  for t in _split_after(ss)[0]:
-                    children.setdefault(t, []).append(k)
-                    in_degree[k] += 1
-            case Ops.BUFFER | Ops.PARAM | Ops.BIND:
-              pass  # BUFFER/PARAM is already realized, BIND is a bound variable (not a buffer dependency)
-            case _:
-              raise RuntimeError(f"input to kernel must be AFTER, BUFFER, PARAM, MSELECT, MSTACK, or BIND, not {s.op}")
+        read_states = [st for s in kernel_deps for st in _states(s)]
+        reads += [(u, k, st) for st in read_states]
+        # RAW deps: a kernel runs after the kernels that produced the states it reads or joins
+        for st in read_states + [st for s in after_deps for st in _states(s)]:
+          if st.op is Ops.AFTER:
+            for t in _split_after(st)[0]:
+              children.setdefault(t, []).append(k)
+              in_degree[k] += 1
+    # WAR deps: a kernel reading buffer state S must run before another write that supersedes S. an AFTER only
+    # supersedes its immediate prior state; join members already present in that prior state are ordering deps, not writes
+    for u, k, s in reads:
+      for a, prev_state, write_kernels in writes.get(s.buf_uop, []):
+        if a is u or prev_state is not s: continue
+        for t in write_kernels:
+          if t is not k and t not in k.backward_slice:
+            children.setdefault(k, []).append(t)
+            in_degree[t] += 1
 
   with cpu_profile(TracingKey("linearize schedule")):
     queue: deque[UOp] = deque(k for k,v in in_degree.items() if v == 0)
@@ -65,6 +76,7 @@ def create_schedule(sched_sink:UOp) -> UOp:
       for x in children.get(rk, []):
         in_degree[x] -= 1
         if in_degree[x] == 0: queue.append(x)
+    if any(in_degree.values()): raise RuntimeError("cycle detected in assign graph")
   return UOp(Ops.LINEAR, src=tuple(linearized))
 
 from tinygrad.schedule.memory import memory_plan_rewrite

tinygrad/schedule/indexing.py

@@ -8,7 +8,7 @@ from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_claus
 from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC
 
 ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.COPY, Ops.BUFFER, Ops.SLICE,
-                      Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.PARAM,
+                      Ops.CONST, Ops.BIND, Ops.MSELECT, Ops.MSTACK, Ops.PARAM,
                       Ops.LOAD, Ops.CALL, Ops.FUNCTION}
 
 def realize(ctx:dict[UOp, None], tr:UOp) -> None: ctx[tr] = None
@@ -162,8 +162,6 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
   # explicit rangeify
   ending_ranges: dict[UOp, list[UOp]] = {}
   for x in reversed(tsink_toposort):
-    if x.op is Ops.DEVICE: continue
-
     # no ranges on kernels, they are internal
     if x.op in {Ops.CALL, Ops.FUNCTION, Ops.LINEAR}: continue
 
@@ -254,7 +252,13 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
 
     # REDUCE creates ranges for the axes it is reducing
     if x.op is Ops.REDUCE and len(x.arg[1]):
-      rngs = tuple(rctx.new_range(s, axistype=AxisType.REDUCE) if i in x.arg[1] else r for i,(r,s) in enumerate(zip(rngs, x.src[0].shape)))
+      out_i, in_rngs = 0, []
+      for i,s in enumerate(x.src[0].shape):
+        if i in x.arg[1]: in_rngs.append(rctx.new_range(s, axistype=AxisType.REDUCE))
+        else:
+          in_rngs.append(out_rngs[out_i])
+          out_i += 1
+      rngs = tuple(in_rngs)
 
     if debug:
       realized_ranges = rctx.realize_map.get(x, None)

tinygrad/schedule/multi.py

@@ -75,7 +75,8 @@ def reduce_multi(root:UOp, multi:UOp):
       return local.cast(orig_dtype).allreduce(op, multi.device).cast(local.dtype)
     return local.allreduce(op, multi.device)
   # reduce on non sharded axes, piecewise is fine. if axis is None this is also correct
-  return multi.src[0]._rop(op, axis).multi(axis=multi.axis)
+  new_axis = multi.axis - sum(1 for a in axis if a < multi.axis) if multi.axis is not None else None
+  return multi.src[0]._rop(op, axis).multi(axis=new_axis)
 
 def reshape_multi(root:UOp, multi:UOp):
   if prod(multi.shape) != prod(new_shape:=root.marg): raise RuntimeError("reshape must maintain prod(shape)")

tinygrad/schedule/rangeify.py

@@ -2,7 +2,7 @@ from dataclasses import dataclass, field, replace
 from typing import cast
 import itertools
 from tinygrad.dtype import dtypes, PtrDType, AddrSpace, Invalid
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, KernelInfo, ParamArg, shape_to_shape_arg
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg
 from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, identity_element
 from tinygrad.uop.symbolic import symbolic
 from tinygrad.helpers import prod, all_same, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS
@@ -32,7 +32,7 @@ def lower_shaped_wmma(ctx, x):
   wmma_arg = (name, dims, dtype_in, dtype_out, device, threads, tc_upcast_axes, ())
   wmma = UOp(Ops.WMMA, dtype_out.vec(x.src[2].shape[-1]), tuple(s[u].contract(u) for s, u in upcasts), arg=wmma_arg)
   tmp = UOp.placeholder((x.src[2].shape[-1],), dtype_out, slot=next(ctx), addrspace=AddrSpace.REG)
-  return tmp.after(UOp.group(*[tmp[e].store(wmma.gep(e)) for e in range(x.src[2].shape[-1])]))
+  return tmp.after(UOp.group(*[tmp[e].store(wmma.index(e)) for e in range(x.src[2].shape[-1])]))
 
 pm_store_ranges = PatternMatcher([
   (UPat(Ops.STORE, name="x"), add_ranges_to_store),
@@ -125,6 +125,12 @@ def split_reduceop(reduce:UOp, x:UOp):
 mop_cleanup = PatternMatcher([
   # merge adjacent RESHAPES
   (UPat(Ops.RESHAPE, src=(UPat(Ops.RESHAPE, name="x2"), UPat()), name="x"), lambda x,x2: x.replace(src=(x2.src[0], x.src[1]))),
+  # remove noop RESHAPEs
+  (UPat(Ops.RESHAPE, src=(UPat(name="x2"), UPat()), name="x"), lambda x,x2: x2 if x2._shape is not None and x2.shape == x.shape else None),
+  # merge PERMUTEs
+  (UPat(Ops.PERMUTE, src=(UPat(Ops.PERMUTE, name="x2"),), name="x"), lambda x,x2: x2.replace(arg=tuple(x2.arg[i] for i in x.arg))),
+  # remove noop PERMUTEs
+  (UPat(Ops.PERMUTE, name="x"), lambda x: x.src[0] if list(x.arg) == list(range(len(x.arg))) else None),
 ])
 
 pm_gather_params = PatternMatcher([ (UPat(Ops.PARAM, name="p"), lambda ctx, p: ctx.append(p) if p.arg.slot >= 0 else None), ])
@@ -479,10 +485,6 @@ pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
   (UPat(Ops.AFTER, name="x"), remove_noop_afters),
 ])
 
-pm_add_buffers_local = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
-  (UPat(Ops.STAGE, src=(UPat(), UPat(name="idx")), name="x"), bufferize_to_store),
-])
-
 # *****************
 # 5. split into kernels
 
@@ -615,18 +617,5 @@ def get_kernel_graph(sink:UOp) -> UOp:
   tsink = graph_rewrite(tsink, pm_add_buffers+pm_add_range_tags, ctx=itertools.count(paramarg_start), bottom_up=True, name="stage to store")
   tsink = graph_rewrite(tsink, split_kernels, bottom_up=True, name="split kernels")
 
-  # WAR deps: if kernel U reads buffer S, and S is also written by another kernel, S's write must wait for U to finish
-  afters = [u for u in tsink.toposort() if u.op is Ops.AFTER]
-  kernel_assign: dict[UOp, UOp] = {u.buf_uop:u for u in afters}
-  assign_rep: dict[UOp, UOp] = {}
-  for u in afters:
-    for s in u.src[1].src:
-      # TODO: this is probably broken for MSELECT/MSTACK
-      if s.op not in {Ops.BUFFER, Ops.PARAM} or s is u.buf_uop or (a:=kernel_assign.get(s)) is None: continue
-      if a.src[1] is u.src[1]: continue  # same kernel (multi-output custom kernels)
-      if any(x.op is Ops.AFTER and x.buf_uop is s for x in kernel_assign[u.buf_uop].backward_slice):
-        raise RuntimeError(f"cycle detected in assign graph, buffers {s} and {u.buf_uop} have circular dependency")
-      assign_rep[a] = kernel_assign[s] = a.replace(src=a.src+(u,))
-  if assign_rep: tsink = graph_rewrite(tsink, _substitute, ctx=assign_rep, bottom_up=True, name="fix_assign")
   if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
   return tsink

tinygrad/tensor.py

@@ -1,10 +1,10 @@
 # inspired by https://github.com/karpathy/micrograd/blob/master/micrograd/engine.py
 from __future__ import annotations
-import time, math, functools, sys, inspect, pathlib, hashlib, weakref
+import time, functools, sys, inspect, pathlib, hashlib, weakref
 from typing import Any, Callable, Sequence, cast, get_args, ParamSpec, TypeVar, Generic, TYPE_CHECKING
 if TYPE_CHECKING: import numpy
 from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, to_dtype, _from_np_dtype, _to_np_dtype, PyConst
-from tinygrad.helpers import argfix, prod, all_int, getenv, fully_flatten, ceildiv, fetch, Metadata, TRACEMETA, is_numpy_ndarray, TracingKey
+from tinygrad.helpers import prod, all_int, getenv, fully_flatten, ceildiv, fetch, Metadata, TRACEMETA, is_numpy_ndarray, TracingKey
 from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc
 from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, _index_to_concrete_int, Variable, _broadcast_shape
 from tinygrad.mixin.rand import RandMixin
@@ -16,16 +16,16 @@ from tinygrad.callify import transform_to_call
 # *** all in scope Tensors are here. this gets relevant UOps ***
 
 all_tensors: dict[weakref.ref[Tensor], None] = {}
-def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str, walk:bool=False) -> None:
+def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str) -> None:
   with cpu_profile(TracingKey(name), "TINY"):
     # get tensors in scope
     in_scope: dict[UOp, bool] = {}
     def visitor(node: UOp) -> bool: return True if node in applied_map else any(in_scope.get(s, False) for s in node.src)
     scope_tensors: list[Tensor] = [t for tref in list(all_tensors) if (t:=tref()) is not None and t.uop.topovisit(visitor, in_scope)]
 
-    # get all Tensors and apply the map
+    # get all Tensors and apply the map. always walk: replace exactly the nodes the map names, values are final
     sink = UOp.sink(*[t.uop for t in scope_tensors])
-    new_sink = sink.substitute(applied_map, name=f"substitute {name}", walk=walk)
+    new_sink = sink.substitute(applied_map, name=f"substitute {name}", walk=True)
 
     # set the relevant uop to the realized UOps
     for t,s,ns in zip(scope_tensors, sink.src, new_sink.src):
@@ -120,7 +120,8 @@ class Tensor(RandMixin):
   def alu(self, op: Ops, *src: Tensor) -> Tensor: return self._apply_uop(lambda *u: u[0].alu(op, *u[1:]), *src)
   @property
   def _uop(self) -> UOp: return self.uop
-  def _wrap_uop(self, u:UOp) -> Tensor: return Tensor(u)
+  @classmethod
+  def _wrap_uop(cls, u:UOp) -> Tensor: return cls(u)
   @staticmethod
   def const(dtype:DType, b:ConstType|UOp) -> Tensor: return Tensor(UOp.const(dtype, b))
 
@@ -225,7 +226,7 @@ class Tensor(RandMixin):
       ib = self.uop
       while not ib.has_buffer_identity() and ib is not base: ib = ib.src[0]
       assigned_ib = ib.after(assign)
-      _apply_map_to_tensors({ib: assigned_ib}, name="Embed View Assign", walk=True)
+      _apply_map_to_tensors({ib: assigned_ib}, name="Embed View Assign")
     else:
       # simple assign
       self.uop = assign
@@ -348,75 +349,8 @@ class Tensor(RandMixin):
     if isinstance(y.device, str): return self.to(y.device)
     return self if isinstance(self.device, tuple) and (y.device, y.uop.axis) == (self.device, self.uop.axis) else self.shard(y.device, y.uop.axis)
 
-  CHUNK_SIZE = 2**20
-  def fs_load(self, size:int) -> Tensor:
-    """
-    Load a tensor from storage.
-
-    self should be a tensor of the hash to load
-    """
-    # TODO: this should work locally as well
-    assert self.dtype == dtypes.uint8, "hash is expected to be uint8"
-    h = self.contiguous().flatten()
-    assert h.shape[0] == 16, "expected hash"
-
-    base_chunks = math.ceil(size / Tensor.CHUNK_SIZE)
-    tree_depth = math.ceil(math.log(base_chunks, Tensor.CHUNK_SIZE // 16))
-    data, level_chunks = h, 0
-    for i in reversed(range(tree_depth + 1)):
-      data = data.to("tinyfs:load")
-
-      # if not last level, its still hashes
-      if i > 0 or tree_depth == 0:
-        level_chunks = max(1, math.ceil(base_chunks / (Tensor.CHUNK_SIZE // 16)**(i-1)))
-        pad_amt = 16 * level_chunks
-      else: pad_amt = Tensor.CHUNK_SIZE * level_chunks
-      if (tsize := data.shape[0]) < pad_amt: data = data.pad((0, pad_amt - tsize))
-      data = data[:pad_amt].contiguous()
-      if i != 0: data = data.to(self.device)
-
-    return data[:size]
-
-  def fs_store(self) -> Tensor:
-    """
-    Store a tensor to storage.
-    """
-    # TODO: this should work locally as well
-    data = self.contiguous().flatten().bitcast(dtypes.uint8)
-
-    # pad to a multiple of 1mb
-    if (tsize := data.shape[0]) % Tensor.CHUNK_SIZE != 0: data = data.pad((0, Tensor.CHUNK_SIZE - tsize % Tensor.CHUNK_SIZE))
-    size = data.shape[0]
-
-    base_chunks = math.ceil(size / Tensor.CHUNK_SIZE)
-    tree_depth = math.ceil(math.log(base_chunks, Tensor.CHUNK_SIZE // 16))
-
-    to_device = "CPU" if isinstance(self.device, str) and self.device.startswith("DISK") else self.device
-
-    level_chunks = base_chunks
-    for _ in range(tree_depth + 1):
-      data = data.to("tinyfs:store")[:level_chunks * 16].contiguous().to(to_device)
-      if (tsize := data.shape[0]) % Tensor.CHUNK_SIZE != 0: data = data.pad((0, Tensor.CHUNK_SIZE - tsize % Tensor.CHUNK_SIZE))
-      level_chunks = math.ceil(data.shape[0] / Tensor.CHUNK_SIZE)
-
-    return data[:16].contiguous()
-
   # ***** creation entrypoint *****
 
-  @staticmethod
-  def empty(*shape, device:str|tuple[str, ...]|None=None, dtype:DTypeLike|None=None) -> Tensor:
-    """
-    Creates an empty tensor with the given shape.
-
-    You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor.empty(2, 3)
-    print(t.shape)
-    ```
-    """
-    return Tensor(UOp.empty(argfix(*shape), dtype, device))
-
   @staticmethod
   def from_blob(ptr:int, shape:tuple[int, ...], **kwargs) -> Tensor:
     """

tinygrad/uop/__init__.py

@@ -38,7 +38,7 @@ class Ops(FastEnum):
   SINK = auto(); AFTER = auto(); GROUP = auto()
 
   # vector creation / item selection
-  GEP = auto(); STACK = auto()
+  STACK = auto()
 
   # tuple/gettuple for function with multiple returns
   TUPLE = auto(); GETTUPLE = auto()
@@ -89,9 +89,6 @@ class Ops(FastEnum):
 
   # ** 6 -- ops that don't exist in programs **
 
-  # tensor graph ops
-  DEVICE = auto()
-
   # ops that adjust the behavior of the scheduler
   CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto()
 

tinygrad/uop/ops.py

@@ -85,7 +85,7 @@ def multirange_str(rngs:Iterable[UOp], color=False, pad=None) -> str:
 def shape_to_shape_arg(arg:tuple[sint, ...]) -> UOp:
   if len(arg) == 0: return UOp(Ops.STACK)
   elif len(arg) == 1: return UOp.const(dtypes.weakint, arg[0])
-  else: return UOp(Ops.STACK, dtypes.weakint.vec(len(arg)), tuple(UOp.const(dtypes.weakint, x) if isinstance(x, int) else x for x in arg))
+  else: return UOp(Ops.STACK, dtypes.weakint, tuple(UOp.const(dtypes.weakint, x) if isinstance(x, int) else x for x in arg))
 
 def consumer_map_from_toposort(lst:Iterable[UOp]):
   ret: dict[UOp, dict[UOp, None]] = {}
@@ -226,8 +226,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   def _shape(self) -> tuple[sint, ...]|None:
     match self.op:
       # late ops don't have shape
-      case Ops.DEVICE | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
-           Ops.SINK | Ops.REWRITE_ERROR | Ops.ENDIF | \
+      case Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | Ops.SINK | Ops.REWRITE_ERROR | Ops.ENDIF | \
            Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.INS | Ops.TUPLE | Ops.CALL | Ops.FUNCTION:
         return None
 
@@ -250,20 +249,11 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
           return tuple(graph_rewrite(s, _pm_resolve_params, self.src[0].src[1:], walk=True) if isinstance(s, UOp) else s for s in inner_shape)
         return inner_shape
 
-      case Ops.CAST:
-        # if it has a vec dtype, set the shape
-        if self.dtype.count > 1: return (self.dtype.count,)
-        # when PTX casts from ptr to non ptr, remove the shape of the buffer
-        if isinstance(self.src[0].dtype, PtrDType) and not isinstance(self.src[0].dtype, ImageDType) and not isinstance(self.dtype, PtrDType):
-          return ()
-
       case Ops.INDEX:
         shp:list[sint] = []
         for s in self.src[1:]: shp.extend(list(s.shape))
         return tuple(shp) + self.src[0].shape[len(self.src[1:]):]
 
-      case Ops.GEP:
-        return (len(self.arg),) if len(self.arg) > 1 else ()
       case Ops.STACK:
         if len(self.src) == 0: return ()
         if isinstance(self.dtype, PtrDType):
@@ -297,16 +287,16 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
         return self.src[0].as_shape if len(self.src) >= 1 else None
 
       # wmma output shape = accumulator shape (src[2])
-      case Ops.WMMA | Ops.SHAPED_WMMA: return self.src[2]._shape
+      case Ops.WMMA:
+        in0, in1, out0 = self.arg[6]
+        wmma_b = _broadcast_shape(self.src[0].shape[:-1], self.src[1].shape[:-1], self.src[2].shape[:-1])
+        return wmma_b + (prod([x for _,x in out0]),)
+      case Ops.SHAPED_WMMA: return self.src[2]._shape
 
       # passthrough ops
       case Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.LOAD | \
            Ops.COPY | Ops.ALLREDUCE | Ops.STORE | Ops.END:
         return self.src[0]._shape
-      # REDUCE with empty axis is passthrough (lowered form)
-      case Ops.REDUCE if len(self.arg[1]) == 0:
-        # these can mismatch if there's a horizonal reduce
-        return (self.dtype.count,) if self.dtype.count > 1 else ()
 
       # TODO: disallow shape changing bitcast
       case Ops.BITCAST:
@@ -354,7 +344,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
           axis_arg = self.arg[1]
           if not isinstance(axis_arg, tuple) or not all(isinstance(x, int) and x>=0 and x<len(ps) for x in axis_arg):
             raise ValueError(f"invalid type for axis: {axis_arg}")
-          return tuple(1 if i in axis_arg else s for i,s in enumerate(ps))
+          return tuple(s for i,s in enumerate(ps) if i not in axis_arg)
 
     if self.op in GroupOp.Unary.union({Ops.CAST}):
       assert len(self.src) == 1, "unary ops must have 1 src"
@@ -473,10 +463,12 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   def _stack(self, *srcs):
     # TODO: this should become the real stack
     return UOp(Ops.STACK, self.dtype, (self,)+srcs)
-  def vectorize(self, *srcs):
-    return UOp(Ops.STACK, self.dtype.vec(len(srcs)+1), (self,)+srcs)
-  def index(self, *srcs:UOp|None, ptr=False, **kwargs):
-    return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype if ptr else self.dtype.base), (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
+  def vectorize(self, *srcs): return self._stack(*srcs)
+  def index(self, *srcs:UOp|int|None, ptr=False, **kwargs):
+    new_srcs: list[UOp] = [UOp.const(dtypes.weakint, x) if isinstance(x, int) else x for x in srcs if x is not None]
+    if not ptr and len(new_srcs) == 1 and new_srcs[0].op is Ops.CONST and self.op is Ops.STACK:
+      return self.src[new_srcs[0].arg]
+    return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype if ptr else self.dtype.base.scalar()), (self,)+tuple(new_srcs), **kwargs)
   def __getitem__(self, idx):
     # pointers index into INDEX UOps (scalar lookup); everything else uses the shared mixin view path
     if not isinstance(self.dtype, PtrDType): return super(UOp, self).__getitem__(idx)
@@ -493,7 +485,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     return self.index(*[UOp.const(dtypes.weakint, x) if isinstance(x, int) else x for x in idx])
   @property
   def _uop(self) -> UOp: return self
-  def _wrap_uop(self, u:UOp) -> UOp: return u
+  @classmethod
+  def _wrap_uop(cls, u:UOp) -> UOp: return u
   def const_like(self, b:ConstLike, dtype:DType|None=None):
     return UOp.const(dtype or self.dtype.base, b, shape=self._shape)
   def vconst_like(self, b:ConstLike, dtype:DType|None=None):
@@ -520,14 +513,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   def bitcast(self, dtype:DTypeLike):
     dtype = to_dtype(dtype)
     return self if self.dtype == dtype else UOp(Ops.BITCAST, dtype, (self,))
-  def gep(self, i:tuple[int, ...]|int):
-    if isinstance(i, tuple) and len(i) == 1: return self.gep(i[0])
-    if isinstance(i, int):
-      # NOTE: these are just shortcuts to not have to create and fold later
-      if self.op is Ops.STACK: return self.src[i]
-      if self.op is Ops.CONST: return UOp.const(self.dtype.scalar(), self.arg)
-      i = (i,)
-    return UOp(Ops.GEP, self.dtype.scalar().vec(len(i)) if len(i) > 1 else self.dtype.scalar(), (self,), i)
   def load(self, *src:UOp, **kwargs): return UOp(Ops.LOAD, dtype=kwargs.pop("dtype", self.dtype.base), src=(self,)+src, **kwargs)
   def store(self, src:UOp|ConstType, gate:UOp|None=None, **kwargs):
     srcs = (self, self.const_like(src) if not isinstance(src, UOp) else src) + ((gate,) if gate is not None else ())
@@ -566,8 +551,11 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   @staticmethod
   def special(end:sint, name:str, dtype=dtypes.weakint): return UOp(Ops.SPECIAL, dtype=dtype, src=(sint_to_uop(end, dtype),), arg=name)
   def _rop(self, op:Ops, axis:tuple[int, ...]):
-    axis = tuple(sorted([x for x in axis if resolve(self.shape[x] != 1)]))
-    return UOp(Ops.REDUCE, self.dtype, (self,), (op, axis)) if len(axis) else self
+    # NOTE: we don't allow reduce on 1s axis
+    axis = tuple(sorted(axis))
+    reduce_axis = tuple(x for x in axis if resolve(self.shape[x] != 1))
+    ret = UOp(Ops.REDUCE, self.dtype, (self,), (op, reduce_axis)) if len(reduce_axis) else self
+    return ret.reshape(tuple(s for i,s in enumerate(self.shape) if i not in axis)) if axis != reduce_axis else ret
   @staticmethod
   def invalid(count=1): return UOp(Ops.CONST, dtypes.weakint.vec(count), src=(), arg=Invalid)
   def valid(self, cond):
@@ -636,7 +624,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     src_axis = self.src[0].axis
     if self.op is Ops.SHRINK and src_axis is not None and self.marg[src_axis] != (0, self.src[0].shape[src_axis]):
       return None # SHRINK will remove the sharding if it's on axis
-    if self.op is Ops.REDUCE: return None if src_axis is not None and src_axis in self.arg[1] else src_axis
+    if self.op is Ops.REDUCE:
+      if src_axis is None: return None
+      if src_axis in self.arg[1]: return None
+      return src_axis - sum(1 for a in self.arg[1] if a < src_axis)
     if self.op is Ops.RESHAPE:
       if src_axis is None: return None
       arg_acc:list[sint] = list(itertools.accumulate(self.marg, operator.mul, initial=1))
@@ -675,13 +666,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   @property
   def base(self) -> UOp:
-    if self.op in GroupOp.Movement: return self.src[0].base
-    if self.op is Ops.MULTI: return self.src[0].base  # MULTI is really a VIEW
-    if self.op is Ops.DETACH: return self.src[0].base  # DETACH can't change base
-    return self
-
-  @property
-  def multibase(self) -> UOp:
     if self.op in GroupOp.Movement: return self.src[0].base
     if self.op is Ops.DETACH: return self.src[0].base  # DETACH can't change base
     return self
@@ -735,16 +719,11 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     if (uop:=UOp.new_buffer(device or opaque.device, opaque.size, opaque.dtype, num=-id(opaque))) not in buffers: buffers[uop] = opaque.ref(1)
     else: assert buffers[uop] is opaque
     return uop
-  @staticmethod
-  def empty(shape:tuple[sint, ...], dtype:DTypeLike|None=None, device:str|tuple[str, ...]|None=None, axis:int|None=None, num=None) -> UOp:
-    dtype, device = to_dtype(dtype) if dtype is not None else dtypes.default_float, canonicalize_device(device)
-    max_shape = to_max_shape(shape)
-    ret = UOp.new_buffer(device, prod(max_shape), dtype, num).reshape(max_shape).shrink_to(shape)
-    return ret.multi(axis) if isinstance(device, tuple) and axis is not None else ret
   def empty_like(self, dtype:DTypeLike|None=None, device:str|tuple[str, ...]|None=None) -> UOp:
     device = canonicalize_device(self.device if device is None else device)
     axis = self.axis if isinstance(device, tuple) else None
-    return UOp.empty(self.shard_shape if axis is not None else self.shape, self.dtype if dtype is None else dtype, device, axis)
+    ret = UOp.empty(self.shard_shape if axis is not None else self.shape, dtype=self.dtype if dtype is None else dtype, device=device)
+    return ret.multi(axis) if axis is not None else ret
   @staticmethod
   def _frompy(x:list|tuple|bytes, dtype:DType, device:str|tuple[str, ...]|None=None) -> UOp:
     device = canonicalize_device(device)
@@ -753,7 +732,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
       # bfloat16 and fp8 have no struct format, so pack a float32 buffer and cast
       bdtype = dtypes.float32 if dtype in [dtypes.bfloat16, *dtypes.fp8s] else dtype
       assert bdtype.fmt is not None, f"{bdtype=} has None fmt"
-      ret = UOp.empty(shape:=get_shape(x), bdtype, "PYTHON")
+      ret = UOp.empty(shape:=get_shape(x), dtype=bdtype, device="PYTHON")
       data = struct.pack(f"{prod(shape)}{bdtype.fmt}", *[truncate[bdtype](bdtype.const(xi)) for xi in fully_flatten(x)])
     # fake realize. if target device is PYTHON it needs bytearray to be writable
     ret.buffer.allocate(memoryview(data if device != "PYTHON" else bytearray(data)))
@@ -767,7 +746,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   @recursive_property
   def device(self) -> str|tuple[str, ...]|None:
     if self.op is Ops.PARAM: return self.arg.device
-    if self.op is Ops.DEVICE: return self.arg
     if self.op is Ops.STAGE: return self.arg.device
     if self.op is Ops.AFTER: return self.src[0].device
     if self.op is Ops.MSELECT:
@@ -786,7 +764,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     if self.op is Ops.BUFFER: return self.arg.addrspace
     if self.op in {Ops.SPECIAL, Ops.RANGE}: return AddrSpace.ALU
     if self.op is Ops.LOAD: return AddrSpace.ALU # LOAD brings things into the ALU
-    if self.op in {Ops.INDEX, Ops.CAST, Ops.AFTER, Ops.REDUCE, Ops.GEP, Ops.STORE, Ops.MSTACK, Ops.MSELECT}:
+    if self.op in {Ops.INDEX, Ops.CAST, Ops.AFTER, Ops.REDUCE, Ops.STORE, Ops.MSTACK, Ops.MSELECT}:
       return self.src[0].addrspace
     if self.op in GroupOp.Movement: return self.src[0].addrspace
     if self.op in {Ops.STACK, Ops.WMMA} or self.op in GroupOp.Elementwise:
@@ -812,19 +790,17 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     out = graph_rewrite(self.flatten().index(UOp.range(numel, 0)), pm_mops+symbolic, name="contiguous_view_offset")
     if out.op is not Ops.INDEX: return None
     if len(out.src) == 1: return 0 if resolve(numel == 1, False) else None
-    if out.src[1].op is Ops.CONST and resolve(numel == 1, False):
-      if not isinstance(out.src[1].arg, int): return None  # masked/padded regions produce InvalidType
-      return out.src[1].arg
-    if out.src[1].op is Ops.RANGE: return 0
-    if out.src[1].op is Ops.ADD and out.src[1].src[0].op is Ops.RANGE and out.src[1].src[1].op is Ops.CONST:
-      if not isinstance(out.src[1].src[1].arg, int): return None  # masked/padded regions produce InvalidType
-      return out.src[1].src[1].arg
+    idx, has_range = out.src[1], False
+    if idx.op is Ops.RANGE: return 0
+    if idx.op is Ops.ADD and idx.src[0].op is Ops.RANGE: idx, has_range = idx.src[1], True
+    if idx.op is Ops.CONST and (has_range or resolve(numel == 1, False)):
+      if not isinstance(idx.arg, int): return None  # masked/padded regions produce InvalidType
+      return idx.arg
     return None
 
   def has_buffer_identity(self):
     """Check if this UOp has a concrete buffer identity in the graph (RESHAPE/MULTI -> BUFFER chain)."""
     if self.op in {Ops.RESHAPE, Ops.MULTI}: return self.src[0].has_buffer_identity()
-    if self.op is Ops.GETTUPLE and self.src[0].op is Ops.TUPLE: return self.src[0].src[self.arg].has_buffer_identity()
     return self.op in {Ops.BUFFER, Ops.SLICE, Ops.PARAM}
 
   def _base_buffer_is_realized(self) -> bool:
@@ -835,7 +811,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   @property
   def buffer(self) -> Buffer|MultiBuffer:
-    if self.op in {Ops.CONTIGUOUS, Ops.RESHAPE, Ops.DETACH, Ops.AFTER}: return self.src[0].buffer
+    if self.op in {Ops.CONTIGUOUS, Ops.RESHAPE, Ops.MULTI, Ops.DETACH, Ops.AFTER}: return self.src[0].buffer
     # this buffer can process disk tensors and simple movement ops
     if self is not self.base:
       buf = self.base.buffer
@@ -877,6 +853,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     return ret
   @property
   def realized(self) -> Buffer|MultiBuffer|None:
+    if self.op is Ops.MULTI: return self.src[0].realized
     # only these can be realized
     if self.op not in (Ops.BUFFER, Ops.MSTACK): return None
     # LOCAL/REG scratch buffers are never realized
@@ -1004,7 +981,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     if self.op is Ops.BIND: return self.src[0]._min_max # ignore the bound value
     if self.op in {Ops.UNROLL, Ops.STACK}: return min(x.vmin for x in self.src), max(x.vmax for x in self.src)
     if self.op is Ops.CONST and self.arg is not Invalid: return self.arg, self.arg
-    if self.op is Ops.GEP: return self.src[0]._min_max
+    if self.op is Ops.INDEX and not isinstance(self.src[0].dtype, PtrDType): return self.src[0]._min_max
     # TODO: CAST to bool/unsigned is not monotone, still some case can be simplified
     if self.op is Ops.CAST and self.dtype in dtypes.floats+dtypes.sints+(dtypes.weakint,):
       return max(self.dtype.min, self.src[0].vmin), min(self.src[0].vmax, self.dtype.max)
@@ -1041,18 +1018,20 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   # *** uop high level syntactic sugar ***
 
   @staticmethod
-  def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL):
+  def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL, is_ptr=True):
     if addrspace is AddrSpace.GLOBAL:
-      ret = UOp(Ops.PARAM, dtype.ptr(prod(shape), addrspace), arg=ParamArg(slot, addrspace=addrspace))
+      # TODO: this should have a shape
+      ret = UOp(Ops.PARAM, dtype.ptr(prod(shape), addrspace) if is_ptr else dtype, arg=ParamArg(slot, addrspace=addrspace))
     else:
       assert addrspace in (AddrSpace.LOCAL, AddrSpace.REG)
       buf_shape = (prod(shape),) + ((dtype.count,) if dtype.count > 1 else ())
-      ret = UOp(Ops.BUFFER, dtype.ptr(prod(shape), addrspace), src=(shape_to_shape_arg(buf_shape),), arg=ParamArg(slot, addrspace=addrspace))
+      ret = UOp(Ops.BUFFER, dtype.ptr(prod(shape), addrspace) if is_ptr else dtype,
+                src=(shape_to_shape_arg(buf_shape),), arg=ParamArg(slot, addrspace=addrspace))
     if len(shape) > 1: ret = ret.reshape(shape + ((dtype.count,) if addrspace in (AddrSpace.LOCAL, AddrSpace.REG) and dtype.count > 1 else ()))
     return ret
-  def placeholder_like(self, slot:int, addrspace=AddrSpace.GLOBAL):
+  def placeholder_like(self, slot:int, addrspace=AddrSpace.GLOBAL, is_ptr=True):
     assert all_int(self.shape), "no placeholder-like on symbolic shape"
-    return UOp.placeholder(self.max_shard_shape, self.dtype, slot, addrspace)
+    return UOp.placeholder(self.max_shard_shape, self.dtype, slot, addrspace, is_ptr=is_ptr)
 
   # set is store+end+after
   def set(self:UOp, val:UOp|ConstType, end:UOp|tuple[UOp, ...]|list[UOp]=()) -> UOp:
@@ -1072,6 +1051,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
       return UOp.param(slot, self.dtype, self._shape, self.device, cast(tuple[int, int], self._min_max), self.src[0].expr, addrspace)
     return UOp.param(slot, self.dtype, self.shard_shape if self.axis is not None else self._shape, self.device, addrspace=addrspace, axis=self.axis)
 
+  @staticmethod
+  def custom_function(name:str, *src:UOp, dtype:DType=dtypes.void) -> UOp: return UOp(Ops.CUSTOM_FUNCTION, dtype, src=src, arg=name)
+
   # opaque bodies stay as Ops.CALL; value-producing bodies become Ops.FUNCTION (wrapped in TUPLE)
   _OPAQUE_CALL_BODIES = {Ops.SINK, Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.SLICE, Ops.CUSTOM_FUNCTION}
   def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=(),
@@ -1279,7 +1261,7 @@ class UPat(OpMixin):
     if dtype is not None and self.match_dtype == (dtype,): return self
     return UPat(Ops.CAST, dtype, (self,), **kwargs)
   def bitcast(self, dtype=None): return UPat(Ops.BITCAST, dtype, (self,))
-  def gep(self, i:int|None=None, **kwargs): return UPat(Ops.GEP, None, (self,), (i,) if i is not None else None, **kwargs)
+  def gep(self, i:int|None=None, **kwargs): return UPat(Ops.INDEX, None, (self, UPat.cvar("i") if i is not None else UPat()), **kwargs)
   def load(self, *src:UPat, **kwargs): return UPat(Ops.LOAD, src=(self,)+src, **kwargs)
   def store(self, *src:UPat, **kwargs): return UPat(Ops.STORE, src=(self,)+src, **kwargs)
   def reduce(self, *src:UPat, **kwargs):
@@ -1648,7 +1630,7 @@ def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=N
 def sint_to_uop(x:sint, dtype=dtypes.weakint) -> UOp: return UOp.const(dtype, x) if isinstance(x, int) else x.cast(dtype)
 def to_max_shape(shape:tuple[sint, ...]) -> tuple[int, ...]: return tuple(int(x.vmax) if isinstance(x, UOp) else x for x in shape)
 
-def select_dtype(u): return (dtypes.long if u.overflows(dtypes.int32) else dtypes.int).vec(u.dtype.count)
+def select_dtype(u): return dtypes.long if u.overflows(dtypes.int32) else dtypes.int
 pm_lower_index_dtype = PatternMatcher([
   # There are no Unary ops at this point in symbolic, those are introduced later
   (UPat(GroupOp.Binary, name="u", src=(UPat.var("x").cast(dtypes.weakint), UPat.var("y").cast(dtypes.weakint))), lambda u,x,y:
@@ -1658,7 +1640,7 @@ pm_lower_index_dtype = PatternMatcher([
     cond.where(x.cast(dt:=least_upper_dtype(x.dtype, y.dtype)), y.cast(dt)).cast(dtypes.weakint)),
   (UPat(Ops.RANGE, src=(UPat.var("end").cast(dtypes.weakint)), name="r"), lambda r,end: r.replace(dtype=end.dtype, src=(end,)).cast(dtypes.weakint)),
   (UPat(Ops.STACK, src=UPat().cast(dtypes.weakint), name="v"),
-    lambda v: v.replace(dtype=(dt:=select_dtype(v)), src=tuple(s.src[0].cast(dt.scalar()) for s in v.src)).cast(dtypes.weakint)),
+    lambda v: v.replace(dtype=(dt:=select_dtype(v)), src=tuple(s.src[0].cast(dt) for s in v.src)).cast(dtypes.weakint)),
   # special can only be int32
   (UPat(Ops.SPECIAL, src=(UPat.var("var").cast(dtypes.weakint),), name="u"),
     lambda u,var: u.replace(dtype=dtypes.int, src=(var,)).cast(dtypes.weakint)),

tinygrad/uop/render.py

@@ -133,7 +133,7 @@ def pyrender(ast:UOp) -> str:
   lst = list(ast.toposort())
 
   cmap = consumer_map_from_toposort(lst)
-  not_rendered = {Ops.CONST, Ops.DEVICE}
+  not_rendered = {Ops.CONST}
   always_rendered = {Ops.PARAM, Ops.LOAD, Ops.SPECIAL, Ops.RANGE, Ops.CONTIGUOUS, Ops.STACK,
                      Ops.BUFFER, Ops.COPY, Ops.CALL, Ops.FUNCTION, Ops.WHERE, Ops.END}
 

tinygrad/uop/spec.py

@@ -23,10 +23,10 @@ def validate_index(uidx:UOp, gate:UOp|None=None):
 
   # TODO: validate these
   # WEBGPU has a BITCAST in the index, PTX casts pointer to long
-  # VECTORIZE/GEP can't be properly modeled in z3 since it doesn't support vectors
+  # VECTORIZE can't be properly modeled in z3 since it doesn't support vectors
   # don't descend into PARAM shape metadata; only the PARAM value participates in index arithmetic
   for x in idx.toposort(gate=lambda x: x.op is not Ops.PARAM) | gate.toposort(gate=lambda x: x.op is not Ops.PARAM):
-    if x.op in {Ops.BITCAST, Ops.STACK, Ops.GEP} or (x.op is Ops.CAST and isinstance(x.src[0].dtype, PtrDType)): return True
+    if x.op in {Ops.BITCAST, Ops.STACK} or (x.op is Ops.CAST and isinstance(x.src[0].dtype, PtrDType)): return True
 
   # if all is good and CHECK_OOB=1, validate with z3
   from tinygrad.uop.validate import validate_index_with_z3
@@ -118,11 +118,11 @@ spec_shared = PatternMatcher([
 
 def is_device(d): return isinstance(d, str) or (isinstance(d, tuple) and all(isinstance(s, str) for s in d))
 
+def valid_gettuple(g:UOp, t:UOp):
+  return isinstance(g.arg, int) and 0 <= g.arg < len(t.src) and g.dtype == t.src[g.arg].dtype
+
 # these ops can exist in tensor but not programs. example: movement
 spec_tensor = PatternMatcher([
-  # DEVICE
-  (UPat(Ops.DEVICE, dtypes.void, (), name="d"), lambda d: is_device(d.arg)),
-
   # BUFFER
   (UPat(Ops.BUFFER, src=(UPat(),), name="buf"), lambda buf:
    (isinstance(buf.dtype, DType) and buf.src[0].dtype.scalar() == dtypes.weakint and is_device(buf.arg.device))
@@ -135,12 +135,13 @@ spec_tensor = PatternMatcher([
   (UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda x: isinstance(x.arg, str)),
 
   # CALL
-  (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.LINEAR, Ops.PROGRAM, Ops.COPY, Ops.CUSTOM_FUNCTION)),), allow_any_len=True), lambda: True),
+  (UPat(Ops.CALL, dtypes.void, src=(UPat((Ops.SINK, Ops.LINEAR, Ops.PROGRAM, Ops.COPY, Ops.CUSTOM_FUNCTION)),), allow_any_len=True), lambda: True),
 
   # FUNCTION + TUPLE must have void dtype, GETTUPLE can only appear on FUNCTION or TUPLE
   (UPat(Ops.FUNCTION, dtypes.void, src=(UPat(Ops.TUPLE),), allow_any_len=True), lambda: True),
   (UPat(Ops.TUPLE, dtypes.void), lambda: True),
-  (UPat(Ops.GETTUPLE, src=(UPat((Ops.FUNCTION, Ops.TUPLE)),), name="g"), lambda g: isinstance(g.arg, int)),
+  (UPat(Ops.GETTUPLE, src=(UPat(Ops.FUNCTION, src=(UPat(Ops.TUPLE, name="t"),), allow_any_len=True),), name="g"), valid_gettuple),
+  (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), valid_gettuple),
 
   # inputs to movement ops
   (UPat({Ops.ADD, Ops.MUL, Ops.CDIV, Ops.FLOORDIV}, dtype=dtypes.weakint), lambda: True),
@@ -188,9 +189,6 @@ spec_tensor = PatternMatcher([
 
 # these ops can exist in programs but not the tensor spec. example: LOAD
 spec_program = PatternMatcher([
-  # no more of these in programs
-  (UPat(Ops.GEP), lambda: False),
-
   # weakint is not allowed in programs
   (UPat(GroupOp.All, dtypes.weakint), lambda: False),
 
@@ -222,7 +220,7 @@ spec_full = PatternMatcher([
                         UPat(Ops.CONST, dtype=dtypes.weakint)), allow_any_len=True, name="bv"),
    lambda bv: isinstance(bv.arg, int)),
 
-  (UPat(Ops.CALL, src=(UPat((Ops.SLICE,)),), allow_any_len=True), lambda: True),
+  (UPat(Ops.CALL, dtypes.void, src=(UPat((Ops.SLICE,)),), allow_any_len=True), lambda: True),
 
   # codegen may end ranges after gpudims has replaced RANGE with SPECIAL.
   (UPat(Ops.END, src=(UPat(), UPat()), allow_any_len=True), lambda: True),
@@ -230,12 +228,9 @@ spec_full = PatternMatcher([
   # allow any AFTER
   (UPat(Ops.AFTER, src=(UPat(),), allow_any_len=True), lambda: True),
 
-  # expander: unroll/contract/gep/cat
+  # expander: unroll/contract
   (UPat((Ops.UNROLL, Ops.CONTRACT), src=(UPat(),)), lambda: True),
 
-  # GEP multi is supported here
-  (UPat(Ops.GEP, name="gep"), lambda gep: gep.dtype is dtypes.void or gep.dtype.vcount == len(gep.arg)),
-
   # all loads/stores
   (UPat((Ops.LOAD, Ops.STORE)), lambda: True),
 

tinygrad/uop/symbolic.py

@@ -2,13 +2,12 @@
 import math, struct
 from collections import defaultdict
 from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
-from tinygrad.dtype import PyConst, ConstType, dtypes, PtrDType, can_lossless_cast, Invalid
-from tinygrad.helpers import partition, all_same, prod, flatten, get_single_element, unwrap, IMAGE, dedup
+from tinygrad.dtype import PyConst, ConstType, dtypes, can_lossless_cast, Invalid
+from tinygrad.helpers import partition, all_same, prod, flatten, unwrap, IMAGE, dedup
+from tinygrad.uop.divandmod import div_and_mod_symbolic
 
 # TODO: symbolic shouldn't be importing from codegen
-from tinygrad.codegen.decomp.op import threefry2x32
 from tinygrad.codegen.decomp.transcendental import xpow
-from tinygrad.codegen.decomp.divandmod import div_and_mod_symbolic
 
 # ******** phase 1 of symbolic used to live in ops, it's the most generic folding rules ********
 
@@ -34,61 +33,72 @@ def fold_const_alu(a:UOp) -> UOp|None:
   vals = [const_arg(s) for s in a.src]
   return None if any(v is None for v in vals) else a.const_like(exec_alu(a.op, a.dtype, vals, False))
 
-invalid_pat = UPat(Ops.CONST, arg=Invalid, name="i")
-invalid_gate = UPat.var("cond").where(UPat.var("x"), invalid_pat)
+def _quotient_base(q:UOp, base:UOp, div:int) -> UOp|None:
+  # the B with q == B//div and B%div == base%div, or None. only such congruence is needed to recombine, and canonicalization
+  # moves consts freely: the quotient may be merged ((x//c + a)//div -> (x + a*c)//(c*div) for div>0) and shifted ((y + k*D)//D == y//D + k)
+  (q, s), (num, a) = q.pop_const(), base.pop_const()
+  if q.op is not Ops.FLOORDIV or q.src[1].op is not Ops.CONST: return None
+  if div > 0 and num.op is Ops.FLOORDIV and num.src[1].op is Ops.CONST and q.src[1].arg == (c:=num.src[1].arg)*div: num, a, D = num.src[0], a*c, c*div
+  elif q.src[1].arg == div: D = div
+  else: return None
+  (x, xa), (p, pa) = num.pop_const(), q.src[0].pop_const()
+  if p is not x or (t:=xa + a - pa) % D: return None
+  return base - k*div if (k:=t//D - s) else base
 
 def fold_add_divmod_recombine(x:UOp) -> UOp|None:
+  # a scaled mod (base%div)*mul recombines with a partner q*(div*mul) carrying the quotient of a b == base (mod div):
+  #   q == b//div     -> b*mul              (full recombine)
+  #   q == (b//div)%d -> (b%(div*d))*mul    (partial recombine into a wider mod, needs d>0)
   terms = list(x.split_uop(Ops.ADD))
   for i,u in enumerate(terms):
-    if u.op is Ops.FLOORMOD and u.src[1].op is Ops.CONST: base, div, mul = u.src[0], u.src[1].arg, 1
-    elif u.op is Ops.MUL and u.src[1].op is Ops.CONST and (m:=u.src[0]).op is Ops.FLOORMOD and m.src[1].op is Ops.CONST:
-      base, div, mul = m.src[0], m.src[1].arg, u.src[1].arg
-    else: continue
+    mod, mul = u.pop_const(Ops.MUL)
+    if mod.op is not Ops.FLOORMOD or mod.src[1].op is not Ops.CONST: continue
+    base, div = mod.src[0], mod.src[1].arg
     for j,v in enumerate(terms):
-      if i == j: continue
-      if v.op is not Ops.MUL or v.src[1].op is not Ops.CONST or v.src[1].arg != div*mul: continue
-      q, exact = v.src[0], False
-      # (base%div)*mul + (base//div)*(div*mul) -> base*mul
-      if q.op is Ops.FLOORDIV and q.src[1].op is Ops.CONST and q.src[1].arg == div: exact = q.src[0] is base
-      # ((base//d)%div)*mul + (base//(d*div))*(div*mul) -> (base//d)*mul if div>0
-      if not exact and div > 0 and base.op is Ops.FLOORDIV and base.src[1].op is Ops.CONST:
-        exact = q.op is Ops.FLOORDIV and q.src[1].op is Ops.CONST and q.src[0] is base.src[0] and q.src[1].arg == base.src[1].arg*div
-      if exact: return (base*mul).usum(*[t for k,t in enumerate(terms) if k not in (i,j)])
-      # ((base//div)%d)*(div*mul) + (base%div)*mul -> (base%(div*d))*mul
-      if div > 0 and q.op is Ops.FLOORMOD and q.src[1].op is Ops.CONST and (d:=q.src[1].arg) > 0 and q.src[0].op is Ops.FLOORDIV:
-        if q.src[0].src[0] is base and q.src[0].src[1].op is Ops.CONST and q.src[0].src[1].arg == div:
-          return ((base % (div*d))*mul).usum(*[t for k,t in enumerate(terms) if k not in (i,j)])
+      q, scale = v.pop_const(Ops.MUL)
+      if i == j or scale != div*mul: continue
+      rest = [t for k,t in enumerate(terms) if k not in (i,j)]
+      if (b:=_quotient_base(q, base, div)) is not None: return (b*mul).usum(*rest)
+      if q.op is Ops.FLOORMOD and q.src[1].op is Ops.CONST and (d:=q.src[1].arg) > 0 and (b:=_quotient_base(q.src[0], base, div)) is not None:
+        return ((b % (div*d))*mul).usum(*rest)
   return None
 
+# an invalid index is cond.where(idx, Invalid) in weakint. the consumer reads cond back off the WHERE with UOp.get_valid,
+# so casts and comparisons of a gated index can drop the gate: when the index is invalid the result is never used
+invalid_idx_gate = UPat().where(UPat.var("x"), UPat(Ops.CONST, dtypes.weakint, arg=Invalid))
+pm_index_invalid = PatternMatcher([
+  (invalid_idx_gate.cast(name="cast"), lambda x,cast: x.cast(cast.dtype)),
+  (UPat(GroupOp.Comparison, src=(invalid_idx_gate, UPat.var("y")), name="alu"), lambda x,y,alu: x.alu(alu.op,y)),
+  (UPat(GroupOp.Comparison, src=(UPat.var("y"), invalid_idx_gate), name="alu"), lambda x,y,alu: y.alu(alu.op,x)),
+])
+
+# everywhere else Invalid poisons the value: ops move inside the gate so the Invalid reaches the LOAD/STORE and folds there.
 # this needs to be before symbolic so that 0*something_that_might_be_invalid doesnt become 0
-propagate_invalid = PatternMatcher([
-  # propagate invalid, push it past children
-  (invalid_gate.cast(name="cast"), lambda i,x,cond,cast: x.cast(cast.dtype) if i.dtype is dtypes.weakint else None),
-  (UPat(GroupOp.Unary, src=(invalid_gate,), name="alu"), lambda cond,x,alu,i: cond.where(x.alu(alu.op), i)),
-  (UPat(GroupOp.Binary-GroupOp.Comparison, src=(invalid_gate, UPat.var("y")), name="alu"), lambda cond,x,y,alu,i: cond.where(x.alu(alu.op,y), i)),
-  (UPat(GroupOp.Binary-GroupOp.Comparison, src=(UPat.var("y"), invalid_gate), name="alu"), lambda cond,x,y,alu,i: cond.where(y.alu(alu.op,x), i)),
-  # TODO: when can this happen? and is it always safe to just drop invalid?
-  (UPat(GroupOp.Comparison, src=(invalid_gate, UPat.var("y")), name="alu"), lambda cond,x,y,alu,i:
-     x.alu(alu.op,y) if i.dtype is dtypes.weakint else cond.where(x.alu(alu.op,y), i.cast(dtypes.bool))),
-  (UPat(GroupOp.Comparison, src=(UPat.var("y"), invalid_gate), name="alu"), lambda cond,x,y,alu,i:
-     y.alu(alu.op,x) if i.dtype is dtypes.weakint else cond.where(y.alu(alu.op,x), i.cast(dtypes.bool))),
-  # alu with invalid -> invalid
-  (UPat(GroupOp.Unary, src=(invalid_pat,)), lambda i: i),
+invalid_pat = UPat(Ops.CONST, arg=Invalid, name="i")
+invalid_gate = UPat.var("cond").where(UPat.var("x"), invalid_pat)
+pm_data_invalid = PatternMatcher([
+  (UPat(GroupOp.Unary|{Ops.BITCAST}, src=(invalid_pat,), name="op"), lambda i,op: i.cast(op.dtype)),
+  (UPat(GroupOp.Unary|{Ops.BITCAST}, src=(invalid_gate,), name="op"), lambda cond,x,op,i: cond.where(op.replace(src=(x,)), i.cast(op.dtype))),
+  # binary ops move inside the gate, with Invalid cast to the result dtype (bool for comparisons)
+  (UPat(GroupOp.Binary, src=(invalid_gate, UPat.var("y")), name="alu"), lambda cond,x,y,alu,i: cond.where(x.alu(alu.op,y), i.cast(alu.dtype))),
+  (UPat(GroupOp.Binary, src=(UPat.var("y"), invalid_gate), name="alu"), lambda cond,x,y,alu,i: cond.where(y.alu(alu.op,x), i.cast(alu.dtype))),
   (UPat(GroupOp.Binary-GroupOp.Comparison, src=[invalid_pat, UPat()]), lambda i: i),
   # normalize where(cond, Invalid, val) -> where(~cond, val, Invalid)
   (UPat.var("cond").where(invalid_pat, UPat.var("val")), lambda cond, i, val: cond.logical_not().where(val, i) if val.arg != Invalid else i),
-  # lift Invalid out  # TODO: this `a is cond` is asymmetric to preserve the pattern
+  # lift Invalid out: a.where(cond.where(x, Invalid), c) -> (~a|cond).where(a.where(x, c), Invalid)
+  # when a is cond, ~a|cond is True and would drop the Invalid gate (losing the valid), so keep cond as the gate
   (UPat.var("a").where(invalid_gate, UPat.var("c")), lambda cond,i,x,a,c:
    (cond if a is cond else (a.logical_not()|cond)).where(a.where(x,c), i) if c.arg != Invalid else None),
   (UPat.var("a").where(UPat.var("b"), invalid_gate), lambda cond,i,x,a,b: (a|cond).where(a.where(b, x), i) if b.arg != Invalid else None),
-  (UPat(Ops.BITCAST, src=(invalid_pat,), name="bc"), lambda bc,i: i.cast(bc.dtype)),
-  (UPat(Ops.BITCAST, src=(invalid_gate,), name="bc"), lambda bc,cond,x,i: cond.where(x.bitcast(bc.dtype), i.bitcast(bc.dtype))),
   # fold gated LOAD/STORE
   (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat(), invalid_pat), allow_any_len=True).or_casted(), UPat())), lambda i: UOp(Ops.NOOP)),
   (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), invalid_pat), allow_any_len=True).or_casted(),), allow_any_len=True, name="x"),
     lambda x,i: x.src[1] if len(x.src) > 1 else x.const_like(0)),
 ])
 
+propagate_invalid = pm_index_invalid + pm_data_invalid
+
+# TODO: this does nothing now
 pm_remove_invalid = PatternMatcher([
   (UPat(Ops.CONST, arg=Invalid, name="i"), lambda i: i.const_like(0) if i.dtype.scalar() is not dtypes.weakint else None),
 ])
@@ -116,7 +126,7 @@ symbolic_simple = propagate_invalid + PatternMatcher([
    lambda x,c: x if c.arg == 0 else x.logical_not() if c.arg == 1 else x.const_like(True)),
   (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)).trunc(), lambda x: x),
   # ** zero folding **
-  (UPat.var("x") < UPat.var("x"), lambda x: x.const_like(False).cast(dtypes.bool.vec(x.dtype.count))), # x < x -> False
+  (UPat.var("x") < UPat.var("x"), lambda x: x.vconst_like(False, dtypes.bool)), # x < x -> False
   (UPat.var("x") % UPat.var("x"), lambda x: x.const_like(0)), # x%x -> 0
   (UPat.var("x") ^ UPat.var("x"), lambda x: x.const_like(0)), # x^x -> 0
   (UPat.var("x") & 0, lambda x: x.const_like(0)), # x&0 -> 0
@@ -127,12 +137,11 @@ symbolic_simple = propagate_invalid + PatternMatcher([
   ((UPat.var("x") & UPat.cvar("mask")) // UPat.cvar("c"),
    lambda x,mask,c: x // c.arg if c.arg > 0 and c.arg & (c.arg-1) == 0 and mask.arg | (c.arg-1) == -1 else None),
   (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)) != UPat.var("x"),
-   lambda x: x.const_like(False).cast(dtypes.bool.vec(x.dtype.count))), # x != x -> False (only ints)
+   lambda x: x.vconst_like(False, dtypes.bool)), # x != x -> False (only ints)
   # ** constant folding **
   (UPat(GroupOp.Unary, src=(UPat((Ops.CONST, Ops.STACK)),), name="a"), fold_const_alu),
+  # NOTE: THREEFRY(const,const) folds via its decomposition
   (UPat(GroupOp.Binary-{Ops.THREEFRY}, src=(UPat((Ops.CONST, Ops.STACK)),)*2, name="a"), fold_const_alu),
-  (UPat(Ops.THREEFRY, src=(UPat.cvar("x"), UPat.cvar("key")), name="a"),
-   lambda a, x, key: a.const_like(threefry2x32(x, key).simplify().arg)),
   (UPat(GroupOp.Ternary, src=(UPat((Ops.CONST, Ops.STACK)),)*3, name="a"), fold_const_alu),
   # bool MUL is AND, ADD/MAX is OR. prevents other rules to rewrite bool ADD/MUL incorrectly
   (UPat.var('x', dtype=dtypes.bool) * UPat.var('y', dtype=dtypes.bool), lambda x,y: x&y),
@@ -173,7 +182,7 @@ symbolic_simple = propagate_invalid + PatternMatcher([
   (UPat.cvar("gate").where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1),
   # a.where(b.where(c, d), d) -> (a & b).where(c, d)
   (UPat.var("a").where(UPat.var("b").where(UPat.var("c"), UPat.var("d")), UPat.var("d")), lambda a,b,c,d: (a&b).where(c,d)),
-  # STACK on INDEX CONST (TODO: remove all the GEP crap)
+  # STACK on INDEX CONST
   (UPat(Ops.STACK, src=UPat(Ops.INDEX, src=(UPat.var("src"), UPat(Ops.CONST))), name="stk"),
    lambda src,stk: src if stk.shape == src.shape and list(range(len(stk.src))) == [x.src[1].arg for x in stk.src] else None),
   # INDEX on STACK
@@ -201,40 +210,6 @@ def canonicalize_simplex(X:UOp) -> UOp|None:
     ret.append(u)
   return UOp.usum(*ret) if changed else None
 
-def gep_through_wmma(gep:UOp, wmma:UOp) -> UOp|None:
-  out_sz = prod(x[1] for x in wmma.arg[6][-1])
-  wmma_idxs = gep.arg[::out_sz]
-  for i in range(out_sz):
-    if tuple(x-i for x in gep.arg[i::out_sz]) != wmma_idxs: return None
-  tsrcs = []
-  for s,sz in zip(wmma.src, wmma.arg[6]):
-    src_args = []
-    ssz = prod(x[1] for x in sz)
-    for w in wmma_idxs: src_args += list(range((w//out_sz)*ssz, (w//out_sz)*ssz + ssz))
-    tsrcs.append(s.gep(tuple(src_args)))
-  return UOp(Ops.WMMA, gep.dtype, tuple(tsrcs), wmma.arg)
-
-gep_pushing = PatternMatcher([
-  # GEP/VECTORIZE, GEP/GEP, GEP/CONST
-  (UPat(Ops.GEP, name='g2').f(Ops.GEP, name='g1'),
-   lambda g1, g2: g2.src[0].gep(tuple(g2.arg[g1.arg[i]] for i in range(len(g1.arg))))),
-  (UPat(Ops.STACK, name='vec').f(Ops.GEP, name='gep'),
-   lambda gep, vec: UOp(Ops.STACK, gep.dtype, tuple(vec.src[i] for i in gep.arg)) if len(gep.arg) > 1 else vec.src[gep.arg[0]]),
-  (UPat.cvar("c").f(Ops.GEP, name="gep"), lambda gep, c: gep.const_like(c.arg)),
-  # GEP on void is skipped
-  (UPat(Ops.GEP, src=(UPat(dtype=dtypes.void, name="x"),)), lambda x: x),
-  # GEP in order is removed
-  (UPat(Ops.GEP, name="g"), lambda g: g.src[0] if not isinstance(g.dtype, PtrDType) and g.arg == tuple(range(g.src[0].dtype.count)) else None),
-  # push all GEPs through ALUs for index (TODO: remove this)
-  (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name='alu').f(Ops.GEP, dtype=dtypes.weakint, name='gep'),
-   lambda gep,alu: UOp(alu.op, alu.dtype.scalar().vec(gep.dtype.count), tuple(x.gep(gep.arg) for x in alu.src), alu.arg) \
-     if not isinstance(gep.dtype, PtrDType) and not isinstance(alu.dtype, PtrDType) else None),
-  # VECTORIZE on same GEP
-  (UPat(Ops.STACK, name="v", src=UPat(Ops.GEP, src=(UPat.var("x"),))), lambda v,x: x.gep(tuple(get_single_element(i.arg) for i in v.src))),
-  # push some GEPs through WMMAs
-  (UPat(Ops.WMMA, name="wmma").f(Ops.GEP, name="gep"), gep_through_wmma),
-])
-
 commutative = PatternMatcher([
   # ** COMMUTATIVE flipping (only for index) **
   # NOTE: this can break merging vector math by only flipping some of them
@@ -317,7 +292,7 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
                         else y.src for y in x.src[1:]]))))),
   # after with 1 src is just src[0]
   (UPat(Ops.AFTER, src=(UPat.var("s"),)), lambda s: s),
-])+div_and_mod_symbolic+gep_pushing
+])+div_and_mod_symbolic
 
 # ******** we take a small aside to "simplify_valid" to rewrite valids ********
 

tinygrad/uop/upat.py

@@ -33,7 +33,7 @@ def _get_clause(self:UPat, base:UOp, depth=0) -> UOp:
   if self.src is not None:
     # single match
     if len(self.src) == 1 and isinstance(self.src[0], tuple):
-      and_clause += [_get_clause(s, base.gep(i), depth) for i,s in enumerate(self.src[0])]
+      and_clause += [_get_clause(s, base.index(i), depth) for i,s in enumerate(self.src[0])]
     # repeat match
     elif len(self.src) == 1 and isinstance(self.src[0], itertools.repeat):
       it = UOp(Ops.NOOP, arg=f"ituop{depth}")
@@ -41,7 +41,7 @@ def _get_clause(self:UPat, base:UOp, depth=0) -> UOp:
       and_clause.append(UOp(Ops.RANGE, src=(match, it, base), arg="all([{0} for {1} in {2}.src])"))
     # multi match (fork)
     elif len(self.src) > 1 and all(isinstance(x, tuple) for x in self.src):
-      fork_cond = [UOp(Ops.AND, src=tuple([_get_clause(s, base.gep(i), depth) for i,s in enumerate(ss)])) for ss in self.src]
+      fork_cond = [UOp(Ops.AND, src=tuple([_get_clause(s, base.index(i), depth) for i,s in enumerate(ss)])) for ss in self.src]
       and_clause.append(UOp(Ops.OR, src=tuple(fork_cond)))
     else: raise RuntimeError("broken")
   return UOp(Ops.AND, src=tuple(and_clause))
@@ -114,7 +114,7 @@ pm_renderer = PatternMatcher([
     lambda r,x: r.replace(op=Ops.CUSTOM, src=(UOp(Ops.NOOP, arg="(" + ' and '.join(y.arg for y in x.src) + ")"),)+r.src[1:])),
 
   (UPat(Ops.CUSTOM, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=x.arg.format(*[y.arg for y in x.src]))),
-  (UPat(Ops.GEP, src=UPat(Ops.NOOP, name="x"), name="g"), lambda x,g: x.replace(arg=x.arg+f".src[{g.arg[0]}]"))
+  (UPat(Ops.INDEX, src=(UPat(Ops.NOOP, name="x"), UPat(Ops.CONST, name="c")), name="g"), lambda x,c,g: x.replace(arg=x.arg+f".src[{c.arg}]"))
 ], compiled=False)
 
 def _final_render(x:UOp, has_ctx:bool, depth=1) -> list[str]:

tinygrad/viz/serve.py

@@ -7,7 +7,7 @@ from dataclasses import dataclass, field
 from urllib.parse import parse_qs, urlparse
 from http.server import BaseHTTPRequestHandler
 from typing import Any, TypedDict, TypeVar, Generator, Callable
-from tinygrad.helpers import colored, getenv, tqdm, unwrap, word_wrap, TRACEMETA, ProfileEvent, ProfileRangeEvent, TracingKey, ProfilePointEvent, temp
+from tinygrad.helpers import colored, getenv, unwrap, word_wrap, TRACEMETA, ProfileEvent, ProfileRangeEvent, TracingKey, ProfilePointEvent, temp
 from tinygrad.helpers import printable, Context, START_TIME, NO_COLOR, ansistrip
 from tinygrad.renderer.amd.dsl import Inst
 from tinygrad.renderer.amd import detect_format
@@ -119,8 +119,8 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]:
   graph: dict[int, dict] = {}
   excluded: set[UOp] = set()
   for u in (toposort:=x.toposort()):
-    # always exclude DEVICE/CONST
-    if u.op in {Ops.DEVICE, Ops.CONST} and u is not x: excluded.add(u)
+    # always exclude CONST
+    if u.op is Ops.CONST and u is not x: excluded.add(u)
     if u.op is Ops.STACK and len(u.src) == 0: excluded.add(u)
     # exclude RESHAPE/EXPAND that only serve to broadcast a CONST
     if u.op in {Ops.RESHAPE, Ops.EXPAND} and len(u.src) >= 1 and u.src[0] in excluded and u is not x: excluded.add(u)
@@ -149,7 +149,7 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]:
       if u.op in {Ops.CALL, Ops.FUNCTION}:
         label += f"\n{u.src[0].key.hex()[:8]}"
       if u.op in {Ops.INDEX, Ops.STAGE}:
-        if len(u.toposort()) < 30: label += f"\n{u.render()}"
+        label += f"\n{u.render()}" if sum(len(s.toposort()) for s in u.src[1:]) < 30 else "\nINDEX TOO LARGE"
         ranges: list[UOp] = []
         for us in u.src[1:]: ranges += [s for s in us.toposort() if s.op in {Ops.RANGE, Ops.SPECIAL}]
         if ranges: label += "\n"+' '.join([f"{s.render()}={s.vmax+1}" for s in ranges])
@@ -182,7 +182,7 @@ def get_full_rewrite(data:VizData, ctx:TrackedGraphRewrite, depth:int|None=None)
   next_sink = _reconstruct(data, ctx.sink, depth=depth)
   yield {"graph":uop_to_json(data, next_sink), "uop":pystr(next_sink), "change":None, "diff":None, "upat":None, "_sink":next_sink}
   replaces: dict[UOp, UOp] = {}
-  for u0_num,u1_num,upat_loc,dur in tqdm(ctx.matches, disable=not ctx.matches):
+  for u0_num,u1_num,upat_loc,dur in ctx.matches:
     replaces[u0:=_reconstruct(data, u0_num, depth=depth)] = u1 = _reconstruct(data, u1_num, depth=depth)
     try: new_sink = next_sink.substitute(replaces, walk=ctx.walk, enter_calls=ctx.enter_calls)
     except RuntimeError as e: new_sink = UOp(Ops.NOOP, arg=str(e))