Merge branch 'master' into rewrite_rangeify2

.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 "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 "comgr.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, comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
+        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.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, bnxt"
+        python3 -c "from tinygrad.runtime.autogen import mlx5"
         python3 -c "from tinygrad.runtime.autogen import ggml_common"
         REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
     - name: Check for differences
@@ -102,3 +102,42 @@ 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

@@ -446,55 +446,14 @@ jobs:
     - name: UsbGPU (USB4/TB) tiny tests
       run: DEV=PCI+NV:NAK python3.11 test/test_tiny.py
 
-  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 }})
+  testcommalatest:
+    name: comma Benchmark (0.11.2)
     runs-on: [self-hosted, Linux, comma]
-    timeout-minutes: 5
+    timeout-minutes: 12
     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
@@ -505,10 +464,45 @@ 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: 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: 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: 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/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
+        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
         language: system
         always_run: true
         pass_filenames: false

examples/mlperf/model_train.py

@@ -1674,7 +1674,8 @@ def train_gptoss():
   config = {}
   BASEDIR            = config["BASEDIR"]                = Path(getenv("BASEDIR", "/raid/datasets/c4-8b/"))
   BS                 = config["BS"]                     = getenv("BS", 16)
-  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS
+  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
+  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
   SEED               = config["SEED"]                   = getenv("SEED", 5760)
   DATA_SEED          = config["DATA_SEED"]              = getenv("DATA_SEED", SEED)
   SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
@@ -1736,8 +1737,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=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),
+    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),
   )
 
   for p in optim.params:
@@ -1769,7 +1770,7 @@ def train_gptoss():
 
   @TinyJit
   @Context(TRAINING=1)
-  def step(tokens:Tensor):
+  def minibatch(tokens:Tensor):
     if is_dp: tokens = tokens.to(None).shard(device, 0)
     if not is_sharding: tokens = tokens.to(None)
 
@@ -1783,20 +1784,22 @@ def train_gptoss():
     for g, new_g in zip(grads, loss.gradient(*optim.params)):
       apply_grad(g, new_g.uop)
 
-    Tensor.realize(loss, *grads)
+    loss_cpu = loss.flatten().float().to("CPU")
+    return loss_cpu.realize(*grads)
 
-    grad_norm = clip_grads(grads, 1, 1.0)
+  @TinyJit
+  def optim_step():
+    grad_norm = clip_grads(grads, grad_acc, 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(loss_cpu, lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales)
+    Tensor.realize(lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales)
 
-    return loss_cpu, lr_cpu, grad_norm_cpu
+    return lr_cpu, grad_norm_cpu
 
   @TinyJit
   @Context(TRAINING=0)
@@ -1845,20 +1848,30 @@ def train_gptoss():
       profile_marker(f"train @ {i}")
       st = time.perf_counter()
 
-      ist = 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
 
-      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()
+      gt = time.perf_counter()
+      ret = optim_step()
+      lr, grad_norm = ret[0].item(), ret[1].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
@@ -1868,7 +1881,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, {dev_time:.3f} s dev, {data_time:.3f} s data, {loss:.4f} loss, " \
+          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"{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())))
 
@@ -1878,6 +1891,8 @@ 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=0
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
 
 export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
 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=0
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
 
 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=0
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
 
 export LOGMLPERF=1
 

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, BENCHMARK_LOG.value))
+    _events[self.event]["wall"].append(self.time)
     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, BENCHMARK_LOG.value))
+    _events[self.event]["kernel"].append(GlobalCounters.time_sum_s - self.start)
     return False
 
 def log_event_instant(event:InstantBenchEvent, value:float):
-  _events[event].append((value, BENCHMARK_LOG.value))
+  _events[event].append(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, log_name, run):
-    point = Point(log_name.replace(':', '_').replace('.', '_')).tag("id", run_id).tag("index", i)
+  def _create_point(run_id, i, attempt, ref, commit, name, value, run):
+    point = Point(BENCHMARK_LOG.value).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, log_name) in enumerate(values):
-            point = _create_point(run_id, i, attempt, ref, commit, f"{event.value}_{event_type}", value, log_name, run)
+          for i, value in enumerate(values):
+            point = _create_point(run_id, i, attempt, ref, commit, f"{event.value}_{event_type}", value, run)
             points.append(point)
       else:
-        for i, (value, log_name) in enumerate(_events[event]):
-          point = _create_point(run_id, i, attempt, ref, commit, event.value, value, log_name, run)
+        for i, value in enumerate(_events[event]):
+          point = _create_point(run_id, i, attempt, ref, commit, event.value, value, 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

@@ -1,238 +0,0 @@
-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

@@ -1,118 +0,0 @@
-#!/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

@@ -1,46 +0,0 @@
-#!/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)

test/backend/test_ops.py

@@ -6,7 +6,6 @@ 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:
@@ -809,8 +808,6 @@ 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)
@@ -1813,7 +1810,6 @@ class TestOps(unittest.TestCase):
     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, Device, dtypes, Variable
-from tinygrad.helpers import Context, GlobalCounters, getenv, DEBUG
+from tinygrad import Tensor, nn, Device, dtypes, Variable
+from tinygrad.helpers import Context, GlobalCounters, getenv, PCONTIG, 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(DEBUG=max(2, DEBUG.value)):
+    with Context(PCONTIG=2, DEBUG=max(2, DEBUG.value)):
       out = (self.a @ self.b @ self.c).contiguous(arg=opts).realize()
 
     with Context(DEBUG=0):
@@ -88,15 +88,16 @@ class TestRangeifyEdgeCase(unittest.TestCase):
     res = Tensor.cat(a, c, dim=0)
     self.assertEqual(res.numpy()[-1, :16].tolist(), [512] * 16)
 
-  def test_multi_gather(self):
+  def test_pcontig_multi_gather(self):
     # regression test: local bufferize must have device set for const_like to work
-    # 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()
+    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()
     self.assertEqual(result.tolist(), [1, 1])
 
 if getenv("BIG") > 2:
@@ -117,6 +118,65 @@ 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_schedule.py

@@ -365,6 +365,16 @@ 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))
@@ -374,6 +384,14 @@ 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/mockgpu/amd/amddriver.py

@@ -1,15 +1,20 @@
-import ctypes, mmap, collections, functools, copy, os
+import pathlib, re, 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 _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)}
+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()
 
 class KFDFileDesc(VirtFileDesc):
   def __init__(self, fd, driver):
@@ -111,44 +116,42 @@ class AMDDriver(VirtDriver):
 
   def kfd_ioctl(self, req, argp):
     nr = req & 0xFF
-    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)
+    struct = kfd_headers[nr].from_address(argp)
 
-    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):
+    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:
       struct.major_version = 1
       struct.minor_version = 14
-    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_ALLOC_MEMORY_OF_GPU):
+    elif nr == kfd_ioctls.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 == _ioctl_nr(kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU):
+    elif nr == kfd_ioctls.AMDKFD_IOC_FREE_MEMORY_OF_GPU:
       self.object_by_handle.pop(struct.handle)
-    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU):
+    elif nr == kfd_ioctls.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 == _ioctl_nr(kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU):
+    elif nr == kfd_ioctls.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 == _ioctl_nr(kfd.AMDKFD_IOC_CREATE_EVENT):
+    elif nr == kfd_ioctls.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 == _ioctl_nr(kfd.AMDKFD_IOC_CREATE_QUEUE):
+    elif nr == kfd_ioctls.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)
@@ -159,7 +162,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 == _ioctl_nr(kfd.AMDKFD_IOC_WAIT_EVENTS):
+    elif nr == kfd_ioctls.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:
@@ -167,7 +170,11 @@ class AMDDriver(VirtDriver):
           ev.memory_exception_data.va = int(os.environ["MOCKGPU_EMU_FAULTADDR"], 16)
           ev.memory_exception_data.failure.NotPresent = 1
     else:
-      raise RuntimeError(f"unsupported kfd ioctl, {nr} {name}")
+      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)
     return 0
 
   def _emulate_execute(self):

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], 0)
+      self.assertGreater(_events[event]["wall"][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], 0)
-      self.assertGreater(_events[event]["wall"][1][0], 0)
+      self.assertGreater(_events[event]["wall"][0], 0)
+      self.assertGreater(_events[event]["wall"][1], 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][0], wall_times[0])
-      self.assertGreater(_events[event]["kernel"][0][0], 0)
+      self.assertLess(_events[event]["kernel"][0], wall_times[0])
+      self.assertGreater(_events[event]["kernel"][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][0], wall_times[0])
-      self.assertGreater(_events[event]["kernel"][0][0], 0)
+      self.assertLess(_events[event]["kernel"][0], wall_times[0])
+      self.assertGreater(_events[event]["kernel"][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][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)
+      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)
 
   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][0], 1000)
+      self.assertEqual(_events[event][0], 1000)
 
 if __name__ == '__main__':
   unittest.main()

test/unit/test_allreduce.py

@@ -33,8 +33,8 @@ class TestRingAllReduce(unittest.TestCase):
       # 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)
+      # N*N shrinks becoming contigs, N ALU, N extra contig, reassembly (cat), and mul
+      sink_count = (N*N)+(N)+(N)+(1)+(1)
       if len(sinks) != sink_count: raise KernelCountException(sink_count, len(sinks))
       # correctness
       run_linear(linear, var_vals)

test/unit/test_assign.py

@@ -1,10 +1,10 @@
 #!/usr/bin/env python
 import unittest
 import numpy as np
-from tinygrad import Device, dtypes, Tensor, TinyJit, GlobalCounters, Variable
+from tinygrad import 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, needs_second_gpu
+from test.helpers import assert_kernel_count
 
 N = 200  # has to be bigger than the cache to fail
 
@@ -1079,80 +1079,5 @@ 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/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, 53), (23, 136, 34)):
+    for ggml_type,packed_size,words in ((13, 176, 44), (14, 210, 210), (23, 136, 34)):
       packed = Tensor.empty(packed_size+4, dtype=dtypes.uint8, device="CPU")[4:]
       decoded = ggml_data_to_tensor(packed, 256, ggml_type).reshape(1, 256)
       linear = Linear(256, 1, bias=False)

test/unit/test_bnxt.py

@@ -1,115 +0,0 @@
-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

@@ -1,40 +0,0 @@
-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_llm_amd.py

@@ -20,13 +20,11 @@ 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, gsum = q8_quantize(Tensor(x), 2, 32)
+    quant, scale = q8_quantize(Tensor(x), 2, 32)
     scale_np = np.maximum(np.max(np.abs(x), axis=-1, keepdims=True) / 127, 1e-8)
     expected = np.clip(np.rint(x / scale_np), -127, 127).astype(np.int8)
     np.testing.assert_array_equal(quant.bitcast(dtypes.int8).reshape(2, 32).numpy(), expected)
     np.testing.assert_allclose(scale.numpy(), scale_np, rtol=1e-6)
-    # 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")
@@ -38,9 +36,7 @@ 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()))
-    # 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)
+    self.assertEqual(linear.weight.uop.buf_uop.buffer.offset, 4)
 
   def test_q4_k_linear(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")

test/unit/test_multitensor.py

@@ -852,6 +852,82 @@ 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/late/coalesce.py

@@ -103,7 +103,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
   if getenv("DMC"): return sink
 
   # collect
-  memory: defaultdict[tuple[Ops, UOp, UOp|str, UOp, object], dict[int, list[UOp]]] = defaultdict(dict)
+  memory: defaultdict[tuple[Ops, UOp, UOp|str, UOp], 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}:
@@ -118,12 +118,11 @@ 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
-      # 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)
+      memory[(u.op, buf, root_src, valid)].setdefault(arg, []).append(u)
 
   # build replacements
   replacements = {}
-  for (op,buf,base,valid,ld_arg),offsets in memory.items():
+  for (op,buf,base,valid),offsets in memory.items():
     # allowed lengths (copied in)
     lengths = []
     must_divide = True
@@ -158,7 +157,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(arg=ld_arg)
+          ld = idx.load()
           for i,g in enumerate(grp):
             for oo in offsets[g]:
               replacements[oo] = ld.index(i) if len(grp) > 1 else ld

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_written_bufs
+from tinygrad.engine.realize import unwrap_multi, resolve_params, get_call_arg_uops, get_call_outs_ins
 from tinygrad.schedule.memory import memory_plan_rewrite, _collect_bufs
 from tinygrad.nn.state import get_parameters
 from tinygrad.uop.movement import mop_cleanup
@@ -173,7 +173,13 @@ class CapturedJit(Generic[ReturnType]):
 
   @functools.cached_property
   def _written_uops(self) -> set[UOp]:
-    return {b for call in self.linear.toposort() if call.op is Ops.CALL for b in get_call_written_bufs(call)}
+    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
 
   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, dedup
+from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansipad, all_int, prod, flatten, Context, getenv, to_tuple, tqdm
 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,10 +26,6 @@ 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]
@@ -217,12 +213,14 @@ 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 = [cast(Buffer, _resolve(u, ctx.input_uops).buffer).get_buf(d).va_addr for d, u in info.input_addrs]
+  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]
   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:
-    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))})
+    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)})
   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/function.py

@@ -46,7 +46,7 @@ class _function(Generic[ReturnType]):
     params = get_state_dict((args, kwargs), tensor_type=(Tensor, UOp)).values()
 
     # deduplicate input_uops, keeping the first occurrence index for each unique uop
-    call_uops: list[UOp] = dedup([u for t in params if (u:=t._uop).device is not None])
+    call_uops: list[UOp] = dedup([u for t in params if (u:=t._uop).device is not None or u.is_bound_var])
 
     # disable realize/schedule while this is running
     # run it and do surgery later
@@ -64,7 +64,10 @@ class _function(Generic[ReturnType]):
       raise RuntimeError(f"function return type {type(ret)} not supported")
 
     # replace the known inputs with params (using deduplicated slots)
-    subs = {x:x.param_like(i) for i,x in enumerate(call_uops)}
+    def make_param(x:UOp, i:int) -> UOp:
+      p = x.param_like(i)
+      return p.replace(arg=replace(p.arg, name=f"p{i}")) if x.is_bound_var else p
+    subs = {x:make_param(x, i) for i,x in enumerate(call_uops)}
     uret = uret.substitute(subs)
 
     # the BUFFERs that are left are the implicit inputs

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 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)
+  # 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)
   try:
     with open("/sys/fs/cgroup/cpu.max") as f:
       quota, period = f.read().strip().split()
-      if quota != "max": count = min(count, max(1, int(quota) // int(period)))
+      if quota != "max": return max(1, int(quota) // int(period))
   except (FileNotFoundError, ValueError, ZeroDivisionError): pass
-  return count
+  # 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))
 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,6 +271,7 @@ 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/kernels/amd.py

@@ -7,13 +7,12 @@ from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.helpers import prod
 from tinygrad.uop.ops import AxisType, KernelInfo, Ops, resolve
 
-BLOCK_M, BLOCK_N, WARP_SIZE = 32, 32, 32
+BLOCK_M, BLOCK_N, DECODE_HEAD_TILE, WARP_SIZE = 32, 32, 8, 32
 WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
 WAVES_M, WAVES_N, LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 2, 2, 16
 WMMA_ACC, THREADS_PER_BLOCK = WMMA_M // LANES_PER_WAVE_M, WARP_SIZE * WAVES_M * WAVES_N
 LDS_PAD, WMMA_ARG, LOG2E = 4, ((WMMA_M, WMMA_N, WMMA_K), 'AMD', 32), math.log2(math.e)
 Q4_K, Q5_K, Q6_K, IQ4_XS, GGML_BLOCK_SIZE, Q8_GROUP_SIZE, Q4_WORDS, Q5_WORDS, Q6_BYTES, IQ4_WORDS = 12, 13, 14, 23, 256, 32, 36, 44, 210, 34
-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:
@@ -55,43 +54,21 @@ 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()}
-    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)
+    raw = next((u for u in decoded.uop.toposort() if u.op is Ops.SHRINK and u.dtype == dtypes.uint8 and prod(u.shape) in packed_sizes), None)
     if raw is None: return
-    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 = ggml_type
+    self.ggml_type = packed_sizes[prod(raw.shape)]
     # store a typed buffer view: a lazy BITCAST is decomposed into byte-combining ALU before custom-kernel
     # scheduling and would copy the entire packed weight on every JIT graph
-    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)))
+    packed_dtype = dtypes.uint8 if self.ggml_type == Q6_K else dtypes.uint32
+    self.weight = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer)
+      .view(raw.max_numel() * raw.dtype.itemsize // packed_dtype.itemsize, packed_dtype, raw_offset)))
   def __call__(self, x:Tensor) -> Tensor:
     supported = self.use_custom_quant and amd_custom_kernels_supported(self.weight.device)
     if self.ggml_type is None and supported:
       self.set_quantized(self.weight)
-      if self.ggml_type is None:
-        # 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 is None: self.use_custom_quant = supported = False  # not a supported quant format
     if self.ggml_type in (Q4_K, Q5_K, Q6_K, IQ4_XS) and supported:
       if isinstance(x.numel(), int): return q8_linear(self, x)
       # symbolic token count: pad to the max chunk size so the kernels see static shapes, garbage rows are sliced off
@@ -100,20 +77,14 @@ class Linear(nn.Linear):
     return super().__call__(x)
 
 def _amd_dp4a(a:UOp, b:UOp, c:UOp) -> UOp:
-  # 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
+  return UOp(Ops.CUSTOMI, src=(a.int(), b.int(), c), arg=("__builtin_amdgcn_sudot4(true, {}, true, {}, {}, false)", dtypes.int32))
 
 def _amd_byte_perm(a:UOp, b:UOp, selectors:UOp) -> UOp:
   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
-  # nontemporal scalar load: streamed weights must not evict the activations/KV cache from L2
-  if lanes is None: return ptr.load(arg="nontemporal")
+  if lanes is None: return UOp(Ops.CUSTOMI, src=(ptr,), arg=("__builtin_nontemporal_load({0})", ptr.dtype))
   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)
@@ -122,19 +93,16 @@ 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]:
-  # 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))
+  scale = (subgroup < 4).where(_load_byte(raw, base, 4 + subgroup) & 63,
+    (_load_byte(raw, base, 8 + subgroup) & 15) | ((_load_byte(raw, base, subgroup) >> 6) << 4))
+  minimum = (subgroup < 4).where(_load_byte(raw, base, 8 + subgroup) & 63,
+    (_load_byte(raw, base, 8 + subgroup) >> 4) | ((_load_byte(raw, base, 4 + subgroup) >> 6) << 4))
   d, dmin = (raw[base] & 0xffff).cast(dtypes.uint16), (raw[base] >> 16).cast(dtypes.uint16)
   return _half(d), _half(dmin), scale.float(), minimum.float()
 
@@ -150,66 +118,52 @@ 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, xsum:UOp, x:UOp, tokens:int, in_features:int) -> UOp:
+def _q8_quantize_kernel(q:UOp, scale:UOp, x:UOp, tokens:int, in_features:int) -> UOp:
   groups = in_features//Q8_GROUP_SIZE
-  token_group, lane = UOp.range(tokens*groups, 0, axis_type=AxisType.GLOBAL), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  token_group, lane = UOp.range(tokens*groups, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
   token, group = token_group//groups, token_group%groups
   x = x.reshape(tokens, groups, 32)
   group_scale = (warp_reduce(x[token, group, lane].float().abs(), maximum=True, full_wave=True) / 127).maximum(1e-8)
   word_lane = lane.minimum(7)
   xs = tuple(x[token, group, word_lane*4+i].float() for i in range(4))
-  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()))))
+  word = sum(((v/group_scale).round().clip(-127, 127).cast(dtypes.int8).cast(dtypes.uint8).cast(dtypes.uint32) << (i*8)
+              for i,v in enumerate(xs)), UOp.const(0, dtypes.uint32))
+  stores = (q[token, group, lane.valid(lane < 8)].store(word), scale[token, group.valid(lane.eq(0))].store(group_scale))
   return UOp.group(*stores).end(token_group, lane).sink(arg=KernelInfo(name="q8_quantize", opts_to_apply=()))
 
-def q8_quantize(x:Tensor, tokens:int, in_features:int) -> tuple[Tensor, Tensor, Tensor]:
+def q8_quantize(x:Tensor, tokens:int, in_features:int) -> tuple[Tensor, Tensor]:
   groups = in_features//Q8_GROUP_SIZE
   q = Tensor.empty(tokens, groups, 8, dtype=dtypes.uint32, device=x.device)
   scale = Tensor.empty(tokens, groups, dtype=dtypes.float32, device=x.device)
-  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
+  q, scale = Tensor.custom_kernel(q, scale, x, fxn=functools.partial(_q8_quantize_kernel, tokens=tokens, in_features=in_features))[:2]
+  return q, scale
 
 def _decode_linear(out:UOp, out_features:int, group_count:int, group_dot, name:str) -> UOp:
-  chunks = 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))
+  chunks = (group_count+31)//32
+  token_output_chunk, lane = UOp.range(out.shape[0]*out_features*chunks, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  token, output, chunk = token_output_chunk // (out_features*chunks), (token_output_chunk//chunks) % out_features, token_output_chunk % chunks
+  group = lane+chunk*32
+  value = group_dot(token, output, group) if group_count % 32 == 0 else \
+    (group < group_count).where(group_dot(token, output, group.minimum(group_count-1)), UOp.const(0, dtypes.float32))
   total = warp_reduce(value, full_wave=True)
-  return out[token, output, chunk.valid(lane.eq(0))].store(total.cast(out.dtype)).end(token_output, chunk, lane).sink(
+  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, xs:UOp, out_features:int, in_features:int, ggml_type:int) -> UOp:
+def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, out_features:int, in_features:int, ggml_type:int) -> UOp:
   group_count = in_features // Q8_GROUP_SIZE
   def group_dot(token:UOp, output:UOp, group:UOp) -> UOp:
     block, subgroup = group // 8, group % 8
     xwords = _amd_load(xq[token, group, 0], 8)
     if ggml_type in (Q4_K, Q5_K):
       base = (output * in_features//GGML_BLOCK_SIZE + block) * (Q4_WORDS if ggml_type == Q4_K else Q5_WORDS)
-      qs_base, dot = 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))
+      qs_base, dot, qsum = base + (4 if ggml_type == Q4_K else 12) + (subgroup//2)*8, UOp.const(0, dtypes.int32), UOp.const(0, dtypes.int32)
       for word_idx in range(8):
-        word = (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)
+        word = (raw[qs_base+word_idx] >> ((subgroup&1)*4).cast(dtypes.uint32)) & 0x0f0f0f0f
+        if ggml_type == Q5_K: word |= ((raw[base+4+word_idx] >> subgroup.cast(dtypes.uint32)) & 0x01010101) << 4
+        dot, qsum = _amd_dp4a(word, xwords[word_idx], dot), _amd_dp4a(UOp.const(0x01010101, dtypes.uint32), xwords[word_idx], qsum)
       d, dmin, scale, minimum = _q5_scales(raw, base, subgroup)
-      gsum = xs[token, group, 0].load() + xs[token, group, 1].load()
-      return (dot.float()*d*scale - gsum*dmin*minimum) * xd[token, group]
+      return (dot.float()*d*scale - qsum.float()*dmin*minimum) * xd[token, group]
     if ggml_type == IQ4_XS:
       base = (output * in_features//GGML_BLOCK_SIZE + block) * IQ4_WORDS
       dot = UOp.const(0, dtypes.int32)
@@ -218,57 +172,45 @@ def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, xs:UOp, out_features:
         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
-    # 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))
+    base = (output*in_features//GGML_BLOCK_SIZE+block)*Q6_BYTES
     dots = [UOp.const(0, dtypes.int32)] * 2
     for word_idx in range(8):
-      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)
+      pos, within = subgroup*32 + word_idx*4, (subgroup*32 + word_idx*4)%128
+      low = _amd_load(raw[base + (pos//128)*64 + within%64], 4) >> ((within//64)*4).cast(dtypes.uint8)
+      high = _amd_load(raw[base + 128 + (pos//128)*32 + within%32], 4) >> ((within//32)*2).cast(dtypes.uint8)
+      quant = ((low & 15) | ((high & 3) << 4)).bitcast(dtypes.int8) - 32
+      word = sum((quant[i].cast(dtypes.uint8).cast(dtypes.uint32) << (i*8) for i in range(4)), UOp.const(0, dtypes.uint32))
       dots[word_idx//4] = _amd_dp4a(word, xwords[word_idx], dots[word_idx//4])
-    scales = [((raw[base + 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)
+    scales = [raw[base + 192 + subgroup*2+i].cast(dtypes.uint8).bitcast(dtypes.int8).float() for i in range(2)]
+    dbits = raw[base+208].cast(dtypes.uint16) | (raw[base+209].cast(dtypes.uint16) << 8)
+    return (dots[0].float()*scales[0] + dots[1].float()*scales[1]) * xd[token, group] * _half(dbits)
   names = {Q4_K: "linear_q4_k", Q5_K: "linear_q5_k", IQ4_XS: "linear_iq4_xs", Q6_K: "linear_q6"}
   return _decode_linear(out, out_features, group_count, group_dot, names[ggml_type])
 
 def _wmma_layout(out:UOp, out_features:int, token_tile:int, output_tiles:int):
   output_waves = 2 if out_features % (32*output_tiles) == 0 else 1
   token_block, output_block = UOp.range(out.shape[0]//token_tile, 0), UOp.range(out_features//(16*output_tiles*output_waves), 1)
-  # lane 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
+  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, src=(lane.int(),), arg=("__builtin_amdgcn_mbcnt_lo(-1, 0)", dtypes.int32)).cast(dtypes.weakint)
+  col, half = hw_lane % 16, hw_lane // 16
   outputs = tuple((output_block*output_waves+wave)*(16*output_tiles) + tile*16 + col for tile in range(output_tiles))
   inputs = tuple(token_block*token_tile + tile*16 + col for tile in range(token_tile//16))
   tokens = tuple(tuple(token_block*token_tile + tile*16 + half*8 + i for i in range(8)) for tile in range(token_tile//16))
   return output_waves, token_block, output_block, lane, wave, half, outputs, inputs, tokens
 
-def _wmma_stores(out, outputs, tokens, accs, update, half, 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))
+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, src=(value,),
+      arg=("__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {0}), 50688))", dtypes.float32)) for value in vals)
     low = half.eq(0)
-    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))]
+    return tuple(low.where(vals[i], swapped[i+4]) if j == 0 else low.where(swapped[i], vals[i+4]) for i in range(4) for j in range(2))
+  return [out[token, output].store(value) for output,output_accs in zip(outputs, accs)
+          for tile_tokens,acc in zip(tokens, output_accs) for token,value in zip(tile_tokens, values(acc))]
 
 def _quant_linear_wmma(out, x, out_features, in_features, type_words, layout, dequant, name):
   x = x.reshape(out.shape[0], in_features)
-  output_waves, token_block, output_block, lane, wave, physical_half, outputs, input_tokens, tokens = layout
+  _, 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)
@@ -287,8 +229,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)
-  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=()))
+  return UOp.group(*_wmma_stores(out, outputs, tokens, accs, update, physical_half)).end(token_block, output_block, lane, wave).sink(
+    arg=KernelInfo(name=name, opts_to_apply=()))
 
 @functools.cache
 def _q5_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, out_features:int, in_features:int, ggml_type:int) -> UOp:
@@ -322,10 +264,9 @@ 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)
-    # 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))
+    def nibble(packed:UOp, index:int): return (packed >> (8*index+4*half)) & 15
+    lut_pairs = (lut[(nibble(packed, i) | nibble(packed, i+1)<<4).cast(dtypes.weakint)]
+                 for packed in (raw[base+2+subgroup*4+i] for i in range(4)) for i in (0, 2))
     return tuple((_half((pair >> (i*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in lut_pairs for i in range(2))
   return _quant_linear_wmma(out, x, out_features, in_features, IQ4_WORDS, layout, dequant, "linear_iq4_xs_f16_wmma")
 
@@ -346,163 +287,67 @@ 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, xs = q8_quantize(x, tokens, in_features)
+  xq, xd = q8_quantize(x, tokens, in_features)
   decode = functools.partial(_quant_decode_kernel, ggml_type=layer.ggml_type)
   out = Tensor.empty(tokens, out_features, (in_features+1023)//1024, dtype=dtypes.float32, device=x.device).uop
-  return run(decode, out, raw, xq_.uop, xd.uop, 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)
+  return run(decode, out, raw, xq.uop, xd.uop)
 
 # ******** 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, waves=4):
+def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, valid_kv_len, max_kv_len, block_n):
   valid_kv_len = _unbind(valid_kv_len)
   _, B, H_KV, N, D = cast(tuple[int, int, int, int, int], cache_kv.shape)
   _, H, M, _ = cast(tuple[int, int, int, int], q.shape)
   assert M == 1 and H % H_KV == 0 and D % WARP_SIZE == 0 and max_kv_len <= N and max_kv_len % block_n == 0
-  G, CHUNK, 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=()))
+  G, CHUNK, DV, heads_per_wave = H // H_KV, block_n, D // WARP_SIZE, 2
+  head_tile = min(DECODE_HEAD_TILE, G)  # share each KV stream across two GQA heads per wave
+  assert G % head_tile == 0 and head_tile % heads_per_wave == 0
+  decode_waves, decode_group = head_tile // heads_per_wave, 4
+  block_bhkv = UOp.range(B*H_KV*(G//head_tile), 0, AxisType.GLOBAL)
+  valid_chunks = (valid_kv_len+CHUNK-1)//CHUNK
+  group_count = min(valid_chunks, out.shape[2]) if isinstance(valid_chunks, int) else valid_chunks.minimum(out.shape[2])
+  block_n, lane = UOp.range(group_count, 1, AxisType.GLOBAL), UOp.range(WARP_SIZE, 2, axis_type=AxisType.LOCAL)
+  wave = UOp.range(decode_waves, 3, axis_type=AxisType.LOCAL)
+  head_group, bhkv = block_bhkv % (G//head_tile), block_bhkv // (G//head_tile)
+  b, kv_head = bhkv // H_KV, bhkv % H_KV
+  dims = tuple(lane + i*WARP_SIZE for i in range(DV))
+  acc, row_max, row_sum = _reg((heads_per_wave, DV), 0, 0), _reg((heads_per_wave,), 1, -math.inf), _reg((heads_per_wave,), 2, 0)
+  groups_per_chunk, offset = CHUNK // decode_group, UOp.range(((valid_chunks+group_count-1)//group_count)*(CHUNK//decode_group), 100, AxisType.REDUCE)
+  chunk = block_n + (offset // groups_per_chunk) * group_count
+  keys = tuple(chunk*CHUNK + (offset % groups_per_chunk)*decode_group + i for i in range(decode_group))
+  valid = tuple(key < valid_kv_len for key in keys)
+  kvals, vvals = (tuple(tuple(is_valid.where(cache_kv[kv, b, kv_head, key, d].float(), UOp.const(0, dtypes.float)) for d in dims)
+    for key,is_valid in zip(keys, valid)) for kv in range(2))
+  q_heads = tuple(kv_head*G + head_group*head_tile + wave*heads_per_wave + head for head in range(heads_per_wave))
+  updates:list[UOp] = []
+  for head,q_head in enumerate(q_heads):
+    scores = tuple(warp_reduce(sum((q[b, q_head, 0, d].float()*k for d,k in zip(dims, key_kvals)),
+                                   UOp.const(0, dtypes.float)), full_wave=True) / math.sqrt(D) for key_kvals in kvals)
+    prev_acc, prev_max, prev_sum = acc.after(offset)[head], row_max.after(offset)[head], row_sum.after(offset)[head]
+    new_max = functools.reduce(lambda a,vs:a.maximum(vs[0].where(vs[1], UOp.const(-math.inf, dtypes.float))), zip(valid, scores), prev_max)
+    alpha = ((prev_max-new_max)*LOG2E).exp2()
+    betas = tuple(is_valid.where(((score-new_max)*LOG2E).exp2(), UOp.const(0, dtypes.float)) for is_valid,score in zip(valid, scores))
+    updates += [acc[head].store(prev_acc*alpha + sum((UOp.stack(*value)*beta for value,beta in zip(vvals, betas)), acc[head].const_like(0))),
+                row_sum[head].store(prev_sum*alpha + sum(betas, UOp.const(0, dtypes.float))), row_max[head].store(new_max)]
+  update = UOp.group(*updates).end(offset)
+  acc, row_max, row_sum = acc.after(update), row_max.after(update), row_sum.after(update)
+  stores = [out[b, q_head, block_n, d].store(acc[head, i]) for head,q_head in enumerate(q_heads) for i,d in enumerate(dims)] + \
+    [stats[b, q_head.valid(lane.eq(0)), block_n, i].store(x[head]) for head,q_head in enumerate(q_heads) for i,x in enumerate((row_max, row_sum))]
+  return UOp.group(*stores).end(lane, wave, block_n, block_bhkv).sink(arg=KernelInfo(name="flash_decode_partial", opts_to_apply=()))
 
 def amd_flash_attention_decode(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, max_kv_len:int) -> Tensor:
   B, H, D = cache_kv.shape[1], q.shape[1], cache_kv.shape[4]
-  chunks = min(256, max_kv_len // 64)
+  chunks = min(64, max_kv_len // 128)
   partial = Tensor.empty(B, H, chunks, D, dtype="float32", device=q.device)
   stats = Tensor.empty(B, H, chunks, 2, dtype="float32", device=q.device)
-  fxn = functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len, block_n=64, waves=16)
+  fxn = functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len, block_n=128)
   partial, stats = Tensor.custom_kernel(partial, stats, q, cache_kv, fxn=fxn)[:2]
-  live = (valid_kv_len+63)//64
+  live = (valid_kv_len+127)//128
   live = min(live, chunks) if isinstance(live, int) else live.minimum(chunks)
-  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]
+  partial, stats = partial[:, :, :live], stats[:, :, :live]
+  weights = ((stats[..., 0]-stats[..., 0].max(2, keepdim=True))*LOG2E).exp2()
+  return ((partial*weights.unsqueeze(-1)).sum(2) / (stats[..., 1]*weights).sum(2, keepdim=True)).unsqueeze(2)
 
 @functools.cache
 def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, valid_kv_len:int|UOp, q_start:int|UOp|None=None) -> UOp:

tinygrad/mixin/elementwise.py

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

tinygrad/mixin/op.py

@@ -549,6 +549,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     squares = (self - self.mean(axis=axis, keepdim=True)).square()
     n = prod([si for si, so in zip(self.shape, squares.sum(axis=axis, keepdim=True).shape) if resolve(si != so)])
     numerator = squares.cast(sum_acc_dtype(self.dtype)).sum(axis=axis, keepdim=keepdim)
+    if resolve(n == 1, False) and correction >= 1: return self.sum(axis=axis, keepdim=keepdim).cast(output_dtype) * math.nan
     return numerator.div(smax(n - correction, 0)).cast(output_dtype)
 
   def var_mean(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> tuple[Self, Self]:

tinygrad/renderer/__init__.py

@@ -1,17 +1,12 @@
 from __future__ import annotations
 from typing import Callable, cast
-from dataclasses import dataclass, replace
+from dataclasses import dataclass
 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

@@ -188,12 +188,10 @@ 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_ptr(self, u:UOp):
-    # the address of an access, vector-cast if the access reads/writes more lanes than the pointer's scalar type
+  def render_access(self, u:UOp):
     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]}"
-  def render_access(self, u:UOp): return f"*{self.render_ptr(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_cast(self, u:UOp, val:str) -> str: return f"({self.render_type(u)})({val})"
 
   # LEGACY
@@ -511,9 +509,6 @@ 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, with_storage
+from tinygrad.renderer import Renderer
 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,12 +123,11 @@ class NIRRenderer(Renderer):
   extra_matcher = PatternMatcher([
     # from ptx
     (UPat.var('x', dtype=dtypes.bool)<UPat.var('y'), lambda x,y: (x^True)&y),
-    # 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
+    # load/store bool -> uint8
     (UPat(Ops.LOAD, dtypes.bool, name="x"),
-     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:])),
+     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:])),
     # 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, with_storage
+from tinygrad.renderer import Renderer
 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))),
-  # 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
+  # load/store bool -> uint8 (only for memory, not registers)
   (UPat(Ops.LOAD, dtypes.bool, src=(UPat(name="idx"),), name="x", allow_any_len=True),
-   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),
+   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),
   (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:]) if idx.addrspace != AddrSpace.REG else None),
+   lambda x,idx: UOp(x.op, src=(x.src[0], x.src[1].cast(dtypes.uint8))+x.src[2:]) if idx.addrspace != AddrSpace.REG else None),
   # ptx shr and shl instructions require y to be uint
   (UPat.var("x") << UPat.var("y"), lambda x,y: UOp(Ops.SHL, src=(x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
   (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, src=(x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),

tinygrad/runtime/autogen/__init__.py

@@ -10,14 +10,10 @@ 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 = (
@@ -54,8 +50,6 @@ 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":
@@ -105,11 +99,14 @@ 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=_extract_deb)
+                  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))
     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=_extract_deb)
-    case "vfio": return load("vfio", ["{}/usr/include/linux/vfio.h"], args=["-I{}/usr/include"], srcs=linux_headers_deb, preprocess=_extract_deb)
+    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))
     # 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"])
@@ -119,10 +116,9 @@ 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":
-      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 {}))
+      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"])
     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",
@@ -185,17 +181,6 @@ 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=_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)
+                  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))
     case _: raise AttributeError(f"no such autogen: {nm}")

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, field
+from dataclasses import replace, dataclass
 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,8 +30,9 @@ class HCQInfo:
   device:tuple[str, ...]
   estimates:Estimates = Estimates()
 
-  inputs:int|None = None
-  input_addrs:tuple[tuple[str, UOp], ...] = () # (device, lane arg uop)
+  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)
   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
@@ -43,8 +44,6 @@ 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:
@@ -155,131 +154,138 @@ 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)
 
-@dataclass(frozen=True)
-class Dep: dev:str; queue:str; tag:int; lane:int # producer submit (dev, queue, tag) awaited by consumer lane # noqa: E702
-
-@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))]
+  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))]
 
-def _wait_ins(ctx:BatchCtx, bufs_by_lane:list[list[Any]], write, devices:tuple[str, ...], queue:str, tag:int) -> list[UOp]:
-  deps:list[Dep] = []
+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]] = []
   for lane, bufs in enumerate(bufs_by_lane):
     written = write if write is not None else list(range(len(bufs)))
-    deps += [Dep(d, q, t, lane) for d, q, t in ctx.tracker.access_resources(bufs, written, (devices[lane], queue, tag)) if t < tag]
+    dep_lanes += [(dep, dlane, lane) for dep, dlane in ctx.access_resources(bufs, written, (key, lane))]
+  return dep_lanes
 
-  # same-queue submits are fifo-ordered, no wait needed
+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
   if devices[0].split(":")[0] in {"AMD", "QCOM", "CPU"} or queue.startswith("COPY"):
-    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)}
+    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)
 
-  # 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 (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
+  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}
 
-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()))
-
-def _merge_queues(submits:list[UOp]) -> list[UOp]:
-  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 = 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 _make_finalizers(ctx:BatchCtx) -> tuple[list[UOp], list[UOp], list[UOp]]:
+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 ctx.batch:
+  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 = len(ctx.batch_info), [], [], []
+  n, fences, resets, fins, signal_tags = len(batch_info), [], [], [], set()
   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")
+    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")
-
-    # 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)))
-
-    # 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)):
+    # 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?
+  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()))
+
+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
+  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()]
+
+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]
+
+  # 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
+
+  # build fences and finalizers
+  fences, fins, finalizer_signal_tags = _build_finalizers(batch, batch_info, deps_tracker, slots)
+  signal_tags |= finalizer_signal_tags
+
+  src, kerns = [], []
+  for tag, ((call, _), (devices, queue), q) in enumerate(zip(batch, batch_info, call_waits)):
     # first queue use, sync prior device work with the device timeline
-    if (devices, queue) not in seen_queues:
-      seen_queues.add((devices, queue))
+    if batch_info.index((devices, queue)) == tag:
       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 ctx.profile else []
+    ts_ids = [next(UOp.unique_num) for _ in range(2)] if 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 ctx.signal_tags: q += [UOp(Ops.INS, arg="store", src=(make_buf(devices, ctx.slots[queue]), UOp.const(tag + 1, dtypes.uint64)))]
+    if tag in signal_tags: q += [UOp(Ops.INS, arg="store", src=(make_buf(devices, 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 + resets + _merge_queues(submits) + fins
+  return fences + _merge_queues(src) + fins
 
-@rewrite_group(new_ctx=False)
-def sched_batches(l:UOp, profile:bool) -> UOp:
+def sched_hcq_batches(l:UOp, profile:bool) -> UOp:
   srcs:list[UOp] = []
   batch:list[tuple[UOp, tuple[str, ...]]] = []
   for call in l.src:
@@ -287,6 +293,9 @@ def sched_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
 
@@ -316,20 +325,18 @@ 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 _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]]:
+def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[UOp, dict[UOp, UOp], tuple[UOp, ...], dict[UOp, int]]:
   bare = {g: g.replace(src=(g.src[0].without_after,)) for g in gaddrs}
 
-  # 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)
+  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)
 
-  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 ()
+  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 ()
   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], stride:int) -> dict[UOp, UOp]:
+def make_gather_loop(patches:list[UOp], table:UOp, slots:dict[UOp, int], lt_patches:list[UOp]) -> 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)]
 
@@ -337,13 +344,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()-(stride-1))))
+  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())))
   # 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 + _dnum(stride)).load()).end(r)
+    .store(table.index(slot).load()).end(r)
   return {p: UOp(Ops.NOOP) for p in patches} | {patches[0]: patch}
 
-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 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 split_patches(call:UOp) -> UOp|None:
   rt_patches:list[UOp] = []
@@ -351,22 +358,20 @@ 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
-  lanes = len(to_tuple(call.arg.aux.device))
-  inputs, internals = partition(dedup(get_getaddrs(UOp.sink(body, *rt_patches))), is_input_addr)
+  inputs, internals = partition(dedup([g for p in rt_patches for g in get_getaddrs(p)] + get_getaddrs(body)), 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, lanes if n == "inputs" else 1) for gs,n in ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems"))]
+  tables = [make_addr_table(call, gs, n) 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, lanes) if ipatches else {}
-  body = body.substitute({p:p.substitute(gathers | reads) for p in rt_patches} | reads, walk=True)
+  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)
 
   lt_srcs = collections.defaultdict(list)
   for p in lt_patches: lt_srcs[p.buf_uop].append(p)
-
-  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))
+  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)))
 pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), split_patches)])
 
 # *****************
@@ -435,12 +440,13 @@ 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)
-  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
+  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
   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, off in zip(batch, offs) for j in range(len(c.arg.aux.device))]
+          for (c, j, _), off in zip(lanes, offs)]
 
   info = HCQInfo((HCQ_RUNTIME_DEV.value,), sum((c.arg.aux.estimates for c in batch), start=Estimates()).simplify(),
-                 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))
+                 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))
   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)
 
@@ -475,8 +481,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 on real buffers
-    linear = sched_batches(linear, profile).substitute({s:p for p,s in back_map.items()}, walk=True, enter_calls=True)
+    # 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")
 
     # 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, ctypes, functools, itertools
+import struct, random, socket, 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, ipv4_to_gid
+from tinygrad.runtime.support.system import PCIDevice
 from tinygrad.runtime.autogen import mlx5, pci
 
 MLX_DEBUG = getenv("MLX_DEBUG", 0)
@@ -11,6 +11,8 @@ 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,8 +10,6 @@ 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

@@ -80,8 +80,9 @@ def create_schedule(sched_sink:UOp) -> UOp:
 
 from tinygrad.schedule.memory import memory_plan_rewrite
 from tinygrad.engine.realize import capturing, pm_flatten_linear
+#from tinygrad.schedule.rangeify import get_kernel_graph
+from tinygrad.schedule.rangeify2 import get_kernel_graph
 from tinygrad.schedule.prepare import prepare_rangeify
-from tinygrad.schedule.rangeify import get_kernel_graph
 from tinygrad.helpers import CAPTURING
 from tinygrad.uop.ops import PatternMatcher, UPat, ParamArg
 from tinygrad.dtype import AddrSpace

tinygrad/schedule/allreduce.py

@@ -58,8 +58,11 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
   return UOp.usum(*[c.pad(((s,numel-e),)) for (s,e),c in zip(chunks, copied_chunks)]).reshape(shape)
 
 def create_allreduce_function(buf:UOp, red:UOp, output:UOp|None=None) -> UOp|None:
-  if output is None: output = UOp.invalids(red.shape, dtype=red.dtype, device=red.device)
+  if output is None:
+    call_output = UOp.invalids(red.max_shape, dtype=red.dtype, device=red.device)
+    output = call_output.shrink_to(red.shape)
+  else: call_output = output
   to = red.param_like(0)
   src = buf.param_like(1)
   red = src.allreduce(*red.arg)
-  return output.after(to.after(to.store(handle_allreduce(src, red))).sink().call(output, buf.contiguous(), name="allreduce", precompile=True))
+  return output.after(to.after(to.store(handle_allreduce(src, red))).sink().call(call_output, buf.contiguous(), name="allreduce", precompile=True))

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, colored, Context, SPEC
+from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC
 
 @dataclass
 class IndexingContext:
@@ -248,12 +248,13 @@ 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:
+        if all_all_same or (PCONTIG and all_same(local_rngs)):
           # 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"))
@@ -267,11 +268,15 @@ 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 = list(range(len(out_rngs)))
+      _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)
       ending_ranges[x] = []
       if len(_realize_axis):
         rctx.realize_map[x] = _realize_axis
-        out_rngs = tuple(rctx.new_range(x.shape[i]) for i in range(len(out_rngs)))
+        out_rngs = tuple([(rctx.new_range(x.shape[i]) if i in _realize_axis else r) for i,r in enumerate(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

@@ -15,8 +15,7 @@ def mstack_early_shrink(ms:UOp, shrink:UOp):
   ret:list[UOp] = []
   for i, x in enumerate(ms.src):
     if x.op is Ops.COPY:
-      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))
+      ret.append(_apply_shrink(shrink.marg, x.src[0], i).copy_to_device(x.device))
     else:
       ret.append(_apply_shrink(shrink.marg, x, i).contiguous())
   return ms.replace(src=tuple(ret))
@@ -34,6 +33,8 @@ replace_allreduce = PatternMatcher([
     x.mselect(0).copy_to_device(c.device) if isinstance(c.device, str) and isinstance(x.device, tuple) else None),
   # MSELECT on MSTACK is replaced with nothing
   (UPat(Ops.MSELECT, src=(UPat(Ops.MSTACK, name="mstack"),), name="ms"), lambda mstack, ms: mstack.src[ms.arg]),
+  # Identical device-less values do not need a multi wrapper.
+  (UPat(Ops.MSTACK, src=(UPat.var("x"),), allow_any_len=True, name="m"), lambda m,x: x if x.device is None and all_same(m.src) else None),
   # 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/ALU ops

tinygrad/schedule/prepare.py

@@ -9,10 +9,15 @@ 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:])
+  if u.op in GroupOp.Movement or u.op in {Ops.INDEX, Ops.UNSHARD}: return walk_mop(u.src[0])
   return u
 
+def has_buffer_view(u:UOp) -> bool:
+  # CALL argument lowering currently passes the base allocation, so only an
+  # offset-zero contiguous view backed by a real buffer can avoid a copy.
+  if u.has_buffer_identity(after_ok=True): return True
+  return (cv:=u.contiguous_view()) is not None and cv[1] == 0 and cv[0].has_buffer_identity(after_ok=True)
+
 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:
@@ -140,6 +145,10 @@ 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),
 
@@ -147,9 +156,9 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   (UPat(Ops.COPY, src=(UPat(Ops.RESHAPE, name="shp"),), name="cpy"), lambda shp,cpy: shp.src[0].copy_to_device(cpy.device).reshape(shp.shape)),
 
   # reshaping on STORE can be a NOOP
-  (UPat(Ops.STORE, src=(UPat(Ops.RESHAPE, src=(UPat.var("dst",),), allow_any_len=True),
-                        UPat(Ops.RESHAPE, src=(UPat.var("src",),), allow_any_len=True))),
-   lambda dst,src: dst.store(src) if dst.shape == src.shape else None),
+  #(UPat(Ops.STORE, src=(UPat(Ops.RESHAPE, src=(UPat.var("dst",),), allow_any_len=True),
+  #                      UPat(Ops.RESHAPE, src=(UPat.var("src",),), allow_any_len=True))),
+  # lambda dst,src: dst.store(src) if dst.shape == src.shape else None),
 
   # ** store rules **
 
@@ -176,6 +185,19 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   # 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),
 
+  # ** new prepare **
+
+  # CALL inputs need buffer identity (and to be flat)
+  (UPat(Ops.CALL, name="c"),
+   lambda c: c.replace(src=c.src[0:1]+tuple(x if has_buffer_view(x) else x.contiguous() for x in c.src[1:]))),
+
+  # MSTACK inputs need buffer identity
+  (UPat(Ops.MSTACK, name="c"),
+   lambda c: c.replace(src=tuple(x.contiguous() if not x.has_buffer_identity(after_ok=True) else x for x in c.src))),
+
+  # STORE to () is reshaped to (1,)
+  (UPat(Ops.STORE, name="s"), lambda s: s.src[0].reshape((1,)).store(s.src[1].reshape((1,))) if s.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))),
@@ -193,12 +215,13 @@ def convert_copy_to_store(ctx, copy:UOp, existing_buf:UOp|None=None):
     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 = buf.shrink_to(input_src.shape)
   # reshape back to input
   return buf.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),
-  (UPat(Ops.COPY, name="copy"), convert_copy_to_store),
+  (UPat((Ops.COPY, Ops.CONTIGUOUS), name="copy"), convert_copy_to_store),
 ])
 
 @rewrite_group(new_ctx=False)

tinygrad/schedule/rangeify.py

@@ -6,7 +6,7 @@ from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, K
 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 get_single_element
+from tinygrad.helpers import PCONTIG, partition, 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: return None
+  if len(accessed_buffers) > 3 and not (PCONTIG > 2): return None
 
   # if any reduces access a buffer, don't remove this buffer
   buffer_in_reduce = False
@@ -94,7 +94,22 @@ 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:
-    return None
+    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
 
   # if it makes it here, the bufferize is removed
   # this is the ranges replaced

tinygrad/schedule/rangeify2.py

@@ -0,0 +1,426 @@
+from dataclasses import dataclass, field
+import itertools
+from tinygrad.dtype import AddrSpace, Invalid, strong_dtype
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, GroupOp, KernelInfo
+from tinygrad.uop.ops import graph_rewrite, AxisType, rewrite_group, remove_all_tags, resolve, shape_to_shape_arg
+from tinygrad.helpers import all_int, prod, VIZ, SPEC, Context, panic
+from tinygrad.schedule.indexing import BufferizeOpts, apply_movement_op
+from tinygrad.schedule.prepare import has_buffer_view
+from tinygrad.uop.symbolic import symbolic
+from tinygrad.codegen.simplify import pm_reduce_simplify
+
+# *** preparation ***
+
+fix_mselect_mstack = PatternMatcher([
+  # 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)),
+])
+
+from tinygrad.helpers import all_same
+from tinygrad.uop.ops import _broadcast_shape
+
+def expand_broadcast(x:UOp):
+  shapes = [u._shape for u in x.src]
+  if any(s is None for s in shapes) or all_same(shapes): return None
+  shape = _broadcast_shape(*shapes)
+  return x.replace(src=tuple([u.expand(shape) for u in x.src]))
+
+pm_lil_prepare_graph = PatternMatcher([
+  # expand broadcasts first
+  (UPat(GroupOp.Binary|GroupOp.Ternary|{Ops.STORE}, name="x"), expand_broadcast),
+])+fix_mselect_mstack
+
+# *** RANGE creation ***
+
+def rangeify_on_reduce(ctx, inp:UOp, red:UOp, idx:UOp|None=None):
+  if red.arg[1] == 0: return None
+  if idx is None and len(red.shape) > 0: return None
+  # TODO: is AxisType.REDUCE a real thing?
+  rngs = [UOp.range(s, next(ctx), AxisType.REDUCE) for s in inp.shape[:red.arg[1]]]
+  return inp.index(*rngs, *(idx.src[1:] if idx is not None else ())).reduce(*rngs, arg=(red.arg[0], 0))
+
+def rangeify_on_store(ctx, x:UOp):
+  if x.shape == (): return None
+  rngs = [UOp.range(s, next(ctx)) for s in x.shape]
+  return x.src[0].index(*rngs).store(x.src[1].index(*rngs)).end(*rngs)
+
+def rangeify_on_stage(ctx, x:UOp):
+  if x.src[0].shape == (): return None
+  # size 1 dims don't get ranges, they are reshaped out and back in
+  if all_int(x.shape) and 0 < len(sq := tuple(s for s in x.shape if s != 1)) < len(x.shape):
+    return rangeify_on_stage(ctx, x.src[0].reshape(sq).bufferize(arg=x.arg)).reshape(x.shape)
+  rngs = [UOp.range(s, next(ctx)) for s in x.shape]
+  return x.replace(src=(x.src[0].index(*rngs), *rngs))
+
+pm_range_creation = PatternMatcher([
+  # reduce/store are what creates ranges
+  (UPat(Ops.REDUCE, src=(UPat.var('inp'),), name="red").index(name="idx", allow_any_len=True), rangeify_on_reduce),
+  (UPat(Ops.REDUCE, src=(UPat.var('inp'),), name="red"), rangeify_on_reduce),
+  (UPat(Ops.STORE, name="x"), rangeify_on_store),
+  (UPat(Ops.STAGE, name="x"), rangeify_on_stage),
+])
+
+# *** RANGE migration ***
+
+# movement op on INDEX as a PatternMatcher
+def _mop_index(r:UOp, idx:UOp):
+  idxs = idx.src[1:]
+  if len(idxs) == len(r.shape):
+    ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idxs), dtype=idx.dtype, arg=idx.arg)
+    if r.op is Ops.PAD:
+      # NOTE: neither 0 or ret.const_like(0) is correct here.
+      # const_like breaks because it adds casts, and 0 is wrong if ret is a bool
+      invalid_value = UOp.const(ret.dtype.const(0))
+      # insert invalid_value for PAD with where
+      a = UOp.const(True)
+      for s in UOp.sink(*ret.src[1:]).simplify().src:
+        if s.is_invalid: return invalid_value
+        if s.op is Ops.WHERE and s.src[2].op is Ops.CONST and s.src[2].arg == Invalid: a = a & s.src[0]
+      ret = a.where(ret, invalid_value)
+    return ret
+  if r.op is Ops.RESHAPE:
+    src_prefix = len(r.src[0].shape) - len(r.shape[len(idxs):])
+    if src_prefix >= 0 and r.src[0].shape[src_prefix:] == r.shape[len(idxs):]:
+      if src_prefix == 0: return r.src[0] if r.src[0].dtype == idx.dtype else None
+      ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape[:src_prefix], r.shape[:len(idxs)], idxs), dtype=idx.dtype, arg=idx.arg)
+      return ret if ret.shape == idx.shape else None
+
+# TODO: this should be in _mop_index
+def index_on_stack(stack:UOp, idx:UOp):
+  srcs = [s.index(*idx.src[2:]) for s in stack.src]
+  r0 = idx.src[1]
+  ret = srcs[-1]
+  for k in range(len(srcs)-2, -1, -1): ret = r0.eq(k).where(srcs[k], ret)
+  return ret
+
+pm_range_migration = PatternMatcher([
+  # STAGE on shape () is nothing
+  (UPat(Ops.STAGE, src=(UPat.var('x'),)), lambda x: x if x.shape == () else None),
+  # 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),
+  # handle movement ops on INDEX
+  (UPat(GroupOp.Movement, name="r").index(name="idx", allow_any_len=True), _mop_index),
+  (UPat(Ops.STACK, name="stack").index(name="idx", allow_any_len=True), index_on_stack),
+  # move movement ops and INDEX after AFTER
+  (UPat(GroupOp.Movement|{Ops.INDEX}, name="r").after(name="a", allow_any_len=True),
+   lambda r,a: UOp(r.op, src=(a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], arg=r.arg)),
+  # block bitcast that changes shape
+  (UPat(Ops.BITCAST, name="b").index(allow_any_len=True),
+   lambda b: panic(RuntimeError, "shape changing bitcast not allowed in rangeify") if b.src[0].shape != b.shape else None),
+  # pass index through elementwise
+  (UPat(GroupOp.Elementwise, name="b").index(name="idx", allow_any_len=True),
+   lambda b,idx: b.replace(src=tuple(s.index(*idx.src[1:]) for s in b.src))),
+  # INDEX without src is nothing (must be at the bottom)
+  (UPat(Ops.INDEX, src=(UPat.var('x'),)), lambda x: x),
+])
+
+# *** split into kernels ***
+
+@dataclass
+class SplitCtx:
+  call_args:list[UOp] = field(default_factory=list)
+  buffers:dict[UOp, int] = field(default_factory=dict)
+  range_number:int = -1
+  addrspace:AddrSpace = AddrSpace.GLOBAL
+
+def _split_graph(ctx:SplitCtx, u:UOp) -> UOp|None:
+  if u.tag is not None: return None
+  if u.addrspace != ctx.addrspace: return None
+  if u.addrspace == AddrSpace.ALU: return u.param_like(-1).rtag().reshape(u.shape)
+
+  # A kernel takes each underlying buffer state once. In particular, AFTER and its buffer must use the same slot, with AFTER kept as the call
+  # argument so its dependencies are preserved.
+  key = u.buf_uop if u.op is Ops.AFTER else u
+  if (slot:=ctx.buffers.get(key)) is None:
+    slot = ctx.buffers[key] = len(ctx.call_args)
+    ctx.call_args.append(u)
+  elif u.op is Ops.AFTER:
+    ctx.call_args[slot] = u
+
+  # Parameters describe the max-sized physical allocation. A symbolic logical shape is a view of that allocation, not part of the PARAM itself.
+  param = u.param_like(slot).rtag().replace(src=(shape_to_shape_arg((u.max_numel(),)),))
+  return param.reshape(u.max_shape).shrink_to(u.shape)
+
+def _renumber_range(ctx:SplitCtx, u:UOp) -> UOp|None:
+  if u.tag is not None: return None
+  ctx.range_number += 1
+  return u.replace(arg=(ctx.range_number, u.arg[-1])).rtag()
+
+pm_split_graph = pm_range_migration+PatternMatcher([
+  (UPat((Ops.PARAM, Ops.AFTER, Ops.BUFFER, Ops.MSELECT, Ops.MSTACK), name="u"), _split_graph),
+  (UPat(Ops.RANGE, name="u"), _renumber_range),
+])
+
+def _is_fully_invalid_state(x:UOp) -> bool:
+  while x.op in GroupOp.Movement|{Ops.INDEX}: x = x.src[0]
+  if x.op is not Ops.AFTER or len(x.src) != 2: return False
+  end = x.src[1]
+  st = end.src[0] if end.op is Ops.END else end
+  if st.op is not Ops.STORE or not st.src[1].base.is_invalid: return False
+  if end.op is not Ops.END: return st.src[0].max_numel() == x.max_numel()
+  covered = 1
+  for r in end.src[1:]:
+    if r.op is not Ops.RANGE or r.src[0].op is not Ops.CONST or not isinstance(r.src[0].val, int): return False
+    covered *= r.src[0].val
+  return covered == x.max_numel()
+
+def split_store(x:UOp) -> UOp|None:
+  st = x.src[0] if x.op is Ops.END else x
+  if st.op is Ops.STORE and st.src[0].is_variable: return None
+  if st.op is Ops.STORE and st.src[0] is st.src[1]: return UOp(Ops.NOOP)
+  # A directly-invalid value makes this store a no-op. An AFTER carrying an
+  # invalid partial store is still a valid buffer state: uncovered elements
+  # must continue to read from the previous state.
+  if st.op is Ops.STORE and (st.src[1].base.is_invalid or _is_fully_invalid_state(st.src[1])): return UOp(Ops.NOOP)
+  ret = graph_rewrite(x, pm_split_graph, ctx:=SplitCtx(), name="split kernel", bottom_up=True)
+  # TODO: params and args should be able to be in any order
+  ctx.addrspace = AddrSpace.ALU
+  ret = graph_rewrite(ret, pm_split_graph, ctx, name="split kernel (vars)", bottom_up=True)
+  ret = graph_rewrite(ret, remove_all_tags, name="remove split tags", bottom_up=True)
+  return ret.sink(arg=KernelInfo()).call(*ctx.call_args)
+
+split_kernels = PatternMatcher([
+  (UPat((Ops.STORE, Ops.END), name="x"), split_store),
+])
+
+# cleanups
+
+def strip_zero_offset_shrink(x:UOp) -> UOp:
+  return x.src[0] if x.op is Ops.SHRINK and all(resolve(start == 0, False) for start,_ in x.marg) else x
+
+def no_indexing_calls(u:UOp):
+  new_srcs = []
+  for x in u.src:
+    if x.op is Ops.INDEX:
+      # sometimes if call srcs have children the call will get an INDEX. we remove it here.
+      # TODO: we should add safety checks here for contiguous
+      new_srcs.append(x.src[0])
+    elif x.op is Ops.SHRINK:
+      # SHRINK with offset 0 is fine
+      new_srcs.append(strip_zero_offset_shrink(x))
+    elif x.op is Ops.MSTACK:
+      new_srcs.append(x.replace(src=tuple(strip_zero_offset_shrink(s) for s in x.src)))
+    else:
+      # everything else we pass through
+      new_srcs.append(x)
+  return u.replace(src=tuple(new_srcs))
+
+pm_no_indexing_calls = PatternMatcher([
+  (UPat(Ops.CALL, name="u"), no_indexing_calls),
+  (UPat(Ops.AFTER, name="u"), lambda u: u.replace(src=tuple(s for s in u.src if s.op is not Ops.NOOP))),
+])
+
+# *** main rangeify ***
+
+debug_tag_factor = PatternMatcher([
+  (UPat(GroupOp.All, name="x"), lambda ctx,x: x.rtag(ctx[0][x] if x not in ctx[1] else 'REAL') if x.tag is None else None),
+])
+
+def remove_stage(ctx, x:UOp) -> UOp:
+  dtype = strong_dtype(x.dtype)
+  buf = UOp.new_buffer(x.arg.device, x.max_numel(), dtype, num=next(ctx))
+  val = x.src[0] if x.src[0].dtype == dtype else x.src[0].cast(dtype)
+  return buf.after(buf.reshape(x.shape).index(*x.src[1:]).store(val).end(*x.src[1:])).reshape(x.shape)
+
+pm_remove_stage = PatternMatcher([
+  (UPat(Ops.STAGE, name="x"), remove_stage),
+])+fix_mselect_mstack
+
+def remove_selected_stage(ctx:set[UOp], stage:UOp, idx:UOp) -> UOp|None:
+  return stage.src[0] if stage in ctx and stage.src[1:] == idx.src[1:] else None
+
+pm_remove_selected_stage = PatternMatcher([
+  (UPat(Ops.STAGE, name="stage").index(name="idx", allow_any_len=True), remove_selected_stage),
+])
+
+def inline_stage_index(stage:UOp, idx:UOp) -> UOp:
+  replacements, cache = dict(zip(stage.src[1:], idx.src[1:])), {}
+  def replace(x:UOp) -> UOp:
+    if x in replacements: return replacements[x]
+    if x.has_buffer_identity(after_ok=True) or x.op is Ops.STAGE: return x
+    if x not in cache: cache[x] = x.replace(src=tuple(replace(s) for s in x.src))
+    return cache[x]
+  return replace(stage.src[0])
+
+MAX_RECOMPUTE = 8
+MAX_SCALAR_RECOMPUTE = 64
+
+def recompute_cost(x:UOp, seen:set[UOp]|None=None) -> int|None:
+  if seen is None: seen = set()
+  if x in seen or x.op is Ops.STAGE or x.has_buffer_identity(after_ok=True): return 0
+  seen.add(x)
+  if x.op is Ops.REDUCE: return None
+  costs = [recompute_cost(s, seen) for s in (x.src[:1] if x.op is Ops.INDEX else x.src)]
+  return None if any(c is None for c in costs) else sum(c for c in costs if c is not None) + (x.op in GroupOp.Elementwise)
+
+def materialize_call_args(c:UOp) -> UOp:
+  srcs:list[UOp] = []
+  for x in c.src[1:]:
+    device = x.device or c.device
+    srcs.append(x if x.op is Ops.STAGE or x.is_bound_var or has_buffer_view(x) or x.shape == () or device is None
+                else x.bufferize(arg=BufferizeOpts(device=device)))
+  return c.replace(src=(c.src[0], *srcs))
+
+def materialize_mselect(m:UOp, x:UOp) -> UOp|None:
+  if x.device is None or x.op is Ops.STAGE or (x.op not in GroupOp.ALU and x.has_buffer_identity(after_ok=True)): return None
+  return m.replace(src=(x.bufferize(arg=BufferizeOpts(device=x.device)),))
+
+pm_materialize_call_args = PatternMatcher([
+  (UPat(Ops.CALL, name="c"), materialize_call_args),
+  (UPat(Ops.MSELECT, src=(UPat(name="x"),), name="m"), materialize_mselect),
+])
+
+@rewrite_group(new_ctx=False)
+def get_kernel_graph(sink:UOp) -> UOp:
+  tsink = graph_rewrite(sink, pm_lil_prepare_graph, bottom_up=True, name="prepare graph")
+  # Calls can only receive buffer states. Materialize lazy constants/computations instead of silently unwrapping them to a nonexistent base buffer.
+  tsink = graph_rewrite(tsink, pm_materialize_call_args, name="materialize call args")
+
+  read_cache:dict[UOp, set[UOp]] = {}
+  def read_buffers(x:UOp) -> set[UOp]:
+    if x not in read_cache:
+      read_cache[x] = {x.buf_uop} if x.has_buffer_identity(after_ok=True) else set().union(*(read_buffers(s) for s in x.src))
+    return read_cache[x]
+  stores = [u for u in tsink.toposort() if u.op is Ops.STORE and not u.src[0].is_variable]
+  dests = [u.src[0].buf_uop for u in stores]
+  reads = [read_buffers(u.src[1]) for u in stores]
+  force_stage:set[UOp] = set()
+  param_writes = [i for i,dest in enumerate(dests) if dest.op is Ops.PARAM]
+  if len(param_writes) == 2:
+    i, j = param_writes
+    if stores[j].src[1].op_in_backward_slice_with_self(Ops.REDUCE) and dests[i] is not dests[j] and \
+       dests[i] in reads[j] and dests[j] in reads[i]: force_stage.add(stores[j].src[1])
+
+  # add safe STAGEs to never duplicate compute
+  # we compute the number of times a buffer is consumed. if > 1, we realize
+  realize = {}
+  consumes = {tsink:0}
+  for u in reversed(tsink.toposort()):
+    assert u in consumes, f"{u.op} not in consumes"
+    if u in force_stage:
+      realize[u] = u.rtag(1).bufferize(arg=BufferizeOpts(device=u.device, removable=False))
+      consumes[u] = 1
+    elif (u.op in GroupOp.ALU or u.op is Ops.REDUCE) and consumes[u] > 1 and u.device is not None:
+      # TODO: rename to stage
+      realize[u] = u.rtag(1).bufferize(arg=BufferizeOpts(device=u.device))
+      consumes[u] = 1
+    if u.op is Ops.STORE: consumes[u] = 1
+    if u.op is Ops.EXPAND: consumes[u] *= u.max_numel() // u.src[0].max_numel()
+    for i,s in enumerate(u.src):
+      if s not in consumes: consumes[s] = 0
+      if u.op is not Ops.STORE or i > 0:
+        consumes[s] += consumes[u]
+  if VIZ:
+    with Context(TRACK_MATCH_STATS=0): ctags = graph_rewrite(tsink, debug_tag_factor, ctx=(consumes, realize), bottom_up=True)
+    graph_rewrite(ctags, PatternMatcher([]), name="View Consumes")
+
+  # add stages
+  tsink = graph_rewrite(tsink.substitute(realize), remove_all_tags, name="untag")
+
+  # simple rangeify
+  tsink = graph_rewrite(tsink, pm_range_creation+pm_range_migration, ctx=itertools.count(0), bottom_up=True, name="simple rangeify")
+
+  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Rangeify")
+
+  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify, name="pre-fusion reduce simplify")
+
+  # remove stage boundaries when this doesn't duplicate expensive compute or nest reductions
+  while 1:
+    staged:dict[UOp, list[UOp]] = {}
+    children:dict[UOp, list[UOp]] = {}
+    for u in tsink.toposort():
+      for s in u.src: children.setdefault(s, []).append(u)
+      if u.op is Ops.INDEX and u.src[0].op is Ops.STAGE: staged.setdefault(u.src[0], []).append(u)
+
+    boundary_cache:dict[UOp, set[UOp]] = {}
+    def boundaries(x:UOp) -> set[UOp]:
+      if x not in boundary_cache:
+        boundary_cache[x] = set().union(*({c} if c.op in {Ops.STAGE, Ops.STORE, Ops.CALL} else boundaries(c)
+                                          for c in children.get(x, [])))
+      return boundary_cache[x]
+
+    reduce_cache:dict[UOp, bool] = {}
+    def feeds_reduce(x:UOp) -> bool:
+      if x not in reduce_cache:
+        reduce_cache[x] = any(c.op is Ops.REDUCE or (c.op not in {Ops.STAGE, Ops.STORE, Ops.CALL} and feeds_reduce(c))
+                              for c in children.get(x, []))
+      return reduce_cache[x]
+
+    def boundary_work(boundary:UOp) -> int|None:
+      value = boundary.src[1] if boundary.op is Ops.STORE else boundary.src[0]
+      if any(r.src[0].op is not Ops.CONST or not isinstance(r.src[0].val, int) for r in value.ranges): return None
+      return prod(r.src[0].val for r in value.ranges)
+
+    def recompute_work(idxs:list[UOp]) -> int|None:
+      works = [boundary_work(next(iter(bs))) for idx in idxs if len(bs:=boundaries(idx)) == 1]
+      return sum(x for x in works if x is not None) if len(works) == len(idxs) and all(x is not None for x in works) else None
+
+    replacements:dict[UOp, UOp] = {}
+    range_replacements:dict[UOp, UOp] = {}
+    selected_stages:set[UOp] = set()
+    for stage,idxs in staged.items():
+      if not stage.arg.removable or children.get(stage) != idxs: continue
+      inlinable = stage.src[0].op in GroupOp.ALU or stage.src[0].op is Ops.REDUCE
+      cost = recompute_cost(stage.src[0])
+
+      # passthrough stages don't duplicate compute when indexing is unchanged
+      if not inlinable:
+        if len(idxs) == 1 and stage.src[1:] == idxs[0].src[1:]:
+          replacements[idxs[0]] = stage.src[0]
+          break
+        continue
+
+      # duplicate cheap elementwise stages; reductions require identical indexing into one output boundary
+      if len(idxs) > 1:
+        work = recompute_work(idxs)
+        if cost is not None and cost <= MAX_RECOMPUTE and work is not None and work <= stage.max_numel() * len(idxs):
+          replacements.update((idx, inline_stage_index(stage, idx)) for idx in idxs)
+        elif cost is None and all(idx.src[1:] == idxs[0].src[1:] for idx in idxs) and not any(feeds_reduce(idx) for idx in idxs):
+          stage_boundaries = set().union(*(boundaries(idx) for idx in idxs))
+          if len(stage_boundaries) == 1 and boundary_work(next(iter(stage_boundaries))) == stage.max_numel():
+            range_replacements.update(zip(stage.src[1:], idxs[0].src[1:]))
+            selected_stages.add(stage)
+        if replacements or range_replacements: break
+        continue
+
+      idx = idxs[0]
+      stage_boundaries = boundaries(idx)
+      work = boundary_work(next(iter(stage_boundaries))) if len(stage_boundaries) == 1 else None
+      scalar = stage.max_numel() == 1 and cost is not None and cost <= MAX_SCALAR_RECOMPUTE and all(r.op is Ops.CONST for r in idx.src[1:])
+      small = cost is not None and cost <= MAX_SCALAR_RECOMPUTE and stage.max_numel() <= 8 and idx.max_numel() <= 8
+      if cost is not None:
+        if stage.src[0].op_in_backward_slice_with_self(Ops.THREEFRY) or scalar or small or \
+           (cost <= MAX_RECOMPUTE and work is not None and work <= stage.max_numel()):
+          replacements[idx] = inline_stage_index(stage, idx)
+      elif len(stage_boundaries) == 1 and not feeds_reduce(idx) and work == stage.max_numel():
+        if all(r.op is Ops.RANGE for r in idx.src[1:]):
+          range_replacements.update(zip(stage.src[1:], idx.src[1:]))
+          selected_stages.add(stage)
+        else: replacements[idx] = inline_stage_index(stage, idx)
+      if replacements or range_replacements: break
+
+    if not replacements and not range_replacements: break
+    tsink = tsink.substitute(replacements).substitute(range_replacements)
+    if selected_stages:
+      selected_stages = {stage.substitute(range_replacements) for stage in selected_stages}
+      tsink = graph_rewrite(tsink, pm_remove_selected_stage, ctx=selected_stages, name="remove selected stage")
+
+  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify, name="reduce simplify")
+
+  # ***** MERGING AND SPLITTING (should be totally optional) *****
+
+  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Merged Rangeify")
+
+  next_buffer_num = itertools.count(1000)
+  tsink = graph_rewrite(tsink, symbolic+pm_remove_stage, ctx=next_buffer_num, bottom_up=True, name="remove stage")
+  tsink = graph_rewrite(tsink, split_kernels, bottom_up=True, name="split kernels")
+  tsink = graph_rewrite(tsink, pm_no_indexing_calls, name="remove indexing from call args")
+
+  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
+  if SPEC:
+    # validate the kernel graph
+    from tinygrad.uop.spec import type_verify, spec_kernel_graph
+    type_verify(tsink, spec_kernel_graph, enter_calls=False)
+  return tsink
+

tinygrad/tensor.py

@@ -114,7 +114,8 @@ def transform_precompiled_call(c:UOp) -> UOp|None:
   # add the outputs to the call
   srcs = c.src[0].src
   resolved = [c.gettuple(i) for i in range(len(srcs))]
-  outs = tuple(r.empty_like() for r in resolved)
+  # CALL outputs are max-sized physical buffers. Keep symbolic shapes as views so writable arguments and returned buffers stay identical.
+  outs = tuple(r.pad_to(r.max_shape).empty_like() for r in resolved)
   targets = [o.param_like(len(c.src)-1+i).shrink_to(s.shape) for i,(o,s) in enumerate(zip(outs, srcs))]
 
   subs:dict[UOp, UOp] = {}
@@ -133,11 +134,8 @@ def transform_precompiled_call(c:UOp) -> UOp|None:
 
   # body switches from TUPLE to SINK, so the node becomes an opaque CALL (not FUNCTION)
   new_call = UOp(Ops.CALL, src=(fxn, *input_buffers, *outs), arg=c.arg)
-  rets = tuple(o.after(new_call) for o in outs)
-
-  # if the CALL has symbolic shapes, shrink the max-sized output to the actual symbolic shape
-  # NOTE: must use resolved shapes from the FUNCTION (which substitutes PARAMs with external args), not raw body shapes
-  rets = tuple(r.shrink_to(rs.shape) for r,rs in zip(rets, resolved))
+  # NOTE: use resolved shapes from the FUNCTION (which substitutes PARAMs with external args), not raw body shapes.
+  rets = tuple(o.after(new_call).shrink_to(rs.shape) for o,rs in zip(outs, resolved))
 
   return UOp.maketuple(*rets)
 
@@ -220,6 +218,32 @@ pm_replace_buf = PatternMatcher([
   (UPat(Ops.AFTER, name="b"), lambda ctx,b: replace_input_buffer(ctx, b) if b.is_bound_var else None),
 ])
 
+def _check_state_cycles(sink:UOp):
+  # Track only the first storage state on each path. Combining a raw read with an assigned state is stale, while distinct AFTERs can be valid
+  # independent snapshots (for example RNG state). Memoizing the two state sets avoids repeatedly walking large training graphs.
+  states:dict[UOp, tuple[frozenset[UOp], frozenset[UOp]]] = {}
+  state:tuple[frozenset[UOp], frozenset[UOp]]
+  for u in sink.toposort(enter_calls=False):
+    if u.op is Ops.BUFFER and u.addrspace == AddrSpace.GLOBAL: state = (frozenset((u,)), frozenset())
+    elif u.op is Ops.AFTER and u.addrspace == AddrSpace.GLOBAL:
+      key = u.buf_uop
+      stores = [x for x in u.src[1:] if x.op is Ops.STORE and x.src[0].buf_uop is key]
+      # Ordering dependencies can STORE to another buffer without changing this state. Self-dependent updates continue the existing state lineage;
+      # only a write independent of the old value creates a conflicting state.
+      self_update = any(key in states[x.src[1]][0] or key in states[x.src[1]][1] for x in stores)
+      state = states[u.src[0]] if not stores or self_update else (frozenset(), frozenset((key,)))
+    else:
+      srcs = u.src[1:] if u.op in {Ops.CALL, Ops.FUNCTION} else u.src
+      raw = frozenset().union(*(states[x][0] for x in srcs))
+      assigned = frozenset().union(*(states[x][1] for x in srcs))
+      # CONTIGUOUS snapshots stale raw reads before a pending assignment, but cannot make a post-assignment read happen earlier.
+      state = (frozenset() if u.op is Ops.CONTIGUOUS else raw, assigned)
+    states[u] = state
+    if u.op in GroupOp.ALU:
+      branches = [states[x] for x in u.src]
+      if any((a[0] & b[1]) or (a[1] & b[0]) for i,a in enumerate(branches) for b in branches[i+1:]):
+        raise RuntimeError("cycle detected while combining buffer states")
+
 @rewrite_group(lambda _,ret: f"Callify {pluralize('Buffer', len(ret[1]))}")
 def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
   if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Tensor Graph")
@@ -406,7 +430,9 @@ class Tensor(RandMixin):
     # weakness ends where storage begins
     if any(t.dtype in dtypes.weaks and t.uop.device is not None for t in (self,)+lst):
       raise RuntimeError("cannot realize a weak dtype; cast to a concrete dtype first")
-    big_sink, becomes_map = transform_to_call(UOp.sink(*[x.uop for x in (self,)+lst]))
+    big_sink = UOp.sink(*[x.uop for x in (self,)+lst])
+    _check_state_cycles(big_sink)
+    big_sink, becomes_map = transform_to_call(big_sink)
     _apply_map_to_tensors(becomes_map, name="buffers")
     return create_linear_with_vars(big_sink)
 

tinygrad/uop/ops.py

@@ -325,8 +325,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   @functools.cached_property
   def tuplize(self:UOp) -> tuple:
-    # 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])
+    return (self.op.value, self.arg, self.dtype,)+tuple([x.tuplize for x in self.src])
 
   # *** uop shape stuff ***
 
@@ -1174,12 +1173,13 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     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))
   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}"))
-    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)
+    # if it's a PARAM or BUFFER, we just replace the slot
+    buf = self
+    while buf.op is Ops.AFTER: buf = buf.src[0]
+    if buf.op in {Ops.PARAM, Ops.BUFFER}: return UOp(Ops.PARAM, src=buf.src, arg=replace(buf.arg, slot=slot))
+    # otherwise we create a new param
+    addrspace = buf.addrspace if buf.addrspace is not None else AddrSpace.GLOBAL
+    return UOp.param(slot, buf.dtype, buf.shard_shape if buf.axis is not None else buf._shape, self.device, addrspace=addrspace, axis=buf.axis)
 
   @staticmethod
   def custom_function(name:str, *src:UOp) -> UOp: return UOp(Ops.CUSTOM_FUNCTION, src=src, arg=name)

tinygrad/uop/spec.py

@@ -257,12 +257,18 @@ spec_kernel_graph = PatternMatcher([
   # const + stack to make vconsts and shape args. a 0-size/bound reduce keeps its const casted
   (UPat(Ops.CONST, src=()), lambda: True),
   (UPat(Ops.CAST, src=(UPat(Ops.CONST, src=()),)), lambda: True),
+  # symbolic shape expressions can remain outside kernels (for example, flattening (n, 4) produces n*4)
+  (UPat(GroupOp.ALU, name="x"), lambda x: x.dtype in dtypes.ints+(dtypes.weakint,) or None),
   (UPat(Ops.STACK, name="s"), lambda s: all(x.op in (Ops.CONST, Ops.PARAM) or x.is_variable or x.is_bound_var for x in s.src) or None),
   # linear for more kernels (TODO: we should enter non sink calls)
   #(UPat(Ops.LINEAR), lambda: True),
   # param is outside buffer, buffer is local buffer
   (UPat(Ops.PARAM, name="x"), lambda x: isinstance(x.arg, ParamArg)),
   (UPat(Ops.BUFFER, name="x"), lambda x: isinstance(x.arg, ParamArg) and x.addrspace in (AddrSpace.GLOBAL, AddrSpace.ALU)),
+  # indexing/movement views on kernel buffers/call results are allowed to carry symbolic shape expressions
+  (UPat({Ops.INDEX}|GroupOp.Movement,
+        src=(UPat({Ops.INDEX}|GroupOp.Movement|{Ops.PARAM, Ops.AFTER, Ops.BUFFER, Ops.MSTACK, Ops.MSELECT, Ops.BITCAST}),),
+        allow_any_len=True), lambda: True),
   # RESHAPE/BITCAST are NOOPs in the kernel graph (do we need them?)
   (UPat((Ops.RESHAPE, Ops.BITCAST)), lambda: True),
   # mstack/mselect