openpilot: no NOLOCALS=1 (kimi) (#17762)

.github/workflows/benchmark.yml

@@ -458,31 +458,27 @@ jobs:
           - version: '0.11.0'
             model: vision
             url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
-            timing: 17
+            timing: 18
           - 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
+            timing: 3.4
           - version: '0.11.0'
             model: dmonitoring
             url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
-            timing: 11
+            timing: 12.5
           - version: '0.11.2'
             model: supercombo
             url: https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
-            timing: 26
+            timing: 28
           - dev: QCOM:IR3
             version: '0.11.2'
             model: supercombo
-            timing: 41
+            timing: 29
           - 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
+            timing: 12.5
       fail-fast: false
     name: openpilot ${{ matrix.version }} compile3 ${{ matrix.model }} (DEV=${{ matrix.dev }})
     runs-on: [self-hosted, Linux, comma]
@@ -506,7 +502,7 @@ jobs:
     - 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 }}
+      run: FLOAT16=1 IMAGE=1 taskset -c 4-7 python3 examples/openpilot/compile3.py ${{ matrix.url }}
     - name: run pickle
       run: BENCHMARK_LOG="${BENCHMARK_LOG}_run_pickle" RUN_PICKLE=1 taskset -c 4-7 python3 examples/openpilot/compile3.py
     - name: Run process replay tests

tinygrad/codegen/opt/heuristic.py

@@ -78,7 +78,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
             return k
 
   # are we grouping? (requires local shape support)
-  if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= (240 if NOLOCALS else 2048), False):
+  if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= (240 if NOLOCALS or k.ren.target.device == "QCOM" else 2048), False):
     for axis, sz in itertools.product((0, 1, 2), (16,)):
       try:
         k.apply_opt(Opt(OptOps.GROUPTOP, axis, sz))
@@ -162,6 +162,19 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
   if k.ren.has_local:
     if NOLOCALS:
       k.apply_opt(Opt(OptOps.NOLOCALS))
+    elif k.ren.target.device == "QCOM":
+      # for openpilot: use 32..128 threads per workgroup, at most 8 on the innermost axis
+      # apply innermost global axes first so the leading hardware local dims hold the trailing global axes, like gidx
+      workgroup = 1
+      opts: list[tuple[int, int]] = []
+      for axis in [a for a in k.axes_of(AxisType.GLOBAL, AxisType.WEAK) if k.rngs[a].src[0].op is Ops.CONST][-3:][::-1]:
+        if (sz:=max(x for x in range(1, min(int(k.full_shape[axis]), 128 // workgroup if opts else 8) + 1) if int(k.full_shape[axis]) % x == 0)) > 1:
+          opts.append((axis, sz))
+          workgroup *= sz
+      if opts and workgroup < 32:  # fill at least one wave: grow the innermost local as much as possible
+        axis, sz = opts[0]
+        opts[0] = axis, max(x for x in range(1, min(int(k.full_shape[axis]), 128 * sz // workgroup) + 1) if int(k.full_shape[axis]) % x == 0)
+      for axis, sz in opts: k.apply_opt(Opt(OptOps.LOCAL, axis, sz))
     else:
       # prioritize making expand axes local
       local_axis_ranking = [(any(k.rngs[axis] not in b.src[1].get_idx().backward_slice for b in k.bufs), axis) \