move shape into arg for param/buffer

* UOp.param accepts single sint as shape (int,)

Change UOp.param signature from shape:tuple[sint, ...]|None to shape:tuple[sint, ...]|sint|None.
A single sint means 1D shape (sint,). Update all callers from (n,) to n syntax.
No param_from_shape — multi-dim shapes stay as tuples.

* use single int syntax in hcq2 copy_with_kernel

.github/workflows/autogen.yml

@@ -40,10 +40,10 @@ jobs:
       run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev liburing-dev
     - name: Regenerate autogen files
       run: |
-        find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.py" -not -name "comgr.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
+        find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
         python3 -c "from tinygrad.runtime.autogen import opencl"
         python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv_610, nv"
-        python3 -c "from tinygrad.runtime.autogen import comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
+        python3 -c "from tinygrad.runtime.autogen import comgr, comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
         python3 -c "from tinygrad.runtime.autogen.am import *"
         python3 -c "from tinygrad.runtime.autogen.nv_regs import *"
         python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, pci, vfio"
@@ -54,7 +54,7 @@ jobs:
         python3 -c "from tinygrad.runtime.autogen import mesa"
         python3 -c "from tinygrad.runtime.autogen import avcodec"
         python3 -c "from tinygrad.runtime.autogen import llvm_qcom"
-        python3 -c "from tinygrad.runtime.autogen import mlx5"
+        python3 -c "from tinygrad.runtime.autogen import mlx5, bnxt"
         python3 -c "from tinygrad.runtime.autogen import ggml_common"
         REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
     - name: Check for differences
@@ -102,42 +102,3 @@ jobs:
       with:
         name: autogen-macos-patch
         path: autogen-macos.patch
-
-  autogen-comgr-2:
-    name: In-tree Autogen (comgr 2)
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: 'autogen-comgr'
-    - name: Install autogen support packages
-      run: |
-        wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
-        sudo tee /etc/apt/sources.list.d/rocm.list <<EOF
-        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.2 $(lsb_release -cs) main
-        EOF
-        echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
-        sudo apt -qq update || true
-        sudo apt-get install -y --no-install-recommends libclang-20-dev comgr
-    - name: Regenerate autogen files
-      run: |
-        rm tinygrad/runtime/autogen/comgr.py
-        python3 -c "from tinygrad.runtime.autogen import comgr"
-    - name: Check for differences
-      run: |
-        if ! git diff --quiet; then
-          git diff
-          git diff > autogen-comgr2.patch
-          echo "Autogen mismatch detected. Patch available at: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
-          exit 1
-        fi
-    - name: Upload patch artifact
-      if: failure()
-      uses: actions/upload-artifact@v7
-      with:
-        name: autogen-comgr2-patch
-        path: autogen-comgr2.patch

.github/workflows/benchmark.yml

@@ -431,32 +431,70 @@ jobs:
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
     - name: Kill stale pids
       run: |
-        PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
-        PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
-    # since sudo is required for usbgpu on macos, do not write bytecode, as some of the files are owned by root
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
     - name: UsbGPU boot time
-      run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEBUG=2 AM_RESET=1 DEV=USB+AMD time python3.11 test/test_tiny.py TestTiny.test_plus
+      run: GMMU=0 DEBUG=2 AM_RESET=1 DEV=USB+AMD time python3.11 test/test_tiny.py TestTiny.test_plus
     - name: UsbGPU tiny tests
-      run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
+      run: GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
     - name: UsbGPU copy speeds
-      run: sudo -E PYTHONDONTWRITEBYTECODE=1 SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
-    #- name: UsbGPU openpilot test
-    #  run: sudo -E PYTHONPATH=. GMMU=0 DEV=USB+AMD GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
     - name: UsbGPU (USB4/TB) install script
-      run: PYTHONPATH=. sh extra/setup_tinygpu_osx.sh
+      run: sh extra/setup_tinygpu_osx.sh
     - name: UsbGPU (USB4/TB) boot time
-      run: PYTHONPATH=. DEBUG=3 DEV=PCI+NV:NAK time python3.11 test/test_tiny.py TestTiny.test_plus
+      run: DEBUG=3 DEV=PCI+NV:NAK time python3.11 test/test_tiny.py TestTiny.test_plus
     - name: UsbGPU (USB4/TB) tiny tests
-      run: PYTHONPATH=. DEV=PCI+NV:NAK python3.11 test/test_tiny.py
+      run: DEV=PCI+NV:NAK python3.11 test/test_tiny.py
 
-  testcommalatest:
-    name: comma Benchmark (0.11.2)
+  testcomma:
+    strategy:
+      matrix:
+        dev: ['QCOM', 'QCOM:IR3']
+        version: ['0.11.0', '0.11.2']
+        model: ['vision', 'policy', 'supercombo', 'dmonitoring']
+        # exclude non-existent models
+        exclude: [{ version: '0.11.0', model: supercombo }, { version: '0.11.2', model: vision }, { version: '0.11.2', model: policy }]
+        include:
+          - version: '0.11.0'
+            model: vision
+            url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
+            timing: 17
+          - version: '0.11.0'
+            model: policy
+            url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
+            timing: 3.2
+          - version: '0.11.0'
+            model: dmonitoring
+            url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
+            timing: 11
+          - version: '0.11.2'
+            model: supercombo
+            url: https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
+            timing: 26
+          - dev: QCOM:IR3
+            version: '0.11.2'
+            model: supercombo
+            timing: 41
+          - version: '0.11.2'
+            model: dmonitoring
+            url: https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/3e7b31dfbc0a5234f1baf196513b77fc6af12204b8a8ffe8ee0417e48352f316
+            timing: 11
+          # IR3 dmonitoring is slightly slower
+          - dev: QCOM:IR3
+            model: dmonitoring
+            timing: 12
+      fail-fast: false
+    name: openpilot ${{ matrix.version }} compile3 ${{ matrix.model }} (DEV=${{ matrix.dev }})
     runs-on: [self-hosted, Linux, comma]
-    timeout-minutes: 12
+    timeout-minutes: 5
     defaults:
       run:
         shell: bash -e -o pipefail {0}
     if: github.repository_owner == 'tinygrad'
+    env:
+      DEV: ${{ matrix.dev }}
+      ASSERT_MIN_STEP_TIME: ${{ matrix.timing }}
+      BENCHMARK_LOG: ${{ matrix.dev == 'QCOM:IR3' && 'ir3_' || '' }}openpilot_${{ matrix.version }}_${{ matrix.model }}
     steps:
     - name: Checkout Code
       uses: actions/checkout@v6
@@ -467,45 +505,10 @@ jobs:
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
     - name: reset process replay
       run: test/external/process_replay/reset.py
-    - name: openpilot compile3 0.11.2 supercombo
-      run: BENCHMARK_LOG=openpilot_0_11_2_supercombo PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
-    - name: openpilot compile3 0.11.2 supercombo (from pickle)
-      run: BENCHMARK_LOG=openpilot_0_11_2_supercombo_run_pickle RUN_PICKLE=1 PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 DEV=QCOM taskset -c 4-7 python3 examples/openpilot/compile3.py
-    - name: IR3 openpilot compile3 0.11.2 supercombo
-      run: BENCHMARK_LOG=ir3_openpilot_0_11_2_supercombo PYTHONPATH="." ASSERT_MIN_STEP_TIME=41 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
-    - name: openpilot compile3 0.11.2 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_11_2_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/3e7b31dfbc0a5234f1baf196513b77fc6af12204b8a8ffe8ee0417e48352f316
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-  testcommaold:
-    name: comma Benchmark (0.11.0)
-    runs-on: [self-hosted, Linux, comma]
-    timeout-minutes: 12
-    defaults:
-      run:
-        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: openpilot compile3 0.11.0 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot compile3 0.11.0 driving_vision (from pickle)
-      run: BENCHMARK_LOG=openpilot_0_11_0_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM taskset -c 4-7 python3 examples/openpilot/compile3.py
-    - name: IR3 openpilot compile3 0.11.0 driving_vision
-      run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=18 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot compile3 0.11.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3.2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
-    - name: openpilot compile3 0.11.0 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_11_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: compile
+      run: FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py ${{ matrix.url }}
+    - name: run pickle
+      run: BENCHMARK_LOG="${BENCHMARK_LOG}_run_pickle" RUN_PICKLE=1 taskset -c 4-7 python3 examples/openpilot/compile3.py
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 

.pre-commit-config.yaml

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

examples/mlperf/model_train.py

@@ -1674,8 +1674,7 @@ def train_gptoss():
   config = {}
   BASEDIR            = config["BASEDIR"]                = Path(getenv("BASEDIR", "/raid/datasets/c4-8b/"))
   BS                 = config["BS"]                     = getenv("BS", 16)
-  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
-  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
+  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS
   SEED               = config["SEED"]                   = getenv("SEED", 5760)
   DATA_SEED          = config["DATA_SEED"]              = getenv("DATA_SEED", SEED)
   SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
@@ -1737,8 +1736,8 @@ def train_gptoss():
   params_wd = [p for p in params if p.ndim >= 3]
   params_no_wd = [p for p in params if p.ndim < 3]
   optim = GradAccClipAdamWGroup(
-    GradAccClipAdamW(params_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device),
-    GradAccClipAdamW(params_no_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=0.0, grad_acc=grad_acc, device=optim_device),
+    GradAccClipAdamW(params_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=1, device=optim_device),
+    GradAccClipAdamW(params_no_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=0.0, grad_acc=1, device=optim_device),
   )
 
   for p in optim.params:
@@ -1770,7 +1769,7 @@ def train_gptoss():
 
   @TinyJit
   @Context(TRAINING=1)
-  def minibatch(tokens:Tensor):
+  def step(tokens:Tensor):
     if is_dp: tokens = tokens.to(None).shard(device, 0)
     if not is_sharding: tokens = tokens.to(None)
 
@@ -1784,22 +1783,20 @@ def train_gptoss():
     for g, new_g in zip(grads, loss.gradient(*optim.params)):
       apply_grad(g, new_g.uop)
 
-    loss_cpu = loss.flatten().float().to("CPU")
-    return loss_cpu.realize(*grads)
+    Tensor.realize(loss, *grads)
 
-  @TinyJit
-  def optim_step():
-    grad_norm = clip_grads(grads, grad_acc, 1.0)
+    grad_norm = clip_grads(grads, 1, 1.0)
     optim.fstep(grads, grad_norm)
     scheduler.step()
 
     for g in grads: g.assign(0)
 
+    loss_cpu = loss.flatten().float().to("CPU")
     lr_cpu = optim.lr.float().to("CPU")
     grad_norm_cpu = grad_norm.float().to("CPU")
-    Tensor.realize(lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales)
+    Tensor.realize(loss_cpu, lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales)
 
-    return lr_cpu, grad_norm_cpu
+    return loss_cpu, lr_cpu, grad_norm_cpu
 
   @TinyJit
   @Context(TRAINING=0)
@@ -1848,30 +1845,20 @@ def train_gptoss():
       profile_marker(f"train @ {i}")
       st = time.perf_counter()
 
-      stopped = False
-      losses, data_time, dev_time = [], 0, 0
-      for _ in range(grad_acc if i >= 2 else 1):
-        ist = time.perf_counter()
-        try: tokens = next(train_iter)
-        except StopIteration:
-          stopped = True
-          break
-        mst = time.perf_counter()
-        data_time += mst - ist
-        losses.append(minibatch(tokens).item())
-        dev_time += time.perf_counter() - mst
-      if stopped: break
+      ist = time.perf_counter()
 
-      gt = time.perf_counter()
-      ret = optim_step()
-      lr, grad_norm = ret[0].item(), ret[1].item()
+      try: tokens = next(train_iter)
+      except StopIteration: break
+      mst = time.perf_counter()
+      data_time = mst - ist
+
+      ret = step(tokens)
+      dev_time = time.perf_counter() - mst
+
+      loss, lr, grad_norm = ret[0].item(), ret[1].item(), ret[2].item()
       et = time.perf_counter()
 
-      loss = sum(losses) / len(losses)
-      optim_time = et - gt
-      dev_time += optim_time
       step_time = et - st
-      gbs_time = gt - st
       if BENCHMARK: step_times.append(step_time)
 
       i += 1
@@ -1881,7 +1868,7 @@ def train_gptoss():
       gflops = GlobalCounters.global_ops / 1e9 / dev_time
       mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 4.6e15)) * 100
       tqdm.write(
-          f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
+          f"{i:5} {step_time:.3f} s step, {dev_time:.3f} s dev, {data_time:.3f} s data, {loss:.4f} loss, " \
           f"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
       if DEBUG >= 1: tqdm.write("  mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
 
@@ -1891,8 +1878,6 @@ def train_gptoss():
           "train/lr": lr,
           "train/grad_norm": grad_norm,
           "train/step_time": step_time,
-          "train/gbs_time": gbs_time,
-          "train/optim_time": optim_time,
           "train/dev_time": dev_time,
           "train/data_time": data_time,
           "train/mem": mem_gb,

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

@@ -44,7 +44,7 @@ export SEED=${SEED:-5760}
 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 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 
 export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then

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

@@ -44,6 +44,6 @@ export SEED=${SEED:-$RANDOM}
 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 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 
 python3 examples/mlperf/model_train.py

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

@@ -44,7 +44,7 @@ export SEED=$RANDOM
 export DATA_SEED=$SEED
 
 export JITBEAM=3
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 
 export LOGMLPERF=1
 

extra/amdpci/am_smi.py

@@ -84,7 +84,8 @@ class AMSMI(AMDev):
     with open(f"/sys/bus/pci/devices/{self.pcibus}/power_state", "r") as f: return f.read().strip().rstrip()
 
 class SMICtx:
-  def __init__(self):
+  def __init__(self, dev_filter=None):
+    self.dev_filter = dev_filter
     self.devs = []
     self.opened_pcidevs = []
     self.opened_pci_resources = {}
@@ -135,6 +136,7 @@ class SMICtx:
     pattern = os.path.join('/tmp', 'am_*.lock')
     for d in [f[8:-5] for f in glob.glob(pattern)]:
       if d.startswith("usb"): continue
+      if self.dev_filter is not None and d != self.dev_filter: continue
       if d not in self.opened_pcidevs:
         self._open_am_device(d)
 
@@ -406,7 +408,7 @@ if __name__ == "__main__":
 
   try:
     if not args.list: os.system('clear')
-    smi_ctx = SMICtx()
+    smi_ctx = SMICtx(args.dev)
     while True:
       smi_ctx.rescan_devs()
       smi_ctx.draw(args.list)

extra/bench_log.py

@@ -35,7 +35,7 @@ class WallTimeEvent:
     return self
   def __exit__(self, *_):
     self.time = time.monotonic() - self.start
-    _events[self.event]["wall"].append(self.time)
+    _events[self.event]["wall"].append((self.time, BENCHMARK_LOG.value))
     return False
 
 class KernelTimeEvent:
@@ -47,19 +47,19 @@ class KernelTimeEvent:
     self.start = GlobalCounters.time_sum_s
     return self
   def __exit__(self, *_):
-    _events[self.event]["kernel"].append(GlobalCounters.time_sum_s - self.start)
+    _events[self.event]["kernel"].append((GlobalCounters.time_sum_s - self.start, BENCHMARK_LOG.value))
     return False
 
 def log_event_instant(event:InstantBenchEvent, value:float):
-  _events[event].append(value)
+  _events[event].append((value, BENCHMARK_LOG.value))
 
 if BENCHMARK_LOG:
   INFLUXDB_HOST = getenv("INFLUXDB_HOST", "")
   INFLUXDB_ORG = getenv("INFLUXDB_ORG", "tiny")
   INFLUXDB_TOKEN = getenv("INFLUXDB_TOKEN", "")
 
-  def _create_point(run_id, i, attempt, ref, commit, name, value, run):
-    point = Point(BENCHMARK_LOG.value).tag("id", run_id).tag("index", i)
+  def _create_point(run_id, i, attempt, ref, commit, name, value, log_name, run):
+    point = Point(log_name.replace(':', '_').replace('.', '_')).tag("id", run_id).tag("index", i)
     point = point.tag("device", Device.DEFAULT)
     point = point.tag("attempt", attempt).tag("ref", ref).tag("commit", commit)
     point = point.field(name, value).field("x", run)
@@ -91,12 +91,12 @@ if BENCHMARK_LOG:
       run_id = str(uuid.uuid4())
       if isinstance(event, BenchEvent):
         for event_type, values in _events[event].items():
-          for i, value in enumerate(values):
-            point = _create_point(run_id, i, attempt, ref, commit, f"{event.value}_{event_type}", value, run)
+          for i, (value, log_name) in enumerate(values):
+            point = _create_point(run_id, i, attempt, ref, commit, f"{event.value}_{event_type}", value, log_name, run)
             points.append(point)
       else:
-        for i, value in enumerate(_events[event]):
-          point = _create_point(run_id, i, attempt, ref, commit, event.value, value, run)
+        for i, (value, log_name) in enumerate(_events[event]):
+          point = _create_point(run_id, i, attempt, ref, commit, event.value, value, log_name, run)
           points.append(point)
 
     write_options = WriteOptions(write_type=WriteType.synchronous, retry_interval=5000, max_retries=5, max_retry_delay=30000, exponential_base=2)

extra/bnxt_driver/bnxtdev.py

@@ -0,0 +1,238 @@
+import ctypes, struct
+from tinygrad.helpers import ceildiv, getenv, wait_cond, DEBUG
+from tinygrad.runtime.autogen import bnxt, pci
+from tinygrad.runtime.support.system import PCIDevice, System, ipv4_to_gid
+
+BNXT_DEBUG = getenv("BNXT_DEBUG", 0)
+BNXT_ACCESS, BNXT_INIT_MASK, BNXT_RTR_MASK, BNXT_RTS_MASK = 3, 0xd, 0x41515ad, 0xae005
+BNXT_CHIMP_COMM, BNXT_CHIMP_COMM_TRIGGER = 0x0, 0x100
+BNXT_BACKING_STORE = ((0, 2), (1, 0), (2, 2), (3, 0), (4, 2), (5, 0), (6, 0), (14, 2), (15, 0))
+
+def db_value(xid, typ, index, epoch):
+  return (xid & bnxt.DBC_DBC_XID_MASK | bnxt.DBC_DBC_PATH_ROCE | typ | bnxt.BNXT_QPLIB_DBR_VALID) << 32 | \
+         index & bnxt.DBC_DBC_INDEX_MASK | epoch << bnxt.BNXT_QPLIB_DBR_EPOCH_SHIFT
+
+def _pbl(dev, paddrs, queue=False):
+  if len(paddrs) == 1: return 0, paddrs[0]
+  values = [p | bnxt.PTU_PTE_VALID for p in paddrs]
+  if queue:
+    values[-1] |= bnxt.PTU_PTE_LAST
+    if len(values) > 1: values[-2] |= bnxt.PTU_PTE_NEXT_TO_LAST
+  table, table_paddrs = dev.pci_dev.alloc_sysmem(ceildiv(len(values), 512) * 0x1000)
+  table[:len(values) * 8] = struct.pack(f"<{len(values)}Q", *values)
+  if len(table_paddrs) == 1: return 1, table_paddrs[0]
+  top, top_paddrs = dev.pci_dev.alloc_sysmem(0x1000)
+  top[:len(table_paddrs) * 8] = struct.pack(f"<{len(table_paddrs)}Q", *(p | bnxt.PTU_PTE_VALID for p in table_paddrs))
+  return 2, top_paddrs[0]
+
+def _queue(dev, stride:int=16, aux=False):
+  mem, paddrs = dev.pci_dev.alloc_sysmem(0x1000 + aux * 0x400)
+  level, base = _pbl(dev, paddrs, queue=True)
+  return {"mem":mem, "paddrs":paddrs, "stride":stride, "prod":0, "cons":0, "level":level, "base":base}
+
+def _qread(q, i):
+  off = (i & 15) * q["stride"]
+  return q["mem"][off:off + q["stride"]]
+
+def _qwrite(q, i, data, aux=False):
+  off = 0x1000 + i % 128 * 8 if aux else (i & 15) * q["stride"]
+  q["mem"][off:off + len(data)] = data
+
+class BNXTDev:
+  def __init__(self, pci_dev:PCIDevice, ip:str=getenv("BNXT_IP", "10.0.0.1")):
+    self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
+    self.bar0, self.db = pci_dev.map_bar(0, fmt='I'), pci_dev.map_bar(2, fmt='Q')
+    pci_dev.write_config(pci.PCI_COMMAND, pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
+    self.resp, self.resp_pa = pci_dev.alloc_sysmem(0x1000)
+    self.seq = 0
+
+    ver = self.hwrm("ver_get")
+    if DEBUG >= 2: print(f"bnxt {self.devfmt}: firmware {ver.hwrm_fw_maj_8b}.{ver.hwrm_fw_min_8b}.{ver.hwrm_fw_bld_8b}")
+    self.hwrm("func_reset", timeout_ms=40000)
+    caps = self.hwrm("func_qcaps", fid=0xffff)
+    self.mac, self.port_id = int.from_bytes(bytes(caps.mac_address), 'big'), caps.port_id
+    self.hwrm("func_drv_rgtr")
+    self.db_off = self.hwrm("func_qcfg", fid=0xffff).legacy_l2_db_size_kb * 1024
+
+    self.setup_backing_store()
+    self._open_rcfw()
+    self._open_l2()
+    self.local_gid = ipv4_to_gid(ip)
+    gids, mac = (ctypes.c_uint32 * 4)(*(int.from_bytes(self.local_gid[i:i + 4], 'big') for i in (12, 8, 4, 0))), self.mac.to_bytes(6, 'big')
+    smac = (ctypes.c_uint16 * 3)(*(int.from_bytes(mac[i:i + 2], 'big') for i in (0, 2, 4)))
+    self.gid_id = self.rcfw("add_gid", gid=gids, src_mac=smac).xid
+
+    if DEBUG >= 2: print(f"bnxt {self.devfmt}: booted mac={self.mac.to_bytes(6, 'big').hex(':')} gid={self.local_gid.hex()}")
+
+  def hwrm(self, name, timeout_ms=10000, **fields):
+    inp, out = getattr(bnxt, f"struct_hwrm_{name}_input"), getattr(bnxt, f"struct_hwrm_{name}_output")
+    opcode = getattr(bnxt, f"HWRM_{name.upper()}")
+    self.seq = (self.seq + 1) & 0xffff
+    data = bytes(inp(req_type=opcode, cmpl_ring=bnxt.BNXT_HWRM_NO_CMPL_RING, seq_id=self.seq, target_id=bnxt.BNXT_HWRM_TARGET,
+                     resp_addr=self.resp_pa[0], **fields))
+    self.resp[:] = bytes(len(self.resp))
+    System.memory_barrier()
+    for i, w in enumerate(memoryview(bytearray(data.ljust(bnxt.HWRM_MAX_REQ_LEN, b'\0'))).cast('I')):
+      self.bar0[BNXT_CHIMP_COMM // 4 + i] = w
+    self.bar0[BNXT_CHIMP_COMM_TRIGGER // 4] = 1
+    def hdr(): return bnxt.struct_hwrm_resp_hdr.from_buffer_copy(bytes(self.resp[:8]))
+    wait_cond(lambda: (n := hdr().resp_len) and hdr().seq_id == self.seq and self.resp[n - 1], timeout_ms=timeout_ms, msg=f"HWRM {name}")
+    ret = out.from_buffer_copy(bytes(self.resp[:ctypes.sizeof(out)]))
+    assert ret.error_code == 0, f"HWRM {name}: {ret.error_code}"
+    return ret
+
+  def setup_backing_store(self):
+    counts: dict[int, int] = {}
+    for typ, extra in BNXT_BACKING_STORE:
+      caps = self.hwrm("func_backing_store_qcaps_v2", type=typ)
+      size, splits = caps.entry_size, tuple(getattr(caps, f"split_entry_{j}") for j in range(caps.subtype_valid_cnt))
+      counts[typ] = n = counts[0] if typ == 15 else max(caps.min_num_entries, sum(splits) + extra)
+      # a zero bitmap means the type has a single instance 0
+      for instance in [i for i in range(8) if caps.instance_bit_map >> i & 1] or [0]:
+        mem, paddrs = self.pci_dev.alloc_sysmem(ceildiv(n * size, 0x1000) * 0x1000)
+        if caps.ctx_init_value:
+          for off in range(caps.ctx_init_offset, len(mem), size): mem[off] = caps.ctx_init_value
+        lvl, base = _pbl(self, paddrs)
+        self.hwrm("func_backing_store_cfg_v2", type=typ, instance=instance, entry_size=size, num_entries=n, page_dir=base,
+          page_size_pbl_level=lvl, subtype_valid_cnt=len(splits),
+          flags=bnxt.FUNC_BACKING_STORE_CFG_V2_REQ_FLAGS_BS_CFG_ALL_DONE if typ == 15 else 0,
+          **{f"split_entry_{j}": v for j, v in enumerate(splits)})
+
+  def _open_rcfw(self):
+    self.rcfw_first = True
+
+    self.creq = _queue(self)
+    self.creq_id = self.hwrm("ring_alloc", ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_NQ, page_tbl_addr=self.creq["base"],
+      page_size=12, page_tbl_depth=self.creq["level"], length=16, int_mode=bnxt.RING_ALLOC_REQ_INT_MODE_MSIX).ring_id
+
+    self.cmdq = _queue(self)
+    self.doorbell(self.creq_id, bnxt.DBC_DBC_TYPE_NQ_ARM, 0, 0)
+    init = bnxt.struct_cmdq_init(cmdq_pbl=self.cmdq["base"], creq_ring_id=self.creq_id,
+                                 cmdq_size_cmdq_lvl=16 << bnxt.CMDQ_INIT_CMDQ_SIZE_SFT)
+
+    System.memory_barrier()
+    for i, w in enumerate(memoryview(bytearray(bytes(init))).cast('I')): self.bar0[bnxt.RCFW_COMM_BASE_OFFSET // 4 + i] = w
+
+    _, p = self.pci_dev.alloc_sysmem(0x1000)
+    self.rcfw("initialize_fw", stat_ctx_id=self.hwrm("stat_ctx_alloc", stats_dma_addr=p[0], stats_dma_length=176).stat_ctx_id,
+      flags=bnxt.CMDQ_INITIALIZE_FW_FLAGS_HW_REQUESTER_RETX_SUPPORTED)
+
+    # RoCE notification ring: never armed or serviced, but CQ and L2 ring allocation require one
+    nq = _queue(self)
+    self.nq_id = self.hwrm("ring_alloc", ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_NQ, page_tbl_addr=nq["base"],
+      page_size=12, page_tbl_depth=nq["level"], length=16, logical_id=1, int_mode=bnxt.RING_ALLOC_REQ_INT_MODE_MSIX).ring_id
+
+  def rcfw(self, name, timeout_ms=20000, **fields):
+    req_t, resp_t = getattr(bnxt, f"struct_cmdq_{name}"), getattr(bnxt, f"struct_creq_{name}_resp")
+    op = getattr(bnxt, f"CMDQ_BASE_OPCODE_{name.upper()}")
+    data = bytes(req_t(opcode=op, cmd_size=(slots := ceildiv(ctypes.sizeof(req_t), 16)), **fields)).ljust(slots * 16, b'\0')
+    for i in range(slots): _qwrite(self.cmdq, self.cmdq["prod"] + i, data[i * 16:(i + 1) * 16])
+
+    self.cmdq["prod"] += slots
+    prod = self.cmdq["prod"] & 0xffff
+    if self.rcfw_first: prod, self.rcfw_first = prod | 1 << bnxt.FIRMWARE_FIRST_FLAG, False
+
+    System.memory_barrier()
+
+    self.bar0[(bnxt.RCFW_COMM_BASE_OFFSET + bnxt.RCFW_PF_VF_COMM_PROD_OFFSET) // 4] = prod
+    self.bar0[(bnxt.RCFW_COMM_BASE_OFFSET + bnxt.RCFW_COMM_TRIG_OFFSET) // 4] = bnxt.RCFW_CMDQ_TRIG_VAL
+
+    def poll():
+      h = bnxt.struct_creq_base.from_buffer_copy(bytes(_qread(self.creq, self.creq["cons"])))
+      return bool(h.v & bnxt.CREQ_BASE_V) != bool((self.creq["cons"] // 16) & 1)
+    wait_cond(poll, timeout_ms=timeout_ms, msg=f"RCFW {name}")
+
+    ret = resp_t.from_buffer_copy(bytes(_qread(self.creq, self.creq["cons"])))
+    self.creq["cons"] += 1
+
+    # NQ_ARM also publishes the CREQ consumer index, which is what frees ring space for the next command
+    self.doorbell(self.creq_id, bnxt.DBC_DBC_TYPE_NQ_ARM, self.creq["cons"] & 15, (self.creq["cons"] // 16) & 1)
+    assert ret.status == 0, f"RCFW {name}: {ret.status}"
+
+    if BNXT_DEBUG >= 1: print(f"bnxt {self.devfmt}: rcfw {name} xid={getattr(ret, 'xid', 0):#x}")
+    return ret
+
+  def doorbell(self, xid, typ, index, epoch):
+    System.memory_barrier()
+    self.db[self.db_off // 8] = db_value(xid, typ, index, epoch)
+
+  # L2 receive path, required for RoCE ingress even though no ethernet receive buffers are posted
+  def _open_l2(self):
+    cq = _queue(self)
+    ci = self.hwrm("ring_alloc", enables=bnxt.RING_ALLOC_REQ_ENABLES_NQ_RING_ID_VALID, ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_L2_CMPL,
+      page_tbl_addr=cq["base"], page_size=12, page_tbl_depth=cq["level"], length=16, nq_ring_id=self.nq_id).ring_id
+    rx = _queue(self)
+    ri = self.hwrm("ring_alloc", enables=bnxt.RING_ALLOC_REQ_ENABLES_NQ_RING_ID_VALID |
+      bnxt.RING_ALLOC_REQ_ENABLES_RX_BUF_SIZE_VALID, ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_RX, page_tbl_addr=rx["base"],
+      page_size=12, page_tbl_depth=rx["level"], length=16, rx_buf_size=640, nq_ring_id=self.nq_id).ring_id
+    vi = self.hwrm("vnic_alloc").vnic_id
+    self.hwrm("vnic_cfg", enables=bnxt.VNIC_CFG_REQ_ENABLES_MRU | bnxt.VNIC_CFG_REQ_ENABLES_DEFAULT_RX_RING_ID |
+      bnxt.VNIC_CFG_REQ_ENABLES_DEFAULT_CMPL_RING_ID, vnic_id=vi, mru=9018,
+      default_rx_ring_id=ri, default_cmpl_ring_id=ci)
+    self.hwrm("cfa_l2_filter_alloc", flags=bnxt.CFA_L2_FILTER_ALLOC_REQ_FLAGS_PATH_RX,
+      enables=bnxt.CFA_L2_FILTER_ALLOC_REQ_ENABLES_L2_ADDR | bnxt.CFA_L2_FILTER_ALLOC_REQ_ENABLES_L2_ADDR_MASK |
+      bnxt.CFA_L2_FILTER_ALLOC_REQ_ENABLES_DST_ID, l2_addr=tuple(self.mac.to_bytes(6, 'big')), l2_addr_mask=(0xff,) * 6, dst_id=vi)
+
+  def register_mem(self, paddrs:list[int], size:int, log_page_size:int=12) -> int:
+    level, base = _pbl(self, paddrs[:ceildiv(size, 1 << log_page_size)])
+    return self.rcfw("register_mr", flags=bnxt.CMDQ_REGISTER_MR_FLAGS_ALLOC_MR,
+      log2_pg_size_lvl=level << bnxt.CMDQ_REGISTER_MR_LVL_SFT | log_page_size << bnxt.CMDQ_REGISTER_MR_LOG2_PG_SIZE_SFT,
+      access=bnxt.CMDQ_REGISTER_MR_ACCESS_LOCAL_WRITE | bnxt.CMDQ_REGISTER_MR_ACCESS_REMOTE_WRITE,
+      log2_pbl_pg_size=12, pbl=base, va=paddrs[0], mr_size=size).xid
+
+class BNXTQP:
+  def __init__(self, dev:BNXTDev):
+    self.dev, self.sq_psn, self.msn = dev, 0, 0
+
+    self.cqq = _queue(dev, ctypes.sizeof(bnxt.struct_cq_base))
+    self.cq_id = dev.rcfw("create_cq", cq_size=16, pbl=self.cqq["base"],
+                          pg_size_lvl=self.cqq["level"], cq_fco_cnq_id=dev.nq_id).xid
+
+    self.sq = _queue(dev, aux=True)
+    self.qpn = dev.rcfw("create_qp", type=bnxt.CMDQ_CREATE_QP_TYPE_RC,
+      sq_size=16, sq_fwo_sq_sge=1, scq_cid=self.cq_id, rcq_cid=self.cq_id,
+      sq_pbl=self.sq["base"], sq_pg_size_sq_lvl=self.sq["level"]).xid
+    self.qp_op(1, BNXT_INIT_MASK, access=BNXT_ACCESS, pkey=0xffff)
+
+  def qp_op(self, state, mask, network_type=0, **fields):
+    self.dev.rcfw("modify_qp", qp_cid=self.qpn, modify_mask=mask,
+      network_type_en_sqd_async_notify_new_state=state | network_type, **fields)
+
+  def connect(self, qpn:int, gid:bytes, mac:int):
+    network_type = bnxt.CMDQ_MODIFY_QP_NETWORK_TYPE_ROCEV2_IPV4
+    dgid = (ctypes.c_uint32 * 4)(*(int.from_bytes(gid[i:i + 4], 'little') for i in (0, 4, 8, 12)))
+    dmac = (ctypes.c_uint16 * 3)(*(int.from_bytes(mac.to_bytes(6, 'big')[i:i + 2], 'little') for i in (0, 2, 4)))
+
+    self.qp_op(2, BNXT_RTR_MASK, network_type=network_type, qp_type=bnxt.CMDQ_MODIFY_QP_QP_TYPE_RC, access=BNXT_ACCESS,
+      pkey=0xffff, dgid=dgid, sgid_index=self.dev.gid_id, hop_limit=64, dest_mac=dmac,
+      path_mtu_pingpong_push_enable=bnxt.CMDQ_MODIFY_QP_PATH_MTU_MTU_1024, max_dest_rd_atomic=4,
+      dest_qp_id=qpn)
+    self.qp_op(3, BNXT_RTS_MASK, network_type=network_type, qp_type=bnxt.CMDQ_MODIFY_QP_QP_TYPE_RC, access=BNXT_ACCESS,
+      max_rd_atomic=1)
+
+    if BNXT_DEBUG >= 1: print(f"bnxt: QP {self.qpn:#x} connected (remote={qpn:#x})")
+
+  def _poll(self, timeout):
+    def poll():
+      base = bnxt.struct_cq_base.from_buffer_copy(bytes(_qread(self.cqq, self.cqq["cons"])))
+      return bool(base.cqe_type_toggle & bnxt.CQ_BASE_TOGGLE) == (not bool((self.cqq["cons"] // 16) & 1))
+    wait_cond(poll, timeout_ms=timeout, msg="BNXT CQ")
+    raw = bytes(_qread(self.cqq, self.cqq["cons"]))
+    self.cqq["cons"] += 1
+    self.dev.doorbell(self.cq_id, bnxt.DBC_DBC_TYPE_CQ, self.cqq["cons"] & 15, (self.cqq["cons"] // 16) & 1)
+    return raw
+
+  def rdma_write(self, rva, rkey, lva, lkey, size, timeout_ms=20000):
+    start = self.sq["prod"] & 15
+    hdr = bytes(bnxt.struct_sq_rdma_hdr(wqe_type=bnxt.SQ_RDMA_HDR_WQE_TYPE_WRITE_WQE,
+                flags=bnxt.SQ_SEND_FLAGS_SIGNAL_COMP, wqe_size=3, length=size, remote_va=rva, remote_key=rkey))
+    for i, data in enumerate((hdr[:16], hdr[16:32], bytes(bnxt.struct_sq_sge(va_or_pa=lva, l_key=lkey, size=size)))):
+      _qwrite(self.sq, start + i, data)
+    nxt = (self.sq_psn + max(1, ceildiv(size, 1024))) & 0xffffff
+    value = start << bnxt.SQ_MSN_SEARCH_START_IDX_SFT | nxt << bnxt.SQ_MSN_SEARCH_NEXT_PSN_SFT | self.sq_psn
+    _qwrite(self.sq, self.msn, struct.pack("<Q", value), aux=True)
+
+    self.msn, self.sq_psn, self.sq["prod"] = (self.msn + 1) % 128, nxt, self.sq["prod"] + 3
+    self.dev.doorbell(self.qpn, bnxt.DBC_DBC_TYPE_SQ, self.sq["prod"] & 15, (self.sq["prod"] // 16) & 1)
+    cqe = bnxt.struct_cq_req.from_buffer_copy(self._poll(timeout_ms))
+    assert cqe.status == 0

extra/bnxt_driver/connect.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python3
+"""Send and validate one RDMA WRITE between two Broadcom BNXT hosts.
+
+This follows ``extra/mlx_driver/connect.py``: sync the driver, start the remote
+endpoint over SSH, exchange QP/GID/MAC/MR metadata, move both RC QPs to RTS,
+write bytes into the remote MR, and verify the bytes on the remote host.
+
+Both PCI functions must be unbound from bnxt_en/bnxt_re first.
+"""
+import json
+import os
+import subprocess
+import sys
+from typing import Any, IO
+
+TINYGRAD = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
+sys.path.insert(0, TINYGRAD)
+
+from extra.bnxt_driver.bnxtdev import BNXTDev, BNXTQP
+from tinygrad.runtime.support.system import PCIDevice
+
+REMOTE_HOST = os.getenv("REMOTE_HOST", "192.168.52.213")
+REMOTE_USER = os.getenv("REMOTE_USER", "nimlgen")
+LOCAL_PCI = os.getenv("BNXT_PCI", "0000:41:00.0")
+REMOTE_PCI = os.getenv("REMOTE_PCI", "0000:41:00.0")
+LOCAL_IP = os.getenv("LOCAL_IP", "10.0.200.5")
+REMOTE_IP = os.getenv("REMOTE_IP", "10.0.200.6")
+MESSAGE = os.getenv("RDMA_MESSAGE", "Test message, rdma works!").encode()
+REMOTE = f"{REMOTE_USER}@{REMOTE_HOST}"
+SSH = ["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=10", "-o", "StrictHostKeyChecking=accept-new", REMOTE]
+SYNC_FILES = ("tinygrad/runtime/autogen/bnxt.py", "tinygrad/runtime/support/system.py",
+              "extra/bnxt_driver/bnxtdev.py", "extra/bnxt_driver/connect.py")
+
+def read_json(stream:IO[str], what:str) -> dict[str, Any]:
+  for line in iter(stream.readline, ""):
+    print(f"  [remote] {line}", end="")
+    try: value = json.loads(line)
+    except json.JSONDecodeError: continue
+    if isinstance(value, dict): return value
+  raise RuntimeError(f"remote exited before publishing {what}")
+
+def wait_line(stream:IO[str], text:str) -> str:
+  for line in iter(stream.readline, ""):
+    print(f"  [remote] {line}", end="")
+    if text in line: return line
+  raise RuntimeError(f"remote exited before reporting {text!r}")
+
+def send_line(stream:IO[str], value:str|dict[str, Any]):
+  stream.write((json.dumps(value) if isinstance(value, dict) else value) + "\n")
+  stream.flush()
+
+def qp_info(dev:BNXTDev, qp:BNXTQP) -> dict[str, Any]:
+  return {"qpn":qp.qpn, "mac":dev.mac.to_bytes(6, "big").hex(), "gid":dev.local_gid.hex()}
+
+def server():
+  dev = BNXTDev(PCIDevice("bnxt", os.getenv("BNXT_PCI", "0000:41:00.0")), ip=os.getenv("BNXT_IP", REMOTE_IP))
+  qp = BNXTQP(dev)
+  print(json.dumps(qp_info(dev, qp)), flush=True)
+
+  peer = json.loads(sys.stdin.readline())
+  qp.connect(peer["qpn"], bytes.fromhex(peer["gid"]), int(peer["mac"], 16))
+  print("connected", flush=True)
+
+  target, target_paddrs = dev.pci_dev.alloc_sysmem(0x1000)
+  target[:0x1000] = bytes(0x1000)
+  rkey = dev.register_mem(target_paddrs, 0x1000)
+  print(json.dumps({"target_addr":target_paddrs[0], "rkey":rkey}), flush=True)
+
+  assert sys.stdin.readline().strip() == "done"
+  received = bytes(target).rstrip(b"\0")
+  print(f"AS TEXT: {received.decode(errors='replace')!r}", flush=True)
+  print(json.dumps({"data":received.hex()}), flush=True)
+
+def sync_remote():
+  if os.getenv("SYNC", "1") == "0": return
+  print("syncing BNXT driver to remote")
+  subprocess.run(["rsync", "-azR", *SYNC_FILES, f"{REMOTE}:~/tinygrad/"], cwd=TINYGRAD, check=True)
+
+def start_remote() -> subprocess.Popen[str]:
+  print("booting remote")
+  command = (f"cd ~/tinygrad && sudo env PYTHONPATH=. PYTHONUNBUFFERED=1 BNXT_DEBUG={os.getenv('BNXT_DEBUG', '0')} "
+             f"BNXT_PCI={REMOTE_PCI} BNXT_IP={REMOTE_IP} python3 extra/bnxt_driver/connect.py --server")
+  return subprocess.Popen(SSH + [command], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=sys.stderr, text=True)
+
+def client():
+  assert 0 < len(MESSAGE) <= 0x1000
+  sync_remote()
+  remote = start_remote()
+  assert remote.stdin is not None and remote.stdout is not None
+  remote_info = read_json(remote.stdout, "QP information")
+  print("booting local")
+  dev = BNXTDev(PCIDevice("bnxt", LOCAL_PCI), ip=LOCAL_IP)
+  qp = BNXTQP(dev)
+
+  send_line(remote.stdin, qp_info(dev, qp))
+  wait_line(remote.stdout, "connected")
+  qp.connect(remote_info["qpn"], bytes.fromhex(remote_info["gid"]), int(remote_info["mac"], 16))
+  print("both QPs in RTS")
+
+  remote_target = read_json(remote.stdout, "MR information")
+  source, source_paddrs = dev.pci_dev.alloc_sysmem(0x1000)
+  source[:len(MESSAGE)] = MESSAGE
+  lkey = dev.register_mem(source_paddrs, 0x1000)
+  print(f"RDMA WRITE {len(MESSAGE)}B to remote phys 0x{remote_target['target_addr']:x}")
+  qp.rdma_write(remote_target["target_addr"], remote_target["rkey"], source_paddrs[0], lkey, len(MESSAGE))
+
+  send_line(remote.stdin, "done")
+  wait_line(remote.stdout, "AS TEXT")
+  result = read_json(remote.stdout, "RDMA result")
+  assert bytes.fromhex(result["data"]) == MESSAGE
+  print("RDMA WRITE data verified")
+
+  remote.stdin.close()
+  assert remote.wait() == 0
+  print("RDMA WRITE test complete")
+
+if __name__ == "__main__":
+  server() if "--server" in sys.argv else client()

extra/bnxt_driver/loopback.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python3
+"""Local BNXT RoCEv2 RDMA WRITE loopback using the firmware's PHY loopback mode.
+
+The kernel bnxt_en/bnxt_re modules must be unloaded first.
+
+  sudo PYTHONPATH=. BNXT_PCI=0000:41:00.0 BNXT_IP=10.0.200.5 python3 extra/bnxt_driver/loopback.py
+"""
+import os
+import sys
+import time
+
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
+
+from extra.bnxt_driver.bnxtdev import BNXTDev, BNXTQP
+from tinygrad.runtime.autogen import bnxt
+from tinygrad.runtime.support.system import PCIDevice
+
+BUF_SIZE = 0x1000
+BNXT_PCI = os.getenv("BNXT_PCI", "0000:41:00.0")
+BNXT_IP = os.getenv("BNXT_IP", "10.0.200.5")
+
+if __name__ == "__main__":
+  print(f"[init] BNXT at {BNXT_PCI}")
+  dev = BNXTDev(PCIDevice("bnxt", BNXT_PCI), ip=BNXT_IP)
+  tx_qp, rx_qp = BNXTQP(dev), BNXTQP(dev)
+  print(f"[init] loopback-connect TX QP 0x{tx_qp.qpn:x} <-> RX QP 0x{rx_qp.qpn:x}")
+  tx_qp.connect(rx_qp.qpn, dev.local_gid, dev.mac)
+  rx_qp.connect(tx_qp.qpn, dev.local_gid, dev.mac)
+
+  src, src_paddrs = dev.pci_dev.alloc_sysmem(BUF_SIZE)
+  dst, dst_paddrs = dev.pci_dev.alloc_sysmem(BUF_SIZE)
+  message = b"Hello from BNXT RoCE PHY loopback!"
+  src[:BUF_SIZE], dst[:BUF_SIZE] = bytes(BUF_SIZE), bytes(BUF_SIZE)
+  src[:len(message)] = message
+  lkey = dev.register_mem(src_paddrs, BUF_SIZE)
+  rkey = dev.register_mem(dst_paddrs, BUF_SIZE)
+
+  print("[loopback] enabling local PHY loopback")
+  dev.hwrm("port_phy_cfg", port_id=dev.port_id, enables=bnxt.PORT_PHY_CFG_REQ_ENABLES_LPBK, lpbk=bnxt.PORT_PHY_CFG_REQ_LPBK_LOCAL)
+  time.sleep(1)
+  tx_qp.rdma_write(dst_paddrs[0], rkey, src_paddrs[0], lkey, len(message))
+  got = bytes(dst[:len(message)])
+  print(f"[result] {got!r}")
+  assert got == message
+  print("BNXT RoCE PHY loopback RDMA WRITE passed")
+  dev.hwrm("port_phy_cfg", port_id=dev.port_id, enables=bnxt.PORT_PHY_CFG_REQ_ENABLES_LPBK, lpbk=bnxt.PORT_PHY_CFG_REQ_LPBK_NONE)

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, src=(gtab.index(g, row, j), val), arg=(atomic_str, dtypes.void))
     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/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, src=(amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=(atomic_arg, dtypes.void))
   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, src=(o.base.index(0),), arg=("", dtypes.void)) 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()

test/amd/test_emu2_pcode.py

@@ -152,7 +152,7 @@ class TestDSPcodePatterns(unittest.TestCase):
 
   def test_global_atomic_add_f32_parsing(self):
     """Test GLOBAL_ATOMIC_ADD_F32 keeps memory values in float dtype."""
-    vmem = UOp.param(2, dtypes.uint32, (1024,))
+    vmem = UOp.param(2, dtypes.uint32, 1024)
     srcs = {
       'ADDR': UOp.const(0, dtypes.uint64),
       'DATA': UOp.const(0x3f800000, dtypes.uint32),
@@ -183,7 +183,7 @@ class TestDSPcodePatterns(unittest.TestCase):
   def test_mem_read_parsing(self):
     """Test MEM[addr].type read expression parsing."""
     # Create a mock LDS buffer
-    lds = UOp.param(3, dtypes.uint32, (16384,))
+    lds = UOp.param(3, dtypes.uint32, 16384)
     addr = UOp.const(0, dtypes.uint32)
     vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(0, dtypes.uint32)}
 
@@ -218,7 +218,7 @@ class TestDSPcodePatterns(unittest.TestCase):
     pcode = PCODE.get(DSOp.DS_LOAD_2ADDR_B32)
     self.assertIsNotNone(pcode)
     assert pcode is not None
-    lds = UOp.param(3, dtypes.uint32, (16384,))
+    lds = UOp.param(3, dtypes.uint32, 16384)
     srcs = {
       'ADDR': UOp.const(0, dtypes.uint32),
       'OFFSET0': UOp.const(0, dtypes.uint32),
@@ -299,7 +299,7 @@ class TestConcatWidthParsing(unittest.TestCase):
       self.assertIs(parsed.simplify(), UOp.const(expected, dtypes.uint32))
 
   def test_permlane64_wave64_pcode_indices(self):
-    vgpr = UOp.param(0, dtypes.uint32, (256,))
+    vgpr = UOp.param(0, dtypes.uint32, 256)
     srcs = {
       'SRC0': UOp.const(0, dtypes.uint32),
       'VDST': UOp.const(1, dtypes.uint32),
@@ -330,7 +330,7 @@ class TestAllPcode(unittest.TestCase):
   def _make_srcs(self):
     """Create dummy source variables for pcode parsing."""
     u32, u64 = lambda v=0: UOp.const(v, dtypes.uint32), lambda v=0: UOp.const(v, dtypes.uint64)
-    lds = UOp.param(3, dtypes.uint32, (16384,))
+    lds = UOp.param(3, dtypes.uint32, 16384)
     return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(),
             'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(),
             'VCC': u64(), 'VCCZ': u32(), 'EXEC': u64(), 'EXEC_LO': u32(), 'EXECZ': u32(), 'SCC': u32(),

test/backend/test_asm_gemm.py

@@ -188,7 +188,7 @@ class TestMXFP4(unittest.TestCase):
     M, N, K = getenv("M", 16384), getenv("N", 4096), getenv("K", 14336)
     a = Tensor.empty(M, K, dtype=dtypes.bfloat16)
     b = Tensor.empty(N, K, dtype=dtypes.bfloat16)
-    asm_gemm(a, b.T, mxfp4=True).realize()
+    for _ in range(getenv("CNT", 1)): asm_gemm(a, b.T, mxfp4=True).realize()
 
 # test the Asm GEMM with Llama shapes, only run on the real machine for speed
 

test/backend/test_assign.py

@@ -1,10 +1,10 @@
 #!/usr/bin/env python
 import unittest
 import numpy as np
-from tinygrad import dtypes, Tensor, TinyJit, GlobalCounters, Variable
+from tinygrad import Device, dtypes, Tensor, TinyJit, GlobalCounters, Variable
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.helpers import temp, DEV, Context
-from test.helpers import assert_kernel_count
+from test.helpers import assert_kernel_count, needs_second_gpu
 
 N = 200  # has to be bigger than the cache to fail
 
@@ -1079,5 +1079,80 @@ class TestBatchNormRunningStats(unittest.TestCase):
     with Context(TRAINING=1): bn(x).realize()
     self.assertTrue(bn.running_mean.uop.base.is_realized)
 
+class TestMultiAssign(unittest.TestCase):
+  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
+
+  @needs_second_gpu
+  def setUp(self): pass
+
+  def test_multi_assign_realized(self):
+    out = Tensor.zeros(4).shard(self.device, 0).contiguous().realize()
+    ones = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
+    out.assign(ones).realize()
+    self.assertListEqual(out.tolist(), [1,1,1,1])
+
+  def test_multi_assign_unrealized(self):
+    out = Tensor.zeros(4).contiguous().realize().shard(self.device, 0)
+    ones = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
+    out.assign(ones).realize()
+    self.assertListEqual(out.tolist(), [1,1,1,1])
+
+  def test_multi_assign_both_unrealized(self):
+    out = Tensor.zeros(4).contiguous().realize().shard(self.device, 0)
+    ones = Tensor.ones(4).contiguous().realize().shard(self.device, 0)
+    out.assign(ones).realize()
+    self.assertListEqual(out.tolist(), [1,1,1,1])
+
+  def test_multi_assign_scalar(self):
+    out = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
+    out.assign(0).realize()
+    self.assertListEqual(out.tolist(), [0,0,0,0])
+
+  def test_multi_assign_const_like(self):
+    out = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
+    out.assign(out.const_like(7)).realize()
+    self.assertListEqual(out.tolist(), [7,7,7,7])
+
+  def test_multi_assign_piece(self):
+    out = Tensor.zeros(4,4).shard(self.device, 0).contiguous().realize()
+    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
+    out[:, 2:3].assign(ones).realize()
+    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
+
+  def test_multi_assign_piece_noncontig(self):
+    out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0).realize()
+    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
+    out[:, 2:3].assign(ones).realize()
+    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
+
+  @unittest.expectedFailure
+  def test_multi_assign_piece_unrealized(self):
+    out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0)
+    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
+    out[:, 2:3].assign(ones).realize()
+    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
+
+  def test_multi_assign_var_offset(self):
+    out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0).realize()
+    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
+    vi = Variable("i", 0, 3).bind(2)
+    out[:, vi:vi+1].assign(ones).realize()
+    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
+
+  def test_multi_assign_var_offset_jit_none(self): self.test_multi_assign_var_offset_jit(None)
+  def test_multi_assign_var_offset_jit(self, shard_axis=0):
+    out = Tensor.zeros(4,6).contiguous().realize().shard(self.device, shard_axis).realize()
+    ones = Tensor.ones(4,1).shard(self.device, shard_axis).contiguous().realize()
+
+    @TinyJit
+    def f(out:Tensor, vi):
+      out[:, vi:vi+1].assign(ones).realize()
+      ones.assign(ones+1).realize()
+
+    vi = Variable("i", 0, 5)
+    for i in range(1,5):
+      GlobalCounters.reset()
+      f(out, vi.bind(i))
+    self.assertListEqual(out.tolist(), [[0,1,2,3,4,0]]*4)
 if __name__ == "__main__":
   unittest.main()

test/backend/test_isel.py

@@ -46,7 +46,7 @@ class TestIselX86(unittest.TestCase):
   # complex address is [base + index*scale + displacement]
   def test_complex_address(self):
     a = UOp.variable("a", 0, 0, dtypes.int32)
-    load = UOp.param(0, dtypes.int32, (16,)).index(a + UOp.cconst(1, dtypes.int32)).load()
+    load = UOp.param(0, dtypes.int32, 16).index(a + UOp.cconst(1, dtypes.int32)).load()
     n = self.isel_rewrite(load)
     # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32
     self.assertTrue(n.src[2].dtype is dtypes.int8 and n.src[2].src[0].op is Ops.CONST and n.src[2].src[0].val == 4)

test/backend/test_linearizer_dumb.py

@@ -11,16 +11,16 @@ from tinygrad.codegen import to_program
 class TestLinearizerFailure(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
   def test_failure_beam_mnist(self):
-    c0 = UOp.param(0, dtypes.uchar, (4014080,))
+    c0 = UOp.param(0, dtypes.uchar, 4014080)
     c1 = UOp.range(UOp.const(512), 0, AxisType.GLOBAL)
     c2 = UOp.range(UOp.const(784), 1, AxisType.GLOBAL)
     c3 = UOp.range(UOp.const(10), 3, AxisType.GLOBAL)
-    c4 = UOp.param(1, dtypes.int, (512,))
+    c4 = UOp.param(1, dtypes.int, 512)
     c5 = c4.index(c1.valid(UOp.const(True)))
     c6 = UOp.range(UOp.const(6000), 1004, AxisType.REDUCE)
     c7 = UOp.range(UOp.const(3750), 2006, AxisType.REDUCE)
     c8 = UOp.range(UOp.const(16), 2007, AxisType.GROUP_REDUCE)
-    c9 = UOp.param(2, dtypes.uchar, (47040000,))
+    c9 = UOp.param(2, dtypes.uchar, 47040000)
     c10 = c9.index((((c3*UOp.const(4704000))+c2)+(c6*UOp.const(784))).valid(UOp.const(True)))
     c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(6000))+c6)+((c7*UOp.const(16))+c8)).alu(Ops.CMPLT, UOp.const(59999)).where(UOp.const(0).cast(dtypes.int), UOp.const(1).cast(dtypes.int)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(-1).cast(dtypes.int))).where(UOp.const(0).cast(dtypes.uchar), c10).reduce(c6, arg=Ops.ADD)
     c12 = c0.index((((c1*UOp.const(7840))+(c2*UOp.const(10)))+c3).valid(UOp.const(True))).store(c11).end(c1, c2, c3)

test/backend/test_multitensor.py

@@ -58,6 +58,11 @@ class TestMultiTensor(unittest.TestCase):
     assert X.uop.ended_ranges == X.uop.src[1:]
     (X + X).realize()
 
+  def test_shard_invalids_contiguous(self):
+    # every store is Invalid, so none of them should become a (empty) kernel
+    t = Tensor.invalids(8).shard(devices_2, axis=0).contiguous()
+    self.assertEqual(len([c for c in t.schedule_linear().src if c.src[0].op is Ops.SINK]), 1)
+
   @unittest.expectedFailure # TODO: fix
   def test_shard_empty(self):
     GlobalCounters.reset()

test/backend/test_ops.py

@@ -6,6 +6,7 @@ from tinygrad.helpers import getenv, DEBUG, DEV, IMAGE, Context
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.renderer.nir import NIRRenderer
+from tinygrad.renderer.isa.x86 import X86Renderer
 
 TINY_BACKEND = getenv("TINY_BACKEND")
 if TINY_BACKEND:
@@ -808,6 +809,8 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: tor^0x1337, lambda: ten^0x1337, forward_only=True)
     helper_test_op([], lambda: 0x1337^tor, lambda: 0x1337^ten, forward_only=True)
 
+  # TODO: x86 PARAM dtype fails SPEC=2
+  @Context(SPEC=1 if isinstance(Device[Device.DEFAULT].renderer, X86Renderer) else 2)
   def test_and(self):
     data = [[1,-8,1],[32,1,6]]
     tor = torch.tensor(data, dtype=torch.int)
@@ -1807,9 +1810,10 @@ class TestOps(unittest.TestCase):
       helper_test_op([()], lambda x: torch.nn.functional.hardtanh(x, -val, val), lambda x: x.hardtanh(-val, val), grad_atol=1e-6)
   def test_asinh(self):
     helper_test_op([(45,65)], lambda x: x.asinh(), grad_atol=1e-6)
-    # TODO: this one has larger tol?
-    helper_test_op([(45,65)], lambda x: x.asinh(), atol=1e-2, rtol=2e-2, grad_rtol=2e-2, low=-300, high=-297)
+    helper_test_op([(45,65)], lambda x: x.asinh(), grad_atol=1e-6, low=-300, high=-297)
     helper_test_op([(45,65)], lambda x: x.asinh(), grad_atol=1e-6, low=300, high=303)
+    helper_test_op([(45,65)], lambda x: x.asinh(), grad_atol=1e-6, low=-1e10, high=-1e9)
+    helper_test_op(None, lambda x: x.asinh(), grad_atol=1e-6, vals=[[-1.0, 0.0, 1.0]])
   def test_acosh(self):
     helper_test_op([(45,65)], lambda x: x.acosh(), grad_atol=1e-6)
     helper_test_op([(45,65)], lambda x: x.acosh(), grad_atol=1e-3, grad_rtol=1e-2, low=-300, high=-297)

test/backend/test_rangeify.py

@@ -1,6 +1,6 @@
 import unittest
-from tinygrad import Tensor, nn, Device, dtypes, Variable
-from tinygrad.helpers import Context, GlobalCounters, getenv, PCONTIG, DEBUG
+from tinygrad import Tensor, Device, dtypes, Variable
+from tinygrad.helpers import Context, GlobalCounters, getenv, DEBUG
 from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops, UOp
 from tinygrad.codegen.opt import OptOps, Opt
 from tinygrad.renderer.ptx import PTXRenderer
@@ -14,7 +14,7 @@ class TestDoubleMatmul(unittest.TestCase):
       self.ref = (self.a @ self.b @ self.c).realize()
 
   def _test(self, opts):
-    with Context(PCONTIG=2, DEBUG=max(2, DEBUG.value)):
+    with Context(DEBUG=max(2, DEBUG.value)):
       out = (self.a @ self.b @ self.c).contiguous(arg=opts).realize()
 
     with Context(DEBUG=0):
@@ -88,16 +88,15 @@ class TestRangeifyEdgeCase(unittest.TestCase):
     res = Tensor.cat(a, c, dim=0)
     self.assertEqual(res.numpy()[-1, :16].tolist(), [512] * 16)
 
-  def test_pcontig_multi_gather(self):
+  def test_multi_gather(self):
     # regression test: local bufferize must have device set for const_like to work
-    with Context(PCONTIG=2):
-      # NOTE: with uint type, this will become a long and fail on WEBGPU
-      forest = Tensor(list(range(8)), dtype='int')
-      idx = Tensor([0, 0], dtype='int')
-      node_val = forest.gather(0, idx)
-      idx2 = idx * 2 + 1
-      node_val2 = forest.gather(0, idx2)
-      result = (node_val + node_val2).numpy()
+    # NOTE: with uint type, this will become a long and fail on WEBGPU
+    forest = Tensor(list(range(8)), dtype='int')
+    idx = Tensor([0, 0], dtype='int')
+    node_val = forest.gather(0, idx)
+    idx2 = idx * 2 + 1
+    node_val2 = forest.gather(0, idx2)
+    result = (node_val + node_val2).numpy()
     self.assertEqual(result.tolist(), [1, 1])
 
 if getenv("BIG") > 2:
@@ -118,65 +117,6 @@ def fa():
   GlobalCounters.reset()
   return q.scaled_dot_product_attention(k, v)
 
-def fa_bw():
-  Tensor.manual_seed(1337)
-  with Context(DEBUG=0):
-    q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
-    attn_output = nn.Linear(HEADS*EMB, HEADS*EMB, bias=False)
-    attn_output.weight.realize()
-    target = Tensor.rand(BS, SEQLEN, HEADS*EMB).contiguous().realize()
-
-  GlobalCounters.reset()
-  attn = q.scaled_dot_product_attention(k, v).contiguous().contiguous_backward()
-  attn = attn.transpose(1, 2).reshape(BS, SEQLEN, -1)
-  out = attn_output(attn)
-  loss = (out - target).square().mean()
-  loss.backward()
-  #ret = [out, Tensor.stack(q.grad, k.grad, v.grad, dim=-1)]
-  #ret = [out, Tensor.stack(q.grad, k.grad, dim=-1), v.grad]
-  ret = [out, q.grad, k.grad, v.grad]
-  Tensor.realize(*ret)
-  return ret
-
-@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "broken in LVP and PTX")
-class TestPcontig(unittest.TestCase):
-  def test_flash_attention_bw(self):
-    with Context(PCONTIG=max(2, PCONTIG.value), DEBUG=2):
-      grads = fa_bw()
-      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
-
-    with Context(PCONTIG=0, DEBUG=2):
-      cmp_grads = fa_bw()
-      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
-
-    with Context(DEBUG=0):
-      mses = [((x-y)**2).sum().item() for x,y in zip(grads, cmp_grads)]
-    mse = sum(mses)
-    print(f"mse: {mse}")
-    self.assertLessEqual(mse, 1e-6)
-
-  def test_flash_attention(self, opts=None):
-    with Context(PCONTIG=2, DEBUG=max(2, DEBUG.value)):
-      ret = fa().realize() if opts is None else fa().contiguous(arg=opts).realize()
-      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
-    with Context(DEBUG=2):
-      cmp = fa().realize()
-      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
-    with Context(DEBUG=0):
-      mse = ((cmp-ret)**2).sum().item()
-    print(f"mse: {mse}")
-    self.assertLessEqual(mse, 1e-6)
-
-  def test_flash_attention_opt(self):
-    opts = ()
-    # columns in top matrix
-    opts += (Opt(OptOps.UPCAST, 0, 4),)
-    # columns in bottom matrix
-    opts += (Opt(OptOps.UPCAST, 3, 4),)
-    # rows in all the matrix
-    opts += (Opt(OptOps.UPCAST, 4, 4),)
-    self.test_flash_attention(opts)
-
 # contiguous + reduce can support ranges?
 
 @unittest.skip("pm_rangeify no longer exists. test this in a different way")

test/backend/test_renderer_failures.py

@@ -23,8 +23,8 @@ def _test_uop_result(inputs:list[Tensor], sink:UOp, local_size=None):
 
 def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
   dtype = alu_src_uops[0].dtype
-  a = UOp.param(0, dtype, (1,))
-  b = UOp.param(1, dtype, (1,))
+  a = UOp.param(0, dtype, 1)
+  b = UOp.param(1, dtype, 1)
   idx = UOp.const(0)
   ld = b.index(idx).load()
   alu = ld.alu(alu_op, *alu_src_uops)
@@ -34,7 +34,7 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
 class TestRendererFailures(unittest.TestCase):
   @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
   def test_gated_store_with_alu(self):
-    a = UOp.param(0, dtypes.int, (4,))
+    a = UOp.param(0, dtypes.int, 4)
     gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
     gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0.valid(gate_alu)), UOp.const(1).cast(dtypes.int)))
     sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo())
@@ -43,7 +43,7 @@ class TestRendererFailures(unittest.TestCase):
 
   @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
   def test_gated_store_with_alu_2d(self):
-    a = UOp.param(0, dtypes.int, (8,))
+    a = UOp.param(0, dtypes.int, 8)
     gate_alu_0 = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
     gate_alu_1 = (lidx1:=UOp.special(2, 'lidx1')).ne(0)
     gated_alu_store = UOp(Ops.STORE, src=(a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(1).cast(dtypes.int)))
@@ -78,7 +78,7 @@ 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,))
+    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())
@@ -93,9 +93,9 @@ class TestWGSLFailures(unittest.TestCase):
 
   # WGSL has a specific select(alt, val, gate) ternary operator instead of gate?val:alt
   def test_gated_load(self):
-    a = UOp.param(0, dtypes.int, (4,))
-    b = UOp.param(1, dtypes.int, (4,))
-    c = UOp.param(2, dtypes.int, (4,))
+    a = UOp.param(0, dtypes.int, 4)
+    b = UOp.param(1, dtypes.int, 4)
+    c = UOp.param(2, dtypes.int, 4)
     lidx0 = UOp.special(4, "lidx0")
     gate = lidx0.ne(0)
     alt = c.index(lidx0).load()
@@ -110,7 +110,7 @@ class TestWGSLFailures(unittest.TestCase):
 class TestPTXFailures(unittest.TestCase):
   @unittest.skip("INDEX can only have a gate ALU parent, not an IF")
   def test_gated_store_with_if(self):
-    a = UOp.param(0, dtypes.int, (4,))
+    a = UOp.param(0, dtypes.int, 4)
     gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
     val = UOp.const(1).cast(dtypes.int)
     if_uop = UOp(Ops.IF, src=(gate_alu,))

test/backend/test_schedule.py

@@ -147,7 +147,7 @@ class TestSchedule(unittest.TestCase):
     devs = ("CPU:0", "CPU:1")
     x = Tensor.ones(2, device="CPU").shard(devs, axis=0).realize()
     out = (x.sum()*2).reshape(1).to("CPU")
-    run_linear(*check_schedule(out, 5))
+    run_linear(*check_schedule(out, 3))
     np.testing.assert_equal(out.numpy(), [4.])
 
 class TestLimitBufs(unittest.TestCase):
@@ -365,16 +365,6 @@ class TestCopyFolding(unittest.TestCase):
     b = a.to("CPU")
     self.assertListEqual(b.tolist(), [2.])
 
-  def test_copy_to_same_device(self):
-    a = Tensor.empty(4).uop
-    b = a.copy_to_device(a.device)
-    check_schedule(b, 1, filter_sink=False) # TODO: 0?
-
-  def test_copy_to_same_device_alt(self):
-    a = Tensor.empty(4, 4).uop
-    b = a.copy_to_device(a.device)
-    check_schedule(b, 1, filter_sink=False) # TODO: 0?
-
   def test_copy_to_same_device_sched(self):
     a = Tensor.ones(4).contiguous().realize().uop.buf_uop
     t = Tensor(a.copy_to_device(a.device))
@@ -384,14 +374,6 @@ class TestCopyFolding(unittest.TestCase):
     assert t.uop.is_realized, f"didn't realize Tensor {t}"
     self.assertListEqual(t.tolist(), [1.,1.,1.,1.])
 
-  @unittest.skip("same-device copies are no-ops")
-  def test_self_assign_same_device_copy(self):
-    a = Tensor.ones(4, 4).contiguous().realize()
-    # use copy_to_device to bypass Tensor.to() shortcircuit and force a real same-device COPY in the graph
-    a.assign(Tensor(a.uop.copy_to_device(a.device), a.device))
-    run_linear(*check_schedule(a, 2, filter_sink=False))
-    self.assertListEqual(a.tolist(), [[1.]*4]*4)
-
   def test_clone(self):
     a = Tensor.empty(4)
     check_schedule(a.clone(), 1, filter_sink=False)

test/backend/test_setitem.py

@@ -301,6 +301,14 @@ class TestSetitem(unittest.TestCase):
     self.assertListEqual(z[2:5].tolist(), [2, 2, 2])
     self.assertListEqual(z[6:7].tolist(), [3])
 
+class TestAssignBitcast(unittest.TestCase):
+  def test_assign_through_bitcast(self):
+    # the dest is unrealized, so callify cannot fold the BITCAST into a buffer view and the STORE keeps a
+    # BITCAST dest; the bitcast has to move to the value side or the store never reaches the buffer
+    a = Tensor.full((4,), 1.0, dtype=dtypes.float32).contiguous()
+    a.bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.uint32)).realize()
+    np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])
+
 class TestWithGrad(unittest.TestCase):
   def test_basic_setitem_works(self):
     z = Tensor.rand(8, 8)

test/backend/test_uops.py

@@ -20,7 +20,7 @@ def run_uops(uops_list:list[UOp], bufs:list[Buffer]):
 
 def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
   if op is Ops.CONST: uops.append(UOp.const(arg).cast(dtype))
-  elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, shape=(1,)))
+  elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, 1))
   else: uops.append(UOp(op, dtype, tuple(src), arg))
   return uops[-1]
 
@@ -62,7 +62,7 @@ def _test_uops_result(output_dtype, uops, res):
 class TestBitcastBufferView(unittest.TestCase):
   @Context(SPEC=2)
   def test_render(self):
-    buf = UOp.param(0, dtypes.uint32, (4,))
+    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)
@@ -71,7 +71,7 @@ class TestBitcastBufferView(unittest.TestCase):
   @Context(SPEC=2)
   def test_load(self):
     val = 0x1122334455667788
-    src, out = UOp.param(0, dtypes.uint32, (4,)), UOp.param(1, dtypes.uint64, (1,))
+    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])
@@ -80,7 +80,7 @@ class TestBitcastBufferView(unittest.TestCase):
   @Context(SPEC=2)
   def test_store(self):
     val = 0x1122334455667788
-    dst = UOp.param(0, dtypes.uint32, (6,))
+    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])
@@ -249,8 +249,8 @@ class TestLocalAccess(unittest.TestCase):
 @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "This only tests assembly backends")
 class TestAssembly(unittest.TestCase):
   def test_bitshift_left(self):
-    g1 = UOp.param(0, dtypes.int32, shape=(3,))
-    out = UOp.param(1, dtypes.int32, shape=(2,))
+    g1 = UOp.param(0, dtypes.int32, 3)
+    out = UOp.param(1, dtypes.int32, 2)
     c1 = UOp.const(2)
     c2 = UOp.const(3)
     l1 = g1.index(c1)
@@ -278,7 +278,7 @@ class TestAssembly(unittest.TestCase):
     self.assertGreaterEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4)
 
   def test_mulacc_shl(self):
-    g1 = UOp.param(0, dtypes.int32, shape=(2,))
+    g1 = UOp.param(0, dtypes.int32, 2)
     c1 = UOp.const(0)
     c2 = UOp.const(1)
     expr = g1.index(c1) * UOp.const(4096) + g1.index(c2)
@@ -287,7 +287,7 @@ class TestAssembly(unittest.TestCase):
     self.assertIn(Ops.MULACC, [x.op for x in uops])
 
   def test_use_cmpeq(self):
-    g = UOp.param(0, dtypes.uint32, shape=(8,))
+    g = UOp.param(0, dtypes.uint32, 8)
     c = UOp.const(7)
     comp = g.index(c).ne(c).ne(True)
     uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer)

test/backend/test_wait_loop.py

@@ -43,7 +43,7 @@ def nested_loop_kernel(C:UOp) -> UOp:
   return C[0].store(i[0].load()).sink(arg=KernelInfo(name="nested_loop", opts_to_apply=()))
 
 def wait_ext_kernel() -> UOp:
-  sig = UOp.param(0, dtypes.int, (1,), volatile=True)
+  sig = UOp.param(0, dtypes.int, 1, volatile=True)
   l = UOp.loop(0)
   v = sig.after(l)[0].load()
   e = v.end(l, v < 1)

test/external/external_benchmark_op_conv.py

@@ -12,7 +12,7 @@ from tinygrad.dtype import Invalid
 # PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
 
 def vision_conv_143():
-  c0 = UOp.param(0, dtypes.half, shape=(16, 1024, 4))
+  c0 = UOp.param_from_shape(0, (16, 1024, 4), dtypes.half)
   c2 = UOp.range(32, 3)
   c5 = UOp.range(128, 4)
   c8 = UOp.range(16, 2)
@@ -22,13 +22,13 @@ def vision_conv_143():
   c26 = UOp.range(7, 1, AxisType.REDUCE)
   c27 = c2*2+c26
   c32 = ((c27<3)!=True)&(c27<67)
-  c34 = UOp.param(1, dtypes.half, shape=(32, 1024, 4))
+  c34 = UOp.param_from_shape(1, (32, 1024, 4), dtypes.half)
   c38 = c5//2
   c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(Invalid, dtypes.weakint))
   c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
-  c49 = UOp.param(2, dtypes.half, shape=(64, 49, 4))
+  c49 = UOp.param_from_shape(2, (64, 49, 4), dtypes.half)
   c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp.param(3, dtypes.float, (128,))
+  c63 = UOp.param(3, dtypes.float, 128)
   c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
   c67 = c0.index((c2*128+c5+c8*4096)).store(c65).end(c8, c2, c5)
 
@@ -38,7 +38,7 @@ def vision_conv_143():
   return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
 
 def vision_conv_153():
-  c0 = UOp.param(0, dtypes.half, shape=(8, 1024, 4))
+  c0 = UOp.param_from_shape(0, (8, 1024, 4), dtypes.half)
   c2 = UOp.range(16, 3)
   c5 = UOp.range(256, 4)
   c8 = UOp.range(8, 2)
@@ -48,13 +48,13 @@ def vision_conv_153():
   c26 = UOp.range(7, 1, AxisType.REDUCE)
   c27 = c2*2+c26
   c32 = ((c27<3)!=True)&(c27<35)
-  c34 = UOp.param(1, dtypes.half, shape=(16, 1024, 4))
+  c34 = UOp.param_from_shape(1, (16, 1024, 4), dtypes.half)
   c38 = c5//2
   c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(Invalid, dtypes.weakint))
   c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
-  c49 = UOp.param(2, dtypes.half, shape=(128, 49, 4))
+  c49 = UOp.param_from_shape(2, (128, 49, 4), dtypes.half)
   c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp.param(3, dtypes.float, (256,))
+  c63 = UOp.param(3, dtypes.float, 256)
   c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
   c67 = c0.index((c2*256+c5+c8*4096)).store(c65).end(c8, c2, c5)
 
@@ -64,16 +64,16 @@ def vision_conv_153():
   return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
 
 def dm_conv_172():
-  c0 = UOp.param(0, dtypes.half, shape=(1, 240, 4))
+  c0 = UOp.param_from_shape(0, (1, 240, 4), dtypes.half)
   c2 = UOp.range(960, 4)
-  c5 = UOp.param(1, dtypes.half, shape=(8, 384, 4))
+  c5 = UOp.param_from_shape(1, (8, 384, 4), dtypes.half)
   c7 = UOp.range(32, 0, AxisType.REDUCE)
   c10 = UOp.range(4, 1, AxisType.REDUCE)
   c13 = UOp.range(12, 3, AxisType.REDUCE)
   c18 = UOp.range(8, 2, AxisType.REDUCE)
-  c23 = UOp.param(2, dtypes.half, shape=(240, 128, 4))
+  c23 = UOp.param_from_shape(2, (240, 128, 4), dtypes.half)
   c35 = c5.index((c7*4+c10+c13*128+c18*1536))*c23.index((c10*4+c2%4+c7*16+c2//4*512))
-  c37 = UOp.param(3, dtypes.float, (960,))
+  c37 = UOp.param(3, dtypes.float, 960)
   c39 = c35.reduce(c7, c10, arg=Ops.ADD)+c37.index(c2)
   c50 = (1.0+((c39+0.044708251953125*(c39*(c39*c39)))*-2.3021129851685216).exp2()).reciprocal()*c39
   c53 = c50.reduce(c18, c13, arg=Ops.ADD)*0.010416666666666666

test/external/external_test_gpu_crash.py

@@ -36,7 +36,7 @@ class TestGPUCrash(unittest.TestCase):
 
   def _run_insts(self, insts: list[Inst]):
     buf = UOp.new_buffer("AMD", 64, dtypes.uint8)
-    sink = UOp.sink(UOp.param(0, dtypes.uint8, (64,), device="AMD"), UOp.special(1, "lidx0"), arg=KernelInfo("test"))
+    sink = UOp.sink(UOp.param(0, dtypes.uint8, 64, device="AMD"), UOp.special(1, "lidx0"), arg=KernelInfo("test"))
     prg = UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=i) for i in insts))))
     run_linear(UOp(Ops.LINEAR, src=(prg.call(buf),)), wait=True)
 

test/helpers.py

@@ -122,7 +122,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple
   for buf_dt, data in inputs or []:
     bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
     allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
-  g = UOp.param(0, uop.dtype, (1,))
+  g = UOp.param(0, uop.dtype, 1)
   prg = to_program(UOp.store(g.index(UOp.const(0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
   prog = dev.runtime(prg.to_elf())
   prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs, vals=vals)

test/mockgpu/amd/amddriver.py

@@ -1,20 +1,15 @@
-import pathlib, re, ctypes, mmap, collections, functools, copy, os
+import ctypes, mmap, collections, functools, copy, os
 from tinygrad.runtime.autogen import kfd, amdgpu_drm, libc
 import tinygrad.runtime.autogen.am.am as am
 from tinygrad.helpers import from_mv
 from test.mockgpu.driver import VirtDriver, VirtFileDesc, TextFileDesc, DirFileDesc, VirtFile
 from test.mockgpu.amd.amdgpu import AMDGPU, gpu_props, GFX_TARGET_VERSION, MOCKGPU_ARCH
 
-def ioctls_from_header():
-  # hdrpy = (pathlib.Path(__file__).parent.parent.parent.parent / "tinygrad" / "runtime" / "autogen" / "kfd.py").read_text()
-  # pattern = r'# (AMDKFD_IOC_[A-Z0-9_]+)\s=\s_(IOW?R?).*\(( 0x[0-9a-fA-F]+) ,\s+struct\s([A-Za-z0-9_]+)\s+\)'
-  # matches = re.findall(pattern, hdrpy, re.MULTILINE)
-  hdr = (pathlib.Path(__file__).parent.parent.parent.parent / "extra" / "hip_gpu_driver" / "kfd_ioctl.h").read_text().replace("\\\n", "")
-  pattern = r'#define\s+(AMDKFD_IOC_[A-Z0-9_]+)\s+AMDKFD_(IOW?R?)\((0x[0-9a-fA-F]+),\s+struct\s([A-Za-z0-9_]+)\)'
-  matches = re.findall(pattern, hdr, re.MULTILINE)
-  return type("KFD_IOCTLS", (object, ), {name: int(nr, 0x10) for name, _, nr, _ in matches}), \
-         {int(nr, 0x10): getattr(kfd, "struct_"+sname, None) for name, idir, nr, sname in matches}
-kfd_ioctls, kfd_headers = ioctls_from_header()
+def _ioctl_nr(ioctl: functools.partial) -> int: return ioctl.args[2]
+
+kfd_ioctl_info = {
+  _ioctl_nr(ioctl): (name, ioctl.args[3]) for name, ioctl in vars(kfd).items()
+  if name.startswith("AMDKFD_IOC_") and isinstance(ioctl, functools.partial)}
 
 class KFDFileDesc(VirtFileDesc):
   def __init__(self, fd, driver):
@@ -116,42 +111,44 @@ class AMDDriver(VirtDriver):
 
   def kfd_ioctl(self, req, argp):
     nr = req & 0xFF
-    struct = kfd_headers[nr].from_address(argp)
+    if nr not in kfd_ioctl_info: raise RuntimeError(f"unknown kfd ioctl, {nr} unknown")
+    name, struct_type = kfd_ioctl_info[nr]
+    struct = struct_type.from_address(argp)
 
-    if nr == kfd_ioctls.AMDKFD_IOC_ACQUIRE_VM: pass
-    elif nr == kfd_ioctls.AMDKFD_IOC_RUNTIME_ENABLE: pass
-    elif nr == kfd_ioctls.AMDKFD_IOC_GET_VERSION:
+    if nr == _ioctl_nr(kfd.AMDKFD_IOC_ACQUIRE_VM): pass
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_RUNTIME_ENABLE): pass
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_GET_VERSION):
       struct.major_version = 1
       struct.minor_version = 14
-    elif nr == kfd_ioctls.AMDKFD_IOC_ALLOC_MEMORY_OF_GPU:
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_ALLOC_MEMORY_OF_GPU):
       if struct.gpu_id not in self.gpus: return -1
       struct.handle = self._alloc_handle()
       self.object_by_handle[struct.handle] = copy.deepcopy(struct) # save memory struct to know what mem it is
       # Track signal memory (uncached + coherent) - progress queues when written to
       if struct.flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED:
         self.track_address(struct.va_addr, struct.va_addr + struct.size, lambda mv,off: None, lambda mv, off: self._emulate_execute())
-    elif nr == kfd_ioctls.AMDKFD_IOC_FREE_MEMORY_OF_GPU:
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU):
       self.object_by_handle.pop(struct.handle)
-    elif nr == kfd_ioctls.AMDKFD_IOC_MAP_MEMORY_TO_GPU:
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU):
       dev_ids = (ctypes.c_int32 * struct.n_devices).from_address(struct.device_ids_array_ptr)
       for i in range(struct.n_devices):
         gpu = self.gpus[dev_ids[i]]
         mem_obj = self.object_by_handle[struct.handle]
         gpu.map_range(mem_obj.va_addr, mem_obj.size)
         struct.n_success = i + 1
-    elif nr == kfd_ioctls.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU:
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU):
       dev_ids = (ctypes.c_int32 * struct.n_devices).from_address(struct.device_ids_array_ptr)
       for i in range(struct.n_devices):
         gpu = self.gpus[dev_ids[i]]
         mem_obj = self.object_by_handle[struct.handle]
         gpu.unmap_range(mem_obj.va_addr, mem_obj.size)
         struct.n_success = i + 1
-    elif nr == kfd_ioctls.AMDKFD_IOC_CREATE_EVENT:
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_CREATE_EVENT):
       struct.event_slot_index = self._alloc_next_event_slot()
       struct.event_id = struct.event_slot_index
 
       if struct.event_type == kfd.KFD_IOC_EVENT_MEMORY: self.mmu_event_ids.append(struct.event_id)
-    elif nr == kfd_ioctls.AMDKFD_IOC_CREATE_QUEUE:
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_CREATE_QUEUE):
       gpu = self.gpus[struct.gpu_id]
       if struct.queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
         gpu.add_sdma_queue(struct.ring_base_address, struct.ring_size, struct.read_pointer_address, struct.write_pointer_address)
@@ -162,7 +159,7 @@ class AMDDriver(VirtDriver):
       # Track writes to doorbell, calling callback
       struct.doorbell_offset = self._alloc_doorbell(struct.gpu_id)
       self.track_address(struct.doorbell_offset, struct.doorbell_offset + 8, lambda mv,off: None, lambda mv, off: self._emulate_execute())
-    elif nr == kfd_ioctls.AMDKFD_IOC_WAIT_EVENTS:
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_WAIT_EVENTS):
       evs = (kfd.struct_kfd_event_data * struct.num_events).from_address(struct.events_ptr)
       for ev in evs:
         if ev.event_id in self.mmu_event_ids and "MOCKGPU_EMU_FAULTADDR" in os.environ:
@@ -170,11 +167,7 @@ class AMDDriver(VirtDriver):
           ev.memory_exception_data.va = int(os.environ["MOCKGPU_EMU_FAULTADDR"], 16)
           ev.memory_exception_data.failure.NotPresent = 1
     else:
-      name = "unknown"
-      for k,v in kfd_ioctls.__dict__.items():
-        if nr == v: name = k
-      assert False, f"unknown kfd ioctl, {nr} {name}"
-      exit(1)
+      raise RuntimeError(f"unsupported kfd ioctl, {nr} {name}")
     return 0
 
   def _emulate_execute(self):

test/mockgpu/amd/emu.py

@@ -321,10 +321,10 @@ def _int_clamp(op_name: str, srcs: dict) -> UOp | None:
 class _Ctx:
   """Context for instruction compilation - holds buffers and helpers."""
   __slots__ = ('inst_size', 'dyn_fields', '_axis_id', 'wave_size', 'vgpr', 'accvgpr')
-  sgpr = UOp.param(0, dtypes.uint32, (SGPR_COUNT,))
-  vmem = UOp.param(2, dtypes.uint32, (1 << 46,))
-  lds = UOp.param(3, dtypes.uint32, (16384,))
-  scratch = UOp.param(4, dtypes.uint8, (1 << 30,))
+  sgpr = UOp.param(0, dtypes.uint32, SGPR_COUNT)
+  vmem = UOp.param(2, dtypes.uint32, 1 << 46)
+  lds = UOp.param(3, dtypes.uint32, 16384)
+  scratch = UOp.param(4, dtypes.uint8, 1 << 30)
   # Cache PARAM UOps by wave_size so all _Ctx instances with same wave_size share identical UOp references
   _vgpr_cache: dict[int, UOp] = {}
   _accvgpr_cache: dict[int, UOp] = {}
@@ -332,10 +332,10 @@ class _Ctx:
   def __init__(self, inst_size: int, wave_size: int = 32):
     self.inst_size, self._axis_id, self.wave_size = inst_size, 0, wave_size
     self.dyn_fields: list[tuple[int, int]] = []  # (lo, hi) of fields read dynamically
-    if wave_size not in _Ctx._vgpr_cache: _Ctx._vgpr_cache[wave_size] = UOp.param(1, dtypes.uint32, (256 * wave_size,))
+    if wave_size not in _Ctx._vgpr_cache: _Ctx._vgpr_cache[wave_size] = UOp.param(1, dtypes.uint32, 256 * wave_size)
     self.vgpr = _Ctx._vgpr_cache[wave_size]
     if wave_size == 64:
-      if wave_size not in _Ctx._accvgpr_cache: _Ctx._accvgpr_cache[wave_size] = UOp.param(5, dtypes.uint32, (256 * wave_size,))
+      if wave_size not in _Ctx._accvgpr_cache: _Ctx._accvgpr_cache[wave_size] = UOp.param(5, dtypes.uint32, 256 * wave_size)
       self.accvgpr = _Ctx._accvgpr_cache[wave_size]
     else:
       self.accvgpr = self.vgpr

test/null/test_gpudims.py

@@ -107,7 +107,7 @@ class TestGroupedDims(unittest.TestCase):
 
   def test_global_prod_max(self):
     g, l = UOp.range(256, 0, AxisType.GLOBAL), UOp.range(256, 1, AxisType.LOCAL)
-    sink = UOp.param(0, dtypes.float, (512,)).index(g + l).store(UOp.const(1.0)).end(g, l).sink(arg=KernelInfo())
+    sink = UOp.param(0, dtypes.float, 512).index(g + l).store(UOp.const(1.0)).end(g, l).sink(arg=KernelInfo())
     class R(Renderer): global_max, local_max, global_prod_max = (256, 256, 256), (128, 128, 128), (128, 128, 128)
     specials = [u for u in add_gpudims(R(Target()), sink).toposort() if u.op is Ops.SPECIAL]
     self.assertGreater(len([s for s in specials if "lidx" in s.arg]), 1)

test/null/test_linearizer_failures.py

@@ -7,14 +7,14 @@ from tinygrad.codegen import to_program
 
 class TestLinearizerFailures(unittest.TestCase):
   def test_fail_1(self):
-    c0 = UOp.param(0, dtypes.float, (64,))
+    c0 = UOp.param(0, dtypes.float, 64)
     c1 = UOp.range(UOp.const(2), 1, AxisType.WEAK)
     c2 = UOp.range(UOp.const(32), 2, AxisType.WEAK)
     c3 = ((c1*UOp.const(32))+c2)
-    c4 = UOp.param(1, dtypes.float, (163840,))
+    c4 = UOp.param(1, dtypes.float, 163840)
     c5 = UOp.range(UOp.const(2560), 0, AxisType.REDUCE)
     c6 = c4.index(((((((c5//UOp.const(8))%UOp.const(8))*UOp.const(8))+(c5%UOp.const(8)))+(((c2*UOp.const(40))+(c5//UOp.const(64)))*UOp.const(64)))+(c1*UOp.const(81920))))
-    c7 = UOp.param(2, dtypes.float, (64,))
+    c7 = UOp.param(2, dtypes.float, 64)
     c8 = c7.index(c3)
     c9 = ((((c6+(c8*UOp.const(-1.0)))*(c6+(c8*UOp.const(-1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(0.000390625))+UOp.const(1e-05)).sqrt().reciprocal()
     c10 = c0.index(c3).store(c9).end(c1, c2)

test/null/test_llm_tokenizer.py

@@ -95,6 +95,26 @@ class TestLLMTokenizer(unittest.TestCase):
     self.assertEqual(template.end_turn(), "[/INST]")
     self.assertEqual(template.role("assistant"), "")
 
+  def test_tekken_gpt4o_split(self):
+    split = {p: SimpleTokenizer({}, {}, p)._split_to_word.findall for p in ("tekken", "gpt-4o")}
+    shared = {
+      "HelloWorld": ["Hello", "World"],
+      " ÜNICODE": [" ÜNICODE"],  # Ü: non-ascii upper joins the run
+      "é café": ["é", " café"],  # first é is e + U+0301 combining acute (NFD)
+      "เพื่อน วิ": ["เพื่อน", " วิ"],  # thai vowel marks stay in the word
+      "a/b\r\n x": ["a", "/b", "\r\n", " x"],  # punct tail eats /
+    }
+    for s, want in shared.items():
+      self.assertEqual(split["tekken"](s), want, f"tekken {s!r}")
+      self.assertEqual(split["gpt-4o"](s), want, f"gpt-4o {s!r}")
+    differ = [
+      ("12345", list("12345"), ["123", "45"]),  # digits: tekken single, o200k groups {1,3}
+      ("it's I'M don'T", ["it", "'s", " I", "'M", " don", "'T"], ["it's", " I'M", " don'T"]),  # contraction: o200k inline suffix
+    ]
+    for s, tk, go in differ:
+      self.assertEqual(split["tekken"](s), tk, f"tekken {s!r}")
+      self.assertEqual(split["gpt-4o"](s), go, f"gpt-4o {s!r}")
+
   def test_stream_decoder(self):
     """stream_decoder buffers incomplete UTF-8: token 25677 has 3/4 of emoji, token 138 completes it."""
     bs = [*range(33, 127), *range(161, 173), *range(174, 256)]

test/null/test_simplify_valid_idx.py

@@ -15,10 +15,10 @@ def simplify_valid_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move
 def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+indexing_simplify, name="simplify_image_idx")
 
 def get_gated_load_uop(valid:UOp, idx:UOp):
-  return UOp.param(0, dtypes.float, (1024,)).index(idx.valid(valid)).load()
+  return UOp.param(0, dtypes.float, 1024).index(idx.valid(valid)).load()
 
 def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
-  return UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)).load()
+  return UOp.param_from_shape(0, image_shape, dtypes.float).index(idx[1].valid(valid), idx[0].valid(valid)).load()
 
 def Special(expr, nmax): return UOp.special(nmax, expr)
 def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax, param=True)
@@ -505,7 +505,7 @@ class TestDropTrueGate(unittest.TestCase):
     from tinygrad.codegen.late.coalesce import indexing_simplify
     from tinygrad.uop.ops import graph_rewrite
     from tinygrad.uop.symbolic import sym
-    buf = UOp.param(0, dtypes.int, (1,))
+    buf = UOp.param(0, dtypes.int, 1)
     idx = UOp.const(0)
     true_gate = UOp.const(True)
     index_with_gate = buf.index(idx.valid(true_gate))
@@ -551,7 +551,7 @@ class TestRangeShrink(unittest.TestCase):
     # one load guards r < 4, but another load uses r without a gate -> no shrink
     r = Range(0, 204)
     load1 = get_gated_load_uop(r < UOp.const(4), r)
-    load2 = UOp.param(1, dtypes.float, (204,)).index(r).load()
+    load2 = UOp.param(1, dtypes.float, 204).index(r).load()
     ranges = self.get_ranges(UOp.sink(load1, load2))
     self.assert_range_end(ranges, 204)
 
@@ -575,7 +575,7 @@ class TestRangeShrink(unittest.TestCase):
     from tinygrad.dtype import Invalid
     r = Range(0, 204)
     x = (r < 4).where(UOp.const(1.0), Invalid)
-    ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink())
+    ranges = self.get_ranges(UOp.param(0, dtypes.float, 204).index(r).store((r < 4).where(x, Invalid)).sink())
     self.assert_range_end(ranges, 4)
 
   def test_range_shrink_store_where_invalid_flipped(self):
@@ -583,7 +583,7 @@ class TestRangeShrink(unittest.TestCase):
     from tinygrad.dtype import Invalid
     r = Range(0, 204)
     x = (r < 4).where(UOp.const(1.0), Invalid)
-    ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink())
+    ranges = self.get_ranges(UOp.param(0, dtypes.float, 204).index(r).store((r >= 4).where(Invalid, x)).sink())
     self.assert_range_end(ranges, 4)
 
 if __name__ == '__main__':

test/null/test_transcendental_helpers.py

@@ -9,7 +9,7 @@ class TestTranscendentalFunctions(unittest.TestCase):
   def test_payne_hanek_reduction(self):
     # TODO: Test constant input when constant folding is fixed (or maybe test both variants)
     # Load input value from a buffer to prevent constant folding
-    input_buf = UOp.param(1, dtypes.double, (1,))
+    input_buf = UOp.param(1, dtypes.double, 1)
     loaded_value = input_buf.index(UOp.const(0)).load()
     def eval_payne_hanek_reduction(v:float) -> tuple[float, int]:
       return tuple(eval_uop(u, [(dtypes.float64, [v])]) for u in payne_hanek_reduction(loaded_value))

test/null/test_uop_graph.py

@@ -202,6 +202,11 @@ class TestUOpGraph(unittest.TestCase):
     invalid_lane_mul = next(u for u in out.src[0].toposort() if u.op is Ops.MUL)
     self.assertIs(invalid_lane_mul.dtype, dtypes.bool)
 
+  def test_devectorize_zero_sized_scalar_expand(self):
+    from tinygrad.codegen import devectorizer2
+    expanded = UOp.const(1.0).reshape(1, 1).expand(0, 3)
+    self.assertEqual(graph_rewrite(expanded, devectorizer2).shape, (0, 3))
+
   def test_gep_vec_const_fold(self):
     for vec_size in [2, 4, 8]:
       consts = [UOp.const(float(i), dtypes.float) for i in range(vec_size)]
@@ -209,8 +214,8 @@ class TestUOpGraph(unittest.TestCase):
       for i, const in enumerate(consts): self.assertIs(vec.index(i), const)
 
   def test_cast_alu_fold(self):
-    d0 = UOp.param(0, dtypes.bool, (1,))
-    d1 = UOp.param(1, dtypes.int, (1,))
+    d0 = UOp.param(0, dtypes.bool, 1)
+    d1 = UOp.param(1, dtypes.int, 1)
     idx = UOp.const(0)
     ld = d1.index(idx)
     alu = (ld<1).cast(dtypes.bool)
@@ -219,8 +224,8 @@ class TestUOpGraph(unittest.TestCase):
     self.assertEqual(len([x for x in uops if x.op is Ops.CAST and x.src[0].op is not Ops.CONST]), 0)
 
   def test_double_cast_fold(self):
-    d0 = UOp.param(0, dtypes.float, (1,))
-    d1 = UOp.param(1, dtypes.int, (1,))
+    d0 = UOp.param(0, dtypes.float, 1)
+    d1 = UOp.param(1, dtypes.int, 1)
     idx = UOp.const(0, dtypes.int)
     ld = d1.index(idx)
     alu = ld.cast(dtypes.float).cast(dtypes.float)
@@ -238,7 +243,7 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_bitcast_to_same_dtype_fold(self):
     for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
-      d0 = UOp.param(0, dt, (1,))
+      d0 = UOp.param(0, dt, 1)
       v = d0.index(UOp.const(0))
       uops = to_uops_list([v.bitcast(dt)])
       self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST and x.dtype is dt]), 0, f"dtype = {dt}")
@@ -250,10 +255,10 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_where_on_gated_load_fold(self):
     ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.long, (100,))
+    d0 = UOp.param(0, dtypes.long, 100)
     ld = d0.index(ridx0.valid(ridx0<50))
     w = (ridx0<50).where(ld, 5)
-    out = UOp.param(1, dtypes.long, (100,))
+    out = UOp.param(1, dtypes.long, 100)
     uops = to_uops_list([out.index(ridx0).store(w)])
     expected = full_rewrite(UOp.const(5, dtypes.long).sink()).src[0]
     for u in uops:
@@ -262,7 +267,7 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_where_on_gated_load_folds_swapped_branches(self):
     ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.long, (100,))
+    d0 = UOp.param(0, dtypes.long, 100)
     ld = d0.index(ridx0.valid((ridx0<50).logical_not()))
     w = (ridx0<50).where(5, ld)
     uops = to_uops_list([w])
@@ -273,11 +278,11 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_where_on_gated_load_with_cast(self):
     ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.int, (100,))
+    d0 = UOp.param(0, dtypes.int, 100)
     gate_idx = ridx0.valid((ridx0<50))
     ld = d0.index(gate_idx).cast(dtypes.float)
     w = (ridx0<50).where(ld, 5.0)
-    out = UOp.param(1, dtypes.float, (100,))
+    out = UOp.param(1, dtypes.float, 100)
     uops = to_uops_list([out.index(ridx0).store(w)])
     expected = full_rewrite(UOp.const(5, dtypes.int).sink()).src[0]
     for u in uops:
@@ -286,27 +291,27 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_where_on_casted_gated_load_extra_cond(self):
     ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.float, (100,))
+    d0 = UOp.param(0, dtypes.float, 100)
     ld = d0.index(ridx0.valid(ridx0<50))
     w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(0.0)).cast(dtypes.half)
-    out = UOp.param(1, dtypes.half, (100,))
+    out = UOp.param(1, dtypes.half, 100)
     uops = to_uops_list([out.index(ridx0).store(w)])
     for u in uops:
       assert u.op is not Ops.WHERE
 
   def test_where_on_casted_gated_load_extra_cond_swapped(self):
     ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.float, (100,))
+    d0 = UOp.param(0, dtypes.float, 100)
     ld = d0.index(ridx0.valid(ridx0<50))
     w = ((ridx0<50) & (ridx0>30)).where(UOp.const(0.0), ld).cast(dtypes.half)
-    out = UOp.param(1, dtypes.half, (100,))
+    out = UOp.param(1, dtypes.half, 100)
     uops = to_uops_list([out.index(ridx0).store(w)])
     for u in uops:
       assert u.op is not Ops.WHERE
 
   def test_where_in_store_becomes_gate(self):
     ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.long, (100,))
+    d0 = UOp.param(0, dtypes.long, 100)
     idx = d0.index(ridx0)
     ld = idx.load()
     val = (ridx0<50).where(5, ld)
@@ -320,14 +325,14 @@ class TestUOpGraph(unittest.TestCase):
   def test_load_idx_becomes_int(self):
     # mnist indexing with split reduceop
     # Make sure we are not doign math on the loaded index, which would promote it to long
-    c0 = UOp.param(0, dtypes.uchar, (128000,))
+    c0 = UOp.param(0, dtypes.uchar, 128000)
     c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK)
     c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK)
-    c3 = UOp.param(1, dtypes.int, (512,))
+    c3 = UOp.param(1, dtypes.int, 512)
     c4 = c3.index(c1)
     c5 = UOp.range(UOp.const(240), 0, AxisType.REDUCE)
     c6 = ((c2*UOp.const(240))+c5)
-    c7 = UOp.param(2, dtypes.uchar, (60000,))
+    c7 = UOp.param(2, dtypes.uchar, 60000)
     c8 = c7.index(c6)
     c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD)
     c10 = c0.index(((c1*UOp.const(250))+c2)).store(c9).end(c1, c2)
@@ -337,14 +342,14 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_load_idx_no_math_on_loaded(self):
     # test the (x+y)<c pattern where x has loads - we shouldn't do math on loaded indices
-    c0 = UOp.param(0, dtypes.uchar, (128000,))
+    c0 = UOp.param(0, dtypes.uchar, 128000)
     c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK)
     c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK)
-    c3 = UOp.param(1, dtypes.int, (512,))
+    c3 = UOp.param(1, dtypes.int, 512)
     c4 = c3.index(c1)  # c4 is a load
     c5 = UOp.range(UOp.const(240), 0, AxisType.REDUCE)
     c6 = ((c2*UOp.const(240))+c5)
-    c7 = UOp.param(2, dtypes.uchar, (60000,))
+    c7 = UOp.param(2, dtypes.uchar, 60000)
     c8 = c7.index(c6)
     # (loaded + range) < const pattern - loaded value shouldn't be promoted to long
     loaded_idx = c4.cast(dtypes.weakint)
@@ -356,9 +361,9 @@ class TestUOpGraph(unittest.TestCase):
       self.assertNotEqual(u.dtype, dtypes.long)
 
   def test_fold_gated_load(self):
-    glbl0 = UOp.param(0, dtypes.int, (1,))
-    glbl1 = UOp.param(1, dtypes.int, (1,))
-    glbl2 = UOp.param(2, dtypes.int, (1,))
+    glbl0 = UOp.param(0, dtypes.int, 1)
+    glbl1 = UOp.param(1, dtypes.int, 1)
+    glbl2 = UOp.param(2, dtypes.int, 1)
     idx = UOp.const(0)
     ld0 = glbl1.index(UOp.invalid())
     ld1 = glbl2.index(idx.valid(UOp.const(True)))
@@ -367,7 +372,7 @@ class TestUOpGraph(unittest.TestCase):
     self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 1)
 
   def test_fold_gated_load_local(self):
-    glbl0 = UOp.param(0, dtypes.int, (16,))
+    glbl0 = UOp.param(0, dtypes.int, 16)
     smem = UOp.placeholder((18,), dtypes.int, slot=0, addrspace=AddrSpace.LOCAL)
     lidx = UOp.special(16, "lidx0")
     st = smem.index(lidx).store(glbl0.index(lidx).load())
@@ -380,7 +385,7 @@ class TestUOpGraph(unittest.TestCase):
     self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 2)
 
   def test_fold_gated_store(self):
-    glbl = UOp.param(0, dtypes.int, (1,))
+    glbl = UOp.param(0, dtypes.int, 1)
     idx0 = UOp.const(0)
     val = UOp.const(42)
     st0 = glbl.index(UOp.invalid()).store(val)
@@ -420,7 +425,7 @@ class TestReduceCollapse(unittest.TestCase):
 
   def test_reduce_shapeless_const_unroll(self):
     """a REDUCE over a shapeless CONST (e.g. x*0 folded late in codegen) must collapse before the expander"""
-    out = UOp.param(0, dtypes.float, (1,))
+    out = UOp.param(0, dtypes.float, 1)
     red = UOp.const(3.0).cast(dtypes.float).reduce(UOp.range(4, 0, AxisType.UNROLL), arg=(Ops.ADD, 0))
     ast = UOp.sink(out.index(UOp.const(0)).store(red)).replace(arg=KernelInfo())
     uops = full_rewrite_to_sink(ast, Device["CPU"].renderer, optimize=False).toposort()
@@ -430,17 +435,18 @@ 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
-    src = UOp.param(0, dtypes.float, shape=(32, 4))
+    src = UOp.param_from_shape(0, (32, 4), dtypes.float)
     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")
     self.assertEqual(result.op, Ops.INDEX)
     self.assertIs(result.src[0], src)
     self.assertEqual(result.shape, (4,))
-    self.assertNotIn(Ops.RESHAPE, [u.op for u in result.toposort()])
+    # the only RESHAPE is src itself: the view of the flat param, the extra reshape was folded into the INDEX
+    self.assertEqual([u for u in result.toposort() if u.op is Ops.RESHAPE], [src])
 
   def test_pm_mops_partial_reshape_index_suffix_mismatch_does_nothing(self):
     from tinygrad.schedule.prepare import pm_mops
-    src = UOp.param(0, dtypes.float, shape=(2, 6))
+    src = UOp.param_from_shape(0, (2, 6), dtypes.float)
     result = graph_rewrite(src.reshape((2, 3, 2)).index(UOp.range(2, 0)), pm_mops, name="test")
     self.assertEqual(result.op, Ops.INDEX)
     self.assertEqual(result.src[0].op, Ops.RESHAPE)

test/null/test_uop_symbolic.py

@@ -5,7 +5,7 @@ import z3
 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.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load, symbolic_simple
 from tinygrad.uop.validate import uops_to_z3
 
 def check_uop_against_string(self, v:UOp, s:str):
@@ -448,10 +448,20 @@ class TestSymbolic(unittest.TestCase):
   def test_and_remove(self):
     self.helper_test_variable(uand([uconst(1), Variable("a", 0, 1)]), 0, 1, "a")
 
+  def test_zero_div_zero_bottom_up(self):
+    # codegen runs symbolic_simple bottom_up, so the 0/0 is rewritten before its consts fold.
+    # without the guard the unsound x/x -> 1 below it claims this one.
+    z = UOp.const(0.0)
+    self.assertTrue(math.isnan(graph_rewrite(z/z, symbolic_simple, bottom_up=True).arg))
+
   def test_masked_shr_fold(self):
     x = UOp.variable('x', 0, 255, dtype=dtypes.uint32, param=True)
     self.helper_test_variable((x & -4) >> 2, 0, 63, "(x>>2)")
 
+  def test_masked_idiv_fold(self):
+    x = UOp.variable('x', 0, 255, dtype=dtypes.uint32, param=True)
+    self.helper_test_variable((x & -4) // 4, 0, 63, "(x//4)")
+
   def test_bool_or_not_tautology(self):
     a = Variable("a", 0, 10)
     c = a<10
@@ -1382,7 +1392,7 @@ class TestInvalidIndex(unittest.TestCase):
 
   def test_gated_load_keeps_index_valid(self):
     # the load executes even on gated-off iterations: gated_given_valid must not erase its mask (PADTO OOB shape)
-    buf = UOp.param(0, dtypes.bool, (17,))
+    buf = UOp.param(0, dtypes.bool, 17)
     ridx = Variable("ridx", 0, 31)
     cond = ridx < 17
     load = buf.index(ridx.valid(cond))
@@ -1394,7 +1404,7 @@ class TestStoreLoadFolding(unittest.TestCase):
   """Tests for store(index, load(index)) -> NOOP rule. This rule matches patterns that EMERGE during simplification."""
   def test_store_load_folding(self):
     # store(idx, load(idx)) -> NOOP, including emergent patterns like store(idx, load(idx) + 0)
-    buf = UOp.param(0, dtypes.int, (1,))
+    buf = UOp.param(0, dtypes.int, 1)
     index = buf.index(UOp.const(0))
     # Direct: store(idx, load(idx)) -> NOOP
     self.assertEqual(graph_rewrite(index.store(index.load()), sym).op, Ops.NOOP)
@@ -1407,7 +1417,7 @@ class TestStoreLoadFolding(unittest.TestCase):
 
 class TestMoveWhereOnLoad(unittest.TestCase):
   def test_bool_index_preserves_dtype(self):
-    buf = UOp.param(0, dtypes.bool, (8,))
+    buf = UOp.param(0, dtypes.bool, 8)
     a = Variable("a", 0, 7)
     r = UOp.range(8, 0)
     # cond has a range that the rewrite can move into the valid: gate (a<4) goes into load valid
@@ -1455,12 +1465,17 @@ class TestGatedUopGivenValid(unittest.TestCase):
     self.assertEqual(idx, (r0 < 3).where(expected_vec, UOp.invalid()))
 
 class TestRangeSplitting(unittest.TestCase):
+  def test_end_preserves_constant_backedge(self):
+    loop, backedge = UOp.loop(0), UOp.const(False)
+    end = graph_rewrite(UOp(Ops.NOOP).end(loop, backedge), sym)
+    self.assertEqual(end.src, (UOp(Ops.NOOP), loop, backedge))
+
   def test_range_split_on_mod(self):
     # test that mark_range_mod splits RANGE(8) into RANGE(4)*2 + RANGE(2) when used with %2
     from tinygrad.codegen.simplify import pm_split_ranges, pm_flatten_range
     r0 = UOp.range(uconst(8), 0)
     # create a simple expression using the range with mod: store range%2 to a buffer
-    buf = UOp.param(0, dtypes.int, (1,))
+    buf = UOp.param(0, dtypes.int, 1)
     val = (r0 % uconst(2)).cast(dtypes.int)
     store = UOp(Ops.STORE, src=(buf.index(uconst(0)), val))
     sink = UOp(Ops.SINK, src=(UOp(Ops.END, src=(store, r0)),))

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))
     uop = x * y
     # TODO: these should be 0, but definitely should not be nan
     self.assertEqual(uop.vmin, -math.inf)
@@ -332,7 +332,7 @@ class TestVminVmaxVConst(unittest.TestCase):
 
   def test_vmin_vmax_vector_with_gep(self):
     # vmin and vmax for a vector constant of bool values
-    d1 = UOp.param(1, dtypes.int, (1,))
+    d1 = UOp.param(1, dtypes.int, 1)
     idx = UOp.const(0)
     val = UOp(Ops.LOAD, src=(d1.index(idx),))
     uop = (val // 32)

test/null/test_uops.py

@@ -55,11 +55,11 @@ class TestDTypeFromUOp(unittest.TestCase):
     self.assertEqual((scratch.dtype, next(u.dtype for u in scratch.uop.toposort() if u.op is Ops.BUFFER), next(u.dtype for u in scratch.uop.toposort()
       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.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)),
               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, ())
+              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)
     self.assertIs((out:=graph_rewrite(gate.where(value, UOp.invalid()), pm_remove_invalid)).src[2], UOp.const(0, dtypes.float))
     type_verify(out.sink(), spec_program)
 
@@ -73,7 +73,7 @@ class TestLowerIndexDtype(unittest.TestCase):
   def test_gated_shrink_lowers_to_selected_width(self):
     # coalesce builds gated SHRINKs for masked vectorized loads; lowering must resolve them at the
     # width the offset bounds select (this one needs long)
-    buf = UOp.param(0, dtypes.float, (2**31+64,))
+    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)
@@ -84,11 +84,10 @@ class TestLowerIndexDtype(unittest.TestCase):
 
   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)
+    self.assertEqual(reg.arg.size, 4)
     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")
-    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):
   def test_cast_folds(self):
@@ -221,7 +220,7 @@ class TestExecALU(unittest.TestCase):
 
 class TestGatedStoreRewrite(unittest.TestCase):
   def test_tiny_gate_store(self):
-    gmem = UOp.param(0, dtypes.float, (8,))
+    gmem = UOp.param(0, dtypes.float, 8)
     gidx0 = UOp.special(4, 'gidx0')
     gate = gidx0<UOp.const(1)
     idx = UOp(Ops.INDEX, src=(gmem, (gidx0 * UOp.const(2)).valid(gate)))
@@ -237,8 +236,8 @@ class TestGatedStoreRewrite(unittest.TestCase):
     self.assertEqual(len(gated_uops[-1].src), 2)
 
   def test_gate_some_stores(self):
-    gmem0 = UOp.param(0, dtypes.float, (8,))
-    gmem1 = UOp.param(1, dtypes.float, (8,))
+    gmem0 = UOp.param(0, dtypes.float, 8)
+    gmem1 = UOp.param(1, dtypes.float, 8)
     gidx0 = UOp.special(4, 'gidx0')
     idx = gidx0 * UOp.const(2)
     idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gidx0<UOp.const(1))))
@@ -257,8 +256,8 @@ class TestGatedStoreRewrite(unittest.TestCase):
   # scaled down version of TestLinearizerDumb.test_unmerged_ifs
   @unittest.skip("we don't merge ifs anymore")
   def test_merge_ifs_alt(self):
-    gmem0 = UOp.param(0, dtypes.float, (8,))
-    gmem1 = UOp.param(1, dtypes.float, (8,))
+    gmem0 = UOp.param(0, dtypes.float, 8)
+    gmem1 = UOp.param(1, dtypes.float, 8)
     gidx0 = UOp.special(4, 'gidx0')
     idx = gidx0*UOp.const(2)
     gate = gidx0<UOp.const(1)
@@ -281,7 +280,7 @@ class TestGatedStoreRewrite(unittest.TestCase):
 class TestFastIdiv(unittest.TestCase):
   def test_division_power_of_two(self):
     for dt in (dtypes.int32, dtypes.uint32):
-      g = UOp.param(0, dt, (3,))
+      g = UOp.param(0, dt, 3)
       c = UOp.const(2)
       l = g.index(c)
       a = UOp(Ops.CDIV, src=(l, c))
@@ -294,7 +293,7 @@ class TestFastIdiv(unittest.TestCase):
   def test_floormod_power_of_two(self):
     # 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,))
+      g = UOp.param(0, dt, 9)
       c = UOp.const(8)
       a = UOp(Ops.FLOORMOD, src=(g.index(c), c))
       uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -306,7 +305,7 @@ class TestFastIdiv(unittest.TestCase):
   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):
-      g = UOp.param(0, dt, (3,))
+      g = UOp.param(0, dt, 3)
       c = UOp.const(2)
       a = UOp(Ops.FLOORDIV, src=(g.index(c), c))
       uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -319,7 +318,7 @@ class TestFastIdiv(unittest.TestCase):
   @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,))
+    g = UOp.param(0, dtypes.uint32, 4)
     c = UOp.const(3)
     l = g.index(c)
     a = UOp(Ops.CDIV, src=(l, c))
@@ -353,7 +352,7 @@ class TestFastIdiv(unittest.TestCase):
   @unittest.expectedFailure
   def test_fast_idiv_overflow(self):
     # This will be possible with a slightly different method for fast_idiv
-    g = UOp.param(0, dtypes.uint32, (8,))
+    g = UOp.param(0, dtypes.uint32, 8)
     c = UOp.const(7).cast(dtypes.uint)
     l = UOp(Ops.LOAD, src=(g.index(c),))
     a = UOp(Ops.CDIV, src=(l, c))
@@ -364,7 +363,7 @@ class TestFastIdiv(unittest.TestCase):
     self.assertNotIn(Ops.CDIV, ops)
 
   def test_disable_fast_idiv(self):
-    g = UOp.param(0, dtypes.uint32, (4,))
+    g = UOp.param(0, dtypes.uint32, 4)
     c = UOp.const(3)
     l = g.index(c)
     a = UOp(Ops.CDIV, src=(l, c))
@@ -407,8 +406,8 @@ class TestUOpMethod(unittest.TestCase):
       self.assertIs((x != x).simplify(), x.const_like(False, dtypes.bool))
 
   def test_replace(self):
-    x = UOp.param(0, dtypes.int, (1,))
-    self.assertEqual(x.replace(arg=UOp.param(1, dtypes.int, (1,)).arg).arg.slot, 1)
+    x = UOp.param(0, dtypes.int, 1)
+    self.assertEqual(x.replace(arg=UOp.param(1, dtypes.int, 1).arg).arg.slot, 1)
     with self.assertRaises(AssertionError): x.replace(field="a")
 
   def test_const_zero_neg_zero_different(self):
@@ -463,7 +462,7 @@ 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)))
+    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")

test/null/test_uops_stats.py

@@ -146,7 +146,7 @@ class TestUOpsStats(unittest.TestCase):
 
   #MULACC should have the same stats as MUL + ADD
   def test_mulacc(self):
-    globl = UOp.param(0, dtypes.int, (3,))
+    globl = UOp.param(0, dtypes.int, 3)
     o1 = UOp.const(1, dtypes.int)
     o2 = UOp.const(2, dtypes.int)
     u1 = globl.index(o1)
@@ -156,7 +156,7 @@ class TestUOpsStats(unittest.TestCase):
     u5 = UOp(Ops.ADD, src=(u4,u3))
     uops = tuple(u5.toposort())
 
-    globl = UOp.param(0, dtypes.int, (3,))
+    globl = UOp.param(0, dtypes.int, 3)
     o1 = UOp.const(1, dtypes.int)
     o2 = UOp.const(2, dtypes.int)
     u1 = globl.index(o1)

test/null/test_validate_oob.py

@@ -13,7 +13,7 @@ class TestValidateOOB(unittest.TestCase):
   # basic index patterns
   def test_const_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      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)
       with self.assertRaises(RuntimeError):
@@ -23,7 +23,7 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_variable_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      buf = UOp.param(0, dtypes.int, 16)
       to_uops_list([buf.index(Variable("i", 0, 15)).load()])  # valid
       with self.assertRaises(RuntimeError):
         to_uops_list([buf.index(Variable("i", 0, 20)).load()])  # oob
@@ -32,7 +32,7 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_range_with_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      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
       with self.assertRaises(RuntimeError):
@@ -40,7 +40,7 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_variable_with_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      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
       with self.assertRaises(RuntimeError):
@@ -48,7 +48,7 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_gated_store(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      buf = UOp.param(0, dtypes.int, 16)
       v = Variable("v", 0, 20)
       to_uops_list([buf.index(v.valid(v < 16)).store(0)])  # valid
       with self.assertRaises(RuntimeError):
@@ -57,14 +57,14 @@ class TestValidateOOB(unittest.TestCase):
   # ALU ops in index
   def test_floordiv(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      buf = UOp.param(0, dtypes.int, 16)
       to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2).load()])  # 0..15 valid
       with self.assertRaises(RuntimeError):
         to_uops_list([buf.index(UOp.range(34, 0, AxisType.GLOBAL) // 2).load()])  # 0..16 oob
 
   def test_mod(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      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
       with self.assertRaises(RuntimeError):
@@ -72,14 +72,14 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_shr(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      buf = UOp.param(0, dtypes.int, 16)
       to_uops_list([buf.index(UOp.range(64, 0, AxisType.GLOBAL) >> 2).load()])  # 0..15 valid
       with self.assertRaises(RuntimeError):
         to_uops_list([buf.index(UOp.range(128, 0, AxisType.GLOBAL) >> 2).load()])  # 0..31 oob
 
   def test_shl(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (64,))
+      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
       with self.assertRaises(RuntimeError):
@@ -87,7 +87,7 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_and(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      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
       with self.assertRaises(RuntimeError):
@@ -102,14 +102,14 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_max(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      buf = UOp.param(0, dtypes.int, 16)
       to_uops_list([buf.index(Variable("v", -10, 15).maximum(0)).load()])  # 0..15 valid
       with self.assertRaises(RuntimeError):
         to_uops_list([buf.index(Variable("v2", -10, 20).maximum(0)).load()])  # 0..20 oob
 
   def test_xor_in_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      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
       with self.assertRaises(RuntimeError):
@@ -118,22 +118,22 @@ class TestValidateOOB(unittest.TestCase):
   # cast patterns
   def test_float_cast_in_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      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()])
 
   def test_bool_cast_in_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (1,))
+      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
 
   # load result as index/mask
   def test_load_as_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf0 = UOp.param(0, dtypes.int, (16,))
-      buf1 = UOp.param(1, dtypes.int, (64,))
+      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
@@ -142,16 +142,16 @@ class TestValidateOOB(unittest.TestCase):
 
   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,))
+      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()])
 
   def test_load_bool_as_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf_bool = UOp.param(0, dtypes.bool, (16,))
-      buf_int = UOp.param(1, dtypes.int, (8,))
+      buf_bool = UOp.param(0, dtypes.bool, 16)
+      buf_int = UOp.param(1, dtypes.int, 8)
       gidx = UOp(Ops.SPECIAL, src=(UOp.const(16),), arg="gidx0")
       ld_bool = buf_bool.index(gidx).load()
       with self.assertRaises(RuntimeError):
@@ -162,7 +162,7 @@ class TestValidateOOB(unittest.TestCase):
   def test_in_bounds_access_gated_local(self):
     with Context(CHECK_OOB=1):
       # Define buffers
-      gbuf = UOp.param(0, dtypes.uint, (400,))
+      gbuf = UOp.param(0, dtypes.uint, 400)
       sbuf = UOp.placeholder((8,), dtypes.uint, slot=0, addrspace=AddrSpace.LOCAL)
 
       # Define indices, valids and barrier
@@ -186,8 +186,8 @@ class TestValidateOOB(unittest.TestCase):
   @unittest.skip("Bool load is not supported yet")
   def test_load_mask(self):
     with Context(CHECK_OOB=1):
-      glbl0 = UOp.param(0, dtypes.int, (16,))
-      mask = UOp.param(0, dtypes.bool, (16,))
+      glbl0 = UOp.param(0, dtypes.int, 16)
+      mask = UOp.param(0, dtypes.bool, 16)
       ridx = UOp.range(20, 0)
       ld0 = UOp(Ops.LOAD, src=(glbl0.index(UOp.const(ridx<16&mask, ridx))))
       to_uops_list([ld0])

test/null/test_viz.py

@@ -185,18 +185,18 @@ class TestViz(unittest.TestCase):
     @dataclass(frozen=True)
     class TestStruct:
       colored_field: str
-    a = UOp(Ops.CUSTOM, arg=TestStruct(colored("xyz", "magenta")+colored("12345", "blue")))
+    a = UOp(Ops.PYLITERAL, arg=TestStruct(colored("xyz", "magenta")+colored("12345", "blue")))
     a2 = uop_to_json(VizData(), a)[id(a)]
-    self.assertEqual(ansistrip(a2["label"]), f"CUSTOM\n{TestStruct.__qualname__}(colored_field='xyz12345')")
+    self.assertEqual(ansistrip(a2["label"]), f"PYLITERAL\n{TestStruct.__qualname__}(colored_field='xyz12345')")
 
   def test_colored_label_multiline(self):
     with save_viz() as viz:
       arg = colored("x", "green")+"\n"+colored("y", "red")+colored("z", "yellow")+colored("ww\nw", "magenta")
       src = [Tensor.empty(1).uop for _ in range(10)]
-      a = UOp(Ops.CUSTOM, src=tuple(src), arg=arg)
+      a = UOp(Ops.PYLITERAL, src=tuple(src), arg=arg)
       exec_rewrite(a, [PatternMatcher([])])
     a2 = next(viz.get_details(0, 0))["graph"][id(a)]
-    self.assertEqual(ansistrip(a2["label"]), "CUSTOM\nx\nyzww\nw")
+    self.assertEqual(ansistrip(a2["label"]), "PYLITERAL\nx\nyzww\nw")
 
   def test_inf_loop(self):
     a = UOp.const(3)
@@ -244,7 +244,7 @@ class TestViz(unittest.TestCase):
     self.assertEqual(len(lst), 1)
     graphs = [x["graph"] for x in viz.get_details(0, 0)]
     # const is always in the graph, client side hides exclude=True nodes by default
-    self.assertEqual(list(graphs[0]), [id(a.src[0]), id(a), id(z), id(alu), id(y), id(sink)])
+    self.assertEqual(list(graphs[0]), [id(a), id(z), id(alu), id(y), id(sink)])
     self.assertTrue(graphs[0][id(z)]["exclude"])
     self.assertTrue(graphs[0][id(y)]["exclude"])
     self.assertFalse(graphs[0][id(alu)]["exclude"])
@@ -347,7 +347,7 @@ class TestVizGC(unittest.TestCase):
       init = bufs_allocated()
       a = UOp.new_buffer("NULL", 10, dtypes.char)
       a.buffer.allocate()
-      exec_rewrite(UOp(Ops.CUSTOM, src=(a,), arg=a), [PatternMatcher([])])
+      exec_rewrite(UOp(Ops.PYLITERAL, src=(a,), arg=a), [PatternMatcher([])])
       del a
       self.assertEqual(bufs_allocated()-init, 0)
     lst = viz.list_items()
@@ -474,7 +474,7 @@ class TestVizIntegration(unittest.TestCase):
     def custom_fn(X:UOp):
       X = X.flatten()
       i = UOp.range(X.numel(), 0)
-      custom_op = UOp(Ops.CUSTOMI, src=(X[i],), arg="{} + undeclared_name")
+      custom_op = UOp(Ops.CUSTOMI, src=(X[i],), arg=("{} + undeclared_name", X.dtype))
       return X[i].store(custom_op).end(i).sink(arg=KernelInfo(name=f"custom_fn_{X.numel()}"))
     x = Tensor.custom_kernel(Tensor.empty(1, device="CPU"), fxn=custom_fn)[0]
     with save_viz() as viz:

test/testextra/test_bench_log.py

@@ -21,7 +21,7 @@ class TestBenchLog(unittest.TestCase):
     # check event list
     for event in BenchEvent:
       self.assertEqual(len(_events[event]["wall"]), 1)
-      self.assertGreater(_events[event]["wall"][0], 0)
+      self.assertGreater(_events[event]["wall"][0][0], 0)
 
   def test_log_double_wall_time(self):
     for event in BenchEvent:
@@ -35,8 +35,8 @@ class TestBenchLog(unittest.TestCase):
     # check event list
     for event in BenchEvent:
       self.assertEqual(len(_events[event]["wall"]), 2)
-      self.assertGreater(_events[event]["wall"][0], 0)
-      self.assertGreater(_events[event]["wall"][1], 0)
+      self.assertGreater(_events[event]["wall"][0][0], 0)
+      self.assertGreater(_events[event]["wall"][1][0], 0)
 
   @skipIf(_SKIP_KERNEL_TIMING, "ci timing is not accurate")
   def test_log_single_kernel_time(self):
@@ -52,8 +52,8 @@ class TestBenchLog(unittest.TestCase):
     # check event list
     for event in BenchEvent:
       self.assertEqual(len(_events[event]["kernel"]), 1)
-      self.assertLess(_events[event]["kernel"][0], wall_times[0])
-      self.assertGreater(_events[event]["kernel"][0], 0)
+      self.assertLess(_events[event]["kernel"][0][0], wall_times[0])
+      self.assertGreater(_events[event]["kernel"][0][0], 0)
 
   @skipIf(_SKIP_KERNEL_TIMING, "ci cuda timing is not accurate")
   def test_interleaved_wall_kernel_time(self):
@@ -74,8 +74,8 @@ class TestBenchLog(unittest.TestCase):
     for event in BenchEvent:
       self.assertEqual(len(_events[event]["wall"]), 1)
       self.assertEqual(len(_events[event]["kernel"]), 1)
-      self.assertLess(_events[event]["kernel"][0], wall_times[0])
-      self.assertGreater(_events[event]["kernel"][0], 0)
+      self.assertLess(_events[event]["kernel"][0][0], wall_times[0])
+      self.assertGreater(_events[event]["kernel"][0][0], 0)
 
   @skipIf(_SKIP_KERNEL_TIMING, "ci cuda timing is not accurate")
   def test_stacked_wall_kernel_time(self):
@@ -93,10 +93,10 @@ class TestBenchLog(unittest.TestCase):
     for event in BenchEvent:
       self.assertEqual(len(_events[event]["wall"]), 2)
       self.assertEqual(len(_events[event]["kernel"]), 2)
-      self.assertLess(_events[event]["kernel"][0], _events[event]["wall"][0])
-      self.assertGreater(_events[event]["kernel"][0], 0)
-      self.assertLess(_events[event]["kernel"][1], _events[event]["wall"][1])
-      self.assertGreater(_events[event]["kernel"][1], 0)
+      self.assertLess(_events[event]["kernel"][0][0], _events[event]["wall"][0][0])
+      self.assertGreater(_events[event]["kernel"][0][0], 0)
+      self.assertLess(_events[event]["kernel"][1][0], _events[event]["wall"][1][0])
+      self.assertGreater(_events[event]["kernel"][1][0], 0)
 
   def test_log_instant_event(self):
     for event in InstantBenchEvent:
@@ -105,7 +105,7 @@ class TestBenchLog(unittest.TestCase):
     # check event list
     for event in InstantBenchEvent:
       self.assertEqual(len(_events[event]), 1)
-      self.assertEqual(_events[event][0], 1000)
+      self.assertEqual(_events[event][0][0], 1000)
 
 if __name__ == '__main__':
   unittest.main()

test/unit/test_allreduce.py

@@ -3,6 +3,7 @@ from tinygrad import Tensor, UOp, dtypes
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import Ops
 from test.helpers import KernelCountException
+from tinygrad.engine.realize import run_linear
 
 class TestRingAllReduce(unittest.TestCase):
   def test_schedule_ring(self):
@@ -21,13 +22,26 @@ class TestRingAllReduce(unittest.TestCase):
   def test_schedule_all2all(self):
     with Context(ALL2ALL=2):
       N = 4
+      M = N*100
       ds = tuple(f"CPU:{i}" for i in range(N))
-      t = Tensor.empty(N, N*100).shard(ds, axis=0).realize()
-      linear = t.sum(0).mul(2.0).contiguous().linear_with_vars()[0]
+      x = Tensor.arange(N*M, dtype=dtypes.float).reshape(N, M)
+      t = (x*x).clone().shard(ds, axis=0).realize()
+      out = t.sum(0).mul(2.).contiguous()
+      linear, var_vals = out.linear_with_vars()
       copies = [si for si in linear.src if si.src[0].op is Ops.COPY]
       sinks = [si for si in linear.src if si.src[0].op is Ops.SINK]
-      if len(copies) != 24: raise KernelCountException(24, len(copies))
-      if len(sinks) != 26: raise KernelCountException(26, len(sinks))
+      # N*(N-1) copies for input and output
+      copy_count = N*(N-1)*2
+      if len(copies) != copy_count: raise KernelCountException(copy_count, len(copies))
+      # N*(N-1) shrinks from other devices becoming contigs, N ALU, N extra contig, reassembly (cat), and mul
+      sink_count = (N*(N-1))+(N)+(N)+(1)+(1)
+      if len(sinks) != sink_count: raise KernelCountException(sink_count, len(sinks))
+      # correctness
+      run_linear(linear, var_vals)
+      expected = [2*sum((d*M+i)**2 for d in range(N)) for i in range(M)]
+      dev_nums = Tensor.arange(1, N+1, dtype=dtypes.float).reshape(N, 1).expand(N, M).shard(ds, axis=0)
+      shards = out.reshape(1, M).expand(N, M)+dev_nums
+      self.assertListEqual(shards.tolist(), [[x+d+1 for x in expected] for d in range(N)])
 
   @Context(RING=0, ALL2ALL=0)
   def test_schedule_naive(self):

test/unit/test_attention.py

@@ -16,7 +16,7 @@ def apply_rope(x:Tensor, start_pos:int):
 
 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)):
+    for ggml_type,packed_size,words in ((13, 176, 44), (14, 210, 53), (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)

test/unit/test_bnxt.py

@@ -0,0 +1,115 @@
+import struct, unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from tinygrad.runtime.autogen import bnxt
+from extra.bnxt_driver.bnxtdev import BNXT_BACKING_STORE, BNXTDev, BNXTQP, _queue, _qwrite, ipv4_to_gid
+
+class FakePCI:
+  def __init__(self): self.next_addr, self.allocations = 0x100000, []
+  def alloc_sysmem(self, size, contiguous=False):
+    pages = [self.next_addr+i*0x1000 for i in range((size+0xfff)//0x1000)]
+    self.next_addr += len(pages)*0x1000
+    self.allocations.append(mem := bytearray(size))
+    return mem, pages
+
+class FakeDev:
+  def __init__(self): self.pci_dev, self.calls = FakePCI(), []
+  def hwrm(self, name, **fields):
+    self.calls.append((name, fields))
+    typ = fields.get("type", 0)
+    return SimpleNamespace(ctx_init_value=0x5a, ctx_init_offset=4, entry_size=16 if typ == 0 else 4,
+                           subtype_valid_cnt=typ == 0, split_entry_0=2, instance_bit_map=5 if typ == 0 else 1, min_num_entries=0)
+
+class FakeRCFW:
+  def __init__(self): self.calls, self.doorbells = [], []
+  def exec(self, name, **fields):
+    self.calls.append((name, fields))
+    return SimpleNamespace(xid={"create_cq":77, "create_qp":88, "register_mr":0x5678}.get(name, 0))
+  def doorbell(self, *args, **kwargs): self.doorbells.append((args, kwargs))
+
+class FakeQPDev:
+  def __init__(self): self.pci_dev, self.fw, self.gid_id, self.nq_id = FakePCI(), FakeRCFW(), 9, 41
+  def rcfw(self, *args, **kwargs): return self.fw.exec(*args, **kwargs)
+  def doorbell(self, *args, **kwargs): self.fw.doorbell(*args, **kwargs)
+
+class TestMemory(unittest.TestCase):
+  def test_cmdq_and_sq_aux(self):
+    dev = FakeDev()
+    cmdq, sq = _queue(dev), _queue(dev, aux=True)
+    self.assertEqual((cmdq["level"], cmdq["base"]), (0, 0x100000))
+    _qwrite(sq, 3, b"ABCDEFGH", aux=True)
+    self.assertEqual(bytes(sq["mem"][0x1018:0x1020]), b"ABCDEFGH")
+
+  def test_f320_backing_layout_and_final_marker(self):
+    self.assertEqual(len(BNXT_BACKING_STORE), 9)
+    dev = FakeDev()
+    small = ((0, 6), (15, 0))
+    with patch("extra.bnxt_driver.bnxtdev.BNXT_BACKING_STORE", small): BNXTDev.setup_backing_store(dev)
+    cfg = [fields for name, fields in dev.calls if name == "func_backing_store_cfg_v2"]
+    self.assertEqual([(x["type"], x["instance"]) for x in cfg], [(0, 0), (0, 2), (15, 0)])
+    self.assertTrue(all(not x["flags"] for x in cfg[:-1]))
+    self.assertEqual(cfg[-1]["flags"], bnxt.FUNC_BACKING_STORE_CFG_V2_REQ_FLAGS_BS_CFG_ALL_DONE)
+    self.assertEqual((dev.pci_dev.allocations[0][4], dev.pci_dev.allocations[0][20]), (0x5a, 0x5a))
+
+class TestRCFW(unittest.TestCase):
+  def setUp(self):
+    patch("extra.bnxt_driver.bnxtdev.System.memory_barrier").start()
+    self.addCleanup(patch.stopall)
+
+  def test_doorbell_encodes_xid_type_and_index(self):
+    dev = BNXTDev.__new__(BNXTDev)
+    dev.db, dev.db_off = [0]*1024, 0x1000
+    dev.doorbell(0x123456, bnxt.DBC_DBC_TYPE_CQ_ARMALL, 0x456, epoch=1)
+    key = dev.db[0x1000//8]
+    self.assertEqual(key >> 32,
+      0x123456 & bnxt.DBC_DBC_XID_MASK | bnxt.DBC_DBC_PATH_ROCE | bnxt.DBC_DBC_TYPE_CQ_ARMALL | bnxt.BNXT_QPLIB_DBR_VALID)
+    self.assertEqual(key & 0xffffffff, 0x456 | 1<<bnxt.BNXT_QPLIB_DBR_EPOCH_SHIFT)
+
+  def test_command_uses_first_flag(self):
+    dev = BNXTDev.__new__(BNXTDev)
+    dev.bar0, dev.cmdq, dev.creq = [0]*1024, _queue(FakeDev()), _queue(FakeDev())
+    dev.rcfw_first, dev.creq_id = True, 23
+    dev.doorbell = lambda *args: None
+    _qwrite(dev.creq, 0, bytes(bnxt.struct_creq_query_version_resp(type=bnxt.CREQ_BASE_TYPE_QP_EVENT, cookie=0, v=1)))
+    ret = dev.rcfw("query_version")
+    req = bnxt.struct_cmdq_query_version.from_buffer_copy(bytes(dev.cmdq["mem"][:16]))
+    prod = dev.bar0[(bnxt.RCFW_COMM_BASE_OFFSET+bnxt.RCFW_PF_VF_COMM_PROD_OFFSET)//4]
+    self.assertEqual((req.cookie, ret.cookie, prod), (0, 0, 1 | 1<<bnxt.FIRMWARE_FIRST_FLAG))
+
+class TestFastPath(unittest.TestCase):
+  def test_unified_mr(self):
+    dev = BNXTDev.__new__(BNXTDev)
+    fw = FakeRCFW()
+    dev.pci_dev, dev.rcfw = FakePCI(), fw.exec
+    self.assertEqual(dev.register_mem([0x800000, 0x900000], 0x2000), 0x5678)
+    mr = fw.calls[-1][1]
+    self.assertEqual((mr["flags"], mr["va"], mr["mr_size"], mr["log2_pg_size_lvl"]),
+      (bnxt.CMDQ_REGISTER_MR_FLAGS_ALLOC_MR, 0x800000, 0x2000,
+       1<<bnxt.CMDQ_REGISTER_MR_LVL_SFT | 12<<bnxt.CMDQ_REGISTER_MR_LOG2_PG_SIZE_SFT))
+
+  def test_qp_creation_and_connect_use_f320_layout(self):
+    dev = FakeQPDev()
+    qp = BNXTQP(dev)
+    create = next(fields for name, fields in dev.fw.calls if name == "create_qp")
+    self.assertEqual((create["sq_size"], "rq_size" in create, qp.qpn), (16, False, 88))
+    qp.connect(0x123, ipv4_to_gid("10.0.0.2"), 0x001122334455)
+    rtr, rts = dev.fw.calls[-2][1], dev.fw.calls[-1][1]
+    self.assertEqual((bytes(rtr["dgid"]), bytes(rtr["dest_mac"])),
+      (ipv4_to_gid("10.0.0.2"), bytes.fromhex("001122334455")))
+    self.assertEqual((rtr["modify_mask"], rts["modify_mask"]), (0x41515ad, 0xae005))
+
+  def test_rdma_write_builds_three_slots_and_host_msn(self):
+    qp = BNXTQP.__new__(BNXTQP)
+    qp.dev, qp.qpn = FakeQPDev(), 88
+    qp.sq, qp.sq_psn, qp.msn = _queue(FakeDev(), aux=True), 5, 0
+    qp._poll = lambda timeout: bytes(bnxt.struct_cq_req())
+    qp.rdma_write(0x1122334455667788, 0x99aa, 0x12345000, 0x55aa, 100)
+    hdr = bnxt.struct_sq_rdma_hdr.from_buffer_copy(bytes(qp.sq["mem"][:32]))
+    sge = bnxt.struct_sq_sge.from_buffer_copy(bytes(qp.sq["mem"][32:48]))
+    self.assertEqual((hdr.remote_va, hdr.remote_key, hdr.length, sge.va_or_pa, sge.l_key, sge.size),
+      (0x1122334455667788, 0x99aa, 100, 0x12345000, 0x55aa, 100))
+    self.assertEqual(struct.unpack_from("<Q", qp.sq["mem"], 0x1000)[0], 6<<24 | 5)
+    self.assertEqual(qp.dev.fw.doorbells, [((88, bnxt.DBC_DBC_TYPE_SQ, 3, 0), {})])
+
+if __name__ == "__main__": unittest.main()

test/unit/test_bnxt_transport.py

@@ -0,0 +1,40 @@
+import ctypes, unittest
+from unittest.mock import patch
+
+from tinygrad.runtime.autogen import bnxt
+from extra.bnxt_driver.bnxtdev import BNXT_CHIMP_COMM, BNXT_CHIMP_COMM_TRIGGER, BNXTDev
+
+class Mailbox:
+  def __init__(self, trigger): self.words, self.trigger = {}, trigger
+  def __setitem__(self, idx, val):
+    self.words[idx] = val
+    if idx == BNXT_CHIMP_COMM_TRIGGER//4: self.trigger()
+  def request(self):
+    base = BNXT_CHIMP_COMM//4
+    return b"".join(self.words.get(base+i, 0).to_bytes(4, "little") for i in range(bnxt.HWRM_MAX_REQ_LEN//4))
+
+def fake_dev():
+  dev = BNXTDev.__new__(BNXTDev)
+  dev.resp, dev.resp_pa, dev.seq = bytearray(0x1000), [0x6789a000], 0
+  return dev
+
+def reply(dev, out_type):
+  req = bnxt.struct_hwrm_cmd_hdr.from_buffer_copy(dev.bar0.request())
+  out = out_type(req_type=req.req_type, seq_id=req.seq_id, resp_len=ctypes.sizeof(out_type), valid=1)
+  dev.resp[:ctypes.sizeof(out_type)] = bytes(out)
+
+class TestHWRM(unittest.TestCase):
+  def setUp(self):
+    self.barrier = patch("extra.bnxt_driver.bnxtdev.System.memory_barrier").start()
+    self.addCleanup(patch.stopall)
+
+  def test_request(self):
+    dev = fake_dev()
+    dev.bar0 = Mailbox(lambda: reply(dev, bnxt.struct_hwrm_func_qcaps_output))
+    dev.hwrm("func_qcaps", fid=0xffff)
+    req = bnxt.struct_hwrm_func_qcaps_input.from_buffer_copy(dev.bar0.request())
+    self.assertEqual((req.req_type, req.seq_id, req.resp_addr, req.fid),
+      (bnxt.HWRM_FUNC_QCAPS, 1, dev.resp_pa[0], 0xffff))
+    self.barrier.assert_called_once_with()
+
+if __name__ == "__main__": unittest.main()

test/unit/test_call.py

@@ -11,7 +11,7 @@ class TestCall(unittest.TestCase):
     Tensor.realize(a,b)
 
     # we define a plus function
-    plus_fxn = UOp.param(0, dtypes.float, (10,10)) + UOp.param(1, dtypes.float, (10,10))
+    plus_fxn = UOp.param_from_shape(0, (10,10), dtypes.float) + UOp.param_from_shape(1, (10,10), dtypes.float)
 
     c = Tensor.call(a, b, fxn=plus_fxn)
     np.testing.assert_equal(c.numpy(), (a+b).numpy())
@@ -29,7 +29,7 @@ class TestCall(unittest.TestCase):
     def grad_fxn(grad:UOp, call:UOp): return (grad, grad)
 
     # we define a plus function
-    plus_fxn = UOp.param(0, dtypes.float, (10,10)) + UOp.param(1, dtypes.float, (10,10))
+    plus_fxn = UOp.param_from_shape(0, (10,10), dtypes.float) + UOp.param_from_shape(1, (10,10), dtypes.float)
     c = Tensor.call(a, b, fxn=plus_fxn, grad_fxn=grad_fxn)
     c.mean().backward()
 
@@ -45,7 +45,7 @@ class TestCall(unittest.TestCase):
     gt_b_grad = b.grad.numpy()
     a.grad, b.grad = None, None
 
-    plus_fxn = UOp.param(0, dtypes.float, (10,10)) + UOp.param(1, dtypes.float, (10,10))
+    plus_fxn = UOp.param_from_shape(0, (10,10), dtypes.float) + UOp.param_from_shape(1, (10,10), dtypes.float)
     c = Tensor.call(a, b, fxn=plus_fxn)
     c.mean().backward()
 
@@ -67,8 +67,8 @@ class TestCall(unittest.TestCase):
     Tensor.realize(a, b)
 
     # we define a gemm function
-    x = UOp.param(0, dtypes.float, shape=(M, K))
-    y = UOp.param(1, dtypes.float, shape=(K, N))
+    x = UOp.param_from_shape(0, (M, K), dtypes.float)
+    y = UOp.param_from_shape(1, (K, N), dtypes.float)
     c = Tensor.call(a, b, fxn=x@y)
 
     np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5, atol=1e-6)
@@ -83,7 +83,7 @@ class TestCall(unittest.TestCase):
     gt_a_grad, gt_b_grad = a.grad.numpy(), b.grad.numpy()
     a.grad, b.grad = None, None
 
-    p0, p1 = UOp.param(0, dtypes.float, (10,10)), UOp.param(1, dtypes.float, (10,10))
+    p0, p1 = UOp.param_from_shape(0, (10,10), dtypes.float), UOp.param_from_shape(1, (10,10), dtypes.float)
     complex_fxn = (p0*p1 + p0).exp2() * p1.reciprocal()
     c = Tensor.call(a, b, fxn=complex_fxn)
     c.mean().backward()

test/unit/test_dtype_weak.py

@@ -116,7 +116,7 @@ class TestWeakPromotion(unittest.TestCase):
 
   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))
+      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)
     # a bare weak CONST commits directly: the pass runs without symbolic, so a CAST here would survive it
@@ -132,7 +132,7 @@ class TestWeakPromotion(unittest.TestCase):
 
   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()
+    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)

test/unit/test_llm_amd.py

@@ -20,11 +20,13 @@ class TestQ8Quantize(unittest.TestCase):
   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)
+    quant, scale, gsum = 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)
+    # xsum holds the two per-16 sums per 32-wide group
+    np.testing.assert_array_equal(gsum.numpy().reshape(2, 2), expected.reshape(2, 2, 16).sum(-1).astype(np.float32))
 
   def test_q6_linear_compiles(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
@@ -36,7 +38,9 @@ class TestQ8Quantize(unittest.TestCase):
     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)
+    # the Q6 weight is repacked: 210-byte blocks padded to 212 (one block = 53 words)
+    self.assertEqual(linear.weight.uop.buf_uop.buffer.nbytes, 53*4)
+    self.assertEqual(linear.weight.dtype, dtypes.uint32)
 
   def test_q4_k_linear(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")

test/unit/test_multitensor.py

@@ -852,82 +852,6 @@ class TestMultiFromUnrenderable(unittest.TestCase):
     np.testing.assert_equal(ll.numpy(), np.arange(100)+1)
 
 @unittest.skipIf(not_support_multi_device(), "need multi")
-class TestMultiAssign(unittest.TestCase):
-  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
-
-  @needs_second_gpu
-  def setUp(self): pass
-
-  def test_multi_assign_realized(self):
-    out = Tensor.zeros(4).shard(self.device, 0).contiguous().realize()
-    ones = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
-    out.assign(ones).realize()
-    self.assertListEqual(out.tolist(), [1,1,1,1])
-
-  def test_multi_assign_unrealized(self):
-    out = Tensor.zeros(4).contiguous().realize().shard(self.device, 0)
-    ones = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
-    out.assign(ones).realize()
-    self.assertListEqual(out.tolist(), [1,1,1,1])
-
-  def test_multi_assign_both_unrealized(self):
-    out = Tensor.zeros(4).contiguous().realize().shard(self.device, 0)
-    ones = Tensor.ones(4).contiguous().realize().shard(self.device, 0)
-    out.assign(ones).realize()
-    self.assertListEqual(out.tolist(), [1,1,1,1])
-
-  def test_multi_assign_scalar(self):
-    out = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
-    out.assign(0).realize()
-    self.assertListEqual(out.tolist(), [0,0,0,0])
-
-  def test_multi_assign_const_like(self):
-    out = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
-    out.assign(out.const_like(7)).realize()
-    self.assertListEqual(out.tolist(), [7,7,7,7])
-
-  def test_multi_assign_piece(self):
-    out = Tensor.zeros(4,4).shard(self.device, 0).contiguous().realize()
-    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
-    out[:, 2:3].assign(ones).realize()
-    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
-
-  def test_multi_assign_piece_noncontig(self):
-    out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0).realize()
-    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
-    out[:, 2:3].assign(ones).realize()
-    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
-
-  @unittest.expectedFailure
-  def test_multi_assign_piece_unrealized(self):
-    out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0)
-    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
-    out[:, 2:3].assign(ones).realize()
-    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
-
-  def test_multi_assign_var_offset(self):
-    out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0).realize()
-    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
-    vi = Variable("i", 0, 3).bind(2)
-    out[:, vi:vi+1].assign(ones).realize()
-    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
-
-  def test_multi_assign_var_offset_jit_none(self): self.test_multi_assign_var_offset_jit(None)
-  def test_multi_assign_var_offset_jit(self, shard_axis=0):
-    out = Tensor.zeros(4,6).contiguous().realize().shard(self.device, shard_axis).realize()
-    ones = Tensor.ones(4,1).shard(self.device, shard_axis).contiguous().realize()
-
-    @TinyJit
-    def f(out:Tensor, vi):
-      out[:, vi:vi+1].assign(ones).realize()
-      ones.assign(ones+1).realize()
-
-    vi = Variable("i", 0, 5)
-    for i in range(1,5):
-      GlobalCounters.reset()
-      f(out, vi.bind(i))
-    self.assertListEqual(out.tolist(), [[0,1,2,3,4,0]]*4)
-
 @unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiSetitem(unittest.TestCase):
   device = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))

tinygrad/codegen/__init__.py

@@ -26,7 +26,7 @@ from tinygrad.schedule.prepare 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
-from tinygrad.helpers import all_same, flatten, argsort, partition
+from tinygrad.helpers import all_same, all_int, flatten, argsort, partition
 from tinygrad.uop.ops import _broadcast_shape, identity_element
 from tinygrad.schedule.rangeify import BufferizeOpts
 
@@ -162,9 +162,10 @@ devectorizer2 = mop_cleanup+pm_mops+PatternMatcher([
   (UPat(Ops.RESHAPE, dtype=dtypes.void, name="x"), lambda x: x.src[0]),
   # reshape of a single element shaped value to scalar is an index
   (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(0) if x.marg == () and x.src[0].shape == (1,) else None),
-  # EXPAND on scalar -> STACK
+  # EXPAND on scalar -> nested STACKs with the same shape
   (UPat(Ops.EXPAND, src=(UPat.var("x"), UPat()), name="out"),
-   lambda x,out: UOp.stack(*([x]*out.max_numel())) if x.shape == () and out.shape == (out.max_numel(),) else None),
+   lambda x,out: functools.reduce(lambda x,s: UOp.stack(*([x]*s)), reversed(out.shape), x)
+   if x.shape == () and all_int(out.shape) and 0 not in out.shape else None),
 ])
 
 def fix_group_for_reduce(x:UOp):

tinygrad/codegen/gpudims.py

@@ -73,7 +73,7 @@ def add_gpudims(ctx:Renderer, s:UOp):
   subs = {}
   for r in s_topo:
     # look for local INDEXes that are not used in the GLOBAL store, then add them as an INVALID
-    if r.op is Ops.STORE and (idx := r.src[0]).src[0].addrspace == AddrSpace.GLOBAL:
+    if r.op is Ops.STORE and len((idx := r.src[0]).src) and idx.src[0].addrspace == AddrSpace.GLOBAL:
       missing_locals = [all_ranges[rng] for rng in local_dims if all_ranges[rng] not in idx.ranges]
       if len(missing_locals):
         assert len(idx.src) == 2, "index has 2 sources"

tinygrad/codegen/late/coalesce.py

@@ -1,7 +1,8 @@
 import itertools, functools
 from collections import defaultdict
+from dataclasses import replace
 from tinygrad.dtype import dtypes, AddrSpace, Invalid, DType
-from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, shape_to_shape_arg, graph_rewrite
+from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, graph_rewrite
 from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate, sym
 from tinygrad.helpers import getenv, IMAGE, OSX, ceildiv, is_image_shape
 from tinygrad.renderer import Renderer
@@ -85,7 +86,8 @@ def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
   if len(cands) == 0: return None
   # and tiebreak with indexing complexity (ie. number of nodes)
   h, w, cidx = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].index(1).simplify().backward_slice))
-  buf = buf.replace(src=(shape_to_shape_arg((h, w, 4)),))
+  # the image dims are stored in the arg, the size in the arg stays the flat buffer len
+  buf = buf.replace(arg=replace(buf.arg, size=h*w*4, image=(h, w)))
   shapes[buf.arg.slot] = (h, w)
   if valid.op is not Ops.CONST or valid.val is not True:
     return buf.index(cidx.src[1].valid(valid), cidx.src[0].valid(valid))
@@ -95,7 +97,6 @@ def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
 pm_simplify_add_image = PatternMatcher([
   (UPat(Ops.SHRINK, src=(UPat(Ops.PARAM, name="buf"), UPat(name="x"), UPat(arg=4))), transform_to_image),
   # image load/store is always float
-  (UPat(Ops.INDEX, dtype=dtypes.float, name="x").load(dtype=dtypes.half), lambda x: x.load().cast(dtypes.half)),
   (UPat(Ops.INDEX, dtype=dtypes.float, name="x").store(UPat(name="d", dtype=dtypes.half)), lambda x,d: x.store(d.cast(dtypes.float))),
   (UPat.var("x", dtype=dtypes.float).cast(dtypes.half).cast(dtypes.float), lambda x: x),
 ])
@@ -104,7 +105,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
   if getenv("DMC"): return sink
 
   # collect
-  memory: defaultdict[tuple[Ops, UOp, UOp|str, UOp], dict[int, list[UOp]]] = defaultdict(dict)
+  memory: defaultdict[tuple[Ops, UOp, UOp|str, UOp, object], dict[int, list[UOp]]] = defaultdict(dict)
   for u in sink.toposort():
     # TODO: this should handle images too, it's just memory coalescing
     if u.op in {Ops.LOAD, Ops.STORE}:
@@ -119,11 +120,12 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
       elif idx.op is Ops.CONST and idx.val is Invalid: root_src, arg = "INVALID", 0
       elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.val
       else: root_src, arg = idx, 0
-      memory[(u.op, buf, root_src, valid)].setdefault(arg, []).append(u)
+      # loads/stores only coalesce with others carrying the same arg (e.g. the nontemporal flag)
+      memory[(u.op, buf, root_src, valid, u.arg)].setdefault(arg, []).append(u)
 
   # build replacements
   replacements = {}
-  for (op,buf,base,valid),offsets in memory.items():
+  for (op,buf,base,valid,ld_arg),offsets in memory.items():
     # allowed lengths (copied in)
     lengths = []
     must_divide = True
@@ -158,7 +160,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
           store = idx.store(UOp.stack(*datas) if len(datas) > 1 else datas[0])
           for i,g in enumerate(grp): replacements[offsets[g][0]] = store
         else:
-          ld = idx.load()
+          ld = idx.load(arg=ld_arg)
           for i,g in enumerate(grp):
             for oo in offsets[g]:
               replacements[oo] = ld.index(i) if len(grp) > 1 else ld

tinygrad/codegen/simplify.py

@@ -35,10 +35,10 @@ def simplify_merge_adjacent(u:UOp) -> UOp|None:
         nidx = graph_rewrite(u, _substitute+symbolic+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
                              name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")
 
-        # check if it simplifies
-        if count_divmod(nidx) <= count_divmod(u):
-          u = nidx
-  return u
+        # check if it simplifies. return after one merge so the next rewrite uses the new ranges,
+        # rather than continuing with stale pairs from the original ended_ranges.
+        if count_divmod(nidx) <= count_divmod(u): return nidx
+  return None
 
 def mark_gated(ctx, idx):
   if len(idx.src) > 1 and idx.src[1].op is Ops.WHERE:

tinygrad/engine/jit.py

@@ -7,7 +7,7 @@ from tinygrad.dtype import DType
 from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, rewrite_group, graph_rewrite
 from tinygrad.renderer import Estimates
 from tinygrad.engine.realize import capturing, compile_linear, link_linear, run_linear, graph_cache, estimate_uop, get_runtime
-from tinygrad.engine.realize import unwrap_multi, resolve_params, get_call_arg_uops, get_call_outs_ins
+from tinygrad.engine.realize import unwrap_multi, resolve_params, get_call_arg_uops, get_call_written_bufs
 from tinygrad.schedule.memory import memory_plan_rewrite, _collect_bufs
 from tinygrad.nn.state import get_parameters
 from tinygrad.uop.movement import mop_cleanup
@@ -67,7 +67,7 @@ def jit_lower(linear:UOp, held_bufs:set[UOp], input_uops:list[UOp]) -> UOp:
   if VIZ: graph_rewrite(linear, PatternMatcher([]), name="View captured linear")
 
   # parametrize input buffers: map each input buffer UOp to a PARAM with the correct slot index
-  linear = linear.substitute({u: UOp.param(i, u.dtype, u.shape, u.device) for i,u in enumerate(input_uops)}, walk=True)
+  linear = linear.substitute({u: UOp.param_from_shape(i, u.shape, u.dtype, u.device) for i,u in enumerate(input_uops)}, walk=True)
   linear = memory_plan_rewrite(linear, held_bufs)
   linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value))
   if JIT < 2: linear = graph_split_rewrite(linear, max_batch_size=JIT_BATCH_SIZE.value)
@@ -173,13 +173,7 @@ class CapturedJit(Generic[ReturnType]):
 
   @functools.cached_property
   def _written_uops(self) -> set[UOp]:
-    out: set[UOp] = set()
-    for call in self.linear.toposort():
-      if call.op is not Ops.CALL: continue
-      arg_uops = get_call_arg_uops(call)
-      outs, ins = get_call_outs_ins(call)
-      out |= {b for k in set(outs) - set(ins) if (b:=u if (cv:=(u:=arg_uops[k]).contiguous_view()) is None else cv[0]).op is Ops.BUFFER}
-    return out
+    return {b for call in self.linear.toposort() if call.op is Ops.CALL for b in get_call_written_bufs(call)}
 
   def __call__(self, input_uops:list[UOp], var_vals:dict[str, int]) -> ReturnType:
     concrete = tuple(_copy_input(u) if u in self._written_uops else u for u in input_uops)

tinygrad/engine/realize.py

@@ -2,7 +2,7 @@ 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, tqdm, dedup
 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.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
@@ -26,6 +26,10 @@ 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_written_bufs(call:UOp) -> list[UOp]:
+  arg_uops, (outs, ins) = get_call_arg_uops(call), get_call_outs_ins(call)
+  return dedup([b for k in outs if k not in ins and (b:=u if (cv:=(u:=arg_uops[k]).contiguous_view()) is None else cv[0]).op is Ops.BUFFER])
+
 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]
@@ -213,14 +217,12 @@ def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
 
 def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
   dev = cast(Any, Device[(info:= call.arg.aux).device[0]])
-  addrs = [(b.bufs[j] if isinstance(b:=_resolve(ctx.input_uops[k], ctx.input_uops).buffer, MultiBuffer) else b).get_buf(dev_name).va_addr
-           for devs, idxs in info.input_idxs for j, dev_name in enumerate(devs) for k in idxs]
+  addrs = [cast(Buffer, _resolve(u, ctx.input_uops).buffer).get_buf(d).va_addr for d, u in info.input_addrs]
   dev.rt_buffer()._buf.cpu_view().view(offset=(base:=dev.rt_allocator.alloc(len(addrs) * 8)), fmt='Q')[:len(addrs)] = array.array('Q', addrs)
 
   if info.inputs is not None:
-    tables = [UOp.from_buffer(dev.rt_buffer().view(len(idxs), dtypes.uint64, base + j*len(idxs)*8), HCQ_RUNTIME_DEV.value)
-              for devs, idxs in info.input_idxs for j in range(len(devs))]
-    call = call.substitute({call.src[1+info.inputs]: UOp.mstack(*tables)})
+    table = UOp.from_buffer(dev.rt_buffer().view(len(info.input_addrs), dtypes.uint64, base), HCQ_RUNTIME_DEV.value)
+    call = call.substitute({call.src[1+info.inputs]: UOp.mstack(*[table]*len(info.device))})
   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:

tinygrad/helpers.py

@@ -250,16 +250,16 @@ EMULATED_DTYPES = ContextVar("EMULATED_DTYPES", "")
 DEFAULT_FLOAT, DEFAULT_INT = ContextVar("DEFAULT_FLOAT", "float32"), ContextVar("DEFAULT_INT", "int32")
 CAPTURE_PROCESS_REPLAY = ContextVar("CAPTURE_PROCESS_REPLAY", 0)
 def _get_cpu_count() -> int:
-  # os.process_cpu_count (3.13+) respects cgroup limits
-  if hasattr(os, "process_cpu_count"): return max(1, os.process_cpu_count() or 1)
-  # cgroup v2 (containers with --cpus=N)
+  # os.process_cpu_count is available in 3.13+, then try affinity, then fallback to cpu_count
+  count = (os.process_cpu_count() if hasattr(os, "process_cpu_count") else
+           len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count()) or 1
+  # limit with cgroup v2 (containers with --cpus=N)
   try:
     with open("/sys/fs/cgroup/cpu.max") as f:
       quota, period = f.read().strip().split()
-      if quota != "max": return max(1, int(quota) // int(period))
+      if quota != "max": count = min(count, max(1, int(quota) // int(period)))
   except (FileNotFoundError, ValueError, ZeroDivisionError): pass
-  # fall back to affinity (respects taskset but not cgroup quota)
-  return max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1))
+  return count
 NUM_CPU_THREADS = ContextVar("NUM_CPU_THREADS", _get_cpu_count())
 NULL_ALLOW_COPYOUT = ContextVar("NULL_ALLOW_COPYOUT", 0)
 # VIZ implies PROFILE, but you can run PROFILE without VIZ
@@ -271,7 +271,6 @@ PROFILE = ContextVar("PROFILE", abs(VIZ.value))
 SPEC = ContextVar("SPEC", 1)
 # TODO: disable by default due to speed
 CHECK_OOB = ContextVar("CHECK_OOB", 0)
-PCONTIG = ContextVar("PCONTIG", 0)  # partial contiguous in rangeify
 DEBUG_RANGEIFY = ContextVar("DEBUG_RANGEIFY", 0)
 # set to 1, this uses tuplize in the linearizer sort order
 TUPLE_ORDER = ContextVar("TUPLE_ORDER", 1)

tinygrad/llm/cli.py

@@ -12,22 +12,29 @@ class SimpleTokenizer:
   def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int], preset:str="llama3",
                bos_id:int|None=None, eos_id:int=0, eot_id:int|None=None):
     preset = {"qwen35":"qwen2","qwen35moe":"qwen2"}.get(preset, preset)
-    if preset not in ("llama3","llama-v3","llama-bpe","qwen2","olmo","kimi-k2","tekken","glm4"):
+    if preset not in ("llama3","llama-v3","llama-bpe","qwen2","olmo","kimi-k2","tekken","glm4","gpt-4o"):
       raise ValueError(f"Invalid tokenizer preset '{preset}'")
     # https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
     bs = [*range(33, 127), *range(161, 173), *range(174, 256)]  # bytes that map to themselves
     self._byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)}
 
     # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
-    # 0x323b0 is one past the max codepoint in unicode categories L/N/Z (0x323af is max L)
+    # each limit is one past the category's max codepoint (Z→U+3000, N→U+1FBF9, L→U+323AF, M→U+E01EF)
     # compact adjacent codepoints into ranges: listing them all makes re spend seconds on large prompts
-    def ucat_range(pre:str) -> str:
-      cps = enumerate(cp for cp in range(0x323b0) if unicodedata.category(chr(cp)).startswith(pre))
+    def ucat_range(pre:str|tuple[str, ...]) -> str:
+      limits = {"Z": 0x3001, "N": 0x1fbfa, "L": 0x323b0, "M": 0xe01f0}
+      limit = max(limits[p if p in limits else p[0]] for p in (pre if isinstance(pre, tuple) else (pre,)))
+      cps = enumerate(cp for cp in range(limit) if unicodedata.category(chr(cp)).startswith(pre))
       runs = [list(g) for _, g in itertools.groupby(cps, lambda e: e[1]-e[0])]
       return "".join(re.escape(chr(g[0][1])) + (f"-{re.escape(chr(g[-1][1]))}" if len(g) > 1 else "") for g in runs)
     r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L")
-    self._split_to_word = re.compile("(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + \
-      f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+")
+    contr, r_l, r_n = "(?i:'s|'t|'re|'ve|'m|'ll|'d)", f"[^\\r\\n{r_p_N}{r_p_L}]?", f"[{r_p_N}]" if preset == "tekken" else f"[{r_p_N}]{{1,3}}"
+    r_p, r_w, r_t = f" ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*", f"{contr}|{r_l}[{r_p_L}]+", f"[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+"
+    if preset in ("tekken", "gpt-4o"):
+      r_up, r_lo = ucat_range(("Lu","Lt","Lm","Lo","M")), ucat_range(("Ll","Lm","Lo","M"))
+      sfx = f"{contr}?" if preset == "gpt-4o" else ""
+      r_p, r_w = f" ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n/]*", f"{r_l}[{r_up}]*[{r_lo}]+{sfx}|{r_l}[{r_up}]+[{r_lo}]*{sfx}"
+    self._split_to_word = re.compile(f"{r_w}|{r_n}|{r_p}|{r_t}")
     self._split_to_sentence = re.compile("|".join(re.escape(tok) for tok in special_tokens.keys()) if special_tokens else r"(?!)")
 
     self._normal_tokens = {bytes(self._byte_decoder[c] for c in tok): tid for tok, tid in normal_tokens.items()}

tinygrad/llm/kernels/amd.py

@@ -7,12 +7,13 @@ 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
+BLOCK_M, BLOCK_N, WARP_SIZE = 32, 32, 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
+Q6_PADDED, Q6_WORDS = 212, 53  # the 210-byte Q6 blocks are padded to 212 bytes so they are word-addressable
 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:
@@ -35,8 +36,8 @@ def amd_custom_kernels_supported(device:str|tuple[str, ...]|None) -> bool:
 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}))")
+    other = UOp(Ops.CUSTOM, src=(val,), arg=
+      (f"__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {{0}}), {0x1f | offset<<10}))", dtypes.float))
     val = val.maximum(other) if maximum else val + other
   return val
 
@@ -54,21 +55,43 @@ class Linear(nn.Linear):
     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)
+    graph = decoded.uop.toposort()
+    raw = next((u for u in graph if u.op is Ops.SHRINK and u.dtype == dtypes.uint8 and prod(u.shape) in packed_sizes), None)
     if raw is None: return
+    ggml_type = packed_sizes[prod(raw.shape)]
+    # the packed byte rate alone can't distinguish same-rate formats (Q4_0 vs Q4_K, Q5_0 vs Q5_K, MXFP4 vs IQ4_XS).
+    # the supported formats are 256-wide superblocks: their decode views the packed bytes at the superblock width
+    # (ggml_data_to_tensor reshapes to (-1, QUANT_SIZES[type])), while same-rate 32-wide formats reshape to 17-22
+    if not any(u.op is Ops.RESHAPE and u.shape[-1:] == (QUANT_SIZES[ggml_type],) for u in graph): 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)]
+    self.ggml_type = ggml_type
     # 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)))
+    if self.ggml_type == Q6_K:
+      # Q6 blocks are 210 bytes, so consecutive blocks are only 2-byte aligned. pad each block to 212 bytes
+      # (a one-time copy at load) so the kernel can do all its reads as aligned u32 words
+      nbytes, nblocks = raw.max_numel(), raw.max_numel() // Q6_BYTES
+      byte_view = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer).view(nbytes, dtypes.uint8, raw_offset)))
+      padded = byte_view.reshape((nblocks, Q6_BYTES)).pad_to((nblocks, Q6_PADDED)).contiguous().realize()
+      self.weight = Tensor(UOp.from_buffer(cast(Buffer, padded.uop.buf_uop.buffer).view(nblocks * Q6_WORDS, dtypes.uint32, 0)))
+    else:
+      self.weight = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer)
+        .view(raw.max_numel() * raw.dtype.itemsize // dtypes.uint32.itemsize, dtypes.uint32, 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 is None:
+        # tiny dense fp16 matmul (e.g. the ssm beta/alpha head rows): single fp16 gemv kernel instead of a
+        # generic matmul schedule, and realize the densely packed weight once if it is still a lazy ggml view
+        if self.weight.dtype in (dtypes.half, dtypes.float, dtypes.bfloat16) and self.out_features <= 2048 \
+          and self.in_features % (WARP_SIZE*4) == 0:
+          numel, max_shape = x.numel(), x.max_shape
+          if isinstance(numel, int) or prod(max_shape) // self.in_features <= 32:
+            out = f16_gemv(self, x if isinstance(numel, int) else x.pad_to(max_shape))
+            return out if isinstance(numel, int) else out.shrink(tuple((0, s) for s in (*x.shape[:-1], self.out_features)))
+        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
@@ -77,14 +100,20 @@ class Linear(nn.Linear):
     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)")
+  # int8 4-wide dot, widened to scalar multiply-adds (2% decode slower than the sudot4 builtin, but portable)
+  for i in range(4):
+    av = ((a >> (8*i)) & 255).cast(dtypes.uint8).bitcast(dtypes.int8).int()
+    bv = ((b >> (8*i)) & 255).cast(dtypes.uint8).bitcast(dtypes.int8).int()
+    c = c + av*bv
+  return c
 
 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({}, {}, {})")
+  return UOp(Ops.CUSTOMI, src=tuple(x.cast(dtypes.uint32) for x in (a, b, selectors)), arg=("__builtin_amdgcn_perm({}, {}, {})", dtypes.uint32))
 
 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})")
+  # nontemporal scalar load: streamed weights must not evict the activations/KV cache from L2
+  if lanes is None: return ptr.load(arg="nontemporal")
   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)
@@ -93,16 +122,19 @@ def _load_byte(raw:UOp, base:UOp, offset:UOp) -> UOp: return (raw[base + offset/
 def _half(value:UOp) -> UOp: return value.cast(dtypes.uint16).bitcast(dtypes.float16).float()
 
 def _iq4_bytes(packed:UOp, shift:int) -> UOp:
+  # the non-linear iq4nl table as a byte lookup: 3 byte_perms beat any arithmetic/select-tree form (~60% decode)
   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))
+  # scales/mins (6-bit each) live in block bytes 4-15: three words total, same for the whole super-block's lanes
+  w1, w2, w3 = _amd_load(raw[base+1]), _amd_load(raw[base+2]), _amd_load(raw[base+3])
+  sb = (subgroup & 3) * 8  # byte within word
+  byte1, byte2, byte3 = (w1 >> sb) & 255, (w2 >> sb) & 255, (w3 >> sb) & 255
+  scale = (subgroup < 4).where(byte1 & 63, (byte3 & 15) | ((byte1 >> 6) << 4))
+  minimum = (subgroup < 4).where(byte2 & 63, (byte3 >> 4) | ((byte2 >> 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()
 
@@ -118,52 +150,66 @@ def iq4_half_lut(device:str) -> Tensor:
                 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:
+def _q8_quantize_kernel(q:UOp, scale:UOp, xsum: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, lane = UOp.range(tokens*groups, 0, axis_type=AxisType.GLOBAL), 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))
+  qs = tuple((v/group_scale).round().clip(-127, 127).cast(dtypes.int8) for v in xs)
+  word = sum((v.cast(dtypes.uint8).cast(dtypes.uint32) << (i*8) for i, v in enumerate(qs)), UOp.const(0, dtypes.uint32))
+  # per-16 sums of the quantized values (lanes 0-3 / 4-7): Q4_K/Q5_K need the 32-sum, Q6_K the 16-sums
+  part = (lane < 8).where(sum((v.cast(dtypes.int32) for v in qs), UOp.const(0, dtypes.int32)), UOp.const(0, dtypes.int32))
+  gsum = [warp_reduce(((lane & 4).eq(h*4)).where(part, UOp.const(0, dtypes.int32)), full_wave=True) for h in range(2)]
+  store_half = (lane & 4) >> 2
+  stores = (q[token, group, lane.valid(lane < 8)].store(word),
+            UOp.group(scale[token, group.valid(lane.eq(0))].store(group_scale),
+                      xsum[token, group, store_half.valid(lane.eq(0) | lane.eq(4))].store(
+                        store_half.eq(0).where(gsum[0].float(), gsum[1].float()))))
   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]:
+def q8_quantize(x:Tensor, tokens:int, in_features:int) -> tuple[Tensor, 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
+  xsum = Tensor.empty(tokens, groups, 2, dtype=dtypes.float32, device=x.device)
+  q, scale, xsum = Tensor.custom_kernel(q, scale, xsum, x, fxn=functools.partial(_q8_quantize_kernel, tokens=tokens, in_features=in_features))[:3]
+  return q, scale, xsum
 
 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))
+  chunks = out.shape[2]
+  # two-dim global grid instead of one flat grid: no div/mods needed to decompose the gid
+  token_output = UOp.range(out.shape[0]*out_features, 0, axis_type=AxisType.GLOBAL)
+  chunk, lane = UOp.range(chunks, 1, axis_type=AxisType.GLOBAL), UOp.range(32, 2, axis_type=AxisType.LOCAL)
+  token, output = token_output // out_features, token_output % out_features
+  group = (lane+chunk*32).minimum(group_count-1)
+  value = group_dot(token, output, group) if chunks*32 == group_count else \
+    (lane+chunk*32 < group_count).where(group_dot(token, output, group), 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(
+  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:
+def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, xs: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)
+      qs_base, dot = base + (4 if ggml_type == Q4_K else 12) + (subgroup//2)*8, UOp.const(0, dtypes.int32)
+      # vectorize the 8 packed-weight words and (for Q5_K) the 32-byte high-bit bitmap
+      qs_pair = (_amd_load(raw[qs_base], 4), _amd_load(raw[qs_base+4], 4))
+      if ggml_type == Q5_K: qh_pair = (_amd_load(raw[base+4], 4), _amd_load(raw[base+8], 4))
       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)
+        word = (qs_pair[word_idx//4][word_idx%4] >> ((subgroup&1)*4).cast(dtypes.uint32)) & 0x0f0f0f0f
+        if ggml_type == Q5_K: word |= ((qh_pair[word_idx//4][word_idx%4] >> subgroup.cast(dtypes.uint32)) & 0x01010101) << 4
+        dot = _amd_dp4a(word, xwords[word_idx], dot)
       d, dmin, scale, minimum = _q5_scales(raw, base, subgroup)
-      return (dot.float()*d*scale - qsum.float()*dmin*minimum) * xd[token, group]
+      gsum = xs[token, group, 0].load() + xs[token, group, 1].load()
+      return (dot.float()*d*scale - gsum*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)
@@ -172,45 +218,57 @@ def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, out_features:int, in_
         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
+    # the packed rows were padded to 212 bytes (53 words) per 256-block in set_quantized: everything is word-aligned
+    base = (output*in_features//GGML_BLOCK_SIZE+block)*Q6_WORDS
+    # the subgroup's 8 ql words and 8 qh words are contiguous: two 16-byte vector loads each
+    lows = tuple(_amd_load(raw[base + (subgroup//4)*16 + (subgroup%2)*8 + half*4], 4) for half in range(2))
+    highs = tuple(_amd_load(raw[base + 32 + (subgroup//4)*8 + half*4], 4) for half in range(2))
     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))
+      within = (subgroup*32 + word_idx*4)%128
+      low = lows[word_idx//4][word_idx%4] >> ((within//64)*4).cast(dtypes.uint32)
+      high = highs[word_idx//4][word_idx%4] >> ((within//32)*2).cast(dtypes.uint32)
+      # 4 values per word: (low nibble) | (2 high bits << 4). values stay positive, so the int8-bitcast/-32 of the
+      # naive dequant is skipped and the -32 offset is applied later via the per-16 sums of the quantized inputs
+      word = (low & 0x0f0f0f0f) | ((high & 0x03030303) << 4)
       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)
+    scales = [((raw[base + 48 + (subgroup*2+i)//4] >> (((subgroup*2+i)%4)*8).cast(dtypes.uint32)) & 255)
+              .cast(dtypes.uint8).bitcast(dtypes.int8).float() for i in range(2)]
+    gsum = [xs[token, group, i].load() * 32 for i in range(2)]
+    return ((dots[0].float() - gsum[0])*scales[0] + (dots[1].float() - gsum[1])*scales[1]) * xd[token, group] * _half(raw[base+52] & 0xffff)
   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
+  # lane is a hardware WARP range (like the flash kernel): the fragment math stays visible without being
+  # range-split into nested loops, which would scramble the WMMA fragment layout
+  lane, wave = UOp.range(WARP_SIZE, -1, axis_type=AxisType.WARP), UOp.range(output_waves, 3, axis_type=AxisType.LOCAL)
+  col, half = lane % 16, 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)
+def _wmma_stores(out, outputs, tokens, accs, update, half, lane, wave, output_waves):
+  # the accumulator fragment halves are exchanged between lane pairs (l, l^16) through LDS (a ds_swizzle without CUSTOM)
+  flat_accs = [acc for output_accs in accs for acc in output_accs]
+  lds = UOp.placeholder((output_waves, 32, len(flat_accs)*8), dtypes.float32, slot=33, addrspace=AddrSpace.LOCAL)
+  stores = [lds[wave, lane, a*8+i].store(acc.after(update)[i].load()) for a,acc in enumerate(flat_accs) for i in range(8)]
+  lds = lds.after(UOp.barrier(UOp.group(*stores)))
+  def values(ai:int) -> tuple[UOp, ...]:
+    own = tuple(lds[wave, lane, ai*8+i].load() for i in range(8))
+    peer = tuple(lds[wave, lane ^ 16, ai*8+i].load() for i in range(8))
     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))]
+    return tuple(low.where(own[i], peer[i+4]) if j == 0 else low.where(peer[i], own[i+4]) for i in range(4) for j in range(2))
+  tt = len(tokens)
+  return [out[token, output].store(value) for ot,(output,output_accs) in enumerate(zip(outputs, accs))
+          for tile,(tile_tokens,_acc) in enumerate(zip(tokens, output_accs)) for token,value in zip(tile_tokens, values(ot*tt+tile))]
 
 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
+  output_waves, 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)
@@ -229,8 +287,8 @@ def _quant_linear_wmma(out, x, out_features, in_features, type_words, layout, de
         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=()))
+  stores = _wmma_stores(out, outputs, tokens, accs, update, physical_half, lane, wave, output_waves)
+  return UOp.group(*stores).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:
@@ -264,9 +322,10 @@ def _iq4_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, lut:UOp, out_features:i
       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))
+    # a subgroup-half gathers the lo (half=0) or hi (half=1) nibbles of byte pairs of each packed word
+    lut_pairs = (lut[(((raw[base+2+subgroup*4+i] >> (8*j+4*half)) & 15) |
+                      (((raw[base+2+subgroup*4+i] >> (8*j+8+4*half)) & 15) << 4)).cast(dtypes.weakint)]
+                 for i in range(4) for j 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")
 
@@ -287,67 +346,163 @@ def q8_linear(layer:Linear, x:Tensor) -> Tensor:
     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)
+  xq_, xd, xs = 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)
+  return run(decode, out, raw, xq_.uop, xd.uop, xs.uop)
+
+# ******** tiny dense fp16 gemv ********
+
+@functools.cache
+def _amd_f16_gemv_kernel(out:UOp, w:UOp, x:UOp, *rest:UOp, in_features:int, out_features:int, tokens:int) -> UOp:
+  bias: UOp|None = rest[0] if rest else None
+  # one block per (token, output row), 32 lanes accumulate 4-wide chunks of the row
+  lanes, val_chunk = WARP_SIZE, 4
+  token, out_row = UOp.range(tokens, 0, AxisType.GLOBAL), UOp.range(out_features, 1, AxisType.GLOBAL)
+  lane = UOp.range(lanes, 2, axis_type=AxisType.LOCAL)
+  per = in_features // (lanes * val_chunk)
+  assert per * lanes * val_chunk == in_features
+  w = w.reshape((out_features, per, lanes*val_chunk))
+  x = x.reshape((tokens, per, lanes*val_chunk))
+  acc = UOp.const(0, dtypes.float32)
+  for i in range(per):
+    for j in range(val_chunk):
+      acc = acc + w[out_row, i, lane*val_chunk + j].load().float() * x[token, i, lane*val_chunk + j].load().float()
+  total = warp_reduce(acc, full_wave=True)
+  if bias is not None: total = total + bias[token, out_row].load().float()
+  return out[token, out_row.valid(lane.eq(0))].store(total).end(token, out_row, lane).sink(arg=KernelInfo(name="linear_f16_gemv", opts_to_apply=()))
+
+def _view_back(t:Tensor) -> Tensor:
+  """strip top-of-chain CAST(s) from a lazy weight: reading the raw file bytes in the kernel instead of
+  materializing the cast into a fresh buffer every step"""
+  uop = t.uop
+  while uop.op is Ops.CAST: uop = uop.src[0]
+  return Tensor(uop).reshape(t.shape)
+
+def f16_gemv(layer:Linear, x:Tensor) -> Tensor:
+  tokens = prod(x.shape[:-1])
+  assert isinstance(tokens, int)
+  weight = _view_back(layer.weight)
+  x = x.contiguous() if x.dtype == dtypes.half else x.cast(dtypes.half).contiguous()
+  out = Tensor.empty(tokens, layer.out_features, dtype=dtypes.float32, device=x.device)
+  fxn = functools.partial(_amd_f16_gemv_kernel, in_features=layer.in_features, out_features=layer.out_features, tokens=tokens)
+  srcs = (out, weight.reshape(-1), x.reshape(tokens, layer.in_features)) + (() if layer.bias is None else (_view_back(layer.bias),))
+  return Tensor.custom_kernel(*srcs, fxn=fxn)[0].reshape(*x.shape[:-1], layer.out_features)
 
 # ******** flash attention on the KV cache ********
 
+def _vec_load(ptr:UOp, lanes:int) -> tuple[UOp, ...]:
+  if lanes == 1: return (ptr.load().float(),)
+  vec = _amd_load(ptr, lanes)
+  return tuple(vec[i].float() for i in range(lanes))
+
 @functools.cache
-def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, valid_kv_len, max_kv_len, block_n):
+def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, valid_kv_len, max_kv_len, block_n, waves=4):
   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=()))
+  G, CHUNK, DPL, WAVES = H // H_KV, block_n, D // WARP_SIZE, waves
+  assert CHUNK % WAVES == 0
+  SEC = CHUNK // WAVES  # keys each wave scans independently
+  live_chunks = (valid_kv_len+CHUNK-1)//CHUNK
+  live_chunks = min(live_chunks, out.shape[2]) if isinstance(live_chunks, int) else live_chunks.minimum(out.shape[2])
+  block_bhkv, block_chunk = UOp.range(B*H_KV, 0, AxisType.GLOBAL), UOp.range(live_chunks, 1, AxisType.GLOBAL)
+  lane, wave = UOp.range(WARP_SIZE, 2, axis_type=AxisType.LOCAL), UOp.range(WAVES, 3, axis_type=AxisType.LOCAL)
+  b, kv_head = block_bhkv // H_KV, block_bhkv % H_KV
+  # per-lane query fragments for every GQA head, kept packed in registers; unpacked at use
+  qf = tuple(_vec_load(q[b, kv_head*G+h, 0, lane*DPL], DPL) for h in range(G))
+  zerof = UOp.const(0, dtypes.float)
+  valids: list[UOp] = []
+  scores: list[list[UOp]] = [[zerof]*G for _ in range(SEC)]
+  vfrags: list[tuple[UOp, ...]] = [()]*SEC
+  for j in range(SEC):
+    key = block_chunk*CHUNK + wave*SEC + j
+    valid = key < valid_kv_len
+    valids.append(valid)
+    kfrag = _vec_load(cache_kv[0, b, kv_head, key, lane*DPL], DPL)
+    # V is prefetched in the score pass so both streams are in flight together
+    vfrags[j] = _vec_load(cache_kv[1, b, kv_head, key, lane*DPL], DPL)
+    for h in range(G):
+      s = warp_reduce(sum((qf[h][i]*kfrag[i] for i in range(DPL)), UOp.const(0, dtypes.float)), full_wave=True) * (1/math.sqrt(D))
+      scores[j][h] = valid.where(s, UOp.const(-math.inf, dtypes.float))
+  ninf = UOp.const(-math.inf, dtypes.float)
+  row_max = [functools.reduce(UOp.maximum, (scores[j][h] for j in range(SEC)), ninf) for h in range(G)]
+  accs:list[list[UOp]] = [[UOp.const(0, dtypes.float)] * DPL for _ in range(G)]
+  row_sums:list[UOp] = [UOp.const(0, dtypes.float) for _ in range(G)]
+  for j in range(SEC):
+    for h in range(G):
+      beta = valids[j].where(((scores[j][h]-row_max[h])*LOG2E).exp2(), UOp.const(0, dtypes.float))
+      accs[h] = [a + beta*v for a, v in zip(accs[h], vfrags[j])]
+      row_sums[h] = row_sums[h] + beta
+  # exchange across the block's waves through LDS (fp16 halves LDS so more blocks fit per CU)
+  acc_lds = UOp.placeholder((WAVES, G, D), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)
+  ml_lds = UOp.placeholder((WAVES, G, 2), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
+  lds_acc = acc_lds.reshape(WAVES, G, WARP_SIZE, DPL)
+  stores = [lds_acc[wave, h, lane].store(UOp.stack(*accs[h]).cast(dtypes.half)) for h in range(G)]
+  # NOTE: duplicate stores of the same value from every lane are harmless here
+  stores += [ml_lds[wave, h, i].store(x) for h in range(G) for i, x in enumerate((row_max[h], row_sums[h]))]
+  barrier = UOp.barrier(UOp.group(*stores))
+  acc_lds, ml_lds = acc_lds.after(barrier), ml_lds.after(barrier)
+  tid = wave*WARP_SIZE + lane
+  final_stores:list[UOp] = []
+  for i in range(-(-G*D//(WAVES*WARP_SIZE))):
+    flat = tid + i*WAVES*WARP_SIZE
+    h, d = flat // D, flat % D
+    M = functools.reduce(UOp.maximum, (ml_lds[w, h, 0].load() for w in range(WAVES)), ninf)
+    val = sum((((ml_lds[w, h, 0].load()-M)*LOG2E).exp2() * acc_lds[w, h, d].load().float() for w in range(WAVES)), UOp.const(0, dtypes.float))
+    oidx = out[b, kv_head*G + h, block_chunk, d]
+    if G*D % (WAVES*WARP_SIZE): oidx = out[b, (kv_head*G + h).valid(flat < G*D), block_chunk, d]
+    final_stores.append(oidx.store(val))
+  hstat = tid
+  M = functools.reduce(UOp.maximum, (ml_lds[w, hstat, 0].load() for w in range(WAVES)), ninf)
+  L = sum((((ml_lds[w, hstat, 0].load()-M)*LOG2E).exp2() * ml_lds[w, hstat, 1].load() for w in range(WAVES)), UOp.const(0, dtypes.float))
+  q_head = (kv_head*G + hstat).valid(hstat < G) if WAVES*WARP_SIZE > G else kv_head*G + hstat
+  final_stores += [stats[b, q_head, block_chunk, 0].store(M), stats[b, q_head, block_chunk, 1].store(L)]
+  return UOp.group(*final_stores).end(lane, wave, block_chunk, block_bhkv).sink(arg=KernelInfo(name="flash_decode_partial", opts_to_apply=()))
+
+@functools.cache
+def _amd_flash_decode_combine(o:UOp, partial:UOp, stats:UOp, live:int|UOp) -> UOp:
+  # one wave per (batch, head, 64-dim tile): every lane redundantly weights its chunks; no cross-lane traffic
+  live = _unbind(live)
+  B, H, C, D = cast(tuple[int, int, int, int], partial.shape)
+  DT = 64 if D % 64 == 0 else WARP_SIZE  # dims per block
+  assert D % DT == 0
+  block_bh, block_dt = UOp.range(B*H, 0, AxisType.GLOBAL), UOp.range(D//DT, 1, AxisType.GLOBAL)
+  lane = UOp.range(WARP_SIZE, 2, axis_type=AxisType.LOCAL)
+  b, h = block_bh // H, block_bh % H
+  NPD = DT // WARP_SIZE  # output dims per lane
+  dims = tuple(block_dt*DT + lane*NPD + i for i in range(NPD))
+  chunk = UOp.range(live, 100, AxisType.REDUCE)
+  def iloop(ph, val): return ph.store(ph.const_like(val))
+  chunk_max = UOp.placeholder((1,), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  chunk_max_i = chunk_max.after(iloop(chunk_max, -math.inf))
+  update0 = chunk_max_i.store(chunk_max_i.after(chunk).maximum(stats[b, h, chunk, 0].load())).end(chunk)
+  chunk_max = chunk_max_i.after(update0)
+  chunk2 = UOp.range(live, 101, AxisType.REDUCE)
+  acc = UOp.placeholder((NPD,), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+  weight_sum = UOp.placeholder((1,), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  acc_i, weight_sum_i = acc.after(iloop(acc, 0)), weight_sum.after(iloop(weight_sum, 0))
+  w = ((stats[b, h, chunk2, 0].load()-chunk_max)*LOG2E).exp2()
+  update1 = UOp.group(*[acc_i[i].store(acc_i.after(chunk2)[i].load() + w*partial[b, h, chunk2, d].load()) for i, d in enumerate(dims)],
+                      weight_sum_i[0].store(weight_sum_i.after(chunk2)[0].load() + w*stats[b, h, chunk2, 1].load())).end(chunk2)
+  acc, weight_sum = acc_i.after(update1), weight_sum_i.after(update1)
+  inv = 1 / weight_sum[0].load()
+  return UOp.group(*[o[b, h, 0, d].store(acc[i].load() * inv) for i, d in enumerate(dims)]) \
+    .end(lane, block_dt, block_bh).sink(arg=KernelInfo(name="flash_decode_combine", 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)
+  chunks = min(256, max_kv_len // 64)
   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)
+  fxn = functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len, block_n=64, waves=16)
   partial, stats = Tensor.custom_kernel(partial, stats, q, cache_kv, fxn=fxn)[:2]
-  live = (valid_kv_len+127)//128
+  live = (valid_kv_len+63)//64
   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)
+  out = Tensor.empty(B, H, 1, D, dtype="float32", device=q.device)
+  fxn = functools.partial(_amd_flash_decode_combine, live=live)
+  return Tensor.custom_kernel(out, partial, stats, fxn=fxn)[0]
 
 @functools.cache
 def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, valid_kv_len:int|UOp, q_start:int|UOp|None=None) -> UOp:

tinygrad/mixin/elementwise.py

@@ -870,7 +870,7 @@ class ElementwiseMixin(CreationMixin):
     print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).asinh().numpy())
     ```
     """
-    return (self + (self.square() + 1).sqrt()).log()
+    return (sg:=(self<0).where(-1.0, 1.0)) * (self*sg + (self.square() + 1).sqrt()).log()
 
   def acosh(self) -> Self:
     """

tinygrad/mixin/gradient.py

@@ -1,6 +1,6 @@
 from typing import cast
 import math, dataclasses
-from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata, broadcast_axes
+from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata, broadcast_axes, to_max_shape
 from tinygrad.helpers import argsort
 from tinygrad.dtype import sum_acc_dtype
 from tinygrad.function import renumber_invalid_outputs
@@ -31,6 +31,13 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
     return (None,) + k.arg.grad_fxn(on_dev(ctx, 0), k)
   assert fxn.op is Ops.TUPLE, f"expected TUPLE body for gradient, got {fxn.op}"
   params = {x.arg.slot:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM}
+  # grads are collected at the flat param storage: reshape to each arg's view (max view shrunk to symbolic)
+  def shaped_grad(grad:UOp, i:int) -> UOp:
+    a = args[i]
+    if a.axis is not None and isinstance(a.device, tuple): return UOp.shared_view(grad, a.shard_shape, a.axis, count=len(a.device))
+    vshape = a._shape
+    grad = grad.reshape(to_max_shape(vshape))
+    return grad.shrink_to(vshape) if to_max_shape(vshape) != tuple(vshape) else grad
   grad_args = ctx.src
   root_grad = UOp(Ops.TUPLE, src=tuple(UOp(Ops.NOOP) if g.op is Ops.NOOP else
     g if g.device is None else g.param_like(len(args)+i) for i,g in enumerate(grad_args)))
@@ -39,7 +46,7 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
   fwd_subs = {src: src.param_like(len(args)+len(grad_args)+i) for i, src in enumerate(fxn.src)} if k.arg.precompile else {}
   fwd_outs = tuple(k.gettuple(i) for i in range(len(fxn.src))) if k.arg.precompile else ()
   # collect needed gradient bodies, compact unused params, create a single backward CALL
-  grad_bodies = [(i, grads[p]) for i in needed if (p:=params.get(i)) is not None and p in grads]
+  grad_bodies = [(i, shaped_grad(grads[p], i)) for i in needed if (p:=params.get(i)) is not None and p in grads]
   bwd_body = UOp.maketuple(*(gb for _, gb in grad_bodies)).substitute(fwd_subs, walk=True)
   bwd_body = renumber_invalid_outputs(bwd_body)
   bwd_body, compact_args = _compact_params(bwd_body, (*args, *grad_args, *fwd_outs))

tinygrad/nn/__init__.py

@@ -359,7 +359,7 @@ def _embedding_bwd(grad_emb:UOp, call:UOp) -> tuple:
     if device in ("CPU", "NULL"): atomic_arg = "__atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED);"
     elif device == "AMD": atomic_arg = "__hip_atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);"
     else: raise NotImplementedError(f"no atomics for device {device}")
-    atomic = UOp(Ops.CUSTOM, src=(grad_weight.index(local_token_id, j_idx), grad_val), arg = atomic_arg)
+    atomic = UOp(Ops.CUSTOM, src=(grad_weight.index(local_token_id, j_idx), grad_val), arg=(atomic_arg, dtypes.void))
     return atomic.end(i, j_outer, j_inner).sink(arg=KernelInfo(name="embedding_bwd", opts_to_apply=()))
 
   grad_weight_uop = grad_weight_uop.custom_kernel(grad_emb, idx, fxn=_embedding_bwd_kernel)[0]

tinygrad/renderer/__init__.py

@@ -1,12 +1,17 @@
 from __future__ import annotations
 from typing import Callable, cast
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from tinygrad.helpers import prod, Target, EMULATED_DTYPES
 from tinygrad.uop.ops import Ops, UOp, sint, ssimplify, smin, GroupOp, PatternMatcher
 from tinygrad.dtype import AddrSpace, DType, dtypes
 from tinygrad.codegen.opt.tc import TensorCore
 from tinygrad.device import Compiler
 
+# an access takes its dtype from the buffer it indexes, so accessing at another dtype restates the storage on the buffer that owns it
+def with_storage(x:UOp, dt:DType) -> UOp:
+  if x.op in {Ops.PARAM, Ops.BUFFER}: return x.replace(dtype=None, arg=replace(x.arg, dtype=dt))
+  return x.replace(dtype=None, src=(with_storage(x.src[0], dt),)+x.src[1:])
+
 @dataclass(frozen=True)
 class Estimates:
   # number of FLOPS used in the Kernel

tinygrad/renderer/cstyle.py

@@ -71,7 +71,7 @@ base_rewrite = PatternMatcher([
    f"({', '.join(f'({ctx.render_type(y)})({ctx[y]})' for y in x.src[1:])}))" + (";" if x.dtype is dtypes.void else "")),
 
   # custom passes through with format
-  (UPat((Ops.CUSTOM, Ops.CUSTOMI), name="x"), lambda ctx,x: x.arg.format(*[ctx[y] for y in x.src])),
+  (UPat((Ops.CUSTOM, Ops.CUSTOMI), name="x"), lambda ctx,x: x.arg[0].format(*[ctx[y] for y in x.src])),
 ])
 
 def create_non_native_float_pats(dts:tuple[DType, ...], casting:bool=True):
@@ -188,10 +188,12 @@ class CStyleLanguage(Renderer):
     return prefix + self.type_map.get(dtype, dtype.name) + suffix
 
   def render_type(self, u:UOp): return self._render_dtype(u.dtype, u.max_numel(), u.addrspace, shape=u._shape)
-  def render_access(self, u:UOp):
+  def render_ptr(self, u:UOp):
+    # the address of an access, vector-cast if the access reads/writes more lanes than the pointer's scalar type
     if u.max_numel() > 1 or u.dtype != u.src[0].dtype:
-      return f"*(({self._render_dtype(u.dtype, u.max_numel(), u.addrspace, override_ptr=True, shape=u._shape)})({self[u]}))"
-    else: return f"*{self[u]}"
+      return f"(({self._render_dtype(u.dtype, u.max_numel(), u.addrspace, override_ptr=True, shape=u._shape)})({self[u]}))"
+    else: return f"{self[u]}"
+  def render_access(self, u:UOp): return f"*{self.render_ptr(u)}"
   def render_cast(self, u:UOp, val:str) -> str: return f"({self.render_type(u)})({val})"
 
   # LEGACY
@@ -509,6 +511,9 @@ class HIPRenderer(CStyleLanguage):
         (UPat(Ops.CAST, dtypes.float, (UPat.var("y", dtypes.fp8s),), name="x",),
           lambda ctx,x,y: f"__builtin_amdgcn_cvt_f32_{('fp8', 'bf8')[fp8_index(y.dtype)]}((unsigned int){ctx[x.src[0]]}, 0)"),
       ]) + base_rewrite
+    # a LOAD flagged nontemporal renders as the cache-bypassing builtin (only used on global loads)
+    self.string_rewrite = PatternMatcher([(UPat(Ops.LOAD, arg="nontemporal", src=(UPat.var("bidx"),)),
+      lambda ctx,bidx: f"__builtin_nontemporal_load({ctx.render_ptr(bidx)})")]) + self.string_rewrite
 
   # https://clang.llvm.org/docs/AttributeReference.html#amdgpu-flat-work-group-size
   # NOTE: this makes hlb_cifar10 twice as fast, there may be more gains in tweaking these parameters

tinygrad/renderer/nir.py

@@ -1,7 +1,7 @@
 from typing import Callable, Any
 from tinygrad.dtype import AddrSpace, DType, dtypes, truncate
 from tinygrad.helpers import DEBUG, OSX, unwrap, fromimport, Target, is_image_shape, round_up
-from tinygrad.renderer import Renderer
+from tinygrad.renderer import Renderer, with_storage
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str
 from tinygrad.runtime.autogen import mesa, libc
@@ -123,11 +123,12 @@ class NIRRenderer(Renderer):
   extra_matcher = PatternMatcher([
     # from ptx
     (UPat.var('x', dtype=dtypes.bool)<UPat.var('y'), lambda x,y: (x^True)&y),
-    # load/store bool -> uint8
+    # a bool is one bit in NIR but a byte in memory, so every access to a bool buffer goes through a uint8 view of it
     (UPat(Ops.LOAD, dtypes.bool, name="x"),
-     lambda x: x.replace(dtype=dtypes.uint8, src=x.src[0:1]+((x.src[1].cast(dtypes.uint8),) if len(x.src)>=2 else ())+x.src[2:]).cast(dtypes.bool)),
-    (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
-     lambda x: x.replace(src=(x.src[0], x.src[1].cast(dtypes.uint8))+x.src[2:])),
+     lambda x: x.replace(dtype=None, src=(with_storage(x.src[0], dtypes.uint8),)+((x.src[1].cast(dtypes.uint8),) if len(x.src)>=2 else ())
+      +x.src[2:]).cast(dtypes.bool)),
+    (UPat(Ops.STORE, src=(UPat(name="idx"), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
+     lambda x,idx: x.replace(src=(with_storage(idx, dtypes.uint8), x.src[1].cast(dtypes.uint8))+x.src[2:])),
     # NIR requires shift amount to be 32 bit: https://docs.mesa3d.org/nir/alu.html#nir-alu-op-ishl
     (UPat((Ops.SHL, Ops.SHR), name="x"), lambda x: x.replace(src=(x.src[0], x.src[1].cast(dtypes.uint))) if x.src[1].dtype.bitsize != 32 else None),
     # OpConvertFToU is undefined if Result Type is not wide enough, cast through int32

tinygrad/renderer/ptx.py

@@ -4,7 +4,7 @@ from collections import defaultdict
 from tinygrad.codegen.opt import tc
 from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp
 from tinygrad.dtype import dtypes, DType, AddrSpace
-from tinygrad.renderer import Renderer
+from tinygrad.renderer import Renderer, with_storage
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.helpers import flatten, prod, unwrap, Target
 
@@ -45,12 +45,12 @@ ptx_matcher = PatternMatcher([
   # upcast to float32 all the ops that don't support half
   (UPat(doesnt_support_half, dtype=dtypes.half, name="x"),
     lambda x: (UOp(x.op, src=tuple(vv.cast(dtypes.float32) for vv in x.src), arg=x.arg).cast(dtypes.half))),
-  # load/store bool -> uint8 (only for memory, not registers)
+  # a bool is a predicate register in PTX but a byte in memory, so a bool buffer is accessed through a uint8 view of it
   (UPat(Ops.LOAD, dtypes.bool, src=(UPat(name="idx"),), name="x", allow_any_len=True),
-   lambda x,idx: UOp(x.op, dtypes.uint8, x.src[0:1] + ((x.src[1].cast(dtypes.uint8),) if len(x.src) >= 2 else ()) + x.src[2:]).cast(dtypes.bool) \
-     if idx.addrspace != AddrSpace.REG else None),
+   lambda x,idx: x.replace(dtype=None, src=(with_storage(idx, dtypes.uint8),) + ((x.src[1].cast(dtypes.uint8),) if len(x.src) >= 2 else ())
+     + x.src[2:]).cast(dtypes.bool) if idx.addrspace != AddrSpace.REG else None),
   (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),
+   lambda x,idx: x.replace(src=(with_storage(idx, dtypes.uint8), 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),

tinygrad/runtime/autogen/__init__.py

@@ -10,10 +10,14 @@ rocr_src = "https://github.com/ROCm/rocm-systems/archive/refs/tags/rocm-7.1.1.ta
 linux_headers_deb = "https://snapshot.debian.org/archive/debian/20260207T145350Z/pool/main/l/linux/linux-libc-dev_6.18.9-1_all.deb"
 linux_headers_kern_deb = "https://snapshot.debian.org/archive/debian/20260207T145350Z/pool/main/l/linux/linux-headers-6.18.9+deb14-common_6.18.9-1_all.deb"
 liburing_src = "https://raw.githubusercontent.com/axboe/liburing/refs/tags/liburing-2.14/src/include/liburing.h"
+bnxt_src = ["https://raw.githubusercontent.com/torvalds/linux/v6.18/drivers/" + s for s in
+            ("infiniband/hw/bnxt_re/roce_hsi.h", "infiniband/hw/bnxt_re/qplib_rcfw.h", "infiniband/hw/bnxt_re/qplib_res.h",
+             "net/ethernet/broadcom/bnxt/bnxt_hwrm.h")]
 ggml_common_src = "https://raw.githubusercontent.com/ggml-org/ggml/d4fcfe88a8bcf5c9840be14be6c2fbf1f5b3b2db/src/ggml-common.h"
 cudart_src = "https://developer.download.nvidia.com/compute/cuda/redist/cuda_cudart/linux-x86_64/cuda_cudart-linux-x86_64-12.0.146-archive.tar.xz"
 nvrtc_src = "https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvrtc/linux-x86_64/cuda_nvrtc-linux-x86_64-12.0.140-archive.tar.xz"
 opencl_src = "https://github.com/KhronosGroup/OpenCL-Headers/archive/2e30669d48718fd460f085b4b35b160dad51ce9d.tar.gz"
+comgr_2_src = "https://repo.radeon.com/rocm/apt/6.2/pool/main/c/comgr/comgr_2.8.0.60200-66~24.04_amd64.deb"
 macossdk = "/var/db/xcode_select_link/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk"
 
 llvm_lib = (
@@ -50,6 +54,8 @@ def load(name, files, **kwargs):
     if srcs: td.cleanup()
   return importlib.import_module(f"{path}.{name.replace('/', '.')}")
 
+def _extract_deb(path:str): subprocess.run("ar x *.deb && tar xf data.tar.*", cwd=path, shell=True, check=True)
+
 def __getattr__(nm):
   match nm:
     case "libc":
@@ -99,14 +105,11 @@ def __getattr__(nm):
     # this defines all syscall numbers. should probably unify linux autogen?
     case "io_uring":
       return load("io_uring", ["{}/liburing.h", "{}/usr/include/linux/io_uring.h", "{}/usr/include/asm-generic/unistd.h"],
-                  args=["-I{}/usr/include"], srcs=[linux_headers_deb, liburing_src], rules=[('__NR', 'NR')],
-                  preprocess=lambda path: subprocess.run(f"ar x {linux_headers_deb.split('/')[-1]} && tar xf data.tar.xz", cwd=path, shell=True, check=True))
+                  args=["-I{}/usr/include"], srcs=[linux_headers_deb, liburing_src], rules=[('__NR', 'NR')], preprocess=_extract_deb)
     case "llvm": return load("llvm", lambda: [system("llvm-config-20 --includedir")+"/llvm-c/**/*.h"], dll=llvm_lib,
                              args=lambda: system("llvm-config-20 --cflags").split(), recsym=True, prolog=["from tinygrad.helpers import WIN, OSX"])
-    case "pci": return load("pci", ["{}/usr/include/linux/pci_regs.h"], srcs=linux_headers_deb,
-                             preprocess=lambda path: subprocess.run(f"ar x {linux_headers_deb.split('/')[-1]} && tar xf data.tar.xz", cwd=path, shell=True, check=True))
-    case "vfio": return load("vfio", ["{}/usr/include/linux/vfio.h"], args=["-I{}/usr/include"], srcs=linux_headers_deb,
-                             preprocess=lambda path: subprocess.run(f"ar x {linux_headers_deb.split('/')[-1]} && tar xf data.tar.xz", cwd=path, shell=True, check=True))
+    case "pci": return load("pci", ["{}/usr/include/linux/pci_regs.h"], srcs=linux_headers_deb, preprocess=_extract_deb)
+    case "vfio": return load("vfio", ["{}/usr/include/linux/vfio.h"], args=["-I{}/usr/include"], srcs=linux_headers_deb, preprocess=_extract_deb)
     # could add rule: WGPU_COMMA -> ','
     case "webgpu": return load("webgpu", [root/"extra/webgpu/webgpu.h"], dll=webgpu_lib,
                                prolog=["from tinygrad.helpers import WIN, OSX", "import sysconfig, os"])
@@ -116,9 +119,10 @@ def __getattr__(nm):
                             dll="os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamdhip64.so'",
                             args=["-D__HIP_PLATFORM_AMD__", "-I/opt/rocm/include", "-x", "c++"], prolog=["import os"])
     case "comgr" | "comgr_3":
-      return load("comgr_3" if nm == "comgr_3" else "comgr", ["/opt/rocm/include/amd_comgr/amd_comgr.h"],
-                  dll= "[os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamd_comgr.so', 'amd_comgr']",
-                  args=["-D__HIP_PLATFORM_AMD__", "-I/opt/rocm/include", "-x", "c++"], prolog=["import os"])
+      prefix = "{}/opt/rocm-6.2.0" if nm == "comgr" else "/opt/rocm"
+      return load(nm, [f"{prefix}/include/amd_comgr/amd_comgr.h"], dll="[os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamd_comgr.so', 'amd_comgr']",
+                  args=["-D__HIP_PLATFORM_AMD__", f"-I{prefix}/include", "-x", "c++"], prolog=["import os"], srcs=comgr_2_src if nm == "comgr" else None,
+                  **({'preprocess':_extract_deb} if nm == "comgr" else {}))
     case "hsa": return load("hsa", [*[f"{{}}/projects/rocr-runtime/runtime/hsa-runtime/core/inc/{s}.h" for s in ["registers"]],
                                     *[f"{{}}/projects/rocr-runtime/runtime/hsa-runtime/inc/{s}.h" for s in [
                                         "hsa", "hsa_ext_amd", "amd_hsa_signal", "amd_hsa_queue", "amd_hsa_kernel_code",
@@ -181,6 +185,17 @@ def __getattr__(nm):
       return load("mlx5", [root/"extra/mlx_driver/mlx5.h", f"{kh}/mlx5_ifc.h"], srcs=linux_headers_kern_deb,
                   args=["-Du8=unsigned char", "-Du16=unsigned short", "-Du32=unsigned int", "-Du64=unsigned long long",
                         "-D__be16=unsigned short", "-D__be32=unsigned int", "-D__be64=unsigned long long", f"-I{kh}"],
-                  preprocess=lambda path: subprocess.run(f"ar x {linux_headers_kern_deb.split('/')[-1]} && tar xf data.tar.xz",
-                                                         cwd=path, shell=True, check=True))
+                  preprocess=_extract_deb)
+    case "bnxt":
+      kh = "{}/usr/src/linux-headers-6.18.9+deb14-common/include"
+      return load("bnxt", [f"{kh}/linux/bnxt/hsi.h", *[f"{{}}/{s.split('/')[-1]}" for s in bnxt_src]],
+                  srcs=[linux_headers_kern_deb, *bnxt_src],
+                  args=["-Du8=unsigned char", "-Du32=unsigned int", "-Du64=unsigned long long", "-D__le16=unsigned short",
+                        "-D__le32=unsigned int", "-D__le64=unsigned long long", "-D__be16=unsigned short", "-D__be32=unsigned int", f"-I{kh}"],
+                  patterns=[r"hwrm_((ver_get|func_(qcaps|qcfg|reset|drv_rgtr|backing_store_(qcaps|cfg)_v2)|stat_ctx_alloc|ring_alloc"
+                            r"|vnic_(alloc|cfg)|cfa_l2_filter_alloc|port_phy_cfg)_(input|output)|(cmd|resp)_hdr)$",
+                            r"((cmdq|creq)_(base|init|add_gid|create_(cq|qp)|initialize_fw|modify_qp|query_version|register_mr)(_resp)?"
+                            r"|cq_(base|req)|sq_(rdma_hdr|sge))$",
+                            r"(BNXT|CMDQ|CREQ|CQ|SQ|DBC|PTU|RCFW|HWRM|VNIC|RING_ALLOC|STAT_CTX|CFA_L2_FILTER|PORT_PHY_CFG|FIRMWARE_FIRST"
+                            r"|FUNC_(QCAPS|QCFG|RESET|DRV_RGTR|BACKING_STORE))_"], preprocess=_extract_deb)
     case _: raise AttributeError(f"no such autogen: {nm}")

tinygrad/runtime/ops_cpu.py

@@ -25,25 +25,25 @@ MAX_ARGS, CMD_SIZE, RING_SLOTS, FUNCS = 63, 64, (16 << 10), (() if WIN else ('cl
 # 1. workers
 
 def signal_prog():
-  val = UOp.param(1, dtypes.int, (), vmin_vmax=(0, dtypes.int.max), name="value", addrspace=AddrSpace.ALU)
-  return UOp.param(0, dtypes.uint32, (1,))[0].store(val.cast(dtypes.uint32))
+  val = UOp.param(1, dtypes.int, vmin_vmax=(0, dtypes.int.max), name="value", addrspace=AddrSpace.ALU)
+  return UOp.param(0, dtypes.uint32, 1)[0].store(val.cast(dtypes.uint32))
 
 def wait_prog():
-  val = UOp.param(1, dtypes.int, (), vmin_vmax=(0, dtypes.int.max), name="value", addrspace=AddrSpace.ALU)
-  return (v:=UOp.param(0, dtypes.uint32, (1,), volatile=True).after(l:=UOp.loop(0))[0].load()).end(l, v < val.cast(dtypes.uint32))
+  val = UOp.param(1, dtypes.int, vmin_vmax=(0, dtypes.int.max), name="value", addrspace=AddrSpace.ALU)
+  return (v:=UOp.param(0, dtypes.uint32, 1, volatile=True).after(l:=UOp.loop(0))[0].load()).end(l, v < val.cast(dtypes.uint32))
 
 def timestamp_prog():
   if WIN: val = UOp.const(0, dtypes.uint64)
   else:
-    fn, ts = UOp.param(1, dtypes.uint64, (1,)), UOp.placeholder((2,), dtypes.uint64, slot=0, addrspace=AddrSpace.REG)
+    fn, ts = UOp.param(1, dtypes.uint64, 1), UOp.placeholder((2,), dtypes.uint64, slot=0, addrspace=AddrSpace.REG)
     call = fn[0].load().call(UOp.const(6 if OSX else 1, dtypes.int), ts[0], ret_dtype=dtypes.void) # clock_gettime(CLOCK_MONOTONIC, &ts)
     val = ts.after(call)[0].load() * 1_000_000_000 + ts.after(call)[1].load()
-  return UOp.param(0, dtypes.uint64, (1,))[0].store(val)
+  return UOp.param(0, dtypes.uint64, 1)[0].store(val)
 
 def worker_prog():
-  ring = UOp.param(0, dtypes.uint64, (RING_SLOTS * CMD_SIZE,), volatile=True)
-  wait, done = UOp.param(1, dtypes.uint64, (1,), volatile=True), UOp.param(2, dtypes.uint64, (1,), volatile=True)
-  sem, cur = UOp.param(3, dtypes.uint64, (1,)), UOp.range(2**64-1, 0, dtype=dtypes.uint64) # sem is unused on windows, it has to come last
+  ring = UOp.param(0, dtypes.uint64, RING_SLOTS * CMD_SIZE, volatile=True)
+  wait, done = UOp.param(1, dtypes.uint64, 1, volatile=True), UOp.param(2, dtypes.uint64, 1, volatile=True)
+  sem, cur = UOp.param(3, dtypes.uint64, 1), UOp.range(2**64-1, 0, dtype=dtypes.uint64) # sem is unused on windows, it has to come last
 
   # spin on windows, sem_wait to sleep on posix
   if WIN: ready = (v:=wait.after(lw:=UOp.loop(1), cur)[0].load()).end(lw, v <= cur)

tinygrad/runtime/ops_qcom.py

@@ -20,10 +20,11 @@ BUFTYPE_BUF, BUFTYPE_TEX, BUFTYPE_IBO = 0, 1, 2
 def dcache_flush():
   from tinygrad.uop.ops import UOp, Ops, KernelInfo
   from tinygrad.codegen import to_program
-  buf, n = UOp.param(0, dtypes.uint8, shape=(1,)), UOp.param(1, dtypes.int, shape=(), name="n", addrspace=AddrSpace.ALU)
+  buf, n = UOp.param(0, dtypes.uint8, 1), UOp.param(1, dtypes.int, name="n", addrspace=AddrSpace.ALU)
   i = UOp.range(n, 0, dtype=dtypes.int)
-  flush = UOp(Ops.CUSTOM, src=(buf.index(i * 64),), arg='__asm__ volatile("dc cvac, %0" :: "r"({0}) : "memory");')
-  sink = UOp.sink(flush.end(i), UOp(Ops.CUSTOM, arg='__asm__ volatile("dsb sy" ::: "memory");'), arg=KernelInfo(name="dcache_flush"), tag=1)
+  flush = UOp(Ops.CUSTOM, src=(buf.index(i * 64),), arg=('__asm__ volatile("dc cvac, %0" :: "r"({0}) : "memory");', dtypes.void))
+  sink = UOp.sink(flush.end(i), UOp(Ops.CUSTOM, arg=('__asm__ volatile("dsb sy" ::: "memory");', dtypes.void)),
+                  arg=KernelInfo(name="dcache_flush"), tag=1)
   prg = to_program(sink, Device["CPU"].renderer)
   return Device["CPU"].runtime(prg.to_elf())
 

tinygrad/runtime/support/hcq2.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 from typing import cast, TypeVar, Generic, Any, Sequence, Iterable
 import struct, functools, time, collections, itertools, decimal, statistics
-from dataclasses import replace, dataclass
+from dataclasses import replace, dataclass, field
 from tinygrad.helpers import suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE, unwrap, PROFILE
 from tinygrad.helpers import to_tuple, round_up, partition, panic, ContextVar, perf_counter_us, Context
 from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer, DepsTracker
@@ -30,9 +30,8 @@ class HCQInfo:
   device:tuple[str, ...]
   estimates:Estimates = Estimates()
 
-  input_idxs:tuple[tuple[tuple[str, ...], tuple[int, ...]], ...] = () # per inputs table: (devices, indexes into input_uops)
-  inputs:int|None = None # index of the inputs table in call.src
-  # per kernel: (devices, name, estimates, timestamps, profile key)
+  inputs:int|None = None
+  input_addrs:tuple[tuple[str, UOp], ...] = () # (device, lane arg uop)
   kernels:tuple[tuple[tuple[str, ...], str, Estimates, tuple[int, ...], bytes], ...] = ()
 
 def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
@@ -44,6 +43,8 @@ def unwrap_mstack(u:UOp) -> tuple[UOp, ...]:
 def unwrap_view(v:UOp) -> tuple[UOp, int]:
   return unwrap_view(v.src[0]) if v.op is Ops.BITCAST else (v.src[0], v.src[1].val) if v.op is Ops.SHRINK else (v, 0)
 
+def _lane(u:UOp, lane:int) -> UOp: return u.src[lane] if u.op is Ops.MSTACK else u.mselect(lane) if len(to_tuple(u.device)) > 1 else u
+
 # patches
 
 def is_value_known_at_link(val:UOp) -> bool:
@@ -138,7 +139,7 @@ def _get_enqueue_devs(call:UOp) -> Any|None:
 
 def copy_with_kernel(call:UOp, dst:UOp, src:UOp) -> UOp|None:
   if (devs:=_get_enqueue_devs(call)) is None or Device[(dev:=to_tuple(devs)[0])].has_copy_queue: return None
-  d, s = (UOp.param(i, dst.dtype, (n:=dst.max_numel(),), device=devs) for i in range(2))
+  d, s = (UOp.param(i, dst.dtype, n:=dst.max_numel(), device=devs) for i in range(2))
   ast = d.index(r:=UOp.range(n, 0)).store(s.index(r).load()).end(r).sink(arg=KernelInfo(name="copy"), tag=1)
   return call.replace(src=(to_program(ast, Device[dev].renderer), dst, src))
 
@@ -154,138 +155,131 @@ pm_insert_copy_staging = PatternMatcher([
 class HCQDepsTracker(DepsTracker):
   @staticmethod
   def _key(buf:Any) -> tuple[Any, int, int]:
+    if isinstance(buf, UOp) and buf.op is Ops.MSELECT: buf = buf.src[0]
     return (buf.arg.slot, 0, buf.max_numel() * buf.dtype.itemsize) if isinstance(buf, UOp) else DepsTracker._key(buf)
 
-def _get_call_bufs_by_lane(call:UOp, devices:tuple[str, ...]) -> list[list[Any]]:
-  refs = get_call_arg_uops(call)
-  return [[b if b.op is Ops.PARAM else mb.bufs[lane] if isinstance(mb:=b.buffer, MultiBuffer) else mb for b in refs] for lane in range(len(devices))]
+@dataclass(frozen=True)
+class Dep: dev:str; queue:str; tag:int; lane:int # producer submit (dev, queue, tag) awaited by consumer lane # noqa: E702
 
-def _get_deps(ctx:DepsTracker, bufs_by_lane:list[list[Any]], write, key:tuple[tuple[str, ...], str, int]) -> list[tuple[tuple, int, int]]:
-  dep_lanes:list[tuple[tuple, int, int]] = []
+@dataclass
+class BatchCtx:
+  batch:list[tuple[UOp, tuple[str, ...]]]; batch_info:list[tuple[tuple[str, ...], str]]; profile:bool # noqa: E702
+  tracker:HCQDepsTracker = field(default_factory=HCQDepsTracker); signal_tags:set[int] = field(default_factory=set) # noqa: E702
+  slots:dict[str, int] = field(default_factory=lambda: collections.defaultdict(lambda: next(UOp.unique_num)))
+
+def _get_call_bufs_by_lane(call:UOp, devices:tuple[str, ...]) -> list[list[Any]]:
+  return [[b if (b:=_lane(a, lane)).op is Ops.PARAM or (b.op is Ops.MSELECT and b.src[0].op is Ops.PARAM) else b.buffer
+           for a in get_call_arg_uops(call)] for lane in range(len(devices))]
+
+def _wait_ins(ctx:BatchCtx, bufs_by_lane:list[list[Any]], write, devices:tuple[str, ...], queue:str, tag:int) -> list[UOp]:
+  deps:list[Dep] = []
   for lane, bufs in enumerate(bufs_by_lane):
     written = write if write is not None else list(range(len(bufs)))
-    dep_lanes += [(dep, dlane, lane) for dep, dlane in ctx.access_resources(bufs, written, (key, lane))]
-  return dep_lanes
+    deps += [Dep(d, q, t, lane) for d, q, t in ctx.tracker.access_resources(bufs, written, (devices[lane], queue, tag)) if t < tag]
 
-def _build_wait_ins(slots:dict[str, int], dep_lanes:list[tuple[tuple, int, int]], devices:tuple[str, ...], queue:str) -> tuple[list[UOp], set[int]]:
-  # opt1: same-queue ops are fifo-ordered
+  # same-queue submits are fifo-ordered, no wait needed
   if devices[0].split(":")[0] in {"AMD", "QCOM", "CPU"} or queue.startswith("COPY"):
-    dep_lanes = [(dep, dlane, lane) for dep, dlane, lane in dep_lanes if (dep[0][dlane], dep[1]) != (devices[lane], queue)]
-
-  # opt2: keep latest dep per (dep device, queue, cur lane)
-  latest = {((dep[0][dlane], dep[1]), lane): (dep, dlane) for dep, dlane, lane in sorted(dep_lanes, key=lambda x: x[0][2])}
-  deps:dict[tuple, dict[int, list[int]]] = collections.defaultdict(lambda: collections.defaultdict(list))
-  for (_, lane), (dep, dlane) in latest.items(): deps[dep][lane].append(dlane)
+    deps = [d for d in deps if (d.dev, d.queue) != (devices[d.lane], queue)]
+  latest = {(d.dev, d.queue, d.lane): d for d in sorted(deps, key=lambda d: d.tag)}
 
+  # keep only the latest signal
+  rows:dict[tuple[str, int], dict[int, list[str]]] = collections.defaultdict(lambda: collections.defaultdict(list))
+  for d in latest.values(): rows[(d.queue, d.tag)][d.lane].append(d.dev)
   waits = []
-  for (ddevs, dqueue, dtag), by_lane in deps.items():
-    for ls in itertools.zip_longest(*(by_lane[lane] for lane in range(len(devices)))):
-      s = UOp.mstack(*[make_buf(d, tag="sentinel_signal") if dl is None else make_buf(ddevs[dl], slots[dqueue]) for dl, d in zip(ls, devices)])
-      waits.append(UOp(Ops.INS, arg="wait", src=(s, UOp.const(dtag + 1, dtypes.uint64))))
-  return waits, {dtag for _, _, dtag in deps}
+  for (dqueue, dtag), by_lane in rows.items():
+    for ds in itertools.zip_longest(*(by_lane[lane] for lane in range(len(devices)))):
+      sig = UOp.mstack(*[make_buf(d, tag="sentinel_signal") if dd is None else make_buf(dd, ctx.slots[dqueue]) for dd, d in zip(ds, devices)])
+      waits.append(UOp(Ops.INS, arg="wait", src=(sig, UOp.const(dtag + 1, dtypes.uint64))))
+  ctx.signal_tags |= {t for _, t in rows}
+  return waits
 
-def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[tuple[tuple[str, ...], str]],
-                      tracker:HCQDepsTracker, slots:dict[str, int]) -> tuple[list[UOp], list[UOp], set[int]]:
-  # collect all buffers which belong to devices
-  dev_bufs:dict[str, dict[int, Any]] = collections.defaultdict(dict)
-  for call, devices in batch:
-    for b in itertools.chain.from_iterable(_get_call_bufs_by_lane(call, devices)):
-      for bd in to_tuple(b.device): dev_bufs[bd][id(b)] = b
-
-  n, fences, resets, fins, signal_tags = len(batch_info), [], [], [], set()
-  for _, devgroup in itertools.groupby(sorted(dev_bufs), key=lambda d: d.split(":")[0]):
-    devs = tuple(devgroup)
-
-    # to finalize the batch, sync all accesses from other devices to buffers that belong to this device
-    fin_deps = [dl for dl in _get_deps(tracker, [list(dev_bufs[d].values()) for d in devs], None, key=(devs, "COMPUTE:0", n)) if dl[0][2] < n]
-    waits, cur_signal_tags = _build_wait_ins(slots, fin_deps, devs, "COMPUTE:0")
-    signal_tags |= cur_signal_tags
-
-    # wait the syncs and signal the device epoch, then bump the timeline on the host
-    tl_signal, tl_value = make_buf(devs, tag="timeline_signal"), make_buf(devs, tag="timeline_value")
-    fin_submit = make_submit(*waits, UOp(Ops.INS, arg="store", src=(tl_signal, tl_value.index(0))), devs=devs, queue="COMPUTE:0")
-    epoch = (epoch_slot:=tl_value.after(fin_submit).index(0)).load()
-
-    # fence once per device group on this schedule's previous epoch
-    qs = dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)])
-    sched_epoch = make_buf(devs, next(UOp.unique_num), tag="epoch")
-
-    wait_device_epoch = (done:=tl_signal.after(loop:=UOp.loop(0)).index(0).load()).end(loop, done < sched_epoch.index(0).load())
-    fences.append(make_call("hcq_fence", UOp.sink(wait_device_epoch), HCQInfo(devs)))
-
-    # queues of other groups wait on these signals, reset them after every group reached its epoch
-    rst = functools.reduce(lambda a,q: a+(make_buf(devs, slots[q]).after(*a[-1:]).index(0).store(0),), qs, cast(tuple[UOp, ...], ()))
-    if rst: resets.append(make_call("hcq_reset", UOp.sink(*rst), HCQInfo(devs)))
-    fins.append(make_call("hcq_finalizer", UOp.sink(epoch_slot.store(epoch + 1), sched_epoch.after(fin_submit).index(0).store(epoch)), HCQInfo(devs)))
-  return fences + resets, fins, signal_tags
-
-def _merged_hcq_call(calls:list[UOp]) -> UOp: # TODO: simplify?
+def _merge_submits(calls:list[UOp]) -> UOp: # TODO: simplify?
   if len(calls) == 1: return calls[0]
   devs, queue = get_submit(calls[0]).src[0].arg
   body = make_submit(*[cmd for c in calls for cmd in get_submit(c).src[0].src], devs=devs, queue=queue).sink()
-  return make_call(f"submit {queue} ({len(calls)})", body,
-    replace(calls[0].arg.aux, estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates()).simplify()))
+  return make_call(f"submit {queue} ({len(calls)})", body, replace(calls[0].arg.aux,
+    estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates()).simplify()))
 
 def _merge_queues(submits:list[UOp]) -> list[UOp]:
-  new_src:list[UOp] = []
-  opened_qs:dict[tuple[tuple[str, ...], str], list[UOp]] = {} # (devs, queue) -> list of hcq calls, kept in submit order
+  merged:list[UOp] = []
+  opened:dict[tuple[tuple[str, ...], str], list[UOp]] = {} # (devs, queue) -> hcq calls in submit order
   limits:dict[tuple[tuple[str, ...], str], int] = collections.defaultdict(lambda: JIT_BATCH_SIZE.value)
-
   for call in submits:
-    devs, queue = get_submit(call).src[0].arg
-    if (old:=opened_qs.pop(key:=(devs, queue), None)) is not None:
-      if limits[key] and len(old) >= limits[key]: new_src, old, limits[key] = new_src + [_merged_hcq_call(old)], [], limits[key] * 2
-      new_rec = old + [call]
-    else:
-      # no such queue opened: close every open submit on this queue that shares a device, so submit order is kept
-      closing = [k for k in opened_qs if k[1] == queue and set(k[0]) & set(devs)]
-      new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in closing]
-      new_rec = [call]
-    opened_qs[(devs, queue)] = new_rec
-  return new_src + [_merged_hcq_call(c) for c in opened_qs.values()]
+    devs, queue = key = get_submit(call).src[0].arg
+    if (group:=opened.pop(key, None)) is None:
+      # first submit on this queue: close open groups on the same queue with shared devices, so submit order is kept
+      for k in [k for k in opened if k[1] == queue and set(k[0]) & set(devs)]: merged.append(_merge_submits(opened.pop(k)))
+      group = []
+    elif limits[key] and len(group) >= limits[key]: merged, group, limits[key] = merged + [_merge_submits(group)], [], limits[key] * 2
+    opened[key] = group + [call]
+  return merged + [_merge_submits(g) for g in opened.values()]
 
-def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> list[UOp]:
-  batch_info = [(devices, "COMPUTE:0" if call.src[0].op is Ops.PROGRAM else "COPY:0") for call, devices in batch]
+def _make_finalizers(ctx:BatchCtx) -> tuple[list[UOp], list[UOp], list[UOp]]:
+  # collect all buffers which belong to devices
+  dev_bufs:dict[str, dict[int, Any]] = collections.defaultdict(dict)
+  for call, devices in ctx.batch:
+    for b in itertools.chain.from_iterable(_get_call_bufs_by_lane(call, devices)):
+      for bd in to_tuple(b.device): dev_bufs[bd][id(b)] = b
 
-  # schedule deps
-  signal_tags:set[int] = set()
-  slots:dict[str, int] = collections.defaultdict(lambda: next(UOp.unique_num))
-  deps_tracker = HCQDepsTracker()
-  call_waits:list[list[UOp]] = []
-  for tag, ((call, _), (devices, queue)) in enumerate(zip(batch, batch_info)):
-    deps = _get_deps(deps_tracker, _get_call_bufs_by_lane(call, devices), get_call_outs_ins(call)[0], key=(devices, queue, tag))
-    cmds, cur_signal_tags = _build_wait_ins(slots, deps, devices, queue)
-    call_waits.append(cmds)
-    signal_tags |= cur_signal_tags
+  n, fences, resets, fins = len(ctx.batch_info), [], [], []
+  for _, devgroup in itertools.groupby(sorted(dev_bufs), key=lambda d: d.split(":")[0]):
+    sched_epoch = make_buf(devs:=tuple(devgroup), next(UOp.unique_num), tag="epoch")
+    tl_signal, tl_value = make_buf(devs, tag="timeline_signal"), make_buf(devs, tag="timeline_value")
 
-  # build fences and finalizers
-  fences, fins, finalizer_signal_tags = _build_finalizers(batch, batch_info, deps_tracker, slots)
-  signal_tags |= finalizer_signal_tags
+    # fence: spin until the device timeline reaches this schedule's previous epoch
+    done = tl_signal.after(loop:=UOp.loop(0)).index(0).load()
+    fences.append(make_call("hcq_fence", UOp.sink(done.end(loop, done < sched_epoch.index(0).load())), HCQInfo(devs)))
 
-  src, kerns = [], []
-  for tag, ((call, _), (devices, queue), q) in enumerate(zip(batch, batch_info, call_waits)):
+    # reset: queues of other groups wait on this group's signals, zero them only after every group reached its epoch
+    qs = dedup([qn for bdevs, qn in ctx.batch_info if set(bdevs) & set(devs)])
+    rst:tuple[UOp, ...] = ()
+    for q in qs: rst += (make_buf(devs, ctx.slots[q]).after(*rst[-1:]).index(0).store(0),)
+    if rst: resets.append(make_call("hcq_reset", UOp.sink(*rst), HCQInfo(devs)))
+
+    # finalizer: bump the host timeline and remember this schedule's epoch for the next fence
+    waits = _wait_ins(ctx, [list(dev_bufs[d].values()) for d in devs], None, devs, "COMPUTE:0", n)
+    fin_submit = make_submit(*waits, UOp(Ops.INS, arg="store", src=(tl_signal, tl_value.index(0))), devs=devs, queue="COMPUTE:0")
+    epoch = (epoch_slot:=tl_value.after(fin_submit).index(0)).load()
+    fins.append(make_call("hcq_finalizer", UOp.sink(epoch_slot.store(epoch + 1), sched_epoch.after(fin_submit).index(0).store(epoch)), HCQInfo(devs)))
+  return fences, resets, fins
+
+def _emit_submits(ctx:BatchCtx, call_waits:list[list[UOp]]) -> tuple[list[UOp], list[tuple]]:
+  # one submit per call: timeline sync on first queue use, timestamps, the call, and a signal if someone waits on it
+  src, kerns, seen_queues = [], [], set()
+  for tag, ((call, _), (devices, queue), q) in enumerate(zip(ctx.batch, ctx.batch_info, call_waits)):
     # first queue use, sync prior device work with the device timeline
-    if batch_info.index((devices, queue)) == tag:
+    if (devices, queue) not in seen_queues:
+      seen_queues.add((devices, queue))
       epoch = make_buf(devices, tag="timeline_value").index(0) - 1
       q = [UOp(Ops.INS, arg="barrier", src=()), UOp(Ops.INS, arg="wait", src=(make_buf(devices, tag="timeline_signal"), epoch))] + q
 
     # and make hcq call
     name, info = get_call_name(call, get_call_arg_uops(call)), HCQInfo(devices, estimate_uop(call))
-    ts_ids = [next(UOp.unique_num) for _ in range(2)] if profile else []
+    ts_ids = [next(UOp.unique_num) for _ in range(2)] if ctx.profile else []
     kerns.append((devices, name, info.estimates, tuple(ts_ids), make_call(name, call.src[0], info).key))
 
     ts_ins = [UOp(Ops.INS, arg="timestamp", src=(make_buf(devices, s),)) for s in ts_ids]
     q += ts_ins[:1] + [call.replace(arg=replace(call.arg, aux=info))] + ts_ins[1:]
 
     # signal the queue if someone waits for us
-    if tag in signal_tags: q += [UOp(Ops.INS, arg="store", src=(make_buf(devices, slots[queue]), UOp.const(tag + 1, dtypes.uint64)))]
+    if tag in ctx.signal_tags: q += [UOp(Ops.INS, arg="store", src=(make_buf(devices, ctx.slots[queue]), UOp.const(tag + 1, dtypes.uint64)))]
     src.append(make_call(f"submit {name}", make_submit(*q, devs=devices, queue=queue).sink(), info))
+  return src, kerns
+
+def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> list[UOp]:
+  ctx = BatchCtx(batch, [(devices, "COMPUTE:0" if call.src[0].op is Ops.PROGRAM else "COPY:0") for call, devices in batch], profile)
+
+  call_waits = [_wait_ins(ctx, _get_call_bufs_by_lane(call, devices), get_call_outs_ins(call)[0], devices, queue, tag)
+                for tag, ((call, _), (devices, queue)) in enumerate(zip(ctx.batch, ctx.batch_info))]
+  fences, resets, fins = _make_finalizers(ctx)
+  submits, kerns = _emit_submits(ctx, call_waits)
 
   # append batch timestamps to finalizers
   fins = [f.replace(arg=replace(f.arg, aux=replace(a:=f.arg.aux, kernels=tuple(x for x in kerns if set(x[0]) & set(a.device))))) for f in fins]
-  return fences + _merge_queues(src) + fins
+  return fences + resets + _merge_queues(submits) + fins
 
-def sched_hcq_batches(l:UOp, profile:bool) -> UOp:
+@rewrite_group(new_ctx=False)
+def sched_batches(l:UOp, profile:bool) -> UOp:
   srcs:list[UOp] = []
   batch:list[tuple[UOp, tuple[str, ...]]] = []
   for call in l.src:
@@ -293,9 +287,6 @@ def sched_hcq_batches(l:UOp, profile:bool) -> UOp:
     else: srcs, batch = srcs + _finalize_batch(batch, profile) + [call], []
   return l.replace(src=tuple(srcs + _finalize_batch(batch, profile)))
 
-pm_schedule_and_merge = PatternMatcher([(UPat(Ops.LINEAR, name="l"),
-  lambda ctx, l: sched_hcq_batches(l, ctx[1]).substitute(ctx[0], walk=True, enter_calls=True))])
-
 # *****************
 # 4.2. hcq lowering: ops to ir
 
@@ -325,18 +316,20 @@ def trim_link_patches(ctx:tuple[list[UOp], list[UOp]], a:UOp) -> UOp|None:
   return a.src[0].after(*kept, *[d for p in afters for d in p.src[1:]]) if links else None
 pm_trim_link_patches = PatternMatcher([(UPat(Ops.AFTER, src=(UPat((Ops.PARAM, Ops.MSTACK)),), allow_any_len=True, name="a"), trim_link_patches)])
 
-def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[UOp, dict[UOp, UOp], tuple[UOp, ...], dict[UOp, int]]:
+def _dnum(stride:int) -> UOp: return UOp.variable("_device_num", 0, stride - 1, dtypes.int, param=True) if stride > 1 else UOp.const(0, dtypes.int)
+
+def make_addr_table(call:UOp, gaddrs:list[UOp], name:str, stride:int=1) -> tuple[UOp, dict[UOp, UOp], tuple[UOp, ...], dict[UOp, int]]:
   bare = {g: g.replace(src=(g.src[0].without_after,)) for g in gaddrs}
 
-  order = sorted(dedup(bare.values()), key=lambda g: ((b:=unwrap_mstack(g.buf_uop)[0]).arg.slot, repr(b.tag)))
-  slots = {g:i for i,g in enumerate(order)}
-  table = UOp.placeholder((len(order),), dtypes.uint64, next(UOp.unique_num), device=call.arg.aux.device).rtag(name)
+  # slot-major layout: slot i of lane j lives at i*stride+j, every lane reads through the same table base
+  slots = {g:i*stride for i,g in enumerate(sorted(dedup(bare.values()), key=lambda g: g.key))}
+  table = UOp.placeholder((len(slots)*stride,), dtypes.uint64, next(UOp.unique_num), device=call.arg.aux.device).rtag(name)
 
-  reads = {g: table.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(slots[bare[g]], dtypes.int)).load() for g in gaddrs}
-  fills = (table.after(*make_patches(table, [(i*table.dtype.itemsize, addr) for addr, i in slots.items()])),) if slots else ()
+  reads = {g: table.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(_dnum(stride) + slots[bare[g]]).load() for g in gaddrs}
+  fills = (table.after(*make_patches(table, [(i*table.dtype.itemsize, addr) for addr, i in slots.items()])),) if slots and stride == 1 else ()
   return table, reads, fills, {g:slots[bare[g]] for g in gaddrs}
 
-def make_gather_loop(patches:list[UOp], table:UOp, slots:dict[UOp, int], lt_patches:list[UOp]) -> dict[UOp, UOp]:
+def make_gather_loop(patches:list[UOp], table:UOp, slots:dict[UOp, int], lt_patches:list[UOp], stride:int) -> dict[UOp, UOp]:
   (dst,), words = dedup(p.buf_uop for p in patches), [(unwrap_view(p.src[0].src[0])[1] + off.val*(val.dtype.itemsize//p.buf_uop.dtype.itemsize),
                                                        slots[val]) for p in patches for off,val in zip(p.src[0].src[1].src, p.src[1].src)]
 
@@ -344,13 +337,13 @@ def make_gather_loop(patches:list[UOp], table:UOp, slots:dict[UOp, int], lt_patc
   pairs = UOp.placeholder((2*len(words),), dtypes.uint32, next(UOp.unique_num), device=dst.device).rtag("systems")
   lt_patches.append(make_binary_patch(pairs, struct.pack(f'<{2*len(words)}I', *itertools.chain(*words))))
   r = UOp.range(len(words), next(UOp.unique_num), dtype=dtypes.int, src=(pairs, dst))
-  off, slot = ((pairs.index(2*r+i).load() % bound).cast(dtypes.int) for i, bound in ((0, dst.max_numel()-1), (1, table.max_numel())))
+  off, slot = ((pairs.index(2*r+i).load() % bound).cast(dtypes.int) for i, bound in ((0, dst.max_numel()-1), (1, table.max_numel()-(stride-1))))
   # SHRINK(offset, length): a const length keeps the end bound from becoming an expression the program spec rejects
   patch = UOp(Ops.SHRINK, src=(dst, off, off.const_like(table.dtype.itemsize//dst.dtype.itemsize))).bitcast(table.dtype).index(0) \
-    .store(table.index(slot).load()).end(r)
+    .store(table.index(slot + _dnum(stride)).load()).end(r)
   return {p: UOp(Ops.NOOP) for p in patches} | {patches[0]: patch}
 
-def is_input_addr(g:UOp) -> bool: return all(x.op is Ops.PARAM and x.tag is None for x in unwrap_mstack(g.buf_uop))
+def is_input_addr(g:UOp) -> bool: return any(x.op is Ops.PARAM and x.tag is None for x in unwrap_mstack(g.buf_uop))
 
 def split_patches(call:UOp) -> UOp|None:
   rt_patches:list[UOp] = []
@@ -358,20 +351,22 @@ def split_patches(call:UOp) -> UOp|None:
   body = graph_rewrite(call.src[0], pm_trim_link_patches, ctx=(rt_patches, lt_patches), name=f"trim link-time patches ({call.arg.name})")
 
   # split patches. addresses read in the body go through the tables too
-  inputs, internals = partition(dedup([g for p in rt_patches for g in get_getaddrs(p)] + get_getaddrs(body)), is_input_addr)
+  lanes = len(to_tuple(call.arg.aux.device))
+  inputs, internals = partition(dedup(get_getaddrs(UOp.sink(body, *rt_patches))), is_input_addr)
   runtimes, systems = partition(internals, lambda g: any(x.tag in {"program", "kernargs", "cmdbuf"} for x in unwrap_mstack(g.buf_uop)))
-  tables = [make_addr_table(call, gs, n) for gs,n in ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems"))]
+  tables = [make_addr_table(call, gs, n, lanes if n == "inputs" else 1) for gs,n in ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems"))]
   reads, fills = {k:v for _,r,_,_ in tables for k,v in r.items()}, [f for t in tables[1:] for f in t[2]] # inputs table is filled by exec
 
   ipatches = [p for p in rt_patches if p.tag == "inputs" and all(v in tables[0][3] for v in p.src[1].src)] # only getaddrs go to the table
-  gathers = make_gather_loop(ipatches, tables[0][0], tables[0][3], lt_patches) if ipatches else {}
-  body = body.substitute({p:p.substitute(gathers | reads) for p in rt_patches}).substitute(reads)
+  gathers = make_gather_loop(ipatches, tables[0][0], tables[0][3], lt_patches, lanes) if ipatches else {}
+  body = body.substitute({p:p.substitute(gathers | reads) for p in rt_patches} | reads, walk=True)
 
   lt_srcs = collections.defaultdict(list)
   for p in lt_patches: lt_srcs[p.buf_uop].append(p)
-  return call.replace(src=(body, *call.src[1:], *[b.after(*ps) for b,ps in lt_srcs.items()], *fills),
-    arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=((to_tuple(inputs[0].arg),
-      tuple(sorted(dedup(b.arg.slot for g in inputs for b in unwrap_mstack(g.buf_uop))))),) if inputs else call.arg.aux.input_idxs)))
+
+  bufs = [u for _, u in sorted(dedup([(i, g.src[0].without_after) for g, i in tables[0][3].items()]))]
+  aux = replace(call.arg.aux, input_addrs=tuple((d, _lane(u, j)) for u in bufs for j,d in enumerate(call.arg.aux.device))) if inputs else call.arg.aux
+  return call.replace(src=(body, *call.src[1:], *[b.after(*ps) for b,ps in lt_srcs.items()], *fills), arg=replace(call.arg, aux=aux))
 pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), split_patches)])
 
 # *****************
@@ -393,7 +388,7 @@ def replace_params(call:UOp) -> UOp|None:
   addrs = dedup([g.src[0].without_after for g in call.toposort() if g.op is Ops.GETADDR])
   refhold += [a for a in addrs if a not in held and all(b.op is not Ops.PARAM or b.tag is not None for b in unwrap_mstack(a))]
 
-  sub = {(b:=u.without_after): UOp.param(i, u.dtype, shape=b.shape, device=HCQ_RUNTIME_DEV.value, volatile=b.op is Ops.PARAM and b.arg.volatile)
+  sub = {(b:=u.without_after): UOp.param_from_shape(i, b.shape, u.dtype, HCQ_RUNTIME_DEV.value, volatile=b.op is Ops.PARAM and b.arg.volatile)
          for i,u in enumerate(c_args)} | {v: v.replace(arg=replace(v.arg, slot=-1)) for v in variables if v.op is Ops.PARAM} | _rank_ranges(tops)
   info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(c_args + refhold) if u.without_after.tag == "inputs"), None))
   prg_sink = body.src[0].substitute(sub).replace(arg=KernelInfo("hcq_submit"), tag=1)
@@ -440,13 +435,12 @@ def _lane_arg(a:UOp, lane:int, table:UOp) -> UOp: return table if a.tag == "inpu
 
 def merge_batch(batch:list[UOp]) -> UOp:
   tables = UOp.variable("hcq_inputs_ptr", 0, 2**64-1, dtypes.uint64, param=True)
-  lanes = [(c, j, sum(len(idxs) * 8 for _, idxs in c.arg.aux.input_idxs)) for c in batch for j in range(len(c.arg.aux.device))] # (call, lane, bytes)
-  offs = itertools.accumulate((table_bytes for _, _, table_bytes in lanes), initial=0) # every lane owns the next table of the region
+  offs = itertools.accumulate((8 * len(c.arg.aux.input_addrs) for c in batch), initial=0) # every call owns the next table of the region
   cmds = [c.src[0].src[0].call(*[_lane_arg(a.without_after, j, tables + off) for a in c.src[1:]], UOp.variable("_device_num", 0, 1 << 30).bind(j))
-          for (c, j, _), off in zip(lanes, offs)]
+          for c, off in zip(batch, offs) for j in range(len(c.arg.aux.device))]
 
   info = HCQInfo((HCQ_RUNTIME_DEV.value,), sum((c.arg.aux.estimates for c in batch), start=Estimates()).simplify(),
-                 input_idxs=tuple(x for c in batch for x in c.arg.aux.input_idxs), kernels=tuple(k for c in batch for k in c.arg.aux.kernels))
+                 input_addrs=tuple(x for c in batch for x in c.arg.aux.input_addrs), kernels=tuple(k for c in batch for k in c.arg.aux.kernels))
   body = UOp.custom_function("hcq", make_submit(*cmds, devs=HCQ_RUNTIME_DEV.value, queue="SUBMIT:0").sink())
   return body.call(*[s for c in batch for s in c.src[1:] if s.without_after.tag != "inputs"], name=f"hcq_submitter ({len(batch)})", aux=info)
 
@@ -481,8 +475,8 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool) -> UOp:
     linear = linear.substitute(back_map:={s.param_like(i): s for i,s in enumerate(input_uops)} if input_uops is not None else {}, walk=True)
     linear = graph_rewrite(linear, pm_insert_copy_staging+pm_flatten_linear, name="insert copy staging")
 
-    # schedule
-    linear = graph_rewrite(linear, pm_schedule_and_merge, ctx=({s:p for p,s in back_map.items()}, profile), walk=True, name="schedule and merge hcq")
+    # schedule on real buffers
+    linear = sched_batches(linear, profile).substitute({s:p for p,s in back_map.items()}, walk=True, enter_calls=True)
 
     # lower to hcq programs, then pack the programs of every batch into one C submitter (needs a C runtime device for the program addresses)
     linear = hcq_lower(linear, pm_encode_cmdbufs+pm_pack_placeholders)

tinygrad/runtime/support/mlx/mlxdev.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
-import struct, random, socket, ctypes, functools, itertools
+import struct, random, ctypes, functools, itertools
 from tinygrad.helpers import getenv, wait_cond, round_up, next_power2, ceildiv, DEBUG, hi32, lo32, to_be32, to_be64
 from tinygrad.runtime.support.memory import BumpAllocator
-from tinygrad.runtime.support.system import PCIDevice
+from tinygrad.runtime.support.system import PCIDevice, ipv4_to_gid
 from tinygrad.runtime.autogen import mlx5, pci
 
 MLX_DEBUG = getenv("MLX_DEBUG", 0)
@@ -11,8 +11,6 @@ MLX5_CMD_STRUCTS = {v: (getattr(mlx5, f"struct_mlx5_ifc_{n[12:].lower()}_in_bits
   getattr(mlx5, f"struct_mlx5_ifc_{n[12:].lower()}_out_bits", None)) for n, v in mlx5.__dict__.items() if n.startswith("MLX5_CMD_OP_")}
 MLX5_CMD_STRUCTS[mlx5.MLX5_CMD_OP_ACCESS_REG] = (mlx5.struct_mlx5_ifc_access_register_in_bits, mlx5.struct_mlx5_ifc_access_register_out_bits)
 
-def ipv4_to_gid(ip): return bytes(10) + b'\xff\xff' + socket.inet_aton(ip)
-
 def udp_sport(lqpn, rqpn):
   v = (lqpn * rqpn ^ ((lqpn * rqpn) >> 20) ^ ((lqpn * rqpn) >> 40)) & 0xFFFFF
   return ((v & 0x3FFF) ^ ((v & 0xFC000) >> 14)) | 0xC000

tinygrad/runtime/support/system.py

@@ -10,6 +10,8 @@ from tinygrad.runtime.support.usb import USB3, CustomASM24Controller, USBMMIOInt
 MAP_FIXED, MAP_FIXED_NOREPLACE = 0x10, 0x100000
 MAP_LOCKED, MAP_POPULATE, MAP_NORESERVE = 0 if OSX else 0x2000, getattr(mmap, "MAP_POPULATE", 0 if OSX else 0x008000), 0x400
 
+def ipv4_to_gid(ip:str) -> bytes: return bytes(10) + b'\xff\xff' + socket.inet_aton(ip)
+
 class _System:
   def write_sysfs(self, path:str, value:str, msg:str, expected:str|None=None):
     if FileIOInterface(path, os.O_RDONLY).read().splitlines()[0] != (expected or value):

tinygrad/schedule/__init__.py

@@ -94,7 +94,7 @@ pm_post_sched_cache = PatternMatcher([
   # only resolve buffer PARAMs (slot>=0); ALU/shape vars use slot=-1 and must not be swapped for call args
   (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg.slot] if x.arg.slot >= 0 else None),
   # create new BUFFERs
-  (UPat(Ops.BUFFER, src=(UPat(),), name="b"), lambda ctx,b:
+  (UPat(Ops.BUFFER, src=(), name="b"), lambda ctx,b:
    create_new_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
 ])
 

tinygrad/schedule/indexing.py

@@ -5,7 +5,7 @@ from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, rewrite_group, broadcast_axes
 from tinygrad.uop.ops import gate_kernel_sink
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
-from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC
+from tinygrad.helpers import argsort, all_same, cpu_profile, colored, Context, SPEC
 
 @dataclass
 class IndexingContext:
@@ -248,13 +248,12 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> UOp:
         local_rngs, valids = zip(*[(r.get_idx(), r.get_valid()) for r in valid_rngs])
         rngs_valids.append((local_rngs, valids))
 
-      # TODO: in RANGEIFY > 1 all_all_same isn't required
       all_all_same = all(all_same(local_rngs) for local_rngs,_ in rngs_valids)
       _out_rngs = []
       _realize_axis = []
       for i,(local_rngs,valids) in enumerate(rngs_valids):
         # we compare the ranges without their valids
-        if all_all_same or (PCONTIG and all_same(local_rngs)):
+        if all_all_same:
           # the new valid is the OR of all the children valids
           minimum_valid = UOp.const(False).usum(valids)
           _out_rngs.append(graph_rewrite(local_rngs[0].valid(minimum_valid), symbolic, name="minimum_valid"))
@@ -268,15 +267,11 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> UOp:
 
     # if this element is a reduce and there's ended ranges, we might have to end some other ranges
     if len(ending_ranges[x]) and x.op in GroupOp.Elementwise.union({Ops.REDUCE}):
-      _realize_axis = rctx.realize_map.get(x) or []
-      for i,r in enumerate(out_rngs):
-        if i in _realize_axis: continue
-        if not (PCONTIG > 1) or any(any(rr.arg > e.arg for e in ending_ranges[x]) for rr in r.ranges):
-          _realize_axis.append(i)
+      _realize_axis = list(range(len(out_rngs)))
       ending_ranges[x] = []
       if len(_realize_axis):
         rctx.realize_map[x] = _realize_axis
-        out_rngs = tuple([(rctx.new_range(x.shape[i]) if i in _realize_axis else r) for i,r in enumerate(out_rngs)])
+        out_rngs = tuple(rctx.new_range(x.shape[i]) for i in range(len(out_rngs)))
     ending_ranges[x] += broadcast_ending_ranges
 
     # TODO: some ops don't have shape, enable this after the `.st` property is removed

tinygrad/schedule/multi.py

@@ -1,6 +1,7 @@
+from typing import cast
 from tinygrad.helpers import all_same, prod, getenv, ALLREDUCE_CAST
 from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, AxisType, graph_rewrite, broadcast_axes, _broadcast_shape, sint_to_uop
-from tinygrad.uop.ops import sint, ssimplify
+from tinygrad.uop.ops import sint, ssimplify, ParamArg, to_max_shape
 from tinygrad.dtype import dtypes
 from tinygrad.schedule.allreduce import handle_allreduce
 
@@ -15,7 +16,8 @@ def mstack_early_shrink(ms:UOp, shrink:UOp):
   ret:list[UOp] = []
   for i, x in enumerate(ms.src):
     if x.op is Ops.COPY:
-      ret.append(_apply_shrink(shrink.marg, x.src[0], i).copy_to_device(x.device))
+      src = _apply_shrink(shrink.marg, x.src[0], i)
+      ret.append(src.contiguous() if src.device == x.device else src.copy_to_device(x.device))
     else:
       ret.append(_apply_shrink(shrink.marg, x, i).contiguous())
   return ms.replace(src=tuple(ret))
@@ -35,9 +37,11 @@ replace_allreduce = PatternMatcher([
   (UPat(Ops.MSELECT, src=(UPat(Ops.MSTACK, name="mstack"),), name="ms"), lambda mstack, ms: mstack.src[ms.arg]),
   # move shrink before MSTACK
   (UPat(Ops.SHRINK, src=(UPat(Ops.MSTACK, name="ms"),), allow_any_len=True, name="shrink"), mstack_early_shrink),
-  # move MSELECT before movement ops
+  # move MSELECT before movement/ALU ops
   (UPat(Ops.MSELECT, src=(UPat(GroupOp.Movement, src=(UPat.var("s"),), allow_any_len=True, name="v"),), name="ms"),
    lambda s,v,ms: v.replace(src=(s.mselect(ms.arg),)+v.src[1:])),
+  (UPat(Ops.MSELECT, src=(UPat(GroupOp.ALU, name="a"),), name="ms"), lambda a,ms:
+   a.replace(src=tuple(s.mselect(ms.arg) if isinstance(s.device, tuple) else s for s in a.src))),
 ])
 
 _early_allreduce = PatternMatcher([
@@ -278,12 +282,33 @@ def rewrite_into_function(call:UOp):
   return call.replace(src=(new_body,)+new_args)
 
 def param_to_multi(p:UOp):
+  # a bare param with an axis is flat storage: the per-shard param is the storage divided by the device count.
+  # params with a view (RESHAPE/SHRINK on top) are handled by param_view_to_multi instead
+  if p.axis is None or p.axis != 0: return None
+  count = len(cast(tuple, p.device))
+  param = UOp(Ops.PARAM, arg=ParamArg(p.arg.slot, p.dtype, cast(int, p.arg.size)//count, p.arg.vmin_vmax, p.arg.multiple_of, p.arg.name,
+                                      p.arg.addrspace, device=p.device))
+  return param.unshard(p.axis)
+
+def param_view_to_multi(v:UOp, p:UOp):
+  # the movement view on the param carries the unpacked (multiplied on axis) shape: shard the view, not the flat storage
   if p.axis is None: return None
-  return UOp.param(p.arg.slot, p.dtype, p.shard_shape, p.device, p.arg.vmin_vmax, p.arg.multiple_of, p.arg.name, p.arg.addrspace).unshard(p.axis)
+  count = len(cast(tuple, p.device))
+  shp = tuple(s//count if i == p.axis else s for i,s in enumerate(v.shape))
+  max_shp, max_shp_of_v = to_max_shape(shp), to_max_shape(v.shape)
+  param = UOp(Ops.PARAM, arg=ParamArg(p.arg.slot, p.dtype, prod(max_shp), p.arg.vmin_vmax, p.arg.multiple_of, p.arg.name,
+                                      p.arg.addrspace, device=p.device))
+  # the unshard has to sit on the per-shard view so its axis is in range
+  view = param.reshape(max_shp) if len(max_shp) > 1 else param
+  if max_shp != shp: view = view.shrink_to(shp)
+  ret = view.unshard(p.axis)
+  if len(max_shp_of_v) > 1: ret = ret.reshape(max_shp_of_v)
+  return ret.shrink_to(v.shape) if max_shp_of_v != tuple(v.shape) else ret
 
 # NOTE: this is the same pattern as unrolled ranges
 multi_pm = PatternMatcher([
   (UPat(Ops.PARAM, name="p"), param_to_multi),
+  (UPat(GroupOp.Movement, src=(UPat(Ops.PARAM, name="p"),), name="v"), param_view_to_multi),
   (UPat(GroupOp.ALU, name="root", custom_early_reject=set([Ops.UNSHARD])), alu_multi),
   (UPat(Ops.REDUCE, src=(UPat(Ops.UNSHARD, name="multi"), ), name="root"), reduce_multi),
   (UPat(Ops.RESHAPE, src=(UPat(Ops.UNSHARD, name="multi"), UPat()), name="root"), reshape_multi),

tinygrad/schedule/prepare.py

@@ -1,13 +1,18 @@
 import itertools
 from tinygrad.dtype import dtypes, to_dtype
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp
-from tinygrad.uop.ops import graph_rewrite, rewrite_group, shape_to_shape_arg, ParamArg, identity_element
+from tinygrad.uop.ops import graph_rewrite, rewrite_group, ParamArg, identity_element
 from tinygrad.uop.movement import mop_cleanup
 from tinygrad.helpers import prod, getenv, all_int, DEBUG, SPLIT_REDUCEOP, OPENPILOT_HACKS, FLOAT16, argsort
 from tinygrad.schedule.indexing import apply_movement_op
 from tinygrad.schedule.allreduce import create_allreduce_function
 from tinygrad.schedule.multi import multi_pm
 
+def walk_mop(u:UOp):
+  if u.op in GroupOp.Movement or u.op in {Ops.INDEX, Ops.UNSHARD, Ops.BITCAST}: return walk_mop(u.src[0])
+  if u.op is Ops.AFTER and (b:=walk_mop(u.src[0])) is not u.src[0]: return b.after(*u.src[1:])
+  return u
+
 def found_after(ctx:dict[UOp, UOp], after:UOp, src:UOp):
   if (x:=src).op is Ops.CAST and x.dtype == dtypes.half and FLOAT16: x, after = x.src[0], after.cast(dtypes.float)
   while True:
@@ -95,11 +100,20 @@ def resolve_function(c:UOp, allow_param_mismatch=True) -> UOp|None:
     if [x.arg.slot for x in params] != list(range(len(params))): raise RuntimeError(f"params not in order: {[x.arg.slot for x in params]}")
     if len(params) != len(args): raise TypeError(f"expected {len(params)} args, got {len(args)}")
 
-  dict_map = {x:args[x.arg.slot] for x in params}
-  for i, (p, a) in enumerate(dict_map.items()):
-    if p.axis != a.axis: raise TypeError(f"arg {i} axis mismatch: expected {p.axis}, got {a.axis}")
-    if p.max_shape != a.max_shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}")
+  # params have a flat storage size in the arg, the logical shape is a view (RESHAPE/SHRINK/UNSHARD) on top of it.
+  # substitute args as views of their flat max-shaped storage so the movement views on the params stay valid
+  def storage(x:UOp) -> tuple[tuple[int, ...], int]:
+    shp = x.max_shard_shape if x.axis is not None and isinstance(x.device, tuple) else x.max_shape
+    return shp, prod(shp)
+  pairs = [(p, args[p.arg.slot]) for p in params]
+  for i, (p, a) in enumerate(pairs):
+    if p.arg.size is not None:
+      if p.arg.size != storage(a)[1]: raise TypeError(f"arg {i} shape mismatch: expected size {p.arg.size}, got {a.shape}")
     if p.dtype != a.dtype: raise TypeError(f"arg {i} dtype mismatch: expected {p.dtype}, got {a.dtype}")
+  def as_storage(a:UOp) -> UOp:
+    mshp, n = storage(a)
+    return a if a.shape == (n,) else a.pad_to(mshp).reshape((n,))
+  dict_map = {p: a if p.arg.size is None else as_storage(a) for p, a in pairs}
   return c.src[0].substitute(dict_map, walk=True)
 
 # shape-changing bitcast
@@ -135,10 +149,6 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
 
   # ** copy rules **
 
-  # COPY transfers a contiguous range, so materialize a source that's resized (shrink/pad/expand) or reordered (permute/flip)
-  (UPat(Ops.COPY, src=(UPat(GroupOp.Movement, name="r"),), name="c"),
-   lambda c,r: c.replace(src=(r.contiguous(),)) if resolve(r.numel() != r.base.numel(), False) or r.contiguous_view_offset() is None else None),
-
   # copy to same device is a no-op
   (UPat(Ops.COPY, src=(UPat.var("x"),), name="copy"), lambda x,copy: x if x.device == copy.device else None),
 
@@ -174,11 +184,16 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
    lambda reduce,x: reduce.const_like(identity_element(reduce.arg[0], reduce.dtype)) if 0 in x.shape and 0 not in reduce.shape else None),
   # handle size 0
   (UPat(GroupOp.All-{Ops.SINK}, name="x"), lambda x: x.const_like(0).rtag(x.tag) if x._shape is not None and 0 in x.shape else None),
+
+  # remove movement ops from SINK/AFTER. TODO: should be generic
+  (UPat(Ops.SINK, name="s"), lambda s: s.replace(src=tuple(walk_mop(u) for u in s.src if u.op is not Ops.NOOP))),
+  (UPat(Ops.AFTER, name="s"), lambda s: s.replace(src=(s.src[0],)+tuple(walk_mop(u) for u in s.src[1:] if u.op is not Ops.NOOP))),
 ])
 
 def convert_copy_to_store(ctx, copy:UOp, existing_buf:UOp|None=None):
   input_src = copy.src[0]
-  if not input_src.has_buffer_identity(after_ok=True): input_src = input_src.contiguous()
+  # if it's a COPY, we need to give the input buffer identity
+  if not input_src.has_buffer_identity(after_ok=True) and copy.op is Ops.COPY: input_src = input_src.contiguous()
   input_src = input_src.flatten()
   if existing_buf is not None:
     # if the existing buffer is not a full buffer, we can't use it
@@ -186,9 +201,9 @@ def convert_copy_to_store(ctx, copy:UOp, existing_buf:UOp|None=None):
     # if there's already a buffer, we just use it
     return existing_buf.flatten().store(input_src)
   # create the output buffer
-  buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg(input_src.max_shape),), arg=ParamArg(next(ctx), copy.dtype, device=copy.device))
+  buf = UOp(Ops.BUFFER, arg=ParamArg(next(ctx), copy.dtype, size=prod(input_src.max_shape), device=copy.device))
   # reshape back to input
-  return buf.after(buf.store(input_src)).reshape(copy.shape)
+  return buf.reshape(input_src.max_shape).after(buf.store(input_src)).reshape(copy.shape)
 
 pm_copy_to_store = PatternMatcher([
   (UPat(name="existing_buf").store(UPat(Ops.COPY, name="copy")), convert_copy_to_store),

tinygrad/schedule/rangeify.py

@@ -1,12 +1,12 @@
 from dataclasses import dataclass, field, replace
 from typing import cast
 import itertools
-from tinygrad.dtype import dtypes, AddrSpace, Invalid, strong_dtype
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg
+from tinygrad.dtype import dtypes, AddrSpace, Invalid
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg
 from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group
 from tinygrad.uop.symbolic import symbolic
 from tinygrad.helpers import prod, dedup, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS, SPEC
-from tinygrad.helpers import PCONTIG, partition, get_single_element
+from tinygrad.helpers import get_single_element
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
 from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, apply_movement_op
@@ -83,7 +83,7 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
   accessed_buffers = dedup(accessed_buffers)
 
   # if this is generated from multiple buffers, don't remove this buffer
-  if len(accessed_buffers) > 3 and not (PCONTIG > 2): return None
+  if len(accessed_buffers) > 3: return None
 
   # if any reduces access a buffer, don't remove this buffer
   buffer_in_reduce = False
@@ -94,22 +94,7 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
   UOp.sink(*[x.src[0] for x in reduces]).toposort(gate=buf_gate)
   del buf_gate
   if buffer_in_reduce:
-    if PCONTIG > 2:
-      out_in_ratio = (prod(buf.shape)+1) / (sum([x.numel() for x in accessed_buffers])+1)
-      if out_in_ratio < 10: return None
-      # here we have to check the indexes, we might do a partial contig here
-      local_indexes = [x for x in indexes if x.src[0].op is Ops.STAGE and x.src[0].arg.addrspace == AddrSpace.LOCAL]
-      exclude_ranges = UOp.group(*[UOp.group(*x.src[1:]) for x in local_indexes]).ranges
-      subs = [(k,v) for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST]
-      # if it's bufferized or a reduce, it's pcontig
-      is_pcontig, is_subs = partition(subs, lambda x: x[0] in exclude_ranges or any([r.arg[-1] == AxisType.REDUCE for r in x[1].ranges]))
-      if not len(is_subs):
-        return None
-      if len(is_pcontig):
-        ret = src.substitute(dict(is_subs), extra_pm=pm_gate_substitute)
-        return ret.bufferize(*[x[0] for x in is_pcontig], arg=BufferizeOpts(None, AddrSpace.LOCAL)).index(*[x[1] for x in is_pcontig])
-    else:
-      return None
+    return None
 
   # if it makes it here, the bufferize is removed
   # this is the ranges replaced
@@ -132,8 +117,6 @@ pm_const_buffer_folding = pm_mops+PatternMatcher([
   (UPat(Ops.STAGE, name="b"), cleanup_dead_axes),
   # remove noop buffers. if we look at the next index we can remove even more of these
   (UPat(Ops.INDEX, name="idx").f(Ops.STAGE, allow_any_len=True, name="b2"), remove_noop_bufferize),
-  (UPat(Ops.INDEX, src=(UPat(Ops.STAGE),), allow_any_len=True, name="idx").f(Ops.NOOP).f(Ops.STAGE, allow_any_len=True, name="b2"),
-   remove_noop_bufferize),
   # no buffers for a const, in either spelling
   (UPat.cvar('c').or_casted().f(Ops.STAGE, allow_any_len=True, name="b"), lambda c,b: b.const_like(c.val)),
   # indexing a const is the const
@@ -141,8 +124,6 @@ pm_const_buffer_folding = pm_mops+PatternMatcher([
   # indexing an after with all fully invalid stores is invalid
   (UPat(Ops.INDEX, src=(UPat(Ops.AFTER, name="after"),), allow_any_len=True, name="idx"),
    lambda idx,after: idx.const_like(Invalid) if after_all_invalid(after) else None),
-  # hack if a noop turned to a const
-  (UPat(Ops.NOOP, src=(UPat.cvar().or_casted("c"),)), lambda c: c),
   # a deviceless MSTACK src is the same value on every device, so indexing the stack is just indexing that value
   (UPat(Ops.MSTACK, src=(UPat.var("s"),), allow_any_len=True).f(Ops.INDEX, allow_any_len=True, name="idx"),
    lambda s,idx: idx.replace(src=(s,)+idx.src[1:]) if s.device is None else None),
@@ -221,7 +202,7 @@ pm_limit_bufs = PatternMatcher([(UPat(set.union(GroupOp.Binary, GroupOp.Ternary)
 
 def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True):
   size = prod(x.shape)
-  dtype = strong_dtype(x.dtype)  # a BUFFER is never weak: store at the concrete dtype, the .cast(x.dtype) on the result keeps readers unchanged
+  if x.dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {x.dtype}")
   rngs = sorted(idx.ranges, key=lambda x: x.arg)
   assert size > 0 and isinstance(size, int), f"no zero sized or symbolic sized buffers {size}"
 
@@ -242,15 +223,15 @@ def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True):
 
   # NOTE: the local BUFFER needs to be disambiguated here
   if x.arg.addrspace == AddrSpace.GLOBAL:
-    buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg((size,)),), arg=ParamArg(next(ctx), dtype, device=x.arg.device, addrspace=AddrSpace.GLOBAL))
-    do_store = buf.index(idx).store(x.src[0].cast(dtype)).end(*rngs)
-    return buf.after(do_store).cast(x.dtype)
+    buf = UOp(Ops.BUFFER, arg=ParamArg(next(ctx), x.dtype, size=size, device=x.arg.device, addrspace=AddrSpace.GLOBAL))
+    do_store = buf.index(idx).store(x.src[0]).end(*rngs)
+    return buf.after(do_store)
 
   if allow_locals:
     # handle locals
-    buf = UOp.placeholder((size,), dtype, next(ctx), AddrSpace.LOCAL)
-    do_store = buf.index(idx).store(x.src[0].cast(dtype)).end(*rngs)
-    return buf.after(do_store).cast(x.dtype)
+    buf = UOp.placeholder((size,), x.dtype, next(ctx), AddrSpace.LOCAL)
+    do_store = buf.index(idx).store(x.src[0]).end(*rngs)
+    return buf.after(do_store)
 
 # collapse any BUFFERIZE to single input BUFFERIZE
 def flatten_bufferize(x:UOp):
@@ -275,11 +256,6 @@ def remove_noop_afters(x:UOp) -> UOp|None:
 pm_add_buffers = pm_mops+pm_flatten_bufferize+PatternMatcher([
   (UPat(Ops.STAGE, src=(UPat(), UPat(name="idx")), name="x"), lambda ctx,x,idx: bufferize_to_store(ctx, x, idx, allow_locals=False)),
 
-  # INDEX of a buffer through the weak cast added above: index the buffer directly and cast the loaded value instead.
-  # this must run in the same rewrite that adds the cast, or the expander expands the whole casted buffer into one big VECTORIZE
-  (UPat(Ops.INDEX, src=(UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat.var("buf"),)),), allow_any_len=True, name="u"),
-   lambda u,buf: u.replace(dtype=None, src=(buf,)+u.src[1:]).cast(u.dtype)),
-
   # move RESHAPEs through MSELECT/MSTACK
   (UPat((Ops.MSELECT, Ops.MSTACK), src=UPat(Ops.RESHAPE), name="m"),
    lambda m: m.replace(src=tuple([x.src[0].base for x in m.src])).reshape(m.shape)),
@@ -306,8 +282,7 @@ class LocalAddBufferContext:
 def debuf(ctx:LocalAddBufferContext, buf:UOp):
   # Variables (ALU buffers with a value range) are scalar symbolic values, not real buffers: they become ALU params with no slot
   if buf.is_variable: return buf.replace(op=Ops.PARAM)
-  param = UOp(Ops.PARAM, src=(UOp.const(prod(buf.max_shape)),),
-              arg=ParamArg(ctx.dg, buf.dtype, addrspace=buf.addrspace, device=buf.device))
+  param = UOp(Ops.PARAM, arg=ParamArg(ctx.dg, buf.dtype, prod(buf.max_shape), addrspace=buf.addrspace, device=buf.device))
   ret = param.reshape(buf.max_shape)
   # if the buffer has symbolic shape, shrink the max-sized view to the actual shape
   if buf.max_shape != buf.shape: ret = ret.shrink(tuple((0, s) for s in buf.shape))
@@ -364,10 +339,6 @@ def get_contiguous(ctx:LocalAddBufferContext, x:UOp):
 
 rangeify_codegen = PatternMatcher([
   (UPat(Ops.CONTIGUOUS, name="x"), get_contiguous),
-
-  # no NOOP in the kernel graph
-  # TODO: this can be moved into codegen?
-  (UPat(Ops.NOOP, name="x"), lambda x: x.src[0] if len(x.src) else None),
 ])
 
 pm_add_param_range_tags = PatternMatcher([

tinygrad/tensor.py

@@ -198,9 +198,7 @@ def finalize_after(ctx:AllocCtx, x:UOp):
 
 def replace_input_buffer(ctx:AllocCtx, b:UOp):
   ctx.replacements.append(b)
-  if b.is_bound_var or b.is_variable: return b.param_like(len(ctx.replacements)-1)
-  return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b.device,
-                   addrspace=b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL)
+  return b.param_like(len(ctx.replacements)-1)
 
 def replace_input_view(ctx:AllocCtx, b:UOp): return replace_input_buffer(ctx, b) if b in ctx.views else None
 
@@ -211,7 +209,7 @@ pm_finalize_call = PatternMatcher([
 
 pm_replace_buf = PatternMatcher([
   # replace BUFFER with PARAM for cache key normalization
-  (UPat(Ops.BUFFER, src=(UPat(),), name="b"), lambda ctx,b:
+  (UPat(Ops.BUFFER, src=(), name="b"), lambda ctx,b:
    replace_input_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
   # replace SHRINK with PARAM
   (UPat(Ops.SHRINK, src=(UPat(Ops.BUFFER),), name="b", allow_any_len=True), replace_input_view),
@@ -381,7 +379,7 @@ class Tensor(RandMixin):
   # ***** data handlers ****
 
   def as_param(self, slot:int):
-    return Tensor(UOp.param(slot, self.dtype, self.uop.shard_shape, self.device, axis=self.uop.axis))
+    return Tensor(self.uop.param_like(slot))
 
   def call(self, *lst:Tensor, fxn:Tensor|UOp, grad_fxn:Callable|None=None) -> Tensor:
     fret = fxn._uop.call(*[t.uop for t in (self,)+lst], grad_fxn=grad_fxn)

tinygrad/uop/ops.py

@@ -23,6 +23,8 @@ class AxisType(Enum):
 class ParamArg:
   slot: int
   dtype: DType
+  # number of elements in the buffer. always a concrete int (never symbolic), None for scalars (shape ())
+  size: int|None = None
   vmin_vmax: tuple[PyConst, PyConst]|None = None
   multiple_of: int|None = None
   name: str|None = None
@@ -30,10 +32,13 @@ class ParamArg:
   axis: int|None = None
   device: str|tuple[str, ...]|None = None
   volatile: bool = False
+  # (h, w) if this is an image2d buffer, then size == h*w*4
+  image: tuple[int, int]|None = None
   def __repr__(self):
     fields = (("vmin_vmax", None), ("multiple_of", None), ("name", None), ("addrspace", AddrSpace.GLOBAL), ("axis", None), ("device", None),
-              ("volatile", False))
-    args = [repr(self.slot), repr(self.dtype)] + [f"{k}={v!r}" for k,default in fields if (v:=getattr(self, k)) != default]
+              ("volatile", False), ("image", None))
+    args = [repr(self.slot), repr(self.dtype)] + ([repr(self.size)] if self.size is not None else []) + \
+      [f"{k}={v!r}" for k,default in fields if (v:=getattr(self, k)) != default]
     return f"ParamArg({', '.join(args)})"
 axis_letters = {AxisType.DEVICE: "d", AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.WEAK: "L",
                 AxisType.LOOP: "L", AxisType.UPCAST: "u", AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
@@ -123,7 +128,8 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None:
       # a CALL of an opaque body is void, a CALL of an address can return a value
       return dtypes.void if src[0].dtype is dtypes.void else None
     case Ops.CUSTOM | Ops.CUSTOMI:
-      return None
+      assert isinstance(arg, tuple) and len(arg) == 2 and isinstance(arg[1], DType), f"CUSTOM/CUSTOMI arg must be (str, DType), got {arg}"
+      return arg[1]
     case Ops.INS:
       return None
     case Ops.NOOP:
@@ -324,7 +330,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   @functools.cached_property
   def tuplize(self:UOp) -> tuple:
-    return (self.op.value, self.arg, self.dtype,)+tuple([x.tuplize for x in self.src])
+    # arg goes through repr: args of different types (None, str, tuple) must stay mutually comparable for the sort
+    return (self.op.value, repr(self.arg), self.dtype,)+tuple([x.tuplize for x in self.src])
 
   # *** uop shape stuff ***
 
@@ -379,9 +386,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
       case Ops.GETADDR: return ()
       case Ops.RANGE | Ops.SPECIAL: return ()
       case Ops.BINARY: return (len(self.arg),)
-      case Ops.BUFFER:
-        if len(self.src): return self.src[0].as_shape
-        return ()
+      case Ops.BUFFER | Ops.PARAM:
+        # PARAM/BUFFER don't have a shape input, they have a size in the arg: int gives shape (size,), None gives ()
+        if (img:=self.arg.image) is not None: return (img[0], img[1], 4)
+        return () if self.arg.size is None else (self.arg.size,)
       case Ops.CUSTOM | Ops.CUSTOMI:
         if self.dtype is dtypes.void: return None
         input_shapes = [x._shape for x in self.src if x._shape is not None]
@@ -392,10 +400,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
         # STAGE adds the existing shape to the front, opposite of INDEX
         return tuple([int(r.vmax+1) for r in self.src[1:]])+self.src[0].shape
 
-      # param has shape as the only arg
-      case Ops.PARAM:
-        return self.src[0].as_shape
-
       # wmma output shape = accumulator shape (src[2])
       case Ops.WMMA:
         wmma_b = _broadcast_shape(self.src[0].shape[:-1], self.src[1].shape[:-1], self.src[2].shape[:-1])
@@ -819,8 +823,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   @staticmethod
   def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None):
     if dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {dtype}")
+    assert isinstance(size, int), f"new_buffer size must be a concrete int, got {size}"
     slot = next(UOp.unique_num) if num is None else num
-    return UOp(Ops.BUFFER, src=(shape_to_shape_arg((size,)),), arg=ParamArg(slot, dtype, device=device))
+    return UOp(Ops.BUFFER, arg=ParamArg(slot, dtype, size=size, device=device))
   @staticmethod
   def from_buffer(opaque:Buffer, device:str|tuple[str, ...]|None=None):
     if (uop:=UOp.new_buffer(device or opaque.device, opaque.size, opaque.dtype, num=-id(opaque))) not in buffers: buffers[uop] = opaque.ref(1)
@@ -976,7 +981,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     # a Variable is a 0-d BUFFER in the ALU addrspace; binding it is storing a CONST into it
     # param=True creates the kernel-side form directly: an ALU PARAM (what the BUFFER becomes inside kernels)
     arg = ParamArg(-1, dtype, name=name, vmin_vmax=(min_val, max_val), multiple_of=multiple_of, addrspace=AddrSpace.ALU)
-    return UOp(Ops.PARAM if param else Ops.BUFFER, src=(shape_to_shape_arg(()),), arg=arg)
+    return UOp(Ops.PARAM if param else Ops.BUFFER, arg=arg)
   @property
   def is_variable(self) -> bool:
     # a Variable is a 0-d BUFFER in the ALU addrspace that carries a value range (it becomes a PARAM inside kernels)
@@ -1146,11 +1151,11 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     dtype = strong_dtype(dtype)  # storage is never weak: a placeholder commits the width of what's put in it
     if slot is None: slot = next(UOp.unique_num)
     if addrspace is AddrSpace.GLOBAL:
-      ret = UOp(Ops.PARAM, src=(shape_to_shape_arg((prod(shape),)),), arg=ParamArg(slot, dtype, addrspace=addrspace, device=device,volatile=volatile))
+      ret = UOp(Ops.PARAM, arg=ParamArg(slot, dtype, size=prod(shape), addrspace=addrspace, device=device, volatile=volatile))
     else:
       assert addrspace in (AddrSpace.LOCAL, AddrSpace.REG)
       assert device is None, "LOCAL and REG placeholders cannot have a device"
-      ret = UOp(Ops.BUFFER, src=(shape_to_shape_arg((prod(shape),)),), arg=ParamArg(slot, dtype, addrspace=addrspace))
+      ret = UOp(Ops.BUFFER, arg=ParamArg(slot, dtype, size=prod(shape), addrspace=addrspace))
     if tag is not None: ret = ret.rtag(tag)
     if len(shape) > 1: ret = ret.reshape(shape)
     return ret
@@ -1164,20 +1169,51 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   # TODO: this should replace placeholder
   @staticmethod
-  def param(slot:int, dtype:DType, shape:tuple[sint, ...]|None=None, device=None, vmin_vmax:tuple[PyConst, PyConst]|None=None,
+  def param(slot:int, dtype:DType, size:int|None=None, device=None, vmin_vmax:tuple[PyConst, PyConst]|None=None,
             multiple_of:int|None=None, name=None, addrspace=AddrSpace.GLOBAL, axis:int|None=None, volatile:bool=False):
+    # a PARAM only has a size: a concrete number of elements. None means scalar (shape ()). no symbolic sizes here,
+    # symbolic shapes are expressed as a max-size param shrunk to the real shape, see param_from_shape
     if dtype in dtypes.weaks: raise RuntimeError(f"cannot create param for weak dtype {dtype}")
-    if shape is not None and axis is not None and isinstance(device, tuple):
-      shape = tuple(s*len(device) if i == axis else s for i,s in enumerate(shape))
-    src: tuple[UOp, ...] = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),)
-    return UOp(Ops.PARAM, src=src, arg=ParamArg(slot, dtype, vmin_vmax, multiple_of, name, addrspace, axis, device, volatile))
+    if axis is not None and isinstance(device, tuple):
+      assert size is not None, "param with axis needs a size"
+      size *= len(device)
+    return UOp(Ops.PARAM, arg=ParamArg(slot, dtype, size, vmin_vmax, multiple_of, name, addrspace, axis, device, volatile))
+  @staticmethod
+  def param_from_shape(slot:int, shape:tuple[sint, ...], dtype:DType, device=None, vmin_vmax:tuple[PyConst, PyConst]|None=None,
+                       multiple_of:int|None=None, name=None, addrspace=AddrSpace.GLOBAL, axis:int|None=None, volatile:bool=False):
+    """create a PARAM for a (possibly multi-dim or symbolic) shape: the arg gets the concrete max size,
+    the real shape is laid on top with RESHAPE/SHRINK"""
+    if dtype in dtypes.weaks: raise RuntimeError(f"cannot create param for weak dtype {dtype}")
+    if axis is not None and isinstance(device, tuple): shape = tuple(s*len(device) if i == axis else s for i,s in enumerate(shape))
+    max_shape = to_max_shape(shape)
+    if len(shape) == 3 and shape[2] == 4:  # image2d buffers keep (h, w) in the arg and are always max size
+      return UOp(Ops.PARAM, arg=ParamArg(slot, dtype, prod(max_shape), vmin_vmax, multiple_of, name,
+                                         addrspace, axis, device, volatile, (max_shape[0], max_shape[1])))
+    ret = UOp(Ops.PARAM, arg=ParamArg(slot, dtype, prod(max_shape) if len(shape) else None, vmin_vmax, multiple_of, name,
+                                      addrspace, axis, device, volatile))
+    if len(shape) > 1: ret = ret.reshape(max_shape)
+    if max_shape != shape: ret = ret.shrink_to(shape)
+    return ret
   def param_like(self, slot:int):
     # Variables become ALU params in the call body; the stored value (if bound) stays in the call args
     if self.is_bound_var or self.is_variable:
       b = self.src[0] if self.op is Ops.AFTER else self
-      return UOp(Ops.PARAM, src=b.src, arg=replace(b.arg, slot=slot, name=f"p{slot}"))
+      return UOp(Ops.PARAM, arg=replace(b.arg, slot=slot, name=f"p{slot}"))
     addrspace = self.addrspace if self.addrspace is not None else AddrSpace.GLOBAL
-    return UOp.param(slot, self.dtype, self.shard_shape if self.axis is not None else self._shape, self.device, addrspace=addrspace, axis=self.axis)
+    # multi-device values become a per-shard sized param wrapped in UNSHARD: the sharding lives in the graph, not the arg
+    if self.axis is not None and isinstance(self.device, tuple):
+      return UOp.shared_view(UOp(Ops.PARAM, arg=ParamArg(slot, self.dtype, prod(to_max_shape(self.shard_shape)),
+                                                       addrspace=addrspace, device=self.device)),
+                             self.shard_shape, self.axis, count=len(self.device))
+    return UOp.param_from_shape(slot, self._shape, self.dtype, self.device, addrspace=addrspace)
+  @staticmethod
+  def shared_view(flat:UOp, shard_shape:tuple[sint, ...], axis:int, count:int) -> UOp:
+    """view flat storage (per-shard size) as a value sharded on axis with full shape shard_shape*count on axis"""
+    max_shp = to_max_shape(shard_shape)
+    ret = flat.reshape(max_shp) if len(shard_shape) > 1 else flat
+    if max_shp != shard_shape: ret = ret.shrink_to(shard_shape)
+    # the unshard output shape is already the full multiplied shape
+    return ret.unshard(axis)
 
   @staticmethod
   def custom_function(name:str, *src:UOp) -> UOp: return UOp(Ops.CUSTOM_FUNCTION, src=src, arg=name)

tinygrad/uop/render.py

@@ -57,8 +57,6 @@ renderer = PatternMatcher([
 ])
 
 renderer_infer = PatternMatcher([
-  (UPat(Ops.CMOD, name="x"), lambda ctx,x: f"cmod({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
-  (UPat(Ops.CDIV, name="x"), lambda ctx,x: f"cdiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
   (UPat(Ops.FLOORMOD, name="x"), lambda ctx,x: f"floormod({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
   (UPat(Ops.FLOORDIV, name="x"), lambda ctx,x: f"floordiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
   (UPat(Ops.CAST, name="x"),
@@ -86,7 +84,7 @@ pm_pyrender_extra = PatternMatcher([
   (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.val})"),
   (UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})" if x.dtype != x.src[0].dtype else None),
   (UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].val}, {repr(x.arg)})"),
-  (UPat(Ops.BUFFER, src=(UPat(),), name="x"), lambda x:
+  (UPat(Ops.BUFFER, src=(), name="x"), lambda x:
     f"UOp.new_buffer({repr(x.arg.device)}, {x.max_numel()}, {x.dtype}, {x.arg.slot})"
     if isinstance(x.arg, ParamArg) and x.addrspace is AddrSpace.GLOBAL else None),
   (UPat(Ops.COPY, src=(UPat(name="x"),), name="copy"), lambda ctx,x,copy: f"{ctx[x]}.copy_to_device({repr(copy.arg)})"),

tinygrad/uop/spec.py

@@ -87,10 +87,9 @@ spec_shared = PatternMatcher([
   # a loop-ended END requires a trailing bool condition for the backedge (loop again while true)
   (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE, dtypes.void), UPat(dtype=dtypes.bool))), lambda: True),
 
-  # PARAM
-  (UPat(Ops.PARAM, name="x"), lambda x: isinstance(x.arg, ParamArg)),
-  (UPat(Ops.BUFFER, src=(UPat(),), name="x"), lambda x:
-   isinstance(x.arg, ParamArg) and x.addrspace in (AddrSpace.REG, AddrSpace.LOCAL)),
+  # PARAM/BUFFER have a size in the arg, no shape input
+  (UPat(Ops.PARAM, src=(), name="x"), lambda x: isinstance(x.arg, ParamArg)),
+  (UPat(Ops.BUFFER, src=(), name="x"), lambda x: isinstance(x.arg, ParamArg) and x.addrspace in (AddrSpace.REG, AddrSpace.LOCAL)),
 
   # GROUP of stores (or groups, or NOOPs)
   (UPat(Ops.GROUP, dtypes.void, src=UPat((Ops.GROUP, Ops.STORE, Ops.NOOP, Ops.INS, Ops.END))), lambda: True),
@@ -100,8 +99,9 @@ spec_shared = PatternMatcher([
                                                      Ops.AFTER, Ops.UNSHARD, Ops.BITCAST, Ops.INS})),),
         allow_any_len=True, name="x"), lambda x: matches_dtype(x.src[0], x.dtype)),
 
-  # CUSTOM (inline and non inline)
-  (UPat((Ops.CUSTOMI, Ops.CUSTOM)), lambda: True),
+  # CUSTOM (inline and non inline): the arg is the source string and the dtype it produces, void for a bare statement
+  (UPat((Ops.CUSTOMI, Ops.CUSTOM), name="x"),
+   lambda x: isinstance(x.arg, tuple) and len(x.arg) == 2 and isinstance(x.arg[0], str) and isinstance(x.arg[1], DType)),
 
   # CALL of an external function
   (UPat(Ops.CALL, src=(UPat(),), allow_any_len=True, name="x"),
@@ -140,12 +140,12 @@ spec_tensor = PatternMatcher([
    lambda u: dtypes.is_float(u.dtype) or u.src[0].base.is_invalid),
 
   # BUFFER
-  (UPat(Ops.BUFFER, src=(UPat(),), name="buf"), lambda buf:
-   (isinstance(buf.dtype, DType) and matches_dtype(buf.src[0], dtypes.weakint) and is_device(buf.arg.device))
+  (UPat(Ops.BUFFER, src=(), name="buf"), lambda buf:
+   (isinstance(buf.dtype, DType) and isinstance(buf.arg.size, int) and is_device(buf.arg.device))
    if isinstance(buf.arg, ParamArg) and buf.addrspace is AddrSpace.GLOBAL else None),
 
   # a Variable is a 0-d ALU BUFFER with a value range and no device
-  (UPat(Ops.BUFFER, src=(UPat(),), name="buf"), lambda buf: buf.arg.device is None if buf.is_variable else None),
+  (UPat(Ops.BUFFER, src=(), name="buf"), lambda buf: buf.arg.device is None if buf.is_variable else None),
 
   # custom function
   (UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda x: isinstance(x.arg, str)),

tinygrad/uop/symbolic.py

@@ -189,6 +189,7 @@ symbolic_simple = pm_data_invalid + PatternMatcher([
   # a conditional with the same results either way is a noop, also fold const conditionals
   (UPat.var().where(UPat.var("val"), UPat.var("val")), lambda val: val),
   (UPat.cvar("gate").where(UPat.var("c0"), UPat.var("c1")).named("w"), fold_const_where),
+  (UPat.var("gate").where(UPat.var("x"), 0) != 0, lambda gate,x: gate & (x != 0)),
   # a.where(b.where(c, d), d) -> (a & b).where(c, d)
   (UPat.var("a").where(UPat.var("b").where(UPat.var("c"), UPat.var("d")), UPat.var("d")), lambda a,b,c,d: (a&b).where(c,d)),
   # a.where(c, b.where(c, d)) -> (a | b).where(c, d)
@@ -312,6 +313,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
                         else y.src for y in x.src[1:]]))))),
   # after/end with 1 src is just src[0]
   (UPat((Ops.AFTER, Ops.END), src=(UPat.var("s"),)), lambda s: s),
+  # ranges can be subbed for CONSTs, remove them from ENDs while preserving a constant bool backedge
+  (UPat(Ops.END, name="x"), lambda x: x.replace(src=(x.src[0],)+tuple(r for r in x.src[1:] if r.op is not Ops.CONST or r.dtype is dtypes.bool))),
   # the rules above key on bare CONSTs, so a redundantly committed const has to be uncast in the same fixpoint
 ])+div_and_mod_symbolic+pm_uncast_const
 
@@ -446,9 +449,6 @@ pm_clean_up_group_sink = PatternMatcher([
 ])
 
 sym = symbolic+pm_simplify_valid+PatternMatcher([
-  # reorder ALU/VECTORIZE
-  (UPat(GroupOp.ALU, src=(UPat(Ops.STACK, src=UPat(name='x')), UPat(Ops.STACK, src=UPat(name='y'))), name='alu'),
-   lambda x,y,alu: UOp(Ops.STACK, src=(UOp(alu.op, src=(x,y)),))),
   # ** where **
   # push cast to branches
   (UPat.var("s").where(UPat.var("a"), UPat.var("b")).cast().named("cast"),

tinygrad/uop/upat.py

@@ -2,6 +2,7 @@ from typing import Any, Callable
 import itertools, inspect, functools, types
 from tinygrad.helpers import partition, dedup, Context
 from tinygrad.uop.ops import UPat, UOp, Ops, PatternMatcher, graph_rewrite, deconstruct_function
+from tinygrad.dtype import dtypes
 
 class UPatCompileError(Exception): pass
 
@@ -18,40 +19,42 @@ def _get_clause(self:UPat, base:UOp, depth=0) -> UOp:
   # build the and_clause for acceptance
   and_clause:list[UOp] = []
   if self.op is not None:
-    if len(self.op) > 1: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=tuple(int(x) for x in self.op))), arg="{0}.op in {1}"))
-    else: and_clause.append(UOp(Ops.CUSTOM, src=(base,), arg="{0}.op == "+str(self.op[0].value)))
+    if len(self.op) > 1:
+      and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=tuple(int(x) for x in self.op))), arg=("{0}.op in {1}", dtypes.void)))
+    else: and_clause.append(UOp(Ops.CUSTOM, src=(base,), arg=("{0}.op == "+str(self.op[0].value), dtypes.void)))
   if self.arg is not None:
-    if isinstance(self.arg, int): and_clause.append(UOp(Ops.CUSTOM, src=(base,), arg="{0}.arg == "+str(int(self.arg))))
-    else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=self.arg)), arg="{0}.arg == {1}"))
+    if isinstance(self.arg, int): and_clause.append(UOp(Ops.CUSTOM, src=(base,), arg=("{0}.arg == "+str(int(self.arg)), dtypes.void)))
+    else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=self.arg)), arg=("{0}.arg == {1}", dtypes.void)))
   if self.strict_length or self.required_len > 0:
-    and_clause.append(UOp(Ops.CUSTOM, src=(base,), arg=("len({0}.src)"+(" == " if self.strict_length else " >= ")+str(self.required_len))))
-  if self.name is not None: and_clause.append(UOp(Ops.STORE, src=(UOp(Ops.CUSTOMI, arg=self.name), base)))
+    and_clause.append(UOp(Ops.CUSTOM, src=(base,),
+      arg=("len({0}.src)"+(" == " if self.strict_length else " >= ")+str(self.required_len), dtypes.void)))
+  if self.name is not None: and_clause.append(UOp(Ops.STORE, src=(UOp(Ops.CUSTOMI, arg=(self.name, dtypes.void)), base)))
   if self.match_dtype is not None:
     if len(self.match_dtype) > 1:
       and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=tuple(self.match_dtype))),
-                            arg="{0}.dtype in {1}"))
+                            arg=("{0}.dtype in {1}", dtypes.void)))
     else:
       and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=self.match_dtype[0])),
-                            arg="{0}.dtype == {1}"))
+                            arg=("{0}.dtype == {1}", dtypes.void)))
   if self.match_tag is not None:
     if len(self.match_tag) > 1:
-      and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=tuple(self.match_tag))), arg="{0}.tag in {1}"))
-    else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=self.match_tag[0])), arg="{0}.tag == {1}"))
+      and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=tuple(self.match_tag))), arg=("{0}.tag in {1}", dtypes.void)))
+    else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=self.match_tag[0])), arg=("{0}.tag == {1}", dtypes.void)))
   if self.src is not None:
     # single match
     if len(self.src) == 1 and isinstance(self.src[0], tuple):
       and_clause += [_get_clause(s, base.index(i), depth) for i,s in enumerate(self.src[0])]
     # repeat match
     elif len(self.src) == 1 and isinstance(self.src[0], itertools.repeat):
-      it = UOp(Ops.CUSTOMI, arg=f"ituop{depth}")
+      it = UOp(Ops.CUSTOMI, arg=(f"ituop{depth}", dtypes.void))
       match = _get_clause(next(self.src[0]), it, depth+1)
-      and_clause.append(UOp(Ops.CUSTOM, src=(match, it, base), arg="all([{0} for {1} in {2}.src])"))
+      and_clause.append(UOp(Ops.CUSTOM, src=(match, it, base), arg=("all([{0} for {1} in {2}.src])", dtypes.void)))
     # multi match (fork)
     elif len(self.src) > 1 and all(isinstance(x, tuple) for x in self.src):
       fork_cond = [UOp(Ops.AND, src=tuple([_get_clause(s, base.index(i), depth) for i,s in enumerate(ss)])) for ss in self.src]
       and_clause.append(UOp(Ops.OR, src=tuple(fork_cond)))
     else: raise RuntimeError("broken")
-  return UOp(Ops.AND, src=tuple(and_clause)) if and_clause else UOp(Ops.CUSTOMI, arg="True")
+  return UOp(Ops.AND, src=tuple(and_clause)) if and_clause else UOp(Ops.CUSTOMI, arg=("True", dtypes.void))
 
 # *** pattern matcher ***
 
@@ -91,7 +94,7 @@ def do_process_and(a:UOp) -> UOp|None:
       for store in stores:
         if store.src[0] in dict_stores:
           # duplicate store is an identity compare
-          new_src.append(UOp(Ops.CUSTOM, src=(dict_stores[store.src[0]], store.src[1]), arg="{0} is {1}"))
+          new_src.append(UOp(Ops.CUSTOM, src=(dict_stores[store.src[0]], store.src[1]), arg=("{0} is {1}", dtypes.void)))
           found = True
         else:
           dict_stores[store.src[0]] = store.src[1]
@@ -108,17 +111,18 @@ pm_proc = PatternMatcher([(UPat(Ops.AND, name="a"), do_process_and)], compiled=F
 # renderer
 def wrap(ctx, x) -> UOp:
   ctx[ret:=f"a{len(ctx)}"] = x.arg
-  return UOp(Ops.CUSTOMI, arg=ret)
+  return UOp(Ops.CUSTOMI, arg=(ret, dtypes.void))
 
 pm_renderer = PatternMatcher([
   (UPat(Ops.PYLITERAL, name="x"), wrap),
 
   # AND of CUSTOMI fragments inside a CUSTOM becomes a single CUSTOMI (joined with " and ")
   (UPat(Ops.CUSTOM, src=(UPat(Ops.AND, src=UPat(Ops.CUSTOMI), name="x"), UPat(), UPat()), name="r"),
-    lambda r,x: r.replace(src=(UOp(Ops.CUSTOMI, arg="(" + ' and '.join(y.arg for y in x.src) + ")"),)+r.src[1:])),
+    lambda r,x: r.replace(src=(UOp(Ops.CUSTOMI, arg=("(" + ' and '.join(y.arg[0] for y in x.src) + ")", dtypes.void)),)+r.src[1:])),
 
-  (UPat(Ops.CUSTOM, src=UPat(Ops.CUSTOMI), name="x"), lambda x: UOp(Ops.CUSTOMI, arg=x.arg.format(*[y.arg for y in x.src]))),
-  (UPat(Ops.INDEX, src=(UPat(Ops.CUSTOMI, name="x"), UPat(Ops.CONST, name="c")), name="g"), lambda x,c,g: x.replace(arg=x.arg+f".src[{c.val}]"))
+  (UPat(Ops.CUSTOM, src=UPat(Ops.CUSTOMI), name="x"), lambda x: UOp(Ops.CUSTOMI, arg=(x.arg[0].format(*[y.arg[0] for y in x.src]), dtypes.void))),
+  (UPat(Ops.INDEX, src=(UPat(Ops.CUSTOMI, name="x"), UPat(Ops.CONST, name="c")), name="g"),
+   lambda x,c,g: x.replace(arg=(x.arg[0]+f".src[{c.val}]", dtypes.void)))
 ], compiled=False)
 
 def _final_render(x:UOp, has_ctx:bool, depth=1) -> list[str]:
@@ -131,8 +135,8 @@ def _final_render(x:UOp, has_ctx:bool, depth=1) -> list[str]:
       for ss in s.src: or_pieces.extend(_final_render(ss, has_ctx, depth+1))
     elif s.op is Ops.STORE:
       assert s.src[0].op is Ops.CUSTOMI and s.src[1].op is Ops.CUSTOMI
-      store_pieces.append(f"{s.src[0].arg}={s.src[1].arg}")
-    elif s.op is Ops.CUSTOMI: and_pieces.append(s.arg)
+      store_pieces.append(f"{s.src[0].arg[0]}={s.src[1].arg[0]}")
+    elif s.op is Ops.CUSTOMI: and_pieces.append(s.arg[0])
     else: raise UPatCompileError(f"can't compile this {s}")
   # if we have an or, render it
   if len(or_pieces):
@@ -145,7 +149,7 @@ def _final_render(x:UOp, has_ctx:bool, depth=1) -> list[str]:
   return [f"{'  '*depth}if {and_clause}: return _ret"]
 
 def _get_code(self:UPat, has_ctx:bool):
-  ret = _get_clause(self, UOp(Ops.CUSTOMI, arg="uop"))
+  ret = _get_clause(self, UOp(Ops.CUSTOMI, arg=("uop", dtypes.void)))
   try:
     # TODO: this should be tracked in a "system" rewrite, not untracked or tracked with kernel
     with Context(TRACK_MATCH_STATS=0):