Merge branch 'master' into call_inline

cl_compiler instead of compiler because overriding Compiled.compiler seems more confusing

.github/workflows/benchmark.yml

@@ -16,6 +16,48 @@ on:
   workflow_dispatch:
 
 jobs:
+  # the goal of this test is to replicate a normal person on a laptop running the test
+  # no process replay, no benchmarks, no CI, just a normal laptop person
+  # the 3 minute timeout should not be raised
+  testmacpytest:
+    name: Mac pytest
+    runs-on: [self-hosted, macOS]
+    timeout-minutes: 3
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    # brew install uv
+    - name: setup python environment
+      run: |
+        rm -rf /tmp/tinygrad_pytest_ci
+        uv venv /tmp/tinygrad_pytest_ci
+        source /tmp/tinygrad_pytest_ci/bin/activate
+        uv pip install .[testing]
+    - name: setup staging db
+      run: |
+        echo "CACHEDB=/tmp/pytest-db-ci.db" >> $GITHUB_ENV
+        rm -f /tmp/pytest-db-ci*
+    # TODO: remove this step once all old caches are migrated
+    - name: Migrate old huggingface cache (symlinks break onnxruntime 1.24+)
+      run: |
+        cd ~/Library/Caches/tinygrad/downloads/models 2>/dev/null || exit 0
+        for old_dir in models--*; do
+          [ -d "$old_dir" ] || continue
+          repo_id=$(echo "$old_dir" | sed 's/models--//; s/--/\//g')
+          snapshot=$(ls -1 "$old_dir/snapshots" 2>/dev/null | head -1)
+          [ -n "$snapshot" ] || continue
+          mkdir -p "$repo_id"
+          cp -RLn "$old_dir/snapshots/$snapshot/"* "$repo_id/" 2>/dev/null || true
+        done
+    - name: Run pytest -nauto
+      run: |
+        source /tmp/tinygrad_pytest_ci/bin/activate
+        pytest -nauto --durations=20
+
   testmacbenchmark:
     name: Mac Benchmark
     env:
@@ -483,6 +525,8 @@ jobs:
       run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
     - name: Run full CIFAR training steps w 6 GPUS
       run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Test full tinyfs load
+      run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 

.github/workflows/test.yml

@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '15'
+  CACHE_VERSION: '16'
   CAPTURE_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
   PYTHONPATH: ${{ github.workspace }}
@@ -29,9 +29,9 @@ jobs:
         deps: testing_unit
         llvm: 'true'
     - name: Speed Test
-      run: CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
+      run: CPU=1 CPU_LLVM=1 THREADS=0 python3 test/speed/external_test_speed_v_torch.py
     - name: Speed Test (BEAM=2)
-      run: BEAM=2 CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
+      run: BEAM=2 CPU=1 CPU_LLVM=1 THREADS=0 python3 test/speed/external_test_speed_v_torch.py
 
   docs:
     name: Docs
@@ -207,11 +207,11 @@ jobs:
       run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
     - name: Test device flop counts
       run: |
-        DEBUG=2 EMULATE=METAL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=AMD PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=CUDA PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=INTEL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 AMX=1 EMULATE=AMX PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStats.test_simple_matmul
+        DEBUG=2 EMULATE=METAL PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=AMD PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=CUDA PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=INTEL PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 AMX=1 EMULATE=AMX PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStats.test_simple_matmul
 
   linter:
     name: Linters
@@ -266,17 +266,19 @@ jobs:
       run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
     - name: Run unit tests
       run: |
-        CPU=1 python test/unit/test_device.py TestRunAsModule.test_module_runs
-        CPU=1 python -m pytest -n=auto test/unit/ --durations=20 --deselect=test/unit/test_device.py::TestRunAsModule::test_module_runs
+        CPU=1 python test/null/test_device.py TestRunAsModule.test_module_runs
+        CPU=1 python -m pytest -n=auto test/unit/ --durations=20
+    - name: Run NULL backend tests
+      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
     - name: Run targetted tests on NULL backend
-      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
+      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
     # TODO: too slow
     # - name: Run SDXL on NULL backend
     #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
     - name: Run Clip tests for SD MLPerf on NULL backend
       run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
     - name: Run AMD emulated BERT training on NULL backend
-      run: EMULATE=AMD_RDNA4 NULL=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+      run: EMULATE=AMD_RDNA4 NULL=1 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
     # TODO: support fake weights
     #- name: Run LLaMA 7B on 4 fake devices
     #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
@@ -314,7 +316,7 @@ jobs:
         deps: testing_unit
         python-version: '3.14'
     - name: Test SPEC=2
-      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --ignore test/unit/test_autogen.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
+      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore=test/null --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
 
   fuzzing:
     name: Fuzzing
@@ -465,11 +467,11 @@ jobs:
       - name: Test MLPerf stuff
         run: CL=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
       - name: NULL=1 beautiful_mnist_multigpu
-        run: NULL=1 python examples/beautiful_mnist_multigpu.py
+        run: NULL=1 NULL_ALLOW_COPYOUT=1 python examples/beautiful_mnist_multigpu.py
       - name: Test Bert training
-        run: NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+        run: NULL=1 NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
       - name: Test llama 3 training
-        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
+        run: NULL=1 NULL_ALLOW_COPYOUT=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -567,7 +569,7 @@ jobs:
       with:
         key: dsp-minimal
         deps: testing_unit
-        pydeps: "onnx==1.18.0 onnxruntime"
+        pydeps: "onnx==1.18.0 onnxruntime ml_dtypes"
         llvm: "true"
     - name: Set up Docker Buildx
       uses: docker/setup-buildx-action@v3
@@ -579,7 +581,7 @@ jobs:
         load: true
         tags: qemu-hexagon:latest
         cache-from: type=gha
-        cache-to: type=gha,mode=min
+        cache-to: ${{ github.event_name != 'pull_request' && 'type=gha,mode=min' || '' }}
     - name: Set MOCKDSP env
       run: printf "MOCKDSP=1" >> $GITHUB_ENV
     - name: Run test_tiny on DSP
@@ -609,7 +611,7 @@ jobs:
         WEBGPU=1 DEBUG=4 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
     - name: Run selected webgpu tests
       run: |
-          WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
+          WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --durations=20
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
@@ -737,7 +739,7 @@ jobs:
           DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
       - name: Run pytest (cuda)
         # skip multitensor because it's slow
-        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore test/test_gc.py --ignore test/test_multitensor.py --durations=20
+        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --ignore test/test_multitensor.py --durations=20
       - name: Run TestOps.test_add with PMA
         run: VIZ=-1 PMA=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
       - name: Run process replay tests
@@ -770,7 +772,7 @@ jobs:
           python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
           DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
       - name: Run pytest (${{ matrix.backend }})
-        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
+        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --durations=20
       - 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 process replay tests
@@ -797,6 +799,8 @@ jobs:
         llvm: 'true'
     - name: Run unit tests
       run: METAL=1 python -m pytest -n=auto test/unit/ --durations=20
+    - name: Run NULL backend tests
+      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
     - name: Run ONNX
       run: METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
     - name: Test tensor core ops (fake)
@@ -898,7 +902,7 @@ jobs:
           python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU','LVP':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
           DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
       - name: Run pytest (${{ matrix.backend }})
-        run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
+        run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --durations=20
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
       - name: Run macOS-specific unit test
@@ -931,7 +935,11 @@ jobs:
       - name: Run unit tests
         if: matrix.backend=='llvm'
         # test_newton_schulz hits RecursionError
-        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_elf.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
+        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
+      - name: Run NULL backend tests
+        if: matrix.backend=='llvm'
+        shell: bash
+        run: CPU=0 CPU_LLVM=0 NULL=1 python -m pytest -n=auto test/null/ --ignore=test/null/test_elf.py --durations=20
       - name: Run pytest (${{ matrix.backend }})
         shell: bash
         run: |
@@ -957,10 +965,10 @@ jobs:
           key: compile-${{ matrix.backend }}
           deps: testing_unit
           mesa: ${{ (matrix.backend == 'ir3' || matrix.backend == 'nak') && 'true' }}
-          python-version: '3.14'
+          python-version: '3.12'
       - name: Set env
         shell: bash
-        run: printf "NULL=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
+        run: printf "NULL=1\nNULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
       - name: Run test_ops
         shell: bash
         run: |

.pre-commit-config.yaml

@@ -28,7 +28,7 @@ repos:
         pass_filenames: false
       - id: tests
         name: comprehensive test suite
-        entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_schedule.py test/unit/test_assign.py test/test_tensor.py test/test_jit.py test/unit/test_schedule_cache.py test/unit/test_pattern_matcher.py test/unit/test_uop_symbolic.py test/unit/test_helpers.py
+        entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_schedule.py test/unit/test_assign.py test/test_tensor.py test/test_jit.py test/unit/test_schedule_cache.py test/null/test_pattern_matcher.py test/null/test_uop_symbolic.py test/unit/test_helpers.py
         language: system
         always_run: true
         pass_filenames: false

examples/anthropic_challenge.py

@@ -72,7 +72,7 @@ vliw_prepare = PatternMatcher([
   # cast is fake
   (UPat(Ops.CAST, name="c"), lambda c: c.src[0]),
   # rewrites to hardcode the addresses in memory
-  (UPat(Ops.DEFINE_GLOBAL, name="dg"), lambda dg: UOp.const(dtypes.uint, global_addrs[dg.arg])),
+  (UPat(Ops.PARAM, name="dg"), lambda dg: UOp.const(dtypes.uint, global_addrs[dg.arg])),
   # INDEX is just plus
   (UPat(Ops.INDEX, name="i"), lambda i: i.src[0]+i.src[1]),
 ])+symbolic

examples/mlperf/dataloader.py

@@ -72,7 +72,7 @@ def loader_process(q_in, q_out, X:Tensor, seed):
       #storage_tensor._copyin(img_tensor.numpy())
 
       # faster
-      X[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
+      X[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = img.tobytes()
 
       # ideal
       #X[idx].assign(img.tobytes())   # NOTE: this is slow!
@@ -264,8 +264,8 @@ def load_unet3d_data(preprocessed_dataset_dir, seed, queue_in, queue_out, X:Tens
       x = random_brightness_augmentation(x)
       x = gaussian_noise(x)
 
-    X[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = x.tobytes()
-    Y[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = y.tobytes()
+    X[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = x.tobytes()
+    Y[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = y.tobytes()
 
     queue_out.put(idx)
   queue_out.put(None)
@@ -379,12 +379,12 @@ def load_retinanet_data(base_dir:Path, val:bool, queue_in:Queue, queue_out:Queue
       clipped_match_idxs = np.clip(match_idxs, 0, None)
       clipped_boxes, clipped_labels = tgt["boxes"][clipped_match_idxs], tgt["labels"][clipped_match_idxs]
 
-      boxes[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_boxes.tobytes()
-      labels[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_labels.tobytes()
-      matches[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = match_idxs.tobytes()
-      anchors[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = anchor.tobytes()
+      boxes[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = clipped_boxes.tobytes()
+      labels[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = clipped_labels.tobytes()
+      matches[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = match_idxs.tobytes()
+      anchors[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = anchor.tobytes()
 
-    imgs[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
+    imgs[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = img.tobytes()
 
     queue_out.put(idx)
   queue_out.put(None)

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

@@ -2,15 +2,18 @@
 
 export PYTHONPATH="."
 export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 
 export DEBUG=${DEBUG:-2}
 export FLASH_ATTENTION=${FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=8 BS=8 EVAL_BS=8 GRADIENT_ACC_STEPS=2
+export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"
@@ -18,13 +21,14 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="2.5e-4" END_LR="2.5e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
 
-export SEED=5760
+export SEED=${SEED:-5760}
 
-export JITBEAM=3
+export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 
 export FAKEDATA=1 BENCHMARK=10 LLAMA_LAYERS=2

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

@@ -13,7 +13,7 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"
@@ -21,9 +21,10 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="2.5e-4" END_LR="2.5e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
 
 export SEED=${SEED:-5760}
 

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+export BENCHMARK=5
+export EVAL_BS=0
+export VIZ=${VIZ:--1}
+examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
+PYTHONPATH="." extra/viz/cli.py --profile --device "AMD" --top 20

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/startup_walltime.sh

@@ -2,8 +2,9 @@
 export BENCHMARK=5
 export EVAL_BS=0
 export FAKEDATA=1
+export NULL_ALLOW_COPYOUT=1
 export HIP_VISIBLE_DEVICES=""
 export DEV=NULL
 export JITBEAM=0
 export LLAMA_LAYERS=${LLAMA_LAYERS:-"2"}
-time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
+time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh

examples/stable_diffusion.py

@@ -6,7 +6,6 @@ import argparse, time
 from collections import namedtuple
 from typing import Dict, Any
 
-from PIL import Image
 import numpy as np
 from tinygrad import Device, GlobalCounters, dtypes, Tensor, TinyJit
 from tinygrad.helpers import Timing, Context, getenv, fetch, colored, tqdm, flatten, profile_marker
@@ -336,6 +335,7 @@ if __name__ == "__main__":
   print(x.shape)
 
   profile_marker("save image")
+  from PIL import Image
   im = Image.fromarray(x.numpy())
   print(f"saving {args.out}")
   im.save(args.out)

examples/tinychat/tinychat-browser/compile.py

@@ -48,7 +48,7 @@ def prepare_browser_chunks(model):
         weight_metadata = metadata.get(name, default)
         weight_metadata["parts"][part_num] = {"file": i, "file_start_pos": cursor, "size": size}
         metadata[name] = weight_metadata
-        data = bytes(state_dict[name].uop.base.realized.as_buffer())
+        data = bytes(state_dict[name].uop.base.realized.as_memoryview())
         data = data if not offsets else data[offsets[0]:offsets[1]]
         writer.write(data)
         cursor += size

extra/amdpci/am_smi.py

@@ -92,7 +92,7 @@ class SMICtx:
     self.prev_terminal_width = 0
     self.prev_terminal_height = 0
 
-    remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:"]
+    remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:", "Processing accelerators:"]
     lspci = subprocess.check_output(["lspci"]).decode("utf-8").splitlines()
     self.lspci = {l.split()[0]: l.split(" ", 1)[1] for l in lspci}
     for k,v in self.lspci.items():
@@ -153,7 +153,8 @@ class SMICtx:
     tables = {}
     for dev in self.devs:
       match dev.ip_ver[am.MP1_HWIP]:
-        case (13,0,6)|(13,0,12): table_t = dev.smu.smu_mod.MetricsTableX_t
+        case (13,0,6): table_t = dev.smu.smu_mod.MetricsTableV0_t
+        case (13,0,12): table_t = dev.smu.smu_mod.MetricsTableV2_t
         case _: table_t = dev.smu.smu_mod.SmuMetricsExternal_t
       tables[dev] = dev.smu.read_table(table_t, dev.smu.smu_mod.SMU_TABLE_SMU_METRICS) if dev.pci_state == "D0" else None
     return tables
@@ -279,7 +280,7 @@ class SMICtx:
       device_line = [f"{bold(dev.pcibus)} {trim(self.lspci[dev.pcibus[5:]], col_size - 20)}"] + [pad("", col_size)]
       activity_line = [f"GFX Activity {draw_bar(self.get_gfx_activity(dev, metrics) / 100, activity_line_width)}"] \
                     + [f"MEM Activity {draw_bar(self.get_mem_activity(dev, metrics) / 100, activity_line_width)}"] \
-                    + [f"MEM Usage    {draw_bar((mem_used / mem_total) / 100, activity_line_width, opt_text=mem_fmt)}"] \
+                    + [f"MEM Usage    {draw_bar(mem_used / mem_total, activity_line_width, opt_text=mem_fmt)}"] \
 
       temps_data, temps_data_compact = self.get_temps(dev, metrics), self.get_temps(dev, metrics, compact=True)
       temps_table = ["=== Temps (°C) ==="] + [f"{name:<16}: {color_temp(val)}" for name, val in temps_data.items()]

extra/amdpci/hive_reset.py

@@ -1,12 +1,18 @@
 #!/usr/bin/env python3
 
+import os
 from tinygrad.helpers import Context
 from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
+from tinygrad.runtime.support.hcq import FileIOInterface
 from tinygrad.runtime.support.am.amdev import AMDev
 
 if __name__ == "__main__":
   gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1, 0x75a0])])
-  pcidevs = [PCIDevice(f"reset:{gpu}", gpu, bars=[0, 2, 5]) for gpu in gpus]
+  for gpu in gpus:
+    drv_path = f"/sys/bus/pci/devices/{gpu}/driver"
+    if FileIOInterface.exists(drv_path) and os.path.basename(os.readlink(drv_path)) == "amdgpu":
+      raise RuntimeError(f"amdgpu is bound to {gpu}. Stopping...")
+  pcidevs = [PCIDevice("AM", gpu, bars=[0, 2, 5]) for gpu in gpus]
   amdevs = []
   with Context(DEBUG=2):
     for pcidev in pcidevs:

extra/amdpci/proclogs.py

@@ -19,8 +19,9 @@ amdev = importlib.import_module("tinygrad.runtime.support.am.amdev")
 amdev.AMDev = AMDFake
 from tinygrad.runtime.ops_amd import PCIIface
 
-def parse_amdgpu_logs(log_content, register_names=None, *, only_xcc0: bool = False):
+def parse_amdgpu_logs(log_content, register_names=None, register_objects=None, *, only_xcc0: bool = False):
   register_map = register_names or {}
+  register_objs = register_objects or {}
 
   def replace_register(match):
     reg = match.group(1)
@@ -37,6 +38,28 @@ def parse_amdgpu_logs(log_content, register_names=None, *, only_xcc0: bool = Fal
   # remove timing prefix
   processed_log = re.sub(r'^\[\s*\d+(?:\.\d+)?\]\s*', '', processed_log, flags=re.MULTILINE)
 
+  # decode register values into field dicts
+  def decode_value(match):
+    reg_name = match.group(1)
+    xcc_part = match.group(2)  # "xcc=0 " or ""
+    val_str = match.group(3)
+    val = int(val_str, 16)
+
+    reg_obj = register_objs.get(reg_name)
+    if reg_obj is not None and reg_obj.fields:
+      fields = reg_obj.decode(val)
+      # show raw for unaccounted bits
+      accounted = 0
+      for name, (start, end) in reg_obj.fields.items():
+        accounted |= (((1 << (end - start + 1)) - 1) << start)
+      unaccounted = val & ~accounted
+      parts = {k: v for k, v in fields.items() if v != 0}
+      if unaccounted: parts['_raw_unaccounted'] = hex(unaccounted)
+      return f"register {reg_name}, {xcc_part}with value {val_str} {parts}"
+    return match.group(0)
+
+  processed_log = re.sub(r'register (reg\w+), ((?:xcc=\d+ )?)with value (0x[0-9a-fA-F]+)', decode_value, processed_log)
+
   # keep only xcc=0 lines (but keep lines with no xcc at all)
   if only_xcc0:
     kept = []
@@ -50,16 +73,18 @@ def main():
   only_xcc0 = bool(getenv("ONLY_XCC0", 0))
 
   reg_names = {}
+  reg_objs = {}
   dev = PCIIface(None, 0)
   for x, y in dev.dev_impl.__dict__.items():
     if isinstance(y, AMRegister):
       for xcc, addr in y.addr.items():
         reg_names[addr] = f"{x}, xcc={xcc}"
+      reg_objs[x] = y
 
   with open(sys.argv[1], 'r') as f:
     log_content = f.read()
 
-  processed_log = parse_amdgpu_logs(log_content, reg_names, only_xcc0=only_xcc0)
+  processed_log = parse_amdgpu_logs(log_content, reg_names, reg_objs, only_xcc0=only_xcc0)
 
   with open(sys.argv[2], 'w') as f:
     f.write(processed_log)

extra/assembly/amd/emu.py

@@ -135,14 +135,6 @@ def _val_to_u32(val: UOp) -> UOp:
   if val.dtype in (dtypes.uint16, dtypes.int16): return val.cast(dtypes.uint32)
   return val.cast(dtypes.uint32)
 
-def _apply_clamp(val: UOp, clmp: int | UOp) -> UOp:
-  """Apply VOP3 clamp modifier: clamp float results to [0.0, 1.0] range."""
-  if isinstance(clmp, int) and clmp == 0: return val
-  if val.dtype not in (dtypes.float32, dtypes.half, dtypes.float64): return val
-  zero, one = UOp.const(val.dtype, 0.0), UOp.const(val.dtype, 1.0)
-  clamped = val.maximum(zero).minimum(one)
-  return clmp.ne(_c(0)).where(clamped, val) if isinstance(clmp, UOp) else clamped
-
 _pcode_fixes = {
   'V_DIV_FMAS_F32': ('D0.f32 = 2.0F ** 32 * fma(S0.f32, S1.f32, S2.f32)',
     'D0.f32 = (exponent(S2.f32) > 127) ? (2.0F ** 64 * fma(S0.f32, S1.f32, S2.f32)) : (2.0F ** -64 * fma(S0.f32, S1.f32, S2.f32))'),
@@ -265,11 +257,11 @@ def _collect_data_slices(assigns: list[tuple[str, UOp]], data_prefix: str, pcode
 class _Ctx:
   """Context for instruction compilation - holds buffers and helpers."""
   __slots__ = ('inst_size', 'dyn_fields', '_axis_id')
-  sgpr = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(SGPR_COUNT), arg=0)
-  vgpr = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(VGPR_SIZE), arg=1)
-  vmem = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(1 << 46), arg=2)
-  lds = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(16384), arg=3)
-  scratch = UOp(Ops.DEFINE_GLOBAL, dtypes.uint8.ptr(1 << 30), arg=4)
+  sgpr = UOp(Ops.PARAM, dtypes.uint32.ptr(SGPR_COUNT), arg=0)
+  vgpr = UOp(Ops.PARAM, dtypes.uint32.ptr(VGPR_SIZE), arg=1)
+  vmem = UOp(Ops.PARAM, dtypes.uint32.ptr(1 << 46), arg=2)
+  lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
+  scratch = UOp(Ops.PARAM, dtypes.uint8.ptr(1 << 30), arg=4)
 
   def __init__(self, inst_size: int):
     self.inst_size, self._axis_id = inst_size, 0
@@ -354,7 +346,7 @@ class _Ctx:
     offset = reg.cast(dtypes.int) * _c(32, dtypes.int) + lane.cast(dtypes.int)
     return buf.index(offset, _lane_active(exec_mask, lane)).store(val.cast(dtypes.uint32))
 
-  def rsrc_dyn(self, off: UOp, lane: UOp | None, bits: int = 32, literal: UOp | None = None, is_f64: bool = False) -> UOp:
+  def rsrc_dyn(self, off: UOp, lane: UOp | None, bits: int = 32, literal: UOp | None = None, is_f64: bool = False, do_cast: bool = True) -> UOp:
     """Read source operand with dynamic offset. Handles SGPR/inline constants (<256), VGPR (>=256).
     If lane is None, only scalar access is supported (off must be < 256).
     is_f64: True for F64 operations where 64-bit literals go in high 32 bits."""
@@ -385,7 +377,7 @@ class _Ctx:
     else:
       scalar_val = sgpr_lo
       if literal is not None: scalar_val = off.eq(_c(255)).where(literal, scalar_val)
-      if bits == 16:  # Float constants: cast F32 to F16
+      if bits == 16 and do_cast:  # Float constants: cast F32 to F16
         scalar_val = is_float_const.where(scalar_val.bitcast(dtypes.float32).cast(dtypes.half).bitcast(dtypes.uint16).cast(dtypes.uint32), scalar_val)
 
     return is_vgpr.where(vgpr_val, scalar_val) if lane is not None else scalar_val
@@ -445,7 +437,7 @@ class _Ctx:
     return UOp.sink(*stores, *self.inc_pc())
 
   def compile_vop_pcode(self, op, srcs: dict[str, UOp], lane: UOp, vdst_reg: UOp, exec_mask: UOp,
-                        opsel_dst_hi: bool | UOp = False, sdst_reg: int | None = None, clmp: int | UOp = 0) -> UOp:
+                        opsel_dst_hi: bool | UOp = False, sdst_reg: int | None = None, clmp: int = 0) -> UOp:
     """Compile VOP instruction. Returns sink with stores and inc_pc."""
     pcode = get_pcode(op)
     vcc_reg = sdst_reg if sdst_reg is not None else VCC_LO.offset
@@ -454,6 +446,24 @@ class _Ctx:
                  'ROUND_MODE': _c(0), 'ROUND_TOWARD_ZERO': _c(0)})  # rounding mode: 0=RNE, RTZ constant
     _, assigns = parse_pcode(pcode, srcs)
 
+    # For integer ops with clamp, compute overflow using wide arithmetic
+    # NOTE: MUL_LO ops don't saturate - they always return the low bits
+    int_saturate = None
+    if clmp and any(p in op.name for p in ('_NC_U', '_MAD_U', '_NC_I', '_MAD_I')):
+      is_signed, is_16bit = '_I' in op.name and '_U' not in op.name, '16' in op.name
+      if not (is_16bit and is_signed):  # Skip 16-bit signed ops due to codegen issues
+        s0, s1, s2 = srcs.get('S0'), srcs.get('S1'), srcs.get('S2')
+        if s0 is not None and s1 is not None:
+          narrow_dt = dtypes.uint16 if is_16bit else (dtypes.int32 if is_signed else dtypes.uint32)
+          wide_dt = dtypes.int32 if is_16bit else dtypes.int64
+          narrow_max, narrow_min = (0xFFFF, 0) if is_16bit else ((0x7FFFFFFF, -0x80000000) if is_signed else (0xFFFFFFFF, 0))
+          def to_wide(x): return (x.bitcast(narrow_dt) if x.dtype.itemsize == narrow_dt.itemsize else x.cast(narrow_dt)).cast(wide_dt)
+          is_sub, is_mad = 'SUB' in op.name, 'MAD' in op.name
+          full = (to_wide(s0) * to_wide(s1) + to_wide(s2)) if is_mad and s2 is not None else \
+                 (to_wide(s1) - to_wide(s0)) if is_sub and 'SUBREV' in op.name else \
+                 (to_wide(s0) - to_wide(s1)) if is_sub else (to_wide(s0) + to_wide(s1))
+          int_saturate = full.clamp(narrow_min, narrow_max).cast(narrow_dt)
+
     raw_stores: list = []
     vcc_val, exec_val = None, None
     for dest, val in assigns:
@@ -468,7 +478,10 @@ class _Ctx:
                        val.cast(dtypes.uint32) if val.dtype in (dtypes.uint16, dtypes.int16) else val.cast(dtypes.uint32) & UOp.const(dtypes.uint32, slice_mask)
             raw_stores.append(('vgpr_slice', (lo_bit, width, val_bits)))
             continue
-        val = _apply_clamp(val, clmp)
+        # For integer ops with clamp, use pre-computed saturated value; for floats, clamp to [0,1]
+        if int_saturate is not None: val = int_saturate
+        elif clmp and val.dtype in (dtypes.float32, dtypes.half, dtypes.float64):
+          val = val.maximum(UOp.const(val.dtype, 0.0)).minimum(UOp.const(val.dtype, 1.0))
         if val.dtype in (dtypes.uint64, dtypes.int64, dtypes.float64):
           lo, hi = _split64(val)
           raw_stores.extend([('vgpr', self.wvgpr_dyn(vdst_reg, lane, lo, exec_mask)), ('vgpr', self.wvgpr_dyn(vdst_reg + _c(1), lane, hi, exec_mask))])
@@ -732,6 +745,7 @@ def _compile_vop3sd(inst: ir3.VOP3SD | ir4.VOP3SD, ctx: _Ctx) -> UOp:
   _, assigns = parse_pcode(pcode, srcs)
 
   has_per_lane_vcc = any('[laneId]' in dest for dest, _ in assigns if dest.startswith('VCC') or dest.startswith('D0.u64'))
+  clmp = getattr(inst, 'clmp', 0)
   if has_per_lane_vcc:
     # VCC computation: RANGE+REDUCE gets axis ID first (lower ID = runs first)
     # This ensures VCC reads source values BEFORE VGPR stores modify them
@@ -743,11 +757,17 @@ def _compile_vop3sd(inst: ir3.VOP3SD | ir4.VOP3SD, ctx: _Ctx) -> UOp:
     final_vcc = ctx.unroll_lanes(get_vcc_bit, exec_mask)
     # VGPR stores: RANGE gets axis ID second (higher ID = runs after VCC loop)
     lane3 = ctx.range()
-    d0_val = None
+    d0_val, vcc_per_lane = None, None
     for dest, val in parse_pcode(pcode, load_srcs(lane3))[1]:
       if dest.startswith('D0') and '[laneId]' not in dest: d0_val = val
+      if dest.startswith('VCC') or (dest.startswith('D0.u64') and '[laneId]' in dest): vcc_per_lane = val
     vgpr_stores = []
     if d0_val is not None:
+      # Apply clamp using carry/borrow bit: ADD overflow->0xFFFFFFFF, SUB underflow->0
+      if clmp and vcc_per_lane is not None:
+        is_sub = 'SUB' in inst.op.name
+        sat_val = _c(0) if is_sub else _c(0xFFFFFFFF)
+        d0_val = vcc_per_lane.cast(dtypes.bool).where(sat_val, d0_val.cast(dtypes.uint32))
       if d0_val.dtype in (dtypes.uint64, dtypes.int64, dtypes.float64):
         lo, hi = _split64(d0_val)
         vgpr_stores.extend([ctx.wvgpr_dyn(vdst_reg, lane3, lo, exec_mask), ctx.wvgpr_dyn(vdst_reg + _c(1), lane3, hi, exec_mask)])
@@ -800,9 +820,10 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P, ctx: _Ctx) -> UOp:
   lane = ctx.range()
   exec_mask = ctx.rsgpr_dyn(_c(EXEC_LO.offset))
   vdst_reg = ctx.inst_field(type(inst).vdst)
-  src0 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src0), lane, 16)
-  src1 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src1), lane, 16)
-  src2 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src2), lane, 16)
+  do_cast = any(x in op_name for x in ('F16', 'F32', 'BF16')) and 'IU' not in op_name
+  src0 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src0), lane, 16, do_cast=do_cast)
+  src1 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src1), lane, 16, do_cast=do_cast)
+  src2 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src2), lane, 16, do_cast=do_cast)
   opsel, opsel_hi = getattr(inst, 'opsel', 0) or 0, getattr(inst, 'opsel_hi', 3) if getattr(inst, 'opsel_hi', 3) is not None else 3
   opsel_hi2 = getattr(inst, 'opsel_hi2', 1) if getattr(inst, 'opsel_hi2', 1) is not None else 1
   neg, neg_hi = getattr(inst, 'neg', 0) or 0, getattr(inst, 'neg_hi', 0) or 0
@@ -1120,8 +1141,8 @@ class WaveState:
     self.n_lanes = n_lanes
     self.vgpr_buf = Buffer('CPU', VGPR_SIZE, dtypes.uint32).ensure_allocated()
     self.sgpr_buf = Buffer('CPU', SGPR_COUNT, dtypes.uint32).ensure_allocated()
-    self._vgpr_mv = self.vgpr_buf.as_buffer(force_zero_copy=True).cast('I')
-    self._sgpr_mv = self.sgpr_buf.as_buffer(force_zero_copy=True).cast('I')
+    self._vgpr_mv = self.vgpr_buf.as_memoryview(force_zero_copy=True).cast('I')
+    self._sgpr_mv = self.sgpr_buf.as_memoryview(force_zero_copy=True).cast('I')
     # Zero memory using ctypes memset (much faster than Python loops)
     ctypes.memset(self.vgpr_buf._buf.va_addr, 0, VGPR_SIZE * 4)
     ctypes.memset(self.sgpr_buf._buf.va_addr, 0, SGPR_COUNT * 4)

extra/assembly/amd/pcode.py

@@ -440,7 +440,7 @@ class Parser:
       self.eat('COMMA')
       lo = self.parse()
       self.eat('RBRACE')
-      return (hi.cast(dtypes.uint64) << _u64(32)) | lo.cast(dtypes.uint64)
+      return (hi.cast(dt:=_BITS_DT.get((s:=lo.dtype.bitsize) * 2, dtypes.uint64)) << _const(dt, s)) | lo.cast(dt)
     if self.at('NUM'):
       num = self.eat('NUM').val
       if self.try_eat('QUOTE'):

extra/assembly/amd/test/hw/test_sop.py

@@ -87,6 +87,50 @@ class TestBasicScalar(unittest.TestCase):
     self.assertEqual(st.sgpr[2], f2i(2.0))
 
 
+class TestPack(unittest.TestCase):
+  """Tests for S_PACK instructions."""
+
+  def test_s_pack_ll_b32_b16(self):
+    """S_PACK_LL_B32_B16 packs low 16 bits of two sources into one 32-bit result."""
+    instructions = [
+      s_mov_b32(s[0], 0xDEADAAAA),
+      s_mov_b32(s[1], 0xDEADBBBB),
+      s_pack_ll_b32_b16(s[2], s[0], s[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.sgpr[2], 0xBBBBAAAA)
+
+  def test_s_pack_lh_b32_b16(self):
+    """S_PACK_LH_B32_B16: D0 = { S1[31:16], S0[15:0] }."""
+    instructions = [
+      s_mov_b32(s[0], 0xDEADAAAA),
+      s_mov_b32(s[1], 0xDEADBBBB),
+      s_pack_lh_b32_b16(s[2], s[0], s[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.sgpr[2], 0xDEADAAAA)
+
+  def test_s_pack_hh_b32_b16(self):
+    """S_PACK_HH_B32_B16: D0 = { S1[31:16], S0[31:16] }."""
+    instructions = [
+      s_mov_b32(s[0], 0xDEADAAAA),
+      s_mov_b32(s[1], 0xDEADBBBB),
+      s_pack_hh_b32_b16(s[2], s[0], s[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.sgpr[2], 0xDEADDEAD)
+
+  def test_s_pack_hl_b32_b16(self):
+    """S_PACK_HL_B32_B16: D0 = { S1[15:0], S0[31:16] }."""
+    instructions = [
+      s_mov_b32(s[0], 0xDEADAAAA),
+      s_mov_b32(s[1], 0xDEADBBBB),
+      s_pack_hl_b32_b16(s[2], s[0], s[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.sgpr[2], 0xBBBBDEAD)
+
+
 class TestQuadmaskWqm(unittest.TestCase):
   """Tests for S_QUADMASK_B32 and S_WQM_B32."""
 

extra/assembly/amd/test/hw/test_vop3.py

@@ -2949,6 +2949,306 @@ class TestVOP3Clamp(unittest.TestCase):
     self.assertAlmostEqual(i2f(st.vgpr[3][1]), 1.0, places=5, msg="lane 3: 2.5 should clamp to 1.0")
 
 
+class TestVOP3ClampUint32(unittest.TestCase):
+  """Tests for VOP3 clamp modifier on unsigned 32-bit integer operations."""
+
+  def test_v_sub_nc_u32_e64_clamp_underflow(self):
+    """V_SUB_NC_U32_E64 with clamp: 0 - 1 should saturate to 0."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0),
+      v_mov_b32_e32(v[1], 1),
+      v_sub_nc_u32_e64(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_sub_nc_u32_e64_clamp_no_underflow(self):
+    """V_SUB_NC_U32_E64 with clamp: 100 - 50 = 50 (no saturation needed)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 100),
+      v_mov_b32_e32(v[1], 50),
+      v_sub_nc_u32_e64(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 50, f"expected 50, got {st.vgpr[0][2]}")
+
+  def test_v_add_nc_u32_e64_clamp_overflow(self):
+    """V_ADD_NC_U32_E64 with clamp: 0xFFFFFFFF + 1 should saturate to 0xFFFFFFFF."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0xFFFFFFFF),
+      v_mov_b32_e32(v[1], 1),
+      v_add_nc_u32_e64(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xFFFFFFFF, f"expected 0xFFFFFFFF, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_add_nc_u32_e64_clamp_no_overflow(self):
+    """V_ADD_NC_U32_E64 with clamp: 100 + 50 = 150 (no saturation needed)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 100),
+      v_mov_b32_e32(v[1], 50),
+      v_add_nc_u32_e64(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 150, f"expected 150, got {st.vgpr[0][2]}")
+
+
+class TestVOP3ClampUint16(unittest.TestCase):
+  """Tests for VOP3 clamp modifier on unsigned 16-bit integer operations."""
+
+  def test_v_sub_nc_u16_clamp_underflow(self):
+    """V_SUB_NC_U16 with clamp: 0 - 1 should saturate to 0."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0),
+      v_mov_b32_e32(v[1], 1),
+      v_sub_nc_u16(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2] & 0xFFFF, 0, f"expected 0, got 0x{st.vgpr[0][2] & 0xFFFF:04x}")
+
+  def test_v_sub_nc_u16_clamp_no_underflow(self):
+    """V_SUB_NC_U16 with clamp: 100 - 50 = 50 (no saturation needed)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 100),
+      v_mov_b32_e32(v[1], 50),
+      v_sub_nc_u16(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2] & 0xFFFF, 50, f"expected 50, got {st.vgpr[0][2] & 0xFFFF}")
+
+  def test_v_add_nc_u16_clamp_overflow(self):
+    """V_ADD_NC_U16 with clamp: 0xFFFF + 1 should saturate to 0xFFFF."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0xFFFF),
+      v_mov_b32_e32(v[1], 1),
+      v_add_nc_u16(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2] & 0xFFFF, 0xFFFF, f"expected 0xFFFF, got 0x{st.vgpr[0][2] & 0xFFFF:04x}")
+
+  def test_v_add_nc_u16_clamp_no_overflow(self):
+    """V_ADD_NC_U16 with clamp: 100 + 50 = 150 (no saturation needed)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 100),
+      v_mov_b32_e32(v[1], 50),
+      v_add_nc_u16(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2] & 0xFFFF, 150, f"expected 150, got {st.vgpr[0][2] & 0xFFFF}")
+
+
+class TestVOP3ClampInt32(unittest.TestCase):
+  """Tests for VOP3 clamp modifier on signed 32-bit integer operations."""
+
+  def test_v_add_nc_i32_clamp_overflow(self):
+    """V_ADD_NC_I32 with clamp: INT_MAX + 1 should saturate to INT_MAX."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0x7FFFFFFF),  # S0 = INT_MAX
+      v_mov_b32_e32(v[1], 1),            # S1 = 1
+      v_add_nc_i32(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x7FFFFFFF, f"expected 0x7FFFFFFF, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_add_nc_i32_clamp_underflow(self):
+    """V_ADD_NC_I32 with clamp: INT_MIN + (-1) should saturate to INT_MIN."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0x80000000),  # S0 = INT_MIN
+      v_mov_b32_e32(v[1], 0xFFFFFFFF),  # S1 = -1
+      v_add_nc_i32(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x80000000, f"expected 0x80000000, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_sub_nc_i32_clamp_underflow(self):
+    """V_SUB_NC_I32 with clamp: INT_MIN - 1 should saturate to INT_MIN."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0x80000000),  # S0 = INT_MIN
+      v_mov_b32_e32(v[1], 1),            # S1 = 1
+      v_sub_nc_i32(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x80000000, f"expected 0x80000000, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_sub_nc_i32_clamp_overflow(self):
+    """V_SUB_NC_I32 with clamp: INT_MAX - (-1) should saturate to INT_MAX."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0x7FFFFFFF),  # S0 = INT_MAX
+      v_mov_b32_e32(v[1], 0xFFFFFFFF),  # S1 = -1
+      v_sub_nc_i32(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x7FFFFFFF, f"expected 0x7FFFFFFF, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_add_nc_i32_no_saturation_positive(self):
+    """V_ADD_NC_I32 with clamp: 100 + 200 = 300 (no saturation needed)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 100),
+      v_mov_b32_e32(v[1], 200),
+      v_add_nc_i32(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 300, f"expected 300, got {st.vgpr[0][2]}")
+
+  def test_v_add_nc_i32_no_saturation_negative(self):
+    """V_ADD_NC_I32 with clamp: -100 + -200 = -300 (no saturation needed)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0xFFFFFF9C),  # -100
+      v_mov_b32_e32(v[1], 0xFFFFFF38),  # -200
+      v_add_nc_i32(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    expected = 0xFFFFFED4  # -300
+    self.assertEqual(st.vgpr[0][2], expected, f"expected 0x{expected:08x}, got 0x{st.vgpr[0][2]:08x}")
+
+
+class TestVOP3ClampCarry(unittest.TestCase):
+  """Tests for VOP3 clamp modifier on carry operations (VOP3SD)."""
+
+  def test_v_add_co_u32_clamp_overflow(self):
+    """V_ADD_CO_U32 with clamp: 0xFFFFFFFF + 1 should saturate to 0xFFFFFFFF."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0xFFFFFFFF),
+      v_mov_b32_e32(v[1], 1),
+      v_add_co_u32(v[2], VCC, v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xFFFFFFFF, f"expected 0xFFFFFFFF, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_add_co_u32_clamp_no_overflow(self):
+    """V_ADD_CO_U32 with clamp: 100 + 200 = 300 (no saturation)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 100),
+      v_mov_b32_e32(v[1], 200),
+      v_add_co_u32(v[2], VCC, v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 300, f"expected 300, got {st.vgpr[0][2]}")
+
+  def test_v_sub_co_u32_clamp_underflow(self):
+    """V_SUB_CO_U32 with clamp: 0 - 1 should saturate to 0."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0),
+      v_mov_b32_e32(v[1], 1),
+      v_sub_co_u32(v[2], VCC, v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_sub_co_u32_clamp_no_underflow(self):
+    """V_SUB_CO_U32 with clamp: 300 - 100 = 200 (no saturation)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 300),
+      v_mov_b32_e32(v[1], 100),
+      v_sub_co_u32(v[2], VCC, v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 200, f"expected 200, got {st.vgpr[0][2]}")
+
+  def test_v_subrev_co_u32_clamp_underflow(self):
+    """V_SUBREV_CO_U32 with clamp: 1 - 0 reversed = 0 - 1 should saturate to 0."""
+    instructions = [
+      v_mov_b32_e32(v[0], 1),  # This becomes the subtrahend
+      v_mov_b32_e32(v[1], 0),  # This becomes the minuend (0 - 1)
+      v_subrev_co_u32(v[2], VCC, v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_add_co_ci_u32_clamp_overflow(self):
+    """V_ADD_CO_CI_U32 with clamp: 0xFFFFFFFF + 1 + 0 should saturate to 0xFFFFFFFF."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0xFFFFFFFF),
+      v_mov_b32_e32(v[1], 1),
+      s_mov_b64(VCC, 0),  # No carry in
+      v_add_co_ci_u32(v[2], VCC, v[0], v[1], VCC, clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xFFFFFFFF, f"expected 0xFFFFFFFF, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_add_co_ci_u32_clamp_overflow_with_carry(self):
+    """V_ADD_CO_CI_U32 with clamp: 0xFFFFFFFE + 1 + 1 should saturate to 0xFFFFFFFF."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0xFFFFFFFE),
+      v_mov_b32_e32(v[1], 1),
+      s_mov_b64(VCC, 1),  # Carry in = 1
+      v_add_co_ci_u32(v[2], VCC, v[0], v[1], VCC, clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xFFFFFFFF, f"expected 0xFFFFFFFF, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_sub_co_ci_u32_clamp_underflow(self):
+    """V_SUB_CO_CI_U32 with clamp: 0 - 1 - 0 should saturate to 0."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0),
+      v_mov_b32_e32(v[1], 1),
+      s_mov_b64(VCC, 0),  # No borrow in
+      v_sub_co_ci_u32(v[2], VCC, v[0], v[1], VCC, clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_subrev_co_ci_u32_clamp_underflow(self):
+    """V_SUBREV_CO_CI_U32 with clamp: reversed 1 - 0 - 0 = 0 - 1 should saturate to 0."""
+    instructions = [
+      v_mov_b32_e32(v[0], 1),
+      v_mov_b32_e32(v[1], 0),
+      s_mov_b64(VCC, 0),
+      v_subrev_co_ci_u32(v[2], VCC, v[0], v[1], VCC, clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
+
+
+class TestVOP3ClampMAD(unittest.TestCase):
+  """Tests for VOP3 clamp modifier on MAD (multiply-add) operations."""
+
+  def test_v_mad_u16_clamp_overflow(self):
+    """V_MAD_U16 with clamp: 0xFFFF * 2 + 0 should saturate to 0xFFFF."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0xFFFF),
+      v_mov_b32_e32(v[1], 2),
+      v_mov_b32_e32(v[2], 0),
+      v_mad_u16(v[3], v[0], v[1], v[2], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][3] & 0xFFFF, 0xFFFF, f"expected 0xFFFF, got 0x{st.vgpr[0][3] & 0xFFFF:04x}")
+
+  def test_v_mad_u16_clamp_overflow_with_add(self):
+    """V_MAD_U16 with clamp: 0x8000 * 2 + 0x1000 should saturate to 0xFFFF."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0x8000),  # 32768
+      v_mov_b32_e32(v[1], 2),        # * 2 = 65536
+      v_mov_b32_e32(v[2], 0x1000),   # + 4096 = 69632 > 0xFFFF
+      v_mad_u16(v[3], v[0], v[1], v[2], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][3] & 0xFFFF, 0xFFFF, f"expected 0xFFFF, got 0x{st.vgpr[0][3] & 0xFFFF:04x}")
+
+  def test_v_mad_u16_no_overflow(self):
+    """V_MAD_U16 with clamp: 100 * 100 + 50 = 10050 (no saturation)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 100),
+      v_mov_b32_e32(v[1], 100),
+      v_mov_b32_e32(v[2], 50),
+      v_mad_u16(v[3], v[0], v[1], v[2], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][3] & 0xFFFF, 10050, f"expected 10050, got {st.vgpr[0][3] & 0xFFFF}")
+
+  def test_v_mad_u16_no_clamp(self):
+    """V_MAD_U16 without clamp: 0xFFFF * 2 + 0 should wrap to 0xFFFE."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0xFFFF),
+      v_mov_b32_e32(v[1], 2),
+      v_mov_b32_e32(v[2], 0),
+      v_mad_u16(v[3], v[0], v[1], v[2], clmp=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # 0xFFFF * 2 = 0x1FFFE, low 16 bits = 0xFFFE
+    self.assertEqual(st.vgpr[0][3] & 0xFFFF, 0xFFFE, f"expected 0xFFFE, got 0x{st.vgpr[0][3] & 0xFFFF:04x}")
+
+
 class TestCvtPkF16(unittest.TestCase):
   """Tests for V_CVT_PK_RTZ_F16_F32 - pack two f32 to f16 with round toward zero."""
 

extra/assembly/amd/test/hw/test_vop3p.py

@@ -390,6 +390,24 @@ class TestVOP3P(unittest.TestCase):
     self.assertAlmostEqual(lo, 6.0, places=1)
     self.assertAlmostEqual(hi, 0.0, places=1)
 
+  def test_v_pk_add_u16_float_inline_const_opsel(self):
+    """V_PK_ADD_U16 with float inline constant 2.0
+    Regression test: for integer packed ops, do not perform the f32->f16 conversion.
+    """
+    # src1 = inline float constant 2.0
+    instructions = [
+      s_mov_b32(s[0], 0x00030005),  # packed u16: hi=3, lo=5
+      v_mov_b32_e32(v[0], s[0]),
+      v_pk_add_u16(v[1], v[0], SrcEnum.POS_TWO, opsel_hi=3, opsel_hi2=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][1]
+    lo = result & 0xffff
+    hi = (result >> 16) & 0xffff
+    # lo = 5 + 0x0000 = 0x0005, hi = 3 + 0x4000 = 0x4003
+    self.assertEqual(lo, 0x0005, f"lo: expected 0x0005, got 0x{lo:04x}")
+    self.assertEqual(hi, 0x4003, f"hi: expected 0x4003, got 0x{hi:04x}")
+
 
 class TestWMMAF16(unittest.TestCase):
   """Tests for WMMA F16 output variant (V_WMMA_F16_16X16X16_F16).

extra/assembly/amd/test/test_emu2_pcode.py

@@ -180,7 +180,7 @@ class TestDSPcodePatterns(unittest.TestCase):
   def test_mem_read_parsing(self):
     """Test MEM[addr].type read expression parsing."""
     # Create a mock LDS buffer
-    lds = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(16384), arg=3)
+    lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
     addr = UOp.const(dtypes.uint32, 0)
     vars = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(dtypes.uint32, 0)}
 
@@ -213,7 +213,7 @@ class TestDSPcodePatterns(unittest.TestCase):
     """Test DS_LOAD_2ADDR_B32 pcode parsing produces RETURN_DATA assignments."""
     pcode = PCODE.get(DSOp.DS_LOAD_2ADDR_B32)
     self.assertIsNotNone(pcode)
-    lds = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(16384), arg=3)
+    lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
     srcs = {
       'ADDR': UOp.const(dtypes.uint32, 0),
       'OFFSET0': UOp.const(dtypes.uint32, 0),
@@ -286,7 +286,7 @@ class TestAllPcode(unittest.TestCase):
   def _make_srcs(self):
     """Create dummy source variables for pcode parsing."""
     u32, u64 = lambda v=0: UOp.const(dtypes.uint32, v), lambda v=0: UOp.const(dtypes.uint64, v)
-    lds = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(16384), arg=3)
+    lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
     return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(),
             'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(),
             'VCC': u64(), 'VCCZ': u32(), 'EXEC': u64(), 'EXEC_LO': u32(), 'EXECZ': u32(), 'SCC': u32(),

extra/fp8/fp8_linear.py

@@ -28,7 +28,7 @@ def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
   return store_op.sink(arg=KernelInfo(name=f"fp8_matmul_{inp.shape}x{weight.shape}"))
 
 def custom_matmul_backward(gradient: UOp, kernel: UOp) -> tuple[UOp, UOp]:
-  _, input_uop, weight_uop = kernel.src
+  _, input_uop, weight_uop = kernel.src[1:]
   input_tensor = Tensor(input_uop, device=input_uop.device)
   grad_tensor = Tensor(gradient, device=gradient.device)
   weight_tensor = Tensor(weight_uop, device=weight_uop.device)

extra/gemm/asm/cdna/asm.py

@@ -9,11 +9,17 @@ GEMM_ARGS = {
   (8192, 4096, 4096): (256, 64, 32768),
   (8192, 14336, 4096): (256, 64, 114688),
   (8192, 4096, 14336): (256, 224, 114688),
-  (8192, 128256, 4096): (16032, 64, 1026048),
+  # TODO: get a fast gemm for this shape
+  #(8192, 128256, 4096): (16032, 64, 1026048),
   (8192, 8192, 8192): (256, 128, 131072),
   (4096, 4096, 4096): (256, 64, 16384),
   (4096, 14336, 4096): (256, 64, 57344),
+  (4096, 14336, 8192): (256, 128, 114688),
   (4096, 4096, 14336): (256, 224, 57344),
+  (14336, 4096, 8192): (256, 128, 114688),
+  (4096, 8192, 14336): (256, 224, 114688),
+  (4096, 4096, 8192): (256, 128, 32768),
+  (4096, 8192, 4096): (256, 64, 32768),
 }
 ITERS_ARGS = {64: (67108864, 0), 128: (33554432, 0), 224: (613566757, 2147483656)}
 

extra/gemm/asm/cdna/gemm.py

@@ -30,14 +30,18 @@ atexit.register(lambda: print(f'asm_gemm: {counters["used"]} used, {len(counters
 def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
   if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
   if a.dtype not in {dtypes.bfloat16, dtypes.float16}: return todo(f"only bfloat16/float16, got {a.dtype}")
-  # only sharding on the batch is tested, others might work too
-  if isinstance(a.device, tuple) and not (a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None):
-    return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
   batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
   N = b.shape[1]
-  if isinstance(a.device, tuple): batch //= len(a.device)
+  # only sharding on the batch or K is tested, others might work too
+  if isinstance(a.device, tuple):
+    if a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
+    elif a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None: batch //= len(a.device)
+    else: return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
+    dname = a.device[0]
+  else: dname = a.device
+  arch = getattr(Device[dname].renderer, "arch", "")
   if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
-  if (key:=(M, N, K)) not in GEMM_ARGS: return todo(f"GEMM shape not supported {key}")
+  if (key:=(M, N, K)) not in GEMM_ARGS and arch == "gfx950": return todo(f"GEMM shape not supported {key} on {arch}")
   return True
 
 # ** UOp gemm to test Tensor.custom_kernel multi and backward correctness on non cdna4
@@ -58,14 +62,13 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
 # ** backward gemm, might use the asm gemm
 
 def custom_gemm_bw(gradient:UOp, kernel:UOp):
-  out, a, b = kernel.src
+  out, a, b = kernel.src[1:]
   assert all_same([gradient.device, a.device, b.device, out.device])
   a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
+  # TODO: this needs to be cleaned up and done properly, the batch dim of grad and a multi need to align
+  g_t = g_t[:a.shape[0]]
   grad_a = (g_t @ b_t.T).uop
-  a_T = a_t.transpose(-2, -1)
-  a_T = a_T.reshape(*a_T.shape[:-1], 1, a_T.shape[-1])
-  g_r = g_t.reshape(*g_t.shape[:-2], 1, *g_t.shape[-2:]).transpose(-1, -2)
-  grad_b = (a_T * g_r).sum((-1, 0)).uop
+  grad_b = (a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1) @ g_t.reshape(-1, g_t.shape[-1])).uop
   return (None, grad_a, grad_b)
 
 # ** main gemm function
@@ -79,17 +82,19 @@ def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
   batch, M, K = a.shape
   N = b.shape[1]
   is_multi = isinstance(a.device, tuple)
+  if (k_sharded:=is_multi and a.uop.axis == 2): K //= len(a.device)
 
   if is_multi:
-    out = Tensor(Tensor.empty(batch//len(a.device), M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
+    out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
   else:
     out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
 
   dname = a.device[0] if is_multi else a.device
-  arch = getattr(Device[dname].renderer, "arch", None)
+  arch = getattr(Device[dname].renderer, "arch", "")
   if arch.startswith("gfx950") and getenv("USE_ASM", 1):
     numWG = GEMM_ARGS[(M, N, K)][0]
     out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname, wg=numWG, arch=arch), grad_fxn=custom_gemm_bw)[0]
   else:
     out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
+  if k_sharded: out = out.sum(0)
   return out.squeeze(0) if squeeze else out

extra/gemm/asm/cdna/test_asm_gemm.py

@@ -1,46 +0,0 @@
-import unittest
-from tinygrad import Tensor, Device, dtypes, Context
-from tinygrad.helpers import getenv
-from extra.gemm.asm.cdna.gemm import asm_gemm
-
-def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.bfloat16, multi=False) -> None:
-  Tensor.manual_seed(0)
-  a_rand = Tensor.randn((batch, M, K), dtype=dtypes.float).sub(0.5).cast(dtype)
-  b_rand = Tensor.randn((K, N), dtype=dtypes.float).sub(0.5).cast(dtype)
-  with Context(DEBUG=0):
-    Tensor.realize(a_rand, b_rand)
-
-  devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(8)) if multi else None
-
-  a, b = Tensor(a_rand.numpy(), requires_grad=True).cast(dtype), Tensor(b_rand.numpy(), requires_grad=True).cast(dtype)
-  if multi: a, b = a.shard(devs, axis=0), b.shard(devs, axis=None)
-  tst = asm_gemm(a, b)
-  tst.sum().backward()
-  Tensor.realize(tst, a.grad, b.grad)
-
-  a_ref, b_ref = Tensor(a_rand.numpy(), requires_grad=True).cast(dtype), Tensor(b_rand.numpy(), requires_grad=True).cast(dtype)
-  if multi: a_ref, b_ref = a_ref.shard(devs, axis=0), b_ref.shard(devs, axis=None)
-  with Context(ASM_GEMM=0): ref = a_ref @ b_ref
-  ref.sum().backward()
-  Tensor.realize(ref, a_ref.grad, b_ref.grad)
-
-  with Context(DEBUG=0):
-    assert (tst - ref).square().max().float().item() < 1e-6, "forward mismatch"
-    assert (a.grad - a_ref.grad).square().max().float().item() < 1e-3, "grad_a mismatch"
-    assert (b.grad - b_ref.grad).square().max().float().item() < 1e-3, "grad_b mismatch"
-
-class TestGemm(unittest.TestCase):
-  def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 4096), N, N, dtype=dtypes.half)
-
-  def test_gemm1(self): verify_asm_gemm(8, 8192, 4096, 14336, multi=True)
-  def test_gemm2(self): verify_asm_gemm(8, 8192, 128256, 4096, multi=True)
-  def test_gemm3(self): verify_asm_gemm(8, 8192, 14336, 4096, multi=True)
-  def test_gemm4(self): verify_asm_gemm(8, 4096, 14336, 4096, multi=True)
-  def test_gemm5(self): verify_asm_gemm(8, 4096, 4096, 14336, multi=True)
-  def test_gemm6(self): verify_asm_gemm(16, 4096, 4096, 14336, multi=True)
-  def test_gemm_unsupported(self):
-    with self.assertRaisesRegex(AssertionError, "shape not supported"):
-      verify_asm_gemm(8, 8192, 1024, 4096, multi=True)
-
-if __name__ == "__main__":
-  unittest.main()

extra/gemm/intel_xmx.py

@@ -37,7 +37,7 @@ b.copyin(row.data)
 c.copyin(mat.data)
 ret = prog(a._buf, b._buf, c._buf, global_size=[1,1,1], local_size=[8,1,1], wait=True)
 print(ret)
-out = np.frombuffer(a.as_buffer(), np.float32)
+out = np.frombuffer(a.as_memoryview(), np.float32)
 real = row.astype(np.float32)@mat.T.astype(np.float32)
 print("out:", out)
 print("real", real)

extra/gemm/triton_nv_matmul.py

@@ -98,10 +98,10 @@ if __name__ == "__main__":
   # check correctness
   if getenv("VERIFY"):
     from tinygrad.engine.realize import run_schedule
-    triton_buf = np.frombuffer(si.bufs[0].as_buffer(), np.float16).reshape(M,N)
+    triton_buf = np.frombuffer(si.bufs[0].as_memoryview(), np.float16).reshape(M,N)
     print(triton_buf)
     run_schedule(sched)
-    tinygrad_buf = np.frombuffer(si.bufs[0].as_buffer(), np.float16).reshape(M,N)
+    tinygrad_buf = np.frombuffer(si.bufs[0].as_memoryview(), np.float16).reshape(M,N)
     print(tinygrad_buf)
     np.testing.assert_allclose(triton_buf, tinygrad_buf)
     print("correct!")

extra/hip_large_kernel.py

@@ -18,7 +18,7 @@ prg = dev.runtime("write_ones", mbin)
 prg(buf0._buf, global_size=(1,65537,1), local_size=(1,1,1), wait=True)
 
 import numpy as np
-def to_np(buf): return np.frombuffer(buf.as_buffer().cast(buf.dtype.base.fmt), dtype=_to_np_dtype(buf.dtype.base))
+def to_np(buf): return np.frombuffer(buf.as_memoryview().cast(buf.dtype.base.fmt), dtype=_to_np_dtype(buf.dtype.base))
 
 big = to_np(buf0)
 print(big)

extra/huggingface_onnx/huggingface_manager.py

@@ -8,14 +8,14 @@ from tinygrad.helpers import _ensure_downloads_dir
 DOWNLOADS_DIR = _ensure_downloads_dir() / "models"
 from tinygrad.helpers import tqdm
 
-def snapshot_download_with_retry(*, repo_id: str, allow_patterns: list[str]|tuple[str, ...]|None=None, cache_dir: str|Path|None=None,
+def snapshot_download_with_retry(*, repo_id: str, allow_patterns: list[str]|tuple[str, ...]|None=None, local_dir: str|Path|None=None,
                                  tries: int=2, **kwargs) -> Path:
   for attempt in range(tries):
     try:
       return Path(snapshot_download(
         repo_id=repo_id,
         allow_patterns=allow_patterns,
-        cache_dir=str(cache_dir) if cache_dir is not None else None,
+        local_dir=str(local_dir) if local_dir is not None else None,
         **kwargs
       ))
     except Exception as e:
@@ -144,14 +144,14 @@ class HuggingFaceONNXManager:
         root_path = snapshot_download_with_retry(
           repo_id=model_id,
           allow_patterns=allow_patterns,
-          cache_dir=str(self.models_dir)
+          local_dir=str(self.models_dir / model_id)
         )
 
         # Download config files (usually small)
         snapshot_download_with_retry(
           repo_id=model_id,
           allow_patterns=["*config.json"],
-          cache_dir=str(self.models_dir)
+          local_dir=str(self.models_dir / model_id)
         )
 
         model_data["download_path"] = str(root_path)

extra/huggingface_onnx/run_models.py

@@ -88,8 +88,8 @@ if __name__ == "__main__":
       # repo id
       # validates all onnx models inside repo
       repo_id = "/".join(path)
-      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], cache_dir=DOWNLOADS_DIR)
-      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
+      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], local_dir=DOWNLOADS_DIR / repo_id)
+      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], local_dir=DOWNLOADS_DIR / repo_id)
       config = get_config(root_path)
       for onnx_model in root_path.rglob("*.onnx"):
         rtol, atol = get_tolerances(onnx_model.name)
@@ -101,8 +101,8 @@ if __name__ == "__main__":
       onnx_model = path[-1]
       assert path[-1].endswith(".onnx")
       repo_id, relative_path = "/".join(path[:2]), "/".join(path[2:])
-      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=[relative_path], cache_dir=DOWNLOADS_DIR)
-      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
+      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=[relative_path], local_dir=DOWNLOADS_DIR / repo_id)
+      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], local_dir=DOWNLOADS_DIR / repo_id)
       config = get_config(root_path)
       rtol, atol = get_tolerances(onnx_model)
       print(f"validating {relative_path} with truncate={args.truncate}, {rtol=}, {atol=}")

extra/models/clip.py

@@ -5,7 +5,6 @@ from tinygrad.nn import Linear, LayerNorm, Embedding, Conv2d
 from typing import List, Optional, Union, Tuple, Dict
 from abc import ABC, abstractmethod
 from functools import lru_cache
-from PIL import Image
 import numpy as np
 import re, gzip
 
@@ -444,7 +443,8 @@ class OpenClipEncoder:
   # TODO:
   # Should be doable in pure tinygrad, would just require some work and verification.
   # This is very desirable since it would allow for full generation->evaluation in a single JIT call.
-  def prepare_image(self, image:Image.Image) -> Tensor:
+  def prepare_image(self, image) -> Tensor:
+    from PIL import Image
     SIZE = 224
     w, h = image.size
     scale = min(SIZE / h, SIZE / w)

extra/models/llama.py

@@ -55,12 +55,15 @@ class Attention:
       xqkv = x @ self.wqkv.T
       xq, xk, xv = xqkv.split([self.wq.weight.shape[0], self.wk.weight.shape[0], self.wv.weight.shape[0]], dim=2)
     else:
-      xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+      xq, xk, xv = self.wq(x), self.wk(x.contiguous_backward()), self.wv(x)
 
     if self.q_norm is not None and self.k_norm is not None:
       xq = self.q_norm(xq)
       xk = self.k_norm(xk)
 
+    # cast_float_to_bf16 is expensive in reduction loops, break it out
+    if x.dtype == dtypes.bfloat16: xq, xk = xq.contiguous_backward(), xk.contiguous_backward()
+
     xq = xq.reshape(xq.shape[0], xq.shape[1], self.n_heads, self.head_dim)
     xk = xk.reshape(xk.shape[0], xk.shape[1], self.n_kv_heads, self.head_dim)
     xv = xv.reshape(xv.shape[0], xv.shape[1], self.n_kv_heads, self.head_dim)
@@ -86,20 +89,23 @@ class Attention:
       assert start_pos == 0
       keys, values = xk, xv
 
-    keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
-    xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
-    attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2)
+    if self.max_context:
+      keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
+      xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
+      attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2)
+    else:
+      xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
+      attn = xq.scaled_dot_product_attention(keys, values, is_causal=True, enable_gqa=True).transpose(1, 2)
     if getenv("STUB_ATTENTION"):
-      # TODO: do we need mask?
       from tinygrad.uop.ops import UOp, KernelInfo
       def fa_custom_forward(attn:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
         return UOp.sink(arg=KernelInfo(name="fa_custom_forward"))
       def fa_custom_backward(out_q:UOp, out_k:UOp, out_v:UOp, grad:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
         return UOp.sink(arg=KernelInfo(name="fa_custom_backward"))
       def fa_backward(grad:UOp, kernel:UOp) -> tuple[None, UOp, UOp, UOp]:
-        grad_q = Tensor.empty_like(q:=Tensor(kernel.src[1]))
-        grad_k = Tensor.empty_like(k:=Tensor(kernel.src[2]))
-        grad_v = Tensor.empty_like(v:=Tensor(kernel.src[3]))
+        grad_q = Tensor.empty_like(q:=Tensor(kernel.src[2]))
+        grad_k = Tensor.empty_like(k:=Tensor(kernel.src[3]))
+        grad_v = Tensor.empty_like(v:=Tensor(kernel.src[4]))
         ck = Tensor.custom_kernel(grad_q, grad_k, grad_v, Tensor(grad), q, k, v, fxn=fa_custom_backward)[:3]
         return (None, ck[0].uop, ck[1].uop, ck[2].uop)
       attn = Tensor.empty_like(attn).custom_kernel(xq, keys, values, fxn=fa_custom_forward, grad_fxn=fa_backward)[0]
@@ -197,7 +203,9 @@ class Transformer:
     h = self.tok_embeddings(tokens)
     freqs_cis = self.freqs_cis.cast(h.dtype)[:, 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) if seqlen > 1 else None
+    if self.max_context != 0 and seqlen > 1:
+      mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1)
+    else: mask = None
     for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask)
     logits = self.output(self.norm(h))
     if math.isnan(temperature): return logits

extra/optimization/generate_dataset.sh

@@ -7,7 +7,7 @@ export CAPTURE_PROCESS_REPLAY=1
 rm "$LOGOPS" 2>/dev/null || true
 test/external/process_replay/reset.py
 
-CI=1 python3 -m pytest -n=auto test/test_ops.py test/test_nn.py test/unit/test_winograd.py test/models/test_real_world.py --durations=20
+CI=1 python3 -m pytest -n=auto test/test_ops.py test/test_nn.py test/unit/test_winograd.py test/null/test_real_world.py --durations=20
 CL=1 python3 -m pytest test/test_tiny.py
 
 # extract, sort and uniq

extra/optimization/get_action_space.py

@@ -1,31 +0,0 @@
-import random
-from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from tinygrad.codegen.opt.search import actions
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
-from tinygrad.helpers import tqdm
-
-tactions = set()
-def test_rebuild(lin):
-  linr = Kernel(lin.ast)
-  for o in lin.applied_opts:
-    assert o in actions, f"{o} is not in actions"
-    tactions.add(o)
-    linr.apply_opt(o)
-
-  assert len(lin.sts) == len(linr.sts)
-  for st1,st2 in zip(lin.sts, linr.sts):
-    assert st1 == st2, f"{st1} != {st2}"
-
-if __name__ == "__main__":
-  ast_strs = load_worlds(False, False, False)
-  random.shuffle(ast_strs)
-  ast_strs = ast_strs[:2000]
-  for ast_str in tqdm(ast_strs):
-    lin = ast_str_to_lin(ast_str)
-    #if not lin.apply_tensor_cores():
-    lin.apply_opts(hand_coded_optimizations(lin))
-    test_rebuild(lin)
-
-  print(len(tactions), len(actions))
-  print(sorted(list(tactions)))

extra/optimization/rl.py

@@ -1,76 +0,0 @@
-import os
-import numpy as np
-import math, random
-from tinygrad.tensor import Tensor
-from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-from tinygrad.codegen.opt.search import actions, bufs_from_lin, get_kernel_actions
-from tinygrad.nn.optim import Adam
-from extra.optimization.extract_policynet import PolicyNet
-from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, time_linearizer
-
-if __name__ == "__main__":
-  net = PolicyNet()
-  if os.path.isfile("/tmp/policynet.safetensors"): load_state_dict(net, safe_load("/tmp/policynet.safetensors"))
-  optim = Adam(get_parameters(net))
-
-  ast_strs = load_worlds()
-
-  # select a world
-  all_feats, all_acts, all_rews = [], [], []
-  while 1:
-    Tensor.training = False
-    lin = ast_str_to_lin(random.choice(ast_strs))
-    rawbufs = bufs_from_lin(lin)
-    tm = last_tm = base_tm = time_linearizer(lin, rawbufs)
-
-    # take actions
-    feats, acts, rews = [], [], []
-    while 1:
-      feat = lin_to_feats(lin)
-      feats.append(feat)
-      probs = net(Tensor([feat])).exp()[0].numpy()
-
-      # mask valid actions
-      valid_action_mask = np.zeros((len(actions)+1), dtype=np.float32)
-      for x in get_kernel_actions(lin): valid_action_mask[x] = 1
-      probs *= valid_action_mask
-      probs /= sum(probs)
-
-      act = np.random.choice(len(probs), p=probs)
-      acts.append(act)
-      if act == 0:
-        rews.append(0)
-        break
-      try:
-        lin.apply_opt(actions[act-1])
-        tm = time_linearizer(lin, rawbufs)
-        if math.isinf(tm): raise Exception("failed")
-        rews.append(((last_tm-tm)/base_tm))
-        last_tm = tm
-      except Exception:
-        rews.append(-0.5)
-        break
-      #print(f"{tm*1e6:10.2f}", lin.colored_shape())
-
-    assert len(feats) == len(acts) and len(acts) == len(rews)
-    #print(rews)
-    print(f"***** EPISODE {len(rews)} steps, {sum(rews):5.2f} reward, {base_tm*1e6:12.2f} -> {tm*1e6:12.2f} : {lin.colored_shape()}")
-    all_feats += feats
-    all_acts += acts
-    # rewards to go
-    for i in range(len(rews)-2, -1, -1): rews[i] += rews[i+1]
-    all_rews += rews
-
-    BS = 32
-    if len(all_feats) >= BS:
-      Tensor.training = True
-      x = Tensor(all_feats[:BS])
-      mask = np.zeros((BS, len(actions)+1), dtype=np.float32)
-      mask[range(BS), all_acts[:BS]] = all_rews[:BS]
-      loss = -(net(x) * Tensor(mask)).mean()
-      optim.zero_grad()
-      loss.backward()
-      optim.step()
-      all_feats = all_feats[BS:]
-      all_acts = all_acts[BS:]
-      all_rews = all_rews[BS:]

extra/optimization/run_qnet.py

@@ -1,32 +0,0 @@
-from typing import List, Tuple
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt.search import get_kernel_actions, actions
-
-_net = None
-def beam_q_estimate(beam:List[Tuple[Kernel, float]]) -> List[Tuple[Kernel, float]]:
-  global _net
-  if _net is None:
-    from tinygrad.nn.state import load_state_dict, safe_load
-    from extra.optimization.pretrain_valuenet import ValueNet
-    _net = ValueNet(1021+len(actions), 2)
-    load_state_dict(_net, safe_load("/tmp/qnet.safetensors"), verbose=False)
-  from tinygrad.tensor import Tensor
-  from tinygrad.helpers import Context
-  from extra.optimization.helpers import lin_to_feats
-  import numpy as np
-  feats = []
-  lins = []
-  base_tms = []
-  for lin,tm in beam:
-    lin_feats = lin_to_feats(lin)
-    for a,v in get_kernel_actions(lin, include_0=False).items():
-      acts = np.zeros(len(actions))
-      acts[a-1] = 1.0
-      feats.append(np.concatenate([lin_feats, acts]))
-      lins.append(v)
-      base_tms.append(tm)
-  with Context(BEAM=0):
-    with Tensor.train(False):
-      preds = _net(Tensor(feats)).numpy()
-  pred_time = np.array(base_tms) / np.exp(preds[:, 0])
-  return sorted(zip(lins, pred_time), key=lambda x: x[1])

extra/optimization/search.py

@@ -1,34 +0,0 @@
-import argparse
-from extra.optimization.helpers import ast_str_to_lin, time_linearizer
-
-from tinygrad import dtypes
-from tinygrad.helpers import BEAM, getenv
-from tinygrad.device import Device, Compiled
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
-
-
-if __name__ == '__main__':
-  parser = argparse.ArgumentParser(description="Run a search for the optimal opts for a kernel", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-  parser.add_argument("--ast", type=str, default=None, help="the ast for the kernel to be optimized")
-  parser.add_argument("--file", type=str, default=None, help="a file containing asts to be optimized, one per line")
-  args = parser.parse_args()
-
-  device: Compiled = Device[Device.DEFAULT]
-  print(f"optimizing for {Device.DEFAULT}")
-
-  if args.ast is not None:
-    ast_strs = [args.ast]
-  elif args.file is not None:
-    with open(args.file, 'r') as file:
-      ast_strs = file.readlines()
-
-  for i, ast_str in enumerate(ast_strs):
-    print(f"optimizing {i}/{len(ast_strs)}\nast={ast_str}")
-    lin = ast_str_to_lin(ast_str, opts=device.renderer)
-    rawbufs = bufs_from_lin(lin)
-    lin = beam_search(lin, rawbufs, getenv("BEAM", 8), bool(getenv("BEAM_ESTIMATE", 1)))
-
-    tm = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10)
-    print(f"final time {tm*1e6:9.0f} us: {lin.colored_shape()}")
-    print(lin.applied_opts)

extra/optimization/test_helpers.py

@@ -1,19 +0,0 @@
-import unittest
-
-from extra.optimization.helpers import load_worlds
-
-class TestKernelDataset(unittest.TestCase):
-  def test_load_worlds_filters(self):
-    all_kernels = load_worlds(filter_reduce=False, filter_noimage=False, filter_novariable=False)
-
-    reduce_kernels = load_worlds(filter_reduce=True, filter_noimage=False, filter_novariable=False)
-    self.assertGreater(len(all_kernels), len(reduce_kernels))
-
-    image_kernels = load_worlds(filter_reduce=False, filter_noimage=True, filter_novariable=False)
-    self.assertGreater(len(all_kernels), len(image_kernels))
-
-    variable_kernels = load_worlds(filter_reduce=False, filter_noimage=False, filter_novariable=True)
-    self.assertGreater(len(all_kernels), len(variable_kernels))
-
-if __name__ == '__main__':
-  unittest.main()

extra/optimization/test_net.py

@@ -1,67 +0,0 @@
-import numpy as np
-import math
-import random
-np.set_printoptions(suppress=True)
-from copy import deepcopy
-from tinygrad.helpers import getenv, colored
-from tinygrad.tensor import Tensor
-from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-from tinygrad.codegen.opt.search import bufs_from_lin, actions, get_kernel_actions
-from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
-from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, time_linearizer
-from extra.optimization.extract_policynet import PolicyNet
-from extra.optimization.pretrain_valuenet import ValueNet
-
-VALUE = getenv("VALUE")
-
-if __name__ == "__main__":
-  if VALUE:
-    net = ValueNet()
-    load_state_dict(net, safe_load("/tmp/valuenet.safetensors"))
-  else:
-    net = PolicyNet()
-    load_state_dict(net, safe_load("/tmp/policynet.safetensors"))
-
-  ast_strs = load_worlds()
-
-  # real randomness
-  random.seed()
-  random.shuffle(ast_strs)
-
-  wins = 0
-  for ep_num,ast_str in enumerate(ast_strs):
-    print("\nEPISODE", ep_num, f"win {wins*100/max(1,ep_num):.2f}%")
-    lin = ast_str_to_lin(ast_str)
-    rawbufs = bufs_from_lin(lin)
-
-    linhc = deepcopy(lin)
-    linhc.applied_opts(hand_coded_optimizations(linhc))
-    tmhc = time_linearizer(linhc, rawbufs)
-    print(f"{tmhc*1e6:10.2f}     HC    ", linhc.colored_shape())
-
-    pred_time = float('nan')
-    tm = float('inf')
-    while 1:
-      if VALUE:
-        acts,feats = [], []
-        for k,v in get_kernel_actions(lin).items():
-          acts.append(k)
-          feats.append(lin_to_feats(v))
-        preds = net(Tensor(feats))
-        pred_time = math.exp(preds.numpy().min())
-        act = acts[preds.numpy().argmin()]
-      else:
-        probs = net(Tensor([lin_to_feats(lin)]))
-        dist = probs.exp().numpy()
-        act = dist.argmax()
-      if act == 0: break
-      try:
-        lin.apply_opt(actions[act-1])
-      except Exception:
-        print("FAILED")
-        break
-      tm = time_linearizer(lin, rawbufs)
-      print(f"{tm*1e6:10.2f} {pred_time*1e6:10.2f}", lin.colored_shape())
-
-    print(f"{colored('BEAT', 'green') if tm < tmhc else colored('lost', 'red')} hand coded {tmhc/tm:5.2f}x")
-    wins += int(tm < tmhc)

extra/optimization/test_time_linearizer.py

@@ -1,21 +0,0 @@
-from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer
-from tinygrad.codegen.opt.search import bufs_from_lin, get_kernel_actions
-
-if __name__ == "__main__":
-  ast_strs = load_worlds()
-  for i, ast_str in enumerate(ast_strs):
-    lin = ast_str_to_lin(ast_str)
-    rawbufs = bufs_from_lin(lin)
-    test_tm = time_linearizer(lin, rawbufs)
-    if test_tm < 1e-2: continue
-    print(f"EXAMPLE {i}")
-    acted_lins = get_kernel_actions(lin)
-    ok_avg, short_avg = 0, 0
-    for k,v in acted_lins.items():
-      tm1 = time_linearizer(v, rawbufs)
-      tm2 = time_linearizer(v, rawbufs)
-      tm3 = time_linearizer(v, rawbufs, False)
-      print(v.colored_shape(50), f"{tm1*1e3:10.2f} {tm2*1e3:10.2f} {tm3*1e3:10.2f}  :  {((tm1-tm2)/tm1)*100:5.2f}%  vs  {((tm1-tm3)/tm1)*100:5.2f}%")
-      ok_avg += (tm1-tm2)/tm1
-      short_avg += (tm1-tm3)/tm1
-    print(f"{ok_avg/len(acted_lins)*100:5.2f}% vs {short_avg/len(acted_lins)*100:5.2f}%")

extra/perfetto/to_perfetto.py

@@ -2,27 +2,26 @@ import sys, pickle, decimal, json
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent
 from tinygrad.helpers import tqdm, temp, ProfileEvent, ProfileRangeEvent, TracingKey
 
-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)}
+devices:dict[str, tuple[decimal.Decimal, int]] = {}
+def prep_ts(device:str, ts:decimal.Decimal): return int(decimal.Decimal(ts) + devices[device][0])
+def dev_to_pid(device:str): return {"pid": devices[device][1], "tid": 0}
 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))
+  devices[ev.device] = (ev.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"}}]
+          {"name": "thread_name", "ph": "M", "pid": dev_to_pid(ev.device)['pid'], "tid": 0, "args": {"name": ev.device}}]
 def range_ev_to_perfetto_json(ev:ProfileRangeEvent):
   name = ev.name.display_name if isinstance(ev.name, TracingKey) else ev.name
-  return [{"name": 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)}]
+  return [{"name": name, "ph": "X", "ts": prep_ts(ev.device, ev.st), "dur": float(ev.en-ev.st), **dev_to_pid(ev.device)}]
 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]
     name = e.name.display_name if isinstance(e.name, TracingKey) else e.name
-    ret += [{"name": name, "ph": "X", "ts": prep_ts(e.device, st, e.is_copy), "dur": float(en-st), **dev_to_pid(e.device, e.is_copy)}]
+    ret += [{"name": name, "ph": "X", "ts": prep_ts(e.device, st), "dur": float(en-st), **dev_to_pid(e.device)}]
     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"}]
+      ret += [{"ph": "s", **dev_to_pid(d.device), "id": reccnt+len(ret), "ts": prep_ts(d.device, ev.sigs[d.en_id]), "bp": "e"}]
+      ret += [{"ph": "f", **dev_to_pid(e.device), "id": reccnt+len(ret)-1, "ts": prep_ts(e.device, st), "bp": "e"}]
   return ret
 def to_perfetto(profile:list[ProfileEvent]):
   # Start json with devices.

extra/sqtt/rgptool.py

@@ -145,7 +145,8 @@ class RGP:
   @staticmethod
   def from_profile(profile_pickled, device:str|None=None):
     profile: list[ProfileEvent] = pickle.loads(profile_pickled)
-    device_events = {x.device:x for x in profile if isinstance(x, ProfileDeviceEvent) and x.device.startswith('AMD')}
+    def _is_base_dev(d): return all(p.isdigit() for p in d.split(":")[1:])
+    device_events = {x.device:x for x in profile if isinstance(x, ProfileDeviceEvent) and x.device.startswith('AMD') and _is_base_dev(x.device)}
     if device is None:
       if len(device_events) == 0: raise RuntimeError('No supported devices found in profile')
       if len(device_events) > 1: raise RuntimeError(f"More than one supported device found, select which one to export: {', '.join(device_events.keys())}")

extra/thunder/tiny/fa.py

@@ -2,7 +2,7 @@ import math
 
 from tinygrad import Tensor, dtypes
 from tinygrad.helpers import DEBUG
-from tinygrad.uop.ops import UOp
+from tinygrad.uop.ops import UOp, Ops
 
 from extra.thunder.tiny.tk import WARP_THREADS
 from extra.thunder.tiny.tk.kernel import Kernel
@@ -43,20 +43,18 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
   B_local = B // num_devices
   if DEBUG >= 2: print(f"Flash Attention {B=} {B_local=} {N=} {H=} {D=} {H_KV=} {GROUP_SIZE=}")
 
-  def custom_forward(ou:UOp, l_vecu:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp) -> UOp:
+  def _custom_forward_impl(ou:UOp, l_vecu:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp|None) -> UOp:
     with Kernel("fa_custom_forward", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B_local), NUM_WORKERS * WARP_THREADS) as ker:
       warp = ker.warp
 
-      o, q, k, v, mask, l_vec = GL(ou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker), GL(l_vecu, ker)
+      o, q, k, v, l_vec = GL(ou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(l_vecu, ker)
+      mask = GL(masku, ker) if masku is not None else None
 
       head = ker.blockIdx_x
       head_kv = head // GROUP_SIZE
       batch = ker.blockIdx_z
       q_seq = ker.blockIdx_y * NUM_WORKERS + ker.warpid
 
-      k_smem = ker.st((KV_BLOCK_SIZE, D), dtypes.bfloat16)
-      v_smem = ker.st((KV_BLOCK_SIZE, D), dtypes.bfloat16)
-
       q_reg_fl = ker.rt((Q_BLOCK_SIZE, D), dtypes.float32)
       q_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
       q_reg_transposed = ker.rt((D, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
@@ -86,12 +84,10 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
       q_reg = warp.copy(q_reg, q_reg_fl)
       q_reg_transposed = warp.transpose(q_reg_transposed, q_reg)
 
-      for kv_idx in ker.range(N // KV_BLOCK_SIZE):
-        k_smem = warp.load(k_smem, k, (), (batch, kv_idx, head_kv, 0), axis=1)
-        v_smem = warp.load(v_smem, v, (), (batch, kv_idx, head_kv, 0), axis=1)
-
-        k_reg = warp.load(k_reg, k_smem)
-        v_reg = warp.load(v_reg, v_smem)
+      num_kv_blocks = (q_seq + 1) if is_causal else (N // KV_BLOCK_SIZE)
+      for kv_idx in ker.range(num_kv_blocks):
+        k_reg = warp.load(k_reg, k, (), (batch, kv_idx, head_kv, 0), axis=1)
+        v_reg = warp.load(v_reg, v, (), (batch, kv_idx, head_kv, 0), axis=1)
 
         # mma qk^t
         att_block = warp.zero(att_block.after(kv_idx))
@@ -99,9 +95,16 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
         att_block = warp.mma_AtB(att_block, k_reg_transposed, q_reg_transposed)
 
         # apply attention mask
-        mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_seq, kv_idx), axis=2)
-        mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
-        att_block += mask_reg_transposed
+        if is_causal:
+          bs_rows, bs_cols, bs_stride = att_block.base_shape.rows, att_block.base_shape.cols, att_block.base_shape.stride
+          q_base = q_seq * Q_BLOCK_SIZE + (warp.laneid % bs_cols)
+          kv_base = kv_idx * KV_BLOCK_SIZE + (warp.laneid // bs_cols) * bs_stride
+          att_block = warp.map(att_block,
+            lambda x, idx: ((kv_base + idx[0]*bs_rows + idx[2]) > (q_base + idx[1]*bs_cols)).alu(Ops.WHERE, UOp.ufix(x._uop, -math.inf), x))
+        elif mask is not None:
+          mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_seq, kv_idx), axis=2)
+          mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
+          att_block += mask_reg_transposed
 
         # softmax
         max_vec_last = warp.copy(max_vec_last.after(kv_idx), max_vec)
@@ -141,11 +144,18 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
 
       return ker.finish()
 
-  def custom_backward_q(dqu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
+  def custom_forward_causal(ou:UOp, l_vecu:UOp, qu:UOp, ku:UOp, vu:UOp) -> UOp:
+    return _custom_forward_impl(ou, l_vecu, qu, ku, vu, None)
+
+  def custom_forward_masked(ou:UOp, l_vecu:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp) -> UOp:
+    return _custom_forward_impl(ou, l_vecu, qu, ku, vu, masku)
+
+  def _custom_backward_q_impl(dqu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp|None, l_vecu:UOp, delta_vecu:UOp) -> UOp:
     with Kernel("fa_custom_backward_q", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B_local), NUM_WORKERS * WARP_THREADS) as ker:
       warp = ker.warp
 
-      dq, do, q, k, v, mask = GL(dqu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker)
+      dq, do, q, k, v = GL(dqu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker)
+      mask = GL(masku, ker) if masku is not None else None
       l_vec, delta_vec = GL(l_vecu, ker), GL(delta_vecu, ker)
 
       head = ker.blockIdx_x
@@ -153,9 +163,6 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
       batch = ker.blockIdx_z
       q_seq = ker.blockIdx_y * NUM_WORKERS + ker.warpid
 
-      k_smem = ker.st((KV_BLOCK_SIZE, D), dtypes.bfloat16)
-      v_smem = ker.st((KV_BLOCK_SIZE, D), dtypes.bfloat16)
-
       q_reg_fl = ker.rt((Q_BLOCK_SIZE, D), dtypes.float32)
       q_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
       q_reg_t = ker.rt((D, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
@@ -194,24 +201,30 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
       l_vec_reg *= 1.0 / math.log(2)
       delta_vec_reg = warp.load(delta_vec_reg, delta_vec, (), (batch, head, 0, q_seq), axis=2)
 
-      for kv_idx in ker.range(N // KV_BLOCK_SIZE):
-        k_smem = warp.load(k_smem, k, (), (batch, kv_idx, head_kv, 0), axis=1)
-        v_smem = warp.load(v_smem, v, (), (batch, kv_idx, head_kv, 0), axis=1)
+      num_kv_blocks = (q_seq + 1) if is_causal else (N // KV_BLOCK_SIZE)
+      for kv_idx in ker.range(num_kv_blocks):
+        k_reg = warp.load(k_reg, k, (), (batch, kv_idx, head_kv, 0), axis=1)
+        k_reg_col = warp.load(k_reg_col, k, (), (batch, kv_idx, head_kv, 0), axis=1)
+        v_reg = warp.load(v_reg, v, (), (batch, kv_idx, head_kv, 0), axis=1)
 
-        k_reg = warp.load(k_reg, k_smem)
         k_reg_t = warp.transpose(k_reg_t, k_reg)
-        k_reg_col = warp.load(k_reg_col, k_smem)
         k_reg_col_t = warp.transpose(k_reg_col_t, k_reg_col)
-        v_reg = warp.load(v_reg, v_smem)
 
         # mma qk^t
         att_block = warp.zero(att_block.after(kv_idx))
         att_block = warp.mma_AtB(att_block, k_reg_t, q_reg_t)
 
         # apply attention mask
-        mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_seq, kv_idx), axis=2)
-        mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
-        att_block += mask_reg_transposed
+        if is_causal:
+          bs_rows, bs_cols, bs_stride = att_block.base_shape.rows, att_block.base_shape.cols, att_block.base_shape.stride
+          q_base = q_seq * Q_BLOCK_SIZE + (warp.laneid % bs_cols)
+          kv_base = kv_idx * KV_BLOCK_SIZE + (warp.laneid // bs_cols) * bs_stride
+          att_block = warp.map(att_block,
+            lambda x, idx: ((kv_base + idx[0]*bs_rows + idx[2]) > (q_base + idx[1]*bs_cols)).alu(Ops.WHERE, UOp.ufix(x._uop, -math.inf), x))
+        elif mask is not None:
+          mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_seq, kv_idx), axis=2)
+          mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
+          att_block += mask_reg_transposed
 
         att_block -= l_vec_reg
         att_block = att_block.exp2()
@@ -231,19 +244,24 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
 
       return ker.finish()
 
-  def custom_backward_kv(dku:UOp, dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp):
+  def custom_backward_q_causal(dqu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
+    return _custom_backward_q_impl(dqu, dou, qu, ku, vu, None, l_vecu, delta_vecu)
+
+  def custom_backward_q_masked(dqu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
+    return _custom_backward_q_impl(dqu, dou, qu, ku, vu, masku, l_vecu, delta_vecu)
+
+  def _custom_backward_kv_impl(dku:UOp, dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp|None, l_vecu:UOp, delta_vecu:UOp):
     with Kernel("fa_custom_backward_kv", (H_KV, N // (KV_BLOCK_SIZE*NUM_WORKERS), B_local), NUM_WORKERS * WARP_THREADS) as ker:
       warp = ker.warp
 
-      dk, dv, do, q, k, v, mask = GL(dku, ker), GL(dvu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker)
+      dk, dv, do, q, k, v = GL(dku, ker), GL(dvu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker)
+      mask = GL(masku, ker) if masku is not None else None
       l_vec, delta_vec = GL(l_vecu, ker), GL(delta_vecu, ker)
 
       head_kv = ker.blockIdx_x
       batch = ker.blockIdx_z
       kv_seq = ker.blockIdx_y * NUM_WORKERS + ker.warpid
 
-      q_smem = ker.st((Q_BLOCK_SIZE, D), dtypes.bfloat16)
-      do_smem = ker.st((Q_BLOCK_SIZE, D), dtypes.bfloat16)
       att_smem = ker.st((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.bfloat16)
 
       q_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
@@ -277,19 +295,17 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
       k_reg_t = warp.transpose(k_reg_t, k_reg)
       v_reg = warp.load(v_reg, v, (), (batch, kv_seq, head_kv, 0), axis=1)
 
-      for q_idx in ker.range(N // Q_BLOCK_SIZE):
+      q_start = kv_seq if is_causal else 0
+      for q_idx in ker.range(q_start, N // Q_BLOCK_SIZE):
         for g in ker.range(GROUP_SIZE):
           head_q = head_kv * GROUP_SIZE + g
 
-          # load q and do
-          q_smem = warp.load(q_smem, q, (), (batch, q_idx, head_q, 0), axis=1)
-          do_smem = warp.load(do_smem, do, (), (batch, q_idx, head_q, 0), axis=1)
+          q_reg = warp.load(q_reg, q, (), (batch, q_idx, head_q, 0), axis=1)
+          q_reg_col = warp.load(q_reg_col, q, (), (batch, q_idx, head_q, 0), axis=1)
+          do_reg = warp.load(do_reg, do, (), (batch, q_idx, head_q, 0), axis=1)
+          do_reg_col = warp.load(do_reg_col, do, (), (batch, q_idx, head_q, 0), axis=1)
 
-          q_reg = warp.load(q_reg, q_smem)
           q_reg_t = warp.transpose(q_reg_t, q_reg)
-          q_reg_col = warp.load(q_reg_col, q_smem)
-          do_reg = warp.load(do_reg, do_smem)
-          do_reg_col = warp.load(do_reg_col, do_smem)
 
           # load l_vec and delta_vec
           l_vec_reg = warp.load(l_vec_reg, l_vec, (), (batch, head_q, 0, q_idx), axis=2)
@@ -302,9 +318,16 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
           att_block *= (1.0 / math.sqrt(D)) * (1.0 / math.log(2))
 
           # apply attention mask
-          mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_idx, kv_seq), axis=2)
-          mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
-          att_block += mask_reg_transposed
+          if is_causal:
+            bs_rows, bs_cols, bs_stride = att_block.base_shape.rows, att_block.base_shape.cols, att_block.base_shape.stride
+            q_base = q_idx * Q_BLOCK_SIZE + (warp.laneid % bs_cols)
+            kv_base = kv_seq * KV_BLOCK_SIZE + (warp.laneid // bs_cols) * bs_stride
+            att_block = warp.map(att_block,
+              lambda x, idx: ((kv_base + idx[0]*bs_rows + idx[2]) > (q_base + idx[1]*bs_cols)).alu(Ops.WHERE, UOp.ufix(x._uop, -math.inf), x))
+          elif mask is not None:
+            mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_idx, kv_seq), axis=2)
+            mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
+            att_block += mask_reg_transposed
 
           att_block -= l_vec_reg
           att_block = att_block.exp2()
@@ -336,24 +359,31 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
 
       return ker.finish(2)
 
+  def custom_backward_kv_causal(dku:UOp, dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, l_vecu:UOp, delta_vecu:UOp):
+    return _custom_backward_kv_impl(dku, dvu, dou, qu, ku, vu, None, l_vecu, delta_vecu)
+
+  def custom_backward_kv_masked(dku:UOp, dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp):
+    return _custom_backward_kv_impl(dku, dvu, dou, qu, ku, vu, masku, l_vecu, delta_vecu)
+
   single_device = xq.device[0] if isinstance(xq.device, tuple) else xq.device
 
   if is_causal:
     if attn_mask is not None: raise RuntimeError("cannot set attn_mask when is_causal=True")
-    attn_mask = Tensor.ones((B, 1, N, N), requires_grad=False, device=single_device, dtype=dtypes.bool).tril()
-  if attn_mask is not None:
+  elif attn_mask is not None:
     if attn_mask.dtype == dtypes.bool: attn_mask = attn_mask.where(0, -float("inf"))
+    if attn_mask.shape != (B, 1, N, N):
+      attn_mask = attn_mask.expand(B, 1, N, N)
+    if isinstance(xq.device, tuple) and not isinstance(attn_mask.device, tuple):
+      attn_mask = attn_mask.shard(xq.device, axis=0)
   else:
     attn_mask = Tensor.zeros((B, 1, N, N), requires_grad=False, device=single_device, dtype=dtypes.float32)
-  if attn_mask.shape != (B, 1, N, N):
-    attn_mask = attn_mask.expand(B, 1, N, N)
-  if isinstance(xq.device, tuple) and not isinstance(attn_mask.device, tuple):
-    attn_mask = attn_mask.shard(xq.device, axis=0)
+    if isinstance(xq.device, tuple):
+      attn_mask = attn_mask.shard(xq.device, axis=0)
 
   attn = _sharded_empty_like(xq, axis=0)
   l_vec = _sharded_empty((B, H, 1, N), xq, axis=0)
 
-  def grad(gradu:UOp, _) -> tuple[None, None, UOp, UOp, UOp, None]:
+  def grad_causal(gradu:UOp, _) -> tuple[None, None, UOp, UOp, UOp]:
     grad = Tensor(gradu, device=gradu.device)
     grad_q = _sharded_empty_like(xq, axis=0)
     grad_k = _sharded_empty_like(xk, axis=0)
@@ -361,11 +391,26 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
 
     delta_vec = (grad * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
 
-    grad_q = Tensor.custom_kernel(grad_q, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_q)[0]
-    grad_k, grad_v = Tensor.custom_kernel(grad_k, grad_v, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_kv)[:2]
+    grad_q = Tensor.custom_kernel(grad_q, grad, xq, xk, xv, l_vec, delta_vec, fxn=custom_backward_q_causal)[0]
+    grad_k, grad_v = Tensor.custom_kernel(grad_k, grad_v, grad, xq, xk, xv, l_vec, delta_vec, fxn=custom_backward_kv_causal)[:2]
+    return (None, None, grad_q.uop, grad_k.uop, grad_v.uop)
+
+  def grad_masked(gradu:UOp, _) -> tuple[None, None, UOp, UOp, UOp, None]:
+    grad = Tensor(gradu, device=gradu.device)
+    grad_q = _sharded_empty_like(xq, axis=0)
+    grad_k = _sharded_empty_like(xk, axis=0)
+    grad_v = _sharded_empty_like(xv, axis=0)
+
+    delta_vec = (grad * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
+
+    grad_q = Tensor.custom_kernel(grad_q, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_q_masked)[0]
+    grad_k, grad_v = Tensor.custom_kernel(grad_k, grad_v, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_kv_masked)[:2]
     return (None, None, grad_q.uop, grad_k.uop, grad_v.uop, None)
 
-  attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, attn_mask, fxn=custom_forward, grad_fxn=grad)[:2]
+  if is_causal:
+    attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, fxn=custom_forward_causal, grad_fxn=grad_causal)[:2]
+  else:
+    attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, attn_mask, fxn=custom_forward_masked, grad_fxn=grad_masked)[:2]
   attn_ = attn[:, :N_, :, :D_]
 
   return attn_.transpose(1, 2).cast(odtype)

extra/thunder/tiny/tk/kernel.py

@@ -12,7 +12,7 @@ class _tk_range:
   def __next__(self):
     if not self.done:
       self.done = True
-      self._rng = UOp.range(self.end // self.step, self.rid, axis_type=self.axis_type) * self.step + self.start
+      self._rng = UOp.range((self.end - self.start) // self.step, self.rid, axis_type=self.axis_type) * self.step + self.start
       return self._rng
     raise StopIteration
 

extra/tinyfs/fetch_file.py

@@ -1,11 +1,36 @@
 from tinygrad.tensor import Tensor
-import argparse
+import argparse, math, hashlib
+
+def _python_hash_1mb(data:bytes|bytearray):
+  chunks = [data[i:i+4096] for i in range(0, len(data), 4096)]
+  chunk_hashes = [hashlib.shake_128(chunk).digest(16) for chunk in chunks]
+  return hashlib.shake_128(b''.join(chunk_hashes)).digest(16)
+
+def hash_file(data: bytes|bytearray):
+  if len(data) % Tensor.CHUNK_SIZE != 0: data += bytes(Tensor.CHUNK_SIZE - len(data) % Tensor.CHUNK_SIZE)
+  base_chunks = math.ceil(len(data) / Tensor.CHUNK_SIZE)
+  tree_depth = math.ceil(math.log(base_chunks, Tensor.CHUNK_SIZE // 16))
+
+  for _ in range(tree_depth + 1):
+    data_chunks = [data[i:i+Tensor.CHUNK_SIZE] for i in range(0, len(data), Tensor.CHUNK_SIZE)]
+    data_chunk_hashes = [_python_hash_1mb(chunk) for chunk in data_chunks]
+    data = b''.join(data_chunk_hashes)
+    if len(data) % Tensor.CHUNK_SIZE != 0: data += bytes(Tensor.CHUNK_SIZE - len(data) % Tensor.CHUNK_SIZE)
+
+  return data[:16]
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
   parser.add_argument("--hash", type=str, required=True, help="file hash to fetch")
   parser.add_argument("--len", type=int, required=True, help="file length to fetch")
   parser.add_argument("--dest", type=str, required=True, help="destination path to save the file")
+  parser.add_argument("--check", action="store_true", help="verify the file hash after fetching")
   args = parser.parse_args()
 
   Tensor(bytes.fromhex(args.hash), device="CPU").fs_load(args.len).to(f"disk:{args.dest}").realize()
+
+  if args.check:
+    with open(args.dest, "rb") as f:
+      data = f.read()
+      assert hash_file(data) == bytes.fromhex(args.hash), "Hash mismatch after fetching file"
+    print("File hash verified successfully!")

extra/torch_backend/backend.py

@@ -29,6 +29,10 @@ def wrap(x: Tensor) -> torch.Tensor:
   x._strides = strides_for_shape(x.shape) # always recalculate
   if (not hasattr(x, '_storage_offset')) or (not x.uop.is_realized): x._storage_offset = calculate_storage_offset(x)
   return mod.wrap(x, _to_torch_dtype(x.dtype), _to_torch_device(x.device).index)
+def _update_torch_metadata(tensor: torch.Tensor, tiny: Tensor) -> None:
+  tiny._strides = strides_for_shape(tiny.shape)
+  tiny._storage_offset = calculate_storage_offset(tiny)
+  mod.update_metadata(tensor, tiny.shape, tiny._strides, tiny._storage_offset)
 def unwrap(x:torch.Tensor) -> Tensor:
   assert isinstance(x, torch.Tensor), f"x isn't {type(x)}"
   return mod.unwrap(x)
@@ -344,7 +348,7 @@ def scatter_add(self, dim, index, src, out):
 def _copy_between_devices(src, dest, cast_dtype, to_device, non_blocking=False):
   if src.is_tiny and dest.is_tiny:
     src_t, dest_t = unwrap(src), unwrap(dest)
-    if dest_t.uop.is_contiguous() or dest_t.uop.is_realized: src_t = src_t.contiguous()
+    if dest_t.uop.has_buffer_identity() or dest_t.uop.is_realized: src_t = src_t.contiguous()
     _apply_inplace(dest_t, src_t.cast(cast_dtype).to(to_device))
   elif src.is_tiny and dest.is_cpu:
     dest.resize_(src.numel()).resize_(src.shape)
@@ -611,7 +615,10 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.fill_.Tensor": lambda self, value: Tensor.full(self.shape, value.reshape(()).item(), device=self.device, dtype=self.dtype),
   "aten.flip": Tensor.flip,
   "aten.scatter_reduce.two": Tensor.scatter_reduce,
-  "aten.squeeze_.dim": lambda self, dim: self.replace(self.squeeze(dim), allow_shape_mismatch=True), # TODO: inplace view op, here?
+  "aten.squeeze_.dim": Tensor.squeeze,
+  "aten.unsqueeze_": Tensor.unsqueeze,
+  "aten.transpose_": Tensor.transpose,
+  "aten.t_": Tensor.transpose,
   "aten.add.Tensor": lambda input,other,alpha=1: input+alpha*other,
   "aten.linspace": lambda start, stop, steps, dtype=None, **kwargs:
     Tensor.linspace(start, stop, steps, **({"dtype": _from_torch_dtype(dtype)} if dtype is not None else {})),
@@ -655,6 +662,13 @@ inplace_ops = {
   "aten.masked_fill_.Tensor",
 }
 
+inplace_view_ops = {
+  "aten.squeeze_.dim",
+  "aten.unsqueeze_",
+  "aten.transpose_",
+  "aten.t_",
+}
+
 def wrap_fxn(k,f):
   def nf(*args, **kwargs):
     if TORCH_DEBUG:
@@ -675,8 +689,42 @@ def wrap_inplace(k,f):
     return orig
   return nf
 
+def wrap_inplace_view_op(k,f):
+  def nf(*args, **kwargs):
+    orig = args[0]
+    args, kwargs = unwrap_args(args, kwargs)
+    target = args[0]
+    new_view = f(*args, **kwargs)
+    if new_view is target or new_view.uop is target.uop:
+      _update_torch_metadata(orig, target)
+      return orig
+    base = canonical_base(target)
+    op = (f, args[1:], kwargs)
+    if target is base:
+      views = derived_views(base)
+      if views:
+        old_base = Tensor(base.uop, device=base.device)
+        old_base.requires_grad = base.requires_grad
+        old_base._views = getattr(base, "_views", set())
+        for v in views: v._view_base = old_base
+        base._views = set()
+        base._view_base = old_base
+        base._view_ops = [op]
+        old_base._views.add(weakref.ref(base))
+    else:
+      target._view_base = base
+      base._views = getattr(base, "_views", set())
+      base._views.add(weakref.ref(target))
+      target._view_ops = _get_view_ops(target) + [op]
+    target.uop = new_view.uop
+    _update_torch_metadata(orig, target)
+    return orig
+  return nf
+
 for k,v in tiny_backend.items():
-  wrapper = wrap_inplace if k in inplace_ops else wrap_fxn
+  if k in inplace_view_ops: wrapper = wrap_inplace_view_op
+  elif k in inplace_ops: wrapper = wrap_inplace
+  else: wrapper = wrap_fxn
   torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrapper(k,v))
 
 @torch.library.impl("aten::equal", "privateuseone")

extra/torch_backend/test_inplace.py

@@ -67,5 +67,37 @@ class TestTorchBackendInplace(unittest.TestCase):
     d += torch.arange(4)
     np.testing.assert_array_equal(a.cpu(), torch.arange(4).cpu())
 
+  def test_inplace_view_metadata(self):
+    a = torch.arange(6, dtype=torch.float32).reshape(1, 2, 3)
+    ret = a.squeeze_(0)
+    self.assertIs(ret, a)
+    self.assertEqual(a.shape, torch.Size([2, 3]))
+    ret = a.unsqueeze_(1)
+    self.assertIs(ret, a)
+    self.assertEqual(a.shape, torch.Size([2, 1, 3]))
+    ret = a.transpose_(0, 2)
+    self.assertIs(ret, a)
+    self.assertEqual(a.shape, torch.Size([3, 1, 2]))
+
+  def test_t_inplace_metadata(self):
+    a = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+    ret = a.t_()
+    self.assertIs(ret, a)
+    self.assertEqual(a.shape, torch.Size([3, 2]))
+    expected = torch.arange(6, dtype=torch.float32).reshape(2, 3).t()
+    np.testing.assert_array_equal(a.cpu().numpy(), expected.cpu().numpy())
+
+  def test_squeeze_matmul(self):
+    # squeeze_ is used internally by PyTorch for vector-matrix matmul (unsqueeze -> mm -> squeeze_)
+    a = torch.arange(65, dtype=torch.float32)
+    b = torch.arange(65*45, dtype=torch.float32).reshape(65, 45)
+    result = a.matmul(b)
+    self.assertEqual(result.shape, torch.Size([45]))
+    # verify correctness
+    a_cpu = torch.arange(65, dtype=torch.float32, device='cpu')
+    b_cpu = torch.arange(65*45, dtype=torch.float32, device='cpu').reshape(65, 45)
+    expected = a_cpu.matmul(b_cpu)
+    np.testing.assert_allclose(result.cpu().numpy(), expected.numpy(), rtol=1e-4, atol=1e-4)
+
 if __name__ == "__main__":
   unittest.main()

extra/torch_backend/wrapped_tensor.cpp

@@ -131,7 +131,15 @@ py::object unwrap_tensor(const at::Tensor &tensor) {
   return py::reinterpret_borrow<py::object>(tiny->ptr(getPyInterpreter()));
 }
 
+void update_metadata(const at::Tensor &tensor, const std::vector<int64_t> &sizes,
+                     const std::vector<int64_t> &strides, int64_t storage_offset) {
+  auto* impl = tensor.unsafeGetTensorImpl();
+  impl->set_allow_tensor_metadata_change(true);
+  impl->set_sizes_and_strides(sizes, strides, storage_offset);
+}
+
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("wrap", &wrap_tensor);
   m.def("unwrap", &unwrap_tensor);
+  m.def("update_metadata", &update_metadata);
 }

extra/viz/README

@@ -0,0 +1,18 @@
+A command line tool for exploring the VIZ trace.
+
+After running with VIZ=-1, use `PYTHONPATH=. extra/viz/cli.py` to explore the saved trace files.
+
+## Inspect runtime profiling
+
+Use `PYTHONPATH=. extra/viz/cli.py --profile` to list all traced devices.
+
+List top slowest kernels on a device: `--profile --device "AMD"`
+List samples of a kernel on a device: `--profile --device "AMD" --kernel E_3`
+
+## Inspect codegen and PatternMatcher
+
+Use `PYTHONPATH=. extra/viz/cli.py --rewrites` to list all traced kernels.
+
+List all codegen steps for a kernel: `--rewrites --kernel E_3`
+Get source code: `--rewrites --kernel E_3 --select "View Program"`
+Inspect a graph rewrite: `--rewrites --kernel E_3 --select "initial symbolic"`

extra/viz/cli.py

@@ -3,7 +3,8 @@ import argparse, pathlib
 from typing import Iterator
 from tinygrad.viz import serve as viz
 from tinygrad.uop.ops import RewriteTrace
-from tinygrad.helpers import temp, ansistrip, colored
+from tinygrad.helpers import temp, ansistrip, colored, time_to_str, ansilen
+from test.null.test_viz import load_profile
 
 def optional_eq(val:dict, arg:str|None) -> bool: return arg is None or ansistrip(val["name"]) == arg
 
@@ -25,13 +26,61 @@ def print_data(data:dict) -> None:
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
-  parser.add_argument('--kernel', type=str, default=None, metavar="NAME", help='Select a kernel by name (optional name, default: only list names)')
-  parser.add_argument('--select', type=str, default=None, metavar="NAME",
-                      help='Select an item within the chosen kernel (optional name, default: only list names)')
+  g_mode = parser.add_argument_group("mode")
+  g_mode.add_argument("--profile", action="store_true", help="View profile trace")
+  g_mode.add_argument("--rewrites", action="store_true", help="View rewrites trace")
+  g_profile = parser.add_argument_group("profile options")
+  g_profile.add_argument("--device", type=str, default=None, metavar="NAME", help="Select a device (optional name, default: only list names)")
+  g_profile.add_argument("--top", type=int, default=10, metavar="N", help="Number of top kernels to show (-1 for all, default: 10)")
+  g_rewrites = parser.add_argument_group("rewrites options")
+  g_rewrites.add_argument("--select", type=str, default=None, metavar="NAME",
+                          help="Select an item within the chosen kernel (optional name, default: only list names)")
+  g_common = parser.add_argument_group("common options")
+  g_common.add_argument("--kernel", type=str, default=None, metavar="NAME", help="Select a kernel by name (optional name, default: only list names)")
+  parser.add_argument("--profile-path", type=pathlib.Path, metavar="PATH", help="Path to profile (optional file, default: latest profile)",
+                        default=pathlib.Path(temp("profile.pkl", append_user=True)))
+  parser.add_argument("--rewrites-path", type=pathlib.Path, metavar="PATH", help="Path to rewrites (optional file, default: latest rewrites)",
+                        default=pathlib.Path(temp("rewrites.pkl", append_user=True)))
   args = parser.parse_args()
+  if not args.profile and not args.rewrites:
+    parser.print_help()
+    exit(0)
 
-  viz.trace = viz.load_pickle(pathlib.Path(temp("rewrites.pkl", append_user=True)), default=RewriteTrace([], [], {}))
+  viz.trace = viz.load_pickle(args.rewrites_path, default=RewriteTrace([], [], {}))
   viz.ctxs = viz.get_rewrites(viz.trace)
+
+  if args.profile:
+    from tabulate import tabulate
+    profile = load_profile(viz.load_pickle(args.profile_path, default=[]))
+    agg, total, n = {}, 0, 0
+    if args.device is None: print("Select a device:")
+    for k,v in profile["layout"].items():
+      if not optional_eq({"name":k}, args.device): continue
+      print(f"  {k}")
+      if args.device is None: continue
+      for e in v.get("events", []):
+        et = e["dur"]*1e-6
+        if args.kernel is not None:
+          if ansistrip(e["name"]) == args.kernel and n < 10:
+            ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
+            name = e["name"]+(" " * (46 - ansilen(e["name"])))
+            print(f"{name} {ptm}/{(et or 0)*1e3:9.2f}ms  "+e['fmt'].replace('\n', ' | ')+"  ")
+            n += 1
+        else:
+          a = agg.setdefault(e["name"], [0.0, 0])
+          a[0] += et
+          a[1] += 1
+          total += et
+    if agg and total > 0:
+      items = sorted(agg.items(), key=lambda kv:kv[1][0], reverse=True)
+      sel = items if args.top == -1 else items[:args.top]
+      table = [[name, time_to_str(t, w=9), c, f"{(t/total*100.0):.2f}%"] for name,(t,c) in sel]
+      if args.top != -1 and (other:=items[len(sel):]):
+        other_t = total-sum(t for _, (t, _) in sel)
+        table.append([f"Other ({len(other)} unique)", time_to_str(other_t, w=9), sum(c for _,(_,c) in other), f"{other_t/total*100.0:.2f}%"])
+      print(tabulate(table, headers=["name", "total", "count", "pct"], tablefmt="github"))
+    exit(0)
+
   for k in viz.ctxs:
     if not optional_eq(k, args.kernel): continue
     print(k["name"])

pyproject.toml

@@ -67,7 +67,7 @@ testing_minimal = [
   "pytest-timeout",
   "pytest-split",
   "hypothesis>=6.148.9",
-  "z3-solver",
+  "z3-solver<4.15.4",  # 4.15.4 has a segfault when creating many z3.Context()
 ]
 testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai", "ggml-python"]
 testing = [

test/device/test_hcq.py

@@ -119,7 +119,7 @@ class TestHCQ(unittest.TestCase):
     TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
 
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 1.0, f"got val {val}"
 
   def test_exec_2_kernels_100_times(self):
@@ -135,7 +135,7 @@ class TestHCQ(unittest.TestCase):
       q.submit(TestHCQ.d0, {virt_val.expr: TestHCQ.d0.timeline_value})
       TestHCQ.d0.timeline_value += 1
 
-    val = TestHCQ.a.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 200.0, f"got val {val}"
 
   @unittest.skipIf(Device.DEFAULT in {"CPU"}, "No globals/locals on LLVM/CPU")
@@ -151,9 +151,9 @@ class TestHCQ(unittest.TestCase):
     TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
 
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 1.0, f"got val {val}"
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[1]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
     assert val == 0.0, f"got val {val}, should not be updated"
 
   @unittest.skipIf(Device.DEFAULT in {"CPU"}, "No globals/locals on LLVM/CPU")
@@ -186,7 +186,7 @@ class TestHCQ(unittest.TestCase):
           TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
           TestHCQ.d0.timeline_value += 1
 
-          res_sum = sum(x for x in zt.as_buffer().cast("I"))
+          res_sum = sum(x for x in zt.as_memoryview().cast("I"))
           assert x * y * z == res_sum, f"want {x * y * z}, got {res_sum}"
 
   # Test copy
@@ -200,7 +200,7 @@ class TestHCQ(unittest.TestCase):
     TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
 
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[1]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
     assert val == 1.0, f"got val {val}"
 
   def test_copy_long(self):
@@ -218,7 +218,7 @@ class TestHCQ(unittest.TestCase):
     TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
 
-    mv_buf1 = buf1.as_buffer().cast('Q')
+    mv_buf1 = buf1.as_memoryview().cast('Q')
     assert libc.memcmp(mv_address(mv_buf1), buf2._buf.va_addr, sz) == 0
 
   @slow
@@ -242,7 +242,7 @@ class TestHCQ(unittest.TestCase):
       TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1
 
-      mv_buf1 = buf1.as_buffer()
+      mv_buf1 = buf1.as_memoryview()
       assert libc.memcmp(mv_address(mv_buf1), buf2._buf.va_addr, sz) == 0
 
   def test_update_copy(self):
@@ -260,7 +260,7 @@ class TestHCQ(unittest.TestCase):
     TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
 
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[1]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
     assert val == 1.0, f"got val {val}"
 
   def test_update_copy_long(self):
@@ -283,7 +283,7 @@ class TestHCQ(unittest.TestCase):
     TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
 
-    mv_buf1 = buf1.as_buffer().cast('Q')
+    mv_buf1 = buf1.as_memoryview().cast('Q')
     for i in range(sz//8): assert mv_buf1[i] == 0x0101010101010101, f"offset {i*8} differs, not all copied, got {hex(mv_buf1[i])}"
 
   # Test bind api
@@ -421,7 +421,7 @@ class TestHCQ(unittest.TestCase):
       TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1
 
-      assert buf1.as_buffer()[0] == i
+      assert buf1.as_memoryview()[0] == i
 
   def test_small_copies_from_host_buf_intercopy(self):
     if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
@@ -440,7 +440,7 @@ class TestHCQ(unittest.TestCase):
       TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1
 
-      assert buf2.as_buffer()[0] == i
+      assert buf2.as_memoryview()[0] == i
 
   def test_small_copies_from_host_buf_transfer(self):
     if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
@@ -463,7 +463,7 @@ class TestHCQ(unittest.TestCase):
       TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1
 
-      assert buf2.as_buffer()[0] == i
+      assert buf2.as_memoryview()[0] == i
 
   def test_memory_barrier(self):
     a = Tensor([0, 1], device=Device.DEFAULT, dtype=dtypes.int8).realize()
@@ -486,7 +486,7 @@ class TestHCQ(unittest.TestCase):
       TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1
 
-      assert buf1.as_buffer()[0] == (i + 1), f"has {buf1.as_buffer()[0]}, need {i + 1}"
+      assert buf1.as_memoryview()[0] == (i + 1), f"has {buf1.as_memoryview()[0]}, need {i + 1}"
 
   def test_memory_barrier_before_copy(self):
     if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
@@ -511,7 +511,7 @@ class TestHCQ(unittest.TestCase):
       TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1
 
-      assert buf2.as_buffer()[0] == i
+      assert buf2.as_memoryview()[0] == i
 
   def test_map_cpu_buffer_to_device(self):
     if Device[Device.DEFAULT].hw_copy_queue_t is None: self.skipTest("skip device without copy queue")

test/device/test_metal.py

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad.device import CompileError, Device, Compiler
+from tinygrad.device import CompileError, Device
 if Device.DEFAULT=="METAL":
   from tinygrad.runtime.ops_metal import MetalDevice, MetalCompiler, MetalProgram
 @unittest.skipIf(Device.DEFAULT!="METAL", "Metal support required")
@@ -48,28 +48,4 @@ kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgr
 """)
     with self.assertRaises(RuntimeError):
       compiled = compiled[:40] # corrupt the compiled program
-      MetalProgram(device, "r_5", compiled)
-
-  def test_program_w_empty_compiler(self):
-    device = MetalDevice("metal")
-    compiler = Compiler(device)
-    compiled = compiler.compile("""
-#include <metal_stdlib>
-kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]){
-  data0[0] = 0;
-}
-""")
-    MetalProgram(device, "r_5", compiled)
-
-  def test_bad_program_w_empty_compiler(self):
-    device = MetalDevice("metal")
-    compiler = Compiler(device)
-    # this does not raise
-    compiled = compiler.compile("""
-#include <metal_stdlib>
-kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]){
-  invalid codes;
-}
-""")
-    with self.assertRaises(RuntimeError):
       MetalProgram(device, "r_5", compiled)

test/device/test_ocl.py

@@ -1,9 +1,19 @@
 import unittest
+from unittest.mock import patch
 from tinygrad import Device
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes
 from tinygrad.runtime.ops_cl import CLDevice, CLAllocator, CLCompiler, CLProgram
 
+@unittest.skipUnless(Device.DEFAULT == "CL", "Runs only on OpenCL")
+class TestCLCompileCache(unittest.TestCase):
+  def test_compile_cached(self):
+    device = Device[Device.DEFAULT]
+    src = "__kernel void cached_test(__global int* a) { a[0] = 1; }"
+    CLProgram(device, name="cached_test", lib=src.encode())
+    with patch.object(CLCompiler, 'compile', side_effect=RuntimeError("compile should not be called on cache hit")):
+      CLProgram(device, name="cached_test", lib=src.encode())
+
 @unittest.skipUnless(Device.DEFAULT == "CL", "Runs only on OpenCL")
 class TestCLError(unittest.TestCase):
   @unittest.skip("allocates tons of memory")
@@ -17,7 +27,7 @@ class TestCLError(unittest.TestCase):
   def test_invalid_kernel_name(self):
     device = Device[Device.DEFAULT]
     with self.assertRaises(RuntimeError) as err:
-      CLProgram(device, name="", lib=CLCompiler(device, "test").compile("__kernel void test(__global int* a) { a[0] = 1; }"))
+      CLProgram(device, name="", lib="__kernel void test(__global int* a) { a[0] = 1; }".encode())
     assert str(err.exception) == "OpenCL Error -46: CL_INVALID_KERNEL_NAME"
 
   def test_unaligned_copy(self):

test/external/external_benchmark_hcopt.py

@@ -1,43 +0,0 @@
-import random
-from tinygrad.helpers import getenv
-from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
-from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
-from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer
-
-def optimize_kernel(k):
-  # TODO: update this
-  return hand_coded_optimizations(k)
-
-if __name__ == '__main__':
-  hcopt_wins = beam_wins = tie = 0
-  hcopt_total = beam_total = 0.0
-
-  worlds = load_worlds(filter_reduce=False, filter_noimage=True, filter_novariable=False)
-  random.seed(0)
-  random.shuffle(worlds)
-
-  for world in worlds[:500]:
-    k = ast_str_to_lin(world)
-    rawbufs = bufs_from_lin(k)
-
-    k_hcopt = k.copy()
-    k_hcopt.apply_opts(optimize_kernel(k_hcopt))
-    k_beam = beam_search(k.copy(), rawbufs, getenv("BEAM", 2))
-
-    disable_cache = bool(getenv("NOCACHE", 0))
-    t_hcopt = time_linearizer(k_hcopt, rawbufs, allow_test_size=False, cnt=10, disable_cache=disable_cache, clear_l2=True) * 1e6
-    t_beam = time_linearizer(k_beam, rawbufs, allow_test_size=False, cnt=10, disable_cache=disable_cache, clear_l2=True) * 1e6
-
-    if t_hcopt == t_beam: tie += 1
-    elif t_hcopt < t_beam: hcopt_wins += 1
-    else: beam_wins += 1
-    hcopt_total += t_hcopt
-    beam_total += t_beam
-
-    print(f"{t_hcopt=:5.2f} {k_hcopt.applied_opts=}")
-    print("")
-    print(f"{t_beam=:5.2f} {k_beam.applied_opts=}")
-    print("*"*20)
-
-  print(f"{hcopt_wins=}, {beam_wins=}, {tie=}")
-  print(f"{hcopt_total=:.2f}, {beam_total=:.2f}")

test/external/external_benchmark_op_conv.py

@@ -11,7 +11,7 @@ from tinygrad.dtype import ImageDType, Invalid
 # PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
 
 def vision_conv_143():
-  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 1024, 4)), (), 0)
+  c0 = UOp(Ops.PARAM, dtypes.imageh((16, 1024, 4)), (), 0)
   c2 = UOp.range(32, 3, AxisType.LOOP)
   c5 = UOp.range(128, 4, AxisType.LOOP)
   c8 = UOp.range(16, 2, AxisType.LOOP)
@@ -21,13 +21,13 @@ def vision_conv_143():
   c26 = UOp.range(7, 1, AxisType.REDUCE)
   c27 = c2*2+c26
   c32 = ((c27<3)!=True)&(c27<67)
-  c34 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 1)
+  c34 = UOp(Ops.PARAM, dtypes.imageh((32, 1024, 4)), (), 1)
   c38 = c5//2
   c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.index, Invalid))
   c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
-  c49 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((64, 49, 4)), (), 2)
+  c49 = UOp(Ops.PARAM, dtypes.imageh((64, 49, 4)), (), 2)
   c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(128), (), 3)
+  c63 = UOp(Ops.PARAM, dtypes.float.ptr(128), (), 3)
   c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
   c67 = c0.index((c2*128+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)
 
@@ -37,7 +37,7 @@ def vision_conv_143():
   return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
 
 def vision_conv_153():
-  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((8, 1024, 4)), (), 0)
+  c0 = UOp(Ops.PARAM, dtypes.imageh((8, 1024, 4)), (), 0)
   c2 = UOp.range(16, 3, AxisType.LOOP)
   c5 = UOp.range(256, 4, AxisType.LOOP)
   c8 = UOp.range(8, 2, AxisType.LOOP)
@@ -47,13 +47,13 @@ def vision_conv_153():
   c26 = UOp.range(7, 1, AxisType.REDUCE)
   c27 = c2*2+c26
   c32 = ((c27<3)!=True)&(c27<35)
-  c34 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 1024, 4)), (), 1)
+  c34 = UOp(Ops.PARAM, dtypes.imageh((16, 1024, 4)), (), 1)
   c38 = c5//2
   c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.index, Invalid))
   c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
-  c49 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((128, 49, 4)), (), 2)
+  c49 = UOp(Ops.PARAM, dtypes.imageh((128, 49, 4)), (), 2)
   c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(256), (), 3)
+  c63 = UOp(Ops.PARAM, dtypes.float.ptr(256), (), 3)
   c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
   c67 = c0.index((c2*256+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)
 
@@ -63,16 +63,16 @@ def vision_conv_153():
   return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
 
 def dm_conv_172():
-  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 240, 4)), (), 0)
+  c0 = UOp(Ops.PARAM, dtypes.imageh((1, 240, 4)), (), 0)
   c2 = UOp.range(960, 4, AxisType.LOOP)
-  c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((8, 384, 4)), (), 1)
+  c5 = UOp(Ops.PARAM, dtypes.imageh((8, 384, 4)), (), 1)
   c7 = UOp.range(32, 0, AxisType.REDUCE)
   c10 = UOp.range(4, 1, AxisType.REDUCE)
   c13 = UOp.range(12, 3, AxisType.REDUCE)
   c18 = UOp.range(8, 2, AxisType.REDUCE)
-  c23 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((240, 128, 4)), (), 2)
+  c23 = UOp(Ops.PARAM, dtypes.imageh((240, 128, 4)), (), 2)
   c35 = c5.index((c7*4+c10+c13*128+c18*1536))*c23.index((c10*4+c2%4+c7*16+c2//4*512))
-  c37 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(960), (), 3)
+  c37 = UOp(Ops.PARAM, dtypes.float.ptr(960), (), 3)
   c39 = c35.reduce(c7, c10, arg=Ops.ADD)+c37.index(c2)
   c50 = (1.0+((c39+0.044708251953125*(c39*(c39*c39)))*-2.3021129851685216).exp2()).reciprocal()*c39
   c53 = c50.reduce(c18, c13, arg=Ops.ADD)*0.010416666666666666
@@ -91,7 +91,7 @@ allocator = Device.default.allocator
 ps = get_program(ast, renderer)
 cr = CompiledRunner(replace(ps, device=Device.DEFAULT))
 
-gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.DEFINE_GLOBAL]), key=lambda u: u.arg)
+gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.PARAM]), key=lambda u: u.arg)
 # print(len(gs))
 # print([g.dtype for g in gs])
 bufs = [Buffer(ps.device, g.size, g.dtype if isinstance(g.dtype, ImageDType) else g.dtype._base).ensure_allocated() for g in gs]

test/external/external_fuzz_qcom_cpu_cache.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python3
+from tinygrad.tensor import Tensor
+import numpy as np
+
+while True:
+  arr = np.ones(1000000, dtype=np.uint8)
+  print(f"numpy: {(arr + 1)[:10]}")
+
+  ptr = arr.ctypes.data
+  tensor = Tensor.from_blob(ptr, arr.shape, dtype='uint8', device='QCOM').realize() + 1
+  print(f"from_blob: {tensor.numpy()[:10]}")

test/external/external_test_amd.py

@@ -20,7 +20,7 @@ class TestAMD(unittest.TestCase):
                               global_size=TestAMD.d0_runner.global_size, local_size=TestAMD.d0_runner.local_size)
       TestAMD.d0_runner.clprg(TestAMD.a.uop.buffer._buf, TestAMD.b.uop.buffer._buf,
                               global_size=TestAMD.d0_runner.global_size, local_size=TestAMD.d0_runner.local_size)
-    val = TestAMD.a.uop.buffer.as_buffer().cast("f")[0]
+    val = TestAMD.a.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 4000.0, f"got val {val}"
 
 if __name__ == "__main__":

test/external/external_test_gpu_crash.py

@@ -0,0 +1,126 @@
+# ruff: noqa: F405
+"""Tests for GPU crash scenarios using AMD assembly to trigger invalid operations.
+
+These tests intentionally cause GPU faults to verify error handling.
+Run with: AMD=1 python -m pytest test/external/external_test_gpu_crash.py -v
+"""
+import unittest, re
+from tinygrad.device import Device
+from extra.assembly.amd.autogen.rdna3.ins import *  # noqa: F403
+from extra.assembly.amd.dsl import s, v, Inst, NULL
+
+def assemble(code:str, name:str="test") -> str:
+  kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "wavefront_size32": 1, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8}
+  return f".text\n.globl {name}\n.p2align 8\n.type {name},@function\n{name}:\n{code}\n.rodata\n.p2align 6\n.amdhsa_kernel {name}\n" + \
+         "\n".join(f".amdhsa_{k} {v}" for k,v in kd.items()) + "\n.end_amdhsa_kernel"
+
+@unittest.skipIf(Device.DEFAULT != "AMD", "AMD required")
+class TestGPUCrash(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    from tinygrad.runtime.support.compiler_amd import HIPCompiler
+    cls.dev = Device["AMD"]
+    cls.compiler = HIPCompiler(cls.dev.arch)
+
+  def setUp(self):
+    # Verify device works before each test
+    from tinygrad import Tensor
+    try:
+      t = Tensor([1.0, 2.0], device="AMD").realize()
+      assert (t + 1).numpy().tolist() == [2.0, 3.0]
+    except Exception:
+      self.fail("Device not working before test")
+
+  def _run(self, code: str):
+    from tinygrad.runtime.ops_amd import AMDProgram
+    prg = AMDProgram(self.dev, "test", self.compiler.compile(assemble(code)))
+    prg(self.dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True)
+
+  def _run_insts(self, insts: list[Inst]): self._run("\n".join(i.disasm() for i in insts))
+
+  def _assert_gpu_fault(self, func):
+    """Assert that func raises a RuntimeError indicating a GPU fault (not a setup error)."""
+    with self.assertRaises(RuntimeError) as cm:
+      func()
+    err_msg = str(cm.exception).lower()
+    # Verify it's a GPU fault, not a setup/device initialization error
+    self.assertTrue(
+      re.search(r'fault|hang|timeout|illegal|memviol', err_msg),
+      f"Expected GPU fault error, got: {cm.exception}"
+    )
+
+
+class TestOutOfBoundsMemoryAccess(TestGPUCrash):
+  """Tests for out-of-bounds memory accesses."""
+
+  def test_global_load_null_ptr(self):
+    """Global load from NULL pointer."""
+    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0),
+             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_global_store_null_ptr(self):
+    """Global store to NULL pointer."""
+    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 0xDEADBEEF),
+             global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_global_load_unmapped_high_address(self):
+    """Global load from high unmapped address (0xDEAD00000000)."""
+    insts = [v_mov_b32_e32(v[0], 0x00000000), v_mov_b32_e32(v[1], 0xDEAD),
+             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_global_store_unmapped_high_address(self):
+    """Global store to high unmapped address."""
+    insts = [v_mov_b32_e32(v[0], 0x00000000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 0x12345678),
+             global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_global_atomic_unmapped(self):
+    """Atomic operation on unmapped memory."""
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 1),
+             global_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+
+class TestSMEMFaults(TestGPUCrash):
+  """Tests for scalar memory (SMEM) faults."""
+
+  def test_smem_load_null(self):
+    """SMEM load from NULL base."""
+    insts = [s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
+             s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_smem_load_unmapped(self):
+    """SMEM load from unmapped address."""
+    insts = [s_mov_b32(s[2], 0xBEEF0000), s_mov_b32(s[3], 0xDEAD),
+             s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+
+class TestFlatMemoryFaults(TestGPUCrash):
+  """Tests for FLAT memory instruction faults."""
+
+  def test_flat_load_null(self):
+    """FLAT load from NULL address."""
+    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0),
+             flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_flat_store_null(self):
+    """FLAT store to NULL address."""
+    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 0xDEADBEEF),
+             flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_flat_atomic_null(self):
+    """FLAT atomic on NULL address."""
+    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 1),
+             flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+
+if __name__ == "__main__":
+  unittest.main()

test/external/external_test_hcq.py

@@ -65,7 +65,7 @@ class TestHCQ(unittest.TestCase):
     q.submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.a.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 2000.0, f"got val {val}"
 
   def test_run_1000_times(self):
@@ -81,7 +81,7 @@ class TestHCQ(unittest.TestCase):
       TestHCQ.compute_queue().signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
       TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.a.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 2000.0, f"got val {val}"
 
   def test_run_to_3(self):
@@ -95,7 +95,7 @@ class TestHCQ(unittest.TestCase):
     q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 3.0, f"got val {val}"
 
   def test_update_exec(self):
@@ -106,9 +106,9 @@ class TestHCQ(unittest.TestCase):
     q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 1.0, f"got val {val}"
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[1]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
     assert val == 0.0, f"got val {val}, should not be updated"
 
   @unittest.skipUnless(Device.DEFAULT == "NV", "Only NV supports bind")
@@ -126,7 +126,7 @@ class TestHCQ(unittest.TestCase):
       TestHCQ.compute_queue().signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
       TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.a.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 2000.0, f"got val {val}"
 
   @unittest.skipUnless(Device.DEFAULT == "NV", "Only NV supports bind")
@@ -141,9 +141,9 @@ class TestHCQ(unittest.TestCase):
     q.submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 1.0, f"got val {val}"
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[1]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
     assert val == 0.0, f"got val {val}, should not be updated"
 
   @unittest.skipIf(CI, "Can't handle async update on CPU")
@@ -174,7 +174,7 @@ class TestHCQ(unittest.TestCase):
     q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 1.0, f"got val {val}"
 
   def test_submit_empty_queues(self):
@@ -206,7 +206,7 @@ class TestHCQ(unittest.TestCase):
     q.submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 1.0, f"got val {val}"
 
   def test_copy_1000_times(self):
@@ -221,7 +221,7 @@ class TestHCQ(unittest.TestCase):
     # confirm the signal didn't exceed the put value
     with self.assertRaises(RuntimeError):
       TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value + 1, timeout=50)
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[1]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
     assert val == 0.0, f"got val {val}"
 
   def test_copy(self):
@@ -231,7 +231,7 @@ class TestHCQ(unittest.TestCase):
     q.submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[1]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
     assert val == 1.0, f"got val {val}"
 
   @unittest.skipUnless(Device.DEFAULT == "NV", "Only NV supports bind")
@@ -248,7 +248,7 @@ class TestHCQ(unittest.TestCase):
     # confirm the signal didn't exceed the put value
     with self.assertRaises(RuntimeError):
       TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value + 1, timeout=50)
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[1]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
     assert val == 0.0, f"got val {val}"
 
   def test_copy_bandwidth(self):
@@ -288,7 +288,7 @@ class TestHCQ(unittest.TestCase):
     q.submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.a.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 1.0, f"got val {val}"
 
   def test_cross_device_signal(self):
@@ -319,7 +319,7 @@ class TestHCQ(unittest.TestCase):
       q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
       TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1
-      val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[0]
+      val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
       assert val == 1.0, f"got val {val}"
 
 if __name__ == "__main__":

test/external/fuzz_graph.py

@@ -30,7 +30,7 @@ def alloc_rawbuffer(device, fill=False):
   if fill:
     with Context(DEBUG=0):
       data = np.random.randint(-10000, 10000, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
-      rawbuf.copyin(Tensor(data).realize().uop.base.realized.as_buffer())
+      rawbuf.copyin(Tensor(data).realize().uop.base.realized.as_memoryview())
   return rawbuf
 
 def gen_kernel_ji(device, deps):
@@ -93,7 +93,7 @@ def run_jit(jis, all_buffers, input_buffers, var_vals):
 
   with Context(DEBUG=0):
     res_buffers = []
-    for rawbuf in all_buffers: res_buffers.append(rawbuf.as_buffer())
+    for rawbuf in all_buffers: res_buffers.append(rawbuf.as_memoryview())
     return res_buffers
 
 def fuzz_graph(jis, all_buffers, input_buffers):

test/external/fuzz_linearizer.py

@@ -1,339 +0,0 @@
-import random, traceback, ctypes, argparse, os
-from typing import Any
-import numpy as np
-from collections import defaultdict
-from extra.optimization.helpers import load_worlds, ast_str_to_lin, kern_str_to_lin
-from tinygrad.engine.realize import get_program
-
-# We need to insert ioctl before opening devices.
-if os.getenv("VALIDATE_HCQ", 0) != 0:
-  try:
-    import extra.nv_gpu_driver.nv_ioctl
-    from tinygrad import Device
-    _, _ = Device["NV"], Device["CUDA"]
-  except Exception: pass
-
-  try:
-    import extra.qcom_gpu_driver.opencl_ioctl
-    from tinygrad import Device
-    _, _ = Device["QCOM"], Device["CL"]
-  except Exception: pass
-
-from tinygrad import Tensor, Device, dtypes
-from tinygrad.tensor import _to_np_dtype
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.codegen.opt.search import get_kernel_actions, bufs_from_lin
-from tinygrad.engine.realize import CompiledRunner
-from tinygrad.helpers import getenv, from_mv, prod, colored, Context, DEBUG, Timing
-from tinygrad.uop.ops import UOp, Ops
-from tinygrad.device import is_dtype_supported
-
-def on_linearizer_will_run(): pass
-def on_linearizer_did_run(): pass
-def compare_states(x, y): return (True, "")
-
-if getenv("VALIDATE_HCQ"):
-  if Device.DEFAULT == "NV":
-    print("VALIDATE_HCQ: Comparing NV to CUDA")
-    import extra.nv_gpu_driver.nv_ioctl
-    validate_device = Device["CUDA"]
-    on_linearizer_will_run = extra.nv_gpu_driver.nv_ioctl.before_launch
-    on_linearizer_did_run = extra.nv_gpu_driver.nv_ioctl.collect_last_launch_state
-    compare_states = extra.nv_gpu_driver.nv_ioctl.compare_launch_state
-  elif Device.DEFAULT == "QCOM":
-    print("VALIDATE_HCQ: Comparing QCOM to CL")
-    import extra.qcom_gpu_driver.opencl_ioctl
-    validate_device = Device["CL"]
-    on_linearizer_will_run = extra.qcom_gpu_driver.opencl_ioctl.before_launch
-    on_linearizer_did_run = extra.qcom_gpu_driver.opencl_ioctl.collect_last_launch_state
-    compare_states = extra.qcom_gpu_driver.opencl_ioctl.compare_launch_state
-  else:
-    print(colored("VALIDATE_HCQ options is ignored", 'red'))
-
-def tuplize_uops(uops:list[UOp]) -> tuple:
-  return tuple([(x.op, x.dtype, tuple(uops.index(x) for x in x.src), x.arg) for x in uops])
-
-def get_fuzz_rawbufs(lin):
-  rawbufs = bufs_from_lin(lin)
-
-  # Reallocate output buffer with additional area to detect out-of-bounds writes.
-  RED_AREA_SIZE = 1024
-  # setting output  # TODO: multi-output kernel
-  rawbufs[0] = get_fuzz_rawbuf_like(rawbufs[0], zero=True, size=rawbufs[0].size+RED_AREA_SIZE)
-  # setting inputs
-  with Context(DEBUG=0):
-    for rawbuf in rawbufs[1:]:
-      if dtypes.is_unsigned(rawbuf.dtype):
-        data = np.random.randint(0, 100, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
-      elif dtypes.is_int(rawbuf.dtype):
-        data = np.random.randint(-100, 100, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
-      elif rawbuf.dtype == dtypes.bool:
-        data = np.random.choice([True, False], size=rawbuf.size)
-      elif rawbuf.dtype == dtypes.half:
-        data = np.random.uniform(-1, 1, size=rawbuf.size).astype(dtype=_to_np_dtype(rawbuf.dtype))
-      else:
-        data = np.random.uniform(-10, 10, size=rawbuf.size).astype(dtype=_to_np_dtype(rawbuf.dtype))
-      rawbuf.copyin(Tensor(data, device=lin.opts.device).realize().uop.base.realized.as_buffer())
-  return rawbufs
-
-def get_fuzz_rawbuf_like(old_rawbuf, zero=False, copy=False, size=None, force_device=None):
-  rawbuf = type(old_rawbuf)(force_device or old_rawbuf.device, old_rawbuf.size if size is None else size, old_rawbuf.dtype).allocate()
-  if copy:
-    with Context(DEBUG=0): rawbuf.copyin(old_rawbuf.as_buffer())
-  elif zero:
-    with Context(DEBUG=0):
-      mv = memoryview(bytearray(rawbuf.size * rawbuf.dtype.itemsize))
-      ctypes.memset(from_mv(mv), 0, len(mv))
-      rawbuf.copyin(mv)
-  return rawbuf
-
-def run_linearizer(lin: Kernel, rawbufs=None, var_vals=None) -> tuple[str, Any]: # (error msg, run state)
-  if rawbufs is None: rawbufs = bufs_from_lin(lin)
-  if var_vals is None: var_vals = {v.expr: v.min for v in lin.vars}
-
-  # TODO: images needs required_optimization
-  try:
-    prg = CompiledRunner(get_program(lin.get_optimized_ast(), lin.opts))
-  except KeyboardInterrupt: raise
-  except Exception:
-    traceback.print_exc()
-    return "COMPILE_ERROR", None
-
-  if getenv("VALIDATE_HCQ"): on_linearizer_will_run()
-  try:
-    prg(rawbufs, var_vals, wait=True)
-  except KeyboardInterrupt: raise
-  except Exception:
-    traceback.print_exc()
-    return "EXEC_ERROR", None
-
-  if getenv("VALIDATE_HCQ"): run_state = on_linearizer_did_run()
-  else: run_state = None
-
-  return "PASS", run_state
-
-def compare_linearizer(lin: Kernel, rawbufs=None, var_vals=None, ground_truth=None, rtol=1e-2, atol=1e-2):
-  # TODO: for bfloat16 it compiles linearizer, but it does not run because numpy cannot generate bf16 buffer.
-  has_bf16 = any(b.dtype.base == dtypes.bfloat16 for b in lin.bufs)
-
-  # TODO: raise specific fuzzing errors instead of str, and propagate the error message
-  try:
-    if rawbufs is None:
-      rawbufs = get_fuzz_rawbufs(lin)
-    else:
-      rawbufs[0] = get_fuzz_rawbuf_like(rawbufs[0], zero=True) # get a new output buffer
-  except KeyboardInterrupt: raise
-  except BaseException:
-    return ("RAWBUFS_ERROR", rawbufs, var_vals, ground_truth, None)
-
-  if var_vals is None:
-    # TODO: handle symbolic max case
-    var_vals = {v.expr: random.randint(v.vmin, v.vmax) for v in lin.ast.variables()}
-
-  if ground_truth is None and not has_bf16:
-    unoptimized = Kernel(lin.ast)
-    if run_linearizer(unoptimized, rawbufs, var_vals)[0] != "PASS":
-      return ("BASELINE_ERROR", rawbufs, var_vals, ground_truth, None)
-    ground_truth = np.frombuffer(rawbufs[0].as_buffer(), _to_np_dtype(rawbufs[0].dtype)).copy()
-
-  rawbufs[0] = get_fuzz_rawbuf_like(rawbufs[0], zero=True) # get a new output buffer
-  run_msg, run_state = run_linearizer(lin, rawbufs, var_vals)
-  if run_msg != "PASS": return (run_msg, rawbufs, var_vals, ground_truth, run_state)
-
-  try:
-    if not has_bf16:
-      result = np.frombuffer(rawbufs[0].as_buffer(), _to_np_dtype(rawbufs[0].dtype))
-      np.testing.assert_allclose(result, ground_truth, rtol=rtol, atol=atol)
-  except KeyboardInterrupt: raise
-  except AssertionError as e:
-    if DEBUG >= 2:
-      print(f"COMPARE_ERROR details: {e}")
-      if getenv("DEBUG_VALUES") > 0:
-        mismatch_indices = np.where(~np.isclose(result, ground_truth, rtol=rtol, atol=atol))
-        mismatched_result = result[mismatch_indices]
-        mismatched_ground_truth = ground_truth[mismatch_indices]
-        for i, idx in enumerate(mismatch_indices[0]):
-          print(f"mismatch at {idx=}: result={mismatched_result[i]} <> ground_truth={mismatched_ground_truth[i]}")
-    return ("COMPARE_ERROR", rawbufs, var_vals, ground_truth, run_state)
-
-  return ("PASS", rawbufs, var_vals, ground_truth, run_state)
-
-def fuzz_linearizer(lin: Kernel, rtol=1e-2, atol=1e-2, opts_list=None):
-  SEED = getenv("SEED", 42)
-  random.seed(SEED)
-  np.random.seed(SEED)
-  print(lin.ast)
-  print(lin.colored_shape())
-  seen_uops = {}
-  last_lins = [lin]
-  failures:defaultdict[str, list[tuple[tuple[UOp, ...], list[Opt]]]] = defaultdict(list)
-  rawbufs, var_vals, ground_truth, validate_rawbufs = None, None, None, None
-
-  FUZZ_ALL_ACTIONS = getenv("FUZZ_ALL_ACTIONS", 0)
-  FUZZ_MAX_SIZE = getenv("FUZZ_MAX_SIZE", 0)
-  FUZZ_IGNORE_SIMPLE_OPS = getenv("FUZZ_IGNORE_SIMPLE_OPS", 1)
-
-  if FUZZ_MAX_SIZE > 0 and prod(lin.full_shape) > FUZZ_MAX_SIZE:
-    print("skipping large kernel")
-    return failures
-  if FUZZ_IGNORE_SIMPLE_OPS and _is_simple(lin):
-    print("skipping simple kernel")
-    return failures
-
-  test_depth = 1 if opts_list is not None else getenv("DEPTH", 1 if FUZZ_ALL_ACTIONS else 10)
-  for depth in range(test_depth):
-    next_lins = []
-    for lin in last_lins:
-      if opts_list is None: actions = get_kernel_actions(lin, include_0=False)
-      else:
-        actions = {}
-        for oi,opts in enumerate(opts_list):
-          lin2 = lin.copy()
-          for o in opts: lin2.apply_opt(o)
-          actions[oi] = lin2
-
-      if not actions: continue
-      if depth == 0 and getenv("FUZZ_REQUIRE_TC", 0):
-        tc_acts = {i: k for k in actions.values() if k.applied_opts[0].op == OptOps.TC}
-        if len(tc_acts) == 0: return failures
-        else: actions = tc_acts
-
-      test_lins = list(actions.values())
-      if FUZZ_ALL_ACTIONS: print(f"testing {lin.applied_opts=} with {len(actions)} actions")
-      elif opts_list is None: test_lins = [random.choice(test_lins)]
-
-      for test_lin in test_lins:
-        if not FUZZ_ALL_ACTIONS and test_lin.applied_opts: print(f"applied opts: {test_lin.applied_opts}")
-
-        # stop if kernel uops repeat
-        try: tuops = tuplize_uops(get_program(test_lin.get_optimized_ast(), test_lin.ren).uops)
-        except KeyboardInterrupt: raise
-        except BaseException as e:
-          print(test_lin.ast)
-          print(test_lin.applied_opts)
-          print(e)
-          failures["LINEARIZE_ERROR"].append((test_lin.ast, test_lin.applied_opts))
-          continue
-
-        if tuops in seen_uops: continue
-        seen_uops[tuops] = tuple(test_lin.applied_opts)
-
-        if not FUZZ_ALL_ACTIONS: print(test_lin.colored_shape())
-
-        (msg, rawbufs, var_vals, ground_truth, state1) = compare_linearizer(test_lin, rawbufs, var_vals, ground_truth, rtol=rtol, atol=atol)
-        if state1 is not None and validate_device is not None:
-          validate_lin = test_lin.copy()
-          validate_lin.ren = validate_device.renderer
-          if validate_rawbufs is None:
-            validate_rawbufs = [get_fuzz_rawbuf_like(x, copy=True, force_device=validate_device.device) for x in rawbufs]
-          (_msg, _, _, _, state2) = compare_linearizer(validate_lin, validate_rawbufs, var_vals, ground_truth, rtol=rtol, atol=atol)
-
-          if _msg != "PASS": failures[f"VALIDATE_DEV_{_msg}"].append((validate_lin.ast, validate_lin.applied_opts))
-
-          ok, err_msg = compare_states(state1, state2)
-          if not ok: failures["HCQ_COMPARE_FAILURE"].append((err_msg, test_lin.ast, test_lin.applied_opts, state1, state2))
-
-        if msg != "PASS":
-          print(test_lin.ast)
-          print(test_lin.applied_opts)
-          print(msg)
-          failures[msg].append((test_lin.ast, test_lin.applied_opts))
-          continue
-
-        next_lins.append(test_lin)
-
-    last_lins = next_lins
-    if FUZZ_ALL_ACTIONS: print(f"depth={depth} total_lins={len(last_lins)} {failures=}")
-  return failures
-
-def _is_simple(lin: Kernel) -> bool:
-  if len(lin.ast.src) > 1: return False
-  ast:UOp = lin.ast.src[0]
-  if ast.src[0].op is Ops.CAST and ast.src[0].src[0].op is Ops.LOAD: return True
-  return False
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser(description="Run a fuzz testing on one or more kernels", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-  parser.add_argument("--ast", type=str, default=None, help="the ast for the kernel to be optimized")
-  parser.add_argument("--file", type=str, default=None, help="a file containing asts to be optimized, one per line")
-  parser.add_argument("--beamreplay", type=str, default=None, help="replay asts and opts got from beam with CAPTURE_BEAM")
-  parser.add_argument("--logfile", type=str, default=None, help="a file containing a tuple of ast and applied_opts, one per line")
-  parser.add_argument("--expected-failures", type=int, default=0, help="the number of expected failed kernels")
-  parser.add_argument("--rtol", type=float, default=1e-2, help="relative tolerance for numerical comparison")
-  parser.add_argument("--atol", type=float, default=1e-2, help="absolute tolerance for numerical comparison")
-  args = parser.parse_args()
-
-  opts_list = None
-  if args.ast is not None:
-    print("loaded AST from CLI")
-    ast_strs = [args.ast]
-  elif args.file is not None:
-    print(f"loading ASTs from file '{args.file}'")
-    with open(args.file, 'r') as file:
-      ast_strs = file.readlines()
-  elif args.beamreplay is not None:
-    print(f"loading BEAM replay from file '{args.beamreplay}'")
-    with open(args.beamreplay, 'r') as file: fdata = file.readlines()
-    ast_strs, opts_list = [x.split(' :: ')[0] for x in fdata if not x.startswith("#")], [x.split(' :: ')[1] for x in fdata if not x.startswith("#")]
-
-    # dedup ast_strs and opts_list
-    dct = defaultdict(list)
-    for i in range(len(ast_strs)): dct[ast_strs[i]].append(eval(opts_list[i]))
-    ast_strs_items = list(dct.keys())
-    opts_list = [dct[c] for c in ast_strs_items]
-  elif args.logfile is not None:
-    print(f"loading ASTs from LOGKERNS file '{args.file}'")
-    with open(args.logfile, 'r') as file:
-      kern_strs = file.readlines()
-      test_lins = [kern_str_to_lin(kern_str) for kern_str in kern_strs]
-      ast_strs = [f"{lin.ast}" for lin in test_lins]
-  else:
-    print("loading ASTs from world")
-    ast_strs = load_worlds(filter_reduce=False, filter_novariable=False)
-
-  print(f"{len(ast_strs)=}")
-  tested = 0
-  failed_ids = []
-  failures = defaultdict(list)
-  seen_ast_strs = set()
-
-  try:
-    for i, ast in enumerate(ast_strs[:getenv("FUZZ_N", len(ast_strs))]):
-      if (nth := getenv("FUZZ_NTH", -1)) != -1 and i != nth: continue
-      if getenv("FUZZ_IMAGEONLY") and "dtypes.image" not in ast: continue
-      if "dtypes.image" in ast and Device.DEFAULT not in {"CL", "QCOM"}: continue  # IMAGE is only for CL
-      if ast in seen_ast_strs: continue
-      seen_ast_strs.add(ast)
-
-      lin = ast_str_to_lin(ast)
-      if not all(is_dtype_supported(buf.dtype) for buf in lin.bufs):
-        print("skipping kernel due to not supported dtype")
-        continue
-
-      with Timing(f"tested ast {i}: "):
-        tested += 1
-        fuzz_failures = fuzz_linearizer(lin, rtol=args.rtol, atol=args.atol, opts_list=(opts_list[i] if opts_list else None))
-        if fuzz_failures: failed_ids.append(i)
-        for k, v in fuzz_failures.items():
-          for f in v:
-            failures[k].append(f)
-  except KeyboardInterrupt: print(colored("STOPPING...", 'red'))
-
-  for msg, errors in failures.items():
-    for i, payload in enumerate(errors):
-      print(f"{msg} {i} kernel: {payload}") # easier to use with output with verify_kernel.py
-
-  print(f"{tested=}")
-  if failures:
-    print(f"{failed_ids=}")
-    for msg, errors in failures.items():
-      print(f"{msg}: {len(errors)}")
-    if len(failed_ids) == args.expected_failures:
-      print(colored(f"{len(failed_ids)} failed as expected", "yellow"))
-  if len(failed_ids) != args.expected_failures:
-    print(colored(f"failed on {len(failed_ids)} kernels, expected {args.expected_failures}", "red"))
-    # TODO: fix this
-    # raise RuntimeError(f"failed on {len(failed_ids)} kernels, expected {args.expected_failures}")
-  else:
-    print(colored("all passed", "green"))

test/external/fuzz_symbolic.py

@@ -1,12 +1,16 @@
-import random, operator
+# NOTE: z3-solver 4.15.4 segfaults (exit code 139) when creating many z3.Context() with complex expressions.
+# Reproduces consistently with seed=74 around iteration 1767. Versions <=4.15.3 are fine.
+# Workaround: reuse a single z3.Context, or pin z3-solver<4.15.4 (see pyproject.toml).
+# To repro: pip install z3-solver==4.15.4.0 && python test/external/fuzz_symbolic.py 74
+import random, operator, sys
 import z3
 from tinygrad import Variable, dtypes
 from tinygrad.uop.ops import UOp
 from tinygrad.uop.validate import uops_to_z3
 from tinygrad.helpers import DEBUG, Context
 
-seed = random.randint(0, 100)
-print(f"Seed: {seed}")
+seed = int(sys.argv[1]) if len(sys.argv) > 1 else random.randint(0, 100)
+print(f"Seed: {seed}", flush=True)
 random.seed(seed)
 
 unary_ops = [lambda a:a+random.randint(-4, 4), lambda a: a*random.randint(-4, 4),

test/external/fuzz_symbolic_symbolic_div.py

@@ -1,11 +1,11 @@
-import random
+import random, sys
 import z3
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.uop.validate import uops_to_z3
 from tinygrad.helpers import DEBUG, Context, colored
 
-seed = random.randint(0, 100)
-print(f"Seed: {seed}")
+seed = int(sys.argv[1]) if len(sys.argv) > 1 else random.randint(0, 100)
+print(f"Seed: {seed}", flush=True)
 random.seed(seed)
 
 def get_random_term(ranges, factors):

test/external/process_replay/process_replay.py

@@ -47,13 +47,15 @@ def replay_get_rangeify_map(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str,
   UOp.unique_num = itertools.count(max([u.arg for u in big_sink.toposort() if u.op is Ops.UNIQUE], default=0)+1)
   new_sink = big_sink.substitute(get_rangeify_map(big_sink))
   def to_str(ret:UOp) -> str:
-    asts = [repr(u.arg.ast) for u in ret.toposort() if u.op is Ops.KERNEL]
+    asts = [repr(u.arg.ast) for u in ret.toposort() if u.op is Ops.CALL]
     return "\n".join([f"{len(asts)} kernels", *asts])
   return to_str(new_sink), to_str(big_sink.substitute(ret)), (big_sink,)
 
 def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
   # the ast.arg is non None if we are inside of search.py
-  sink_arg = ast.arg or KernelInfo(opts_to_apply=tuple(opts) if opts is not None else p.applied_opts if BEAM>=1 else None)
+  sink_arg = ast.arg or KernelInfo()
+  if opts is not None: sink_arg = replace(sink_arg, opts_to_apply=tuple(opts))
+  elif BEAM >= 1 and sink_arg.opts_to_apply is None: sink_arg = replace(sink_arg, opts_to_apply=p.applied_opts)
   input_ast = ast if ast.op is Ops.PROGRAM else ast.replace(arg=replace(sink_arg, name=p.name))
   p2 = get_program(input_ast, renderer=renderer)
   def to_str(ret:ProgramSpec) -> str:

test/external/speed_compare_amd_am.py

@@ -1,157 +0,0 @@
-from tinygrad import Device, dtypes
-from tinygrad.helpers import getenv, colorize_float, DEBUG
-from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from test.external.fuzz_linearizer import get_fuzz_rawbufs
-from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
-from tinygrad.codegen.opt.search import bufs_from_lin
-from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.tensor import _to_np_dtype
-from tinygrad.runtime.ops_amd import AMDDevice
-from contextlib import contextmanager
-import numpy as np
-import os, random, statistics
-
-am_signal_pages, am_signal_pool, am_devices = [], [], []
-amd_signal_pages, amd_signal_pool, amd_devices = [], [], []
-
-def rebind_vfio(pcibus="0000:44:00.0"):
-  print("rebind ", pcibus)
-  os.system("sudo rmmod amdgpu")
-  os.system("sudo modprobe vfio-pci")
-
-  base = f"/sys/bus/pci/devices/{pcibus}"
-  if os.path.exists(f"{base}/driver"):
-    with open(f"{base}/driver/unbind", "w") as f: f.write(pcibus)
-  with open(f"{base}/driver_override", "w") as f: f.write("vfio-pci")
-  with open("/sys/bus/pci/drivers_probe", "w") as f: f.write(pcibus)
-
-  os.system("sudo modprobe amdgpu")
-  os.system("rocm-smi --setprofile compute")
-  os.system("rocm-smi --setperflevel high")
-
-@contextmanager
-def run_amd():
-  global amd_signal_pages, amd_signal_pool, amd_devices
-  AMDDevice.driverless = False
-  AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices = amd_signal_pages, amd_signal_pool, amd_devices
-  yield
-  amd_signal_pages, amd_signal_pool, amd_devices = AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices
-  AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices = [], [], []
-
-@contextmanager
-def run_am():
-  global am_signal_pages, am_signal_pool, am_devices
-  AMDDevice.driverless = True
-  AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices = am_signal_pages, am_signal_pool, am_devices
-  yield
-  am_signal_pages, am_signal_pool, am_devices = AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices
-  AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices = [], [], []
-
-if __name__ == "__main__":
-  CHECK_CPU = getenv("CHECK_CPU", 0)
-  SEED = getenv("SEED", 42)
-  CNT = getenv("CNT", 7)
-  random.seed(SEED)
-  np.random.seed(SEED)
-
-  # TODO: NUM=780 is super slow
-  # NUM=1907 is broken on AMD and AM have some mismatches (0 vs 1)
-  # kfd feels so bad when taking gpu out while it's running... Need hacks to rebind it before running.
-  rebind_vfio(pcibus="0000:44:00.0")
-
-  ast_strs = load_worlds(filter_reduce=False, filter_novariable=True)
-
-  with run_am():
-    amdev = Device["AMD:1"]
-
-  with run_amd():
-    amddev = Device["AMD"]
-
-  if CHECK_CPU: cpudev = Device["CPU"]
-
-  single = getenv("NUM", -1)
-  if single != -1: ast_strs = ast_strs[single:single+1]
-
-  average_tm_amd, average_tm_am = 0, 0
-  for num,ast in enumerate(ast_strs):
-    with run_amd():
-      amdlin = ast_str_to_lin(ast, opts=amddev.renderer)
-      amdlin.apply_opts(hand_coded_optimizations(amdlin))
-      has_bf16 = any(b.dtype == dtypes.bfloat16 for b in amdlin.bufs)
-
-      amd_prg = CompiledRunner(get_program(amdlin.get_optimized_ast(), amdlin.opts))
-      amdbufs = bufs_from_lin(amdlin)
-      test_amdbufs = get_fuzz_rawbufs(amdlin) if not has_bf16 else amdbufs
-      if not has_bf16: contents = [buf.as_buffer() for buf in test_amdbufs]
-
-    with run_am():
-      rdr = amdev.renderer
-      rdr.device = "AMD:1"
-      amlin = ast_str_to_lin(ast, opts=amdev.renderer)
-      amlin.apply_opts(hand_coded_optimizations(amlin))
-      am_prg = CompiledRunner(get_program(amlin.get_optimized_ast(), amlin.opts))
-      ambufs = bufs_from_lin(amlin)
-      test_ambufs = get_fuzz_rawbufs(amlin) if not has_bf16 else ambufs
-      if not has_bf16:
-        for i,rawbuf in enumerate(test_ambufs): rawbuf.copyin(contents[i])
-
-    if CHECK_CPU:
-      cpu_rdr = cpudev.renderer
-      cpu_rdr.device = "CPU"
-      cpulin = ast_str_to_lin(ast, opts=cpu_rdr)
-      cpulin.apply_opts(hand_coded_optimizations(cpulin))
-      cpu_prg = CompiledRunner(get_program(cpulin.get_optimized_ast(), cpulin.opts))
-      cpubufs = bufs_from_lin(cpulin)
-      test_cpubufs = get_fuzz_rawbufs(cpulin) if not has_bf16 else ambufs
-      if not has_bf16:
-        for i,rawbuf in enumerate(test_cpubufs): rawbuf.copyin(contents[i])
-
-    # warmup
-    tm_amd, tm_am, failed = [], [], False
-    with run_amd():
-      try:
-        amd_prg(test_amdbufs, {}, wait=True)
-        for i in range(CNT): tm_amd.append(amd_prg(amdbufs, {}, wait=True))
-      except RuntimeError:
-        print("AMD FAILED")
-        tm_amd = [1e9]
-        failed = True
-
-    with run_am():
-      try:
-        am_prg(test_ambufs, {}, wait=True)
-        for i in range(CNT): tm_am.append(am_prg(ambufs, {}, wait=True))
-      except RuntimeError:
-        print("AM FAILED")
-        tm_am = [1e9]
-        failed = True
-
-    if CHECK_CPU:
-      cpu_prg(test_cpubufs, {}, wait=True)
-      for i in range(1): cpu_prg(cpubufs, {}, wait=True)
-
-    if not failed and not has_bf16:
-      with run_amd():
-        curesult = np.frombuffer(test_amdbufs[0].as_buffer(), _to_np_dtype(test_amdbufs[0].dtype))
-
-      with run_am():
-        amresult = np.frombuffer(test_ambufs[0].as_buffer(), _to_np_dtype(test_ambufs[0].dtype))
-
-      if CHECK_CPU:
-        cpuresult = np.frombuffer(test_cpubufs[0].as_buffer(), _to_np_dtype(test_cpubufs[0].dtype))
-        np.testing.assert_allclose(amresult, cpuresult, rtol=1e-2, atol=1e-2)
-        np.testing.assert_allclose(curesult, cpuresult, rtol=1e-2, atol=1e-2)
-
-      try:
-        np.testing.assert_allclose(curesult, amresult, rtol=1e-2, atol=1e-2)
-      except AssertionError as e:
-        print("AM and AMD results do not match")
-        print(e)
-
-    bam = statistics.median(tm_am)
-    bamd = statistics.median(tm_amd)
-    average_tm_amd += bamd
-    average_tm_am += bam
-    ratio = bam/bamd
-    print(f"{average_tm_am/average_tm_amd:5.2f}x -- {num:4d} {colorize_float(ratio)} {bam*1e6:7.2f} vs {bamd*1e6:7.2f} us", amlin.name)
-    if DEBUG > 3 and ratio > 1.04: print(f"AM slower {ratio}", amlin.ast, amlin.applied_opts)

test/external/speed_compare_cuda_nv.py

@@ -1,69 +0,0 @@
-from tinygrad import Device, dtypes
-from tinygrad.helpers import getenv, colorize_float
-from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from test.external.fuzz_linearizer import get_fuzz_rawbufs
-from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
-from tinygrad.codegen.opt.search import bufs_from_lin
-from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.tensor import _to_np_dtype
-import numpy as np
-
-if __name__ == "__main__":
-  ast_strs = load_worlds(filter_reduce=False, filter_novariable=True)
-  cudev = Device["CUDA"]
-  nvdev = Device["NV"]
-
-  # NUM=112 python3 test/external/speed_compare_cuda_nv.py
-
-  single = getenv("NUM", -1)
-  if single != -1: ast_strs = ast_strs[single:single+1]
-
-  average_tm_cuda, average_tm_nv = 0, 0
-  for num,ast in enumerate(ast_strs):
-    # cuda compile
-    culin = ast_str_to_lin(ast, opts=cudev.renderer)
-    culin.apply_opts(hand_coded_optimizations(culin))
-    has_bf16 = any(b.dtype == dtypes.bfloat16 for b in culin.bufs)
-
-    cuda_prg = CompiledRunner(get_program(culin.get_optimized_ast(), culin.opts))
-    cubufs = bufs_from_lin(culin)
-    test_cubufs = get_fuzz_rawbufs(culin) if not has_bf16 else cubufs
-
-    rdr = nvdev.renderer
-    rdr.device = "NV"
-    nvlin = ast_str_to_lin(ast, opts=rdr)
-    nvlin.apply_opts(hand_coded_optimizations(nvlin))
-    nv_prg = CompiledRunner(get_program(nvlin.get_optimized_ast(), nvlin.opts))
-    nvbufs = bufs_from_lin(nvlin)
-    test_nvbufs = get_fuzz_rawbufs(nvlin) if not has_bf16 else nvbufs
-    if not has_bf16:
-      for i,rawbuf in enumerate(test_nvbufs): rawbuf.copyin(test_cubufs[i].as_buffer())
-
-    # warmup
-    tm_cuda, tm_nv, failed = [], [], False
-    try:
-      cuda_prg(test_cubufs, {}, wait=True)
-      for i in range(5): tm_cuda.append(cuda_prg(cubufs, {}, wait=True))
-    except RuntimeError:
-      print("CUDA FAILED")
-      tm_cuda = [1e9]
-      failed = True
-
-    try:
-      nv_prg(test_nvbufs, {}, wait=True)
-      for i in range(5): tm_nv.append(nv_prg(nvbufs, {}, wait=True))
-    except RuntimeError:
-      print("NV FAILED")
-      tm_nv = [1e9]
-      failed = True
-
-    if not failed and not has_bf16:
-      curesult = np.frombuffer(test_cubufs[0].as_buffer(), _to_np_dtype(test_cubufs[0].dtype))
-      nvresult = np.frombuffer(test_nvbufs[0].as_buffer(), _to_np_dtype(test_nvbufs[0].dtype))
-      np.testing.assert_allclose(curesult, nvresult, rtol=1e-2, atol=1e-2)
-
-    average_tm_cuda += min(tm_cuda)
-    average_tm_nv += min(tm_nv)
-    ratio = min(tm_nv)/min(tm_cuda)
-    print(f"{average_tm_nv/average_tm_cuda:5.2f}x -- {num:4d} {colorize_float(ratio)}  {min(tm_nv)*1e6:7.2f} us", nvlin.name)
-    if ratio > 1.04: print(f"NV slower {ratio}", nvlin.ast, nvlin.applied_opts)

test/external/speed_compare_cuda_ptx.py

@@ -1,60 +0,0 @@
-import itertools
-from tinygrad import Device
-from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
-from tinygrad.helpers import getenv, colorize_float
-from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from tinygrad.codegen.opt.search import bufs_from_lin
-from tinygrad.runtime.ops_cuda import PTXCompiler, PTXRenderer, CUDACompiler
-
-if __name__ == "__main__":
-  ast_strs = load_worlds(filter_reduce=False, filter_novariable=True)
-  # no bfloat16 for ptx at the moment
-  ast_strs = [x for x in ast_strs if "dtypes.bfloat16" not in x]
-  dev = Device["CUDA"]
-  ptx = PTXRenderer(dev.arch)
-
-  # NUM=112 python3 test/external/speed_compare_cuda_ptx.py
-
-  single = getenv("NUM", -1)
-  if single != -1: ast_strs = ast_strs[single:single+1]
-
-  average_tm_cuda, average_tm_ptx = 0, 0
-  for num,ast in enumerate(ast_strs):
-    # cuda compile
-    dev.compiler = CUDACompiler(dev.arch)
-    lin = ast_str_to_lin(ast, opts=dev.renderer)
-    lin.apply_opts(hand_coded_optimizations(lin))
-    cuda_prg = CompiledRunner(get_program(lin.get_optimized_ast(), lin.opts))
-
-    bufs = bufs_from_lin(lin)
-
-    # ptx compile
-    dev.compiler = PTXCompiler(dev.arch)
-    lin = ast_str_to_lin(ast, opts=ptx)
-    lin.apply_opts(hand_coded_optimizations(lin))
-    ptx_prg = CompiledRunner(get_program(lin.get_optimized_ast(), lin.opts))
-
-    # warmup
-    try:
-      cuda_prg(bufs, {}, wait=True)
-    except RuntimeError:
-      print("cuda failed ast:", num)
-      continue
-    ptx_prg(bufs, {}, wait=True)
-
-    tm_cuda, tm_ptx = [], []
-    for i in range(5):
-      tm_cuda.append(cuda_prg(bufs, {}, wait=True))
-      tm_ptx.append(ptx_prg(bufs, {}, wait=True))
-    average_tm_cuda += min(tm_cuda)
-    average_tm_ptx += min(tm_ptx)
-    ratio = min(tm_ptx)/min(tm_cuda)
-    print(f"{average_tm_ptx/average_tm_cuda:5.2f}x -- {num:4d} {colorize_float(ratio)}  {min(tm_ptx)*1e6:7.2f} us", lin.name)
-    if ratio > 1.5:
-      def fix(x): return x.replace('\t', ' ').strip()
-      ll1, ll2 = cuda_prg.lib.decode().split('\n'), ptx_prg.lib.decode().split('\n')
-      if single != -1:
-        for ln, (l1, l2) in enumerate(itertools.zip_longest(ll1, ll2, fillvalue='')):
-          print(f"{ln:5d} | {fix(l1):80s} | {fix(l2):80s}")
-      print(len(ll1), len(ll2), "RATIO", ratio, "us", min(tm_ptx)*1e6)

test/external/verify_kernel.py

@@ -1,78 +0,0 @@
-import argparse
-from collections import defaultdict
-from extra.optimization.helpers import kern_str_to_lin, time_linearizer
-from test.external.fuzz_linearizer import compare_linearizer
-from tinygrad.helpers import colored
-from tinygrad.codegen.opt.kernel import Kernel
-
-# Use this with the LOGKERNS options to verify that all executed kernels are valid and evaluate to the same ground truth results
-
-# Example for GPT2:
-# 1) Run the model to log all kernels: `PYTHONPATH=. LOGKERNS=/tmp/gpt2_kerns.txt JIT=1 HALF=1 BEAM=2 CACHELEVEL=0 python3 examples/gpt2.py --count 10 --temperature 0 --timing`   # noqa: E501
-# 2) Validate the kernel correctness: `PYTHONPATH=. python3 ./test/external/verify_kernel.py --file /tmp/gpt2_kerns.txt`
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser(description="Verify the correctness of one or more kernel", formatter_class=argparse.ArgumentDefaultsHelpFormatter)    # noqa: E501
-  parser.add_argument("--kernel", type=str, default=None, help="a string of a tuple of (ast, applied_opts,)")
-  parser.add_argument("--file", type=str, default=None, help="a file containing a tuple of ast and applied_opts, one per line")
-  parser.add_argument("--pkl", type=str, default=None, help="a pickle file containing a single tuple of ast and applied_opts")
-  parser.add_argument("--rtol", type=float, default=1e-2, help="relative tolerance for numerical comparison")
-  parser.add_argument("--atol", type=float, default=1e-2, help="absolute tolerance for numerical comparison")
-  parser.add_argument("--timing", action='store_true', help="show final timing for the kernel")
-  parser.add_argument("--expected-failures", type=int, default=0, help="the number of expected failed kernels")
-  args = parser.parse_args()
-
-  if args.kernel is not None:
-    print("loading kernel from args")
-    test_lins = [kern_str_to_lin(args.kernel)]
-  elif args.file is not None:
-    print(f"loading kernel from file '{args.file}'")
-    with open(args.file, 'r') as file:
-      kern_strs = file.readlines()
-      test_lins = [kern_str_to_lin(kern_str) for kern_str in kern_strs]
-  elif args.pkl is not None:
-    print(f"loading kernel from pickle file '{args.file}'")
-    import pickle
-    with open(args.pkl, 'rb') as file:
-      (ast, applied_opts,) = pickle.load(file)
-      lin = Kernel(ast)
-      lin.apply_opts(applied_opts)
-      test_lins = [lin]
-
-  else:
-    raise RuntimeError("no kernel specified; use --kernel, --file, or --pkl options")
-
-  print(f"verifying {len(test_lins)} kernels")
-
-  failed_ids = []
-  failures = defaultdict(list)
-  for i, test_lin in enumerate(test_lins):
-    print(f"testing kernel {i}")
-    print(test_lin.ast)
-    print(test_lin.applied_opts)
-    unoptimized_lin = Kernel(test_lin.ast)
-    print(f"{unoptimized_lin.colored_shape()} -> {test_lin.colored_shape()}")
-    (msg,rb,vv,gt) = compare_linearizer(test_lin, None, None, None, rtol=args.rtol, atol=args.atol)
-    if msg != "PASS":
-      failed_ids.append(i)
-      failures[msg].append((test_lin.ast, test_lin.applied_opts))
-    if args.timing:
-      tm = time_linearizer(test_lin, rb, allow_test_size=False, cnt=10)
-      print(f"final time {tm*1e6:9.0f} us")
-
-  for msg, errors in failures.items():
-    for i, (ast, opts) in enumerate(errors):
-      print(f"{msg} {i} AST: {ast}")
-      print(f"{msg} {i} OPTS: {opts}\n")
-
-  print(f"tested {len(test_lins)} kernels")
-  if failures:
-    print(f"{failed_ids=}")
-    for msg, errors in failures.items():
-      print(f"{msg}: {len(errors)}")
-    if len(failed_ids) == args.expected_failures:
-      print(colored(f"{len(failed_ids)} failed as expected", "yellow"))
-  if len(failed_ids) != args.expected_failures:
-    raise RuntimeError(f"failed on {len(failed_ids)} kernels, expected {args.expected_failures}")
-  else:
-    print(colored("all passed", "green"))

test/helpers.py

@@ -61,8 +61,8 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None):
   for buf_dt, data in inputs or []:
     bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
     allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
-  g = UOp(Ops.DEFINE_GLOBAL, uop.dtype.ptr(), arg=0, src=())
-  prg = get_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(), PythonRenderer())
+  g = UOp(Ops.PARAM, uop.dtype.ptr(), arg=0, src=())
+  prg = get_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(arg=KernelInfo()), PythonRenderer())
   prog = PythonProgram("run", PythonCompiler().compile(prg.src))
   prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs)
   return out_buf.cast(uop.dtype.fmt or "").tolist()[0]

test/mockgpu/nv/nvdriver.py

@@ -130,15 +130,18 @@ class NVDriver(VirtDriver):
       struct.hObjectNew = self._alloc_handle()
       self.object_by_handle[struct.hObjectNew] = NVContextShare(self.object_by_handle[struct.hObjectParent])
     elif struct.hClass == nv_gpu.AMPERE_CHANNEL_GPFIFO_A:
-      assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVChannelGroup)
+      parent = self.object_by_handle.get(struct.hObjectParent)
+      assert parent is not None and isinstance(parent, (NVChannelGroup, NVGPU))
       struct.hObjectNew = self._alloc_handle()
       params = nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS.from_address(params_ptr)
-      gpu = self.object_by_handle[struct.hObjectParent].device
+      gpu = parent.device if isinstance(parent, NVChannelGroup) else parent
       gpfifo_token = gpu.add_gpfifo(params.gpFifoOffset, params.gpFifoEntries)
       self.object_by_handle[struct.hObjectNew] = NVGPFIFO(gpu, gpfifo_token)
-    elif struct.hClass == nv_gpu.AMPERE_DMA_COPY_B or struct.hClass == nv_gpu.ADA_COMPUTE_A:
+    elif struct.hClass in (nv_gpu.AMPERE_DMA_COPY_B, nv_gpu.ADA_COMPUTE_A, nv_gpu.NVC9B0_VIDEO_DECODER, nv_gpu.NVCFB0_VIDEO_DECODER):
       assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVGPFIFO)
       struct.hObjectNew = self._alloc_handle()
+      gpfifo = self.object_by_handle[struct.hObjectParent]
+      gpfifo.device.queues[gpfifo.token].bound_engines.add(struct.hClass)
     elif struct.hClass == nv_gpu.GT200_DEBUGGER:
       struct.hObjectNew = self._alloc_handle()
     elif struct.hClass == nv_gpu.MAXWELL_PROFILER_DEVICE:
@@ -194,7 +197,7 @@ class NVDriver(VirtDriver):
       params = nv_gpu.NVC36F_CTRL_CMD_GPFIFO_GET_WORK_SUBMIT_TOKEN_PARAMS.from_address(params_ptr)
       gpu_fifo = self.object_by_handle[struct.hObject]
       params.workSubmitToken = gpu_fifo.token
-    elif struct.cmd == nv_gpu.NVA06C_CTRL_CMD_GPFIFO_SCHEDULE: pass
+    elif struct.cmd in (nv_gpu.NVA06C_CTRL_CMD_GPFIFO_SCHEDULE, nv_gpu.NVA06F_CTRL_CMD_BIND, nv_gpu.NVA06F_CTRL_CMD_GPFIFO_SCHEDULE): pass
     elif struct.cmd == nv_gpu.NV2080_CTRL_CMD_PERF_BOOST: pass
     elif struct.cmd == nv_gpu.NV2080_CTRL_CMD_FB_FLUSH_GPU_CACHE: pass
     elif struct.cmd == nv_gpu.NV83DE_CTRL_CMD_DEBUG_READ_ALL_SM_ERROR_STATES:

test/mockgpu/nv/nvgpu.py

@@ -26,6 +26,7 @@ class GPFIFO:
     self.gpfifo = to_mv(self.base, self.entries_cnt * 8).cast("Q")
     self.ctrl = nv_gpu.AmpereAControlGPFifo.from_address(self.base + self.entries_cnt * 8)
     self.state = {}
+    self.bound_engines: set[int] = set()
 
     # Buf exec state
     self.buf = None
@@ -115,9 +116,11 @@ class GPFIFO:
   def execute_cmd(self, cmd) -> SchedResult:
     if cmd == nv_gpu.NVC56F_SEM_EXECUTE: return self._exec_signal()
     elif cmd == nv_gpu.NVC6C0_LAUNCH_DMA: return self._exec_nvc6c0_dma()
-    elif cmd == nv_gpu.NVC6B5_LAUNCH_DMA: return self._exec_nvc6b5_dma()
+    elif cmd == nv_gpu.NVC6B5_LAUNCH_DMA: # NOTE: NVC6B5_LAUNCH_DMA == NVC9B0_EXECUTE == 0x300
+      return self._exec_vid_decode() if self.bound_engines & {nv_gpu.NVC9B0_VIDEO_DECODER, nv_gpu.NVCFB0_VIDEO_DECODER} else self._exec_nvc6b5_dma()
     elif cmd == nv_gpu.NVC6C0_SEND_SIGNALING_PCAS2_B: return self._exec_pcas2()
     elif cmd == 0x0320: return self._exec_load_inline_qmd() # NVC6C0_LOAD_INLINE_QMD_DATA
+    elif cmd == nv_gpu.NVC9B0_SEMAPHORE_D: return self._exec_vid_semaphore()
     else: self.state[cmd] = self._next_dword() # just state update
     return SchedResult.CONT
 
@@ -136,6 +139,27 @@ class GPFIFO:
     else: raise RuntimeError(f"Unsupported type={typ} in exec wait/signal")
     return SchedResult.CONT
 
+  def _exec_vid_decode(self) -> SchedResult:
+    self._next_dword() # consume execute flags
+    # validate that all required decode state was set up correctly
+    assert self._state(nv_gpu.NVC9B0_SET_APPLICATION_ID) == nv_gpu.NVC9B0_SET_APPLICATION_ID_ID_HEVC
+    pic_desc_addr = self._state(nv_gpu.NVC9B0_SET_DRV_PIC_SETUP_OFFSET) << 8
+    pic = nv_gpu.nvdec_hevc_pic_s.from_address(pic_desc_addr)
+    assert pic.stream_len > 0 and pic.pic_width_in_luma_samples > 0 and pic.pic_height_in_luma_samples > 0
+    assert self._state(nv_gpu.NVC9B0_SET_IN_BUF_BASE_OFFSET) != 0
+    assert self._state(nv_gpu.NVC9B0_SET_COLOC_DATA_OFFSET) != 0
+    assert self._state(nv_gpu.NVC9B0_SET_NVDEC_STATUS_OFFSET) != 0
+    assert self._state(nv_gpu.NVC9B0_HEVC_SET_FILTER_BUFFER_OFFSET) != 0
+    return SchedResult.CONT
+
+  def _exec_vid_semaphore(self) -> SchedResult:
+    signal = self._state64(nv_gpu.NVC9B0_SEMAPHORE_A)
+    val = self._state(nv_gpu.NVC9B0_SEMAPHORE_C)
+    self._next_dword() # flags
+    to_mv(signal, 8).cast('Q')[0] = val
+    to_mv(signal + 8, 8).cast('Q')[0] = int(time.perf_counter() * 1e9)
+    return SchedResult.CONT
+
   def _exec_load_inline_qmd(self):
     qmd_addr = self._state64(nv_gpu.NVC6C0_SET_INLINE_QMD_ADDRESS_A) << 8
     assert qmd_addr != 0x0, f"invalid qmd address {qmd_addr}"

test/models/test_onnx.py

@@ -73,7 +73,7 @@ class TestHuggingFaceOnnxModels(unittest.TestCase):
     onnx_model_path = snapshot_download_with_retry(
       repo_id=repo_id,
       allow_patterns=["*.onnx", "*.onnx_data"],
-      cache_dir=str(DOWNLOADS_DIR)
+      local_dir=DOWNLOADS_DIR / repo_id
     )
     onnx_model_path = onnx_model_path / model_file
     file_size = onnx_model_path.stat().st_size

test/null/test_attention.py

@@ -0,0 +1,39 @@
+import unittest
+from tinygrad import Tensor, dtypes, TinyJit, UOp
+from tinygrad.apps.llm import apply_rope as apply_rope_new, precompute_freqs_cis
+
+def apply_rope(x:Tensor, start_pos:int):
+  B, H, T, Hd = x.shape
+  precompute_freqs_cis.cache_clear()
+  freqs_cis = precompute_freqs_cis(Hd, start_pos+T)[start_pos:start_pos+T]
+  return apply_rope_new(x, freqs_cis)
+
+class TestAttention(unittest.TestCase):
+  def test_half_qkv_buffers(self):
+    BS, seqlen, dim = 10, 4, 100
+    q = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
+    k = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
+    v = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
+    attn = q.scaled_dot_product_attention(k, v)
+    sched = attn.schedule()
+    # attention has 4 kernels now
+    self.assertEqual(len(sched), 4)
+
+  def test_apply_rope_jit_prune(self):
+    def rope_fn(x_in, pos): return apply_rope(x_in, pos)
+    rope_noprune = TinyJit(rope_fn)
+    rope_prune = TinyJit(rope_fn, prune=True)
+
+    v_pos = UOp.variable("start_pos", 0, 100)
+    for _ in range(3):
+      rope_noprune(Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32), v_pos.bind(1))
+      rope_prune(Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32), v_pos.bind(1))
+    noprune_size = len(rope_noprune.captured.jit_cache)
+    prune_size = len(rope_prune.captured.jit_cache)
+
+    self.assertGreater(noprune_size, prune_size)
+    self.assertGreaterEqual(noprune_size, 2)
+    self.assertEqual(prune_size, 1)
+
+if __name__ == '__main__':
+  unittest.main()

test/null/test_const_folding.py

@@ -0,0 +1,149 @@
+import unittest, itertools, math
+from tinygrad import Tensor, dtypes, Context
+from tinygrad.dtype import DType, ConstType
+from tinygrad.uop.ops import Ops, UOp
+from tinygrad.codegen import full_rewrite_to_sink
+import numpy as np
+
+def _check_ast_count(desired_count:int, t:Tensor):
+  # NOTE: this has side effect because everything can be scheduled only once
+  schedule = t.schedule()
+  asts = [s for s in schedule if s.ast.op is Ops.SINK]
+  assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
+
+class TestUnaryOpsConstFolding(unittest.TestCase):
+  def test_all_consts_ops(self):
+    _check_ast_count(0, Tensor.ones(4).exp())
+    _check_ast_count(0, Tensor.ones(4).sqrt())
+    _check_ast_count(0, Tensor.ones(4) + Tensor.ones(4))
+    _check_ast_count(0, Tensor.ones(4) / Tensor.ones(4))
+
+  def test_cast(self):
+    _check_ast_count(0, Tensor.ones(4).cast(dtypes.int16))
+    _check_ast_count(0, Tensor.full(4, fill_value=-1).cast(dtypes.uint16))
+
+  def test_neg_folding(self):
+    _check_ast_count(0, Tensor([1, 2, 3]).mul(-1).neg())
+    _check_ast_count(0, Tensor([1, 2, 3]).neg().mul(-1))
+    _check_ast_count(0, Tensor([1, 2, 3]).neg().neg())
+
+  def test_neg_realized_no_fold(self):
+    x = Tensor.randn(32, 32)
+    x = x.clip(0, 1).realize()
+    _check_ast_count(1, x.neg())
+
+class TestBinaryOpsConstFolding(unittest.TestCase):
+  def test_add_literal_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + 0)
+  def test_add_tensor_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(4))
+  def test_literal_zero_add(self):
+    _check_ast_count(0, 0 + Tensor([1.0, 2, 3, 4]))
+  def test_tensor_zero_add(self):
+    _check_ast_count(0, Tensor.zeros(4) + Tensor([1.0, 2, 3, 4]))
+
+  def test_sub_literal_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) - 0)
+  def test_sub_tensor_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) - Tensor.zeros(4))
+
+  def test_mul_literal_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * 0)
+  def test_mul_tensor_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.zeros(4))
+  def test_literal_zero_mul(self):
+    _check_ast_count(0, 0 * Tensor([1.0, 2, 3, 4]) * 0)
+  def test_tensor_zero_mul(self):
+    _check_ast_count(0, Tensor.zeros(4) * Tensor([1.0, 2, 3, 4]))
+
+  def test_mul_literal_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * 1)
+  def test_mul_tensor_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.ones(4))
+  def test_literal_one_mul(self):
+    _check_ast_count(0, 1 * Tensor([1.0, 2, 3, 4]))
+  def test_tensor_one_mul(self):
+    _check_ast_count(0, Tensor.ones(4) * Tensor([1.0, 2, 3, 4]))
+
+  def test_bool_tensor_mul_bool(self):
+    _check_ast_count(0, Tensor([True, False]) * True)
+    _check_ast_count(0, Tensor([True, False]) * False)
+  def test_bool_mul_bool_tensor(self):
+    _check_ast_count(0, True * Tensor([True, False]))
+    _check_ast_count(0, False * Tensor([True, False]))
+
+  def test_div_literal_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) / 1)
+  def test_div_tensor_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) / Tensor.ones(4))
+
+  def test_idiv_literal_one(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) // 1)
+  def test_idiv_tensor_one(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) // Tensor.ones(4, dtype=dtypes.int32))
+
+  def test_pow_literal_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** 0)
+  def test_pow_tensor_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** Tensor.zeros(4))
+
+  def test_pow_literal_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** 1)
+  def test_pow_tensor_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** Tensor.ones(4))
+  def test_literal_one_pow(self):
+    _check_ast_count(0, 1 ** Tensor([1.0, 2, 3, 4]))
+  def test_tensor_one_pow(self):
+    _check_ast_count(0, Tensor.ones(4) ** Tensor([1.0, 2, 3, 4]))
+
+class TestBitcastConstFolding(unittest.TestCase):
+  def test_scalar_bitcast(self):
+    def t(cases: dict[DType, ConstType]):
+      for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()):
+        if not math.isnan(from_v):
+          r = full_rewrite_to_sink(UOp.const(from_dt, from_v).bitcast(to_dt).sink()).src[0]
+          self.assertEqual(r.op, Ops.CONST, msg:=f"{from_dt} -> {to_dt} ({from_v} -> {to_v})")
+          self.assertEqual(r.dtype, to_dt, msg)
+          np.testing.assert_equal(r.arg, to_v, msg)
+
+    t({dtypes.int8: 0, dtypes.uint8: 0, dtypes.bool: False})
+    t({dtypes.int8: 1, dtypes.uint8: 1, dtypes.bool: True})
+
+    t({dtypes.int8:  -1, dtypes.uint8:  2**8-1})
+    t({dtypes.int16: -1, dtypes.uint16: 2**16-1, dtypes.float16: float('nan')})
+    t({dtypes.int32: -1, dtypes.uint32: 2**32-1, dtypes.float32: float('nan')})
+    t({dtypes.int64: -1, dtypes.uint64: 2**64-1, dtypes.float64: float('nan')})
+
+    t({dtypes.int8:  -2**7,  dtypes.uint8:  2**7})
+    t({dtypes.int16: -2**15, dtypes.uint16: 2**15})
+    t({dtypes.int32: -2**31, dtypes.uint32: 2**31})
+    t({dtypes.int64: -2**63, dtypes.uint64: 2**63})
+
+    t({dtypes.int16: 13496, dtypes.uint16: 13496, dtypes.float16: 0.294921875})
+    t({dtypes.int32: 1050081145, dtypes.uint32: 1050081145, dtypes.float32: 0.29485681653022766})
+    t({dtypes.int64: 4598983288165178391, dtypes.uint64: 4598983288165178391, dtypes.float64: 0.29485681936461233})
+
+  def test_vec_bitcast(self):
+    with Context(SPEC=0):
+      r = full_rewrite_to_sink(UOp.const(dtypes.int32.vec(3), (-1, -2**31, 75)).bitcast(dtypes.uint32.vec(3)).sink()).src[0]
+    self.assertEqual(r.op, Ops.VECTORIZE)
+    self.assertEqual(r.dtype, dtypes.uint32.vec(3))
+    self.assertEqual(tuple(x.arg for x in r.src), (2**32-1, 2**31, 75))
+
+# folds advance indexing into basic indexing
+class TestIndexingConstFolding(unittest.TestCase):
+  def test_scalar_index(self):
+    t = Tensor.arange(16).float().reshape(1,1,4,4).realize()
+    _check_ast_count(1, t[:,:,Tensor(1),:])
+    _check_ast_count(1, t[:,:,Tensor(1)+2,:])
+    _check_ast_count(1, t[:,:,Tensor(1),Tensor(0)])
+
+  def test_const_tensor_index(self):
+    # TODO: these can be 0, implement const tensor folded indexing
+    t = Tensor.arange(16).float().reshape(1,1,4,4).realize()
+    _check_ast_count(1, t[:,:,Tensor.ones(2,1,dtype=dtypes.int),:])
+    _check_ast_count(1, t[:,:,Tensor.ones(1,2,dtype=dtypes.int)+2,:])
+    _check_ast_count(1, t[:,:,Tensor.ones(1,1,dtype=dtypes.int),Tensor.zeros(2,1,2,dtype=dtypes.int)])
+
+if __name__ == '__main__':
+  unittest.main()

test/null/test_device.py

@@ -102,8 +102,8 @@ class TestCompiler(unittest.TestCase):
 
 class TestRunAsModule(unittest.TestCase):
   def test_module_runs(self):
-    out = '\n'.join(enumerate_devices_str())
-    self.assertIn("CPU", out) # for sanity check
+    cpu_line = [l for l in enumerate_devices_str() if "CPU" in l][0]
+    self.assertIn("PASS", cpu_line, f"expected CPU to PASS, got: {cpu_line}")
 
 if __name__ == "__main__":
   unittest.main()

test/null/test_dtype_spec.py

@@ -0,0 +1,413 @@
+import unittest, math, struct, operator
+from tinygrad.tensor import Tensor, dtypes
+from tinygrad.dtype import DTYPES_DICT, truncate, float_to_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float
+from tinygrad.device import is_dtype_supported
+
+from tinygrad.helpers import getenv
+from hypothesis import given, settings, strategies as strat
+import numpy as np
+import torch
+
+settings.register_profile("my_profile", max_examples=50, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
+settings.load_profile("my_profile")
+
+core_dtypes = list(DTYPES_DICT.values())
+dtype_ints = [dt for dt in core_dtypes if dtypes.is_int(dt) and is_dtype_supported(dt)]
+dtype_floats = [dt for dt in core_dtypes if dtypes.is_float(dt) and is_dtype_supported(dt)]
+
+FP8E4M3_MAX = 448.0
+FP8E5M2_MAX = 57344.0
+
+def u32_to_f32(u): return struct.unpack('f', struct.pack('I', u))[0]
+def f32_to_u32(f): return struct.unpack('I', struct.pack('f', f))[0]
+
+class TestHelpers(unittest.TestCase):
+  signed_ints = (dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64)
+  uints = (dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64)
+  floats = (dtypes.float16, dtypes.float32, dtypes.float64)
+
+  @given(strat.sampled_from(signed_ints+uints), strat.integers(min_value=1, max_value=8))
+  def test_is_int(self, dtype, amt):
+    assert dtypes.is_int(dtype.vec(amt) if amt > 1 else dtype)
+    assert not dtypes.is_float(dtype.vec(amt) if amt > 1 else dtype)
+
+  @given(strat.sampled_from(uints), strat.integers(min_value=1, max_value=8))
+  def test_is_unsigned_uints(self, dtype, amt):
+    assert dtypes.is_unsigned(dtype.vec(amt) if amt > 1 else dtype)
+
+  @given(strat.sampled_from(signed_ints), strat.integers(min_value=1, max_value=8))
+  def test_is_unsigned_signed_ints(self, dtype, amt):
+    assert not dtypes.is_unsigned(dtype.vec(amt) if amt > 1 else dtype)
+
+  @given(strat.sampled_from(floats), strat.integers(min_value=1, max_value=8))
+  def test_is_float(self, dtype, amt):
+    assert dtypes.is_float(dtype.vec(amt) if amt > 1 else dtype)
+    assert not dtypes.is_int(dtype.vec(amt) if amt > 1 else dtype)
+    assert not dtypes.is_unsigned(dtype.vec(amt) if amt > 1 else dtype)
+
+  def test_bf16_is_float(self):
+    assert dtypes.is_float(dtypes.bfloat16)
+
+  def test_fp8s_are_float(self):
+    assert dtypes.is_float(dtypes.fp8e4m3)
+    assert dtypes.is_float(dtypes.fp8e5m2)
+
+  @given(strat.sampled_from([d for d in DTYPES_DICT.values() if dtypes.is_float(d) or dtypes.is_int(d)]), strat.integers(min_value=2, max_value=8))
+  def test_scalar(self, dtype, amt):
+    assert dtype.vec(amt).scalar() == dtype
+
+  def test_from_py(self):
+    assert dtypes.from_py(True) == dtypes.bool
+    assert dtypes.from_py(2) == dtypes.default_int
+    assert dtypes.from_py(3.0) == dtypes.default_float
+    assert dtypes.from_py([]) == dtypes.default_float
+    assert dtypes.from_py(()) == dtypes.default_float
+    assert dtypes.from_py([True]) == dtypes.bool
+    assert dtypes.from_py([True, 2]) == dtypes.default_int
+    assert dtypes.from_py([True, 3.0]) == dtypes.default_float
+    assert dtypes.from_py([2, 3.0]) == dtypes.default_float
+    assert dtypes.from_py([True, 2, 3.0]) == dtypes.default_float
+    with self.assertRaises(RuntimeError): dtypes.from_py(None)
+    with self.assertRaises(RuntimeError): dtypes.from_py([None])
+    with self.assertRaises(RuntimeError): dtypes.from_py({})
+    with self.assertRaises(RuntimeError): dtypes.from_py(set())
+
+  def test_dtype_range(self):
+    for dt in core_dtypes:
+      if dtypes.is_float(dt):
+        np.testing.assert_equal(dtypes.min(dt), -math.inf)
+        np.testing.assert_equal(dtypes.max(dt), math.inf)
+        np.testing.assert_equal(dt.min, -math.inf)
+        np.testing.assert_equal(dt.max, math.inf)
+      elif dtypes.is_int(dt):
+        info = np.iinfo(_to_np_dtype(dt))
+        np.testing.assert_equal(dtypes.min(dt), info.min)
+        np.testing.assert_equal(dtypes.max(dt), info.max)
+        np.testing.assert_equal(dt.min, info.min)
+        np.testing.assert_equal(dt.max, info.max)
+      else:
+        assert dt == dtypes.bool, dt
+        np.testing.assert_equal(dtypes.min(dt), False)
+        np.testing.assert_equal(dtypes.max(dt), True)
+        np.testing.assert_equal(dt.min, False)
+        np.testing.assert_equal(dt.max, True)
+
+  def test_dtype_range_vec(self):
+    for dt in core_dtypes:
+      self.assertEqual(dt.min, dt.vec(4).min)
+      self.assertEqual(dt.max, dt.vec(4).max)
+
+  def test_float_to_fp16(self):
+    self.assertEqual(float_to_fp16(1), 1)
+    self.assertEqual(float_to_fp16(65504), 65504)
+    self.assertEqual(float_to_fp16(65519.999), 65504)
+    self.assertEqual(float_to_fp16(65520), math.inf)
+    self.assertEqual(float_to_fp16(1e-8), 0.0)
+    self.assertEqual(float_to_fp16(-65504), -65504)
+    self.assertEqual(float_to_fp16(-65519.999), -65504)
+    self.assertEqual(float_to_fp16(-65520), -math.inf)
+    self.assertTrue(math.isnan(float_to_fp16(math.nan)))
+
+  def test_float_to_bf16(self):
+    max_bf16 = torch.finfo(torch.bfloat16).max
+    for a in [1, 1.1, 1234, 23456, -777.777, max_bf16, max_bf16 * 1.00001, -max_bf16, -max_bf16 * 1.00001, math.inf, -math.inf]:
+      self.assertEqual(float_to_bf16(a), torch.tensor([a], dtype=torch.bfloat16).item())
+    self.assertTrue(math.isnan(float_to_bf16(math.nan)))
+
+  def test_float_to_bf16_nan(self):
+    patterns = [0x7FC00001, 0xFFC00001, 0x7F800001, 0xFF800001, 0x7FFFFFFF, 0xFFFFFFFF]
+    for u in patterns:
+      x = u32_to_f32(u)
+      y = float_to_bf16(x)
+      t = torch.tensor([x], dtype=torch.bfloat16).item()
+      self.assertTrue(math.isnan(y))
+      self.assertTrue(math.isnan(t))
+
+  def test_float_to_bf16_round(self):
+    uppers = [0x3f800000, 0x41230000, 0xC1460000]
+    for upper in uppers:
+      base = upper & 0xFFFF0000
+      base_f32 = u32_to_f32(base)
+      base_f32_round_up = u32_to_f32(base + 0x00010000)
+
+      x = u32_to_f32(base | 0x00007000)
+      self.assertEqual(float_to_bf16(x), base_f32)
+      self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32)
+
+      x = u32_to_f32(base | 0x0000C000)
+      self.assertEqual(float_to_bf16(x), base_f32_round_up)
+      self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32_round_up)
+
+      if ((upper >> 16) & 1) == 0:
+        x = u32_to_f32(base | 0x00008000)
+        self.assertEqual(float_to_bf16(x), base_f32)
+        self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32)
+      else:
+        x = u32_to_f32(base | 0x00008000)
+        self.assertEqual(float_to_bf16(x), base_f32_round_up)
+        self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32_round_up)
+
+  def test_float_to_bf16_boundary(self):
+    base = 0x7F7F0000
+    inf_u32 = 0x7F800000
+
+    x = u32_to_f32(base | 0x00007FFF)
+    self.assertEqual(f32_to_u32(float_to_bf16(x)), base)
+    self.assertEqual(f32_to_u32(torch.tensor([x], dtype=torch.bfloat16).item()), base)
+
+    x = u32_to_f32(base | 0x0000C000)
+    self.assertEqual(f32_to_u32(float_to_bf16(x)), inf_u32)
+    self.assertEqual(f32_to_u32(torch.tensor([x], dtype=torch.bfloat16).item()), inf_u32)
+
+    x = u32_to_f32(base | 0x00008000)
+    self.assertEqual(f32_to_u32(float_to_bf16(x)), inf_u32)
+    self.assertEqual(f32_to_u32(torch.tensor([x], dtype=torch.bfloat16).item()), inf_u32)
+
+  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=True, allow_infinity=True))
+  def test_truncate_fp8e4m3(self, x):
+    if math.isnan(x): np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), x)
+    elif math.isinf(x): np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), math.copysign(math.nan, x))
+    elif x > FP8E4M3_MAX: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), FP8E4M3_MAX)
+    elif x < -FP8E4M3_MAX: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), -FP8E4M3_MAX)
+    else: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), torch.tensor(x, dtype=torch.float8_e4m3fn).float().item())
+
+  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=True, allow_infinity=True))
+  def test_truncate_fp8e5m2(self, x):
+    if math.isnan(x): np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), x)
+    elif math.isinf(x): np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), x)
+    elif x > FP8E5M2_MAX: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), FP8E5M2_MAX)
+    elif x < -FP8E5M2_MAX: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), -FP8E5M2_MAX)
+    else: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), torch.tensor(x, dtype=torch.float8_e5m2).float().item())
+
+class TestTypePromotion(unittest.TestCase):
+  @given(strat.sampled_from(core_dtypes))
+  def test_self_promo_to_self(self, dtype):
+    assert least_upper_dtype(dtype) == dtype
+    assert least_upper_dtype(dtype, dtype) == dtype
+    assert least_upper_dtype(dtype, dtype, dtype) == dtype
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
+  def test_promo_resulted_higher_than_inputs(self, dtype1, dtype2):
+    result = least_upper_dtype(dtype1, dtype2)
+    assert not (result < dtype1) and not (result < dtype2)
+
+  def test_dtype_promo(self):
+    assert least_upper_dtype(dtypes.bool, dtypes.int8) == dtypes.int8
+    assert least_upper_dtype(dtypes.int8, dtypes.uint8) == dtypes.int16
+    assert least_upper_dtype(dtypes.uint8, dtypes.int16) == dtypes.int16
+    assert least_upper_dtype(dtypes.int16, dtypes.uint16) == dtypes.int32
+    assert least_upper_dtype(dtypes.uint16, dtypes.int32) == dtypes.int32
+    assert least_upper_dtype(dtypes.int32, dtypes.uint32) == dtypes.int64
+    assert least_upper_dtype(dtypes.uint32, dtypes.int64) == dtypes.int64
+    assert least_upper_dtype(dtypes.int64, dtypes.uint64) == dtypes.uint64
+    assert least_upper_dtype(dtypes.float16, dtypes.float32) == dtypes.float32
+    assert least_upper_dtype(dtypes.float32, dtypes.float64) == dtypes.float64
+
+    assert least_upper_dtype(dtypes.bool, dtypes.float32) == dtypes.float32
+    assert least_upper_dtype(dtypes.bool, dtypes.float64) == dtypes.float64
+    assert least_upper_dtype(dtypes.float16, dtypes.int64) == dtypes.float16
+    assert least_upper_dtype(dtypes.float16, dtypes.uint64) == dtypes.float16
+    assert least_upper_dtype(dtypes.fp8e4m3, dtypes.fp8e5m2) == dtypes.half
+    assert least_upper_dtype(dtypes.fp8e4m3, dtypes.bfloat16) == dtypes.bfloat16
+    assert least_upper_dtype(dtypes.fp8e5m2, dtypes.bfloat16) == dtypes.bfloat16
+    assert least_upper_dtype(dtypes.fp8e4m3, dtypes.float16) == dtypes.float16
+    assert least_upper_dtype(dtypes.fp8e5m2, dtypes.float16) == dtypes.float16
+    assert least_upper_dtype(dtypes.fp8e4m3, dtypes.int64) == dtypes.fp8e4m3
+    assert least_upper_dtype(dtypes.fp8e4m3, dtypes.uint64) == dtypes.fp8e4m3
+    assert least_upper_dtype(dtypes.fp8e5m2, dtypes.int64) == dtypes.fp8e5m2
+    assert least_upper_dtype(dtypes.fp8e5m2, dtypes.uint64) == dtypes.fp8e5m2
+
+class TestTypeSpec(unittest.TestCase):
+  def setUp(self):
+    self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float
+  def tearDown(self):
+    dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
+
+  def test_set_dtype_default(self):
+    for default_int in [dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64]:
+      dtypes.default_int = default_int
+      assert dtypes.default_int == default_int
+
+    for default_float in [*dtypes.fp8s, dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]:
+      dtypes.default_float = default_float
+      assert dtypes.default_float == default_float
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from([operator.gt, operator.ge, operator.le, operator.lt, operator.eq, operator.ne]))
+  def test_bool_ops(self, dtype, op):
+    assert op(Tensor.ones(4, 4, dtype=dtype), Tensor.ones(4, 4, dtype=dtype)).dtype == dtypes.bool
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
+  def test_functions_return_index(self, dtype, default_int, default_float):
+    dtypes.default_int, dtypes.default_float = default_int, default_float
+    assert Tensor([0, 1], dtype=dtype).argmax().dtype == dtypes.int32
+    assert Tensor([0, 1], dtype=dtype).argmin().dtype == dtypes.int32
+    assert Tensor([0, 1], dtype=dtype).multinomial().dtype == dtypes.int32
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints))
+  def test_tensor_indexing_returns_same_dtype(self, data_dtype, indices_dtype):
+    X_data =  Tensor.ones(60000, 1, 28, 28, dtype=data_dtype)
+    indices =  Tensor.randint(512, high=X_data.shape[0]).cast(indices_dtype)
+    assert X_data[indices].dtype == X_data.dtype
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints))
+  def test_gather_returns_same_dtype(self, data_dtype, indices_dtype):
+    X_data = Tensor([[1, 0], [0, 1]], dtype=data_dtype)
+    indices = Tensor([[0, 0], [1, 0]], dtype=indices_dtype)
+    assert X_data.gather(0, indices).dtype == X_data.dtype
+    assert X_data.gather(1, indices).dtype == X_data.dtype
+
+  @given(strat.sampled_from(dtype_floats), strat.sampled_from(dtype_floats))
+  def test_attention_returns_same_dtype(self, data_dtype, default_float):
+    dtypes.default_float = default_float
+    query = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
+    key = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
+    value = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
+    mask = (Tensor.rand(32, 8, 128, 128) < 0.5)
+    assert query.scaled_dot_product_attention(key, value, is_causal=True).dtype == data_dtype
+    assert query.scaled_dot_product_attention(key, value, is_causal=True, dropout_p=0.3).dtype == data_dtype
+    assert query.scaled_dot_product_attention(key, value, is_causal=False).dtype == data_dtype
+    assert query.scaled_dot_product_attention(key, value, attn_mask=mask).dtype == data_dtype
+
+class TestAutoCastType(unittest.TestCase):
+  def setUp(self):
+    self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float
+  def tearDown(self):
+    dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
+
+  @given(strat.sampled_from(dtype_floats), strat.sampled_from(dtype_floats))
+  def test_least_upper_float_input_is_float(self, input_dtype, default_float):
+    dtypes.default_float = default_float
+    self.assertEqual(least_upper_float(input_dtype), input_dtype)
+
+  @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
+  def test_least_upper_float_input_is_int(self, input_dtype, default_float):
+    dtypes.default_float = default_float
+    self.assertEqual(least_upper_float(input_dtype), default_float)
+
+  @given(strat.sampled_from(core_dtypes))
+  def test_broadcast_scalar(self, dt):
+    assert (Tensor.ones(4, 4, dtype=dt) + 2.3).dtype == (dt if dtypes.is_float(dt) else dtypes.default_float)
+    assert (Tensor.ones(4, 4, dtype=dt) + 2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int)
+    assert (Tensor.ones(4, 4, dtype=dt) + True).dtype == dt
+
+  @given(strat.sampled_from(dtype_floats))
+  def test_int_div_int(self, default_float):
+    dtypes.default_float = default_float
+    self.assertEqual(Tensor([1]).div(Tensor([2])).dtype, default_float)
+
+  def test_sum(self):
+    assert (Tensor([0, 1], dtype=dtypes.bool)).sum().dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int8)).sum().dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int16)).sum().dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int32)).sum().dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int64)).sum().dtype == dtypes.int64
+    assert (Tensor([0, 1], dtype=dtypes.uint8)).sum().dtype == dtypes.uint32
+    assert (Tensor([0, 1], dtype=dtypes.uint16)).sum().dtype == dtypes.uint32
+    assert (Tensor([0, 1], dtype=dtypes.uint32)).sum().dtype == dtypes.uint32
+    assert (Tensor([0, 1], dtype=dtypes.uint64)).sum().dtype == dtypes.uint64
+    assert (Tensor([0, 1], dtype=dtypes.fp8e4m3)).sum().dtype == dtypes.fp8e4m3
+    assert (Tensor([0, 1], dtype=dtypes.fp8e5m2)).sum().dtype == dtypes.fp8e5m2
+    assert (Tensor([0, 1], dtype=dtypes.float16)).sum().dtype == dtypes.float16
+    assert (Tensor([0, 1], dtype=dtypes.bfloat16)).sum().dtype == dtypes.bfloat16
+    assert (Tensor([0, 1], dtype=dtypes.float32)).sum().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.float64)).sum().dtype == dtypes.float64
+
+  def test_mean(self):
+    assert (Tensor([0, 1], dtype=dtypes.bool)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.int8)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.int16)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.int32)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.int64)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.uint8)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.uint16)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.uint32)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.uint64)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.fp8e4m3)).mean().dtype == dtypes.fp8e4m3
+    assert (Tensor([0, 1], dtype=dtypes.fp8e5m2)).mean().dtype == dtypes.fp8e5m2
+    assert (Tensor([0, 1], dtype=dtypes.float16)).mean().dtype == dtypes.float16
+    assert (Tensor([0, 1], dtype=dtypes.bfloat16)).mean().dtype == dtypes.bfloat16
+    assert (Tensor([0, 1], dtype=dtypes.float32)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.float64)).mean().dtype == dtypes.float64
+
+  def test_cumsum(self):
+    assert (Tensor([0, 1], dtype=dtypes.bool)).cumsum(0).dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int8)).cumsum(0).dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int16)).cumsum(0).dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int32)).cumsum(0).dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int64)).cumsum(0).dtype == dtypes.int64
+    assert (Tensor([0, 1], dtype=dtypes.uint8)).cumsum(0).dtype == dtypes.uint32
+    assert (Tensor([0, 1], dtype=dtypes.uint16)).cumsum(0).dtype == dtypes.uint32
+    assert (Tensor([0, 1], dtype=dtypes.uint32)).cumsum(0).dtype == dtypes.uint32
+    assert (Tensor([0, 1], dtype=dtypes.uint64)).cumsum(0).dtype == dtypes.uint64
+    assert (Tensor([0, 1], dtype=dtypes.fp8e4m3)).cumsum(0).dtype == dtypes.fp8e4m3
+    assert (Tensor([0, 1], dtype=dtypes.fp8e5m2)).cumsum(0).dtype == dtypes.fp8e5m2
+    assert (Tensor([0, 1], dtype=dtypes.float16)).cumsum(0).dtype == dtypes.float16
+    assert (Tensor([0, 1], dtype=dtypes.bfloat16)).cumsum(0).dtype == dtypes.bfloat16
+    assert (Tensor([0, 1], dtype=dtypes.float32)).cumsum(0).dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.float64)).cumsum(0).dtype == dtypes.float64
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
+  def test_matmul(self, dt1, dt2, acc_dt):
+    t1 = Tensor([0, 1], dtype=dt1)
+    t2 = Tensor([0, 1], dtype=dt2)
+    self.assertEqual(t1.matmul(t2).dtype, least_upper_dtype(t1.dtype, t2.dtype))
+    self.assertEqual(t1.matmul(t2, dtype=acc_dt).dtype, acc_dt)
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
+  def test_linear(self, dt1, dt2, dt3, acc_dt):
+    x = Tensor([0, 1], dtype=dt1)
+    w = Tensor([0, 1], dtype=dt2)
+    b = Tensor([0, 1], dtype=dt3)
+    self.assertEqual(x.linear(w).dtype, least_upper_dtype(x.dtype, w.dtype))
+    self.assertEqual(x.linear(w, b).dtype, least_upper_dtype(least_upper_dtype(x.dtype, w.dtype), b.dtype))
+    self.assertEqual(x.linear(w, dtype=acc_dt).dtype, acc_dt)
+    self.assertEqual(x.linear(w, b, dtype=acc_dt).dtype, acc_dt)
+
+  @staticmethod
+  def check_where_alternate_input_other(input_, other, data_type):
+    assert (Tensor([True, False]).where(input_, other)).dtype == data_type
+    assert (Tensor([True, False]).where(other, input_)).dtype == data_type
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
+  def test_where_no_scalar(self, dt1, dt2):
+    self.check_where_alternate_input_other(Tensor(2, dtype=dt1), Tensor(3, dtype=dt2), least_upper_dtype(dt1, dt2))
+
+  @given(strat.sampled_from(core_dtypes))
+  def test_where_one_scalar(self, dt):
+    t = Tensor(2, dtype=dt)
+    self.check_where_alternate_input_other(t, 3.2, (dt if dtypes.is_float(dt) else dtypes.default_float))
+    self.check_where_alternate_input_other(t, 3, (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int))
+    self.check_where_alternate_input_other(t, True, dt)
+
+  def test_where_two_scalars(self):
+    self.check_where_alternate_input_other(3.1, 3.2, dtypes.default_float)
+    self.check_where_alternate_input_other(3.1, 3, dtypes.default_float)
+    self.check_where_alternate_input_other(3.1, True, dtypes.default_float)
+    self.check_where_alternate_input_other(3, 2, dtypes.default_int)
+    self.check_where_alternate_input_other(3, True, dtypes.default_int)
+    self.check_where_alternate_input_other(False, True, dtypes.bool)
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
+  def test_maximum(self, dt1, dt2):
+    assert Tensor([0, 1, 2], dtype=dt1).maximum(Tensor([2, 0, 5], dtype=dt2)).dtype == least_upper_dtype(dt1, dt2)
+
+  @given(strat.sampled_from(core_dtypes))
+  def test_maximum_const(self, dt):
+    assert Tensor([1, 2], dtype=dt).maximum(3.1).dtype == (dt if dtypes.is_float(dt) else dtypes.default_float)
+    assert Tensor([1, 2], dtype=dt).maximum(3).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int)
+    assert Tensor([1, 2], dtype=dt).maximum(True).dtype == dt
+
+  def test_div(self):
+    assert (Tensor([1, 2], dtype=dtypes.int32) / Tensor([2, 2], dtype=dtypes.int32)).dtype == dtypes.default_float
+    assert (Tensor([1, 2], dtype=dtypes.int16) / Tensor([2, 2], dtype=dtypes.int32)).dtype == dtypes.default_float
+    assert (Tensor([1, 2], dtype=dtypes.float32) / Tensor([2, 2], dtype=dtypes.float16)).dtype == dtypes.float32
+    assert (Tensor([1, 2], dtype=dtypes.int32) / Tensor([2, 2], dtype=dtypes.float16)).dtype == dtypes.float16
+
+  def test_div_const(self):
+    assert (Tensor([1, 2], dtype=dtypes.int32) / 2).dtype == dtypes.default_float
+    assert (Tensor([1, 2], dtype=dtypes.int32) / 2.0).dtype == dtypes.default_float
+    assert (Tensor([1, 2], dtype=dtypes.float16) / 2).dtype == dtypes.float16
+    assert (Tensor([1, 2], dtype=dtypes.float16) / 2.0).dtype == dtypes.float16
+
+if __name__ == '__main__':
+  unittest.main()

test/null/test_gc.py

@@ -72,7 +72,6 @@ class TestGC(unittest.TestCase):
     ys = y.schedule()
     del x
     run_schedule(ys)
-    np.testing.assert_equal(y.numpy(), np.full((256,), 2))
     self.assertEqual(bufs_allocated()-init, 1)
     del y
     self.assertEqual(bufs_allocated()-init, 0)

test/null/test_gradient.py

@@ -0,0 +1,83 @@
+from typing import Callable
+import unittest, math
+import torch
+from tinygrad import Tensor
+from tinygrad.dtype import dtypes
+from tinygrad.uop.ops import UOp
+from tinygrad.gradient import compute_gradient
+
+class TestGradient(unittest.TestCase):
+  def _cmp_nan_okay(self, x, y):
+    if math.isnan(x) and math.isnan(y): return
+    self.assertAlmostEqual(x, y, places=5)
+
+  def _test_one_input_function(self, f:Callable, jf:Callable|None=None):
+    if jf is None: jf = f
+    x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float)
+    gx = compute_gradient(f(x), UOp.const(dtypes.float, 1.0), set([x]))[x]
+
+    for val in [-5., -2.0, 0.0, 2.0, 5.]:
+      tg_out = gx.substitute({x: x.const_like(val)}).ssimplify()
+      tx = torch.tensor([val], dtype=torch.float, requires_grad=True)
+      torch_out = torch.autograd.grad(jf(tx), tx)[0].item()
+      self._cmp_nan_okay(tg_out, torch_out)
+
+  def _test_two_input_function(self, f:Callable, jf:Callable|None=None):
+    if jf is None: jf = f
+    x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float)
+    y = UOp.variable('y', -math.inf, math.inf, dtype=dtypes.float)
+    grads = compute_gradient(f(x, y), UOp.const(dtypes.float, 1.0), set([x, y]))
+    gx, gy = grads[x], grads[y]
+
+    for valx in [-5., -2.0, 0.0, 2.0, 5.]:
+      for valy in [-5., -2.0, 0.0, 2.0, 5.]:
+        # Substitute the values into the gradient expressions
+        substitutions = {x: x.const_like(valx), y: y.const_like(valy)}
+        tg_out_x = gx.substitute(substitutions).ssimplify()
+        tg_out_y = gy.substitute(substitutions).ssimplify()
+
+        tx = torch.tensor([valx], dtype=torch.float, requires_grad=True)
+        ty = torch.tensor([valy], dtype=torch.float, requires_grad=True)
+        torch_grad = torch.autograd.grad(jf(tx, ty), [tx, ty])
+        torch_out_x, torch_out_y = [x.item() for x in torch_grad]
+
+        self._cmp_nan_okay(tg_out_x, torch_out_x)
+        self._cmp_nan_okay(tg_out_y, torch_out_y)
+
+  # unary ops unit
+  def test_recip(self): self._test_one_input_function(lambda x: 1.0/x)
+  def test_sin(self): self._test_one_input_function(lambda x: x.sin())
+  def test_sqrt(self): self._test_one_input_function(lambda x: x.sqrt())
+  def test_log2(self): self._test_one_input_function(lambda x: x.log2())
+  def test_exp2(self): self._test_one_input_function(lambda x: x.exp2())
+
+  # binary ops unit
+  def test_add(self): self._test_two_input_function(lambda x,y: x+y)
+  def test_mul(self): self._test_two_input_function(lambda x,y: x*y)
+
+  # chain rule
+  def test_chain(self): self._test_one_input_function(lambda x: x.sin().sqrt())
+  def test_chain_binop(self): self._test_two_input_function(lambda x,y: (x*y)+x*y)
+  def test_big_add_sin(self): self._test_two_input_function(lambda x,y: x.sin()+3.0/y)
+  def test_big_chain(self): self._test_two_input_function(lambda x,y: (1.0/x*y)+x*y)
+  def test_where(self): self._test_two_input_function(lambda x,y: (x<y).where(x,y), lambda x,y: torch.where(x<y,x,y))
+
+class TestRealizeMeansRealize(unittest.TestCase):
+  def test_randn_realizes(self):
+    x = Tensor.randn(2, 3, 64, 64, requires_grad=True).realize()
+    assert x.uop is not x.uop.base
+    assert x.uop.is_realized
+
+  def test_uniform_realizes(self):
+    x = Tensor.uniform(16, 3, 3, 3, requires_grad=True).realize()
+    print(x.uop)
+    assert x.uop is not x.uop.base
+    assert x.uop.is_realized
+
+  def test_uniform_gradient(self):
+    x = Tensor.uniform(16, 3, 3, 3, requires_grad=True).realize()
+    y = x * 2
+    y.sum().gradient(x)[0].realize()
+
+if __name__ == '__main__':
+  unittest.main()

test/null/test_helpers.py

@@ -0,0 +1,469 @@
+import ctypes, gzip, unittest, timeit, pickle
+from tinygrad import Variable
+from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, get_contraction, count, all_same
+from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
+from tinygrad.helpers import ceildiv
+from tinygrad.tensor import Tensor, get_shape
+import numpy as np
+
+VARIABLE = ContextVar("VARIABLE", 0)
+
+class TestContextVars(unittest.TestCase):
+  # Ensuring that the test does not modify variables outside the tests.
+  ctx = Context()
+  def setUp(self): TestContextVars.ctx.__enter__()
+  def tearDown(self): TestContextVars.ctx.__exit__()
+
+  def test_initial_value_is_set(self):
+    _TMP = ContextVar("_TMP", 5)
+    self.assertEqual(_TMP.value, 5)
+
+  def test_cannot_recreate(self):
+    _TMP2 = ContextVar("_TMP2", 1)
+    with self.assertRaises(RuntimeError):
+      _TMP2 = ContextVar("_TMP2", 2)
+
+  def test_new_var_inside_context(self):
+    with Context(VARIABLE=1):
+      _TMP3 = ContextVar("_TMP3", 1)
+    with self.assertRaises(RuntimeError):
+      _TMP3 = ContextVar("_TMP3", 2)
+
+  def test_value_across_modules(self):
+    # Mocking module import by invoking the code but not in our globals().
+    exec('from tinygrad.helpers import ContextVar;C = ContextVar("C", 13)', {}) # pylint:disable=exec-used
+    # It should not matter that the first creation was in another module.
+    with self.assertRaises(RuntimeError):
+      _C = ContextVar("C", 0)
+
+  def test_assignment_across_modules(self):
+    B = ContextVar("B", 1)
+    # local assignment
+    B.value = 2
+    self.assertEqual(B.value, 2)
+    with self.assertRaises(RuntimeError):
+      # Assignment in another module.
+      exec('from tinygrad.helpers import ContextVar;B = ContextVar("B", 0);B.value = 3;', {}) # pylint:disable=exec-used
+
+  def test_context_assignment(self):
+    with Context(VARIABLE=1):
+      self.assertEqual(VARIABLE.value, 1)
+    self.assertEqual(VARIABLE.value, 0)
+
+  def test_unknown_param_to_context(self):
+    with self.assertRaises(KeyError):
+      with Context(SOMETHING_ELSE=1):
+        pass
+
+  def test_nested_context(self):
+    with Context(VARIABLE=1):
+      with Context(VARIABLE=2):
+        MORE = ContextVar("MORE", 2)
+        with Context(VARIABLE=3, MORE=3):
+          self.assertEqual(VARIABLE.value, 3)
+          self.assertEqual(MORE.value, 3)
+        self.assertEqual(VARIABLE.value, 2)
+        self.assertEqual(MORE.value, 2)
+      self.assertEqual(VARIABLE.value, 1)
+      self.assertEqual(MORE.value, 2)  # TODO: should this raise?
+    self.assertEqual(VARIABLE.value, 0)
+
+  def test_decorator(self):
+    @Context(VARIABLE=1, DEBUG=4)
+    def test():
+      self.assertEqual(VARIABLE.value, 1)
+
+    self.assertEqual(VARIABLE.value, 0)
+    test()
+    self.assertEqual(VARIABLE.value, 0)
+
+  def test_context_exit_reverts_updated_values(self):
+    D = ContextVar("D", 1)
+    D.value = 2
+    with Context(D=3):
+      ...
+    assert D.value == 2, f"Expected D to be 2, but was {D.value}. Indicates that Context.__exit__ did not restore to the correct value."
+
+class TestAllSame(unittest.TestCase):
+  def test_empty(self): self.assertTrue(all_same([]))
+  def test_single(self): self.assertTrue(all_same([1]))
+  def test_same(self): self.assertTrue(all_same([1, 1, 1]))
+  def test_different(self): self.assertFalse(all_same([1, 2, 1]))
+
+class TestMergeDicts(unittest.TestCase):
+  def test_merge_dicts(self):
+    a = {"a": 1, "b": 2}
+    b = {"a": 1, "c": 3}
+    c = {}
+    d = {"a": 2, "b": 2}
+    assert merge_dicts([a, b]) == {"a": 1, "b": 2, "c": 3}
+    assert merge_dicts([a, c]) == a
+    assert merge_dicts([a, b, c]) == {"a": 1, "b": 2, "c": 3}
+    with self.assertRaises(RuntimeError):
+      merge_dicts([a, d])
+
+class TestStripParens(unittest.TestCase):
+  def test_simple(self): self.assertEqual("1+2", strip_parens("(1+2)"))
+  def test_nested(self): self.assertEqual("1+(2+3)", strip_parens("(1+(2+3))"))
+  def test_casted_no_strip(self): self.assertEqual("(int)(1+2)", strip_parens("(int)(1+2)"))
+  def test_unmatched_parens(self): self.assertEqual("((c35+c39>>23&255)+-127).cast(dtypes.float)",
+    strip_parens("((c35+c39>>23&255)+-127).cast(dtypes.float)"))
+  def test_single_paren_left(self): self.assertEqual("(abc", strip_parens("(abc"))
+  def test_single_paren_right(self): self.assertEqual("abc)", strip_parens("abc)"))
+  def test_parens_at_different_depths(self): self.assertEqual("(a+(b))*(c)", strip_parens("(a+(b))*(c)"))
+
+class TestProd(unittest.TestCase):
+  def test_empty(self): self.assertEqual(1, prod(tuple()))
+  def test_ints(self): self.assertEqual(30, prod((2, 3, 5)))
+  def test_variable(self): self.assertEqual("(a*12)", prod((Variable("a", 1, 5), 3, 4)).render())
+  def test_variable_order(self): self.assertEqual("(a*12)", prod((3, 4, Variable("a", 1, 5))).render())
+
+class TestRoundUp(unittest.TestCase):
+  def test_round_up(self):
+    self.assertEqual(round_up(-3,4), 0)
+    self.assertEqual(round_up(-4,4), -4)
+    self.assertEqual(round_up(6,4), 8)
+    self.assertEqual(round_up(8,4), 8)
+    self.assertEqual(round_up(232, 24984), 24984)
+    self.assertEqual(round_up(24984, 232), 25056)
+
+class TestCeilDiv(unittest.TestCase):
+  def test_int(self):
+    self.assertEqual(ceildiv(10, 3), 4)
+    self.assertEqual(ceildiv(9, 3), 3)
+    self.assertEqual(ceildiv(0, 5), 0)
+    self.assertEqual(ceildiv(1, 5), 1)
+  def test_symbolic(self):
+    # tests that ceildiv with UOp uses (num + amt - 1) // amt formula for non-negative num
+    v = Variable('v', 0, 100)
+    result = ceildiv(v, 6)
+    self.assertEqual(result.render(), "((v+5)//6)")
+  def test_symbolic_negative_offset(self):
+    # tests ceildiv(v-5, 6) which is used in conv2d output shape
+    # old implementation incorrectly simplified -(x//-y) to ((v+1)//6-1) for v-5
+    # new implementation uses (v-5+5)//6 = v//6 which is correct
+    v = Variable('v', 11, 100)
+    result = ceildiv(v - 5, 6)
+    self.assertEqual(result.render(), "(v//6)")
+
+class TestCount(unittest.TestCase):
+  def test_count_basic(self):
+    c = count(3)
+    self.assertEqual(next(c), 3)
+    self.assertEqual(next(c), 4)
+
+  def test_count_step_pickle(self):
+    c = count(1, 2)
+    self.assertEqual(next(c), 1)
+    c2 = pickle.loads(pickle.dumps(c))
+    self.assertEqual(next(c2), 3)
+
+@unittest.skip("no fetch tests because they need internet")
+class TestFetch(unittest.TestCase):
+  def test_fetch_bad_http(self):
+    self.assertRaises(Exception, fetch, 'http://www.google.com/404', allow_caching=False)
+
+  def test_fetch_small(self):
+    assert (len(fetch('https://google.com', allow_caching=False).read_bytes())>0)
+
+  def test_fetch_img(self):
+    from PIL import Image
+    img = fetch("https://avatars.githubusercontent.com/u/132956020", allow_caching=False)
+    with Image.open(img) as pimg:
+      assert pimg.size == (77, 77), pimg.size
+
+  def test_fetch_subdir(self):
+    from PIL import Image
+    img = fetch("https://avatars.githubusercontent.com/u/132956020", allow_caching=False, subdir="images")
+    with Image.open(img) as pimg:
+      assert pimg.size == (77, 77), pimg.size
+    assert img.parent.name == "images"
+
+  def test_fetch_gunzip_valid(self):
+    # compare fetch(gunzip=True) to fetch(gunzip=False) plus decompressing afterwards
+    gzip_url: str = 'https://ftp.gnu.org/gnu/gzip/gzip-1.13.tar.gz'
+    fp_gz = fetch(gzip_url, gunzip=True)
+    fp_no_gz = fetch(gzip_url, gunzip=False)
+    with open(fp_gz, 'rb') as f: content_gz = f.read()
+    with open(fp_no_gz, 'rb') as f: content_no_gz = gzip.decompress(f.read())
+    assert fp_gz.stat().st_size > fp_no_gz.stat().st_size
+    assert isinstance(content_gz, bytes) and isinstance(content_no_gz, bytes)
+    assert len(content_gz) == len(content_no_gz)
+    assert content_gz == content_no_gz
+
+  def test_fetch_gunzip_invalid(self):
+    # given a non-gzipped file, fetch(gunzip=True) fails
+    no_gzip_url: str = 'https://ftp.gnu.org/gnu/gzip/gzip-1.13.zip'
+    with self.assertRaises(gzip.BadGzipFile):
+      fetch(no_gzip_url, gunzip=True)
+
+  def test_fetch_user_agent(self):
+    fetch("https://csrc.nist.gov/CSRC/media/Projects/lightweight-cryptography/documents/finalist-round/updated-submissions/sparkle.zip",
+          allow_caching=False)
+
+  def test_fetch_half_and_full_file(self):
+    x = fetch("https://csrc.nist.gov/CSRC/media/Projects/lightweight-cryptography/documents/finalist-round/updated-submissions/sparkle.zip",
+          headers={"Range": "bytes=0-10"}).read_bytes()
+    assert len(x) == 11, f"{len(x) != 11}"
+    x = fetch("https://csrc.nist.gov/CSRC/media/Projects/lightweight-cryptography/documents/finalist-round/updated-submissions/sparkle.zip",
+          headers={"Range": "bytes=0-100"}).read_bytes()
+    assert len(x) == 101, f"{len(x) != 101}"
+
+class TestFullyFlatten(unittest.TestCase):
+  def test_fully_flatten(self):
+    self.assertEqual(fully_flatten([[1, 3], [1, 2]]), [1, 3, 1, 2])
+    self.assertEqual(fully_flatten(((1, 3), (1, 2))), [1, 3, 1, 2])
+    self.assertEqual(fully_flatten([[[1], [3]], [[1], [2]]]), [1, 3, 1, 2])
+    self.assertEqual(fully_flatten([[[[1], 2], 3], 4]), [1, 2, 3, 4])
+    self.assertEqual(fully_flatten([[1, 2, [3, 4]], [5, 6], 7]), [1, 2, 3, 4, 5, 6, 7])
+    self.assertEqual(fully_flatten([[1, "ab"], [True, None], [3.14, [5, "b"]]]), [1, "ab", True, None, 3.14, 5, "b"])
+
+  def test_fully_flatten_numpy(self):
+    self.assertEqual(fully_flatten([np.array([])]), [])
+    self.assertEqual(fully_flatten([np.array(3)]), [3])
+    self.assertEqual(fully_flatten([np.array([3])]), [3])
+    self.assertEqual(fully_flatten([np.array([[3]])]), [3])
+    self.assertEqual(fully_flatten([np.array([1, 3]), np.array([1, 2])]), [1, 3, 1, 2])
+    self.assertEqual(fully_flatten((np.array([1, 3]), np.array([1, 2]))), [1, 3, 1, 2])
+    self.assertEqual(fully_flatten([np.array([[1], [3]]), np.array([[1], [2]])]), [1, 3, 1, 2])
+    self.assertEqual(fully_flatten([[1, "ab"], [True, None], np.array([[3.14], [6.28]])]), [1, "ab", True, None, 3.14, 6.28])
+
+class TestMemoryview(unittest.TestCase):
+  def test_from_mv_to_mv(self):
+    base = memoryview(bytearray(b"\x11\x22\x33"*40))
+    ct = from_mv(base)
+    mv = to_mv(ctypes.addressof(ct), len(base))
+    mv[0] = 2
+    assert base[0] == 2
+
+  @unittest.skip("allocates tons of memory")
+  def test_to_mv(self):
+    sizes = [
+      (16, "16 B"),
+      (64, "64 B"),
+      (256, "256 B"),
+      (1024, "1 KB"),
+      (4 * 1024, "4 KB"),
+      (16 * 1024, "16 KB"),
+      (64 * 1024, "64 KB"),
+      (256 * 1024, "256 KB"),
+      (1 * 1024 * 1024, "1 MB"),
+      (10 * 1024 * 1024, "10 MB"),
+      (200 * 1024 * 1024, "200 MB"),
+    ]
+
+    for sz, label in sizes:
+      buf = np.random.randint(0, 256, sz, dtype=np.uint8)
+      ptr = buf.ctypes.data
+
+      iters = 100_000
+      t_us = timeit.timeit(lambda: to_mv(ptr, sz), number=iters) * 1e6 / iters
+      print(f"Size {label:>9} | Time: {t_us:8.3f} µs")
+
+  def test_speed_from_mv_vs_mv_address(self):
+    x = memoryview(bytearray(1))
+
+    iters = 100000
+    fmv_us = timeit.timeit(lambda: from_mv(x), number=iters) * 1e6 / iters
+    mva_us = timeit.timeit(lambda: mv_address(x), number=iters) * 1e6 / iters
+    print(f"from_mv vs mv_address: {fmv_us:8.3f} µs vs {mva_us:8.3f} µs")
+
+class TestGetContraction(unittest.TestCase):
+  def test_contraction(self):
+    r = get_contraction((1,2,3,4), (2,3,4))
+    self.assertEqual(r, [[0, 1], [2], [3]])
+
+    r = get_contraction((2,1,3,4), (2,3,4))
+    self.assertEqual(r, [[0], [1, 2], [3]])
+
+    r = get_contraction((1,2,3,1,4), (1,2,3,4))
+    self.assertEqual(r, [[], [0, 1], [2], [3, 4]])
+
+    r = get_contraction((1,2,3,1,4,1,1), (2,3,4))
+    self.assertEqual(r, [[0, 1], [2], [3, 4, 5, 6]])
+
+    r = get_contraction((1,2,3,4), (1,2,3*4))
+    self.assertEqual(r, [[], [0, 1], [2, 3]])
+
+    r = get_contraction((1,2,3,4), (2,1,3,4))
+    self.assertEqual(r, [[0, 1], [], [2], [3]])
+
+    r = get_contraction((1,2,3,4), (1,1,2*3*4,1))
+    self.assertEqual(r, [[], [], [0,1,2,3], []])
+
+    r = get_contraction((2,1,3,4), (1,2,3,4))
+    self.assertEqual(r, [[], [0], [1, 2], [3]])
+
+    r = get_contraction((1,2,3,4), (2*3*4,1,1,1))
+    self.assertEqual(r, [[0, 1, 2, 3], [], [], []])
+
+    r = get_contraction((4,4,4,4), (16,1,16))
+    self.assertEqual(r, [[0, 1], [], [2, 3]])
+
+    r = get_contraction((1,2,3,4,1,1,1), (2,3,4))
+    self.assertEqual(r, [[0, 1], [2], [3, 4, 5, 6]])
+
+    r = get_contraction((1,2,3,4), (1,2,3,4,1))
+    self.assertEqual(r, [[], [0, 1], [2], [3], []])
+
+    r = get_contraction((14,1,384,14,1,1,1,1), (1,14,384,14))
+    self.assertEqual(r, [[], [0], [1,2], [3,4,5,6,7]])
+
+    r = get_contraction((14,1,384,1,14,1,1,1,1), (1,14,384,14))
+    self.assertEqual(r, [[], [0], [1,2], [3,4,5,6,7,8]])
+
+    r = get_contraction((512, 512), (1, 1, 512, 1, 1, 1, 1, 512))
+    self.assertEqual(r, [[], [], [0], [], [], [], [], [1]])
+
+    r = get_contraction((1,2,3,4), (1,2,6,2))
+    self.assertEqual(r, None)
+
+  def test_contraction_ones(self):
+    r = get_contraction((1,), (1,1,1))
+    self.assertEqual(r, [[], [], [0]])
+
+    r = get_contraction((1,1), (1,1,1))
+    self.assertEqual(r, [[], [], [0, 1]])
+
+    r = get_contraction((1,1,1,1), (1,))
+    self.assertEqual(r, [[0,1,2,3]])
+
+    r = get_contraction((1,1,1,1), (1,1))
+    self.assertEqual(r, [[], [0,1,2,3]])
+
+    r = get_contraction((1,1,1,1), (1,1,1))
+    self.assertEqual(r, [[], [], [0,1,2,3]])
+
+    r = get_contraction((1,1,1,1), (1,1,1,1))
+    self.assertEqual(r, [[], [], [], [0,1,2,3]])
+
+class TestGetShape(unittest.TestCase):
+  def test_get_shape(self):
+    assert get_shape(2) == ()
+    assert get_shape([]) == (0,)
+    assert get_shape([[]]) == (1, 0)
+    assert get_shape([[1, 2]]) == (1, 2)
+    assert get_shape([[1, 2], (3, 4)]) == (2, 2)
+
+  def test_inhomogeneous_shape(self):
+    with self.assertRaises(ValueError): get_shape([[], [1]])
+    with self.assertRaises(ValueError): get_shape([[1, [2]], [1]])
+
+class TestPolyN(unittest.TestCase):
+  def test_float(self):
+    np.testing.assert_allclose(polyN(1.0, [1.0, -2.0, 1.0]), 0.0)
+    np.testing.assert_allclose(polyN(2.0, [1.0, -2.0, 1.0]), 1.0)
+    np.testing.assert_allclose(polyN(3.0, [1.0, -2.0, 1.0]), 4.0)
+    np.testing.assert_allclose(polyN(4.0, [1.0, -2.0, 1.0]), 9.0)
+
+  def test_uop(self):
+    from tinygrad.dtype import dtypes
+    from tinygrad.uop.ops import UOp
+    from test.helpers import eval_uop
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 1.0), [1.0, -2.0, 1.0])), 0.0)
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 2.0), [1.0, -2.0, 1.0])), 1.0)
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 3.0), [1.0, -2.0, 1.0])), 4.0)
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 4.0), [1.0, -2.0, 1.0])), 9.0)
+
+class TestTimeToStr(unittest.TestCase):
+  def test_seconds(self):           self.assertEqual("   10.01s ", time_to_str(10.01))
+  def test_boundary_sec_ms(self):   self.assertEqual("10000.00ms", time_to_str(10))
+  def test_milliseconds(self):      self.assertEqual("  500.00ms", time_to_str(0.5))
+  def test_boundary_ms_us(self):    self.assertEqual("10000.00us", time_to_str(0.01))
+  def test_microseconds(self):      self.assertEqual("  100.00us", time_to_str(0.0001))
+  def test_zero(self):              self.assertEqual("    0.00us", time_to_str(0))
+  def test_width_formatting(self):  self.assertEqual(" 10.01s ", time_to_str(10.01, w=6))
+
+class TestCStyleDivMod(unittest.TestCase):
+  def test_div_pos(self):
+    self.assertEqual(cdiv(-9, 5), -1)
+    self.assertEqual(cdiv(-4, 5), 0)
+    self.assertEqual(cdiv(0, 5), 0)
+    self.assertEqual(cdiv(4, 5), 0)
+    self.assertEqual(cdiv(9, 5), 1)
+  def test_div_neg(self):
+    self.assertEqual(cdiv(-9, -5), 1)
+    self.assertEqual(cdiv(-4, -5), 0)
+    self.assertEqual(cdiv(0, -5), 0)
+    self.assertEqual(cdiv(4, -5), 0)
+    self.assertEqual(cdiv(9, -5), -1)
+  def test_mod_pos(self):
+    self.assertEqual(cmod(-9, 5), -4)
+    self.assertEqual(cmod(-4, 5), -4)
+    self.assertEqual(cmod(0, 5), 0)
+    self.assertEqual(cmod(4, 5), 4)
+    self.assertEqual(cmod(9, 5), 4)
+  def test_mod_neg(self):
+    self.assertEqual(cmod(-9, -5), -4)
+    self.assertEqual(cmod(-4, -5), -4)
+    self.assertEqual(cmod(0, -5), 0)
+    self.assertEqual(cmod(4, -5), 4)
+    self.assertEqual(cmod(9, -5), 4)
+
+class TestGetBits(unittest.TestCase):
+  def test_low_bits(self):
+    self.assertEqual(getbits(0b11010110, 0, 3), 0b0110)
+
+  def test_high_bits(self):
+    self.assertEqual(getbits(0b11010110, 4, 7), 0b1101)
+
+  def test_middle_bits(self):
+    self.assertEqual(getbits(0b11010110, 3, 5), 0b010)
+
+  def test_full_range(self):
+    self.assertEqual(getbits(0b11010110, 0, 7), 0b11010110)
+
+  def test_single_bit(self):
+    self.assertEqual(getbits(0b100000000, 8, 8), 1)
+
+class TestArgFix(unittest.TestCase):
+  def test_none(self):
+    self.assertEqual(argfix(None), (None, ))
+    self.assertEqual(argfix(None, None), (None, None))
+  def test_positional_arguments(self):
+    self.assertEqual(argfix(1, 2, 3), (1, 2, 3))
+  def test_tuple(self):
+    self.assertEqual(argfix((1., 2., 3.)), (1., 2., 3.))
+  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)
+
+  def test_wrap_explicit_newline(self):
+    wrap = 10
+    st = "\n".join(["x"*wrap, "x"*wrap, "x"*wrap])
+    st2 = word_wrap(st, wrap=wrap)
+    self.assertEqual(len(st2.splitlines()), len(st.splitlines()))
+
+    st = "\n".join(["x"*(wrap+1), "x"*wrap, "x"*wrap])
+    st2 = word_wrap(st, wrap=wrap)
+    self.assertEqual(len(st2.splitlines()), len(st.splitlines())+1)
+
+    st = "\n".join(["x"*(wrap+1), "x"*(wrap+1), "x"*(wrap+1)])
+    st2 = word_wrap(st, wrap=wrap)
+    self.assertEqual(len(st2.splitlines()), len(st.splitlines())+3)
+
+class TestIsNumpyNdarray(unittest.TestCase):
+  def test_ndarray(self):
+    self.assertTrue(is_numpy_ndarray(np.array([1, 2, 3])))
+  def test_ndarray_tolist(self):
+    self.assertFalse(is_numpy_ndarray(np.array([1, 2, 3]).tolist()))
+  def test_list(self):
+    self.assertFalse(is_numpy_ndarray([1, 2, 3]))
+  def test_tensor(self):
+    self.assertFalse(is_numpy_ndarray(Tensor([1, 2, 3])))
+    self.assertFalse(is_numpy_ndarray(Tensor(np.array([1, 2, 3]))))
+
+if __name__ == '__main__':
+  unittest.main()

test/null/test_indexing.py

@@ -0,0 +1,143 @@
+# test cases are modified from pytorch test_indexing.py
+
+import unittest
+import numpy as np
+
+from tinygrad import Tensor, dtypes
+
+def numpy_testing_assert_equal_helper(a, b):
+  if isinstance(a, Tensor): a = a.numpy()
+  if isinstance(b, Tensor): b = b.numpy()
+  np.testing.assert_equal(a, b)
+
+class TestIndexing(unittest.TestCase):
+  def test_single_int(self):
+    v = Tensor.randn(5, 7, 3)
+    numpy_testing_assert_equal_helper(v[4].shape, (7, 3))
+
+  def test_multiple_int(self):
+    v = Tensor.randn(5, 7, 3)
+    numpy_testing_assert_equal_helper(v[4].shape, (7, 3))
+    numpy_testing_assert_equal_helper(v[4, :, 1].shape, (7,))
+
+  def test_none(self):
+    v = Tensor.randn(5, 7, 3)
+    numpy_testing_assert_equal_helper(v[None].shape, (1, 5, 7, 3))
+    numpy_testing_assert_equal_helper(v[:, None].shape, (5, 1, 7, 3))
+    numpy_testing_assert_equal_helper(v[:, None, None].shape, (5, 1, 1, 7, 3))
+    numpy_testing_assert_equal_helper(v[..., None].shape, (5, 7, 3, 1))
+
+  def test_int_indices(self):
+    v = Tensor.randn(5, 7, 3)
+    numpy_testing_assert_equal_helper(v[[0, 4, 2]].shape, (3, 7, 3))
+    numpy_testing_assert_equal_helper(v[:, [0, 4, 2]].shape, (5, 3, 3))
+    numpy_testing_assert_equal_helper(v[:, [[0, 1], [4, 3]]].shape, (5, 2, 2, 3))
+
+  def test_index_src_datatype(self):
+    src = Tensor.ones(3, 2, 4)
+    # test index
+    res = src[[0, 2, 1], :, :]
+    numpy_testing_assert_equal_helper(res.shape, src.shape)
+
+  def test_empty_slice(self):
+    x = Tensor.randn(2, 3, 4, 5)
+    y = x[:, :, :, 1]
+    z = y[:, 1:1, :]
+    numpy_testing_assert_equal_helper((2, 0, 4), z.shape)
+
+  def test_invalid_index(self):
+    x = Tensor.arange(0, 16).reshape(4, 4)
+    self.assertRaises(TypeError, lambda: x["0":"1"])
+
+  def test_out_of_bound_index(self):
+    x = Tensor.arange(0, 100).reshape(2, 5, 10)
+    self.assertRaises(IndexError, lambda: x[0, 5])
+    self.assertRaises(IndexError, lambda: x[4, 5])
+    self.assertRaises(IndexError, lambda: x[0, 1, 15])
+    self.assertRaises(IndexError, lambda: x[:, :, 12])
+
+  def test_take_along_dim(self):
+    # NOTE: the actual test logic is inside _test_against_numpy which is never called
+    # This test effectively does nothing but defines a function
+    def _test_against_numpy(t: Tensor, indices: Tensor, dim):
+      actual = t.gather(dim, indices)
+      t_np = t.numpy()
+      indices_np = indices.numpy()
+      expected = np.take_along_axis(t_np, indices_np, axis=dim)
+      numpy_testing_assert_equal_helper(actual, expected)
+
+      # TODO argsort
+      '''
+      for shape in [(3, 2), (2, 3, 5), (2, 4, 0), (2, 3, 1, 4)]:
+        for noncontiguous in [True, False]:
+          for dtype in (dtypes.float32, dtypes.int64):
+            t = make_tensor(shape, dtype=dtype, noncontiguous=noncontiguous)
+            for dim in list(range(t.ndim)) + [None]:
+              if dim is None:
+                indices = argsort(t.reshape(-1))
+              else:
+                indices = argsort(t, dim=dim)
+
+          _test_against_numpy(t, indices, dim)
+      '''
+
+      # test broadcasting
+      t = Tensor.ones((3, 4, 1))
+      indices = Tensor.ones((1, 2, 5), dtype=dtypes.int64)
+
+      _test_against_numpy(t, indices, 1)
+
+      # test empty indices
+      t = Tensor.ones((3, 4, 5))
+      indices = Tensor.ones((3, 0, 5), dtype=dtypes.int64)
+
+      _test_against_numpy(t, indices, 1)
+
+class TestNumpy(unittest.TestCase):
+  def test_index_no_floats(self):
+    a = Tensor([[[5.]]])
+
+    self.assertRaises(IndexError, lambda: a[0.0])
+    self.assertRaises(IndexError, lambda: a[0, 0.0])
+    self.assertRaises(IndexError, lambda: a[0.0, 0])
+    self.assertRaises(IndexError, lambda: a[0.0, :])
+    self.assertRaises(IndexError, lambda: a[:, 0.0])
+    self.assertRaises(IndexError, lambda: a[:, 0.0, :])
+    self.assertRaises(IndexError, lambda: a[0.0, :, :])
+    self.assertRaises(IndexError, lambda: a[0, 0, 0.0])
+    self.assertRaises(IndexError, lambda: a[0.0, 0, 0])
+    self.assertRaises(IndexError, lambda: a[0, 0.0, 0])
+    self.assertRaises(IndexError, lambda: a[-1.4])
+    self.assertRaises(IndexError, lambda: a[0, -1.4])
+    self.assertRaises(IndexError, lambda: a[-1.4, 0])
+    self.assertRaises(IndexError, lambda: a[-1.4, :])
+    self.assertRaises(IndexError, lambda: a[:, -1.4])
+    self.assertRaises(IndexError, lambda: a[:, -1.4, :])
+    self.assertRaises(IndexError, lambda: a[-1.4, :, :])
+    self.assertRaises(IndexError, lambda: a[0, 0, -1.4])
+    self.assertRaises(IndexError, lambda: a[-1.4, 0, 0])
+    self.assertRaises(IndexError, lambda: a[0, -1.4, 0])
+    # these two trigger slice internal type verification first
+    self.assertRaises(TypeError, lambda: a[0.0:, 0.0])
+    self.assertRaises(TypeError, lambda: a[0.0:, 0.0,:])
+
+  def test_none_index(self):
+    # `None` index adds newaxis
+    a = Tensor([1, 2, 3])
+    numpy_testing_assert_equal_helper(a[None].ndim, a.ndim+1)
+
+  def test_everything_returns_views(self):
+    # Before `...` would return a itself.
+    a = Tensor([5])
+
+    self.assertIs(a, a[()])
+    self.assertIs(a, a[...])
+    self.assertIs(a, a[:])
+
+  def test_broaderrors_indexing(self):
+    a = Tensor.zeros(5, 5)
+    self.assertRaises(IndexError, a.__getitem__, ([0, 1], [0, 1, 2]))
+    self.assertRaises(IndexError, a.contiguous().__setitem__, ([0, 1], [0, 1, 2]), 0)
+
+if __name__ == '__main__':
+  unittest.main()

test/null/test_linearizer_failures.py

@@ -1,24 +1,24 @@
 # ruff: noqa: E501
 import unittest
-from tinygrad.uop.ops import UOp, Ops, AxisType
+from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
 from tinygrad.dtype import dtypes
 from tinygrad.engine.realize import get_program
 from tinygrad.device import Device
 
 class TestLinearizerFailures(unittest.TestCase):
   def test_fail_1(self):
-    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64), arg=0, src=())
+    c0 = UOp(Ops.PARAM, dtypes.float.ptr(64), arg=0, src=())
     c1 = UOp.range(UOp.const(dtypes.index, 2), 1, AxisType.LOOP)
     c2 = UOp.range(UOp.const(dtypes.index, 32), 2, AxisType.LOOP)
     c3 = ((c1*UOp.const(dtypes.index, 32))+c2)
-    c4 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(163840), arg=1, src=())
+    c4 = UOp(Ops.PARAM, dtypes.float.ptr(163840), arg=1, src=())
     c5 = UOp.range(UOp.const(dtypes.index, 2560), 0, AxisType.REDUCE)
     c6 = c4.index(((((((c5//UOp.const(dtypes.index, 8))%UOp.const(dtypes.index, 8))*UOp.const(dtypes.index, 8))+(c5%UOp.const(dtypes.index, 8)))+(((c2*UOp.const(dtypes.index, 40))+(c5//UOp.const(dtypes.index, 64)))*UOp.const(dtypes.index, 64)))+(c1*UOp.const(dtypes.index, 81920))))
-    c7 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64), arg=2, src=())
+    c7 = UOp(Ops.PARAM, dtypes.float.ptr(64), arg=2, src=())
     c8 = c7.index(c3)
     c9 = ((((c6+(c8*UOp.const(dtypes.float, -1.0)))*(c6+(c8*UOp.const(dtypes.float, -1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(dtypes.float, 0.000390625))+UOp.const(dtypes.float, 1e-05)).sqrt().reciprocal()
     c10 = c0.index(c3).store(c9).end(c1, c2)
-    ast = c10.sink()
+    ast = c10.sink(arg=KernelInfo())
     get_program(ast, renderer=Device[Device.DEFAULT].renderer)
 
 if __name__ == '__main__':

test/null/test_linearizer_rewrite.py

@@ -35,7 +35,7 @@ class TestLinearizerRewrite(unittest.TestCase):
     prg = get_program(ast, Device["CPU"].renderer)
     assert prg.applied_opts == (), f"expected no opts, got {prg}"
 
-    prg = get_program(ast.replace(arg=None), Device["CPU"].renderer)
+    prg = get_program(ast.replace(arg=KernelInfo()), Device["CPU"].renderer)
     assert prg.applied_opts != (), f"expected opts to apply, got {prg.applied_opts}"
 
     prg = get_program(ast.replace(arg=KernelInfo(name="custom")), Device["CPU"].renderer)

test/null/test_llm_server.py

@@ -0,0 +1,137 @@
+import unittest, threading, time
+from unittest.mock import Mock
+
+class TestLLMServer(unittest.TestCase):
+  """Integration tests using the real OpenAI client."""
+
+  @classmethod
+  def setUpClass(cls):
+    cls.mock_tok = Mock()
+    cls.mock_tok.role = Mock(return_value=[100, 101])
+    cls.mock_tok.encode = Mock(return_value=[200, 201, 202])
+    cls.mock_tok.decode = Mock(return_value="Hello")
+    cls.mock_tok.end_turn = Mock(return_value=[998])
+
+    cls.mock_model = Mock()
+    cls.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 999]))
+
+    cls.bos_id = 1
+    cls.eos_id = 999
+
+    import tinygrad.apps.llm as llm_module
+    llm_module.model = cls.mock_model
+    llm_module.tok = cls.mock_tok
+    llm_module.bos_id = cls.bos_id
+    llm_module.eos_id = cls.eos_id
+
+    from tinygrad.apps.llm import Handler
+    from tinygrad.viz.serve import TCPServerWithReuse
+
+    cls.server = TCPServerWithReuse(('127.0.0.1', 0), Handler)
+    cls.port = cls.server.server_address[1]
+    cls.server_thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
+    cls.server_thread.start()
+    time.sleep(0.1)
+
+    from openai import OpenAI
+    cls.client = OpenAI(base_url=f"http://127.0.0.1:{cls.port}/v1", api_key="test")
+
+  @classmethod
+  def tearDownClass(cls):
+    cls.server.shutdown()
+    cls.server.server_close()
+
+  def test_chat_completion_stream(self):
+    stream = self.client.chat.completions.create(
+      model="test",
+      messages=[{"role": "user", "content": "Hello"}],
+      stream=True
+    )
+
+    chunks = list(stream)
+    self.assertGreater(len(chunks), 0)
+    self.assertEqual(chunks[0].choices[0].delta.role, "assistant")
+    self.assertEqual(chunks[-1].choices[0].finish_reason, "stop")
+
+  def test_openai_response_structure(self):
+    stream = self.client.chat.completions.create(
+      model="test-model",
+      messages=[{"role": "user", "content": "Test"}],
+      stream=True
+    )
+
+    for chunk in stream:
+      self.assertTrue(chunk.id.startswith("chatcmpl-"))
+      self.assertEqual(chunk.object, "chat.completion.chunk")
+      self.assertIsNotNone(chunk.choices)
+      self.assertIsNotNone(chunk.created)
+      self.assertIsInstance(chunk.created, int)
+      self.assertEqual(chunk.model, "test-model")
+
+  def test_stream_with_usage(self):
+    stream = self.client.chat.completions.create(
+      model="test",
+      messages=[{"role": "user", "content": "Hello"}],
+      stream=True,
+      stream_options={"include_usage": True}
+    )
+
+    chunks = list(stream)
+    last_chunk = chunks[-1]
+
+    self.assertIsNotNone(last_chunk.usage)
+    self.assertIsNotNone(last_chunk.usage.prompt_tokens)
+    self.assertIsNotNone(last_chunk.usage.completion_tokens)
+    self.assertIsNotNone(last_chunk.usage.total_tokens)
+
+  def test_multi_turn_conversation(self):
+    stream = self.client.chat.completions.create(
+      model="test",
+      messages=[
+        {"role": "system", "content": "You are helpful."},
+        {"role": "user", "content": "Hello"},
+        {"role": "assistant", "content": "Hi!"},
+        {"role": "user", "content": "How are you?"}
+      ],
+      stream=True
+    )
+
+    chunks = list(stream)
+    self.assertGreater(len(chunks), 0)
+    self.assertEqual(chunks[-1].choices[0].finish_reason, "stop")
+
+  def test_content_is_streamed(self):
+    stream = self.client.chat.completions.create(
+      model="test",
+      messages=[{"role": "user", "content": "Hello"}],
+      stream=True
+    )
+
+    contents = []
+    for chunk in stream:
+      if chunk.choices and chunk.choices[0].delta.content:
+        contents.append(chunk.choices[0].delta.content)
+
+    self.assertGreater(len(contents), 0)
+
+  def test_non_streaming(self):
+    resp = self.client.chat.completions.create(
+      model="test-model",
+      messages=[{"role": "user", "content": "Hello"}],
+      stream=False
+    )
+
+    self.assertTrue(resp.id.startswith("chatcmpl-"))
+    self.assertEqual(resp.object, "chat.completion")
+    self.assertEqual(resp.model, "test-model")
+    self.assertIsNotNone(resp.created)
+    self.assertEqual(len(resp.choices), 1)
+    self.assertEqual(resp.choices[0].message.role, "assistant")
+    self.assertIsNotNone(resp.choices[0].message.content)
+    self.assertEqual(resp.choices[0].finish_reason, "stop")
+    self.assertIsNotNone(resp.usage)
+    self.assertIsNotNone(resp.usage.prompt_tokens)
+    self.assertIsNotNone(resp.usage.completion_tokens)
+
+if __name__ == '__main__':
+  unittest.main()

test/null/test_multitensor.py

@@ -0,0 +1,58 @@
+import unittest
+from tinygrad import Tensor, GlobalCounters, dtypes
+
+class TestMultiRamUsage(unittest.TestCase):
+  def setUp(self):
+    self.baseline = GlobalCounters.mem_used
+    self.N = 100
+  def assertUsed(self, amt, strict=True):
+    used = GlobalCounters.mem_used - self.baseline
+    print(f"used {used} bytes")
+    if strict: self.assertEqual(used, amt)
+    else: self.assertLessEqual(used, amt)
+
+  def test_zeros(self):
+    _ = Tensor.zeros(self.N, self.N).contiguous().realize()
+    self.assertUsed(self.N*self.N*4)
+
+  def test_zeros_del(self):
+    _ = Tensor.zeros(self.N, self.N).contiguous().realize()
+    del _
+    self.assertUsed(0)
+
+  @unittest.skip("flaky")
+  def test_zeros_copy(self):
+    devices_2 = ("NULL:1", "NULL:2")
+    _ = Tensor.zeros(self.N, self.N).contiguous().to(devices_2).realize()
+    # NOTE: the first one on the DEFAULT device should be freed
+    self.assertUsed(self.N*self.N*4*2)
+
+  @unittest.skip("flaky")
+  def test_zeros_shard(self, devices=("NULL:1", "NULL:2")):
+    _ = Tensor.zeros(self.N, self.N).contiguous().shard(devices, axis=0).realize()
+    self.assertUsed(self.N*self.N*4) # sharding should not increase total ram usage
+  def test_zeros_shard_self(self): self.test_zeros_shard(("NULL:0", "NULL:1"))
+
+  @unittest.skip("flaky")
+  def test_zeros_contiguous_shard(self):
+    devices_2 = ("NULL:1", "NULL:2")
+    _ = Tensor.zeros(self.N, self.N).contiguous().shard(devices_2, axis=0).contiguous().realize()
+    self.assertUsed(self.N*self.N*4) # sharding should not increase total ram usage
+
+  def _test_matmul_half(self, dev_count:int):
+    N = 32
+    total_mem = {}
+    devs = tuple(f"NULL:{i}" for i in range(dev_count))
+    for dtype in {dtypes.float, dtypes.half}:
+      GlobalCounters.reset()
+      a = Tensor.empty((N, N), dtype=dtype, device=devs[0]).shard(devs, axis=0)
+      b = Tensor.empty((N, N), dtype=dtype, device=devs[0]).shard(devs, axis=None)
+      (a @ b).realize()
+      total_mem[dtype] = GlobalCounters.global_mem
+    self.assertEqual(total_mem[dtypes.half], total_mem[dtypes.float] // 2)
+
+  def test_matmul_half(self): self._test_matmul_half(dev_count=2)
+  def test_matmul_half_alt(self): self._test_matmul_half(dev_count=4)
+
+if __name__ == '__main__':
+  unittest.main()