failing tests

* delete the ShapeTracker

* fix tests

* fix more

* fix gc test

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

@@ -41,6 +41,10 @@ inputs:
     description: "Install LLVM?"
     required: false
     default: 'false'
+  mesa:
+    description: "Install mesa"
+    required: false
+    default: 'false'
 runs:
   using: "composite"
   steps:
@@ -289,3 +293,13 @@ runs:
       if: inputs.llvm == 'true' && runner.os == 'macOS'
       shell: bash
       run: brew install llvm@20
+
+    # **** mesa ****
+    - name: Install mesa (linux)
+      if: inputs.mesa == 'true' && runner.os == 'Linux'
+      shell: bash
+      run: sudo curl -L https://github.com/sirhcm/tinymesa/releases/download/tinymesa-32dc66c/libtinymesa_cpu-mesa-25.2.4-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
+    - name: Install mesa (macOS)
+      if: inputs.mesa == 'true' && runner.os == 'macOS'
+      shell: bash
+      run: brew install sirhcm/tinymesa/tinymesa

.github/workflows/autogen.yml

@@ -36,8 +36,9 @@ jobs:
         cuda: 'true'
         webgpu: 'true'
         llvm: 'true'
+        pydeps: 'pyyaml mako'
     - name: Install autogen support packages
-      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev
+      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev llvm-20-dev
     - name: Verify OpenCL autogen
       run: |
         cp tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
@@ -89,3 +90,8 @@ jobs:
         cp tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
         ./autogen_stubs.sh llvm
         diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
+    - name: Verify mesa autogen
+      run: |
+        cp tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
+        ./autogen_stubs.sh mesa
+        diff /tmp/mesa.py.bak tinygrad/runtime/autogen/mesa.py

.github/workflows/benchmark.yml

@@ -51,17 +51,18 @@ jobs:
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
     - name: reset process replay
       run: python3.11 test/external/process_replay/reset.py
+    - name: Print macOS version
+      run: sw_vers
     - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
     - name: Run Stable Diffusion without fp16
-      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=900 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
+      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
     - name: Run Stable Diffusion v2
       # TODO: very slow step time
-      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=10000 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
+      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
     # process replay can't capture this, the graph is too large
-    # TODO: too slow
-    # - name: Run SDXL
-    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+    - name: Run SDXL
+      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
     - name: Run model inference benchmark
       run: METAL=1 python3.11 test/external/external_model_benchmark.py
     - name: Test speed vs torch

.github/workflows/mlperf.yml

@@ -12,7 +12,7 @@ jobs:
   run_script_job:
     runs-on: [self-hosted, Linux, tinybox]
     if: github.repository_owner == 'tinygrad'
-    timeout-minutes: 360
+    timeout-minutes: 720
 
     steps:
     - name: Checkout Code

.github/workflows/test.yml

@@ -89,64 +89,65 @@ jobs:
         clang -O2 recognize.c -lm -o recognize
         cat test/models/efficientnet/Chicken.jpg | ./recognize | grep cock
 
-  torchbackend:
-    name: Torch Backend Tests
-    runs-on: ubuntu-latest
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_minimal
-        pydeps: "pillow torchvision expecttest"
-        llvm: 'true'
-    - name: Install ninja
-      run: |
-        sudo apt update || true
-        sudo apt install -y --no-install-recommends ninja-build
-    - name: Lint with ruff
-      run: |
-        pip3 install --upgrade --force-reinstall ruff==0.11.0
-        python3 -m ruff check extra/torch_backend/backend.py
-    - name: Test one op
-      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
-    - name: Test ResNet-18
-      run: DEBUG=2 python3 extra/torch_backend/example.py
-    - name: My (custom) tests
-      run: python3 extra/torch_backend/test.py
-    - name: Test one op in torch tests
-      run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
-    - name: Test Ops with TINY_BACKEND
-      run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
-    - name: Test in-place operations on views
-      run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
-    - name: Test multi-gpu
-      run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
+  # TODO: fix the torch backend and reenable
+  # torchbackend:
+  #   name: Torch Backend Tests
+  #   runs-on: ubuntu-latest
+  #   timeout-minutes: 15
+  #   steps:
+  #   - name: Checkout Code
+  #     uses: actions/checkout@v4
+  #   - name: Setup Environment
+  #     uses: ./.github/actions/setup-tinygrad
+  #     with:
+  #       key: torch-backend-pillow-torchvision-et-pt
+  #       deps: testing_minimal
+  #       pydeps: "pillow torchvision expecttest"
+  #       llvm: 'true'
+  #   - name: Install ninja
+  #     run: |
+  #       sudo apt update || true
+  #       sudo apt install -y --no-install-recommends ninja-build
+  #   - name: Lint with ruff
+  #     run: |
+  #       pip3 install --upgrade --force-reinstall ruff==0.11.0
+  #       python3 -m ruff check extra/torch_backend/backend.py
+  #   - name: Test one op
+  #     run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
+  #   - name: Test ResNet-18
+  #     run: DEBUG=2 python3 extra/torch_backend/example.py
+  #   - name: My (custom) tests
+  #     run: python3 extra/torch_backend/test.py
+  #   - name: Test one op in torch tests
+  #     run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
+  #   - name: Test Ops with TINY_BACKEND
+  #     run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
+  #   - name: Test in-place operations on views
+  #     run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
+  #   - name: Test multi-gpu
+  #     run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
 
-  torchbackendmore:
-    name: Torch Backend Tests More
-    runs-on: ubuntu-latest
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_minimal
-        llvm: 'true'
-    - name: Install ninja
-      run: |
-        sudo apt update || true
-        sudo apt install -y --no-install-recommends ninja-build
-    - name: Test beautiful_mnist in torch with TINY_BACKEND
-      run: CPU=1 CPU_LLVM=1 TARGET_EVAL_ACC_PCT=96.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
-    - name: Test some torch tests (expect failure)
-      run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
+  # torchbackendmore:
+  #   name: Torch Backend Tests More
+  #   runs-on: ubuntu-latest
+  #   timeout-minutes: 15
+  #   steps:
+  #   - name: Checkout Code
+  #     uses: actions/checkout@v4
+  #   - name: Setup Environment
+  #     uses: ./.github/actions/setup-tinygrad
+  #     with:
+  #       key: torch-backend-pillow-torchvision-et-pt
+  #       deps: testing_minimal
+  #       llvm: 'true'
+  #   - name: Install ninja
+  #     run: |
+  #       sudo apt update || true
+  #       sudo apt install -y --no-install-recommends ninja-build
+  #   - name: Test beautiful_mnist in torch with TINY_BACKEND
+  #     run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+  #   - name: Test some torch tests (expect failure)
+  #     run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
 
   bepython:
     name: Python Backend
@@ -272,6 +273,8 @@ jobs:
     #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
     - name: Run Clip tests for SD MLPerf on NULL backend
       run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
+    - name: Run AMD emulated BERT training on NULL backend
+      run: EMULATE=AMD_RDNA4 NULL=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
     # TODO: support fake weights
     #- name: Run LLaMA 7B on 4 fake devices
     #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
@@ -307,10 +310,6 @@ jobs:
       run: python test/external/fuzz_symbolic.py
     - name: Fuzz Test fast idiv
       run: python test/external/fuzz_fast_idiv.py
-    - name: Fuzz Test shapetracker
-      run: CNT=50 python test/external/fuzz_shapetracker.py
-    - name: Fuzz Test shapetracker math
-      run: CNT=200 python test/external/fuzz_shapetracker_math.py
     - name: Fuzz Test shape ops
       run: python test/external/fuzz_shape_ops.py
 
@@ -384,8 +383,6 @@ jobs:
       #   run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/35ff4f4577002f2685e50c8346addae33fe8da27a41dd4d6a0f14d1f4b1af81b
       - name: Test openpilot LLVM compile
         run: CPU=1 CPU_LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
-      - name: Test openpilot compile4
-        run: NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 DEBUG=2 python3 examples/openpilot/compile4.py
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -675,7 +672,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        backend: [llvm, cpu, opencl]
+        backend: [llvm, cpu, opencl, lvp]
 
     name: Linux (${{ matrix.backend }})
     runs-on: ubuntu-22.04
@@ -689,9 +686,10 @@ jobs:
           key: ${{ matrix.backend }}-minimal
           deps: testing_minimal
           opencl: ${{ matrix.backend == 'opencl' && 'true' }}
-          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
+          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
+          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
       - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'CL=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'CL=1' || matrix.backend == 'lvp' && 'CPU=1\nCPU_LVP=1' }}" >> $GITHUB_ENV
       - name: Check Device.DEFAULT and print some source
         run: |
           python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
@@ -893,7 +891,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        backend: [metal, llvm, cpu]
+        backend: [metal, llvm, cpu, lvp]
     name: MacOS (${{ matrix.backend }})
     runs-on: macos-15
     timeout-minutes: 20
@@ -906,12 +904,13 @@ jobs:
           key: macos-${{ matrix.backend }}-minimal
           deps: testing_minimal
           pydeps: "capstone"
-          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
+          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
+          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
       - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'metal' && 'METAL=1'}}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'metal' && 'METAL=1' || matrix.backend == 'lvp' && 'CPU=1\nCPU_LVP=1' }}" >> $GITHUB_ENV
       - name: Check Device.DEFAULT and print some source
         run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
+          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU','LVP':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
           DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
       - name: Run pytest (${{ matrix.backend }})
         run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20

.gitignore

@@ -38,6 +38,7 @@ extra/huggingface_onnx/models/*
 extra/huggingface_onnx/*.yaml
 extra/weights
 venv
+venv_sd_mlperf
 examples/**/net.*[js,json]
 examples/**/*.safetensors
 node_modules

autogen_stubs.sh

@@ -279,9 +279,9 @@ generate_llvm() {
     --clang-args="$(llvm-config-14 --cflags)" \
     -o "$BASE/llvm.py"
 
-  sed -i "s\import ctypes\import ctypes, tinygrad.runtime.support.llvm as llvm_support\g" "$BASE/llvm.py"
+  sed -i "s\import ctypes\import ctypes, tinygrad.runtime.support.llvm as llvm_support, tinygrad.helpers as helpers\g" "$BASE/llvm.py"
   sed -i "s\FIXME_STUB\llvm\g" "$BASE/llvm.py"
-  sed -i "s\FunctionFactoryStub()\ctypes.CDLL(llvm_support.LLVM_PATH)\g" "$BASE/llvm.py"
+  sed -i "s\FunctionFactoryStub()\ctypes.CDLL(llvm_support.LLVM_PATH, ctypes.RTLD_GLOBAL if helpers.OSX else ctypes.DEFAULT_MODE)\g" "$BASE/llvm.py"
 
   fixup "$BASE/llvm.py"
 }
@@ -461,6 +461,85 @@ generate_libusb() {
   python3 -c "import tinygrad.runtime.autogen.libusb"
 }
 
+generate_mesa() {
+  MESA_TAG="mesa-25.2.4"
+  MESA_SRC=/tmp/mesa-$MESA_TAG
+  TINYMESA_TAG=tinymesa-32dc66c
+  TINYMESA_DIR=/tmp/tinymesa-$MESA_TAG-$TINYMESA_TAG/
+  TINYMESA_SO=$TINYMESA_DIR/libtinymesa_cpu.so
+  if [ ! -d "$MESA_SRC" ]; then
+    git clone --depth 1 --branch $MESA_TAG https://gitlab.freedesktop.org/mesa/mesa.git $MESA_SRC
+    pushd .
+    cd $MESA_SRC
+    git reset --hard $MESA_COMMIT_HASH
+    # clang 14 doesn't support packed enums
+    sed -i "s/enum \w\+ \(\w\+\);$/uint8_t \1;/" $MESA_SRC/src/nouveau/headers/nv_device_info.h
+    sed -i "s/enum \w\+ \(\w\+\);$/uint8_t \1;/" $MESA_SRC/src/nouveau/compiler/nak.h
+    sed -i "s/nir_instr_type \(\w\+\);/uint8_t \1;/" $MESA_SRC/src/compiler/nir/nir.h
+    mkdir -p gen/util/format
+    python3 src/util/format/u_format_table.py src/util/format/u_format.yaml --enums > gen/util/format/u_format_gen.h
+    python3 src/compiler/nir/nir_opcodes_h.py > gen/nir_opcodes.h
+    python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
+    python3 src/compiler/nir/nir_intrinsics_indices_h.py --outdir gen
+    python3 src/compiler/nir/nir_builder_opcodes_h.py > gen/nir_builder_opcodes.h
+    python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
+    python3 src/compiler/builtin_types_h.py gen/builtin_types.h
+    popd
+  fi
+
+  if [ ! -d "$TINYMESA_DIR" ]; then
+    mkdir $TINYMESA_DIR
+    curl -L https://github.com/sirhcm/tinymesa/releases/download/$TINYMESA_TAG/libtinymesa_cpu-$MESA_TAG-linux-amd64.so -o $TINYMESA_SO
+  fi
+
+  clang2py -k cdefstu \
+    $MESA_SRC/src/compiler/nir/nir.h \
+    $MESA_SRC/src/compiler/nir/nir_builder.h \
+    $MESA_SRC/src/compiler/nir/nir_shader_compiler_options.h \
+    $MESA_SRC/src/compiler/nir/nir_serialize.h \
+    $MESA_SRC/gen/nir_intrinsics.h \
+    $MESA_SRC/src/nouveau/headers/nv_device_info.h \
+    $MESA_SRC/src/nouveau/compiler/nak.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_passmgr.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_misc.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_type.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_init.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_nir.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_struct.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_jit_types.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_flow.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_const.h \
+    $MESA_SRC/src/compiler/glsl_types.h \
+    $MESA_SRC/src/util/blob.h \
+    $MESA_SRC/src/util/ralloc.h \
+    --clang-args="-DHAVE_ENDIAN_H -DHAVE_STRUCT_TIMESPEC -DHAVE_PTHREAD -I$MESA_SRC/src -I$MESA_SRC/include -I$MESA_SRC/gen -I$MESA_SRC/src/compiler/nir -I$MESA_SRC/src/gallium/auxiliary -I$MESA_SRC/src/gallium/include -I$(llvm-config-20 --includedir)" \
+    -l $TINYMESA_SO \
+    -o $BASE/mesa.py
+
+  LVP_NIR_OPTIONS=$(./extra/mesa/lvp_nir_options.sh $MESA_SRC)
+
+  fixup $BASE/mesa.py
+  patch_dlopen $BASE/mesa.py tinymesa_cpu "(BASE:=os.getenv('MESA_PATH', f\"/usr{'/local/' if helpers.OSX else '/'}lib\"))+'/libtinymesa_cpu'+(EXT:='.dylib' if helpers.OSX else '.so')" "f'{BASE}/libtinymesa{EXT}'" "f'{brew_prefix()}/lib/libtinymesa_cpu.dylib'"
+  echo "lvp_nir_options = gzip.decompress(base64.b64decode('$LVP_NIR_OPTIONS'))" >> $BASE/mesa.py
+  cat <<EOF | sed -i "/import ctypes.*/r /dev/stdin" $BASE/mesa.py
+def brew_prefix():
+  try: return subprocess.check_output(['brew', '--prefix', 'tinymesa']).decode().strip()
+  except Exception: return ''
+EOF
+  sed -i "/in_dll/s/.*/try: &\nexcept AttributeError: pass/" $BASE/mesa.py
+  sed -i "s/import ctypes/import ctypes, ctypes.util, os, gzip, base64, subprocess, tinygrad.helpers as helpers/" $BASE/mesa.py
+  sed -i "s/ctypes.CDLL('.\+')/(dll := _try_dlopen_tinymesa_cpu())/" $BASE/mesa.py
+  echo "def __getattr__(nm): raise AttributeError() if dll else FileNotFoundError(f'libtinymesa not found (MESA_PATH={BASE}). See https://github.com/sirhcm/tinymesa ($TINYMESA_TAG, $MESA_TAG)')" >> $BASE/mesa.py
+  sed -i "s/ctypes.glsl_base_type/glsl_base_type/" $BASE/mesa.py
+  # bitfield bug in clang2py
+  sed -i "s/('fp_fast_math', ctypes.c_bool, 9)/('fp_fast_math', ctypes.c_uint32, 9)/" $BASE/mesa.py
+  sed -i "s/('\(\w\+\)', pipe_shader_type, 8)/('\1', ctypes.c_ubyte)/" $BASE/mesa.py
+  sed -i "s/\([0-9]\+\)()/\1/" $BASE/mesa.py
+  sed -i "s/\(struct_nir_builder._pack_\) = 1/\1 = 0/" $BASE/mesa.py
+  python3 -c "import tinygrad.runtime.autogen.mesa"
+}
+
 if [ "$1" == "opencl" ]; then generate_opencl
 elif [ "$1" == "hip" ]; then generate_hip
 elif [ "$1" == "comgr" ]; then generate_comgr
@@ -484,6 +563,7 @@ elif [ "$1" == "pci" ]; then generate_pci
 elif [ "$1" == "vfio" ]; then generate_vfio
 elif [ "$1" == "webgpu" ]; then generate_webgpu
 elif [ "$1" == "libusb" ]; then generate_libusb
-elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu
+elif [ "$1" == "mesa" ]; then generate_mesa
+elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu; generate_mesa
 else echo "usage: $0 <type>"
 fi

examples/openpilot/compile4.py

@@ -1,47 +0,0 @@
-import sys
-from tinygrad import Tensor, fetch, GlobalCounters, dtypes
-from tinygrad.uop.ops import UOp
-from tinygrad.nn.onnx import OnnxRunner
-from tinygrad.schedule.rangeify import get_rangeify_map
-from tinygrad.engine.schedule import create_schedule_with_vars
-from tinygrad.engine.realize import run_schedule
-
-# NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 VIZ=1 DEBUG=2 python3 examples/openpilot/compile4.py
-
-OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
-OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
-
-if __name__ == "__main__":
-  onnx_file = fetch(OPENPILOT_MODEL)
-  run_onnx = OnnxRunner(onnx_file)
-
-  inputs = run_onnx.get_empty_input_data("npy", dtypes.float32)
-  out: Tensor = next(iter(run_onnx({k:v.to(None) for k,v in inputs.items()}).values())).to('cpu')
-  root = out.uop
-  targets = [x.uop for x in inputs.values()]
-  print(targets)
-
-  # TODO: abstract this from gradient?
-
-  # compute the target path (top down)
-  in_target_path: dict[UOp, bool] = {}
-  for u in root.toposort(): in_target_path[u] = any(x in targets or in_target_path[x] for x in u.src)
-  independent_set = {}
-  for u in root.toposort():
-    if in_target_path[u]:
-      for s in u.src:
-        if not in_target_path[s]:
-          independent_set[s] = None
-  independent = UOp.sink(*independent_set.keys())
-  kernelized = get_rangeify_map(independent)
-  independent = independent.substitute(kernelized)
-  schedule, var_vals = create_schedule_with_vars(independent)
-  run_schedule(schedule)
-
-  print("**** real ****")
-  GlobalCounters.reset()
-  out.uop = root.substitute(kernelized)
-  out.kernelize()
-
-  # realize
-  out.realize()

examples/stable_diffusion.py

@@ -269,7 +269,8 @@ if __name__ == "__main__":
 
   # load in weights
   with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
-    load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], verbose=False, strict=False, realize=False)
+    model_bin = fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt')
+    load_state_dict(model, torch_load(model_bin)['state_dict'], verbose=False, strict=False, realize=False)
 
     if args.fp16:
       for k,v in get_state_dict(model).items():

extra/mesa/lvp_nir_options.sh

@@ -0,0 +1,23 @@
+#!/bin/sh
+
+if [ "$#" -ne 1 ] || ! [ -d $1 ]; then
+  echo "usage: $0 MESA_PREFIX"
+  exit 1
+fi
+
+TMP=$(mktemp)
+trap 'rm -f "$TMP"' EXIT
+
+(
+  cat <<EOF
+#define HAVE_ENDIAN_H
+#define HAVE_STRUCT_TIMESPEC
+#define HAVE_PTHREAD
+#include <unistd.h>
+#include "nir_shader_compiler_options.h"
+#include "compiler/shader_enums.h"
+EOF
+  sed -n '/struct nir_shader_compiler_options/,/^}/{p;/^}/q}' $1/src/gallium/drivers/llvmpipe/lp_screen.c
+  echo "int main(void) { write(1, &gallivm_nir_options, sizeof(gallivm_nir_options)); }"
+) | cc -x c -o $TMP - -I$1/src/compiler/nir -I$1/src -I$1/include && $TMP | gzip | base64 -w0
+

extra/tinyfs/fetch_file.py

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

extra/tinyfs/fetch_raid.py

@@ -0,0 +1,39 @@
+import json, multiprocessing
+from pathlib import Path
+
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import tqdm, getenv
+
+raid_root = Path(getenv("RAID_ROOT", "/raid"))
+
+def fetch_file(item):
+  path, info = item
+  h, size = info["hash"], info["size"]
+
+  path = raid_root / Path(path)
+  path.parent.mkdir(parents=True, exist_ok=True)
+
+  try:
+    pt = Tensor(bytes.fromhex(h), device="CPU").load(size).to(f"disk:{path.as_posix()}").realize()
+  except Exception as e:
+    print(f"error fetching {path}, {h}, {size}: {e}")
+    raise
+
+  pt.uop.buffer.deallocate()
+
+def fetch_mapping():
+  mapping_tensor = Tensor(bytes.fromhex("d734f5e3be9f1e9d863bfaa4fc6c1ef2")).load(175866113).realize()
+  mapping = mapping_tensor.data().tobytes().decode()
+  mapping = json.loads(mapping)
+  mapped_files = mapping.items()
+  return list(mapped_files)
+
+if __name__ == "__main__":
+  with multiprocessing.Pool(processes=1) as pool:
+    mapped_files = pool.apply(fetch_mapping)
+
+  print(f"fetched mapping for {len(mapped_files)} files")
+
+  with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
+    for _ in tqdm(pool.imap_unordered(fetch_file, mapped_files), total=len(mapped_files)):
+      pass

extra/tinyfs/upload_raid.py

@@ -0,0 +1,31 @@
+from pathlib import Path
+import multiprocessing, json
+
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import tqdm
+
+raid_root = Path("/raid")
+
+def upload_file(path: Path):
+  pt = Tensor(path).realize()
+  h = pt.store().realize()
+  pt.uop.realized.deallocate()
+  return h.data().hex(), path, pt.nbytes()
+
+if __name__ == "__main__":
+  raid_files = sorted([p for p in raid_root.rglob("*") if p.is_file()])
+  print(f"found {len(raid_files)} files in /raid")
+
+  mapping = {}
+  with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
+    for h, p, s in tqdm(pool.imap_unordered(upload_file, raid_files), total=len(raid_files)):
+      mapping[p.relative_to(raid_root).as_posix()] = {"hash": h, "size": s}
+
+  # sort the mapping by key
+  mapping = dict(sorted(mapping.items()))
+
+  mapping = json.dumps(mapping).encode()
+  mapping_tensor = Tensor(mapping, device="CPU")
+  h = mapping_tensor.store().realize()
+
+  print(f"final hash: {h.data().hex()}, size: {len(mapping)}")

extra/torch_backend/backend.py

@@ -642,10 +642,11 @@ def get_real_tinygrad_buffers():
 torch.nn.modules.module.register_module_buffer_registration_hook(register_torch_buffer)
 
 from torch.nn.modules import Module
-def backward_hook(model:Module, _grad_input, _grad_out):
-  grads_to_realize = [unwrap(p.grad) for p in model.parameters() if p.grad is not None]
-  if len(grads_to_realize): Tensor.realize(*grads_to_realize)
-def module_hook(module:Module, _name, _submodule): module.register_backward_hook(backward_hook)
+def param_hook(_grad):
+  if _grad is not None and _grad.is_tiny: Tensor.realize(unwrap(_grad))
+def module_hook(module:Module, _name, _submodule):
+  for param in _submodule.parameters(recurse=False):
+    if param.requires_grad: param.register_hook(param_hook)
 torch.nn.modules.module.register_module_module_registration_hook(module_hook)
 
 def realize_optimizer_step(optimizer: torch.optim.Optimizer, *args, **kwargs):

extra/usbgpu/tbgpu/installer/.gitignore

@@ -0,0 +1,13 @@
+xcuserdata/
+
+**/*.xcodeproj/project.xcworkspace/*
+!**/*.xcodeproj/project.xcworkspace/xcshareddata
+
+**/*.xcodeproj/project.xcworkspace/xcshareddata/*
+!**/*.xcodeproj/project.xcworkspace/xcshareddata/WorkspaceSettings.xcsettings
+
+**/*.playground/playground.xcworkspace/*
+!**/*.playground/playground.xcworkspace/xcshareddata
+
+**/*.playground/playground.xcworkspace/xcshareddata/*
+!**/*.playground/playground.xcworkspace/xcshareddata/WorkspaceSettings.xcsettings

extra/usbgpu/tbgpu/installer/Shared/Assets.xcassets/AccentColor.colorset/Contents.json

@@ -0,0 +1,11 @@
+{
+  "colors" : [
+    {
+      "idiom" : "universal"
+    }
+  ],
+  "info" : {
+    "author" : "xcode",
+    "version" : 1
+  }
+}

extra/usbgpu/tbgpu/installer/Shared/Assets.xcassets/AppIcon.appiconset/Contents.json

@@ -0,0 +1,148 @@
+{
+  "images" : [
+    {
+      "idiom" : "iphone",
+      "scale" : "2x",
+      "size" : "20x20"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "3x",
+      "size" : "20x20"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "2x",
+      "size" : "29x29"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "3x",
+      "size" : "29x29"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "2x",
+      "size" : "40x40"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "3x",
+      "size" : "40x40"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "2x",
+      "size" : "60x60"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "3x",
+      "size" : "60x60"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "1x",
+      "size" : "20x20"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "2x",
+      "size" : "20x20"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "1x",
+      "size" : "29x29"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "2x",
+      "size" : "29x29"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "1x",
+      "size" : "40x40"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "2x",
+      "size" : "40x40"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "1x",
+      "size" : "76x76"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "2x",
+      "size" : "76x76"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "2x",
+      "size" : "83.5x83.5"
+    },
+    {
+      "idiom" : "ios-marketing",
+      "scale" : "1x",
+      "size" : "1024x1024"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "1x",
+      "size" : "16x16"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "2x",
+      "size" : "16x16"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "1x",
+      "size" : "32x32"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "2x",
+      "size" : "32x32"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "1x",
+      "size" : "128x128"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "2x",
+      "size" : "128x128"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "1x",
+      "size" : "256x256"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "2x",
+      "size" : "256x256"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "1x",
+      "size" : "512x512"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "2x",
+      "size" : "512x512"
+    }
+  ],
+  "info" : {
+    "author" : "xcode",
+    "version" : 1
+  }
+}

extra/usbgpu/tbgpu/installer/Shared/Assets.xcassets/Contents.json

@@ -0,0 +1,6 @@
+{
+  "info" : {
+    "author" : "xcode",
+    "version" : 1
+  }
+}

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

@@ -0,0 +1,19 @@
+import AppKit
+import SwiftUI
+
+final class AppDelegate: NSObject, NSApplicationDelegate {
+	func applicationShouldTerminateAfterLastWindowClosed(_ sender: NSApplication) -> Bool {
+		true
+	}
+}
+
+@main
+struct TinyGPUApp: App {
+	@NSApplicationDelegateAdaptor(AppDelegate.self) var appDelegate
+
+    var body: some Scene {
+        WindowGroup {
+			TinyGPUView()
+        }
+    }
+}

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

@@ -0,0 +1,33 @@
+import SwiftUI
+
+struct TinyGPUView: View {
+	@ObservedObject var viewModel = TinyGPUViewModel()
+
+	var body: some View {
+#if os(macOS)
+		VStack(alignment: .center) {
+			Text("TinyGPU Intsaller")
+				.padding()
+				.font(.title)
+			Text(self.viewModel.dextLoadingState)
+				.multilineTextAlignment(.center)
+			HStack {
+				Button(
+					action: {
+						self.viewModel.activateMyDext()
+					}, label: {
+						Text("Install extension")
+					}
+				)
+			}
+		}
+		.frame(width: 500, height: 200, alignment: .center)
+#endif
+	}
+}
+
+struct TinyGPUView_Previews: PreviewProvider {
+    static var previews: some View {
+		TinyGPUView()
+    }
+}

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

@@ -0,0 +1,149 @@
+import Foundation
+import os.log
+import SystemExtensions
+
+class TinyGPUDriverLoadingStateMachine {
+  enum State { case unloaded, activating, needsApproval, activated, activationError }
+}
+
+class TinyGPUViewModel: NSObject {
+
+	@Published private var state: TinyGPUDriverLoadingStateMachine.State = .unloaded
+	
+	override init() {
+		super.init()
+		refreshInitialDextState()
+	}
+	
+	private func refreshInitialDextState() {
+#if os(macOS)
+		Task.detached { [dextIdentifier] in
+		  let newState = Self.queryDextState(bundleID: dextIdentifier)
+		  await MainActor.run { self.state = newState }
+		}
+#endif
+	}
+
+#if os(macOS)
+	private static func queryDextState(bundleID: String) -> TinyGPUDriverLoadingStateMachine.State {
+		let tool = "/usr/bin/systemextensionsctl"
+		let p = Process()
+		p.executableURL = URL(fileURLWithPath: tool)
+		p.arguments = ["list"]
+
+		let pipe = Pipe()
+		p.standardOutput = pipe
+		p.standardError = Pipe()
+
+		do {
+			try p.run()
+			p.waitUntilExit()
+			let data = pipe.fileHandleForReading.readDataToEndOfFile()
+			guard let output = String(data: data, encoding: .utf8) else { return .unloaded }
+
+			// Look for our bundle id line
+			if let line = output.split(separator: "\n").first(where: { $0.contains(bundleID) }) {
+				if line.contains("[activated enabled]") { return .activated }
+				if line.contains("[activated waiting for user]") { return .needsApproval }
+				if line.contains("terminated waiting to uninstall") { return .unloaded }
+				return .activating
+			} else {
+				return .unloaded
+			}
+		} catch {
+			return .unloaded
+		}
+	}
+#endif
+
+	private let dextIdentifier: String = Bundle.main.bundleIdentifier! + ".Driver"
+
+	public var dextLoadingState: String {
+		switch state {
+		case .unloaded:
+			return "TinyGPUDriver isn't loaded."
+		case .activating:
+			return "Activating TinyGPUDriver, please wait."
+		case .needsApproval:
+			return "Please follow the prompt to approve TinyGPUDriver."
+		case .activated:
+			return "TinyGPUDriver has been activated and is ready to use. You can close the installer."
+		case .activationError:
+			return "TinyGPUDriver has experienced an error during activation.\nPlease check the logs to find the error."
+		}
+	}
+}
+
+extension TinyGPUViewModel: ObservableObject {
+
+#if os(macOS)
+	func activateMyDext() {
+		activateExtension(dextIdentifier)
+	}
+	
+	func deactivateMyDext() {
+		deactivateExtension(dextIdentifier)
+	}
+
+	func activateExtension(_ dextIdentifier: String) {
+
+		let request = OSSystemExtensionRequest
+			.activationRequest(forExtensionWithIdentifier: dextIdentifier,
+							   queue: .main)
+		request.delegate = self
+		OSSystemExtensionManager.shared.submitRequest(request)
+
+		self.state = .activating
+	}
+	
+	func deactivateExtension(_ dextIdentifier: String) {
+
+		let request = OSSystemExtensionRequest.deactivationRequest(forExtensionWithIdentifier: dextIdentifier, queue: .main)
+		request.delegate = self
+		OSSystemExtensionManager.shared.submitRequest(request)
+
+		self.state = .unloaded
+	}
+#endif
+}
+
+#if os(macOS)
+extension TinyGPUViewModel: OSSystemExtensionRequestDelegate {
+
+	func request(
+		_ request: OSSystemExtensionRequest,
+		actionForReplacingExtension existing: OSSystemExtensionProperties,
+		withExtension ext: OSSystemExtensionProperties) -> OSSystemExtensionRequest.ReplacementAction {
+
+		var replacementAction: OSSystemExtensionRequest.ReplacementAction
+
+		os_log("sysex actionForReplacingExtension: %@ %@", existing, ext)
+
+		// Add appropriate logic here to determine whether to replace the extension
+		// with the new extension. Common things to check for include
+		// testing whether the new extension's version number is newer than
+		// the current version number, or whether the bundleIdentifier is different.
+		// For simplicity, this sample always replaces the current extension
+		// with the new one.
+		replacementAction = .replace
+
+		self.state = .activating
+		return replacementAction
+	}
+
+	func requestNeedsUserApproval(_ request: OSSystemExtensionRequest) {
+		os_log("sysex requestNeedsUserApproval")
+		self.state = .needsApproval
+	}
+
+	func request(_ request: OSSystemExtensionRequest, didFinishWithResult result: OSSystemExtensionRequest.Result) {
+		os_log("sysex didFinishWithResult: %d", result.rawValue)
+		self.state = .activated
+	}
+
+	func request(_ request: OSSystemExtensionRequest, didFailWithError error: Error) {
+		os_log("sysex didFailWithError: %@", error.localizedDescription)
+		self.state = .activationError
+	}
+}
+#endif

extra/usbgpu/tbgpu/installer/TinyGPUDriverExtension.xcodeproj/project.pbxproj

@@ -0,0 +1,589 @@
+// !$*UTF8*$!
+{
+	archiveVersion = 1;
+	classes = {
+	};
+	objectVersion = 54;
+	objects = {
+
+/* Begin PBXBuildFile section */
+		0ACB55392E9CB880007029EF /* PCIDriverKit.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 0ACB55382E9CB880007029EF /* PCIDriverKit.framework */; };
+		54798269286A3512009785F6 /* CoreAudio.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 54798268286A3512009785F6 /* CoreAudio.framework */; };
+		549EB121286A1A37009D38AB /* TinyGPUViewModel.swift in Sources */ = {isa = PBXBuildFile; fileRef = 549EB11F286A1A37009D38AB /* TinyGPUViewModel.swift */; };
+		549EB123286A1D48009D38AB /* org.tinygrad.tinygpu.edriver.dext in Embed System Extensions */ = {isa = PBXBuildFile; fileRef = C5B7D9BC26128AC50089B4C3 /* org.tinygrad.tinygpu.edriver.dext */; settings = {ATTRIBUTES = (RemoveHeadersOnCopy, ); }; };
+		549EB131286A2B98009D38AB /* IOKit.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 549EB130286A2B98009D38AB /* IOKit.framework */; };
+		54E42BC8286A1697000E1E9A /* TinyGPUApp.swift in Sources */ = {isa = PBXBuildFile; fileRef = 54E42BB8286A1696000E1E9A /* TinyGPUApp.swift */; };
+		54E42BCA286A1697000E1E9A /* TinyGPUView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 54E42BB9286A1696000E1E9A /* TinyGPUView.swift */; };
+		54E42BCC286A1697000E1E9A /* Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = 54E42BBA286A1697000E1E9A /* Assets.xcassets */; };
+		C5B7D9C326128AC50089B4C3 /* TinyGPUDriver.cpp in Sources */ = {isa = PBXBuildFile; fileRef = C5B7D9C226128AC50089B4C3 /* TinyGPUDriver.cpp */; };
+		C5B7D9C526128AC50089B4C3 /* TinyGPUDriver.iig in Sources */ = {isa = PBXBuildFile; fileRef = C5B7D9C426128AC50089B4C3 /* TinyGPUDriver.iig */; };
+		C5C3BBB32612ACDC003C7BFE /* AudioDriverKit.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = C5C3BBB12612ACD3003C7BFE /* AudioDriverKit.framework */; };
+		C5C3BBB52612ACEF003C7BFE /* DriverKit.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = C5C3BBB42612ACEF003C7BFE /* DriverKit.framework */; };
+		C5D787AC261667FC006047E5 /* TinyGPUDriverUserClient.iig in Sources */ = {isa = PBXBuildFile; fileRef = C5D787AB261667FC006047E5 /* TinyGPUDriverUserClient.iig */; };
+		C5D787AE26168E59006047E5 /* TinyGPUDriverUserClient.cpp in Sources */ = {isa = PBXBuildFile; fileRef = C5D787AD26168D1E006047E5 /* TinyGPUDriverUserClient.cpp */; };
+/* End PBXBuildFile section */
+
+/* Begin PBXContainerItemProxy section */
+		549EB126286A1D66009D38AB /* PBXContainerItemProxy */ = {
+			isa = PBXContainerItemProxy;
+			containerPortal = C5B7D9B326128AC50089B4C3 /* Project object */;
+			proxyType = 1;
+			remoteGlobalIDString = C5B7D9BB26128AC50089B4C3;
+			remoteInfo = SimpleAudioDriver;
+		};
+/* End PBXContainerItemProxy section */
+
+/* Begin PBXCopyFilesBuildPhase section */
+		549EB122286A1D3A009D38AB /* Embed System Extensions */ = {
+			isa = PBXCopyFilesBuildPhase;
+			buildActionMask = 2147483647;
+			dstPath = "$(SYSTEM_EXTENSIONS_FOLDER_PATH)";
+			dstSubfolderSpec = 16;
+			files = (
+				549EB123286A1D48009D38AB /* org.tinygrad.tinygpu.edriver.dext in Embed System Extensions */,
+			);
+			name = "Embed System Extensions";
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+/* End PBXCopyFilesBuildPhase section */
+
+/* Begin PBXFileReference section */
+		0ACB55382E9CB880007029EF /* PCIDriverKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = PCIDriverKit.framework; path = System/DriverKit/System/Library/Frameworks/PCIDriverKit.framework; sourceTree = SDKROOT; };
+		54798268286A3512009785F6 /* CoreAudio.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = CoreAudio.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX13.0.sdk/System/Library/Frameworks/CoreAudio.framework; sourceTree = DEVELOPER_DIR; };
+		549EB11F286A1A37009D38AB /* TinyGPUViewModel.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TinyGPUViewModel.swift; sourceTree = "<group>"; usesTabs = 1; };
+		549EB130286A2B98009D38AB /* IOKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = IOKit.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX13.0.sdk/System/Library/Frameworks/IOKit.framework; sourceTree = DEVELOPER_DIR; };
+		549EB132286A2B9D009D38AB /* IOKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = IOKit.framework; path = Platforms/iPhoneOS.platform/Developer/SDKs/iPhoneOS16.0.sdk/System/Library/Frameworks/IOKit.framework; sourceTree = DEVELOPER_DIR; };
+		54E42BB8286A1696000E1E9A /* TinyGPUApp.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TinyGPUApp.swift; sourceTree = "<group>"; };
+		54E42BB9286A1696000E1E9A /* TinyGPUView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TinyGPUView.swift; sourceTree = "<group>"; };
+		54E42BBA286A1697000E1E9A /* Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = Assets.xcassets; sourceTree = "<group>"; };
+		54E42BC4286A1697000E1E9A /* TinyGPU.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = TinyGPU.app; sourceTree = BUILT_PRODUCTS_DIR; };
+		54E42BC6286A1697000E1E9A /* macOS.entitlements */ = {isa = PBXFileReference; lastKnownFileType = text.plist.entitlements; path = macOS.entitlements; sourceTree = "<group>"; };
+		C5B7D9BC26128AC50089B4C3 /* org.tinygrad.tinygpu.edriver.dext */ = {isa = PBXFileReference; explicitFileType = "wrapper.driver-extension"; includeInIndex = 0; path = org.tinygrad.tinygpu.edriver.dext; sourceTree = BUILT_PRODUCTS_DIR; };
+		C5B7D9BF26128AC50089B4C3 /* DriverKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = DriverKit.framework; path = Library/Frameworks/DriverKit.framework; sourceTree = DEVELOPER_DIR; };
+		C5B7D9C226128AC50089B4C3 /* TinyGPUDriver.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; path = TinyGPUDriver.cpp; sourceTree = "<group>"; usesTabs = 1; };
+		C5B7D9C426128AC50089B4C3 /* TinyGPUDriver.iig */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.iig; path = TinyGPUDriver.iig; sourceTree = "<group>"; };
+		C5B7D9C626128AC50089B4C3 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = Info.plist; sourceTree = "<group>"; };
+		C5B7D9CC26128ADA0089B4C3 /* AudioDriverKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = AudioDriverKit.framework; path = System/DriverKit/System/Library/Frameworks/AudioDriverKit.framework; sourceTree = SDKROOT; };
+		C5B7D9CE26128B150089B4C3 /* TinyGPUDriver.entitlements */ = {isa = PBXFileReference; lastKnownFileType = text.plist.entitlements; path = TinyGPUDriver.entitlements; sourceTree = "<group>"; };
+		C5C0063326178F98003345D8 /* AppKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = AppKit.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX12.0.sdk/System/Library/Frameworks/AppKit.framework; sourceTree = DEVELOPER_DIR; };
+		C5C006352617ACB8003345D8 /* CoreAudio.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = CoreAudio.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX12.0.sdk/System/Library/Frameworks/CoreAudio.framework; sourceTree = DEVELOPER_DIR; };
+		C5C3BBB12612ACD3003C7BFE /* AudioDriverKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = AudioDriverKit.framework; path = Platforms/DriverKit.platform/Developer/SDKs/DriverKit.MacOSX21.0.Internal.sdk/System/DriverKit/System/Library/Frameworks/AudioDriverKit.framework; sourceTree = DEVELOPER_DIR; };
+		C5C3BBB42612ACEF003C7BFE /* DriverKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = DriverKit.framework; path = Platforms/DriverKit.platform/Developer/SDKs/DriverKit.MacOSX21.0.Internal.sdk/System/DriverKit/System/Library/Frameworks/DriverKit.framework; sourceTree = DEVELOPER_DIR; };
+		C5D787AB261667FC006047E5 /* TinyGPUDriverUserClient.iig */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.iig; path = TinyGPUDriverUserClient.iig; sourceTree = "<group>"; };
+		C5D787AD26168D1E006047E5 /* TinyGPUDriverUserClient.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; path = TinyGPUDriverUserClient.cpp; sourceTree = "<group>"; };
+		C5D787B026169723006047E5 /* IOKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = IOKit.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX12.0.sdk/System/Library/Frameworks/IOKit.framework; sourceTree = DEVELOPER_DIR; };
+		C5D787B22616973F006047E5 /* SystemExtensions.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = SystemExtensions.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX12.0.sdk/System/Library/Frameworks/SystemExtensions.framework; sourceTree = DEVELOPER_DIR; };
+		C5D787B426169747006047E5 /* Foundation.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = Foundation.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX12.0.sdk/System/Library/Frameworks/Foundation.framework; sourceTree = DEVELOPER_DIR; };
+/* End PBXFileReference section */
+
+/* Begin PBXFrameworksBuildPhase section */
+		54E42BC1286A1697000E1E9A /* Frameworks */ = {
+			isa = PBXFrameworksBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+				54798269286A3512009785F6 /* CoreAudio.framework in Frameworks */,
+				549EB131286A2B98009D38AB /* IOKit.framework in Frameworks */,
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+		C5B7D9B926128AC50089B4C3 /* Frameworks */ = {
+			isa = PBXFrameworksBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+				C5C3BBB32612ACDC003C7BFE /* AudioDriverKit.framework in Frameworks */,
+				C5C3BBB52612ACEF003C7BFE /* DriverKit.framework in Frameworks */,
+				0ACB55392E9CB880007029EF /* PCIDriverKit.framework in Frameworks */,
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+/* End PBXFrameworksBuildPhase section */
+
+/* Begin PBXGroup section */
+		54E42BB7286A1696000E1E9A /* Shared */ = {
+			isa = PBXGroup;
+			children = (
+				54E42BB8286A1696000E1E9A /* TinyGPUApp.swift */,
+				54E42BB9286A1696000E1E9A /* TinyGPUView.swift */,
+				549EB11F286A1A37009D38AB /* TinyGPUViewModel.swift */,
+				54E42BBA286A1697000E1E9A /* Assets.xcassets */,
+			);
+			path = Shared;
+			sourceTree = "<group>";
+		};
+		54E42BC5286A1697000E1E9A /* macOS */ = {
+			isa = PBXGroup;
+			children = (
+				54E42BC6286A1697000E1E9A /* macOS.entitlements */,
+			);
+			path = macOS;
+			sourceTree = "<group>";
+		};
+		C5B7D9B226128AC50089B4C3 = {
+			isa = PBXGroup;
+			children = (
+				C5B7D9C126128AC50089B4C3 /* TinyGPUDriverExtension */,
+				54E42BB7286A1696000E1E9A /* Shared */,
+				54E42BC5286A1697000E1E9A /* macOS */,
+				C5B7D9BE26128AC50089B4C3 /* Frameworks */,
+				C5B7D9BD26128AC50089B4C3 /* Products */,
+			);
+			sourceTree = "<group>";
+			usesTabs = 1;
+		};
+		C5B7D9BD26128AC50089B4C3 /* Products */ = {
+			isa = PBXGroup;
+			children = (
+				C5B7D9BC26128AC50089B4C3 /* org.tinygrad.tinygpu.edriver.dext */,
+				54E42BC4286A1697000E1E9A /* TinyGPU.app */,
+			);
+			name = Products;
+			sourceTree = "<group>";
+		};
+		C5B7D9BE26128AC50089B4C3 /* Frameworks */ = {
+			isa = PBXGroup;
+			children = (
+				0ACB55382E9CB880007029EF /* PCIDriverKit.framework */,
+				54798268286A3512009785F6 /* CoreAudio.framework */,
+				549EB130286A2B98009D38AB /* IOKit.framework */,
+				549EB132286A2B9D009D38AB /* IOKit.framework */,
+				C5C006352617ACB8003345D8 /* CoreAudio.framework */,
+				C5C0063326178F98003345D8 /* AppKit.framework */,
+				C5D787B426169747006047E5 /* Foundation.framework */,
+				C5D787B22616973F006047E5 /* SystemExtensions.framework */,
+				C5D787B026169723006047E5 /* IOKit.framework */,
+				C5C3BBB42612ACEF003C7BFE /* DriverKit.framework */,
+				C5B7D9CC26128ADA0089B4C3 /* AudioDriverKit.framework */,
+				C5C3BBB12612ACD3003C7BFE /* AudioDriverKit.framework */,
+				C5B7D9BF26128AC50089B4C3 /* DriverKit.framework */,
+			);
+			name = Frameworks;
+			sourceTree = "<group>";
+		};
+		C5B7D9C126128AC50089B4C3 /* TinyGPUDriverExtension */ = {
+			isa = PBXGroup;
+			children = (
+				C5B7D9C226128AC50089B4C3 /* TinyGPUDriver.cpp */,
+				C5B7D9C426128AC50089B4C3 /* TinyGPUDriver.iig */,
+				C5D787AD26168D1E006047E5 /* TinyGPUDriverUserClient.cpp */,
+				C5D787AB261667FC006047E5 /* TinyGPUDriverUserClient.iig */,
+				C5B7D9C626128AC50089B4C3 /* Info.plist */,
+				C5B7D9CE26128B150089B4C3 /* TinyGPUDriver.entitlements */,
+			);
+			path = TinyGPUDriverExtension;
+			sourceTree = "<group>";
+		};
+/* End PBXGroup section */
+
+/* Begin PBXHeadersBuildPhase section */
+		C5B7D9B726128AC50089B4C3 /* Headers */ = {
+			isa = PBXHeadersBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+/* End PBXHeadersBuildPhase section */
+
+/* Begin PBXNativeTarget section */
+		54E42BC3286A1697000E1E9A /* TinyGPU Installer (macOS) */ = {
+			isa = PBXNativeTarget;
+			buildConfigurationList = 54E42BD2286A1697000E1E9A /* Build configuration list for PBXNativeTarget "TinyGPU Installer (macOS)" */;
+			buildPhases = (
+				54E42BC0286A1697000E1E9A /* Sources */,
+				54E42BC1286A1697000E1E9A /* Frameworks */,
+				54E42BC2286A1697000E1E9A /* Resources */,
+				549EB122286A1D3A009D38AB /* Embed System Extensions */,
+			);
+			buildRules = (
+			);
+			dependencies = (
+				549EB127286A1D66009D38AB /* PBXTargetDependency */,
+			);
+			name = "TinyGPU Installer (macOS)";
+			productName = "SimpleAudioDriverExtension2 (macOS)";
+			productReference = 54E42BC4286A1697000E1E9A /* TinyGPU.app */;
+			productType = "com.apple.product-type.application";
+		};
+		C5B7D9BB26128AC50089B4C3 /* TinyGPUDriver */ = {
+			isa = PBXNativeTarget;
+			buildConfigurationList = C5B7D9C926128AC50089B4C3 /* Build configuration list for PBXNativeTarget "TinyGPUDriver" */;
+			buildPhases = (
+				C5B7D9B726128AC50089B4C3 /* Headers */,
+				C5B7D9B826128AC50089B4C3 /* Sources */,
+				C5B7D9B926128AC50089B4C3 /* Frameworks */,
+				C5B7D9BA26128AC50089B4C3 /* Resources */,
+			);
+			buildRules = (
+			);
+			dependencies = (
+			);
+			name = TinyGPUDriver;
+			productName = SimpleAudioDriverExtension;
+			productReference = C5B7D9BC26128AC50089B4C3 /* org.tinygrad.tinygpu.edriver.dext */;
+			productType = "com.apple.product-type.driver-extension";
+		};
+/* End PBXNativeTarget section */
+
+/* Begin PBXProject section */
+		C5B7D9B326128AC50089B4C3 /* Project object */ = {
+			isa = PBXProject;
+			attributes = {
+				BuildIndependentTargetsInParallel = YES;
+				DefaultBuildSystemTypeForWorkspace = Latest;
+				LastSwiftUpdateCheck = 1400;
+				LastUpgradeCheck = 1600;
+				ORGANIZATIONNAME = Apple;
+				TargetAttributes = {
+					54E42BC3286A1697000E1E9A = {
+						CreatedOnToolsVersion = 14.0;
+						LastSwiftMigration = 1400;
+					};
+					C5B7D9BB26128AC50089B4C3 = {
+						CreatedOnToolsVersion = 13.0;
+					};
+				};
+			};
+			buildConfigurationList = C5B7D9B626128AC50089B4C3 /* Build configuration list for PBXProject "TinyGPUDriverExtension" */;
+			compatibilityVersion = "Xcode 9.3";
+			developmentRegion = en;
+			hasScannedForEncodings = 0;
+			knownRegions = (
+				en,
+				Base,
+			);
+			mainGroup = C5B7D9B226128AC50089B4C3;
+			productRefGroup = C5B7D9BD26128AC50089B4C3 /* Products */;
+			projectDirPath = "";
+			projectRoot = "";
+			targets = (
+				C5B7D9BB26128AC50089B4C3 /* TinyGPUDriver */,
+				54E42BC3286A1697000E1E9A /* TinyGPU Installer (macOS) */,
+			);
+		};
+/* End PBXProject section */
+
+/* Begin PBXResourcesBuildPhase section */
+		54E42BC2286A1697000E1E9A /* Resources */ = {
+			isa = PBXResourcesBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+				54E42BCC286A1697000E1E9A /* Assets.xcassets in Resources */,
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+		C5B7D9BA26128AC50089B4C3 /* Resources */ = {
+			isa = PBXResourcesBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+/* End PBXResourcesBuildPhase section */
+
+/* Begin PBXSourcesBuildPhase section */
+		54E42BC0286A1697000E1E9A /* Sources */ = {
+			isa = PBXSourcesBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+				549EB121286A1A37009D38AB /* TinyGPUViewModel.swift in Sources */,
+				54E42BCA286A1697000E1E9A /* TinyGPUView.swift in Sources */,
+				54E42BC8286A1697000E1E9A /* TinyGPUApp.swift in Sources */,
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+		C5B7D9B826128AC50089B4C3 /* Sources */ = {
+			isa = PBXSourcesBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+				C5B7D9C526128AC50089B4C3 /* TinyGPUDriver.iig in Sources */,
+				C5D787AE26168E59006047E5 /* TinyGPUDriverUserClient.cpp in Sources */,
+				C5D787AC261667FC006047E5 /* TinyGPUDriverUserClient.iig in Sources */,
+				C5B7D9C326128AC50089B4C3 /* TinyGPUDriver.cpp in Sources */,
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+/* End PBXSourcesBuildPhase section */
+
+/* Begin PBXTargetDependency section */
+		549EB127286A1D66009D38AB /* PBXTargetDependency */ = {
+			isa = PBXTargetDependency;
+			target = C5B7D9BB26128AC50089B4C3 /* TinyGPUDriver */;
+			targetProxy = 549EB126286A1D66009D38AB /* PBXContainerItemProxy */;
+		};
+/* End PBXTargetDependency section */
+
+/* Begin XCBuildConfiguration section */
+		54E42BCF286A1697000E1E9A /* Debug */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
+				ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor;
+				CLANG_ENABLE_MODULES = YES;
+				CODE_SIGN_ENTITLEMENTS = macOS/macOS.entitlements;
+				CODE_SIGN_IDENTITY = "-";
+				"CODE_SIGN_IDENTITY[sdk=macosx*]" = "Apple Development";
+				CODE_SIGN_STYLE = Automatic;
+				COMBINE_HIDPI_IMAGES = YES;
+				CURRENT_PROJECT_VERSION = 1;
+				DEAD_CODE_STRIPPING = YES;
+				DEVELOPMENT_TEAM = 9YG3G8543N;
+				ENABLE_HARDENED_RUNTIME = YES;
+				ENABLE_PREVIEWS = YES;
+				ENABLE_USER_SCRIPT_SANDBOXING = YES;
+				GENERATE_INFOPLIST_FILE = YES;
+				INFOPLIST_KEY_NSHumanReadableCopyright = "";
+				LD_RUNPATH_SEARCH_PATHS = (
+					"$(inherited)",
+					"@executable_path/../Frameworks",
+				);
+				MACOSX_DEPLOYMENT_TARGET = 12.1;
+				MARKETING_VERSION = 1.0;
+				PRODUCT_BUNDLE_IDENTIFIER = org.tinygrad.tinygpu.installer;
+				PRODUCT_NAME = TinyGPU;
+				PROVISIONING_PROFILE_SPECIFIER = "";
+				SDKROOT = macosx;
+				SWIFT_ACTIVE_COMPILATION_CONDITIONS = DEBUG;
+				SWIFT_EMIT_LOC_STRINGS = YES;
+				SWIFT_OPTIMIZATION_LEVEL = "-Onone";
+				SWIFT_VERSION = 5.0;
+			};
+			name = Debug;
+		};
+		54E42BD0286A1697000E1E9A /* Release */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
+				ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor;
+				CLANG_ENABLE_MODULES = YES;
+				CODE_SIGN_ENTITLEMENTS = macOS/macOS.entitlements;
+				CODE_SIGN_IDENTITY = "-";
+				"CODE_SIGN_IDENTITY[sdk=macosx*]" = "Apple Development";
+				CODE_SIGN_STYLE = Automatic;
+				COMBINE_HIDPI_IMAGES = YES;
+				CURRENT_PROJECT_VERSION = 1;
+				DEAD_CODE_STRIPPING = YES;
+				DEVELOPMENT_TEAM = 9YG3G8543N;
+				ENABLE_HARDENED_RUNTIME = YES;
+				ENABLE_PREVIEWS = YES;
+				ENABLE_USER_SCRIPT_SANDBOXING = YES;
+				GENERATE_INFOPLIST_FILE = YES;
+				INFOPLIST_KEY_NSHumanReadableCopyright = "";
+				LD_RUNPATH_SEARCH_PATHS = (
+					"$(inherited)",
+					"@executable_path/../Frameworks",
+				);
+				MACOSX_DEPLOYMENT_TARGET = 12.1;
+				MARKETING_VERSION = 1.0;
+				PRODUCT_BUNDLE_IDENTIFIER = org.tinygrad.tinygpu.installer;
+				PRODUCT_NAME = TinyGPU;
+				PROVISIONING_PROFILE_SPECIFIER = "";
+				SDKROOT = macosx;
+				SWIFT_COMPILATION_MODE = wholemodule;
+				SWIFT_EMIT_LOC_STRINGS = YES;
+				SWIFT_OPTIMIZATION_LEVEL = "-O";
+				SWIFT_VERSION = 5.0;
+			};
+			name = Release;
+		};
+		C5B7D9C726128AC50089B4C3 /* Debug */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				ALWAYS_SEARCH_USER_PATHS = NO;
+				CLANG_ANALYZER_NONNULL = YES;
+				CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE;
+				CLANG_CXX_LANGUAGE_STANDARD = "gnu++17";
+				CLANG_CXX_LIBRARY = "libc++";
+				CLANG_ENABLE_MODULES = YES;
+				CLANG_ENABLE_OBJC_ARC = YES;
+				CLANG_ENABLE_OBJC_WEAK = YES;
+				CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES;
+				CLANG_WARN_BOOL_CONVERSION = YES;
+				CLANG_WARN_COMMA = YES;
+				CLANG_WARN_CONSTANT_CONVERSION = YES;
+				CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES;
+				CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR;
+				CLANG_WARN_DOCUMENTATION_COMMENTS = YES;
+				CLANG_WARN_EMPTY_BODY = YES;
+				CLANG_WARN_ENUM_CONVERSION = YES;
+				CLANG_WARN_INFINITE_RECURSION = YES;
+				CLANG_WARN_INT_CONVERSION = YES;
+				CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES;
+				CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;
+				CLANG_WARN_OBJC_LITERAL_CONVERSION = YES;
+				CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;
+				CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;
+				CLANG_WARN_RANGE_LOOP_ANALYSIS = YES;
+				CLANG_WARN_STRICT_PROTOTYPES = YES;
+				CLANG_WARN_SUSPICIOUS_MOVE = YES;
+				CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;
+				CLANG_WARN_UNREACHABLE_CODE = YES;
+				CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
+				COPY_PHASE_STRIP = NO;
+				DEBUG_INFORMATION_FORMAT = dwarf;
+				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
+				ENABLE_STRICT_OBJC_MSGSEND = YES;
+				ENABLE_TESTABILITY = YES;
+				GCC_C_LANGUAGE_STANDARD = gnu11;
+				GCC_DYNAMIC_NO_PIC = NO;
+				GCC_NO_COMMON_BLOCKS = YES;
+				GCC_OPTIMIZATION_LEVEL = 0;
+				GCC_PREPROCESSOR_DEFINITIONS = (
+					"DEBUG=1",
+					"$(inherited)",
+				);
+				GCC_WARN_64_TO_32_BIT_CONVERSION = YES;
+				GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR;
+				GCC_WARN_UNDECLARED_SELECTOR = YES;
+				GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE;
+				GCC_WARN_UNUSED_FUNCTION = YES;
+				GCC_WARN_UNUSED_VARIABLE = YES;
+				MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE;
+				MTL_FAST_MATH = YES;
+				ONLY_ACTIVE_ARCH = YES;
+				SDKROOT = driverkit;
+			};
+			name = Debug;
+		};
+		C5B7D9C826128AC50089B4C3 /* Release */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				ALWAYS_SEARCH_USER_PATHS = NO;
+				CLANG_ANALYZER_NONNULL = YES;
+				CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE;
+				CLANG_CXX_LANGUAGE_STANDARD = "gnu++17";
+				CLANG_CXX_LIBRARY = "libc++";
+				CLANG_ENABLE_MODULES = YES;
+				CLANG_ENABLE_OBJC_ARC = YES;
+				CLANG_ENABLE_OBJC_WEAK = YES;
+				CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES;
+				CLANG_WARN_BOOL_CONVERSION = YES;
+				CLANG_WARN_COMMA = YES;
+				CLANG_WARN_CONSTANT_CONVERSION = YES;
+				CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES;
+				CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR;
+				CLANG_WARN_DOCUMENTATION_COMMENTS = YES;
+				CLANG_WARN_EMPTY_BODY = YES;
+				CLANG_WARN_ENUM_CONVERSION = YES;
+				CLANG_WARN_INFINITE_RECURSION = YES;
+				CLANG_WARN_INT_CONVERSION = YES;
+				CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES;
+				CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;
+				CLANG_WARN_OBJC_LITERAL_CONVERSION = YES;
+				CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;
+				CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;
+				CLANG_WARN_RANGE_LOOP_ANALYSIS = YES;
+				CLANG_WARN_STRICT_PROTOTYPES = YES;
+				CLANG_WARN_SUSPICIOUS_MOVE = YES;
+				CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;
+				CLANG_WARN_UNREACHABLE_CODE = YES;
+				CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
+				COPY_PHASE_STRIP = NO;
+				DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
+				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
+				ENABLE_NS_ASSERTIONS = NO;
+				ENABLE_STRICT_OBJC_MSGSEND = YES;
+				GCC_C_LANGUAGE_STANDARD = gnu11;
+				GCC_NO_COMMON_BLOCKS = YES;
+				GCC_WARN_64_TO_32_BIT_CONVERSION = YES;
+				GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR;
+				GCC_WARN_UNDECLARED_SELECTOR = YES;
+				GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE;
+				GCC_WARN_UNUSED_FUNCTION = YES;
+				GCC_WARN_UNUSED_VARIABLE = YES;
+				MTL_ENABLE_DEBUG_INFO = NO;
+				MTL_FAST_MATH = YES;
+				SDKROOT = driverkit;
+				SWIFT_COMPILATION_MODE = wholemodule;
+			};
+			name = Release;
+		};
+		C5B7D9CA26128AC50089B4C3 /* Debug */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				AD_HOC_CODE_SIGNING_ALLOWED = YES;
+				CODE_SIGN_ENTITLEMENTS = TinyGPUDriverExtension/TinyGPUDriver.entitlements;
+				CODE_SIGN_IDENTITY = "Apple Development";
+				CODE_SIGN_STYLE = Automatic;
+				CURRENT_PROJECT_VERSION = 1;
+				DEVELOPMENT_TEAM = 9YG3G8543N;
+				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
+				ENABLE_USER_SCRIPT_SANDBOXING = YES;
+				FRAMEWORK_SEARCH_PATHS = (
+					"$(inherited)",
+					"$(SDKROOT)/System/DriverKit/System/Library/Frameworks",
+				);
+				GENERATE_INFOPLIST_FILE = YES;
+				INFOPLIST_FILE = TinyGPUDriverExtension/Info.plist;
+				INFOPLIST_KEY_OSBundleUsageDescription = "Sample Code Audio Driver Kit Extension";
+				MARKETING_VERSION = 1.0;
+				PRODUCT_BUNDLE_IDENTIFIER = org.tinygrad.tinygpu.edriver;
+				PRODUCT_NAME = "$(inherited)";
+				PROVISIONING_PROFILE_SPECIFIER = "";
+				RUN_CLANG_STATIC_ANALYZER = YES;
+				SDKROOT = driverkit;
+				SKIP_INSTALL = YES;
+			};
+			name = Debug;
+		};
+		C5B7D9CB26128AC50089B4C3 /* Release */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				AD_HOC_CODE_SIGNING_ALLOWED = YES;
+				CODE_SIGN_ENTITLEMENTS = TinyGPUDriverExtension/TinyGPUDriver.entitlements;
+				CODE_SIGN_IDENTITY = "Apple Development";
+				CODE_SIGN_STYLE = Automatic;
+				CURRENT_PROJECT_VERSION = 1;
+				DEVELOPMENT_TEAM = 9YG3G8543N;
+				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
+				ENABLE_USER_SCRIPT_SANDBOXING = YES;
+				FRAMEWORK_SEARCH_PATHS = (
+					"$(inherited)",
+					"$(SDKROOT)/System/DriverKit/System/Library/Frameworks",
+				);
+				GENERATE_INFOPLIST_FILE = YES;
+				INFOPLIST_FILE = TinyGPUDriverExtension/Info.plist;
+				INFOPLIST_KEY_OSBundleUsageDescription = "Sample Code Audio Driver Kit Extension";
+				MARKETING_VERSION = 1.0;
+				PRODUCT_BUNDLE_IDENTIFIER = org.tinygrad.tinygpu.edriver;
+				PRODUCT_NAME = "$(inherited)";
+				PROVISIONING_PROFILE_SPECIFIER = "";
+				RUN_CLANG_STATIC_ANALYZER = YES;
+				SDKROOT = driverkit;
+				SKIP_INSTALL = YES;
+			};
+			name = Release;
+		};
+/* End XCBuildConfiguration section */
+
+/* Begin XCConfigurationList section */
+		54E42BD2286A1697000E1E9A /* Build configuration list for PBXNativeTarget "TinyGPU Installer (macOS)" */ = {
+			isa = XCConfigurationList;
+			buildConfigurations = (
+				54E42BCF286A1697000E1E9A /* Debug */,
+				54E42BD0286A1697000E1E9A /* Release */,
+			);
+			defaultConfigurationIsVisible = 0;
+			defaultConfigurationName = Release;
+		};
+		C5B7D9B626128AC50089B4C3 /* Build configuration list for PBXProject "TinyGPUDriverExtension" */ = {
+			isa = XCConfigurationList;
+			buildConfigurations = (
+				C5B7D9C726128AC50089B4C3 /* Debug */,
+				C5B7D9C826128AC50089B4C3 /* Release */,
+			);
+			defaultConfigurationIsVisible = 0;
+			defaultConfigurationName = Release;
+		};
+		C5B7D9C926128AC50089B4C3 /* Build configuration list for PBXNativeTarget "TinyGPUDriver" */ = {
+			isa = XCConfigurationList;
+			buildConfigurations = (
+				C5B7D9CA26128AC50089B4C3 /* Debug */,
+				C5B7D9CB26128AC50089B4C3 /* Release */,
+			);
+			defaultConfigurationIsVisible = 0;
+			defaultConfigurationName = Release;
+		};
+/* End XCConfigurationList section */
+	};
+	rootObject = C5B7D9B326128AC50089B4C3 /* Project object */;
+}

extra/usbgpu/tbgpu/installer/TinyGPUDriverExtension.xcodeproj/project.xcworkspace/xcshareddata/WorkspaceSettings.xcsettings

@@ -0,0 +1,10 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>BuildSystemType</key>
+	<string>Latest</string>
+	<key>DerivedDataLocationStyle</key>
+	<string>Default</string>
+</dict>
+</plist>

extra/usbgpu/tbgpu/installer/TinyGPUDriverExtension/Info.plist

@@ -0,0 +1,37 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>IOKitPersonalities</key>
+	<dict>
+		<key>TinyGPUDriver</key>
+		<dict>
+			<key>CFBundleIdentifier</key>
+			<string>$(PRODUCT_BUNDLE_IDENTIFIER)</string>
+			<key>IOClass</key>
+			<string>IOUserService</string>
+			<key>IOMatchCategory</key>
+			<string>TinyGPUDriver</string>
+			<key>IOPCIPrimaryMatch</key>
+			<string>0x70001002&amp;0xF000FFFF</string>
+			<key>IOPCITunnelCompatible</key>
+			<true/>
+			<key>IOProviderClass</key>
+			<string>IOPCIDevice</string>
+			<key>IOResourceMatch</key>
+			<string>IOKit</string>
+			<key>IOUserClass</key>
+			<string>TinyGPUDriver</string>
+			<key>IOUserServerName</key>
+			<string>org.tinygrad.tinygpu.Driver</string>
+			<key>TinyGPUDriverUserClientProperties</key>
+			<dict>
+				<key>IOClass</key>
+				<string>IOUserUserClient</string>
+				<key>IOUserClass</key>
+				<string>TinyGPUDriverUserClient</string>
+			</dict>
+		</dict>
+	</dict>
+</dict>
+</plist>

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

@@ -0,0 +1,223 @@
+#include "TinyGPUDriver.h"
+#include "TinyGPUDriverUserClient.h"
+#include <AudioDriverKit/AudioDriverKit.h>
+#include <DriverKit/IOUserServer.h>
+#include <DriverKit/IOLib.h>
+#include <DriverKit/OSString.h>
+#include <DriverKit/IOMemoryMap.h>
+#include <DriverKit/IODMACommand.h>
+#include <DriverKit/IODispatchQueue.h>
+#include <PCIDriverKit/PCIDriverKit.h>
+#include <DriverKit/OSAction.h>
+
+struct TinyGPUDriver_IVars
+{
+	IOPCIDevice *pci = nullptr;
+};
+
+bool TinyGPUDriver::init()
+{
+	os_log(OS_LOG_DEFAULT, "tinygpu: init");
+
+	auto answer = super::init();
+	if (!answer) {
+		return false;
+	}
+
+	ivars = new TinyGPUDriver_IVars();
+	if (ivars == nullptr) {
+		return false;
+	}
+
+	return true;
+}
+
+void TinyGPUDriver::free()
+{
+	if (ivars != nullptr) {
+		
+	}
+	IOSafeDeleteNULL(ivars, TinyGPUDriver_IVars, 1);
+	super::free();
+}
+
+kern_return_t TinyGPUDriver::Start_Impl(IOService* in_provider)
+{
+	IOServiceName service_name;
+	os_log(OS_LOG_DEFAULT, "tinygpu: on gpu detected");
+
+	kern_return_t err = Start(in_provider, SUPERDISPATCH);
+	if (err) return err;
+
+	ivars->pci = OSDynamicCast(IOPCIDevice, in_provider);
+	if (!ivars->pci) return kIOReturnNoDevice;
+
+	err = ivars->pci->Open(this, 0);
+	if (err) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: Open() failed 0x%08x", err);
+		ivars->pci = nullptr;
+		return err;
+	}
+
+	uint16_t ven = 0, dev = 0;
+	ivars->pci->ConfigurationRead16(kIOPCIConfigurationOffsetVendorID, &ven);
+	ivars->pci->ConfigurationRead16(kIOPCIConfigurationOffsetDeviceID, &dev);
+	os_log(OS_LOG_DEFAULT, "tinygpu: opened device ven=0x%04x dev=0x%04x", ven, dev);
+
+#if 0
+	uint32_t off = 0x100;
+	while (off) {
+		uint32_t hdr = 0, next = 0, cap_id = 0;
+		ivars->pci->ConfigurationRead32(off, &hdr);
+		cap_id = hdr & 0xFFFFu;
+		next = (hdr >> 20) & 0xFFCu;
+		os_log(OS_LOG_DEFAULT, "tinygpu: cap: %u", cap_id);
+		if (cap_id == 0x15) {
+			uint32_t cap = 0, ctrl = 0;
+			ivars->pci->ConfigurationRead32(off+0x4, &cap);
+			ivars->pci->ConfigurationRead32(off+0x8, &ctrl);
+
+			uint32_t new_bar_size = 31 - __builtin_clz(cap >> 4);
+			uint32_t new_ctrl = (ctrl & ~0x1f00) | (new_bar_size << 8);
+			ivars->pci->ConfigurationWrite32(off+0x8, new_ctrl);
+
+			os_log(OS_LOG_DEFAULT, "tinygpu: rebar: cap=%u ctrl=%u new_bar_size=%u new_ctrl=%u", cap, ctrl, new_bar_size, new_ctrl);
+			ivars->pci->Reset(0);
+			break;
+		}
+		off = next;
+	}
+	ivars->pci->Reset(0);
+#endif
+
+	uint16_t commandRegister;
+	ivars->pci->ConfigurationRead16(kIOPCIConfigurationOffsetCommand, &commandRegister);
+	commandRegister |= (kIOPCICommandIOSpace | kIOPCICommandBusMaster | kIOPCICommandMemorySpace);
+	ivars->pci->ConfigurationWrite16(kIOPCIConfigurationOffsetCommand, commandRegister);
+
+	memcpy((void*)service_name, (void*)"tinygpu\0", 8);
+	SetName(service_name);
+
+	os_log(OS_LOG_DEFAULT, "tinygpu: will register service %s", service_name);
+	RegisterService();
+
+	os_log(OS_LOG_DEFAULT, "tinygpu: service started %s", service_name);
+	return 0;
+}
+
+kern_return_t TinyGPUDriver::Stop_Impl(IOService* in_provider)
+{
+	ivars->pci->Close(this, 0);
+	return 0;
+}
+
+kern_return_t TinyGPUDriver::NewUserClient_Impl(uint32_t in_type, IOUserClient** out_user_client)
+{
+	kern_return_t err = 0;
+
+	IOService* user_client_service = nullptr;
+	err = Create(this, "TinyGPUDriverUserClientProperties", &user_client_service);
+	if (err) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: failed to create NewUserClient");
+		goto error;
+	}
+	*out_user_client = OSDynamicCast(IOUserClient, user_client_service);
+	os_log(OS_LOG_DEFAULT, "tinygpu: NewUserClient created");
+
+error:
+	return err;
+}
+
+kern_return_t TinyGPUDriver::MapBar(uint32_t bar, IOMemoryDescriptor** memory)
+{
+	kern_return_t err = 0;
+	uint8_t barMemoryIndex, barMemoryType;
+	uint64_t barMemorySize;
+	err = ivars->pci->GetBARInfo(bar, &barMemoryIndex, &barMemorySize, &barMemoryType);
+	if (err) return err;
+
+	os_log(OS_LOG_DEFAULT, "tinygpu: requested bar mapping %d, %d", bar, (uint32_t)barMemoryIndex);
+	err = ivars->pci->_CopyDeviceMemoryWithIndex(barMemoryIndex, memory, this);
+	return err;
+}
+
+kern_return_t TinyGPUDriver::CreateDMA(size_t size, TinyGPUCreateDMAResp* dmaDesc)
+{
+	kern_return_t err = 0;
+	IOMemoryMap* memoryMap = nullptr;
+	IOBufferMemoryDescriptor* sharedBuf = nullptr;
+	IODMACommand* dmaCmd = nullptr;
+	uint64_t flags = kIOMemoryDirectionInOut;
+	uint32_t segCount = 32;
+	IOAddressSegment segments[32];
+	IODMACommandSpecification dmaSpec = {
+		.options = 0,
+		.maxAddressBits = 40,
+	};
+
+	err = IOBufferMemoryDescriptor::Create(kIOMemoryDirectionInOut, size, IOVMPageSize, &sharedBuf);
+	if (err) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: failed to alloc user buffer, err=%d", err);
+		goto error;
+	}
+	
+	err = IODMACommand::Create(ivars->pci, kIODMACommandCreateNoOptions, &dmaSpec, &dmaCmd);
+	if (err) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: failed to create dma command, err=%d", err);
+		goto error;
+	}
+
+	err = dmaCmd->PrepareForDMA(kIODMACommandPrepareForDMANoOptions, sharedBuf, 0, size,
+								&flags, &segCount, segments);
+	if (err) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: failed to prepare for dma, err=%d", err);
+		goto error;
+	}
+
+	// pass addresses to userland
+	{
+		// debug
+		for (int i = 0; i < segCount; i++) {
+			os_log(OS_LOG_DEFAULT, "tinygpu: new dma mapping (sz=0x%zx) %d 0x%llx 0x%llx", size, i, segments[i].address, segments[i].length);
+		}
+
+		err = sharedBuf->CreateMapping(0, 0, 0, IOVMPageSize, IOVMPageSize, &memoryMap); // one page should be fine
+		if (err) {
+			os_log(OS_LOG_DEFAULT, "tinygpu: failed to map memory, err=%d", err);
+			goto error;
+		}
+
+		// Send back gpu addresses
+		uint64_t* addr = (uint64_t*)memoryMap->GetAddress();
+		for (int i = 0; i < segCount; i++) {
+			addr[i * 2] = segments[i].address;
+			addr[i * 2 + 1] = segments[i].length;
+		}
+		addr[segCount * 2] = 0;
+		addr[segCount * 2 + 1] = 0;
+
+		// free memoryMap
+		memoryMap->release();
+		memoryMap = nullptr;
+	}
+
+	dmaDesc->sharedBuf = sharedBuf;
+	dmaDesc->dmaCmd = dmaCmd;
+	return 0;
+
+error:
+	if (memoryMap) {
+		memoryMap->release();
+		memoryMap = nullptr;
+	}
+	if (dmaCmd) {
+		dmaCmd->CompleteDMA(kIODMACommandCompleteDMANoOptions);
+		dmaCmd->release();
+		dmaCmd = nullptr;
+	}
+	if (sharedBuf) {
+		sharedBuf->release();
+		sharedBuf = nullptr;
+	}
+	return err;
+}

extra/usbgpu/tbgpu/installer/TinyGPUDriverExtension/TinyGPUDriver.entitlements

@@ -0,0 +1,17 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>com.apple.developer.driverkit.transport.pci</key>
+	<array>
+		<dict>
+			<key>IOPCIMatch</key>
+			<string>0x70001002&amp;0xF000FFFF</string>
+		</dict>
+	</array>
+	<key>com.apple.developer.driverkit.allow-any-userclient-access</key>
+	<true/>
+	<key>com.apple.developer.driverkit</key>
+	<true/>
+</dict>
+</plist>

extra/usbgpu/tbgpu/installer/TinyGPUDriverExtension/TinyGPUDriver.iig

@@ -0,0 +1,33 @@
+#ifndef TinyGPUDriver_h
+#define TinyGPUDriver_h
+
+#include <Availability.h>
+#include <DriverKit/IOService.iig>
+#include <PCIDriverKit/IOPCIDevice.iig>
+#include <DriverKit/IOMemoryMap.iig>
+#include <DriverKit/IODMACommand.iig>
+
+struct TinyGPUCreateDMAResp
+{
+	IOBufferMemoryDescriptor* sharedBuf;
+	IODMACommand* dmaCmd;
+};
+
+class TinyGPUDriver: public IOService
+{
+public:
+	virtual bool init() override;
+
+	virtual void free() override;
+
+	virtual kern_return_t Start(IOService * provider) override;
+
+	virtual kern_return_t Stop(IOService * provider) override;
+
+	virtual kern_return_t NewUserClient(uint32_t in_type, IOUserClient** out_user_client) override;
+
+	kern_return_t MapBar(uint32_t bar, IOMemoryDescriptor** memory) LOCALONLY;
+	kern_return_t CreateDMA(size_t size, TinyGPUCreateDMAResp* dmaDesc) LOCALONLY;
+};
+
+#endif /* TinyGPUDriver_h */

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

@@ -0,0 +1,94 @@
+#include "TinyGPUDriverUserClient.h"
+#include "TinyGPUDriver.h"
+#include <DriverKit/DriverKit.h>
+#include <DriverKit/OSSharedPtr.h>
+#include <PCIDriverKit/PCIDriverKit.h>
+
+struct TinyGPUDriverUserClient_IVars
+{
+	OSSharedPtr<TinyGPUDriver> provider = nullptr;
+};
+
+bool TinyGPUDriverUserClient::init()
+{
+	auto theAnswer = super::init();
+	if (!theAnswer) {
+		return false;
+	}
+
+	ivars = IONewZero(TinyGPUDriverUserClient_IVars, 1);
+	if (ivars == nullptr) {
+		return false;
+	}
+
+	return true;
+}
+
+void TinyGPUDriverUserClient::free()
+{
+	if (ivars != nullptr) {
+		ivars->provider.reset();
+	}
+
+	IOSafeDeleteNULL(ivars, TinyGPUDriverUserClient_IVars, 1);
+	super::free();
+}
+
+kern_return_t TinyGPUDriverUserClient::Start_Impl(IOService* in_provider)
+{
+	kern_return_t err = kIOReturnSuccess;
+	if (!in_provider) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: provider is null");
+		err = kIOReturnBadArgument;
+		goto error;
+	}
+
+	err = Start(in_provider, SUPERDISPATCH);
+	if (err) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: failed to start super (%d)", err);
+		goto error;
+	}
+
+	ivars->provider = OSSharedPtr(OSDynamicCast(TinyGPUDriver, in_provider), OSRetain);
+	return 0;
+
+error:
+	ivars->provider.reset();
+	return err;
+}
+
+kern_return_t TinyGPUDriverUserClient::Stop_Impl(IOService* in_provider)
+{
+	return Stop(in_provider, SUPERDISPATCH);
+}
+
+kern_return_t TinyGPUDriverUserClient::ExternalMethod(uint64_t in_selector, IOUserClientMethodArguments* in_arguments, const IOUserClientMethodDispatch* in_dispatch, OSObject* in_target, void* in_reference)
+{
+	return kIOReturnUnsupported;
+}
+
+kern_return_t IMPL(TinyGPUDriverUserClient, CopyClientMemoryForType)
+{
+	if (!memory) {
+		return kIOReturnBadArgument;
+	}
+
+	if (ivars->provider.get() == nullptr) {
+		return kIOReturnNotAttached;
+	}
+
+	if (type < 6) {
+		uint32_t bar = (uint32_t)type;
+		return ivars->provider->MapBar(bar, memory);
+	}
+
+	// dma page buffer
+	TinyGPUCreateDMAResp buf;
+	kern_return_t err = ivars->provider->CreateDMA(type, &buf);
+	if (err) {
+		return err;
+	}
+
+	*memory = buf.sharedBuf;
+	return 0;
+}

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

@@ -0,0 +1,21 @@
+#ifndef TinyGPUDriverUserClient_h
+#define TinyGPUDriverUserClient_h
+
+#include <DriverKit/IOUserClient.iig>
+
+class TinyGPUDriverUserClient : public IOUserClient
+{
+public:
+	virtual bool init() final;
+	virtual void free() final;
+
+	virtual kern_return_t Start(IOService* in_provider) final;
+	virtual kern_return_t Stop(IOService* in_provider) final;
+
+	virtual kern_return_t ExternalMethod(uint64_t in_selector, IOUserClientMethodArguments* in_arguments, const IOUserClientMethodDispatch* in_dispatch, OSObject* in_target, void* in_reference) final;
+
+	virtual kern_return_t CopyClientMemoryForType(
+		uint64_t type, uint64_t *options, IOMemoryDescriptor **memory) final;
+};
+
+#endif /* TinyGPUDriverUserClient_h */

extra/usbgpu/tbgpu/installer/macOS/macOS.entitlements

@@ -0,0 +1,12 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>com.apple.security.app-sandbox</key>
+	<false/>
+	<key>com.apple.security.files.user-selected.read-only</key>
+	<true/>
+	<key>com.apple.developer.system-extension.install</key>
+	<true/>
+</dict>
+</plist>

extra/usbgpu/tbgpu/main.cpp

@@ -0,0 +1,39 @@
+#include <CoreFoundation/CoreFoundation.h>
+#include <IOKit/IOKitLib.h>
+#include <mach/mach.h>
+#include <stdio.h>
+#include <inttypes.h>
+
+static io_connect_t open_uc_by_name(const char *svc_name) {
+  io_connect_t conn = IO_OBJECT_NULL;
+  io_service_t service = IOServiceGetMatchingService(kIOMasterPortDefault, IOServiceNameMatching(svc_name));
+  if (!service) { fprintf(stderr, "service not found: %s\n", svc_name); return IO_OBJECT_NULL; }
+  kern_return_t kr = IOServiceOpen(service, mach_task_self(), /*type*/0, &conn);
+  IOObjectRelease(service);
+  if (kr) { fprintf(stderr, "IOServiceOpen 0x%x\n", kr); return IO_OBJECT_NULL; }
+  return conn;
+}
+
+int main(int argc, char **argv) {
+  uint32_t bar = (argc > 1) ? (uint32_t)strtoul(argv[1], NULL, 0) : 0;  // pick BAR index
+  io_connect_t conn = open_uc_by_name("tinygpu");
+  if (!conn) return 2;
+
+  mach_vm_address_t addr = 0;
+  mach_vm_size_t    size = 0;
+  kern_return_t kr = IOConnectMapMemory64(conn, bar, mach_task_self(), &addr, &size, kIOMapAnywhere);
+  if (kr) { fprintf(stderr, "Map BAR%u failed 0x%x\n", bar, kr); IOServiceClose(conn); return 3; }
+
+  printf("BAR%u mapped at 0x%llx, size 0x%llx\n", bar, (unsigned long long)addr, (unsigned long long)size);
+
+  // example: read a 32-bit register at offset 0x0 (make sure it’s safe!)
+  volatile uint32_t *mmio = (volatile uint32_t*)(uintptr_t)addr;
+  uint32_t v = mmio[0];
+  printf("mmio[0]=0x%08x\n", v);
+
+  kr = IOConnectUnmapMemory64(conn, bar, mach_task_self(), addr);
+  if (kr) fprintf(stderr, "Unmap failed 0x%x\n", kr);
+
+  IOServiceClose(conn);
+  return 0;
+}

extra/usbgpu/tbgpu/main.py

@@ -0,0 +1,82 @@
+import ctypes, ctypes.util, sys
+
+cf = ctypes.CDLL(ctypes.util.find_library("CoreFoundation"))
+iokit = ctypes.CDLL(ctypes.util.find_library("IOKit"))
+libsys = ctypes.CDLL(ctypes.util.find_library("System"))
+
+kern_return_t = ctypes.c_int
+mach_port_t = ctypes.c_uint
+io_object_t = mach_port_t
+io_service_t = io_object_t
+io_connect_t = mach_port_t
+CFMutableDictionaryRef = ctypes.c_void_p
+CFStringRef = ctypes.c_void_p
+
+kIOMasterPortDefault = mach_port_t(0)
+
+libsys.mach_task_self_.restype = mach_port_t
+
+iokit.IOServiceNameMatching.argtypes = [ctypes.c_char_p]
+iokit.IOServiceNameMatching.restype  = CFMutableDictionaryRef
+
+iokit.IOServiceGetMatchingService.argtypes = [mach_port_t, CFMutableDictionaryRef]
+iokit.IOServiceGetMatchingService.restype  = io_service_t
+
+iokit.IOObjectRelease.argtypes = [io_object_t]
+iokit.IOObjectRelease.restype  = kern_return_t
+
+iokit.IOServiceOpen.argtypes = [io_service_t, mach_port_t, ctypes.c_uint32, ctypes.POINTER(io_connect_t)]
+iokit.IOServiceOpen.restype  = kern_return_t
+
+iokit.IOConnectCallMethod.argtypes = [io_connect_t, ctypes.c_uint32, ctypes.POINTER(ctypes.c_uint64), ctypes.c_uint32, ctypes.c_void_p,
+  ctypes.c_size_t, ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(ctypes.c_uint32), ctypes.c_void_p, ctypes.POINTER(ctypes.c_size_t)]
+iokit.IOConnectCallMethod.restype = kern_return_t
+
+def open_userclient_by_name(name: str, uc_type: int = 0) -> io_connect_t:
+  mdict = iokit.IOServiceNameMatching(name.encode("utf-8"))
+  if not mdict: raise RuntimeError("IOServiceNameMatching returned NULL")
+
+  # Grab the first matching service
+  service = iokit.IOServiceGetMatchingService(kIOMasterPortDefault, mdict)
+  if not service: raise RuntimeError(f'service "{name}" not found')
+
+  # print("lol", service)
+  # print(libsys.mach_task_self_)
+  # cast libsys.mach_task_self_ to uint and print
+  # print("lol", ctypes.cast(libsys.mach_task_self_, ctypes.POINTER(ctypes.c_uint)).contents.value)
+
+  try:
+    # Open user client (type -> passed to NewUserClient_Impl)
+    conn = io_connect_t(0)
+    # print("lol", libsys.mach_task_self_)
+    kr = iokit.IOServiceOpen(service, ctypes.cast(libsys.mach_task_self_, ctypes.POINTER(ctypes.c_uint)).contents.value,
+                             ctypes.c_uint32(uc_type), ctypes.byref(conn))
+    if kr != 0: raise OSError(kr, f"IOServiceOpen failed (0x{kr:08x})")
+    return conn
+  finally: iokit.IOObjectRelease(service)
+
+def external_method(conn: io_connect_t, selector: int = 0) -> int:
+  # no scalars in/out, no struct in/out — just ping selector 0
+  in_scalars = ctypes.POINTER(ctypes.c_uint64)()  # NULL
+  out_scalars = (ctypes.c_uint64 * 1)()       # space if driver returns something
+  out_scalars_cnt = ctypes.c_uint32(0)      # driver can set this
+
+  return iokit.IOConnectCallMethod(conn, ctypes.c_uint32(selector), in_scalars, ctypes.c_uint32(0), None, ctypes.c_size_t(0),
+    out_scalars, ctypes.byref(out_scalars_cnt), None, ctypes.byref(ctypes.c_size_t(0)))
+
+def close_userclient(conn: io_connect_t) -> None:
+  # IOServiceClose is a macro; exported symbol is IOServiceClose in IOKit
+  iokit.IOServiceClose.argtypes = [io_connect_t]
+  iokit.IOServiceClose.restype  = kern_return_t
+  iokit.IOServiceClose(conn)
+
+if __name__ == "__main__":
+  try:
+    conn = open_userclient_by_name("tinygpu", uc_type=0)
+    kr = external_method(conn, selector=0)
+    print(f"ExternalMethod(0) -> 0x{kr:08x}")
+  except Exception as e:
+    print(e)
+    sys.exit(1)
+  finally:
+    if 'conn' in locals() and conn.value: close_userclient(conn)

setup.py

@@ -42,7 +42,6 @@ setup(name='tinygrad',
         'tinygrad.runtime.support.am',
         'tinygrad.runtime.support.nv',
         'tinygrad.schedule',
-        'tinygrad.shape',
         'tinygrad.uop',
         'tinygrad.viz',
       ],

test/external/external_uop_gc.py

@@ -1,6 +1,5 @@
 import gc
 from tinygrad import Tensor, UOp, Device, nn
-from tinygrad.shape.shapetracker import views_to_valid_uop
 from tinygrad.engine.realize import method_cache, get_program
 from tinygrad.schedule.indexing import apply_movement_op
 from test.test_tiny import TestTiny
@@ -69,7 +68,6 @@ if __name__ == "__main__":
 
     # these caches will keep uops alive
     method_cache.clear()
-    views_to_valid_uop.cache_clear()
     apply_movement_op.cache_clear()
     Tensor._device_seeds.clear()
     Tensor._device_rng_counters.clear()

test/opt/test_gen_float4.py

@@ -2,7 +2,6 @@ import unittest
 from tinygrad import Device, Tensor, dtypes
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad.engine.realize import get_program
 from tinygrad.helpers import AMX
 
@@ -149,33 +148,5 @@ class TestFloat4(unittest.TestCase):
 
     assert TestFloat4.count_float4(uops) == (1, 1)
 
-  @unittest.skip("Ops.VIEW no longer exists")
-  def test_half4_load_unrolled(self):
-    # from llama 7B shard 4 gpus
-    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.VIEW, dtypes.float.ptr(96000), arg=ShapeTracker(views=(View(shape=(1, 3, 32000, 1), strides=(0, 32000, 1, 0), offset=0, mask=None, contiguous=True),)), src=( # noqa: E501
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(96000), arg=0, src=()),)),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (3,)), src=(
-          UOp(Ops.CAST, dtypes.float, arg=None, src=(
-            UOp(Ops.MUL, dtypes.half, arg=None, src=(
-              UOp(Ops.LOAD, dtypes.half, arg=None, src=(
-                UOp(Ops.VIEW, dtypes.half.ptr(9216), arg=ShapeTracker(views=(View(shape=(1, 3, 32000, 1024), strides=(0, 4096, 0, 1), offset=0, mask=None, contiguous=False),)), src=( # noqa: E501
-                  UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(9216), arg=1, src=()),)),)),
-              UOp(Ops.LOAD, dtypes.half, arg=None, src=(
-                UOp(Ops.VIEW, dtypes.half.ptr(32768000), arg=ShapeTracker(views=(View(shape=(1, 3, 32000, 1024), strides=(0, 0, 1024, 1), offset=0, mask=None, contiguous=False),)), src=( # noqa: E501
-                  UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(32768000), arg=2, src=()),)),)),)),)),)),)),))
-
-    # TODO: fix this, expected might change but should be positive
-    for expected, opts in [
-      ((7, 0), [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=3), Opt(op=OptOps.UNROLL, axis=0, arg=4)]),
-      ((5, 0), [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UNROLL, axis=0, arg=4)]),
-      ((2, 0), [Opt(op=OptOps.UNROLL, axis=0, arg=4)]),
-    ]:
-      program = get_program(ast, Device[Device.DEFAULT].renderer, opts=opts)
-
-      count = TestFloat4.count_half4(program.uops)
-      assert count == expected, f"{count=}, {expected=}"
-
 if __name__ == '__main__':
   unittest.main()

test/opt/test_kernel_opts.py

@@ -1,6 +1,5 @@
 import unittest
 from tinygrad import Device, Tensor, dtypes
-from tinygrad.helpers import CI
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 
 # TODO: write a clean version of this
@@ -177,9 +176,7 @@ class TestKernelOpts(unittest.TestCase):
     ], apply_tc=True, atol=atol, rtol=rtol)
 
   def test_padto_matmul(self):
-    if (CI and Device.DEFAULT in ["AMD", "NV", "CUDA"]):
-      self.skipTest("super slow on CUDA and AMD because of the big grid dims")
-    N = 17 * 17
+    N = 17
     Tensor.manual_seed(289)
     a = Tensor.rand(N, N)
     b = Tensor.rand(N, N)
@@ -213,7 +210,7 @@ class TestKernelOpts(unittest.TestCase):
       helper_linearizer_opt(a@b, [[Opt(OptOps.UNROLL, 0, 0), Opt(OptOps.PADTO, 2, 8)]])
 
   def test_padto_sum_ok(self):
-    N = 18 * 18
+    N = 18
     # NOTE: this setup prevents 17 * 17 contiguous merged into one dimension
     a = Tensor.rand(N, N).realize().shrink(((0, 17), (0, 17))) * 100
     b = (Tensor.rand(N, N) < 0.5).realize().shrink(((0, 17), (0, 17)))
@@ -244,7 +241,7 @@ class TestKernelOpts(unittest.TestCase):
     helper_linearizer_opt(a.sum(0).exp(), [[Opt(OptOps.PADTO, 1, 32)],])
 
   def test_padto_sum_not_ok(self):
-    N = 18 * 18
+    N = 18
     # NOTE: this setup prevents 17 * 17 contiguous merged into one dimension
     a = Tensor.rand(N, N).shrink(((0, 17), (0, 17))).exp()
     # exp is not safe to pad
@@ -261,7 +258,7 @@ class TestKernelOpts(unittest.TestCase):
       helper_linearizer_opt(b.sum(0), [[Opt(OptOps.PADTO, 1, 32)],])
 
   def test_padto_max(self):
-    N = 18 * 18
+    N = 18
     # NOTE: this setup prevents 17 * 17 contiguous merged into one axis
     a = -Tensor.rand(N, N).shrink(((0, 17), (0, 17))) * 100
 
@@ -282,7 +279,7 @@ class TestKernelOpts(unittest.TestCase):
 
   def test_padto_where(self):
     Tensor.manual_seed(0)
-    N = 17 * 17
+    N = 17
     a = (Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1).where(1, 0)
     helper_linearizer_opt(a.max(0), [
       [Opt(OptOps.PADTO, 0, 32)],
@@ -291,7 +288,7 @@ class TestKernelOpts(unittest.TestCase):
 
   def test_padto_where_multioutput(self):
     Tensor.manual_seed(0)
-    N = 17 * 17
+    N = 17
     r = Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1
     a0 = r.where(1, 0)
     a1 = r.where(2, 0)

test/test_arange.py

@@ -131,8 +131,7 @@ class TestIndexing(unittest.TestCase):
     # llama3 is 128256
     vocab_size, embed_size = (10, 3) if CI else (32000, 4096)
     emb = nn.Embedding(vocab_size, embed_size)
-    # TODO: why is a new realize needed here
-    emb_w = emb.weight.realize().numpy()
+    emb_w = emb.weight.numpy()
     x = Tensor([1,2,3,4])
     with Context(NOOPT=noopt):
       GlobalCounters.reset()

test/test_compile_failures.py

@@ -1,7 +1,7 @@
 import unittest, io
 from contextlib import redirect_stdout
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.helpers import OSX, CPU_LLVM
+from tinygrad.helpers import OSX, CPU_LLVM, CPU_LVP
 from tinygrad.engine.realize import lower_schedule
 from tinygrad.device import is_dtype_supported
 from tinygrad.engine.realize import get_program
@@ -19,7 +19,7 @@ class TestCompileFailures(unittest.TestCase):
 
 class TestDisassembly(unittest.TestCase):
   # TODO: fails on llvm. llvm.LLVMGetHostCPUName() returns "generic"
-  @unittest.skipUnless(Device.DEFAULT in ("CPU",) and not CPU_LLVM and OSX, "m series cpus support fp16 arithmetic")
+  @unittest.skipUnless(Device.DEFAULT in ("CPU",) and not (CPU_LLVM or CPU_LVP) and OSX, "m series cpus support fp16 arithmetic")
   def test_float16_alu(self):
     c = Tensor([1], dtype=dtypes.float16) + Tensor([1], dtype=dtypes.float16)
     s = c.schedule()[-1]

test/test_dtype.py

@@ -6,6 +6,7 @@ from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import getenv, DEBUG, CI
 from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8, _to_np_dtype, _to_torch_dtype, truncate
 from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.nir import NIRRenderer
 from tinygrad import Device, Tensor, dtypes
 from hypothesis import given, settings, strategies as strat
 from test.helpers import rand_for_dtype
@@ -102,7 +103,7 @@ class TestDType(unittest.TestCase):
     ))
 
   @unittest.skipIf(Device.DEFAULT == "PYTHON", "skip for now")
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "skip for now")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "skip for now")
   def test_uint_overflow(self):
     if not dtypes.is_unsigned(self.DTYPE): raise unittest.SkipTest("only for unsigned")
     v = dtypes.max(self.DTYPE)
@@ -261,7 +262,7 @@ class TestFloatDType(TestDType):
 class TestDoubleDType(TestDType):
   DTYPE = dtypes.double
   @unittest.skipIf((CI and Device.DEFAULT in {"CUDA", "NV"}) or \
-                    isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "conversion not supported on CI CUDA and PTX")  # TODO: why not?
+   isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "conversion not supported on CI CUDA, PTX, and NIR")  # TODO: why not?
   def test_float64_increased_precision(self):
     for func in [
       lambda t: t.exp(),

test/test_dtype_alu.py

@@ -6,6 +6,7 @@ from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
 from tinygrad.runtime.ops_python import from_storage_scalar
 from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.nir import NIRRenderer
 import numpy as np
 import pytest
 from hypothesis import assume, given, strategies as strat, settings, HealthCheck
@@ -29,8 +30,8 @@ unary_operations = [(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.
 # TODO: enable this (this is a dtype issue)
 #binary_operations.append(operator.truediv)
 
-# TODO: CI CUDA segfaults on sin, WEBGPU sin is not precise enough for large numbers
-if (getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}) or Device.DEFAULT == "WEBGPU":
+# TODO: CI CUDA segfaults on sin, WEBGPU and NIR sines are not precise enough for large numbers
+if (getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}) or Device.DEFAULT == "WEBGPU" or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer):
   unary_operations.remove((Tensor.sin, np.sin))
   unary_operations.remove((Tensor.cos, np.cos))
 
@@ -184,8 +185,8 @@ class TestDTypeALU(unittest.TestCase):
   @given(ht.int32, ht.int32, ht.float32, strat.sampled_from(integer_binary_operations), strat.sampled_from(binary_operations))
   def test_int32_midcast_float(self, a, b, c, op1, op2): universal_test_midcast(a, b, c, op1, op2, dtypes.int32, dtypes.float32)
 
-  # Metal and CUDA and HIP behave differently than numpy in CI for overflows
-  skip_overflow = CI and Device.DEFAULT in {"AMD", "NV", "CUDA"}
+  # Metal and CUDA and HIP and NIR behave differently than numpy in CI for overflows
+  skip_overflow = (CI and Device.DEFAULT in {"AMD", "NV", "CUDA"}) or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer)
   @given(strat.floats(width=32, min_value=0, max_value=10.0) if skip_overflow else ht.float32,
          strat.floats(width=32, min_value=0, max_value=10.0) if skip_overflow else ht.float32,
          ht.int32, strat.sampled_from(binary_operations), strat.sampled_from(integer_binary_operations))

test/test_edgecases.py

@@ -26,8 +26,9 @@ import unittest
 import numpy as np
 import torch
 from tinygrad import Tensor, dtypes, nn
-from tinygrad.device import is_dtype_supported
+from tinygrad.device import Device, is_dtype_supported
 from tinygrad.helpers import getenv
+from tinygrad.renderer.nir import NIRRenderer
 
 MOCKGPU = getenv("MOCKGPU")
 
@@ -206,7 +207,8 @@ class TestUOpValidationIssue(unittest.TestCase):
   # these fail with UOp verification error.
   # we want more of these with diverse errors!
 
-  @unittest.skipIf((not is_dtype_supported(dtypes.long)) or MOCKGPU, "hangs gpuocelot")
+  @unittest.skipIf((not is_dtype_supported(dtypes.long)) or MOCKGPU or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer),
+                   "hangs gpuocelot, NIR cannot render")
   def test_tensor_index_overflow(self):
     val = Tensor([1])
     big = val.expand(2**31 + 3)

test/test_linearizer_dumb.py

@@ -4,12 +4,9 @@
 
 import unittest
 from tinygrad import Device, dtypes
-from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
-from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad.codegen.opt.search import Opt, OptOps
 from tinygrad.engine.realize import get_program
-from tinygrad.renderer.ptx import PTXRenderer
 
 class TestLinearizerFailure(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
@@ -30,91 +27,5 @@ class TestLinearizerFailure(unittest.TestCase):
     ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None))
     _ = get_program(ast, Device["METAL"].renderer)
 
-class TestLinearizerDumb(unittest.TestCase):
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "need local")
-  @unittest.skip("Ops.VALID no longer exists")
-  def test_max_simplify_and_cancel(self):
-    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1000), arg=0, src=())
-    c1 = c0.view(ShapeTracker(views=(View(shape=(1000, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),)))
-    c2 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1000), arg=1, src=())
-    c3 = c2.view(ShapeTracker(views=(View(shape=(1000, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),)))
-    c4 = c3.load()
-    c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1), arg=2, src=())
-    c6 = c5.view(ShapeTracker(views=(View(shape=(1000, 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)))
-    c7 = c6.load()
-    c8 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1000, 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=())
-    c9 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1001, 1999), strides=(0, 0), offset=0, mask=((0, 1001), (999, 1999)), contiguous=False), View(shape=(1000, 1000), strides=(1, 2000), offset=0, mask=None, contiguous=False))), src=())
-    c10 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1000, 1000), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=())
-    c11 = c1.store((c4.alu(Ops.CMPNE, c7).alu(Ops.CMPNE, UOp.const(dtypes.bool, True, src=c8)).cast(dtypes.int)*(c9.f(Ops.VALID, dtype=dtypes.bool).where(UOp.const(dtypes.int, -1, src=c10), UOp.const(dtypes.int, 0, src=c10)).f(Ops.REDUCE_AXIS, arg=(Ops.ADD, (1,)))+UOp.const(dtypes.int, 1000, src=c8))))
-    ast = c11.sink()
-    #opts = [Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=8)]
-    opts = [Opt(op=OptOps.LOCAL, axis=0, arg=8)]
-    prg = get_program(ast, Device[Device.DEFAULT].renderer, opts)
-    print(prg.src)
-    assert prg.uops is not None and not any(uop.op is Ops.MAX for uop in prg.uops), "leftover MAX"
-
-  # this was a bug in embedding, someday we should fold this anyway
-  @unittest.skipUnless(is_dtype_supported(dtypes.half), f"half dtype not supported on {Device.DEFAULT}")
-  @unittest.skip("UOp.view is no longer supported")
-  def test_llama_embedding(self):
-    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(4096), arg=0, src=())
-    c1 = c0.view(ShapeTracker(views=(View(shape=(4096, 1, 1), strides=(1, 0, 0), offset=0, mask=None, contiguous=True),)))
-    c2 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(32001, 63999), strides=(0, 0), offset=0, mask=((0, 32001), (31999, 63999)), contiguous=False), View(shape=(4096, 32000, 32000), strides=(0, 1, 64000), offset=0, mask=None, contiguous=False))), src=())
-    c3 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4096, 32000, 32000), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)), src=())
-    c4 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4096, 32000, 1), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)), src=())
-    c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=1, src=())
-    c6 = c5.view(ShapeTracker(views=(View(shape=(4096, 32000, 1), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)))
-    c7 = c6.load()
-    c8 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(131072000), arg=2, src=())
-    c9 = c8.view(ShapeTracker(views=(View(shape=(4096, 32000, 1), strides=(1, 4096, 0), offset=0, mask=None, contiguous=False),)))
-    c10 = c9.load()
-    c11 = c1.store(((c2.f(Ops.VALID, dtype=dtypes.bool).where(UOp.const(dtypes.int, 1, src=c3), UOp.const(dtypes.int, 0, src=c3)).f(Ops.REDUCE_AXIS, arg=(Ops.ADD, (2,)))+UOp.const(dtypes.int, -1, src=c4)).alu(Ops.CMPNE, c7).alu(Ops.CMPNE, UOp.const(dtypes.bool, True, src=c4)).cast(dtypes.half)*c10).cast(dtypes.float).f(Ops.REDUCE_AXIS, arg=(Ops.ADD, (1,))).cast(dtypes.half))
-    ast = c11.sink()
-    prg = get_program(ast, Device[Device.DEFAULT].renderer)
-    print(prg.src)
-
-  @unittest.expectedFailure
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "need float4")
-  def test_unrolled_float4_align(self):
-    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1), arg=0, src=())
-    c1 = c0.view(ShapeTracker(views=(View(shape=(1, 1), strides=(0, 0), offset=0, mask=None, contiguous=True),)))
-    c2 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(18), arg=1, src=())
-    c3 = c2.view(ShapeTracker(views=(View(shape=(3, 6), strides=(6, 1), offset=0, mask=None, contiguous=True),)))
-    c4 = c3.load()
-    c5 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(3, 6), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=())
-    c6 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(18), arg=2, src=())
-    c7 = c6.view(ShapeTracker(views=(View(shape=(3, 6), strides=(6, 1), offset=0, mask=None, contiguous=True),)))
-    c8 = c7.load()
-    c9 = c1.store(c4.alu(Ops.CMPNE, UOp.const(dtypes.long, -1, src=c5)).alu(Ops.CMPNE, UOp.const(dtypes.bool, True, src=c5)).where(UOp.const(dtypes.float, 0.0, src=c5), c8).f(Ops.REDUCE_AXIS, arg=(Ops.ADD, (0, 1))))
-    ast = c9.sink()
-    opts = [Opt(op=OptOps.UNROLL, axis=0, arg=0)]
-    prg = get_program(ast, Device[Device.DEFAULT].renderer, opts)
-    print(prg.src)
-    load_idxs = [x.src[1] for x in prg.uops if x.op is Ops.LOAD and x.src[0].arg == 2]
-    assert load_idxs[0] < load_idxs[1], f"first loaded idx {load_idxs[0].arg} then {load_idxs[1].arg}!"
-
-  @unittest.expectedFailure
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "need float4")
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "this is somehow correct in PTX")
-  def test_upcasted_stores_out_of_order(self):
-    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(9360), arg=0, src=())
-    c1 = c0.view(ShapeTracker(views=(View(shape=(4, 5, 13, 1, 1, 1, 1, 1, 4, 3, 3), strides=(2340, 468, 36, 0, 0, 0, 0, 0, 9, 3, 1), offset=0, mask=None, contiguous=True),)))
-    c2 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(144), arg=1, src=())
-    c3 = c2.view(ShapeTracker(views=(View(shape=(4, 5, 13, 1, 1, 1, 4, 1, 4, 3, 3), strides=(0, 0, 0, 0, 0, 0, 1, 0, 4, 48, 16), offset=0, mask=None, contiguous=False),)))
-    c4 = c3.load()
-    c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1040), arg=2, src=())
-    c6 = c5.view(ShapeTracker(views=(View(shape=(4, 5, 13, 1, 1, 1, 4, 1, 4, 3, 3), strides=(260, 13, 1, 0, 0, 0, 65, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)))
-    c7 = c6.load()
-    c8 = c1.store((c4*c7).f(Ops.REDUCE_AXIS, arg=(Ops.ADD, (6,))))
-    ast = c8.sink()
-    opts = [Opt(op=OptOps.UPCAST, axis=3, arg=0), Opt(op=OptOps.UPCAST, axis=2, arg=0)]
-    prg = get_program(ast, Device[Device.DEFAULT].renderer, opts)
-    print(prg.src)
-    store_idxs = [x.src[1] for x in prg.uops if x.op is Ops.STORE]
-    for i in range(len(store_idxs) - 1):
-      first_bounds = store_idxs[i].vmin+store_idxs[i].vmax
-      next_bounds = store_idxs[i+1].vmin+store_idxs[i+1].vmax
-      assert first_bounds < next_bounds, f"first stored (max) idx {first_bounds} then {next_bounds}!"
-
 if __name__ == '__main__':
   unittest.main()

test/test_multitensor.py

@@ -658,7 +658,7 @@ class TestMultiTensor(unittest.TestCase):
 
   # it doesn't work like this anymore
   # NOTE: this never failed in assign_multi, it failed tensor spec because MULTI was never pushed in the graph
-  @unittest.expectedFailure
+  @unittest.skip("this test is broken")
   def test_mlb_assign_change_axis(self):
     t_none = Tensor.zeros((16, 16)).shard(devices_2).contiguous().realize()
     t_zero = Tensor.ones((16, 16)).shard(devices_2, axis=0)

test/test_ops.py

@@ -2,7 +2,7 @@ import time, math, unittest, functools, platform, warnings
 import numpy as np
 from typing import List, Callable
 import torch
-from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, CPU_LLVM, AMD_LLVM
+from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, CPU_LLVM, CPU_LVP, AMD_LLVM
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
@@ -698,8 +698,8 @@ class TestOps(unittest.TestCase):
 
   def test_pow_zero_tensor(self):
     helper_test_op(None, lambda x,y: x**y, vals=[[0.0], [0.0]])
-    # TODO: fix WEBGPU
-    if Device.DEFAULT != "WEBGPU":
+    # TODO: fix WEBGPU and LVP
+    if Device.DEFAULT != "WEBGPU" and not CPU_LVP:
       helper_test_op(None, lambda x,y: x**y, vals=[[0.0], [0.3]])
       helper_test_op(None, lambda x,y: x**y, vals=[[0.0], [-0.3]])
   def test_pow_zero_const(self):
@@ -830,6 +830,7 @@ class TestOps(unittest.TestCase):
     self.assertEqual(a, b)
     self.assertEqual(Tensor(-1).contiguous().idiv(4).item(), 0)  # NOTE this is trunc-div behaviour
 
+  @unittest.skipIf(getenv("NV_NAK"), "MUFU.SIN is not accurate enough")
   def test_sin(self):
     helper_test_op([(45,65)], lambda x: x.sin())
     helper_test_op([()], lambda x: x.sin())
@@ -839,6 +840,7 @@ class TestOps(unittest.TestCase):
       helper_test_op(None, lambda x: x.sin(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
                     atol=3e-3, rtol=3e-3, grad_atol=3e-3, grad_rtol=3e-3)
   @unittest.skipIf(Device.DEFAULT == "WEBGPU" and platform.system() == "Windows", "Not accurate enough with DirectX backend")
+  @unittest.skipIf(getenv("NV_NAK"), "MUFU.SIN is not accurate enough")
   def test_cos(self):
     helper_test_op([(45,65)], lambda x: x.cos())
     helper_test_op([()], lambda x: x.cos())
@@ -847,6 +849,7 @@ class TestOps(unittest.TestCase):
       helper_test_op(None, lambda x: x.cos(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
                     atol=3e-3, rtol=3e-3, grad_atol=3e-3, grad_rtol=3e-3)
   @unittest.skipIf(Device.DEFAULT == "WEBGPU" and platform.system() == "Windows", "Not accurate enough with DirectX backend")
+  @unittest.skipIf(getenv("NV_NAK"), "MUFU.SIN is not accurate enough")
   def test_tan(self):
     # NOTE: backward has much higher diff with input close to pi/2 and -pi/2
     helper_test_op([(45,65)], lambda x: x.tan(), low=-1.5, high=1.5)

test/test_opts.py

@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor, Device
-from tinygrad.helpers import CPU_LLVM
+from tinygrad.helpers import CPU_LLVM, CPU_LVP
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.engine.realize import get_program
 
@@ -12,7 +12,7 @@ class TestOpts(unittest.TestCase):
     out = (a+b).contiguous(arg=opts)
     s = out.schedule()
     self.assertEqual(s[-1].ast.arg.opts_to_apply, opts)
-    if Device.DEFAULT in {"CPU", "CL", "METAL"} and not CPU_LLVM:
+    if Device.DEFAULT in {"CPU", "CL", "METAL"} and not CPU_LLVM and not CPU_LVP:
       prg = get_program(s[-1].ast)
       self.assertIn('float4', prg.src)
 

test/test_profiler.py

@@ -217,7 +217,7 @@ class TestProfiler(unittest.TestCase):
         Tensor.realize(a, b)
     profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
     exec_points = [e for e in profile if isinstance(e, ProfilePointEvent) and e.name == "exec"]
-    range_events = [e for e in profile if isinstance(e, ProfileRangeEvent)]
+    range_events = [e for e in profile if isinstance(e, ProfileRangeEvent) and not e.is_copy]
     self.assertEqual(len(exec_points), len(range_events), 2)
     self.assertEqual(len(dedup(e.key for e in exec_points)), 1)
     self.assertEqual(len(dedup(e.arg['metadata'] for e in exec_points)), 1)

test/test_randomness.py

@@ -6,6 +6,7 @@ from tinygrad.helpers import getenv, CI, OSX
 from tinygrad.device import is_dtype_supported
 from tinygrad.engine.realize import lower_schedule, CompiledRunner
 from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.nir import NIRRenderer
 from test.helpers import not_support_multi_device
 
 import numpy as np
@@ -100,7 +101,7 @@ class TestRandomness(unittest.TestCase):
 
     np.testing.assert_allclose(jr, r)
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "fails with PTX")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "PTX and NIR use pointer arithmetic")
   def test_threefry_doesnt_use_long(self):
     for (_,ei) in lower_schedule(Tensor.rand(20).schedule()):
       if isinstance(ei.prg, CompiledRunner):

test/test_schedule.py

@@ -11,7 +11,6 @@ from hypothesis import assume, given, settings, strategies as strat
 from tinygrad import nn, dtypes, Device, Tensor, Variable
 from tinygrad.device import is_dtype_supported
 from tinygrad.dtype import DType, ImageDType
-from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat
 from tinygrad.helpers import CI, DEBUG, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp
 from tinygrad.schedule.rangeify import get_rangeify_map, Kernel
@@ -2160,8 +2159,8 @@ class TestCopyFolding(unittest.TestCase):
     a = Tensor.ones((4,)).to("CPU")
     b = Tensor.empty(4, device="CPU")
     add = a+b
-    add.kernelize()
     assert all_same([x.device for x in add.uop.src]), f"ALU has different devices! {[x.device for x in add.src]}"
+    add.kernelize()
 
   def test_alu_before_copy(self):
     buf = Tensor.ones(1).contiguous().realize()
@@ -2251,7 +2250,7 @@ class TestBufferUOp(unittest.TestCase):
   def test_buffer_has_buffer(self):
     buf = Tensor.empty(10)
     self.assertIsNotNone(buf.uop.buffer)
-    self.assertEqual(buf.uop.st, ShapeTracker.from_shape((10,)))
+    self.assertEqual(buf.uop.shape, (10,))
     # the device Buffer remains unallocated until it's we run the schedule
     self.assertFalse(buf.uop.buffer.is_allocated())
     add = buf+1

test/test_setitem.py

@@ -52,7 +52,6 @@ class TestSetitem(unittest.TestCase):
 
   def test_setitem_into_noncontiguous(self):
     t = Tensor.ones(4)
-    self.assertFalse(t.uop.st.contiguous)
     with self.assertRaises(RuntimeError): t[1] = 5
 
   @unittest.skip("TODO: flaky")

test/test_symbolic_ops.py

@@ -1,6 +1,5 @@
 import unittest
 from tinygrad import Tensor, Variable, GlobalCounters
-from tinygrad.shape.shapetracker import View
 from tinygrad.uop.ops import sym_infer
 from tinygrad.dtype import dtypes
 from tinygrad.device import is_dtype_supported
@@ -64,14 +63,6 @@ class TestSymbolicOps(unittest.TestCase):
     self.test_attention(imin=4, imax=5, use_symbolic=False)
     self.test_attention(imin=4, imax=5, use_symbolic=True)
 
-  # until this works, symbolic single kernel softmax won't
-  @unittest.expectedFailure
-  def test_attention_simple_view(self):
-    i = Variable("i", 2, 10)
-    v1 = View.create((2,4,1,i,i), ((i*4),i,0,0,1))
-    v2 = View.create((2,4,1,i,i,i), (((i*i)*4),(i*i),0,0,i,1))
-    self.assertIsNotNone(v1+v2)
-
   def test_attention_training(self):
     with Tensor.train():
       self.test_attention(dropout_p=0.0)

test/test_tensor.py

@@ -9,6 +9,7 @@ from hypothesis import given, settings, strategies as strat
 from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.nir import NIRRenderer
 from tinygrad.codegen import full_rewrite
 from tinygrad.dtype import DType
 
@@ -526,6 +527,7 @@ class TestTinygrad(unittest.TestCase):
     self.assertListEqual(t.shrink_to(16).tolist(), list(range(16)))
     t = t.reshape(4, 8).contiguous().realize()
     self.assertListEqual(t.shrink_to(2, 2).tolist(), [[0, 1], [8, 9]])
+    self.assertListEqual(t.shrink_to(None, 2).tolist(), t.shrink_to(4, 2).tolist())
     with self.assertRaises(ValueError): t.shrink_to(2)
     with self.assertRaises(ValueError): t.shrink_to(2, 2, 2)
 
@@ -568,22 +570,6 @@ class TestMoveTensor(unittest.TestCase):
     np.testing.assert_equal(x.grad.numpy(), [[2,2,2],[0,0,0],[-2,-2,-2]])
 
 class TestZeroShapeTensor(unittest.TestCase):
-  def test_shape_is_expanded(self):
-    t = Tensor.empty(3, 2, 0)
-    assert t.shape == (3, 2, 0)
-    # numpy has stride 0, 0, 0; torch has stride 2, 1, 1
-    assert t.uop.st.is_expanded() == (True, True, True)
-
-    t = Tensor.empty(3, 0, 2)
-    assert t.shape == (3, 0, 2)
-    # numpy has stride 0, 0, 0; torch has stride 2, 2, 1
-    assert t.uop.st.is_expanded() == (True, True, True)
-
-    t = Tensor.empty(0, 0, 0)
-    assert t.shape == (0, 0, 0)
-    # numpy has stride 0, 0, 0; torch has stride 1, 1, 1
-    assert t.uop.st.is_expanded() == (True, True, True)
-
   def test_rand(self):
     t = Tensor.rand(3, 2, 0)
     assert t.shape == (3, 2, 0)
@@ -635,8 +621,10 @@ class TestZeroShapeTensor(unittest.TestCase):
 
     np.testing.assert_equal(Tensor([1, 2]).pad_to(4).numpy(), [1, 2, 0, 0])
     np.testing.assert_equal(Tensor([[1, 2]]).pad_to(2, 3).numpy(), [[1, 2, 0], [0, 0, 0]])
-    with self.assertRaises(TypeError): Tensor([1, 2]).pad_to(2, 3)
-    with self.assertRaises(TypeError): Tensor([[1, 2]]).pad_to(3)
+    np.testing.assert_equal(Tensor([[1, 2]]).pad_to(1, 3).numpy(), [[1, 2, 0]])
+    np.testing.assert_equal(Tensor([[1, 2]]).pad_to(None, 3).numpy(), [[1, 2, 0]])
+    with self.assertRaises(ValueError): Tensor([1, 2]).pad_to(2, 3)
+    with self.assertRaises(ValueError): Tensor([[1, 2]]).pad_to(3)
 
   def test_shrink_into_zero(self):
     t = Tensor.rand(3, 4).realize()
@@ -833,7 +821,6 @@ class TestTensorMetadata(unittest.TestCase):
     self.assertEqual(len(si.metadata), 3)
     self.assertEqual(set(m.name for m in si.metadata), {"relu", "sigmoid", "__mul__"})
 
-  @unittest.skip("not accurate")
   def test_complex_backward(self):
     x = Tensor.rand(3, requires_grad=True).realize()
     y = Tensor.rand(3, requires_grad=True).realize()
@@ -845,11 +832,12 @@ class TestTensorMetadata(unittest.TestCase):
     self.assertEqual(y.grad.uop.metadata[0].name, "sigmoid")
     self.assertTrue(y.grad.uop.metadata[0].backward)
     si = Tensor.schedule(out, x.grad, y.grad)[-1]
-    self.assertEqual(len(si.metadata), 3, f"failed with {si.metadata}")
-    self.assertSetEqual(set(m.name for m in si.metadata), {"sigmoid", "relu"})
+    self.assertEqual(len(si.metadata), 4, f"failed with {si.metadata}")
+    self.assertSetEqual(set(m.name for m in si.metadata), {"__mul__", "sigmoid", "relu"})
     bw = [m for m in si.metadata if m.backward]
-    self.assertEqual(len(bw), 1)
-    self.assertEqual(bw[0].name, "sigmoid")
+    self.assertEqual(len(bw), 2)
+    self.assertEqual(bw[0].name, "__mul__")
+    self.assertEqual(bw[1].name, "sigmoid")
 
 class TestIdxUpcast(unittest.TestCase):
   def _find_op(self, ast: UOp, op: Ops):
@@ -871,7 +859,8 @@ class TestIdxUpcast(unittest.TestCase):
     store = next(uop for uop in uops if uop.op is Ops.STORE)
     assert store.op is Ops.STORE
     idx = self._find_op(store, Ops.INDEX)
-    if idx is not None: # PTX turns Ops.INDEX into pointer arithmetic earlier than cstyle, plus it's already cast to int64
+    # PTX and NIR turn Ops.INDEX into pointer arithmetic earlier than cstyle, plus it's already cast to int64
+    if not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)):
       assert idx.op is Ops.INDEX
       idx_val = idx.src[1]
       assert idx_val.dtype is dtype
@@ -895,7 +884,7 @@ class TestIdxUpcast(unittest.TestCase):
   def test_regular_sym(self):
     self.do_op_then_assert(dtypes.int, 2048, 2048, UOp.variable("dim3", 1, 64).bind(32))
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX always convert Ops.INDEX to int64")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "PTX and NIR always converts Ops.INDEX to int64")
   def test_symfold(self):
     # This would cause an overflow, but after sym fold it's within int32
     a = Tensor.arange(65535)

test/test_tensor_uop.py

@@ -11,8 +11,7 @@ class TestTensorUOp(unittest.TestCase):
     def helper(a: np.ndarray):
       print(a.shape, a.strides, a.flags.c_contiguous)
       b = Tensor(a).uop
-      #assert b.st.contiguous == a.flags.c_contiguous
-      assert b.st.shape == a.shape
+      assert b.shape == a.shape
       np.testing.assert_equal(a, Tensor(b).numpy())
 
     for ndims in range(1, 4):

test/test_tensor_variable.py

@@ -93,7 +93,7 @@ class TestTensorVariable(unittest.TestCase):
     vb = v.bind(3)
     t = Tensor.empty(3, vb)
     assert t.uop.base.buffer.size == 30
-    assert t.uop.st.shape == (3, vb)
+    assert t.uop.shape == (3, vb)
 
 
 if __name__ == '__main__':

test/test_transcendental.py

@@ -149,6 +149,7 @@ class TestTranscendentalVectorized(unittest.TestCase):
     for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.log2, np.log2, (0.001, 200), vec_size)
 
   @unittest.skipIf(getenv("DSP"), "requires int division")
+  @unittest.skipIf(getenv("NV_NAK"), "MUFU.SIN is not accurate enough")
   def test_sin_vectorized(self):
     for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.sin, np.sin, (-100, 100), vec_size)
 

test/unit/test_disk_tensor.py

@@ -418,5 +418,32 @@ class TestPathTensor(unittest.TestCase):
       Tensor(pathlib.Path(test_file)).tolist()
     os.chmod(test_file, 0o644)
     assert Tensor(pathlib.Path(test_file)).tolist(), list(range(10))
+
+class TestDiskTensorMovement(unittest.TestCase):
+  def setUp(self):
+    self.fn = pathlib.Path(temp("custom_disk_range"))
+    self.fn.unlink(missing_ok=True)
+    Tensor.arange(100, dtype=dtypes.uint8).to(f"disk:{str(self.fn)}").realize()
+
+  def test_simple_read(self):
+    t = Tensor(self.fn)
+    self.assertTrue(Tensor.all(t.to(None) == Tensor.arange(100, dtype=dtypes.uint8)).item())
+
+  def test_slice_read(self):
+    t = Tensor(self.fn)
+    self.assertListEqual(t[16:18].tolist(), [16,17])
+
+  # TODO: fix this! at least assert on it
+  @unittest.expectedFailure
+  def test_slice_read_cat(self):
+    t = Tensor(self.fn)
+    self.assertListEqual(Tensor.cat(t[16:18], t[20:22]).tolist(), [16,17,20,21])
+
+  # TODO: fix this! at least assert on it
+  @unittest.expectedFailure
+  def test_slice_sum(self):
+    t = Tensor(self.fn)
+    self.assertListEqual((t[16:18]+t[20:22]).tolist(), [16+20,17+21])
+
 if __name__ == "__main__":
   unittest.main()

test/unit/test_elf.py

@@ -24,6 +24,15 @@ class TestElfLoader(unittest.TestCase):
     '''
     with self.assertRaisesRegex(RuntimeError, 'evil_external_function'):
       ClangJITCompiler().compile(src)
+  def test_link(self):
+    src = '''
+      float powf(float, float); // from libm
+      float test(float x, float y) { return powf(x, y); }
+    '''
+    args = ('-x', 'c', '-c', '-target', f'{platform.machine()}-none-unknown-elf', '-march=native', '-fPIC', '-O2', '-ffreestanding', '-nostdlib')
+    obj = subprocess.check_output(('clang',) + args + ('-', '-o', '-'), input=src.encode())
+    with self.assertRaisesRegex(RuntimeError, 'powf'): elf_loader(obj)
+    elf_loader(obj, link_libs=['m'])
 
 if __name__ == '__main__':
   unittest.main()

test/unit/test_graph_rewrite.py

@@ -334,5 +334,19 @@ class TestBidirectional(unittest.TestCase):
     graph_rewrite(c, pm, ctx=ctx_list, bpm=bpm)
     self.assertListEqual(ctx_list, [('+', True), (1, True), (1, False), (2, True), (2, False), ('+', False)])
 
+class TestStopEarly(unittest.TestCase):
+  def test_stop_early(self):
+    a = UOp.const(dtypes.int, 3)
+    b = UOp.const(dtypes.int, 4)
+    c = a+b
+    cn = UOp.const(dtypes.int, 7)
+    d = UOp.const(dtypes.int, 2)
+    def visit_const(c:UOp):
+      print(f"visit {c.arg}")
+      assert c.arg not in (3,4)
+    pm_cvisit = PatternMatcher([(UPat(Ops.CONST, name="c"), visit_const),])
+    ret = (c+d).substitute({c:cn}, extra_pm=pm_cvisit)
+    assert ret == cn+d
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_indexing.py

@@ -5,8 +5,6 @@ import numpy as np
 
 from tinygrad import Tensor, dtypes, Device, TinyJit
 from tinygrad.device import is_dtype_supported
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
 from tinygrad.helpers import CI, all_same, prod
 
 random.seed(42)
@@ -22,11 +20,13 @@ def consec(shape, start=1):
 # creates strided tensor with base set to reference tensor's base, equivalent to torch.set_()
 def set_(reference: Tensor, shape, strides, offset):
   raise NotImplementedError("need to implement without calling uop.view")
+  """
   if reference.uop.base.realized is None: reference.realize()
   assert reference.uop.base.realized, "base has to be realized before setting it to strided's base"
   strided = Tensor(reference.uop.view(ShapeTracker((View.create(shape=shape, strides=strides, offset=offset),))))
   assert strided.uop.st.real_strides() == strides, "real_strides should equal strides for strided"
   return strided
+  """
 
 def clone(original:Tensor): return original.clone()
 def copy_(src:Tensor, other:Tensor) -> Tensor: return src.clone()
@@ -501,10 +501,6 @@ class TestIndexing(unittest.TestCase):
     y = x[:, :, :, 1]
     z = y[:, 1:1, :]
     numpy_testing_assert_equal_helper((2, 0, 4), z.shape)
-    # this isn't technically necessary, but matches NumPy stride calculations.
-    # NOTE: this is empty and shouldn't have strides
-    numpy_testing_assert_equal_helper((True, True, True), z.uop.st.is_expanded())
-    self.assertTrue(z.uop.st.contiguous)
 
   @unittest.skip("bool indexing not supported")
   def test_index_getitem_copy_bools_slices(self):

test/unit/test_shapetracker.py

@@ -1,774 +0,0 @@
-#!/usr/bin/env python
-import unittest
-import numpy as np
-from tinygrad.dtype import dtypes, Invalid
-from tinygrad.helpers import prod
-from tinygrad.shape.shapetracker import ShapeTracker, View, views_to_valid_uop
-from tinygrad import Variable
-from tinygrad.uop.ops import UOp, Ops, graph_rewrite
-from tinygrad.codegen.late.devectorizer import sym
-from itertools import product
-
-def shapetracker_getitem(st:ShapeTracker, val:int):
-  valid_idx = views_to_valid_uop(st.reshape((st.size,)).views, (UOp.const(dtypes.int, val),))
-  idx, valid = valid_idx.get_idx(), valid_idx.get_valid()
-  idx, valid = graph_rewrite(idx, sym), graph_rewrite(valid, sym)
-  assert idx.op is Ops.CONST and valid.op is Ops.CONST
-  return idx.arg, valid.arg
-
-class CheckingShapeTracker:
-  def __init__(self, shape):
-    self.st = ShapeTracker.from_shape(shape)
-    self.t = np.arange(prod(shape), dtype=np.int32).reshape(shape)
-
-  @property
-  def shape(self):
-    return self.t.shape
-
-  def simplify(self):
-    self.st = self.st.simplify()
-    return self
-
-  def reshape(self, new_shape):
-    self.st = self.st.reshape(new_shape)
-    self.t = self.t.reshape(new_shape)
-    return self
-
-  def permute(self, axis):
-    self.st = self.st.permute(axis)
-    self.t = np.transpose(self.t, axis)
-    return self
-
-  def expand(self, new_shape):
-    self.st = self.st.expand(new_shape)
-    self.t = np.broadcast_to(self.t, new_shape)
-    return self
-
-  def flip(self, arg):
-    self.st = self.st.flip(arg)
-    self.t = np.flip(self.t, tuple(i for i in range(len(arg)) if arg[i]))
-    return self
-
-  def shrink(self, arg):
-    self.st = self.st.shrink(arg)
-    self.t = self.t[tuple([slice(x[0], x[1]) for x in arg])]
-    return self
-
-  def pad(self, arg):
-    self.st = self.st.pad(arg)
-    self.t = np.pad(self.t, arg, constant_values=-1)
-    return self
-
-  def __getitem__(self, val):
-    return self.t.flatten()[val]
-
-  @property
-  def views(self): return self.st.views
-
-  @property
-  def contiguous(self): return self.st.contiguous
-
-  def assert_same(self):
-    x = [(v[0] if (v:=shapetracker_getitem(self.st, i))[1] and v[0] is not Invalid else -1) for i in range(prod(self.st.shape))]
-    y = [self[i] for i in range(prod(self.shape))]
-    assert self.st.shape == self.shape
-    assert x == y, f"mismatch shapetracker:{x} real:{y}"
-
-@unittest.skip("don't create shapetrackers with views")
-class TestRealIssues(unittest.TestCase):
-  def test_reshape_doesnt_multiview(self):
-    self.st = ShapeTracker((View.create((256, 256, 2, 2, 2, 2, 2, 256, 8, 2), (0, 8, 0, 4, 0, 0, 2, 16384, 2048, 1), 0, None),))
-    self.st.reshape((128, 2, 256, 2, 2, 2, 2, 2, 256, 8, 2))
-    assert len(self.st.views) == 1
-
-  def test_reshape_stable_diffusion(self):
-    # regression test for https://github.com/tinygrad/tinygrad/pull/2616
-    st = ShapeTracker((View((2, 1920, 32, 32), (1310720, 1024, 32, 1), 0, ((0, 2), (0, 1280), (0, 32), (0, 32)), False),))
-    st = st.reshape((2, 32, 240, 256))
-    assert len(st.views) == 2
-
-  def test_reshape_trailing_invalid_ones(self):
-    st = ShapeTracker((View(shape=(1, 1, 5), strides=(0, 0, 1), offset=-5, mask=((1, 1), (0, 1), (0, 5)), contiguous=False),))
-    st = st.reshape((5,))
-    assert len(st.views) == 1
-    assert st.views[0].mask == ((0,0),)
-
-class TestRealDoesntSimplify(unittest.TestCase):
-  def tearDown(self):
-    self.st = self.st.simplify()
-    assert len(self.st.views) != 1
-
-  def test_1(self):
-    self.st = ShapeTracker((
-      View.create((8, 3, 1, 2, 11, 1), (33, 11, 0, 0, 1, 0), 0, None),
-      View.create((8, 6, 11), (66, 11, 1), 0, None)))
-    self.assertEqual(self.st.is_expanded(), (False, False, False))
-
-  def test_2(self):
-    self.st = ShapeTracker((
-      View.create((2, 2, 4, 3, 3), (72, 9, 18, -3, -1), 8, None),
-      View.create((4, 4, 3, 3), (36, 9, 3, 1), 0, None)))
-    self.assertEqual(self.st.is_expanded(), (False, False, False, False))
-
-class TestRealStrides(unittest.TestCase):
-  def test_1(self):
-    st = ShapeTracker((
-      View.create((2048,), (1,), 0, ((0, 512),)),
-      View.create((16, 32, 4), (128, 4, 1), 0, None),
-    ))
-    self.assertEqual(st.is_expanded(), (False, False, False))
-
-  def test_2(self):
-    # test/test_ops.py::TestOps::test_simple_padding_conv1d
-    st = ShapeTracker((
-      View.create((6, 2, 5, 14), (90, 45, 1, 5), 0, ((0, 6), (0, 2), (0, 5), (0, 9))),
-      View.create((6, 2, 78), (140, 70, 1), 0, ((0, 6), (0, 2), (0, 70))),
-      View.create((6, 2, 13, 6), (156, 78, 1, 13), 0, None),
-    ))
-    self.assertEqual(st.is_expanded(), (False, False, False, False))
-
-  def test_3(self):
-    # test/test_ops.py::TestOps::test_simple_cumsum
-    st = ShapeTracker((
-      View.create((4, 256, 512), (256, 0, 1), 0, ((0, 4), (0, 256), (0, 256))),
-      View.create((4, 131327), (131072, 1), 0, ((0, 4), (0, 131072))),
-      View.create((4, 511, 257), (131327, 1, 511), 0, None),
-    ))
-    self.assertEqual(st.is_expanded(), (False, False, False))
-
-  def test_4(self):
-    # test/test_nn.py::TestNN::test_conv_transpose1d
-    st = ShapeTracker((
-      View.create((4, 16, 56, 2), (896, 56, 1, 0), 0, ((0, 4), (0, 16), (0, 56), (0, 1))),
-      View.create((1, 4, 1, 16, 8, 121), (0, 1792, 0, 112, 0, 1), -5, ((0, 1), (0, 4), (0, 1), (0, 16), (0, 8), (5, 116))),
-      View.create((4, 64, 115, 16, 7), (15488, 0, 1, 968, 122), 0, None),
-    ))
-    self.assertEqual(st.is_expanded(), (False, True, False, False, False))
-
-  def test_5(self):
-    # test/test_ops.py::TestOps::test_conv2d
-    st = ShapeTracker((
-      View.create((1, 3, 1, 12, 2, 8), (0, 132, 0, 12, 1, 2), 0, ((0, 1), (0, 3), (0, 1), (0, 11), (0, 2), (0, 6))),
-      View.create((1, 3, 22, 21), (0, 192, 16, 1), 0, ((0, 1), (0, 3), (0, 12), (0, 16))),
-      View.create((3, 11, 7, 2, 3), (462, 21, 1, 231, 7), 0, None),
-    ))
-    self.assertEqual(st.is_expanded(), (False, False, False, True, False))
-
-class TestIndexExpressions2d(unittest.TestCase):
-  def setUp(self):
-    shapes = [(30, 5), (15, 10), (15, 1), (5, 10), (5, 1)] # Make sure dim0 is a multiple of 5, one of the tests divides this dimension by 5
-    offsets = [0, 1, 15, 28, 10000]
-    self.sts = [ShapeTracker.from_shape((prod(base_shape)+offset,)).shrink(((offset, offset+prod(base_shape)),)).\
-                reshape(base_shape) for base_shape in shapes for offset in offsets]
-    self.offset = [offset for base_shape in shapes for offset in offsets]
-    self.shapes = [shape for shape in shapes for offset in offsets]
-    self.idxs_exprs = []
-
-  def tearDown(self):
-    for st, offset, shape, idxs_expr in zip(self.sts, self.offset, self.shapes, self.idxs_exprs):
-      numel = prod(shape)
-      self.check_bounds(idxs_expr(self.default_idxs(st.shape)), offset, numel)
-      idx0s = [(0,0), (0, min(1, st.shape[0]-1)), (0, st.shape[0]-1), (min(3, st.shape[0]-1), min(6, st.shape[0]-1)), (st.shape[0]-1, st.shape[0]-1)]
-      idx1s = [(0,0), (0, min(1, st.shape[1]-1)), (0, st.shape[1]-1), (min(3, st.shape[1]-1), min(6, st.shape[1]-1)), (st.shape[1]-1, st.shape[1]-1)]
-      idx2s = [(0,0), (0, min(1, st.shape[2]-1)), (0, st.shape[2]-1), (min(3, st.shape[2]-1), min(6, st.shape[2]-1)),
-               (st.shape[2]-1, st.shape[2]-1)] if len(st.shape) == 3 else [None for _ in idx0s]
-      for idx0, idx1, idx2 in product(idx0s, idx1s, idx2s):
-        idxs = [Variable(f"idx{i}", idx[0], idx[1]) for i, idx in enumerate((idx0, idx1, idx2)) if idx is not None]
-        self.check_bounds(idxs_expr(idxs), offset, numel)
-
-  def default_idx(self, shape):
-    return Variable("idx", 0, prod(shape)-1)
-
-  def default_idxs(self, shape):
-    return [Variable(f"idx{i}", 0, d-1) for i,d in enumerate(shape)]
-
-  def check_bounds(self, expr, offset, numel):
-    assert expr.vmin >= offset
-    assert expr.vmax <= offset + numel - 1
-
-  def test_noop(self):
-    for st, base_shape, offset in zip(self.sts, self.shapes, self.offset):
-      self.idxs_exprs.append(lambda idxs, base_shape=base_shape, offset=offset: idxs[0]*base_shape[1] + idxs[1] + offset)
-
-  def test_permute(self):
-    new_st = []
-    for st, base_shape, offset in zip(self.sts, self.shapes, self.offset):
-      st = st.permute((1, 0))
-      self.idxs_exprs.append(lambda idxs, base_shape=base_shape, offset=offset: idxs[0] + idxs[1]*base_shape[1] + offset)
-      new_st.append(st)
-    self.sts = new_st
-
-  def test_reshape(self):
-    new_st = []
-    for st, base_shape, offset in zip(self.sts, self.shapes, self.offset):
-      st = st.reshape((base_shape[0], 1, base_shape[1]))
-      self.idxs_exprs.append(lambda idxs, base_shape=base_shape, offset=offset: idxs[0]*base_shape[1] + idxs[2] + offset)
-      new_st.append(st)
-    self.sts = new_st
-
-  def test_reshape_expand(self):
-    new_st = []
-    for st, base_shape, offset in zip(self.sts, self.shapes, self.offset):
-      st = st.reshape((base_shape[0], 1, base_shape[1]))
-      st = st.expand((base_shape[0], base_shape[1], base_shape[1]))
-      self.idxs_exprs.append(lambda idxs, base_shape=base_shape, offset=offset: idxs[0]*base_shape[1] + idxs[2] + offset)
-      new_st.append(st)
-    self.sts = new_st
-
-  def test_permute_reshape_1(self): # This tests multiple views
-    new_st = []
-    for st, base_shape, offset in zip(self.sts, self.shapes, self.offset):
-      st = st.permute((1, 0))
-      st = st.reshape((base_shape[0]//5, 1, base_shape[1]*5))
-      self.idxs_exprs.append(lambda idxs, base_shape=base_shape, offset=offset: (idxs[0]*(base_shape[1]*5)+idxs[2])%base_shape[0]*base_shape[1] + \
-                             (idxs[0]*(base_shape[1]*5)+idxs[2])//base_shape[0] + offset)
-      new_st.append(st)
-    self.sts = new_st
-
-  def test_permute_reshape_2(self):
-    new_st = []
-    for st, base_shape, offset in zip(self.sts, self.shapes, self.offset):
-      st = st.permute((1, 0))
-      st = st.reshape((1, base_shape[0]//5, base_shape[1]*5))
-      self.idxs_exprs.append(lambda idxs, base_shape=base_shape, offset=offset: (idxs[1]*(base_shape[1]*5)+idxs[2])%base_shape[0]*base_shape[1] + \
-                             (idxs[1]*(base_shape[1]*5)+idxs[2])//base_shape[0] + offset)
-      new_st.append(st)
-    self.sts = new_st
-
-  def test_reshaping_splitting(self):
-    self.st = CheckingShapeTracker((5,10,5,10))
-    self.st.permute((1, 0, 3, 2))
-    self.st.pad(((0,0), (0,5), (0,0), (0,5)))
-    self.st.reshape((10,2,5,10,2,5))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_reshape_splitting_1(self):
-    self.st = CheckingShapeTracker((1,10,1))
-    self.st.pad(((0,4),(0,0),(1,0)))
-    self.st.reshape((5,5,2,2))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_reshape_combining_1(self):
-    self.st = CheckingShapeTracker((2,1,10))
-    self.st.pad(((2,6), (0,0), (0,0)))
-    self.st.reshape((100,))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_reshape_combining_2(self):
-    self.st = CheckingShapeTracker((1,1,5))
-    self.st.pad(((3,6), (0,0), (0,5)))
-    self.st.reshape((100,))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_reshape_combining_3(self):
-    self.st = CheckingShapeTracker((1,1,4))
-    self.st.pad(((3,6), (0,0), (1,5)))
-    self.st.reshape((100,))
-    assert len(self.st.views) == 1
-    assert self.st.views[0].mask[0] == (31, 35)
-    self.st.assert_same()
-
-  def test_reshape_combining_4(self):
-    # interestingly this one is quite slow
-    self.st = CheckingShapeTracker((1,1,5,5,1,1,5))
-    self.st.pad(((2,1), (0,0), (0,2), (0,0), (2,1), (0,0), (0,2)))
-    self.st.reshape((28,5,28))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_reshape_splitting_combining(self):
-    self.st = CheckingShapeTracker((1,5,5))
-    self.st.pad(((0,4), (0,5), (0,0)))
-    self.st.reshape((10,25))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_reshape_only_1s(self):
-    self.st = CheckingShapeTracker((1, 1, 1, 4, 1, 3, 5, 1))
-    self.st.pad(((0,4), (0,0), (0,0), (1,1), (0,0), (0,0), (0,0), (0,0)))
-    self.st.reshape((5, 6, 3, 5))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-    self.st.reshape((1, 1, 5, 6, 3, 5, 1, 1))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-    self.st.reshape((1, 5, 6, 1, 3, 1, 5, 1))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_zero_mask_1(self):
-    self.st = CheckingShapeTracker((1, 3, 2))
-    self.st.pad(((0,0), (0,3), (0,0)))
-    self.st.shrink(((0,1), (3,6), (0,2)))
-    self.st.reshape((3,2))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-    self.st.reshape((1, 3, 1, 2, 1))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_zero_mask_2(self):
-    self.st = CheckingShapeTracker((1, 3, 2))
-    self.st.pad(((0,2), (0,3), (0,0)))
-    self.st.shrink(((2,3), (3,6), (0,2)))
-    self.st.reshape((3,2))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-    self.st.reshape((1, 3, 1, 2, 1))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_expanded_reshaped(self):
-    self.st = CheckingShapeTracker((1, 3, 2, 1))
-    self.st.expand((5, 3, 2, 2))
-    self.st.pad(((0,0), (0,3), (0,0), (0, 0)))
-    self.st.reshape((5, 2, 3, 2, 2))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_splitting_big(self):
-    self.st = CheckingShapeTracker((1, 5, 1, 15, 1))
-    self.st.pad(((0,0), (0,5), (0,0), (0,15), (0,0)))
-    self.st.reshape((10, 1, 30))
-    self.st.permute((2,1,0))
-    self.st.reshape((2,3,5,2,5))
-    assert len(self.st.views) == 1
-    v = self.st.views[-1]
-    assert v.strides == (0, 5, 1, 0, 15) and v.mask == ((0, 1), (0, 3), (0, 5), (0, 1), (0, 5))
-    self.st.assert_same()
-
-  def test_combining_big(self):
-    self.st = CheckingShapeTracker((1,3,1,5,3,1))
-    self.st.pad(((0,0),(2,2),(0,0),(0,0),(0,0),(0,0)))
-    self.st.reshape((1,1,1,105,1,1))
-    assert len(self.st.views) == 1
-    v = self.st.views[-1]
-    assert v.strides == (0, 0, 0, 1, 0, 0) and v.mask == ((0, 1), (0, 1), (0, 1), (30, 75), (0, 1), (0, 1)) and v.offset == -30
-    self.st.assert_same()
-
-  def test_pad_reshape(self):
-    self.st = CheckingShapeTracker((4,))
-    self.st.pad(((2,2),))
-    self.st.reshape((4,2))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-class TestSimplifyingShapeTracker(unittest.TestCase):
-  def setUp(self):
-    self.st = CheckingShapeTracker((1, 10))
-
-  def tearDown(self):
-    self.st.assert_same()
-
-  # multiview simplify
-  def test_expand_contract_simple(self):
-    self.st = self.st.expand((10, 10))
-    self.st = self.st.reshape((100,))
-    print(self.st.views)
-    assert (len(self.st.views) == 2)
-    self.st = self.st.reshape((10, 10))
-    print(self.st.views)
-
-    self.st = self.st.simplify()
-    print(self.st.views)
-    assert (len(self.st.views) == 1)
-
-  # multiview simplify
-  def test_expand_contract_different_shape(self):
-    self.st.expand((10, 10))
-    self.st.reshape((100,))
-    print(self.st.views)
-    assert (len(self.st.views) == 2)
-    self.st.reshape((2, 5, 2, 5))
-    print(self.st.views)
-
-    self.st = self.st.simplify()
-    print(self.st.views)
-    assert (len(self.st.views) == 1)
-
-  # multiview simplify
-  def test_expand_contract_still_complex(self):
-    self.st.expand((10, 10))
-    self.st.reshape((100,))
-    print(self.st.views)
-    assert (len(self.st.views) == 2)
-    self.st.reshape((5, 20))
-
-    self.st = self.st.simplify()
-    print(self.st.views)
-    assert (len(self.st.views) == 2)
-
-# Tensor.zeros(2, 4).permute(1,0).reshape(2, 4)
-# (d1*4 + d0%4), d1=x//4, d0=x%4 = ((x//4)*4) + (x%4)%4
-
-class TestComplexShapeTracker(unittest.TestCase):
-  def test_add_1s(self):
-    self.st = CheckingShapeTracker((4, 4))
-    self.st.permute((1,0))
-    self.st.reshape((1,4,1,4,1))
-    assert not self.st.contiguous
-    self.st.permute((0,3,2,1,4))
-    assert self.st.contiguous
-
-  def test_permute_1s_simple(self):
-    self.st = CheckingShapeTracker((1, 16, 9,9))
-    self.st.permute((1,0,2,3))
-    assert self.st.contiguous
-    self.st = CheckingShapeTracker((2, 16, 9,9))
-    self.st.permute((1,0,2,3))
-    assert not self.st.contiguous
-
-  def test_remove_1s_simple(self):
-    self.st = CheckingShapeTracker((1, 16, 1, 1))
-    self.st.reshape((16,))
-    assert self.st.contiguous
-
-  def test_remove_1s(self):
-    self.st = CheckingShapeTracker((1, 4, 1, 4, 1))
-    self.st.permute((0,3,2,1,4))
-    self.st.reshape((4,4))
-    assert not self.st.contiguous
-    self.st.permute((1,0))
-    assert self.st.contiguous
-
-  def test_permute_reshape(self):
-    self.st = CheckingShapeTracker((4, 4))
-    self.st.permute((1,0))
-    self.st.reshape((2, 2, 2, 2))
-    # TODO: should also be tested by test_super_complex
-    assert len(self.st.views) == 1
-
-  def test_factorize_split(self):
-    self.st = CheckingShapeTracker((4, 4))
-    self.st.permute((1,0))
-    self.st.reshape((2, 2, 2, 2))
-    self.st.permute((2,3,0,1))
-    assert self.st.contiguous
-
-  def test_factorize_combine(self):
-    self.st = CheckingShapeTracker((4, 4, 4))
-    self.st.permute((2, 0, 1))
-    self.st.reshape((4, 16))
-    self.st.permute((1, 0))
-    assert self.st.contiguous
-
-  def test_factorize_combine_add_ones(self):
-    self.st = CheckingShapeTracker((4, 4, 4))
-    self.st.permute((2, 0, 1))
-    self.st.reshape((4, 16, 1, 1))
-    self.st.permute((1, 0, 2, 3))
-    assert self.st.contiguous
-
-  def test_fancy_factorize(self):
-    self.st = CheckingShapeTracker((32, 3, 3, 1))
-    self.st.reshape((8, 4, 3, 3))
-    assert len(self.st.views) == 1
-
-  def test_super_complex_2_fail(self):
-    self.st = CheckingShapeTracker((4, 4, 4))
-    self.st.permute((2, 0, 1))
-    self.st.reshape((16, 4))
-    assert len(self.st.views) != 1
-
-  def test_work(self):
-    self.st = CheckingShapeTracker((64, 1024, 4))
-    self.st.reshape((1, 64, 128, 32))
-    self.st.permute((0, 3, 1, 2))
-    self.st.reshape((1, 32, 1, 64, 128))
-    self.st.permute((0, 3, 4, 1, 2))
-    assert self.st.contiguous
-
-  def test_work2(self):
-    self.st = CheckingShapeTracker((64, 1024, 4))
-    self.st.reshape((1, 64, 128, 32))
-    self.st.permute((0, 3, 1, 2))
-    self.st.reshape((1, 1, 32, 64, 128))
-    self.st.permute((0, 3, 4, 1, 2))
-    self.st.reshape((64, 1024, 4))
-    print(self.st.views)
-    assert self.st.contiguous
-
-class TestShapeTrackerEquality(unittest.TestCase):
-  def test_simple_equals(self):
-    self.assertEqual(ShapeTracker.from_shape((10,10)), ShapeTracker.from_shape((10,10)))
-  def test_other_equals(self):
-    st1 = ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True)))
-    st2 = ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True)))
-    self.assertEqual(st1, st2)
-
-class TestSingleShapeTracker(unittest.TestCase):
-  def setUp(self):
-    self.st = CheckingShapeTracker((7,4))
-
-  def tearDown(self):
-    self.st.assert_same()
-
-  def test_reshape(self):
-    self.st.reshape((7,1,4))
-    assert self.st.contiguous
-
-  def test_permute(self):
-    self.st.permute((1,0))
-    assert not self.st.contiguous
-
-  def test_shrink(self):
-    self.st.shrink(((1,2), (0,4)))
-    assert not self.st.contiguous
-
-  def test_double_permute(self):
-    self.st.permute((1,0))
-    self.st.permute((1,0))
-    assert self.st.contiguous
-
-  def test_reshape_permute(self):
-    self.st.reshape((7,1,4))
-    self.st.permute((0,1,2))
-    assert self.st.contiguous
-
-  def test_reshape_permute_yes(self):
-    self.st.reshape((7,1,4))
-    self.st.permute((0,2,1))
-    assert self.st.contiguous
-
-  def test_reshape_permute_no(self):
-    self.st.reshape((4,7))
-    self.st.permute((1,0))
-    assert not self.st.contiguous
-
-class TestShapeTrackerFuzzFailures(unittest.TestCase):
-  def setUp(self):
-    self.st = CheckingShapeTracker((3,3,3))
-  def tearDown(self):
-    self.st.assert_same()
-  def test_case_1(self):
-    self.st.shrink(((1, 2), (1, 3), (1, 3)))
-    self.st.reshape((1, 4))
-    self.st.shrink(((0, 1), (1, 3)))
-    self.st = self.st.simplify()
-  def test_case_2(self):
-    self.st.flip( (True, False, True) )
-    self.st.reshape( (3, 9) )
-    self.st.shrink( ((1, 2), (1, 5)) )
-    self.st.flip( (True, True) )
-  def test_case_3(self):
-    self.st.shrink( ((0, 2), (0, 2), (0, 1)) )
-    self.st.permute( (1, 0, 2) )
-    self.st.reshape( (4,) )
-    self.st.shrink( ((0, 3),) )
-    self.st.flip( (True, False) )
-  def test_case_4(self):
-    self.st.reshape( (3, 3, 3, 1) )
-    self.st.pad( ((0, 0), (0, 0), (0, 0), (1, 1)) )
-    self.st.shrink( ((0, 2), (1, 2), (0, 2), (0, 1)) )
-    self.st.expand( (2, 1, 2, 3) )
-
-class TestMaskedShapeTracker(unittest.TestCase):
-  def test_pad_1x1(self):
-    self.st = CheckingShapeTracker((1,1))
-    self.st.pad(((1,1), (1,1)))
-    self.st.assert_same()
-
-  def test_pad_2x2(self):
-    self.st = CheckingShapeTracker((2,2))
-    self.st.pad(((1,1), (1,1)))
-    self.st.assert_same()
-
-  def test_pad_reshape(self):
-    st1 = CheckingShapeTracker((1, 2))
-    st1.pad(((1, 0), (0, 1)))
-    st1.reshape((3, 2))
-    st1.assert_same()
-
-    st2 = CheckingShapeTracker((1, 2))
-    st2.pad(((1, 1), (0, 2)))
-    st2.reshape((4, 3))
-    st2.assert_same()
-
-    st3 = CheckingShapeTracker((1, 1, 1, 2))
-    st3.pad(((0, 2), (1, 2), (2, 2), (0, 4)))
-    st3.reshape((4, 3, 6, 5))
-    st3.assert_same()
-
-class TestShapeTracker(unittest.TestCase):
-  def setUp(self):
-    self.st = CheckingShapeTracker((7,4))
-    self.apply = lambda fxn: [fxn(x) for x in [self.st]]
-
-  def tearDown(self):
-    self.st.assert_same()
-
-  def test_noop(self):
-    pass
-
-  def test_simple_split(self):
-    self.test_permute()
-    self.apply(lambda x: x.reshape((prod(self.st.shape), )))
-
-  def test_simple_pad(self):
-    self.st.pad(((1,1), (1,1)))
-
-  def test_pad_shrink(self):
-    self.st.pad(((1,1), (1,1)))
-    self.st.shrink(((0,4), (0,4)))
-
-  def test_pad_one_sided(self):
-    self.st.pad(((0,1), (0,0)))
-
-  def test_pad_reshape(self):
-    self.st.pad(((0,1), (0,0)))
-    self.st.reshape((8*4,))
-
-  def test_pad_pad(self):
-    self.st.pad(((1,1), (1,1)))
-    self.st.pad(((1,1), (1,1)))
-
-  def test_pad_permute(self):
-    self.st.pad(((1,1), (2,2)))
-    self.st.permute((1,0))
-
-  def test_pad_expand(self):
-    self.st.reshape((7,4,1))
-    self.st.pad(((1,1), (1,1), (0,0)))
-    self.st.expand((9,6,4))
-
-  def test_pad_expand_alt(self):
-    self.st.pad(((1,1), (1,1)))
-    self.st.reshape((9,6,1))
-    self.st.expand((9,6,4))
-
-  def test_pad_flip(self):
-    self.st.pad(((1,4), (1,3)))
-    self.st.flip((True, False))
-
-  def test_pad_flip_int(self):
-    self.st.pad(((1,4), (1,3)))
-    self.st.flip((0, 1))
-
-  def test_reshape(self):
-    new_shape = self.st.shape[::-1]
-    self.apply(lambda x: x.reshape(new_shape))
-
-  def test_permute(self):
-    if len(self.st.shape) == 2: self.apply(lambda x: x.permute((1,0)))
-    elif len(self.st.shape) == 3: self.apply(lambda x: x.permute((2,0,1)))
-
-  def test_reshape_with_1(self):
-    new_shape = (self.st.shape[0], 1, self.st.shape[1])
-    self.apply(lambda x: x.reshape(new_shape))
-
-  def test_expand(self):
-    self.test_reshape_with_1()
-    new_shape = list(self.st.shape)
-    new_shape[1] = 2
-    self.apply(lambda x: x.expand(tuple(new_shape)))
-
-  def test_flip_0(self):
-    self.apply(lambda x: x.flip((True, False)))
-
-  def test_flip_1(self):
-    self.apply(lambda x: x.flip((False, True)))
-
-  def test_flip_01(self):
-    self.apply(lambda x: x.flip((True, True)))
-
-  def test_slice_0(self):
-    self.apply(lambda x: x.shrink(((1, x.shape[0]), (0, x.shape[1]))))
-
-  def test_slice_1(self):
-    self.apply(lambda x: x.shrink(((0, x.shape[0]), (1, x.shape[1]))))
-
-  def test_slice_1c1(self):
-    self.apply(lambda x: x.shrink(((0, 1), (0, 1))))
-
-  def test_slice_1c2(self):
-    self.apply(lambda x: x.shrink(((1, 2), (1, 2))))
-
-  def test_double_permute(self):
-    self.apply(lambda x: x.permute((1, 0)))
-    self.apply(lambda x: x.permute((1, 0)))
-
-  def test_slice_permute(self):
-    self.apply(lambda x: x.shrink(((0, 2), (2, 4))))
-    self.apply(lambda x: x.permute((1, 0)))
-
-  def test_slice_expand(self):
-    self.apply(lambda x: x.shrink(((0, 2), (3, 4))))
-    self.apply(lambda x: x.expand((2, 10)))
-
-  def test_double_flip(self):
-    self.apply(lambda x: x.flip((True, False)))
-    self.apply(lambda x: x.flip((True, False)))
-
-  def test_flip(self): self.apply(lambda x: x.flip((True, False)))
-  def test_flip2(self): self.apply(lambda x: x.flip((False, True)))
-  def test_flip3(self): self.apply(lambda x: x.flip((True, True)))
-
-  def test_reshape_then_permute(self):
-    self.test_reshape()
-    self.test_permute()
-
-  def test_reshape_then_expand(self):
-    self.test_reshape()
-    self.test_expand()
-
-  def test_permute_then_reshape(self):
-    self.test_permute()
-    self.test_reshape()
-
-  def test_expand_then_reshape(self):
-    self.test_expand()
-    self.test_reshape()
-
-  def test_combo(self):
-    self.test_permute()
-    self.test_reshape()
-    self.test_slice_1()
-    self.test_expand()
-    self.test_permute()
-
-class TestVariableShrink(unittest.TestCase):
-  def test_shrink(self):
-    st = ShapeTracker.from_shape((10,))
-    st = st.shrink(((0, Variable("i", 1, 10)),))
-    assert len(st.views) == 1
-
-  def test_shrink_bound(self):
-    st = ShapeTracker.from_shape((10,))
-    st = st.shrink(((0, Variable("i", 1, 10).bind(3)),))
-    assert len(st.views) == 1
-
-class TestVariableMerge(unittest.TestCase):
-  def test_add_reshape(self):
-    vi = Variable("i", 1, 10)
-    st1 = ShapeTracker.from_shape((vi,))
-    st2 = ShapeTracker.from_shape((1, vi,))
-    st = st1+st2
-    assert len(st.views) == 1
-
-  def test_add_stride_0(self):
-    st1 = ShapeTracker.from_shape((3,), (0,))
-    st2 = ShapeTracker.from_shape((Variable("i", 1, 10).bind(3),), (0,))
-    st = st1+st2
-    assert len(st.views) == 1, f"multiview {st}"
-
-  def test_add_reshape_bound(self):
-    vi = Variable("i", 1, 10).bind(3)
-    st1 = ShapeTracker.from_shape((vi,))
-    st2 = ShapeTracker.from_shape((1, vi,))
-    st = st1+st2
-    assert len(st.views) == 1
-
-  def test_simplify(self):
-    vi = Variable("i", 1, 10).bind(3)
-    st1 = ShapeTracker.from_shape((vi,))
-    st2 = ShapeTracker.from_shape((1, vi,))
-    st = ShapeTracker((st1.views[0], st2.views[0]))
-    st = st.simplify()
-    assert len(st.views) == 1
-
-if __name__ == '__main__':
-  unittest.main()

test/unit/test_shapetracker_math.py

@@ -1,108 +0,0 @@
-import unittest
-from tinygrad.helpers import prod
-from tinygrad.shape.view import View
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad import Variable
-from test.unit.test_shapetracker import shapetracker_getitem
-
-class MultiShapeTracker:
-  def __init__(self, sts:list[ShapeTracker]): self.sts = sts
-  @property
-  def shape(self): return self.sts[0].shape
-  def reshape(self, arg): self.sts = [x.reshape(arg) for x in self.sts]
-  def permute(self, arg): self.sts = [x.permute(arg) for x in self.sts]
-  def expand(self, arg): self.sts = [x.expand(arg) for x in self.sts]
-  def shrink(self, arg): self.sts = [x.shrink(arg) for x in self.sts]
-  def flip(self, arg): self.sts = [x.flip(arg) for x in self.sts]
-  def pad(self, arg): self.sts = [x.pad(arg) for x in self.sts]
-
-def st_equal(st1:ShapeTracker, st2:ShapeTracker) -> bool:
-  if st1.shape != st2.shape: return False
-  if st1 == st2: return True
-  for i in range(0, prod(st1.shape)):
-    st1_off, st1_v = shapetracker_getitem(st1, i)
-    st2_off, st2_v = shapetracker_getitem(st2, i)
-    if st1_v != st2_v or (st1_off != st2_off and st1_v):
-      print(f"ST MISMATCH @ {i}, {st1_v=} != {st2_v=}, {st1_off=} != {st2_off=}")
-      print(st1)
-      print(st2)
-      return False
-  return True
-
-class TestShapeTrackerBasics(unittest.TestCase):
-  def test_pad_shrink_removes_mask(self):
-    a = ShapeTracker.from_shape((10, 10))
-    a = a.pad(((0,2), (0,2)))
-    a = a.shrink(((0,10), (0,10)))
-    assert len(a.views) == 1 and a.views[-1].mask is None
-
-  def test_pad_shrink_leaves_mask(self):
-    a = ShapeTracker.from_shape((10, 10))
-    a = a.pad(((0,2), (0,2)))
-    a = a.shrink(((0,10), (0,11)))
-    assert len(a.views) == 1 and a.views[-1].mask is not None
-
-  def test_reshape_makes_same(self):
-    a = ShapeTracker.from_shape((2, 5))
-    x = a.pad( ((2, 0), (0, 0)) )
-    x = x.reshape( (2, 2, 5) )
-    x1 = x.reshape( (4, 5) )
-    x1 = x1.reshape( (2, 2, 5) )
-    assert x == x1.simplify()
-
-  def test_simplify_is_correct(self):
-    multiv = ShapeTracker(views=(View(shape=(15, 3), strides=(9, 1), offset=6, mask=None, contiguous=False),
-                                 View(shape=(4, 3), strides=(12, 4), offset=0, mask=None, contiguous=False)))
-    assert st_equal(multiv, multiv.simplify())
-
-class TestShapeTrackerAdd(unittest.TestCase):
-  def test_simple_add_reshape(self):
-    a = ShapeTracker.from_shape((10, 10))
-    a = a.reshape((100,))
-    b = ShapeTracker.from_shape((100,))
-    assert a+b == b
-
-  @unittest.skip("no longer simplifies")
-  def test_simple_add_permute(self):
-    a = ShapeTracker.from_shape((10, 10))
-    a = a.permute((1,0))
-    b = ShapeTracker.from_shape((10, 10))
-    b = b.permute((1,0))
-    assert a+b == ShapeTracker.from_shape((10, 10))
-
-  def test_plus_real1(self):
-    st = MultiShapeTracker([ShapeTracker.from_shape((15, 9))])
-    st.shrink( ((0, 15), (6, 9)) )
-    backup = st.sts[0]
-    st.sts.append(ShapeTracker.from_shape(backup.shape))
-    st.reshape( (45,) )
-    st.flip( (True,) )
-    st.reshape( (15, 3) )
-    assert st_equal(backup + st.sts[1], st.sts[0])
-
-  def test_off_by_one(self):
-    st1 = ShapeTracker(views=(View(shape=(5,), strides=(1,), offset=0, mask=None, contiguous=True),
-                              View(shape=(5,), strides=(1,), offset=0, mask=None, contiguous=True)))
-    st2 = ShapeTracker(views=(View(shape=(4,), strides=(1,), offset=0, mask=None, contiguous=True),
-                              View(shape=(5,), strides=(1,), offset=0, mask=None, contiguous=True)))
-    assert not (st_equal(st1, st2))
-
-class TestShapeTrackerAddVariable(unittest.TestCase):
-  def test_merge_symbolic_views(self):
-    var_i = Variable('i', 1, 10)
-    var_j = Variable('i', 1, 10)
-    vm1 = View(shape=(var_i, var_j, 3), strides=(3, 0, 1), offset=0, mask=None, contiguous=False)
-    vm2 = View(shape=(var_i, var_j, 3), strides=(var_j*3, 3, 1), offset=0, mask=None, contiguous=True)
-    ShapeTracker((vm1,)) + ShapeTracker((vm2,))
-
-  def test_merge_symbolic_views_2(self):
-    var_i = Variable('i', 1, 10)
-    var_j = Variable('j', 1, 10)
-    vm1 = View(shape=(var_i, var_j), strides=(0, 0), offset=0, mask=None, contiguous=False)
-    vm2 = View(shape=(var_i, var_j), strides=(var_j, 1), offset=0, mask=None, contiguous=True)
-    ret = (ShapeTracker((vm1,)) + ShapeTracker((vm2,))).reshape((var_i, var_j, 1))
-    ret_2 = ShapeTracker((vm1,)) + ShapeTracker((vm2,)).reshape((var_i, var_j, 1))
-    assert ret == ret_2
-
-if __name__ == '__main__':
-  unittest.main()

test/unit/test_symbolic_shapetracker.py

@@ -1,5 +1,4 @@
 import unittest
-from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad import Variable
 from tinygrad.tensor import Tensor
 
@@ -7,115 +6,33 @@ class TestSymbolic(unittest.TestCase):
   def assert_tuple_equal(self, x, y):
     for a,b in zip(x,y): self.assertFalse(a != b)
 
-  def test_symbolic_st(self):
-    x = Variable("x", 1, 100)
-    st = ShapeTracker.from_shape((x, 3))
-    self.assert_tuple_equal(st.shape, (x, 3))
-    self.assert_tuple_equal(st.is_expanded(), (False, False))
-
-  def test_is_expanded_0(self):
-    st = ShapeTracker(views=(View(shape=(2, (Variable('start_pos', 1, 8)+1), 1, 1), strides=(8, 1, 0, 0), offset=0, mask=((0, 2), (0, Variable('start_pos', 1, 8)), (0, 1), (0, 1)), contiguous=False), View(shape=(2, (Variable('start_pos', 1, 8)+1)), strides=((Variable('start_pos', 1, 8)+1), 1), offset=0, mask=None, contiguous=True)))   # noqa: E501
-    self.assert_tuple_equal(st.is_expanded(), (False, False))
-
-  def test_is_expanded_1(self):
-    st = ShapeTracker(views=(View(shape=(3, (Variable('i', 1, 10)+2)), strides=(Variable('i', 1, 10), 1), offset=0, mask=((0, 3), (0, Variable('i', 1, 10))), contiguous=False),))   # noqa: E501
-    self.assert_tuple_equal(st.is_expanded(), (False, False))
-
-  def test_is_expanded_2(self):
-    st = ShapeTracker(views=(View(shape=(3, (Variable('i', 1, 10)+Variable('j', 1, 10))), strides=(Variable('i', 1, 10), 1), offset=0, mask=((0, 3), (0, Variable('i', 1, 10))), contiguous=False),))   # noqa: E501
-    self.assert_tuple_equal(st.is_expanded(), (False, False))
-
-  def test_merge_view_recursion_err(self):
-    vm2 = View(shape=(Variable('j', 1, 10),), strides=(0,), offset=0, mask=None, contiguous=False)
-    vm1 = View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True)
-    self.assertEqual(vm2+vm1, None)
-
-  def test_merge_view_recursion_err2(self):
-    vm2 = View(shape=(Variable('a', 1, 10).bind(4),), strides=(0,), offset=0, mask=None, contiguous=False)
-    # NOTE: vm1 is different from what create function would give, and this test vm2+vm1 halts
-    vm1 = View(shape=(Variable('a', 1, 10).bind(4),), strides=(1,), offset=0, mask=((0, Variable('a', 1, 10).bind(4)),), contiguous=False)
-    self.assertEqual(vm2+vm1, None)
-
-    vm3 = View.create(shape=(Variable('a', 1, 10).bind(4),))
-    self.assertEqual(vm3.shape, vm1.shape)
-    self.assertEqual(vm3.strides, vm1.strides)
-    self.assertEqual(vm2+vm3, vm2)
-
   def test_cat_dim0_is_expanded(self):
     i = Variable("i", 1, 5).bind(3)
     j = Variable("j", 1, 5).bind(3)
     k = Variable("k", 1, 5).bind(3)
     t = Tensor.rand(5, 4)[:i].cat(Tensor.rand(5, 4)[:j], dim=0).cat(Tensor.rand(5, 4)[:k], dim=0)
-    st = t.uop.st
-    self.assert_tuple_equal(st.shape, (i+j+k, 4))
-    self.assert_tuple_equal(st.is_expanded(), (False, False))
+    self.assert_tuple_equal(t.shape, (i+j+k, 4))
     t = Tensor.rand(5, 3)[:i].cat(Tensor.rand(5, 3)[:i], dim=0).cat(Tensor.rand(3, 3), dim=0)
-    st = t.uop.st
-    self.assert_tuple_equal(st.shape, (2*i+3, 3))
-    self.assert_tuple_equal(st.is_expanded(), (False, False))
+    self.assert_tuple_equal(t.shape, (2*i+3, 3))
 
   def test_cat_dim1_strides(self):
     i = Variable("i", 1, 5).bind(4)
     j = Variable("j", 1, 5).bind(4)
     k = Variable("k", 1, 5).bind(4)
     t = Tensor.rand(3, 5)[:, :i].cat(Tensor.rand(3, 5)[:, :j], dim=1).cat(Tensor.rand(3, 5)[:, :k], dim=1)
-    st = t.uop.st
-    self.assert_tuple_equal(st.shape, (3, i+j+k))
-    self.assert_tuple_equal(st.is_expanded(), (False, False))
+    self.assert_tuple_equal(t.shape, (3, i+j+k))
 
 class TestSymbolicVarVals(unittest.TestCase):
   def assert_equal(self, x, y): self.assertFalse(x != y)
-  def test_var_vals_empty(self):
-    assert ShapeTracker.from_shape((3, 4, 5)).var_vals == {}
-
-  def test_var_vals_shape(self):
-    x = Variable("x", 1, 100).bind(3)
-    assert ShapeTracker.from_shape((x, 3)).var_vals == {"x": 3}
-
-  def test_var_vals_offset(self):
-    x = Variable("x", 1, 100).bind(3)
-    st = ShapeTracker.from_shape((4, 3)).shrink(((x, x+1), (0, 3)))
-    self.assert_equal(st.views[-1].offset, x * 3)
-    assert st.var_vals == {"x": 3}
-
-  def test_var_vals_mask(self):
-    x = Variable("x", 1, 100).bind(3)
-    view = View.create(shape=(3,4), strides=(4,1), offset=0, mask=((0, x), (0, 4)))
-    st = ShapeTracker(views=(view,))
-    assert st.var_vals == {"x": 3}
-
-  def test_var_vals_complex(self):
-    x = Variable("x", 1, 100).bind(3)
-    y = Variable("y", 1, 100).bind(4)
-    z = Variable("z", 1, 100).bind(5)
-    st = ShapeTracker.from_shape((x, 5, y)).shrink(((0, x), (z, z+1), (0, 3)))
-    self.assert_equal(st.views[-1].offset, y * z)
-    assert st.var_vals == {"x": 3, "y": 4, "z": 5}
-
-  def test_shrink_reshape(self):
-    x = Variable("x", 1, 100).bind(3)
-    st = ShapeTracker.from_shape((10, 10, 10)).shrink(((x, x+3), (3, 7), (2, 5)))
-    st = st.reshape((3*4*3,))
-    assert st.var_vals == {"x": 3}
-
-class TestShapeTrackerUnbind(unittest.TestCase):
-  def test_view_unbind(self):
-    v = Variable("v", 1, 100)
-    bv = Variable("v", 1, 100).bind(3)
-    unbound_view, var_val = View.create(shape=(bv, 4)).unbind()
-    assert unbound_view == View.create(shape=(v, 4))
-    assert var_val == {v: 3}
 
   def test_shrink_unbind(self):
     v = Variable("v", 1, 100)
     bv = Variable("v", 1, 100).bind(2)
     t = Tensor.rand(3, 4).shrink(((0,bv),(0,4)))
-    unbound_st, var_val = t.uop.st.unbind()
-    assert unbound_st == ShapeTracker((View.create(shape=(v, 4)),))
+    unbound_st, var_val = t.uop.unbind_all()
     assert var_val == {v: 2}
     t = Tensor.rand(3, 4).shrink(((bv, bv+1), (0, 4)))
-    unbound_st, var_val = t.uop.st.unbind()
-    assert unbound_st == ShapeTracker((View.create(shape=(1, 4), offset=4*v),))
+    unbound_st, var_val = t.uop.unbind_all()
     assert var_val == {v: 2}
 
 class TestSymbolicReshape(unittest.TestCase):
@@ -145,17 +62,6 @@ class TestSymbolicReshape(unittest.TestCase):
         ret = ret.reshape(1, vi*vj)
         assert ret.shape == (1, vi*vj)
 
-  def test_symbolic_mask(self):
-    # taken from gpt2 single kvcache
-    # these two caused problems in gpt2 if reshape merged views
-    view = View(shape=(1, (Variable('start_pos', 1, 128).bind(2)+1), 16, 64), strides=(0, 0, 64, 1), offset=1024, mask=((0, 1), (Variable('start_pos', 1, 128).bind(2), (Variable('start_pos', 1, 128).bind(2)+1)), (0, 16), (0, 64)), contiguous=False)   # noqa: E501
-    new_shape = (1, 1, (Variable('start_pos', 1, 128).bind(2)+1), 16, 64)
-    assert view.reshape(new_shape) is None
-
-    view = View(shape=(2, 1, (Variable('start_pos', 1, 128)+1), 16, 64), strides=(0, 0, 1024, 64, 1), offset=131072, mask=((1, 2), (0, 1), (0, (Variable('start_pos', 1, 128)+1)), (0, 16), (0, 64)), contiguous=False)   # noqa: E501
-    new_shape = (2, (Variable('start_pos', 1, 128)+1), 16, 64)
-    assert view.reshape(new_shape) is None
-
 class TestSymbolicExpand(unittest.TestCase):
   def test_expand_into_symbols(self):
     vi = Variable("i", 1, 5).bind(3)
@@ -198,6 +104,5 @@ class TestSymbolicPad(unittest.TestCase):
     t = t[:9]
     assert t.tolist() == [0,0,0,0,1,1,1,1,1]
 
-
 if __name__ == '__main__':
   unittest.main()

test/unit/test_tensor_uop_representation.py

@@ -4,8 +4,6 @@ from tinygrad.uop.ops import UPat, Ops, UOp
 
 # NOTE: unlike before base for a realized tensor is always a BUFFER
 realized_pattern = UPat(Ops.BUFFER)
-# after realization, base tensor uops become RESHAPE(BUFFER)
-buffer_view_pattern = UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER),))
 def is_pattern_uop(u:UOp, pat:UPat): assert pat.match(u, {}), f"{u}\nis not\n{pat}"
 def is_pattern(ten:Tensor, pat:UPat): is_pattern_uop(ten.uop, pat)
 
@@ -57,13 +55,5 @@ class TestTensorUopRepresentation(unittest.TestCase):
     is_pattern(c, UPat(Ops.ADD))
     for s in c.uop.src: is_pattern_uop(s.base, realized_pattern)
 
-  def test_empty_buf(self):
-    a = Tensor.empty(3, 3)
-    is_pattern(a, UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER),)))
-    vi = UOp.variable("i", 1, 3).bind(1)
-    a = Tensor.empty(3, vi)
-    is_pattern(a, UPat(Ops.RESHAPE, src=(UPat(Ops.SHRINK, src=(UPat(Ops.BUFFER),))),))
-    self.assertEqual(a.uop.base.buffer.size, 9)
-
 if __name__ == '__main__':
   unittest.main()

test/unit/test_view.py

@@ -1,73 +0,0 @@
-#!/usr/bin/env python
-import unittest
-from tinygrad.shape.view import View, merge_dims
-# from tinygrad.shape.shapetracker import ShapeTracker
-
-class TestView(unittest.TestCase):
-  def test_canonicalize_empty_mask(self):
-    v = View.create(shape=(2,2,2), strides=(4,2,1), mask=((0,2),(0,2),(0,2)))
-    self.assertIsNone(v.mask)
-    v = View.create(shape=(4,3,2), strides=(1,4,10), mask=((0,4),(0,3),(0,2)))
-    self.assertIsNone(v.mask)
-
-  def test_empty_mask_contiguous(self):
-    v1 = View.create(shape=(2,2,2), strides=(4,2,1), mask=None)
-    v2 = View.create(shape=(2,2,2), strides=(4,2,1), mask=((0,2),(0,2),(0,2)))
-    self.assertEqual(v1.contiguous, v2.contiguous)
-    v1 = View.create(shape=(1,1,1,4), strides=(0,0,0,1), offset=0, mask=None)
-    v2 = View.create(shape=(1,1,1,4), strides=(0,0,0,1), offset=0, mask=((0,1),(0,1),(0,1),(0,4)))
-    self.assertEqual(v1.contiguous, v2.contiguous)
-    v = View.create(shape=(2,3,4), mask=((0,2),(0,3),(0,4)))
-    self.assertTrue(v.contiguous)
-
-  def test_reshape_all_invalid(self):
-    v = View.create((4,5), mask=((0,0), (0,0))).reshape((20,))
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View.create((20,), mask=((0,0),)))
-
-  def test_add_0(self):
-    v1 = View.create((2,3,4))
-    v2 = View.create((2,0,4))
-    self.assertEqual(v2, v1+v2)
-
-  def test_add_0_masked(self):
-    v1 = View.create((2,3,4), mask=((0, 0), (0, 0), (0, 0)))
-    v2 = View.create((2,0,4))
-    self.assertEqual(v2, v1+v2)
-
-class TestMergeDims(unittest.TestCase):
-  def test_contiguous(self):
-    shape = (2, 3, 4)
-    strides = (12, 4, 1) #=strides_for_shape(shape)
-    m = merge_dims(shape, strides)
-    self.assertEqual(m, ((24, 1, 24),))
-
-  def test_0_in_strides(self):
-    shape = (2, 3, 4)
-    self.assertEqual(merge_dims(shape, (0, 4, 1)), ((2, 0, 0), (12, 1, 12)))
-    self.assertEqual(merge_dims(shape, (0, 0, 1)), ((6, 0, 0), (4, 1, 4)))
-    self.assertEqual(merge_dims(shape, (3, 1, 0)), ((6, 1, 6), (4, 0, 4)))
-    self.assertEqual(merge_dims(shape, (0, 0, 0)), ((24, 0, 0),))
-
-  def test_pad(self):
-    # print(ShapeTracker.from_shape((1, 2)).pad(((1, 0), (0, 1))).views[-1])
-    self.assertEqual(merge_dims((2, 3), (0, 1), ((1, 2), (0, 2))), ((6, 1, 3),))
-
-    # print(f"{ShapeTracker.from_shape((1, 1, 2)).pad(((1, 0), (1, 0), (0, 1))).views[-1]}")
-    self.assertEqual(merge_dims((2, 2, 3), (0, 0, 1), ((1, 2), (1, 2), (0, 2))), ((12, 1, 3),))
-
-    # print(f"{ShapeTracker.from_shape((1, 1, 2, 2)).pad(((1, 0), (1, 0), (0, 1), (0, 1))).views[-1]}")
-    self.assertEqual(merge_dims((2, 2, 3, 3), (0, 0, 2, 1), ((1, 2), (1, 2), (0, 2), (0, 2))), ((12, 2, 3), (3, 1, 3)))
-
-    # print(f"{ShapeTracker.from_shape((2, 1, 2)).pad(((0, 0), (1, 0), (0, 1))).views[-1]}")
-    self.assertEqual(merge_dims((2, 2, 3), (2, 0, 1), ((0, 2), (1, 2), (0, 2))), ((2, 2, 2), (6, 1, 3)))
-
-  def test_different_1_pad(self):
-    # print(f"{ShapeTracker.from_shape((2, 2, 1)).pad(((0, 0), (0, 0), (0, 1))).views[-1]}")
-    self.assertEqual(merge_dims((2, 2, 2), (2, 1, 0), ((0, 2), (0, 2), (0, 1))), ((4, 1, 4), (2, 0, 2)))
-
-    # print(f"{ShapeTracker.from_shape((2, 1, 1)).pad(((0, 0), (0, 1), (0, 1))).views[-1]}")
-    self.assertEqual(merge_dims((2, 2, 2), (1, 0, 0), ((0, 2), (0, 2), (0, 1))), ((2, 1, 2), (4, 0, 4)))
-
-if __name__ == '__main__':
-  unittest.main()

test/unit/test_viz.py

@@ -2,8 +2,7 @@ import unittest, decimal, json, struct
 from dataclasses import dataclass
 from typing import Generator
 
-from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher
-from tinygrad.uop.ops import graph_rewrite, track_rewrites, TRACK_MATCH_STATS
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, track_rewrites, TRACK_MATCH_STATS
 from tinygrad.uop.symbolic import sym
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import PROFILE, colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker
@@ -15,15 +14,16 @@ def exec_rewrite(sink:UOp, pm_lst:list[PatternMatcher], names:None|list[str]=Non
     sink = graph_rewrite(sink, TrackedPatternMatcher(pm.patterns), name=names[i] if names else None)
   return sink
 
-# real VIZ=1 pickles these tracked values
-from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, _name_cnt
-traces = [(tracked_keys, tracked_ctxs, uop_fields)]
-from tinygrad.viz.serve import get_metadata, uop_to_json, get_details
-def get_viz_list(): return get_metadata(traces)
+# real VIZ=1 loads the trace from a file, we just keep it in memory for tests
+from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, _name_cnt, RewriteTrace
+from tinygrad.viz import serve
+serve.trace = RewriteTrace(tracked_keys, tracked_ctxs, uop_fields)
+from tinygrad.viz.serve import get_rewrites, get_full_rewrite, uop_to_json
+def get_viz_list(): return get_rewrites(serve.trace)
 def get_viz_details(rewrite_idx:int, step:int) -> Generator[dict, None, None]:
   lst = get_viz_list()
   assert len(lst) > rewrite_idx, "only loaded {len(lst)} traces, expecting at least {idx}"
-  return get_details(tracked_ctxs[rewrite_idx][step])
+  return get_full_rewrite(tracked_ctxs[rewrite_idx][step])
 
 class BaseTestViz(unittest.TestCase):
   def setUp(self):
@@ -333,7 +333,7 @@ def load_profile(lst:list[ProfileEvent]) -> dict:
       for _ in range(event_count):
         alloc, ts, key = u("<BII")
         if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
-        else: v["events"].append({"event":"free", "ts":ts, "key":key})
+        else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg":{"users":[u("<I")[0] for _ in range(u("<I")[0])]}})
   return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}
 
 class TestVizProfiler(unittest.TestCase):
@@ -479,5 +479,16 @@ class TestVizMemoryLayout(BaseTestViz):
     self.assertEqual(len(ret["events"]), 6)
     self.assertEqual(len(profile["markers"]), 6)
 
+  def test_producer_simple(self):
+    a = Tensor.empty(10, device="NULL")
+    Tensor.realize(a.add(1), a.add(2))
+    b = Tensor.empty(10, device="NULL")
+    Tensor.realize(b.add(1))
+    profile = load_profile(cpu_events+Buffer.profile_events)
+    buffers = profile["layout"]["NULL Memory"]["events"]
+    programs = profile["layout"]["NULL"]["events"]
+    user_cnt = [len(b["arg"]["users"]) for b in buffers if b["arg"].get("users")]
+    self.assertEqual(len(user_cnt), len(programs))
+
 if __name__ == "__main__":
   unittest.main()

tinygrad/codegen/__init__.py

@@ -9,7 +9,7 @@ from tinygrad.renderer import Renderer
 # import all pattern matchers here
 from tinygrad.codegen.quantize import pm_quant
 from tinygrad.codegen.gpudims import pm_add_gpudims
-from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic
+from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic, pm_move_where_on_load
 from tinygrad.uop.decompositions import get_late_rewrite_patterns
 from tinygrad.codegen.late.expander import migrate_indexing, expander, pm_pre_expander, pm_group_for_reduce
 from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
@@ -62,7 +62,7 @@ def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _Q
     ret.append(RewriteStep(pm_postrange_opt, ctx=lambda _: opts, name="post optimize ast"))
 
   # ** expander (expand_rewrite) **
-  ret.append(RewriteStep(sym+migrate_indexing, name="postopt symbolic"))
+  ret.append(RewriteStep(sym+migrate_indexing+pm_move_where_on_load, name="postopt symbolic"))
 
   # expand
   ret.append(RewriteStep(sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander"))

tinygrad/codegen/quantize.py

@@ -26,14 +26,6 @@ pm_quant = symbolic+PatternMatcher([
   # x*c1 + y*c2 -> (x+y)*c1 (if c1 and c2 are close floats)
   (UPat.var("x")*UPat.cvar("c1", dtype=dtypes.floats) + UPat.var("y")*UPat.cvar("c2", dtype=dtypes.floats),
    lambda x,y,c1,c2: (x+y)*c1 if abs(c1.arg-c2.arg) < 1e-9 else None),
-  # mul 0 * c1 is 0
-  #(UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c1"), UPat(Ops.CONST, arg=0)) *
-  # UPat(Ops.LOAD, src=(UPat().view(name="v"),)).cast(dtypes.int).cast(dtypes.float).named("ld"), lambda ld,v,c1: ld*c1),
-  # mul (with plus) 0 * c1 is 0
-  #(UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c1"), UPat(Ops.CONST, arg=0)) *
-  # (UPat(Ops.LOAD, src=(UPat().view(name="v"),)).cast(dtypes.int) + \
-  #  UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar(), UPat(Ops.CONST, arg=0))).cast(dtypes.float).named("ld"),
-  #    lambda ld,v,c1: ld*c1),
 
   # const push through add
   ((UPat.var("x")*UPat.cvar("c1") + UPat.var("y")*UPat.cvar("c2")) * UPat.cvar("c3"), lambda x,y,c1,c2,c3: (x*c1*c3) + (y*c2*c3)),

tinygrad/codegen/simplify.py

@@ -1,5 +1,5 @@
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, ImageDType
-from tinygrad.uop.symbolic import symbolic_flat, sym, invalid_pat
+from tinygrad.uop.symbolic import symbolic_flat, sym
 from tinygrad.helpers import partition
 from tinygrad.dtype import dtypes
 
@@ -89,9 +89,9 @@ pm_reduce_collapse = PatternMatcher([
   # lift x+y out of reduce on ne
   ((UPat.var("x")+UPat.var("y")).or_casted() != UPat.var("c"), lambda x,y,c: (x != (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
   # fold the range
-  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(0, UPat.cvar("val")).reduce(arg=Ops.ADD, allow_any_len=True),
+  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(0, UPat.cvar("val")).reduce(UPat.var("r"), arg=Ops.ADD),
    lambda r,cut,val: (r.src[0]-cut).maximum(0).minimum(r.src[0]).cast(val.dtype) * val),
-  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(UPat.cvar("val"), 0).reduce(arg=Ops.ADD, allow_any_len=True),
+  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(UPat.cvar("val"), 0).reduce(UPat.var("r"), arg=Ops.ADD),
    lambda r,cut,val: cut.maximum(0).minimum(r.src[0]).cast(val.dtype) * val),
   # REDUCE on ADD
   ((UPat.var("x")+UPat.var("y")).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
@@ -100,9 +100,8 @@ pm_reduce_collapse = PatternMatcher([
   ((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast().or_broadcasted(name="b")),
    lambda x,gate,b=None: gate.broadcast(x.dtype.count).where(x, 0) if b is not None else gate.where(x, 0)),
   # reduce on gated load becomes can substitute the range and remove the reduce
-  (UPat.var("buf").index(UPat.var("idx").eq(UPat(Ops.RANGE, name="r").or_casted()).where(UPat.var("expr"), invalid_pat)).load()
-    .reduce(arg=Ops.ADD, allow_any_len=True), lambda buf,r,idx,expr,i:
-    buf.index(expr.substitute({r:idx.cast(r.dtype)}).valid((idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0]))).load()),
+  ((UPat.var("idx")!=(UPat(Ops.RANGE, name="r").or_casted())).where(0, UPat.var("expr")).reduce(UPat.var("r"), arg=Ops.ADD),
+   lambda r,idx,expr: (v:=(idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])).where(expr.substitute({r:idx.cast(r.dtype).valid(v)}),0)),
   # AND on WHERE
   ((UPat.any(UPat(Ops.DEFINE_VAR, name="x"), UPat(Ops.DEFINE_VAR).gep(name="x")) & UPat.var("y")) \
    .where(UPat.cvar("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),

tinygrad/device.py

@@ -327,8 +327,8 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
   if device is None: device = Device.DEFAULT
   if dtype == dtypes.bfloat16:
     if device == "METAL": return not CI
-    if device in {"CUDA", "NV"}: return not CI and not getenv(f"{device}_PTX")
-    if device in {"CPU"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"}
+    if device in {"CUDA", "NV"}: return not CI and not getenv(f"{device}_PTX") and not getenv("NV_NAK")
+    if device in {"CPU"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and not getenv("CPU_LVP")
     return device in {"AMD", "PYTHON", "NULL"}
   if dtype in dtypes.fp8s: return device in {"PYTHON", "NULL"}
   if device == "WEBGPU": return dtype in [dtypes.bool, dtypes.char, dtypes.uchar, dtypes.short,

tinygrad/engine/jit.py

@@ -5,10 +5,10 @@ from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv,
 from tinygrad.device import Buffer, Compiled, Device, MultiBuffer
 from tinygrad.dtype import DType
 from tinygrad.uop.ops import UOp, Variable, sym_infer, Ops
-from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.engine.realize import ExecItem, capturing, ViewOp, BufferCopy, BufferXfer, CompiledRunner, Runner, Estimates
 from tinygrad.engine.memory import _internal_memory_planner
 from tinygrad.nn.state import get_parameters
+from tinygrad.schedule.rangeify import mop_cleanup
 from dataclasses import dataclass
 from weakref import WeakKeyDictionary
 
@@ -158,7 +158,7 @@ class CapturedJit(Generic[ReturnType]):
   input_replace: dict[tuple[int, int], int]
   extra_view_inputs: list[tuple[int, int, str, int, DType]]
   expected_names: list[int|str]
-  expected_st_vars_dtype_device: list[tuple[ShapeTracker, tuple[Variable, ...], DType, str]]
+  expected_st_vars_dtype_device: list[tuple[UOp, tuple[Variable, ...], DType, str]]
 
   def __reduce__(self):
     # TODO: free_intermediates here? replan_buffers_memory_layout here?
@@ -224,7 +224,7 @@ def _prepare_jit_inputs(args, kwargs):
   input_buffers: list[Buffer] = flatten([rb.bufs if isinstance(rb:=lb.base.realized, MultiBuffer) else [rb]
                                          for lb in lbs if lb.base.realized is not None])
   assert len(set(input_buffers)) == len(input_buffers), "duplicate inputs to JIT"
-  st_varval_dtype_device = [(*unwrap(lb.st).unbind(), lb.dtype, lb.device) for lb in lbs]
+  st_varval_dtype_device = [(*(lb.substitute({lb.base:UOp(Ops.NOOP)}, extra_pm=mop_cleanup).unbind_all()), lb.dtype, lb.device) for lb in lbs]
   _var_vals = merge_dicts([x[1] for x in st_varval_dtype_device] + [dict(v.unbind() for v in (args + tuple(kwargs.values())) if isinstance(v, UOp))])
   var_vals = {k.expr:v for k,v in _var_vals.items()}
   st_vars_dtype_device = [(x[0], tuple(sorted(x[1].keys(), key=lambda v: v.expr)), x[2], x[3]) for x in st_varval_dtype_device]

tinygrad/gradient.py

@@ -31,12 +31,12 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.REDUCE_AXIS, name="ret"), reduce_gradient),
   (UPat((Ops.CONTIGUOUS, Ops.FUSE)), lambda ctx: (ctx,)),
   (UPat(Ops.CONTIGUOUS_BACKWARD), lambda ctx: (ctx.contiguous(),)),
-  (UPat(Ops.RESHAPE, name="ret"), lambda ctx, ret: (ctx.reshape(ret.src[0].shape),)),
-  (UPat(Ops.PERMUTE, name="ret"), lambda ctx, ret: (ctx.permute(argsort(ret.arg)),)),
-  (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.arg)])),)),
-  (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[0], s-p[1]) for s,p in zip(ret.src[0].shape, ret.arg)])),)),
-  (UPat(Ops.FLIP, name="ret"), lambda ctx, ret: (ctx.flip(ret.arg),)),
-  (UPat(Ops.EXPAND, name="ret"), lambda ctx, ret: (ctx.r(Ops.ADD, tuple(i for i,(si,so) in enumerate(zip(ret.src[0].shape, ret.arg)) if si!=so)),)),
+  (UPat(Ops.RESHAPE, name="ret"), lambda ctx, ret: (ctx.reshape(ret.src[0].shape), None)),
+  (UPat(Ops.EXPAND, name="ret"), lambda ctx, ret: (ctx.r(Ops.ADD,tuple(i for i,(s,n) in enumerate(zip(ret.src[0].shape, ret.shape)) if s!=n)), None)),
+  (UPat(Ops.PAD, name="ret"), lambda ctx, ret: (ctx.shrink(tuple([(p[0], s+p[0]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
+  (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[0], s-p[1]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
+  (UPat(Ops.PERMUTE, name="ret"), lambda ctx, ret: (ctx.permute(argsort(ret.marg)),)),
+  (UPat(Ops.FLIP, name="ret"), lambda ctx, ret: (ctx.flip(ret.marg),)),
   (UPat(Ops.MULTI, name="ret"), lambda ctx, ret: ctx.shard(ret.device, ret.axis).src),
   # there's no gradient for bitcast
   (UPat(Ops.BITCAST), lambda: (None,)),

tinygrad/helpers.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass
 import urllib.request, subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools
 from dataclasses import dataclass, field
-from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator
+from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator, cast
 
 T = TypeVar("T")
 U = TypeVar("U")
@@ -84,6 +84,17 @@ def word_wrap(x, wrap=80):
   i = 0
   while len(ansistrip(x[:i])) < wrap and i < len(x): i += 1
   return x[:i] + "\n" + word_wrap(x[i:], wrap)
+def pad_bytes(b:bytes, align:int) -> bytes: return b + b'\x00' * ((align - (len(b) % align)) % align)
+
+@functools.cache
+def canonicalize_strides(shape:tuple[T, ...], strides:tuple[T, ...]) -> tuple[T, ...]:
+  return tuple(cast(T, 0) if s == 1 else st for s, st in zip(shape, strides))
+
+@functools.cache
+def strides_for_shape(shape:tuple[T, ...]) -> tuple[T, ...]:
+  if not shape: return ()
+  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
@@ -154,7 +165,7 @@ ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), Conte
 FUSE_ATTENTION = ContextVar("FUSE_ATTENTION", 0)
 EMULATE = ContextVar("EMULATE", "")
 CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
-CPU_LLVM, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("AMD_LLVM", 1)
+CPU_LLVM, CPU_LVP, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0), ContextVar("AMD_LLVM", 1)
 VIZ = PROFILE = ContextVar("VIZ", 0)
 SPEC = ContextVar("SPEC", 0)
 # TODO: disable by default due to speed

tinygrad/nn/state.py

@@ -1,10 +1,9 @@
 import json, pathlib, zipfile, pickle, tarfile, struct, functools, io
 from collections import OrderedDict
-from typing import Any, Callable, BinaryIO, Iterable
+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, CI, unwrap, GlobalCounters, tqdm, round_up, T
-from tinygrad.shape.view import strides_for_shape
+from tinygrad.helpers import prod, argsort, DEBUG, Timing, CI, unwrap, GlobalCounters, tqdm, round_up, T, strides_for_shape
 
 class TensorIO(io.RawIOBase, BinaryIO):
   def __init__(self, t: Tensor):
@@ -238,11 +237,18 @@ def torch_load(t:Tensor) -> dict[str, Tensor]:
 
   if passthrough_reset(zipfile.is_zipfile(fobj)): # NOTE: passthrough_reset required to support python < 3.14
     myzip = zipfile.ZipFile(fobj, 'r')
-    base_name = myzip.namelist()[0].split('/', 1)[0]
-    for n in myzip.namelist():
-      if n.startswith(f'{base_name}/data/'):
-        with myzip.open(n) as myfile:
-          offsets[n.split("/")[-1]] = myfile._orig_compress_start # type: ignore
+    base_name = None
+    header_offsets = {}
+    for zi in myzip.filelist:
+      if base_name is None: base_name = zi.filename.split('/', 1)[0]
+      if zi.filename.startswith(f'{base_name}/data/'): header_offsets[zi.filename.split("/")[-1]] = zi.header_offset
+    # sadly there's no way to get the start of the file in the zip without reading the header
+    # at least here we read them in parallel
+    header_contents = [t[v+26:v+30].bitcast(dtypes.uint16).to('CPU') for v in header_offsets.values()]
+    Tensor.realize(*header_contents)
+    for (n,o),c in zip(header_offsets.items(), header_contents):
+      # header_offset + sizeFileHeader + File name length + Extra field length : https://en.wikipedia.org/wiki/ZIP_(file_format)
+      offsets[n] = o+30+sum(cast(list[int], c.tolist()))
     with myzip.open(f'{base_name}/data.pkl') as myfile:
       return TorchPickle(myfile).load()
   elif passthrough_reset(tarfile.is_tarfile(fobj)): # NOTE: passthrough_reset required to support python < 3.11

tinygrad/renderer/nir.py

@@ -0,0 +1,237 @@
+from typing import Callable, cast
+from tinygrad.dtype import AddrSpace, DType, PtrDType, dtypes
+from tinygrad.helpers import DEBUG, OSX, unwrap
+from tinygrad.renderer import Renderer
+from tinygrad.renderer.cstyle import CUDARenderer
+from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat
+import tinygrad.runtime.autogen.mesa as mesa
+import base64, ctypes, ctypes.util, struct, functools, inspect
+
+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))
+
+# this is a ridiculous hack, but I can't find a better way to grab the glsl_type objects
+glsl_base = {**{d:g(f"GLSL_TYPE_{'U' if d in dtypes.uints else ''}INT{d.itemsize*8 if d.itemsize != 4 else ''}") for d in dtypes.ints},
+             **{getattr(dtypes,d):g(f"GLSL_TYPE_{d.upper()}") for d in ['double', 'float', 'float16']}, dtypes.bool: mesa.GLSL_TYPE_UINT8}
+def glsl_type(t:DType) -> mesa.struct_glsl_type:
+  if isinstance(t, PtrDType): return mesa.glsl_array_type(glsl_type(t.base), t.size, 0).contents
+  return mesa.glsl_get_base_glsl_type(mesa.glsl_type(base_type=glsl_base[t])).contents
+
+# alu ops, aop[<dtype>][<op>]
+u_aop = { Ops.ADD: "iadd", Ops.MUL: "imul", Ops.IDIV: "udiv", Ops.MOD: "umod", Ops.CMPLT: "ult", Ops.CMPNE: "ine", Ops.CMPEQ: "ieq", Ops.OR: "ior",
+          Ops.AND: "iand", Ops.XOR: "ixor", Ops.WHERE: "bcsel", Ops.MAX: "umax"}
+s_aop = {**u_aop, Ops.CMPLT: "ilt", Ops.IDIV: "idiv", Ops.MOD: "irem", Ops.MAX: "imax"}
+f_aop = { Ops.ADD: "fadd", Ops.MUL: "fmul", Ops.CMPLT: "flt", Ops.CMPNE: "fneu", Ops.CMPEQ: "feq", Ops.FDIV: "fdiv", Ops.RECIP: "frcp",
+          Ops.MAX: "fmax", Ops.TRUNC: "ftrunc", Ops.SIN: "fsin", Ops.EXP2: "fexp2", Ops.LOG2: "flog2"}
+aop = {**{x:u_aop for x in (dtypes.bool,)+dtypes.uints}, **{x:s_aop for x in dtypes.sints}, **{x:f_aop for x in dtypes.floats}}
+
+def c(t:DType, u:bool=True) -> str: return "u" if t in dtypes.uints and u else ("i" if t in dtypes.ints else ("f" if t in dtypes.floats else "b"))
+def ncast(b:mesa.nir_builder, src:mesa.nir_def, it:DType, ot:DType) -> mesa.nir_def:
+  if isinstance(it, PtrDType) and ot == dtypes.long: return src
+  if ot == dtypes.bool: return nalu(b, c(it, False)+'ne'+('u' if c(it) == 'f' else ''), src, nimm(b, 0, it))
+  return nalu(b, f"{c(it)}2{c(it) if it in dtypes.ints and ot in dtypes.ints else c(ot, ot == dtypes.bool)}{ot.itemsize*8}", src)
+
+def nif(b:mesa.nir_builder, cond:mesa.nir_def, then_fn:Callable, else_fn:Callable):
+  nif = mesa.nir_push_if(b, cond)
+  t = then_fn()
+  mesa.nir_push_else(b, nif)
+  e = else_fn()
+  mesa.nir_pop_if(b, nif)
+  return t, e
+
+def nalu(b:mesa.nir_builder, op:str, *srcs:mesa.nir_def) -> mesa.nir_def: return g(f"nir_build_alu{len(srcs)}")(b, g(f"nir_op_{op}"), *srcs).contents
+
+def nir_instr(nc=1, bs=lambda: None, intrins=None, srcs=None, has_def=True, df=None, also=lambda: None, **contents):
+  def dec(f:Callable):
+    @functools.wraps(f)
+    def wrapper(*args, **kwargs) -> mesa.nir_def:
+      (ba:=inspect.signature(f).bind(*args, **kwargs)).apply_defaults()
+      def go(g): return g(**{nm: ba.arguments[nm] for nm in inspect.signature(g).parameters}) if callable(g) else g
+
+      instr = f(*args, **kwargs)
+      if has_def: mesa.nir_def_init(instr.contents.instr, getattr(instr.contents, "def"), go(nc), go(bs))
+      for k, v in go(intrins or {}).items():
+        idx = mesa.nir_intrinsic_infos[instr.contents.intrinsic].index_map[g(f"NIR_INTRINSIC_{k}")]
+        assert idx > 0
+        instr.contents.const_index[idx - 1] = go(v)
+      for i, src in enumerate(go(srcs or [])): ctypes.cast(instr.contents.src, ctypes.POINTER(mesa.nir_src))[i] = go(src)
+      for k,v in {k:vcomp for k,v in contents.items() if (vcomp:=go(v)) is not None}.items(): setattr(instr.contents, k, go(v))
+      mesa.nir_builder_instr_insert(ba.arguments['b'], instr.contents.instr)
+      go(also)
+      return getattr(instr.contents, "def") if has_def else (mesa.nir_def() if df is None else go(df))
+    return wrapper
+  return dec
+
+@nir_instr(nc=1, bs=lambda src: src.bit_size, exact=lambda b:b.exact, fp_fast_math=lambda b:b.fp_fast_math)
+def nchannel(b:mesa.nir_builder, src:mesa.nir_def, c:int):
+  alu_src = mesa.nir_alu_src(src=nsrc(src))
+  alu_src.swizzle[0] = c
+  mov = mesa.nir_alu_instr_create(b.shader, mesa.nir_op_mov)
+  ctypes.cast(mov.contents.src, ctypes.POINTER(mesa.nir_alu_src))[0] = alu_src
+  return mov
+
+@nir_instr(nc=1, bs=lambda dtype: 1 if dtype == dtypes.bool else dtype.itemsize * 8)
+def nimm(b:mesa.nir_builder, x, dtype:DType) -> mesa.nir_def:
+  instr = mesa.nir_load_const_instr_create(b.shader, 1, 1 if dtype == dtypes.bool else dtype.itemsize * 8)
+  struct.pack_into(unwrap(dtype.fmt), (ctypes.c_ubyte * dtype.itemsize).from_address(ctypes.addressof(instr.contents.value)), 0, x)
+  return instr
+
+deref_var = nir_instr(nc=1, bs=32, modes=lambda var:var.data.mode, type=lambda var:var.type, var=lambda var:ctypes.pointer(var))( # pylint: disable=W0108
+  lambda b, var: mesa.nir_deref_instr_create(b.shader, mesa.nir_deref_type_var))
+
+def iointr(space): return {"ALIGN_MUL":lambda dtype:dtype.itemsize} if space != AddrSpace.REG else {}
+def scope(space): return 'global' if space == AddrSpace.GLOBAL else ('shared' if space == AddrSpace.LOCAL else 'deref')
+nstore = nir_instr(has_def=False, df=lambda addr:addr, intrins=lambda space,val: {"WRITE_MASK":(1<<val.num_components)-1, **iointr(space)},
+  num_components=lambda val:val.num_components, srcs=lambda space, addr, val: [nsrc(val), nsrc(addr)][::1 if space != AddrSpace.REG else -1])(
+    lambda b, space, addr, val, dtype: mesa.nir_intrinsic_instr_create(b.shader, g(f"nir_intrinsic_store_{scope(space)}")))
+nload = nir_instr(nc=lambda dtype:dtype.count, bs=lambda dtype:dtype.itemsize*8//dtype.count, num_components=lambda dtype:dtype.count,
+  intrins=lambda space:{**({"ACCESS":mesa.ACCESS_CAN_REORDER} if space==AddrSpace.GLOBAL else {}), **iointr(space)}, srcs=lambda addr: [nsrc(addr)])(
+    lambda b, space, addr, dtype: mesa.nir_intrinsic_instr_create(b.shader, g(f"nir_intrinsic_load_{scope(space)}")))
+
+ngid = nir_instr(nc=3, bs=32)(lambda b: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_workgroup_id))
+nlid = nir_instr(nc=3, bs=32)(lambda b: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_local_invocation_id))
+
+nbarrier = nir_instr(has_def=False, intrins={"EXECUTION_SCOPE":mesa.SCOPE_WORKGROUP})(
+  lambda b: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_barrier))
+
+@nir_instr(has_def=False, target=lambda tgt:tgt and ctypes.pointer(tgt), condition=lambda cond:cond and nsrc(cond),
+           else_target=lambda else_tgt: else_tgt and ctypes.pointer(else_tgt))
+def njump(b:mesa.nir_builder, typ, tgt=None, cond=None, else_tgt=None): return mesa.nir_jump_instr_create(b.shader, typ)
+
+def if_phi(b:mesa.nir_builder, cond, then_fn, else_fn): return mesa.nir_if_phi(b, *nif(b, cond, then_fn, else_fn)).contents
+
+def nidx(b:mesa.nir_builder, buf, off, dtype, gate=None) -> mesa.nir_def:
+  @nir_instr(nc=1, bs=32, modes=lambda buf: buf.data.mode, type=lambda buf: mesa.glsl_get_array_element(buf.type))
+  def reg(b, buf):
+    deref = mesa.nir_deref_instr_create(b.shader, mesa.nir_deref_type_array)
+    deref.contents.parent, deref.contents.arr.index = nsrc(deref_var(b, buf)), nsrc(off)
+    return deref
+  f = (functools.partial(reg, b, buf) if dtype.addrspace == AddrSpace.REG else
+       lambda: nalu(b, "iadd", buf, nalu(b, "imul", off, nimm(b, dtype.itemsize, dtypes.long))))
+  return if_phi(b, gate, f, lambda: buf) if gate is not None else f()
+
+class NIRRenderer(Renderer):
+  suffix = "NAK"
+  global_max, local_max, shared_max = CUDARenderer.global_max, CUDARenderer.local_max, CUDARenderer.shared_max
+  code_for_op = {**{k:lambda:None for k in u_aop.keys()}, **{k:lambda:None for k in s_aop.keys()}, **{k:lambda:None for k in f_aop.keys()}}
+
+  extra_matcher = PatternMatcher([
+    # handle negative unsigned CONST
+    (UPat.cvar("x", dtypes.uints), lambda x: UOp(Ops.CONST, dtype=x.dtype, arg=x.dtype.max+x.arg+1) if x.arg < 0 else None),
+    # from ptx
+    (UPat.var('x', dtype=dtypes.bool)<UPat.var('y'), lambda x,y: (x^True)&y),
+    # load/store bool -> uint8
+    (UPat(Ops.LOAD, dtypes.bool, name="x"),
+     lambda x: x.replace(dtype=dtypes.uint8, src=x.src[0:1]+((x.src[1].cast(dtypes.uint8),) if len(x.src)>=2 else ())+x.src[2:]).cast(dtypes.bool)),
+    (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
+     lambda x: x.replace(src=x.src[0:1] + (x.src[1].cast(dtypes.uint8),) + x.src[2:])),
+    # load/store use pointer arithmetic, and the cast does nothing
+    (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off")), allow_any_len=True, name="x"),
+     lambda x,buf,off: x.replace(src=(buf,off.cast(dtypes.long))+x.src[2:]) if buf.dtype.addrspace != AddrSpace.REG and off.op != Ops.CAST else None),
+    (UPat(Ops.CAST, name="x"), lambda x: x.src[0] if isinstance(x.dtype, PtrDType) or x.src[0].dtype == dtypes.void else None),
+  ])
+
+  def_rewrite = PatternMatcher([
+    (UPat(Ops.CONST, name="x"), lambda ctx,x: nimm(ctx.b, x.arg, x.dtype)),
+    (UPat(Ops.DEFINE_GLOBAL, name="x"), lambda ctx,x: ctx.param(ctx.b, x.dtype, 8)),
+    (UPat(Ops.DEFINE_VAR, name="x"), lambda ctx,x: ctx.param(ctx.b, x.dtype, 4)),
+    (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, ngid(ctx.b) if x.arg[0] == 'g' else nlid(ctx.b), int(x.arg[-1]))),
+    (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"),UPat.var("off")), allow_any_len=True), UPat.var("val")), allow_any_len=True, name="x"),
+     lambda ctx,x,buf,off,val: nstore(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), ctx.r[val], val.dtype)),
+    (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"), UPat.var("gate"))), UPat.var("alt")), allow_any_len=True, name="x"),
+     lambda ctx,x,buf,off,alt,gate: if_phi(ctx.b, ctx.r[gate],
+      lambda: nload(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype, ctx.r[gate]), x.dtype), lambda: ctx.r[alt])),
+    (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))),), allow_any_len=True, name="x"),
+     lambda ctx,x,buf,off: nload(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), x.dtype)),
+    (UPat(Ops.VECTORIZE, name="x"), lambda ctx,x: nalu(ctx.b, f"vec{x.dtype.count}", *[ctx.r[src] for src in x.src])),
+    (UPat(GroupOp.ALU, name="x"), lambda ctx,x: nalu(ctx.b, aop[x.src[0].dtype.scalar()][x.op], *[ctx.r[src] for src in x.src])),
+    (UPat(Ops.CAST, name="x"), lambda ctx,x: ncast(ctx.b, ctx.r[x.src[0]], x.src[0].dtype, x.dtype)),
+    (UPat(Ops.BITCAST, src=(UPat.var("a"),), allow_any_len=True), lambda ctx,a: ctx.r[a]),
+    (UPat(Ops.GEP, src=(UPat.var("a"),), name="x"), lambda ctx,x,a: nchannel(ctx.b, ctx.r[a], x.arg[0])),
+    (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x:mesa.nir_local_variable_create(ctx.b.impl, glsl_type(x.dtype), f"acc{x.arg[0]}".encode()).contents),
+    (UPat(Ops.BARRIER), lambda ctx: nbarrier(ctx.b)),
+    (UPat(Ops.IF, name="x"), lambda ctx,x: mesa.nir_push_if(ctx.b, ctx.r[x.src[0]])),
+    (UPat(Ops.ENDIF, name="x"), lambda ctx,x: (lambda _: mesa.nir_def())(mesa.nir_pop_if(ctx.b, ctx.r[x.src[0]])))
+  ])
+
+  def __init__(self): mesa.glsl_type_singleton_init_or_ref()
+
+  def __del__(self):
+    try: mesa.glsl_type_singleton_decref()
+    except FileNotFoundError: pass
+
+  @property
+  def nir_options(self): raise NotImplementedError("needs nir_options")
+  def param(self, b:mesa.nir_builder, dtype:DType, sz:int) -> mesa.nir_def: raise NotImplementedError("needs param")
+  def prerender(self, uops:list[UOp]):
+    self.b = mesa.nir_builder_init_simple_shader(mesa.MESA_SHADER_COMPUTE, mesa.nir_shader_compiler_options.from_buffer_copy(self.nir_options), None)
+
+  def render(self, uops:list[UOp]):
+    self.prerender(uops)
+    for u in [u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]: self.b.shader.contents.info.workgroup_size[int(u.arg[-1])] = u.src[0].arg
+    self.r, self.param_idx, ranges = {}, 0, []
+
+    for u in uops:
+      if u.op == Ops.NOOP or u.op == Ops.INDEX: pass
+      elif u.op == Ops.SINK:
+        if u.arg is not None: self.b.shader.contents.info.name = mesa.char_pointer_cast(u.arg.function_name)
+      elif u.op == Ops.DEFINE_LOCAL:
+        self.r[u] = nimm(self.b, self.b.shader.contents.info.shared_size, dtypes.long)
+        self.b.shader.contents.info.shared_size += u.dtype.nbytes()
+      elif u.op == Ops.RANGE:
+        ranges.append(i:=deref_var(self.b, mesa.nir_local_variable_create(self.b.impl, glsl_type(u.dtype), f"idx{u.arg[0]}".encode()).contents))
+        nstore(self.b, AddrSpace.REG, i, nimm(self.b, 0, u.dtype), u.dtype)
+        mesa.nir_push_loop(self.b)
+        self.r[u] = nload(self.b, AddrSpace.REG, i, u.dtype)
+      elif u.op == Ops.ENDRANGE:
+        nif(self.b, nalu(self.b, "ilt", x:=nalu(self.b, "iadd", self.r[u.src[0]], nimm(self.b, 1, u.src[0].dtype)), self.r[u.src[0].src[0]]),
+            functools.partial(nstore, self.b, AddrSpace.REG, ranges.pop(), x, u.src[0].dtype), lambda: njump(self.b, mesa.nir_jump_break))
+        mesa.nir_pop_loop(self.b, None)
+      else:
+        if (d:=self.def_rewrite.rewrite(u, ctx=self)) is None: raise RuntimeError(f"failed to render {u.op} srcs {[x.dtype for x in u.src]}")
+        self.r[u] = cast(mesa.nir_def, d)
+
+    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"))
+    mesa.nir_serialize(blob:=mesa.struct_blob(), self.b.shader, False)
+    ret = base64.b64encode(ctypes.string_at(blob.data, blob.size)).decode()
+
+    mesa.ralloc_free(self.b.shader)
+    ctypes.CDLL(None).free(blob.data)
+    del self.b, self.r
+
+    return ret
+
+class NAKRenderer(NIRRenderer):
+  device = "NV"
+  def __init__(self, dev=None, nir_options=None):
+    self.dev, self._nir_options = dev, nir_options
+    super().__init__()
+
+  def __reduce__(self): return NAKRenderer, (None, self.nir_options,)
+
+  @property
+  def nir_options(self):
+    if self._nir_options is None: self._nir_options = self.dev.compiler.nir_options
+    return self._nir_options
+
+  param = nir_instr(nc=1, num_components=1, bs=lambda sz:sz*8, also=lambda self,sz: setattr(self, "param_idx", self.param_idx + sz),
+    intrins={"ALIGN_MUL":lambda sz:sz}, srcs=lambda self,b: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, self.param_idx, dtypes.int))])(
+       lambda self, b, dtype, sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_ldc_nv))
+
+class LVPRenderer(NIRRenderer):
+  device = "CPU"
+  has_local = False
+  has_shared = False
+  global_max = (1, 0, 0)
+  nir_options = mesa.lvp_nir_options
+
+  param = nir_instr(nc=1, bs=lambda sz: sz * 8, num_components=1, intrins={"ALIGN_MUL":lambda sz: sz, "RANGE":lambda self: self.param_sz},
+    srcs=lambda b, self: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, self.param_idx, dtypes.int))], also=lambda self, sz:
+    setattr(self, "param_idx", self.param_idx+sz))(lambda self, b, dtype, sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_ubo))
+
+  def prerender(self, uops:list[UOp]):
+    super().prerender(uops)
+    self.param_sz = sum([8 if u.op == Ops.DEFINE_GLOBAL else u.dtype.itemsize for u in uops if u.op in (Ops.DEFINE_GLOBAL, Ops.DEFINE_VAR)])
+

tinygrad/runtime/autogen/llvm.py

@@ -6,7 +6,7 @@
 # POINTER_SIZE is: 8
 # LONGDOUBLE_SIZE is: 16
 #
-import ctypes, tinygrad.runtime.support.llvm as llvm_support
+import ctypes, tinygrad.runtime.support.llvm as llvm_support, tinygrad.helpers as helpers
 
 
 class AsDictMixin:
@@ -146,7 +146,7 @@ class FunctionFactoryStub:
 # You can either re-run clan2py with -l /path/to/library.so
 # Or manually fix this by comment the ctypes.CDLL loading
 _libraries = {}
-_libraries['llvm'] = ctypes.CDLL(llvm_support.LLVM_PATH) #  ctypes.CDLL('llvm')
+_libraries['llvm'] = ctypes.CDLL(llvm_support.LLVM_PATH, ctypes.RTLD_GLOBAL if helpers.OSX else ctypes.DEFAULT_MODE) #  ctypes.CDLL('llvm')
 c_int128 = ctypes.c_ubyte*16
 c_uint128 = c_int128
 void = None

tinygrad/runtime/ops_cpu.py

@@ -1,11 +1,15 @@
 from __future__ import annotations
 import platform, sys, ctypes, functools, time, mmap, threading, queue
-from tinygrad.helpers import from_mv, to_mv, OSX, WIN, mv_address, wait_cond, cpu_profile, suppress_finalizing, unwrap
-from tinygrad.device import BufferSpec, DMACPURef
+from tinygrad.helpers import from_mv, to_mv, OSX, WIN, mv_address, wait_cond, cpu_profile, suppress_finalizing, unwrap, data64_le
+from tinygrad.device import BufferSpec, DMACPURef, CompilerPairT
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocatorBase, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface
+from tinygrad.runtime.support.hcq import CLikeArgsState
 from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.renderer.llvmir import LLVMRenderer
+from tinygrad.renderer.nir import LVPRenderer
 from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangJITCompiler
+from tinygrad.runtime.support.compiler_mesa import LVPCompiler
+from tinygrad.runtime.support.elf import jit_loader
 from tinygrad.uop.ops import sint
 
 class CPUSignal(HCQSignal):
@@ -46,12 +50,18 @@ class CPUComputeQueue(HWQueue):
 
   def memory_barrier(self): return self
   def exec(self, prg:CPUProgram, args_state:HCQArgsState, global_size, local_size):
+    if isinstance(args_state, LVPArgsState):
+      self.bind_args_state(args_state)
+      return self.cmd(self._exec, prg, 1, args_state.buf.va_addr)
     return self.cmd(self._exec, prg, len(args_state.bufs), *[x.va_addr for x in args_state.bufs], *args_state.vals, threads=(global_size or (1,))[0])
   def wait(self, signal, value=0): return self.cmd(self._wait, signal.value_addr, value)
   def timestamp(self, signal): return self.cmd(self._timestamp, signal.timestamp_addr)
   def signal(self, signal, value:sint=0): return self.cmd(self._signal, signal.value_addr, value)
   def _submit(self, dev): dev.tasks.put(self._q[:])
 
+class LVPArgsState(CLikeArgsState):
+  def __init__(self, buf, prg, bufs, vals=()): super().__init__(buf, prg, bufs, vals, [*data64_le(buf.va_addr + 12), (len(bufs) + len(vals)) * 2])
+
 # NOTE: MAP_JIT is added to mmap module in python 3.13
 MAP_JIT = 0x0800
 
@@ -61,6 +71,7 @@ class CPUProgram(HCQProgram):
   except OSError: pass
 
   def __init__(self, dev, name:str, lib:bytes):
+    LVP = isinstance(dev.compiler, LVPCompiler)
     if sys.platform == "win32": # mypy doesn't understand when WIN is used here
       PAGE_EXECUTE_READWRITE, MEM_COMMIT, MEM_RESERVE = 0x40, 0x1000, 0x2000
       ctypes.windll.kernel32.VirtualAlloc.restype = ctypes.c_void_p
@@ -76,6 +87,7 @@ class CPUProgram(HCQProgram):
       self.mem = mmap.mmap(-1, len(lib), mmap.MAP_ANON|mmap.MAP_PRIVATE|(MAP_JIT if OSX else 0), mmap.PROT_READ|mmap.PROT_WRITE|mmap.PROT_EXEC)
 
       if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(False)
+      if LVP: lib = jit_loader(lib, base=ctypes.addressof(ctypes.c_void_p.from_buffer(self.mem)), link_libs=['m'])
       self.mem.write(lib)
       if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(True)
 
@@ -92,7 +104,7 @@ class CPUProgram(HCQProgram):
 
       self.fxn = ctypes.CFUNCTYPE(None)(mv_address(self.mem))
 
-    super().__init__(HCQArgsState, dev, name, kernargs_alloc_size=0)
+    super().__init__(LVPArgsState if LVP else HCQArgsState, dev, name, kernargs_alloc_size=12+256 if LVP else 0)
 
   @suppress_finalizing
   def __del__(self):
@@ -123,5 +135,5 @@ class CPUDevice(HCQCompiled):
   def __init__(self, device:str=""):
     self.tasks:queue.Queue = queue.Queue()
     CPUWorker(self, self.tasks, thread_id=0).start()
-    compilers = [(ClangRenderer, ClangJITCompiler), (LLVMRenderer, CPULLVMCompiler)]
+    compilers:list[CompilerPairT] = [(ClangRenderer, ClangJITCompiler), (LLVMRenderer, CPULLVMCompiler), (LVPRenderer, LVPCompiler)]
     super().__init__(device, CPUAllocator(self), compilers, functools.partial(CPUProgram, self), CPUSignal, CPUComputeQueue)

tinygrad/runtime/ops_null.py

@@ -1,9 +1,11 @@
 import functools
+from typing import cast
 from tinygrad.device import Compiled, Compiler, Allocator
 from tinygrad.engine.jit import MultiGraphRunner
-from tinygrad.renderer.cstyle import CStyleLanguage
+from tinygrad.renderer.cstyle import Renderer, CStyleLanguage
+from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.uop.ops import Ops
-from tinygrad.helpers import cpu_profile
+from tinygrad.helpers import cpu_profile, EMULATE
 
 class NullRenderer(CStyleLanguage):
   device = "NULL"
@@ -29,5 +31,11 @@ class NullGraph(MultiGraphRunner):
   def __call__(self, input_rawbuffers, var_vals, wait=False) -> float|None: return 1e-3
 
 class NullDevice(Compiled):
-  def __init__(self, device:str): super().__init__(device, NullAllocator(self), [(NullRenderer, Compiler)], functools.partial(NullProgram, device),
-                                                   NullGraph)
+  def __init__(self, device:str):
+    renderer:functools.partial|type[Renderer]
+    match cast(str, EMULATE.value):
+      case "AMD": renderer = functools.partial(AMDLLVMRenderer, "gfx1100")
+      case "AMD_RDNA4": renderer = functools.partial(AMDLLVMRenderer, "gfx1201")
+      case "": renderer = NullRenderer
+      case _: raise RuntimeError(f"can't EMULATE device: {EMULATE.value}")
+    super().__init__(device, NullAllocator(self), [(renderer, Compiler)], functools.partial(NullProgram, device), NullGraph)

tinygrad/runtime/ops_nv.py

@@ -11,10 +11,12 @@ from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, pr
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import NVRenderer
 from tinygrad.runtime.support.compiler_cuda import CUDACompiler, PTXCompiler, NVPTXCompiler, NVCompiler
-from tinygrad.runtime.autogen import nv_gpu, pci
+from tinygrad.runtime.support.compiler_mesa import NAKCompiler
+from tinygrad.runtime.autogen import nv_gpu, pci, mesa
 from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.runtime.support.nv.nvdev import NVDev, NVMemoryManager
 from tinygrad.runtime.support.system import System, PCIIfaceBase, MAP_FIXED
+from tinygrad.renderer.nir import NAKRenderer
 if getenv("IOCTL"): import extra.nv_gpu_driver.nv_ioctl # noqa: F401 # pylint: disable=unused-import
 
 def get_error_str(status): return f"{status}: {nv_gpu.nv_status_codes.get(status, 'Unknown error')}"
@@ -133,7 +135,7 @@ class NVComputeQueue(NVCommandQueue):
     qmd_buf.cpu_view().view(size=prg.qmd.mv.nbytes, fmt='B')[:] = prg.qmd.mv
     assert qmd_buf.va_addr < (1 << 40), f"large qmd addr {qmd_buf.va_addr:x}"
 
-    qmd = QMD(dev=prg.dev, addr=cast(int, qmd_buf.va_addr)) # Save qmd for later update
+    qmd = QMD(dev=prg.dev, addr=qmd_buf.cpu_view().addr) # Save qmd for later update
 
     self.bind_sints_to_mem(*global_size, mem=qmd_buf.cpu_view(), fmt='I', offset=qmd.field_offset('cta_raster_width' if qmd.ver<4 else 'grid_width'))
     self.bind_sints_to_mem(*(local_size[:2]), mem=qmd_buf.cpu_view(), fmt='H', offset=qmd.field_offset('cta_thread_dimension0'))
@@ -185,68 +187,69 @@ class NVCopyQueue(NVCommandQueue):
 
 class NVArgsState(CLikeArgsState):
   def __init__(self, buf:HCQBuffer, prg:NVProgram, bufs:tuple[HCQBuffer, ...], vals:tuple[int, ...]=()):
-    if MOCKGPU: prg.constbuffer_0[80:82] = [len(bufs), len(vals)]
-    super().__init__(buf, prg, bufs, vals=vals, prefix=prg.constbuffer_0)
+    if MOCKGPU: prg.cbuf_0[80:82] = [len(bufs), len(vals)]
+    super().__init__(buf, prg, bufs, vals=vals, prefix=prg.cbuf_0 or None)
 
 class NVProgram(HCQProgram):
   def __init__(self, dev:NVDevice, name:str, lib:bytes):
     self.dev, self.name, self.lib = dev, name, lib
-
-    # For MOCKGPU, the lib is PTX code, so some values are emulated.
-    cbuf0_size = 0 if not MOCKGPU else 0x160
-
-    if MOCKGPU: image, sections, relocs = memoryview(bytearray(lib) + b'\x00' * (4 - len(lib)%4)).cast("I"), [], [] # type: ignore
-    else: image, sections, relocs = elf_loader(self.lib, force_section_align=128)
-
-    # NOTE: Ensure at least 4KB of space after the program to mitigate prefetch memory faults.
-    self.lib_gpu = self.dev.allocator.alloc(round_up(image.nbytes, 0x1000) + 0x1000, buf_spec:=BufferSpec(cpu_access=True))
-
-    self.prog_addr, self.prog_sz, self.regs_usage, self.shmem_usage, self.lcmem_usage = self.lib_gpu.va_addr, image.nbytes, 0, 0x400, 0
     self.constbufs: dict[int, tuple[int, int]] = {0: (0, 0x160)} # dict[constbuf index, tuple[va_addr, size]]
-    for sh in sections:
-      if sh.name == f".nv.shared.{self.name}": self.shmem_usage = round_up(0x400 + sh.header.sh_size, 128)
-      if sh.name == f".text.{self.name}": self.prog_addr, self.prog_sz = self.lib_gpu.va_addr+sh.header.sh_addr, sh.header.sh_size
-      elif m:=re.match(r'\.nv\.constant(\d+)', sh.name): self.constbufs[int(m.group(1))] = (self.lib_gpu.va_addr+sh.header.sh_addr, sh.header.sh_size)
-      elif sh.name.startswith(".nv.info"):
-        for typ, param, data in self._parse_elf_info(sh):
-          if sh.name == f".nv.info.{name}" and param == 0xa: cbuf0_size = struct.unpack_from("IH", data)[1] # EIATTR_PARAM_CBANK
-          elif sh.name == ".nv.info" and param == 0x12: self.lcmem_usage = struct.unpack_from("II", data)[1] + 0x240 # EIATTR_MIN_STACK_SIZE
-          elif sh.name == ".nv.info" and param == 0x2f: self.regs_usage = struct.unpack_from("II", data)[1] # EIATTR_REGCOUNT
+
+    if (NAK:=isinstance(dev.compiler, NAKCompiler)):
+      image, self.cbuf_0 = memoryview(bytearray(lib[ctypes.sizeof(info:=mesa.struct_nak_shader_info.from_buffer_copy(lib)):])), []
+      self.regs_usage, self.shmem_usage, self.lcmem_usage = info.num_gprs, round_up(info.cs.smem_size, 128), round_up(info.slm_size, 16)
+    elif MOCKGPU: image, sections, relocs = memoryview(bytearray(lib) + b'\x00' * (4 - len(lib)%4)).cast("I"), [], [] # type: ignore
+    else: image, sections, relocs = elf_loader(self.lib, force_section_align=128)
+    # NOTE: Ensure at least 4KB of space after the program to mitigate prefetch memory faults.
+    self.lib_gpu = self.dev.allocator.alloc(round_up((prog_sz:=image.nbytes), 0x1000) + 0x1000, buf_spec:=BufferSpec(nolru=True))
+    prog_addr = self.lib_gpu.va_addr
+    if not NAK:
+      # For MOCKGPU, the lib is PTX code, so some values are emulated.
+      self.regs_usage, self.shmem_usage, self.lcmem_usage, cbuf0_size = 0, 0x400, 0x240, 0 if not MOCKGPU else 0x160
+      for sh in sections: # pylint: disable=possibly-used-before-assignment
+        if sh.name == f".nv.shared.{self.name}": self.shmem_usage = round_up(0x400 + sh.header.sh_size, 128)
+        if sh.name == f".text.{self.name}": prog_addr, prog_sz = self.lib_gpu.va_addr+sh.header.sh_addr, sh.header.sh_size
+        elif m:=re.match(r'\.nv\.constant(\d+)', sh.name):
+          self.constbufs[int(m.group(1))] = (self.lib_gpu.va_addr+sh.header.sh_addr, sh.header.sh_size)
+        elif sh.name.startswith(".nv.info"):
+          for typ, param, data in self._parse_elf_info(sh):
+            if sh.name == f".nv.info.{name}" and param == 0xa: cbuf0_size = struct.unpack_from("IH", data)[1] # EIATTR_PARAM_CBANK
+            elif sh.name == ".nv.info" and param == 0x12: self.lcmem_usage = struct.unpack_from("II", data)[1] + 0x240 # EIATTR_MIN_STACK_SIZE
+            elif sh.name == ".nv.info" and param == 0x2f: self.regs_usage = struct.unpack_from("II", data)[1] # EIATTR_REGCOUNT
+
+      # Apply relocs
+      for apply_image_offset, rel_sym_offset, typ, _ in relocs: # pylint: disable=possibly-used-before-assignment
+        # These types are CUDA-specific, applying them here
+        if typ == 2: image[apply_image_offset:apply_image_offset+8] = struct.pack('<Q', self.lib_gpu.va_addr + rel_sym_offset) # R_CUDA_64
+        elif typ == 0x38: image[apply_image_offset+4:apply_image_offset+8] = struct.pack('<I', (self.lib_gpu.va_addr + rel_sym_offset) & 0xffffffff)
+        elif typ == 0x39: image[apply_image_offset+4:apply_image_offset+8] = struct.pack('<I', (self.lib_gpu.va_addr + rel_sym_offset) >> 32)
+        else: raise RuntimeError(f"unknown NV reloc {typ}")
+
+      self.cbuf_0 = [0] * (cbuf0_size // 4)
 
     # Ensure device has enough local memory to run the program
     self.dev._ensure_has_local_memory(self.lcmem_usage)
-
-    # Apply relocs
-    for apply_image_offset, rel_sym_offset, typ, _ in relocs:
-      # These types are CUDA-specific, applying them here
-      if typ == 2: image[apply_image_offset:apply_image_offset+8] = struct.pack('<Q', self.lib_gpu.va_addr + rel_sym_offset) # R_CUDA_64
-      elif typ == 0x38: image[apply_image_offset+4:apply_image_offset+8] = struct.pack('<I', (self.lib_gpu.va_addr + rel_sym_offset) & 0xffffffff)
-      elif typ == 0x39: image[apply_image_offset+4:apply_image_offset+8] = struct.pack('<I', (self.lib_gpu.va_addr + rel_sym_offset) >> 32)
-      else: raise RuntimeError(f"unknown NV reloc {typ}")
-
-    ctypes.memmove(self.lib_gpu.va_addr, mv_address(image), image.nbytes)
-
-    self.constbuffer_0 = [0] * (cbuf0_size // 4)
+    self.dev.allocator._copyin(self.lib_gpu, image)
+    self.dev.synchronize()
 
     if dev.iface.compute_class >= nv_gpu.BLACKWELL_COMPUTE_A:
-      self.constbuffer_0[188:192], self.constbuffer_0[223] = [*data64_le(self.dev.shared_mem_window), *data64_le(self.dev.local_mem_window)], 0xfffdc0
-      qmd = {'qmd_major_version':5, 'qmd_type':nv_gpu.NVCEC0_QMDV05_00_QMD_TYPE_GRID_CTA, 'register_count':self.regs_usage,
-        'program_address_upper_shifted4':hi32(self.prog_addr>>4), 'program_address_lower_shifted4':lo32(self.prog_addr>>4),
-        'shared_memory_size_shifted7':self.shmem_usage>>7, 'shader_local_memory_high_size_shifted4':self.dev.slm_per_thread>>4}
+      if not NAK: self.cbuf_0[188:192], self.cbuf_0[223] = [*data64_le(self.dev.shared_mem_window), *data64_le(self.dev.local_mem_window)], 0xfffdc0
+      qmd = {'qmd_major_version':5, 'qmd_type':nv_gpu.NVCEC0_QMDV05_00_QMD_TYPE_GRID_CTA, 'program_address_upper_shifted4':hi32(prog_addr>>4),
+        'program_address_lower_shifted4':lo32(prog_addr>>4), 'register_count':self.regs_usage, 'shared_memory_size_shifted7':self.shmem_usage>>7,
+        'shader_local_memory_high_size_shifted4':self.lcmem_usage>>4 if NAK else self.dev.slm_per_thread>>4}
     else:
-      self.constbuffer_0[6:12] = [*data64_le(self.dev.shared_mem_window), *data64_le(self.dev.local_mem_window), *data64_le(0xfffdc0)]
-      qmd = {'qmd_major_version':3, 'sm_global_caching_enable':1, 'shader_local_memory_high_size':self.dev.slm_per_thread,
-        'program_address_upper':hi32(self.prog_addr), 'program_address_lower':lo32(self.prog_addr), 'shared_memory_size':self.shmem_usage,
-        'register_count_v':self.regs_usage}
+      if not NAK: self.cbuf_0[6:12] = [*data64_le(self.dev.shared_mem_window), *data64_le(self.dev.local_mem_window), *data64_le(0xfffdc0)]
+      qmd = {'qmd_major_version':3, 'sm_global_caching_enable':1, 'program_address_upper':hi32(prog_addr), 'program_address_lower':lo32(prog_addr),
+        'shared_memory_size':self.shmem_usage, 'register_count_v':self.regs_usage,
+        **({'shader_local_memory_low_size':self.lcmem_usage} if NAK else {'shader_local_memory_high_size':self.dev.slm_per_thread})}
 
     smem_cfg = min(shmem_conf * 1024 for shmem_conf in [32, 64, 100] if shmem_conf * 1024 >= self.shmem_usage) // 4096 + 1
 
     self.qmd:QMD = QMD(dev, **qmd, qmd_group_id=0x3f, invalidate_texture_header_cache=1, invalidate_texture_sampler_cache=1,
       invalidate_texture_data_cache=1, invalidate_shader_data_cache=1, api_visible_call_limit=1, sampler_index=1, barrier_count=1,
-      cwd_membar_type=nv_gpu.NVC6C0_QMDV03_00_CWD_MEMBAR_TYPE_L1_SYSMEMBAR, constant_buffer_invalidate_0=1,
-      min_sm_config_shared_mem_size=smem_cfg, target_sm_config_shared_mem_size=smem_cfg, max_sm_config_shared_mem_size=0x1a,
-      program_prefetch_size=min(self.prog_sz>>8, 0x1ff), sass_version=dev.sass_version,
-      program_prefetch_addr_upper_shifted=self.prog_addr>>40, program_prefetch_addr_lower_shifted=self.prog_addr>>8)
+      cwd_membar_type=nv_gpu.NVC6C0_QMDV03_00_CWD_MEMBAR_TYPE_L1_SYSMEMBAR, constant_buffer_invalidate_0=1, min_sm_config_shared_mem_size=smem_cfg,
+      target_sm_config_shared_mem_size=smem_cfg, max_sm_config_shared_mem_size=0x1a, program_prefetch_size=min(prog_sz>>8, 0x1ff),
+      sass_version=dev.sass_version, program_prefetch_addr_upper_shifted=prog_addr>>40, program_prefetch_addr_lower_shifted=prog_addr>>8)
 
     for i,(addr,sz) in self.constbufs.items():
       self.qmd.set_constant_buf_addr(i, addr)
@@ -516,7 +519,7 @@ class NVDevice(HCQCompiled[HCQSignal]):
 
     self.cmdq_page:HCQBuffer = self.iface.alloc(0x200000, cpu_access=True)
     self.cmdq_allocator = BumpAllocator(size=self.cmdq_page.size, base=cast(int, self.cmdq_page.va_addr), wrap=True)
-    self.cmdq = MMIOInterface(cast(int, self.cmdq_page.va_addr), 0x200000, fmt='I')
+    self.cmdq = self.cmdq_page.cpu_view().view(fmt='I')
 
     self.num_gpcs, self.num_tpc_per_gpc, self.num_sm_per_tpc, self.max_warps_per_sm, self.sm_version = self._query_gpu_info('num_gpcs',
       'num_tpc_per_gpc', 'num_sm_per_tpc', 'max_warps_per_sm', 'sm_version')
@@ -526,7 +529,8 @@ class NVDevice(HCQCompiled[HCQSignal]):
     self.sass_version = ((self.sm_version & 0xf00) >> 4) | (self.sm_version & 0xf)
 
     compilers:list[CompilerPairT] = [(functools.partial(NVRenderer, self.arch),functools.partial(CUDACompiler if MOCKGPU else NVCompiler, self.arch)),
-      (functools.partial(PTXRenderer, self.arch, device="NV"), functools.partial(PTXCompiler if MOCKGPU else NVPTXCompiler, self.arch))]
+      (functools.partial(PTXRenderer, self.arch, device="NV"), functools.partial(PTXCompiler if MOCKGPU else NVPTXCompiler, self.arch)),
+      (functools.partial(NAKRenderer, dev=self), functools.partial(NAKCompiler, self.arch, self.max_warps_per_sm))]
     super().__init__(device, NVAllocator(self), compilers, functools.partial(NVProgram, self), HCQSignal, NVComputeQueue, NVCopyQueue)
 
     self._setup_gpfifos()
@@ -548,8 +552,8 @@ class NVDevice(HCQCompiled[HCQSignal]):
       nv_gpu.NVC36F_CTRL_CMD_GPFIFO_GET_WORK_SUBMIT_TOKEN_PARAMS(workSubmitToken=-1))
     self.iface.setup_gpfifo_vm(gpfifo)
 
-    return GPFifo(ring=MMIOInterface(gpfifo_area.va_addr + offset, entries*8, fmt='Q'), entries_count=entries, token=ws_token_params.workSubmitToken,
-                  controls=nv_gpu.AmpereAControlGPFifo.from_address(gpfifo_area.va_addr + offset + entries * 8))
+    return GPFifo(ring=gpfifo_area.cpu_view().view(offset, entries*8, fmt='Q'), entries_count=entries, token=ws_token_params.workSubmitToken,
+                  controls=nv_gpu.AmpereAControlGPFifo.from_address(gpfifo_area.cpu_view().addr + offset + entries * 8))
 
   def _query_gpu_info(self, *reqs):
     nvrs = [getattr(nv_gpu,'NV2080_CTRL_GR_INFO_INDEX_'+r.upper(), getattr(nv_gpu,'NV2080_CTRL_GR_INFO_INDEX_LITTER_'+r.upper(), None)) for r in reqs]

tinygrad/runtime/ops_tinyfs.py

@@ -32,6 +32,9 @@ class TinyFSDevice(Compiled):
     self.conn_pools: dict[str, asyncio.Queue] = {}
     self.conn_pools_lock = asyncio.Lock()
 
+    # current request
+    self.request_id = uuid.UUID(int=0)
+
   def finalize(self):
     self.sfile.close()
 
@@ -71,9 +74,8 @@ class TinyFSDevice(Compiled):
       await self.conn_pools[loc].put((reader, writer))
 
 class TinyFSBuffer:
-  def __init__(self, device:TinyFSDevice, size:int, offset=0, request_id=None, copyout_queue=None):
+  def __init__(self, device:TinyFSDevice, size:int, offset=0, copyout_queue=None):
     self.device, self.size, self.offset = device, size, offset
-    self.request_id: uuid.UUID|None = request_id
     self.copyout_queue = copyout_queue or []
   def __repr__(self): return f"<TinyFSBuffer size={self.size} offset={self.offset}>"
 
@@ -87,8 +89,8 @@ class TinyFSAllocator(Allocator[TinyFSDevice]):
 
     if dest.device.op == "STORE":
       self.dev.sfile.flush()
-      dest.request_id = uuid.UUID(bytes=self.dev.sfile.read(16))
-      if DEBUG >= 2: print(f"Request ID: {dest.request_id}")
+      self.dev.request_id = uuid.UUID(bytes=self.dev.sfile.read(16))
+      if DEBUG >= 2: print(f"Request ID: {self.dev.request_id}")
 
     self.dev.sfile.write(src)
     self.dev.sfile.flush()
@@ -99,17 +101,15 @@ class TinyFSAllocator(Allocator[TinyFSDevice]):
 
       dest.copyout_queue = []
       for i, loc in enumerate(locs):
-        dest.copyout_queue.append((i, loc, src[i*16:(i+1)*16]))
+        dest.copyout_queue.append((i, loc, src[i*16:(i+1)*16].tobytes()))
 
   def _copyout(self, dest:memoryview, src:TinyFSBuffer):
     if DEBUG >= 2: print(f"Copying out {src.size} bytes from TINYFS:{src.device.op}")
     if src.device.op == "LOAD":
       asyncio.run_coroutine_threadsafe(self._copyout_async(dest, src), src.device.loop).result()
     else:
-      self.dev.sfile.write(f"{src.device.op}_OUT {src.size} {src.request_id}\r\n".encode())
+      self.dev.sfile.write(f"{src.device.op}_OUT {src.size} {self.dev.request_id}\r\n".encode())
       self.dev.sfile.flush()
-      src.request_id = uuid.UUID(bytes=self.dev.sfile.read(16))
-      if DEBUG >= 2: print(f"Request ID: {src.request_id}")
       self.dev.sfile.readinto(dest)
 
   async def _copyout_async(self, dest:memoryview, src:TinyFSBuffer):
@@ -131,7 +131,6 @@ class TinyFSAllocator(Allocator[TinyFSDevice]):
 
     workers = [asyncio.create_task(_worker(item)) for item in src.copyout_queue]
     await asyncio.gather(*workers)
-    src.copyout_queue.clear()
 
   def _offset(self, buf:TinyFSBuffer, size:int, offset:int):
-    return TinyFSBuffer(buf.device, size, offset, buf.request_id, buf.copyout_queue)
+    return TinyFSBuffer(buf.device, size, offset, buf.copyout_queue)

tinygrad/runtime/support/am/amdev.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import ctypes, collections, dataclasses, functools, os, hashlib
+import ctypes, collections, dataclasses, functools, os, hashlib, array
 from tinygrad.helpers import mv_address, getenv, DEBUG, fetch
 from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.hcq import MMIOInterface
@@ -168,7 +168,7 @@ class AMDev(PCIDevImplBase):
     # Memory manager & firmware
     self.mm = AMMemoryManager(self, self.vram_size, boot_size=(32 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39], va_bits=48,
       first_lv=am.AMDGPU_VM_PDB2, va_base=AMMemoryManager.va_allocator.base,
-      palloc_ranges=[(1 << (i + 12), 0x1000) for i in range(9 * (3 - am.AMDGPU_VM_PDB2), -1, -1)])
+      palloc_ranges=[(1 << (i + 12), 0x1000) for i in range(9 * (3 - am.AMDGPU_VM_PDB2), -1, -1)], reserve_ptable=not self.large_bar)
     self.fw = AMFirmware(self)
 
     # Initialize IP blocks
@@ -217,14 +217,25 @@ class AMDev(PCIDevImplBase):
     self.reg("regBIF_BX_PF0_RSMU_INDEX").write(reg)
     self.reg("regBIF_BX_PF0_RSMU_DATA").write(val)
 
+  def _read_vram(self, addr, size) -> bytes:
+    assert addr % 4 == 0 and size % 4 == 0, f"Invalid address {addr:#x} or size {size:#x}"
+    res = []
+    for caddr in range(addr, addr + size, 4):
+      self.wreg(0x06, caddr >> 31)
+      self.wreg(0x00, (caddr & 0x7FFFFFFF) | 0x80000000)
+      res.append(self.rreg(0x01))
+    return bytes(array.array('I', res))
+
   def _run_discovery(self):
     # NOTE: Fixed register to query memory size without known ip bases to find the discovery table.
     #       The table is located at the end of VRAM - 64KB and is 10KB in size.
     mmRCC_CONFIG_MEMSIZE = 0xde3
     self.vram_size = self.rreg(mmRCC_CONFIG_MEMSIZE) << 20
+    self.large_bar = self.vram.nbytes >= self.vram_size
     tmr_offset, tmr_size = self.vram_size - (64 << 10), (10 << 10)
 
-    self.bhdr = am.struct_binary_header.from_buffer(bytearray(self.vram.view(tmr_offset, tmr_size)[:]))
+    disc_tbl = self.vram.view(tmr_offset, tmr_size)[:] if self.large_bar else self._read_vram(tmr_offset, tmr_size)
+    self.bhdr = am.struct_binary_header.from_buffer(bytearray(disc_tbl))
     ihdr = am.struct_ip_discovery_header.from_address(ctypes.addressof(self.bhdr) + self.bhdr.table_list[am.IP_DISCOVERY].offset)
     assert self.bhdr.binary_signature == am.BINARY_SIGNATURE and ihdr.signature == am.DISCOVERY_TABLE_SIGNATURE, "discovery signatures mismatch"
 

tinygrad/runtime/support/am/ip.py

@@ -1,6 +1,6 @@
 import ctypes, time, contextlib, functools
 from typing import Literal
-from tinygrad.helpers import to_mv, data64, lo32, hi32, DEBUG, wait_cond
+from tinygrad.helpers import to_mv, data64, lo32, hi32, DEBUG, wait_cond, pad_bytes
 from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.amd import import_soc
 
@@ -417,7 +417,8 @@ class AM_PSP(AM_IP):
 
   def _prep_msg1(self, data:memoryview):
     assert len(data) <= self.msg1_view.nbytes, f"msg1 buffer is too small {len(data):#x} > {self.msg1_view.nbytes:#x}"
-    self.msg1_view[:len(data)+4] = bytes(data) + b'\x00' * 4
+    padded_data = pad_bytes(bytes(data) + b'\x00' * 4, 16) # HACK: apple's memcpy requires 16-bytes alignment
+    self.msg1_view[:len(padded_data)] = padded_data
     self.adev.gmc.flush_hdp()
 
   def _bootloader_load_component(self, fw:int, compid:int):

tinygrad/runtime/support/compiler_cuda.py

@@ -64,7 +64,8 @@ class PTXCompiler(Compiler):
   def __init__(self, arch:str, cache_key="ptx"):
     self.arch = arch
     super().__init__(f"compile_{cache_key}_{self.arch}")
-  def compile(self, src:str) -> bytes: return src.replace("TARGET", self.arch).replace("VERSION", "7.8" if self.arch >= "sm_89" else "7.5").encode()
+  def compile(self, src:str) -> bytes:
+    return src.replace("TARGET", self.arch).replace("VERSION", "8.7" if (ver:=int(self.arch[3:]))>=120 else ("7.8" if ver>=89 else "7.5")).encode()
   def disassemble(self, lib:bytes): cuda_disassemble(lib, self.arch)
 
 class NVPTXCompiler(PTXCompiler):

tinygrad/runtime/support/compiler_mesa.py

@@ -0,0 +1,86 @@
+import base64, ctypes, pathlib, tempfile, hashlib, subprocess
+from tinygrad.device import Compiler
+from tinygrad.helpers import cpu_objdump
+import tinygrad.runtime.autogen.mesa as mesa
+from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, expect, cerr
+try: import tinygrad.runtime.autogen.llvm as llvm
+except (ImportError, FileNotFoundError): llvm = None #type:ignore[assignment]
+
+def deserialize(enc_src, opts):
+  blobreader = mesa.struct_blob_reader()
+  mesa.blob_reader_init(blobreader, src:=base64.b64decode(enc_src), len(src))
+  return mesa.nir_deserialize(None, ctypes.cast(opts, ctypes.POINTER(mesa.nir_shader_compiler_options)), blobreader)
+
+class NIRCompiler(Compiler):
+  def __init__(self, cache_key):
+    mesa.glsl_type_singleton_init_or_ref()
+    super().__init__(cache_key)
+  def __del__(self): mesa.glsl_type_singleton_decref()
+
+class LVPCompiler(CPULLVMCompiler, NIRCompiler):
+  def __init__(self, cache_key="lvp"):
+    CPULLVMCompiler.__init__(self)
+    NIRCompiler.__init__(self, f"compile_{cache_key}")
+
+  def __del__(self):
+    NIRCompiler.__del__(self)
+    CPULLVMCompiler.__del__(self)
+
+  def compile(self, src) -> bytes:
+    shader, ctx = deserialize(src, mesa.lvp_nir_options), llvm.LLVMGetGlobalContext()
+    gallivm = mesa.gallivm_create(None, mesa.lp_context_ref(ctypes.cast(ctx, ctypes.POINTER(mesa.struct_LLVMOpaqueContext)), True), None).contents
+    module, builder = ctypes.cast(gallivm.module, llvm.LLVMModuleRef), ctypes.cast(gallivm.builder, llvm.LLVMBuilderRef)
+
+    params = mesa.struct_lp_build_tgsi_params(mesa.struct_lp_type(floating=True, sign=True, width=32, length=4),
+      resources_type=mesa.lp_build_jit_resources_type(gallivm), mask=ctypes.pointer(mesa.struct_lp_build_mask_context()))
+
+    pt = llvm.LLVMPointerType(ctypes.cast(params.resources_type, llvm.LLVMTypeRef), 0)
+    fn = llvm.LLVMAddFunction(module, shader.contents.info.name, llvm.LLVMFunctionType(llvm.LLVMVoidTypeInContext(ctx), pt, 1, 0))
+    llvm.LLVMPositionBuilderAtEnd(builder, llvm.LLVMAppendBasicBlockInContext(ctx, fn, b"entry"))
+
+    params.consts_ptr = mesa.lp_build_struct_get_ptr2(gallivm, params.resources_type,
+      ctypes.cast(llvm.LLVMGetParam(fn, 0), mesa.LLVMValueRef), mesa.LP_JIT_RES_CONSTANTS, b"constants")
+    mesa.lp_build_mask_begin(params.mask, gallivm, params.type, mesa.lp_build_one(gallivm, params.type))
+    mesa.lp_build_mask_end(params.mask)
+
+    mesa.lp_build_nir_soa(gallivm, shader, params, None)
+    llvm.LLVMBuildRetVoid(builder)
+    mesa.gallivm_verify_function(gallivm, ctypes.cast(fn, mesa.LLVMValueRef))
+    mesa.lp_passmgr_run(gallivm.passmgr, gallivm.module, ctypes.cast(self.target_machine, mesa.LLVMTargetMachineRef), gallivm.module_name)
+    obj_buf = expect(llvm.LLVMTargetMachineEmitToMemoryBuffer(self.target_machine, module, llvm.LLVMObjectFile, err:=cerr(),
+                                                              ctypes.pointer(buf:=llvm.LLVMMemoryBufferRef())), err, buf)
+    obj = ctypes.string_at(llvm.LLVMGetBufferStart(obj_buf), llvm.LLVMGetBufferSize(obj_buf))
+
+    mesa.gallivm_destroy(gallivm)
+    mesa.ralloc_free(shader)
+    return obj
+
+  def disassemble(self, lib: bytes): cpu_objdump(lib)
+
+class NAKCompiler(NIRCompiler):
+  def __init__(self, arch, warps_per_sm, cache_key="nak"):
+    self.arch, self.warps_per_sm = arch, warps_per_sm
+    self.cc = mesa.nak_compiler_create(mesa.struct_nv_device_info(sm=int(arch[3:]), max_warps_per_mp=warps_per_sm))
+    self.nir_options = bytes(mesa.nak_nir_options(self.cc).contents)
+    super().__init__(f"compile_{cache_key}_{arch}")
+
+  def __del__(self):
+    mesa.nak_compiler_destroy(self.cc)
+    super().__del__()
+
+  def __reduce__(self): return NAKCompiler, (self.arch, self.warps_per_sm)
+
+  def compile(self, src) -> bytes:
+    shader = deserialize(src, self.nir_options)
+    mesa.nak_preprocess_nir(shader, self.cc)
+    ret = bytes((out:=mesa.nak_compile_shader(shader, False, self.cc, 0, None).contents).info) + ctypes.string_at(out.code, out.code_size)
+    mesa.nak_shader_bin_destroy(out)
+    mesa.ralloc_free(shader)
+    return ret
+
+  def disassemble(self, lib: bytes):
+    try:
+      fn = (pathlib.Path(tempfile.gettempdir()) / f"tinynak_{hashlib.md5(lib).hexdigest()}").as_posix()
+      with open(fn, "wb") as f: f.write(lib[ctypes.sizeof(mesa.struct_nak_shader_info):])
+      print(subprocess.check_output(['nvdisasm', "-b", f"SM{self.arch[3:]}", fn]).decode('utf-8'))
+    except Exception as e: print("Failed to generate SASS", str(e), "Make sure your PATH contains nvdisasm binary of compatible version.")

tinygrad/runtime/support/elf.py

@@ -1,12 +1,18 @@
-import struct
+import struct, ctypes, ctypes.util
 from dataclasses import dataclass
-from tinygrad.helpers import getbits, i2u
+from tinygrad.helpers import getbits, i2u, unwrap
 import tinygrad.runtime.autogen.libc as libc
 
 @dataclass(frozen=True)
 class ElfSection: name:str; header:libc.Elf64_Shdr; content:bytes # noqa: E702
 
-def elf_loader(blob:bytes, force_section_align:int=1) -> tuple[memoryview, list[ElfSection], list[tuple]]:
+def link_sym(sym:str, libs:list[str]) -> int:
+  for lib in libs:
+    try: return unwrap(ctypes.cast(getattr(ctypes.CDLL(ctypes.util.find_library(lib)), sym), ctypes.c_void_p).value)
+    except (OSError, AttributeError): pass
+  raise RuntimeError(f'Attempting to relocate against an undefined symbol {sym}')
+
+def elf_loader(blob:bytes, force_section_align:int=1, link_libs:list[str]|None=None) -> tuple[memoryview, list[ElfSection], list[tuple]]:
   def _strtab(blob: bytes, idx: int) -> str: return blob[idx:blob.find(b'\x00', idx)].decode('utf-8')
 
   header = libc.Elf64_Ehdr.from_buffer_copy(blob)
@@ -31,33 +37,42 @@ def elf_loader(blob:bytes, force_section_align:int=1) -> tuple[memoryview, list[
   # Relocations
   relocs = []
   for sh, trgt_sh_name, c_rels in rel + rela:
+    if trgt_sh_name == ".eh_frame": continue
     target_image_off = next(tsh for tsh in sections if tsh.name == trgt_sh_name).header.sh_addr
     rels = [(r.r_offset, symtab[libc.ELF64_R_SYM(r.r_info)], libc.ELF64_R_TYPE(r.r_info), getattr(r, "r_addend", 0)) for r in c_rels]
-    for _, sym, _, _ in rels:
-      if sym.st_shndx == 0: raise RuntimeError(f'Attempting to relocate against an undefined symbol {repr(_strtab(sh_strtab, sym.st_name))}')
-    relocs += [(target_image_off + roff, sections[sym.st_shndx].header.sh_addr + sym.st_value, rtype, raddend) for roff, sym, rtype, raddend in rels]
+    relocs += [(target_image_off + roff, link_sym(_strtab(sh_strtab, sym.st_name), link_libs or []) if sym.st_shndx == 0 else
+                sections[sym.st_shndx].header.sh_addr + sym.st_value, rtype, raddend) for roff, sym, rtype, raddend in rels]
 
   return memoryview(image), sections, relocs
 
-def relocate(instr: int, ploc: int, tgt: int, r_type: int):
-  match r_type:
-    # https://refspecs.linuxfoundation.org/elf/x86_64-abi-0.95.pdf
-    case libc.R_X86_64_PC32: return i2u(32, tgt-ploc)
-    # https://github.com/ARM-software/abi-aa/blob/main/aaelf64/aaelf64.rst for definitions of relocations
-    # https://www.scs.stanford.edu/~zyedidia/arm64/index.html for instruction encodings
-    case libc.R_AARCH64_ADR_PREL_PG_HI21:
-      rel_pg = (tgt & ~0xFFF) - (ploc & ~0xFFF)
-      return instr | (getbits(rel_pg, 12, 13) << 29) | (getbits(rel_pg, 14, 32) << 5)
-    case libc.R_AARCH64_ADD_ABS_LO12_NC: return instr | (getbits(tgt, 0, 11) << 10)
-    case libc.R_AARCH64_LDST16_ABS_LO12_NC: return instr | (getbits(tgt, 1, 11) << 10)
-    case libc.R_AARCH64_LDST32_ABS_LO12_NC: return instr | (getbits(tgt, 2, 11) << 10)
-    case libc.R_AARCH64_LDST64_ABS_LO12_NC: return instr | (getbits(tgt, 3, 11) << 10)
-    case libc.R_AARCH64_LDST128_ABS_LO12_NC: return instr | (getbits(tgt, 4, 11) << 10)
-  raise NotImplementedError(f"Encountered unknown relocation type {r_type}")
+def jit_loader(obj: bytes, base:int=0, link_libs:list[str]|None=None) -> bytes:
+  image_, _, relocs = elf_loader(obj, link_libs=link_libs)
+  image = bytearray(image_)
+
+  def relocate(instr: int, base: int, ploc: int, tgt: int, r_type: int):
+    match r_type:
+      # https://refspecs.linuxfoundation.org/elf/x86_64-abi-0.95.pdf
+      case libc.R_X86_64_PC32: return i2u(32, tgt-ploc)
+      case libc.R_X86_64_PLT32: return i2u(32, tgt-ploc-base)
+      # https://github.com/ARM-software/abi-aa/blob/main/aaelf64/aaelf64.rst for definitions of relocations
+      # https://www.scs.stanford.edu/~zyedidia/arm64/index.html for instruction encodings
+      case libc.R_AARCH64_ADR_PREL_PG_HI21:
+        rel_pg = (tgt & ~0xFFF) - (ploc & ~0xFFF)
+        return instr | (getbits(rel_pg, 12, 13) << 29) | (getbits(rel_pg, 14, 32) << 5)
+      case libc.R_AARCH64_ADD_ABS_LO12_NC: return instr | (getbits(tgt, 0, 11) << 10)
+      case libc.R_AARCH64_LDST16_ABS_LO12_NC: return instr | (getbits(tgt, 1, 11) << 10)
+      case libc.R_AARCH64_LDST32_ABS_LO12_NC: return instr | (getbits(tgt, 2, 11) << 10)
+      case libc.R_AARCH64_LDST64_ABS_LO12_NC: return instr | (getbits(tgt, 3, 11) << 10)
+      case libc.R_AARCH64_LDST128_ABS_LO12_NC: return instr | (getbits(tgt, 4, 11) << 10)
+      case libc.R_AARCH64_CALL26:
+        if -(2**25) <= tgt-ploc-base and tgt-ploc-base <= (2**25 - 1) * 4: return instr | getbits(tgt-ploc-base, 2, 27)
+        nonlocal image
+        # create trampoline:         LDR x17, 8  BR x17
+        image += struct.pack("<IIQ", 0x58000051, 0xD61F0220, tgt)
+        return instr | getbits(len(image)-ploc-16, 2, 27)
+    raise NotImplementedError(f"Encountered unknown relocation type {r_type}")
 
-def jit_loader(obj: bytes) -> bytes:
-  image, _, relocs = elf_loader(obj)
   # This is needed because we have an object file, not a .so that has all internal references (like loads of constants from .rodata) resolved.
   for ploc,tgt,r_type,r_addend in relocs:
-    image[ploc:ploc+4] = struct.pack("<I", relocate(struct.unpack("<I", image[ploc:ploc+4])[0], ploc, tgt+r_addend, r_type))
+    image[ploc:ploc+4] = struct.pack("<I", relocate(struct.unpack("<I", image[ploc:ploc+4])[0], base, ploc, tgt+r_addend, r_type))
   return bytes(image)

tinygrad/runtime/support/hcq.py

@@ -509,7 +509,7 @@ class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
         self.dev.timeline_signal.wait(self.b_timeline[self.b_next])
 
         lsize = min(self.b[self.b_next].size, src.nbytes - i)
-        self.b[self.b_next].cpu_view().view(size=lsize, fmt='B')[:] = src[i:i+lsize]
+        self.b[self.b_next].cpu_view().view(size=lsize, fmt='B')[:] = src.cast('B')[i:i+lsize]
         self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
                                   .copy(dest.va_addr+i, self.b[self.b_next].va_addr, lsize) \
                                   .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
@@ -541,7 +541,7 @@ class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
                                   .copy(self.b[0].va_addr, src.va_addr+i, lsize:=min(cp_size, dest.nbytes-i)) \
                                   .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
         self.dev.timeline_signal.wait(self.dev.timeline_value - 1)
-        dest[i:i+lsize] = self.b[0].cpu_view().view(size=lsize, fmt='B')[:]
+        dest.cast('B')[i:i+lsize] = self.b[0].cpu_view().view(size=lsize, fmt='B')[:]
 
   def _transfer(self, dest:HCQBuffer, src:HCQBuffer, sz:int, src_dev:HCQDeviceType, dest_dev:HCQDeviceType):
     cast(HCQAllocator, src_dev.allocator).map(dest)

tinygrad/runtime/support/memory.py

@@ -28,7 +28,7 @@ class TLSFAllocator:
 
     # self.blocks is more like a linked list, where each entry is a contiguous block.
     self.blocks:dict[int, tuple[int, int|None, int|None, bool]] = {0: (size, None, None, True)} # size, next, prev, is_free
-    self._insert_block(0, size)
+    if size > 0: self._insert_block(0, size)
 
   @functools.cache # pylint: disable=method-cache-max-size-none
   def lv1(self, size): return size.bit_length()
@@ -124,7 +124,7 @@ class PageTableTraverseContext:
 
     if not pt.valid(pte_idx):
       assert self.create_pts, "Not allowed to create new page table"
-      pt.set_entry(pte_idx, self.dev.mm.palloc(0x1000, zero=True, boot=self.boot), table=True, valid=True)
+      pt.set_entry(pte_idx, self.dev.mm.palloc(0x1000, zero=True, boot=self.boot, ptable=True), table=True, valid=True)
 
     assert not pt.is_page(pte_idx), f"Must be table pt={pt.paddr:#x}, {pt.lv=} {pte_idx=} {pt.read_fields(pte_idx)}"
     child_page_table = self.dev.mm.pt_t(self.dev, pt.address(pte_idx), lv=pt.lv+1)
@@ -135,7 +135,7 @@ class PageTableTraverseContext:
   def _try_free_pt(self) -> bool:
     pt, _, _ = self.pt_stack[-1]
     if self.free_pts and pt != self.dev.mm.root_page_table and all(not pt.valid(i) for i in range(self._pt_pte_cnt(self.pt_stack[-1][0].lv))):
-      self.dev.mm.pfree(pt.paddr)
+      self.dev.mm.pfree(pt.paddr, ptable=True)
       parent_pt, parent_pte_idx, _ = self.pt_stack[-2]
       parent_pt.set_entry(parent_pte_idx, 0x0, valid=False)
       return True
@@ -167,13 +167,14 @@ class MemoryManager:
   va_allocator: ClassVar[TLSFAllocator|None] = None
 
   def __init__(self, dev, vram_size:int, boot_size:int, pt_t, va_bits:int, va_shifts:list[int], va_base:int,
-               palloc_ranges:list[tuple[int, int]], first_lv:int=0):
+               palloc_ranges:list[tuple[int, int]], first_lv:int=0, reserve_ptable=False):
     self.dev, self.vram_size, self.va_shifts, self.va_base, lvl_msb = dev, vram_size, va_shifts, va_base, va_shifts + [va_bits + 1]
     self.pte_covers, self.pte_cnt = [1 << x for x in va_shifts][::-1], [1 << (lvl_msb[i+1] - lvl_msb[i]) for i in range(len(lvl_msb) - 1)][::-1]
-    self.pt_t, self.palloc_ranges, self.level_cnt, self.va_bits = pt_t, palloc_ranges, len(va_shifts), va_bits
+    self.pt_t, self.palloc_ranges, self.level_cnt, self.va_bits, self.reserve_ptable = pt_t, palloc_ranges, len(va_shifts), va_bits, reserve_ptable
 
-    self.boot_allocator = TLSFAllocator(boot_size, base=0) # per device
-    self.pa_allocator = TLSFAllocator(vram_size - (64 << 20), base=self.boot_allocator.size) # per device
+    self.boot_allocator = TLSFAllocator(boot_size, base=0)
+    self.ptable_allocator = TLSFAllocator(round_up(vram_size // 512, 1 << 20) if self.reserve_ptable else 0, base=self.boot_allocator.size)
+    self.pa_allocator = TLSFAllocator(vram_size - (64 << 20), base=self.boot_allocator.size + self.ptable_allocator.size)
     self.root_page_table = pt_t(self.dev, self.palloc(0x1000, zero=not self.dev.smi_dev, boot=True), lv=first_lv)
 
   def _frag_size(self, va, sz, must_cover=True):
@@ -250,10 +251,11 @@ class MemoryManager:
     self.va_allocator.free(vm.va_addr)
     for paddr, _ in vm.paddrs: self.pa_allocator.free(paddr)
 
-  def palloc(self, size:int, align:int=0x1000, zero=True, boot=False) -> int:
+  def palloc(self, size:int, align:int=0x1000, zero=True, boot=False, ptable=False) -> int:
     assert self.dev.is_booting == boot, "During booting, only boot memory can be allocated"
-    paddr = (self.boot_allocator if boot else self.pa_allocator).alloc(round_up(size, 0x1000), align)
+    allocator = self.boot_allocator if boot else (self.ptable_allocator if self.reserve_ptable and ptable else self.pa_allocator)
+    paddr = allocator.alloc(round_up(size, 0x1000), align)
     if zero: self.dev.vram[paddr:paddr+size] = bytes(size)
     return paddr
 
-  def pfree(self, paddr:int): self.pa_allocator.free(paddr)
+  def pfree(self, paddr:int, ptable=False): (self.ptable_allocator if self.reserve_ptable and ptable else self.pa_allocator).free(paddr)

tinygrad/runtime/support/nv/nvdev.py

@@ -87,7 +87,7 @@ class NVDev(PCIDevImplBase):
     # 5           PTE_64K / PTE_4K                    20:16 / 20:12
     bits, shifts = (56, [12, 21, 29, 38, 47, 56]) if self.mmu_ver == 3 else (48, [12, 21, 29, 38, 47])
     self.mm = NVMemoryManager(self, self.vram_size, boot_size=(2 << 20), pt_t=NVPageTableEntry, va_bits=bits, va_shifts=shifts, va_base=0,
-      palloc_ranges=[(x, x) for x in [512 << 20, 2 << 20, 4 << 10]])
+      palloc_ranges=[(x, x) for x in [512 << 20, 2 << 20, 4 << 10]], reserve_ptable=not self.large_bar)
     self.flcn:NV_FLCN|NV_FLCN_COT = NV_FLCN_COT(self) if self.fmc_boot else NV_FLCN(self)
     self.gsp:NV_GSP = NV_GSP(self)
 
@@ -114,6 +114,7 @@ class NVDev(PCIDevImplBase):
     self.chip_id = self.reg("NV_PMC_BOOT_0").read()
     self.chip_details = self.reg("NV_PMC_BOOT_42").read_bitfields()
     self.chip_name = {0x17: "GA1", 0x19: "AD1", 0x1b: "GB2"}[self.chip_details['architecture']] + f"{self.chip_details['implementation']:02d}"
+    self.fw_name = {"GB2": "GB202", "AD1": "AD102", "GA1": "GA102"}[self.chip_name[:3]]
     self.mmu_ver, self.fmc_boot = (3, True) if self.chip_details['architecture'] >= 0x1a else (2, False)
 
     self.include("src/common/inc/swref/published/turing/tu102/dev_fb.h")
@@ -133,6 +134,7 @@ class NVDev(PCIDevImplBase):
     self.pte_t, self.pde_t, self.dual_pde_t = tuple([self.__dict__[name] for name in mmu_pd_names])
 
     self.vram_size = self.reg("NV_PGC6_AON_SECURE_SCRATCH_GROUP_42").read() << 20
+    self.large_bar = self.vram.nbytes >= self.vram_size
 
   def _alloc_boot_struct(self, struct:ctypes.Structure) -> tuple[ctypes.Structure, int]:
     va, paddrs = System.alloc_sysmem(sz:=ctypes.sizeof(type(struct)), contiguous=True)
@@ -146,8 +148,8 @@ class NVDev(PCIDevImplBase):
   def extract_fw(self, file:str, dname:str) -> bytes:
     # Extracts the firmware binary from the given header
     tname = file.replace("kgsp", "kgspGet")
-    text = self._download(f"src/nvidia/generated/g_bindata_{tname}_{self.chip_name}.c")
-    info, sl = text[text[:text.index(dnm:=f'{file}_{self.chip_name}_{dname}')].rindex("COMPRESSION:"):][:16], text[text.index(dnm) + len(dnm) + 7:]
+    text = self._download(f"src/nvidia/generated/g_bindata_{tname}_{self.fw_name}.c")
+    info, sl = text[text[:text.index(dnm:=f'{file}_{self.fw_name}_{dname}')].rindex("COMPRESSION:"):][:16], text[text.index(dnm) + len(dnm) + 7:]
     image = bytes.fromhex(sl[:sl.find("};")].strip().replace("0x", "").replace(",", "").replace(" ", "").replace("\n", ""))
     return gzip.decompress(struct.pack("<4BL2B", 0x1f, 0x8b, 8, 0, 0, 0, 3) + image) if "COMPRESSION: YES" in info else image
 

tinygrad/runtime/support/system.py

@@ -1,4 +1,4 @@
-import os, mmap, array, functools, ctypes, select, contextlib, dataclasses, sys, errno
+import os, mmap, array, functools, ctypes, select, contextlib, dataclasses, sys, errno, itertools
 from typing import cast, ClassVar
 from tinygrad.helpers import round_up, to_mv, getenv, OSX, temp
 from tinygrad.runtime.autogen import libc, vfio
@@ -8,9 +8,50 @@ from tinygrad.runtime.support.memory import MemoryManager, VirtMapping
 MAP_FIXED, MAP_LOCKED, MAP_POPULATE, MAP_NORESERVE = 0x10, 0 if OSX else 0x2000, getattr(mmap, "MAP_POPULATE", 0 if OSX else 0x008000), 0x400
 
 class _System:
+  @functools.cached_property
+  def atomic_lib(self): return ctypes.CDLL(ctypes.util.find_library('atomic')) if sys.platform == "linux" else None
+
+  @functools.cached_property
+  def iokit(self): return ctypes.CDLL(ctypes.util.find_library("IOKit"))
+
+  @functools.cached_property
+  def libsys(self): return ctypes.CDLL(ctypes.util.find_library("System"))
+
+  @functools.cached_property
+  def mach_task_self(self): return ctypes.cast(self.libsys.mach_task_self_, ctypes.POINTER(ctypes.c_uint)).contents.value
+
+  @functools.cached_property
+  def pagemap(self) -> FileIOInterface:
+    if FileIOInterface(reloc_sysfs:="/proc/sys/vm/compact_unevictable_allowed", os.O_RDONLY).read()[0] != "0":
+      os.system(cmd:=f"sudo sh -c 'echo 0 > {reloc_sysfs}'")
+      assert FileIOInterface(reloc_sysfs, os.O_RDONLY).read()[0] == "0", f"Failed to disable migration of locked pages. Please run {cmd} manually."
+    return FileIOInterface("/proc/self/pagemap", os.O_RDONLY)
+
+  @functools.cached_property
+  def vfio(self) -> FileIOInterface|None:
+    try:
+      if not FileIOInterface.exists("/sys/module/vfio"): os.system("sudo modprobe vfio-pci disable_idle_d3=1")
+
+      FileIOInterface("/sys/module/vfio/parameters/enable_unsafe_noiommu_mode", os.O_RDWR).write("1")
+      vfio_fd = FileIOInterface("/dev/vfio/vfio", os.O_RDWR)
+      vfio.VFIO_CHECK_EXTENSION(vfio_fd, vfio.VFIO_NOIOMMU_IOMMU)
+
+      return vfio_fd
+    except OSError: return None
+
+  @functools.cached_property
+  def macos_tinygpu_conn(self):
+    self.iokit.IOServiceNameMatching.restype = ctypes.c_void_p # CFMutableDictionaryRef
+    if not (mdict:=self.iokit.IOServiceNameMatching("tinygpu".encode("utf-8"))): raise RuntimeError("IOServiceNameMatching returned NULL")
+    if not (service:=self.iokit.IOServiceGetMatchingService(ctypes.c_uint(0), ctypes.c_void_p(mdict))):
+      raise RuntimeError('Service "tinygpu" is not running')
+    if self.iokit.IOServiceOpen(service, self.mach_task_self, ctypes.c_uint32(0), ctypes.byref(conn:=ctypes.c_uint(0))):
+      raise RuntimeError("IOServiceOpen failed")
+    return conn
+
   def reserve_hugepages(self, cnt): os.system(f"sudo sh -c 'echo {cnt} > /proc/sys/vm/nr_hugepages'")
 
-  def memory_barrier(self): lib.atomic_thread_fence(__ATOMIC_SEQ_CST:=5) if (lib:=self.atomic_lib) is not None else None
+  def memory_barrier(self): lib.atomic_thread_fence(__ATOMIC_SEQ_CST:=5) if (lib:=self.libsys if OSX else self.atomic_lib) is not None else None
 
   def lock_memory(self, addr:int, size:int):
     if libc.mlock(ctypes.c_void_p(addr), size): raise RuntimeError(f"Failed to lock memory at {addr:#x} with size {size:#x}")
@@ -36,28 +77,6 @@ class _System:
       if vendor == target_vendor and device in target_devices: result.append(pcibus)
     return sorted(result)
 
-  @functools.cached_property
-  def atomic_lib(self): return ctypes.CDLL(ctypes.util.find_library('atomic')) if sys.platform == "linux" else None
-
-  @functools.cached_property
-  def pagemap(self) -> FileIOInterface:
-    if FileIOInterface(reloc_sysfs:="/proc/sys/vm/compact_unevictable_allowed", os.O_RDONLY).read()[0] != "0":
-      os.system(cmd:=f"sudo sh -c 'echo 0 > {reloc_sysfs}'")
-      assert FileIOInterface(reloc_sysfs, os.O_RDONLY).read()[0] == "0", f"Failed to disable migration of locked pages. Please run {cmd} manually."
-    return FileIOInterface("/proc/self/pagemap", os.O_RDONLY)
-
-  @functools.cached_property
-  def vfio(self) -> FileIOInterface|None:
-    try:
-      if not FileIOInterface.exists("/sys/module/vfio"): os.system("sudo modprobe vfio-pci disable_idle_d3=1")
-
-      FileIOInterface("/sys/module/vfio/parameters/enable_unsafe_noiommu_mode", os.O_RDWR).write("1")
-      vfio_fd = FileIOInterface("/dev/vfio/vfio", os.O_RDWR)
-      vfio.VFIO_CHECK_EXTENSION(vfio_fd, vfio.VFIO_NOIOMMU_IOMMU)
-
-      return vfio_fd
-    except OSError: return None
-
   def flock_acquire(self, name:str) -> int:
     import fcntl # to support windows
 
@@ -123,13 +142,23 @@ class PCIDevice:
     libc.madvise(loc:=fd.mmap(addr, sz, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | (MAP_FIXED if addr else 0), off), sz, libc.MADV_DONTFORK)
     return MMIOInterface(loc, sz, fmt=fmt)
 
+class APLPCIDevice(PCIDevice):
+  def __init__(self, pcibus:str, bars:list[int], resize_bars:list[int]|None=None): self.pcibus, self.bars = pcibus, {b: self.map_mem(b) for b in bars}
+  def map_mem(self, typ:int) -> MMIOInterface:
+    if System.iokit.IOConnectMapMemory64(System.macos_tinygpu_conn, ctypes.c_uint32(typ), System.mach_task_self,
+      ctypes.byref(addr:=ctypes.c_uint64(0)), ctypes.byref(size:=ctypes.c_uint64(0)), 0x1): raise RuntimeError(f"IOConnectMapMemory64({typ=}) failed")
+    return MMIOInterface(addr.value, size.value)
+  def map_bar(self, bar:int, off:int=0, addr:int=0, size:int|None=None, fmt='B') -> MMIOInterface: return self.bars[bar].view(off, size, fmt)
+  def read_config(self, offset:int, size:int): return 0
+  def write_config(self, offset:int, value:int, size:int): pass
+
 class PCIDevImplBase:
   mm: MemoryManager
 
 @dataclasses.dataclass
 class PCIAllocationMeta: mapping:VirtMapping; has_cpu_mapping:bool; hMemory:int=0 # noqa: E702
 
-class PCIIfaceBase:
+class LNXPCIIfaceBase:
   dev_impl:PCIDevImplBase
   gpus:ClassVar[list[str]] = []
 
@@ -166,9 +195,31 @@ class PCIIfaceBase:
     if b.owner is not None and b.owner._is_cpu():
       System.lock_memory(cast(int, b.va_addr), b.size)
       paddrs, snooped, uncached = [(x, 0x1000) for x in System.system_paddrs(cast(int, b.va_addr), round_up(b.size, 0x1000))], True, True
-    elif (ifa:=getattr(b.owner, "iface", None)) is not None and isinstance(ifa, PCIIfaceBase):
+    elif (ifa:=getattr(b.owner, "iface", None)) is not None and isinstance(ifa, LNXPCIIfaceBase):
       paddrs = [(paddr if b.meta.mapping.system else (paddr + ifa.p2p_base_addr), size) for paddr,size in b.meta.mapping.paddrs]
       snooped, uncached = b.meta.mapping.snooped, b.meta.mapping.uncached
     else: raise RuntimeError(f"map failed: {b.owner} -> {self.dev}")
 
     self.dev_impl.mm.map_range(cast(int, b.va_addr), round_up(b.size, 0x1000), paddrs, system=True, snooped=snooped, uncached=uncached)
+
+class APLPCIIfaceBase(LNXPCIIfaceBase):
+  def __init__(self, dev, dev_id, vendor, devices, bars, vram_bar, va_start, va_size):
+    self.pci_dev, self.dev, self.vram_bar = APLPCIDevice(pcibus=f'usb4:{dev_id}', bars=bars), dev, vram_bar
+
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, **kwargs) -> HCQBuffer:
+    if host or uncached or cpu_access: # cpu access memory goes here, since bar is small.
+      vaddr = self.dev_impl.mm.alloc_vaddr(size:=round_up(size, mmap.PAGESIZE), align=mmap.PAGESIZE)
+      assert size >= mmap.PAGESIZE, "Size must be at least one page"
+
+      sysmem = cast(APLPCIDevice, self.pci_dev).map_mem(size).view(fmt='Q')
+      paddrs = list(itertools.takewhile(lambda p: p[1] != 0, zip(sysmem[0::2], sysmem[1::2])))
+
+      mapping = self.dev_impl.mm.map_range(vaddr, size, paddrs, system=True, snooped=True, uncached=True)
+      return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(mapping, has_cpu_mapping=True), view=sysmem.view(fmt='B'), owner=self.dev)
+
+    mapping = self.dev_impl.mm.valloc(size:=round_up(size, 4 << 10), uncached=uncached, contiguous=cpu_access)
+    return HCQBuffer(mapping.va_addr, size, view=None, meta=PCIAllocationMeta(mapping, has_cpu_mapping=False), owner=self.dev)
+
+  def map(self, b:HCQBuffer): raise RuntimeError(f"map failed: {b.owner} -> {self.dev}")
+
+PCIIfaceBase:type = APLPCIIfaceBase if OSX else LNXPCIIfaceBase

tinygrad/schedule/indexing.py

@@ -147,6 +147,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
   ending_ranges: dict[UOp, bool] = {}
   for x in tsink_reverse_toposort:
     if x.op in {Ops.DEVICE, Ops.UNIQUE}: continue
+    if x.dtype.scalar() == dtypes.index: continue  # TODO: why do I need this?
     ending_ranges[x] = any(ending_ranges[u] for u in consumer_map[x])
 
     # if this element has weight and it's ending a range, we (force) realize it
@@ -210,7 +211,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     rngs = out_rngs  # rngs is the input ranges  # pylint: disable=possibly-used-before-assignment
 
     # apply movement ops
-    if x.op in GroupOp.Movement: rngs = apply_movement_op(x.op, x.src[0].shape, x.arg, rngs)
+    if x.op in GroupOp.Movement: rngs = apply_movement_op(x.op, x.src[0].shape, x.marg, rngs)
     # if the EXPAND is used to inject a range, we don't mark it as ending_ranges. otherwise we do.
     if x.op is Ops.EXPAND and all(isinstance(y, int) or y.op is not Ops.RANGE for y in x.shape): ending_ranges[x] = True
 

tinygrad/schedule/multi.py

@@ -85,7 +85,7 @@ def mstack_early_shrink(ms:UOp, shrink:UOp):
   ret:list[UOp] = []
   def apply_shrink(s:UOp, i:int) -> UOp:
     new_arg = [tuple([x.substitute({dvar[0]:dvar[0].const_like(i)}) if isinstance(x, UOp) and
-                      (dvar:=[v for v in x.vars() if v.op is Ops.DEFINE_VAR and v.arg[0]=='_device_num']) else x for x in ss]) for ss in shrink.arg]
+                      (dvar:=[v for v in x.vars() if v.op is Ops.DEFINE_VAR and v.arg[0]=='_device_num']) else x for x in ss]) for ss in shrink.marg]
     return s.shrink(tuple(new_arg))
   for i, x in enumerate(ms.src):
     if x.op is Ops.COPY:
@@ -111,9 +111,10 @@ replace_allreduce = PatternMatcher([
   # MSELECT on MSTACK is replaced with nothing
   (UPat(Ops.MSELECT, src=(UPat(Ops.MSTACK, name="mstack"),), name="ms"), lambda mstack, ms: mstack.src[ms.arg]),
   # move shrink before MSTACK
-  (UPat(Ops.SHRINK, src=(UPat(Ops.MSTACK, name="ms"),), name="shrink"), mstack_early_shrink),
+  (UPat(Ops.SHRINK, src=(UPat(Ops.MSTACK, name="ms"),), allow_any_len=True, name="shrink"), mstack_early_shrink),
   # move MSELECT before movement ops
-  (UPat(Ops.MSELECT, src=(UPat(GroupOp.Movement, src=(UPat.var("s"),), name="v"),), name="ms"), lambda s,v,ms: v.replace(src=(s.mselect(ms.arg),))),
+  (UPat(Ops.MSELECT, src=(UPat(GroupOp.Movement, src=(UPat.var("s"),), allow_any_len=True, name="v"),), name="ms"),
+   lambda s,v,ms: v.replace(src=(s.mselect(ms.arg),)+v.src[1:])),
 ])
 
 # ***** multi functions *****
@@ -152,40 +153,40 @@ def _shape_to_single_shard(axis, shape:tuple[sint, ...], lb:UOp) -> tuple[sint,
   return tuple(lb.shape[axis] if a == axis else s for a,s in enumerate(shape))
 
 def reshape_multi(root:UOp, multi:UOp):
-  arg = root.arg
+  arg = root.marg
   if (new_axis:=root.axis) is None: return multi.src[0].reshape(arg).multi(new_axis)
   assert prod(multi.shape) == prod(arg), "reshape must maintain prod(shape)"
-  assert prod(multi.src[0].shape[multi.axis:])%prod(arg[new_axis+1:]) == 0, f"reshape cannot move items between shards {multi.shape} -> {root.arg=}"
+  assert prod(multi.src[0].shape[multi.axis:])%prod(arg[new_axis+1:]) == 0, f"reshape cannot move items between shards {multi.shape} -> {arg=}"
   new_shape_axis = prod(multi.src[0].shape[multi.axis:]) // prod(arg[new_axis+1:])
   return multi.src[0].reshape(tuple(s if a!=new_axis else new_shape_axis for a,s in enumerate(arg))).multi(new_axis)
 
 def expand_multi(root:UOp, multi:UOp):
   # NOTE: this assert isn't needed, sharded axis can have dim 1
-  assert multi.axis is None or root.arg[multi.axis] == multi.shape[multi.axis], f"expand not supported on sharded axis {root.arg=}"
-  return multi.src[0].expand(_shape_to_single_shard(multi.axis, root.arg, multi.src[0])).multi(multi.axis)
+  assert multi.axis is None or root.marg[multi.axis] == multi.shape[multi.axis], f"expand not supported on sharded axis {root.marg=}"
+  return multi.src[0].expand(_shape_to_single_shard(multi.axis, root.marg, multi.src[0])).multi(multi.axis)
 
 def pad_multi(root:UOp, multi:UOp):
-  assert multi.axis is None or root.arg[multi.axis] == (0,0), f"padding not supported for {root.arg=}"
-  return multi.src[0].pad(root.arg).multi(multi.axis)
+  assert multi.axis is None or root.marg[multi.axis] == (0,0), f"padding not supported for {root.marg=}"
+  return multi.src[0].pad(root.marg).multi(multi.axis)
 
 def permute_multi(root:UOp, multi:UOp):
   # all permutes supported!
-  return multi.src[0].permute(root.arg).multi(root.axis)
+  return multi.src[0].permute(root.marg).multi(root.axis)
 
 def shrink_multi(root:UOp, multi:UOp):
-  assert multi.axis is None or root.arg[multi.axis] == (0, multi.shape[multi.axis]) or root.arg[multi.axis] in multi.bounds, \
-    f"shrinking not supported for {root.arg=}"
-  if multi.axis is not None and root.arg[multi.axis] in multi.bounds and root.arg[multi.axis] != (0, multi.shape[multi.axis]):
-    assert all(root.arg[i] == (0, s) or i == multi.axis for i,s in enumerate(multi.shape)), \
+  assert multi.axis is None or root.marg[multi.axis] == (0, multi.shape[multi.axis]) or root.marg[multi.axis] in multi.bounds, \
+    f"shrinking not supported for {root.marg=}"
+  if multi.axis is not None and root.marg[multi.axis] in multi.bounds and root.marg[multi.axis] != (0, multi.shape[multi.axis]):
+    assert all(root.marg[i] == (0, s) or i == multi.axis for i,s in enumerate(multi.shape)), \
       "cannot shrink sharded and non-sharded axis at the same time"
     # NOTE: shrink on the shard axis is only allowed when result is a single partition, denoted by the new real
     # we just copy it to all the devices, no real. this will be optimized out later
-    return multi.src[0].copy_to_device(multi.device, arg=multi.bounds.index(root.arg[multi.axis]))
-  return multi.src[0].shrink(tuple((0, multi.src[0].shape[multi.axis]) if a == multi.axis else s for a,s in enumerate(root.arg))).multi(multi.axis)
+    return multi.src[0].copy_to_device(multi.device, arg=multi.bounds.index(root.marg[multi.axis]))
+  return multi.src[0].shrink(tuple((0, multi.src[0].shape[multi.axis]) if a == multi.axis else s for a,s in enumerate(root.marg))).multi(multi.axis)
 
 def flip_multi(root:UOp, multi:UOp):
-  assert multi.axis is None or not root.arg[multi.axis], "flipping not supported on sharded axis"
-  return multi.src[0].flip(root.arg).multi(multi.axis)
+  assert multi.axis is None or not root.marg[multi.axis], "flipping not supported on sharded axis"
+  return multi.src[0].flip(root.marg).multi(multi.axis)
 
 # from multiple devices -> one
 def copy_multi(multi:UOp, device:UOp):
@@ -203,11 +204,11 @@ def passthrough_multi(root:UOp, multi:UOp):
 multi_pm = PatternMatcher([
   (UPat(GroupOp.ALU, name="root", custom_early_reject=set([Ops.MULTI])), alu_multi),
   (UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.MULTI, name="multi"), ), name="root"), reduce_multi),
-  (UPat(Ops.RESHAPE, src=(UPat(Ops.MULTI, name="multi"), ), name="root"), reshape_multi),
-  (UPat(Ops.EXPAND, src=(UPat(Ops.MULTI, name="multi"), ), name="root"), expand_multi),
-  (UPat(Ops.PAD, src=(UPat(Ops.MULTI, name="multi"), ), name="root"), pad_multi),
+  (UPat(Ops.RESHAPE, src=(UPat(Ops.MULTI, name="multi"), UPat()), name="root"), reshape_multi),
+  (UPat(Ops.EXPAND, src=(UPat(Ops.MULTI, name="multi"), UPat()), name="root"), expand_multi),
+  (UPat(Ops.PAD, src=(UPat(Ops.MULTI, name="multi"), UPat(), UPat()), name="root"), pad_multi),
+  (UPat(Ops.SHRINK, src=(UPat(Ops.MULTI, name="multi"), UPat(), UPat()), name="root"), shrink_multi),
   (UPat(Ops.PERMUTE, src=(UPat(Ops.MULTI, name="multi"), ), name="root"), permute_multi),
-  (UPat(Ops.SHRINK, src=(UPat(Ops.MULTI, name="multi"), ), name="root"), shrink_multi),
   (UPat(Ops.FLIP, src=(UPat(Ops.MULTI, name="multi"), ), name="root"), flip_multi),
   (UPat(Ops.ASSIGN, src=(UPat(Ops.MULTI, name="dest"), UPat(Ops.MULTI, name="src"))), assign_multi),
   (UPat(Ops.COPY, src=(UPat(Ops.MULTI, name="multi"), UPat(Ops.DEVICE, name="device"))), copy_multi),

tinygrad/schedule/rangeify.py

@@ -4,7 +4,7 @@ from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, ssimplify, KernelInfo
 from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType
 from tinygrad.uop.symbolic import symbolic_simple
-from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, flatten, dedup, unwrap, all_int, DEBUG, SPLIT_REDUCEOP, Metadata
+from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, Metadata
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_unparented
 from tinygrad.codegen.opt import Opt
 from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, ALWAYS_CONTIGUOUS, IndexingContext, apply_movement_op
@@ -13,6 +13,12 @@ from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, ALWAYS_CONTI
 import sys
 sys.setrecursionlimit(10000)
 
+# movement op on INDEX as a PatternMatcher
+pm_mops = PatternMatcher([
+  (UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
+   lambda r,idx: r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idx.src[1:]), dtype=idx.dtype, arg=idx.arg)),  # type: ignore
+])
+
 # *****************
 # 0. do some cleanup rewrites, mostly copied from the old stuff
 
@@ -31,7 +37,12 @@ def split_reduceop(reduce:UOp, x:UOp):
   #   ~2**10 should be enough if GROUP is used
   # 256 split maximum should be "negligible reduce" for low prod(reduce.shape), 8 split minimum.
   # split is moved to the end to provide maximum locality for the second phase reduce.
-  is_expanded = unwrap(x.st).is_expanded()
+
+  # get expanded by rangeifying the UOp x
+  indexed = x.index(*[UOp.range(s, i) if resolve(s>1) else UOp.const(dtypes.index, 0) for i,s in enumerate(x.shape)])
+  range_nums = [y.arg[0] for y in indexed.substitute({x.base:UOp(Ops.NOOP)}, extra_pm=pm_mops).ranges]
+  is_expanded = [i not in range_nums for i in range(len(x.shape))]
+
   if not (split_candidates:=[(i,d) for i in reduce.arg[1] for d in range(min(256,2**getenv("REDUCEOP_SPLIT_SIZE",22)//prod(reduce.shape)),8-1,-1)
                              if x.shape[i]%d==0 and not is_expanded[i]]): return None
   dim_to_split, divisor = split_candidates[0]
@@ -41,15 +52,18 @@ def split_reduceop(reduce:UOp, x:UOp):
   # reduce original axes, then split
   return splitted.r(*reduce.arg).contiguous().r(reduce.arg[0], (len(reduce.shape),)).reshape(reduce.shape).replace(tag=reduce.tag)
 
-earliest_rewrites = PatternMatcher([
+mop_cleanup = PatternMatcher([
+  # merge adjacent RESHAPES, safe because they are not tagged
+  (UPat(Ops.RESHAPE, name="x2").f(Ops.RESHAPE, allow_any_len=True, name="x"),
+   lambda x,x2: x.replace(src=(x2.src[0], x.src[1])) if x.tag is None and x2.tag is None else None),
+])
+
+earliest_rewrites = mop_cleanup+PatternMatcher([
   # just removing it works...
   (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE), name="x"), lambda x: x.src[0]),
 
-  # merge adjacent RESHAPES, safe because they are not tagged
-  (UPat(Ops.RESHAPE, name="x2").f(Ops.RESHAPE, name="x"), lambda x,x2: x.replace(src=(x2.src[0],)) if x.tag is None and x2.tag is None else None),
-
   # remove CONTIGUOUS if the BUFFER is already contiguous
-  (UPat(Ops.BUFFER).f(Ops.RESHAPE, name="r").f(Ops.CONTIGUOUS, name="c"), lambda r,c: r.replace(tag=c.tag)),
+  (UPat(Ops.BUFFER).f(Ops.RESHAPE, allow_any_len=True, name="r").f(Ops.CONTIGUOUS, name="c"), lambda r,c: r.replace(tag=c.tag)),
 
   # split_reduceop
   (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)), split_reduceop),
@@ -60,7 +74,7 @@ earliest_rewrites = PatternMatcher([
    lambda reduce,x: reduce.const_like(identity_element(reduce.arg[0], reduce.dtype)) if x.size == 0 and reduce.size != 0 else None),
 
   # handle size 0
-  (UPat(GroupOp.All-{Ops.SINK}, name="x"), lambda x: x.const_like(0).rtag(x.tag) if x.st is not None and x.size == 0 else None),
+  (UPat(GroupOp.All-{Ops.SINK}, name="x"), lambda x: x.const_like(0).rtag(x.tag) if x._shape is not None and x.size == 0 else None),
 
   # remove contiguous on movement ops before a copy on disk
   (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.CONTIGUOUS).f(Ops.COPY, allow_any_len=True, name="copy"),
@@ -94,15 +108,6 @@ earliest_rewrites = PatternMatcher([
    (UPat(Ops.ASSIGN, src=(UPat.var("a"), UPat.var("b")), name="assign"), find_permutes),
 ])
 
-# *****************
-# 3a. rangeify (movement)
-
-# movement op on INDEX as a PatternMatcher
-pm_mops = PatternMatcher([
-  (UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
-   lambda r,idx: r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.arg, idx.src[1:]), dtype=idx.dtype, arg=idx.arg)),  # type: ignore
-])
-
 # *****************
 # 3.5 cleanups
 
@@ -271,7 +276,7 @@ def bufferize_to_store(x:UOp):
     mops = []
     walk = assign_mops
     while walk is not assign_mops.base:
-      mops.append((walk.op, walk.arg))
+      mops.append((walk.op, walk.marg))
       walk = walk.src[0]
     for m in mops[::-1]: ret = ret._mop(*m)
     return ret.forced_reshape(shape).replace(tag=x.tag)
@@ -293,14 +298,14 @@ def bufferize_to_store(x:UOp):
   buf = UOp(Ops.DEFINE_LOCAL, sdtype, arg=tag)
   # store has the other dtype here
   # TODO: how is this unified?
-  return buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=sdtype).forced_reshape(shape, dtype=x.dtype)
+  return buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=sdtype).reshape(shape)
 
 pm_add_buffers = pm_mops+to_bufferview+PatternMatcher([
   (UPat(Ops.BUFFERIZE, name="x"), bufferize_to_store),
 
   # move RESHAPEs through MSELECT/MSTACK
   (UPat((Ops.MSELECT, Ops.MSTACK), src=UPat(Ops.RESHAPE), name="m"),
-   lambda m: m.replace(src=tuple([x.src[0].base for x in m.src]), tag=None).reshape(m.src[0].arg).rtag(m.tag)),
+   lambda m: m.replace(src=tuple([x.src[0].base for x in m.src]), tag=None).reshape(m.shape).rtag(m.tag)),
 ])
 
 # *****************
@@ -434,6 +439,7 @@ split_kernels = PatternMatcher([
 
 def tag_uop(ctx:list[UOp], x:UOp):
   if x.tag is not None: return None
+  if x.dtype.scalar() == dtypes.index: return None
   ctx.append(x)
   return x.replace(tag=(len(ctx)-1,))
 add_tags = PatternMatcher([
@@ -449,7 +455,7 @@ add_tags = PatternMatcher([
 def found_contiguous(ctx:dict[UOp, UOp], contig:UOp, src:UOp):
   x = src
   while x is not src.base:
-    if x.op is Ops.PERMUTE: contig = contig.permute(argsort(x.arg))
+    if x.op is Ops.PERMUTE: contig = contig.permute(argsort(x.marg))
     elif x.op is Ops.RESHAPE: contig = contig.reshape(x.src[0].shape)
     else: return None
     x = x.src[0]

tinygrad/shape/shapetracker.py

@@ -1,81 +0,0 @@
-# ShapeTracker allows movement operations to a buffer that don't require a copy to be made.
-from __future__ import annotations
-from dataclasses import dataclass
-import functools
-from typing import Callable
-from tinygrad.helpers import merge_dicts, getenv
-from tinygrad.shape.view import View, unravel
-from tinygrad.uop.symbolic import sym
-from tinygrad.uop.ops import UOp, Ops, graph_rewrite, Variable, sint, sint_to_uop, Context
-
-@functools.cache
-def views_to_valid_uop(views: tuple[View, ...], _idxs:tuple[UOp, ...]|None=None) -> UOp:
-  idx = views[-1].to_valid_uop(_idxs)
-  for view in reversed(views[0:-1]):
-    idx = view.to_valid_uop([sint_to_uop(i) for i in unravel(view.shape, idx)])
-  with Context(TRACK_MATCH_STATS=0):
-    return graph_rewrite(idx, sym, name="indexing sym @ 1")
-
-@functools.cache
-def views_to_is_expanded(views: tuple[View, ...]) -> tuple[bool, ...]:
-  # NOTE: return if each dim is expanded
-  if len(views) == 1 and views[-1].mask is None: return tuple([bool(st==0) for st in views[-1].strides])
-  idx = views_to_valid_uop(views).get_idx()
-  used_ranges = [x.arg[0] for x in idx.toposort() if x.op is Ops.RANGE]
-  return tuple([i not in used_ranges for i in range(len(views[-1].shape))])
-
-@dataclass(frozen=True, order=True)
-class ShapeTracker:
-  views: tuple[View, ...]
-
-  def __add__(self, st:ShapeTracker) -> ShapeTracker:
-    ret = self
-    for v in st.views: ret = ShapeTracker(ret.views + (v,)).simplify() # one view at a time = better simplification
-    return ret
-
-  @staticmethod
-  def from_shape(shape:tuple[sint, ...], strides:tuple[sint, ...]|None=None) -> ShapeTracker: return ShapeTracker((View.create(shape, strides),))
-
-  @property
-  def contiguous(self) -> bool: return len(self.views) == 1 and self.views[0].contiguous
-
-  @property
-  def shape(self) -> tuple[sint, ...]: return self.views[-1].shape
-
-  @property
-  def size(self) -> int: return self.views[-1].size()
-
-  def vars(self) -> set[Variable]: return set().union(*[v.vars() for v in self.views])
-
-  @property
-  def var_vals(self) -> dict[str, int]: return merge_dicts([{(vu:=v.unbind())[0].expr:vu[1]} for v in self.vars()])
-
-  def unbind(self) -> tuple[ShapeTracker, dict[Variable, int]]:
-    unbound_views, var_vals = zip(*[v.unbind() for v in self.views])
-    if all(len(x) == 0 for x in var_vals): return self, {}
-    return ShapeTracker(tuple(unbound_views)), merge_dicts(var_vals)
-
-  def is_expanded(self) -> tuple[bool, ...]:
-    with Context(TRACK_MATCH_STATS=0): return views_to_is_expanded(self.views)
-
-  def simplify(self) -> ShapeTracker:
-    if len(self.views) >= 2 and (new_view := self.views[-2] + self.views[-1]) is not None:
-      return ShapeTracker(self.views[:-2] + (new_view,)).simplify()
-    return self
-
-  # *** under this line are the movement ops ***
-
-  def pad(self, arg: tuple[tuple[sint, sint], ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].pad(arg), ))
-  def shrink(self, arg: tuple[tuple[sint, sint], ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].shrink(arg), ))
-  def expand(self, new_shape: tuple[sint, ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].expand(new_shape), ))
-  def permute(self, axis: tuple[int, ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].permute(axis), ))
-  def flip(self, mul: tuple[int, ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].flip(mul), ))
-
-  def reshape(self, new_shape: tuple[sint, ...]) -> ShapeTracker:
-    if getenv("MERGE_VIEW", 1) and (new_view := self.views[-1].reshape(new_shape)) is not None: return ShapeTracker(self.views[0:-1] + (new_view,))
-    return ShapeTracker(self.views + (View.create(new_shape), ))
-
-  def mop(self, op, arg): return mops[op](self, arg)
-
-mops: dict[Ops, Callable] = {Ops.RESHAPE: ShapeTracker.reshape, Ops.PERMUTE: ShapeTracker.permute, Ops.EXPAND: ShapeTracker.expand,
-                             Ops.SHRINK: ShapeTracker.shrink, Ops.FLIP: ShapeTracker.flip, Ops.PAD: ShapeTracker.pad}

tinygrad/shape/view.py

@@ -1,261 +0,0 @@
-from __future__ import annotations
-import functools, operator, itertools
-from dataclasses import dataclass
-from typing import cast, Sequence
-from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import resolve, UOp, Variable, sint, smax, smin, sint_to_uop, Ops, ssimplify
-from tinygrad.helpers import prod, all_int, flatten
-
-@functools.cache
-def canonicalize_strides(shape:tuple[sint, ...], strides:tuple[sint, ...]) -> tuple[sint, ...]:
-  return tuple(0 if s == 1 else st for s, st in zip(shape, strides))
-
-@functools.cache
-def strides_for_shape(shape:tuple[sint, ...]) -> tuple[sint, ...]:
-  if not shape: return ()
-  strides = tuple(itertools.accumulate(reversed(shape[1:]), operator.mul, initial=1))[::-1]
-  return canonicalize_strides(shape, strides)
-
-@functools.cache
-def merge_dims(shape:tuple[int, ...], strides:tuple[int, ...], mask:tuple[tuple[int, int], ...]|None=None) -> tuple[tuple[int, int, int], ...]:
-  # merge contiguous sub-parts or zero strided dims
-  # any stride 0, masked from dim=1, or contiguous part is merged into next dim.
-  # stride != 0 to stride == 0 starts a new merging block
-  # ret = tuple[(merged_size, stride, merged size w/o zero stride), ...]
-  if not shape: return ()
-  assert len(shape) == len(strides) and (mask is None or len(shape) == len(mask))
-  ret = [(shape[0], strides[0], shape[0] if strides[0] != 0 else 0)]
-  # merge this dim to next dim if size is 1
-  merging = (mask[0][1] - mask[0][0] == 1) if mask is not None else shape[0] == 1
-  for i, (s, st) in enumerate(zip(shape[1:], strides[1:]), start=1):
-    # always merge 1
-    if s == 1: continue
-    last_s, last_st, last_pre_expand_s = ret[-1]
-    # merge last dim with this dim if merging or strides matched
-    if merging or last_st == s * st: ret[-1] = (last_s * s, st, (s if merging else last_pre_expand_s * s))
-    else: ret.append((s, st, s))
-    # merge this dim to next dim if size is 1
-    merging = (mask[i][1] - mask[i][0] == 1) if mask is not None else s == 1
-  return tuple(ret)
-
-@functools.cache
-def _reshape_mask(_mask:tuple[tuple[sint, sint], ...]|None, old_shape:tuple[sint, ...], new_shape:tuple[sint, ...]) \
-  -> tuple[tuple[sint, sint], ...]|None:
-  """Returns the new mask if reshape is possible, and None if not possible."""
-  if _mask is None: return tuple((0, s) for s in new_shape)
-  if not all_int(flatten(_mask)): return None
-
-  new_mask: list[tuple[int, int]] = []
-  # _mask is all int here
-  r_masks, r_shape, r_new_shape = reversed(cast(tuple[tuple[int, int], ...], _mask)), reversed(old_shape), reversed(new_shape)
-  curr_stride, old_dim, new_dim, mask = 1, next(r_shape, 1), next(r_new_shape, 1), next(r_masks, (0,1))
-
-  while len(new_mask) < len(new_shape):
-    (l, r), next_stride = mask, ssimplify(new_dim * curr_stride)
-
-    # need to split mask
-    if old_dim == next_stride: # simply copy the mask and get next batch for merging
-      new_mask.append((l // curr_stride, (r - 1) // curr_stride + 1))
-      curr_stride, old_dim, new_dim, mask = 1, next(r_shape, 1), next(r_new_shape, 1), next(r_masks, (0,1))
-    elif old_dim > next_stride: # mask can only be splitted if reshape doesn't cut across the mask.
-      if old_dim % next_stride != 0: return None
-      if (l % next_stride != 0 or r % next_stride != 0) and l // next_stride != (r - 1) // next_stride: return None
-      new_mask.append((l % next_stride // curr_stride, (r - 1) % next_stride // curr_stride + 1))
-      curr_stride, new_dim = next_stride,  next(r_new_shape, 1) # need to get mask for next dimension
-    else:
-      next_mask = next(r_masks, (0, 1))
-      # combine if the mask can unfold continuously
-      if mask != (0, old_dim) and l != r and next_mask[1] - next_mask[0] != 1: return None
-      mask, old_dim = (next_mask[0] * old_dim + l, (next_mask[1] - 1) * old_dim + r), ssimplify(old_dim * next(r_shape, 1))
-
-  return tuple(reversed(new_mask))
-
-def unravel(shape:tuple[sint, ...], offset:sint) -> list[sint]:
-  # find the position of offset on each dimension based on shape
-  # similar to unravel_index in numpy/torch
-  acc, idxs = 1, []
-  for d in reversed(shape):
-    idxs.append((offset//acc)%d)
-    acc *= d
-  return idxs[::-1]
-
-@dataclass(frozen=True)
-class View:
-  shape:tuple[sint, ...]
-  strides:tuple[sint, ...]
-  offset:sint
-  mask:tuple[tuple[sint, sint], ...]|None
-  contiguous:bool
-
-  def to_valid_uop(self, idxs:Sequence[UOp]|None=None) -> UOp:
-    """valid.where(idx, INVALID)"""
-    if idxs is None: idxs = [UOp.range(s, i) for i,s in enumerate(self.shape)]
-    iexpr = sint_to_uop(self.offset)
-    where = UOp.const(dtypes.bool, True)
-    for idx,sh,st,m in zip(idxs, self.shape, self.strides, self.mask if self.mask is not None else itertools.repeat(None)):
-      iexpr = iexpr + idx*sint_to_uop(st)
-      if m is not None:
-        if resolve(m[0] != 0): where &= (idx >= sint_to_uop(m[0]))
-        if resolve(m[1] != sh): where &= (idx < sint_to_uop(m[1]))
-    return where.where(iexpr, UOp.invalid())
-
-  @functools.cache  # pylint: disable=method-cache-max-size-none
-  def size(self) -> int:
-    ret = prod([x.vmax if isinstance(x, UOp) else x for x in self.shape])
-    assert isinstance(ret, int), f"{ret=} is not int"
-    return ret
-
-  @staticmethod
-  @functools.cache
-  def create(shape:tuple[sint, ...], strides:tuple[sint, ...]|None=None, offset:sint=0, mask:tuple[tuple[sint, sint], ...]|None=None):
-    # TODO: resolve shouldn't be needed here
-    if not all(resolve(s >= 0) for s in shape): raise ValueError(f"Trying to create View with negative dimension: {shape=}")
-    strides = canonicalize_strides(shape, strides) if strides else strides_for_shape(shape)
-    # canonicalize 0 in shape
-    if 0 in shape: return View(shape, (0,) * len(shape), offset=0, mask=None, contiguous=True)
-    # canonicalize no-op mask
-    if mask is not None and all(m == (0,s) for m,s in zip(mask, shape)): mask = None
-    # if any dimension has size >1, but is masked such that only one index in the dimension is unmasked
-    # then its stride can also be set to 0, albeit with a corresponding adjustment required to the offset
-    if mask and any(elim := [not resolve(b+1 < e) for b,e in mask]):
-      if any(not resolve(b < e) for b,e in mask):
-        strides, offset, mask = (0,) * len(shape), 0, ((0,0),) * len(shape)
-      offset += sum((strides[i] * mask[i][0]) if e else 0 for i, e in enumerate(elim))
-      strides = tuple(0 if e else st for st,e in zip(strides, elim))
-    # simplify as we go
-    if isinstance(offset, UOp): offset = cast(sint, offset.ssimplify())
-    shape = tuple(cast(sint, x.ssimplify()) if isinstance(x, UOp) else x for x in shape)
-    # TODO: enabling stride simplification breaks symbolic jit
-    """
-    strides = tuple(x.ssimplify() if isinstance(x, UOp) else x for x in strides)
-    if mask: mask = tuple((s.ssimplify() if isinstance(s, UOp) else s, e.ssimplify() if isinstance(e, UOp) else e) for s,e in mask)
-    """
-    contiguous = offset == 0 and mask is None and strides == strides_for_shape(shape)
-    return View(shape, strides, offset, mask, contiguous)
-
-  @functools.cache  # pylint: disable=method-cache-max-size-none
-  def vars(self) -> set[Variable]:
-    flatten_mask = tuple(x for m in self.mask for x in m) if self.mask is not None else tuple()
-    return functools.reduce(operator.or_, [x.vars() for x in self.shape+self.strides+(self.offset,)+flatten_mask if isinstance(x, UOp)], set())
-
-  @functools.cache  # pylint: disable=method-cache-max-size-none
-  def unbind(self) -> tuple[View, dict[Variable, int]]:
-    var_unboundvar_val = [(v, v.unbind()) for v in self.vars() if v.op is Ops.BIND]
-    unbound_vars = {v:uv for v,(uv,_) in var_unboundvar_val}
-    return self.substitute(unbound_vars), dict(x[1] for x in var_unboundvar_val)
-
-  def substitute(self, dvars:dict[UOp, UOp]):
-    def _substitute(x:sint): return x if isinstance(x, int) else x.substitute(dvars)
-    new_shape = tuple(map(_substitute, self.shape))
-    new_strides = tuple(map(_substitute, self.strides))
-    new_offset = _substitute(self.offset)
-    new_mask = tuple((_substitute(x[0]), _substitute(x[1])) for x in self.mask) if self.mask is not None else None
-    return View.create(new_shape, new_strides, new_offset, new_mask)
-
-  @functools.cache  # pylint: disable=method-cache-max-size-none
-  def __add__(self, vm1:View) -> View|None:
-    vm2 = self
-    if vm2.contiguous or vm1.size() == 0: return vm1
-    if vm1.contiguous and vm1.shape == vm2.shape: return vm2
-    if vm1.contiguous and vm1.size() == vm2.size() and (ret := vm2.reshape(vm1.shape)) is not None: return ret
-    if vm1.mask:
-      if (new_vm1 := vm1.shrink(vm1.mask)) == vm1 or (merged := vm2 + new_vm1) is None: return None
-      return merged.pad(tuple((b,s-e) for (b,e),s in zip(vm1.mask, vm1.shape)))
-    if not all_int(vm1.shape):
-      # if all strides are 0 and vm2 is unmasked, return vm1
-      if all(x == 0 for x in vm2.strides+vm1.strides) and vm2.mask is None: return vm1
-      return None
-
-    # Project vm1's offset and strides on to vm2.
-    origin = [ssimplify(o) for o in unravel(vm2.shape, vm1.offset)]
-    terms: list[list[tuple[int, sint]]] = [[] for _ in vm2.shape]
-    strides: list[sint] = [0] * len(vm1.shape)
-    for d1, st in enumerate(vm1.strides):
-      if st == 0: continue
-      for d2, (o, s1) in enumerate(zip(origin, unravel(vm2.shape, vm1.offset + st))):
-        if not resolve((s1 := s1 - o)!=0): continue  # if s1 can possibly be 0
-        terms[d2].append((d1, s1))
-        strides[d1] += ssimplify(s1 * vm2.strides[d2])
-    return None
-
-  def __unsafe_resize(self, arg: tuple[tuple[sint, sint], ...], mask=None) -> View:
-    offset = sum([s * x[0] for s, x in zip(self.strides,arg)])
-    if self.mask:
-      # move the old mask
-      nmask = tuple([(smax(0, smin(mx-ax,ay-ax)), smax(0, smin(my-ax,ay-ax))) for (mx,my),(ax,ay) in zip(self.mask, arg)])
-      # merge the masks if we have two
-      mask = tuple([(smax(mx1, mx2), smin(my1, my2)) for (mx1, my1), (mx2, my2) in zip(nmask, mask)]) if mask is not None else nmask
-    return View.create(tuple([y-x for x,y in arg]), self.strides, self.offset+offset, mask)
-
-  @functools.cache  # pylint: disable=method-cache-max-size-none
-  def pad(self, arg: tuple[tuple[sint, sint], ...]) -> View:
-    assert len(arg) == len(self.shape), f"invalid pad {arg} for {self.shape}"
-    # NOTE: not checking for symbolic arg
-    for b,e in arg: assert not all_int([b,e]) or b>=0 and e>=0, f"invalid pad {arg} for {self.shape}"
-    if any(resolve(b!=0) or resolve(e!=0) for b, e in arg):
-      zvarg = tuple([(-b,s+e) for s,(b,e) in zip(self.shape, arg)])
-      mask = tuple([(b,s+b) for s,(b,_) in zip(self.shape, arg)])
-      return self.__unsafe_resize(zvarg, mask=mask)
-    return self
-
-  @functools.cache  # pylint: disable=method-cache-max-size-none
-  def shrink(self, arg: tuple[tuple[sint, sint], ...]) -> View:
-    assert len(arg) == len(self.shape), f"invalid shrink {arg} for {self.shape}"
-    # NOTE: not checking for symbolic arg
-    for s,(b,e) in zip(self.shape,arg): assert not all_int([s,b,e]) or (0<=b<=e<=s), f"invalid shrink {arg} for {self.shape}"
-    return self.__unsafe_resize(arg)
-
-  @functools.cache  # pylint: disable=method-cache-max-size-none
-  def expand(self, new_shape: tuple[sint, ...]) -> View:
-    if len(new_shape) != len(self.shape): raise ValueError(f"expand arg {new_shape=} must have same number of dimensions as shape {self.shape=}")
-    # NOTE: does not check multiple of symbolic shape
-    assert all(resolve(s == ns) or s == 1 for s,ns in zip(self.shape, new_shape)), f"can't expand {self.shape} into {new_shape}"
-    if 0 in self.shape: return View.create(new_shape)
-    # TODO: resolve may not be needed, but it's hard because vars need to be canonicalized
-    mask = tuple([(((0,0) if m != (0,1) else (0,ns)) if resolve(s != ns) and resolve(s == 1, False) else m) \
-                  for m,s,ns in zip(self.mask, self.shape, new_shape)]) if self.mask else None
-    return View.create(new_shape, self.strides, self.offset, mask)
-
-  @functools.cache  # pylint: disable=method-cache-max-size-none
-  def permute(self, axis: tuple[int, ...]) -> View:
-    assert sorted(axis) == list(range(len(self.shape))), f"invalid permutation {axis} of len {len(self.shape)}"
-    return View.create(tuple(self.shape[a] for a in axis), tuple(self.strides[a] for a in axis), self.offset,
-                       tuple(self.mask[a] for a in axis) if self.mask is not None else None)
-
-  @functools.cache  # pylint: disable=method-cache-max-size-none
-  def flip(self, arg: tuple[bool, ...]) -> View:
-    offset = sum((s-1)*z for s,z,f in zip(self.shape, self.strides, arg) if f)
-    mask = tuple((s-my,s-mx) if f else (mx,my) for (mx,my),s,f in zip(self.mask, self.shape, arg)) if self.mask is not None else None
-    return View.create(self.shape, tuple(-z if f else z for z,f in zip(self.strides, arg)), self.offset+offset, mask)
-
-  @functools.cache  # pylint: disable=method-cache-max-size-none
-  def reshape(self, new_shape: tuple[sint, ...]) -> View|None:
-    if self.shape == new_shape: return self
-
-    if not all(x >= 0 for x in new_shape): raise ValueError(f"shape can't contain negative numbers {new_shape}")
-    # check for the same size
-    if resolve(prod(self.shape) != prod(new_shape), True): raise ValueError(f"size mismatched, can't reshape {self.shape=} -> {new_shape=}")
-
-    if 0 in self.shape: return View.create(new_shape)
-    if new_shape == () and self.mask and any(mx==my for (mx,my) in self.mask): return None
-
-    # after the asserts, it's okay to check contiguous
-    if self.contiguous: return View.create(new_shape)
-
-    r_strides, r_new_shape = [], reversed(new_shape)
-    for merged_size, new_stride, real_size in reversed(merge_dims(self.shape, self.strides, self.mask)):
-      acc = 1
-      # TODO: third resolve shouldn't be needed
-      while resolve(acc <= merged_size) and resolve(acc != merged_size) and resolve((new_dim := next(r_new_shape, 0)) > 0):
-        r_strides.append(new_stride * acc)
-        acc = acc * new_dim
-        if not resolve(acc < real_size): new_stride = 0
-      if resolve(acc != merged_size): return None
-    new_strides = (0,) * (len(new_shape) - len(r_strides)) + tuple(r_strides[::-1])
-
-    if (new_mask:=_reshape_mask(self.mask, self.shape, new_shape)) is not None:
-      extra_offset = (sum(m[0] * s for m,s in zip(self.mask, self.strides)) if self.mask else 0) - \
-                     (sum(m[0] * s for m,s in zip(new_mask, new_strides)))
-      return View.create(new_shape, new_strides, self.offset + extra_offset, new_mask)
-
-    return None

tinygrad/tensor.py

@@ -6,7 +6,7 @@ from typing import Callable, ClassVar, Sequence, cast, get_args, Literal, Suppor
 from tinygrad.dtype import DType, DTypeLike, dtypes, ImageDType, ConstType, least_upper_float, least_upper_dtype, sum_acc_dtype, to_dtype, truncate
 from tinygrad.dtype import _from_np_dtype, _to_np_dtype
 from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten, dedup
-from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, unwrap, DEBUG, is_numpy_ndarray, FUSE_ATTENTION
+from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, DEBUG, is_numpy_ndarray, FUSE_ATTENTION
 from tinygrad.helpers import suppress_finalizing
 from tinygrad.gradient import compute_gradient
 from tinygrad.uop.mathtraits import MathTrait
@@ -197,7 +197,7 @@ class Tensor(MathTrait):
 
   def __repr__(self):
     ld = self.uop
-    ld_repr = f"<UOp {ld.device} {ld.shape} {str(ld.dtype)[7:]} {ld.st if ld.base is not ld else (ld.op, ld.realized)}>"
+    ld_repr = f"<UOp {ld.device} {ld.shape} {str(ld.dtype)[7:]}>"
     return f"<Tensor {ld_repr} on {self.device} with grad {(self.grad.uop if self.grad is not None else None)!r}>"
 
   # Python has a non moving GC, so this should be okay
@@ -256,7 +256,7 @@ class Tensor(MathTrait):
     # create the schedule
     schedule, var_vals = create_schedule_with_vars(sink)
     schedule = memory_planner(schedule)
-    if DEBUG >= 1 and len(schedule) > 1: print(f"scheduled {len(schedule)} kernels in {(time.perf_counter()-st)*1000:.2f} ms")
+    if (DEBUG >= 1 and len(schedule) > 1) or DEBUG >= 3: print(f"scheduled {len(schedule)} kernels in {(time.perf_counter()-st)*1000:.2f} ms")
     return schedule, var_vals
 
   def schedule(self, *lst:Tensor) -> list[ScheduleItem]:
@@ -267,7 +267,8 @@ class Tensor(MathTrait):
 
   def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor:
     """Triggers the computation needed to create these Tensor(s)."""
-    run_schedule(*self.schedule_with_vars(*lst), do_update_stats=do_update_stats)
+    if len(to_realize:=[x for x in (self,)+lst if not x.uop.is_contiguous()]):
+      run_schedule(*Tensor.schedule_with_vars(*to_realize), do_update_stats=do_update_stats)
     return self
 
   def replace(self, x:Tensor, allow_shape_mismatch=False) -> Tensor:
@@ -1191,8 +1192,11 @@ class Tensor(MathTrait):
     return X.shrink(tuple((-min(pB,0), min(pA+s,s)) for (pB,pA),s in zip(pX, X.shape)))
 
   # convenience
-  def pad_to(self, shape, *args): return self.pad(tuple([(0, ns-s) for s,ns in itertools.zip_longest(self.shape, argfix(shape, *args))]))
-  def shrink_to(self, shape, *args): return self.shrink(tuple([(0, ns) for ns in argfix(shape, *args)]))
+  def pad_to(self, shape, *args):
+    if len(new_shape := argfix(shape, *args)) != self.ndim: raise ValueError(f"dim mismatch, cannot pad {self.shape} to {new_shape}")
+    return self.pad(tuple([None if ns is None else (0, ns-s) for s,ns in zip(self.shape, new_shape)]))
+  def shrink_to(self, shape, *args):
+    return self.shrink(tuple([None if ns is None else (0, ns) for ns in argfix(shape, *args)]))
 
   # ***** movement high level ops *****
 
@@ -1345,12 +1349,12 @@ class Tensor(MathTrait):
       self.realize()._getitem(indices).assign(v)
       return
     # NOTE: check that setitem target is valid first
-    if not unwrap(self.uop.st).contiguous: raise RuntimeError("setitem target needs to be contiguous")
     if isinstance(v, get_args(ConstType)): v = Tensor(v, device=self.device, dtype=self.dtype)
     if not isinstance(v, Tensor): raise TypeError(f"can't set a {type(v).__name__} to a Tensor")
     if self.requires_grad or v.requires_grad: raise NotImplementedError("setitem with requires_grad is not supported")
-
-    res = self.realize()._getitem(indices, v)
+    self.realize()
+    if not self.uop.is_contiguous(): raise RuntimeError("setitem target needs to be contiguous")
+    res = self._getitem(indices, v)
     # if shapes match and data is not shared it's a copy and we assign to self
     if res.shape == self.shape and res.uop is not self.uop:
       self.assign(res).realize()

tinygrad/uop/ops.py

@@ -10,7 +10,6 @@ from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Contex
 from tinygrad.helpers import PICKLE_BUFFERS, PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, VIZ, SPEC
 from tinygrad.helpers import strip_parens
 if TYPE_CHECKING:
-  from tinygrad.shape.shapetracker import ShapeTracker
   from tinygrad.device import Buffer, MultiBuffer
 
 class AxisType(Enum):
@@ -176,59 +175,97 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   # *** uop shape stuff ***
 
   @recursive_property
-  def st(self) -> ShapeTracker|None:
-    if self.op is Ops.INDEX and self.src[0].op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.MSTACK,
-                                                   Ops.MSELECT, Ops.BUFFER, Ops.BUFFERIZE, Ops.VECTORIZE, Ops.STORE}:
-      return None
-    if self.op is Ops.INDEX and self.src[0].op is Ops.ASSIGN and self.src[0].src[1].op is Ops.KERNEL: return None
-    if self.op is Ops.BARRIER: return None
-    if self.op in GroupOp.Block: return None
-    from tinygrad.shape.shapetracker import ShapeTracker
-    # MovementOps define a new ShapeTracker from the arg
-    if self.op is Ops.BUFFERIZE: return ShapeTracker.from_shape(tuple([int(r.vmax+1) for r in self.src[1:]]))
-    # allow reshape from nothing
-    if self.op is Ops.RESHAPE and self.src[0].st is None: return ShapeTracker.from_shape(self.arg)
-    if self.op in GroupOp.Movement: return unwrap(self.src[0].st).mop(self.op, self.arg)
-    # CONST with a DEVICE has a shape of ()
-    if self.op is Ops.CONST and len(self.src) and self.src[0].op is Ops.DEVICE: return ShapeTracker.from_shape(())
-    if self.op is Ops.STORE and isinstance(self.dtype, PtrDType): return ShapeTracker.from_shape((self.dtype.size,))
-    if self.op is Ops.STORE and self.dtype is not dtypes.void: return self.src[0].src[0].st
-    # BufferOps and ASSIGN flow ShapeTracker from a direct edge
-    if self.op in {Ops.STORE, Ops.ASSIGN, Ops.LOAD}: return self.src[0].st
-
-    # BUFFER/BUFFER_VIEW and KERNEL only have a size
-    if self.op in {Ops.BUFFER, Ops.BUFFER_VIEW}: return ShapeTracker.from_shape((self.size,))
-    if self.op is Ops.KERNEL:
-      ast = self.arg.ast
-      return ShapeTracker.from_shape((ast.size,)) if ast.st is not None else None
-    if self.op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}:
-      sz = self.ptrdtype.size
-      return ShapeTracker.from_shape((sz,)) if sz > 0 else None
-
-    # hack for PTX, CASTing the ptr loses the shape
-    if self.op is Ops.CAST and self.src[0].op is Ops.DEFINE_GLOBAL: return None
-
-    # otherwise we get the shape from sources
-    if not (src_sts := [x.st for x in self.src if x.st is not None]): return None
-    assert all_same([x.shape for x in src_sts]), f"UOp sources must have the same shape {self} {[x.shape for x in src_sts]}"
-    shape = src_sts[0].shape
-    # shape changing ops
+  def _shape(self) -> tuple[sint, ...]|None:
     match self.op:
-      case Ops.MULTI: shape = tuple(s*len(self.device) if a == self.axis else s for a,s in enumerate(shape))
+      # late ops don't have shape
+      case Ops.UNIQUE | Ops.DEVICE | Ops.RANGE | Ops.INDEX | Ops.LOAD | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
+           Ops.VECTORIZE | Ops.VCONST | Ops.SUBSTITUTE | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.PRECAST:
+        return None
+
+      # some ops init the shape
+      case Ops.CONST | Ops.DEFINE_VAR | Ops.BIND: return () if self._device is not None else None
+      case Ops.BUFFER: return (self.arg,)
+      case Ops.BUFFER_VIEW: return (self.arg[0],)
+      case Ops.BUFFERIZE: return tuple([int(r.vmax+1) for r in self.src[1:]])
+      case Ops.DEFINE_GLOBAL | Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return (self.ptrdtype.size,)
+
+      # passthrough ops
+      case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.FUSE: return self.src[0]._shape
+
+      # ops with custom handling
+      case Ops.KERNEL: return self.arg.ast._shape
+      case Ops.STORE:
+        if isinstance(self.dtype, PtrDType): return (self.ptrdtype.size,)
+        if self.dtype is not dtypes.void: return self.src[0].src[0].shape
+        return None
+
+      # TODO: disallow shape changing bitcast
       case Ops.BITCAST:
-        if (output_sz:=self.dtype.itemsize) != (input_sz:=self.src[0].dtype.itemsize): shape = shape[:-1]+((shape[-1]*input_sz) // output_sz,)
-      case Ops.REDUCE_AXIS | Ops.WMMA:
-        axis_arg = self.arg[1] if self.op is Ops.REDUCE_AXIS else self.arg[7]
-        assert isinstance(axis_arg, tuple) and all(isinstance(x, int) for x in axis_arg), f"invalid type for axis: {axis_arg}"
-        shape = tuple(1 if i in axis_arg else s for i,s in enumerate(shape))
-    return ShapeTracker.from_shape(shape)
+        ps = self.src[0]._shape
+        if ps is None: return None
+        if (output_sz:=self.dtype.itemsize) != (input_sz:=self.src[0].dtype.itemsize): return ps[:-1]+(ssimplify((ps[-1]*input_sz) // output_sz),)
+        return ps
+
+      # TODO: disallow reshape from nothing. tested by TestOpenClip.test_multigpu_clip_score
+      case Ops.RESHAPE:
+        if self.src[0]._shape is None: return self.marg
+
+    # movement ops change the shape. this is the logic from the old ShapeTracker
+    # NOTE: ssimplify is required because the shape needs to be canonical for broadcasting and same shape checking
+    if self.op in GroupOp.Movement.union({Ops.MULTI, Ops.REDUCE_AXIS, Ops.WMMA}):
+      ps = self.src[0]._shape
+      # TODO: WMMA is used for both axis WMMA and op WMMA. fix this and remove this hack. tested by BERT on AMD LLVM
+      if ps is None and self.op is Ops.WMMA: return None
+      if ps is None: raise RuntimeError(f"movement op {self.op} requires shape")
+      match self.op:
+        case Ops.RESHAPE:
+          if not all(x >= 0 for x in self.marg): raise ValueError(f"shape can't contain negative numbers {self.marg}")
+          if prod(ps) != prod(self.marg): raise ValueError(f"bad reshape: {ps} -> {self.marg}")
+          return self.marg
+        case Ops.EXPAND:
+          if len(ps) != len(self.marg) or not all(s==ns or (s==1 and ns>=0) for s,ns in zip(ps, self.marg)):
+            raise ValueError(f"bad expand: {ps} -> {self.marg}")
+          return self.marg
+        case Ops.PERMUTE:
+          if sorted(self.marg) != list(range(len(ps))): raise ValueError(f"invalid permutation {self.marg} of len {len(ps)}")
+          return tuple(ps[i] for i in self.marg)
+        case Ops.PAD:
+          # TODO: why do i need resolve here?
+          if len(ps) != len(self.marg) or not all(resolve(b>=0) and resolve(e>=0) for b,e in self.marg): raise ValueError(f"invalid pad {self.marg}")
+          return tuple(ssimplify(s+b+e) for s,(b,e) in zip(ps, self.marg))
+        case Ops.SHRINK:
+          # TODO: why do i need resolve here?
+          if len(ps) != len(self.marg) or not all(resolve(0<=b) and resolve(b<=e) and resolve(e<=s) for s,(b,e) in zip(ps, self.marg)):
+            raise ValueError(f"invalid shrink {self.marg} for {ps}")
+          return tuple(ssimplify(e-s) for s,e in self.marg)
+        case Ops.FLIP:
+          if len(ps) != len(self.marg) or not all(isinstance(x, bool) for x in self.marg): raise ValueError(f"bad flip on {ps}, {self.marg}")
+          return ps
+        case Ops.MULTI: return tuple(s*len(self.device) if a == self.axis else s for a,s in enumerate(ps))
+        case Ops.REDUCE_AXIS | Ops.WMMA:
+          axis_arg = self.arg[1] if self.op is Ops.REDUCE_AXIS else self.arg[7]
+          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))
+
+    # elementwise ops keep the shape the same. all inputs with shape must match
+    if self.op in (GroupOp.Elementwise-{Ops.BITCAST}).union({Ops.COPY, Ops.ASSIGN, Ops.NOOP, Ops.SINK, Ops.ALLREDUCE}):
+      # TODO: remove this hack for 3 op assign
+      input_shapes = [x._shape for x in (self.src[:2] if self.op is Ops.ASSIGN else self.src) if x._shape is not None]
+      if len(input_shapes) == 0: return None
+      if not all_same(input_shapes): raise RuntimeError(f"shape mismatch at {self.op}: {input_shapes}")
+      return input_shapes[0]
+
+    # all Ops must be explicitly handled
+    raise NotImplementedError(f"no shape handling for {self.op} with {self.dtype}")
 
   @property
   def shape(self) -> tuple[sint, ...]:
-    assert self.st is not None, f"{self.op} doesn't have a shape"
-    return unwrap(self.st).shape
+    if (ret:=self._shape) is None: raise RuntimeError(f"shape requested, but {self.op} doesn't have a shape")
+    return ret
+
   @property
-  def size(self) -> int: return self.arg[0] if self.op is Ops.BUFFER_VIEW else self.arg if self.op is Ops.BUFFER else unwrap(self.st).size
+  def size(self) -> int: return prod([int(x.vmax) if isinstance(x, UOp) else x for x in self.shape])
 
   # determine what ranges this is in
   @recursive_property
@@ -264,11 +301,11 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   def __bool__(self): return self._eval((dtypes.bool,), bool)
   def __int__(self): return self._eval(dtypes.ints, int)
   def __float__(self): return self._eval(dtypes.floats, float)
-  def substitute(self, dvars:dict[UOp, UOp], name:str|None=None):
+  def substitute(self, dvars:dict[UOp, UOp], name:str|None=None, extra_pm:PatternMatcher|None=None):
     dvars = {k:v for k,v in dvars.items() if k is not v}
     if len(dvars) == 0: return self
     with Context(TRACK_MATCH_STATS=(0 if name is None else TRACK_MATCH_STATS.value)):
-      return graph_rewrite(self, _substitute, dvars, bottom_up=True, name=name)
+      return graph_rewrite(self, (extra_pm+_substitute) if extra_pm is not None else _substitute, dvars, bottom_up=True, name=name)
 
   # *** uop tracing stuff ***
 
@@ -290,7 +327,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   def __getitem__(self, idx): return self.index(idx)
   def const_like(self, b:ConstLike):
     # constants can optionally have a DEVICE source
-    return UOp.const(self.dtype, b, device=self._device, shape=self.shape if self.st is not None else None)
+    return UOp.const(self.dtype, b, device=self._device, shape=self._shape)
   def broadcast(self, count:int):
     assert self.dtype.count == 1
     if count == 1: return self
@@ -352,7 +389,15 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     assert self.dtype.scalar() is dtypes.index, "Can only call get_valid on index dtype"
     return self.src[0] if self.op is Ops.WHERE and self.src[2].arg is Invalid else UOp.const(dtypes.bool, self.arg is not Invalid)
   def reduce(self, *src:UOp, **kwargs): return UOp(Ops.REDUCE, kwargs.pop('dtype', self.dtype), src=(self,)+src, **kwargs)
-  def contiguous(self, *args, **kwargs): return UOp(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs)
+
+  def is_contiguous(self):
+    # TODO: this is is_realized
+    if self.op is Ops.RESHAPE: return self.src[0].is_contiguous()
+    return self.op is Ops.BUFFER
+
+  def contiguous(self, *args, **kwargs):
+    if self.is_contiguous(): return self
+    return UOp(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs)
   def contiguous_backward(self): return self.alu(Ops.CONTIGUOUS_BACKWARD)
   def bufferize(self, *args, **kwargs): return UOp(Ops.BUFFERIZE, dtype=self.dtype, src=(self,)+args, **kwargs)
   def fuse(self): return self.alu(Ops.FUSE)
@@ -390,11 +435,11 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     if self.op is Ops.REDUCE_AXIS: return None if src_axis is not None and src_axis in self.arg[1] else src_axis
     if self.op is Ops.RESHAPE:
       if src_axis is None: return None
-      arg_acc:list[sint] = list(itertools.accumulate(self.arg, operator.mul, initial=1))
+      arg_acc:list[sint] = list(itertools.accumulate(self.marg, operator.mul, initial=1))
       # new_axis is the last one that preserves prod(prior to new_axis) and must not move items between shards
       # TODO: what to do about shrinking to self.shape[self.axis]==1 len(self.real_lbs)==1?
       return len(arg_acc) - arg_acc[::-1].index(prod(self.src[0].shape[:src_axis])) - 1
-    if self.op is Ops.PERMUTE: return self.arg.index(src_axis) if src_axis is not None else None
+    if self.op is Ops.PERMUTE: return self.marg.index(src_axis) if src_axis is not None else None
     return src_axis
 
   def _unshard(self, axis:int) -> UOp:
@@ -428,19 +473,48 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     if self.op is Ops.MULTI: return self.src[0].base  # MULTI is really a VIEW
     return self
 
-  def _mop(self, op:Ops, arg) -> UOp:
-    ret = UOp(op, self.dtype, (self,), arg)
-    if self.st == ret.st: return self  # ignore NOOPs, also check ret.st
+  # like gep, but might return an integer
+  def sgep(self, i:int) -> sint:
+    match self.op:
+      case Ops.CONST: return self.arg
+      case Ops.VCONST: return self.arg[i]
+      case Ops.VECTORIZE: return cast(sint, self.src[i].ssimplify())
+      case _: raise RuntimeError(f"no sgep on {self.op}")
+
+  @functools.cached_property
+  def marg(self):
+    match self.op:
+      case Ops.RESHAPE | Ops.EXPAND: return tuple(self.src[1].sgep(i) for i in range(self.src[1].dtype.count))
+      case Ops.PAD | Ops.SHRINK: return tuple((self.src[1].sgep(i), self.src[2].sgep(i)) for i in range(self.src[1].dtype.count))
+      case Ops.PERMUTE | Ops.FLIP: return self.arg
+      case _: raise RuntimeError(f"{self.op} is not a MovementOp")
+
+  def _mop(self, op:Ops, arg, same_shape_noop:bool=False) -> UOp:
+    match op:
+      case Ops.RESHAPE | Ops.EXPAND: src_args = [arg]
+      case Ops.PAD | Ops.SHRINK: src_args = list(zip(*arg))
+      case Ops.PERMUTE | Ops.FLIP: src_args = []
+      case _: raise RuntimeError(f"{op} is not a MovementOp")
+    usrcs = []
+    for arg in src_args:
+      if len(arg) == 0: usrcs.append(UOp(Ops.VECTORIZE, dtypes.index.vec(0)))
+      elif all(isinstance(x, int) for x in arg): usrcs.append(UOp.const(dtypes.index.vec(len(arg)), arg))
+      else: usrcs.append(UOp(Ops.VECTORIZE, dtypes.index.vec(len(arg)), tuple(UOp.const(dtypes.index, x) if isinstance(x, int) else x for x in arg)))
+    ret = UOp(op, self.dtype, (self,)+tuple(usrcs), arg if len(usrcs) == 0 else None)
+    # for all movement ops, we check shape property
+    if ret.shape == self.shape and same_shape_noop: return self
     return ret
 
-  def forced_reshape(self, arg:tuple[sint, ...], **kwargs): return UOp(Ops.RESHAPE, kwargs.pop("dtype", self.dtype), src=(self,), arg=arg)
+  # in these four, if the shape doesn't change we can return self
+  def forced_reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, same_shape_noop=False)
+  def reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, same_shape_noop=True)
+  def expand(self, arg:tuple[sint, ...]): return self._mop(Ops.EXPAND, arg, same_shape_noop=True)
+  def shrink(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.SHRINK, arg, same_shape_noop=True)
+  def pad(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.PAD, arg, same_shape_noop=True)
 
-  def reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg)
-  def expand(self, arg:tuple[sint, ...]): return self._mop(Ops.EXPAND, arg)
-  def shrink(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.SHRINK, arg)
-  def pad(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.PAD, arg)
-  def permute(self, arg:tuple[int, ...]): return self._mop(Ops.PERMUTE, arg)
-  def flip(self, arg:tuple[bool, ...]): return self._mop(Ops.FLIP, arg)
+  # in these two, we have custom logic to check if they are a no-op
+  def permute(self, arg:tuple[int, ...]): return self._mop(Ops.PERMUTE, arg, same_shape_noop=False) if arg != tuple(range(len(self.shape))) else self
+  def flip(self, arg:tuple[bool, ...]): return self._mop(Ops.FLIP, arg, same_shape_noop=False) if any(arg) and len(arg) == len(self.shape) else self
 
   # *** uop UNIQUE ***
 
@@ -531,6 +605,9 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   def unbind(self) -> tuple[Variable, int]:
     assert self.op is Ops.BIND and self.src[0].op is Ops.DEFINE_VAR and self.src[1].op is Ops.CONST, f"can't unbind {self}"
     return self.src[0], self.src[1].arg
+  def unbind_all(self) -> tuple[UOp, dict[Variable, int]]:
+    ret:dict[Variable, int] = {}
+    return graph_rewrite(self, pm_unbind, ctx=ret), ret
   @property
   def val(self) -> int: return self.unbind()[1]
   def vars(self) -> set[UOp]:
@@ -944,6 +1021,9 @@ class TrackedPatternMatcher(PatternMatcher):
       match_stats[p][2] += time.perf_counter()-st
     return None
 
+@dataclass(frozen=True)
+class RewriteTrace: keys:list[TracingKey]; rewrites:list[list[TrackedGraphRewrite]]; uop_fields:dict[int, tuple] # noqa: E702
+
 if TRACK_MATCH_STATS or PROFILE:
   PatternMatcher = TrackedPatternMatcher  # type: ignore
   import atexit
@@ -952,7 +1032,7 @@ if TRACK_MATCH_STATS or PROFILE:
     if TRACK_MATCH_STATS >= 2:
       with open(fn:=temp("rewrites.pkl", append_user=True), "wb") as f:
         print(f"rewrote {len(tracked_ctxs)} graphs and matched {sum(len(r.matches) for x in tracked_ctxs for r in x)} times, saved to {fn}")
-        pickle.dump([(tracked_keys, tracked_ctxs, uop_fields)], f)
+        pickle.dump(RewriteTrace(tracked_keys, tracked_ctxs, uop_fields), f)
     if VIZ: return launch_viz("VIZ", temp("rewrites.pkl", append_user=True))
     if getenv("PRINT_MATCH_STATS", TRACK_MATCH_STATS.value):
       ret = [0,0,0.0,0.0]
@@ -1101,6 +1181,12 @@ def _index_to_concrete_int(u:UOp): return graph_rewrite(u.sink(), pm_lower_index
 
 _substitute = PatternMatcher([(UPat(tuple(Ops), name="x"), lambda ctx,x: ctx.get(x,None))])
 
+def do_unbind(ctx:dict[Variable, int], x:UOp):
+  v,i = x.unbind()
+  ctx[v] = i
+  return v
+pm_unbind = PatternMatcher([(UPat(Ops.BIND, name="x"), do_unbind)])
+
 # for debug
 syms = { Ops.ADD: "+", Ops.SUB: "-", Ops.IDIV: "//", Ops.MOD: "%", Ops.SHL: "<<", Ops.SHR: ">>",
          Ops.MUL: "*", Ops.CMPLT: "<", Ops.CMPNE: "!=", Ops.AND: "&", Ops.OR: "|", Ops.XOR: "^"}

tinygrad/uop/spec.py

@@ -82,13 +82,13 @@ assign_spec = PatternMatcher([
 # *** this is the spec of a Tensor in UOp ***
 
 tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([
-  (UPat(GroupOp.Movement, name="mv", src=(UPat.var("x"),)),
-   # naturally correct
-   lambda mv,x: (isinstance(mv.arg, tuple) and mv.dtype == x.dtype) or
-   # "make things that can't be images not images" can change the buffer dtype
-   # this is fine as long as it's a realized buffer or const and base dtypes match.
-   ((isinstance(mv.dtype, ImageDType) or isinstance(x.dtype, ImageDType)) and x.dtype.base == mv.dtype.base \
-       and x.base.op in {Ops.BUFFER,Ops.ASSIGN,Ops.CONST})),
+  (UPat((Ops.RESHAPE, Ops.EXPAND), name="mv", src=(UPat.var("x"), UPat(dtype=dtypes.index))), lambda mv,x: True),
+  (UPat((Ops.PAD, Ops.SHRINK), name="mv", src=(UPat.var("x"), UPat(dtype=dtypes.index), UPat(dtype=dtypes.index))), lambda mv,x: True),
+  (UPat((Ops.PERMUTE, Ops.FLIP), name="mv", src=(UPat.var("x"),)), lambda mv,x: isinstance(mv.arg, tuple)),
+
+  # inputs to movement ops
+  (UPat((Ops.VECTORIZE, Ops.VCONST), dtype=dtypes.index), lambda: True),
+  (UPat({Ops.ADD, Ops.MUL, Ops.IDIV}, dtype=dtypes.index), lambda: True),
 
   # Tensor variable bindings
   (UPat(Ops.BIND, (dtypes.int,dtypes.index,), (UPat(Ops.DEFINE_VAR), UPat.cvar(dtype=(dtypes.int,dtypes.index,))), arg=None), lambda: True),

tinygrad/uop/symbolic.py

@@ -473,6 +473,7 @@ def drop_and_clauses(cond:UOp, x:UOp, i:UOp) -> UOp|None:
   if not (dropped_clauses:=[c for c in cond.split_uop(Ops.AND) if not any(r in x.ranges for r in c.ranges)]): return None
   return functools.reduce(operator.and_, [c for c in cond.split_uop(Ops.AND) if c not in dropped_clauses], UOp.const(dtypes.bool, True)).where(x, i)
 pm_drop_and_clauses = PatternMatcher([(UPat.var("cond").where(UPat.var("x", dtype=dtypes.index), invalid_pat), drop_and_clauses)])
+
 def where_on_load(l, c1, buf, x):
   c2 = x.get_valid()
   duplicate_clauses = [c for c in c1.split_uop(Ops.AND) if c in c2.split_uop(Ops.AND)]
@@ -484,6 +485,11 @@ def where_on_load(l, c1, buf, x):
   # aditionally we can drop the clause on the where if it already exists in the load
   remaining_clause = functools.reduce(operator.and_, [c for c in c1.split_uop(Ops.AND) if c not in removed], UOp.const(dtypes.bool, True))
   return remaining_clause.where(UOp.load(buf.index(x.get_idx().valid(functools.reduce(operator.and_, moved_clauses, c2)), *l.src[1:])), 0)
+pm_move_where_on_load = PatternMatcher([
+  (UPat.var("c1").where(UPat(Ops.LOAD, src=(UPat.var("buf").index(UPat.var("x")),), name="l"), 0), where_on_load),
+  (UPat.var("c1").where(0, UPat(Ops.LOAD, src=(UPat.var("buf").index(UPat.var("x")),), name="l")),
+    lambda l,c1,buf,x: where_on_load(l,c1.logical_not(),buf,x)),
+])
 
 pm_simplify_valid = PatternMatcher([
   # simplify valid
@@ -529,9 +535,6 @@ sym = symbolic_flat+pm_simplify_valid+PatternMatcher([
   (UPat((Ops.LOAD, Ops.STORE), src=(UPat().index(UPat.const(dtypes.index, Invalid)).or_casted(),), allow_any_len=True, name="x"),
     lambda x: UOp(Ops.NOOP) if x.op is Ops.STORE else x.const_like(0)), # invalid store does nothing. invalid load produces 0
   # # Where after gated load becomes alt value, TODO: this is sort of duplicated with rules in devectorizer
-  (UPat.var("c1").where(UPat(Ops.LOAD, src=(UPat.var("buf").index(UPat.var("x")),), name="l"), 0), where_on_load),
-  (UPat.var("c1").where(0, UPat(Ops.LOAD, src=(UPat.var("buf").index(UPat.var("x")),), name="l")),
-    lambda l,c1,buf,x: where_on_load(l,c1.logical_not(),buf,x)),
   # remove VECTORIZE from SINK/BARRIER. TODO: SINK/BARRIER are really the same thing at GLOBAL/LOCAL levels
   (UPat(Ops.BARRIER, name="root"),
     lambda root: UOp(Ops.BARRIER, root.dtype, tuple(flatten(x.src if x.op in REMOVE_FROM_BARRIER else (x,) for x in root.src)), root.arg)

tinygrad/viz/index.html

@@ -32,7 +32,10 @@
     font-size: 14px;
     overflow: hidden;
     background-color: #08090e;
+    scrollbar-color: #686977 #151720;
   }
+  ::-webkit-scrollbar-track { background: #151720; }
+  ::-webkit-scrollbar-thumb { background: #686977; }
   a {
     color: #4a90e2;
   }
@@ -142,7 +145,7 @@
     inset: 0;
     z-index: 1;
   }
-  .profiler, .disasm {
+  .profiler, .render {
     flex: 1 1 auto;
     min-width: 0;
     width: 100%;
@@ -332,7 +335,7 @@
     <div id="progress-message"></div>
     <div class="container ctx-list-parent"><div class="ctx-list"></div></div>
     <div class="view profiler"></div>
-    <div class="view disasm"></div>
+    <div class="view render"></div>
     <div class="view graph">
       <svg id="graph-svg" preserveAspectRatio="xMidYMid meet">
         <g id="render">

tinygrad/viz/js/index.js

@@ -259,6 +259,7 @@ async function renderProfiler() {
           x += 1; y += nbytes; valueMap.set(ts, y);
         } else {
           const free = buf_shapes.get(key);
+          free.users = Array.from({ length: u32() }, () => strings[u32()]);
           timestamps.push(ts); valueMap.set(ts, y);
           x += 1; y -= free.nbytes;
           free.x.push(x);
@@ -278,12 +279,21 @@ async function renderProfiler() {
       timestamps.push(dur);
       const height = heightScale(peak);
       const yscale = d3.scaleLinear().domain([0, peak]).range([height, 0]);
-      for (const [num, {dtype, sz, nbytes, y, x:steps}] of buf_shapes) {
+      for (const [num, {dtype, sz, nbytes, y, x:steps, users}] of buf_shapes) {
         const x = steps.map(s => timestamps[s]);
         const dur = x.at(-1)-x[0];
         const html = document.createElement("div");
         const rows = [["DType", dtype], ["Len", formatUnit(sz)], ["Size", formatUnit(nbytes, "B")], ["Lifetime", formatTime(dur)]];
+        if (users != null) rows.push(["Users", users.length]);
         const info = html.appendChild(tabulate(rows).node());
+        for (let u=0; u<users?.length; u++) {
+          const p = html.appendChild(document.createElement("p")); p.style.marginTop = "4px"; p.style.cursor = "pointer";
+          const name = users[u]; p.appendChild(colored(`[${u}] ${name}`));
+          p.onclick = () => {
+            const cid = ctxs.findIndex(c => c.name === name);
+            if (cid != null) setCtxWithHistory(cid-1);
+          }
+        }
         const arg = {tooltipText:info.outerHTML, html, key:`${k}-${num}`};
         shapes.push({ x, y0:y.map(yscale), y1:y.map(y0 => yscale(y0+nbytes)), arg, fillColor:cycleColors(colorScheme.BUFFER, shapes.length) });
       }
@@ -354,7 +364,7 @@ async function renderProfiler() {
           for (let i=x.length-1; i>=0; i--) p.lineTo(x[i], offsetY+e.y1[i]);
           p.closePath();
           ctx.fillStyle = e.fillColor; ctx.fill(p);
-          if (focusedShape && e.arg?.key === focusedShape.key) { paths.push(p); }
+          if (focusedShape?.key && e.arg?.key === focusedShape.key) { paths.push(p); }
           continue;
         }
         // contiguous rect
@@ -515,14 +525,6 @@ function appendTd(tr, value, unit=null) {
   tr.appendChild(document.createElement("td")).innerText = unit == "us" ? formatTime(value) : fmt+(unit ?? "");
 }
 
-function appendRow(table, name, value, unit=null, cls="main-row") {
-  const tr = table.appendChild(document.createElement("tr"));
-  tr.className = cls;
-  tr.appendChild(document.createElement("td")).innerText = name;
-  appendTd(tr, value, unit);
-  return tr;
-}
-
 function setActive(e) {
   if (e == null) return;
   e.classList.add("active");
@@ -630,9 +632,9 @@ async function main() {
     ret = cache[ckey];
   }
   // ** Disassembly view
-  if (ckey.startsWith("/disasm")) {
+  if (ckey.startsWith("/render")) {
     if (!(ckey in cache)) cache[ckey] = ret = await (await fetch(ckey)).json();
-    displayGraph("disasm");
+    displayGraph("render");
     const root = document.createElement("div");
     root.className = "raw-text";
     const metadata = document.querySelector(".metadata");
@@ -666,8 +668,8 @@ async function main() {
         const div = d3.create("div").style("background", cycleColors(colorScheme.CATEGORICAL, s.idx)).style("width", "24px").style("height", "100%");
         return [s.label.trim(), div.node()];
       })).node());
-    } else root.appendChild(codeBlock(ret.src, "x86asm"));
-    return document.querySelector(".disasm").replaceChildren(root);
+    } else root.appendChild(codeBlock(ret.src, ret.lang));
+    return document.querySelector(".render").replaceChildren(root);
   }
   // ** UOp view (default)
   // if we don't have a complete cache yet we start streaming rewrites in this step
@@ -691,8 +693,8 @@ async function main() {
   renderDag(ret[currentRewrite].graph, ret[currentRewrite].changed_nodes ?? [], currentRewrite === 0);
   // ** right sidebar code blocks
   const metadata = document.querySelector(".metadata");
-  const [code, lang] = ctx.fmt != null ? [ctx.fmt, "cpp"] : [ret[currentRewrite].uop, "python"];
-  metadata.replaceChildren(codeBlock(step.code_line, "python", { loc:step.loc, wrap:true }), codeBlock(code, lang, { wrap:false }));
+  metadata.replaceChildren(codeBlock(step.code_line, "python", { loc:step.loc, wrap:true }),
+                           codeBlock(ret[currentRewrite].uop, "python", { wrap:false }));
   // ** rewrite steps
   if (step.match_count >= 1) {
     const rewriteList = metadata.appendChild(document.createElement("div"));