work

the two remaining slow ones

.coveragerc

@@ -0,0 +1,3 @@
+[run]
+source = tinygrad
+branch = True

.github/workflows/benchmark.yml

@@ -630,17 +630,17 @@ jobs:
     - name: openpilot compile3 0.9.9 dmonitoring
       run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
     - name: openpilot compile3 0.10.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=5 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
     - name: openpilot compile3 0.10.0 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=13 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
     - name: openpilot compile3 0.10.1 driving_vision
       # TODO: ASSERT_MIN_STEP_TIME=17
-      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=25 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=21 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
     - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=5 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
     - name: openpilot compile3 0.10.1 dmonitoring
       # TODO: ASSERT_MIN_STEP_TIME=10
-      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=13 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
     - name: benchmark MobileNetV2 on DSP
       run: |
         # generate quantized weights

.github/workflows/test.yml

@@ -290,8 +290,8 @@ jobs:
         python extra/optimization/extract_dataset.py
         gzip -c /tmp/sops > extra/datasets/sops.gz
         #DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
-    - name: Repo line count < 18000 lines
-      run: MAX_LINE_COUNT=18000 python sz.py
+    - name: Repo line count < 18500 lines
+      run: MAX_LINE_COUNT=18500 python sz.py
 
   spec:
     strategy:

.gitignore

@@ -63,3 +63,5 @@ profile_stats
 *.log
 target
 .mypy_cache
+mutants
+.mutmut-cache

autogen_stubs.sh

@@ -188,6 +188,7 @@ nv_status_codes = {}
     extra/nv_gpu_driver/g_rpc-message-header.h \
     extra/nv_gpu_driver/gsp_static_config.h \
     extra/nv_gpu_driver/vbios.h \
+    extra/nv_gpu_driver/pci_exp_table.h \
     --clang-args="-DRPC_MESSAGE_STRUCTURES -DRPC_STRUCTURES -include $NVKERN_SRC/src/common/sdk/nvidia/inc/nvtypes.h -I$NVKERN_SRC/src/nvidia/generated -I$NVKERN_SRC/src/common/inc -I$NVKERN_SRC/src/nvidia/inc -I$NVKERN_SRC/src/nvidia/interface/ -I$NVKERN_SRC/src/nvidia/inc/kernel -I$NVKERN_SRC/src/nvidia/inc/libraries -I$NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc -I$NVKERN_SRC/kernel-open/nvidia-uvm -I$NVKERN_SRC/kernel-open/common/inc -I$NVKERN_SRC/src/common/sdk/nvidia/inc -I$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include -I$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl" \
     -o $BASE/nv/nv.py
 
@@ -549,7 +550,6 @@ elif [ "$1" == "kfd" ]; then generate_kfd
 elif [ "$1" == "nv" ]; then generate_nv
 elif [ "$1" == "amd" ]; then generate_amd
 elif [ "$1" == "am" ]; then generate_am
-elif [ "$1" == "nvdrv" ]; then generate_nvdrv
 elif [ "$1" == "sqtt" ]; then generate_sqtt
 elif [ "$1" == "qcom" ]; then generate_qcom
 elif [ "$1" == "io_uring" ]; then generate_io_uring

examples/openpilot/compile3.py

@@ -4,8 +4,6 @@ import numpy as np
 from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device, dtypes
 from tinygrad.helpers import DEBUG, getenv
 from tinygrad.engine.realize import CompiledRunner
-
-import onnx
 from tinygrad.nn.onnx import OnnxRunner
 
 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
@@ -40,7 +38,7 @@ def compile(onnx_file):
   np.testing.assert_equal(test_val, ret, "JIT run failed")
   print("jit run validated")
 
-  # checks from compile2
+  # check gated read_image usage
   kernel_count = 0
   read_image_count = 0
   gated_read_image_count = 0
@@ -96,6 +94,7 @@ def test_vs_compile(run, inputs, test_val=None):
   return val
 
 def test_vs_onnx(new_inputs, test_val, onnx_file, tol):
+  import onnx
   import onnxruntime as ort
   
   onnx_inputs = {k:v.numpy() for k,v in new_inputs.items()}

extra/models/llama.py

@@ -89,6 +89,20 @@ class Attention:
     keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
     xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
     attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2)
+    if getenv("STUB_ATTENTION"):
+      # TODO: do we need mask?
+      from tinygrad.uop.ops import UOp, KernelInfo
+      def fa_custom_forward(attn:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
+        return UOp.sink(arg=KernelInfo(name="fa_custom_forward"))
+      def fa_custom_backward(out_q:UOp, out_k:UOp, out_v:UOp, grad:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
+        return UOp.sink(arg=KernelInfo(name="fa_custom_backward"))
+      def fa_backward(grad:UOp, kernel:UOp) -> tuple[None, UOp, UOp, UOp]:
+        grad_q = Tensor.empty_like(q:=Tensor(kernel.src[1]))
+        grad_k = Tensor.empty_like(k:=Tensor(kernel.src[2]))
+        grad_v = Tensor.empty_like(v:=Tensor(kernel.src[3]))
+        ck = Tensor.custom_kernel(grad_q, grad_k, grad_v, Tensor(grad), q, k, v, fxn=fa_custom_backward)[:3]
+        return (None, ck[0].uop, ck[1].uop, ck[2].uop)
+      attn = Tensor.empty_like(attn).custom_kernel(xq, keys, values, fxn=fa_custom_forward, grad_fxn=fa_backward)[0]
     attn = attn.reshape(bsz, seqlen, -1)
     return self.wo(attn)
 

extra/nv_gpu_driver/pci_exp_table.h

@@ -0,0 +1,134 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#ifndef PCIEXPTBL_H
+#define PCIEXPTBL_H
+
+#define NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_BASE   0x00
+#define NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_EXT    0xE0
+
+//
+// The VBIOS object comes from walking the PCI expansion code block
+// The following structure holds the expansion code format.
+//
+#define PCI_EXP_ROM_SIGNATURE     0xaa55
+#define PCI_EXP_ROM_SIGNATURE_NV  0x4e56 // "VN" in word format
+#define PCI_EXP_ROM_SIGNATURE_NV2 0xbb77
+#define IS_VALID_PCI_ROM_SIG(sig) ((sig == PCI_EXP_ROM_SIGNATURE) ||    \
+                                   (sig == PCI_EXP_ROM_SIGNATURE_NV) || \
+                                   (sig == PCI_EXP_ROM_SIGNATURE_NV2))
+
+#define OFFSETOF_PCI_EXP_ROM_SIG                    0x0
+#define OFFSETOF_PCI_EXP_ROM_NBSI_DATA_OFFSET       0x16
+#define OFFSETOF_PCI_EXP_ROM_PCI_DATA_STRUCT_PTR    0x18
+
+#pragma pack(1)
+typedef struct _PCI_EXP_ROM_STANDARD
+{
+    NvU16       sig;                //  00h: ROM Signature 0xaa55
+    NvU8        reserved [0x16];    //  02h: Reserved (processor architecture unique data)
+    NvU16       pciDataStrucPtr;    //  18h: Pointer to PCI Data Structure
+    NvU32       sizeOfBlock;        //  1Ah: <NBSI-specific appendage>
+} PCI_EXP_ROM_STANDARD, *PPCI_EXP_ROM_STANDARD;
+#pragma pack()
+
+#pragma pack(1)
+typedef struct _PCI_EXP_ROM_NBSI
+{
+    NvU16       sig;                //  00h: ROM Signature 0xaa55
+    NvU8        reserved [0x14];    //  02h: Reserved (processor architecture unique data)
+    NvU16       nbsiDataOffset;     //  16h: Offset from header to NBSI image
+    NvU16       pciDataStrucPtr;    //  18h: Pointer to PCI Data Structure
+    NvU32       sizeOfBlock;        //  1Ah: <NBSI-specific appendage>
+} PCI_EXP_ROM_NBSI, *PPCI_EXP_ROM_NBSI;
+#pragma pack()
+
+typedef union _PCI_EXP_ROM {
+    PCI_EXP_ROM_STANDARD standard;
+    PCI_EXP_ROM_NBSI nbsi;
+} PCI_EXP_ROM, *PPCI_EXP_ROM;
+
+#define PCI_DATA_STRUCT_SIGNATURE     0x52494350 // "PCIR" in dword format
+#define PCI_DATA_STRUCT_SIGNATURE_NV  0x5344504E // "NPDS" in dword format
+#define PCI_DATA_STRUCT_SIGNATURE_NV2 0x53494752 // "RGIS" in dword format
+#define IS_VALID_PCI_DATA_SIG(sig) ((sig == PCI_DATA_STRUCT_SIGNATURE) ||    \
+                                    (sig == PCI_DATA_STRUCT_SIGNATURE_NV) || \
+                                    (sig == PCI_DATA_STRUCT_SIGNATURE_NV2))
+
+#define PCI_LAST_IMAGE NVBIT(7)
+#define PCI_ROM_IMAGE_BLOCK_SIZE 512U
+
+#define OFFSETOF_PCI_DATA_STRUCT_SIG        0x0
+#define OFFSETOF_PCI_DATA_STRUCT_VENDOR_ID  0x4
+#define OFFSETOF_PCI_DATA_STRUCT_LEN        0xa
+#define OFFSETOF_PCI_DATA_STRUCT_CLASS_CODE 0xd
+#define OFFSETOF_PCI_DATA_STRUCT_CODE_TYPE  0x14
+#define OFFSETOF_PCI_DATA_STRUCT_IMAGE_LEN  0x10
+#define OFFSETOF_PCI_DATA_STRUCT_LAST_IMAGE 0x15
+
+#pragma pack(1)
+typedef struct _PCI_DATA_STRUCT
+{
+    NvU32       sig;                //  00h: Signature, the string "PCIR" or NVIDIA's alternate "NPDS"
+    NvU16       vendorID;           //  04h: Vendor Identification
+    NvU16       deviceID;           //  06h: Device Identification
+    NvU16       deviceListPtr;      //  08h: Device List Pointer
+    NvU16       pciDataStructLen;   //  0Ah: PCI Data Structure Length
+    NvU8        pciDataStructRev;   //  0Ch: PCI Data Structure Revision
+    NvU8        classCode[3];       //  0Dh: Class Code
+    NvU16       imageLen;           //  10h: Image Length (units of 512 bytes)
+    NvU16       vendorRomRev;       //  12h: Revision Level of the Vendor's ROM
+    NvU8        codeType;           //  14h: holds NBSI_OBJ_CODE_TYPE (0x70) and others
+    NvU8        lastImage;          //  15h: Last Image Indicator: bit7=1 is lastImage
+    NvU16       maxRunTimeImageLen; //  16h: Maximum Run-time Image Length (units of 512 bytes)
+} PCI_DATA_STRUCT, *PPCI_DATA_STRUCT;
+#pragma pack()
+
+#define NV_PCI_DATA_EXT_SIG 0x4544504E // "NPDE" in dword format
+#define NV_PCI_DATA_EXT_REV_10 0x100      // 1.0
+#define NV_PCI_DATA_EXT_REV_11 0x101      // 1.1
+
+#define OFFSETOF_PCI_DATA_EXT_STRUCT_SIG            0x0
+#define OFFSETOF_PCI_DATA_EXT_STRUCT_LEN            0x6
+#define OFFSETOF_PCI_DATA_EXT_STRUCT_REV            0x4
+#define OFFSETOF_PCI_DATA_EXT_STRUCT_SUBIMAGE_LEN   0x8
+#define OFFSETOF_PCI_DATA_EXT_STRUCT_LAST_IMAGE     0xa
+#define OFFSETOF_PCI_DATA_EXT_STRUCT_FLAGS          0xb
+
+#define PCI_DATA_EXT_STRUCT_FLAGS_CHECKSUM_DISABLED 0x04
+
+#pragma pack(1)
+typedef struct _NV_PCI_DATA_EXT_STRUCT
+{
+    NvU32   signature;          //  00h: Signature, the string "NPDE"
+    NvU16   nvPciDataExtRev;    //  04h: NVIDIA PCI Data Extension Revision
+    NvU16   nvPciDataExtLen;    //  06h: NVIDIA PCI Data Extension Length
+    NvU16   subimageLen;        //  08h: Sub-image Length
+    NvU8    privLastImage;      //  0Ah: Private Last Image Indicator
+    NvU8    flags;              //  0Bh: Private images enabled if bit0=1
+} NV_PCI_DATA_EXT_STRUCT, *PNV_PCI_DATA_EXT_STRUCT;
+#pragma pack()
+
+#endif // PCIEXPTBL_H
+
+

extra/sqtt/rgptool.py

@@ -4,7 +4,7 @@ import argparse, ctypes, struct, hashlib, pickle, code, typing, functools
 import tinygrad.runtime.autogen.sqtt as sqtt
 from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent
-from tinygrad.helpers import round_up, flatten, all_same
+from tinygrad.helpers import round_up, flatten, all_same, temp
 from dataclasses import dataclass
 
 CHUNK_CLASSES = {
@@ -210,7 +210,7 @@ class RGP:
         flags=0,
         trace_shader_core_clock=0x93f05080,
         trace_memory_clock=0x4a723a40,
-        device_id={110000: 0x744c, 110003: 0x7480, 120001: 0x7550}[device_props['gfx_target_version']],
+        device_id={110000: 0x744c, 110003: 0x7480, 120001: 0x7550, 120000: 0x7550}[device_props['gfx_target_version']],
         device_revision_id=0xc8,
         vgprs_per_simd=1536,
         sgprs_per_simd=128*16,
@@ -324,7 +324,7 @@ class RGP:
 if __name__ == '__main__':
   parser = argparse.ArgumentParser(prog='rgptool', description='A tool to create (from pickled tinygrad profile), inspect and modify Radeon GPU Profiler files')
   parser.add_argument('command')
-  parser.add_argument('input')
+  parser.add_argument('input', nargs='?', default=temp("profile.pkl", append_user=True))
   parser.add_argument('-d', '--device')
   parser.add_argument('-o', '--output')
   args = parser.parse_args()
@@ -346,3 +346,4 @@ if __name__ == '__main__':
 
   if args.output is not None:
     with open(args.output, 'wb+') as fd: fd.write(rgp.to_bytes())
+    print(f"Saved to {args.output}")

extra/sqtt/roc.py

@@ -56,12 +56,19 @@ class PrgExec:
   simd:int
   def __str__(self): return f"{self.name},{self.wave},{self.cu},{self.simd}"
 
+@dataclasses.dataclass(frozen=True)
+class WaveExec:
+  wave_id:int
+  cu:int
+  simd:int
+  insts:list[InstExec]
+
 class _ROCParseCtx:
   def __init__(self, dev_evs:dict[str, ProfileDeviceEvent], sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
     self.dev_evs, self.sqtt_evs, self.prog_evs = dev_evs, iter(sqtt_evs), prog_evs
     self.wave_events:dict[PrgExec, dict[int, InstInfo]] = {}
-    self.disasms:dict[int, tuple[str, int]] = {}
-    self.inst_execs:dict[PrgExec, list[InstExec]] = {}
+    self.disasms:dict[tuple[str, int], tuple[str, int]] = {}
+    self.inst_execs:dict[str, list[WaveExec]] = {}
 
     for prog in prog_evs:
       arch = "gfx%d%x%x" % ((trgt:=unwrap(dev_evs[prog.device].props)['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
@@ -85,14 +92,14 @@ class _ROCParseCtx:
     for j in range(ev.instructions_size):
       inst_ev = ev.instructions_array[j]
       inst_typ = rocprof.rocprofiler_thread_trace_decoder_inst_category_t__enumvalues[inst_ev.category]
-      inst_disasm = self.disasms[(self.active_kern, inst_ev.pc.address)][0]
+      inst_disasm = self.disasms[(unwrap(self.active_kern), unwrap(inst_ev.pc.address))][0]
       asm.setdefault(inst_ev.pc.address, InstInfo(typ=inst_typ, inst=inst_disasm))
       asm[inst_ev.pc.address].on_ev(inst_ev)
       inst_execs.append(InstExec(inst_typ, inst_disasm, inst_ev.stall, inst_ev.duration, inst_ev.time))
 
     if ev.instructions_size > 0:
-      self.wave_events[key:=PrgExec(self.active_kern, ev.wave_id, ev.cu, ev.simd)] = asm
-      self.inst_execs[key] = inst_execs
+      self.wave_events[key:=PrgExec(unwrap(self.active_kern), ev.wave_id, ev.cu, ev.simd)] = asm
+      self.inst_execs.setdefault(key.name, []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, inst_execs))
 
 def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
   dev_events:dict[str, ProfileDeviceEvent] = {}
@@ -125,7 +132,7 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
 
   @rocprof.rocprof_trace_decoder_isa_callback_t
   def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, data_ptr):
-    instr, mem_size_ptr[0] = ROCParseCtx.disasms[(ROCParseCtx.active_kern, pc.address)]
+    instr, mem_size_ptr[0] = ROCParseCtx.disasms[(unwrap(ROCParseCtx.active_kern), pc.address)]
 
     # this is the number of bytes to next instruction, set to 0 for end_pgm
     if instr == "s_endpgm": mem_size_ptr[0] = 0

extra/thunder/tiny/tk/__init__.py

@@ -0,0 +1 @@
+WARP_THREADS = 32

extra/thunder/tiny/tk/group.py

@@ -0,0 +1,272 @@
+import math, functools
+from typing import cast, Callable
+from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
+from tinygrad.uop.ops import AxisType, UOp, KernelInfo, Ops
+from tinygrad.engine.realize import ExecItem, get_runner
+from tinygrad.dtype import AddrSpace, PtrDType
+from tinygrad.helpers import getenv, prod
+
+from extra.thunder.tiny.tk import WARP_THREADS
+from extra.thunder.tiny.tk.tiles import TILE_ROW_DIM, TILE_COL_DIM, RT_BASE_TILE_NEPT, slots
+
+class Group:
+  def __init__(self, warps:int, ker):
+    self.warps = warps
+    self.group_threads = warps * WARP_THREADS
+    self.threadIdx_x = ker.threadIdx_x
+    self.ker = ker
+
+  # helpers
+  @property
+  def laneid(self): return self.threadIdx_x % self.group_threads
+  @property
+  def warpid(self): return self.laneid // WARP_THREADS
+  @property
+  def groupid(self): return self.threadIdx_x // self.group_threads
+
+  # ops that only work on a single warp
+
+  clear_rid = 1000
+  def clear(self, reg:UOp, value:float=0):
+    assert self.warps == 1
+
+    i = UOp.range(reg.size, Group.clear_rid)
+    Group.clear_rid += 1
+    return reg.reshape((reg.size,))[i].set(value, end=i).after(reg).reshape(reg.shape)
+
+  def zero(self, reg:UOp): return self.clear(reg, 0)
+  def neg_inf(self, reg:UOp): return self.clear(reg, -math.inf)
+
+  copy_rid = 300
+  def copy(self, dst:UOp, src:UOp):
+    assert self.warps == 1
+
+    assert dst.shape == src.shape
+    assert cast(PtrDType, dst.dtype).addrspace == AddrSpace.REG
+    assert cast(PtrDType, src.dtype).addrspace == AddrSpace.REG
+
+    rngs_for_shape = tuple(UOp.range(dim, Group.copy_rid + i) for i, dim in enumerate(dst.shape))
+    Group.copy_rid += len(dst.shape)
+
+    dst_store = dst[*rngs_for_shape].store(src[*rngs_for_shape].cast(dst.dtype.base)).end(*rngs_for_shape)
+
+    self.ker.push_store(dst_store, dst)
+    return dst.after(dst_store).reshape(dst.shape)
+
+  mma_rid = 600
+  def mma_AB(self, c:UOp, a:UOp, b:UOp, after=True):
+    assert self.warps == 1
+
+    mma_i_height = UOp.range(c.shape[-3], Group.mma_rid)
+    mma_i_width = UOp.range(c.shape[-2], Group.mma_rid+1)
+    mma_i_inner = UOp.range(a.shape[-2], Group.mma_rid+2, AxisType.REDUCE)
+    Group.mma_rid += 3
+
+    wmma_arg = ("WMMA_8_16_16_bfloat16_float", (8, 16, 16), dtypes.bfloat16, dtypes.float, "CUDA", 32, (((4, 2), (3, 2), (8, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
+
+    a_in = UOp.vectorize(*[a[mma_i_height, mma_i_inner, i] for i in range(8)])
+    b_in1 = UOp.vectorize(*([b[mma_i_inner, mma_i_width, i] for i in range(2)] + [b[mma_i_inner, mma_i_width, 4+i] for i in range(2)]))
+    c_out1 = UOp.vectorize(*[c[mma_i_height, mma_i_width, i] for i in range(4)])
+    b_in2 = UOp.vectorize(*([b[mma_i_inner, mma_i_width, 2+i] for i in range(2)] + [b[mma_i_inner, mma_i_width, 6+i] for i in range(2)]))
+    c_out2 = UOp.vectorize(*[c[mma_i_height, mma_i_width, 4+i] for i in range(4)])
+
+    out1 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in1, c_out1), arg=wmma_arg)
+    out2 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in2, c_out2), arg=wmma_arg)
+    c_i = [c[mma_i_height, mma_i_width, i].store(out1.gep(i)) for i in range(4)] + [c[mma_i_height, mma_i_width, 4+i].store(out2.gep(i)) for i in range(4)]
+    c_store = UOp.group(*c_i).end(mma_i_height, mma_i_width, mma_i_inner)
+
+    self.ker.push_store(c_store, c)
+    return c.after(c_store).reshape(c.shape) if after else c_store
+
+  def mma_ABt(self, c:UOp, a:UOp, b:UOp, after=True):
+    assert self.warps == 1
+
+    mma_i_height = UOp.range(c.shape[-3], Group.mma_rid)
+    mma_i_width = UOp.range(c.shape[-2], Group.mma_rid+1)
+    mma_i_inner = UOp.range(a.shape[-2], Group.mma_rid+2, AxisType.REDUCE)
+    Group.mma_rid += 3
+
+    wmma_arg = ("WMMA_8_16_16_bfloat16_float", (8, 16, 16), dtypes.bfloat16, dtypes.float, "CUDA", 32, (((4, 2), (3, 2), (8, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
+
+    a_in = UOp.vectorize(*[a[mma_i_height, mma_i_inner, i] for i in range(8)])
+    b_in1 = UOp.vectorize(*([b[mma_i_width, mma_i_inner, i] for i in range(2)] + [b[mma_i_width, mma_i_inner, 4+i] for i in range(2)]))
+    c_out1 = UOp.vectorize(*[c[mma_i_height, mma_i_width, i] for i in range(4)])
+    b_in2 = UOp.vectorize(*([b[mma_i_width, mma_i_inner, 2+i] for i in range(2)] + [b[mma_i_width, mma_i_inner, 6+i] for i in range(2)]))
+    c_out2 = UOp.vectorize(*[c[mma_i_height, mma_i_width, 4+i] for i in range(4)])
+
+    out1 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in1, c_out1), arg=wmma_arg)
+    out2 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in2, c_out2), arg=wmma_arg)
+    c_i = [c[mma_i_height, mma_i_width, i].store(out1.gep(i)) for i in range(4)] + [c[mma_i_height, mma_i_width, 4+i].store(out2.gep(i)) for i in range(4)]
+    c_store = UOp.group(*c_i).end(mma_i_height, mma_i_width, mma_i_inner)
+
+    self.ker.push_store(c_store, c)
+    return c.after(c_store).reshape(c.shape) if after else c_store
+
+  map_rid = 400
+  def map(self, a:UOp, op:Callable[[UOp], UOp]|Callable[[UOp, tuple], UOp]):
+    assert self.warps == 1
+
+    rngs_for_shape = tuple(UOp.range(dim, Group.map_rid + i) for i, dim in enumerate(a.shape))
+    Group.map_rid += len(a.shape)
+
+    if op.__code__.co_argcount == 1:
+      to_store = op(a[*rngs_for_shape])
+    else:
+      to_store = op(a[*rngs_for_shape], rngs_for_shape)
+
+    a_store = a[*rngs_for_shape].store(to_store).end(*rngs_for_shape)
+
+    self.ker.push_store(a_store, a)
+    return a.after(a_store).reshape(a.shape)
+
+  def row_reduce(self, vec:UOp, src:UOp, op:Callable[[UOp, UOp], UOp]):
+    assert self.warps == 1
+
+    red_local = UOp.placeholder((self.group_threads, 2), src.dtype.base, addrspace=AddrSpace.LOCAL, slot=slots.shared_slot)
+    slots.shared_slot += 1
+
+    for height in self.ker.range(src.shape[-3], track=False):
+      for i_outer in self.ker.range(2, track=False):
+        for width in self.ker.range(src.shape[-2], AxisType.REDUCE, track=False):
+          for i_inner in self.ker.range(4, AxisType.REDUCE, track=False):
+            elem_index = i_inner + 2 * (i_inner // 2) + i_outer * 2
+            vec_store = vec[height, 0, i_outer].store(op(vec[height, 0, i_outer], src[height, width, elem_index])).end(width, i_inner, i_outer)
+            vec = vec.after(vec_store).reshape(vec.shape)
+
+      # store to shared memory
+      for i_outer in self.ker.range(2, track=False):
+        red_local_store = red_local[self.laneid, i_outer].store(vec[height, 0, i_outer]).end(i_outer)
+        red_local = red_local.after(red_local_store).reshape(red_local.shape)
+
+      # reduce from shared memory
+      for i_outer in self.ker.range(2, track=False):
+        for i_inner in self.ker.range(3, AxisType.REDUCE, track=False):
+          offset = (self.laneid // 4) * 4 + ((self.laneid + 1 + i_inner) % 4)
+          vec_store = vec[height, 0, i_outer].store(op(vec[height, 0, i_outer], red_local[offset, i_outer])).end(i_inner, i_outer)
+
+    self.ker.push_store(vec_store, vec)
+    return vec.after(vec_store).reshape(vec.shape)
+
+  # ops that can work across multiple warps
+
+  LOAD_INNER = 8
+  load_rid = 100
+  def load(self, dst:UOp, src:UOp, dst_idxs:tuple[UOp|int,...]=(), idxs:tuple[UOp|int,...]=(), axis:int=0, transpose:bool=False):
+    assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
+    dst_dtype, src_dtype = cast(PtrDType, dst.dtype), cast(PtrDType, src.dtype)
+    if dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.LOCAL:
+      srcf = src.flatten(-2)
+
+      load_i_height = UOp.range(dst.shape[-3], Group.load_rid)
+      load_i_width = UOp.range(dst.shape[-2], Group.load_rid+1)
+      load_i_inner = UOp.range(RT_BASE_TILE_NEPT, Group.load_rid+2)
+      Group.load_rid += 3
+
+      if self.warps % 4 == 0: local_warpid = (self.warpid // 4) + (self.warpid % 4) * (self.warps // 4)
+      else: local_warpid = self.warpid
+      warp_laneid = self.threadIdx_x % WARP_THREADS
+
+      if not transpose:
+        row = (local_warpid * dst.shape[-3] + load_i_height) * TILE_ROW_DIM + (warp_laneid // 4)
+        col = load_i_width * TILE_COL_DIM + 2 * (warp_laneid % 4)
+
+        row_offset = ((load_i_inner % 4) // 2) * 8
+        col_offset = (load_i_inner % 2) + (load_i_inner // 4) * 8
+      else:
+        row = (local_warpid * dst.shape[-3] + load_i_height) * TILE_ROW_DIM + 2 * (warp_laneid % 4)
+        col = load_i_width * TILE_COL_DIM + (warp_laneid // 4)
+
+        row_offset = (load_i_inner % 2) + (load_i_inner // 4) * 8
+        col_offset = ((load_i_inner % 4) // 2) * 8
+
+      src_i_last = (row + row_offset) * src.shape[-1] + col + col_offset
+
+      dst_store = dst[*dst_idxs, load_i_height, load_i_width, load_i_inner].store(srcf[*idxs[:-2], src_i_last])
+      dst_store = dst_store.end(load_i_height, load_i_width, load_i_inner)
+    elif dst_dtype.addrspace == AddrSpace.LOCAL and src_dtype.addrspace == AddrSpace.GLOBAL:
+      dstf = dst.flatten(-2)
+
+      srcf = src.flatten()
+      row_stride = prod(src.shape[axis+1:])
+
+      idxs = tuple(idx * dst.shape[-2] if i == axis else idx for i, idx in enumerate(idxs))
+      idxs = tuple(idx * dst.shape[-1] if i == 3 else idx for i, idx in enumerate(idxs))
+      src_i = ((idxs[0] * src.shape[-3] + idxs[1]) * src.shape[-2] + idxs[2]) * src.shape[-1] + idxs[3]
+
+      memcpy_per_row = dst.shape[-1] // Group.LOAD_INNER
+      total_calls = prod(dst.shape[-2:]) // (self.group_threads * Group.LOAD_INNER)
+
+      load_i_outer = UOp.range(total_calls, Group.load_rid)
+      load_i_inner = UOp.range(Group.LOAD_INNER, Group.load_rid+1)
+      Group.load_rid += 2
+
+      load_idx = load_i_outer * self.group_threads + self.laneid
+      row = load_idx // memcpy_per_row
+      col = (load_idx * Group.LOAD_INNER) % dst.shape[-1]
+
+      dst_i = row * dst.shape[-1] + col + load_i_inner
+      src_i += row * row_stride + col + load_i_inner
+
+      dst_store = dstf[*dst_idxs, dst_i].store(srcf[src_i]).end(load_i_outer, load_i_inner)
+    else:
+      raise NotImplementedError(f"load from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented")
+
+    return dst.after(dst_store.barrier()).reshape(dst.shape)
+
+  STORE_INNER = 8
+  store_rid = 200
+  def store(self, dst:UOp, src:UOp, idxs:tuple[UOp|int,...]=(), src_idxs:tuple[UOp|int,...]=(), axis=0, after=True):
+    assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
+    dst_dtype, src_dtype = cast(PtrDType, dst.dtype), cast(PtrDType, src.dtype)
+    if src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.LOCAL:
+      dstf = dst.flatten(-2)
+
+      store_i_height = UOp.range(src.shape[-3], Group.store_rid)
+      store_i_width = UOp.range(src.shape[-2], Group.store_rid+1)
+      store_i_inner = UOp.range(RT_BASE_TILE_NEPT, Group.store_rid+2)
+      Group.store_rid += 3
+
+      if self.warps % 4 == 0: local_warpid = (self.warpid // 4) + (self.warpid % 4) * (self.warps // 4)
+      else: local_warpid = self.warpid
+      warp_laneid = self.threadIdx_x % WARP_THREADS
+
+      row = (local_warpid * src.shape[-3] + store_i_height) * TILE_ROW_DIM + (warp_laneid // 4)
+      col = store_i_width * TILE_COL_DIM + 2 * (warp_laneid % 4)
+
+      row_offset = ((store_i_inner % 4) // 2) * 8
+      col_offset = (store_i_inner % 2) + (store_i_inner // 4) * 8
+
+      dst_i_last = (row + row_offset) * dst.shape[-1] + col + col_offset
+
+      dst_store = dstf[*idxs[:-2], dst_i_last].store(src[*src_idxs, store_i_height, store_i_width, store_i_inner])
+      dst_store = dst_store.end(store_i_height, store_i_width, store_i_inner)
+    elif src_dtype.addrspace == AddrSpace.LOCAL and dst_dtype.addrspace == AddrSpace.GLOBAL:
+      dstf = dst.flatten()
+      row_stride = prod(dst.shape[axis+1:])
+
+      idxs = tuple(idx * src.shape[-2] if i == axis else idx for i, idx in enumerate(idxs))
+      idxs = tuple(idx * src.shape[-1] if i == 3 else idx for i, idx in enumerate(idxs))
+      dst_i = ((idxs[0] * dst.shape[-3] + idxs[1]) * dst.shape[-2] + idxs[2]) * dst.shape[-1] + idxs[3]
+
+      srcf = src.flatten(-2)
+
+      memcpy_per_row = src.shape[-1] // Group.STORE_INNER
+      total_calls = prod(src.shape[-2:]) // (self.group_threads * Group.STORE_INNER)
+
+      store_i_outer = UOp.range(total_calls, Group.store_rid)
+      store_i_inner = UOp.range(Group.STORE_INNER, Group.store_rid+1)
+      Group.store_rid += 2
+
+      load_idx = store_i_outer * self.group_threads + self.laneid
+      row = load_idx // memcpy_per_row
+      col = (load_idx * Group.STORE_INNER) % src.shape[-1]
+
+      src_i = row * src.shape[-1] + col + store_i_inner
+      dst_i += row * row_stride + col + store_i_inner
+
+      dst_store = dstf[dst_i].store(srcf[*src_idxs, src_i]).end(store_i_outer, store_i_inner)
+    else:
+      raise NotImplementedError(f"store from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented")
+
+    self.ker.push_store(dst_store, dst)
+    return dst.after(dst_store.barrier()).reshape(dst.shape) if after else dst_store

extra/thunder/tiny/tk/kernel.py

@@ -0,0 +1,57 @@
+from contextlib import AbstractContextManager
+from tinygrad.uop.ops import UOp, KernelInfo, AxisType
+from extra.thunder.tiny.tk import WARP_THREADS
+from extra.thunder.tiny.tk.group import Group
+
+class _tk_range:
+  user_rid = 0
+  def __init__(self, end:int, axis_type:AxisType): self.end, self.axis_type, self.done = end, axis_type, False
+  def __iter__(self): return self
+  def __next__(self):
+    if not self.done:
+      self.done = True
+      _tk_range.user_rid += 1
+      self._rng = UOp.range(self.end, _tk_range.user_rid-1, axis_type=self.axis_type)
+      return self._rng
+    raise StopIteration
+
+class Kernel(AbstractContextManager):
+  def __init__(self, grid_size:tuple[int, int, int], block_size:int):
+    self.blockIdx_x = UOp.special(grid_size[0], "gidx0")
+    self.blockIdx_y = UOp.special(grid_size[1], "gidx1")
+    self.blockIdx_z = UOp.special(grid_size[2], "gidx2")
+    self.threadIdx_x = UOp.special(block_size, "lidx0")
+
+    self.range_stack = []
+    self.store_stack = []
+
+  @property
+  def warpid(self): return self.threadIdx_x // WARP_THREADS
+
+  def __enter__(self): return self
+  def __exit__(self, exc_type, exc_value, traceback): pass
+
+  def group(self, size:int): return Group(size, self)
+  @property
+  def warp(self): return self.group(1)
+  @property
+  def warpgroup(self): return self.group(4)
+
+  def range(self, end:int, axis_type:AxisType=AxisType.LOOP, track:bool=True):
+    rng = _tk_range(end, axis_type)
+    if track: self.range_stack.append(rng)
+    return rng
+
+  def push_store(self, store:UOp, uop:UOp): self.store_stack.append((store, uop))
+
+  def finish(self):
+    # end all ranges
+    rngs = []
+    while self.range_stack: rngs.append(self.range_stack.pop(0)._rng)
+
+    return self.store_stack.pop()[0].end(*rngs).sink(arg=KernelInfo(opts_to_apply=())).simplify()
+
+  def endrange(self):
+    last_store = self.store_stack.pop()
+    last_range = self.range_stack.pop()
+    return last_store[1].after(last_store[0].barrier().end(last_range._rng)).reshape(last_store[1].shape)

extra/thunder/tiny/tk/tiles.py

@@ -0,0 +1,52 @@
+import math
+from typing import cast, Callable
+from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
+from tinygrad.uop.ops import AxisType, UOp, KernelInfo, Ops
+from tinygrad.engine.realize import ExecItem, get_runner
+from tinygrad.dtype import AddrSpace, PtrDType
+from tinygrad.helpers import getenv, prod
+
+from extra.thunder.tiny.tk import WARP_THREADS
+
+class _Slots:
+  def __init__(self):
+    self.global_slot = 0
+    self.shared_slot = 0
+    self.register_slot = 0
+slots = _Slots()
+
+def gl(shape, dtype):
+  slots.global_slot += 1
+  return UOp.placeholder(shape, dtype, slot=slots.global_slot-1)
+
+shared_slot = 0
+def st(shape, dtype):
+  slots.shared_slot += 1
+  return UOp.placeholder(shape, dtype, addrspace=AddrSpace.LOCAL, slot=slots.shared_slot-1)
+
+TILE_ROW_DIM, TILE_COL_DIM = 16, 16
+RT_BASE_TILE_NE = TILE_ROW_DIM * TILE_COL_DIM
+RT_BASE_TILE_NEPT = RT_BASE_TILE_NE // WARP_THREADS
+register_slot = 0
+def rt(shape, dtype):
+  assert len(shape) == 2
+
+  height = shape[0] // TILE_ROW_DIM
+  width = shape[1] // TILE_COL_DIM
+
+  slots.register_slot += 1
+  return UOp.placeholder((height, width, RT_BASE_TILE_NEPT), dtype, addrspace=AddrSpace.REG, slot=slots.register_slot-1)
+
+def rv(length, dtype, layout="naive"):
+  tiles = length // TILE_ROW_DIM
+  match layout:
+    case "naive":
+      inner_dim = 1
+      outer_dim = (tiles + 1) // 2
+    case "ortho":
+      inner_dim = 1
+      outer_dim = tiles
+    case _: raise NotImplementedError(f"rv layout {layout} not implemented")
+
+  slots.register_slot += 1
+  return UOp.placeholder((outer_dim, inner_dim, 2), dtype, addrspace=AddrSpace.REG, slot=slots.register_slot-1)

extra/weekly_commits_table.py

@@ -1,7 +1,7 @@
 # extra/weekly_commits_table.py
 import os, subprocess, datetime as dt
 
-NAMES = ["chenyu","George Hotz","nimlgen","qazal","Sieds Lykles","wozeparrot"]
+NAMES = ["chenyu","George Hotz","nimlgen","qazal","wozeparrot"]
 REPO  = os.environ.get("REPO_PATH",".")
 today = dt.date.today()
 days  = [(today - dt.timedelta(i)).strftime("%Y-%m-%d") for i in range(6,-1,-1)]
@@ -40,4 +40,4 @@ for d in days:
 print("** Commits by day (last 7) **")
 print("```")
 print("\n".join([header, rule] + rows))
-print("```")
+print("```")

setup.cfg

@@ -0,0 +1,21 @@
+[mutmut]
+paths_to_mutate=tinygrad
+do_not_mutate=
+  tinygrad/apps/*
+  tinygrad/codegen/*
+  tinygrad/engine/*
+  tinygrad/nn/*
+  tinygrad/renderer/*
+  tinygrad/runtime/*
+  tinygrad/schedule/*
+  tinygrad/uop/*
+  tinygrad/viz/*
+  tinygrad/device.py
+  tinygrad/dtype.py
+  tinygrad/gradient.py
+  tinygrad/helpers.py
+  tinygrad/tensor.py
+tests_dir=
+  test/test_tiny.py
+  test/test_ops.py
+debug=true

test/external/external_benchmark_op_cat.py

@@ -0,0 +1,163 @@
+# ruff: noqa: E501 E712
+from tinygrad import dtypes, Device
+from tinygrad.uop.ops import UOp, AxisType, Ops, KernelInfo
+from tinygrad.codegen import full_rewrite
+from tinygrad.renderer import ProgramSpec
+from tinygrad.engine.realize import CompiledRunner
+from tinygrad.helpers import dedup
+from tinygrad.device import Buffer
+from tinygrad.dtype import ImageDType, Invalid
+
+c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1576), (), 0)
+c2 = UOp.range(1576, 20, AxisType.LOOP)
+c5 = c2<55
+c6 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 16, 4)), (), 1)
+c8 = UOp.range(16, 0, AxisType.REDUCE)
+c11 = UOp.range(4, 1, AxisType.REDUCE)
+c14 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((14, 64, 4)), (), 2)
+c25 = c5.where((c2%4*4+c11+c8*16+c2//4*256), UOp.const(dtypes.index, Invalid))
+c27 = c6.index((c8*4+c11))*c14.index(c25)
+c29 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(55), (), 3)
+c30 = c5.where(c2, UOp.const(dtypes.index, Invalid))
+c34 = c5.where((c27.reduce(c8, c11, arg=Ops.ADD)+c29.index(c30)), UOp.const(dtypes.float, 0.0))
+c38 = c2<87
+c39 = (c5!=True)&c38
+c40 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 8, 4)), (), 4)
+c42 = UOp.range(8, 2, AxisType.REDUCE)
+c44 = UOp.range(4, 3, AxisType.REDUCE)
+c47 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((8, 32, 4)), (), 5)
+c49 = c2+1
+c51 = c49%4*4
+c57 = c49//4*128
+c61 = c39.where((c51+c44+c42*16+c57+-1792), UOp.const(dtypes.index, Invalid))
+c63 = c40.index((c42*4+c44))*c47.index(c61)
+c65 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(32), (), 6)
+c68 = c39.where((c2+-55), UOp.const(dtypes.index, Invalid))
+c71 = c39.where((c63.reduce(c42, c44, arg=Ops.ADD)+c65.index(c68)), UOp.const(dtypes.float, 0.0))
+c75 = c2<99
+c76 = (c38!=True)&c75
+c77 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 8, 4)), (), 7)
+c78 = UOp.range(8, 4, AxisType.REDUCE)
+c80 = UOp.range(4, 5, AxisType.REDUCE)
+c83 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((3, 32, 4)), (), 8)
+c90 = c76.where((c51+c80+c78*16+c57+-2816), UOp.const(dtypes.index, Invalid))
+c92 = c77.index((c78*4+c80))*c83.index(c90)
+c94 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(12), (), 9)
+c97 = c76.where((c2+-87), UOp.const(dtypes.index, Invalid))
+c100 = c76.where((c92.reduce(c78, c80, arg=Ops.ADD)+c94.index(c97)), UOp.const(dtypes.float, 0.0))
+c104 = c2<105
+c105 = (c75!=True)&c104
+c106 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 8, 4)), (), 10)
+c107 = UOp.range(8, 6, AxisType.REDUCE)
+c109 = UOp.range(4, 7, AxisType.REDUCE)
+c112 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((2, 32, 4)), (), 11)
+c119 = c105.where((c51+c109+c107*16+c57+-3200), UOp.const(dtypes.index, Invalid))
+c121 = c106.index((c107*4+c109))*c112.index(c119)
+c123 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(6), (), 12)
+c126 = c105.where((c2+-99), UOp.const(dtypes.index, Invalid))
+c129 = c105.where((c121.reduce(c107, c109, arg=Ops.ADD)+c123.index(c126)), UOp.const(dtypes.float, 0.0))
+c133 = c2<117
+c134 = (c104!=True)&c133
+c135 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 8, 4)), (), 13)
+c136 = UOp.range(8, 8, AxisType.REDUCE)
+c138 = UOp.range(4, 9, AxisType.REDUCE)
+c141 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((3, 32, 4)), (), 14)
+c143 = c2+3
+c145 = c143%4*4
+c149 = c143//4
+c150 = c149*128
+c154 = c134.where((c145+c138+c136*16+c150+-3456), UOp.const(dtypes.index, Invalid))
+c156 = c135.index((c136*4+c138))*c141.index(c154)
+c158 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(12), (), 15)
+c161 = c134.where((c2+-105), UOp.const(dtypes.index, Invalid))
+c164 = c134.where((c156.reduce(c136, c138, arg=Ops.ADD)+c158.index(c161)), UOp.const(dtypes.float, 0.0))
+c168 = c2<645
+c169 = (c133!=True)&c168
+c170 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 16, 4)), (), 16)
+c171 = UOp.range(16, 10, AxisType.REDUCE)
+c173 = UOp.range(4, 11, AxisType.REDUCE)
+c176 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((132, 64, 4)), (), 17)
+c180 = c149*256
+c184 = c169.where((c145+c173+c171*16+c180+-7680), UOp.const(dtypes.index, Invalid))
+c186 = c170.index((c171*4+c173))*c176.index(c184)
+c188 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(528), (), 18)
+c191 = c169.where((c2+-117), UOp.const(dtypes.index, Invalid))
+c194 = c169.where((c186.reduce(c171, c173, arg=Ops.ADD)+c188.index(c191)), UOp.const(dtypes.float, 0.0))
+c198 = c2<653
+c199 = (c168!=True)&c198
+c200 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 4, 4)), (), 19)
+c201 = UOp.range(4, 12, AxisType.REDUCE)
+c203 = UOp.range(4, 13, AxisType.REDUCE)
+c206 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((2, 16, 4)), (), 20)
+c215 = c199.where((c145+c203+c201*16+c149*64+-10368), UOp.const(dtypes.index, Invalid))
+c217 = c200.index((c201*4+c203))*c206.index(c215)
+c219 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(8), (), 21)
+c222 = c199.where((c2+-645), UOp.const(dtypes.index, Invalid))
+c225 = c199.where((c217.reduce(c201, c203, arg=Ops.ADD)+c219.index(c222)), UOp.const(dtypes.float, 0.0))
+c229 = c2<917
+c230 = (c198!=True)&c229
+c231 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 8, 4)), (), 22)
+c232 = UOp.range(8, 14, AxisType.REDUCE)
+c234 = UOp.range(4, 15, AxisType.REDUCE)
+c237 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((66, 32, 4)), (), 23)
+c244 = c230.where((c145+c234+c232*16+c150+-20992), UOp.const(dtypes.index, Invalid))
+c246 = c231.index((c232*4+c234))*c237.index(c244)
+c248 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(264), (), 24)
+c251 = c230.where((c2+-653), UOp.const(dtypes.index, Invalid))
+c254 = c230.where((c246.reduce(c232, c234, arg=Ops.ADD)+c248.index(c251)), UOp.const(dtypes.float, 0.0))
+c258 = c2<1061
+c259 = (c229!=True)&c258
+c260 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 16, 4)), (), 25)
+c261 = UOp.range(16, 16, AxisType.REDUCE)
+c263 = UOp.range(4, 17, AxisType.REDUCE)
+c266 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((36, 64, 4)), (), 26)
+c273 = c259.where((c145+c263+c261*16+c180+-58880), UOp.const(dtypes.index, Invalid))
+c275 = c260.index((c261*4+c263))*c266.index(c273)
+c277 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(144), (), 27)
+c280 = c259.where((c2+-917), UOp.const(dtypes.index, Invalid))
+c283 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(144), (), 28)
+c286 = c259.where(((c275.reduce(c261, c263, arg=Ops.ADD)+c277.index(c280))*c283.index(c280)), UOp.const(dtypes.float, 0.0))
+c290 = c2<1064
+c291 = (c258!=True)&c290
+c292 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 4, 4)), (), 29)
+c293 = UOp.range(4, 18, AxisType.REDUCE)
+c295 = UOp.range(4, 19, AxisType.REDUCE)
+c298 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 16, 4)), (), 30)
+c305 = c291.where((c2*4+c295+c293*16+-4244), UOp.const(dtypes.index, Invalid))
+c307 = c292.index((c293*4+c295))*c298.index(c305)
+c309 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(3), (), 31)
+c312 = c291.where((c2+-1061), UOp.const(dtypes.index, Invalid))
+c315 = c291.where((c307.reduce(c293, c295, arg=Ops.ADD)+c309.index(c312)), UOp.const(dtypes.float, 0.0))
+c317 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 128, 4)), (), 32)
+c321 = (c290!=True).where((c2+-1064), UOp.const(dtypes.index, Invalid))
+c323 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1), (), 33)
+c328 = c290.where(UOp.const(dtypes.float, 0.0), (c317.index(c321)*c323.index(UOp.const(dtypes.index, 0)).reciprocal()))
+c329 = c34+c71+c100+c129+c164+c194+c225+c254+c286+c315+c328
+c331 = c0.index(c2, ptr=True).store(c329).end(c2)
+ast = c331.sink(arg=KernelInfo(name="cat", opts_to_apply=None))
+
+compiler = Device.default.compiler
+renderer = Device.default.renderer
+allocator = Device.default.allocator
+
+uops = full_rewrite(ast, renderer)
+src = renderer.render(uops)
+
+# NOLOCALS=1 IMAGE=2 DEV=CL
+lib = compiler.compile(src)
+
+ps = ProgramSpec("cat", src, Device.DEFAULT, ast, uops)
+# print(ps.src)
+# print(ps.applied_opts)
+# NOTE: this is faster with no GROUP and with NOLOCALS
+# (Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UNROLL, axis=19, arg=4), Opt(op=OptOps.UNROLL, axis=17, arg=4), Opt(op=OptOps.UNROLL, axis=15, arg=4), Opt(op=OptOps.UNROLL, axis=13, arg=4), Opt(op=OptOps.UNROLL, axis=11, arg=4), Opt(op=OptOps.UNROLL, axis=9, arg=4), Opt(op=OptOps.UNROLL, axis=7, arg=4), Opt(op=OptOps.UNROLL, axis=5, arg=4), Opt(op=OptOps.UNROLL, axis=3, arg=4), Opt(op=OptOps.UNROLL, axis=1, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None))
+cr = CompiledRunner(ps, precompiled=lib)
+
+gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.DEFINE_GLOBAL]), key=lambda u: u.arg)
+print(len(gs))
+print([g.dtype for g in gs])
+
+bufs = [Buffer(ps.device, g.size, g.dtype if isinstance(g.dtype, ImageDType) else g.dtype._base).ensure_allocated() for g in gs]
+
+t = cr(bufs, wait=True)
+print(f"{t*1e6:.2f} us")

test/external/external_benchmark_op_conv.py

@@ -1,237 +1,65 @@
-# ruff: noqa: E501
+# ruff: noqa: E501 E712
 from tinygrad import dtypes, Device
-from tinygrad.uop.ops import UOp, AxisType, Ops
+from tinygrad.uop.ops import UOp, AxisType, Ops, KernelInfo
 from tinygrad.codegen import full_rewrite
+# from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.renderer import ProgramSpec
 from tinygrad.engine.realize import CompiledRunner
-from tinygrad.helpers import dedup
+from tinygrad.helpers import dedup, getenv
 from tinygrad.device import Buffer
-from tinygrad.dtype import ImageDType
+from tinygrad.dtype import ImageDType, Invalid
 
-# PYTHONPATH="." DEBUG=5 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-# kernel 672
-# faster on d59d4cd, 50% slower with the new linearizer
+# PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
 
-""" d59d4cd
-c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), arg=0, src=())
-c1 = UOp.range(UOp.const(dtypes.index, 64), 3, AxisType.LOOP)
-c2 = UOp.range(UOp.const(dtypes.index, 64), 4, AxisType.LOOP)
-c3 = UOp.range(UOp.const(dtypes.index, 32), 2, AxisType.LOOP)
-c4 = (((c1*UOp.const(dtypes.index, 64))+c2)+(c3*UOp.const(dtypes.index, 4096)))
-c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), arg=1, src=())
-c6 = c5.index(c4).load()
-c7 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 3072, 4)), arg=2, src=())
-c8 = UOp.range(UOp.const(dtypes.index, 48), 0, AxisType.REDUCE)
-c9 = UOp.range(UOp.const(dtypes.index, 4), 1, AxisType.REDUCE)
-c10 = c7.index(((((c8*UOp.const(dtypes.index, 4))+c9)+(c1*UOp.const(dtypes.index, 192)))+(c3*UOp.const(dtypes.index, 12288)))).load()
-c11 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 192, 4)), arg=3, src=())
-c12 = c11.index(((((c9*UOp.const(dtypes.index, 4))+(c2%UOp.const(dtypes.index, 4)))+(c8*UOp.const(dtypes.index, 16)))+((c2//UOp.const(dtypes.index, 4))*UOp.const(dtypes.index, 768)))).load()
-c13 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(64), arg=4, src=())
-c14 = c13.index(c2).load()
-c15 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(64), arg=5, src=())
-c16 = c15.index(c2).load()
-c17 = (c6+(((c10*c12.cast(dtypes.float)).cast(dtypes.float).reduce(c8, c9, arg=Ops.ADD)+c14.cast(dtypes.float))*c16.cast(dtypes.float)))
-c18 = c0.index(c4).store(c17, c3, c1, c2)
-ast = c18.sink()
-more upcast axis : [(3, 320, 0, 4)]
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-__kernel void r_512_16_4_4_48_4(write_only image2d_t data0_131072, read_only image2d_t data1_131072, read_only image2d_t data2_393216, read_only image2d_t data3_12288, __global half* data4_64, __global half* data5_64) {
-const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;
-  float acc0[16];
-  int idx0 = get_global_id(0); /* 16 */
-  int idx1 = get_global_id(1); /* 512 */
-  int alu0 = (idx1>>4);
-  *(acc0+0) = 0.0f;
-  *(acc0+1) = 0.0f;
-  *(acc0+2) = 0.0f;
-  *(acc0+3) = 0.0f;
-  *(acc0+4) = 0.0f;
-  *(acc0+5) = 0.0f;
-  *(acc0+6) = 0.0f;
-  *(acc0+7) = 0.0f;
-  *(acc0+8) = 0.0f;
-  *(acc0+9) = 0.0f;
-  *(acc0+10) = 0.0f;
-  *(acc0+11) = 0.0f;
-  *(acc0+12) = 0.0f;
-  *(acc0+13) = 0.0f;
-  *(acc0+14) = 0.0f;
-  *(acc0+15) = 0.0f;
-  for (int Ridx0 = 0; Ridx0 < 48; Ridx0++) {
-    int alu17 = ((idx1*192)+Ridx0);
-    int alu18 = (alu17+48);
-    int alu19 = (alu17+96);
-    int alu20 = (alu17+144);
-    int alu21 = (Ridx0<<2);
-    float4 val0 = read_imagef(data3_12288, smp, (int2)(alu21,idx0));
-    float4 val1 = read_imagef(data3_12288, smp, (int2)((alu21+1),idx0));
-    float4 val2 = read_imagef(data3_12288, smp, (int2)((alu21+2),idx0));
-    float4 val3 = read_imagef(data3_12288, smp, (int2)((alu21+3),idx0));
-    float4 val4 = read_imagef(data2_393216, smp, (int2)((alu18-(3072*(((alu18>>10)*43)>>7))),alu0));
-    float4 val5 = read_imagef(data2_393216, smp, (int2)((alu19-(3072*(((alu19>>10)*43)>>7))),alu0));
-    float4 val6 = read_imagef(data2_393216, smp, (int2)((alu20-(3072*(((alu20>>10)*43)>>7))),alu0));
-    float4 val7 = read_imagef(data2_393216, smp, (int2)((alu17-(3072*(((alu17>>10)*43)>>7))),alu0));
-    *(acc0+1) = ((*(acc0+1))+(val4.x*val0.x)+(val4.y*val1.x)+(val4.z*val2.x)+(val4.w*val3.x));
-    *(acc0+5) = ((*(acc0+5))+(val4.x*val0.y)+(val4.y*val1.y)+(val4.z*val2.y)+(val4.w*val3.y));
-    *(acc0+9) = ((*(acc0+9))+(val4.x*val0.z)+(val4.y*val1.z)+(val4.z*val2.z)+(val4.w*val3.z));
-    *(acc0+13) = ((*(acc0+13))+(val4.x*val0.w)+(val4.y*val1.w)+(val4.z*val2.w)+(val4.w*val3.w));
-    *(acc0+2) = ((*(acc0+2))+(val5.x*val0.x)+(val5.y*val1.x)+(val5.z*val2.x)+(val5.w*val3.x));
-    *(acc0+6) = ((*(acc0+6))+(val5.x*val0.y)+(val5.y*val1.y)+(val5.z*val2.y)+(val5.w*val3.y));
-    *(acc0+10) = ((*(acc0+10))+(val5.x*val0.z)+(val5.y*val1.z)+(val5.z*val2.z)+(val5.w*val3.z));
-    *(acc0+14) = ((*(acc0+14))+(val5.x*val0.w)+(val5.y*val1.w)+(val5.z*val2.w)+(val5.w*val3.w));
-    *(acc0+3) = ((*(acc0+3))+(val6.x*val0.x)+(val6.y*val1.x)+(val6.z*val2.x)+(val6.w*val3.x));
-    *(acc0+7) = ((*(acc0+7))+(val6.x*val0.y)+(val6.y*val1.y)+(val6.z*val2.y)+(val6.w*val3.y));
-    *(acc0+11) = ((*(acc0+11))+(val6.x*val0.z)+(val6.y*val1.z)+(val6.z*val2.z)+(val6.w*val3.z));
-    *(acc0+15) = ((*(acc0+15))+(val6.x*val0.w)+(val6.y*val1.w)+(val6.z*val2.w)+(val6.w*val3.w));
-    *(acc0+0) = ((*(acc0+0))+(val7.x*val0.x)+(val7.y*val1.x)+(val7.z*val2.x)+(val7.w*val3.x));
-    *(acc0+4) = ((*(acc0+4))+(val7.x*val0.y)+(val7.y*val1.y)+(val7.z*val2.y)+(val7.w*val3.y));
-    *(acc0+8) = ((*(acc0+8))+(val7.x*val0.z)+(val7.y*val1.z)+(val7.z*val2.z)+(val7.w*val3.z));
-    *(acc0+12) = ((*(acc0+12))+(val7.x*val0.w)+(val7.y*val1.w)+(val7.z*val2.w)+(val7.w*val3.w));
-  }
-  int alu39 = (idx0<<2);
-  half4 val8 = (*((__global half4*)((data4_64+alu39))));
-  half4 val9 = (*((__global half4*)((data5_64+alu39))));
-  int alu40 = (idx0+(idx1<<6));
-  int2 cast0 = (int2)((alu40&1023),alu0);
-  float4 val10 = read_imagef(data1_131072, smp, cast0);
-  int2 cast1 = (int2)(((alu40+16)&1023),alu0);
-  float4 val11 = read_imagef(data1_131072, smp, cast1);
-  int2 cast2 = (int2)(((alu40+32)&1023),alu0);
-  float4 val12 = read_imagef(data1_131072, smp, cast2);
-  int2 cast3 = (int2)(((alu40+48)&1023),alu0);
-  float4 val13 = read_imagef(data1_131072, smp, cast3);
-  float cast4 = ((float)(val8.x));
-  float cast5 = ((float)(val9.x));
-  float cast6 = ((float)(val8.y));
-  float cast7 = ((float)(val9.y));
-  float cast8 = ((float)(val8.z));
-  float cast9 = ((float)(val9.z));
-  float cast10 = ((float)(val8.w));
-  float cast11 = ((float)(val9.w));
-  write_imagef(data0_131072, cast0, (float4)((val10.x+(((*(acc0+0))+cast4)*cast5)),(val10.y+(((*(acc0+4))+cast6)*cast7)),(val10.z+(((*(acc0+8))+cast8)*cast9)),(val10.w+(((*(acc0+12))+cast10)*cast11))));
-  write_imagef(data0_131072, cast1, (float4)((val11.x+(((*(acc0+1))+cast4)*cast5)),(val11.y+(((*(acc0+5))+cast6)*cast7)),(val11.z+(((*(acc0+9))+cast8)*cast9)),(val11.w+(((*(acc0+13))+cast10)*cast11))));
-  write_imagef(data0_131072, cast2, (float4)((val12.x+(((*(acc0+2))+cast4)*cast5)),(val12.y+(((*(acc0+6))+cast6)*cast7)),(val12.z+(((*(acc0+10))+cast8)*cast9)),(val12.w+(((*(acc0+14))+cast10)*cast11))));
-  write_imagef(data0_131072, cast3, (float4)((val13.x+(((*(acc0+3))+cast4)*cast5)),(val13.y+(((*(acc0+7))+cast6)*cast7)),(val13.z+(((*(acc0+11))+cast8)*cast9)),(val13.w+(((*(acc0+15))+cast10)*cast11))));
-}
-*** QCOM     672 r_512_16_4_4_48_4                              arg  6 mem   0.10 GB tm    322.55us/    77.83ms (    157 GFLOPS    4|160    GB/s) ['mul', '__add__', 'conv2d']
-"""
+def vision_conv_143():
+  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 1024, 4)), (), 0)
+  c2 = UOp.range(32, 3, AxisType.LOOP)
+  c5 = UOp.range(128, 4, AxisType.LOOP)
+  c8 = UOp.range(16, 2, AxisType.LOOP)
+  c16 = UOp.range(7, 0, AxisType.REDUCE)
+  c17 = c8*2+c16
+  c24 = ((c17<3)!=True)&(c17<35)
+  c26 = UOp.range(7, 1, AxisType.REDUCE)
+  c27 = c2*2+c26
+  c32 = ((c27<3)!=True)&(c27<67)
+  c34 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 1)
+  c38 = c5//2
+  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.index, Invalid))
+  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
+  c49 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((64, 49, 4)), (), 2)
+  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
+  c63 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(128), (), 3)
+  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
+  c67 = c0.index((c2*128+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)
 
-""" master 99e76f33a0f4ec84c79c1271dbc955fe6b5a7778
-c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 0)
-c2 = UOp.range(64, 3, AxisType.LOOP)
-c4 = UOp.range(64, 4, AxisType.LOOP)
-c7 = UOp.range(32, 2, AxisType.LOOP)
-c10 = (((c2*64)+c4)+(c7*4096))
-c12 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 1)
-c14 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 3072, 4)), (), 2)
-c16 = UOp.range(48, 0, AxisType.REDUCE)
-c19 = UOp.range(4, 1, AxisType.REDUCE)
-c28 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 192, 4)), (), 3)
-c40 = (c14.index(((((c16*4)+c19)+(c2*192))+(c7*12288)))*c28.index(((((c19*4)+(c4%4))+(c16*16))+((c4//4)*768))))
-c42 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(64), (), 4)
-c46 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(64), (), 5)
-c50 = (c12.index(c10)+((c40.reduce(c16, c19, arg=Ops.ADD)+c42.index(c4).cast(dtypes.float))*c46.index(c4).cast(dtypes.float)))
-c52 = c0.index(c10, ptr=True).store(c50).end(c7, c2, c4)
-ast = c52.sink()
-more upcast axis : [(3, 320, 0, 4)]
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-__kernel void r_512_16_4_4_48_4(write_only image2d_t data0_131072, read_only image2d_t data1_131072, read_only image2d_t data2_393216, read_only image2d_t data3_12288, __global half* data4_64, __global half* data5_64) {
-const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;
-  float acc0[16];
-  int idx0 = get_global_id(0); /* 16 */
-  int idx1 = get_global_id(1); /* 512 */
-  *(acc0+0) = 0.0f;
-  *(acc0+1) = 0.0f;
-  *(acc0+2) = 0.0f;
-  *(acc0+3) = 0.0f;
-  *(acc0+4) = 0.0f;
-  *(acc0+5) = 0.0f;
-  *(acc0+6) = 0.0f;
-  *(acc0+7) = 0.0f;
-  *(acc0+8) = 0.0f;
-  *(acc0+9) = 0.0f;
-  *(acc0+10) = 0.0f;
-  *(acc0+11) = 0.0f;
-  *(acc0+12) = 0.0f;
-  *(acc0+13) = 0.0f;
-  *(acc0+14) = 0.0f;
-  *(acc0+15) = 0.0f;
-  int alu16 = (idx0<<2);
-  half4 val0 = (*((__global half4*)((data4_64+alu16))));
-  half4 val1 = (*((__global half4*)((data5_64+alu16))));
-  int alu17 = (idx0+(idx1<<6));
-  int alu18 = (idx1>>4);
-  int2 cast0 = (int2)((alu17&1023),alu18);
-  float4 val2 = read_imagef(data1_131072, smp, cast0);
-  int2 cast1 = (int2)(((alu17+16)&1023),alu18);
-  float4 val3 = read_imagef(data1_131072, smp, cast1);
-  int2 cast2 = (int2)(((alu17+32)&1023),alu18);
-  float4 val4 = read_imagef(data1_131072, smp, cast2);
-  int2 cast3 = (int2)(((alu17+48)&1023),alu18);
-  float4 val5 = read_imagef(data1_131072, smp, cast3);
-  for (int Ridx0 = 0; Ridx0 < 48; Ridx0++) {
-    int alu19 = ((idx1*192)+Ridx0);
-    int alu20 = (alu19+48);
-    int alu21 = (alu19+96);
-    int alu22 = (alu19+144);
-    int alu23 = (Ridx0<<2);
-    float4 val6 = read_imagef(data3_12288, smp, (int2)(alu23,idx0));
-    float4 val7 = read_imagef(data3_12288, smp, (int2)((alu23+1),idx0));
-    float4 val8 = read_imagef(data3_12288, smp, (int2)((alu23+2),idx0));
-    float4 val9 = read_imagef(data3_12288, smp, (int2)((alu23+3),idx0));
-    float4 val10 = read_imagef(data2_393216, smp, (int2)((alu20-(3072*(((alu20>>10)*43)>>7))),alu18));
-    *(acc0+1) = ((*(acc0+1))+(val10.x*val6.x)+(val10.y*val7.x)+(val10.z*val8.x)+(val10.w*val9.x));
-    *(acc0+5) = ((*(acc0+5))+(val10.x*val6.y)+(val10.y*val7.y)+(val10.z*val8.y)+(val10.w*val9.y));
-    *(acc0+9) = ((*(acc0+9))+(val10.x*val6.z)+(val10.y*val7.z)+(val10.z*val8.z)+(val10.w*val9.z));
-    *(acc0+13) = ((*(acc0+13))+(val10.x*val6.w)+(val10.y*val7.w)+(val10.z*val8.w)+(val10.w*val9.w));
-    float4 val11 = read_imagef(data2_393216, smp, (int2)((alu21-(3072*(((alu21>>10)*43)>>7))),alu18));
-    *(acc0+2) = ((*(acc0+2))+(val11.x*val6.x)+(val11.y*val7.x)+(val11.z*val8.x)+(val11.w*val9.x));
-    *(acc0+6) = ((*(acc0+6))+(val11.x*val6.y)+(val11.y*val7.y)+(val11.z*val8.y)+(val11.w*val9.y));
-    *(acc0+10) = ((*(acc0+10))+(val11.x*val6.z)+(val11.y*val7.z)+(val11.z*val8.z)+(val11.w*val9.z));
-    *(acc0+14) = ((*(acc0+14))+(val11.x*val6.w)+(val11.y*val7.w)+(val11.z*val8.w)+(val11.w*val9.w));
-    float4 val12 = read_imagef(data2_393216, smp, (int2)((alu22-(3072*(((alu22>>10)*43)>>7))),alu18));
-    *(acc0+3) = ((*(acc0+3))+(val12.x*val6.x)+(val12.y*val7.x)+(val12.z*val8.x)+(val12.w*val9.x));
-    *(acc0+7) = ((*(acc0+7))+(val12.x*val6.y)+(val12.y*val7.y)+(val12.z*val8.y)+(val12.w*val9.y));
-    *(acc0+11) = ((*(acc0+11))+(val12.x*val6.z)+(val12.y*val7.z)+(val12.z*val8.z)+(val12.w*val9.z));
-    *(acc0+15) = ((*(acc0+15))+(val12.x*val6.w)+(val12.y*val7.w)+(val12.z*val8.w)+(val12.w*val9.w));
-    float4 val13 = read_imagef(data2_393216, smp, (int2)((alu19-(3072*(((alu19>>10)*43)>>7))),alu18));
-    *(acc0+0) = ((*(acc0+0))+(val13.x*val6.x)+(val13.y*val7.x)+(val13.z*val8.x)+(val13.w*val9.x));
-    *(acc0+4) = ((*(acc0+4))+(val13.x*val6.y)+(val13.y*val7.y)+(val13.z*val8.y)+(val13.w*val9.y));
-    *(acc0+8) = ((*(acc0+8))+(val13.x*val6.z)+(val13.y*val7.z)+(val13.z*val8.z)+(val13.w*val9.z));
-    *(acc0+12) = ((*(acc0+12))+(val13.x*val6.w)+(val13.y*val7.w)+(val13.z*val8.w)+(val13.w*val9.w));
-  }
-  float cast4 = ((float)(val0.x));
-  float cast5 = ((float)(val1.x));
-  float cast6 = ((float)(val0.y));
-  float cast7 = ((float)(val1.y));
-  float cast8 = ((float)(val0.z));
-  float cast9 = ((float)(val1.z));
-  float cast10 = ((float)(val0.w));
-  float cast11 = ((float)(val1.w));
-  write_imagef(data0_131072, cast0, (float4)((val2.x+(((*(acc0+0))+cast4)*cast5)),(val2.y+(((*(acc0+4))+cast6)*cast7)),(val2.z+(((*(acc0+8))+cast8)*cast9)),(val2.w+(((*(acc0+12))+cast10)*cast11))));
-  write_imagef(data0_131072, cast1, (float4)((val3.x+(((*(acc0+1))+cast4)*cast5)),(val3.y+(((*(acc0+5))+cast6)*cast7)),(val3.z+(((*(acc0+9))+cast8)*cast9)),(val3.w+(((*(acc0+13))+cast10)*cast11))));
-  write_imagef(data0_131072, cast2, (float4)((val4.x+(((*(acc0+2))+cast4)*cast5)),(val4.y+(((*(acc0+6))+cast6)*cast7)),(val4.z+(((*(acc0+10))+cast8)*cast9)),(val4.w+(((*(acc0+14))+cast10)*cast11))));
-  write_imagef(data0_131072, cast3, (float4)((val5.x+(((*(acc0+3))+cast4)*cast5)),(val5.y+(((*(acc0+7))+cast6)*cast7)),(val5.z+(((*(acc0+11))+cast8)*cast9)),(val5.w+(((*(acc0+15))+cast10)*cast11))));
-}
-*** QCOM     672 r_512_16_4_4_48_4                              arg  6 mem   0.10 GB tm    527.97us/    78.94ms (     96 GFLOPS    3|98     GB/s) ['conv2d', 'mul', '__add__']
-"""
+  opts = None
+  return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
 
-c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 0)
-c2 = UOp.range(64, 3, AxisType.LOOP)
-c4 = UOp.range(64, 4, AxisType.LOOP)
-c7 = UOp.range(32, 2, AxisType.LOOP)
-c10 = (((c2*64)+c4)+(c7*4096))
-c12 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 1)
-c14 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 3072, 4)), (), 2)
-c16 = UOp.range(48, 0, AxisType.REDUCE)
-c19 = UOp.range(4, 1, AxisType.REDUCE)
-c28 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 192, 4)), (), 3)
-c40 = (c14.index(((((c16*4)+c19)+(c2*192))+(c7*12288)))*c28.index(((((c19*4)+(c4%4))+(c16*16))+((c4//4)*768))))
-c42 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(64), (), 4)
-c46 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(64), (), 5)
-c50 = (c12.index(c10)+((c40.reduce(c16, c19, arg=Ops.ADD)+c42.index(c4).cast(dtypes.float))*c46.index(c4).cast(dtypes.float)))
-c52 = c0.index(c10, ptr=True).store(c50).end(c7, c2, c4)
-ast = c52.sink()
+def vision_conv_153():
+  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((8, 1024, 4)), (), 0)
+  c2 = UOp.range(16, 3, AxisType.LOOP)
+  c5 = UOp.range(256, 4, AxisType.LOOP)
+  c8 = UOp.range(8, 2, AxisType.LOOP)
+  c16 = UOp.range(7, 0, AxisType.REDUCE)
+  c17 = c8*2+c16
+  c24 = ((c17<3)!=True)&(c17<19)
+  c26 = UOp.range(7, 1, AxisType.REDUCE)
+  c27 = c2*2+c26
+  c32 = ((c27<3)!=True)&(c27<35)
+  c34 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 1024, 4)), (), 1)
+  c38 = c5//2
+  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.index, Invalid))
+  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
+  c49 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((128, 49, 4)), (), 2)
+  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
+  c63 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(256), (), 3)
+  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
+  c67 = c0.index((c2*256+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)
+
+  opts = None
+  return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
+
+ast = vision_conv_143() if getenv("NUM", 143) == 143 else vision_conv_153()
 
 compiler = Device.default.compiler
 renderer = Device.default.renderer
@@ -240,20 +68,13 @@ allocator = Device.default.allocator
 uops = full_rewrite(ast, renderer)
 src = renderer.render(uops)
 
-# NOLOCALS=1 IMAGE=2 DEV=CL
 lib = compiler.compile(src)
-# r_64_8_16_4_4_48_4
-# NOLOCALS: r_512_16_4_4_48_4
-ps = ProgramSpec("r_512_16_4_4_48_4", src, Device.DEFAULT, ast, uops)
-print(ps.src)
-print(ps.applied_opts)
-# (Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None))
+ps = ProgramSpec("conv", src, Device.DEFAULT, ast, uops)
 cr = CompiledRunner(ps, precompiled=lib)
 
 gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.DEFINE_GLOBAL]), key=lambda u: u.arg)
-print(len(gs))
-print([g.dtype for g in gs])
-
+# print(len(gs))
+# print([g.dtype for g in gs])
 bufs = [Buffer(ps.device, g.size, g.dtype if isinstance(g.dtype, ImageDType) else g.dtype._base).ensure_allocated() for g in gs]
 
 t = cr(bufs, wait=True)

test/external/external_test_tk.py

@@ -0,0 +1,345 @@
+import unittest
+
+from tinygrad import Tensor, Device, dtypes, Context
+from tinygrad.engine.realize import ExecItem, get_runner
+
+from extra.thunder.tiny.tk import WARP_THREADS
+from extra.thunder.tiny.tk.kernel import Kernel
+from extra.thunder.tiny.tk.tiles import gl, st, rt, rv
+
+class TestTK(unittest.TestCase):
+  @unittest.skip("store from float rt is wrong")
+  def test_simple_matmul(self):
+    N = 32
+    BLOCK_SIZE = 16
+    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      c = gl((1, 1, N, N), dtypes.float32)
+      a = gl((1, 1, N, N), dtypes.bfloat16)
+      b = gl((1, 1, N, N), dtypes.bfloat16)
+
+      a_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      b_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      c_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      b_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      c_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      col, row = ker.blockIdx_x, ker.blockIdx_y
+
+      c_reg = warp.zero(c_reg)
+      for tile in ker.range(N // BLOCK_SIZE):
+        a_smem = warp.load(a_smem, a, (), (0, 0, row, tile), axis=2)
+        b_smem = warp.load(b_smem, b, (), (0, 0, tile, col), axis=2)
+
+        a_reg = warp.load(a_reg, a_smem)
+        b_reg = warp.load(b_reg, b_smem, transpose=True)
+
+        c_reg = warp.mma_AB(c_reg, a_reg, b_reg)
+      c_reg = ker.endrange()
+
+      c_smem = warp.store(c_smem, c_reg)
+      c = warp.store(c, c_smem, (0, 0, row, col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
+      b = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
+      c = Tensor.empty(1, 1, N, N, dtype="float32")
+      Tensor.realize(a, b, c)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (c, a, b)])
+    for _ in range(5): ei.run(wait=True)
+    c = c.float()
+
+    ref = a.matmul(b, dtype=dtypes.float32).float()
+
+    assert ref.allclose(c)
+
+  @unittest.skip("store from float rt is wrong")
+  def test_simple_matmul_transposed(self):
+    N = 32
+    BLOCK_SIZE = 16
+    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      c = gl((1, 1, N, N), dtypes.float32)
+      a = gl((1, 1, N, N), dtypes.bfloat16)
+      b = gl((1, 1, N, N), dtypes.bfloat16)
+
+      a_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      b_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      c_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      b_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      c_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      col, row = ker.blockIdx_x, ker.blockIdx_y
+
+      c_reg = warp.zero(c_reg)
+      for tile in ker.range(N // BLOCK_SIZE):
+        a_smem = warp.load(a_smem, a, (), (0, 0, row, tile), axis=2)
+        b_smem = warp.load(b_smem, b, (), (0, 0, col, tile), axis=2)
+
+        a_reg = warp.load(a_reg, a_smem)
+        b_reg = warp.load(b_reg, b_smem)
+
+        c_reg = warp.mma_ABt(c_reg, a_reg, b_reg)
+      c_reg = ker.endrange()
+
+      c_smem = warp.store(c_smem, c_reg)
+      c = warp.store(c, c_smem, (0, 0, row, col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
+      b = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
+      c = Tensor.empty(1, 1, N, N, dtype="float32")
+      Tensor.realize(a, b, c)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (c, a, b)])
+    for _ in range(5): ei.run(wait=True)
+    c = c.float()
+
+    ref = a.matmul(b.transpose(2, 3), dtype=dtypes.float32).float()
+
+    assert ref.allclose(c)
+
+  def test_load_store(self):
+    N = 32
+    BLOCK_SIZE = 16
+    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = gl((1, 1, N, N), dtypes.float32)
+      a = gl((1, 1, N, N), dtypes.float32)
+
+      a_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      col, row = ker.blockIdx_x, ker.blockIdx_y
+
+      a_smem = warp.load(a_smem, a, (), (0, 0, row, col), axis=2)
+      a_reg = warp.load(a_reg, a_smem)
+      b_reg = warp.copy(b_reg, a_reg)
+      b_smem = warp.store(b_smem, b_reg)
+      b = warp.store(b, b_smem, (0, 0, row, col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
+      b = Tensor.empty(1, 1, N, N, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float()
+
+    assert ref.allclose(b)
+
+  def test_max(self):
+    N = 16
+    BLOCK_SIZE = 16
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = gl((1, 1, N, N), dtypes.float32)
+      a = gl((1, 1, N, N), dtypes.float32)
+
+      a_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      max_reg = rv(BLOCK_SIZE, dtypes.float32, "ortho")
+
+      max_reg = warp.neg_inf(max_reg)
+
+      for tile_row in ker.range(N // BLOCK_SIZE):
+        for tile_col in ker.range(N // BLOCK_SIZE):
+          a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
+          a_reg = warp.load(a_reg, a_smem)
+          max_reg = warp.row_reduce(max_reg, a_reg, lambda a, b: a.maximum(b))
+        sum_reg = ker.endrange()
+
+        b_reg = warp.zero(b_reg).after(tile_row)
+        b_reg = warp.map(b_reg, lambda _, idx: sum_reg[idx[0], 0, (idx[2]%4)//2])
+        b_smem = warp.store(b_smem, b_reg)
+
+        for tile_col in ker.range(N // BLOCK_SIZE):
+          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
+      b = Tensor.empty(1, 1, N, N, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float().max(axis=3, keepdim=True).expand(a.shape)
+
+    assert ref.allclose(b)
+
+  def test_max_nonsquare(self):
+    N, M = 16, 64
+    BLOCK_N, BLOCK_M = 16, 64
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = gl((1, 1, N, M), dtypes.float32)
+      a = gl((1, 1, N, M), dtypes.float32)
+
+      a_smem = st((BLOCK_N, BLOCK_M), dtypes.float32)
+      b_smem = st((BLOCK_N, BLOCK_M), dtypes.float32)
+
+      a_reg = rt((BLOCK_N, BLOCK_M), dtypes.float32)
+      b_reg = rt((BLOCK_N, BLOCK_M), dtypes.float32)
+
+      max_reg = rv(BLOCK_N, dtypes.float32, "ortho")
+
+      max_reg = warp.zero(max_reg)
+
+      for tile_row in ker.range(N // BLOCK_N):
+        for tile_col in ker.range(M // BLOCK_M):
+          a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
+          a_reg = warp.load(a_reg, a_smem)
+          sum_reg = warp.row_reduce(max_reg, a_reg, lambda a, b: a.maximum(b))
+        sum_reg = ker.endrange()
+
+        b_reg = warp.zero(b_reg).after(tile_row)
+        b_reg = warp.map(b_reg, lambda _, idx: sum_reg[idx[0], 0, (idx[2]%4)//2])
+        b_smem = warp.store(b_smem, b_reg)
+
+        for tile_col in ker.range(M // BLOCK_M):
+          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, M, dtype="float32").contiguous()
+      b = Tensor.empty(1, 1, N, M, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float().max(axis=3, keepdim=True).expand(a.shape)
+
+    assert ref.allclose(b)
+
+  def test_sum(self):
+    N = 16
+    BLOCK_SIZE = 16
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = gl((1, 1, N, N), dtypes.float32)
+      a = gl((1, 1, N, N), dtypes.float32)
+
+      a_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      sum_reg = rv(BLOCK_SIZE, dtypes.float32, "ortho")
+
+      for tile_row in ker.range(N // BLOCK_SIZE):
+        sum_reg = warp.zero(sum_reg).after(tile_row)
+
+        for tile_col in ker.range(N // BLOCK_SIZE):
+          a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
+          a_reg = warp.load(a_reg, a_smem)
+          sum_reg = warp.row_reduce(sum_reg, a_reg, lambda a, b: a + b)
+        sum_reg = ker.endrange()
+
+        b_reg = warp.zero(b_reg).after(tile_row)
+        b_reg = warp.map(b_reg, lambda _, idx: sum_reg[idx[0], 0, (idx[2]%4)//2])
+        b_smem = warp.store(b_smem, b_reg)
+
+        for tile_col in ker.range(N // BLOCK_SIZE):
+          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
+      a = Tensor.arange(1 * 1 * N * N).reshape(1, 1, N, N).cast(dtypes.float32).contiguous()
+      b = Tensor.empty(1, 1, N, N, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float().sum(axis=3, keepdim=True).expand(a.shape)
+
+    assert ref.allclose(b)
+
+  def test_sum_nonsquare(self):
+    N, M = 16, 64
+    BLOCK_N, BLOCK_M = 16, 64
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = gl((1, 1, N, M), dtypes.float32)
+      a = gl((1, 1, N, M), dtypes.float32)
+
+      a_smem = st((BLOCK_N, BLOCK_M), dtypes.float32)
+      b_smem = st((BLOCK_N, BLOCK_M), dtypes.float32)
+
+      a_reg = rt((BLOCK_N, BLOCK_M), dtypes.float32)
+      b_reg = rt((BLOCK_N, BLOCK_M), dtypes.float32)
+
+      sum_reg = rv(BLOCK_N, dtypes.float32, "ortho")
+
+      sum_reg = warp.zero(sum_reg)
+
+      for tile_row in ker.range(N // BLOCK_N):
+        for tile_col in ker.range(M // BLOCK_M):
+          a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
+          a_reg = warp.load(a_reg, a_smem)
+          sum_reg = warp.row_reduce(sum_reg, a_reg, lambda a, b: a + b)
+        sum_reg = ker.endrange()
+
+        b_reg = warp.zero(b_reg).after(tile_row)
+        b_reg = warp.map(b_reg, lambda _, idx: sum_reg[idx[0], 0, (idx[2]%4)//2])
+        b_smem = warp.store(b_smem, b_reg)
+
+        for tile_col in ker.range(M // BLOCK_M):
+          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, M, dtype="float32").contiguous()
+      b = Tensor.empty(1, 1, N, M, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float().sum(axis=3, keepdim=True).expand(a.shape)
+
+    assert ref.allclose(b)
+
+if __name__ == "__main__":
+  unittest.main()

test/test_outerworld_transformer.py

@@ -0,0 +1,18 @@
+from tinygrad import Tensor, UOp
+from tinygrad.uop.ops import Ops, AxisType
+import unittest
+# this test is only focused on transformers and using range for the layers
+
+class TestOuterworldTransformer(unittest.TestCase):
+  def test_three_mats(self):
+    w = Tensor.empty(3, 1024, 1024)
+    inp = Tensor.empty(1, 1024)
+    i = UOp.range(3, -1, AxisType.OUTER)
+    inp_after = Tensor(inp.uop.after(i))
+    inp_gemm = inp_after@w[i]
+    inp = inp.uop.after(inp.uop.store(inp_gemm.uop).end(i)).contiguous()
+    inp = Tensor(inp)
+    inp.realize()
+
+if __name__ == "__main__":
+  unittest.main()

test/test_rangeify.py

@@ -300,6 +300,12 @@ class TestRangeify(unittest.TestCase):
     w2 = Tensor.empty(12, 8, 3, 3)
     x.conv2d(w1).conv2d(w2).realize()
 
+  def test_resnet_conv2d(self):
+    x = Tensor.empty(1, 8, 32, 32)
+    w1 = Tensor.empty(8, 8, 3, 3)
+    w2 = Tensor.empty(8, 8, 1, 1)
+    x.conv2d(w1).conv2d(w2).realize()
+
   def test_xception_conv2d(self):
     # NOTE: this fusion is bad, it's recomputing the inner many times
     x = Tensor.empty(1, 4, 32, 32)

test/test_schedule.py

@@ -1573,6 +1573,13 @@ class TestSchedule(unittest.TestCase):
   def test_conv2d(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4)
   def test_conv2d_fused(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4)
 
+  def test_resnet_conv2d(self):
+    x = Tensor.empty(1, 8, 32, 32)
+    w1 = Tensor.empty(8, 8, 3, 3)
+    w2 = Tensor.empty(8, 8, 1, 1)
+    out = x.conv2d(w1).conv2d(w2)
+    check_schedule(out, 2)
+
   @unittest.skipUnless(is_dtype_supported(dtypes.half) and is_dtype_supported(dtypes.ulong), "need half and ulong")
   def test_conv2d_half(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4, dtype=dtypes.half)
   @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")

tinygrad/codegen/__init__.py

@@ -136,7 +136,7 @@ def full_rewrite(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
   """
 
   full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
-  assert len(full_sink.ranges) == 0, "all ranges must end by the sink"
+  assert len(full_sink.ranges) == 0, f"all ranges must end by the sink, {full_sink.ranges}"
   lst = line_rewrite(linearize(full_sink), pm_linearize_cleanups)
   if SPEC: type_verify(lst, program_spec)
   return lst

tinygrad/codegen/gpudims.py

@@ -26,33 +26,35 @@ def _split_dims(dims, max_sizes):
   return tuple(_dims[:2] if _dims[2] == 1 else _dims[0] if _dims[1:3] == [1,1] else _dims)
 
 def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|None, reverse=False) -> list[UOp]:
-  if reverse: dims = dims[::-1]
-  # try to group first: (a, b, c, d) -> (ab, c, d)
-  limited = (grouped if (grouped := _group_dims(dims, max_sizes)) else dims) if max_sizes is not None else dims
-  # check if grouping failed
-  if max_sizes is not None and len(limited) > len(max_sizes): raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
-  # try to split up dims: (a,) -> (b, c)
-  if limited == dims: limited = _split_dims(dims, max_sizes) if max_sizes is not None else dims
-  ret = raw_idxs = [UOp(Ops.SPECIAL, dtypes.index, (sint_to_uop(s),), (f"{prefix}{i}")) for i,s in enumerate(limited)]
+  if reverse: return get_grouped_dims(prefix, dims[::-1], max_sizes)[::-1]
+  if max_sizes is None: limited = dims
+  else:
+    # try to group first: (a, b, c, d) -> (ab, c, d)
+    limited = grouped if (grouped := _group_dims(dims, max_sizes)) else dims
+    # check if grouping failed
+    if len(limited) > len(max_sizes): raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
+    # try to split up dims: (a,) -> (b, c)
+    if limited == dims: limited = _split_dims(dims, max_sizes)
+  raw_idxs = [UOp(Ops.SPECIAL, dtypes.index, (sint_to_uop(s),), (f"{prefix}{i}")) for i,s in enumerate(limited)]
   if len(limited) < len(dims):
     ret = []
-    if (contraction:=get_contraction(dims, limited)) is None: raise AssertionError(f"get_contraction should not be None {dims=} {limited=}")
+    if (contraction:=get_contraction(dims, limited)) is None: raise RuntimeError(f"get_contraction should not be None {dims=} {limited=}")
     for idx, contraction_group in zip(raw_idxs, contraction):
       for c in contraction_group[:-1]:
         ret.append(idx % dims[c])
         idx //= dims[c]
       ret.append(idx)
-  elif len(limited) > len(dims):
-    a, b = len(limited), len(dims)
-    if a == 2 and b == 1: ret = [raw_idxs[0] * limited[1] + raw_idxs[1]]
-    if a == 3 and b == 1: ret = [raw_idxs[0] * (limited[1] * limited[2]) + raw_idxs[1] * limited[2] + raw_idxs[2]]
-    if a == 3 and b == 2: ret = [raw_idxs[0] * limited[1] + raw_idxs[1], raw_idxs[2]]
+    return ret
+  elif (a:=len(limited)) > (b:=len(dims)):
+    if a == 2 and b == 1: return [raw_idxs[0] * limited[1] + raw_idxs[1]]
+    if a == 3 and b == 1: return [(raw_idxs[0] * limited[1] + raw_idxs[1]) * limited[2] + raw_idxs[2]]
+    if a == 3 and b == 2: return [raw_idxs[0] * limited[1] + raw_idxs[1], raw_idxs[2]]
   elif limited != dims:
     # Convert to 1D
     flat = raw_idxs[0]*limited[1]+raw_idxs[1] if len(dims) == 2 else raw_idxs[0]*(limited[1]*limited[2])+raw_idxs[1]*limited[2]+raw_idxs[2]
     # Get back original indices from 1D
-    ret = [flat//dims[1], flat%dims[1]] if len(dims) == 2 else [flat//(dims[2]*dims[1]), (flat//dims[2])%dims[1], flat%dims[2]]
-  return ret[::-1] if reverse else ret
+    return [flat//dims[1], flat%dims[1]] if len(dims) == 2 else [flat//(dims[2]*dims[1]), (flat//dims[2])%dims[1], flat%dims[2]]
+  return raw_idxs
 
 def add_gpudims(ctx:Renderer, s:UOp):
   if s.arg is None: return None

tinygrad/codegen/late/linearizer.py

@@ -41,7 +41,7 @@ def linearize(sink:UOp) -> list[UOp]:
   # number the uops in "ideal" order
   nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: priorities[x]+(x.tuplize if TUPLE_ORDER else ())))}
 
-  # then force then to be toposorted in as close to the ideal order as possible
+  # then force them to be toposorted in as close to the ideal order as possible
   heap = [(-nkey[sink], sink)]
   newlst = []
   while heap:

tinygrad/codegen/opt/heuristic.py

@@ -81,7 +81,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) <= 2048, False):
+  if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= (128 if NOLOCALS else 2048), False):
     for sz in [16]:
       try:
         k.apply_opt(Opt(OptOps.GROUPTOP, 0, sz))

tinygrad/codegen/opt/postrange.py

@@ -2,7 +2,8 @@ from __future__ import annotations
 import math, itertools
 from collections import defaultdict
 from typing import cast, Final
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp, axis_letters, axis_colors
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp
+from tinygrad.uop.ops import axis_letters, axis_colors, axis_to_pos
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes, ImageDType
 from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten
@@ -12,10 +13,6 @@ from tinygrad.renderer import Renderer
 
 remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
 
-# NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
-axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
-               AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
-
 class Scheduler:
   def __init__(self, ast:UOp, ren:Renderer):
     self.ast, self.ren = ast, ren
@@ -219,8 +216,7 @@ class Scheduler:
     return ret
 
   def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> None|list[UOp]:
-    reduceops = [x for x in self.ast.toposort() if x.op is Ops.REDUCE]
-    if not len(reduceops): raise KernelOptError("no reduce ops for TensorCore")
+    if not (reduceops := self.reduceops): raise KernelOptError("no reduce ops for TensorCore")
     reduceop = reduceops[0]
     if use_tensor_cores and reduceop is not None and reduceop.arg is Ops.ADD:
       mul = reduceop.src[0] if reduceop.src[0].op is not Ops.CAST else reduceop.src[0].src[0]
@@ -314,9 +310,10 @@ class Scheduler:
 
   # helpers for hand_coded_optimizations
   @property
+  def reduceops(self) -> list[UOp]: return [x for x in self.ast.backward_slice if x.op is Ops.REDUCE]
+  @property
   def reduceop(self) -> UOp|None:
-    red = [x for x in self.ast.backward_slice if x.op is Ops.REDUCE]
-    if not len(red): return None
+    if not (red := self.reduceops): return None
     return UOp(Ops.REDUCE_AXIS, red[0].dtype, red[0].src, (red[0].arg, ()))
   @property
   def bufs(self) -> list[UOp]: return [x for x in self.ast.toposort() if x.op is Ops.INDEX][::-1]

tinygrad/codegen/simplify.py

@@ -34,7 +34,6 @@ def simplify_merge_adjacent(u:UOp) -> UOp|None:
         # check if it simplifies
         if count_divmod(nidx) <= count_divmod(u):
           u = nidx
-          continue
   return u
 
 pm_simplify_ranges = PatternMatcher([

tinygrad/mixin/movement.py

@@ -322,7 +322,7 @@ class MovementMixin:
     """
     repeats = argfix(repeats, *args)
     base_shape = _align_left(self.shape, repeats)[0]
-    unsqueezed_shape = flatten([[1, s] for s in base_shape])
-    expanded_shape = flatten([[r, s] for r,s in zip(repeats, base_shape)])
+    unsqueezed_shape = flatten([[s] if r == 1 else [1, s] for r,s in zip(repeats, base_shape)])
+    expanded_shape = flatten([[s] if r == 1 else [r, s] for r,s in zip(repeats, base_shape)])
     final_shape = [r*s for r,s in zip(repeats, base_shape)]
     return self.reshape(unsqueezed_shape).expand(expanded_shape).reshape(final_shape)

tinygrad/runtime/autogen/nv/nv.py

@@ -7357,6 +7357,113 @@ class struct_c__SA_FWSECLIC_FRTS_CMD(Structure):
      ]
 
 FWSECLIC_FRTS_CMD = struct_c__SA_FWSECLIC_FRTS_CMD
+PCIEXPTBL_H = True # macro
+NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_BASE = 0x00 # macro
+NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_EXT = 0xE0 # macro
+PCI_EXP_ROM_SIGNATURE = 0xaa55 # macro
+PCI_EXP_ROM_SIGNATURE_NV = 0x4e56 # macro
+PCI_EXP_ROM_SIGNATURE_NV2 = 0xbb77 # macro
+def IS_VALID_PCI_ROM_SIG(sig):  # macro
+   return ((sig==0xaa55) or (sig==0x4e56) or (sig==0xbb77))
+OFFSETOF_PCI_EXP_ROM_SIG = 0x0 # macro
+OFFSETOF_PCI_EXP_ROM_NBSI_DATA_OFFSET = 0x16 # macro
+OFFSETOF_PCI_EXP_ROM_PCI_DATA_STRUCT_PTR = 0x18 # macro
+PCI_DATA_STRUCT_SIGNATURE = 0x52494350 # macro
+PCI_DATA_STRUCT_SIGNATURE_NV = 0x5344504E # macro
+PCI_DATA_STRUCT_SIGNATURE_NV2 = 0x53494752 # macro
+def IS_VALID_PCI_DATA_SIG(sig):  # macro
+   return ((sig==0x52494350) or (sig==0x5344504E) or (sig==0x53494752))
+# PCI_LAST_IMAGE = NVBIT ( 7 ) # macro
+PCI_ROM_IMAGE_BLOCK_SIZE = 512 # macro
+OFFSETOF_PCI_DATA_STRUCT_SIG = 0x0 # macro
+OFFSETOF_PCI_DATA_STRUCT_VENDOR_ID = 0x4 # macro
+OFFSETOF_PCI_DATA_STRUCT_LEN = 0xa # macro
+OFFSETOF_PCI_DATA_STRUCT_CLASS_CODE = 0xd # macro
+OFFSETOF_PCI_DATA_STRUCT_CODE_TYPE = 0x14 # macro
+OFFSETOF_PCI_DATA_STRUCT_IMAGE_LEN = 0x10 # macro
+OFFSETOF_PCI_DATA_STRUCT_LAST_IMAGE = 0x15 # macro
+NV_PCI_DATA_EXT_SIG = 0x4544504E # macro
+NV_PCI_DATA_EXT_REV_10 = 0x100 # macro
+NV_PCI_DATA_EXT_REV_11 = 0x101 # macro
+OFFSETOF_PCI_DATA_EXT_STRUCT_SIG = 0x0 # macro
+OFFSETOF_PCI_DATA_EXT_STRUCT_LEN = 0x6 # macro
+OFFSETOF_PCI_DATA_EXT_STRUCT_REV = 0x4 # macro
+OFFSETOF_PCI_DATA_EXT_STRUCT_SUBIMAGE_LEN = 0x8 # macro
+OFFSETOF_PCI_DATA_EXT_STRUCT_LAST_IMAGE = 0xa # macro
+OFFSETOF_PCI_DATA_EXT_STRUCT_FLAGS = 0xb # macro
+PCI_DATA_EXT_STRUCT_FLAGS_CHECKSUM_DISABLED = 0x04 # macro
+class struct__PCI_EXP_ROM_STANDARD(Structure):
+    pass
+
+struct__PCI_EXP_ROM_STANDARD._pack_ = 1 # source:False
+struct__PCI_EXP_ROM_STANDARD._fields_ = [
+    ('sig', ctypes.c_uint16),
+    ('reserved', ctypes.c_ubyte * 22),
+    ('pciDataStrucPtr', ctypes.c_uint16),
+    ('sizeOfBlock', ctypes.c_uint32),
+]
+
+PCI_EXP_ROM_STANDARD = struct__PCI_EXP_ROM_STANDARD
+PPCI_EXP_ROM_STANDARD = ctypes.POINTER(struct__PCI_EXP_ROM_STANDARD)
+class struct__PCI_EXP_ROM_NBSI(Structure):
+    pass
+
+struct__PCI_EXP_ROM_NBSI._pack_ = 1 # source:False
+struct__PCI_EXP_ROM_NBSI._fields_ = [
+    ('sig', ctypes.c_uint16),
+    ('reserved', ctypes.c_ubyte * 20),
+    ('nbsiDataOffset', ctypes.c_uint16),
+    ('pciDataStrucPtr', ctypes.c_uint16),
+    ('sizeOfBlock', ctypes.c_uint32),
+]
+
+PCI_EXP_ROM_NBSI = struct__PCI_EXP_ROM_NBSI
+PPCI_EXP_ROM_NBSI = ctypes.POINTER(struct__PCI_EXP_ROM_NBSI)
+class union__PCI_EXP_ROM(Union):
+    _pack_ = 1 # source:False
+    _fields_ = [
+    ('standard', PCI_EXP_ROM_STANDARD),
+    ('nbsi', PCI_EXP_ROM_NBSI),
+     ]
+
+PCI_EXP_ROM = union__PCI_EXP_ROM
+PPCI_EXP_ROM = ctypes.POINTER(union__PCI_EXP_ROM)
+class struct__PCI_DATA_STRUCT(Structure):
+    pass
+
+struct__PCI_DATA_STRUCT._pack_ = 1 # source:False
+struct__PCI_DATA_STRUCT._fields_ = [
+    ('sig', ctypes.c_uint32),
+    ('vendorID', ctypes.c_uint16),
+    ('deviceID', ctypes.c_uint16),
+    ('deviceListPtr', ctypes.c_uint16),
+    ('pciDataStructLen', ctypes.c_uint16),
+    ('pciDataStructRev', ctypes.c_ubyte),
+    ('classCode', ctypes.c_ubyte * 3),
+    ('imageLen', ctypes.c_uint16),
+    ('vendorRomRev', ctypes.c_uint16),
+    ('codeType', ctypes.c_ubyte),
+    ('lastImage', ctypes.c_ubyte),
+    ('maxRunTimeImageLen', ctypes.c_uint16),
+]
+
+PCI_DATA_STRUCT = struct__PCI_DATA_STRUCT
+PPCI_DATA_STRUCT = ctypes.POINTER(struct__PCI_DATA_STRUCT)
+class struct__NV_PCI_DATA_EXT_STRUCT(Structure):
+    pass
+
+struct__NV_PCI_DATA_EXT_STRUCT._pack_ = 1 # source:False
+struct__NV_PCI_DATA_EXT_STRUCT._fields_ = [
+    ('signature', ctypes.c_uint32),
+    ('nvPciDataExtRev', ctypes.c_uint16),
+    ('nvPciDataExtLen', ctypes.c_uint16),
+    ('subimageLen', ctypes.c_uint16),
+    ('privLastImage', ctypes.c_ubyte),
+    ('flags', ctypes.c_ubyte),
+]
+
+NV_PCI_DATA_EXT_STRUCT = struct__NV_PCI_DATA_EXT_STRUCT
+PNV_PCI_DATA_EXT_STRUCT = ctypes.POINTER(struct__NV_PCI_DATA_EXT_STRUCT)
 __all__ = \
     ['ACPI_DATA', 'ACPI_DSM_CACHE', 'ACPI_DSM_FUNCTION_COUNT',
     'ACPI_DSM_FUNCTION_CURRENT', 'ACPI_DSM_FUNCTION_GPS',
@@ -7480,12 +7587,16 @@ __all__ = \
     'NVDM_TYPE_UEFI_XTL_DEBUG_INTR',
     'NVGPU_ENGINE_CAPS_MASK_ARRAY_MAX', 'NVGPU_ENGINE_CAPS_MASK_BITS',
     'NV_ACPI_GENERIC_FUNC_COUNT',
+    'NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_BASE',
+    'NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_EXT',
     'NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_FLAGS_VERSION_AVAILABLE',
     'NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_FLAGS_VERSION_UNAVAILABLE',
     'NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V1',
     'NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V2',
     'NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V3',
     'NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V4',
+    'NV_PCI_DATA_EXT_REV_10', 'NV_PCI_DATA_EXT_REV_11',
+    'NV_PCI_DATA_EXT_SIG', 'NV_PCI_DATA_EXT_STRUCT',
     'NV_RPC_UPDATE_PDE_BAR_1', 'NV_RPC_UPDATE_PDE_BAR_2',
     'NV_RPC_UPDATE_PDE_BAR_INVALID', 'NV_RPC_UPDATE_PDE_BAR_TYPE',
     'NV_RPC_UPDATE_PDE_BAR_TYPE__enumvalues',
@@ -7775,8 +7886,31 @@ __all__ = \
     'NV_VGPU_PTE_64_INDEX_SHIFT', 'NV_VGPU_PTE_64_PAGE_SIZE',
     'NV_VGPU_PTE_64_SIZE', 'NV_VGPU_PTE_INDEX_MASK',
     'NV_VGPU_PTE_INDEX_SHIFT', 'NV_VGPU_PTE_PAGE_SIZE',
-    'NV_VGPU_PTE_SIZE', 'PACKED_REGISTRY_ENTRY',
-    'PACKED_REGISTRY_TABLE', 'REGISTRY_TABLE_ENTRY_TYPE_BINARY',
+    'NV_VGPU_PTE_SIZE', 'OFFSETOF_PCI_DATA_EXT_STRUCT_FLAGS',
+    'OFFSETOF_PCI_DATA_EXT_STRUCT_LAST_IMAGE',
+    'OFFSETOF_PCI_DATA_EXT_STRUCT_LEN',
+    'OFFSETOF_PCI_DATA_EXT_STRUCT_REV',
+    'OFFSETOF_PCI_DATA_EXT_STRUCT_SIG',
+    'OFFSETOF_PCI_DATA_EXT_STRUCT_SUBIMAGE_LEN',
+    'OFFSETOF_PCI_DATA_STRUCT_CLASS_CODE',
+    'OFFSETOF_PCI_DATA_STRUCT_CODE_TYPE',
+    'OFFSETOF_PCI_DATA_STRUCT_IMAGE_LEN',
+    'OFFSETOF_PCI_DATA_STRUCT_LAST_IMAGE',
+    'OFFSETOF_PCI_DATA_STRUCT_LEN', 'OFFSETOF_PCI_DATA_STRUCT_SIG',
+    'OFFSETOF_PCI_DATA_STRUCT_VENDOR_ID',
+    'OFFSETOF_PCI_EXP_ROM_NBSI_DATA_OFFSET',
+    'OFFSETOF_PCI_EXP_ROM_PCI_DATA_STRUCT_PTR',
+    'OFFSETOF_PCI_EXP_ROM_SIG', 'PACKED_REGISTRY_ENTRY',
+    'PACKED_REGISTRY_TABLE', 'PCIEXPTBL_H',
+    'PCI_DATA_EXT_STRUCT_FLAGS_CHECKSUM_DISABLED', 'PCI_DATA_STRUCT',
+    'PCI_DATA_STRUCT_SIGNATURE', 'PCI_DATA_STRUCT_SIGNATURE_NV',
+    'PCI_DATA_STRUCT_SIGNATURE_NV2', 'PCI_EXP_ROM',
+    'PCI_EXP_ROM_NBSI', 'PCI_EXP_ROM_SIGNATURE',
+    'PCI_EXP_ROM_SIGNATURE_NV', 'PCI_EXP_ROM_SIGNATURE_NV2',
+    'PCI_EXP_ROM_STANDARD', 'PCI_ROM_IMAGE_BLOCK_SIZE',
+    'PNV_PCI_DATA_EXT_STRUCT', 'PPCI_DATA_STRUCT', 'PPCI_EXP_ROM',
+    'PPCI_EXP_ROM_NBSI', 'PPCI_EXP_ROM_STANDARD',
+    'REGISTRY_TABLE_ENTRY_TYPE_BINARY',
     'REGISTRY_TABLE_ENTRY_TYPE_DWORD',
     'REGISTRY_TABLE_ENTRY_TYPE_STRING',
     'REGISTRY_TABLE_ENTRY_TYPE_UNKNOWN', 'RM_ENGINE_TYPE',
@@ -8343,6 +8477,8 @@ __all__ = \
     'struct_UpdateBarPde_v15_00',
     'struct_VIRTUAL_DISPLAY_GET_MAX_RESOLUTION_PARAMS',
     'struct_VIRTUAL_DISPLAY_GET_NUM_HEADS_PARAMS',
+    'struct__NV_PCI_DATA_EXT_STRUCT', 'struct__PCI_DATA_STRUCT',
+    'struct__PCI_EXP_ROM_NBSI', 'struct__PCI_EXP_ROM_STANDARD',
     'struct_alloc_object_FERMI_CONTEXT_SHARE_A_v04_00',
     'struct_alloc_object_FERMI_VASPACE_A_v03_00',
     'struct_alloc_object_GF100_DISP_SW_v03_00',
@@ -8595,7 +8731,7 @@ __all__ = \
     'union_NV2080_CTRL_FB_FS_INFO_QUERY_DATA_v26_04',
     'union_NV2080_CTRL_GRMGR_GR_FS_INFO_QUERY_DATA_v1A_1D',
     'union_NV2080_CTRL_INTERNAL_PFM_REQ_HNDLR_STATE_SYNC_DATA_type_v21_04',
-    'union_alloc_object_params_v25_08',
+    'union__PCI_EXP_ROM', 'union_alloc_object_params_v25_08',
     'union_alloc_object_params_v26_00',
     'union_alloc_object_params_v27_00',
     'union_alloc_object_params_v29_06', 'union_c__SA_GspFwWprMeta_0',

tinygrad/runtime/ops_nv.py

@@ -389,12 +389,12 @@ class NVKIface:
 
   def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, map_flags=0, cpu_addr=None, **kwargs) -> HCQBuffer:
     # Uncached memory is "system". Use huge pages only for gpu memory.
-    page_size = (4 << (12 if OSX else 10)) if uncached or host else ((2 << 20) if size >= (8 << 20) else (4 << (12 if OSX else 10)))
+    page_size = mmap.PAGESIZE if uncached or host else ((2 << 20) if size >= (8 << 20) else (mmap.PAGESIZE if MOCKGPU else 4 << 10))
     size = round_up(size, page_size)
-    va_addr = self._alloc_gpu_vaddr(size, alignment=page_size, force_low=cpu_access)
+    va_addr = self._alloc_gpu_vaddr(size, alignment=page_size, force_low=cpu_access) if (alloced:=cpu_addr is None) else cpu_addr
 
     if host:
-      va_addr = cpu_addr or FileIOInterface.anon_mmap(va_addr, size, mmap.PROT_READ|mmap.PROT_WRITE, MAP_FIXED|mmap.MAP_SHARED|mmap.MAP_ANONYMOUS, 0)
+      if alloced: va_addr = FileIOInterface.anon_mmap(va_addr, size, mmap.PROT_READ|mmap.PROT_WRITE, MAP_FIXED|mmap.MAP_SHARED|mmap.MAP_ANONYMOUS, 0)
 
       flags = (nv_gpu.NVOS02_FLAGS_PHYSICALITY_NONCONTIGUOUS << 4) | (nv_gpu.NVOS02_FLAGS_COHERENCY_CACHED << 12) \
             | (nv_gpu.NVOS02_FLAGS_MAPPING_NO_MAP << 30)
@@ -471,7 +471,7 @@ class PCIIface(PCIIfaceBase):
   def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, **kwargs) -> HCQBuffer:
     # Force use of huge pages for large allocations. NVDev will attempt to use huge pages in any case,
     # but if the size is not aligned, the tail will be allocated with 4KB pages, increasing TLB pressure.
-    page_size = (2 << 20) if size >= (8 << 20) and not uncached and not host else (4 << 10)
+    page_size = mmap.PAGESIZE if uncached or host else ((2 << 20) if size >= (8 << 20) else (4 << 10))
     return super().alloc(round_up(size, page_size), host=host, uncached=uncached, cpu_access=cpu_access, contiguous=contiguous, **kwargs)
 
   def setup_usermode(self): return 0xce000000, self.pci_dev.map_bar(bar=0, fmt='I', off=0xbb0000, size=0x10000)

tinygrad/runtime/ops_qcom.py

@@ -45,6 +45,10 @@ class QCOMSignal(HCQSignal):
       kgsl.IOCTL_KGSL_DEVICE_WAITTIMESTAMP_CTXTID(self.owner.fd, context_id=self.owner.ctx, timestamp=self.owner.last_cmd, timeout=0xffffffff)
 
 class QCOMComputeQueue(HWQueue):
+  def __init__(self, dev:QCOMDevice):
+    self.dev = dev
+    super().__init__()
+
   def __del__(self):
     if self.binded_device is not None: self.binded_device.allocator.free(self.hw_page, self.hw_page.size, BufferSpec(cpu_access=True, nolru=True))
 
@@ -54,7 +58,7 @@ class QCOMComputeQueue(HWQueue):
 
   def _cache_flush(self, write_back=True, invalidate=False, sync=True, memsync=False):
     # TODO: 7xx support.
-    if write_back: self.cmd(adreno.CP_EVENT_WRITE, adreno.CACHE_FLUSH_TS, *data64_le(QCOMDevice.dummy_addr), 0) # dirty cache write-back.
+    if write_back: self.cmd(adreno.CP_EVENT_WRITE, adreno.CACHE_FLUSH_TS, *data64_le(self.dev.dummy_addr), 0) # dirty cache write-back.
     if invalidate: self.cmd(adreno.CP_EVENT_WRITE, adreno.CACHE_INVALIDATE) # invalidate cache lines (following reads from RAM).
     if memsync: self.cmd(adreno.CP_WAIT_MEM_WRITES)
     if sync: self.cmd(adreno.CP_WAIT_FOR_IDLE)
@@ -65,7 +69,7 @@ class QCOMComputeQueue(HWQueue):
 
   def signal(self, signal:QCOMSignal, value=0, ts=False):
     self.cmd(adreno.CP_WAIT_FOR_IDLE)
-    if QCOMDevice.gpu_id < 700:
+    if self.dev.gpu_id[:2] < (7, 3):
       self.cmd(adreno.CP_EVENT_WRITE, qreg.cp_event_write_0(event=adreno.CACHE_FLUSH_TS, timestamp=ts),
                *data64_le(signal.timestamp_addr if ts else signal.value_addr), qreg.cp_event_write_3(value & 0xFFFFFFFF))
       self._cache_flush(write_back=True, invalidate=False, sync=False, memsync=False)
@@ -314,12 +318,9 @@ class QCOMAllocator(HCQAllocatorBase):
     self.dev._gpu_free(opaque)
 
 class QCOMDevice(HCQCompiled):
-  gpu_id: int = 0
-  dummy_addr: int = 0
-
   def __init__(self, device:str=""):
     self.fd = FileIOInterface('/dev/kgsl-3d0', os.O_RDWR)
-    QCOMDevice.dummy_addr = cast(int, self._gpu_alloc(0x1000).va_addr)
+    self.dummy_addr = cast(int, self._gpu_alloc(0x1000).va_addr)
 
     flags = kgsl.KGSL_CONTEXT_PREAMBLE | kgsl.KGSL_CONTEXT_PWR_CONSTRAINT | kgsl.KGSL_CONTEXT_NO_FAULT_TOLERANCE | kgsl.KGSL_CONTEXT_NO_GMEM_ALLOC \
               | kgsl.KGSL_CONTEXT_PRIORITY(getenv("QCOM_PRIORITY", 8)) | kgsl.KGSL_CONTEXT_PREEMPT_STYLE(kgsl.KGSL_CONTEXT_PREEMPT_STYLE_FINEGRAIN)
@@ -339,11 +340,14 @@ class QCOMDevice(HCQCompiled):
     # Load info about qcom device
     info = kgsl.struct_kgsl_devinfo()
     kgsl.IOCTL_KGSL_DEVICE_GETPROPERTY(self.fd, type=kgsl.KGSL_PROP_DEVICE_INFO, value=ctypes.addressof(info), sizebytes=ctypes.sizeof(info))
-    QCOMDevice.gpu_id = ((info.chip_id >> 24) & 0xFF) * 100 + ((info.chip_id >> 16) & 0xFF) * 10 + ((info.chip_id >>  8) & 0xFF)
-    if QCOMDevice.gpu_id >= 700: raise RuntimeError(f"Unsupported GPU: {QCOMDevice.gpu_id}")
+    self.gpu_id = (info.chip_id >> 24, (info.chip_id >> 16) & 0xFF, (info.chip_id >> 8) & 0xFF)
+
+    # a7xx start with 730x or 'Cxxx', a8xx starts 'Exxx'
+    if self.gpu_id[:2] >= (7, 3): raise RuntimeError(f"Unsupported GPU: chip_id={info.chip_id:#x}")
 
     compilers = [(QCOMRenderer, functools.partial(QCOMCompiler, device))]
-    super().__init__(device, QCOMAllocator(self), compilers, functools.partial(QCOMProgram, self), QCOMSignal, QCOMComputeQueue, None)
+    super().__init__(device, QCOMAllocator(self), compilers, functools.partial(QCOMProgram, self), QCOMSignal,
+                     functools.partial(QCOMComputeQueue, self), None)
 
   def _gpu_alloc(self, size:int, flags:int=0, uncached=False, fill_zeroes=False) -> HCQBuffer:
     flags |= kgsl.KGSL_MEMALIGN(alignment_hint:=12) | kgsl.KGSL_MEMFLAGS_USE_CPU_MAP

tinygrad/runtime/support/nv/ip.py

@@ -89,10 +89,18 @@ class NV_FLCN(NV_IP):
     self.prep_booter()
 
   def prep_ucode(self):
-    expansion_rom_off, bit_addr = {"GA": 0x16600, "AD": 0x14e00}[self.nvdev.chip_name[:2]], 0x1b0
-    vbios_bytes = bytes(array.array('I', self.nvdev.mmio[0x00300000//4:(0x00300000+0x98e00)//4]))
+    vbios_bytes, vbios_off = memoryview(bytes(array.array('I', self.nvdev.mmio[0x00300000//4:(0x00300000+0x100000)//4]))), 0
+    while True:
+      pci_blck = vbios_bytes[vbios_off + nv.OFFSETOF_PCI_EXP_ROM_PCI_DATA_STRUCT_PTR:].cast('H')[0]
+      imglen = vbios_bytes[vbios_off + pci_blck + nv.OFFSETOF_PCI_DATA_STRUCT_IMAGE_LEN:].cast('H')[0] * nv.PCI_ROM_IMAGE_BLOCK_SIZE
+      match vbios_bytes[vbios_off + pci_blck + nv.OFFSETOF_PCI_DATA_STRUCT_CODE_TYPE]:
+        case nv.NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_BASE: block_size = imglen
+        case nv.NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_EXT:
+          expansion_rom_off = vbios_off - block_size
+          break
+      vbios_off += imglen
 
-    bit_header = nv.BIT_HEADER_V1_00.from_buffer_copy(vbios_bytes[bit_addr:bit_addr + ctypes.sizeof(nv.BIT_HEADER_V1_00)])
+    bit_header = nv.BIT_HEADER_V1_00.from_buffer_copy(vbios_bytes[(bit_addr:=0x1b0):bit_addr + ctypes.sizeof(nv.BIT_HEADER_V1_00)])
     assert bit_header.Signature == 0x00544942, f"Invalid BIT header signature {hex(bit_header.Signature)}"
 
     for i in range(bit_header.TokenEntries):

tinygrad/runtime/support/system.py

@@ -285,6 +285,7 @@ class LNXPCIIfaceBase:
 class APLPCIIfaceBase(LNXPCIIfaceBase):
   def __init__(self, dev, dev_id, vendor, devices, bars, vram_bar, va_start, va_size):
     self.pci_dev, self.dev, self.vram_bar = APLPCIDevice(dev.__class__.__name__[:2], pcibus=f'usb4:{dev_id}', bars=bars), dev, vram_bar
+    assert (read_vendor:=self.pci_dev.read_config(0x00, 2)) == vendor, f"Vendor ID mismatch: expected {vendor:#x}, got {read_vendor:#x}"
   def map(self, b:HCQBuffer): raise RuntimeError(f"map failed: {b.owner} -> {self.dev}")
 
 PCIIfaceBase:type = APLPCIIfaceBase if OSX else LNXPCIIfaceBase

tinygrad/schedule/indexing.py

@@ -16,6 +16,9 @@ def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None:
   for s in rb.src:
     if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None
 
+def realize_store(ctx:dict[UOp, None], a:UOp) -> None:
+  if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
+
 def realize_assign(ctx:dict[UOp, None], a:UOp) -> None:
   if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
   # if it's a kernel, we don't realize it
@@ -30,6 +33,8 @@ pm_generate_realize_map = PatternMatcher([
   (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
   # realize ASSIGN and input to assign (might be optimized out)
   (UPat(Ops.ASSIGN, name="a"), realize_assign),
+  # realize STORE
+  (UPat(Ops.STORE, name="a"), realize_store),
 ])
 
 @dataclass(frozen=True)
@@ -50,13 +55,14 @@ class IndexingContext:
     # if a range has a 1 src, it's the same as UOp.const(dtypes.index, 0)
     return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(dtypes.index, 0)
 
+ops_allowed_after = (Ops.KERNEL, Ops.RANGE)
 def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
   if x.op in {Ops.BUFFERIZE, Ops.INDEX}: return None
-  if x.op is Ops.AFTER and x.src[1].op is Ops.KERNEL: return None
+  if x.op is Ops.AFTER and x.src[1].op in ops_allowed_after: return None
   new_srcs = []
   for s in x.src:
     new_src = s
-    if s.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT} or (s.op is Ops.AFTER and s.src[1].op is Ops.KERNEL):
+    if s.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT} or (s.op is Ops.AFTER and s.src[1].op in ops_allowed_after):
       if x in ctx.range_map: new_src = new_src.index(*ctx.range_map[x][0])
     elif s in ctx.realize_map:
       realized_ranges = ctx.realize_map[s]
@@ -176,7 +182,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
       # mark all ranges as ended
       assert rctx.realize_map[x] is None
       rctx.realize_map[x] = list(range(len(x.shape)))
-    elif x.op in {Ops.MSTACK, Ops.MSELECT}:
+    elif x.op in {Ops.MSTACK, Ops.MSELECT, Ops.END}:
       # treat MSTACK/MSELECT like SINK
       continue
     elif len(consumer_rngs) == 0:
@@ -239,7 +245,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     # if the EXPAND is used to inject a range, we don't mark it as ending_ranges. otherwise we do.
     # NOTE: this doesn't actually always end a range, but this is why convs are realized, so for now we need it
     if x.op is Ops.EXPAND and all(isinstance(y, int) or y.op is not Ops.RANGE for y in x.shape):
-      ending_ranges[x] = list(UOp.sink(*[ro for ri, ro in zip(rngs, out_rngs) if ri is not ro]).ranges.keys())
+      ending_ranges[x] += list(UOp.sink(*[ro for ri, ro in zip(rngs, out_rngs) if ri is not ro]).ranges.keys())
 
     # REDUCE_AXIS creates ranges for the axes it is reducing
     if x.op is Ops.REDUCE_AXIS:
@@ -247,15 +253,23 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
 
     if debug:
       realized_ranges = rctx.realize_map.get(x, None)
-      disp = []
-      for i, (ri, ro) in enumerate(zip([r.render() for r in rngs], [r.render() for r in out_rngs])):
-        rng = f"{ri}" if ri == ro else f"{ri} -> {ro}"
-        if realized_ranges is not None and i in realized_ranges: rng = colored(rng, "yellow")
-        disp.append("["+rng+"]")
-      print("***" if x in rctx.realize_map else "   ", len(consumer_map[x]), f"{str(x.op):20s}", ''.join(disp))
+      if x.op is Ops.RESHAPE or len(rngs) != len(out_rngs):
+        disp = render_ranges(rngs, realized=realized_ranges) + " -> " + render_ranges(out_rngs, realized=realized_ranges)
+      else:
+        disp = render_ranges(rngs, out_rngs, realized=realized_ranges)
+      print("***" if x in rctx.realize_map else "   ",
+            f"{len(consumer_map[x]):2d} {str(x.op):20s} {str(x.shape):35s} {len(ending_ranges[x]):2d}", disp)
 
     # assign to the range map. rngs are the input ranges, out_rngs are the output ranges, from the x op.
     rctx.range_map[x] = (rngs, out_rngs)
 
   tsink = graph_rewrite(tsink, pm_apply_rangeify, ctx=rctx, bottom_up=True, name="apply rangeify")
   return tsink, rctx
+
+def render_ranges(*rngs_list, realized) -> str:
+  disp = []
+  for i, rs in enumerate(zip(*[[r.render() for r in rngs] for rngs in rngs_list])):
+    rng = rs[0] if all_same(rs) else " -> ".join(rs)
+    if realized is not None and i in realized: rng = colored(rng, "yellow")
+    disp.append("["+rng+"]")
+  return ''.join(disp)

tinygrad/schedule/rangeify.py

@@ -396,6 +396,7 @@ def handle_after(ctx:LocalAddBufferContext, after:UOp):
   return buf
 
 def renumber_range(ctx:LocalAddBufferContext, r:UOp):
+  if r.arg[-1] == AxisType.OUTER: return None
   if r.tag != (): return None
   ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=None)
   ctx.range += 1
@@ -469,7 +470,10 @@ pm_add_range_tags = PatternMatcher([
 ])
 
 def split_store(ctx:list[UOp], x:UOp) -> UOp|None:
-  if len(x.ranges): return None
+  if len([r for r in x.ranges if r.arg[-1] != AxisType.OUTER]): return None
+
+  # ends of outer range don't go in kernels
+  if x.op is Ops.END and x.src[1].op is Ops.RANGE and x.src[1].arg[-1] == AxisType.OUTER: return None
 
   # local kernel rewrite
   lctx = LocalAddBufferContext()
@@ -538,6 +542,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
 
   tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding, name="symbolic+reduce_collapse")  # this does const folding
   tsink = graph_rewrite(tsink, pm_remove_bufferize, bottom_up=True, name="remove bufferize with cost function")
+  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding, name="symbolic+reduce_collapse pt 2")
   tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
 
   # rebuild the sink with all the BUFFERIZEs with tags, this is what's ending up in the tensor graph

tinygrad/tensor.py

@@ -2100,23 +2100,17 @@ class Tensor(OpMixin):
     noop, i_ = [None] * (self.ndim-len(k_)), self.shape[-len(k_):]
     assert all(resolve(d*(k-1)+1 <= i) for k,d,i in zip(k_,d_,i_)), "kernel size cannot be greater than actual input size"
     o_ = [ceildiv(i-d*(k-1), s) for i,d,k,s in zip(i_,d_,k_,s_)]
-    if any(resolve(k > s) for k,s in zip(k_,s_)) or any(d != 1 for d in d_):
-      # input size scaling factor to make sure shrink for stride is possible
-      f_ = [1 + int(resolve(o*s > (i - d*(k-1)))) for o,s,i,d,k in zip(o_,s_,i_,d_,k_)]
-      # # repeats such that we don't need padding
-      x = self.repeat([1]*len(noop) + [ceildiv(k*(i*f+d),i) for k,i,d,f in zip(k_,i_,d_,f_)])
-      # handle dilation
-      x = x.shrink(tuple(noop + [(0,k*(i*f+d)) for k,i,d,f in zip(k_,i_,d_,f_)])).reshape(noop + flatten((k,(i*f+d)) for k,i,d,f in zip(k_,i_,d_,f_)))
-      # handle stride
-      x = x.shrink(tuple(noop + flatten(((0,k), (0,o*s)) for k,o,s in zip(k_,o_,s_)))).reshape(noop + flatten((k,o,s) for k,o,s in zip(k_,o_,s_)))
-      x = x.shrink(tuple(noop + flatten(((0,k), (0,o), (0,1)) for k,o in zip(k_,o_)))).reshape(noop + flatten((k,o) for k,o in zip(k_,o_)))
-      # permute to move reduce to the end
-      return x.permute(*range(len(noop)), *[len(noop)+i*2+1 for i in range(len(i_))], *[len(noop)+i*2 for i in range(len(i_))])
-    # TODO: once the shapetracker can optimize well, remove this alternative implementation
-    x = self.pad(tuple(noop + [(0, max(0,o*s-i)) for i,o,s in zip(i_,o_,s_)])).shrink(tuple(noop + [(0,o*s) for o,s in zip(o_,s_)]))
-    x = x.reshape(noop + flatten(((o,s) for o,s in zip(o_,s_))))
-    x = x.shrink(tuple(noop + flatten(((0,o), (0,k)) for o,k in zip(o_,k_))))
-    return x.permute(*range(len(noop)), *[len(noop)+i*2 for i in range(len(i_))], *[len(noop)+i*2+1 for i in range(len(i_))])
+    # input size scaling factor to make sure shrink for stride is possible
+    f_ = [smax(1, ceildiv(o*s - d, i)) for o,s,i,d in zip(o_,s_,i_,d_)]
+    # repeats such that we don't need padding
+    x = self.repeat([1]*len(noop) + [ceildiv(k*(i*f+d),i) for k,i,d,f in zip(k_,i_,d_,f_)])
+    # handle dilation
+    x = x.shrink_to(noop + [k*(i*f+d) for k,i,d,f in zip(k_,i_,d_,f_)]).reshape(noop + flatten((k,(i*f+d)) for k,i,d,f in zip(k_,i_,d_,f_)))
+    # handle stride
+    x = x.shrink_to(noop + flatten((k,o*s) for k,o,s in zip(k_,o_,s_))).reshape(noop + flatten((k,o,s) for k,o,s in zip(k_,o_,s_)))
+    x = x.shrink_to(noop + flatten((k,o,1) for k,o in zip(k_,o_))).reshape(noop + flatten((k,o) for k,o in zip(k_,o_)))
+    # permute to move reduce to the end
+    return x.permute(*range(len(noop)), *[len(noop)+i*2+1 for i in range(len(i_))], *[len(noop)+i*2 for i in range(len(i_))])
 
   def _resolve_pool_pads(self, padding:int|Sequence[int], dims:int) -> Sequence[int]:
     if not isinstance(padding, int) and not (len(padding) == 2*dims or len(padding) == dims):
@@ -4166,18 +4160,18 @@ class Tensor(OpMixin):
     x, w = x.contiguous(), w.contiguous()
 
     # expand out
-    rcin_hi, rcin_lo = cin//4 if cin >= 4 else 1, 4 if cin >= 4 else 1
-    cout_expand = [groups//4 if cin == 1 else groups, 4 if cin == 1 else 1, rcout//4 if rcout >= 4 else 1, 4 if rcout >= 4 else 1]
+    rcin_hi, rcin_lo = (cin//4, 4) if cin >= 4 else (1, 1)
+    group_shape, rcout_expand = (groups//4, 4) if cin == 1 else (groups, 1), (rcout//4, 4) if rcout >= 4 else (1, 1)
     x = x.reshape(bs, iy, ix, groups, rcin_hi, rcin_lo)
     if cin_last: w = w.reshape(cout//4, H, rcin_hi, W, 4, rcin_lo)
     else: w = w.reshape(cout//4, H, rcin_hi, W, rcin_lo, 4).permute(0,1,2,3,5,4)
 
     # prepare input
     x = x.permute(0,3,4,5,1,2).pad(self._resolve_pool_pads(padding,2))._pool((H,W), stride, dilation)# -> (bs, groups, rcin_hi, rcin_lo, oy, ox, H, W)
-    x = x.permute(0,4,5,1,2,3,6,7).reshape(bs, (oy := x.shape[4]), (ox := x.shape[5]), *cout_expand[0:2], 1, 1, rcin_hi, rcin_lo, H, W)
+    x = x.permute(0,4,5,1,2,3,6,7).reshape(bs, (oy := x.shape[4]), (ox := x.shape[5]), *group_shape, 1, 1, rcin_hi, rcin_lo, H, W)
 
     # prepare weights
-    w = w.permute(0,4,2,5,1,3).reshape((1, 1, 1, *cout_expand, rcin_hi, rcin_lo, H, W))
+    w = w.permute(0,4,2,5,1,3).reshape((1, 1, 1, *group_shape, *rcout_expand, rcin_hi, rcin_lo, H, W))
 
     # the conv!
     ret = (x*w).cast(base_image_type((bs*oy, ox*cout//4, 4)) if IMAGE >= 2 else dtypes.float32).sum((-4, -3, -2, -1), dtype=dtype)

tinygrad/uop/ops.py

@@ -15,11 +15,15 @@ if TYPE_CHECKING:
 class AxisType(Enum):
   def __repr__(self): return str(self)
   GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto(); UNROLL = auto() # noqa: E702
-  THREAD = auto()
+  THREAD = auto(); OUTER = auto() # noqa: E702
 axis_letters = {AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u",
-                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
+                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r", AxisType.OUTER: ("O")}
 axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE",
-               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}
+               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta", AxisType.OUTER: "green"}
+
+# NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
+axis_to_pos = {AxisType.OUTER: -2, AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2,
+               AxisType.UPCAST: 3, AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
 
 range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3, Ops.END: 1}
 

tinygrad/viz/js/index.js

@@ -488,12 +488,15 @@ async function renderProfiler() {
     }
   }
 
-  canvas.addEventListener("click", e => {
+  const clickShape = (e) => {
     e.preventDefault();
     const foundRect = findRectAtPosition(e.clientX, e.clientY);
-    if (foundRect?.step != null && foundRect?.key == null) { return switchCtx(foundRect.ctx, foundRect.step); }
+    if (foundRect?.step != null && (foundRect?.key == null || e.type == "dblclick")) { return switchCtx(foundRect.ctx, foundRect.step); }
     if (foundRect?.key != focusedShape) { focusShape(foundRect); }
-  });
+  }
+  canvas.addEventListener("click", clickShape);
+
+  canvas.addEventListener("dblclick", clickShape);
 
   canvas.addEventListener("mousemove", e => {
     const foundRect = findRectAtPosition(e.clientX, e.clientY);
@@ -650,7 +653,7 @@ async function main() {
         u.li = list.appendChild(document.createElement("ul"));
         u.li.id = `step-${i}-${j}`;
         const p = u.li.appendChild(document.createElement("p"));
-        p.innerText = `${u.name}`+(u.match_count ? ` - ${u.match_count}` : '');
+        p.appendChild(colored(`${u.name}`+(u.match_count ? ` - ${u.match_count}` : '')));
         p.onclick = (e) => {
           e.stopPropagation();
           const subrewrites = getSubrewrites(e.currentTarget.parentElement);
@@ -661,7 +664,7 @@ async function main() {
       }
       for (const l of ul.querySelectorAll("ul > ul > p")) {
         const subrewrites = getSubrewrites(l.parentElement);
-        if (subrewrites.length > 0) { l.innerText += ` (${subrewrites.length})`; l.parentElement.classList.add("has-children"); }
+        if (subrewrites.length > 0) { l.appendChild(d3.create("span").text(` (${subrewrites.length})`).node()); l.parentElement.classList.add("has-children"); }
       }
     }
     return setState({ currentCtx:-1 });
@@ -716,8 +719,8 @@ async function main() {
         }
       }
       metadata.appendChild(tabulate(ret.summary.map(s => {
-        const div = d3.create("div").style("background", cycleColors(colorScheme.CATEGORICAL, s.idx)).style("width", "24px").style("height", "100%");
-        return [s.label.trim(), div.node()];
+        const div = d3.create("div").style("background", cycleColors(colorScheme.CATEGORICAL, s.idx)).style("width", "100%").style("height", "100%");
+        return [s.label.trim(), div.text(s.value.toLocaleString()).node()];
       })).node());
     } else root.appendChild(codeBlock(ret.src, ret.lang || "txt"));
     return document.querySelector("#custom").replaceChildren(root);

tinygrad/viz/serve.py

@@ -204,14 +204,20 @@ def load_sqtt(profile:list[ProfileEvent]) -> None:
     return ctxs.append({"name":"Counters", "steps":[step]})
   try: from extra.sqtt.roc import decode
   except Exception: return err("DECODER IMPORT ISSUE")
-  try:
-    rctx = decode(profile)
-    steps = [{"name":str(x[0]), "depth":0, "data":{"rows":[(e.inst, e.time, e.time-x[1][i-1].time if i else 0, e.dur, e.stall,
-                                                            str(e.typ).split("_")[-1]) for i,e in enumerate(x[1])],
-                                              "cols":["Instruction", "Clk", "Wait", "Duration", "Stall", "Type"], "summary":[]},
-              "query":f"/render?ctx={len(ctxs)}&step={i}&fmt=counters"} for i,x in enumerate(rctx.inst_execs.items())]
-    if not steps: return err("EMPTY SQTT OUTPUT", f"{len(sqtt_events)} SQTT events recorded, none got decoded")
+  try: rctx = decode(profile)
   except Exception: return err("DECODER ERROR")
+  if not rctx.inst_execs: return err("EMPTY SQTT OUTPUT", f"{len(sqtt_events)} SQTT events recorded, none got decoded")
+  steps:list[dict] = []
+  for name,waves in rctx.inst_execs.items():
+    if (r:=ref_map.get(name)): name = ctxs[r]["name"]
+    steps.append({"name":name, "depth":0, "query":f"/render?ctx={len(ctxs)}&step={len(steps)}&fmt=counters",
+                  "data":{"src":trace.keys[r].ret.src if r else name, "lang":"cpp"}})
+    for w in waves:
+      rows = [(e.inst, e.time, e.time-(w.insts[i-1].time if i else 0), e.dur, e.stall, str(e.typ).split("_")[-1]) for i,e in enumerate(w.insts)]
+      summary = [{"label":"Total Cycles", "value":w.insts[-1].time-w.insts[0].time if w.insts else 0}, {"label":"CU", "value":w.cu},
+                 {"label":"SIMD", "value":w.simd}]
+      steps.append({"name":f"Wave {w.wave_id}", "depth":1, "query":f"/render?ctx={len(ctxs)}&step={len(steps)}&fmt=counters",
+                    "data":{"rows":rows, "cols":["Instruction", "Clk", "Wait", "Duration", "Stall", "Type"], "summary":summary}})
   ctxs.append({"name":"Counters", "steps":steps})
 
 def get_profile(profile:list[ProfileEvent]) -> bytes|None: