restore that

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

@@ -121,7 +121,7 @@ runs:
         echo 'Acquire::GzipIndexes "true";' | sudo tee /etc/apt/apt.conf.d/gzip
         echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
         echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' | sudo tee -a /etc/apt/apt.conf.d/99keep-debs
-
+    
     - name: Add OpenCL Repo
       if: inputs.opencl == 'true' && runner.os == 'Linux'
       shell: bash
@@ -174,7 +174,7 @@ runs:
         if [[ "${{ inputs.llvm }}" == "true" ]]; then
           pkgs+=" libllvm20 clang-20 lld-20"
         fi
-
+        
         echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
         echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"
 
@@ -183,21 +183,21 @@ runs:
       uses: actions/cache@v4
       with:
         path: /var/cache/apt/archives/
-        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.APT_CACHE_VERSION }}
+        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}
 
     - name: Run apt Update + Install
       if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
       shell: bash
       run: |
         sudo apt -qq update || true
-
+        
         # ******** do install ********
         if [[ -n "${{ steps.apt-pkgs.outputs.pkgs }}" ]]; then
           sudo apt-get -y --allow-unauthenticated --no-install-recommends install ${{ steps.apt-pkgs.outputs.pkgs }}
         fi
-
+        
         sudo chown -R $USER:$USER /var/cache/apt/archives/
-
+    
     # **** AMD ****
     - name: Setup AMD (Linux)
       if: inputs.amd == 'true' && runner.os == 'Linux'
@@ -234,7 +234,7 @@ runs:
         cache-name: cache-gpuocelot-build
       with:
         path: ${{ github.workspace }}/gpuocelot/ocelot
-        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.BUILD_CACHE_VERSION }}
+        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-0
     - name: Clone/compile gpuocelot
       if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
       shell: bash

.github/workflows/benchmark.yml

@@ -62,8 +62,10 @@ jobs:
       run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
     - name: Run model inference benchmark
       run: METAL=1 python3.11 test/external/external_model_benchmark.py
+    - name: Run huggingface_onnx test
+      run: METAL=1 python3.11 extra/huggingface_onnx/run_models.py test --debug FacebookAI/xlm-roberta-large
     - name: Test speed vs torch
-      run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
+      run: BIG=2 MPS=1 python3.11 test/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
     - name: Test AMX tensor cores
@@ -187,7 +189,7 @@ jobs:
     - name: Run model inference benchmark
       run: NV=1 CAPTURE_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
     - name: Test speed vs torch
-      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
+      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/test_speed_v_torch.py | tee torch_speed.txt
     - name: Test speed vs theoretical
       run: NV=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
     - name: Test benchmark allreduce
@@ -389,7 +391,7 @@ jobs:
     #- name: Test speed vs torch
     #  run: |
     #    python3 -c "import torch; print(torch.__version__)"
-    #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
+    #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/test_speed_v_torch.py | tee torch_speed.txt
     - name: Test speed vs theoretical
       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
@@ -603,12 +605,12 @@ jobs:
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
     - name: reset process replay
       run: test/external/process_replay/reset.py
-    - name: benchmark openpilot 0.9.9 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
-    - name: benchmark openpilot 0.9.9 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
-    - name: benchmark openpilot 0.9.9 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: validate openpilot 0.9.7
+      run: PYTHONPATH=. FLOAT16=0 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
+    - name: benchmark openpilot 0.9.7
+      run: BENCHMARK_LOG=openpilot_0_9_7 PYTHONPATH=. QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_0_9_7.txt
+    - name: benchmark openpilot w IMAGE=2 0.9.7
+      run: BENCHMARK_LOG=openpilot_0_9_7_image PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
     - name: openpilot compile3 0.9.9 driving_vision
       run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
     - name: openpilot compile3 0.9.9 driving_policy

.github/workflows/test.yml

@@ -1,10 +1,8 @@
 name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
-  DOWNLOAD_CACHE_VERSION: '12'
-  PYTHON_CACHE_VERSION: '3'
-  APT_CACHE_VERSION: '1'
-  BUILD_CACHE_VERSION: '1'
+  DOWNLOAD_CACHE_VERSION: '10'
+  PYTHON_CACHE_VERSION: '2'
   CAPTURE_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
 
@@ -32,9 +30,9 @@ jobs:
     - name: External Benchmark Schedule
       run: PYTHONPATH="." python3 test/external/external_benchmark_schedule.py
     - name: Speed Test
-      run: LLVM=1 python3 test/speed/external_test_speed_v_torch.py
+      run: LLVM=1 python3 test/test_speed_v_torch.py
     - name: Speed Test (BEAM=2)
-      run: BEAM=2 LLVM=1 python3 test/speed/external_test_speed_v_torch.py
+      run: BEAM=2 LLVM=1 python3 test/test_speed_v_torch.py
 
   docs:
     name: Docs
@@ -48,11 +46,6 @@ jobs:
       with:
         deps: docs
         pydeps: "capstone"
-    - name: Build wheel and show size
-      run: |
-        pip install build
-        python -m build --wheel --outdir dist
-        ls -lh dist/*.whl
     - name: Use as an external package
       run: |
         mkdir $HOME/test_external_dir
@@ -336,6 +329,7 @@ jobs:
       run: |
         pip3 install --upgrade --force-reinstall ruff==0.11.0
         python3 -m ruff check .
+        python3 -m ruff check extra/onnx.py
         python3 -m ruff check examples/mlperf/ --ignore E501
     - name: Lint tinygrad with pylint
       run: python -m pylint tinygrad/
@@ -343,8 +337,7 @@ jobs:
       run: |
         python -m mypy --strict-equality --lineprecision-report .
         cat lineprecision.txt
-    - name: Run TYPED=1
-      run: TYPED=1 python -c "import tinygrad"
+        python -m mypy --strict-equality extra/onnx.py
 
   unittest:
     name: Unit Tests
@@ -382,8 +375,8 @@ jobs:
         PYTHONPATH=. 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
+    - name: Repo line count < 16000 lines
+      run: MAX_LINE_COUNT=16000 python sz.py
 
   fuzzing:
     name: Fuzzing
@@ -462,7 +455,7 @@ jobs:
   testopenpilot:
     name: 'openpilot Compile Tests'
     runs-on: ubuntu-22.04
-    timeout-minutes: 15
+    timeout-minutes: 10
     env:
       IGNORE_OOB: 0
     steps:
@@ -593,56 +586,6 @@ jobs:
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
-  testrangeify:
-    name: Linux (rangeify)
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: rangeify-minimal-llvm
-        deps: testing_minimal
-        llvm: "true"
-    - name: Test CPU=1 RANGEIFY=1
-      # TODO: add more passing tests here
-      # test_symbolic_arange_sym_step is passing now
-      # test_threefry_doesnt_use_long is because there's a contig after the long now
-      run: |
-        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
-    - 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)
-      run: LLVM=1 RANGEIFY=1 python3 -m pytest -n auto test/models/test_mnist.py --durations 20
-
-  testdevectorize:
-    name: Linux (devectorize)
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
-    env:
-      IGNORE_OOB: 0
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: devectorize-minimal
-        deps: testing_minimal
-        pydeps: "pillow"
-        llvm: "true"
-    - 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
-    - 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"
-
   testdsp:
     name: Linux (DSP)
     runs-on: ubuntu-24.04
@@ -678,6 +621,12 @@ jobs:
       run: CC=clang-20 PYTHONPATH="." 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
+    - 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
+    - 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"
 
   testwebgpu:
     name: Linux (WebGPU)
@@ -737,9 +686,9 @@ jobs:
           DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
       - name: Run LLVM test
         if: matrix.backend=='amdllvm'
-        run: python test/device/test_amd_llvm.py
+        run: python test/test_amd_llvm.py
       - name: Run pytest (amd)
-        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/device/test_hcq.py --durations=20
+        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/test_hcq.py --durations=20
       - name: Run pytest (amd)
         run: python -m pytest test/external/external_test_am.py --durations=20
       - name: Run TRANSCENDENTAL math
@@ -864,14 +813,14 @@ jobs:
         AMD: 1
         FORWARD_ONLY: 1
       run: |
-        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
+        python3 -m pytest -n=auto test/test_hcq.py test/test_tiny.py --durations=20
     - name: Run pytest (amd with llvm backend)
       env:
         MOCKGPU: 1
         AMD: 1
         FORWARD_ONLY: 1
       run: |
-        python -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py test/device/test_amd_llvm.py --durations=20
+        python -m pytest -n=auto test/test_hcq.py test/test_tiny.py test/test_amd_llvm.py --durations=20
     - name: Run pytest (ptx)
       env:
         MOCKGPU: 1
@@ -879,7 +828,7 @@ jobs:
         NV: 1
         FORWARD_ONLY: 1
       run: |
-        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
+        python3 -m pytest -n=auto test/test_hcq.py test/test_tiny.py --durations=20
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
@@ -990,18 +939,18 @@ jobs:
         env:
           HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
         run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py --durations 20
+          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py
       - name: Run REMOTE=1 Test (GPU)
         env:
           HOST: 127.0.0.1:7667*6
         run: |
-          python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py --durations 20
+          python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py
           IMAGE=2 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
       - name: Run REMOTE=1 Test (CPU)
         env:
           HOST: 127.0.0.1:8667*6
         run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py --durations 20
+          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py
       - name: Show remote server logs
         if: always()
         run: |

autogen_stubs.sh

@@ -198,7 +198,11 @@ generate_amd() {
   clang2py -k cdefstum \
     extra/hip_gpu_driver/sdma_registers.h \
     extra/hip_gpu_driver/nvd.h \
+    extra/hip_gpu_driver/kfd_pm4_headers_ai.h \
+    extra/hip_gpu_driver/soc21_enum.h \
+    extra/hip_gpu_driver/sdma_v6_0_0_pkt_open.h \
     extra/hip_gpu_driver/gc_11_0_0_offset.h \
+    extra/hip_gpu_driver/gc_10_3_0_offset.h \
     extra/hip_gpu_driver/sienna_cichlid_ip_offset.h \
     --clang-args="-I/opt/rocm/include -x c++" \
     -o $BASE/amd_gpu.py
@@ -372,6 +376,26 @@ generate_am() {
     -o $BASE/am/pm4_nv.py
   fixup $BASE/am/pm4_nv.py
 
+  clang2py -k cdefstum \
+    $AMKERN_INC/vega10_enum.h \
+    -o $BASE/am/vega10.py
+  fixup $BASE/am/vega10.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/navi10_enum.h \
+    -o $BASE/am/navi10.py
+  fixup $BASE/am/navi10.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/soc21_enum.h \
+    -o $BASE/am/soc21.py
+  fixup $BASE/am/soc21.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/soc24_enum.h \
+    -o $BASE/am/soc24.py
+  fixup $BASE/am/soc24.py
+
   clang2py -k cdefstum \
     extra/hip_gpu_driver/sdma_registers.h \
     $AMKERN_AMD/amdgpu/vega10_sdma_pkt_open.h \

docs/tensor/elementwise.md

@@ -78,7 +78,6 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.minimum
 ::: tinygrad.Tensor.where
 ::: tinygrad.Tensor.copysign
-::: tinygrad.Tensor.logaddexp
 
 ## Casting Ops
 

docs/tinybox.md

@@ -6,7 +6,7 @@ If you don't have a tinybox and you want one, see [tinygrad.org](https://tinygra
 
 ## Welcome
 
-Welcome to your tinybox! The tinybox is the universal system purpose-built for all AI infrastructure and workloads, from training to inference. The red box includes six 7900XTX GPUs, the green box includes six 4090 GPUs, and the green v2 box includes four 5090 GPUs. Whether you bought a red one or a green one, we want you to love it.
+Welcome to your tinybox! The tinybox is the universal system purpose-built for all AI infrastructure and workloads, from training to inference. The red box includes six 7900XTX GPUs, and the green box includes six 4090 GPUs. Whether you bought a red one or a green one, we want you to love it.
 
 We don't have a stupid cloud service, you don't have to create a tiny account to set it up, and we aren't tracking how you use the box. We're just happy you bought one. This petaflop is your petaflop.
 

examples/beautiful_mnist.py

@@ -21,7 +21,7 @@ if __name__ == "__main__":
   X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))
 
   model = Model()
-  opt = (nn.optim.Adam if not getenv("MUON") else nn.optim.Muon)(nn.state.get_parameters(model))
+  opt = nn.optim.Adam(nn.state.get_parameters(model))
 
   @TinyJit
   @Tensor.train()
@@ -29,7 +29,8 @@ if __name__ == "__main__":
     opt.zero_grad()
     samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
     loss = model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
-    return loss.realize(*opt.schedule_step())
+    opt.step()
+    return loss
 
   @TinyJit
   def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100

examples/hlb_cifar10.py

@@ -118,7 +118,7 @@ class SpeedyResNet:
 # hyper-parameters were exactly the same as the original repo
 bias_scaler = 58
 hyp = {
-  'seed' : 201,
+  'seed' : 200,
   'opt': {
     'bias_lr':            1.76 * bias_scaler/512,
     'non_bias_lr':        1.76 / 512,

examples/mlperf/model_train.py

@@ -1297,9 +1297,6 @@ def train_llama3():
   SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
   TRAIN_ON_VAL       = config["TRAIN_ON_VAL"]           = getenv("TRAIN_ON_VAL", 0)
   SAMPLES            = config["SAMPLES"]                = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000 * 1152)
-  EVAL_FREQ          = config["EVAL_FREQ"]              = getenv("EVAL_FREQ", 46080)
-  EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 16)
-  EVAL_TARGET        = config["EVAL_TARGET"]            = getenv("EVAL_TARGET", 5.6)
 
   # LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 FUSE_ARANGE=1 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
   # trains to 7
@@ -1321,10 +1318,6 @@ def train_llama3():
   if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
   model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
 
-  if getenv("FAKEDATA"):
-    for v in get_parameters(model):
-      v = v.assign(Tensor.empty(v.shape))
-
   if (DP := getenv("DP", 1)) > 1:
     device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
     for v in get_parameters(model):
@@ -1346,8 +1339,6 @@ def train_llama3():
       else:
         # attention_norm, ffn_norm, norm
         v.shard_(device, axis=None)
-      # prevents memory spike on device 0
-      v.realize()
 
   optim = AdamW(get_parameters(model), lr=0.0,
                 b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
@@ -1378,7 +1369,7 @@ def train_llama3():
         total_norm += p.grad.float().square().sum()
       total_norm = total_norm.sqrt().contiguous()
       for p in optim.params:
-        p.grad = p.grad * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)
+        p.grad = p.grad * opt_gradient_clip_norm / (total_norm + 1e-6)
 
     optim.step()
     scheduler.step()
@@ -1387,40 +1378,16 @@ def train_llama3():
     loss.realize(lr)
     return loss, lr
 
-  @TinyJit
-  @Tensor.train(False)
-  def eval_step(model, tokens:Tensor):
-    if (DP := getenv("DP", 1)) > 1:
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.shard(device, 0)
-    if (MP := getenv("MP", 1)) > 1:
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
-      tokens = tokens.shard(device)
-    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    return loss.flatten().float()
+  if getenv("FAKEDATA", 0):
+    def fake_data():
+      for _ in range(SAMPLES // GBS):
+        yield Tensor.randint(GBS, SEQLEN + 1, low=0, high=32000, dtype=dtypes.int32, device=Device.DEFAULT)
+    iter = fake_data()
+  else:
+    from examples.mlperf.dataloader import batch_load_llama3
+    iter = batch_load_llama3(GBS, SAMPLES, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=bool(TRAIN_ON_VAL))
 
-  # ** 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)
-
-  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))
-
-  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)
-
-  iter = get_train_iter()
-  i, sequences_seen = 0, 0
+  i = 0
   for tokens in tqdm(iter, total=SAMPLES//GBS):
     t = time.perf_counter()
     GlobalCounters.reset()
@@ -1435,33 +1402,9 @@ def train_llama3():
     if getenv("CKPT") and (i % 200 == 0 or i == 10):
       tqdm.write("saving checkpoint")
       if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
-      fn = f"{ckpt_dir}/llama3_{i}.safe"
+      fn = f"{ckpt_dir}/{i}.safe"
       safe_save(get_state_dict(model), fn)
-
     i += 1
-    sequences_seen += tokens.shape[0]
-
-    if sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1):
-      tqdm.write(f"evaluating after {sequences_seen} sequences")
-
-      # run eval
-      eval_losses = []
-      eval_iter = get_eval_iter()
-      tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")
-
-      for tokens in tqdm(eval_iter, total=5760//EVAL_BS):
-        eval_losses += eval_step(model, tokens).tolist()
-      log_perplexity = Tensor(eval_losses).mean().float().item()
-
-      tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")
-
-      if log_perplexity < EVAL_TARGET:
-        tqdm.write(f"target achieved after {sequences_seen} sequences")
-        if getenv("CKPT"):
-          if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
-          fn = f"{ckpt_dir}/llama3.safe"
-          safe_save(get_state_dict(model), fn)
-        break
 
 if __name__ == "__main__":
   multiprocessing.set_start_method('spawn')

extra/amdpci/proclogs.py

@@ -1,16 +1,6 @@
-import re, ctypes, sys, importlib
+import re, ctypes, sys
 
-from tinygrad.runtime.support.am.amdev import AMDev, AMRegister
-class AMDFake(AMDev):
-  def __init__(self, devfmt, vram, doorbell, mmio, dma_regions=None):
-    self.devfmt, self.vram, self.doorbell64, self.mmio, self.dma_regions = devfmt, vram, doorbell, mmio, dma_regions
-    self._run_discovery()
-    self._build_regs()
-
-amdev = importlib.import_module("tinygrad.runtime.support.am.amdev")
-amdev.AMDev = AMDFake
-
-from tinygrad.runtime.ops_amd import PCIIface
+from tinygrad.runtime.autogen.am import am, mp_11_0, mp_13_0_0, nbio_4_3_0, mmhub_3_0_0, gc_11_0_0, osssys_6_0_0
 
 def parse_amdgpu_logs(log_content, register_names=None):
   register_map = register_names
@@ -33,11 +23,16 @@ def parse_amdgpu_logs(log_content, register_names=None):
   return processed_log
 
 def main():
+  regs_offset = {13: {0: [3072, 37784576]}, 28: {0: [93184, 37754880], 1: [201327616, 201461760], 2: [209716224, 209850368], 3: [218104832, 218238976], 4: [226493440, 226627584], 5: [234882048, 235016192], 6: [243270656, 243404800]}, 21: {0: [28672, 12582912, 37795840, 130023424, 306184192], 1: [201326592, 201463808, 201465856, 204210176, 204472320], 2: [209715200, 209852416, 209854464, 212598784, 212860928], 3: [218103808, 218241024, 218243072, 220987392, 221249536], 4: [226492416, 226629632, 226631680, 229376000, 229638144], 5: [234881024, 235018240, 235020288, 237764608, 238026752], 6: [243269632, 243406848, 243408896, 246153216, 246415360]}, 22: {0: [18, 192, 13504, 36864, 37764096]}, 1: {0: [4704, 40960, 114688, 37760000]}, 2: {0: [3872, 37790720]}, 11: {0: [70656, 38103040]}, 12: {0: [106496, 37783552]}, 15: {0: [90112, 14417920, 14680064, 14942208, 38009856]}, 16: {0: [90112, 14417920, 14680064, 14942208, 38009856]}, 14: {0: [0, 20, 3360, 66560, 37859328, 67371008]}, 26: {0: [0, 20, 3360, 66560, 37859328, 67371008]}, 23: {0: [4256, 37789696]}, 33: {0: [0, 20, 3360, 66560, 37859328, 67371008]}, 25: {0: []}, 3: {0: [4704, 40960, 114688, 37760000]}, 4: {0: [4704, 40960, 114688, 37760000]}, 24: {0: [92160, 92672, 37752832, 54788096]}, 27: {0: [91648, 37751808], 1: [201339904, 201458176], 2: [209728512, 209846784], 3: [218117120, 218235392], 4: [226505728, 226624000], 5: [234894336, 235012608], 6: [243282944, 243401216]}, 29: {0: [201342976, 201344000, 205520896, 205537280], 1: [209731584, 209732608, 213909504, 213925888], 2: [218120192, 218121216, 222298112, 222314496], 3: [226508800, 226509824, 230686720, 230703104], 4: [234897408, 234898432, 239075328, 239091712], 5: [243286016, 243287040, 247463936, 247480320]}, 17: {0: [30720, 32256], 1: [31488, 73728]}}
+
   reg_names = {}
-  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}"
+  def _prepare_registers(modules):
+    for base, m in modules:
+      for k, regval in m.__dict__.items():
+        if k.startswith("reg") and not k.endswith("_BASE_IDX") and (base_idx:=getattr(m, f"{k}_BASE_IDX", None)) is not None:
+          reg_names[regs_offset[am.__dict__.get(f"{base}_HWIP")][0][base_idx] + regval] = k
+
+  _prepare_registers([("MP0", mp_13_0_0), ("NBIO", nbio_4_3_0), ("MMHUB", mmhub_3_0_0), ("GC", gc_11_0_0), ("OSSSYS", osssys_6_0_0)])
 
   with open(sys.argv[1], 'r') as f:
     log_content = log_content_them = f.read()

extra/gemm/amd_uop_matmul.py

@@ -5,7 +5,7 @@ from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import getenv, colored, prod, unwrap
 from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad.shape.view import strides_for_shape
-from tinygrad.codegen.opt.kernel import axis_colors, Opt, OptOps
+from tinygrad.codegen.opt.kernel import axis_colors
 from tinygrad.codegen.opt.swizzler import merge_views, view_left
 
 def to_colored(full_shape, axis_types): return '_'.join([colored(str(s), axis_colors[at]) for s,at in zip(full_shape, axis_types)])
@@ -44,21 +44,6 @@ pm = PatternMatcher([
   (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), swizzle_reduceop),
 ])
 
-def rangeify_kernel3():
-  a = Tensor.empty(N,N)
-  b = Tensor.empty(N,N)
-  c = a@b
-  #c = c.reshape((32,2,16,4,32,2,16,4)).contiguous()
-  with Context(RANGEIFY=1):
-    sink = c.schedule()[-1].ast
-  #print(sink)
-
-  opts  = [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.UPCAST, 0, 2)]
-  opts += [Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 1, 16), Opt(OptOps.UPCAST, 1, 2)]
-  opts += [Opt(OptOps.UNROLL, 0, 8)]
-
-  return sink.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
-
 def top_spec_kernel3():
   a = Tensor.empty(N,N)
   b = Tensor.empty(N,N)
@@ -324,15 +309,10 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
 
 if __name__ == "__main__":
   HL = getenv("HL")
-  if HL == 3: hprg = rangeify_kernel3()
-  elif HL == 2: hprg = top_spec_kernel3()
+  if HL == 2: hprg = top_spec_kernel3()
   elif HL == 1: hprg = hl_spec_kernel3()
   else: hprg = hand_spec_kernel3()
-  if HL == 3:
-    with Context(RANGEIFY=1, BLOCK_REORDER=0):
-      prg = get_program(hprg, Device.default.renderer)
-  else:
-    prg = get_program(hprg, Device.default.renderer)
+  prg = get_program(hprg, Device.default.renderer)
   print(prg.src)
   if getenv("SRC"): exit(0)
   hrunner = CompiledRunner(prg)

extra/huggingface_onnx/README.md

@@ -1,61 +0,0 @@
-# HuggingFace ONNX
-
-Tool for discovering, downloading, and validating ONNX models from HuggingFace.
-
-## Extra Dependencies
-
-```bash
-pip install huggingface_hub pyyaml requests onnx onnxruntime numpy
-```
-
-## Huggingface Manager (discovering and downloading)
-
-The `huggingface_manager.py` script discovers top ONNX models from HuggingFace, collects metadata, and optionally downloads them.
-
-```bash
-# Download top 50 models sorted by downloads
-python huggingface_manager.py --limit 50 --download
-
-# Just collect metadata (no download)
-python huggingface_manager.py --limit 100
-
-# Sort by likes instead of downloads
-python huggingface_manager.py --limit 20 --sort likes --download
-
-# Custom output file
-python huggingface_manager.py --limit 10 --output my_models.yaml
-```
-
-### Output Format
-
-The tool generates a YAML file with the following structure:
-
-```yaml
-repositories:
-  "model-name":
-    url: "https://huggingface.co/model-name"
-    download_path: "/path/to/models/..."  # when --download used
-    files:
-      - file: "model.onnx"
-        size: "90.91MB"
-total_size: "2.45GB"
-created_at: "2024-01-15T10:30:00Z"
-```
-
-## Run Models (validation)
-
-The `run_models.py` script validates ONNX models against ONNX Runtime for correctness.
-
-```bash
-# Validate models from a YAML configuration file
-python run_models.py --validate huggingface_repos.yaml
-
-# Debug specific repository (downloads and validates all ONNX models)
-python run_models.py --debug sentence-transformers/all-MiniLM-L6-v2
-
-# Debug specific model file
-python run_models.py --debug sentence-transformers/all-MiniLM-L6-v2/onnx/model.onnx
-
-# Debug with model truncation for debugging and validating intermediate results
-DEBUGONNX=1 python run_models.py --debug sentence-transformers/all-MiniLM-L6-v2/onnx/model.onnx --truncate 10
-```

extra/huggingface_onnx/collect_metadata.py

@@ -0,0 +1,85 @@
+import yaml, time, requests, argparse
+from pathlib import Path
+from huggingface_hub import list_models, HfApi
+from tinygrad.helpers import tqdm
+
+HUGGINGFACE_URL = "https://huggingface.co"
+SKIPPED_FILES = [
+  "fp16", "int8", "uint8", "quantized",      # numerical accuracy issues
+  "avx2", "arm64", "avx512", "avx512_vnni",  # numerical accuracy issues
+  "q4", "q4f16", "bnb4",                     # unimplemented quantization
+  "model_O4",                                # requires non cpu ort runner and MemcpyFromHost op
+  "merged",                                  # TODO implement attribute with graph type and Loop op
+]
+SKIPPED_REPO_PATHS = [
+  # Invalid model-index
+  "AdamCodd/vit-base-nsfw-detector",
+  # TODO: implement attribute with graph type and Loop op
+  "minishlab/potion-base-8M", "minishlab/M2V_base_output", "minishlab/potion-retrieval-32M",
+  # TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, GroupQueryAttention
+  "HuggingFaceTB/SmolLM2-360M-Instruct",
+  # TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, RotaryEmbedding, MultiHeadAttention
+  "HuggingFaceTB/SmolLM2-1.7B-Instruct",
+  # TODO: implmement RandomNormalLike
+  "stabilityai/stable-diffusion-xl-base-1.0", "stabilityai/sdxl-turbo", 'SimianLuo/LCM_Dreamshaper_v7',
+  # TODO: implement NonZero
+  "mangoapps/fb_zeroshot_mnli_onnx",
+  # TODO huge Concat in here with 1024 (1, 3, 32, 32) Tensors, and maybe a MOD bug with const folding
+  "briaai/RMBG-2.0",
+]
+
+def get_top_repos(n: int, sort: str) -> list[str]: # list["FacebookAI/xlm-roberta-large", ...]
+  print(f"** Getting top {n} models sorted by {sort} **")
+  repos = []
+  i = 0
+  for model in list_models(filter="onnx", sort=sort):
+    if model.id in SKIPPED_REPO_PATHS: continue
+    print(f"{i+1}/{n}: {model.id} ({getattr(model, sort)})")
+    repos.append(model.id)
+    i += 1
+    if i == n: break
+  return repos
+
+def get_metadata(repos:list[str]) -> dict:
+  api = HfApi()
+  repos_metadata = {"repositories": {}}
+  total_size = 0
+
+  # TODO: speed head requests up with async?
+  for repo in tqdm(repos, desc="Getting metadata"):
+    files_metadata = []
+    model_info = api.model_info(repo)
+
+    for file in model_info.siblings:
+      filename = file.rfilename
+      if not (filename.endswith('.onnx') or filename.endswith('.onnx_data')): continue
+      if any(skip_str in filename for skip_str in SKIPPED_FILES): continue
+      head = requests.head(f"{HUGGINGFACE_URL}/{repo}/resolve/main/{filename}", allow_redirects=True)
+      file_size = file.size or int(head.headers.get('Content-Length', 0))
+      files_metadata.append({"file": filename, "size": f"{file_size/1e6:.2f}MB"})
+      total_size += file_size
+
+    repos_metadata["repositories"][repo] = {
+      "url": f"{HUGGINGFACE_URL}/{repo}",
+      "download_path": None,
+      "files": files_metadata,
+    }
+  repos_metadata['total_size'] = f"{total_size/1e9:.2f}GB"
+  repos_metadata['created_at'] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
+  return repos_metadata
+
+if __name__ == "__main__":
+  sort = "downloads" # recent 30 days downloads
+  huggingface_onnx_dir = Path(__file__).parent
+
+  parser = argparse.ArgumentParser(description="Produces a YAML file with metadata of top huggingface onnx models")
+  parser.add_argument("--limit", type=int, required=True, help="Number of top repositories to process (e.g., 100)")
+  parser.add_argument("--output", type=str, default="huggingface_repos.yaml", help="Output YAML file name to save the report")
+  args = parser.parse_args()
+
+  top_repos = get_top_repos(args.limit, sort)
+  metadata = get_metadata(top_repos)
+  yaml_path = huggingface_onnx_dir / args.output
+  with open(yaml_path, 'w') as f:
+    yaml.dump(metadata, f, sort_keys=False)
+    print(f"YAML saved to: {str(yaml_path)}")

extra/huggingface_onnx/download_models.py

@@ -0,0 +1,29 @@
+import yaml, argparse
+from pathlib import Path
+from huggingface_hub import snapshot_download
+
+def download_models(yaml_file: str, download_dir: str) -> None:
+  with open(yaml_file, 'r') as f: metadata = yaml.safe_load(f)
+  n = len(metadata["repositories"])
+
+  for i, (model_id, model_data) in enumerate(metadata["repositories"].items()):
+    print(f"Downloading {i+1}/{n}: {model_id}...")
+    allow_patterns = [file_info["file"] for file_info in model_data["files"]]
+    root_path = Path(snapshot_download(repo_id=model_id, allow_patterns=allow_patterns, cache_dir=download_dir))
+    # download configs too (the sizes are small)
+    snapshot_download(repo_id=model_id, allow_patterns=["*config.json"], cache_dir=download_dir)
+    print(f"Downloaded model files to: {root_path}")
+    model_data["download_path"] = str(root_path)
+
+  # Save the updated metadata back to the YAML file
+  with open(yaml_file, 'w') as f: yaml.dump(metadata, f, sort_keys=False)
+  print("Download completed according to YAML file.")
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser(description="Download models from Huggingface Hub based on a YAML configuration file.")
+  parser.add_argument("input", type=str, help="Path to the input YAML configuration file containing model information.")
+  args = parser.parse_args()
+
+  models_folder = Path(__file__).parent / "models"
+  models_folder.mkdir(parents=True, exist_ok=True)
+  download_models(args.input, str(models_folder))

extra/huggingface_onnx/huggingface_manager.py

@@ -1,230 +0,0 @@
-import yaml
-import time
-import requests
-import argparse
-from pathlib import Path
-from huggingface_hub import list_models, HfApi, snapshot_download
-from tinygrad.helpers import _ensure_downloads_dir
-DOWNLOADS_DIR = _ensure_downloads_dir() / "models"
-from tinygrad.helpers import tqdm
-
-def snapshot_download_with_retry(*, repo_id: str, allow_patterns: list[str]|tuple[str, ...]|None=None, cache_dir: str|Path|None=None,
-                                 tries: int=2, **kwargs) -> Path:
-  for attempt in range(tries):
-    try:
-      return Path(snapshot_download(
-        repo_id=repo_id,
-        allow_patterns=allow_patterns,
-        cache_dir=str(cache_dir) if cache_dir is not None else None,
-        **kwargs
-      ))
-    except Exception as e:
-      if attempt == tries-1: raise
-      time.sleep(1)
-
-# Constants for filtering models
-HUGGINGFACE_URL = "https://huggingface.co"
-SKIPPED_FILES = [
-  "fp16", "int8", "uint8", "quantized",      # numerical accuracy issues
-  "avx2", "arm64", "avx512", "avx512_vnni",  # numerical accuracy issues
-  "q4", "q4f16", "bnb4",                     # unimplemented quantization
-  "model_O4",                                # requires non cpu ort runner and MemcpyFromHost op
-  "merged",                                  # TODO implement attribute with graph type and Loop op
-]
-
-SKIPPED_REPO_PATHS = [
-  # Invalid model-index
-  "AdamCodd/vit-base-nsfw-detector",
-  # TODO: implement attribute with graph type and Loop op
-  "minishlab/potion-base-8M", "minishlab/M2V_base_output", "minishlab/potion-retrieval-32M",
-  # TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, GroupQueryAttention
-  "HuggingFaceTB/SmolLM2-360M-Instruct",
-  # TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, RotaryEmbedding, MultiHeadAttention
-  "HuggingFaceTB/SmolLM2-1.7B-Instruct",
-  # TODO: implement RandomNormalLike
-  "stabilityai/stable-diffusion-xl-base-1.0", "stabilityai/sdxl-turbo", 'SimianLuo/LCM_Dreamshaper_v7',
-  # TODO: implement NonZero
-  "mangoapps/fb_zeroshot_mnli_onnx",
-  # TODO huge Concat in here with 1024 (1, 3, 32, 32) Tensors, and maybe a MOD bug with const folding
-  "briaai/RMBG-2.0",
-]
-
-
-class HuggingFaceONNXManager:
-  def __init__(self):
-    self.base_dir = Path(__file__).parent
-    self.models_dir = DOWNLOADS_DIR
-    self.api = HfApi()
-
-  def discover_models(self, limit: int, sort: str = "downloads") -> list[str]:
-    print(f"Discovering top {limit} ONNX models sorted by {sort}...")
-    repos = []
-    i = 0
-
-    for model in list_models(filter="onnx", sort=sort):
-      if model.id in SKIPPED_REPO_PATHS:
-        continue
-
-      print(f"  {i+1}/{limit}: {model.id} ({getattr(model, sort)})")
-      repos.append(model.id)
-      i += 1
-      if i == limit:
-        break
-
-    print(f"Found {len(repos)} suitable ONNX models")
-    return repos
-
-  def collect_metadata(self, repos: list[str]) -> dict:
-    print(f"Collecting metadata for {len(repos)} repositories...")
-    metadata = {"repositories": {}}
-    total_size = 0
-
-    for repo in tqdm(repos, desc="Collecting metadata"):
-      try:
-        files_metadata = []
-        model_info = self.api.model_info(repo)
-
-        for file in model_info.siblings:
-          filename = file.rfilename
-          if not (filename.endswith('.onnx') or filename.endswith('.onnx_data')):
-            continue
-          if any(skip_str in filename for skip_str in SKIPPED_FILES):
-            continue
-
-          # Get file size from API or HEAD request
-          try:
-            head = requests.head(
-              f"{HUGGINGFACE_URL}/{repo}/resolve/main/{filename}",
-              allow_redirects=True,
-              timeout=10
-            )
-            file_size = file.size or int(head.headers.get('Content-Length', 0))
-          except requests.RequestException:
-            file_size = file.size or 0
-
-          files_metadata.append({
-            "file": filename,
-            "size": f"{file_size/1e6:.2f}MB"
-          })
-          total_size += file_size
-
-        if files_metadata:  # Only add repos with valid ONNX files
-          metadata["repositories"][repo] = {
-            "url": f"{HUGGINGFACE_URL}/{repo}",
-            "download_path": None,
-            "files": files_metadata,
-          }
-
-      except Exception as e:
-        print(f"WARNING: Failed to collect metadata for {repo}: {e}")
-        continue
-
-    metadata['total_size'] = f"{total_size/1e9:.2f}GB"
-    metadata['created_at'] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
-
-    print(f"Collected metadata for {len(metadata['repositories'])} repositories")
-    print(f"Total estimated download size: {metadata['total_size']}")
-
-    return metadata
-
-  def download_models(self, metadata: dict) -> dict:
-    self.models_dir.mkdir(parents=True, exist_ok=True)
-
-    repos = metadata["repositories"]
-    n = len(repos)
-
-    print(f"Downloading {n} repositories to {self.models_dir}...")
-
-    for i, (model_id, model_data) in enumerate(repos.items()):
-      print(f"  Downloading {i+1}/{n}: {model_id}...")
-
-      try:
-        # Download ONNX model files
-        allow_patterns = [file_info["file"] for file_info in model_data["files"]]
-        root_path = snapshot_download_with_retry(
-          repo_id=model_id,
-          allow_patterns=allow_patterns,
-          cache_dir=str(self.models_dir)
-        )
-
-        # Download config files (usually small)
-        snapshot_download_with_retry(
-          repo_id=model_id,
-          allow_patterns=["*config.json"],
-          cache_dir=str(self.models_dir)
-        )
-
-        model_data["download_path"] = str(root_path)
-        print(f"    Downloaded to: {root_path}")
-
-      except Exception as e:
-        print(f"    ERROR: Failed to download {model_id}: {e}")
-        model_data["download_path"] = None
-        continue
-
-    successful_downloads = sum(1 for repo in repos.values() if repo["download_path"] is not None)
-    print(f"Successfully downloaded {successful_downloads}/{n} repositories")
-    print(f"All models saved to: {self.models_dir}")
-
-    return metadata
-
-  def save_metadata(self, metadata: dict, output_file: str):
-    yaml_path = self.base_dir / output_file
-    with open(yaml_path, 'w') as f:
-      yaml.dump(metadata, f, sort_keys=False)
-    print(f"Metadata saved to: {yaml_path}")
-
-  def discover_and_download(self, limit: int, output_file: str = "huggingface_repos.yaml",
-                          sort: str = "downloads", download: bool = True):
-    print(f"Starting HuggingFace ONNX workflow...")
-    print(f"   Limit: {limit} models")
-    print(f"   Sort by: {sort}")
-    print(f"   Download: {'Yes' if download else 'No'}")
-    print(f"   Output: {output_file}")
-    print("-" * 50)
-
-    repos = self.discover_models(limit, sort)
-
-    metadata = self.collect_metadata(repos)
-
-    if download:
-      metadata = self.download_models(metadata)
-
-    self.save_metadata(metadata, output_file)
-
-    print("-" * 50)
-    print("Workflow completed successfully!")
-    if download:
-      successful = sum(1 for repo in metadata["repositories"].values()
-                     if repo["download_path"] is not None)
-      print(f"{successful}/{len(metadata['repositories'])} models downloaded")
-
-    return metadata
-
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser(
-    description="HuggingFace ONNX Model Manager - Discover, collect metadata, and download ONNX models",
-  )
-
-  parser.add_argument("--limit", type=int, help="Number of top repositories to process")
-  parser.add_argument("--output", type=str, default="huggingface_repos.yaml",
-                     help="Output YAML file name (default: huggingface_repos.yaml)")
-  parser.add_argument("--sort", type=str, default="downloads",
-                     choices=["downloads", "likes", "created", "modified"],
-                     help="Sort criteria for model discovery (default: downloads)")
-
-  parser.add_argument("--download", action="store_true", default=False,
-                     help="Download models after collecting metadata")
-
-  args = parser.parse_args()
-
-  if not args.limit: parser.error("--limit is required")
-
-  manager = HuggingFaceONNXManager()
-  manager.discover_and_download(
-    limit=args.limit,
-    output_file=args.output,
-    sort=args.sort,
-    download=args.download
-  )

extra/huggingface_onnx/run_models.py

@@ -1,11 +1,10 @@
-import onnx, yaml, tempfile, time, argparse, json
+import onnx, yaml, tempfile, time, collections, pprint, argparse, json
 from pathlib import Path
-from typing import Any
 from tinygrad.frontend.onnx import OnnxRunner
+from extra.onnx import get_onnx_ops
 from extra.onnx_helpers import validate, get_example_inputs
-from extra.huggingface_onnx.huggingface_manager import DOWNLOADS_DIR, snapshot_download_with_retry
 
-def get_config(root_path: Path) -> dict[str, Any]:
+def get_config(root_path: Path):
   ret = {}
   for path in root_path.rglob("*config.json"):
     config = json.load(path.open())
@@ -13,19 +12,19 @@ def get_config(root_path: Path) -> dict[str, Any]:
       ret.update(config)
   return ret
 
-def get_tolerances(file_name: str) -> tuple[float, float]:
+def run_huggingface_validate(onnx_model_path, config, rtol, atol):
+  onnx_runner = OnnxRunner(onnx_model_path)
+  inputs = get_example_inputs(onnx_runner.graph_inputs, config)
+  validate(onnx_model_path, inputs, rtol=rtol, atol=atol)
+
+def get_tolerances(file_name): # -> rtol, atol
   # TODO very high rtol atol
   if "fp16" in file_name: return 9e-2, 9e-2
   if any(q in file_name for q in ["int8", "uint8", "quantized"]): return 4, 4
   return 4e-3, 3e-2
 
-def run_huggingface_validate(onnx_model_path: str | Path, config: dict[str, Any], rtol: float, atol: float):
-  onnx_runner = OnnxRunner(onnx_model_path)
-  inputs = get_example_inputs(onnx_runner.graph_inputs, config)
-  validate(onnx_model_path, inputs, rtol=rtol, atol=atol)
-
 def validate_repos(models:dict[str, tuple[Path, Path]]):
-  print(f"** Validating {len(models)} models **")
+  print(f"** Validating {len(model_paths)} models **")
   for model_id, (root_path, relative_path) in models.items():
     print(f"validating model {model_id}")
     model_path = root_path / relative_path
@@ -37,6 +36,25 @@ def validate_repos(models:dict[str, tuple[Path, Path]]):
     et = time.time() - st
     print(f"passed, took {et:.2f}s")
 
+def retrieve_op_stats(models:dict[str, tuple[Path, Path]]) -> dict:
+  ret = {}
+  op_counter = collections.Counter()
+  unsupported_ops = collections.defaultdict(set)
+  supported_ops = get_onnx_ops()
+  print(f"** Retrieving stats from {len(model_paths)} models **")
+  for model_id, (root_path, relative_path) in models.items():
+    print(f"examining {model_id}")
+    model_path = root_path / relative_path
+    onnx_runner = OnnxRunner(model_path)
+    for node in onnx_runner.graph_nodes:
+      op_counter[node.op] += 1
+      if node.op not in supported_ops:
+        unsupported_ops[node.op].add(model_id)
+    del onnx_runner
+  ret["unsupported_ops"] = {k:list(v) for k, v in unsupported_ops.items()}
+  ret["op_counter"] = op_counter.most_common()
+  return ret
+
 def debug_run(model_path, truncate, config, rtol, atol):
   if truncate != -1:
     model = onnx.load(model_path)
@@ -53,9 +71,12 @@ def debug_run(model_path, truncate, config, rtol, atol):
     run_huggingface_validate(model_path, config, rtol, atol)
 
 if __name__ == "__main__":
-  parser = argparse.ArgumentParser(description="Huggingface ONNX Model Validator")
-  parser.add_argument("--validate", type=str, default="",
-                      help="Validate correctness of models from the specified YAML configuration file")
+  parser = argparse.ArgumentParser(description="Huggingface ONNX Model Validator and Ops Checker")
+  parser.add_argument("input", type=str, help="Path to the input YAML configuration file containing model information.")
+  parser.add_argument("--check_ops", action="store_true", default=False,
+                      help="Check support for ONNX operations in models from the YAML file")
+  parser.add_argument("--validate", action="store_true", default=False,
+                      help="Validate correctness of models from the YAML file")
   parser.add_argument("--debug", type=str, default="",
                       help="""Validates without explicitly needing a YAML or models pre-installed.
                       provide repo id (e.g. "minishlab/potion-base-8M") to validate all onnx models inside the repo
@@ -64,13 +85,13 @@ if __name__ == "__main__":
   parser.add_argument("--truncate", type=int, default=-1, help="Truncate the ONNX model so intermediate results can be validated")
   args = parser.parse_args()
 
-  if not (args.validate or args.debug):
-    parser.error("Please provide either --validate <yaml_file> or --debug <repo_id>.")
+  if not (args.check_ops or args.validate or args.debug):
+    parser.error("Please provide either --validate, --check_ops, or --debug.")
   if args.truncate != -1 and not args.debug:
     parser.error("--truncate and --debug should be used together for debugging")
 
-  if args.validate:
-    with open(args.validate, 'r') as f:
+  if args.check_ops or args.validate:
+    with open(args.input, 'r') as f:
       data = yaml.safe_load(f)
       assert all(repo["download_path"] is not None for repo in data["repositories"].values()), "please run `download_models.py` for this yaml"
       model_paths = {
@@ -80,16 +101,22 @@ if __name__ == "__main__":
         if model["file"].endswith(".onnx")
       }
 
-    validate_repos(model_paths)
+    if args.check_ops:
+      pprint.pprint(retrieve_op_stats(model_paths))
+
+    if args.validate:
+      validate_repos(model_paths)
 
   if args.debug:
+    from huggingface_hub import snapshot_download
+    download_dir = Path(__file__).parent / "models"
     path:list[str] = args.debug.split("/")
     if len(path) == 2:
       # repo id
       # validates all onnx models inside repo
       repo_id = "/".join(path)
-      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], cache_dir=DOWNLOADS_DIR)
-      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
+      root_path = Path(snapshot_download(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], cache_dir=download_dir))
+      snapshot_download(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=download_dir)
       config = get_config(root_path)
       for onnx_model in root_path.rglob("*.onnx"):
         rtol, atol = get_tolerances(onnx_model.name)
@@ -101,8 +128,8 @@ if __name__ == "__main__":
       onnx_model = path[-1]
       assert path[-1].endswith(".onnx")
       repo_id, relative_path = "/".join(path[:2]), "/".join(path[2:])
-      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=[relative_path], cache_dir=DOWNLOADS_DIR)
-      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
+      root_path = Path(snapshot_download(repo_id=repo_id, allow_patterns=[relative_path], cache_dir=download_dir))
+      snapshot_download(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=download_dir)
       config = get_config(root_path)
       rtol, atol = get_tolerances(onnx_model)
       print(f"validating {relative_path} with truncate={args.truncate}, {rtol=}, {atol=}")

extra/onnx_helpers.py

@@ -1,6 +1,7 @@
 from tinygrad import Tensor
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.frontend.onnx import OnnxRunner, OnnxValue
+from tinygrad.frontend.onnx import OnnxRunner
+from extra.onnx import OnnxValue
 import numpy as np
 import onnxruntime as ort
 

extra/torch_backend/backend.py

@@ -381,7 +381,6 @@ decomps = [
   aten.elu,  # elu has a scale + input_scale param
   aten.elu_backward,
   aten.softplus,
-  aten.logaddexp,
   aten.threshold,
   aten.nll_loss_forward,
   aten.nll_loss_backward,

extra/torch_muon.py

@@ -1,75 +0,0 @@
-import torch
-
-#credit to KellerJordan at https://github.com/KellerJordan/Muon/tree/master
-#some changes: classic momentum instead of weighting gradient
-#added ns_steps, ns_params, nesterov as hyperparams
-def zeropower_via_newtonschulz5(G:torch.tensor, steps:int, params:tuple[int, ...]):
-  """
-  Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
-  quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
-  of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
-  zero even beyond the point where the iteration no longer converges all the way to one everywhere
-  on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
-  where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
-  performance at all relative to UV^T, where USV^T = G is the SVD.
-  """
-  assert G.ndim >= 2 # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng
-
-  a, b, c = params
-  X = G
-  if G.size(-2) > G.size(-1):
-    X = X.mT
-
-  # Ensure spectral norm is at most 1
-  X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
-  # Perform the NS iterations
-  for _ in range(steps):
-    A = X @ X.mT
-    B = b * A + c * A @ A # quintic computation strategy adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng
-    X = a * X + B @ X
-
-  if G.size(-2) > G.size(-1):
-    X = X.mT
-
-  return X
-
-def muon_update(grad, momentum, beta=0.95, ns_steps=5, ns_params=(3.4445, -4.7750,  2.0315), nesterov=True):
-  if beta:
-    momentum.mul_(beta).add_(grad)
-    update = grad.add(momentum,alpha=beta) if nesterov else momentum
-  else: update = grad
-  if update.ndim == 4: # for the case of conv filters
-    update = update.view(len(update), -1)
-  update = zeropower_via_newtonschulz5(update, steps=ns_steps, params=ns_params)
-  return update
-
-class SingleDeviceMuon(torch.optim.Optimizer):
-  """
-  Muon variant for usage in non-distributed settings.
-  """
-  def __init__(self, params, lr=0.02, weight_decay=0.0, momentum=0.95, ns_steps=5, ns_params=(3.4445, -4.7750,  2.0315), nesterov=True):
-    defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum, ns_steps=ns_steps, ns_params=ns_params, nesterov=nesterov)
-    super().__init__(params, defaults)
-
-  @torch.no_grad()
-  def step(self, closure=None):
-
-    loss = None
-    if closure is not None:
-      with torch.enable_grad():
-        loss = closure()
-
-    for group in self.param_groups:
-      for p in group["params"]:
-        if p.grad is None:
-          p.grad = torch.zeros_like(p)  # Force synchronization
-        state = self.state[p]
-        if len(state) == 0:
-          state["momentum_buffer"] = torch.zeros_like(p)
-        update = muon_update(p.grad, state["momentum_buffer"], beta=group["momentum"], ns_steps=group["ns_steps"],
-                             ns_params=group["ns_params"], nesterov=group["nesterov"])
-        p.mul_(1.0 - group["lr"] * group["weight_decay"])
-
-        p.add_(update.reshape(p.shape), alpha=-group["lr"])
-
-    return loss

pytest.ini

@@ -1,2 +0,0 @@
-[pytest]
-norecursedirs = extra

ruff.toml

@@ -35,7 +35,6 @@ lint.select = [
 line-length = 150
 
 exclude = [
-  ".git/",
   "docs/",
   "extra/",
   "tinygrad/runtime/autogen",

setup.py

@@ -18,35 +18,16 @@ testing_minimal = [
 ]
 
 setup(name='tinygrad',
-      version='0.11.0',
+      version='0.10.3',
       description='You like pytorch? You like micrograd? You love tinygrad! <3',
       author='George Hotz',
       license='MIT',
       long_description=long_description,
       long_description_content_type='text/markdown',
-      packages = [
-        'tinygrad',
-        'tinygrad.apps',
-        'tinygrad.codegen',
-        'tinygrad.codegen.opt',
-        'tinygrad.codegen.late',
-        'tinygrad.engine',
-        'tinygrad.frontend',
-        'tinygrad.nn',
-        'tinygrad.renderer',
-        'tinygrad.runtime',
-        'tinygrad.runtime.autogen',
-        'tinygrad.runtime.autogen.am',
-        'tinygrad.runtime.autogen.nv',
-        'tinygrad.runtime.graph',
-        'tinygrad.runtime.support',
-        'tinygrad.runtime.support.am',
-        'tinygrad.runtime.support.nv',
-        'tinygrad.schedule',
-        'tinygrad.shape',
-        'tinygrad.uop',
-        'tinygrad.viz',
-      ],
+      packages = ['tinygrad', 'tinygrad.runtime.autogen', 'tinygrad.runtime.autogen.am', 'tinygrad.codegen', 'tinygrad.nn',
+                  'tinygrad.renderer', 'tinygrad.engine', 'tinygrad.viz', 'tinygrad.runtime', 'tinygrad.runtime.support', 'tinygrad.schedule',
+                  'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop', 'tinygrad.codegen.opt',
+                  'tinygrad.runtime.support.nv', 'tinygrad.apps'],
       package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'assets/**/*', 'js/*']},
       classifiers=[
         "Programming Language :: Python :: 3",
@@ -64,7 +45,6 @@ setup(name='tinygrad',
             "pre-commit",
             "ruff",
             "numpy",
-            "typeguard",
         ],
         #'mlperf': ["mlperf-logging @ git+https://github.com/mlperf/[email protected]"],
         'testing_minimal': testing_minimal,

test/external/external_benchmark_kernel_launch.py

@@ -2,7 +2,7 @@ import time
 from tinygrad import Tensor, TinyJit, Device, Context
 from tinygrad.helpers import Profiling, Timing, GlobalCounters
 
-# python3 test/speed/external_test_speed_v_torch.py TestSpeed.test_add_a
+# python3 test/test_speed_v_torch.py TestSpeed.test_add_a
 
 @TinyJit
 def plus(a:Tensor, b:Tensor): return a+b

test/external/external_test_onnx_backend.py

@@ -134,6 +134,7 @@ 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')
@@ -182,8 +183,6 @@ 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,25 +100,6 @@ 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)
@@ -313,6 +294,7 @@ class TestTrainingOnnxOps(TestOnnxOps):
         outputs = ["X_out", "V_out"]
         self._validate_training("Momentum", onnx_fxn, inputs, attributes, outputs)
 
+  @unittest.expectedFailure  # TODO: regression from removing StrEnum in Domain
   def test_adam_t_greater_than_zero(self):
     from onnx.backend.test.case.node.adam import apply_adam
     for t in [1, 3, 100]:

test/external/external_test_onnx_runner.py

@@ -3,7 +3,8 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.uop.ops import Ops
 from tinygrad.device import is_dtype_supported
-from tinygrad.frontend.onnx import OnnxRunner, OnnxDataType
+from extra.onnx import OnnxDataType
+from tinygrad.frontend.onnx import OnnxRunner
 from hypothesis import given, strategies as st
 
 # copied from test_const_folding.py

test/external/fuzz_fast_idiv.py

@@ -1,9 +1,9 @@
 import random
 import z3
 from tinygrad import dtypes
-from tinygrad.uop.spec import uops_to_z3, z3_cdiv
-from tinygrad.uop.ops import UOp
-from tinygrad.uop.decompositions import fast_idiv
+from tinygrad.uop.spec import z3_renderer, z3_cdiv
+from tinygrad.uop.ops import UOp, graph_rewrite
+from tinygrad.uop.transcendental import fast_idiv
 random.seed(42)
 
 powers_of_two = [2**i for i in range(64)]
@@ -19,7 +19,8 @@ if __name__ == "__main__":
     if expr is None: continue
 
     solver = z3.Solver()
-    z3_expr, x =uops_to_z3(solver, expr, u)
+    z3_sink = graph_rewrite(expr.sink(u), z3_renderer, ctx=(solver, {}))
+    z3_expr, x = z3_sink.src[0].arg, z3_sink.src[1].arg
 
     if solver.check(z3_expr != z3_cdiv(x, d)) == z3.sat:
       assert False, f"Failed: {expr.render()} != x//{d} at x={solver.model()}\nx={u}\nd={d}\n{z3_expr=}\n{x/d=}"

test/external/fuzz_symbolic.py

@@ -1,8 +1,8 @@
 import random, operator
 import z3
 from tinygrad import Variable, dtypes
-from tinygrad.uop.ops import UOp
-from tinygrad.uop.spec import uops_to_z3
+from tinygrad.uop.ops import UOp, graph_rewrite
+from tinygrad.uop.spec import z3_renderer
 from tinygrad.helpers import DEBUG, Context
 
 seed = random.randint(0, 100)
@@ -57,7 +57,8 @@ if __name__ == "__main__":
 
     solver = z3.Solver()
     solver.set(timeout=5000)  # some expressions take very long verify, but its very unlikely they actually return sat
-    z3_expr, z3_simplified_expr, v1, v2, v3 = uops_to_z3(solver, expr, simplified_expr, u1, u2, u3)
+    z3_sink = graph_rewrite(expr.sink(simplified_expr, u1, u2, u3), z3_renderer, ctx=(solver, {}))
+    z3_expr, z3_simplified_expr = z3_sink.src[0].arg, z3_sink.src[1].arg
     check = solver.check(z3_simplified_expr != z3_expr)
     if check == z3.unknown and DEBUG>=1:
       skipped += 1
@@ -68,6 +69,7 @@ if __name__ == "__main__":
             f"expr = {expr.render(simplify=False)}\n")
     elif check == z3.sat:
       m = solver.model()
+      v1, v2, v3 = z3_sink.src[2].arg, z3_sink.src[3].arg, z3_sink.src[4].arg
       n1, n2, n3 = m[v1], m[v2], m[v3]
       u1_val, u2_val, u3_val = u1.const_like(n1.as_long()), u2.const_like(n2.as_long()), u3.const_like(n3.as_long())
       with Context(CORRECT_DIVMOD_FOLDING=1):

test/external/process_replay/process_replay.py

@@ -99,6 +99,7 @@ def diff(offset:int, fxns:dict[str, Callable[..., tuple|None]]) -> None:
     except Exception as e:
       changed += 1
       warnings.warn(f"{name=} {loc=} {e=}", ProcessReplayWarning)
+  conn.commit()
   cur.close()
 
 # *** generic runner to map rows of a table to a function in parallel
@@ -110,11 +111,12 @@ def _pmap(fxns:dict[str, Callable]) -> None:
   except sqlite3.OperationalError:
     raise RuntimeError(f"{TABLE_NAME} isn't accessible in master, did DB_VERSION change?")
   finally:
+    conn.commit()
     cur.close()
 
   with multiprocessing.get_context("spawn").Pool(multiprocessing.cpu_count()) as pool:
-    bar = tqdm(total=row_count)
-    for _ in pool.imap_unordered(functools.partial(diff, fxns=fxns), range(0, row_count, PAGE_SIZE)): bar.update(PAGE_SIZE)
+    inputs = list(range(0, row_count, PAGE_SIZE))
+    list(tqdm(pool.imap_unordered(functools.partial(diff, fxns=fxns), inputs), total=len(inputs)))
     pool.close()
     pool.join()
     pool.terminate()

test/mockgpu/amd/amddriver.py

@@ -87,19 +87,16 @@ class AMDDriver(VirtDriver):
         functools.partial(TextFileDesc, text=gpu_props.format(drm_render_minor=gpu_id))),
       VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0',
                functools.partial(DirFileDesc, child_names=[str(am.GC_HWID), str(am.SDMA0_HWID), str(am.NBIF_HWID)])),
-      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}', functools.partial(DirFileDesc, child_names=['0'])),
       VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}/0/major', functools.partial(TextFileDesc, text='11')),
       VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}/0/minor', functools.partial(TextFileDesc, text='0')),
       VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}/0/revision', functools.partial(TextFileDesc, text='0')),
       VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}/0/base_addr',
                functools.partial(TextFileDesc, text='0x00001260\n0x0000A000\n0x0001C000\n0x02402C00')),
-      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}', functools.partial(DirFileDesc, child_names=['0'])),
       VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}/0/major', functools.partial(TextFileDesc, text='6')),
       VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}/0/minor', functools.partial(TextFileDesc, text='0')),
       VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}/0/revision', functools.partial(TextFileDesc, text='0')),
       VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}/0/base_addr',
                functools.partial(TextFileDesc, text='0x00001260\n0x0000A000\n0x0001C000\n0x02402C00')),
-      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.NBIF_HWID}', functools.partial(DirFileDesc, child_names=['0'])),
       VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.NBIF_HWID}/0/major', functools.partial(TextFileDesc, text='4')),
       VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.NBIF_HWID}/0/minor', functools.partial(TextFileDesc, text='3')),
       VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.NBIF_HWID}/0/revision', functools.partial(TextFileDesc, text='0')),

test/mockgpu/amd/amdgpu.py

@@ -1,7 +1,7 @@
 import ctypes, time
 from test.mockgpu.gpu import VirtGPU
 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
+import tinygrad.runtime.autogen.amd_gpu as amd_gpu, tinygrad.runtime.autogen.am.pm4_nv as pm4, tinygrad.runtime.autogen.am.soc21 as soc21
 
 SDMA_MAX_COPY_SIZE = 0x400000
 
@@ -14,9 +14,6 @@ regSQ_THREAD_TRACE_BUF0_SIZE = 0x39e9 + amd_gpu.GC_BASE__INST0_SEG1
 regSQ_THREAD_TRACE_WPTR = 0x39ef + amd_gpu.GC_BASE__INST0_SEG1
 regSQ_THREAD_TRACE_STATUS = 0x39f4 + amd_gpu.GC_BASE__INST0_SEG1
 
-class SQTT_EVENTS:
-  THREAD_TRACE_FINISH = 0x00000037
-
 CACHE_FLUSH_AND_INV_TS_EVENT = 0x14
 
 WAIT_REG_MEM_FUNCTION_ALWAYS = 0
@@ -211,7 +208,7 @@ class PM4Executor(AMDQueue):
     assert n == 0
     event_dw = self._next_dword()
     match (event_dw & 0xFF): # event type
-      case SQTT_EVENTS.THREAD_TRACE_FINISH:
+      case soc21.THREAD_TRACE_FINISH:
         old_idx = self.gpu.regs.grbm_index
         for se in range(self.gpu.regs.n_se):
           self.gpu.regs.grbm_index = 0b011 << 29 | se << 16 # select se, broadcast sa and instance

test/models/test_onnx.py

@@ -9,15 +9,7 @@ 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
-
-try:
-  from extra.onnx_helpers import validate
-  from extra.huggingface_onnx.huggingface_manager import DOWNLOADS_DIR, snapshot_download_with_retry
-  HUGGINGFACE_AVAILABLE = True
-except ModuleNotFoundError:
-  HUGGINGFACE_AVAILABLE = False
+from tinygrad.helpers import CI, fetch, temp
 
 def run_onnx_torch(onnx_model, inputs):
   import torch
@@ -32,7 +24,6 @@ OPENPILOT_MODEL = "https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/mod
 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)
@@ -146,36 +137,5 @@ class TestOnnxModel(unittest.TestCase):
     print(cls, _LABELS[cls])
     assert "car" in _LABELS[cls] or _LABELS[cls] == "convertible"
 
-@unittest.skipUnless(HUGGINGFACE_AVAILABLE and Device.DEFAULT == "METAL", "only run on METAL")
-class TestHuggingFaceOnnxModels(unittest.TestCase):
-  @classmethod
-  def setUpClass(cls):
-    cls._ctx = Context(MAX_BUFFER_SIZE=0)
-    cls._ctx.__enter__()
-
-  @classmethod
-  def tearDownClass(cls):
-    cls._ctx.__exit__()
-
-  def _validate(self, repo_id, model_file, custom_inputs, rtol=1e-4, atol=1e-4):
-    onnx_model_path = snapshot_download_with_retry(
-      repo_id=repo_id,
-      allow_patterns=["*.onnx", "*.onnx_data"],
-      cache_dir=str(DOWNLOADS_DIR)
-    )
-    onnx_model_path = onnx_model_path / model_file
-    file_size = onnx_model_path.stat().st_size
-    print(f"Validating model: {repo_id}/{model_file} ({file_size/1e6:.2f}M)")
-    validate(onnx_model_path, custom_inputs, rtol=rtol, atol=atol)
-
-  def test_xlm_roberta_large(self):
-    repo_id = "FacebookAI/xlm-roberta-large"
-    model_file = "onnx/model.onnx"
-    custom_inputs = {
-      "input_ids": np.random.randint(0, 250002, (1, 11), dtype=np.int64),
-      "attention_mask": np.ones((1, 11), dtype=np.int64),
-    }
-    self._validate(repo_id, model_file, custom_inputs)
-
 if __name__ == "__main__":
   unittest.main()

test/models/test_whisper.py

@@ -16,7 +16,7 @@ TRANSCRIPTION_2 = "a slightly longer audio file so that we can test batch transc
 TEST_FILE_3_URL = 'https://homepage.ntu.edu.tw/~karchung/miniconversations/mc45.mp3'
 TRANSCRIPTION_3 = "Just lie back and relax. Is the level of pressure about right? Yes, it's fine, and I'd like conditioner please. Sure. I'm going to start the second lathering now. Would you like some Q-tips? How'd you like it cut? I'd like my bangs and the back trimmed, and I'd like the rest thinned out a bit and layered. Where would you like the part? On the left, right about here. Here, have a look. What do you think? It's fine. Here's a thousand anti-dollars. It's 30-ant extra for the rants. Here's your change and receipt. Thank you, and please come again. So how do you like it? It could have been worse, but you'll notice that I didn't ask her for her card. Hmm, yeah. Maybe you can try that place over there next time."   # noqa: E501
 
-@unittest.skipIf(Device.DEFAULT in ["CPU", "LLVM"], "slow")
+@unittest.skipIf(CI and Device.DEFAULT in ["CPU"], "slow")
 @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need float16 support")
 class TestWhisper(unittest.TestCase):
   @classmethod

test/test_compile_failures.py

@@ -1,10 +1,7 @@
-import unittest, io
-from contextlib import redirect_stdout
+import unittest
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.helpers import OSX
 from tinygrad.engine.realize import lower_schedule
 from tinygrad.device import is_dtype_supported
-from tinygrad.engine.realize import get_program
 
 class TestCompileFailures(unittest.TestCase):
   def compile(self, out:Tensor):
@@ -17,17 +14,5 @@ class TestCompileFailures(unittest.TestCase):
   def test_add_max_uchar(self):
     self.compile((Tensor.empty(1024, dtype='uint8') + Tensor.empty(1024, dtype='uint8')).max())
 
-class TestDisassembly(unittest.TestCase):
-  # TODO: fails on llvm. llvm.LLVMGetHostCPUName() returns "generic"
-  @unittest.skipUnless(Device.DEFAULT in ("CPU",) and OSX, "m series cpus support fp16 arithmetic")
-  def test_float16_alu(self):
-    c = Tensor([1], dtype=dtypes.float16) + Tensor([1], dtype=dtypes.float16)
-    s = c.schedule()[-1]
-    p = get_program(s.ast, Device[Device.DEFAULT].renderer)
-    lib = Device[Device.DEFAULT].compiler.compile(p.src)
-    out = io.StringIO()
-    with redirect_stdout(out): Device[Device.DEFAULT].compiler.disassemble(lib)
-    assert "fcvt" not in out.getvalue()
-
 if __name__ == '__main__':
   unittest.main()

test/test_const_folding.py

@@ -1,10 +1,11 @@
 import unittest, itertools, math
+from typing import Any
 from tinygrad import Tensor, Device, dtypes
-from tinygrad.dtype import DType, ConstType
+from tinygrad.dtype import DType
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.codegen import full_rewrite_to_sink
-from tinygrad.device import is_dtype_supported
 import numpy as np
+from tinygrad.device import is_dtype_supported
 from test.helpers import not_support_multi_device
 
 def _check_ast_count(desired_count:int, t:Tensor):
@@ -24,7 +25,7 @@ class TestUnaryOpsConstFolding(unittest.TestCase):
     _check_ast_count(0, Tensor.ones(4).cast(dtypes.int16))
     _check_ast_count(0, Tensor.full(4, fill_value=-1).cast(dtypes.uint16))
 
-  @unittest.expectedFailure  # no two level fold
+  @unittest.expectedFailure  # no two level fold at lazybuffer
   def test_neg_folding(self):
     _check_ast_count(0, Tensor([1, 2, 3]).mul(-1).neg())
     _check_ast_count(0, Tensor([1, 2, 3]).neg().mul(-1))
@@ -103,7 +104,7 @@ class TestBinaryOpsConstFolding(unittest.TestCase):
 
 class TestBitcastConstFolding(unittest.TestCase):
   def test_scalar_bitcast(self):
-    def t(cases: dict[DType, ConstType]):
+    def t(cases: dict[DType, Any]):
       for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()):
         if not math.isnan(from_v):
           r = full_rewrite_to_sink(UOp.const(from_dt, from_v).bitcast(to_dt).sink()).src[0]
@@ -142,12 +143,13 @@ class TestIndexingConstFolding(unittest.TestCase):
     _check_ast_count(1, t[:,:,Tensor(1)+2,:])
     _check_ast_count(1, t[:,:,Tensor(1),Tensor(0)])
 
+  @unittest.expectedFailure
   def test_const_tensor_index(self):
-    # TODO: these can be 0, implement const tensor folded indexing
+    # TODO: implement const tensor folded indexing
     t = Tensor.arange(16).float().reshape(1,1,4,4).realize()
-    _check_ast_count(1, t[:,:,Tensor.ones(2,1,dtype=dtypes.int),:])
-    _check_ast_count(1, t[:,:,Tensor.ones(1,2,dtype=dtypes.int)+2,:])
-    _check_ast_count(1, t[:,:,Tensor.ones(1,1,dtype=dtypes.int),Tensor.zeros(2,1,2,dtype=dtypes.int)])
+    _check_ast_count(0, t[:,:,Tensor.ones(2,1),:])
+    _check_ast_count(0, t[:,:,Tensor.ones(1,2)+2,:])
+    _check_ast_count(0, t[:,:,Tensor.ones(1,1),Tensor.zeros(2,1,2)])
 
 class TestMovedConstFolding(unittest.TestCase):
   def test_add_shrunk_zero(self):
@@ -164,6 +166,7 @@ class TestMovedConstFolding(unittest.TestCase):
     _check_ast_count(1, Tensor([1.0, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),)))
 
   def test_cast_padded(self):
+    # NOTE: this is folded due to CAST_BEFORE_VIEW
     if is_dtype_supported(dtypes.int16):
       _check_ast_count(0, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16))
       np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16).numpy(), [0, 1, 1, 1, 1, 0])

test/test_copy_speed.py

@@ -3,7 +3,7 @@ from tinygrad import Tensor, Device, TinyJit
 from tinygrad.helpers import Timing, CI, OSX
 import multiprocessing.shared_memory as shared_memory
 
-N = 256
+N = 256 if CI else 4096
 class TestCopySpeed(unittest.TestCase):
   @classmethod
   def setUpClass(cls): Device[Device.DEFAULT].synchronize()

test/test_define_reg.py

@@ -0,0 +1,32 @@
+import unittest
+from tinygrad import dtypes, Device, Tensor, Context
+from tinygrad.dtype import AddrSpace
+from tinygrad.helpers import getenv
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.engine.realize import get_program, ExecItem, CompiledRunner
+
+class TestDefineReg(unittest.TestCase):
+  def test_simple(self, at=AxisType.UPCAST):
+    N = 16
+    bout = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0).view(ShapeTracker.from_shape((N,N)))
+    a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1).view(ShapeTracker.from_shape((N,N)))
+    a_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(N, AddrSpace.REG), arg=0).view(ShapeTracker.from_shape((N,N), (0,1)))
+
+    out = a_col.load(a_col.store(a.load()))
+    sink = bout.store(out).sink(arg=KernelInfo(name="regcopy", axis_types=(AxisType.LOOP, at)))
+    prg = get_program(sink, Device.default.renderer)
+
+    with Context(DEBUG=0):
+      a = Tensor.randn(N, N).realize()
+      b = Tensor.empty(N, N).realize()
+    hrunner = CompiledRunner(prg)
+    ExecItem(hrunner, [b.uop.buffer, a.uop.buffer]).run(wait=True)
+    with Context(DEBUG=0):
+      self.assertEqual((b-a).mean().item(), 0.0)
+
+  @unittest.skipIf(getenv("PTX"), "ptx needs regs to be unrolled")
+  def test_simple_loop(self): self.test_simple(AxisType.LOOP)
+
+if __name__ == '__main__':
+  unittest.main()

test/test_disassembly.py

@@ -0,0 +1,21 @@
+import unittest, io
+from tinygrad import Tensor, dtypes
+from contextlib import redirect_stdout
+from tinygrad.device import Device
+from tinygrad.helpers import OSX
+from tinygrad.engine.realize import get_program
+
+class TestDisassembly(unittest.TestCase):
+  # TODO: fails on llvm. llvm.LLVMGetHostCPUName() returns "generic"
+  @unittest.skipUnless(Device.DEFAULT in ("CPU",) and OSX, "m series cpus support fp16 arithmetic")
+  def test_float16_alu(self):
+    c = Tensor([1], dtype=dtypes.float16) + Tensor([1], dtype=dtypes.float16)
+    s = c.schedule()[-1]
+    p = get_program(s.ast, Device[Device.DEFAULT].renderer)
+    lib = Device[Device.DEFAULT].compiler.compile(p.src)
+    out = io.StringIO()
+    with redirect_stdout(out): Device[Device.DEFAULT].compiler.disassemble(lib)
+    assert "fcvt" not in out.getvalue()
+
+if __name__ == "__main__":
+  unittest.main()

test/test_dtype.py

@@ -414,11 +414,11 @@ 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
-    # TODO: wrong output with GPU=1 on mac
+    # TODO: wrong output with GPU=1 / PYTHON=1 on mac
     data = [60000.0, 70000.0, 80000.0]
     np.testing.assert_allclose(Tensor(data).cast("bfloat16").numpy(), torch.tensor(data).type(torch.bfloat16).float().numpy())
 

test/test_dtype_alu.py

@@ -1,13 +1,16 @@
-import unittest, operator, math
+import unittest
+
 from tinygrad import Tensor, dtypes, Device
+import operator
+import numpy as np
+from hypothesis import given, strategies as strat, settings, HealthCheck
 from tinygrad.dtype import DType
 from tinygrad.helpers import CI, getenv
+from tinygrad.engine.realize import run_schedule
+from tinygrad.uop.ops import GroupOp
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
-import numpy as np
-import pytest
-from hypothesis import given, strategies as strat, settings, HealthCheck
-
+import pytest, math
 pytestmark = pytest.mark.filterwarnings("ignore")
 
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
@@ -60,21 +63,25 @@ def universal_test(a, b, dtype, op):
   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())
-  if dtype in dtypes.floats:
-    atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2)}.get(dtype, (1e-10, 1e-7))
-    np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
+  if dtype == dtypes.bfloat16: np.testing.assert_allclose(tensor_value, numpy_value, atol=1e-3, rtol=1e-2)
+  elif dtype in dtypes_float: np.testing.assert_allclose(tensor_value, numpy_value, atol=1e-10)
   else: np.testing.assert_equal(tensor_value, numpy_value)
 
 def universal_test_unary(a, dtype, op):
   if not isinstance(op, tuple): op = (op, op)
   ta = Tensor([a], dtype=dtype)
   out: Tensor = op[0](ta)
+  sched = out.schedule()
+  ast = sched[-1].ast
+  run_schedule(sched)
   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))
-    np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
+  if dtype in (dtypes.float16, dtypes.bfloat16): np.testing.assert_allclose(tensor_value, numpy_value, atol=1e-3, rtol=1e-2)
+  elif dtype in dtypes_float: np.testing.assert_allclose(tensor_value, numpy_value, atol=1e-6, rtol=1e-5)
   else: np.testing.assert_equal(tensor_value, numpy_value)
+  if op[0] != Tensor.reciprocal: # reciprocal is not supported in most backends
+    op = [x for x in ast.toposort() if x.op in GroupOp.Unary][0]
+    assert op.dtype == dtype
 
 def universal_test_cast(a, in_dtype, dtype):
   tensor_value = Tensor([a], dtype=in_dtype).cast(dtype)
@@ -92,44 +99,45 @@ def universal_test_midcast(a, b, c, op1, op2, d1:DType, d2:DType):
   np.testing.assert_allclose(tensor_value, numpy_value, rtol=1e-6 if getenv("PTX") else 1e-7)
 
 class TestDTypeALU(unittest.TestCase):
-  @unittest.skipUnless(is_dtype_supported(dtypes.float64), f"no float64 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.float64, Device.DEFAULT), f"no float64 on {Device.DEFAULT}")
   @given(ht.float64, ht.float64, strat.sampled_from(binary_operations))
   def test_float64(self, a, b, op): universal_test(a, b, dtypes.float64, op)
 
   @given(ht.float32, ht.float32, strat.sampled_from(binary_operations))
   def test_float32(self, a, b, op): universal_test(a, b, dtypes.float32, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.float16), f"no float16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.float16, Device.DEFAULT), f"no float16 on {Device.DEFAULT}")
   @given(ht.float16, ht.float16, strat.sampled_from(binary_operations))
   def test_float16(self, a, b, op): universal_test(a, b, dtypes.float16, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16, Device.DEFAULT), 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)
 
   @given(ht.float32, strat.sampled_from(unary_operations))
   def test_float32_unary(self, a, op): universal_test_unary(a, dtypes.float32, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.float16), f"no float16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.float16, Device.DEFAULT), f"no float16 on {Device.DEFAULT}")
   @given(ht.float16, strat.sampled_from(unary_operations))
   def test_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16, Device.DEFAULT), f"no bfloat16 on {Device.DEFAULT}")
   @given(ht.bfloat16, strat.sampled_from(unary_operations))
+  @unittest.skipIf(Device.DEFAULT in ["AMD"], "broken on AMD?")
   def test_bfloat16_unary(self, a, op): universal_test_unary(a, 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)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.uint16), f"no uint16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.uint16, Device.DEFAULT), f"no uint16 on {Device.DEFAULT}")
   @given(ht.uint16, ht.uint16, strat.sampled_from(integer_binary_operations))
   def test_uint16(self, a, b, op): universal_test(a, b, dtypes.uint16, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.uint32), f"no uint32 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.uint32, Device.DEFAULT), f"no uint32 on {Device.DEFAULT}")
   @given(ht.uint32, ht.uint32, strat.sampled_from(integer_binary_operations))
   def test_uint32(self, a, b, op): universal_test(a, b, dtypes.uint32, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.uint64), f"no uint64 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.uint64, Device.DEFAULT), f"no uint64 on {Device.DEFAULT}")
   @given(ht.uint64, ht.uint64, strat.sampled_from(integer_binary_operations))
   def test_uint64(self, a, b, op): universal_test(a, b, dtypes.uint64, op)
 
@@ -142,7 +150,7 @@ class TestDTypeALU(unittest.TestCase):
   @given(ht.int32, ht.int32, strat.sampled_from(integer_binary_operations))
   def test_int32(self, a, b, op): universal_test(a, b, dtypes.int32, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.int64), f"no int64 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.int64, Device.DEFAULT), f"no int64 on {Device.DEFAULT}")
   @given(ht.int64, ht.int64, strat.sampled_from(integer_binary_operations))
   def test_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op)
 
@@ -172,7 +180,7 @@ class TestDTypeALU(unittest.TestCase):
   @settings(suppress_health_check=[HealthCheck.filter_too_much])
   @given(strat.data(), strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
   def test_float_cast_to_unsigned(self, a, float_dtype, unsigned_dtype):
-    if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
+    if not is_dtype_supported(float_dtype, Device.DEFAULT): float_dtype = dtypes.float32
     float_strat = {dtypes.float16: ht.float16, dtypes.float32: ht.float32, dtypes.float64: ht.float64}[float_dtype]
     float_strat = float_strat.filter(lambda x: 0 < x < dtypes.max(unsigned_dtype))
     universal_test_cast(a.draw(float_strat), float_dtype, unsigned_dtype)
@@ -180,7 +188,7 @@ class TestDTypeALU(unittest.TestCase):
   @settings(suppress_health_check=[HealthCheck.filter_too_much])
   @given(strat.data(), strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
   def test_float_cast_to_unsigned_overflow(self, a, float_dtype, unsigned_dtype):
-    if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
+    if not is_dtype_supported(float_dtype, Device.DEFAULT): float_dtype = dtypes.float32
     float_strat = {dtypes.float16: ht.float16, dtypes.float32: ht.float32, dtypes.float64: ht.float64}[float_dtype]
     overflow_strat = float_strat.filter(lambda x: x > dtypes.max(unsigned_dtype) and x <= dtypes.max(dtypes.int32))
     universal_test_cast(a.draw(overflow_strat), float_dtype, unsigned_dtype)
@@ -188,7 +196,7 @@ class TestDTypeALU(unittest.TestCase):
   @settings(suppress_health_check=[HealthCheck.filter_too_much])
   @given(strat.data(), strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
   def test_float_cast_to_unsigned_underflow(self, a, float_dtype, unsigned_dtype):
-    if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
+    if not is_dtype_supported(float_dtype, Device.DEFAULT): float_dtype = dtypes.float32
     float_strat = {dtypes.float16: ht.float16, dtypes.float32: ht.float32, dtypes.float64: ht.float64}[float_dtype]
     underflow_strat = float_strat.filter(lambda x: x < 0 and x >= dtypes.min(dtypes.int32))
     universal_test_cast(a.draw(underflow_strat), float_dtype, unsigned_dtype)

test/test_linalg.py

@@ -62,15 +62,5 @@ class TestLinAlg(unittest.TestCase):
       orthogonality_helper(Q)
       reconstruction_helper([Q,R],a)
 
-  def test_newton_schulz(self):
-    coefficients = [(2, -1.5, 0.5), (2.0, -1.4, 0.2, 0.2)]#these params map to the sign function
-    sizes = [(2,2), (3,2), (2,3), (2,2,2)]
-    for coefs in coefficients:
-      for size in sizes:
-        a = Tensor.randn(size)
-        b = Tensor.newton_schulz(a, steps=20, params=coefs, eps=0.0)
-        # ns(A) = U @ Vt -> (U @ Vt) @ (U @ Vt)t = I
-        orthogonality_helper(b if size[-1] > size[-2] else b.transpose(-2, -1), tolerance=1e-1)
-
 if __name__ == "__main__":
   unittest.main()

test/test_linearizer.py

@@ -117,7 +117,6 @@ class TestLinearizer(unittest.TestCase):
         if skip and i in skip: continue
         assert ranges[i-1] != u, f"multireduce nested the ranges! {ranges[i-1], {u}}"
 
-  @unittest.skip("broken. should not depends on push_views and implementation details of getitem")
   @unittest.skipIf(CI and Device.DEFAULT in {"PTX", "AMD", "NV"}, "very slow")
   def test_indexing_multireduce(self):
     dataset = Tensor.rand(16384, 256).realize()

test/test_memory_planner.py

@@ -120,19 +120,5 @@ class TestMemoryPlanner(unittest.TestCase):
     ]
     check_assign(bs)
 
-  def test_very_small_buffers(self):
-    bs = [
-      [b(0, pin=True), b(1, size=32)],
-      [b(3, size=4), b(4, size=6)],
-    ]
-    check_assign(bs)
-
-  def test_very_big_buffers(self):
-    bs = [
-      [b(0, pin=True), b(1, size=34359738368000)],
-      [b(3, size=1 << 128), b(4, size=1 << 64)],
-    ]
-    check_assign(bs)
-
 if __name__ == "__main__":
   unittest.main()

test/test_multitensor.py

@@ -373,11 +373,11 @@ class TestMultiTensor(unittest.TestCase):
     np.testing.assert_allclose(y.numpy(), y_shard.numpy(), atol=1e-6, rtol=1e-6)
 
   # NOTE: this is failing on LLVM CI, no idea why. Works locally.
-  @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "LLVM", "CPU", "AMD"), "slow, and flaky on LLVM/CPU")
+  @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "LLVM", "CPU"), "slow, and flaky on LLVM/CPU")
   def test_data_parallel_resnet(self):
     from extra.models.resnet import ResNet18
 
-    fake_image = Tensor.rand((2, 3, 224//16, 224//16))
+    fake_image = Tensor.rand((2, 3, 224//8, 224//8))
     fake_image_sharded = fake_image.shard(devices_2, axis=0)
     m = ResNet18()
     m.load_from_pretrained()
@@ -409,10 +409,10 @@ class TestMultiTensor(unittest.TestCase):
     # sometimes there is zeros in these grads... why?
     np.testing.assert_allclose(grad, shard_grad, atol=1e-5, rtol=1e-5)
 
-  @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "LLVM", "CPU", "AMD"), "slow, and flaky on LLVM/CPU")
+  @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "LLVM", "CPU"), "slow, and flaky on LLVM/CPU")
   def test_data_parallel_resnet_train_step(self):
     from extra.models.resnet import ResNet18
-    fake_image = Tensor.rand((2, 3, 224//16, 224//16))
+    fake_image = Tensor.rand((2, 3, 224//8, 224//8))
     labels = Tensor.randint(2, low=0, high=1000)
     m = ResNet18()
     self._test_model_train_step(m, fake_image, labels)
@@ -1128,7 +1128,6 @@ class TestMultiRamUsage(unittest.TestCase):
     self.assertUsed(self.N*self.N*4) # sharding should not increase total ram usage
   def test_zeros_shard_self(self): self.test_zeros_shard((d0, d1))
 
-  @unittest.skip("flaky")
   def test_zeros_contiguous_shard(self):
     _ = Tensor.zeros(self.N, self.N).contiguous().shard(devices_2, axis=0).contiguous().realize()
     self.assertUsed(self.N*self.N*4) # sharding should not increase total ram usage

test/test_net_speed.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python
+import time
+import unittest
+import torch
+from tinygrad import Tensor, Device
+from tinygrad.helpers import Profiling, CI
+
+@unittest.skipIf(CI and Device.DEFAULT in {"CUDA", "NV"}, "slow")
+class TestConvSpeed(unittest.TestCase):
+
+  def test_mnist(self):
+    # https://keras.io/examples/vision/mnist_convnet/
+    conv = 3
+    inter_chan, out_chan = 32, 64
+
+    # ****** torch baseline *******
+
+    torch.backends.mkldnn.enabled = False
+
+    conv = 3
+    inter_chan, out_chan = 32, 64
+    c1 = torch.randn(inter_chan,1,conv,conv, requires_grad=True)
+    c2 = torch.randn(out_chan,inter_chan,conv,conv, requires_grad=True)
+    l1 = torch.randn(out_chan*5*5, 10, requires_grad=True)
+
+    c2d = torch.nn.functional.conv2d
+    mp = torch.nn.MaxPool2d((2,2))
+    lsm = torch.nn.LogSoftmax(dim=1)
+
+    cnt = 5
+    fpt, bpt = 0.0, 0.0
+    for i in range(cnt):
+      et0 = time.time()
+      x = torch.randn(128, 1, 28, 28, requires_grad=True)
+      x = mp(c2d(x,c1).relu())
+      x = mp(c2d(x,c2).relu())
+      x = x.reshape(x.shape[0], -1)
+      out = lsm(x.matmul(l1))
+      out = out.mean()
+      et1 = time.time()
+      out.backward()
+      et2 = time.time()
+      fpt += (et1-et0)
+      bpt += (et2-et1)
+
+    fpt_baseline = (fpt*1000/cnt)
+    bpt_baseline = (bpt*1000/cnt)
+    print("torch forward pass:  %.3f ms" % fpt_baseline)
+    print("torch backward pass: %.3f ms" % bpt_baseline)
+
+    # ****** tinygrad compare *******
+
+    c1 = Tensor(c1.detach().numpy(), requires_grad=True)
+    c2 = Tensor(c2.detach().numpy(), requires_grad=True)
+    l1 = Tensor(l1.detach().numpy(), requires_grad=True)
+
+    cnt = 5
+    fpt, bpt = 0.0, 0.0
+    for i in range(1+cnt):
+      et0 = time.time()
+      x = Tensor.randn(128, 1, 28, 28)
+      x = x.conv2d(c1).relu().avg_pool2d()
+      x = x.conv2d(c2).relu().max_pool2d()
+      x = x.reshape(shape=(x.shape[0], -1))
+      out = x.dot(l1).log_softmax()
+      out = out.mean()
+      out.backward()  # NOTE: we have to now compute this here, but it doesn't realize
+      out.realize()
+      et1 = time.time()
+      [x.grad.realize() for x in [c1, c2, l1]]
+      et2 = time.time()
+      if i == 0:
+        pr = Profiling(sort='time', frac=0.2)
+        pr.__enter__()
+      else:
+        fpt += (et1-et0)
+        bpt += (et2-et1)
+
+    pr.__exit__()
+    fpt = (fpt*1000/cnt)
+    bpt = (bpt*1000/cnt)
+    print("forward pass:  %.3f ms, %.2fx off baseline %.3f ms" % (fpt, fpt/fpt_baseline, fpt_baseline))
+    print("backward pass: %.3f ms, %.2fx off baseline %.3f ms" % (bpt, bpt/bpt_baseline, bpt_baseline))
+
+
+if __name__ == '__main__':
+  unittest.main()

test/test_nn.py

@@ -210,27 +210,6 @@ class TestNN(unittest.TestCase):
       np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
       np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
 
-  def test_layernorm_forward(self):
-    N, C, H, W = 20, 5, 10, 10
-
-    # create in torch
-    torch_layer = torch.nn.LayerNorm([H, W]).eval()
-
-    # create in tinygrad
-    layer = LayerNorm([H, W])
-    layer.weight = Tensor(torch_layer.weight.detach().numpy(), requires_grad=True)
-    layer.bias = Tensor(torch_layer.bias.detach().numpy(), requires_grad=True)
-
-    x = Tensor.empty(N, C, H, W, requires_grad=True)
-    z = layer(x)
-    z.realize()
-
-    torch_x = torch.tensor(x.numpy(), requires_grad=True)
-    torch_z = torch_layer(torch_x)
-    torch_z.sum().backward()
-
-    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
-
   def test_layernorm(self):
     N, C, H, W = 20, 5, 10, 10
 

test/test_ops.py

@@ -393,7 +393,7 @@ class TestOps(unittest.TestCase):
   def test_trunc(self):
     helper_test_op([()], lambda x: x.trunc(), forward_only=True)
     helper_test_op([(45,35)], lambda x: x.trunc(), forward_only=True)
-    helper_test_op(None, lambda x: x.trunc(), vals=[[1.499, 1.5, 1.501, 1.0, 2.1, 0.0, -5.0, -2.499, -2.5, -2.501, 1e12, -1e12]], forward_only=True)
+    helper_test_op(None, lambda x: x.trunc(), vals=[[1.499, 1.5, 1.501, 1.0, 2.1, 0.0, -5.0, -2.499, -2.5, -2.501]], forward_only=True)
   def test_floor(self):
     helper_test_op([()], lambda x: x.floor(), forward_only=True)
     helper_test_op([(45,35)], lambda x: x.floor(), forward_only=True)
@@ -928,12 +928,6 @@ class TestOps(unittest.TestCase):
       for j in [-1., 0., 1.]:
         helper_test_op(None, torch.copysign, Tensor.copysign, vals=[[i], [j]])
 
-  def test_logaddexp(self):
-    helper_test_op([(45,65), (45,65)], torch.logaddexp, Tensor.logaddexp)
-    helper_test_op(None, torch.logaddexp, Tensor.logaddexp, vals=[[-1.], [-1.0, 2, 3]])
-    helper_test_op(None, torch.logaddexp, Tensor.logaddexp, vals=[[-100.0, -200, -300], [-1.0, 2, 3]])
-    helper_test_op(None, torch.logaddexp, Tensor.logaddexp, vals=[[1.0, 2000, 30000], [-1.0, 2, 3]])
-
   def test_softsign(self):
     helper_test_op([(45,65)], torch.nn.functional.softsign, Tensor.softsign)
     helper_test_op([()], torch.nn.functional.softsign, Tensor.softsign)
@@ -971,6 +965,8 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6)
     helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=3), lambda t: Tensor.softplus(t, beta=3), grad_atol=1e-6)
     helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=1/3), lambda t: Tensor.softplus(t, beta=1/3), grad_atol=1e-6)
+    helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=3, threshold=0.5),
+                              lambda t: Tensor.softplus(t, beta=3, threshold=0.5), grad_atol=1e-6)
     helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=300, high=400)
     helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=-400, high=-300)
     helper_test_op([()], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6)
@@ -2465,20 +2461,6 @@ class TestOps(unittest.TestCase):
       lambda x: Tensor.max_unpool2d(*Tensor.max_pool2d(x, kernel_size=(2,2), return_indices=True),
                                     kernel_size=(2,2), output_size=(99,99,7,6)), forward_only=True)
 
-  def test_max_unpool2d_inf(self):
-    data = [[[[math.inf, -math.inf, math.nan], [1.0, 2.0, 3.0]]]]
-    ksz = (2,2)
-    helper_test_op((),
-      lambda: torch.nn.functional.max_unpool2d(
-        *torch.nn.functional.max_pool2d(torch.tensor(data), kernel_size=ksz, return_indices=True),
-        kernel_size=ksz
-      ),
-      lambda: Tensor.max_unpool2d(
-        *Tensor.max_pool2d(Tensor(data), kernel_size=ksz, return_indices=True),
-        kernel_size=ksz
-      ),
-      forward_only=True)
-
   def test_avg_pool2d(self):
     shape = (32,2,111,28)
     for ksz in [(2,2), (3,3), (3,2), (5,5), (5,1)]:
@@ -2712,10 +2694,6 @@ class TestOps(unittest.TestCase):
     i, j, k, o, p = [Tensor(tor.detach().cpu().numpy().astype(np.int32), requires_grad=False) for tor in [a,b,c,d,e]]
     return a,b,c,d,e,i,j,k,o,p
 
-  def test_fancy_indexing_inf(self):
-    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])])
-
   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
@@ -2826,7 +2804,11 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x: x.gather(dim=0, index=torch.tensor([2, 1, 0, 1, 2], requires_grad=False)),
                          lambda x: x.gather(dim=0, index=Tensor([2, 1, 0, 1, 2])),
                          vals=[[1., 2., 3.]])
-    # gather with inf values
+
+  @unittest.expectedFailure
+  @unittest.skipIf(torch._C._get_privateuse1_backend_name() == "tiny", 'results in a success instead of a failure')
+  def test_gather_failure(self):
+    # gather with inf values do not work, other values results in nan
     helper_test_op(None, lambda x: x.gather(dim=0, index=torch.tensor([2, 1, 0, 1, 2], requires_grad=False)),
                          lambda x: x.gather(dim=0, index=Tensor([2, 1, 0, 1, 2])),
                          vals=[[-float("inf"), 2., 3.]])

test/test_optim.py

@@ -2,10 +2,9 @@ import numpy as np
 import torch
 import unittest
 from tinygrad import Tensor, Device, dtypes
-from tinygrad.nn.optim import Adam, SGD, AdamW, Muon
+from tinygrad.nn.optim import Adam, SGD, AdamW
 from tinygrad.helpers import CI
 from tinygrad.device import is_dtype_supported
-from extra.torch_muon import SingleDeviceMuon as TorchMuon
 
 np.random.seed(1337)
 x_init = np.random.randn(1,4).astype(np.float32)
@@ -58,12 +57,9 @@ class TestOptim(unittest.TestCase):
   def _test_sgd(self, steps, opts, atol, rtol): self._test_optim(SGD, torch.optim.SGD, steps, opts, atol, rtol)
   def _test_adam(self, steps, opts, atol, rtol): self._test_optim(Adam, torch.optim.Adam, steps, opts, atol, rtol)
   def _test_adamw(self, steps, opts, atol, rtol): self._test_optim(AdamW, torch.optim.AdamW, steps, opts, atol, rtol)
-  #TODO: use torch.muon when it comes out
-  def _test_muon(self, steps, opts, atol, rtol): self._test_optim(Muon, TorchMuon, steps, opts, atol, rtol)
 
   def test_multistep_sgd_high_lr_teeny(self): self._test_sgd(2, {'lr': 1.1, 'teeny': True}, 1e-6, 1e-5)
   def test_multistep_adam_high_lr_teeny(self): self._test_adam(2, {'lr': 1.1, 'teeny': True}, 2e-4, 5e-4)
-  def test_multistep_muon_high_lr_teeny(self): self._test_muon(2, {'lr': 1.1, 'teeny': True}, 2e-4, 5e-4)
 
   def test_sgd(self): self._test_sgd(1, {'lr': 0.001}, 1e-6, 0)
   def test_sgd_high_lr(self): self._test_sgd(1, {'lr': 10}, 1e-6, 1e-5)
@@ -87,28 +83,6 @@ class TestOptim(unittest.TestCase):
   def test_multistep_sgd_high_lr_nesterov_momentum_wd(self):
     self._test_sgd(10, {'lr': 9, 'momentum': 0.9, 'nesterov': True, 'weight_decay': 0.1}, 1e-5, 3e-4)
 
-  def test_muon(self): self._test_muon(1, {'lr': 0.001}, 1e-6, 0)
-  def test_muon_high_lr(self): self._test_muon(1, {'lr': 10}, 1e-6, 3e-4)
-  def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 0.01}, 1e-6, 0)
-  def test_muon_high_lr_wd(self): self._test_muon(1, {'lr': 10, 'weight_decay': 0.01}, 1e-6, 3e-4)
-
-  # NOTE: momentum set to 0.95 by default, nesterov set to True by default
-  def test_multistep_muon_momentum_wd(self): self._test_muon(10, {'lr': 0.001, 'weight_decay': 0.01}, 1e-5, 0)
-  # ns defaults are numerically unstable, but it is tolerable in real training (see nsteps/nparam tests)
-  def test_multistep_muon_high_lr_momentum_wd(self): self._test_muon(10, {'lr': 10, 'weight_decay': 0.01}, 1e-1, 3e-4)
-  def test_multistep_muon_no_nesterov_momentum(self): self._test_muon(10, {'lr': 0.001, 'nesterov': False}, 1e-5, 0)
-  def test_multistep_muon_high_lr_no_nesterov_momentum(self): self._test_muon(10, {'lr': 10, 'nesterov': False}, 0.5e-1, 1e-1)
-
-  def test_muon_ns_steps(self): self._test_muon(1, {'lr': 0.001, 'ns_steps': 3}, 1e-6, 0)
-  def test_muon_high_lr_ns_steps(self): self._test_muon(1, {'lr': 10, 'ns_steps': 3}, 1e-5, 3e-4)
-  def test_muon_ns_params(self): self._test_muon(1, {'lr': 0.001,'ns_params': (2.0,-1.5,0.5)}, 1e-6, 0)
-  def test_muon_high_lr_ns_params(self): self._test_muon(1, {'lr': 10,'ns_params': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
-
-  def test_muon_momentum_wd_ns_steps_ns_params(self):
-    self._test_muon(10, {'lr': 0.001, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_params': (2.0,-1.5,0.5)}, 1e-5, 0)
-  def test_multistep_muon_high_lr_momentum_wd_ns_steps_ns_params(self):
-    self._test_muon(10, {'lr': 10, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_params': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
-
   def test_adam(self): self._test_adam(1, {'lr': 0.001}, 1e-5, 0)
   def test_adam_high_lr(self): self._test_adam(1, {'lr': 10}, 1e-4, 1e-4)
   def test_adamw(self): self._test_adamw(1, {'lr': 0.001}, 1e-5, 0)

test/test_pickle.py

@@ -20,7 +20,7 @@ class TestPickle(unittest.TestCase):
     self.assertEqual(pm2.rewrite(sink).key, tt.key)
 
   def test_pickle_main_pattern_matcher(self):
-    from tinygrad.codegen.late.devectorizer import sym
+    from tinygrad.codegen.devectorizer import sym
     ssym = pickle.dumps(sym)
     dsym = pickle.loads(ssym)
     self.assertEqual(dsym.patterns[0][0].location, sym.patterns[0][0].location)

test/test_profiler.py

@@ -1,6 +1,6 @@
 import unittest, struct, contextlib, statistics, time, gc
 from tinygrad import Device, Tensor, dtypes, TinyJit
-from tinygrad.helpers import CI, getenv, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
+from tinygrad.helpers import CI, getenv, Context, ProfileRangeEvent, cpu_profile, cpu_events
 from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileDeviceEvent, ProfileGraphEvent
 from tinygrad.runtime.support.hcq import HCQCompiled
 from tinygrad.engine.realize import get_runner
@@ -209,18 +209,5 @@ class TestProfiler(unittest.TestCase):
     for ge in graphs:
       self.assertEqual(len(ge.ents), len(graphs))
 
-  def test_trace_metadata(self):
-    with Context(TRACEMETA=1):
-      a = Tensor.empty(1)+2
-      b = Tensor.empty(1)+2
-      with helper_collect_profile(TestProfiler.d0) as profile:
-        Tensor.realize(a, b)
-    profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
-    exec_points = [e for e in profile if isinstance(e, ProfilePointEvent) and e.name == "exec"]
-    range_events = [e for e in profile if isinstance(e, ProfileRangeEvent)]
-    self.assertEqual(len(exec_points), len(range_events), 2)
-    self.assertEqual(len(dedup(e.key for e in exec_points)), 1)
-    self.assertEqual(len(dedup(e.arg['metadata'] for e in exec_points)), 1)
-
 if __name__ == "__main__":
   unittest.main()

test/test_rangeify.py

@@ -1,198 +0,0 @@
-import unittest
-from tinygrad import Tensor
-from tinygrad.helpers import RANGEIFY, Context, GlobalCounters
-
-N = 256
-
-@unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY")
-class TestRangeify(unittest.TestCase):
-  def test_expand_children(self):
-    A = Tensor.empty(N, N).sum(axis=1)
-    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)
-    C = Tensor.empty(N, N)
-    (A@B@C).realize()
-
-  def test_double_gemm_exp(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    (((A@B).exp()@C).exp()).realize()
-
-  def test_double_gemm_relu(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    (((A@B).relu()@C).relu()).realize()
-
-  def test_double_gemm_relu_half_contig(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    (((A@B).relu().contiguous(arg=(1,))@C).relu()).realize()
-
-  def test_double_gemm_half_contig(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    ((A@B).contiguous(arg=(1,))@C).realize()
-
-  def test_double_gemm_contig(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    ((A@B).contiguous()@C).realize()
-
-  def test_many_gemm(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    D = Tensor.empty(N, N)
-    E = Tensor.empty(N, N)
-    F = Tensor.empty(N, N)
-    (A@B@C@D@E@F).realize()
-
-  def test_conv2d(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    x.conv2d(w1).realize()
-
-  def test_conv2d_t(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    (x*2).conv2d(w1).realize()
-
-  def test_double_conv2d(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    w2 = Tensor.empty(12, 8, 3, 3)
-    x.conv2d(w1).conv2d(w2).realize()
-
-  def test_double_conv2d_half_contig(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    w2 = Tensor.empty(12, 8, 3, 3)
-    # NOTE: this contiguous doesn't help
-    x.conv2d(w1).contiguous(arg=(1,)).conv2d(w2).permute(0,2,3,1).contiguous().realize()
-
-  def test_double_conv2d_contig(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    w2 = Tensor.empty(12, 8, 3, 3)
-    x.conv2d(w1).contiguous().conv2d(w2).realize()
-
-  def test_transformer_ffn(self):
-    from tinygrad.apps.llm import TransformerBlock
-    from tinygrad import nn
-    blk = TransformerBlock(1024, 4096, 1, 1, 1e-5)
-    for p in nn.state.get_parameters(blk): p.replace(Tensor.empty(p.shape))
-
-    x = Tensor.empty(128, 1024)
-    out = blk._feed_forward(x)
-    out.realize()
-
-  def test_flash_attention(self):
-    #BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
-
-    # bigger
-    BS, HEADS, SEQLEN, EMB = 4, 32, 1024, 64
-
-    # llama 8B
-    #BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128
-
-    def fa():
-      Tensor.manual_seed(1337)
-      with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
-      return q.scaled_dot_product_attention(k, v).realize()
-
-    with Context(DEBUG=4):
-      GlobalCounters.reset()
-      ret = fa()
-    with Context(RANGEIFY=0):
-      with Context(DEBUG=2):
-        GlobalCounters.reset()
-        cmp = fa()
-      with Context(DEBUG=0):
-        mse = ((cmp-ret)**2).sum().item()
-    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")
-class TestOuterworld(unittest.TestCase):
-  def test_passthrough_range(self):
-    t = Tensor.rand(10, 10).realize()
-
-    # passthrough ranges
-    a = UOp.range(dtypes.int, 10, -1)
-    sel = t[a]
-    cpy = sel.contiguous(a).realize()
-
-    self.assertTrue((t==cpy).all().item())
-
-  def test_flip_range(self):
-    t = Tensor.rand(10, 10).realize()
-
-    # passthrough ranges
-    a = UOp.range(dtypes.int, 10, -1)
-    sel = t[9-a]
-    cpy = sel.contiguous(a).realize()
-
-    self.assertTrue((t.flip(0)==cpy).all().item())
-
-  def test_vmap(self):
-    def f(x): return x.sum(axis=0)*2
-
-    x = Tensor.ones(3, 10, 2).contiguous()
-
-    # vmap across axis 0
-    a = UOp.range(dtypes.int, 3, -1)
-    out = f(x[a])
-    out = out.contiguous(a)
-
-    # 3x2 grid of 20
-    out.realize()
-    print(out.numpy())
-
-  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)
-    x = x.assign(x @ W[a])
-    out = x.contiguous(a)[-1].contiguous().realize()
-
-    self.assertTrue((manual==out).all().item())
-
-if __name__ == '__main__':
-  unittest.main()

test/test_schedule.py

@@ -1050,14 +1050,6 @@ 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_softmax_fusion.py

@@ -117,14 +117,6 @@ class TestFuse(unittest.TestCase):
     c = (a.sum(axis=1) + b.sum(axis=1)).fuse()
     self.assertListEqual(c.tolist(), [30]*16)
 
-  @unittest.skipUnless(Device.DEFAULT == "METAL", "METAL TC")
-  def test_fuse_and_tc_opt(self):
-    A = Tensor.randn(8, 8).realize()
-    B = Tensor.randn(8, 8).realize()
-    C = Tensor.ones(1, 8, 8).pad(((1,1), None, None),).sum(0)
-    out = (C + (A @ B)).fuse()
-    out.realize()
-
 class TestSoftmaxFusion(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
@@ -163,7 +155,7 @@ class TestSoftmaxFusion(unittest.TestCase):
       out = single_kernel_softmax(self.test)
       out.realize()
 
-    np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
+    np.testing.assert_allclose(sout.numpy(), out.numpy())
 
   def test_auto_softmax(self):
     print("*** softmax ***")
@@ -176,7 +168,7 @@ class TestSoftmaxFusion(unittest.TestCase):
       out = self.test.contiguous().softmax(-1).fuse()
       run_one_schedule_item(out)
 
-    np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
+    np.testing.assert_allclose(sout.numpy(), out.numpy())
 
   @unittest.skip("recursion error no longer raised")
   def test_softmax_bw(self):

test/test_symbolic_ops.py

@@ -229,12 +229,12 @@ class TestSymbolicOps(unittest.TestCase):
           np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_var(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]:
-        expected = a[:i, :].var(axis).numpy()
-        symbolic = a[:vi, :].var(axis).reshape(expected.shape).numpy()
+        a = Tensor.rand(i, 3)
+        expected = a.var(axis).numpy()
+        symbolic = a.reshape(vi, 3).var(axis).reshape(expected.shape).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_var_2d(self):

test/test_tensor.py

@@ -415,21 +415,6 @@ class TestTinygrad(unittest.TestCase):
         data = _generate_data(depth)
         np.testing.assert_allclose(Tensor(data).numpy(), np.array(data))
 
-  def test_tensor_list_implicit_cast(self):
-    data = [True, False]
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.int).numpy(), torch.tensor(data, dtype=torch.int).numpy())
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.uint8).numpy(), torch.tensor(data, dtype=torch.uint8).numpy())
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())
-    data = [-1, 0, 1, 2, 3]
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.int).numpy(), torch.tensor(data, dtype=torch.int).numpy())
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.uint8).numpy(), torch.tensor(data, dtype=torch.uint8).numpy())
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())
-    data = [-3.5, -2.5, -1.5, 0, 1.5, 2.5, 3.5]
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.int).numpy(), torch.tensor(data, dtype=torch.int).numpy())
-    # NOTE: torch and jax raise OverflowError: Python integer -3 out of bounds for uint8
-    # np.testing.assert_equal(Tensor(data, dtype=dtypes.uint8).numpy(), torch.tensor(data, dtype=torch.uint8).numpy())
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())
-
   def test_tensor_list_special_values(self):
     if is_dtype_supported(dtypes.float16):
       data = [math.nan, -math.inf, 65504, 65519, 65519.999, 65520, 65520.1]

test/test_tiny.py

@@ -30,10 +30,7 @@ class TestTiny(unittest.TestCase):
   def test_gemm(self, N=64, out_dtype=dtypes.float):
     a = Tensor.ones(N,N).contiguous()
     b = Tensor.eye(N).contiguous()
-    lst = (out:=a@b).tolist()
-    for y in range(N):
-      for x in range(N):
-        self.assertEqual(lst[y][x], 1.0, msg=f"mismatch at ({y},{x})")
+    self.assertListEqual((out:=a@b).flatten().tolist(), [1.0]*(N*N))
     if IMAGE < 2: self.assertEqual(out.dtype, out_dtype)
 
   # *** randomness ***
@@ -76,17 +73,17 @@ class TestTiny(unittest.TestCase):
 
   def test_symbolic(self):
     i = Variable('i', 1, 10)
-    ones = Tensor.ones(10).contiguous()
-    for s in [2,5]:
-      ret = ones[:i.bind(s)] + 1
-      self.assertListEqual(ret.contiguous().reshape(s).tolist(), [2.0]*s)
+    with Context(IGNORE_OOB=1):
+      for s in [2,5]:
+        ret = Tensor.ones(s).contiguous().reshape(i.bind(s)) + 1
+        self.assertListEqual(ret.reshape(s).tolist(), [2.0]*s)
 
   def test_symbolic_reduce(self):
     i = Variable('i', 1, 10)
-    ones = Tensor.ones(10).contiguous()
-    for s in [2,5]:
-      ret = ones[:i.bind(s)].sum()
-      self.assertEqual(ret.item(), s)
+    with Context(IGNORE_OOB=1):
+      for s in [2,5]:
+        ret = Tensor.ones(s).contiguous().reshape(i.bind(s)).sum()
+        self.assertEqual(ret.item(), s)
 
   # *** a model ***
 
@@ -109,24 +106,6 @@ class TestTiny(unittest.TestCase):
     probs = Tensor.rand(1, 1, 28, 28).sequential(layers).tolist()
     self.assertEqual(len(probs[0]), 10)
 
-  # TODO: this is failing because of how swizzling rewrites the ShapeTracker of the final STORE
-  @unittest.skipIf(IMAGE>0 or (CI and Device.DEFAULT == "DSP"), "failing because of make things that can't be images not images")
-  def test_mnist_backward(self):
-    # NOTE: we don't have the whole model here for speed
-    layers = [
-      nn.Conv2d(1, 32, 5), Tensor.relu,
-      nn.Conv2d(32, 32, 5), Tensor.relu]
-
-    # replace random weights with ones
-    # TODO: there's a bug here where it's tying two of the biases together. we need UNIQUE const
-    #Tensor.realize(*[p.replace(Tensor.ones_like(p).contiguous()) for p in nn.state.get_parameters(layers)])
-    for p in nn.state.get_parameters(layers): p.replace(Tensor.empty(p.shape))
-
-    # realize gradients
-    for x in nn.state.get_parameters(layers): x.requires_grad_()
-    Tensor.empty(4, 1, 28, 28).sequential(layers).sum().backward()
-    Tensor.realize(*[x.grad for x in nn.state.get_parameters(layers) if x.grad is not None])
-
   # *** image ***
 
   @unittest.skipIf(Device.DEFAULT != "GPU", "image only supported on GPU")

test/test_uop_graph.py

@@ -6,7 +6,7 @@ from tinygrad.helpers import DEBUG, Context
 from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp
 from tinygrad.uop.symbolic import sym
 from tinygrad.codegen import full_rewrite, full_rewrite_to_sink
-from tinygrad.codegen.late.expander import expander
+from tinygrad.codegen.expander import expander
 
 simple_pm = PatternMatcher([
   (UPat.cvar('x', dtypes.int), lambda x: UOp.const(dtypes.float, 1.0) + UOp.const(dtypes.float, 2.0)),
@@ -441,16 +441,18 @@ class TestUOpGraph(unittest.TestCase):
       ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(Variable("i", 0, 20)),))
       with self.assertRaises(RuntimeError): to_uops_list([ld0])
 
+  @unittest.skip("outdated")
   def test_in_out_of_bounds_access_gated_store(self):
     with Context(IGNORE_OOB=0):
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), src=(), arg=0)
+      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 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, (glbl0.index(v), UOp.const(dtypes.int, 0), v<16))
       to_uops_list([st0])
 
       st1 = UOp(Ops.STORE, dtypes.void, (glbl0.index(v), v, v<20))
       with self.assertRaises(RuntimeError): to_uops_list([st1])
 
+  @unittest.skip("outdated")
   def test_in_bounds_access_gated_local(self):
     with Context(IGNORE_OOB=0):
       # Define buffers
@@ -463,7 +465,7 @@ class TestUOpGraph(unittest.TestCase):
 
       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), UOp.const(dtypes.uint, 1), lidx<8))
 
       barrier = UOp(Ops.BARRIER, dtypes.void, (local_store,))
       if_barrier = UOp(Ops.IF, dtypes.void, (gate, barrier))
@@ -475,34 +477,6 @@ class TestUOpGraph(unittest.TestCase):
       global_store = UOp(Ops.STORE, dtypes.void, (gbuf.index(gidx), local_load))
       to_uops_list([global_store])
 
-  def test_load_with_float_in_index(self):
-    with Context(IGNORE_OOB=0):
-      ridx = UOp.range(dtypes.int, 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))),))
-      to_uops_list([ld0])
-      glblfloat = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(20), (), 0)
-      ldfloat = UOp(Ops.LOAD, dtypes.float, (glblfloat.index(ridx),))
-      i = (ldfloat+3.14).cast(dtypes.int)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i, ((0<=i)&(i<16))),))
-
-  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)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(ridx, ridx.cast(dtypes.bool).logical_not()),))
-      to_uops_list([ld0])
-
-  @unittest.skip("Bool load is not supported yet")
-  def test_load_mask(self):
-    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)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(ridx, ridx<16&mask),)))
-      to_uops_list([ld0])
-
   def test_out_of_bounds_off_by_one_access(self):
     with Context(IGNORE_OOB=0):
       glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
@@ -591,9 +565,10 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_switched_range_order(self):
     glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
+    c2 = UOp.const(dtypes.int, 2)
     cf = UOp.const(dtypes.float, 0.0)
-    r1 = UOp.range(dtypes.int, 2, 0)
-    r2 = UOp.range(dtypes.int, 2, 1)
+    r1 = UOp(Ops.RANGE, dtypes.int, (c2,), 0)
+    r2 = UOp(Ops.RANGE, dtypes.int, (c2,), 1)
     alu = UOp(Ops.MUL, dtypes.int, (r2, r1))
     store = UOp(Ops.STORE, dtypes.void, (glbl.index(alu), cf))
     uops = to_uops_list([store])

test/test_uops.py

@@ -402,14 +402,6 @@ class TestAssembly(unittest.TestCase):
     self.assertIn(Ops.SHR, ops)
     self.assertNotIn(Ops.IDIV, ops)
 
-  def test_fast_idiv_remove_powers_of_two(self):
-    ridx = UOp.range(dtypes.int, 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
-    # (((ridx0>>6)*18725)>>17) instead of (int)((((long)(ridx0)*1198373)>>29))
-    self.assertNotIn(Ops.CAST, ops)
-
   def test_mulacc_unrolled(self):
     # test that     acc = acc + a0*b0 + a1*b1 + a2*b2 + a3*b3
     # is not        acc = acc + (a0*b0 + a1*b1 + a2*b2 + a3*b3)

test/test_uops_stats.py

@@ -20,7 +20,6 @@ def get_stats(x:Tensor):
   ei = lower_schedule_item(si)
   return ei.prg.estimates.ops, ei.prg.estimates.mem
 
-@unittest.skipIf(Device.DEFAULT == "WEBGPU", "webgpu does extra load/store for packed types")
 class TestMemoryCount(unittest.TestCase):
   def test_add(self):
     a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)

test/unit/test_block_reorder.py

@@ -1,7 +1,7 @@
 import unittest, random
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import print_uops, UOp, Ops
-from tinygrad.codegen.late.linearize import block_reorder
+from tinygrad.codegen.linearize import block_reorder
 from tinygrad.renderer.cstyle import OpenCLRenderer
 
 def is_toposorted(lst:list[UOp]):

test/unit/test_dtype.py

@@ -53,14 +53,5 @@ class TestCastConvenienceMethod(unittest.TestCase):
       self.assertEqual(t.float().dtype, dtypes.float)
       self.assertEqual(t.double().dtype, dtypes.double)
 
-class TestDtypeTolist(unittest.TestCase):
-  def test_bfloat16(self):
-    self.assertEqual(Tensor([-60000, 1.5, 3.1, 60000], device="PYTHON", dtype=dtypes.bfloat16).tolist(), [-59904.0, 1.5, 3.09375, 59904.0])
-  def test_fp8(self):
-    # 448
-    self.assertEqual(Tensor([-30000, 1.5, 3.1, 30000], device="PYTHON", dtype=dtypes.fp8e4m3).tolist(), [-448.0, 1.5, 3.0, 448.0])
-    # 57344
-    self.assertEqual(Tensor([-30000, 1.5, 3.1, 30000], device="PYTHON", dtype=dtypes.fp8e5m2).tolist(), [-28672.0, 1.5, 3.0, 28672.0])
-
 if __name__ == "__main__":
   unittest.main()

test/unit/test_dtype_spec.py

@@ -1,6 +1,6 @@
-import unittest, math, operator, subprocess, struct
+import unittest, math, operator, subprocess
 from tinygrad.tensor import Tensor, dtypes, Device
-from tinygrad.dtype import DType, DTYPES_DICT, truncate, truncate_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float
+from tinygrad.dtype import DType, DTYPES_DICT, truncate, truncate_fp16, truncate_bf16, _to_np_dtype, least_upper_dtype, least_upper_float
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import getenv, CI, DEBUG
 from hypothesis import given, settings, strategies as strat
@@ -26,9 +26,6 @@ def _assert_eq(tensor:Tensor, target_dtype:DType, target, tol_target_dtype:float
   except AssertionError as e:
     raise AssertionError(f"\ntensor {tensor.numpy()} dtype {tensor.dtype} does not match target {target} with dtype {target_dtype}") from e
 
-def u32_to_f32(u): return struct.unpack('f', struct.pack('I', u))[0]
-def f32_to_u32(f): return struct.unpack('I', struct.pack('f', f))[0]
-
 class TestHelpers(unittest.TestCase):
   signed_ints = (dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64)
   uints = (dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64)
@@ -105,79 +102,18 @@ class TestHelpers(unittest.TestCase):
     self.assertEqual(truncate_fp16(65504), 65504)
     self.assertEqual(truncate_fp16(65519.999), 65504)
     self.assertEqual(truncate_fp16(65520), math.inf)
-    self.assertEqual(truncate_fp16(1e-8), 0.0)
-    self.assertEqual(truncate_fp16(-65504), -65504)
-    self.assertEqual(truncate_fp16(-65519.999), -65504)
-    self.assertEqual(truncate_fp16(-65520), -math.inf)
-    self.assertTrue(math.isnan(truncate_fp16(math.nan)))
 
-  def test_float_to_bf16(self):
-    # TODO: fuzz this better
+  def test_truncate_bf16(self):
+    self.assertEqual(truncate_bf16(1), 1)
+    self.assertAlmostEqual(truncate_bf16(1.1), 1.09375, places=7)
+    for a in [1234, 23456, -777.777]:
+      self.assertEqual(truncate_bf16(a), torch.tensor([a], dtype=torch.bfloat16).item())
+    # TODO: torch bfloat 1.1 gives 1.1015625 instead of 1.09375
     max_bf16 = torch.finfo(torch.bfloat16).max
-    for a in [1, 1.1, 1234, 23456, -777.777, max_bf16, max_bf16 * 1.00001, -max_bf16, -max_bf16 * 1.00001, math.inf, -math.inf]:
-      self.assertEqual(float_to_bf16(a), torch.tensor([a], dtype=torch.bfloat16).item())
-    self.assertTrue(math.isnan(float_to_bf16(math.nan)))
-
-  def test_float_to_bf16_nan(self):
-    # In f32, NaN = exp 0xFF and mantissa ≠ 0. Quiet-vs-signaling is bit 22 of the mantissa: 1 = qNaN, 0 = sNaN.
-    # qNaN(+/-), sNaN(+/-) overflow(+/-)
-    patterns = [0x7FC00001, 0xFFC00001, 0x7F800001, 0xFF800001, 0x7FFFFFFF, 0xFFFFFFFF]
-    for u in patterns:
-      x = u32_to_f32(u)
-      y = float_to_bf16(x)
-      t = torch.tensor([x], dtype=torch.bfloat16).item()
-      self.assertTrue(math.isnan(y))
-      self.assertTrue(math.isnan(t))
-
-  def test_float_to_bf16_round(self):
-    # round_to_nearest_even
-    uppers = [0x3f800000, 0x41230000, 0xC1460000] # 1.0, 10.1875, -12.375
-    for upper in uppers:
-      base = upper & 0xFFFF0000
-      base_f32 = u32_to_f32(base)
-      base_f32_round_up = u32_to_f32(base + 0x00010000)
-
-      # low < 0x8000(0.5ULP) -> round down
-      x = u32_to_f32(base | 0x00007000)
-      self.assertEqual(float_to_bf16(x), base_f32)
-      self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32)
-
-      # low > 0x8000(0.5ULP) -> round up
-      x = u32_to_f32(base | 0x0000C000)
-      self.assertEqual(float_to_bf16(x), base_f32_round_up)
-      self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32_round_up)
-
-      # low == 0x8000(0.5ULP) and LSB even -> round down
-      if ((upper >> 16) & 1) == 0:
-        x = u32_to_f32(base | 0x00008000)
-        self.assertEqual(float_to_bf16(x), base_f32)
-        self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32)
-      # low == 0x8000(0.5ULP) and LSB odd -> round up
-      else:
-        x = u32_to_f32(base | 0x00008000)
-        self.assertEqual(float_to_bf16(x), base_f32_round_up)
-        self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32_round_up)
-
-  def test_float_to_bf16_boundary(self):
-    # bf16 max finite: exp=0xFE, faction=0x7F => 0x7F7F0000(f32)
-    # bf16 inf(+/-):   exp=0xFF
-    base = 0x7F7F0000
-    inf_u32 = 0x7F800000
-
-    # low < 0.5ULP
-    x = u32_to_f32(base | 0x00007FFF)
-    self.assertEqual(f32_to_u32(float_to_bf16(x)), base)
-    self.assertEqual(f32_to_u32(torch.tensor([x], dtype=torch.bfloat16).item()), base)
-
-    # low > 0.5ULP -> overflows to +inf
-    x = u32_to_f32(base | 0x0000C000)
-    self.assertEqual(f32_to_u32(float_to_bf16(x)), inf_u32)
-    self.assertEqual(f32_to_u32(torch.tensor([x], dtype=torch.bfloat16).item()), inf_u32)
-
-    # low == 0.5ULP and LSB odd -> overflows to +inf
-    x = u32_to_f32(base | 0x00008000)
-    self.assertEqual(f32_to_u32(float_to_bf16(x)), inf_u32)
-    self.assertEqual(f32_to_u32(torch.tensor([x], dtype=torch.bfloat16).item()), inf_u32)
+    self.assertEqual(truncate_bf16(max_bf16), max_bf16)
+    self.assertEqual(truncate_bf16(min_bf16:=-max_bf16), min_bf16)
+    self.assertEqual(truncate_bf16(max_bf16 * 1.00001), math.inf)
+    self.assertEqual(truncate_bf16(min_bf16 * 1.00001), -math.inf)
 
   @given(strat.floats(width=32, allow_subnormal=True, allow_nan=True, allow_infinity=True))
   def test_truncate_fp8e4m3(self, x):

test/unit/test_gguf.py

@@ -53,37 +53,11 @@ class TestGGUF(unittest.TestCase):
   def test_load_tinyllama_q4_0(self): self._test_gguf_load("https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories15M-q4_0.gguf?download=true")
   def test_load_gpt2_q4_1(self): self._test_gguf_load("https://huggingface.co/PrunaAI/gpt2-GGUF-smashed/resolve/main/gpt2.Q4_1.gguf?download=true")
   def test_load_sample_q6_k(self): self._test_gguf_load("https://huggingface.co/Isotr0py/test-gguf-sample/resolve/main/Quant_Q6_K_1024.gguf?download=true")
-  def test_load_sample_mxfp4(self): self._test_gguf_load("https://huggingface.co/ngxson/boring-testing-tiny/resolve/main/stories260K-mxfp4.gguf?download=true")
 
   def test_dequantization_q4_0(self): self._test_dequantization(ggml.GGML_TYPE_Q4_0)
   def test_dequantization_q4_1(self): self._test_dequantization(ggml.GGML_TYPE_Q4_1)
   def test_dequantization_q8_0(self): self._test_dequantization(ggml.GGML_TYPE_Q8_0)
   def test_dequantization_q6_k(self): self._test_dequantization(ggml.GGML_TYPE_Q6_K)
-  def test_dequantization_mxfp4(self):
-    MXFP4 = 39
-
-    def encode(nibbles, E):
-      packed = [(low & 0xF) | ((high & 0xF) << 4) for low, high in zip(nibbles[:16], nibbles[16:])]
-      return np.array([E] + packed, dtype=np.uint8)
-
-    def decode(code, E):
-      sign = -1.0 if code * 0b1000 else 1.0
-      exp = (code >> 1) & 0b11
-      mant = code & 0b1
-      val = (1.0 + 0.5 * mant) * np.exp2(exp - 1) if exp else 0.5 * mant
-      scale = np.exp2(E - 128) if E >= 2 else np.exp2(-127 if E == 1 else -128)
-      return sign * val * scale
-
-    blocks, expected = [], []
-    rng = np.random.default_rng(42)
-    for _ in range(4):
-      E = rng.integers(0, 256)
-      codes = rng.integers(0, 16, size=32, dtype=np.uint8)
-      blocks.append(encode(codes, E))
-      expected.extend(decode(c, E) for c in codes)
-    tensor = Tensor(np.concatenate(blocks))
-    out = ggml_data_to_tensor(tensor, len(expected), MXFP4)
-    self.assertListEqual(out.numpy().tolist(), np.array(expected, dtype=np.float32).tolist())
 
   def test_expected_failure_unknown_type(self):
     with self.assertRaises(ValueError):

test/unit/test_graph_rewrite.py

@@ -303,8 +303,8 @@ class TestRecurse(unittest.TestCase):
   def test_inf_loop(self):
     a = UOp.variable('a', 0, 10)
     pm = PatternMatcher([
-      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.CONST)),
-      (UPat(Ops.CONST, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
+      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.DEFINE_REG)),
+      (UPat(Ops.DEFINE_REG, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
     ])
     with self.assertRaises(RuntimeError):
       graph_rewrite(a, pm)
@@ -312,8 +312,8 @@ class TestRecurse(unittest.TestCase):
   def test_inf_loop_bottom_up(self):
     a = UOp.variable('a', 0, 10)
     pm = PatternMatcher([
-      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.CONST)),
-      (UPat(Ops.CONST, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
+      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.DEFINE_REG)),
+      (UPat(Ops.DEFINE_REG, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
     ])
     with self.assertRaises(RuntimeError):
       graph_rewrite(a, pm, bottom_up=True)

test/unit/test_hashing.py

@@ -2,6 +2,7 @@ from typing_extensions import Callable
 import hashlib, random, unittest
 from tinygrad import Tensor, Device, getenv, dtypes
 from tinygrad.device import is_dtype_supported
+from tinygrad.helpers import CI
 
 @unittest.skipUnless(is_dtype_supported(dtypes.uint8) and is_dtype_supported(dtypes.uint64), "Device must support uint8 and uint64")
 @unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "NV", "crashes in NV CI")
@@ -11,7 +12,7 @@ class TestHashing(unittest.TestCase):
     chunk_hashes = [hashlib.shake_128(chunk).digest(16) for chunk in chunks]
     return hashlib.shake_128(b''.join(chunk_hashes)).digest(16)
 
-  @unittest.skip("very slow")
+  @unittest.skipIf(CI, "very slow")
   def test_abc(self):
     expected = self._python_hash_1mb(b"abc" + b"\x00" * (2**20 - 3))
     out = Tensor(b"abc").hash()
@@ -64,7 +65,7 @@ class TestKeccak(unittest.TestCase):
     data = b"\x00" * 4
     self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))
 
-    data = b"\x00" * 1000
+    data = b"\x00" * (1000 if CI else 4096)
     self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))
 
 if __name__ == "__main__":

test/unit/test_rewrite_tracked_childen.py

@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor
 from tinygrad.uop.ops import PatternMatcher, Ops, UPat, graph_rewrite, RewriteContext, UOp
-from tinygrad.schedule.kernelize import kernelize_sym, merge_views
+from tinygrad.schedule.kernelize import sym, merge_views
 
 class TestRewriteTrackedChildren(unittest.TestCase):
   @unittest.skip("track_children no longer supported")
@@ -57,7 +57,7 @@ class TestRewriteTrackedChildren(unittest.TestCase):
     extra = PatternMatcher([(UPat(Ops.REDUCE_AXIS, name="r"), print_children)])
     a = Tensor.empty(3, 3)
     r = (a+0).sum()
-    graph_rewrite(r.uop, merge_views+kernelize_sym+extra, track_children=True)
+    graph_rewrite(r.uop, merge_views+sym+extra, track_children=True)
 
 if __name__ == '__main__':
   unittest.main()

test/unit/test_shapetracker.py

@@ -6,7 +6,7 @@ from tinygrad.helpers import prod
 from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad import Variable
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite
-from tinygrad.codegen.late.devectorizer import sym
+from tinygrad.codegen.devectorizer import sym
 from itertools import product
 
 def shapetracker_getitem(st:ShapeTracker, val:int):

test/unit/test_simplify_valid_idx.py

@@ -19,7 +19,7 @@ def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UO
 
 def Special(expr, nmax): return UOp(Ops.SPECIAL, dtypes.int, (), (expr, nmax))
 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(Ops.RANGE, dtypes.int, arg=n, src=(UOp.const(dtypes.int, nmax),))
 
 class TestHelpers(unittest.TestCase):
   def test_is_increasing(self):

test/unit/test_transcendental_helpers.py

@@ -2,8 +2,8 @@ import unittest, math
 import numpy as np
 from tinygrad import dtypes
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.uop.decompositions import TRANSCENDENTAL_DTYPES, payne_hanek_reduction, cody_waite_reduction
-from tinygrad.uop.decompositions import frexp, rintk, xpow, xexp2, xlog2, trig_poly, pow2if
+from tinygrad.uop.transcendental import TRANSCENDENTAL_SUPPORTED_DTYPES, payne_hanek_reduction, cody_waite_reduction
+from tinygrad.uop.transcendental import frexp, rintk, xpow, xexp2, xlog2, trig_poly, pow2if
 from test.helpers import eval_uop
 
 class TestTranscendentalFunctions(unittest.TestCase):
@@ -89,7 +89,7 @@ class TestTranscendentalVectorizedFunctions(unittest.TestCase):
       assert u1.op == u2.op, f'expected {u1.op=} but got {u2.op=} for UOps\n{u1=}\n{u2}'
     [self._check_uops_match(x1, x2) for x1, x2 in zip((u1 if isinstance(u1, tuple) else u1.src), (u2 if isinstance(u2, tuple) else u2.src))]
 
-  def _test_vectorized(self, fxn, scalar_dtypes=TRANSCENDENTAL_DTYPES, vals=[-2,1.3,194], vcounts=[1,4,19]):
+  def _test_vectorized(self, fxn, scalar_dtypes=TRANSCENDENTAL_SUPPORTED_DTYPES, vals=[-2,1.3,194], vcounts=[1,4,19]):
     for scalar_dtype in scalar_dtypes:
       for val in vals:
         for vcount in vcounts:

test/unit/test_uop_spec.py

@@ -81,16 +81,5 @@ 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

@@ -4,11 +4,11 @@ import z3
 
 from tinygrad.dtype import dtypes, ConstType
 from tinygrad.codegen import full_rewrite
-from tinygrad.codegen.late.devectorizer import sym
+from tinygrad.codegen.devectorizer import sym
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad import Variable
-from tinygrad.uop.spec import uops_to_z3
+from tinygrad.uop.spec import z3_renderer
 
 def render(self) -> tuple[str, ConstType, ConstType]:
   # NOTE: we need STORE so the ALU op has children
@@ -30,15 +30,16 @@ class TestSymbolicPickle(unittest.TestCase):
 
 class TestSymbolic(unittest.TestCase):
   def helper_test_variable(self, v, n, m, s, test_z3:bool=True):
-    if test_z3:
-      solver = z3.Solver()
-      expr, expr_simplified = uops_to_z3(solver, v, v.simplify())
-      self.assertEqual(solver.check(expr != expr_simplified), z3.unsat, "simplified expression not equal to original")
     rendered, nmin, nmax = render(v)
     if isinstance(s, tuple): self.assertIn(rendered, s)
     else: self.assertEqual(rendered, s)
     self.assertEqual(nmin, n)
     self.assertEqual(nmax, m)
+    if test_z3:
+      solver = z3.Solver()
+      z3_sink = graph_rewrite(v.sink(v.simplify()), z3_renderer, ctx=(solver, {}))
+      expr, epxr_simplified = z3_sink.src[0].arg, z3_sink.src[1].arg
+      self.assertEqual(solver.check(expr != epxr_simplified), z3.unsat, "simplified expression not equal to original")
 
   def test_cmp_simple(self):
     self.helper_test_variable(Variable("a", 3, 8) < 4, 0, 1, "(a<4)")
@@ -127,8 +128,6 @@ class TestSymbolic(unittest.TestCase):
     b = Variable("b", 0, 8)
     self.helper_test_variable(a+a, 0, 16, "(a*2)")
     self.helper_test_variable((a+b)+b, 0, 24, "(a+(b*2))")
-    self.helper_test_variable((a*3+b)+a, 0, 40, "(b+(a*4))")
-    self.helper_test_variable((a+b)+a*3, 0, 40, "(b+(a*4))")
 
   def test_sub_self(self):
     a = Variable("a", 0, 8)
@@ -163,6 +162,10 @@ class TestSymbolic(unittest.TestCase):
   def test_div_remove(self):
     self.helper_test_variable(Variable("a", 0, 7) // 20, 0, 0, "0")
 
+  def test_div_min_max(self):
+    self.helper_test_variable(Variable("a", 1, 7) // 2, 0, 3, "(a//2)")
+    self.helper_test_variable(Variable("a", 0, 6) // 2, 0, 3, "(a//2)")
+
   def test_div_neg_min_max(self):
     self.helper_test_variable(Variable("a", 1, 7) // -2, -3, 0, "((a//2)*-1)")
     self.helper_test_variable(Variable("a", 0, 6) // -2, -3, 0, "((a//2)*-1)")
@@ -208,18 +211,6 @@ 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_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)")
-
-    self.helper_test_variable(Variable("x", 0, 10)//Variable("y", 1, 10), 0, 10, "(x//y)")
-    self.helper_test_variable(Variable("x", -10, 0)//Variable("y", 1, 10), -10, 0, "(((x*-1)//y)*-1)")
-    self.helper_test_variable(Variable("x", 0, 10)//Variable("y", -10, -1), -10, 0, "((x//(y*-1))*-1)")
-    self.helper_test_variable(Variable("x", -10, 0)//Variable("y", -10, -1), 0, 10, "((x*-1)//(y*-1))")
-
-    self.helper_test_variable(Variable("x", -10, 10)//Variable("y", 1, 10), -10, 10, "(x//y)")
-    self.helper_test_variable(Variable("x", -10, 10)//Variable("y", -10, -1), -10, 10, "((x//(y*-1))*-1)")
-
   def test_mod_factor(self):
     self.helper_test_variable(usum([Variable("a", 0, 7)*100, Variable("b", 0, 3)*50]) % 100, 0, 50, "((b%2)*50)")
 
@@ -275,16 +266,6 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(((5*Variable("a", 0, 31)) % 12) % 5, 0, 4, "(((a*5)%12)%5)")
     self.helper_test_variable((Variable("a", 0, 31) % 4) % 12, 0, 3, "(a%4)")
 
-  def test_mod_mod_wrong_sign(self):
-    v1=Variable("v1", 0, 128)
-    v3=Variable("v3", 0, 7)
-    self.helper_test_variable((((((v1%2)*2)+((v3+-1)%5))+-2)%5), -4, 4, "(((((v1%2)*2)+((v3+-1)%5))+-2)%5)")
-
-  def test_mod_mod_wrong_sign2(self):
-    v2=Variable("v2", 0, 8)
-    v3=Variable("v3", 0, 4)
-    self.helper_test_variable((((((v3+3)%7)+(v2+-2))%7)%7), -6, 6, "(((v2+((v3+3)%7))+-2)%7)")
-
   def test_mul_mul(self):
     self.helper_test_variable((Variable("a", 0, 5)*10)*9, 0, 5*10*9, "(a*90)")
 
@@ -394,17 +375,6 @@ class TestSymbolic(unittest.TestCase):
   def test_mul_div(self):
     self.helper_test_variable((Variable("a", 0, 10)*4)//4, 0, 10, "a")
 
-  def test_div_drop_small_terms(self):
-    # from openpilot, shouldnt simplify
-    gidx0 = UOp.variable("gidx0", 0, 10)
-    gidx1 = UOp.variable("gidx1", 0, 10)
-    lidx0 = UOp.variable("lidx0", 0, 1)
-    lidx1 = UOp.variable("lidx1", 0, 1)
-    ridx1005 = UOp.variable("ridx1005", 0, 2)
-    ridx1006 = UOp.variable("ridx1006", 0, 2)
-    self.helper_test_variable((lidx1+((gidx1*18)+(ridx1005*18)+(lidx0*162))+(gidx0*2)+(ridx1006*2)+-40)//18, -2, 20,
-      "(((((lidx1+(((gidx1*18)+(ridx1005*18))+(lidx0*162)))+(gidx0*2))+(ridx1006*2))+-40)//18)")
-
   def test_add_div(self):
     # careful about the lower bounds and upper bounds
     self.helper_test_variable((Variable("a", 0, 5)-2)//4, 0, 0, "0")
@@ -449,13 +419,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((-Variable("a", 10, 10))%7, -3, -3, "-3")
 
   def test_div_numerator_negative(self):
-    with Context(CORRECT_DIVMOD_FOLDING=1):
-      self.helper_test_variable((Variable("idx", 0, 9)*-10)//11, -8, 0, "(((idx*10)//11)*-1)")
-
-  def test_nest_div_negative_factor(self):
-    ridx0=UOp.variable("ridx0", 0, 9)
-    ridx1=UOp.variable("ridx1", 0, 6)
-    self.helper_test_variable(((((ridx0*-7)+ridx1)+63)//35), 0, 1, "(((ridx0//5)*-1)+1)")
+    self.helper_test_variable((Variable("idx", 0, 9)*-10)//11, -8, 0, "(((idx*10)//11)*-1)")
 
   def test_div_into_mod(self):
     self.helper_test_variable((Variable("idx", 0, 16)*4)%8//4, 0, 1, "(idx%2)")
@@ -632,23 +596,20 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(cond, 0, 1, "(a<2)")
     self.helper_test_variable(cond.where(u1, u0), 0, 1, "(a<2)")
     self.helper_test_variable(cond.where(u1, u0).where(u1, u0), 0, 1, "(a<2)")
-    self.helper_test_variable(cond.where(u0, u1), 0, 1, "((a<2)!=True)")
-    self.helper_test_variable(cond.where(u0, u1).where(u0, u1), 0, 1, "(a<2)")
 
   def test_where_combine(self):
     cond = Variable("x", 0, 3) < 2
     a = Variable("a", 0, 3)
     b = Variable("b", 0, 3)
-    c = Variable("c", 0, 3)
     aa = cond.where(a, a.ufix(0))
     bb = cond.where(b, b.ufix(1))
     self.helper_test_variable(aa, 0, 3, "(a if (x<2) else 0)")
     self.helper_test_variable(bb, 0, 3, "(b if (x<2) else 1)")
     self.helper_test_variable(aa+bb, 0, 6, "((a+b) if (x<2) else 1)")
     self.helper_test_variable(aa.maximum(bb), 0, 3, "(max(a, b) if (x<2) else 1)")
-    self.helper_test_variable((c+aa)+bb, 0, 9, "(c+((a+b) if (x<2) else 1))")
 
     # not combining because it increased total ALU
+    c = Variable("c", 0, 3)
     cc = cond.where(c, c+1)
     self.helper_test_variable(bb+cc, 0, 7, "((b if (x<2) else 1)+(c if (x<2) else (c+1)))")
 
@@ -718,10 +679,6 @@ class TestSymbolic(unittest.TestCase):
     # TODO: should z3 work?
     self.helper_test_variable(2*(2*a).reciprocal(), -math.inf, math.inf, "(1/a)", test_z3=False)
 
-  def test_trunc_noop(self):
-    a = Variable("a", 1, 10, dtypes.int)
-    self.helper_test_variable(a.trunc(), 1, 10, "a", test_z3=False)
-
 class TestSymbolicNumeric(unittest.TestCase):
   def helper_test_numeric(self, f):
     MIN, MAX = 0, 10

test/unit/test_viz.py

@@ -1,11 +1,11 @@
-import unittest, decimal, json, struct
+import unittest, decimal, json
 from dataclasses import dataclass
 
 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
 from tinygrad.device import Buffer
 
 @track_rewrites(name=True)
@@ -76,7 +76,7 @@ class TestViz(BaseTestViz):
     self.assertEqual(lineno, inner.__code__.co_firstlineno)
 
   def test_exceptions(self):
-    # VIZ tracks rewrites up to and including the error
+    # VIZ tracks rewrites up to the error
     def count_3(x:UOp):
       assert x.arg <= 3
       return x.replace(arg=x.arg+1)
@@ -85,7 +85,7 @@ class TestViz(BaseTestViz):
     with self.assertRaises(AssertionError): exec_rewrite(a, [err_pm])
     lst = get_viz_list()
     err_step = lst[0]["steps"][0]
-    self.assertEqual(err_step["match_count"], 4) # 3 successful rewrites + 1 err
+    self.assertEqual(err_step["match_count"], 3)
 
   def test_default_name(self):
     a = UOp.variable("a", 1, 10)
@@ -124,10 +124,10 @@ class TestViz(BaseTestViz):
 
   def test_inf_loop(self):
     a = UOp.variable('a', 0, 10)
-    b = a.replace(op=Ops.CONST)
+    b = a.replace(op=Ops.DEFINE_REG)
     pm = PatternMatcher([
-      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.CONST)),
-      (UPat(Ops.CONST, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
+      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.DEFINE_REG)),
+      (UPat(Ops.DEFINE_REG, 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]))
@@ -240,84 +240,34 @@ class TestVizIntegration(BaseTestViz):
     self.assertEqual(lst[0]["name"], "Schedule 1 Kernel n1")
     self.assertEqual(lst[1]["name"], prg.name)
 
-  def test_metadata_tracing(self):
-    with Context(TRACEMETA=2):
-      a = Tensor.empty(1)
-      b = Tensor.empty(1)
-      metadata = (alu:=a+b).uop.metadata
-      alu.kernelize()
-      graph = next(get_details(tracked_ctxs[0][0]))["graph"]
-    self.assertEqual(len([n for n in graph.values() if repr(metadata) in n["label"]]), 1)
-
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry
 from tinygrad.viz.serve import get_profile
 
-class TinyUnpacker:
-  def __init__(self, buf): self.buf, self.offset = buf, 0
-  def __call__(self, fmt:str) -> tuple:
-    ret = struct.unpack_from(fmt, self.buf, self.offset)
-    self.offset += struct.calcsize(fmt)
-    return ret
-
-# 0 means None, otherwise it's an enum value
-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()
-  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":[]}
-    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)})
-    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}
-
 class TestVizProfiler(unittest.TestCase):
   def test_perfetto_node(self):
     prof = [ProfileRangeEvent(device='NV', name='E_2', st=decimal.Decimal(1000), en=decimal.Decimal(1010), is_copy=False),
             ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100))]
 
-    j = load_profile(prof)
+    j = json.loads(get_profile(prof))
 
-    dev_events = j['layout']['NV']['shapes']
+    dev_events = j['layout']['NV']['timeline']['shapes']
     self.assertEqual(len(dev_events), 1)
     event = dev_events[0]
     self.assertEqual(event['name'], 'E_2')
     self.assertEqual(event['st'], 0)
     self.assertEqual(event['dur'], 10)
-    assert event['ref'] is None
 
   def test_perfetto_copy_node(self):
     prof = [ProfileRangeEvent(device='NV', name='COPYxx', st=decimal.Decimal(1000), en=decimal.Decimal(1010), is_copy=True),
-            ProfileRangeEvent(device='NV:2', name='COPYxx', st=decimal.Decimal(1000), en=decimal.Decimal(1010), is_copy=True),
-            ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100)),
-            ProfileDeviceEvent(device='NV:2', comp_tdiff=decimal.Decimal(-800), copy_tdiff=decimal.Decimal(-80))]
+            ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100))]
 
-    j = load_profile(prof)
+    j = json.loads(get_profile(prof))
 
-    event = j['layout']['NV']['shapes'][0]
+    event = j['layout']['NV']['timeline']['shapes'][0]
     self.assertEqual(event['name'], 'COPYxx')
-    self.assertEqual(event['st'], 0)   # first event
+    self.assertEqual(event['st'], 900) # diff clock
     self.assertEqual(event['dur'], 10)
 
-    event2 = j['layout']['NV:2']['shapes'][0]
-    self.assertEqual(event2['st'], 20) # second event, diff clock
-
   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)),
@@ -326,44 +276,28 @@ class TestVizProfiler(unittest.TestCase):
                               deps=[[], [0]],
                               sigs=[decimal.Decimal(1000), decimal.Decimal(1002), decimal.Decimal(1004), decimal.Decimal(1008)])]
 
-    j = load_profile(prof)
+    j = json.loads(get_profile(prof))
 
-    tracks = list(j['layout'])
-    self.assertEqual(tracks[0], 'NV Graph')
-    self.assertEqual(tracks[1], 'NV')
-    self.assertEqual(tracks[2], 'NV:1')
+    devices = list(j['layout'])
+    self.assertEqual(devices[0], 'NV Graph')
+    self.assertEqual(devices[1], 'NV')
+    self.assertEqual(devices[2], 'NV:1')
 
-    nv_events = j['layout']['NV']['shapes']
+    nv_events = j['layout']['NV']['timeline']['shapes']
     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']['timeline']['shapes']
     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']['timeline']['shapes']
     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'])
 
-  def test_bytes_per_kernel(self):
-    step = 10
-    n_events = 1_000
-    prof = [ProfileRangeEvent("CPU", name="k_test", st=decimal.Decimal(ts:=i*step), en=decimal.Decimal(ts)+step) for i in range(n_events)]
-    sz = len(get_profile(prof))
-    self.assertLessEqual(sz/n_events, 26)
-
-  # can pack up to 1hr 11 min of trace events
-  def test_trace_duration(self):
-    dur_mins = 72
-    n_events = 1_000
-    step = decimal.Decimal(dur_mins*60*1e6//n_events)
-    prof = [ProfileRangeEvent("CPU", name="k_test", st=decimal.Decimal(ts:=i*step), en=decimal.Decimal(ts)+step) for i in range(n_events)]
-    with self.assertRaises(struct.error):
-      get_profile(prof)
-
 def _alloc(b:int):
   a = Tensor.empty(b, device="NULL", dtype=dtypes.char)
   a.uop.buffer.allocate()
@@ -373,29 +307,38 @@ class TestVizMemoryLayout(BaseTestViz):
   def test_double_alloc(self):
     a = _alloc(1)
     _b = _alloc(1)
-    profile_ret = load_profile(Buffer.profile_events)
-    ret = profile_ret["layout"][f"{a.device} Memory"]
+    profile_ret = json.loads(get_profile(Buffer.profile_events))
+    ret = profile_ret["layout"][a.device]["mem"]
     self.assertEqual(ret["peak"], 2)
-    self.assertEqual(len(ret["shapes"]), 2)
+    self.assertEqual(ret["shapes"][0]["x"], [0, 2])
+    self.assertEqual(ret["shapes"][1]["x"], [1, 2])
 
   def test_del_once(self):
     a = _alloc(1)
     del a
     b = _alloc(1)
-    profile_ret = load_profile(Buffer.profile_events)
-    ret = profile_ret["layout"][f"{b.device} Memory"]
+    profile_ret = json.loads(get_profile(Buffer.profile_events))
+    ret = profile_ret["layout"][b.device]["mem"]
     self.assertEqual(ret["peak"], 1)
-    self.assertEqual(len(ret["shapes"]), 3)
+    self.assertEqual(ret["shapes"][0]["x"], [0, 2])
+    self.assertEqual(ret["shapes"][1]["x"], [2, 3])
+    self.assertEqual(ret["shapes"][0]["y"], [0, 0])
+    self.assertEqual(ret["shapes"][1]["y"], [0, 0])
 
   def test_alloc_free(self):
     a = _alloc(1)
     _b = _alloc(1)
     del a
     c = _alloc(1)
-    profile_ret = load_profile(Buffer.profile_events)
-    ret = profile_ret["layout"][f"{c.device} Memory"]
+    profile_ret = json.loads(get_profile(Buffer.profile_events))
+    ret = profile_ret["layout"][c.device]["mem"]
     self.assertEqual(ret["peak"], 2)
-    self.assertEqual(len(ret["shapes"]), 4)
+    self.assertEqual(ret["shapes"][0]["x"], [0, 3])
+    self.assertEqual(ret["shapes"][1]["x"], [1, 3, 3, 4])
+    self.assertEqual(ret["shapes"][0]["y"], [0, 0])
+    self.assertEqual(ret["shapes"][1]["y"], [1, 1, 0, 0])
+    self.assertEqual(ret["shapes"][2]["x"], [3, 4])
+    self.assertEqual(ret["shapes"][2]["y"], [1, 1])
 
 if __name__ == "__main__":
   unittest.main()

test_driven_development.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+python3 test/external/process_replay/reset.py
+CAPTURE_PROCESS_REPLAY=1 pytest -n auto test/test_tiny.py test/test_uop_graph.py test/test_ops.py test/test_linearizer.py
+while true; do
+  if python3 test/test_tiny.py; then
+    PYTHONPATH="." python3 test/external/process_replay/process_replay.py
+  fi
+done

tinygrad/codegen/__init__.py

@@ -1,7 +1,7 @@
 from typing import Any, Callable
 import functools
 from dataclasses import dataclass
-from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL, RANGEIFY, POSTOPT
+from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp
 from tinygrad.uop.spec import type_verify
 from tinygrad.renderer import Renderer
@@ -9,18 +9,15 @@ from tinygrad.renderer import Renderer
 # import all pattern matchers here
 from tinygrad.codegen.lowerer import pm_lowerer, get_index
 from tinygrad.codegen.quantize import pm_quant
-from tinygrad.codegen.gpudims import pm_add_gpudims, pm_tensor_cores, pm_group_for_reduce, pm_fix_locals, pm_bufferize_loop
+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, pm_pre_expander
-from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
+from tinygrad.uop.optional import get_late_rewrite_patterns
+from tinygrad.codegen.expander import migrate_indexing, expander
+from tinygrad.codegen.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
   ReduceContext, correct_load_store, pm_render
-from tinygrad.codegen.late.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
-from tinygrad.codegen.opt import pm_get_optimization, pm_do_optimize, pm_postrange_opt
+from tinygrad.codegen.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
+from tinygrad.codegen.opt import pm_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
-from tinygrad.codegen.opt.postrange import pm_flatten_range
 
 @dataclass
 class RewriteStep:
@@ -47,10 +44,10 @@ rewrites_for_linearizer = [
 
 def get_rewrites_for_renderer(opts:Renderer, linearizer:bool=True) -> list[RewriteStep]:
   # cache with the values of the context vars
-  return _get_rewrites_for_renderer(opts, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value, RANGEIFY.value, POSTOPT.value)
+  return _get_rewrites_for_renderer(opts, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value)
 
 @functools.cache
-def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL, _RANGEIFY, _POSTOPT) -> list[RewriteStep]:
+def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL) -> list[RewriteStep]:
   # ** lowerer (rewrite_shapetracker_with_index) **
   ret: list[RewriteStep] = []
 
@@ -58,39 +55,22 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
   ret.extend(rewrites_for_views)
 
   # this is kernel.py
-  if not _RANGEIFY: ret.append(RewriteStep(pm_get_optimization, ctx=lambda _: opts, name="get optimization"))
-
-  if not _POSTOPT and not _RANGEIFY: ret.append(RewriteStep(pm_do_optimize, ctx=lambda _: opts, name="optimize ast"))
+  ret.append(RewriteStep(pm_optimize, ctx=lambda _: opts, name="optimize ast"))
 
   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))
-  
-  # add tensor cores
-  if _RANGEIFY:
-    ret.append(RewriteStep(pm_bufferize_loop, name="bufferize loop"))
-    ret.append(RewriteStep(pm_tensor_cores, lambda _: ({}, opts), name="tensor cores", bottom_up=True))
-
-  if _POSTOPT or _RANGEIFY: 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"))
 
-  ret.append(RewriteStep(pm_fix_locals, name="fix locals"))
-
-  # add locals
-  ret.append(RewriteStep(pm_flatten_range+pm_add_buffers_local+rangeify_codegen+pm_group_for_reduce, name="add local buffers"))
-
-  # 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+pm_pre_expander+expander, name="expander"))
+  ret.append(RewriteStep(sym+expander, name="expander"))
 
   # ** devectorizer (full_graph_rewrite) **
   # 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
+  # add gpu dims (late)
   ret.append(RewriteStep(pm_add_gpudims, lambda _: opts, name="add gpudims"))
 
   # devectorize (TODO: does this need opts?)
@@ -102,12 +82,12 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
   supported_ops = tuple(opts.code_for_op.keys())
   extra_matcher = opts.extra_matcher if opts.extra_matcher is not None else PatternMatcher([])
 
-  # optional pre matcher
-  if opts.pre_matcher is not None: ret.append(RewriteStep(opts.pre_matcher, name="pre_matcher"))
-
   # decompositions
   pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, _TRANSCENDENTAL>=2)
-  ret.append(RewriteStep(pm_decomp, lambda _: opts.device, name="decompositions"))
+  ret.append(RewriteStep(pm_decomp, name="decompositions"))
+
+  # optional pre matcher
+  if opts.pre_matcher is not None: ret.append(RewriteStep(opts.pre_matcher, name="pre_matcher"))
 
   # final rules for the renderer (without sym)
   pm_final_rewrite = pm_decomp+pm_render+extra_matcher

tinygrad/codegen/devectorizer.py

@@ -232,21 +232,17 @@ def no_vectorized_alu(alu:UOp):
   alus = tuple(UOp(alu.op, alu.dtype.scalar(), tuple(s.gep(i) for s in alu.src), alu.arg) for i in range(alu.dtype.vcount))
   return UOp(Ops.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)
-
-def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp):
-  cnt = cast.dtype.count
-  assert idx.dtype.count == 1, f"idx dtype must be 1 {idx.dtype}"
-  return buf.broadcast(cnt).index(idx.broadcast(cnt)*cnt+UOp.const(dtypes.int.vec(cnt), tuple(range(cnt))))
+def no_vectorized_acc(acc:UOp, c:UOp):
+  if acc.dtype.count == 1: return None
+  assert c.arg == 0, "this only supports index 0"
+  new_acc = acc.replace(dtype=acc.dtype.base.scalar().ptr(acc.dtype.count, cast(PtrDType, acc.dtype).addrspace))
+  return UOp(Ops.PTRCAT, acc.dtype, tuple([new_acc.index(UOp.const(dtypes.int, i)) for i in range(acc.dtype.count)]))
 
 devectorize = PatternMatcher([
   # no ALU on vectorized dtypes
   (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name="alu"), no_vectorized_alu),
   (UPat(Ops.WMMA, name="wmma"), no_vectorized_wmma),
-  (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG), name="buf"), no_vectorized_buf),
-  (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG), name="buf").cast(name="cast").index(UPat.var("idx")), no_vectorized_index),
+  (UPat(Ops.DEFINE_REG, name="acc").index(UPat.cvar("c")), no_vectorized_acc),
 ])
 
 pm_render = PatternMatcher([
@@ -289,7 +285,7 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp):
     topo = inp.toposort()
     stored_ranges = flatten([x.src[2:] for x in topo if x.op is Ops.STORE])
     input_ranges = tuple([x for x in topo if x.op is Ops.RANGE and x not in reduce_range and x not in stored_ranges])
-    identity = red.const(red.dtype, identity_element(red.arg, red.dtype.scalar()))
+    identity = red.const_like(identity_element(red.arg, red.dtype.scalar()))
     acc = UOp(Ops.DEFINE_REG, red.dtype.ptr(size=1, addrspace=AddrSpace.REG), arg=(ctx.acc_num,)).index(UOp.const(dtypes.int, 0))
     do_store = acc.store(identity, UOp(Ops.NOOP, src=input_ranges)) if len(input_ranges) else acc.store(identity)
     lst = [acc.load(do_store, *reduce_range)] + lst  # put acc as the first element
@@ -368,7 +364,8 @@ def reduce_collapse(red:UOp):
         replaces[s] = UOp(Ops.DEFINE_VAR, dtype=s.dtype, arg=(f'in{len(replaces)}', s.vmin, s.vmax))
   collapse_fxn = red.substitute(replaces)
   sink = graph_rewrite(collapse_fxn, pm_reduce_collapse, name="reduce_collapse")
-  if any(x.op is Ops.RANGE for x in sink.toposort()): return None
+  # TODO: why is REDUCE needed here and just RANGE isn't enough?
+  if any(x.op in {Ops.REDUCE, Ops.RANGE} for x in sink.toposort()): return None
   return sink.substitute({v:k for k,v in replaces.items()})
 
 def reduce_unparented(red:UOp):

tinygrad/codegen/expander.py

@@ -1,9 +1,8 @@
 # this converts a lowerer program into a vectorized program
 
 import functools, itertools, operator
-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
+from tinygrad.helpers import AMX, dedup, flatten, all_same, prod
+from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp
 
 def _expand_arg_to_idx(args:tuple[tuple[int, int], ...], rpk:dict[int, int]) -> int:
   idx, mul = 0, 1
@@ -47,13 +46,11 @@ def do_expand(root:UOp):
         new_srcs.append(src.src[0].gep(tuple(lst)))
     else:
       # non-UNROLL input
-      if root.op is Ops.IF or src.op is Ops.IF:
+      if root.op is Ops.IF:
         # for the first arg of IF, just pass them through ignoring UNROLLS
         new_srcs.append(src)
-      elif (root.op is Ops.STORE and i >= 2) or (root.op in {Ops.REDUCE, Ops.BUFFERIZE} and i >= 1) or (root.op is Ops.WMMA and i >= 3):
-        # for any range args of STORE/REDUCE, pass them through
-        new_srcs.append(src)
-      elif root.op is Ops.INDEX and i >= 1 and not isinstance(root.dtype, PtrDType):
+      elif root.op in {Ops.REDUCE, Ops.STORE} and src.op is Ops.RANGE:
+        # for any range args of REDUCE, pass them through
         new_srcs.append(src)
       elif src.dtype.count > 1:
         # put any input dtype > 1 grouped together
@@ -75,7 +72,7 @@ def do_contract(con:UOp):
   # CONTRACT without UNROLL repeats the element VECTORIZED
   if ex.op is not Ops.UNROLL: return UOp(Ops.VECTORIZE, con.dtype, con.src*con.dtype.count)
   # CONTRACT may remove several axes from UNROLL
-  assert con.dtype == dtypes.void or con.dtype.count == prod([x[1] for x in con.arg]), "dtype is wrong"
+  assert con.dtype.count == prod([x[1] for x in con.arg]), "dtype is wrong"
   idxs = []
   for rpk in _choices_from_args(new_ex_args:=tuple(x for x in ex.arg if x not in con.arg)):
     idxs += [_expand_arg_to_idx(ex.arg, {**rpk, **lrpk}) for lrpk in _choices_from_args(con.arg)]
@@ -86,7 +83,7 @@ expander = PatternMatcher([
   (UPat(Ops.UNROLL, name="outer", src=(UPat(Ops.UNROLL, name="inner"),)),
    lambda outer, inner: UOp(Ops.UNROLL, outer.dtype, (inner.src[0],), inner.arg+outer.arg)),
   # do expansion
-  (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.GEP, Ops.WMMA, Ops.LOAD, Ops.STORE, Ops.INDEX, Ops.BUFFERIZE,
+  (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.GEP, Ops.WMMA, Ops.LOAD, Ops.STORE, Ops.INDEX,
          Ops.VECTORIZE, Ops.IF, Ops.REDUCE), name="root", custom_early_reject=set([Ops.UNROLL])), do_expand),
   (UPat(Ops.CONTRACT, name="con"), do_contract),
   # BARRIERs aren't actually expanded
@@ -114,49 +111,3 @@ 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/gpudims.py

@@ -1,10 +1,9 @@
-import math, functools, operator
-from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType, graph_rewrite
-from tinygrad.helpers import all_int, partition, flatten, prod, dedup, USE_TC, DEBUG
-from tinygrad.dtype import dtypes, AddrSpace
+import math
+from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType
+from tinygrad.helpers import all_int
+from tinygrad.dtype import dtypes
 from tinygrad.shape.view import get_contraction
 from tinygrad.renderer import Renderer
-from tinygrad.codegen.opt.tc import TensorCore
 
 def _group_dims(dims:tuple[sint, ...], max_sizes:tuple[int, ...]):
   # TODO: symbolic shape
@@ -53,24 +52,20 @@ def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|No
 
 def add_gpudims(ctx:Renderer, s:UOp):
   if s.arg is None: return None
+  ki: KernelInfo = s.arg
+  global_dims = [i for i,x in enumerate(ki.axis_types) if x is AxisType.GLOBAL]
+  local_dims = [i for i,x in enumerate(ki.axis_types) if x in (AxisType.LOCAL, AxisType.GROUP_REDUCE)]
+  if not global_dims and not local_dims: return None
   s_topo = list(s.toposort())
   if any(x.op is Ops.SPECIAL for x in s_topo): return None
 
-  # get ranges
-  all_ranges = {x.arg[0:-1]:x for x in s_topo if x.op is Ops.RANGE}
-
-  # extract global/local dims
-  global_dims = sorted(dedup([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] is AxisType.GLOBAL]))
-  local_dims = sorted(dedup([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] in (AxisType.LOCAL, AxisType.GROUP_REDUCE)]))
-  if not global_dims and not local_dims: return None
-
   # get global and local shape
+  all_ranges = {x.arg%1000:x for x in s_topo if x.op is Ops.RANGE}
   ranges = [all_ranges[r] for r in global_dims+local_dims if r in all_ranges]
-  global_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg[0:-1] in global_dims])
-  local_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg[0:-1] in local_dims])
+  global_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg%1000 in global_dims])
+  local_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg%1000 in local_dims])
 
   # get the idxs
-  ki: KernelInfo = s.arg
   if ki.dont_use_locals:
     assert not local_dims, "can't use locals if there's no local dims"
     idxs = get_grouped_dims("idx", global_shape, ctx.global_max, reverse=True)
@@ -83,187 +78,12 @@ 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:-1])
-      if r.arg[1] == AxisType.REDUCE: continue
+      ii = (global_dims+local_dims).index(r.arg%1000)
+      if r.arg < 2000 and ki.axis_types[r.arg%1000] == AxisType.GROUP_REDUCE: continue
       subs[r] = idxs[ii]
     except ValueError: continue
   return s.substitute(subs)
 
-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).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_group_for_reduce = PatternMatcher([
-  # fix group for reduce
-  (UPat(Ops.REDUCE, name="x"), fix_group_for_reduce),
-])
-
 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),
 ])
-
-def apply_tensor_cores(ctx:tuple[dict, Renderer], in0:UOp, in1:UOp, r_range:UOp, reduceop:UOp):
-  if not USE_TC: return None
-  # 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])
-  #print(len(in0_ranges), len(in1_ranges))
-  if not len(in0_ranges) or not len(in1_ranges): return None
-  in0_range, in1_range = in0_ranges[0], in1_ranges[0]
-  if DEBUG >= 2: print('TC', in0_range.arg, in1_range.arg, r_range.arg)
-
-  # confirm the dtype and size is good
-  tc_opts: list[TensorCore] = []
-  for tc in ctx[1].tensor_cores:
-    if reduceop.dtype == tc.dtype_out and in0.dtype == tc.dtype_in and in1.dtype == tc.dtype_in:
-      if all(i <= j for i,j in zip(tc.dims, [in0_range.vmax+1, in1_range.vmax+1, r_range.vmax+1])):
-        tc_opts.append(tc)
-  if len(tc_opts) == 0: return None
-  tc = tc_opts[0]
-
-  # create the new ranges as speced by the tensor core
-  old_range = [in0_range, in1_range, r_range]
-  new_range = [r.replace(src=(r.src[0]//tc.dims[i],)) for i,r in enumerate(old_range)]
-  new_reduce_range = new_range[2]
-  tc_range = 9050 #+ r_range.arg[0]*100
-  red_ranges = []
-
-  ne: list[UOp] = []
-  for o in tc.opts:
-    lrange = UOp.range(dtypes.int, 2, tc_range, AxisType.UPCAST if o[0] == "u" else AxisType.LOCAL)
-    ne.append(lrange)
-    tc_range += 1
-    new_range[1-int(o[1])] = (2 * new_range[1-int(o[1])]) + lrange
-  for _, amt in tc.get_reduce_axes():
-    lrange = UOp.range(dtypes.int, amt, tc_range, AxisType.UNROLL)
-    ne.append(lrange)
-    red_ranges.append(lrange)
-    tc_range += 1
-    new_range[2] = (amt * new_range[2]) + lrange
-  tne = [x.replace(tag=1) for x in ne]
-
-  # replace ranges in other parts of the graph
-  for x,y in zip(old_range, new_range): ctx[0][x] = y
-
-  # apply the swizzled ranges to the srcs
-  srcs = [s.substitute(dict(zip(old_range, new_range))).substitute(dict(zip(ne, tne))) for s in (in0, in1)]
-  srcs = [x.substitute(dict(zip(tne, [ne[i] for i in p]))) for x,p in zip(srcs, tc.permutes_for_shape_str(tc.base_shape_str()))]
-
-  ned = dict(zip(tc.base_shape_str(), ne))
-  tc_reduce_axes = tuple([ned[f"r{i}"].arg[0] for i in range(len(tc.get_reduce_axes()))])
-  base_upcast_axes = tuple([(ned[s].arg[0], 2) for s in tc.base_upcast_axes()])
-  tc_upcast_axes = tuple([base_upcast_axes[:int(math.log2(tc.elements_per_thread[i]))] for i in range(3)])
-
-  # construct the op
-  # TODO: remove tc_upcast_axes from the arg
-  wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, ctx[1].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])
-  ret = tc_uop.reduce(new_reduce_range, arg=Ops.ADD)
-  # confirm the UNROLLs aren't actually used, these need to be broadcast MUL
-  assert all(u not in red_ranges for u in ret.toposort()), "UNROLLs in TC"
-  return ret
-
-from tinygrad.codegen.opt.postrange import pm_flatten_range
-
-pm_tensor_cores = PatternMatcher([
-  ((UPat.var("in0")*UPat.var("in1")).reduce(UPat(Ops.RANGE, name="r_range"), name="reduceop", arg=Ops.ADD), apply_tensor_cores),
-
-  # replace range
-  #(UPat(Ops.RANGE, name="r"), lambda ctx,r: ctx[0].get(r, None)),
-  (UPat(Ops.SINK, name="s"), lambda ctx,s: graph_rewrite(s.substitute(ctx[0]), pm_flatten_range, name="flatten")),
-])
-
-def fix_bufferize(x:UOp):
-  if x.arg != AddrSpace.LOCAL: return None
-  locals_left = [r for r in x.ranges if r.arg[1] == AxisType.LOCAL]
-  if not len(locals_left): return None
-  acc = x.size
-  st = []
-  for l in locals_left:
-    st.append(l*acc)
-    acc *= l.vmax+1
-  return x.replace(src=(x.src[0],) + tuple(locals_left[::-1]) + x.src[1:]).index(sum(st))
-
-
-pm_double_index = PatternMatcher([
-  (UPat(Ops.INDEX, src=(UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("x"))), UPat.var("y"))), lambda b,x,y: b.index(x+y)),
-])
-
-pm_fix_locals = pm_double_index+PatternMatcher([
-  (UPat(Ops.BUFFERIZE, name="x"), fix_bufferize),
-])
-
-B = 8
-
-def loop_store(x:UOp):
-  r_maybe = [r for r in x.src[0].ranges if r.arg[0] == 2 and r.tag is None]
-  ur_maybe = [r for r in x.ranges if r.arg[0] < 0]
-  #print("store", len(r_maybe), len(ur_maybe))
-  if not len(r_maybe) or not len(ur_maybe): return None
-
-  r = r_maybe[0]
-  ur = ur_maybe[0]
-  rr = r.replace(src=(r.src[0]//B,), tag=1)
-  return x.substitute({r:rr*B+ur})
-
-def loop_bufferize(x:UOp):
-  if x.arg != AddrSpace.LOCAL: return None
-  r_maybe = [r for r in x.ranges if r.arg[0] == 2 and r.tag is None]
-  ur_maybe = [r for r in x.ranges if r.arg[0] < 0]
-
-  if len(ur_maybe):
-    ur = ur_maybe[0]
-  else:
-    if len(r_maybe) == 0: return None
-    r = r_maybe[0]
-    ur = UOp.range(dtypes.int, B, -2010)
-    rr = r.replace(src=(r.src[0]//B,), tag=1)
-    x = x.substitute({r:rr*B + ur})
-
-  ur1 = UOp.range(dtypes.int, B, ur.arg[0]+1)
-  return x.replace(src=(x.src[0],ur)+x.src[1:]).index(x.size*ur1)
-
-pm_bufferize_loop = pm_double_index+PatternMatcher([
-  (UPat(Ops.BUFFERIZE, name="x"), loop_bufferize),
-  (UPat(Ops.STORE, name="x"), loop_store),
-])

tinygrad/codegen/linearize.py

@@ -3,7 +3,7 @@ import heapq
 from collections import defaultdict
 from dataclasses import dataclass, replace
 from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp
-from tinygrad.helpers import dedup, all_same, flatten, BLOCK_REORDER
+from tinygrad.helpers import dedup, all_same, flatten, getenv
 
 # NOTE: any toposort should be valid here, unlike last time this isn't required, it's just for speed
 def block_reorder(lst:list[UOp]) -> list[UOp]:
@@ -150,7 +150,7 @@ def make_block_bottom_up(ctx:BlockContext, x:UOp):
     srcs.append(add_blockends(base_block, new_ctx, current_ctx))
 
   lst = lst[::-1]
-  if BLOCK_REORDER: lst = block_reorder(lst)
+  if getenv("BLOCK_REORDER", 1): lst = block_reorder(lst)
   bb = BasicBlock(tuple(lst), ctx=current_ctx, cnt=child_count, child_ctx=child_ctx)
   return UOp(Ops.BLOCK, src=tuple(srcs), arg=bb)
 
@@ -227,7 +227,7 @@ block_merge = PatternMatcher([
 
 def finalize(sink:UOp) -> UOp:
   if sink.op is not Ops.BLOCK or not all(x.op in DONT_PLACE_IN_BLOCK for x in sink.src):
-    raise RuntimeError(f"linearize failure {sink.op} {[x.op for x in sink.src if x.op not in DONT_PLACE_IN_BLOCK]}")
+    raise RuntimeError("linearize failure")
 
   # place the early things
   lst = sorted(dedup(sink.src), key=lambda x: x.tuplize) + list(sink.arg.lst)

tinygrad/codegen/lowerer.py

@@ -1,7 +1,10 @@
 # the job of the lowerer is to do indexing
+import functools, operator
+from typing import cast
 from dataclasses import dataclass
-from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint_to_uop, AxisType, graph_rewrite, resolve
+from tinygrad.dtype import dtypes, AddrSpace, PtrDType
+from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint_to_uop, AxisType, graph_rewrite
+from tinygrad.helpers import prod, partition, flatten
 
 # ***** indexing *****
 
@@ -12,12 +15,20 @@ 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, at) for i, (s, at) in enumerate(zip(s, axis_types))]
+  # indexes
+  idxs = []
+  for i, (s, at) in enumerate(zip(s, axis_types)):
+    if at in (AxisType.UPCAST, AxisType.UNROLL):
+      assert isinstance(s, int), "needs to be int to upcast/unroll"
+      idxs.append(UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(s), tuple(range(s))),), ((i,s),), tag=1))
+    else:
+      # all others are RANGES
+      idxs.append(UOp(Ops.RANGE, dtypes.int, (sint_to_uop(s),), start+i))
+  return idxs
 
 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 resolve(s != fs) else AxisType.LOOP for s,fs in zip(ast.shape, ast.full_shape)])
+  if len(ast.full_shape) != len(axis_types): axis_types = (AxisType.LOOP,)*len(ast.full_shape)
   return IndexContext(axis_types, [], 0)
 
 # ***** lowering (given index) *****
@@ -31,8 +42,16 @@ def lower_reduce_axis(ctx: IndexContext, x: UOp):
   new_idxs = shape_to_idx(x.src[0].shape, ctx.axis_types, ctx.start)
   full_new_idx = list(ctx.idxs)
   for a in x.axis_arg: full_new_idx[a] = new_idxs[a]
+
   ret = subblock(ctx, full_new_idx, x.src[0])
-  return UOp(Ops.REDUCE, x.dtype, (ret,)+tuple([full_new_idx[i] for i in x.axis_arg]), x.arg[0])
+
+  # NOTE: always using ridxs is fine here
+  reduce_range, reduce_expand = partition([full_new_idx[i] for i in x.axis_arg], lambda y: y.op is Ops.RANGE)
+  assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand} for {x.axis_arg}"
+  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 UOp(Ops.REDUCE, x.dtype, (ret,)+tuple(reduce_range), x.arg[0])
 
 def lower_store(ctx: IndexContext, x: UOp, buf: UOp):
   # TODO: reenable after REDUCE_AXIS is fixed
@@ -48,7 +67,15 @@ def lower_store(ctx: IndexContext, x: UOp, buf: UOp):
 
   stored = subblock(ctx, real_new_idxs, x.src[1])
   used_ranges = [x for x in used_idxs if x.op is Ops.RANGE]
-  return buf.index(idx, valid).store(stored, *used_ranges)
+  ret = buf.index(idx, valid).store(stored, *used_ranges)
+
+  # insert BARRIER if we are ending a LOCAL, IF if we are ending a GROUP_REDUCE
+  if cast(PtrDType, buf.dtype).addrspace == AddrSpace.LOCAL and \
+      any(ctx.axis_types[x.arg%1000] in {AxisType.GROUP_REDUCE, AxisType.LOCAL} for x in used_ranges):
+    ret = ret.barrier()
+    range_gates = [x.eq(0) for x in used_ranges if ctx.axis_types[x.arg%1000] == AxisType.GROUP_REDUCE]
+    if len(range_gates): ret = UOp(Ops.IF, src=(functools.reduce(operator.and_, range_gates), ret))
+  return ret
 
 def fixup_wmma(ctx:IndexContext, x:UOp):
   if x.tag is not None: return None
@@ -59,8 +86,8 @@ def fixup_wmma(ctx:IndexContext, x:UOp):
   srcs = subblock(ctx, full_new_idx, UOp.sink(*x.src)).src
 
   # NOTE: this assumes these are expanded. which now shouldn't change anything
-  new_x_arg_m2 = tuple([tuple([(full_new_idx[a].arg[0], sz) for a,sz in v]) for v in x.arg[-2]])
-  new_x_arg_m1 = tuple([full_new_idx[a].arg[0] for a in x.arg[-1]])
+  new_x_arg_m2 = tuple([tuple([(full_new_idx[a].arg[0][0], sz) for a,sz in v]) for v in x.arg[-2]])
+  new_x_arg_m1 = tuple([full_new_idx[a].arg[0][0] for a in x.arg[-1]])
   return x.replace(src=srcs, arg=x.arg[:-2]+(new_x_arg_m2, new_x_arg_m1), tag=1)
 
 pm_lowerer = PatternMatcher([
@@ -83,5 +110,5 @@ pm_lowerer = PatternMatcher([
 
   # axis fixups for WMMA
   (UPat((Ops.CONTRACT, Ops.UNROLL), name="x"),
-   lambda ctx,x: x.replace(tag=1, arg=tuple([(ctx.idxs[a].arg[0], sz) for a,sz in x.arg])) if x.tag is None else None),
+   lambda ctx,x: x.replace(tag=1, arg=tuple([(ctx.idxs[a].arg[0][0], sz) for a,sz in x.arg])) if x.tag is None else None),
 ])

tinygrad/codegen/opt/__init__.py

@@ -1,9 +1,8 @@
 # opt opinionatedly transforms an ast into an optimized ast using either heuristics or beam search
 
 from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt.postrange import RKernel, pm_flatten_range
 from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
-from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, KernelInfo
+from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops
 from tinygrad.helpers import NOOPT, BEAM, USE_TC, getenv
 from tinygrad.renderer import Renderer
 from tinygrad.uop.spec import type_verify
@@ -20,40 +19,20 @@ def get_optimized_ast(ast:UOp, renderer:Renderer) -> UOp:
     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 ast.arg is not None and ast.arg.opts_to_apply is not None: k.apply_opts(ast.arg.opts_to_apply)
+  elif 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),
-])
-
-# ** postrange **
-
-def apply_ropt(ast:UOp, renderer:Renderer):
-  k = RKernel(ast, opts=renderer)
-  if ast.arg is not None: k.apply_opts(ast.arg.opts_to_apply)
-  return k.get_optimized_ast()
-
-pm_postrange_opt = pm_flatten_range+PatternMatcher([
-  (UPat(Ops.SINK, name="ast"), lambda ctx,ast: apply_ropt(ast, ctx) if ast.arg is None or \
-   (ast.arg is not None and ast.arg.opts_to_apply is not None) else None),
+pm_optimize = PatternMatcher([
+  (UPat(Ops.SINK, name="ast"), lambda ctx,ast:
+   get_optimized_ast(ast, ctx) if (ast.arg is None or ast.arg.opts_to_apply is not None) and ast.src[0].st is not None else None),
 ])

tinygrad/codegen/opt/heuristic.py

@@ -28,7 +28,7 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
           return k.applied_opts
 
   # are we grouping? (requires local shape support)
-  if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= 2048, False):
+  if resolve(prod(k.sts[0].shape[i] for i in k.upcastable_dims) <= 2048, False):
     for sz in [16]:
       try:
         k.apply_opt(Opt(OptOps.GROUPTOP, 0, sz))
@@ -62,7 +62,7 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
   # potentially do more upcasts of non reduce axes based on a heuristic
   is_dsp = k.opts is not None and k.opts.device == "DSP"
   upcasted_axis: set[int] = set()
-  while resolve(prod(k.output_shape[i] for i in k.upcastable_dims) >= 1024):
+  while resolve(prod(k.sts[0].shape[i] for i in k.upcastable_dims) >= 1024):
     xb_choices = []
     # consider all upcastable axes with 3 or 4 upcast (128 on the DSP)
     for axis, upcast_amount in itertools.product(k.upcastable_dims, ([128] if not len(upcasted_axis) else []) if is_dsp else [3,4]):

tinygrad/codegen/opt/kernel.py

@@ -10,7 +10,7 @@ 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.dtype import ImageDType, AddrSpace
 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
@@ -60,7 +60,7 @@ class Kernel:
 
     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]
+    self.bufs: list[UOp] = [x for x in self.ast.toposort() if x.op in GroupOp.Buffer][::-1]
 
     # create new shapetrackers inside this kernel, we will permute them
     self.sts: list[ShapeTracker] = [x.st_arg for x in self.bufs]
@@ -92,6 +92,10 @@ class Kernel:
     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))
 
@@ -118,7 +122,7 @@ class Kernel:
   @property
   def output_shape(self) -> tuple[sint, ...]: return self.sts[0].shape
   @property
-  def shape_len(self) -> int: return len(self.full_shape)
+  def shape_len(self) -> int: return len(self.sts[0].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
@@ -170,7 +174,7 @@ class Kernel:
   # 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:
+  def shift_to(self, axis:int, amount:int, new_type:AxisType, top:bool=False, insert_at:int|None=None):
     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
@@ -179,7 +183,6 @@ class Kernel:
     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 ********************
 
@@ -241,11 +244,11 @@ class Kernel:
       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=}")
+      check(axis < self.shape_len, "invalid axis")
       return axis
     except IndexError as e: raise KernelOptError from e
 
-  def apply_opt(self, opt:Opt, append_opt:bool=True) -> int|None:
+  def apply_opt(self, opt:Opt, append_opt:bool=True):
     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")
 
@@ -259,7 +262,7 @@ class Kernel:
       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
+      return
 
     axis = self.real_axis(opt.op, opt.axis)
 
@@ -282,30 +285,28 @@ class Kernel:
       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)
+      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)))
+      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)
+      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)
+      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")
@@ -335,7 +336,6 @@ class Kernel:
     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)
@@ -378,9 +378,9 @@ class Kernel:
       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)
+        if tensor_core_opts[0] is None: continue
         self.tensor_core_opts = tc_opts = tensor_core_opts[0]
 
         # attempt to pad the tensor axes that require it
@@ -460,7 +460,8 @@ class Kernel:
       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)
+        axes = tuple(i for i in self.axes_of(AxisType.REDUCE, AxisType.UNROLL) if i in changed)
+        grouped_axes = tuple(i for i in self.axes_of(AxisType.GROUP_REDUCE) 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()))])
@@ -485,6 +486,23 @@ class Kernel:
           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))
+        if self.group_for_reduces and grouped_axes:
+          local_axes = tuple([i for i,t in enumerate(self.axis_types) if t in (AxisType.LOCAL, AxisType.UPCAST) or i in grouped_axes])
+          slocal, supcast, sgroup = sorted(self.axes_of(AxisType.LOCAL)), sorted(self.axes_of(AxisType.UPCAST)), sorted(grouped_axes)
+          # NOTE: start with UPCAST at the end so it has stride 1 and can merge
+          base_shape = tuple([self.full_shape[i] for i in slocal] + [self.full_shape[i] for i in sgroup] + [self.full_shape[i] for i in supcast])
+          permute_axes = tuple([local_axes.index(i) for i in slocal+sgroup+supcast])
+          local_shape = tuple([s if i in local_axes else 1 for i,s in enumerate(self.full_shape)])
+          local_src_shape = tuple([self.full_shape[i] if i in self.axes_of(AxisType.GLOBAL) else s for i,s in enumerate(local_shape)])
+          st = ShapeTracker.from_shape(base_shape).permute(permute_axes).reshape(local_shape).expand(local_src_shape)
+          local_size = st.real_size()
+          local_buffer = UOp(Ops.DEFINE_LOCAL, op.dtype.ptr(local_size, addrspace=AddrSpace.LOCAL), (), f"temp{self.reduceops.index(op)}")
+          local_load = local_buffer.view(st).load(local_buffer.view(st).store(ret))
+          grouped_reduce = UOp(Ops.REDUCE_AXIS, op.dtype, (local_load,), arg=(op.arg[0], grouped_axes))
+          if op is self.reduceops[-1]: return grouped_reduce
+          st = ShapeTracker.from_shape(tuple([1 if i in grouped_axes else s for i,s in enumerate(local_shape)]))
+          return local_buffer.view(st).load(local_buffer.view(st).store(grouped_reduce))
+
       return ret
     self.finalized = True
     fixed_ast = fixup_ast(self.ast)

tinygrad/codegen/opt/postrange.py

@@ -1,183 +0,0 @@
-import math
-from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import UOp, Ops, sint, ssimplify, AxisType, KernelInfo, PatternMatcher, UPat, graph_rewrite, _substitute
-from tinygrad.uop.symbolic import symbolic
-from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
-from tinygrad.renderer import Renderer
-from tinygrad.dtype import dtypes
-
-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 RKernel(Kernel):
-  def __init__(self, ast:UOp, opts:Renderer|None=None):
-    self.rng = sorted([u for u in ast.toposort() if u.op is Ops.RANGE and u.vmax > 0], key=lambda x: x.arg)
-    super().__init__(ast, opts)
-    self.sts.clear()
-
-    # convert LOOP to GLOBAL
-    self.replaces = {}
-    if self.opts.has_local:
-      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].dtype.addrspace == AddrSpace.LOCAL) \
-                          or (x.op is Ops.BUFFERIZE and x.arg == AddrSpace.LOCAL)]
-      for ls in local_store_rngs: store_rngs = [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.rng]
-      self.replaces.update(dict(zip(self.rng, rng)))
-      self.rng = rng
-
-  def simplify_merge_adjacent(self):
-    return
-    # NOTE: this one is better than the one in kernel.py, which is kind of a problem
-    terminators = [u for u in self.ast.toposort() if u.op in {Ops.REDUCE, Ops.STORE}]
-    termination = {}
-    for t in terminators:
-      for u in t.src[1 if t.op is Ops.REDUCE else 2:]: termination[u] = t
-
-    replaces = {}
-    i = 0
-    while i < len(self.rng)-1:
-      r0, r1 = self.rng[i], self.rng[i+1]
-      # same axistype and same termination
-      if r0.arg[1] == r1.arg[1] and termination[r0] == termination[r1]:
-        s0, s1 = r0.src[0], r1.src[0]
-        new_range = r0.replace(src=(s0*s1,)).simplify()
-        # this checks the legality of a merge
-        oidx = self.ast.simplify()
-        nidx = graph_rewrite(oidx, _substitute+symbolic+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1}, name=f"check_merge_{i}_{i+1}")
-        # it simplifies
-        if count_divmod(nidx) <= count_divmod(oidx):
-          # it is correct
-          midx = graph_rewrite(nidx, _substitute+symbolic+pm_flatten_range, ctx={new_range:r0*s1+r1}, name=f"correct_merge_{i}_{i+1}")
-          if oidx is midx:
-            termination[new_range] = termination[r0]
-            replaces[r0] = new_range//s1
-            replaces[r1] = new_range%s1
-            self.rng[i] = new_range
-            del self.rng[i+1]
-            continue
-      i += 1
-    self.ast = self.ast.substitute(replaces, name="simplify_merge_adjacent")
-
-  def shift_to(self, axis:int, amount:int, new_type:AxisType, top:bool=False, insert_at:int|None=None):
-    old_sz = self.rng[axis].src[0].arg // amount
-    assert old_sz > 0, f"bad old_sz on {axis} {amount} {self.rng[axis]}"
-
-    maxarg = max([x.arg[0] for x in self.rng])
-    new_rng = UOp.range(dtypes.int, amount, maxarg+1, new_type)
-
-    if old_sz == 1:
-      self.replaces[self.rng[axis]] = new_rng
-      self.rng.insert(insert_at if insert_at is not None else len(self.rng), new_rng)
-      del self.rng[axis]
-    else:
-      replaced_rng = self.rng[axis].replace(src=(UOp.const(dtypes.int, old_sz),))
-      self.replaces[self.rng[axis]] = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng)
-      self.rng[axis] = replaced_rng
-      self.rng.insert(insert_at if insert_at is not None else len(self.rng), new_rng)
-    return new_rng
-
-  @property
-  def axis_types(self) -> list[AxisType]: return [x.arg[1] for x in self.rng]
-  @property
-  def shape_len(self): return len(self.rng)
-
-  @property
-  def full_shape(self) -> tuple[sint, ...]: return tuple([ssimplify(x.src[0]) for x in self.rng])
-  @property
-  def output_shape(self) -> tuple[sint, ...]: return tuple([ssimplify(x.src[0]) for x in self.ast.src[0].src[2:]])
-
-  def get_optimized_ast(self, name_override:str|None=None) -> UOp:
-    ret = self.ast
-    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
-    rarg = KernelInfo(kernel_name, tuple(self.axis_types), self.dont_use_locals, tuple(self.applied_opts))
-    return ret.substitute(self.replaces).replace(arg=rarg)
-
-  # does nothing
-  @axis_types.setter
-  def axis_types(self, value): pass
-
-  def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> bool:
-    reduceop = [x for x in self.ast.toposort() if x.op is Ops.REDUCE][0]
-    if use_tensor_cores and reduceop is not None and reduceop.arg 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:
-        if tc.dtype_in == dtypes.float and tc.dtype_out == dtypes.float:
-          axes = [1,0]
-
-          # do optimizations and save the ranges
-          ne: list[UOp] = []
-          for opt in tc.opts:
-            ne.append(self.apply_opt(Opt({"u":OptOps.UPCAST, "l":OptOps.LOCAL}[opt[0]], axes[int(opt[1])], 2), append_opt=False))
-          for _, amt in tc.get_reduce_axes():
-            ne.append(self.apply_opt(Opt(OptOps.UNROLL, 0, amt), append_opt=False)) # TODO: this should be the reduce, not 0
-
-          # early realize for TC
-          self.ast = self.ast.substitute(self.replaces)
-          self.replaces = {}
-
-          # 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 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.rng[a].arg[0], sz) for a,sz in v]) for v in tc_upcast_axes])
-          tc_reduce_axes = tuple([self.rng[a].arg[0] for a in tc_reduce_axes])
-
-          # construct the op
-          # TODO: remove tc_upcast_axes from the arg
-          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 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 True
-    return False
-
-from dataclasses import replace
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo
-from tinygrad.helpers import colored
-from tinygrad.codegen.opt.kernel import axis_colors
-
-def rename_sink(s:UOp):
-  if s.arg is not None and s.arg.name != "test": return None
-
-  # get all ranges (sorted)
-  rngs = sorted([u for u in s.parents if u.op is Ops.RANGE], key=lambda x: x.arg[0:-1])
-
-  # add name to kernel
-  name = "k" + colored('_', 'BLACK').join(['']+[colored(x.src[0].render(), axis_colors[x.arg[-1]]) for x in rngs])
-  return s.replace(arg=KernelInfo(name=name) if s.arg is None else replace(s.arg, name=name))
-
-pm_postrange_opt = PatternMatcher([
-  (UPat(Ops.SINK, name="s"), rename_sink),
-])

tinygrad/codegen/opt/swizzler.py

@@ -128,8 +128,7 @@ fix_kernel_ops = view_left_through_load+PatternMatcher([
   (UPat(Ops.VIEW, src=(UPat.cvar(),), name="self"),
    lambda self: UOp.where(UOp(Ops.VALID, dtypes.bool, (UOp(Ops.VIEW, arg=self.st),)), self.const_like(self.base.arg), 0)),
   # no ImageDType after index
-  (UPat(GroupOp.All-{Ops.DEFINE_GLOBAL, Ops.VIEW, Ops.INDEX}, name="x"),
-   lambda x: x.replace(dtype=x.dtype.base) if isinstance(x.dtype, ImageDType) else None),
+  (UPat(GroupOp.All-{Ops.DEFINE_GLOBAL, Ops.VIEW}, name="x"), lambda x: x.replace(dtype=x.dtype.base) if isinstance(x.dtype, ImageDType) else None),
   # if this kernel also assigns to the loaded buffer, ensure we can index it correctly
   (UPat(Ops.LOAD, src=(UPat.var("glbl").view(name="view"),)), check_load_st),
 ])

tinygrad/codegen/opt/tc.py

@@ -22,15 +22,6 @@ class TensorCore: # D = A * B + C, A is (M x K), B is (K x N), C and D are (M x
   def permutes_for_shape_str(self, shape_str:list[str]) -> tuple[tuple[int, ...], tuple[int, ...]]:
     ret = [[shape_str.index(remap[ss]) if ss in remap else i for i,ss in enumerate(shape_str)] for remap in self._remaps()]
     return tuple(ret[0]), tuple(ret[1])
-  @functools.cache  # pylint: disable=method-cache-max-size-none
-  def base_shape_str(self) -> list[str]:
-    ret = []
-    cnt = {'u': 0, 'l': 0}
-    for opt in self.opts:
-      ret.append(f"{opt[0]}{cnt[opt[0]]}")
-      cnt[opt[0]] += 1
-    # assumes you do the UNROLL after the opts
-    return ret + [f"r{i}" for i in range(len(self.get_reduce_axes()))]
   def get_reduce_axes(self): return [(i, 2) for i in range(int(math.log2(self.dims[2])))]
   def get_upcast_axes(self): return [opt for opt in self.opts if opt[0] == "u"]
   def get_local_axes(self): return [opt for opt in self.opts if opt[0] == "l"]

tinygrad/device.py

@@ -3,8 +3,8 @@ 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
-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.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, \
+                             colored, Context, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, cpu_events, ProfileEvent, ProfilePointEvent, dedup
 from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
 from tinygrad.renderer import Renderer
 
@@ -124,7 +124,7 @@ class Buffer:
   def allocate(self, opaque=None, external_ptr=None) -> Buffer:
     assert not self.is_initialized(), "can't allocate already allocated buffer"
     if DEBUG >= 7: print(f"buffer: allocate {self.nbytes} bytes on {self.device}")
-    if MAX_BUFFER_SIZE > 0 and self.size > MAX_BUFFER_SIZE: raise RuntimeError(f"buffer of size {self.size/1e6:.2f}M is too large")
+    if (mbs:=getenv("MAX_BUFFER_SIZE", 0)) > 0 and self.size > mbs: raise RuntimeError(f"buffer of size {self.size/1e6:.2f}M is too large")
     self.allocator:Allocator = Device[self.device].allocator
     if external_ptr is not None:
       self.options = replace(self.options, external_ptr=external_ptr) if self.options else BufferSpec(external_ptr=external_ptr)
@@ -139,7 +139,7 @@ class Buffer:
       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", ts, num, {"dtype":str(self.dtype),"sz":self.size,"nbytes":self.nbytes}))
     return self
   def deallocate(self):
     assert hasattr(self, '_buf'), "buffer must be allocated to deallocate"

tinygrad/dtype.py

@@ -108,6 +108,7 @@ class dtypes:
     if isinstance(val, tuple):
       assert len(val) == dtype.count, f"mismatch {val} {dtype}"
       return tuple(dtypes.as_const(x, dtype) for x in val)
+    # TODO: should truncate here
     return int(val) if dtypes.is_int(dtype) else float(val) if dtypes.is_float(dtype) else bool(val)
   @staticmethod
   @functools.cache
@@ -214,14 +215,15 @@ def sum_acc_dtype(dt:DType):
   return least_upper_dtype(dt, to_dtype(getenv("SUM_DTYPE", "float32")))
 
 def truncate_fp16(x):
-  try: return struct.unpack('e', struct.pack('e', float(x)))[0]
+  try: return struct.unpack("@e", struct.pack("@e", float(x)))[0]
   except OverflowError: return math.copysign(math.inf, x)
 
-def float_to_bf16(x):
-  if not math.isfinite(x): return x
-  u = struct.unpack('I', struct.pack('f', x))[0]
-  u = (u + 0x7FFF + ((u >> 16) & 1)) & 0xFFFF0000
-  return struct.unpack('f', struct.pack('I', u))[0]
+def truncate_bf16(x):
+  max_bf16 = struct.unpack('f', struct.pack('I', 0x7f7f0000))[0]
+  if abs(x) > max_bf16: return math.copysign(math.inf, x)
+  f32_int = struct.unpack('I', struct.pack('f', x))[0]
+  bf = struct.unpack('f', struct.pack('I', f32_int & 0xFFFF0000))[0]
+  return bf
 
 # fp8-float conversions based on https://gitlab.com/nvidia/headers/cuda-individual/cudart/-/blob/main/cuda_fp8.hpp
 def float_to_fp8(x: float, dtype: DType) -> int:
@@ -286,7 +288,7 @@ def fp8_to_float(x: int, dtype: DType) -> float:
   return float(float32_val)
 
 truncate: dict[DType, Callable] = {dtypes.bool: bool,
-  dtypes.float16: truncate_fp16, dtypes.bfloat16: lambda x: float_to_bf16(float(x)),
+  dtypes.float16: truncate_fp16, dtypes.bfloat16: truncate_bf16,
   **{fp8: (lambda x, dtype=fp8: fp8_to_float(float_to_fp8(x, dtype), dtype)) for fp8 in dtypes.fp8s},
   dtypes.float32: lambda x: ctypes.c_float(x).value, dtypes.float64: lambda x: ctypes.c_double(x).value,
   dtypes.uint8: lambda x: ctypes.c_uint8(x).value, dtypes.uint16: lambda x: ctypes.c_uint16(x).value,

tinygrad/engine/memory.py

@@ -23,13 +23,12 @@ def _internal_memory_planner(buffers:list[list[Buffer]], noopt_buffers=None, ign
   # Sort buffer operations in timeline order. Two events: buffer is allocated or buffer is freed.
   buffer_requests = sorted([((first_appearance[buf], True), buf) for buf in first_appearance.keys()] + \
                            [((last_appearance[buf] + 1, False), buf) for buf in first_appearance.keys()], key=lambda x: x[0])
-  total_memory = sum(round_up(buf.nbytes, min_block_size:=0x1000) for buf in first_appearance.keys()) * 2 # *2 for fragmentation (which is about 15%)
 
   # Try to suballocate from a shared buffer managed by global_planner using TLSFAllocator.
   # Also track buffer replacements for buffers that do not support suballocation.
   buffer_replace:dict[Buffer, tuple[Buffer|None, int|None]] = {}
   reuse_buffers:dict[tuple, list[Buffer]] = defaultdict(list)
-  global_planner:dict[str, tuple[int, TLSFAllocator]] = defaultdict(lambda: (0, TLSFAllocator(total_memory, block_size=min_block_size, lv2_cnt=32)))
+  global_planner:dict[str, tuple[int, TLSFAllocator]] = defaultdict(lambda: (0, TLSFAllocator(1 << 44, block_size=0x1000, lv2_cnt=32)))
   for (_, is_open_ev), buf in buffer_requests:
     # Check if suballocation is possible for the given buffer and device.
     if hasattr(Device[buf.device].allocator, "_offset") and not isinstance(buf.dtype, ImageDType):

tinygrad/engine/realize.py

@@ -1,8 +1,8 @@
 from typing import cast, Generator
-import time, pprint, decimal
+import time, pprint
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA, TracingKey
-from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, cpu_profile, PROFILE, ProfilePointEvent, cpu_events
+from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, cpu_profile
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, Variable, sym_infer, graph_rewrite, print_uops, track_rewrites, KernelInfo
 from tinygrad.device import Device, Buffer
 from tinygrad.renderer import Renderer, ProgramSpec, Estimates
@@ -149,8 +149,6 @@ class ExecItem:
   def run(self, _var_vals:dict[Variable, int]|None=None, wait=False, jit=False, do_update_stats=True) -> float|None:
     var_vals = self.fixedvars if _var_vals is None else (_var_vals|self.fixedvars)
     bufs = [cast(Buffer, x) for x in self.bufs] if jit else [cast(Buffer, x).ensure_allocated() for x in self.bufs]
-    if PROFILE: cpu_events.append(ProfilePointEvent(self.prg.device, "exec", decimal.Decimal(time.perf_counter_ns())/1000, self.prg.display_name,
-                                                    {"metadata":self.metadata, "var_vals":var_vals}))
     et = self.prg(bufs, var_vals, wait=wait or DEBUG >= 2)
     if do_update_stats:
       GlobalCounters.kernel_count += 1
@@ -160,15 +158,10 @@ class ExecItem:
       if DEBUG >= 2:
         lds_est = sym_infer(self.prg.estimates.lds, var_vals)
         mem_est = min(mem_est, lds_est)   # there can't be more memory accessed than loads/stores. remove this when symbolic is fixed
-        header_color = 'magenta' if jit else ('green' if self.prg.first_run else None)
         ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
-        flops, membw, ldsbw = op_est/(et or 1e-20), mem_est/(et or 1e-20), lds_est/(et or 1e-20)
-        flops_str = f"{flops*1e-9:9.2f} GFLOPS" if flops < 1e14 else colored(f"{flops*1e-12:9.2f} TFLOPS", 'green')
-        mem_str = f"{membw*1e-9:6.1f}|{ldsbw*1e-9:<7.1f} GB/s" if membw < 1e13 else colored(f"{membw*1e-12:6.1f}|{ldsbw*1e-12:<7.1f} TB/s", 'green')
-        print(f"{colored(f'*** {self.prg.device[:7]:7s} {GlobalCounters.kernel_count:4d}', header_color)}"+
-          f" {self.prg.display_name+' '*(44-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:5.2f} GB"+
-          ("" if et is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})")+
-          f" {[repr(m) if TRACEMETA >= 2 else str(m) for m in self.metadata] if self.metadata else ''}")
+        print(f"{colored(f'*** {self.prg.device[:7]:7s} {GlobalCounters.kernel_count:4d}', 'magenta' if jit else ('green' if self.prg.first_run else None))} {self.prg.display_name+' '*(44-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:5.2f} GB " +  # noqa: E501
+              (str() if et is None else f"tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({op_est/((et or 1e-20)*1e9):9.2f} GFLOPS {mem_est/((et or 1e-20)*1e9):6.1f}|{lds_est/((et or 1e-20)*1e9):<7.1f} GB/s)" +  # noqa: E501
+               f" {[repr(m) if TRACEMETA >= 2 else str(m) for m in self.metadata] if self.metadata else ''}"))
       self.prg.first_run = False
     return et
 

tinygrad/engine/schedule.py

@@ -33,7 +33,7 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
         for ss in s.src:
           if ss.op is Ops.MSELECT: ss = ss.src[0]
           if ss.op is not Ops.BUFFER:
-            assert ss.op is Ops.ASSIGN, f"ss.op is not ASSIGN, it's {ss.op}"
+            assert ss.op is Ops.ASSIGN
             children[ss.src[1]].append(k)
             in_degree[k] += 1
       elif s.op is Ops.BUFFER: