notes

* tc3 shape expand [pr]

* remove unused stuff in lowerer

.github/workflows/benchmark.yml

@@ -67,11 +67,11 @@ jobs:
     - name: Test speed vs torch
       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_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestKernelOpts.test_tensor_core_opts
+      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
       run: |
-        DEBUG=2 CPU=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
-        DEBUG=2 LLVM=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        DEBUG=2 CPU=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        DEBUG=2 LLVM=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
     - name: Run Tensor Core GEMM (float)
       run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
     - name: Run Tensor Core GEMM (half)
@@ -123,7 +123,7 @@ jobs:
     - name: UsbGPU copy speeds
       run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
     - name: UsbGPU openpilot test
-      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB AMD_LLVM=1 NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
     - uses: actions/upload-artifact@v4
       with:
         name: Speed (Mac)
@@ -196,8 +196,8 @@ jobs:
       run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
     - name: Test tensor cores
       run: |
-        NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestKernelOpts.test_tensor_core_opts
-        PTX=1 ALLOW_TF32=1 NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestKernelOpts.test_tensor_core_opts
+        NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        PTX=1 ALLOW_TF32=1 NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
     - name: Run Tensor Core GEMM (CUDA)
       run: |
         CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
@@ -396,8 +396,8 @@ jobs:
       run: AMD=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
     - name: Test tensor cores
       run: |
-        AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops TestKernelOpts.test_tensor_core_opts
-        AMD=1 AMD_LLVM=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops TestKernelOpts.test_tensor_core_opts
+        AMD=1 AMD_LLVM=0 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+        AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
         AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
     - name: Run Tensor Core GEMM (AMD)
       run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt

.github/workflows/benchmark_search.yml

@@ -27,7 +27,7 @@ jobs:
         BENCHMARK_LOG=search_sdxl_cached PYTHONPATH=. AMD=1 JITBEAM=2 python examples/sdxl.py --noshow --timing --seed 0
     - name: Run winograd cifar with new search
       run: |
-        BENCHMARK_LOG=search_wino_cifar WINO=1 DEFAULT_FLOAT=HALF FUSE_ARANGE=1 JITBEAM=4 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 BS=1024 STEPS=500 python examples/hlb_cifar10.py
+        BENCHMARK_LOG=search_wino_cifar WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 BS=1024 STEPS=500 python examples/hlb_cifar10.py
     - name: Run winograd cifar with cached search
       run: |
-        BENCHMARK_LOG=search_wino_cifar_cached WINO=1 DEFAULT_FLOAT=HALF FUSE_ARANGE=1 JITBEAM=4 BS=1024 STEPS=500 python examples/hlb_cifar10.py
+        BENCHMARK_LOG=search_wino_cifar_cached WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 BS=1024 STEPS=500 python examples/hlb_cifar10.py

.github/workflows/test.yml

@@ -272,14 +272,6 @@ jobs:
         PYTHONPATH=. DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
         PYTHONPATH=. DEBUG=2 EMULATE_INTEL=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
         PYTHONPATH=. DEBUG=2 AMX=1 EMULATE_AMX=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
-    - name: Test tensor cores (TC=3)
-      run: |
-        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_METAL=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
-        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_AMD=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
-        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_AMD_MFMA=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
-        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_CUDA=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
-        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_INTEL=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
-        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_AMX=1 AMX=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
     - name: Test device flop counts
       run: |
         PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
@@ -310,8 +302,6 @@ jobs:
       run: PYTHON=1 python3 -m pytest test/test_uops.py --durations=20
     - name: Test symbolic with Python emulator
       run: PYTHONPATH=. PYTHON=1 python3 test/test_symbolic_ops.py
-    - name: test_linearizer_failures with Python emulator
-      run: PYTHONPATH=. PYTHON=1 python3 -m pytest -rA test/test_linearizer_failures.py::TestLinearizerFailures::test_failure_1
     - name: test_renderer_failures with Python emulator
       run: PYTHONPATH=. PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures
 
@@ -336,7 +326,7 @@ jobs:
       run: |
         pip3 install --upgrade --force-reinstall ruff==0.11.0
         python3 -m ruff check .
-        python3 -m ruff check examples/mlperf/model_train.py --ignore E501
+        python3 -m ruff check examples/mlperf/ --ignore E501
     - name: Lint tinygrad with pylint
       run: python -m pylint tinygrad/
     - name: Run mypy
@@ -371,8 +361,15 @@ jobs:
       run: PYTHONPATH="." python test/external/external_uop_gc.py
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
-    - name: Repo line count < 14000 lines
-      run: MAX_LINE_COUNT=14000 python sz.py
+    - name: Regen dataset on test_tiny
+      run: |
+        test/external/process_replay/reset.py
+        CAPTURE_PROCESS_REPLAY=1 python test/test_tiny.py TestTiny.test_plus
+        PYTHONPATH=. python extra/optimization/extract_dataset.py
+        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 < 14600 lines
+      run: MAX_LINE_COUNT=14600 python sz.py
 
   fuzzing:
     name: Fuzzing
@@ -476,6 +473,10 @@ jobs:
         run: CPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
       - name: Test ONNX (LLVM)
         run: LLVM=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+      - name: Test ONNX Runner (CPU)
+        run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py
+      - name: Test ONNX Runner (WEBGPU)
+        run: WEBGPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py
       - name: Test Additional ONNX Ops (CPU)
         run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_ops.py
       - name: Test Quantize ONNX
@@ -512,8 +513,8 @@ jobs:
           REMOTEDEV=GPU IMAGE=2 REMOTE=1 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
       - name: Test Optimization Helpers
         run: PYTHONPATH="." DEBUG=1 python3 extra/optimization/test_helpers.py
-      - name: Test Action Space
-        run: PYTHONPATH="." DEBUG=1 GPU=1 python3 extra/optimization/get_action_space.py
+      #- name: Test Action Space
+      #  run: PYTHONPATH="." DEBUG=1 GPU=1 python3 extra/optimization/get_action_space.py
       - name: Test Beam Search
         run: PYTHONPATH="." GPU=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
       - name: Test MLPerf stuff
@@ -612,7 +613,7 @@ jobs:
       run: |
           WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit \
           --ignore=test/test_copy_speed.py --ignore=test/test_rearrange_einops.py \
-          --ignore=test/test_fuzz_shape_ops.py --ignore=test/test_linearizer_failures.py --durations=20
+          --ignore=test/test_fuzz_shape_ops.py --durations=20
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
@@ -763,8 +764,8 @@ jobs:
       run: PYTHONPATH="." METAL=1 python test/external/external_test_speed_llama.py
     - name: Test Beam Search
       run: PYTHONPATH="." METAL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
-    - name: Fuzz Test linearizer
-      run: PYTHONPATH="." METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
+    #- name: Fuzz Test linearizer
+    #  run: PYTHONPATH="." METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
     - name: Run TRANSCENDENTAL math
       run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
     - name: Run pytest (amd)
@@ -778,7 +779,6 @@ jobs:
       env:
         MOCKGPU: 1
         AMD: 1
-        AMD_LLVM: 1
         FORWARD_ONLY: 1
       run: |
         python -m pytest -n=auto test/test_hcq.py test/test_tiny.py test/test_amd_llvm.py --durations=20
@@ -814,8 +814,9 @@ jobs:
       run: npm cache clean --force
     - name: Install Puppeteer
       run: npm install puppeteer
-    - name: Run WEBGPU Efficientnet
-      run: node test/web/test_webgpu.js
+    # this is also flaky
+    #- name: Run WEBGPU Efficientnet
+    #  run: node test/web/test_webgpu.js
     # this is flaky
     #- name: Run VIZ tests as external package
     #  run: |

autogen_stubs.sh

@@ -118,7 +118,9 @@ generate_nv() {
   clang2py -k cdefstum \
     extra/nv_gpu_driver/clc6c0qmd.h \
     extra/nv_gpu_driver/clcec0qmd.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl0000.h \
     $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl0080.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl2080.h \
     $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl2080_notification.h \
     $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc56f.h \
     $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc86f.h \
@@ -165,6 +167,26 @@ generate_nv() {
 nv_status_codes = {}
 /^NV_STATUS_CODE/ { s/^NV_STATUS_CODE(\([^,]*\), *\([^,]*\), *"\([^"]*\)") *.*$/\1 = \2\nnv_status_codes[\1] = "\3"/; p }' $NVKERN_SRC/src/common/sdk/nvidia/inc/nvstatuscodes.h >> $BASE/nv_gpu.py
 
+  clang2py -k cdefstum \
+    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/fsp/kern_fsp_cot_payload.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gspifpub.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gsp_fw_wpr_meta.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gsp_fw_sr_meta.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/gsp/gsp_init_args.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/gsp/gsp_init_args.h \
+    $NVKERN_SRC/src/common/uproc/os/common/include/libos_init_args.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/rmRiscvUcode.h \
+    $NVKERN_SRC/src/common/shared/msgq/inc/msgq/msgq_priv.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/vgpu/rpc_headers.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/vgpu/rpc_global_enums.h \
+    $NVKERN_SRC/src/nvidia/generated/g_rpc-structures.h \
+    extra/nv_gpu_driver/g_rpc-message-header.h \
+    extra/nv_gpu_driver/gsp_static_config.h \
+    extra/nv_gpu_driver/vbios.h \
+    --clang-args="-DRPC_MESSAGE_STRUCTURES -DRPC_STRUCTURES -include $NVKERN_SRC/src/common/sdk/nvidia/inc/nvtypes.h -I$NVKERN_SRC/src/nvidia/generated -I$NVKERN_SRC/src/common/inc -I$NVKERN_SRC/src/nvidia/inc -I$NVKERN_SRC/src/nvidia/interface/ -I$NVKERN_SRC/src/nvidia/inc/kernel -I$NVKERN_SRC/src/nvidia/inc/libraries -I$NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc -I$NVKERN_SRC/kernel-open/nvidia-uvm -I$NVKERN_SRC/kernel-open/common/inc -I$NVKERN_SRC/src/common/sdk/nvidia/inc -I$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include -I$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl" \
+    -o $BASE/nv/nv.py
+
+  fixup $BASE/nv/nv.py
   python3 -c "import tinygrad.runtime.autogen.nv_gpu"
 }
 
@@ -436,6 +458,7 @@ elif [ "$1" == "kfd" ]; then generate_kfd
 elif [ "$1" == "nv" ]; then generate_nv
 elif [ "$1" == "amd" ]; then generate_amd
 elif [ "$1" == "am" ]; then generate_am
+elif [ "$1" == "nvdrv" ]; then generate_nvdrv
 elif [ "$1" == "sqtt" ]; then generate_sqtt
 elif [ "$1" == "qcom" ]; then generate_qcom
 elif [ "$1" == "io_uring" ]; then generate_io_uring

docs/abstractions2.py

@@ -60,7 +60,7 @@ s = UOp(Ops.SINK, dtypes.void, (st_0,))
 
 # convert the computation to a "linearized" format (print the format)
 from tinygrad.engine.realize import get_program, CompiledRunner
-program = get_program(Device[DEVICE].renderer, s)
+program = get_program(s, Device[DEVICE].renderer)
 
 # compile a program (and print the source)
 fxn = CompiledRunner(program)
@@ -78,7 +78,7 @@ print("******** third, the UOp ***********")
 
 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.schedule import create_schedule_with_vars
-from tinygrad.engine.kernelize import get_kernelize_map
+from tinygrad.kernelize.kernelize import get_kernelize_map
 
 # allocate some values + load in values
 a = UOp.new_buffer(DEVICE, 1, dtypes.int32)

docs/developer/layout.md

@@ -0,0 +1,66 @@
+# tinygrad directory layout
+
+This explains the flow of a big graph down to programs.
+
+Directories are listed in order of how they are processed.
+
+---
+
+## tinygrad/kernelize
+
+Group UOps into kernels.
+
+::: tinygrad.kernelize.kernelize.get_kernelize_map
+    options:
+        members: false
+        show_labels: false
+        show_source: false
+
+---
+
+## tinygrad/opt
+
+Transforms the ast into an optimized ast. This is where BEAM search and heuristics live.
+
+::: tinygrad.opt.get_optimized_ast
+    options:
+        members: false
+        show_labels: false
+        show_source: false
+
+---
+
+## tinygrad/codegen
+
+Transform the optimized ast into a linearized list of UOps.
+
+::: tinygrad.codegen.full_rewrite
+    options:
+        members: false
+        show_labels: false
+        show_source: false
+
+---
+
+## tinygrad/renderer
+
+Transform the linearized list of UOps into a program, represented as a string.
+
+::: tinygrad.renderer.Renderer
+    options:
+        members:
+            - render
+        show_labels: false
+        show_source: false
+
+---
+
+## tinygrad/engine
+
+Abstracted high level interface to the runtimes.
+
+::: tinygrad.engine.realize.get_program
+    options:
+        members: false
+        show_labels: false
+        show_source: false

docs/tensor/elementwise.md

@@ -35,6 +35,7 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 
 ::: tinygrad.Tensor.relu
 ::: tinygrad.Tensor.sigmoid
+::: tinygrad.Tensor.logsigmoid
 ::: tinygrad.Tensor.hardsigmoid
 ::: tinygrad.Tensor.elu
 ::: tinygrad.Tensor.celu

examples/benchmark_onnx.py

@@ -1,4 +1,4 @@
-import sys, time, pickle
+import sys, time
 from tinygrad import TinyJit, GlobalCounters, fetch, getenv
 from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 from extra.onnx_helpers import get_example_inputs, validate

examples/coder.py

@@ -4,7 +4,7 @@ sys.path.append(os.getcwd())
 
 from io import StringIO
 from contextlib import redirect_stdout
-from tinygrad import Tensor, nn, Device, dtypes
+from tinygrad import Tensor, nn
 from tinygrad.helpers import Timing, colored, getenv, fetch
 from extra.models.llama import Transformer, convert_from_huggingface, fix_bf16
 from sentencepiece import SentencePieceProcessor

examples/hlb_cifar10.py

@@ -7,8 +7,8 @@ import random, time
 import numpy as np
 from typing import Optional
 from extra.lr_scheduler import OneCycleLR
-from tinygrad import nn, dtypes, Tensor, Device, GlobalCounters, TinyJit
-from tinygrad.nn.state import get_state_dict, get_parameters
+from tinygrad import nn, dtypes, Tensor, Device, GlobalCounters, TinyJit, Variable
+from tinygrad.nn.state import get_state_dict
 from tinygrad.nn import optim
 from tinygrad.helpers import Context, BEAM, WINO, getenv, colored, prod
 from extra.bench_log import BenchEvent, WallTimeEvent
@@ -145,6 +145,7 @@ hyp = {
   },
 }
 
+@Context(FUSE_ARANGE=getenv("FUSE_ARANGE", 1))
 def train_cifar():
 
   def set_seed(seed):
@@ -201,24 +202,37 @@ def train_cifar():
     idx_y = Tensor.arange(H, dtype=dtypes.int32).reshape((1,1,H,1))
     return (idx_x >= low_x) * (idx_x < (low_x + mask_size)) * (idx_y >= low_y) * (idx_y < (low_y + mask_size))
 
-  def random_crop(X:Tensor, crop_size=32):
-    mask = make_square_mask(X.shape, crop_size)
-    mask = mask.expand((-1,3,-1,-1))
-    X_cropped = Tensor(X.numpy()[mask.numpy()])
-    return X_cropped.reshape((-1, 3, crop_size, crop_size))
+  # Similar, but different enough.
+  def make_random_crop_indices(shape, mask_size) -> Tensor:
+    BS, _, H, W = shape
+    low_x = Tensor.randint(BS, low=0, high=W-mask_size).reshape(BS,1,1,1)
+    low_y = Tensor.randint(BS, low=0, high=H-mask_size).reshape(BS,1,1,1)
+    idx_x = Tensor.arange(mask_size, dtype=dtypes.int32).reshape((1,1,1,mask_size))
+    idx_y = Tensor.arange(mask_size, dtype=dtypes.int32).reshape((1,1,mask_size,1))
+    return low_x, low_y, idx_x, idx_y
 
-  def cutmix(X:Tensor, Y:Tensor, mask_size=3):
-    # fill the square with randomly selected images from the same batch
+  def random_crop(X:Tensor, crop_size=32):
+    Xs, Ys, Xi, Yi = make_random_crop_indices(X.shape, crop_size)
+    return X.gather(-1, (Xs + Xi).expand(-1, 3, X.shape[2], -1)).gather(-2, ((Ys+Yi).expand(-1, 3, crop_size, crop_size)))
+
+  def cutmix(X, Y, order, mask_size=3):
     mask = make_square_mask(X.shape, mask_size)
-    order = list(range(0, X.shape[0]))
-    random.shuffle(order)
-    X_patch = Tensor(X.numpy()[order], device=X.device, dtype=X.dtype)
-    Y_patch = Tensor(Y.numpy()[order], device=Y.device, dtype=Y.dtype)
+    X_patch, Y_patch = X[order], Y[order]
     X_cutmix = mask.where(X_patch, X)
     mix_portion = float(mask_size**2)/(X.shape[-2]*X.shape[-1])
     Y_cutmix = mix_portion * Y_patch + (1. - mix_portion) * Y
     return X_cutmix, Y_cutmix
 
+  @TinyJit
+  def augmentations(X:Tensor, Y:Tensor):
+    perms = Tensor.randperm(X.shape[0], device=X.device) # We reuse perms for cutmix, because they are expensivne to generate
+    if getenv("RANDOM_CROP", 1):
+      X = random_crop(X, crop_size=32)
+    if getenv("RANDOM_FLIP", 1):
+      X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X) # flip LR
+    X, Y = X[perms], Y[perms]
+    return X, Y, *cutmix(X, Y, perms, mask_size=hyp['net']['cutmix_size'])
+
   # the operations that remain inside batch fetcher is the ones that involves random operations
   def fetch_batches(X_in:Tensor, Y_in:Tensor, BS:int, is_train:bool):
     step, epoch = 0, 0
@@ -226,28 +240,16 @@ def train_cifar():
       st = time.monotonic()
       X, Y = X_in, Y_in
       if is_train:
-        # TODO: these are not jitted
-        if getenv("RANDOM_CROP", 1):
-          X = random_crop(X, crop_size=32)
-        if getenv("RANDOM_FLIP", 1):
-          X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X) # flip LR
-        if getenv("CUTMIX", 1):
-          if step >= hyp['net']['cutmix_steps']:
-            X, Y = cutmix(X, Y, mask_size=hyp['net']['cutmix_size'])
-        order = list(range(0, X.shape[0]))
-        random.shuffle(order)
-        X, Y = X.numpy()[order], Y.numpy()[order]
-      else:
-        X, Y = X.numpy(), Y.numpy()
+        X, Y, X_cm, Y_cm = augmentations(X, Y)
+        if getenv("CUTMIX", 1) and step >= hyp['net']['cutmix_steps']: X, Y = X_cm, Y_cm
       et = time.monotonic()
       print(f"shuffling {'training' if is_train else 'test'} dataset in {(et-st)*1e3:.2f} ms ({epoch=})")
-      for i in range(0, X.shape[0], BS):
-        # pad the last batch  # TODO: not correct for test
-        batch_end = min(i+BS, Y.shape[0])
-        x = Tensor(X[batch_end-BS:batch_end], device=X_in.device, dtype=X_in.dtype)
-        y = Tensor(Y[batch_end-BS:batch_end], device=Y_in.device, dtype=Y_in.dtype)
+
+      vi = Variable("i", 0, (full_batches := (X.shape[0] // BS) * BS) - BS)
+      for i in range(0, full_batches, BS):
         step += 1
-        yield x, y
+        vib = vi.bind(i)
+        yield X[vib:vib+BS], Y[vib:vib+BS]
       epoch += 1
       if not is_train: break
 

examples/llama3.py

@@ -248,7 +248,7 @@ if __name__ == "__main__":
   parser.add_argument("--port", type=int, default=7776, help="Web server port")
   parser.add_argument("--debug", action="store_true", help="Enable debug mode")
   parser.add_argument("--seed", type=int, help="Random seed")
-  parser.add_argument("--temperature", type=int, default=0.85, help="Temperature")
+  parser.add_argument("--temperature", type=float, default=0.85, help="Temperature")
   parser.add_argument("--benchmark", action="store_true", help="Run a benchmark")
   parser.add_argument("--timing", action="store_true", help="Print timing per token")
   parser.add_argument("--profile", action="store_true", help="Output profile data")

examples/llm.c/export.py

@@ -1,11 +1,11 @@
 #!/usr/bin/env python3
 import os
 if "NOOPT" not in os.environ: os.environ["NOOPT"] = "1"
-from tinygrad import Device, nn, Tensor, dtypes, Variable
+from tinygrad import Device, nn, Tensor, dtypes
 Device.DEFAULT = "CPU"
 from train_gpt2 import GPT, GPTConfig
-from tinygrad.helpers import dedup, to_function_name, flatten, getenv, GlobalCounters, ansilen, to_function_name
-from tinygrad.engine.realize import get_kernel, run_schedule
+from tinygrad.helpers import dedup, flatten, getenv, GlobalCounters, to_function_name
+from tinygrad.engine.realize import get_kernel
 from tinygrad.engine.memory import memory_planner
 from tinygrad.uop.ops import Ops
 

examples/mlperf/helpers.py

@@ -212,7 +212,7 @@ def get_mlperf_bert_model():
   from examples.mlperf.initializers import LinearBert, EmbeddingBert, LayerNormBert
 
   bert.Linear = LinearBert
-  bert.Embedding = EmbeddingBert 
+  bert.Embedding = EmbeddingBert
   bert.LayerNorm = LayerNormBert
 
   from extra.models.bert import BertForPretraining

examples/mlperf/initializers.py

@@ -39,7 +39,7 @@ class LinearBert(nn.Linear):
   def __init__(self, in_features, out_features, bias=True, std=0.02):
     self.weight = std * rand_truncn(out_features, in_features, dtype=dtypes.float32)
     self.bias = Tensor.zeros(out_features, dtype=dtypes.float32) if bias else None
-  
+
   def __call__(self, x:Tensor):
     return x.cast(dtypes.default_float).linear(self.weight.cast(dtypes.default_float).transpose(), self.bias.cast(dtypes.default_float) if self.bias is not None else None)
 

examples/mlperf/lr_schedulers.py

@@ -1,4 +1,5 @@
-from tinygrad import Tensor, dtypes
+import math
+from tinygrad import dtypes
 from tinygrad.nn.optim import Optimizer
 
 from extra.lr_scheduler import LR_Scheduler
@@ -20,3 +21,19 @@ class PolynomialDecayWithWarmup(LR_Scheduler):
     warmup_lr = (self.epoch_counter * (1.0 / self.warmup)) * self.initial_lr
     x = (1 - (self.epoch_counter - self.warmup) / (self.epochs - self.warmup + 1))
     return (self.epoch_counter <= self.warmup).where(warmup_lr, (self.initial_lr - self.end_lr) * x ** self.power + self.end_lr).cast(self.optimizer.lr.dtype)
+
+class CosineAnnealingLRWithWarmup(LR_Scheduler):
+  def __init__(self, optimizer:Optimizer, base_lr, end_lr, warmup_steps:int, decay_steps:int):
+    assert warmup_steps > 0 and decay_steps > 0
+    super().__init__(optimizer)
+    self.base_lr = base_lr
+    self.end_lr = end_lr
+    self.warmup_steps = warmup_steps
+    self.decay_steps = decay_steps
+    # set lr for first warmup step
+    self.optimizer.lr.assign(self.get_lr()).realize()
+
+  def get_lr(self):
+    warmup_lr = ((self.epoch_counter+1) / self.warmup_steps) * self.base_lr
+    decay_lr = self.end_lr + 0.5 * (self.base_lr-self.end_lr) * (1 + (((self.epoch_counter+1-self.warmup_steps)/self.decay_steps) * math.pi).cos())
+    return (self.epoch_counter < self.warmup_steps).where(warmup_lr, decay_lr).cast(self.optimizer.lr.dtype)

examples/mlperf/metrics.py

@@ -1,6 +1,6 @@
-import re
-import string
+import re, string
 from collections import Counter
+from tinygrad import Tensor
 
 def levenshtein(a, b):
   n, m = len(a), len(b)
@@ -59,3 +59,11 @@ def f1_score(x, y):
   p = ns / len(xt)
   r = ns / len(yt)
   return 2 * p * r / (p + r)
+
+def log_perplexity(logit:Tensor, target:Tensor, ignore_index:int|None=None):
+  # logit has shape (n_samples, seq_len, vocab_size), target has shape (n_samples, seq_len)
+  assert logit.ndim == 3, logit.ndim
+  assert target.ndim == 2, target.ndim
+  assert logit.shape[:2] == target.shape, f"{logit.shape[:2]=}, {target.shape=}"
+  log_prob = logit.log_softmax(axis=-1)
+  return log_prob.transpose(1, 2).nll_loss(target, ignore_index=ignore_index)

examples/mlperf/model_train.py

@@ -5,7 +5,7 @@ import multiprocessing
 from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
 from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, FUSE_CONV_BW, Profiling
 from tinygrad.nn.state import get_parameters, get_state_dict, safe_load, safe_save
-from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam
+from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW
 
 from extra.lr_scheduler import LRSchedulerGroup
 from examples.mlperf.helpers import get_training_state, load_training_state
@@ -933,7 +933,7 @@ def train_step_bert(model, optimizer, scheduler, loss_scaler:float, GPUS, grad_a
     # TODO: OOM without this realize with large grad_acc
     Tensor.realize(*[p.grad for p in optimizer.params])
 
-  global_norm = Tensor([0.0], dtype=dtypes.float32, device=optimizer[0].device)
+  global_norm = Tensor(0.0, dtype=dtypes.float32, device=optimizer[0].device)
   for p in optimizer.params:
     p.grad = p.grad / loss_scaler
     global_norm += p.grad.float().square().sum()
@@ -1284,6 +1284,76 @@ def train_bert():
         MLLOGGER.start(key=mllog_constants.BLOCK_START, value=None, metadata={"first_epoch_num": 1, "epoch_num": 1, "epoch_count": 1, "samples_count": i * GBS, "step_num": i, "first_step_num": i+1})
         previous_step = i
 
+def train_llama3():
+  from extra.models.llama import Transformer
+  from examples.llama3 import MODEL_PARAMS
+  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
+
+  config = {}
+  BS                 = config["BS"]                     = getenv("BS", 4)
+  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
+  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
+
+  opt_adamw_beta_1 = 0.9
+  opt_adamw_beta_2 = 0.95
+  opt_adamw_epsilon = 1e-5
+  opt_adamw_weight_decay = 0.1
+
+  opt_gradient_clip_norm = 1.0
+  sequence_length = 8192
+  opt_learning_rate_warmup_steps = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
+  opt_learning_rate_decay_steps = getenv("DECAY_STEPS", math.ceil(1_200_000 * 1152 / GBS) - opt_learning_rate_warmup_steps)
+  opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152)  # NOTE: cannot change for benchmark
+  opt_end_learning_rate = 8e-7
+
+  # TODO: confirm weights are in bf16
+  # vocab_size from the mixtral tokenizer
+  model = Transformer(**(MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}), max_context=sequence_length, jit=False, disable_kv_cache=True)
+
+  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)
+  scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
+
+  @TinyJit
+  @Tensor.train()
+  def train_step(model, x, y):
+    optim.zero_grad()
+    logits:Tensor = model(x, start_pos=0, temperature=math.nan)
+    loss = logits.cross_entropy(y)
+    loss.backward()
+
+    # L2 norm grad clip
+    # https://github.com/NVIDIA/NeMo/blob/3368c3fc0b4a186ab33a1d68a504315100c0b2a6/nemo/collections/nlp/modules/common/megatron/clip_grads.py#L57
+    # https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html
+    if not getenv("DISABLE_GRAD_CLIP_NORM"):
+      total_norm = Tensor(0.0, dtype=dtypes.float32, device=optim.params[0].device)
+      for p in optim.params:
+        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)
+
+    optim.step()
+    scheduler.step()
+
+    lr = optim.lr
+    loss.realize(lr)
+    return loss, lr
+
+  # overfitting this example should give cross_entropy log(BS)
+  fake_input = Tensor([list(range(getenv("SEQLEN", 10)))], dtype="int16").expand(BS, -1)
+  fake_label = Tensor(list(range(BS)), dtype="int16")
+
+  for _ in range(100):
+    GlobalCounters.reset()
+    loss, lr = train_step(model, fake_input, fake_label)
+    # BS=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=8B WARMUP_STEPS=2 DECAY_STEPS=300 PYTHONPATH=. AMD=1 MODEL=llama3 python3 examples/mlperf/model_train.py
+    # uses 43% ~= 83GB
+    # 8B bf16 = 16GB. model + grad + optim m and v = 64GB
+    # TODO: this OOM
+    # BS=1 SEQLEN=4000 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=8B WARMUP_STEPS=2 DECAY_STEPS=300 PYTHONPATH=. AMD=1 MODEL=llama3 python3 examples/mlperf/model_train.py
+    print(loss.item(), lr.item(), f"{GlobalCounters.global_mem//10**9=}")
+
 if __name__ == "__main__":
   multiprocessing.set_start_method('spawn')
 

examples/openpilot/compile4.py

@@ -1,8 +1,8 @@
 import sys, onnx
-from tinygrad import Tensor, fetch, GlobalCounters
+from tinygrad import Tensor, fetch, GlobalCounters, dtypes
 from tinygrad.uop.ops import UOp
 from tinygrad.frontend.onnx import OnnxRunner
-from tinygrad.engine.kernelize import get_kernelize_map
+from tinygrad.kernelize.kernelize import get_kernelize_map
 from tinygrad.engine.schedule import create_schedule_with_vars
 from tinygrad.engine.realize import run_schedule
 
@@ -17,7 +17,7 @@ if __name__ == "__main__":
   onnx_model = onnx.load(onnx_file)
   run_onnx = OnnxRunner(onnx_model)
 
-  inputs = run_onnx.get_empty_input_data("npy")
+  inputs = run_onnx.get_empty_input_data("npy", dtypes.float32)
   out: Tensor = next(iter(run_onnx({k:v.to(None) for k,v in inputs.items()}).values())).to('cpu')
   root = out.uop
   targets = [x.uop for x in inputs.values()]

examples/other_mnist/beautiful_mnist_torch.py

@@ -27,7 +27,7 @@ class Model(nn.Module):
 
 if __name__ == "__main__":
   if getenv("TINY_BACKEND"):
-    import tinygrad.frontend.torch
+    import tinygrad.frontend.torch  # noqa: F401
     device = torch.device("tiny")
   else:
     device = torch.device({"METAL":"mps","NV":"cuda"}.get(Device.DEFAULT, "cpu"))

examples/sdxl.py

@@ -5,7 +5,7 @@
 
 from tinygrad import Tensor, TinyJit, dtypes, GlobalCounters
 from tinygrad.nn import Conv2d, GroupNorm
-from tinygrad.nn.state import safe_load, load_state_dict, get_state_dict
+from tinygrad.nn.state import safe_load, load_state_dict
 from tinygrad.helpers import fetch, trange, colored, Timing
 from extra.models.clip import Embedder, FrozenClosedClipEmbedder, FrozenOpenClipEmbedder
 from extra.models.unet import UNetModel, Upsample, Downsample, timestep_embedding

examples/stunning_mnist.py

@@ -7,7 +7,7 @@
 from examples.beautiful_mnist import Model
 from tinygrad import Tensor, nn, getenv, GlobalCounters, Variable
 from tinygrad.nn.datasets import mnist
-from tinygrad.helpers import trange, DEBUG
+from tinygrad.helpers import trange
 
 # STEPS=70 python3 examples/stunning_mnist.py
 # NOTE: it's broken with STACK=1, why?

examples/torch_cuda_kernel.py

@@ -2,10 +2,9 @@
 #!POPCORN gpu A100
 # not a stable API, but works
 
-import torch, functools
+import torch
 from tinygrad import Tensor, TinyJit, Device
-from tinygrad.engine.realize import CompiledRunner
-from tinygrad.helpers import get_single_element, Context, OSX
+from tinygrad.helpers import Context, OSX
 from tinygrad.dtype import _from_torch_dtype
 
 @TinyJit

examples/vgg7.py

@@ -2,8 +2,6 @@ import sys
 import random
 import json
 import numpy
-from pathlib import Path
-from PIL import Image
 from tinygrad.tensor import Tensor
 from tinygrad.nn.optim import SGD
 from tinygrad.nn.state import safe_save, safe_load, get_state_dict, load_state_dict

examples/whisper.py

@@ -5,7 +5,7 @@ from typing import Optional, Union, Literal, List
 
 from tinygrad import Tensor, TinyJit, Variable, nn
 from tinygrad.nn.state import torch_load, load_state_dict
-from tinygrad.helpers import getenv, DEBUG, fetch
+from tinygrad.helpers import getenv, fetch
 
 import numpy as np
 import librosa

examples/yolov8-onnx.py

@@ -4,7 +4,6 @@ from ultralytics import YOLO
 from pathlib import Path
 from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 from extra.onnx_helpers import get_example_inputs
-from tinygrad.tensor import Tensor
 
 os.chdir("/tmp")
 if not Path("yolov8n-seg.onnx").is_file():

examples/yolov8.py

@@ -1,7 +1,5 @@
 from tinygrad.nn import Conv2d, BatchNorm2d
 from tinygrad.tensor import Tensor
-from tinygrad.device import is_dtype_supported
-from tinygrad import dtypes
 import numpy as np
 from itertools import chain
 from pathlib import Path

extra/assembly/assembly.py

@@ -156,7 +156,7 @@ def uops_to_asmstyle(lang, function_name:str, uops:List[UOp]):
         lang.ins.append(AssemblyInstruction(Ops.ALU, out, [tmp], args))
       else:
         lang.ins.append(AssemblyInstruction(Ops.ALU, out, [lang.tor[x] for x in vin], args))
-    elif uop == Ops.DEFINE_ACC:
+    elif uop == Ops.DEFINE_REG:
       reg = lang.newreg(u, dtype=dtype)
       lang.ins.append(AssemblyInstruction(Ops.LOAD, reg, [], args))
     elif uop == Ops.SPECIAL:

extra/backends/triton.py

@@ -88,7 +88,7 @@ def uops_to_triton(function_name:str, uops:List[UOp]):
       assert dtype is not None
       if len(vin) == 2: kk(f"{ssa(u, 'val')} = {render_cast(f'tl.load({r[vin[0]]} + { fill_dims_for_idx(r[vin[1]], dims)}, mask = {render_valid(valid)})', dtype)}")
       else: kk(f"{ssa(u, 'val')} = {render_cast(f'tl.where({r[vin[2]]}, tl.load({r[vin[0]]}+{fill_dims_for_idx(r[vin[1]],dims)} , mask={render_valid(valid+[r[vin[2]]])}), 0.0)', dtype)}")
-    elif uop == Ops.DEFINE_ACC: kk(f"{ssa(u, 'acc')} = {define_scalar(local_size, dtype, args).replace('//', '/')}")
+    elif uop == Ops.DEFINE_REG: kk(f"{ssa(u, 'acc')} = {define_scalar(local_size, dtype, args).replace('//', '/')}")
     elif uop == Ops.CONST: r[u] = define_scalar([], dtype, args)
     elif uop == Ops.ASSIGN:
       kk(f"{r[vin[0]]} = {r[vin[1]].replace('//', '/')}")

extra/gemm/halide_gemm.py

@@ -0,0 +1,90 @@
+import numpy as np
+import halide as hl
+from tinygrad.helpers import Timing, getenv
+
+# HL_DEBUG_CODEGEN=1
+N = getenv("N", 1024)
+
+def gemm_pipeline(gpu=False):
+  # ---------------- Vars & Parameters ----------------
+  i, j = hl.Var("i"), hl.Var("j")  # output tile coordinates
+
+  A = hl.InputBuffer(hl.Float(32), 2)  # [M, K]
+  B = hl.InputBuffer(hl.Float(32), 2)  # [K, N]
+
+  A.dim(0).set_bounds(0, N)
+  A.dim(1).set_bounds(0, N)
+  B.dim(0).set_bounds(0, N)
+  B.dim(1).set_bounds(0, N)
+
+  # ---------------- Definition ----------------
+
+  k = hl.RDom([(0, N)])
+
+  partial = hl.Func("partial")
+  partial[i, j] = 0.0
+  partial[i, j] += A[i, k] * B[k, j]
+
+  C = hl.Func("C")
+  C[i, j] = partial[i, j]
+
+  if not gpu:
+    # ---------------- Schedule ----------------
+    VEC = 16
+    TILE_I = 64
+    TILE_J = 64
+
+    io, jo, ii, ji = hl.Var("io"), hl.Var("jo"), hl.Var("ii"), hl.Var("ji")
+    C.update().tile(i, j, io, jo, ii, ji, TILE_I, TILE_J).fuse(io, jo, io).parallel(io).vectorize(ji, VEC)
+  else:
+    # ---------------- Schedule ----------------
+    GRP_I    = 8     # output tile size
+    GRP_J    = 16
+
+    #partial.store_in(hl.MemoryType.Register)
+    #partial.update().unroll(k, 4)
+
+    io, jo, ii, ji = hl.Var(), hl.Var(), hl.Var(), hl.Var()
+    C.gpu_tile(i, j, io, jo, ii, ji, GRP_I, GRP_J, hl.TailStrategy.RoundUp)
+
+  return C, A, B
+
+if __name__ == "__main__":
+  pipe, A, B = gemm_pipeline(gpu=True)
+
+  # NOTE: meteal does nothing
+  target = hl.get_host_target().with_feature(hl.TargetFeature.Metal)
+
+  a_np = np.random.randn(N, N).astype(np.float32)
+  b_np = np.random.randn(N, N).astype(np.float32)
+
+  # reverse order is correct!
+  a_hal = hl.Buffer(b_np)
+  b_hal = hl.Buffer(a_np)
+  A.set(a_hal)
+  B.set(b_hal)
+
+  pipe.compile_to_lowered_stmt("/tmp/my_function.html", [A, B], hl.StmtOutputFormat.HTML, target=target)
+  #exit(0)
+
+  c_hal = hl.Buffer(hl.Float(32), [N,N])
+  with Timing("halide gemm "):
+    pipe.realize(c_hal, target)
+    c_hal.copy_to_host()
+    c_out = np.array(c_hal)
+  print(c_out)
+
+  # tinygrad gets 60 ms with no BEAM, 20 ms with BEAM on CPU
+  with Timing("halide gemm "):
+    pipe.realize(c_hal, target)
+    c_hal.copy_to_host()
+
+  # Check correctness
+  with Timing("numpy gemm "):
+    ref = a_np @ b_np
+  max_err = np.abs(ref - c_out).max()
+  print("Max absolute error:", max_err)
+  assert max_err < 1e-4, "GEMM result incorrect!"
+
+  print("Pipeline ran on", target)
+  print("Success - GEMM Halide-Python output matches NumPy.")

extra/hcqfuzz/readme

@@ -4,9 +4,9 @@ To add a new test, define a `TestSpec`-based class in a file in the `tests/` fol
 
 You can choose which tests to load from which file:
 ```bash
-RUN_FILES="hcq,allocator" python3 extra/hcqfuzz/fuzzer.py
+PYTHONPATH=. RUN_FILES="hcq,allocator" python3 extra/hcqfuzz/fuzzer.py
 ```
 Or skip tests from any file:
 ```bash
-SKIP_FILES="allocator" python3 extra/hcqfuzz/fuzzer.py
+PYTHONPATH=. SKIP_FILES="allocator" python3 extra/hcqfuzz/fuzzer.py
 ```

extra/models/llama.py

@@ -1,5 +1,5 @@
 from typing import Union, Optional, Any
-import collections
+import collections, math
 from tinygrad import Tensor, Variable, TinyJit, dtypes, nn, Device
 from tinygrad.helpers import getenv, DEBUG
 
@@ -166,27 +166,29 @@ def sample(logits: Tensor, temp: float, k: int, p: float, af: float, ap: float):
 
 class Transformer:
   def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size, linear=nn.Linear, embedding=nn.Embedding,
-               n_kv_heads=None, rope_theta=10000, max_context=1024, jit=True, feed_forward=FeedForward, qk_norm=None):
-    self.layers = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, max_context, linear, feed_forward=feed_forward, qk_norm=qk_norm) for _ in range(n_layers)]
+               n_kv_heads=None, rope_theta=10000, max_context=1024, jit=True, feed_forward=FeedForward, qk_norm=None, disable_kv_cache=False):
+    self.layers = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, 0 if disable_kv_cache else max_context,
+                                    linear, feed_forward=feed_forward, qk_norm=qk_norm) for _ in range(n_layers)]
     self.norm = nn.RMSNorm(dim, norm_eps)
     self.tok_embeddings = embedding(vocab_size, dim)
     self.output = nn.Linear(dim, vocab_size, bias=False) if embedding == nn.Embedding else linear(dim, vocab_size, bias=False)
     self.max_context = max_context
-    self.freqs_cis = precompute_freqs_cis(dim // n_heads, self.max_context * 2, rope_theta).contiguous()
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, self.max_context * 2, rope_theta).contiguous().requires_grad_(False)
     self.forward_jit = TinyJit(self.forward) if jit else None
 
   def forward(self, tokens:Tensor, start_pos:Union[Variable,int], temperature:float, top_k:int, top_p:float, alpha_f:float, alpha_p:float):
     _bsz, seqlen = tokens.shape
     h = self.tok_embeddings(tokens)
 
-    self.freqs_cis = self.freqs_cis.cast(h.dtype).kernelize()
+    self.freqs_cis = self.freqs_cis.cast(h.dtype).contiguous()
     freqs_cis = self.freqs_cis[:, start_pos:start_pos+seqlen, :, :, :]
 
-    mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1).kernelize() if seqlen > 1 else None
+    mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1) if seqlen > 1 else None
     for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask)
     logits = self.output(self.norm(h)).float()[:, -1, :]
+    if math.isnan(temperature): return logits
 
-    return sample(logits.flatten(), temperature, top_k, top_p, alpha_f, alpha_p).kernelize()
+    return sample(logits.flatten(), temperature, top_k, top_p, alpha_f, alpha_p)
 
   def __call__(self, tokens:Tensor, start_pos:int, temperature:float=0.0, top_k:int=0, top_p:float=0.8, alpha_f:float=0.0, alpha_p:float=0.0):
     # TODO: better way to handle the first call v.s. the rest?

extra/nv_gpu_driver/g_rpc-message-header.h

@@ -0,0 +1,77 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2008-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+/*
+ * WARNING: This is an autogenerated file. DO NOT EDIT.
+ * This file is generated using below files:
+ * template file: inc/kernel/vgpu/gt_rpc-message.h
+ * definition file: inc/kernel/vgpu/rpc-message-header.def
+ */
+
+typedef struct GSP_MSG_QUEUE_ELEMENT
+{
+    NvU8  authTagBuffer[16];         // Authentication tag buffer.
+    NvU8  aadBuffer[16];             // AAD buffer.
+    NvU32 checkSum;                  // Set to value needed to make checksum always zero.
+    NvU32 seqNum;                    // Sequence number maintained by the message queue.
+    NvU32 elemCount;                 // Number of message queue elements this message has.
+    NvU32 padding;                  // Reserved for future use.
+} GSP_MSG_QUEUE_ELEMENT;
+
+#ifdef RPC_MESSAGE_STRUCTURES
+typedef union rpc_message_rpc_union_field_v03_00
+{
+    NvU32      spare;
+    NvU32      cpuRmGfid;
+} rpc_message_rpc_union_field_v03_00;
+
+typedef rpc_message_rpc_union_field_v03_00 rpc_message_rpc_union_field_v;
+
+typedef struct rpc_message_header_v03_00
+{
+    NvU32      header_version;
+    NvU32      signature;
+    NvU32      length;
+    NvU32      function;
+    NvU32      rpc_result;
+    NvU32      rpc_result_private;
+    NvU32      sequence;
+    rpc_message_rpc_union_field_v u;
+    // rpc_generic_union rpc_message_data[];
+} rpc_message_header_v03_00;
+
+typedef rpc_message_header_v03_00 rpc_message_header_v;
+
+
+#endif
+
+#ifdef RPC_MESSAGE_GENERIC_UNION
+// This is a generic union, that will be used for the communication between the vmioplugin & guest RM.
+typedef union rpc_message_generic_union {
+    rpc_message_rpc_union_field_v03_00 rpc_union_field_v03_00;
+    rpc_message_rpc_union_field_v rpc_union_field_v;
+    rpc_message_header_v03_00 header_v03_00;
+    rpc_message_header_v header_v;
+} rpc_message_generic_union;
+
+#endif

extra/nv_gpu_driver/gsp_static_config.h

@@ -0,0 +1,455 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#ifndef GSP_STATIC_CONFIG_H
+#define GSP_STATIC_CONFIG_H
+
+//
+// This header describes the set of static GPU configuration information
+// that is collected during GSP RM init and made available to the
+// CPU RM (aka GSP client) via NV_RM_RPC_GET_GSP_STATIC_INFO() call.
+
+#include "ctrl/ctrl0080/ctrl0080gpu.h"
+#include "ctrl/ctrl2080/ctrl2080bios.h"
+#include "ctrl/ctrl2080/ctrl2080fb.h"
+#include "ctrl/ctrl2080/ctrl2080gpu.h"
+
+#include "vgpu/rpc_headers.h"
+#include "nvacpitypes.h"
+
+#include "ctrl/ctrl0073/ctrl0073system.h"
+
+#define MAX_DSM_SUPPORTED_FUNCS_RTN_LEN 8 // # bytes to store supported functions
+#define NV_ACPI_GENERIC_FUNC_COUNT                  8
+
+#define REGISTRY_TABLE_ENTRY_TYPE_UNKNOWN  0
+#define REGISTRY_TABLE_ENTRY_TYPE_DWORD    1
+#define REGISTRY_TABLE_ENTRY_TYPE_BINARY   2
+#define REGISTRY_TABLE_ENTRY_TYPE_STRING   3
+typedef struct PACKED_REGISTRY_ENTRY
+{
+    NvU32                   nameOffset;
+    NvU8                    type;
+    NvU32                   data;
+    NvU32                   length;
+} PACKED_REGISTRY_ENTRY;
+
+typedef struct PACKED_REGISTRY_TABLE
+{
+    NvU32                   size;
+    NvU32                   numEntries;
+} PACKED_REGISTRY_TABLE;
+
+/* Indicates the current state of mux */
+typedef enum
+{
+    dispMuxState_None = 0,
+    dispMuxState_IntegratedGPU,
+    dispMuxState_DiscreteGPU,
+} DISPMUXSTATE;
+
+typedef struct {
+    // supported function status and cache
+    NvU32  suppFuncStatus;
+    NvU8   suppFuncs[MAX_DSM_SUPPORTED_FUNCS_RTN_LEN];
+    NvU32  suppFuncsLen;
+    NvBool bArg3isInteger;
+    // callback status and cache
+    NvU32  callbackStatus;
+    NvU32  callback;
+} ACPI_DSM_CACHE;
+
+typedef struct {
+
+    ACPI_DSM_CACHE                   dsm[ACPI_DSM_FUNCTION_COUNT];
+    ACPI_DSM_FUNCTION                dispStatusHotplugFunc;
+    ACPI_DSM_FUNCTION                dispStatusConfigFunc;
+    ACPI_DSM_FUNCTION                perfPostPowerStateFunc;
+    ACPI_DSM_FUNCTION                stereo3dStateActiveFunc;
+    NvU32                            dsmPlatCapsCache[ACPI_DSM_FUNCTION_COUNT];
+    NvU32                            MDTLFeatureSupport;
+
+    // cache of generic func/subfunction remappings.
+    ACPI_DSM_FUNCTION                dsmCurrentFunc[NV_ACPI_GENERIC_FUNC_COUNT];
+    NvU32                            dsmCurrentSubFunc[NV_ACPI_GENERIC_FUNC_COUNT];
+    NvU32                            dsmCurrentFuncSupport;
+
+} ACPI_DATA;
+
+typedef struct DOD_METHOD_DATA
+{
+    NV_STATUS status;
+    NvU32     acpiIdListLen;
+    NvU32     acpiIdList[NV0073_CTRL_SYSTEM_ACPI_ID_MAP_MAX_DISPLAYS];
+} DOD_METHOD_DATA;
+
+typedef struct JT_METHOD_DATA
+{
+    NV_STATUS status;
+    NvU32     jtCaps;
+    NvU16     jtRevId;
+    NvBool    bSBIOSCaps;
+} JT_METHOD_DATA;
+
+typedef struct MUX_METHOD_DATA_ELEMENT
+{
+    NvU32       acpiId;
+    NvU32       mode;
+    NV_STATUS   status;
+} MUX_METHOD_DATA_ELEMENT;
+
+typedef struct MUX_METHOD_DATA
+{
+    NvU32                       tableLen;
+    MUX_METHOD_DATA_ELEMENT     acpiIdMuxModeTable[NV0073_CTRL_SYSTEM_ACPI_ID_MAP_MAX_DISPLAYS];
+    MUX_METHOD_DATA_ELEMENT     acpiIdMuxPartTable[NV0073_CTRL_SYSTEM_ACPI_ID_MAP_MAX_DISPLAYS];
+    MUX_METHOD_DATA_ELEMENT     acpiIdMuxStateTable[NV0073_CTRL_SYSTEM_ACPI_ID_MAP_MAX_DISPLAYS];    
+} MUX_METHOD_DATA;
+
+typedef struct CAPS_METHOD_DATA
+{
+    NV_STATUS status;
+    NvU32     optimusCaps;
+} CAPS_METHOD_DATA;
+
+typedef struct ACPI_METHOD_DATA
+{
+    NvBool                                               bValid;
+    DOD_METHOD_DATA                                      dodMethodData;
+    JT_METHOD_DATA                                       jtMethodData;
+    MUX_METHOD_DATA                                      muxMethodData;
+    CAPS_METHOD_DATA                                     capsMethodData;
+} ACPI_METHOD_DATA;
+
+#define MAX_GROUP_COUNT 2
+
+// #include "gpu/nvbitmask.h"
+typedef enum
+{
+    RM_ENGINE_TYPE_NULL                            =       (0x00000000),
+    RM_ENGINE_TYPE_GR0                             =       (0x00000001),
+    RM_ENGINE_TYPE_GR1                             =       (0x00000002),
+    RM_ENGINE_TYPE_GR2                             =       (0x00000003),
+    RM_ENGINE_TYPE_GR3                             =       (0x00000004),
+    RM_ENGINE_TYPE_GR4                             =       (0x00000005),
+    RM_ENGINE_TYPE_GR5                             =       (0x00000006),
+    RM_ENGINE_TYPE_GR6                             =       (0x00000007),
+    RM_ENGINE_TYPE_GR7                             =       (0x00000008),
+    RM_ENGINE_TYPE_COPY0                           =       (0x00000009),
+    RM_ENGINE_TYPE_COPY1                           =       (0x0000000a),
+    RM_ENGINE_TYPE_COPY2                           =       (0x0000000b),
+    RM_ENGINE_TYPE_COPY3                           =       (0x0000000c),
+    RM_ENGINE_TYPE_COPY4                           =       (0x0000000d),
+    RM_ENGINE_TYPE_COPY5                           =       (0x0000000e),
+    RM_ENGINE_TYPE_COPY6                           =       (0x0000000f),
+    RM_ENGINE_TYPE_COPY7                           =       (0x00000010),
+    RM_ENGINE_TYPE_COPY8                           =       (0x00000011),
+    RM_ENGINE_TYPE_COPY9                           =       (0x00000012),
+    RM_ENGINE_TYPE_COPY10                          =       (0x00000013),
+    RM_ENGINE_TYPE_COPY11                          =       (0x00000014),
+    RM_ENGINE_TYPE_COPY12                          =       (0x00000015),
+    RM_ENGINE_TYPE_COPY13                          =       (0x00000016),
+    RM_ENGINE_TYPE_COPY14                          =       (0x00000017),
+    RM_ENGINE_TYPE_COPY15                          =       (0x00000018),
+    RM_ENGINE_TYPE_COPY16                          =       (0x00000019),
+    RM_ENGINE_TYPE_COPY17                          =       (0x0000001a),
+    RM_ENGINE_TYPE_COPY18                          =       (0x0000001b),
+    RM_ENGINE_TYPE_COPY19                          =       (0x0000001c),
+    RM_ENGINE_TYPE_NVDEC0                          =       (0x0000001d),
+    RM_ENGINE_TYPE_NVDEC1                          =       (0x0000001e),
+    RM_ENGINE_TYPE_NVDEC2                          =       (0x0000001f),
+    RM_ENGINE_TYPE_NVDEC3                          =       (0x00000020),
+    RM_ENGINE_TYPE_NVDEC4                          =       (0x00000021),
+    RM_ENGINE_TYPE_NVDEC5                          =       (0x00000022),
+    RM_ENGINE_TYPE_NVDEC6                          =       (0x00000023),
+    RM_ENGINE_TYPE_NVDEC7                          =       (0x00000024),
+    RM_ENGINE_TYPE_NVENC0                          =       (0x00000025),
+    RM_ENGINE_TYPE_NVENC1                          =       (0x00000026),
+    RM_ENGINE_TYPE_NVENC2                          =       (0x00000027),
+    // Bug 4175886 - Use this new value for all chips once GB20X is released
+    RM_ENGINE_TYPE_NVENC3                          =       (0x00000028),
+    RM_ENGINE_TYPE_VP                              =       (0x00000029),
+    RM_ENGINE_TYPE_ME                              =       (0x0000002a),
+    RM_ENGINE_TYPE_PPP                             =       (0x0000002b),
+    RM_ENGINE_TYPE_MPEG                            =       (0x0000002c),
+    RM_ENGINE_TYPE_SW                              =       (0x0000002d),
+    RM_ENGINE_TYPE_TSEC                            =       (0x0000002e),
+    RM_ENGINE_TYPE_VIC                             =       (0x0000002f),
+    RM_ENGINE_TYPE_MP                              =       (0x00000030),
+    RM_ENGINE_TYPE_SEC2                            =       (0x00000031),
+    RM_ENGINE_TYPE_HOST                            =       (0x00000032),
+    RM_ENGINE_TYPE_DPU                             =       (0x00000033),
+    RM_ENGINE_TYPE_PMU                             =       (0x00000034),
+    RM_ENGINE_TYPE_FBFLCN                          =       (0x00000035),
+    RM_ENGINE_TYPE_NVJPEG0                         =       (0x00000036),
+    RM_ENGINE_TYPE_NVJPEG1                         =       (0x00000037),
+    RM_ENGINE_TYPE_NVJPEG2                         =       (0x00000038),
+    RM_ENGINE_TYPE_NVJPEG3                         =       (0x00000039),
+    RM_ENGINE_TYPE_NVJPEG4                         =       (0x0000003a),
+    RM_ENGINE_TYPE_NVJPEG5                         =       (0x0000003b),
+    RM_ENGINE_TYPE_NVJPEG6                         =       (0x0000003c),
+    RM_ENGINE_TYPE_NVJPEG7                         =       (0x0000003d),
+    RM_ENGINE_TYPE_OFA0                            =       (0x0000003e),
+    RM_ENGINE_TYPE_OFA1                            =       (0x0000003f),
+    RM_ENGINE_TYPE_RESERVED40                      =       (0x00000040),
+    RM_ENGINE_TYPE_RESERVED41                      =       (0x00000041),
+    RM_ENGINE_TYPE_RESERVED42                      =       (0x00000042),
+    RM_ENGINE_TYPE_RESERVED43                      =       (0x00000043),
+    RM_ENGINE_TYPE_RESERVED44                      =       (0x00000044),
+    RM_ENGINE_TYPE_RESERVED45                      =       (0x00000045),
+    RM_ENGINE_TYPE_RESERVED46                      =       (0x00000046),
+    RM_ENGINE_TYPE_RESERVED47                      =       (0x00000047),
+    RM_ENGINE_TYPE_RESERVED48                      =       (0x00000048),
+    RM_ENGINE_TYPE_RESERVED49                      =       (0x00000049),
+    RM_ENGINE_TYPE_RESERVED4a                      =       (0x0000004a),
+    RM_ENGINE_TYPE_RESERVED4b                      =       (0x0000004b),
+    RM_ENGINE_TYPE_RESERVED4c                      =       (0x0000004c),
+    RM_ENGINE_TYPE_RESERVED4d                      =       (0x0000004d),
+    RM_ENGINE_TYPE_RESERVED4e                      =       (0x0000004e),
+    RM_ENGINE_TYPE_RESERVED4f                      =       (0x0000004f),
+    RM_ENGINE_TYPE_RESERVED50                      =       (0x00000050),
+    RM_ENGINE_TYPE_RESERVED51                      =       (0x00000051),
+    RM_ENGINE_TYPE_RESERVED52                      =       (0x00000052),
+    RM_ENGINE_TYPE_RESERVED53                      =       (0x00000053),
+    RM_ENGINE_TYPE_LAST                            =       (0x00000054),
+} RM_ENGINE_TYPE;
+
+//
+// The duplicates in the RM_ENGINE_TYPE. Using define instead of putting them
+// in the enum to make sure that each item in the enum has a unique number.
+//
+#define RM_ENGINE_TYPE_GRAPHICS                 RM_ENGINE_TYPE_GR0
+#define RM_ENGINE_TYPE_BSP                      RM_ENGINE_TYPE_NVDEC0
+#define RM_ENGINE_TYPE_MSENC                    RM_ENGINE_TYPE_NVENC0
+#define RM_ENGINE_TYPE_CIPHER                   RM_ENGINE_TYPE_TSEC
+#define RM_ENGINE_TYPE_NVJPG                    RM_ENGINE_TYPE_NVJPEG0
+
+#define RM_ENGINE_TYPE_COPY_SIZE 20
+// Bug 4175886 - Use this new value for all chips once GB20X is released
+#define RM_ENGINE_TYPE_NVENC_SIZE 4
+#define RM_ENGINE_TYPE_NVJPEG_SIZE 8
+#define RM_ENGINE_TYPE_NVDEC_SIZE 8
+#define RM_ENGINE_TYPE_OFA_SIZE 2
+#define RM_ENGINE_TYPE_GR_SIZE 8
+
+#define NVGPU_ENGINE_CAPS_MASK_BITS                32
+#define NVGPU_ENGINE_CAPS_MASK_ARRAY_MAX           ((RM_ENGINE_TYPE_LAST-1)/NVGPU_ENGINE_CAPS_MASK_BITS + 1)
+#define NVGPU_GET_ENGINE_CAPS_MASK(caps, id)       (caps[(id)/NVGPU_ENGINE_CAPS_MASK_BITS] & NVBIT((id) % NVGPU_ENGINE_CAPS_MASK_BITS))
+#define NVGPU_SET_ENGINE_CAPS_MASK(caps, id)       (caps[(id)/NVGPU_ENGINE_CAPS_MASK_BITS] |= NVBIT((id) % NVGPU_ENGINE_CAPS_MASK_BITS))
+
+
+// #include "gpu/gpu.h" // COMPUTE_BRANDING_TYPE
+// #include "gpu/gpu_acpi_data.h" // ACPI_METHOD_DATA
+// #include "vgpu/rpc_headers.h" // MAX_GPC_COUNT
+// #include "platform/chipset/chipset.h" // BUSINFO
+// #include "gpu/nvbitmask.h" // NVGPU_ENGINE_CAPS_MASK_ARRAY_MAX
+
+typedef struct
+{
+    NvU16               deviceID;           // deviceID
+    NvU16               vendorID;           // vendorID
+    NvU16               subdeviceID;        // subsystem deviceID
+    NvU16               subvendorID;        // subsystem vendorID
+    NvU8                revisionID;         // revision ID
+} BUSINFO;
+
+// VF related info for GSP-RM
+typedef struct GSP_VF_INFO
+{
+    NvU32  totalVFs;
+    NvU32  firstVFOffset;
+    NvU64  FirstVFBar0Address;
+    NvU64  FirstVFBar1Address;
+    NvU64  FirstVFBar2Address;
+    NvBool b64bitBar0;
+    NvBool b64bitBar1;
+    NvBool b64bitBar2;
+} GSP_VF_INFO;
+
+// Cache config registers from pcie space
+typedef struct
+{
+    // Link capabilities
+    NvU32 linkCap;
+} GSP_PCIE_CONFIG_REG;
+
+typedef struct
+{
+    NvU32 ecidLow;
+    NvU32 ecidHigh;
+    NvU32 ecidExtended;
+} EcidManufacturingInfo;
+
+typedef struct
+{
+    NvU64 nonWprHeapOffset;
+    NvU64 frtsOffset;
+} FW_WPR_LAYOUT_OFFSET;
+
+// Fetched from GSP-RM into CPU-RM
+typedef struct GspStaticConfigInfo_t
+{
+    NvU8 grCapsBits[NV0080_CTRL_GR_CAPS_TBL_SIZE];
+    NV2080_CTRL_GPU_GET_GID_INFO_PARAMS gidInfo;
+    NV2080_CTRL_BIOS_GET_SKU_INFO_PARAMS SKUInfo;
+    NV2080_CTRL_CMD_FB_GET_FB_REGION_INFO_PARAMS fbRegionInfoParams;
+
+    NV0080_CTRL_GPU_GET_SRIOV_CAPS_PARAMS sriovCaps;
+    NvU32 sriovMaxGfid;
+
+    NvU32 engineCaps[NVGPU_ENGINE_CAPS_MASK_ARRAY_MAX];
+
+    NvBool poisonFuseEnabled;
+
+    NvU64 fb_length;
+    NvU64 fbio_mask;
+    NvU32 fb_bus_width;
+    NvU32 fb_ram_type;
+    NvU64 fbp_mask;
+    NvU32 l2_cache_size;
+
+    NvU8 gpuNameString[NV2080_GPU_MAX_NAME_STRING_LENGTH];
+    NvU8 gpuShortNameString[NV2080_GPU_MAX_NAME_STRING_LENGTH];
+    NvU16 gpuNameString_Unicode[NV2080_GPU_MAX_NAME_STRING_LENGTH];
+    NvBool bGpuInternalSku;
+    NvBool bIsQuadroGeneric;
+    NvBool bIsQuadroAd;
+    NvBool bIsNvidiaNvs;
+    NvBool bIsVgx;
+    NvBool bGeforceSmb;
+    NvBool bIsTitan;
+    NvBool bIsTesla;
+    NvBool bIsMobile;
+    NvBool bIsGc6Rtd3Allowed;
+    NvBool bIsGc8Rtd3Allowed;
+    NvBool bIsGcOffRtd3Allowed;
+    NvBool bIsGcoffLegacyAllowed;
+    NvBool bIsMigSupported;
+
+    /* "Total Board Power" refers to power requirement of GPU,
+     * while in GC6 state. Majority of this power will be used
+     * to keep V-RAM active to preserve its content.
+     * Some energy maybe consumed by Always-on components on GPU chip.
+     * This power will be provided by 3.3v voltage rail.
+     */
+    NvU16  RTD3GC6TotalBoardPower;
+
+    /* PERST# (i.e. PCI Express Reset) is a sideband signal
+     * generated by the PCIe Host to indicate the PCIe devices,
+     * that the power-rails and the reference-clock are stable.
+     * The endpoint device typically uses this signal as a global reset.
+     */
+    NvU16  RTD3GC6PerstDelay;
+
+    NvU64 bar1PdeBase;
+    NvU64 bar2PdeBase;
+
+    NvBool bVbiosValid;
+    NvU32 vbiosSubVendor;
+    NvU32 vbiosSubDevice;
+
+    NvBool bPageRetirementSupported;
+
+    NvBool bSplitVasBetweenServerClientRm;
+
+    NvBool bClRootportNeedsNosnoopWAR;
+
+    VIRTUAL_DISPLAY_GET_NUM_HEADS_PARAMS displaylessMaxHeads;
+    VIRTUAL_DISPLAY_GET_MAX_RESOLUTION_PARAMS displaylessMaxResolution;
+    NvU64 displaylessMaxPixels;
+
+    // Client handle for internal RMAPI control.
+    NvHandle hInternalClient;
+
+    // Device handle for internal RMAPI control.
+    NvHandle hInternalDevice;
+
+    // Subdevice handle for internal RMAPI control.
+    NvHandle hInternalSubdevice;
+
+    NvBool bSelfHostedMode;
+    NvBool bAtsSupported;
+
+    NvBool bIsGpuUefi;
+    NvBool bIsEfiInit;
+
+    EcidManufacturingInfo ecidInfo[MAX_GROUP_COUNT];
+
+    FW_WPR_LAYOUT_OFFSET fwWprLayoutOffset;
+} GspStaticConfigInfo;
+
+// Pushed from CPU-RM to GSP-RM
+typedef struct GspSystemInfo
+{
+    NvU64 gpuPhysAddr;
+    NvU64 gpuPhysFbAddr;
+    NvU64 gpuPhysInstAddr;
+    NvU64 gpuPhysIoAddr;
+    NvU64 nvDomainBusDeviceFunc;
+    NvU64 simAccessBufPhysAddr;
+    NvU64 notifyOpSharedSurfacePhysAddr;
+    NvU64 pcieAtomicsOpMask;
+    NvU64 consoleMemSize;
+    NvU64 maxUserVa;
+    NvU32 pciConfigMirrorBase;
+    NvU32 pciConfigMirrorSize;
+    NvU32 PCIDeviceID;
+    NvU32 PCISubDeviceID;
+    NvU32 PCIRevisionID;
+    NvU32 pcieAtomicsCplDeviceCapMask;
+    NvU8 oorArch;
+    NvU64 clPdbProperties;
+    NvU32 Chipset;
+    NvBool bGpuBehindBridge;
+    NvBool bFlrSupported;
+    NvBool b64bBar0Supported;
+    NvBool bMnocAvailable;
+    NvU32  chipsetL1ssEnable;
+    NvBool bUpstreamL0sUnsupported;
+    NvBool bUpstreamL1Unsupported;
+    NvBool bUpstreamL1PorSupported;
+    NvBool bUpstreamL1PorMobileOnly;
+    NvBool bSystemHasMux;
+    NvU8   upstreamAddressValid;
+    BUSINFO FHBBusInfo;
+    BUSINFO chipsetIDInfo;
+    ACPI_METHOD_DATA acpiMethodData;
+    NvU32 hypervisorType;
+    NvBool bIsPassthru;
+    NvU64 sysTimerOffsetNs;
+    GSP_VF_INFO gspVFInfo;
+    NvBool bIsPrimary;
+    NvBool isGridBuild;
+    GSP_PCIE_CONFIG_REG pcieConfigReg;
+    NvU32 gridBuildCsp;
+    NvBool bPreserveVideoMemoryAllocations;
+    NvBool bTdrEventSupported;
+    NvBool bFeatureStretchVblankCapable;
+    NvBool bEnableDynamicGranularityPageArrays;
+    NvBool bClockBoostSupported;
+    NvBool bRouteDispIntrsToCPU;
+    NvU64  hostPageSize;
+} GspSystemInfo;
+
+
+#endif /* GSP_STATIC_CONFIG_H */

extra/nv_gpu_driver/vbios.h

@@ -0,0 +1,209 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#ifndef VBIOS_H
+#define VBIOS_H
+
+#include "gpu/vbios/bios_types.h"
+#define FALCON_APPLICATION_INTERFACE_ENTRY_ID_DMEMMAPPER     (0x4)
+
+typedef struct
+{
+    NvU32 signature;
+    NvU16 version;
+    NvU16 size;
+    NvU32 cmd_in_buffer_offset;
+    NvU32 cmd_in_buffer_size;
+    NvU32 cmd_out_buffer_offset;
+    NvU32 cmd_out_buffer_size;
+    NvU32 nvf_img_data_buffer_offset;
+    NvU32 nvf_img_data_buffer_size;
+    NvU32 printfBufferHdr;
+    NvU32 ucode_build_time_stamp;
+    NvU32 ucode_signature;
+    NvU32 init_cmd;
+    NvU32 ucode_feature;
+    NvU32 ucode_cmd_mask0;
+    NvU32 ucode_cmd_mask1;
+    NvU32 multiTgtTbl;
+} __attribute__((packed)) FALCON_APPLICATION_INTERFACE_DMEM_MAPPER_V3;
+
+#define FALCON_APPLICATION_INTERFACE_DMEM_MAPPER_V3_CMD_FRTS (0x15)
+#define FALCON_APPLICATION_INTERFACE_DMEM_MAPPER_V3_CMD_SB   (0x19)
+
+#define BIT_HEADER_ID                     0xB8FF
+#define BIT_HEADER_SIGNATURE              0x00544942  // "BIT\0"
+#define BIT_HEADER_SIZE_OFFSET            8
+
+struct __attribute__((packed)) BIT_HEADER_V1_00
+{
+    unsigned short Id;
+    unsigned int Signature;
+    unsigned short BCD_Version;
+    unsigned char HeaderSize;
+    unsigned char TokenSize;
+    unsigned char TokenEntries;
+    unsigned char HeaderChksum;
+};
+#define BIT_HEADER_V1_00_FMT "1w1d1w4b"
+typedef struct BIT_HEADER_V1_00 BIT_HEADER_V1_00;
+
+struct __attribute__((packed)) BIT_TOKEN_V1_00
+{
+    unsigned char TokenId;
+    unsigned char DataVersion;
+    unsigned short DataSize;
+    unsigned int DataPtr;
+};
+
+#define BIT_TOKEN_V1_00_SIZE_6     6U
+#define BIT_TOKEN_V1_00_SIZE_8     8U
+
+#define BIT_TOKEN_V1_00_FMT_SIZE_6 "2b2w"
+#define BIT_TOKEN_V1_00_FMT_SIZE_8 "2b1w1d"
+typedef struct BIT_TOKEN_V1_00 BIT_TOKEN_V1_00;
+
+#define BIT_TOKEN_BIOSDATA          0x42
+
+// structure for only version info from BIT_DATA_BIOSDATA_V1 and BIT_DATA_BIOSDATA_V2
+typedef struct
+{
+    unsigned int Version;     // BIOS Binary Version Ex. 5.40.00.01.12 = 0x05400001
+    unsigned char OemVersion;  // OEM Version Number  Ex. 5.40.00.01.12 = 0x12
+} __attribute__((packed)) BIT_DATA_BIOSDATA_BINVER;
+
+#define BIT_DATA_BIOSDATA_VERSION_1         0x1
+#define BIT_DATA_BIOSDATA_VERSION_2         0x2
+
+#define BIT_DATA_BIOSDATA_BINVER_FMT "1d1b"
+#define BIT_DATA_BIOSDATA_BINVER_SIZE_5    5
+
+#define BIT_TOKEN_FALCON_DATA       0x70
+
+typedef struct
+{
+    unsigned int FalconUcodeTablePtr;
+} __attribute__((packed)) BIT_DATA_FALCON_DATA_V2;
+
+#define BIT_DATA_FALCON_DATA_V2_4_FMT       "1d"
+#define BIT_DATA_FALCON_DATA_V2_SIZE_4      4
+
+typedef struct
+{
+    unsigned char Version;
+    unsigned char HeaderSize;
+    unsigned char EntrySize;
+    unsigned char EntryCount;
+    unsigned char DescVersion;
+    unsigned char DescSize;
+} __attribute__((packed)) FALCON_UCODE_TABLE_HDR_V1;
+
+#define FALCON_UCODE_TABLE_HDR_V1_VERSION   1
+#define FALCON_UCODE_TABLE_HDR_V1_SIZE_6    6
+#define FALCON_UCODE_TABLE_HDR_V1_6_FMT     "6b"
+
+typedef struct
+{
+    unsigned char ApplicationID;
+    unsigned char TargetID;
+    unsigned int DescPtr;
+} __attribute__((packed)) FALCON_UCODE_TABLE_ENTRY_V1;
+
+#define FALCON_UCODE_TABLE_ENTRY_V1_VERSION             1
+#define FALCON_UCODE_TABLE_ENTRY_V1_SIZE_6              6
+#define FALCON_UCODE_TABLE_ENTRY_V1_6_FMT               "2b1d"
+
+#define FALCON_UCODE_ENTRY_APPID_FIRMWARE_SEC_LIC       0x05
+#define FALCON_UCODE_ENTRY_APPID_FWSEC_DBG              0x45
+#define FALCON_UCODE_ENTRY_APPID_FWSEC_PROD             0x85
+
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_FLAGS_VERSION                0:0
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_FLAGS_VERSION_UNAVAILABLE    0x00
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_FLAGS_VERSION_AVAILABLE      0x01
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_FLAGS_RESERVED               1:1
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_FLAGS_ENCRYPTED              2:2
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_RESERVED                     7:3
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION                      15:8
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V1                   0x01
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V2                   0x02
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V3                   0x03
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V4                   0x04
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_SIZE                         31:16
+
+typedef struct
+{
+    unsigned int vDesc;
+} __attribute__((packed)) FALCON_UCODE_DESC_HEADER;
+#define FALCON_UCODE_DESC_HEADER_FORMAT   "1d"
+
+typedef struct {
+    FALCON_UCODE_DESC_HEADER Hdr;
+    unsigned int StoredSize;
+    unsigned int PKCDataOffset;
+    unsigned int InterfaceOffset;
+    unsigned int IMEMPhysBase;
+    unsigned int IMEMLoadSize;
+    unsigned int IMEMVirtBase;
+    unsigned int DMEMPhysBase;
+    unsigned int DMEMLoadSize;
+    unsigned short EngineIdMask;
+    unsigned char UcodeId;
+    unsigned char SignatureCount;
+    unsigned short SignatureVersions;
+    unsigned short Reserved;
+} FALCON_UCODE_DESC_V3;
+
+#define FALCON_UCODE_DESC_V3_SIZE_44    44
+#define FALCON_UCODE_DESC_V3_44_FMT     "9d1w2b2w"
+#define BCRT30_RSA3K_SIG_SIZE 384
+
+typedef struct
+{
+    NvU32 version;
+    NvU32 size;
+    NvU64 gfwImageOffset;
+    NvU32 gfwImageSize;
+    NvU32 flags;
+} __attribute__((packed)) FWSECLIC_READ_VBIOS_DESC;
+
+#define FWSECLIC_READ_VBIOS_STRUCT_FLAGS (2)
+
+typedef struct
+{
+    NvU32 version;
+    NvU32 size;
+    NvU32 frtsRegionOffset4K;
+    NvU32 frtsRegionSize;
+    NvU32 frtsRegionMediaType;
+} __attribute__((packed)) FWSECLIC_FRTS_REGION_DESC;
+
+#define FWSECLIC_FRTS_REGION_MEDIA_FB (2)
+#define FWSECLIC_FRTS_REGION_SIZE_1MB_IN_4K (0x100)
+
+typedef struct
+{
+    FWSECLIC_READ_VBIOS_DESC readVbiosDesc;
+    FWSECLIC_FRTS_REGION_DESC frtsRegionDesc;
+} __attribute__((packed)) FWSECLIC_FRTS_CMD;
+
+#endif /* VBIOS_H */

extra/nvpci/gpu_reset.sh

@@ -0,0 +1,2 @@
+GPU="$1"
+echo 1 | sudo tee /sys/bus/pci/devices/$GPU/reset 2>/dev/null

extra/onnx.py

@@ -1,11 +1,28 @@
 from types import SimpleNamespace
 from typing import Any, Sequence, cast, Literal, Callable
-import dataclasses, functools, io, math, types, warnings
+import dataclasses, functools, io, math, types, warnings, sys
 from tinygrad.tensor import Tensor, _broadcast_shape, ReductionStr
 from tinygrad.helpers import getenv, DEBUG, all_same, prod, flatten, make_tuple, argsort
-from tinygrad.dtype import DType, ConstType, dtypes, ImageDType
+from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype
 from tinygrad.device import is_dtype_supported, Device
 
+# https://github.com/onnx/onnx/blob/rel-1.17.0/onnx/onnx.proto3#L500-L544
+data_types: dict[int, DType] = {
+  1:dtypes.float32, 2:dtypes.uint8, 3:dtypes.int8, 4:dtypes.uint16, 5:dtypes.int16, 6:dtypes.int32, 7:dtypes.int64,
+  9:dtypes.bool, 10:dtypes.float16, 11:dtypes.double, 12:dtypes.uint32, 13:dtypes.uint64, 16:dtypes.bfloat16,
+}
+
+# https://github.com/onnx/onnx/blob/rel-1.17.0/onnx/onnx.proto3#L128-L145
+attribute_types: dict[int, Callable] = {
+  1: lambda a: float(a.f),
+  2: lambda a: int(a.i),
+  3: lambda a: a.s.data().tobytes().decode("utf8") if isinstance(a.s, Tensor) else a.s.decode("utf8"),
+  4: lambda a: buffer_parse(a.t),
+  6: lambda a: tuple(float(x) for x in a.floats),
+  7: lambda a: tuple(int(x) for x in a.ints),
+  8: lambda a: tuple(x.data().tobytes().decode("utf8") for x in a.strings)
+}
+
 # ***** protobuf parsing ******
 from onnx import AttributeProto, ModelProto, TensorProto, TypeProto, helper
 import numpy as np
@@ -15,40 +32,19 @@ def has_field(onnx_type: TypeProto|SimpleNamespace, field):
   return hasattr(onnx_type, field)
 
 def dtype_parse(onnx_dtype: int, fallback_context: str | None = None) -> DType:
-  supported: dict[int, DType] = {
-    TensorProto.FLOAT:dtypes.float32, TensorProto.UINT8:dtypes.uint8, TensorProto.INT8:dtypes.int8,
-    TensorProto.UINT16:dtypes.uint16, TensorProto.INT16:dtypes.int16, TensorProto.INT32:dtypes.int32, TensorProto.INT64:dtypes.int64,
-    TensorProto.BOOL:dtypes.bool, TensorProto.FLOAT16:dtypes.float32, TensorProto.DOUBLE:dtypes.double, TensorProto.UINT32:dtypes.uint32,
-    TensorProto.UINT64:dtypes.uint64, TensorProto.BFLOAT16:dtypes.bfloat16,
-  }
-  unsupported = {
-    TensorProto.UNDEFINED, TensorProto.STRING, TensorProto.COMPLEX64, TensorProto.COMPLEX128, TensorProto.FLOAT8E4M3FN, TensorProto.FLOAT8E4M3FNUZ,
-    TensorProto.FLOAT8E5M2, TensorProto.FLOAT8E5M2FNUZ, TensorProto.UINT4, TensorProto.INT4
-  }
-  if onnx_dtype in unsupported: raise NotImplementedError(f"onnx dtype {TensorProto.DataType.Name(onnx_dtype)} is not supported")
-  if is_dtype_supported(dtype := supported[onnx_dtype]): return dtype
+  if onnx_dtype not in data_types: raise NotImplementedError(f"onnx dtype id {onnx_dtype} is not supported")
+  if is_dtype_supported(dtype := data_types[onnx_dtype]): return dtype
   # if fallback_context is provided, we can fall back to a default dtype
   if fallback_context is not None:
-    default_dtype = dtypes.float
+    default_dtype = dtypes.default_int if dtypes.is_int(dtype) else dtypes.default_float
     warnings.warn(f"dtype {dtype} on {Device.DEFAULT} from {fallback_context} is not supported, falling back to {default_dtype}")
+    assert is_dtype_supported(default_dtype), f"dtype {default_dtype} must be supported on {Device.DEFAULT}"
     return default_dtype
   raise RuntimeError(f"dtype {dtype} on device {Device.DEFAULT} is not supported")
 
 def attribute_parse(onnx_attribute: AttributeProto):
-  supported: dict[AttributeProto.AttributeType, Callable[[AttributeProto], Any]] = {
-    AttributeProto.FLOAT: lambda a: float(a.f), AttributeProto.INT: lambda a: int(a.i),
-    AttributeProto.STRING: lambda a: a.s.data().tobytes().decode("utf8") if isinstance(a.s, Tensor) else a.s.decode("utf8"),
-    AttributeProto.TENSOR: lambda a: buffer_parse(a.t),
-    AttributeProto.FLOATS: lambda a: tuple(float(x) for x in a.floats), AttributeProto.INTS: lambda a: tuple(int(x) for x in a.ints),
-    AttributeProto.STRINGS: lambda a: tuple(x.data().tobytes().decode("utf8") for x in a.strings)
-  }
-  unsupported = {
-    AttributeProto.UNDEFINED, AttributeProto.GRAPH, AttributeProto.SPARSE_TENSOR, AttributeProto.TYPE_PROTO, AttributeProto.TENSORS,
-    AttributeProto.GRAPHS, AttributeProto.SPARSE_TENSORS, AttributeProto.TYPE_PROTOS
-  }
-  if onnx_attribute.type in unsupported:
-    raise NotImplementedError(f"attribute with type {AttributeProto.AttributeType.Name(onnx_attribute.type)} is not supported")
-  return supported[onnx_attribute.type](onnx_attribute)
+  if onnx_attribute.type not in attribute_types: raise NotImplementedError(f"attribute type {onnx_attribute.type} is not supported")
+  return attribute_types[onnx_attribute.type](onnx_attribute)
 
 def buffer_parse(onnx_tensor: TensorProto) -> Tensor:
   if onnx_tensor.string_data: raise NotImplementedError("Parsing for buffer with string data is not implemented.")
@@ -65,13 +61,15 @@ def buffer_parse(onnx_tensor: TensorProto) -> Tensor:
   if has_field(onnx_tensor, "raw_data"):
     raw_data = onnx_tensor.raw_data
     if not isinstance(raw_data, Tensor): raw_data = Tensor(raw_data)
-    if onnx_tensor.data_type == TensorProto.FLOAT16:
+    if not is_dtype_supported(data_types[onnx_tensor.data_type]):
       np_buffer = np.frombuffer(raw_data.data().tobytes(),
                                 dtype=helper.tensor_dtype_to_np_dtype(onnx_tensor.data_type)).copy().reshape(shape)
       if np_buffer.size == 1: return Tensor(np_buffer.item(), dtype=dtype).reshape(shape)
       return Tensor(np_buffer, dtype=dtype)
     ret = raw_data.bitcast(dtype).reshape(shape).to(Device.DEFAULT)
-    if shape == (): ret = Tensor(ret.item(), dtype=dtype).reshape(shape)
+    if shape == ():
+      if ret.dtype is dtypes.float16 and sys.version_info < (3, 12): ret = ret.cast(dtypes.float32)
+      ret = Tensor(ret.item(), dtype=dtype).reshape(shape)
     return ret
   return Tensor(None)
 
@@ -84,7 +82,7 @@ def type_parse(onnx_type: TypeProto):
   if has_field(elem_type, "tensor_type"):
     shape = tuple(getattr(d, "dim_param", None) or getattr(d, "dim_value") for d in elem_type.tensor_type.shape.dim) \
       if has_field(elem_type.tensor_type, "shape") else None # test_identity_sequence_cpu
-    dtype = dtype_parse(elem_type.tensor_type.elem_type, "input type spec parse")
+    dtype = data_types[elem_type.tensor_type.elem_type]
     return OnnxValue(shape, dtype, is_optional, is_sequence)
   raise RuntimeError(f"TypeProto was not parsed properly: {onnx_type=}")
 
@@ -151,13 +149,15 @@ class OnnxRunner:
 
   def _parse_input(self, name: str, value: Any, spec: OnnxValue):
     if spec.is_optional and value is None: return None
-    # TODO: need true float16 for dtype checking
     if spec.is_sequence:
       if not isinstance(value, Sequence): raise RuntimeError(f"input {name} received {value}, expected a sequence type")
       sequence = [Tensor(v, dtype=spec.dtype, requires_grad=self.is_training) if not isinstance(v, Tensor) else v for v in value]
       if not all_same(tuple(t.shape for t in sequence)): raise RuntimeError(f"Shapes for input {name} sequence must be homogeneous")
+      if not all(t.dtype is spec.dtype for t in sequence): warnings.warn(f"Dtypes for input {name} sequence aren't all {spec.dtype}")
       return sequence
-    tensor = Tensor(value, dtype=spec.dtype, requires_grad=self.is_training) if not isinstance(value, Tensor) else value
+    dtype = _from_np_dtype(value.dtype) if str(type(value)) == "<class 'numpy.ndarray'>" else spec.dtype
+    tensor = Tensor(value, dtype=dtype, requires_grad=self.is_training) if not isinstance(value, Tensor) else value
+    if tensor.dtype is not spec.dtype: warnings.warn(f"input {name} has mismatch on dtype. Expected {spec.dtype}, received {tensor.dtype}.")
     for dim, (onnx_dim, user_dim_input) in enumerate(zip(spec.shape, tensor.shape, strict=True)):
       if isinstance(onnx_dim, str):
         onnx_dim = self.variable_dims[onnx_dim] if onnx_dim in self.variable_dims else self.variable_dims.setdefault(onnx_dim, int(user_dim_input))
@@ -175,8 +175,8 @@ class OnnxRunner:
       return real_fxn(*inps, **opts)
     raise NotImplementedError(f"{op=} not supported")
 
-  def get_empty_input_data(self, device:str|None=None) -> dict[str, Tensor]:
-    return {name:Tensor.empty(*spec.shape, device=device, dtype=spec.dtype) for name, spec in self.graph_inputs.items()}
+  def get_empty_input_data(self, device:str|None=None, dtype:DType|None=None) -> dict[str, Tensor]:
+    return {name:Tensor.empty(*spec.shape, device=device, dtype=dtype or spec.dtype) for name, spec in self.graph_inputs.items()}
 
   def __call__(self, inputs:dict[str, Any], debug=debug):
     for name, input_spec in self.graph_inputs.items():
@@ -325,7 +325,7 @@ def get_onnx_ops():
   def Binarizer(x:Tensor, threshold:float=0.0): return (x > threshold).float()
 
   # ***** Unary Ops (broadcasted) *****
-  def Add(x:Tensor,y:Tensor, broadcast=None, axis=None): return x + y if x.dtype == dtypes.float or isinstance(x.dtype, ImageDType) else (x + y).cast(x.dtype)
+  def Add(x:Tensor,y:Tensor, broadcast=None, axis=None): return x + y
   def Sub(x:Tensor|int,y:Tensor): return x - y # some test has input as int
   def Div(x:Tensor,y:Tensor): return x.div(y, rounding_mode='trunc' if dtypes.is_int(x.dtype) else None)
   def Less(x:Tensor,y:Tensor): return x < y

extra/onnx_parser.py

@@ -25,6 +25,7 @@ class PBType: FLOAT = 1; INT = 2; STRING = 3; FLOATS = 4; INTS = 5; STRINGS = 6;
 PB_INFOS = {
   "OperatorSetIdProto": {1: ("domain", PBType.STRING), 2: ("version", PBType.INT)},
   "StringStringEntryProto": {1: ("key", PBType.STRING), 2: ("value", PBType.STRING)},
+  # TODO: support uint64 parsing (11: "uint64_data") and double parsing (10: "double_data")
   "TensorProto": {1: ("dims", PBType.INT, True), 2: ("data_type", PBType.INT), 4: ("float_data", PBType.FLOATS),
     13: ("external_data", PBType.SUB, True, "StringStringEntryProto"), 14: ("data_location", PBType.INT),
     5: ("int32_data", PBType.INTS), 7: ("int64_data", PBType.INTS), 8: ("name", PBType.STRING), 9: ("raw_data", PBType.BYTES)},

extra/optimization/extract_dataset.py

@@ -6,8 +6,8 @@ from test.external.process_replay.process_replay import _pmap
 LOGOPS = os.getenv("LOGOPS", "/tmp/sops")
 
 def extract_ast(*args) -> None:
-  open(LOGOPS, "a").write(str(args[0]).replace("\n", "").replace(" ", "")+"\n")
+  open(LOGOPS, "a").write(str(args[1]).replace("\n", "").replace(" ", "")+"\n")
   return None
 
 if __name__ == "__main__":
-  _pmap("kernel", extract_ast)
+  _pmap({"get_program":extract_ast})

extra/optimization/generate_dataset.sh

@@ -6,24 +6,12 @@ export CAPTURE_PROCESS_REPLAY=1
 rm $LOGOPS
 test/external/process_replay/reset.py
 
-python3 -m pytest -n=auto test/ --ignore=test/unit --durations=20
-STEPS=3 python3 examples/hlb_cifar10.py
-WINO=1 STEPS=3 python3 examples/hlb_cifar10.py
-python3 examples/stable_diffusion.py --noshow
-python3 examples/llama.py --prompt "hello" --count 5
-python3 examples/gpt2.py --count 5
-HALF=1 python3 examples/gpt2.py --count 5
-python3 examples/beautiful_mnist.py
-python3 examples/beautiful_cartpole.py
-python3 examples/mlperf/model_spec.py
-python3 examples/yolov8.py ./test/models/efficientnet/Chicken.jpg
-examples/openpilot/go.sh
-JIT=2 BIG=1 MPS=1 pytest -n=auto test/ --ignore=test/test_fusion_op.py --ignore=test/test_linearizer_failures.py --ignore=test/test_gc.py --ignore=test/test_speed_v_torch.py --ignore=test/test_jit.py
-JIT=2 BIG=1 MPS=1 python -m pytest test/test_gc.py
-JIT=2 BIG=1 MPS=1 python -m pytest test/test_jit.py
-JIT=2 BIG=1 MPS=1 python -m pytest test/test_speed_v_torch.py
+CI=1 python3 -m pytest -n=auto test/test_ops.py test/test_nn.py test/test_winograd.py test/models/test_real_world.py --durations=20
+GPU=1 python3 -m pytest test/test_tiny.py
 
 # extract, sort and uniq
 extra/optimization/extract_dataset.py
 sort -u /tmp/ops > /tmp/sops
-ls -lh /tmp/ops /tmp/sops	
+ls -lh /tmp/ops /tmp/sops
+# gzip -k /tmp/sops
+# mv /tmp/sops.gz extra/datasets/

extra/optimization/helpers.py

@@ -5,6 +5,7 @@ from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.dtype import dtypes, PtrDType
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
+from tinygrad.helpers import getenv
 inf, nan = float('inf'), float('nan')
 UOps = Ops
 
@@ -26,7 +27,7 @@ from tinygrad.helpers import dedup, DEBUG
 def load_worlds(filter_reduce=True, filter_noimage=True, filter_novariable=True):
   fn = Path(__file__).parent.parent / "datasets/sops.gz"
   ast_strs = dedup(gzip.open(fn).read().decode('utf-8').strip().split("\n"))
-  assert len(ast_strs) > 5000, f"dataset size = {len(ast_strs)} is too small"
+  assert len(ast_strs) >= getenv("MIN_ASTS", 1000), f"dataset size = {len(ast_strs)} is too small"
   if DEBUG >= 1: print(f"loaded {len(ast_strs)=} before filters")
   if filter_reduce: ast_strs = [x for x in ast_strs if "REDUCE_AXIS" in x]
   if filter_noimage: ast_strs = [x for x in ast_strs if "dtypes.image" not in x]

extra/perfetto/to_perfetto.py

@@ -0,0 +1,38 @@
+import sys, pickle, decimal, json
+from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileRangeEvent, ProfileGraphEvent
+from tinygrad.helpers import tqdm, temp
+
+devices:dict[str, tuple[decimal.Decimal, decimal.Decimal, int]] = {}
+def prep_ts(device:str, ts:decimal.Decimal, is_copy): return int(decimal.Decimal(ts) + devices[device][is_copy])
+def dev_to_pid(device:str, is_copy=False): return {"pid": devices[device][2], "tid": int(is_copy)}
+def dev_ev_to_perfetto_json(ev:ProfileDeviceEvent):
+  devices[ev.device] = (ev.comp_tdiff, ev.copy_tdiff if ev.copy_tdiff is not None else ev.comp_tdiff, len(devices))
+  return [{"name": "process_name", "ph": "M", "pid": dev_to_pid(ev.device)['pid'], "args": {"name": ev.device}},
+          {"name": "thread_name", "ph": "M", "pid": dev_to_pid(ev.device)['pid'], "tid": 0, "args": {"name": "COMPUTE"}},
+          {"name": "thread_name", "ph": "M", "pid": dev_to_pid(ev.device)['pid'], "tid": 1, "args": {"name": "COPY"}}]
+def range_ev_to_perfetto_json(ev:ProfileRangeEvent):
+  return [{"name": ev.name, "ph": "X", "ts": prep_ts(ev.device, ev.st, ev.is_copy), "dur": float(ev.en-ev.st), **dev_to_pid(ev.device, ev.is_copy)}]
+def graph_ev_to_perfetto_json(ev:ProfileGraphEvent, reccnt):
+  ret = []
+  for i,e in enumerate(ev.ents):
+    st, en = ev.sigs[e.st_id], ev.sigs[e.en_id]
+    ret += [{"name": e.name, "ph": "X", "ts": prep_ts(e.device, st, e.is_copy), "dur": float(en-st), **dev_to_pid(e.device, e.is_copy)}]
+    for dep in ev.deps[i]:
+      d = ev.ents[dep]
+      ret += [{"ph": "s", **dev_to_pid(d.device, d.is_copy), "id": reccnt+len(ret), "ts": prep_ts(d.device, ev.sigs[d.en_id], d.is_copy), "bp": "e"}]
+      ret += [{"ph": "f", **dev_to_pid(e.device, e.is_copy), "id": reccnt+len(ret)-1, "ts": prep_ts(e.device, st, e.is_copy), "bp": "e"}]
+  return ret
+def to_perfetto(profile:list[ProfileEvent]):
+  # Start json with devices.
+  prof_json = [x for ev in profile if isinstance(ev, ProfileDeviceEvent) for x in dev_ev_to_perfetto_json(ev)]
+  for ev in tqdm(profile, desc="preparing profile"):
+    if isinstance(ev, ProfileRangeEvent): prof_json += range_ev_to_perfetto_json(ev)
+    elif isinstance(ev, ProfileGraphEvent): prof_json += graph_ev_to_perfetto_json(ev, reccnt=len(prof_json))
+  return {"traceEvents": prof_json}
+
+if __name__ == "__main__":
+  fp = sys.argv[1]
+  with open(fp, "rb") as f: profile = pickle.load(f)
+  ret = to_perfetto(profile)
+  with open(fp:=temp("perfetto.json", append_user=True), "w") as f: json.dump(ret, f)
+  print(f"Saved perfetto output to {fp}. You can use upload this to the perfetto UI or Chrome devtools.")

extra/remu/README.md

@@ -22,7 +22,7 @@ This will produce a binary in the `extra/remu/target/release` directory.
 
 The latest binaries are released in https://github.com/Qazalin/remu/releases. Alternatively, you can [build locally](#build-locally).
 
-Tinygrad does not yet output RDNA3 kernels directly. You can either install comgr or use `AMD_LLVM=1` if you have [LLVM@19](https://github.com/tinygrad/tinygrad/blob/e2ed673c946c8f1774d816c75e52a994c2dd8a88/.github/actions/setup-tinygrad/action.yml#L208).
+Tinygrad does not yet output RDNA3 kernels directly. You can either install comgr or use `AMD_LLVM=1` (default) if you have [LLVM@19](https://github.com/tinygrad/tinygrad/blob/e2ed673c946c8f1774d816c75e52a994c2dd8a88/.github/actions/setup-tinygrad/action.yml#L208).
 
 `PYTHONPATH="." MOCKGPU=1 AMD=1 python test/test_tiny.py TestTiny.test_plus` runs an emulated RDNA3 kernel with Remu.
 

extra/torch_backend/backend.py

@@ -59,6 +59,7 @@ view_ops = {
   "aten.squeeze.dim": Tensor.squeeze,
   "aten.unsqueeze": Tensor.unsqueeze,
   "aten.detach": Tensor.detach,
+  "aten.select.int": lambda self, dim, idx: self[(slice(None),) * (dim%self.ndim) + (idx,)],
 }
 
 for k,v in view_ops.items(): torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrap_view_op(v))
@@ -368,6 +369,7 @@ decomps = [
   aten.threshold,
   aten.nll_loss_forward,
   aten.nll_loss_backward,
+  aten.nll_loss2d_backward,
   # AttributeError: 'int' object has no attribute '_broadcasted'
   aten.sigmoid_backward,
   aten.tanh_backward,
@@ -376,6 +378,7 @@ decomps = [
   aten.softshrink,
   aten.hardshrink,
   aten.log_sigmoid_forward,
+  aten.log_sigmoid_backward,
   aten.isneginf,
   aten.isposinf,
   aten.nan_to_num,

extra/torch_backend/test.py

@@ -206,5 +206,16 @@ class TestTorchBackend(unittest.TestCase):
       X.cpu(), Y.cpu()
       self.assertLessEqual(GlobalCounters.global_ops, 10_000_000)
 
+  def _test_diagonal(self, *shape):
+    a = torch.randn(*shape, dtype=torch.float32, device=device)
+    ref = np.diagonal(a.cpu().numpy(), axis1=-2, axis2=-1)
+    diag = torch.linalg.diagonal(a)
+    np.testing.assert_equal(diag.cpu().numpy(), ref)
+    np.testing.assert_equal(diag[-1].cpu().numpy(), ref[-1])
+
+  def test_diagonal_cube(self): self._test_diagonal(3, 3, 3)
+  def test_diagonal_rectangular(self): self._test_diagonal(4, 5, 6)
+  def test_diagonal_4d(self): self._test_diagonal(2, 3, 4, 5)
+
 if __name__ == "__main__":
   unittest.main()

extra/upcasted_warps.py

@@ -0,0 +1,73 @@
+# play with upcasted warps
+from tinygrad import Tensor, Device
+from tinygrad.uop.ops import KernelInfo
+from tinygrad.opt import get_optimized_ast
+from tinygrad.opt.kernel import OptOps, Opt
+from tinygrad.engine.realize import get_program
+
+if __name__ == "__main__":
+  renderer = Device.default.renderer
+  N = 64
+
+  """
+  a = Tensor.empty(N,N)
+
+  out = (a + 1) #.sum(axis=2)
+  ast = out.schedule()[-1].ast
+  opts = tuple()
+  opts += (Opt(OptOps.UPCAST, 0, 32),)
+  ast = ast.replace(arg=KernelInfo(opts_to_apply=opts))
+  ast = get_optimized_ast(ast, renderer)
+  prg = get_program(ast, renderer)
+  print(prg.src)
+  """
+
+  # how you split the store determines everything if you don't allow cross warp comms.
+  # actually not everything, there's also the split before the horizontal (unrolled) reduces
+
+  # new flow
+  #  - pull out any dimensions from the store that you want to upcast.
+  #  - decide how you want to assign them to registers. GPUs have a 512-byte memory LOAD/STORE which loads into 4 regs. see BUFFER_LOAD_B128
+  #    - the loads and stores can be shuffled, but only in restrictive ways. in kernels without reduces, the store determines everything
+  #    - it loads 16 bytes from up 32 different places = 512 bytes
+  #  - in kernels with reduces, you now have more flexibility. the final target of the reduce must be what is stored
+  #    - warp dimensions can be in the reduce (this is GROUP)
+
+  # every dimension can be assigned to <global, local, loop, upcast, warp>
+
+  """
+  out = a.sum(axis=1)
+  ast = out.schedule()[-1].ast
+  opts = tuple()
+  opts += (Opt(OptOps.UPCAST, 0, 8),)
+  opts += (Opt(OptOps.UNROLL, 0, 8),)
+  ast = ast.replace(arg=KernelInfo(opts_to_apply=opts))
+  ast = get_optimized_ast(ast, renderer)
+  prg = get_program(ast, renderer)
+  print(prg.src)
+
+  out = a.sum(axis=1)
+  ast = out.schedule()[-1].ast
+  opts = tuple()
+  opts += (Opt(OptOps.UNROLL, 0, 8),)
+  opts += (Opt(OptOps.UPCAST, 0, 8),)
+  ast = ast.replace(arg=KernelInfo(opts_to_apply=opts))
+  ast = get_optimized_ast(ast, renderer)
+  prg = get_program(ast, renderer)
+  print(prg.src)
+  """
+
+  # gemm
+  b = Tensor.empty(N,N)
+  # metal TC
+  #opts = (Opt(OptOps.UPCAST, 0, 2), # not the warp
+  #        Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UPCAST, 1, 2), Opt(OptOps.UPCAST, 1, 2),
+  #        Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UPCAST, 1, 2))
+  # new TC should just be able to extract from this and swizzle as needed
+  opts = (Opt(OptOps.UPCAST, 0, 8), Opt(OptOps.UPCAST, 1, 8), Opt(OptOps.UNROLL, 0, 8))
+  c = (a@b)
+  ast = c.schedule()[-1].ast
+  ast = ast.replace(arg=KernelInfo(opts_to_apply=opts))
+  ast = get_optimized_ast(ast, renderer)
+  prg = get_program(ast, renderer)
+  print(prg.src)

mkdocs.yml

@@ -22,6 +22,7 @@ nav:
     - Runtime: runtime.md
   - Developer:
     - Intro: developer/developer.md
+    - Layout: developer/layout.md
     - Speed: developer/speed.md
     - UOp: developer/uop.md
     - Grouper:

ruff.toml

@@ -36,9 +36,17 @@ line-length = 150
 
 exclude = [
   "docs/",
-  "examples/",
   "extra/",
   "tinygrad/runtime/autogen",
   "test/external/mlperf_resnet",
   "test/external/mlperf_unet3d",
 ]
+
+# detect unused imports in examples
+[lint.per-file-ignores]
+"examples/**/*.py" = [
+  "W6", "E71", "E72", "E112", "E113", "E203", "E272", "E275",
+  "E303", "E304", "E501", "E702", "E703", "E731", "W191",
+  "W291", "W293", "UP039", "C416", "RET506", "RET507", "A",
+  "FURB110", "RUF018", "F541", "F841"
+]

setup.py

@@ -25,9 +25,10 @@ setup(name='tinygrad',
       long_description=long_description,
       long_description_content_type='text/markdown',
       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.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop', 'tinygrad.opt'],
-      package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'perfetto.html', 'assets/**/*', 'js/*']},
+                  'tinygrad.renderer', 'tinygrad.engine', 'tinygrad.viz', 'tinygrad.runtime', 'tinygrad.runtime.support', 'tinygrad.kernelize',
+                  'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop', 'tinygrad.opt',
+                  'tinygrad.runtime.support.nv'],
+      package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'assets/**/*', 'js/*']},
       classifiers=[
         "Programming Language :: Python :: 3",
         "License :: OSI Approved :: MIT License"

test/external/external_benchmark_llama_schedule.py

@@ -0,0 +1,35 @@
+from tinygrad import nn, Tensor, Device, dtypes
+from tinygrad.helpers import Timing
+
+from extra.models.llama import Transformer
+from examples.llama3 import MODEL_PARAMS
+
+if __name__ == "__main__":
+  Device.DEFAULT = "NULL"
+  Tensor.training = True
+  #model_size = "8B"
+  model_size = "405B"
+
+  with Timing("total "):
+    with Timing("***** create model in    "):
+      model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=nn.Linear, embedding=nn.Embedding,
+                          max_context=1024, jit=True, disable_kv_cache=True)
+
+    with Timing("***** fake state in      "):
+      Tensor.realize(*[p.assign(Tensor.empty(*p.shape, device=p.device, dtype=p.dtype)) for p in nn.state.get_parameters(model)])
+
+    with Timing("***** create optim in    "):
+      opt = nn.optim.AdamW(nn.state.get_parameters(model))
+
+    with Timing("***** run model in       "):
+      toks = Tensor.empty(1, 1024, dtype=dtypes.int)
+      out = model(toks, 0, temperature=float('nan'))
+
+    with Timing("***** backward in        "):
+      out.mean().backward()
+
+    with Timing("***** realize in        "):
+      out.realize()
+
+    with Timing("***** step in           "):
+      opt.step()

test/external/external_fuzz_ampt.py

@@ -1,4 +1,5 @@
 import random
+from typing import Optional
 from tinygrad.helpers import round_up
 from tinygrad.runtime.support.am.amdev import AMPageTableTraverseContext
 from test.external.external_test_am import helper_read_entry_components, FakeAM
@@ -59,7 +60,7 @@ class AMPTFuzzer:
 
     return True
 
-  def random_alloc(self):
+  def random_alloc(self) -> Optional[int]:
     if self.total_size - self.alloc_payload < self.min_alloc_size: return None
 
     size = random.randint(self.min_alloc_size, min(self.max_alloc_size, self.total_size - self.alloc_payload))

test/external/external_fuzz_tlsf.py

@@ -1,7 +1,7 @@
 import random
 from typing import Dict, Optional
 from tinygrad.helpers import getenv
-from tinygrad.runtime.support.allocator import TLSFAllocator
+from tinygrad.runtime.support.memory import TLSFAllocator
 
 class AllocatorFuzzer:
   def __init__(self, total_size):
@@ -30,6 +30,8 @@ class AllocatorFuzzer:
     return True
 
   def random_alloc(self) -> Optional[int]:
+    if self.total_size - self.alloc_payload < self.min_alloc_size: return None
+
     size = random.randint(self.min_alloc_size, min(self.max_alloc_size, self.total_size - self.alloc_payload))
 
     try:

test/external/external_llama_eval.py

@@ -1,102 +0,0 @@
-from lm_eval.base import BaseLM
-from lm_eval import evaluator, tasks
-import torch, json, argparse
-
-from examples.llama import LLaMa
-from tinygrad.tensor import Tensor
-from tinygrad import Device
-
-class LLaMaAdaptor(BaseLM):
-  def __init__(
-    self,
-    model_size="7B",
-    model_gen=1,
-    device="",
-    quantize=False,
-    batch_size=1,
-    max_batch_size=1,
-    do_sample=False,
-    temperature=1.0,
-    checkpoint_path="",
-    tokenizer_path="",
-  ):
-    super().__init__()
-
-    if batch_size is None:
-      batch_size = 1
-    self.do_sample = do_sample
-    self.temperature = temperature
-    self._device = device
-
-    assert isinstance(model_gen, int)
-    assert isinstance(model_size, str)
-    assert isinstance(batch_size, int)
-    assert isinstance(checkpoint_path, str)
-    assert isinstance(tokenizer_path, str)
-
-    self.llama = LLaMa.build(checkpoint_path, tokenizer_path, model_gen, model_size, quantize)
-
-  @classmethod
-  def create_from_arg_string(cls, arg_string, additional_config=None):
-    kwargs = {el.split("=")[0]: el.split("=")[1] for el in arg_string.split(",")}
-    return cls(**kwargs, **additional_config)
-
-  @property
-  def eot_token_id(self):
-    # we use EOT because end of *text* is more accurate for what we're doing than end of *sentence*
-    return self.llama.tokenizer.eos_id()
-
-  @property
-  def max_length(self):
-    return 1024
-
-  @property
-  def max_gen_toks(self):
-    return 256
-
-  @property
-  def batch_size(self):
-    return 1
-
-  @property
-  def device(self):
-    return self._device
-
-  def tok_encode(self, string: str):
-    return [self.llama.tokenizer.bos_id()] + self.llama.tokenizer.encode(string)
-
-  def tok_decode(self, tokens):
-    return self.llama.tokenizer.decode(tokens)
-
-  def _model_call(self, inps):
-    return torch.Tensor(self.llama.model(Tensor(inps.numpy()), 0).numpy())
-
-  def greedy_until(self, requests):
-    continuations = []
-    for request in requests:
-      prompt, until = request[0], request[1]['until']
-      output = self.llama.greedy_until(prompt, until, max_length=128, temperature=0.0)
-      continuations.append(output[len(prompt):])
-    return continuations
-
-  def _model_generate(self, context, max_length, eos_token_id):
-    raise NotImplementedError()
-
-if __name__ == '__main__':
-  print(f"using {Device.DEFAULT} backend")
-
-  parser = argparse.ArgumentParser(description='Run LLaMA evals in tinygrad', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-  parser.add_argument('--size', type=str, default="7B", help="Size of model to use [7B, 13B, 30B, 65B] for Gen 1, [7B, 13B] for Gen 2")
-  parser.add_argument('--gen', type=int, default="1", help="Generation of the model to use [1, 2]")
-  parser.add_argument('--quantize', action='store_true', help="Quantize the weights to int8 in memory")
-  parser.add_argument('--eval', type=str, default="arc_easy", help="Run in evaluation mode")
-  parser.add_argument('--limit', type=int, default=None, help="Limit tests in eval")
-  parser.add_argument('--weights', type=str, default="./weights/LLaMa/", help="Location of the weights")
-  parser.add_argument('--tokenizer', type=str, default="./weights/LLaMa/tokenizer.model", help="Location of the tokenizer")
-  args = parser.parse_args()
-
-  # run eval and exit
-  adaptor = LLaMaAdaptor(model_gen=args.gen, model_size=args.size, quantize=args.quantize,
-                         checkpoint_path=args.weights, tokenizer_path=args.tokenizer, device="cpu")
-  results = evaluator.evaluate(adaptor, tasks.get_task_dict(args.eval.split(",")), False, 0, args.limit)
-  print(json.dumps(results, indent=2))

test/external/external_model_benchmark.py

@@ -7,7 +7,7 @@ import onnxruntime as ort
 from onnx2torch import convert
 from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 from tinygrad.helpers import OSX, DEBUG, fetch, getenv
-from tinygrad import Tensor, Device
+from tinygrad import Tensor, Device, dtypes
 
 MODELS = {
   "resnet50": "https://github.com/onnx/models/raw/main/validated/vision/classification/resnet/model/resnet50-caffe2-v1-9.onnx",
@@ -27,6 +27,7 @@ MODELS = {
   # really slow
   # "resnet18": "https://github.com/onnx/models/raw/main/archive/vision/classification/resnet/model/resnet18-v2-7.onnx",
 }
+half_models = ["openpilot", "commavq"]
 
 CSV = {}
 open_csv = None
@@ -54,7 +55,6 @@ def benchmark_model(m, devices, validate_outs=False):
   excluded = {inp.name for inp in onnx_model.graph.initializer}
   input_shapes = {inp.name:tuple(x.dim_value if hasattr(x, "dim_value") and x.dim_value != 0 else 1 for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input if inp.name not in excluded}  # noqa: E501
   input_types = {inp.name: tensor_dtype_to_np_dtype(inp.type.tensor_type.elem_type) for inp in onnx_model.graph.input if inp.name not in excluded}
-  #input_types = {k:v if v!=np.float16 else np.float32 for k,v in input_types.items()}  # cast
   np_inputs = {k:torch.randn(shp).numpy().astype(input_types[k]) for k,shp in input_shapes.items()}
   assert len(input_shapes) < 30, f"too many input shapes {len(input_shapes)}"
 
@@ -106,7 +106,12 @@ def benchmark_model(m, devices, validate_outs=False):
     for device in devices:
       rtol, atol = 2e-3, 2e-3  # tolerance for fp16 models
       Device.DEFAULT = device
-      inputs = {k:Tensor(inp) for k,inp in np_inputs.items()}
+      # force half inputs to float for numerical stability when validating
+      # this will reply on automatic dtype promotion for converting half weights inside the graph
+      if m in half_models:
+        inputs = {k:Tensor(inp, dtype=dtypes.float32) if inp.dtype == np.float16 else Tensor(inp) for k,inp in np_inputs.items()}
+      else:
+        inputs = {k:Tensor(inp) for k,inp in np_inputs.items()}
       tinygrad_model = OnnxRunner(onnx_model)
       tinygrad_out = tinygrad_model(inputs)
 

test/external/external_test_am.py

@@ -1,7 +1,8 @@
 import unittest
-from tinygrad.runtime.support.am.amdev import AMMemoryManager, AMPageTableTraverseContext
+from tinygrad.runtime.support.am.amdev import AMMemoryManager, AMPageTableEntry
 from tinygrad.runtime.support.am.ip import AM_GMC
 from tinygrad.runtime.support.hcq import MMIOInterface
+from tinygrad.runtime.support.memory import PageTableTraverseContext
 from tinygrad.runtime.autogen.am import am
 from tinygrad.helpers import mv_address
 
@@ -23,7 +24,9 @@ class FakeAM:
     self.vram_mv = memoryview(bytearray(4 << 30))
     self.vram = MMIOInterface(mv_address(self.vram_mv), self.vram_mv.nbytes)
     self.gmc = FakeGMC(self)
-    self.mm = AMMemoryManager(self, vram_size=4 << 30)
+    self.mm = AMMemoryManager(self, 4 << 30, boot_size=(32 << 20), pt_t=AMPageTableEntry, pte_cnt=[512, 512, 512, 512],
+      pte_covers=[(1 << ((9 * (3-lv)) + 12)) for lv in range(4)], first_lv=am.AMDGPU_VM_PDB1, first_page_lv=am.AMDGPU_VM_PDB2,
+      va_base=AMMemoryManager.va_allocator.base)
     self.is_booting = False
     self.ip_ver = {am.GC_HWIP: (11, 0, 0)}
   def paddr2cpu(self, paddr:int) -> int: return paddr + mv_address(self.vram)
@@ -65,7 +68,7 @@ class TestAMPageTable(unittest.TestCase):
       exteranl_va = va + AMMemoryManager.va_allocator.base
       mm.map_range(vaddr=exteranl_va, size=sz, paddrs=[(va, sz)])
 
-      ctx = AMPageTableTraverseContext(self.d[0], mm.root_page_table, exteranl_va)
+      ctx = PageTableTraverseContext(self.d[0], mm.root_page_table, exteranl_va)
       results = list(ctx.next(sz))
 
       total_covered = 0
@@ -126,7 +129,7 @@ class TestAMPageTable(unittest.TestCase):
       # Finally can map and check paddrs
       mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(0xdead0000, 0x1000), (0xdead1000, 0xff000)])
 
-      ctx = AMPageTableTraverseContext(self.d[0], mm0.root_page_table, exteranl_va + 0x2000)
+      ctx = PageTableTraverseContext(self.d[0], mm0.root_page_table, exteranl_va + 0x2000)
       for tup in ctx.next(0x100000):
         _offset, _pt, _pte_idx, _n_ptes, _pte_covers = tup
         for i in range(_n_ptes):

test/external/external_test_datasets.py

@@ -64,6 +64,7 @@ class TestKiTS19Dataset(ExternalTestDatasets):
 
     return iter(dataset)
 
+  @unittest.skip("flaky")
   def test_training_set(self):
     preproc_pth, preproc_img_pths, preproc_lbl_pths = self._create_samples(False)
     ref_dataset = self._create_ref_dataloader(preproc_img_pths, preproc_lbl_pths, False)

test/external/external_test_metrics.py

@@ -1,10 +1,10 @@
 from tinygrad import Tensor
 from test.external.mlperf_unet3d.dice import DiceScore
-from examples.mlperf.metrics import dice_score
+from examples.mlperf.metrics import dice_score, log_perplexity
 
 import numpy as np
 import torch
-import unittest
+import unittest, math
 
 class ExternalTestMetrics(unittest.TestCase):
   def _test_metrics(self, tinygrad_metrics, orig_metrics, pred, label, atol=1e-8, rtol=1e-7):
@@ -16,5 +16,27 @@ class ExternalTestMetrics(unittest.TestCase):
     pred, label = np.random.rand(1, 3, 128, 128, 128).astype(np.float32), np.ones((1, 1, 128, 128, 128)).astype(np.uint8)
     self._test_metrics(dice_score, DiceScore(), pred, label)
 
+  def test_log_perplexity(self):
+    # equally likely
+    np.testing.assert_allclose(log_perplexity(Tensor([[[1.0, 1, 1, 1]]]), Tensor([[2]])).numpy(), math.log(4))
+    np.testing.assert_allclose(log_perplexity(Tensor([[[1.0]*256]*32]), Tensor([[2]*32])).numpy(), math.log(256), rtol=1e-6)
+    # pretty correct and incorrect
+    np.testing.assert_allclose(log_perplexity(Tensor([[[10000., 0, 0, 0]]]), Tensor([[0]])).numpy(), 0)
+    np.testing.assert_allclose(log_perplexity(Tensor([[[0.0, 10000, 10000, 10000]]]), Tensor([[0]])).numpy(), 10000, rtol=1e-3)
+    # higher logit -> lower loss
+    x = Tensor([[[4.0, 3, 2, 1]]])
+    for i in range(x.numel()-1): self.assertLess(log_perplexity(x, Tensor([[i]])).item(), log_perplexity(x, Tensor([[i+1]])).item())
+    # torch eval examples
+    np.testing.assert_allclose(
+      log_perplexity(Tensor([[[0.3659, 0.7025, 0.3104], [0.0097, 0.6577, 0.1947]]]), Tensor([[2, 1]])).exp().numpy(),
+      2.7593, rtol=1e-5)
+    np.testing.assert_allclose(
+      log_perplexity(Tensor([[[0.3, 0.7, 0.3, 0.1], [0.5, 0.4, 0.1, 0.4],[0.1, 0.1, 0.2, 0.5]],
+                             [[0.1, 0.6, 0.1, 0.5], [0.3, 0.7, 0.3, 0.4], [0.3, 0.7, 0.3, 0.4]]]), Tensor([[2, 1, 3],  [1, 0, 1]])).exp().numpy(),
+      3.6216, rtol=1e-5)
+    np.testing.assert_allclose(
+      log_perplexity(Tensor([[[0.3659, 0.7025, 0.3104], [0.0097, 0.6577, 0.1947]]]), Tensor([[2, 1]]), ignore_index=1).exp().numpy(),
+      3.5372, rtol=1e-4)
+
 if __name__ == '__main__':
   unittest.main()

test/external/external_test_onnx_backend.py

@@ -94,13 +94,6 @@ backend_test.exclude('FLOAT8')
 backend_test.exclude('INT4')
 backend_test.exclude('UINT4')
 backend_test.exclude('BFLOAT16')  # not supported in numpy
-# TODO: fix these with true onnx float16
-backend_test.exclude('to_FLOAT16')
-backend_test.exclude('cast_no_saturate')
-backend_test.exclude('test_dequantizelinear_e4m3fn_float16_cpu')
-backend_test.exclude('test_max_float16_cpu')
-backend_test.exclude('test_min_float16_cpu')
-backend_test.exclude('test_mod_mixed_sign_float16_cpu')
 
 backend_test.exclude('test_dequantizelinear_int4_cpu')
 backend_test.exclude('test_dequantizelinear_uint4_cpu')
@@ -114,6 +107,7 @@ backend_test.exclude('test_quantizelinear_e4m3fn_cpu')
 backend_test.exclude('test_quantizelinear_e5m2_cpu')
 backend_test.exclude('test_dequantizelinear_e4m3fn_cpu')
 backend_test.exclude('test_dequantizelinear_e4m3fn_zero_point_cpu')
+backend_test.exclude('test_dequantizelinear_e4m3fn_float16_cpu')
 backend_test.exclude('test_dequantizelinear_e5m2_cpu')
 
 # we don't support indexes

test/external/external_test_onnx_runner.py

@@ -0,0 +1,77 @@
+import unittest, onnx, tempfile
+from tinygrad import dtypes
+from tinygrad.frontend.onnx import OnnxRunner, onnx_load
+from tinygrad.device import is_dtype_supported
+from extra.onnx import data_types
+from hypothesis import given, settings, strategies as st
+import numpy as np
+
+data_types.pop(16) # TODO: this is bf16, need to support double parsing first.
+device_supported_dtypes = [odt for odt, dtype in data_types.items() if is_dtype_supported(dtype)]
+device_unsupported_dtypes = [odt for odt, dtype in data_types.items() if not is_dtype_supported(dtype)]
+
+class TestOnnxRunnerDtypes(unittest.TestCase):
+  def _test_input_spec_dtype(self, onnx_data_type, tinygrad_dtype):
+    input_tensor = onnx.helper.make_tensor_value_info('input', onnx_data_type, ())
+    output_tensor = onnx.helper.make_tensor_value_info('output', onnx_data_type, ())
+    node = onnx.helper.make_node('Identity', inputs=['input'], outputs=['output'])
+    graph = onnx.helper.make_graph([node], 'identity_test', [input_tensor], [output_tensor])
+    model = onnx.helper.make_model(graph)
+    tmp = tempfile.NamedTemporaryFile(suffix='.onnx')
+    onnx.save(model, tmp.name)
+    tmp.flush()
+    model = onnx_load(tmp.name)
+    runner = OnnxRunner(model)
+    self.assertEqual(len(runner.graph_inputs), 1)
+    self.assertEqual(runner.graph_inputs['input'].dtype, tinygrad_dtype)
+
+  def _test_initializer_dtype(self, onnx_data_type, tinygrad_dtype):
+    arr = np.array([0, 1], dtype=onnx.helper.tensor_dtype_to_np_dtype(onnx_data_type))
+    initializer = onnx.helper.make_tensor('initializer', onnx_data_type, arr.shape, arr.tobytes(), raw=True)
+    input_tensor = onnx.helper.make_tensor_value_info('input', onnx_data_type, ())
+    output_tensor = onnx.helper.make_tensor_value_info('output', onnx_data_type, ())
+    node = onnx.helper.make_node('Identity', inputs=['input'], outputs=['output'])
+    graph = onnx.helper.make_graph([node], 'identity_test', [input_tensor], [output_tensor], [initializer])
+    model = onnx.helper.make_model(graph)
+    tmp = tempfile.NamedTemporaryFile(suffix='.onnx')
+    onnx.save(model, tmp.name)
+    tmp.flush()
+    model = onnx_load(tmp.name)
+    runner = OnnxRunner(model)
+    self.assertEqual(len(runner.graph_inputs), 1)
+    self.assertEqual(runner.graph_values['initializer'].dtype, tinygrad_dtype)
+
+  def _test_node_attribute_dtype(self, onnx_data_type, tinygrad_dtype):
+    arr = np.array([0, 1], dtype=onnx.helper.tensor_dtype_to_np_dtype(onnx_data_type))
+    output_tensor = onnx.helper.make_tensor_value_info('output', onnx_data_type, arr.shape)
+    value_tensor = onnx.helper.make_tensor('value', onnx_data_type, arr.shape, arr.tobytes(), raw=True)
+    node = onnx.helper.make_node('Constant', inputs=[], outputs=['output'], value=value_tensor)
+    graph = onnx.helper.make_graph([node], 'attribute_test', [], [output_tensor])
+    model = onnx.helper.make_model(graph)
+    tmp = tempfile.NamedTemporaryFile(suffix='.onnx')
+    tmp.flush()
+    onnx.save(model, tmp.name)
+    model = onnx_load(tmp.name)
+    runner = OnnxRunner(model)
+    self.assertEqual(runner.graph_nodes[0].opts['value'].dtype, tinygrad_dtype)
+
+  @settings(deadline=1000) # TODO investigate unreliable timing
+  @given(onnx_data_type=st.sampled_from(device_supported_dtypes))
+  def test_supported_dtype_spec(self, onnx_data_type):
+    tinygrad_dtype = data_types[onnx_data_type]
+    self._test_input_spec_dtype(onnx_data_type, tinygrad_dtype)
+    self._test_initializer_dtype(onnx_data_type, tinygrad_dtype)
+    self._test_node_attribute_dtype(onnx_data_type, tinygrad_dtype)
+
+  @unittest.skipUnless(device_unsupported_dtypes, "No unsupported dtypes for this device to test.")
+  @settings(deadline=1000) # TODO investigate unreliable timing
+  @given(onnx_data_type=st.sampled_from(device_unsupported_dtypes))
+  def test_unsupported_dtype_spec(self, onnx_data_type):
+    true_dtype = data_types[onnx_data_type]
+    default_dtype = dtypes.default_int if dtypes.is_int(true_dtype) else dtypes.default_float
+    self._test_input_spec_dtype(onnx_data_type, true_dtype)
+    self._test_initializer_dtype(onnx_data_type, default_dtype)
+    self._test_node_attribute_dtype(onnx_data_type, default_dtype)
+
+if __name__ == '__main__':
+  unittest.main()

test/external/external_test_optim.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-import unittest
+import unittest, math
 import numpy as np
 import tensorflow as tf
 import tensorflow_addons as tfa
@@ -7,11 +7,11 @@ from tensorflow.python.ops import math_ops
 from extra.lr_scheduler import LRSchedulerGroup
 
 from tinygrad.tensor import Tensor
-from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup
+from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, AdamW
 
 from test.external.mlperf_resnet.lars_optimizer import LARSOptimizer
 
-from examples.mlperf.lr_schedulers import PolynomialDecayWithWarmup
+from examples.mlperf.lr_schedulers import PolynomialDecayWithWarmup, CosineAnnealingLRWithWarmup
 from test.external.mlperf_resnet.lars_util import PolynomialDecayWithWarmup as PolynomialDecayWithWarmup_tf
 
 np.random.seed(1337)
@@ -171,5 +171,26 @@ class ExternalTestOptim(unittest.TestCase):
       'warmup': steps_per_epoch * warmup_epochs,
     }, 1e-5, 1e-5, do_optim=False)
 
+
+class TestCosineAnnealingLRWithWarmup(unittest.TestCase):
+  # only tests the lr
+  def _test_lr(self, base_lr, end_lr, warmup_steps, decay_steps):
+    net = TinyNet()
+    optim = AdamW([net.W], lr=0.0)
+    tiny_lr = CosineAnnealingLRWithWarmup(optim, base_lr, end_lr, warmup_steps, decay_steps)
+    lr = []
+    for _ in range(warmup_steps+decay_steps):
+      lr.append(optim.lr.item())
+      tiny_lr.step()
+    # reimplemented in python
+    expected = []
+    for i in range(warmup_steps): expected.append((i+1)/warmup_steps*base_lr)
+    for i in range(decay_steps): expected.append(end_lr+(base_lr-end_lr)*(1+math.cos((i+1)/decay_steps*math.pi))/2)
+    np.testing.assert_allclose(lr, expected, rtol=1e-5)
+
+  def test_lr_0(self): self._test_lr(3e-4, 8e-5, 3, 5)
+  def test_lr_1(self): self._test_lr(3e-4, 8e-5, 10, 20)
+  def test_lr_llama3(self): self._test_lr(8e-5, 8e-7, 20, 100)
+
 if __name__ == '__main__':
   unittest.main()

test/external/external_test_tlsf.py

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad.runtime.support.allocator import TLSFAllocator
+from tinygrad.runtime.support.memory import TLSFAllocator
 
 class TestTLSFAllocator(unittest.TestCase):
   def setUp(self):

test/external/external_uop_gc.py

@@ -14,10 +14,10 @@ def two_plus_two(): Tensor([2])+Tensor([2])
 def two_plus_two_schedule(): (Tensor([2])+Tensor([2])).schedule()
 def two_plus_two_kernel():
   si = (Tensor([2])+Tensor([2])).schedule()[-1]
-  get_program(Device.default.renderer, si.ast)
+  get_program(si.ast, Device.default.renderer)
 def two_plus_two_linearize():
   si = (Tensor([2])+Tensor([2])).schedule()[-1]
-  get_program(Device.default.renderer, si.ast)
+  get_program(si.ast, Device.default.renderer)
 def two_plus_two_realize(): (Tensor([2])+Tensor([2])).realize()
 def two_plus_two_item(): (Tensor([2])+Tensor([2])).item()
 def gradient_test():
@@ -34,7 +34,7 @@ def kernel_matmul():
   y = Tensor([[2.0,0,-2.0]], requires_grad=True)
   z = y.matmul(x)
   si = z.schedule()[-1]
-  get_program(Device.default.renderer, si.ast)
+  get_program(si.ast, Device.default.renderer)
 def realized_matmul():
   x = Tensor.eye(3, requires_grad=True)
   y = Tensor([[2.0,0,-2.0]], requires_grad=True)

test/external/process_replay/process_replay.py

@@ -1,11 +1,12 @@
 #!/usr/bin/env python3
 # compare kernels created by HEAD against master
-import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, itertools
+import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, itertools, functools
 from typing import Callable, Any
-from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm, to_function_name
-from tinygrad.engine.kernelize import get_kernelize_map
-from tinygrad.opt.kernel import Kernel
-from tinygrad.uop.ops import UOp, Ops
+from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm
+from tinygrad.kernelize.kernelize import get_kernelize_map
+from tinygrad.renderer import Renderer, ProgramSpec
+from tinygrad.engine.realize import get_program
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
 
 # *** process replay settings
 
@@ -39,21 +40,16 @@ def replay_kernelize(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str, tuple[
     return "\n".join([f"{len(asts)} kernels", *asts])
   return to_str(new_sink), to_str(ret[big_sink]), (big_sink,)
 
-def replay_linearize(k:Kernel, _:Kernel, name_override=None, ast_transform=None) -> tuple[str, str, tuple[Any, ...]]:
-  # create a copy because the Kernel class contains optimization parameters (other than applied_opts) in its state
-  # this should be made fully functional. It's fine for process replay since copy returns a fresh instance
-  k2 = k.copy()
-  k2.linearize(name_override=name_override or to_function_name(k.name), ast_transform=ast_transform)
-  def to_str(ret:Kernel) -> str:
-    try: return ret.opts.render(ret.uops)
-    except NotImplementedError: return "" # NULL backend doesn't have a renderer, this is okay
-  return to_str(k2), to_str(k), (k.ast, k.opts, k.applied_opts)
+def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer) -> tuple[str, str, tuple[Any, ...]]:
+  p2 = get_program(ast.replace(arg=KernelInfo(opts_to_apply=p.applied_opts, name=p.name)) if ast.arg is None else ast, renderer)
+  def to_str(ret:ProgramSpec) -> str: return ret.src
+  return to_str(p2), to_str(p), (p.ast, renderer, p.applied_opts)
 
-replayers: dict[str, Callable[..., tuple[str, str, tuple[Any, ...]]]] = {"get_kernelize_map":replay_kernelize, "linearize":replay_linearize}
+replayers: dict[str, Callable[..., tuple[str, str, tuple[Any, ...]]]] = {"get_kernelize_map":replay_kernelize, "get_program":replay_get_program}
 
 # *** run replayers on captured rows and print diffs
 
-def diff(offset:int) -> None:
+def diff(offset:int, fxns:dict[str, Callable[..., tuple|None]]) -> None:
   if ASSERT_DIFF: warnings.filterwarnings("error", category=ProcessReplayWarning)
   if early_stop.is_set(): return None
   conn = db_connection()
@@ -68,8 +64,10 @@ def diff(offset:int) -> None:
     try:
       name, args, kwargs, ctx_vals, loc, ret = pickle.loads(row[0])
       ctx_vars = {k:v.value for k,v in ctx_vals.items() if k != "DEBUG" and (var:=ContextVar._cache.get(k)) is not None and var.value != v.value}
-      if (replayer:=replayers.get(name)) is None: continue
-      with Context(**ctx_vars): good, compare, metadata = replayer(ret, *args, **kwargs)
+      if (replayer:=fxns.get(name)) is None: continue
+      with Context(**ctx_vars):
+        if (ret:=replayer(ret, *args, **kwargs)) is None: continue
+        good, compare, metadata = ret
       if good != compare:
         for m in metadata: trunc_log(m)
         logging.info(loc)
@@ -83,6 +81,25 @@ def diff(offset:int) -> None:
   conn.commit()
   cur.close()
 
+# *** generic runner to map rows of a table to a function in parallel
+
+def _pmap(fxns:dict[str, Callable]) -> None:
+  conn = db_connection()
+  cur = conn.cursor()
+  try: row_count = cur.execute(f"select count(*) from '{TABLE_NAME}'").fetchone()[0]
+  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:
+    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()
+
 # *** main loop
 
 if __name__ == "__main__":
@@ -90,20 +107,8 @@ if __name__ == "__main__":
     logging.info("skipping process replay.")
     exit(0)
 
-  conn = db_connection()
-  cur = conn.cursor()
-  try: row_count = cur.execute(f"select count(*) from '{TABLE_NAME}'").fetchone()[0]
-  except sqlite3.OperationalError:
-    warnings.warn(f"{TABLE_NAME} isn't accessible in master, did DB_VERSION change?", ProcessReplayWarning)
-    exit(int(ASSERT_DIFF))
-  finally:
-    conn.commit()
-    cur.close()
-
   logging.info(f"running process replay with {ASSERT_DIFF=}")
-  with multiprocessing.get_context("spawn").Pool(multiprocessing.cpu_count()) as pool:
-    inputs = list(range(0, row_count, PAGE_SIZE))
-    list(tqdm(pool.imap_unordered(diff, inputs), total=len(inputs)))
-    pool.close()
-    pool.join()
-    pool.terminate()
+  try: _pmap(replayers)
+  except Exception as e:
+    logging.info("process replay err", e)
+    exit(int(ASSERT_DIFF))

test/external/sglang_llama/external_llama_eval.py

@@ -0,0 +1,100 @@
+from lm_eval import simple_evaluate
+from lm_eval.api.instance import Instance
+from lm_eval.api.model import LM
+from lm_eval.tasks import TaskManager
+from pathlib import Path
+import json, argparse
+
+from examples.llama3 import build_transformer, Tokenizer, MODEL_PARAMS
+from tinygrad import Tensor, Device
+from tinygrad.helpers import tqdm
+
+class LLaMaAdaptor(LM):
+  def __init__(
+    self,
+    model_size: str,
+    checkpoint_path: Path,
+    max_length: int,
+    quantize: str | None,
+  ):
+    super().__init__()
+    self.max_length = max_length
+    self.tokenizer = Tokenizer(str((checkpoint_path if checkpoint_path.is_dir() else checkpoint_path.parent) / "tokenizer.model"))
+    self.model = build_transformer(checkpoint_path, model_size=model_size, quantize=quantize, max_context=self.max_length)
+    self.last_seen_toks = []
+  def _prefill(self, toks, temperature) -> int:
+    start_pos = 0
+    # we can skip part of the prompt if it is the same as last
+    for i, (a, b) in enumerate(zip(toks, self.last_seen_toks)):
+      if a != b: break
+    else: i = min(len(toks), len(self.last_seen_toks))
+    start_pos += i
+    self.last_seen_toks = toks
+    toks = toks[i:]
+
+    # prefill the model
+    for tok in toks:
+      self.model(Tensor([[tok]]), start_pos, temperature).realize()
+      start_pos += 1
+    return start_pos
+
+  @property
+  def tokenizer_name(self) -> str: pass
+  def chat_template(self, chat_template: bool | str = False) -> str: pass
+  def apply_chat_template(self, chat_history: list[dict[str, str]], add_generation_prompt: bool = True) -> str:
+    ret = ""
+    for message in chat_history:
+      ret += f"<|start_header_id|>{message['role']}<|end_header_id|>\n\n{message['content'].strip()}<|eot_id|>"
+    if add_generation_prompt: ret += "<|start_header_id|>assistant<|end_header_id|>\n\n"
+    return ret
+
+  def generate_until(self, requests: list[Instance]) -> list[str]:
+    continuations = []
+    for request in tqdm(requests):
+      prompt, args = request.args
+      until = [self.tokenizer.encode(tok) for tok in args.get("until", [])]
+      toks = [self.tokenizer.bos_id] + self.tokenizer.encode(prompt,allow_special=True)
+      prompt_len = len(toks)
+      max_gen_toks = args.get("max_gen_toks") or args.get("max_length") or self.max_length-prompt_len
+      assert self.max_length >= max_gen_toks, "This eval needs a longer context length"
+      temperature = args.get("temperature", 0.0)
+      start_pos = self._prefill(toks[:-1], temperature)
+
+      for _ in range(max_gen_toks):
+        next_tok = self.model(Tensor([toks[start_pos:]]), start_pos, temperature).item()
+        if next_tok in self.tokenizer.stop_tokens or next_tok in until: break
+        toks.append(next_tok)
+        start_pos += 1
+
+      continuations.append(self.tokenizer.decode(toks[prompt_len:]))
+    return continuations
+
+  def loglikelihood(self, requests: list[Instance]) -> list[tuple[float, bool]]: raise NotImplementedError() # needs changes to extra/models/llama.py
+  def loglikelihood_rolling(self, requests: list[Instance]) -> list[tuple[float, bool]]: raise NotImplementedError()
+
+if __name__ == '__main__':
+  print(f"using {Device.DEFAULT} backend")
+
+  parser = argparse.ArgumentParser(description='Run LLaMA evals in tinygrad', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+  parser.add_argument('--size', type=str, default="8B", help=f"Size of model to use [{', '.join(list(MODEL_PARAMS.keys()))}]")
+  parser.add_argument('--chat', action='store_true', help="Use chat model")
+  parser.add_argument('--ctx', type=int, default=8192, help="Max context length")
+  parser.add_argument('--quantize', type=str, default=None, help="Quantize the weights to int8 or int4 in memory")
+  parser.add_argument('--eval', type=str, default="mgsm_en_cot_sglang", help="Run in evaluation mode")
+  parser.add_argument('--limit', type=int, default=None, help="Limit tests in eval")
+  parser.add_argument('--num_fewshot', type=int, default=None, help="Number of examples to add to context")
+  parser.add_argument('--model', type=Path, default="./weights/LLaMa/", help="Location of the weights")
+  parser.add_argument('--output_path', type=Path, default=None, help="Location of the log file")
+  args = parser.parse_args()
+
+  # run eval and exit
+  adaptor = LLaMaAdaptor(model_size=args.size, quantize=args.quantize,
+                         checkpoint_path=args.model, max_length=args.ctx)
+  task_manager = TaskManager(include_path="./")
+  results = simple_evaluate(model=adaptor, tasks=args.eval.split(","), task_manager=task_manager, apply_chat_template=args.chat,
+                            num_fewshot=args.num_fewshot, limit=args.limit)
+
+  if args.output_path: args.output_path.write_text(json.dumps(results, indent=2))
+  for task_name, val in results["results"].items():
+    print(f"{task_name}:")
+    print("\n".join(f"\t{k}: {v}" for k, v in val.items() if k != "alias"))

test/external/sglang_llama/mgsm.yaml

@@ -0,0 +1,36 @@
+# https://github.com/sgl-project/sglang/blob/main/python/sglang/test/simple_eval_mgsm.py#L41
+task: mgsm_en_cot_sglang
+dataset_path: juletxara/mgsm
+dataset_name: en
+output_type: generate_until
+training_split: train
+test_split: test
+doc_to_target: '{{answer[21:] if answer is not none else answer_number|string}}'
+doc_to_text: >-
+  {{'Solve this math problem. Give the reasoning steps before giving the final answer on the last line by itself in the format of "Answer:".
+  Do not add anything other than the integer answer after "Answer:".\n\n'
+  +(question[10:] if answer is not none else question)}}
+generation_kwargs:
+  do_sample: false
+  temperature: 0.0
+  until: []
+metric_list:
+  - metric: exact_match
+    aggregation: mean
+    higher_is_better: true
+    ignore_case: true
+    ignore_punctuation: true
+filter_list:
+  - name: "strict-match"
+    filter:
+      - function: "regex"
+        regex_pattern: 'Answer:\s*([\-]?[0-9\.\,]+)'
+      - function: "take_first"
+  - filter:
+    - function: regex
+      group_select: -1
+      regex_pattern: (-?[$0-9.,]{2,})|(-?[0-9]+)
+    - function: take_first
+    name: flexible-extract
+metadata:
+  version: 3.0

test/external/speed_compare_cuda_ptx.py

@@ -33,7 +33,6 @@ if __name__ == "__main__":
     dev.compiler = PTXCompiler(dev.arch)
     lin = ast_str_to_lin(ast, opts=ptx)
     lin.apply_opts(hand_coded_optimizations(lin))
-    lin.linearize()
     ptx_prg = CompiledRunner(lin.to_program())
 
     # warmup

test/helpers.py

@@ -1,15 +1,13 @@
 import time, struct
-from typing import Any, Callable, Optional
+from typing import Any, Callable
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.uop.ops import UOp, Ops, sint, graph_rewrite
-from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.uop.ops import UOp, Ops
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.engine.realize import Runner
-from tinygrad.engine.kernelize import view_left
-from tinygrad.dtype import ConstType, DType
+from tinygrad.dtype import DType
 from tinygrad.nn.state import get_parameters
-from tinygrad.helpers import T, unwrap, CI
+from tinygrad.helpers import T, CI
 from tinygrad.codegen import full_rewrite
 from tinygrad.runtime.ops_python import PythonProgram, PythonRenderer, PythonCompiler
 
@@ -40,13 +38,6 @@ def rand_for_dtype(dt:DType, size:int):
     return np.random.choice([True, False], size=size)
   return np.random.uniform(-10, 10, size=size).astype(_to_np_dtype(dt))
 
-def ast_const(dtype:DType, val:ConstType, shape:tuple[sint, ...]=(), st:Optional[ShapeTracker]=None, st_src:Optional[tuple[UOp]]=None) -> UOp:
-  if st_src is None:
-    st_src = (st.to_uop() if st is not None else ShapeTracker.from_shape(()).reshape((1,)*len(shape)).expand(shape).to_uop(),)
-  st = unwrap(st_src[0].st)
-  if all(v.mask is None for v in st.views): return UOp.const(dtype, val).replace(src=(st.to_uop(),))
-  return graph_rewrite(UOp.const(dtype, val).view(st).valid(), view_left)
-
 def timeit(fxn:Callable[..., T], *args, **kwargs) -> tuple[T, float]:
   st = time.perf_counter_ns()
   ret = fxn(*args, **kwargs)

test/models/test_real_world.py

@@ -165,7 +165,7 @@ class TestRealWorld(unittest.TestCase):
       for v in data.values(): v.to_(Device.DEFAULT)
 
       helper_test("train_bert", lambda: (data["input_ids"], data["segment_ids"], data["input_mask"], data["masked_lm_positions"], \
-          data["masked_lm_ids"], data["masked_lm_weights"], data["next_sentence_labels"]), train, 0.25, 346)
+          data["masked_lm_ids"], data["masked_lm_weights"], data["next_sentence_labels"]), train, 0.25, 347)
 
   def test_bert_fuse_arange(self):
     with Context(FUSE_ARANGE=1):

test/test_multitensor.py

@@ -1,4 +1,4 @@
-import unittest, functools, random, os
+import unittest, functools, random
 from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable
 from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import Ops, UOp
@@ -1101,7 +1101,6 @@ class TestTensorOps(unittest.TestCase):
   def test_bitcast(self):
     helper_test_shard_op([(256,), (256,)], lambda x: x.bitcast(dtypes.int))
 
-# TODO: make these tests pass with VIZ=1
 @unittest.skipIf(not_support_multi_device(), "no multi")
 class TestMultiRamUsage(unittest.TestCase):
   def setUp(self):
@@ -1129,13 +1128,11 @@ class TestMultiRamUsage(unittest.TestCase):
 
   def test_zeros_shard(self, devices=(d1, d2)):
     _ = Tensor.zeros(self.N, self.N).contiguous().shard(devices, axis=0).realize()
-    assert int(os.getenv("VIZ", "0")) == 0
     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))
 
   def test_zeros_contiguous_shard(self):
     _ = Tensor.zeros(self.N, self.N).contiguous().shard(devices_2, axis=0).contiguous().realize()
-    assert int(os.getenv("VIZ", "0")) == 0
     self.assertUsed(self.N*self.N*4) # sharding should not increase total ram usage
 
 @unittest.skipIf(not_support_multi_device(), "need multi")

test/test_ops.py

@@ -1,9 +1,8 @@
-import time, math, unittest, functools
+import time, math, unittest, functools, warnings
 import numpy as np
 from typing import List, Callable
 import torch
-import warnings
-from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, DEVECTORIZE, OSX
+from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, OSX, AMD_LLVM
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
@@ -547,7 +546,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65), (45,65)], lambda x,y: x/y)
     helper_test_op([(), ()], lambda x,y: x/y)
 
-  @unittest.skipIf(getenv("AMD_LLVM", 0), "AMD with LLVM backend generate rcp in FP division causes trunc/floor errors")
+  @unittest.skipIf(AMD_LLVM, "AMD with LLVM backend generate rcp in FP division causes trunc/floor errors")
   def test_div_rounding_mode(self):
     for denominator in [-10, -5, -3, -2, -1, 1, 2, 3, 5, 10]:
       # int numerator
@@ -946,6 +945,11 @@ class TestOps(unittest.TestCase):
     self.assertAlmostEqual(sigmoid(x)[0].gradient(x)[0].item(), 0.0)
     x = Tensor([-300.0])
     self.assertAlmostEqual(sigmoid(x)[0].gradient(x)[0].item(), 0.0)
+
+  def test_logsigmoid(self):
+    helper_test_op([(45,65)], torch.nn.functional.logsigmoid, Tensor.logsigmoid)
+    helper_test_op([()], torch.nn.functional.logsigmoid, Tensor.logsigmoid)
+
   def test_hardsigmoid(self):
     helper_test_op([(45,65)], torch.nn.functional.hardsigmoid, Tensor.hardsigmoid)
     helper_test_op([()], torch.nn.functional.hardsigmoid, Tensor.hardsigmoid)
@@ -1534,6 +1538,9 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65)], lambda x: torch.logsumexp(x, dim=0), lambda x: x.logsumexp(0), atol=1e-7, grad_atol=1e-7)
     helper_test_op([(45,65)], lambda x: torch.logsumexp(x, dim=0, keepdim=True), lambda x: x.logsumexp(0, True), atol=1e-7, grad_atol=1e-7)
     helper_test_op([(45,65)], lambda x: torch.logsumexp(x, dim=1), lambda x: x.logsumexp(1), atol=1e-7, grad_atol=1e-7)
+    helper_test_op([(6,6,6)], lambda x: torch.logsumexp(x, dim=2), lambda x: x.logsumexp(2), atol=1e-7, grad_atol=1e-7)
+    helper_test_op([(6,6,6,6)], lambda x: torch.logsumexp(x, dim=2), lambda x: x.logsumexp(2), atol=1e-7, grad_atol=1e-7)
+    helper_test_op([(6,6,6,6)], lambda x: torch.logsumexp(x, dim=3), lambda x: x.logsumexp(3), atol=1e-7, grad_atol=1e-7)
     helper_test_op([(45)], lambda x: torch.logsumexp(x, dim=0), lambda x: x.logsumexp(0), atol=1e-7, grad_atol=1e-7)
     helper_test_op([()], lambda x: torch.logsumexp(x, dim=0), lambda x: x.logsumexp(0), atol=1e-7, grad_atol=1e-7)
     helper_test_op([()], lambda x: torch.logsumexp(x, dim=-1), lambda x: x.logsumexp(-1), atol=1e-7, grad_atol=1e-7)
@@ -1541,12 +1548,14 @@ class TestOps(unittest.TestCase):
   def test_logcumsumexp(self):
     helper_test_op([(45,65)], lambda x: torch.logcumsumexp(x, dim=0), lambda x: x.logcumsumexp(0), atol=1e-7, grad_atol=1e-7)
     helper_test_op([(45,65)], lambda x: torch.logcumsumexp(x, dim=1), lambda x: x.logcumsumexp(1), atol=1e-7, grad_atol=1e-7)
+    helper_test_op([(6,6,6)], lambda x: torch.logcumsumexp(x, dim=2), lambda x: x.logcumsumexp(2), atol=1e-7, grad_atol=1e-7)
+    helper_test_op([(6,6,6,6)], lambda x: torch.logcumsumexp(x, dim=2), lambda x: x.logcumsumexp(2), atol=1e-7, grad_atol=1e-7)
+    helper_test_op([(6,6,6,6)], lambda x: torch.logcumsumexp(x, dim=3), lambda x: x.logcumsumexp(3), atol=1e-7, grad_atol=1e-7)
     helper_test_op([(45)], lambda x: torch.logcumsumexp(x, dim=0), lambda x: x.logcumsumexp(0), atol=1e-7, grad_atol=1e-7)
     helper_test_op([()], lambda x: torch.logcumsumexp(x, dim=0), lambda x: x.logcumsumexp(0), atol=1e-7, grad_atol=1e-7)
     helper_test_op([()], lambda x: torch.logcumsumexp(x, dim=0), lambda x: x.logcumsumexp(), atol=1e-7, grad_atol=1e-7)
     helper_test_op([()], lambda x: torch.logcumsumexp(x, dim=-1), lambda x: x.logcumsumexp(-1), atol=1e-7, grad_atol=1e-7)
 
-  @unittest.skipIf(not DEVECTORIZE, "broken without DEVECTORIZE. TODO: fix this")
   def test_logcumsumexp_numerical(self):
     helper_test_op(None, lambda x: torch.logcumsumexp(x, dim=0), lambda x: x.logcumsumexp(), atol=1e-7, grad_atol=1e-7, vals=[[0.0, 100.0]])
 
@@ -2889,25 +2898,37 @@ class TestOps(unittest.TestCase):
       expected=RuntimeError)
 
   def test_binary_crossentropy(self):
-    helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy(x.sigmoid(),torch.clip(y,0,1)),
+    helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy(x.sigmoid(),y.clip(0,1)),
                                        lambda x,y: x.sigmoid().binary_crossentropy(y.clip(0,1)))
-    helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy_with_logits(x,torch.clip(y,0,1)),
+    helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy_with_logits(x,y.clip(0,1)),
                                        lambda x,y: x.binary_crossentropy_logits(y.clip(0,1)))
-    helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy_with_logits(x,torch.clip(y,0,1)),
+    helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy_with_logits(x,y.clip(0,1)),
                                        lambda x,y: x.sigmoid().binary_crossentropy(y.clip(0,1)))
-    helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy(x.sigmoid(),torch.clip(y,0,1)),
+    helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy(x.sigmoid(),y.clip(0,1)),
                                        lambda x,y: x.binary_crossentropy_logits(y.clip(0,1)))
   def test_binary_crossentropy_reductions(self):
     for r in ("mean", "sum", "none"):
-      helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy(x.sigmoid(), torch.clip(y,0,1), reduction=r),
+      helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy(x.sigmoid(), y.clip(0,1), reduction=r),
                                          lambda x,y: x.sigmoid().binary_crossentropy(y.clip(0,1), reduction=r))
-      helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy_with_logits(x, torch.clip(y,0,1), reduction=r),
+      helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy_with_logits(x, y.clip(0,1), reduction=r),
                                          lambda x,y: x.binary_crossentropy_logits(y.clip(0,1), reduction=r))
-  def test_cross_entropy(self):
-    helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.cross_entropy(x, y),
-                                       lambda x,y: x.cross_entropy(y))
-    helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.cross_entropy(x, torch.argmax(y, dim=1)),
-                                       lambda x,y: x.cross_entropy(y.argmax(axis=1)), forward_only=True)
+  def test_binary_crossentropy_logits_pos_weights(self):
+    pos_weight = [0.25, 0.5, 0.75, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]
+    helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy_with_logits(x,y.clip(0,1),
+                                                                                                        pos_weight=torch.tensor(pos_weight)),
+                                       lambda x,y: x.binary_crossentropy_logits(y.clip(0,1),pos_weight=Tensor(pos_weight)))
+  def test_cross_entropy_class_probabilities(self):
+    helper_test_op([(32,), (32,)], lambda x,y: torch.nn.functional.cross_entropy(x, y), lambda x,y: x.cross_entropy(y))
+    helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.cross_entropy(x, y), lambda x,y: x.cross_entropy(y))
+    helper_test_op([(32,4,4,4), (32,4,4,4)], lambda x,y: torch.nn.functional.cross_entropy(x, y), lambda x,y: x.cross_entropy(y))
+
+  def test_cross_entropy_class_indices(self):
+    classes = np.random.randint(0, 10, (32,), dtype=np.int32).tolist()
+    helper_test_op([(32,10)], lambda x: torch.nn.functional.cross_entropy(x, torch.tensor(classes)),
+                              lambda x: x.cross_entropy(Tensor(classes)))
+    self.helper_test_exception([(32,10), (32,1)], lambda x,y: torch.nn.functional.cross_entropy(x, y),
+                                                  lambda x,y: x.cross_entropy(y), expected=(AssertionError, RuntimeError))
+
   def test_cross_entropy_reductions(self):
     for r in ("mean", "sum", "none"):
       helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.cross_entropy(x, y, reduction=r),
@@ -2919,49 +2940,56 @@ class TestOps(unittest.TestCase):
     for ls in (0., 0.3, 0.7, 1.):
       helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.cross_entropy(x, y, label_smoothing=ls),
                                          lambda x,y: x.cross_entropy(y, label_smoothing=ls))
+      classes = np.random.randint(0, 10, (32,), dtype=np.int32).tolist()
+      helper_test_op([(32,10)], lambda x: torch.nn.functional.cross_entropy(x, torch.tensor(classes), label_smoothing=ls),
+                                lambda x: x.cross_entropy(Tensor(classes), label_smoothing=ls))
 
   def test_nll_loss(self):
-    helper_test_op([(32,10), (32)],
-                   lambda x,y: torch.nn.functional.nll_loss(torch.nn.functional.log_softmax(x, dim=1), torch.clip(y,0).type(torch.long)),
-                   lambda x,y: x.log_softmax(axis=1).nll_loss(y.clip(0).cast(dtypes.int32)), forward_only=True)
+    target = np.random.randint(0, 10, (32,), dtype=np.int32).tolist()
+    helper_test_op([(32,10)],
+                   lambda x: torch.nn.functional.nll_loss(torch.nn.functional.log_softmax(x, dim=1), torch.tensor(target)),
+                   lambda x: x.log_softmax(axis=1).nll_loss(Tensor(target)))
 
   def test_nll_loss_3d(self):
-    helper_test_op([(32,10,3,3,3), (32,3,3,3)],
-                   lambda x,y: torch.nn.functional.nll_loss(torch.nn.functional.log_softmax(x, dim=1), torch.clip(y,0).type(torch.long)),
-                   lambda x,y: x.log_softmax(axis=1).nll_loss(y.clip(0).cast(dtypes.int32)), forward_only=True)
+    target = np.random.randint(0, 10, (32,3,3,3), dtype=np.int32).tolist()
+    helper_test_op([(32,10,3,3,3)],
+                   lambda x: torch.nn.functional.nll_loss(torch.nn.functional.log_softmax(x, dim=1), torch.tensor(target)),
+                   lambda x: x.log_softmax(axis=1).nll_loss(Tensor(target)))
 
   def test_nll_loss_reductions(self):
+    target = np.random.randint(0, 10, (32,), dtype=np.int32).tolist()
     for r in ("mean", "sum", "none"):
-      helper_test_op([(32,10), (32)],
-        lambda x,y: torch.nn.functional.nll_loss(torch.nn.functional.log_softmax(x, dim=1), torch.clip(y,0).type(torch.long), reduction=r),
-        lambda x,y: x.log_softmax(axis=1).nll_loss(y.clip(0).cast(dtypes.int32), reduction=r), forward_only=True)
-    self.helper_test_exception([(32,10), (32)],
-      lambda x,y: torch.nn.functional.nll_loss(x, torch.clip(y,0).type(torch.long), reduction="typo"),
-      lambda x,y: x.nll_loss(y.clip(0).cast(dtypes.int32), reduction="typo"), expected=ValueError)
+      helper_test_op([(32,10)],
+        lambda x: torch.nn.functional.nll_loss(torch.nn.functional.log_softmax(x, dim=1), torch.tensor(target), reduction=r),
+        lambda x: x.log_softmax(axis=1).nll_loss(Tensor(target), reduction=r))
+    self.helper_test_exception([(32,10)],
+      lambda x: torch.nn.functional.nll_loss(x, torch.tensor(target), reduction="typo"),
+      lambda x: x.nll_loss(Tensor(target), reduction="typo"), expected=ValueError)
 
   def test_nll_loss_weight(self):
+    target = np.random.randint(0, 10, (32,), dtype=np.int32).tolist()
+    weight = np.random.normal(0, 1, (10,)).astype(np.float32).tolist()
     for r in ("mean", "sum", "none"):
-      helper_test_op([(32,10), (32), (10)],
-        lambda x,y,z: torch.nn.functional.nll_loss(torch.nn.functional.log_softmax(x, dim=1), torch.clip(y,0).type(torch.long),
-                                                   weight=z, reduction=r),
-        lambda x,y,z: x.log_softmax(axis=1).nll_loss(y.clip(0).cast(dtypes.int32), weight=z, reduction=r), forward_only=True)
+      helper_test_op([(32,10)],
+        lambda x: torch.nn.functional.nll_loss(torch.nn.functional.log_softmax(x, dim=1), torch.tensor(target), torch.tensor(weight), reduction=r),
+        lambda x: x.log_softmax(axis=1).nll_loss(Tensor(target), Tensor(weight), reduction=r))
 
   def test_nll_loss_3d_weight(self):
+    target = np.random.randint(0, 10, (32,3,3,3), dtype=np.int32).tolist()
+    weight = np.random.normal(0, 1, (10,)).astype(np.float32).tolist()
     for r in ("mean", "sum", "none"):
-      helper_test_op([(32,10,3,3,3), (32,3,3,3), (10)],
-          lambda x,y,z: torch.nn.functional.nll_loss(torch.nn.functional.log_softmax(x, dim=1), torch.clip(y,0).type(torch.long),
-                                                    weight=z, reduction=r),
-          lambda x,y,z: x.log_softmax(axis=1).nll_loss(y.clip(0).cast(dtypes.int32), weight=z, reduction=r), forward_only=True)
+      helper_test_op([(32,10,3,3,3)],
+          lambda x: torch.nn.functional.nll_loss(torch.nn.functional.log_softmax(x, dim=1), torch.tensor(target), torch.tensor(weight), reduction=r),
+          lambda x: x.log_softmax(axis=1).nll_loss(Tensor(target), Tensor(weight), reduction=r))
 
   def test_nll_loss_ignore_index(self):
     logits = [[2.0, 0.5, -1.0],
               [1.5, 2.5, -0.5],
               [0.0, -2.0, 1.0]]
-    targets = [0, 1, 2]
-    helper_test_op(None, lambda x,y: torch.nn.functional.nll_loss(torch.nn.functional.log_softmax(x, dim=1),
-                                                                  torch.clip(y,0).type(torch.long), ignore_index=1),
-                         lambda x,y: x.log_softmax().nll_loss(y.clip(0), ignore_index=1),
-                         forward_only=True, vals=[logits, targets])
+    target = [0, 1, 2]
+    helper_test_op(None, lambda x: torch.nn.functional.nll_loss(torch.nn.functional.log_softmax(x, dim=1), torch.tensor(target), ignore_index=1),
+                         lambda x: x.log_softmax().nll_loss(Tensor(target), ignore_index=1),
+                         vals=[logits])
 
   def test_one_hot(self):
     data = [1, 2, 4]

test/test_quantize_onnx.py

@@ -78,7 +78,7 @@ class TestQuantizeOnnxCPU(unittest.TestCase):
     with Context(DONT_REALIZE_EXPAND=1, QUANTIZE=1):
       sched = run_onnx({"input":inp})["output"].schedule()
       ei = lower_schedule_item(sched[-2])
-      daccs = [u for u in ei.prg.p.uops if u.op is Ops.DEFINE_ACC]
+      daccs = [u for u in ei.prg.p.uops if u.op is Ops.DEFINE_REG]
       assert all(u.dtype.scalar() is dtypes.int for u in daccs)
 
 @unittest.skipIf(Device.DEFAULT != "DSP", "only tests for DSP")

test/test_schedule.py

@@ -15,7 +15,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.uop.ops import PatternMatcher, UOp, Ops, GroupOp, UPat, graph_rewrite, track_rewrites
 from tinygrad.uop.symbolic import symbolic_simple
 from tinygrad.helpers import CI, DEBUG, FUSE_ARANGE, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp
-from tinygrad.engine.kernelize import merge_views, get_kernelize_map, Kernel
+from tinygrad.kernelize.kernelize import merge_views, get_kernelize_map, Kernel
 from tinygrad.engine.schedule import ScheduleItem, create_schedule_with_vars
 from tinygrad.engine.realize import CompiledRunner, run_schedule, lower_schedule
 
@@ -1622,7 +1622,7 @@ class TestSchedule(unittest.TestCase):
     run_schedule(check_schedule(out, 3)) # TODO: push a reduceop through a reshape
 
   def test_conv2d(self): _test_conv2d(7)
-  def test_conv2d_fused(self): _test_conv2d(6, FUSE_CONV_BW=1)
+  def test_conv2d_fused(self): _test_conv2d(5, FUSE_CONV_BW=1)
 
   @unittest.skipUnless(is_dtype_supported(dtypes.half) and is_dtype_supported(dtypes.ulong), "need half and ulong")
   def test_conv2d_half(self): _test_conv2d(7, dtype=dtypes.half)
@@ -1631,7 +1631,6 @@ class TestSchedule(unittest.TestCase):
   @unittest.expectedFailure
   def test_conv2d_fused_half(self): _test_conv2d(5, dtype=dtypes.half)
 
-  @unittest.skipIf(getenv("VIZ"), "TODO: VIZ blocks gc")
   def test_schedule_mem_used(self):
     base = GlobalCounters.mem_used
     Tensor.ones(256).contiguous().realize()

test/test_search.py

@@ -83,6 +83,7 @@ class TestBEAM(unittest.TestCase):
     actions_after = actions.copy()
     assert actions_after == actions_before, "actions state was not preserved"
 
+  @unittest.skip("invalid reduce now")
   def test_filter_global_buffer(self):
     # taken from https://github.com/tinygrad/tinygrad/issues/4612
     ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(

test/test_subbuffer.py

@@ -9,6 +9,7 @@ class TestSubBuffer(unittest.TestCase):
   def setUp(self):
     self.buf = Buffer(Device.DEFAULT, 10, dtypes.uint8).ensure_allocated()
     self.buf.copyin(memoryview(bytearray(range(10))))
+    self.buf_unalloc = Buffer(Device.DEFAULT, 10, dtypes.uint8)
 
   def test_subbuffer(self):
     vbuf = self.buf.view(2, dtypes.uint8, offset=3).ensure_allocated()
@@ -64,5 +65,119 @@ class TestSubBuffer(unittest.TestCase):
       tst = vbuf.as_buffer().tolist()
       assert tst == [13, 14]
 
+  def test_subbuffer_is_allocated(self):
+    buf = self.buf_unalloc
+    sub_buf = buf.view(3, dtypes.uint8, offset=4)
+    self.assertFalse(buf.is_allocated())
+    self.assertFalse(buf.is_initialized())
+    self.assertFalse(sub_buf.is_allocated())
+    self.assertFalse(sub_buf.is_initialized())
+
+    # base buffer alloc
+    buf.allocate()
+    self.assertTrue(buf.is_allocated())
+    self.assertTrue(buf.is_initialized())
+    self.assertTrue(sub_buf.is_allocated())
+    self.assertFalse(sub_buf.is_initialized())
+
+    # sub buffer alloc
+    sub_buf.allocate()
+    self.assertTrue(sub_buf.is_initialized())
+
+    # sub buffer dealloc
+    sub_buf.deallocate()
+    self.assertTrue(buf.is_allocated())
+    self.assertTrue(buf.is_initialized())
+    self.assertTrue(sub_buf.is_allocated())
+    self.assertFalse(sub_buf.is_initialized())
+
+    # base buffer dealloc
+    buf.deallocate()
+    self.assertFalse(buf.is_allocated())
+    self.assertFalse(buf.is_initialized())
+    self.assertFalse(sub_buf.is_allocated())
+    self.assertFalse(sub_buf.is_initialized())
+
+    # sub buffer alloc
+    sub_buf.ensure_allocated()
+    self.assertTrue(buf.is_allocated())
+    self.assertTrue(buf.is_initialized())
+    self.assertTrue(sub_buf.is_allocated())
+    self.assertTrue(sub_buf.is_initialized())
+
+  def test_subbuffer_copy_in_out(self):
+    sub_buf = self.buf.view(3, dtypes.uint8, offset=3).ensure_allocated() # [3:6]
+    data_out_sub = bytearray([0]*3)
+    sub_buf.copyout(memoryview(data_out_sub))
+    assert data_out_sub == bytearray(range(3, 6))
+    sub_buf.copyin(memoryview(bytearray(range(3))))
+    assert sub_buf.as_buffer().tolist() == list(range(3))
+    assert self.buf.as_buffer().tolist()[3:6] == list(range(3))
+    sub_buf.copyout(memoryview(data_out_sub))
+    assert data_out_sub == bytearray(range(3))
+    data_out_base = bytearray([0]*10)
+    self.buf.copyout(memoryview(data_out_base))
+    assert data_out_base[0:3] == bytearray(range(0, 3))
+    assert data_out_base[3:6] == data_out_sub
+    assert data_out_base[6:10] == bytearray(range(6, 10))
+
+  def test_subbuffer_copy_in_out_view_of_view(self):
+    view1 = self.buf.view(7, dtypes.uint8, offset=2).ensure_allocated() # [2:9]
+    view2 = view1.view(3, dtypes.uint8, offset=2).ensure_allocated()   # [4:7]
+    self.assertTrue(view1.is_allocated())
+    self.assertTrue(view2.is_allocated())
+
+    data_in = bytearray([7, 8, 9])
+    view2.copyin(memoryview(data_in))
+    data_out_v2 = bytearray([0]*3)
+    view2.copyout(memoryview(data_out_v2))
+    assert data_in == data_out_v2
+
+    expected_base_data = memoryview(bytearray(range(10)))
+    expected_base_data[4:7] = data_in
+
+    data_out_base = bytearray([0]*10)
+    self.buf.copyout(memoryview(data_out_base))
+    assert expected_base_data == data_out_base
+
+  def test_subbuffer_alloc(self):
+    sub_buf = self.buf.view(4, dtypes.int8, offset=3)
+    sub_buf.allocate()
+    sub_buf.copyin(memoryview(bytearray(range(10, 14))))
+    assert self.buf.as_buffer().tolist()[3:7] == sub_buf.as_buffer().tolist()
+
+    sub_buf = self.buf_unalloc.view(4, dtypes.int8, offset=3)
+    sub_buf.allocate()
+    sub_buf.copyin(memoryview(bytearray(range(10, 14))))
+    assert self.buf_unalloc.as_buffer().tolist()[3:7] == sub_buf.as_buffer().tolist()
+
+  def test_subbuffer_dealloc(self):
+    sub_buf = self.buf.view(4, dtypes.int8, offset=3).ensure_allocated()
+    sub_buf.deallocate()
+    assert self.buf.as_buffer().tolist() == list(range(10))
+
+  def test_subbuffer_double_dealloc(self):
+    sub_buf = self.buf.view(3, dtypes.uint8, offset=4).ensure_allocated()
+    self.buf.deallocate()
+    with self.assertRaises(AssertionError):
+      self.buf.deallocate()
+    sub_buf.deallocate()
+    with self.assertRaises(AssertionError):
+      sub_buf.deallocate()
+
+  def test_subbuffer_uaf(self):
+    sub_buf = self.buf.view(4, dtypes.int8, offset=3).ensure_allocated()
+    assert self.buf.as_buffer().tolist(), list(range(10))
+    sub_buf.deallocate()
+    with self.assertRaises(AssertionError):
+      sub_buf.as_buffer().tolist()
+    assert self.buf.as_buffer().tolist(), list(range(10))
+
+    sub_buf = self.buf.view(4, dtypes.int8, offset=3).ensure_allocated()
+    assert sub_buf.as_buffer().tolist(), list(range(3, 7))
+    self.buf.deallocate()
+    with self.assertRaises(AssertionError):
+      sub_buf.as_buffer().tolist()
+
 if __name__ == '__main__':
   unittest.main()

test/test_tiny.py

@@ -100,7 +100,7 @@ class TestTiny(unittest.TestCase):
       lambda x: x.flatten(1), nn.Linear(576, 10)]
 
     # replace random weights with ones
-    for p in nn.state.get_parameters(layers): p.replace(Tensor.ones_like(p).contiguous()).realize()
+    Tensor.realize(*[p.replace(Tensor.ones_like(p).contiguous()) for p in nn.state.get_parameters(layers)])
 
     # run model inference
     probs = Tensor.rand(1, 1, 28, 28).sequential(layers).tolist()

test/test_uop_graph.py

@@ -423,6 +423,40 @@ class TestUOpGraph(unittest.TestCase):
       ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(Variable("i", 0, 20)),))
       with self.assertRaises(RuntimeError): to_uops_list([ld0])
 
+  def test_in_out_of_bounds_access_gated_store(self):
+    with Context(IGNORE_OOB=0):
+      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      v = Variable("v", 0, 20)
+      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])
+
+  def test_in_bounds_access_gated_local(self):
+    with Context(IGNORE_OOB=0):
+      # Define buffers
+      gbuf = UOp(Ops.DEFINE_GLOBAL, dtypes.uint.ptr(400), (), 0)
+      sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.uint.ptr(8, local=True), (), "temp0")
+
+      # Define indices, valids and barrier
+      gidx = UOp(Ops.SPECIAL, dtypes.int, (), ("gidx0", 416))
+      lidx = UOp(Ops.SPECIAL, dtypes.int, (), ("lidx0", 10))
+
+      gate = (gidx<400) & (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))
+
+      # Load from local memory (after the IF/barrier)
+      local_load = UOp(Ops.LOAD, dtypes.uint, (sbuf.index(lidx), if_barrier))
+
+      # Store to global memory
+      global_store = UOp(Ops.STORE, dtypes.void, (gbuf.index(gidx), local_load))
+      to_uops_list([global_store])
+
   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)

test/test_uops.py

@@ -10,12 +10,12 @@ from tinygrad.device import Buffer, Device
 from tinygrad.uop.ops import Ops, UOp, UPat, KernelInfo, exec_alu # noqa F401
 from tinygrad.uop.spec import spec
 from tinygrad.renderer import ProgramSpec
-from tinygrad.engine.kernelize import fix_kernel_ops
-from tinygrad.engine.realize import CompiledRunner
+from tinygrad.kernelize.kernelize import fix_kernel_ops
+from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.codegen import full_rewrite
 from tinygrad.uop.symbolic import sym
 from tinygrad.device import is_dtype_supported
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Opt, OptOps
 
 def to_uops_list(u:list[UOp], opts=None, skip_check=False) -> list[UOp]: return full_rewrite(UOp.sink(*u), opts)
 
@@ -409,9 +409,11 @@ class TestAssembly(unittest.TestCase):
     a = Tensor.empty(1024)
     b = Tensor.empty(1024)
     c = (a*b).sum()
-    k = Kernel(c.schedule()[-1].ast)
-    k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
-    uops = k.linearize().uops
+    ast = c.schedule()[-1].ast
+    opts_to_apply = [Opt(OptOps.UNROLL, 0, 4)]
+    ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
+    program = get_program(ast, Device[Device.DEFAULT].renderer)
+    uops = program.uops
     self.assertEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4)
 
 class TestUOpMethod(unittest.TestCase):

test/test_winograd.py

@@ -44,7 +44,6 @@ class TestWinograd(unittest.TestCase):
       with Timing(f"linearize {i} with {len(ops):4d} ops: "):
         l = Kernel(s.ast)
         l.apply_opts(hand_coded_optimizations(l))
-        l.linearize()
       assert len(l.sts) <= 256  # just the current value to prevent regression
       if DEBUG >= 2: print(f"{len(l.sts):4d} shapetrackers with max {max(len(x.views) for x in l.sts)} views")
       for st in l.sts:

test/unit/test_device.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-import unittest
+import unittest, os, subprocess, sys
 from tinygrad import Tensor
 from tinygrad.device import Device, Compiler
 from tinygrad.helpers import diskcache_get, diskcache_put, getenv, Context
@@ -53,5 +53,13 @@ class TestCompiler(unittest.TestCase):
       a = Tensor([0.,1.], device=Device.DEFAULT).realize()
       (a + 1).realize()
 
+class TestRunAsModule(unittest.TestCase):
+  def test_module_runs(self):
+    p = subprocess.run([sys.executable, "-m", "tinygrad.device"],stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+      env={**os.environ, "DEBUG": "1"}, timeout=10,)
+    out = (p.stdout + p.stderr).decode()
+    self.assertEqual(p.returncode, 0, msg=out)
+    self.assertIn("CPU", out) # for sanity check
+
 if __name__ == "__main__":
   unittest.main()

test/unit/test_graph_rewrite.py

@@ -290,5 +290,40 @@ class TestRecurse(unittest.TestCase):
     pm = PatternMatcher([(UPat(Ops.DEFINE_VAR, name="x"), lambda x: x)])
     graph_rewrite(a, pm)
 
+  def test_no_inf_loop_bottom_up(self):
+    a = UOp.variable('a', 0, 10)
+    pm = PatternMatcher([(UPat(Ops.DEFINE_VAR, name="x"), lambda x: x)])
+    graph_rewrite(a, pm, bottom_up=True)
+
+  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.DEFINE_REG)),
+      (UPat(Ops.DEFINE_REG, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
+    ])
+    with self.assertRaises(RuntimeError):
+      graph_rewrite(a, pm)
+
+  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.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)
+
+def bidir_append(ctx, x, b): ctx.append((x.arg if x.op is Ops.CONST else "+", b))
+class TestBidirectional(unittest.TestCase):
+  def test_simple(self):
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
+    c = a + b
+    pm = PatternMatcher([ (UPat(GroupOp.All, name="x"), lambda ctx,x: bidir_append(ctx, x, False)) ])
+    bpm = PatternMatcher([ (UPat(GroupOp.All, name="x"), lambda ctx,x: bidir_append(ctx, x, True)) ])
+    ctx_list = []
+    graph_rewrite(c, pm, ctx=ctx_list, bpm=bpm)
+    self.assertListEqual(ctx_list, [('+', True), (1, True), (1, False), (2, True), (2, False), ('+', False)])
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_helpers.py

@@ -1,9 +1,9 @@
 import ctypes, gzip, unittest
 from tinygrad import Variable
-from tinygrad.helpers import Context, ContextVar, argfix
+from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap
 from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
 from tinygrad.tensor import get_shape
-from tinygrad.codegen.lowerer import get_contraction, get_contraction_with_reduce
+from tinygrad.shape.view import get_contraction, get_contraction_with_reduce
 import numpy as np
 
 VARIABLE = ContextVar("VARIABLE", 0)
@@ -363,5 +363,18 @@ class TestArgFix(unittest.TestCase):
   def test_list(self):
     self.assertEqual(argfix([True, False]), (True, False))
 
+class TestWordWrap(unittest.TestCase):
+  def test_wrap_simple(self):
+    wrap = 10
+    st = "x"*wrap*2
+    st2 = word_wrap(st, wrap)
+    self.assertEqual(len(st2.splitlines()), 2)
+
+  def test_wrap_colored(self):
+    wrap = 10
+    st = colored("x"*wrap*2, "red")
+    st2 = word_wrap(st, wrap=wrap)
+    self.assertEqual(len(st2.splitlines()), 2)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_linearizer_rewrite.py

@@ -1,6 +1,8 @@
 import unittest
 from tinygrad import Tensor, Context, Device
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.engine.realize import get_program
+from tinygrad.renderer import Opt, OptOps
+from tinygrad.uop.ops import KernelInfo
 
 class TestLinearizerRewrite(unittest.TestCase):
   def test_reduction(self):
@@ -8,20 +10,22 @@ class TestLinearizerRewrite(unittest.TestCase):
     out = (t*2).sum(axis=1)
     with Context(SPLIT_REDUCEOP=0, DEVECTORIZE=0):
       si = out.schedule()[-1]
-      k = Kernel(si.ast, Device["CPU"].renderer)
-      k.apply_opt(Opt(OptOps.UPCAST, 0, 4))
-      k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
-      prg = k.to_program()
+      opts_to_apply = []
+      opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4))
+      opts_to_apply.append(Opt(OptOps.UNROLL, 0, 4))
+      ast = si.ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
+      prg = get_program(ast, Device["CPU"].renderer)
       print(prg.src)
 
   def test_arange(self):
     out = Tensor.arange(32, device="NULL")
     with Context(SPLIT_REDUCEOP=0, DEVECTORIZE=0):
       si = out.schedule()[-1]
-      k = Kernel(si.ast, Device["CPU"].renderer)
-      k.apply_opt(Opt(OptOps.UPCAST, 0, 4))
-      k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
-      prg = k.to_program()
+      opts_to_apply = []
+      opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4))
+      opts_to_apply.append(Opt(OptOps.UNROLL, 0, 4))
+      ast = si.ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
+      prg = get_program(ast, Device["CPU"].renderer)
       print(prg.src)
 
 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.engine.kernelize import sym, merge_views
+from tinygrad.kernelize.kernelize import sym, merge_views
 
 class TestRewriteTrackedChildren(unittest.TestCase):
   @unittest.skip("track_children no longer supported")

test/unit/test_tensor_uop_representation.py

@@ -78,42 +78,19 @@ class TestTensorUopRepresentation(unittest.TestCase):
     is_pattern(a, const_pattern)
     self.assertEqual(a.uop.shape, (10, 10))
 
-  # currently, CONSTs have a "fake" BUFFER. this should be fixed
-  # current:
-  # UOp(Ops.EXPAND, dtypes.float, arg=(10, 10), src=(
-  #   UOp(Ops.RESHAPE, dtypes.float, arg=(1, 1), src=(
-  #     UOp(Ops.VIEW, dtypes.float, arg=ShapeTracker(views=(View(shape=(), strides=(), offset=0, mask=None, contiguous=True),)), src=(
-  #       UOp(Ops.BUFFER, dtypes.float, arg=(-1, 'METAL', 1), src=()),
-  #       UOp(Ops.CONST, dtypes.float, arg=1.0, src=()),)),)),))
-  # expected:
-  # UOp(Ops.EXPAND, dtypes.float, arg=(10, 10), src=(
-  #   UOp(Ops.RESHAPE, dtypes.float, arg=(1, 1), src=(
-  #     UOp(Ops.VIEW, dtypes.float, arg=ShapeTracker(views=(View(shape=(), strides=(), offset=0, mask=None, contiguous=True),)), src=(
-  #       UOp(Ops.CONST, dtypes.float, arg=1.0, src=(
-  #         UOp(Ops.DEVICE, dtypes.void, arg="METAL", src=()),)),)),))
+  # CONST is EXPAND -> RESHAPE -> CONST -> DEVICE
   def test_consts_dont_have_buffers(self):
     a = Tensor.ones(10, 10)
-    print(a.uop)
     buffers_in_parents = [x.op for x in a.uop.toposort() if x.op is Ops.BUFFER]
     self.assertEqual(len(buffers_in_parents), 0)
+    is_pattern(a, UPat(Ops.EXPAND, src=(UPat(Ops.RESHAPE, src=(const_pattern,)),)))
 
-  # currently, COPY has an extra BUFFER on the output
-  # current:
-  # UOp(Ops.VIEW, dtypes.float, arg=ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True),)), src=(
-  #   UOp(Ops.BUFFER, dtypes.float, arg=(2, 'TEST', 3), src=()),
-  #   UOp(Ops.COPY, dtypes.float, arg=('TEST', False), src=(
-  #     UOp(Ops.VIEW, dtypes.float, arg=ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True),)), src=(
-  #       UOp(Ops.BUFFER, dtypes.float, arg=(1, 'METAL', 3), src=()),)),)),))
-  # expected:
-  # UOp(Ops.COPY, dtypes.float, arg=('TEST', False), src=(
-  #   UOp(Ops.VIEW, dtypes.float, arg=ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True),)), src=(
-  #     UOp(Ops.BUFFER, dtypes.float, arg=(1, 'METAL', 3), src=()),))
-  # update: now the arg is just a single bool, the first source is a device.
+  # COPY has a copyin source and a device.
   def test_copyin(self):
     a = Tensor([1.,2,3]).realize()
     c = a.to("TEST")   # NOTE: this isn't checked
     print(c.uop)
-    is_pattern(c, UPat(Ops.COPY, src=(realized_pattern, UPat(Ops.DEVICE))))
+    is_pattern(c, UPat(Ops.COPY, src=(realized_pattern, UPat(Ops.DEVICE)), arg=None))
 
   def test_empty_buf(self):
     a = Tensor.empty(3, 3)

test/unit/test_uop_spec.py

@@ -2,26 +2,26 @@ from __future__ import annotations
 import unittest
 
 from tinygrad import Tensor
-from tinygrad.opt.kernel import Kernel
 from tinygrad.helpers import DEBUG
 from tinygrad.uop.ops import UOp, Ops, print_uops
 from tinygrad.uop.spec import type_verify, ast_spec, tensor_uop_spec
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad import dtypes
 from tinygrad.shape.view import View
+from tinygrad.engine.realize import get_program
+from tinygrad.device import Device
 
 class InvalidASTException(Exception): pass
-def helper_test_verify_ast(*stores:UOp) -> Kernel:
+def helper_test_verify_ast(*stores:UOp):
   sink = UOp(Ops.SINK, dtypes.void, stores)
   if DEBUG >= 3:
     for op in stores: print(op)
   try: type_verify(list(sink.toposort()), ast_spec)
   except RuntimeError as e: raise InvalidASTException(e.args)
-  k = Kernel(sink)
-  k.linearize()
-  if DEBUG >= 6: print_uops(k.uops)
-  if DEBUG >= 4: print(k.to_program().src)
-  return k
+  program = get_program(sink, Device[Device.DEFAULT].renderer)
+
+  if DEBUG >= 6: print_uops(program.uops)
+  if DEBUG >= 4: print(program.src)
 
 class TestUOpSpec(unittest.TestCase):
   def test_tiny_add(self):

test/unit/test_uop_symbolic.py

@@ -193,6 +193,17 @@ class TestSymbolic(unittest.TestCase):
   def test_sum_div_no_factor(self):
     self.helper_test_variable(usum([Variable("a", 0, 7)*5, Variable("b", 0, 3)*5]) // 2, 0, 25, "(((a*5)+(b*5))//2)")
 
+  def test_mod_min_max(self):
+    self.helper_test_variable(Variable("x", 0, 10)%Variable("y", 1, 10), 0, 9, "(x%y)")
+    self.helper_test_variable(Variable("x", -10, 0)%Variable("y", 1, 10), -9, 0, "(((x*-1)%y)*-1)")
+    self.helper_test_variable(Variable("x", 0, 10)%Variable("y", -10, -1), 0, 9, "(x%y)")
+    self.helper_test_variable(Variable("x", -10, 0)%Variable("y", -10, -1), -9, 0, "(((x*-1)%y)*-1)")
+    self.helper_test_variable(Variable("x", -10, 10)%Variable("y", -10, -1), -9, 9, "(x%y)")
+
+    # test _min_max directly without the rewrite taking out the sign
+    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_mod_factor(self):
     self.helper_test_variable(usum([Variable("a", 0, 7)*100, Variable("b", 0, 3)*50]) % 100, 0, 50, "((b%2)*50)")
 

test/unit/test_viz.py

@@ -1,271 +1,242 @@
 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.uop.ops import TRACK_MATCH_STATS, TrackedPatternMatcher, UOp, graph_rewrite, track_rewrites, UPat, Ops
-from tinygrad.uop.symbolic import symbolic, symbolic_simple
-from tinygrad.uop.ops import tracked_ctxs as contexts, tracked_keys as keys, _name_cnt, _substitute
-from tinygrad.device import ProfileDeviceEvent, ProfileRangeEvent, ProfileGraphEvent, ProfileGraphEntry
-from tinygrad.viz.serve import get_metadata, get_details, uop_to_json, to_perfetto
+from tinygrad.helpers import PROFILE, colored, ansistrip
+from tinygrad.device import Buffer
 
-# NOTE: VIZ tests always use the tracked PatternMatcher instance
-symbolic = TrackedPatternMatcher(symbolic.patterns)
-substitute = TrackedPatternMatcher(_substitute.patterns)
+@track_rewrites(name=True)
+def exec_rewrite(sink:UOp, pm_lst:list[PatternMatcher], names:None|list[str]=None) -> UOp:
+  for i,pm in enumerate(pm_lst):
+    sink = graph_rewrite(sink, TrackedPatternMatcher(pm.patterns), name=names[i] if names else None)
+  return sink
 
-inner_rewrite = TrackedPatternMatcher([
-  (UPat.cvar("x"), lambda x: None if x.dtype == dtypes.float32 else UOp.const(dtypes.float32, x.arg)),
-])
-
-l2 = TrackedPatternMatcher([(UPat(Ops.CUSTOM, arg=2, name="x"), lambda x: x.replace(arg=3))])
-l1 = TrackedPatternMatcher([(UPat(Ops.CUSTOM, arg=1, name="x"), lambda x: graph_rewrite(x.replace(arg=2), l2))])
-l0 = TrackedPatternMatcher([(UPat(Ops.CUSTOM, arg=0, name="x"), lambda x: graph_rewrite(x.replace(arg=1), l1))])
+# real VIZ=1 pickles these tracked values
+from tinygrad.viz.serve import get_metadata, uop_to_json
+from tinygrad.uop.ops import tracked_keys, tracked_ctxs, active_rewrites, _name_cnt
+def get_viz_list(): return get_metadata(tracked_keys, tracked_ctxs)
 
 class TestViz(unittest.TestCase):
   def setUp(self):
     # clear the global context
-    contexts.clear()
-    keys.clear()
-    _name_cnt.clear()
+    for lst in [tracked_keys, tracked_ctxs, active_rewrites, _name_cnt]: lst.clear()
+    Buffer.profile_events.clear()
     self.tms = TRACK_MATCH_STATS.value
+    self.profile = PROFILE.value
     TRACK_MATCH_STATS.value = 2
-  def tearDown(self): TRACK_MATCH_STATS.value = self.tms
+    PROFILE.value = 1
+  def tearDown(self):
+    TRACK_MATCH_STATS.value = self.tms
+    PROFILE.value = self.profile
 
-  def test_viz_simple(self):
+  def test_simple(self):
     a = UOp.variable("a", 0, 10)
-    @track_rewrites(name=True)
-    def test(sink): return graph_rewrite(sink, symbolic)
-    test(a*1)
-    ret = get_metadata(keys, contexts)
-    self.assertEqual(len(ret), 1)
-    key, val = ret[0]["name"], ret[0]["steps"]
-    self.assertEqual(key, "test n1")
-    self.assertEqual(val[0]["match_count"], 1)
+    exec_rewrite((a+0)*1, [sym])
+    lst = get_viz_list()
+    # VIZ displays rewrites in groups of tracked functions
+    self.assertEqual(len(lst), 1)
+    # each group has a list of steps
+    self.assertEqual(len(lst[0]["steps"]), 1)
+    # each step has a list of matches
+    self.assertEqual(lst[0]["steps"][0]["match_count"], 2)
 
-  def test_track_two_rewrites(self):
+  def test_rewrites(self):
     a = UOp.variable("a", 0, 10)
-    @track_rewrites(name=True)
-    def test(sink): return graph_rewrite(sink, symbolic)
-    test((a+a)*1)
-    ret = get_metadata(keys, contexts)
-    key, val = ret[0]["name"], ret[0]["steps"]
-    self.assertEqual(len(ret), 1)              # one context
-    self.assertEqual(len(val), 1)              # one graph_rewrite call in context
-    self.assertEqual(key, "test n1")
-    self.assertEqual(val[0]["match_count"], 2) # two upats applied
+    exec_rewrite(a*1, [sym])
+    exec_rewrite(a*2, [sym])
+    lst = get_viz_list()
+    self.assertEqual(len(lst), 2)
+    # names dedup using a counter
+    self.assertEqual(lst[0]["name"], "exec_rewrite n1")
+    self.assertEqual(lst[1]["name"], "exec_rewrite n2")
 
-  def test_track_multiple_calls_one_ctx(self):
+  def test_steps(self):
     a = UOp.variable("a", 0, 10)
-    @track_rewrites(name=True)
-    def test(a, b):
-      a = graph_rewrite(a, symbolic)
-      b = graph_rewrite(b, symbolic)
-    test(a*1, a*5)
-    ret = get_metadata(keys, contexts)
-    key, val = ret[0]["name"], ret[0]["steps"]
-    self.assertEqual(len(ret), 1)              # one context
-    self.assertEqual(len(val), 2)              # two graph_rewrite calls in context
-    self.assertEqual(key, "test n1")
-    self.assertEqual(val[0]["match_count"], 1) # one rewrite for a*0
-    self.assertEqual(val[1]["match_count"], 0) # no rewrites for a*5
+    exec_rewrite(a+1, [PatternMatcher([]), PatternMatcher([])], ["x", "y"])
+    steps = get_viz_list()[0]["steps"]
+    # steps can optionally have a name
+    self.assertEqual(steps[0]["name"], "x")
+    self.assertEqual(steps[1]["name"], "y")
 
-  def test_track_rewrites(self):
+  def test_rewrite_location(self):
+    def inner(sink): return graph_rewrite(sink, PatternMatcher([]))
     @track_rewrites(name=True)
-    def do_rewrite(x:UOp): return graph_rewrite(x, symbolic)
-    a = UOp.variable("a", 0, 10)
-    b = UOp.variable("b", 0, 4)
-    do_rewrite(a*1)
-    do_rewrite(a*b)
-    ret = get_metadata(keys, contexts)
-    self.assertEqual(len(ret), 2)
-    key, m = ret[0]["name"], ret[0]["steps"]
-    self.assertEqual(key, "do_rewrite n1")
-    self.assertEqual(m[0]["match_count"], 1)
-    key, m = ret[1]["name"], ret[1]["steps"]
-    self.assertEqual(key, "do_rewrite n2")
-    self.assertEqual(m[0]["match_count"], 0)
+    def outer(sink): return inner(sink)
+    outer(UOp.variable("a", 1, 10))
+    lst = get_viz_list()
+    # step location comes from inner rewrite
+    fp, lineno = lst[0]["steps"][0]["loc"]
+    self.assertEqual(fp, inner.__code__.co_filename)
+    self.assertEqual(lineno, inner.__code__.co_firstlineno)
 
-  def test_track_rewrites_with_exception(self):
+  def test_exceptions(self):
+    # VIZ tracks rewrites up to the error
+    def count_3(x:UOp):
+      assert x.arg <= 3
+      return x.replace(arg=x.arg+1)
+    err_pm = PatternMatcher([(UPat.cvar("x"), count_3),])
+    a = UOp.const(dtypes.int, 1)
+    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"], 3)
+
+  def test_default_name(self):
+    a = UOp.variable("a", 1, 10)
     @track_rewrites()
-    def do_rewrite(x:UOp):
-      x = graph_rewrite(x, symbolic) # NOTE: viz tracks this
-      raise Exception("test")
-    a = UOp.variable("a", 0, 10)
-    with self.assertRaises(Exception): do_rewrite(a*1)
-    ret = get_metadata(keys, contexts)
-    self.assertEqual(len(ret), 1)
+    def name_default(): return graph_rewrite(a, PatternMatcher([]))
+    name_default()
+    lst = get_viz_list()
+    self.assertEqual(lst[0]["name"], "name_default n1")
 
-  def test_track_rewrites_name_fxn(self):
-    @track_rewrites(name=lambda _,ret: f"output_{ret}")
-    def do_rewrite(x:UOp):
-      x = graph_rewrite(x, symbolic)
-      return x.render()
-    expr = UOp.variable("a",0,10)*UOp.variable("b",0,10)
-    do_rewrite(expr)
-    key = get_metadata(keys, contexts)[0]["name"]
-    self.assertEqual(key, "output_(a*b) n1")
-
-    expr2 = UOp.variable("a",0,10)+UOp.variable("b",0,10)
-    do_rewrite(expr2)
-    key = get_metadata(keys, contexts)[1]["name"]
-    self.assertEqual(key, "output_(a+b) n2")
-
-  @unittest.expectedFailure
-  def test_name_in_positional_arg(self):
-    @track_rewrites(name=True)
-    def test(sink): return graph_rewrite(sink, symbolic, None, False, "name")
-    test(UOp.variable("a", 0, 1))
-    self.assertEqual(contexts[0].pop().name, "name")
-
-  # NOTE: CONST UOps do not get nodes in the graph
-  def test_dont_create_const_nodes(self):
-    a = UOp.variable("a", 0, 10)
-    b = UOp.variable("b", 0, 4)
-    self.assertEqual(len(uop_to_json(a*1)), 2)
-    self.assertEqual(len(uop_to_json(a*b)), 3)
-
-  def test_bottom_up_rewrite(self):
-    a = UOp.variable("a", 0, 10)
-    b = UOp.variable("b", 0, 10)
-    c = UOp.variable("c", 0, 10)
-    @track_rewrites(name=True)
-    def fxn(sink): return graph_rewrite(sink, substitute, ctx={a+b:c}, bottom_up=True)
-    fxn(a+b)
-    #UOp.substitute(a+b, {a+b:c})
-    ret = get_metadata(keys, contexts)
-    self.assertEqual(len(ret), 1)
-    m = ret[0]["steps"]
-    self.assertEqual(m[0]["match_count"], 1)
-
-  # NOTE: calling graph_rewrite when the function isn't decorated with track_rewrites should not VIZ
-  def test_rewrite_without_context(self):
-    def untracked_graph_rewrite(sink): return graph_rewrite(sink, symbolic)
-    @track_rewrites(name=True)
-    def tracked_graph_rewrite(sink): return graph_rewrite(sink, symbolic)
-    # test
-    add = UOp.const(dtypes.int, 2) + UOp.const(dtypes.int, 1)
-    untracked_graph_rewrite(add)
-    self.assertEqual(len(contexts), 0)
-    tracked_graph_rewrite(add)
-    self.assertEqual(len(contexts), 1)
-
-  def test_inner_rewrite_location(self):
-    # inner rewrite gets tracked in another context
-    def inner_rewrite(sink): return graph_rewrite(sink, symbolic)
-    @track_rewrites(name=True)
-    def tracked_graph_rewrite(sink): return inner_rewrite(sink)
-    # test
-    add = UOp.const(dtypes.int, 2) + UOp.const(dtypes.int, 1)
-    tracked_graph_rewrite(add)
-    self.assertEqual(len(contexts), 1)
-    # location of context is inner_rewrite
-    fp, lineno = contexts[0][0].loc
-    self.assertEqual(lineno, inner_rewrite.__code__.co_firstlineno)
-    self.assertEqual(fp, inner_rewrite.__code__.co_filename)
-
-  def test_upat_location(self):
-    for (pat, fn) in symbolic_simple.patterns:
-      self.assertIn("symbolic.py", pat.location[0])
-
-  def test_nested_rewrite(self):
-    def make_float(x:UOp, y:UOp):
-      if x.dtype == dtypes.float: return None
-      x2 = graph_rewrite(x, inner_rewrite, name="inner_x")
-      y2 = graph_rewrite(y, inner_rewrite, name="inner_y")
-      return None if (x2 is x and y2 is y) else x2+y2
-    outer_rewrite = TrackedPatternMatcher([(UPat.cvar("x")+UPat.cvar("y"), make_float),])
-    @track_rewrites(name=True)
-    def rewrite(u:UOp): return graph_rewrite(u, outer_rewrite, name="outer")
-    a = UOp.const(dtypes.int, 1)+UOp.const(dtypes.int, 2)
-    rewrite(a)
-    self.assertEqual(len(contexts), 1)
-    tracked = contexts[0]
-    self.assertEqual(len(tracked), 3)
-    self.assertEqual(tracked[0].depth, 0)
-    self.assertEqual(tracked[1].depth, 1)
-    self.assertEqual(tracked[2].depth, 1)
-    # NOTE: this is sorted by the time called, maybe it should be by depth
-    self.assertEqual([x.name for x in tracked], ["outer", "inner_x", "inner_y"])
-    self.assertEqual([len(x.matches) for x in tracked], [1, 1, 1])
-
-  def test_depth_level(self):
-    @track_rewrites(name=True)
-    def fxn(u:UOp): return graph_rewrite(u, l0)
-    ret = fxn(UOp(Ops.CUSTOM, arg=0))
-    assert ret is UOp(Ops.CUSTOM, arg=3)
-    self.assertEqual(len(contexts), 1)
-    tracked = contexts[0]
-    self.assertEqual(tracked[0].depth, 0)
-    self.assertEqual(tracked[1].depth, 1)
-    self.assertEqual(tracked[2].depth, 2)
-
-  def test_shape_label(self):
-    a = UOp.new_buffer("CPU", 1, dtypes.uint8).expand((4,))
-    b = UOp.new_buffer("CPU", 1, dtypes.uint8).expand((8,))
-    n = a+b
-    ser = uop_to_json(n)
-    self.assertIn("(4,)", ser[id(a)]["label"])
-    self.assertIn("(8,)", ser[id(b)]["label"])
-    with self.assertRaises(AssertionError): n.st
-    _  = ser[id(n)]["label"] # VIZ should not crash
-
-  def test_default_named(self):
-    test = UOp(Ops.NOOP)
+  # name can also be the first arg
+  def test_self_name(self):
     @track_rewrites()
-    def test_fxn(): return graph_rewrite(test, l0)
-    assert test_fxn() is test
-    self.assertEqual(keys[0], "test_fxn n1")
+    def name_is_self(s:UOp): return graph_rewrite(s, PatternMatcher([]))
+    name_is_self(arg:=UOp.variable("a", 1, 10))
+    lst = get_viz_list()
+    self.assertEqual(lst[0]["name"], str(arg))
 
-  @unittest.skip("TODO: doesn't work")
-  def test_recursion_err(self):
-    inf = TrackedPatternMatcher([
-      (UPat.const(dtypes.int, 0).name("a"), lambda a: a.const_like(1)),
-      (UPat.const(dtypes.int, 1).name("b"), lambda b: b.const_like(0)),
-    ])
-    @track_rewrites(name=True)
-    def func(u): return graph_rewrite(u, inf)
-    with self.assertRaises(RecursionError): func(UOp.const(dtypes.int, 0))
-    _ = list(get_details(keys[0], contexts[0][0]))
+  # name can also come from a function
+  def test_dyn_name_fxn(self):
+    @track_rewrites(name=lambda a,ret: a.render())
+    def name_from_fxn(s:UOp): return graph_rewrite(s, PatternMatcher([]))
+    name_from_fxn(UOp.variable("a", 1, 10)+1)
+    lst = get_viz_list()
+    self.assertEqual(lst[0]["name"], "(a+1) n1")
 
-class TextVizProfiler(unittest.TestCase):
+  def test_colored_label(self):
+    # NOTE: dataclass repr prints literal escape codes instead of unicode chars
+    @dataclass(frozen=True)
+    class TestStruct:
+      colored_field: str
+    a = UOp(Ops.CUSTOM, arg=TestStruct(colored("xyz", "magenta")+colored("12345", "blue")))
+    a2 = uop_to_json(a)[id(a)]
+    self.assertEqual(ansistrip(a2["label"]), f"CUSTOM\n{TestStruct.__qualname__}(colored_field='xyz12345')")
+
+# VIZ displays nested graph_rewrites in a tree view
+
+def leaf_rewrite(x:UOp): return x.rtag(1) if x.tag is None else None
+leaf = TrackedPatternMatcher([(UPat(Ops.DEFINE_VAR, name="x"), leaf_rewrite)])
+
+def branch_rewrite(x:UOp, y:UOp):
+  if x.tag is not None: return
+  x2 = graph_rewrite(x, leaf, name="leaf_left")
+  y2 = graph_rewrite(y, leaf, name="leaf_right")
+  return x2 * y2
+branch = TrackedPatternMatcher([(UPat.var("x")+UPat.var("y"), branch_rewrite)])
+
+def root_rewrite(root:UOp):
+  new_src = tuple(graph_rewrite(b, branch, name=f"branch_{i}") for i,b in enumerate(root.src))
+  return root.replace(src=new_src)
+root = TrackedPatternMatcher([(UPat(Ops.SINK, src=UPat(Ops.ADD), name="root"), root_rewrite),])
+
+class TestVizTree(TestViz):
+  def assertStepEqual(self, step:dict, want:dict):
+    for k,v in want.items():
+      self.assertEqual(step[k], v, f"failed at '{k}': {v} != {step[k]}\n{step=}")
+
+  def test_tree_view(self):
+    a = UOp.variable("a",0,10)
+    b = UOp.variable("b",0,10)
+    c = UOp.variable("c",0,10)
+    d = UOp.variable("d",0,10)
+    sink = UOp.sink(a+b, c+d)
+    @track_rewrites()
+    def tree_rewrite(): return graph_rewrite(sink, root, name="root")
+    tree_rewrite()
+    lst = get_viz_list()
+    steps = lst[0]["steps"]
+    self.assertEqual(len(steps), 1+2+4)
+    self.assertStepEqual(steps[0], {"name":"root", "depth":0, "match_count":1})
+    self.assertStepEqual(steps[1], {"name":"branch_0", "depth":1, "match_count":1})
+    self.assertStepEqual(steps[2], {"name":"leaf_left", "depth":2, "match_count":1})
+    self.assertStepEqual(steps[3], {"name":"leaf_right", "depth":2, "match_count":1})
+    self.assertStepEqual(steps[4], {"name":"branch_1", "depth":1, "match_count":1})
+    self.assertStepEqual(steps[5], {"name":"leaf_left", "depth":2, "match_count":1})
+    self.assertStepEqual(steps[6], {"name":"leaf_right", "depth":2, "match_count":1})
+
+import gc
+
+def bufs_allocated() -> int:
+  gc.collect()
+  return sum([isinstance(x, Buffer) for x in gc.get_objects()])
+
+class TestVizGC(TestViz):
+  def test_gc(self):
+    init = bufs_allocated()
+    a = UOp.new_buffer("NULL", 10, dtypes.char)
+    a.buffer.allocate()
+    exec_rewrite(a, [PatternMatcher([])])
+    del a
+    self.assertEqual(bufs_allocated()-init, 0)
+    lst = get_viz_list()
+    self.assertEqual(len(lst), 1)
+
+  @unittest.skip("it's not generic enough to handle arbitrary UOps in arg")
+  def test_gc_uop_in_arg(self):
+    init = bufs_allocated()
+    a = UOp.new_buffer("NULL", 10, dtypes.char)
+    a.buffer.allocate()
+    exec_rewrite(UOp(Ops.CUSTOM, src=(a,), arg=a), [PatternMatcher([])])
+    del a
+    self.assertEqual(bufs_allocated()-init, 0)
+    lst = get_viz_list()
+    self.assertEqual(len(lst), 1)
+
+# VIZ integrates with other parts of tinygrad
+
+from tinygrad import Tensor, Device
+from tinygrad.engine.realize import get_program
+
+class TestVizIntegration(TestViz):
+  # kernelize has a custom name function in VIZ
+  def test_kernelize_tracing(self):
+    a = Tensor.empty(4, 4)
+    Tensor.kernelize(a+1, a+2)
+    lst = get_viz_list()
+    self.assertEqual(len(lst), 1)
+    self.assertEqual(lst[0]["name"], "Schedule 2 Kernels n1")
+
+  # codegen supports rendering of code blocks
+  def test_codegen_tracing(self):
+    ast = Tensor.schedule(Tensor.empty(4)+Tensor.empty(4))[0].ast
+    prg = get_program(ast, Device[Device.DEFAULT].renderer)
+    lst = get_viz_list()
+    self.assertEqual(len(lst), 2)
+    self.assertEqual(lst[0]["name"], "Schedule 1 Kernel n1")
+    self.assertEqual(lst[1]["name"], prg.name)
+
+from tinygrad.device import ProfileDeviceEvent, ProfileRangeEvent, ProfileGraphEvent, ProfileGraphEntry
+from tinygrad.viz.serve import get_profile
+
+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 = json.loads(to_perfetto(prof))
+    j = json.loads(get_profile(prof))
 
-    # Device regs always first
-    self.assertEqual(j['traceEvents'][0]['name'], 'process_name')
-    self.assertEqual(j['traceEvents'][0]['ph'], 'M')
-    self.assertEqual(j['traceEvents'][0]['args']['name'], 'NV')
-
-    self.assertEqual(j['traceEvents'][1]['name'], 'thread_name')
-    self.assertEqual(j['traceEvents'][1]['ph'], 'M')
-    self.assertEqual(j['traceEvents'][1]['pid'], j['traceEvents'][0]['pid'])
-    self.assertEqual(j['traceEvents'][1]['tid'], 0)
-    self.assertEqual(j['traceEvents'][1]['args']['name'], 'COMPUTE')
-
-    self.assertEqual(j['traceEvents'][2]['name'], 'thread_name')
-    self.assertEqual(j['traceEvents'][2]['ph'], 'M')
-    self.assertEqual(j['traceEvents'][2]['pid'], j['traceEvents'][0]['pid'])
-    self.assertEqual(j['traceEvents'][2]['tid'], 1)
-    self.assertEqual(j['traceEvents'][2]['args']['name'], 'COPY')
-
-    self.assertEqual(j['traceEvents'][3]['name'], 'E_2')
-    self.assertEqual(j['traceEvents'][3]['ts'], 0)
-    self.assertEqual(j['traceEvents'][3]['dur'], 10)
-    self.assertEqual(j['traceEvents'][3]['ph'], 'X')
-    self.assertEqual(j['traceEvents'][3]['pid'], j['traceEvents'][0]['pid'])
-    self.assertEqual(j['traceEvents'][3]['tid'], 0)
+    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)
 
   def test_perfetto_copy_node(self):
     prof = [ProfileRangeEvent(device='NV', 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))]
 
-    j = json.loads(to_perfetto(prof))
+    j = json.loads(get_profile(prof))
 
-    self.assertEqual(j['traceEvents'][3]['name'], 'COPYxx')
-    self.assertEqual(j['traceEvents'][3]['ts'], 900) # diff clock
-    self.assertEqual(j['traceEvents'][3]['dur'], 10)
-    self.assertEqual(j['traceEvents'][3]['ph'], 'X')
-    self.assertEqual(j['traceEvents'][3]['tid'], 1)
+    event = j['layout']['NV']['timeline']['shapes'][0]
+    self.assertEqual(event['name'], 'COPYxx')
+    self.assertEqual(event['st'], 900) # diff clock
+    self.assertEqual(event['dur'], 10)
 
   def test_perfetto_graph(self):
     prof = [ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100)),
@@ -275,25 +246,64 @@ class TextVizProfiler(unittest.TestCase):
                               deps=[[], [0]],
                               sigs=[decimal.Decimal(1000), decimal.Decimal(1002), decimal.Decimal(1004), decimal.Decimal(1008)])]
 
-    j = json.loads(to_perfetto(prof))
+    j = json.loads(get_profile(prof))
 
-    # Device regs always first
-    self.assertEqual(j['traceEvents'][0]['args']['name'], 'NV')
-    self.assertEqual(j['traceEvents'][1]['args']['name'], 'COMPUTE')
-    self.assertEqual(j['traceEvents'][2]['args']['name'], 'COPY')
-    self.assertEqual(j['traceEvents'][3]['args']['name'], 'NV:1')
-    self.assertEqual(j['traceEvents'][4]['args']['name'], 'COMPUTE')
-    self.assertEqual(j['traceEvents'][5]['args']['name'], 'COPY')
+    devices = list(j['layout'])
+    self.assertEqual(devices[0], 'NV')
+    self.assertEqual(devices[1], 'NV:1')
 
-    self.assertEqual(j['traceEvents'][6]['name'], 'E_25_4n2')
-    self.assertEqual(j['traceEvents'][6]['ts'], 0)
-    self.assertEqual(j['traceEvents'][6]['dur'], 2)
-    self.assertEqual(j['traceEvents'][6]['pid'], j['traceEvents'][0]['pid'])
+    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'])
 
-    self.assertEqual(j['traceEvents'][7]['name'], 'NV -> NV:1')
-    self.assertEqual(j['traceEvents'][7]['ts'], 954)
-    self.assertEqual(j['traceEvents'][7]['dur'], 4)
-    self.assertEqual(j['traceEvents'][7]['pid'], j['traceEvents'][3]['pid'])
+    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'])
+
+def _alloc(b:int):
+  a = Tensor.empty(b, device="NULL", dtype=dtypes.char)
+  a.uop.buffer.allocate()
+  return a
+
+class TestVizMemoryLayout(TestViz):
+  def test_double_alloc(self):
+    a = _alloc(1)
+    _b = _alloc(1)
+    profile_ret = json.loads(get_profile(Buffer.profile_events))
+    ret = profile_ret["layout"][a.device]["mem"]
+    self.assertEqual(ret["peak"], 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 = json.loads(get_profile(Buffer.profile_events))
+    ret = profile_ret["layout"][b.device]["mem"]
+    self.assertEqual(ret["peak"], 1)
+    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 = json.loads(get_profile(Buffer.profile_events))
+    ret = profile_ret["layout"][c.device]["mem"]
+    self.assertEqual(ret["peak"], 2)
+    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()

tinygrad/codegen/__init__.py

@@ -7,7 +7,9 @@ from tinygrad.uop.spec import type_verify
 from tinygrad.renderer import Renderer
 
 # import all pattern matchers here
-from tinygrad.codegen.lowerer import pm_quant, pm_lowerer, get_index
+from tinygrad.codegen.lowerer import pm_lowerer, get_index
+from tinygrad.codegen.quantize import pm_quant
+from tinygrad.codegen.gpudims import pm_add_gpudims
 from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing
 from tinygrad.codegen.expander import migrate_indexing, expander
 from tinygrad.codegen.devectorizer import load_store_folding, load_store_indexing, devectorize, \
@@ -35,12 +37,23 @@ def get_rewrites_for_renderer(opts:Renderer, linearizer:bool=True) -> list[Rewri
   # cache with the values of the context vars
   return _get_rewrites_for_renderer(opts, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value)
 
+# tensor cores
+
+from tinygrad.uop.ops import PatternMatcher, UPat, UOp
+
+def tensor_cores(a:UOp, b:UOp, r:UOp):
+  print("use tensor cores")
+
+pm_tensor_cores = PatternMatcher([
+  ((UPat.var().gep(name='a') * UPat.var().gep(name='b')).reduce(name='r', allow_any_len=True), tensor_cores),
+])
+
 @functools.cache
 def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL) -> list[RewriteStep]:
   # ** lowerer (rewrite_shapetracker_with_index) **
   ret: list[RewriteStep] = []
   if _QUANTIZE and opts.device in {"CPU", "DSP"}: ret.append(RewriteStep(pm_quant, name="quantize"))
-  ret.append(RewriteStep(pm_lowerer, lambda ast: get_index(ast, opts), name="lowerer"))
+  ret.append(RewriteStep(pm_lowerer, get_index, name="lowerer", bottom_up=True))
 
   # ** expander (expand_rewrite) **
   ret.append(RewriteStep(sym+migrate_indexing, name="initial symbolic"))
@@ -48,10 +61,19 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
   # expand
   ret.append(RewriteStep(sym+expander, name="expander"))
 
+  # use tensor cores
+  ret.append(RewriteStep(pm_tensor_cores, name="tensor cores"))
+
   # ** devectorizer (full_graph_rewrite) **
   # remove reduce
   ret.append(RewriteStep(pm_reduce+gep_pushing, lambda _: ReduceContext(), name="remove_reduce"))
 
+  # factorize warp (before gpu dims)
+  #ret.append(RewriteStep(pm_warp, name="warpcast"))
+
+  # add gpu dims (late)
+  ret.append(RewriteStep(pm_add_gpudims, lambda _: opts, name="add gpudims"))
+
   # devectorize (TODO: does this need opts?)
   if _DEVECTORIZE >= 2: pm_devectorize = sym+load_store_folding+load_store_indexing
   elif _DEVECTORIZE: pm_devectorize = sym+devectorize+load_store_folding+correct_load_store+load_store_indexing
@@ -75,6 +97,17 @@ def full_rewrite_to_sink(sink:UOp, opts:Renderer|None=None, linearizer:bool=Fals
   return apply_rewrites(sink, get_rewrites_for_renderer(opts if opts is not None else Renderer(), linearizer))
 
 def full_rewrite(sink:UOp, opts:Renderer|None=None) -> list[UOp]:
+  """
+  Function to transform the Kernel UOp graph into a linearized program.
+
+  Args:
+    sink: The Ops.SINK rooting the Kernel graph.
+    opts: The Renderer (can change how things are processed, fix this).
+
+  Returns:
+    Linear program in UOps.
+  """
+
   lst = list(full_rewrite_to_sink(sink, opts, linearizer=True).arg.lst)
   if __debug__: type_verify(lst)
   return lst

tinygrad/codegen/devectorizer.py

@@ -290,7 +290,7 @@ devectorize = PatternMatcher([
   # no ALU on vectorized dtypes
   (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.ASSIGN), name="alu"), no_vectorized_alu),
   (UPat(Ops.WMMA, name="wmma"), no_vectorized_wmma),
-  (UPat(Ops.DEFINE_ACC, name="acc"), no_vectorized_acc),
+  (UPat(Ops.DEFINE_REG, name="acc"), no_vectorized_acc),
 ])
 
 pm_render = PatternMatcher([
@@ -329,7 +329,7 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp):
   assert all(x.dtype == red.dtype for x in lst), f"horizontal reduction mismatch {lst[0].dtype} != {red.dtype}"
   # if we have a range
   if len(reduce_range) != 0:
-    acc = UOp(Ops.DEFINE_ACC, red.dtype, (red.const_like(identity_element(red.arg, red.dtype.scalar())),) + tuple(reduce_range), (ctx.acc_num,))
+    acc = UOp(Ops.DEFINE_REG, red.dtype, (red.const_like(identity_element(red.arg, red.dtype.scalar())),) + tuple(reduce_range), (ctx.acc_num,))
     lst = [acc] + lst  # put acc as the first element
     ctx.acc_num += 1
   ret = functools.reduce(lambda x,y: x.alu(red.arg, y), lst)

tinygrad/codegen/expander.py

@@ -87,7 +87,7 @@ expander = PatternMatcher([
          Ops.VECTORIZE, Ops.IF, Ops.REDUCE), name="root", custom_early_reject=set([Ops.UNROLL])), do_expand),
   (UPat(Ops.CONTRACT, name="con"), do_contract),
   # vectorize DEFINE_ACC
-  (UPat(Ops.VECTORIZE, src=UPat(Ops.DEFINE_ACC, name="acc"), name="v"),
+  (UPat(Ops.VECTORIZE, src=UPat(Ops.DEFINE_REG, name="acc"), name="v"),
     lambda acc,v: acc.replace(dtype=v.dtype, src=(acc.src[0].broadcast(v.dtype.count),)+acc.src[1:])),
   # BARRIERs aren't actually expanded
   (UPat(Ops.BARRIER, src=(UPat(Ops.UNROLL, name="ex"),)),

tinygrad/codegen/gpudims.py

@@ -0,0 +1,73 @@
+import math
+from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify
+from tinygrad.helpers import all_int
+from tinygrad.dtype import dtypes
+from tinygrad.shape.view import get_contraction
+from tinygrad.renderer import Renderer
+
+def _group_dims(dims:tuple[sint, ...], max_sizes:tuple[int, ...]):
+  # TODO: symbolic shape
+  if not all_int(dims): return dims
+  while len(dims) > len(max_sizes) or any(d > m for d,m in zip(dims, max_sizes)):
+    for i,m in enumerate(max_sizes):
+      if i < (len(dims)-1) and dims[i] * dims[i+1] <= m:
+        dims = dims[:i] + (dims[i]*dims[i+1],) + dims[i+2:]
+        break
+    else: return None
+  return dims
+
+def _split_dims(dims, max_sizes):
+  if all(d <= m for d,m in zip(dims, max_sizes)): return dims
+  _dims = list(dims) + [1]*(3-len(dims))
+  for i in range(len(_dims)):
+    while _dims[i] > max_sizes[i]:
+      div = next((d for d in range(2, math.ceil(math.sqrt(_dims[i])) + 1) if (_dims[i] % d) == 0), 1)
+      if div == 1: raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
+      _dims[i], _dims[(i+1)%len(_dims)] = _dims[i]//div, _dims[(i+1)%len(_dims)]*div
+  return tuple(_dims[:2] if _dims[2] == 1 else _dims[0] if _dims[1:3] == [1,1] else _dims)
+
+def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|None, reverse=False) -> list[UOp]:
+  if reverse: dims = dims[::-1]
+  # try to group first: (a, b, c, d) -> (ab, c, d)
+  limited = (grouped if (grouped := _group_dims(dims, max_sizes)) else dims) if max_sizes is not None else dims
+  # check if grouping failed
+  if max_sizes is not None and len(limited) > len(max_sizes): raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
+  # try to split up dims: (a,) -> (b, c)
+  if limited == dims: limited = _split_dims(dims, max_sizes) if max_sizes is not None else dims
+  ret = raw_idxs = [UOp(Ops.SPECIAL, dtypes.int, (), (f"{prefix}{i}", s)) for i,s in enumerate(limited)]
+  if len(limited) < len(dims):
+    ret = []
+    if (contraction:=get_contraction(dims, limited)) is None: raise AssertionError(f"get_contraction should not be None {dims=} {limited=}")
+    for idx, contraction_group in zip(raw_idxs, contraction):
+      for c in contraction_group[:-1]:
+        ret.append(idx % dims[c])
+        idx //= dims[c]
+      ret.append(idx)
+  elif len(limited) > len(dims):
+    a, b = len(limited), len(dims)
+    if a == 2 and b == 1: ret = [raw_idxs[0] * limited[1] + raw_idxs[1]]
+    if a == 3 and b == 1: ret = [raw_idxs[0] * (limited[1] * limited[2]) + raw_idxs[1] * limited[2] + raw_idxs[2]]
+    if a == 3 and b == 2: ret = [raw_idxs[0] * limited[1] + raw_idxs[1], raw_idxs[2]]
+  return ret[::-1] if reverse else ret
+
+def add_gpudims(ctx:Renderer, s:UOp):
+  if s.arg is None: return None
+  ki: KernelInfo = s.arg
+  if ki.global_dims == 0 and ki.local_dims == 0: return None
+  s_topo = list(s.toposort())
+  if any(x.op is Ops.SPECIAL for x in s_topo): return None
+  ranges = sorted([x for x in s_topo if x.op is Ops.RANGE and x.arg < (ki.global_dims+ki.local_dims)], key=lambda x: x.arg)
+  if not len(ranges): return None
+  global_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg < ki.global_dims])
+  local_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg >= ki.global_dims])
+  if ki.dont_use_locals:
+    assert ki.local_dims == 0, "can't use locals if there's no local dims"
+    idxs = get_grouped_dims("idx", global_shape, ctx.global_max, reverse=True)
+  else:
+    # define indexes for GPU-like execution
+    idxs = get_grouped_dims("gidx", global_shape, ctx.global_max, reverse=True) + get_grouped_dims("lidx", local_shape, ctx.local_max)
+  return s.substitute(dict(zip(ranges, idxs)))
+
+pm_add_gpudims = PatternMatcher([
+  (UPat(Ops.SINK, name="s"), add_gpudims),
+])

tinygrad/codegen/linearize.py

@@ -99,7 +99,7 @@ class BlockContext:
           ctx.child_ctxs[u] = tuple([y for y in store_context if y not in idx_context and y.op is Ops.RANGE])
         else: ctx.child_ctxs[u] = ()
       elif u.op is Ops.ASSIGN:
-        assert u.src[0].op is Ops.DEFINE_ACC
+        assert u.src[0].op is Ops.DEFINE_REG
         ctx.child_ctxs[u] = tuple([y for y in ctx.last_ctx(u.src[1]) if y not in u.src[0].src[1:]])
     return ctx
 
@@ -216,7 +216,7 @@ def remove_blockend(x:UOp):
     parent_block = parent_blocks[0]
     assert len(parent_blocks) == parent_block.arg.cnt
     # range needs DEFINE_ACC to be before the range (never in DEFINE_ACC for if)
-    early_ops, late_ops = partition(x.arg.lst, lambda y: y.op is Ops.DEFINE_ACC and x.arg.end in y.src)
+    early_ops, late_ops = partition(x.arg.lst, lambda y: y.op is Ops.DEFINE_REG and x.arg.end in y.src)
     # NOTE: we have to add a barrier at the start if barrier is used in the range
     if x.op is Ops.BLOCKEND and any(y.op is Ops.BARRIER for y in late_ops) and late_ops[-1].op is Ops.ENDRANGE:
       late_ops = [UOp(Ops.BARRIER)] + late_ops

tinygrad/codegen/lowerer.py

@@ -1,116 +1,25 @@
 # the job of the lowerer is to do indexing
-import itertools, operator, math
 from dataclasses import dataclass
 from typing import cast
-from tinygrad.dtype import dtypes, PtrDType, least_upper_dtype
-from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint, sint_to_uop
-from tinygrad.renderer import Renderer
-from tinygrad.helpers import all_int, prod, partition, flatten, unwrap
-from tinygrad.uop.symbolic import symbolic
-
-# returns the axes to create new_shape if new_shape can be created by combining axis from old_shape
-def get_contraction(old_shape:tuple[sint, ...], new_shape:tuple[sint, ...]) -> list[list[int]]|None:
-  acc_old, acc_new = list(itertools.accumulate(old_shape, operator.mul)), list(itertools.accumulate(new_shape, operator.mul))
-  try: split = [acc_old.index(acc)+1 if acc != 1 else 0 for acc in acc_new]
-  except ValueError: return None
-  return [list(range(st,ed)) for st,ed in zip([0]+split[:-1], split[:-1]+[len(old_shape)])]
-
-def get_contraction_with_reduce(old_shape:tuple[sint, ...], new_shape:tuple[sint, ...], reduce_axis:tuple[int, ...]) -> list[list[int]]|None:
-  if (contraction:=get_contraction(old_shape, new_shape)) is None: return None
-  # contraction returns the 1s as right justified as possible
-  # normally this contraction is good, but sometimes the reduce dim is empty. borrow from the next one, leaving one
-  # this ensures there's always ones available in the reduce dimension. this is also a valid contraction
-  for i in range(len(contraction)):
-    if i in reduce_axis and len(contraction[i]) == 0:
-      take_from = i+1
-      while take_from < len(contraction) and len(contraction[take_from]) == 0:
-        assert new_shape[take_from] == 1
-        take_from += 1
-      if take_from == len(contraction) or new_shape[take_from] != 1: return None # nothing to take
-      for j in range(take_from, i, -1):
-        assert len(contraction[j]) > 0
-        contraction[j-1] = contraction[j][:-1]
-        contraction[j] = contraction[j][-1:]
-  return contraction
+from tinygrad.dtype import dtypes, PtrDType
+from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint_to_uop
+from tinygrad.helpers import prod, partition, flatten
 
 # ***** indexing *****
-def _group_dims(dims:tuple[sint, ...], max_sizes:tuple[int, ...]):
-  # TODO: symbolic shape
-  if not all_int(dims): return dims
-  while len(dims) > len(max_sizes) or any(d > m for d,m in zip(dims, max_sizes)):
-    for i,m in enumerate(max_sizes):
-      if i < (len(dims)-1) and dims[i] * dims[i+1] <= m:
-        dims = dims[:i] + (dims[i]*dims[i+1],) + dims[i+2:]
-        break
-    else: return None
-  return dims
-
-def _split_dims(dims, max_sizes):
-  if all(d <= m for d,m in zip(dims, max_sizes)): return dims
-  _dims = list(dims) + [1]*(3-len(dims))
-  for i in range(len(_dims)):
-    while _dims[i] > max_sizes[i]:
-      div = next((d for d in range(2, math.ceil(math.sqrt(_dims[i])) + 1) if (_dims[i] % d) == 0), 1)
-      if div == 1: raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
-      _dims[i], _dims[(i+1)%len(_dims)] = _dims[i]//div, _dims[(i+1)%len(_dims)]*div
-  return tuple(_dims[:2] if _dims[2] == 1 else _dims[0] if _dims[1:3] == [1,1] else _dims)
-
-def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|None, reverse=False) -> list[UOp]:
-  if reverse: dims = dims[::-1]
-  # try to group first: (a, b, c, d) -> (ab, c, d)
-  limited = (grouped if (grouped := _group_dims(dims, max_sizes)) else dims) if max_sizes is not None else dims
-  # check if grouping failed
-  if max_sizes is not None and len(limited) > len(max_sizes): raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
-  # try to split up dims: (a,) -> (b, c)
-  if limited == dims: limited = _split_dims(dims, max_sizes) if max_sizes is not None else dims
-  ret = raw_idxs = [UOp(Ops.SPECIAL, dtypes.int, (), (f"{prefix}{i}", s)) for i,s in enumerate(limited)]
-  if len(limited) < len(dims):
-    ret = []
-    if (contraction:=get_contraction(dims, limited)) is None: raise AssertionError(f"get_contraction should not be None {dims=} {limited=}")
-    for idx, contraction_group in zip(raw_idxs, contraction):
-      for c in contraction_group[:-1]:
-        ret.append(idx % dims[c])
-        idx //= dims[c]
-      ret.append(idx)
-  elif len(limited) > len(dims):
-    a, b = len(limited), len(dims)
-    if a == 2 and b == 1: ret = [raw_idxs[0] * limited[1] + raw_idxs[1]]
-    if a == 3 and b == 1: ret = [raw_idxs[0] * (limited[1] * limited[2]) + raw_idxs[1] * limited[2] + raw_idxs[2]]
-    if a == 3 and b == 2: ret = [raw_idxs[0] * limited[1] + raw_idxs[1], raw_idxs[2]]
-  return ret[::-1] if reverse else ret
 
 @dataclass
 class IndexContext:
   idxs: list[UOp]
   ridxs: list[UOp]
 
-def get_index(ast:UOp, opts:Renderer) -> IndexContext:
+def get_index(ast:UOp) -> IndexContext:
   ki = ast.arg if isinstance(ast.arg, KernelInfo) else KernelInfo()
   # NOTE: assumes the shape is <global dims> <local dims> <group_for_reduces> <reduces> <upcasts/unrolls>
   full_shape = ast.full_shape
   first_upcasted = len(full_shape)-ki.upcasted
-  # if there's no reduce, this is first_upcasted. assumes reduces are at the end
-  first_reduce = min([first_upcasted]+flatten(x.axis_arg for x in ast.toposort() if x.op is Ops.REDUCE_AXIS))
-  local_loads = [x for x in ast.toposort() if x.op is Ops.LOAD and x.src[0].base.op is Ops.DEFINE_LOCAL]
-  # NOTE: sum up the reduced axes looking across all local loads, yields the number of grouped reduces
-  group_for_reduces = sum([any(l.st_arg.shape[i]!=ast.src[0].st_arg.shape[i] for l in local_loads) for i in range(first_reduce,first_upcasted)])
-  global_dims = first_reduce-ki.local_dims
 
-  if opts.has_local:
-    if ki.dont_use_locals:
-      assert ki.local_dims == 0, "can't use locals if there's no local dims"
-      idxs = get_grouped_dims("idx", full_shape[:global_dims], opts.global_max, reverse=True)
-    else:
-      # define indexes for GPU-like execution
-      idxs = get_grouped_dims("gidx", full_shape[:global_dims], opts.global_max, reverse=True) + \
-             get_grouped_dims("lidx", full_shape[global_dims:first_reduce+group_for_reduces], opts.local_max)
-  else:
-    # all loops are RANGES
-    idxs = [UOp(Ops.RANGE, dtypes.int, (sint_to_uop(g),), i) for i,g in enumerate(full_shape[:first_reduce])]
-
-  # reduce loops
-  idxs += [UOp(Ops.RANGE, dtypes.int, (sint_to_uop(g),), i)
-    for i,g in enumerate(full_shape[first_reduce+group_for_reduces:first_upcasted], start=first_reduce+group_for_reduces)]
+  # all loops are RANGES
+  idxs = [UOp(Ops.RANGE, dtypes.int, (sint_to_uop(g),), i) for i,g in enumerate(full_shape[:first_upcasted])]
 
   # upcast loops
   for i,g in enumerate(full_shape[first_upcasted:], start=first_upcasted):
@@ -118,6 +27,11 @@ def get_index(ast:UOp, opts:Renderer) -> IndexContext:
     idxs.append(UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(g), tuple(range(g))),), ((i,g),)))
 
   # late indexes (group for reduce)
+  # if there's no reduce, this is first_upcasted. assumes reduces are at the end
+  first_reduce = min([first_upcasted]+flatten(x.axis_arg for x in ast.toposort() if x.op is Ops.REDUCE_AXIS))
+  local_loads = [x for x in ast.toposort() if x.op is Ops.LOAD and x.src[0].base.op is Ops.DEFINE_LOCAL]
+  # NOTE: sum up the reduced axes looking across all local loads, yields the number of grouped reduces
+  group_for_reduces = sum([any(l.st_arg.shape[i]!=ast.src[0].st_arg.shape[i] for l in local_loads) for i in range(first_reduce,first_upcasted)])
   ridxs = idxs[:]
   for a in range(first_reduce, first_reduce+group_for_reduces):
     ridxs[a] = UOp(Ops.RANGE, dtypes.int, (sint_to_uop(full_shape[a]),), 1000+a)
@@ -130,102 +44,41 @@ def lower_reduce_axis(ctx: IndexContext, x: UOp):
   # NOTE: always using ridxs is fine here
   reduce_range, reduce_expand = partition([ctx.ridxs[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}"
-  alu_op: Ops = x.arg[0]
   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))
   # 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), alu_op)
+  return UOp(Ops.REDUCE, x.dtype, (ret,)+tuple(reduce_range), x.arg[0])
 
-def lower_load_store(ctx: IndexContext, x: UOp, buf: UOp):
-  idx, valid = x.st_arg.to_indexed_uops(ctx.ridxs if x.op is Ops.LOAD and buf.op is Ops.DEFINE_LOCAL else ctx.idxs)
-  if x.op is Ops.LOAD:
-    barrier = (UOp(Ops.BARRIER, dtypes.void, (x.src[1],)),) if buf.op is Ops.DEFINE_LOCAL else ()
-    return UOp(Ops.LOAD, x.dtype, (buf.index(idx, valid),) + barrier)
-  # NOTE: only store the local reduceop in the threads that are actually doing the reduce
-  if cast(PtrDType, buf.dtype).local and x.src[1].op is Ops.REDUCE:
-    reduce_input = x.src[1].src[0]
-    store_back = reduce_input.op is Ops.LOAD and cast(PtrDType, reduce_input.src[0].dtype).local
-  else: store_back = False
-  # NOTE: If we're storing the reduced value back into each thread, need to zero-out the reduced axes
-  if store_back: idx, _ = x.st_arg.to_indexed_uops([u.const_like(0) if u in x.src[1].src else u for u in ctx.idxs])
-  if (not cast(PtrDType, buf.dtype).local) or store_back:
+def lower_load(ctx: IndexContext, x: UOp, buf: UOp):
+  idx, valid = x.st_arg.to_indexed_uops(ctx.ridxs if buf.op is Ops.DEFINE_LOCAL else ctx.idxs)
+  barrier = (UOp(Ops.BARRIER, dtypes.void, (x.src[1],)),) if buf.op is Ops.DEFINE_LOCAL else ()
+  return UOp(Ops.LOAD, x.dtype, (buf.index(idx, valid),) + barrier)
+
+def lower_store(ctx: IndexContext, x: UOp, buf: UOp):
+  idx, valid = x.st_arg.to_indexed_uops(ctx.idxs)
+  if not cast(PtrDType, buf.dtype).local:
+    # NOTE: only store the local reduceop in the threads that are actually doing the reduce
     for oidx, ridx in zip(ctx.idxs, ctx.ridxs):
       if oidx is not ridx: valid = valid * oidx.eq(0)
   return UOp(Ops.STORE, dtypes.void, (buf.index(idx, valid), x.src[1]))
 
-def lower_const(x:UOp):
-  assert all(v.mask is None for v in unwrap(x.st).views), f"VIEW in CONST/DEFINE_VAR source must be unmasked, got {x.st}"
-  return x.replace(src=())
+def lower_const(ctx:IndexContext, view:UOp, c:UOp):
+  if all(x.mask is None for x in view.arg.views): return c
+  _, valid = view.arg.to_indexed_uops(ctx.idxs)
+  return valid.where(c, c.const_like(0))
 
 pm_lowerer = PatternMatcher([
+  # TODO: remove these hacks
+  # hack for old style CONST(VIEW) (now it's just VIEW(CONST))
+  (UPat((Ops.DEFINE_VAR, Ops.CONST), src=(UPat(Ops.VIEW, name="v"),), name="c"), lambda c,v: c.replace(src=()).view(v.arg)),
+  # hack for old style VALID (now it's just VIEW(CONST))
+  (UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c"), UPat(Ops.CONST, arg=0)), lambda c,v: c.replace(src=()).view(v.arg)),
+
+  # reduce/view_const
   (UPat(Ops.REDUCE_AXIS, name="x"), lower_reduce_axis),
-  (UPat((Ops.CONST, Ops.DEFINE_VAR), src=(UPat(Ops.VIEW),), name="x"), lower_const),
-  (UPat(Ops.VALID, src=(UPat(Ops.VIEW),), name="x"), lambda ctx,x: x.st_arg.to_indexed_uops(ctx.idxs)[1]),
+  (UPat(Ops.VIEW, src=(UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"),), name="view"), lower_const),
   # rewrite LOAD/STORE VIEW to LOAD/STORE with indexed
-  (UPat((Ops.LOAD, Ops.STORE), src=(UPat.var("buf").view(),), allow_any_len=True, name="x"), lower_load_store),
-  (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx"), UPat.const(dtypes.bool, True))), lambda b, idx: b.index(idx)),
-])
-
-# **** this is the "quantization preprocessor", it makes ONNX quantized models, and probably also others, actually use ints ****
-
-FP = (1 << 15)
-pm_quant = symbolic+PatternMatcher([
-  # cast after add/mul
-  (UPat.var("x").cast(dtypes.float32) + UPat.var("y").cast(dtypes.float32),
-   lambda x,y: (x.cast(least_upper_dtype(x.dtype, y.dtype))+y.cast(least_upper_dtype(x.dtype, y.dtype))).cast(dtypes.float32)),
-  (UPat.var("x").cast(dtypes.float32) * UPat.var("y").cast(dtypes.float32),
-   lambda x,y: (x.cast(least_upper_dtype(x.dtype, y.dtype))*y.cast(least_upper_dtype(x.dtype, y.dtype))).cast(dtypes.float32)),
-
-  # masked MUL after masked ADD
-  ((UPat.var("x") + UPat.var("v").where(UPat.var('cadd'), UPat(Ops.CONST, arg=0))) * UPat.var("v").where(UPat.var('cmul'), UPat(Ops.CONST, arg=0)),
-   lambda x,v,cadd,cmul: x*v.where(cmul, 0)+v.where(cadd*cmul, 0)),
-
-  # MUL after reduce
-  (UPat(Ops.REDUCE_AXIS, src=(UPat.var("x") * UPat.cvar("c"),), name="r"), lambda x,c,r: r.replace(src=(x,))*c.arg),
-  # CAST after reduce (doesn't work if it's a size change)
-  (UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.CAST, src=(UPat.var("x"),)),), name="r"),
-    lambda x,r: r.replace(dtype=x.dtype, src=(x,)).cast(r.dtype) if dtypes.is_float(r.dtype) else None),
-
-  # x*c1 + y*c2 -> (x+y)*c1 (if c1 and c2 are close floats)
-  (UPat.var("x")*UPat.cvar("c1", dtype=dtypes.floats) + UPat.var("y")*UPat.cvar("c2", dtype=dtypes.floats),
-   lambda x,y,c1,c2: (x+y)*c1 if abs(c1.arg-c2.arg) < 1e-9 else None),
-  # mul 0 * c1 is 0
-  (UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c1"), UPat(Ops.CONST, arg=0)) *
-   UPat(Ops.LOAD, src=(UPat().view(name="v"),)).cast(dtypes.int).cast(dtypes.float).named("ld"), lambda ld,v,c1: ld*c1),
-  # mul (with plus) 0 * c1 is 0
-  (UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c1"), UPat(Ops.CONST, arg=0)) *
-   (UPat(Ops.LOAD, src=(UPat().view(name="v"),)).cast(dtypes.int) + \
-    UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar(), UPat(Ops.CONST, arg=0))).cast(dtypes.float).named("ld"),
-      lambda ld,v,c1: ld*c1),
-
-  # const push through add
-  ((UPat.var("x")*UPat.cvar("c1") + UPat.var("y")*UPat.cvar("c2")) * UPat.cvar("c3"), lambda x,y,c1,c2,c3: (x*c1*c3) + (y*c2*c3)),
-
-  # fixed point mult, replace (x.float()*c1+c2).int() with an int expression
-  ((UPat.var("x").cast(dtypes.float)*UPat.var("c1")+UPat.var("cc")).cast(dtypes.int),
-   lambda x,c1,cc: ((x*(c1*FP).cast(x.dtype) + (cc*FP).cast(x.dtype)) // FP).cast(dtypes.int)),
-  # fixed point mult, replace (x.float()*c1 + y.float()*c2)*cc.int() with an int expression
-  ((UPat.var("x").cast(dtypes.float)*UPat.var("c1")+UPat.var("y").cast(dtypes.float)*UPat.var("c2")+UPat.var("cc")).cast(dtypes.int),
-   lambda x,c1,y,c2,cc: ((x*(c1*FP).cast(x.dtype) + y.cast(x.dtype)*(c2*FP).cast(x.dtype) + (cc*FP).cast(x.dtype)) // FP).cast(dtypes.int)),
-
-  # where move
-  (UPat.var("valid").where(UPat.var("yes"), UPat(Ops.CONST, arg=0))*UPat.var("mul"), lambda valid, yes, mul:
-    (yes*mul*valid.where(UOp.const(mul.dtype, 1), UOp.const(mul.dtype, 0))) if yes.op is not Ops.CONST or yes.arg != 1 else None),
-  ((UPat.var("x")*UPat.cvar("c"))*(UPat.var().where(UPat(Ops.CONST, arg=1), UPat(Ops.CONST, arg=0)).named("v")), lambda x,c,v: (x*v)*c),
-  (UPat.var("x").cast().named('c') * UPat.var('valid').where(UPat(Ops.CONST, arg=1), UPat(Ops.CONST, arg=0)), lambda x,c,valid:
-    (x*valid.where(UOp.const(x.dtype, 1), UOp.const(x.dtype, 0))).cast(c.dtype)),
-  ((UPat.var('x') * UPat.var('v1').where(UPat(Ops.CONST, arg=1), UPat(Ops.CONST, arg=0)) *
-    UPat.var('v2').where(UPat(Ops.CONST, arg=1), UPat(Ops.CONST, arg=0))).named("mul"), lambda x, mul, v1, v2:
-    x * (v1&v2).where(UOp.const(mul.dtype, 1), UOp.const(mul.dtype, 0))),
-
-  # where on two adds
-  (UPat.var("x") + UPat.var("v").where(UPat.var("a0"), UPat.var("a1")) + UPat.var("v").where(UPat.var("b0"), UPat.var("b1")),
-    lambda x,v,a0,a1,b0,b1: x + v.where(a0+b0, a1+b1)),
-
-  # split REDUCE into multiple reduces (who remembers FOIL?)
-  (UPat(Ops.REDUCE_AXIS, src=((UPat(Ops.CAST, name="v1")+UPat.var("c1")) * UPat(Ops.CAST, name="v2"),), name="r"),
-    lambda v1,v2,c1,r: r.replace(src=(v1*v2,)) + r.replace(src=(c1*v2,))),
-  (UPat(Ops.REDUCE_AXIS, src=((UPat(Ops.CAST, name="v1")+UPat.var("c1")) * (UPat(Ops.CAST, name="v2",)+UPat.var("c2")),), name="r"),
-    lambda v1,v2,c1,c2,r: r.replace(src=(v1*v2,)) + r.replace(src=(c2*v1,)) + r.replace(src=(c1*v2,)) + r.replace(src=(c1*c2,))),
+  (UPat(Ops.LOAD, src=(UPat.var("buf").view(),), allow_any_len=True, name="x"), lower_load),
+  (UPat(Ops.STORE, src=(UPat.var("buf").view(),), allow_any_len=True, name="x"), lower_store),
 ])

tinygrad/codegen/quantize.py

@@ -0,0 +1,67 @@
+from tinygrad.dtype import dtypes, least_upper_dtype
+from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
+from tinygrad.uop.symbolic import symbolic
+
+# **** this is the "quantization preprocessor", it makes ONNX quantized models, and probably also others, actually use ints ****
+# this is badly tested and low quality. remove it?
+
+FP = (1 << 15)
+pm_quant = symbolic+PatternMatcher([
+  # cast after add/mul
+  (UPat.var("x").cast(dtypes.float32) + UPat.var("y").cast(dtypes.float32),
+   lambda x,y: (x.cast(least_upper_dtype(x.dtype, y.dtype))+y.cast(least_upper_dtype(x.dtype, y.dtype))).cast(dtypes.float32)),
+  (UPat.var("x").cast(dtypes.float32) * UPat.var("y").cast(dtypes.float32),
+   lambda x,y: (x.cast(least_upper_dtype(x.dtype, y.dtype))*y.cast(least_upper_dtype(x.dtype, y.dtype))).cast(dtypes.float32)),
+
+  # masked MUL after masked ADD
+  ((UPat.var("x") + UPat.var("v").where(UPat.var('cadd'), UPat(Ops.CONST, arg=0))) * UPat.var("v").where(UPat.var('cmul'), UPat(Ops.CONST, arg=0)),
+   lambda x,v,cadd,cmul: x*v.where(cmul, 0)+v.where(cadd*cmul, 0)),
+
+  # MUL after reduce
+  (UPat(Ops.REDUCE_AXIS, src=(UPat.var("x") * UPat.cvar("c"),), name="r"), lambda x,c,r: r.replace(src=(x,))*c.arg),
+  # CAST after reduce (doesn't work if it's a size change)
+  (UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.CAST, src=(UPat.var("x"),)),), name="r"),
+    lambda x,r: r.replace(dtype=x.dtype, src=(x,)).cast(r.dtype) if dtypes.is_float(r.dtype) else None),
+
+  # x*c1 + y*c2 -> (x+y)*c1 (if c1 and c2 are close floats)
+  (UPat.var("x")*UPat.cvar("c1", dtype=dtypes.floats) + UPat.var("y")*UPat.cvar("c2", dtype=dtypes.floats),
+   lambda x,y,c1,c2: (x+y)*c1 if abs(c1.arg-c2.arg) < 1e-9 else None),
+  # mul 0 * c1 is 0
+  (UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c1"), UPat(Ops.CONST, arg=0)) *
+   UPat(Ops.LOAD, src=(UPat().view(name="v"),)).cast(dtypes.int).cast(dtypes.float).named("ld"), lambda ld,v,c1: ld*c1),
+  # mul (with plus) 0 * c1 is 0
+  (UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c1"), UPat(Ops.CONST, arg=0)) *
+   (UPat(Ops.LOAD, src=(UPat().view(name="v"),)).cast(dtypes.int) + \
+    UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar(), UPat(Ops.CONST, arg=0))).cast(dtypes.float).named("ld"),
+      lambda ld,v,c1: ld*c1),
+
+  # const push through add
+  ((UPat.var("x")*UPat.cvar("c1") + UPat.var("y")*UPat.cvar("c2")) * UPat.cvar("c3"), lambda x,y,c1,c2,c3: (x*c1*c3) + (y*c2*c3)),
+
+  # fixed point mult, replace (x.float()*c1+c2).int() with an int expression
+  ((UPat.var("x").cast(dtypes.float)*UPat.var("c1")+UPat.var("cc")).cast(dtypes.int),
+   lambda x,c1,cc: ((x*(c1*FP).cast(x.dtype) + (cc*FP).cast(x.dtype)) // FP).cast(dtypes.int)),
+  # fixed point mult, replace (x.float()*c1 + y.float()*c2)*cc.int() with an int expression
+  ((UPat.var("x").cast(dtypes.float)*UPat.var("c1")+UPat.var("y").cast(dtypes.float)*UPat.var("c2")+UPat.var("cc")).cast(dtypes.int),
+   lambda x,c1,y,c2,cc: ((x*(c1*FP).cast(x.dtype) + y.cast(x.dtype)*(c2*FP).cast(x.dtype) + (cc*FP).cast(x.dtype)) // FP).cast(dtypes.int)),
+
+  # where move
+  (UPat.var("valid").where(UPat.var("yes"), UPat(Ops.CONST, arg=0))*UPat.var("mul"), lambda valid, yes, mul:
+    (yes*mul*valid.where(UOp.const(mul.dtype, 1), UOp.const(mul.dtype, 0))) if yes.op is not Ops.CONST or yes.arg != 1 else None),
+  ((UPat.var("x")*UPat.cvar("c"))*(UPat.var().where(UPat(Ops.CONST, arg=1), UPat(Ops.CONST, arg=0)).named("v")), lambda x,c,v: (x*v)*c),
+  (UPat.var("x").cast().named('c') * UPat.var('valid').where(UPat(Ops.CONST, arg=1), UPat(Ops.CONST, arg=0)), lambda x,c,valid:
+    (x*valid.where(UOp.const(x.dtype, 1), UOp.const(x.dtype, 0))).cast(c.dtype)),
+  ((UPat.var('x') * UPat.var('v1').where(UPat(Ops.CONST, arg=1), UPat(Ops.CONST, arg=0)) *
+    UPat.var('v2').where(UPat(Ops.CONST, arg=1), UPat(Ops.CONST, arg=0))).named("mul"), lambda x, mul, v1, v2:
+    x * (v1&v2).where(UOp.const(mul.dtype, 1), UOp.const(mul.dtype, 0))),
+
+  # where on two adds
+  (UPat.var("x") + UPat.var("v").where(UPat.var("a0"), UPat.var("a1")) + UPat.var("v").where(UPat.var("b0"), UPat.var("b1")),
+    lambda x,v,a0,a1,b0,b1: x + v.where(a0+b0, a1+b1)),
+
+  # split REDUCE into multiple reduces (who remembers FOIL?)
+  (UPat(Ops.REDUCE_AXIS, src=((UPat(Ops.CAST, name="v1")+UPat.var("c1")) * UPat(Ops.CAST, name="v2"),), name="r"),
+    lambda v1,v2,c1,r: r.replace(src=(v1*v2,)) + r.replace(src=(c1*v2,))),
+  (UPat(Ops.REDUCE_AXIS, src=((UPat(Ops.CAST, name="v1")+UPat.var("c1")) * (UPat(Ops.CAST, name="v2",)+UPat.var("c2")),), name="r"),
+    lambda v1,v2,c1,c2,r: r.replace(src=(v1*v2,)) + r.replace(src=(c2*v1,)) + r.replace(src=(c1*v2,)) + r.replace(src=(c1*c2,))),
+])

tinygrad/device.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from dataclasses import dataclass, replace
+from dataclasses import dataclass, replace, field
 from collections import defaultdict
 from typing import Optional, Any, Generic, TypeVar, Iterator, Generator
 import importlib, inspect, functools, pathlib, os, ctypes, ctypes.util, platform, contextlib, sys, re, atexit, pickle, decimal, time
@@ -23,7 +23,7 @@ class _Device:
   @functools.cache  # this class is a singleton, pylint: disable=method-cache-max-size-none
   def __get_canonicalized_item(self, ix:str) -> Compiled:
     assert ALLOW_DEVICE_USAGE or ix.split(":")[0] in ["DISK", "NPY", "PYTHON"], f"usage of device {ix} disallowed"
-    base = __name__.split('.')[0]  # tinygrad
+    base = (__package__ or __name__).split('.')[0]  # tinygrad
     x = ix.split(":")[0].lower()
     ret = [cls for cname, cls in inspect.getmembers(importlib.import_module(f'{base}.runtime.ops_{x}')) \
            if (cname.lower() == x + "device")][0](ix)
@@ -59,6 +59,9 @@ class ProfileDeviceEvent(ProfileEvent):
 @dataclass(frozen=True)
 class ProfileRangeEvent(ProfileEvent): device:str; name:str; st:decimal.Decimal; en:decimal.Decimal; is_copy:bool # noqa: E702
 
+@dataclass(frozen=True)
+class ProfilePointEvent(ProfileEvent): device:str; name:str; st:decimal.Decimal; ref:int; arg:dict=field(default_factory=dict) # noqa: E702
+
 @dataclass(frozen=True)
 class ProfileProgramEvent(ProfileEvent): device:str; name:str; lib:bytes|None; base:int|None # noqa: E702
 
@@ -105,6 +108,7 @@ class MultiBuffer:
   def __repr__(self): return f"<multibuf real:{self.is_allocated()} device:{tuple(x.device for x in self.bufs)} size:{self.size} dtype:{self.dtype}>"
 
 class Buffer:
+  profile_events:list[ProfileEvent] = []
   def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:Optional[BufferSpec]=None, initial_value:Optional[bytes]=None,
                uop_refcount=0, base:Optional[Buffer]=None, offset:int=0, preallocate=False):
     if isinstance(dtype, ImageDType): options = BufferSpec(image=dtype) # TODO: image hack shouldn't be here. where should it be?
@@ -130,10 +134,13 @@ class Buffer:
   def ref(self, cnt):
     self.base._uop_refcount += cnt
     return self
-  def is_allocated(self) -> bool: return hasattr(self, '_buf')
-  def ensure_allocated(self) -> Buffer: return self.allocate() if not self.is_allocated() else self
+  # check if the underlying buffer is allocated and the current buffer/view is initialized
+  def is_initialized(self) -> bool: return self.is_allocated() and hasattr(self, '_buf')
+  # check if the underlying buffer is allocated, possibly from the base object
+  def is_allocated(self) -> bool: return self.base.is_allocated() if self._base is not None else hasattr(self, '_buf')
+  def ensure_allocated(self) -> Buffer: return self.allocate() if not self.is_initialized() else self
   def allocate(self, opaque=None, external_ptr=None) -> Buffer:
-    assert not self.is_allocated(), "can't allocate already allocated 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 (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
@@ -147,12 +154,17 @@ class Buffer:
     else:
       self._buf = opaque if opaque is not None else self.allocator.alloc(self.nbytes, self.options)
       if not self.device.startswith("DISK"): GlobalCounters.mem_used += self.nbytes
+      if PROFILE:
+        self._prof_num = num = len(Buffer.profile_events)
+        ts = decimal.Decimal(time.perf_counter_ns())/1000
+        Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", ts, num, {"dtype":str(self.dtype),"sz":self.size,"nbytes":self.nbytes}))
     return self
   def deallocate(self):
-    assert self.is_allocated(), "buffer must be allocated to deallocate"
+    assert hasattr(self, '_buf'), "buffer must be allocated to deallocate"
     if DEBUG is not None and DEBUG >= 7: print(f"buffer: deallocate {self.nbytes} bytes on {self.device}")
     if self._base is None and (self.options is None or self.options.external_ptr is None):
       if GlobalCounters is not None and not self.device.startswith("DISK"): GlobalCounters.mem_used -= self.nbytes
+      if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", decimal.Decimal(time.perf_counter_ns())/1000, self._prof_num))
       self.allocator.free(self._buf, self.nbytes, self.options)
     elif self._base is not None: self._base.allocated_views -= 1
     del self._buf
@@ -167,7 +179,7 @@ class Buffer:
     return self.__class__, (self.device, self.size, self.dtype, None, self.options, buf, self.uop_refcount)
   @property
   def nbytes(self): return self.size*self.dtype.itemsize
-  def __del__(self): (not self.is_allocated()) or self.deallocate()
+  def __del__(self): (not hasattr(self, '_buf')) or self.deallocate()
   def __repr__(self):
     return f"<buf real:{self.is_allocated()} device:{self.device} size:{self.size} dtype:{self.dtype}" + \
            (f" offset:{self.offset}" if self._base is not None else "") + (f" {self.options=}" if self.options is not None else "") + ">"
@@ -188,13 +200,13 @@ class Buffer:
   def copyin(self, mv:memoryview):
     mv = flat_mv(mv)
     assert len(mv) == self.nbytes, f"size mismatch, {len(mv)=} != {self.dtype=} {self.size=}"
-    assert self.is_allocated(), "can't copyin to unallocated buffer"
+    assert self.is_initialized(), "can't copyin to unallocated buffer"
     self.allocator._copyin(self._buf, mv)
     return self
   def copyout(self, mv:memoryview) -> memoryview:
     mv = flat_mv(mv)
     assert len(mv) == self.nbytes, f"size mismatch, {len(mv)=} != {self.dtype=} {self.size=}"
-    assert self.is_allocated(), "can't copyout unallocated buffer"
+    assert self.is_initialized(), "can't copyout unallocated buffer"
     self.allocator._copyout(mv, self._buf)
     return mv
   def view(self, size:int, dtype:DType, offset:int) -> Buffer:
@@ -270,7 +282,6 @@ MAP_JIT = 0x0800
 # CPUProgram is a jit/shellcode program that can be just mmapped and jumped to
 class CPUProgram:
   rt_lib = ctypes.CDLL(ctypes.util.find_library('System' if OSX else 'kernel32') if OSX or sys.platform == "win32" else 'libgcc_s.so.1')
-  atomic_lib = ctypes.CDLL(ctypes.util.find_library('atomic')) if sys.platform == "linux" else None
 
   def __init__(self, name:str, lib:bytes):
     if sys.platform == "win32":
@@ -386,7 +397,7 @@ if PROFILE:
     for dev in devs: dev.synchronize()
     for dev in devs: dev._at_profile_finalize()
 
-    with open(fn:=temp("profile.pkl", append_user=True), "wb") as f: pickle.dump(Compiled.profile_events, f)
+    with open(fn:=temp("profile.pkl", append_user=True), "wb") as f: pickle.dump(Compiled.profile_events+Buffer.profile_events, f)
 
     if not getenv("SQTT", 0):
       from tinygrad.uop.ops import launch_viz

tinygrad/engine/memory.py

@@ -5,7 +5,7 @@ from tinygrad.device import Device, Buffer
 from tinygrad.helpers import NO_MEMORY_PLANNER, dedup, DEBUG, round_up
 from tinygrad.uop.ops import Ops
 from tinygrad.dtype import dtypes, ImageDType
-from tinygrad.runtime.support.allocator import TLSFAllocator
+from tinygrad.runtime.support.memory import TLSFAllocator
 
 # **************** memory planning ****************
 

tinygrad/engine/realize.py

@@ -1,29 +1,51 @@
 from typing import Optional, cast, Generator
 import time, pprint
 from dataclasses import dataclass, replace, field
-from tinygrad.helpers import all_same, colored, getenv, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA
-from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU
-from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, Variable, sym_infer
+from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA
+from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, Variable, sym_infer, graph_rewrite, print_uops, track_rewrites
 from tinygrad.device import Device, Buffer
 from tinygrad.renderer import Renderer, ProgramSpec, Estimates
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.engine.schedule import ScheduleItem
+from tinygrad.opt import get_optimized_ast
+from tinygrad.codegen import full_rewrite
+from tinygrad.uop.spec import type_verify
 
 # **************** Program Creation ****************
 
-logkerns, logkerns_level = open(getenv("LOGKERNS", ""), "a") if getenv("LOGKERNS", "") else None, getenv("LOGKERNS_LEVEL", 1)
-def get_program(renderer:Renderer, ast:UOp) -> ProgramSpec:
-  k = Kernel(ast, opts=renderer)
-  if not NOOPT:
-    if not k.apply_tensor_cores(getenv("TC", 1)): k.apply_opts(hand_coded_optimizations(k))
-    if BEAM >= 1:
-      from tinygrad.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)))
-  if logkerns is not None: logkerns.writelines([f"{(k.ast, k.applied_opts)}\n"])
-  return k.to_program()
+@track_rewrites(name=lambda _ast,_renderer,ret:ret)
+def get_program(ast:UOp, renderer:Renderer) -> ProgramSpec:
+  """
+  Transform an AST into a ProgramSpec. May trigger BEAM search.
+
+  Args:
+    ast: The Ops.SINK rooted AST
+    renderer: The renderer used to generate the code
+
+  Returns:
+    The ProgramSpec of the program.
+  """
+
+  if getenv("VIZ"): graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
+  modified_ast = get_optimized_ast(ast, renderer) if ast.arg is None or ast.arg.opts_to_apply is not None else ast
+  if __debug__: type_verify(list(modified_ast.toposort()))
+
+  # linearize
+  try:
+    uops = full_rewrite(modified_ast, renderer)
+  except RuntimeError:
+    print("***** LINEARIZE FAILURE *****")
+    print(f"ast = {ast}")
+    print(f"opts = {modified_ast.arg.applied_opts}")
+    raise
+  assert uops[-1].op is Ops.SINK, "last uop must be sink"
+
+  # print and render
+  if DEBUG >= 6: print_uops(uops)
+  src = renderer.render(uops)
+
+  return ProgramSpec(uops[-1].arg.name, src, renderer.device, ast, uops,
+                     global_size=[1,1,1] if renderer.has_local else None, local_size=[1,1,1] if renderer.has_local else None)
 
 # **************** Runners ****************
 
@@ -109,7 +131,7 @@ def get_runner(device:str, ast:UOp) -> CompiledRunner:
   if bret:=method_cache.get(bkey):
     method_cache[ckey] = ret = CompiledRunner(replace(bret.p, device=device), bret.lib)
   else:
-    prg: ProgramSpec = get_program(Device[device].renderer, ast)
+    prg: ProgramSpec = get_program(ast, Device[device].renderer)
     method_cache[ckey] = method_cache[bkey] = ret = CompiledRunner(replace(prg, device=device))
   return ret