add ansipad and PARALLEL contextvar

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

@@ -61,7 +61,6 @@ runs:
         echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
         if [[ "$RUNNER_OS" == "Linux" ]]; then
           echo "VIRTUAL_ENV=/opt/venv/${{ inputs.python-version }}" >> "$GITHUB_ENV"
-          echo "UV_PYTHON_INSTALL_DIR=/opt/python" >> "$GITHUB_ENV"
         else
           echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
         fi
@@ -71,6 +70,11 @@ runs:
       with:
         enable-cache: 'false' # see below for manual caching
 
+    - name: Set up Python ${{ inputs.python-version }}
+      uses: actions/setup-python@v6
+      with:
+        python-version: ${{ inputs.python-version }}
+
     # **** Caching packages ****
 
     - name: Cache Python packages (PR)

.github/workflows/benchmark.yml

@@ -88,13 +88,13 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 30
+    timeout-minutes: 60
     defaults:
       run:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: "0"
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -102,11 +102,12 @@ jobs:
     - name: Setup (AMD)
       if: ${{ matrix.dev == 'AMD' }}
       run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
-      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -116,10 +117,10 @@ jobs:
       run: python3 test/external/process_replay/reset.py
     - name: Run llama3.2
       run: BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
-    - name: Run qwen3.8
-      # qwen3.8:27b doesn't fit on mac
+    - name: Run qwen3.6
+      # qwen3.6:35b-a3b doesn't fit on mac
       if: ${{ matrix.dev != 'METAL' }}
-      run: BENCHMARK_LOG=qwen38_27b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.8:27b --benchmark --warmup
+      run: BENCHMARK_LOG=qwen36_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.6:35b-a3b --benchmark --warmup
     - name: Run olmoe
       # just metal for now
       if: ${{ matrix.dev == 'METAL' }}
@@ -134,13 +135,13 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 10
+    timeout-minutes: 60
     defaults:
       run:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: "0"
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -148,11 +149,12 @@ jobs:
     - name: Setup (AMD)
       if: ${{ matrix.dev == 'AMD' }}
       run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
-      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -182,13 +184,13 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['AMD', 'NV']
-    timeout-minutes: 5
+    timeout-minutes: 60
     defaults:
       run:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: "0"
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -196,11 +198,12 @@ jobs:
     - name: Setup (AMD)
       if: ${{ matrix.dev == 'AMD' }}
       run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
-      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
     - name: Symlink models and datasets
       run: |
         mkdir -p extra/datasets
@@ -224,13 +227,13 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 15
+    timeout-minutes: 60
     defaults:
       run:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: "0"
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -238,11 +241,12 @@ jobs:
     - name: Setup (AMD)
       if: ${{ matrix.dev == 'AMD' }}
       run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
-      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -269,13 +273,13 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['AMD', 'NV']
-    timeout-minutes: 20
+    timeout-minutes: 60
     defaults:
       run:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: "0"
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -283,11 +287,12 @@ jobs:
     - name: Setup (AMD)
       if: ${{ matrix.dev == 'AMD' }}
       run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
-      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
     - name: Symlink models and datasets
       run: |
         mkdir -p weights
@@ -322,7 +327,7 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 10
+    timeout-minutes: 60
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -335,11 +340,12 @@ jobs:
     - name: Setup (AMD)
       if: ${{ matrix.dev == 'AMD' }}
       run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
-      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -416,7 +422,7 @@ jobs:
   testusbgpu:
     name: UsbGPU Benchmark
     runs-on: [self-hosted, macOS]
-    timeout-minutes: 3
+    timeout-minutes: 10
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -518,6 +524,15 @@ jobs:
         shell: bash -e -o pipefail {0}
     if: github.repository_owner == 'tinygrad'
     steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: reset process replay
+      run: test/external/process_replay/reset.py
     - name: Checkout Code
       uses: actions/checkout@v6
     - name: setup staging db
@@ -541,7 +556,7 @@ jobs:
   testcommausbgpubenchmark:
     name: UsbGPU Benchmark (comma)
     runs-on: [self-hosted, Linux, comma4]
-    timeout-minutes: 10
+    timeout-minutes: 20
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -570,7 +585,7 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['AMD', 'NV']
-    timeout-minutes: 5
+    timeout-minutes: 20
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -582,8 +597,9 @@ jobs:
       uses: actions/checkout@v6
     - name: Setup
       run: |
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} rmmod --expect
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} rmmod
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} kill_pids
         mkdir -p extra/datasets
         ln -s /raid/datasets/imagenet extra/datasets/imagenet
     - name: setup staging db
@@ -618,9 +634,6 @@ jobs:
       run: |
         GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
         GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
-    - name: HEVC Decode Benchmark
-      if: ${{ matrix.dev == 'NV' }}
-      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
     - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
       if: ${{ matrix.dev == 'NV' }}
       run: BENCHMARK_LOG=resnet_10steps MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
@@ -643,7 +656,7 @@ jobs:
   llvmspeed:
     name: LLVM Speed
     runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 5
+    timeout-minutes: 20
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code

.github/workflows/test.yml

@@ -233,7 +233,7 @@ jobs:
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
     - name: Repo line count <= 26000 lines
-      run: MAX_LINE_COUNT=26500 python sz.py
+      run: MAX_LINE_COUNT=26000 python sz.py
 
   spec:
     strategy:
@@ -504,7 +504,7 @@ jobs:
       - name: Run AMD renderer tests (AMD:LLVM)
         run: DEV=MOCKKFD+AMD:LLVM python -m pytest -n=auto test/amd/ --durations 20
       - name: Run SQTT profiling tests
-        run: VIZ=-2 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
+        run: PROFILE=1 SQTT=1 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
       - name: Run AMD emulated tests on NULL backend
         env:
           AMD: 0
@@ -679,5 +679,4 @@ jobs:
         run: |
           python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
           DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
-          # QCOMCL compiles in qemu, too slow for parallel workers
-          ${{ contains(matrix.dev, 'QCOMCL') && 'PARALLEL=0' || '' }} python -m pytest -n=auto test/backend/test_ops.py --durations=20
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20

.gitignore

@@ -69,4 +69,3 @@ mutants
 dagre/
 graphlib/
 uv.lock
-pi_session_window0.jsonl

AGENTS.md

@@ -4,4 +4,3 @@
 - Run `python -m mypy tinygrad/` to typecheck
 - Run `python -m ruff check .` to lint
 - Read `./tinygrad/viz/README.md` for profiling and debugging rewrite rules
-- Do not do amend commits. Always do a new commit if a force push to origin would be required.

examples/mlperf/model_train.py

@@ -1773,13 +1773,8 @@ def train_gptoss():
   def minibatch(tokens:Tensor):
     if is_dp: tokens = tokens.to(None).shard(device, 0)
     if not is_sharding: tokens = tokens.to(None)
-
     logits:Tensor = model(tokens[:, :-1], save=True)
-    if getenv("FUSED_CE", 0):
-      from extra.llama_kernels.fused_ce import fused_ce_loss
-      loss = fused_ce_loss(logits.cast(dtypes.bfloat16), tokens[:, 1:], label_smoothing=0.0)
-    else:
-      loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
 
     for g, new_g in zip(grads, loss.gradient(*optim.params)):
       apply_grad(g, new_g.uop)

examples/mlperf/optim.py

@@ -15,7 +15,7 @@ def stochastic_round_bf16(x:Tensor) -> Tensor:
   bits = x.bitcast(dtypes.uint32)
   if isinstance(x.device, tuple):
     shape = x.uop.shard_shape if x.uop.axis is not None else x.shape
-    noise = Tensor(UOp(Ops.MSTACK, src=tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
+    noise = Tensor(UOp(Ops.MSTACK, dtypes.default_float, tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
   else:
     noise = x.rand_like()
   noise = (noise * 0xFFFF).cast(dtypes.uint32)

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

@@ -35,7 +35,7 @@ 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="1e-3" END_LR="1e-4" WARMUP_SAMPLES=2048 MAX_STEPS=1200000
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}
@@ -46,7 +46,7 @@ export DATA_SEED=${DATA_SEED:-5760}
 export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
 
-export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+export FAKEDATA=${FAKEDATA:-$([[ "$DEV" == NULL:* ]] && echo 1 || echo 0)} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
   export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
 fi

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

@@ -35,7 +35,7 @@ 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="1e-3" END_LR="1e-4" WARMUP_SAMPLES=2048 MAX_STEPS=1200000
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/run_and_time.sh

@@ -17,7 +17,7 @@ export USE_ATOMICS=1
 export ASM_GEMM=1
 export WQKV=1
 export MASTER_WEIGHTS=1
-export MXFP4=1
+export FP8=1
 export ALLREDUCE_CAST=1
 export FAST_CE=1
 export FUSED_INPUT_QUANTIZE=1
@@ -26,7 +26,7 @@ export FUSED_ADD_NORM_MUL_QUANTIZE=1
 export FUSED_SILU_W13=1
 export SPLIT_W13=0
 
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32"
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=8 MP=1 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=2
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 

examples/openpilot/compile3.py

@@ -107,21 +107,14 @@ def compile(onnx_file):
   return inputs, test_val
 
 def test_vs_compile(run, inputs, test_val=None):
-  if (log:=bool(getenv("BENCHMARK_LOG", ""))): from extra.bench_log import WallTimeEvent, BenchEvent
 
   # run 20 times
   step_times = []
   for _ in range(20):
     st = time.perf_counter()
-    if log:
-      with WallTimeEvent(BenchEvent.STEP):
-        out = run(**inputs)
-        mt = time.perf_counter()
-        val = out.numpy()
-    else:
-      out = run(**inputs)
-      mt = time.perf_counter()
-      val = out.numpy()
+    out = run(**inputs)
+    mt = time.perf_counter()
+    val = out.numpy()
     et = time.perf_counter()
     step_times.append((et-st)*1e3)
     print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
@@ -167,6 +160,12 @@ def test_vs_onnx(new_inputs, test_val, onnx_file, tol):
   print("test vs onnx passed")
   return timings
 
+def bench(run, inputs):
+  from extra.bench_log import WallTimeEvent, BenchEvent
+  for _ in range(10):
+    with WallTimeEvent(BenchEvent.STEP):
+      run(**inputs).numpy()
+
 if __name__ == "__main__":
   if getenv("RUN_PICKLE"):
     with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)
@@ -182,3 +181,6 @@ if __name__ == "__main__":
     test_vs_compile(pickle_loaded, inputs, outputs)
     if getenv("SELFTEST"):
       test_vs_onnx(inputs, outputs, onnx_file, 1e-4)
+
+  if getenv("BENCHMARK_LOG", ""):
+    bench(pickle_loaded, inputs)

extra/gemm/moe_routing.py

@@ -53,7 +53,7 @@ def _ggather_bwd(gradient:UOp, kernel:UOp) -> tuple:
     g, m, j, jo, ji = _kv_ranges(Gk, M, Dk, _blk_for(Dk))
     row = idx.index(g, m).cast(dtypes.weakint)
     val = gout.index(g, m, j).load().cast(dtypes.float32)
-    atomic = UOp(Ops.CUSTOM, src=(gtab.index(g, row, j), val), arg=atomic_str)
+    atomic = UOp(Ops.CUSTOM, dtypes.void, (gtab.index(g, row, j), val), arg=atomic_str)
     return atomic.end(g, m, jo, ji).sink(arg=KernelInfo(name=f"ggather_bwd_{M}_{Dk}", opts_to_apply=()))
   grad_table = Tensor.custom_kernel(gt, go, Tensor(idx_u, device=dev), fxn=_bwd_kernel)[0]
   return (None, grad_table.cast(table_u.dtype).uop, None)

extra/hcq/hcq_smi.py

@@ -16,12 +16,9 @@ def _do_reset_device(pci_bus): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/dev
 def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")
 
 def cmd_remove_module(args):
-  modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia"] if args.backend == "nv" else ["amdgpu"]
+  modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia", "ast"] if args.backend == "nv" else ["amdgpu"]
   to_unload = [m for m in modules if _is_module_loaded(m)]
   if not to_unload: print("Kernel modules are not loaded")
-  elif getattr(args, "expect", False):
-    print(f"Kernel modules are loaded: {to_unload}")
-    sys.exit(1)
   else:
     print("Removing kernel modules:", ", ".join(to_unload))
     try: subprocess.run(["sudo", "modprobe", "-r", *to_unload], check=True)
@@ -63,19 +60,17 @@ def cmd_show_pids(args):
 
 def cmd_kill_pids(args):
   devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
-  use_sudo = not getattr(args, "sudoless", False)
 
   for dev in devs:
     for i in range(128):
       if i > 0: time.sleep(0.2)
 
       try:
-        try: pid = subprocess.check_output((['sudo'] if use_sudo else []) +
-                                           ['lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
+        try: pid = subprocess.check_output(['sudo', 'lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
         except subprocess.CalledProcessError: break
 
         print(f"Killing process {pid} (which uses {dev})")
-        subprocess.run((['sudo'] if use_sudo else []) + ['kill', '-9', pid], check=True)
+        subprocess.run(['sudo', 'kill', '-9', pid], check=True)
       except subprocess.CalledProcessError as e:
         print(f"Failed to kill process for device {dev}: {e}", file=sys.stderr)
 
@@ -84,7 +79,6 @@ def add_common_commands(parent_subparsers):
   p_insmod.set_defaults(func=cmd_insert_module)
 
   p_rmmod = parent_subparsers.add_parser("rmmod", help="Remove a kernel module")
-  p_rmmod.add_argument("--expect", action="store_true", help="Just assert that module is already unloaded")
   p_rmmod.set_defaults(func=cmd_remove_module)
 
   p_reset = parent_subparsers.add_parser("reset", help="Reset a device")
@@ -97,7 +91,6 @@ def add_common_commands(parent_subparsers):
 
   p_reset = parent_subparsers.add_parser("kill_pids", help="Kill pids of processes using the device")
   p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
-  p_reset.add_argument("--sudoless", action="store_true", help="Do not use sudo when detecting or killing pids")
   p_reset.set_defaults(func=cmd_kill_pids)
 
 if __name__ == "__main__":

extra/hcq2/ops_amd2.py

@@ -19,7 +19,7 @@ from tinygrad.runtime.support.hcq import FileIOInterface, HCQBuffer, MMIOInterfa
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
 from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc
 from tinygrad.runtime.support.system import PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
-from tinygrad.runtime.support.usb import USB3, usb_ib, usb_push, usb_arm_bytes, pm_usb_stage, pm_usb_hostio, pm_usb_bufferize
+from tinygrad.runtime.support.usb import USB3
 from tinygrad.runtime.support.memory import AddrSpace, BumpAllocator
 from tinygrad.runtime.ops_amd import SQTT, SQTT_ITRACE_SE_MASK, SQTT_LIMIT_SE, SQTT_SIMD_SEL, SQTT_TOKEN_EXCLUDE, PMC
 from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_EQ, WAIT_REG_MEM_FUNCTION_NEQ, WAIT_REG_MEM_FUNCTION_GEQ
@@ -87,7 +87,7 @@ def release_mem(ctx, address=0x0, value=0, data_sel=0, int_sel=2, ctxid=0, cache
 
 def memory_barrier(ctx):
   pf = '' if ctx.nbio.version[0] == 2 else '0' if ctx.nbio.version[:2] != (7, 11) else '1'
-  return UOp(Ops.LINEAR, src=(
+  return UOp(Ops.LINEAR, dtypes.void, (
     wait_reg_mem(ctx, reg=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_REQ').addr[0],
                  reg_done=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff),
     acquire_mem(ctx)))
@@ -135,7 +135,7 @@ def pm4_program(ctx, call, prg):
     wreg(ctx, ctx.gc.regCOMPUTE_START_X, 0, 0, 0, *(info.local_size or (1, 1, 1)), 0, 0),
     pkt3(ctx, PM4Ops.DISPATCH_DIRECT, *info.global_size, dispatch_init),
     pkt3(ctx, PM4Ops.EVENT_WRITE, ctx.pm4.EVENT_TYPE(ctx.soc.CS_PARTIAL_FLUSH) | ctx.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))]
-  return UOp(Ops.LINEAR, src=tuple(ins))
+  return UOp(Ops.LINEAR, dtypes.void, tuple(ins))
 
 pm_pm4_opsel = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), pm4_program),
@@ -146,14 +146,11 @@ pm_pm4_opsel = PatternMatcher([
   (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
 ])
 
-def queue_ptrs(devs, qname:str, q:AMDQueueDesc) -> tuple[UOp, ...]:
-  return tuple(UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"{qname}_{n}")
-               for n, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
-
 def pm4_submit(ctx, lin):
   # ensure compute queues are allocated
   for d in (devs:=ctx.devs): q = Device[d].compute_queue
-  ring, wptr, doorbell, put_ptr = queue_ptrs(devs, "COMPUTE:0", q)
+  ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COMPUTE:0_{name}")
+    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
 
   # the host fence at the start of the batch guarantees the ib is free to reuse
   size_dw = sum(len(ins.src) for ins in lin.src)
@@ -207,7 +204,7 @@ def sdma_timestamp(ctx, ins, dst):
 pm_sdma_opsel = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),
 
-  (UPat(Ops.INS, arg="barrier"), lambda: UOp(Ops.NOOP)),
+  (UPat(Ops.INS, arg="barrier"), lambda: UOp(Ops.NOOP, dtypes.void, ())),
   (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val")), name="ins"), sdma_wait),
   (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),), name="ins"), sdma_timestamp),
   (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val")), name="ins"), sdma_store),
@@ -219,7 +216,8 @@ def sdma_submit(cmdbuf, devs):
 
   # the sdma queue's ring and its host-side ring/write/put pointers
   for d in devs: q = Device[d].sdma_queue(0)
-  ring, wptr, doorbell, put_ptr = queue_ptrs(devs, "COPY:0", q)
+  ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COPY:0_{name}")
+    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
 
   # sdma needs the cmdbuf contiguous: if it won't fit before the ring end, restart at 0 and zero the tail
   put_b = put_ptr.index(zero)
@@ -246,32 +244,15 @@ def sdma_submit(cmdbuf, devs):
 pm_sdma_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
   lambda ctx, lin: sdma_submit(make_cmdbuf(lin, ctx.devs), ctx.devs))])
 
-# *****************
-# USB submit
-
-def amd_usb_submit(ctx, lin):
-  for d in ctx.devs: q = Device[d].compute_queue if (comp:=ctx.qname.startswith("COMPUTE")) else Device[d].sdma_queue(0)
-
-  if nb:=usb_arm_bytes(ctx.pre, Device[ctx.devs[0]].iface.usb_sram):
-    poke = (ctx.sdma.SDMA_OP_WRITE, *data64_le(Device[ctx.devs[0]].iface.cq_buf.va_addr + 12), 0, 0)
-    lin = lin.replace(src=lin.src + (UOp(Ops.INS, arg="poke", src=tuple(UOp.const(x, dtypes.uint32) for x in poke)),))
-
-  ib_host, ib_gpu, pkt_dw = usb_ib(ctx.devs, lin, 32 if comp else 0x100, nb)
-  pkt = (ctx.pm4.PACKET3(ctx.pm4.PACKET3_INDIRECT_BUFFER,2),*data64_le(ib_gpu.getaddr(ctx.devs)),pkt_dw|ctx.pm4.INDIRECT_BUFFER_VALID) if comp else ()
-  return usb_push(ctx.devs, *queue_ptrs(ctx.devs, ctx.qname, q), ib_host, ib_gpu, pkt, 4 if comp else 1)
-
-pm_usb_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"), amd_usb_submit)])
-
 @dataclass(frozen=True)
 class AMDEncodeCtx:  # encode-time constants for one queue: devs (every cmdbuf address resolves into these) + gfx version + packet/ip modules
   devs: tuple[str, ...]; target: tuple[int, ...]; pm4: Any; sdma: Any; soc: Any  # noqa: E702
-  gc: AMDIP; nbio: AMDIP; xccs: int; max_copy_size: int; tmpring_size: Callable; qname: str; pre: UOp # pre: the queue before opsel
+  gc: AMDIP; nbio: AMDIP; xccs: int; max_copy_size: int; tmpring_size: Callable  # noqa: E702
 
 def encode_queue(q:UOp) -> UOp|None:
   d = Device[(devs:=to_tuple(q.arg[0]))[0]]
-  ctx = AMDEncodeCtx(devs, d.target, d.pm4, d.sdma, d.soc, d.gc, d.nbio, d.xccs, d.max_copy_size, d.tmpring_size, q.arg[1], q)
-  opsel = pm_pm4_opsel if (comp:=q.arg[1].startswith("COMPUTE")) else pm_sdma_opsel
-  submit = d.pm_submit if d.pm_submit is not None else (pm_pm4_submit if comp else pm_sdma_submit)
+  ctx = AMDEncodeCtx(devs, d.target, d.pm4, d.sdma, d.soc, d.gc, d.nbio, d.xccs, d.max_copy_size, d.tmpring_size)
+  opsel, submit = (pm_pm4_opsel, pm_pm4_submit) if q.arg[1].startswith("COMPUTE") else (pm_sdma_opsel, pm_sdma_submit)
   return submit.rewrite(graph_rewrite(q, opsel + pm_flatten_linear, walk=True, ctx=ctx, name=f"{q.arg[1]} opsel"), ctx)
 
 @dataclass(frozen=True)
@@ -301,14 +282,13 @@ def amd_build_program(prg:UOp) -> UOp:
       wave32=bool(desc.kernel_code_properties & 0x400), private_segment_size=desc.private_segment_fixed_size, kernargs_segment_size=desc.kernarg_size,
       kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0), enable_dispatch_ptr=edp,
       enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER)
-    image = bytes(image).ljust(round_up(len(image), 4), b"\x00") # the program is uploaded as whole dwords
     buf = UOp.placeholder((len(image),), dtypes.uint8, next(UOp.unique_num), device=prg.device).rtag("program")
-    cached = _amd_program_cache[key] = prg.replace(src=(buf.after(make_binary_patch(buf, image)),), arg=(data, prg.arg))
+    cached = _amd_program_cache[key] = prg.replace(src=(buf.after(make_binary_patch(buf, bytes(image))),), arg=(data, prg.arg))
   return cached
 
 class AMDAllocator(HCQAllocator['AMDDevice']):
   def __init__(self, dev:AMDDevice):
-    super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb)
+    super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb())
 
   def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
     return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_copy_queue)
@@ -544,7 +524,8 @@ class PCIIface(PCIIfaceBase):
       cq = d.compute_queue
       for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0
       d.iface.dev_impl.gfx.setup_ring(*cq.params)
-      d.signal('timeline')._buf.cpu_view().view(fmt='Q')[0] = d.signal('value', 1, device="CPU")._buf.cpu_view().view(fmt='Q')[0] - 1
+      d.signal('timeline')._buf.cpu_view().mv.cast('Q')[0] = \
+        d.signal('value', 1).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] - 1
 
   def sleep(self, timeout):
     if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
@@ -558,32 +539,6 @@ class PCIIface(PCIIfaceBase):
 
   def device_fini(self): self.dev_impl.fini()
 
-class USBIface(PCIIface):
-  def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
-    if dev_id >= len(visible:=hcq_filter_visible_devices(USB3.list_devices(0xADD1, 0x0001) + USB3.list_devices(0x3801, 0x0001), "AMD")):
-      raise RuntimeError(f"AMD:{dev_id} does not exist ({pluralize('device', len(visible))} available)")
-    self.dev, self.pci_dev, self.vram_bar, self.count = dev, USBPCIDevice("AM", *visible[dev_id]), 0, len(visible)
-    self.dev_impl = AMDev(self.pci_dev)
-    self._compute_props()
-    self.sram = self._dma_region(ctrl_addr=0xf000, sys_addr=0x200000, size=0x80000)
-    self.cq_buf = self._dma_region(ctrl_addr=0xb800, sys_addr=0x822000, size=0x1000) # +12 is the dword that releases an armed read
-    self.usb_handle = unwrap(ctypes.cast(self.pci_dev.usb.usb.handle, ctypes.c_void_p).value)
-
-  def _dma_region(self, ctrl_addr, sys_addr, size):
-    region = self.dev_impl.mm.map_range(vaddr:=self.dev_impl.mm.alloc_vaddr(size=size), size, [(sys_addr, size)], aspace=AddrSpace.SYS, uncached=True)
-    return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(region, has_cpu_mapping=False), view=self.pci_dev.dma_view(ctrl_addr, size), owner=self.dev)
-
-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
-    # everything, even host-style signals, lives in vram: gpu writes into the bridge's own memory collide with an armed 0xF2 read stream
-    return super().alloc(size, host=False, uncached=uncached, cpu_access=cpu_access or host, contiguous=contiguous, force_devmem=True, **kwargs)
-
-  def sleep(self, timeout): pass
-
-  # we don't own the sram region, so the buffer never frees it
-  @functools.cached_property
-  def usb_sram(self) -> Buffer:
-    return Buffer(self.dev.device, (b:=self.sram).size, dtypes.uint8, options=BufferSpec(external_ptr=b.va_addr, nolru=True)).allocate(opaque=b)
-
 def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {})
 
 class AMDDevice(HCQ2Compiled):
@@ -594,21 +549,19 @@ class AMDDevice(HCQ2Compiled):
     # encoding of cmdbuf
     (UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),)), encode_queue),
   ])
-  pm_submit: PatternMatcher|None = None
 
   timestamp_divider = 100.0  # AMD GPU clock: ticks/us
   max_scratch_psize = 0
 
-  ifaces = [KFDIface, PCIIface, USBIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface), _mock(USBIface)]
+  ifaces = [KFDIface, PCIIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface)]
 
   def device_props(self): return self.iface.props
 
   def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,))
+  def is_usb(self) -> bool: return False
 
   def __init__(self, device:str=""):
     self.iface = self._select_iface(device)
-    self.is_usb = isinstance(self.iface, USBIface)
-    if self.is_usb: self.rt_nbytes = 4 << 20
 
     self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
     self.arch = "gfx%d%x%x" % self.target
@@ -633,7 +586,7 @@ class AMDDevice(HCQ2Compiled):
 
     self.is_aql = getenv("AMD_AQL", int(self.xccs > 1))
     if self.is_aql:
-      self.pm4_ibs = self.iface.alloc(0x2000 if self.is_usb else (16 << 20), uncached=True, cpu_access=True)
+      self.pm4_ibs = self.iface.alloc(0x2000 if self.is_usb() else (16 << 20), uncached=True, cpu_access=True)
       self.pm4_ib_alloc = BumpAllocator(self.pm4_ibs.size, wrap=True)
 
     self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
@@ -646,10 +599,6 @@ class AMDDevice(HCQ2Compiled):
     self.max_private_segment_size = 0
     self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag="scratch", name="b"), lambda ctx, b: ctx[0].scratch_buffer(b.max_numel()))]) + self.pm_bufferize
 
-    if self.is_usb:
-      self.pm_bufferize = pm_usb_bufferize + self.pm_bufferize
-      self.pm_stage_copy, self.pm_host_lower, self.pm_submit = pm_usb_stage, pm_usb_hostio, pm_usb_submit
-
     self.pmc_enabled:bool = PROFILE > 0 and PMC > 0
     if self.pmc_enabled:
       self.iface.require_profile_mode()
@@ -710,7 +659,7 @@ class AMDDevice(HCQ2Compiled):
     wg_data_size = round_up((vgpr_size_per_cu + sgrp_size_per_cu + lds_size_per_cu + hwreg_size_per_cu) * self.cu_cnt, mmap.PAGESIZE)
     ctl_stack_size = round_up((12 if self.target[0] != 9 else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE)
     return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL if self.is_aql else kfd.KFD_IOC_QUEUE_TYPE_COMPUTE,
-      0x2000 if self.is_usb else (16 << 20), eop_buffer_size=0x1000,
+      0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000,
       ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size,
       debug_memory_size=round_up(self.wave_cnt * 32, 64))
 
@@ -718,7 +667,7 @@ class AMDDevice(HCQ2Compiled):
     if getenv("AMD_DISABLE_SDMA"): return None
     if idx in self.sdma_queues: return self.sdma_queues[idx]
     with contextlib.suppress(OSError):
-      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x2000 if self.is_usb else (16 << 20), idx=idx)
+      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20), idx=idx)
     return self.sdma_queues.get(idx, None)
 
   def tmpring_size(self, private_segment_size):

extra/llama_kernels/quantize_fp8_delayed/__init__.py

@@ -50,7 +50,7 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:
   else: raise NotImplementedError(f"no atomic max for device {device}")
   amax_idx = amax_out.reshape((1,)).index(UOp.const(0))
   max_val = lds[0].load()
-  atomic = UOp(Ops.CUSTOM, src=(amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=atomic_arg)
+  atomic = UOp(Ops.CUSTOM, dtypes.void, (amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=atomic_arg)
   return atomic.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))
 
 @functools.cache

extra/llama_kernels/quantize_mxfp4/__init__.py

@@ -12,7 +12,7 @@ def _custom_quantize_mxfp4(row_fp4:UOp, row_scale:UOp, col_fp4:UOp, col_scale:UO
   mem = M*N*2 + M*N + M*N//16  # read bf16, write row+col fp4 + e8m0
   outputs = (row_fp4, row_scale, col_fp4, col_scale)
   sink = UOp.sink(*(o.base for o in outputs), x.base,
-                  *(UOp(Ops.CUSTOM, src=(o.base.index(0),), arg="") for o in outputs),
+                  *(UOp(Ops.CUSTOM, dtypes.void, (o.base.index(0),), arg="") for o in outputs),
                   UOp.special(256, "lidx0"), UOp.special(M//128, "gidx0"), UOp.special(N//64, "gidx1"),
                   arg=KernelInfo(name, estimates=Estimates(ops=12*M*N, mem=mem)))
   src = (pathlib.Path(__file__).parent/"quantize_mxfp4.cpp").read_text()

extra/models/llama.py

@@ -5,9 +5,9 @@ from tinygrad.helpers import getenv, DEBUG
 
 # https://github.com/facebookresearch/llama/blob/1076b9c51c77ad06e9d7ba8a4c6df775741732bd/llama/model.py#L47
 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
-  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2, dtype=dtypes.float32)[:(dim // 2)] / dim))
-  freqs = Tensor.arange(end, dtype=dtypes.float32).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
-  return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).cast(dtypes.default_float).reshape(1, end, 1, dim//2, 2)
+  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2)[:(dim // 2)] / dim))
+  freqs = Tensor.arange(end).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
+  return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).reshape(1, end, 1, dim//2, 2)
 
 # matches meta, non hugging face weights
 # (a+i*b) * (c+i*d) = (ac-bd) + i*(ad+bc)

extra/sqtt/examples/generate_examples.py

@@ -1,4 +1,4 @@
-import os, subprocess, sys, shlex, pickle
+import os, subprocess, sys, shlex
 from pathlib import Path
 from tinygrad.helpers import temp, getenv
 
@@ -23,8 +23,5 @@ if __name__ == "__main__":
       # AM_RESET=1 gets a clear trace, does not work on mi300 machines
       subprocess.run([sys.executable, *shlex.split(test)], cwd=EXAMPLES_DIR.parent.parent.parent,
                      env={**os.environ, "DEV":"AMD", "AM_RESET":"1" if not arch.startswith("gfx9") else "0", "VIZ":"-2", "PYTHONPATH":"."})
-      with open(PROFILE_PATH, "rb") as f: events = pickle.load(f)
-      with open(PROFILE_PATH, "wb") as f:
-        pickle.dump([e for e in events if type(e).__name__ in {"ProfilePMCEvent", "ProfileSQTTEvent", "ProfileProgramEvent"}], f)
       PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{name}_run_{i}.pkl")
       print(f"saved SQTT trace to {dest}")

extra/sqtt/roc.py

@@ -192,7 +192,7 @@ def unpack_insts(viz_data, i:int, j:int, data:dict) -> dict:
     prev_instr = max(prev_instr, e.time + e.dur)
   summary = [{"label":"Total Cycles", "value":w.end_time-w.begin_time}, {"label":"SE", "value":w.se}, {"label":"CU", "value":w.cu},
              {"label":"SIMD", "value":w.simd}, {"label":"Wave ID", "value":w.wave_id}, {"label":"Run number", "value":data["run_number"]}]
-  return {"rows":[tuple(v.values()) for v in rows.values()], "cols":columns, "metadata":[summary],"ref":viz_data.ref_map.get(data["prg"].profile_key)}
+  return {"rows":[tuple(v.values()) for v in rows.values()], "cols":columns, "metadata":[summary], "ref":viz_data.ref_map.get(data["prg"].name)}
 
 def print_data(data:dict) -> None:
   from tabulate import tabulate

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tinygrad"
-version = "0.14.0"
+version = "0.13.0"
 description = "You like pytorch? You like micrograd? You love tinygrad! <3"
 authors = [{ name = "George Hotz" }]
 
@@ -84,7 +84,7 @@ testing = [
   "pillow",
   "onnx==1.19.0",
   "onnx2torch",
-  "onnxruntime==1.24.1",
+  "onnxruntime",
   "opencv-python",
   "transformers",
   "sentencepiece",

test/amd/hw/test_sop.py

@@ -1002,39 +1002,6 @@ class TestBarrier(unittest.TestCase):
     for tid in range(64):
       self.assertEqual(st.vgpr[tid][0], tid + 100 + 1000, f"tid={tid}")
 
-class TestSMaxMinSCCRegressions(unittest.TestCase):
-  """Regression test: S_MAX sets SCC only on strict inequality (equal operands -> SCC=0)."""
-
-  def test_s_max_i32_equal_scc(self):
-    st = run_program([s_mov_b32(s[4], 64), s_mov_b32(s[5], 64), s_max_i32(s[6], s[4], s[5])], n_lanes=1)
-    self.assertEqual(st.scc, 0)
-    self.assertEqual(st.sgpr[6], 64)
-    st = run_program([s_mov_b32(s[4], 65), s_mov_b32(s[5], 64), s_max_i32(s[6], s[4], s[5])], n_lanes=1)
-    self.assertEqual(st.scc, 1)  # still set when strictly greater
-
-  def test_s_max_u32_equal_scc(self):
-    st = run_program([s_mov_b32(s[4], 64), s_mov_b32(s[5], 64), s_max_u32(s[6], s[4], s[5])], n_lanes=1)
-    self.assertEqual(st.scc, 0)
-
-class TestAbsdiffOverflowRegressions(unittest.TestCase):
-  """Regression test: S_ABSDIFF_I32 computes abs on the WRAPPED 32-bit difference (found by random difftest vs hardware)."""
-
-  def test_s_absdiff_wrapped(self):
-    # |45 - (-2147483647)| overflows int32; hardware takes abs of the wrapped 32-bit difference
-    instructions = [s_mov_b32(s[4], 45), s_mov_b32(s[5], 0x80000001), s_absdiff_i32(s[6], s[4], s[5])]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[6], 0x7FFFFFD4)
-    self.assertEqual(st.scc, 1)
-    # INT_MIN - 1 wraps to +2147483647, already positive
-    instructions = [s_mov_b32(s[4], 0x80000000), s_mov_b32(s[5], 1), s_absdiff_i32(s[6], s[4], s[5])]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[6], 0x7FFFFFFF)
-    # equality -> 0 and SCC=0
-    instructions = [s_mov_b32(s[4], 7), s_mov_b32(s[5], 7), s_absdiff_i32(s[6], s[4], s[5])]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[6], 0)
-    self.assertEqual(st.scc, 0)
-
 
 if __name__ == '__main__':
   unittest.main()

test/amd/hw/test_vop1.py

@@ -1629,66 +1629,5 @@ class TestSwap(unittest.TestCase):
     self.assertEqual(st.vgpr[0][1], 0x55555555)
 
 
-class TestCvtFrexpRegressions(unittest.TestCase):
-  """Regression tests for float<->int conversion and FREXP corner cases (found by random difftest vs hardware)."""
-
-  def test_cvt_i32_f32_nan_is_zero(self):
-    """v_cvt_i32_f32 of NaN is 0, not INT_MIN (x86 cvttss2si returns INT_MIN)."""
-    for nan in (0x7FC00000, 0xFFC00000, 0x7F800001):
-      st = run_program([v_mov_b32_e32(v[0], nan), v_cvt_i32_f32_e32(v[1], v[0])], n_lanes=1)
-      self.assertEqual(st.vgpr[0][1], 0, f"nan=0x{nan:08x}")
-
-  def test_cvt_i32_f32_positive_overflow(self):
-    """v_cvt_i32_f32 saturates positive overflow/inf to INT_MAX, not INT_MIN."""
-    for bits in (0x7F800000, 0x4F000000, 0x4F800000):  # +inf, 2^31, ~2^32
-      st = run_program([v_mov_b32_e32(v[0], bits), v_cvt_i32_f32_e32(v[1], v[0])], n_lanes=1)
-      self.assertEqual(st.vgpr[0][1], 0x7FFFFFFF, f"bits=0x{bits:08x}")
-
-  def test_cvt_i32_f32_negative_overflow(self):
-    """v_cvt_i32_f32 saturates negative overflow/-inf to INT_MIN."""
-    for bits in (0xFF800000, 0xCF000001):  # -inf, below -2^31
-      st = run_program([v_mov_b32_e32(v[0], bits), v_cvt_i32_f32_e32(v[1], v[0])], n_lanes=1)
-      self.assertEqual(st.vgpr[0][1], 0x80000000, f"bits=0x{bits:08x}")
-
-  def test_cvt_u32_f32_nan_is_zero(self):
-    """v_cvt_u32_f32 of NaN is 0, not UINT_MAX."""
-    for nan in (0x7FC00000, 0xFFC00000, 0x7F800001):
-      st = run_program([v_mov_b32_e32(v[0], nan), v_cvt_u32_f32_e32(v[1], v[0])], n_lanes=1)
-      self.assertEqual(st.vgpr[0][1], 0, f"nan=0x{nan:08x}")
-
-  def test_cvt_i32_f64_nan_and_overflow(self):
-    """v_cvt_i32_f64: NaN -> 0, positive overflow/+inf -> INT_MAX."""
-    st = run_program([v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0x7FF80000), v_cvt_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0)
-    st = run_program([v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0x41F00000), v_cvt_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x7FFFFFFF)  # 2^32 -> INT_MAX
-
-  def test_frexp_f32_denormal(self):
-    """v_frexp_exp/mant_f32 of denormal/zero inputs is (0, signed zero) on hardware."""
-    for bits in (0x00000001, 0x007FFFFF, 0x00000000):
-      st = run_program([v_mov_b32_e32(v[0], bits), v_frexp_exp_i32_f32_e32(v[1], v[0]), v_frexp_mant_f32_e32(v[2], v[0])], n_lanes=1)
-      self.assertEqual(st.vgpr[0][1] & 0xFFFFFFFF, 0, f"exp bits=0x{bits:08x}")
-      self.assertEqual(st.vgpr[0][2], bits & 0x80000000, f"mant bits=0x{bits:08x}")
-    # negative denormal: mant is -0.0
-    st = run_program([v_mov_b32_e32(v[0], 0x80000001), v_frexp_mant_f32_e32(v[2], v[0])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x80000000)
-
-  def test_frexp_f64_denormal(self):
-    """v_frexp_exp_f64 of a denormal returns the normalized exponent (-1073 for min-denormal); zero -> 0."""
-    st = run_program([v_mov_b32_e32(v[0], 1), v_mov_b32_e32(v[1], 0), v_frexp_exp_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2] & 0xFFFFFFFF, 0xFFFFFBCF)  # -1073
-    st = run_program([v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_frexp_exp_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0)
-
-  def test_frexp_exp_inf_nan(self):
-    """v_frexp_exp of +/-inf and NaN is 0 on hardware (host frexp gives 129/1024), for both f32 and f64."""
-    for bits in (0x7F800000, 0xFF800000, 0x7FC00000):
-      st = run_program([v_mov_b32_e32(v[0], bits), v_frexp_exp_i32_f32_e32(v[1], v[0])], n_lanes=1)
-      self.assertEqual(st.vgpr[0][1] & 0xFFFFFFFF, 0, f"f32 bits=0x{bits:08x}")
-    for lo, hi in ((0, 0x7FF00000), (0, 0xFFF00000), (0, 0x7FF80000), (1, 0x7FF00000)):
-      st = run_program([v_mov_b32_e32(v[0], lo), v_mov_b32_e32(v[1], hi), v_frexp_exp_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
-      self.assertEqual(st.vgpr[0][2] & 0xFFFFFFFF, 0, f"f64 bits=0x{hi:08x}{lo:08x}")
-
-
 if __name__ == '__main__':
   unittest.main()

test/amd/hw/test_vop2.py

@@ -989,53 +989,6 @@ class TestCarryOps(unittest.TestCase):
     self.assertEqual(st.vgpr[0][0], 0)  # 0xFFFFFFFF + 1 + 0 = 0 (overflow)
     self.assertEqual(st.vcc, 0xDEADBEEF)  # VCC unchanged - carry was discarded
 
-class TestSelectFlushRegressions(unittest.TestCase):
-  """Regression tests: f32 MIN/MAX flush denormal inputs to signed zero (select-style ops propagate inputs bitwise)."""
-
-  def test_v_min_f32_denormal_flush(self):
-    """min(denormal, 1.0) is +0, min(-denormal, -1.0) is -0."""
-    st = run_program([v_mov_b32_e32(v[0], 0x00000001), v_mov_b32_e32(v[1], 0x3F800000), v_min_f32_e32(v[2], v[0], v[1])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x00000000)
-    # flush(-denormal) = -0.0 > -1.0, so the result is -1.0 (both operand orders)
-    st = run_program([v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0xBF800000), v_min_f32_e32(v[2], v[0], v[1])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0xBF800000)
-    st = run_program([v_mov_b32_e32(v[1], 0xBF800000), v_mov_b32_e32(v[2], 0x80000001), v_min_f32_e32(v[3], v[1], v[2])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][3], 0xBF800000)
-
-  def test_v_max_f32_denormal_flush(self):
-    """max(-denormal, -1.0) is -0; max(+denormal, -0) is +0."""
-    st = run_program([v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0xBF800000), v_max_f32_e32(v[2], v[0], v[1])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x80000000)
-    st = run_program([v_mov_b32_e32(v[0], 0x00000001), v_mov_b32_e32(v[1], 0x80000000), v_max_f32_e32(v[2], v[0], v[1])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x00000000)
-
-
-class TestCarryExecRegressions(unittest.TestCase):
-  """Regression tests: per-lane VCC writes (carry ops) zero inactive lane bits - VCC = mask & EXEC, never preserved."""
-
-  def test_co_ci_e32_vcc_masked_by_exec(self):
-    """v_sub_co_ci_u32_e32 with EXEC=0xFFFF0000: hw clears inactive VCC bits instead of preserving them."""
-    instructions = [
-      s_mov_b32(EXEC_LO, 0xFFFF0000),
-      s_mov_b32(VCC_LO, 0xFFFFFFFF),  # preset all bits
-      v_mov_b32_e32(v[0], 0xFFFFFFFE), v_mov_b32_e32(v[1], 0x80000000),
-      v_sub_co_ci_u32_e32(v[2], v[0], v[1]),  # active lanes: no borrow
-    ]
-    st = run_program(instructions, n_lanes=32)
-    self.assertEqual(st.vcc, 0x00000000)
-
-  def test_co_ci_e32_vcc_masked_by_exec_ones(self):
-    """Same with all-ones carry: VCC = borrow_mask & EXEC."""
-    instructions = [
-      s_mov_b32(EXEC_LO, 0x0F0F0F0F),
-      s_mov_b32(VCC_LO, 0),
-      v_mov_b32_e32(v[0], 0xFFFFFFFF), v_mov_b32_e32(v[1], 1),
-      v_add_co_ci_u32_e32(v[2], v[0], v[1]),  # all lanes would carry if active
-    ]
-    st = run_program(instructions, n_lanes=32)
-    self.assertEqual(st.vcc, 0x0F0F0F0F)
-    self.assertEqual(st.vgpr[31][2], 0)  # 0xFFFFFFFF + 1 wraps to 0 in active lanes
-
 
 if __name__ == '__main__':
   unittest.main()

test/amd/hw/test_vop3.py

@@ -4,7 +4,6 @@ Includes: v_fma_f32, v_div_scale_f32, v_div_fmas_f32, v_div_fixup_f32,
           v_alignbit_b32, v_bfe_i32, v_mad_u64_u32, v_readlane_b32, v_writelane_b32
 """
 import unittest
-from tinygrad.helpers import OSX
 from test.amd.hw.helpers import *
 
 class TestFMA(unittest.TestCase):
@@ -3265,23 +3264,6 @@ class TestVOP3ClampMAD(unittest.TestCase):
     # 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 TestMadNarrowClampRegressions(unittest.TestCase):
-  """Regression tests: mad i16/i24 with clamp saturate to narrow output range (found by random difftest vs hardware)."""
-
-  def test_mad_i16_clamp_sat_max(self):
-    # neg/src-floggled 16-bit mul operands are sign-extended after toggling bit15; sum > INT_MAX saturates
-    instructions = [s_mov_b32(s[4], 1232348160), v_mov_b32_e32(v[3], 0x80000000),
-                    v_mov_b32_e32(v[1], 0x7F7FFFFF), v_mad_i32_i16(v[0], s[4], v[3], v[1], 0, 3, 5, 1)]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0x7FFFFFFF)
-
-  def test_mad_i24_clamp_sat_min(self):
-    # sext24(-6344704) * sext24(+4210688) << -2^31 saturates to INT_MIN
-    instructions = [s_mov_b32(s[7], 4290772992), v_mov_b32_e32(v[1], 1077936128),
-                    v_mad_i32_i24(v[0], s[7], v[1], v[1], 1, 0, 0, 1)]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0x80000000)
-
 
 class TestCvtPkF16(unittest.TestCase):
   """Tests for V_CVT_PK_RTZ_F16_F32 - pack two f32 to f16 with round toward zero."""
@@ -3669,80 +3651,6 @@ class TestPermlane(unittest.TestCase):
     self.assertEqual(st.vgpr[21][1], 5)
     self.assertEqual(st.vgpr[31][1], 15)
 
-class TestClampLdExpRegressions(unittest.TestCase):
-  """Regression tests for f32 clamp (-0 -> +0) and ldexp input passthrough."""
-
-  def test_clamp_negative_zero(self):
-    """clmp=1 maps -0.0 to +0.0 (found by random difftest vs hardware)."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0x80000000), v_mov_b32_e32(v[1], 0x80000000),
-      v_add_f32_e64(v[2], v[0], v[1], clmp=1),  # -0 + -0 = -0, clamp -> +0
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x00000000)
-    instructions = [
-      v_mov_b32_e32(v[0], 0x3F800000), v_mov_b32_e32(v[1], 0x80000000),
-      v_min_f32_e64(v[2], v[0], v[1], clmp=1),  # min(1.0, -0) = -0, clamp -> +0
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x00000000)
-
-  def test_ldexp_special_inputs(self):
-    """v_ldexp_f32 of 0/-0/inf/NaN propagates the input instead of computing val * 2**exp (0*inf = NaN on host)."""
-    # -0.0 * 2^INT_MIN = -0.0 (src1 as integer exponent; huge negative)
-    instructions = [v_mov_b32_e32(v[0], 0x80000000), v_mov_b32_e32(v[1], 0x80000000), v_ldexp_f32(v[2], v[0], v[1])]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x80000000)
-    # inf stays inf even with negative exponent
-    instructions = [v_mov_b32_e32(v[0], 0x7F800000), v_mov_b32_e32(v[1], 0xFFFFFF80), v_ldexp_f32(v[2], v[0], v[1])]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x7F800000)
-
-  def test_ldexp_denormal_flush(self):
-    """v_ldexp_f32/f64 flush denormal inputs to signed zero (found by random difftest vs hardware)."""
-    # ldexp(+denorm, 1) = +0, ldexp(-denorm, 250) = -0
-    for src, exp_val, want in [(0x00000001, 1, 0x00000000), (0x80000001, 250, 0x80000000)]:
-      st = run_program([v_mov_b32_e32(v[0], src), v_mov_b32_e32(v[1], exp_val), v_ldexp_f32(v[2], v[0], v[1])], n_lanes=1)
-      self.assertEqual(st.vgpr[0][2], want)
-
-  def test_v_mul_neg_modifier_nan_sign(self):
-    """neg modifier is a pure sign-bit toggle on a NaN operand; result keeps that sign (found by random difftest)."""
-    # mul(normal, NEG(ABS(qNaN))): NaN payload negated in the operand stays negative qNaN
-    instructions = [v_mov_b32_e32(v[0], 0xC96CF47F), v_mov_b32_e32(v[1], 0x7FC00000),
-                    v_mul_f32_e64(v[2], v[0], v[1], s[0], 0, 7, 6)]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0xFFC00000)
-    # plain neg modifier still applies to non-NaN values: mul(-1.0, NEG(2.0)) = +2.0
-    st = run_program([v_mov_b32_e32(v[0], 0xBF800000), v_mov_b32_e32(v[1], 0x40000000),
-                      v_mul_f32_e64(v[2], v[0], v[1], s[0], 0, 2, 0)], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x40000000)
-
-
-class TestNaNPropagationRegressions(unittest.TestCase):
-  """Regression tests: float arithmetic propagates a NaN from the FIRST NaN operand, quieted with its own sign/payload."""
-
-  @unittest.skipIf(OSX, "broken on mac, TODO: why?")
-  def test_mul_nan_priority(self):
-    # first NaN operand wins (sign+payload), not x86's second-source propagation
-    for a, b, want in [(0x7FC00001, 0x7F800003, 0x7FC00001), (0xFFC00005, 0x7F800003, 0xFFC00005),
-                       (0x7F800001, 0xFFC00005, 0x7FC00001), (0xFF9F1800, 0x7F800001, 0xFFDF1800)]:
-      st = run_program([v_mov_b32_e32(v[0], a), v_mov_b32_e32(v[1], b),
-                        v_mul_f32_e32(v[2], v[0], v[1])], n_lanes=1)
-      self.assertEqual(st.vgpr[0][2], want, f"mul({a:#x}, {b:#x})")
-
-class TestMinMaxFlushE64Regressions(unittest.TestCase):
-  """Regression tests: f32 min/max/median flush denormal inputs to signed zero (e64 forms)."""
-
-  def test_v_min3_f32_denormal_flush(self):
-    st = run_program([v_mov_b32_e32(v[0], 0x00000001), v_mov_b32_e32(v[1], 0x3F800000), v_mov_b32_e32(v[2], 0x40000000),
-                      v_min3_f32(v[3], v[0], v[1], v[2])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][3], 0x00000000)  # min(+denorm, 1, 2) = +0
-
-  def test_v_med3_f32_denormal_flush(self):
-    st = run_program([v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0x3F800000), v_mov_b32_e32(v[2], 0x40000000),
-                      v_med3_f32(v[3], v[0], v[1], v[2])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][3], 0x3F800000)  # med(-0, 1, 2) = 1
-
 
 if __name__ == '__main__':
   unittest.main()

test/amd/hw/test_vopc.py

@@ -973,71 +973,6 @@ class TestCmpxPartialWavefront(unittest.TestCase):
     self.assertEqual(st.sgpr[EXEC_LO.offset] & 0xFFFFFFFF, 0x4,
                      "Only lane 2 should be active after v_cmpx_eq_u32_e64")
 
-class TestClassDenormalRegressions(unittest.TestCase):
-  """Regression tests: V_CMP_CLASS classifies denormals as DENORMAL (raw bits), not as zero class."""
-
-  def test_class_pos_denormal(self):
-    for bits in (0x00000001, 0x007FFFFF):
-      instructions = [v_mov_b32_e32(v[0], bits), v_mov_b32_e32(v[1], 0x80), v_cmp_class_f32_e64(VCC_LO, v[0], v[1])]
-      st = run_program(instructions, n_lanes=1)
-      self.assertEqual(st.vcc, 1, f"bits=0x{bits:08x}")  # n_lanes=1
-      # ...and it is not the zero class
-      instructions = [v_mov_b32_e32(v[0], bits), v_mov_b32_e32(v[1], 0x40), v_cmp_class_f32_e64(VCC_LO, v[0], v[1])]
-      st = run_program(instructions, n_lanes=1)
-      self.assertEqual(st.vcc, 0, f"bits=0x{bits:08x}")
-
-  def test_class_neg_denormal(self):
-    instructions = [v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0x10), v_cmp_class_f32_e64(VCC_LO, v[0], v[1])]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc, 1)  # n_lanes=1
-    instructions = [v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0x20), v_cmp_class_f32_e64(VCC_LO, v[0], v[1])]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc, 0)  # not the negative-zero class
-
-
-class TestIntCmpModRegressions(unittest.TestCase):
-  """Regression tests: int compares (i32/u32) honor abs/neg as bit-level sign clear/flip (not integer abs/negate)."""
-
-  def test_cmp_i32_abs_neg_bit_level(self):
-    # abs(0x80000001) = 1 -> 1 > 1 is false (integer abs would give 2147483647 > 1)
-    instructions = [v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 1), v_cmp_gt_i32_e64(VCC_LO, v[0], v[1], abs=1)]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc, 0)
-    # neg(0x80000001) flips the sign bit -> 1 > 2 is false (integer negate would give 2147483647 > 2)
-    instructions = [v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 2), v_cmp_gt_i32_e64(VCC_LO, v[0], v[1], neg=1)]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc, 0)
-
-  def test_cmp_u32_abs_bit_level(self):
-    # abs(0x80000000) = 0 -> 0 < 1 is true
-    instructions = [v_mov_b32_e32(v[0], 0x80000000), v_mov_b32_e32(v[1], 1), v_cmp_lt_u32_e64(VCC_LO, v[0], v[1], abs=1)]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc, 1)  # n_lanes=1
-
-
-class TestCmpxSdstRegressions(unittest.TestCase):
-  """Regression tests: V_CMPX_*_E64 writes EXEC only, never SDST (hardware verified)."""
-
-  def test_cmpx_e64_no_sdst(self):
-    instructions = [
-      s_mov_b32(VCC_LO, 0),  # preset VCC to 0
-      v_mov_b32_e32(v[0], 0x3F800000), v_mov_b32_e32(v[1], 0x40000000),
-      v_cmpx_lt_f32_e64(VCC_LO, v[0], v[1]),  # 1.0 < 2.0
-    ]
-    st = run_program(instructions, n_lanes=32)
-    self.assertEqual(st.sgpr[EXEC_LO.offset], 0xFFFFFFFF)  # EXEC updated
-    self.assertEqual(st.vcc, 0)  # but VCC untouched
-
-  def test_cmpx_e64_partial_exec(self):
-    instructions = [
-      s_mov_b32(EXEC_LO, 0x0F0F0F0F),
-      s_mov_b32(VCC_LO, 0xFFFFFFFF),
-      v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0x3F800000),
-      v_cmpx_lt_f32_e64(VCC_LO, v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=32)
-    self.assertEqual(st.sgpr[EXEC_LO.offset], 0x0F0F0F0F)  # EXEC = computed & old EXEC
-
 
 if __name__ == '__main__':
   unittest.main()

test/amd/test_sqtt_examples.py

@@ -88,6 +88,7 @@ def run_rocprof_decoder(blobs: list[bytes], lib: bytes, base: int, target: str):
   if t.is_alive(): raise RuntimeError("rocprof decoder timeout")
   return occupancy_records, wave_insts
 
+@unittest.skip("TODO: fix to not require unpickling UOps.")
 class SQTTExamplesTestBase(unittest.TestCase):
   target: str
   examples: dict

test/amd/test_sqtt_profiler.py

@@ -1,28 +1,30 @@
 import unittest, contextlib
 from tinygrad import Device, Tensor, Context, TinyJit
-from tinygrad.device import Compiled, ProfileProgramEvent
+from tinygrad.device import Compiled, ProfileProgramEvent, ProfileDeviceEvent
 from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
 from tinygrad.viz.serve import load_amd_counters, VizData
 
 @contextlib.contextmanager
 def save_sqtt():
-  Device[Device.DEFAULT].synchronize()
-  profile_start = len(Compiled.profile_events)
   data = VizData()
   yield data.ctxs
   Device[Device.DEFAULT].synchronize()
   Device[Device.DEFAULT]._at_profile_finalize()
-  load_amd_counters(data, [e for e in Compiled.profile_events[:profile_start] if isinstance(e, ProfileProgramEvent)] +
-                          Compiled.profile_events[profile_start:])
+  load_amd_counters(data, Compiled.profile_events)
   data.ctxs[:] = [r for r in data.ctxs if r["name"].startswith("SQTT")]
 
 @unittest.skipUnless(Device.DEFAULT == "AMD", "only runs on AMD")
 class TestSQTTProfiler(unittest.TestCase):
+  # TODO: can we enable SQTT profiling in context?
   @classmethod
   def setUpClass(cls):
     if not Device[Device.DEFAULT].sqtt_enabled: raise unittest.SkipTest("device must be in SQTT profiling mode")
 
+  def setUp(self):
+    Device[Device.DEFAULT].synchronize()
+    Compiled.profile_events[:] = [e for e in Compiled.profile_events if isinstance(e, (ProfileProgramEvent, ProfileDeviceEvent))]
+
   def test_simple(self):
     t = Tensor.empty(1) + 1
     with save_sqtt() as sqtt:

test/backend/test_isel.py

@@ -7,7 +7,7 @@ from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
 from tinygrad.renderer.isa import IselContext
 
 # INDEX on a register value with a constant index extracts a single element (the old GEP)
-def lane(y:UOp, i:int) -> UOp: return y.index(UOp.cconst(i, dtypes.int))
+def lane(y:UOp, i:int) -> UOp: return y.index(UOp.cconst(i, dtypes.int), dtype=y.dtype)
 
 @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86")
 class TestIselX86(unittest.TestCase):

test/backend/test_llama_kernels.py

@@ -9,7 +9,7 @@ from extra.llama_kernels.swiglu import swiglu
 from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
 from extra.thunder.amd.fa import custom_fused_qkv_rope_backward, fused_qkv_rope
 from test.helpers import needs_second_gpu, assert_kernel_count
-from test.backend.test_asm_gemm import has_hipcc, is_cdna4
+from test.backend.test_asm_gemm import has_hipcc
 
 def run_fused_ce(bs:int, seqlen:int, vocab:int, label_smoothing:float=0.0) -> None:
   Tensor.manual_seed(0)
@@ -129,7 +129,7 @@ class TestFusedQKVRoPE(unittest.TestCase):
       self.assertTrue(k.allclose(k_ref, atol=2e-2, rtol=0).item(), "K forward mismatch")
       self.assertTrue(v.allclose(v_ref, atol=0, rtol=0).item(), "V forward mismatch")
 
-  @unittest.skipUnless(has_hipcc() and is_cdna4(), "backward kernel requires hipcc to compile")
+  @unittest.skipUnless(has_hipcc(), "backward kernel requires hipcc to compile")
   def test_llama31_8b(self):
     Tensor.manual_seed(1)
     B, N, H, H_KV, D = self.SHAPE

test/backend/test_multitensor.py

@@ -3,7 +3,7 @@ from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variab
 from tinygrad.uop.ops import Ops, UOp, AxisType, graph_rewrite
 from tinygrad.helpers import getenv, prod, Context
 from tinygrad.nn.state import get_parameters
-from tinygrad.engine.realize import run_linear, compile_linear, lower_and_compile, pm_beam
+from tinygrad.engine.realize import run_linear, compile_linear, pm_beam, pm_compile
 import numpy as np
 from hypothesis import given, strategies as strat, settings
 from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count, KernelCountException
@@ -80,7 +80,7 @@ class TestMultiTensor(unittest.TestCase):
     cpu_2 = ("CPU:1", "CPU:2")
     src = Tensor.ones(16).shard(cpu_2, 0).realize()
     lin = UOp(Ops.LINEAR, src=(src.to(cpu_2[::-1]).schedule_linear().src[0],))
-    with Context(BEAM=1, IGNORE_BEAM_CACHE=1): call = lower_and_compile(graph_rewrite(lin, pm_beam, ctx=1, walk=True)).src[0]
+    with Context(BEAM=1, IGNORE_BEAM_CACHE=1): call = graph_rewrite(graph_rewrite(lin, pm_beam, ctx=1, walk=True), pm_compile, walk=True).src[0]
     self.assertNotEqual(call.src[0].src[0].arg.applied_opts, ())
 
   def test_shard_same_device(self):

test/backend/test_ops.py

@@ -865,9 +865,9 @@ class TestOps(unittest.TestCase):
                    lambda: (ten << Tensor([0,2,4], dtype=dtypes.uint32)).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor.__lshift__(2), lambda: ten.__lshift__(2).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor.bitwise_left_shift(2), lambda: ten.lshift(2).cast(dtypes.int32), forward_only=True)
-    self.helper_test_exception([], lambda: torch.tensor([1.0]) << 2, lambda: (Tensor([1.0]) << 2).realize(), expected=RuntimeError)
-    self.helper_test_exception([], lambda: tor << torch.tensor([1.0]), lambda: (ten << Tensor([1.0])).realize(), expected=RuntimeError)
-    self.helper_test_exception([], lambda: tor << 1.0, lambda: (ten << 1.0).realize(), expected=RuntimeError)
+    self.helper_test_exception([], lambda: torch.tensor([1.0]) << 2, lambda: Tensor([1.0]) << 2, expected=RuntimeError)
+    self.helper_test_exception([], lambda: tor << torch.tensor([1.0]), lambda: ten << Tensor([1.0]), expected=RuntimeError)
+    self.helper_test_exception([], lambda: tor << 1.0, lambda: ten << 1.0, expected=RuntimeError)
 
   def test_rshift(self):
     data = [[0,1,2],[1<<8,1<<16,1<<31-1]]
@@ -881,8 +881,8 @@ class TestOps(unittest.TestCase):
                    lambda: (ten >> Tensor([0,2,4], dtype=dtypes.uint32)).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor.__rshift__(2), lambda: ten.__rshift__(2).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor.bitwise_right_shift(2), lambda: ten.rshift(2).cast(dtypes.int32), forward_only=True)
-    self.helper_test_exception([], lambda: torch.tensor([4.0]) >> 1, lambda: (Tensor([4.0]) >> 1).realize(), expected=RuntimeError)
-    self.helper_test_exception([], lambda: tor >> torch.tensor([1.0]), lambda: (ten >> Tensor([1.0])).realize(), expected=RuntimeError)
+    self.helper_test_exception([], lambda: torch.tensor([4.0]) >> 1, lambda: Tensor([4.0]) >> 1, expected=RuntimeError)
+    self.helper_test_exception([], lambda: tor >> torch.tensor([1.0]), lambda: ten >> Tensor([1.0]), expected=RuntimeError)
 
   def test_lshift_signed(self):
     data = [[-1, -3, 1, 7], [0, -2147483648, 2147483647, -1]]

test/backend/test_optim.py

@@ -87,8 +87,7 @@ class TestOptim(unittest.TestCase):
   def test_muon(self): self._test_muon(1, {'lr': 0.001}, 1e-3, 0)
   # TODO: disabled due to big atol
   # def test_muon_high_lr(self): self._test_muon(1, {'lr': 10}, 1e-6, 3e-4)
-  # NOTE: big weight_decay so a missing wd would be way over atol
-  def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 10}, 1e-3, 3e-4)
+  def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 0.01}, 1e-3, 3e-4)
   # TODO: disabled due to big atol
   # def test_muon_high_lr_wd(self): self._test_muon(1, {'lr': 10, 'weight_decay': 0.01}, 1e-6, 5e-4)
 

test/backend/test_renderer_failures.py

@@ -77,14 +77,6 @@ class TestCStyleFailures(unittest.TestCase):
 
 @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "tests for wgsl renderer")
 class TestWGSLFailures(unittest.TestCase):
-  def test_folded_packed_store(self):
-    b = UOp.param(0, dtypes.char, (4,))
-    idx = b.index(UOp.const(0).cast(dtypes.int))
-    store = UOp.store(idx, idx.cast(dtypes.uint32).load() & UOp.const(0xffffff00).cast(dtypes.uint32))
-    src = Device[Device.DEFAULT].renderer.render(UOp.sink(store, arg=KernelInfo()).toposort())
-    self.assertIn("atomicAnd(&data0_4[0],4294967040u);", src)
-    self.assertNotIn("atomicAdd", src)
-
   def test_multiply_infinity(self):
     # multiplying a positive constant by infinity should return infinity
     # WGSL pipelines do not handle this reliably, some of which return zero, unless infinity always comes from a read on a dynamic buffer

test/backend/test_uops.py

@@ -2,7 +2,7 @@ from typing import Optional, Any
 import unittest, math
 import numpy as np
 from tinygrad.tensor import Tensor, _to_np_dtype
-from tinygrad.helpers import Context, ceildiv
+from tinygrad.helpers import Context
 from tinygrad.dtype import dtypes, DType, AddrSpace, ConstFloat  # noqa: F401
 from tinygrad.device import Buffer, Device
 from tinygrad.uop.ops import Ops, UOp, KernelInfo, AxisType, buffers
@@ -57,35 +57,6 @@ def _test_uops_result(output_dtype, uops, res):
   run_uops([out], [buf])
   return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]
 
-@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage) and
-                     dtypes.uint64 in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires C-style pointer bitcast and 64-bit ints")
-class TestBitcastBufferView(unittest.TestCase):
-  @Context(SPEC=2)
-  def test_render(self):
-    buf = UOp.param(0, dtypes.uint32, (4,))
-    uops = to_uops_list([buf.shrink(((1, 3),)).bitcast(dtypes.uint64).index(0).store(1)], ren=Device[Device.DEFAULT].renderer)
-    idx = next(u for u in uops if u.op is Ops.INDEX and u.src[0].op is Ops.BITCAST)
-    self.assertEqual(idx.src[0].src[0].op, Ops.SHRINK)
-    Device[Device.DEFAULT].renderer.render(uops)
-
-  @Context(SPEC=2)
-  def test_load(self):
-    val = 0x1122334455667788
-    src, out = UOp.param(0, dtypes.uint32, (4,)), UOp.param(1, dtypes.uint64, (1,))
-    ibuf = Buffer(Device.DEFAULT, 4, dtypes.uint32, initial_value=np.array([0, 0x55667788, 0x11223344, 0], dtype=np.uint32).tobytes())
-    obuf = Buffer(Device.DEFAULT, 1, dtypes.uint64).allocate()
-    run_uops([out.index(0).store(src.shrink(((1, 3),)).bitcast(dtypes.uint64).index(0))], [ibuf, obuf])
-    self.assertEqual(np.frombuffer(obuf.as_memoryview(), dtype=np.uint64)[0], val)
-
-  @Context(SPEC=2)
-  def test_store(self):
-    val = 0x1122334455667788
-    dst = UOp.param(0, dtypes.uint32, (6,))
-    buf = Buffer(Device.DEFAULT, 6, dtypes.uint32, initial_value=bytes(24))
-    view = dst.shrink(((1, 5),)).bitcast(dtypes.uint64)  # two stores through one view: it must inline, not get a declared vector-pointer
-    run_uops([view.index(0).store(val ^ 0xff), view.index(1).store(val)], [buf])
-    self.assertEqual(np.frombuffer(buf.as_memoryview(), dtype=np.uint64, count=2, offset=4).tolist(), [val ^ 0xff, val])
-
 class TestUOps(unittest.TestCase):
   def _equal(self, v1, v2):
     assert isinstance(v2, (float, int, bool))
@@ -222,16 +193,15 @@ class TestLocalAccess(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "WEBGPU", "Test local memory size for packed data types")
   def test_packed_smem_size(self):
     _dtypes = [dtypes.char, dtypes.uchar, dtypes.short, dtypes.ushort, dtypes.half]
-    # a partial word still needs a whole word, so sizes that don't fill one must round up
-    for size in (16, 5):
-      for dtype in _dtypes:
-        temp = UOp.placeholder((size,), dtype, slot=0, addrspace=AddrSpace.LOCAL)
-        uops = to_uops_list([temp], ren=Device[Device.DEFAULT].renderer)
-        out = Device[Device.DEFAULT].renderer.render(uops)
-        # half is supported in wgsl, so it doesn't have to be packed
-        corrected_size = ceildiv(size, 4//dtype.itemsize) if dtype != dtypes.half else size
-        # temp0: array<{Device[Device.DEFAULT].renderer.buf_map(dtype)},{corrected_size}>;
-        self.assertIn(f",{corrected_size}>;", out)
+    size = 16
+    for dtype in _dtypes:
+      temp = UOp.placeholder((size,), dtype, slot=0, addrspace=AddrSpace.LOCAL)
+      uops = to_uops_list([temp], ren=Device[Device.DEFAULT].renderer)
+      out = Device[Device.DEFAULT].renderer.render(uops)
+      # half is supported in wgsl, so it doesn't have to be packed
+      corrected_size = size//(4//dtype.itemsize) if dtype != dtypes.half else size
+      # temp0: array<{Device[Device.DEFAULT].renderer.buf_map(dtype)},{corrected_size}>;
+      self.assertIn(f",{corrected_size}>;", out)
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared memory")
   @unittest.skip("tinygrad doesn't support this behavior")

test/external/external_test_gpu_crash.py

@@ -23,7 +23,7 @@ class TestGPUCrash(unittest.TestCase):
     cls.is_cdna = cls.dev.target[0] < 10
     ins = importlib.import_module('tinygrad.runtime.autogen.amd.' + ('cdna' if cls.is_cdna else 'rdna3') + '.ins')
     for rdna3_name, cdna3_name in RDNA3_CDNA3_MAP.items():
-      setattr(cls, rdna3_name, staticmethod(getattr(ins, cdna3_name if cls.is_cdna else rdna3_name)))
+      setattr(cls, rdna3_name, getattr(ins, cdna3_name if cls.is_cdna else rdna3_name))
 
   def setUp(self):
     # Verify device works before each test

test/mockgpu/amd/pcode.py

@@ -1,20 +1,5 @@
 # Tokenizer-based expression parser for AMD pcode
-import ast, itertools, operator, re
 from typing import Any, Callable
-
-_BINOPS = {ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, ast.FloorDiv: operator.floordiv,
-           ast.Mod: operator.mod, ast.LShift: operator.lshift, ast.RShift: operator.rshift,
-           ast.BitAnd: operator.and_, ast.BitOr: operator.or_, ast.BitXor: operator.xor}
-def _const_int(expr: str) -> int:
-  """Evaluate a compile-time integer expression (integer literals and basic arithmetic only)."""
-  def ev(node: ast.AST) -> int:
-    if isinstance(node, ast.Expression): return ev(node.body)
-    if isinstance(node, ast.Constant) and isinstance(node.value, int): return node.value
-    if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.USub, ast.UAdd)):
-      return (-1 if isinstance(node.op, ast.USub) else 1) * ev(node.operand)
-    if isinstance(node, ast.BinOp) and type(node.op) in _BINOPS: return _BINOPS[type(node.op)](ev(node.left), ev(node.right))
-    raise ValueError(f"not a constant integer expression: {expr!r}")
-  return ev(ast.parse(expr.strip(), mode='eval'))
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.codegen.decomp.dtype import f2f
@@ -23,7 +8,6 @@ from tinygrad.codegen.decomp.dtype import f2f
 VarVal = UOp | tuple[str, list[str], str]
 
 def _const(dt, v): return UOp.const(v, dt)
-def _single_value(v: UOp): return v.vmin if v.vmin == v.vmax else None
 def _u32(v): return _const(dtypes.uint32, v)
 def _u64(v): return _const(dtypes.uint64, v)
 def _to_u32(v): return v if v.dtype == dtypes.uint32 else v.bitcast(dtypes.uint32) if v.dtype.itemsize == 4 else v.cast(dtypes.uint32)
@@ -71,8 +55,8 @@ def _expr_bits(v: UOp) -> int:
   if v.op in (Ops.AND, Ops.XOR):
     widths: list[int] = []
     for src in v.src:
-      if isinstance(sv:=_single_value(src), int) and sv > 0 and (sv & (sv + 1)) == 0:
-        widths.append(sv.bit_length())
+      if src.op == Ops.CONST and isinstance(src.val, int) and src.val > 0 and (src.val & (src.val + 1)) == 0:
+        widths.append(src.val.bit_length())
     if widths: return max(widths)
   return v.dtype.bitsize
 
@@ -160,9 +144,9 @@ def _minmax_reduce(is_max: bool, dt, *args: UOp) -> UOp:
 def _find_two_pi_mul(x):
   if x.op != Ops.MUL or len(x.src) != 2: return None
   for i, s in enumerate(x.src):
-    if (sv:=_single_value(s)) is not None and abs(sv - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
+    if s.op == Ops.CONST and abs(s.val - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
     if s.op == Ops.MUL and len(s.src) == 2:
-      vals = [sv for ss in s.src if (sv:=_single_value(ss)) is not None]
+      vals = [ss.val for ss in s.src if ss.op == Ops.CONST] + [ss.src[0].val for ss in s.src if ss.op == Ops.CAST and ss.src[0].op == Ops.CONST]
       if len(vals) == 2 and abs(vals[0] * vals[1] - 6.283185307179586) < 1e-5: return (x.src[1-i], vals[0] * vals[1])
   return None
 
@@ -179,7 +163,7 @@ def _trig_reduce(x, phase=0.0):
 
 def _signext(val: UOp) -> UOp:
   for bits, mask, ext in [(4, 0xF, 0xFFFFFFF0), (8, 0xFF, 0xFFFFFF00), (16, 0xFFFF, 0xFFFF0000)]:
-    if (val.op == Ops.AND and len(val.src) == 2 and _single_value(val.src[1]) == mask) or val.dtype.itemsize == bits // 8:
+    if (val.op == Ops.AND and len(val.src) == 2 and val.src[1].op == Ops.CONST and val.src[1].val == mask) or val.dtype.itemsize == bits // 8:
       v32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
       sb = (v32 >> _u32(bits - 1)) & _u32(1)
       return sb.ne(_u32(0)).where(v32 | _u32(ext), v32).cast(dtypes.int)
@@ -201,20 +185,7 @@ def _abs(val: UOp) -> UOp:
 def _f_to_u(f, dt):
   clamped = (f < _const(f.dtype, 0.0)).where(_const(f.dtype, 0.0), f)
   truncated = UOp(Ops.TRUNC, src=(clamped,))
-  res = (truncated >= _const(f.dtype, 2**(dt.itemsize*8))).where(_const(dt, dt.max), truncated.cast(dt))
-  return _isnan(f).where(_const(dt, 0), res)  # float->uint conversion of NaN is 0 on hardware
-
-def _f_to_i32(a: UOp) -> UOp:
-  """v_cvt_i32_f32/f64: truncate toward zero, saturate to [INT_MIN, INT_MAX], NaN -> 0.
-  (x86 cvttss2si returns 0x80000000 for all of these, which matches hardware only for negative overflow.)"""
-  res = (a >= _const(a.dtype, 2147483648.0)).where(_const(dtypes.int, 0x7FFFFFFF), UOp(Ops.TRUNC, src=(a,)).cast(dtypes.int))
-  return _isnan(a).where(_const(dtypes.int, 0), res)
-
-def _ftz_f32(v: UOp) -> UOp:
-  """Flush f32 denormals to signed zero (RDNA default float mode flushes denormal f32 inputs on select-style ops)."""
-  bits = v.bitcast(dtypes.uint32) if v.dtype == dtypes.float32 else v
-  return ((bits & _u32(0x7FFFFFFF)) < _u32(0x00800000)).where((bits & _u32(0x80000000)).bitcast(dtypes.float32),
-                                                              v if v.dtype == dtypes.float32 else v.bitcast(dtypes.float32))
+  return (truncated >= _const(f.dtype, 2**(dt.itemsize*8))).where(_const(dt, dt.max), truncated.cast(dt))
 
 def _cvt_quiet(val: UOp) -> UOp:
   bits, _, _, qb, _ = _float_info(val)
@@ -259,51 +230,18 @@ def _ldexp(val: UOp, exp: UOp) -> UOp:
   if val.dtype == dtypes.uint32: val = val.bitcast(dtypes.float32)
   elif val.dtype == dtypes.uint64: val = val.bitcast(dtypes.float64)
   if exp.dtype in (dtypes.uint32, dtypes.uint64): exp = exp.cast(dtypes.int if exp.dtype == dtypes.uint32 else dtypes.int64)
-  bits = val.bitcast(dtypes.uint32) if val.dtype == dtypes.float32 else val.bitcast(dtypes.uint64)
-  abs_max = _const(bits.dtype, 0x7F800000 if val.dtype == dtypes.float32 else 0x7FF0000000000000)
-  sign_mask = _const(bits.dtype, 0x80000000 if val.dtype == dtypes.float32 else 0x8000000000000000)
-  # hardware flushes denormal inputs to signed zero
-  magn_mask = _const(bits.dtype, 0x7FFFFFFF if val.dtype == dtypes.float32 else 0x7FFFFFFFFFFFFFFF)
-  is_denorm = ((bits & abs_max).eq(_const(bits.dtype, 0))) & ((bits & magn_mask).ne(_const(bits.dtype, 0)))
-  val = is_denorm.where((bits & sign_mask).bitcast(val.dtype), val)
-  # hardware propagates 0/+-inf/NaN unchanged (avoids 0*inf = NaN on the host)
-  res = val * UOp(Ops.EXP2, src=(exp.cast(val.dtype),))
-  is_special = (bits & abs_max).eq(_const(bits.dtype, 0)) | ((bits & abs_max) >= abs_max)
-  return is_special.where(val, res)
+  return val * UOp(Ops.EXP2, src=(exp.cast(val.dtype),))
 
 def _frexp_mant(val: UOp) -> UOp:
   val = val.bitcast(dtypes.float32) if val.dtype == dtypes.uint32 else val.bitcast(dtypes.float64) if val.dtype == dtypes.uint64 else val
-  if val.dtype == dtypes.float32:
-    bits = val.bitcast(dtypes.uint32)
-    # denormal/zero inputs (exponent field == 0) return signed zero on hardware
-    return ((bits & _u32(0x7F800000)).ne(_u32(0))).where(((bits & _u32(0x807FFFFF)) | _u32(0x3F000000)).bitcast(dtypes.float32),
-                                                          (bits & _u32(0x80000000)).bitcast(dtypes.float32))
-  bits = val.bitcast(dtypes.uint64)
-  return ((bits & _const(dtypes.uint64, 0x7FF0000000000000)).ne(_const(dtypes.uint64, 0))).where(
-    ((bits & _const(dtypes.uint64, 0x800FFFFFFFFFFFFF)) | _const(dtypes.uint64, 0x3fe0000000000000)).bitcast(dtypes.float64),
-    (bits & _const(dtypes.uint64, 0x8000000000000000)).bitcast(dtypes.float64))
-
-def _msb(val: UOp, bits: int) -> UOp:
-  """Index of the highest set bit, or -1 if val == 0."""
-  dt = dtypes.uint64 if bits > 32 else dtypes.uint32
-  val = val.cast(dt) if val.dtype != dt else val
-  result = _const(dtypes.int, -1)
-  for i in range(bits - 1, -1, -1):
-    cond = ((val >> _const(dt, i)) & _const(dt, 1)).ne(_const(dt, 0)) & result.eq(_const(dtypes.int, -1))
-    result = cond.where(_const(dtypes.int, i), result)
-  return result
+  if val.dtype == dtypes.float32: return ((val.bitcast(dtypes.uint32) & _u32(0x807FFFFF)) | _u32(0x3f000000)).bitcast(dtypes.float32)
+  return ((val.bitcast(dtypes.uint64) & _const(dtypes.uint64, 0x800FFFFFFFFFFFFF)) |
+    _const(dtypes.uint64, 0x3fe0000000000000)).bitcast(dtypes.float64)
 
 def _frexp_exp(val: UOp) -> UOp:
   val = val.bitcast(dtypes.float32) if val.dtype == dtypes.uint32 else val.bitcast(dtypes.float64) if val.dtype == dtypes.uint64 else val
-  if val.dtype == dtypes.float32:
-    e = (val.bitcast(dtypes.uint32) >> _u32(23)) & _u32(0xFF)
-    return e.ne(_u32(0)).where(e.cast(dtypes.int) - _const(dtypes.int, 126), _const(dtypes.int, 0))  # f32 denormals -> 0 (hardware verified)
-  bits = val.bitcast(dtypes.uint64)
-  e = (bits >> _const(dtypes.uint64, 52)) & _const(dtypes.uint64, 0x7FF)
-  mant = bits & _const(dtypes.uint64, 0xFFFFFFFFFFFFF)
-  # f64 denormals: normalized exponent = highest set mantissa bit - 1073, zero -> 0 (hardware verified)
-  denorm = mant.ne(_const(dtypes.uint64, 0)).where(_msb(mant, 52) - _const(dtypes.int, 1073), _const(dtypes.int, 0))
-  return e.ne(_const(dtypes.uint64, 0)).where(e.cast(dtypes.int) - _const(dtypes.int, 1022), denorm)
+  if val.dtype == dtypes.float32: return ((val.bitcast(dtypes.uint32) >> _u32(23)) & _u32(0xFF)).cast(dtypes.int) - _const(dtypes.int, 126)
+  return ((val.bitcast(dtypes.uint64) >> _const(dtypes.uint64, 52)) & _const(dtypes.uint64, 0x7FF)).cast(dtypes.int) - _const(dtypes.int, 1022)
 
 TWO_OVER_PI = int(
   "0145f306dc9c882a53f84eafa3ea69bb81b6c52b3278872083fca2c757bd778ac36e48dc74849ba5c00c925dd413a32439fc3bd"
@@ -361,9 +299,9 @@ _FUNCS: dict[str, Callable[..., UOp]] = {
   'fma': lambda a, b, c: a * b + c,
   'i32_to_f32': lambda a: a.cast(dtypes.int).cast(dtypes.float32),
   'u32_to_f32': lambda a: a.cast(dtypes.uint32).cast(dtypes.float32),
-  'f32_to_i32': lambda a: _f_to_i32(a.bitcast(dtypes.float32)),
+  'f32_to_i32': lambda a: UOp(Ops.TRUNC, src=(a.bitcast(dtypes.float32),)).cast(dtypes.int),
   'f32_to_u32': lambda a: _f_to_u(a.bitcast(dtypes.float32), dtypes.uint32),
-  'f64_to_i32': lambda a: _f_to_i32(a.bitcast(dtypes.float64)),
+  'f64_to_i32': lambda a: UOp(Ops.TRUNC, src=(a.bitcast(dtypes.float64),)).cast(dtypes.int),
   'f64_to_u32': lambda a: _f_to_u(a.bitcast(dtypes.float64), dtypes.uint32),
   'f16_to_f32': lambda a: _f16_extract(a).cast(dtypes.float32),
   'f32_to_f16': lambda a: a.cast(dtypes.half),
@@ -422,13 +360,22 @@ _FUNCS: dict[str, Callable[..., UOp]] = {
   'fp8_to_f32': _fp8_to_f32, 'bf8_to_f32': _bf8_to_f32, 'f32_to_fp8': _f32_to_fp8, 'f32_to_bf8': _f32_to_bf8,
   'f32_to_bf16': _f32_to_bf16, 'f32_to_bf16_SR': _f32_to_bf16_sr, 'f32_to_bf16_sr': _f32_to_bf16_sr,
 }
-# min/max family: min/max + 3-input (x3), IEEE num variants (f16/f32 only), and long names minimum/maximum (f16/f32 only)
-for is_max, name, full in [(False, 'min', 'minimum'), (True, 'max', 'maximum')]:
-  for dt, sfx, pre in [(dtypes.float32, 'f32', None), (dtypes.int, 'i32', None), (dtypes.uint32, 'u32', None),
-                       (dtypes.int16, 'i16', None), (dtypes.uint16, 'u16', None), (dtypes.half, 'f16', _f16_extract)]:
-    def mm(*a, im=is_max, d=dt, p=pre): return _minmax_reduce(im, d, *(a if p is None else [p(x) for x in a]))
-    extra = (f'v_{name}_num_{sfx}', f'v_{name}3_num_{sfx}', f'v_{full}_{sfx}', f'v_{full}3_{sfx}') if dt in (dtypes.float32, dtypes.half) else ()
-    for fn in (f'v_{name}_{sfx}', f'v_{name}3_{sfx}', *extra): _FUNCS[fn] = mm
+for is_max, name in [(False, 'min'), (True, 'max')]:
+  for dt, sfx in [(dtypes.float32, 'f32'), (dtypes.int, 'i32'), (dtypes.uint32, 'u32'), (dtypes.int16, 'i16'), (dtypes.uint16, 'u16')]:
+    _FUNCS[f'v_{name}_{sfx}'] = lambda *a, im=is_max, d=dt: _minmax_reduce(im, d, *a)
+    _FUNCS[f'v_{name}3_{sfx}'] = lambda *a, im=is_max, d=dt: _minmax_reduce(im, d, *a)
+# f16 min/max/min3/max3/med3
+for is_max, name in [(False, 'min'), (True, 'max')]:
+  _FUNCS[f'v_{name}_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}3_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}_num_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}_num_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
+  _FUNCS[f'v_{name}3_num_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}3_num_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
+  _FUNCS[f'v_{name}imum_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}imum_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
+  _FUNCS[f'v_{name}imum3_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}imum3_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # TOKENIZER/PARSER
@@ -550,7 +497,7 @@ class Parser:
         if not dtypes.is_int(right.dtype): right = right.cast(dtypes.uint32)
         return (left >> right) if op == '>>' else (left << right)
       case '+' | '-':
-        if op == '-' and (lv:=_single_value(left)) is not None and (rv:=_single_value(right)) is not None: return _const(left.dtype, lv - rv)
+        if op == '-' and left.op == Ops.CONST and right.op == Ops.CONST: return _const(left.dtype, left.val - right.val)
         return (left + right) if op == '+' else (left - right)
       case '*' | '/':
         # Integer promotion: promote 16-bit integers to 32-bit before multiply to avoid overflow
@@ -560,7 +507,7 @@ class Parser:
           left, right = left.cast(pdt), right.cast(pdt)
         if op == '*': return left * right
         return (left // right) if dtypes.is_int(left.dtype) else (left / right)
-      case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if _single_value(left) == 2.0 else left
+      case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if left.op == Ops.CONST and left.val == 2.0 else left
 
   _PREC = [('||',), ('&&',), ('|',), ('^',), ('&',), ('==', '!=', '<>'), ('>=', '<=', '>', '<'), ('>>', '<<'), ('+', '-'), ('*', '/'), ('**',)]
 
@@ -582,8 +529,8 @@ class Parser:
       return inner.eq(_const(inner.dtype, 0))
     if self.try_eat_val('-', 'OP'):
       inner = self.unary()
-      if (v:=_single_value(inner)) is not None:
-        return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -v)
+      if inner.op == Ops.CONST:
+        return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -inner.val)
       return inner.neg()
     if self.try_eat_val('+', 'OP'): return self.unary()
     return self.postfix()
@@ -722,13 +669,15 @@ class Parser:
       self.eat('OP')
       width = self.parse()
       self.eat('RBRACKET')
-      if isinstance(w:=_single_value(width), int):
+      if width.op == Ops.CONST:
+        w = int(width.val)
         return (base >> _to_u32(first)) & _const(base.dtype, (1 << w) - 1)
       return base
     if self.try_eat('COLON'):
       second = self.parse()
       self.eat('RBRACKET')
-      if isinstance(a:=_single_value(first), int) and isinstance(b:=_single_value(second), int):
+      if first.op == Ops.CONST and second.op == Ops.CONST:
+        a, b = int(first.val), int(second.val)
         if a < b: return _bitreverse(base, b - a + 1)
         hi, lo = a, b
         if lo >= base.dtype.itemsize * 8:
@@ -749,7 +698,8 @@ class Parser:
       dt_suffix = DTYPES.get(self.eat('IDENT').val, dtypes.uint32)
     if var_name is None:
       var_name = self._find_var_name(base)
-    if isinstance(idx:=_single_value(first), int):
+    if first.op == Ops.CONST:
+      idx = int(first.val)
       # Check for array element (var@idx)
       if var_name and f'{var_name}@{idx}' in self.vars:
         v = self.vars[f'{var_name}@{idx}']
@@ -922,7 +872,7 @@ class Parser:
 
   def _coerce_cmp(self, l: UOp, r: UOp) -> tuple[UOp, UOp]:
     if l.dtype != r.dtype:
-      if r.dtype == dtypes.int and isinstance(rv:=_single_value(r), int) and rv < 0: l = l.cast(dtypes.int)
+      if r.dtype == dtypes.int and r.op == Ops.CONST and r.val < 0: l = l.cast(dtypes.int)
       else: r = r.cast(l.dtype)
     return l, r
 
@@ -940,8 +890,6 @@ class Parser:
       return result & _isnan(l).logical_not() & _isnan(r).logical_not()
     return result
 
-_break_var_ids = itertools.count()  # unique names for per-loop break-tracking variables
-
 def _match_bracket(toks: list[Token], start: int) -> tuple[int, list[Token]]:
   """Match brackets from start, return (end_idx, inner_tokens)."""
   j, depth = start + 1, 1
@@ -1020,7 +968,9 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
           p.eat('NUM')
           p.eat('QUOTE')
         if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
-        return int(p.parse())
+        expr = p.parse().simplify()
+        assert expr.op == Ops.CONST, f"loop bound must be constant, got {expr}"
+        return int(expr.val)
       start_val = parse_bound()
       p.eat('COLON')
       end_val = parse_bound()
@@ -1037,7 +987,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
         i += 1
       # Execute loop with break support
       has_break = any('break' in bl.lower() for bl in body_lines)
-      found_var = f'_found_{next(_break_var_ids)}' if has_break else None
+      found_var = f'_found_{id(body_lines)}' if has_break else None
       if found_var: env[found_var] = block_assigns[found_var] = _const(dtypes.bool, False)
       for loop_i in range(start_val, end_val + 1):
         subst_lines = [_subst_loop_var(bl, loop_var, loop_i) for bl in body_lines if not (has_break and bl.strip().lower() == 'break')]
@@ -1137,7 +1087,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
           j, slice_toks = _match_bracket(toks, j)
           slice_str = _tok_str(slice_toks)
           hi_str, lo_str = slice_str.split(':')
-          hi_val, lo_val = _const_int(hi_str), _const_int(lo_str)
+          hi_val, lo_val = int(eval(hi_str.strip())), int(eval(lo_str.strip()))
           if j < len(toks) and toks[j].type == 'DOT': j += 2  # skip .type suffix
           if j < len(toks) and toks[j].type == 'EQUALS': j += 1
           ln = parse_tokens(lane_toks, env, funcs)
@@ -1195,7 +1145,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
         hi_str = ' '.join(t.val for t in toks[bracket_start:colon_pos] if t.type != 'EOF')
         lo_str = ' '.join(t.val for t in toks[colon_pos+1:j] if t.type != 'EOF')
         try:
-          hi_val, lo_val = _const_int(hi_str), _const_int(lo_str)
+          hi_val, lo_val = int(eval(hi_str)), int(eval(lo_str))
           hi, lo = max(hi_val, lo_val), min(hi_val, lo_val)
           j += 1
           if j < len(toks) and toks[j].type == 'DOT': j += 2
@@ -1209,7 +1159,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
           block_assigns[var] = env[var] = _set_bits(old, _val_to_bits(val), hi - lo + 1, lo)
           i += 1
           continue
-        except (ValueError, SyntaxError): pass  # non-constant slice bounds - fall through to other statement forms
+        except Exception: pass
       elif toks[1].type == 'LBRACKET':  # bit index: var[expr] (only for var[...], not var.type[...])
         existing = block_assigns.get(var, env.get(var))
         if existing is not None and isinstance(existing, UOp) and \
@@ -1410,25 +1360,3 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
 
 def parse_expr(expr: str, env: dict[str, VarVal], funcs: dict | None = None) -> UOp:
   return parse_tokens(tokenize(expr.strip().rstrip(';')), env, funcs)
-
-def parse_pcode(pcode: str, srcs: dict[str, UOp | int] | None = None) -> tuple[dict, list]:
-  env: dict = srcs.copy() if srcs else {}
-  assigns: list[tuple[str, UOp]] = []
-  raw_lines = [l.strip().rstrip(';') for l in pcode.split('\n') if l.strip() and not l.strip().startswith('//')]
-  # TODO: pcode.py should tokenize full pcode string instead of line-by-line, then this hack can be removed
-  lines: list[str] = []
-  for l in raw_lines:
-    if lines and re.search(r'(&&|\|\||[&|+\-*/^])\s*$', lines[-1]): lines[-1] = lines[-1] + ' ' + l
-    else: lines.append(l)
-  _, final, _ = parse_block(lines, 0, env, assigns=assigns)
-  sliced = set(d.split('[')[0] for d, _ in assigns if '[' in d)
-  for var, val in final.items():
-    if var in ['D0', 'S0', 'SCC', 'VCC', 'EXEC', 'PC', 'RETURN_DATA', 'VDATA'] and isinstance(val, UOp):
-      if var in sliced and not any(re.match(rf'{var}\.\w+\s*=', l) for l in lines): continue
-      for l in lines:
-        if (m := re.match(rf'{var}\.(\w+(?:\[\w+\])?)', l)):
-          assigns.append((f'{var}.{m.group(1)}', val))
-          break
-      else: assigns.append((var, val))
-  return env, assigns
-

test/mockgpu/amd/sqtt_enc.py

@@ -1,100 +0,0 @@
-# SQTT trace encoder for the emulator (the decoder lives in tinygrad/renderer/amd/sqtt.py).
-# run_asm emits packets inline as instructions execute; finished traces end up in emu.sqtt_traces.
-from __future__ import annotations
-from tinygrad.renderer.amd.dsl import Inst
-from tinygrad.renderer.amd.sqtt import (_build_decode_tables, PACKET_TYPES_RDNA3, PacketType, InstOp,
-                                        LAYOUT_HEADER, WAVESTART, WAVEEND, INST, IMMEDIATE, VALUINST)
-
-_NIB_COUNTS = {cls: nc for _, (cls, nc, *_) in _build_decode_tables(PACKET_TYPES_RDNA3)[0].items()}
-
-def _emit_nibbles(nibbles: list[int], pkt_cls: type[PacketType], **kwargs):
-  raw = pkt_cls.encoding.default
-  for k, v in kwargs.items(): raw = pkt_cls.__dict__[k].set(raw, v)
-  nibbles.extend((raw >> (i * 4)) & 0xF for i in range(_NIB_COUNTS[pkt_cls]))
-
-def make_encoder():
-  """Build an SQTT trace encoder for the emulator. Returns (emit, finish, finalize)."""
-  from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp as SOPPOp3
-  from tinygrad.runtime.autogen.amd.rdna4.enum import SOPPOp as SOPPOp4
-  from tinygrad.runtime.autogen.amd.rdna3 import ins as ir3
-  from tinygrad.runtime.autogen.amd.rdna4 import ins as ir4
-  from tinygrad.runtime.autogen.amd.cdna import ins as irc
-  import re
-
-  def _kinds(*names: str) -> tuple[type[Inst], ...]:
-    return tuple(getattr(m, n) for m in (ir3, ir4, irc) for n in names if hasattr(m, n))
-  _SOPP, _SMEM, _DS = _kinds('SOPP'), _kinds('SMEM'), _kinds('DS')
-  _GLOBAL, _FLAT, _SCRATCH = _kinds('GLOBAL', 'VGLOBAL'), _kinds('FLAT', 'VFLAT'), _kinds('SCRATCH', 'VSCRATCH')
-  _VALU = _kinds('VOP1', 'VOP2', 'VOP3', 'VOP3P', 'VOP3PX2', 'VOPC', 'VOPD', 'VOP3SD', 'VOP3_SDST', 'VOP1_SDST')
-
-  # SOPP classification sets
-  _SOPP_SKIP = {SOPPOp3.S_ENDPGM.value, SOPPOp3.S_ENDPGM_SAVED.value, SOPPOp3.S_ENDPGM_ORDERED_PS_DONE.value, SOPPOp3.S_DELAY_ALU.value}
-  _SOPP_IMMEDIATE = {SOPPOp3.S_NOP.value, SOPPOp3.S_CLAUSE.value, SOPPOp3.S_WAITCNT.value, SOPPOp3.S_WAITCNT_DEPCTR.value,
-                     SOPPOp3.S_WAIT_IDLE.value, SOPPOp3.S_WAIT_EVENT.value, SOPPOp3.S_SLEEP.value, SOPPOp3.S_SET_INST_PREFETCH_DISTANCE.value}
-  for _op in (SOPPOp4.S_WAIT_ALU, SOPPOp4.S_WAIT_LOADCNT, SOPPOp4.S_WAIT_STORECNT, SOPPOp4.S_WAIT_SAMPLECNT,
-              SOPPOp4.S_WAIT_BVHCNT, SOPPOp4.S_WAIT_EXPCNT, SOPPOp4.S_WAIT_DSCNT, SOPPOp4.S_WAIT_KMCNT,
-              SOPPOp4.S_WAIT_LOADCNT_DSCNT, SOPPOp4.S_WAIT_STORECNT_DSCNT):
-    _SOPP_IMMEDIATE.add(_op.value)
-  _SOPP_BARRIER = {SOPPOp3.S_BARRIER.value}
-  if hasattr(SOPPOp4, 'S_BARRIER_WAIT'): _SOPP_BARRIER.add(SOPPOp4.S_BARRIER_WAIT.value)
-  if hasattr(SOPPOp4, 'S_BARRIER_LEAVE'): _SOPP_BARRIER.add(SOPPOp4.S_BARRIER_LEAVE.value)
-  _SOPP_BRANCH = {SOPPOp3.S_BRANCH.value, SOPPOp3.S_CBRANCH_SCC0.value, SOPPOp3.S_CBRANCH_SCC1.value,
-                  SOPPOp3.S_CBRANCH_VCCZ.value, SOPPOp3.S_CBRANCH_VCCNZ.value,
-                  SOPPOp3.S_CBRANCH_EXECZ.value, SOPPOp3.S_CBRANCH_EXECNZ.value}
-
-  # VALU sub-classification patterns
-  _VALUT_4_RE = re.compile(r'V_(EXP|LOG|RCP|RSQ|SQRT|SIN|COS|CEIL|FLOOR|TRUNC|RNDNE|FRACT|FREXP)_')
-  _VALUB_2_RE = re.compile(r'V_(LSHLREV|LSHRREV|ASHRREV)_(B|I)64')
-  _VALUB_4_RE = re.compile(r'V_MAD_(U|I)64')
-  _VALUB_16_RE = re.compile(r'V_\w+_F64')
-
-  def _valu_op(op_name: str) -> InstOp|None:
-    if 'CMPX' in op_name: return InstOp.VALU1_WR_EXEC
-    if _VALUB_2_RE.search(op_name): return InstOp.VALUB_2
-    if _VALUB_4_RE.search(op_name): return InstOp.VALUB_4
-    if _VALUB_16_RE.search(op_name): return InstOp.VALUB_16
-    if _VALUT_4_RE.search(op_name): return InstOp.VALUT_4
-    return None
-
-  def _mem_op(t: type[Inst], op_name: str) -> InstOp:
-    is_store = "STORE" in op_name
-    if issubclass(t, _DS): return InstOp.LDS_WR_2 if is_store else InstOp.LDS_RD
-    if issubclass(t, _GLOBAL): return InstOp.SGMEM_WR_2 if is_store else InstOp.SGMEM_RD_1
-    if issubclass(t, _FLAT) or issubclass(t, _SCRATCH): return InstOp.FLAT_WR_3 if is_store else InstOp.FLAT_RD_2
-    return InstOp.SALU
-
-  nibbles: list[int] = []
-  started: set[int] = set()
-  _emit_nibbles(nibbles, LAYOUT_HEADER, layout=3, sel_a=6)
-
-  def emit(wave_id: int, inst: Inst, branch_taken: bool|None):
-    """Emit an SQTT packet for one executed instruction."""
-    w = wave_id & 0x1F
-    if wave_id not in started:
-      _emit_nibbles(nibbles, WAVESTART, delta=1, simd=0, wgp=0, wave=w, id7=wave_id)
-      started.add(wave_id)
-    inst_type, inst_op, op_name = type(inst), inst.op.value if hasattr(inst, 'op') else 0, inst.op.name if hasattr(inst, 'op') else ""
-    if issubclass(inst_type, _SOPP):
-      if inst_op in _SOPP_SKIP: return
-      if inst_op in _SOPP_IMMEDIATE: _emit_nibbles(nibbles, IMMEDIATE, delta=1, wave=w)
-      elif inst_op in _SOPP_BARRIER: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.BARRIER)
-      elif inst_op in _SOPP_BRANCH: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.JUMP if branch_taken else InstOp.JUMP_NO)
-      else: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.SALU)
-    elif issubclass(inst_type, _VALU):
-      if (op := _valu_op(op_name)) is None: _emit_nibbles(nibbles, VALUINST, delta=1, wave=w)
-      else: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=op)
-    elif issubclass(inst_type, _SMEM): _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.SMEM_RD)
-    else: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=_mem_op(inst_type, op_name))
-
-  def finish(wave_id: int):
-    """Emit WAVEEND for a completed wave."""
-    if wave_id in started: _emit_nibbles(nibbles, WAVEEND, delta=1, simd=0, wgp=0, wave=wave_id & 0x1F)
-
-  def finalize() -> bytes:
-    """Pad and return the encoded SQTT blob."""
-    while len(nibbles) % 2 != 0: nibbles.append(0)
-    nibbles.extend([0] * 32)
-    while len(nibbles) % 64 != 0: nibbles.append(0)
-    return bytes(nibbles[i] | ((nibbles[i + 1] if i + 1 < len(nibbles) else 0) << 4) for i in range(0, len(nibbles), 2))
-
-  return emit, finish, finalize

test/mockgpu/usb.py

@@ -160,7 +160,7 @@ class MockUSB3:
     elif request == 0xE5:
       self.state._xram_write_byte(value, index)
     elif request == 0xF2:
-      op = ("sram_read" if value & 0x8000 else "sram_write", 0xF000 + (index & 0xFF) * 0x4000, (value & 0x7FFF) * 512)
+      op = ("sram_read" if value & 0x8000 else "sram_write", 0xF000, (value & 0x7FFF) * 512)
       if value & 0x8000: self._bulk_read_op = op
       else: self._bulk_write_op = op
     elif request == 0xF0:
@@ -193,33 +193,19 @@ class MockUSB3:
     op, address, size = self._bulk_write_op
     assert len(data) == size
     if op == "sram_write":
-      ctrl, (host_addr, region_size) = next((ca, r) for ca, r in self.state._dma_regions.items() if ca <= address < ca + r[1])
-      ctypes.memmove(host_addr + (address - ctrl), data, min(len(data), region_size - (address - ctrl)))
-      self.state.driver._emulate_execute()  # landed data may un-stall a ring polling on it (e.g. copyin sentinels)
+      host_addr, region_size = self.state._dma_regions[address]
+      ctypes.memmove(host_addr, data, min(len(data), region_size))
     elif op == "pcie_write": self.state._pcie_write(address, data)
     else: raise RuntimeError(f"cannot bulk write for {op}")
     self._bulk_write_op = None
 
-  def bulk_write_async(self, payload:memoryview, timeout:int=10000) -> int:  # the mock completes transfers synchronously
-    self.bulk_write(bytes(payload), timeout)
-    return 0
-
-  def control_write_async(self, request:int, value:int=0, index:int=0, data:bytes=b"", timeout:int=1000) -> int:
-    self.control_write(request, value, index, data, timeout)
-    return 0
-
-  def control_read_async(self, request:int, length:int, value:int=0, index:int=0, timeout:int=1000) -> tuple[int, memoryview]:
-    return 0, self.control_read(request, length, value, index, timeout)
-
-  def bulk_wait(self, tag:int): pass
-
   def bulk_read(self, length:int, timeout:int=1000) -> memoryview:
     assert self._bulk_read_op is not None
     op, address, size = self._bulk_read_op
     assert length == size
     if op == "sram_read":
-      ctrl, (host_addr, region_size) = next((ca, r) for ca, r in self.state._dma_regions.items() if ca <= address < ca + r[1])
-      data = bytes((ctypes.c_ubyte * min(length, region_size - (address - ctrl))).from_address(host_addr + (address - ctrl)))
+      host_addr, region_size = self.state._dma_regions[address]
+      data = bytes((ctypes.c_ubyte * min(length, region_size)).from_address(host_addr))
     elif op == "pcie_read": data = self.state._pcie_read(address, length)
     else: raise RuntimeError(f"cannot bulk read for {op}")
     self._bulk_read_op = None

test/null/test_const_folding.py

@@ -3,6 +3,7 @@ from tinygrad import dtypes, Context
 from tinygrad.dtype import DType, ConstType
 from tinygrad.uop.ops import Ops, UOp
 from test.helpers import full_rewrite
+import numpy as np
 
 class TestWeakConstFolding(unittest.TestCase):
   def test_weakint_math(self):
@@ -26,14 +27,16 @@ class TestBitcastConstFolding(unittest.TestCase):
     for val, src_dt, dst_dt, bits in ((3000000000, dtypes.int32, dtypes.uint32, 3000000000),
                                       (70000, dtypes.int16, dtypes.uint16, 4464),
                                       (-5, dtypes.uint32, dtypes.int32, -5)):
-      self.assertIs(UOp.const(val, src_dt).bitcast(dst_dt).simplify(), UOp.const(bits, dst_dt))
+      self.assertEqual(UOp.const(val, src_dt).bitcast(dst_dt).simplify().val, bits)
 
   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 = UOp.const(from_v, from_dt).bitcast(to_dt).simplify()
-          self.assertIs(r, UOp.const(to_v, to_dt), f"{from_dt} -> {to_dt} ({from_v} -> {to_v})")
+          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.val, to_v, msg)
 
     t({dtypes.int8: 0, dtypes.uint8: 0, dtypes.bool: False})
     t({dtypes.int8: 1, dtypes.uint8: 1, dtypes.bool: True})

test/null/test_dtype_spec.py

@@ -419,7 +419,7 @@ class TestAutoCastType(unittest.TestCase):
     self.check_where_alternate_input_other(3, True, dtypes.weakint)
 
   def test_where_non_bool_cond_raises(self):
-    with self.assertRaises(RuntimeError): Tensor([1, 0, 2]).where(1, 0).dtype
+    with self.assertRaises(RuntimeError): Tensor([1, 0, 2]).where(1, 0)
     self.check_where_alternate_input_other(False, True, dtypes.bool)
 
   @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))

test/null/test_simplify_valid_idx.py

@@ -3,7 +3,7 @@ import unittest, itertools
 from tinygrad.codegen.late.coalesce import indexing_simplify
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite
-from tinygrad.uop.weak import pm_commit_weak
+from tinygrad.uop.weak import pm_lower_index_dtype
 from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load
 from tinygrad.helpers import Context
 from test.helpers import full_rewrite
@@ -496,7 +496,7 @@ class TestImageSimplification(unittest.TestCase):
     idx_y = (f + UOp.const(1.0)).cast(dtypes.int)
     load = get_load_image_uop((10, 10, 4), (UOp.const(-1) < idx_y) & (idx_y < UOp.const(10)),
                               (Special("gidx0", 10), idx_y))
-    off = graph_rewrite(load.sink(), pm_commit_weak+indexing_simplify).src[0].src[0]
+    off = graph_rewrite(load.sink(), pm_lower_index_dtype+indexing_simplify, ctx={}).src[0].src[0]
     self.assertEqual(off.src[1].get_valid(), UOp.const(True))
 
 class TestDropTrueGate(unittest.TestCase):

test/null/test_uop_graph.py

@@ -230,11 +230,11 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_depth_2_const_fold(self):
     v = UOp.variable("tmp", 0, 1, dtypes.int, param=True)
-    c2 = UOp.const(2)
-    c4 = UOp.const(4)
+    c2 = UOp.const(2, dtypes.int)
+    c4 = UOp.const(4, dtypes.int)
     vc = v+c2
     out = vc+c4
-    self.assertIs(out.simplify(), (v+UOp.const(6)).simplify())
+    self.assertIs(out.simplify(), (v+UOp.const(6, dtypes.int)).simplify())
 
   def test_bitcast_to_same_dtype_fold(self):
     for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
@@ -245,7 +245,7 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_sub_with_cast_folds(self):
     a = Variable("a", 0, 5)
-    out = a+(-a)
+    out = a.cast(dtypes.int)+(-a).cast(dtypes.int)
     self.assertIs(full_rewrite(out.sink()).src[0], full_rewrite(UOp.const(0, dtypes.int).sink()).src[0])
 
   def test_where_on_gated_load_fold(self):
@@ -429,7 +429,7 @@ class TestReduceCollapse(unittest.TestCase):
 
 class TestMovementOps(unittest.TestCase):
   def test_pm_mops_partial_reshape_index_removes_reshape(self):
-    from tinygrad.schedule.prepare import pm_mops
+    from tinygrad.schedule.rangeify import pm_mops
     src = UOp.param(0, dtypes.float, shape=(32, 4))
     r0, r1 = UOp.range(4, 0), UOp.range(8, 1)
     result = graph_rewrite(src.reshape((4, 8, 4)).index(r0, r1), pm_mops, name="test")
@@ -439,7 +439,7 @@ class TestMovementOps(unittest.TestCase):
     self.assertNotIn(Ops.RESHAPE, [u.op for u in result.toposort()])
 
   def test_pm_mops_partial_reshape_index_suffix_mismatch_does_nothing(self):
-    from tinygrad.schedule.prepare import pm_mops
+    from tinygrad.schedule.rangeify import pm_mops
     src = UOp.param(0, dtypes.float, shape=(2, 6))
     result = graph_rewrite(src.reshape((2, 3, 2)).index(UOp.range(2, 0)), pm_mops, name="test")
     self.assertEqual(result.op, Ops.INDEX)

test/null/test_uop_symbolic.py

@@ -6,6 +6,7 @@ from tinygrad.dtype import dtypes, ConstType, DType, Invalid
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad.uop.spec import spec_shared, type_verify
 from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load
+from tinygrad.uop.weak import pm_cast_weak
 from tinygrad.uop.validate import uops_to_z3
 
 def check_uop_against_string(self, v:UOp, s:str):
@@ -35,7 +36,7 @@ class TestSymbolic(unittest.TestCase):
     self.assertEqual(solver.check(expr1 != expr2), z3.unsat, "simplified expression not equal to original")
 
   def helper_test_variable(self, v, n, m, s, test_z3:bool=True):
-    v_simplified = graph_rewrite(v, sym, name="simplify symbolic uop")
+    v_simplified = graph_rewrite(v, sym+pm_cast_weak, name="simplify symbolic uop")
     if test_z3: self.check_equal_z3(v, v_simplified)
     nmin, nmax = v_simplified.vmin, v_simplified.vmax
     check_uop_against_string(self, v_simplified, s)
@@ -148,13 +149,6 @@ class TestSymbolic(unittest.TestCase):
   def test_xor_0(self):
     self.helper_test_variable(Variable("a", 0, 8, dtypes.int) ^ 0, 0, 8, "a", test_z3=False)
 
-  def test_or_0(self):
-    self.helper_test_variable(Variable("a", 0, 8, dtypes.int) | 0, 0, 8, "a", test_z3=False)
-
-  def test_shift_0(self):
-    self.helper_test_variable(Variable("a", 0, 8, dtypes.int) << 0, 0, 8, "a")
-    self.helper_test_variable(Variable("a", 0, 8, dtypes.int) >> 0, 0, 8, "a")
-
   def test_xor_self_inverse(self):
     self.helper_test_variable((Variable("a", 0, 8, dtypes.int) ^ 5) ^ 5, 0, 8, "a", test_z3=False)
 
@@ -1023,7 +1017,7 @@ class TestSymbolic(unittest.TestCase):
     cond = Variable("s", 0, 3, dtypes.int) < 2
     a = Variable("a", 0, 3, dtypes.int)
     self.assertIs(graph_rewrite(cond.where(a, a+1).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), (a+1).cast(dtypes.half)))
-    self.assertIs(graph_rewrite(cond.where(a, uconst(2)).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), uconst(2.0)))
+    self.assertIs(graph_rewrite(cond.where(a, uconst(2)).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.const(2, dtypes.half)))
     self.assertIs(graph_rewrite(cond.where(a, UOp.invalid()).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.invalid()))
 
   def test_where_const_gate_keeps_stated_width(self):

test/null/test_uop_vmin_vmax.py

@@ -82,7 +82,7 @@ class TestVminVmaxProperties(unittest.TestCase):
   def test_vmin_vmax_multiplication_0_inf(self):
     # vmin and vmax for multiplication with a variable
     x = UOp.const(0.0)
-    y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(0))
+    y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(0), dtype=dtypes.float)
     uop = x * y
     # TODO: these should be 0, but definitely should not be nan
     self.assertEqual(uop.vmin, -math.inf)

test/null/test_uops.py

@@ -5,8 +5,8 @@ from tinygrad.tensor import Tensor
 from tinygrad.helpers import Timing, Context, cdiv
 from tinygrad.dtype import dtypes, AddrSpace, ConstFloat, Invalid  # noqa: F401
 from tinygrad.device import Device
-from tinygrad.uop.ops import Ops, AxisType, ParamArg, PatternMatcher, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
-from tinygrad.uop.weak import pm_lower_weak
+from tinygrad.uop.ops import Ops, ParamArg, PatternMatcher, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
+from tinygrad.uop.weak import pm_lower_index_dtype
 from tinygrad.uop.spec import spec_program, spec_shared, type_verify
 from tinygrad.uop.symbolic import sym, pm_remove_invalid
 from test.helpers import eval_uop, to_uops_list
@@ -56,7 +56,7 @@ class TestDTypeFromUOp(unittest.TestCase):
       if u.is_invalid)), (dtypes.float32, dtypes.float32, dtypes.bool))
     invalid, value = UOp.invalid(), UOp.const(1, dtypes.float32)
     for u in (UOp.param(0, dtypes.bool, ()).where(value, invalid), value+invalid, UOp.stack(value, invalid)): self.assertIs(u.src[-1], invalid)
-    for u in (UOp(Ops.STACK, src=(value, invalid)), UOp(Ops.ADD, src=(value, invalid)),
+    for u in (UOp(Ops.STACK, dtypes.float32, src=(value, invalid)), UOp(Ops.ADD, dtypes.float32, src=(value, invalid)),
               UOp.const(True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)),
               UOp.param(0, dtypes.float32, (4,)).index(invalid)): type_verify(u, spec_shared)
     gate, value = UOp.param(0, dtypes.bool, ()), UOp.param(1, dtypes.float, ())
@@ -64,7 +64,7 @@ class TestDTypeFromUOp(unittest.TestCase):
     type_verify(out.sink(), spec_program)
 
   def test_remove_invalid_stack_lanes(self):
-    stack = UOp(Ops.STACK, src=(UOp.const(1, dtypes.half), UOp.invalid()))
+    stack = UOp(Ops.STACK, dtypes.half, (UOp.const(1, dtypes.half), UOp.invalid()))
     out = graph_rewrite(stack, pm_remove_invalid)
     self.assertEqual(out.src, (UOp.const(1, dtypes.half), UOp.const(0, dtypes.half)))
     type_verify(out.sink(), spec_program)
@@ -76,18 +76,16 @@ class TestLowerIndexDtype(unittest.TestCase):
     buf = UOp.param(0, dtypes.float, (2**31+64,))
     i = UOp.variable("i", 0, 2**28)
     shrink = UOp(Ops.SHRINK, src=(buf, (i*24).valid(i < 2**28), UOp.const(4)))
-    lowered = graph_rewrite(shrink.sink(), pm_lower_weak)
-    self.assertTrue(all(u.op is Ops.CONST for u in lowered.backward_slice_with_self if u.dtype in dtypes.weaks),
-                    "lowering must resolve every weak width, except a typed literal's value half")
+    lowered = graph_rewrite(shrink.sink(), pm_lower_index_dtype)
+    self.assertTrue(all(u.dtype != dtypes.weakint for u in lowered.backward_slice_with_self), "lowering must resolve all weakint")
     sh = next(u for u in lowered.backward_slice_with_self if u.op is Ops.SHRINK)
     self.assertEqual(sh.src[1].dtype, dtypes.long)
 
   def test_reg_buffer_size_lowers(self):
     reg = UOp.placeholder((4,), dtypes.float, 0, addrspace=AddrSpace.REG)
     self.assertEqual(reg.src[0].dtype, dtypes.weakint)
-    lowered = graph_rewrite(reg.sink(), pm_lower_weak)
-    self.assertTrue(all(u.op is Ops.CONST for u in lowered.backward_slice_with_self if u.dtype in dtypes.weaks),
-                    "lowering must resolve every weak width, except a typed literal's value half")
+    lowered = graph_rewrite(reg.sink(), pm_lower_index_dtype)
+    self.assertTrue(all(u.dtype != dtypes.weakint for u in lowered.backward_slice_with_self), "lowering must resolve all weakint")
     self.assertEqual(next(u for u in lowered.backward_slice_with_self if u.op is Ops.BUFFER).src[0].dtype, dtypes.int)
 
 class TestSafeCast(unittest.TestCase):
@@ -282,9 +280,9 @@ class TestFastIdiv(unittest.TestCase):
   def test_division_power_of_two(self):
     for dt in (dtypes.int32, dtypes.uint32):
       g = UOp.param(0, dt, (3,))
-      c = UOp.const(2)
+      c = UOp.const(2).cast(dt)
       l = g.index(c)
-      a = UOp(Ops.CDIV, src=(l, c))
+      a = UOp(Ops.CDIV, dt, (l, c))
       uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
       Device[Device.DEFAULT].renderer.render(uops)
       ops = [x.op for x in uops]
@@ -295,32 +293,31 @@ class TestFastIdiv(unittest.TestCase):
     # FLOORMOD by a power of two lowers to AND (correct floor mod for any sign in two's complement)
     for dt in (dtypes.int32, dtypes.uint32):
       g = UOp.param(0, dt, (9,))
-      c = UOp.const(8)
-      a = UOp(Ops.FLOORMOD, src=(g.index(c), c))
+      c = UOp.const(8).cast(dt)
+      a = UOp(Ops.FLOORMOD, dt, (g.index(c), c))
       uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
       ops = [x.op for x in uops]
       self.assertIn(Ops.AND, ops, f"For dtype={dt} FLOORMOD by pow2 did not simplify to AND")
       self.assertNotIn(Ops.CMOD, ops, f"For dtype={dt} FLOORMOD by pow2 left a MOD")
       self.assertNotIn(Ops.FLOORMOD, ops, f"For dtype={dt} FLOORMOD survived past late rewrite")
 
-  def test_floordiv_power_of_two(self):
-    # FLOORDIV by a power of two lowers to a shift, with no round toward zero correction (a shift is exactly floor division)
-    for dt in (dtypes.int32, dtypes.uint32, dtypes.int64, dtypes.uint64):
+  def test_floordiv_power_of_two_uint(self):
+    # uint FLOORDIV by a power of two lowers to a shift, leaving no IDIV/FLOORDIV in the kernel
+    for dt in (dtypes.uint32, dtypes.uint64):
       g = UOp.param(0, dt, (3,))
-      c = UOp.const(2)
-      a = UOp(Ops.FLOORDIV, src=(g.index(c), c))
+      c = UOp.const(2).cast(dt)
+      a = UOp(Ops.FLOORDIV, dt, (g.index(c), c))
       uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
       ops = [x.op for x in uops]
       self.assertIn(Ops.SHR, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
       self.assertNotIn(Ops.CDIV, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
-      self.assertNotIn(Ops.CMOD, ops, f"For dtype={dt} FLOORDIV by pow2 kept the round toward zero correction")
       self.assertNotIn(Ops.FLOORDIV, ops, f"For dtype={dt} FLOORDIV survived past late rewrite")
 
   @Context(DISABLE_FAST_IDIV=0)
   @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
   def test_fast_idiv_and_mod(self):
     g = UOp.param(0, dtypes.uint32, (4,))
-    c = UOp.const(3)
+    c = UOp.const(3).cast(dtypes.uint)
     l = g.index(c)
     a = UOp(Ops.CDIV, src=(l, c))
     uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -340,7 +337,7 @@ class TestFastIdiv(unittest.TestCase):
   def test_fast_idiv_bounded_numerator_zero(self):
     x = UOp.variable("x", 0, 1, dtype=dtypes.int32)
     for val in range(2):
-      self.assertEqual(eval_uop(x.alu(Ops.CDIV, UOp.const(3)), vals=(val,)), cdiv(val, 3))
+      self.assertEqual(eval_uop(x.alu(Ops.CDIV, UOp.const(3).cast(x.dtype)), vals=(val,)), cdiv(val, 3))
 
   @Context(DISABLE_FAST_IDIV=0)
   def test_fast_idiv_remove_powers_of_two(self):
@@ -365,7 +362,7 @@ class TestFastIdiv(unittest.TestCase):
 
   def test_disable_fast_idiv(self):
     g = UOp.param(0, dtypes.uint32, (4,))
-    c = UOp.const(3)
+    c = UOp.const(3).cast(dtypes.uint)
     l = g.index(c)
     a = UOp(Ops.CDIV, src=(l, c))
     with Context(DISABLE_FAST_IDIV=1):
@@ -460,19 +457,11 @@ class TestUopsObject(unittest.TestCase):
     self.assertEqual(a.device, Device.DEFAULT)
 
 class TestUOpRender(unittest.TestCase):
-  def test_render_ssimplified_marg_outside_toposort(self):
-    r = UOp.range(UOp.const(16, dtypes.int), 2, AxisType.WEAK, dtype=dtypes.int)
-    offset = (r * 2) + (r * 2)
-    shrink = UOp(Ops.SHRINK, src=(UOp.param(0, dtypes.uint, (32,)), offset, UOp.const(2, dtypes.int)))
-    self.assertIsNot(shrink.src[1], shrink.marg[0][0])
-    self.assertEqual(shrink.render(simplify=False), "p0.shrink((((r2*4), 2),))")
-    self.assertEqual(UOp.range(1, 0, src=(shrink,), dtype=dtypes.int).render(simplify=False), "r0")
-
   def test_render_vectorize_empty(self):
-    u = UOp(Ops.STACK, src=())
+    u = UOp(Ops.STACK, dtype=dtypes.void, src=())
     self.assertEqual(u.render(simplify=False), "{}")
   def test_render_vectorize_empty_simplified(self):
-    u = UOp(Ops.STACK, src=())
+    u = UOp(Ops.STACK, dtype=dtypes.void, src=())
     self.assertEqual(u.render(), "{}")
   def test_render_vectorize_same(self):
     u = UOp(Ops.STACK, src=(UOp.const(0),)*3)

test/null/test_validate_oob.py

@@ -14,37 +14,37 @@ class TestValidateOOB(unittest.TestCase):
   def test_const_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
       buf = UOp.param(0, dtypes.int, (16,))
-      to_uops_list([buf.index(UOp.const(0)).load()])  # valid
-      to_uops_list([buf.index(UOp.const(15)).load()])  # valid (last element)
+      to_uops_list([buf.index(UOp.const(0)).load(dtype=dtypes.int)])  # valid
+      to_uops_list([buf.index(UOp.const(15)).load(dtype=dtypes.int)])  # valid (last element)
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(UOp.const(16)).load()])  # off by one
+        to_uops_list([buf.index(UOp.const(16)).load(dtype=dtypes.int)])  # off by one
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(UOp.const(42)).load()])  # way out
+        to_uops_list([buf.index(UOp.const(42)).load(dtype=dtypes.int)])  # way out
 
   def test_variable_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
       buf = UOp.param(0, dtypes.int, (16,))
-      to_uops_list([buf.index(Variable("i", 0, 15)).load()])  # valid
+      to_uops_list([buf.index(Variable("i", 0, 15)).load(dtype=dtypes.int)])  # valid
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(Variable("i", 0, 20)).load()])  # oob
+        to_uops_list([buf.index(Variable("i", 0, 20)).load(dtype=dtypes.int)])  # oob
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(Variable("i", -5, 10)).load()])  # negative
+        to_uops_list([buf.index(Variable("i", -5, 10)).load(dtype=dtypes.int)])  # negative
 
   def test_range_with_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
       buf = UOp.param(0, dtypes.int, (16,))
       r = UOp.range(42, 0, AxisType.GLOBAL)
-      to_uops_list([buf.index(r.valid(r < 16)).load()])  # valid
+      to_uops_list([buf.index(r.valid(r < 16)).load(dtype=dtypes.int)])  # valid
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(r.valid(r < 17)).load()])  # oob
+        to_uops_list([buf.index(r.valid(r < 17)).load(dtype=dtypes.int)])  # oob
 
   def test_variable_with_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
       buf = UOp.param(0, dtypes.int, (16,))
       v = Variable("v", -5, 80)
-      to_uops_list([buf.index(v.valid((v >= 0) & (v < 16))).load()])  # valid
+      to_uops_list([buf.index(v.valid((v >= 0) & (v < 16))).load(dtype=dtypes.int)])  # valid
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(v.valid(v < 20)).load()])  # negative not masked
+        to_uops_list([buf.index(v.valid(v < 20)).load(dtype=dtypes.int)])  # negative not masked
 
   def test_gated_store(self):
     with Context(CHECK_OOB=1, SPEC=2):
@@ -58,62 +58,62 @@ class TestValidateOOB(unittest.TestCase):
   def test_floordiv(self):
     with Context(CHECK_OOB=1, SPEC=2):
       buf = UOp.param(0, dtypes.int, (16,))
-      to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2).load()])  # 0..15 valid
+      to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2).load(dtype=dtypes.int)])  # 0..15 valid
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(UOp.range(34, 0, AxisType.GLOBAL) // 2).load()])  # 0..16 oob
+        to_uops_list([buf.index(UOp.range(34, 0, AxisType.GLOBAL) // 2).load(dtype=dtypes.int)])  # 0..16 oob
 
   def test_mod(self):
     with Context(CHECK_OOB=1, SPEC=2):
       buf = UOp.param(0, dtypes.int, (16,))
       r = UOp.range(100, 0, AxisType.GLOBAL)
-      to_uops_list([buf.index(r % 16).load()])  # 0..15 valid
+      to_uops_list([buf.index(r % 16).load(dtype=dtypes.int)])  # 0..15 valid
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(r % 20).load()])  # 0..19 oob
+        to_uops_list([buf.index(r % 20).load(dtype=dtypes.int)])  # 0..19 oob
 
   def test_shr(self):
     with Context(CHECK_OOB=1, SPEC=2):
       buf = UOp.param(0, dtypes.int, (16,))
-      to_uops_list([buf.index(UOp.range(64, 0, AxisType.GLOBAL) >> 2).load()])  # 0..15 valid
+      to_uops_list([buf.index(UOp.range(64, 0, AxisType.GLOBAL) >> 2).load(dtype=dtypes.int)])  # 0..15 valid
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(UOp.range(128, 0, AxisType.GLOBAL) >> 2).load()])  # 0..31 oob
+        to_uops_list([buf.index(UOp.range(128, 0, AxisType.GLOBAL) >> 2).load(dtype=dtypes.int)])  # 0..31 oob
 
   def test_shl(self):
     with Context(CHECK_OOB=1, SPEC=2):
       buf = UOp.param(0, dtypes.int, (64,))
       r = UOp.range(8, 0, AxisType.GLOBAL)
-      to_uops_list([buf.index(r << 2).load()])  # 0..28 valid
+      to_uops_list([buf.index(r << 2).load(dtype=dtypes.int)])  # 0..28 valid
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(r << 4).load()])  # 0..112 oob
+        to_uops_list([buf.index(r << 4).load(dtype=dtypes.int)])  # 0..112 oob
 
   def test_and(self):
     with Context(CHECK_OOB=1, SPEC=2):
       buf = UOp.param(0, dtypes.int, (16,))
       r = UOp.range(100, 0, AxisType.GLOBAL)
-      to_uops_list([buf.index(r & 15).load()])  # 0..15 valid
+      to_uops_list([buf.index(r & 15).load(dtype=dtypes.int)])  # 0..15 valid
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(r & 31).load()])  # 0..31 oob
+        to_uops_list([buf.index(r & 31).load(dtype=dtypes.int)])  # 0..31 oob
       # align masks round down to a multiple of 2^k
-      to_uops_list([buf.index((r & -4).valid(r < 16)).load()])  # 0..12 valid
+      to_uops_list([buf.index((r & -4).valid(r < 16)).load(dtype=dtypes.int)])  # 0..12 valid
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(r & -2).load()])  # 0..100 oob
+        to_uops_list([buf.index(r & -2).load(dtype=dtypes.int)])  # 0..100 oob
       # other masks can't be modeled as mod
       with self.assertRaisesRegex(RuntimeError, "z3 int AND only supports"):
-        to_uops_list([buf.index(r & 21).load()])
+        to_uops_list([buf.index(r & 21).load(dtype=dtypes.int)])
 
   def test_max(self):
     with Context(CHECK_OOB=1, SPEC=2):
       buf = UOp.param(0, dtypes.int, (16,))
-      to_uops_list([buf.index(Variable("v", -10, 15).maximum(0)).load()])  # 0..15 valid
+      to_uops_list([buf.index(Variable("v", -10, 15).maximum(0)).load(dtype=dtypes.int)])  # 0..15 valid
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(Variable("v2", -10, 20).maximum(0)).load()])  # 0..20 oob
+        to_uops_list([buf.index(Variable("v2", -10, 20).maximum(0)).load(dtype=dtypes.int)])  # 0..20 oob
 
   def test_xor_in_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
       buf = UOp.param(0, dtypes.int, (16,))
       r = UOp.range(32, 0, AxisType.GLOBAL)
-      to_uops_list([buf.index(r.valid((r < 8) ^ ((r >= 8) & (r < 16)))).load()])  # 0..15 valid
+      to_uops_list([buf.index(r.valid((r < 8) ^ ((r >= 8) & (r < 16)))).load(dtype=dtypes.int)])  # 0..15 valid
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(r.valid((r < 10) ^ (r >= 20))).load()])  # 0..9,20..31 oob
+        to_uops_list([buf.index(r.valid((r < 10) ^ (r >= 20))).load(dtype=dtypes.int)])  # 0..9,20..31 oob
 
   # cast patterns
   def test_float_cast_in_index(self):
@@ -121,13 +121,13 @@ class TestValidateOOB(unittest.TestCase):
       buf = UOp.param(0, dtypes.int, (16,))
       r = UOp.range(20, 0)
       i = (r.cast(dtypes.float) * 0.68).trunc().cast(dtypes.int)
-      to_uops_list([buf.index(i.valid((i >= 0) & (i < 16))).load()])
+      to_uops_list([buf.index(i.valid((i >= 0) & (i < 16))).load(dtype=dtypes.int)])
 
   def test_bool_cast_in_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
       buf = UOp.param(0, dtypes.int, (1,))
       r = UOp.range(20, 0)
-      to_uops_list([buf.index(r.valid(r.cast(dtypes.bool).logical_not())).load()])  # only r=0 valid
+      to_uops_list([buf.index(r.valid(r.cast(dtypes.bool).logical_not())).load(dtype=dtypes.int)])  # only r=0 valid
 
   # load result as index/mask
   def test_load_as_index(self):
@@ -135,18 +135,18 @@ class TestValidateOOB(unittest.TestCase):
       buf0 = UOp.param(0, dtypes.int, (16,))
       buf1 = UOp.param(1, dtypes.int, (64,))
       r = UOp.range(42, 0, AxisType.GLOBAL)
-      ld0 = buf0.index(r.valid(r < 8)).load().cast(dtypes.weakint)
-      to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32))).load()])  # valid
+      ld0 = buf0.index(r.valid(r < 8)).load(dtype=dtypes.int).cast(dtypes.weakint)
+      to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32))).load(dtype=dtypes.int)])  # valid
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 64))).load()])  # oob
+        to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 64))).load(dtype=dtypes.int)])  # oob
 
   def test_load_from_shrink_as_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
       buf0 = UOp.param(0, dtypes.int, (16,))
       buf1 = UOp.param(1, dtypes.int, (64,))
       shrink = UOp(Ops.SHRINK, src=(buf0, UOp.const(0, dtypes.int), UOp.const(4)))
-      ld0 = shrink.load().index(0)
-      to_uops_list([buf1.index(ld0.valid((ld0 >= 0) & (ld0 < 64))).load()])
+      ld0 = shrink.load(dtype=dtypes.int).index(0)
+      to_uops_list([buf1.index(ld0.valid((ld0 >= 0) & (ld0 < 64))).load(dtype=dtypes.int)])
 
   def test_load_bool_as_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):

test/null/test_viz.py

@@ -1,5 +1,5 @@
-import unittest
-import decimal, sys, json, contextlib, tempfile, pickle, io, math, pathlib
+import unittest, decimal, sys, json, contextlib, tempfile, pickle, io, math
+from pathlib import Path
 from dataclasses import dataclass
 from typing import Generator
 
@@ -516,22 +516,6 @@ class TestVizIntegration(unittest.TestCase):
     src_render = get_render(viz.data, steps[src_idx]["query"])["src"]
     self.assertEqual(src, src_render)
 
-  def test_profiler_duplicate_name(self):
-    kernel_name = "duplicate_name"
-    def one(A:UOp): return A[0].store(UOp.const(1.0, dtypes.float)).sink(arg=KernelInfo(kernel_name))
-    def zero(A:UOp): return A[0].store(UOp.const(0.0, dtypes.float)).sink(arg=KernelInfo(kernel_name))
-    with save_viz() as viz:
-      @TinyJit
-      def f(a:Tensor, b:Tensor): return Tensor.custom_kernel(a, fxn=one)[0], Tensor.custom_kernel(b, fxn=zero)[0]
-      a, b = Tensor.empty(4, device="NULL"), Tensor.empty(4, device="NULL")
-      # warmup
-      for _ in range(2): Tensor.realize(*f(a, b))
-      Tensor.realize(*f(a, b))
-    kernels = {i for i,c in enumerate(viz.list_items()) if c["name"] == kernel_name}
-    profile = decode_profile(unwrap(get_profile(viz.data, cpu_events)))
-    events = [e for e in profile["layout"]["NULL"]["events"] if e["name"] == kernel_name]
-    self.assertEqual({e["ref"] for e in events}, kernels)
-
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry
 from tinygrad.viz.serve import get_profile
 from tinygrad.viz.cli import decode_profile
@@ -835,6 +819,8 @@ from extra.gemm.amd_asm_matmul import Kernel
 
 @needs_tracked_pm
 class TestCfg(unittest.TestCase):
+  def setUp(self): self.arch = "gfx1100"
+
   def get_cfg(self, name:str, k:Kernel):
     insts = k.finalize()
     def fxn(out:UOp) -> UOp:
@@ -843,7 +829,7 @@ class TestCfg(unittest.TestCase):
       sink = UOp.sink(out.base, lidx, gidx, arg=KernelInfo(name=name))
       return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
     with save_viz() as viz:
-      with Context(DEV="NULL::gfx1100"):
+      with Context(DEV=f"NULL::{self.arch}"):
         out = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
         _ = do_to_program(out.schedule_linear().src[-1].src[0], Device[out.device].renderer)
     codegen_rewrites = next(s for s in viz.list_items() if s["name"] == name)
@@ -1025,8 +1011,8 @@ def run_cli(*cli_args) -> list[dict]:
 @contextlib.contextmanager
 def write_files(viz) -> list[str]:
   with tempfile.TemporaryDirectory() as tmpdir:
-    (r:=pathlib.Path(tmpdir)/"rewrites.pkl").write_bytes(pickle.dumps(viz.data.trace))
-    (p:=pathlib.Path(tmpdir)/"profile.pkl").write_bytes(pickle.dumps(cpu_events))
+    (r:=Path(tmpdir)/"rewrites.pkl").write_bytes(pickle.dumps(viz.data.trace))
+    (p:=Path(tmpdir)/"profile.pkl").write_bytes(pickle.dumps(cpu_events))
     yield ["--rewrites-path", str(r), "--profile-path", str(p)]
 
 class TestCLI(unittest.TestCase):

test/unit/test_attention.py

@@ -5,8 +5,6 @@ from tinygrad.llm.model import (
   GatedDeltaNetBlock, SSMConfig, TransformerBlock, TransformerConfig,
   apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk,
 )
-from tinygrad.llm.kernels.amd import Linear, gated_delta_prefill, amd_custom_kernels_supported
-from tinygrad.llm.gguf import ggml_data_to_tensor
 
 def apply_rope(x:Tensor, start_pos:int):
   B, H, T, Hd = x.shape
@@ -14,15 +12,6 @@ def apply_rope(x:Tensor, start_pos:int):
   freqs_cis = precompute_freqs_cis(Hd, start_pos+T)[start_pos:start_pos+T]
   return apply_rope_new(x, freqs_cis)
 
-class TestLinear(unittest.TestCase):
-  def test_recovers_packed_ggml_weight(self):
-    for ggml_type,packed_size,words in ((13, 176, 44), (14, 210, 210), (23, 136, 34)):
-      packed = Tensor.empty(packed_size+4, dtype=dtypes.uint8, device="CPU")[4:]
-      decoded = ggml_data_to_tensor(packed, 256, ggml_type).reshape(1, 256)
-      linear = Linear(256, 1, bias=False)
-      linear.set_quantized(decoded)
-      self.assertEqual((linear.ggml_type, linear.weight.numel()), (ggml_type, words))
-
 class TestAttention(unittest.TestCase):
   def test_apply_rope(self):
     x = Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32)
@@ -52,23 +41,6 @@ class TestAttention(unittest.TestCase):
     np.testing.assert_allclose(block.cache_kv[0, :, :, :seqlen, :].numpy(), expected.numpy(), rtol=1e-5, atol=1e-5)
 
 class TestGatedDeltaNetBlock(unittest.TestCase):
-  def test_gated_delta_rectangular_state_and_row_decay(self):
-    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
-    rng = np.random.default_rng(42)
-    q, k = (rng.normal(size=(1, 1, 3, 32)).astype(np.float32) for _ in range(2))
-    v, beta = rng.normal(size=(1, 1, 3, 4)).astype(np.float32), rng.uniform(size=(1, 1, 3)).astype(np.float32)
-    alpha, initial = rng.uniform(0.8, 1, size=(1, 1, 3, 4)).astype(np.float32), rng.normal(size=(1, 1, 4, 32)).astype(np.float32)
-    expected_state, expected_out = initial.copy(), np.empty_like(v)
-    for t in range(3):
-      previous, av = expected_state.copy(), alpha[:, :, t, :, None]
-      delta = (v[:, :, t] - (previous*k[:, :, t, None]).sum(-1)*alpha[:, :, t]) * beta[:, :, t, None]
-      expected_state = previous*av + delta[..., None]*k[:, :, t, None, :]
-      expected_out[:, :, t] = (previous*q[:, :, t, None]).sum(-1)*alpha[:, :, t] + delta*(q[:, :, t]*k[:, :, t]).sum(-1)
-    state = Tensor(initial).contiguous().realize()
-    out = gated_delta_prefill(Tensor(q), Tensor(k), Tensor(v), Tensor(beta), Tensor(alpha), state).realize()
-    np.testing.assert_allclose(out.numpy(), expected_out, rtol=1e-4, atol=1e-4)
-    np.testing.assert_allclose(state.numpy(), expected_state, rtol=1e-4, atol=1e-4)
-
   def _tensor_linspace(self, start:float, stop:float, shape:tuple[int, ...]) -> Tensor:
     return Tensor.linspace(start, stop, int(np.prod(shape)), dtype=dtypes.float32).reshape(*shape)
 

test/unit/test_call.py

@@ -224,21 +224,6 @@ class TestCallSchedule(unittest.TestCase):
     np.testing.assert_equal(x.numpy(), [2, 2, 2])
     np.testing.assert_equal(y.numpy(), [3, 3, 3])
 
-  def test_precompile_nested_scope_collision(self):
-    # a precompiled function body gets its own positional p{slot} params; they must not be renumbered when the call is
-    # scheduled inside an enclosing realize with a different slot ordering. the store must use this call's Variable
-    cache = Tensor.zeros(16)
-    @function(precompile=True, allow_implicit=True)
-    def store(x:Tensor, sp:UOp) -> Tensor:
-      # update a cache at a symbolic offset, like an attention KV cache update
-      return Tensor(cache.uop.after(cache[sp:sp+x.shape[0]].uop.store(x.uop)))[:sp+x.shape[0]].sum()
-    sp_v, nt_v = UOp.variable("sp", 0, 8), UOp.variable("nt", 1, 8)
-    t = Tensor.arange(16).float().realize()
-    sp, nt = sp_v.bind(0), nt_v.bind(8)
-    store(t[sp:sp+nt].clone().realize(), sp).realize()
-    np.testing.assert_equal(cache.numpy()[:8], t[:8].numpy())
-    np.testing.assert_equal(cache.numpy()[8:], np.zeros(8))
-
   def test_precompile_schedule_cache_hit(self):
     """two instances of the same @function should produce identical function body keys (schedule cache hit)"""
     @function(precompile=True)

test/unit/test_dtype_weak.py

@@ -4,7 +4,7 @@ from tinygrad import Tensor, dtypes, TinyJit
 from tinygrad.helpers import Context
 from tinygrad.dtype import least_upper_float
 from tinygrad.uop.ops import UOp, Ops, GroupOp, dtype_from_uop, graph_rewrite
-from tinygrad.uop.weak import pm_commit_weak
+from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak
 from tinygrad.uop.symbolic import symbolic_simple
 from tinygrad.uop.spec import spec_shared, type_verify
 from tinygrad.engine.jit import JitError
@@ -74,7 +74,7 @@ class TestWeakPromotion(unittest.TestCase):
     recips = [u for u in (x / y)._uop.toposort() if u.op is Ops.RECIPROCAL]
     self.assertEqual([(u.dtype, u.src[0].dtype) for u in recips], [(dtypes.float32, dtypes.float32)])
     with Context(DEFAULT_FLOAT=dtypes.float16):
-      committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_commit_weak)
+      committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_lower_index_dtype, ctx={})
     self.assertEqual([u.dtype for u in committed.toposort() if u.op is Ops.ADD], [dtypes.float32])
 
   def test_div_sub_operand_kept_weak(self):
@@ -85,7 +85,7 @@ class TestWeakPromotion(unittest.TestCase):
   def test_cast_weak_expression_commits_at_cast_floor(self):
     # the floor never narrows: a cast BELOW the default does not pull the compute width down with it
     with Context(DEFAULT_FLOAT=dtypes.float32):
-      narrowed = graph_rewrite((UOp.const(1.0) + UOp.const(2.0)).cast(dtypes.float16), pm_commit_weak)
+      narrowed = graph_rewrite((UOp.const(1.0) + UOp.const(2.0)).cast(dtypes.float16), pm_lower_index_dtype, ctx={})
     self.assertEqual((narrowed.dtype, narrowed.src[0].dtype), (dtypes.float16, dtypes.float32))
 
   def test_cast_weak_expression_value_uses_cast_floor(self):
@@ -114,35 +114,27 @@ class TestWeakPromotion(unittest.TestCase):
     self.assertIsInstance((x + 2).src[1].val, float)
     self.assertIs(x + UOp.const(2), x + 2)
 
+  def test_index_dtype_ignores_weakness(self):
+    with Context(SPEC=2):
+      idx = UOp.const(0).cast(dtypes.int32)
+      weak = UOp.const(1.0).expand((1,))
+      self.assertEqual(UOp(Ops.INDEX, dtypes.float32, (weak, idx)).dtype, dtypes.float32)
+      with self.assertRaisesRegex(RuntimeError, "bad dtype"): UOp(Ops.INDEX, dtypes.int32, (weak, idx))
+
   def test_store_weak_value_uses_destination_dtype(self):
     with Context(DEFAULT_FLOAT=dtypes.float16):
       dst = UOp.param(0, dtypes.bfloat16, (1,)).index(UOp.const(0).cast(dtypes.int32))
       gate = UOp.const(True)
-      out = graph_rewrite(dst.store(UOp.const(5.0), gate), pm_commit_weak)
+      out = graph_rewrite(dst.store(UOp.const(5.0), gate), pm_lower_index_dtype, ctx={})
     # a bare weak CONST commits directly: the pass runs without symbolic, so a CAST here would survive it
     self.assertEqual((out.src[1], out.src[2]), (UOp.const(5.0, dtypes.bfloat16), gate))
 
   def test_weak_srcs_commit_only_at_a_concrete_lub(self):
     weak_lub = UOp(Ops.ADD, src=(UOp.const(1), UOp.const(1.0)))
-    self.assertIs(graph_rewrite(weak_lub, pm_commit_weak), weak_lub)
+    self.assertIs(graph_rewrite(weak_lub, pm_lower_index_dtype, ctx={}), weak_lub)
     concrete = UOp.const(2.0).cast(dtypes.float16)
-    # the weak arm stays bare: its sibling states the width, so the WHERE already derives float16 for it
-    where = graph_rewrite(UOp(Ops.WHERE, src=(UOp.const(True), concrete, UOp.const(1.0))), pm_commit_weak)
-    self.assertEqual((where.dtype, tuple(x.dtype for x in where.src)), (dtypes.float16, (dtypes.bool, dtypes.float16, dtypes.weakfloat)))
-
-  def test_derivable_const_rounds_at_the_derived_width(self):
-    # re-rounds a derivable const in place (still bare) so value-keyed folds (x*1 -> x, x*-1 -> NEG) still fire
-    x = UOp.param(0, dtypes.float32, (1,)).index(UOp.const(0).cast(dtypes.int32)).load()
-    mul = graph_rewrite(x * UOp.const(-0.9999999893980771), symbolic_simple+pm_commit_weak)
-    self.assertIs(mul.src[1], UOp.const(-1.0))
-    self.assertIs(graph_rewrite(x * UOp.const(1.0000000106), symbolic_simple+pm_commit_weak), x)
-
-  def test_committed_const_conversion_folds_for_native_format(self):
-    folded = graph_rewrite(UOp.const(16256, dtypes.ushort).cast(dtypes.uint), symbolic_simple)
-    self.assertIs(folded, UOp.const(16256, dtypes.uint))
-    # fmt-less targets are lowered by renderer rewrites, where collapsing this pair would cycle with float-intermediate insertion.
-    emulated = UOp.const(1.0, dtypes.float).cast(dtypes.bfloat16)
-    self.assertIs(graph_rewrite(emulated, symbolic_simple), emulated)
+    where = graph_rewrite(UOp(Ops.WHERE, src=(UOp.const(True), concrete, UOp.const(1.0))), pm_lower_index_dtype, ctx={})
+    self.assertEqual(tuple(x.dtype for x in where.src), (dtypes.bool, dtypes.float16, dtypes.float16))
 
   def test_weak_shift_lhs_commits_the_node(self):
     # a shift derives its lhs's dtype, so committing the lhs restates the root (WGSL's packed store writes `mask << shift_am`)
@@ -183,11 +175,11 @@ class TestWeakPromotion(unittest.TestCase):
     self.assertEqual(dtype_from_uop(Ops.SHL, (UOp.const(1, dtypes.int8), UOp.const(1, dtypes.uint32)), None), dtypes.int8)
     self.assertEqual(UOp.const(1).alu(Ops.SHL, UOp.const(1, dtypes.uint)).dtype, dtypes.weakint)
     self.assertEqual((v & 3).dtype, dtypes.weakint)
-    with self.assertRaises(RuntimeError): (Tensor.const(1.0) << Tensor.const(1.0)).dtype
-    with self.assertRaises(RuntimeError): UOp.const(1, dtypes.int32).alu(Ops.SHL, UOp.const(1, dtypes.float64)).dtype
+    with self.assertRaises(RuntimeError): Tensor.const(1.0) << Tensor.const(1.0)
+    with self.assertRaises(RuntimeError): UOp.const(1, dtypes.int32).alu(Ops.SHL, UOp.const(1, dtypes.float64))
     for op in (Ops.SHL, Ops.SHR):
       with self.assertRaises(RuntimeError):
-        UOp.const(1, dtypes.float32).alu(op, UOp.const(1, dtypes.int32)).dtype
+        UOp.const(1, dtypes.float32).alu(op, UOp.const(1, dtypes.int32))
     # float bitwise builds, the spec rejects it
     with Context(SPEC=1):
       f32, wf = UOp.const(1.0, dtypes.float32), UOp.const(1.0)

test/unit/test_llm_amd.py

@@ -1,107 +0,0 @@
-import unittest
-import numpy as np
-from tinygrad import Tensor, UOp, dtypes, nn
-from tinygrad.llm.kernels.amd import Linear, amd_custom_kernels_supported, q8_quantize, flash_attention
-from tinygrad.llm.gguf import ggml_data_to_tensor
-
-class TestQ8Quantize(unittest.TestCase):
-  def test_word_quant_weights_use_typed_buffer_view(self):
-    for ggml_type, type_size in ((13, 176), (23, 136)):
-      with self.subTest(ggml_type=ggml_type):
-        raw = Tensor(np.zeros(type_size + 4, dtype=np.uint8), device="CPU").contiguous().realize()[4:]
-        decoded = ggml_data_to_tensor(raw, 256, ggml_type).reshape(1, 256)
-        linear = Linear(256, 1, bias=False)
-        linear.set_quantized(decoded)
-        self.assertEqual(linear.ggml_type, ggml_type)
-        self.assertEqual(linear.weight.dtype, dtypes.uint32)
-        self.assertEqual(linear.weight.nbytes(), type_size)
-        self.assertEqual(linear.weight.uop.buf_uop.buffer.offset, 4)
-
-  def test_values_and_scales(self):
-    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
-    x = np.linspace(-3.1, 2.7, 64, dtype=np.float32).reshape(2, 32)
-    quant, scale = q8_quantize(Tensor(x), 2, 32)
-    scale_np = np.maximum(np.max(np.abs(x), axis=-1, keepdims=True) / 127, 1e-8)
-    expected = np.clip(np.rint(x / scale_np), -127, 127).astype(np.int8)
-    np.testing.assert_array_equal(quant.bitcast(dtypes.int8).reshape(2, 32).numpy(), expected)
-    np.testing.assert_allclose(scale.numpy(), scale_np, rtol=1e-6)
-
-  def test_q6_linear_compiles(self):
-    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
-    rng = np.random.default_rng(42)
-    packed = rng.integers(0, 256, 210, dtype=np.uint8)
-    packed[-2:] = np.array([0.01], dtype=np.float16).view(np.uint8)
-    raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
-    decoded = ggml_data_to_tensor(raw, 256, 14).reshape(1, 256)
-    linear = Linear(256, 1, bias=False)
-    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
-    self.assertTrue(np.isfinite(linear(Tensor.randn(1, 256)).realize().item()))
-    self.assertEqual(linear.weight.uop.buf_uop.buffer.offset, 4)
-
-  def test_q4_k_linear(self):
-    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
-    rng = np.random.default_rng(42)
-    in_features, blocks = 2048, 16*2048//256
-    packed = rng.integers(0, 256, blocks*144, dtype=np.uint8)
-    for i in range(blocks): packed[i*144:i*144+4] = np.array([0.01, 0.002], dtype=np.float16).view(np.uint8)
-    raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
-    decoded = ggml_data_to_tensor(raw, 16*in_features, 12).reshape(16, in_features)
-    weight = decoded.numpy()
-    linear = Linear(in_features, 16, bias=False)
-    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
-    x = rng.normal(size=(3, in_features)).astype(np.float32)
-    scale = np.maximum(np.abs(x).reshape(3, in_features//32, 32).max(-1, keepdims=True) / 127, 1e-8)
-    xq = np.clip(np.rint(x.reshape(3, in_features//32, 32) / scale), -127, 127) * scale
-    np.testing.assert_allclose(linear(Tensor(x)).numpy(), xq.reshape(3, in_features) @ weight.T, rtol=2e-3, atol=2e-2)
-    self.assertEqual(linear.ggml_type, 12)
-
-  def test_q6_linear_multiple_tokens(self):
-    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
-    rng = np.random.default_rng(42)
-    in_features, blocks = 2048, 16*2048//256
-    packed = rng.integers(0, 256, blocks*210, dtype=np.uint8)
-    for i in range(blocks): packed[i*210+208:i*210+210] = np.array([0.01], dtype=np.float16).view(np.uint8)
-    raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
-    decoded = ggml_data_to_tensor(raw, 16*in_features, 14).reshape(16, in_features)
-    weight = decoded.numpy()
-    linear = Linear(in_features, 16, bias=False)
-    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
-    x = rng.normal(size=(3, in_features)).astype(np.float32)
-    scale = np.maximum(np.abs(x).reshape(3, in_features//32, 32).max(-1, keepdims=True) / 127, 1e-8)
-    xq = np.clip(np.rint(x.reshape(3, in_features//32, 32) / scale), -127, 127) * scale
-    np.testing.assert_allclose(linear(Tensor(x)).numpy(), xq.reshape(3, in_features) @ weight.T, rtol=2e-3, atol=2e-2)
-    self.assertEqual(linear.ggml_type, 14)
-
-    # symbolic token counts take the padded kernel path and give the same results
-    generic = Linear(in_features, 16, bias=False)
-    nn.state.load_state_dict(generic, {"weight":decoded}, verbose=False, realize=False)
-    sym = Tensor(np.concatenate([x, np.zeros((1, in_features), np.float32)])).contiguous()[:UOp.variable("tokens", 1, 4).bind(3)]
-    np.testing.assert_allclose(generic(sym)[:3].numpy(), xq.reshape(3, in_features) @ weight.T, rtol=2e-3, atol=2e-2)
-    self.assertTrue(generic.use_custom_quant)
-    self.assertEqual(generic.ggml_type, 14)
-
-  def test_attention_uses_physical_cache_length(self):
-    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
-    q, k, v = Tensor.zeros(1, 2, 1, 32), Tensor.randn(1, 1, 1, 32), Tensor.randn(1, 1, 1, 32)
-    cache = Tensor.empty(2, 1, 1, 256, 32, dtype=dtypes.half).contiguous()
-    assigned = Tensor(cache.uop.after(cache[:, :, :, 0:1, :].uop.store(Tensor.stack(k, v).cast(dtypes.half).uop)))
-    out = flash_attention(q, assigned, 1).realize()
-    np.testing.assert_allclose(out.numpy(), v.expand(1, 2, 1, 32).numpy(), rtol=2e-2, atol=2e-2)
-
-  def test_prefill_attention_unaligned_start(self):
-    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
-    rng = np.random.default_rng(42)
-    start_pos = 1718
-    q = Tensor.zeros(1, 8, 32, 128)
-    old_kv = rng.normal(size=(2, 1, 1, start_pos, 128)).astype(np.float32)
-    new_kv = rng.normal(size=(2, 1, 1, 32, 128)).astype(np.float32)
-    cache = Tensor.zeros(2, 1, 1, 2048, 128, dtype=dtypes.half).contiguous()
-    Tensor.realize(cache[:, :, :, :start_pos].assign(Tensor(old_kv).cast(dtypes.half)))
-    sp = UOp.variable("start_pos", 0, 2047).bind(start_pos)
-    assigned = Tensor(cache.uop.after(cache[:, :, :, sp:sp+32, :].uop.store(Tensor(new_kv).cast(dtypes.half).uop)))
-    out = flash_attention(q, assigned, sp+32).realize()
-    values = np.concatenate([old_kv[1, 0, 0], new_kv[1, 0, 0]]).astype(np.float16).astype(np.float32)
-    expected = np.stack([values[:start_pos+i+1].mean(0) for i in range(32)])[None, None].repeat(8, axis=1)
-    np.testing.assert_allclose(out.numpy(), expected, rtol=2e-3, atol=2e-3)
-
-if __name__ == "__main__": unittest.main()

test/unit/test_llm_server.py

@@ -1,8 +1,6 @@
 import unittest
-import numpy as np
 from unittest.mock import patch
 from tinygrad import Tensor, UOp
-from tinygrad.nn.state import get_state_dict
 from tinygrad.schedule import schedule_cache
 from tinygrad.llm.model import Transformer, TransformerConfig
 from tinygrad.llm.serve import StreamRouter
@@ -44,10 +42,7 @@ class TestTransformerGenerate(unittest.TestCase):
       return Tensor([[42]])
     with patch.object(Transformer, '__call__', mock_call):
       next(model.generate([1, 2, 3, 4, 5, 42, 10]))
-    # resumes from the reused state at position 5 and consumes the 2 new tokens (one chunk or two decode steps)
-    self.assertEqual(calls[0][1], V_START_POS.bind(5))
-    def ntok(shape): return shape[1] if isinstance(shape[1], int) else shape[1].unbind()[1]
-    self.assertEqual(sum(ntok(c[0]) for c in calls), 2)
+    self.assertEqual(calls, [((1, 1), V_START_POS.bind(5)), ((1, 1), V_START_POS.bind(6))])
 
   def test_recurrent_divergent_prompt_restarts(self):
     model, calls = Transformer(TEST_CONFIG), []
@@ -157,22 +152,6 @@ class TestTransformerGenerate(unittest.TestCase):
     # 4 tokens, chunk_size=4 -> 1 prefill chunk
     self.assertEqual(get_prefill_flags(list(range(4)), 4), [True, False, False])
 
-  def test_chunked_prefill_kv_cache_matches_single_chunk(self):
-    config = TransformerConfig(num_blocks=1, dim=8, hidden_dim=16, n_heads=1, n_kv_heads=1, norm_eps=1e-5,
-      vocab_size=32, head_dim=4, rope_theta=1000000, rope_dim=4, qk_norm=4, v_head_dim=4, max_context=16)
-    def model():
-      m = Transformer(config)
-      rng = np.random.RandomState(1234)
-      for t in get_state_dict(m).values():
-        t.assign(Tensor(rng.uniform(-1, 1, t.shape).astype(np.float32))).realize()
-      return m
-    def prefill(m, chunk_size):
-      gen = m.generate(list(range(1, 9)), chunk_size=chunk_size, temperature=0.0)
-      next(gen)
-      return [b.cache_kv.numpy() for b in m.blk]
-    for g, r in zip(prefill(model(), 4), prefill(model(), 8)):
-      np.testing.assert_allclose(g[:, :, :, :8, :], r[:, :, :, :8, :], atol=1e-5)
-
   def test_kv_cache_resume_matches_fresh(self):
     model = Transformer(TEST_CONFIG)
 

tinygrad/codegen/__init__.py

@@ -1,9 +1,9 @@
 from dataclasses import replace, dataclass
 import itertools, functools
 from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
-from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, NUM_CPU_THREADS, TC_SELECT, TC_OPT, TracingKey, Context, panic
+from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, TracingKey, Context, panic
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, Ops, UPat, rewrite_group, KernelInfo, ProgramInfo, GroupOp, AxisType
-from tinygrad.uop.weak import pm_lower_weak, pm_commit_weak, pm_cast_const
+from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak, pm_cast_weak
 from tinygrad.uop.render import pyrender
 from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
 from tinygrad.renderer import Renderer, Estimates
@@ -22,7 +22,7 @@ from tinygrad.codegen.opt.postrange import apply_opts
 from tinygrad.codegen.late.gater import pm_move_gates_from_index
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse, pm_reduce_unparented
 from tinygrad.schedule.multi import multi_pm
-from tinygrad.schedule.prepare import pm_mops
+from tinygrad.schedule.rangeify import pm_mops
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite
 from tinygrad.codegen.late.coalesce import memory_coalescing, pm_simplify_add_image
@@ -233,11 +233,10 @@ pm_reduce_local = pm_wmma_add+PatternMatcher([
   (UPat(Ops.SINK, name="sink"), merge_reduce_ends),
 ])+pm_clean_up_group_sink
 
-def is_shape_changing_bitcast(u:UOp): return u.op is Ops.BITCAST and u.shape != u.src[0].shape
 def maybe_load(u:UOp): return u.load() if u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL, AddrSpace.REG) else u
 pm_add_loads = PatternMatcher([
-  (UPat(GroupOp.Elementwise|{Ops.REDUCE,Ops.WMMA,Ops.STACK}, name="x"),
-   lambda x: None if is_shape_changing_bitcast(x) else x.replace(src=tuple(map(maybe_load, x.src)))),
+  # BITCAST?
+  (UPat(GroupOp.Elementwise|{Ops.REDUCE,Ops.WMMA,Ops.STACK}, name="x"), lambda x: x.replace(src=tuple([maybe_load(u) for u in x.src]))),
   (UPat(Ops.STORE, name="x"), lambda x: x.replace(src=(x.src[0], maybe_load(x.src[1]))+x.src[2:])),
 ])
 
@@ -282,6 +281,10 @@ pm_implicit_barriers = PatternMatcher([
   (UPat(Ops.END, name="end"), add_war_barrier),
 ])
 
+pm_casted_consts = PatternMatcher([
+  (UPat(Ops.CONST, dtypes.all, name="c"), lambda c: UOp.cconst(c.val, c.dtype)),
+])
+
 def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
   if DEBUG >= 5: print(pyrender(ast))
@@ -343,13 +346,11 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
 
   # extra symbolic before decomp. crashes without this?
   # NOTE: also run indexing_simplify here, while the index is still weakint and (x+y)*c -> x*c+y*c applies
-  # commit widths minted in this fixpoint before lowering inspects INDEX shapes
-  sink = graph_rewrite(sink, sym+indexing_simplify+pm_commit_weak, name="extra symbolic")
+  sink = graph_rewrite(sink, sym+indexing_simplify, name="extra symbolic")
 
-  # the boundary: required compute dtypes settle here; derivable const edges may stay bare
+  # lower index dtype
   # NOTE: we need indexing_simplify to remove the cast to long using the Invalid
-  # NOTE: symbolic must NOT be composed here -- pm_data_invalid pushes the weak result CAST into a gated WHERE, remaking the weak node, and it cycles
-  sink = graph_rewrite(sink, pm_lower_weak+indexing_simplify, name="lower all index dtypes")
+  sink = graph_rewrite(sink, symbolic_simple+pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes")
 
   # final symbolic before decomp
   sink = graph_rewrite(sink, symbolic, name="final symbolic")
@@ -373,12 +374,9 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
 
   # final rules for the renderer (without sym)
   extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([])
-  pm_final_rewrite = pm_commit_weak+pm_decomp+extra_matcher+pm_split_ends
+  pm_final_rewrite = pm_commit_weak+pm_cast_weak+pm_decomp+extra_matcher+pm_split_ends
   sink = graph_rewrite(sink, pm_final_rewrite+pm_remove_invalid, ctx=ren, name="final rewrite")
 
-  # commit every const still bare so no renderer reads one
-  sink = graph_rewrite(sink, pm_cast_const, name="cast consts")
-
   # add implicit barriers (stores/loads through LOCAL memory ordered by AFTER or across loop iterations need workgroup barriers)
   sink = graph_rewrite(sink, pm_implicit_barriers, name="add implicit barriers")
 
@@ -389,6 +387,10 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   num_params = len([x for x in sink.toposort() if x.op is Ops.PARAM and x.arg.slot != -1])
   sink = graph_rewrite(sink, pm_number_params, ctx=[num_params], name="number params with -1", walk=True)
 
+  # spell every literal as a casted const CAST(dt, CONST(value))
+  # TODO: remove once consts are always weak
+  sink = graph_rewrite(sink, pm_casted_consts, name="casted consts", walk=True)
+
   if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Output AST")
   if SPEC: type_verify(sink, spec_program)
 
@@ -457,7 +459,7 @@ pm_to_program = PatternMatcher([
   (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.LINEAR), UPat(Ops.SOURCE, name="source")), name="prg"), do_compile),
 ])
 
-@rewrite_group(name=lambda ast,renderer,ret,**_: TracingKey((k:=ret.src[0].arg).name,(k.function_name, ast, ret.key),ret=renderer), replay=True)
+@rewrite_group(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.src[0].arg.name,(ret.src[0].arg.function_name, ast), ret=renderer), replay=True)
 @Context(ALLOW_DEVICE_USAGE=0)
 def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
   """
@@ -486,14 +488,9 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
   if VIZ: graph_rewrite(prg, PatternMatcher([]), name="View Program")
   return prg
 
-# config affects generated programs and cache keys; context also carries compile-only behavior to workers
-to_program_config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32,
-                     DEFAULT_FLOAT, DEFAULT_INT, NUM_CPU_THREADS, TC_SELECT, TC_OPT)
-to_program_context = (*to_program_config, SPEC, DEBUG)
-def to_program_key(ast:UOp, renderer:Renderer) -> tuple:
-  return (ast.key, type(renderer), renderer.target, *[x.value for x in to_program_config])
-
 to_program_cache: dict[tuple, UOp] = {}
 def to_program(ast:UOp, renderer:Renderer) -> UOp:
-  if (prg:=to_program_cache.get(key:=to_program_key(ast, renderer))) is None: to_program_cache[key] = prg = do_to_program(ast, renderer)
+  config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT)
+  key = (ast.key, type(renderer), renderer.target, *[x.value for x in config])
+  if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer)
   return prg

tinygrad/codegen/decomp/dtype.py

@@ -1,9 +1,8 @@
 from dataclasses import replace
 from tinygrad.dtype import dtypes, DType, truncate
-from tinygrad.helpers import flatten, DEBUG, EMULATED_DTYPES
+from tinygrad.helpers import flatten, DEBUG, EMULATED_DTYPES, Context, SPEC
 from tinygrad.uop import GroupOp
-from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite
-from tinygrad.uop.weak import commit_weak_consts
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite, ParamArg
 from tinygrad.renderer import Renderer
 from tinygrad.codegen.decomp.transcendental import exponent_bias, shl, shr
 
@@ -26,10 +25,10 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
   match op:
     case Ops.NEG: return l2i(Ops.SUB, dt, zero, zero, *uops)
     case Ops.CAST if dt in (dtypes.long, dtypes.ulong) and uops[0].dtype not in dtypes.floats:
-      # the high word is the sign extension, and unsigned and bool sources zero extend
+      # the high word is the sign extension; bool has no sign, test the already-cast low word instead (bool < 0 would promote to weakint)
       x, lo = uops[0], uops[0].cast(l2i_dt[dt])
-      if x.dtype is dtypes.bool or x.dtype in dtypes.uints: return lo, lo.const_like(0)
-      return lo, (x < x.const_like(0)).where(lo.const_like(-1), lo.const_like(0))
+      sign = lo if x.dtype is dtypes.bool else x
+      return lo, (sign < sign.const_like(0)).where(lo.const_like(-1), lo.const_like(0))
     case Ops.CAST if dt in (dtypes.long, dtypes.ulong):
       return (lo:=uops[0].cast(l2i_dt[dt])), (uops[0] / 2**32).cast(l2i_dt[dt]) - ((uops[0] < 0) & lo.ne(0))
     case Ops.CAST if dt in dtypes.floats:
@@ -128,22 +127,19 @@ def f2f_clamp(val:UOp, dt:DType, sat=True) -> UOp:
   return val.ne(val).where(val, (val < -mx).where(-sat, (mx < val).where(sat, val)))
 
 def f2f_load(x: UOp, fr:DType, to:DType) -> UOp:
-  storage_idx = graph_rewrite(x.src[0], pm_float_decomp, ctx=(fr, to), bottom_up=True)
-  if (n:=x.max_numel()) == 1: return f2f(storage_idx.load(*x.src[1:]), fr, to)
-  return UOp(Ops.STACK, src=tuple(f2f(reindex(storage_idx, i, 1).load(*x.src[1:]), fr, to) for i in range(n)))
+  if (n:=x.max_numel()) == 1: return f2f(x.replace(dtype=f2f_dt[fr]), fr, to)
+  return UOp(Ops.STACK, src=tuple(f2f(x.replace(dtype=f2f_dt[fr], src=(reindex(x.src[0], i, 1),)), fr, to) for i in range(n)))
 
 def f2f_store(st, idx, val, fr:DType, to:DType):
   if (n:=val.max_numel()) == 1: return st.replace(src=(idx, f2f(val.bitcast(f2f_dt[to]), to, fr)))
   return UOp.group(*(st.replace(src=(reindex(idx, i, 1), f2f(val.index(i).bitcast(f2f_dt[to]), to, fr))) for i in range(n)))
 
 # tag is the 32-bit word this node becomes - (0 for the low word, 1 for the high, the dtype the consumer wants)
-pm_long_decomp: PatternMatcher = PatternMatcher([
-  # the decomp's own bottom-up rewrite can mint bare consts mid-flight: word splitting commits them at the long sibling's dtype
-  (UPat(GroupOp.All, name='x'), lambda x: commit_weak_consts(x, next((s.dtype for s in x.src if s.dtype in l2i_dt), None))),
-  (UPat(GroupOp.Defines, tuple(l2i_dt.keys()), src=(UPat.var("sz"),), name="x"), lambda x,sz:
-   UOp(x.op, src=(sz*2,), arg=replace(x.arg, dtype=l2i_dt[x.dtype]), tag=x.tag)),
+pm_long_decomp = PatternMatcher([
+  (UPat(GroupOp.Defines, src=(UPat.var("sz"),), name="x"), lambda x,sz:
+   x.replace(dtype=l2i_dt[x.dtype], arg=replace(x.arg, dtype=l2i_dt[x.dtype]), src=(sz*2,)) if x.dtype in l2i_dt else None),
   (UPat(Ops.INDEX, tuple(l2i_dt.keys()), name='x'), lambda x:
-   reindex(x, x.tag[0]).replace(tag=None) if x.tag is not None else None),
+   reindex(x, x.tag[0]).replace(dtype=x.tag[1], tag=None) if x.tag is not None else None),
   (UPat(Ops.STORE, src=(UPat.var('idx', tuple(l2i_dt.keys())), UPat.var('val')), name='st'), lambda st,idx,val:
    st.replace(src=(idx.rtag((0, dt:=l2i_dt[idx.dtype])), val.rtag((0, dt)))).group(
      st.replace(src=(idx.rtag((1, dt)), val.rtag((1, dt))))) if val.tag is None else None),
@@ -151,9 +147,6 @@ pm_long_decomp: PatternMatcher = PatternMatcher([
    split_l2i(ctx, x.op, dt:=l2i_dt[a.dtype], *flatten((s.rtag((0, dt)), s.rtag((1, dt))) for s in x.src))),
   (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda ctx,a,x:
    split_l2i(ctx, Ops.BITCAST, l2i_dt[x.dtype], a.rtag((0, dt:=l2i_dt[a.dtype])), a.rtag((1, dt)))[x.tag[0]]),
-  # a const splits by value; the general CAST arm below would drop its high word
-  (UPat(Ops.CAST, src=(UPat(Ops.CONST, name='c'),), tag={(w, dt) for w in (0, 1) for dt in l2i_dt.values()}, name='x'),
-   lambda x,c: UOp.const(truncate[x.tag[1]](c.val >> (32*x.tag[0])), x.tag[1])),
   (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a'),), name="x"), lambda ctx,a,x:
    split_l2i(ctx, x.op, x.dtype, a)[x.tag[0]] if x.tag is not None else None),
   (UPat(Ops.CAST, src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda ctx,a,x:
@@ -166,22 +159,21 @@ pm_long_decomp: PatternMatcher = PatternMatcher([
   (UPat((*(GroupOp.ALU - GroupOp.Comparison - {Ops.SHL, Ops.SHR, Ops.WHERE}), Ops.BITCAST), tuple(l2i_dt.keys()), name="x"), lambda ctx,x:
    split_l2i(ctx, x.op, l2i_dt[x.dtype], *flatten((a.rtag((0, l2i_dt[x.dtype])), a.rtag((1, l2i_dt[x.dtype]))) for a in x.src))[x.tag[0]]
    if x.tag is not None else None),
-  (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda ctx,x,idx:
-   reindex(graph_rewrite(idx, pm_long_decomp, ctx=ctx, bottom_up=True), x.tag[0]).replace(tag=None).load() if x.tag is not None else None)
+  (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx:
+   x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag[0]).replace(dtype=l2i_dt[x.dtype], tag=None),), tag=None) if x.tag is not None else None),
+  (UPat(Ops.CONST, tag={(w, dt) for w in (0, 1) for dt in l2i_dt.values()}, name='x'), lambda x:
+   UOp.const(truncate[x.tag[1]]((x.val >> 32) if x.tag[0] == 1 else (x.val & 0xFFFFFFFF)), x.tag[1]))
 ])
 
 # float decomposition patterns - ctx is (fr, to) tuple
-pm_float_decomp: PatternMatcher = PatternMatcher([
-  (UPat(GroupOp.Defines, name="x"), lambda ctx,x:
-   UOp(x.op, src=x.src, arg=replace(x.arg, dtype=f2f_dt[ctx[0]]), tag=ctx[0]) if x.dtype == ctx[0] else None),
-  # INDEX into a LOAD/STACK selects a lane of an already converted value, the load rules below own those
-  (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat(GroupOp.All-{Ops.LOAD, Ops.STACK}),), allow_any_len=True, name="x"), lambda ctx,x:
-   UOp(x.op, src=(graph_rewrite(x.src[0], pm_float_decomp, ctx=ctx, bottom_up=True), *x.src[1:]), arg=x.arg, tag=ctx[0])
-   if x.dtype == ctx[0] else None),
+pm_float_decomp = PatternMatcher([
+  (UPat((*GroupOp.Defines, Ops.INDEX, Ops.SHRINK), name="x"), lambda ctx,x:
+   x.replace(dtype=f2f_dt[ctx[0]], arg=replace(x.arg, dtype=f2f_dt[ctx[0]]) if isinstance(x.arg, ParamArg) else x.arg, tag=ctx[0])
+   if x.dtype == ctx[0] and (x.op is not Ops.INDEX or x.src[0].op not in {Ops.LOAD, Ops.STACK}) else None),
   (UPat(Ops.LOAD, dtypes.floats, name="x"), lambda ctx,x: f2f_load(x, *ctx) if x.dtype == ctx[0] else None),
   # bitcasted load should just replace load
   (UPat(Ops.BITCAST, src=(UPat(Ops.LOAD, name="ld"),), name="bc"), lambda ctx,bc,ld:
-   graph_rewrite(ld.src[0], pm_float_decomp, ctx=ctx, bottom_up=True).load(*ld.src[1:]).bitcast(bc.dtype) if ld.dtype == ctx[0] else None),
+   ld.replace(dtype=f2f_dt[ctx[0]]).bitcast(bc.dtype) if ld.dtype == ctx[0] else None),
   # bitcast from
   (UPat(Ops.BITCAST, src=(UPat.var("x", dtypes.floats),), name="bc"), lambda ctx,bc,x:
    bc.replace(src=(f2f(x.bitcast(f2f_dt[ctx[1]]), ctx[1], ctx[0]),)) if x.dtype == ctx[1] and bc.dtype.bitsize == ctx[0].bitsize else None),
@@ -190,21 +182,26 @@ pm_float_decomp: PatternMatcher = PatternMatcher([
    f2f(x.bitcast(f2f_dt[ctx[0]]), ctx[0], ctx[1]) if bc.dtype == ctx[0] else None),
   (UPat(Ops.CAST, dtypes.floats, src=(UPat.var("val"),), name="x"), lambda ctx,x,val:
    f2f_clamp(val.cast(ctx[1]), ctx[0]) if x.dtype == ctx[0] else None),
+  # a CONST has no srcs to cast, it restates its value at the emulating dtype
+  (UPat(Ops.CONST, dtypes.floats, name="x"), lambda ctx,x: UOp.const(x.val, ctx[1]) if x.dtype == ctx[0] else None),
   (UPat(GroupOp.All-GroupOp.Defines-{Ops.CAST, Ops.BITCAST, Ops.CONST}, dtypes.floats, name="x"), lambda ctx,x:
-   UOp(x.op, src=tuple(s.cast(ctx[1]) if s.dtype == ctx[0] else s for s in x.src), arg=x.arg, tag=x.tag) if x.dtype == ctx[0] else None),
+   x.replace(dtype=ctx[1], src=tuple(s.cast(ctx[1]) if s.dtype == ctx[0] else s for s in x.src))
+   if x.dtype == ctx[0] else None),
   (UPat(Ops.STORE, src=(UPat.var("idx"), UPat(Ops.BITCAST, dtypes.floats, name="val")), name='st'), lambda ctx,st,idx,val:
-   st.replace(src=(idx, val.src[0].bitcast(f2f_dt[ctx[0]]))) if val.dtype == ctx[0] and idx.tag == ctx[0] else None),
-  (UPat(Ops.STORE, src=(UPat.var("idx").or_casted(), UPat.var("val", dtypes.floats)), name='st'), lambda ctx,st,idx,val:
-   f2f_store(st, idx, val, *ctx) if val.dtype == ctx[1] and idx.tag == ctx[0] else None),
+   st.replace(src=(idx, val.replace(dtype=f2f_dt[ctx[0]]))) if val.dtype == ctx[0] and idx.tag == ctx[0] else None),
+  (UPat(Ops.STORE, src=(UPat.var("idx"), UPat.var("val", dtypes.floats)), name='st'), lambda ctx,st,idx,val:
+   f2f_store(st, idx, val, *ctx) if val.dtype == ctx[1] and (idx:=idx.src[0] if idx.op == Ops.CAST else idx).tag == ctx[0] else None),
 ])
 
 def do_dtype_decomps(sink:UOp, ctx:tuple[set[DType], Renderer]) -> UOp:
   def _should_emulate(dt): return dt in EMULATED_DTYPES.tolist(dtypes) or dt not in ctx[1].supported_dtypes()
-  for fr in sorted(filter(_should_emulate, ctx[0])):
-    to = dtypes.int if fr == dtypes.long else dtypes.half if not _should_emulate(dtypes.half) and fr in dtypes.fp8s else dtypes.float
-    if DEBUG >= 2: print(f"emulating {fr} as {to}")
-    pm = pm_float_decomp if fr in dtypes.floats else pm_long_decomp
-    sink = graph_rewrite(sink, pm, name=f"decomp {fr} -> {to}", ctx={} if pm is pm_long_decomp else (fr, to), bottom_up=True)
+  # NOTE: dtype decomp creates intermediate UOps that don't follow the spec (e.g. half LOAD on ushort BUFFER)
+  with Context(SPEC=min(SPEC.value, 1)):
+    for fr in sorted(filter(_should_emulate, ctx[0])):
+      to = dtypes.int if fr == dtypes.long else dtypes.half if not _should_emulate(dtypes.half) and fr in dtypes.fp8s else dtypes.float
+      if DEBUG >= 2: print(f"emulating {fr} as {to}")
+      pm = pm_float_decomp if fr in dtypes.floats else pm_long_decomp
+      sink = graph_rewrite(sink, pm, name=f"decomp {fr} -> {to}", ctx={} if pm is pm_long_decomp else (fr, to), bottom_up=True)
   ctx[0].clear()
   return sink
 

tinygrad/codegen/decomp/op.py

@@ -75,11 +75,7 @@ powers_of_two: dict[int, int] = {2**i:i for i in range(64)}
 @functools.cache
 def get_simplifying_rewrite_patterns(ops:tuple[Ops, ...]) -> PatternMatcher:
   # these are rewrites that make things simpler
-  pat: list[tuple[UPat, Callable]] = []
-  # FLOORDIV by 2**y -> x >> y (an arithmetic shift is exactly floor division for any sign); fires before floordiv_to_idiv
-  if Ops.SHR in ops: pat.append((UPat.var("x", dtypes.ints)//UPat.cvar("c"),
-    lambda x,c: x >> v if (v:=powers_of_two.get(c.val, 0)) else None))
-  pat.append((UPat.var("a")//UPat.var("b"), floordiv_to_idiv))
+  pat: list[tuple[UPat, Callable]] = [(UPat.var("a")//UPat.var("b"), floordiv_to_idiv)]
   # FLOORMOD by 2**y -> x & (2**y-1) (correct floor mod for any sign in two's complement); fires before floormod_to_mod
   if Ops.AND in ops: pat.append((UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.val-1) if c.val in powers_of_two else None))
   pat.append((UPat.var("a")%UPat.var("b"), floormod_to_mod))
@@ -132,6 +128,6 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa
     if Ops.SHL in ops: pat += [(UPat.var('x').alu(Ops.SHL, UPat.cvar('n'))+UPat.var('c'), lambda x,n,c: x.alu(Ops.MULACC, x.const_like(1<<n.val), c))]
   # some backends emit FDIV for RECIP, in that case: a*(1/b) -> a/b
   if Ops.FDIV in ops:
-    pat += [(UPat.var("x").reciprocal(), lambda x: UOp.const(1.0).alu(Ops.FDIV, x))]
-    pat += [(UPat.var("a") * UPat(Ops.FDIV, dtypes.floats, src=(UPat.const(1), UPat.var("b"))), lambda a,b: a.alu(Ops.FDIV, b))]
+    pat += [(UPat.var("x").reciprocal(), lambda x: x.const_like(1).alu(Ops.FDIV, x))]
+    pat += [(UPat.var("a", dtypes.floats) * UPat(Ops.FDIV, dtypes.floats, src=(UPat.const(1), UPat.var("b"))), lambda a,b: a.alu(Ops.FDIV, b))]
   return PatternMatcher(pat)

tinygrad/codegen/decomp/transcendental.py

@@ -90,8 +90,8 @@ def payne_hanek_reduction(d:UOp) -> tuple[UOp, UOp]:
     if count+offset < len(two_over_pi_f) - 1:
       an = i.ne(count).where(_take(an, offset, count=count+1), an.const_like(two_over_pi_f[count+offset]))
     return an
-  def _shl_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) << y.cast(dtypes.uint64)).cast(dtypes.uint32)
-  def _shr_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) >> y.cast(dtypes.uint64)).cast(dtypes.uint32)
+  def _shl_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) * pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32)
+  def _shr_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) // pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32)
 
   a = [_take(UOp.const(0, dtypes.uint32), i) for i in range(4)]
   #  (two_over_pi_f[Int(i) + n] << e) | (two_over_pi_f[Int(i) + n+1] >> (nbits - e))

tinygrad/codegen/opt/postrange.py

@@ -193,7 +193,7 @@ class Scheduler:
       for b in self.bufs:
         if rng in (i:=b.src[1].get_idx()).backward_slice_with_self:
           nb = b.replace(src=(b.src[0], i.valid(valid&b.src[1].get_valid())))
-          replaces[b] = nb if b in store_targets else valid.where(nb, UOp.const(Invalid))
+          replaces[b] = nb if b in store_targets else valid.where(nb, UOp.const(Invalid, b.dtype))
       self.ast = self.ast.substitute(replaces, f"padto {rng.arg[:-1]} {opt.arg}")
     elif opt.op is OptOps.SWAP:
       try:

tinygrad/codegen/opt/search.py

@@ -1,13 +1,12 @@
-import math, time, traceback, signal
+import math, time, multiprocessing, traceback, signal, atexit
 from dataclasses import replace
 from tinygrad.uop.ops import sym_infer, AxisType, UOp, Ops
 from tinygrad.uop.render import pyrender
 from tinygrad.device import Device, Buffer
-from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, colored, time_to_str
+from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str
 from tinygrad.helpers import IGNORE_BEAM_CACHE
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.engine.realize import time_call
-from tinygrad.engine.worker import get_worker_pool, terminate_worker_pool
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt.postrange import Scheduler
 
@@ -79,6 +78,11 @@ def _try_compile(x:tuple[int,Scheduler]) -> tuple[int, tuple[UOp, float]|None]:
     if hasattr(signal, "alarm"): signal.alarm(0)
   return x[0], ret
 
+# workers should not open devices and should ignore ctrl c and should not launch VIZ
+def _init_worker():
+  Context(ALLOW_DEVICE_USAGE=0, VIZ=0, TRACK_MATCH_STATS=0).__enter__()
+  signal.signal(signal.SIGINT, signal.SIG_IGN)
+
 def _ensure_buffer_alloc(bufs:list[Buffer]) -> list[Buffer]: return [buf.ensure_allocated() if buf is not None else buf for buf in bufs]
 
 # *** external API ***
@@ -107,8 +111,9 @@ def get_kernel_actions(s:Scheduler, include_0=True, max_up:int|None=None) -> dic
     except KernelOptError: pass
   return acted
 
-BEAM_DEBUG = getenv("BEAM_DEBUG")
+beam_pool, BEAM_DEBUG = None, getenv("BEAM_DEBUG")
 def beam_search(s:Scheduler, rawbufs:list[Buffer], var_vals:dict[str,int], amt:int, allow_test_size=True, disable_cache=IGNORE_BEAM_CACHE.value):
+  global beam_pool
   key = {"ast": s.ast.key, "amt": amt, "allow_test_size": allow_test_size, "device": s.ren.target.device, "suffix": s.ren.suffix}
   if not disable_cache and CACHELEVEL >= 1 and (val:=diskcache_get("beam_search", key)) is not None:
     ret = s.copy()
@@ -118,7 +123,11 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], var_vals:dict[str,int], amt:i
   beam: list[tuple[Scheduler, float]] = [(s, float("inf"))]
   seen_libs = set()
 
-  pool = get_worker_pool()
+  default_parallel = multiprocessing.cpu_count() if s.ren.target.device in {"CUDA", "AMD", "NV", "METAL", "HIP"} else 0
+  if beam_pool is None and (workers := getenv("PARALLEL", default_parallel)):
+    beam_pool = multiprocessing.get_context("spawn").Pool(workers, _init_worker, (), getenv("BEAM_MAX_TASKS_PER_CHILD", 16))
+    @atexit.register
+    def close_pool(): beam_pool.close()
 
   min_progress = getenv("BEAM_MIN_PROGRESS", 0.01)/1e6
   if BEAM_DEBUG:
@@ -134,7 +143,7 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], var_vals:dict[str,int], amt:i
       candidates: list[Scheduler] = flatten([get_kernel_actions(si, include_0=False).values() for si,_ in beam])
       timed: list[tuple[Scheduler, float]] = []
       least_compute_ops = math.inf
-      for i, proc in ((map if pool is None else pool.imap_unordered)(_try_compile, enumerate(candidates))):
+      for i, proc in ((map if beam_pool is None else beam_pool.imap_unordered)(_try_compile, enumerate(candidates))):
         if proc is None: continue
         prg, compile_et = proc
         if (lib:=prg.src[3].arg) in seen_libs: continue
@@ -170,7 +179,7 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], var_vals:dict[str,int], amt:i
         print(f"\r{time.perf_counter() - st:7.2f}s:", colored(time_to_str(beam[0][1], w=12), "green" if exiting else None),
               f"from {len(candidates):3d} -> {len(opts):3d} actions\033[K", beam[0][0].colored_shape())
   except KeyboardInterrupt as e:
-    terminate_worker_pool()
+    if beam_pool is not None: beam_pool.terminate()
     raise e
 
   if CACHELEVEL >= 1: diskcache_put("beam_search", key, beam[0][0].applied_opts)

tinygrad/codegen/opt/tc.py

@@ -153,6 +153,7 @@ pm_validate_wmma_rdna3 = PatternMatcher([
 
 pm_validate_wmma_rdna4 = PatternMatcher([
   (UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: x.replace(
+    dtype=dtypes.uint16,
     src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16)))
       .bitcast(dtypes.bfloat16) if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None),
   (UPat(Ops.WMMA, name="x", dtype=dtypes.float),

tinygrad/device.py

@@ -66,10 +66,10 @@ def canonicalize_device(device:str|tuple|list|None) -> str|tuple[str, ...]:
 class ProfileDeviceEvent(ProfileEvent): device:str; tdiff:decimal.Decimal=decimal.Decimal(0); props:dict[str,Any]|None=None # noqa: E702
 
 @dataclass(frozen=True)
-class ProfileProgramEvent(ProfileEvent): device:str; name:str; lib:bytes|None; base:int|None; tag:int|None=None; profile_key:bytes|None=None # noqa: E702
+class ProfileProgramEvent(ProfileEvent): device:str; name:str; lib:bytes|None; base:int|None; tag:int|None=None # noqa: E702
 
 @dataclass(frozen=True)
-class ProfileGraphEntry: device:str; name:str|TracingKey; st_id:int; en_id:int; profile_key:bytes|None=None # noqa: E702
+class ProfileGraphEntry: device:str; name:str|TracingKey; st_id:int; en_id:int # noqa: E702
 
 @dataclass(frozen=True)
 class ProfileGraphEvent(ProfileEvent): ents:list[ProfileGraphEntry]; deps:list[list[int]]; sigs:list[decimal.Decimal] # noqa: E702
@@ -83,7 +83,6 @@ class BufferSpec:
   cpu_access: bool = False
   host: bool = False
   nolru: bool = False
-  zero: bool = False
   external_ptr: int|None = None
 
 class MultiBuffer:
@@ -266,7 +265,7 @@ class LRUAllocator(Allocator, Generic[DeviceType]):
       for opaque in opaques: super().free(opaque, sz, options)
       opaques.clear()
   def free(self, opaque:Any, size:int, options:BufferSpec|None=None):
-    if LRU and (options is None or (not (options.nolru or options.zero) and options.external_ptr is None)): self.cache[(size, options)].append(opaque)
+    if LRU and (options is None or (not options.nolru and options.external_ptr is None)): self.cache[(size, options)].append(opaque)
     else: super().free(opaque, size, options)
 
 class DepsTracker:
@@ -327,7 +326,6 @@ class TinyELF:
   target: Target
   # tuple of (name, slot, dtype, shape)
   signature: tuple[tuple[str|None, int, DType, tuple], ...]
-  profile_key: bytes|None = None
 
   @staticmethod
   def iter_sig(signature:tuple[tuple[str|None, int, DType, tuple], ...], offset:int=0) -> Generator[tuple[int, DType], None, None]:

tinygrad/engine/jit.py

@@ -211,7 +211,7 @@ def _prepare_jit_inputs(args, kwargs):
   # collect buffer UOps (including MultiBuffer)
   input_buf_uops: list[UOp] = [u.base for u in input_uops if u.base.realized is not None]
   if len(set(input_buf_uops)) != len(input_buf_uops): raise JitError("duplicate inputs to JIT")
-  inputs = [(*(u.substitute({u.base:UOp(Ops.NOOP)}, extra_pm=mop_cleanup).unbind_all()), u.dtype, u.device) for u in input_uops]
+  inputs = [(*(u.substitute({u.base:UOp(Ops.NOOP, u.base.dtype)}, extra_pm=mop_cleanup).unbind_all()), u.dtype, u.device) for u in input_uops]
   _var_vals = merge_dicts([x[1] for x in inputs] + [dict(v.unbind() for v in (args + tuple(kwargs.values())) if isinstance(v, UOp))])
   var_vals = {k.expr:v for k,v in _var_vals.items()}
   expected_input_info = [(x[0], tuple(sorted(x[1].keys(), key=lambda v: v.expr)), x[2], x[3]) for x in inputs]

tinygrad/engine/realize.py

@@ -2,15 +2,14 @@ from __future__ import annotations
 from typing import cast, Iterator, Any, Sequence
 import random, itertools, math, weakref, array, decimal
 from dataclasses import dataclass, replace, field
-from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansipad, all_int, prod, flatten, Context, getenv, to_tuple, tqdm
+from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansipad, all_int, prod, flatten, Context, getenv, to_tuple
 from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us
-from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite, ProgramInfo
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite
 from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
 from tinygrad.dtype import dtypes
-from tinygrad.renderer import Estimates, Renderer
-from tinygrad.codegen import to_program, to_program_cache, to_program_key, to_program_context
+from tinygrad.renderer import Estimates
+from tinygrad.codegen import to_program
 from tinygrad.codegen.opt.postrange import args_from_ast
-from tinygrad.engine.worker import get_worker_pool, terminate_worker_pool
 
 # **************** Helpers ****************
 
@@ -26,12 +25,11 @@ def get_call_outs_ins(call:UOp) -> tuple[tuple[int, ...], tuple[int, ...]]:
   if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return (0,), tuple(range(1, len(get_call_arg_uops(call))))
   return (), ()
 
-def get_call_kernels(call:UOp) -> list[tuple[str, UOp, tuple[str, Estimates, bytes]|None]]:
-  if (ast:=call.src[0]).op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq":
-    return [(d, call, (name, estimates, profile_key)) for devices,name,estimates,_,profile_key in call.arg.aux.kernels for d in devices]
-  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return [(to_tuple(ast.device)[0], call, None)]
+def get_call_kernels(call:UOp) -> list[tuple[str, UOp]]:
+  if (ast:=call.src[0]).op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq": return [(d, k) for devs, k, _ in call.arg.aux.kernels for d in devs]
+  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return [(to_tuple(ast.device)[0], call)]
   if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "validate": return []
-  return [(d, call, None) for d in to_tuple(call.src[1].device)]
+  return [(d, call) for d in to_tuple(call.src[1].device)]
 
 def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|None=None) -> str:
   def _uop_sz_to_str(uop:UOp) -> str: return size_to_str(sym_infer(prod(uop.shape) * uop.dtype.itemsize, var_vals or {}))
@@ -67,25 +65,24 @@ def track_stats(ctx:ExecContext, call:UOp, st:decimal.Decimal, ets:list[float|No
   if DEBUG < 2 and not PROFILE: return
 
   kernels = get_call_kernels(call) # everything below is the per kernel display: exec events for the profiler and DEBUG=2 lines
-  args = resolve_params(call, ctx.input_uops) if kernels and kernels[0][2] is None else []
+  args = resolve_params(call, ctx.input_uops) if kernels and kernels[0][1] is call else []
   lanes = list(unwrap_multi(call, [args[g] for g in call.src[0].arg.globals] if call.src[0].op is Ops.PROGRAM else args)) if args else []
-  for i, (device, kcall, stats) in enumerate(kernels):
+  for i, (device, kcall) in enumerate(kernels):
     et, bufs = ets[i] if i < len(ets) else None, lanes[i][0] if i < len(lanes) else []
-    display_name = get_call_name(kcall, bufs, ctx.var_vals) if stats is None else stats[0]
     if PROFILE: # backdate the event to the start of the call, the viz matches a device range with the exec event before it
       outputs, inputs = get_call_outs_ins(kcall)
       cpu_events.append(ProfilePointEvent(device, "exec", len(cpu_events), {"var_vals": ctx.var_vals,
-        "bufs": [b.trace_num for b in bufs], "name": display_name, "outputs": outputs, "inputs": inputs}, ts=st))
+        "bufs": [b.trace_num for b in bufs], "name": get_call_name(kcall, bufs, ctx.var_vals), "outputs": outputs, "inputs": inputs}, ts=st))
     if DEBUG < 2 or not ctx.update_stats: continue
     if et is None:
       Device[device].synchronize()
       et, st = float(perf_counter_us() - st)*1e-6, perf_counter_us()
       GlobalCounters.time_sum_s += et
 
-    estimates = estimate_uop(kcall) if stats is None else stats[1]
+    estimates = estimate_uop(kcall)
+    display_name = get_call_name(kcall, bufs, ctx.var_vals)
     op_est, mem_est, lds_est = (sym_infer(x, ctx.var_vals) for x in (estimates.ops, estimates.mem, estimates.lds))
-    key = kcall.src[0].key if stats is None else stats[2]
-    header_color = 'magenta' if ctx.jit else ('green' if key not in first_run_cache else None)
+    header_color = 'magenta' if ctx.jit else ('green' if kcall.src[0].key not in first_run_cache else None)
     ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
     flops, membw, ldsbw = op_est/(et or 1e-20), mem_est/(et or 1e-20), lds_est/(et or 1e-20)
     flops_str = f"{flops*1e-9:7.0f} GFLOPS" if flops < 1e14 else colored(f"{flops*1e-12:7.0f} TFLOPS", 'green')
@@ -94,7 +91,7 @@ def track_stats(ctx:ExecContext, call:UOp, st:decimal.Decimal, ets:list[float|No
     print(f"{colored(f'*** {device[:7]:7s} {GlobalCounters.kernel_count:4d}', header_color)}"+
       f" {ansipad(display_name, 46)} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:6.2f} GB"+
       ("" if et is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})"))
-    first_run_cache.add(key)
+    first_run_cache.add(kcall.src[0].key)
 
 local_size_cache: dict[bytes, tuple[int, ...]] = {}
 def optimize_local_size(call:UOp, prg:UOp) -> UOp|None:
@@ -223,14 +220,13 @@ def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
     call = call.substitute({call.src[1+info.inputs]: UOp.mstack(*tables)})
   exec_kernel(replace(ctx, var_vals={**ctx.var_vals, "hcq_inputs_ptr": dev.rt_buffer()._buf.va_addr + base}), call, ast)
 
-  def _prof_tm(device:str, name:str, prof:tuple[int, ...], profile_key:bytes) -> float|None:
-    (d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, name, prof[0], prof[1], profile_key)
+  def _prof_tm(device:str, stat_call:UOp, prof:tuple[int, ...]) -> float|None:
+    (d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, stat_call.arg.name, *prof)
     if not ctx.wait: return None
     d.synchronize(timeout=ctx.timeout)
     st, en = (d.signal(x)._buf.cpu_view().view(fmt='Q')[0] for x in prof)
     return float(en-st)/d.timestamp_divider/1e6
-  return [_prof_tm(device, name, prof, profile_key) for devices,name,_,prof,profile_key in info.kernels
-          if prof for device in devices] if PROFILE or ctx.wait else []
+  return [_prof_tm(device, k, prof) for devices, k, prof in info.kernels if prof for device in devices] if PROFILE or ctx.wait else []
 
 # flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
 pm_flatten_linear = PatternMatcher([
@@ -251,44 +247,10 @@ pm_beam = PatternMatcher([
    lambda ctx,call,sink: call.replace(src=(sink.replace(arg=replace(sink.arg, beam=ctx)), *call.src[1:])) if sink.arg.beam == 0 else None),
 ])
 
-# **************** parallel lowering + compilation ****************
-
-def _compile_kernel(x:tuple[int, tuple[UOp, Renderer], dict]) -> tuple[int, UOp]:
-  with Context(**x[2]): return x[0], to_program(*x[1])
-
-def _get_call_to_compile(c:UOp) -> tuple[UOp, Renderer]|None:
-  ast = a0.src[0] if (a0:=c.src[0]).op is Ops.CUSTOM_FUNCTION and a0.arg == "hcq" else a0
-  # a PROGRAM with a ProgramInfo and a BINARY is already compiled
-  if ast.op is Ops.SINK or (ast.op is Ops.PROGRAM and not (isinstance(ast.arg, ProgramInfo) and ast.src[-1].op is Ops.BINARY)):
-    return ast, Device[c.device if isinstance(c.device, str) else c.device[0]].renderer
-  return None
-
-def lower_and_compile(linear:UOp) -> UOp:
-  # collect the kernels to lower and compile, deduped by their compile cache key
-  if not len(ar:={c: a for c in linear.toposort() if c.op is Ops.CALL and (a:=_get_call_to_compile(c)) is not None}): return linear
-
-  # lower and compile what's not cached, in parallel if there's a worker pool
-  keys = {c: to_program_key(*a) for c, a in ar.items()}
-  todo = list({keys[c]: a for c, a in ar.items() if keys[c] not in to_program_cache}.items())
-  if len(todo):
-    # kernels that beam search must compile in the parent, beam needs device access to time candidates
-
-    pool = None if len(todo) == 1 or any(getattr(c.src[0].arg, "beam", 0) for c in ar) else get_worker_pool()
-    ctx = {v.key: v.value for v in to_program_context}
-    tasks = ((i, ast_ren, ctx) for i, (_, ast_ren) in enumerate(todo))
-    try:
-      with tqdm(total=len(todo), desc="compiling", disable=DEBUG<1) as pbar:
-        for i, prg in (map if pool is None else pool.imap_unordered)(_compile_kernel, tasks):
-          pbar.set_description(f"compiling {ansipad(prg.src[0].arg.name, 40)}")
-          to_program_cache[todo[i][0]] = prg
-          pbar.update(1)
-    except KeyboardInterrupt:
-      if pool is not None: terminate_worker_pool()
-      raise
-
-  # swap the compiled PROGRAMs into the calls
-  return linear.substitute({c: c.replace(src=(c.src[0].substitute({a[0]: to_program_cache[keys[c]]}), *c.src[1:])) for c, a in ar.items()},
-                           name="precompile kernels")
+pm_compile = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.PROGRAM), name="ast"),), name="call", allow_any_len=True), lambda call,ast:
+    call.replace(src=(to_program(ast, Device[call.device if isinstance(call.device, str) else call.device[0]].renderer), *call.src[1:]))),
+])
 
 pm_optimize_local_size = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), optimize_local_size),
@@ -308,7 +270,7 @@ if getenv("HCQ2"): from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_li
 def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, profile:bool|None=None) -> UOp:
   if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
   if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
-  linear = lower_and_compile(linear)
+  linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True)
   linear = graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
   if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
   return linear

tinygrad/engine/worker.py

@@ -1,49 +0,0 @@
-import multiprocessing, atexit, signal, sys, threading, contextlib
-from multiprocessing.context import SpawnContext, SpawnProcess
-from tinygrad.helpers import Context, getenv, PARALLEL
-
-# generic pool of worker processes for parallel compilation, shared by kernel lowering and BEAM search
-
-# workers should not open devices and should ignore ctrl c and should not launch VIZ
-def _init_worker():
-  Context(ALLOW_DEVICE_USAGE=0, VIZ=0, TRACK_MATCH_STATS=0).__enter__()
-  signal.signal(signal.SIGINT, signal.SIG_IGN)
-
-# spawn normally reimports the user's __main__ before _init_worker. This replays top-level code and can recursively create pools. There is no public
-# multiprocessing switch to skip that import, so hide the two attributes used to locate __main__ while each worker (including replacements) starts.
-_spawn_lock, _missing = threading.Lock(), object()
-@contextlib.contextmanager
-def _without_main():
-  main = sys.modules.get("__main__")
-  if main is None:
-    yield
-    return
-  with _spawn_lock:
-    saved = {name:getattr(main, name, _missing) for name in ("__file__", "__spec__")}
-    try:
-      for name in saved: setattr(main, name, None)
-      yield
-    finally:
-      for name,value in saved.items(): delattr(main, name) if value is _missing else setattr(main, name, value)
-
-class _WorkerProcess(SpawnProcess):
-  @staticmethod
-  def _Popen(process_obj):
-    with _without_main(): return SpawnProcess._Popen(process_obj)
-
-class _WorkerContext(SpawnContext): Process = _WorkerProcess
-
-worker_pool = None
-def get_worker_pool():
-  global worker_pool
-  if multiprocessing.current_process().daemon or PARALLEL == 0: return None
-  if worker_pool is None:
-    worker_pool = _WorkerContext().Pool(PARALLEL.value, _init_worker, (), getenv("BEAM_MAX_TASKS_PER_CHILD", 16))
-    @atexit.register
-    def close_pool(pool=worker_pool): pool.close()
-  return worker_pool
-
-def terminate_worker_pool():
-  global worker_pool
-  if worker_pool is not None: worker_pool.terminate()
-  worker_pool = None

tinygrad/helpers.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 import time
 START_TIME = time.perf_counter()
-import os, functools, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc, threading
+import os, functools, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc
 from collections import defaultdict
 import shutil, math, types, copyreg, inspect, importlib, decimal, itertools, difflib
 from dataclasses import dataclass, field, replace
@@ -364,8 +364,7 @@ class TracingKey:
 class ProfileEvent: pass
 
 @dataclass
-class ProfileRangeEvent(ProfileEvent):
-  device:str; name:str|TracingKey; st:decimal.Decimal; en:decimal.Decimal|None=None; profile_key:bytes|None=None # noqa: E702
+class ProfileRangeEvent(ProfileEvent): device:str; name:str|TracingKey; st:decimal.Decimal; en:decimal.Decimal|None=None # noqa: E702
 
 @dataclass(frozen=True)
 class ProfilePointEvent(ProfileEvent):
@@ -373,8 +372,8 @@ class ProfilePointEvent(ProfileEvent):
 
 cpu_events:list[ProfileEvent] = []
 @contextlib.contextmanager
-def cpu_profile(name:str|TracingKey, device="TINY", display=True, profile_key:bytes|None=None) -> Generator[ProfileRangeEvent, None, None]:
-  res = ProfileRangeEvent(device, name, perf_counter_us(), profile_key=profile_key)
+def cpu_profile(name:str|TracingKey, device="TINY", display=True) -> Generator[ProfileRangeEvent, None, None]:
+  res = ProfileRangeEvent(device, name, perf_counter_us())
   try: yield res
   finally:
     res.en = perf_counter_us()
@@ -398,17 +397,18 @@ cache_dir: str = os.path.join(getenv("XDG_CACHE_HOME", os.path.expanduser("~/Lib
 CACHEDB: str = getenv("CACHEDB", os.path.abspath(os.path.join(cache_dir, "cache.db")))
 
 VERSION = 22
-_db_connection = threading.local()
+_db_connection = None
 def db_connection():
-  if (conn:=getattr(_db_connection, "conn", None)) is None:
+  global _db_connection
+  if _db_connection is None:
     os.makedirs(CACHEDB.rsplit(os.sep, 1)[0], exist_ok=True)
-    conn = _db_connection.conn = sqlite3.connect(CACHEDB, timeout=60, isolation_level="IMMEDIATE")
+    _db_connection = sqlite3.connect(CACHEDB, timeout=60, isolation_level="IMMEDIATE")
     # another connection has set it already or is in the process of setting it
     # that connection will lock the database
-    with contextlib.suppress(sqlite3.OperationalError): conn.execute("PRAGMA journal_mode=WAL").fetchone()
-    conn.execute("PRAGMA synchronous=NORMAL")
-    if DEBUG >= 8: conn.set_trace_callback(print)
-  return conn
+    with contextlib.suppress(sqlite3.OperationalError): _db_connection.execute("PRAGMA journal_mode=WAL").fetchone()
+    _db_connection.execute("PRAGMA synchronous=NORMAL")
+    if DEBUG >= 8: _db_connection.set_trace_callback(print)
+  return _db_connection
 
 def diskcache_clear():
   cur = db_connection().cursor()
@@ -464,18 +464,16 @@ def _ensure_downloads_dir() -> pathlib.Path:
   return pathlib.Path(cache_dir) / "downloads"
 
 def fetch(url:str, name:pathlib.Path|str|None=None, subdir:str|None=None, gunzip:bool=False, allow_caching=not getenv("DISABLE_HTTP_CACHE"),
-          headers:dict[str, str]={}, sha256:str|None=None, extract:bool=False) -> pathlib.Path:
+          headers:dict[str, str]={}, sha256:str|None=None) -> pathlib.Path:
   import urllib.request
   if url.startswith(("/", ".")): return pathlib.Path(url)
   if name is not None and (isinstance(name, pathlib.Path) or '/' in name): fp = pathlib.Path(name)
   else:
     hh = "_"+hashlib.md5(("\n".join(f"{k.strip()}:{v.strip()}" for k,v in sorted(headers.items()))).encode("utf-8")).hexdigest() if headers else ""
     fp = _ensure_downloads_dir() / (subdir or "") / ((name or hashlib.md5(url.encode('utf-8')).hexdigest()) + hh + (".gunzip" if gunzip else ""))
-  extract_dir = fp.parent / f"{fp.name}.extract"
   if not fp.is_file() or not allow_caching or (sha256 and hashlib.sha256(fp.read_bytes()).hexdigest() != sha256):
-    if extract: shutil.rmtree(extract_dir, ignore_errors=True)
     (_dir := fp.parent).mkdir(parents=True, exist_ok=True)
-    with urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": "tinygrad 0.14.0", **headers}), timeout=10) as r:
+    with urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": "tinygrad 0.13.0", **headers}), timeout=10) as r:
       assert r.status in {200, 206}, r.status
       length = int(r.headers.get('content-length', 0)) if not gunzip else None
       readfile = gzip.GzipFile(fileobj=r) if gunzip else r
@@ -490,17 +488,6 @@ def fetch(url:str, name:pathlib.Path|str|None=None, subdir:str|None=None, gunzip
         pathlib.Path(f.name).rename(fp)
       progress_bar.update(close=True)
       if length and (file_size:=os.stat(fp).st_size) < length: raise RuntimeError(f"fetch size incomplete, {file_size} < {length}")
-  if extract:
-    if not extract_dir.is_dir():
-      import tarfile
-      tmpdir = tempfile.mkdtemp(dir=fp.parent)
-      try:
-        with tarfile.open(fp) as t: t.extractall(tmpdir, filter="data")
-        try: os.rename(tmpdir, extract_dir)  # rename is atomic, so concurrent fetches can't see a partial extraction
-        except OSError:
-          if not extract_dir.is_dir(): raise
-      finally: shutil.rmtree(tmpdir, ignore_errors=True)
-    return extract_dir
   return fp
 
 def fetch_fw(path:str, name:str, sha256:str) -> bytes:

tinygrad/llm/cli.py

@@ -88,8 +88,6 @@ models = {
   "qwen3.5:9b": "https://huggingface.co/unsloth/Qwen3.5-9B-GGUF/resolve/main/Qwen3.5-9B-Q4_K_M.gguf",
   "qwen3.6:27b": "https://huggingface.co/unsloth/Qwen3.6-27B-GGUF/resolve/main/Qwen3.6-27B-Q4_K_M.gguf",
   "qwen3.6:35b-a3b": "https://huggingface.co/unsloth/Qwen3.6-35B-A3B-GGUF/resolve/main/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf",
-  # pinned to the last revision with the plain IQ4_XS quant: the UD replacement uses Q3_K tensors the loader doesn't support
-  "qwen3.8:27b": "https://huggingface.co/unsloth/Qwen3.8-27B-GGUF/resolve/b62a80264f8b0c1bb849ee1c9c487415ebeca194/Qwen3.8-27B-IQ4_XS.gguf",
   "olmoe": "https://huggingface.co/allenai/OLMoE-1B-7B-0924-Instruct-GGUF/resolve/main/olmoe-1b-7b-0924-instruct-q4_k_m.gguf",
   "moonlight": "https://huggingface.co/gabriellarson/Moonlight-16B-A3B-Instruct-GGUF/resolve/main/Moonlight-16B-A3B-Instruct-Q4_K_M.gguf",
   "glm-4.7-flash": "https://huggingface.co/unsloth/GLM-4.7-Flash-GGUF/resolve/main/GLM-4.7-Flash-Q4_K_M.gguf",
@@ -141,8 +139,7 @@ def main():
   args = parser.parse_args()
 
   # load the model
-  with Context(DEBUG=max(DEBUG.value, 2 if args.serve else 0)):
-    model, kv = Transformer.from_gguf(fetch(models.get(args.model, args.model)), args.max_context)
+  model, kv = Transformer.from_gguf(fetch(models.get(args.model, args.model)), args.max_context)
   model_name = kv.get('general.name') or kv.get('general.basename') or args.model
   file_sizes = [y.nbytes() for y in UOp.sink(*[x.uop for x in nn.state.get_parameters(model)]).toposort() if y.op is Ops.BUFFER]
   print(f"using model \"{model_name}\" with {sum(file_sizes):,} bytes and {sum(x.numel() for x in nn.state.get_parameters(model)):,} params, "

tinygrad/llm/kernels/amd.py

@@ -1,499 +0,0 @@
-from __future__ import annotations
-import functools, math
-from typing import Callable, cast
-from tinygrad import Tensor, UOp, nn, Device, Context
-from tinygrad.device import Buffer
-from tinygrad.dtype import AddrSpace, dtypes
-from tinygrad.helpers import prod
-from tinygrad.uop.ops import AxisType, KernelInfo, Ops, resolve
-
-BLOCK_M, BLOCK_N, DECODE_HEAD_TILE, WARP_SIZE = 32, 32, 8, 32
-WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
-WAVES_M, WAVES_N, LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 2, 2, 16
-WMMA_ACC, THREADS_PER_BLOCK = WMMA_M // LANES_PER_WAVE_M, WARP_SIZE * WAVES_M * WAVES_N
-LDS_PAD, WMMA_ARG, LOG2E = 4, ((WMMA_M, WMMA_N, WMMA_K), 'AMD', 32), math.log2(math.e)
-Q4_K, Q5_K, Q6_K, IQ4_XS, GGML_BLOCK_SIZE, Q8_GROUP_SIZE, Q4_WORDS, Q5_WORDS, Q6_BYTES, IQ4_WORDS = 12, 13, 14, 23, 256, 32, 36, 44, 210, 34
-QUANT_SIZES = {Q4_K: Q4_WORDS*4, Q5_K: Q5_WORDS*4, Q6_K: Q6_BYTES, IQ4_XS: IQ4_WORDS*4}  # bytes per 256-weight block
-
-def kernel_var(x:UOp) -> UOp:
-  # a Variable is a 0-d ALU BUFFER in the tensor graph; inside kernels it takes the ALU PARAM form (same name keeps the value binding)
-  return x.substitute({v: UOp.variable(v.expr, v.vmin, v.vmax, dtype=v.dtype, multiple_of=v.arg.multiple_of, param=True)
-                       for v in x.toposort() if v.is_variable})
-
-def _unbind(v:int|UOp) -> int|UOp: return kernel_var(v.unbind_all()[0]) if isinstance(v, UOp) else v
-
-@functools.cache
-def amd_custom_kernels_supported(device:str|tuple[str, ...]|None) -> bool:
-  # the custom kernels are tuned for RDNA3 (gfx11): the WMMA register layouts don't match gfx12 (RDNA4)
-  # or CDNA (MFMA-only, wave64), and the dp4a builtins and 32-lane wave ops aren't portable either.
-  if isinstance(device, tuple): device = device[0]
-  if device is None or device.split(":")[0] != "AMD": return False
-  # @function contexts set ALLOW_DEVICE_USAGE=0 (scheduling must not open devices); the device is always open here
-  with Context(ALLOW_DEVICE_USAGE=1):
-    return (t:=getattr(Device[device], "target", None)) is not None and t[0] == 11
-
-def warp_reduce(val:UOp, maximum:bool=False, full_wave:bool=False) -> UOp:
-  for offset in ((16, 8, 4, 2, 1) if full_wave else (8, 4, 2, 1)):
-    if val.op is Ops.INDEX and val.addrspace == AddrSpace.REG: val = val.load()
-    other = UOp(Ops.CUSTOM, dtypes.float, (val,), arg=
-      f"__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {{0}}), {0x1f | offset<<10}))")
-    val = val.maximum(other) if maximum else val + other
-  return val
-
-def _reg(shape:tuple[int, ...], slot:int, value:float, dep:UOp|None=None) -> UOp:
-  ret = UOp.placeholder(shape, dtypes.float, slot=slot, addrspace=AddrSpace.REG)
-  return ret.after((ret if dep is None else ret.after(dep)).store(ret.const_like(value)))
-
-# ******** quant linear: q8-activation kernels over packed ggml weights (Q4_K/Q5_K/Q6_K/IQ4_XS) ********
-
-class Linear(nn.Linear):
-  ggml_type:int|None = None
-  use_custom_quant = True
-  def __init__(self, in_features:int, out_features:int, bias=True):
-    super().__init__(in_features, out_features, bias)
-    self.in_features, self.out_features = in_features, out_features
-  def set_quantized(self, decoded:Tensor):
-    packed_sizes = {decoded.numel() // 256 * type_size:typ for typ,type_size in QUANT_SIZES.items()}
-    raw = next((u for u in decoded.uop.toposort() if u.op is Ops.SHRINK and u.dtype == dtypes.uint8 and prod(u.shape) in packed_sizes), None)
-    if raw is None: return
-    raw_offset = raw.contiguous_view_offset()
-    assert raw_offset is not None and raw_offset % 4 == 0 and raw.buf_uop.dtype == dtypes.uint8
-    self.ggml_type = packed_sizes[prod(raw.shape)]
-    # store a typed buffer view: a lazy BITCAST is decomposed into byte-combining ALU before custom-kernel
-    # scheduling and would copy the entire packed weight on every JIT graph
-    packed_dtype = dtypes.uint8 if self.ggml_type == Q6_K else dtypes.uint32
-    self.weight = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer)
-      .view(raw.max_numel() * raw.dtype.itemsize // packed_dtype.itemsize, packed_dtype, raw_offset)))
-  def __call__(self, x:Tensor) -> Tensor:
-    supported = self.use_custom_quant and amd_custom_kernels_supported(self.weight.device)
-    if self.ggml_type is None and supported:
-      self.set_quantized(self.weight)
-      if self.ggml_type is None: self.use_custom_quant = supported = False  # not a supported quant format
-    if self.ggml_type in (Q4_K, Q5_K, Q6_K, IQ4_XS) and supported:
-      if isinstance(x.numel(), int): return q8_linear(self, x)
-      # symbolic token count: pad to the max chunk size so the kernels see static shapes, garbage rows are sliced off
-      out = q8_linear(self, x.pad_to(x.max_shape))
-      return out.shrink(tuple((0, s) for s in (*x.shape[:-1], self.out_features)))
-    return super().__call__(x)
-
-def _amd_dp4a(a:UOp, b:UOp, c:UOp) -> UOp:
-  return UOp(Ops.CUSTOMI, dtypes.int32, (a.int(), b.int(), c), arg="__builtin_amdgcn_sudot4(true, {}, true, {}, {}, false)")
-
-def _amd_byte_perm(a:UOp, b:UOp, selectors:UOp) -> UOp:
-  return UOp(Ops.CUSTOMI, dtypes.uint32, tuple(x.cast(dtypes.uint32) for x in (a, b, selectors)), arg="__builtin_amdgcn_perm({}, {}, {})")
-
-def _amd_load(ptr:UOp, lanes:int|None=None) -> UOp:
-  assert ptr.op is Ops.INDEX
-  if lanes is None: return UOp(Ops.CUSTOMI, ptr.dtype, (ptr,), arg="__builtin_nontemporal_load({0})")
-  buf, coords = ptr.src[0], ptr.src[1:]
-  idx = sum((coord*math.prod(buf.shape[i+1:]) for i,coord in enumerate(coords)), UOp.const(0))
-  return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes))).load(dtype=ptr.dtype)
-
-def _load_byte(raw:UOp, base:UOp, offset:UOp) -> UOp: return (raw[base + offset//4] >> ((offset&3)*8).cast(dtypes.uint32)) & 255
-def _half(value:UOp) -> UOp: return value.cast(dtypes.uint16).bitcast(dtypes.float16).float()
-
-def _iq4_bytes(packed:UOp, shift:int) -> UOp:
-  selectors = (packed >> shift) & 0x0f0f0f0f
-  low = _amd_byte_perm(UOp.const(0xf6eaddcf, dtypes.uint32), UOp.const(0xbfad9881, dtypes.uint32), selectors)
-  high = _amd_byte_perm(UOp.const(0x71594535, dtypes.uint32), UOp.const(0x26190d01, dtypes.uint32), selectors & 0x07070707)
-  return _amd_byte_perm(high, low, 0x03020100 | ((selectors & 0x08080808) >> 1))
-
-def _q5_scales(raw:UOp, base:UOp, subgroup:UOp) -> tuple[UOp, UOp, UOp, UOp]:
-  scale = (subgroup < 4).where(_load_byte(raw, base, 4 + subgroup) & 63,
-    (_load_byte(raw, base, 8 + subgroup) & 15) | ((_load_byte(raw, base, subgroup) >> 6) << 4))
-  minimum = (subgroup < 4).where(_load_byte(raw, base, 8 + subgroup) & 63,
-    (_load_byte(raw, base, 8 + subgroup) >> 4) | ((_load_byte(raw, base, 4 + subgroup) >> 6) << 4))
-  d, dmin = (raw[base] & 0xffff).cast(dtypes.uint16), (raw[base] >> 16).cast(dtypes.uint16)
-  return _half(d), _half(dmin), scale.float(), minimum.float()
-
-def _iq4_scales(raw:UOp, base:UOp, subgroup:UOp) -> tuple[UOp, UOp]:
-  low = _load_byte(raw, base, 4 + subgroup//2)
-  scale = ((low >> (4*(subgroup%2)).cast(dtypes.uint32)) & 15) | ((((raw[base] >> 16) >> (2*subgroup).cast(dtypes.uint32)) & 3) << 4)
-  return _half(raw[base] & 0xffff), (scale.cast(dtypes.uint8).bitcast(dtypes.int8)-32).float()
-
-@functools.cache
-def iq4_half_lut(device:str) -> Tensor:
-  from tinygrad.runtime.autogen.ggml_common import kvalues_iq4nl
-  return Tensor([x for j in range(16) for i in range(16) for x in (kvalues_iq4nl[i], kvalues_iq4nl[j])],
-                dtype=dtypes.float16, device=device).bitcast(dtypes.uint32).contiguous()
-
-@functools.cache
-def _q8_quantize_kernel(q:UOp, scale:UOp, x:UOp, tokens:int, in_features:int) -> UOp:
-  groups = in_features//Q8_GROUP_SIZE
-  token_group, lane = UOp.range(tokens*groups, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
-  token, group = token_group//groups, token_group%groups
-  x = x.reshape(tokens, groups, 32)
-  group_scale = (warp_reduce(x[token, group, lane].float().abs(), maximum=True, full_wave=True) / 127).maximum(1e-8)
-  word_lane = lane.minimum(7)
-  xs = tuple(x[token, group, word_lane*4+i].float() for i in range(4))
-  word = sum(((v/group_scale).round().clip(-127, 127).cast(dtypes.int8).cast(dtypes.uint8).cast(dtypes.uint32) << (i*8)
-              for i,v in enumerate(xs)), UOp.const(0, dtypes.uint32))
-  stores = (q[token, group, lane.valid(lane < 8)].store(word), scale[token, group.valid(lane.eq(0))].store(group_scale))
-  return UOp.group(*stores).end(token_group, lane).sink(arg=KernelInfo(name="q8_quantize", opts_to_apply=()))
-
-def q8_quantize(x:Tensor, tokens:int, in_features:int) -> tuple[Tensor, Tensor]:
-  groups = in_features//Q8_GROUP_SIZE
-  q = Tensor.empty(tokens, groups, 8, dtype=dtypes.uint32, device=x.device)
-  scale = Tensor.empty(tokens, groups, dtype=dtypes.float32, device=x.device)
-  q, scale = Tensor.custom_kernel(q, scale, x, fxn=functools.partial(_q8_quantize_kernel, tokens=tokens, in_features=in_features))[:2]
-  return q, scale
-
-def _decode_linear(out:UOp, out_features:int, group_count:int, group_dot, name:str) -> UOp:
-  chunks = (group_count+31)//32
-  token_output_chunk, lane = UOp.range(out.shape[0]*out_features*chunks, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
-  token, output, chunk = token_output_chunk // (out_features*chunks), (token_output_chunk//chunks) % out_features, token_output_chunk % chunks
-  group = lane+chunk*32
-  value = group_dot(token, output, group) if group_count % 32 == 0 else \
-    (group < group_count).where(group_dot(token, output, group.minimum(group_count-1)), UOp.const(0, dtypes.float32))
-  total = warp_reduce(value, full_wave=True)
-  return out[token, output, chunk.valid(lane.eq(0))].store(total.cast(out.dtype)).end(token_output_chunk, lane).sink(
-    arg=KernelInfo(name=name, opts_to_apply=()))
-
-@functools.cache
-def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, out_features:int, in_features:int, ggml_type:int) -> UOp:
-  group_count = in_features // Q8_GROUP_SIZE
-  def group_dot(token:UOp, output:UOp, group:UOp) -> UOp:
-    block, subgroup = group // 8, group % 8
-    xwords = _amd_load(xq[token, group, 0], 8)
-    if ggml_type in (Q4_K, Q5_K):
-      base = (output * in_features//GGML_BLOCK_SIZE + block) * (Q4_WORDS if ggml_type == Q4_K else Q5_WORDS)
-      qs_base, dot, qsum = base + (4 if ggml_type == Q4_K else 12) + (subgroup//2)*8, UOp.const(0, dtypes.int32), UOp.const(0, dtypes.int32)
-      for word_idx in range(8):
-        word = (raw[qs_base+word_idx] >> ((subgroup&1)*4).cast(dtypes.uint32)) & 0x0f0f0f0f
-        if ggml_type == Q5_K: word |= ((raw[base+4+word_idx] >> subgroup.cast(dtypes.uint32)) & 0x01010101) << 4
-        dot, qsum = _amd_dp4a(word, xwords[word_idx], dot), _amd_dp4a(UOp.const(0x01010101, dtypes.uint32), xwords[word_idx], qsum)
-      d, dmin, scale, minimum = _q5_scales(raw, base, subgroup)
-      return (dot.float()*d*scale - qsum.float()*dmin*minimum) * xd[token, group]
-    if ggml_type == IQ4_XS:
-      base = (output * in_features//GGML_BLOCK_SIZE + block) * IQ4_WORDS
-      dot = UOp.const(0, dtypes.int32)
-      for word_idx in range(8):
-        packed = _amd_load(raw[base + 2 + subgroup*4 + word_idx%4])
-        dot = _amd_dp4a(_iq4_bytes(packed, 4*(word_idx//4)), xwords[word_idx], dot)
-      d, scale = _iq4_scales(raw, base, subgroup)
-      return dot.float() * xd[token, group] * d * scale
-    base = (output*in_features//GGML_BLOCK_SIZE+block)*Q6_BYTES
-    dots = [UOp.const(0, dtypes.int32)] * 2
-    for word_idx in range(8):
-      pos, within = subgroup*32 + word_idx*4, (subgroup*32 + word_idx*4)%128
-      low = _amd_load(raw[base + (pos//128)*64 + within%64], 4) >> ((within//64)*4).cast(dtypes.uint8)
-      high = _amd_load(raw[base + 128 + (pos//128)*32 + within%32], 4) >> ((within//32)*2).cast(dtypes.uint8)
-      quant = ((low & 15) | ((high & 3) << 4)).bitcast(dtypes.int8) - 32
-      word = sum((quant[i].cast(dtypes.uint8).cast(dtypes.uint32) << (i*8) for i in range(4)), UOp.const(0, dtypes.uint32))
-      dots[word_idx//4] = _amd_dp4a(word, xwords[word_idx], dots[word_idx//4])
-    scales = [raw[base + 192 + subgroup*2+i].cast(dtypes.uint8).bitcast(dtypes.int8).float() for i in range(2)]
-    dbits = raw[base+208].cast(dtypes.uint16) | (raw[base+209].cast(dtypes.uint16) << 8)
-    return (dots[0].float()*scales[0] + dots[1].float()*scales[1]) * xd[token, group] * _half(dbits)
-  names = {Q4_K: "linear_q4_k", Q5_K: "linear_q5_k", IQ4_XS: "linear_iq4_xs", Q6_K: "linear_q6"}
-  return _decode_linear(out, out_features, group_count, group_dot, names[ggml_type])
-
-def _wmma_layout(out:UOp, out_features:int, token_tile:int, output_tiles:int):
-  output_waves = 2 if out_features % (32*output_tiles) == 0 else 1
-  token_block, output_block = UOp.range(out.shape[0]//token_tile, 0), UOp.range(out_features//(16*output_tiles*output_waves), 1)
-  lane, wave = UOp.range(WARP_SIZE, 2, axis_type=AxisType.LOCAL), UOp.range(output_waves, 3, axis_type=AxisType.LOCAL)
-  hw_lane = UOp(Ops.CUSTOM, dtypes.int32, (lane.int(),), arg="__builtin_amdgcn_mbcnt_lo(-1, 0)").cast(dtypes.weakint)
-  col, half = hw_lane % 16, hw_lane // 16
-  outputs = tuple((output_block*output_waves+wave)*(16*output_tiles) + tile*16 + col for tile in range(output_tiles))
-  inputs = tuple(token_block*token_tile + tile*16 + col for tile in range(token_tile//16))
-  tokens = tuple(tuple(token_block*token_tile + tile*16 + half*8 + i for i in range(8)) for tile in range(token_tile//16))
-  return output_waves, token_block, output_block, lane, wave, half, outputs, inputs, tokens
-
-def _wmma_stores(out, outputs, tokens, accs, update, half):
-  def values(acc:UOp) -> tuple[UOp, ...]:
-    vals = tuple(acc.after(update)[i].load() for i in range(8))
-    swapped = tuple(UOp(Ops.CUSTOM, dtypes.float32, (value,),
-      arg="__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {0}), 50688))") for value in vals)
-    low = half.eq(0)
-    return tuple(low.where(vals[i], swapped[i+4]) if j == 0 else low.where(swapped[i], vals[i+4]) for i in range(4) for j in range(2))
-  return [out[token, output].store(value) for output,output_accs in zip(outputs, accs)
-          for tile_tokens,acc in zip(tokens, output_accs) for token,value in zip(tile_tokens, values(acc))]
-
-def _quant_linear_wmma(out, x, out_features, in_features, type_words, layout, dequant, name):
-  x = x.reshape(out.shape[0], in_features)
-  _, token_block, output_block, lane, wave, physical_half, outputs, input_tokens, tokens = layout
-  token_tile, output_tiles = len(tokens)*16, len(outputs)
-  output_words = in_features // GGML_BLOCK_SIZE * type_words
-  accs = tuple(tuple(UOp.placeholder((8,), dtypes.float32, slot=ot*(token_tile//16)+tile, addrspace=AddrSpace.REG)
-                     for tile in range(token_tile // 16)) for ot in range(output_tiles))
-  accs = tuple(tuple(acc.after(acc.store(acc.const_like(0))) for acc in output_accs) for output_accs in accs)
-  group = UOp.range(in_features // Q8_GROUP_SIZE, 4, AxisType.REDUCE)
-  block, subgroup = group // 8, group % 8
-  wmma_accs = [list(output_accs) for output_accs in accs]
-  for half in range(2):
-    afrags = tuple(UOp.stack(*(x[input_token, group*32 + half*16 + i].cast(dtypes.float16) for i in range(16)))
-                   for input_token in input_tokens)
-    for output_tile,output in enumerate(outputs):
-      bfrag = UOp.stack(*dequant(output*output_words + block*type_words, subgroup, half))
-      for tile,afrag in enumerate(afrags):
-        previous = accs[output_tile][tile].after(group) if half == 0 else wmma_accs[output_tile][tile]
-        wmma_accs[output_tile][tile] = UOp.wmma(afrag, bfrag, previous, *WMMA_ARG)
-  update = UOp.group(*(acc.store(value) for output_accs,output_values in zip(accs, wmma_accs)
-                       for acc,value in zip(output_accs, output_values))).end(group)
-  return UOp.group(*_wmma_stores(out, outputs, tokens, accs, update, physical_half)).end(token_block, output_block, lane, wave).sink(
-    arg=KernelInfo(name=name, opts_to_apply=()))
-
-@functools.cache
-def _q5_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, out_features:int, in_features:int, ggml_type:int) -> UOp:
-  token_tile, output_tiles = (64, 1) if out_features <= 1024 and out.shape[0] % 64 == 0 else \
-    (64, 2) if out.shape[0] % 64 == 0 else (32 if out.shape[0] % 32 == 0 else 16, 2)
-  def dequant(base:UOp, subgroup:UOp, half:int) -> tuple[UOp, ...]:
-    d, dmin, scale, minimum = _q5_scales(raw, base, subgroup)
-    qs_base = base + (4 if ggml_type == Q4_K else 12) + (subgroup // 2)*8 + half*4
-    words = tuple((raw[qs_base+i] >> ((subgroup&1)*4).cast(dtypes.uint32) & 0x0f0f0f0f) |
-      (((raw[base+4+half*4+i] >> subgroup.cast(dtypes.uint32) & 0x01010101) << 4) if ggml_type == Q5_K else 0) for i in range(4))
-    return tuple(((word >> (byte*8) & 255).float()*d*scale-dmin*minimum).cast(dtypes.float16) for word in words for byte in range(4))
-  return _quant_linear_wmma(out, x, out_features, in_features, Q4_WORDS if ggml_type == Q4_K else Q5_WORDS,
-                            _wmma_layout(out, out_features, token_tile, output_tiles), dequant,
-                            f"linear_q{4 if ggml_type == Q4_K else 5}_k_f16_wmma")
-
-@functools.cache
-def _iq4_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, lut:UOp, out_features:int, in_features:int) -> UOp:
-  token_tile = 32 if out_features <= 1024 and out.shape[0] % 32 == 0 else 64 if out.shape[0] % 64 == 0 and \
-    (out_features <= 6144 or out_features == 5120 and in_features > 8192) else 128 if out.shape[0] % 128 == 0 else \
-    32 if out.shape[0] % 32 == 0 else 16
-  output_tiles = 1 if out_features <= 1024 else 2 if out_features <= 6144 else 1 if out_features < 8192 else 2
-  layout = _wmma_layout(out, out_features, token_tile, output_tiles)
-  output_waves, _, _, lane, wave, _, _, _, _ = layout
-  local_lut = UOp.placeholder((256,), dtypes.uint32, slot=32, addrspace=AddrSpace.LOCAL)
-  tid, lut_items = wave*32+lane, 256//(32*output_waves)
-  lut = local_lut.after(UOp.group(*(local_lut[tid*lut_items+i].store(lut[tid*lut_items+i]) for i in range(lut_items))).barrier())
-  def dequant(base:UOp, subgroup:UOp, half:int) -> tuple[UOp, ...]:
-    d, scale = _iq4_scales(raw, base, subgroup)
-    scale = scale * d
-    if out_features <= 6144:
-      pairs = tuple(lut[((raw[base + 2 + subgroup*4 + word] >> (byte*8)) & 255).cast(dtypes.weakint)]
-                    for word in range(4) for byte in range(4))
-      return tuple((_half((pair >> (half*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in pairs)
-    def nibble(packed:UOp, index:int): return (packed >> (8*index+4*half)) & 15
-    lut_pairs = (lut[(nibble(packed, i) | nibble(packed, i+1)<<4).cast(dtypes.weakint)]
-                 for packed in (raw[base+2+subgroup*4+i] for i in range(4)) for i in (0, 2))
-    return tuple((_half((pair >> (i*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in lut_pairs for i in range(2))
-  return _quant_linear_wmma(out, x, out_features, in_features, IQ4_WORDS, layout, dequant, "linear_iq4_xs_f16_wmma")
-
-def q8_linear(layer:Linear, x:Tensor) -> Tensor:
-  assert layer.ggml_type in (Q4_K, Q5_K, Q6_K, IQ4_XS)
-  tokens = int(x.numel()) // layer.in_features
-  raw, out_features, in_features = layer.weight.uop.buf_uop, layer.out_features, layer.in_features
-  def run(fxn:Callable[..., UOp], out:UOp, *srcs:UOp) -> Tensor:
-    all_srcs = (out,)+srcs
-    params = tuple(UOp.placeholder_like(src, slot=i) for i,src in enumerate(all_srcs))
-    kernel = fxn(*params, out_features=out_features, in_features=in_features).call(*all_srcs)
-    result = Tensor(out.after(kernel))
-    if len(result.shape) == 3: result = result.sum(-1)
-    result = result.reshape(*x.shape[:-1], out_features)
-    return result if layer.bias is None else result + layer.bias
-  out = Tensor.empty(tokens, out_features, dtype=dtypes.float32, device=x.device).uop
-  if tokens % 16 == 0 and out_features % 16 == 0 and layer.ggml_type in (Q4_K, Q5_K, IQ4_XS):
-    fxn = _iq4_linear_f16_wmma_kernel if layer.ggml_type == IQ4_XS else functools.partial(_q5_linear_f16_wmma_kernel, ggml_type=layer.ggml_type)
-    extra = (iq4_half_lut(str(x.device)).uop,) if layer.ggml_type == IQ4_XS else ()
-    return run(fxn, out, raw, x.cast(dtypes.float16).contiguous().uop, *extra)
-  xq, xd = q8_quantize(x, tokens, in_features)
-  decode = functools.partial(_quant_decode_kernel, ggml_type=layer.ggml_type)
-  out = Tensor.empty(tokens, out_features, (in_features+1023)//1024, dtype=dtypes.float32, device=x.device).uop
-  return run(decode, out, raw, xq.uop, xd.uop)
-
-# ******** flash attention on the KV cache ********
-
-@functools.cache
-def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, valid_kv_len, max_kv_len, block_n):
-  valid_kv_len = _unbind(valid_kv_len)
-  _, B, H_KV, N, D = cast(tuple[int, int, int, int, int], cache_kv.shape)
-  _, H, M, _ = cast(tuple[int, int, int, int], q.shape)
-  assert M == 1 and H % H_KV == 0 and D % WARP_SIZE == 0 and max_kv_len <= N and max_kv_len % block_n == 0
-  G, CHUNK, DV, heads_per_wave = H // H_KV, block_n, D // WARP_SIZE, 2
-  head_tile = min(DECODE_HEAD_TILE, G)  # share each KV stream across two GQA heads per wave
-  assert G % head_tile == 0 and head_tile % heads_per_wave == 0
-  decode_waves, decode_group = head_tile // heads_per_wave, 4
-  block_bhkv = UOp.range(B*H_KV*(G//head_tile), 0, AxisType.GLOBAL)
-  valid_chunks = (valid_kv_len+CHUNK-1)//CHUNK
-  group_count = min(valid_chunks, out.shape[2]) if isinstance(valid_chunks, int) else valid_chunks.minimum(out.shape[2])
-  block_n, lane = UOp.range(group_count, 1, AxisType.GLOBAL), UOp.range(WARP_SIZE, 2, axis_type=AxisType.LOCAL)
-  wave = UOp.range(decode_waves, 3, axis_type=AxisType.LOCAL)
-  head_group, bhkv = block_bhkv % (G//head_tile), block_bhkv // (G//head_tile)
-  b, kv_head = bhkv // H_KV, bhkv % H_KV
-  dims = tuple(lane + i*WARP_SIZE for i in range(DV))
-  acc, row_max, row_sum = _reg((heads_per_wave, DV), 0, 0), _reg((heads_per_wave,), 1, -math.inf), _reg((heads_per_wave,), 2, 0)
-  groups_per_chunk, offset = CHUNK // decode_group, UOp.range(((valid_chunks+group_count-1)//group_count)*(CHUNK//decode_group), 100, AxisType.REDUCE)
-  chunk = block_n + (offset // groups_per_chunk) * group_count
-  keys = tuple(chunk*CHUNK + (offset % groups_per_chunk)*decode_group + i for i in range(decode_group))
-  valid = tuple(key < valid_kv_len for key in keys)
-  kvals, vvals = (tuple(tuple(is_valid.where(cache_kv[kv, b, kv_head, key, d].float(), UOp.const(0, dtypes.float)) for d in dims)
-    for key,is_valid in zip(keys, valid)) for kv in range(2))
-  q_heads = tuple(kv_head*G + head_group*head_tile + wave*heads_per_wave + head for head in range(heads_per_wave))
-  updates:list[UOp] = []
-  for head,q_head in enumerate(q_heads):
-    scores = tuple(warp_reduce(sum((q[b, q_head, 0, d].float()*k for d,k in zip(dims, key_kvals)),
-                                   UOp.const(0, dtypes.float)), full_wave=True) / math.sqrt(D) for key_kvals in kvals)
-    prev_acc, prev_max, prev_sum = acc.after(offset)[head], row_max.after(offset)[head], row_sum.after(offset)[head]
-    new_max = functools.reduce(lambda a,vs:a.maximum(vs[0].where(vs[1], UOp.const(-math.inf, dtypes.float))), zip(valid, scores), prev_max)
-    alpha = ((prev_max-new_max)*LOG2E).exp2()
-    betas = tuple(is_valid.where(((score-new_max)*LOG2E).exp2(), UOp.const(0, dtypes.float)) for is_valid,score in zip(valid, scores))
-    updates += [acc[head].store(prev_acc*alpha + sum((UOp.stack(*value)*beta for value,beta in zip(vvals, betas)), acc[head].const_like(0))),
-                row_sum[head].store(prev_sum*alpha + sum(betas, UOp.const(0, dtypes.float))), row_max[head].store(new_max)]
-  update = UOp.group(*updates).end(offset)
-  acc, row_max, row_sum = acc.after(update), row_max.after(update), row_sum.after(update)
-  stores = [out[b, q_head, block_n, d].store(acc[head, i]) for head,q_head in enumerate(q_heads) for i,d in enumerate(dims)] + \
-    [stats[b, q_head.valid(lane.eq(0)), block_n, i].store(x[head]) for head,q_head in enumerate(q_heads) for i,x in enumerate((row_max, row_sum))]
-  return UOp.group(*stores).end(lane, wave, block_n, block_bhkv).sink(arg=KernelInfo(name="flash_decode_partial", opts_to_apply=()))
-
-def amd_flash_attention_decode(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, max_kv_len:int) -> Tensor:
-  B, H, D = cache_kv.shape[1], q.shape[1], cache_kv.shape[4]
-  chunks = min(64, max_kv_len // 128)
-  partial = Tensor.empty(B, H, chunks, D, dtype="float32", device=q.device)
-  stats = Tensor.empty(B, H, chunks, 2, dtype="float32", device=q.device)
-  fxn = functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len, block_n=128)
-  partial, stats = Tensor.custom_kernel(partial, stats, q, cache_kv, fxn=fxn)[:2]
-  live = (valid_kv_len+127)//128
-  live = min(live, chunks) if isinstance(live, int) else live.minimum(chunks)
-  partial, stats = partial[:, :, :live], stats[:, :, :live]
-  weights = ((stats[..., 0]-stats[..., 0].max(2, keepdim=True))*LOG2E).exp2()
-  return ((partial*weights.unsqueeze(-1)).sum(2) / (stats[..., 1]*weights).sum(2, keepdim=True)).unsqueeze(2)
-
-@functools.cache
-def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, valid_kv_len:int|UOp, q_start:int|UOp|None=None) -> UOp:
-  valid_kv_len, q_start = _unbind(valid_kv_len), _unbind(q_start) if q_start is not None else None
-  BH, M, D = q.shape
-  _, B, H_KV, physical_n, cache_dim = cache.shape
-  k, v = cache[0].reshape(B*H_KV, physical_n, cache_dim), cache[1].reshape(B*H_KV, physical_n, cache_dim)
-  assert k.shape == v.shape and BH % k.shape[0] == 0 and k.shape[2] == D
-  gqa_group = BH // k.shape[0]
-  if isinstance(M, int) and isinstance(valid_kv_len, int): assert M % BLOCK_M == 0 and valid_kv_len % BLOCK_N == 0
-  assert isinstance(D, int) and D % WMMA_K == 0 and D % LANES_PER_WAVE_N == 0
-  TM, TN, TD, SCALE = BLOCK_M//(WAVES_M*LANES_PER_WAVE_M), BLOCK_N//LANES_PER_WAVE_N, D//(WAVES_N*LANES_PER_WAVE_N), 1/math.sqrt(D)
-  # query row 0 sits at sequence position q_base (the queries may be padded beyond valid_kv_len - q_base rows)
-  q_base = valid_kv_len - M if q_start is None else q_start
-  block_bh, block_m = UOp.range(BH, 0, AxisType.GLOBAL), UOp.range(M // BLOCK_M, 1, AxisType.GLOBAL)
-  kv_head = block_bh // gqa_group
-  q, o = (x.reshape(BH, M//BLOCK_M, BLOCK_M, D)[block_bh, block_m] for x in (q, o))
-  k, v = k[kv_head], v[kv_head]
-  wave_m, wave_n, lane = UOp.range(WAVES_M, 2, AxisType.LOCAL), UOp.range(WAVES_N, 3, AxisType.LOCAL), UOp.range(WARP_SIZE, -1, AxisType.WARP)
-  tid, lane_m, lane_n = (wave_m * WAVES_N + wave_n) * WARP_SIZE + lane, lane // LANES_PER_WAVE_N, lane % LANES_PER_WAVE_N
-  Q_ELEMS_PER_THREAD, KV_ELEMS_PER_THREAD = BLOCK_M * D // THREADS_PER_BLOCK, BLOCK_N * D // THREADS_PER_BLOCK
-  QP_lds = UOp.placeholder((BLOCK_M, D + LDS_PAD), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)
-  KV_lds = UOp.placeholder((BLOCK_N, D + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :D]
-  acc, m_i, l_i = _reg((TM, TD), 2, 0), _reg((TM,), 3, -math.inf), _reg((TM,), 4, 0)
-  n_tile = UOp.range((q_base + (block_m + 1) * BLOCK_M + BLOCK_N - 1) // BLOCK_N, 100, AxisType.REDUCE)
-  Q_lds = QP_lds[:, :D]
-  Q_store = Q_lds.after(n_tile).reshape(THREADS_PER_BLOCK, Q_ELEMS_PER_THREAD)[tid].store(q.reshape(THREADS_PER_BLOCK, Q_ELEMS_PER_THREAD)[tid])
-  load_k = UOp.range(KV_ELEMS_PER_THREAD, 90)
-  kval = k.reshape(physical_n*D)[n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_k].float()
-  K_store = KV_lds.reshape(THREADS_PER_BLOCK, KV_ELEMS_PER_THREAD)[tid, load_k].store(kval).end(load_k)
-  qk_load_barrier = UOp.barrier(UOp.group(Q_store, K_store))
-  Q_lds, KV_lds_k = Q_lds.after(qk_load_barrier), KV_lds.after(qk_load_barrier)
-  S_reg = _reg((TM, TN), 6, 0, n_tile)
-  k_qk, tm1, tn1 = UOp.range(D//WMMA_K, 101, AxisType.REDUCE), UOp.range(TM//WMMA_ACC, 200), UOp.range(TN, 201)
-  S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
-  q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
-  k_frag = KV_lds_k.reshape(TN, WMMA_N, D // WMMA_K, WMMA_K)[tn1, lane_n, k_qk]
-  qk_done = S_frag.store(UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), *WMMA_ARG)).end(tm1, tn1).end(k_qk)
-  S_reg = S_reg.after(qk_done, S_reg.store(S_reg * SCALE))
-  rm, rn = UOp.range(TM, 250), UOp.range(TN, 251)
-  q_idx = q_base + block_m * BLOCK_M + wave_m * WMMA_M + rm * LANES_PER_WAVE_M + lane_m
-  k_idx = n_tile * BLOCK_N + rn * LANES_PER_WAVE_N + lane_n
-  S_reg = S_reg.after(S_reg[rm, rn].store((k_idx <= q_idx).where(S_reg[rm, rn], S_reg[rm, rn].const_like(-math.inf))).end(rm, rn))
-  m_ij, rm2 = _reg((TM,), 7, -math.inf, n_tile), UOp.range(TN, 261, AxisType.REDUCE)
-  m_ij = m_ij.after(m_ij.store(m_ij.after(rm2).maximum(S_reg[:, rm2])).end(rm2))
-  ri_w = UOp.range(TM, 270)
-  m_ij = m_ij.after(m_ij[ri_w].store(warp_reduce(m_ij[ri_w], maximum=True)).end(ri_w))
-  tile_max = m_ij.reshape(TM, 1).expand(TM, TN).maximum(-1e30)
-  S_reg = S_reg.after(S_reg.store(((S_reg - tile_max) * LOG2E).exp2()))
-  p_local, ri_ws = _reg((TM,), 8, 0, n_tile), UOp.range(TM, 295)
-  p_sum = p_local.after(p_local[ri_ws].store(sum((warp_reduce(S_reg[ri_ws, rn]) for rn in range(TN)), S_reg.const_like(0))).end(ri_ws))
-  P_lds = QP_lds.flatten()[:WAVES_N * BLOCK_M * BLOCK_N].reshape(WAVES_N, BLOCK_M, BLOCK_N)
-  P_write = P_lds.reshape(WAVES_N, WAVES_M, TM, LANES_PER_WAVE_M, 1, TN, LANES_PER_WAVE_N, 1).permute((1, 0, 3, 6, 2, 4, 5, 7)) \
-    .reshape(THREADS_PER_BLOCK, TM, TN)
-  P_store = P_write[tid].store(S_reg.cast(dtypes.half))
-  beta_i, ri4, rj4 = UOp.placeholder((TM,), dtypes.float, slot=9, addrspace=AddrSpace.REG), UOp.range(TM, 330), UOp.range(TD, 331)
-  m_new = m_i[ri4].maximum(m_ij[ri4])
-  alpha_val, beta_val = ((m_i[ri4] - m_new) * LOG2E).exp2(), ((m_ij[ri4] - m_new) * LOG2E).exp2()
-  correction = UOp.group(acc[ri4, rj4].store(alpha_val * acc[ri4, rj4]).end(rj4),
-                         l_i[ri4].store(alpha_val * l_i[ri4] + beta_val * p_sum[ri4]),
-                         m_i[ri4].store(m_new), beta_i[ri4].store(beta_val)).end(ri4)
-  acc, l_i, m_i, beta_i = acc.after(correction), l_i.after(correction), m_i.after(correction), beta_i.after(correction)
-  V_lds = UOp.placeholder((D, BLOCK_N + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :BLOCK_N]
-  V_copy, load_v = V_lds.after(qk_done).permute(1, 0), UOp.range(KV_ELEMS_PER_THREAD, 390)
-  vval = v.reshape(physical_n*D)[n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_v].float()
-  V_store = V_copy.reshape(THREADS_PER_BLOCK, KV_ELEMS_PER_THREAD)[tid, load_v].store(vval).end(load_v)
-  pv_barrier = UOp.barrier(UOp.group(P_store, V_store))
-  P_lds, V_lds = P_lds.after(pv_barrier), V_lds.after(pv_barrier)
-  pv_acc = _reg((TM, TD), 10, 0, n_tile).after(pv_barrier)
-  k_pv, tm2, tn2 = UOp.range(BLOCK_N//WMMA_K, 400, AxisType.REDUCE), UOp.range(TM//WMMA_ACC, 401), UOp.range(TD, 402)
-  pv_frag = pv_acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2]
-  p_frag = P_lds[wave_n].reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv]
-  v_frag = V_lds.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv]
-  pv_done = pv_frag.store(UOp.wmma(p_frag, v_frag, pv_frag.after(k_pv), *WMMA_ARG)).end(tm2, tn2).end(k_pv)
-  pv_acc = pv_acc.after(pv_done)
-  ri5, rj5 = UOp.range(TM, 410), UOp.range(TD, 411)
-  n_tile_end = acc[ri5, rj5].store(acc[ri5, rj5] + beta_i[ri5] * pv_acc[ri5, rj5]).end(ri5, rj5).barrier().end(n_tile)
-  acc, l_i, m_i = acc.after(n_tile_end), l_i.after(n_tile_end), m_i.after(n_tile_end)
-  acc = acc.after(acc.store(acc * (1 / l_i).reshape(TM, 1).expand(TM, TD)))
-  o = o.reshape(WAVES_M, TM, LANES_PER_WAVE_M, 1, WAVES_N, TD, LANES_PER_WAVE_N, 1) \
-    .permute((0, 4, 2, 6, 1, 3, 5, 7)).reshape(THREADS_PER_BLOCK, TM, TD)
-  return o[tid].store(acc).end(wave_m, wave_n, lane).end(block_m, block_bh).sink(arg=KernelInfo(opts_to_apply=()))
-
-def flash_attention(q:Tensor, assigned_kv:Tensor, valid_end:int|UOp) -> Tensor:
-  # cached flash attention on the half KV cache (already written through assigned_kv); valid_end stays bound at the graph level
-  T_real, q_start = q.shape[2], None
-  if resolve(T_real == 1): return amd_flash_attention_decode(q.half(), assigned_kv, valid_end, cast(int, assigned_kv.shape[3]))
-  if isinstance(T_real, UOp):
-    # symbolic chunk: pad the queries to the static tile; garbage rows are sliced off
-    T_pad = q.max_shape[2]
-    assert T_pad % BLOCK_M == 0, "chunk_size must be a multiple of 32"
-    q, q_start = q.pad_to((*q.shape[:2], T_pad, q.shape[3])), valid_end - T_real
-  B, H, T, D = q.shape
-  out = Tensor.empty(B*H, T, D, dtype="float32", device=q.device)
-  fxn = functools.partial(_amd_flash_attention, valid_kv_len=valid_end, q_start=q_start)
-  out = Tensor.custom_kernel(out, q.half().reshape(B*H, T, D), assigned_kv, fxn=fxn)[0].reshape(B, H, T, D)
-  return out if q_start is None else out[:, :, :T_real]
-
-# ******** gated delta net: fused recurrent scan ********
-
-@functools.cache
-def _gated_delta_prefill_kernel(core:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp, kq:UOp, start_pos:UOp|None=None) -> UOp:
-  batch, heads, tokens, value_dim, row_tile = *core.shape, 4
-  key_dim, alpha_dim = q.shape[-1], alpha.shape[-1] if len(alpha.shape) == 4 else 1
-  assert all(isinstance(x, int) for x in (batch, heads, tokens, value_dim, key_dim)) and key_dim % 32 == 0 and value_dim % row_tile == 0
-  batch, heads, tokens, value_dim, key_dim = cast(tuple[int, int, int, int, int], (batch, heads, tokens, value_dim, key_dim))
-  core, v = (x.reshape(batch*heads, tokens, value_dim) for x in (core, v))
-  q, k = (x.reshape(batch*heads, tokens, key_dim) for x in (q, k))
-  beta, kq = (x.reshape(batch*heads, tokens) for x in (beta, kq))
-  alpha, state = alpha.reshape(batch*heads, tokens, alpha_dim), state.reshape(batch*heads, value_dim, key_dim)
-  bh_row, lane = UOp.range(batch*heads*value_dim//row_tile, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
-  bh, row_base = bh_row // (value_dim//row_tile), (bh_row % (value_dim//row_tile))*row_tile
-  rows, cols = tuple(row_base+i for i in range(row_tile)), tuple(lane + i*32 for i in range(key_dim//32))
-  current = UOp.placeholder((row_tile*key_dim//32,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
-  initial = None if start_pos is None else start_pos.eq(0)
-  current = current.after(current.store(UOp.stack(*(state[bh, row, col].float() if initial is None else
-    initial.where(0, state[bh, row, col].float()) for row in rows for col in cols))))
-  token = UOp.range(tokens, 2, AxisType.REDUCE)
-  keys = tuple(k[bh, token, col].load() for col in cols)
-  queries = tuple(q[bh, token, col].load() for col in cols)
-  updates, stores = [], []
-  for row_idx,row in enumerate(rows):
-    previous = tuple(current.after(token)[row_idx*key_dim//32+i].load() for i in range(key_dim//32))
-    av, bv = alpha[bh, token, row if alpha_dim > 1 else 0].load(), beta[bh, token].load()
-    state_k = warp_reduce(sum((x*y for x,y in zip(previous, keys)), UOp.const(0, dtypes.float32)), full_wave=True)
-    state_q = warp_reduce(sum((x*y for x,y in zip(previous, queries)), UOp.const(0, dtypes.float32)), full_wave=True)
-    delta = (v[bh, token, row].load() - state_k*av) * bv
-    updates += [x*av + delta*y for x,y in zip(previous, keys)]
-    stores.append(core[bh, token, row.valid(lane.eq(0))].store(state_q*av + delta*kq[bh, token]))
-  step = UOp.group(*stores, current.store(UOp.stack(*updates))).end(token)
-  state_stores = (state[bh, row, col].store(current.after(step)[row_idx*key_dim//32+i].load().cast(state.dtype))
-                  for row_idx,row in enumerate(rows) for i,col in enumerate(cols))
-  return UOp.group(*state_stores).end(lane, bh_row).sink(arg=KernelInfo(name="gated_delta_prefill", opts_to_apply=()))
-
-def gated_delta_prefill(q:Tensor, k:Tensor, v:Tensor, beta:Tensor, alpha:Tensor, state:Tensor, start_pos:Tensor|None=None) -> Tensor:
-  batch, heads, tokens, key_dim = q.shape
-  value_dim = v.shape[-1]
-  assert q.shape == k.shape and v.shape[:3] == beta.shape == (batch, heads, tokens) and state.shape == (batch, heads, value_dim, key_dim)
-  assert alpha.shape[:3] == (batch, heads, tokens) and (len(alpha.shape) == 3 or alpha.shape[-1] in (1, value_dim))
-  assert key_dim % 32 == 0 and value_dim % 4 == 0
-  core, kq = Tensor.empty_like(v), (q*k).sum(-1).contiguous()
-  srcs = (core, q.contiguous(), k.contiguous(), v.contiguous(), beta.contiguous(), alpha.contiguous(), state, kq)
-  if start_pos is None: return Tensor.custom_kernel(*srcs, fxn=_gated_delta_prefill_kernel)[0]
-  contig = tuple(x.uop if x.uop.op is Ops.AFTER else x.uop.contiguous() for x in srcs)
-  params = tuple(UOp.placeholder_like(x, slot=i) for i,x in enumerate(contig))
-  assert start_pos.uop.is_bound_var
-  # the bound start_pos reaches the graph through the state AFTER chain, like the flash kernels' valid_end
-  call = _gated_delta_prefill_kernel(*params, kernel_var(start_pos.uop.src[0])).call(*contig)
-  return Tensor(contig[0].after(call))

tinygrad/llm/model.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 import enum, functools, itertools, pathlib
 from dataclasses import dataclass, replace
 from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, dtypes
-from tinygrad.llm.kernels.amd import Linear, gated_delta_prefill, flash_attention, amd_custom_kernels_supported
+from tinygrad.nn import Linear
 from tinygrad.llm.gguf import gguf_load
 from tinygrad.uop.ops import resolve
 
@@ -181,13 +181,7 @@ class TransformerBlock(FFNBlock):
     k = apply_rope(k[..., :self.config.rope_dim], self.freqs_cis[start_pos:start_pos+T]).cat(k[..., self.config.rope_dim:], dim=-1)
 
     # NOTE: we don't want to change self.cache_kv, the function API doesn't support this well
-    store = self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(Tensor.stack(k, v).cast(dtypes.half).uop)
-    assigned_kv = Tensor(self.cache_kv.uop.after(store))
-    # on RDNA3, hybrid models use custom flash attention kernels on the KV cache
-    if amd_custom_kernels_supported(x.device) and self.config.ssm is not None:
-      attn = flash_attention(q, assigned_kv, start_pos+T)
-      attn = attn.transpose(1, 2).reshape(B, T, -1)                                    # back to (B,T,D)
-      return self.attn_output(attn if not self.config.attn_output_gate else (attn * gate.sigmoid()))
+    assigned_kv = Tensor(self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(Tensor.stack(k, v).uop)))
     k = assigned_kv[0, :, :, 0:start_pos+T, :]
     v = assigned_kv[1, :, :, 0:start_pos+T, :]
 
@@ -205,9 +199,8 @@ class TransformerBlock(FFNBlock):
 
   def _init_state(self, x:Tensor):
     if not hasattr(self, "cache_kv"):
-      # zeroed so the flash kernels can safely read whole tiles past the valid region (masked lanes multiply by 0)
-      self.cache_kv = Tensor.zeros(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim,
-                                   dtype=dtypes.half, device=x.device)
+      self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim,
+                                   dtype=dtypes.default_float, device=x.device)
       self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)
 
 class MLATransformerBlock(FFNBlock):
@@ -318,28 +311,21 @@ class GatedDeltaNetBlock(FFNBlock):
     v = v.reshape(B, T_pad, self.num_v_heads, self.head_v_dim)
     # layout the per-step operands to broadcast against the (B, H, V, K) state
     q, k, v, beta = (z.transpose(1, 2).float() for z in (q, k, v, beta))
-    q = q * self.head_k_dim**-0.5
-    alpha = log_alpha.transpose(1, 2).exp()  # per-channel decay for kda, per-head otherwise (B, H, T, V|1)
+    q, k, v, beta = q.unsqueeze(-2) * self.head_k_dim**-0.5, k.unsqueeze(-2), v.unsqueeze(-1), beta.unsqueeze(-1).unsqueeze(-1)
+    alpha = log_alpha.transpose(1, 2).exp().unsqueeze(-1)  # per-channel decay for kda, per-head otherwise (B, H, T, V|1, 1)
 
     # recurrent: scan over the (padded) tokens, updating the recurrent state. collect the per-step outputs
-    state = Tensor(self.recurrent_state.uop.after(conv_state_store))  # carry the conv write into this graph
-    if self.head_k_dim % 32 == 0 and self.head_v_dim % 4 == 0 and amd_custom_kernels_supported(x.device):
-      # one fused kernel for the whole scan; it resets and updates the recurrent state in place (RDNA3)
-      core = gated_delta_prefill(q, k, v, beta, alpha, state, Tensor(start_pos)).transpose(1, 2)
-    else:
-      q, k, v, beta = q.unsqueeze(-2), k.unsqueeze(-2), v.unsqueeze(-1), beta.unsqueeze(-1).unsqueeze(-1)
-      alpha = alpha.unsqueeze(-1)
-      state = initial.where(0, state.float())
-      outs = []
-      for t in range(T_pad):
-        s1 = state * alpha[:, :, t]  # decay the state
-        delta = (v[:, :, t] - (s1*k[:, :, t]).sum(-1, keepdim=True)) * beta[:, :, t]  # the delta rule update
-        state = s1 + delta * k[:, :, t]
-        outs.append((state * q[:, :, t]).sum(-1))
+    state = Tensor(self.recurrent_state.uop.after(conv_state_store)).float()  # carry the conv write into this graph
+    state = initial.where(0, state)
+    outs = []
+    for t in range(T_pad):
+      s1 = state * alpha[:, :, t]  # decay the state
+      delta = (v[:, :, t] - (s1*k[:, :, t]).sum(-1, keepdim=True)) * beta[:, :, t]  # the delta rule update
+      state = s1 + delta * k[:, :, t]
+      outs.append((state * q[:, :, t]).sum(-1))
 
-      # store the updated recurrent state in place, then read the stacked outputs after the write
-      state_store = self.recurrent_state.uop.store(state.cast(self.recurrent_state.dtype).uop)
-      core = Tensor(outs[0].stack(*outs[1:], dim=1).contiguous().uop.after(state_store))
+    # store the updated recurrent state in place, then read the stacked outputs after the write
+    core = Tensor(outs[0].stack(*outs[1:], dim=1).contiguous().uop.after(self.recurrent_state.uop.store(state.cast(self.recurrent_state.dtype).uop)))
 
     # output; undo the padding before the output projection
     z = (self.ssm_norm(core) * (out_gate.sigmoid() if is_kda else out_gate.silu())).cast(x.dtype).contiguous()
@@ -476,7 +462,7 @@ class Transformer:
     return min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk)
 
   def generate(self, tokens:list[int], chunk_size:int=32, temperature:float=0.0):
-    if self.has_recurrent_block and not amd_custom_kernels_supported(self.token_embd.weight.device): chunk_size = 1
+    if self.has_recurrent_block: chunk_size = 1
     v_start_pos = UOp.variable("start_pos", 0, self.max_context-1)
     v_toks = UOp.variable("toks", 1, chunk_size)
     # TODO: use UOp.variable for temperature once float variables are supported

tinygrad/nn/optim.py

@@ -121,10 +121,10 @@ class LARS(Optimizer):
         self.b[i].assign(self.momentum * self.b[i] + g)  # NOTE: self.b[i] is zero on the first run, no if required
         g = (g + self.momentum * self.b[i]) if self.nesterov else self.b[i]
       if self.ns_coefficients: g = g.reshape(g.shape[0], -1).newton_schulz(self.ns_steps, self.ns_coefficients).reshape(g.shape)
+      # muon does post momentum weight decay
+      if not self.pre_wd and self.wd > 0: t = t.detach() * (1.0 - self.wd * self.lr)
       # popular momentum does pre learning rate update
       if not self.classic: g = g * r * self.lr
-      # muon does post momentum weight decay
-      if not self.pre_wd and self.wd > 0: g = g + self.wd * self.lr * t.detach()
       ret.append(g.cast(t.dtype))
     return ret, self.b
 

tinygrad/renderer/cstyle.py

@@ -3,7 +3,6 @@ import math, sys, struct
 from collections import defaultdict, Counter
 from tinygrad.codegen.opt import tc
 from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str, axis_letters
-from tinygrad.uop.weak import commit_weak_consts
 from tinygrad.helpers import strip_parens, getenv, prod, dedup, Target, NUM_CPU_THREADS, IMAGE, FLOAT16, is_image_shape
 from tinygrad.dtype import dtypes, DType, AddrSpace, truncate, float_to_bf16
 from tinygrad.renderer import Renderer
@@ -39,8 +38,7 @@ base_rewrite = PatternMatcher([
   (UPat(Ops.CAST, name="x"), lambda ctx,x: f"__builtin_convertvector({ctx[x.src[0]]}, {ctx.render_type(x)})" \
     if x.max_numel() > 1 and x.addrspace is AddrSpace.REG else None),
   (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({ctx.render_cast(x, ctx[x.src[0]])})"),
-  (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"(({ctx._render_dtype(x.dtype, addrspace=x.addrspace)})({ctx[x.src[0]]}))"
-   if x.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None),
+  (UPat(Ops.BITCAST, name="x"), lambda ctx,x: ctx[x.src[0]] if x.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None),
   (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"__builtin_bit_cast({ctx.render_type(x)}, ({ctx.render_type(x.src[0])})({ctx[x.src[0]]}))"),
 
   # GPU stuff
@@ -76,8 +74,6 @@ base_rewrite = PatternMatcher([
 
 def create_non_native_float_pats(dts:tuple[DType, ...], casting:bool=True):
   patterns = PatternMatcher([
-    # a weak CONST states no width and cannot be restated: commit it at the emulated dtype a sibling src states
-    (UPat(GroupOp.ALU, name="x"), lambda x, dts=dts: commit_weak_consts(x, next((s.dtype for s in x.src if s.dtype in dts), None))),
     (UPat(Ops.WHERE, dtype=dts, src=(UPat.var("b"), UPat.var("x"), UPat.var("y")), name="w"),
      lambda w,b,x,y: b.where(x.cast(dtypes.float), y.cast(dtypes.float)).cast(w.dtype)),
     (UPat(GroupOp.ALU-{Ops.WHERE}, dtype=dts, name="x"),
@@ -241,7 +237,7 @@ class CStyleLanguage(Renderer):
       if (u.op is not Ops.CAST or u.max_numel() == 1) and ((u.op is Ops.CAST and u.src[0].op is Ops.CONST) or \
         u.op in {Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \
         (u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG and child_count[u] == 1) or \
-        (u.op in {Ops.CAST, Ops.BITCAST} and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)) or \
+        (u.op is Ops.CAST and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)) or \
         (u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))):
         r[u] = l
       else:
@@ -322,8 +318,7 @@ class OpenCLRenderer(CStyleLanguage):
   extra_matcher = create_non_native_float_pats((dtypes.bfloat16,)) + pm_manual_bf16_cast
 
   string_rewrite = PatternMatcher([
-    (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_{ctx.render_dtype(x.dtype)}(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"
-     if x.addrspace not in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None),
+    (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_{ctx.render_dtype(x.dtype)}(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"),
     # bfloat16 constants need to be rendered as their bit pattern since bf16 is stored as ushort
     (UPat.cvar("c").cast(dtypes.bfloat16), lambda ctx,c: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(c.val)))[0] >> 16)}u"),
     # load/store image (OpenCL)
@@ -374,8 +369,7 @@ class MetalRenderer(CStyleLanguage):
   ]) + pm_manual_bf16_cast
 
   string_rewrite = PatternMatcher([
-    (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_type<{ctx.render_dtype(x.dtype)}>(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"
-     if x.addrspace not in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None),
+    (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_type<{ctx.render_dtype(x.dtype)}>(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"),
   ]) + base_rewrite
 
   def render_kernel(self, function_name, kernel, bufs, uops, prefix=None):
@@ -431,8 +425,7 @@ class CUDARenderer(CStyleLanguage):
     (UPat(Ops.CAST, dtypes.fp8s, UPat.var("x", dtypes.fp8s), name='y'), lambda x,y: x.cast(dtypes.float).cast(y.dtype) if x.dtype!=y.dtype else None),
   ])
   string_rewrite = PatternMatcher([
-    (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"tg_bitcast<{ctx.render_dtype(x.dtype)}>(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"
-     if x.addrspace not in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None),
+    (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"tg_bitcast<{ctx.render_dtype(x.dtype)}>(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"),
   ]) + base_rewrite
 
   def render_vector_prefix(self, dt:DType, count:int) -> str:
@@ -527,6 +520,8 @@ class HIPRenderer(CStyleLanguage):
     (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
       lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2]))
       if x.src[0].max_numel() == 8 and x.src[0].dtype in dtypes.fp8_ocp else None),
+    # bfloat16 constant casting
+    (UPat.cvar('x', dtypes.bfloat16), lambda x: cast_float_to_bf16(UOp.const(x.val, dtypes.float))),
   ])
 
   def asm(self, prg:UOp, lin:UOp) -> bytes:

tinygrad/renderer/isa/x86.py

@@ -143,16 +143,16 @@ extra_matcher = PatternMatcher([
   (UPat.var("m").where(UPat.var("a", (dtypes.bool,)+dtypes.int8s), UPat.var("b")),
    lambda m,a,b: m.where(a.cast(dtypes.int16), b.cast(dtypes.int16)).cast(a.dtype) if a.max_numel() == 1 else None),
   # float16 alus are done in float32
-  (UPat(GroupOp.ALU, dtypes.float16, name="x"), lambda x: UOp(x.op,
-   src=tuple(s.cast(dtypes.float) if s.dtype != dtypes.bool else s for s in x.src)).cast(x.dtype)),
+  (UPat(GroupOp.ALU, dtypes.float16, name="x"), lambda x: UOp(x.op, dtypes.float,
+   tuple(s.cast(dtypes.float) if s.dtype != dtypes.bool else s for s in x.src)).cast(x.dtype)),
   (UPat(GroupOp.Comparison, src=(UPat.var("a", dtypes.float16), UPat.var("b")), name="x"),
    lambda x,a,b: UOp(x.op, src=(a.cast(dtypes.float32), b.cast(dtypes.float32))).cast(x.dtype)),
   # no cmpne for packed ints, y != x => !(y==x)
   (UPat(Ops.CMPNE, src=(UPat.var("y", dtypes.ints), UPat.var("x")), name="cmp"),
    lambda y,x,cmp: UOp(Ops.CMPEQ, src=(y,x))^True if y.max_numel() > 1 else None),
-  # float WHERE needs a mask unless its comparison already has a float operand
-  (UPat.var("m", dtypes.bool).where(UPat.var("a", dtypes.floats+(dtypes.weakfloat,)), UPat.var("b")).named("w"),
-   lambda m,a,b,w: m.cast(w.dtype).ne(0).where(a, b) if w.dtype in dtypes.floats and not dtypes.is_float(m.src[0].dtype) else None),
+  # float where expects a mask
+  (UPat.var("m", dtypes.bool).where(UPat.var("a", dtypes.floats), UPat.var("b")),
+   lambda m,a,b: m.cast(a.dtype).ne(0).where(a, b) if m.src[0].dtype not in dtypes.floats else None),
   # rewrite -x -> 0 - x
   (UPat(Ops.NEG, name="x"), lambda x: UOp(Ops.SUB, src=(x.const_like(0),) + x.src)),
   # TODO: add support for mod, requires support for accessing the 2nd+ reg of a multi output instruction
@@ -653,7 +653,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
     # 0b10 -- signals memory access with 32bit displacement
     # 0b11 -- signals no memory access
     if disp_uop is not None:
-      assert disp_uop.op is Ops.CAST, "displacement must be a const"
+      assert disp_uop.op is Ops.CAST, "displacement must be a literal"
       assert disp_uop.dtype in (dtypes.int8, dtypes.int32), "displacement can only be 1 or 4 byte signed int"
       # rbp/r13 always require a displacement
       if disp_uop.src[0].val != 0 or rm == 0b101: mod = 0b01 if disp_uop.dtype.itemsize == 1 else 0b10
@@ -836,7 +836,7 @@ class X86Renderer(ISARenderer):
     if x.op is Ops.BUFFER: x = x.replace(dtype=dtypes.uint64)
     is_xmm = isinstance(x.tag, tuple) and x.tag[0].cons[0].size == 16
     op = X86Ops.VMOVUPSm if is_xmm else X86Ops.MOVm
-    return UOp(Ops.INS, src=fold_address(self.stack_pointer().index(disp)) + (x,), arg=op, tag=x.tag)
+    return UOp(Ops.INS, dtypes.void, fold_address(self.stack_pointer().index(disp)) + (x,), op, x.tag)
 
   def fill(self, disp:UOp, x:UOp, reg:Register) -> UOp:
     is_xmm = reg.cons[0].size == 16

tinygrad/renderer/nir.py

@@ -121,6 +121,8 @@ class NIRRenderer(Renderer):
   code_for_op = {**{k:lambda:None for k in u_aop.keys()}, **{k:lambda:None for k in s_aop.keys()}, **{k:lambda:None for k in f_aop.keys()}}
 
   extra_matcher = PatternMatcher([
+    # handle negative unsigned CONST
+    (UPat.cvar("x", dtypes.uints), lambda x: UOp.const(x.dtype.max+x.val+1, x.dtype) if x.val < 0 else None),
     # from ptx
     (UPat.var('x', dtype=dtypes.bool)<UPat.var('y'), lambda x,y: (x^True)&y),
     # load/store bool -> uint8
@@ -134,10 +136,9 @@ class NIRRenderer(Renderer):
     # ref: https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpConvertFToU
     (UPat(Ops.CAST, (dtypes.uchar, dtypes.ushort), src=(UPat.var("x", dtypes.floats),), name="c"), lambda x,c: x.cast(dtypes.int32).cast(c.dtype)),
     # load/store use pointer arithmetic, and the cast does nothing. NOTE: this doesn't apply to image indexing cause it's 1-D
-    # nor to REG/ALU register picks, which keep their own index dtype
     (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("off")), allow_any_len=True, name="x"), lambda x,buf,off: x.replace(
       src=(buf,UOp.const(off.val, dtypes.long) if off.op is Ops.CONST else off.cast(dtypes.long))+x.src[2:])
-      if buf.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL) and not is_image_shape(buf._shape) else None),
+      if buf.addrspace != AddrSpace.REG and not is_image_shape(buf._shape) else None),
     # images need index to be int for nir (coordinates only: the INDEX keeps its access dtype)
     (UPat.var("buf").index(UPat.var("idx_y"), UPat.var("idx_x"), name="x"),
      lambda x,buf,idx_y,idx_x: x.replace(src=(buf, idx_y.cast(dtypes.int), idx_x.cast(dtypes.int)))),

tinygrad/renderer/ptx.py

@@ -52,8 +52,8 @@ ptx_matcher = PatternMatcher([
   (UPat(Ops.STORE, src=(UPat(name="idx"), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
    lambda x,idx: UOp(x.op, src=(x.src[0], x.src[1].cast(dtypes.uint8))+x.src[2:]) if idx.addrspace != AddrSpace.REG else None),
   # ptx shr and shl instructions require y to be uint
-  (UPat.var("x") << UPat.var("y"), lambda x,y: UOp(Ops.SHL, src=(x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
-  (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, src=(x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
+  (UPat.var("x") << UPat.var("y"), lambda x,y: UOp(Ops.SHL, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
+  (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
 ])
 
 def mem_type(x:UOp) -> str: return 'shared' if x.addrspace == AddrSpace.LOCAL else 'global'

tinygrad/renderer/wgsl.py

@@ -1,52 +1,55 @@
 from tinygrad.dtype import DType, dtypes, truncate, AddrSpace
 from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
 from tinygrad.renderer.cstyle import CStyleLanguage, base_rewrite
-from tinygrad.helpers import strip_parens, ceildiv
+from tinygrad.helpers import strip_parens
 
-# a field of `width` bits sitting in the low bits of val: shift it up to the sign bit, then let the arithmetic shift fill
-def sign_extend(val:UOp, width:int): return (val << (32-width)).bitcast(dtypes.int) >> (32-width)
+def _mask(dt:DType): return 0xFF if dt.itemsize == 1 else 0xFFFF
 
-# a packed field of dt: the word it lives in, its offset in that word, and its mask. width is 8*itemsize, bool is one bit in a byte
-def packed_field(bidx:UOp, dt:DType) -> tuple[UOp, UOp, int]:
-  elems, width = 4//dt.itemsize, 8*dt.itemsize
-  return bidx.src[0].index(bidx.src[1] // elems), (bidx.src[1].cast(dtypes.uint32) % elems) * width, (1 << width)-1
+def sign_extend(val:UOp, sext_am:int):
+  return (UOp.where((val >> (sext_am - 1)) > 0, UOp.const(0xffffffff << sext_am, dtypes.uint32), UOp.const(0, dtypes.uint32)) \
+        | val.bitcast(dtypes.uint32)).bitcast(dtypes.int)
 
 # store for char: buf[idx/4] <- (var << (idx%4)*8))
-def packed_store(s:UOp):
-  bidx, var, *gate = s.src
-  idx, shift_am, mask = packed_field(bidx, var.dtype)
+def packed_store(bidx:UOp, var:UOp, gate:UOp|None=None):
+  elems, mask = 4//var.dtype.itemsize, _mask(var.dtype)
+  shift_am, div_idx = (bidx.src[1].cast(dtypes.uint32) % elems) * (8*var.dtype.itemsize), bidx.src[1] // elems
   # bool does its mask math at int32: renderer rewrites run after weak dtypes are lowered, and bool & 0xFF would create a weakint const
   if var.dtype == dtypes.bool: var = var.cast(dtypes.int32)
   new_v, wmask = (var & mask).cast(dtypes.uint32) << shift_am, ((mask << shift_am) ^ 0xFFFFFFFF).cast(dtypes.uint32)
-  buf = idx.cast(dtypes.uint32).load(*((UOp.const(0, dtypes.uint32), *gate) if gate else ()))
-  return idx.store((buf & wmask) | new_v, *gate)
+  idx = UOp(Ops.INDEX, src=(bidx.src[0], div_idx))
+  buf = UOp.load(idx, *((UOp.const(0, dtypes.uint32), gate) if gate is not None else ()), dtype=dtypes.uint32)
+  return UOp.store(idx, (buf & wmask) | new_v, *((gate,) if gate is not None else ()))
 
 # load for char: sign_extend(buf[idx/4] >> ((idx%4)*8))
-def packed_load(root:UOp):
-  bidx, *alt = root.src
-  idx, shift_am, mask = packed_field(bidx, dtype:=root.dtype)
-  load = idx.cast(dtypes.uint32).load(*((alt[0].cast(dtypes.uint32), *alt[1:]) if alt else ()), arg=root.arg)
-  val = (load >> shift_am) & mask
+def packed_load(root:UOp, bidx:UOp, dtype:DType, var:UOp|None=None, gate:UOp|None=None):
+  elems, mask = 4//dtype.itemsize, _mask(dtype)
+  shift_am, div_idx = (bidx.src[1].cast(dtypes.uint32) % elems) * (8*dtype.itemsize), bidx.src[1] // elems
+  idx = UOp(Ops.INDEX, src=(bidx.src[0], div_idx))
+  load = UOp.load(idx, *((var, gate) if var is not None and gate is not None else root.src[1:]), dtype=dtypes.uint32, arg=root.arg)
+  val = (load.cast(dtypes.uint32) >> shift_am) & mask
   return sign_extend(val, 8*dtype.itemsize).cast(dtype) if dtype in [dtypes.char, dtypes.short] else val.cast(dtype)
 
 def is_packed(x:UOp):
-  dt = x.src[1].dtype if x.op is Ops.STORE else x.buf_uop.dtype
-  return dt.itemsize < 4 and dt != dtypes.half and x.buf_uop.addrspace != AddrSpace.REG
-def _packed_size(u:UOp): return ceildiv(u.max_numel(), 4//u.dtype.itemsize) if is_packed(u) else u.max_numel()
+  if x.op is Ops.LOAD: dt, addrspace = x.dtype, x.src[0].addrspace
+  elif x.op is Ops.STORE: dt, addrspace = x.src[1].dtype, x.src[0].addrspace
+  else: dt, addrspace = x.dtype, x.addrspace
+  return dt.itemsize < 4 and dt != dtypes.half and addrspace != AddrSpace.REG
+def _packed_size(u:UOp): return u.max_numel() // (4//u.dtype.itemsize) if is_packed(u) else u.max_numel()
 def is_nan(a):
   bs, (exp, mant) = a.dtype.bitsize, dtypes.finfo(a.dtype)
   return (a.bitcast(getattr(dtypes, f"uint{bs}")) & ((1 << (bs - 1)) - 1)) > (((1 << exp) - 1) << mant)
 
-# the read-modify-write packed_store emits: a load of the very index being stored to, masked (a gated store loads with 3 srcs)
-packed_rmw = UPat(Ops.LOAD, src=(UPat(Ops.CAST, dtype=dtypes.uint32, src=(UPat.var("b"),)),), allow_any_len=True) & UPat.var("wmask")
-
 wgsl_matcher = PatternMatcher([
   (UPat((Ops.CMPLT, Ops.XOR), src=(UPat(name="a", dtype=dtypes.bool), UPat.var("b")), name="c"),
    lambda a,b,c: a.cast(dtypes.int).alu(c.op, b.cast(dtypes.int)).cast(dtypes.bool)),
-  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX),), allow_any_len=True, name="l"), lambda l: packed_load(l) if is_packed(l) else None),
-  (UPat(Ops.STORE, name="s"), lambda s: packed_store(s) if is_packed(s) else None),
+  (UPat.load(UPat.var("b"), UPat.var("c"), UPat.var("gate"), name="l"),
+   lambda l,b,c,gate: packed_load(l,b,l.dtype,c.cast(dtypes.uint32),gate) if is_packed(l) else None),
+  (UPat.load(UPat.var("b"), name='l'), lambda l,b: packed_load(l,b,l.dtype) if is_packed(l) else None),
+  (UPat.store(UPat.var("b"), UPat.var("var"), UPat.var("gate"), name="s"),
+   lambda b,var,gate,s: packed_store(b,var,gate) if is_packed(s) else None),
+  (UPat.store(UPat.var("b"), UPat.var("var"), name="s"), lambda b,var,s: packed_store(b,var) if is_packed(s) else None),
   (UPat.var("a") << UPat.var("b"),lambda a,b:(a.bitcast(dtypes.uint32)<<b.cast(dtypes.uint32)).bitcast(a.dtype) if b.dtype!=dtypes.uint32 else None),
-  (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, src=(x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
+  (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
   # fix nan check: 'a != a -> is_nan()'. the decomp rewrites (a != a).logical_not() to CMPEQ, so match both forms
   (UPat.var("a", dtypes.floats) != UPat.var("a"), is_nan),
   (UPat.var("a", dtypes.floats).alu(Ops.CMPEQ, UPat.var("a")), lambda a: is_nan(a).ne(True)),
@@ -65,13 +68,11 @@ class WGSLRenderer(CStyleLanguage):
               dtypes.char: "i32", dtypes.int32: "i32", dtypes.uint32: "u32", dtypes.bool: "bool", dtypes.half: "f16" }
 
   string_rewrite = PatternMatcher([
-    (UPat(Ops.CAST, dtype=dtypes.uint32, src=(UPat(Ops.INDEX, name="x"),)), lambda ctx,x: ctx[x] if is_packed(x) else None),
     (UPat(Ops.NEG, dtypes.uints, src=(UPat.var('x'))), lambda ctx,x: f"(0-{ctx[x]})"),
     (UPat.cvar("c").cast(dtypes.bool), lambda c: "true" if c.val else "false"),
     (UPat.cvar("c").cast((dtypes.uchar, dtypes.ushort, dtypes.uint32)),
      lambda c: f"bitcast<u32>({c.val})" if c.val < 0 else f"{c.val&0xFFFFFFFF}u"),
-    # a negative const must state its type: contextual conversion of a bare abstract int rejects it in a u32 position
-    (UPat.cvar("c").cast(dtypes.int32, name="x"), lambda ctx,x,c: f"i32({v})" if (v:=truncate[x.dtype](c.val)) < 0 else f"{v}"),
+    (UPat.cvar("c").cast(dtypes.int32, name="x"), lambda ctx,x,c: f"{truncate[x.dtype](c.val)}"),
     (UPat(Ops.BUFFER, name="x"), lambda ctx,x:
      f"var{'<workgroup>' if x.addrspace == AddrSpace.LOCAL else ''} {ctx[x]}: array<{ctx.buf_map(x)},{_packed_size(x)}>;"),
     (UPat(Ops.BITCAST, dtype=dtypes.half, name="x", src=(UPat(dtype=(dtypes.short, dtypes.ushort, dtypes.uint32),),)),
@@ -86,10 +87,10 @@ class WGSLRenderer(CStyleLanguage):
     (UPat.load(UPat.var("b"), UPat.var("v"), UPat.var("gate")),
       lambda ctx,b,v,gate: f"select({ctx[v]}, {ctx.render_load(ctx[b], b.src[0])}, {ctx[gate]})"),
     (UPat.load(UPat.var("b")), lambda ctx, b: ctx.render_load(ctx[b], b)),
-    # packed_store writes (load & wmask) | new_v: atomicAnd clears the field, atomicAdd sets it. new_v is gone when it is 0
-    (UPat.store(UPat.var("b"), UPat.any(packed_rmw, packed_rmw | UPat.var("nv"))), lambda ctx,b,wmask,nv=None:
-     f"atomicAnd(&{ctx[b]},{ctx[wmask]});"+(f"\n  atomicAdd(&{ctx[b]},{ctx[nv]});" if nv is not None else "") if is_packed(b) else None),
-    (UPat.store(UPat.var("b"), UPat.var("v")), lambda ctx,b,v: f"{ctx[b]} = {ctx[v]};"),
+    (UPat.store(UPat.var("b"), UPat.var("v")), lambda ctx,b,v:\
+     # (load & mask) | var -> mask = v.src[0].src[1], var = v.src[1]
+     f"atomicAnd(&{ctx[b]},{ctx[v.src[0].src[1]]});\n  atomicAdd(&{ctx[b]},{ctx[v.src[1]]});" if is_packed(b) \
+      else f"{ctx[b]} = {ctx[v]};"),
     (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx"))),
      lambda ctx,b,idx: f"{ctx[b]}[{strip_parens(ctx[idx]) if idx.arg is Ops.ADD else ctx[idx]}]"),
   ]) + base_rewrite

tinygrad/runtime/graph/hcq.py

@@ -139,8 +139,7 @@ class HCQGraph(MultiGraphRunner):
         prof_ji_desc = runtime.name if runtime is not None else TracingKey(f"{bufs[1].device} -> {bufs[0].device}", ret=bufs[0].nbytes)
 
         prof_name = enqueue_dev.device if runtime is not None else f"{enqueue_dev.device}:SDMA:{queue_idx}"
-        self.prof_graph_entries.append(ProfileGraphEntry(prof_name, prof_ji_desc, sig_st, j * 2 + 1,
-                                                         runtime.profile_key if runtime is not None else None))
+        self.prof_graph_entries.append(ProfileGraphEntry(prof_name, prof_ji_desc, sig_st, j * 2 + 1))
         self.prof_graph_deps.append([d - 1 for _, d in rdeps])
 
       self.last_j[enqueue_queue] = j