not a part of regen

* delete more tests

* delete and simplify

* flaky on windows

* a few more, those remained

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

@@ -225,13 +225,22 @@ runs:
     - name: Install gpuocelot dependencies (MacOS)
       if: inputs.ocelot == 'true' && runner.os == 'macOS'
       shell: bash
-      run: brew install --quiet cmake ninja llvm@15 zlib glew flex bison boost zstd ncurses
+      run: |
+        pkgs=(cmake ninja llvm@15 zlib glew flex bison [email protected] zstd ncurses)
+        for f in "${pkgs[@]}"; do
+          brew ls --versions "$f" >/dev/null 2>&1 || brew install --quiet "$f"
+        done
+
+        # Fix boost 1.85 for gpuocelot
+        ln -s /opt/homebrew/opt/[email protected] /opt/homebrew/opt/boost || true
+        ln -s /opt/homebrew/opt/boost/lib/libboost_atomic-mt.dylib /opt/homebrew/opt/boost/lib/libboost_atomic.dylib || true
+        ln -s /opt/homebrew/opt/boost/lib/libboost_thread-mt.dylib /opt/homebrew/opt/boost/lib/libboost_thread.dylib || true
     - name: Cache gpuocelot
       if: inputs.ocelot == 'true'
       id: cache-build
       uses: actions/cache@v4
       env:
-        cache-name: cache-gpuocelot-build
+        cache-name: cache-gpuocelot-build-1
       with:
         path: ${{ github.workspace }}/gpuocelot/ocelot
         key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.BUILD_CACHE_VERSION }}
@@ -244,7 +253,8 @@ runs:
         git checkout b16039dc940dc6bc4ea0a98380495769ff35ed99
         mkdir build
         cd build
-        cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5
+        cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF \
+          -DBoost_INCLUDE_DIR=$(brew --prefix boost)/include -DBoost_LIBRARY_DIR=$(brew --prefix boost)/lib -DCMAKE_POLICY_VERSION_MINIMUM=3.5
         ninja
     - name: Install gpuocelot
       if: inputs.ocelot == 'true'

.github/workflows/benchmark.yml

@@ -65,11 +65,13 @@ jobs:
     - name: Test speed vs torch
       run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
     - name: Test tensor cores
-      run: METAL=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+      run: METAL=1 python3.11 test/opt/test_tensor_cores.py
     - name: Test AMX tensor cores
       run: |
-        DEBUG=2 CPU=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
-        DEBUG=2 LLVM=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
+        DEBUG=2 CPU=1 AMX=1 python3.11 test/opt/test_tensor_cores.py
+        DEBUG=2 LLVM=1 AMX=1 python3.11 test/opt/test_tensor_cores.py
+        DEBUG=2 CPU=1 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
+        DEBUG=2 LLVM=1 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
     - name: Run Tensor Core GEMM (float)
       run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
     - name: Run Tensor Core GEMM (half)
@@ -120,8 +122,8 @@ jobs:
       run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
     - name: UsbGPU copy speeds
       run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
-    - name: UsbGPU openpilot test
-      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+    #- name: UsbGPU openpilot test
+    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
     - uses: actions/upload-artifact@v4
       with:
         name: Speed (Mac)
@@ -194,8 +196,8 @@ jobs:
       run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
     - name: Test tensor cores
       run: |
-        NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
-        PTX=1 ALLOW_TF32=1 NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        NV=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
+        PTX=1 ALLOW_TF32=1 NV=1 python3 test/opt/test_tensor_cores.py
     - name: Run Tensor Core GEMM (CUDA)
       run: |
         CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
@@ -317,10 +319,10 @@ jobs:
       run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
     - name: Run MLPerf resnet eval on training data
       run: time BENCHMARK_LOG=resnet_eval NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
-    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
-    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
+    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+    #  run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
+    #- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+    #  run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
     - name: Run 10 MLPerf Bert training steps (6 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
       run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
@@ -394,8 +396,8 @@ jobs:
       run: AMD=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
     - name: Test tensor cores
       run: |
-        AMD=1 AMD_LLVM=0 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
-        AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+        AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
+        AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
         AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
     - name: Run Tensor Core GEMM (AMD)
       run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
@@ -515,10 +517,10 @@ jobs:
       run: BENCHMARK_LOG=cifar_10steps_half_wino AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
     - name: Run full CIFAR training w 1 GPU
       run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
-    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
-    - name: Run full CIFAR training steps w 6 GPUS (REMOTE)
-      run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
+    #- name: Run full CIFAR training steps w 6 GPUS
+    #  run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+    #- name: Run full CIFAR training steps w 6 GPUS (REMOTE)
+    #  run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
     - uses: actions/upload-artifact@v4
       with:
         name: Speed (AMD Training)
@@ -568,10 +570,10 @@ jobs:
       run: test/external/process_replay/reset.py
     - name: Run MLPerf resnet eval
       run: time BENCHMARK_LOG=resnet_eval AMD=1 MODEL=resnet python3 examples/mlperf/model_eval.py
-    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
-    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
+    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+    #  run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
+    #- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+    #  run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
     - name: Run 10 MLPerf Bert training steps (6 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
       run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
@@ -679,8 +681,8 @@ jobs:
     # Fails on 9070
     # - name: Test tensor cores
     #   run: |
-    #     AMD=1 AMD_LLVM=0 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
-    #     AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+    #     AMD=1 AMD_LLVM=0 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
+    #     AMD=1 AMD_LLVM=1 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
     #     AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
     - name: Run Tensor Core GEMM (AMD)
       run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee am_matmul_amd.txt
@@ -688,6 +690,10 @@ jobs:
       run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
     - name: Test DISK copy time
       run: AMD=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
+    - name: Test CPU copy time
+      run: |
+        AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
+        AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
     - name: Run full CIFAR training w 1 GPU
       run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
     # TODO: enable
@@ -742,15 +748,19 @@ jobs:
     - name: Test driver start time
       run: time DEBUG=3 NV=1 python3 test/test_tiny.py TestTiny.test_plus
     - name: Test tensor cores
-      run: NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+      run: NV=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
     - name: Test DISK copy time
       run: NV=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
+    - name: Test CPU copy time
+      run: |
+        NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
+        NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
     - name: Test LLAMA-3
       run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
     - name: Run full CIFAR training w 1 GPU
       run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
-    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
+    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+    #  run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
     - name: Run 10 MLPerf Bert training steps (1 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
       run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt

.github/workflows/test.yml

@@ -7,6 +7,7 @@ env:
   BUILD_CACHE_VERSION: '1'
   CAPTURE_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+  PYTHONPATH: ${{ github.workspace }}
 
 on:
   push:
@@ -79,7 +80,7 @@ jobs:
         python docs/abstractions2.py
         python docs/abstractions3.py
     - name: Test Quickstart
-      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' docs/quickstart.md > quickstart.py &&  PYTHONPATH=. python quickstart.py
+      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' docs/quickstart.md > quickstart.py &&  python quickstart.py
     - name: Test DEBUG
       run: DEBUG=100 python3 -c "from tinygrad import Tensor; N = 1024; a, b = Tensor.rand(N, N), Tensor.rand(N, N); c = (a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2); print((c.numpy() - (a.numpy() @ b.numpy())).mean())"
     - name: Compile EfficientNet to C and test it
@@ -182,19 +183,19 @@ jobs:
         pip3 install --upgrade --force-reinstall ruff==0.11.0
         python3 -m ruff check extra/torch_backend/backend.py
     - name: Test one op
-      run: PYTHONPATH=. FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
+      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
     - name: Test ResNet-18
-      run: PYTHONPATH=. DEBUG=2 python3 extra/torch_backend/example.py
+      run: DEBUG=2 python3 extra/torch_backend/example.py
     - name: My (custom) tests
-      run: PYTHONPATH=. python3 extra/torch_backend/test.py
+      run: python3 extra/torch_backend/test.py
     - name: Test one op in torch tests
-      run: PYTHONPATH=. DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
+      run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
     - name: Test Ops with TINY_BACKEND
-      run: PYTHONPATH=. LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
+      run: 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: PYTHONPATH=. TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
+      run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
     - name: Test multi-gpu
-      run: PYTHONPATH=. LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
+      run: LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
 
   torchbackendmore:
     name: Torch Backend Tests More
@@ -216,9 +217,9 @@ jobs:
         sudo apt update || true
         sudo apt install -y --no-install-recommends ninja-build
     - name: Test beautiful_mnist in torch with TINY_BACKEND
-      run: SPLIT_REDUCEOP=0 FUSE_ARANGE=1 PYTHONPATH=. LLVM=1 TARGET_EVAL_ACC_PCT=96.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+      run: SPLIT_REDUCEOP=0 FUSE_ARANGE=1 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: PYTHONPATH=. python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
+      run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
 
   tc:
     name: Tensor Core tests
@@ -240,55 +241,45 @@ jobs:
         IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_simple_conv2d
     - name: Test emulated METAL tensor cores
       run: |
-        DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_big_gemm
-        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
-        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_big_gemm
+        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
     - name: Test emulated AMX tensor cores
-      run: PYTHONPATH=. DEBUG=2 AMX=1 EMULATE_AMX=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
+      run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
     - name: Test emulated AMD tensor cores
       run: |
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
     - name: Test emulated AMD MFMA tensor cores
       run: |
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_MFMA=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_MFMA=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_MFMA=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        DEBUG=2 EMULATE=AMD_MFMA FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD_MFMA FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
     - name: Test emulated AMD RDNA4 tensor cores
       run: |
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
     - name: Test emulated CUDA tensor cores
       run: |
-        DEBUG=2 EMULATE_CUDA=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
-        DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
-        DEBUG=2 EMULATE_CUDA_SM75=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
-        PYTHONPATH="." DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
-        PYTHONPATH="." DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        DEBUG=2 EMULATE=CUDA FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
+        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
+        DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
+        PYTHONPATH="." DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
     - name: Test emulated INTEL OpenCL tensor cores
-      run: DEBUG=2 EMULATE_INTEL=1 FORWARD_ONLY=1 PYTHON=1 HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
-    - name: Full test tensor cores
-      run: |
-        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
-        PYTHONPATH=. DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
-        PYTHONPATH=. DEBUG=2 EMULATE_INTEL=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
-        PYTHONPATH=. DEBUG=2 AMX=1 EMULATE_AMX=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
+      run: DEBUG=2 EMULATE=INTEL FORWARD_ONLY=1 PYTHON=1 HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
+    - name: Test emulated AMX tensor cores
+      run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
     - name: Test device flop counts
       run: |
-        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        PYTHONPATH=. DEBUG=2 EMULATE_CUDA=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        PYTHONPATH=. DEBUG=2 EMULATE_INTEL=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        PYTHONPATH=. DEBUG=2 AMX=1 EMULATE_AMX=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStats.test_simple_matmul
+        DEBUG=2 EMULATE=METAL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=AMD PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=CUDA PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=INTEL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 AMX=1 EMULATE=AMX PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStats.test_simple_matmul
 
   bepython:
     name: Python Backend
@@ -305,15 +296,15 @@ jobs:
         key: be-minimal
         deps: testing_minimal
     - name: Test dtype with Python emulator
-      run: DEBUG=1 PYTHONPATH=. PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
+      run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
     - name: Test ops with Python emulator
-      run: DEBUG=2 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py -k "not (test_split or test_simple_cumsum or test_cumsum or test_einsum or test_dot or test_dot_1d or test_big_gemm or test_broadcastdot or test_multidot or test_var_axis or test_std_axis or test_broadcast_full or test_broadcast_partial or test_simple_conv3d or test_dilated_conv_transpose2d or test_simple_conv_transpose3d or test_large_input_conv2d or test_max_pool2d or test_max_pool2d_simple or test_max_pool2d_bigger_stride or test_avg_pool2d or test_cat or test_scaled_product_attention or test_scaled_product_attention_causal or test_slice_fancy_indexing_dim_inject_none or test_slice_fancy_indexing_list_indices or test_slice_fancy_indexing_no_dim_collapse or test_slice_fancy_indexing_tuple_indices or test_slice_fancy_indexing_list_with_tensors or test_slice_fancy_indexing_dim_collapse_int or test_interpolate_bilinear or test_interpolate_bilinear_corners_aligned or test_scaled_dot_product_attention or test_cummax or test_simple_cummax or test_logcumsumexp or test_sort or test_cumprod)" --durations=20
+      run: DEBUG=2 SKIP_SLOW_TEST=1 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py --durations=20
     - name: Test uops with Python emulator
       run: PYTHON=1 python3 -m pytest test/test_uops.py --durations=20
     - name: Test symbolic with Python emulator
-      run: PYTHONPATH=. PYTHON=1 python3 test/test_symbolic_ops.py
+      run: PYTHON=1 python3 test/test_symbolic_ops.py
     - name: test_renderer_failures with Python emulator
-      run: PYTHONPATH=. PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures
+      run: PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures
 
   linter:
     name: Linters
@@ -361,7 +352,7 @@ jobs:
         pydeps: "pillow"
         deps: testing_unit
     - name: Test README
-      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py &&  PYTHONPATH=. python README.py
+      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py &&  python README.py
     - name: Run unit tests
       run: PYTHONPATH="." python -m pytest -n=auto test/unit/ --durations=20
     - name: Run targetted tests on NULL backend
@@ -379,11 +370,11 @@ jobs:
       run: |
         test/external/process_replay/reset.py
         CAPTURE_PROCESS_REPLAY=1 python test/test_tiny.py TestTiny.test_plus
-        PYTHONPATH=. python extra/optimization/extract_dataset.py
+        python extra/optimization/extract_dataset.py
         gzip -c /tmp/sops > extra/datasets/sops.gz
-        DEBUG=1 MIN_ASTS=1 PYTHONPATH=. python extra/optimization/get_action_space.py
-    - name: Repo line count < 17500 lines
-      run: MAX_LINE_COUNT=17500 python sz.py
+        #DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
+    - name: Repo line count < 18000 lines
+      run: MAX_LINE_COUNT=18000 python sz.py
 
   fuzzing:
     name: Fuzzing
@@ -477,7 +468,7 @@ jobs:
           llvm: 'true'
       - name: Test openpilot model kernel count and gate usage
         run: |
-          PYTHONPATH="." ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2134 ALLOWED_GATED_READ_IMAGE=13 FLOAT16=0 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
+          PYTHONPATH="." ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2175 ALLOWED_GATED_READ_IMAGE=16 FLOAT16=0 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
       - name: Test openpilot alt model correctness (float32)
         run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
       - name: Test openpilot fastvits model correctness (float32)
@@ -513,11 +504,11 @@ jobs:
       - name: Test ONNX (LLVM)
         run: LLVM=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
       - name: Test ONNX Runner (CPU)
-        run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py
+        run: CPU=1 python3 test/external/external_test_onnx_runner.py
       - name: Test Additional ONNX Ops (CPU)
-        run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_ops.py
+        run: CPU=1 python3 test/external/external_test_onnx_ops.py
       - name: Test Quantize ONNX
-        run: CPU=1 PYTHONPATH=. python3 test/test_quantize_onnx.py
+        run: CPU=1 python3 test/test_quantize_onnx.py
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -541,16 +532,16 @@ jobs:
           opencl: 'true'
       - name: Test ONNX (GPU)
         run: GPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
-      - name: Test Optimization Helpers
-        run: PYTHONPATH="." DEBUG=1 python3 extra/optimization/test_helpers.py
+      #- name: Test Optimization Helpers
+      #  run: DEBUG=1 python3 extra/optimization/test_helpers.py
       #- name: Test Action Space
-      #  run: PYTHONPATH="." DEBUG=1 GPU=1 python3 extra/optimization/get_action_space.py
+      #  run: DEBUG=1 GPU=1 python3 extra/optimization/get_action_space.py
       - name: Test Beam Search
-        run: PYTHONPATH="." GPU=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
+        run: GPU=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
       - name: Test MLPerf stuff
         run: GPU=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
       - name: Test llama 3 training
-        run: MAX_BUFFER_SIZE=0 PYTHONPATH="." DEV=NULL SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
+        run: MAX_BUFFER_SIZE=0 DEV=NULL SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -614,7 +605,7 @@ jobs:
         CPU=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 \
           -k "not test_symbolic_arange_sym_step and not test_threefry_doesnt_use_long" \
           test/test_tiny.py test/test_rangeify.py test/test_ops.py test/test_tensor_variable.py \
-          test/test_outerworld_range.py test/test_sample.py test/test_randomness.py test/test_tensor_data.py
+          test/test_outerworld_range.py test/test_sample.py test/test_randomness.py
     - name: Test CPU=1 RANGEIFY=2
       run: CPU=1 RANGEIFY=2 python3 -m pytest -n auto test/test_tiny.py test/test_rangeify.py test/test_ops.py --durations 20
     - name: Test LLVM=1 RANGEIFY=1 (slow tests)
@@ -639,7 +630,7 @@ jobs:
     - name: Test LLVM=1 DEVECTORIZE=0
       run: LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
     - name: Test LLVM=1 DEVECTORIZE=0 for model
-      run: PYTHONPATH="." LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
+      run: LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
     - name: Test CPU=1 DEVECTORIZE=0
       run: CPU=1 DEVECTORIZE=0 FUSE_ARANGE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
 
@@ -675,9 +666,9 @@ jobs:
     - name: Run test_tiny on DSP
       run: DEBUG=2 DSP=1 python test/test_tiny.py
     - name: Test transcendentals
-      run: CC=clang-20 PYTHONPATH="." DEBUG=2 DSP=1 python test/test_transcendental.py TestTranscendentalVectorized
+      run: CC=clang-20 DEBUG=2 DSP=1 python test/test_transcendental.py TestTranscendentalVectorized
     - name: Test quantize onnx
-      run: PYTHONPATH="." DEBUG=2 DSP=1 python3 test/test_quantize_onnx.py
+      run: DEBUG=2 DSP=1 python3 test/test_quantize_onnx.py
 
   testwebgpu:
     name: Linux (WebGPU)
@@ -720,7 +711,6 @@ jobs:
       MOCKGPU: 1
       FORWARD_ONLY: 1
       AMD_LLVM: ${{ matrix.backend == 'amdllvm' && '1' || matrix.backend != 'amdllvm' && '0' }}
-      PYTHONPATH: ${{ github.workspace }}
     steps:
       - name: Checkout Code
         uses: actions/checkout@v4
@@ -760,7 +750,9 @@ jobs:
     name: Linux (${{ matrix.backend }})
     runs-on: ubuntu-22.04
     timeout-minutes: 20
-
+    env:
+      MOCKGPU: 1
+      FORWARD_ONLY: 1
     steps:
       - name: Checkout Code
         uses: actions/checkout@v4
@@ -772,11 +764,11 @@ jobs:
           cuda: 'true'
           ocelot: 'true'
       - name: Set env
-        run: printf "${{ matrix.backend == 'PTX' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nCUDA=1\nPTX=1\nMOCKGPU=1' || matrix.backend == 'nv' && 'NV=1\nMOCKGPU=1\nFORWARD_ONLY=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'PTX' && 'CUDA=1\nPTX=1' || matrix.backend == 'nv' && 'NV=1\nSKIP_SLOW_TEST=1' }}" >> $GITHUB_ENV
       - name: Check Device.DEFAULT and print some source
         run: |
-          PYTHONPATH=${{ github.workspace }} python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"
-          DEBUG=5 PYTHONPATH=${{ github.workspace }} FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"
+          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
       - name: Run pytest (cuda)
         # skip multitensor because it's slow
         run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore test/test_gc.py --ignore test/test_multitensor.py --durations=20
@@ -809,8 +801,8 @@ jobs:
         run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'gpu' && 'GPU=1' }}" >> $GITHUB_ENV
       - name: Check Device.DEFAULT and print some source
         run: |
-          PYTHONPATH=${{ github.workspace }} python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['LLVM','CPU','GPU'], Device.DEFAULT"
-          DEBUG=5 PYTHONPATH=${{ github.workspace }} FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['LLVM','CPU','GPU'], Device.DEFAULT"
+          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
       - name: Run pytest (not cuda)
         run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
       - name: Run TRANSCENDENTAL math
@@ -851,11 +843,11 @@ jobs:
     - name: Test tensor core ops (real)
       run: METAL=1 DEBUG=3 python test/test_ops.py TestOps.test_big_gemm
     - name: Test LLaMA compile speed
-      run: PYTHONPATH="." METAL=1 python test/external/external_test_speed_llama.py
+      run: METAL=1 python test/external/external_test_speed_llama.py
     - name: Test Beam Search
-      run: PYTHONPATH="." METAL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
+      run: METAL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
     #- name: Fuzz Test linearizer
-    #  run: PYTHONPATH="." METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
+    #  run: METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
     - name: Run TRANSCENDENTAL math
       run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
     - name: Run pytest (amd)
@@ -918,7 +910,7 @@ jobs:
     #    cp $GITHUB_WORKSPACE/test/web/test_viz.js .
     #    node test_viz.js
     - name: Test ONNX Runner (WEBGPU)
-      run: WEBGPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py
+      run: WEBGPU=1 python3 test/external/external_test_onnx_runner.py
 
   osxremote:
    name: MacOS (remote metal)
@@ -950,7 +942,6 @@ jobs:
     timeout-minutes: 20
     env:
       REMOTE: 1
-      PYTHONPATH: ${{ github.workspace }}
     steps:
       - name: Checkout Code
         uses: actions/checkout@v4
@@ -1071,7 +1062,8 @@ jobs:
         run:  printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'webgpu' && 'WEBGPU=1'}}" >> $GITHUB_ENV
       - name: Run unit tests
         if: matrix.backend=='llvm'
-        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_elf.py --ignore=test/unit/test_tar.py
+        # test_newton_schulz hits RecursionError
+        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_elf.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
       - name: Run pytest (${{ matrix.backend }})
         shell: bash
         run: |

.pre-commit-config.yaml

@@ -20,12 +20,6 @@ repos:
         language: system
         always_run: true
         pass_filenames: false
-      - id: devicetests
-        name: select GPU tests
-        entry: env GPU=1 PYTHONPATH="." python3 -m pytest test/test_uops.py test/test_search.py
-        language: system
-        always_run: true
-        pass_filenames: false
       - id: tests
         name: subset of tests
         entry: env PYTHONPATH="." python3 -m pytest -n=4 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py

.pylintrc

@@ -54,11 +54,12 @@ confidence=
 # --enable=similarities". If you want to run only the classes checker, but have
 # no Warning level messages displayed, use"--disable=all --enable=classes
 # --disable=W"
-disable=C,R,W0613,W0511,W0212,W0201,W0106,W0603,W0621,W0703,W1201,W1203,E1136,W1514,E1101,W0221,W0105,E0401,abstract-method
+disable=C,R,W0613,W0511,W0212,W0201,W0106,W0603,W0621,W0703,W1201,W1203,E1136,W1514,E1101,W0221,W0105,E0401,abstract-method,W0707
 # E1101 for function binding
 # W0221 for Function class
 # W0105 for comment strings
 # E0401 for missing imports
+# W0707 for not reraising
 
 # Enable the message, report, category or checker with the given id(s). You can
 # either give multiple identifier separated by comma (,) or put this option

docs/developer/layout.md

@@ -22,12 +22,6 @@ Group UOps into kernels.
 
 Transforms the ast into an optimized ast. This is where BEAM search and heuristics live.
 
-::: tinygrad.codegen.opt.get_optimized_ast
-    options:
-        members: false
-        show_labels: false
-        show_source: false
-
 ---
 
 ## tinygrad/codegen

examples/beautiful_cifar.py

@@ -2,7 +2,6 @@ import time
 start_tm = time.perf_counter()
 import math
 from typing import Tuple, cast
-import numpy as np
 from tinygrad import Tensor, nn, GlobalCounters, TinyJit, dtypes, Device
 from tinygrad.helpers import partition, trange, getenv, Context
 from extra.lr_scheduler import OneCycleLR
@@ -150,13 +149,12 @@ if __name__ == "__main__":
       acc.append((out.argmax(-1) == Y).sum() / eval_batchsize)
     return Tensor.stack(*loss).mean() / (batchsize*loss_batchsize_scaler), Tensor.stack(*acc).mean()
 
-  np.random.seed(1337)
+  Tensor.manual_seed(1337)
+  num_train_samples = X_train.shape[0]
+
   for epoch in range(math.ceil(hyp['misc']['train_epochs'])):
-    # TODO: move to tinygrad
     gst = time.perf_counter()
-    idxs = np.arange(X_train.shape[0])
-    np.random.shuffle(idxs)
-    tidxs = Tensor(idxs, dtype='int')[:num_steps_per_epoch*batchsize].reshape(num_steps_per_epoch, batchsize)  # NOTE: long doesn't fold
+    tidxs = Tensor.randperm(num_train_samples, dtype='int')[:num_steps_per_epoch*batchsize].reshape(num_steps_per_epoch, batchsize)
     train_loss:float = 0
     for epoch_step in (t:=trange(num_steps_per_epoch)):
       st = time.perf_counter()

examples/mlperf/dataloader.py

@@ -758,6 +758,27 @@ def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0
       batch.append(tokens)
     yield Tensor.stack(batch, dim=0)
 
+def batch_load_llama3_small(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
+  if val:
+    dataset = BlendedGPTDataset([
+      base_dir / "c4-validation-91205-samples.en_text_document",
+    ], [
+      1.0
+    ], samples, seqlen, seed, False)
+  else:
+    dataset = BlendedGPTDataset([
+      base_dir / "c4-train.en_6_text_document",
+    ], [
+      1.0
+    ], samples, seqlen, seed, True)
+
+  for b in range(math.ceil(samples / bs)):
+    batch = []
+    for i in range(bs):
+      tokens = dataset.get(b * bs + i)
+      batch.append(tokens)
+    yield Tensor.stack(batch, dim=0)
+
 if __name__ == "__main__":
   def load_unet3d(val):
     assert not val, "validation set is not supported due to different sizes on inputs"

examples/mlperf/model_eval.py

@@ -243,31 +243,49 @@ def eval_mrcnn():
 
 def eval_llama3():
   from extra.models.llama import Transformer
-  from examples.llama3 import MODEL_PARAMS
+  from examples.llama3 import MODEL_PARAMS, load, convert_from_huggingface
   from tinygrad.helpers import tqdm
 
-  bs = 4
-  sequence_length = 512
+  BASEDIR = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
+  BS = getenv("BS", 4)
+  SMALL = getenv("SMALL", 0)
+  SEQLEN = getenv("SEQLEN", 8192)
+  MODEL_PATH = Path(getenv("MODEL_PATH", "/raid/weights/llama31_8b/"))
 
-  model = Transformer(**(MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}), max_context=sequence_length, jit=False, disable_kv_cache=True)
+  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
+  params = params | {"vocab_size": 32000} if not SMALL else params
+  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
+  model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
+
+  # load weights
+  weights = load(str(MODEL_PATH / "model.safetensors.index.json"))
+  if "model.embed_tokens.weight" in weights:
+    print("converting from huggingface format")
+    weights = convert_from_huggingface(weights, params["n_layers"], params["n_heads"], params["n_kv_heads"])
+
+  load_state_dict(model, weights, strict=False, consume=True)
 
   @TinyJit
   def eval_step(model, tokens):
     logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
     loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    return loss.flatten()
+    return loss.flatten().float()
 
-  from examples.mlperf.dataloader import batch_load_llama3
-  iter = batch_load_llama3(bs, 5760, sequence_length, Path(getenv("BASEDIR", "/raid/datasets/c4/")), True)
+  if SMALL:
+    from examples.mlperf.dataloader import batch_load_llama3_small
+    iter = batch_load_llama3_small(BS, 5760, SEQLEN, BASEDIR, val=True)
+  else:
+    from examples.mlperf.dataloader import batch_load_llama3
+    iter = batch_load_llama3(BS, 5760, SEQLEN, BASEDIR, val=True)
 
   losses = []
-  for tokens in tqdm(iter, total=5760//bs):
+  for tokens in tqdm(iter, total=5760//BS):
     GlobalCounters.reset()
     losses += eval_step(model, tokens).tolist()
     tqdm.write(f"loss: {np.mean(losses)}")
 
-  log_perplexity = Tensor(losses).mean()
-  print(f"Log Perplexity: {log_perplexity.item()}")
+  log_perplexity = np.mean(losses)
+  print(f"Log Perplexity: {log_perplexity}")
 
 if __name__ == "__main__":
   # inference only

examples/mlperf/model_train.py

@@ -4,7 +4,7 @@ import multiprocessing
 
 from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
 from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, FUSE_CONV_BW, Profiling
-from tinygrad.nn.state import get_parameters, get_state_dict, safe_load, safe_save
+from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
 from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW
 
 from extra.lr_scheduler import LRSchedulerGroup
@@ -1290,12 +1290,14 @@ def train_llama3():
   from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
 
   config = {}
+  BASEDIR            = config["BASEDIR"]                = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
   BS                 = config["BS"]                     = getenv("BS", 16)
   grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
   GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
   SEED               = config["SEED"]                   = getenv("SEED", 5760)
   SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
   TRAIN_ON_VAL       = config["TRAIN_ON_VAL"]           = getenv("TRAIN_ON_VAL", 0)
+  SMALL              = config["SMALL"]                  = getenv("SMALL", 0)
   SAMPLES            = config["SAMPLES"]                = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000 * 1152)
   EVAL_FREQ          = config["EVAL_FREQ"]              = getenv("EVAL_FREQ", 46080)
   EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 16)
@@ -1311,13 +1313,14 @@ def train_llama3():
 
   opt_gradient_clip_norm = 1.0
   opt_learning_rate_warmup_steps = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
-  opt_learning_rate_decay_steps = getenv("DECAY_STEPS", math.ceil(1_200_000 * 1152 / GBS) - opt_learning_rate_warmup_steps)
+  opt_learning_rate_decay_steps = getenv("MAX_STEPS", math.ceil(1_200_000 * 1152 / GBS)) - opt_learning_rate_warmup_steps
   opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152)  # NOTE: cannot change for benchmark
-  opt_end_learning_rate = 8e-7
+  opt_end_learning_rate = getenv("END_LR", 8e-7)
 
   # TODO: confirm weights are in bf16
   # vocab_size from the mixtral tokenizer
-  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}
+  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
+  params = params | {"vocab_size": 32000} if not SMALL else params
   if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
   model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
 
@@ -1353,6 +1356,15 @@ def train_llama3():
                 b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
   scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
 
+  if resume_ckpt := getenv("RESUME_CKPT"):
+    fn = f"./ckpts/llama3_{resume_ckpt}.safe"
+    print(f"loading initial checkpoint from {fn}")
+    load_state_dict(model, safe_load(fn), realize=False)
+
+    fn = f"./ckpts/llama3_{resume_ckpt}_optim.safe"
+    print(f"loading optim checkpoint from {fn}")
+    load_state_dict(scheduler, safe_load(fn), realize=False)
+
   @TinyJit
   @Tensor.train()
   def train_step(model, tokens:Tensor, grad_acc:int):
@@ -1403,43 +1415,55 @@ def train_llama3():
   # ** data iters **
   def fake_data(bs, samples):
     for _ in range(samples // bs):
-      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=32000, dtype=dtypes.int32, device=Device.DEFAULT)
+      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
 
   def get_train_iter():
     if getenv("FAKEDATA", 0):
       return fake_data(GBS, SAMPLES)
     else:
-      from examples.mlperf.dataloader import batch_load_llama3
-      return batch_load_llama3(GBS, SAMPLES, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=bool(TRAIN_ON_VAL))
+      if SMALL:
+        from examples.mlperf.dataloader import batch_load_llama3_small
+        return batch_load_llama3_small(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
+      else:
+        from examples.mlperf.dataloader import batch_load_llama3
+        return batch_load_llama3(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
 
   def get_eval_iter():
     if getenv("FAKEDATA", 0):
       return fake_data(EVAL_BS, 5760)
     else:
-      from examples.mlperf.dataloader import batch_load_llama3
-      return batch_load_llama3(EVAL_BS, 5760, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=True)
+      if SMALL:
+        from examples.mlperf.dataloader import batch_load_llama3_small
+        return batch_load_llama3_small(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
+      else:
+        from examples.mlperf.dataloader import batch_load_llama3
+        return batch_load_llama3(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
 
   iter = get_train_iter()
-  i, sequences_seen = 0, 0
+  i, sequences_seen = resume_ckpt, 0
   for tokens in tqdm(iter, total=SAMPLES//GBS):
     t = time.perf_counter()
     GlobalCounters.reset()
     loss, lr = train_step(model, tokens, grad_acc)
     loss = loss.float().item()
-    # above as tqdm.write f-string
+
+    i += 1
+    sequences_seen += tokens.shape[0]
+
     tqdm.write(f"{loss:.4f} loss, {lr.item():.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {time.perf_counter()-t:.2f} s")
     if (fname:=getenv("LOSS_FILE", "")):
       with open(fname, "a") as f:
         f.write(f"{i} {loss:.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")
 
-    if getenv("CKPT") and (i % 200 == 0 or i == 10):
+    if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
       tqdm.write("saving checkpoint")
       if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
       fn = f"{ckpt_dir}/llama3_{i}.safe"
       safe_save(get_state_dict(model), fn)
 
-    i += 1
-    sequences_seen += tokens.shape[0]
+      tqdm.write("saving optim checkpoint")
+      fn = f"{ckpt_dir}/llama3_{i}_optim.safe"
+      safe_save(get_state_dict(scheduler), fn)
 
     if sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1):
       tqdm.write(f"evaluating after {sequences_seen} sequences")

extra/amdpci/proclogs.py

@@ -37,7 +37,7 @@ def main():
   dev = PCIIface(None, 0)
   for x, y in dev.dev_impl.__dict__.items():
     if isinstance(y, AMRegister):
-      for inst, addr in y.addr.keys(): reg_names[addr] = f"{x}, xcc={inst}"
+      for inst, addr in y.addr.items(): reg_names[addr] = f"{x}, xcc={inst}"
 
   with open(sys.argv[1], 'r') as f:
     log_content = log_content_them = f.read()

extra/backends/rdna.py

@@ -29,10 +29,10 @@ def uops_to_rdna(function_name:str, uops:UOpGraph) -> str:
   r: Dict[UOp, str] = {}
   for u in uops:
     if u.uop == UOps.SPECIAL:
-      if u.arg[1].startswith("lidx"):
-        r[u] = f'v{u.arg[0]}'
-      elif u.arg[1].startswith("gidx"):
-        r[u] = f's{2+u.arg[0]}'
+      if u.arg.startswith("lidx"):
+        r[u] = f'v{u.src[0].arg}'
+      elif u.arg.startswith("gidx"):
+        r[u] = f's{2+u.src[0].arg}'
       else:
         raise NotImplementedError
     elif u.uop == UOps.CONST:

extra/gemm/amd_uop_matmul.py

@@ -65,7 +65,7 @@ def top_spec_kernel3():
   c = a@b
   sink = c.schedule()[-1].ast
   L = 16
-  sink = sink.reshape((N//L, L, N//L, L)) #.lift({0:UOp.range(dtypes.int, N//BM, 0), 2:UOp.range(dtypes.int, N//BN, 1)})
+  sink = sink.reshape((N//L, L, N//L, L)) #.lift({0:UOp.range(N//BM, 0), 2:UOp.range(N//BN, 1)})
   sink = graph_rewrite(sink, view_left+pm)
   axis_types = (AxisType.GLOBAL, AxisType.LOCAL, AxisType.GLOBAL, AxisType.LOCAL, AxisType.REDUCE)
   return sink.replace(arg=KernelInfo(name="top_"+to_colored(sink.full_shape, axis_types), axis_types=axis_types))
@@ -186,7 +186,7 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
 
   c_regs = UOp(Ops.DEFINE_REG, dtypes.float.ptr(TM * nbIterWaveM * TN * nbIterWaveN), arg=2)
 
-  i = UOp.range(dtypes.int, c_regs.dtype.size, 16)
+  i = UOp.range(c_regs.dtype.size, 16)
   init_store = c_regs[i].store(UOp.const(dtypes.float, 0.0), i)
 
   if kernel4:
@@ -197,53 +197,53 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
     kId = 0
 
     # load from globals into locals
-    i = UOp.range(dtypes.int, nbReadsB, 0)
+    i = UOp.range(nbReadsB, 0)
     index_x = BN * blockIdx_x + rBIdx
     index_y = rBIdy + i * strideReadB + kId
     Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(b[N * index_y + index_x].load(), i)
 
-    i = UOp.range(dtypes.int, nbReadsA, 1)
+    i = UOp.range(nbReadsA, 1)
     index_x = rAIdx + kId
     index_y = BM * blockIdx_y + rAIdy + i * strideReadA
     As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)
 
     # iterate over the middle chunk
-    kId_range = UOp.range(dtypes.int, N//BK-1, 2)
+    kId_range = UOp.range(N//BK-1, 2)
     kId = kId_range*BK
 
     barrier = UOp.barrier(As_store, Bs_store)
 
     # load from globals into registers (next round)
-    i = UOp.range(dtypes.int, nbReadsB, 3)
+    i = UOp.range(nbReadsB, 3)
     index_x = BN * blockIdx_x + rBIdx
     index_y = rBIdy + i * strideReadB + kId + BK
     regB_store = regB[i].store(b[N * index_y + index_x].load(), i)
 
-    i = UOp.range(dtypes.int, nbReadsA, 4)
+    i = UOp.range(nbReadsA, 4)
     index_x = rAIdx + kId + BK
     index_y = BM * blockIdx_y + rAIdy + i * strideReadA
     regA_store = regA[i].store(a[N * index_y + index_x].load(), i)
 
     def inner_loop(first_range, inp_dep=()):
       # inner unroll
-      k = UOp.range(dtypes.int, BK, first_range+0)
+      k = UOp.range(BK, first_range+0)
 
       # load from locals into registers
-      iterWave = UOp.range(dtypes.int, nbIterWaveN, first_range+1)
-      i = UOp.range(dtypes.int, TN, first_range+2)
+      iterWave = UOp.range(nbIterWaveN, first_range+1)
+      i = UOp.range(TN, first_range+2)
       index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
       B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(*inp_dep), iterWave, i)
 
-      iterWave = UOp.range(dtypes.int, nbIterWaveM, first_range+3)
-      i = UOp.range(dtypes.int, TM, first_range+4)
+      iterWave = UOp.range(nbIterWaveM, first_range+3)
+      i = UOp.range(TM, first_range+4)
       index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
       A_col_store = A_col[iterWave*TM + i].store(As[k*BM_As_stride + index].load(*inp_dep), iterWave, i)
 
       # do the GEMM math
-      iterWaveM = UOp.range(dtypes.int, nbIterWaveM, first_range+5)
-      yt = UOp.range(dtypes.int, TM, first_range+6)
-      iterWaveN = UOp.range(dtypes.int, nbIterWaveN, first_range+7)
-      xt = UOp.range(dtypes.int, TN, first_range+8)
+      iterWaveM = UOp.range(nbIterWaveM, first_range+5)
+      yt = UOp.range(TM, first_range+6)
+      iterWaveN = UOp.range(nbIterWaveN, first_range+7)
+      xt = UOp.range(TN, first_range+8)
       x = iterWaveN * TN + xt
       y = iterWaveM * TM + yt
       c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
@@ -256,12 +256,12 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
     sink = inner_loop(5, (barrier, regB_store, regA_store)).barrier()
 
     # load from registers into locals
-    i = UOp.range(dtypes.int, nbReadsB, 14)
+    i = UOp.range(nbReadsB, 14)
     index_x = BN * blockIdx_x + rBIdx
     index_y = rBIdy + i * strideReadB + kId + BK
     Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(regB[i].load(sink), i, kId_range)
 
-    i = UOp.range(dtypes.int, nbReadsA, 15)
+    i = UOp.range(nbReadsA, 15)
     index_x = rAIdx + kId + BK
     index_y = BM * blockIdx_y + rAIdy + i * strideReadA
     As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(regA[i].load(sink), i, kId_range)
@@ -269,40 +269,40 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
     # final iteration without the copy
     sink = inner_loop(16, (UOp.barrier(Bs_store, As_store),))
   else:
-    kId_range = UOp.range(dtypes.int, N//BK, 0)
+    kId_range = UOp.range(N//BK, 0)
     kId = kId_range*BK
 
     # load from globals into locals
-    i = UOp.range(dtypes.int, nbReadsB, 1)
+    i = UOp.range(nbReadsB, 1)
     index_x = BN * blockIdx_x + rBIdx
     index_y = rBIdy + i * strideReadB + kId
     Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(b[N * index_y + index_x].load(), i)
 
-    i = UOp.range(dtypes.int, nbReadsA, 2)
+    i = UOp.range(nbReadsA, 2)
     index_x = rAIdx + kId
     index_y = BM * blockIdx_y + rAIdy + i * strideReadA
     As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)
 
     barrier = UOp.barrier(As_store, Bs_store)
 
-    k = UOp.range(dtypes.int, BK, 3)
+    k = UOp.range(BK, 3)
 
     # load from locals into registers
-    iterWave = UOp.range(dtypes.int, nbIterWaveN, 4)
-    i = UOp.range(dtypes.int, TN, 5)
+    iterWave = UOp.range(nbIterWaveN, 4)
+    i = UOp.range(TN, 5)
     index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
     B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(barrier), iterWave, i)
 
-    iterWave = UOp.range(dtypes.int, nbIterWaveM, 6)
-    i = UOp.range(dtypes.int, TM, 7)
+    iterWave = UOp.range(nbIterWaveM, 6)
+    i = UOp.range(TM, 7)
     index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
     A_col_store = A_col[iterWave*TM + i].store(As[k*BM_As_stride + index].load(barrier), iterWave, i)
 
     # do the GEMM math
-    iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 8)
-    yt = UOp.range(dtypes.int, TM, 9)
-    iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 10)
-    xt = UOp.range(dtypes.int, TN, 12)
+    iterWaveM = UOp.range(nbIterWaveM, 8)
+    yt = UOp.range(TM, 9)
+    iterWaveN = UOp.range(nbIterWaveN, 10)
+    xt = UOp.range(TN, 12)
     x = iterWaveN * TN + xt
     y = iterWaveM * TM + yt
     c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
@@ -310,10 +310,10 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
                             iterWaveM, iterWaveN, yt, xt, k, kId_range)
 
   # store c_regs into c
-  iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 1000)
-  yt = UOp.range(dtypes.int, TM, 1001)
-  iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 1002)
-  xt = UOp.range(dtypes.int, TN, 1003)
+  iterWaveM = UOp.range(nbIterWaveM, 1000)
+  yt = UOp.range(TM, 1001)
+  iterWaveN = UOp.range(nbIterWaveN, 1002)
+  xt = UOp.range(TN, 1003)
   xOut = blockIdx_x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave
   yOut = blockIdx_y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave
   indexC = N * (yOut + yt) + xOut + xt

extra/gemm/max_matmul.py

@@ -56,7 +56,7 @@ def randoms():
 def ast_to_cuda_prog(compiler, ast, opts):
   k = Kernel(ast)
   k.apply_opts(opts)
-  p = get_program(k.get_optimized_ast(), k.opts)
+  p = get_program(k.ast, k.opts, k.applied_opts)
   return CUDAProgram(device, p.function_name, compiler.compile(p.src))
 
 if __name__ == "__main__":

extra/gemm/simple_matmul.py

@@ -2,7 +2,7 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, get_single_element
 from tinygrad.dtype import _to_np_dtype
-from tinygrad.codegen.opt.kernel import OptOps
+from tinygrad.codegen.opt import OptOps
 from tinygrad.engine.realize import lower_schedule
 
 dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float

extra/gemm/tinygrad_nv_matmul.py

@@ -29,7 +29,7 @@ if __name__ == "__main__":
       Opt(op=OptOps.LOCAL, axis=0, amt=2),
     ]
   k.apply_opts(opts)
-  prg = get_program(k.get_optimized_ast(), k.opts)
+  prg = get_program(k.ast, k.opts, k.applied_opts)
   new_src = prg.src
   # can mod source here
   prg = replace(prg, src=new_src)

extra/optimization/helpers.py

@@ -1,6 +1,6 @@
 # stuff needed to unpack a kernel
 from tinygrad import Variable
-from tinygrad.codegen.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.dtype import dtypes, PtrDType
 from tinygrad.shape.shapetracker import ShapeTracker

extra/optimization/test_beam_search.py

@@ -16,9 +16,9 @@ class TestBeamSearch(unittest.TestCase):
     BEAM.value = self.old_beam
 
   def test_variable_ast_beam(self):
-    with Context(IGNORE_OOB=1):
-      a = rand(3, 3).reshape((Variable("a", 1, 10).bind(3), 3))
-      a = (a+1).realize()
+    vi = Variable("a", 1, 10).bind(3)
+    a = rand(10, 3)[:vi]
+    a = (a+1).realize()
 
   def test_big_prime_number(self):
     a = rand(367, 367)
@@ -42,18 +42,16 @@ class TestBeamSearch(unittest.TestCase):
 
   def test_variable_big_prime_number(self):
     v = Variable("v", 1, 400).bind(367)
-    a = rand(367, 367)
-    b = rand(367, 367)
-    with Context(IGNORE_OOB=1):
-      c = (a.reshape(367, v) @ b.reshape(v, 367)).realize()
-      np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), atol=1e-4, rtol=1e-4)
+    a = rand(367, 400)
+    b = rand(400, 367)
+    c = (a[:, :v] @ b[:v, :]).realize()
+    np.testing.assert_allclose(c.numpy(), a[:, :367].numpy() @ b[:367, :].numpy(), atol=1e-4, rtol=1e-4)
 
   def test_variable_shrink_prime_number(self):
     v = Variable("v", 1, 400).bind(367)
     a = rand(400, 367)
-    with Context(IGNORE_OOB=1):
-      b = (a.shrink(((0,v), None))+1).reshape(367,367).realize()
-      np.testing.assert_allclose(b.numpy(), a.numpy()[:367]+1, atol=1e-4, rtol=1e-4)
+    b = (a.shrink(((0,v), None))+1).reshape(367,367).realize()
+    np.testing.assert_allclose(b.numpy(), a.numpy()[:367]+1, atol=1e-4, rtol=1e-4)
 
   def test_no_mutate_rawbuffers(self):
     a = rand(3, 3).realize()

extra/remu/src/thread.rs

@@ -673,6 +673,7 @@ impl<'a> Thread<'a> {
                                         39 => f32::log2(s0),
                                         42 => 1.0 / s0,
                                         43 => 1.0 / s0,
+                                        46 => 1.0 / f32::sqrt(s0),
                                         51 => f32::sqrt(s0),
                                         _ => todo_instr!(instruction)?,
                                     }
@@ -1246,7 +1247,7 @@ impl<'a> Thread<'a> {
                             }
 
                             let ret = match op {
-                                257 | 259 | 299 | 260 | 261 | 264 | 272 | 392 | 426 | 531 | 537 | 540 | 551 | 567 | 796 => {
+                                257 | 259 | 299 | 260 | 261 | 264 | 272 | 392 | 426 | 430 | 531 | 537 | 540 | 551 | 567 | 796 => {
                                     let s0 = f32::from_bits(s0).negate(0, neg).absolute(0, abs);
                                     let s1 = f32::from_bits(s1).negate(1, neg).absolute(1, abs);
                                     let s2 = f32::from_bits(s2).negate(2, neg).absolute(2, abs);
@@ -1258,6 +1259,7 @@ impl<'a> Thread<'a> {
                                         272 => f32::max(s0, s1),
                                         299 => f32::mul_add(s0, s1, f32::from_bits(self.vec_reg[vdst])),
                                         426 => s0.recip(),
+                                        430 => 1.0 / f32::sqrt(s0),
                                         531 => f32::mul_add(s0, s1, s2),
                                         537 => f32::min(f32::min(s0, s1), s2),
                                         540 => f32::max(f32::max(s0, s1), s2),
@@ -2625,6 +2627,14 @@ mod test_vop1 {
         assert_eq!(thread.vec_reg[3], 1071644672);
     }
 
+    #[test]
+    fn test_v_rsq_f32() {
+        let mut thread = _helper_test_thread();
+        thread.vec_reg[0] = f32::to_bits(4.0);
+        r(&vec![0x7E005D00, END_PRG], &mut thread);
+        assert_eq!(f32::from_bits(thread.vec_reg[0]), 0.5);
+    }
+
     #[test]
     fn test_v_frexp_exp_i32_f64() {
         [(3573412790272.0, 42), (69.0, 7), (2.0, 2), (f64::NEG_INFINITY, 0)]

extra/replay_pkl.py

@@ -58,7 +58,7 @@ if __name__ == "__main__":
             GlobalCounters.kernel_count -= 1
 
         if not getenv("NOOPT"): k.apply_opts(hand_coded_optimizations(k))
-        p2 = get_program(k.get_optimized_ast(), k.opts)
+        p2 = get_program(k.ast, k.opts, k.applied_opts)
         new_ei = replace(ei, prg=CompiledRunner(p2))
         new_ei.run()
         new_jit.append(new_ei)

extra/test_hcopt.py

@@ -0,0 +1,40 @@
+import time
+from extra.optimization.helpers import load_worlds, ast_str_to_ast
+from tinygrad import Device
+from tinygrad.codegen.lowerer import pm_lowerer, get_index
+from tinygrad.uop.ops import graph_rewrite
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.postrange import Scheduler
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
+from tinygrad.helpers import getenv
+
+if __name__ == "__main__":
+  renderer = Device.default.renderer
+  ast_strs = load_worlds()
+  if (n:=getenv("N", -1)) != -1: ast_strs = ast_strs[n:n+1]
+  good = 0
+  for i, ast_str in enumerate(ast_strs):
+    ast = ast_str_to_ast(ast_str)
+
+    st = time.perf_counter()
+    lin = Kernel(ast, renderer)
+    opt1 = hand_coded_optimizations(lin)
+    et_lin = time.perf_counter() - st
+
+    lowered = graph_rewrite(ast, pm_lowerer, ctx=get_index(ast), bottom_up=True)
+    st = time.perf_counter()
+    sch = Scheduler(lowered, renderer)
+    sch.convert_loop_to_global()
+    sch.simplify_merge_adjacent()
+    opt2 = hand_coded_optimizations(sch)
+    et_sch = time.perf_counter() - st
+
+    if opt1 != opt2:
+      print(f"******* {i:6d}")
+      print("Kernel:    ", lin.colored_shape(), "->", lin.apply_opts(opt1).colored_shape())
+      print("Scheduler: ", sch.colored_shape(), "->", sch.apply_opts(opt2).colored_shape())
+      print(opt1)
+      print(opt2)
+    else:
+      good += 1
+      print(f"******* {i:6d} MATCH {good/(i+1)*100:.2f}% -- {et_lin/et_sch:4.2f}x speedup")

extra/test_pyrender.py

@@ -0,0 +1,20 @@
+from extra.optimization.helpers import load_worlds, ast_str_to_ast
+from tinygrad.helpers import tqdm
+from tinygrad.uop.ops import pyrender, UOp, Ops
+from tinygrad import dtypes
+from tinygrad.shape.shapetracker import ShapeTracker, View
+inf, nan = float('inf'), float('nan')
+
+if __name__ == "__main__":
+  ast_strs = load_worlds()
+  for i, ast_str in enumerate(tqdm(ast_strs)):
+    good_ast = ast_str_to_ast(ast_str)
+    code = '\n'.join(pyrender(good_ast))
+    print("\n***************\n\n"+code)
+    exec(code)
+    if str(good_ast) != str(ast):
+      print(code)
+      print("MISMATCH")
+      print(good_ast)
+      print(ast)
+      break

extra/to_movement_ops.py

@@ -2,7 +2,6 @@ import itertools
 from enum import Enum, auto
 from collections import defaultdict
 from typing import List, Tuple, DefaultDict
-from extra.optimization.helpers import load_worlds, ast_str_to_ast
 from tinygrad.helpers import prod, tqdm
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.shape.shapetracker import ShapeTracker
@@ -147,6 +146,7 @@ def test_rebuild_bufferop_st(ast:UOp):
   for src in ast.src: test_rebuild_bufferop_st(src)
 
 if __name__ == "__main__":
+  from extra.optimization.helpers import load_worlds, ast_str_to_ast
   ast_strs = load_worlds(False, False, True)[:2000]
   for ast_str in tqdm(ast_strs):
     test_rebuild_bufferop_st(ast_str_to_ast(ast_str))

test/device/test_hcq.py

@@ -6,7 +6,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.system import PCIIfaceBase
 from tinygrad.engine.realize import get_runner, CompiledRunner, get_program
-from tinygrad.codegen.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad import Variable
 
 MOCKGPU = getenv("MOCKGPU")

test/external/external_benchmark_sdxl_softmax.py

@@ -24,5 +24,5 @@ if __name__ == "__main__":
   #k.apply_opt(Opt(OptOps.GROUP, 1, 32))
   #k.apply_opt(Opt(OptOps.GROUP, 0, 32))
   from tinygrad.engine.realize import CompiledRunner, ExecItem
-  run = CompiledRunner(prg:=get_program(k.get_optimized_ast(), k.opts))
+  run = CompiledRunner(prg:=get_program(k.ast, k.opts, k.applied_opts))
   ExecItem(run, si.bufs).run()

test/external/external_debug_metal_sd_conv.py

@@ -35,7 +35,7 @@ k = Kernel(ast)
 k.apply_opts(opts)
 bufs = bufs_from_lin(k)
 
-prg = CompiledRunner(get_program(k.get_optimized_ast(), k.opts))
+prg = CompiledRunner(get_program(k.ast, k.opts, k.applied_opts))
 
 for i in range(10):
   speed = prg(bufs, var_vals={}, wait=True)

test/external/external_metal_compile_slow.py

@@ -0,0 +1,56 @@
+# ruff: noqa: E501
+from tinygrad import dtypes
+from tinygrad.helpers import Timing, getenv
+from tinygrad.codegen.opt.kernel import Opt, OptOps
+from tinygrad.engine.realize import get_program, CompiledRunner
+from tinygrad.uop.ops import UOp, Ops, AxisType
+
+if __name__ == "__main__":
+  if getenv("TC", 0) == 0:
+    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1179648), arg=0, src=())
+    c1 = UOp.range(UOp.const(dtypes.int, 512), 0, AxisType.GLOBAL)
+    c2 = UOp.range(UOp.const(dtypes.int, 64), 1, AxisType.GLOBAL)
+    c3 = UOp.range(UOp.const(dtypes.int, 6), 2, AxisType.GLOBAL)
+    c4 = UOp.range(UOp.const(dtypes.int, 6), 3, AxisType.GLOBAL)
+    c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(2097152), arg=1, src=())
+    c6 = UOp.range(UOp.const(dtypes.int, 64), 1004, AxisType.REDUCE)
+    c7 = UOp.range(UOp.const(dtypes.int, 3), 1005, AxisType.REDUCE)
+    c8 = UOp.range(UOp.const(dtypes.int, 3), 1006, AxisType.REDUCE)
+    c9 = c5.index(((((((c1*UOp.const(dtypes.int, 4096))+(c3*UOp.const(dtypes.int, 8)))+c4)+(c6*UOp.const(dtypes.int, 64)))+(c7*UOp.const(dtypes.int, 8)))+c8), UOp.const(dtypes.bool, True)).load()
+    c10 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(36864), arg=2, src=())
+    c11 = c10.index(((((c2*UOp.const(dtypes.int, 576))+(c6*UOp.const(dtypes.int, 9)))+(c7*UOp.const(dtypes.int, 3)))+c8), UOp.const(dtypes.bool, True)).load()
+    c12 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64), arg=3, src=())
+    c13 = c12.index(c2, UOp.const(dtypes.bool, True)).load()
+    c14 = ((c9*c11).reduce(c6, c7, c8, arg=Ops.ADD)+c13)
+    c15 = c0.index(((((c1*UOp.const(dtypes.int, 2304))+(c2*UOp.const(dtypes.int, 36)))+(c3*UOp.const(dtypes.int, 6)))+c4), UOp.const(dtypes.bool, True)).store(c14, c1, c2, c3, c4)
+    ast = c15.sink()
+
+    # this does have tons of locals
+    opts = [Opt(op=OptOps.LOCAL, axis=1, arg=16), Opt(op=OptOps.UPCAST, axis=3, arg=0),
+            Opt(op=OptOps.LOCAL, axis=0, arg=16), Opt(op=OptOps.UPCAST, axis=3, arg=2),
+            Opt(op=OptOps.GROUPTOP, axis=0, arg=16)]
+  else:
+    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(10616832), arg=0, src=())
+    c1 = UOp.range(UOp.const(dtypes.int, 512), 0, AxisType.GLOBAL)
+    c2 = UOp.range(UOp.const(dtypes.int, 64), 1, AxisType.GLOBAL)
+    c3 = UOp.range(UOp.const(dtypes.int, 36), 2, AxisType.GLOBAL)
+    c4 = UOp.range(UOp.const(dtypes.int, 9), 3, AxisType.GLOBAL)
+    c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(36864), arg=1, src=())
+    c6 = UOp.range(UOp.const(dtypes.int, 64), 1004, AxisType.REDUCE)
+    c7 = c5.index((((c2*UOp.const(dtypes.int, 9))+c4)+(c6*UOp.const(dtypes.int, 576))), UOp.const(dtypes.bool, True)).load()
+    c8 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1179648), arg=2, src=())
+    c9 = c8.index((((c1*UOp.const(dtypes.int, 2304))+c3)+(c6*UOp.const(dtypes.int, 36))), UOp.const(dtypes.bool, True)).load()
+    c10 = (c7*c9).reduce(c6, arg=Ops.ADD)
+    c11 = c0.index(((((c1*UOp.const(dtypes.int, 20736))+(c2*UOp.const(dtypes.int, 324)))+(c3*UOp.const(dtypes.int, 9)))+c4), UOp.const(dtypes.bool, True)).store(c10, c1, c2, c3, c4)
+    ast = c11.sink()
+
+    opts = [Opt(op=OptOps.TC, axis=0, arg=(0, 0, 1)), Opt(op=OptOps.UPCAST, axis=2, arg=4),
+            Opt(op=OptOps.UPCAST, axis=3, arg=0), Opt(op=OptOps.GROUP, axis=0, arg=0)]
+
+  prg = get_program(ast, opts=opts)
+  print(prg.src)
+  for i in range(10):
+    with Timing(f"try {i}: "):
+      # NOTE: this doesn't even run the kernel
+      try: CompiledRunner(prg)
+      except RuntimeError: pass

test/external/external_test_onnx_backend.py

@@ -134,7 +134,6 @@ backend_test.exclude('test_simple_rnn_*')
 
 # no control flow
 # control flow uses AttributeProto.GRAPH
-backend_test.exclude('test_if_*')
 backend_test.exclude('test_loop*')
 backend_test.exclude('test_range_float_type_positive_delta_expanded_cpu') # requires loop
 backend_test.exclude('test_affine_grid_2d_align_corners_expanded_cpu')
@@ -183,6 +182,8 @@ backend_test.exclude('test_resize_downsample_scales_cubic_antialias_cpu') # anti
 backend_test.exclude('test_resize_downsample_sizes_cubic_antialias_cpu') # antialias not implemented
 backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_value_only_mapping_cpu') # bad data type string
 backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_mapping_cpu') # bad data type string
+backend_test.exclude('test_if_opt_cpu') # ValueError: 13 is not a valid AttributeType
+backend_test.exclude('test_if_seq_cpu') # NotImplementedError: op='SequenceConstruct' is not supported
 
 backend_test.exclude('test_scatternd_min_cpu') # min not yet supported
 backend_test.exclude('test_scatternd_max_cpu') # max not yet supported

test/external/external_test_onnx_ops.py

@@ -100,6 +100,25 @@ class TestMainOnnxOps(TestOnnxOps):
     self._test_resize_scales([0.01, 0.25, 0.5, 0.51, 0.6, 1.0, 1.5, 2.0, 3.5, 20.0], mode="cubic", exclude_outside=1)
     self._test_resize_scales([0.01, 0.25, 0.5, 0.51, 0.6, 1.0, 1.5, 2.0, 3.5, 20.0], mode="cubic", exclude_outside=0)
 
+  def _test_if(self, then_value, else_value):
+    then_out = onnx.helper.make_tensor_value_info("res", onnx.TensorProto.FLOAT, then_value.shape)
+    else_out = onnx.helper.make_tensor_value_info("res", onnx.TensorProto.FLOAT, else_value.shape)
+
+    then_const_node = onnx.helper.make_node("Constant", inputs=[], outputs=["res"], value=onnx.numpy_helper.from_array(then_value))
+    else_const_node = onnx.helper.make_node("Constant", inputs=[], outputs=["res"], value=onnx.numpy_helper.from_array(else_value))
+
+    then_body = onnx.helper.make_graph([then_const_node], "then_body", [], [then_out])
+    else_body = onnx.helper.make_graph([else_const_node], "else_body", [], [else_out])
+
+    self.helper_test_single_op("If", {"cond": np.array(False).astype(bool)}, {"then_branch": then_body, "else_branch": else_body}, ["res"])
+    self.helper_test_single_op("If", {"cond": np.array(True).astype(bool)}, {"then_branch": then_body, "else_branch": else_body}, ["res"])
+
+  def test_if_different_shapes_broadcastable(self):
+    self._test_if(np.array([[1], [2]]).astype(np.float32), np.array([[6, 5, 4, 3, 2, 1]]).astype(np.float32))
+
+  def test_if_different_shapes_not_broadcastable(self):
+    self._test_if(np.array([[1, 2, 3], [4, 5, 6]]).astype(np.float32), np.array([[6, 5, 4, 3, 2, 1]]).astype(np.float32))
+
   def test_resize_downsample_scales_linear_align_corners(self):
     # https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-131
     X = np.array([[[[1, 2, 3, 4], [5, 6, 7, 8]]]], dtype=np.float32)

test/external/external_test_train_gpt2.py

@@ -2,7 +2,7 @@
 import unittest
 
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.codegen.opt.search import Opt, OptOps
+from .search import Opt, OptOps
 from tinygrad.dtype import dtypes
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View

test/external/fuzz_linearizer.py

@@ -22,7 +22,7 @@ if os.getenv("VALIDATE_HCQ", 0) != 0:
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.codegen.opt.search import get_kernel_actions, bufs_from_lin
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.helpers import getenv, from_mv, prod, colored, Context, DEBUG, Timing

test/external/process_replay/process_replay.py

@@ -12,7 +12,7 @@ try:
   from tinygrad.renderer import Renderer, ProgramSpec
   from tinygrad.engine.realize import get_program
   from tinygrad.uop.ops import UOp, Ops, KernelInfo
-  from tinygrad.codegen.opt.kernel import Opt
+  from tinygrad.codegen.opt import Opt
   from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm
   from tinygrad.device import Device
 except ImportError as e:

test/external/speed_beam_v_hcopt.py

@@ -1,41 +0,0 @@
-from tinygrad import Device
-from tinygrad.helpers import getenv, DEBUG, BEAM
-from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
-from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
-from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer
-
-if __name__ == "__main__":
-  filter_reduce = bool(getenv("FILTER_REDUCE"))
-  ast_strs = load_worlds(filter_reduce=filter_reduce, filter_novariable=True)
-  dev = Device[Device.DEFAULT]
-
-  test_n = getenv("TEST_N", 10)
-  single = getenv("NUM", -1)
-  if single != -1: ast_strs = ast_strs[single:single+1]
-
-  beam_won, tested = 0, 0
-
-  for num, ast in enumerate(ast_strs[:test_n]):
-    def new_lin(): return ast_str_to_lin(ast, opts=dev.renderer)
-
-    k = new_lin()
-
-    if not (used_tensor_cores:=k.apply_tensor_cores(getenv("TC", 1))): k.apply_opts(hand_coded_optimizations(k))
-
-    assert BEAM > 0
-
-    lins = [(("tc" if used_tensor_cores else "hc"), k)]
-    if used_tensor_cores:
-      lins.append(("hc", new_lin()))
-      lins[-1][1].apply_opts(hand_coded_optimizations(lins[-1][1]))
-    kb = new_lin()
-    test_rawbuffers = bufs_from_lin(kb)    # allocate scratch buffers for optimization
-    lins.append((f"beam{BEAM.value}", beam_search(kb, test_rawbuffers, BEAM.value, bool(getenv("BEAM_ESTIMATE", 1)))))
-    timed = sorted([(nm, tk, time_linearizer(tk, test_rawbuffers, allow_test_size=False, clear_l2=True)) for nm, tk in lins], key=lambda x: x[2])
-    if DEBUG >= 1: print("  <  ".join(f"{nm:6s} : {lin.colored_shape(30, dense=True)} : {tm*1e6:8.2f} us" for nm, lin, tm in timed))
-
-    tested += 1
-    if timed[0][0].startswith("beam"):
-      beam_won += 1
-
-  print(f"{beam_won=} / {tested=} = {beam_won/tested:.3f}")

test/mockgpu/amd/amdgpu.py

@@ -1,5 +1,6 @@
 import ctypes, time
 from test.mockgpu.gpu import VirtGPU
+from test.mockgpu.helpers import _try_dlopen_remu
 from tinygrad.helpers import getbits, to_mv, init_c_struct_t
 import tinygrad.runtime.autogen.amd_gpu as amd_gpu, tinygrad.runtime.autogen.am.pm4_nv as pm4
 
@@ -24,19 +25,6 @@ WAIT_REG_MEM_FUNCTION_EQ  = 3 # ==
 WAIT_REG_MEM_FUNCTION_NEQ = 4 # !=
 WAIT_REG_MEM_FUNCTION_GEQ = 5 # >=
 
-REMU_PATHS = ["extra/remu/target/release/libremu.so", "libremu.so", "/usr/local/lib/libremu.so",
-              "extra/remu/target/release/libremu.dylib", "libremu.dylib", "/usr/local/lib/libremu.dylib", "/opt/homebrew/lib/libremu.dylib"]
-def _try_dlopen_remu():
-  for path in REMU_PATHS:
-    try:
-      remu = ctypes.CDLL(path)
-      remu.run_asm.restype = ctypes.c_int32
-      remu.run_asm.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32,
-        ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_void_p]
-    except OSError: pass
-    else: return remu
-  print("Could not find libremu.so")
-  return None
 remu = _try_dlopen_remu()
 
 def create_sdma_packets():

test/mockgpu/cuda/cuda.py

@@ -2,16 +2,14 @@ from __future__ import annotations
 from typing import Any
 import ctypes, time
 from tinygrad.runtime.autogen import cuda as orig_cuda
+from test.mockgpu.helpers import _try_dlopen_gpuocelot
 from tinygrad.helpers import mv_address
 
 for attr in dir(orig_cuda):
   if not attr.startswith('__'):
     globals()[attr] = getattr(orig_cuda, attr)
 
-try:
-  gpuocelot_lib = ctypes.CDLL(ctypes.util.find_library("gpuocelot"))
-  gpuocelot_lib.ptx_run.argtypes = [ctypes.c_char_p, ctypes.c_int, ctypes.POINTER(ctypes.c_void_p), ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int]  # noqa: E501
-except Exception: pass
+gpuocelot_lib = _try_dlopen_gpuocelot()
 
 # Global state
 class CUDAState:
@@ -130,7 +128,10 @@ def cuModuleUnload(hmod) -> int:
 def cuLaunchKernel(f, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, sharedMemBytes: int,
                    hStream: Any, kernelParams: Any, extra: Any) -> int:
   cargs = [ctypes.cast(getattr(extra, field[0]), ctypes.c_void_p) for field in extra._fields_]
-  gpuocelot_lib.ptx_run(ctypes.cast(f.value, ctypes.c_char_p), len(cargs), (ctypes.c_void_p*len(cargs))(*cargs), lx, ly, lz, gx, gy, gz, 0)
+  try: gpuocelot_lib.ptx_run(ctypes.cast(f.value, ctypes.c_char_p), len(cargs), (ctypes.c_void_p*len(cargs))(*cargs), lx, ly, lz, gx, gy, gz, 0)
+  except Exception as e:
+    print("Error in cuLaunchKernel:", e)
+    return orig_cuda.CUDA_ERROR_LAUNCH_FAILED
   return orig_cuda.CUDA_SUCCESS
 
 def cuDeviceComputeCapability(major, minor, dev: int) -> int:

test/mockgpu/helpers.py

@@ -0,0 +1,29 @@
+import ctypes, ctypes.util
+
+def _try_dlopen_gpuocelot():
+  GPUOCELOT_PATHS = [ctypes.util.find_library("gpuocelot")] if ctypes.util.find_library("gpuocelot") is not None else []
+  GPUOCELOT_PATHS += ["libgpuocelot.so", "/usr/local/lib/libgpuocelot.so",
+                      "libgpuocelot.dylib", "/usr/local/lib/libgpuocelot.dylib", "/opt/homebrew/lib/libgpuocelot.dylib"]
+  for path in GPUOCELOT_PATHS:
+    try:
+      gpuocelot_lib = ctypes.CDLL(path)
+      gpuocelot_lib.ptx_run.argtypes = [ctypes.c_char_p, ctypes.c_int, ctypes.POINTER(ctypes.c_void_p), ctypes.c_int, ctypes.c_int,
+        ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int]
+    except OSError: pass
+    else: return gpuocelot_lib
+  print("Could not find libgpuocelot.so")
+  return None
+
+def _try_dlopen_remu():
+  REMU_PATHS = ["extra/remu/target/release/libremu.so", "libremu.so", "/usr/local/lib/libremu.so",
+               "extra/remu/target/release/libremu.dylib", "libremu.dylib", "/usr/local/lib/libremu.dylib", "/opt/homebrew/lib/libremu.dylib"]
+  for path in REMU_PATHS:
+    try:
+      remu = ctypes.CDLL(path)
+      remu.run_asm.restype = ctypes.c_int32
+      remu.run_asm.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32,
+        ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_void_p]
+    except OSError: pass
+    else: return remu
+  print("Could not find libremu.so")
+  return None

test/mockgpu/nv/nvgpu.py

@@ -2,6 +2,7 @@ import ctypes, ctypes.util, time
 import tinygrad.runtime.autogen.nv_gpu as nv_gpu
 from enum import Enum, auto
 from test.mockgpu.gpu import VirtGPU
+from test.mockgpu.helpers import _try_dlopen_gpuocelot
 from tinygrad.helpers import to_mv, init_c_struct_t
 
 def make_qmd_struct_type():
@@ -16,10 +17,7 @@ def make_qmd_struct_type():
 qmd_struct_t = make_qmd_struct_type()
 assert ctypes.sizeof(qmd_struct_t) == 0x40 * 4
 
-try:
-  gpuocelot_lib = ctypes.CDLL(ctypes.util.find_library("gpuocelot"))
-  gpuocelot_lib.ptx_run.argtypes = [ctypes.c_char_p, ctypes.c_int, ctypes.POINTER(ctypes.c_void_p), ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int]  # noqa: E501
-except Exception: pass
+gpuocelot_lib = _try_dlopen_gpuocelot()
 
 class SchedResult(Enum): CONT = auto(); YIELD = auto() # noqa: E702
 
@@ -99,7 +97,10 @@ class GPFIFO:
     cargs = [ctypes.cast(args[i], ctypes.c_void_p) for i in range(args_cnt)] + [ctypes.cast(vals[i], ctypes.c_void_p) for i in range(vals_cnt)]
     gx, gy, gz = qmd.cta_raster_width, qmd.cta_raster_height, qmd.cta_raster_depth
     lx, ly, lz = qmd.cta_thread_dimension0, qmd.cta_thread_dimension1, qmd.cta_thread_dimension2
-    gpuocelot_lib.ptx_run(ctypes.cast(prg_addr, ctypes.c_char_p), args_cnt+vals_cnt, (ctypes.c_void_p*len(cargs))(*cargs), lx, ly, lz, gx, gy, gz, 0)
+    try:
+      gpuocelot_lib.ptx_run(ctypes.cast(prg_addr, ctypes.c_char_p), args_cnt+vals_cnt,
+        (ctypes.c_void_p*len(cargs))(*cargs), lx, ly, lz, gx, gy, gz, 0)
+    except Exception as e: print("failed to execute:", e)
     if qmd.release0_enable:
       rel0 = to_mv(qmd.release0_address_lower + (qmd.release0_address_upper << 32), 0x10).cast('Q')
       rel0[0] = qmd.release0_payload_lower + (qmd.release0_payload_upper << 32)

test/models/test_onnx.py

@@ -1,16 +1,9 @@
 #!/usr/bin/env python
-import os
-import time
 import unittest
 import numpy as np
-try:
-  import onnx
-except ModuleNotFoundError:
-  raise unittest.SkipTest("onnx not installed, skipping onnx test")
 from tinygrad.frontend.onnx import OnnxRunner
-from tinygrad.tensor import Tensor
 from tinygrad.device import Device
-from tinygrad.helpers import CI, fetch, temp, Context
+from tinygrad.helpers import fetch, Context
 
 try:
   from extra.onnx_helpers import validate
@@ -27,86 +20,9 @@ def run_onnx_torch(onnx_model, inputs):
     torch_out = torch_model(*[torch.tensor(x) for x in inputs.values()])
   return torch_out
 
-OPENPILOT_MODEL = "https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx"
-
 np.random.seed(1337)
 
 class TestOnnxModel(unittest.TestCase):
-  @unittest.skip("this isn't a test, it can't fail")
-  def test_benchmark_openpilot_model(self):
-    onnx_model = fetch(OPENPILOT_MODEL)
-    run_onnx = OnnxRunner(onnx_model)
-    def get_inputs():
-      np_inputs = {
-        "input_imgs": np.random.randn(*(1, 12, 128, 256)),
-        "big_input_imgs": np.random.randn(*(1, 12, 128, 256)),
-        "desire": np.zeros((1, 100, 8)),
-        "traffic_convention": np.array([[1., 0.]]),
-        "nav_features": np.zeros((1, 256)),
-        "features_buffer": np.zeros((1, 99, 128)),
-    }
-      inputs = {k:Tensor(v.astype(np.float32), requires_grad=False) for k,v in np_inputs.items()}
-      return inputs
-
-    for _ in range(7):
-      inputs = get_inputs()
-      st = time.monotonic()
-      tinygrad_out = run_onnx(inputs)['outputs']
-      mt = time.monotonic()
-      tinygrad_out.realize()
-      mt2 = time.monotonic()
-      tinygrad_out = tinygrad_out.numpy()
-      et = time.monotonic()
-      if not CI:
-        print(f"ran openpilot model in {(et-st)*1000.0:.2f} ms, waited {(mt2-mt)*1000.0:.2f} ms for realize, {(et-mt2)*1000.0:.2f} ms for GPU queue")
-
-    if not CI:
-      import cProfile
-      import pstats
-      inputs = get_inputs()
-      pr = cProfile.Profile(timer=time.perf_counter_ns, timeunit=1e-6)
-      pr.enable()
-    tinygrad_out = run_onnx(inputs)['outputs']
-    tinygrad_out.realize()
-    tinygrad_out = tinygrad_out.numpy()
-    if not CI:
-      pr.disable()
-      stats = pstats.Stats(pr)
-      stats.dump_stats(temp("net.prof"))
-      os.system(f"flameprof {temp('net.prof')} > {temp('prof.svg')}")
-      ps = stats.sort_stats(pstats.SortKey.TIME)
-      ps.print_stats(30)
-
-  def test_openpilot_model(self):
-    onnx_model = fetch(OPENPILOT_MODEL)
-    run_onnx = OnnxRunner(onnx_model)
-    print("got run_onnx")
-    inputs = {
-      "input_imgs": np.random.randn(*(1, 12, 128, 256)),
-      "big_input_imgs": np.random.randn(*(1, 12, 128, 256)),
-      "desire": np.zeros((1, 100, 8)),
-      "traffic_convention": np.array([[1., 0.]]),
-      "nav_features": np.zeros((1, 256)),
-      "features_buffer": np.zeros((1, 99, 128)),
-    }
-    inputs = {k:v.astype(np.float32) for k,v in inputs.items()}
-
-    st = time.monotonic()
-    print("****** run onnx ******")
-    tinygrad_out = run_onnx(inputs)['outputs']
-    mt = time.monotonic()
-    print("****** realize ******")
-    tinygrad_out.realize()
-    mt2 = time.monotonic()
-    tinygrad_out = tinygrad_out.numpy()
-    et = time.monotonic()
-    print(f"ran openpilot model in {(et-st)*1000.0:.2f} ms, waited {(mt2-mt)*1000.0:.2f} ms for realize, {(et-mt2)*1000.0:.2f} ms for GPU queue")
-
-    onnx_model = onnx.load(fetch(OPENPILOT_MODEL))
-    torch_out = run_onnx_torch(onnx_model, inputs).numpy()
-    print(tinygrad_out, torch_out)
-    np.testing.assert_allclose(tinygrad_out, torch_out, atol=1e-4, rtol=1e-2)
-
   @unittest.skip("slow")
   def test_efficientnet(self):
     input_name, input_new = "images:0", True

test/opt/test_gen_float4.py

@@ -0,0 +1,180 @@
+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
+
+@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "need backends that support float4")
+class TestFloat4(unittest.TestCase):
+  @staticmethod
+  def count_float4(uops: list[UOp], n=4):
+    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.float.vec(n)]),
+            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.float.vec(n)]))
+  @staticmethod
+  def count_half4(uops: list[UOp]):
+    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.half.vec(4)]),
+            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.half.vec(4)]))
+
+  def test_float4_basic(self):
+    a = Tensor.empty(2, 8).realize()
+    b = Tensor.empty(2, 8).realize()
+    c = a + b
+
+    s = c.schedule()[0]
+    realized_ast = s.ast
+    opts_to_apply = [Opt(op=OptOps.UPCAST, axis=0, arg=4)]
+    program = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts_to_apply)
+
+    assert TestFloat4.count_float4(program.uops) == (2, 1)
+
+  @unittest.skipIf(Device.DEFAULT in {"CPU", "LLVM"} and AMX, "CPU with AMX upcasts float up to size 16")
+  def test_float4_multidim(self):
+    a = Tensor.empty(2, 8).realize()
+    b = Tensor.empty(2, 8).realize()
+    c = a + b
+
+    s = c.schedule()[0]
+    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=2)]).uops
+    assert TestFloat4.count_float4(uops) == (4, 2)
+
+  @unittest.skipUnless(Device.DEFAULT in {"CPU", "LLVM"} and AMX, "Only CPU with AMX upcasts float up to size 16")
+  def test_float4_multidim_amx(self):
+    def kernel_for_shape(size, shift):
+      a = Tensor.empty(2, size).realize()
+      b = Tensor.empty(2, size).realize()
+      c = a + b
+
+      s = c.schedule()[0]
+      return get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=shift)]).uops
+
+    sizes = [12, 8, 16]
+    shifts = [3, 2, 4]
+    expected_upcast_size = [4, 8, 16]
+    expected_output = [(6,3), (2,1), (2,1)]
+
+    for i in range(len(sizes)):
+      assert TestFloat4.count_float4(kernel_for_shape(sizes[i], shifts[i]), expected_upcast_size[i]) == expected_output[i]
+
+  def test_float4_unaligned_load(self):
+    a = Tensor.empty(9).realize().shrink(((1, 9),))
+    b = Tensor.empty(9).realize().shrink(((1, 9),))
+    c = a + b
+
+    s = c.schedule()[0]
+    realized_ast = s.ast
+    opts_to_apply = [Opt(op=OptOps.UPCAST, axis=0, arg=4)]
+    program = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts_to_apply)
+
+    assert TestFloat4.count_float4(program.uops) == (0, 1)
+
+  @unittest.skipIf(Device.DEFAULT in {"CPU", "LLVM"} and AMX, "CPU with AMX upcasts float up to size 16")
+  def test_float4_multidim_unaligned_load(self):
+    a = Tensor.empty(2, 9).realize().shrink(((0, 2), (1, 9),))
+    b = Tensor.empty(2, 9).realize().shrink(((0, 2), (1, 9),))
+    c = a + b
+
+    s = c.schedule()[0]
+    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]).uops
+
+    assert TestFloat4.count_float4(uops) == (0, 2)
+
+  @unittest.skipUnless(Device.DEFAULT in {"CPU", "LLVM"} and AMX, "Only CPU with AMX upcasts float up to size 16")
+  def test_float4_multidim_unaligned_load_amx(self):
+    def kernel_for_shape(size, shift):
+      a = Tensor.empty(2, size).realize().shrink(((0, 2), (1, size),))
+      b = Tensor.empty(2, size).realize().shrink(((0, 2), (1, size),))
+      c = a + b
+
+      s = c.schedule()[0]
+      return get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=shift)]).uops
+
+    sizes = [13, 9, 17]
+    shifts = [3, 2, 4]
+    expected_upcast_size = [4, 8, 16]
+    expected_output = [(0,3), (0,1), (0,1)]
+
+    for i in range(len(sizes)):
+      assert TestFloat4.count_float4(kernel_for_shape(sizes[i], shifts[i]), expected_upcast_size[i]) == expected_output[i]
+
+  def test_float4_sometimes_unaligned(self):
+    a = Tensor.empty(1, 1, 8).realize()
+    b = Tensor.empty(1, 1, 5).realize().shrink(((0, 1), (0, 1), (1, 5)))
+    c = a.conv2d(b)
+    # only the first and last conv dot products are aligned in a, and b is never aligned, so no
+    # float4 should be emitted (the reduce axis of size 4 is the float4 axis here)
+
+    s = c.schedule()[0]
+    uops = get_program(s.ast, opts=[Opt(op=OptOps.UNROLL, axis=0, arg=4)]).uops
+
+    assert TestFloat4.count_float4(uops) == (0, 0)
+
+  def test_float4_multidim_sometimes_unaligned(self):
+    a = Tensor.empty(1, 1, 7).realize()
+    b = Tensor.empty(1, 1, 5).realize().shrink(((0, 1), (0, 1), (1, 5)))
+    c = a.conv2d(b)
+    # the first conv dot product is aligned in a. If we upcast the output and reduce
+    # dimension, then we could do float4 for only that one set of loads, but we currently
+    # don't.
+    # UPDATE: now we do this fusion
+
+    s = c.schedule()[0]
+    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]).uops
+
+    assert TestFloat4.count_float4(uops) in {(0,1), (1,1)}
+
+  def test_float4_expand(self):
+    a = Tensor.empty(9).realize().shrink(((1, 9),))
+    b = Tensor.empty(2).realize().reshape((2, 1)).expand((2,4)).reshape((8,))
+    c = a + b
+
+    # we will upcast the top axis of sz 4. they should not be coalesced into float4,
+    # since the top axis is not contiguous.
+
+    s = c.schedule()[0]
+    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4)]).uops
+
+    assert TestFloat4.count_float4(uops) == (0, 1)
+
+  def test_float4_heterogeneous(self):
+    a = Tensor.empty(8).realize()
+    b = Tensor.empty(9).realize().shrink(((1, 9),))
+    c = a + b
+
+    # should float4 b but not a
+
+    s = c.schedule()[0]
+    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4)]).uops
+
+    assert TestFloat4.count_float4(uops) == (1, 1)
+
+  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

@@ -0,0 +1,340 @@
+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
+from test.test_linearizer import helper_linearizer_opt
+
+class TestKernelOpts(unittest.TestCase):
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
+  def test_local_and_grouped_reduce(self):
+    N = 128
+    Tensor.manual_seed(1882)
+    a = Tensor.rand(4, 4, N, N)
+    b = Tensor.rand(4, 4, N)
+    r = (b.sqrt() + ((a+1).sum(axis=3).exp()))
+    helper_linearizer_opt(r, [
+      [Opt(OptOps.LOCAL, 0, 2)],
+      [Opt(OptOps.LOCAL, 0, 8)],
+      [Opt(OptOps.LOCAL, 0, 16)], # Checking how it works with locals
+      [Opt(OptOps.GROUPTOP, 0, 2)],
+      [Opt(OptOps.GROUPTOP, 0, 32)],
+      [Opt(OptOps.GROUPTOP, 0, 64)], # Checking how it works with grouped reduce
+      [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.GROUPTOP, 0, 2)],
+      [Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.GROUPTOP, 0, 16)],
+      [Opt(OptOps.LOCAL, 0, 32), Opt(OptOps.GROUPTOP, 0, 2)],
+      # Checking how it works with locals + grouped reduce
+      [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.GROUPTOP, 0, 64)],
+      # Checking how it works with locals + grouped reduce + upcasts
+      [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.GROUPTOP, 0, 2), Opt(OptOps.UPCAST, 0, 8), Opt(OptOps.UNROLL, 1, 4)],
+      # many local + many group
+      [Opt(OptOps.GROUP, 0, 2)] * 4,
+      [Opt(OptOps.LOCAL, 0, 2)] * 4,
+      [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.GROUP, 0, 2)] * 4,
+    ])
+
+  def test_upcasts(self):
+    N = 16
+    Tensor.manual_seed(1772)
+    a = Tensor.rand(N, N)
+    b = Tensor.rand(N, N)
+    r = (a+b).sqrt() * ((a+1).exp())
+    helper_linearizer_opt(r, [
+      [Opt(OptOps.UPCAST, 0, 2)],
+      [Opt(OptOps.UPCAST, 0, 4)],
+      [Opt(OptOps.UPCAST, 0, 8)], # Checking how it works with upcasts
+    ])
+
+  def test_full_upcast(self):
+    Tensor.manual_seed(1772)
+    a = Tensor.rand(4)
+    b = Tensor.rand(4)
+    r = (a+b).sqrt() * ((a+1).exp())
+    helper_linearizer_opt(r, [
+      [Opt(OptOps.UPCAST, 0, 4)], # Checking how it works with upcasts
+    ])
+
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
+  def test_matmul(self):
+    N = 128
+    Tensor.manual_seed(1552)
+    a = Tensor.rand(N, N)
+    b = Tensor.rand(N, N)
+    r = a@b
+    helper_linearizer_opt(r, [
+      [Opt(OptOps.UPCAST, 0, 2)],
+      [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4)], # Checking how it works with upcasts
+      [Opt(OptOps.LOCAL, 0, 2)],
+      [Opt(OptOps.LOCAL, 1, 32)],
+      [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 4)],
+      [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 32)],
+      [Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.LOCAL, 1, 8)], # Checking how it works with locals
+      [Opt(OptOps.GROUPTOP, 0, 2)],
+      [Opt(OptOps.GROUPTOP, 0, 32)],
+      [Opt(OptOps.GROUPTOP, 0, 32), Opt(OptOps.UNROLL, 0, 4)], # Checking how it works with grouped_reduce
+      [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 1, 2), Opt(OptOps.GROUPTOP, 0, 32)],
+      [Opt(OptOps.LOCAL, 0, 8), Opt(OptOps.GROUPTOP, 0, 32)],
+      [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 0, 8), Opt(OptOps.GROUPTOP, 0, 4)], # Checking how it works with local+grouped_reduce
+      # Checking all together
+      [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.GROUPTOP, 0, 8), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4),
+       Opt(OptOps.UPCAST, 1, 2)],
+      # Full global upcast + local
+      [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.GROUPTOP, 0, 8), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 8)],
+    ])
+
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
+  def test_double_reduce(self):
+    N = 128
+    Tensor.manual_seed(1552)
+    a = Tensor.rand(8, N, 8, N)
+    r = a.sum(axis=(1,3))
+    helper_linearizer_opt(r, [
+      # openCL / GPU=1 is 256 max threads
+      [Opt(OptOps.GROUPTOP, 0, 2)], [Opt(OptOps.GROUPTOP, 0, 32)],
+      [Opt(OptOps.GROUPTOP, 1, 2)], [Opt(OptOps.GROUPTOP, 1, 32)], # Checking how it works with 1 grouped_reduce.
+      [Opt(OptOps.GROUPTOP, 0, 2), Opt(OptOps.GROUPTOP, 1, 2)],
+      [Opt(OptOps.GROUPTOP, 0, 16), Opt(OptOps.GROUPTOP, 1, 2)],
+      [Opt(OptOps.GROUPTOP, 0, 4), Opt(OptOps.GROUPTOP, 1, 64)], # Checking how it works with 2 grouped_reduces.
+      [Opt(OptOps.GROUPTOP, 0, 16), Opt(OptOps.GROUPTOP, 1, 2), Opt(OptOps.UNROLL, 0, 4)],
+      [Opt(OptOps.GROUPTOP, 0, 2), Opt(OptOps.GROUPTOP, 1, 32), Opt(OptOps.UNROLL, 2, 4)], # Checking how it works with 2 grouped_reduces + upcasts.
+      [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 4), Opt(OptOps.GROUPTOP, 0, 4), Opt(OptOps.GROUPTOP, 1, 4)],
+      # Checking how it works with 2 grouped_reduces + upcasts + locals.
+      [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 4), Opt(OptOps.GROUPTOP, 0, 2), Opt(OptOps.GROUPTOP, 1, 32), Opt(OptOps.UNROLL, 1, 4)],
+      [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 1, 2), Opt(OptOps.GROUPTOP, 0, 8), Opt(OptOps.GROUPTOP, 1, 4), Opt(OptOps.UPCAST, 0, 2)],
+      [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 1, 2), Opt(OptOps.GROUPTOP, 0, 8), Opt(OptOps.GROUPTOP, 1, 4), Opt(OptOps.UPCAST, 0, 2),
+       Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4)], # Checking how it works with 2 grouped_reduces + upcasts + locals.
+      [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 4), Opt(OptOps.GROUPTOP, 0, 4), Opt(OptOps.GROUPTOP, 1, 4), Opt(OptOps.UPCAST, 0, 2),
+       Opt(OptOps.UPCAST, 0, 2)], # No globals
+    ])
+
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
+  @unittest.skipUnless(any(tc.dtype_in == tc.dtype_out == dtypes.half for tc in Device[Device.DEFAULT].renderer.tensor_cores),
+                      "test requires tensor cores with accumulation in half") # testing with half suffices.
+  def test_tensor_core_opts(self):
+    N = 128
+    Tensor.manual_seed(1552)
+    a, b = Tensor.rand(N, N, dtype=dtypes.half), Tensor.rand(N, N, dtype=dtypes.half)
+    r = a.matmul(b, dtype=dtypes.half)
+    atol, rtol = 0.25, 0.01
+    helper_linearizer_opt(r, [
+      [],
+      [Opt(OptOps.UPCAST, 0, 4)],
+      [Opt(OptOps.UPCAST, 1, 4)],
+      [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4)], # check upcasts
+      [Opt(OptOps.UNROLL, 0, 2)], # check unroll
+      [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 0, 2)], # check combo of unroll and local
+      [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 2)],
+      [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)],
+      [Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 0, 4)], # check permutations
+      [Opt(OptOps.UNROLL, 0, 2), Opt(OptOps.UPCAST, 0, 4)],
+      [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 0, 2), Opt(OptOps.UPCAST, 1, 4)],
+      [Opt(OptOps.UNROLL, 0, 2), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 0, 4)],
+    ], apply_tc=True, atol=atol, rtol=rtol)
+
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
+  @unittest.skipUnless(any(tc.dtype_in == tc.dtype_out == dtypes.half for tc in Device[Device.DEFAULT].renderer.tensor_cores),
+                      "test requires tensor cores with accumulation in half") # testing with half suffices.
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  def test_tensor_core_opts_locals(self):
+    N = 128
+    Tensor.manual_seed(1552)
+    a, b = Tensor.rand(N, N, dtype=dtypes.half), Tensor.rand(N, N, dtype=dtypes.half)
+    r = a.matmul(b, dtype=dtypes.half)
+    atol, rtol = 0.25, 0.01
+    helper_linearizer_opt(r, [
+      [Opt(OptOps.UNROLL, 0, 0)], # check full unroll of reduce with locals
+      [Opt(OptOps.LOCAL, 0, 4)], # check local
+      [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.LOCAL, 0, 2)],
+      [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 2), Opt(OptOps.UPCAST, 0, 4)],
+    ], apply_tc=True, atol=atol, rtol=rtol)
+
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared memory")
+  @unittest.skipUnless(any(tc.dtype_in == tc.dtype_out == dtypes.half for tc in Device[Device.DEFAULT].renderer.tensor_cores),
+                      "test requires tensor cores with accumulation in half") # testing with half suffices.
+  # NOTE: the METAL test is broken, likely due to a compiler bug. passes on CI with -O0 and with default opt level locally on M3
+  @unittest.skipIf(Device.DEFAULT == "METAL", "broken for METAL")
+  @unittest.skip("feature was removed")
+  def test_tensor_core_opts_group(self):
+    N = 128
+    Tensor.manual_seed(1552)
+    a, b = Tensor.rand(N, N, dtype=dtypes.half), Tensor.rand(N, N, dtype=dtypes.half)
+    r = a.matmul(b, dtype=dtypes.half)
+    atol, rtol = 0.25, 0.01
+    helper_linearizer_opt(r, [
+      [Opt(OptOps.GROUP, 0, 2)],
+      [Opt(OptOps.GROUPTOP, 0, 4)],
+      [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.GROUP, 0, 2)],
+      [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.GROUP, 0, 2)],
+      [Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.GROUP, 0, 2)],
+      [Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.GROUP, 0, 2)],
+      [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.GROUPTOP, 0, 8), Opt(OptOps.UNROLL, 0, 2), Opt(OptOps.UPCAST, 1, 2)],
+    ], 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
+    Tensor.manual_seed(289)
+    a = Tensor.rand(N, N)
+    b = Tensor.rand(N, N)
+    helper_linearizer_opt(a@b, [
+      [Opt(OptOps.PADTO, 0, 32)],
+      [Opt(OptOps.PADTO, 1, 32)],
+      [Opt(OptOps.PADTO, 2, 32)],
+      [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.PADTO, 1, 32)],
+      [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.PADTO, 1, 32), Opt(OptOps.PADTO, 2, 32)],
+      # can optimize further post PADTO
+      [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.PADTO, 1, 32), Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UPCAST, 1, 2),],
+    ])
+
+  def test_padto_upcasted_not_ok(self):
+    N = 4
+    a = Tensor.rand(N, N)
+    b = Tensor.rand(N, N)
+    helper_linearizer_opt(a@b, [
+      [Opt(OptOps.UPCAST, 0, 0)],
+      [Opt(OptOps.UPCAST, 1, 0)],
+      [Opt(OptOps.UNROLL, 0, 0)],
+      [Opt(OptOps.PADTO, 0, 8)],
+      [Opt(OptOps.PADTO, 1, 8)],
+      [Opt(OptOps.PADTO, 2, 8)],
+    ])
+    with self.assertRaises(KernelOptError):
+      helper_linearizer_opt(a@b, [[Opt(OptOps.UPCAST, 0, 0), Opt(OptOps.PADTO, 1, 8)]])
+    with self.assertRaises(KernelOptError):
+      helper_linearizer_opt(a@b, [[Opt(OptOps.UPCAST, 1, 0), Opt(OptOps.PADTO, 1, 8)]])
+    with self.assertRaises(KernelOptError):
+      helper_linearizer_opt(a@b, [[Opt(OptOps.UNROLL, 0, 0), Opt(OptOps.PADTO, 2, 8)]])
+
+  def test_padto_sum_ok(self):
+    N = 18 * 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)))
+
+    helper_linearizer_opt(a.sum(0), [
+      [Opt(OptOps.PADTO, 0, 32)],
+      [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.UPCAST, 0, 8),],
+    ])
+    helper_linearizer_opt(a.sum(1), [
+      [Opt(OptOps.PADTO, 0, 32)],
+      [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.UPCAST, 0, 8),],
+    ])
+
+    # can pad sum reduce axis if there's no unsafe ops prior to sum
+    for axis in (0, 1):
+      helper_linearizer_opt(a.sum(), [[Opt(OptOps.PADTO, axis, 32)],])
+      helper_linearizer_opt(a.sum(0), [[Opt(OptOps.PADTO, axis, 32)],])
+      helper_linearizer_opt(b.sum(), [[Opt(OptOps.PADTO, axis, 32)],])
+      helper_linearizer_opt(b.sum(0), [[Opt(OptOps.PADTO, axis, 32)],])
+      helper_linearizer_opt(b.sum(dtype=dtypes.bool), [[Opt(OptOps.PADTO, axis, 32)],])
+      # TODO: why?
+      if Device.DEFAULT != "WEBGPU":
+        helper_linearizer_opt(b.sum(0, dtype=dtypes.bool), [[Opt(OptOps.PADTO, axis, 32)],])
+        helper_linearizer_opt(b.sum(1, dtype=dtypes.bool), [[Opt(OptOps.PADTO, axis, 32)],])
+
+    # having unsafe ops after sum is fine
+    helper_linearizer_opt(a.sum().exp(), [[Opt(OptOps.PADTO, 0, 32)],])
+    helper_linearizer_opt(a.sum(0).exp(), [[Opt(OptOps.PADTO, 1, 32)],])
+
+  def test_padto_sum_not_ok(self):
+    N = 18 * 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
+    with self.assertRaises(KernelOptError):
+      helper_linearizer_opt(a.exp().sum(), [[Opt(OptOps.PADTO, 0, 32)],])
+    with self.assertRaises(KernelOptError):
+      helper_linearizer_opt(a.exp().sum(0), [[Opt(OptOps.PADTO, 1, 32)],])
+
+    b = a < 1
+    # lt is not safe to pad
+    with self.assertRaises(KernelOptError):
+      helper_linearizer_opt(b.sum(), [[Opt(OptOps.PADTO, 0, 32)],])
+    with self.assertRaises(KernelOptError):
+      helper_linearizer_opt(b.sum(0), [[Opt(OptOps.PADTO, 1, 32)],])
+
+  def test_padto_max(self):
+    N = 18 * 18
+    # NOTE: this setup prevents 17 * 17 contiguous merged into one axis
+    a = -Tensor.rand(N, N).shrink(((0, 17), (0, 17))) * 100
+
+    helper_linearizer_opt(a.max(0), [
+      [Opt(OptOps.PADTO, 0, 32)],
+      [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.UPCAST, 0, 8),],
+    ])
+    helper_linearizer_opt(a.max(1), [
+      [Opt(OptOps.PADTO, 0, 32)],
+      [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.UPCAST, 0, 8),],
+    ])
+
+    # cannot pad max kernel on reduce
+    with self.assertRaises(KernelOptError):
+      helper_linearizer_opt(a.max(), [[Opt(OptOps.PADTO, 0, 32)],])
+    with self.assertRaises(KernelOptError):
+      helper_linearizer_opt(a.max(0), [[Opt(OptOps.PADTO, 1, 32)],])
+
+  def test_padto_where(self):
+    Tensor.manual_seed(0)
+    N = 17 * 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)],
+      [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.UPCAST, 0, 8),],
+    ])
+
+  def test_padto_where_multioutput(self):
+    Tensor.manual_seed(0)
+    N = 17 * 17
+    r = Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1
+    a0 = r.where(1, 0)
+    a1 = r.where(2, 0)
+    helper_linearizer_opt([a0.max(0), a1.max(0)], [
+      [Opt(OptOps.PADTO, 0, 32)],
+      [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.UPCAST, 0, 8),],
+    ])
+
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
+  def test_color_shapes_with_local(self):
+    N = 32
+    Tensor.manual_seed(1552)
+    a = Tensor.rand(N, N)
+    b = Tensor.rand(N, N)
+    r = a@b
+    opts_shapes = [
+      ([Opt(OptOps.LOCAL, 0, 2)], [("blue",16),("blue",32),("cyan",2),("red",32)]),
+      ([Opt(OptOps.LOCAL, 0, 2),Opt(OptOps.GROUP, 0, 2)], [("blue",16),("blue",32),("cyan",2),("green",2),("red",16)]),
+      # check to ensure local_dims are stable for full UNROLL of the first reduce
+      ([Opt(OptOps.LOCAL, 0, 2),Opt(OptOps.UNROLL, 0, 0)], [("blue",16),("blue",32),("cyan",2),("magenta",32)]),
+      ([Opt(OptOps.UNROLL, 0, 0),Opt(OptOps.LOCAL, 0, 2)], [("blue",16),("blue",32),("cyan",2),("magenta",32)]),
+      # check behavior for full UNROLL on an existing GROUP
+      ([Opt(OptOps.LOCAL, 0, 2),Opt(OptOps.GROUP, 0, 0),Opt(OptOps.UNROLL, 0, 2)], [("blue",16),("blue",32),("cyan",2),("green",16),("magenta",2)]),
+      ([Opt(OptOps.LOCAL, 0, 2),Opt(OptOps.GROUP, 0, 0),Opt(OptOps.UNROLL, 0, 0)], [("blue",16),("blue",32),("cyan",2),("magenta",32)]),
+      ([Opt(OptOps.GROUP, 0, 0),Opt(OptOps.LOCAL, 0, 2),Opt(OptOps.UNROLL, 0, 0)], [("blue",16),("blue",32),("cyan",2),("magenta",32)]),
+      ([Opt(OptOps.GROUP, 0, 2),Opt(OptOps.UNROLL, 0, 0)], [("blue",32),("blue",32),("red",16),("magenta",2)]),
+    ]
+    helper_linearizer_opt(r, [x[0] for x in opts_shapes], color_sizes=[x[1] for x in opts_shapes])
+
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
+  def test_arange_opts(self):
+    a = Tensor.arange(128)
+    helper_linearizer_opt(a, [
+      [Opt(OptOps.GROUP, 0, 32)],
+      [Opt(OptOps.GROUPTOP, 0, 32)],
+      [Opt(op=OptOps.LOCAL, axis=0, arg=8)],
+      [Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UPCAST, axis=0, arg=0)],
+      [Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.GROUP, axis=0, arg=8)],
+      [Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.UNROLL, axis=1, arg=4)], # noqa: E501
+    ])
+
+if __name__ == '__main__':
+  unittest.main()

test/opt/test_tensor_cores.py

@@ -0,0 +1,191 @@
+import numpy as np
+import unittest
+from dataclasses import replace
+
+from tinygrad import Device, Tensor, dtypes
+from tinygrad.tensor import _to_np_dtype
+from tinygrad.uop.ops import Ops
+from tinygrad.dtype import DType
+from tinygrad.device import is_dtype_supported
+from tinygrad.helpers import AMX, CI, AMD_LLVM
+from tinygrad.engine.realize import CompiledRunner, get_program
+from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
+
+# TODO: write a clean version of this
+from test.test_linearizer import helper_realized_ast, helper_linearizer_opt
+
+def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, dtype_out:DType, axis:int=0, tc_select:int=-1, tc_opt:int=0,
+                                         ensure_triggered:bool=True):
+  a, b = Tensor.rand(M, K, dtype=dtype_in), Tensor.rand(K, N, dtype=dtype_in)
+  r = a.matmul(b, dtype=dtype_out)
+  sched = r.schedule()
+  realized_ast = sched[-1].ast
+  opts_to_apply = [Opt(OptOps.TC, axis, (tc_select, tc_opt, 1))]
+
+  if ensure_triggered:
+    program = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts_to_apply)
+    wmmas = len([uop for uop in program.uops if uop.op is Ops.WMMA])
+    tcs = len([x for x in program.applied_opts if x.op is OptOps.TC])
+    assert wmmas > 0, "tensor core not triggered"
+    assert tcs == 1, "tensor core opt not included"
+  else:
+    try:
+      program = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts_to_apply)
+      assert False, "OptOps.TC triggered, expected KernelOptError"
+    except KernelOptError: pass
+
+def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axis:int=0, tc_select:int=-1, tc_opt:int=0, use_tensor_cores:int=1):
+  a, b = Tensor.rand(M, K, dtype=dtype_in), Tensor.rand(K, N, dtype=dtype_in)
+  np_a, np_b = a.numpy(), b.numpy()
+  r = a.matmul(b, dtype=dtype_out)
+  if dtype_in == dtypes.bfloat16: r = r.float()
+  realized_ast, bufs = helper_realized_ast(r)
+  opts = [Opt(op=OptOps.TC, axis=axis, arg=(tc_select, tc_opt, use_tensor_cores))]
+  prg = CompiledRunner(replace(get_program(realized_ast, opts=opts), device=Device.DEFAULT))
+  if use_tensor_cores == 1: assert len([uop for uop in prg.p.uops if uop.op is Ops.WMMA]) > 0, "wmma not triggered"
+  assert len([x for x in prg.p.uops[-1].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
+  prg.exec(bufs)
+  if dtype_in == dtypes.half: tc_atol, tc_rtol = 1e-2, 1e-3
+  elif dtype_in == dtypes.bfloat16: tc_atol, tc_rtol = 1e-2, 1e-2
+  else: tc_atol, tc_rtol = 5e-3, 1e-4
+  c = bufs[0].numpy().reshape((M,N))
+  np.testing.assert_allclose(c, np_a @ np_b, atol=tc_atol, rtol=tc_rtol)
+
+class TestTensorCores(unittest.TestCase):
+  # TODO: don't skip bf16 for real device (METAL, AMD)
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
+  def test_tensor_cores(self):
+    for tc in Device[Device.DEFAULT].renderer.tensor_cores:
+      if not is_dtype_supported(tc.dtype_in) or not is_dtype_supported(tc.dtype_out): continue
+      # for AMX, tc.dims[2] == 1 so reduceop is None thus tensor_cores are not triggered
+      helper_tc_allclose(tc.dims[0], tc.dims[1], 2 if AMX else tc.dims[2], tc.dtype_in, tc.dtype_out, axis=0, tc_opt=0)
+
+  @unittest.skipIf(Device.DEFAULT == "PYTHON", "not generated on EMULATED device")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
+  def test_tensor_cores_codegen(self):
+    for tc in Device[Device.DEFAULT].renderer.tensor_cores:
+      if not is_dtype_supported(tc.dtype_in) or not is_dtype_supported(tc.dtype_out): continue
+      n, m, k = tc.dims[0], tc.dims[1], 2 if AMX else tc.dims[2]
+      a, b = Tensor.rand(m, k, dtype=tc.dtype_in), Tensor.rand(k, n, dtype=tc.dtype_in)
+      r = a.matmul(b, dtype=tc.dtype_out)
+      prg = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.TC, axis=0, arg=(-1, 2, 1))])
+      if Device.DEFAULT == "LLVM":
+        assert "0x201000" in prg.src
+      elif Device.DEFAULT == "AMD" and AMD_LLVM:
+        assert "@llvm.amdgcn.wmma" in prg.src
+      elif Device[Device.DEFAULT].renderer.suffix == "PTX":
+        assert "mma.sync.aligned" in prg.src
+      else:
+        assert "__WMMA_" in prg.src
+
+  @unittest.skipIf((Device.DEFAULT == "AMD") or (Device.DEFAULT == "PYTHON" and Device.default.renderer.device == "AMD"), "broken for AMD")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
+  def test_tensor_cores_padded(self):
+    for tc in Device[Device.DEFAULT].renderer.tensor_cores:
+      if not is_dtype_supported(tc.dtype_in) or not is_dtype_supported(tc.dtype_out): continue
+      helper_tc_allclose(tc.dims[0]+(pad:=1), tc.dims[1]+pad, tc.dims[2]+pad, tc.dtype_in, tc.dtype_out, tc_opt=2)
+
+  # AMD compiler bug: AMD miscompiles non-zero padded tc kernels with -O3, producing wrong results, nans or hang (see #9606)
+  # Internal bug: zero-stride dimensions combined with a mask may produce wrong index/valid for pad == 1 on AMD
+  @unittest.skipUnless((Device.DEFAULT == "AMD") or (Device.DEFAULT == "PYTHON" and Device.default.renderer.device == "AMD"), "test for AMD's tc")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
+  @unittest.skip("warp elements not duplicated properly across lanes")
+  def test_tensor_cores_padded_amd(self):
+    for tc in Device[Device.DEFAULT].renderer.tensor_cores:
+      if not is_dtype_supported(tc.dtype_in) or not is_dtype_supported(tc.dtype_out): continue
+      helper_tc_allclose(tc.dims[0]+(pad:=1), tc.dims[1]+pad, tc.dims[2]+pad, tc.dtype_in, tc.dtype_out, tc_opt=2)
+
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
+  def test_tensor_cores_padded_uops(self):
+    for tc in Device[Device.DEFAULT].renderer.tensor_cores:
+      pad = 1
+
+      # check that TC is triggered for TC_OPT=2
+      helper_tc_ensure_uops_and_opts_count(tc.dims[0]+pad, tc.dims[1]+pad, tc.dims[2]+pad,
+                                           tc.dtype_in, tc.dtype_out, tc_opt=2, ensure_triggered=True)
+
+      # check that TC is not triggered for TC_OPT<2
+      helper_tc_ensure_uops_and_opts_count(tc.dims[0]+pad, tc.dims[1]+pad, tc.dims[2]+pad,
+                                           tc.dtype_in, tc.dtype_out, tc_opt=1, ensure_triggered=False)
+      helper_tc_ensure_uops_and_opts_count(tc.dims[0]+pad, tc.dims[1]+pad, tc.dims[2]+pad,
+                                           tc.dtype_in, tc.dtype_out, tc_opt=0, ensure_triggered=False)
+
+      # check excessive padding doesn't trigger padded TC in TC_OPT=2
+      helper_tc_ensure_uops_and_opts_count(tc.dims[0]//4, tc.dims[1], tc.dims[2], tc.dtype_in, tc.dtype_out, tc_opt=2, ensure_triggered=False)
+      helper_tc_ensure_uops_and_opts_count(tc.dims[0], tc.dims[1]//4, tc.dims[2], tc.dtype_in, tc.dtype_out, tc_opt=2, ensure_triggered=False)
+      if not AMX: # AMX tc.dims[2] == 1
+        helper_tc_ensure_uops_and_opts_count(tc.dims[0], tc.dims[1], tc.dims[2]//8, tc.dtype_in, tc.dtype_out, tc_opt=2, ensure_triggered=False)
+
+  @unittest.skipIf(Device.DEFAULT == "PYTHON", "not generated on EMULATED device")
+  @unittest.skipIf(CI and Device.DEFAULT in {"AMD"}, "AMD CI is really slow here")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
+  def test_tensor_cores_multi_reduce(self):
+    for tc in Device[Device.DEFAULT].renderer.tensor_cores:
+      if not is_dtype_supported(tc.dtype_in) or not is_dtype_supported(tc.dtype_out): continue
+      if tc.dtype_in is dtypes.bfloat16: continue # <-- broken with numpy
+      # this will be a M=G16, N=G32, M=G16, M=G16, K=R16, K=R16, K=R16 with 9 choices of TC MNK axes
+      golden_result = None
+      for axis in range(9):
+        a = Tensor.rand(16, 16, 29, 29, dtype=tc.dtype_in).realize()
+        b = Tensor.rand(32, 16, 16, 16, dtype=tc.dtype_in).realize()
+        c = a.conv2d(b, padding=1, dtype=tc.dtype_out)
+        realized_ast, real_bufs = helper_realized_ast(c)
+
+        program = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.TC, axis, (-1, 2, 1))])
+        assert len([uop for uop in program.uops if uop.op is Ops.WMMA]) > 0, "tensor core not triggered"
+        assert len([x for x in program.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
+
+        prg = CompiledRunner(program)
+        # TODO: support this even if numpy doesn't
+        if _to_np_dtype(real_bufs[0].dtype) is None: continue
+        real_bufs[0].copyin(np.zeros((real_bufs[0].size, ), dtype=_to_np_dtype(real_bufs[0].dtype)).data) # Zero to check that all values are filled
+        prg.exec(real_bufs)
+        result = np.frombuffer(real_bufs[0].as_buffer(), _to_np_dtype(real_bufs[0].dtype))
+
+        # ensure the results for each choice of axis matches
+        if golden_result is None: golden_result = np.frombuffer(real_bufs[0].as_buffer(), _to_np_dtype(real_bufs[0].dtype))
+        np.testing.assert_allclose(result, golden_result, atol=0.1, rtol=0.2)
+
+  @unittest.skipIf(Device.DEFAULT == "PYTHON", "slow on EMULATED device")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
+  def test_tensor_cores_unroll_phi(self):
+    tc = Device[Device.DEFAULT].renderer.tensor_cores[0]
+    x, y = Tensor.rand(128, 128, dtype=tc.dtype_in), Tensor.rand(128, 128, dtype=tc.dtype_in)
+    r = x.matmul(y, dtype=tc.dtype_out)
+    opts = [Opt(OptOps.UNROLL, 0, 4)]
+    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
+    for u in get_program(ast, opts=opts).uops:
+      if u.op is Ops.WMMA:
+        assert u.src[-1].src[0].op != Ops.STORE
+
+  @unittest.skipIf(Device.DEFAULT == "PYTHON", "slow on EMULATED device")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
+  @unittest.skipIf(Device.DEFAULT in {"CPU", "LLVM"}, "CPU does not support using a different type for accumulation")
+  def test_tensor_cores_unroll_casted_phi(self):
+    tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out][0]
+    x, y = Tensor.rand(128, 128, dtype=tc.dtype_in), Tensor.rand(128, 128, dtype=tc.dtype_in)
+    r = x.matmul(y, dtype=tc.dtype_out)
+    opts = [Opt(OptOps.UNROLL, 0, 4)]
+    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
+    for u in get_program(ast, opts=opts).uops:
+      if u.op is Ops.WMMA:
+        #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
+        assert u.src[-1].src[0].op != Ops.STORE
+
+  @unittest.skipIf(Device.DEFAULT == "PYTHON", "slow on EMULATED device")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
+  @unittest.skipIf(Device.DEFAULT in {"CPU", "LLVM"}, "CPU does not support using a different type for accumulation")
+  def test_tensor_cores_unroll_casted_phi_with_children(self):
+    # all STORE children are outside the loop
+    tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out][0]
+    x, y = Tensor.rand(128, 128, dtype=tc.dtype_in), Tensor.rand(128, 128, dtype=tc.dtype_in)
+    r = x.matmul(y, dtype=tc.dtype_out).relu()
+    opts = [Opt(OptOps.UNROLL, 0, 4)]
+    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
+    for u in get_program(ast, opts=opts).uops:
+      if u.op is Ops.WMMA:
+        #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
+        assert u.src[-1].src[0].op != Ops.STORE
+
+if __name__ == '__main__':
+  unittest.main()

test/speed/external_test_copy_speed.py

@@ -1,9 +1,9 @@
 import unittest, numpy as np
 from tinygrad import Tensor, Device, TinyJit
-from tinygrad.helpers import Timing, CI, OSX
+from tinygrad.helpers import Timing, CI, OSX, getenv
 import multiprocessing.shared_memory as shared_memory
 
-N = 256
+N = getenv("NSZ", 256)
 class TestCopySpeed(unittest.TestCase):
   @classmethod
   def setUpClass(cls): Device[Device.DEFAULT].synchronize()
@@ -54,22 +54,24 @@ class TestCopySpeed(unittest.TestCase):
     @TinyJit
     def _do_copy(t): return t.to('CPU').realize()
 
-    t = Tensor.randn(N, N, 4).contiguous().realize()
+    t = Tensor.randn(N, N).contiguous().realize()
+    Device[Device.DEFAULT].synchronize()
     for _ in range(5):
-      with Timing("sync:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
+      with Timing(f"copy {Device.DEFAULT} -> CPU {t.nbytes()/(1024**2)}M:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
         x = _do_copy(t)
         Device[Device.DEFAULT].synchronize()
       np.testing.assert_equal(t.numpy(), x.numpy())
 
-  def testCopytoCPUtoDefaultJit(self):
+  def testCopyCPUtoDefaultJit(self):
     if Device.DEFAULT == "CPU": return unittest.skip("CPU to CPU copy is a no-op")
 
     @TinyJit
-    def _do_copy(x): return t.to(Device.DEFAULT).realize()
+    def _do_copy(x): return x.to(Device.DEFAULT).realize()
 
     for _ in range(5):
-      t = Tensor.randn(N, N, 4, device="CPU").contiguous().realize()
-      with Timing("sync:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
+      t = Tensor.randn(N, N, device="CPU").contiguous().realize()
+      Device["CPU"].synchronize()
+      with Timing(f"copy CPU -> {Device.DEFAULT} {t.nbytes()/(1024**2)}M:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
         x = _do_copy(t)
         Device[Device.DEFAULT].synchronize()
       np.testing.assert_equal(t.numpy(), x.numpy())

test/test_arange.py

@@ -1,13 +1,11 @@
-import unittest, contextlib
+import unittest
 import numpy as np
 from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
 from tinygrad.helpers import CI, Context, getenv
 from tinygrad.engine.realize import run_schedule
-from tinygrad.codegen.opt.kernel import Opt, OptOps, Kernel, KernelOptError
+from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
-from tinygrad.codegen.opt.search import get_kernel_actions
 from tinygrad.uop.ops import Ops
-from tinygrad.codegen import apply_rewrites, rewrites_for_views
 
 class TestArange(unittest.TestCase):
   def _get_flops(self, N, opts=None):
@@ -49,28 +47,6 @@ class TestArange(unittest.TestCase):
     @unittest.skip("doesn't work yet")
     def test_complexity_w_local_and_padto(self): return self.test_complexity([Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.PADTO, axis=1, arg=32)])
 
-  def test_all_opts(self, opts=None, exclude=None):
-    k = Kernel(apply_rewrites(Tensor.arange(256).schedule()[-1].ast, rewrites_for_views))
-    if opts is not None:
-      for o in opts: k.apply_opt(o)
-    all_opts_256 = [kk.applied_opts for kk in get_kernel_actions(k, include_0=False).values()]
-    k = Kernel(apply_rewrites(Tensor.arange(2560).schedule()[-1].ast, rewrites_for_views))
-    if opts is not None:
-      for o in opts: k.apply_opt(o)
-    all_opts_2560 = [kk.applied_opts for kk in get_kernel_actions(k, include_0=False).values()]
-    all_opts = [x for x in all_opts_256 if x in all_opts_2560]
-    for opts in all_opts:
-      if exclude is not None and opts[-1] in exclude: continue
-      print(opts)
-      self.test_complexity(opts)
-  def test_all_opts_w_local(self):
-    with contextlib.suppress(KernelOptError):
-      return self.test_all_opts([Opt(OptOps.LOCAL, 0, 16)], [Opt(op=OptOps.PADTO, axis=1, arg=32)])
-  def test_all_opts_w_upcast(self): return self.test_all_opts([Opt(OptOps.UPCAST, 0, 4)])
-  def test_all_opts_w_unroll(self): return self.test_all_opts([Opt(OptOps.UNROLL, 0, 4)], [Opt(op=OptOps.GROUP, axis=0, arg=0)])
-  def test_all_opts_w_upcast_and_unroll(self):
-    return self.test_all_opts([Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 0, 4)], [Opt(op=OptOps.GROUP, axis=0, arg=0)])
-
 class TestRand(unittest.TestCase):
   def test_fused_rand_less_ops(self, noopt=1):
     GlobalCounters.reset()
@@ -102,7 +78,6 @@ class TestIndexing(unittest.TestCase):
       run_schedule(sched)
     self.assertEqual(out.item(), 1337)
 
-  @unittest.skipIf(getenv("PTX"), "broken on ptx for some reason")
   def test_manual_index(self):
     dataset = Tensor.rand(DSET, DDIM).realize()
     idxs = Tensor([0,3,5,6]).realize()
@@ -172,7 +147,6 @@ class TestIndexing(unittest.TestCase):
       X = dataset[idxs]
       np.testing.assert_equal(X.numpy(), 0)
 
-  @unittest.skipIf(getenv("PTX"), "broken on ptx for some reason")
   def test_index_mnist(self, noopt=1, op_limit=512*784*13, split_reduceop=0):
     # WEBGPU generates more ops due to bitpacking of < 4-byte dtypes
     if Device.DEFAULT == "WEBGPU": op_limit *= 15
@@ -191,7 +165,6 @@ class TestIndexing(unittest.TestCase):
   def test_index_mnist_split(self): self.test_index_mnist(1, split_reduceop=1)
   def test_index_mnist_opt_split(self): self.test_index_mnist(0, split_reduceop=1)
 
-  @unittest.skipIf(getenv("PTX"), "broken on ptx for some reason")
   def test_llama_embedding(self, noopt=1, op_limit=65536):
     # llama3 is 128256
     vocab_size, embed_size = (10, 3) if CI else (32000, 4096)

test/test_dtype.py

@@ -4,9 +4,8 @@ import torch
 from typing import Any, List
 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
+from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8, _to_np_dtype, _to_torch_dtype
 from tinygrad import Device, Tensor, dtypes
-from tinygrad.tensor import _to_np_dtype
 from hypothesis import assume, given, settings, strategies as strat
 from test.helpers import rand_for_dtype
 from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX
@@ -24,6 +23,10 @@ def get_available_cast_dtypes(dtype: DType) -> List[DType]:
   # dont cast internal dtypes
   return [v for k, v in DTYPES_DICT.items() if v != dtype and is_dtype_supported(v) and not k.startswith("_")]
 
+def _to_torch_storage_type(dtype:DType):
+  if dtype == dtypes.bfloat16: return torch.float32
+  return _to_torch_dtype(dtype)
+
 def _test_to_np(a:Tensor, np_dtype, target):
   if DEBUG >= 2: print(a)
   na = a.numpy()
@@ -46,10 +49,10 @@ def _test_cast(a:Tensor, target_dtype:DType):
 
   _test_op(lambda: a.cast(target_dtype), target_dtype, list(a.numpy().astype(_to_np_dtype(target_dtype))))
 def _test_bitcast(a:Tensor, target_dtype:DType, target=None):
-  if target_dtype == dtypes.bfloat16: raise unittest.SkipTest("no test for bf16 bitcast yet")
   if getenv("PTX") and a.dtype == dtypes.int8 and target_dtype.itemsize != a.dtype.itemsize:
     raise unittest.SkipTest("shape changing bitcast of int8 broken on PTX")
-  _test_op(lambda: a.bitcast(target_dtype), target_dtype, target or a.numpy().view(_to_np_dtype(target_dtype)).tolist())
+  expected = torch.tensor(a.tolist(), dtype=_to_torch_storage_type(a.dtype)).view(_to_torch_dtype(target_dtype))
+  _test_op(lambda: a.bitcast(target_dtype), target_dtype, target or expected.tolist())
 
 class TestDType(unittest.TestCase):
   DTYPE: Any = None
@@ -126,7 +129,7 @@ class TestDType(unittest.TestCase):
 
   def test_finfo(self):
     if self.DTYPE not in [dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]: return
-    info = np.finfo(_to_np_dtype(self.DTYPE))
+    info = ml_dtypes.finfo(ml_dtypes.bfloat16 if self.DTYPE is dtypes.bfloat16 else _to_np_dtype(self.DTYPE))
     assert info.bits == self.DTYPE.itemsize*8
     assert info.nexp == dtypes.finfo(self.DTYPE)[0]
     assert info.nmant == dtypes.finfo(self.DTYPE)[1]
@@ -294,15 +297,14 @@ class TestUint8DType(TestDType):
   def test_uint8_to_int8_overflow(self):
     _test_op(lambda: Tensor([255, 254, 253, 252], dtype=dtypes.uint8).cast(dtypes.int8), dtypes.int8, [-1, -2, -3, -4])
 
-@unittest.skipIf(Device.DEFAULT == "WEBGL", "No bitcast on WebGL")
 class TestBitCast(unittest.TestCase):
   @given(strat.sampled_from(dtype_ints + dtype_floats), strat.sampled_from(dtype_ints + dtype_floats))
   def test_shape_change_bitcast(self, dt1, dt2):
     # NOTE: this has to be assume to prevent hypothesis from skipping all samples
-    assume(dt2 != dtypes.bfloat16 and dt1 != dtypes.bfloat16) # no test for bf16 bitcast yet
     assume(not (getenv("PTX") and dt1 == dtypes.int8)) # TODO: bitcasting int8 fails in PTX
     data = rand_for_dtype(dt1, 32).reshape(2, 2, 8)
-    _test_op(lambda: Tensor(data, dtype=dt1).bitcast(dt2), dt2, data.view(_to_np_dtype(dt2)).tolist())
+    expected = torch.tensor(data.tolist(), dtype=_to_torch_storage_type(dt1)).view(_to_torch_dtype(dt2))
+    _test_op(lambda: Tensor(data, dtype=dt1).bitcast(dt2), dt2, expected.tolist())
 
   def test_shape_change_bitcast_exceptions(self):
     with self.assertRaises(RuntimeError):
@@ -342,6 +344,9 @@ class TestUint64DType(TestDType):
 
 class TestBoolDType(TestDType): DTYPE = dtypes.bool
 
+@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
+class TestBFloat16Type(TestDType): DTYPE = dtypes.bfloat16
+
 class TestPtrDType(unittest.TestCase):
   def test_vec_double(self):
     dt1 = dtypes.float.vec(4).ptr().vec(4)
@@ -414,7 +419,7 @@ class TestDtypeUsage(unittest.TestCase):
         t = Tensor([[1, 2], [3, 4]], dtype=d)
         (t*t).max().item()
 
-@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16) or Device.DEFAULT == "PYTHON", f"no bfloat16 on {Device.DEFAULT}")
+@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
 class TestOpsBFloat16(unittest.TestCase):
   def test_cast(self):
     # TODO: helper_test_op breaks in unrelated part
@@ -422,5 +427,10 @@ class TestOpsBFloat16(unittest.TestCase):
     data = [60000.0, 70000.0, 80000.0]
     np.testing.assert_allclose(Tensor(data).cast("bfloat16").numpy(), torch.tensor(data).type(torch.bfloat16).float().numpy())
 
+  def test_no_approximation(self):
+    data = [326.0, 339.0, 10603200512.0]
+    expected = torch.tensor(data, dtype=torch.bfloat16).sqrt().float().numpy()
+    np.testing.assert_allclose(Tensor(data, dtype=dtypes.bfloat16).sqrt().numpy(), expected)
+
 if __name__ == '__main__':
   unittest.main()

test/test_dtype_alu.py

@@ -4,6 +4,7 @@ from tinygrad.dtype import DType
 from tinygrad.helpers import CI, getenv
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
+from tinygrad.runtime.ops_python import from_storage_scalar
 import numpy as np
 import pytest
 from hypothesis import given, strategies as strat, settings, HealthCheck
@@ -19,29 +20,23 @@ dtypes_int = (dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64, dtypes.uint
 dtypes_bool = (dtypes.bool,)
 binary_operations = [operator.add, operator.sub, operator.mul, operator.lt, operator.eq]
 
-# TODO: LLVM comparing with nan is incorrect
-if Device.DEFAULT == "LLVM" or getenv("AMD_LLVM", 0):
-  binary_operations.remove(operator.lt)
-
 integer_binary_operations = binary_operations + [(Tensor.bitwise_xor, np.bitwise_xor), (Tensor.bitwise_and, np.bitwise_and),
-                                                 (Tensor.bitwise_or, np.bitwise_or), operator.mod]
+                                                 (Tensor.bitwise_or, np.bitwise_or), (Tensor.maximum, np.maximum), operator.mod]
 unary_operations = [(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.sin),
-                    (Tensor.sqrt, np.sqrt), (Tensor.reciprocal, np.reciprocal)]
+                    (Tensor.sqrt, np.sqrt), (Tensor.reciprocal, np.reciprocal), (Tensor.cos, np.cos)]
 
 # TODO: enable this (this is a dtype issue)
 #binary_operations.append(operator.truediv)
 
-# TODO: (a+b)/2 in tensor.py's maximum can overflow. This requires a new implementation of maximum that can be backpropagated
-#binary_operations += [(Tensor.maximum, np.maximum)]
-
 # 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": unary_operations.remove((Tensor.sin, np.sin))
+if (getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}) or Device.DEFAULT == "WEBGPU":
+  unary_operations.remove((Tensor.sin, np.sin))
+  unary_operations.remove((Tensor.cos, np.cos))
 
 class ht:
   float64 = strat.floats(width=64, allow_subnormal=False)
   float32 = strat.floats(width=32, allow_subnormal=False)
   float16 = strat.floats(width=16, allow_subnormal=False)
-  bfloat16 = strat.floats(width=16, allow_subnormal=False)
   uint8 = strat.integers(0, 255)
   uint16 = strat.integers(0, 65535)
   uint32 = strat.integers(0, 2**32-1)
@@ -51,12 +46,13 @@ class ht:
   int32 = strat.integers(-2147483648, 2147483647)
   int64 = strat.integers(-9223372036854775808, 9223372036854775807)
   bool = strat.booleans()
+ht.bfloat16 = ht.uint16
 
 def universal_test(a, b, dtype, op):
-  # The 'nan' cases only fail with Vulkan WebGPU backend (CI)
-  if (math.isnan(a) or math.isnan(b)) and Device.DEFAULT == "WEBGPU" and CI: return
   if not isinstance(op, tuple): op = (op, op)
   if op[0] == operator.mod and b == 0: return
+  # lt and max with nan is undefined in tinygrad
+  if op[0] in (operator.lt, Tensor.maximum) and (math.isnan(a) or math.isnan(b)): return
   ta, tb = Tensor([a], dtype=dtype), Tensor([b], dtype=dtype)
   tensor_value = (op[0](ta, tb)).numpy()
   numpy_value = op[1](ta.numpy(), tb.numpy())
@@ -68,11 +64,14 @@ def universal_test(a, b, dtype, op):
 def universal_test_unary(a, dtype, op):
   if not isinstance(op, tuple): op = (op, op)
   ta = Tensor([a], dtype=dtype)
+  # TODO: cos does not match for large input
+  if op[0] == Tensor.cos and abs(a) > 30: return
+  if op[0] == Tensor.log and a <= 0: return
   out: Tensor = op[0](ta)
   tensor_value = out.numpy()
   numpy_value = op[1](ta.numpy())
   if dtype in dtypes.floats:
-    atol, rtol = {dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 1e-2)}.get(dtype, (1e-6, 1e-5))
+    atol, rtol = {dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 2e-2)}.get(dtype, (1e-6, 1e-5))
     np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
   else: np.testing.assert_equal(tensor_value, numpy_value)
 
@@ -85,6 +84,9 @@ def universal_test_cast(a, in_dtype, dtype):
 def universal_test_midcast(a, b, c, op1, op2, d1:DType, d2:DType):
   if not isinstance(op1, tuple): op1 = (op1, op1)
   if not isinstance(op2, tuple): op2 = (op2, op2)
+  # lt and max with nan is undefined in tinygrad
+  if op1[0] in (operator.lt, Tensor.maximum) and (math.isnan(a) or math.isnan(b)): return
+  if op2[0] in (operator.lt, Tensor.maximum) and math.isnan(c): return
   at, bt, ct = Tensor([a], dtype=d1), Tensor([b], dtype=d1), Tensor([c], dtype=d2)
   an, bn, cn = np.array([a]).astype(_to_np_dtype(d1)), np.array([b]).astype(_to_np_dtype(d1)), np.array([c]).astype(_to_np_dtype(d2))
   tensor_value = op2[0](op1[0](at, bt).cast(d2), ct).numpy()
@@ -105,7 +107,8 @@ class TestDTypeALU(unittest.TestCase):
 
   @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
   @given(ht.bfloat16, ht.bfloat16, strat.sampled_from(binary_operations))
-  def test_bfloat16(self, a, b, op): universal_test(a, b, dtypes.bfloat16, op)
+  def test_bfloat16(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)
 
   @given(ht.float32, strat.sampled_from(unary_operations))
   def test_float32_unary(self, a, op): universal_test_unary(a, dtypes.float32, op)
@@ -116,7 +119,7 @@ class TestDTypeALU(unittest.TestCase):
 
   @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
   @given(ht.bfloat16, strat.sampled_from(unary_operations))
-  def test_bfloat16_unary(self, a, op): universal_test_unary(a, dtypes.bfloat16, op)
+  def test_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)
 
   @given(ht.uint8, ht.uint8, strat.sampled_from(integer_binary_operations))
   def test_uint8(self, a, b, op): universal_test(a, b, dtypes.uint8, op)

test/test_linearizer_dumb.py

@@ -5,193 +5,110 @@
 import unittest
 from tinygrad import Device, dtypes
 from tinygrad.device import is_dtype_supported
-from tinygrad.uop.ops import UOp, Ops
+from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
 from tinygrad.helpers import getenv
 from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad.codegen.opt.search import Opt, OptOps
-from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.engine.realize import get_program
 
-class TestLinearizerDumb(unittest.TestCase):
+class TestLinearizerFailure(unittest.TestCase):
+  @unittest.expectedFailure
   @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
-  def test_unmerged_ifs(self):
-    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.VIEW, dtypes.half.ptr(1605632), arg=ShapeTracker(views=(View(shape=(64, 1, 512, 7, 7, 1, 1, 1), strides=(25088, 0, 49, 7, 1, 0, 0, 0), offset=0, mask=None, contiguous=True),)), src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(1605632), arg=0, src=()),)),
-        UOp(Ops.MAX, dtypes.half, arg=None, src=(
-          UOp(Ops.MUL, dtypes.half, arg=None, src=(
-            UOp(Ops.CAST, dtypes.half, arg=None, src=(
-              UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (5, 6, 7)), 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(1605632), arg=ShapeTracker(views=(View(shape=(1, 64, 1, 512, 4, 9, 4, 9), strides=(0, 25088, 0, 49, 0, 7, 0, 1), offset=-8, mask=((0, 1), (0, 64), (0, 1), (0, 512), (0, 4), (1, 8), (0, 4), (1, 8)), contiguous=False), View(shape=(64, 1, 512, 7, 7, 512, 3, 3), strides=(663552, 0, 0, 36, 1, 1296, 360, 10), offset=0, mask=None, contiguous=False))), src=(
-                        UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(1605632), arg=1, src=()),)),)),
-                    UOp(Ops.LOAD, dtypes.half, arg=None, src=(
-                      UOp(Ops.VIEW, dtypes.half.ptr(2359296), arg=ShapeTracker(views=(View(shape=(64, 1, 512, 7, 7, 512, 3, 3), strides=(0, 0, 4608, 0, 0, 9, 3, 1), offset=0, mask=None, contiguous=False),)), src=(
-                        UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(2359296), arg=2, src=()),)),)),)),)),)),)),
-            UOp(Ops.CONST, dtypes.half, arg=0.9999950000374996, src=(
-              x16:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(64, 1, 512, 7, 7, 1, 1, 1), strides=(0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
-          UOp(Ops.CONST, dtypes.half, arg=0.0, src=(
-             x16,)),)),)),))
-    opts = [Opt(op=OptOps.TC, axis=2, arg=(-1, 2, 1)), Opt(op=OptOps.UPCAST, axis=2, arg=0), Opt(op=OptOps.UNROLL, axis=1, arg=0)]
-    k = Kernel(ast, opts=Device["METAL"].renderer)
-    k.apply_opts(opts)
-    prg = get_program(k.get_optimized_ast(), k.opts)
-    print(prg.src)
-    Device[Device.DEFAULT].compiler.compile_cached(prg.src)
-    gate_count = len([x for x in prg.src.splitlines() if "if" in x])
-    assert gate_count == 1, f"must have only one gate {gate_count} != 1"
-    assert len([u for u in prg.uops if u.op is Ops.IF]) == 1, "must have a single IF"
+  def test_failure_beam_mnist(self):
+    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(4014080), arg=0, src=())
+    c1 = UOp.range(UOp.const(dtypes.int, 512), 0, AxisType.GLOBAL)
+    c2 = UOp.range(UOp.const(dtypes.int, 784), 1, AxisType.GLOBAL)
+    c3 = UOp.range(UOp.const(dtypes.int, 10), 3, AxisType.GLOBAL)
+    c4 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(512), arg=1, src=())
+    c5 = c4.index(c1, UOp.const(dtypes.bool, True)).load()
+    c6 = UOp.range(UOp.const(dtypes.int, 6000), 1004, AxisType.REDUCE)
+    c7 = UOp.range(UOp.const(dtypes.int, 3750), 2006, AxisType.REDUCE)
+    c8 = UOp.range(UOp.const(dtypes.int, 16), 2007, AxisType.GROUP_REDUCE)
+    c9 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(47040000), arg=2, src=())
+    c10 = c9.index((((c3*UOp.const(dtypes.int, 4704000))+c2)+(c6*UOp.const(dtypes.int, 784))), UOp.const(dtypes.bool, True)).load()
+    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.int, 6000))+c6)+((c7*UOp.const(dtypes.int, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.int, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
+    c12 = c0.index((((c1*UOp.const(dtypes.int, 7840))+(c2*UOp.const(dtypes.int, 10)))+c3), UOp.const(dtypes.bool, True)).store(c11, c1, c2, c3)
+    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")
   def test_max_simplify_and_cancel(self):
-    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.VIEW, dtypes.int.ptr(1000), arg=ShapeTracker(views=(View(shape=(1000, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),)), src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1000), arg=0, src=()),)),
-        UOp(Ops.MUL, dtypes.int, arg=None, src=(
-          UOp(Ops.CAST, dtypes.int, arg=None, src=(
-            UOp(Ops.CMPNE, dtypes.bool, arg=None, src=(
-              UOp(Ops.CMPNE, dtypes.bool, arg=None, src=(
-                UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                  UOp(Ops.VIEW, dtypes.float.ptr(1000), arg=ShapeTracker(views=(View(shape=(1000, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),)), src=(
-                    UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1000), arg=1, src=()),)),)),
-                UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                  UOp(Ops.VIEW, dtypes.float.ptr(1), arg=ShapeTracker(views=(View(shape=(1000, 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=(
-                    UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1), arg=2, src=()),)),)),)),
-              UOp(Ops.CONST, dtypes.bool, arg=True, src=(
-                x14:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1000, 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),
-          UOp(Ops.ADD, dtypes.int, arg=None, src=(
-            UOp(Ops.REDUCE_AXIS, dtypes.int, arg=(Ops.ADD, (1,)), src=(
-              UOp(Ops.WHERE, dtypes.int, arg=None, src=(
-                UOp(Ops.VALID, dtypes.bool, arg=None, src=(
-                  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=()),)),
-                UOp(Ops.CONST, dtypes.int, arg=-1, src=(
-                  x21:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1000, 1000), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),
-                UOp(Ops.CONST, dtypes.int, arg=0, src=(
-                   x21,)),)),)),
-            UOp(Ops.CONST, dtypes.int, arg=1000, src=(
-               x14,)),)),)),)),))
+    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)]
-    k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
-    k.apply_opts(opts)
-    prg = get_program(k.get_optimized_ast(), k.opts)
+    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"
 
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "need local")
-  @unittest.skip("not applicable")
-  def test_expander_new_srcs(self):
-    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.VIEW, dtypes.float.ptr(25), arg=ShapeTracker(views=(View(shape=(25, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),)), src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(25), arg=0, src=()),)),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (1,)), src=(
-          UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-            UOp(Ops.VIEW, dtypes.float.ptr(25), arg=ShapeTracker(views=(View(shape=(26, 49), strides=(0, -1), offset=48, mask=((0, 26), (24, 49)), contiguous=False), View(shape=(25, 25), strides=(1, 50), offset=0, mask=None, contiguous=False))), src=(
-              UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(25), arg=1, src=()),)),)),)),)),))
-    opts = [Opt(op=OptOps.GROUP, axis=0, arg=0), Opt(op=OptOps.PADTO, axis=0, arg=32), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=0)]
-    k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
-    k.apply_opts(opts)
-    prg = get_program(k.get_optimized_ast(), k.opts)
-    print(prg.src)
-    if_uops = [u for u in prg.uops if u.op is Ops.IF]
-    self.assertIn(len(if_uops), {1,2,3})
-    conditions = if_uops[0].src[0].toposort()
-    self.assertLessEqual(len(conditions), 9)
-
   # 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}")
   def test_llama_embedding(self):
-    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.VIEW, dtypes.half.ptr(4096), arg=ShapeTracker(views=(View(shape=(4096, 1, 1), strides=(1, 0, 0), offset=0, mask=None, contiguous=True),)), src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(4096), arg=0, src=()),)),
-        UOp(Ops.CAST, dtypes.half, arg=None, src=(
-          UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (1,)), src=(
-            UOp(Ops.CAST, dtypes.float, arg=None, src=(
-              UOp(Ops.MUL, dtypes.half, arg=None, src=(
-                UOp(Ops.CAST, dtypes.half, arg=None, src=(
-                  UOp(Ops.CMPNE, dtypes.bool, arg=None, src=(
-                    UOp(Ops.CMPNE, dtypes.bool, arg=None, src=(
-                      UOp(Ops.ADD, dtypes.int, arg=None, src=(
-                        UOp(Ops.REDUCE_AXIS, dtypes.int, arg=(Ops.ADD, (2,)), src=(
-                          UOp(Ops.WHERE, dtypes.int, arg=None, src=(
-                            UOp(Ops.VALID, dtypes.bool, arg=None, src=(
-                              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=()),)),
-                            UOp(Ops.CONST, dtypes.int, arg=1, src=(
-                              x16:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4096, 32000, 32000), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),
-                            UOp(Ops.CONST, dtypes.int, arg=0, src=(
-                               x16,)),)),)),
-                        UOp(Ops.CONST, dtypes.int, arg=-1, src=(
-                          x19:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4096, 32000, 1), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
-                      UOp(Ops.LOAD, dtypes.int, arg=None, src=(
-                        UOp(Ops.VIEW, dtypes.int.ptr(1), arg=ShapeTracker(views=(View(shape=(4096, 32000, 1), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)), src=(
-                          UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=1, src=()),)),)),)),
-                    UOp(Ops.CONST, dtypes.bool, arg=True, src=(
-                       x19,)),)),)),
-                UOp(Ops.LOAD, dtypes.half, arg=None, src=(
-                  UOp(Ops.VIEW, dtypes.half.ptr(131072000), arg=ShapeTracker(views=(View(shape=(4096, 32000, 1), strides=(1, 4096, 0), offset=0, mask=None, contiguous=False),)), src=(
-                    UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(131072000), arg=2, src=()),)),)),)),)),)),)),)),))
-    k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
-    prg = get_program(k.get_optimized_ast(), k.opts)
+    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):
-    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.VIEW, dtypes.float.ptr(1), arg=ShapeTracker(views=(View(shape=(1, 1), strides=(0, 0), offset=0, mask=None, contiguous=True),)), src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1), arg=0, src=()),)),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (0, 1)), src=(
-          UOp(Ops.WHERE, dtypes.float, arg=None, src=(
-            UOp(Ops.CMPNE, dtypes.bool, arg=None, src=(
-              UOp(Ops.CMPNE, dtypes.bool, arg=None, src=(
-                UOp(Ops.LOAD, dtypes.long, arg=None, src=(
-                  UOp(Ops.VIEW, dtypes.long.ptr(18), arg=ShapeTracker(views=(View(shape=(3, 6), strides=(6, 1), offset=0, mask=None, contiguous=True),)), src=(
-                    UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(18), arg=1, src=()),)),)),
-                UOp(Ops.CONST, dtypes.long, arg=-1, src=(
-                  x11:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(3, 6), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
-              UOp(Ops.CONST, dtypes.bool, arg=True, src=(
-                 x11,)),)),
-            UOp(Ops.CONST, dtypes.float, arg=0.0, src=(
-               x11,)),
-            UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-              UOp(Ops.VIEW, dtypes.float.ptr(18), arg=ShapeTracker(views=(View(shape=(3, 6), strides=(6, 1), offset=0, mask=None, contiguous=True),)), src=(
-                UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(18), arg=2, src=()),)),)),)),)),)),))
+    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)]
-    k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
-    k.apply_opts(opts)
-    prg = get_program(k.get_optimized_ast(), k.opts)
+    prg = get_program(ast, Device[Device.DEFAULT].renderer, opts)
     print(prg.src)
-    load_idxs = [x.src[1] for x in k.uops if x.op is Ops.LOAD and x.src[0].arg == 2]
+    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(getenv("PTX"), "this is somehow correct in PTX")
   def test_upcasted_stores_out_of_order(self):
-    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.VIEW, dtypes.float.ptr(9360), arg=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),)), src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(9360), arg=0, src=()),)),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (6,)), src=(
-          UOp(Ops.MUL, dtypes.float, arg=None, src=(
-            UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-              UOp(Ops.VIEW, dtypes.float.ptr(144), arg=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),)), src=(
-                UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(144), arg=1, src=()),)),)),
-            UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-              UOp(Ops.VIEW, dtypes.float.ptr(1040), arg=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),)), src=(
-                UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1040), arg=2, src=()),)),)),)),)),)),))
+    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)]
-    k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
-    k.apply_opts(opts)
-    prg = get_program(k.get_optimized_ast(), k.opts)
+    prg = get_program(ast, Device[Device.DEFAULT].renderer, opts)
     print(prg.src)
-    store_idxs = [x.src[1] for x in k.uops if x.op is Ops.STORE]
+    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

test/test_linearizer_overflows.py

@@ -1,165 +0,0 @@
-# ruff: noqa: E501
-import unittest
-from tinygrad import dtypes
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt.search import Opt, OptOps, bufs_from_lin
-from extra.optimization.helpers import time_linearizer
-
-# stuff needed to unpack a kernel
-from tinygrad.uop.ops import UOp, Ops
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
-
-def _test_overflow(ast, opts):
-  lin = Kernel(ast)
-  lin.apply_opts(opts)
-  bufs = bufs_from_lin(lin)
-  print(bufs)
-  time_linearizer(lin, bufs)
-
-# NOTE: if you want these to trigger, set launch bounds on HIP kernels
-@unittest.skip("unneeded without launch bounds")
-class TestLinearizerOverflow(unittest.TestCase):
-  def test_overflow_1(self):
-    ast = UOp(Ops.SINK, None, arg=None, src=(
-      UOp(Ops.STORE, None, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(51380224), arg=0, src=()),
-        UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(64, 1, 64, 112, 112, 1, 1, 1), strides=(802816, 0, 12544, 112, 1, 0, 0, 0), offset=0, mask=None, contiguous=True),)), src=()),
-        UOp(Ops.MAX, dtypes.float, arg=None, src=(
-          UOp(Ops.ADD, dtypes.float, arg=None, src=(
-            UOp(Ops.MUL, dtypes.float, arg=None, src=(
-              UOp(Ops.MUL, dtypes.float, arg=None, src=(
-                UOp(Ops.ADD, dtypes.float, arg=None, src=(
-                  UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (7, 6, 5)), src=(
-                    UOp(Ops.MUL, dtypes.float, arg=None, src=(
-                      UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                        UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(9633792), arg=1, src=()),
-                        UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(1, 64, 1, 3, 8, 230, 8, 230), strides=(0, 150528, 0, 50176, 0, 224, 0, 1), offset=-675, mask=((0, 1), (0, 64), (0, 1), (0, 3), (0, 8), (3, 227), (0, 8), (3, 227)), contiguous=False), View(shape=(64, 1, 64, 112, 112, 3, 7, 7), strides=(10156800, 0, 0, 3680, 2, 3385600, 425040, 231), offset=0, mask=None, contiguous=False))), src=()),)),
-                      UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                        UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(9408), arg=2, src=()),
-                        UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(64, 1, 64, 112, 112, 3, 7, 7), strides=(0, 0, 147, 0, 0, 49, 7, 1), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),
-                  x16:=UOp(Ops.CONST, dtypes.float, arg=0.0, src=(
-                    x17:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(64, 1, 64, 112, 112, 1, 1, 1), strides=(0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
-                UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                  UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64), arg=3, src=()),
-                  x20:=UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(64, 1, 64, 112, 112, 1, 1, 1), strides=(0, 0, 1, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
-              UOp(Ops.SQRT, dtypes.float, arg=None, src=(
-                UOp(Ops.MUL, dtypes.float, arg=None, src=(
-                  x23:=UOp(Ops.CONST, dtypes.float, arg=1.0, src=(
-                     x17,)),
-                  UOp(Ops.RECIP, dtypes.float, arg=None, src=(
-                    UOp(Ops.ADD, dtypes.float, arg=None, src=(
-                       x23,
-                      UOp(Ops.CONST, dtypes.float, arg=1e-05, src=(
-                         x17,)),)),)),)),)),)),
-            UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-              UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64), arg=4, src=()),
-               x20,)),)),
-           x16,)),)),))
-    opts = [Opt(op=OptOps.LOCAL, axis=3, arg=16), Opt(op=OptOps.LOCAL, axis=2, arg=16), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=2, arg=0)]
-    _test_overflow(ast, opts)
-
-  # From BEAM on hlb_cifar.py
-  def test_overflow_2(self):
-    ast = UOp(Ops.SINK, None, arg=None, src=(
-      UOp(Ops.STORE, None, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(33554432), arg=0, src=()),
-        UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(512, 1, 64, 32, 32, 1, 1, 1), strides=(65536, 0, 1024, 32, 1, 0, 0, 0), offset=0, mask=None, contiguous=True),)), src=()),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (7, 6, 5)), src=(
-          UOp(Ops.MUL, dtypes.float, arg=None, src=(
-            UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-              UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(16777216), arg=1, src=()),
-              UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(1, 512, 1, 32, 4, 34, 4, 34), strides=(0, 32768, 0, 1024, 0, 32, 0, 1), offset=-33, mask=((0, 1), (0, 512), (0, 1), (0, 32), (0, 4), (1, 33), (0, 4), (1, 33)), contiguous=False), View(shape=(512, 1, 64, 32, 32, 32, 3, 3), strides=(591872, 0, 0, 136, 1, 18496, 4760, 35), offset=0, mask=None, contiguous=False))), src=()),)),
-            UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-              UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(18432), arg=2, src=()),
-              UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(512, 1, 64, 32, 32, 32, 3, 3), strides=(0, 0, 288, 0, 0, 9, 3, 1), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),))
-    opts = [Opt(op=OptOps.LOCAL, axis=3, arg=16), Opt(op=OptOps.LOCAL, axis=2, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=2, arg=0), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UNROLL, axis=0, arg=0)]
-    _test_overflow(ast, opts)
-
-  # from BEAM on default simple_conv.py (which is quite large):
-  def test_overflow_3(self):
-    ast = UOp(Ops.SINK, None, arg=None, src=(
-      UOp(Ops.STORE, None, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(33554432), arg=0, src=()),
-        UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(16, 1, 128, 128, 128, 1, 1, 1), strides=(2097152, 0, 16384, 128, 1, 0, 0, 0), offset=0, mask=None, contiguous=True),)), src=()),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (7, 6, 5)), src=(
-          UOp(Ops.MUL, dtypes.float, arg=None, src=(
-            UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-              UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(33554432), arg=1, src=()),
-              UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(1, 16, 1, 128, 4, 130, 4, 130), strides=(0, 2097152, 0, 16384, 0, 128, 0, 1), offset=-129, mask=((0, 1), (0, 16), (0, 1), (0, 128), (0, 4), (1, 129), (0, 4), (1, 129)), contiguous=False), View(shape=(16, 1, 128, 128, 128, 128, 3, 3), strides=(34611200, 0, 0, 520, 1, 270400, 68120, 131), offset=0, mask=None, contiguous=False))), src=()),)),
-            UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-              UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(147456), arg=2, src=()),
-              UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(16, 1, 128, 128, 128, 128, 3, 3), strides=(0, 0, 1152, 0, 0, 9, 3, 1), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),))
-    opts = [Opt(op=OptOps.LOCAL, axis=3, arg=16), Opt(op=OptOps.LOCAL, axis=2, arg=8), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=3, arg=0), Opt(op=OptOps.UPCAST, axis=1, arg=2), Opt(op=OptOps.UPCAST, axis=2, arg=2)]
-    _test_overflow(ast, opts)
-
-  # from BEAM on BS=4 simple_conv.py:
-  def test_overflow_4(self):
-    ast = UOp(Ops.SINK, None, arg=None, src=(
-      UOp(Ops.STORE, None, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(8388608), arg=0, src=()),
-        UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(4, 1, 128, 128, 128, 1, 1, 1), strides=(2097152, 0, 16384, 128, 1, 0, 0, 0), offset=0, mask=None, contiguous=True),)), src=()),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (7, 6, 5)), src=(
-          UOp(Ops.MUL, dtypes.float, arg=None, src=(
-            UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-              UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(8388608), arg=1, src=()),
-              UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(1, 4, 1, 128, 4, 130, 4, 130), strides=(0, 2097152, 0, 16384, 0, 128, 0, 1), offset=-129, mask=((0, 1), (0, 4), (0, 1), (0, 128), (0, 4), (1, 129), (0, 4), (1, 129)), contiguous=False), View(shape=(4, 1, 128, 128, 128, 128, 3, 3), strides=(34611200, 0, 0, 520, 1, 270400, 68120, 131), offset=0, mask=None, contiguous=False))), src=()),)),
-            UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-              UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(147456), arg=2, src=()),
-              UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(4, 1, 128, 128, 128, 128, 3, 3), strides=(0, 0, 1152, 0, 0, 9, 3, 1), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),))
-    opts = [Opt(op=OptOps.UPCAST, axis=3, arg=4), Opt(op=OptOps.LOCAL, axis=3, arg=16), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=2, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2), Opt(op=OptOps.UPCAST, axis=2, arg=4)]
-    _test_overflow(ast, opts)
-
-  # from BEAM on BS=2 simple_conv.py:
-  def test_overflow_5(self):
-    ast = UOp(Ops.SINK, None, arg=None, src=(
-      UOp(Ops.STORE, None, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(4194304), arg=0, src=()),
-        UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(2, 1, 128, 128, 128, 1, 1, 1), strides=(2097152, 0, 16384, 128, 1, 0, 0, 0), offset=0, mask=None, contiguous=True),)), src=()),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (7, 6, 5)), src=(
-          UOp(Ops.MUL, dtypes.float, arg=None, src=(
-            UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-              UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(4194304), arg=1, src=()),
-              UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(1, 2, 1, 128, 4, 130, 4, 130), strides=(0, 2097152, 0, 16384, 0, 128, 0, 1), offset=-129, mask=((0, 1), (0, 2), (0, 1), (0, 128), (0, 4), (1, 129), (0, 4), (1, 129)), contiguous=False), View(shape=(2, 1, 128, 128, 128, 128, 3, 3), strides=(34611200, 0, 0, 520, 1, 270400, 68120, 131), offset=0, mask=None, contiguous=False))), src=()),)),
-            UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-              UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(147456), arg=2, src=()),
-              UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(2, 1, 128, 128, 128, 128, 3, 3), strides=(0, 0, 1152, 0, 0, 9, 3, 1), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),))
-    opts = [Opt(op=OptOps.LOCAL, axis=3, arg=16), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=3, arg=0), Opt(op=OptOps.LOCAL, axis=2, arg=2), Opt(op=OptOps.UPCAST, axis=1, arg=2), Opt(op=OptOps.UPCAST, axis=2, arg=2)]
-    _test_overflow(ast, opts)
-
-  # from BEAM on BS=3 simple_conv.py:
-  def test_overflow_6(self):
-    ast = UOp(Ops.SINK, None, arg=None, src=(
-      UOp(Ops.STORE, None, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(6291456), arg=0, src=()),
-        UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(3, 1, 128, 128, 128, 1, 1, 1), strides=(2097152, 0, 16384, 128, 1, 0, 0, 0), offset=0, mask=None, contiguous=True),)), src=()),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (7, 6, 5)), src=(
-          UOp(Ops.MUL, dtypes.float, arg=None, src=(
-            UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-              UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(6291456), arg=1, src=()),
-              UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(1, 3, 1, 128, 4, 130, 4, 130), strides=(0, 2097152, 0, 16384, 0, 128, 0, 1), offset=-129, mask=((0, 1), (0, 3), (0, 1), (0, 128), (0, 4), (1, 129), (0, 4), (1, 129)), contiguous=False), View(shape=(3, 1, 128, 128, 128, 128, 3, 3), strides=(34611200, 0, 0, 520, 1, 270400, 68120, 131), offset=0, mask=None, contiguous=False))), src=()),)),
-            UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-              UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(147456), arg=2, src=()),
-              UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(3, 1, 128, 128, 128, 128, 3, 3), strides=(0, 0, 1152, 0, 0, 9, 3, 1), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),))
-    opts = [Opt(op=OptOps.LOCAL, axis=3, arg=16), Opt(op=OptOps.UPCAST, axis=3, arg=0), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=2, arg=8), Opt(op=OptOps.UPCAST, axis=1, arg=2), Opt(op=OptOps.UPCAST, axis=3, arg=2)]
-    _test_overflow(ast, opts)
-
-  # from BEAM on BS=3 simple_conv.py: (alt)
-  def test_overflow_7(self):
-    ast = UOp(Ops.SINK, None, arg=None, src=(
-      UOp(Ops.STORE, None, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(6291456), arg=0, src=()),
-        UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(3, 1, 128, 128, 128, 1, 1, 1), strides=(2097152, 0, 16384, 128, 1, 0, 0, 0), offset=0, mask=None, contiguous=True),)), src=()),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (7, 6, 5)), src=(
-          UOp(Ops.MUL, dtypes.float, arg=None, src=(
-            UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-              UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(6291456), arg=1, src=()),
-              UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(1, 3, 1, 128, 4, 130, 4, 130), strides=(0, 2097152, 0, 16384, 0, 128, 0, 1), offset=-129, mask=((0, 1), (0, 3), (0, 1), (0, 128), (0, 4), (1, 129), (0, 4), (1, 129)), contiguous=False), View(shape=(3, 1, 128, 128, 128, 128, 3, 3), strides=(34611200, 0, 0, 520, 1, 270400, 68120, 131), offset=0, mask=None, contiguous=False))), src=()),)),
-            UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-              UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(147456), arg=2, src=()),
-              UOp(Ops.VIEW, None, arg=ShapeTracker(views=(View(shape=(3, 1, 128, 128, 128, 128, 3, 3), strides=(0, 0, 1152, 0, 0, 9, 3, 1), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),))
-    opts = [Opt(op=OptOps.UPCAST, axis=3, arg=4), Opt(op=OptOps.LOCAL, axis=3, arg=16), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=2, arg=8), Opt(op=OptOps.UPCAST, axis=1, arg=2), Opt(op=OptOps.UPCAST, axis=2, arg=4)]
-    _test_overflow(ast, opts)
-
-if __name__ == '__main__':
-  unittest.main()

test/test_nn.py

@@ -2,7 +2,7 @@
 import unittest
 import numpy as np
 import torch
-from tinygrad import Tensor, Device, TinyJit
+from tinygrad import Tensor, Device, TinyJit, dtypes
 from tinygrad.uop.ops import Ops
 from tinygrad.helpers import GlobalCounters, CI, Context
 from tinygrad.nn import Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, Linear, Embedding
@@ -465,7 +465,7 @@ class TestNN(unittest.TestCase):
     # used to fail bounds check
     with Context(FUSE_ARANGE=1):
       embedding = Embedding(100, 1024)
-      input_ids = Tensor.empty(16, 16)
+      input_ids = Tensor.empty(16, 16, dtype=dtypes.int)
       embedding(input_ids).realize()
 
   def test_load_state_dict(self):

test/test_ops.py

@@ -17,6 +17,9 @@ if CI:
 FORWARD_ONLY = getenv("FORWARD_ONLY", 0)
 PRINT_TENSORS = getenv("PRINT_TENSORS", 0)
 
+def slow_test(test_func):
+  return unittest.skipIf(getenv("SKIP_SLOW_TEST"), "Skipping slow test")(test_func)
+
 def helper_test_op(shps, torch_fxn, tinygrad_fxn=None, atol=1e-6, rtol=1e-3, grad_atol=1e-4, grad_rtol=1e-3,
                    forward_only=False, vals=None, low=-2, high=2):
   if tinygrad_fxn is None: tinygrad_fxn = torch_fxn
@@ -1009,9 +1012,11 @@ class TestOps(unittest.TestCase):
 
   def test_small_cumsum(self):
     helper_test_op([(10)], lambda x: torch.cumsum(x, dim=0), lambda x: Tensor.cumsum(x, axis=0))
+  @slow_test
   def test_simple_cumsum(self):
     helper_test_op([(512)], lambda x: torch.cumsum(x, dim=0), lambda x: Tensor.cumsum(x, axis=0))
     helper_test_op([(1022)], lambda x: torch.cumsum(x, dim=0), lambda x: Tensor.cumsum(x, axis=0))
+  @slow_test
   def test_cumsum(self):
     helper_test_op([()], lambda x: torch.cumsum(x, dim=0), lambda x: Tensor.cumsum(x, axis=0))
     self.helper_test_exception([()], lambda x: torch.cumsum(x, dim=1), lambda x: Tensor.cumsum(x, axis=1), expected=IndexError)
@@ -1029,9 +1034,11 @@ class TestOps(unittest.TestCase):
 
   def test_small_cumprod(self):
     helper_test_op([(10)],lambda x: torch.cumprod(x, dim=0),lambda x: Tensor.cumprod(x, axis=0))
+  @slow_test
   def test_simple_cumprod(self):
     helper_test_op([(512)],lambda x: torch.cumprod(x, dim=0),lambda x: Tensor.cumprod(x, axis=0))
     helper_test_op([(1022)],lambda x: torch.cumprod(x, dim=0),lambda x: Tensor.cumprod(x, axis=0))
+  @slow_test
   def test_cumprod(self):
     helper_test_op([()],lambda x: torch.cumprod(x, dim=0),lambda x: Tensor.cumprod(x, axis=0))
     self.helper_test_exception([()],lambda x: torch.cumprod(x, dim=1),lambda x: Tensor.cumprod(x, axis=1),expected=IndexError)
@@ -1049,9 +1056,11 @@ class TestOps(unittest.TestCase):
 
   def test_small_cummax(self):
     helper_test_op([(10)], lambda x: torch.cummax(x, dim=0).values, lambda x: Tensor.cummax(x, axis=0))
+  @slow_test
   def test_simple_cummax(self):
     helper_test_op([(512)], lambda x: torch.cummax(x, dim=0).values, lambda x: Tensor.cummax(x, axis=0))
     helper_test_op([(1022)], lambda x: torch.cummax(x, dim=0).values, lambda x: Tensor.cummax(x, axis=0))
+  @slow_test
   def test_cummax(self):
     helper_test_op([()], lambda x: torch.cummax(x, dim=0).values, lambda x: Tensor.cummax(x, axis=0))
     # TODO: torch allows this?
@@ -1128,12 +1137,12 @@ class TestOps(unittest.TestCase):
                                   lambda x: x.argsort(dim, descending), forward_only=True)
 
   def test_topk(self):
-    helper_test_op([(10)], lambda x: x.topk(3).values, lambda x: x.topk(3)[0], forward_only=True)
-    helper_test_op([(10)], lambda x: x.topk(3).indices.type(torch.int32), lambda x: x.topk(3)[1], forward_only=True)
+    helper_test_op([(8)], lambda x: x.topk(3).values, lambda x: x.topk(3)[0], forward_only=True)
+    helper_test_op([(8)], lambda x: x.topk(3).indices.type(torch.int32), lambda x: x.topk(3)[1], forward_only=True)
     for dim in [0, 1, -1]:
       for largest in [True, False]:
         for sorted_ in [True]: # TODO support False
-          helper_test_op([(6,5,4)],
+          helper_test_op([(5,5,4)],
                           lambda x: x.topk(4, dim, largest, sorted_).values,
                           lambda x: x.topk(4, dim, largest, sorted_)[0], forward_only=True)
           helper_test_op([(5,5,4)],
@@ -1148,53 +1157,55 @@ class TestOps(unittest.TestCase):
     np.testing.assert_equal(indices.numpy(), [2, 4, 6])
     self.helper_test_exception([(4)], lambda x: x.topk(5), expected=(RuntimeError, ValueError))
 
+  @slow_test
   def test_einsum(self):
     # matrix transpose
-    helper_test_op([(150,150)], lambda a: torch.einsum('ij->ji', a), lambda a: Tensor.einsum('ij->ji', a))
-    helper_test_op([(150,150)], lambda a: torch.einsum('ij -> ji', a), lambda a: Tensor.einsum('ij -> ji', a))
-    helper_test_op([(150,150)], lambda a: torch.einsum('ji', a), lambda a: Tensor.einsum('ji', a))
-    helper_test_op([(20,30,40)], lambda a: torch.einsum('jki', a), lambda a: Tensor.einsum('jki', a))
-    helper_test_op([(20,30,40)], lambda a: torch.einsum('dog', a), lambda a: Tensor.einsum('dog', a))
+    helper_test_op([(10,10)], lambda a: torch.einsum('ij->ji', a), lambda a: Tensor.einsum('ij->ji', a))
+    helper_test_op([(10,10)], lambda a: torch.einsum('ij -> ji', a), lambda a: Tensor.einsum('ij -> ji', a))
+    helper_test_op([(10,10)], lambda a: torch.einsum('ji', a), lambda a: Tensor.einsum('ji', a))
+    helper_test_op([(4,6,8)], lambda a: torch.einsum('jki', a), lambda a: Tensor.einsum('jki', a))
+    helper_test_op([(4,6,8)], lambda a: torch.einsum('dog', a), lambda a: Tensor.einsum('dog', a))
     # no -> and empty rhs
-    helper_test_op([(20,30),(30,40)], lambda a, b: torch.einsum('ij,jk', a, b), lambda a, b: Tensor.einsum('ij,jk', a, b))
+    helper_test_op([(4,6),(6,8)], lambda a, b: torch.einsum('ij,jk', a, b), lambda a, b: Tensor.einsum('ij,jk', a, b))
     # sum all elements
-    helper_test_op([(20,30,40)], lambda a: torch.einsum('ijk->', a), lambda a: Tensor.einsum('ijk->', a))
+    helper_test_op([(4,6,8)], lambda a: torch.einsum('ijk->', a), lambda a: Tensor.einsum('ijk->', a))
     # column sum
-    helper_test_op([(50,50)], lambda a: torch.einsum('ij->j', a), lambda a: Tensor.einsum('ij->j', a))
+    helper_test_op([(5,5)], lambda a: torch.einsum('ij->j', a), lambda a: Tensor.einsum('ij->j', a))
     # row sum
-    helper_test_op([(15,15)], lambda a: torch.einsum('ij->i', a), lambda a: Tensor.einsum('ij->i', a))
+    helper_test_op([(5,5)], lambda a: torch.einsum('ij->i', a), lambda a: Tensor.einsum('ij->i', a))
     # matrix-vector multiplication
-    helper_test_op([(15,20), (20,)], lambda a,b: torch.einsum('ik,k->i', a,b), lambda a,b: Tensor.einsum('ik,k->i', a, b))
+    helper_test_op([(3,4), (4,)], lambda a,b: torch.einsum('ik,k->i', a,b), lambda a,b: Tensor.einsum('ik,k->i', a, b))
     # matrix-matrix multiplication
-    helper_test_op([(15,20), (20,30)], lambda a,b: torch.einsum('ik,kj->ij', a,b), lambda a,b: Tensor.einsum('ik,kj->ij', a, b))
+    helper_test_op([(3,4), (4,5)], lambda a,b: torch.einsum('ik,kj->ij', a,b), lambda a,b: Tensor.einsum('ik,kj->ij', a, b))
     # matrix-matrix multiplication, different letter order
-    helper_test_op([(15,20), (20,30)], lambda a,b: torch.einsum('jk,ki->ji', a,b), lambda a,b: Tensor.einsum('jk,ki->ji', a, b))
+    helper_test_op([(3,4), (4,5)], lambda a,b: torch.einsum('jk,ki->ji', a,b), lambda a,b: Tensor.einsum('jk,ki->ji', a, b))
     # dot product
-    helper_test_op([(30),(30)], lambda a,b: torch.einsum('i,i->i', [a,b]), lambda a,b: Tensor.einsum('i,i->i', [a,b]))
+    helper_test_op([(5),(5)], lambda a,b: torch.einsum('i,i->i', [a,b]), lambda a,b: Tensor.einsum('i,i->i', [a,b]))
     # hadamard product
-    helper_test_op([(30,40),(30,40)], lambda a,b: torch.einsum('ij,ij->ij', a,b), lambda a,b: Tensor.einsum('ij,ij->ij', a,b))
+    helper_test_op([(5,6),(5,6)], lambda a,b: torch.einsum('ij,ij->ij', a,b), lambda a,b: Tensor.einsum('ij,ij->ij', a,b))
     # outer product
-    helper_test_op([(15,), (15,)], lambda a,b: torch.einsum('i,j->ij', a,b), lambda a,b: Tensor.einsum('i,j->ij',a,b))
+    helper_test_op([(5,), (5,)], lambda a,b: torch.einsum('i,j->ij', a,b), lambda a,b: Tensor.einsum('i,j->ij',a,b))
     # batch matrix multiplication
-    helper_test_op([(10,20,30),(10,30,40)], lambda a,b: torch.einsum('ijk,ikl->ijl', [a, b]), lambda a,b: Tensor.einsum('ijk,ikl->ijl', [a, b]))
+    helper_test_op([(2,4,6),(2,6,8)], lambda a,b: torch.einsum('ijk,ikl->ijl', [a, b]), lambda a,b: Tensor.einsum('ijk,ikl->ijl', [a, b]))
     # batch matrix multiplication, result permuted
-    helper_test_op([(10,20,25),(10,25,32)], lambda a,b: torch.einsum('ijk,ikl->jil', [a, b]), lambda a,b: Tensor.einsum('ijk,ikl->jil', [a, b]))
+    helper_test_op([(2,4,5),(2,5,7)], lambda a,b: torch.einsum('ijk,ikl->jil', [a, b]), lambda a,b: Tensor.einsum('ijk,ikl->jil', [a, b]))
     # batch matrix multiplication, result & input permuted
-    helper_test_op([(20,10,25),(10,25,32)], lambda a,b: torch.einsum('jik,ikl->jil', [a, b]), lambda a,b: Tensor.einsum('jik,ikl->jil', [a, b]))
+    helper_test_op([(4,2,5),(2,5,7)], lambda a,b: torch.einsum('jik,ikl->jil', [a, b]), lambda a,b: Tensor.einsum('jik,ikl->jil', [a, b]))
     # batch matrix multiplication, result with different letters
-    helper_test_op([(10,20,30),(10,30,40)], lambda a,b: torch.einsum('ijk,ika->ija', [a, b]), lambda a,b: Tensor.einsum('ijk,ika->ija', [a, b]))
+    helper_test_op([(2,4,6),(2,6,8)], lambda a,b: torch.einsum('ijk,ika->ija', [a, b]), lambda a,b: Tensor.einsum('ijk,ika->ija', [a, b]))
     # tensor contraction
-    helper_test_op([(3,5,8,10),(11,13,5,16,8)], lambda a,b: torch.einsum('pqrs,tuqvr->pstuv', a,b),
+    helper_test_op([(3,5,8,10),(11,7,5,13,8)], lambda a,b: torch.einsum('pqrs,tuqvr->pstuv', a,b),
                                                 lambda a,b: Tensor.einsum('pqrs,tuqvr->pstuv', a,b), atol=1e-5)
     # tensor contraction, input permuted
-    helper_test_op([(3,8,10,5),(11,5,13,16,8)], lambda a,b: torch.einsum('prsq,tquvr->pstuv', a,b),
+    helper_test_op([(3,8,10,5),(11,5,7,13,8)], lambda a,b: torch.einsum('prsq,tquvr->pstuv', a,b),
                                                 lambda a,b: Tensor.einsum('prsq,tquvr->pstuv', a,b), atol=1e-5)
     # tensor contraction, result with different letters
-    helper_test_op([(3,5,8,10),(11,13,5,16,8)], lambda a,b: torch.einsum('zqrs,tuqvr->zstuv', a,b),
+    helper_test_op([(3,5,8,10),(11,7,5,13,8)], lambda a,b: torch.einsum('zqrs,tuqvr->zstuv', a,b),
                                                 lambda a,b: Tensor.einsum('zqrs,tuqvr->zstuv', a,b), atol=1e-5)
     # bilinear transformation
     helper_test_op([(2,3),(5,3,7),(2,7)], lambda a,b,c: torch.einsum('ik,jkl,il->ij', [a,b,c]), lambda a,b,c: Tensor.einsum('ik,jkl,il->ij', [a,b,c]))
 
+  @slow_test
   def test_einsum_ellipsis(self):
     """The expected behavior for einsum is described in the PyTorch docs: https://pytorch.org/docs/stable/generated/torch.einsum.html"""
     # test ellipsis
@@ -1209,32 +1220,24 @@ class TestOps(unittest.TestCase):
     # match torch ellipsis handling
     helper_test_op([(32, 7, 24, 24, 24), (32, 7, 24, 24, 24)], lambda a, b: torch.einsum('ij...,ij...->ij', [a, b]),
                lambda a, b: Tensor.einsum('ij...,ij...->ij', [a, b]))
-    # multiple ellipsis in one operand are not allowed. This test shall raise an exception.
-    with self.assertRaises(RuntimeError):
-      helper_test_op([(2, 3, 4), (2, 3, 4)], lambda a, b: torch.einsum('...ik..., ...jk ->', [a, b]),
-                       lambda a, b: Tensor.einsum('...ik..., ...jk ->', [a, b]))
-    # multiple ellipsis must broadcast together. This test shall raise an exception.
-    with self.assertRaises(RuntimeError):
-      helper_test_op([(2, 3, 4, 5), (5, 2, 7)], lambda a, b: torch.einsum('i...j,ji...->...', [a, b]),
-                       lambda a, b: Tensor.einsum('i...j,ji...->...', [a, b]))
+    # multiple ellipsis in one operand are not allowed
+    self.helper_test_exception([(2, 3, 4), (2, 3, 4)], lambda a, b: torch.einsum('...ik..., ...jk ->', [a, b]),
+                lambda a, b: Tensor.einsum('...ik..., ...jk ->', [a, b]), expected=(RuntimeError, IndexError))
+    # multiple ellipsis must broadcast together
+    self.helper_test_exception([(2, 3, 4), (2, 3, 4)], lambda a, b: torch.einsum('i...j,ji...->...', [a, b]),
+                lambda a, b: Tensor.einsum('i...j,ji...->...', [a, b]), expected=RuntimeError)
 
   def test_einsum_shape_check(self):
-    a = Tensor.zeros(3,8,10,5)
-    b = Tensor.zeros(11,5,13,16,8)
-    with self.assertRaises(AssertionError):
-      Tensor.einsum('pqrs,tuqvr->pstuv',a,b)
+    self.helper_test_exception([(3,8,10,5), (11,5,13,16,8)], lambda a, b: torch.einsum('pqrs,tuqvr->pstuv', [a, b]),
+                lambda a, b: Tensor.einsum('pqrs,tuqvr->pstuv', [a, b]), expected=RuntimeError)
 
   def test_einsum_arity_check1(self):
-    a = Tensor.zeros(10,15)
-    b = Tensor.zeros(15,20)
-    c = Tensor.zeros(20,10)
-    with self.assertRaises(AssertionError):
-      Tensor.einsum('ij,jk->ij', a,b,c)
+    self.helper_test_exception([(10,15), (15,20), (20,10)], lambda a, b, c: torch.einsum('ij,jk->ij', [a, b, c]),
+                lambda a, b, c: Tensor.einsum('ij,jk->ij', [a, b, c]), expected=(ValueError, RuntimeError))
 
   def test_einsum_arity_check2(self):
-    a = Tensor.zeros(10,10)
-    with self.assertRaises(AssertionError):
-      Tensor.einsum('ij,jk->ij', a)
+    self.helper_test_exception([(10,10)], lambda a: torch.einsum('ij,jk->ij', a),
+                lambda a: Tensor.einsum('ij,jk->ij', a), expected=(ValueError, RuntimeError))
 
   @unittest.skipIf(IMAGE>0, "no 1d dot for images")
   def test_dot_1d(self):
@@ -1246,6 +1249,7 @@ class TestOps(unittest.TestCase):
     self.helper_test_exception([(4), (1,2)], lambda x, y: x.matmul(y), Tensor.dot, expected=RuntimeError)
     self.helper_test_exception([(2,1), (4)], lambda x, y: x.matmul(y), Tensor.dot, expected=RuntimeError)
     self.helper_test_exception([(1), (4)], lambda x, y: x.matmul(y), Tensor.dot, expected=RuntimeError)
+  @slow_test
   def test_dot(self):
     helper_test_op([(45,65), (65,100)], lambda x,y: x.matmul(y), Tensor.dot, atol=1e-5)
     helper_test_op([(8,45,65), (8,65,100)], lambda x,y: x.matmul(y), Tensor.dot, atol=1e-5)
@@ -1306,6 +1310,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(64,64), (64,64)], lambda x,y: x.half().matmul(y.half()), atol=5e-3, rtol=5e-3)
   def test_gemm(self):
     helper_test_op([(64,64), (64,64)], lambda x,y: x.matmul(y))
+  @slow_test
   def test_big_gemm(self):
     helper_test_op([(256,256), (256,256)], lambda x,y: x.matmul(y), atol=1e-4)
   @unittest.skipIf(IMAGE>0, "no 0 in shape matmul on images")
@@ -1317,12 +1322,14 @@ class TestOps(unittest.TestCase):
     helper_test_op([(0,0), (0,0)], lambda x,y: x.matmul(y), Tensor.dot, atol=1e-7)
     helper_test_op([(0), (0,8)], lambda x,y: x.matmul(y), Tensor.dot, atol=1e-7)
     helper_test_op([(0), (0)], lambda x,y: x.matmul(y), Tensor.dot, atol=1e-7)
+  @slow_test
   def test_broadcastdot(self):
     helper_test_op([(10,45,65), (65,45)], lambda x,y: x @ y, Tensor.dot, atol=1e-4)
     with self.assertRaises(RuntimeError):
       a = Tensor(3.14)
       b = Tensor.ones(3,3)
       a @ b
+  @slow_test
   def test_multidot(self):
     helper_test_op([(10,45,65), (10,65,45)], lambda x,y: x @ y, Tensor.dot, atol=1e-4)
     helper_test_op([(3,3,45,65), (3,3,65,45)], lambda x,y: x @ y, Tensor.dot, atol=1e-4)
@@ -1457,6 +1464,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(15, 25, 35)], lambda x: x.var(correction=5))
     # TODO: fix this
     # helper_test_op([(10, 2)], lambda x: x.var(correction=50))
+  @slow_test
   def test_var_axis(self):
     helper_test_op([(15, 25, 35)], lambda x: x.var(0))
     helper_test_op([(15, 25, 35)], lambda x: x.var(2))
@@ -1489,6 +1497,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(15, 25, 35)], lambda x: x.std())
     helper_test_op([(15, 25, 35)], lambda x: x.std(correction=0))
     helper_test_op([(15, 25, 35)], lambda x: x.std(correction=5))
+  @slow_test
   def test_std_axis(self):
     helper_test_op([(15, 25, 35)], lambda x: x.std(0))
     helper_test_op([(15, 25, 35)], lambda x: x.std(2))
@@ -1565,6 +1574,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([()], lambda x: torch.logsumexp(x, dim=0), lambda x: x.logsumexp(0), atol=1e-7, grad_atol=1e-7)
     helper_test_op([()], lambda x: torch.logsumexp(x, dim=-1), lambda x: x.logsumexp(-1), atol=1e-7, grad_atol=1e-7)
 
+  @slow_test
   def test_logcumsumexp(self):
     helper_test_op([(45,65)], lambda x: torch.logcumsumexp(x, dim=0), lambda x: x.logcumsumexp(0), atol=1e-7, grad_atol=1e-7)
     helper_test_op([(45,65)], lambda x: torch.logcumsumexp(x, dim=1), lambda x: x.logcumsumexp(1), atol=1e-7, grad_atol=1e-7)
@@ -1634,6 +1644,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65), (45,1)], lambda x,y: x/y)
     helper_test_op([(45,65), ()], lambda x,y: x/y)
 
+  @slow_test
   def test_broadcast_partial(self):
     for torch_op, tinygrad_op in [(torch.add, Tensor.add), (torch.sub, Tensor.sub), (torch.mul, Tensor.mul),
                                   (torch.div, Tensor.div), (torch.pow, Tensor.pow)]:
@@ -2056,6 +2067,7 @@ class TestOps(unittest.TestCase):
       lambda x,w: torch.nn.functional.conv2d(x,w),
       lambda x,w: Tensor.conv2d(x,w), grad_rtol=1e-5)
 
+  @slow_test
   def test_nested_conv2d(self):
     helper_test_op([(1,32,9,9), (32,32,3,3), (32,32,3,3)],
       lambda x,w1,w2: torch.nn.functional.conv2d(torch.nn.functional.conv2d(x,w1).relu(), w2),
@@ -2116,6 +2128,7 @@ class TestOps(unittest.TestCase):
         lambda x,w,b: torch.nn.functional.conv_transpose2d(x,w,b,output_padding=output_padding,stride=stride),
         lambda x,w,b: Tensor.conv_transpose2d(x,w,b,output_padding=output_padding,stride=stride), grad_rtol=1e-5)
 
+  @slow_test
   @unittest.skipIf(IMAGE>0, "no conv3d on images")
   def test_simple_conv_transpose3d(self):
     helper_test_op([(2,4,9,9,9), (4,4,3,3,3)],
@@ -2185,6 +2198,7 @@ class TestOps(unittest.TestCase):
     self.helper_test_exception([(2,16,2,2), (32,16,3,3)], lambda x,w:torch.nn.functional.conv2d(x,w,padding=(1,1,1)),
                                lambda x,w: Tensor.conv2d(x,w,padding=(1,1,1)), expected=(RuntimeError, ValueError))
 
+  @slow_test
   def test_large_input_conv2d(self):
     bs = 4
     cin = 16
@@ -2337,6 +2351,7 @@ class TestOps(unittest.TestCase):
       lambda x: torch.nn.functional.max_pool2d(x, kernel_size=ksz),
       lambda x: Tensor.max_pool2d(x, kernel_size=ksz))
 
+  @slow_test
   def test_max_pool2d(self):
     for ksz in [(2,2), (3,3), 2, 3, (3,2), (5,5), (5,1)]:
       with self.subTest(kernel_size=ksz):
@@ -2344,41 +2359,45 @@ class TestOps(unittest.TestCase):
           lambda x: torch.nn.functional.max_pool2d(x, kernel_size=ksz),
           lambda x: Tensor.max_pool2d(x, kernel_size=ksz))
 
+  @slow_test
   def test_max_pool2d_padding(self):
     for ksz in [(2,2), (3,3), 2, 3, (3,2)]:
       for p in [1, (1,0), (0,1)]:
         with self.subTest(kernel_size=ksz, padding=p):
-          helper_test_op([(32,2,11,28)],
+          helper_test_op([(4,2,11,28)],
             lambda x: torch.nn.functional.max_pool2d(x, kernel_size=ksz, padding=p),
             lambda x: Tensor.max_pool2d(x, kernel_size=ksz, padding=p))
-    self.helper_test_exception([(32,2,110,28)], lambda x: torch.nn.functional.max_pool2d(x, kernel_size=(2,2), padding=(1,1,1)),
+    self.helper_test_exception([(4,2,110,28)], lambda x: torch.nn.functional.max_pool2d(x, kernel_size=(2,2), padding=(1,1,1)),
                    lambda x: Tensor.max_pool2d(x, kernel_size=(2,2), padding=(1,1,1)), expected=(RuntimeError, ValueError))
 
+  @slow_test
   def test_max_pool2d_asymmetric_padding(self):
-    shape = (32,2,111,28)
     for p in [(0,1,0,1), (2,1,2,1), (2,0,2,1)]:
       with self.subTest(padding=p):
-        helper_test_op([shape],
+        helper_test_op([(4,2,111,28)],
           lambda x: torch.nn.functional.max_pool2d(torch.nn.functional.pad(x, p, value=float("-inf")), kernel_size=(5,5)),
           lambda x: Tensor.max_pool2d(x, kernel_size=(5,5), padding=p))
 
+  @slow_test
   def test_max_pool2d_padding_int(self):
     ksz = (2,2)
-    helper_test_op([(32,2,11,28)],
+    helper_test_op([(4,2,11,28)],
       lambda x: torch.nn.functional.max_pool2d(x.int(), kernel_size=ksz, padding=1),
       lambda x: Tensor.max_pool2d(x.int(), kernel_size=ksz, padding=1), forward_only=True)
 
+  @slow_test
   def test_max_pool2d_bigger_stride(self):
     for stride in [(2,3), (3,2), 2, 3]:
       with self.subTest(stride=stride):
-        helper_test_op([(32,2,11,28)],
+        helper_test_op([(4,2,11,28)],
           lambda x: torch.nn.functional.max_pool2d(x, kernel_size=(2,2), stride=stride),
           lambda x: Tensor.max_pool2d(x, kernel_size=(2,2), stride=stride))
 
+  @slow_test
   def test_max_pool2d_bigger_stride_dilation(self):
     for stride, dilation in zip([(2,3), (3,2), 2, 3, 4], [(3,2), (2,3), 2, 3, 6]):
       with self.subTest(stride=stride):
-        helper_test_op([(32,2,11,28)],
+        helper_test_op([(4,2,11,28)],
           lambda x: torch.nn.functional.max_pool2d(x, kernel_size=(2,2), stride=stride, dilation=dilation),
           lambda x: Tensor.max_pool2d(x, kernel_size=(2,2), stride=stride, dilation=dilation))
 
@@ -2388,6 +2407,7 @@ class TestOps(unittest.TestCase):
       lambda x: torch.nn.functional.max_pool2d(x, kernel_size=(5,5), stride=1),
       lambda x: Tensor.max_pool2d(x, kernel_size=(5,5), stride=1))
 
+  @slow_test
   def test_max_pool2d_smaller_stride(self):
     for stride in [(2,3), (3,2), 2, 3]:
       with self.subTest(stride=stride):
@@ -2395,6 +2415,7 @@ class TestOps(unittest.TestCase):
           lambda x: torch.nn.functional.max_pool2d(x, kernel_size=(5,5), stride=stride),
           lambda x: Tensor.max_pool2d(x, kernel_size=(5,5), stride=stride))
 
+  @slow_test
   def test_max_pool2d_dilation(self):
     for dilation in [(2, 3), (3, 2), 2, 3]:
       helper_test_op([(3, 2, 17, 14)],
@@ -2449,6 +2470,7 @@ class TestOps(unittest.TestCase):
       lambda x: Tensor.max_pool2d(x, kernel_size=(2,2), stride=1, return_indices=True)[1],
       vals=[[[[[1,2]*3]*6]]], forward_only=True)  # Tensor([1,2,1,2,1,2]).expand(1,1,6,6)
 
+  @slow_test
   def test_max_unpool2d(self):
     args = {"kernel_size":(5,5), "stride":(6,5)}
     helper_test_op([(8,3,50,50)],
@@ -2492,6 +2514,7 @@ class TestOps(unittest.TestCase):
       lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=(1,2), padding=(0,1), stride=(5,1)),
       lambda x: Tensor.avg_pool2d(x, kernel_size=(1,2), padding=(0,1), stride=(5,1)), rtol=1e-5)
 
+  @slow_test
   def test_avg_pool2d_padding(self):
     shape = (32,2,111,28)
     for ksz in [(2,2), (3,3), 2, 3, (3,2)]:
@@ -2513,6 +2536,7 @@ class TestOps(unittest.TestCase):
     self.helper_test_exception([shape], lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=(2,2), padding=(1,1,1)),
                                lambda x: Tensor.avg_pool2d(x, kernel_size=(2,2), padding=(1,1,1)), expected=(RuntimeError, ValueError))
 
+  @slow_test
   def test_avg_pool2d_padding_not_counted(self):
     shape = (32,2,111,28)
     for ksz in [(2,2), (3,3), 2, 3, (3,2)]:
@@ -2594,12 +2618,14 @@ class TestOps(unittest.TestCase):
 
   def test_interpolate_nearest_exact(self): self.test_interpolate_nearest("nearest-exact")
 
+  @slow_test
   def test_interpolate_bilinear(self):
     for in_sz, out_sz in [((12,20),(9,31)), ((12,9),(31,20)), ((9,31),(20,12))]:
       helper_test_op([(2,3)+in_sz],
         lambda x: torch.nn.functional.interpolate(x, size=out_sz, mode="bilinear"),
         lambda x: Tensor.interpolate(x, size=out_sz, mode="linear"), atol=1e-4)
 
+  @slow_test
   def test_interpolate_bilinear_corners_aligned(self):
     for in_sz, out_sz in [((12,20),(9,31)), ((12,9),(31,20)), ((9,31),(20,12))]:
       helper_test_op([(2,3)+in_sz],
@@ -2618,6 +2644,7 @@ class TestOps(unittest.TestCase):
         lambda x: torch.nn.functional.interpolate(x, size=out_sz, mode="trilinear", align_corners=True),
         lambda x: Tensor.interpolate(x, size=out_sz, mode="linear", align_corners=True), atol=1e-4)
 
+  @slow_test
   def test_cat(self):
     for dim in range(-2, 3):
       helper_test_op([(45,65,9), (45,65,9), (45,65,9)], lambda x,y,z: torch.cat((x,y,z), dim), lambda x,y,z: x.cat(y, z, dim=dim))
@@ -2716,6 +2743,7 @@ class TestOps(unittest.TestCase):
     data = [math.inf, -math.inf, math.nan]
     helper_test_op((), lambda: torch.tensor(data)[torch.tensor([0, 1, 2])], lambda: Tensor(data)[Tensor([0, 1, 2])])
 
+  @slow_test
   def test_slice_fancy_indexing_no_dim_collapse(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     # no dim collapse from int or dim injection from None
@@ -2725,6 +2753,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,...,e], lambda x: x[i,...,p])
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[...,c,:,e], lambda x: x[...,k,:,p])
 
+  @slow_test
   def test_slice_fancy_indexing_dim_collapse_int(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     # dim collapse from int
@@ -2734,6 +2763,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,2,2,2,e], lambda x: x[i,2,2,2,p])
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[1,:,3:11:2,d,0:2], lambda x: x[1,:,3:11:2,o,0:2])
 
+  @slow_test
   def test_slice_fancy_indexing_dim_inject_none(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     # dim injection from None
@@ -2767,6 +2797,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,3)], lambda x: x[torch.tensor([[0,1,-1],[-1,-2,0]]), torch.tensor([2,1,-1])],
                             lambda x: x[Tensor([[0,1,-1],[-1,-2,0]]), Tensor([2,1,-1])])
 
+  @slow_test
   def test_slice_fancy_indexing_list_indices(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[[[0]]], lambda x: x[[[0]]])
@@ -2777,6 +2808,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,b,c,[[1],[2],[3]],...], lambda x: x[i,j,k,[[1],[2],[3]],...])
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,[2,1,0],c,[-2,1,0],e], lambda x: x[i,[2,1,0],k,[-2,1,0],p])
 
+  @slow_test
   def test_slice_fancy_indexing_tuple_indices(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[(((0,),),)], lambda x: x[(((0,),),)])
@@ -2786,6 +2818,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,((2,),(1,),(0,)),c,(2,1,0)], lambda x: x[i,((2,),(1,),(0,)),k,(2,1,0)])
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[1,(2,1,0),None,c,(2,1,0),e], lambda x: x[1,(2,1,0),None,k,(2,1,0),p])
 
+  @slow_test
   def test_slice_fancy_indexing_list_with_tensors(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[[a]], lambda x: x[[i]])
@@ -2892,6 +2925,7 @@ class TestOps(unittest.TestCase):
     with self.assertRaises(TypeError):
       Tensor.ones(4).scatter(dim=1, index=Tensor([0]), src=Tensor.ones(4), reduce="add")
 
+  @slow_test
   def test_scatter_reduce(self):
     b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False)
     a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32, requires_grad=False)
@@ -2927,15 +2961,18 @@ class TestOps(unittest.TestCase):
       lambda x,src: x.half().scatter_reduce(dim=0, index=a, src=src, reduce="sum"),
       RuntimeError)
 
+  @slow_test
   def test_scaled_dot_product_attention(self):
     helper_test_op([(32,8,16,64), (32,8,16,64), (32,8,16,64)], torch.nn.functional.scaled_dot_product_attention, Tensor.scaled_dot_product_attention)
     helper_test_op([(32,8,16,64), (32,8,16,64), (32,8,16,64), (32,8,16,16)],
                    lambda x,y,z,m: torch.nn.functional.scaled_dot_product_attention(x,y,z,attn_mask=m),
                    lambda x,y,z,m: Tensor.scaled_dot_product_attention(x,y,z,attn_mask=m))
 
+  @slow_test
   def test_scaled_dot_product_attention_mismatch_ls(self):
     helper_test_op([(32,8,4,64), (32,8,16,64), (32,8,16,64)], torch.nn.functional.scaled_dot_product_attention, Tensor.scaled_dot_product_attention)
 
+  @slow_test
   def test_scaled_dot_product_attention_causal(self):
     helper_test_op([(32,8,16,64), (32,8,16,64), (32,8,16,64)],
                    lambda x,y,z: torch.nn.functional.scaled_dot_product_attention(x,y,z,is_causal=True),
@@ -2946,6 +2983,7 @@ class TestOps(unittest.TestCase):
       lambda x,y,z,m: Tensor.scaled_dot_product_attention(x,y,z,is_causal=True,attn_mask=m),
       expected=RuntimeError)
 
+  @slow_test
   def test_scaled_dot_product_attention_gqa(self):
     helper_test_op([(32,32,16,64), (32,8,16,64), (32,8,16,64)],
                    lambda x,y,z: torch.nn.functional.scaled_dot_product_attention(x,y,z,enable_gqa=True),

test/test_opt_gemm.py

@@ -2,7 +2,7 @@ import numpy as np
 import unittest
 from tinygrad import Tensor
 from tinygrad.helpers import get_single_element
-from tinygrad.codegen.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 
 class TestOptGemm(unittest.TestCase):

test/test_opts.py

@@ -0,0 +1,22 @@
+import unittest
+from tinygrad import Tensor, Device
+from tinygrad.helpers import RANGEIFY
+from tinygrad.codegen.opt import Opt, OptOps
+from tinygrad.engine.realize import get_program
+
+@unittest.skipIf(RANGEIFY>0, "arg is partial contig in rangeify")
+class TestOpts(unittest.TestCase):
+  def test_opt_upcast(self):
+    opts = (Opt(OptOps.UPCAST, 0, 4),)
+    a = Tensor.empty(16)
+    b = Tensor.empty(16)
+    out = (a+b).contiguous(arg=opts)
+    s = out.schedule()
+    self.assertEqual(s[-1].ast.arg.opts_to_apply, opts)
+    if Device.DEFAULT in {"CPU", "GPU", "METAL"}:
+      prg = get_program(s[-1].ast)
+      self.assertIn('float4', prg.src)
+
+if __name__ == '__main__':
+  unittest.main()
+

test/test_outerworld_range.py

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad import Tensor, nn, Variable, UOp, dtypes
+from tinygrad import Tensor, nn, Variable, UOp
 
 # outerworld range should support three things
 # 1. full optimizer steps (test_model_bound_range)
@@ -136,7 +136,7 @@ class TestOuterworldRange(unittest.TestCase):
   def test_model_bound_range(self):
     m, opt = get_model_and_opt()
     # TODO: should ranges be unique so you don't have to pass in the -1?
-    rng = UOp.range(dtypes.int, self.STEPS, -1)
+    rng = UOp.range(self.STEPS, -1)
     vib = Variable('i', 0, self.STEPS-1).bind(rng)
     loss = (m(self.X[vib]) - self.Y[vib]).square().mean()
     loss.backward()

test/test_quantize_onnx.py

@@ -3,11 +3,9 @@ import numpy as np
 import unittest
 from dataclasses import replace
 from tinygrad import Tensor, Context, Device, dtypes
-from tinygrad.uop.ops import Ops, UOp # noqa: F401 # pylint: disable=unused-import
-from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.uop.ops import Ops
+from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem, lower_schedule_item, get_program
-from tinygrad.codegen.opt.search import bufs_from_lin
-from tinygrad.shape.shapetracker import ShapeTracker, View # noqa: F401 # pylint: disable=unused-import
 
 N = 512
 
@@ -236,129 +234,5 @@ class TestQuantizeOnnx(unittest.TestCase):
     opts = [Opt(op=OptOps.UPCAST, axis=0, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
     sexec(out, opts)
 
-@unittest.skipIf(Device.DEFAULT != "DSP", "only tests for DSP")
-class TestDSPCache(unittest.TestCase):
-  def test_cache_speed(self):
-    # string becuase this breaks Python language server for syntax highlight for some reason
-    ast = eval("""UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.VIEW, dtypes.uchar.ptr(25088), arg=ShapeTracker(views=(View(shape=(1, 28, 28, 32, 1), strides=(0, 896, 32, 1, 0), offset=0, mask=None, contiguous=True),)), src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(25088), arg=0, src=()),)),
-        UOp(Ops.CAST, dtypes.uchar, arg=None, src=(
-          UOp(Ops.XOR, dtypes.int, arg=None, src=(
-            UOp(Ops.MAX, dtypes.int, arg=None, src=(
-              UOp(Ops.XOR, dtypes.int, arg=None, src=(
-                UOp(Ops.MAX, dtypes.int, arg=None, src=(
-                  UOp(Ops.CAST, dtypes.int, arg=None, src=(
-                    UOp(Ops.ADD, dtypes.float, arg=None, src=(
-                      UOp(Ops.ADD, dtypes.float, arg=None, src=(
-                        UOp(Ops.MUL, dtypes.float, arg=None, src=(
-                          UOp(Ops.ADD, dtypes.float, arg=None, src=(
-                            UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (4,)), src=(
-                              UOp(Ops.MUL, dtypes.float, arg=None, src=(
-                                UOp(Ops.MUL, dtypes.float, arg=None, src=(
-                                  UOp(Ops.CAST, dtypes.float, arg=None, src=(
-                                    UOp(Ops.CAST, dtypes.int, arg=None, src=(
-                                      UOp(Ops.LOAD, dtypes.uchar, arg=None, src=(
-                                        UOp(Ops.VIEW, dtypes.uchar.ptr(150528), arg=ShapeTracker(views=(View(shape=(1, 28, 28, 32, 192), strides=(0, 5376, 192, 0, 1), offset=0, mask=None, contiguous=False),)), src=(
-                                          UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(150528), arg=1, src=()),)),)),)),)),
-                                  UOp(Ops.CONST, dtypes.float, arg=0.012368360534310341, src=(
-                                    x22:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 28, 28, 32, 192), strides=(0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
-                                UOp(Ops.MUL, dtypes.float, arg=None, src=(
-                                  UOp(Ops.CAST, dtypes.float, arg=None, src=(
-                                    UOp(Ops.CAST, dtypes.int, arg=None, src=(
-                                      UOp(Ops.LOAD, dtypes.char, arg=None, src=(
-                                        UOp(Ops.VIEW, dtypes.char.ptr(6144), arg=ShapeTracker(views=(View(shape=(32, 48, 4), strides=(4, 128, 1), offset=0, mask=None, contiguous=False), View(shape=(1, 28, 28, 32, 192), strides=(0, 0, 0, 192, 1), offset=0, mask=None, contiguous=False))), src=(
-                                          UOp(Ops.DEFINE_GLOBAL, dtypes.char.ptr(6144), arg=2, src=()),)),)),)),)),
-                                  UOp(Ops.CONST, dtypes.float, arg=0.007441135589033365, src=(
-                                    x22,)),)),)),)),
-                            UOp(Ops.MUL, dtypes.float, arg=None, src=(
-                              UOp(Ops.CAST, dtypes.float, arg=None, src=(
-                                UOp(Ops.LOAD, dtypes.int, arg=None, src=(
-                                  UOp(Ops.VIEW, dtypes.int.ptr(32), arg=ShapeTracker(views=(View(shape=(1, 28, 28, 32, 1), strides=(0, 0, 0, 1, 0), offset=0, mask=None, contiguous=False),)), src=(
-                                    UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(32), arg=3, src=()),)),)),)),
-                              UOp(Ops.CONST, dtypes.float, arg=9.203465015161783e-05, src=(
-                                x36:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 28, 28, 32, 1), strides=(0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),
-                          UOp(Ops.CONST, dtypes.float, arg=33.812857328652136, src=(
-                            x36,)),)),
-                        UOp(Ops.CONST, dtypes.float, arg=0.4999999, src=(
-                          x36,)),)),
-                      UOp(Ops.CONST, dtypes.float, arg=136.0, src=(
-                        x36,)),)),)),
-                  UOp(Ops.CONST, dtypes.int, arg=0, src=(
-                    x36,)),)),
-                x41:=UOp(Ops.CONST, dtypes.int, arg=-1, src=(
-                  x36,)),)),
-              UOp(Ops.CONST, dtypes.int, arg=-256, src=(
-                x36,)),)),
-            x41,)),)),)),))""")
-    opts = [Opt(op=OptOps.UNROLL, axis=0, arg=8), Opt(op=OptOps.UPCAST, axis=1, arg=32), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
-    with Context(DEVECTORIZE=0, QUANTIZE=1):
-      prg = get_program(ast, opts=opts)
-
-    new_src = """
-typedef int int32 __attribute__((aligned(128),vector_size(128)));
-typedef signed char signed_char128 __attribute__((aligned(128),vector_size(128)));
-typedef unsigned char unsigned_char8 __attribute__((aligned(8),vector_size(8)));
-typedef unsigned char unsigned_char4 __attribute__((aligned(4),vector_size(4)));
-typedef unsigned char unsigned_char128 __attribute__((aligned(128),vector_size(128)));
-__attribute__((noinline)) void r_196_32_4_24_8(unsigned char* restrict __attribute__((align_value(128))) data0, unsigned char* restrict __attribute__((align_value(128))) data1, signed char* restrict __attribute__((align_value(
-128))) data2, int* restrict __attribute__((align_value(128))) data3) {
-  int32 cast0 = (int32){0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0};
-  int32 val0 = *((int32*)((data3+0)));
-  for (int ridx0 = 0; ridx0 < 196; ridx0++) {
-    int32 acc0 = cast0;
-    int32 acc1 = cast0;
-    int32 acc2 = cast0;
-    int32 acc3 = cast0;
-    __builtin_HEXAGON_Y2_dcfetch(data1+ridx0*768);
-    __builtin_HEXAGON_Y2_dcfetch(data1+ridx0*768+192);
-    __builtin_HEXAGON_Y2_dcfetch(data1+ridx0*768+384);
-    __builtin_HEXAGON_Y2_dcfetch(data1+ridx0*768+576);
-    for (int ridx1 = 0; ridx1 < 24; ridx1++) {
-      signed_char128 val1 = *((signed_char128*)((data2+(ridx1<<8))));
-      signed_char128 val2 = *((signed_char128*)((data2+((1+(ridx1<<1))<<7))));
-
-      int alu0 = ((ridx0*768)+(ridx1<<3));
-
-      unsigned_char8 val3 = *((unsigned_char8*)((data1+alu0)));
-      __builtin_HEXAGON_Y2_dcfetch(((data1+alu0)+16));
-      unsigned_char8 val4 = *((unsigned_char8*)((data1+(alu0+192))));
-      __builtin_HEXAGON_Y2_dcfetch(((data1+(alu0+192))+16));
-      unsigned_char8 val5 = *((unsigned_char8*)((data1+(alu0+384))));
-      __builtin_HEXAGON_Y2_dcfetch(((data1+(alu0+384))+16));
-      unsigned_char8 val6 = *((unsigned_char8*)((data1+(alu0+576))));
-      __builtin_HEXAGON_Y2_dcfetch(((data1+(alu0+576))+16));
-
-      unsigned_char4 alu5 = __builtin_shufflevector(val3, val3,    0,   1,   2,   3);
-      unsigned_char4 alu6 = __builtin_shufflevector(val4, val4,    0,   1,   2,   3);
-      unsigned_char4 alu7 = __builtin_shufflevector(val5, val5,    0,   1,   2,   3);
-      unsigned_char4 alu8 = __builtin_shufflevector(val6, val6,    0,   1,   2,   3);
-      acc0 = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc0, val1, (*((unsigned int*)&alu5)));
-      acc1 = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc1, val1, (*((unsigned int*)&alu6)));
-      acc2 = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc2, val1, (*((unsigned int*)&alu7)));
-      acc3 = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc3, val1, (*((unsigned int*)&alu8)));
-
-      unsigned_char4 alu9 = __builtin_shufflevector(val3, val3,    4,   5,   6,   7);
-      unsigned_char4 alu10 = __builtin_shufflevector(val4, val4,    4,   5,   6,   7);
-      unsigned_char4 alu11 = __builtin_shufflevector(val5, val5,    4,   5,   6,   7);
-      unsigned_char4 alu12 = __builtin_shufflevector(val6, val6,    4,   5,   6,   7);
-      acc0 = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc0, val2, (*((unsigned int*)&alu9)));
-      acc1 = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc1, val2, (*((unsigned int*)&alu10)));
-      acc2 = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc2, val2, (*((unsigned int*)&alu11)));
-      acc3 = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc3, val2, (*((unsigned int*)&alu12)));
-    }
-    unsigned_char128 alu18 = __builtin_HEXAGON_V6_vpackhub_sat_128B(__builtin_HEXAGON_V6_vpackwh_sat_128B((((((acc3+val0)*203)+32767)/65536)+136), (((((acc2+val0)*203)+32767)/65536)+136)), __builtin_HEXAGON_V6_vpackwh_sat_128B((((((acc1+val0)*203)+32767)/65536)+136), (((((acc0+val0)*203)+32767)/65536)+136)));
-    *((unsigned_char128*)((data0+(ridx0<<7)))) = alu18;
-  }
-}
-"""
-    prg = replace(prg, src=new_src+prg.src.split("/* DSP boilerplate */ ")[1])
-    rt = CompiledRunner(prg)
-    #Device.default.compiler.disassemble(rt.lib)
-    ei = ExecItem(rt, bufs_from_lin(Kernel(ast)))
-    tm = ei.run(wait=True)
-    print(f"final time {tm*1e6:.2f} us")
-
 if __name__ == "__main__":
   unittest.main()

test/test_randomness.py

@@ -323,9 +323,9 @@ class TestRandomness(unittest.TestCase):
         torch_res = torch_res.unsqueeze(0)
       for i in range(torch_res.shape[0]):
         self.assertTrue(equal_distribution(lambda *_: tiny_res[i], lambda _: torch_res[i]))
-    _check_with_torch(w=[0.231, 0., 1., 0.5], num_samples=2000, replacement=True)
-    _check_with_torch(w=[[0.2, 0.8]], num_samples=2000, replacement=True)  # 2D but only 1 row
-    _check_with_torch(w=[[0.453, 0., 1., 0.81], [0.1, 0.8, 0., 0.1]], num_samples=2000, replacement=True)
+    _check_with_torch(w=[0.231, 0., 1., 0.5], num_samples=300, replacement=True)
+    _check_with_torch(w=[[0.2, 0.8]], num_samples=300, replacement=True)  # 2D but only 1 row
+    _check_with_torch(w=[[0.453, 0., 1., 0.81], [0.1, 0.8, 0., 0.1]], num_samples=300, replacement=True)
     # no-replacement isn't supported, unless taking only one sample
     w = [0.1, 0.9]
     self.assertRaises(AssertionError, lambda: Tensor(w).multinomial(100, replacement=False))

test/test_rangeify.py

@@ -1,6 +1,7 @@
 import unittest
 from tinygrad import Tensor
 from tinygrad.helpers import RANGEIFY, Context, GlobalCounters
+from tinygrad.uop.ops import UOp
 
 N = 256
 
@@ -11,6 +12,26 @@ class TestRangeify(unittest.TestCase):
     ba = A.expand(N, N)
     ((ba+1).sum(axis=1) + (ba+2).sum(axis=0)).realize()
 
+  def test_partial_contig(self):
+    A = Tensor.empty(64, 64, 64)
+    ret = A.sum(axis=2).contiguous(arg=(1,)).sum(axis=1)
+    ret.realize()
+
+  def test_double_gemm_real(self):
+    def go():
+      with Context(DEBUG=0):
+        Tensor.manual_seed(1337)
+        A,B,C = [Tensor.randn(N, N) for _ in range(3)]
+        Tensor.realize(A, B, C)
+      GlobalCounters.reset()
+      return (A@B@C).realize()
+    rng = go()
+    with Context(RANGEIFY=0, DEBUG=2):
+      ref = go()
+      mse = ((rng-ref)**2).sum().item()
+    print(f"mse: {mse}")
+    self.assertLessEqual(mse, 1e-2)
+
   def test_double_gemm(self):
     A = Tensor.empty(N, N)
     B = Tensor.empty(N, N)
@@ -99,7 +120,7 @@ class TestRangeify(unittest.TestCase):
     BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
 
     # bigger
-    #BS, HEADS, SEQLEN, EMB = 4, 16, 128, 64
+    #BS, HEADS, SEQLEN, EMB = 4, 32, 1024, 64
 
     # llama 8B
     #BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128
@@ -121,9 +142,6 @@ class TestRangeify(unittest.TestCase):
     print(f"mse: {mse}")
     self.assertLessEqual(mse, 1e-6)
 
-from tinygrad import dtypes
-from tinygrad.uop.ops import UOp
-
 # contiguous + reduce can support ranges?
 
 @unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY")
@@ -132,7 +150,7 @@ class TestOuterworld(unittest.TestCase):
     t = Tensor.rand(10, 10).realize()
 
     # passthrough ranges
-    a = UOp.range(dtypes.int, 10, -1)
+    a = UOp.range(10, -1)
     sel = t[a]
     cpy = sel.contiguous(a).realize()
 
@@ -142,7 +160,7 @@ class TestOuterworld(unittest.TestCase):
     t = Tensor.rand(10, 10).realize()
 
     # passthrough ranges
-    a = UOp.range(dtypes.int, 10, -1)
+    a = UOp.range(10, -1)
     sel = t[9-a]
     cpy = sel.contiguous(a).realize()
 
@@ -154,7 +172,7 @@ class TestOuterworld(unittest.TestCase):
     x = Tensor.ones(3, 10, 2).contiguous()
 
     # vmap across axis 0
-    a = UOp.range(dtypes.int, 3, -1)
+    a = UOp.range(3, -1)
     out = f(x[a])
     out = out.contiguous(a)
 
@@ -162,17 +180,39 @@ class TestOuterworld(unittest.TestCase):
     out.realize()
     print(out.numpy())
 
+  @unittest.skip("opts don't work")
   def test_triple_gemm(self):
     x = Tensor.rand(1, 16).realize()
     W = Tensor.rand(3, 16, 16).realize()
 
     manual = (x @ W[0] @ W[1] @ W[2]).contiguous().realize()
 
-    a = UOp.range(dtypes.int, 3, -1)
+    a = UOp.range(3, -1)
     x = x.assign(x @ W[a])
     out = x.contiguous(a)[-1].contiguous().realize()
 
     self.assertTrue((manual==out).all().item())
 
+  def test_setitem_pyrange(self):
+    with Context(DEBUG=0):
+      t = Tensor.rand(10).realize()
+      o = Tensor.empty(10)
+    GlobalCounters.reset()
+    for i in range(10):
+      o[i] = t[i]
+    o.realize()
+    self.assertTrue((t==o).all().item())
+
+  @unittest.skip("TODO: fix this")
+  def test_setitem(self):
+    with Context(DEBUG=0):
+      t = Tensor.rand(10).realize()
+      o = Tensor.empty(10)
+    GlobalCounters.reset()
+    i = UOp.range(10, -1)
+    o[i] = t[i]
+    o.contiguous(i).realize()
+    self.assertTrue((t==o).all().item())
+
 if __name__ == '__main__':
   unittest.main()

test/test_renderer_failures.py

@@ -25,7 +25,8 @@ def _test_uop_result(inputs:List[Tensor], stores:List[UOp], local_size=None):
       initial_value=np.zeros(sz, dtype=_to_np_dtype(dtype)).data) for u in uops if u.op is Ops.STORE]
   inbufs = [cast(UOp,x.uop).base.buffer for x in inputs]
   src = Device[Device.DEFAULT].renderer.render(uops)
-  ei = CompiledRunner(ProgramSpec("test", src, Device.DEFAULT, uops[-1], uops=uops, local_size=local_size))
+  ei = CompiledRunner(ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test",
+                                  src, Device.DEFAULT, uops[-1], uops=uops, local_size=local_size))
   ei.exec(outbufs+inbufs)
   return [np.frombuffer(x.as_buffer(), _to_np_dtype(x.dtype)) for x in outbufs]
 
@@ -45,7 +46,7 @@ class TestRendererFailures(unittest.TestCase):
   @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
   def test_gated_store_with_alu(self):
     a = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
-    gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (), ('lidx0', 4))).ne(0)
+    gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0, gate_alu), UOp.const(dtypes.int, 1)))
     sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
     uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
@@ -55,8 +56,8 @@ class TestRendererFailures(unittest.TestCase):
   @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
   def test_gated_store_with_alu_2d(self):
     a = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
-    gate_alu_0 = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (), ('lidx0', 4))).ne(0)
-    gate_alu_1 = (lidx1:=UOp(Ops.SPECIAL, dtypes.int, (), ('lidx1', 2))).ne(0)
+    gate_alu_0 = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
+    gate_alu_1 = (lidx1:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 2),), 'lidx1')).ne(0)
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0+lidx1*4, gate_alu_0&gate_alu_1), UOp.const(dtypes.int, 1)))
     sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
     uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
@@ -100,7 +101,7 @@ class TestPTXFailures(unittest.TestCase):
   @unittest.skip("INDEX can only have a gate ALU parent, not an IF")
   def test_gated_store_with_if(self):
     a = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
-    gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (), ('lidx0', 4))).ne(0)
+    gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
     val = UOp.const(dtypes.int, 1)
     if_uop = UOp(Ops.IF, dtypes.void, (gate_alu,))
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0, if_uop), val))

test/test_schedule.py

@@ -1050,6 +1050,14 @@ class TestSchedule(unittest.TestCase):
       compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
       np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)
 
+    with Context(FUSE_ATTENTION=1):
+      out = Tensor.scaled_dot_product_attention(q,k,v)
+      run_schedule(check_schedule(out, 1))
+      if getenv("CHECK", 1):
+        import torch
+        compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
+        np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)
+
   def test_ugly_reduceop_pairing(self):
     Tensor.manual_seed(0)
     a = Tensor.randn(4, 32).realize()

test/test_search.py

@@ -1,146 +0,0 @@
-import unittest
-
-from tinygrad.codegen.opt.kernel import Opt, OptOps, Kernel
-from tinygrad.uop.ops import UOp, Ops
-from tinygrad.codegen.opt.search import bufs_from_lin, actions, beam_search
-from tinygrad.device import Device
-from tinygrad.tensor import Tensor
-from tinygrad.dtype import dtypes
-from tinygrad.helpers import Context, GlobalCounters
-from tinygrad.engine.realize import capturing
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
-from extra.optimization.helpers import time_linearizer
-
-class TestBEAM(unittest.TestCase):
-  def test_dynamic_beam(self):
-    # TODO: make this infra globally usable
-    class Capture:
-      def __init__(self): self.captured = []
-      def add(self, x): self.captured.append(x)
-
-    capturing.append(Capture())
-    kernel_count = GlobalCounters.kernel_count
-    with Context(BEAM=1): Tensor.zeros(16).contiguous().realize()
-    assert GlobalCounters.kernel_count == kernel_count + 1
-    k_beam_1 = capturing[0].captured
-    capturing.clear()
-
-    capturing.append(Capture())
-    kernel_count = GlobalCounters.kernel_count
-    with Context(BEAM=0): Tensor.zeros(16).contiguous().realize()
-    assert GlobalCounters.kernel_count == kernel_count + 1
-    k_beam_0 = capturing[0].captured
-    capturing.clear()
-    self.assertNotEqual(k_beam_0[-1].prg.p.src, k_beam_1[-1].prg.p.src)
-
-  def test_get_kernel_actions_dedup(self):
-    from test.test_linearizer import helper_realized_ast
-    from tinygrad.codegen.opt.search import get_kernel_actions
-    a = Tensor.empty(4, 3)
-    b = Tensor.empty(3)
-    realized_ast, _ = helper_realized_ast(a @ b)
-    candidates = [
-      Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UPCAST, axis=0, arg=4),
-      Opt(op=OptOps.LOCAL, axis=0, arg=0), Opt(op=OptOps.LOCAL, axis=0, arg=4),
-      Opt(op=OptOps.UNROLL, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=3),
-      Opt(op=OptOps.GROUP, axis=0, arg=0), Opt(op=OptOps.GROUP, axis=0, arg=3),
-      Opt(op=OptOps.GROUPTOP, axis=0, arg=0), Opt(op=OptOps.GROUPTOP, axis=0, arg=3),
-    ]
-    lins = get_kernel_actions(Kernel(realized_ast), include_0=False, candidates=candidates).values()
-
-    # ensure amt=0 are not duplicated
-    assert all(len(x.applied_opts) == 1 for x in lins)
-    kernel_actions = [x.applied_opts[0] for x in lins]
-    assert Opt(OptOps.UPCAST, axis=0, arg=4) not in kernel_actions, "did not de-dup UPCAST"
-    assert Opt(OptOps.LOCAL, axis=0, arg=4) not in kernel_actions, "did not de-dup LOCAL"
-    assert Opt(OptOps.UNROLL, axis=0, arg=3) not in kernel_actions, "did not de-dup UNROLL"
-    assert Opt(OptOps.GROUP, axis=0, arg=3) not in kernel_actions, "did not de-dup GROUP"
-    assert Opt(OptOps.GROUPTOP, axis=0, arg=3) not in kernel_actions, "did not de-dup GROUPTOP"
-
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
-  def test_search_over_shape(self):
-    from test.test_linearizer import helper_realized_ast
-    from tinygrad.codegen.opt.search import get_kernel_actions
-
-    dtype_pairs = [(tc.dtype_in, tc.dtype_out) for tc in Device[Device.DEFAULT].renderer.tensor_cores]
-    multi_shape_dtype_pairs = [dts for dts in dtype_pairs if dtype_pairs.count(dts) > 1]
-
-    if len(multi_shape_dtype_pairs) == 0: raise unittest.SkipTest("only one tc available per dtype pair to search over")
-
-    for (dtype_in, dtype_out) in multi_shape_dtype_pairs:
-      a = Tensor.rand(16, 16, dtype=dtype_in)
-      b = Tensor.rand(16, 16, dtype=dtype_in)
-      realized_ast, _ = helper_realized_ast(a.matmul(b, dtype=dtype_out))
-
-      lins = get_kernel_actions(Kernel(realized_ast)).values()
-      assert len(set(lin.tensor_core.dims for lin in lins if lin.tensor_core is not None)) > 1
-
-  def test_get_kernel_actions_preserves_actions_state(self):
-    from test.test_linearizer import helper_realized_ast
-    from tinygrad.codegen.opt.search import get_kernel_actions
-    a = Tensor.rand(16, 16)
-    b = Tensor.rand(16, 16)
-    realized_ast, _ = helper_realized_ast(a @ b)
-    actions_before = actions.copy()
-    get_kernel_actions(Kernel(realized_ast))
-    actions_after = actions.copy()
-    assert actions_after == actions_before, "actions state was not preserved"
-
-  @unittest.skip("invalid reduce now")
-  def test_filter_global_buffer(self):
-    # taken from https://github.com/tinygrad/tinygrad/issues/4612
-    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.VIEW, dtypes.float.ptr(256), arg=ShapeTracker(views=(View(shape=(1, 1, 256), strides=(0, 0, 1), offset=0, mask=None, contiguous=True),)), src=( # noqa: E501
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(256), arg=0, src=()),)),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.MAX, (1,)), src=(
-          UOp(Ops.MUL, dtypes.float, arg=None, src=(
-            UOp(Ops.ADD, dtypes.float, arg=None, src=(
-              UOp(Ops.ADD, dtypes.float, arg=None, src=(
-                UOp(Ops.ADD, dtypes.float, arg=None, src=(
-                  UOp(Ops.ADD, dtypes.float, arg=None, src=(
-                    UOp(Ops.ADD, dtypes.float, arg=None, src=(
-                      UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                        UOp(Ops.VIEW, dtypes.float.ptr(64128), arg=ShapeTracker(views=(View(shape=(384768,), strides=(1,), offset=0, mask=((0, 64128),), contiguous=False), View(shape=(1, 501, 256), strides=(0, 1, 501), offset=256512, mask=None, contiguous=False))), src=( # noqa: E501
-                          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64128), arg=1, src=()),)),)),
-                      UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                        UOp(Ops.VIEW, dtypes.float.ptr(64128), arg=ShapeTracker(views=(View(shape=(384768,), strides=(1,), offset=-64128, mask=((64128, 128256),), contiguous=False), View(shape=(1, 501, 256), strides=(0, 1, 501), offset=256512, mask=None, contiguous=False))), src=( # noqa: E501
-                          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64128), arg=2, src=()),)),)),)),
-                    UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                      UOp(Ops.VIEW, dtypes.float.ptr(64128), arg=ShapeTracker(views=(View(shape=(384768,), strides=(1,), offset=-128256, mask=((128256, 192384),), contiguous=False), View(shape=(1, 501, 256), strides=(0, 1, 501), offset=256512, mask=None, contiguous=False))), src=( # noqa: E501
-                        UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64128), arg=3, src=()),)),)),)),
-                  UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                    UOp(Ops.VIEW, dtypes.float.ptr(64128), arg=ShapeTracker(views=(View(shape=(384768,), strides=(1,), offset=-192384, mask=((192384, 256512),), contiguous=False), View(shape=(1, 501, 256), strides=(0, 1, 501), offset=256512, mask=None, contiguous=False))), src=( # noqa: E501
-                      UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64128), arg=4, src=()),)),)),)),
-                UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                  UOp(Ops.VIEW, dtypes.float.ptr(64128), arg=ShapeTracker(views=(View(shape=(384768,), strides=(1,), offset=-256512, mask=((256512, 320640),), contiguous=False), View(shape=(1, 501, 256), strides=(0, 1, 501), offset=256512, mask=None, contiguous=False))), src=( # noqa: E501
-                    UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64128), arg=5, src=()),)),)),)),
-              UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                UOp(Ops.VIEW, dtypes.float.ptr(64128), arg=ShapeTracker(views=(View(shape=(384768,), strides=(1,), offset=-320640, mask=((320640, 384768),), contiguous=False), View(shape=(1, 501, 256), strides=(0, 1, 501), offset=256512, mask=None, contiguous=False))), src=( # noqa: E501
-                  UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64128), arg=6, src=()),)),)),)),
-            UOp(Ops.CONST, dtypes.float, arg=1.4285714285714286, src=(
-              UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 501, 256), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),)) # noqa: E501
-    lin = Kernel(ast)
-
-    bufs = bufs_from_lin(lin)
-    best_lin = beam_search(lin, bufs, 2)
-    assert best_lin
-    # need disable_cache to trigger.
-    tm = time_linearizer(best_lin, bufs, allow_test_size=False, cnt=2, disable_cache=True)
-    assert tm
-
-  def test_beam_unnamed_kernels(self):
-    from test.test_linearizer import push_views
-    a = Tensor.rand(100)
-    b = Tensor.rand(100)
-    si = (a+b).schedule()[-1]
-    lin = Kernel(push_views(si.ast))
-    bufs = bufs_from_lin(lin)
-    # TODO: beam should have better instrumentation so we don't have to check this indirect thing
-    kcount = len(Kernel.kernel_cnt)
-    beam_search(lin, bufs, 3, disable_cache=True)
-    self.assertEqual(kcount, len(Kernel.kernel_cnt))
-
-if __name__ == '__main__':
-  unittest.main()

test/test_symbolic_jit.py

@@ -2,50 +2,41 @@ import unittest
 
 from test.helpers import assert_jit_cache_len
 from tinygrad import Variable, Tensor, TinyJit
-from tinygrad.helpers import Context
 import numpy as np
 
 class TestSymbolicJit(unittest.TestCase):
-  def setUp(self):
-    # A lot of these test are out of bounds, so we ignore the bounds check
-    self.context = Context(IGNORE_OOB=1)
-    self.context.__enter__()
-
-  def tearDown(self):
-    self.context.__exit__(None, None, None)
-
   def test_plus1(self):
     def f(a): return (a+1).realize()
     jf = TinyJit(f)
+    a = Tensor.rand(3, 10)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      a = Tensor.rand(3, i)
-      symbolic = jf(a.reshape(3, vi)).reshape(3, i).numpy()
-      expected = f(a).numpy()
+      symbolic = jf(a[:, :vi]).reshape(3, i).numpy()
+      expected = f(a[:, :i]).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
   def test_add(self):
     def f(a, b): return (a+b).realize()
     jf = TinyJit(f)
+    a = Tensor.rand(3, 10)
+    b = Tensor.rand(3, 10)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      a = Tensor.rand(3, i)
-      b = Tensor.rand(3, i)
-      symbolic = jf(a.reshape(3, vi), b.reshape(3, vi)).reshape(3, i).numpy()
-      expected = f(a, b).numpy()
+      symbolic = jf(a[:, :vi], b[:, :vi]).reshape(3, i).numpy()
+      expected = f(a[:, :i], b[:, :i]).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
   def test_matmul(self):
     def f(a, b): return (a@b).realize()
     jf = TinyJit(f)
+    a = Tensor.rand(3, 10)
+    b = Tensor.rand(10, 5)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      a = Tensor.rand(3, i)
-      b = Tensor.rand(i, 5)
-      symbolic = jf(a.reshape(3, vi), b.reshape(vi, 5)).numpy()
-      expected = f(a, b).numpy()
+      symbolic = jf(a[:, :vi], b[:vi, :]).numpy()
+      expected = f(a[:, :i], b[:i, :]).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
@@ -55,119 +46,119 @@ class TestSymbolicJit(unittest.TestCase):
       s = (s+s).realize() # this one does not have symbols in input
       return s
     jf = TinyJit(f)
+    a = Tensor.rand(3, 10)
+    b = Tensor.rand(10, 5)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      a = Tensor.rand(3, i)
-      b = Tensor.rand(i, 5)
-      symbolic = jf(a.reshape(3, vi), b.reshape(vi, 5)).numpy()
-      expected = f(a, b).numpy()
+      symbolic = jf(a[:, :vi], b[:vi, :]).numpy()
+      expected = f(a[:, :i], b[:i, :]).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 2)
 
   def test_attention(self):
     def f(q, k, v): return Tensor.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).realize()
     jf = TinyJit(f)
+    q = Tensor.rand(2, 1, 4, 8)
+    k = Tensor.rand(2, 10, 4, 8)
+    v = Tensor.rand(2, 10, 4, 8)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      q = Tensor.rand(2, 1, 4, 8)
-      k = Tensor.rand(2, i, 4, 8)
-      v = Tensor.rand(2, i, 4, 8)
-      symbolic = jf(q, k.reshape(2, vi, 4, 8), v.reshape(2, vi, 4, 8)).reshape(2, 4, 1, 8).numpy()
-      expected = f(q, k, v).numpy()
+      symbolic = jf(q, k[:, :vi], v[:, :vi]).reshape(2, 4, 1, 8).numpy()
+      expected = f(q, k[:, :i], v[:, :i]).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 5)
 
   def test_cat_dim0(self):
     def f(a, b): return a.cat(b, dim=0).realize()
     jf = TinyJit(f)
+    a = Tensor.rand(10, 3)
+    b = Tensor.rand(2, 3)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      a = Tensor.rand(i, 3)
-      b = Tensor.rand(2, 3)
-      symbolic = jf(a.reshape(vi, 3), b).reshape(i+2, 3).numpy()
-      expected = f(a, b).numpy()
+      symbolic = jf(a[:vi], b).reshape(i+2, 3).numpy()
+      expected = f(a[:i], b).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
   def test_cat_dim1(self):
     def f(a, b): return a.cat(b, dim=1).realize()
     jf = TinyJit(f)
+    a = Tensor.rand(3, 10)
+    b = Tensor.rand(3, 2)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      a = Tensor.rand(3, i)
-      b = Tensor.rand(3, 2)
-      symbolic = jf(a.reshape(3, vi), b).reshape(3, i+2).numpy()
-      expected = f(a, b).numpy()
+      symbolic = jf(a[:, :vi], b).reshape(3, i+2).numpy()
+      expected = f(a[:, :i], b).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
   def test_cat_dim0_two_vars(self):
     def f(a, b): return a.cat(b, dim=0).realize()
     jf = TinyJit(f)
-    for i in range(1, 5):
-      for j in range(1, 5):
+    a = Tensor.rand(10, 3)
+    b = Tensor.rand(10, 3)
+    for i in range(2, 5):
+      for j in range(2, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        a = Tensor.rand(i, 3)
-        b = Tensor.rand(j, 3)
-        symbolic = jf(a.reshape(vi, 3), b.reshape(vj, 3)).reshape(i+j, 3).numpy()
-        expected = f(a, b).numpy()
+        symbolic = jf(a[:vi], b[:vj]).reshape(i+j, 3).numpy()
+        expected = f(a[:i], b[:j]).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
   def test_cat_dim1_two_vars(self):
     def f(a, b): return a.cat(b, dim=1).realize()
     jf = TinyJit(f)
-    for i in range(1, 5):
-      for j in range(1, 5):
+    a = Tensor.rand(3, 10)
+    b = Tensor.rand(3, 10)
+    for i in range(2, 5):
+      for j in range(2, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        a = Tensor.rand(3, i)
-        b = Tensor.rand(3, j)
-        symbolic = jf(a.reshape(3, vi), b.reshape(3, vj)).reshape(3, i+j).numpy()
-        expected = f(a, b).numpy()
+        symbolic = jf(a[:, :vi], b[:, :vj]).reshape(3, i+j).numpy()
+        expected = f(a[:, :i], b[:, :j]).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
   def test_two_vars_plus1_ij(self):
     def f(a, b): return (a@b+1).realize()
     jf = TinyJit(f)
-    for i in range(1, 5):
-      for j in range(1, 5):
+    a = Tensor.rand(10, 3)
+    b = Tensor.rand(3, 10)
+    for i in range(2, 5):
+      for j in range(2, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        a = Tensor.rand(i, 3)
-        b = Tensor.rand(3, j)
-        symbolic = jf(a.reshape(vi, 3), b.reshape(3, vj)).reshape(i, j).numpy()
-        expected = f(a, b).numpy()
+        symbolic = jf(a[:vi, :], b[:, :vj]).reshape(i, j).numpy()
+        expected = f(a[:i, :], b[:, :j]).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
   def test_two_vars_plus1_ji(self):
     def f(a, b): return (a@b+1).realize()
     jf = TinyJit(f)
-    for i in range(1, 5):
-      for j in range(1, 5):
+    a = Tensor.rand(10, 3)
+    b = Tensor.rand(3, 10)
+    for i in range(2, 5):
+      for j in range(2, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        a = Tensor.rand(j, 3)
-        b = Tensor.rand(3, i)
-        symbolic = jf(a.reshape(vj, 3), b.reshape(3, vi)).reshape(j, i).numpy()
-        expected = f(a, b).numpy()
+        symbolic = jf(a[:vj, :], b[:, :vi]).reshape(j, i).numpy()
+        expected = f(a[:j, :], b[:, :i]).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
   def test_jit_symbolic_shape_mismatch(self):
     @TinyJit
     def add(a, b): return (a+b).realize()
+    a = Tensor.rand(3, 10)
+    b = Tensor.rand(3, 10)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      a = Tensor.rand(3, i).reshape(3, vi)
-      b = Tensor.rand(3, i).reshape(3, vi)
-      add(a, b)
+      add(a[:, :vi], b[:, :vi])
     vi2 = Variable("i", 1, 10).bind(7)
-    a = Tensor.rand(3, 7).reshape(3, vi2)
-    bad = Tensor.rand(4, 7).reshape(4, vi2)
+    a = Tensor.rand(3, 7)[:, :vi2]
+    bad = Tensor.rand(4, 7)[:, :vi2]
     with self.assertRaises(AssertionError):
       add(a, bad)
 
@@ -175,9 +166,9 @@ class TestSymbolicJit(unittest.TestCase):
     # shrink is a movement, so we pair it with a simple function to test the JIT interaction
     def f(a): return (a+1).realize()
     jf = TinyJit(f)
+    a = Tensor.rand(7, 11)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      a = Tensor.rand(7, 11)
       symbolic = a.shrink(((3,5),(vi,vi+2)))
       symbolic = jf(symbolic).numpy()
       expected = f(a.shrink(((3,5),(i,i+2)))).numpy()
@@ -188,9 +179,9 @@ class TestSymbolicJit(unittest.TestCase):
     # slice is a movement, so we pair it with a simple function to test the JIT interaction
     def f(a): return (a+1).realize()
     jf = TinyJit(f)
+    a = Tensor.rand(7, 11)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      a = Tensor.rand(7, 11)
       symbolic = a[3:5, vi:vi+2]
       symbolic = jf(symbolic).numpy()
       expected = f(a[3:5, i:i+2]).numpy()
@@ -212,11 +203,11 @@ class TestSymbolicJit(unittest.TestCase):
   def test_ones_sum(self):
     def f(a): return a.sum().realize()
     jf = TinyJit(f)
+    t = Tensor.ones(10)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      t = Tensor.ones(i)
-      symbolic = jf(t.reshape(vi)).item()
-      expected = f(t).item()
+      symbolic = jf(t[:vi]).item()
+      expected = f(t[:i]).item()
       np.testing.assert_equal(symbolic, expected)
 
   def test_mean(self):
@@ -226,22 +217,22 @@ class TestSymbolicJit(unittest.TestCase):
     jf = TinyJit(f)
     jf0 = TinyJit(f0)
     jf1 = TinyJit(f1)
+    a = Tensor.rand(10, 3)
+    b = Tensor.rand(10, 3)
+    c = Tensor.rand(10, 3)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      # aixs = None
-      a = Tensor.rand(i, 3)
-      symbolic = jf(a.reshape(vi, 3)).numpy()
-      expected = a.mean().numpy()
+      # axis = None
+      symbolic = jf(a[:vi]).numpy()
+      expected = a[:i].mean().numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-      # aixs = 0
-      a = Tensor.rand(i, 3)
-      symbolic = jf0(a.reshape(vi, 3)).numpy()
-      expected = a.mean(0).numpy()
+      # axis = 0
+      symbolic = jf0(b[:vi]).numpy()
+      expected = b[:i].mean(0).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-      # aixs = 1
-      a = Tensor.rand(i, 3)
-      symbolic = jf1(a.reshape(vi, 3)).reshape(i).numpy()
-      expected = a.mean(1).numpy()
+      # axis = 1
+      symbolic = jf1(c[:vi]).reshape(i).numpy()
+      expected = c[:i].mean(1).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_mean_2d(self):
@@ -251,24 +242,24 @@ class TestSymbolicJit(unittest.TestCase):
     jf = TinyJit(f)
     jf0 = TinyJit(f0)
     jf1 = TinyJit(f1)
-    for i in range(1, 5):
-      for j in range(1, 5):
+    a = Tensor.rand(10, 10)
+    b = Tensor.rand(10, 10)
+    c = Tensor.rand(10, 10)
+    for i in range(2, 5):
+      for j in range(2, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        # aixs = None
-        a = Tensor.rand(i, j)
-        symbolic = jf(a.reshape(vi, vj)).numpy()
-        expected = a.mean().numpy()
+        # axis = None
+        symbolic = jf(a[:vi, :vj]).numpy()
+        expected = a[:i, :j].mean().numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-        # aixs = 0
-        a = Tensor.rand(i, j)
-        symbolic = jf0(a.reshape(vi, vj)).reshape(j).numpy()
-        expected = a.mean(0).numpy()
+        # axis = 0
+        symbolic = jf0(b[:vi, :vj]).reshape(j).numpy()
+        expected = b[:i, :j].mean(0).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-        # aixs = 1
-        a = Tensor.rand(i, j)
-        symbolic = jf1(a.reshape(vi, vj)).reshape(i).numpy()
-        expected = a.mean(1).numpy()
+        # axis = 1
+        symbolic = jf1(c[:vi, :vj]).reshape(i).numpy()
+        expected = c[:i, :j].mean(1).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_var(self):
@@ -278,22 +269,22 @@ class TestSymbolicJit(unittest.TestCase):
     jf = TinyJit(f)
     jf0 = TinyJit(f0)
     jf1 = TinyJit(f1)
+    a = Tensor.rand(10, 3)
+    b = Tensor.rand(10, 3)
+    c = Tensor.rand(10, 3)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      # aixs = None
-      a = Tensor.rand(i, 3)
-      symbolic = jf(a.reshape(vi, 3)).numpy()
-      expected = a.var().numpy()
+      # axis = None
+      symbolic = jf(a[:vi]).numpy()
+      expected = a[:i].var().numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-      # aixs = 0
-      a = Tensor.rand(i, 3)
-      symbolic = jf0(a.reshape(vi, 3)).numpy()
-      expected = a.var(0).numpy()
+      # axis = 0
+      symbolic = jf0(b[:vi]).numpy()
+      expected = b[:i].var(0).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-      # aixs = 1
-      a = Tensor.rand(i, 3)
-      symbolic = jf1(a.reshape(vi, 3)).reshape(i).numpy()
-      expected = a.var(1).numpy()
+      # axis = 1
+      symbolic = jf1(c[:vi]).reshape(i).numpy()
+      expected = c[:i].var(1).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_var_2d(self):
@@ -303,24 +294,24 @@ class TestSymbolicJit(unittest.TestCase):
     jf = TinyJit(f)
     jf0 = TinyJit(f0)
     jf1 = TinyJit(f1)
-    for i in range(1, 5):
-      for j in range(1, 5):
+    a = Tensor.rand(10, 10)
+    b = Tensor.rand(10, 10)
+    c = Tensor.rand(10, 10)
+    for i in range(2, 5):
+      for j in range(2, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        # aixs = None
-        a = Tensor.rand(i, j)
-        symbolic = jf(a.reshape(vi, vj)).numpy()
-        expected = a.var().numpy()
+        # axis = None
+        symbolic = jf(a[:vi, :vj]).numpy()
+        expected = a[:i, :j].var().numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-        # aixs = 0
-        a = Tensor.rand(i, j)
-        symbolic = jf0(a.reshape(vi, vj)).reshape(j).numpy()
-        expected = a.var(0).numpy()
+        # axis = 0
+        symbolic = jf0(b[:vi, :vj]).reshape(j).numpy()
+        expected = b[:i, :j].var(0).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-        # aixs = 1
-        a = Tensor.rand(i, j)
-        symbolic = jf1(a.reshape(vi, vj)).reshape(i).numpy()
-        expected = a.var(1).numpy()
+        # axis = 1
+        symbolic = jf1(c[:vi, :vj]).reshape(i).numpy()
+        expected = c[:i, :j].var(1).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
 if __name__ == '__main__':

test/test_symbolic_ops.py

@@ -1,62 +1,53 @@
 import unittest
-from tinygrad import Tensor, Variable
+from tinygrad import Tensor, Variable, GlobalCounters
 from tinygrad.shape.shapetracker import View
-from tinygrad.helpers import Context, GlobalCounters
 from tinygrad.uop.ops import sym_infer
 from tinygrad.dtype import dtypes
-from tinygrad.device import Device
+from tinygrad.device import is_dtype_supported
 from examples.gpt2 import Attention
 import numpy as np
 
 class TestSymbolicOps(unittest.TestCase):
-  def setUp(self):
-    # A lot of these test are out of bounds, so we ignore the bounds check
-    self.context = Context(IGNORE_OOB=1)
-    self.context.__enter__()
-
-  def tearDown(self):
-    self.context.__exit__(None, None, None)
-
   def test_plus1(self):
     def f(a): return (a+1).realize()
+    a = Tensor.rand(3, 10)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      a = Tensor.rand(3, i)
-      symbolic = f(a.reshape(3, vi)).reshape(3, i).numpy()
-      expected = f(a).numpy()
+      symbolic = f(a[:, :vi]).reshape(3, i).numpy()
+      expected = f(a[:, :i]).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_add(self):
     def f(a, b): return (a+b).realize()
+    a = Tensor.rand(3, 10)
+    b = Tensor.rand(3, 10)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      a = Tensor.rand(3, i)
-      b = Tensor.rand(3, i)
-      symbolic = f(a.reshape(3, vi), b.reshape(3, vi)).reshape(3, i).numpy()
-      expected = f(a, b).numpy()
+      symbolic = f(a[:, :vi], b[:, :vi]).reshape(3, i).numpy()
+      expected = f(a[:, :i], b[:, :i]).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_matmul(self):
     def f(a, b): return (a@b).realize()
+    a = Tensor.rand(3, 10)
+    b = Tensor.rand(10, 5)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      a = Tensor.rand(3, i)
-      b = Tensor.rand(i, 5)
-      symbolic = f(a.reshape(3, vi), b.reshape(vi, 5)).numpy()
-      expected = f(a, b).numpy()
+      symbolic = f(a[:, :vi], b[:vi, :]).numpy()
+      expected = f(a[:, :i], b[:i, :]).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_attention(self, dropout_p=0.0, imin=1, imax=5, use_symbolic=True):
     def f(q, k, v): return Tensor.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), dropout_p=dropout_p).realize()
+    q = Tensor.rand(2, 1, 4, 8)
+    k = Tensor.rand(2, 10, 4, 8)
+    v = Tensor.rand(2, 10, 4, 8)
     for i in range(imin, imax):
       vi = Variable("i", 1, 10).bind(i) if use_symbolic else i
-      q = Tensor.rand(2, 1, 4, 8)
-      k = Tensor.rand(2, i, 4, 8)
-      v = Tensor.rand(2, i, 4, 8)
       Tensor.realize(q, k, v)
       GlobalCounters.reset()
-      symbolic = f(q, k.reshape(2, vi, 4, 8), v.reshape(2, vi, 4, 8)).reshape(2, 4, 1, 8).numpy()
-      expected = f(q, k, v).numpy()
+      symbolic = f(q, k[:, :vi, :, :], v[:, :vi, :, :]).reshape(2, 4, 1, 8).numpy()
+      expected = f(q, k[:, :i, :, :], v[:, :i, :, :]).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_attention_cmp_symbolic(self):
@@ -90,73 +81,89 @@ class TestSymbolicOps(unittest.TestCase):
 
   def test_cat_dim0(self):
     def f(a, b): return a.cat(b, dim=0).realize()
+    a = Tensor.rand(10, 3)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      a = Tensor.rand(i, 3)
       b = Tensor.rand(2, 3)
-      symbolic = f(a.reshape(vi, 3), b).reshape(i+2, 3).numpy()
-      expected = f(a, b).numpy()
+      symbolic = f(a[:vi, :], b).reshape(i+2, 3).numpy()
+      expected = f(a[:i, :], b).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_cat_dim1(self):
     def f(a, b): return a.cat(b, dim=1).realize()
+    a = Tensor.rand(3, 10)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      a = Tensor.rand(3, i)
       b = Tensor.rand(3, 2)
-      symbolic = f(a.reshape(3, vi), b).reshape(3, i+2).numpy()
-      expected = f(a, b).numpy()
+      symbolic = f(a[:, :vi], b).reshape(3, i+2).numpy()
+      expected = f(a[:, :i], b).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_cat_dim0_two_vars(self):
     def f(a, b): return a.cat(b, dim=0).realize()
-    for i in range(1, 5):
-      for j in range(1, 5):
+    a = Tensor.rand(10, 3)
+    b = Tensor.rand(10, 3)
+    for i in range(2, 5):
+      for j in range(2, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        a = Tensor.rand(i, 3)
-        b = Tensor.rand(j, 3)
-        symbolic = f(a.reshape(vi, 3), b.reshape(vj, 3)).reshape(i+j, 3).numpy()
-        expected = f(a, b).numpy()
+        symbolic = f(a[:vi, :], b[:vj, :]).reshape(i+j, 3).numpy()
+        expected = f(a[:i, :], b[:j, :]).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_cat_dim1_two_vars(self):
     def f(a, b): return a.cat(b, dim=1).realize()
-    for i in range(1, 5):
-      for j in range(1, 5):
+    a = Tensor.rand(3, 10)
+    b = Tensor.rand(3, 10)
+    for i in range(2, 5):
+      for j in range(2, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        a = Tensor.rand(3, i)
-        b = Tensor.rand(3, j)
-        symbolic = f(a.reshape(3, vi), b.reshape(3, vj)).reshape(3, i+j).numpy()
-        expected = f(a, b).numpy()
+        symbolic = f(a[:, :vi], b[:, :vj]).reshape(3, i+j).numpy()
+        expected = f(a[:, :i], b[:, :j]).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_two_vars_plus1_ij(self):
     def f(a, b): return (a@b+1).realize()
-    for i in range(1, 5):
-      for j in range(1, 5):
+    a = Tensor.rand(10, 3)
+    b = Tensor.rand(3, 10)
+    for i in range(2, 5):
+      for j in range(2, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        a = Tensor.rand(i, 3)
-        b = Tensor.rand(3, j)
-        symbolic = f(a.reshape(vi, 3), b.reshape(3, vj)).reshape(i, j).numpy()
-        expected = f(a, b).numpy()
+        symbolic = f(a[:vi, :], b[:, :vj]).reshape(i, j).numpy()
+        expected = f(a[:i, :], b[:, :j]).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_two_vars_plus1_ji(self):
     # reverse the order of variables
     def f(a, b): return (a@b+1).realize()
-    for i in range(1, 5):
-      for j in range(1, 5):
+    a = Tensor.rand(10, 3)
+    b = Tensor.rand(3, 10)
+    for i in range(2, 5):
+      for j in range(2, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        a = Tensor.rand(j, 3)
-        b = Tensor.rand(3, i)
-        symbolic = f(a.reshape(vj, 3), b.reshape(3, vi)).reshape(j, i).numpy()
-        expected = f(a, b).numpy()
+        symbolic = f(a[:vj, :], b[:, :vi]).reshape(j, i).numpy()
+        expected = f(a[:j, :], b[:, :i]).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
+  def test_reshape_from_symbolic(self):
+    a = Tensor.rand(30)
+    for i in range(3, 5):
+      vi = Variable("i", 3, 10).bind(i)
+      symbolic = a[:vi*3].reshape((3, 3)).numpy()
+      # To match symbolic reshape (potential implicit shrink), we need a shrink
+      expected = a[:i*3].shrink(((0, 9),)).reshape((3, 3)).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_invalid_symbolic_reshape(self):
+    a = Tensor.rand(30)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      # Cannot reshape into symbolic from non-symbolic
+      with self.assertRaises(AssertionError): a.reshape((3, vi))
+
   def test_shrink(self):
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
@@ -176,11 +183,10 @@ class TestSymbolicOps(unittest.TestCase):
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_slice_no_start(self):
+    a = Tensor.rand(7, 11)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      a = Tensor.rand(7, 11)
-      symbolic = a[3:5, :vi:1].reshape(2,i)
-      symbolic = symbolic.numpy()
+      symbolic = a[3:5, :vi:1].reshape(2, i).numpy()
       expected = a[3:5, :i:1].numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
@@ -201,31 +207,31 @@ class TestSymbolicOps(unittest.TestCase):
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_ones_sum(self):
+    t = Tensor.ones(10)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      t = Tensor.ones(i)
-      symbolic = t.reshape(vi).sum().item()
-      expected = t.sum().item()
+      symbolic = t[:vi].sum().item()
+      expected = t[:i].sum().item()
       np.testing.assert_equal(symbolic, expected)
 
   def test_mean(self):
+    a = Tensor.rand(10, 3)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
       for axis in [None, 0, 1]:
-        a = Tensor.rand(i, 3)
-        expected = a.mean(axis).numpy()
-        symbolic = a.reshape(vi, 3).mean(axis).reshape(expected.shape).numpy()
+        expected = a[:i].mean(axis).numpy()
+        symbolic = a[:vi].mean(axis).reshape(expected.shape).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_mean_2d(self):
-    for i in range(1, 5):
-      for j in range(1, 5):
+    a = Tensor.rand(10, 10)
+    for i in range(2, 5):
+      for j in range(2, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
         for axis in [None, 0, 1]:
-          a = Tensor.rand(i, j)
-          expected = a.mean(axis).numpy()
-          symbolic = a.reshape(vi, vj).mean(axis).reshape(expected.shape).numpy()
+          expected = a[:i, :j].mean(axis).numpy()
+          symbolic = a[:vi, :vj].mean(axis).reshape(expected.shape).numpy()
           np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_var(self):
@@ -233,43 +239,43 @@ class TestSymbolicOps(unittest.TestCase):
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
       for axis in [None, 0, 1]:
-        expected = a[:i, :].var(axis).numpy()
-        symbolic = a[:vi, :].var(axis).reshape(expected.shape).numpy()
+        expected = a[:i].var(axis).numpy()
+        symbolic = a[:vi].var(axis).reshape(expected.shape).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_var_2d(self):
-    for i in range(1, 5):
-      for j in range(1, 5):
+    a = Tensor.rand(10, 10)
+    for i in range(2, 5):
+      for j in range(2, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
         for axis in [None, 0, 1]:
-          a = Tensor.rand(i, j)
-          expected = a.var(axis).numpy()
-          symbolic = a.reshape(vi, vj).var(axis).reshape(expected.shape).numpy()
+          expected = a[:i, :j].var(axis).numpy()
+          symbolic = a[:vi, :vj].var(axis).reshape(expected.shape).numpy()
           np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_bitcast_down(self):
+    a = Tensor.rand(10, 3)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      a = Tensor.rand(i, 3)
-      expected = a.bitcast(dtypes.uint8).numpy()
-      symbolic = a.reshape(vi, 3).bitcast(dtypes.uint8).reshape(expected.shape).numpy()
+      expected = a[:i].bitcast(dtypes.uint8).numpy()
+      symbolic = a[:vi].bitcast(dtypes.uint8).reshape(expected.shape).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=0)
 
-  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "no uint64")
+  @unittest.skipUnless(is_dtype_supported(dtypes.uint64), "no uint64")
   def test_bitcast_up(self):
+    a = Tensor.rand(10, 4)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      a = Tensor.rand(i, 4)
-      expected = a.bitcast(dtypes.uint64).numpy()
-      symbolic = a.reshape(vi, 4).bitcast(dtypes.uint64).reshape(expected.shape).numpy()
+      expected = a[:i].bitcast(dtypes.uint64).numpy()
+      symbolic = a[:vi].bitcast(dtypes.uint64).reshape(expected.shape).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=0)
 
   @unittest.expectedFailure
   def test_conv2d_ceildiv_edge_case(self):
     v = Variable('v', 11, 50_000)
     val = 39601
-    x = Tensor.randn(1, 22, 39601).reshape(1, 22, v.bind(val))
+    x = Tensor.randn(1, 22, 50_000)[:, :, :v.bind(val)]
     weight = Tensor.randn(256, 22, 12)
 
     result = x.conv2d(weight=weight, groups=1, stride=6, dilation=1, padding=(3, 3))

test/test_tensor_variable.py

@@ -1,7 +1,6 @@
 import unittest
 import numpy as np
 from tinygrad import Tensor, Variable
-from tinygrad.helpers import Context
 
 class TestTensorVariable(unittest.TestCase):
   def test_add_tvar(self):
@@ -23,43 +22,38 @@ class TestTensorVariable(unittest.TestCase):
     assert (Tensor(3) * (vv * 4)).item() == 24
 
   def test_symbolic_mean(self):
-    with Context(IGNORE_OOB=1):
-      vv = Variable("a", 1, 10).bind(2)
-      t = Tensor.ones(2, 2).contiguous().reshape(2, vv)
-      ret = t.mean().item()
-      assert ret == 1
+    vv = Variable("a", 1, 10).bind(2)
+    t = Tensor.ones(2, 10).contiguous()[:, :vv]
+    ret = t.mean().item()
+    assert ret == 1
 
   def test_symbolic_mean_2d(self):
-    with Context(IGNORE_OOB=1):
-      vv = Variable("a", 1, 10).bind(2)
-      vv2 = Variable("b", 1, 10).bind(2)
-      t = Tensor.ones(2, 2).contiguous().reshape(vv2, vv)
-      ret = t.mean().item()
-      assert ret == 1
+    vv = Variable("a", 1, 10).bind(2)
+    vv2 = Variable("b", 1, 10).bind(2)
+    t = Tensor.ones(10, 10).contiguous()[:vv2, :vv]
+    ret = t.mean().item()
+    assert ret == 1
 
   def test_symbolic_mean_2d_axis_1(self):
-    with Context(IGNORE_OOB=1):
-      vv = Variable("a", 1, 10).bind(2)
-      vv2 = Variable("b", 1, 10).bind(2)
-      t = Tensor.ones(2, 2).contiguous().reshape(vv2, vv)
-      ret = t.mean(axis=1).reshape(2, 1).numpy()
-      assert np.all(ret == 1)
+    vv = Variable("a", 1, 10).bind(2)
+    vv2 = Variable("b", 1, 10).bind(2)
+    t = Tensor.ones(10, 10).contiguous()[:vv2, :vv]
+    ret = t.mean(axis=1).reshape(2, 1).numpy()
+    assert np.all(ret == 1)
 
   def test_symbolic_mean_2d_add(self):
-    with Context(IGNORE_OOB=1):
-      add_term = Variable("c", 0, 10).bind(1)
-      vv = Variable("a", 1, 10).bind(1)
-      vv2 = Variable("b", 1, 10).bind(1)
-      t = Tensor.ones(2, 2).contiguous().reshape(vv2+add_term, vv+add_term)
-      ret = t.mean().item()
-      assert ret == 1
+    add_term = Variable("c", 0, 10).bind(1)
+    vv = Variable("a", 1, 10).bind(1)
+    vv2 = Variable("b", 1, 10).bind(1)
+    t = Tensor.ones(20, 20).contiguous()[:vv2+add_term, :vv+add_term]
+    ret = t.mean().item()
+    assert ret == 1
 
   def test_symbolic_var(self):
-    with Context(IGNORE_OOB=1):
-      vv = Variable("a", 1, 10).bind(2)
-      t = Tensor.ones(2, 2).contiguous().reshape(2, vv)
-      ret = t.var().item()
-      assert ret == 0
+    vv = Variable("a", 1, 10).bind(2)
+    t = Tensor.ones(2, 10).contiguous()[:, :vv]
+    ret = t.var().item()
+    assert ret == 0
 
   def test_symbolic_pad(self):
     vv = Variable("a", 1, 10).bind(2)
@@ -92,5 +86,15 @@ class TestTensorVariable(unittest.TestCase):
     ret = Tensor.arange(begin.bind(4), end.bind(7))
     self.assertListEqual(ret.reshape(3).tolist(), [4,5,6])
 
+  def test_variable_empty(self):
+    v = Variable("i", 1, 10)
+    # TODO: Tensor creation from unbound variable should assert
+    # with self.assertRaises(AssertionError): t = Tensor.empty(3, v)
+    vb = v.bind(3)
+    t = Tensor.empty(3, vb)
+    assert t.uop.base.buffer.size == 30
+    assert t.uop.st.shape == (3, vb)
+
+
 if __name__ == '__main__':
   unittest.main()

test/test_tiny.py

@@ -1,7 +1,7 @@
 # basic self-contained tests of the external functionality of tinygrad
 import unittest, random
 from tinygrad import Tensor, Context, Variable, TinyJit, dtypes, Device, nn
-from tinygrad.helpers import IMAGE, CI
+from tinygrad.helpers import IMAGE, CI, getenv
 
 class TestTiny(unittest.TestCase):
 
@@ -27,7 +27,7 @@ class TestTiny(unittest.TestCase):
     out = Tensor.ones(256).contiguous().sum()
     self.assertEqual(out.item(), 256)
 
-  def test_gemm(self, N=64, out_dtype=dtypes.float):
+  def test_gemm(self, N=getenv("GEMM_N", 64), out_dtype=dtypes.float):
     a = Tensor.ones(N,N).contiguous()
     b = Tensor.eye(N).contiguous()
     lst = (out:=a@b).tolist()
@@ -36,6 +36,14 @@ class TestTiny(unittest.TestCase):
         self.assertEqual(lst[y][x], 1.0, msg=f"mismatch at ({y},{x})")
     if IMAGE < 2: self.assertEqual(out.dtype, out_dtype)
 
+  def test_gemv(self, N=getenv("GEMV_N", 64), out_dtype=dtypes.float):
+    a = Tensor.ones(1,N).contiguous()
+    b = Tensor.eye(N).contiguous()
+    lst = (out:=a@b).tolist()
+    for x in range(N):
+      self.assertEqual(lst[0][x], 1.0, msg=f"mismatch at {x}")
+    if IMAGE < 2: self.assertEqual(out.dtype, out_dtype)
+
   # *** randomness ***
 
   def test_random(self):

test/test_uop_graph.py

@@ -3,7 +3,7 @@ import unittest, pytest
 from tinygrad import dtypes, Variable
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import DEBUG, Context
-from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp
+from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp, KernelInfo
 from tinygrad.uop.symbolic import sym
 from tinygrad.codegen import full_rewrite, full_rewrite_to_sink
 from tinygrad.codegen.late.expander import expander
@@ -17,7 +17,7 @@ simple_pm = PatternMatcher([
 
 def to_uops_list(u:List[UOp]) -> List[UOp]:
   # we strip the SINK here for legacy reasons
-  ret = full_rewrite(UOp.sink(*u))
+  ret = full_rewrite(UOp.sink(*u, arg=KernelInfo(opts_to_apply=())))
   assert ret[-1].op is Ops.SINK
   return ret[:-1]
 
@@ -445,7 +445,7 @@ class TestUOpGraph(unittest.TestCase):
     with Context(IGNORE_OOB=0):
       glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), src=(), arg=0)
       v = Variable("v", 0, 20)
-      st0 = UOp(Ops.STORE, dtypes.void, src=(glbl0.index(v), UOp.const(dtypes.int, 0), UOp(Ops.IF, src=(v<16,))))
+      st0 = UOp(Ops.STORE, dtypes.void, src=(glbl0.index(v, v<16), UOp.const(dtypes.int, 0)))
       to_uops_list([st0])
 
       st1 = UOp(Ops.STORE, dtypes.void, (glbl0.index(v), v, v<20))
@@ -458,12 +458,12 @@ class TestUOpGraph(unittest.TestCase):
       sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.uint.ptr(8, addrspace=AddrSpace.LOCAL), (), "temp0")
 
       # Define indices, valids and barrier
-      gidx = UOp(Ops.SPECIAL, dtypes.int, (), ("gidx0", 416))
-      lidx = UOp(Ops.SPECIAL, dtypes.int, (), ("lidx0", 10))
+      gidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 416),), "gidx0")
+      lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 10),), "lidx0")
 
       gate = (gidx<400) & (lidx<8)
 
-      local_store = UOp(Ops.STORE, dtypes.void, (sbuf.index(lidx), UOp.const(dtypes.uint, 1), UOp(Ops.IF, src=(lidx<8,))))
+      local_store = UOp(Ops.STORE, dtypes.void, (sbuf.index(lidx, lidx<8), UOp.const(dtypes.uint, 1)))
 
       barrier = UOp(Ops.BARRIER, dtypes.void, (local_store,))
       if_barrier = UOp(Ops.IF, dtypes.void, (gate, barrier))
@@ -477,7 +477,7 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_load_with_float_in_index(self):
     with Context(IGNORE_OOB=0):
-      ridx = UOp.range(dtypes.int, 20, 0)
+      ridx = UOp.range(20, 0)
       glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
       i = (ridx.cast(dtypes.float)*0.68).trunc().cast(dtypes.int)
       ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i, ((0<=i)&(i<16))),))
@@ -490,7 +490,7 @@ class TestUOpGraph(unittest.TestCase):
   def test_load_cast_to_bool(self):
     with Context(IGNORE_OOB=0):
       glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), (), 0)
-      ridx = UOp.range(dtypes.int, 20, 0)
+      ridx = UOp.range(20, 0)
       ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(ridx, ridx.cast(dtypes.bool).logical_not()),))
       to_uops_list([ld0])
 
@@ -499,7 +499,7 @@ class TestUOpGraph(unittest.TestCase):
     with Context(IGNORE_OOB=0):
       glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
       mask = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(16), (), 0)
-      ridx = UOp.range(dtypes.int, 20, 0)
+      ridx = UOp.range(20, 0)
       ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(ridx, ridx<16&mask),)))
       to_uops_list([ld0])
 
@@ -512,7 +512,7 @@ class TestUOpGraph(unittest.TestCase):
   def test_in_out_bounds_access_with_mask(self):
     with Context(IGNORE_OOB=0):
       glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      gidx0 = UOp(Ops.SPECIAL, dtype=dtypes.int, arg=("gidx0", 42))
+      gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 42),), "gidx0")
       ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0, (5<gidx0)&(gidx0<16)),))
       ld1 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0, gidx0<16),))
       to_uops_list([ld0, ld1])
@@ -536,7 +536,7 @@ class TestUOpGraph(unittest.TestCase):
     with Context(IGNORE_OOB=0):
       glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
       glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(64), (), 0)
-      gidx0 = UOp(Ops.SPECIAL, dtype=dtypes.int, arg=("gidx0", 42))
+      gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 42),), "gidx0")
       ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0, gidx0<8),))
       ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index(ld0*2, (ld0>=0)&(ld0<32)),))
       to_uops_list([ld1])
@@ -559,7 +559,7 @@ class TestUOpGraph(unittest.TestCase):
   def test_fold_gated_load_local(self):
     glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
     smem = UOp(Ops.DEFINE_LOCAL, dtypes.int.ptr(size=18, addrspace=AddrSpace.LOCAL), (), "temp")
-    lidx = UOp(Ops.SPECIAL, dtypes.int, (), ("lidx0", 16))
+    lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 16),), "lidx0")
     st = UOp(Ops.STORE, dtypes.void, (smem.index(lidx), UOp.load(glbl0.index(lidx), dtype=dtypes.int)))
     barrier = UOp(Ops.BARRIER, dtypes.void, (st, ))
     ld0 = UOp(Ops.LOAD, dtypes.int, (smem.index(lidx+1, UOp.const(dtypes.bool, False)), barrier))
@@ -592,8 +592,8 @@ class TestUOpGraph(unittest.TestCase):
   def test_switched_range_order(self):
     glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
     cf = UOp.const(dtypes.float, 0.0)
-    r1 = UOp.range(dtypes.int, 2, 0)
-    r2 = UOp.range(dtypes.int, 2, 1)
+    r1 = UOp.range(2, 0)
+    r2 = UOp.range(2, 1)
     alu = UOp(Ops.MUL, dtypes.int, (r2, r1))
     store = UOp(Ops.STORE, dtypes.void, (glbl.index(alu), cf))
     uops = to_uops_list([store])
@@ -756,8 +756,8 @@ class TestIFUOps(unittest.TestCase):
   def test_create_ifs(self):
     gbuf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
     sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(size=4, addrspace=AddrSpace.LOCAL), (), "smem")
-    valid = UOp(Ops.SPECIAL, dtypes.int, (), ("gidx0", 10))<5
-    lidx = UOp(Ops.SPECIAL, dtypes.int, (), ("lidx0", 4))
+    valid = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 10),), "gidx0")<5
+    lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), "lidx0")
     gate = valid&(lidx.ne(2))
     idx = UOp.const(dtypes.int, 0)
     st = UOp(Ops.STORE, dtypes.void, (sbuf.index(idx), UOp.const(dtypes.float, 42)))
@@ -775,8 +775,8 @@ class TestIFUOps(unittest.TestCase):
   def test_expand_ifs_one_gate(self):
     gbuf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
     sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(size=16, addrspace=AddrSpace.LOCAL), (), "smem")
-    valid = UOp(Ops.SPECIAL, dtypes.int, (), ("gidx0", 4))<1
-    lidx = UOp(Ops.SPECIAL, dtypes.int, (), ("lidx0", 16))
+    valid = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), "gidx0")<1
+    lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 16),), "lidx0")
     gate = valid&(lidx.ne(2))
     st = UOp(Ops.STORE, dtypes.void, (sbuf, lidx, UOp.const(dtypes.float, 42)))
     barrier = UOp(Ops.BARRIER, dtypes.void, (st,))
@@ -794,8 +794,8 @@ class TestIFUOps(unittest.TestCase):
   @unittest.expectedFailure
   def test_expand_ifs_dumb(self):
     buf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    valid = UOp(Ops.SPECIAL, dtypes.int, (), ("gidx0", 10))<5
-    lidx = UOp(Ops.SPECIAL, dtypes.int, (), ("lidx0", 4))
+    valid = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 10),), "gidx0")<5
+    lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), "lidx0")
     gate = valid&(lidx.ne(2))
     stores = [UOp(Ops.STORE, dtypes.void, (buf, UOp.const(dtypes.int, i), UOp.const(dtypes.float, i), gate)) for i in range(4)]
     sink = UOp(Ops.SINK, dtypes.void, tuple(stores))

test/test_uops.py

@@ -14,7 +14,7 @@ from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.codegen import full_rewrite
 from tinygrad.uop.symbolic import sym
 from tinygrad.device import is_dtype_supported
-from tinygrad.codegen.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt import Opt, OptOps
 
 def to_uops_list(u:list[UOp], opts=None, skip_check=False) -> list[UOp]: return full_rewrite(UOp.sink(*u), opts)
 
@@ -22,7 +22,7 @@ def _uops_to_prg(uops_list):
   uops = full_rewrite(ast:=UOp.sink(*uops_list), opts=Device[Device.DEFAULT].renderer)
   src = Device[Device.DEFAULT].renderer.render(uops)
   has_local = Device[Device.DEFAULT].renderer.has_local
-  return CompiledRunner(ProgramSpec("test", src, Device.DEFAULT, ast, uops=uops,
+  return CompiledRunner(ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test", src, Device.DEFAULT, ast, uops=uops,
                                 global_size=[1,1,1] if has_local else None, local_size=[1,1,1] if has_local else None))
 
 def uop(uops:list[UOp], uop:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
@@ -177,6 +177,32 @@ class TestBoolUOps(TestUOps):
   def test_cmplt_bool(self): self._test_bop_bool_fxn(Ops.CMPLT, lambda a,b: a < b)
   def test_where_bool(self): self._test_top_bool_fxn(Ops.WHERE, lambda a,b,c: b if a else c)
 
+class TestSafeCast(TestUOps):
+  def test_cast_folds(self):
+    a = UOp.variable("a", 1, 10, dtype=dtypes.int32)
+    self.assertEqual(a.cast(dtypes.int64).cast(dtypes.int32).simplify(), a)
+    self.assertEqual(a.cast(dtypes.double).cast(dtypes.int32).simplify(), a)
+    a = UOp.variable("a", 1, 10, dtype=dtypes.uint8)
+    self.assertEqual(a.cast(dtypes.int64).cast(dtypes.uint8).simplify(), a)
+    self.assertEqual(a.cast(dtypes.uint32).cast(dtypes.uint8).simplify(), a)
+
+  def test_remove_intermediate_cast(self):
+    a = UOp.variable("a", 0., 100., dtype=dtypes.half)
+    self.assertEqual(a.cast(dtypes.double).cast(dtypes.float).simplify(), a.cast(dtypes.float))
+    a = UOp.variable("a", 1, 10, dtype=dtypes.int32)
+    # TODO: double preserves certain int dtypes
+    self.assertEqual(a.cast(dtypes.double).cast(dtypes.float).simplify(), a.cast(dtypes.float))
+    self.assertEqual(a.cast(dtypes.int64).cast(dtypes.int16).simplify(), a.cast(dtypes.int16))
+    a = UOp.variable("a", 1, 10, dtype=dtypes.uint8)
+    self.assertEqual(a.cast(dtypes.int64).cast(dtypes.int32).simplify(), a.cast(dtypes.int32))
+
+  def test_safe_cast_using_bounds(self):
+    a = UOp.variable("a", 1, 10, dtype=dtypes.uint64)
+    self.assertEqual(a.cast(dtypes.int16).cast(dtypes.int).simplify(), a.cast(dtypes.int))
+    a = UOp.variable("a", -10, 10, dtype=dtypes.int32)
+    self.assertEqual(a.cast(dtypes.int8).cast(dtypes.int64).simplify(), a.cast(dtypes.int64))
+    self.assertEqual(a.cast(dtypes.int8).cast(dtypes.float).simplify(), a.cast(dtypes.float))
+
 class TestExecALU(TestUOps):
   def test_sqrt(self):
     self.assertEqual(exec_alu(Ops.SQRT, dtypes.float, (0.0,)), 0.0)
@@ -244,7 +270,7 @@ class TestConstantFolding(unittest.TestCase):
 class TestGatedStoreRewrite(unittest.TestCase):
   def test_tiny_gate_store(self):
     gmem = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (), ('gidx0', 4))
+    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0')
     gate = gidx0<UOp.const(dtypes.int, 1)
     idx = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem, gidx0 * UOp.const(dtypes.int, 2), gate))
     val = UOp.const(dtypes.float, 42.0)
@@ -261,7 +287,7 @@ class TestGatedStoreRewrite(unittest.TestCase):
   def test_gate_some_stores(self):
     gmem0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
     gmem1 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 1)
-    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (), ('gidx0', 4))
+    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0')
     idx = gidx0 * UOp.const(dtypes.int, 2)
     idx0 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem0, idx, gidx0<UOp.const(dtypes.int, 1)))
     idx1 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem1, idx))
@@ -280,7 +306,7 @@ class TestGatedStoreRewrite(unittest.TestCase):
   def test_merge_ifs_alt(self):
     gmem0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
     gmem1 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 1)
-    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (), ('gidx0', 4))
+    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0')
     idx = gidx0*UOp.const(dtypes.int, 2)
     gate = gidx0<UOp.const(dtypes.int, 1)
     idx0 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem0, idx, gate))
@@ -403,7 +429,7 @@ class TestAssembly(unittest.TestCase):
     self.assertNotIn(Ops.IDIV, ops)
 
   def test_fast_idiv_remove_powers_of_two(self):
-    ridx = UOp.range(dtypes.int, 2**20, 0)
+    ridx = UOp.range(2**20, 0)
     uops = to_uops_list([ridx//(7*64)], opts=Device[Device.DEFAULT].renderer)
     ops = [x.op for x in uops]
     # this requires shifting out the powers of two before doing fast_idiv
@@ -447,13 +473,13 @@ class TestUOpMethod(unittest.TestCase):
   def test_uop_variables(self):
     a = UOp.variable("a", 1, 10)
     uop_var = Tensor(a.bind(1))
-    st_var = Tensor.empty((2, 1)).reshape((2, a.bind(1)))
+    st_var = Tensor.empty((2, 10))[:, :a.bind(1)]
     _, var_vals = (uop_var+st_var).schedule_with_vars()
     self.assertEqual(len(var_vals), 1)
     self.assertEqual(list(var_vals)[0], a)
 
   def test_const_factor(self):
-    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (), ('gidx0', 8))
+    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 8),), 'gidx0')
     self.assertEqual(UOp(Ops.CONST, dtypes.int, (), 17).const_factor(), 17)
     self.assertEqual(gidx0.const_factor(), 1)
     self.assertEqual((gidx0*3).const_factor(), 3)

test/test_uops_stats.py

@@ -1,12 +1,12 @@
 import unittest
 from tinygrad import Tensor
-from tinygrad.helpers import getenv, GlobalCounters
+from tinygrad.helpers import getenv, GlobalCounters, EMULATE
 from tinygrad.engine.realize import lower_schedule_item, ProgramSpec, get_program
 from tinygrad.renderer import Estimates
 from tinygrad.codegen import full_rewrite
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.dtype import dtypes
-from tinygrad.codegen.opt.kernel import Opt, OptOps, KernelOptError
+from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.device import Device
 
 def flops_mem(uops, ignore_indexing=False):
@@ -87,6 +87,7 @@ class TestUOpsStatsMatmulHalf(unittest.TestCase):
     expected_ops = N ** 3 * 2
     self.assertEqual(expected_ops, GlobalCounters.global_ops)
 
+  @unittest.skipIf(EMULATE.value=="INTEL", "intel gets 524288 != 524352")
   def test_bigger_matmul_half(self): self.test_simple_matmul_half(64)
 
   def test_batched_matmul_half(self, N=16):
@@ -98,7 +99,6 @@ class TestUOpsStatsMatmulHalf(unittest.TestCase):
     self.assertEqual(expected_ops, GlobalCounters.global_ops)
 
 class TestUOpsStats(unittest.TestCase):
-  @unittest.skipIf(getenv("PTX"), "wrong in PTX")
   def test_simple_add(self):
     a = Tensor.empty(100,100)
     b = Tensor.empty(100,100)
@@ -110,7 +110,6 @@ class TestUOpsStats(unittest.TestCase):
     # NOTE; ops also include indexing ops
     assert expected_ops <= ops and ops <= expected_ops * 2
 
-  @unittest.skipIf(getenv("PTX"), "wrong in PTX")
   def test_simple_add_sq(self):
     a = Tensor.empty(100,100)
     b = Tensor.empty(100,100)

test/unit/test_attention.py

@@ -1,5 +1,6 @@
 import unittest
-from tinygrad import Tensor, dtypes
+from tinygrad import Tensor, dtypes, TinyJit, UOp
+from tinygrad.apps.llm import apply_rope
 
 # TODO: test_scheduler, but just in uint
 class TestAttention(unittest.TestCase):
@@ -16,5 +17,29 @@ class TestAttention(unittest.TestCase):
     for si in softmax_inputs:
       assert all(b.dtype == dtypes.half for b in si.bufs), f"non half {si.bufs=}"
 
+  def test_apply_rope(self):
+    x = Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32)
+    result = apply_rope(x, 0)
+    self.assertEqual(result.shape, x.shape)
+    self.assertEqual(result.dtype, x.dtype)
+    self.assertGreater((result - apply_rope(x, 5)).abs().max().item(), 1e-6)
+    with self.assertRaises(AssertionError): apply_rope(Tensor.randn(1, 1, 4, 7, dtype=dtypes.float32), 0)
+
+  def test_apply_rope_jit_prune(self):
+    def rope_fn(x_in, pos): return apply_rope(x_in, pos)
+    rope_noprune = TinyJit(rope_fn)
+    rope_prune = TinyJit(rope_fn, prune=True)
+
+    v_pos = UOp.variable("start_pos", 0, 100)
+    for _ in range(3):
+      rope_noprune(Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32), v_pos.bind(1))
+      rope_prune(Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32), v_pos.bind(1))
+    noprune_size = len(rope_noprune.captured.jit_cache)
+    prune_size = len(rope_prune.captured.jit_cache)
+
+    self.assertGreater(noprune_size, prune_size)
+    self.assertGreaterEqual(noprune_size, 3)
+    self.assertEqual(prune_size, 1)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_block_reorder.py

@@ -37,8 +37,8 @@ class TestBlockReorder(unittest.TestCase):
     a = UOp(Ops.DEFINE_GLOBAL, dtype=dtypes.float.ptr(), arg=0)
     b = UOp(Ops.DEFINE_GLOBAL, dtype=dtypes.float.ptr(), arg=1)
     c = UOp(Ops.DEFINE_GLOBAL, dtype=dtypes.float.ptr(), arg=2)
-    v1 = UOp(Ops.SPECIAL, dtype=dtypes.int, arg=("gidx0", 4))
-    v2 = UOp(Ops.SPECIAL, dtype=dtypes.int, arg=("gidx1", 4))
+    v1 = UOp(Ops.SPECIAL, dtype=dtypes.int, src=(UOp.const(dtypes.int, 4),), arg="gidx0")
+    v2 = UOp(Ops.SPECIAL, dtype=dtypes.int, src=(UOp.const(dtypes.int, 4),), arg="gidx1")
     v1 = v1*27
     v2 = v2*4
     loads = [

test/unit/test_dtype.py

@@ -21,6 +21,10 @@ class TestEqStrDType(unittest.TestCase):
   def test_ptr_eq(self):
     assert dtypes.float32.ptr() == dtypes.float32.ptr()
     assert not (dtypes.float32.ptr() != dtypes.float32.ptr())
+  def test_ptr_nbytes(self):
+    assert dtypes.float16.ptr(32).nbytes() == 32 * dtypes.float16.itemsize
+  def test_ptr_nbytes_unlimited(self):
+    self.assertRaises(RuntimeError, lambda: dtypes.float32.ptr().nbytes())
   def test_strs(self):
     if PtrDType is None: raise unittest.SkipTest("no PtrDType support")
     self.assertEqual(str(dtypes.imagef((1,2,4))), "dtypes.imagef((1, 2, 4))")

test/unit/test_dtype_spec.py

@@ -100,6 +100,11 @@ class TestHelpers(unittest.TestCase):
         np.testing.assert_equal(dt.min, False)
         np.testing.assert_equal(dt.max, True)
 
+  def test_dtype_range_vec(self):
+    for dt in core_dtypes:
+      self.assertEqual(dt.min, dt.vec(4).min)
+      self.assertEqual(dt.max, dt.vec(4).max)
+
   def test_truncate_fp16(self):
     self.assertEqual(truncate_fp16(1), 1)
     self.assertEqual(truncate_fp16(65504), 65504)
@@ -613,4 +618,4 @@ class TestAutoCastType(unittest.TestCase):
     np.testing.assert_allclose(out.numpy(), tt.log_softmax(0).numpy(), rtol=1e-3)
     out = t.log_softmax(0, dtype=dtypes.float)
     self.assertEqual(out.dtype, dtypes.float)
-    np.testing.assert_allclose(out.numpy(), tt.log_softmax(0, dtype=torch.float).numpy(), rtol=1e-3)
+    np.testing.assert_allclose(out.numpy(), tt.log_softmax(0, dtype=torch.float).numpy(), rtol=1e-3)

test/unit/test_graph_rewrite.py

@@ -65,21 +65,21 @@ class TestFoldingAndReduction(unittest.TestCase):
   def test_full_graph_rewrite_reduction_with_unused_range(self):
     const1 = UOp.const(dtypes.int32, 15)
     const2 = UOp.const(dtypes.int32, 25)
-    rng = UOp.range(dtypes.int32, 10, idx=0)
+    rng = UOp.range(10, idx=0)
     optimized_sink = apply_rewrite((const1 + const2).reduce(Ops.ADD, rng))
     expected_sum = 10 * (15 + 25)
     self.assertEqual(optimized_sink.arg, expected_sum)
 
   @unittest.skip("currently failing")
   def test_full_graph_rewrite_range_reduction(self):
-    simple_range = UOp.range(dtypes.int32, 5, idx=0)
+    simple_range = UOp.range(5, idx=0)
     optimized_sink = apply_rewrite(simple_range.reduce(Ops.ADD, simple_range))
     expected_sum = sum(range(5))
     self.assertEqual(optimized_sink.arg, expected_sum)
 
   @unittest.skip("currently failing")
   def test_full_graph_rewrite_simple_reduction_folding(self):
-    simple_range = UOp.range(dtypes.int32, 4, idx=0)
+    simple_range = UOp.range(4, idx=0)
     add_uop = simple_range + UOp.const(dtypes.int32, 1)
     optimized_sink = apply_rewrite(add_uop.reduce(Ops.ADD, simple_range))
     expected_sum = sum(i + 1 for i in range(4))
@@ -87,8 +87,8 @@ class TestFoldingAndReduction(unittest.TestCase):
 
   @unittest.skip("currently failing")
   def test_full_graph_rewrite_nested_loop_collapse(self):
-    outer_range = UOp.range(dtypes.int32, 8, 0)
-    inner_range = UOp.range(dtypes.int32, 4, 1)
+    outer_range = UOp.range(8, 0)
+    inner_range = UOp.range(4, 1)
     expr = (outer_range * 10) + inner_range
     optimized_reduce_uop = apply_rewrite(expr.reduce(Ops.ADD, outer_range, inner_range))
     self.assertEqual(optimized_reduce_uop.op, Ops.CONST)
@@ -116,7 +116,7 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
 
   def test_graph_rewrite_div_folding_bug(self):
     lhs = UOp(Ops.ADD, dtypes.int.vec(4), src=(
-      UOp(Ops.VECTORIZE, dtypes.int.vec(4), arg=None, src=(UOp(Ops.SPECIAL, dtypes.int, arg=('lidx0', 32), src=()),)*4),
+      UOp(Ops.VECTORIZE, dtypes.int.vec(4), arg=None, src=(UOp(Ops.SPECIAL, dtypes.int, arg='lidx0', src=(UOp.const(dtypes.int, 32),)),)*4),
       UOp(Ops.VCONST, dtypes.int.vec(4), arg=(0, 256, 512, 768), src=())))
     rhs = UOp.const(dtypes.int.vec(4), 2)
     unopt = lhs<rhs

test/unit/test_helpers.py

@@ -93,7 +93,7 @@ class TestMergeDicts(unittest.TestCase):
     assert merge_dicts([a, b]) == {"a": 1, "b": 2, "c": 3}
     assert merge_dicts([a, c]) == a
     assert merge_dicts([a, b, c]) == {"a": 1, "b": 2, "c": 3}
-    with self.assertRaises(AssertionError):
+    with self.assertRaises(RuntimeError):
       merge_dicts([a, d])
 
 class TestStripParens(unittest.TestCase):

test/unit/test_linearizer_rewrite.py

@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor, Context, Device
 from tinygrad.engine.realize import get_program
-from tinygrad.codegen.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.uop.ops import KernelInfo
 
 class TestLinearizerRewrite(unittest.TestCase):

test/unit/test_rewrite_bottom_up_gate.py

@@ -0,0 +1,28 @@
+import unittest
+from tinygrad.uop.ops import PatternMatcher, UOp, graph_rewrite, Ops, UPat, BottomUpGate
+
+def assert_not_reached(): assert False, "This function should not be reached"
+def gate(): raise BottomUpGate
+
+class TestBottomUpGate(unittest.TestCase):
+  def test_basic_bottom_up_gate(self):
+    """Test that BottomUpGate stops bottom-up"""
+    pm = PatternMatcher([
+      (UPat(Ops.ADD), gate),
+      (UPat(Ops.MUL), assert_not_reached)
+    ])
+
+    a,b,c = UOp.variable("a",0,10), UOp.variable("b",0,10), UOp.variable("c",0,10)
+    graph_rewrite((a*a)+(b*c), pm, bottom_up=True)
+
+  def test_bottom_up_gate_with_rewriting(self):
+    pm = PatternMatcher([
+      (UPat.var("a")+UPat.var("a"), lambda a: 2*a),
+      (UPat(Ops.MUL), gate),
+      (UPat(Ops.CONST), assert_not_reached)
+    ])
+    a = UOp.variable("a",0,10)
+    graph_rewrite(a+a, pm, bottom_up=True)
+
+if __name__ == "__main__":
+  unittest.main()

test/unit/test_search.py

@@ -1,56 +0,0 @@
-import unittest
-from tinygrad import Tensor, Device
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.device import Buffer
-from tinygrad.codegen.opt.search import get_test_global_size, bufs_from_lin
-from tinygrad.helpers import GlobalCounters
-from extra.optimization.helpers import time_linearizer
-from test.test_linearizer import push_views
-
-class TestSearchUtil(unittest.TestCase):
-  def test_get_test_global_size(self):
-    self.assertEqual(get_test_global_size([256, 256, 256], 65536, {}), ([256, 16, 16], 256.0))
-    self.assertEqual(get_test_global_size([65536, 1, 1], 256, {}), ([256, 1, 1], 256.0))
-    self.assertEqual(get_test_global_size([77, 1, 1], 16, {}), ([9, 1, 1], 77/9))
-
-  def test_bufs_from_lin(self):
-    a = Tensor([1,2,3,4]).realize()
-    si = (a+1).schedule()[0]
-    rawbufs = bufs_from_lin(Kernel(si.ast))
-    assert len(rawbufs) == 2
-    assert all(r is not None for r in rawbufs)
-    assert all(isinstance(r, Buffer) for r in rawbufs)
-    assert all(r.size > 0 for r in rawbufs)
-
-  def test_bufs_from_lin_alt(self):
-    a = Tensor.randn(4, 4).realize()
-    b = a+a[0]
-    si = b.schedule()[0]
-    rawbufs = bufs_from_lin(Kernel(push_views(si.ast)))
-    assert len(rawbufs) == 2
-    assert all(r is not None for r in rawbufs)
-    assert all(isinstance(r, Buffer) for r in rawbufs)
-    assert all(r.size > 0 for r in rawbufs)
-
-class TestTimeLinearizer(unittest.TestCase):
-  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WebGPU timestamps are low precision, tm is 0")
-  def test_reasonable_time(self):
-    a = Tensor([1,2,3,4]).realize()
-    si = (a+1).schedule()[0]
-    # create fresh empty buffers
-    rawbufs = [Buffer(b.device, b.size, b.dtype).allocate() for b in si.bufs]
-    tm = time_linearizer(Kernel(push_views(si.ast)), rawbufs, allow_test_size=False, cnt=10, disable_cache=True)
-    assert tm > 0 and tm != float('inf')
-
-  # Ensure that the kernel count is not incremented by time_linearizer when clearing l2
-  def test_kernel_count(self):
-    ast = Tensor.zeros(16).contiguous().kernelize().uop.src[1].arg.ast
-    lin = Kernel(push_views(ast))
-    bufs = bufs_from_lin(lin)
-
-    kernel_count = GlobalCounters.kernel_count
-    time_linearizer(lin, bufs, allow_test_size=False, cnt=2, disable_cache=True, clear_l2=True)
-    assert GlobalCounters.kernel_count == kernel_count, "kernel count was incremented by time_linearizer"
-
-if __name__ == "__main__":
-  unittest.main()

test/unit/test_shapetracker.py

@@ -839,25 +839,22 @@ class TestRender(unittest.TestCase):
     self.assertEqual(idx.render(), "((ridx0*3)+ridx1)")
     self.assertEqual(valid.render(), "(ridx0<2)")
 
-class TestVariableReshape(unittest.TestCase):
-  def test_reshape(self):
-    st = ShapeTracker.from_shape((3,))
-    st = st.reshape((Variable("i", 1, 10),))
+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_reshape_stride_0(self):
-    st = ShapeTracker.from_shape((3,), (0,))
-    st = st.reshape((Variable("i", 1, 10).bind(3),))
-    assert len(st.views) == 1, f"multiview {st}"
-
-  def test_reshape_bound(self):
-    st = ShapeTracker.from_shape((3,))
-    st = st.reshape((Variable("i", 1, 10).bind(3),))
+  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
 
-  def test_add(self):
-    st1 = ShapeTracker.from_shape((3,))
-    st2 = ShapeTracker.from_shape((Variable("i", 1, 10),))
+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
 
@@ -867,15 +864,17 @@ class TestVariableReshape(unittest.TestCase):
     st = st1+st2
     assert len(st.views) == 1, f"multiview {st}"
 
-  def test_add_bound(self):
-    st1 = ShapeTracker.from_shape((3,))
-    st2 = ShapeTracker.from_shape((Variable("i", 1, 10).bind(3),))
+  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):
-    st1 = ShapeTracker.from_shape((3,))
-    st2 = ShapeTracker.from_shape((Variable("i", 1, 10).bind(3),))
+    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

test/unit/test_shapetracker_math.py

@@ -87,20 +87,6 @@ class TestShapeTrackerAdd(unittest.TestCase):
     assert not (st_equal(st1, st2))
 
 class TestShapeTrackerAddVariable(unittest.TestCase):
-  def test_self_add(self):
-    j = Variable("j", 0, 20).bind(10)
-    a = ShapeTracker.from_shape((10,10))
-    x = a.reshape((10, j))
-    out = x + x
-    assert out == x
-
-  def test_self_add_reshape(self):
-    j = Variable("j", 0, 20).bind(10)
-    a = ShapeTracker.from_shape((10,10))
-    x = a.reshape((10, j))
-    out = x.reshape((5, 2, j)) + x
-    assert out == x
-
   def test_merge_symbolic_views(self):
     var_i = Variable('i', 1, 10)
     var_j = Variable('i', 1, 10)

test/unit/test_simplify_valid_idx.py

@@ -4,6 +4,7 @@ from tinygrad.codegen import full_rewrite_to_sink
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.uop.symbolic import simplify_valid
+from tinygrad.helpers import Context
 
 def get_gated_load_uop(valid:UOp, idx:UOp):
   return UOp(Ops.LOAD, dtypes.float, (
@@ -17,9 +18,9 @@ def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UO
     UOp(Ops.VECTORIZE, dtypes.float.vec(4), src=(UOp.const(dtypes.float, 0.0),) * 4)
   ))
 
-def Special(expr, nmax): return UOp(Ops.SPECIAL, dtypes.int, (), (expr, nmax))
+def Special(expr, nmax): return UOp(Ops.SPECIAL, dtypes.int32, (UOp.const(dtypes.int, nmax),), expr)
 def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax)
-def Range(n, nmax): return UOp.range(dtypes.int, nmax, n)
+def Range(n, nmax): return UOp.range(nmax, n)
 
 class TestHelpers(unittest.TestCase):
   def test_is_increasing(self):
@@ -45,7 +46,8 @@ class TestHelpers(unittest.TestCase):
 
 class TestValidIdxSimplification(unittest.TestCase):
   def check(self, load, sidx, svalid):
-    load = full_rewrite_to_sink(load.sink()).src[0]
+    with Context(NOOPT=1):
+      load = full_rewrite_to_sink(load.sink()).src[0]
     idx, valid = load.src[0].src[1], load.src[0].src[2]
     self.assertEqual(idx.render(simplify=False), sidx)
     self.assertEqual(valid.render(simplify=False), svalid)
@@ -195,9 +197,21 @@ class TestValidIdxSimplification(unittest.TestCase):
       "1",
       "((((ridx0+ridx1)<1)!=True)&(((ridx2+ridx3)<1)!=True))")
 
+  def test_valid_with_non_const_rhs(self):
+    ridx0 = Range(0, 2**16)
+    ridx1 = Range(1, 4)
+    ridx2 = Range(2, 4)
+    valid = (ridx0<(ridx1*4 + ridx2))&(ridx0<-1).ne(True)
+    idx = ridx0%1024
+    load = get_gated_load_uop(valid, idx)
+    self.check(load,
+      "ridx0",
+      "(ridx0<((ridx1*4)+ridx2))")
+
 class TestImageSimplification(unittest.TestCase):
   def check(self, load, svalid, sidx0, sidx1):
-    load = full_rewrite_to_sink(load.sink()).src[0]
+    with Context(NOOPT=1):
+      load = full_rewrite_to_sink(load.sink()).src[0]
     idx = load.src[0].src[1]
     self.assertEqual(idx.op, Ops.VECTORIZE)
     self.assertEqual(len(idx.src), 2)

test/unit/test_symbolic_shapetracker.py

@@ -48,11 +48,11 @@ class TestSymbolic(unittest.TestCase):
     i = Variable("i", 1, 5).bind(3)
     j = Variable("j", 1, 5).bind(3)
     k = Variable("k", 1, 5).bind(3)
-    t = Tensor.rand(3, 4).reshape(i, 4).cat(Tensor.rand(3, 4).reshape(j, 4), dim=0).cat(Tensor.rand(3, 4).reshape(k, 4), dim=0)
+    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))
     assert st.real_strides() == (4, 1)
-    t = Tensor.rand(3, 3).reshape(i, 3).cat(Tensor.rand(3, 3).reshape(i, 3), dim=0).cat(Tensor.rand(3, 3), dim=0)
+    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))
     assert st.real_strides() == (3, 1)
@@ -61,7 +61,7 @@ class TestSymbolic(unittest.TestCase):
     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, 4).reshape(3, i).cat(Tensor.rand(3, 4).reshape(3, j), dim=1).cat(Tensor.rand(3, 4).reshape(3, k), dim=1)
+    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.real_strides(), (i+j+k, 1))
@@ -109,60 +109,44 @@ class TestShapeTrackerUnbind(unittest.TestCase):
     assert unbound_view == View.create(shape=(v, 4))
     assert var_val == {v: 3}
 
-  def test_reshape_unbind(self):
-    v = Variable("v", 1, 100)
-    bv = Variable("v", 1, 100).bind(3)
-    t = Tensor.rand(3, 4).reshape(bv, 4)
-    unbound_st, var_val = t.uop.st.unbind()
-    assert unbound_st == ShapeTracker((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)),))
+    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),))
     assert var_val == {v: 2}
 
-class TestSymbolicReshapeFromContiguous(unittest.TestCase):
-  def test_reshape_into_symbols_simple(self):
+class TestSymbolicReshape(unittest.TestCase):
+  def test_reshape(self):
+    a = Tensor.rand(5, 4)
+    b = Tensor.rand(5, 6)
     for i in range(1, 6):
       vi = Variable("i", 1, 5).bind(i)
-      t = Tensor.rand(i, 4).reshape(vi, 4)
-      assert t.shape == (vi, 4)
-      t = Tensor.rand(i, 6).reshape(vi, 2, 3)
-      assert t.shape == (vi, 2, 3)
-
-  def test_reshape_symbols_reshape_ints(self):
-    for i in range(1, 6):
-      vi = Variable("i", 1, 5).bind(i)
-      t = Tensor.rand(i, 4).reshape(vi, 4)
-      assert t.shape == (vi, 4)
-      t = t.reshape(i, 4)
-      assert t.shape == (i, 4)
-
-  @unittest.skip("works now")
-  def test_reshape_into_symbols_bad_shape(self):
-    vi = Variable("i", 1, 10).bind(4)
-    # TODO: this never actually worked, it relied on lazy
-    #with self.assertRaises(ValueError):
-    #  Tensor.rand(4, 6).reshape(vi, 6).reshape(1, 77) # reshape to a different size new shape through symbolic shape
-    with self.assertRaises(AssertionError):
-      Tensor.rand(3, 4).reshape(3, (vi+1)) # reshape into non-Variable Node
+      ret = a[:vi]
+      ret = ret.reshape((vi, 4))
+      assert ret.shape == (vi, 4)
+      ret = b[:vi]
+      ret = ret.reshape((vi, 2, 3))
+      assert ret.shape == (vi, 2, 3)
 
   def test_two_symbol_reshape(self):
+    t = Tensor.rand(5, 5)
     for i in range(1, 6):
       for j in range(1, 6):
         vi = Variable("i", 1, 5).bind(i)
         vj = Variable("j", 1, 5).bind(j)
-        t = Tensor.rand(i, j).reshape(vi, vj)
-        assert t.shape == (vi, vj)
-        # NOTE: this is currently not allowed
-        # t = t.reshape(1, vi*vj)
-        # assert t.shape == (1, vi*vj)
-        t = t.reshape(vj, vi)
-        assert t.shape == (vj, vi)
+        ret = t[:vi, :vj]
+        ret = ret.reshape(vj, vi)
+        assert ret.shape == (vj, vi)
+        ret = ret.reshape(vi, vj)
+        assert ret.shape == (vi, vj)
+        ret = ret.reshape(1, vi*vj)
+        assert ret.shape == (1, vi*vj)
 
   def test_symbolic_mask(self):
     # taken from gpt2 single kvcache
@@ -175,41 +159,6 @@ class TestSymbolicReshapeFromContiguous(unittest.TestCase):
     new_shape = (2, (Variable('start_pos', 1, 128)+1), 16, 64)
     assert view.reshape(new_shape) is None
 
-class TestSymbolicReshapeFromNonContiguous(unittest.TestCase):
-  def test_reshape_from_const(self):
-    vi = Variable("i", 1, 5).bind(4)
-    t = Tensor.ones(3, 4).reshape(3, vi)
-    assert t.shape == (3, vi)
-    assert not t.uop.st.contiguous
-    assert len(t.uop.st.views) == 1
-
-  def test_reshape_not_allowed(self):
-    vi = Variable("i", 1, 5).bind(4)
-    with self.assertRaises(ValueError):
-      # different shape length  # TODO: cases where contractions matched might be fine
-      Tensor.ones(3, 4, 1).reshape(3, vi)
-    with self.assertRaises(ValueError):
-      # size matched, but dimensions do not match
-      Tensor.ones(4, 3).reshape(3, vi)
-
-  def test_reshape_from_padded(self):
-    vi = Variable("i", 1, 5).bind(4)
-    t = Tensor.ones(3, 4).contiguous().expand(2, 3, 4).pad(((1, 1), None, None)).shrink((None, None, (1, 3)))
-    st = t.uop.st
-    assert len(st.views) == 1
-    view = st.views[0]
-    assert view.shape == (4, 3, 2)
-    t = t.reshape(vi, 3, 2)
-    st2 = t.uop.st
-    assert len(st2.views) == 1
-    view2 = st2.views[0]
-    # check only shape changed. strides, offset, mask, contiguous remained the same
-    assert view2.shape == (vi, 3, 2)
-    assert view.strides == view2.strides == (0, 4, 1)
-    assert view.offset == view2.offset == 1
-    assert view.mask == view2.mask == ((1, 3), (0, 3), (0, 2))
-    assert not view.contiguous and not view2.contiguous
-
 class TestSymbolicExpand(unittest.TestCase):
   def test_expand_into_symbols(self):
     vi = Variable("i", 1, 5).bind(3)
@@ -220,11 +169,12 @@ class TestSymbolicExpand(unittest.TestCase):
     assert a.shape == (3, vi, vj)
 
   def test_plus_expands_constant(self):
+    a = Tensor.rand(3, 5)
     for i in range(1, 6):
       vi = Variable("i", 1, 5).bind(i)
-      a = Tensor.rand(3, i).reshape(3, vi)
-      a = a + 1
-      self.assertTupleEqual(a.shape, (3, vi))
+      ret = a[:, :vi]
+      ret = ret + 1
+      self.assertTupleEqual(ret.shape, (3, vi))
 
   def test_pad_then_expand_into_symbols(self):
     vi = Variable("i", 1, 10).bind(3)
@@ -234,6 +184,11 @@ class TestSymbolicExpand(unittest.TestCase):
     self.assertEqual(a.reshape(vi*25).shape, (vi*25,))
 
 class TestSymbolicShrink(unittest.TestCase):
+  def test_shrink_symbols_simple(self):
+    vi = Variable("i", 1, 5)
+    t = Tensor.rand(5, 5).shrink(((0, 5),(0,vi)))
+    assert t.shape == (5, vi)
+
   def test_shrink_symbols(self):
     vi = Variable("i", 1, 5)
     t = Tensor.rand(3, 5).shrink(((0, 2), (vi, vi+1)))
@@ -242,10 +197,10 @@ class TestSymbolicShrink(unittest.TestCase):
 class TestSymbolicPad(unittest.TestCase):
   def test_pad(self):
     v = Variable("v", 1, 100).bind(5)
-    t = Tensor.ones(5).reshape(v).pad(((4, 0),)).reshape(9)
-    assert t.shape == (9,)
-    st = t.uop.st
-    print(st)
+    t = Tensor.ones(100)[:v].pad(((4, 0),))
+    t = t.reshape(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

@@ -97,7 +97,7 @@ class TestTensorUopRepresentation(unittest.TestCase):
     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.BUFFER),)))
+    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__':

test/unit/test_uop_spec.py

@@ -81,5 +81,16 @@ class TestUOpSpec(unittest.TestCase):
     with self.assertRaisesRegex(RuntimeError, "UOp verification failed"):
       type_verify([a], tensor_uop_spec)
 
+class TestUOpSink(unittest.TestCase):
+  def test_0(self):
+    s = UOp.sink()
+    self.assertEqual(len(s.src), 0)
+
+  def test_1(self):
+    a = UOp.const(dtypes.int, 0)
+    s1 = UOp.sink(a)
+    s2 = a.sink()
+    self.assertIs(s1, s2)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_uop_symbolic.py

@@ -208,6 +208,16 @@ class TestSymbolic(unittest.TestCase):
     self.assertEqual((Variable("x", -10, 0)%Variable("y", -10, -1))._min_max, (-9, 0))
     self.assertEqual((Variable("x", -10, 0)%Variable("y", 1, 10))._min_max, (-9, 0))
 
+  def test_range_div_its_symbolic_bound(self):
+    a = Variable("a", 1, 10)
+    ridx0 = UOp.range(a+2, 0)
+    self.helper_test_variable(ridx0//(a+2), 0, 0, "0")
+
+  def test_range_mod_its_symbolic_bound(self):
+    a = Variable("a", 1, 10)
+    ridx = UOp.range(a+2, 0)
+    self.helper_test_variable(ridx%(a+2), 0, 11, "ridx0")
+
   def test_div_min_max(self):
     self.helper_test_variable(Variable("a", 2, 7) // 2, 1, 3, "(a//2)")
     self.helper_test_variable(Variable("a", 0, 6) // 2, 0, 3, "(a//2)")
@@ -550,8 +560,13 @@ class TestSymbolic(unittest.TestCase):
 
   def test_div_mod_recombine(self):
     gidx = Variable("gidx", 0, 124)
+    lidx = Variable("lidx", 0, 124)
     self.helper_test_variable(gidx%4+(gidx//4)*4, 0, 124, "gidx")
     self.helper_test_variable((gidx//4)*4+gidx%4, 0, 124, "gidx")
+    self.helper_test_variable(lidx+gidx%4+(gidx//4)*4, 0, 248, "(gidx+lidx)")
+    self.helper_test_variable(lidx+(gidx//4)*4+gidx%4, 0, 248, "(gidx+lidx)")
+    self.helper_test_variable(lidx+(gidx//4)*8+2*(gidx%4), 0, 372, "(lidx+(gidx*2))")
+    self.helper_test_variable(lidx+2*(gidx%4)+(gidx//4)*8, 0, 372, "(lidx+(gidx*2))")
 
   def test_div_mod_recombine_folded_mod(self):
     a = Variable("a", 0, 2)
@@ -722,6 +737,12 @@ class TestSymbolic(unittest.TestCase):
     a = Variable("a", 1, 10, dtypes.int)
     self.helper_test_variable(a.trunc(), 1, 10, "a", test_z3=False)
 
+  def test_do_math_in_int32(self):
+    a = Variable("a", 1, 10)
+    b = Variable("b", 1, 10)
+    self.helper_test_variable(a.cast(dtypes.long)+b.cast(dtypes.long), 2, 20, "(long)((a+b))")
+    self.helper_test_variable(a.cast(dtypes.long)*b.cast(dtypes.long), 1, 100, "(long)((a*b))")
+
 class TestSymbolicNumeric(unittest.TestCase):
   def helper_test_numeric(self, f):
     MIN, MAX = 0, 10

test/unit/test_uop_vmin_vmax.py

@@ -58,9 +58,9 @@ class TestVminVmaxProperties(unittest.TestCase):
     self.assertEqual(uop.vmax, 8)
 
   def test_vmin_vmax_variable_inside_special(self):
-    uop = UOp(Ops.SPECIAL, dtypes.int, arg=('gidx0', UOp(Ops.DEFINE_VAR, dtypes.int, arg=('i', 1, 10))))
+    uop = UOp(Ops.SPECIAL, dtypes.int, arg='gidx0', src=(UOp(Ops.DEFINE_VAR, dtypes.int, arg=('i', 1, 10)),))
     self.assertEqual(uop.vmin, 0)
-    self.assertEqual(uop.vmax, 10)
+    self.assertEqual(uop.vmax, 9)
 
   def test_vmin_vmax_multiplication_0_inf(self):
     # vmin and vmax for multiplication with a variable
@@ -251,6 +251,15 @@ class TestVminVmaxVConst(unittest.TestCase):
     self.assertIs(uop.vmin, False)
     self.assertIs(uop.vmax, True)
 
+  def test_vmin_vmax_vector_with_gep(self):
+    # vmin and vmax for a vector constant of bool values
+    d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 1)
+    idx = UOp.const(dtypes.int, 0)
+    val = UOp(Ops.LOAD, dtypes.int.vec(2), (d1.index(idx),))
+    uop = (val // 32).gep(0)
+    self.assertEqual(uop.vmin, -67108864)
+    self.assertEqual(uop.vmax, 67108863)
+
 class TestConstFactor(unittest.TestCase):
   def test_const_factor_constant(self):
     # const_factor for a constant

test/unit/test_viz.py

@@ -1,11 +1,12 @@
 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.symbolic import sym
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import PROFILE, colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context
+from tinygrad.helpers import PROFILE, colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker
 from tinygrad.device import Buffer
 
 @track_rewrites(name=True)
@@ -16,10 +17,13 @@ def exec_rewrite(sink:UOp, pm_lst:list[PatternMatcher], names:None|list[str]=Non
 
 # real VIZ=1 pickles these tracked values
 from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, _name_cnt
-from tinygrad.viz import serve
-serve.contexts = (tracked_keys, tracked_ctxs, uop_fields)
+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(tracked_keys, tracked_ctxs)
+def get_viz_list(): return get_metadata(traces)
+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])
 
 class BaseTestViz(unittest.TestCase):
   def setUp(self):
@@ -66,7 +70,6 @@ class TestViz(BaseTestViz):
 
   def test_rewrite_location(self):
     def inner(sink): return graph_rewrite(sink, PatternMatcher([]))
-    @track_rewrites(name=True)
     def outer(sink): return inner(sink)
     outer(UOp.variable("a", 1, 10))
     lst = get_viz_list()
@@ -130,7 +133,7 @@ class TestViz(BaseTestViz):
       (UPat(Ops.CONST, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
     ])
     with self.assertRaises(RuntimeError): exec_rewrite(a, [pm])
-    graphs = flatten(x["graph"].values() for x in get_details(tracked_ctxs[0][0]))
+    graphs = flatten(x["graph"].values() for x in get_viz_details(0, 0))
     self.assertEqual(graphs[0], uop_to_json(a)[id(a)])
     self.assertEqual(graphs[1], uop_to_json(b)[id(b)])
     # fallback to NOOP with the error message
@@ -142,7 +145,9 @@ class TestViz(BaseTestViz):
     z = UOp.const(dtypes.int, 0)
     alu = a*z
     exec_rewrite(alu, [sym])
-    graphs = [x["graph"] for x in get_details(tracked_ctxs[0][0])]
+    lst = get_viz_list()
+    self.assertEqual(len(lst), 1)
+    graphs = [x["graph"] for x in get_viz_details(0, 0)]
     # embed const in the parent node when possible
     self.assertEqual(list(graphs[0]), [id(a), id(alu)])
     self.assertEqual(list(graphs[1]), [id(z)])
@@ -175,7 +180,6 @@ class TestVizTree(BaseTestViz):
     c = UOp.variable("c",0,10)
     d = UOp.variable("d",0,10)
     sink = UOp.sink(a+b, c+d)
-    @track_rewrites()
     def tree_rewrite(): return graph_rewrite(sink, root, name="root")
     tree_rewrite()
     lst = get_viz_list()
@@ -246,9 +250,46 @@ class TestVizIntegration(BaseTestViz):
       b = Tensor.empty(1)
       metadata = (alu:=a+b).uop.metadata
       alu.kernelize()
-      graph = next(get_details(tracked_ctxs[0][0]))["graph"]
+      graph = next(get_viz_details(0, 0))["graph"]
     self.assertEqual(len([n for n in graph.values() if repr(metadata) in n["label"]]), 1)
 
+  # tracing also works without a track_rewrites context
+  # all graph_rewrites get put into the default group
+  def test_default_tracing(self):
+    def test(root):
+      return graph_rewrite(root, sym)
+    test(c:=UOp.const(dtypes.int, 1))
+    test(c+1)
+    ls = get_viz_list()
+    self.assertEqual(len(ls), 1)
+    self.assertEqual(ls[0]["name"], "default graph_rewrite")
+
+  # using @track_rewrites organizes function calls into groups
+  # and nicely counts function calls.
+  def test_group_traces(self):
+    @track_rewrites()
+    def test(root):
+      return graph_rewrite(root, sym)
+    test(c:=UOp.const(dtypes.int, 1))
+    test(c+1)
+    ls = get_viz_list()
+    self.assertEqual(len(ls), 2)
+    for i in range(2): self.assertEqual(ls[i]["name"], f"test n{i+1}")
+
+  # @track_rewrites always starts a new group.
+  def test_group_combined(self):
+    def default_test(root): return graph_rewrite(root, sym)
+    tracked_test = track_rewrites()(default_test)
+    c = UOp.const(dtypes.int, 1)
+    default_test(c+1) # goes to the default group
+    tracked_test(c)   # all rewrites after this go inside the second group.
+    default_test(c+2)
+    ls = get_viz_list()
+    self.assertEqual(len(ls), 2)
+    self.assertEqual(list(next(get_viz_details(0, 0))["graph"]), [id(c+1)])
+    self.assertEqual(list(next(get_viz_details(1, 0))["graph"]), [id(c)])
+    self.assertEqual(list(next(get_viz_details(1, 1))["graph"]), [id(c+2)])
+
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry
 from tinygrad.viz.serve import get_profile
 
@@ -265,27 +306,27 @@ def option(i:int) -> int|None: return None if i == 0 else i-1
 def load_profile(lst:list[ProfileEvent]) -> dict:
   ret = get_profile(lst)
   u = TinyUnpacker(ret)
-  dur, global_peak, index_len, layout_len = u("<IQII")
-  strings, dtypes = json.loads(ret[u.offset:u.offset+index_len]).values()
+  total_dur, global_peak, index_len, layout_len = u("<IQII")
+  strings, dtypes, markers = json.loads(ret[u.offset:u.offset+index_len]).values()
   u.offset += index_len
   layout:dict[str, dict] = {}
   for _ in range(layout_len):
     klen = u("<B")[0]
     k = ret[u.offset:u.offset+klen].decode()
     u.offset += klen
-    layout[k] = v = {"shapes":[]}
+    layout[k] = v = {"events":[]}
     event_type, event_count = u("<BI")
     if event_type == 0:
       for _ in range(event_count):
-        name, ref, st, dur, cat, _ = u("<IIIfBI")
-        v["shapes"].append({"name":strings[name], "ref":option(ref), "st":st, "dur":dur, "cat":option(cat)})
+        name, ref, st, dur, _ = u("<IIIfI")
+        v["events"].append({"name":strings[name], "ref":option(ref), "st":st, "dur":dur})
     else:
       v["peak"] = u("<Q")[0]
       for _ in range(event_count):
         alloc, ts, key = u("<BII")
-        if alloc: v["shapes"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
-        else: v["shapes"].append({"event":"free", "ts":ts, "key":key})
-  return {"dur":dur, "peak":global_peak, "layout":layout}
+        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})
+  return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}
 
 class TestVizProfiler(unittest.TestCase):
   def test_perfetto_node(self):
@@ -294,7 +335,7 @@ class TestVizProfiler(unittest.TestCase):
 
     j = load_profile(prof)
 
-    dev_events = j['layout']['NV']['shapes']
+    dev_events = j['layout']['NV']['events']
     self.assertEqual(len(dev_events), 1)
     event = dev_events[0]
     self.assertEqual(event['name'], 'E_2')
@@ -310,14 +351,16 @@ class TestVizProfiler(unittest.TestCase):
 
     j = load_profile(prof)
 
-    event = j['layout']['NV']['shapes'][0]
+    event = j['layout']['NV']['events'][0]
     self.assertEqual(event['name'], 'COPYxx')
     self.assertEqual(event['st'], 0)   # first event
     self.assertEqual(event['dur'], 10)
 
-    event2 = j['layout']['NV:2']['shapes'][0]
+    event2 = j['layout']['NV:2']['events'][0]
     self.assertEqual(event2['st'], 20) # second event, diff clock
 
+    self.assertEqual(j["dur"], (event2["st"]+event2["dur"])-event["st"])
+
   def test_perfetto_graph(self):
     prof = [ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100)),
             ProfileDeviceEvent(device='NV:1', comp_tdiff=decimal.Decimal(-500), copy_tdiff=decimal.Decimal(-50)),
@@ -333,18 +376,18 @@ class TestVizProfiler(unittest.TestCase):
     self.assertEqual(tracks[1], 'NV')
     self.assertEqual(tracks[2], 'NV:1')
 
-    nv_events = j['layout']['NV']['shapes']
+    nv_events = j['layout']['NV']['events']
     self.assertEqual(nv_events[0]['name'], 'E_25_4n2')
     self.assertEqual(nv_events[0]['st'], 0)
     self.assertEqual(nv_events[0]['dur'], 2)
     #self.assertEqual(j['devEvents'][6]['pid'], j['devEvents'][0]['pid'])
 
-    nv1_events = j['layout']['NV:1']['shapes']
+    nv1_events = j['layout']['NV:1']['events']
     self.assertEqual(nv1_events[0]['name'], 'NV -> NV:1')
     self.assertEqual(nv1_events[0]['st'], 954)
     #self.assertEqual(j['devEvents'][7]['pid'], j['devEvents'][3]['pid'])
 
-    graph_events = j['layout']['NV Graph']['shapes']
+    graph_events = j['layout']['NV Graph']['events']
     self.assertEqual(graph_events[0]['st'], nv_events[0]['st'])
     self.assertEqual(graph_events[0]['st']+graph_events[0]['dur'], nv1_events[0]['st']+nv1_events[0]['dur'])
 
@@ -364,6 +407,21 @@ class TestVizProfiler(unittest.TestCase):
     with self.assertRaises(struct.error):
       get_profile(prof)
 
+  def test_python_marker(self):
+    with Context(PROFILE=1):
+      a = Tensor.empty(1, device="NULL")
+      b = Tensor.empty(1, device="NULL")
+      (a+b).realize()
+      profile_marker("test 1")
+      (a*b).realize()
+      profile_marker("test 2")
+    profile_ret = load_profile(cpu_events)
+    markers = profile_ret["markers"]
+    kernels = profile_ret["layout"]["NULL"]["events"]
+    self.assertEqual(len(markers), 2)
+    assert kernels[0]["st"] <= markers[0]["ts"] <= kernels[1]["st"]
+    assert markers[1]["ts"] >= kernels[1]["st"]+kernels[1]["dur"]
+
 def _alloc(b:int):
   a = Tensor.empty(b, device="NULL", dtype=dtypes.char)
   a.uop.buffer.allocate()
@@ -376,7 +434,7 @@ class TestVizMemoryLayout(BaseTestViz):
     profile_ret = load_profile(Buffer.profile_events)
     ret = profile_ret["layout"][f"{a.device} Memory"]
     self.assertEqual(ret["peak"], 2)
-    self.assertEqual(len(ret["shapes"]), 2)
+    self.assertEqual(len(ret["events"]), 2)
 
   def test_del_once(self):
     a = _alloc(1)
@@ -385,7 +443,7 @@ class TestVizMemoryLayout(BaseTestViz):
     profile_ret = load_profile(Buffer.profile_events)
     ret = profile_ret["layout"][f"{b.device} Memory"]
     self.assertEqual(ret["peak"], 1)
-    self.assertEqual(len(ret["shapes"]), 3)
+    self.assertEqual(len(ret["events"]), 3)
 
   def test_alloc_free(self):
     a = _alloc(1)
@@ -395,7 +453,7 @@ class TestVizMemoryLayout(BaseTestViz):
     profile_ret = load_profile(Buffer.profile_events)
     ret = profile_ret["layout"][f"{c.device} Memory"]
     self.assertEqual(ret["peak"], 2)
-    self.assertEqual(len(ret["shapes"]), 4)
+    self.assertEqual(len(ret["events"]), 4)
 
 if __name__ == "__main__":
   unittest.main()

test/unit/test_winograd.py

@@ -1,8 +1,9 @@
-import unittest
+import unittest, sys
 import numpy as np
 from tinygrad import Tensor, GlobalCounters, dtypes, Context, nn
-from tinygrad.helpers import CI, Profiling, WINO, getenv
+from tinygrad.helpers import CI, Profiling, WINO
 
+@unittest.skipIf(sys.platform.startswith("win"), "flaky on Windows")
 class TestWinogradClose(unittest.TestCase):
   def test_close(self):
     inp = Tensor.rand(1, 16, 16, 16)
@@ -18,6 +19,7 @@ class TestWinogradClose(unittest.TestCase):
       test = conv(inp).realize()
     np.testing.assert_allclose(cmp.numpy(), test.numpy(), atol=1e-5)
 
+@unittest.skipIf(sys.platform.startswith("win"), "flaky on Windows")
 class TestWinograd(unittest.TestCase):
   def setUp(self):
     self.old = WINO.value
@@ -28,17 +30,20 @@ class TestWinograd(unittest.TestCase):
   def test_profile(self):
     x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize()
     with Profiling(enabled=not CI, sort='time'):
-      out = Tensor.conv2d(x,w).realize()
-    out.numpy()
+      Tensor.conv2d(x,w).realize()
 
-  def test_four_kernels(self):
+  def test_forward_kernels(self):
     x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize()
-    GlobalCounters.reset()
-    out = Tensor.conv2d(x,w).realize()
-    assert GlobalCounters.kernel_count == 4
-    out.numpy()
+    out = Tensor.conv2d(x,w)
+    self.assertEqual(len(out.schedule()), 4)
+
+  def test_backward_kernels(self):
+    x,w = Tensor.empty(1,4,9,9,requires_grad=True).realize(), Tensor.empty(4,4,3,3,requires_grad=True).realize()
+    out = Tensor.conv2d(x,w, padding=1)
+    out.mean().backward()
+    backward_schedule = Tensor.schedule(x.grad, w.grad)
+    self.assertEqual(len(backward_schedule), 9)
 
-  @unittest.skipIf(getenv("PTX"), "winograd uses too much in PTX")
   def test_counters(self):
     IC, OC, X, Y = 4,4,9,9
     #OC, IC, X, Y = 512, 256, 8, 8

tinygrad/apps/llm.py

@@ -53,17 +53,15 @@ class SimpleTokenizer:
     try: return [ self._normal_tokens[p] for p in parts ]
     except KeyError: raise RuntimeError("token not found")
 
-def apply_rope(x:Tensor, start_pos:int|UOp, base:int=10000):
+def apply_rope(x:Tensor, start_pos:int|UOp, base:float = 10000.0) -> Tensor:
   B, H, T, Hd = x.shape
-  # NOTE: this is usually in a RoPE cache, but tinygrad JIT should prune it outside the kernel
-  # TODO: make it do that
-  freq = base ** (-Tensor.arange(0, 1, 2/Hd, dtype='float32'))
-  angles = Tensor.arange(start_pos, start_pos+T, dtype='float32')[None, None, :, None] * freq
-  cos, sin = angles.cos(), angles.sin()
-  x = x.reshape(B, H, T, Hd // 2, 2)    # split into pairs
-  y1 = x[..., 0] * cos - x[..., 1] * sin
-  y2 = x[..., 0] * sin + x[..., 1] * cos
-  return Tensor.stack(y1, y2, dim=-1).reshape(B, H, T, Hd)
+  assert (Hd & 1) == 0, "RoPE requires an even head dimension"
+  half = Hd // 2
+  angles = (Tensor.arange(T, dtype="float32") + start_pos)[:, None] * (base ** (-(Tensor.arange(half, dtype="float32") / half)))[None, :]
+  cos, sin = angles.cos().reshape(1, 1, T, half).cast(x.dtype), angles.sin().reshape(1, 1, T, half).cast(x.dtype)
+  x_pairs = x.reshape(B, H, T, half, 2)
+  return Tensor.stack(x_pairs[..., 0] * cos - x_pairs[..., 1] * sin,
+                      x_pairs[..., 0] * sin + x_pairs[..., 1] * cos, dim=-1).reshape(B, H, T, Hd)
 
 class TransformerBlock:
   def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int, norm_eps:float, max_context:int=0):

tinygrad/codegen/__init__.py

@@ -12,12 +12,12 @@ from tinygrad.codegen.quantize import pm_quant
 from tinygrad.codegen.gpudims import pm_add_gpudims
 from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing
 from tinygrad.uop.decompositions import get_late_rewrite_patterns
-from tinygrad.codegen.late.expander import migrate_indexing, expander
+from tinygrad.codegen.late.expander import migrate_indexing, expander, pm_pre_expander
 from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
   ReduceContext, correct_load_store, pm_render
 from tinygrad.codegen.late.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
-from tinygrad.codegen.opt import pm_get_optimization, pm_do_optimize
 from tinygrad.codegen.opt.swizzler import view_left, view_right, fix_kernel_ops
+from tinygrad.codegen.opt.postrange import pm_postrange_opt
 from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen
 
 @dataclass
@@ -55,21 +55,18 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
   # view pushing
   ret.extend(rewrites_for_views)
 
-  # this is kernel.py
-  ret.append(RewriteStep(pm_get_optimization, ctx=lambda _: opts, name="get optimization"))
-  ret.append(RewriteStep(pm_do_optimize, ctx=lambda _: opts, name="optimize ast"))
-
+  # lowerer first
   if _QUANTIZE and opts.device in {"CPU", "DSP"}: ret.append(RewriteStep(pm_quant, name="quantize"))
   ret.append(RewriteStep(pm_lowerer, get_index, name="lowerer", bottom_up=True))
 
+  # optimize (schedule) the AST
+  ret.append(RewriteStep(pm_postrange_opt, ctx=lambda _: opts, name="post optimize ast"))
+
   # ** expander (expand_rewrite) **
   ret.append(RewriteStep(sym+migrate_indexing, name="initial symbolic"))
 
-  # add gpu dims (late). this also handles UNROLL range
-  ret.append(RewriteStep(pm_add_gpudims, lambda _: opts, name="add gpudims"))
-
   # expand
-  ret.append(RewriteStep(sym+expander, name="expander"))
+  ret.append(RewriteStep(sym+pm_pre_expander+expander, name="expander"))
 
   # add locals
   ret.append(RewriteStep(pm_add_buffers_local+rangeify_codegen, name="add local buffers"))
@@ -78,6 +75,9 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
   # remove reduce
   ret.append(RewriteStep(pm_reduce+gep_pushing, lambda _: ReduceContext(), name="remove_reduce"))
 
+  # add gpu dims (late). this works after devectorize, but it's faster here
+  ret.append(RewriteStep(pm_add_gpudims, lambda _: opts, name="add gpudims"))
+
   # devectorize (TODO: does this need opts?)
   if _DEVECTORIZE >= 2: pm_devectorize = sym+load_store_folding+load_store_indexing
   elif _DEVECTORIZE: pm_devectorize = sym+devectorize+load_store_folding+correct_load_store+load_store_indexing

tinygrad/codegen/gpudims.py

@@ -1,7 +1,7 @@
-import math, functools, operator
-from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType
-from tinygrad.helpers import all_int, partition, flatten, prod, dedup
-from tinygrad.dtype import dtypes, AddrSpace
+import math
+from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType, sint_to_uop
+from tinygrad.helpers import all_int, dedup
+from tinygrad.dtype import dtypes
 from tinygrad.shape.view import get_contraction
 from tinygrad.renderer import Renderer
 
@@ -34,7 +34,7 @@ def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|No
   if max_sizes is not None and len(limited) > len(max_sizes): raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
   # try to split up dims: (a,) -> (b, c)
   if limited == dims: limited = _split_dims(dims, max_sizes) if max_sizes is not None else dims
-  ret = raw_idxs = [UOp(Ops.SPECIAL, dtypes.int, (), (f"{prefix}{i}", s)) for i,s in enumerate(limited)]
+  ret = raw_idxs = [UOp(Ops.SPECIAL, dtypes.int, (sint_to_uop(s),), (f"{prefix}{i}")) for i,s in enumerate(limited)]
   if len(limited) < len(dims):
     ret = []
     if (contraction:=get_contraction(dims, limited)) is None: raise AssertionError(f"get_contraction should not be None {dims=} {limited=}")
@@ -56,17 +56,17 @@ def add_gpudims(ctx:Renderer, s:UOp):
   if any(x.op is Ops.SPECIAL for x in s_topo): return None
 
   # get ranges
-  all_ranges = {x.arg[0]%1000:x for x in s_topo if x.op is Ops.RANGE}
+  all_ranges = {x.arg[0:-1]:x for x in s_topo if x.op is Ops.RANGE}
 
   # extract global/local dims
-  global_dims = sorted(dedup([x.arg[0]%1000 for x in all_ranges.values() if x.arg[1] is AxisType.GLOBAL]))
-  local_dims = sorted(dedup([x.arg[0]%1000 for x in all_ranges.values() if x.arg[1] in (AxisType.LOCAL, AxisType.GROUP_REDUCE)]))
+  global_dims = sorted(dedup([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] is AxisType.GLOBAL]))
+  local_dims = sorted(dedup([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] in (AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE)]))
   if not global_dims and not local_dims: return None
 
   # get global and local shape
   ranges = [all_ranges[r] for r in global_dims+local_dims if r in all_ranges]
-  global_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg[0]%1000 in global_dims])
-  local_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg[0]%1000 in local_dims])
+  global_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg[0:-1] in global_dims])
+  local_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg[0:-1] in local_dims])
 
   # get the idxs
   ki: KernelInfo = s.arg
@@ -82,54 +82,13 @@ def add_gpudims(ctx:Renderer, s:UOp):
   for r in s_topo:
     if r.op is not Ops.RANGE: continue
     try:
-      ii = (global_dims+local_dims).index(r.arg[0]%1000)
+      ii = (global_dims+local_dims).index(r.arg[0:-1])
       if r.arg[1] == AxisType.REDUCE: continue
       subs[r] = idxs[ii]
     except ValueError: continue
   return s.substitute(subs)
 
-def fix_reduce_unroll(x:UOp):
-  reduce_range, reduce_expand = partition(x.src[1:], lambda y: y.op is Ops.RANGE)
-  if len(reduce_expand) == 0: return None
-  reduce_expand = [x for x in reduce_expand if x.op is not Ops.CONST]
-  assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand}"
-  ret = x.src[0]
-  if len(contract_axis:=flatten(x.arg for x in reduce_expand)):
-    ret = UOp(Ops.CONTRACT, x.dtype.vec(prod(x[1] for x in contract_axis)), (ret,), tuple(contract_axis), tag=1)
-  # REDUCE supports both "horizontal" reduction and range reduction. the horizontal elements are taken in the nearest group
-  return x.replace(src=(ret,)+tuple(reduce_range))
-
-def fix_store_unroll(x:UOp):
-  store_expand, store_range = partition(x.src[2:], lambda y: y.op is Ops.UNROLL)
-  if len(store_expand) == 0: return None
-  return UOp(Ops.CONTRACT, dtypes.void, (x.replace(src=x.src[:2]+tuple(store_range)),), tuple(flatten(x.arg for x in store_expand)), tag=1)
-
-def fix_group_for_reduce(x:UOp):
-  reduce_gfr, reduce_r = partition(x.src[1:], lambda u: u.op is Ops.RANGE and u.arg[1] == AxisType.GROUP_REDUCE)
-  if len(reduce_gfr) == 0: return None
-
-  # NOTE: if there's other locals here, we need them in the buffer too
-  upstream_locals = [u for u in x.toposort() if u.op is Ops.RANGE and u.arg[1] == AxisType.LOCAL]
-
-  # do only the non grouped reduces early
-  ret = x.replace(src=(x.src[0],)+tuple(reduce_r))
-  reduce_loop = [x.replace(arg=(x.arg[0]+100, AxisType.REDUCE)) for x in reduce_gfr]
-  buf = ret.bufferize(*upstream_locals, *reduce_gfr, arg=(AddrSpace.LOCAL, reduce_gfr[0].arg[0])).index(*upstream_locals, *reduce_loop)
-
-  # gate with an if on the store + do the final reduce
-  buf = UOp(Ops.IF, dtype=buf.dtype, src=(functools.reduce(operator.and_, [x.eq(0) for x in reduce_gfr]), buf))
-  return buf.reduce(*reduce_loop, arg=x.arg)
-
 pm_add_gpudims = PatternMatcher([
   # add gpudims must be last
   (UPat(Ops.SINK, name="s"), add_gpudims),
-  # rewrite UPCAST/UNROLL range to something to be expanded
-  (UPat(Ops.RANGE, name="r"),
-   lambda r: UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(s:=r.vmax+1), tuple(range(s))),), ((r.arg[0],s),)) \
-    if r.arg[1] in {AxisType.UNROLL, AxisType.UPCAST} else None),
-  # fix REDUCEs with UNROLLs
-  (UPat(Ops.REDUCE, name="x"), fix_reduce_unroll),
-  (UPat(Ops.STORE, name="x"), fix_store_unroll),
-  # fix group for reduce
-  (UPat(Ops.REDUCE, name="x"), fix_group_for_reduce),
 ])

tinygrad/codegen/late/devectorizer.py

@@ -2,9 +2,9 @@ from typing import Any, cast
 import functools, operator, itertools
 from collections import defaultdict
 from dataclasses import dataclass
-from tinygrad.dtype import dtypes, ImageDType, PtrDType, DType, AddrSpace
+from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace
 from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, graph_rewrite, GroupOp, identity_element
-from tinygrad.uop.symbolic import split_uop, uop_given_valid, parse_valid, simplify_valid, sym, symbolic_flat
+from tinygrad.uop.symbolic import uop_given_valid, parse_valid, simplify_valid, sym, symbolic_flat
 from tinygrad.helpers import getenv, flatten, AMX, prod, partition
 from tinygrad.renderer import Renderer
 
@@ -19,13 +19,13 @@ def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
 
   # can drop valid if idx is out of bound when valid is False
   drop_stmt = []
-  for stmt in split_uop(valid, Ops.AND):
+  for stmt in valid.split_uop(Ops.AND):
     try: X, is_upper_bound, c = parse_valid(stmt)
     except ValueError: return None
 
     # for X0 + X1 + ... >= 1, check if it's out of bound when Xi = 0 for all i
-    if not is_upper_bound and c == 1 and all(u.op in GroupOp.Irreducible and u.vmin == 0 for u in split_uop(X, Ops.ADD)):
-      testidx = functools.reduce(lambda nowidx,u: nowidx.substitute({u:u.const_like(0)}), split_uop(X, Ops.ADD), idx)
+    if not is_upper_bound and c == 1 and all(u.op in GroupOp.Irreducible and u.vmin == 0 for u in X.split_uop(Ops.ADD)):
+      testidx = functools.reduce(lambda nowidx,u: nowidx.substitute({u:u.const_like(0)}), X.split_uop(Ops.ADD), idx)
       testidx = testidx.simplify()
       if testidx.gep(0).vmax < 0 or testidx.gep(1).vmax < 0:
         drop_stmt.append(stmt)
@@ -42,7 +42,7 @@ def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
           break
 
   if not drop_stmt and idx is start_idx: return None
-  new_valid = functools.reduce(operator.and_, ss) if (ss:=[s for s in split_uop(valid, Ops.AND) if s not in drop_stmt]) else None
+  new_valid = functools.reduce(operator.and_, ss) if (ss:=[s for s in valid.split_uop(Ops.AND) if s not in drop_stmt]) else None
   return buf.index(idx, new_valid)
 
 def delete_redundant_gates(store:UOp, buf:UOp, idx:UOp, val:UOp, store_gate:UOp, cast:UOp|None=None) -> UOp|None:
@@ -80,9 +80,6 @@ def expand_index(buf:UOp, vec:UOp, mask:UOp|None=None):
     if len(midx.src[i].src) == 3: root_src = (midx.src[i].src[2], root_src)
     offsets_rootsrc[root_src].setdefault(arg, []).append(i)
 
-  # the buf.dtype is always a pointer
-  ptrdtype = cast(PtrDType, buf.dtype)
-
   # then rewrite everything we can into groups
   ret = []
   idxs: list[int|None] = [None]*vec.dtype.count
@@ -92,7 +89,7 @@ def expand_index(buf:UOp, vec:UOp, mask:UOp|None=None):
     for grp in grouped_offsets:
       # get the index offset for this element. using [0] is okay, because they are the same
       lidx = midx.src[offsets[grp[0]][0]]
-      if len(grp) > 1: lidx = lidx.cast(ptrdtype.base.vec(len(grp)).ptr(size=ptrdtype.size, addrspace=ptrdtype.addrspace))
+      if len(grp) > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(len(grp)).ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace))
       # set the idxs of the output
       for i,g in enumerate(grp):
         for oo in offsets[g]: idxs[oo] = global_offset+i
@@ -101,7 +98,7 @@ def expand_index(buf:UOp, vec:UOp, mask:UOp|None=None):
       global_offset += len(grp)
   assert None not in idxs, f"some idxs are missing {idxs}"
   # this base thing is for image, we want the CAT to be a normal pointer
-  post_cat = UOp(Ops.PTRCAT, ptrdtype.base.ptr(size=ptrdtype.size, addrspace=ptrdtype.addrspace).vec(vec.dtype.count), tuple(ret))
+  post_cat = UOp(Ops.PTRCAT, buf.ptrdtype.base.ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace).vec(vec.dtype.count), tuple(ret))
   return post_cat.gep(tuple(cast(list[int], idxs)))
 
 def cat_after_store(cat:UOp, data:UOp, sto:UOp):
@@ -154,7 +151,7 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
     must_divide = False
   elif buf.dtype.base != dtypes.float and buf.dtype.base != dtypes.half and not isinstance(buf.dtype, ImageDType):
     pass
-  elif cast(PtrDType, buf.dtype).addrspace == AddrSpace.REG:
+  elif buf.ptrdtype.addrspace == AddrSpace.REG:
     pass
   elif isinstance(buf.dtype, ImageDType):
     lengths = [4]
@@ -169,13 +166,12 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
   # split based on the fold lengths
   global_offset = 0
   ret = []
-  ptrdtype = cast(PtrDType, buf.dtype)
   while global_offset < sz:
     # with 1 at the end of the lengths list, this will always hit
     for fold_length in lengths:
       if global_offset+fold_length > sz: continue
       lidx = buf.index(idx.src[1] + global_offset, idx.src[2] if len(idx.src) > 2 else None)
-      if fold_length > 1: lidx = lidx.cast(ptrdtype.base.vec(fold_length).ptr(size=ptrdtype.size, addrspace=ptrdtype.addrspace))
+      if fold_length > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(fold_length).ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace))
       if ls.op is Ops.STORE: ret.append(ls.replace(src=(lidx,ls.src[1].gep(tuple(range(global_offset, global_offset+fold_length))))+ls.src[2:]))
       else: ret.append(ls.replace(src=(lidx,)+ls.src[1:], dtype=ls.dtype.scalar().vec(fold_length)))
       global_offset += fold_length
@@ -233,8 +229,7 @@ def no_vectorized_alu(alu:UOp):
   return UOp(Ops.VECTORIZE, alu.dtype, alus)
 
 def no_vectorized_buf(buf:UOp):
-  dtype = cast(PtrDType, buf.dtype)
-  return buf.replace(dtype=dtype.base.scalar().ptr(dtype.size*dtype.count, dtype.addrspace)).cast(dtype)
+  return buf.replace(dtype=buf.ptrdtype.base.scalar().ptr(buf.ptrdtype.size*buf.ptrdtype.count, buf.ptrdtype.addrspace)).cast(buf.dtype)
 
 def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp):
   cnt = cast.dtype.count

tinygrad/codegen/late/expander.py

@@ -1,9 +1,9 @@
 # this converts a lowerer program into a vectorized program
 
 import functools, itertools, operator
-from tinygrad.dtype import dtypes, PtrDType
-from tinygrad.helpers import AMX, dedup, flatten, all_same, prod
-from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp
+from tinygrad.dtype import dtypes, PtrDType, AddrSpace
+from tinygrad.helpers import AMX, dedup, flatten, all_same, prod, partition
+from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, AxisType
 
 def _expand_arg_to_idx(args:tuple[tuple[int, int], ...], rpk:dict[int, int]) -> int:
   idx, mul = 0, 1
@@ -114,3 +114,49 @@ migrate_indexing = PatternMatcher([
   # create gate MUST BE BEFORE expander
   (UPat(Ops.STORE, name="root"), create_gate),
 ])
+
+# ****
+
+def fix_reduce_unroll(x:UOp):
+  reduce_range, reduce_expand = partition(x.src[1:], lambda y: y.op is Ops.RANGE)
+  if len(reduce_expand) == 0: return None
+  reduce_expand = [x for x in reduce_expand if x.op is not Ops.CONST]
+  assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand}"
+  ret = x.src[0]
+  if len(contract_axis:=flatten(x.arg for x in reduce_expand)):
+    ret = UOp(Ops.CONTRACT, x.dtype.vec(prod(x[1] for x in contract_axis)), (ret,), tuple(contract_axis), tag=1)
+  # REDUCE supports both "horizontal" reduction and range reduction. the horizontal elements are taken in the nearest group
+  return x.replace(src=(ret,)+tuple(reduce_range))
+
+def fix_store_unroll(x:UOp):
+  store_expand, store_range = partition(x.src[2:], lambda y: y.op is Ops.UNROLL)
+  if len(store_expand) == 0: return None
+  return UOp(Ops.CONTRACT, dtypes.void, (x.replace(src=x.src[:2]+tuple(store_range)),), tuple(flatten(x.arg for x in store_expand)), tag=1)
+
+def fix_group_for_reduce(x:UOp):
+  reduce_gfr, reduce_r = partition(x.src[1:], lambda u: u.op is Ops.RANGE and u.arg[1] == AxisType.GROUP_REDUCE)
+  if len(reduce_gfr) == 0: return None
+
+  # NOTE: if there's other locals here, we need them in the buffer too
+  upstream_locals = [u for u in x.toposort() if u.op is Ops.RANGE and u.arg[1] == AxisType.LOCAL]
+
+  # do only the non grouped reduces early
+  ret = x.replace(src=(x.src[0],)+tuple(reduce_r))
+  reduce_loop = [x.replace(arg=(x.arg[0]+100, AxisType.REDUCE)) for x in reduce_gfr]
+  buf = ret.bufferize(*upstream_locals, *reduce_gfr, arg=(AddrSpace.LOCAL, reduce_gfr[0].arg[0])).index(*upstream_locals, *reduce_loop)
+
+  # gate with an if on the store + do the final reduce
+  buf = UOp(Ops.IF, dtype=buf.dtype, src=(functools.reduce(operator.and_, [x.eq(0) for x in reduce_gfr]), buf))
+  return buf.reduce(*reduce_loop, arg=x.arg)
+
+pm_pre_expander = PatternMatcher([
+  # rewrite UPCAST/UNROLL range to something to be expanded
+  (UPat(Ops.RANGE, name="r"),
+   lambda r: UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(s:=r.vmax+1), tuple(range(s))),), ((r.arg[0],s),)) \
+    if r.arg[1] in {AxisType.UNROLL, AxisType.UPCAST} else None),
+  # fix REDUCEs with UNROLLs
+  (UPat(Ops.REDUCE, name="x"), fix_reduce_unroll),
+  (UPat(Ops.STORE, name="x"), fix_store_unroll),
+  # fix group for reduce
+  (UPat(Ops.REDUCE, name="x"), fix_group_for_reduce),
+])

tinygrad/codegen/late/linearize.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 import heapq
 from collections import defaultdict
 from dataclasses import dataclass, replace
-from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp
+from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, BottomUpGate
 from tinygrad.helpers import dedup, all_same, flatten, BLOCK_REORDER
 
 # NOTE: any toposort should be valid here, unlike last time this isn't required, it's just for speed
@@ -76,12 +76,13 @@ class BlockContext:
   def from_sink(sink:UOp) -> BlockContext:
     # get children and all block contexts
     ctx = BlockContext({}, {}, {})
-    for u in sink.toposort():
+    for u in sink.toposort(gate=lambda u:u.op is not Ops.SPECIAL):
       this_block_ctx: list[UOp] = []
       ctx.child_count[u] = 0
 
       # get children and accumulate the last_ctx
       for s in u.src:
+        if s.op is Ops.SPECIAL: continue
         # NOTE: if a parent appears multiple times in the src, it counts multiple times as a child
         ctx.child_count[s] += 1
         this_block_ctx += ctx.last_ctx(s)
@@ -142,7 +143,7 @@ def make_block_bottom_up(ctx:BlockContext, x:UOp):
 
   # add unmergables to sources
   srcs = []
-  for u,cnt in unmergable.items(): srcs += [add_blockends(u, ctx.block_ctxs[u], current_ctx, cnt=cnt)]*cnt
+  for u,cnt in unmergable.items(): srcs += [add_blockends(u, ctx.block_ctxs.get(u,()), current_ctx, cnt=cnt)]*cnt
 
   # add blockseeds, with blockends as needed
   for (new_ctx, new_child_ctx), v in blockseeds.items():
@@ -154,8 +155,12 @@ def make_block_bottom_up(ctx:BlockContext, x:UOp):
   bb = BasicBlock(tuple(lst), ctx=current_ctx, cnt=child_count, child_ctx=child_ctx)
   return UOp(Ops.BLOCK, src=tuple(srcs), arg=bb)
 
+# we prevent the source of the SPECIAL from being linearized since its not part of the kernel
+def raise_bottom_up_gate(): raise BottomUpGate()
+
 block_create = PatternMatcher([
   (UPat(GroupOp.All-DONT_PLACE_IN_BLOCK.union({Ops.BLOCK, Ops.BLOCKEND}), name="x"), make_block_bottom_up),
+  (UPat(Ops.SPECIAL), raise_bottom_up_gate)
 ])
 
 # ***** blockend merging ****

tinygrad/codegen/lowerer.py

@@ -1,7 +1,6 @@
 # the job of the lowerer is to do indexing
 from dataclasses import dataclass
-from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint_to_uop, AxisType, graph_rewrite
+from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint_to_uop, AxisType, graph_rewrite, resolve
 
 # ***** indexing *****
 
@@ -12,12 +11,12 @@ class IndexContext:
   start: int = 0
 
 def shape_to_idx(s, axis_types, start=0):
-  return [UOp.range(dtypes.int, sint_to_uop(s), start+i, axistype=at) for i, (s, at) in enumerate(zip(s, axis_types))]
+  return [UOp.range(sint_to_uop(s), start+i, at) for i, (s, at) in enumerate(zip(s, axis_types))]
 
 def get_index(ast:UOp) -> IndexContext:
   axis_types = ast.arg.axis_types if isinstance(ast.arg, KernelInfo) else ()
-  if len(ast.full_shape) != len(axis_types):
-    axis_types = tuple([AxisType.REDUCE if s is not fs else AxisType.LOOP for s,fs in zip(ast.shape, ast.full_shape)])
+  if len(ast.full_shape) != len(axis_types) and ast.st is not None:
+    axis_types = tuple([AxisType.REDUCE if resolve(s != fs) else AxisType.LOOP for s,fs in zip(ast.shape, ast.full_shape)])
   return IndexContext(axis_types, [], 0)
 
 # ***** lowering (given index) *****

tinygrad/codegen/opt/__init__.py

@@ -1,46 +1,26 @@
 # opt opinionatedly transforms an ast into an optimized ast using either heuristics or beam search
+from __future__ import annotations
+from enum import Enum, auto
+from dataclasses import dataclass
+from tinygrad.uop.ops import AxisType
 
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
-from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, KernelInfo
-from tinygrad.helpers import NOOPT, BEAM, USE_TC, getenv
-from tinygrad.renderer import Renderer
-from tinygrad.uop.spec import type_verify
+class OptOps(Enum):
+  TC = auto(); UPCAST = auto(); UNROLL = auto(); LOCAL = auto() # noqa: E702
+  GROUP = auto(); GROUPTOP = auto(); NOLOCALS = auto(); PADTO = auto(); SWAP = auto() # noqa: E702
+  def __lt__(self, x:OptOps): return self.value < x.value
 
-def get_optimized_ast(ast:UOp, renderer:Renderer) -> UOp:
-  """
-  Optimize an AST based on heuristics or BEAM search.
+@dataclass(frozen=True, order=True)
+class Opt:
+  op: OptOps
+  axis: int|None = None
+  arg: int|tuple|None = None
+  def __repr__(self): return f"Opt(op={self.op}, axis={self.axis}, arg={self.arg})"
 
-  Args:
-    ast: The Ops.SINK rooted AST
-    renderer: The renderer used to generate the code
+axis_letters = {AxisType.GLOBAL: "g", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u",
+                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
+axis_colors = {AxisType.GLOBAL: "blue", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE", AxisType.UPCAST: "yellow",
+               AxisType.GROUP_REDUCE: "green", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}
 
-  Returns:
-    The Ops.SINK rooted AST transformed to apply the opts and with a KernelInfo in the arg.
-  """
-
-  assert ast.arg is None, "no opt if there's an arg"
-  k = Kernel(ast, opts=renderer)
-  if not NOOPT:
-    if not k.apply_tensor_cores(USE_TC.value): k.apply_opts(hand_coded_optimizations(k))
-    if BEAM >= 1:
-      from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
-      kb = Kernel(ast, opts=renderer)
-      rawbufs = bufs_from_lin(kb, allocate=False)
-      k = beam_search(kb, rawbufs, BEAM.value, bool(getenv("BEAM_ESTIMATE", 1)))
-  return ast.replace(arg=KernelInfo(opts_to_apply=tuple(k.applied_opts)))
-
-pm_get_optimization = PatternMatcher([
-  (UPat(Ops.SINK, name="ast"), lambda ctx,ast: get_optimized_ast(ast, ctx) if ast.arg is None and ast.src[0].st is not None else None),
-])
-
-def apply_opt(ast:UOp, renderer:Renderer):
-  k = Kernel(ast, opts=renderer)
-  k.apply_opts(ast.arg.opts_to_apply)
-  ret = k.get_optimized_ast()
-  if __debug__: type_verify(list(ret.toposort()))
-  return ret
-
-pm_do_optimize = PatternMatcher([
-  (UPat(Ops.SINK, name="ast"), lambda ctx,ast: apply_opt(ast, ctx) if ast.arg is not None and ast.arg.opts_to_apply is not None else None),
-])
+class KernelOptError(Exception): pass
+def check(cond:bool, msg:str=""):
+  if not cond: raise KernelOptError(msg)

tinygrad/codegen/opt/heuristic.py

@@ -1,10 +1,50 @@
 import itertools
-from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps, KernelOptError, AxisType
-from tinygrad.helpers import getenv, DEBUG, prod, NOLOCALS
+from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
+from tinygrad.helpers import getenv, DEBUG, prod, NOLOCALS, TC_OPT, TC_SELECT, USE_TC, AMX
 from tinygrad.dtype import ImageDType
-from tinygrad.uop.ops import Ops, resolve
+from tinygrad.uop.ops import Ops, resolve, AxisType
+from tinygrad.codegen.opt.postrange import Scheduler
+
+def hand_coded_optimizations(k:Scheduler) -> list[Opt]:
+  # first try the tensor cores
+  """ Attempts to apply a tensor core optimization to the kernel. If one exists and applies properly, return true, otherwise return false.
+  Tensor cores are optimized instructions that matrix multiply-accumulate across a wave of threads: D(M, N) = A(M, K) * B(K, N) + C(M, N).
+
+  Keyword arguments:
+  use_tensor_cores -- controls how tensor cores are applied (default 1)
+    0: will disable any tensor core matching
+    1: enable tensor cores
+    2: apply tensor core shape but don't use UOp.WMMA
+  extra_opts -- additional Opt's to apply after the tensor core instead of the hand-coded additional Opt's (default None)
+  tc_select -- specifies which tensor core(s) to use for optimization (default -1)
+    -1: iterates through all available tensor cores in order and uses the first one that matches the requirements (dims and dtypes)
+    [0-N]: uses only the n'th tensor core available; useful for search
+  tc_opt -- controls which kinds of kernels may be eligible for tensor cores application (default 2 during BEAM, 0 otherwise)
+    0: applies to only kernels with a single reduce axis and direct Ops.LOAD into Ops.MUL
+    1: allows kernels with multiple reduce axes and also multiplication of Ops.CAST'd buffers
+    2: allows kernels with M, N, K axes that are not multiples of the tensor core dimensions by applying padding those axes as needed
+  """
+  # NOTE: unless TC_OPT is > 0, we only trigger tensor cores if there's only one reduce axis
+  if USE_TC > 0 and (len(k.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE)) == 1 or (TC_OPT.value >= 1)):
+    good_tc_opt = False
+    try: # check TC first and apply hand-coded opts if successful
+      tk = k.copy()
+      rngs = tk.apply_opt(Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, USE_TC.value)))
+      good_tc_opt = True
+    except KernelOptError:
+      pass
+    if good_tc_opt:
+      # skip hand-coded TC opts if AMX, upcasting will make kernel slower
+      if rngs is not None and not AMX:
+        for tc_dim in [1,0]: # attempt to upcast M and N
+          szs = [sz for sz in [5,4,3,2] if rngs[tc_dim].src[0].divides(sz) is not None]
+          if szs:
+            # set it to the replaced range
+            rngs[tc_dim] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[tc_dim]), szs[0]))[0]
+        if (szs := [sz for sz in [4,2] if rngs[0].src[0].divides(sz) is not None]): # attempt to local N
+          tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[0]), szs[0]))
+      return tk.applied_opts
 
-def hand_coded_optimizations(k:Kernel) -> list[Opt]:
   # make a copy so it does not mutate the input
   k = k.copy()
 
@@ -13,15 +53,13 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
   if k.opts.has_local and getenv("MV",1) != 0 and (MV_BLOCKSIZE > 1 or MV_THREADS_PER_ROW > 1 or MV_ROWS_PER_THREAD > 1) and  \
     k.reduceop is not None and k.reduceop.arg[0] is Ops.ADD and len(k.full_shape) >= 2 and k.opts.has_shared and \
     (mulop:=k.reduceop.src[0]).op is Ops.MUL and mulop.src[0].op is Ops.LOAD and mulop.src[1].op is Ops.LOAD:
-    st0, st1 = k.sts[k.bufs.index(mulop.src[0])], k.sts[k.bufs.index(mulop.src[1])]
-    strides0, strides1 = st0.real_strides(), st1.real_strides()
-    def has_expanded_axis(shape, strides): return any(resolve(s > 1) and not resolve(st != 0) for s,st in zip(shape,strides))
-    if strides0[first_reduce:=(k.axes_of(AxisType.REDUCE)[0])] == 1 and \
-      not (has_expanded_axis(st0.shape, strides0) and has_expanded_axis(st1.shape, strides1)):
+    idx0, idx1 = mulop.src[0].src[0].src[1], mulop.src[1].src[0].src[1]
+    first_reduce_rng = k.ranges_of(AxisType.REDUCE)[0]
+    if any(u is first_reduce_rng for u in idx0.split_uop(Ops.ADD)) and all(r in idx1.ranges for r in idx0.ranges):
       for global_idx in k.axes_of(AxisType.GLOBAL):
-        if k.full_shape[first_reduce]%MV_THREADS_PER_ROW == 0 and k.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0:
+        if first_reduce_rng.src[0].divides(MV_THREADS_PER_ROW) is not None and k.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0:
           if DEBUG >= 3:
-            print(f"MATVEC: {k.full_shape=} {first_reduce=} {strides0=} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}")
+            print(f"MATVEC: {k.full_shape=} {first_reduce_rng.render()} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}")
           if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW))
           if MV_BLOCKSIZE > 1: k.apply_opt(Opt(OptOps.LOCAL, global_idx, MV_BLOCKSIZE))
           if MV_ROWS_PER_THREAD > 1: k.apply_opt(Opt(OptOps.UPCAST, global_idx, MV_ROWS_PER_THREAD))
@@ -38,7 +76,9 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
   # upcast float4 images
   for buf_index,buf in enumerate(k.bufs):
     if isinstance(buf.src[0].dtype, ImageDType):
-      if (unit_stride_axes_mul_4 := [i for i in k.sts[buf_index].unit_stride_axes(ignore_valid=True) if k.sts[buf_index].shape[i]%4 == 0]):
+      # part of real_strides
+      unit_stride_axes_mul_4 = [k.rngs.index(c) for c in k.bufs[buf_index].src[1].split_uop(Ops.ADD) if c.op is Ops.RANGE and (c.vmax+1)%4 == 0]
+      if len(unit_stride_axes_mul_4):
         if (axis:=unit_stride_axes_mul_4[0]) in k.upcastable_dims:
           k.apply_opt(Opt(OptOps.UPCAST, axis, 4))
         elif axis in k.unrollable_dims:
@@ -53,8 +93,10 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
   to_upcast: list[int] = []
   # upcast leading axes first (hack-ish for winograd; we actually want to upcast masked axes with low stride first)
   for axis in k.upcastable_dims:
-    if k.full_shape[axis] <= 7 and any(st.axis_is_masked(axis) for st in k.sts) and \
-      prod(k.full_shape[j] for j in to_upcast) * k.full_shape[axis] <= 7 * 7:
+    # for Schedule, we check if the range is used in INDEX gates or WHERE gates
+    is_masked = any(len(st.src) > 2 and k.rngs[axis] in st.src[2].parents for st in k.bufs) or \
+                any(any(o is k.rngs[axis] for o in u.src[0].parents) for u in k.ast.parents if u.op is Ops.WHERE)
+    if k.full_shape[axis] <= 7 and is_masked and prod(k.full_shape[j] for j in to_upcast) * k.full_shape[axis] <= 7 * 7:
       if DEBUG >= 4: print(f"upcasting masked axis : {axis}")
       to_upcast.append(axis)
   for axis in to_upcast[::-1]: k.apply_opt(Opt(OptOps.UPCAST, axis, 0))
@@ -68,10 +110,16 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
     for axis, upcast_amount in itertools.product(k.upcastable_dims, ([128] if not len(upcasted_axis) else []) if is_dsp else [3,4]):
       # if we haven't upcasted it, it mods, and buffer has stride 0 on axis while having no stride 0 in the upcasted axis already
       if axis in upcasted_axis or k.full_shape[axis]%upcast_amount != 0: continue
-      if any(st.views[-1].strides[axis] == 0 and \
-             all(x != 0 for t,x in zip(k.axis_types, st.real_strides()) if t in (AxisType.UPCAST, AxisType.UNROLL)) for st in k.sts):
-        xb_choices.append((sum(st.views[-1].strides[axis]>0 for st in k.sts),
-                           sum(st.views[-1].strides[axis] for st in k.sts), axis, upcast_amount))
+      rng = k.rngs[axis]
+      if any(rng not in b.src[1].parents and all(r2 in b.src[1].parents for r2 in k.ranges_of(AxisType.UPCAST, AxisType.UNROLL)) for b in k.bufs):
+        num_strides, sum_strides = 0, 0
+        for b in k.bufs:
+          if rng in b.src[1].parents: num_strides += 1
+          for c in b.src[1].split_uop(Ops.ADD):
+            if c is rng: sum_strides += 1
+            if c.op is Ops.MUL and c.src[0] is rng and c.src[1].op is Ops.CONST: sum_strides += c.src[1].arg
+            if c.op is Ops.MUL and c.src[1] is rng and c.src[0].op is Ops.CONST: sum_strides += c.src[0].arg
+        xb_choices.append((num_strides, sum_strides, axis, upcast_amount))
     if xb_choices:
       xb_choices = sorted(xb_choices)
       if DEBUG >= 4: print(f"more upcast axis : {xb_choices}")
@@ -109,7 +157,8 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
       k.apply_opt(Opt(OptOps.NOLOCALS))
     else:
       # prioritize making expand axes local
-      local_axis_ranking = [(any(st.views[-1].strides[axis] == 0 for st in k.sts), axis) for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP)]
+      local_axis_ranking = [(any(k.rngs[axis] not in b.src[1].parents for b in k.bufs), axis) \
+                              for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP) if k.rngs[axis].src[0].op is Ops.CONST]
       to_local: list[tuple[int, int]] = []
       for _, axis in sorted(local_axis_ranking, key=lambda x: (-x[0], -x[1])):
         local_size = prod(sz for _, sz in to_local)

tinygrad/codegen/opt/kernel.py

@@ -1,496 +0,0 @@
-from __future__ import annotations
-import itertools, functools, math
-from dataclasses import dataclass
-from collections import defaultdict
-from typing import cast, Final, Callable, Sequence
-from enum import Enum, auto
-
-from tinygrad.uop.ops import GroupOp, KernelInfo, UOp, Ops, can_pad, resolve, Variable, sint, graph_rewrite, AxisType
-from tinygrad.uop.spec import type_verify, ast_spec
-from tinygrad.device import Device
-from tinygrad.codegen.opt.tc import TensorCore
-from tinygrad.renderer import Renderer
-from tinygrad.dtype import ImageDType
-from tinygrad.helpers import all_same, colored, ansilen, dedup, prod, round_up, to_function_name, unwrap, argfix, DEBUG, TC_SELECT, TC_OPT, AMX
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import strides_for_shape, get_contraction
-from tinygrad.codegen.opt.swizzler import view_left, view_left_through_load
-
-class OptOps(Enum):
-  TC = auto(); UPCAST = auto(); UNROLL = auto(); LOCAL = auto() # noqa: E702
-  GROUP = auto(); GROUPTOP = auto(); NOLOCALS = auto(); PADTO = auto(); SWAP = auto() # noqa: E702
-  def __lt__(self, x:OptOps): return self.value < x.value
-
-@dataclass(frozen=True, order=True)
-class Opt:
-  op: OptOps
-  axis: int|None = None
-  arg: int|tuple|None = None
-  def __repr__(self): return f"Opt(op={self.op}, axis={self.axis}, arg={self.arg})"
-
-axis_letters = {AxisType.GLOBAL: "g", AxisType.LOCAL: "l", AxisType.LOOP: "L", AxisType.UPCAST: "u",
-                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
-axis_colors = {AxisType.GLOBAL: "blue", AxisType.LOCAL: "cyan", AxisType.LOOP: "WHITE", AxisType.UPCAST: "yellow",
-               AxisType.GROUP_REDUCE: "green", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}
-
-class KernelOptError(Exception): pass
-def check(cond:bool, msg:str=""):
-  if not cond: raise KernelOptError(msg)
-
-@dataclass
-class TensorCoreOptions:
-  axes: tuple[int, ...] # the location of the original N and M axes if still in the shape
-  axes_exist: tuple[bool, ...] # true if the original N and M axes are still in the shape
-  axis_pads: tuple[tuple[int, int], ...]
-  def fix_axes(self, removed_axis:int): # adjust the TC axes if necessary when a dimension is removed
-    axes, axes_exist = list(self.axes), list(self.axes_exist)
-    for tc_dim in [i for i in range(2) if axes_exist[i]]:
-      if removed_axis < axes[tc_dim]: axes[tc_dim] -= 1
-      elif removed_axis == axes[tc_dim]: axes_exist[tc_dim] = False
-    self.axes, self.axes_exist = tuple(axes), tuple(axes_exist)
-
-class Kernel:
-  def __init__(self, ast:UOp, opts:Renderer|None=None):
-    assert ast.op is Ops.SINK, ast.op
-    self.ast = ast
-
-    self.opts = opts if opts is not None else Device[Device.DEFAULT].renderer
-    # verify AST matches the spec
-    if __debug__: type_verify(list(self.ast.toposort()), ast_spec)
-
-    self.vars: list[Variable] = self.ast.variables()
-    # NOTE: this requires a specific order with the [::-1], this is likely a bug
-    self.bufs: list[UOp] = [x for x in self.ast.toposort() if x.op in GroupOp.Buffer and x.st is not None][::-1]
-
-    # create new shapetrackers inside this kernel, we will permute them
-    self.sts: list[ShapeTracker] = [x.st_arg for x in self.bufs]
-
-    # add the shapetrackers for each reduce
-    # we use this to track which axes are reduced in each reduce
-    self.reduceops = [x for x in self.ast.toposort() if x.op is Ops.REDUCE_AXIS]
-    for x in self.reduceops:
-      self.sts.append(unwrap(x.st))
-      self.sts.append(unwrap(x.src[0].st))
-
-    # add a shapetracker to the end to track the full shape, with 0 strides so it can merge
-    full_shape = ast.full_shape
-    self.sts.append(ShapeTracker.from_shape(full_shape, (0,)*len(full_shape)))
-
-    # parameters for optimization
-    self.tensor_core: TensorCore|None = None
-    self.tensor_core_opts: TensorCoreOptions|None = None
-    self.use_tensor_cores: int = 0
-    self.applied_opts: list[Opt] = []
-    self.dont_use_locals = False
-    self.finalized: bool = False
-
-    # group simplifies
-    self.simplify_ones()
-    self.simplify_merge_adjacent()
-
-    # axis types
-    global_loops = AxisType.GLOBAL if self.opts.has_local else AxisType.LOOP
-    self.axis_types: list[AxisType] = [AxisType.REDUCE if resolve(x!=y) else global_loops for x,y in zip(self.output_shape, self.full_shape)]
-
-    # confirm all reduce axes are at the end
-    if (final_reduces := [x for x in self.axis_types if x == AxisType.REDUCE]) and final_reduces != self.axis_types[-len(final_reduces):]:
-      raise RuntimeError(f"reduces are not at the end of the shape {self.full_shape} -> {self.output_shape}")
-
-  def copy(self):
-    ret = type(self).__new__(type(self))
-
-    # base linearizer params
-    ret.opts, ret.ast = self.opts, self.ast
-
-    # things downstream of the AST
-    ret.reduceops, ret.vars, ret.bufs = self.reduceops, self.vars, self.bufs
-    ret.sts = self.sts[:]
-    ret.axis_types = self.axis_types[:]
-
-    # parameters for optimizations
-    ret.applied_opts, ret.dont_use_locals = self.applied_opts[:], self.dont_use_locals
-    ret.tensor_core, ret.tensor_core_opts, ret.use_tensor_cores = self.tensor_core, self.tensor_core_opts, self.use_tensor_cores
-    ret.finalized = self.finalized
-
-    return ret
-
-  @property
-  def reduceop(self) -> UOp|None: return self.reduceops[0] if len(self.reduceops) > 0 else None
-  @property
-  def full_shape(self) -> tuple[sint, ...]: return self.sts[-1].shape
-
-  @property
-  def output_shape(self) -> tuple[sint, ...]: return self.sts[0].shape
-  @property
-  def shape_len(self) -> int: return len(self.full_shape)
-
-  def axes_of(self, *axis_type:AxisType) -> list[int]: return [i for i,t in enumerate(self.axis_types) if t in argfix(axis_type)]
-  @property
-  def upcasted(self) -> int: return len(self.axes_of(AxisType.UPCAST, AxisType.UNROLL))
-  @property
-  def group_for_reduces(self) -> int: return len(self.axes_of(AxisType.GROUP_REDUCE))
-
-  # heuristic helpers
-  @property
-  def upcastable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP) \
-                                                  if isinstance(s:=self.full_shape[i], int) and s > 1]
-  @property
-  def unrollable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE) \
-                                                  if isinstance(s:=self.full_shape[i], int) and s > 1]
-
-  # ******************** colors and names ********************
-
-  def colors(self) -> list[str]:
-    assert len(self.axis_types) == self.shape_len, "colors size mismatch"
-    return [axis_colors[x] if not self.dont_use_locals or not x == AxisType.GLOBAL else "BLUE" for x in self.axis_types]
-
-  def colored_shape(self, pad:int|None=None, dense=False) -> str:
-    shape_strs = [(s if dense else f"{s:4d}") if isinstance(s, int) else s.render() for s in self.full_shape]
-    ret = ' '.join(colored(s, color) for s,color in zip(shape_strs, self.colors()))
-    if pad: ret += ' '*(pad-ansilen(ret))
-    return ret
-
-  kernel_cnt: Final[defaultdict[str, int]] = defaultdict(int)
-  @functools.cached_property
-  def name(self) -> str:
-    # kernel name (before late upcast)
-    kernel_type = "r" if self.reduceop is not None else ("C" if all(x.op is Ops.SINK or x.op in GroupOp.Buffer for x in self.ast.toposort()) else "E")
-    suffix = colored('_', 'BLACK').join([colored(x.render() if isinstance(x, UOp) else str(x), c) for x,c in zip(self.full_shape, self.colors())])
-    name = kernel_type + (f"{len(self.ast.src)}" if len(self.ast.src) > 1 else "") + "_" + suffix
-
-    # name the function something unique
-    Kernel.kernel_cnt[(function_name := to_function_name(name))] += 1
-    num = f"n{Kernel.kernel_cnt[function_name]-1}" if Kernel.kernel_cnt[function_name] > 1 else ""
-    return name + colored(num, 'BLACK')
-
-  # ******************** base simplifiers ********************
-
-  # apply reshape and permute to all shapetrackers
-  def reshape(self, new_shape_fxn:Callable[[tuple[sint, ...]], Sequence[sint]]):
-    self.sts = [st.reshape(tuple(new_shape_fxn(st.shape))) for st in self.sts]
-  def permute(self, new_axes:Sequence[int]): self.sts = [st.permute(tuple(new_axes)) for st in self.sts]
-
-  # axis : the axis to pull from
-  # amount : the amount to take
-  # top : if you want to pull that amount from the top
-  # insert_at : place to insert the new stuff
-  def shift_to(self, axis:int, amount:int, new_type:AxisType, top:bool=False, insert_at:int|None=None) -> int:
-    if insert_at is None: insert_at = self.shape_len
-    self.axis_types.insert(insert_at, new_type)
-    move_axis = axis if top else axis+1
-    if move_axis < insert_at: insert_at += 1
-    def new_shape_fxn(x): return x[0:axis] + (((amount,x[axis]//amount) if top else (x[axis]//amount,amount)) if x[axis] > 1 else (1,1)) + x[axis+1:]
-    new_axes = [i for i in range(insert_at) if i != move_axis]+[move_axis]+[i for i in range(insert_at, self.shape_len+1) if i != move_axis]
-    self.reshape(new_shape_fxn)
-    self.permute(new_axes)
-    return insert_at
-
-  # ******************** complex simplifiers ********************
-
-  def simplify_ones(self) -> bool:
-    # remove places where the shape is all ones
-    if any(all_ones:=[s==1 for s in self.full_shape]):
-      if hasattr(self, 'axis_types'):
-        self.axis_types = [x for i,x in enumerate(self.axis_types) if not all_ones[i]]
-      self.reshape(lambda shape: [x for i,x in enumerate(shape) if not all_ones[i]])
-      return True
-    return False
-
-  def simplify_merge_adjacent(self):
-    assert not hasattr(self, 'axis_types'), "don't call this after init"
-    if self.shape_len == 0: return
-    shapes, strides = [x.shape for x in self.sts], [x.real_strides() for x in self.sts]
-    # NOTE: we can't use self.first_reduce yet
-    first_reduce = [resolve(x!=y) for x,y in zip(self.output_shape+(0,), self.full_shape+(1,))].index(True)
-
-    # if it's an image, insert fake strides such that this fusion doesn't happen across image axes
-    # TODO: remove membufs
-    membufs = dedup([x.src[0].base for x in self.bufs if x.op in {Ops.LOAD, Ops.STORE}])
-    if isinstance(membufs[0].base.dtype, ImageDType):
-      base_shape = membufs[0].base.dtype.shape
-      if shape_idx_groups := get_contraction(self.output_shape, base_shape):
-        special_strides: tuple[sint, ...] = tuple()
-        for i,g in enumerate(shape_idx_groups):
-          shape_piece = tuple(self.output_shape[x] for x in g)
-          assert prod(shape_piece) == base_shape[i], f"get_contraction was wrong? {shape_piece} != {base_shape[i]}"
-          special_strides += strides_for_shape(shape_piece)
-        # adding the fake image shape
-        shapes.append(self.output_shape)
-        strides.append(special_strides)
-
-    # merge dimensions if we can, multi _merge_dims
-    # NOTE: this does not always preserve the reduce dimension
-    # TODO: move this into shapetracker, with tests!
-    # TODO: how does this work with multi-reduce?
-    rets = [[(s[0], st[0])] for s,st in zip(shapes, strides)]
-    for i in range(1, len(shapes[0])):
-      can_merge = []
-      for s,st,ret in zip(shapes, strides, rets):
-        # TODO: added the always mergeability of 1s, is this right? if so, add to shapetracker in the 1 case
-        si, sti, last_st = s[i], st[i], ret[-1][1]
-        can_merge.append((sti is not None) and ((sti != 0 and last_st == si*sti) or (sti == 0 and last_st == 0)))
-      # more can merge than this
-      mergeable = all(can_merge) and i != first_reduce
-      for j,(s,st) in enumerate(zip(shapes, strides)):
-        if mergeable: rets[j][-1] = (rets[j][-1][0] * s[i], st[i])
-        else: rets[j].append((s[i], st[i]))
-
-    # do the reshapes
-    for i,x in enumerate(rets[:len(self.sts)]): self.sts[i] = self.sts[i].reshape(tuple([y[0] for y in x]))
-
-  # ******************** apply optimizations ********************
-
-  def real_axis(self, op:OptOps, axis:int|None):
-    try:
-      if axis is None: return -1
-      if op is OptOps.UNROLL: return self.unrollable_dims[axis]
-      if op in {OptOps.GROUP, OptOps.GROUPTOP}: return self.axes_of(AxisType.REDUCE)[axis]
-      check(axis < self.shape_len, f"invalid axis on {axis=} {op=} {self.shape_len=}")
-      return axis
-    except IndexError as e: raise KernelOptError from e
-
-  def apply_opt(self, opt:Opt, append_opt:bool=True) -> int|None:
-    if self.finalized: raise RuntimeError("can't optimize Kernel after it's finalized")
-    if self.dont_use_locals: check(opt.op not in {OptOps.LOCAL, OptOps.GROUP, OptOps.GROUPTOP}, "not using locals")
-
-    if opt.op is OptOps.TC:
-      check(len(self.applied_opts) == 0, "tensor core opts must be first") # TODO: things like PADTO might be fine
-      check(len(self.opts.tensor_cores) > 0, "must have tensor cores")
-      check(opt.axis is not None, "tensor core opts must have an axis")
-      check(opt.arg is not None and isinstance(opt.arg, tuple) and len(opt.arg) == 3, "tensor core opts must have valid arg")
-      check(-1 <= (tc_select:=cast(tuple, opt.arg)[0]) < len(self.opts.tensor_cores), "tensor core opts must have valid tc_select")
-      check(0 <= (tc_opt:=cast(tuple, opt.arg)[1]) <= 2, "tensor core opts must have valid tc_opt")
-      check(0 < (use_tensor_cores:=cast(tuple, opt.arg)[2]) <= 2, "use_tensor_cores value is not valid")
-      check(self._apply_tc_opt(use_tensor_cores, cast(int, opt.axis), tc_select, tc_opt), "no tensor core available")
-      self.applied_opts.append(opt)
-      return None
-
-    axis = self.real_axis(opt.op, opt.axis)
-
-    if opt.op is OptOps.SWAP: amt = self.real_axis(opt.op, cast(int, opt.arg))  # arg is an axis in the SWAPs
-    elif opt.arg is not None:
-      check(isinstance(opt.arg, int), "arg should be int")
-      amt = arg if (arg:=cast(int, opt.arg)) != 0 else self.full_shape[axis]
-      check(isinstance(amt, int) and amt != 1, f"shift/padto of {amt=}, 1 or symbolic amount is meaningless")
-      if opt.op is not OptOps.PADTO:
-        # we check both the full_shape and each shape
-        check(self.full_shape[axis] % amt == 0, f"no longer valid shift {self.full_shape[axis]=}, {amt=}")
-        for st in self.sts: check(st.shape[axis] == 1 or st.shape[axis] % amt == 0, f"no longer valid shift {st.shape[axis]=}, {amt=}")
-    else: amt = -1
-
-    if self.reduceop is not None and (opt.op in {OptOps.GROUP, OptOps.GROUPTOP} or \
-                                      (self.group_for_reduces and opt.op not in {OptOps.NOLOCALS, OptOps.PADTO})):
-      acc_sz = self.reduceop.dtype.itemsize
-      upcast_sz = prod([self.full_shape[a] for a in self.axes_of(AxisType.UPCAST)])
-      local_sz = prod([self.full_shape[a] for a in self.axes_of(AxisType.LOCAL)])
-      smem_sz = amt*acc_sz*upcast_sz*local_sz
-      check(smem_sz <= self.opts.shared_max, f"exceeds maximum shared memory size: needs {smem_sz}, max {self.opts.shared_max}")
-
-    new_axis = None
-    if opt.op is OptOps.LOCAL:    # cyan
-      # NOTE: LLVM/CPU can use locals too, but they are treated the same as globals (still helpful for L1 cache)
-      # it's disabled for now since it makes BEAM slow for little gain
-      check(self.opts.has_local, "target does not support local")
-      check(self.axis_types[axis] is AxisType.GLOBAL, "local is for globals")
-      new_axis = self.shift_to(axis, amt, AxisType.LOCAL, insert_at=max(self.axes_of(AxisType.GLOBAL, AxisType.LOCAL))+1)
-    elif opt.op in {OptOps.GROUP, OptOps.GROUPTOP}:   # green
-      check(self.opts.has_local and self.opts.has_shared, "target does not support local or shared mem")
-      check(self.axis_types[axis] is AxisType.REDUCE, "must be reduce axis to group")
-      check(not self.tensor_core, "can't group with tensor cores")
-      check(len(reduce_axes:=[i for r in self.reduceops for i in r.axis_arg]) == len(set(reduce_axes)), "can't group with parallel reduces")
-      new_axis = self.shift_to(axis, amt, AxisType.GROUP_REDUCE, top=(opt.op is OptOps.GROUPTOP), insert_at=min(self.axes_of(AxisType.REDUCE)))
-    elif opt.op is OptOps.UNROLL:                     # purple
-      check(self.axis_types[axis] not in (AxisType.UPCAST, AxisType.UNROLL), "can't upcasted already upcasted")
-      check(amt <= 32, "don't unroll more than 32")
-      new_axis = self.shift_to(axis, amt, AxisType.UNROLL, insert_at=None)
-    elif opt.op is OptOps.UPCAST:                     # yellow
-      check(axis in self.upcastable_dims, f"{axis=} not in {self.upcastable_dims=}")
-      # NOTE: assume the first get_local_axes() LOCAL are for TC
-      check(not (self.tensor_core and axis in self.axes_of(AxisType.LOCAL)[:len(self.tensor_core.get_local_axes())]), "can't upcast TC locals")
-      check((self.opts is not None and self.opts.device == "DSP") or amt <= 16, "don't upcast more than 16")
-      new_axis = self.shift_to(axis, amt, AxisType.UPCAST,
-                               insert_at=max(self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP, AxisType.UPCAST))+1)
-    elif opt.op is OptOps.NOLOCALS:
-      check(self.opts.has_local and not self.dont_use_locals, "NOLOCALS is meaningless if target does not support local or already not using locals")
-      check(AxisType.LOCAL not in self.axis_types and self.group_for_reduces == 0, "can't have no locals with locals")
-      self.dont_use_locals = True
-    elif opt.op is OptOps.SWAP:
-      check(axis < amt, f"swap is only for axis < amt, getting {amt=}, {axis=}")
-      check(self.axis_types[axis]==self.axis_types[amt]==AxisType.GLOBAL, f"swap is for globals {self.axis_types[axis]=}, {self.axis_types[amt]=}")
-      permute = list(range(self.shape_len))
-      permute[axis], permute[amt] = permute[amt], permute[axis]
-      self.permute(tuple(permute))
-    elif opt.op is OptOps.PADTO:
-      check(not self.vars, "does not work with symbolic shape")
-      check(self.axis_types[axis] not in (AxisType.UPCAST, AxisType.UNROLL), "cannot pad upcasted")
-      # ok to pad SUM if all parent ALU ops have f(0) = 0
-      if (r:=self.reduceop) is not None and self.axis_types[axis] in (AxisType.GROUP_REDUCE, AxisType.REDUCE):
-        check(r.arg[0] is Ops.ADD and can_pad(r, {}), f"cannot pad {r}")
-      padded = False
-      for i,st in enumerate(self.sts):
-        if (s:=st.shape[axis]) == 1: continue  # reduced
-        check(s > amt//4, f"pad adds more than quadruple the work {st.shape[axis]=} > {amt//4=}")
-        if (ru := round_up(cast(int, s), amt) - s):
-          # pad right seems to be faster
-          self.sts[i] = st.pad(((0,0),) * axis + ((0,ru),) + ((0,0),) * (len(st.shape)-axis-1))
-          padded = True
-      check(padded, "nothing was padded")
-
-    if append_opt: self.applied_opts.append(opt)
-    if self.simplify_ones() and self.tensor_core_opts:
-      self.tensor_core_opts.fix_axes(axis) # fix up axes in TC opts if required after simplify_ones()
-    return new_axis
-
-  def apply_opts(self, opts:Sequence[Opt]) -> Kernel:
-    for opt in opts: self.apply_opt(opt)
-    return self
-
-  # **** kernel outputs, mostly tensor cores ****
-
-  def _create_tc_opts(self, reduceop:UOp, tc:TensorCore, axis:int, opt_level:int) -> TensorCoreOptions|None:
-    has_cast = tc.dtype_in != tc.dtype_out
-    if has_cast and not (reduceop.src[0].op is Ops.CAST and reduceop.src[0].dtype == tc.dtype_out): return None
-
-    mul_op = reduceop.src[0].src[0] if has_cast else reduceop.src[0]
-    if mul_op.op is not Ops.MUL: return None
-
-    def buf_index(src:UOp) -> int|None:
-      # TODO: apply tc even if the sources are not from LOAD
-      if src.op is Ops.LOAD and src.dtype == tc.dtype_in: return self.bufs.index(src)
-      try:
-        if opt_level >= 1 and src.op is Ops.CAST and src.dtype == tc.dtype_in: return self.bufs.index(src.src[0])
-      except ValueError: return None
-      return None
-    if (buf0:=buf_index(mul_op.src[0])) is None or (buf1:=buf_index(mul_op.src[1])) is None: return None
-
-    buf0_strides, buf1_strides = self.sts[buf0].real_strides(), self.sts[buf1].real_strides()
-    axis_buf0 = [(i,self.full_shape[i],buf1_strides[i]) for i in self.upcastable_dims if buf0_strides[i] == 0]
-    axis_buf1 = [(i,self.full_shape[i],buf0_strides[i]) for i in self.upcastable_dims if buf1_strides[i] == 0]
-    if not (axis_buf0 and axis_buf1 and (len(self.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE)) == 1 or (opt_level >= 1))): return None
-
-    axis_choices = list(itertools.product(axis_buf0, axis_buf1, self.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE)))
-    if not (axis < len(axis_choices)): return None
-
-    s0, s1, s2 = axis_choices[-(axis+1)][0][0], axis_choices[-(axis+1)][1][0], axis_choices[-(axis+1)][2]  # s0 is n, s1 is m, s2 is k
-    axis_pads = tuple((x, tc.dims[i]) for i, x in enumerate([s0, s1, s2]) if resolve(self.full_shape[x]%tc.dims[i] != 0))
-    if axis_pads and (opt_level < 2): return None
-    if DEBUG >= 3: print("TENSOR CORES", axis_buf0, axis_buf1, tc)
-    return TensorCoreOptions(axes=(s0, s1, s2), axes_exist=(True, True), axis_pads=axis_pads)
-
-  def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> bool:
-    if use_tensor_cores and self.reduceop is not None and self.reduceop.arg[0] is Ops.ADD:
-      tensor_cores = self.opts.tensor_cores if tc_select == -1 else [self.opts.tensor_cores[tc_select]]
-      for tc in tensor_cores:
-        tensor_core_opts = [self._create_tc_opts(reduceop, tc, axis, opt_level) for reduceop in self.reduceops]
-        if tensor_core_opts[0] is None: continue
-        # can only fuse reduces with the same tc options
-        assert all_same(tensor_core_opts)
-        self.tensor_core_opts = tc_opts = tensor_core_opts[0]
-
-        # attempt to pad the tensor axes that require it
-        try:
-          for axis, dim in tc_opts.axis_pads: self.apply_opt(Opt(OptOps.PADTO, axis, dim), append_opt=False) # PADTO might fail
-        except KernelOptError: continue
-        # tensor core -- unroll the reduce dim (K), upcast and local the inner and outer dims (N, M)
-        for opt in tc.opts: self.apply_opt(Opt({"u":OptOps.UPCAST, "l":OptOps.LOCAL}[opt[0]], tc_opts.axes[int(opt[1])], 2), append_opt=False)
-        for dim, amt in tc.get_reduce_axes(): self.apply_opt(Opt(OptOps.UNROLL, 0, amt), append_opt=False) # TODO: this should be the reduce, not 0
-        self.tensor_core = tc
-        self.use_tensor_cores = use_tensor_cores  # TC=2 will do the shape ops without the WMMA
-        return True
-    return False
-
-  def apply_tensor_cores(self, use_tensor_cores=1, extra_opts:list[Opt]|None=None, axis:int=0, tc_select:int|None=None, tc_opt:int|None=None) -> bool:
-    """ Attempts to apply a tensor core optimization to the kernel. If one exists and applies properly, return true, otherwise return false.
-    Tensor cores are optimized instructions that matrix multiply-accumulate across a wave of threads: D(M, N) = A(M, K) * B(K, N) + C(M, N).
-
-    Keyword arguments:
-    use_tensor_cores -- controls how tensor cores are applied (default 1)
-      0: will disable any tensor core matching
-      1: enable tensor cores
-      2: apply tensor core shape but don't use UOp.WMMA
-    extra_opts -- additional Opt's to apply after the tensor core instead of the hand-coded additional Opt's (default None)
-    tc_select -- specifies which tensor core(s) to use for optimization (default -1)
-      -1: iterates through all available tensor cores in order and uses the first one that matches the requirements (dims and dtypes)
-      [0-N]: uses only the n'th tensor core available; useful for search
-    tc_opt -- controls which kinds of kernels may be eligible for tensor cores application (default 2 during BEAM, 0 otherwise)
-      0: applies to only kernels with a single reduce axis and direct Ops.LOAD into Ops.MUL
-      1: allows kernels with multiple reduce axes and also multiplication of Ops.CAST'd buffers
-      2: allows kernels with M, N, K axes that are not multiples of the tensor core dimensions by applying padding those axes as needed
-    """
-    if tc_select is None: tc_select = TC_SELECT.value
-    if tc_opt is None: tc_opt = TC_OPT.value
-    if not self.opts.tensor_cores: return False
-    try: # check TC first and apply hand-coded opts if successful
-      self.apply_opt(Opt(OptOps.TC, axis, (tc_select, tc_opt, use_tensor_cores)))
-
-      if (tc_opts:=self.tensor_core_opts) is not None:
-        if extra_opts is not None: self.apply_opts(extra_opts)
-        else:
-          if AMX: return True # skip hand-coded TC opts if AMX, upcasting will make kernel slower
-          # hand-coded TC opts
-          for tc_dim in [tc_dim for tc_dim in [1,0] if tc_opts.axes_exist[tc_dim]]: # attempt to upcast M and N
-            szs = [sz for sz in [5,4,3,2] if self.full_shape[tc_opts.axes[tc_dim]] % sz == 0]
-            if szs: self.apply_opt(Opt(OptOps.UPCAST, tc_opts.axes[tc_dim], szs[0]))
-
-          if tc_opts.axes_exist[0] and (szs := [sz for sz in [4,2] if self.full_shape[tc_opts.axes[0]] % sz == 0]): # attempt to local N
-            self.apply_opt(Opt(OptOps.LOCAL, tc_opts.axes[0], szs[0]))
-      return True
-    except KernelOptError:
-      return False
-
-  # strings like ['g0', 'g1', 'l0', 'l1', 'l2', 'l3', 'l4', 'l5', 'R0', 'r0', 'r1', 'r2', 'u0', 'u1', 'u2']
-  def shape_str(self) -> list[str]:
-    ret: list[str] = []
-    cnt: dict[AxisType, int] = {}
-    for x in self.axis_types:
-      cnt[x] = (cnt[x] + 1) if x in cnt else 0
-      ret.append(f"{axis_letters[x]}{cnt[x]}")
-    return ret
-  def shape_str_to_axis(self, nms:list[str]) -> tuple[int, ...]: return tuple([self.shape_str().index(x) for x in nms])
-
-  def get_optimized_ast(self, name_override:str|None=None) -> UOp:
-    @functools.cache
-    def fixup_ast(op:UOp) -> UOp:
-      ret = op.replace(src=tuple(fixup_ast(x) for x in op.src)) # noqa: F821
-      if op.op in GroupOp.Buffer and op in self.bufs:
-        st = self.sts[self.bufs.index(op)]
-        # replace the VIEW source
-        return ret.replace(src=(ret.src[0].replace(arg=st),)+ret.src[1:])
-      if op.op is Ops.SINK:
-        # NOTE: should group_for_reduces be added to the local_dims?
-        # TODO: arg.name should be able to be None
-        kernel_name = ret.arg.name if ret.arg is not None and ret.arg.name != "test" else self.name if name_override is None else name_override
-        return ret.replace(arg=KernelInfo(kernel_name, tuple(self.axis_types), self.dont_use_locals, tuple(self.applied_opts)))
-      if op.op is Ops.REDUCE_AXIS:
-        reduce_idx = len(self.bufs) + self.reduceops.index(op) * 2
-        changed = tuple(i for i in range(self.shape_len) if resolve(self.sts[reduce_idx].shape[i] != self.sts[reduce_idx + 1].shape[i]))
-        axes = tuple(i for i in self.axes_of(AxisType.REDUCE, AxisType.GROUP_REDUCE, AxisType.UNROLL) if i in changed)
-        if (tc := self.tensor_core) and self.use_tensor_cores == 1:
-          # get reduce/upcast axes for the tensor cores
-          tc_reduce_axes = self.shape_str_to_axis([f"r{i}" for i in range(len(tc.get_reduce_axes()))])
-          base_upcast_axes = tuple([(s,2) for s in self.shape_str_to_axis(tc.base_upcast_axes())])
-          tc_upcast_axes = tuple([base_upcast_axes[:int(math.log2(tc.elements_per_thread[i]))] for i in range(3)])
-
-          # permute the srcs
-          srcs = list((ret.src[0] if ret.src[0].op is not Ops.CAST else ret.src[0].src[0]).src)
-          for i, (src, permaxis) in enumerate(zip(srcs, tc.permutes_for_shape_str(self.shape_str()))):
-            src_st = (src if src.op is Ops.LOAD else src.src[0]).st_arg
-            srcs[i] = src.view(ShapeTracker.from_shape(src_st.shape).permute(permaxis))
-
-          # construct the op
-          wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, self.opts.device, tc.threads, tc_upcast_axes, tc_reduce_axes)
-          wmma = UOp(Ops.WMMA, dtype=tc.dtype_out.vec(tc.elements_per_thread[2]), src=(
-            UOp(Ops.CONTRACT, dtype=srcs[0].dtype.vec(tc.elements_per_thread[0]), src=(srcs[0],), arg=tc_upcast_axes[0]),
-            UOp(Ops.CONTRACT, dtype=srcs[1].dtype.vec(tc.elements_per_thread[1]), src=(srcs[1],), arg=tc_upcast_axes[1]),
-            UOp.const(tc.dtype_out.vec(tc.elements_per_thread[2]), 0.0)), arg=wmma_arg)
-          tc_uop = UOp(Ops.UNROLL, tc.dtype_out, (wmma,), arg=tc_upcast_axes[2])
-
-          # preserve any other reduce
-          return ret.replace(src=(tc_uop,), arg=(Ops.ADD, new_axes)) if (new_axes := tuple(i for i in axes if i not in tc_reduce_axes)) else tc_uop
-
-        ret = ret.replace(arg = (op.arg[0], axes))
-      return ret
-    self.finalized = True
-    fixed_ast = fixup_ast(self.ast)
-    del fixup_ast
-    return graph_rewrite(fixed_ast, view_left+view_left_through_load, name="fixup optimized AST")

tinygrad/codegen/opt/postrange.py

@@ -0,0 +1,366 @@
+from __future__ import annotations
+import math, itertools
+from collections import defaultdict
+from typing import cast, Final, Sequence
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, _substitute, AxisType, ssimplify, can_pad
+from tinygrad.uop.symbolic import symbolic_flat
+from tinygrad.device import Buffer
+from tinygrad.dtype import AddrSpace, dtypes, ImageDType
+from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod
+from tinygrad.codegen.opt import axis_colors, Opt, OptOps, KernelOptError, check, axis_letters
+from tinygrad.renderer import Renderer
+from tinygrad.schedule.rangeify import remove_tags
+
+# NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
+axis_to_pos = {AxisType.LOOP: -1, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
+               AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
+
+def flatten_range(r:UOp):
+  off = 2 if r.op is Ops.STORE else 1
+  rngs = r.src[off:]
+  if not len(rngs): return None
+  new_rngs = [x for x in UOp.sink(*rngs).toposort() if x.op is Ops.RANGE]
+  return r.replace(src=r.src[:off]+tuple(new_rngs))
+
+pm_flatten_range = PatternMatcher([
+  # real ranges only
+  (UPat((Ops.REDUCE, Ops.STORE), name="r"), flatten_range),
+])
+
+def count_divmod(x:UOp): return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
+
+class Scheduler:
+  def __init__(self, ast:UOp, opts:Renderer):
+    self.ast, self.opts = ast, opts
+    self.dont_use_locals = self.ast.arg.dont_use_locals if self.ast.arg is not None else False
+    self.applied_opts = list(self.ast.arg.applied_opts) if self.ast.arg is not None else []
+
+  @property
+  def rngs(self):
+    # always in order by axistype
+    return sorted([u for u in self.ast.parents if u.op is Ops.RANGE and u.vmax > 0], key=lambda x: (axis_to_pos[x.arg[-1]],) + x.arg[0:-1])
+  @property
+  def shape_len(self): return len(self.rngs)
+  @property
+  def full_shape(self): return [ssimplify(x.src[0]) for x in self.rngs]
+  @property
+  def axis_types(self): return [x.arg[-1] for x in self.rngs]
+  @property
+  def maxarg(self): return max([x.arg[0] for x in self.rngs], default=0)
+
+  # strings like ['g0', 'g1', 'l0', 'l1', 'l2', 'l3', 'l4', 'l5', 'R0', 'r0', 'r1', 'r2', 'u0', 'u1', 'u2']
+  def shape_str(self) -> list[str]:
+    ret: list[str] = []
+    cnt: dict[AxisType, int] = {}
+    for x in self.axis_types:
+      cnt[x] = (cnt[x] + 1) if x in cnt else 0
+      ret.append(f"{axis_letters[x]}{cnt[x]}")
+    return ret
+  def shape_str_to_axis(self, nms:list[str]) -> tuple[int, ...]: return tuple([self.shape_str().index(x) for x in nms])
+
+  @property
+  def termination(self):
+    terminators = [u for u in self.ast.parents if u.op in {Ops.REDUCE, Ops.STORE}]
+    termination = {}
+    for t in terminators:
+      # works without pm_flatten_range
+      for u in UOp.sink(*t.src[1 if t.op is Ops.REDUCE else 2:]).parents:
+        if u.op is Ops.RANGE: termination[u] = t
+    return termination
+
+  def copy(self): return Scheduler(self.get_optimized_ast(), self.opts)
+
+  kernel_cnt: Final[defaultdict[str, int]] = defaultdict(int)
+  def get_optimized_ast(self, name_override:str|None=None):
+    if name_override is not None: name = name_override
+    else:
+      kernel_type = "r" if self.reduceop is not None else "E"
+      name = kernel_type + colored('_', 'BLACK').join(['']+[colored(x.src[0].render(), color) for x,color in zip(self.rngs, self.colors())])
+      Scheduler.kernel_cnt[(function_name := to_function_name(name))] += 1
+      num = f"n{Scheduler.kernel_cnt[function_name]-1}" if Scheduler.kernel_cnt[function_name] > 1 else ""
+      name += colored(num, 'BLACK')
+    self.ast = graph_rewrite(self.ast, pm_flatten_range, name="flatten range")
+    return self.ast.replace(arg=KernelInfo(name=name, applied_opts=tuple(self.applied_opts), dont_use_locals=self.dont_use_locals), tag=1)
+
+  def convert_loop_to_global(self):
+    if not self.opts.has_local: return None
+    store_rngs = self.ast.src[0].src[2:]
+
+    # filter any not in local stores
+    local_store_rngs = [x.ranges for x in self.ast.toposort() if (x.op is Ops.STORE and x.src[0].ptrdtype.addrspace == AddrSpace.LOCAL) \
+                        or (x.op is Ops.BUFFERIZE and x.arg == AddrSpace.LOCAL)]
+    for ls in local_store_rngs: store_rngs = tuple([x for x in store_rngs if x in ls])
+
+    store_rng = [x for x in UOp.sink(*store_rngs).toposort() if x.op is Ops.RANGE] if store_rngs else []
+    rng = [x.replace(arg=(x.arg[0], AxisType.GLOBAL)) if x.arg[1] == AxisType.LOOP and x in store_rng else x for x in self.rngs]
+
+    self.ast = self.ast.substitute(dict(zip(self.rngs, rng)))
+
+  def simplify_merge_adjacent(self):
+    i = 0
+    while i < len(self.rngs)-1:
+      r0, r1 = self.rngs[i], self.rngs[i+1]
+      # same axistype and same termination
+      termination = self.termination
+      if r0.arg[1] == r1.arg[1] and r0 in termination and r1 in termination and termination[r0] == termination[r1]:
+        s0, s1 = r0.src[0], r1.src[0]
+        # do the merge
+        oidx = self.ast.simplify()
+        new_range = r0.replace(src=(s0*s1,))
+        nidx = graph_rewrite(oidx, _substitute+symbolic_flat+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1}, name=f"check_merge_{i}_{i+1}")
+        # check if it simplifies
+        if count_divmod(nidx) <= count_divmod(oidx):
+          self.ast = nidx
+          continue
+      i += 1
+
+  def colors(self) -> list[str]: return [axis_colors[x] if not self.dont_use_locals or not x == AxisType.GLOBAL else "BLUE" for x in self.axis_types]
+  def colored_shape(self) -> str: return ' '.join([colored(f'{x.src[0].render():>4s}', color) for x,color in zip(self.rngs, self.colors())])
+
+  def shift_to(self, rng:UOp, amount:int, new_type:AxisType, top:bool=False, input_new_rng=None):
+    if (old_sz:=rng.src[0].divides(amount)) is None:
+      raise KernelOptError(f"{amount} can't divide {rng.src[0]} in {self.colored_shape()}")
+    new_rng = UOp.range(amount, self.maxarg+1, new_type) if input_new_rng is None else input_new_rng
+    replaced_rng = rng.replace(src=(UOp.const(dtypes.int, old_sz),))
+    sub_axis = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng)
+    self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[0]} {amount}")
+    return replaced_rng, new_rng
+
+  def ranges_of(self, *axis_type:AxisType) -> list[UOp]: return [r for r in self.rngs if r.arg[-1] in axis_type]
+  def axes_of(self, *axis_type:AxisType) -> list[int]: return [i for i,t in enumerate(self.axis_types) if t in axis_type]
+
+  # copied from kernel.py
+  @property
+  def upcastable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP) \
+                                                  if isinstance(s:=self.full_shape[i], int) and s > 1]
+  @property
+  def unrollable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE) \
+                                                  if isinstance(s:=self.full_shape[i], int) and s > 1]
+
+  def real_axis(self, op:OptOps, axis:int|None):
+    try:
+      if axis is None or op is OptOps.TC: return -1
+      if op is OptOps.UNROLL: return self.unrollable_dims[axis]
+      if op in {OptOps.GROUP, OptOps.GROUPTOP}: return self.axes_of(AxisType.REDUCE)[axis]
+      check(axis < self.shape_len, f"invalid axis on {axis=} {op=} {self.shape_len=}")
+      return axis
+    except IndexError as e: raise KernelOptError from e
+
+  def apply_opts(self, opts:Sequence[Opt]) -> Scheduler:
+    for opt in opts: self.apply_opt(opt)
+    return self
+
+  def apply_opt(self, opt:Opt, append_opt:bool=True):
+    if opt.op is OptOps.NOLOCALS:
+      check(all(x not in {AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE} for x in self.axis_types), "no locals can't have locals")
+      if append_opt: self.applied_opts.append(opt)
+      self.dont_use_locals = True
+      return
+
+    if opt.op in {OptOps.LOCAL, OptOps.GROUP, OptOps.GROUPTOP}:
+      check(self.opts.has_local, "locals needed for opt")
+
+    rng = self.rngs[real_axis] if (real_axis:=self.real_axis(opt.op, opt.axis)) >= 0 else UOp(Ops.NOOP)
+
+    opt_to_at = {
+      OptOps.LOCAL: AxisType.LOCAL, OptOps.UPCAST: AxisType.UPCAST,
+      OptOps.UNROLL: AxisType.UNROLL, OptOps.GROUP: AxisType.GROUP_REDUCE,
+      OptOps.GROUPTOP: AxisType.GROUP_REDUCE}
+
+    ret = None
+    if opt.op in opt_to_at:
+      amt:int = int(rng.vmax+1) if opt.arg == 0 else cast(int, opt.arg)
+
+      # copied from kernel.py. prevents METAL compiler hangs
+      if self.reduceop is not None and (opt.op in {OptOps.GROUP, OptOps.GROUPTOP} or \
+                                        (self.group_for_reduces and opt.op not in {OptOps.NOLOCALS, OptOps.PADTO})):
+        upcast_local_sz = prod([self.full_shape[a] for a in self.axes_of(AxisType.UPCAST, AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE)])
+        smem_sz = amt*upcast_local_sz*self.reduceop.dtype.itemsize
+        check(smem_sz <= self.opts.shared_max, f"exceeds maximum shared memory size: needs {smem_sz}, max {self.opts.shared_max}")
+
+      if opt.op is OptOps.UNROLL:
+        check(amt <= 32, "don't unroll more than 32")
+        check(rng.arg[-1] in {AxisType.GROUP_REDUCE, AxisType.REDUCE}, "unroll is for GROUP_REDUCE/REDUCE")
+      if opt.op is OptOps.UPCAST:
+        check((self.opts is not None and self.opts.device == "DSP") or amt <= 16, "don't upcast more than 16")
+        check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP}, f"upcast is for GLOBAL/LOCAL/LOOP, not {rng.arg[-1]}")
+      if opt.op is OptOps.LOCAL:
+        check(not self.dont_use_locals, "can't use locals")
+        check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOOP}, "local is for globals")
+      if opt.op in {OptOps.GROUP, OptOps.GROUPTOP}:
+        check(all(x.op is not OptOps.TC for x in self.applied_opts), "no grouping with tensor cores")  # TODO: why is this wrong?
+        check(not self.dont_use_locals, "can't use locals")
+        check(rng.arg[-1] == AxisType.REDUCE, "group is for reduce")
+      ret = self.shift_to(rng, amt, opt_to_at[opt.op], top=opt.op==OptOps.GROUPTOP)
+    elif opt.op is OptOps.TC:
+      check(len(self.applied_opts) == 0, "tensor core opts must be first") # TODO: remove the need for this by having warps
+      check(opt.axis is not None, "tensor core opts must have an axis")
+      check(opt.arg is not None and isinstance(opt.arg, tuple) and len(opt.arg) == 3, "tensor core opts must have valid arg")
+      check(-1 <= (tc_select:=cast(tuple, opt.arg)[0]) < len(self.opts.tensor_cores), "tensor core opts must have valid tc_select")
+      check(0 <= (tc_opt:=cast(tuple, opt.arg)[1]) <= 2, "tensor core opts must have valid tc_opt")
+      check(0 < (use_tensor_cores:=cast(tuple, opt.arg)[2]) <= 2, "use_tensor_cores value is not valid")
+      ret = self._apply_tc_opt(use_tensor_cores, cast(int, opt.axis), tc_select, tc_opt)
+      check(ret is not None, "no tensor core available")
+    elif opt.op is OptOps.PADTO:
+      check(rng.src[0].op is Ops.CONST, "only pad const axes")
+      check(rng.arg[-1] not in {AxisType.UPCAST, AxisType.UNROLL}, "cannot pad upcasted") # TODO: why is this wrong?
+      # ok to pad SUM if all parent ALU ops have f(0) = 0
+      if (r:=self.reduceop) is not None and rng.arg[-1] in (AxisType.GROUP_REDUCE, AxisType.REDUCE):
+        check(r.arg[0] is Ops.ADD and can_pad(r, {}), f"cannot pad {r}")
+      new_sz = round_up(int(rng.vmax+1), cast(int, opt.arg))
+      check(rng.vmax+1 > new_sz//4, "pad adds more than quadruple the work")
+      replaced_rng = UOp.range(new_sz, *rng.arg)
+      replaces = {rng:replaced_rng}
+      for b in self.bufs:
+        if rng in b.src[1].sparents:
+          valid = replaced_rng < rng.vmax+1
+          if len(b.src) > 2: valid = b.src[2] & valid
+          replaces[b] = b.replace(src=b.src[0:2]+(valid,))
+      self.ast = self.ast.substitute(replaces, f"padto {rng.arg[:-1]} {opt.arg}")
+    elif opt.op is OptOps.SWAP:
+      try:
+        altrng = self.rngs[opt.arg]
+      except IndexError:
+        raise KernelOptError
+      check(rng.arg[-1] == AxisType.GLOBAL and altrng.arg[-1] == AxisType.GLOBAL, "swap only for globals")
+      self.ast = self.ast.substitute({rng:rng.replace(arg=(*altrng.arg[0:-1], rng.arg[-1]), tag=1),
+                                      altrng:altrng.replace(arg=(*rng.arg[0:-1], altrng.arg[-1]), tag=1)})
+      self.ast = graph_rewrite(self.ast, remove_tags)
+    else:
+      raise KernelOptError(f"unsupported opt {opt.op}")
+
+    if append_opt: self.applied_opts.append(opt)
+    return ret
+
+  def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> None|list[UOp]:
+    reduceops = [x for x in self.ast.toposort() if x.op is Ops.REDUCE]
+    if not len(reduceops): raise KernelOptError("no reduce ops for TensorCore")
+    reduceop = reduceops[0]
+    if use_tensor_cores and reduceop is not None and reduceop.arg is Ops.ADD:
+      mul = reduceop.src[0] if reduceop.src[0].op is not Ops.CAST else reduceop.src[0].src[0]
+      if mul.op is not Ops.MUL: return None
+      in0, in1 = mul.src
+      try:
+        tensor_cores = self.opts.tensor_cores if tc_select == -1 else [self.opts.tensor_cores[tc_select]]
+      except IndexError:
+        raise KernelOptError(f"invalid tensor core choice {tc_select}")
+      for tc in tensor_cores:
+        if tc.dtype_in == in0.dtype.scalar() and tc.dtype_in == in1.dtype.scalar() and tc.dtype_out == reduceop.dtype.scalar():
+          # tensor cores have three ranges. X, Y, and REDUCE
+          in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: -x.arg[0])
+          in1_ranges = sorted([u for u in in1.ranges if u not in in0.ranges], key=lambda x: -x.arg[0])
+          red_ranges = sorted(reduceop.src[1:], key=lambda x: -x.arg[0])
+          if DEBUG >= 3:
+            print(f"TC({axis}): {[(x.arg[0],x.vmax+1) for x in in0_ranges]}",
+                              f"{[(x.arg[0],x.vmax+1) for x in in1_ranges]} {[(x.arg[0],x.vmax+1) for x in red_ranges]}")
+          if not len(in0_ranges) or not len(in1_ranges) or not len(red_ranges): continue
+
+          # pick ranges
+          # NOTE: why are in1 and in0 switched?
+          axis_choices = list(itertools.product(in1_ranges, in0_ranges, red_ranges))
+          if not (axis < len(axis_choices)): continue
+          axes = list(axis_choices[axis])
+
+          # do optimizations and save the ranges
+          try:
+            for i,a in enumerate(axes):
+              idx = self.rngs.index(a)
+              if (a.vmax+1) % tc.dims[i] != 0:
+                if opt_level < 2: raise KernelOptError("tc padding requires opt_level >= 2")
+                # apply_opt should return the updated range?
+                self.apply_opt(Opt(OptOps.PADTO, idx, tc.dims[i]), append_opt=False) # PADTO might fail
+                axes[i] = self.rngs[idx]
+          except KernelOptError: continue
+
+          # we create the warp as a whole thing, in case some of these ranges are moved/removed later
+          warp = UOp.range(tc.threads, -1, AxisType.WARP)
+          ne: list[UOp] = []
+          for opt in tc.opts:
+            if opt[0] == "l":
+              axes[int(opt[1])], new_range = self.shift_to(axes[int(opt[1])], 2, AxisType.LOCAL, input_new_rng=warp%2)
+              warp //= 2
+            elif opt[0] == "u":
+              axes[int(opt[1])], new_range = self.shift_to(axes[int(opt[1])], 2, AxisType.UPCAST)
+            else: raise RuntimeError(f"unsupported opt {opt[0]} in tensor cores")
+            ne.append(new_range)
+
+          for _, amt in tc.get_reduce_axes():
+            axes[2], new_range = self.shift_to(axes[2], amt, AxisType.UNROLL)
+            ne.append(new_range)
+
+          if use_tensor_cores != 2:
+            # fix the srcs
+            reduceop = [x for x in self.ast.toposort() if x.op is Ops.REDUCE][0]
+            tne = [x.replace(tag=1) for x in ne]
+            ret = reduceop.substitute(dict(zip(ne, tne)))
+            srcs = list((ret.src[0] if ret.src[0].op is not Ops.CAST else ret.src[0].src[0]).src)
+            srcs = [x.substitute(dict(zip(tne, [ne[i] for i in argsort(p)]))) for x,p in zip(srcs, tc.permutes_for_shape_str(tc.base_shape_str()))]
+
+            # get reduce/upcast axes for the tensor cores
+            tc_reduce_axes = self.shape_str_to_axis([f"r{i}" for i in range(len(tc.get_reduce_axes()))])
+            base_upcast_axes = tuple([(s,2) for s in self.shape_str_to_axis(tc.base_upcast_axes())])
+            tc_upcast_axes = tuple([base_upcast_axes[:int(math.log2(tc.elements_per_thread[i]))] for i in range(3)])
+
+            # axes to range number (was done in lowerer)
+            tc_upcast_axes = tuple([tuple([(self.rngs[a].arg[0], sz) for a,sz in v]) for v in tc_upcast_axes])
+            tc_reduce_axes = tuple([self.rngs[a].arg[0] for a in tc_reduce_axes])
+
+            # construct the op
+            # TODO: remove tc_upcast_axes from the arg
+            # do the reduce_axes always disappear? i think they don't
+            # they need to be moved into the WMMA srcs
+            wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, self.opts.device, tc.threads, tc_upcast_axes, ()) #, tc_reduce_axes)
+            wmma = UOp(Ops.WMMA, dtype=tc.dtype_out.vec(tc.elements_per_thread[2]), src=(
+              UOp(Ops.CONTRACT, dtype=srcs[0].dtype.vec(tc.elements_per_thread[0]), src=(srcs[0],), arg=tc_upcast_axes[0], tag=1),
+              UOp(Ops.CONTRACT, dtype=srcs[1].dtype.vec(tc.elements_per_thread[1]), src=(srcs[1],), arg=tc_upcast_axes[1], tag=1),
+              UOp.const(tc.dtype_out.vec(tc.elements_per_thread[2]), 0.0)), arg=wmma_arg, tag=1)
+            tc_uop = UOp(Ops.UNROLL, tc.dtype_out, (wmma,), arg=tc_upcast_axes[2], tag=1)
+
+            # preserve extra reduces
+            reduce_ranges = [x for x in UOp.sink(*reduceop.src[1:]).toposort() if x.op is Ops.RANGE and x.arg[0] not in tc_reduce_axes]
+            if len(reduce_ranges): tc_uop = UOp(Ops.REDUCE, tc_uop.dtype, (tc_uop,)+tuple(reduce_ranges), Ops.ADD)
+            self.ast = self.ast.substitute({reduceop: tc_uop})
+          return axes
+    return None
+
+  # helpers for hand_coded_optimizations
+  @property
+  def reduceop(self) -> UOp|None:
+    red = [x for x in self.ast.parents if x.op is Ops.REDUCE]
+    if not len(red): return None
+    return UOp(Ops.REDUCE_AXIS, red[0].dtype, red[0].src, (red[0].arg, ()))
+  @property
+  def bufs(self) -> list[UOp]: return [x for x in self.ast.toposort() if x.op is Ops.INDEX][::-1]
+  @property
+  def output_shape(self):
+    return [s if at not in {AxisType.REDUCE, AxisType.UNROLL, AxisType.GROUP_REDUCE} else 1 for s,at in zip(self.full_shape, self.axis_types)]
+  @property
+  def upcasted(self) -> int: return len(self.axes_of(AxisType.UPCAST, AxisType.UNROLL))
+  @property
+  def group_for_reduces(self) -> int: return len(self.axes_of(AxisType.GROUP_REDUCE))
+
+def bufs_from_ast(ast:UOp, dname:str) -> list[Buffer]:
+  glbls = sorted([x for x in ast.parents if x.op is Ops.DEFINE_GLOBAL], key=lambda x: x.arg)
+  return [Buffer(dname, x.ptrdtype.size, x.dtype.base if not isinstance(x.dtype, ImageDType) else x.dtype) for x in glbls]
+
+def apply_opts(ctx:Renderer, ast:UOp):
+  if ast.tag is not None: return None
+  k = Scheduler(ast, ctx)
+  k.convert_loop_to_global()
+  k.simplify_merge_adjacent()
+  if BEAM >= 1:
+    from tinygrad.codegen.opt.search import beam_search
+    rawbufs = bufs_from_ast(ast, ctx.device)
+    k = beam_search(k, rawbufs, BEAM.value, bool(getenv("BEAM_ESTIMATE", 1)))
+  elif ast.arg is not None and ast.arg.opts_to_apply is not None:
+    for opt in ast.arg.opts_to_apply: k.apply_opt(opt)
+  elif not NOOPT and (ast.arg is None or ast.arg.applied_opts == ()):
+    from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
+    # NOTE: hand_coded_optimizations doesn't support multiblock opts yet
+    if all(len(u.src) == 1 for u in ast.parents if u.op is Ops.LOAD):
+      for opt in hand_coded_optimizations(k): k.apply_opt(opt)
+  return k.get_optimized_ast(name_override=ast.arg.name if ast.arg is not None and ast.arg.name != "test" else None)
+
+pm_postrange_opt = PatternMatcher([
+  (UPat(Ops.SINK, name="ast"), apply_opts),
+])

tinygrad/codegen/opt/search.py

@@ -1,17 +1,19 @@
-from typing import cast, Callable
-import itertools, functools, random, math, time, multiprocessing, traceback, signal, atexit
+from typing import cast
+import functools, math, time, multiprocessing, traceback, signal, atexit
 from collections import defaultdict
 from dataclasses import replace
-from tinygrad.uop.ops import UOp, Ops, Variable, sym_infer, AxisType
+from tinygrad.uop.ops import UOp, Ops, Variable, sym_infer, AxisType, pyrender
 from tinygrad.device import Device, Buffer, Compiler
 from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str
-from tinygrad.helpers import IGNORE_BEAM_CACHE, TC_SEARCH_OVER_SHAPE
+from tinygrad.helpers import IGNORE_BEAM_CACHE
 from tinygrad.dtype import ImageDType, PtrDType
-from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps, KernelOptError
+from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.tensor import Tensor
 from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.renderer import ProgramSpec
 
+from tinygrad.codegen.opt.postrange import Scheduler
+
 actions = [Opt(op=OptOps.UPCAST, axis=axis, arg=amt) for amt in [0,2,3,4,5,7] for axis in range(8)]
 actions += [Opt(op=OptOps.UNROLL, axis=axis, arg=amt) for amt in [0,4,7] for axis in range(5)]
 actions += [Opt(op=OptOps.LOCAL, axis=axis, arg=amt) for amt in [2,3,4,8,13,16,29] for axis in range(6)]
@@ -55,9 +57,11 @@ def _time_program(p:ProgramSpec, lib:bytes, var_vals:dict[Variable, int], rawbuf
   return tms
 
 class TimeoutException(Exception): pass
-def timeout_handler(signum, frame): raise TimeoutException()
+def timeout_handler(signum, frame):
+  if DEBUG >= 2: print("*** BEAM COMPILE TIMEOUT")
+  raise TimeoutException()
 
-def _try_compile_linearized_w_idx(x:tuple[int,Kernel], compiler:Compiler) -> tuple[int, tuple[ProgramSpec, bytes, float]|None]:
+def _try_compile_linearized_w_idx(x:tuple[int,Scheduler], compiler:Compiler) -> tuple[int, tuple[ProgramSpec, bytes, float]|None]:
   if hasattr(signal, "alarm"):
     signal.signal(getattr(signal, 'SIGALRM'), timeout_handler)
     # set timeout
@@ -91,7 +95,8 @@ def _ensure_buffer_alloc(bufs:list[Buffer]) -> list[Buffer]: return [buf.ensure_
 # *** external API ***
 
 # get (scrap) buffers for timing the linearizer
-def bufs_from_lin(lin:Kernel, allocate:bool=True) -> list[Buffer]:
+# NOTE: there's also bufs_from_ast in postrange
+def bufs_from_lin(lin:Scheduler, allocate:bool=True) -> list[Buffer]:
   bufsts: defaultdict[int, list[UOp]] = defaultdict(list)
   for x in lin.bufs:
     if x.src[0].base.op is Ops.DEFINE_GLOBAL: bufsts[x.src[0].base.arg].append(x)
@@ -108,17 +113,10 @@ def bufs_from_lin(lin:Kernel, allocate:bool=True) -> list[Buffer]:
   return cast(list[Buffer], rawbufs)
 
 # get dictionary of all possible actions
-def get_kernel_actions(lin:Kernel, include_0=True, candidates:list[Opt]|None=None) -> dict[int, Kernel]:
+def get_kernel_actions(lin:Scheduler, include_0=True, candidates:list[Opt]|None=None) -> dict[int, Scheduler]:
   acted_lins, max_up, max_lcl = {0:lin} if include_0 else {}, getenv("BEAM_UPCAST_MAX", 256), getenv("BEAM_LOCAL_MAX", 1024)
   kernel_actions = (actions if candidates is None else candidates).copy()
 
-  if TC_SEARCH_OVER_SHAPE and len(lin.applied_opts) == 0: # tensor core opts must be first
-    for i, action in enumerate(kernel_actions):
-      if action.op == OptOps.TC and (tc_arg := cast(tuple, action.arg))[0] == -1:
-        # replace every tc_action with default tc with one tc_action for each available tc
-        kernel_actions[i:i+1] = \
-          [Opt(op=OptOps.TC, axis=action.axis, arg=(tc_select, tc_arg[1], tc_arg[2])) for tc_select,_ in enumerate(lin.opts.tensor_cores)]
-
   for i,a in enumerate(kernel_actions):
     if a.axis is not None and a.op is not OptOps.TC:
       try: ax = lin.real_axis(a.op, a.axis)
@@ -127,10 +125,10 @@ def get_kernel_actions(lin:Kernel, include_0=True, candidates:list[Opt]|None=Non
     lin2 = lin.copy()
     try:
       lin2.apply_opt(a)
-      up, lcl, tc_up = 1, 1, prod(tc.dims)//tc.threads if (tc:=lin2.tensor_core) else 1
+      up, lcl, tc_up = 1, 1, prod(tc.dims)//tc.threads if hasattr(lin2, 'tensor_core') and (tc:=lin2.tensor_core) else 1
       for s,c in zip(lin2.full_shape, lin2.axis_types):
         if c in (AxisType.UPCAST, AxisType.UNROLL): up *= s
-        elif c in (AxisType.LOCAL, AxisType.GROUP_REDUCE): lcl *= s
+        elif c in (AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE): lcl *= s
       if up//tc_up > max_up or lcl > max_lcl:
         if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too many upcast/local. {up//tc_up=}, {max_up=}, {lcl=}, {max_lcl=}")
         continue
@@ -139,7 +137,7 @@ def get_kernel_actions(lin:Kernel, include_0=True, candidates:list[Opt]|None=Non
   return acted_lins
 
 beam_pool, BEAM_DEBUG = None, getenv("BEAM_DEBUG")
-def beam_search(lin:Kernel, rawbufs:list[Buffer], amt:int, allow_test_size=True, disable_cache=IGNORE_BEAM_CACHE.value) -> Kernel:
+def beam_search(lin:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True, disable_cache=IGNORE_BEAM_CACHE.value):
   global beam_pool
   key = {"ast": lin.ast.key, "amt": amt, "allow_test_size": allow_test_size, "device": lin.opts.device, "suffix": lin.opts.suffix}
   if not disable_cache and CACHELEVEL >= 1 and (val:=diskcache_get("beam_search", key)) is not None:
@@ -147,7 +145,7 @@ def beam_search(lin:Kernel, rawbufs:list[Buffer], amt:int, allow_test_size=True,
     for o in val[len(lin.applied_opts):]: ret.apply_opt(o)
     return ret
 
-  beam: list[tuple[Kernel, float]] = [(lin, float("inf"))]
+  beam: list[tuple[Scheduler, float]] = [(lin, float("inf"))]
   seen_libs = set()
 
   default_parallel = multiprocessing.cpu_count() if lin.opts.device in {"CUDA", "AMD", "NV", "METAL", "HIP"} else 0
@@ -157,7 +155,9 @@ def beam_search(lin:Kernel, rawbufs:list[Buffer], amt:int, allow_test_size=True,
     def close_pool(): beam_pool.close()
 
   min_progress = getenv("BEAM_MIN_PROGRESS", 0.01)/1e6
-  if BEAM_DEBUG: print(f"BEAM_SEARCH:\n{lin.ast}")
+  if BEAM_DEBUG:
+    print("BEAM_SEARCH:")
+    print('\n'.join(pyrender(lin.ast.replace(arg=None))))
   if DEBUG >= 2: print(f"   0.00s:                from   1 ->   1 actions {lin.colored_shape()}")
 
   try:
@@ -166,8 +166,8 @@ def beam_search(lin:Kernel, rawbufs:list[Buffer], amt:int, allow_test_size=True,
     exiting, st = False, time.perf_counter()
     dev = Device[lin.opts.device]
     while not exiting:
-      acted_lins: list[Kernel] = flatten([get_kernel_actions(lin, include_0=False).values() for lin,_ in beam])
-      timed_lins: list[tuple[Kernel, float]] = []
+      acted_lins: list[Scheduler] = flatten([get_kernel_actions(lin, include_0=False).values() for lin,_ in beam])
+      timed_lins: list[tuple[Scheduler, float]] = []
       _compile_fn = functools.partial(_try_compile_linearized_w_idx, compiler=dev.compiler)
       least_compute_ops = math.inf
       for i,proc in (map(_compile_fn, enumerate(acted_lins)) if beam_pool is None else beam_pool.imap_unordered(_compile_fn, enumerate(acted_lins))):
@@ -201,15 +201,3 @@ def beam_search(lin:Kernel, rawbufs:list[Buffer], amt:int, allow_test_size=True,
   if CACHELEVEL >= 1: diskcache_put("beam_search", key, beam[0][0].applied_opts)
   if BEAM_DEBUG: print(f"BEAM_SEARCH: final tm={time_to_str(beam[0][1], w=0)}, applied_opts={beam[0][0].applied_opts}")
   return beam[0][0]
-
-def optimize_local_size(_prg:Callable, global_size:list[int], rawbufs:list[Buffer]) -> list[int]:
-  test_rawbuffers = [Buffer(rawbufs[0].device, rawbufs[0].size, rawbufs[0].dtype).allocate(), *rawbufs[1:]] if rawbufs[0] in rawbufs[1:] else rawbufs
-  MAX_WORKGROUP = 1024
-  local_dims = [[x for x in set([sz, 1, 2, 4, 8, 16, 32, 64, 128, 256, MAX_WORKGROUP]) if x<=sz] for sz in global_size]
-  local_sizes = [list(x) for x in itertools.product(*local_dims) if prod(x) <= MAX_WORKGROUP] * 2  # try each valid size twice
-  def try_exec(local_size):
-    try: return _prg(*[x._buf for x in test_rawbuffers], global_size=[g//l if g%l == 0 else g/l for g,l in zip(global_size, local_size)], local_size=local_size, wait=True)  # noqa: E501
-    except Exception: return float('inf')
-  ret = min([(try_exec(local_size), local_size) for local_size in random.sample(local_sizes, len(local_sizes))])
-  assert not math.isinf(ret[0]), "all optimize_local_size exec failed"
-  return ret[1]

tinygrad/device.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 from dataclasses import dataclass, replace
 from collections import defaultdict
 from typing import Any, Generic, TypeVar, Iterator
-import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal, time
+import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
 from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored, \
                              Context, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup
 from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
@@ -138,15 +138,14 @@ class Buffer:
       if not self.device.startswith("DISK"): GlobalCounters.mem_used += self.nbytes
       if PROFILE:
         self._prof_num = num = len(Buffer.profile_events)
-        ts = decimal.Decimal(time.perf_counter_ns())/1000
-        Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", ts, num, {"dtype":self.dtype, "sz":self.size}))
+        Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", num, {"dtype":self.dtype, "sz":self.size}))
     return self
   def deallocate(self):
     assert hasattr(self, '_buf'), "buffer must be allocated to deallocate"
     if DEBUG is not None and DEBUG >= 7: print(f"buffer: deallocate {self.nbytes} bytes on {self.device}")
     if self._base is None and (self.options is None or self.options.external_ptr is None):
       if GlobalCounters is not None and not self.device.startswith("DISK"): GlobalCounters.mem_used -= self.nbytes
-      if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", decimal.Decimal(time.perf_counter_ns())/1000, self._prof_num))
+      if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", self._prof_num))
       self.allocator.free(self._buf, self.nbytes, self.options)
     elif self._base is not None: self._base.allocated_views -= 1
     del self._buf
@@ -304,7 +303,7 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
     if device == "METAL": return not CI
     if device in {"CUDA", "NV"}: return not CI and not getenv("PTX")
     if device in {"CPU", "LLVM"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"}
-    return device == "AMD"
+    return device in {"AMD", "PYTHON"}
   if dtype in dtypes.fp8s:
     # not supported yet - in progress
     return False

tinygrad/dtype.py

@@ -67,7 +67,7 @@ class PtrDType(DType):
     return type(self)(self.priority, self.itemsize, self.name, self.fmt, self.count, self, self._base, self.addrspace, sz, self.size)
   def ptr(self, size=-1, addrspace=AddrSpace.GLOBAL): raise RuntimeError("can't make a pointer from a pointer")
   def nbytes(self) -> int:
-    if self.size == -1: return 0  # TODO: this should be an exception
+    if self.size == -1: raise RuntimeError("can't get nbytes of a pointer with unlimited size")
     return self.size*self.itemsize
   @property
   def vcount(self): return self.v
@@ -112,12 +112,12 @@ class dtypes:
   @staticmethod
   @functools.cache
   def min(dtype:DType):
-    if dtypes.is_int(dtype): return 0 if dtypes.is_unsigned(dtype) else -2**(dtype.itemsize*8-1)
+    if dtypes.is_int(dtype): return 0 if dtypes.is_unsigned(dtype) else -2**(dtype.scalar().itemsize*8-1)
     return -float("inf") if dtypes.is_float(dtype) else False
   @staticmethod
   @functools.cache
   def max(dtype:DType):
-    if dtypes.is_int(dtype): return 2**(dtype.itemsize*8)-1+dtypes.min(dtype)
+    if dtypes.is_int(dtype): return 2**(dtype.scalar().itemsize*8)-1+dtypes.min(dtype)
     return float("inf") if dtypes.is_float(dtype) else True
   @staticmethod
   def finfo(dtype:DType) -> tuple[int, int]:
@@ -198,8 +198,9 @@ def can_safe_cast(dt0:DType, dt1:DType) -> bool:
   # https://numpy.org/doc/stable/reference/generated/numpy.can_cast.html
   if dt0 == dt1 or dt0 == dtypes.bool: return True
   match dt1:
-    case dtypes.double: return dt0 in (dtypes.float, dtypes.half, dtypes.bfloat16)
-    case dtypes.float: return dt0 in (dtypes.half, dtypes.bfloat16)
+    case dtypes.double: return dt0 in (dtypes.float, dtypes.half, dtypes.bfloat16,
+      dtypes.uint32, dtypes.uint16, dtypes.uint8, dtypes.int32, dtypes.int16, dtypes.int8)
+    case dtypes.float: return dt0 in (dtypes.half, dtypes.bfloat16, dtypes.uint16, dtypes.uint8, dtypes.int16, dtypes.int8)
     case dtypes.uint64: return dt0 in (dtypes.uint32, dtypes.uint16, dtypes.uint8)
     case dtypes.uint32: return dt0 in (dtypes.uint16, dtypes.uint8)
     case dtypes.int64: return dt0 in (dtypes.uint32, dtypes.uint16, dtypes.uint8, dtypes.int32, dtypes.int16, dtypes.int8)
@@ -298,6 +299,7 @@ truncate: dict[DType, Callable] = {dtypes.bool: bool,
 
 def _to_np_dtype(dtype:DType) -> type|None:
   import numpy as np
+  if dtype == dtypes.bfloat16: return np.float32
   return np.dtype(dtype.fmt).type if dtype.fmt is not None else None
 def _from_np_dtype(npdtype:'np.dtype') -> DType: # type: ignore [name-defined] # noqa: F821
   import numpy as np
@@ -306,9 +308,11 @@ def _from_np_dtype(npdtype:'np.dtype') -> DType: # type: ignore [name-defined] #
 @functools.cache
 def _to_torch_dtype(dtype:DType) -> 'torch.dtype'|None:  # type: ignore [name-defined] # noqa: F821
   import numpy as np, torch
+  if dtype == dtypes.uint64: return torch.uint64
+  if dtype == dtypes.bfloat16: return torch.bfloat16
   # NOTE: torch doesn't expose this mapping with a stable API
   try: return torch.from_numpy(np.array([], dtype=_to_np_dtype(dtype))).dtype
   except TypeError: return None
 @functools.cache
 def _from_torch_dtype(torchdtype:'torch.dtype') -> DType: # type: ignore [name-defined] # noqa: F821
-  return {v:k for k in dtypes.all if (v:=_to_torch_dtype(k)) is not None}[torchdtype]
+  return {v:k for k in dtypes.all if (v:=_to_torch_dtype(k)) is not None}[torchdtype]