IQ.Pilot Prebuilt Release @ ab07000

tinygrad_repo/extra/amdpci/am_smi.py

@@ -0,0 +1,416 @@
+#!/usr/bin/env python3
+
+import time, mmap, sys, shutil, os, glob, subprocess, argparse, collections
+from tinygrad.helpers import DEBUG, colored, ansilen
+from tinygrad.runtime.autogen import libc
+from tinygrad.runtime.autogen.am import am
+from tinygrad.runtime.support.hcq import MMIOInterface
+from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager, AMPageTableEntry
+from tinygrad.runtime.support.am.ip import AM_SOC, AM_GMC, AM_IH, AM_PSP, AM_SMU, AM_GFX, AM_SDMA
+
+def bold(s): return f"\033[1m{s}\033[0m"
+
+def trim(s:str, length:int) -> str:
+  if len(s) > length: return s[:length-3] + "..."
+  return s
+
+def pad(x:str, length:int) -> str:
+  if len(x) < length: return x + " " * (length - len(x))
+  return x
+
+def color_temp(temp):
+  if temp >= 87: return colored(f"{temp:>3}", "red")
+  elif temp >= 80: return colored(f"{temp:>3}", "yellow")
+  return f"{temp:>3}"
+
+def color_voltage(voltage): return colored(f"{voltage/1000:>5.3f}V", "cyan")
+
+def draw_bar(percentage, width=40, fill='|', empty=' ', opt_text='', color='cyan'):
+  percentage = 0.0 if percentage != percentage else percentage  # NaN guard
+  percentage = max(0.0, min(1.0, float(percentage)))
+  filled_width = int(width * percentage)
+  if not opt_text: opt_text = f'{percentage*100:.1f}%'
+
+  bar = fill * filled_width + empty * (width - filled_width)
+  if opt_text and len(opt_text) <= len(bar): bar = (bar[:-len(opt_text)] + opt_text)
+  bar = colored(bar[:filled_width], color) + bar[filled_width:]
+  return f'[{bar}]'
+
+def same_line(strs:list[list[str]|None], split=8) -> list[str]:
+  strs = [s for s in strs if s is not None]
+  if len(strs) == 0: return []
+
+  ret = []
+  max_width_in_block = [max(ansilen(line) for line in block) for block in strs]
+  max_height = max(len(block) for block in strs)
+  for i in range(max_height):
+    line = []
+    for bid, block in enumerate(strs):
+      if i < len(block): line.append(block[i] + (' ' * (split + max_width_in_block[bid] - ansilen(block[i])) if bid != len(strs) - 1 else ''))
+      else: line.append(' ' * (split + max_width_in_block[bid]))
+    ret.append(' '.join(line))
+  return ret
+
+def get_bar0_size(pcibus):
+  resource_file = f"/sys/bus/pci/devices/{pcibus}/resource"
+  if not os.path.exists(resource_file): raise FileNotFoundError(f"Resource file not found: {resource_file}")
+
+  with open(resource_file, "r") as f: lines = f.readlines()
+  bar0_info = lines[0].split()
+  if len(bar0_info) < 3: raise ValueError("Unexpected resource file format for BAR0.")
+
+  start_hex, end_hex, _flags = bar0_info
+  return int(end_hex, 16) - int(start_hex, 16) + 1
+
+class AMSMI(AMDev):
+  def __init__(self, pcibus, vram_bar:MMIOInterface, doorbell_bar:MMIOInterface, mmio_bar:MMIOInterface):
+    self.pcibus, self.devfmt = pcibus, pcibus
+    self.vram, self.doorbell64, self.mmio = vram_bar, doorbell_bar, mmio_bar
+    self.pci_state = self.read_pci_state()
+    if self.pci_state == "D0": self._init_from_d0()
+
+  def _init_from_d0(self):
+    self._run_discovery()
+    self._build_regs()
+
+    if self.reg("regSCRATCH_REG7").read() != AMDev.Version:
+      raise Exception(f"Unsupported AM version: {self.reg('regSCRATCH_REG7').read():x}")
+
+    self.is_booting = True
+    self.init_sw(smi_dev=True)
+    self.partial_boot = True # do not init anything
+
+  def read_pci_state(self):
+    with open(f"/sys/bus/pci/devices/{self.pcibus}/power_state", "r") as f: return f.read().strip().rstrip()
+
+class SMICtx:
+  def __init__(self):
+    self.devs = []
+    self.opened_pcidevs = []
+    self.opened_pci_resources = {}
+    self.prev_lines_cnt = 0
+    self.prev_terminal_width = 0
+    self.prev_terminal_height = 0
+    self.prev_metrics = {}
+
+    remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:", "Processing accelerators:"]
+    lspci = subprocess.check_output(["lspci"]).decode("utf-8").splitlines()
+    self.lspci = {l.split()[0]: l.split(" ", 1)[1] for l in lspci}
+    for k,v in self.lspci.items():
+      for part in remove_parts: self.lspci[k] = self.lspci[k].replace(part, "").strip().rstrip()
+
+  def _smuq10_round(self, v:int) -> int:
+    v = int(v)
+    return (v + 512) >> 10  # SMUQ10_ROUND
+
+  def _fmt_kb(self, kb:int) -> str:
+    kb = int(kb)
+    if kb < 1024: return f"{kb}KB"
+    mb = kb / 1024.0
+    if mb < 1024: return f"{mb:.1f}MB"
+    gb = mb / 1024.0
+    if gb < 1024: return f"{gb:.2f}GB"
+    tb = gb / 1024.0
+    return f"{tb:.2f}TB"
+
+  def _open_am_device(self, pcibus):
+    if pcibus not in self.opened_pci_resources:
+      bar_fds = {bar: os.open(f"/sys/bus/pci/devices/{pcibus}/resource{bar}", os.O_RDWR | os.O_SYNC) for bar in [0, 2, 5]}
+      bar_size = {0: get_bar0_size(pcibus), 2: os.fstat(bar_fds[2]).st_size, 5: os.fstat(bar_fds[5]).st_size}
+
+      def map_pci_range(bar, fmt='B'):
+        return MMIOInterface(libc.mmap(0, bar_size[bar], mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, bar_fds[bar], 0), bar_size[bar], fmt)
+      self.opened_pci_resources[pcibus] = (map_pci_range(0), None, map_pci_range(5, 'I'))
+
+    try:
+      self.devs.append(AMSMI(pcibus, *self.opened_pci_resources[pcibus]))
+    except Exception as e:
+      if DEBUG >= 2: print(f"Failed to open AM device {pcibus}: {e}")
+      return
+
+    self.opened_pcidevs.append(pcibus)
+    if DEBUG >= 2: print(f"Opened AM device {pcibus}")
+
+  def rescan_devs(self):
+    pattern = os.path.join('/tmp', 'am_*.lock')
+    for d in [f[8:-5] for f in glob.glob(pattern)]:
+      if d.startswith("usb"): continue
+      if d not in self.opened_pcidevs:
+        self._open_am_device(d)
+
+    for d in self.devs:
+      if d.read_pci_state() != d.pci_state:
+        d.pci_state = d.read_pci_state()
+        if d.pci_state == "D0": d._init_from_d0()
+        os.system('clear')
+
+      if d.pci_state == "D0" and d.reg("regSCRATCH_REG7").read() != AMDev.Version:
+        self.devs.remove(d)
+        self.opened_pcidevs.remove(d.pcibus)
+        os.system('clear')
+        if DEBUG >= 2: print(f"Removed AM device {d.pcibus}")
+
+  def collect(self):
+    tables = {}
+    for dev in self.devs:
+      match dev.ip_ver[am.MP1_HWIP]:
+        case (13,0,6): table_t = dev.smu.smu_mod.MetricsTableV0_t
+        case (13,0,12): table_t = dev.smu.smu_mod.MetricsTable_t
+        case _: table_t = dev.smu.smu_mod.SmuMetricsExternal_t
+      tables[dev] = dev.smu.read_table(table_t, dev.smu.smu_mod.SMU_TABLE_SMU_METRICS) if dev.pci_state == "D0" else None
+    return tables
+
+  def _pick_nonzero_avg(self, vals) -> int:
+    xs = [x for x in vals if x > 0]
+    return int(sum(xs) / len(xs)) if xs else 0
+
+  def get_gfx_activity(self, dev, metrics):
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6)|(13,0,12): return max(0, min(100, self._smuq10_round(metrics.SocketGfxBusy)))
+      case _: return metrics.SmuMetrics.AverageGfxActivity
+
+  def get_mem_activity(self, dev, metrics):
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6)|(13,0,12): return max(0, min(100, self._smuq10_round(metrics.DramBandwidthUtilization)))
+      case _: return metrics.SmuMetrics.AverageUclkActivity
+
+  def get_temps(self, dev, metrics, compact=False):
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6)|(13,0,12):
+        temps = {
+          "Hotspot": self._smuq10_round(metrics.MaxSocketTemperature),
+          "HBM": self._smuq10_round(metrics.MaxHbmTemperature),
+          "VR": self._smuq10_round(metrics.MaxVrTemperature),
+        }
+        if compact: return {k: temps[k] for k in ("Hotspot", "HBM") if temps.get(k, 0) != 0}
+        return {k: v for k, v in temps.items() if v != 0}
+      case _:
+        temps_keys = [(k, name) for k, name in dev.smu.smu_mod.TEMP_e.items()
+                      if k < dev.smu.smu_mod.TEMP_COUNT and metrics.SmuMetrics.AvgTemperature[k] != 0]
+        if compact: temps_keys = [(k, name) for k, name in temps_keys if k in (dev.smu.smu_mod.TEMP_HOTSPOT, dev.smu.smu_mod.TEMP_MEM)]
+        return {name: metrics.SmuMetrics.AvgTemperature[k] for k, name in temps_keys}
+
+  def get_voltage(self, dev, metrics, compact=False):
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6)|(13,0,12): return {}
+      case _:
+        voltage_keys = [(k, name) for k, name in dev.smu.smu_mod.SVI_PLANE_e.items()
+                        if k < dev.smu.smu_mod.SVI_PLANE_COUNT and metrics.SmuMetrics.AvgVoltage[k] != 0]
+        return {name: metrics.SmuMetrics.AvgVoltage[k] for k, name in voltage_keys}
+
+  def get_busy_threshold(self, dev):
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (14, 0, 2): return 5
+      case _: return 15
+
+  def get_gfx_freq(self, dev, metrics):
+    if metrics is None: return 0
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6)|(13,0,12): return self._smuq10_round(metrics.GfxclkFrequency[0])
+      case _:
+        return metrics.SmuMetrics.AverageGfxclkFrequencyPostDs if self.get_gfx_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
+               metrics.SmuMetrics.AverageGfxclkFrequencyPreDs
+
+  def get_mem_freq(self, dev, metrics):
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6)|(13,0,12): return self._smuq10_round(metrics.UclkFrequency)
+      case _:
+        return metrics.SmuMetrics.AverageMemclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
+               metrics.SmuMetrics.AverageMemclkFrequencyPreDs
+
+  def get_fckl_freq(self, dev, metrics):
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6)|(13,0,12): return self._smuq10_round(metrics.FclkFrequency)
+      case _:
+        return metrics.SmuMetrics.AverageFclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
+               metrics.SmuMetrics.AverageFclkFrequencyPreDs
+
+  def get_fan_rpm_pwm(self, dev, metrics):
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6)|(13,0,12): return None, None
+      case _: return metrics.SmuMetrics.AvgFanRpm, metrics.SmuMetrics.AvgFanPwm
+
+  def get_power(self, dev, metrics):
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.MaxSocketPowerLimit)
+      case (13,0,12): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.SocketPowerLimit)
+      case _: return metrics.SmuMetrics.AverageSocketPower, metrics.SmuMetrics.dGPU_W_MAX
+
+  def get_throttle_info(self, dev, metrics):
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6)|(13,0,12):
+        throttle_fields = [('ProchotResidencyAcc', 'Prochot'), ('PptResidencyAcc', 'PPT'),
+                           ('SocketThmResidencyAcc', 'Socket Thm'), ('VrThmResidencyAcc', 'VR Thm'), ('HbmThmResidencyAcc', 'HBM Thm')]
+        prev = self.prev_metrics.get(dev.pcibus)
+        active = []
+        if prev is not None:
+          acc_delta = metrics.AccumulationCounter - prev.AccumulationCounter
+          if acc_delta > 0:
+            for field, name in throttle_fields:
+              delta = getattr(metrics, field) - getattr(prev, field)
+              if delta > 0 and (pct := min(100, (delta * 100 + acc_delta // 2) // acc_delta)) > 0: active.append((name, pct))
+        return active
+      case _:
+        smu_mod = dev.smu.smu_mod
+        throttler_names = {getattr(smu_mod, a): a[len('THROTTLER_'):-len('_BIT')]
+                           for a in dir(smu_mod) if a.startswith('THROTTLER_') and a.endswith('_BIT')}
+        active = []
+        for i, pct in enumerate(metrics.SmuMetrics.ThrottlingPercentage):
+          if pct > 0: active.append((throttler_names.get(i, f"UNK_{i}"), int(pct)))
+        return active
+
+  def get_mem_usage(self, dev):
+    usage = 0
+    pt_stack = [dev.mm.root_page_table]
+    while len(pt_stack) > 0:
+      pt = pt_stack.pop()
+      for i in range(512):
+        entry = pt.entries[i]
+
+        if (entry & am.AMDGPU_PTE_VALID) == 0: continue
+        if pt.lv < am.AMDGPU_VM_PDB0 and not dev.gmc.is_pte_huge_page(pt.lv, entry):
+          pt_stack.append(AMPageTableEntry(dev, dev.xgmi2paddr(entry & 0x0000FFFFFFFFF000), lv=pt.lv+1))
+          continue
+        if (entry & am.AMDGPU_PTE_SYSTEM) != 0: continue
+        usage += (1 << ((9 * (3-pt.lv)) + 12))
+    return usage
+
+  def draw(self, once):
+    terminal_width, terminal_height = shutil.get_terminal_size()
+    if not once and (self.prev_terminal_width != terminal_width or self.prev_terminal_height != terminal_height):
+      os.system('clear')
+    self.prev_terminal_width, self.prev_terminal_height = terminal_width, terminal_height
+
+    padding = 8
+    col_size = (terminal_width) // 2 - padding - 2
+    activity_line_width = 50 if terminal_width > 170 else \
+                         (30 if terminal_width > 130 else \
+                         (16 if terminal_width > 92 else \
+                         max(0, terminal_width - 77)))
+
+    dev_metrics = self.collect()
+    dev_content = []
+    for dev, metrics in dev_metrics.items():
+      if dev.pci_state != "D0":
+        dev_content.append([f"{colored('(sleep)', 'yellow')} {bold(dev.pcibus)}: {trim(self.lspci[dev.pcibus[5:]], col_size - 20)}"] +
+                           [pad(f"PCI State: {dev.pci_state}", col_size)])
+        continue
+
+      mem_used = self.get_mem_usage(dev)
+      mem_total = dev.vram_size
+      mem_fmt = f"{mem_used/1024**3:.1f}/{mem_total/1024**3:.1f}G"
+
+      device_line = [f"{bold(dev.pcibus)} {trim(self.lspci[dev.pcibus[5:]], col_size - 20)}"] + [pad("", col_size)]
+      activity_line = [f"GFX Activity {draw_bar(self.get_gfx_activity(dev, metrics) / 100, activity_line_width)}"] \
+                    + [f"MEM Activity {draw_bar(self.get_mem_activity(dev, metrics) / 100, activity_line_width)}"] \
+                    + [f"MEM Usage    {draw_bar(mem_used / mem_total, activity_line_width, opt_text=mem_fmt)}"] \
+
+      throttle_info = self.get_throttle_info(dev, metrics)
+      if throttle_info:
+        throttle_text = colored(', '.join(f"{name} {pct}%" for name, pct in throttle_info), "red")
+      else:
+        throttle_text = colored("None", "green")
+      activity_line += [f"Throttle     {throttle_text}" + " " * (activity_line_width + 2)]
+
+      temps_data, temps_data_compact = self.get_temps(dev, metrics), self.get_temps(dev, metrics, compact=True)
+      temps_table = ["=== Temps (°C) ==="] + [f"{name:<16}: {color_temp(val)}" for name, val in temps_data.items()]
+      temps_table_compact = ["Temps (°C):" + '/'.join([f"{color_temp(val)} {name}" for name, val in temps_data_compact.items()])]
+
+      fan_rpm, fan_pwm = self.get_fan_rpm_pwm(dev, metrics)
+      power_table = ["=== Power ==="]
+      power_table += ["Fan: N/A"] if fan_rpm is None or fan_pwm is None else [f"Fan Speed: {fan_rpm} RPM", f"Fan Power: {fan_pwm}%"]
+
+      total_power, max_power = self.get_power(dev, metrics)
+      if max_power > 0:
+        power_line = [f"Power: " + draw_bar(total_power / max_power, 16, opt_text=f"{total_power}/{max_power}W")]
+        power_line_compact = [f"Power:       " + draw_bar(total_power / max_power, activity_line_width, opt_text=f"{total_power}/{max_power}W")]
+      else:
+        power_line = ["Power: N/A"]
+        power_line_compact = ["Power: N/A"]
+
+      voltage_data = self.get_voltage(dev, metrics)
+      voltage_table = None if not voltage_data else (["=== Voltages ==="] + [f"{name:<20}: {color_voltage(voltage)}" for name, voltage in voltage_data.items()])
+
+      gfx_freq = self.get_gfx_freq(dev, metrics)
+      mclk_freq = self.get_mem_freq(dev, metrics)
+      fclk_freq = self.get_fckl_freq(dev, metrics)
+      frequency_table = ["=== Frequencies ===", f"GFXCLK: {gfx_freq:>4} MHz", f"FCLK  : {fclk_freq:>4} MHz", f"MCLK  : {mclk_freq:>4} MHz"]
+
+      if self.prev_terminal_width >= 231:
+        power_table += power_line
+        if voltage_table is not None: power_table += [""] + voltage_table
+        activity_line += [""]
+      elif self.prev_terminal_width >= 171:
+        power_table += power_line + [""] + frequency_table
+        activity_line += [""]
+        frequency_table = None
+      elif self.prev_terminal_width >= 121:
+        temps_table = None
+        activity_line += power_line_compact
+      else:
+        temps_table = None
+        power_table = None
+        frequency_table = None
+        activity_line += power_line_compact
+
+      dev_content.append(device_line + activity_line + same_line([temps_table, power_table, frequency_table]))
+
+    self.prev_metrics = {dev.pcibus: m for dev, m in dev_metrics.items() if m is not None}
+
+    raw_text = 'AM Monitor'.center(terminal_width) + "\n" + "=" * terminal_width + "\n\n"
+    for i in range(0, len(dev_content), 2):
+      if i + 1 < len(dev_content): raw_text += '\n'.join(same_line([dev_content[i], dev_content[i+1]], split=padding))
+      else: raw_text += '\n'.join(dev_content[i])
+      if i + 2 < len(dev_content): raw_text += "\n" + "=" * terminal_width + "\n\n"
+
+    sys.stdout.write(f'\033[{self.prev_lines_cnt}A')
+    sys.stdout.flush()
+    print(raw_text)
+
+    self.prev_lines_cnt = len(raw_text.splitlines()) + 2
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--list", action="store_true", help="Run once and exit")
+  parser.add_argument("--pids", action="store_true", help="Print pids for all AM devices")
+  parser.add_argument("--kill", action="store_true", help="Kill all pids associated with AM devices. Valid only with --pids")
+  parser.add_argument("--dev", type=str, default=None, help="PCI bus ID of the AM device to monitor (e.g., 0000:01:00.0)")
+  args = parser.parse_args()
+
+  if args.pids:
+    for dev in glob.glob('/tmp/am_*.lock'):
+      if args.dev and not dev.endswith(f"{args.dev}.lock"):
+        print(f"{dev[8:-5]}: skipping")
+        continue
+
+      try:
+        if args.kill:
+          stopped_pids = collections.defaultdict(int)
+          while True:
+            try: pid = subprocess.check_output(['sudo', 'lsof', '-t', dev]).decode('utf-8').split('\n')[0]
+            except subprocess.CalledProcessError: break
+            if stopped_pids[pid] > 0: time.sleep(0.1)
+            if stopped_pids[pid] == 64:
+              print(f"{dev[8:-5]}: can't stop process {pid}, exitting")
+              exit(1)
+
+            print(f"{dev[8:-5]}: killing process {pid}")
+            os.system(f'sudo pkill -g -9 {pid}')
+            stopped_pids[pid] += 1
+        else:
+          pid = subprocess.check_output(['sudo', 'lsof', dev]).decode('utf-8').strip().split('\n')[1].split()[1]
+          print(f"{dev[8:-5]}: {pid}")
+      except subprocess.CalledProcessError:
+        print(f"{dev[8:-5]}: no process found")
+    sys.exit(0)
+
+  try:
+    if not args.list: os.system('clear')
+    smi_ctx = SMICtx()
+    while True:
+      smi_ctx.rescan_devs()
+      smi_ctx.draw(args.list)
+      if args.list: break
+      time.sleep(1)
+  except KeyboardInterrupt:
+    print("Exiting...")

tinygrad_repo/extra/amdpci/hive_reset.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+
+import os
+from tinygrad.helpers import Context
+from tinygrad.runtime.support.system import System, PCIDevice
+from tinygrad.runtime.support.hcq import FileIOInterface
+from tinygrad.runtime.support.am.amdev import AMDev
+
+if __name__ == "__main__":
+  gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1, 0x75a0])])
+  for gpu in gpus:
+    drv_path = f"/sys/bus/pci/devices/{gpu}/driver"
+    if FileIOInterface.exists(drv_path) and os.path.basename(os.readlink(drv_path)) == "amdgpu":
+      raise RuntimeError(f"amdgpu is bound to {gpu}. Stopping...")
+  pcidevs = [PCIDevice("AM", gpu) for gpu in gpus]
+  amdevs = []
+  with Context(DEBUG=2):
+    for pcidev in pcidevs:
+      amdevs.append(AMDev(pcidev, reset_mode=True))
+    for amdev in amdevs: amdev.smu.mode1_reset()

tinygrad_repo/extra/amdpci/proclogs.py

@@ -0,0 +1,96 @@
+import re, ctypes, sys, importlib
+from tinygrad.helpers import getenv
+
+from tinygrad.runtime.support.am.amdev import AMDev, AMRegister
+
+class GFXFake:
+  def __init__(self): self.xccs = 8
+
+class AMDFake(AMDev):
+  def __init__(self, pci_dev):
+    self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
+    self.vram, self.doorbell64, self.mmio = self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')
+    self._run_discovery()
+    self._build_regs()
+
+    self.gfx = GFXFake()
+
+amdev = importlib.import_module("tinygrad.runtime.support.am.amdev")
+amdev.AMDev = AMDFake
+from tinygrad.runtime.ops_amd import PCIIface
+
+def parse_amdgpu_logs(log_content, register_names=None, register_objects=None, *, only_xcc0: bool = False):
+  register_map = register_names or {}
+  register_objs = register_objects or {}
+
+  def replace_register(match):
+    reg = match.group(1)
+    return f"Reading register {register_map.get(int(reg, 16), reg)}"
+
+  processed_log = re.sub(r'Reading register (0x[0-9a-fA-F]+)', replace_register, log_content)
+
+  def replace_register_2(match):
+    reg = match.group(1)
+    return f"Writing register {register_map.get(int(reg, 16), reg)}"
+
+  processed_log = re.sub(r'Writing register (0x[0-9a-fA-F]+)', replace_register_2, processed_log)
+
+  # remove timing prefix
+  processed_log = re.sub(r'^\[\s*\d+(?:\.\d+)?\]\s*', '', processed_log, flags=re.MULTILINE)
+
+  # decode register values into field dicts
+  def decode_value(match):
+    reg_name = match.group(1)
+    xcc_part = match.group(2)  # "xcc=0 " or ""
+    val_str = match.group(3)
+    val = int(val_str, 16)
+
+    reg_obj = register_objs.get(reg_name)
+    if reg_obj is not None and reg_obj.fields:
+      fields = reg_obj.decode(val)
+      # show raw for unaccounted bits
+      accounted = 0
+      for name, (start, end) in reg_obj.fields.items():
+        accounted |= (((1 << (end - start + 1)) - 1) << start)
+      unaccounted = val & ~accounted
+      parts = {k: v for k, v in fields.items() if v != 0}
+      if unaccounted: parts['_raw_unaccounted'] = hex(unaccounted)
+      return f"register {reg_name}, {xcc_part}with value {val_str} {parts}"
+    return match.group(0)
+
+  processed_log = re.sub(r'register (reg\w+), ((?:xcc=\d+ )?)with value (0x[0-9a-fA-F]+)', decode_value, processed_log)
+
+  # keep only xcc=0 lines (but keep lines with no xcc at all)
+  if only_xcc0:
+    kept = []
+    for line in processed_log.splitlines(True):
+      if "xcc=" not in line or re.search(r'\bxcc=0\b', line): kept.append(line)
+    processed_log = "".join(kept)
+
+  return processed_log
+
+def main():
+  only_xcc0 = bool(getenv("ONLY_XCC0", 0))
+
+  reg_names = {}
+  reg_objs = {}
+  dev = PCIIface(None, 0)
+  for x, y in dev.dev_impl.__dict__.items():
+    if isinstance(y, AMRegister):
+      for xcc, addr in y.addr.items():
+        reg_names[addr] = f"{x}, xcc={xcc}"
+      reg_objs[x] = y
+
+  with open(sys.argv[1], 'r') as f:
+    log_content = f.read()
+
+  processed_log = parse_amdgpu_logs(log_content, reg_names, reg_objs, only_xcc0=only_xcc0)
+
+  with open(sys.argv[2], 'w') as f:
+    f.write(processed_log)
+
+if __name__ == '__main__':
+  if len(sys.argv) != 3:
+    print("Usage: <input_file_path> <output_file_path>")
+    sys.exit(1)
+  main()

tinygrad_repo/extra/amdpci/setup_python_cap.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+PYTHON_PATH=$(readlink -f $(which python3))
+sudo setcap 'cap_dac_override,cap_sys_rawio,cap_sys_admin,cap_ipc_lock=ep' $PYTHON_PATH

tinygrad_repo/extra/amdpci/setup_vfio.sh

@@ -0,0 +1,2 @@
+#!/bin/bash
+sudo modprobe vfio-pci disable_idle_d3=1

tinygrad_repo/extra/archprobe.py

@@ -0,0 +1,154 @@
+# copying the kernels from https://github.com/microsoft/ArchProbe into Python
+import numpy as np
+import pickle
+from tinygrad.runtime.ops_cl import CLProgram, CLBuffer
+from tinygrad import dtypes
+from tqdm import trange, tqdm
+from matplotlib import pyplot as plt
+
+tests = {}
+def register_test(fxn):
+  tests[fxn.__name__] = fxn
+
+def warp_size2(nthread):
+  prg = """__kernel void warp_size2(
+    __global float* src,
+    __global int* dst,
+    const int niter,
+    const int prime_number
+  ) {
+    int drain = 0;
+    for (int j = 0; j < niter; ++j) {
+      drain += j / prime_number;
+      barrier(0);
+    }
+    dst[get_local_id(0)] = drain;
+  }"""
+  src_buf = CLBuffer(1, dtypes.float32)
+  dst_buf = CLBuffer(1, dtypes.int32)
+  cl = CLProgram("warp_size2", prg, argdtypes=[None, None, np.int32, np.int32])
+  return min([cl([nthread, 1024, 1], [nthread, 1, 1], src_buf, dst_buf, 10, 3, wait=True) for _ in range(5)])*1e9
+
+@register_test
+def test_warp_size():
+  return [(nthread, warp_size2(nthread)) for nthread in trange(1,256)]
+
+def reg_count(nthread, ngrp, nreg):
+  reg_declr = ''.join([f"float reg_data{i} = (float)niter + {i};\n" for i in range(nreg)])
+  reg_comp = ''.join([f"reg_data{i} *= {(i-1)%nreg};\n" for i in range(nreg)])
+  reg_reduce = ''.join([f"out_buf[{i}] = reg_data{i};\n" for i in range(nreg)])
+  prg = f"""__kernel void reg_count(
+    __global float* out_buf,
+    __private const int niter
+  ) {{
+    {reg_declr}
+    int i = 0;
+    for (; i < niter; ++i) {{
+      {reg_comp}
+    }}
+    i = i >> 31;
+    {reg_reduce}
+  }}"""
+  out_buf = CLBuffer(1, dtypes.float32)
+  cl = CLProgram("reg_count", prg, argdtypes=[None, np.int32])
+  return min([cl([nthread, ngrp, 1], [nthread, 1, 1], out_buf, 20, wait=True) for _ in range(10)])*1e9
+
+@register_test
+def test_reg_count(nthread=1, ngrp=1):
+  base = reg_count(nthread, ngrp, 1)
+  return [(nreg, (reg_count(nthread, ngrp, nreg)-base)/nreg) for nreg in trange(4, 513, 4)]
+
+def buf_cache_hierarchy_pchase(ndata, stride=1, NCOMP=1, steps=65536):
+  ndata //= NCOMP*4  # ptr size
+  prg = f"""__kernel void buf_cache_hierarchy_pchase(
+    __global int{str(NCOMP) if NCOMP > 1 else ''}* src,
+    __global int* dst,
+    const int niter
+  ) {{
+    int idx = 0;
+    for (int i = 0; i < niter; ++i) {{
+      idx = src[idx]{'.x' if NCOMP > 1 else ''};
+    }}
+    *dst = idx;
+  }}"""
+  idx_buf = np.zeros(ndata*NCOMP, dtype=np.int32)
+  for i in range(ndata): idx_buf[i*NCOMP] = (i + stride) % ndata
+  in_buf = CLBuffer.fromCPU(idx_buf)
+  out_buf = CLBuffer(1, dtypes.int32)
+  cl = CLProgram("buf_cache_hierarchy_pchase", prg, argdtypes=[None, None, np.int32])
+  return min([cl([1, 1, 1], [1, 1, 1], in_buf, out_buf, steps, wait=True)/steps for _ in range(5)])*1e9
+
+@register_test
+def test_memory_latency():
+  # requires cacheline < 16
+  szs = [int(1.3**x) for x in range(20, 70)]
+  return [(ndata, buf_cache_hierarchy_pchase(ndata, NCOMP=16, steps=128*1024)) for ndata in tqdm(szs)]
+
+@register_test
+def test_cacheline_size():
+  # TODO: this buffer must be at least 2x the L1 cache for this test to work
+  return [(stride, buf_cache_hierarchy_pchase(4*65536, stride, steps=65536)) for stride in trange(1,64)]
+
+def cl_read(sz, niter=1):
+  prg = f"""__kernel void copy(
+    __global float4* src,
+    __global float* dst) {{
+      int gid = get_global_id(0);
+      if (src[gid].x == 99+get_global_id(1)) *dst = 1;
+  }}"""
+
+  in_buf = CLBuffer(sz//4, dtypes.float32)
+  out_buf = CLBuffer(1, dtypes.float32)
+  cl = CLProgram("copy", prg)
+  # NOTE: if nay of the niters form a local group, this is wrong
+  return min([cl([sz//16, niter, 1], [1, 1, 1], in_buf, out_buf, wait=True) for _ in range(10)])*1e9
+
+@register_test
+def test_read_bandwidth():
+  szs = list(range(128*1024, 20*1024*1024, 128*1024))
+  NITER = 8
+  base = cl_read(16, niter=NITER)
+  return [(sz, (sz*NITER)/(cl_read(sz, niter=NITER)-base)) for sz in tqdm(szs)]
+
+
+def gflops(niter=4, nroll=4, ngroups=4096):
+  NCOMP = 8
+  prg = f"""__kernel void gflops(
+    __global float* out_buf
+    ) {{
+      float{NCOMP} x = (float{NCOMP})({",".join(f"get_local_id(0)+{i}" for i in range(NCOMP))});
+      float{NCOMP} y = (float{NCOMP})({",".join(f"get_local_id(1)+{i}" for i in range(NCOMP))});
+
+      for (int i = 0; i < {niter}; i++) {{
+        {''.join(['x = mad(y, y, x); y = mad(x, x, y);'+chr(10)]*nroll)}
+      }}
+
+      out_buf[get_global_id(0) >> 31] = {'+'.join(f"y.s{'0123456789abcdef'[i]}" for i in range(NCOMP))};
+  }}"""
+  out_buf = CLBuffer(1, dtypes.float32)
+  cl = CLProgram("gflops", prg, options="-cl-mad-enable -cl-fast-relaxed-math")
+  FLOPS = NCOMP*2*2 * niter * nroll * ngroups * 32
+  # NOTE: if nay of the niters form a local group, this is wrong
+  return FLOPS/(min([cl([32, ngroups, 1], [32, 1, 1], out_buf, wait=True) for _ in range(10)])*1e9)
+
+@register_test
+def test_gflops():
+  return [(niter, gflops(niter=niter, nroll=32)) for niter in trange(1, 32, 1)]
+
+if __name__ == "__main__":
+  cache = {}
+  #cache = pickle.load(open("/tmp/cache.pkl", "rb"))
+  #tests = {"test_cacheline_size": tests["test_cacheline_size"]}
+  plt.figure(figsize=(16, 9))
+  for i,(k,test) in enumerate(tests.items()):
+    print(f"running {k}")
+    plt.subplot(2, (len(tests)+1)//2, i+1)
+    plt.title(k)
+    if k == "test_memory_latency": plt.xscale('log')
+    if k not in cache: cache[k] = test()
+    plt.plot(*zip(*cache[k]))
+  #pickle.dump(cache, open("/tmp/cache.pkl", "wb"))
+
+  plt.tight_layout(pad=0.5)
+  plt.savefig("/tmp/results.png")
+  plt.show()

tinygrad_repo/extra/bench_log.py

@@ -0,0 +1,111 @@
+import time, atexit, uuid
+from enum import Enum
+
+from tinygrad.device import Device
+from tinygrad.helpers import DEBUG, ContextVar, getenv, GlobalCounters
+
+BENCHMARK_LOG = ContextVar("BENCHMARK_LOG", "")
+
+if BENCHMARK_LOG:
+  from influxdb_client_3 import InfluxDBClient3, Point, WriteOptions, write_client_options
+  from influxdb_client_3.write_client.client.write_api import WriteType
+
+class BenchEvent(Enum):
+  LOAD_WEIGHTS = "load_weights"
+  STEP = "step"
+  FULL = "full"
+  MLPERF_INIT = "mlperf_init"
+  MLPERF_RUN = "mlperf_run"
+class InstantBenchEvent(Enum):
+  GFLOPS = "gflops"
+
+_events = {}
+def clear_events():
+  for event in BenchEvent:
+    _events[event] = {"wall": [], "kernel": []}
+  for event in InstantBenchEvent:
+    _events[event] = []
+clear_events()
+
+class WallTimeEvent:
+  def __init__(self, event:BenchEvent):
+    self.event = event
+  def __enter__(self):
+    self.start = time.monotonic()
+    return self
+  def __exit__(self, *_):
+    self.time = time.monotonic() - self.start
+    _events[self.event]["wall"].append(self.time)
+    return False
+
+class KernelTimeEvent:
+  def __init__(self, event:BenchEvent):
+    if DEBUG < 2:
+      raise Exception("KernelTimeEvent should only be used in DEBUG >= 2")
+    self.event = event
+  def __enter__(self):
+    self.start = GlobalCounters.time_sum_s
+    return self
+  def __exit__(self, *_):
+    _events[self.event]["kernel"].append(GlobalCounters.time_sum_s - self.start)
+    return False
+
+def log_event_instant(event:InstantBenchEvent, value:float):
+  _events[event].append(value)
+
+if BENCHMARK_LOG:
+  INFLUXDB_HOST = getenv("INFLUXDB_HOST", "")
+  INFLUXDB_ORG = getenv("INFLUXDB_ORG", "tiny")
+  INFLUXDB_TOKEN = getenv("INFLUXDB_TOKEN", "")
+
+  def _create_point(run_id, i, attempt, ref, commit, name, value, run):
+    point = Point(BENCHMARK_LOG.value).tag("id", run_id).tag("index", i)
+    point = point.tag("device", Device.DEFAULT)
+    point = point.tag("attempt", attempt).tag("ref", ref).tag("commit", commit)
+    point = point.field(name, value).field("x", run)
+    return point
+
+  @atexit.register
+  def write_events():
+    # see if there are any events to write
+    have_events = False
+    for event in _events:
+      if isinstance(event, BenchEvent):
+        for event_type, values in _events[event].items():
+          if len(values) > 0:
+            have_events = True
+      else:
+        if len(_events[event]) > 0:
+          have_events = True
+    if not have_events:
+      return
+
+    # pull from github envvars
+    ref = getenv("GITHUB_REF_NAME", "")
+    commit = getenv("GITHUB_SHA", "")
+    run = getenv("GITHUB_RUN_NUMBER", "")
+    attempt = getenv("GITHUB_RUN_ATTEMPT", "")
+
+    points = []
+    for event in _events:
+      run_id = str(uuid.uuid4())
+      if isinstance(event, BenchEvent):
+        for event_type, values in _events[event].items():
+          for i, value in enumerate(values):
+            point = _create_point(run_id, i, attempt, ref, commit, f"{event.value}_{event_type}", value, run)
+            points.append(point)
+      else:
+        for i, value in enumerate(_events[event]):
+          point = _create_point(run_id, i, attempt, ref, commit, event.value, value, run)
+          points.append(point)
+
+    write_options = WriteOptions(write_type=WriteType.synchronous, retry_interval=5000, max_retries=5, max_retry_delay=30000, exponential_base=2)
+    wco = write_client_options(write_options=write_options)
+    with InfluxDBClient3(
+        host=INFLUXDB_HOST,
+        org=INFLUXDB_ORG,
+        token=INFLUXDB_TOKEN,
+        auth_scheme="Basic",
+        database="benchmarks",
+        write_client_options=wco) as client:
+      client.write(points)

tinygrad_repo/extra/cl_android.sh

@@ -0,0 +1,4 @@
+# source extra/cl_android.sh
+export LD_LIBRARY_PATH=/data/data/com.termux/files/usr/lib:/system/vendor/lib64
+export LD_PRELOAD=/system/vendor/lib64/libOpenCL.so
+

tinygrad_repo/extra/datasets/.gitignore

@@ -0,0 +1,4 @@
+imagenet
+imagenet_bak
+mnist
+open-images-v6TEST

tinygrad_repo/extra/datasets/__init__.py

@@ -0,0 +1,43 @@
+import os, gzip, tarfile, pickle
+import numpy as np
+from tinygrad import Tensor, dtypes
+from tinygrad.helpers import fetch
+
+def fetch_mnist(tensors=False):
+  parse = lambda file: np.frombuffer(gzip.open(file).read(), dtype=np.uint8).copy()
+  BASE_URL = "https://storage.googleapis.com/cvdf-datasets/mnist/"   # http://yann.lecun.com/exdb/mnist/ lacks https
+  X_train = parse(fetch(f"{BASE_URL}train-images-idx3-ubyte.gz"))[0x10:].reshape((-1, 28*28)).astype(np.float32)
+  Y_train = parse(fetch(f"{BASE_URL}train-labels-idx1-ubyte.gz"))[8:].astype(np.int8)
+  X_test = parse(fetch(f"{BASE_URL}t10k-images-idx3-ubyte.gz"))[0x10:].reshape((-1, 28*28)).astype(np.float32)
+  Y_test = parse(fetch(f"{BASE_URL}t10k-labels-idx1-ubyte.gz"))[8:].astype(np.int8)
+  if tensors: return Tensor(X_train).reshape(-1, 1, 28, 28), Tensor(Y_train), Tensor(X_test).reshape(-1, 1, 28, 28), Tensor(Y_test)
+  else: return X_train, Y_train, X_test, Y_test
+
+cifar_mean = [0.4913997551666284, 0.48215855929893703, 0.4465309133731618]
+cifar_std = [0.24703225141799082, 0.24348516474564, 0.26158783926049628]
+
+def fetch_cifar():
+  X_train = Tensor.empty(50000, 3*32*32, device=f'disk:/tmp/cifar_train_x', dtype=dtypes.uint8)
+  Y_train = Tensor.empty(50000, device=f'disk:/tmp/cifar_train_y', dtype=dtypes.int64)
+  X_test = Tensor.empty(10000, 3*32*32, device=f'disk:/tmp/cifar_test_x', dtype=dtypes.uint8)
+  Y_test = Tensor.empty(10000, device=f'disk:/tmp/cifar_test_y', dtype=dtypes.int64)
+
+  if not os.path.isfile("/tmp/cifar_extracted"):
+    def _load_disk_tensor(X, Y, db_list):
+      idx = 0
+      for db in db_list:
+        x, y = db[b'data'], np.array(db[b'labels'])
+        assert x.shape[0] == y.shape[0]
+        X[idx:idx+x.shape[0]].assign(x)
+        Y[idx:idx+x.shape[0]].assign(y)
+        idx += x.shape[0]
+      assert idx == X.shape[0] and X.shape[0] == Y.shape[0]
+
+    print("downloading and extracting CIFAR...")
+    fn = fetch('https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz')
+    tt = tarfile.open(fn, mode='r:gz')
+    _load_disk_tensor(X_train, Y_train, [pickle.load(tt.extractfile(f'cifar-10-batches-py/data_batch_{i}'), encoding="bytes") for i in range(1,6)])
+    _load_disk_tensor(X_test, Y_test, [pickle.load(tt.extractfile('cifar-10-batches-py/test_batch'), encoding="bytes")])
+    open("/tmp/cifar_extracted", "wb").close()
+
+  return X_train, Y_train, X_test, Y_test

tinygrad_repo/extra/datasets/fake_imagenet_from_mnist.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python3
+import pathlib, json
+from tinygrad.helpers import trange
+from extra.datasets import fetch_mnist
+from PIL import Image
+import numpy as np
+from multiprocessing import Pool
+
+X_train, Y_train, X_test, Y_test = fetch_mnist()
+
+def act(arg):
+  (basedir, i, train) = arg
+  if train:
+    img = np.uint8(X_train[i]).reshape(28, 28)
+    nm = f"train/{Y_train[i]}/{i}.jpg"
+  else:
+    img = np.uint8(X_test[i]).reshape(28, 28)
+    nm = f"val/{Y_test[i]}/{i}.jpg"
+  Image.fromarray(img).resize((224, 224)).convert('RGB').save(basedir / nm)
+
+def create_fake_mnist_imagenet(basedir:pathlib.Path):
+  print(f"creating mock MNIST dataset at {basedir}")
+  basedir.mkdir(exist_ok=True)
+
+  with (basedir / "imagenet_class_index.json").open('w') as f:
+    f.write(json.dumps({str(i):[str(i), str(i)] for i in range(10)}))
+
+  for i in range(10):
+    (basedir / f"train/{i}").mkdir(parents=True, exist_ok=True)
+    (basedir / f"val/{i}").mkdir(parents=True, exist_ok=True)
+
+  def gen(train):
+    for idx in trange(X_train.shape[0] if train else X_test.shape[0]):
+      yield (basedir, idx, train)
+
+  with Pool(64) as p:
+    for _ in p.imap_unordered(act, gen(True)): pass
+    for _ in p.imap_unordered(act, gen(False)): pass
+
+if __name__ == "__main__":
+  create_fake_mnist_imagenet(pathlib.Path("./mnist"))

tinygrad_repo/extra/datasets/imagenet.py

@@ -0,0 +1,91 @@
+# for imagenet download prepare.sh and run it
+import glob, random, json, math
+import numpy as np
+from PIL import Image
+import functools, pathlib
+from tinygrad.helpers import diskcache, getenv
+
+@functools.cache
+def get_imagenet_categories():
+  ci = json.load(open(BASEDIR / "imagenet_class_index.json"))
+  return {v[0]: int(k) for k,v in ci.items()}
+
+if getenv("MNISTMOCK"):
+  BASEDIR = pathlib.Path(__file__).parent / "mnist"
+
+  @functools.cache
+  def get_train_files():
+    if not BASEDIR.exists():
+      from extra.datasets.fake_imagenet_from_mnist import create_fake_mnist_imagenet
+      create_fake_mnist_imagenet(BASEDIR)
+
+    if not (files:=glob.glob(p:=str(BASEDIR / "train/*/*"))): raise FileNotFoundError(f"No training files in {p}")
+    return files
+else:
+  BASEDIR = pathlib.Path(__file__).parent / "imagenet"
+
+  @diskcache
+  def get_train_files():
+    if not (files:=glob.glob(p:=str(BASEDIR / "train/*/*"))): raise FileNotFoundError(f"No training files in {p}")
+    return files
+
+@functools.cache
+def get_val_files():
+  if not (files:=glob.glob(p:=str(BASEDIR / "val/*/*"))): raise FileNotFoundError(f"No validation files in {p}")
+  return files
+
+def image_resize(img, size, interpolation):
+  w, h = img.size
+  w_new = int((w / h) * size) if w > h else size
+  h_new = int((h / w) * size) if h > w else size
+  return img.resize([w_new, h_new], interpolation)
+
+def rand_flip(img):
+  if random.random() < 0.5:
+    img = np.flip(img, axis=1).copy()
+  return img
+
+def center_crop(img):
+  rescale = min(img.size) / 256
+  crop_left = (img.width - 224 * rescale) / 2.0
+  crop_top = (img.height - 224 * rescale) / 2.0
+  img = img.resize((224, 224), Image.BILINEAR, box=(crop_left, crop_top, crop_left + 224 * rescale, crop_top + 224 * rescale))
+  return img
+
+# we don't use supplied imagenet bounding boxes, so scale min is just min_object_covered
+# https://github.com/tensorflow/tensorflow/blob/e193d8ea7776ef5c6f5d769b6fb9c070213e737a/tensorflow/core/kernels/image/sample_distorted_bounding_box_op.cc
+def random_resized_crop(img, size, scale=(0.10, 1.0), ratio=(3/4, 4/3)):
+  w, h = img.size
+  area = w * h
+
+  # Crop
+  random_solution_found = False
+  for _ in range(100):
+    aspect_ratio = random.uniform(ratio[0], ratio[1])
+    max_scale = min(min(w * aspect_ratio / h, h / aspect_ratio / w), scale[1])
+    target_area = area * random.uniform(scale[0], max_scale)
+
+    w_new = int(round(math.sqrt(target_area * aspect_ratio)))
+    h_new = int(round(math.sqrt(target_area / aspect_ratio)))
+
+    if 0 < w_new <= w and 0 < h_new <= h:
+      crop_left = random.randint(0, w - w_new)
+      crop_top = random.randint(0, h - h_new)
+
+      img = img.crop((crop_left, crop_top, crop_left + w_new, crop_top + h_new))
+      random_solution_found = True
+      break
+
+  if not random_solution_found:
+    # Center crop
+    img = center_crop(img)
+  else:
+    # Resize
+    img = img.resize([size, size], Image.BILINEAR)
+
+  return img
+
+def preprocess_train(img):
+  img = random_resized_crop(img, 224)
+  img = rand_flip(np.array(img))
+  return img

tinygrad_repo/extra/datasets/imagenet_download.py

@@ -0,0 +1,51 @@
+# Python version of https://gist.github.com/antoinebrl/7d00d5cb6c95ef194c737392ef7e476a
+from tinygrad.helpers import fetch
+from pathlib import Path
+from tqdm import tqdm
+import tarfile, os
+
+def imagenet_extract(file, path, small=False):
+  with tarfile.open(name=file) as tar:
+    if small: # Show progressbar only for big files
+      for member in tar.getmembers(): tar.extract(path=path, member=member)
+    else:
+      for member in tqdm(iterable=tar.getmembers(), total=len(tar.getmembers())): tar.extract(path=path, member=member)
+    tar.close()
+
+def imagenet_prepare_val():
+  # Read in the labels file
+  with open(Path(__file__).parent / "imagenet" / "imagenet_2012_validation_synset_labels.txt", 'r') as f:
+    labels = f.read().splitlines()
+  f.close()
+  # Get a list of images
+  images = os.listdir(Path(__file__).parent / "imagenet" / "val")
+  images.sort()
+  # Create folders and move files into those
+  for co,dir in enumerate(labels):
+    os.makedirs(Path(__file__).parent / "imagenet" / "val" / dir, exist_ok=True)
+    os.replace(Path(__file__).parent / "imagenet" / "val" / images[co], Path(__file__).parent / "imagenet" / "val" / dir / images[co])
+  os.remove(Path(__file__).parent / "imagenet" / "imagenet_2012_validation_synset_labels.txt")
+
+def imagenet_prepare_train():
+  images = os.listdir(Path(__file__).parent / "imagenet" / "train")
+  for co,tarf in enumerate(images):
+    # for each tar file found. Create a folder with its name. Extract into that folder. Remove tar file
+    if Path(Path(__file__).parent / "imagenet" / "train" / images[co]).is_file():
+      images[co] = tarf[:-4] # remove .tar from extracted tar files
+      os.makedirs(Path(__file__).parent / "imagenet" / "train" / images[co], exist_ok=True)
+      imagenet_extract(Path(__file__).parent / "imagenet" / "train" / tarf, Path(__file__).parent/ "imagenet" / "train" / images[co], small=True)
+      os.remove(Path(__file__).parent / "imagenet" / "train" / tarf)
+
+if __name__ == "__main__":
+  os.makedirs(Path(__file__).parent / "imagenet", exist_ok=True)
+  os.makedirs(Path(__file__).parent / "imagenet" / "val", exist_ok=True)
+  os.makedirs(Path(__file__).parent / "imagenet" / "train", exist_ok=True)
+  fetch("https://raw.githubusercontent.com/raghakot/keras-vis/master/resources/imagenet_class_index.json", Path(__file__).parent / "imagenet" / "imagenet_class_index.json")
+  fetch("https://raw.githubusercontent.com/tensorflow/models/master/research/slim/datasets/imagenet_2012_validation_synset_labels.txt", Path(__file__).parent / "imagenet"/ "imagenet_2012_validation_synset_labels.txt")
+  fetch("https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_val.tar", Path(__file__).parent / "imagenet" / "ILSVRC2012_img_val.tar") # 7GB
+  imagenet_extract(Path(__file__).parent / "imagenet" / "ILSVRC2012_img_val.tar", Path(__file__).parent / "imagenet" / "val")
+  imagenet_prepare_val()
+  if os.getenv('IMGNET_TRAIN', None) is not None:
+    fetch("https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_train.tar", Path(__file__).parent / "imagenet" / "ILSVRC2012_img_train.tar") #138GB!
+    imagenet_extract(Path(__file__).parent / "imagenet" / "ILSVRC2012_img_train.tar", Path(__file__).parent / "imagenet" / "train")
+    imagenet_prepare_train()

tinygrad_repo/extra/datasets/kits19.py

@@ -0,0 +1,219 @@
+import random
+import functools
+from pathlib import Path
+import numpy as np
+import nibabel as nib
+from scipy import signal, ndimage
+import os
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import fetch
+
+BASEDIR = Path(__file__).parent / "kits19" / "data"
+TRAIN_PREPROCESSED_DIR =  Path(__file__).parent / "kits19" / "preprocessed" / "train"
+VAL_PREPROCESSED_DIR =  Path(__file__).parent / "kits19" / "preprocessed" / "val"
+
+@functools.cache
+def get_train_files():
+  return sorted([x for x in BASEDIR.iterdir() if x.stem.startswith("case") and int(x.stem.split("_")[-1]) < 210 and x not in get_val_files()])
+
+@functools.cache
+def get_val_files():
+  data = fetch("https://raw.githubusercontent.com/mlcommons/training/master/retired_benchmarks/unet3d/pytorch/evaluation_cases.txt").read_text()
+  return sorted([x for x in BASEDIR.iterdir() if x.stem.split("_")[-1] in data.split("\n")])
+
+def load_pair(file_path):
+  image, label = nib.load(file_path / "imaging.nii.gz"), nib.load(file_path / "segmentation.nii.gz")
+  image_spacings = image.header["pixdim"][1:4].tolist()
+  image, label = image.get_fdata().astype(np.float32), label.get_fdata().astype(np.uint8)
+  image, label = np.expand_dims(image, 0), np.expand_dims(label, 0)
+  return image, label, image_spacings
+
+def resample3d(image, label, image_spacings, target_spacing=(1.6, 1.2, 1.2)):
+  if image_spacings != target_spacing:
+    spc_arr, targ_arr, shp_arr = np.array(image_spacings), np.array(target_spacing), np.array(image.shape[1:])
+    new_shape = (spc_arr / targ_arr * shp_arr).astype(int).tolist()
+    image = F.interpolate(torch.from_numpy(np.expand_dims(image, axis=0)), size=new_shape, mode="trilinear", align_corners=True)
+    label = F.interpolate(torch.from_numpy(np.expand_dims(label, axis=0)), size=new_shape, mode="nearest")
+    image = np.squeeze(image.numpy(), axis=0)
+    label = np.squeeze(label.numpy(), axis=0)
+  return image, label
+
+def normal_intensity(image, min_clip=-79.0, max_clip=304.0, mean=101.0, std=76.9):
+  image = np.clip(image, min_clip, max_clip)
+  image = (image - mean) / std
+  return image
+
+def pad_to_min_shape(image, label, roi_shape=(128, 128, 128)):
+  current_shape = image.shape[1:]
+  bounds = [max(0, roi_shape[i] - current_shape[i]) for i in range(3)]
+  paddings = [(0, 0)] + [(bounds[i] // 2, bounds[i] - bounds[i] // 2) for i in range(3)]
+  image = np.pad(image, paddings, mode="edge")
+  label = np.pad(label, paddings, mode="edge")
+  return image, label
+
+def preprocess(file_path):
+  image, label, image_spacings = load_pair(file_path)
+  image, label = resample3d(image, label, image_spacings)
+  image = normal_intensity(image.copy())
+  image, label = pad_to_min_shape(image, label)
+  return image, label
+
+def preprocess_dataset(filenames, preprocessed_dir, val):
+  if not preprocessed_dir.is_dir(): os.makedirs(preprocessed_dir)
+  for fn in tqdm(filenames, desc=f"preprocessing {'validation' if val else 'training'}"):
+    case = os.path.basename(fn)
+    image, label = preprocess(fn)
+    image, label = image.astype(np.float32), label.astype(np.uint8)
+    np.save(preprocessed_dir / f"{case}_x.npy", image, allow_pickle=False)
+    np.save(preprocessed_dir / f"{case}_y.npy", label, allow_pickle=False)
+
+def iterate(files, preprocessed_dir=None, val=True, shuffle=False, bs=1):
+  order = list(range(0, len(files)))
+  if shuffle: random.shuffle(order)
+  for i in range(0, len(files), bs):
+    samples = []
+    for i in order[i:i+bs]:
+      if preprocessed_dir is not None:
+        x_cached_path, y_cached_path = preprocessed_dir / f"{os.path.basename(files[i])}_x.npy", preprocessed_dir / f"{os.path.basename(files[i])}_y.npy"
+        if x_cached_path.exists() and y_cached_path.exists():
+          samples += [(np.load(x_cached_path), np.load(y_cached_path))]
+      else: samples += [preprocess(files[i])]
+    X, Y = [x[0] for x in samples], [x[1] for x in samples]
+    if val:
+      yield X[0][None], Y[0]
+    else:
+      X_preprocessed, Y_preprocessed = [], []
+      for x, y in zip(X, Y):
+        x, y = rand_balanced_crop(x, y)
+        x, y = rand_flip(x, y)
+        x, y = x.astype(np.float32), y.astype(np.uint8)
+        x = random_brightness_augmentation(x)
+        x = gaussian_noise(x)
+        X_preprocessed.append(x)
+        Y_preprocessed.append(y)
+      yield np.stack(X_preprocessed, axis=0), np.stack(Y_preprocessed, axis=0)
+
+def gaussian_kernel(n, std):
+  gaussian_1d = signal.windows.gaussian(n, std)
+  gaussian_2d = np.outer(gaussian_1d, gaussian_1d)
+  gaussian_3d = np.outer(gaussian_2d, gaussian_1d)
+  gaussian_3d = gaussian_3d.reshape(n, n, n)
+  gaussian_3d = np.cbrt(gaussian_3d)
+  gaussian_3d /= gaussian_3d.max()
+  return gaussian_3d
+
+def pad_input(volume, roi_shape, strides, padding_mode="constant", padding_val=-2.2, dim=3):
+  bounds = [(strides[i] - volume.shape[2:][i] % strides[i]) % strides[i] for i in range(dim)]
+  bounds = [bounds[i] if (volume.shape[2:][i] + bounds[i]) >= roi_shape[i] else bounds[i] + strides[i] for i in range(dim)]
+  paddings = [bounds[2]//2, bounds[2]-bounds[2]//2, bounds[1]//2, bounds[1]-bounds[1]//2, bounds[0]//2, bounds[0]-bounds[0]//2, 0, 0, 0, 0]
+  return F.pad(torch.from_numpy(volume), paddings, mode=padding_mode, value=padding_val).numpy(), paddings
+
+def sliding_window_inference(model, inputs, labels, roi_shape=(128, 128, 128), overlap=0.5, gpus=None):
+  from tinygrad.engine.jit import TinyJit
+  mdl_run = TinyJit(lambda x: model(x).realize())
+  image_shape, dim = list(inputs.shape[2:]), len(inputs.shape[2:])
+  strides = [int(roi_shape[i] * (1 - overlap)) for i in range(dim)]
+  bounds = [image_shape[i] % strides[i] for i in range(dim)]
+  bounds = [bounds[i] if bounds[i] < strides[i] // 2 else 0 for i in range(dim)]
+  inputs = inputs[
+    ...,
+    bounds[0]//2:image_shape[0]-(bounds[0]-bounds[0]//2),
+    bounds[1]//2:image_shape[1]-(bounds[1]-bounds[1]//2),
+    bounds[2]//2:image_shape[2]-(bounds[2]-bounds[2]//2),
+  ]
+  labels = labels[
+    ...,
+    bounds[0]//2:image_shape[0]-(bounds[0]-bounds[0]//2),
+    bounds[1]//2:image_shape[1]-(bounds[1]-bounds[1]//2),
+    bounds[2]//2:image_shape[2]-(bounds[2]-bounds[2]//2),
+  ]
+  inputs, paddings = pad_input(inputs, roi_shape, strides)
+  padded_shape = inputs.shape[2:]
+  size = [(inputs.shape[2:][i] - roi_shape[i]) // strides[i] + 1 for i in range(dim)]
+  result = np.zeros((1, 3, *padded_shape), dtype=np.float32)
+  norm_map = np.zeros((1, 3, *padded_shape), dtype=np.float32)
+  norm_patch = gaussian_kernel(roi_shape[0], 0.125 * roi_shape[0])
+  norm_patch = np.expand_dims(norm_patch, axis=0)
+  for i in range(0, strides[0] * size[0], strides[0]):
+    for j in range(0, strides[1] * size[1], strides[1]):
+      for k in range(0, strides[2] * size[2], strides[2]):
+        out = mdl_run(Tensor(inputs[..., i:roi_shape[0]+i,j:roi_shape[1]+j, k:roi_shape[2]+k], device=gpus)).numpy()
+        result[..., i:roi_shape[0]+i, j:roi_shape[1]+j, k:roi_shape[2]+k] += out * norm_patch
+        norm_map[..., i:roi_shape[0]+i, j:roi_shape[1]+j, k:roi_shape[2]+k] += norm_patch
+  result /= norm_map
+  result = result[..., paddings[4]:image_shape[0]+paddings[4], paddings[2]:image_shape[1]+paddings[2], paddings[0]:image_shape[2]+paddings[0]]
+  return result, labels
+
+def rand_flip(image, label, axis=(1, 2, 3)):
+  prob = 1 / len(axis)
+  for ax in axis:
+    if random.random() < prob:
+      image = np.flip(image, axis=ax).copy()
+      label = np.flip(label, axis=ax).copy()
+  return image, label
+
+def random_brightness_augmentation(image, low=0.7, high=1.3, prob=0.1):
+  if random.random() < prob:
+    factor = np.random.uniform(low=low, high=high, size=1)
+    image = (image * (1 + factor)).astype(image.dtype)
+  return image
+
+def gaussian_noise(image, mean=0.0, std=0.1, prob=0.1):
+  if random.random() < prob:
+    scale = np.random.uniform(low=0.0, high=std)
+    noise = np.random.normal(loc=mean, scale=scale, size=image.shape).astype(image.dtype)
+    image += noise
+  return image
+
+def _rand_foreg_cropb(image, label, patch_size):
+  def adjust(foreg_slice, label, idx):
+    diff = patch_size[idx - 1] - (foreg_slice[idx].stop - foreg_slice[idx].start)
+    sign = -1 if diff < 0 else 1
+    diff = abs(diff)
+    ladj = 0 if diff == 0 else random.randrange(diff)
+    hadj = diff - ladj
+    low = max(0, foreg_slice[idx].start - sign * ladj)
+    high = min(label.shape[idx], foreg_slice[idx].stop + sign * hadj)
+    diff = patch_size[idx - 1] - (high - low)
+    if diff > 0 and low == 0: high += diff
+    elif diff > 0: low -= diff
+    return low, high
+
+  cl = np.random.choice(np.unique(label[label > 0]))
+  foreg_slices = ndimage.find_objects(ndimage.label(label==cl)[0])
+  foreg_slices = [x for x in foreg_slices if x is not None]
+  slice_volumes = [np.prod([s.stop - s.start for s in sl]) for sl in foreg_slices]
+  slice_idx = np.argsort(slice_volumes)[-2:]
+  foreg_slices = [foreg_slices[i] for i in slice_idx]
+  if not foreg_slices: return _rand_crop(image, label)
+  foreg_slice = foreg_slices[random.randrange(len(foreg_slices))]
+  low_x, high_x = adjust(foreg_slice, label, 1)
+  low_y, high_y = adjust(foreg_slice, label, 2)
+  low_z, high_z = adjust(foreg_slice, label, 3)
+  image = image[:, low_x:high_x, low_y:high_y, low_z:high_z]
+  label = label[:, low_x:high_x, low_y:high_y, low_z:high_z]
+  return image, label
+
+def _rand_crop(image, label, patch_size):
+  ranges = [s - p for s, p in zip(image.shape[1:], patch_size)]
+  cord = [0 if x == 0 else random.randrange(x) for x in ranges]
+  low_x, high_x = cord[0], cord[0] + patch_size[0]
+  low_y, high_y = cord[1], cord[1] + patch_size[1]
+  low_z, high_z = cord[2], cord[2] + patch_size[2]
+  image = image[:, low_x:high_x, low_y:high_y, low_z:high_z]
+  label = label[:, low_x:high_x, low_y:high_y, low_z:high_z]
+  return image, label
+
+def rand_balanced_crop(image, label, patch_size=(128, 128, 128), oversampling=0.4):
+  if random.random() < oversampling:
+    image, label = _rand_foreg_cropb(image, label, patch_size)
+  else:
+    image, label = _rand_crop(image, label, patch_size)
+  return image, label
+
+if __name__ == "__main__":
+  for X, Y in iterate(get_val_files()):
+    print(X.shape, Y.shape)

tinygrad_repo/extra/datasets/librispeech.py

@@ -0,0 +1,82 @@
+import json
+import pathlib
+import numpy as np
+import librosa
+import soundfile
+
+"""
+The dataset has to be downloaded manually from https://www.openslr.org/12/ and put in `extra/datasets/librispeech`.
+For mlperf validation the dev-clean dataset is used.
+
+Then all the flacs have to be converted to wav using something like:
+```fish
+for file in $(find * | grep flac); do ffmpeg -i $file -ar 16k "$(dirname $file)/$(basename $file .flac).wav"; done
+```
+
+Then this [file](https://github.com/mlcommons/inference/blob/master/speech_recognition/rnnt/dev-clean-wav.json) has to also be put in `extra/datasets/librispeech`.
+"""
+BASEDIR = pathlib.Path(__file__).parent / "librispeech"
+with open(BASEDIR / "dev-clean-wav.json") as f:
+  ci = json.load(f)
+
+FILTER_BANK = np.expand_dims(librosa.filters.mel(sr=16000, n_fft=512, n_mels=80, fmin=0, fmax=8000), 0)
+WINDOW = librosa.filters.get_window("hann", 320)
+
+def feature_extract(x, x_lens):
+  x_lens = np.ceil((x_lens / 160) / 3).astype(np.int32)
+
+  # pre-emphasis
+  x = np.concatenate((np.expand_dims(x[:, 0], 1), x[:, 1:] - 0.97 * x[:, :-1]), axis=1)
+
+  # stft
+  x = librosa.stft(x, n_fft=512, window=WINDOW, hop_length=160, win_length=320, center=True, pad_mode="reflect")
+  x = np.stack((x.real, x.imag), axis=-1)
+
+  # power spectrum
+  x = (x**2).sum(-1)
+
+  # mel filter bank
+  x = np.matmul(FILTER_BANK, x)
+
+  # log
+  x = np.log(x + 1e-20)
+
+  # feature splice
+  seq = [x]
+  for i in range(1, 3):
+    tmp = np.zeros_like(x)
+    tmp[:, :, :-i] = x[:, :, i:]
+    seq.append(tmp)
+  features = np.concatenate(seq, axis=1)[:, :, ::3]
+
+  # normalize
+  features_mean = np.zeros((features.shape[0], features.shape[1]), dtype=np.float32)
+  features_std = np.zeros((features.shape[0], features.shape[1]), dtype=np.float32)
+  for i in range(features.shape[0]):
+    features_mean[i, :] = features[i, :, :x_lens[i]].mean(axis=1)
+    features_std[i, :] = features[i, :, :x_lens[i]].std(axis=1, ddof=1)
+  features_std += 1e-5
+  features = (features - np.expand_dims(features_mean, 2)) / np.expand_dims(features_std, 2)
+
+  return features.transpose(2, 0, 1), x_lens.astype(np.float32)
+
+def load_wav(file):
+  sample = soundfile.read(file)[0].astype(np.float32)
+  return sample, sample.shape[0]
+
+def iterate(bs=1, start=0):
+  print(f"there are {len(ci)} samples in the dataset")
+  for i in range(start, len(ci), bs):
+    samples, sample_lens = zip(*[load_wav(BASEDIR / v["files"][0]["fname"]) for v in ci[i : i + bs]])
+    samples = list(samples)
+    # pad to same length
+    max_len = max(sample_lens)
+    for j in range(len(samples)):
+      samples[j] = np.pad(samples[j], (0, max_len - sample_lens[j]), "constant")
+    samples, sample_lens = np.array(samples), np.array(sample_lens)
+
+    yield feature_extract(samples, sample_lens), np.array([v["transcript"] for v in ci[i : i + bs]])
+
+if __name__ == "__main__":
+  X, Y = next(iterate())
+  print(X[0].shape, Y.shape)

tinygrad_repo/extra/datasets/openimages.py

@@ -0,0 +1,209 @@
+import glob
+import sys
+import json
+import numpy as np
+from PIL import Image
+from pathlib import Path
+import boto3, botocore
+from tinygrad import Tensor, dtypes
+from tinygrad.helpers import fetch, tqdm, getenv
+import pandas as pd
+import concurrent.futures
+
+BASEDIR = Path(__file__).parent / "open-images-v6-mlperf"
+BUCKET_NAME = "open-images-dataset"
+TRAIN_BBOX_ANNOTATIONS_URL = "https://storage.googleapis.com/openimages/v6/oidv6-train-annotations-bbox.csv"
+VALIDATION_BBOX_ANNOTATIONS_URL = "https://storage.googleapis.com/openimages/v5/validation-annotations-bbox.csv"
+MAP_CLASSES_URL = "https://storage.googleapis.com/openimages/v5/class-descriptions-boxable.csv"
+MLPERF_CLASSES = ['Airplane', 'Antelope', 'Apple', 'Backpack', 'Balloon', 'Banana',
+  'Barrel', 'Baseball bat', 'Baseball glove', 'Bee', 'Beer', 'Bench', 'Bicycle',
+  'Bicycle helmet', 'Bicycle wheel', 'Billboard', 'Book', 'Bookcase', 'Boot',
+  'Bottle', 'Bowl', 'Bowling equipment', 'Box', 'Boy', 'Brassiere', 'Bread',
+  'Broccoli', 'Bronze sculpture', 'Bull', 'Bus', 'Bust', 'Butterfly', 'Cabinetry',
+  'Cake', 'Camel', 'Camera', 'Candle', 'Candy', 'Cannon', 'Canoe', 'Carrot', 'Cart',
+  'Castle', 'Cat', 'Cattle', 'Cello', 'Chair', 'Cheese', 'Chest of drawers', 'Chicken',
+  'Christmas tree', 'Coat', 'Cocktail', 'Coffee', 'Coffee cup', 'Coffee table', 'Coin',
+  'Common sunflower', 'Computer keyboard', 'Computer monitor', 'Convenience store',
+  'Cookie', 'Countertop', 'Cowboy hat', 'Crab', 'Crocodile', 'Cucumber', 'Cupboard',
+  'Curtain', 'Deer', 'Desk', 'Dinosaur', 'Dog', 'Doll', 'Dolphin', 'Door', 'Dragonfly',
+  'Drawer', 'Dress', 'Drum', 'Duck', 'Eagle', 'Earrings', 'Egg (Food)', 'Elephant',
+  'Falcon', 'Fedora', 'Flag', 'Flowerpot', 'Football', 'Football helmet', 'Fork',
+  'Fountain', 'French fries', 'French horn', 'Frog', 'Giraffe', 'Girl', 'Glasses',
+  'Goat', 'Goggles', 'Goldfish', 'Gondola', 'Goose', 'Grape', 'Grapefruit', 'Guitar',
+  'Hamburger', 'Handbag', 'Harbor seal', 'Headphones', 'Helicopter', 'High heels',
+  'Hiking equipment', 'Horse', 'House', 'Houseplant', 'Human arm', 'Human beard',
+  'Human body', 'Human ear', 'Human eye', 'Human face', 'Human foot', 'Human hair',
+  'Human hand', 'Human head', 'Human leg', 'Human mouth', 'Human nose', 'Ice cream',
+  'Jacket', 'Jeans', 'Jellyfish', 'Juice', 'Kitchen & dining room table', 'Kite',
+  'Lamp', 'Lantern', 'Laptop', 'Lavender (Plant)', 'Lemon', 'Light bulb', 'Lighthouse',
+  'Lily', 'Lion', 'Lipstick', 'Lizard', 'Man', 'Maple', 'Microphone', 'Mirror',
+  'Mixing bowl', 'Mobile phone', 'Monkey', 'Motorcycle', 'Muffin', 'Mug', 'Mule',
+  'Mushroom', 'Musical keyboard', 'Necklace', 'Nightstand', 'Office building',
+  'Orange', 'Owl', 'Oyster', 'Paddle', 'Palm tree', 'Parachute', 'Parrot', 'Pen',
+  'Penguin', 'Personal flotation device', 'Piano', 'Picture frame', 'Pig', 'Pillow',
+  'Pizza', 'Plate', 'Platter', 'Porch', 'Poster', 'Pumpkin', 'Rabbit', 'Rifle',
+  'Roller skates', 'Rose', 'Salad', 'Sandal', 'Saucer', 'Saxophone', 'Scarf', 'Sea lion',
+  'Sea turtle', 'Sheep', 'Shelf', 'Shirt', 'Shorts', 'Shrimp', 'Sink', 'Skateboard',
+  'Ski', 'Skull', 'Skyscraper', 'Snake', 'Sock', 'Sofa bed', 'Sparrow', 'Spider', 'Spoon',
+  'Sports uniform', 'Squirrel', 'Stairs', 'Stool', 'Strawberry', 'Street light',
+  'Studio couch', 'Suit', 'Sun hat', 'Sunglasses', 'Surfboard', 'Sushi', 'Swan',
+  'Swimming pool', 'Swimwear', 'Tank', 'Tap', 'Taxi', 'Tea', 'Teddy bear', 'Television',
+  'Tent', 'Tie', 'Tiger', 'Tin can', 'Tire', 'Toilet', 'Tomato', 'Tortoise', 'Tower',
+  'Traffic light', 'Train', 'Tripod', 'Truck', 'Trumpet', 'Umbrella', 'Van', 'Vase',
+  'Vehicle registration plate', 'Violin', 'Wall clock', 'Waste container', 'Watch',
+  'Whale', 'Wheel', 'Wheelchair', 'Whiteboard', 'Window', 'Wine', 'Wine glass', 'Woman',
+  'Zebra', 'Zucchini',
+]
+
+
+def openimages(base_dir:Path, subset:str, ann_file:Path):
+  valid_subsets = ['train', 'validation']
+  if subset not in valid_subsets:
+    raise ValueError(f"{subset=} must be one of {valid_subsets}")
+
+  fetch_openimages(ann_file, base_dir, subset)
+
+# this slows down the conversion a lot!
+# maybe use https://raw.githubusercontent.com/scardine/image_size/master/get_image_size.py
+def extract_dims(path): return Image.open(path).size[::-1]
+
+def export_to_coco(class_map, annotations, image_list, dataset_path, output_path, subset, classes=MLPERF_CLASSES):
+  output_path.parent.mkdir(parents=True, exist_ok=True)
+  cats = [{"id": i, "name": c, "supercategory": None} for i, c in enumerate(classes)]
+  categories_map = pd.DataFrame([(i, c) for i, c in enumerate(classes)], columns=["category_id", "category_name"])
+  class_map = class_map.merge(categories_map, left_on="DisplayName", right_on="category_name", how="inner")
+  annotations = annotations[annotations["ImageID"].isin(image_list)]
+  annotations = annotations.merge(class_map, on="LabelName", how="inner")
+  annotations["image_id"] = pd.factorize(annotations["ImageID"].tolist())[0]
+  annotations[["height", "width"]] = annotations.apply(lambda x: extract_dims(dataset_path / f"{x['ImageID']}.jpg"), axis=1, result_type="expand")
+
+  # Images
+  imgs = [{"id": int(id + 1), "file_name": f"{image_id}.jpg", "height": row["height"], "width": row["width"], "subset": subset, "license": None, "coco_url": None}
+    for (id, image_id), row in (annotations.groupby(["image_id", "ImageID"]).first().iterrows())
+  ]
+
+  # Annotations
+  annots = []
+  for i, row in annotations.iterrows():
+    xmin, ymin, xmax, ymax, img_w, img_h = [row[k] for k in ["XMin", "YMin", "XMax", "YMax", "width", "height"]]
+    x, y, w, h = xmin * img_w, ymin * img_h, (xmax - xmin) * img_w, (ymax - ymin) * img_h
+    coco_annot = {"id": int(i) + 1, "image_id": int(row["image_id"] + 1), "category_id": int(row["category_id"]), "bbox": [x, y, w, h], "area": w * h}
+    coco_annot.update({k: row[k] for k in ["IsOccluded", "IsInside", "IsDepiction", "IsTruncated", "IsGroupOf"]})
+    coco_annot["iscrowd"] = int(row["IsGroupOf"])
+    annots.append(coco_annot)
+
+  info = {"dataset": "openimages_mlperf", "version": "v6"}
+  coco_annotations = {"info": info, "licenses": [], "categories": cats, "images": imgs, "annotations": annots}
+  with open(output_path, "w") as fp:
+    json.dump(coco_annotations, fp)
+
+def get_image_list(class_map, annotations, classes=MLPERF_CLASSES):
+  labels = class_map[class_map["DisplayName"].isin(classes)]["LabelName"]
+  image_ids = annotations[annotations["LabelName"].isin(labels)]["ImageID"].unique()
+  return image_ids
+
+def download_image(bucket, subset, image_id, data_dir):
+  try:
+    bucket.download_file(f"{subset}/{image_id}.jpg", f"{data_dir}/{image_id}.jpg")
+  except botocore.exceptions.ClientError as exception:
+    sys.exit(f"ERROR when downloading image `validation/{image_id}`: {str(exception)}")
+
+def fetch_openimages(output_fn:str, base_dir:Path, subset:str):
+  bucket = boto3.resource("s3", config=botocore.config.Config(signature_version=botocore.UNSIGNED)).Bucket(BUCKET_NAME)
+
+  annotations_dir, data_dir = base_dir / "annotations", base_dir / f"{subset}/data"
+  annotations_dir.mkdir(parents=True, exist_ok=True)
+  data_dir.mkdir(parents=True, exist_ok=True)
+
+  if subset == "train":
+    annotations_fn = annotations_dir / TRAIN_BBOX_ANNOTATIONS_URL.split('/')[-1]
+    fetch(TRAIN_BBOX_ANNOTATIONS_URL, annotations_fn)
+  else:  # subset == validation
+    annotations_fn = annotations_dir / VALIDATION_BBOX_ANNOTATIONS_URL.split('/')[-1]
+    fetch(VALIDATION_BBOX_ANNOTATIONS_URL, annotations_fn)
+
+  annotations = pd.read_csv(annotations_fn)
+
+  classmap_fn = annotations_dir / MAP_CLASSES_URL.split('/')[-1]
+  fetch(MAP_CLASSES_URL, classmap_fn)
+  class_map = pd.read_csv(classmap_fn, names=["LabelName", "DisplayName"])
+
+  image_list = get_image_list(class_map, annotations)
+
+  with concurrent.futures.ThreadPoolExecutor() as executor:
+    futures = [executor.submit(download_image, bucket, subset, image_id, data_dir) for image_id in image_list]
+    for future in (t := tqdm(concurrent.futures.as_completed(futures), total=len(image_list))):
+      t.set_description(f"Downloading images")
+      future.result()
+
+  print("Converting annotations to COCO format...")
+  export_to_coco(class_map, annotations, image_list, data_dir, output_fn, subset)
+
+def image_load(base_dir, subset, fn):
+  img_folder = base_dir / f"{subset}/data"
+  return Image.open(img_folder / fn).convert('RGB')
+
+def prepare_target(annotations, img_id, img_size):
+  boxes = [annot["bbox"] for annot in annotations]
+  boxes = np.array(boxes, dtype=np.float32).reshape(-1, 4)
+  boxes[:, 2:] += boxes[:, :2]
+  boxes[:, 0::2] = boxes[:, 0::2].clip(0, img_size[1])
+  boxes[:, 1::2] = boxes[:, 1::2].clip(0, img_size[0])
+  keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
+  boxes = boxes[keep]
+  classes = [annot["category_id"] for annot in annotations]
+  classes = np.array(classes, dtype=np.int64)
+  classes = classes[keep]
+  return {"boxes": boxes, "labels": classes, "image_id": img_id, "image_size": img_size}
+
+def download_dataset(base_dir:Path, subset:str) -> Path:
+  if (ann_file:=base_dir / f"{subset}/labels/openimages-mlperf.json").is_file(): print(f"{subset} dataset is already available")
+  else:
+    print(f"Downloading {subset} dataset...")
+    openimages(base_dir, subset, ann_file)
+    print("Done")
+
+  return ann_file
+
+def random_horizontal_flip(img, tgt, prob=0.5):
+  import torch
+  import torchvision.transforms.functional as F
+  if torch.rand(1) < prob:
+    w = img.size[0]
+    img = F.hflip(img)
+    tgt["boxes"][:, [0, 2]] = w - tgt["boxes"][:, [2, 0]]
+  return img, tgt
+
+def resize(img:Image, tgt:dict[str, np.ndarray|tuple]|None=None, size:tuple[int, int]=(800, 800)) -> tuple[np.ndarray, np.ndarray, tuple]|tuple[np.ndarray, tuple]:
+  import torchvision.transforms.functional as F
+  img_size = img.size[::-1]
+  img = F.resize(img, size=size)
+  img = np.array(img)
+
+  if tgt is not None:
+    ratios = [s / s_orig for s, s_orig in zip(size, img_size)]
+    ratio_h, ratio_w = ratios
+    x_min, y_min, x_max, y_max = [tgt["boxes"][:, i] for i in range(tgt["boxes"].shape[-1])]
+    x_min = x_min * ratio_w
+    x_max = x_max * ratio_w
+    y_min = y_min * ratio_h
+    y_max = y_max * ratio_h
+
+    tgt["boxes"] = np.stack([x_min, y_min, x_max, y_max], axis=1)
+    return img, tgt, img_size
+
+  return img, img_size
+
+def normalize(img:Tensor, device:list[str]|None = None):
+  mean = Tensor([0.485, 0.456, 0.406], device=device, dtype=dtypes.float32).reshape(1, -1, 1, 1)
+  std = Tensor([0.229, 0.224, 0.225], device=device, dtype=dtypes.float32).reshape(1, -1, 1, 1)
+  img = ((img.permute([0, 3, 1, 2]) / 255.0) - mean) / std
+  return img.cast(dtypes.default_float)
+
+def get_dataset_count(base_dir:Path, val:bool) -> int:
+  if not (files:=glob.glob(p:=str(base_dir / f"{'validation' if val else 'train'}/data/*.jpg"))): raise FileNotFoundError(f"No files in {p}")
+  return len(files)
+
+if __name__ == "__main__":
+  download_dataset(base_dir:=getenv("BASEDIR", BASEDIR), "train")
+  download_dataset(base_dir, "validation")

tinygrad_repo/extra/datasets/preprocess_imagenet.py

@@ -0,0 +1,21 @@
+from tinygrad import Tensor, dtypes
+from extra.datasets.imagenet import iterate, get_val_files
+
+if __name__ == "__main__":
+  #sz = len(get_val_files())
+  sz = 32*100
+  X,Y = None, None
+
+  idx = 0
+  for x,y in iterate(shuffle=False):
+    print(x.shape, y.shape, x.dtype, y.dtype)
+    assert x.shape[0] == y.shape[0]
+    bs = x.shape[0]
+    if X is None:
+      X = Tensor.empty(sz, *x.shape[1:], device="disk:/tmp/imagenet_x", dtype=dtypes.uint8)
+      Y = Tensor.empty(sz, *y.shape[1:], device="disk:/tmp/imagenet_y", dtype=dtypes.int64)
+      print(X.shape, Y.shape)
+    X[idx:idx+bs].assign(x)
+    Y[idx:idx+bs].assign(y)
+    idx += bs
+    if idx >= sz: break

tinygrad_repo/extra/datasets/squad.py

@@ -0,0 +1,148 @@
+import json
+import os
+from pathlib import Path
+from transformers import BertTokenizer
+import numpy as np
+from tinygrad.helpers import fetch
+
+BASEDIR = Path(__file__).parent / "squad"
+def init_dataset():
+  os.makedirs(BASEDIR, exist_ok=True)
+  fetch("https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json", BASEDIR / "dev-v1.1.json")
+  with open(BASEDIR / "dev-v1.1.json") as f:
+    data = json.load(f)["data"]
+
+  examples = []
+  for article in data:
+    for paragraph in article["paragraphs"]:
+      text = paragraph["context"]
+      doc_tokens = []
+      prev_is_whitespace = True
+      for c in text:
+        if c == " " or c == "\t" or c == "\r" or c == "\n" or ord(c) == 0x202F:
+          prev_is_whitespace = True
+        else:
+          if prev_is_whitespace:
+            doc_tokens.append(c)
+          else:
+            doc_tokens[-1] += c
+          prev_is_whitespace = False
+
+      for qa in paragraph["qas"]:
+        qa_id = qa["id"]
+        q_text = qa["question"]
+
+        examples.append({
+          "id": qa_id,
+          "question": q_text,
+          "context": doc_tokens,
+          "answers": list(map(lambda x: x["text"], qa["answers"]))
+        })
+  return examples
+
+def _check_is_max_context(doc_spans, cur_span_index, position):
+  best_score, best_span_index = None, None
+  for di, (doc_start, doc_length) in enumerate(doc_spans):
+    end = doc_start + doc_length - 1
+    if position < doc_start:
+      continue
+    if position > end:
+      continue
+    num_left_context = position - doc_start
+    num_right_context = end - position
+    score = min(num_left_context, num_right_context) + 0.01 * doc_length
+    if best_score is None or score > best_score:
+      best_score = score
+      best_span_index = di
+  return cur_span_index == best_span_index
+
+def convert_example_to_features(example, tokenizer):
+  query_tokens = tokenizer.tokenize(example["question"])
+
+  if len(query_tokens) > 64:
+    query_tokens = query_tokens[:64]
+
+  tok_to_orig_index = []
+  orig_to_tok_index = []
+  all_doc_tokens = []
+  for i, token in enumerate(example["context"]):
+    orig_to_tok_index.append(len(all_doc_tokens))
+    sub_tokens = tokenizer.tokenize(token)
+    for sub_token in sub_tokens:
+      tok_to_orig_index.append(i)
+      all_doc_tokens.append(sub_token)
+
+  max_tokens_for_doc = 384 - len(query_tokens) - 3
+
+  doc_spans = []
+  start_offset = 0
+  while start_offset < len(all_doc_tokens):
+    length = len(all_doc_tokens) - start_offset
+    length = min(length, max_tokens_for_doc)
+    doc_spans.append((start_offset, length))
+    if start_offset + length == len(all_doc_tokens):
+      break
+    start_offset += min(length, 128)
+
+  outputs = []
+  for di, (doc_start, doc_length) in enumerate(doc_spans):
+    tokens = []
+    token_to_orig_map = {}
+    token_is_max_context = {}
+    segment_ids = []
+    tokens.append("[CLS]")
+    segment_ids.append(0)
+    for token in query_tokens:
+      tokens.append(token)
+      segment_ids.append(0)
+    tokens.append("[SEP]")
+    segment_ids.append(0)
+
+    for i in range(doc_length):
+      split_token_index = doc_start + i
+      token_to_orig_map[len(tokens)] = tok_to_orig_index[split_token_index]
+      token_is_max_context[len(tokens)] = _check_is_max_context(doc_spans, di, split_token_index)
+      tokens.append(all_doc_tokens[split_token_index])
+      segment_ids.append(1)
+    tokens.append("[SEP]")
+    segment_ids.append(1)
+
+    input_ids = tokenizer.convert_tokens_to_ids(tokens)
+    input_mask = [1] * len(input_ids)
+
+    while len(input_ids) < 384:
+      input_ids.append(0)
+      input_mask.append(0)
+      segment_ids.append(0)
+
+    assert len(input_ids) == 384
+    assert len(input_mask) == 384
+    assert len(segment_ids) == 384
+
+    outputs.append({
+      "input_ids": np.expand_dims(np.array(input_ids), 0).astype(np.float32),
+      "input_mask": np.expand_dims(np.array(input_mask), 0).astype(np.float32),
+      "segment_ids": np.expand_dims(np.array(segment_ids), 0).astype(np.float32),
+      "token_to_orig_map": token_to_orig_map,
+      "token_is_max_context": token_is_max_context,
+      "tokens": tokens,
+    })
+
+  return outputs
+
+def iterate(tokenizer, start=0):
+  examples = init_dataset()
+  print(f"there are {len(examples)} pairs in the dataset")
+
+  for i in range(start, len(examples)):
+    example = examples[i]
+    features = convert_example_to_features(example, tokenizer)
+    # we need to yield all features here as the f1 score is the maximum over all features
+    yield features, example
+
+if __name__ == "__main__":
+  tokenizer = BertTokenizer(str(Path(__file__).parents[2] / "weights" / "bert_vocab.txt"))
+
+  X, Y = next(iterate(tokenizer))
+  print(" ".join(X[0]["tokens"]))
+  print(X[0]["input_ids"].shape, Y)

tinygrad_repo/extra/datasets/wikipedia.py

@@ -0,0 +1,398 @@
+# Preprocessing of downloaded text from Wikipedia for MLPerf BERT training
+# This is a modified version of the original script:
+# https://github.com/mlcommons/training/blob/master/language_model/tensorflow/bert/cleanup_scripts/create_pretraining_data.py
+# ENV VARS:
+# MAX_SEQ_LENGTH          - Maximum sequence length
+# MAX_PREDICTIONS_PER_SEQ - Maximum number of masked LM predictions per sequence
+# RANDOM_SEED             - Random seed
+# DUPE_FACTOR             - Number of times to duplicate the input data with different masks
+# MASKED_LM_PROB          - Probability of masking a token
+# SHORT_SEQ_PROB          - Probability of picking a sequence shorter than MAX_SEQ_LENGTH
+
+import os, sys, pickle, random, unicodedata
+from pathlib import Path
+import numpy as np
+from tqdm import tqdm
+from tqdm.contrib.concurrent import process_map
+
+from tinygrad.helpers import diskcache, getenv
+
+BASEDIR = getenv('BASEDIR', Path(__file__).parent / "wiki")
+
+################### Tokenization #####################
+
+def _is_whitespace(char:str) -> bool:
+  if char == " " or char == "\t" or char == "\n" or char == "\r":
+    return True
+  return unicodedata.category(char) == "Zs"
+
+def _is_control(char:str) -> bool:
+  if char == "\t" or char == "\n" or char == "\r":
+    return False
+  return unicodedata.category(char).startswith("C")
+
+def _is_punctuation(char:str) -> bool:
+  # range(33, 48) -> ! " # $ % & ' ( ) * + , - . /
+  # range(58, 65) -> : ; < = > ? @
+  # range(91, 97) -> [ \ ] ^ _
+  # range(123, 127) -> { | } ~
+  if (cp := ord(char)) in range(33, 48) or cp in range(58, 65) or cp in range(91, 97) or cp in range(123, 127):
+    return True
+  return unicodedata.category(char).startswith("P")
+
+def _is_chinese_char(cp:int) -> bool:
+  if ((cp >= 0x4E00 and cp <= 0x9FFF) or
+      (cp >= 0x3400 and cp <= 0x4DBF) or
+      (cp >= 0x20000 and cp <= 0x2A6DF) or
+      (cp >= 0x2A700 and cp <= 0x2B73F) or
+      (cp >= 0x2B740 and cp <= 0x2B81F) or
+      (cp >= 0x2B820 and cp <= 0x2CEAF) or
+      (cp >= 0xF900 and cp <= 0xFAFF) or
+      (cp >= 0x2F800 and cp <= 0x2FA1F)):
+    return True
+  return False
+
+def _run_split_on_punc(text:str) -> list[str]:
+  if text in ("[UNK]", "[SEP]", "[PAD]", "[CLS]", "[MASK]"):
+    return [text]
+  start_new_word = True
+  output = []
+  for i in range(len(text)):
+    if _is_punctuation(char := text[i]):
+      output.append([char])
+      start_new_word = True
+    else:
+      if start_new_word:
+        output.append([])
+      start_new_word = False
+      output[-1].append(char)
+  return ["".join(x) for x in output]
+
+def _run_strip_accents(text:str) -> str:
+  output = []
+  for char in unicodedata.normalize("NFD", text):
+    if unicodedata.category(char) != "Mn":
+      output.append(char)
+  return "".join(output)
+
+def _clean_text(text:str) -> str:
+  output = []
+  for char in text:
+    if not ((cp := ord(char)) == 0 or cp == 0xfffd or _is_control(char)):
+      output.append(" " if _is_whitespace(char) else char)
+  return "".join(output)
+
+def _tokenize_chinese_chars(text:str) -> str:
+  output = []
+  for char in text:
+    cp = ord(char)
+    if _is_chinese_char(cp):
+      output.append(" ")
+      output.append(char)
+      output.append(" ")
+    else:
+      output.append(char)
+  return "".join(output)
+
+def whitespace_tokenize(text):
+  if not (text := text.strip()): return []
+  return text.split()
+
+def _wordpiece_tokenize(text:str, vocab:dict[str, int]) -> list[str]:
+  text = text.decode("utf-8", "ignore") if isinstance(text, bytes) else text
+  output_tokens = []
+  for token in text.strip().split():
+    chars = list(token)
+    if len(chars) > 200:
+      output_tokens.append("[UNK]")
+      continue
+
+    is_bad = False
+    start = 0
+    sub_tokens = []
+    while start < len(chars):
+      end = len(chars)
+      cur_substr = None
+      while start < end:
+        substr = "".join(chars[start:end])
+        if start > 0: substr = "##" + substr
+        if substr in vocab:
+          cur_substr = substr
+          break
+        end -= 1
+      if cur_substr is None:
+        is_bad = True
+        break
+      sub_tokens.append(cur_substr)
+      start = end
+
+    if is_bad: output_tokens.append("[UNK]")
+    else: output_tokens.extend(sub_tokens)
+  return output_tokens
+
+class Tokenizer:
+  def __init__(self, vocab_file):
+    self.vocab = {}
+    with open(vocab_file) as f:
+      for line in f:
+        line = line.decode("utf-8", "ignore") if isinstance(line, bytes) else line
+        if (token := line.strip()) and token not in self.vocab: self.vocab[token] = len(self.vocab)
+    self.inv_vocab = {v: k for k, v in self.vocab.items()}
+
+  def tokenize(self, text:str) -> list[str]:
+    # BasicTokenizer
+    split_tokens = []
+    for token in whitespace_tokenize(_tokenize_chinese_chars(_clean_text(text.decode("utf-8", "ignore") if isinstance(text, bytes) else text))):
+      split_tokens.extend(_run_split_on_punc(_run_strip_accents(token.lower())))
+    split_tokens = " ".join(split_tokens).strip().split()
+    # WordpieceTokenizer
+    tokens = []
+    for token in split_tokens:
+      tokens.extend(_wordpiece_tokenize(token, self.vocab))
+    return tokens
+
+  def convert_tokens_to_ids(self, tokens:list[str]) -> list[int]: return [self.vocab[token] for token in tokens]
+  def convert_ids_to_tokens(self, ids:list[int]) -> list[str]: return [self.inv_vocab[id] for id in ids]
+
+##################### Feature transformation #####################
+
+def truncate_seq_pair(tokens_a:list[str], tokens_b:list[str], max_num_tokens:int, rng:random.Random) -> None:
+  while True:
+    total_length = len(tokens_a) + len(tokens_b)
+    if total_length <= max_num_tokens:
+      break
+
+    trunc_tokens = tokens_a if len(tokens_a) > len(tokens_b) else tokens_b
+    assert len(trunc_tokens) >= 1
+
+    if rng.random() < 0.5:
+      del trunc_tokens[0]
+    else:
+      trunc_tokens.pop()
+
+def create_masked_lm_predictions(tokens:list[str], tokenizer:Tokenizer, rng:random.Random, vocab_words:list[str]) -> tuple[list[str], list[int], list[str]]:
+  cand_indices = []
+  for i, token in enumerate(tokens):
+    if token == "[CLS]" or token == "[SEP]":
+      continue
+    cand_indices.append(i)
+
+  rng.shuffle(cand_indices)
+  output_tokens = list(tokens)
+  num_to_predict = min(getenv('MAX_PREDICTIONS_PER_SEQ', 76), max(1, int(round(len(tokens) * getenv("MASKED_LM_PROB", 0.15)))))
+
+  masked_lms = []
+  covered_indices = set()
+  for index in cand_indices:
+    if len(masked_lms) >= num_to_predict:
+      break
+    if index in covered_indices:
+      continue
+    covered_indices.add(index)
+
+    masked_token = None
+    if rng.random() < 0.8:
+      masked_token = "[MASK]"
+    else:
+      if rng.random() < 0.5:
+        masked_token = tokens[index]
+      else:
+        masked_token = vocab_words[rng.randint(0, len(tokenizer.vocab) - 1)]
+
+    output_tokens[index] = masked_token
+    masked_lms.append((index, tokens[index]))
+  masked_lms = sorted(masked_lms, key=lambda x: x[0])
+
+  masked_lm_positions = []
+  masked_lm_labels = []
+  for p in masked_lms:
+    masked_lm_positions.append(p[0])
+    masked_lm_labels.append(p[1])
+
+  return output_tokens, masked_lm_positions, masked_lm_labels
+
+def create_instances_from_document(rng:random.Random, tokenizer:Tokenizer, doc:list[str], di:int, documents:list[list[str]]) -> list[dict]:
+  max_num_tokens = getenv('MAX_SEQ_LENGTH', 512) - 3 # [CLS] + 2 * [SEP]
+
+  target_seq_length = max_num_tokens
+  if rng.random() < getenv("SHORT_SEQ_PROB", 0.1):
+    target_seq_length = rng.randint(2, max_num_tokens)
+
+  instances = []
+  current_chunk = []
+  current_length = 0
+  i = 0
+  while i < len(doc):
+    segment = doc[i]
+    current_chunk.append(segment)
+    current_length += len(segment)
+    if i == len(doc) - 1 or current_length >= target_seq_length:
+      if current_chunk:
+        a_end = 1
+        if len(current_chunk) >= 2:
+          a_end = rng.randint(1, len(current_chunk) - 1)
+
+        tokens_a = []
+        for j in range(a_end):
+          tokens_a.extend(current_chunk[j])
+
+        tokens_b = []
+        is_random_next = False
+        if len(current_chunk) == 1 or rng.random() < 0.5:
+          is_random_next = True
+          target_b_length = target_seq_length - len(tokens_a)
+
+          for _ in range(10):
+            random_document_index = rng.randint(0, len(documents) - 1)
+            if random_document_index != di:
+              break
+
+          random_document = documents[random_document_index]
+          random_start = rng.randint(0, len(random_document) - 1)
+          for j in range(random_start, len(random_document)):
+            tokens_b.extend(random_document[j])
+            if len(tokens_b) >= target_b_length:
+              break
+
+          num_unused_segments = len(current_chunk) - a_end
+          i -= num_unused_segments
+        else:
+          is_random_next = False
+          for j in range(a_end, len(current_chunk)):
+            tokens_b.extend(current_chunk[j])
+        truncate_seq_pair(tokens_a, tokens_b, max_num_tokens, rng)
+
+        assert len(tokens_a) >= 1
+        assert len(tokens_b) >= 1
+
+        tokens = []
+        segment_ids = []
+        tokens.append("[CLS]")
+        segment_ids.append(0)
+        for token in tokens_a:
+          tokens.append(token)
+          segment_ids.append(0)
+        tokens.append("[SEP]")
+        segment_ids.append(0)
+        for token in tokens_b:
+          tokens.append(token)
+          segment_ids.append(1)
+        tokens.append("[SEP]")
+        segment_ids.append(1)
+
+        tokens, masked_lm_positions, masked_lm_labels = create_masked_lm_predictions(tokens, tokenizer, rng, list(tokenizer.vocab.keys()))
+        instances.append({
+          "tokens": tokens,
+          "segment_ids": segment_ids,
+          "masked_lm_positions": masked_lm_positions,
+          "masked_lm_labels": masked_lm_labels,
+          "is_random_next": is_random_next
+        })
+      current_chunk = []
+      current_length = 0
+    i += 1
+  return instances
+
+def get_documents(rng:random.Random, tokenizer:Tokenizer, fn:str) -> list[list[str]]:
+  documents = [[]]
+  with open(BASEDIR / fn) as f:
+    for line in f.readlines():
+      if not (line := line.decode("utf-8", "ignore") if isinstance(line, bytes) else line): break
+      if not (line := line.strip()): documents.append([])
+      if (tokens := tokenizer.tokenize(line)): documents[-1].append(tokens)
+  documents = [x for x in documents if x]
+  rng.shuffle(documents)
+  return documents
+
+def get_instances(rng:random.Random, tokenizer:Tokenizer, documents:list[list[str]]) -> list[dict]:
+  instances = []
+  for _ in range(getenv('DUPE_FACTOR', 10)):
+    for di, doc in enumerate(documents):
+      instances.extend(create_instances_from_document(rng, tokenizer, doc, di, documents))
+  rng.shuffle(instances)
+  return instances
+
+def instance_to_features(instance:dict, tokenizer:Tokenizer) -> dict:
+  input_ids = tokenizer.convert_tokens_to_ids(instance["tokens"])
+  input_mask = [1] * len(input_ids)
+  segment_ids = instance["segment_ids"]
+
+  max_seq_length = getenv('MAX_SEQ_LENGTH', 512)
+
+  assert len(input_ids) <= max_seq_length
+  while len(input_ids) < max_seq_length:
+    input_ids.append(0)
+    input_mask.append(0)
+    segment_ids.append(0)
+  assert len(input_ids) == max_seq_length
+  assert len(input_mask) == max_seq_length
+  assert len(segment_ids) == max_seq_length
+
+  masked_lm_positions = instance["masked_lm_positions"]
+  masked_lm_ids = tokenizer.convert_tokens_to_ids(instance["masked_lm_labels"])
+  masked_lm_weights = [1.0] * len(masked_lm_ids)
+
+  while len(masked_lm_positions) < getenv("MAX_PREDICTIONS_PER_SEQ", 76):
+    masked_lm_positions.append(0)
+    masked_lm_ids.append(0)
+    masked_lm_weights.append(0.0)
+
+  next_sentence_label = 1 if instance["is_random_next"] else 0
+
+  return {
+    "input_ids": np.expand_dims(np.array(input_ids, dtype=np.int32), 0),
+    "input_mask": np.expand_dims(np.array(input_mask, dtype=np.int32), 0),
+    "segment_ids": np.expand_dims(np.array(segment_ids, dtype=np.int32), 0),
+    "masked_lm_positions": np.expand_dims(np.array(masked_lm_positions, dtype=np.int32), 0),
+    "masked_lm_ids": np.expand_dims(np.array(masked_lm_ids, dtype=np.int32), 0),
+    "masked_lm_weights": np.expand_dims(np.array(masked_lm_weights, dtype=np.float32), 0),
+    "next_sentence_labels": np.expand_dims(np.array([next_sentence_label], dtype=np.int32), 0),
+  }
+
+def process_part(part:int):
+  tokenizer = Tokenizer(getenv("BASEDIR", Path(__file__).parent / "wiki") / "vocab.txt")
+  os.makedirs(BASEDIR / "train", exist_ok=True)
+
+  if os.path.exists(BASEDIR / f"train/{str(part)}.pkl"): return
+  features = get_features_from_part(tokenizer, val=False, part=part)
+  with open(BASEDIR / f"train/{str(part)}.pkl", "wb") as f:
+    pickle.dump(features, f)
+
+def get_features_from_part(tokenizer:Tokenizer, val:bool=False, part:int=0) -> list[dict]: # Convert raw text to masked NSP samples
+  rng = random.Random(getenv('RANDOM_SEED', 12345))
+
+  if val:
+    tqdm.write("Getting samples from dataset")
+    documents = get_documents(rng, tokenizer, "results4/eval.txt")
+    instances = get_instances(rng, tokenizer, documents)
+
+    tqdm.write(f"There are {len(instances)} samples in the dataset")
+    tqdm.write(f"Picking 10000 samples")
+
+    pick_ratio = len(instances) / 10000
+    return [instance_to_features(instances[int(inst*pick_ratio)], tokenizer) for inst in range(10000)]
+  else:
+    documents = get_documents(rng, tokenizer, f"results4/part-{part:05d}-of-00500")
+    instances = get_instances(rng, tokenizer, documents)
+    return [instance_to_features(instance, tokenizer) for instance in instances]
+
+##################### Load files #####################
+
+@diskcache
+def get_wiki_train_files(): return sorted(list((BASEDIR / "train/").glob("*.pkl")))
+
+if __name__ == "__main__":
+  tokenizer = Tokenizer(getenv("BASEDIR", Path(__file__).parent / "wiki") / "vocab.txt")
+
+  assert len(sys.argv) > 1, "Usage: python wikipedia.py pre-eval|pre-train [part]|all"
+
+  if sys.argv[1] == "pre-eval": # Generate 10000 eval samples
+    with open(BASEDIR / "eval.pkl", "wb") as f:
+      pickle.dump(get_features_from_part(tokenizer, val=True), f)
+  elif sys.argv[1] == "pre-train":
+    if sys.argv[2] == "all": # Use all 500 parts for training generation
+      process_map(process_part, [part for part in range(500)], max_workers=getenv('NUM_WORKERS', min(os.cpu_count(), 32)), chunksize=1)
+    else: # Use a specific part for training generation
+      part = sys.argv[2]
+      print(f"Processing part {part}...")
+      process_part(int(part))

tinygrad_repo/extra/datasets/wikipedia_download.py

@@ -0,0 +1,54 @@
+# pip install gdown
+# Downloads the 2020 wikipedia dataset used for MLPerf BERT training
+import os, hashlib
+from pathlib import Path
+import tarfile
+import gdown
+from tqdm import tqdm
+from tinygrad.helpers import getenv
+
+def gdrive_download(url:str, path:str):
+  if not os.path.exists(path): gdown.download(url, path)
+
+def wikipedia_uncompress_and_extract(file:str, path:str, small:bool=False):
+  if not os.path.exists(os.path.join(path, "results4")):
+    print("Uncompressing and extracting file...")
+    with tarfile.open(file, 'r:gz') as tar:
+      tar.extractall(path=path)
+      os.remove(file)
+      if small:
+        for member in tar.getmembers(): tar.extract(path=path, member=member)
+      else:
+        for member in tqdm(iterable=tar.getmembers(), total=len(tar.getmembers())): tar.extract(path=path, member=member)
+
+def verify_checksum(folder_path:str, checksum_path:str):
+  print("Verifying checksums...")
+  with open(checksum_path, 'r') as f:
+    for line in f:
+      expected_checksum, folder_name = line.split()
+      file_path = os.path.join(folder_path, folder_name[2:]) # remove './' from the start of the folder name
+      hasher = hashlib.md5()
+      with open(file_path, 'rb') as f:
+        for buf in iter(lambda: f.read(4096), b''): hasher.update(buf)
+      if hasher.hexdigest() != expected_checksum:
+        raise ValueError(f"Checksum does not match for file: {file_path}")
+  print("All checksums match.")
+
+def download_wikipedia(path:str):
+  # Links from: https://github.com/mlcommons/training/blob/master/language_model/tensorflow/bert/dataset.md
+  os.makedirs(path, exist_ok=True)
+  gdrive_download("https://drive.google.com/uc?id=1fbGClQMi2CoMv7fwrwTC5YYPooQBdcFW", os.path.join(path, "bert_config.json"))
+  gdrive_download("https://drive.google.com/uc?id=1USK108J6hMM_d27xCHi738qBL8_BT1u1", os.path.join(path, "vocab.txt"))
+  gdrive_download("https://drive.google.com/uc?id=1chiTBljF0Eh1U5pKs6ureVHgSbtU8OG_", os.path.join(path, "model.ckpt-28252.data-00000-of-00001"))
+  gdrive_download("https://drive.google.com/uc?id=1Q47V3K3jFRkbJ2zGCrKkKk-n0fvMZsa0", os.path.join(path, "model.ckpt-28252.index"))
+  gdrive_download("https://drive.google.com/uc?id=1vAcVmXSLsLeQ1q7gvHnQUSth5W_f_pwv", os.path.join(path, "model.ckpt-28252.meta"))
+  with open(os.path.join(path, "checkpoint"), "w") as f: f.write('model_checkpoint_path: "model.ckpt-28252"\nall_model_checkpoint_paths: "model.ckpt-28252"')
+  if getenv("WIKI_TRAIN", 0):
+    gdrive_download("https://drive.google.com/uc?id=1tmMgLwoBvbEJEHXh77sqrXYw5RpqT8R_", os.path.join(path, "bert_reference_results_text_md5.txt"))
+    gdrive_download("https://drive.google.com/uc?id=14xV2OUGSQDG_yDBrmbSdcDC-QGeqpfs_", os.path.join(path, "results_text.tar.gz"))
+    wikipedia_uncompress_and_extract(os.path.join(path, "results_text.tar.gz"), path)
+    if getenv("VERIFY_CHECKSUM", 0):
+      verify_checksum(os.path.join(path, "results4"), os.path.join(path, "bert_reference_results_text_md5.txt"))
+
+if __name__ == "__main__":
+  download_wikipedia(getenv("BASEDIR", os.path.join(Path(__file__).parent / "wiki")))

tinygrad_repo/extra/dsp/Dockerfile

@@ -0,0 +1,38 @@
+# Use a recent Ubuntu base image.
+FROM ubuntu:22.04
+
+# Install required packages.
+RUN apt-get update && apt-get install -y \
+    git \
+    build-essential \
+    python3 \
+    python3-pip \
+    python3-tomli \
+    pkg-config \
+    libglib2.0-dev \
+    libfdt-dev \
+    libpixman-1-dev \
+    zlib1g-dev \
+    ninja-build \
+    meson \
+    wget
+
+# Clone QEMU source (you can pin a specific version if desired)
+RUN wget https://download.qemu.org/qemu-9.2.0.tar.xz && tar xvJf qemu-9.2.0.tar.xz
+
+WORKDIR /qemu-9.2.0
+
+RUN apt-get install -y flex bison
+
+# Configure QEMU to build the hexagon user-mode emulator.
+RUN ./configure --target-list=hexagon-linux-user && make -j$(nproc)
+
+# Optionally, install QEMU into /usr/local (or leave it in place).
+RUN make install
+
+# delete the source (for space)
+RUN cd .. && rm -rf /qemu-9.2.0
+
+# The QEMU binaries will be in /usr/local/bin.
+# Set the entrypoint to bash so you can interact with the container.
+ENTRYPOINT ["/bin/bash"]

tinygrad_repo/extra/dsp/compile.py

@@ -0,0 +1,125 @@
+#!/usr/bin/env python3
+import os, ctypes, time, fcntl, mmap
+import llvmlite.binding as llvm
+from tinygrad.helpers import getenv, to_mv
+from tinygrad.runtime.support.elf import elf_loader
+from hexdump import hexdump
+from tinygrad.runtime.autogen import libc
+if getenv("IOCTL"): import run # noqa: F401 # pylint: disable=unused-import
+
+adsp = ctypes.CDLL(ctypes.util.find_library("adsprpc"))
+import adsprpc
+import ion
+import msm_ion
+ION_IOC_ALLOC = 0
+ION_IOC_MAP = 2
+ION_IOC_SHARE = 4
+ION_IOC_CUSTOM = 6
+ION_ADSP_HEAP_ID = 22
+ION_IOMMU_HEAP_ID = 25
+
+def ion_iowr(fd, nr, args):
+  ret = fcntl.ioctl(fd, (3 << 30) | (ctypes.sizeof(args) & 0x1FFF) << 16 | (ord(ion.ION_IOC_MAGIC) & 0xFF) << 8 | (nr & 0xFF), args)
+  if ret != 0: raise RuntimeError(f"ioctl returned {ret}")
+
+if __name__ == "__main__":
+  # TODO: mmap tensors to the DSP
+  # call the target function with the mmaped tensors
+  ion_fd = os.open("/dev/ion", os.O_RDWR | os.O_CLOEXEC)
+  arg3 = ion.struct_ion_allocation_data(len=0x1000, align=0x1000, heap_id_mask=1<<msm_ion.ION_SYSTEM_HEAP_ID, flags=ion.ION_FLAG_CACHED)
+  ion_iowr(ion_fd, ION_IOC_ALLOC, arg3)
+  print(arg3.handle)
+  arg2 = ion.struct_ion_fd_data(handle=arg3.handle)
+  ion_iowr(ion_fd, ION_IOC_SHARE, arg2)
+  print(arg2.fd)
+
+  res = libc.mmap(0, 0x1000, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED, arg2.fd, 0)
+  print("mmapped", hex(res))
+  to_mv(res, 0x10)[1] = 0xaa
+
+  from tinygrad.runtime.ops_dsp import ClangCompiler
+  cc = ClangCompiler(args=["--target=hexagon", "-mcpu=hexagonv65", "-fuse-ld=lld", "-nostdlib"])
+
+  obj = cc.compile("""
+    typedef unsigned long long remote_handle64;
+    typedef struct { void *pv; unsigned int len; } remote_buf;
+    typedef struct { int fd; unsigned int offset; } remote_dma_handle;
+    typedef union { remote_buf buf; remote_handle64 h64; remote_dma_handle dma; } remote_arg;
+    void* HAP_mmap(void *addr, int len, int prot, int flags, int fd, long offset);
+    int HAP_munmap(void *addr, int len);
+    #define HAP_MEM_CACHE_WRITETHROUGH 0x40
+
+    int entry(unsigned long long handle, unsigned int sc, remote_arg* pra) {
+      if (sc>>24 == 1) {
+        //void *mmaped = *((void**)pra[0].buf.pv);
+        void *a = HAP_mmap(0, 0x1000, 3, 0, pra[1].dma.fd, 0);
+        ((char*)a)[0] = 0x55;
+        ((char*)a)[4] = 0x55;
+        ((char*)a)[8] = 0x99;
+        //((char*)a)[1] = 0x9b;
+        //char ret = ((char*)a)[1];
+        HAP_munmap(a, 0x1000);
+        return 0;
+
+        //return ((int)mmaped)&0xFFFF;
+        //return ((char*)mmaped)[1];
+        //return sizeof(void*);
+
+        //((char*)mmaped)[0] = 55;
+        //return ((int)mmaped)&0xFFFF;
+
+        //void addr = *((void**)pra[1])
+        //return sizeof(remote_buf);
+        //((char*)pra[1].h64)[0] = 55;
+
+        //return ((char*)mmaped)[1];
+        //((char*)mmaped)[0] = 55;
+        // NOTE: you have to return 0 for outbufs to work
+        //return ((int)pra[1].h64)&0xFFFF;
+      }
+      return 0;
+    }
+  """)
+  with open("/tmp/swag.so", "wb") as f: f.write(obj)
+
+  handle = ctypes.c_int64(-1)
+  adsp.remote_handle64_open(ctypes.create_string_buffer(b"file:////tmp/swag.so?entry&_modver=1.0&_dom=cdsp"), ctypes.byref(handle))
+  print("HANDLE", handle.value)
+  #print(adsp.remote_handle64_invoke(handle, 0, None))
+
+  #rem = adsp.remote_register_buf(res, 0x1000, arg2.fd, 4)
+  #rem = adsp.remote_register_dma_handle(arg2.fd, 0x1000)
+  #print("remote_register_buf_attr", rem)
+
+  #out = ctypes.c_uint64(0)
+  #ret = adsp.remote_mmap(arg2.fd, 0, 0, 0x1000, ctypes.byref(out))
+  #print(ret)
+  #print("mapped at", hex(out.value))
+
+  #arg_2 = ctypes.c_int64(out.value)
+  arg_2 = ctypes.c_int64(arg2.fd)
+  pra = (adsprpc.union_remote_arg64 * 3)()
+  pra[0].buf.pv = ctypes.addressof(arg_2)
+  pra[0].buf.len = 8
+  pra[1].dma.fd = arg2.fd
+  pra[1].dma.len = 0x1000
+  print("invoke")
+  ret = adsp.remote_handle64_invoke(handle, (1<<24) | (1<<16) | (1 << 4), pra)
+  print("return value", ret, hex(ret))
+  #print(hex(arg_2.value), arg_2.value)
+  #time.sleep(0.1)
+
+  # flush the cache
+  """
+  flush_data = msm_ion.struct_ion_flush_data(handle=arg3.handle, vaddr=res, offset=0, length=0x1000)
+  # ION_IOC_CLEAN_INV_CACHES
+  cd = ion.struct_ion_custom_data(
+    cmd=(3 << 30) | (ctypes.sizeof(flush_data) & 0x1FFF) << 16 | (ord(msm_ion.ION_IOC_MSM_MAGIC) & 0xFF) << 8 | (2 & 0xFF),
+    arg=ctypes.addressof(flush_data))
+  ret = ion_iowr(ion_fd, ION_IOC_CUSTOM, cd)
+
+  res2 = libc.mmap(0, 0x1000, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED, arg2.fd, 0)
+  """
+  hexdump(to_mv(res, 0x10))
+  os._exit(0)
+

tinygrad_repo/extra/dsp/gen.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+clang2py adsprpc_shared.h -k cdefstum -o adsprpc.py
+

tinygrad_repo/extra/dsp/hook.py

@@ -0,0 +1,101 @@
+import os
+print("from import")
+del os.environ["LD_PRELOAD"]
+import ctypes, ctypes.util
+from extra.dsp.run import install_hook, ioctl, libc, get_struct, qcom_dsp, format_struct, to_mv, hexdump
+
+@ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_long)
+def _mmap(addr, length, prot, flags, fd, offset):
+  mmap_type = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_long)
+  orig_mmap = mmap_type(ctypes.addressof(orig_mmap_mv))
+  ret = orig_mmap(addr, length, prot, flags, fd, offset)
+  # ll = os.readlink(f"/proc/self/fd/{fd}") if fd >= 0 else ""
+  print(f"mmap {addr=}, {length=}, {prot=}, {flags=}, {fd=}, {offset=} {ret=}")
+  return ret
+
+#install_hook(libc.ioctl, ioctl)
+#orig_mmap_mv = install_hook(libc.mmap, _mmap)
+print("import done")
+import mmap
+
+alloc_sizes = {}
+mmaped = {}
+
+def handle_ioctl(fd, request, argp, ret):
+  fn = os.readlink(f"/proc/self/fd/{fd}")
+  idir, size, itype, nr = (request>>30), (request>>16)&0x3FFF, (request>>8)&0xFF, request&0xFF
+
+  if fn == "/dev/ion":
+    if nr == 0:
+      st = get_struct(argp, qcom_dsp.struct_ion_allocation_data)
+      print(ret, "ION_IOC_ALLOC", format_struct(st))
+      alloc_sizes[st.handle] = st.len
+    elif nr == 1:
+      st = get_struct(argp, qcom_dsp.struct_ion_handle_data)
+      print(ret, "ION_IOC_FREE", format_struct(st))
+      if st.handle in alloc_sizes: del alloc_sizes[st.handle]
+      if st.handle in mmaped: del mmaped[st.handle]
+    elif nr == 2:
+      st = get_struct(argp, qcom_dsp.struct_ion_fd_data)
+      print(ret, "ION_IOC_MAP", format_struct(st))
+      mmaped[st.handle] = mmap.mmap(st.fd, alloc_sizes[st.handle])
+  elif fn == "/dev/adsprpc-smd":
+    assert chr(itype) == 'R'
+    if nr == 8:
+      st = ctypes.c_uint32.from_address(argp)
+      print(ret, "FASTRPC_IOCTL_GETINFO", st.value)
+    elif nr == 2:
+      st = get_struct(argp, qcom_dsp.struct_fastrpc_ioctl_mmap)
+      print(ret, "FASTRPC_IOCTL_MMAP", format_struct(st))
+    elif nr == 1:
+      # https://research.checkpoint.com/2021/pwn2own-qualcomm-dsp/
+      st = get_struct(argp, qcom_dsp.struct_fastrpc_ioctl_invoke)
+      print(ret, "FASTRPC_IOCTL_INVOKE", format_struct(st))
+      # 0xFF000000 = Method index and attribute (the highest byte)
+      # 0x00FF0000 = Number of input arguments
+      # 0x0000FF00 = Number of output arguments
+      # 0x000000F0 = Number of input handles
+      # 0x0000000F = Number of output handles
+
+      method = (st.sc>>24) & 0xFF
+      in_args = (st.sc>>16) & 0xFF
+      out_args = (st.sc>>8) & 0xFF
+      in_h = (st.sc>>4) & 0xF
+      out_h = (st.sc>>0) & 0xF
+      print(f"\tm:{method} ia:{in_args} oa:{out_args} ih:{in_h} oh:{out_h}")
+      """
+      if in_args or out_args:
+        for arg in range(in_args+out_args):
+          print(arg, format_struct(st.pra[arg]))
+          if st.pra[arg].buf.pv is not None:
+            ww = to_mv(st.pra[arg].buf.pv, st.pra[arg].buf.len)
+            hexdump(to_mv(st.pra[arg].buf.pv, st.pra[arg].buf.len)[:0x40])
+      """
+    elif nr == 6:
+      print(ret, "FASTRPC_IOCTL_INIT", format_struct(ini:=get_struct(argp, qcom_dsp.struct_fastrpc_ioctl_init)))
+      print(os.readlink(f"/proc/self/fd/{ini.filefd}"))
+      # print(bytearray(to_mv(ini.file, ini.filelen)))
+    elif nr == 7:
+      print(ret, "FASTRPC_IOCTL_INVOKE_ATTRS", format_struct(ini:=get_struct(argp, qcom_dsp.struct_fastrpc_ioctl_invoke_attrs)))
+    elif nr == 12: print(ret, "FASTRPC_IOCTL_CONTROL", format_struct(get_struct(argp, qcom_dsp.struct_fastrpc_ioctl_control)))
+    elif nr == 4:
+      st_fd = get_struct(argp, qcom_dsp.struct_fastrpc_ioctl_invoke_fd)
+      st = st_fd.inv
+      print(ret, "FASTRPC_IOCTL_INVOKE_FD", format_struct(st))
+
+      method = (st.sc>>24) & 0xFF
+      in_args = (st.sc>>16) & 0xFF
+      out_args = (st.sc>>8) & 0xFF
+      in_h = (st.sc>>4) & 0xF
+      out_h = (st.sc>>0) & 0xF
+      print(f"\tm:{method} ia:{in_args} oa:{out_args} ih:{in_h} oh:{out_h}")
+
+      if st.sc in [0x2030200, 0x3040300]:
+        for handle, mapped in mmaped.items():
+          print(f" buffer {handle} {alloc_sizes[handle]:X}")
+          with open(f"/tmp/buf_{st.sc:X}_{handle}_{alloc_sizes[handle]:X}", "wb") as f: f.write(mapped)
+    else:
+      print(f"{ret} UNPARSED {nr}")
+  else:
+    print("ioctl", f"{idir=} {size=} {itype=} {nr=} {fd=} {ret=}", fn)
+

tinygrad_repo/extra/dsp/invoke_bug.py

@@ -0,0 +1,326 @@
+from tinygrad.runtime.ops_dsp import DSPDevice
+
+kernel = """__attribute__((noinline)) void r_6_10_13_4_4_29(float* restrict __attribute__((align_value(128))) data0, const float* restrict __attribute__((align_value(128))) data1, const float* restrict __attribute__((align_value(128))) data2, const float* restrict __attribute__((align_value(128))) data3) {
+  float val0 = data1[0];
+  float val1 = data1[1];
+  float val2 = data1[2];
+  float val3 = data1[3];
+  float val4 = data1[4];
+  float val5 = data1[5];
+  float val6 = data1[6];
+  float val7 = data1[7];
+  float val8 = data1[8];
+  float val9 = data1[9];
+  float val10 = data1[10];
+  float val11 = data1[11];
+  float val12 = data1[12];
+  float val13 = data1[13];
+  float val14 = data1[14];
+  float val15 = data1[15];
+  float val16 = data1[16];
+  float val17 = data1[17];
+  float val18 = data1[18];
+  float val19 = data1[19];
+  float val20 = data1[20];
+  float val21 = data1[21];
+  float val22 = data1[22];
+  float val23 = data1[23];
+  float val24 = data1[24];
+  float val25 = data1[25];
+  float val26 = data1[26];
+  float val27 = data1[27];
+  float val28 = data1[28];
+  for (int ridx0 = 0; ridx0 < 6; ridx0++) {
+    for (int ridx1 = 0; ridx1 < 10; ridx1++) {
+      int alu0 = ((ridx0*1160)+(ridx1*4));
+      float val29 = data3[alu0+1];
+      float val30 = data3[alu0+2];
+      float val31 = data3[alu0+3];
+      float val32 = data3[alu0+40];
+      float val33 = data3[alu0+41];
+      float val34 = data3[alu0+42];
+      float val35 = data3[alu0+43];
+      float val36 = data3[alu0+80];
+      float val37 = data3[alu0+81];
+      float val38 = data3[alu0+82];
+      float val39 = data3[alu0+83];
+      float val40 = data3[alu0+120];
+      float val41 = data3[alu0+121];
+      float val42 = data3[alu0+122];
+      float val43 = data3[alu0+123];
+      float val44 = data3[alu0+160];
+      float val45 = data3[alu0+161];
+      float val46 = data3[alu0+162];
+      float val47 = data3[alu0+163];
+      float val48 = data3[alu0+200];
+      float val49 = data3[alu0+201];
+      float val50 = data3[alu0+202];
+      float val51 = data3[alu0+203];
+      float val52 = data3[alu0+240];
+      float val53 = data3[alu0+241];
+      float val54 = data3[alu0+242];
+      float val55 = data3[alu0+243];
+      float val56 = data3[alu0+280];
+      float val57 = data3[alu0+281];
+      float val58 = data3[alu0+282];
+      float val59 = data3[alu0+283];
+      float val60 = data3[alu0+320];
+      float val61 = data3[alu0+321];
+      float val62 = data3[alu0+322];
+      float val63 = data3[alu0+323];
+      float val64 = data3[alu0+360];
+      float val65 = data3[alu0+361];
+      float val66 = data3[alu0+362];
+      float val67 = data3[alu0+363];
+      float val68 = data3[alu0+400];
+      float val69 = data3[alu0+401];
+      float val70 = data3[alu0+402];
+      float val71 = data3[alu0+403];
+      float val72 = data3[alu0+440];
+      float val73 = data3[alu0+441];
+      float val74 = data3[alu0+442];
+      float val75 = data3[alu0+443];
+      float val76 = data3[alu0+480];
+      float val77 = data3[alu0+481];
+      float val78 = data3[alu0+482];
+      float val79 = data3[alu0+483];
+      float val80 = data3[alu0+520];
+      float val81 = data3[alu0+521];
+      float val82 = data3[alu0+522];
+      float val83 = data3[alu0+523];
+      float val84 = data3[alu0+560];
+      float val85 = data3[alu0+561];
+      float val86 = data3[alu0+562];
+      float val87 = data3[alu0+563];
+      float val88 = data3[alu0+600];
+      float val89 = data3[alu0+601];
+      float val90 = data3[alu0+602];
+      float val91 = data3[alu0+603];
+      float val92 = data3[alu0+640];
+      float val93 = data3[alu0+641];
+      float val94 = data3[alu0+642];
+      float val95 = data3[alu0+643];
+      float val96 = data3[alu0+680];
+      float val97 = data3[alu0+681];
+      float val98 = data3[alu0+682];
+      float val99 = data3[alu0+683];
+      float val100 = data3[alu0+720];
+      float val101 = data3[alu0+721];
+      float val102 = data3[alu0+722];
+      float val103 = data3[alu0+723];
+      float val104 = data3[alu0+760];
+      float val105 = data3[alu0+761];
+      float val106 = data3[alu0+762];
+      float val107 = data3[alu0+763];
+      float val108 = data3[alu0+800];
+      float val109 = data3[alu0+801];
+      float val110 = data3[alu0+802];
+      float val111 = data3[alu0+803];
+      float val112 = data3[alu0+840];
+      float val113 = data3[alu0+841];
+      float val114 = data3[alu0+842];
+      float val115 = data3[alu0+843];
+      float val116 = data3[alu0+880];
+      float val117 = data3[alu0+881];
+      float val118 = data3[alu0+882];
+      float val119 = data3[alu0+883];
+      float val120 = data3[alu0+920];
+      float val121 = data3[alu0+921];
+      float val122 = data3[alu0+922];
+      float val123 = data3[alu0+923];
+      float val124 = data3[alu0+960];
+      float val125 = data3[alu0+961];
+      float val126 = data3[alu0+962];
+      float val127 = data3[alu0+963];
+      float val128 = data3[alu0+1000];
+      float val129 = data3[alu0+1001];
+      float val130 = data3[alu0+1002];
+      float val131 = data3[alu0+1003];
+      float val132 = data3[alu0+1040];
+      float val133 = data3[alu0+1041];
+      float val134 = data3[alu0+1042];
+      float val135 = data3[alu0+1043];
+      float val136 = data3[alu0+1080];
+      float val137 = data3[alu0+1081];
+      float val138 = data3[alu0+1082];
+      float val139 = data3[alu0+1083];
+      float val140 = data3[alu0+1120];
+      float val141 = data3[alu0+1121];
+      float val142 = data3[alu0+1122];
+      float val143 = data3[alu0+1123];
+      float val144 = data3[alu0];
+      for (int ridx2 = 0; ridx2 < 13; ridx2++) {
+        int alu1 = (ridx2*4);
+        int alu2 = ((ridx0*2080)+(ridx1*208)+alu1);
+        float val145 = data2[alu1+1];
+        float cast0 = (float)(((val0!=val145)!=1));
+        float cast1 = (float)(((val1!=val145)!=1));
+        float cast2 = (float)(((val2!=val145)!=1));
+        float cast3 = (float)(((val3!=val145)!=1));
+        float cast4 = (float)(((val4!=val145)!=1));
+        float cast5 = (float)(((val5!=val145)!=1));
+        float cast6 = (float)(((val6!=val145)!=1));
+        float cast7 = (float)(((val7!=val145)!=1));
+        float cast8 = (float)(((val8!=val145)!=1));
+        float cast9 = (float)(((val9!=val145)!=1));
+        float cast10 = (float)(((val10!=val145)!=1));
+        float cast11 = (float)(((val11!=val145)!=1));
+        float cast12 = (float)(((val12!=val145)!=1));
+        float cast13 = (float)(((val13!=val145)!=1));
+        float cast14 = (float)(((val14!=val145)!=1));
+        float cast15 = (float)(((val15!=val145)!=1));
+        float cast16 = (float)(((val16!=val145)!=1));
+        float cast17 = (float)(((val17!=val145)!=1));
+        float cast18 = (float)(((val18!=val145)!=1));
+        float cast19 = (float)(((val19!=val145)!=1));
+        float cast20 = (float)(((val20!=val145)!=1));
+        float cast21 = (float)(((val21!=val145)!=1));
+        float cast22 = (float)(((val22!=val145)!=1));
+        float cast23 = (float)(((val23!=val145)!=1));
+        float cast24 = (float)(((val24!=val145)!=1));
+        float cast25 = (float)(((val25!=val145)!=1));
+        float cast26 = (float)(((val26!=val145)!=1));
+        float cast27 = (float)(((val27!=val145)!=1));
+        float cast28 = (float)(((val28!=val145)!=1));
+        data0[alu2+1] = ((cast0*val144)+(cast1*val32)+(cast2*val36)+(cast3*val40)+(cast4*val44)+(cast5*val48)+(cast6*val52)+(cast7*val56)+(cast8*val60)+(cast9*val64)+(cast10*val68)+(cast11*val72)+(cast12*val76)+(cast13*val80)+(cast14*val84)+(cast15*val88)+(cast16*val92)+(cast17*val96)+(cast18*val100)+(cast19*val104)+(cast20*val108)+(cast21*val112)+(cast22*val116)+(cast23*val120)+(cast24*val124)+(cast25*val128)+(cast26*val132)+(cast27*val136)+(cast28*val140));
+        data0[alu2+53] = ((cast0*val29)+(cast1*val33)+(cast2*val37)+(cast3*val41)+(cast4*val45)+(cast5*val49)+(cast6*val53)+(cast7*val57)+(cast8*val61)+(cast9*val65)+(cast10*val69)+(cast11*val73)+(cast12*val77)+(cast13*val81)+(cast14*val85)+(cast15*val89)+(cast16*val93)+(cast17*val97)+(cast18*val101)+(cast19*val105)+(cast20*val109)+(cast21*val113)+(cast22*val117)+(cast23*val121)+(cast24*val125)+(cast25*val129)+(cast26*val133)+(cast27*val137)+(cast28*val141));
+        data0[alu2+105] = ((cast0*val30)+(cast1*val34)+(cast2*val38)+(cast3*val42)+(cast4*val46)+(cast5*val50)+(cast6*val54)+(cast7*val58)+(cast8*val62)+(cast9*val66)+(cast10*val70)+(cast11*val74)+(cast12*val78)+(cast13*val82)+(cast14*val86)+(cast15*val90)+(cast16*val94)+(cast17*val98)+(cast18*val102)+(cast19*val106)+(cast20*val110)+(cast21*val114)+(cast22*val118)+(cast23*val122)+(cast24*val126)+(cast25*val130)+(cast26*val134)+(cast27*val138)+(cast28*val142));
+        data0[alu2+157] = ((cast0*val31)+(cast1*val35)+(cast2*val39)+(cast3*val43)+(cast4*val47)+(cast5*val51)+(cast6*val55)+(cast7*val59)+(cast8*val63)+(cast9*val67)+(cast10*val71)+(cast11*val75)+(cast12*val79)+(cast13*val83)+(cast14*val87)+(cast15*val91)+(cast16*val95)+(cast17*val99)+(cast18*val103)+(cast19*val107)+(cast20*val111)+(cast21*val115)+(cast22*val119)+(cast23*val123)+(cast24*val127)+(cast25*val131)+(cast26*val135)+(cast27*val139)+(cast28*val143));
+        float val146 = data2[alu1+2];
+        float cast29 = (float)(((val0!=val146)!=1));
+        float cast30 = (float)(((val1!=val146)!=1));
+        float cast31 = (float)(((val2!=val146)!=1));
+        float cast32 = (float)(((val3!=val146)!=1));
+        float cast33 = (float)(((val4!=val146)!=1));
+        float cast34 = (float)(((val5!=val146)!=1));
+        float cast35 = (float)(((val6!=val146)!=1));
+        float cast36 = (float)(((val7!=val146)!=1));
+        float cast37 = (float)(((val8!=val146)!=1));
+        float cast38 = (float)(((val9!=val146)!=1));
+        float cast39 = (float)(((val10!=val146)!=1));
+        float cast40 = (float)(((val11!=val146)!=1));
+        float cast41 = (float)(((val12!=val146)!=1));
+        float cast42 = (float)(((val13!=val146)!=1));
+        float cast43 = (float)(((val14!=val146)!=1));
+        float cast44 = (float)(((val15!=val146)!=1));
+        float cast45 = (float)(((val16!=val146)!=1));
+        float cast46 = (float)(((val17!=val146)!=1));
+        float cast47 = (float)(((val18!=val146)!=1));
+        float cast48 = (float)(((val19!=val146)!=1));
+        float cast49 = (float)(((val20!=val146)!=1));
+        float cast50 = (float)(((val21!=val146)!=1));
+        float cast51 = (float)(((val22!=val146)!=1));
+        float cast52 = (float)(((val23!=val146)!=1));
+        float cast53 = (float)(((val24!=val146)!=1));
+        float cast54 = (float)(((val25!=val146)!=1));
+        float cast55 = (float)(((val26!=val146)!=1));
+        float cast56 = (float)(((val27!=val146)!=1));
+        float cast57 = (float)(((val28!=val146)!=1));
+        data0[alu2+2] = ((cast29*val144)+(cast30*val32)+(cast31*val36)+(cast32*val40)+(cast33*val44)+(cast34*val48)+(cast35*val52)+(cast36*val56)+(cast37*val60)+(cast38*val64)+(cast39*val68)+(cast40*val72)+(cast41*val76)+(cast42*val80)+(cast43*val84)+(cast44*val88)+(cast45*val92)+(cast46*val96)+(cast47*val100)+(cast48*val104)+(cast49*val108)+(cast50*val112)+(cast51*val116)+(cast52*val120)+(cast53*val124)+(cast54*val128)+(cast55*val132)+(cast56*val136)+(cast57*val140));
+        data0[alu2+54] = ((cast29*val29)+(cast30*val33)+(cast31*val37)+(cast32*val41)+(cast33*val45)+(cast34*val49)+(cast35*val53)+(cast36*val57)+(cast37*val61)+(cast38*val65)+(cast39*val69)+(cast40*val73)+(cast41*val77)+(cast42*val81)+(cast43*val85)+(cast44*val89)+(cast45*val93)+(cast46*val97)+(cast47*val101)+(cast48*val105)+(cast49*val109)+(cast50*val113)+(cast51*val117)+(cast52*val121)+(cast53*val125)+(cast54*val129)+(cast55*val133)+(cast56*val137)+(cast57*val141));
+        data0[alu2+106] = ((cast29*val30)+(cast30*val34)+(cast31*val38)+(cast32*val42)+(cast33*val46)+(cast34*val50)+(cast35*val54)+(cast36*val58)+(cast37*val62)+(cast38*val66)+(cast39*val70)+(cast40*val74)+(cast41*val78)+(cast42*val82)+(cast43*val86)+(cast44*val90)+(cast45*val94)+(cast46*val98)+(cast47*val102)+(cast48*val106)+(cast49*val110)+(cast50*val114)+(cast51*val118)+(cast52*val122)+(cast53*val126)+(cast54*val130)+(cast55*val134)+(cast56*val138)+(cast57*val142));
+        data0[alu2+158] = ((cast29*val31)+(cast30*val35)+(cast31*val39)+(cast32*val43)+(cast33*val47)+(cast34*val51)+(cast35*val55)+(cast36*val59)+(cast37*val63)+(cast38*val67)+(cast39*val71)+(cast40*val75)+(cast41*val79)+(cast42*val83)+(cast43*val87)+(cast44*val91)+(cast45*val95)+(cast46*val99)+(cast47*val103)+(cast48*val107)+(cast49*val111)+(cast50*val115)+(cast51*val119)+(cast52*val123)+(cast53*val127)+(cast54*val131)+(cast55*val135)+(cast56*val139)+(cast57*val143));
+        float val147 = data2[alu1+3];
+        float cast58 = (float)(((val0!=val147)!=1));
+        float cast59 = (float)(((val1!=val147)!=1));
+        float cast60 = (float)(((val2!=val147)!=1));
+        float cast61 = (float)(((val3!=val147)!=1));
+        float cast62 = (float)(((val4!=val147)!=1));
+        float cast63 = (float)(((val5!=val147)!=1));
+        float cast64 = (float)(((val6!=val147)!=1));
+        float cast65 = (float)(((val7!=val147)!=1));
+        float cast66 = (float)(((val8!=val147)!=1));
+        float cast67 = (float)(((val9!=val147)!=1));
+        float cast68 = (float)(((val10!=val147)!=1));
+        float cast69 = (float)(((val11!=val147)!=1));
+        float cast70 = (float)(((val12!=val147)!=1));
+        float cast71 = (float)(((val13!=val147)!=1));
+        float cast72 = (float)(((val14!=val147)!=1));
+        float cast73 = (float)(((val15!=val147)!=1));
+        float cast74 = (float)(((val16!=val147)!=1));
+        float cast75 = (float)(((val17!=val147)!=1));
+        float cast76 = (float)(((val18!=val147)!=1));
+        float cast77 = (float)(((val19!=val147)!=1));
+        float cast78 = (float)(((val20!=val147)!=1));
+        float cast79 = (float)(((val21!=val147)!=1));
+        float cast80 = (float)(((val22!=val147)!=1));
+        float cast81 = (float)(((val23!=val147)!=1));
+        float cast82 = (float)(((val24!=val147)!=1));
+        float cast83 = (float)(((val25!=val147)!=1));
+        float cast84 = (float)(((val26!=val147)!=1));
+        float cast85 = (float)(((val27!=val147)!=1));
+        float cast86 = (float)(((val28!=val147)!=1));
+        data0[alu2+3] = ((cast58*val144)+(cast59*val32)+(cast60*val36)+(cast61*val40)+(cast62*val44)+(cast63*val48)+(cast64*val52)+(cast65*val56)+(cast66*val60)+(cast67*val64)+(cast68*val68)+(cast69*val72)+(cast70*val76)+(cast71*val80)+(cast72*val84)+(cast73*val88)+(cast74*val92)+(cast75*val96)+(cast76*val100)+(cast77*val104)+(cast78*val108)+(cast79*val112)+(cast80*val116)+(cast81*val120)+(cast82*val124)+(cast83*val128)+(cast84*val132)+(cast85*val136)+(cast86*val140));
+        data0[alu2+55] = ((cast58*val29)+(cast59*val33)+(cast60*val37)+(cast61*val41)+(cast62*val45)+(cast63*val49)+(cast64*val53)+(cast65*val57)+(cast66*val61)+(cast67*val65)+(cast68*val69)+(cast69*val73)+(cast70*val77)+(cast71*val81)+(cast72*val85)+(cast73*val89)+(cast74*val93)+(cast75*val97)+(cast76*val101)+(cast77*val105)+(cast78*val109)+(cast79*val113)+(cast80*val117)+(cast81*val121)+(cast82*val125)+(cast83*val129)+(cast84*val133)+(cast85*val137)+(cast86*val141));
+        data0[alu2+107] = ((cast58*val30)+(cast59*val34)+(cast60*val38)+(cast61*val42)+(cast62*val46)+(cast63*val50)+(cast64*val54)+(cast65*val58)+(cast66*val62)+(cast67*val66)+(cast68*val70)+(cast69*val74)+(cast70*val78)+(cast71*val82)+(cast72*val86)+(cast73*val90)+(cast74*val94)+(cast75*val98)+(cast76*val102)+(cast77*val106)+(cast78*val110)+(cast79*val114)+(cast80*val118)+(cast81*val122)+(cast82*val126)+(cast83*val130)+(cast84*val134)+(cast85*val138)+(cast86*val142));
+        data0[alu2+159] = ((cast58*val31)+(cast59*val35)+(cast60*val39)+(cast61*val43)+(cast62*val47)+(cast63*val51)+(cast64*val55)+(cast65*val59)+(cast66*val63)+(cast67*val67)+(cast68*val71)+(cast69*val75)+(cast70*val79)+(cast71*val83)+(cast72*val87)+(cast73*val91)+(cast74*val95)+(cast75*val99)+(cast76*val103)+(cast77*val107)+(cast78*val111)+(cast79*val115)+(cast80*val119)+(cast81*val123)+(cast82*val127)+(cast83*val131)+(cast84*val135)+(cast85*val139)+(cast86*val143));
+        float val148 = data2[alu1];
+        float cast87 = (float)(((val0!=val148)!=1));
+        float cast88 = (float)(((val1!=val148)!=1));
+        float cast89 = (float)(((val2!=val148)!=1));
+        float cast90 = (float)(((val3!=val148)!=1));
+        float cast91 = (float)(((val4!=val148)!=1));
+        float cast92 = (float)(((val5!=val148)!=1));
+        float cast93 = (float)(((val6!=val148)!=1));
+        float cast94 = (float)(((val7!=val148)!=1));
+        float cast95 = (float)(((val8!=val148)!=1));
+        float cast96 = (float)(((val9!=val148)!=1));
+        float cast97 = (float)(((val10!=val148)!=1));
+        float cast98 = (float)(((val11!=val148)!=1));
+        float cast99 = (float)(((val12!=val148)!=1));
+        float cast100 = (float)(((val13!=val148)!=1));
+        float cast101 = (float)(((val14!=val148)!=1));
+        float cast102 = (float)(((val15!=val148)!=1));
+        float cast103 = (float)(((val16!=val148)!=1));
+        float cast104 = (float)(((val17!=val148)!=1));
+        float cast105 = (float)(((val18!=val148)!=1));
+        float cast106 = (float)(((val19!=val148)!=1));
+        float cast107 = (float)(((val20!=val148)!=1));
+        float cast108 = (float)(((val21!=val148)!=1));
+        float cast109 = (float)(((val22!=val148)!=1));
+        float cast110 = (float)(((val23!=val148)!=1));
+        float cast111 = (float)(((val24!=val148)!=1));
+        float cast112 = (float)(((val25!=val148)!=1));
+        float cast113 = (float)(((val26!=val148)!=1));
+        float cast114 = (float)(((val27!=val148)!=1));
+        float cast115 = (float)(((val28!=val148)!=1));
+        data0[alu2+52] = ((cast87*val29)+(cast88*val33)+(cast89*val37)+(cast90*val41)+(cast91*val45)+(cast92*val49)+(cast93*val53)+(cast94*val57)+(cast95*val61)+(cast96*val65)+(cast97*val69)+(cast98*val73)+(cast99*val77)+(cast100*val81)+(cast101*val85)+(cast102*val89)+(cast103*val93)+(cast104*val97)+(cast105*val101)+(cast106*val105)+(cast107*val109)+(cast108*val113)+(cast109*val117)+(cast110*val121)+(cast111*val125)+(cast112*val129)+(cast113*val133)+(cast114*val137)+(cast115*val141));
+        data0[alu2+104] = ((cast87*val30)+(cast88*val34)+(cast89*val38)+(cast90*val42)+(cast91*val46)+(cast92*val50)+(cast93*val54)+(cast94*val58)+(cast95*val62)+(cast96*val66)+(cast97*val70)+(cast98*val74)+(cast99*val78)+(cast100*val82)+(cast101*val86)+(cast102*val90)+(cast103*val94)+(cast104*val98)+(cast105*val102)+(cast106*val106)+(cast107*val110)+(cast108*val114)+(cast109*val118)+(cast110*val122)+(cast111*val126)+(cast112*val130)+(cast113*val134)+(cast114*val138)+(cast115*val142));
+        data0[alu2+156] = ((cast87*val31)+(cast88*val35)+(cast89*val39)+(cast90*val43)+(cast91*val47)+(cast92*val51)+(cast93*val55)+(cast94*val59)+(cast95*val63)+(cast96*val67)+(cast97*val71)+(cast98*val75)+(cast99*val79)+(cast100*val83)+(cast101*val87)+(cast102*val91)+(cast103*val95)+(cast104*val99)+(cast105*val103)+(cast106*val107)+(cast107*val111)+(cast108*val115)+(cast109*val119)+(cast110*val123)+(cast111*val127)+(cast112*val131)+(cast113*val135)+(cast114*val139)+(cast115*val143));
+        data0[alu2] = ((cast87*val144)+(cast88*val32)+(cast89*val36)+(cast90*val40)+(cast91*val44)+(cast92*val48)+(cast93*val52)+(cast94*val56)+(cast95*val60)+(cast96*val64)+(cast97*val68)+(cast98*val72)+(cast99*val76)+(cast100*val80)+(cast101*val84)+(cast102*val88)+(cast103*val92)+(cast104*val96)+(cast105*val100)+(cast106*val104)+(cast107*val108)+(cast108*val112)+(cast109*val116)+(cast110*val120)+(cast111*val124)+(cast112*val128)+(cast113*val132)+(cast114*val136)+(cast115*val140));
+      }
+    }
+  }
+}"""
+
+entry = """typedef union { struct { void *pv; unsigned int len; } buf; struct { int fd; unsigned int offset; } dma; } remote_arg;
+void* HAP_mmap(void *addr, int len, int prot, int flags, int fd, long offset);
+int HAP_munmap(void *addr, int len);
+int HAP_mmap_get(int fd, void **vaddr, void **paddr);
+int HAP_mmap_put(int fd);
+unsigned long long HAP_perf_get_time_us(void);
+int entry(unsigned long long handle, unsigned int sc, remote_arg* pra) {
+if ((sc>>24) != 2) return 0;
+
+unsigned long long start = HAP_perf_get_time_us();
+for (int i = 0; i < 50; i++) {
+  void* buf = HAP_mmap(0, 1, 3, 0, pra[2].dma.fd, 0);
+  HAP_munmap(buf, 1);
+}
+*(unsigned long long *)(pra[1].buf.pv) = HAP_perf_get_time_us() - start;
+
+return 0; }
+"""
+
+if __name__ == "__main__":
+  dev = DSPDevice()
+
+  bufs = [dev.allocator.alloc(0x60000) for _ in range(4)]
+
+  only_entry = dev.compiler.compile(entry)
+  app1 = dev.runtime("test", only_entry)
+  x = app1(*bufs)
+
+  entry_n_unsued_code = dev.compiler.compile(kernel + "\n" + entry)
+  app2 = dev.runtime("test", entry_n_unsued_code)
+  x = app2(*bufs)

tinygrad_repo/extra/dsp/invoke_bug_2.py

@@ -0,0 +1,279 @@
+from tinygrad.runtime.ops_dsp import DSPDevice
+
+kernel = """__attribute__((noinline)) void r_64_4_4_64_4_4_4(float* restrict __attribute__((align_value(128))) data0, const float* restrict __attribute__((align_value(128))) data1, const float* restrict __attribute__((align_value(128))) data2, const float* restrict __attribute__((align_value(128))) data3) {
+  for (int ridx0 = 0; ridx0 < 64; ridx0++) {
+    int alu0 = (ridx0*4096);
+    for (int ridx1 = 0; ridx1 < 4; ridx1++) {
+      int alu1 = (ridx1*64);
+      for (int ridx2 = 0; ridx2 < 4; ridx2++) {
+        int alu2 = (ridx2*4);
+        int alu3 = ((ridx0*1024)+alu1+alu2);
+        int alu4 = (alu1+alu2);
+        float val0 = data3[alu4+1];
+        float val1 = data3[alu4+2];
+        float val2 = data3[alu4+3];
+        float val3 = data3[alu4+16];
+        float val4 = data3[alu4+17];
+        float val5 = data3[alu4+18];
+        float val6 = data3[alu4+19];
+        float val7 = data3[alu4+32];
+        float val8 = data3[alu4+33];
+        float val9 = data3[alu4+34];
+        float val10 = data3[alu4+35];
+        float val11 = data3[alu4+48];
+        float val12 = data3[alu4+49];
+        float val13 = data3[alu4+50];
+        float val14 = data3[alu4+51];
+        float val15 = data3[alu4];
+        float acc0 = 0.0f;
+        float acc1 = 0.0f;
+        float acc2 = 0.0f;
+        float acc3 = 0.0f;
+        float acc4 = 0.0f;
+        float acc5 = 0.0f;
+        float acc6 = 0.0f;
+        float acc7 = 0.0f;
+        float acc8 = 0.0f;
+        float acc9 = 0.0f;
+        float acc10 = 0.0f;
+        float acc11 = 0.0f;
+        float acc12 = 0.0f;
+        float acc13 = 0.0f;
+        float acc14 = 0.0f;
+        float acc15 = 0.0f;
+        float acc16 = 0.0f;
+        float acc17 = 0.0f;
+        float acc18 = 0.0f;
+        float acc19 = 0.0f;
+        float acc20 = 0.0f;
+        float acc21 = 0.0f;
+        float acc22 = 0.0f;
+        float acc23 = 0.0f;
+        float acc24 = 0.0f;
+        float acc25 = 0.0f;
+        float acc26 = 0.0f;
+        float acc27 = 0.0f;
+        float acc28 = 0.0f;
+        float acc29 = 0.0f;
+        float acc30 = 0.0f;
+        float acc31 = 0.0f;
+        float acc32 = 0.0f;
+        float acc33 = 0.0f;
+        float acc34 = 0.0f;
+        float acc35 = 0.0f;
+        float acc36 = 0.0f;
+        float acc37 = 0.0f;
+        float acc38 = 0.0f;
+        float acc39 = 0.0f;
+        float acc40 = 0.0f;
+        float acc41 = 0.0f;
+        float acc42 = 0.0f;
+        float acc43 = 0.0f;
+        float acc44 = 0.0f;
+        float acc45 = 0.0f;
+        float acc46 = 0.0f;
+        float acc47 = 0.0f;
+        float acc48 = 0.0f;
+        float acc49 = 0.0f;
+        float acc50 = 0.0f;
+        float acc51 = 0.0f;
+        float acc52 = 0.0f;
+        float acc53 = 0.0f;
+        float acc54 = 0.0f;
+        float acc55 = 0.0f;
+        float acc56 = 0.0f;
+        float acc57 = 0.0f;
+        float acc58 = 0.0f;
+        float acc59 = 0.0f;
+        float acc60 = 0.0f;
+        float acc61 = 0.0f;
+        float acc62 = 0.0f;
+        float acc63 = 0.0f;
+        for (int ridx3 = 0; ridx3 < 64; ridx3++) {
+          int alu5 = (alu0+(ridx2*256)+ridx3);
+          float val16 = data2[alu5+64];
+          float val17 = data2[alu5+128];
+          float val18 = data2[alu5+192];
+          float val19 = data2[alu5+1024];
+          float val20 = data2[alu5+1088];
+          float val21 = data2[alu5+1152];
+          float val22 = data2[alu5+1216];
+          float val23 = data2[alu5+2048];
+          float val24 = data2[alu5+2112];
+          float val25 = data2[alu5+2176];
+          float val26 = data2[alu5+2240];
+          float val27 = data2[alu5+3072];
+          float val28 = data2[alu5+3136];
+          float val29 = data2[alu5+3200];
+          float val30 = data2[alu5+3264];
+          float val31 = data2[alu5];
+          int alu6 = (alu0+(ridx1*256)+ridx3);
+          float val32 = data1[alu6+64];
+          float val33 = data1[alu6+128];
+          float val34 = data1[alu6+192];
+          float val35 = data1[alu6+1024];
+          float val36 = data1[alu6+1088];
+          float val37 = data1[alu6+1152];
+          float val38 = data1[alu6+1216];
+          float val39 = data1[alu6+2048];
+          float val40 = data1[alu6+2112];
+          float val41 = data1[alu6+2176];
+          float val42 = data1[alu6+2240];
+          float val43 = data1[alu6+3072];
+          float val44 = data1[alu6+3136];
+          float val45 = data1[alu6+3200];
+          float val46 = data1[alu6+3264];
+          float val47 = data1[alu6];
+          acc0 = (acc0+(val47*val31));
+          acc1 = (acc1+(val35*val19));
+          acc2 = (acc2+(val39*val23));
+          acc3 = (acc3+(val43*val27));
+          acc4 = (acc4+(val32*val31));
+          acc5 = (acc5+(val36*val19));
+          acc6 = (acc6+(val40*val23));
+          acc7 = (acc7+(val44*val27));
+          acc8 = (acc8+(val33*val31));
+          acc9 = (acc9+(val37*val19));
+          acc10 = (acc10+(val41*val23));
+          acc11 = (acc11+(val45*val27));
+          acc12 = (acc12+(val34*val31));
+          acc13 = (acc13+(val38*val19));
+          acc14 = (acc14+(val42*val23));
+          acc15 = (acc15+(val46*val27));
+          acc16 = (acc16+(val47*val16));
+          acc17 = (acc17+(val35*val20));
+          acc18 = (acc18+(val39*val24));
+          acc19 = (acc19+(val43*val28));
+          acc20 = (acc20+(val32*val16));
+          acc21 = (acc21+(val36*val20));
+          acc22 = (acc22+(val40*val24));
+          acc23 = (acc23+(val44*val28));
+          acc24 = (acc24+(val33*val16));
+          acc25 = (acc25+(val37*val20));
+          acc26 = (acc26+(val41*val24));
+          acc27 = (acc27+(val45*val28));
+          acc28 = (acc28+(val34*val16));
+          acc29 = (acc29+(val38*val20));
+          acc30 = (acc30+(val42*val24));
+          acc31 = (acc31+(val46*val28));
+          acc32 = (acc32+(val47*val17));
+          acc33 = (acc33+(val35*val21));
+          acc34 = (acc34+(val39*val25));
+          acc35 = (acc35+(val43*val29));
+          acc36 = (acc36+(val32*val17));
+          acc37 = (acc37+(val36*val21));
+          acc38 = (acc38+(val40*val25));
+          acc39 = (acc39+(val44*val29));
+          acc40 = (acc40+(val33*val17));
+          acc41 = (acc41+(val37*val21));
+          acc42 = (acc42+(val41*val25));
+          acc43 = (acc43+(val45*val29));
+          acc44 = (acc44+(val34*val17));
+          acc45 = (acc45+(val38*val21));
+          acc46 = (acc46+(val42*val25));
+          acc47 = (acc47+(val46*val29));
+          acc48 = (acc48+(val47*val18));
+          acc49 = (acc49+(val35*val22));
+          acc50 = (acc50+(val39*val26));
+          acc51 = (acc51+(val43*val30));
+          acc52 = (acc52+(val32*val18));
+          acc53 = (acc53+(val36*val22));
+          acc54 = (acc54+(val40*val26));
+          acc55 = (acc55+(val44*val30));
+          acc56 = (acc56+(val33*val18));
+          acc57 = (acc57+(val37*val22));
+          acc58 = (acc58+(val41*val26));
+          acc59 = (acc59+(val45*val30));
+          acc60 = (acc60+(val34*val18));
+          acc61 = (acc61+(val38*val22));
+          acc62 = (acc62+(val42*val26));
+          acc63 = (acc63+(val46*val30));
+        }
+        data0[alu3] = ((acc0*0.125f)+val15);
+        data0[alu3+256] = ((acc1*0.125f)+val15);
+        data0[alu3+512] = ((acc2*0.125f)+val15);
+        data0[alu3+768] = ((acc3*0.125f)+val15);
+        data0[alu3+16] = ((acc4*0.125f)+val3);
+        data0[alu3+272] = ((acc5*0.125f)+val3);
+        data0[alu3+528] = ((acc6*0.125f)+val3);
+        data0[alu3+784] = ((acc7*0.125f)+val3);
+        data0[alu3+32] = ((acc8*0.125f)+val7);
+        data0[alu3+288] = ((acc9*0.125f)+val7);
+        data0[alu3+544] = ((acc10*0.125f)+val7);
+        data0[alu3+800] = ((acc11*0.125f)+val7);
+        data0[alu3+48] = ((acc12*0.125f)+val11);
+        data0[alu3+304] = ((acc13*0.125f)+val11);
+        data0[alu3+560] = ((acc14*0.125f)+val11);
+        data0[alu3+816] = ((acc15*0.125f)+val11);
+        data0[alu3+1] = ((acc16*0.125f)+val0);
+        data0[alu3+257] = ((acc17*0.125f)+val0);
+        data0[alu3+513] = ((acc18*0.125f)+val0);
+        data0[alu3+769] = ((acc19*0.125f)+val0);
+        data0[alu3+17] = ((acc20*0.125f)+val4);
+        data0[alu3+273] = ((acc21*0.125f)+val4);
+        data0[alu3+529] = ((acc22*0.125f)+val4);
+        data0[alu3+785] = ((acc23*0.125f)+val4);
+        data0[alu3+33] = ((acc24*0.125f)+val8);
+        data0[alu3+289] = ((acc25*0.125f)+val8);
+        data0[alu3+545] = ((acc26*0.125f)+val8);
+        data0[alu3+801] = ((acc27*0.125f)+val8);
+        data0[alu3+49] = ((acc28*0.125f)+val12);
+        data0[alu3+305] = ((acc29*0.125f)+val12);
+        data0[alu3+561] = ((acc30*0.125f)+val12);
+        data0[alu3+817] = ((acc31*0.125f)+val12);
+        data0[alu3+2] = ((acc32*0.125f)+val1);
+        data0[alu3+258] = ((acc33*0.125f)+val1);
+        data0[alu3+514] = ((acc34*0.125f)+val1);
+        data0[alu3+770] = ((acc35*0.125f)+val1);
+        data0[alu3+18] = ((acc36*0.125f)+val5);
+        data0[alu3+274] = ((acc37*0.125f)+val5);
+        data0[alu3+530] = ((acc38*0.125f)+val5);
+        data0[alu3+786] = ((acc39*0.125f)+val5);
+        data0[alu3+34] = ((acc40*0.125f)+val9);
+        data0[alu3+290] = ((acc41*0.125f)+val9);
+        data0[alu3+546] = ((acc42*0.125f)+val9);
+        data0[alu3+802] = ((acc43*0.125f)+val9);
+        data0[alu3+50] = ((acc44*0.125f)+val13);
+        data0[alu3+306] = ((acc45*0.125f)+val13);
+        data0[alu3+562] = ((acc46*0.125f)+val13);
+        data0[alu3+818] = ((acc47*0.125f)+val13);
+        data0[alu3+3] = ((acc48*0.125f)+val2);
+        data0[alu3+259] = ((acc49*0.125f)+val2);
+        data0[alu3+515] = ((acc50*0.125f)+val2);
+        data0[alu3+771] = ((acc51*0.125f)+val2);
+        data0[alu3+19] = ((acc52*0.125f)+val6);
+        data0[alu3+275] = ((acc53*0.125f)+val6);
+        data0[alu3+531] = ((acc54*0.125f)+val6);
+        data0[alu3+787] = ((acc55*0.125f)+val6);
+        data0[alu3+35] = ((acc56*0.125f)+val10);
+        data0[alu3+291] = ((acc57*0.125f)+val10);
+        data0[alu3+547] = ((acc58*0.125f)+val10);
+        data0[alu3+803] = ((acc59*0.125f)+val10);
+        data0[alu3+51] = ((acc60*0.125f)+val14);
+        data0[alu3+307] = ((acc61*0.125f)+val14);
+        data0[alu3+563] = ((acc62*0.125f)+val14);
+        data0[alu3+819] = ((acc63*0.125f)+val14);
+      }
+    }
+  }
+}
+"""
+
+entry = """unsigned long long HAP_perf_get_time_us(void);
+int entry(unsigned long long handle, unsigned int sc, void* pra) {
+return HAP_perf_get_time_us() == 1 ? 4 : 0;
+}
+"""
+
+if __name__ == "__main__":
+  dev = DSPDevice()
+
+  bufs = [dev.allocator.alloc(0x60000) for _ in range(4)]
+
+  only_entry = dev.compiler.compile(entry)
+  app1 = dev.runtime("test", only_entry)
+  x = app1(*bufs)
+
+  entry_n_unsued_code = dev.compiler.compile(kernel + "\n" + entry)
+  app2 = dev.runtime("test", entry_n_unsued_code)
+  x = app2(*bufs)

tinygrad_repo/extra/dsp/opt.py

@@ -0,0 +1,27 @@
+from tinygrad import Device
+
+# PATH=/opt/homebrew/opt/llvm/bin:$PATH python3 extra/dsp/opt.py
+
+if __name__ == "__main__":
+  compiler = Device["DSP"].compiler
+
+  lib = compiler.compile("""
+typedef long HVX_Vector __attribute__((__vector_size__(128))) __attribute__ ((aligned(128)));
+typedef long HVX_VectorPair __attribute__((__vector_size__(256))) __attribute__ ((aligned(256)));
+
+void test(unsigned char *c, unsigned char *a, unsigned char *b) {
+  HVX_Vector t0 = *(HVX_Vector*)a;
+  //HVX_VectorPair t1 = *((HVX_VectorPair*)b);
+  HVX_Vector acc = __builtin_HEXAGON_V6_vd0_128B();
+  for (int i = 0; i < 128; i++) {
+    //__builtin_HEXAGON_V6_lvsplatb_128B(t0[i])
+    //acc += __builtin_HEXAGON_V6_lvsplatb_128B(t0[i]) * t1;
+    //acc += t0[i] * t1;
+    unsigned int t1 = ((unsigned int *)b)[i];
+    //acc = __builtin_HEXAGON_V6_vrmpyub_acc_128B(acc, t0, t1);
+    acc = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc, t0, t1);
+  }
+  *((HVX_Vector*)c) = acc;
+}""")
+
+  compiler.disassemble(lib)

tinygrad_repo/extra/dsp/run.py

@@ -0,0 +1,152 @@
+# mypy: ignore-errors
+#!/usr/bin/env python3
+import os, ctypes, ctypes.util, struct, platform, time
+from tinygrad.runtime.autogen import libc, qcom_dsp
+def to_mv(ptr, sz) -> memoryview: return memoryview(ctypes.cast(ptr, ctypes.POINTER(ctypes.c_uint8 * sz)).contents).cast("B")
+from hexdump import hexdump
+
+def get_struct(argp, stype):
+  return ctypes.cast(ctypes.c_void_p(argp), ctypes.POINTER(stype)).contents
+
+def format_struct(s):
+  sdats = []
+  for field in s._fields_:
+    dat = getattr(s, field[0])
+    if isinstance(dat, int): sdats.append(f"{field[0]}:0x{dat:X}")
+    elif hasattr(dat, "_fields_"): sdats.append((field[0], format_struct(dat)))
+    elif field[0] == "PADDING_0": pass
+    else: sdats.append(f"{field[0]}:{dat}")
+  return sdats
+
+@ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int, ctypes.c_ulong, ctypes.c_void_p)
+def ioctl(fd, request, argp):
+  fn = os.readlink(f"/proc/self/fd/{fd}")
+  idir, size, itype, nr = (request>>30), (request>>16)&0x3FFF, (request>>8)&0xFF, request&0xFF
+
+  if fn == "/dev/adsprpc-smd":
+    if nr == 1:
+      st = get_struct(argp, qcom_dsp.struct_fastrpc_ioctl_invoke)
+      method = (st.sc>>24) & 0xFF
+      in_args = (st.sc>>16) & 0xFF
+      out_args = (st.sc>>8) & 0xFF
+      if out_args:
+        for arg in range(in_args, in_args+out_args):
+          ctypes.memset(st.pra[arg].buf.pv, 0, st.pra[arg].buf.len)
+
+  # print("enter", libc.gettid())
+  ret = libc.syscall(0x1d, ctypes.c_int(fd), ctypes.c_ulong(request), ctypes.c_void_p(argp))
+  # print("done", libc.gettid())
+  if fn == "/dev/ion":
+    if nr == 0:
+      st = get_struct(argp, qcom_dsp.struct_ion_allocation_data)
+      print(ret, "ION_IOC_ALLOC", format_struct(st))
+    elif nr == 1:
+      st = get_struct(argp, qcom_dsp.struct_ion_handle_data)
+      print(ret, "ION_IOC_FREE", format_struct(st))
+    elif nr == 2:
+      st = get_struct(argp, qcom_dsp.struct_ion_fd_data)
+      print(ret, "ION_IOC_MAP", format_struct(st))
+  elif fn == "/dev/adsprpc-smd":
+    assert chr(itype) == 'R'
+    if nr == 8:
+      st = ctypes.c_uint32.from_address(argp)
+      print(ret, "FASTRPC_IOCTL_GETINFO", st.value)
+    elif nr == 2:
+      st = get_struct(argp, qcom_dsp.struct_fastrpc_ioctl_mmap)
+      print(ret, "FASTRPC_IOCTL_MMAP", format_struct(st))
+    elif nr == 1:
+      # https://research.checkpoint.com/2021/pwn2own-qualcomm-dsp/
+      st = get_struct(argp, qcom_dsp.struct_fastrpc_ioctl_invoke)
+      print(ret, "FASTRPC_IOCTL_INVOKE", format_struct(st))
+      # 0xFF000000 = Method index and attribute (the highest byte)
+      # 0x00FF0000 = Number of input arguments
+      # 0x0000FF00 = Number of output arguments
+      # 0x000000F0 = Number of input handles
+      # 0x0000000F = Number of output handles
+
+      method = (st.sc>>24) & 0xFF
+      in_args = (st.sc>>16) & 0xFF
+      out_args = (st.sc>>8) & 0xFF
+      in_h = (st.sc>>4) & 0xF
+      out_h = (st.sc>>0) & 0xF
+      print(f"\tm:{method} ia:{in_args} oa:{out_args} ih:{in_h} oh:{out_h}")
+      if in_args or out_args:
+        for arg in range(in_args+out_args):
+          print(arg, format_struct(st.pra[arg]))
+          if st.pra[arg].buf.pv is not None:
+            ww = to_mv(st.pra[arg].buf.pv, st.pra[arg].buf.len)
+            hexdump(to_mv(st.pra[arg].buf.pv, st.pra[arg].buf.len)[:0x40])
+    elif nr == 6:
+      print(ret, "FASTRPC_IOCTL_INIT", format_struct(ini:=get_struct(argp, qcom_dsp.struct_fastrpc_ioctl_init)))
+      print(os.readlink(f"/proc/self/fd/{ini.filefd}"))
+      # print(bytearray(to_mv(ini.file, ini.filelen)))
+    elif nr == 7:
+      print(ret, "FASTRPC_IOCTL_INVOKE_ATTRS", format_struct(ini:=get_struct(argp, qcom_dsp.struct_fastrpc_ioctl_invoke_attrs)))
+    elif nr == 12: print(ret, "FASTRPC_IOCTL_CONTROL", format_struct(get_struct(argp, qcom_dsp.struct_fastrpc_ioctl_control)))
+    else:
+      print(f"{ret} UNPARSED {nr}")
+  else:
+    print("ioctl", f"{idir=} {size=} {itype=} {nr=} {fd=} {ret=}", fn)
+  return ret
+
+def install_hook(c_function, python_function):
+  orig_func = (ctypes.c_char*4096)()
+  python_function_addr = ctypes.cast(ctypes.byref(python_function), ctypes.POINTER(ctypes.c_ulong)).contents.value
+  # AARCH64 trampoline to ioctl
+  # 0x0000000000000000:  70 00 00 10    adr x16, #0xc
+  # 0x0000000000000004:  10 02 40 F9    ldr x16, [x16]
+  # 0x0000000000000008:  00 02 1F D6    br  x16
+  tramp = b"\x70\x00\x00\x10\x10\x02\x40\xf9\x00\x02\x1f\xd6"
+  tramp += struct.pack("Q", python_function_addr)
+
+  # get real ioctl address
+  ioctl_address = ctypes.cast(ctypes.byref(c_function), ctypes.POINTER(ctypes.c_ulong))
+
+  # hook ioctl
+  ret = libc.mprotect(ctypes.c_ulong((ioctl_address.contents.value//0x1000)*0x1000), 0x2000, 7)
+  assert ret == 0
+  ret = libc.mprotect(ctypes.c_ulong((ctypes.addressof(orig_func)//0x1000)*0x1000), 0x3000, 7)
+  assert ret == 0
+  libc.memcpy(orig_func, ioctl_address.contents, 0x1000)
+  libc.memcpy(ioctl_address.contents, ctypes.create_string_buffer(tramp), len(tramp))
+  return orig_func
+
+libc = ctypes.CDLL(ctypes.util.find_library("libc"))
+#install_hook(libc.ioctl, ioctl)
+adsp = ctypes.CDLL(ctypes.util.find_library("adsprpc"))
+
+def send_rpc_invoke(filename):
+  pass
+
+if __name__ == "__main__":
+  print("calculator_open")
+  # /dsp/cdsp/fastrpc_shell_3
+
+  handle = ctypes.c_int64(-1)
+  z = adsp.remote_handle64_open(ctypes.create_string_buffer(b"file:///libcalculator_skel.so?calculator_skel_handle_invoke&_modver=1.0&_dom=cdsp"),
+                            ctypes.byref(handle))
+
+  print("handle", z, hex(handle.value))
+  assert handle.value != -1
+  test = (ctypes.c_int32 * 100)()
+  for i in range(100): test[i] = i
+  print("calculator_sum")
+  pra = (qcom_dsp.union_remote_arg64 * 3)()
+  #arg_0 = ctypes.c_int32(100)
+  arg_0 = ctypes.c_int32(100)
+  arg_2 = ctypes.c_int64(-1)
+  pra[0].buf.pv = ctypes.addressof(arg_0)
+  pra[0].buf.len = 4
+  pra[1].buf.pv = ctypes.addressof(test)
+  pra[1].buf.len = 0x190
+  pra[2].buf.pv = ctypes.addressof(arg_2)
+  pra[2].buf.len = 8
+  adsp.remote_handle64_invoke(handle, (2<<24) | (2<<16) | (1<<8), pra)
+  print(arg_2.value)
+  print("done")
+
+  print("closing")
+  x = adsp.remote_handle64_close(handle)
+  print(x)
+  print("dun")
+  os._exit(0)

tinygrad_repo/extra/dsp/run_3.py

@@ -0,0 +1,312 @@
+#!/usr/bin/env python3
+import os, ctypes, ctypes.util, struct, platform, pathlib, contextlib, mmap, array
+from threading import Thread
+from tinygrad.runtime.autogen import qcom_dsp
+from tinygrad.helpers import round_up, mv_address, to_mv
+from hexdump import hexdump
+
+def get_struct(argp, stype):
+  return ctypes.cast(ctypes.c_void_p(argp), ctypes.POINTER(stype)).contents
+
+def format_struct(s):
+  sdats = []
+  for field in s._fields_:
+    dat = getattr(s, field[0])
+    if isinstance(dat, int): sdats.append(f"{field[0]}:0x{dat:X}")
+    elif hasattr(dat, "_fields_"): sdats.append((field[0], format_struct(dat)))
+    elif field[0] == "PADDING_0": pass
+    else: sdats.append(f"{field[0]}:{dat}")
+  return sdats
+
+@ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int, ctypes.c_ulong, ctypes.c_void_p)
+def ioctl(fd, request, argp):
+  fn = os.readlink(f"/proc/self/fd/{fd}")
+  idir, size, itype, nr = (request>>30), (request>>16)&0x3FFF, (request>>8)&0xFF, request&0xFF
+
+  # print("enter", libc.gettid())
+  ret = libc.syscall(0x1d, ctypes.c_int(fd), ctypes.c_ulong(request), ctypes.c_void_p(argp))
+  # print("done", libc.gettid())
+  if fn == "/dev/ion":
+    if nr == 0:
+      st = get_struct(argp, qcom_dsp.struct_ion_allocation_data)
+      print(ret, "ION_IOC_ALLOC", format_struct(st))
+    elif nr == 1:
+      st = get_struct(argp, qcom_dsp.struct_ion_handle_data)
+      print(ret, "ION_IOC_FREE", format_struct(st))
+    elif nr == 2:
+      st = get_struct(argp, qcom_dsp.struct_ion_fd_data)
+      print(ret, "ION_IOC_MAP", format_struct(st))
+  elif fn == "/dev/adsprpc-smd":
+    assert chr(itype) == 'R'
+    if nr == 8:
+      st = ctypes.c_uint32.from_address(argp)
+      print(ret, "FASTRPC_IOCTL_GETINFO", st.value)
+    elif nr == 2:
+      st = get_struct(argp, qcom_dsp.struct_fastrpc_ioctl_mmap)
+      print(ret, "FASTRPC_IOCTL_MMAP", format_struct(st))
+    elif nr == 1:
+      # https://research.checkpoint.com/2021/pwn2own-qualcomm-dsp/
+      st = get_struct(argp, qcom_dsp.struct_fastrpc_ioctl_invoke)
+      print(ret, "FASTRPC_IOCTL_INVOKE", format_struct(st))
+      # 0xFF000000 = Method index and attribute (the highest byte)
+      # 0x00FF0000 = Number of input arguments
+      # 0x0000FF00 = Number of output arguments
+      # 0x000000F0 = Number of input handles
+      # 0x0000000F = Number of output handles
+
+      method = (st.sc>>24) & 0xFF
+      in_args = (st.sc>>16) & 0xFF
+      out_args = (st.sc>>8) & 0xFF
+      in_h = (st.sc>>4) & 0xF
+      out_h = (st.sc>>0) & 0xF
+      print(f"\tm:{method} ia:{in_args} oa:{out_args} ih:{in_h} oh:{out_h}")
+      if in_args or out_args:
+        for arg in range(in_args+out_args):
+          print(arg, format_struct(st.pra[arg]))
+          # print(arg, f"arg (0x{st.pra[arg].buf.pv:X} len=0x{st.pra[arg].buf.len:X})")
+          # print("input" if arg < in_args else "output", f"arg (0x{st.pra[arg].buf.pv:X} len=0x{st.pra[arg].buf.len:X})")
+          if st.pra[arg].buf.pv is not None:
+            # if st.pra[arg].buf.len == 0x258:
+            #   print(bytearray(to_mv(st.pra[arg].buf.pv, st.pra[arg].buf.len)))
+            if st.pra[arg].buf.len == 0x68:
+              print(bytearray(to_mv(st.pra[arg].buf.pv, st.pra[arg].buf.len)))
+
+            cut = 0x2000 if st.pra[arg].buf.len == 0x2000 or st.pra[arg].buf.len == 0x258 else 0x100
+            ww = to_mv(st.pra[arg].buf.pv, st.pra[arg].buf.len)
+            hexdump(to_mv(st.pra[arg].buf.pv, st.pra[arg].buf.len)[:cut])
+
+            # if st.pra[arg].buf.len == 0x1000 and ww[0x30] == 0x6e:
+            #   z = ww.cast('Q')[1] + 0x7F00000000
+            #   print("DOO")
+            #   hexdump(to_mv(z, 0x200))
+      #print(format_struct(st.pra)))
+    elif nr == 6:
+      print(ret, "FASTRPC_IOCTL_INIT", format_struct(ini:=get_struct(argp, qcom_dsp.struct_fastrpc_ioctl_init)))
+      print(os.readlink(f"/proc/self/fd/{ini.filefd}"))
+      # print(bytearray(to_mv(ini.file, ini.filelen)))
+    elif nr == 7:
+      print(ret, "FASTRPC_IOCTL_INVOKE_ATTRS", format_struct(ini:=get_struct(argp, qcom_dsp.struct_fastrpc_ioctl_invoke_attrs)))
+    elif nr == 12: print(ret, "FASTRPC_IOCTL_CONTROL", format_struct(get_struct(argp, qcom_dsp.struct_fastrpc_ioctl_control)))
+    else:
+      print(f"{ret} UNPARSED {nr}")
+  else:
+    print("ioctl", f"{idir=} {size=} {itype=} {nr=} {fd=} {ret=}", fn)
+  return ret
+
+def install_hook(c_function, python_function):
+  orig_func = (ctypes.c_char*4096)()
+  python_function_addr = ctypes.cast(ctypes.byref(python_function), ctypes.POINTER(ctypes.c_ulong)).contents.value
+  # AARCH64 trampoline to ioctl
+  # 0x0000000000000000:  70 00 00 10    adr x16, #0xc
+  # 0x0000000000000004:  10 02 40 F9    ldr x16, [x16]
+  # 0x0000000000000008:  00 02 1F D6    br  x16
+  tramp = b"\x70\x00\x00\x10\x10\x02\x40\xf9\x00\x02\x1f\xd6"
+  tramp += struct.pack("Q", python_function_addr)
+
+  # get real ioctl address
+  ioctl_address = ctypes.cast(ctypes.byref(c_function), ctypes.POINTER(ctypes.c_ulong))
+
+  # hook ioctl
+  ret = libc.mprotect(ctypes.c_ulong((ioctl_address.contents.value//0x1000)*0x1000), 0x2000, 7)
+  assert ret == 0
+  ret = libc.mprotect(ctypes.c_ulong((ctypes.addressof(orig_func)//0x1000)*0x1000), 0x3000, 7)
+  assert ret == 0
+  libc.memcpy(orig_func, ioctl_address.contents, 0x1000)
+  libc.memcpy(ioctl_address.contents, ctypes.create_string_buffer(tramp), len(tramp))
+  return orig_func
+
+libc = ctypes.CDLL(ctypes.util.find_library("libc"))
+install_hook(libc.ioctl, ioctl)
+from tinygrad.runtime.autogen import libc
+
+# adsp = ctypes.CDLL(ctypes.util.find_library("adsprpc"))
+# print(adsp)
+
+def rpc_invoke(rpcfd, handle, method, ins=None, outs=None):
+  if ins or outs:
+    ins = ins or list()
+    outs = outs or list()
+    pra = (qcom_dsp.union_remote_arg * (len(ins) + len(outs)))()
+    for i,mv in enumerate(ins + outs):
+      if isinstance(mv, memoryview):
+        pra[i].buf.pv = mv_address(mv) if mv.nbytes > 0 else 0
+        pra[i].buf.len = mv.nbytes
+      else: assert False, "not supported"
+    # pra = (qcom_dsp.union_remote_arg * (len(ins) + len(outs))).from_address(ctypes.addressof(pra))
+  else:
+    pra = None
+    ins = ins or list()
+    outs = outs or list()
+
+  sc = (method << 24) | (len(ins) << 16) | (len(outs) << 8)
+  return qcom_dsp.FASTRPC_IOCTL_INVOKE(rpcfd, handle=handle, sc=sc, pra=pra)
+
+def listner_worker():
+  context = 0
+  handle = 0xffffffff
+  msg_send = memoryview(bytearray(0x10)).cast('I')
+  msg_recv = memoryview(bytearray(0x10)).cast('I')
+  out_buf = memoryview(bytearray(0x1000)).cast('I')
+  in_buf = memoryview(bytearray(0x1000)).cast('I')
+
+  prev_res = 0xffffffff
+  out_buf_size = 0
+
+  req_args = (qcom_dsp.union_remote_arg * 4)()
+  req_args[0].buf = qcom_dsp.struct_remote_buf(pv=mv_address(msg_send), len=0x10)
+  req_args[1].buf = qcom_dsp.struct_remote_buf(pv=mv_address(out_buf), len=0x1000)
+  req_args[2].buf = qcom_dsp.struct_remote_buf(pv=mv_address(msg_recv), len=0x10)
+  req_args[3].buf = qcom_dsp.struct_remote_buf(pv=mv_address(in_buf), len=0x1000)
+
+  while True:
+    msg_send[0] = context
+    msg_send[1] = prev_res
+    msg_send[2] = out_buf_size
+    msg_send[3] = 0x1000
+
+    req_args[1].buf.len = out_buf_size
+    qcom_dsp.FASTRPC_IOCTL_INVOKE(rpcfd, handle=0x3, sc=0x04020200, pra=req_args) # listener
+
+    context = msg_recv[0]
+    handle = msg_recv[1]
+    sc = msg_recv[2]
+    inbufs = (sc >> 16) & 0xff
+    outbufs = (sc >> 8) & 0xff
+
+    in_args, out_args = [], []
+    ptr = mv_address(in_buf)
+    for i in range(inbufs):
+      sz = to_mv(ptr, 4).cast('I')[0]
+      obj_ptr = round_up(ptr + 4, 8)
+      in_args.append(to_mv(obj_ptr, sz))
+      ptr = obj_ptr + sz
+
+    ctypes.memset(mv_address(out_buf), 0, 0x1000)
+    ptr_out = mv_address(out_buf)
+    for i in range(outbufs):
+      sz = to_mv(ptr, 4).cast('I')[0]
+      ptr += 4
+
+      to_mv(ptr_out, 4).cast('I')[0] = sz
+      obj_ptr = round_up(ptr_out + 4, 8)
+
+      out_args.append(to_mv(obj_ptr, sz))
+      ptr_out = obj_ptr + sz
+
+    out_buf_size = ptr_out - mv_address(out_buf)
+
+    if sc == 0x20200: # greating?
+      prev_res = 0
+    elif sc == 0x13050100: # open
+      # for a in in_args: hexdump(a)
+      try:
+        fd = os.open(in_args[3].tobytes()[:-1].decode(), os.O_RDONLY)
+        out_args[0].cast('I')[0] = fd
+        prev_res = 0
+      except: prev_res = 2
+    elif sc == 0x9010000: # seek
+      res = os.lseek(in_args[0].cast('I')[0], in_args[0].cast('I')[1], in_args[0].cast('I')[2])
+      prev_res = 0 if res >= 0 else res
+    elif sc == 0x4010200: # read
+      buf = os.read(in_args[0].cast('I')[0], in_args[0].cast('I')[1])
+      out_args[1][:len(buf)] = buf
+      out_args[0].cast('I')[0] = len(buf)
+      out_args[0].cast('I')[1] = int(len(buf) == 0)
+      prev_res = 0
+    elif sc == 0x3010000: # close
+      os.close(in_args[0].cast('I')[0])
+      prev_res = 0
+    elif sc == 0x1f020100: # stat
+      # try:
+      stat = os.stat(in_args[1].tobytes()[:-1].decode())
+      out_stat = out_args[0].cast('Q')
+      out_stat[1] = stat.st_dev
+      out_stat[2] = stat.st_ino
+      out_stat[3] = stat.st_mode | (stat.st_nlink << 32)
+      out_stat[4] = stat.st_rdev
+      out_stat[5] = stat.st_size
+      # print(stat, stat.st_rdev)
+      # assert False
+      prev_res = 0
+      # except: prev_res = 2
+    elif sc == 0x2010100:
+      heapid = in_args[0].cast('I')[0]
+      lflags = in_args[0].cast('I')[1]
+      rflags = in_args[0].cast('I')[2]
+      assert rflags == 0x1000
+
+      # print(in_args[0])
+
+      # print("WOOW", in_args[0].cast('Q')[2])
+      # print("WOOW2", in_args[0].cast('Q')[2])
+      # print("WOOW3", in_args[0].cast('Q')[3])
+      # print("WOOW3", in_args[0].cast('Q')[3])
+
+      vin = in_args[0].cast('Q')[2]
+      sz = in_args[0].cast('Q')[3]
+      # vin = to_mv(in_args[0].cast('Q')[2], 8).cast('Q')[0]
+      # sz = to_mv(in_args[0].cast('Q')[3], 8).cast('Q')[0]
+
+      st = qcom_dsp.FASTRPC_IOCTL_MMAP(rpcfd, fd=-1, flags=rflags, vaddrin=0, size=sz)
+      out_args[0].cast('Q')[0] = 0
+      out_args[0].cast('Q')[1] = st.vaddrout
+      prev_res = 0
+    else: raise RuntimeError(f"Unknown {sc=:X}")
+
+if __name__ == "__main__":
+  ionfd = os.open('/dev/ion', os.O_RDONLY)
+  rpcfd = os.open('/dev/adsprpc-smd', os.O_RDONLY | os.O_NONBLOCK)
+
+  with contextlib.suppress(RuntimeError, OSError): qcom_dsp.ION_IOC_FREE(ionfd, handle=0)
+  info = qcom_dsp.FASTRPC_IOCTL_GETINFO(rpcfd, 3)
+  # x = qcom_dsp.FASTRPC_IOCTL_SETMODE(rpcfd, 0, __force_as_val=True)
+
+  # init shell?
+  fastrpc_shell = memoryview(bytearray(pathlib.Path('/vendor/dsp/cdsp/fastrpc_shell_3').read_bytes()))
+  shell_mem = qcom_dsp.ION_IOC_ALLOC(ionfd, len=round_up(fastrpc_shell.nbytes, 0x1000), align=0x1000, heap_id_mask=0x2000000, flags=0x1)
+  shell_mapped = qcom_dsp.ION_IOC_MAP(ionfd, handle=shell_mem.handle)
+  fastrpc_shell_addr = libc.mmap(0, shell_mem.len, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED, shell_mapped.fd, 0)
+
+  ctypes.memmove(fastrpc_shell_addr, mv_address(fastrpc_shell), fastrpc_shell.nbytes)
+  # ctypes.memset(fastrpc_shell_addr, 0x0, 0xd6000)
+  # print(hex(fastrpc_shell_addr))
+
+  ctrls = qcom_dsp.FASTRPC_IOCTL_CONTROL(rpcfd, req=0x3)
+
+  init = qcom_dsp.FASTRPC_IOCTL_INIT(rpcfd, flags=0x1, file=fastrpc_shell_addr, filelen=fastrpc_shell.nbytes, filefd=shell_mapped.fd)
+  print("init shell done", shell_mapped.fd)
+
+  # TODO: unmap here
+  # qcom_dsp.ION_IOC_FREE(ionfd, handle=shell_mem.handle)
+
+  rpc_invoke(rpcfd, handle=3, method=3)
+
+  thread = Thread(target=listner_worker)
+  thread.start()
+
+  a1 = memoryview(bytearray(b'\x52\x00\x00\x00\xFF\x00\x00\x00'))
+  a2 = memoryview(bytearray(b"file:///libcalculator_skel.so?calculator_skel_handle_invoke&_modver=1.0&_dom=cdsp\0"))
+  o1 = memoryview(bytearray(0x8))
+  o2 = memoryview(bytearray(0xff))
+  z = rpc_invoke(rpcfd, handle=0, method=0, ins=[a1, a2], outs=[o1, o2])
+  prg_handle = o1.cast('I')[0]
+
+  # test
+  test = (ctypes.c_int32 * 100)()
+  for i in range(100): test[i] = i
+  print("calculator_sum")
+  pra = (qcom_dsp.union_remote_arg * 3)()
+  #arg_0 = ctypes.c_int32(100)
+  arg_0 = ctypes.c_int32(100)
+  arg_2 = ctypes.c_int64(-1)
+  pra[0].buf.pv = ctypes.addressof(arg_0)
+  pra[0].buf.len = 4
+  pra[1].buf.pv = ctypes.addressof(test)
+  pra[1].buf.len = 0x190
+  pra[2].buf.pv = ctypes.addressof(arg_2)
+  pra[2].buf.len = 8
+  qcom_dsp.FASTRPC_IOCTL_INVOKE(rpcfd, handle=prg_handle, sc=(2<<24) | (2<<16) | (1<<8), pra=pra)
+
+  print(arg_2.value)
+  print("done")
+  os._exit(0)

tinygrad_repo/extra/dsp/snpe.sh

@@ -0,0 +1,11 @@
+#!/bin/bash -e
+echo "building"
+gcc -shared -fPIC -o preload_python.so preload.c -L/usr/local/pyenv/versions/3.11.4/lib -lpython3.11 -I/usr/local/pyenv/versions/3.11.4/include/python3.11
+echo "compiled"
+export LD_LIBRARY_PATH="/usr/local/pyenv/versions/3.11.4/lib;/data/snpe"
+export LD_PRELOAD="$PWD/preload_python.so"
+export PYTHONPATH="/data/tinygrad"
+cd /data/snpe
+#ADSP_LIBRARY_PATH="." strace -f -e ioctl ./snpe-net-run --container MobileNetV2.dlc --input_list hello --use_dsp
+ADSP_LIBRARY_PATH="." ./snpe-net-run --container MobileNetV2.dlc --input_list hello --use_dsp
+

tinygrad_repo/extra/dsp/snpe_logs/dlc_info_2

@@ -0,0 +1,715 @@
+DLC info for: /home/batman/xx/ml_tools/snpe/snpe-1.61.0.3358/mobilenetv2-7.dlc
+Model Version: N/A
+Model Copyright:N/A
+-----------------------------------------------------------------------------------------------------------------------------------------
+| Id  | Name                 | Type                  | Inputs  | Outputs | Out Dims     | Runtimes | Parameters                         |
+-----------------------------------------------------------------------------------------------------------------------------------------
+| 0   | input                | data                  | input   | input   | 1x224x224x3  | A D G C  | input_preprocessing: passthrough   |
+|     |                      |                       |         |         |              |          | input_type: image                  |
+| 1   | Conv_0               | convolutional         | input   | 474     | 1x112x112x32 | A D G C  | padding x: 1                       |
+|     |                      |                       |         |         |              |          | padding y: 1                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 2                        |
+|     |                      |                       |         |         |              |          | stride y: 2                        |
+|     |                      |                       |         |         |              |          | num filters: 32                    |
+|     |                      |                       |         |         |              |          | kernel: 3x3                        |
+|     |                      |                       |         |         |              |          | param count: 896 (0.0257%)         |
+|     |                      |                       |         |         |              |          | MACs per inference: 10M (3.6%)     |
+| 2   | Clip_1               | neuron                | 474     | 317     | 1x112x112x32 | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 3   | Conv_2               | convolutional         | 317     | 477     | 1x112x112x32 | A D G C  | padding x: 1                       |
+|     |                      |                       |         |         |              |          | padding y: 1                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 32                    |
+|     |                      |                       |         |         |              |          | kernel: 3x3                        |
+|     |                      |                       |         |         |              |          | groups: 32                         |
+|     |                      |                       |         |         |              |          | param count: 320 (0.00917%)        |
+|     |                      |                       |         |         |              |          | MACs per inference: 3M (1.2%)      |
+| 4   | Clip_3               | neuron                | 477     | 320     | 1x112x112x32 | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 5   | Conv_4               | convolutional         | 320     | 480     | 1x112x112x16 | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 16                    |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 528 (0.0151%)         |
+|     |                      |                       |         |         |              |          | MACs per inference: 6M (2.13%)     |
+| 6   | Conv_5               | convolutional         | 480     | 483     | 1x112x112x96 | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 96                    |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 1k (0.0468%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 19M (6.4%)     |
+| 7   | Clip_6               | neuron                | 483     | 325     | 1x112x112x96 | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 8   | Conv_7               | convolutional         | 325     | 486     | 1x56x56x96   | A D G C  | padding x: 1                       |
+|     |                      |                       |         |         |              |          | padding y: 1                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 2                        |
+|     |                      |                       |         |         |              |          | stride y: 2                        |
+|     |                      |                       |         |         |              |          | num filters: 96                    |
+|     |                      |                       |         |         |              |          | kernel: 3x3                        |
+|     |                      |                       |         |         |              |          | groups: 96                         |
+|     |                      |                       |         |         |              |          | param count: 960 (0.0275%)         |
+|     |                      |                       |         |         |              |          | MACs per inference: 2M (0.9%)      |
+| 9   | Clip_8               | neuron                | 486     | 328     | 1x56x56x96   | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 10  | Conv_9               | convolutional         | 328     | 489     | 1x56x56x24   | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 24                    |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 2k (0.0667%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 7M (2.4%)      |
+| 11  | Conv_10              | convolutional         | 489     | 492     | 1x56x56x144  | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 144                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 3k (0.103%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 10M (3.6%)     |
+| 12  | Clip_11              | neuron                | 492     | 333     | 1x56x56x144  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 13  | Conv_12              | convolutional         | 333     | 495     | 1x56x56x144  | A D G C  | padding x: 1                       |
+|     |                      |                       |         |         |              |          | padding y: 1                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 144                   |
+|     |                      |                       |         |         |              |          | kernel: 3x3                        |
+|     |                      |                       |         |         |              |          | groups: 144                        |
+|     |                      |                       |         |         |              |          | param count: 1k (0.0413%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 4M (1.35%)     |
+| 14  | Clip_13              | neuron                | 495     | 336     | 1x56x56x144  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 15  | Conv_14              | convolutional         | 336     | 498     | 1x56x56x24   | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 24                    |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 3k (0.0998%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 10M (3.6%)     |
+| 16  | Add_15               | elementwise_binary_op | 489     | 339     | 1x56x56x24   | A D G C  | operation: sum                     |
+|     |                      |                       | 498     |         |              |          | MACs per inference: 75k (0.025%)   |
+| 17  | Conv_16              | convolutional         | 339     | 501     | 1x56x56x144  | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 144                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 3k (0.103%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 10M (3.6%)     |
+| 18  | Clip_17              | neuron                | 501     | 342     | 1x56x56x144  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 19  | Conv_18              | convolutional         | 342     | 504     | 1x28x28x144  | A D G C  | padding x: 1                       |
+|     |                      |                       |         |         |              |          | padding y: 1                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 2                        |
+|     |                      |                       |         |         |              |          | stride y: 2                        |
+|     |                      |                       |         |         |              |          | num filters: 144                   |
+|     |                      |                       |         |         |              |          | kernel: 3x3                        |
+|     |                      |                       |         |         |              |          | groups: 144                        |
+|     |                      |                       |         |         |              |          | param count: 1k (0.0413%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 1M (0.338%)    |
+| 20  | Clip_19              | neuron                | 504     | 345     | 1x28x28x144  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 21  | Conv_20              | convolutional         | 345     | 507     | 1x28x28x32   | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 32                    |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 4k (0.133%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 3M (1.2%)      |
+| 22  | Conv_21              | convolutional         | 507     | 510     | 1x28x28x192  | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 192                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 6k (0.182%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 4M (1.6%)      |
+| 23  | Clip_22              | neuron                | 510     | 350     | 1x28x28x192  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 24  | Conv_23              | convolutional         | 350     | 513     | 1x28x28x192  | A D G C  | padding x: 1                       |
+|     |                      |                       |         |         |              |          | padding y: 1                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 192                   |
+|     |                      |                       |         |         |              |          | kernel: 3x3                        |
+|     |                      |                       |         |         |              |          | groups: 192                        |
+|     |                      |                       |         |         |              |          | param count: 1k (0.055%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 1M (0.45%)     |
+| 25  | Clip_24              | neuron                | 513     | 353     | 1x28x28x192  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 26  | Conv_25              | convolutional         | 353     | 516     | 1x28x28x32   | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 32                    |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 6k (0.177%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 4M (1.6%)      |
+| 27  | Add_26               | elementwise_binary_op | 507     | 356     | 1x28x28x32   | A D G C  | operation: sum                     |
+|     |                      |                       | 516     |         |              |          | MACs per inference: 25k (0.00833%) |
+| 28  | Conv_27              | convolutional         | 356     | 519     | 1x28x28x192  | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 192                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 6k (0.182%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 4M (1.6%)      |
+| 29  | Clip_28              | neuron                | 519     | 359     | 1x28x28x192  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 30  | Conv_29              | convolutional         | 359     | 522     | 1x28x28x192  | A D G C  | padding x: 1                       |
+|     |                      |                       |         |         |              |          | padding y: 1                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 192                   |
+|     |                      |                       |         |         |              |          | kernel: 3x3                        |
+|     |                      |                       |         |         |              |          | groups: 192                        |
+|     |                      |                       |         |         |              |          | param count: 1k (0.055%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 1M (0.45%)     |
+| 31  | Clip_30              | neuron                | 522     | 362     | 1x28x28x192  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 32  | Conv_31              | convolutional         | 362     | 525     | 1x28x28x32   | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 32                    |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 6k (0.177%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 4M (1.6%)      |
+| 33  | Add_32               | elementwise_binary_op | 356     | 365     | 1x28x28x32   | A D G C  | operation: sum                     |
+|     |                      |                       | 525     |         |              |          | MACs per inference: 25k (0.00833%) |
+| 34  | Conv_33              | convolutional         | 365     | 528     | 1x28x28x192  | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 192                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 6k (0.182%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 4M (1.6%)      |
+| 35  | Clip_34              | neuron                | 528     | 368     | 1x28x28x192  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 36  | Conv_35              | convolutional         | 368     | 531     | 1x14x14x192  | A D G C  | padding x: 1                       |
+|     |                      |                       |         |         |              |          | padding y: 1                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 2                        |
+|     |                      |                       |         |         |              |          | stride y: 2                        |
+|     |                      |                       |         |         |              |          | num filters: 192                   |
+|     |                      |                       |         |         |              |          | kernel: 3x3                        |
+|     |                      |                       |         |         |              |          | groups: 192                        |
+|     |                      |                       |         |         |              |          | param count: 1k (0.055%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 338k (0.113%)  |
+| 37  | Clip_36              | neuron                | 531     | 371     | 1x14x14x192  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 38  | Conv_37              | convolutional         | 371     | 534     | 1x14x14x64   | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 64                    |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 12k (0.354%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 2M (0.8%)      |
+| 39  | Conv_38              | convolutional         | 534     | 537     | 1x14x14x384  | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 384                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 24k (0.716%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 4M (1.6%)      |
+| 40  | Clip_39              | neuron                | 537     | 376     | 1x14x14x384  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 41  | Conv_40              | convolutional         | 376     | 540     | 1x14x14x384  | A D G C  | padding x: 1                       |
+|     |                      |                       |         |         |              |          | padding y: 1                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 384                   |
+|     |                      |                       |         |         |              |          | kernel: 3x3                        |
+|     |                      |                       |         |         |              |          | groups: 384                        |
+|     |                      |                       |         |         |              |          | param count: 3k (0.11%)            |
+|     |                      |                       |         |         |              |          | MACs per inference: 677k (0.225%)  |
+| 42  | Clip_41              | neuron                | 540     | 379     | 1x14x14x384  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 43  | Conv_42              | convolutional         | 379     | 543     | 1x14x14x64   | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 64                    |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 24k (0.706%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 4M (1.6%)      |
+| 44  | Add_43               | elementwise_binary_op | 534     | 382     | 1x14x14x64   | A D G C  | operation: sum                     |
+|     |                      |                       | 543     |         |              |          | MACs per inference: 12k (0.00417%) |
+| 45  | Conv_44              | convolutional         | 382     | 546     | 1x14x14x384  | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 384                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 24k (0.716%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 4M (1.6%)      |
+| 46  | Clip_45              | neuron                | 546     | 385     | 1x14x14x384  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 47  | Conv_46              | convolutional         | 385     | 549     | 1x14x14x384  | A D G C  | padding x: 1                       |
+|     |                      |                       |         |         |              |          | padding y: 1                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 384                   |
+|     |                      |                       |         |         |              |          | kernel: 3x3                        |
+|     |                      |                       |         |         |              |          | groups: 384                        |
+|     |                      |                       |         |         |              |          | param count: 3k (0.11%)            |
+|     |                      |                       |         |         |              |          | MACs per inference: 677k (0.225%)  |
+| 48  | Clip_47              | neuron                | 549     | 388     | 1x14x14x384  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 49  | Conv_48              | convolutional         | 388     | 552     | 1x14x14x64   | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 64                    |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 24k (0.706%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 4M (1.6%)      |
+| 50  | Add_49               | elementwise_binary_op | 382     | 391     | 1x14x14x64   | A D G C  | operation: sum                     |
+|     |                      |                       | 552     |         |              |          | MACs per inference: 12k (0.00417%) |
+| 51  | Conv_50              | convolutional         | 391     | 555     | 1x14x14x384  | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 384                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 24k (0.716%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 4M (1.6%)      |
+| 52  | Clip_51              | neuron                | 555     | 394     | 1x14x14x384  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 53  | Conv_52              | convolutional         | 394     | 558     | 1x14x14x384  | A D G C  | padding x: 1                       |
+|     |                      |                       |         |         |              |          | padding y: 1                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 384                   |
+|     |                      |                       |         |         |              |          | kernel: 3x3                        |
+|     |                      |                       |         |         |              |          | groups: 384                        |
+|     |                      |                       |         |         |              |          | param count: 3k (0.11%)            |
+|     |                      |                       |         |         |              |          | MACs per inference: 677k (0.225%)  |
+| 54  | Clip_53              | neuron                | 558     | 397     | 1x14x14x384  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 55  | Conv_54              | convolutional         | 397     | 561     | 1x14x14x64   | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 64                    |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 24k (0.706%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 4M (1.6%)      |
+| 56  | Add_55               | elementwise_binary_op | 391     | 400     | 1x14x14x64   | A D G C  | operation: sum                     |
+|     |                      |                       | 561     |         |              |          | MACs per inference: 12k (0.00417%) |
+| 57  | Conv_56              | convolutional         | 400     | 564     | 1x14x14x384  | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 384                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 24k (0.716%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 4M (1.6%)      |
+| 58  | Clip_57              | neuron                | 564     | 403     | 1x14x14x384  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 59  | Conv_58              | convolutional         | 403     | 567     | 1x14x14x384  | A D G C  | padding x: 1                       |
+|     |                      |                       |         |         |              |          | padding y: 1                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 384                   |
+|     |                      |                       |         |         |              |          | kernel: 3x3                        |
+|     |                      |                       |         |         |              |          | groups: 384                        |
+|     |                      |                       |         |         |              |          | param count: 3k (0.11%)            |
+|     |                      |                       |         |         |              |          | MACs per inference: 677k (0.225%)  |
+| 60  | Clip_59              | neuron                | 567     | 406     | 1x14x14x384  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 61  | Conv_60              | convolutional         | 406     | 570     | 1x14x14x96   | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 96                    |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 36k (1.06%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 7M (2.4%)      |
+| 62  | Conv_61              | convolutional         | 570     | 573     | 1x14x14x576  | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 576                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 55k (1.6%)            |
+|     |                      |                       |         |         |              |          | MACs per inference: 10M (3.6%)     |
+| 63  | Clip_62              | neuron                | 573     | 411     | 1x14x14x576  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 64  | Conv_63              | convolutional         | 411     | 576     | 1x14x14x576  | A D G C  | padding x: 1                       |
+|     |                      |                       |         |         |              |          | padding y: 1                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 576                   |
+|     |                      |                       |         |         |              |          | kernel: 3x3                        |
+|     |                      |                       |         |         |              |          | groups: 576                        |
+|     |                      |                       |         |         |              |          | param count: 5k (0.165%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 1M (0.338%)    |
+| 65  | Clip_64              | neuron                | 576     | 414     | 1x14x14x576  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 66  | Conv_65              | convolutional         | 414     | 579     | 1x14x14x96   | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 96                    |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 55k (1.59%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 10M (3.6%)     |
+| 67  | Add_66               | elementwise_binary_op | 570     | 417     | 1x14x14x96   | A D G C  | operation: sum                     |
+|     |                      |                       | 579     |         |              |          | MACs per inference: 18k (0.00625%) |
+| 68  | Conv_67              | convolutional         | 417     | 582     | 1x14x14x576  | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 576                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 55k (1.6%)            |
+|     |                      |                       |         |         |              |          | MACs per inference: 10M (3.6%)     |
+| 69  | Clip_68              | neuron                | 582     | 420     | 1x14x14x576  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 70  | Conv_69              | convolutional         | 420     | 585     | 1x14x14x576  | A D G C  | padding x: 1                       |
+|     |                      |                       |         |         |              |          | padding y: 1                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 576                   |
+|     |                      |                       |         |         |              |          | kernel: 3x3                        |
+|     |                      |                       |         |         |              |          | groups: 576                        |
+|     |                      |                       |         |         |              |          | param count: 5k (0.165%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 1M (0.338%)    |
+| 71  | Clip_70              | neuron                | 585     | 423     | 1x14x14x576  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 72  | Conv_71              | convolutional         | 423     | 588     | 1x14x14x96   | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 96                    |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 55k (1.59%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 10M (3.6%)     |
+| 73  | Add_72               | elementwise_binary_op | 417     | 426     | 1x14x14x96   | A D G C  | operation: sum                     |
+|     |                      |                       | 588     |         |              |          | MACs per inference: 18k (0.00625%) |
+| 74  | Conv_73              | convolutional         | 426     | 591     | 1x14x14x576  | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 576                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 55k (1.6%)            |
+|     |                      |                       |         |         |              |          | MACs per inference: 10M (3.6%)     |
+| 75  | Clip_74              | neuron                | 591     | 429     | 1x14x14x576  | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 76  | Conv_75              | convolutional         | 429     | 594     | 1x7x7x576    | A D G C  | padding x: 1                       |
+|     |                      |                       |         |         |              |          | padding y: 1                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 2                        |
+|     |                      |                       |         |         |              |          | stride y: 2                        |
+|     |                      |                       |         |         |              |          | num filters: 576                   |
+|     |                      |                       |         |         |              |          | kernel: 3x3                        |
+|     |                      |                       |         |         |              |          | groups: 576                        |
+|     |                      |                       |         |         |              |          | param count: 5k (0.165%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 254k (0.0844%) |
+| 77  | Clip_76              | neuron                | 594     | 432     | 1x7x7x576    | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 78  | Conv_77              | convolutional         | 432     | 597     | 1x7x7x160    | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 160                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 92k (2.65%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 4M (1.5%)      |
+| 79  | Conv_78              | convolutional         | 597     | 600     | 1x7x7x960    | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 960                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 154k (4.43%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 7M (2.5%)      |
+| 80  | Clip_79              | neuron                | 600     | 437     | 1x7x7x960    | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 81  | Conv_80              | convolutional         | 437     | 603     | 1x7x7x960    | A D G C  | padding x: 1                       |
+|     |                      |                       |         |         |              |          | padding y: 1                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 960                   |
+|     |                      |                       |         |         |              |          | kernel: 3x3                        |
+|     |                      |                       |         |         |              |          | groups: 960                        |
+|     |                      |                       |         |         |              |          | param count: 9k (0.275%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 423k (0.141%)  |
+| 82  | Clip_81              | neuron                | 603     | 440     | 1x7x7x960    | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 83  | Conv_82              | convolutional         | 440     | 606     | 1x7x7x160    | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 160                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 153k (4.41%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 7M (2.5%)      |
+| 84  | Add_83               | elementwise_binary_op | 597     | 443     | 1x7x7x160    | A D G C  | operation: sum                     |
+|     |                      |                       | 606     |         |              |          | MACs per inference: 7k (0.0026%)   |
+| 85  | Conv_84              | convolutional         | 443     | 609     | 1x7x7x960    | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 960                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 154k (4.43%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 7M (2.5%)      |
+| 86  | Clip_85              | neuron                | 609     | 446     | 1x7x7x960    | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 87  | Conv_86              | convolutional         | 446     | 612     | 1x7x7x960    | A D G C  | padding x: 1                       |
+|     |                      |                       |         |         |              |          | padding y: 1                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 960                   |
+|     |                      |                       |         |         |              |          | kernel: 3x3                        |
+|     |                      |                       |         |         |              |          | groups: 960                        |
+|     |                      |                       |         |         |              |          | param count: 9k (0.275%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 423k (0.141%)  |
+| 88  | Clip_87              | neuron                | 612     | 449     | 1x7x7x960    | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 89  | Conv_88              | convolutional         | 449     | 615     | 1x7x7x160    | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 160                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 153k (4.41%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 7M (2.5%)      |
+| 90  | Add_89               | elementwise_binary_op | 443     | 452     | 1x7x7x160    | A D G C  | operation: sum                     |
+|     |                      |                       | 615     |         |              |          | MACs per inference: 7k (0.0026%)   |
+| 91  | Conv_90              | convolutional         | 452     | 618     | 1x7x7x960    | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 960                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 154k (4.43%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 7M (2.5%)      |
+| 92  | Clip_91              | neuron                | 618     | 455     | 1x7x7x960    | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 93  | Conv_92              | convolutional         | 455     | 621     | 1x7x7x960    | A D G C  | padding x: 1                       |
+|     |                      |                       |         |         |              |          | padding y: 1                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 960                   |
+|     |                      |                       |         |         |              |          | kernel: 3x3                        |
+|     |                      |                       |         |         |              |          | groups: 960                        |
+|     |                      |                       |         |         |              |          | param count: 9k (0.275%)           |
+|     |                      |                       |         |         |              |          | MACs per inference: 423k (0.141%)  |
+| 94  | Clip_93              | neuron                | 621     | 458     | 1x7x7x960    | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 95  | Conv_94              | convolutional         | 458     | 624     | 1x7x7x320    | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 320                   |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 307k (8.82%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 15M (5%)       |
+| 96  | Conv_95              | convolutional         | 624     | 627     | 1x7x7x1280   | A D G C  | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | padding mode: zero                 |
+|     |                      |                       |         |         |              |          | stride x: 1                        |
+|     |                      |                       |         |         |              |          | stride y: 1                        |
+|     |                      |                       |         |         |              |          | num filters: 1280                  |
+|     |                      |                       |         |         |              |          | kernel: 1x1                        |
+|     |                      |                       |         |         |              |          | param count: 410k (11.8%)          |
+|     |                      |                       |         |         |              |          | MACs per inference: 20M (6.67%)    |
+| 97  | Clip_96              | neuron                | 627     | 463     | 1x7x7x1280   | A D G C  | a: 0                               |
+|     |                      |                       |         |         |              |          | b: 0                               |
+|     |                      |                       |         |         |              |          | min_clamp: 0                       |
+|     |                      |                       |         |         |              |          | max_clamp: 6                       |
+|     |                      |                       |         |         |              |          | func: relu_min_max                 |
+| 98  | GlobalAveragePool_97 | pooling               | 463     | 464     | 1x1x1x1280   | A D G C  | pool size x: 7                     |
+|     |                      |                       |         |         |              |          | pool size y: 7                     |
+|     |                      |                       |         |         |              |          | stride x: 7                        |
+|     |                      |                       |         |         |              |          | stride y: 7                        |
+|     |                      |                       |         |         |              |          | padding x: 0                       |
+|     |                      |                       |         |         |              |          | padding y: 0                       |
+|     |                      |                       |         |         |              |          | pool_type: POOL_AVG                |
+|     |                      |                       |         |         |              |          | MACs per inference: 62k (0.0208%)  |
+| 99  | 464.ncs              | permute               | 464     | 464.ncs | 1x1280x1x1   | A D G C  | permute_order: [0, 3, 1, 2]        |
+| 100 | Gemm_104             | fully_connected       | 464.ncs | output  | 1x1000       | A D G C  | param count: 1M (36.7%)            |
+|     |                      |                       |         |         |              |          | MACs per inference: 1M (0.425%)    |
+-----------------------------------------------------------------------------------------------------------------------------------------
+Note: The supported runtimes column assumes a processor target of Snapdragon 835 (8998)
+Key : A:AIP
+      D:DSP
+      G:GPU
+      C:CPU
+
+Total parameters: 3487816 (13 MB assuming single precision float)
+Total MACs per inference: 301M (100%)
+Converter command: snpe-onnx-to-dlc adjust_nms_features_dims=False align_matmul_ranks=True copyright_file=None custom_op_config_paths=None debug=-1 disable_batchnorm_folding=False disable_chaining_eltwise_ops=False dry_run=None dumpIR=False dump_inferred_model=False dump_value_info=False enable_strict_validation=False extract_color_transform=False force_prune_cast_ops=True handle_gather_negative_indices=False inject_cast_for_gather=False input_dim=[['input', '1,3,224,224']] input_dtype=[] input_encoding=[] input_layout=[] input_type=[['input', 'image']] keep_disconnected_nodes=False keep_quant_nodes=False match_caffe_ssd_to_tf=False model_version=None no_simplification=False out_names=['output'] perform_axes_to_spatial_first_order=True prepare_inputs_as_params=True preprocess_lstm_ops=False preprocess_roi_pool_inputs=False quantization_overrides= squash_box_decoder=False unroll_lstm_time_steps=False use_convert_quantization_nodes=True validation_target=[]
+Quantizer command: N/A
+DLC created with converter version: 1.61.0.3358
+Layers used by DLC: CONVOLUTIONAL, DATA, ELEMENTWISE_BINARY_OP_SUM, FULLY_CONNECTED, NEURON_RELU_MIN_MAX, PERMUTE, POOLING
+Est. Steady-State Memory Needed to Run: 164.3 MiB
+-----------------------------------------------------------------------------------------------------------------------------------------
+

tinygrad_repo/extra/dsp/snpe_logs/high_perf_2

@@ -0,0 +1,131 @@
+Log File Created: Tue Mar 18 01:33:12 2025
+Time Scale: 1e-06
+Epoch Timestamp: 1742286792883569 Steady Clock Timestamp: 75586845756
+Software library version: 1.61.0.3358
+
+Dnn Runtime Load/Deserialize/Create/De-Init Statistics:
+--------------------------------------------------
+Load: 333 us
+Deserialize: 32452 us
+Create: 143084 us
+
+Init: 178071 us
+De-Init: 16710 us
+
+Create Network(s): 86850 us
+RPC Init Time: 43213 us
+Snpe Accelerator Init Time: 42154 us
+Accelerator Init Time: 39189 us
+
+Average SNPE Statistics:
+------------------------------
+Total Inference Time: 11868 us
+Forward Propagate Time: 11816 us
+RPC Execute Time: 9810 us
+Snpe Accelerator Time: 9129 us
+Accelerator Time: 8701 us
+Misc Accelerator Time: 10 us
+
+Layer Times: 
+---------------
+0: 42 us : DSP
+1: 0 us : DSP
+2: 254 us : DSP
+3: 0 us : DSP
+4: 153 us : DSP
+5: 295 us : DSP
+6: 0 us : DSP
+7: 287 us : DSP
+8: 0 us : DSP
+9: 162 us : DSP
+10: 210 us : DSP
+11: 0 us : DSP
+12: 138 us : DSP
+13: 0 us : DSP
+14: 176 us : DSP
+15: 293 us : DSP
+16: 60 us : DSP
+17: 0 us : DSP
+18: 157 us : DSP
+19: 0 us : DSP
+20: 112 us : DSP
+21: 134 us : DSP
+22: 0 us : DSP
+23: 81 us : DSP
+24: 0 us : DSP
+25: 104 us : DSP
+26: 130 us : DSP
+27: 37 us : DSP
+28: 0 us : DSP
+29: 81 us : DSP
+30: 0 us : DSP
+31: 87 us : DSP
+32: 124 us : DSP
+33: 30 us : DSP
+34: 0 us : DSP
+35: 87 us : DSP
+36: 0 us : DSP
+37: 63 us : DSP
+38: 74 us : DSP
+39: 0 us : DSP
+40: 102 us : DSP
+41: 0 us : DSP
+42: 82 us : DSP
+43: 95 us : DSP
+44: 29 us : DSP
+45: 0 us : DSP
+46: 112 us : DSP
+47: 0 us : DSP
+48: 88 us : DSP
+49: 96 us : DSP
+50: 25 us : DSP
+51: 0 us : DSP
+52: 103 us : DSP
+53: 0 us : DSP
+54: 80 us : DSP
+55: 100 us : DSP
+56: 26 us : DSP
+57: 0 us : DSP
+58: 102 us : DSP
+59: 0 us : DSP
+60: 85 us : DSP
+61: 129 us : DSP
+62: 0 us : DSP
+63: 155 us : DSP
+64: 0 us : DSP
+65: 113 us : DSP
+66: 194 us : DSP
+67: 34 us : DSP
+68: 0 us : DSP
+69: 157 us : DSP
+70: 0 us : DSP
+71: 120 us : DSP
+72: 198 us : DSP
+73: 34 us : DSP
+74: 0 us : DSP
+75: 155 us : DSP
+76: 0 us : DSP
+77: 101 us : DSP
+78: 121 us : DSP
+79: 0 us : DSP
+80: 256 us : DSP
+81: 0 us : DSP
+82: 134 us : DSP
+83: 159 us : DSP
+84: 31 us : DSP
+85: 0 us : DSP
+86: 199 us : DSP
+87: 0 us : DSP
+88: 142 us : DSP
+89: 152 us : DSP
+90: 26 us : DSP
+91: 0 us : DSP
+92: 202 us : DSP
+93: 0 us : DSP
+94: 143 us : DSP
+95: 278 us : DSP
+96: 0 us : DSP
+97: 316 us : DSP
+98: 40 us : DSP
+99: 12 us : DSP
+100: 199 us : DSP

tinygrad_repo/extra/dsp/snpe_logs/parse.py

@@ -0,0 +1,21 @@
+di = open("dlc_info_2").read().split("\n")
+layers = {}
+for l in di:
+  if not l.startswith("| "): continue
+  if l.startswith("|     |"): continue
+  ll = [x.strip() for x in l.split("|")]
+  if ll[1] == "Id": continue
+  layers[int(ll[1])] = (ll[2], ll[6])
+hp = open("high_perf_2").read().split("Layer Times:")[1].strip().split("\n")[2:]
+sl = 1
+tms = 0
+for l in hp:
+  kk, tm, _ = l.split(" ", 2)
+  tm = int(tm)
+  lnum = int(kk.strip(":"))
+  if int(tm) != 0:
+    print(f"{sl:2d} {tm:4d} us {layers[lnum]}")
+    tms += tm
+    sl += 1
+print(f"total time, {tms/1000:.2f} ms")
+

tinygrad_repo/extra/export_model.py

@@ -0,0 +1,302 @@
+from typing import Tuple, Dict, List, Optional
+from tinygrad.dtype import DType, dtypes
+from tinygrad.tensor import Tensor
+from tinygrad.device import Device, Buffer
+from tinygrad.engine.jit import TinyJit
+from tinygrad.nn.state import get_state_dict
+from tinygrad.helpers import Context, to_mv, prod
+from tinygrad.uop.ops import Ops, UOp
+from tinygrad.codegen import to_program
+import json
+from collections import OrderedDict
+
+EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "CL"]
+
+_KERNEL_ASTS = {Ops.SINK, Ops.PROGRAM}
+def iter_kernel_calls(linear:UOp):
+  """Yield kernel CALLs from a LINEAR UOp. Toposort descends naturally into CUSTOM_FUNCTION graph batches; gate stops at kernel ASTs."""
+  return (u for u in linear.toposort(gate=lambda x: x.op not in _KERNEL_ASTS) if u.op is Ops.CALL and u.src[0].op in _KERNEL_ASTS)
+
+def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], List, Dict[str,Tuple[int,DType,int]], Dict[str,Buffer]]:
+  output_name = {id(b): f"output{i}" for i, b in enumerate(output_bufs)}
+  functions, bufs, bufs_to_save, statements, n = {}, {}, {}, [], 0
+
+  def name_of(bu:UOp, is_out:bool) -> str:
+    nonlocal n
+    if bu.op is Ops.PARAM: key, name, size = ("in", bu.arg.slot), f"input{bu.arg.slot}", prod(bu.shape)*bu.dtype.itemsize
+    else:
+      b = bu.buffer
+      key, size = (id(b.base), b.offset, b.size, b.dtype), b.size*b.dtype.itemsize
+      if key in bufs: return bufs[key][0]
+      if (name:=output_name.get(id(b))) is None:
+        name, n = f"buf_{n}", n+1
+        if not is_out: bufs_to_save[name] = b
+    bufs[key] = (name, size, bu.dtype, key)
+    return name
+
+  for call in iter_kernel_calls(linear):
+    arg_uops = [b for b in call.src[1:] if b.op is not Ops.BIND]
+    prg = to_program(call.src[0], Device[arg_uops[0].device].renderer)
+    info = prg.arg
+    functions[info.function_name] = prg.src[3].arg
+    cargs = [name_of(bu, i == 0) for i, bu in enumerate(arg_uops)] + [v for v in info.vars if v.op is Ops.DEFINE_VAR]
+    statements.append((info.function_name, cargs, info.global_size, info.local_size))
+
+  return functions, statements, {name:(size, dtype, key) for name, size, dtype, key in bufs.values()}, bufs_to_save
+
+def jit_model(model, *args) -> Tuple[UOp, List[Buffer]]:
+  assert hasattr(model, "forward") or callable(model), "model needs a forward function"
+  @TinyJit
+  def run(*x):
+    out = model.forward(*x) if hasattr(model, "forward") else model(*x)
+    assert isinstance(out, (tuple, list, Tensor)), "model output must be a Tensor, tuple, or a list of Tensors for export"
+    out = [out] if isinstance(out, Tensor) else out
+    return [o.realize() for o in out]
+
+  # run twice to trigger JIT capture
+  for _ in range(2): the_output = run(*args)
+  assert run.captured is not None
+  return run.captured.linear, [o.uop.base.realized for o in the_output]
+
+def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,int,int]], bufs:Dict[str,Tuple[str,int,int]],
+  bufs_to_save:Dict[str,Tensor], input_names:List[str], output_names:List[str], weight_names={}, model_name="model", symbolic_vars={}, wasm=False) -> str:
+  headers = ["#include <tgmath.h>"]
+  cprog = list(functions.values())
+  dtype_map = {dtypes.int: "int", dtypes.float: "float", dtypes.uchar: "unsigned char", dtypes.char: "signed char", dtypes.half: "__fp16", dtypes.uint: "unsigned int"}
+  inputs = [(name, dtype_map[bufs[name][1]], bufs[name][0]) for name in input_names + list(symbolic_vars.values())]
+  outputs = [(name, dtype_map[bufs[name][1]], bufs[name][0]) for name in output_names]
+  forward_args = ",".join(f"{dtype}{'*' if name not in symbolic_vars.values() else ''} {name}" for name,dtype,_ in (outputs+inputs if wasm else inputs+outputs))
+
+  if not wasm:
+    for name,cl in bufs_to_save.items():
+      weight = ''.join(["\\x%02X"%x for x in bytes(to_mv(cl._buf.va_addr, cl._buf.size))])
+      cprog.append(f"unsigned char {name}_data[] = \"{weight}\";")
+    cprog += [f"{dtype_map[dtype]} {name}[{len}];" if name not in bufs_to_save else f"{dtype_map[dtype]} *{name} = ({dtype_map[dtype]} *){name}_data;" for name,(len,dtype,_key) in bufs.items() if name not in input_names+output_names]
+    cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)});" for (name, args, _global_size, _local_size) in statements] + ["}"]
+    return '\n'.join(headers + cprog)
+  else:
+    if bufs_to_save:
+      headers += ["#include <stddef.h>"]
+      bufs_to_save = {k:v for k,v in bufs.items() if v[2] in weight_names} # causes random seeds to be set as zeroes, not exported as a model weight
+      buf_to_name = OrderedDict((buf_name, {"name": weight_names[data[2]], "idx": i}) for i, (buf_name, data) in enumerate(bufs_to_save.items()))
+      cprog.append(f"void* bufs[{len(buf_to_name)}];")
+      cprog.append(f"""void set_buf(size_t index, void* ptr) {{\n  bufs[index] = ptr;\n}}""")
+
+    for name in set(bufs.keys()) - set(bufs_to_save.keys()) - set(input_names + output_names):
+      n_bytes, dtype, _ = bufs[name]
+      cprog += [f"{dtype_map[dtype]} {name}[{n_bytes // dtype.itemsize}];"]
+
+    cprog += [f"void net({forward_args})"] + ["{"]
+    get_weight_ptr = lambda x: f"({dtype_map[bufs_to_save[x][1]]} *)bufs[{buf_to_name[x]['idx']}]" if x in bufs_to_save else x
+    cprog += [f"  {name}({', '.join(map(get_weight_ptr, args))});" for (name, args, _global_size, _local_size) in statements] + ["}"]
+    weightMapping = "" if not bufs_to_save else f"""\nconst weightNames = [{", ".join([f'"{weight_name}"' for weight_name in [v["name"] for v in buf_to_name.values()]])}];
+const {model_name}_name_to_id = Object.fromEntries(weightNames.map((name, index) => [name, index]));\n"""
+    top = f"""import {model_name}Module from './{model_name}.js'{weightMapping}"""
+    whitespace = "\n  "
+    js_wrapper = f"""{top}\nvar {model_name} = async function() {{
+  const wasm = await {model_name}Module();
+  {whitespace.join(f"const {name}Ptr = wasm._malloc({n_bytes});" for name, _, n_bytes in outputs+inputs if name not in symbolic_vars.values())}
+  return {{
+    run: ({",".join(name for name,_,_ in inputs)}) => {{
+      {(whitespace + "    ").join(f"wasm.HEAPU8.set({name}, {name}Ptr);" for name,_,_ in inputs if name not in symbolic_vars.values())}
+      wasm._net({", ".join(f"{name}{'Ptr' if name not in symbolic_vars.values() else ''}" for name,_,_ in outputs+inputs)});
+      {(whitespace + "    ").join(f"const {name} = wasm.HEAPU8.slice({name}Ptr, {name}Ptr + {n_bytes});" for name,_,n_bytes in outputs)}
+      return [{", ".join(f"{name}" for name,_,_ in outputs)}];
+    }},
+    wasm: wasm
+  }}
+}}\nexport {{ {model_name}, {model_name}_name_to_id }};"""
+
+    return '\n'.join(headers + cprog), js_wrapper
+
+def dtype_to_js_type(dtype: DType) -> str:
+  return f"{'Uint' if dtype in dtypes.uints else 'Int' if (dtype in dtypes.sints or dtype == dtypes.bool) else 'Float'}{8*dtype.itemsize}Array"
+
+def export_model_webgpu(functions, statements, bufs, weight_names, input_names, output_names, model_name, symbolic_vars={}, stream_weights=False) -> Tuple[str,int,int]:
+  kernel_code = '\n\n'.join([f"const {key} = `{code.replace(key, 'main')}`;" for key, code in functions.items()])
+  kernel_names = ', '.join([name for (name, _, _, _) in statements])
+  input_names += list(symbolic_vars.values())
+  input_buffer_types = [dtype_to_js_type(bufs[inp_name][1]) for inp_name in input_names]
+  output_buffer_types = [dtype_to_js_type(bufs[out_name][1]) for out_name in output_names]
+
+  buf_type = lambda x: "uniform" if x in set(symbolic_vars.values()) else "storage"
+  create_bind_group_layouts = ",".join([
+    "device.createBindGroupLayout({{entries: [{{binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: {{ type: 'uniform' }}}}, {}]}})".format(
+        ",".join([f"{{binding: {argIdx+1}, visibility: GPUShaderStage.COMPUTE, buffer: {{ type: '{buf_type(argName)}' }} }}" for argIdx, argName in enumerate(args)])
+    )
+    for _, (_, args, _, _) in enumerate(statements)
+  ])
+  layouts = f"const layouts=[{create_bind_group_layouts}]"
+  kernel_calls = '\n        '.join([f"addComputePass(device, commandEncoder, pipelines[{i}], layouts[{i}], infinityBuf, [{', '.join(args)}], [{', '.join(str(x) for x in global_size)}]);" for i, (_name, args, global_size, _local_size) in enumerate(statements) ])
+
+  buf_type = lambda x: "createUniformBuf" if x in set(uop.arg[0] for uop in symbolic_vars) else "createEmptyBuf"
+  map_to_external_weight = lambda _key: f"state_dict['{weight_names[_key]}']" if stream_weights else f"getTensorBuffer(safetensor, metadata['{weight_names[_key]}'])"
+  _bufs =  '\n    '.join([f"const {name} = " + (f"{buf_type(_key)}(device, {size});" if _key not in weight_names else f"createWeightBuf(device, {size}, {map_to_external_weight(_key)})") + ";"  for name,(size,dtype,_key) in bufs.items()])
+  gpu_write_bufs =  '\n    '.join([f"const gpuWriteBuffer{i} = device.createBuffer({{size:{input_name}.size, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE }});" for i,input_name in enumerate(input_names)])
+  input_writers = '\n    '.join([f"await gpuWriteBuffer{i}.mapAsync(GPUMapMode.WRITE);\n        new {input_buffer_types[i]}(gpuWriteBuffer{i}.getMappedRange()).set(" + f'_{inp_name});' + f"\n        gpuWriteBuffer{i}.unmap();\n        commandEncoder.copyBufferToBuffer(gpuWriteBuffer{i}, 0, {inp_name}, 0, gpuWriteBuffer{i}.size);"  for i,inp_name in enumerate(input_names)])
+  gpu_read_bufs = '\n    '.join([f"const gpuReadBuffer{i} = device.createBuffer({{size:{output_name}.size, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ }});" for i,output_name in enumerate(output_names)])
+  outbuf_copies = '\n        '.join([f"commandEncoder.copyBufferToBuffer({output_name}, 0, gpuReadBuffer{i}, 0, output{i}.size);" for i,output_name in enumerate(output_names)])
+  output_readers = '\n        '.join([f"await gpuReadBuffer{i}.mapAsync(GPUMapMode.READ);\n        const resultBuffer{i} = new {output_buffer_types[i]}(gpuReadBuffer{i}.size/{bufs[output_names[i]][1].itemsize});\n        resultBuffer{i}.set(new {output_buffer_types[i]}(gpuReadBuffer{i}.getMappedRange()));\n        gpuReadBuffer{i}.unmap();" for i in range(len(output_names))])
+  output_return = '[{}]'.format(",".join([f'resultBuffer{i}' for i in range(len(output_names))]))
+  getTensorMetadata = f"""\nconst getTensorMetadata = (safetensorBuffer) => {{
+    const metadataLength = Number(new DataView(safetensorBuffer.buffer).getBigUint64(0, true));
+    const metadata = JSON.parse(new TextDecoder("utf8").decode(safetensorBuffer.subarray(8, 8 + metadataLength)));
+    return Object.fromEntries(Object.entries(metadata).filter(([k, v]) => k !== "__metadata__").map(([k, v]) => [k, {{...v, data_offsets: v.data_offsets.map(x => 8 + metadataLength + x)}}]));
+}};\n""" if not stream_weights else ""
+  return f"""
+const {model_name} = (() => {{
+const getTensorBuffer = (safetensorBuffer, tensorMetadata) => {{
+  return safetensorBuffer.subarray(...tensorMetadata.data_offsets);
+}};
+{getTensorMetadata}
+const createEmptyBuf = (device, size) => {{
+    return device.createBuffer({{size, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST }});
+}};
+
+const createUniformBuf = (device, size) => {{
+  return device.createBuffer({{size, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST}})
+}}
+
+const createInfinityUniformBuf = (device) => {{
+  const size = 4;
+  const buf = device.createBuffer({{
+    mappedAtCreation: true,
+    size,
+    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST
+  }});
+  new Float32Array(buf.getMappedRange())[0] = Infinity;
+  buf.unmap();
+  return buf;
+}};
+
+const createWeightBuf = (device, size, data) => {{
+  const buf = device.createBuffer({{ size, usage: GPUBufferUsage.STORAGE{" | GPUBufferUsage.COPY_DST" if stream_weights else ", mappedAtCreation: true"} }});
+  {"data.bytes = buf;" if stream_weights else "new Uint8Array(buf.getMappedRange()).set(data); buf.unmap();"}
+  return buf;
+}};
+
+const addComputePass = (device, commandEncoder, pipeline, layout, infinityUniformBuf, bufs, workgroup) => {{
+  const bindGroup = device.createBindGroup({{
+    layout: layout,
+    entries: [
+      {{ binding: 0, resource: {{ buffer: infinityUniformBuf }} }},
+      ...bufs.map((buffer, index) => ({{ binding: index + 1, resource: {{ buffer }} }}))
+    ]
+  }});
+
+  const passEncoder = commandEncoder.beginComputePass();
+  passEncoder.setPipeline(pipeline);
+  passEncoder.setBindGroup(0, bindGroup);
+  passEncoder.dispatchWorkgroups(...workgroup);
+  passEncoder.end();
+}};
+
+{kernel_code}
+
+const setupNet = async (device, {"state_dict" if stream_weights else "safetensor"}) => {{
+    {"const metadata = getTensorMetadata(safetensor);" if not stream_weights else ""}
+    const infinityBuf = createInfinityUniformBuf(device);
+
+    {layouts}
+
+    {_bufs}
+
+    {gpu_write_bufs}
+
+    {gpu_read_bufs}
+
+    const kernels = [{kernel_names}];
+    const pipelines = await Promise.all(kernels.map(async (name, i) => {{
+      return await device.createComputePipelineAsync({{
+          layout: device.createPipelineLayout({{
+              bindGroupLayouts: [layouts[i]],
+          }}),
+          compute: {{
+              module: device.createShaderModule({{
+                  code: name,
+              }}),
+              entryPoint: "main",
+          }},
+      }});
+  }}))
+
+    return async ({",".join([f"_{input_name}" for input_name in input_names])}) => {{
+        const commandEncoder = device.createCommandEncoder();
+        {input_writers}
+        {kernel_calls}
+        {outbuf_copies}
+        const gpuCommands = commandEncoder.finish();
+        device.queue.submit([gpuCommands]);
+
+        {output_readers}
+        return {output_return};
+    }}
+}}
+const load = async (device, weight_path) => {{ return await fetch(weight_path).then(x => x.arrayBuffer()).then(x => setupNet(device, new Uint8Array(x))); }}
+return {{ load, setupNet }};
+}})();
+export default {model_name};
+"""
+
+def export_model(model, target:str, *inputs, model_name: Optional[str] = "model", stream_weights=False):
+  assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"only {', '.join(EXPORT_SUPPORTED_DEVICE)} are supported"
+
+  # NOTE: CPU_COUNT=1, since export does not support threading
+  with Context(JIT=2, CPU_COUNT=1): linear, output_bufs = jit_model(model, *inputs)
+  functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
+  state = get_state_dict(model)
+  weight_names = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}
+  input_names = [f"input{i}" for i in range(len(inputs))]
+  output_names = [f"output{i}" for i in range(len(output_bufs))]
+
+  # handle symbolic variables; TODO: refactor to fix some of this stuff upstream in tinygrad
+  symbolic_vars = OrderedDict()
+  for i, (_, args, global_size, _) in enumerate(statements):
+    for j, var in enumerate(args):
+      if getattr(var, "op", None) is Ops.DEFINE_VAR and isinstance(getattr(var, "arg", None), tuple) and isinstance(var.arg[0], str):
+        if var not in symbolic_vars:
+          symbolic_vars[var] = var.arg[0]
+          bufs[symbolic_vars[var]] = (var.dtype.itemsize, var.dtype, symbolic_vars[var])
+        statements[i][1][j] = symbolic_vars[var]
+
+    if global_size:
+      for j, dim in enumerate(global_size):
+        if getattr(dim, "op", None) is Ops.ADD and len(dim.src) == 2 and {dim.src[0].op, dim.src[1].op} == {Ops.DEFINE_VAR, Ops.CONST}:
+          name, val = dim.src if dim.src[1].op is Ops.CONST else reversed(dim.src)
+          global_size[j] = f"_{name.arg[0]}[0] + {val.arg}"
+
+  prg = ""
+  if target == "clang":
+    prg = export_model_clang(functions, statements, bufs, bufs_to_save, input_names, output_names)
+  elif target == "wasm":
+    return export_model_clang(functions, statements, bufs, bufs_to_save, input_names, output_names, weight_names, model_name, symbolic_vars, wasm=True)
+  elif target == "webgpu":
+    prg = export_model_webgpu(functions, statements, bufs, weight_names, input_names, output_names, model_name, symbolic_vars, stream_weights)
+  else:
+    prg = json.dumps({
+      "backend": Device.DEFAULT,
+      "inputs": [{
+        "size": bufs[name][0],
+        "dtype": bufs[name][1].name
+      } for name in input_names],
+      "outputs": [{
+        "size": bufs[name][0],
+        "dtype": bufs[name][1].name
+      } for name in output_names],
+      "functions": functions,
+      "statements": [{
+        "kernel": kernel,
+        "args": args,
+        "global_size": global_size,
+        "local_size": local_size
+      } for (kernel, args, global_size, local_size) in statements],
+      "buffers": {
+        name: {
+          "size": size,
+          "dtype": dtype.name,
+          "id": weight_names[_key] if _key in weight_names else ""
+        } for name, (size,dtype,_key) in bufs.items() if name not in ["input", "outputs"]
+      }
+    })
+
+  return prg, {input:bufs[input][0] for input in input_names}, {output:bufs[output][0] for output in output_names}, state

tinygrad_repo/extra/f16_decompress.py

@@ -0,0 +1,16 @@
+from tinygrad import Tensor
+
+def bit_extract(x: Tensor, e: int, s: int) -> Tensor:
+  mask = (1 << (e - s + 1)) - 1
+  return (x >> s) & mask
+
+def u16_to_f16(x: Tensor) -> Tensor:
+  sign = bit_extract(x, 15, 15).float()
+  exponent = bit_extract(x, 14, 10).float()
+  fraction = bit_extract(x, 9, 0).float()
+  return sign.where(-1, 1) * exponent.where((exponent - 15.0).exp2() * (1 + fraction / 1024.0), 6.103515625e-5 * (fraction / 1024.0))
+
+def u32_to_f16(oo: Tensor) -> Tensor:
+  f1 = u16_to_f16(oo>>16)
+  f2 = u16_to_f16(oo&0xFFFF)
+  return Tensor.cat(f2.reshape(-1, 1), f1.reshape(-1, 1), dim=1).flatten()

tinygrad_repo/extra/fp8/fp8_linear.py

@@ -0,0 +1,101 @@
+from typing import Callable, Any
+from tinygrad import Tensor, dtypes, nn, UOp
+from tinygrad.uop.ops import KernelInfo, AxisType, Ops
+
+def quantize_to_fp8(x: Tensor, dtype=dtypes.fp8e4m3):
+  fp8_min = -448.0 if dtype == dtypes.fp8e4m3 else -57344.0
+  fp8_max = 448.0 if dtype == dtypes.fp8e4m3 else 57344.0
+  x_abs_max = x.abs().max().detach()
+  scale = fp8_max / (x_abs_max + 1e-8)
+  x_scaled = x * scale
+  x_det = x_scaled.detach()
+  x_clamped = x_det.clamp(fp8_min, fp8_max)
+  x_clamped_ste = x_scaled + (x_clamped - x_det)
+  res = x_clamped_ste.cast(dtype)
+  return res, scale.float().reciprocal()
+
+def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
+  SEQ = inp.shape[1]
+  OUT = weight.shape[0]
+  IN = weight.shape[-1]
+  seq_idx = UOp.range(SEQ, 2, AxisType.LOOP)
+  out_idx = UOp.range(OUT, 3, AxisType.LOOP)
+  batch_idx = UOp.range(output.size//SEQ//OUT, 1, AxisType.LOOP)
+  reduce_idx = UOp.range(IN, 0, AxisType.REDUCE)
+  product = (inp.index((seq_idx*IN+reduce_idx+batch_idx*IN*SEQ)) * weight.index((out_idx*IN+reduce_idx))).cast(dtypes.float)
+  reduced = product.reduce(reduce_idx, arg=Ops.ADD)
+  store_op = output.index((seq_idx*OUT+out_idx+batch_idx*OUT*SEQ), ptr=True).store(reduced).end(batch_idx, seq_idx, out_idx)
+  return store_op.sink(arg=KernelInfo(name=f"fp8_matmul_{inp.shape}x{weight.shape}"))
+
+def custom_matmul_backward(gradient: UOp, kernel: UOp) -> tuple[UOp, UOp]:
+  _, input_uop, weight_uop = kernel.src[1:]
+  input_tensor = Tensor(input_uop, device=input_uop.device)
+  grad_tensor = Tensor(gradient, device=gradient.device)
+  weight_tensor = Tensor(weight_uop, device=weight_uop.device)
+  grad_quantized, scale = quantize_to_fp8(grad_tensor)
+  scale_scalar = scale.reshape(())
+  grad_weight = Tensor.einsum("bso,bsi->oi", grad_quantized, input_tensor, dtype=dtypes.float)
+  grad_weight = grad_weight * scale_scalar
+  grad_2d = grad_quantized.reshape(grad_tensor.shape[0] * grad_tensor.shape[1], grad_tensor.shape[-1])
+  grad_input = (grad_2d.dot(weight_tensor, dtype=dtypes.float)).contiguous().reshape(input_tensor.shape) * scale
+  return (None, grad_input.uop, grad_weight.uop)
+
+class FP8Linear:
+  def __init__(self, in_features:int, out_features:int, bias:bool=True):
+    self.weight = Tensor.empty(out_features, in_features, dtype=dtypes.float32)
+    self.bias = Tensor.empty(out_features, dtype=dtypes.float32) if bias else None
+
+  def __call__(self, x: Tensor) -> Tensor:
+    original_ndim = len(x.shape)
+    if original_ndim == 2: x = x.reshape(x.shape[0], 1, x.shape[1])
+    batch, seq, _ = x.shape
+    w_fp8, w_scale = quantize_to_fp8(self.weight)
+    x_fp8, x_scale = quantize_to_fp8(x)
+    GPUS = self.weight.device
+    if isinstance(GPUS, tuple) and len(GPUS) > 1:
+      y = Tensor(Tensor.empty((batch//len(GPUS), seq, self.weight.shape[0]), dtype=dtypes.float, device=GPUS).uop.multi(0), device=GPUS)
+    else:
+      y = Tensor.empty((batch, seq, self.weight.shape[0]), dtype=dtypes.float)
+    y = Tensor.custom_kernel(y, x_fp8, w_fp8, fxn=custom_matmul, grad_fxn=custom_matmul_backward)[0]
+    y = y * w_scale * x_scale
+    if self.bias is not None: y = y + self.bias
+    if original_ndim == 2: y = y.reshape(batch, self.weight.shape[0])
+    return y.cast(x.dtype)
+
+def _replace_linear(layer: nn.Linear):
+  fp8_linear = FP8Linear(layer.weight.shape[1], layer.weight.shape[0], layer.bias is not None)
+  fp8_linear.weight = layer.weight
+  if layer.bias is not None: fp8_linear.bias = layer.bias
+  return fp8_linear
+
+def _swap_linear_with_fp8(model, module_filter_fn:Callable[[Any, str],bool]|None=None, fqn:str="", parent:Any|None=None,
+                          attr_name:str="", visited:set|None=None):
+  if visited is None: visited = set()
+  if id(model) in visited: return
+  visited.add(id(model))
+  if isinstance(model, (str, int, float, bool, type(None), Tensor, UOp)): return
+  elif isinstance(model, nn.Linear):
+    if module_filter_fn is not None and not module_filter_fn(model, fqn): return
+    fp8_linear = _replace_linear(model)
+    if parent is not None and attr_name:
+      setattr(parent, attr_name, fp8_linear)
+  elif isinstance(model, list):
+    for i, item in enumerate(model):
+      child_fqn = f"{fqn}.{i}" if fqn else str(i)
+      if isinstance(item, nn.Linear) and (module_filter_fn is None or module_filter_fn(item, child_fqn)): model[i] = _replace_linear(item)
+      else: _swap_linear_with_fp8(item, module_filter_fn, child_fqn, None, "", visited)
+  elif isinstance(model, dict):
+    for key, item in list(model.items()):
+      child_fqn = f"{fqn}.{key}" if fqn else str(key)
+      if isinstance(item, nn.Linear) and (module_filter_fn is None or module_filter_fn(item, child_fqn)): model[key] = _replace_linear(item)
+      else: _swap_linear_with_fp8(item, module_filter_fn, child_fqn, None, "", visited)
+  elif hasattr(model, "__dict__"):
+    for attr_key in list(vars(model).keys()):
+      try: attr = getattr(model, attr_key)
+      except Exception: continue
+      child_fqn = f"{fqn}.{attr_key}" if fqn else attr_key
+      _swap_linear_with_fp8(attr, module_filter_fn, child_fqn, model, attr_key, visited)
+
+def convert_to_float8_training(model, module_filter_fn:Callable[[Any,str],bool]|None=None):
+  _swap_linear_with_fp8(model, module_filter_fn, "", None, "")
+  return model

tinygrad_repo/extra/gemm/.gitignore

@@ -0,0 +1,2 @@
+*.ll
+fp32_sgemm_amd

tinygrad_repo/extra/gemm/amd_asm_matmul.py

@@ -0,0 +1,500 @@
+# RDNA3 128x128 tiled GEMM kernel - DSL version
+# Computes C = A @ B for NxN float32 matrices using 128x128 tiles
+#
+# Architecture: RDNA3 (gfx1100)
+# Tile size: 128x128 (each workgroup computes one tile of C)
+# Workgroup: 128 threads (arranged as 32x4 for coalesced memory access)
+# Inner loop: 8 iterations per K-block, processing 8 columns of A and 8 rows of B
+#
+# Accumulators: 128 vgprs (v[2-129])
+
+import numpy as np
+from tinygrad import Tensor, Device, Context, GlobalCounters
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.helpers import getenv, colored
+from tinygrad.dtype import dtypes, AddrSpace
+from tinygrad.engine.realize import Estimates, run_linear
+from tinygrad.renderer.amd.dsl import s, v, VCC_LO, NULL
+from tinygrad.runtime.autogen.amd.rdna3.ins import *
+
+# =============================================================================
+# Kernel constants
+# =============================================================================
+LDS_SIZE = 8320       # Local data share size in bytes
+LDS_A_STRIDE = 0x210  # LDS stride for A tile (528 bytes)
+LDS_B_STRIDE = 0x200  # LDS stride for B tile (512 bytes)
+LDS_BASE_OFFSET = 0x1080  # Base LDS offset for tiles
+ADDR_MASK = 0x3fffff80    # Address alignment mask
+
+# =============================================================================
+# Named register assignments (VGPRs)
+# =============================================================================
+V_LANE_ID = 0             # lane_id set on startup
+# Use tile gaps (v146-159) for named regs to minimize max VGPR
+V_LANE_ID_MOD8 = 146      # lane_id & 7
+V_LANE_MOD8_X4 = 147      # (lane_id & 7) << 2
+V_LANE_DIV8_X4 = 150      # ((lane_id >> 3) & 3) << 2
+V_LDS_B_BASE = 151        # LDS B-tile base address for inner loop
+V_LDS_A_BASE = 154        # LDS A-tile base address for inner loop
+V_GLOBAL_A_ADDR = 155     # global memory A prefetch address
+V_GLOBAL_B_ADDR = 158     # global memory B prefetch address
+V_LDS_A_ADDR = 159        # single base register for A stores
+V_LDS_B_ADDR = 162        # single base register for B stores
+
+# LDS tile register destinations - SEPARATE from DATA to avoid overlap
+# A on banks 2-3, B on banks 0-1 to avoid bank conflicts in VOPD
+V_A_TILE_REGS = [130, 134, 138, 142]  # A tile: banks 2,2,2,2 (130%4=2, etc.)
+V_B_TILE_REGS = [132, 136, 140, 144, 148, 152, 156, 160]  # B tile: banks 0,0,0,0,0,0,0,0
+
+# =============================================================================
+# Named register assignments (SGPRs)
+# =============================================================================
+S_OUT_PTR = (0, 1)        # output C matrix base pointer
+S_WORKGROUP_X = 2         # workgroup_id_x (system SGPR, follows user SGPRs)
+S_WORKGROUP_Y = 3         # workgroup_id_y (system SGPR)
+S_DIM_N = 4               # matrix dimension N
+S_LOOP_BOUND = 7          # K-8 (loop termination bound)
+S_LOOP_CTR = 12           # loop counter (increments by 8)
+S_PREFETCH_FLAG = 13      # prefetch condition flag / row stride in epilogue
+S_TILE_X = 14             # workgroup_x << 7
+S_TILE_Y = 15             # workgroup_y << 7
+# Kernarg load destinations
+S_KERNARG_A = (20, 21)    # A pointer from kernarg
+S_KERNARG_B = (22, 23)    # B pointer from kernarg
+# Prefetch base pointers (8 pairs each, B: N*4 bytes apart, A: N*64 bytes apart)
+S_PREFETCH_B = 24         # s[24:39] - 8 B tile pointers
+S_PREFETCH_A = 40         # s[40:55] - 8 A tile pointers
+
+# =============================================================================
+# Data tables
+# =============================================================================
+
+# Accumulator grid: ACC_GRID[a_idx][b_idx] = vgpr for C[a,b]
+# a_idx: which A value (0-7), b_idx: which B value (0-15)
+# Scattered due to VOPD bank constraints (vdst_x % 4 != vdst_y % 4)
+# Range is from v2 - v129
+ACC_GRID = [
+  [  5,  3,  9,  8,   37, 35, 41, 40,   69, 67, 73, 72,  101, 99,105,104],  # a0
+  [  4,  2,  7,  6,   36, 34, 39, 38,   68, 66, 71, 70,  100, 98,103,102],  # a1
+  [ 17, 16, 13, 11,   49, 48, 45, 43,   81, 80, 77, 75,  113,112,109,107],  # a2
+  [ 15, 14, 12, 10,   47, 46, 44, 42,   79, 78, 76, 74,  111,110,108,106],  # a3
+  [ 21, 19, 25, 24,   53, 51, 57, 56,   85, 83, 89, 88,  117,115,121,120],  # a4
+  [ 20, 18, 23, 22,   52, 50, 55, 54,   84, 82, 87, 86,  116,114,123,122],  # a5
+  [125,128, 29, 27,   33, 32, 61, 59,   65, 64, 93, 91,   97, 96,129,127],  # a6
+  [119,118, 28, 26,   31, 30, 60, 58,   63, 62, 92, 90,   95, 94,124,126],  # a7
+]
+
+# Optimized (a_pair, b_pair) iteration order for better GPU scheduling
+# Interleaves A and B pairs to maximize instruction-level parallelism
+FMAC_PAIR_ORDER = [
+  (0,0),(0,1),(1,1),(1,0), (2,0),(2,1),(3,1),(3,2), (0,2),(0,3),(1,3),(1,2), (2,2),(2,3),(3,3),(3,4),
+  (0,4),(0,5),(1,5),(1,4), (2,4),(2,5),(3,5),(3,6), (0,6),(0,7),(1,7),(1,6), (2,6),(2,7),(3,7),(3,0),
+]
+
+def derive_fmac_pattern(acc_grid, a_tile_regs=None, b_tile_regs=None):
+  """Generate 64 dual FMAC ops from accumulator grid with optimized iteration order."""
+  pattern = []
+  for idx, (a_pair, b_pair) in enumerate(FMAC_PAIR_ORDER):
+    a_even, a_odd = a_pair * 2, a_pair * 2 + 1
+    b_even, b_odd = b_pair * 2, b_pair * 2 + 1
+    a_base, b_base = a_tile_regs[a_pair], b_tile_regs[b_pair]
+    # Op 1: normal order -> C[a_even, b_even] + C[a_odd, b_odd]
+    pattern.append((acc_grid[a_even][b_even], acc_grid[a_odd][b_odd],
+                   a_base, b_base, a_base+1, b_base+1))
+    # Op 2: alternate swapping A vs B to vary register banks
+    if idx % 2 == 0:  # swap B
+      pattern.append((acc_grid[a_even][b_odd], acc_grid[a_odd][b_even],
+                     a_base, b_base+1, a_base+1, b_base))
+    else:  # swap A
+      pattern.append((acc_grid[a_odd][b_even], acc_grid[a_even][b_odd],
+                     a_base+1, b_base, a_base, b_base+1))
+  return pattern
+
+# Derived: 64 dual FMAC operations
+FMAC_PATTERN = derive_fmac_pattern(ACC_GRID, V_A_TILE_REGS, V_B_TILE_REGS)
+
+def derive_permute_swaps(acc_grid, out_regs):
+  """Derive swap sequence to permute accumulators from FMAC layout to output order.
+
+  After FMAC loop: acc_grid[a][b] holds C[a,b]
+  Output order: for row_half in 0,1; col_group in 0-3; row_in_group in 0-3; b_off in 0-3
+    -> need C[row_half*4 + row_in_group, col_group*4 + b_off] in specified reg order
+  """
+  def target_ab(i):
+    row_half, col_group = i // 64, (i // 16) % 4
+    row_in_group, b_off = (i // 4) % 4, i % 4
+    return (row_half * 4 + row_in_group, col_group * 4 + b_off)
+
+  reg_contents = {acc_grid[a][b]: (a, b) for a in range(8) for b in range(16)}
+  ab_location = {ab: r for r, ab in reg_contents.items()}
+
+  swaps = []
+  for i in range(128):
+    target_reg, needed_ab = out_regs[i], target_ab(i)
+    current_reg = ab_location[needed_ab]
+    if current_reg != target_reg:
+      swaps.append((current_reg, target_reg))
+      ab_at_target = reg_contents.get(target_reg)
+      reg_contents[target_reg], ab_location[needed_ab] = needed_ab, target_reg
+      if ab_at_target is not None:
+        reg_contents[current_reg], ab_location[ab_at_target] = ab_at_target, current_reg
+  return swaps
+
+# Derived: swap sequence to arrange accumulators for output
+# Each group of 4 registers is ascending for direct global_store_b128
+OUT_REGS = [r for i in range(32) for r in range(126 - i*4, 130 - i*4)]
+PERMUTE_SWAPS = derive_permute_swaps(ACC_GRID, OUT_REGS)
+
+# =============================================================================
+# LDS tile staging registers
+# =============================================================================
+# DATA regs receive contiguous global prefetch, then write to LDS
+# TILE regs receive scattered LDS loads (ds_load_b64 pairs), then feed FMACs
+# Contiguous layout with mod4=[3,0,1,2,3,0,1,2] for bank conflict avoidance
+V_LDS_A_DATA = [163, 164, 165, 166, 167, 168, 169, 170]
+V_LDS_B_DATA = [171, 172, 173, 174, 175, 176, 177, 178]
+
+# Initial tile prefetch: (vdst, saddr_lo) - load into A data regs using B prefetch pointers (s[24:31])
+INIT_PREFETCH = [(V_LDS_A_DATA[i], S_PREFETCH_B+2*i) for i in range(4)]
+
+# Global memory prefetch schedule: (vdst1, vdst2, addr_vreg, saddr_lo1, saddr_lo2)
+# First 2 pairs from B prefetch pointers (s[32:39]), next 4 pairs from A prefetch pointers (s[40:55])
+PREFETCH_LOADS = [(V_LDS_A_DATA[4+2*i], V_LDS_A_DATA[4+2*i+1], V_GLOBAL_B_ADDR, S_PREFETCH_B+8+4*i, S_PREFETCH_B+10+4*i) for i in range(2)] + \
+                 [(V_LDS_B_DATA[2*(i-2)], V_LDS_B_DATA[2*(i-2)+1], V_GLOBAL_A_ADDR, S_PREFETCH_A+4*(i-2), S_PREFETCH_A+2+4*(i-2)) for i in range(2, 6)]
+
+# =============================================================================
+# Kernel class
+# =============================================================================
+
+class Kernel:
+  def __init__(self): self.instructions, self.labels, self.pos = [], {}, 0
+  def label(self, name): self.labels[name] = self.pos
+
+  def emit(self, inst, target=None):
+    self.instructions.append(inst)
+    inst._target, inst._pos = target, self.pos
+    self.pos += inst.size()
+    return inst
+
+  def waitcnt(self, lgkm=None, vm=None):
+    """Wait for memory operations. lgkm=N waits until N lgkm ops remain, vm=N waits until N vmem ops remain."""
+    vmcnt, lgkmcnt, expcnt = vm if vm is not None else 63, lgkm if lgkm is not None else 63, 7
+    waitcnt = (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
+    self.emit(s_waitcnt(simm16=waitcnt))
+
+  def finalize(self):
+    """Patch branch offsets and return the finalized instruction list."""
+    for inst in self.instructions:
+      if inst._target is None: continue
+      offset_dwords = (self.labels[inst._target] - inst._pos - inst.size()) // 4
+      if not -32768 <= offset_dwords <= 32767: raise ValueError(f"branch to '{inst._target}' offset {offset_dwords} exceeds simm16 range")
+      inst.simm16 = offset_dwords
+    return self.instructions
+
+
+# =============================================================================
+# Kernel builder
+# =============================================================================
+
+def build_kernel(N):
+  assert N % 128 == 0, f"N must be a multiple of 128 (tile size), got {N}"
+  assert N >= 256, f"N must be >= 256 (prefetch pipeline requires at least 2 K-blocks), got {N}"
+  k = Kernel()
+
+  # ===========================================================================
+  # PROLOGUE: Load kernel arguments, compute tile coordinates and addresses
+  # ===========================================================================
+  k.emit(s_load_b128(sdata=s[S_KERNARG_A[0]:S_KERNARG_B[1]], sbase=s[0:1], offset=0x0, soffset=NULL))
+  k.emit(s_load_b64(sdata=s[S_OUT_PTR[0]:S_OUT_PTR[1]], sbase=s[0:1], offset=0x10, soffset=NULL))
+  k.emit(s_mov_b32(s[S_DIM_N], N))
+  k.emit(s_mov_b32(s[S_LOOP_CTR], 0))  # used by LDS swizzle, always 0 for valid workgroups
+  k.emit(s_lshl_b32(s[S_TILE_X], s[S_WORKGROUP_X], 7))
+  k.emit(s_lshl_b32(s[S_TILE_Y], s[S_WORKGROUP_Y], 7))
+
+  # Lane-derived values
+  k.emit(v_and_b32_e32(v[V_LANE_ID_MOD8], 7, v[V_LANE_ID]))
+  k.emit(v_lshrrev_b32_e32(v[4], 3, v[V_LANE_ID]))
+  k.emit(v_or_b32_e32(v[1], s[S_TILE_X], v[V_LANE_ID]))
+  k.emit(v_or_b32_e32(v[22], s[S_TILE_Y], v[4]))
+  k.emit(v_lshlrev_b32_e32(v[V_LANE_MOD8_X4], 2, v[V_LANE_ID_MOD8]))
+  k.waitcnt(lgkm=0)
+
+  # Compute 8 A and B matrix tile base pointers for prefetch
+  k.emit(s_mov_b64(s[S_PREFETCH_B:S_PREFETCH_B+1], s[S_KERNARG_B[0]:S_KERNARG_B[1]]))  # B[0]: no offset
+  for i in range(1, 8):  # B: each pointer 1 row of B apart (N*4 bytes)
+    k.emit(s_add_u32(s[S_PREFETCH_B+i*2], s[S_KERNARG_B[0]], i * N * 4))
+    k.emit(s_addc_u32(s[S_PREFETCH_B+i*2+1], s[S_KERNARG_B[1]], 0))
+  k.emit(s_mov_b64(s[S_PREFETCH_A:S_PREFETCH_A+1], s[S_KERNARG_A[0]:S_KERNARG_A[1]]))  # A[0]: no offset
+  for i in range(1, 8):  # A: each pointer 16 rows of A apart (16*N*4 bytes)
+    k.emit(s_add_u32(s[S_PREFETCH_A+i*2], s[S_KERNARG_A[0]], i * N * 64))
+    k.emit(s_addc_u32(s[S_PREFETCH_A+i*2+1], s[S_KERNARG_A[1]], 0))
+
+  # Global prefetch addresses: B = (tile_x + lane_id) * 4, A = (tile_y*N + (lane_id/8)*N + lane_id%8) * 4
+  k.emit(v_add_nc_u32_e32(v[V_GLOBAL_B_ADDR], s[S_TILE_X], v[V_LANE_ID]))
+  k.emit(v_lshlrev_b32_e32(v[V_GLOBAL_B_ADDR], 2, v[V_GLOBAL_B_ADDR]))
+  k.emit(s_mul_i32(s[19], s[S_TILE_Y], N))
+  k.emit(v_mul_lo_u32(v[V_GLOBAL_A_ADDR], v[4], N))                          # (lane_id/8)*N
+  k.emit(v_add_nc_u32_e32(v[V_GLOBAL_A_ADDR], v[V_LANE_ID_MOD8], v[V_GLOBAL_A_ADDR]))  # + lane_id%8
+  k.emit(v_add_nc_u32_e32(v[V_GLOBAL_A_ADDR], s[19], v[V_GLOBAL_A_ADDR]))
+  k.emit(v_lshlrev_b32_e32(v[V_GLOBAL_A_ADDR], 2, v[V_GLOBAL_A_ADDR]))
+
+  # Do initial loads
+  for vdst, saddr_lo in INIT_PREFETCH:
+    k.emit(global_load_b32(vdst=v[vdst], addr=v[V_GLOBAL_B_ADDR], saddr=s[saddr_lo:saddr_lo+1]))
+  for iter in range(6):
+    vdst1, vdst2, addr, slo1, slo2 = PREFETCH_LOADS[iter]
+    k.emit(global_load_b32(vdst=v[vdst1], addr=v[addr], saddr=s[slo1:slo1+1]))
+    k.emit(global_load_b32(vdst=v[vdst2], addr=v[addr], saddr=s[slo2:slo2+1]))
+
+  # ===========================================================================
+  # LDS store address computation (bank-conflict-avoiding swizzle)
+  # ===========================================================================
+  # This section computes LDS store addresses with a swizzle pattern to avoid bank conflicts.
+  # The swizzle ensures that threads in the same wavefront write to different LDS banks.
+  # Formula: swizzled_addr = base + (lane_id & 7) * LDS_A_STRIDE + swizzle_offset
+  # where swizzle_offset depends on (lane_id >> 3) to distribute across banks.
+  k.emit(v_add_nc_u32_e32(v[9], s[S_LOOP_CTR], v[22]))  # row 0 base
+  k.emit(v_and_b32_e32(v[9], ADDR_MASK, v[9]))
+  k.emit(v_sub_nc_u32_e32(v[9], v[22], v[9]))  # row 0 swizzle offset
+  k.emit(v_lshlrev_b32_e32(v[9], 2, v[9]))  # * 4
+  k.emit(v_mad_u32_u24(v[V_LDS_B_ADDR], LDS_A_STRIDE, v[V_LANE_ID_MOD8], v[9]))
+
+  # For V_LDS_A_BASE and epilogue
+  k.emit(v_bfe_u32(v[2], v[V_LANE_ID], 3, 2))  # v[2] = (lane_id >> 3) & 3
+  k.emit(v_lshlrev_b32_e32(v[V_LANE_DIV8_X4], 2, v[2]))
+
+  # Compute LDS load/store base addresses for inner loop
+  k.emit(v_lshlrev_b32_e32(v[2], 4, v[2]))
+  k.emit(v_and_b32_e32(v[3], 0x7F, v[1]))  # simplified from 3 lines
+  k.emit(v_lshl_or_b32(v[V_LDS_B_BASE], v[V_LANE_ID_MOD8], 4, LDS_BASE_OFFSET))
+  k.emit(v_lshl_add_u32(v[V_LDS_A_ADDR], v[3], 2, LDS_BASE_OFFSET))
+  k.emit(v_lshlrev_b32_e32(v[3], 2, v[V_LANE_ID]))
+  k.emit(v_and_or_b32(v[V_LDS_A_BASE], 0x180, v[3], v[2]))
+
+  # Do initial stores
+  k.waitcnt(vm=0)
+  for i in range(4):  # A tile: 8 values via 4 stride64 stores
+    k.emit(ds_store_2addr_stride64_b32(addr=v[V_LDS_A_ADDR], data0=v[V_LDS_A_DATA[i*2]], data1=v[V_LDS_A_DATA[i*2+1]], offset0=i*4, offset1=i*4+2))
+  for i in range(8):  # B tile: 8 values via 8 scalar stores with 64-byte spacing
+    offset = i * 64
+    k.emit(ds_store_b32(addr=v[V_LDS_B_ADDR], data0=v[V_LDS_B_DATA[i]], offset0=offset & 0xFF, offset1=offset >> 8))
+
+  # Zero all 128 accumulators using VOPD dual moves (64 instructions instead of 128)
+  for i in range(0, len(OUT_REGS), 2):
+    k.emit(VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[OUT_REGS[i]], vdsty=v[OUT_REGS[i+1]], srcx0=0, srcy0=0))
+  k.emit(s_add_i32(s[S_LOOP_BOUND], s[S_DIM_N], -8))
+
+  # S_LOOP_CTR is already 0 from prologue initialization
+  k.emit(s_branch(), target='LOOP_ENTRY')
+
+  # ===========================================================================
+  # MAIN GEMM LOOP
+  # ===========================================================================
+
+  NO_ALU, NO_DS, NO_GLOBAL = getenv("NO_ALU", 0), getenv("NO_DS", 0), getenv("NO_GLOBAL", 0)
+
+  k.label('LOOP_INC')
+  k.emit(s_add_i32(s[S_LOOP_CTR], s[S_LOOP_CTR], 8))
+  k.emit(s_cmp_ge_i32(s[S_LOOP_CTR], s[S_DIM_N]))
+  k.emit(s_cbranch_scc1(), target='EPILOGUE')
+
+  k.label('LOOP_ENTRY')
+  k.emit(s_cmp_lt_i32(s[S_LOOP_CTR], s[S_LOOP_BOUND]))
+  k.emit(s_cselect_b32(s[S_PREFETCH_FLAG], -1, 0))  # s_cselect doesn't modify SCC
+  k.emit(s_cbranch_scc0(), target='SKIP_PREFETCH')  # branch if loop_ctr >= loop_bound
+
+  if not NO_GLOBAL:
+    # Advance prefetch pointers (VGPR)
+    #k.emit(v_add_nc_u32_e32(v[V_GLOBAL_B_ADDR], N * 32, v[V_GLOBAL_B_ADDR]))
+    #k.emit(v_add_nc_u32_e32(v[V_GLOBAL_A_ADDR], 0x20, v[V_GLOBAL_A_ADDR]))
+
+    # Advance prefetch pointers (64-bit adds): B advances 8 rows (8*N*4 bytes), A advances 8 cols (8*4 bytes)
+    k.emit(s_clause(simm16=31))
+    for i in range(8):
+      k.emit(s_add_u32(s[S_PREFETCH_B+i*2], s[S_PREFETCH_B+i*2], N * 32))
+      k.emit(s_addc_u32(s[S_PREFETCH_B+i*2+1], s[S_PREFETCH_B+i*2+1], 0))
+    for i in range(8):
+      k.emit(s_add_u32(s[S_PREFETCH_A+i*2], s[S_PREFETCH_A+i*2], 0x20))
+      k.emit(s_addc_u32(s[S_PREFETCH_A+i*2+1], s[S_PREFETCH_A+i*2+1], 0))
+
+    # do the fetch
+    for vdst, saddr_lo in INIT_PREFETCH:
+      k.emit(global_load_b32(vdst=v[vdst], addr=v[V_GLOBAL_B_ADDR], saddr=s[saddr_lo:saddr_lo+1]))
+
+  k.label('SKIP_PREFETCH')
+
+  # wait for local stores to finish (either initial or loop)
+  # then sync the warp so it's safe to load local
+  k.waitcnt(lgkm=0)
+  k.emit(s_barrier())
+
+  # 8 inner loop iterations
+  for iter in range(8):
+    # Load A tile (4 pairs) and B tile (8 pairs) from LDS
+    if not NO_DS:
+      k.emit(s_clause(simm16=len(V_A_TILE_REGS) + len(V_B_TILE_REGS) - 1))  # 12 loads total: 4 A + 8 B
+      # A tile: 4 ds_load_b64
+      for i, vdst in enumerate(V_A_TILE_REGS):
+        a_off = (i & 1) * 8 + (i >> 1) * 64 + iter * LDS_A_STRIDE
+        k.emit(ds_load_b64(vdst=v[vdst:vdst+1], addr=v[V_LDS_A_BASE], offset0=a_off & 0xFF, offset1=a_off >> 8))
+      # B tile: 8 ds_load_b64
+      for i, vdst in enumerate(V_B_TILE_REGS):
+        b_off = (i & 1) * 8 + (i & 2) * 64 + (i >> 2) * 256 + iter * LDS_B_STRIDE
+        k.emit(ds_load_b64(vdst=v[vdst:vdst+1], addr=v[V_LDS_B_BASE], offset0=b_off & 0xFF, offset1=b_off >> 8))
+
+    # Issue global prefetch (first 6 iterations only)
+    if iter < 6 and not NO_GLOBAL:
+      vdst1, vdst2, addr, slo1, slo2 = PREFETCH_LOADS[iter]
+      k.emit(global_load_b32(vdst=v[vdst1], addr=v[addr], saddr=s[slo1:slo1+1]))
+      k.emit(global_load_b32(vdst=v[vdst2], addr=v[addr], saddr=s[slo2:slo2+1]))
+
+    # 64 dual FMACs
+    k.waitcnt(lgkm=0)
+    if not NO_ALU:
+      k.emit(s_clause(simm16=len(FMAC_PATTERN)-1))
+      for i, (vdst_x, vdst_y, ax, bx, ay, by) in enumerate(FMAC_PATTERN):
+        k.emit(VOPD(VOPDOp.V_DUAL_FMAC_F32, VOPDOp.V_DUAL_FMAC_F32,
+                    vdstx=v[vdst_x], vdsty=v[vdst_y], srcx0=v[ax], vsrcx1=v[bx], srcy0=v[ay], vsrcy1=v[by]))
+
+  # wait for all global loads to finish
+  # then sync the warp so it's safe to store local
+  k.waitcnt(vm=0)
+  k.emit(s_barrier())
+
+  # Store prefetched data to LDS
+  # NOTE: Register naming reflects LDS tile organization, not source matrix:
+  #   V_LDS_A_DATA (v155-162) holds data that goes to LDS A-tile region
+  #   V_LDS_B_DATA (v163-170) holds data that goes to LDS B-tile region
+  # The data sources are swapped: A-tile receives B matrix rows, B-tile receives A matrix columns
+  if not NO_DS:
+    for i in range(4):  # A tile: 8 values via 4 stride64 stores
+      k.emit(ds_store_2addr_stride64_b32(addr=v[V_LDS_A_ADDR], data0=v[V_LDS_A_DATA[i*2]], data1=v[V_LDS_A_DATA[i*2+1]], offset0=i*4, offset1=i*4+2))
+    for i in range(8):  # B tile: 8 values via 8 scalar stores with 64-byte spacing
+      offset = i * 64
+      k.emit(ds_store_b32(addr=v[V_LDS_B_ADDR], data0=v[V_LDS_B_DATA[i]], offset0=offset & 0xFF, offset1=offset >> 8))
+
+  k.emit(s_branch(), target='LOOP_INC')
+
+  # ===========================================================================
+  # EPILOGUE: Permute and store results
+  # ===========================================================================
+  k.label('EPILOGUE')
+
+  # Rearrange accumulators from FMAC layout to contiguous output order
+  for a, b in PERMUTE_SWAPS:
+    k.emit(v_swap_b32_e32(v[a], v[b]))
+
+  # Compute output base coordinates
+  # v[130] = col_base = tile_x + (lane_id & 7) * 4
+  # v[131] = row_base = tile_y + (lane_id & 0x60) + ((lane_id >> 3) & 3) * 4
+  # v[132] = 0 (for 64-bit address high part)
+  k.emit(v_add_nc_u32_e32(v[130], s[S_TILE_X], v[V_LANE_MOD8_X4]))
+  k.emit(v_and_b32_e32(v[131], 0x60, v[V_LANE_ID]))
+  k.emit(v_add_nc_u32_e32(v[131], s[S_TILE_Y], v[131]))
+  k.emit(v_add_nc_u32_e32(v[131], v[V_LANE_DIV8_X4], v[131]))
+  k.emit(v_mov_b32_e32(v[132], 0))
+
+  # Precompute row offsets: v[133-136] for rows 0-3, v[137-140] for rows 16-19
+  for base, row_off in [(133, 0), (137, 16)]:
+    if row_off: k.emit(v_add_nc_u32_e32(v[141], row_off, v[131]))
+    k.emit(v_mul_lo_u32(v[base], v[141] if row_off else v[131], s[S_DIM_N]))
+    for j in range(3): k.emit(v_add_nc_u32_e32(v[base + 1 + j], s[S_DIM_N], v[base + j]))
+
+  # s[S_PREFETCH_FLAG] = row stride in bytes (N * 4)
+  k.emit(s_lshl_b32(s[S_PREFETCH_FLAG], s[S_DIM_N], 2))
+
+  # Store 128 output values as 32 groups of 4 (128-bit stores)
+  # Layout: 2 row halves (0-3, 16-19) x 4 col groups x 4 rows = 32 stores of 4 floats
+  for i, (row_half, col_off, row_in_group) in enumerate([(rh, co, ri)
+      for rh in range(2) for co in [0, 32, 64, 96] for ri in range(4)]):
+    row = row_half * 16 + row_in_group
+    src = OUT_REGS[i*4]  # first reg of ascending group of 4
+
+    if row_in_group == 0:
+      # First row of group: compute full address
+      if col_off == 0: k.emit(v_mov_b32_e32(v[141], v[130]))
+      else: k.emit(v_add_nc_u32_e32(v[141], col_off, v[130]))
+      row_base = 133 + row if row < 4 else 137 + row - 16
+      k.emit(v_add_nc_u32_e32(v[141], v[row_base], v[141]))
+      k.emit(v_lshlrev_b32_e32(v[141], 2, v[141]))
+      k.emit(v_add_co_u32(v[141], VCC_LO, s[S_OUT_PTR[0]], v[141]))
+      k.emit(v_add_co_ci_u32_e32(v[142], s[S_OUT_PTR[1]], v[132]))
+    else:
+      # Subsequent rows: add stride
+      k.emit(v_add_co_u32(v[141], VCC_LO, s[S_PREFETCH_FLAG], v[141]))
+      k.emit(v_add_co_ci_u32_e32(v[142], v[142], v[132]))
+
+    k.emit(global_store_b128(addr=v[141:142], data=v[src:src+3], saddr=NULL))
+
+  k.emit(s_sendmsg(simm16=3))  # DEALLOC_VGPRS
+  k.emit(s_endpgm())
+
+  return k.finalize()
+
+# =============================================================================
+# Test harness
+# =============================================================================
+
+N = getenv("N", 4096)
+BLOCK_M, BLOCK_N = 128, 128
+THREADS = 128
+
+def test_matmul():
+  dev = Device[Device.DEFAULT]
+  print(f"Device arch: {dev.renderer.target.arch}")
+
+  insts = build_kernel(N)
+
+  rng = np.random.default_rng(42)
+  a = Tensor(rng.random((N, N), dtype=np.float32) - 0.5)
+  b = Tensor(rng.random((N, N), dtype=np.float32) - 0.5)
+  c = Tensor.empty(N, N)
+  Tensor.realize(a, b, c)
+
+  grid, local = (N // BLOCK_N, N // BLOCK_M, 1), (THREADS, 1, 1)
+  print(f"Grid: {grid}, Local: {local}")
+
+  dname:str = Device.DEFAULT
+  def asm_kernel(A:UOp, B:UOp, C:UOp) -> UOp:
+    gidxs = [UOp.special(n, f"gidx{i}") for i,n in enumerate(grid)]
+    lidxs = [UOp.special(n, f"lidx{i}") for i,n in enumerate(local)]
+    lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=max(LDS_SIZE, 65536//getenv("LIMIT_OCC", 65536)), addrspace=AddrSpace.LOCAL), (), 'lds')
+    sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs, arg=KernelInfo(name=colored("kernel", "cyan"),
+                                                                                  estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+  c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
+  linear = c.schedule_linear()
+
+  ets = []
+  with Context(DEBUG=2):
+    for _ in range(getenv("CNT", 5)):
+      start = GlobalCounters.time_sum_s
+      run_linear(linear)
+      ets.append(GlobalCounters.time_sum_s - start)
+  print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")
+
+  if getenv("VERIFY", 1):
+    GlobalCounters.reset()
+    with Context(DEBUG=2): tc = (a @ b).realize()
+    with Context(DEBUG=0): err = (c - tc).square().mean().item()
+    print(f"mean squared error {err}")
+    if err != err or err > 1e-06:
+      c_np, tc_np = c.numpy(), tc.numpy()
+      for bi in range(N // 128):
+        for bj in range(N // 128):
+          blk_c = c_np[bi*128:(bi+1)*128, bj*128:(bj+1)*128]
+          blk_ref = tc_np[bi*128:(bi+1)*128, bj*128:(bj+1)*128]
+          blk_diff = blk_c - blk_ref
+          zero_rows = [i for i in range(128) if np.all(np.abs(blk_c[i,:]) < 1e-10)]
+          nz_rows = [i for i in range(128) if i not in zero_rows]
+          nz_mse = float(np.mean(blk_diff[nz_rows,:]**2)) if nz_rows else 0
+          print(f"Block ({bi},{bj}): zero_rows={zero_rows}, nz_rows_mse={nz_mse:.2e}")
+          # show first few non-zero row comparisons
+          if nz_rows and nz_mse > 1e-6:
+            for r in nz_rows[:3]:
+              print(f"  row {r} asm[0:8]:  {blk_c[r,:8]}")
+              print(f"  row {r} ref[0:8]:  {blk_ref[r,:8]}")
+      raise RuntimeError("matmul is wrong!")
+
+if __name__ == "__main__":
+  test_matmul()

tinygrad_repo/extra/gemm/amd_copy_matmul.py

@@ -0,0 +1,118 @@
+from tinygrad import Device, UOp, getenv
+from tinygrad.uop.ops import AxisType, KernelInfo, Ops
+from tinygrad.dtype import AddrSpace, dtypes
+
+N = getenv("N", 4096)
+M = getenv("M", N)
+K = getenv("K", N)
+
+WARP_SIZE = 32
+BLOCK_M, BLOCK_N = 128, 128
+BLOCK_K = getenv("BK", 16)
+assert N % BLOCK_N == 0 and M % BLOCK_M == 0 and K % BLOCK_K == 0
+
+use_wmma = getenv("WMMA")
+if use_wmma:
+  is_rdna4 = Device[Device.DEFAULT].renderer.target.arch.startswith("gfx12")
+
+  WAVES_M, WAVES_N = 2, 2
+  LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 16
+
+  # wmma params
+  WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
+  WMMA_ACC = WMMA_M // LANES_PER_WAVE_M
+  UNROLL_M, UNROLL_N = (WMMA_ACC, 1) if is_rdna4 else (1, 1)
+else:
+  WAVES_M, WAVES_N = 4, 1
+  LANES_PER_WAVE_M, LANES_PER_WAVE_N = 4, 8
+  UNROLL_M, UNROLL_N = 4, 4
+
+# total lanes must be the warp size
+assert LANES_PER_WAVE_M*LANES_PER_WAVE_N == WARP_SIZE
+
+# WARP_SIZE * total waves
+THREADS_PER_BLOCK = WARP_SIZE * WAVES_M * WAVES_N
+
+# accumulator size
+TM = BLOCK_M // (WAVES_M * LANES_PER_WAVE_M)
+TN = BLOCK_N // (WAVES_N * LANES_PER_WAVE_N)
+
+def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
+  wave_m = UOp.range(WAVES_M, 2, AxisType.LOCAL)
+  wave_n = UOp.range(WAVES_N, 3, AxisType.LOCAL)
+  lane = UOp.range(WARP_SIZE, -1, AxisType.WARP)
+  tid = (wave_m * WAVES_N + wave_n) * WARP_SIZE + lane
+
+  # -- GLOBAL -> LOCAL --
+  # wmma: spatial outer, k inner (k contiguous for vectorized WMMA tile loads)
+  # gemm: k outer, spatial inner
+  A_local = UOp.placeholder((BLOCK_M, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_M), a.dtype.base, slot=0, addrspace=AddrSpace.LOCAL)
+  B_local = UOp.placeholder((BLOCK_N, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_N), b.dtype.base, slot=1, addrspace=AddrSpace.LOCAL)
+
+  a = a.reshape(K // BLOCK_K, BLOCK_K, BLOCK_M)
+  b = b.reshape(K // BLOCK_K, BLOCK_K, BLOCK_N)
+  k_tile = UOp.range(K // BLOCK_K, 100, AxisType.REDUCE)
+
+  # copy with transpose for wmma (input is k×spatial, LDS is spatial×k)
+  A_copy = A_local.permute((1,0)) if use_wmma else A_local
+  B_copy = B_local.permute((1,0)) if use_wmma else B_local
+  A_store = A_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(a[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
+  B_store = B_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(b[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
+  barrier = UOp.barrier(A_store, B_store)
+  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
+
+  # -- COMPUTE --
+  lane_m, lane_n = lane // LANES_PER_WAVE_N, lane % LANES_PER_WAVE_N
+
+  # accumulator (unified: both paths use (TM, TN) with scalar dtypes.float)
+  acc = UOp.placeholder((TM, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  acc = acc.after(acc.store(acc.zeros_like()))
+
+  if use_wmma:
+    k = UOp.range(BLOCK_K // WMMA_K, 101, AxisType.REDUCE)
+    tile_m = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
+    tile_n = UOp.range(TN, 201, AxisType.LOOP)
+
+    acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0,2,1)[tile_m, tile_n]
+    a_frag = A_local.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_K // WMMA_K, WMMA_K)[wave_m, tile_m, lane_n, k]
+    b_frag = B_local.reshape(WAVES_N, TN, WMMA_N, BLOCK_K // WMMA_K, WMMA_K)[wave_n, tile_n, lane_n, k]
+    if is_rdna4:
+      # NOTE: since this is part of K, these 2 can be anywhere in the frags and long as a and b match
+      a_frag = a_frag.reshape(2, 8)[lane_m, :]
+      b_frag = b_frag.reshape(2, 8)[lane_m, :]
+    wmma = UOp(Ops.SHAPED_WMMA, dtypes.float, (a_frag, b_frag, acc_frag.after(k)), arg=((16, 16, 16), 'AMD', 32))
+    acc_store = acc_frag.store(wmma).end(tile_m, tile_n)
+  else:
+    # registers for LOCAL -> REG
+    a_frag = UOp.placeholder((TM//UNROLL_M, UNROLL_M), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+    b_frag = UOp.placeholder((TN//UNROLL_N, UNROLL_N), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+
+    k = UOp.range(BLOCK_K, 101, AxisType.REDUCE)
+    a_frag = a_frag.after(a_frag.store(A_local[k].reshape(WAVES_M, TM//UNROLL_M, LANES_PER_WAVE_M, UNROLL_M)[wave_m, :, lane_m, :]))
+    b_frag = b_frag.after(b_frag.store(B_local[k].reshape(WAVES_N, TN//UNROLL_N, LANES_PER_WAVE_N, UNROLL_N)[wave_n, :, lane_n, :]))
+
+    # FMA
+    a_frag = a_frag.reshape(TM, 1).expand(TM, TN)
+    b_frag = b_frag.reshape(1, TN).expand(TM, TN)
+    acc_store = acc.store(acc.after(k) + (a_frag * b_frag))
+
+  # store accumulator and loop
+  acc = acc.after(acc_store.end(k).barrier().end(k_tile))
+
+  # store accumulator to output (unified)
+  c = c.reshape(WAVES_M, TM//UNROLL_M, LANES_PER_WAVE_M, UNROLL_M,
+                WAVES_N, TN//UNROLL_N, LANES_PER_WAVE_N, UNROLL_N)
+  c = c.permute((0,4,2,6, 1,3,5,7)).reshape(THREADS_PER_BLOCK, TM, TN)
+  return c[tid].store(acc).end(wave_m, wave_n, lane)
+
+def amd_copy_matmul(c:UOp, a:UOp, b:UOp) -> UOp:
+  block_id_m = UOp.range(M // BLOCK_M, 0, AxisType.GLOBAL)
+  block_id_n = UOp.range(N // BLOCK_N, 1, AxisType.GLOBAL)
+  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[block_id_m, :, block_id_n, :]
+  a = a.T.reshape(K, M // BLOCK_M, BLOCK_M)[:, block_id_m, :]
+  b = b.reshape(K, N // BLOCK_N, BLOCK_N)[:, block_id_n, :]
+  return block_128x128_gemm(c, a, b).end(block_id_n, block_id_m).sink(arg=KernelInfo(opts_to_apply=()))
+
+if __name__ == "__main__":
+  from amd_uop_matmul import eval_custom_matmul
+  eval_custom_matmul(amd_copy_matmul, dtypes.half if use_wmma else dtypes.float)

tinygrad_repo/extra/gemm/amd_flash_attention.py

@@ -0,0 +1,205 @@
+from tinygrad import Tensor, UOp, getenv
+from tinygrad.uop.ops import AxisType, KernelInfo, Ops
+from tinygrad.dtype import AddrSpace, dtypes
+from tinygrad.helpers import DEBUG, GlobalCounters, Context
+import math
+
+BLOCK_M, BLOCK_N = 64, 64
+WARP_SIZE = 32
+WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
+WAVES_M, WAVES_N = 4, 1
+LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 16
+WMMA_ACC = WMMA_M // LANES_PER_WAVE_M
+THREADS_PER_BLOCK = WARP_SIZE * WAVES_M * WAVES_N
+LDS_PAD = 4  # pad LDS rows to reduce bank conflicts
+
+WMMA_ARG = ((WMMA_M, WMMA_N, WMMA_K), 'AMD', 32)
+LOG2E = math.log2(math.e)
+
+def warp_shfl_xor(val, offset, lane):
+  """Read val from lane ^ offset using ds_bpermute."""
+  idx = ((lane ^ offset) * 4).cast(dtypes.int)
+  return UOp(Ops.CUSTOM, dtypes.float, (idx, val),
+             arg="__builtin_bit_cast(float, __builtin_amdgcn_ds_bpermute({0}, __builtin_bit_cast(int, {1})))")
+
+def warp_reduce_max(val, lane):
+  """Tree reduce MAX across LANES_PER_WAVE_N=16 lanes."""
+  for offset in [8, 4, 2, 1]:
+    val = UOp(Ops.MAX, dtypes.float, (val, warp_shfl_xor(val, offset, lane)))
+  return val
+
+def warp_reduce_sum(val, lane):
+  """Tree reduce SUM across LANES_PER_WAVE_N=16 lanes."""
+  for offset in [8, 4, 2, 1]:
+    val = val + warp_shfl_xor(val, offset, lane)
+  return val
+
+def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
+  # inputs are (B*H, N, D)
+  BH, N, D = q.shape
+  assert N % BLOCK_M == 0 and N % BLOCK_N == 0, f"N={N} must be divisible by BLOCK_M={BLOCK_M} and BLOCK_N={BLOCK_N}"
+  assert D % WMMA_K == 0 and D % LANES_PER_WAVE_N == 0, f"D={D} must be divisible by WMMA_K={WMMA_K} and LANES_PER_WAVE_N={LANES_PER_WAVE_N}"
+  assert BLOCK_M % (WAVES_M * WMMA_M) == 0 and BLOCK_N % LANES_PER_WAVE_N == 0
+  TM = BLOCK_M // (WAVES_M * LANES_PER_WAVE_M)
+  TN = BLOCK_N // (WAVES_N * LANES_PER_WAVE_N)
+  TD = D // (WAVES_N * LANES_PER_WAVE_N)
+  SCALE = 1.0 / math.sqrt(D)
+
+  block_bh = UOp.range(BH, 0, AxisType.GLOBAL)
+  block_m = UOp.range(N // BLOCK_M, 1, AxisType.GLOBAL)
+
+  q = q.reshape(BH, N//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
+  k = k.reshape(BH, N//BLOCK_N, BLOCK_N, D)[block_bh]
+  v = v.reshape(BH, N//BLOCK_N, BLOCK_N, D)[block_bh]
+  o = o.reshape(BH, N//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
+
+  wave_m = UOp.range(WAVES_M, 2, AxisType.LOCAL)
+  wave_n = UOp.range(WAVES_N, 3, AxisType.LOCAL)
+  lane = UOp.range(WARP_SIZE, -1, AxisType.WARP)
+  tid = (wave_m * WAVES_N + wave_n) * WARP_SIZE + lane
+  lane_m = lane // LANES_PER_WAVE_N
+  lane_n = lane % LANES_PER_WAVE_N
+
+  # LDS allocation: slot 0 = Q then P (shared), slot 1 = K then V
+  # TODO: the memory planner should be able to find this reuse
+  ELEMS_PER_THREAD = BLOCK_M * D // THREADS_PER_BLOCK
+  QP_lds = UOp.placeholder((BLOCK_M, D + LDS_PAD), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)
+  KV_lds = UOp.placeholder((BLOCK_N, D + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :D]
+
+  # register state
+  acc = UOp.placeholder((TM, TD), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  m_i = UOp.placeholder((TM,), dtypes.float, slot=3, addrspace=AddrSpace.REG)
+  l_i = UOp.placeholder((TM,), dtypes.float, slot=4, addrspace=AddrSpace.REG)
+  acc = acc.after(acc.store(acc.const_like(0)))
+  m_i = m_i.after(m_i.store(m_i.const_like(-math.inf)))
+  l_i = l_i.after(l_i.store(l_i.const_like(0)))
+
+  # ====== KV tile loop ======
+  n_tile = UOp.range(N // BLOCK_N, 100, AxisType.REDUCE)
+
+  # load Q + K into LDS (Q reloaded each iteration since P overwrites slot 0)
+  Q_lds = QP_lds[:, :D]
+  Q_store = Q_lds.after(n_tile).reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
+    q.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
+  K_store = KV_lds.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
+    k[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
+  qk_load_barrier = UOp.barrier(UOp.group(Q_store, K_store))
+  Q_lds = Q_lds.after(qk_load_barrier)
+  KV_lds_k = KV_lds.after(qk_load_barrier)
+
+  # -- S = Q @ K^T via WMMA (re-init each n_tile) --
+  S_reg = UOp.placeholder((TM, TN), dtypes.float, slot=6, addrspace=AddrSpace.REG)
+  S_reg = S_reg.after(S_reg.after(n_tile).store(S_reg.const_like(0)))
+  k_qk = UOp.range(D // WMMA_K, 101, AxisType.REDUCE)
+  tm1 = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
+  tn1 = UOp.range(TN, 201, AxisType.LOOP)
+  S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
+  q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
+  k_frag = KV_lds_k.reshape(WAVES_N, TN, WMMA_N, D // WMMA_K, WMMA_K)[wave_n, tn1, lane_n, k_qk]
+  qk = UOp(Ops.SHAPED_WMMA, dtypes.float, (q_frag, k_frag, S_frag.after(k_qk)), arg=WMMA_ARG)
+  qk_done = S_frag.store(qk).end(tm1, tn1).end(k_qk)
+  S_reg = S_reg.after(qk_done)
+
+  # -- softmax in registers with warp shuffles --
+  S_reg = S_reg.after(S_reg.store(S_reg * SCALE))
+
+  # per-thread local row max over TN=4 elements, then warp reduce across 16 lanes
+  m_ij = UOp.placeholder((TM,), dtypes.float, slot=7, addrspace=AddrSpace.REG)
+  m_ij = m_ij.after(m_ij.after(n_tile).store(m_ij.const_like(-math.inf)))
+  rm2 = UOp.range(TN, 261, AxisType.REDUCE)
+  m_ij = m_ij.after(m_ij.store(m_ij.after(rm2).maximum(S_reg[:, rm2])).end(rm2))
+  # warp reduce max (in-place)
+  ri_w = UOp.range(TM, 270, AxisType.LOOP)
+  m_ij = m_ij.after(m_ij[ri_w].store(warp_reduce_max(m_ij[ri_w], lane)).end(ri_w))
+
+  # compute P = exp(S - m_ij) in S_reg
+  S_reg = S_reg.after(S_reg.store(((S_reg - m_ij.reshape(TM, 1).expand(TM, TN)) * LOG2E).exp2()))
+
+  p_local = UOp.placeholder((TM,), dtypes.float, slot=8, addrspace=AddrSpace.REG)
+  p_local = p_local.after(p_local.after(n_tile).store(p_local.const_like(0)))
+  rp2 = UOp.range(TN, 291, AxisType.REDUCE)
+  p_local = p_local.after(p_local.store(p_local.after(rp2) + S_reg[:, rp2]).end(rp2))
+  ri_ws = UOp.range(TM, 295, AxisType.LOOP)
+  p_sum = p_local.after(p_local[ri_ws].store(warp_reduce_sum(p_local[ri_ws], lane)).end(ri_ws))
+
+  # write P = exp(S - m_ij) to P_lds (reuses slot 0, Q no longer needed)
+  P_lds = QP_lds[:, :BLOCK_N]
+  P_write = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_ACC, LANES_PER_WAVE_M, WAVES_N, TN, LANES_PER_WAVE_N)
+  P_write = P_write.permute((0, 4, 3, 6, 1, 2, 5)).reshape(THREADS_PER_BLOCK, TM, TN)
+  # TODO: P_write[tid].store(S_reg.cast(dtypes.half)) — shaped store fails due to RESHAPE(DEFINE_LOCAL) surviving linearization
+  rw1 = UOp.range(TM, 296, AxisType.LOOP)
+  rw2 = UOp.range(TN, 297, AxisType.LOOP)
+  P_store = P_write[tid, rw1, rw2].store(S_reg[rw1, rw2].cast(dtypes.half)).end(rw1, rw2)
+
+  # -- online softmax correction --
+  ri4 = UOp.range(TM, 330, AxisType.LOOP)
+  m_new_val = m_i[ri4].maximum(m_ij[ri4])
+  alpha_val = ((m_i[ri4] - m_new_val) * LOG2E).exp2()
+  beta_val = ((m_ij[ri4] - m_new_val) * LOG2E).exp2()
+  rj4 = UOp.range(TD, 331, AxisType.LOOP)
+  correction = UOp.group(
+    acc[ri4, rj4].store(alpha_val * acc[ri4, rj4]).end(rj4),
+    l_i[ri4].store(alpha_val * l_i[ri4] + beta_val * p_sum[ri4]),
+    m_i[ri4].store(m_new_val),
+  ).end(ri4)
+  acc = acc.after(correction)
+  l_i = l_i.after(correction)
+  m_i = m_i.after(correction)
+
+  # load V into KV_lds (must wait for QK WMMA to finish reading K from KV_lds)
+  V_store = KV_lds.after(qk_done).reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
+    v[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
+  pv_barrier = UOp.barrier(UOp.group(P_store, V_store))
+  P_lds = P_lds.after(pv_barrier)
+  KV_lds_v = KV_lds.after(pv_barrier)
+
+  # -- acc += P @ V via WMMA --
+  k_pv = UOp.range(BLOCK_N // WMMA_K, 400, AxisType.REDUCE)
+  tm2 = UOp.range(TM // WMMA_ACC, 401, AxisType.LOOP)
+  tn2 = UOp.range(TD, 402, AxisType.LOOP)
+  acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2]
+  p_frag = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv]
+  v_frag = KV_lds_v.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv]
+  pv = UOp(Ops.SHAPED_WMMA, dtypes.float, (p_frag, v_frag, acc_frag.after(k_pv)), arg=WMMA_ARG)
+
+  # end KV tile loop
+  n_tile_end = acc_frag.store(pv).end(tm2, tn2).end(k_pv).barrier().end(n_tile)
+  acc = acc.after(n_tile_end)
+  l_i = l_i.after(n_tile_end)
+  m_i = m_i.after(n_tile_end)
+
+  # normalize: acc /= l_i
+  acc = acc.after(acc.store(acc * (1 / l_i).reshape(TM, 1).expand(TM, TD)))
+
+  # store output
+  o = o.reshape(WAVES_M, TM // WMMA_ACC, WMMA_ACC, LANES_PER_WAVE_M, WAVES_N, TD, LANES_PER_WAVE_N)
+  o = o.permute((0, 4, 3, 6, 1, 2, 5)).reshape(THREADS_PER_BLOCK, TM, TD)
+  return o[tid].store(acc).end(wave_m, wave_n, lane).end(block_m, block_bh).sink(arg=KernelInfo(opts_to_apply=()))
+
+if __name__ == "__main__":
+  B, H, N, D = getenv("B", 1), getenv("H", 32), getenv("N", 1024), getenv("D", 64)
+  q = Tensor.rand(B, H, N, D).cast(dtypes.half)
+  k = Tensor.rand(B, H, N, D).cast(dtypes.half)
+  v = Tensor.rand(B, H, N, D).cast(dtypes.half)
+  o = Tensor.empty(B, H, N, D, dtype=dtypes.float)
+  with Context(DEBUG=0): Tensor.realize(q, k, v)
+
+  q_flat, k_flat, v_flat, o_flat = q.reshape(B*H, N, D), k.reshape(B*H, N, D), v.reshape(B*H, N, D), o.reshape(B*H, N, D)
+  NUM_RUNS = getenv("CNT", 5)
+  ets = []
+  with Context(DEBUG=2):
+    for _ in range(NUM_RUNS):
+      GlobalCounters.reset()
+      tst = Tensor.custom_kernel(o_flat, q_flat, k_flat, v_flat, fxn=amd_flash_attention)[0].realize()
+      ets.append(GlobalCounters.time_sum_s)
+  print(f"best time: {min(ets)*1e3:.2f}ms")
+
+  if getenv("VERIFY", 1):
+    with Context(DEBUG=0):
+      ref = q.float().scaled_dot_product_attention(k.float(), v.float()).reshape(B*H, N, D).realize()
+      err = (ref - tst).square().mean().item()
+    print(f"mean squared error {err}")
+    if err > 1e-2:
+      raise RuntimeError("flash attention is wrong!")
+    else:
+      print("flash attention is correct!")

tinygrad_repo/extra/gemm/amd_matmul.py

@@ -0,0 +1,50 @@
+# kernel8_batched_gmem.s from https://seb-v.github.io/optimization/update/2025/01/20/Fast-GPU-Matrix-multiplication.html
+# sudo PATH=/opt/homebrew/Cellar/llvm/20.1.6/bin:$PATH AMD_LLVM=0 AMD=1 DEBUG=2 python3 extra/gemm/amd_matmul.py
+import pathlib
+from tinygrad import Tensor, Device, Context, GlobalCounters
+from tinygrad.helpers import getenv
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from tinygrad.engine.realize import run_linear
+
+N = 4096
+run_count = 5
+
+def make_matmul_kernel(name:str, src:str, local_size:int):
+  def fxn(a:UOp, b:UOp, c:UOp) -> UOp:
+    threads = UOp.special(local_size, "lidx0")
+    wg_x = UOp.special(N//128, "gidx0")
+    wg_y = UOp.special(N//128, "gidx1")
+    sink = UOp.sink(a.base, b.base, c.base, threads, wg_x, wg_y, arg=KernelInfo(name, estimates=Estimates(ops=2*N**3, mem=3*N*N*4)))
+    lib = Device[Device.DEFAULT].compiler.compile_cached(src)
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                                 UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
+  return fxn
+
+if __name__ == "__main__":
+  if getenv("ASM") == 1:
+    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel8_batched_gmem.s").read_text()
+    name, local_size = "kernel", 128
+  elif getenv("ASM") == -1:
+    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel3_registers.cpp").read_text()
+    name, local_size = "kernel3_registers", 256
+  elif getenv("ASM") == -2:
+    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel4_gmem_df.cpp").read_text()
+    name, local_size = "kernel4_gmem_db", 256
+  else:
+    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel5_lds_optim.cpp").read_text()
+    name, local_size = "kernel5_lds_optim", 128
+
+  a = Tensor.randn(N, N).realize()
+  b = Tensor.randn(N, N).realize()
+  c = Tensor.zeros(N, N).contiguous().realize()
+
+  GlobalCounters.reset()
+  with Context(DEBUG=2):
+    for _ in range(run_count): tc = (a@b).realize()
+
+  linear = Tensor.custom_kernel(a, b, c, fxn=make_matmul_kernel(name, src, local_size))[2].schedule_linear()
+  GlobalCounters.reset()
+  with Context(DEBUG=2):
+    for _ in range(run_count): run_linear(linear)
+  print(f"custom  {(c-tc).square().mean().item()}")

tinygrad_repo/extra/gemm/amd_uop_matmul.py

@@ -0,0 +1,143 @@
+from tinygrad import Tensor, Context, GlobalCounters, dtypes
+from tinygrad.uop.ops import UOp, KernelInfo, sint, AxisType
+from tinygrad.dtype import AddrSpace
+from tinygrad.helpers import DEBUG, getenv
+
+N = getenv("N", 4096)
+M = getenv("M", N)
+K = getenv("K", N)
+NUM_RUNS = getenv("CNT", 5)
+
+# ---------------------------
+# launch/config constants
+# ---------------------------
+
+WARP_SIZE = 32
+BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 8
+TM, TN = 4, 4
+LANES_PER_WAVE_M, LANES_PER_WAVE_N = 4, 8
+assert N % BLOCK_N == 0 and M % BLOCK_M == 0 and K % BLOCK_K == 0
+
+is_kernel5 = getenv("K5", 0)
+THREADS_PER_BLOCK = 128 if is_kernel5 else 256
+WAVES_PER_BLOCK_N = 1 if is_kernel5 else 2
+WAVES_PER_BLOCK_M = THREADS_PER_BLOCK // WARP_SIZE // WAVES_PER_BLOCK_N
+REG_TILES_PER_WAVE_N = BLOCK_N // (WAVES_PER_BLOCK_N * LANES_PER_WAVE_N * TN)
+REG_TILES_PER_WAVE_M = BLOCK_M // (WAVES_PER_BLOCK_M * LANES_PER_WAVE_M * TM)
+
+assert WAVES_PER_BLOCK_M*REG_TILES_PER_WAVE_M*LANES_PER_WAVE_M*TM == BLOCK_M, "M reshape is wrong"
+assert WAVES_PER_BLOCK_N*REG_TILES_PER_WAVE_N*LANES_PER_WAVE_N*TN == BLOCK_N, "N reshape is wrong"
+
+def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.LOOP): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
+def copy(dest:UOp, src:UOp, rng:int, upcast=False):
+  assert dest.shape == src.shape
+  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
+  return dest[*rngs].store(src[*rngs]).end(*rngs)
+
+def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
+  # ---------------------------
+  # block indices
+  # ---------------------------
+  block_id_n = UOp.special(N // BLOCK_N, "gidx0")
+  block_id_m = UOp.special(M // BLOCK_M, "gidx1")
+
+  # index the output with the globals
+  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[block_id_m, :, block_id_n, :]
+
+  # open the main reduction range
+  k_tile_range = UOp.range(K // BLOCK_K, 0, AxisType.REDUCE)
+  a = a.reshape(M // BLOCK_M, BLOCK_M, K // BLOCK_K, BLOCK_K)[block_id_m, :, k_tile_range, :]
+  b = b.reshape(K // BLOCK_K, BLOCK_K, N // BLOCK_N, BLOCK_N)[k_tile_range, :, block_id_n, :]
+
+  # globals are no longer used, they are already in the indexes
+  del block_id_m, block_id_n
+
+  # ---------------------------
+  # GLOBAL -> LOCAL (A_local, B_local)
+  # ---------------------------
+  tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
+
+  # A: read BM x BK tiles (permute on store into locals)
+  BM_A_local_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
+  A_local = UOp.placeholder((BLOCK_K, BM_A_local_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
+  A_local_store = copy(A_local.permute((1,0)).reshape(-1, THREADS_PER_BLOCK)[:, tid], a.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=100)
+
+  # B: read BK x BN tiles
+  B_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
+  B_local_store = copy(B_local.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
+
+  # TODO: can we automate barrier?
+  barrier = UOp.barrier(A_local_store, B_local_store)
+  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
+
+  # open inner k range
+  k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
+
+  # ---------------------------
+  # LOCAL -> REG (per-wave tiles)
+  # ---------------------------
+  warp, lane = tid // WARP_SIZE, tid % WARP_SIZE
+  waveIdx, waveIdy = warp % WAVES_PER_BLOCK_N, warp // WAVES_PER_BLOCK_N
+  laneIdx, laneIdy = lane % LANES_PER_WAVE_N, lane // LANES_PER_WAVE_N
+  assert waveIdy.vmax+1 == WAVES_PER_BLOCK_M and laneIdy.vmax+1 == LANES_PER_WAVE_M
+
+  A_col = UOp.placeholder((REG_TILES_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  A_local_slice = A_local[k, :].reshape(WAVES_PER_BLOCK_M, REG_TILES_PER_WAVE_M, LANES_PER_WAVE_M, TM)[waveIdy, :, laneIdy, :]
+  A_col = A_col.after(copy(A_col, A_local_slice, 300, upcast=True))
+
+  B_row = UOp.placeholder((REG_TILES_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+  B_local_slice = B_local[k, :].reshape(WAVES_PER_BLOCK_N, REG_TILES_PER_WAVE_N, LANES_PER_WAVE_N, TN)[waveIdx, :, laneIdx, :]
+  B_row = B_row.after(copy(B_row, B_local_slice, 400, upcast=True))
+
+  # ---------------------------
+  # FMA: c_regs += A_col * B_row
+  # ---------------------------
+  c_regs = UOp.placeholder((REG_TILES_PER_WAVE_M, TM, REG_TILES_PER_WAVE_N, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  i = UOp.range(c_regs.size, 16)
+  c_regs = c_regs.after(c_regs.flatten()[i].store(0.0).end(i))
+
+  # TODO: why don't these work as upcast?
+  # why if the ranges merge is it slow?!? (if you change the order on end, they will merge. big slowdown on METAL)
+  iter_m, t_m, iter_n, t_n = rngs = rngs_for_shape(c_regs.shape, 500)
+  sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iter_m, t_m] * B_row[iter_n, t_n]).end(iter_m, iter_n, t_m, t_n)
+
+  # Close k, sync, and close K tiles
+  sink = sink.end(k).barrier().end(k_tile_range)
+
+  # ---------------------------
+  # REG -> GLOBAL (epilogue)
+  # ---------------------------
+  c = c.reshape(WAVES_PER_BLOCK_M, REG_TILES_PER_WAVE_M, LANES_PER_WAVE_M, TM,
+                WAVES_PER_BLOCK_N, REG_TILES_PER_WAVE_N, LANES_PER_WAVE_N, TN)
+  c = c[waveIdy, :, laneIdy, :,
+        waveIdx, :, laneIdx, :]
+  sink = copy(c, c_regs.after(sink), rng=600)
+
+  return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()
+
+def eval_custom_matmul(fxn, dt=dtypes.float):
+  a = Tensor.randn(M, K, dtype=dt)
+  b = Tensor.randn(K, N, dtype=dt)
+  c = Tensor.empty(M, N, dtype=dtypes.float)
+  with Context(DEBUG=0): Tensor.realize(a, b)
+
+  ets = []
+  with Context(DEBUG=max(2, DEBUG.value)):
+    for _ in range(NUM_RUNS):
+      GlobalCounters.reset()
+      tst = Tensor.custom_kernel(c, a, b, fxn=fxn)[0].realize()
+      ets.append(GlobalCounters.time_sum_s)
+  print(f"REAL TFLOPS {M * N * K * 2 / min(ets) * 1e-12:.2f}")
+
+  if getenv("VERIFY", 1):
+    GlobalCounters.reset()
+    with Context(DEBUG=2):
+      tc = (a.float() @ b.float()).realize()
+    with Context(DEBUG=0):
+      err = (tc - tst).square().mean().item()
+    print(f"mean squared error {err}")
+    if err > (1e-2 if dt == dtypes.half else 1e-6):
+      raise RuntimeError("matmul is wrong!")
+
+if __name__ == "__main__":
+  eval_custom_matmul(hand_spec_kernel3)

tinygrad_repo/extra/gemm/amx.py

@@ -0,0 +1,180 @@
+#!/usr/bin/env python3
+import numpy as np
+import time
+import sys
+np.set_printoptions(linewidth=160)
+np.set_printoptions(linewidth=1000, threshold=10000000000, suppress=False)
+from tinygrad.runtime.ops_llvm import LLVMDevice, LLVMProgram, LLVMCompiler
+from llvmlite import ir  # type: ignore
+from tinygrad.helpers import flat_mv
+from tinygrad.device import MallocAllocator
+
+# https://github.com/corsix/amx/blob/main/Instructions.md
+# 12 lines for AMX support
+from functools import partialmethod
+class AMX:
+  @staticmethod
+  def nop_op_imm5(op, imm5, builder): builder.asm(ir.FunctionType(ir.VoidType(), []), f".word (0x201000 + ({op} << 5) + {imm5}); amx op {op} imm {imm5}", "", tuple(), True)
+  @staticmethod
+  def op_gpr(op, builder, gpr): builder.asm(ir.FunctionType(ir.VoidType(), [ir.IntType(64)]), f".word (0x201000 + ({op} << 5) + 0$0 - ((0$0 >> 4) * 6)); amx op {op} reg $0", "r", (gpr,), True)
+  set, clr = partialmethod(nop_op_imm5, 17, 0), partialmethod(nop_op_imm5, 17, 1)
+  ldx, ldy, stx, sty = partialmethod(op_gpr, 0), partialmethod(op_gpr, 1), partialmethod(op_gpr, 2), partialmethod(op_gpr, 3)
+  ldz, stz, ldzi, stzi = partialmethod(op_gpr, 4), partialmethod(op_gpr, 5), partialmethod(op_gpr, 6), partialmethod(op_gpr, 7)
+  extrx, extry = partialmethod(op_gpr, 8), partialmethod(op_gpr, 9)
+  fma64, fms64, fma32, fms32 = partialmethod(op_gpr, 10), partialmethod(op_gpr, 11), partialmethod(op_gpr, 12), partialmethod(op_gpr, 13)
+  mac16, fma16, fms16 = partialmethod(op_gpr, 14), partialmethod(op_gpr, 15), partialmethod(op_gpr, 16)
+  vecint, vecfp, matint, matfp, genlut = partialmethod(op_gpr, 18), partialmethod(op_gpr, 19), partialmethod(op_gpr, 20), partialmethod(op_gpr, 21), partialmethod(op_gpr, 22)
+
+def int_const(x): return ir.Constant(ir.IntType(64), x)
+
+
+N = 4096
+# N = 1024
+# N = 64
+
+BW = N*N*4
+
+# matrix is 64M, max load bandwidth is 57 GB/s
+# cache line looks like 256 bytes (64 floats)
+
+na = np.zeros((256), dtype=np.float32)
+# na = np.zeros((N, N), dtype=np.float32)
+nb = np.random.randn(N, N).astype(np.float32)
+nc = np.random.randn(N, N).astype(np.float32)
+
+ns = nb.reshape(-1, 32).sum(axis=0)
+
+a = MallocAllocator.alloc(na.nbytes)
+b = MallocAllocator.alloc(nb.nbytes)
+c = MallocAllocator.alloc(nc.nbytes)
+
+MallocAllocator._copyin(b, flat_mv(nb.data))
+MallocAllocator._copyin(c, flat_mv(nc.data))
+
+module = ir.Module(name=__file__)
+func = ir.Function(module, ir.FunctionType(ir.IntType(64), [ir.FloatType().as_pointer()]*3), name='exec')
+
+# load all
+entry = ir.IRBuilder(func.append_basic_block(name="entry"))
+zm, xm, ym = [entry.ptrtoint(func.args[i], ir.IntType(64)) for i in range(3)]
+
+loop_1 = ir.IRBuilder(func.append_basic_block(name="loop_y"))
+loop_1_exit = ir.IRBuilder(func.append_basic_block(name="loop_y_exit"))
+exit = ir.IRBuilder(func.append_basic_block(name="exit"))
+
+y = loop_1.phi(ir.IntType(64), name="y")
+y.add_incoming(int_const(0), entry._block)
+yp = loop_1_exit.add(y, int_const(32*2))
+y.add_incoming(yp, loop_1_exit._block)
+
+prefetch_function = ir.Function(module, ir.FunctionType(ir.VoidType(), [ir.PointerType(ir.FloatType()), ir.IntType(32), ir.IntType(32), ir.IntType(32)]), name="llvm.prefetch")
+
+xptr = y
+addr = loop_1_exit.add(xm, loop_1_exit.mul(int_const(4), xptr))
+
+#prefetch_ptr = loop_1_exit.inttoptr(loop_1_exit.add(addr, int_const(128)), ir.PointerType(ir.FloatType()))
+#loop_1_exit.call(prefetch_function, [prefetch_ptr, ir.IntType(32)(0), ir.IntType(32)(2), ir.IntType(32)(1)])
+
+AMX.ldx(loop_1_exit, loop_1_exit.add(int_const(1<<62), addr))
+xptr = loop_1_exit.add(xptr, int_const(32))
+AMX.ldy(loop_1_exit, loop_1_exit.add(int_const(1<<62), loop_1_exit.add(xm, loop_1_exit.mul(int_const(4), xptr))))
+
+AMX.fma32(loop_1_exit, int_const(1 << 63 | 1 << 28))
+AMX.fma32(loop_1_exit, int_const(1 << 63 | 1 << 28 | 1 << 20 | (16*4)<<10))
+AMX.fma32(loop_1_exit, int_const(1 << 63 | 1 << 29))
+AMX.fma32(loop_1_exit, int_const(1 << 63 | 1 << 29 | 1 << 20 | (16*4)))
+
+AMX.set(entry)
+
+AMX.stz(exit, exit.add(zm, int_const(1 << 62 | (0 << 56) | 0)))
+AMX.clr(exit)
+
+entry.branch(loop_1._block)
+loop_1.branch(loop_1_exit._block)
+loop_1_exit.cbranch(loop_1_exit.icmp_unsigned("==", yp, int_const(N*N)), exit._block, loop_1._block)
+exit.ret(int_const(0))
+
+device = LLVMDevice("llvm")
+prog = LLVMProgram(device, "exec", LLVMCompiler(device).compile(str(module)))
+
+"""
+loop_1 = ir.IRBuilder(func.append_basic_block(name="loop_y"))
+loop_2 = ir.IRBuilder(func.append_basic_block(name="loop_x"))
+loop_3 = ir.IRBuilder(func.append_basic_block(name="loop_k"))
+loop_3_exit = ir.IRBuilder(func.append_basic_block(name="loop_k_exit"))
+loop_2_exit = ir.IRBuilder(func.append_basic_block(name="loop_x_exit"))
+loop_1_exit = ir.IRBuilder(func.append_basic_block(name="loop_y_exit"))
+
+y = loop_1.phi(ir.IntType(64), name="y")
+x = loop_2.phi(ir.IntType(64), name="x")
+k = loop_3.phi(ir.IntType(64), name="k")
+
+exit = ir.IRBuilder(func.append_basic_block(name="exit"))
+
+AMX.set(loop_2)
+
+# stride
+xptr = loop_3_exit.add(x, loop_3_exit.mul(k, int_const(N)))
+yptr = loop_3_exit.add(y, loop_3_exit.mul(k, int_const(N)))
+
+# if you are okay with the wrong answer, this is faster
+#xptr = loop_3_exit.add(x, loop_3_exit.mul(k, int_const(32)))
+#yptr = loop_3_exit.add(y, loop_3_exit.mul(k, int_const(32)))
+
+# double loads load 32 floats
+AMX.ldx(loop_3_exit, loop_3_exit.add(int_const(1<<62), loop_3_exit.add(xm, loop_3_exit.mul(int_const(4), xptr))))
+AMX.ldy(loop_3_exit, loop_3_exit.add(int_const(1<<62), loop_3_exit.add(ym, loop_3_exit.mul(int_const(4), yptr))))
+
+# <Z row> <X offset> <Y offset>
+AMX.fma32(loop_3_exit, int_const(0<<20 | (0*16*4)<<10 | (0*16*4)))
+AMX.fma32(loop_3_exit, int_const(1<<20 | (1*16*4)<<10 | (0*16*4)))
+AMX.fma32(loop_3_exit, int_const(2<<20 | (0*16*4)<<10 | (1*16*4)))
+AMX.fma32(loop_3_exit, int_const(3<<20 | (1*16*4)<<10 | (1*16*4)))
+
+# store
+gptr = loop_2_exit.mul(loop_2_exit.add(loop_2.mul(y, int_const(N)), x), int_const(4))
+zmp = loop_2_exit.add(zm, gptr)
+for j in range(2):
+  for r in range(16):
+    z_row = j*2
+    ptr = ((j*16)+r)*N
+    AMX.stz(loop_2_exit, loop_2_exit.add(zmp, int_const(1 << 62 | ((r*4+z_row) << 56) | ptr*4)))
+AMX.clr(loop_2_exit)
+
+yp = loop_1_exit.add(y, int_const(32))
+xp = loop_2_exit.add(x, int_const(32))
+kp = loop_3_exit.add(k, int_const(1))
+
+y.add_incoming(int_const(0), entry._block)
+x.add_incoming(int_const(0), loop_1._block)
+k.add_incoming(int_const(0), loop_2._block)
+y.add_incoming(yp, loop_1_exit._block)
+x.add_incoming(xp, loop_2_exit._block)
+k.add_incoming(kp, loop_3_exit._block)
+
+entry.branch(loop_1._block)
+loop_1.branch(loop_2._block)
+loop_2.branch(loop_3._block)
+loop_3.branch(loop_3_exit._block)
+loop_3_exit.cbranch(loop_3_exit.icmp_unsigned("==", kp, int_const(N)), loop_2_exit._block, loop_3._block)
+loop_2_exit.cbranch(loop_2_exit.icmp_unsigned("==", xp, int_const(N)), loop_1_exit._block, loop_2._block)
+loop_1_exit.cbranch(loop_1_exit.icmp_unsigned("==", yp, int_const(N)), exit._block, loop_1._block)
+exit.ret(int_const(0))
+
+device = LLVMDevice("llvm")
+prog = LLVMProgram(device, "exec", LLVMCompiler(device).compile(str(module)))
+"""
+
+def timeit(fxn):
+  st = time.perf_counter()
+  et = fxn()
+  return time.perf_counter() - st
+
+tm = min([timeit(lambda: prog(a, b, c, N**2)) for _ in range(20)])
+MallocAllocator._copyout(flat_mv(na.data), a)
+print(f"{N*N:10d} {tm*1e6:9.2f} us, {BW*1e-9/tm:.2f} GB/s")
+
+np.testing.assert_allclose(na[:ns.shape[0]], ns, atol=1e-4, rtol=1e-4)
+
+# comp = (nb.T @ nc).T
+# np.testing.assert_allclose(na, comp, atol=1e-4, rtol=1e-5)

tinygrad_repo/extra/gemm/cuda_matmul.py

@@ -0,0 +1,107 @@
+import os
+import numpy as np
+os.environ["CUDA"] = "1"
+from tinygrad.runtime.ops_cuda import CUDAAllocator, CUDADevice, CUDAProgram, CUDACompiler
+from tinygrad.helpers import flat_mv
+
+FLOAT16 = True
+ACC_FLOAT16 = False
+N = 4096
+
+na = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32)
+nb = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32)
+nc = np.empty(N*N, np.float32)
+
+if FLOAT16:
+  na = na.astype(np.float16)
+  nb = nb.astype(np.float16)
+
+device = CUDADevice("cuda:0")
+cudaalloc = CUDAAllocator(device)
+
+a = cudaalloc.alloc(N*N*2 if FLOAT16 else N*N*4)
+b = cudaalloc.alloc(N*N*2 if FLOAT16 else N*N*4)
+c = cudaalloc.alloc(N*N*4)
+
+cudaalloc._copyin(a, bytearray(na))
+cudaalloc._copyin(b, bytearray(nb))
+
+FLOPS = N*N*N*2
+BW = N*N*3*4
+
+print(device.arch)
+compiler = CUDACompiler(device.arch)
+prog = CUDAProgram(device, "wmma_example", compiler.compile(f"""
+#include <mma.h>
+using namespace nvcuda;
+
+const int WMMA_M = 16;
+const int WMMA_N = 16;
+const int WMMA_K = {'16' if FLOAT16 else '8'};
+
+extern "C" __global__ void wmma_example({'half' if FLOAT16 else 'float'} *a, {'half' if FLOAT16 else 'float'} *b, float *c)
+{{
+  int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;
+  int warpN = (blockIdx.y * blockDim.y + threadIdx.y);
+  warpM *= 4;
+  warpN *= 4;
+
+  wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, {'half' if FLOAT16 else 'wmma::precision::tf32'}, wmma::col_major> a_frag[4];
+  wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, {'half' if FLOAT16 else 'wmma::precision::tf32'}, wmma::col_major> b_frag[4];
+  wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, {'half' if ACC_FLOAT16 else 'float'}> acc_frag[4][4];
+  for (int j = 0; j < 4; j++) {{
+    for (int i = 0; i < 4; i++) {{
+      wmma::fill_fragment(acc_frag[i][j], 0.0f);
+    }}
+  }}
+
+  for (int k = 0; k < {N}; k += WMMA_K) {{
+    int aRow = warpM * WMMA_M;
+    int aCol = k;
+    int bRow = k;
+    int bCol = warpN * WMMA_N;
+
+    wmma::load_matrix_sync(a_frag[0], a + aRow + 0 * WMMA_M + aCol * {N}, {N});
+    wmma::load_matrix_sync(a_frag[1], a + aRow + 1 * WMMA_M + aCol * {N}, {N});
+    wmma::load_matrix_sync(a_frag[2], a + aRow + 2 * WMMA_M + aCol * {N}, {N});
+    wmma::load_matrix_sync(a_frag[3], a + aRow + 3 * WMMA_M + aCol * {N}, {N});
+
+    wmma::load_matrix_sync(b_frag[0], b + bRow + (0 * WMMA_N + bCol) * {N}, {N});
+    wmma::load_matrix_sync(b_frag[1], b + bRow + (1 * WMMA_N + bCol) * {N}, {N});
+    wmma::load_matrix_sync(b_frag[2], b + bRow + (2 * WMMA_N + bCol) * {N}, {N});
+    wmma::load_matrix_sync(b_frag[3], b + bRow + (3 * WMMA_N + bCol) * {N}, {N});
+
+    #pragma unroll
+    for (int i = 0; i < {'0' if FLOAT16 else '4'}; i++) {{
+      #pragma unroll
+      for (int t = 0; t < a_frag[i].num_elements; t++) {{ a_frag[i].x[t] =  wmma::__float_to_tf32(a_frag[i].x[t]); }}
+      #pragma unroll
+      for (int t = 0; t < b_frag[i].num_elements; t++) {{ b_frag[i].x[t] =  wmma::__float_to_tf32(b_frag[i].x[t]); }}
+    }}
+
+    #pragma unroll
+    for (int j = 0; j < 4; j++) {{
+      #pragma unroll
+      for (int i = 0; i < 4; i++) {{
+        wmma::mma_sync(acc_frag[i][j], a_frag[i], b_frag[j], acc_frag[i][j]);
+      }}
+    }}
+  }}
+
+  for (int j = 0; j < 4; j++) {{
+    for (int i = 0; i < 4; i++) {{
+      wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float> acc_store;
+      for (int t = 0; t < acc_frag[i][j].num_elements; t++) acc_store.x[t] = acc_frag[i][j].x[t];
+      int cRow = (warpM + i) * WMMA_M;
+      int cCol = (warpN + j) * WMMA_N;
+      wmma::store_matrix_sync(c + cRow + cCol * {N}, acc_store, {N}, wmma::mem_col_major);
+    }}
+  }}
+}}
+"""))
+
+global_size, local_size = [(N//16)//4, (N//16)//4, 1], [32, 1, 1]
+tm = min([prog(a, b, c, global_size=global_size, local_size=local_size, wait=True) for _ in range(20)])
+print(f"{N*N:10d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS matmul, {BW*1e-9/tm:.2f} GB/s")
+cudaalloc._copyout(flat_mv(nc.data), c)
+np.testing.assert_allclose(na.T.astype(np.float32) @ nb.T.astype(np.float32), nc.reshape(N,N).T, atol=1e-2)

tinygrad_repo/extra/gemm/fuzz_matmul.py

@@ -0,0 +1,44 @@
+import numpy as np
+from tinygrad.helpers import getenv
+from tinygrad import dtypes, Tensor
+dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float
+acc_dtype = dtypes.half if getenv("ACC_HALF") else dtypes.bfloat16 if getenv("ACC_BFLOAT16") else None
+N_START = getenv("N_START", 1)
+M_START = getenv("M_START", 1)
+K_START = getenv("K_START", 1)
+N_STOP = getenv("N_STOP", 32)
+M_STOP = getenv("M_STOP", N_STOP)
+K_STOP = getenv("K_STOP", N_STOP)
+N_STEP = getenv("N_STEP", 1)
+M_STEP = getenv("M_STEP", 1)
+K_STEP = getenv("K_STEP", 1)
+ATOL = getenv("ATOL", 1e-4)
+RTOL = getenv("RTOL", 3e-2)
+
+if __name__ == "__main__":
+  failed = []
+  for M in range(M_START, M_STOP+1, M_STEP):
+    for N in range(N_START, N_STOP+1, N_STEP):
+      for K in range(K_START, K_STOP+1, K_STEP):
+        print(f"testing {M=} {N=} {K=}")
+        a, b = Tensor.rand(M, K, dtype=dtype_in).realize(), Tensor.rand(K, N, dtype=dtype_in).realize()
+        c = a.matmul(b, dtype=acc_dtype).realize()
+        comp = a.numpy().astype(np.float32) @ b.numpy().astype(np.float32)
+        nc = c.numpy()
+        try:
+          np.testing.assert_allclose(nc, comp, atol=ATOL, rtol=RTOL)
+        except AssertionError as e:
+          failed.append((M,N,K,))
+          if getenv("DEBUG_VALUES") > 0:
+            indices = np.where(~np.isclose(nc, comp, rtol=RTOL, atol=ATOL))
+            non_matching_elements_nc = nc[indices]
+            non_matching_elements_comp = comp[indices]
+            print(indices)
+            print("result      :", non_matching_elements_nc)
+            print("ground truth:", non_matching_elements_comp)
+          print(e)
+          pass
+  print(f"failed sizes: {failed}")
+  print(f"num failures: {len(failed)}")
+  if len(failed) > 0:
+    raise RuntimeError(f"failed on {len(failed)} kernels")

tinygrad_repo/extra/gemm/gemm.py

@@ -0,0 +1,28 @@
+#!/usr/bin/env python3
+import os
+#os.environ['OMP_NUM_THREADS'] = '1'
+import time
+import numpy as np
+
+N = 512
+if __name__ == "__main__":
+  # N^2
+  A = np.random.randn(N, N).astype(np.float32)
+  # N^2
+  B = np.random.randn(N, N).astype(np.float32)
+
+  # 2N compute in N^2 output cells
+  flop = 2*N*N*N
+  #print(f"{flop / 1e9:.2f} GFLOP")
+
+  for i in range(10):
+    st = time.monotonic()
+    C = A @ B.T
+    et = time.monotonic()
+    s = et-st
+    print(f"{flop/s * 1e-9:.2f} GFLOP/S, {s*1e3:.2f} ms")
+
+  with open("/tmp/matmul", "wb") as f:
+    f.write(A.data)
+    f.write(B.data)
+    f.write(C.data)

tinygrad_repo/extra/gemm/halide_gemm.py

@@ -0,0 +1,90 @@
+import numpy as np
+import halide as hl
+from tinygrad.helpers import Timing, getenv
+
+# HL_DEBUG_CODEGEN=1
+N = getenv("N", 1024)
+
+def gemm_pipeline(gpu=False):
+  # ---------------- Vars & Parameters ----------------
+  i, j = hl.Var("i"), hl.Var("j")  # output tile coordinates
+
+  A = hl.InputBuffer(hl.Float(32), 2)  # [M, K]
+  B = hl.InputBuffer(hl.Float(32), 2)  # [K, N]
+
+  A.dim(0).set_bounds(0, N)
+  A.dim(1).set_bounds(0, N)
+  B.dim(0).set_bounds(0, N)
+  B.dim(1).set_bounds(0, N)
+
+  # ---------------- Definition ----------------
+
+  k = hl.RDom([(0, N)])
+
+  partial = hl.Func("partial")
+  partial[i, j] = 0.0
+  partial[i, j] += A[i, k] * B[k, j]
+
+  C = hl.Func("C")
+  C[i, j] = partial[i, j]
+
+  if not gpu:
+    # ---------------- Schedule ----------------
+    VEC = 16
+    TILE_I = 64
+    TILE_J = 64
+
+    io, jo, ii, ji = hl.Var("io"), hl.Var("jo"), hl.Var("ii"), hl.Var("ji")
+    C.update().tile(i, j, io, jo, ii, ji, TILE_I, TILE_J).fuse(io, jo, io).parallel(io).vectorize(ji, VEC)
+  else:
+    # ---------------- Schedule ----------------
+    GRP_I    = 8     # output tile size
+    GRP_J    = 16
+
+    #partial.store_in(hl.MemoryType.Register)
+    #partial.update().unroll(k, 4)
+
+    io, jo, ii, ji = hl.Var(), hl.Var(), hl.Var(), hl.Var()
+    C.gpu_tile(i, j, io, jo, ii, ji, GRP_I, GRP_J, hl.TailStrategy.RoundUp)
+
+  return C, A, B
+
+if __name__ == "__main__":
+  pipe, A, B = gemm_pipeline(gpu=True)
+
+  # NOTE: meteal does nothing
+  target = hl.get_host_target().with_feature(hl.TargetFeature.Metal)
+
+  a_np = np.random.randn(N, N).astype(np.float32)
+  b_np = np.random.randn(N, N).astype(np.float32)
+
+  # reverse order is correct!
+  a_hal = hl.Buffer(b_np)
+  b_hal = hl.Buffer(a_np)
+  A.set(a_hal)
+  B.set(b_hal)
+
+  pipe.compile_to_lowered_stmt("/tmp/my_function.html", [A, B], hl.StmtOutputFormat.HTML, target=target)
+  #exit(0)
+
+  c_hal = hl.Buffer(hl.Float(32), [N,N])
+  with Timing("halide gemm "):
+    pipe.realize(c_hal, target)
+    c_hal.copy_to_host()
+    c_out = np.array(c_hal)
+  print(c_out)
+
+  # tinygrad gets 60 ms with no BEAM, 20 ms with BEAM on CPU
+  with Timing("halide gemm "):
+    pipe.realize(c_hal, target)
+    c_hal.copy_to_host()
+
+  # Check correctness
+  with Timing("numpy gemm "):
+    ref = a_np @ b_np
+  max_err = np.abs(ref - c_out).max()
+  print("Max absolute error:", max_err)
+  assert max_err < 1e-4, "GEMM result incorrect!"
+
+  print("Pipeline ran on", target)
+  print("Success - GEMM Halide-Python output matches NumPy.")

tinygrad_repo/extra/gemm/hip_matmul.py

@@ -0,0 +1,142 @@
+import time
+import numpy as np
+from tinygrad.helpers import getenv, prod, flat_mv
+from tinygrad.runtime.ops_amd import AMDAllocator, AMDDevice, AMDProgram
+
+# AMD_LOG_LEVEL=3 ./MIOpenDriver gemm --iter 1000 --time 1 --a_w 2048 --a_h 2048 --b_w 2048
+# 5.5: Cijk_Ailk_Bljk_HHS_BH_MT128x128x16_MI16x16x16x1_SN_1LDSB0_APM1_ABV0_ACED0_AF0EM1_AF1EM1_AMAS3_ASE_ASGT_ASAE01_ASCE01_ASEM1_AAC0_BL1_BS1_DTL0_DTVA0_DVO0_ETSP_EPS1_FL0_GRVW8_GSU1_GSUASB_GLS0_ISA1100_IU1_K1_KLA_LBSPP128_LPA0_LPB8_LDL1_LRVW16_LWPMn1_LDW0_FMA_MIAV1_MDA2_NTA0_NTB0_NTC0_NTD0_NEPBS0_NLCA1_NLCB1_ONLL1_OPLV0_PK0_PAP0_PGR1_PLR1_RK0_SIA1_SS1_SU32_SUM0_SUS128_SCIUI1_SPO0_SRVW0_SSO0_SVW4_SNLL0_TT4_64_TLDS1_USFGROn1_VAW2_VSn1_VW4_WSGRA1_WSGRB1_WS32_WG32_4_1_WGM4
+# 5.6: Cijk_Ailk_Bljk_HHS_BH_MT128x128x16_MI16x16x16x1_SN_1LDSB0_APM1_ABV0_ACED0_AF0EM1_AF1EM1_AMAS3_ASE_ASGT_ASLT_ASAE01_ASCE01_ASEM1_AAC0_BL1_BS1_DTL0_DTVA0_DVO0_ETSP_EPS1_FL0_GRPM1_GRVW8_GSU1_GSUASB_GLS0_ISA1100_IU1_K1_KLA_LBSPP128_LPA0_LPB8_LDL1_LRVW16_LWPMn1_LDW0_FMA_MIAV1_MDA2_MO40_NTA0_NTB0_NTC0_NTD0_NEPBS0_NLCA1_NLCB1_ONLL1_OPLV0_PK0_PAP0_PGR1_PLR1_RK0_SIA1_SS1_SU32_SUM0_SUS128_SCIUI1_SPO0_SRVW0_SSO0_SVW4_SNLL0_TT4_64_TLDS1_USFGROn1_VAW2_VSn1_VW4_WSGRA1_WSGRB1_WS32_WG32_4_1_WGM4
+# gets ~100
+# hipExtModuleLaunchKernel ( 0x0x16ccde0, 2048, 16, 1, 128, 1, 1,
+# 161.60 us = 106.31 TFLOPS
+# with --batch_count 8 / 1.258128 ms / (8*2048*2048*2048*2)/(1.258128)*1e-9 / 109.24 TFLOPS
+
+# we only get ~53
+# KY=2 KX=2 N=2048 python3 extra/gemm/hip_matmul.py
+#   4194304    324.76 us, would be  52899.88 GFLOPS matmul, 154.98 GB/s
+
+DEBUG = getenv("DEBUG", 0)
+RAND = getenv("RAND", 0)
+CNT = getenv("CNT", 128)
+
+N = getenv("N", 4096)
+KX = getenv("KX", 4)
+KY = getenv("KY", 4)
+assert N%(16*KX) == 0, f"N must be multiple of {16*KX}"
+assert N%(16*KY) == 0, f"N must be multiple of {16*KY}"
+FLOPS = N*N*N*2
+BW = N*N*3*4
+
+local_size = [32, 1, 1]
+global_size = [N//(KX*16), N//(KY*16), 1]
+num_threads = prod(local_size)
+
+# Can AMDAllocator initialized as device=0 by default?
+device = AMDDevice()
+hipallocator = AMDAllocator(device)
+a = hipallocator.alloc(N*N*4)
+b = hipallocator.alloc(N*N*2)
+c = hipallocator.alloc(N*N*2)
+na = np.empty(N*N, np.float32)
+nb = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32).astype(np.float16)
+nc = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32).astype(np.float16)
+hipallocator._copyin(b, memoryview(bytearray(nb)))
+hipallocator._copyin(c, memoryview(bytearray(nc)))
+
+prog_str = f"""
+#define F32
+typedef long unsigned int size_t;
+#define half _Float16
+typedef float float8 __attribute__((ext_vector_type(8)));
+typedef _Float16 half4 __attribute__((ext_vector_type(4)));
+typedef _Float16 half8 __attribute__((ext_vector_type(8)));
+typedef _Float16 half16 __attribute__((ext_vector_type(16)));
+extern "C" __attribute__((device)) __attribute__((const)) size_t __ockl_get_local_id(unsigned int);
+extern "C" __attribute__((device)) __attribute__((const)) size_t __ockl_get_group_id(unsigned int);
+extern "C" __attribute__((device)) __attribute__((const)) size_t __ockl_get_local_size(unsigned int);
+extern "C" __attribute__((global))void __attribute__((amdgpu_flat_work_group_size(1, {num_threads}))) test(float* c, half* a, half* b) {{
+
+  const int gx = __ockl_get_group_id(0) + __ockl_get_local_id(2);
+  const int gy = __ockl_get_group_id(1) + __ockl_get_local_id(3);
+
+  const int lIdx = __ockl_get_local_id(0);
+  const int lane = lIdx%16;
+
+  c += gx*{KX*16}*{N} + gy*{KY*16} + (lIdx/16)*{N} + lane;
+  a += gx*{KX*16}*{N};
+  b += gy*{KY*16};
+
+  half16 a_frag[{KX}];
+  half16 b_frag[{KY}];
+  #ifdef F32
+    float8 c_frag[{KY}][{KX}] = {{}};
+  #else
+    half16 c_frag[{KY}][{KX}] = {{}};
+  #endif
+
+  for (int k = 0; k < {N}; k += 16) {{
+    __builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");
+    __builtin_amdgcn_s_barrier();
+    __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");
+
+    for (int ele = 0; ele < 16; ++ele) {{
+      for (int x = 0; x < {KX}; x++) {{
+        a_frag[x][ele] = a[(k+ele) + x*{16*N} + {N}*lane];
+      }}
+    }}
+    for (int ele = 0; ele < 16; ++ele) {{
+      for (int y = 0; y < {KY}; y++) {{
+        b_frag[y][ele] = b[(k+ele)*{N} + y*16 + lane];
+      }}
+    }}
+    for (int y = 0; y < {KY}; y++) {{
+      for (int x = 0; x < {KX}; x++) {{
+        #ifdef F32
+          c_frag[y][x] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a_frag[x], b_frag[y], c_frag[y][x]);
+        #else
+          c_frag[y][x] = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(a_frag[x], b_frag[y], c_frag[y][x], false);
+        #endif
+      }}
+    }}
+  }}
+
+  for (int ele = 0; ele < 8; ++ele) {{
+    for (int y = 0; y < {KY}; y++) {{
+      for (int x = 0; x < {KX}; x++) {{
+        #ifdef F32
+          c[ele*{2*N} + y*16 + x*{16*N}] = c_frag[y][x][ele];
+        #else
+          c[ele*{2*N} + y*16 + x*{16*N}] = c_frag[y][x][ele*2];
+        #endif
+      }}
+    }}
+  }}
+}}"""
+
+if DEBUG > 1: print(prog_str)
+lib = device.compiler.compile(prog_str)
+prog = AMDProgram(device, "test", lib)
+
+def timeit(fxn):
+  st = time.perf_counter()
+  et = fxn()
+  ret = time.perf_counter() - st # NOTE: et doesn't contain the launch overhead
+  if DEBUG > 0: print(f"{ret*1e6:.2f} us")
+  # rerun rand
+  if RAND:
+    nb = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32).astype(np.float16)
+    nc = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32).astype(np.float16)
+    hipallocator._copyin(b, memoryview(bytearray(nb)))
+    hipallocator._copyin(c, memoryview(bytearray(nc)))
+  return et
+
+print("global/local size", global_size, local_size, f"local_size:{prod(local_size)} total_size:{prod(global_size+local_size)}")
+tm = min([timeit(lambda: prog(a, b, c, global_size=global_size, local_size=local_size, wait=True)) for _ in range(CNT)])
+hipallocator._copyout(flat_mv(na.data),a)
+na = na.reshape(N,N)
+comp = nb.astype(np.float32) @ nc.astype(np.float32)
+print(f"{N*N:10d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS matmul, {BW*1e-9/tm:.2f} GB/s")
+if DEBUG > 2: print(f"which  nan={np.where(np.isnan(na))} len={len(np.where(np.isnan(na))[0])}")
+if DEBUG > 2: print(f"which diff={np.where(abs(na-comp) > 2e-2)} len={len(np.where(abs(na-comp) > 2e-2)[0])}")
+if DEBUG > 2: print(f"which zero={np.where(abs(na) < 2e-2)} len={len(np.where(abs(na) < 2e-2)[0])}")
+np.testing.assert_allclose(na, comp, atol=1e-2, rtol=1e-2)

tinygrad_repo/extra/gemm/intel_xmx.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+import numpy as np
+from tinygrad.runtime.ops_cl import CLProgram, CLCompiler
+from tinygrad import Device, dtypes
+from tinygrad.device import Buffer
+from hexdump import hexdump
+
+# https://github.com/intel/intel-graphics-compiler/blob/master/documentation/visa/instructions/DPAS.md
+# https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroups.html
+# https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_matrix_multiply_accumulate.html
+# https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_split_matrix_multiply_accumulate.html
+# https://hc34.hotchips.org/assets/program/conference/day1/GPU%20HPC/Intel_s%20Ponte%20Vecchio%20GPU%20-%20Architecture%20Systems%20and%20Software%20FINAL.pdf
+
+device = Device["CL"]
+
+# NOTE: only the subgroup type 8 ones work
+prog = CLProgram(device, "test", CLCompiler(device, "test").compile(f"""
+__attribute__((intel_reqd_sub_group_size(8)))
+__kernel void test(__global float* data0, const __global int* data1, const __global int8* data2) {{
+  int lidx0 = get_local_id(0);
+  int a = data1[lidx0];
+  int8 b = data2[lidx0];
+  float out = intel_sub_group_f16_f16_matrix_mad_k16(a, b, 0.0f);
+  data0[lidx0] = out;
+}}
+"""))
+#with open("/tmp/test.elf", "wb") as f: f.write(prog.lib)
+
+a = Buffer("CL", 8, dtypes.float32).allocate()
+b = Buffer("CL", 0x10, dtypes.float16).allocate()
+c = Buffer("CL", 8*0x10, dtypes.float16).allocate()
+
+row = np.array([1,2,3,4,5,6,7,8,1,2,3,4,5,6,7,8], np.float16)
+mat = np.random.random((8, 0x10)).astype(np.float16)
+
+b.copyin(row.data)
+c.copyin(mat.data)
+ret = prog(a._buf, b._buf, c._buf, global_size=[1,1,1], local_size=[8,1,1], wait=True)
+print(ret)
+out = np.frombuffer(a.as_memoryview(), np.float32)
+real = row.astype(np.float32)@mat.T.astype(np.float32)
+print("out:", out)
+print("real", real)

tinygrad_repo/extra/gemm/max_kernels/nv.fp16_fp16_fp16.2_stage.cu

@@ -0,0 +1,508 @@
+#define INFINITY (__int_as_float(0x7f800000))
+#define NAN (__int_as_float(0x7fffffff))
+#include <cuda_fp16.h>
+#include <cuda_pipeline.h>
+#define N_PAD 132
+
+struct __align__(8) half4 { half x, y, z, w; };
+__device__ half4 make_half4(half x, half y, half z, half w) { half4 r={x, y, z, w}; return r; }
+
+struct __align__(16) half8 { half x, y, z, w, a, b, c, d; };
+__device__ half8 make_half8(half x, half y, half z, half w, half a, half b, half c, half d) { half8 r={x, y, z, w, a, b, c, d}; return r; }
+
+__device__ void __ldmatrix_a_elems(half8 *regs, half *smem) {
+    uint32_t reg0, reg1, reg2, reg3;
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(reg0), "=r"(reg1), "=r"(reg2), "=r"(reg3)
+        : "l"(__cvta_generic_to_shared(smem))
+    );
+    uint32_t *addr = reinterpret_cast<uint32_t*>(regs);
+    addr[0] = reg0;
+    addr[1] = reg1;
+    addr[2] = reg2;
+    addr[3] = reg3;
+}
+
+__device__ void __ldmatrix_b_elems(half4 *regs_lo, half4 *regs_hi, half *smem) {
+    uint32_t reg0, reg1, reg2, reg3;
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(reg0), "=r"(reg1), "=r"(reg2), "=r"(reg3)
+        : "l"(__cvta_generic_to_shared(smem))
+    );
+    uint32_t *addr_lo = reinterpret_cast<uint32_t*>(regs_lo);
+    uint32_t *addr_hi = reinterpret_cast<uint32_t*>(regs_hi);
+    addr_lo[0] = reg0;
+    addr_lo[1] = reg1;
+    addr_hi[0] = reg2;
+    addr_hi[1] = reg3;
+}
+
+__device__ half4 __WMMA_8_16_16_half_half(half8 a, half4 b, half4 c) {
+    int *a_pk = (int *) (&a), *b_pk = (int *) (&b), *c_pk = (int *) (&c);
+    asm( "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 { %0, %1 }, { %2, %3, %4, %5 }, { %6, %7 }, { %0, %1 };"
+        : "+r"(c_pk[0]), "+r"(c_pk[1]): "r"(a_pk[0]), "r"(a_pk[1]), "r"(a_pk[2]),  "r"(a_pk[3]), "r"(b_pk[0]), "r"(b_pk[1]) );
+    return c;
+}
+
+extern "C" __global__ void __launch_bounds__(128) wmma_example(half* data0, const half* data1, const half* data2, int N, int K) {
+    int grid_m = blockIdx.x;        /* M//64 */
+    int grid_n = blockIdx.y;        /* N//128 */
+    int threads = threadIdx.x;      /* 128 */
+    int wg_m = (threads/64);        // 0 or 1 for 1st and 3rd blocks of b_m=16xb_k=16 vs 2nd and 4th blocks
+    int wg_n = (threads/32)%2;      // 0 or 1 for 1st, 3rd, 5th, 7th blocks of b_n=16xb_k=16 vs 2nd, 4th, 6th, 8th blocks - differs from triton
+    int wg_threads = threads%32;
+    int num_k_blocks = K / 64;
+
+    // load indexes
+    size_t global_a_off = ((grid_m * 64) * K) + ((threads %  8) * 8) + ((threads /  8) * K);
+    size_t global_b_off = (grid_n * 128)      + ((threads % 16) * 8) + ((threads / 16) * N);
+
+    // swizzled smem store offsets - columns of smem are swizzled
+    // here's a link to a description of the triton: https://github.com/triton-lang/triton/discussions/2026#discussioncomment-6746579
+    // see also the thunderkittens impl: https://github.com/HazyResearch/ThunderKittens/blob/main/include/types/shared/st.cuh
+    size_t store_smem_a_off  = ((threads /  8) *  64) + (((threads * 8) ^ threads) & 56);               // r15
+    size_t store_smem_b_off  = ((threads / 16) * 128) + (((threads / 16) * 8) ^ ((threads % 16) * 8));  // r19\
+
+    // ldmatrix indices
+    // threads 0-7 are row starts for A, 8-15 for B, 16-23 for C, 24-31 for D
+    // [ A | C ]
+    // [ - + - ]
+    // [ B | D ]
+
+    // swizzled ldmatrix
+    size_t load_smem_a_row   = ((wg_m * 16) + (threads % 16)) * 64; // r293
+    size_t load_smem_a_phase = (threads / 16) % 2;                  // r4
+
+    size_t load_smem_b_row   = (threads % 16) * 128;                // r299
+    size_t load_smem_b_phase = (wg_n * 2) + (((threads / 16) % 2)); // r297 -- this differs from the generated triton kernel (swapped order)
+
+    size_t load_smem_a_0_k_0 = load_smem_a_row + (((load_smem_a_phase +  0) ^ (threads % 8)) * 8); // r38
+    size_t load_smem_a_1_k_0 = load_smem_a_0_k_0 + (32 * 64);
+    size_t load_smem_b_0_k_0 = load_smem_b_row + (((load_smem_b_phase +  0) ^ (threads % 8)) * 8);
+    size_t load_smem_b_1_k_0 = load_smem_b_row + (((load_smem_b_phase +  4) ^ (threads % 8)) * 8);
+    size_t load_smem_b_2_k_0 = load_smem_b_row + (((load_smem_b_phase +  8) ^ (threads % 8)) * 8);
+    size_t load_smem_b_3_k_0 = load_smem_b_row + (((load_smem_b_phase + 12) ^ (threads % 8)) * 8);
+
+    size_t load_smem_a_0_k_1 = load_smem_a_row + (((load_smem_a_phase +  2) ^ (threads % 8)) * 8); // r58 = r293 + r316;
+    size_t load_smem_a_1_k_1 = load_smem_a_0_k_1 + (32 * 64);
+    size_t load_smem_b_0_k_1 = load_smem_b_0_k_0 + (16 * 128);
+    size_t load_smem_b_1_k_1 = load_smem_b_1_k_0 + (16 * 128);
+    size_t load_smem_b_2_k_1 = load_smem_b_2_k_0 + (16 * 128);
+    size_t load_smem_b_3_k_1 = load_smem_b_3_k_0 + (16 * 128);
+
+    size_t load_smem_a_0_k_2 = load_smem_a_row + (((load_smem_a_phase +  4) ^ (threads % 8)) * 8); // r59 = r293 + r319;
+    size_t load_smem_a_1_k_2 = load_smem_a_0_k_2 + (32 * 64);
+    size_t load_smem_b_0_k_2 = load_smem_b_0_k_0 + (32 * 128);
+    size_t load_smem_b_1_k_2 = load_smem_b_1_k_0 + (32 * 128);
+    size_t load_smem_b_2_k_2 = load_smem_b_2_k_0 + (32 * 128);
+    size_t load_smem_b_3_k_2 = load_smem_b_3_k_0 + (32 * 128);
+
+    size_t load_smem_a_0_k_3 = load_smem_a_row + (((load_smem_a_phase +  6) ^ (threads % 8)) * 8); // r60 = r293 + r322;
+    size_t load_smem_a_1_k_3 = load_smem_a_0_k_3 + (32 * 64);
+    size_t load_smem_b_0_k_3 = load_smem_b_0_k_0 + (48 * 128);
+    size_t load_smem_b_1_k_3 = load_smem_b_1_k_0 + (48 * 128);
+    size_t load_smem_b_2_k_3 = load_smem_b_2_k_0 + (48 * 128);
+    size_t load_smem_b_3_k_3 = load_smem_b_3_k_0 + (48 * 128);
+
+    // create shared mem (A_1 8192 bytes, A_2 8192 bytes, B_1 16384 bytes, B2_16384 bytes)
+    __shared__ alignas(16) char smem[49152];
+
+    // create accs (16 WMMAs and 4 output elements each) and zero
+    half4 acc_frag_0_0 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_1 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_2 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_3 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_4 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_5 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_6 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_7 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_0 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_1 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_2 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_3 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_4 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_5 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_6 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_7 = make_half4(0.0f,0.0f,0.0f,0.0f);
+
+    // create registers for block A elements (2)
+    half8 a_frag_0;
+    half8 a_frag_1;
+
+    // create register for block B elements (8)
+    half4 b_frag_0;
+    half4 b_frag_1;
+    half4 b_frag_2;
+    half4 b_frag_3;
+    half4 b_frag_4;
+    half4 b_frag_5;
+    half4 b_frag_6;
+    half4 b_frag_7;
+
+    half *smem_a_even = (half *)(smem);
+    half *smem_a_odd  = (half *)(smem + 8192);
+    half *smem_b_even = (half *)(smem + 16384);
+    half *smem_b_odd  = (half *)(smem + 32768);
+
+    // https://developer.nvidia.com/blog/controlling-data-movement-to-boost-performance-on-ampere-architecture/
+    // https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#asynchronous-data-copies
+
+    // start first pre-fetch load A
+    __pipeline_memcpy_async(&smem_a_even[store_smem_a_off +  (    0)], &data1[global_a_off + (   0)], 16);
+    __pipeline_memcpy_async(&smem_a_even[store_smem_a_off +  (16*64)], &data1[global_a_off + (16*K)], 16);
+    __pipeline_memcpy_async(&smem_a_even[store_smem_a_off +  (32*64)], &data1[global_a_off + (32*K)], 16);
+    __pipeline_memcpy_async(&smem_a_even[store_smem_a_off +  (48*64)], &data1[global_a_off + (48*K)], 16);
+
+    // start first pre-fetch load B
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (     0)], &data2[global_b_off + (   0)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + ( 8*128)], &data2[global_b_off + ( 8*N)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (16*128)], &data2[global_b_off + (16*N)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (24*128)], &data2[global_b_off + (24*N)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (32*128)], &data2[global_b_off + (32*N)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (40*128)], &data2[global_b_off + (40*N)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (48*128)], &data2[global_b_off + (48*N)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (56*128)], &data2[global_b_off + (56*N)], 16);
+    __pipeline_commit();
+
+    global_a_off += 64;
+    global_b_off += 64 * N;
+    __syncthreads();
+
+    // start second pre-fetch load A
+    __pipeline_memcpy_async(&smem_a_odd[store_smem_a_off +  (    0)], &data1[global_a_off + (   0)], 16);
+    __pipeline_memcpy_async(&smem_a_odd[store_smem_a_off +  (16*64)], &data1[global_a_off + (16*K)], 16);
+    __pipeline_memcpy_async(&smem_a_odd[store_smem_a_off +  (32*64)], &data1[global_a_off + (32*K)], 16);
+    __pipeline_memcpy_async(&smem_a_odd[store_smem_a_off +  (48*64)], &data1[global_a_off + (48*K)], 16);
+
+    // start second pre-fetch load B
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (     0)], &data2[global_b_off + (   0)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + ( 8*128)], &data2[global_b_off + ( 8*N)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (16*128)], &data2[global_b_off + (16*N)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (24*128)], &data2[global_b_off + (24*N)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (32*128)], &data2[global_b_off + (32*N)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (40*128)], &data2[global_b_off + (40*N)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (48*128)], &data2[global_b_off + (48*N)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (56*128)], &data2[global_b_off + (56*N)], 16);
+    __pipeline_commit();
+
+    global_a_off += 64;
+    global_b_off += 64 * N;
+
+    // wait on needed prefetch value
+    __pipeline_wait_prior(0); // TODO: this enables fast iterations, but incorrect results with 1 (it shouldn't)
+    __syncthreads();
+
+    for (int block_k = 0; block_k < num_k_blocks; block_k++) {
+        // BLOCK_K==4: unroll 4 iterations of ldmatrix/wmma
+        half *smem_a_curr = (block_k % 2) ? smem_a_even : smem_a_odd;
+        half *smem_b_curr = (block_k % 2) ? smem_b_even : smem_b_odd;
+
+        // first load 16 K elements and 16 WMMAs: BLOCK_M==2 * BLOCK_N==8
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_0]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_0]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_0]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_0]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_0]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_0]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // next 16 K elements
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_1]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_1]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_1]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_1]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_1]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_1]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // next 16 K elements
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_2]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_2]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_2]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_2]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_2]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_2]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // last 16 K elements
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_3]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_3]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_3]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_3]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_3]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_3]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_half(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_half(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // prefetch next iteration if needed
+        __syncthreads();
+        if (block_k < (num_k_blocks-2)) {
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (    0)], &data1[global_a_off + (   0)], 16);
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (16*64)], &data1[global_a_off + (16*K)], 16);
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (32*64)], &data1[global_a_off + (32*K)], 16);
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (48*64)], &data1[global_a_off + (48*K)], 16);
+
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (     0)], &data2[global_b_off + (   0)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + ( 8*128)], &data2[global_b_off + ( 8*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (16*128)], &data2[global_b_off + (16*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (24*128)], &data2[global_b_off + (24*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (32*128)], &data2[global_b_off + (32*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (40*128)], &data2[global_b_off + (40*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (48*128)], &data2[global_b_off + (48*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (56*128)], &data2[global_b_off + (56*N)], 16);
+
+            global_a_off += 64;
+            global_b_off += 64 * N;
+        }
+        __pipeline_commit();
+
+        if (block_k < num_k_blocks-1) {
+            __pipeline_wait_prior(1);
+            __syncthreads();
+        }
+    }
+
+    // write accumulators to output
+    __pipeline_wait_prior(0);
+    __syncthreads();
+
+    // // store registers to smem first, then read back to do float4 writes to global
+    // float *smem_d = (float *)(smem);
+    // size_t smem_d_off = (wg_m * 16 * N_PAD) + (wg_n * 16) + ((wg_threads % 4) * 2) + (((wg_threads / 4) % 8) * N_PAD);
+    // smem_d[smem_d_off + 0 + ( 0*8)            ] = acc_frag_0_0.x;
+    // smem_d[smem_d_off + 1 + ( 0*8)            ] = acc_frag_0_0.y;
+    // smem_d[smem_d_off + 0 + ( 0*8) + (8*N_PAD)] = acc_frag_0_0.z;
+    // smem_d[smem_d_off + 1 + ( 0*8) + (8*N_PAD)] = acc_frag_0_0.w;
+    // smem_d[smem_d_off + 0 + ( 1*8)            ] = acc_frag_0_1.x;
+    // smem_d[smem_d_off + 1 + ( 1*8)            ] = acc_frag_0_1.y;
+    // smem_d[smem_d_off + 0 + ( 1*8) + (8*N_PAD)] = acc_frag_0_1.z;
+    // smem_d[smem_d_off + 1 + ( 1*8) + (8*N_PAD)] = acc_frag_0_1.w;
+    // smem_d[smem_d_off + 0 + ( 4*8)            ] = acc_frag_0_2.x;
+    // smem_d[smem_d_off + 1 + ( 4*8)            ] = acc_frag_0_2.y;
+    // smem_d[smem_d_off + 0 + ( 4*8) + (8*N_PAD)] = acc_frag_0_2.z;
+    // smem_d[smem_d_off + 1 + ( 4*8) + (8*N_PAD)] = acc_frag_0_2.w;
+    // smem_d[smem_d_off + 0 + ( 5*8)            ] = acc_frag_0_3.x;
+    // smem_d[smem_d_off + 1 + ( 5*8)            ] = acc_frag_0_3.y;
+    // smem_d[smem_d_off + 0 + ( 5*8) + (8*N_PAD)] = acc_frag_0_3.z;
+    // smem_d[smem_d_off + 1 + ( 5*8) + (8*N_PAD)] = acc_frag_0_3.w;
+    // smem_d[smem_d_off + 0 + ( 8*8)            ] = acc_frag_0_4.x;
+    // smem_d[smem_d_off + 1 + ( 8*8)            ] = acc_frag_0_4.y;
+    // smem_d[smem_d_off + 0 + ( 8*8) + (8*N_PAD)] = acc_frag_0_4.z;
+    // smem_d[smem_d_off + 1 + ( 8*8) + (8*N_PAD)] = acc_frag_0_4.w;
+    // smem_d[smem_d_off + 0 + ( 9*8)            ] = acc_frag_0_5.x;
+    // smem_d[smem_d_off + 1 + ( 9*8)            ] = acc_frag_0_5.y;
+    // smem_d[smem_d_off + 0 + ( 9*8) + (8*N_PAD)] = acc_frag_0_5.z;
+    // smem_d[smem_d_off + 1 + ( 9*8) + (8*N_PAD)] = acc_frag_0_5.w;
+    // smem_d[smem_d_off + 0 + (12*8)            ] = acc_frag_0_6.x;
+    // smem_d[smem_d_off + 1 + (12*8)            ] = acc_frag_0_6.y;
+    // smem_d[smem_d_off + 0 + (12*8) + (8*N_PAD)] = acc_frag_0_6.z;
+    // smem_d[smem_d_off + 1 + (12*8) + (8*N_PAD)] = acc_frag_0_6.w;
+    // smem_d[smem_d_off + 0 + (13*8)            ] = acc_frag_0_7.x;
+    // smem_d[smem_d_off + 1 + (13*8)            ] = acc_frag_0_7.y;
+    // smem_d[smem_d_off + 0 + (13*8) + (8*N_PAD)] = acc_frag_0_7.z;
+    // smem_d[smem_d_off + 1 + (13*8) + (8*N_PAD)] = acc_frag_0_7.w;
+
+    // __syncthreads();
+    // size_t load_smem_d_off = ((threads % 32) * 4) + ((threads / 32) * N_PAD);
+    // float4 d_0_0 = *((float4 *)(smem_d + load_smem_d_off + ( 0 * N_PAD)));
+    // float4 d_0_1 = *((float4 *)(smem_d + load_smem_d_off + ( 4 * N_PAD)));
+    // float4 d_0_2 = *((float4 *)(smem_d + load_smem_d_off + ( 8 * N_PAD)));
+    // float4 d_0_3 = *((float4 *)(smem_d + load_smem_d_off + (12 * N_PAD)));
+    // float4 d_0_4 = *((float4 *)(smem_d + load_smem_d_off + (16 * N_PAD)));
+    // float4 d_0_5 = *((float4 *)(smem_d + load_smem_d_off + (20 * N_PAD)));
+    // float4 d_0_6 = *((float4 *)(smem_d + load_smem_d_off + (24 * N_PAD)));
+    // float4 d_0_7 = *((float4 *)(smem_d + load_smem_d_off + (28 * N_PAD)));
+
+    // __syncthreads();
+    // smem_d[smem_d_off + 0 + ( 0*8)            ] = acc_frag_1_0.x;
+    // smem_d[smem_d_off + 1 + ( 0*8)            ] = acc_frag_1_0.y;
+    // smem_d[smem_d_off + 0 + ( 0*8) + (8*N_PAD)] = acc_frag_1_0.z;
+    // smem_d[smem_d_off + 1 + ( 0*8) + (8*N_PAD)] = acc_frag_1_0.w;
+    // smem_d[smem_d_off + 0 + ( 1*8)            ] = acc_frag_1_1.x;
+    // smem_d[smem_d_off + 1 + ( 1*8)            ] = acc_frag_1_1.y;
+    // smem_d[smem_d_off + 0 + ( 1*8) + (8*N_PAD)] = acc_frag_1_1.z;
+    // smem_d[smem_d_off + 1 + ( 1*8) + (8*N_PAD)] = acc_frag_1_1.w;
+    // smem_d[smem_d_off + 0 + ( 4*8)            ] = acc_frag_1_2.x;
+    // smem_d[smem_d_off + 1 + ( 4*8)            ] = acc_frag_1_2.y;
+    // smem_d[smem_d_off + 0 + ( 4*8) + (8*N_PAD)] = acc_frag_1_2.z;
+    // smem_d[smem_d_off + 1 + ( 4*8) + (8*N_PAD)] = acc_frag_1_2.w;
+    // smem_d[smem_d_off + 0 + ( 5*8)            ] = acc_frag_1_3.x;
+    // smem_d[smem_d_off + 1 + ( 5*8)            ] = acc_frag_1_3.y;
+    // smem_d[smem_d_off + 0 + ( 5*8) + (8*N_PAD)] = acc_frag_1_3.z;
+    // smem_d[smem_d_off + 1 + ( 5*8) + (8*N_PAD)] = acc_frag_1_3.w;
+    // smem_d[smem_d_off + 0 + ( 8*8)            ] = acc_frag_1_4.x;
+    // smem_d[smem_d_off + 1 + ( 8*8)            ] = acc_frag_1_4.y;
+    // smem_d[smem_d_off + 0 + ( 8*8) + (8*N_PAD)] = acc_frag_1_4.z;
+    // smem_d[smem_d_off + 1 + ( 8*8) + (8*N_PAD)] = acc_frag_1_4.w;
+    // smem_d[smem_d_off + 0 + ( 9*8)            ] = acc_frag_1_5.x;
+    // smem_d[smem_d_off + 1 + ( 9*8)            ] = acc_frag_1_5.y;
+    // smem_d[smem_d_off + 0 + ( 9*8) + (8*N_PAD)] = acc_frag_1_5.z;
+    // smem_d[smem_d_off + 1 + ( 9*8) + (8*N_PAD)] = acc_frag_1_5.w;
+    // smem_d[smem_d_off + 0 + (12*8)            ] = acc_frag_1_6.x;
+    // smem_d[smem_d_off + 1 + (12*8)            ] = acc_frag_1_6.y;
+    // smem_d[smem_d_off + 0 + (12*8) + (8*N_PAD)] = acc_frag_1_6.z;
+    // smem_d[smem_d_off + 1 + (12*8) + (8*N_PAD)] = acc_frag_1_6.w;
+    // smem_d[smem_d_off + 0 + (13*8)            ] = acc_frag_1_7.x;
+    // smem_d[smem_d_off + 1 + (13*8)            ] = acc_frag_1_7.y;
+    // smem_d[smem_d_off + 0 + (13*8) + (8*N_PAD)] = acc_frag_1_7.z;
+    // smem_d[smem_d_off + 1 + (13*8) + (8*N_PAD)] = acc_frag_1_7.w;
+
+    // __syncthreads();
+    // float4 d_1_0 = *((float4 *)(smem_d + load_smem_d_off + ( 0 * N_PAD)));
+    // float4 d_1_1 = *((float4 *)(smem_d + load_smem_d_off + ( 4 * N_PAD)));
+    // float4 d_1_2 = *((float4 *)(smem_d + load_smem_d_off + ( 8 * N_PAD)));
+    // float4 d_1_3 = *((float4 *)(smem_d + load_smem_d_off + (12 * N_PAD)));
+    // float4 d_1_4 = *((float4 *)(smem_d + load_smem_d_off + (16 * N_PAD)));
+    // float4 d_1_5 = *((float4 *)(smem_d + load_smem_d_off + (20 * N_PAD)));
+    // float4 d_1_6 = *((float4 *)(smem_d + load_smem_d_off + (24 * N_PAD)));
+    // float4 d_1_7 = *((float4 *)(smem_d + load_smem_d_off + (28 * N_PAD)));
+
+    // __syncthreads();
+    // float *global_d = &data0[((grid_m * 64) * N) + (grid_n * 128) + ((threads % 32) * 4) + ((threads / 32) * N)];
+    // *((float4 *)(global_d +  0*N)) = d_0_0;
+    // *((float4 *)(global_d +  4*N)) = d_0_1;
+    // *((float4 *)(global_d +  8*N)) = d_0_2;
+    // *((float4 *)(global_d + 12*N)) = d_0_3;
+    // *((float4 *)(global_d + 16*N)) = d_0_4;
+    // *((float4 *)(global_d + 20*N)) = d_0_5;
+    // *((float4 *)(global_d + 24*N)) = d_0_6;
+    // *((float4 *)(global_d + 28*N)) = d_0_7;
+    // *((float4 *)(global_d + 32*N)) = d_1_0;
+    // *((float4 *)(global_d + 36*N)) = d_1_1;
+    // *((float4 *)(global_d + 40*N)) = d_1_2;
+    // *((float4 *)(global_d + 44*N)) = d_1_3;
+    // *((float4 *)(global_d + 48*N)) = d_1_4;
+    // *((float4 *)(global_d + 52*N)) = d_1_5;
+    // *((float4 *)(global_d + 56*N)) = d_1_6;
+    // *((float4 *)(global_d + 60*N)) = d_1_7;
+
+    // slower way: write floats one by one to data0
+    size_t wg_c_off     = ((grid_m * 64) * N) + (grid_n * 128) + (wg_m * 16 * N) + (wg_n * 16);
+    size_t thread_c_off = ((wg_threads % 4) * 2) + (((wg_threads / 4) % 8) * N);
+    data0[wg_c_off + thread_c_off           + 0 + ( 0*8)] = acc_frag_0_0.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 0*8)] = acc_frag_0_0.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 0*8)] = acc_frag_0_0.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 0*8)] = acc_frag_0_0.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 1*8)] = acc_frag_0_1.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 1*8)] = acc_frag_0_1.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 1*8)] = acc_frag_0_1.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 1*8)] = acc_frag_0_1.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 4*8)] = acc_frag_0_2.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 4*8)] = acc_frag_0_2.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 4*8)] = acc_frag_0_2.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 4*8)] = acc_frag_0_2.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 5*8)] = acc_frag_0_3.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 5*8)] = acc_frag_0_3.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 5*8)] = acc_frag_0_3.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 5*8)] = acc_frag_0_3.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 8*8)] = acc_frag_0_4.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 8*8)] = acc_frag_0_4.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 8*8)] = acc_frag_0_4.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 8*8)] = acc_frag_0_4.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 9*8)] = acc_frag_0_5.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 9*8)] = acc_frag_0_5.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 9*8)] = acc_frag_0_5.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 9*8)] = acc_frag_0_5.w;
+    data0[wg_c_off + thread_c_off           + 0 + (12*8)] = acc_frag_0_6.x;
+    data0[wg_c_off + thread_c_off           + 1 + (12*8)] = acc_frag_0_6.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (12*8)] = acc_frag_0_6.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (12*8)] = acc_frag_0_6.w;
+    data0[wg_c_off + thread_c_off           + 0 + (13*8)] = acc_frag_0_7.x;
+    data0[wg_c_off + thread_c_off           + 1 + (13*8)] = acc_frag_0_7.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (13*8)] = acc_frag_0_7.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (13*8)] = acc_frag_0_7.w;
+    wg_c_off += 32*N;
+    data0[wg_c_off + thread_c_off           + 0 + ( 0*8)] = acc_frag_1_0.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 0*8)] = acc_frag_1_0.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 0*8)] = acc_frag_1_0.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 0*8)] = acc_frag_1_0.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 1*8)] = acc_frag_1_1.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 1*8)] = acc_frag_1_1.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 1*8)] = acc_frag_1_1.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 1*8)] = acc_frag_1_1.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 4*8)] = acc_frag_1_2.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 4*8)] = acc_frag_1_2.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 4*8)] = acc_frag_1_2.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 4*8)] = acc_frag_1_2.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 5*8)] = acc_frag_1_3.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 5*8)] = acc_frag_1_3.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 5*8)] = acc_frag_1_3.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 5*8)] = acc_frag_1_3.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 8*8)] = acc_frag_1_4.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 8*8)] = acc_frag_1_4.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 8*8)] = acc_frag_1_4.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 8*8)] = acc_frag_1_4.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 9*8)] = acc_frag_1_5.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 9*8)] = acc_frag_1_5.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 9*8)] = acc_frag_1_5.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 9*8)] = acc_frag_1_5.w;
+    data0[wg_c_off + thread_c_off           + 0 + (12*8)] = acc_frag_1_6.x;
+    data0[wg_c_off + thread_c_off           + 1 + (12*8)] = acc_frag_1_6.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (12*8)] = acc_frag_1_6.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (12*8)] = acc_frag_1_6.w;
+    data0[wg_c_off + thread_c_off           + 0 + (13*8)] = acc_frag_1_7.x;
+    data0[wg_c_off + thread_c_off           + 1 + (13*8)] = acc_frag_1_7.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (13*8)] = acc_frag_1_7.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (13*8)] = acc_frag_1_7.w;
+}

tinygrad_repo/extra/gemm/max_kernels/nv.fp16_fp16_fp16.3_stage.cu

@@ -0,0 +1,465 @@
+#define INFINITY (__int_as_float(0x7f800000))
+#define NAN (__int_as_float(0x7fffffff))
+#include <cuda_fp16.h>
+#include <cuda_pipeline.h>
+#define N_PAD 132
+
+struct __align__(8) half4 { half x, y, z, w; };
+__device__ half4 make_half4(half x, half y, half z, half w) { half4 r={x, y, z, w}; return r; }
+
+struct __align__(16) half8 { half x, y, z, w, a, b, c, d; };
+__device__ half8 make_half8(half x, half y, half z, half w, half a, half b, half c, half d) { half8 r={x, y, z, w, a, b, c, d}; return r; }
+
+__device__ void __ldmatrix_a_elems(half8 *regs, half *smem) {
+    uint32_t reg0, reg1, reg2, reg3;
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(reg0), "=r"(reg1), "=r"(reg2), "=r"(reg3)
+        : "l"(__cvta_generic_to_shared(smem))
+    );
+    uint32_t *addr = reinterpret_cast<uint32_t*>(regs);
+    addr[0] = reg0;
+    addr[1] = reg1;
+    addr[2] = reg2;
+    addr[3] = reg3;
+}
+
+__device__ void __ldmatrix_b_elems(half4 *regs_lo, half4 *regs_hi, half *smem) {
+    uint32_t reg0, reg1, reg2, reg3;
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(reg0), "=r"(reg1), "=r"(reg2), "=r"(reg3)
+        : "l"(__cvta_generic_to_shared(smem))
+    );
+    uint32_t *addr_lo = reinterpret_cast<uint32_t*>(regs_lo);
+    uint32_t *addr_hi = reinterpret_cast<uint32_t*>(regs_hi);
+    addr_lo[0] = reg0;
+    addr_lo[1] = reg1;
+    addr_hi[0] = reg2;
+    addr_hi[1] = reg3;
+}
+
+__device__ half4 __WMMA_8_16_16_half_half(half8 a, half4 b, half4 c) {
+    int *a_pk = (int *) (&a), *b_pk = (int *) (&b), *c_pk = (int *) (&c);
+    asm( "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 { %0, %1 }, { %2, %3, %4, %5 }, { %6, %7 }, { %0, %1 };"
+        : "+r"(c_pk[0]), "+r"(c_pk[1]): "r"(a_pk[0]), "r"(a_pk[1]), "r"(a_pk[2]),  "r"(a_pk[3]), "r"(b_pk[0]), "r"(b_pk[1]) );
+    return c;
+}
+
+extern "C" __global__ void __launch_bounds__(256) wmma_example(half* data0, const half* data1, const half* data2, int N, int K) {
+    extern __shared__ char smem[];
+    half *smem_a_0 = (half *)(smem);
+    half *smem_a_1 = (half *)(smem + 16384);
+    half *smem_a_2 = (half *)(smem + 32768);
+    half *smem_b_0 = (half *)(smem + 49152);
+    half *smem_b_1 = (half *)(smem + 57344);
+    half *smem_b_2 = (half *)(smem + 65536);
+
+    int grid_m = blockIdx.x;        /* M//256 */
+    int grid_n = blockIdx.y;        /* N//128 */
+    int wg_threads = threadIdx.x;   // 32
+    int wg_m = threadIdx.y;         // 4
+    int wg_n = threadIdx.z;         // 2
+    int threads = threadIdx.x + (threadIdx.y * 32) + (threadIdx.z * 128); /* 256 */
+    int num_k_blocks = K / 32;
+
+    // load indexes
+    size_t global_a_off = ((grid_m * 256) * K) + ((threads %  4) * 8) + ((threads /  4) * K);
+    size_t global_b_off = (grid_n * 128)       + ((threads % 16) * 8) + ((threads / 16) * N);
+
+    // unswizzed smem store
+    size_t store_smem_a_off = ((threads %  4) * 8) + ((threads /  4) *  32);    // 64 rows /  32 cols per copy
+    size_t store_smem_b_off = ((threads % 16) * 8) + ((threads / 16) * 128);    // 16 rows / 128 cols per copy
+
+    // ldmatrix indices
+    // threads 0-7 are row starts for A, 8-15 for B, 16-23 for C, 24-31 for D
+    // [ A | C ]
+    // [ - + - ]
+    // [ B | D ]
+
+    // unswizzed ldmatrix
+    size_t load_smem_a_0_k_0 = (wg_m * 16 * 32) + ((wg_threads % 16) * 32) + ((wg_threads / 16) * 8);
+    size_t load_smem_a_1_k_0 = load_smem_a_0_k_0 + ( 64 * 32);
+    size_t load_smem_a_2_k_0 = load_smem_a_0_k_0 + (128 * 32);
+    size_t load_smem_a_3_k_0 = load_smem_a_0_k_0 + (192 * 32);
+
+    size_t load_smem_a_0_k_1 = load_smem_a_0_k_0 + 16;
+    size_t load_smem_a_1_k_1 = load_smem_a_0_k_1 + ( 64 * 32);
+    size_t load_smem_a_2_k_1 = load_smem_a_0_k_1 + (128 * 32);
+    size_t load_smem_a_3_k_1 = load_smem_a_0_k_1 + (192 * 32);
+
+    size_t load_smem_b_0_k_0 = (wg_n * 16) + ((wg_threads % 16) * 128) + ((wg_threads / 16) * 8);
+    size_t load_smem_b_1_k_0 = load_smem_b_0_k_0 + 32;
+    size_t load_smem_b_2_k_0 = load_smem_b_0_k_0 + 64;
+    size_t load_smem_b_3_k_0 = load_smem_b_0_k_0 + 96;
+
+    size_t load_smem_b_0_k_1 = load_smem_b_0_k_0 + (16 * 128);
+    size_t load_smem_b_1_k_1 = load_smem_b_0_k_1 + 32;
+    size_t load_smem_b_2_k_1 = load_smem_b_0_k_1 + 64;
+    size_t load_smem_b_3_k_1 = load_smem_b_0_k_1 + 96;
+
+    // create accs (M=4, N=8)
+    half4 acc_frag_0_0 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_1 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_2 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_3 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_4 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_5 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_6 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_7 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_0 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_1 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_2 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_3 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_4 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_5 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_6 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_7 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_0 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_1 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_2 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_3 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_4 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_5 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_6 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_7 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_0 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_1 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_2 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_3 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_4 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_5 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_6 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_7 = make_half4(0.0f,0.0f,0.0f,0.0f);
+
+    // create registers for block A elements
+    half8 a_frag_0_k_0;
+    half8 a_frag_1_k_0;
+    half8 a_frag_2_k_0;
+    half8 a_frag_3_k_0;
+    half8 a_frag_0_k_1;
+    half8 a_frag_1_k_1;
+    half8 a_frag_2_k_1;
+    half8 a_frag_3_k_1;
+
+    // create register for block B elements
+    half4 b_frag_0_k_0;
+    half4 b_frag_1_k_0;
+    half4 b_frag_2_k_0;
+    half4 b_frag_3_k_0;
+    half4 b_frag_4_k_0;
+    half4 b_frag_5_k_0;
+    half4 b_frag_6_k_0;
+    half4 b_frag_7_k_0;
+    half4 b_frag_0_k_1;
+    half4 b_frag_1_k_1;
+    half4 b_frag_2_k_1;
+    half4 b_frag_3_k_1;
+    half4 b_frag_4_k_1;
+    half4 b_frag_5_k_1;
+    half4 b_frag_6_k_1;
+    half4 b_frag_7_k_1;
+
+    __syncthreads();
+
+    // load first tile
+    __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + (     0)], &data1[global_a_off + (    0)], 16);
+    __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + ( 64*32)], &data1[global_a_off + ( 64*K)], 16);
+    __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + (128*32)], &data1[global_a_off + (128*K)], 16);
+    __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + (192*32)], &data1[global_a_off + (192*K)], 16);
+    __pipeline_memcpy_async(&smem_b_0[store_smem_b_off + (     0)], &data2[global_b_off + (    0)], 16);
+    __pipeline_memcpy_async(&smem_b_0[store_smem_b_off + (16*128)], &data2[global_b_off + ( 16*N)], 16);
+    __pipeline_commit();
+
+    global_a_off += 32;
+    global_b_off += 32 * N;
+
+    // load second tile
+    __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + (     0)], &data1[global_a_off + (    0)], 16);
+    __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + ( 64*32)], &data1[global_a_off + ( 64*K)], 16);
+    __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + (128*32)], &data1[global_a_off + (128*K)], 16);
+    __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + (192*32)], &data1[global_a_off + (192*K)], 16);
+    __pipeline_memcpy_async(&smem_b_1[store_smem_b_off + (     0)], &data2[global_b_off + (    0)], 16);
+    __pipeline_memcpy_async(&smem_b_1[store_smem_b_off + (16*128)], &data2[global_b_off + ( 16*N)], 16);
+    __pipeline_commit();
+
+    global_a_off += 32;
+    global_b_off += 32 * N;
+
+    // wait on first pre-fetch load
+    __pipeline_wait_prior(1);
+    __syncthreads();
+
+    // load K=0 for the first tile
+    __ldmatrix_a_elems(&a_frag_0_k_0,                &smem_a_0[load_smem_a_0_k_0]);
+    __ldmatrix_a_elems(&a_frag_1_k_0,                &smem_a_0[load_smem_a_1_k_0]);
+    __ldmatrix_a_elems(&a_frag_2_k_0,                &smem_a_0[load_smem_a_2_k_0]);
+    __ldmatrix_a_elems(&a_frag_3_k_0,                &smem_a_0[load_smem_a_3_k_0]);
+    __ldmatrix_b_elems(&b_frag_0_k_0, &b_frag_1_k_0, &smem_b_0[load_smem_b_0_k_0]);
+    __ldmatrix_b_elems(&b_frag_2_k_0, &b_frag_3_k_0, &smem_b_0[load_smem_b_1_k_0]);
+    __ldmatrix_b_elems(&b_frag_4_k_0, &b_frag_5_k_0, &smem_b_0[load_smem_b_2_k_0]);
+    __ldmatrix_b_elems(&b_frag_6_k_0, &b_frag_7_k_0, &smem_b_0[load_smem_b_3_k_0]);
+
+    for (int block_k = 0; block_k < num_k_blocks; block_k++) {
+        int phase_k = block_k % 3;
+        half *smem_a_curr = (phase_k == 0) ? smem_a_0 : ((phase_k == 1) ? smem_a_1 : smem_a_2);
+        half *smem_b_curr = (phase_k == 0) ? smem_b_0 : ((phase_k == 1) ? smem_b_1 : smem_b_2);
+
+        int next_phase_k = (block_k+1) % 3;
+        half *smem_a_next = (next_phase_k == 0) ? smem_a_0 : ((next_phase_k == 1) ? smem_a_1 : smem_a_2);
+        half *smem_b_next = (next_phase_k == 0) ? smem_b_0 : ((next_phase_k == 1) ? smem_b_1 : smem_b_2);
+
+        int store_phase_k = (block_k+2) % 3;
+        half *smem_a_store = (store_phase_k == 0) ? smem_a_0 : ((store_phase_k == 1) ? smem_a_1 : smem_a_2);
+        half *smem_b_store = (store_phase_k == 0) ? smem_b_0 : ((store_phase_k == 1) ? smem_b_1 : smem_b_2);
+
+        // load K=1 elements for the current tile
+        __ldmatrix_a_elems(&a_frag_0_k_1,                &smem_a_curr[load_smem_a_0_k_1]);
+        __ldmatrix_a_elems(&a_frag_1_k_1,                &smem_a_curr[load_smem_a_1_k_1]);
+        __ldmatrix_a_elems(&a_frag_2_k_1,                &smem_a_curr[load_smem_a_2_k_1]);
+        __ldmatrix_a_elems(&a_frag_3_k_1,                &smem_a_curr[load_smem_a_3_k_1]);
+        __ldmatrix_b_elems(&b_frag_0_k_1, &b_frag_1_k_1, &smem_b_curr[load_smem_b_0_k_1]);
+        __ldmatrix_b_elems(&b_frag_2_k_1, &b_frag_3_k_1, &smem_b_curr[load_smem_b_1_k_1]);
+        __ldmatrix_b_elems(&b_frag_4_k_1, &b_frag_5_k_1, &smem_b_curr[load_smem_b_2_k_1]);
+        __ldmatrix_b_elems(&b_frag_6_k_1, &b_frag_7_k_1, &smem_b_curr[load_smem_b_3_k_1]);
+
+        // MMA K=0, (M=4 x N=8)
+        acc_frag_0_0 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_0_k_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_1_k_0, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_2_k_0, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_3_k_0, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_4_k_0, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_5_k_0, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_6_k_0, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_7_k_0, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_0_k_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_1_k_0, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_2_k_0, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_3_k_0, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_4_k_0, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_5_k_0, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_6_k_0, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_7_k_0, acc_frag_1_7);
+        acc_frag_2_0 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_0_k_0, acc_frag_2_0);
+        acc_frag_2_1 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_1_k_0, acc_frag_2_1);
+        acc_frag_2_2 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_2_k_0, acc_frag_2_2);
+        acc_frag_2_3 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_3_k_0, acc_frag_2_3);
+        acc_frag_2_4 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_4_k_0, acc_frag_2_4);
+        acc_frag_2_5 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_5_k_0, acc_frag_2_5);
+        acc_frag_2_6 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_6_k_0, acc_frag_2_6);
+        acc_frag_2_7 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_7_k_0, acc_frag_2_7);
+        acc_frag_3_0 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_0_k_0, acc_frag_3_0);
+        acc_frag_3_1 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_1_k_0, acc_frag_3_1);
+        acc_frag_3_2 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_2_k_0, acc_frag_3_2);
+        acc_frag_3_3 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_3_k_0, acc_frag_3_3);
+        acc_frag_3_4 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_4_k_0, acc_frag_3_4);
+        acc_frag_3_5 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_5_k_0, acc_frag_3_5);
+        acc_frag_3_6 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_6_k_0, acc_frag_3_6);
+        acc_frag_3_7 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_7_k_0, acc_frag_3_7);
+
+        // load next tile
+        if (block_k < (num_k_blocks-2)) {
+            __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + (     0)], &data1[global_a_off + (    0)], 16);
+            __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + ( 64*32)], &data1[global_a_off + ( 64*K)], 16);
+            __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + (128*32)], &data1[global_a_off + (128*K)], 16);
+            __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + (192*32)], &data1[global_a_off + (192*K)], 16);
+            __pipeline_memcpy_async(&smem_b_store[store_smem_b_off + (     0)], &data2[global_b_off + (    0)], 16);
+            __pipeline_memcpy_async(&smem_b_store[store_smem_b_off + (16*128)], &data2[global_b_off + ( 16*N)], 16);
+            global_a_off += 32;
+            global_b_off += 32 * N;
+        }
+        __pipeline_commit();
+
+        // wait next tile
+        __pipeline_wait_prior(1);
+        __syncthreads();
+
+        // load K=0 for the next tile
+        __ldmatrix_a_elems(&a_frag_0_k_0,                &smem_a_next[load_smem_a_0_k_0]);
+        __ldmatrix_a_elems(&a_frag_1_k_0,                &smem_a_next[load_smem_a_1_k_0]);
+        __ldmatrix_a_elems(&a_frag_2_k_0,                &smem_a_next[load_smem_a_2_k_0]);
+        __ldmatrix_a_elems(&a_frag_3_k_0,                &smem_a_next[load_smem_a_3_k_0]);
+        __ldmatrix_b_elems(&b_frag_0_k_0, &b_frag_1_k_0, &smem_b_next[load_smem_b_0_k_0]);
+        __ldmatrix_b_elems(&b_frag_2_k_0, &b_frag_3_k_0, &smem_b_next[load_smem_b_1_k_0]);
+        __ldmatrix_b_elems(&b_frag_4_k_0, &b_frag_5_k_0, &smem_b_next[load_smem_b_2_k_0]);
+        __ldmatrix_b_elems(&b_frag_6_k_0, &b_frag_7_k_0, &smem_b_next[load_smem_b_3_k_0]);
+
+        // MMA K=1, (M=4 x N=8)
+        acc_frag_0_0 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_0_k_1, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_1_k_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_2_k_1, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_3_k_1, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_4_k_1, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_5_k_1, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_6_k_1, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_7_k_1, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_0_k_1, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_1_k_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_2_k_1, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_3_k_1, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_4_k_1, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_5_k_1, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_6_k_1, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_7_k_1, acc_frag_1_7);
+        acc_frag_2_0 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_0_k_1, acc_frag_2_0);
+        acc_frag_2_1 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_1_k_1, acc_frag_2_1);
+        acc_frag_2_2 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_2_k_1, acc_frag_2_2);
+        acc_frag_2_3 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_3_k_1, acc_frag_2_3);
+        acc_frag_2_4 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_4_k_1, acc_frag_2_4);
+        acc_frag_2_5 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_5_k_1, acc_frag_2_5);
+        acc_frag_2_6 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_6_k_1, acc_frag_2_6);
+        acc_frag_2_7 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_7_k_1, acc_frag_2_7);
+        acc_frag_3_0 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_0_k_1, acc_frag_3_0);
+        acc_frag_3_1 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_1_k_1, acc_frag_3_1);
+        acc_frag_3_2 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_2_k_1, acc_frag_3_2);
+        acc_frag_3_3 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_3_k_1, acc_frag_3_3);
+        acc_frag_3_4 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_4_k_1, acc_frag_3_4);
+        acc_frag_3_5 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_5_k_1, acc_frag_3_5);
+        acc_frag_3_6 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_6_k_1, acc_frag_3_6);
+        acc_frag_3_7 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_7_k_1, acc_frag_3_7);
+    }
+
+    // write accumulators to output
+    __pipeline_wait_prior(0);
+    __syncthreads();
+
+    // slower way: write accs one by one to data0
+
+    size_t wg_c_off     = ((grid_m * 256) * N) + (grid_n * 128) + (wg_m * 16 * N) + (wg_n * 16);
+    size_t thread_c_off = ((wg_threads % 4) * 2) + (((wg_threads / 4) % 8) * N);
+    data0[wg_c_off + thread_c_off           + 0 + ( 0*8)] = acc_frag_0_0.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 0*8)] = acc_frag_0_0.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 0*8)] = acc_frag_0_0.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 0*8)] = acc_frag_0_0.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 1*8)] = acc_frag_0_1.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 1*8)] = acc_frag_0_1.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 1*8)] = acc_frag_0_1.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 1*8)] = acc_frag_0_1.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 4*8)] = acc_frag_0_2.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 4*8)] = acc_frag_0_2.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 4*8)] = acc_frag_0_2.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 4*8)] = acc_frag_0_2.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 5*8)] = acc_frag_0_3.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 5*8)] = acc_frag_0_3.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 5*8)] = acc_frag_0_3.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 5*8)] = acc_frag_0_3.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 8*8)] = acc_frag_0_4.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 8*8)] = acc_frag_0_4.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 8*8)] = acc_frag_0_4.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 8*8)] = acc_frag_0_4.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 9*8)] = acc_frag_0_5.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 9*8)] = acc_frag_0_5.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 9*8)] = acc_frag_0_5.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 9*8)] = acc_frag_0_5.w;
+    data0[wg_c_off + thread_c_off           + 0 + (12*8)] = acc_frag_0_6.x;
+    data0[wg_c_off + thread_c_off           + 1 + (12*8)] = acc_frag_0_6.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (12*8)] = acc_frag_0_6.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (12*8)] = acc_frag_0_6.w;
+    data0[wg_c_off + thread_c_off           + 0 + (13*8)] = acc_frag_0_7.x;
+    data0[wg_c_off + thread_c_off           + 1 + (13*8)] = acc_frag_0_7.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (13*8)] = acc_frag_0_7.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (13*8)] = acc_frag_0_7.w;
+
+    wg_c_off += 64*N;
+    data0[wg_c_off + thread_c_off           + 0 + ( 0*8)] = acc_frag_1_0.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 0*8)] = acc_frag_1_0.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 0*8)] = acc_frag_1_0.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 0*8)] = acc_frag_1_0.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 1*8)] = acc_frag_1_1.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 1*8)] = acc_frag_1_1.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 1*8)] = acc_frag_1_1.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 1*8)] = acc_frag_1_1.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 4*8)] = acc_frag_1_2.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 4*8)] = acc_frag_1_2.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 4*8)] = acc_frag_1_2.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 4*8)] = acc_frag_1_2.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 5*8)] = acc_frag_1_3.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 5*8)] = acc_frag_1_3.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 5*8)] = acc_frag_1_3.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 5*8)] = acc_frag_1_3.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 8*8)] = acc_frag_1_4.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 8*8)] = acc_frag_1_4.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 8*8)] = acc_frag_1_4.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 8*8)] = acc_frag_1_4.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 9*8)] = acc_frag_1_5.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 9*8)] = acc_frag_1_5.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 9*8)] = acc_frag_1_5.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 9*8)] = acc_frag_1_5.w;
+    data0[wg_c_off + thread_c_off           + 0 + (12*8)] = acc_frag_1_6.x;
+    data0[wg_c_off + thread_c_off           + 1 + (12*8)] = acc_frag_1_6.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (12*8)] = acc_frag_1_6.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (12*8)] = acc_frag_1_6.w;
+    data0[wg_c_off + thread_c_off           + 0 + (13*8)] = acc_frag_1_7.x;
+    data0[wg_c_off + thread_c_off           + 1 + (13*8)] = acc_frag_1_7.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (13*8)] = acc_frag_1_7.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (13*8)] = acc_frag_1_7.w;
+
+    wg_c_off += 64*N;
+    data0[wg_c_off + thread_c_off           + 0 + ( 0*8)] = acc_frag_2_0.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 0*8)] = acc_frag_2_0.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 0*8)] = acc_frag_2_0.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 0*8)] = acc_frag_2_0.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 1*8)] = acc_frag_2_1.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 1*8)] = acc_frag_2_1.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 1*8)] = acc_frag_2_1.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 1*8)] = acc_frag_2_1.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 4*8)] = acc_frag_2_2.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 4*8)] = acc_frag_2_2.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 4*8)] = acc_frag_2_2.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 4*8)] = acc_frag_2_2.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 5*8)] = acc_frag_2_3.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 5*8)] = acc_frag_2_3.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 5*8)] = acc_frag_2_3.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 5*8)] = acc_frag_2_3.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 8*8)] = acc_frag_2_4.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 8*8)] = acc_frag_2_4.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 8*8)] = acc_frag_2_4.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 8*8)] = acc_frag_2_4.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 9*8)] = acc_frag_2_5.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 9*8)] = acc_frag_2_5.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 9*8)] = acc_frag_2_5.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 9*8)] = acc_frag_2_5.w;
+    data0[wg_c_off + thread_c_off           + 0 + (12*8)] = acc_frag_2_6.x;
+    data0[wg_c_off + thread_c_off           + 1 + (12*8)] = acc_frag_2_6.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (12*8)] = acc_frag_2_6.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (12*8)] = acc_frag_2_6.w;
+    data0[wg_c_off + thread_c_off           + 0 + (13*8)] = acc_frag_2_7.x;
+    data0[wg_c_off + thread_c_off           + 1 + (13*8)] = acc_frag_2_7.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (13*8)] = acc_frag_2_7.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (13*8)] = acc_frag_2_7.w;
+
+    wg_c_off += 64*N;
+    data0[wg_c_off + thread_c_off           + 0 + ( 0*8)] = acc_frag_3_0.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 0*8)] = acc_frag_3_0.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 0*8)] = acc_frag_3_0.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 0*8)] = acc_frag_3_0.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 1*8)] = acc_frag_3_1.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 1*8)] = acc_frag_3_1.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 1*8)] = acc_frag_3_1.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 1*8)] = acc_frag_3_1.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 4*8)] = acc_frag_3_2.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 4*8)] = acc_frag_3_2.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 4*8)] = acc_frag_3_2.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 4*8)] = acc_frag_3_2.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 5*8)] = acc_frag_3_3.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 5*8)] = acc_frag_3_3.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 5*8)] = acc_frag_3_3.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 5*8)] = acc_frag_3_3.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 8*8)] = acc_frag_3_4.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 8*8)] = acc_frag_3_4.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 8*8)] = acc_frag_3_4.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 8*8)] = acc_frag_3_4.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 9*8)] = acc_frag_3_5.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 9*8)] = acc_frag_3_5.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 9*8)] = acc_frag_3_5.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 9*8)] = acc_frag_3_5.w;
+    data0[wg_c_off + thread_c_off           + 0 + (12*8)] = acc_frag_3_6.x;
+    data0[wg_c_off + thread_c_off           + 1 + (12*8)] = acc_frag_3_6.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (12*8)] = acc_frag_3_6.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (12*8)] = acc_frag_3_6.w;
+    data0[wg_c_off + thread_c_off           + 0 + (13*8)] = acc_frag_3_7.x;
+    data0[wg_c_off + thread_c_off           + 1 + (13*8)] = acc_frag_3_7.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (13*8)] = acc_frag_3_7.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (13*8)] = acc_frag_3_7.w;
+
+}

tinygrad_repo/extra/gemm/max_kernels/nv.fp16_fp16_fp16.3_stage_swizzled.cu

@@ -0,0 +1,517 @@
+#define INFINITY (__int_as_float(0x7f800000))
+#define NAN (__int_as_float(0x7fffffff))
+#include <cuda_fp16.h>
+#include <cuda_pipeline.h>
+#define N_PAD 132
+
+struct __align__(8) half4 { half x, y, z, w; };
+__device__ half4 make_half4(half x, half y, half z, half w) { half4 r={x, y, z, w}; return r; }
+
+struct __align__(16) half8 { half x, y, z, w, a, b, c, d; };
+__device__ half8 make_half8(half x, half y, half z, half w, half a, half b, half c, half d) { half8 r={x, y, z, w, a, b, c, d}; return r; }
+
+__device__ void __ldmatrix_a_elems(half8 *regs, half *smem) {
+    uint32_t reg0, reg1, reg2, reg3;
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(reg0), "=r"(reg1), "=r"(reg2), "=r"(reg3)
+        : "l"(__cvta_generic_to_shared(smem))
+    );
+    uint32_t *addr = reinterpret_cast<uint32_t*>(regs);
+    addr[0] = reg0;
+    addr[1] = reg1;
+    addr[2] = reg2;
+    addr[3] = reg3;
+}
+
+__device__ void __ldmatrix_b_elems(half4 *regs_lo, half4 *regs_hi, half *smem) {
+    uint32_t reg0, reg1, reg2, reg3;
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(reg0), "=r"(reg1), "=r"(reg2), "=r"(reg3)
+        : "l"(__cvta_generic_to_shared(smem))
+    );
+    uint32_t *addr_lo = reinterpret_cast<uint32_t*>(regs_lo);
+    uint32_t *addr_hi = reinterpret_cast<uint32_t*>(regs_hi);
+    addr_lo[0] = reg0;
+    addr_lo[1] = reg1;
+    addr_hi[0] = reg2;
+    addr_hi[1] = reg3;
+}
+
+__device__ half4 __WMMA_8_16_16_half_half(half8 a, half4 b, half4 c) {
+    int *a_pk = (int *) (&a), *b_pk = (int *) (&b), *c_pk = (int *) (&c);
+    asm( "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 { %0, %1 }, { %2, %3, %4, %5 }, { %6, %7 }, { %0, %1 };"
+        : "+r"(c_pk[0]), "+r"(c_pk[1]): "r"(a_pk[0]), "r"(a_pk[1]), "r"(a_pk[2]),  "r"(a_pk[3]), "r"(b_pk[0]), "r"(b_pk[1]) );
+    return c;
+}
+
+extern "C" __global__ void __launch_bounds__(256) wmma_example(half* data0, const half* data1, const half* data2, int N, int K) {
+    extern __shared__ char smem[];
+    half *smem_a_0 = (half *)(smem);
+    half *smem_a_1 = (half *)(smem + 16384);
+    half *smem_a_2 = (half *)(smem + 32768);
+    half *smem_b_0 = (half *)(smem + 49152);
+    half *smem_b_1 = (half *)(smem + 57344);
+    half *smem_b_2 = (half *)(smem + 65536);
+
+    int grid_m = blockIdx.x;        /* M//256 */
+    int grid_n = blockIdx.y;        /* N//128 */
+    int wg_threads = threadIdx.x;   // 32
+    int wg_m = threadIdx.y;         // 4
+    int wg_n = threadIdx.z;         // 2
+    int threads = threadIdx.x + (threadIdx.y * 32) + (threadIdx.z * 128); /* 256 */
+    int num_k_blocks = K / 32;
+
+    // ldmatrix indices
+    // threads 0-7 are row starts for A, 8-15 for B, 16-23 for C, 24-31 for D
+    // [ A | C ]
+    // [ - + - ]
+    // [ B | D ]
+
+    // unswizzled A - SMEM_A is 256 rows x 32 cols
+    // size_t global_a_off = ((grid_m * 256) * K) + ((threads %  4) * 8) + ((threads /  4) * K);
+    // size_t store_smem_a_off = ((threads %  4) * 8) + ((threads /  4) * 32); // 64 rows /  32 cols per copy
+    // size_t load_smem_a_0_k_0 = (wg_m * 16 * 32) + ((wg_threads % 16) * 32) + ((wg_threads / 16) * 8);
+    // size_t load_smem_a_1_k_0 = load_smem_a_0_k_0 + ( 64 * 32);
+    // size_t load_smem_a_2_k_0 = load_smem_a_0_k_0 + (128 * 32);
+    // size_t load_smem_a_3_k_0 = load_smem_a_0_k_0 + (192 * 32);
+    // size_t load_smem_a_0_k_1 = load_smem_a_0_k_0 + 16;
+    // size_t load_smem_a_1_k_1 = load_smem_a_0_k_1 + ( 64 * 32);
+    // size_t load_smem_a_2_k_1 = load_smem_a_0_k_1 + (128 * 32);
+    // size_t load_smem_a_3_k_1 = load_smem_a_0_k_1 + (192 * 32);
+
+    // unswizzled reshaped A - SMEM_A is 128 rows x 64 cols, [ (M=0, K=0), (M=0, K=1), (M=8, K=0), (M=8, K=1) ], etc.
+    // size_t global_a_off = ((grid_m * 256) * K) + ((threads %  4) * 8) + (((threads /  4) % 2) * 8 * 16 * K) + ((threads / 8) * K);
+    // size_t store_smem_a_off = ((threads %  8) * 8) + ((threads /  8) * 64); // 32 rows / 64 cols per copy
+    // size_t load_smem_a_0_k_0 = (wg_m * 16 * 64) + ((wg_threads % 16) * 64) + ((wg_threads / 16) * 8);
+    // size_t load_smem_a_1_k_0 = load_smem_a_0_k_0 + (64 * 64);
+    // size_t load_smem_a_2_k_0 = load_smem_a_0_k_0 +           + 32;
+    // size_t load_smem_a_3_k_0 = load_smem_a_0_k_0 + (64 * 64) + 32;
+    // size_t load_smem_a_0_k_1 = load_smem_a_0_k_0 + 16;
+    // size_t load_smem_a_1_k_1 = load_smem_a_1_k_0 + 16;
+    // size_t load_smem_a_2_k_1 = load_smem_a_2_k_0 + 16;
+    // size_t load_smem_a_3_k_1 = load_smem_a_3_k_0 + 16;
+
+    // swizzled A
+    size_t global_a_off = ((grid_m * 256) * K) + ((threads %  4) * 8) + (((threads /  4) % 2) * 8 * 16 * K) + ((threads / 8) * K);
+    size_t store_smem_a_off  = ((threads /  8) *  64) + (((threads * 8) ^ threads) & 56); // 32 rows / 64 cols per copy
+    size_t load_smem_a_row   = ((wg_m * 16) + (threads % 16)) * 64;
+    size_t load_smem_a_phase = (threads / 16) % 2;
+    size_t load_smem_a_0_k_0 = load_smem_a_row + ( 0 * 64) + (((load_smem_a_phase + 0) ^ (threads % 8)) * 8);
+    size_t load_smem_a_1_k_0 = load_smem_a_row + (64 * 64) + (((load_smem_a_phase + 0) ^ (threads % 8)) * 8);
+    size_t load_smem_a_2_k_0 = load_smem_a_row + ( 0 * 64) + (((load_smem_a_phase + 4) ^ (threads % 8)) * 8);
+    size_t load_smem_a_3_k_0 = load_smem_a_row + (64 * 64) + (((load_smem_a_phase + 4) ^ (threads % 8)) * 8);
+    size_t load_smem_a_0_k_1 = load_smem_a_row + ( 0 * 64) + (((load_smem_a_phase + 2) ^ (threads % 8)) * 8);
+    size_t load_smem_a_1_k_1 = load_smem_a_row + (64 * 64) + (((load_smem_a_phase + 2) ^ (threads % 8)) * 8);
+    size_t load_smem_a_2_k_1 = load_smem_a_row + ( 0 * 64) + (((load_smem_a_phase + 6) ^ (threads % 8)) * 8);
+    size_t load_smem_a_3_k_1 = load_smem_a_row + (64 * 64) + (((load_smem_a_phase + 6) ^ (threads % 8)) * 8);
+
+    // unswizzed B
+    // size_t global_b_off = (grid_n * 128) + ((threads % 16) * 8) + ((threads / 16) * N);
+    // size_t store_smem_b_off = ((threads % 16) * 8) + ((threads / 16) * 128); // 16 rows / 128 cols per copy
+    // size_t load_smem_b_0_k_0 = (wg_n * 16) + ((wg_threads % 16) * 128) + ((wg_threads / 16) * 8);
+    // size_t load_smem_b_1_k_0 = load_smem_b_0_k_0 + 32;
+    // size_t load_smem_b_2_k_0 = load_smem_b_0_k_0 + 64;
+    // size_t load_smem_b_3_k_0 = load_smem_b_0_k_0 + 96;
+    // size_t load_smem_b_0_k_1 = load_smem_b_0_k_0 + (16 * 128);
+    // size_t load_smem_b_1_k_1 = load_smem_b_0_k_1 + 32;
+    // size_t load_smem_b_2_k_1 = load_smem_b_0_k_1 + 64;
+    // size_t load_smem_b_3_k_1 = load_smem_b_0_k_1 + 96;
+
+    // swizzled B
+    size_t global_b_off = (grid_n * 128) + ((threads % 16) * 8) + ((threads / 16) * N);
+    size_t store_smem_b_off  = ((threads / 16) * 128) + ((((threads / 16) % 8) * 8) ^ ((threads % 16) * 8)); // 16 rows / 128 cols per copy
+    size_t load_smem_b_row   = (threads % 16) * 128;
+    size_t load_smem_b_phase = (wg_n * 2) + (wg_threads / 16);
+    size_t load_smem_b_0_k_0 = load_smem_b_row + (((load_smem_b_phase +  0) ^ (threads % 8)) * 8);
+    size_t load_smem_b_1_k_0 = load_smem_b_row + (((load_smem_b_phase +  4) ^ (threads % 8)) * 8);
+    size_t load_smem_b_2_k_0 = load_smem_b_row + (((load_smem_b_phase +  8) ^ (threads % 8)) * 8);
+    size_t load_smem_b_3_k_0 = load_smem_b_row + (((load_smem_b_phase + 12) ^ (threads % 8)) * 8);
+    size_t load_smem_b_0_k_1 = load_smem_b_0_k_0 + (16 * 128);
+    size_t load_smem_b_1_k_1 = load_smem_b_1_k_0 + (16 * 128);
+    size_t load_smem_b_2_k_1 = load_smem_b_2_k_0 + (16 * 128);
+    size_t load_smem_b_3_k_1 = load_smem_b_3_k_0 + (16 * 128);
+
+    // create accs (M=4, N=8)
+    half4 acc_frag_0_0 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_1 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_2 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_3 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_4 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_5 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_6 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_7 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_0 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_1 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_2 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_3 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_4 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_5 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_6 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_7 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_0 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_1 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_2 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_3 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_4 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_5 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_6 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_7 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_0 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_1 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_2 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_3 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_4 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_5 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_6 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_7 = make_half4(0.0f,0.0f,0.0f,0.0f);
+
+    // create registers for block A elements
+    half8 a_frag_0_k_0;
+    half8 a_frag_1_k_0;
+    half8 a_frag_2_k_0;
+    half8 a_frag_3_k_0;
+    half8 a_frag_0_k_1;
+    half8 a_frag_1_k_1;
+    half8 a_frag_2_k_1;
+    half8 a_frag_3_k_1;
+
+    // create register for block B elements
+    half4 b_frag_0_k_0;
+    half4 b_frag_1_k_0;
+    half4 b_frag_2_k_0;
+    half4 b_frag_3_k_0;
+    half4 b_frag_4_k_0;
+    half4 b_frag_5_k_0;
+    half4 b_frag_6_k_0;
+    half4 b_frag_7_k_0;
+    half4 b_frag_0_k_1;
+    half4 b_frag_1_k_1;
+    half4 b_frag_2_k_1;
+    half4 b_frag_3_k_1;
+    half4 b_frag_4_k_1;
+    half4 b_frag_5_k_1;
+    half4 b_frag_6_k_1;
+    half4 b_frag_7_k_1;
+
+    __syncthreads();
+
+    // load first tile
+    // unswizzled 256 x 32
+    // __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + (     0)], &data1[global_a_off + (    0)], 16);
+    // __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + ( 64*32)], &data1[global_a_off + ( 64*K)], 16);
+    // __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + (128*32)], &data1[global_a_off + (128*K)], 16);
+    // __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + (192*32)], &data1[global_a_off + (192*K)], 16);
+    // unswizzled 128 x 64
+    __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + (     0)], &data1[global_a_off + (    0)], 16);
+    __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + ( 32*64)], &data1[global_a_off + ( 32*K)], 16);
+    __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + ( 64*64)], &data1[global_a_off + ( 64*K)], 16);
+    __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + ( 96*64)], &data1[global_a_off + ( 96*K)], 16);
+    __pipeline_memcpy_async(&smem_b_0[store_smem_b_off + (     0)], &data2[global_b_off + (    0)], 16);
+    __pipeline_memcpy_async(&smem_b_0[store_smem_b_off + (16*128)], &data2[global_b_off + ( 16*N)], 16);
+    __pipeline_commit();
+    global_a_off += 32;
+    global_b_off += 32 * N;
+
+    // load second tile
+    // unswizzled 256 x 32
+    // __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + (     0)], &data1[global_a_off + (    0)], 16);
+    // __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + ( 64*32)], &data1[global_a_off + ( 64*K)], 16);
+    // __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + (128*32)], &data1[global_a_off + (128*K)], 16);
+    // __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + (192*32)], &data1[global_a_off + (192*K)], 16);
+    // unswizzled 128 x 64
+    __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + (     0)], &data1[global_a_off + (    0)], 16);
+    __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + ( 32*64)], &data1[global_a_off + ( 32*K)], 16);
+    __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + ( 64*64)], &data1[global_a_off + ( 64*K)], 16);
+    __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + ( 96*64)], &data1[global_a_off + ( 96*K)], 16);
+    __pipeline_memcpy_async(&smem_b_1[store_smem_b_off + (     0)], &data2[global_b_off + (    0)], 16);
+    __pipeline_memcpy_async(&smem_b_1[store_smem_b_off + (16*128)], &data2[global_b_off + ( 16*N)], 16);
+    __pipeline_commit();
+
+    global_a_off += 32;
+    global_b_off += 32 * N;
+
+    // wait on first pre-fetch load
+    __pipeline_wait_prior(1);
+    __syncthreads();
+
+    // load K=0 for the first tile
+    __ldmatrix_a_elems(&a_frag_0_k_0,                &smem_a_0[load_smem_a_0_k_0]);
+    __ldmatrix_a_elems(&a_frag_1_k_0,                &smem_a_0[load_smem_a_1_k_0]);
+    __ldmatrix_a_elems(&a_frag_2_k_0,                &smem_a_0[load_smem_a_2_k_0]);
+    __ldmatrix_a_elems(&a_frag_3_k_0,                &smem_a_0[load_smem_a_3_k_0]);
+    __ldmatrix_b_elems(&b_frag_0_k_0, &b_frag_1_k_0, &smem_b_0[load_smem_b_0_k_0]);
+    __ldmatrix_b_elems(&b_frag_2_k_0, &b_frag_3_k_0, &smem_b_0[load_smem_b_1_k_0]);
+    __ldmatrix_b_elems(&b_frag_4_k_0, &b_frag_5_k_0, &smem_b_0[load_smem_b_2_k_0]);
+    __ldmatrix_b_elems(&b_frag_6_k_0, &b_frag_7_k_0, &smem_b_0[load_smem_b_3_k_0]);
+
+    for (int block_k = 0; block_k < num_k_blocks; block_k++) {
+        int phase_k = block_k % 3;
+        half *smem_a_curr = (phase_k == 0) ? smem_a_0 : ((phase_k == 1) ? smem_a_1 : smem_a_2);
+        half *smem_b_curr = (phase_k == 0) ? smem_b_0 : ((phase_k == 1) ? smem_b_1 : smem_b_2);
+
+        int next_phase_k = (block_k+1) % 3;
+        half *smem_a_next = (next_phase_k == 0) ? smem_a_0 : ((next_phase_k == 1) ? smem_a_1 : smem_a_2);
+        half *smem_b_next = (next_phase_k == 0) ? smem_b_0 : ((next_phase_k == 1) ? smem_b_1 : smem_b_2);
+
+        int store_phase_k = (block_k+2) % 3;
+        half *smem_a_store = (store_phase_k == 0) ? smem_a_0 : ((store_phase_k == 1) ? smem_a_1 : smem_a_2);
+        half *smem_b_store = (store_phase_k == 0) ? smem_b_0 : ((store_phase_k == 1) ? smem_b_1 : smem_b_2);
+
+        // load K=1 elements for the current tile
+        __ldmatrix_a_elems(&a_frag_0_k_1,                &smem_a_curr[load_smem_a_0_k_1]);
+        __ldmatrix_a_elems(&a_frag_1_k_1,                &smem_a_curr[load_smem_a_1_k_1]);
+        __ldmatrix_a_elems(&a_frag_2_k_1,                &smem_a_curr[load_smem_a_2_k_1]);
+        __ldmatrix_a_elems(&a_frag_3_k_1,                &smem_a_curr[load_smem_a_3_k_1]);
+        __ldmatrix_b_elems(&b_frag_0_k_1, &b_frag_1_k_1, &smem_b_curr[load_smem_b_0_k_1]);
+        __ldmatrix_b_elems(&b_frag_2_k_1, &b_frag_3_k_1, &smem_b_curr[load_smem_b_1_k_1]);
+        __ldmatrix_b_elems(&b_frag_4_k_1, &b_frag_5_k_1, &smem_b_curr[load_smem_b_2_k_1]);
+        __ldmatrix_b_elems(&b_frag_6_k_1, &b_frag_7_k_1, &smem_b_curr[load_smem_b_3_k_1]);
+
+        // MMA K=0, (M=4 x N=8)
+        acc_frag_0_0 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_0_k_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_1_k_0, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_2_k_0, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_3_k_0, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_4_k_0, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_5_k_0, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_6_k_0, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_7_k_0, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_0_k_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_1_k_0, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_2_k_0, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_3_k_0, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_4_k_0, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_5_k_0, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_6_k_0, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_7_k_0, acc_frag_1_7);
+        acc_frag_2_0 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_0_k_0, acc_frag_2_0);
+        acc_frag_2_1 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_1_k_0, acc_frag_2_1);
+        acc_frag_2_2 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_2_k_0, acc_frag_2_2);
+        acc_frag_2_3 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_3_k_0, acc_frag_2_3);
+        acc_frag_2_4 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_4_k_0, acc_frag_2_4);
+        acc_frag_2_5 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_5_k_0, acc_frag_2_5);
+        acc_frag_2_6 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_6_k_0, acc_frag_2_6);
+        acc_frag_2_7 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_7_k_0, acc_frag_2_7);
+        acc_frag_3_0 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_0_k_0, acc_frag_3_0);
+        acc_frag_3_1 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_1_k_0, acc_frag_3_1);
+        acc_frag_3_2 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_2_k_0, acc_frag_3_2);
+        acc_frag_3_3 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_3_k_0, acc_frag_3_3);
+        acc_frag_3_4 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_4_k_0, acc_frag_3_4);
+        acc_frag_3_5 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_5_k_0, acc_frag_3_5);
+        acc_frag_3_6 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_6_k_0, acc_frag_3_6);
+        acc_frag_3_7 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_7_k_0, acc_frag_3_7);
+
+        // load next tile
+        if (block_k < (num_k_blocks-2)) {
+            // unswizzled 256 x 32
+            // __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + (     0)], &data1[global_a_off + (    0)], 16);
+            // __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + ( 64*32)], &data1[global_a_off + ( 64*K)], 16);
+            // __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + (128*32)], &data1[global_a_off + (128*K)], 16);
+            // __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + (192*32)], &data1[global_a_off + (192*K)], 16);
+            // unswizzled 128 x 64
+            __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + (     0)], &data1[global_a_off + (    0)], 16);
+            __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + ( 32*64)], &data1[global_a_off + ( 32*K)], 16);
+            __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + ( 64*64)], &data1[global_a_off + ( 64*K)], 16);
+            __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + ( 96*64)], &data1[global_a_off + ( 96*K)], 16);
+            __pipeline_memcpy_async(&smem_b_store[store_smem_b_off + (     0)], &data2[global_b_off + (    0)], 16);
+            __pipeline_memcpy_async(&smem_b_store[store_smem_b_off + (16*128)], &data2[global_b_off + ( 16*N)], 16);
+            global_a_off += 32;
+            global_b_off += 32 * N;
+        }
+        __pipeline_commit();
+
+        // wait next tile
+        __pipeline_wait_prior(1);
+        __syncthreads();
+
+        // load K=0 for the next tile
+        __ldmatrix_a_elems(&a_frag_0_k_0,                &smem_a_next[load_smem_a_0_k_0]);
+        __ldmatrix_a_elems(&a_frag_1_k_0,                &smem_a_next[load_smem_a_1_k_0]);
+        __ldmatrix_a_elems(&a_frag_2_k_0,                &smem_a_next[load_smem_a_2_k_0]);
+        __ldmatrix_a_elems(&a_frag_3_k_0,                &smem_a_next[load_smem_a_3_k_0]);
+        __ldmatrix_b_elems(&b_frag_0_k_0, &b_frag_1_k_0, &smem_b_next[load_smem_b_0_k_0]);
+        __ldmatrix_b_elems(&b_frag_2_k_0, &b_frag_3_k_0, &smem_b_next[load_smem_b_1_k_0]);
+        __ldmatrix_b_elems(&b_frag_4_k_0, &b_frag_5_k_0, &smem_b_next[load_smem_b_2_k_0]);
+        __ldmatrix_b_elems(&b_frag_6_k_0, &b_frag_7_k_0, &smem_b_next[load_smem_b_3_k_0]);
+
+        // MMA K=1, (M=4 x N=8)
+        acc_frag_0_0 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_0_k_1, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_1_k_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_2_k_1, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_3_k_1, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_4_k_1, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_5_k_1, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_6_k_1, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_7_k_1, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_0_k_1, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_1_k_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_2_k_1, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_3_k_1, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_4_k_1, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_5_k_1, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_6_k_1, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_7_k_1, acc_frag_1_7);
+        acc_frag_2_0 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_0_k_1, acc_frag_2_0);
+        acc_frag_2_1 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_1_k_1, acc_frag_2_1);
+        acc_frag_2_2 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_2_k_1, acc_frag_2_2);
+        acc_frag_2_3 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_3_k_1, acc_frag_2_3);
+        acc_frag_2_4 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_4_k_1, acc_frag_2_4);
+        acc_frag_2_5 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_5_k_1, acc_frag_2_5);
+        acc_frag_2_6 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_6_k_1, acc_frag_2_6);
+        acc_frag_2_7 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_7_k_1, acc_frag_2_7);
+        acc_frag_3_0 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_0_k_1, acc_frag_3_0);
+        acc_frag_3_1 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_1_k_1, acc_frag_3_1);
+        acc_frag_3_2 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_2_k_1, acc_frag_3_2);
+        acc_frag_3_3 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_3_k_1, acc_frag_3_3);
+        acc_frag_3_4 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_4_k_1, acc_frag_3_4);
+        acc_frag_3_5 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_5_k_1, acc_frag_3_5);
+        acc_frag_3_6 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_6_k_1, acc_frag_3_6);
+        acc_frag_3_7 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_7_k_1, acc_frag_3_7);
+    }
+
+    // write accumulators to output
+    __pipeline_wait_prior(0);
+    __syncthreads();
+
+    // slower way: write accs one by one to data0
+
+    size_t wg_c_off     = ((grid_m * 256) * N) + (grid_n * 128) + (wg_m * 16 * N) + (wg_n * 16);
+    size_t thread_c_off = ((wg_threads % 4) * 2) + (((wg_threads / 4) % 8) * N);
+    data0[wg_c_off + thread_c_off           + 0 + ( 0*8)] = acc_frag_0_0.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 0*8)] = acc_frag_0_0.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 0*8)] = acc_frag_0_0.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 0*8)] = acc_frag_0_0.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 1*8)] = acc_frag_0_1.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 1*8)] = acc_frag_0_1.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 1*8)] = acc_frag_0_1.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 1*8)] = acc_frag_0_1.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 4*8)] = acc_frag_0_2.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 4*8)] = acc_frag_0_2.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 4*8)] = acc_frag_0_2.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 4*8)] = acc_frag_0_2.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 5*8)] = acc_frag_0_3.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 5*8)] = acc_frag_0_3.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 5*8)] = acc_frag_0_3.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 5*8)] = acc_frag_0_3.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 8*8)] = acc_frag_0_4.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 8*8)] = acc_frag_0_4.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 8*8)] = acc_frag_0_4.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 8*8)] = acc_frag_0_4.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 9*8)] = acc_frag_0_5.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 9*8)] = acc_frag_0_5.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 9*8)] = acc_frag_0_5.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 9*8)] = acc_frag_0_5.w;
+    data0[wg_c_off + thread_c_off           + 0 + (12*8)] = acc_frag_0_6.x;
+    data0[wg_c_off + thread_c_off           + 1 + (12*8)] = acc_frag_0_6.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (12*8)] = acc_frag_0_6.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (12*8)] = acc_frag_0_6.w;
+    data0[wg_c_off + thread_c_off           + 0 + (13*8)] = acc_frag_0_7.x;
+    data0[wg_c_off + thread_c_off           + 1 + (13*8)] = acc_frag_0_7.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (13*8)] = acc_frag_0_7.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (13*8)] = acc_frag_0_7.w;
+
+    wg_c_off += 64*N;
+    data0[wg_c_off + thread_c_off           + 0 + ( 0*8)] = acc_frag_1_0.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 0*8)] = acc_frag_1_0.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 0*8)] = acc_frag_1_0.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 0*8)] = acc_frag_1_0.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 1*8)] = acc_frag_1_1.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 1*8)] = acc_frag_1_1.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 1*8)] = acc_frag_1_1.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 1*8)] = acc_frag_1_1.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 4*8)] = acc_frag_1_2.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 4*8)] = acc_frag_1_2.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 4*8)] = acc_frag_1_2.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 4*8)] = acc_frag_1_2.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 5*8)] = acc_frag_1_3.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 5*8)] = acc_frag_1_3.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 5*8)] = acc_frag_1_3.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 5*8)] = acc_frag_1_3.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 8*8)] = acc_frag_1_4.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 8*8)] = acc_frag_1_4.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 8*8)] = acc_frag_1_4.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 8*8)] = acc_frag_1_4.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 9*8)] = acc_frag_1_5.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 9*8)] = acc_frag_1_5.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 9*8)] = acc_frag_1_5.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 9*8)] = acc_frag_1_5.w;
+    data0[wg_c_off + thread_c_off           + 0 + (12*8)] = acc_frag_1_6.x;
+    data0[wg_c_off + thread_c_off           + 1 + (12*8)] = acc_frag_1_6.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (12*8)] = acc_frag_1_6.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (12*8)] = acc_frag_1_6.w;
+    data0[wg_c_off + thread_c_off           + 0 + (13*8)] = acc_frag_1_7.x;
+    data0[wg_c_off + thread_c_off           + 1 + (13*8)] = acc_frag_1_7.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (13*8)] = acc_frag_1_7.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (13*8)] = acc_frag_1_7.w;
+
+    wg_c_off += 64*N;
+    data0[wg_c_off + thread_c_off           + 0 + ( 0*8)] = acc_frag_2_0.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 0*8)] = acc_frag_2_0.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 0*8)] = acc_frag_2_0.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 0*8)] = acc_frag_2_0.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 1*8)] = acc_frag_2_1.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 1*8)] = acc_frag_2_1.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 1*8)] = acc_frag_2_1.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 1*8)] = acc_frag_2_1.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 4*8)] = acc_frag_2_2.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 4*8)] = acc_frag_2_2.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 4*8)] = acc_frag_2_2.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 4*8)] = acc_frag_2_2.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 5*8)] = acc_frag_2_3.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 5*8)] = acc_frag_2_3.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 5*8)] = acc_frag_2_3.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 5*8)] = acc_frag_2_3.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 8*8)] = acc_frag_2_4.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 8*8)] = acc_frag_2_4.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 8*8)] = acc_frag_2_4.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 8*8)] = acc_frag_2_4.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 9*8)] = acc_frag_2_5.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 9*8)] = acc_frag_2_5.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 9*8)] = acc_frag_2_5.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 9*8)] = acc_frag_2_5.w;
+    data0[wg_c_off + thread_c_off           + 0 + (12*8)] = acc_frag_2_6.x;
+    data0[wg_c_off + thread_c_off           + 1 + (12*8)] = acc_frag_2_6.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (12*8)] = acc_frag_2_6.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (12*8)] = acc_frag_2_6.w;
+    data0[wg_c_off + thread_c_off           + 0 + (13*8)] = acc_frag_2_7.x;
+    data0[wg_c_off + thread_c_off           + 1 + (13*8)] = acc_frag_2_7.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (13*8)] = acc_frag_2_7.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (13*8)] = acc_frag_2_7.w;
+
+    wg_c_off += 64*N;
+    data0[wg_c_off + thread_c_off           + 0 + ( 0*8)] = acc_frag_3_0.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 0*8)] = acc_frag_3_0.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 0*8)] = acc_frag_3_0.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 0*8)] = acc_frag_3_0.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 1*8)] = acc_frag_3_1.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 1*8)] = acc_frag_3_1.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 1*8)] = acc_frag_3_1.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 1*8)] = acc_frag_3_1.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 4*8)] = acc_frag_3_2.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 4*8)] = acc_frag_3_2.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 4*8)] = acc_frag_3_2.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 4*8)] = acc_frag_3_2.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 5*8)] = acc_frag_3_3.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 5*8)] = acc_frag_3_3.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 5*8)] = acc_frag_3_3.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 5*8)] = acc_frag_3_3.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 8*8)] = acc_frag_3_4.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 8*8)] = acc_frag_3_4.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 8*8)] = acc_frag_3_4.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 8*8)] = acc_frag_3_4.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 9*8)] = acc_frag_3_5.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 9*8)] = acc_frag_3_5.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 9*8)] = acc_frag_3_5.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 9*8)] = acc_frag_3_5.w;
+    data0[wg_c_off + thread_c_off           + 0 + (12*8)] = acc_frag_3_6.x;
+    data0[wg_c_off + thread_c_off           + 1 + (12*8)] = acc_frag_3_6.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (12*8)] = acc_frag_3_6.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (12*8)] = acc_frag_3_6.w;
+    data0[wg_c_off + thread_c_off           + 0 + (13*8)] = acc_frag_3_7.x;
+    data0[wg_c_off + thread_c_off           + 1 + (13*8)] = acc_frag_3_7.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (13*8)] = acc_frag_3_7.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (13*8)] = acc_frag_3_7.w;
+
+}

tinygrad_repo/extra/gemm/max_kernels/nv.fp16_fp16_fp16.max.cu

@@ -0,0 +1,482 @@
+#define INFINITY (__int_as_float(0x7f800000))
+#define NAN (__int_as_float(0x7fffffff))
+#include <cuda_fp16.h>
+#include <cuda_pipeline.h>
+#define SMEM_N_WIDTH 136
+
+struct __align__(8) half4 { half x, y, z, w; };
+__device__ half4 make_half4(half x, half y, half z, half w) { half4 r={x, y, z, w}; return r; }
+
+struct __align__(16) half8 { half x, y, z, w, a, b, c, d; };
+__device__ half8 make_half8(half x, half y, half z, half w, half a, half b, half c, half d) { half8 r={x, y, z, w, a, b, c, d}; return r; }
+
+__device__ void __ldmatrix_a_elems(half8 *regs, half *smem) {
+    uint32_t reg0, reg1, reg2, reg3;
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(reg0), "=r"(reg1), "=r"(reg2), "=r"(reg3)
+        : "l"(__cvta_generic_to_shared(smem))
+    );
+    uint32_t *addr = reinterpret_cast<uint32_t*>(regs);
+    addr[0] = reg0;
+    addr[1] = reg1;
+    addr[2] = reg2;
+    addr[3] = reg3;
+}
+
+__device__ void __ldmatrix_b_elems(half4 *regs_lo, half4 *regs_hi, half *smem) {
+    uint32_t reg0, reg1, reg2, reg3;
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(reg0), "=r"(reg1), "=r"(reg2), "=r"(reg3)
+        : "l"(__cvta_generic_to_shared(smem))
+    );
+    uint32_t *addr_lo = reinterpret_cast<uint32_t*>(regs_lo);
+    uint32_t *addr_hi = reinterpret_cast<uint32_t*>(regs_hi);
+    addr_lo[0] = reg0;
+    addr_lo[1] = reg1;
+    addr_hi[0] = reg2;
+    addr_hi[1] = reg3;
+}
+
+__device__ half4 __WMMA_8_16_16_half_half(half8 a, half4 b, half4 c) {
+    int *a_pk = (int *) (&a), *b_pk = (int *) (&b), *c_pk = (int *) (&c);
+    asm( "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 { %0, %1 }, { %2, %3, %4, %5 }, { %6, %7 }, { %0, %1 };"
+        : "+r"(c_pk[0]), "+r"(c_pk[1]): "r"(a_pk[0]), "r"(a_pk[1]), "r"(a_pk[2]),  "r"(a_pk[3]), "r"(b_pk[0]), "r"(b_pk[1]) );
+    return c;
+}
+
+extern "C" __global__ void __launch_bounds__(256) wmma_example(half* data0, const half* data1, const half* data2, int N, int K) {
+    extern __shared__ char smem[];
+    half *smem_a_0 = (half *)(smem);
+    half *smem_a_1 = (half *)(smem + 16384);
+    half *smem_a_2 = (half *)(smem + 32768);
+    half *smem_b_0 = (half *)(smem + 49152);
+    half *smem_b_1 = (half *)(smem + 57344);
+    half *smem_b_2 = (half *)(smem + 65536);
+
+    int grid_m = blockIdx.x;        /* M//256 */
+    int grid_n = blockIdx.y;        /* N//128 */
+    int wg_threads = threadIdx.x;   // 32
+    int wg_m = threadIdx.y;         // 4
+    int wg_n = threadIdx.z;         // 2
+    int threads = threadIdx.x + (threadIdx.y * 32) + (threadIdx.z * 128); /* 256 */
+    int num_k_blocks = K / 32;
+
+    // ldmatrix indices - 4x loads of 8x8 matrices by 32 threads
+    // threads 0-7 are row starts for A, 8-15 for B, 16-23 for C, 24-31 for D
+    // [ A | C ]
+    // [ - + - ]
+    // [ B | D ]
+
+    // swizzled A - SMEM_A is 128 rows x 64 cols
+    size_t global_a_off = ((grid_m * 256) * K) + ((threads %  4) * 8) + (((threads /  4) % 2) * 8 * 16 * K) + ((threads / 8) * K);
+    size_t store_smem_a_off  = ((threads /  8) *  64) + (((threads * 8) ^ threads) & 56); // 32 rows / 64 cols per copy
+    size_t load_smem_a_row   = ((wg_m * 16) + (threads % 16)) * 64;
+    size_t load_smem_a_phase = (threads / 16) % 2;
+    size_t load_smem_a_0_k_0 = load_smem_a_row + ( 0 * 64) + (((load_smem_a_phase + 0) ^ (threads % 8)) * 8);
+    size_t load_smem_a_1_k_0 = load_smem_a_row + (64 * 64) + (((load_smem_a_phase + 0) ^ (threads % 8)) * 8);
+    size_t load_smem_a_2_k_0 = load_smem_a_row + ( 0 * 64) + (((load_smem_a_phase + 4) ^ (threads % 8)) * 8);
+    size_t load_smem_a_3_k_0 = load_smem_a_row + (64 * 64) + (((load_smem_a_phase + 4) ^ (threads % 8)) * 8);
+    size_t load_smem_a_0_k_1 = load_smem_a_row + ( 0 * 64) + (((load_smem_a_phase + 2) ^ (threads % 8)) * 8);
+    size_t load_smem_a_1_k_1 = load_smem_a_row + (64 * 64) + (((load_smem_a_phase + 2) ^ (threads % 8)) * 8);
+    size_t load_smem_a_2_k_1 = load_smem_a_row + ( 0 * 64) + (((load_smem_a_phase + 6) ^ (threads % 8)) * 8);
+    size_t load_smem_a_3_k_1 = load_smem_a_row + (64 * 64) + (((load_smem_a_phase + 6) ^ (threads % 8)) * 8);
+
+    // swizzled B - SMEM_B is 32 rows x 128 cols
+    size_t global_b_off = (grid_n * 128) + ((threads % 16) * 8) + ((threads / 16) * N);
+    size_t store_smem_b_off  = ((threads / 16) * 128) + ((((threads / 16) % 8) * 8) ^ ((threads % 16) * 8)); // 16 rows / 128 cols per copy
+    size_t load_smem_b_row   = (threads % 16) * 128;
+    size_t load_smem_b_phase = (wg_n * 2) + (wg_threads / 16);
+    size_t load_smem_b_0_k_0 = load_smem_b_row + (((load_smem_b_phase +  0) ^ (threads % 8)) * 8);
+    size_t load_smem_b_1_k_0 = load_smem_b_row + (((load_smem_b_phase +  4) ^ (threads % 8)) * 8);
+    size_t load_smem_b_2_k_0 = load_smem_b_row + (((load_smem_b_phase +  8) ^ (threads % 8)) * 8);
+    size_t load_smem_b_3_k_0 = load_smem_b_row + (((load_smem_b_phase + 12) ^ (threads % 8)) * 8);
+    size_t load_smem_b_0_k_1 = load_smem_b_0_k_0 + (16 * 128);
+    size_t load_smem_b_1_k_1 = load_smem_b_1_k_0 + (16 * 128);
+    size_t load_smem_b_2_k_1 = load_smem_b_2_k_0 + (16 * 128);
+    size_t load_smem_b_3_k_1 = load_smem_b_3_k_0 + (16 * 128);
+
+    // create accs (M=4, N=8)
+    half4 acc_frag_0_0 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_1 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_2 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_3 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_4 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_5 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_6 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_7 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_0 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_1 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_2 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_3 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_4 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_5 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_6 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_7 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_0 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_1 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_2 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_3 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_4 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_5 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_6 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_7 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_0 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_1 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_2 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_3 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_4 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_5 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_6 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_7 = make_half4(0.0f,0.0f,0.0f,0.0f);
+
+    // create registers for block A elements
+    half8 a_frag_0_k_0;
+    half8 a_frag_1_k_0;
+    half8 a_frag_2_k_0;
+    half8 a_frag_3_k_0;
+    half8 a_frag_0_k_1;
+    half8 a_frag_1_k_1;
+    half8 a_frag_2_k_1;
+    half8 a_frag_3_k_1;
+
+    // create register for block B elements
+    half4 b_frag_0_k_0;
+    half4 b_frag_1_k_0;
+    half4 b_frag_2_k_0;
+    half4 b_frag_3_k_0;
+    half4 b_frag_4_k_0;
+    half4 b_frag_5_k_0;
+    half4 b_frag_6_k_0;
+    half4 b_frag_7_k_0;
+    half4 b_frag_0_k_1;
+    half4 b_frag_1_k_1;
+    half4 b_frag_2_k_1;
+    half4 b_frag_3_k_1;
+    half4 b_frag_4_k_1;
+    half4 b_frag_5_k_1;
+    half4 b_frag_6_k_1;
+    half4 b_frag_7_k_1;
+
+    __syncthreads();
+
+    // load first tile
+    __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + (     0)], &data1[global_a_off + (    0)], 16);
+    __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + ( 32*64)], &data1[global_a_off + ( 32*K)], 16);
+    __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + ( 64*64)], &data1[global_a_off + ( 64*K)], 16);
+    __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + ( 96*64)], &data1[global_a_off + ( 96*K)], 16);
+    __pipeline_memcpy_async(&smem_b_0[store_smem_b_off + (     0)], &data2[global_b_off + (    0)], 16);
+    __pipeline_memcpy_async(&smem_b_0[store_smem_b_off + (16*128)], &data2[global_b_off + ( 16*N)], 16);
+    __pipeline_commit();
+    global_a_off += 32;
+    global_b_off += 32 * N;
+
+    // load second tile
+    __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + (     0)], &data1[global_a_off + (    0)], 16);
+    __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + ( 32*64)], &data1[global_a_off + ( 32*K)], 16);
+    __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + ( 64*64)], &data1[global_a_off + ( 64*K)], 16);
+    __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + ( 96*64)], &data1[global_a_off + ( 96*K)], 16);
+    __pipeline_memcpy_async(&smem_b_1[store_smem_b_off + (     0)], &data2[global_b_off + (    0)], 16);
+    __pipeline_memcpy_async(&smem_b_1[store_smem_b_off + (16*128)], &data2[global_b_off + ( 16*N)], 16);
+    __pipeline_commit();
+
+    global_a_off += 32;
+    global_b_off += 32 * N;
+
+    // wait on first pre-fetch load
+    __pipeline_wait_prior(1);
+    __syncthreads();
+
+    // load K=0 elements for the first tile
+    __ldmatrix_a_elems(&a_frag_0_k_0,                &smem_a_0[load_smem_a_0_k_0]);
+    __ldmatrix_a_elems(&a_frag_1_k_0,                &smem_a_0[load_smem_a_1_k_0]);
+    __ldmatrix_a_elems(&a_frag_2_k_0,                &smem_a_0[load_smem_a_2_k_0]);
+    __ldmatrix_a_elems(&a_frag_3_k_0,                &smem_a_0[load_smem_a_3_k_0]);
+    __ldmatrix_b_elems(&b_frag_0_k_0, &b_frag_1_k_0, &smem_b_0[load_smem_b_0_k_0]);
+    __ldmatrix_b_elems(&b_frag_2_k_0, &b_frag_3_k_0, &smem_b_0[load_smem_b_1_k_0]);
+    __ldmatrix_b_elems(&b_frag_4_k_0, &b_frag_5_k_0, &smem_b_0[load_smem_b_2_k_0]);
+    __ldmatrix_b_elems(&b_frag_6_k_0, &b_frag_7_k_0, &smem_b_0[load_smem_b_3_k_0]);
+
+    for (int block_k = 0; block_k < num_k_blocks; block_k++) {
+        int phase_k = block_k % 3;
+        half *smem_a_curr = (phase_k == 0) ? smem_a_0 : ((phase_k == 1) ? smem_a_1 : smem_a_2);
+        half *smem_b_curr = (phase_k == 0) ? smem_b_0 : ((phase_k == 1) ? smem_b_1 : smem_b_2);
+
+        int next_phase_k = (block_k+1) % 3;
+        half *smem_a_next = (next_phase_k == 0) ? smem_a_0 : ((next_phase_k == 1) ? smem_a_1 : smem_a_2);
+        half *smem_b_next = (next_phase_k == 0) ? smem_b_0 : ((next_phase_k == 1) ? smem_b_1 : smem_b_2);
+
+        int store_phase_k = (block_k+2) % 3;
+        half *smem_a_store = (store_phase_k == 0) ? smem_a_0 : ((store_phase_k == 1) ? smem_a_1 : smem_a_2);
+        half *smem_b_store = (store_phase_k == 0) ? smem_b_0 : ((store_phase_k == 1) ? smem_b_1 : smem_b_2);
+
+        // load K=1 elements for the current tile
+        __ldmatrix_a_elems(&a_frag_0_k_1,                &smem_a_curr[load_smem_a_0_k_1]);
+        __ldmatrix_a_elems(&a_frag_1_k_1,                &smem_a_curr[load_smem_a_1_k_1]);
+        __ldmatrix_a_elems(&a_frag_2_k_1,                &smem_a_curr[load_smem_a_2_k_1]);
+        __ldmatrix_a_elems(&a_frag_3_k_1,                &smem_a_curr[load_smem_a_3_k_1]);
+        __ldmatrix_b_elems(&b_frag_0_k_1, &b_frag_1_k_1, &smem_b_curr[load_smem_b_0_k_1]);
+        __ldmatrix_b_elems(&b_frag_2_k_1, &b_frag_3_k_1, &smem_b_curr[load_smem_b_1_k_1]);
+        __ldmatrix_b_elems(&b_frag_4_k_1, &b_frag_5_k_1, &smem_b_curr[load_smem_b_2_k_1]);
+        __ldmatrix_b_elems(&b_frag_6_k_1, &b_frag_7_k_1, &smem_b_curr[load_smem_b_3_k_1]);
+
+        // MMA K=0, (M=4 x N=8)
+        acc_frag_0_0 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_0_k_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_1_k_0, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_2_k_0, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_3_k_0, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_4_k_0, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_5_k_0, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_6_k_0, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_7_k_0, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_0_k_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_1_k_0, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_2_k_0, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_3_k_0, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_4_k_0, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_5_k_0, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_6_k_0, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_7_k_0, acc_frag_1_7);
+        acc_frag_2_0 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_0_k_0, acc_frag_2_0);
+        acc_frag_2_1 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_1_k_0, acc_frag_2_1);
+        acc_frag_2_2 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_2_k_0, acc_frag_2_2);
+        acc_frag_2_3 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_3_k_0, acc_frag_2_3);
+        acc_frag_2_4 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_4_k_0, acc_frag_2_4);
+        acc_frag_2_5 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_5_k_0, acc_frag_2_5);
+        acc_frag_2_6 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_6_k_0, acc_frag_2_6);
+        acc_frag_2_7 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_7_k_0, acc_frag_2_7);
+        acc_frag_3_0 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_0_k_0, acc_frag_3_0);
+        acc_frag_3_1 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_1_k_0, acc_frag_3_1);
+        acc_frag_3_2 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_2_k_0, acc_frag_3_2);
+        acc_frag_3_3 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_3_k_0, acc_frag_3_3);
+        acc_frag_3_4 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_4_k_0, acc_frag_3_4);
+        acc_frag_3_5 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_5_k_0, acc_frag_3_5);
+        acc_frag_3_6 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_6_k_0, acc_frag_3_6);
+        acc_frag_3_7 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_7_k_0, acc_frag_3_7);
+
+        // load next tile if needed
+        if (block_k < (num_k_blocks-2)) {
+            __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + (     0)], &data1[global_a_off + (    0)], 16);
+            __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + ( 32*64)], &data1[global_a_off + ( 32*K)], 16);
+            __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + ( 64*64)], &data1[global_a_off + ( 64*K)], 16);
+            __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + ( 96*64)], &data1[global_a_off + ( 96*K)], 16);
+            __pipeline_memcpy_async(&smem_b_store[store_smem_b_off + (     0)], &data2[global_b_off + (    0)], 16);
+            __pipeline_memcpy_async(&smem_b_store[store_smem_b_off + (16*128)], &data2[global_b_off + ( 16*N)], 16);
+            global_a_off += 32;
+            global_b_off += 32 * N;
+        }
+        __pipeline_commit();
+
+        // wait next tile
+        __pipeline_wait_prior(1);
+        __syncthreads();
+
+        // load K=0 elements for the next tile
+        __ldmatrix_a_elems(&a_frag_0_k_0,                &smem_a_next[load_smem_a_0_k_0]);
+        __ldmatrix_a_elems(&a_frag_1_k_0,                &smem_a_next[load_smem_a_1_k_0]);
+        __ldmatrix_a_elems(&a_frag_2_k_0,                &smem_a_next[load_smem_a_2_k_0]);
+        __ldmatrix_a_elems(&a_frag_3_k_0,                &smem_a_next[load_smem_a_3_k_0]);
+        __ldmatrix_b_elems(&b_frag_0_k_0, &b_frag_1_k_0, &smem_b_next[load_smem_b_0_k_0]);
+        __ldmatrix_b_elems(&b_frag_2_k_0, &b_frag_3_k_0, &smem_b_next[load_smem_b_1_k_0]);
+        __ldmatrix_b_elems(&b_frag_4_k_0, &b_frag_5_k_0, &smem_b_next[load_smem_b_2_k_0]);
+        __ldmatrix_b_elems(&b_frag_6_k_0, &b_frag_7_k_0, &smem_b_next[load_smem_b_3_k_0]);
+
+        // MMA K=1, (M=4 x N=8)
+        acc_frag_0_0 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_0_k_1, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_1_k_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_2_k_1, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_3_k_1, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_4_k_1, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_5_k_1, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_6_k_1, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_7_k_1, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_0_k_1, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_1_k_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_2_k_1, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_3_k_1, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_4_k_1, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_5_k_1, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_6_k_1, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_7_k_1, acc_frag_1_7);
+        acc_frag_2_0 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_0_k_1, acc_frag_2_0);
+        acc_frag_2_1 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_1_k_1, acc_frag_2_1);
+        acc_frag_2_2 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_2_k_1, acc_frag_2_2);
+        acc_frag_2_3 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_3_k_1, acc_frag_2_3);
+        acc_frag_2_4 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_4_k_1, acc_frag_2_4);
+        acc_frag_2_5 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_5_k_1, acc_frag_2_5);
+        acc_frag_2_6 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_6_k_1, acc_frag_2_6);
+        acc_frag_2_7 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_7_k_1, acc_frag_2_7);
+        acc_frag_3_0 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_0_k_1, acc_frag_3_0);
+        acc_frag_3_1 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_1_k_1, acc_frag_3_1);
+        acc_frag_3_2 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_2_k_1, acc_frag_3_2);
+        acc_frag_3_3 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_3_k_1, acc_frag_3_3);
+        acc_frag_3_4 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_4_k_1, acc_frag_3_4);
+        acc_frag_3_5 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_5_k_1, acc_frag_3_5);
+        acc_frag_3_6 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_6_k_1, acc_frag_3_6);
+        acc_frag_3_7 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_7_k_1, acc_frag_3_7);
+    }
+
+    // write accumulators to output
+    __pipeline_wait_prior(0);
+    __syncthreads();
+
+    // faster epilogue: write each 8x8 TC accs to SMEM first
+    //   - SMEM_N_WIDTH 8 larger than 128 required to deconflict bank access
+    //   - around 14 micros
+    //   - check bank conflict with in sudo with: "PYTHONPATH=. CUDA=1 GEMM_VARIATION="max" DTYPE_IN=half DTYPE_OUT=half DTYPE_ACC=half CNT=8 INPUT=ONES /usr/local/cuda/bin/ncu --section MemoryWorkloadAnalysis --metrics l1tex__data_bank_conflicts_pipe_lsu_mem_shared_op_ld.sum,l1tex__data_bank_conflicts_pipe_lsu_mem_shared_op_st.sum python3 ./extra/gemm/max_matmul.py"
+
+    // epilogue chunk with 256 threads / WG_M=4 / WG_N=2: split into 8 chunks (hi/lo for each in TC M)
+    //   1) write 32 rows of 128 cols (rows 0-7, 16-23, 32-39, 48-53 in acc_frag_0.lo, then acc_frag_0.hi, etc.)
+    //   2) read/write 16 rows of 128 elements in 8 elem (16B) chunks
+
+    half2 *smem32_d  = (half2 *)(smem);
+    half8 *smem128_d = (half8 *)(smem);
+    half8 *out128_d  = (half8 *)(data0);
+    size_t smem32_d_write_off  = (wg_m * 8 * (SMEM_N_WIDTH / 2)) + (wg_n * (16 / 2));
+    size_t smem32_d_thread_off = ((wg_threads / 4) * (SMEM_N_WIDTH / 2)) + (wg_threads % 4);
+    size_t smem128_d_read_off = ((threads / 16) * (SMEM_N_WIDTH / 8)) + (threads % 16);
+    size_t out128_d_off = ((grid_m * 256) * (N / 8)) + (grid_n * (128 / 8)) +
+        ((threads / 128) * 16 * (N / 8)) + (((threads / 16) % 8) * (N / 8)) + (threads % 16);
+
+    // write acc_frag_0_*
+
+    // write 32 rows of 128 N elements to SMEM
+    __syncthreads();
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 0*4)] = half2(acc_frag_0_0.x, acc_frag_0_0.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 1*4)] = half2(acc_frag_0_1.x, acc_frag_0_1.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 4*4)] = half2(acc_frag_0_2.x, acc_frag_0_2.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 5*4)] = half2(acc_frag_0_3.x, acc_frag_0_3.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 8*4)] = half2(acc_frag_0_4.x, acc_frag_0_4.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 9*4)] = half2(acc_frag_0_5.x, acc_frag_0_5.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (12*4)] = half2(acc_frag_0_6.x, acc_frag_0_6.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (13*4)] = half2(acc_frag_0_7.x, acc_frag_0_7.y);
+
+    // each thread reads and writes two 8 element chunks
+    __syncthreads();
+    out128_d[out128_d_off + ( 0 * (N / 8))] = smem128_d[smem128_d_read_off];
+    out128_d[out128_d_off + (32 * (N / 8))] = smem128_d[smem128_d_read_off + (16 * (SMEM_N_WIDTH / 8))];
+
+    // write 32 rows of 128 N elements to SMEM
+    __syncthreads();
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 0*4)] = half2(acc_frag_0_0.z, acc_frag_0_0.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 1*4)] = half2(acc_frag_0_1.z, acc_frag_0_1.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 4*4)] = half2(acc_frag_0_2.z, acc_frag_0_2.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 5*4)] = half2(acc_frag_0_3.z, acc_frag_0_3.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 8*4)] = half2(acc_frag_0_4.z, acc_frag_0_4.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 9*4)] = half2(acc_frag_0_5.z, acc_frag_0_5.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (12*4)] = half2(acc_frag_0_6.z, acc_frag_0_6.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (13*4)] = half2(acc_frag_0_7.z, acc_frag_0_7.w);
+
+    // each thread reads and writes two 8 element chunks
+    __syncthreads();
+    out128_d[out128_d_off + ( 8 * (N / 8))] = smem128_d[smem128_d_read_off];
+    out128_d[out128_d_off + (40 * (N / 8))] = smem128_d[smem128_d_read_off + (16 * (SMEM_N_WIDTH / 8))];
+
+    // write acc_frag_1_*
+    out128_d_off += (64 * (N / 8));
+
+    // write 32 rows of 128 N elements to SMEM
+    __syncthreads();
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 0*4)] = half2(acc_frag_1_0.x, acc_frag_1_0.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 1*4)] = half2(acc_frag_1_1.x, acc_frag_1_1.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 4*4)] = half2(acc_frag_1_2.x, acc_frag_1_2.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 5*4)] = half2(acc_frag_1_3.x, acc_frag_1_3.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 8*4)] = half2(acc_frag_1_4.x, acc_frag_1_4.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 9*4)] = half2(acc_frag_1_5.x, acc_frag_1_5.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (12*4)] = half2(acc_frag_1_6.x, acc_frag_1_6.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (13*4)] = half2(acc_frag_1_7.x, acc_frag_1_7.y);
+
+    // each thread reads and writes two 8 element chunks
+    __syncthreads();
+    out128_d[out128_d_off + ( 0 * (N / 8))] = smem128_d[smem128_d_read_off];
+    out128_d[out128_d_off + (32 * (N / 8))] = smem128_d[smem128_d_read_off + (16 * (SMEM_N_WIDTH / 8))];
+
+    // write 32 rows of 128 N elements to SMEM
+    __syncthreads();
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 0*4)] = half2(acc_frag_1_0.z, acc_frag_1_0.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 1*4)] = half2(acc_frag_1_1.z, acc_frag_1_1.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 4*4)] = half2(acc_frag_1_2.z, acc_frag_1_2.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 5*4)] = half2(acc_frag_1_3.z, acc_frag_1_3.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 8*4)] = half2(acc_frag_1_4.z, acc_frag_1_4.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 9*4)] = half2(acc_frag_1_5.z, acc_frag_1_5.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (12*4)] = half2(acc_frag_1_6.z, acc_frag_1_6.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (13*4)] = half2(acc_frag_1_7.z, acc_frag_1_7.w);
+
+    // each thread reads and writes two 8 element chunks
+    __syncthreads();
+    out128_d[out128_d_off + ( 8 * (N / 8))] = smem128_d[smem128_d_read_off];
+    out128_d[out128_d_off + (40 * (N / 8))] = smem128_d[smem128_d_read_off + (16 * (SMEM_N_WIDTH / 8))];
+
+    // write acc_frag_2_*
+    out128_d_off += (64 * (N / 8));
+
+    // write 32 rows of 128 N elements to SMEM
+    __syncthreads();
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 0*4)] = half2(acc_frag_2_0.x, acc_frag_2_0.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 1*4)] = half2(acc_frag_2_1.x, acc_frag_2_1.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 4*4)] = half2(acc_frag_2_2.x, acc_frag_2_2.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 5*4)] = half2(acc_frag_2_3.x, acc_frag_2_3.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 8*4)] = half2(acc_frag_2_4.x, acc_frag_2_4.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 9*4)] = half2(acc_frag_2_5.x, acc_frag_2_5.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (12*4)] = half2(acc_frag_2_6.x, acc_frag_2_6.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (13*4)] = half2(acc_frag_2_7.x, acc_frag_2_7.y);
+
+    // each thread reads and writes two 8 element chunks
+    __syncthreads();
+    out128_d[out128_d_off + ( 0 * (N / 8))] = smem128_d[smem128_d_read_off];
+    out128_d[out128_d_off + (32 * (N / 8))] = smem128_d[smem128_d_read_off + (16 * (SMEM_N_WIDTH / 8))];
+
+    // write 32 rows of 128 N elements to SMEM
+    __syncthreads();
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 0*4)] = half2(acc_frag_2_0.z, acc_frag_2_0.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 1*4)] = half2(acc_frag_2_1.z, acc_frag_2_1.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 4*4)] = half2(acc_frag_2_2.z, acc_frag_2_2.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 5*4)] = half2(acc_frag_2_3.z, acc_frag_2_3.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 8*4)] = half2(acc_frag_2_4.z, acc_frag_2_4.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 9*4)] = half2(acc_frag_2_5.z, acc_frag_2_5.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (12*4)] = half2(acc_frag_2_6.z, acc_frag_2_6.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (13*4)] = half2(acc_frag_2_7.z, acc_frag_2_7.w);
+
+    // each thread reads and writes two 8 element chunks
+    __syncthreads();
+    out128_d[out128_d_off + ( 8 * (N / 8))] = smem128_d[smem128_d_read_off];
+    out128_d[out128_d_off + (40 * (N / 8))] = smem128_d[smem128_d_read_off + (16 * (SMEM_N_WIDTH / 8))];
+
+    // write acc_frag_3_*
+    out128_d_off += (64 * (N / 8));
+
+    // write 32 rows of 128 N elements to SMEM
+    __syncthreads();
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 0*4)] = half2(acc_frag_3_0.x, acc_frag_3_0.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 1*4)] = half2(acc_frag_3_1.x, acc_frag_3_1.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 4*4)] = half2(acc_frag_3_2.x, acc_frag_3_2.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 5*4)] = half2(acc_frag_3_3.x, acc_frag_3_3.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 8*4)] = half2(acc_frag_3_4.x, acc_frag_3_4.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 9*4)] = half2(acc_frag_3_5.x, acc_frag_3_5.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (12*4)] = half2(acc_frag_3_6.x, acc_frag_3_6.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (13*4)] = half2(acc_frag_3_7.x, acc_frag_3_7.y);
+
+    // each thread reads and writes two 8 element chunks
+    __syncthreads();
+    out128_d[out128_d_off + ( 0 * (N / 8))] = smem128_d[smem128_d_read_off];
+    out128_d[out128_d_off + (32 * (N / 8))] = smem128_d[smem128_d_read_off + (16 * (SMEM_N_WIDTH / 8))];
+
+    // write 32 rows of 128 N elements to SMEM
+    __syncthreads();
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 0*4)] = half2(acc_frag_3_0.z, acc_frag_3_0.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 1*4)] = half2(acc_frag_3_1.z, acc_frag_3_1.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 4*4)] = half2(acc_frag_3_2.z, acc_frag_3_2.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 5*4)] = half2(acc_frag_3_3.z, acc_frag_3_3.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 8*4)] = half2(acc_frag_3_4.z, acc_frag_3_4.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 9*4)] = half2(acc_frag_3_5.z, acc_frag_3_5.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (12*4)] = half2(acc_frag_3_6.z, acc_frag_3_6.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (13*4)] = half2(acc_frag_3_7.z, acc_frag_3_7.w);
+
+    // each thread reads and writes two 8 element chunks
+    __syncthreads();
+    out128_d[out128_d_off + ( 8 * (N / 8))] = smem128_d[smem128_d_read_off];
+    out128_d[out128_d_off + (40 * (N / 8))] = smem128_d[smem128_d_read_off + (16 * (SMEM_N_WIDTH / 8))];
+
+    __syncthreads();
+}

tinygrad_repo/extra/gemm/max_kernels/nv.fp16_fp16_fp16.no_xor.cu

@@ -0,0 +1,486 @@
+#define INFINITY (__int_as_float(0x7f800000))
+#define NAN (__int_as_float(0x7fffffff))
+#include <cuda_fp16.h>
+#include <cuda_pipeline.h>
+#define SMEM_N_WIDTH 136
+
+struct __align__(8) half4 { half x, y, z, w; };
+__device__ half4 make_half4(half x, half y, half z, half w) { half4 r={x, y, z, w}; return r; }
+
+struct __align__(16) half8 { half x, y, z, w, a, b, c, d; };
+__device__ half8 make_half8(half x, half y, half z, half w, half a, half b, half c, half d) { half8 r={x, y, z, w, a, b, c, d}; return r; }
+
+__device__ void __ldmatrix_a_elems(half8 *regs, half *smem) {
+    uint32_t reg0, reg1, reg2, reg3;
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(reg0), "=r"(reg1), "=r"(reg2), "=r"(reg3)
+        : "l"(__cvta_generic_to_shared(smem))
+    );
+    uint32_t *addr = reinterpret_cast<uint32_t*>(regs);
+    addr[0] = reg0;
+    addr[1] = reg1;
+    addr[2] = reg2;
+    addr[3] = reg3;
+}
+
+__device__ void __ldmatrix_b_elems(half4 *regs_lo, half4 *regs_hi, half *smem) {
+    uint32_t reg0, reg1, reg2, reg3;
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(reg0), "=r"(reg1), "=r"(reg2), "=r"(reg3)
+        : "l"(__cvta_generic_to_shared(smem))
+    );
+    uint32_t *addr_lo = reinterpret_cast<uint32_t*>(regs_lo);
+    uint32_t *addr_hi = reinterpret_cast<uint32_t*>(regs_hi);
+    addr_lo[0] = reg0;
+    addr_lo[1] = reg1;
+    addr_hi[0] = reg2;
+    addr_hi[1] = reg3;
+}
+
+__device__ half4 __WMMA_8_16_16_half_half(half8 a, half4 b, half4 c) {
+    int *a_pk = (int *) (&a), *b_pk = (int *) (&b), *c_pk = (int *) (&c);
+    asm( "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 { %0, %1 }, { %2, %3, %4, %5 }, { %6, %7 }, { %0, %1 };"
+        : "+r"(c_pk[0]), "+r"(c_pk[1]): "r"(a_pk[0]), "r"(a_pk[1]), "r"(a_pk[2]),  "r"(a_pk[3]), "r"(b_pk[0]), "r"(b_pk[1]) );
+    return c;
+}
+
+extern "C" __global__ void __launch_bounds__(256) wmma_example(half* data0, const half* data1, const half* data2, int N, int K) {
+    extern __shared__ char smem[];
+    half *smem_a_0 = (half *)(smem);
+    half *smem_a_1 = (half *)(smem + 16384);
+    half *smem_a_2 = (half *)(smem + 32768);
+    half *smem_b_0 = (half *)(smem + 49152);
+    half *smem_b_1 = (half *)(smem + 57344);
+    half *smem_b_2 = (half *)(smem + 65536);
+
+    int grid_m = blockIdx.x;        /* M//256 */
+    int grid_n = blockIdx.y;        /* N//128 */
+    int wg_threads = threadIdx.x;   // 32
+    int wg_m = threadIdx.y;         // 4
+    int wg_n = threadIdx.z;         // 2
+    int threads = threadIdx.x + (threadIdx.y * 32) + (threadIdx.z * 128); /* 256 */
+    int num_k_blocks = K / 32;
+
+    // ldmatrix indices - 4x loads of 8x8 matrices by 32 threads
+    // threads 0-7 are row starts for A, 8-15 for B, 16-23 for C, 24-31 for D
+    // [ A | C ]
+    // [ - + - ]
+    // [ B | D ]
+
+    // swizzled A - SMEM_A is 128 rows x 64 cols
+    size_t global_a_off = ((grid_m * 256) * K) + ((threads %  4) * 8) + (((threads /  4) % 2) * 8 * 16 * K) + ((threads / 8) * K);
+    size_t store_smem_a_off  = ((threads /  8) *  64) + (((threads * 8) ^ threads) & 56); // 32 rows / 64 cols per copy
+    size_t load_smem_a_row   = ((wg_m * 16) + (threads % 16)) * 64;
+    size_t load_smem_a_phase = (threads / 16) % 2;
+    size_t load_smem_a_0_k_0 = load_smem_a_row + ( 0 * 64) + (((load_smem_a_phase + 0) ^ (threads % 8)) * 8);
+    size_t load_smem_a_1_k_0 = load_smem_a_row + (64 * 64) + (((load_smem_a_phase + 0) ^ (threads % 8)) * 8);
+    size_t load_smem_a_2_k_0 = load_smem_a_row + ( 0 * 64) + (((load_smem_a_phase + 4) ^ (threads % 8)) * 8);
+    size_t load_smem_a_3_k_0 = load_smem_a_row + (64 * 64) + (((load_smem_a_phase + 4) ^ (threads % 8)) * 8);
+    size_t load_smem_a_0_k_1 = load_smem_a_row + ( 0 * 64) + (((load_smem_a_phase + 2) ^ (threads % 8)) * 8);
+    size_t load_smem_a_1_k_1 = load_smem_a_row + (64 * 64) + (((load_smem_a_phase + 2) ^ (threads % 8)) * 8);
+    size_t load_smem_a_2_k_1 = load_smem_a_row + ( 0 * 64) + (((load_smem_a_phase + 6) ^ (threads % 8)) * 8);
+    size_t load_smem_a_3_k_1 = load_smem_a_row + (64 * 64) + (((load_smem_a_phase + 6) ^ (threads % 8)) * 8);
+
+    // swizzled B - SMEM_B is 64 rows x 64 cols
+    size_t global_b_off = (grid_n * 128) + ((threads % 16) * 8) + ((threads / 16) * N);
+    size_t store_smem_b_off = // 32 rows of 64 cols per copy
+        ((threads / 128) * (64)) +                  // [A,C] vs [B,D] in ldmatrix
+        ((threads % 2) * (2 * 64)) +                // [A vs C] or [B vs. D]
+        (((threads / 2) % 2) * (4 * 64)) +          // WG_N in [0, 1]
+        (((threads / 4) % 4) * (8 * 64)) +          // B in [0, 1, 2, 3]
+        (((threads / 16) % 8) * (8));               // cols in SMEM_B i.e. rows of 8x8
+
+    size_t load_smem_b_0_k_0 = (wg_n * 4 * 64) + ((wg_threads / 8) * 64) + ((wg_threads % 8) * 8);
+    size_t load_smem_b_1_k_0 = load_smem_b_0_k_0 + ( 8 * 64);
+    size_t load_smem_b_2_k_0 = load_smem_b_0_k_0 + (16 * 64);
+    size_t load_smem_b_3_k_0 = load_smem_b_0_k_0 + (24 * 64);
+    size_t load_smem_b_0_k_1 = load_smem_b_0_k_0 + (32 * 64);
+    size_t load_smem_b_1_k_1 = load_smem_b_1_k_0 + (32 * 64);
+    size_t load_smem_b_2_k_1 = load_smem_b_2_k_0 + (32 * 64);
+    size_t load_smem_b_3_k_1 = load_smem_b_3_k_0 + (32 * 64);
+
+    // create accs (M=4, N=8)
+    half4 acc_frag_0_0 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_1 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_2 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_3 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_4 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_5 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_6 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_0_7 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_0 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_1 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_2 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_3 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_4 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_5 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_6 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_1_7 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_0 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_1 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_2 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_3 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_4 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_5 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_6 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_2_7 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_0 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_1 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_2 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_3 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_4 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_5 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_6 = make_half4(0.0f,0.0f,0.0f,0.0f);
+    half4 acc_frag_3_7 = make_half4(0.0f,0.0f,0.0f,0.0f);
+
+    // create registers for block A elements
+    half8 a_frag_0_k_0;
+    half8 a_frag_1_k_0;
+    half8 a_frag_2_k_0;
+    half8 a_frag_3_k_0;
+    half8 a_frag_0_k_1;
+    half8 a_frag_1_k_1;
+    half8 a_frag_2_k_1;
+    half8 a_frag_3_k_1;
+
+    // create register for block B elements
+    half4 b_frag_0_k_0;
+    half4 b_frag_1_k_0;
+    half4 b_frag_2_k_0;
+    half4 b_frag_3_k_0;
+    half4 b_frag_4_k_0;
+    half4 b_frag_5_k_0;
+    half4 b_frag_6_k_0;
+    half4 b_frag_7_k_0;
+    half4 b_frag_0_k_1;
+    half4 b_frag_1_k_1;
+    half4 b_frag_2_k_1;
+    half4 b_frag_3_k_1;
+    half4 b_frag_4_k_1;
+    half4 b_frag_5_k_1;
+    half4 b_frag_6_k_1;
+    half4 b_frag_7_k_1;
+
+    __syncthreads();
+
+    // load first tile
+    __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + (     0)], &data1[global_a_off + (    0)], 16);
+    __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + ( 32*64)], &data1[global_a_off + ( 32*K)], 16);
+    __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + ( 64*64)], &data1[global_a_off + ( 64*K)], 16);
+    __pipeline_memcpy_async(&smem_a_0[store_smem_a_off + ( 96*64)], &data1[global_a_off + ( 96*K)], 16);
+    __pipeline_memcpy_async(&smem_b_0[store_smem_b_off + (     0)], &data2[global_b_off + (    0)], 16);
+    __pipeline_memcpy_async(&smem_b_0[store_smem_b_off + ( 32*64)], &data2[global_b_off + ( 16*N)], 16);
+    __pipeline_commit();
+    global_a_off += 32;
+    global_b_off += 32 * N;
+
+    // load second tile
+    __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + (     0)], &data1[global_a_off + (    0)], 16);
+    __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + ( 32*64)], &data1[global_a_off + ( 32*K)], 16);
+    __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + ( 64*64)], &data1[global_a_off + ( 64*K)], 16);
+    __pipeline_memcpy_async(&smem_a_1[store_smem_a_off + ( 96*64)], &data1[global_a_off + ( 96*K)], 16);
+    __pipeline_memcpy_async(&smem_b_1[store_smem_b_off + (     0)], &data2[global_b_off + (    0)], 16);
+    __pipeline_memcpy_async(&smem_b_1[store_smem_b_off + ( 32*64)], &data2[global_b_off + ( 16*N)], 16);
+    __pipeline_commit();
+
+    global_a_off += 32;
+    global_b_off += 32 * N;
+
+    // wait on first pre-fetch load
+    __pipeline_wait_prior(1);
+    __syncthreads();
+
+    // load K=0 elements for the first tile
+    __ldmatrix_a_elems(&a_frag_0_k_0,                &smem_a_0[load_smem_a_0_k_0]);
+    __ldmatrix_a_elems(&a_frag_1_k_0,                &smem_a_0[load_smem_a_1_k_0]);
+    __ldmatrix_a_elems(&a_frag_2_k_0,                &smem_a_0[load_smem_a_2_k_0]);
+    __ldmatrix_a_elems(&a_frag_3_k_0,                &smem_a_0[load_smem_a_3_k_0]);
+    __ldmatrix_b_elems(&b_frag_0_k_0, &b_frag_1_k_0, &smem_b_0[load_smem_b_0_k_0]);
+    __ldmatrix_b_elems(&b_frag_2_k_0, &b_frag_3_k_0, &smem_b_0[load_smem_b_1_k_0]);
+    __ldmatrix_b_elems(&b_frag_4_k_0, &b_frag_5_k_0, &smem_b_0[load_smem_b_2_k_0]);
+    __ldmatrix_b_elems(&b_frag_6_k_0, &b_frag_7_k_0, &smem_b_0[load_smem_b_3_k_0]);
+
+    for (int block_k = 0; block_k < num_k_blocks; block_k++) {
+        int phase_k = block_k % 3;
+        half *smem_a_curr = (phase_k == 0) ? smem_a_0 : ((phase_k == 1) ? smem_a_1 : smem_a_2);
+        half *smem_b_curr = (phase_k == 0) ? smem_b_0 : ((phase_k == 1) ? smem_b_1 : smem_b_2);
+
+        int next_phase_k = (block_k+1) % 3;
+        half *smem_a_next = (next_phase_k == 0) ? smem_a_0 : ((next_phase_k == 1) ? smem_a_1 : smem_a_2);
+        half *smem_b_next = (next_phase_k == 0) ? smem_b_0 : ((next_phase_k == 1) ? smem_b_1 : smem_b_2);
+
+        int store_phase_k = (block_k+2) % 3;
+        half *smem_a_store = (store_phase_k == 0) ? smem_a_0 : ((store_phase_k == 1) ? smem_a_1 : smem_a_2);
+        half *smem_b_store = (store_phase_k == 0) ? smem_b_0 : ((store_phase_k == 1) ? smem_b_1 : smem_b_2);
+
+        // load K=1 elements for the current tile
+        __ldmatrix_a_elems(&a_frag_0_k_1,                &smem_a_curr[load_smem_a_0_k_1]);
+        __ldmatrix_a_elems(&a_frag_1_k_1,                &smem_a_curr[load_smem_a_1_k_1]);
+        __ldmatrix_a_elems(&a_frag_2_k_1,                &smem_a_curr[load_smem_a_2_k_1]);
+        __ldmatrix_a_elems(&a_frag_3_k_1,                &smem_a_curr[load_smem_a_3_k_1]);
+        __ldmatrix_b_elems(&b_frag_0_k_1, &b_frag_1_k_1, &smem_b_curr[load_smem_b_0_k_1]);
+        __ldmatrix_b_elems(&b_frag_2_k_1, &b_frag_3_k_1, &smem_b_curr[load_smem_b_1_k_1]);
+        __ldmatrix_b_elems(&b_frag_4_k_1, &b_frag_5_k_1, &smem_b_curr[load_smem_b_2_k_1]);
+        __ldmatrix_b_elems(&b_frag_6_k_1, &b_frag_7_k_1, &smem_b_curr[load_smem_b_3_k_1]);
+
+        // MMA K=0, (M=4 x N=8)
+        acc_frag_0_0 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_0_k_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_1_k_0, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_2_k_0, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_3_k_0, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_4_k_0, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_5_k_0, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_6_k_0, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_half(a_frag_0_k_0, b_frag_7_k_0, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_0_k_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_1_k_0, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_2_k_0, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_3_k_0, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_4_k_0, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_5_k_0, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_6_k_0, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_half(a_frag_1_k_0, b_frag_7_k_0, acc_frag_1_7);
+        acc_frag_2_0 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_0_k_0, acc_frag_2_0);
+        acc_frag_2_1 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_1_k_0, acc_frag_2_1);
+        acc_frag_2_2 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_2_k_0, acc_frag_2_2);
+        acc_frag_2_3 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_3_k_0, acc_frag_2_3);
+        acc_frag_2_4 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_4_k_0, acc_frag_2_4);
+        acc_frag_2_5 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_5_k_0, acc_frag_2_5);
+        acc_frag_2_6 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_6_k_0, acc_frag_2_6);
+        acc_frag_2_7 = __WMMA_8_16_16_half_half(a_frag_2_k_0, b_frag_7_k_0, acc_frag_2_7);
+        acc_frag_3_0 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_0_k_0, acc_frag_3_0);
+        acc_frag_3_1 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_1_k_0, acc_frag_3_1);
+        acc_frag_3_2 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_2_k_0, acc_frag_3_2);
+        acc_frag_3_3 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_3_k_0, acc_frag_3_3);
+        acc_frag_3_4 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_4_k_0, acc_frag_3_4);
+        acc_frag_3_5 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_5_k_0, acc_frag_3_5);
+        acc_frag_3_6 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_6_k_0, acc_frag_3_6);
+        acc_frag_3_7 = __WMMA_8_16_16_half_half(a_frag_3_k_0, b_frag_7_k_0, acc_frag_3_7);
+
+        // load next tile if needed
+        if (block_k < (num_k_blocks-2)) {
+            __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + (     0)], &data1[global_a_off + (    0)], 16);
+            __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + ( 32*64)], &data1[global_a_off + ( 32*K)], 16);
+            __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + ( 64*64)], &data1[global_a_off + ( 64*K)], 16);
+            __pipeline_memcpy_async(&smem_a_store[store_smem_a_off + ( 96*64)], &data1[global_a_off + ( 96*K)], 16);
+            __pipeline_memcpy_async(&smem_b_store[store_smem_b_off + (     0)], &data2[global_b_off + (    0)], 16);
+            __pipeline_memcpy_async(&smem_b_store[store_smem_b_off + ( 32*64)], &data2[global_b_off + ( 16*N)], 16);
+            global_a_off += 32;
+            global_b_off += 32 * N;
+        }
+        __pipeline_commit();
+
+        // wait next tile
+        __pipeline_wait_prior(1);
+        __syncthreads();
+
+        // load K=0 elements for the next tile
+        __ldmatrix_a_elems(&a_frag_0_k_0,                &smem_a_next[load_smem_a_0_k_0]);
+        __ldmatrix_a_elems(&a_frag_1_k_0,                &smem_a_next[load_smem_a_1_k_0]);
+        __ldmatrix_a_elems(&a_frag_2_k_0,                &smem_a_next[load_smem_a_2_k_0]);
+        __ldmatrix_a_elems(&a_frag_3_k_0,                &smem_a_next[load_smem_a_3_k_0]);
+        __ldmatrix_b_elems(&b_frag_0_k_0, &b_frag_1_k_0, &smem_b_next[load_smem_b_0_k_0]);
+        __ldmatrix_b_elems(&b_frag_2_k_0, &b_frag_3_k_0, &smem_b_next[load_smem_b_1_k_0]);
+        __ldmatrix_b_elems(&b_frag_4_k_0, &b_frag_5_k_0, &smem_b_next[load_smem_b_2_k_0]);
+        __ldmatrix_b_elems(&b_frag_6_k_0, &b_frag_7_k_0, &smem_b_next[load_smem_b_3_k_0]);
+
+        // MMA K=1, (M=4 x N=8)
+        acc_frag_0_0 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_0_k_1, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_1_k_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_2_k_1, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_3_k_1, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_4_k_1, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_5_k_1, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_6_k_1, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_half(a_frag_0_k_1, b_frag_7_k_1, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_0_k_1, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_1_k_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_2_k_1, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_3_k_1, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_4_k_1, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_5_k_1, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_6_k_1, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_half(a_frag_1_k_1, b_frag_7_k_1, acc_frag_1_7);
+        acc_frag_2_0 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_0_k_1, acc_frag_2_0);
+        acc_frag_2_1 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_1_k_1, acc_frag_2_1);
+        acc_frag_2_2 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_2_k_1, acc_frag_2_2);
+        acc_frag_2_3 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_3_k_1, acc_frag_2_3);
+        acc_frag_2_4 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_4_k_1, acc_frag_2_4);
+        acc_frag_2_5 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_5_k_1, acc_frag_2_5);
+        acc_frag_2_6 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_6_k_1, acc_frag_2_6);
+        acc_frag_2_7 = __WMMA_8_16_16_half_half(a_frag_2_k_1, b_frag_7_k_1, acc_frag_2_7);
+        acc_frag_3_0 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_0_k_1, acc_frag_3_0);
+        acc_frag_3_1 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_1_k_1, acc_frag_3_1);
+        acc_frag_3_2 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_2_k_1, acc_frag_3_2);
+        acc_frag_3_3 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_3_k_1, acc_frag_3_3);
+        acc_frag_3_4 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_4_k_1, acc_frag_3_4);
+        acc_frag_3_5 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_5_k_1, acc_frag_3_5);
+        acc_frag_3_6 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_6_k_1, acc_frag_3_6);
+        acc_frag_3_7 = __WMMA_8_16_16_half_half(a_frag_3_k_1, b_frag_7_k_1, acc_frag_3_7);
+    }
+
+    // write accumulators to output
+    __pipeline_wait_prior(0);
+    __syncthreads();
+
+    // faster epilogue: write each 8x8 TC accs to SMEM first
+    //   - SMEM_N_WIDTH 8 larger than 128 required to deconflict bank access
+    //   - around 14 micros
+    //   - check bank conflict with in sudo with: "PYTHONPATH=. CUDA=1 GEMM_VARIATION="max" DTYPE_IN=half DTYPE_OUT=half DTYPE_ACC=half CNT=8 INPUT=ONES /usr/local/cuda/bin/ncu --section MemoryWorkloadAnalysis --metrics l1tex__data_bank_conflicts_pipe_lsu_mem_shared_op_ld.sum,l1tex__data_bank_conflicts_pipe_lsu_mem_shared_op_st.sum python3 ./extra/gemm/max_matmul.py"
+
+    // epilogue chunk with 256 threads / WG_M=4 / WG_N=2: split into 8 chunks (hi/lo for each in TC M)
+    //   1) write 32 rows of 128 cols (rows 0-7, 16-23, 32-39, 48-53 in acc_frag_0.lo, then acc_frag_0.hi, etc.)
+    //   2) read/write 16 rows of 128 elements in 8 elem (16B) chunks
+
+    half2 *smem32_d  = (half2 *)(smem);
+    half8 *smem128_d = (half8 *)(smem);
+    half8 *out128_d  = (half8 *)(data0);
+    size_t smem32_d_write_off  = (wg_m * 8 * (SMEM_N_WIDTH / 2)) + (wg_n * (16 / 2));
+    size_t smem32_d_thread_off = ((wg_threads / 4) * (SMEM_N_WIDTH / 2)) + (wg_threads % 4);
+    size_t smem128_d_read_off = ((threads / 16) * (SMEM_N_WIDTH / 8)) + (threads % 16);
+    size_t out128_d_off = ((grid_m * 256) * (N / 8)) + (grid_n * (128 / 8)) +
+        ((threads / 128) * 16 * (N / 8)) + (((threads / 16) % 8) * (N / 8)) + (threads % 16);
+
+    // write acc_frag_0_*
+
+    // write 32 rows of 128 N elements to SMEM
+    __syncthreads();
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 0*4)] = half2(acc_frag_0_0.x, acc_frag_0_0.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 1*4)] = half2(acc_frag_0_1.x, acc_frag_0_1.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 4*4)] = half2(acc_frag_0_2.x, acc_frag_0_2.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 5*4)] = half2(acc_frag_0_3.x, acc_frag_0_3.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 8*4)] = half2(acc_frag_0_4.x, acc_frag_0_4.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 9*4)] = half2(acc_frag_0_5.x, acc_frag_0_5.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (12*4)] = half2(acc_frag_0_6.x, acc_frag_0_6.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (13*4)] = half2(acc_frag_0_7.x, acc_frag_0_7.y);
+
+    // each thread reads and writes two 8 element chunks
+    __syncthreads();
+    out128_d[out128_d_off + ( 0 * (N / 8))] = smem128_d[smem128_d_read_off];
+    out128_d[out128_d_off + (32 * (N / 8))] = smem128_d[smem128_d_read_off + (16 * (SMEM_N_WIDTH / 8))];
+
+    // write 32 rows of 128 N elements to SMEM
+    __syncthreads();
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 0*4)] = half2(acc_frag_0_0.z, acc_frag_0_0.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 1*4)] = half2(acc_frag_0_1.z, acc_frag_0_1.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 4*4)] = half2(acc_frag_0_2.z, acc_frag_0_2.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 5*4)] = half2(acc_frag_0_3.z, acc_frag_0_3.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 8*4)] = half2(acc_frag_0_4.z, acc_frag_0_4.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 9*4)] = half2(acc_frag_0_5.z, acc_frag_0_5.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (12*4)] = half2(acc_frag_0_6.z, acc_frag_0_6.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (13*4)] = half2(acc_frag_0_7.z, acc_frag_0_7.w);
+
+    // each thread reads and writes two 8 element chunks
+    __syncthreads();
+    out128_d[out128_d_off + ( 8 * (N / 8))] = smem128_d[smem128_d_read_off];
+    out128_d[out128_d_off + (40 * (N / 8))] = smem128_d[smem128_d_read_off + (16 * (SMEM_N_WIDTH / 8))];
+
+    // write acc_frag_1_*
+    out128_d_off += (64 * (N / 8));
+
+    // write 32 rows of 128 N elements to SMEM
+    __syncthreads();
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 0*4)] = half2(acc_frag_1_0.x, acc_frag_1_0.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 1*4)] = half2(acc_frag_1_1.x, acc_frag_1_1.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 4*4)] = half2(acc_frag_1_2.x, acc_frag_1_2.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 5*4)] = half2(acc_frag_1_3.x, acc_frag_1_3.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 8*4)] = half2(acc_frag_1_4.x, acc_frag_1_4.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 9*4)] = half2(acc_frag_1_5.x, acc_frag_1_5.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (12*4)] = half2(acc_frag_1_6.x, acc_frag_1_6.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (13*4)] = half2(acc_frag_1_7.x, acc_frag_1_7.y);
+
+    // each thread reads and writes two 8 element chunks
+    __syncthreads();
+    out128_d[out128_d_off + ( 0 * (N / 8))] = smem128_d[smem128_d_read_off];
+    out128_d[out128_d_off + (32 * (N / 8))] = smem128_d[smem128_d_read_off + (16 * (SMEM_N_WIDTH / 8))];
+
+    // write 32 rows of 128 N elements to SMEM
+    __syncthreads();
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 0*4)] = half2(acc_frag_1_0.z, acc_frag_1_0.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 1*4)] = half2(acc_frag_1_1.z, acc_frag_1_1.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 4*4)] = half2(acc_frag_1_2.z, acc_frag_1_2.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 5*4)] = half2(acc_frag_1_3.z, acc_frag_1_3.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 8*4)] = half2(acc_frag_1_4.z, acc_frag_1_4.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 9*4)] = half2(acc_frag_1_5.z, acc_frag_1_5.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (12*4)] = half2(acc_frag_1_6.z, acc_frag_1_6.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (13*4)] = half2(acc_frag_1_7.z, acc_frag_1_7.w);
+
+    // each thread reads and writes two 8 element chunks
+    __syncthreads();
+    out128_d[out128_d_off + ( 8 * (N / 8))] = smem128_d[smem128_d_read_off];
+    out128_d[out128_d_off + (40 * (N / 8))] = smem128_d[smem128_d_read_off + (16 * (SMEM_N_WIDTH / 8))];
+
+    // write acc_frag_2_*
+    out128_d_off += (64 * (N / 8));
+
+    // write 32 rows of 128 N elements to SMEM
+    __syncthreads();
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 0*4)] = half2(acc_frag_2_0.x, acc_frag_2_0.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 1*4)] = half2(acc_frag_2_1.x, acc_frag_2_1.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 4*4)] = half2(acc_frag_2_2.x, acc_frag_2_2.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 5*4)] = half2(acc_frag_2_3.x, acc_frag_2_3.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 8*4)] = half2(acc_frag_2_4.x, acc_frag_2_4.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 9*4)] = half2(acc_frag_2_5.x, acc_frag_2_5.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (12*4)] = half2(acc_frag_2_6.x, acc_frag_2_6.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (13*4)] = half2(acc_frag_2_7.x, acc_frag_2_7.y);
+
+    // each thread reads and writes two 8 element chunks
+    __syncthreads();
+    out128_d[out128_d_off + ( 0 * (N / 8))] = smem128_d[smem128_d_read_off];
+    out128_d[out128_d_off + (32 * (N / 8))] = smem128_d[smem128_d_read_off + (16 * (SMEM_N_WIDTH / 8))];
+
+    // write 32 rows of 128 N elements to SMEM
+    __syncthreads();
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 0*4)] = half2(acc_frag_2_0.z, acc_frag_2_0.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 1*4)] = half2(acc_frag_2_1.z, acc_frag_2_1.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 4*4)] = half2(acc_frag_2_2.z, acc_frag_2_2.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 5*4)] = half2(acc_frag_2_3.z, acc_frag_2_3.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 8*4)] = half2(acc_frag_2_4.z, acc_frag_2_4.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 9*4)] = half2(acc_frag_2_5.z, acc_frag_2_5.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (12*4)] = half2(acc_frag_2_6.z, acc_frag_2_6.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (13*4)] = half2(acc_frag_2_7.z, acc_frag_2_7.w);
+
+    // each thread reads and writes two 8 element chunks
+    __syncthreads();
+    out128_d[out128_d_off + ( 8 * (N / 8))] = smem128_d[smem128_d_read_off];
+    out128_d[out128_d_off + (40 * (N / 8))] = smem128_d[smem128_d_read_off + (16 * (SMEM_N_WIDTH / 8))];
+
+    // write acc_frag_3_*
+    out128_d_off += (64 * (N / 8));
+
+    // write 32 rows of 128 N elements to SMEM
+    __syncthreads();
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 0*4)] = half2(acc_frag_3_0.x, acc_frag_3_0.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 1*4)] = half2(acc_frag_3_1.x, acc_frag_3_1.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 4*4)] = half2(acc_frag_3_2.x, acc_frag_3_2.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 5*4)] = half2(acc_frag_3_3.x, acc_frag_3_3.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 8*4)] = half2(acc_frag_3_4.x, acc_frag_3_4.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 9*4)] = half2(acc_frag_3_5.x, acc_frag_3_5.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (12*4)] = half2(acc_frag_3_6.x, acc_frag_3_6.y);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (13*4)] = half2(acc_frag_3_7.x, acc_frag_3_7.y);
+
+    // each thread reads and writes two 8 element chunks
+    __syncthreads();
+    out128_d[out128_d_off + ( 0 * (N / 8))] = smem128_d[smem128_d_read_off];
+    out128_d[out128_d_off + (32 * (N / 8))] = smem128_d[smem128_d_read_off + (16 * (SMEM_N_WIDTH / 8))];
+
+    // write 32 rows of 128 N elements to SMEM
+    __syncthreads();
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 0*4)] = half2(acc_frag_3_0.z, acc_frag_3_0.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 1*4)] = half2(acc_frag_3_1.z, acc_frag_3_1.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 4*4)] = half2(acc_frag_3_2.z, acc_frag_3_2.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 5*4)] = half2(acc_frag_3_3.z, acc_frag_3_3.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 8*4)] = half2(acc_frag_3_4.z, acc_frag_3_4.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + ( 9*4)] = half2(acc_frag_3_5.z, acc_frag_3_5.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (12*4)] = half2(acc_frag_3_6.z, acc_frag_3_6.w);
+    smem32_d[smem32_d_write_off + smem32_d_thread_off + (13*4)] = half2(acc_frag_3_7.z, acc_frag_3_7.w);
+
+    // each thread reads and writes two 8 element chunks
+    __syncthreads();
+    out128_d[out128_d_off + ( 8 * (N / 8))] = smem128_d[smem128_d_read_off];
+    out128_d[out128_d_off + (40 * (N / 8))] = smem128_d[smem128_d_read_off + (16 * (SMEM_N_WIDTH / 8))];
+
+    __syncthreads();
+}

tinygrad_repo/extra/gemm/max_kernels/nv.fp16_fp32_fp16.hcopt.cu

@@ -0,0 +1,157 @@
+#define INFINITY (__int_as_float(0x7f800000))
+#define NAN (__int_as_float(0x7fffffff))
+#include <cuda_fp16.h>
+struct __align__(8) half4 { half x, y, z, w; }; __device__ half4 make_half4(half x, half y, half z, half w) { half4 r={x, y, z, w}; return r; }
+struct __align__(16) half8 { half x, y, z, w, a, b, c, d; }; __device__ half8 make_half8(half x, half y, half z, half w, half a, half b, half c, half d) { half8 r={x, y, z, w, a, b, c, d}; return r; }
+__device__ float4 __WMMA_8_16_16_half_float(half8 a, half4 b, float4 c) { int *a_pk = (int *) (&a), *b_pk = (int *) (&b);
+asm( "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 { %0, %1, %2, %3 }, { %4, %5, %6, %7 }, { %8, %9 }, { %0, %1, %2, %3 };"
+  : "+f"(c.x), "+f"(c.y), "+f"(c.z), "+f"(c.w) : "r"(a_pk[0]), "r"(a_pk[1]), "r"(a_pk[2]),  "r"(a_pk[3]), "r"(b_pk[0]), "r"(b_pk[1]) );
+return c;}
+extern "C" __global__ void __launch_bounds__(128) wmma_example(half* data0, const half* data1, const half* data2) {
+  int gidx0 = blockIdx.x; /* 32 */
+  int gidx1 = blockIdx.y; /* 64 */
+  int lidx0 = threadIdx.x; /* 16 */
+  int lidx1 = threadIdx.y; /* 2 */
+  int lidx2 = threadIdx.z; /* 4 */
+  float4 cast0 = make_float4(0.0f,0.0f,0.0f,0.0f);
+  int alu0 = (gidx0*128);
+  int alu1 = (gidx1*262144);
+  int alu2 = (lidx1*32768);
+  int alu3 = (lidx2*32);
+  int alu4 = (lidx0/8);
+  int alu5 = (alu4*16384);
+  int alu6 = (lidx0%2);
+  int alu7 = (alu6*2);
+  int alu8 = ((lidx0/2)%2);
+  int alu9 = (alu8*4);
+  int alu10 = ((lidx0/4)%2);
+  int alu11 = (alu10*8192);
+  int alu12 = (alu1+alu0+alu7+alu9+alu11+alu5+alu2+alu3);
+  int alu13 = (alu1+alu7+alu9+alu11+alu5+alu2);
+  float4 acc0 = cast0;
+  float4 acc1 = cast0;
+  float4 acc2 = cast0;
+  float4 acc3 = cast0;
+  float4 acc4 = cast0;
+  float4 acc5 = cast0;
+  float4 acc6 = cast0;
+  float4 acc7 = cast0;
+  float4 acc8 = cast0;
+  float4 acc9 = cast0;
+  float4 acc10 = cast0;
+  float4 acc11 = cast0;
+  float4 acc12 = cast0;
+  float4 acc13 = cast0;
+  float4 acc14 = cast0;
+  float4 acc15 = cast0;
+  for (int ridx0 = 0; ridx0 < 256; ridx0++) {
+    int alu14 = (ridx0*16);
+    int alu15 = (alu13+alu14);
+    int alu16 = (alu14+alu13);
+    int alu17 = (alu0+(alu6*8192)+(alu8*16384)+alu10+(alu4*2)+(lidx1*4)+alu3+(ridx0*65536));
+    half val0 = data2[alu17+8];
+    half val1 = data2[alu17+16];
+    half val2 = data2[alu17+24];
+    half val3 = data2[alu17+4096];
+    half val4 = data2[alu17+4104];
+    half val5 = data2[alu17+4112];
+    half val6 = data2[alu17+4120];
+    half val7 = data2[alu17+32768];
+    half val8 = data2[alu17+32776];
+    half val9 = data2[alu17+32784];
+    half val10 = data2[alu17+32792];
+    half val11 = data2[alu17+36864];
+    half val12 = data2[alu17+36872];
+    half4 cast1 = make_half4(val0,val4,val8,val12);
+    half val13 = data2[alu17+36880];
+    half4 cast2 = make_half4(val1,val5,val9,val13);
+    half val14 = data2[alu17+36888];
+    half4 cast3 = make_half4(val2,val6,val10,val14);
+    half val15 = data2[alu17];
+    half4 cast4 = make_half4(val15,val3,val7,val11);
+    half2 val16 = *((half2*)(data1+alu15+4096));
+    half2 val17 = *((half2*)(data1+alu15+65536));
+    half2 val18 = *((half2*)(data1+alu15+69632));
+    half2 val19 = *((half2*)(data1+alu15+131072));
+    half2 val20 = *((half2*)(data1+alu15+135168));
+    half2 val21 = *((half2*)(data1+alu15+196608));
+    half2 val22 = *((half2*)(data1+alu15+200704));
+    half2 val23 = *((half2*)(data1+alu15));
+    half2 val24 = *((half2*)(data1+alu16+8));
+    half2 val25 = *((half2*)(data1+alu16+4104));
+    half8 cast5 = make_half8(val23.x,val23.y,val16.x,val16.y,val24.x,val24.y,val25.x,val25.y);
+    float4 wmma0 = __WMMA_8_16_16_half_float(cast5, cast1, acc1);
+    float4 wmma1 = __WMMA_8_16_16_half_float(cast5, cast2, acc2);
+    float4 wmma2 = __WMMA_8_16_16_half_float(cast5, cast3, acc3);
+    float4 wmma3 = __WMMA_8_16_16_half_float(cast5, cast4, acc0);
+    half2 val26 = *((half2*)(data1+alu16+65544));
+    half2 val27 = *((half2*)(data1+alu16+69640));
+    half8 cast6 = make_half8(val17.x,val17.y,val18.x,val18.y,val26.x,val26.y,val27.x,val27.y);
+    float4 wmma4 = __WMMA_8_16_16_half_float(cast6, cast1, acc5);
+    float4 wmma5 = __WMMA_8_16_16_half_float(cast6, cast2, acc6);
+    float4 wmma6 = __WMMA_8_16_16_half_float(cast6, cast3, acc7);
+    float4 wmma7 = __WMMA_8_16_16_half_float(cast6, cast4, acc4);
+    half2 val28 = *((half2*)(data1+alu16+131080));
+    half2 val29 = *((half2*)(data1+alu16+135176));
+    half8 cast7 = make_half8(val19.x,val19.y,val20.x,val20.y,val28.x,val28.y,val29.x,val29.y);
+    float4 wmma8 = __WMMA_8_16_16_half_float(cast7, cast1, acc9);
+    float4 wmma9 = __WMMA_8_16_16_half_float(cast7, cast2, acc10);
+    float4 wmma10 = __WMMA_8_16_16_half_float(cast7, cast3, acc11);
+    float4 wmma11 = __WMMA_8_16_16_half_float(cast7, cast4, acc8);
+    half2 val30 = *((half2*)(data1+alu16+196616));
+    half2 val31 = *((half2*)(data1+alu16+200712));
+    half8 cast8 = make_half8(val21.x,val21.y,val22.x,val22.y,val30.x,val30.y,val31.x,val31.y);
+    float4 wmma12 = __WMMA_8_16_16_half_float(cast8, cast1, acc13);
+    float4 wmma13 = __WMMA_8_16_16_half_float(cast8, cast2, acc14);
+    float4 wmma14 = __WMMA_8_16_16_half_float(cast8, cast3, acc15);
+    float4 wmma15 = __WMMA_8_16_16_half_float(cast8, cast4, acc12);
+    acc0 = wmma3;
+    acc1 = wmma0;
+    acc2 = wmma1;
+    acc3 = wmma2;
+    acc4 = wmma7;
+    acc5 = wmma4;
+    acc6 = wmma5;
+    acc7 = wmma6;
+    acc8 = wmma11;
+    acc9 = wmma8;
+    acc10 = wmma9;
+    acc11 = wmma10;
+    acc12 = wmma15;
+    acc13 = wmma12;
+    acc14 = wmma13;
+    acc15 = wmma14;
+  }
+  *((half2*)(data0+alu12+8)) = make_half2((half)(acc1.x),(half)(acc1.y));
+  *((half2*)(data0+alu12+16)) = make_half2((half)(acc2.x),(half)(acc2.y));
+  *((half2*)(data0+alu12+24)) = make_half2((half)(acc3.x),(half)(acc3.y));
+  *((half2*)(data0+alu12+4096)) = make_half2((half)(acc0.z),(half)(acc0.w));
+  *((half2*)(data0+alu12+4104)) = make_half2((half)(acc1.z),(half)(acc1.w));
+  *((half2*)(data0+alu12+4112)) = make_half2((half)(acc2.z),(half)(acc2.w));
+  *((half2*)(data0+alu12+4120)) = make_half2((half)(acc3.z),(half)(acc3.w));
+  *((half2*)(data0+alu12+65536)) = make_half2((half)(acc4.x),(half)(acc4.y));
+  *((half2*)(data0+alu12+65544)) = make_half2((half)(acc5.x),(half)(acc5.y));
+  *((half2*)(data0+alu12+65552)) = make_half2((half)(acc6.x),(half)(acc6.y));
+  *((half2*)(data0+alu12+65560)) = make_half2((half)(acc7.x),(half)(acc7.y));
+  *((half2*)(data0+alu12+69632)) = make_half2((half)(acc4.z),(half)(acc4.w));
+  *((half2*)(data0+alu12+69640)) = make_half2((half)(acc5.z),(half)(acc5.w));
+  *((half2*)(data0+alu12+69648)) = make_half2((half)(acc6.z),(half)(acc6.w));
+  *((half2*)(data0+alu12+69656)) = make_half2((half)(acc7.z),(half)(acc7.w));
+  *((half2*)(data0+alu12+131072)) = make_half2((half)(acc8.x),(half)(acc8.y));
+  *((half2*)(data0+alu12+131080)) = make_half2((half)(acc9.x),(half)(acc9.y));
+  *((half2*)(data0+alu12+131088)) = make_half2((half)(acc10.x),(half)(acc10.y));
+  *((half2*)(data0+alu12+131096)) = make_half2((half)(acc11.x),(half)(acc11.y));
+  *((half2*)(data0+alu12+135168)) = make_half2((half)(acc8.z),(half)(acc8.w));
+  *((half2*)(data0+alu12+135176)) = make_half2((half)(acc9.z),(half)(acc9.w));
+  *((half2*)(data0+alu12+135184)) = make_half2((half)(acc10.z),(half)(acc10.w));
+  *((half2*)(data0+alu12+135192)) = make_half2((half)(acc11.z),(half)(acc11.w));
+  *((half2*)(data0+alu12+196608)) = make_half2((half)(acc12.x),(half)(acc12.y));
+  *((half2*)(data0+alu12+196616)) = make_half2((half)(acc13.x),(half)(acc13.y));
+  *((half2*)(data0+alu12+196624)) = make_half2((half)(acc14.x),(half)(acc14.y));
+  *((half2*)(data0+alu12+196632)) = make_half2((half)(acc15.x),(half)(acc15.y));
+  *((half2*)(data0+alu12+200704)) = make_half2((half)(acc12.z),(half)(acc12.w));
+  *((half2*)(data0+alu12+200712)) = make_half2((half)(acc13.z),(half)(acc13.w));
+  *((half2*)(data0+alu12+200720)) = make_half2((half)(acc14.z),(half)(acc14.w));
+  *((half2*)(data0+alu12+200728)) = make_half2((half)(acc15.z),(half)(acc15.w));
+  *((half2*)(data0+alu12)) = make_half2((half)(acc0.x),(half)(acc0.y));
+}

tinygrad_repo/extra/gemm/max_kernels/nv.fp16_fp32_fp32.2_stage_swizzled_smem_input.cu

@@ -0,0 +1,398 @@
+#define INFINITY (__int_as_float(0x7f800000))
+#define NAN (__int_as_float(0x7fffffff))
+#include <cuda_fp16.h>
+#include <cuda_pipeline.h>
+#define N_PAD 132
+
+struct __align__(8) half4 { half x, y, z, w; };
+__device__ half4 make_half4(half x, half y, half z, half w) { half4 r={x, y, z, w}; return r; }
+
+struct __align__(16) half8 { half x, y, z, w, a, b, c, d; };
+__device__ half8 make_half8(half x, half y, half z, half w, half a, half b, half c, half d) { half8 r={x, y, z, w, a, b, c, d}; return r; }
+
+__device__ void __ldmatrix_a_elems(half8 *regs, half *smem) {
+    uint32_t reg0, reg1, reg2, reg3;
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(reg0), "=r"(reg1), "=r"(reg2), "=r"(reg3)
+        : "l"(__cvta_generic_to_shared(smem))
+    );
+    uint32_t *addr = reinterpret_cast<uint32_t*>(regs);
+    addr[0] = reg0;
+    addr[1] = reg1;
+    addr[2] = reg2;
+    addr[3] = reg3;
+}
+
+__device__ void __ldmatrix_b_elems(half4 *regs_lo, half4 *regs_hi, half *smem) {
+    uint32_t reg0, reg1, reg2, reg3;
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(reg0), "=r"(reg1), "=r"(reg2), "=r"(reg3)
+        : "l"(__cvta_generic_to_shared(smem))
+    );
+    uint32_t *addr_lo = reinterpret_cast<uint32_t*>(regs_lo);
+    uint32_t *addr_hi = reinterpret_cast<uint32_t*>(regs_hi);
+    addr_lo[0] = reg0;
+    addr_lo[1] = reg1;
+    addr_hi[0] = reg2;
+    addr_hi[1] = reg3;
+}
+
+__device__ float4 __WMMA_8_16_16_half_float(half8 a, half4 b, float4 c) {
+    int *a_pk = (int *) (&a), *b_pk = (int *) (&b);
+    asm( "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 { %0, %1, %2, %3 }, { %4, %5, %6, %7 }, { %8, %9 }, { %0, %1, %2, %3 };"
+        : "+f"(c.x), "+f"(c.y), "+f"(c.z), "+f"(c.w) : "r"(a_pk[0]), "r"(a_pk[1]), "r"(a_pk[2]),  "r"(a_pk[3]), "r"(b_pk[0]), "r"(b_pk[1]) );
+    return c;
+}
+
+extern "C" __global__ void __launch_bounds__(128) wmma_example(float* data0, const half* data1, const half* data2, int N, int K) {
+    int grid_m = blockIdx.x;        /* M//64 */
+    int grid_n = blockIdx.y;        /* N//128 */
+    int threads = threadIdx.x;      /* 128 */
+    int wg_m = (threads/64);        // 0 or 1 for 1st and 3rd blocks of b_m=16xb_k=16 vs 2nd and 4th blocks
+    int wg_n = (threads/32)%2;      // 0 or 1 for 1st, 3rd, 5th, 7th blocks of b_n=16xb_k=16 vs 2nd, 4th, 6th, 8th blocks - differs from triton
+    int wg_threads = threads%32;
+    int num_k_blocks = K / 64;
+
+    // load indexes
+    size_t global_a_off = ((grid_m * 64) * K) + ((threads %  8) * 8) + ((threads /  8) * K);
+    size_t global_b_off = (grid_n * 128)      + ((threads % 16) * 8) + ((threads / 16) * N);
+
+    // swizzled smem store offsets - columns of smem are swizzled
+    // here's a link to a description of the triton: https://github.com/triton-lang/triton/discussions/2026#discussioncomment-6746579
+    // see also the thunderkittens impl: https://github.com/HazyResearch/ThunderKittens/blob/main/include/types/shared/st.cuh
+    size_t store_smem_a_off  = ((threads /  8) *  64) + (((threads * 8) ^ threads) & 56);               // r15
+    size_t store_smem_b_off  = ((threads / 16) * 128) + (((threads / 16) * 8) ^ ((threads % 16) * 8));  // r19
+
+    // ldmatrix indices
+    // threads 0-7 are row starts for A, 8-15 for B, 16-23 for C, 24-31 for D
+    // [ A | C ]
+    // [ - + - ]
+    // [ B | D ]
+
+    // swizzled ldmatrix
+    size_t load_smem_a_row   = ((wg_m * 16) + (threads % 16)) * 64; // r293
+    size_t load_smem_a_phase = (threads / 16) % 2;                  // r4
+
+    size_t load_smem_b_row   = (threads % 16) * 128;                // r299
+    size_t load_smem_b_phase = (wg_n * 2) + (((threads / 16) % 2)); // r297 -- this differs from the generated triton kernel (swapped order)
+
+    size_t load_smem_a_0_k_0 = load_smem_a_row + (((load_smem_a_phase +  0) ^ (threads % 8)) * 8); // r38
+    size_t load_smem_a_1_k_0 = load_smem_a_0_k_0 + (32 * 64);
+    size_t load_smem_b_0_k_0 = load_smem_b_row + (((load_smem_b_phase +  0) ^ (threads % 8)) * 8);
+    size_t load_smem_b_1_k_0 = load_smem_b_row + (((load_smem_b_phase +  4) ^ (threads % 8)) * 8);
+    size_t load_smem_b_2_k_0 = load_smem_b_row + (((load_smem_b_phase +  8) ^ (threads % 8)) * 8);
+    size_t load_smem_b_3_k_0 = load_smem_b_row + (((load_smem_b_phase + 12) ^ (threads % 8)) * 8);
+
+    size_t load_smem_a_0_k_1 = load_smem_a_row + (((load_smem_a_phase +  2) ^ (threads % 8)) * 8); // r58 = r293 + r316;
+    size_t load_smem_a_1_k_1 = load_smem_a_0_k_1 + (32 * 64);
+    size_t load_smem_b_0_k_1 = load_smem_b_0_k_0 + (16 * 128);
+    size_t load_smem_b_1_k_1 = load_smem_b_1_k_0 + (16 * 128);
+    size_t load_smem_b_2_k_1 = load_smem_b_2_k_0 + (16 * 128);
+    size_t load_smem_b_3_k_1 = load_smem_b_3_k_0 + (16 * 128);
+
+    size_t load_smem_a_0_k_2 = load_smem_a_row + (((load_smem_a_phase +  4) ^ (threads % 8)) * 8); // r59 = r293 + r319;
+    size_t load_smem_a_1_k_2 = load_smem_a_0_k_2 + (32 * 64);
+    size_t load_smem_b_0_k_2 = load_smem_b_0_k_0 + (32 * 128);
+    size_t load_smem_b_1_k_2 = load_smem_b_1_k_0 + (32 * 128);
+    size_t load_smem_b_2_k_2 = load_smem_b_2_k_0 + (32 * 128);
+    size_t load_smem_b_3_k_2 = load_smem_b_3_k_0 + (32 * 128);
+
+    size_t load_smem_a_0_k_3 = load_smem_a_row + (((load_smem_a_phase +  6) ^ (threads % 8)) * 8); // r60 = r293 + r322;
+    size_t load_smem_a_1_k_3 = load_smem_a_0_k_3 + (32 * 64);
+    size_t load_smem_b_0_k_3 = load_smem_b_0_k_0 + (48 * 128);
+    size_t load_smem_b_1_k_3 = load_smem_b_1_k_0 + (48 * 128);
+    size_t load_smem_b_2_k_3 = load_smem_b_2_k_0 + (48 * 128);
+    size_t load_smem_b_3_k_3 = load_smem_b_3_k_0 + (48 * 128);
+
+    // create shared mem (A_1 8192 bytes, A_2 8192 bytes, B_1 16384 bytes, B2_16384 bytes)
+    __shared__ alignas(16) char smem[49152];
+
+    // create accs (16 WMMAs and 4 output elements each) and zero
+    float4 acc_frag_0_0 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_1 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_2 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_3 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_4 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_5 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_6 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_7 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_0 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_1 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_2 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_3 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_4 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_5 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_6 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_7 = make_float4(0.0f,0.0f,0.0f,0.0f);
+
+    // create registers for block A elements (2)
+    half8 a_frag_0;
+    half8 a_frag_1;
+
+    // create register for block B elements (8)
+    half4 b_frag_0;
+    half4 b_frag_1;
+    half4 b_frag_2;
+    half4 b_frag_3;
+    half4 b_frag_4;
+    half4 b_frag_5;
+    half4 b_frag_6;
+    half4 b_frag_7;
+
+    half *smem_a_even = (half *)(smem);
+    half *smem_a_odd  = (half *)(smem + 8192);
+    half *smem_b_even = (half *)(smem + 16384);
+    half *smem_b_odd  = (half *)(smem + 32768);
+
+    // https://developer.nvidia.com/blog/controlling-data-movement-to-boost-performance-on-ampere-architecture/
+    // https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#asynchronous-data-copies
+
+    // start first pre-fetch load A
+    __pipeline_memcpy_async(&smem_a_even[store_smem_a_off +  (    0)], &data1[global_a_off + (   0)], 16);
+    __pipeline_memcpy_async(&smem_a_even[store_smem_a_off +  (16*64)], &data1[global_a_off + (16*K)], 16);
+    __pipeline_memcpy_async(&smem_a_even[store_smem_a_off +  (32*64)], &data1[global_a_off + (32*K)], 16);
+    __pipeline_memcpy_async(&smem_a_even[store_smem_a_off +  (48*64)], &data1[global_a_off + (48*K)], 16);
+
+    // start first pre-fetch load B
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (     0)], &data2[global_b_off + (   0)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + ( 8*128)], &data2[global_b_off + ( 8*N)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (16*128)], &data2[global_b_off + (16*N)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (24*128)], &data2[global_b_off + (24*N)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (32*128)], &data2[global_b_off + (32*N)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (40*128)], &data2[global_b_off + (40*N)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (48*128)], &data2[global_b_off + (48*N)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (56*128)], &data2[global_b_off + (56*N)], 16);
+    __pipeline_commit();
+
+    global_a_off += 64;
+    global_b_off += 64 * N;
+    __syncthreads();
+
+    // start second pre-fetch load A
+    __pipeline_memcpy_async(&smem_a_odd[store_smem_a_off +  (    0)], &data1[global_a_off + (   0)], 16);
+    __pipeline_memcpy_async(&smem_a_odd[store_smem_a_off +  (16*64)], &data1[global_a_off + (16*K)], 16);
+    __pipeline_memcpy_async(&smem_a_odd[store_smem_a_off +  (32*64)], &data1[global_a_off + (32*K)], 16);
+    __pipeline_memcpy_async(&smem_a_odd[store_smem_a_off +  (48*64)], &data1[global_a_off + (48*K)], 16);
+
+    // start second pre-fetch load B
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (     0)], &data2[global_b_off + (   0)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + ( 8*128)], &data2[global_b_off + ( 8*N)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (16*128)], &data2[global_b_off + (16*N)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (24*128)], &data2[global_b_off + (24*N)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (32*128)], &data2[global_b_off + (32*N)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (40*128)], &data2[global_b_off + (40*N)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (48*128)], &data2[global_b_off + (48*N)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (56*128)], &data2[global_b_off + (56*N)], 16);
+    __pipeline_commit();
+
+    global_a_off += 64;
+    global_b_off += 64 * N;
+
+    // wait on needed prefetch value
+    __pipeline_wait_prior(0); // TODO: this enables fast iterations, but incorrect results with 1 (it shouldn't)
+    __syncthreads();
+
+    for (int block_k = 0; block_k < num_k_blocks; block_k++) {
+        // BLOCK_K==4: unroll 4 iterations of ldmatrix/wmma
+        half *smem_a_curr = (block_k % 2) ? smem_a_even : smem_a_odd;
+        half *smem_b_curr = (block_k % 2) ? smem_b_even : smem_b_odd;
+
+        // first load 16 K elements and 16 WMMAs: BLOCK_M==2 * BLOCK_N==8
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_0]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_0]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_0]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_0]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_0]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_0]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // next 16 K elements
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_1]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_1]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_1]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_1]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_1]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_1]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // next 16 K elements
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_2]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_2]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_2]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_2]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_2]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_2]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // last 16 K elements
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_3]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_3]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_3]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_3]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_3]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_3]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // prefetch next iteration if needed
+        __syncthreads();
+        if (block_k < (num_k_blocks-2)) {
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (    0)], &data1[global_a_off + (   0)], 16);
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (16*64)], &data1[global_a_off + (16*K)], 16);
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (32*64)], &data1[global_a_off + (32*K)], 16);
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (48*64)], &data1[global_a_off + (48*K)], 16);
+
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (     0)], &data2[global_b_off + (   0)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + ( 8*128)], &data2[global_b_off + ( 8*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (16*128)], &data2[global_b_off + (16*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (24*128)], &data2[global_b_off + (24*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (32*128)], &data2[global_b_off + (32*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (40*128)], &data2[global_b_off + (40*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (48*128)], &data2[global_b_off + (48*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (56*128)], &data2[global_b_off + (56*N)], 16);
+
+            global_a_off += 64;
+            global_b_off += 64 * N;
+        }
+        __pipeline_commit();
+
+        if (block_k < num_k_blocks-1) {
+            __pipeline_wait_prior(1);
+            __syncthreads();
+        }
+    }
+
+    // write accumulators to output
+    __pipeline_wait_prior(0);
+    __syncthreads();
+
+    // slower way: write floats one by one to data0
+    size_t wg_c_off     = ((grid_m * 64) * N) + (grid_n * 128) + (wg_m * 16 * N) + (wg_n * 16);
+    size_t thread_c_off = ((wg_threads % 4) * 2) + (((wg_threads / 4) % 8) * N);
+    data0[wg_c_off + thread_c_off           + 0 + ( 0*8)] = acc_frag_0_0.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 0*8)] = acc_frag_0_0.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 0*8)] = acc_frag_0_0.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 0*8)] = acc_frag_0_0.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 1*8)] = acc_frag_0_1.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 1*8)] = acc_frag_0_1.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 1*8)] = acc_frag_0_1.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 1*8)] = acc_frag_0_1.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 4*8)] = acc_frag_0_2.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 4*8)] = acc_frag_0_2.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 4*8)] = acc_frag_0_2.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 4*8)] = acc_frag_0_2.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 5*8)] = acc_frag_0_3.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 5*8)] = acc_frag_0_3.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 5*8)] = acc_frag_0_3.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 5*8)] = acc_frag_0_3.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 8*8)] = acc_frag_0_4.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 8*8)] = acc_frag_0_4.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 8*8)] = acc_frag_0_4.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 8*8)] = acc_frag_0_4.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 9*8)] = acc_frag_0_5.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 9*8)] = acc_frag_0_5.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 9*8)] = acc_frag_0_5.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 9*8)] = acc_frag_0_5.w;
+    data0[wg_c_off + thread_c_off           + 0 + (12*8)] = acc_frag_0_6.x;
+    data0[wg_c_off + thread_c_off           + 1 + (12*8)] = acc_frag_0_6.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (12*8)] = acc_frag_0_6.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (12*8)] = acc_frag_0_6.w;
+    data0[wg_c_off + thread_c_off           + 0 + (13*8)] = acc_frag_0_7.x;
+    data0[wg_c_off + thread_c_off           + 1 + (13*8)] = acc_frag_0_7.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (13*8)] = acc_frag_0_7.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (13*8)] = acc_frag_0_7.w;
+    wg_c_off += 32*N;
+    data0[wg_c_off + thread_c_off           + 0 + ( 0*8)] = acc_frag_1_0.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 0*8)] = acc_frag_1_0.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 0*8)] = acc_frag_1_0.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 0*8)] = acc_frag_1_0.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 1*8)] = acc_frag_1_1.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 1*8)] = acc_frag_1_1.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 1*8)] = acc_frag_1_1.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 1*8)] = acc_frag_1_1.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 4*8)] = acc_frag_1_2.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 4*8)] = acc_frag_1_2.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 4*8)] = acc_frag_1_2.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 4*8)] = acc_frag_1_2.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 5*8)] = acc_frag_1_3.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 5*8)] = acc_frag_1_3.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 5*8)] = acc_frag_1_3.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 5*8)] = acc_frag_1_3.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 8*8)] = acc_frag_1_4.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 8*8)] = acc_frag_1_4.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 8*8)] = acc_frag_1_4.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 8*8)] = acc_frag_1_4.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 9*8)] = acc_frag_1_5.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 9*8)] = acc_frag_1_5.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 9*8)] = acc_frag_1_5.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 9*8)] = acc_frag_1_5.w;
+    data0[wg_c_off + thread_c_off           + 0 + (12*8)] = acc_frag_1_6.x;
+    data0[wg_c_off + thread_c_off           + 1 + (12*8)] = acc_frag_1_6.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (12*8)] = acc_frag_1_6.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (12*8)] = acc_frag_1_6.w;
+    data0[wg_c_off + thread_c_off           + 0 + (13*8)] = acc_frag_1_7.x;
+    data0[wg_c_off + thread_c_off           + 1 + (13*8)] = acc_frag_1_7.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (13*8)] = acc_frag_1_7.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (13*8)] = acc_frag_1_7.w;
+}

tinygrad_repo/extra/gemm/max_kernels/nv.fp16_fp32_fp32.flat_smem_input.cu

@@ -0,0 +1,363 @@
+#define INFINITY (__int_as_float(0x7f800000))
+#define NAN (__int_as_float(0x7fffffff))
+#include <cuda_fp16.h>
+#include <cuda_pipeline.h>
+#define N_PAD 132
+
+struct __align__(8) half4 { half x, y, z, w; };
+__device__ half4 make_half4(half x, half y, half z, half w) { half4 r={x, y, z, w}; return r; }
+
+struct __align__(16) half8 { half x, y, z, w, a, b, c, d; };
+__device__ half8 make_half8(half x, half y, half z, half w, half a, half b, half c, half d) { half8 r={x, y, z, w, a, b, c, d}; return r; }
+
+__device__ void __ldmatrix_a_elems(half8 *regs, half *smem) {
+    uint32_t reg0, reg1, reg2, reg3;
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(reg0), "=r"(reg1), "=r"(reg2), "=r"(reg3)
+        : "l"(__cvta_generic_to_shared(smem))
+    );
+    uint32_t *addr = reinterpret_cast<uint32_t*>(regs);
+    addr[0] = reg0;
+    addr[1] = reg1;
+    addr[2] = reg2;
+    addr[3] = reg3;
+}
+
+__device__ void __ldmatrix_b_elems(half4 *regs_lo, half4 *regs_hi, half *smem) {
+    uint32_t reg0, reg1, reg2, reg3;
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(reg0), "=r"(reg1), "=r"(reg2), "=r"(reg3)
+        : "l"(__cvta_generic_to_shared(smem))
+    );
+    uint32_t *addr_lo = reinterpret_cast<uint32_t*>(regs_lo);
+    uint32_t *addr_hi = reinterpret_cast<uint32_t*>(regs_hi);
+    addr_lo[0] = reg0;
+    addr_lo[1] = reg1;
+    addr_hi[0] = reg2;
+    addr_hi[1] = reg3;
+}
+
+__device__ float4 __WMMA_8_16_16_half_float(half8 a, half4 b, float4 c) {
+    int *a_pk = (int *) (&a), *b_pk = (int *) (&b);
+    asm( "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 { %0, %1, %2, %3 }, { %4, %5, %6, %7 }, { %8, %9 }, { %0, %1, %2, %3 };"
+        : "+f"(c.x), "+f"(c.y), "+f"(c.z), "+f"(c.w) : "r"(a_pk[0]), "r"(a_pk[1]), "r"(a_pk[2]),  "r"(a_pk[3]), "r"(b_pk[0]), "r"(b_pk[1]) );
+    return c;
+}
+
+extern "C" __global__ void __launch_bounds__(128) wmma_example(float* data0, const half* data1, const half* data2, int N, int K) {
+    int grid_m = blockIdx.x;        /* M//64 */
+    int grid_n = blockIdx.y;        /* N//128 */
+    int threads = threadIdx.x;      /* 128 */
+    int wg_m = (threads/64);        // 0 or 1 for 1st and 3rd blocks of b_m=16xb_k=16 vs 2nd and 4th blocks
+    int wg_n = (threads/32)%2;      // 0 or 1 for 1st, 3rd, 5th, 7th blocks of b_n=16xb_k=16 vs 2nd, 4th, 6th, 8th blocks - differs from triton
+    int wg_threads = threads%32;
+    int num_k_blocks = K / 64;
+
+    // load indexes
+    size_t global_a_off = ((grid_m * 64) * K) + ((threads %  8) * 8) + ((threads /  8) * K);
+    size_t global_b_off = (grid_n * 128)      + ((threads % 16) * 8) + ((threads / 16) * N);
+
+    // non-swizzled - should work slowly with bank conflicts
+    size_t store_smem_a_off = ((threads %  8) * 8) + ((threads /  8) *  64);
+    size_t store_smem_b_off = ((threads % 16) * 8) + ((threads / 16) * 128);
+
+    // ldmatrix indices
+    // threads 0-7 are row starts for A, 8-15 for B, 16-23 for C, 24-31 for D
+    // [ A | C ]
+    // [ - + - ]
+    // [ B | D ]
+
+    // unswizzled ldmatrix
+    size_t load_smem_a_0_k_0 = (wg_m * 16 * 64) + ((wg_threads % 8) *  64) + (((wg_threads / 8) % 2) *  64 * 8) + ((wg_threads / 16) * 8);
+    size_t load_smem_a_1_k_0 = load_smem_a_0_k_0 + (32*64);
+    size_t load_smem_b_0_k_0 = (wg_n * 16)      + ((wg_threads % 8) * 128) + (((wg_threads / 8) % 2) * 128 * 8) + ((wg_threads / 16) * 8);
+    size_t load_smem_b_1_k_0 = load_smem_b_0_k_0 + 32;
+    size_t load_smem_b_2_k_0 = load_smem_b_0_k_0 + 64;
+    size_t load_smem_b_3_k_0 = load_smem_b_0_k_0 + 96;
+
+    size_t load_smem_a_0_k_1 = load_smem_a_0_k_0 + 16;
+    size_t load_smem_a_1_k_1 = load_smem_a_1_k_0 + 16;
+    size_t load_smem_b_0_k_1 = load_smem_b_0_k_0 + (16 * 128);
+    size_t load_smem_b_1_k_1 = load_smem_b_1_k_0 + (16 * 128);
+    size_t load_smem_b_2_k_1 = load_smem_b_2_k_0 + (16 * 128);
+    size_t load_smem_b_3_k_1 = load_smem_b_3_k_0 + (16 * 128);
+
+    size_t load_smem_a_0_k_2 = load_smem_a_0_k_0 + 32;
+    size_t load_smem_a_1_k_2 = load_smem_a_1_k_0 + 32;
+    size_t load_smem_b_0_k_2 = load_smem_b_0_k_0 + (32 * 128);
+    size_t load_smem_b_1_k_2 = load_smem_b_1_k_0 + (32 * 128);
+    size_t load_smem_b_2_k_2 = load_smem_b_2_k_0 + (32 * 128);
+    size_t load_smem_b_3_k_2 = load_smem_b_3_k_0 + (32 * 128);
+
+    size_t load_smem_a_0_k_3 = load_smem_a_0_k_0 + 48;
+    size_t load_smem_a_1_k_3 = load_smem_a_1_k_0 + 48;
+    size_t load_smem_b_0_k_3 = load_smem_b_0_k_0 + (48 * 128);
+    size_t load_smem_b_1_k_3 = load_smem_b_1_k_0 + (48 * 128);
+    size_t load_smem_b_2_k_3 = load_smem_b_2_k_0 + (48 * 128);
+    size_t load_smem_b_3_k_3 = load_smem_b_3_k_0 + (48 * 128);
+
+    // create shared mem (A 8192 bytes, B 16384 bytes)
+    __shared__ alignas(16) char smem[24576];
+
+    // create accs (16 WMMAs and 4 output elements each) and zero
+    float4 acc_frag_0_0 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_1 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_2 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_3 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_4 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_5 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_6 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_7 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_0 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_1 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_2 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_3 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_4 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_5 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_6 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_7 = make_float4(0.0f,0.0f,0.0f,0.0f);
+
+    // create registers for block A elements (2)
+    half8 a_frag_0;
+    half8 a_frag_1;
+
+    // create register for block B elements (8)
+    half4 b_frag_0;
+    half4 b_frag_1;
+    half4 b_frag_2;
+    half4 b_frag_3;
+    half4 b_frag_4;
+    half4 b_frag_5;
+    half4 b_frag_6;
+    half4 b_frag_7;
+
+    half *smem_a = (half *)(smem);
+    half *smem_b = (half *)(smem + 8192);
+
+    // https://developer.nvidia.com/blog/controlling-data-movement-to-boost-performance-on-ampere-architecture/
+    // https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#asynchronous-data-copies
+
+    // start first pre-fetch load A
+    __pipeline_memcpy_async(&smem_a[store_smem_a_off +  (    0)], &data1[global_a_off + (   0)], 16);
+    __pipeline_memcpy_async(&smem_a[store_smem_a_off +  (16*64)], &data1[global_a_off + (16*K)], 16);
+    __pipeline_memcpy_async(&smem_a[store_smem_a_off +  (32*64)], &data1[global_a_off + (32*K)], 16);
+    __pipeline_memcpy_async(&smem_a[store_smem_a_off +  (48*64)], &data1[global_a_off + (48*K)], 16);
+
+    // start first pre-fetch load B
+    __pipeline_memcpy_async(&smem_b[store_smem_b_off + (     0)], &data2[global_b_off + (   0)], 16);
+    __pipeline_memcpy_async(&smem_b[store_smem_b_off + ( 8*128)], &data2[global_b_off + ( 8*N)], 16);
+    __pipeline_memcpy_async(&smem_b[store_smem_b_off + (16*128)], &data2[global_b_off + (16*N)], 16);
+    __pipeline_memcpy_async(&smem_b[store_smem_b_off + (24*128)], &data2[global_b_off + (24*N)], 16);
+    __pipeline_memcpy_async(&smem_b[store_smem_b_off + (32*128)], &data2[global_b_off + (32*N)], 16);
+    __pipeline_memcpy_async(&smem_b[store_smem_b_off + (40*128)], &data2[global_b_off + (40*N)], 16);
+    __pipeline_memcpy_async(&smem_b[store_smem_b_off + (48*128)], &data2[global_b_off + (48*N)], 16);
+    __pipeline_memcpy_async(&smem_b[store_smem_b_off + (56*128)], &data2[global_b_off + (56*N)], 16);
+    __pipeline_commit();
+
+    global_a_off += 64;
+    global_b_off += 64 * N;
+    __syncthreads();
+
+    for (int block_k = 0; block_k < num_k_blocks; block_k++) {
+        // wait on needed prefetch value
+        __pipeline_wait_prior(0);
+        __syncthreads();
+
+        // BLOCK_K==4: unroll 4 iterations of ldmatrix/wmma
+        half *smem_a_curr = smem_a;
+        half *smem_b_curr = smem_b;
+
+        // first load 16 K elements and 16 WMMAs: BLOCK_M==2 * BLOCK_N==8
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_0]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_0]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_0]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_0]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_0]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_0]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // next 16 K elements
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_1]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_1]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_1]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_1]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_1]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_1]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // next 16 K elements
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_2]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_2]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_2]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_2]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_2]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_2]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // last 16 K elements
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_3]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_3]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_3]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_3]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_3]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_3]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // prefetch next iteration if needed
+        __syncthreads();
+        if (block_k < (num_k_blocks-1)) {
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (    0)], &data1[global_a_off + (   0)], 16);
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (16*64)], &data1[global_a_off + (16*K)], 16);
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (32*64)], &data1[global_a_off + (32*K)], 16);
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (48*64)], &data1[global_a_off + (48*K)], 16);
+
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (     0)], &data2[global_b_off + (   0)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + ( 8*128)], &data2[global_b_off + ( 8*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (16*128)], &data2[global_b_off + (16*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (24*128)], &data2[global_b_off + (24*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (32*128)], &data2[global_b_off + (32*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (40*128)], &data2[global_b_off + (40*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (48*128)], &data2[global_b_off + (48*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (56*128)], &data2[global_b_off + (56*N)], 16);
+
+            global_a_off += 64;
+            global_b_off += 64 * N;
+        }
+        __pipeline_commit();
+    }
+
+    // write accumulators to output
+    __pipeline_wait_prior(0);
+    __syncthreads();
+
+    // slower way: write floats one by one to data0
+    size_t wg_c_off     = ((grid_m * 64) * N) + (grid_n * 128) + (wg_m * 16 * N) + (wg_n * 16);
+    size_t thread_c_off = ((wg_threads % 4) * 2) + (((wg_threads / 4) % 8) * N);
+    data0[wg_c_off + thread_c_off           + 0 + ( 0*8)] = acc_frag_0_0.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 0*8)] = acc_frag_0_0.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 0*8)] = acc_frag_0_0.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 0*8)] = acc_frag_0_0.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 1*8)] = acc_frag_0_1.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 1*8)] = acc_frag_0_1.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 1*8)] = acc_frag_0_1.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 1*8)] = acc_frag_0_1.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 4*8)] = acc_frag_0_2.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 4*8)] = acc_frag_0_2.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 4*8)] = acc_frag_0_2.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 4*8)] = acc_frag_0_2.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 5*8)] = acc_frag_0_3.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 5*8)] = acc_frag_0_3.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 5*8)] = acc_frag_0_3.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 5*8)] = acc_frag_0_3.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 8*8)] = acc_frag_0_4.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 8*8)] = acc_frag_0_4.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 8*8)] = acc_frag_0_4.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 8*8)] = acc_frag_0_4.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 9*8)] = acc_frag_0_5.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 9*8)] = acc_frag_0_5.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 9*8)] = acc_frag_0_5.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 9*8)] = acc_frag_0_5.w;
+    data0[wg_c_off + thread_c_off           + 0 + (12*8)] = acc_frag_0_6.x;
+    data0[wg_c_off + thread_c_off           + 1 + (12*8)] = acc_frag_0_6.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (12*8)] = acc_frag_0_6.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (12*8)] = acc_frag_0_6.w;
+    data0[wg_c_off + thread_c_off           + 0 + (13*8)] = acc_frag_0_7.x;
+    data0[wg_c_off + thread_c_off           + 1 + (13*8)] = acc_frag_0_7.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (13*8)] = acc_frag_0_7.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (13*8)] = acc_frag_0_7.w;
+    wg_c_off += 32*N;
+    data0[wg_c_off + thread_c_off           + 0 + ( 0*8)] = acc_frag_1_0.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 0*8)] = acc_frag_1_0.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 0*8)] = acc_frag_1_0.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 0*8)] = acc_frag_1_0.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 1*8)] = acc_frag_1_1.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 1*8)] = acc_frag_1_1.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 1*8)] = acc_frag_1_1.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 1*8)] = acc_frag_1_1.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 4*8)] = acc_frag_1_2.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 4*8)] = acc_frag_1_2.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 4*8)] = acc_frag_1_2.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 4*8)] = acc_frag_1_2.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 5*8)] = acc_frag_1_3.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 5*8)] = acc_frag_1_3.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 5*8)] = acc_frag_1_3.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 5*8)] = acc_frag_1_3.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 8*8)] = acc_frag_1_4.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 8*8)] = acc_frag_1_4.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 8*8)] = acc_frag_1_4.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 8*8)] = acc_frag_1_4.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 9*8)] = acc_frag_1_5.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 9*8)] = acc_frag_1_5.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 9*8)] = acc_frag_1_5.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 9*8)] = acc_frag_1_5.w;
+    data0[wg_c_off + thread_c_off           + 0 + (12*8)] = acc_frag_1_6.x;
+    data0[wg_c_off + thread_c_off           + 1 + (12*8)] = acc_frag_1_6.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (12*8)] = acc_frag_1_6.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (12*8)] = acc_frag_1_6.w;
+    data0[wg_c_off + thread_c_off           + 0 + (13*8)] = acc_frag_1_7.x;
+    data0[wg_c_off + thread_c_off           + 1 + (13*8)] = acc_frag_1_7.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (13*8)] = acc_frag_1_7.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (13*8)] = acc_frag_1_7.w;
+}

tinygrad_repo/extra/gemm/max_kernels/nv.fp16_fp32_fp32.max.cu

@@ -0,0 +1,439 @@
+#define INFINITY (__int_as_float(0x7f800000))
+#define NAN (__int_as_float(0x7fffffff))
+#include <cuda_fp16.h>
+#include <cuda_pipeline.h>
+#define N_PAD 132
+
+struct __align__(8) half4 { half x, y, z, w; };
+__device__ half4 make_half4(half x, half y, half z, half w) { half4 r={x, y, z, w}; return r; }
+
+struct __align__(16) half8 { half x, y, z, w, a, b, c, d; };
+__device__ half8 make_half8(half x, half y, half z, half w, half a, half b, half c, half d) { half8 r={x, y, z, w, a, b, c, d}; return r; }
+
+__device__ void __ldmatrix_a_elems(half8 *regs, half *smem) {
+    uint32_t reg0, reg1, reg2, reg3;
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(reg0), "=r"(reg1), "=r"(reg2), "=r"(reg3)
+        : "l"(__cvta_generic_to_shared(smem))
+    );
+    uint32_t *addr = reinterpret_cast<uint32_t*>(regs);
+    addr[0] = reg0;
+    addr[1] = reg1;
+    addr[2] = reg2;
+    addr[3] = reg3;
+}
+
+__device__ void __ldmatrix_b_elems(half4 *regs_lo, half4 *regs_hi, half *smem) {
+    uint32_t reg0, reg1, reg2, reg3;
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(reg0), "=r"(reg1), "=r"(reg2), "=r"(reg3)
+        : "l"(__cvta_generic_to_shared(smem))
+    );
+    uint32_t *addr_lo = reinterpret_cast<uint32_t*>(regs_lo);
+    uint32_t *addr_hi = reinterpret_cast<uint32_t*>(regs_hi);
+    addr_lo[0] = reg0;
+    addr_lo[1] = reg1;
+    addr_hi[0] = reg2;
+    addr_hi[1] = reg3;
+}
+
+__device__ float4 __WMMA_8_16_16_half_float(half8 a, half4 b, float4 c) {
+    int *a_pk = (int *) (&a), *b_pk = (int *) (&b);
+    asm( "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 { %0, %1, %2, %3 }, { %4, %5, %6, %7 }, { %8, %9 }, { %0, %1, %2, %3 };"
+        : "+f"(c.x), "+f"(c.y), "+f"(c.z), "+f"(c.w) : "r"(a_pk[0]), "r"(a_pk[1]), "r"(a_pk[2]),  "r"(a_pk[3]), "r"(b_pk[0]), "r"(b_pk[1]) );
+    return c;
+}
+
+extern "C" __global__ void __launch_bounds__(128) wmma_example(float* data0, const half* data1, const half* data2, int N, int K) {
+    int grid_m = blockIdx.x;        /* M//64 */
+    int grid_n = blockIdx.y;        /* N//128 */
+    int threads = threadIdx.x;      /* 128 */
+    int wg_m = (threads/64);        // 0 or 1 for 1st and 3rd blocks of b_m=16xb_k=16 vs 2nd and 4th blocks
+    int wg_n = (threads/32)%2;      // 0 or 1 for 1st, 3rd, 5th, 7th blocks of b_n=16xb_k=16 vs 2nd, 4th, 6th, 8th blocks - differs from triton
+    int wg_threads = threads%32;
+    int num_k_blocks = K / 64;
+
+    // load indexes
+    size_t global_a_off = ((grid_m * 64) * K) + ((threads %  8) * 8) + ((threads /  8) * K);
+    size_t global_b_off = (grid_n * 128)      + ((threads % 16) * 8) + ((threads / 16) * N);
+
+    // swizzled smem store offsets - columns of smem are swizzled
+    // here's a link to a description of the triton: https://github.com/triton-lang/triton/discussions/2026#discussioncomment-6746579
+    // see also the thunderkittens impl: https://github.com/HazyResearch/ThunderKittens/blob/main/include/types/shared/st.cuh
+    size_t store_smem_a_off  = ((threads /  8) *  64) + (((threads * 8) ^ threads) & 56);               // r15
+    size_t store_smem_b_off  = ((threads / 16) * 128) + (((threads / 16) * 8) ^ ((threads % 16) * 8));  // r19
+
+    // ldmatrix indices - 4x loads of 8x8 matrices by 32 threads
+    // threads 0-7 are row starts for A, 8-15 for B, 16-23 for C, 24-31 for D
+    // [ A | C ]
+    // [ - + - ]
+    // [ B | D ]
+
+    // swizzled ldmatrix
+    size_t load_smem_a_row   = ((wg_m * 16) + (threads % 16)) * 64; // r293
+    size_t load_smem_a_phase = (threads / 16) % 2;                  // r4
+
+    size_t load_smem_b_row   = (threads % 16) * 128;                // r299
+    size_t load_smem_b_phase = (wg_n * 2) + (((threads / 16) % 2)); // r297 -- this differs from the generated triton kernel (swapped order)
+
+    size_t load_smem_a_0_k_0 = load_smem_a_row + (((load_smem_a_phase +  0) ^ (threads % 8)) * 8); // r38
+    size_t load_smem_a_1_k_0 = load_smem_a_0_k_0 + (32 * 64);
+    size_t load_smem_b_0_k_0 = load_smem_b_row + (((load_smem_b_phase +  0) ^ (threads % 8)) * 8);
+    size_t load_smem_b_1_k_0 = load_smem_b_row + (((load_smem_b_phase +  4) ^ (threads % 8)) * 8);
+    size_t load_smem_b_2_k_0 = load_smem_b_row + (((load_smem_b_phase +  8) ^ (threads % 8)) * 8);
+    size_t load_smem_b_3_k_0 = load_smem_b_row + (((load_smem_b_phase + 12) ^ (threads % 8)) * 8);
+
+    size_t load_smem_a_0_k_1 = load_smem_a_row + (((load_smem_a_phase +  2) ^ (threads % 8)) * 8); // r58 = r293 + r316;
+    size_t load_smem_a_1_k_1 = load_smem_a_0_k_1 + (32 * 64);
+    size_t load_smem_b_0_k_1 = load_smem_b_0_k_0 + (16 * 128);
+    size_t load_smem_b_1_k_1 = load_smem_b_1_k_0 + (16 * 128);
+    size_t load_smem_b_2_k_1 = load_smem_b_2_k_0 + (16 * 128);
+    size_t load_smem_b_3_k_1 = load_smem_b_3_k_0 + (16 * 128);
+
+    size_t load_smem_a_0_k_2 = load_smem_a_row + (((load_smem_a_phase +  4) ^ (threads % 8)) * 8); // r59 = r293 + r319;
+    size_t load_smem_a_1_k_2 = load_smem_a_0_k_2 + (32 * 64);
+    size_t load_smem_b_0_k_2 = load_smem_b_0_k_0 + (32 * 128);
+    size_t load_smem_b_1_k_2 = load_smem_b_1_k_0 + (32 * 128);
+    size_t load_smem_b_2_k_2 = load_smem_b_2_k_0 + (32 * 128);
+    size_t load_smem_b_3_k_2 = load_smem_b_3_k_0 + (32 * 128);
+
+    size_t load_smem_a_0_k_3 = load_smem_a_row + (((load_smem_a_phase +  6) ^ (threads % 8)) * 8); // r60 = r293 + r322;
+    size_t load_smem_a_1_k_3 = load_smem_a_0_k_3 + (32 * 64);
+    size_t load_smem_b_0_k_3 = load_smem_b_0_k_0 + (48 * 128);
+    size_t load_smem_b_1_k_3 = load_smem_b_1_k_0 + (48 * 128);
+    size_t load_smem_b_2_k_3 = load_smem_b_2_k_0 + (48 * 128);
+    size_t load_smem_b_3_k_3 = load_smem_b_3_k_0 + (48 * 128);
+
+    // create shared mem (A_1 8192 bytes, A_2 8192 bytes, B_1 16384 bytes, B2_16384 bytes)
+    __shared__ alignas(16) char smem[49152];
+
+    // create accs (16 WMMAs and 4 output elements each) and zero
+    float4 acc_frag_0_0 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_1 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_2 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_3 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_4 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_5 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_6 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_7 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_0 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_1 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_2 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_3 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_4 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_5 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_6 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_7 = make_float4(0.0f,0.0f,0.0f,0.0f);
+
+    // create registers for block A elements (2)
+    half8 a_frag_0;
+    half8 a_frag_1;
+
+    // create register for block B elements (8)
+    half4 b_frag_0;
+    half4 b_frag_1;
+    half4 b_frag_2;
+    half4 b_frag_3;
+    half4 b_frag_4;
+    half4 b_frag_5;
+    half4 b_frag_6;
+    half4 b_frag_7;
+
+    half *smem_a_even = (half *)(smem);
+    half *smem_a_odd  = (half *)(smem + 8192);
+    half *smem_b_even = (half *)(smem + 16384);
+    half *smem_b_odd  = (half *)(smem + 32768);
+
+    // https://developer.nvidia.com/blog/controlling-data-movement-to-boost-performance-on-ampere-architecture/
+    // https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#asynchronous-data-copies
+
+    // start first pre-fetch load A
+    __pipeline_memcpy_async(&smem_a_even[store_smem_a_off +  (    0)], &data1[global_a_off + (   0)], 16);
+    __pipeline_memcpy_async(&smem_a_even[store_smem_a_off +  (16*64)], &data1[global_a_off + (16*K)], 16);
+    __pipeline_memcpy_async(&smem_a_even[store_smem_a_off +  (32*64)], &data1[global_a_off + (32*K)], 16);
+    __pipeline_memcpy_async(&smem_a_even[store_smem_a_off +  (48*64)], &data1[global_a_off + (48*K)], 16);
+
+    // start first pre-fetch load B
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (     0)], &data2[global_b_off + (   0)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + ( 8*128)], &data2[global_b_off + ( 8*N)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (16*128)], &data2[global_b_off + (16*N)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (24*128)], &data2[global_b_off + (24*N)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (32*128)], &data2[global_b_off + (32*N)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (40*128)], &data2[global_b_off + (40*N)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (48*128)], &data2[global_b_off + (48*N)], 16);
+    __pipeline_memcpy_async(&smem_b_even[store_smem_b_off + (56*128)], &data2[global_b_off + (56*N)], 16);
+    __pipeline_commit();
+
+    global_a_off += 64;
+    global_b_off += 64 * N;
+    __syncthreads();
+
+    // start second pre-fetch load A
+    __pipeline_memcpy_async(&smem_a_odd[store_smem_a_off +  (    0)], &data1[global_a_off + (   0)], 16);
+    __pipeline_memcpy_async(&smem_a_odd[store_smem_a_off +  (16*64)], &data1[global_a_off + (16*K)], 16);
+    __pipeline_memcpy_async(&smem_a_odd[store_smem_a_off +  (32*64)], &data1[global_a_off + (32*K)], 16);
+    __pipeline_memcpy_async(&smem_a_odd[store_smem_a_off +  (48*64)], &data1[global_a_off + (48*K)], 16);
+
+    // start second pre-fetch load B
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (     0)], &data2[global_b_off + (   0)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + ( 8*128)], &data2[global_b_off + ( 8*N)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (16*128)], &data2[global_b_off + (16*N)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (24*128)], &data2[global_b_off + (24*N)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (32*128)], &data2[global_b_off + (32*N)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (40*128)], &data2[global_b_off + (40*N)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (48*128)], &data2[global_b_off + (48*N)], 16);
+    __pipeline_memcpy_async(&smem_b_odd[store_smem_b_off + (56*128)], &data2[global_b_off + (56*N)], 16);
+    __pipeline_commit();
+
+    global_a_off += 64;
+    global_b_off += 64 * N;
+
+    // wait on needed prefetch value
+    __pipeline_wait_prior(0); // TODO: this enables fast iterations, but incorrect results with 1 (it shouldn't)
+    __syncthreads();
+
+    for (int block_k = 0; block_k < num_k_blocks; block_k++) {
+        // BLOCK_K==4: unroll 4 iterations of ldmatrix/wmma
+        half *smem_a_curr = (block_k % 2) ? smem_a_even : smem_a_odd;
+        half *smem_b_curr = (block_k % 2) ? smem_b_even : smem_b_odd;
+
+        // first load 16 K elements and 16 WMMAs: BLOCK_M==2 * BLOCK_N==8
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_0]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_0]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_0]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_0]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_0]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_0]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // next 16 K elements
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_1]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_1]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_1]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_1]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_1]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_1]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // next 16 K elements
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_2]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_2]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_2]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_2]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_2]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_2]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // last 16 K elements
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_3]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_3]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_3]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_3]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_3]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_3]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // prefetch next iteration if needed
+        __syncthreads();
+        if (block_k < (num_k_blocks-2)) {
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (    0)], &data1[global_a_off + (   0)], 16);
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (16*64)], &data1[global_a_off + (16*K)], 16);
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (32*64)], &data1[global_a_off + (32*K)], 16);
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (48*64)], &data1[global_a_off + (48*K)], 16);
+
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (     0)], &data2[global_b_off + (   0)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + ( 8*128)], &data2[global_b_off + ( 8*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (16*128)], &data2[global_b_off + (16*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (24*128)], &data2[global_b_off + (24*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (32*128)], &data2[global_b_off + (32*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (40*128)], &data2[global_b_off + (40*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (48*128)], &data2[global_b_off + (48*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (56*128)], &data2[global_b_off + (56*N)], 16);
+
+            global_a_off += 64;
+            global_b_off += 64 * N;
+        }
+        __pipeline_commit();
+
+        if (block_k < num_k_blocks-1) {
+            __pipeline_wait_prior(1);
+            __syncthreads();
+        }
+    }
+
+    // write accumulators to output
+    __pipeline_wait_prior(0);
+    __syncthreads();
+
+    // store registers to smem first, then read back to do float4 writes to global
+    float *smem_d = (float *)(smem);
+    size_t smem_d_off = (wg_m * 16 * N_PAD) + (wg_n * 16) + ((wg_threads % 4) * 2) + (((wg_threads / 4) % 8) * N_PAD);
+    smem_d[smem_d_off + 0 + ( 0*8)            ] = acc_frag_0_0.x;
+    smem_d[smem_d_off + 1 + ( 0*8)            ] = acc_frag_0_0.y;
+    smem_d[smem_d_off + 0 + ( 0*8) + (8*N_PAD)] = acc_frag_0_0.z;
+    smem_d[smem_d_off + 1 + ( 0*8) + (8*N_PAD)] = acc_frag_0_0.w;
+    smem_d[smem_d_off + 0 + ( 1*8)            ] = acc_frag_0_1.x;
+    smem_d[smem_d_off + 1 + ( 1*8)            ] = acc_frag_0_1.y;
+    smem_d[smem_d_off + 0 + ( 1*8) + (8*N_PAD)] = acc_frag_0_1.z;
+    smem_d[smem_d_off + 1 + ( 1*8) + (8*N_PAD)] = acc_frag_0_1.w;
+    smem_d[smem_d_off + 0 + ( 4*8)            ] = acc_frag_0_2.x;
+    smem_d[smem_d_off + 1 + ( 4*8)            ] = acc_frag_0_2.y;
+    smem_d[smem_d_off + 0 + ( 4*8) + (8*N_PAD)] = acc_frag_0_2.z;
+    smem_d[smem_d_off + 1 + ( 4*8) + (8*N_PAD)] = acc_frag_0_2.w;
+    smem_d[smem_d_off + 0 + ( 5*8)            ] = acc_frag_0_3.x;
+    smem_d[smem_d_off + 1 + ( 5*8)            ] = acc_frag_0_3.y;
+    smem_d[smem_d_off + 0 + ( 5*8) + (8*N_PAD)] = acc_frag_0_3.z;
+    smem_d[smem_d_off + 1 + ( 5*8) + (8*N_PAD)] = acc_frag_0_3.w;
+    smem_d[smem_d_off + 0 + ( 8*8)            ] = acc_frag_0_4.x;
+    smem_d[smem_d_off + 1 + ( 8*8)            ] = acc_frag_0_4.y;
+    smem_d[smem_d_off + 0 + ( 8*8) + (8*N_PAD)] = acc_frag_0_4.z;
+    smem_d[smem_d_off + 1 + ( 8*8) + (8*N_PAD)] = acc_frag_0_4.w;
+    smem_d[smem_d_off + 0 + ( 9*8)            ] = acc_frag_0_5.x;
+    smem_d[smem_d_off + 1 + ( 9*8)            ] = acc_frag_0_5.y;
+    smem_d[smem_d_off + 0 + ( 9*8) + (8*N_PAD)] = acc_frag_0_5.z;
+    smem_d[smem_d_off + 1 + ( 9*8) + (8*N_PAD)] = acc_frag_0_5.w;
+    smem_d[smem_d_off + 0 + (12*8)            ] = acc_frag_0_6.x;
+    smem_d[smem_d_off + 1 + (12*8)            ] = acc_frag_0_6.y;
+    smem_d[smem_d_off + 0 + (12*8) + (8*N_PAD)] = acc_frag_0_6.z;
+    smem_d[smem_d_off + 1 + (12*8) + (8*N_PAD)] = acc_frag_0_6.w;
+    smem_d[smem_d_off + 0 + (13*8)            ] = acc_frag_0_7.x;
+    smem_d[smem_d_off + 1 + (13*8)            ] = acc_frag_0_7.y;
+    smem_d[smem_d_off + 0 + (13*8) + (8*N_PAD)] = acc_frag_0_7.z;
+    smem_d[smem_d_off + 1 + (13*8) + (8*N_PAD)] = acc_frag_0_7.w;
+
+    __syncthreads();
+    size_t load_smem_d_off = ((threads % 32) * 4) + ((threads / 32) * N_PAD);
+    float4 d_0_0 = *((float4 *)(smem_d + load_smem_d_off + ( 0 * N_PAD)));
+    float4 d_0_1 = *((float4 *)(smem_d + load_smem_d_off + ( 4 * N_PAD)));
+    float4 d_0_2 = *((float4 *)(smem_d + load_smem_d_off + ( 8 * N_PAD)));
+    float4 d_0_3 = *((float4 *)(smem_d + load_smem_d_off + (12 * N_PAD)));
+    float4 d_0_4 = *((float4 *)(smem_d + load_smem_d_off + (16 * N_PAD)));
+    float4 d_0_5 = *((float4 *)(smem_d + load_smem_d_off + (20 * N_PAD)));
+    float4 d_0_6 = *((float4 *)(smem_d + load_smem_d_off + (24 * N_PAD)));
+    float4 d_0_7 = *((float4 *)(smem_d + load_smem_d_off + (28 * N_PAD)));
+
+    __syncthreads();
+    smem_d[smem_d_off + 0 + ( 0*8)            ] = acc_frag_1_0.x;
+    smem_d[smem_d_off + 1 + ( 0*8)            ] = acc_frag_1_0.y;
+    smem_d[smem_d_off + 0 + ( 0*8) + (8*N_PAD)] = acc_frag_1_0.z;
+    smem_d[smem_d_off + 1 + ( 0*8) + (8*N_PAD)] = acc_frag_1_0.w;
+    smem_d[smem_d_off + 0 + ( 1*8)            ] = acc_frag_1_1.x;
+    smem_d[smem_d_off + 1 + ( 1*8)            ] = acc_frag_1_1.y;
+    smem_d[smem_d_off + 0 + ( 1*8) + (8*N_PAD)] = acc_frag_1_1.z;
+    smem_d[smem_d_off + 1 + ( 1*8) + (8*N_PAD)] = acc_frag_1_1.w;
+    smem_d[smem_d_off + 0 + ( 4*8)            ] = acc_frag_1_2.x;
+    smem_d[smem_d_off + 1 + ( 4*8)            ] = acc_frag_1_2.y;
+    smem_d[smem_d_off + 0 + ( 4*8) + (8*N_PAD)] = acc_frag_1_2.z;
+    smem_d[smem_d_off + 1 + ( 4*8) + (8*N_PAD)] = acc_frag_1_2.w;
+    smem_d[smem_d_off + 0 + ( 5*8)            ] = acc_frag_1_3.x;
+    smem_d[smem_d_off + 1 + ( 5*8)            ] = acc_frag_1_3.y;
+    smem_d[smem_d_off + 0 + ( 5*8) + (8*N_PAD)] = acc_frag_1_3.z;
+    smem_d[smem_d_off + 1 + ( 5*8) + (8*N_PAD)] = acc_frag_1_3.w;
+    smem_d[smem_d_off + 0 + ( 8*8)            ] = acc_frag_1_4.x;
+    smem_d[smem_d_off + 1 + ( 8*8)            ] = acc_frag_1_4.y;
+    smem_d[smem_d_off + 0 + ( 8*8) + (8*N_PAD)] = acc_frag_1_4.z;
+    smem_d[smem_d_off + 1 + ( 8*8) + (8*N_PAD)] = acc_frag_1_4.w;
+    smem_d[smem_d_off + 0 + ( 9*8)            ] = acc_frag_1_5.x;
+    smem_d[smem_d_off + 1 + ( 9*8)            ] = acc_frag_1_5.y;
+    smem_d[smem_d_off + 0 + ( 9*8) + (8*N_PAD)] = acc_frag_1_5.z;
+    smem_d[smem_d_off + 1 + ( 9*8) + (8*N_PAD)] = acc_frag_1_5.w;
+    smem_d[smem_d_off + 0 + (12*8)            ] = acc_frag_1_6.x;
+    smem_d[smem_d_off + 1 + (12*8)            ] = acc_frag_1_6.y;
+    smem_d[smem_d_off + 0 + (12*8) + (8*N_PAD)] = acc_frag_1_6.z;
+    smem_d[smem_d_off + 1 + (12*8) + (8*N_PAD)] = acc_frag_1_6.w;
+    smem_d[smem_d_off + 0 + (13*8)            ] = acc_frag_1_7.x;
+    smem_d[smem_d_off + 1 + (13*8)            ] = acc_frag_1_7.y;
+    smem_d[smem_d_off + 0 + (13*8) + (8*N_PAD)] = acc_frag_1_7.z;
+    smem_d[smem_d_off + 1 + (13*8) + (8*N_PAD)] = acc_frag_1_7.w;
+
+    __syncthreads();
+    float4 d_1_0 = *((float4 *)(smem_d + load_smem_d_off + ( 0 * N_PAD)));
+    float4 d_1_1 = *((float4 *)(smem_d + load_smem_d_off + ( 4 * N_PAD)));
+    float4 d_1_2 = *((float4 *)(smem_d + load_smem_d_off + ( 8 * N_PAD)));
+    float4 d_1_3 = *((float4 *)(smem_d + load_smem_d_off + (12 * N_PAD)));
+    float4 d_1_4 = *((float4 *)(smem_d + load_smem_d_off + (16 * N_PAD)));
+    float4 d_1_5 = *((float4 *)(smem_d + load_smem_d_off + (20 * N_PAD)));
+    float4 d_1_6 = *((float4 *)(smem_d + load_smem_d_off + (24 * N_PAD)));
+    float4 d_1_7 = *((float4 *)(smem_d + load_smem_d_off + (28 * N_PAD)));
+
+    __syncthreads();
+    float *global_d = &data0[((grid_m * 64) * N) + (grid_n * 128) + ((threads % 32) * 4) + ((threads / 32) * N)];
+    *((float4 *)(global_d +  0*N)) = d_0_0;
+    *((float4 *)(global_d +  4*N)) = d_0_1;
+    *((float4 *)(global_d +  8*N)) = d_0_2;
+    *((float4 *)(global_d + 12*N)) = d_0_3;
+    *((float4 *)(global_d + 16*N)) = d_0_4;
+    *((float4 *)(global_d + 20*N)) = d_0_5;
+    *((float4 *)(global_d + 24*N)) = d_0_6;
+    *((float4 *)(global_d + 28*N)) = d_0_7;
+    *((float4 *)(global_d + 32*N)) = d_1_0;
+    *((float4 *)(global_d + 36*N)) = d_1_1;
+    *((float4 *)(global_d + 40*N)) = d_1_2;
+    *((float4 *)(global_d + 44*N)) = d_1_3;
+    *((float4 *)(global_d + 48*N)) = d_1_4;
+    *((float4 *)(global_d + 52*N)) = d_1_5;
+    *((float4 *)(global_d + 56*N)) = d_1_6;
+    *((float4 *)(global_d + 60*N)) = d_1_7;
+}

tinygrad_repo/extra/gemm/max_kernels/nv.fp16_fp32_fp32.swizzled_smem_input.cu

@@ -0,0 +1,371 @@
+#define INFINITY (__int_as_float(0x7f800000))
+#define NAN (__int_as_float(0x7fffffff))
+#include <cuda_fp16.h>
+#include <cuda_pipeline.h>
+#define N_PAD 132
+
+struct __align__(8) half4 { half x, y, z, w; };
+__device__ half4 make_half4(half x, half y, half z, half w) { half4 r={x, y, z, w}; return r; }
+
+struct __align__(16) half8 { half x, y, z, w, a, b, c, d; };
+__device__ half8 make_half8(half x, half y, half z, half w, half a, half b, half c, half d) { half8 r={x, y, z, w, a, b, c, d}; return r; }
+
+__device__ void __ldmatrix_a_elems(half8 *regs, half *smem) {
+    uint32_t reg0, reg1, reg2, reg3;
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(reg0), "=r"(reg1), "=r"(reg2), "=r"(reg3)
+        : "l"(__cvta_generic_to_shared(smem))
+    );
+    uint32_t *addr = reinterpret_cast<uint32_t*>(regs);
+    addr[0] = reg0;
+    addr[1] = reg1;
+    addr[2] = reg2;
+    addr[3] = reg3;
+}
+
+__device__ void __ldmatrix_b_elems(half4 *regs_lo, half4 *regs_hi, half *smem) {
+    uint32_t reg0, reg1, reg2, reg3;
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(reg0), "=r"(reg1), "=r"(reg2), "=r"(reg3)
+        : "l"(__cvta_generic_to_shared(smem))
+    );
+    uint32_t *addr_lo = reinterpret_cast<uint32_t*>(regs_lo);
+    uint32_t *addr_hi = reinterpret_cast<uint32_t*>(regs_hi);
+    addr_lo[0] = reg0;
+    addr_lo[1] = reg1;
+    addr_hi[0] = reg2;
+    addr_hi[1] = reg3;
+}
+
+__device__ float4 __WMMA_8_16_16_half_float(half8 a, half4 b, float4 c) {
+    int *a_pk = (int *) (&a), *b_pk = (int *) (&b);
+    asm( "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 { %0, %1, %2, %3 }, { %4, %5, %6, %7 }, { %8, %9 }, { %0, %1, %2, %3 };"
+        : "+f"(c.x), "+f"(c.y), "+f"(c.z), "+f"(c.w) : "r"(a_pk[0]), "r"(a_pk[1]), "r"(a_pk[2]),  "r"(a_pk[3]), "r"(b_pk[0]), "r"(b_pk[1]) );
+    return c;
+}
+
+extern "C" __global__ void __launch_bounds__(128) wmma_example(float* data0, const half* data1, const half* data2, int N, int K) {
+    int grid_m = blockIdx.x;        /* M//64 */
+    int grid_n = blockIdx.y;        /* N//128 */
+    int threads = threadIdx.x;      /* 128 */
+    int wg_m = (threads/64);        // 0 or 1 for 1st and 3rd blocks of b_m=16xb_k=16 vs 2nd and 4th blocks
+    int wg_n = (threads/32)%2;      // 0 or 1 for 1st, 3rd, 5th, 7th blocks of b_n=16xb_k=16 vs 2nd, 4th, 6th, 8th blocks - differs from triton
+    int wg_threads = threads%32;
+    int num_k_blocks = K / 64;
+
+    // load indexes
+    size_t global_a_off = ((grid_m * 64) * K) + ((threads %  8) * 8) + ((threads /  8) * K);
+    size_t global_b_off = (grid_n * 128)      + ((threads % 16) * 8) + ((threads / 16) * N);
+
+    // swizzled smem store offsets - columns of smem are swizzled
+    // here's a link to a description of the triton: https://github.com/triton-lang/triton/discussions/2026#discussioncomment-6746579
+    // see also the thunderkittens impl: https://github.com/HazyResearch/ThunderKittens/blob/main/include/types/shared/st.cuh
+    size_t store_smem_a_off  = ((threads /  8) *  64) + (((threads * 8) ^ threads) & 56);               // r15
+    size_t store_smem_b_off  = ((threads / 16) * 128) + (((threads / 16) * 8) ^ ((threads % 16) * 8));  // r19
+
+    // ldmatrix indices
+    // threads 0-7 are row starts for A, 8-15 for B, 16-23 for C, 24-31 for D
+    // [ A | C ]
+    // [ - + - ]
+    // [ B | D ]
+
+    // swizzled ldmatrix
+    size_t load_smem_a_row   = ((wg_m * 16) + (threads % 16)) * 64; // r293
+    size_t load_smem_a_phase = (threads / 16) % 2;                  // r4
+
+    size_t load_smem_b_row   = (threads % 16) * 128;                // r299
+    size_t load_smem_b_phase = (wg_n * 2) + (((threads / 16) % 2)); // r297 -- this differs from the generated triton kernel (swapped order)
+
+    size_t load_smem_a_0_k_0 = load_smem_a_row + (((load_smem_a_phase +  0) ^ (threads % 8)) * 8); // r38
+    size_t load_smem_a_1_k_0 = load_smem_a_0_k_0 + (32 * 64);
+    size_t load_smem_b_0_k_0 = load_smem_b_row + (((load_smem_b_phase +  0) ^ (threads % 8)) * 8);
+    size_t load_smem_b_1_k_0 = load_smem_b_row + (((load_smem_b_phase +  4) ^ (threads % 8)) * 8);
+    size_t load_smem_b_2_k_0 = load_smem_b_row + (((load_smem_b_phase +  8) ^ (threads % 8)) * 8);
+    size_t load_smem_b_3_k_0 = load_smem_b_row + (((load_smem_b_phase + 12) ^ (threads % 8)) * 8);
+
+    size_t load_smem_a_0_k_1 = load_smem_a_row + (((load_smem_a_phase +  2) ^ (threads % 8)) * 8); // r58 = r293 + r316;
+    size_t load_smem_a_1_k_1 = load_smem_a_0_k_1 + (32 * 64);
+    size_t load_smem_b_0_k_1 = load_smem_b_0_k_0 + (16 * 128);
+    size_t load_smem_b_1_k_1 = load_smem_b_1_k_0 + (16 * 128);
+    size_t load_smem_b_2_k_1 = load_smem_b_2_k_0 + (16 * 128);
+    size_t load_smem_b_3_k_1 = load_smem_b_3_k_0 + (16 * 128);
+
+    size_t load_smem_a_0_k_2 = load_smem_a_row + (((load_smem_a_phase +  4) ^ (threads % 8)) * 8); // r59 = r293 + r319;
+    size_t load_smem_a_1_k_2 = load_smem_a_0_k_2 + (32 * 64);
+    size_t load_smem_b_0_k_2 = load_smem_b_0_k_0 + (32 * 128);
+    size_t load_smem_b_1_k_2 = load_smem_b_1_k_0 + (32 * 128);
+    size_t load_smem_b_2_k_2 = load_smem_b_2_k_0 + (32 * 128);
+    size_t load_smem_b_3_k_2 = load_smem_b_3_k_0 + (32 * 128);
+
+    size_t load_smem_a_0_k_3 = load_smem_a_row + (((load_smem_a_phase +  6) ^ (threads % 8)) * 8); // r60 = r293 + r322;
+    size_t load_smem_a_1_k_3 = load_smem_a_0_k_3 + (32 * 64);
+    size_t load_smem_b_0_k_3 = load_smem_b_0_k_0 + (48 * 128);
+    size_t load_smem_b_1_k_3 = load_smem_b_1_k_0 + (48 * 128);
+    size_t load_smem_b_2_k_3 = load_smem_b_2_k_0 + (48 * 128);
+    size_t load_smem_b_3_k_3 = load_smem_b_3_k_0 + (48 * 128);
+
+    // create shared mem (A 8192 bytes, B 16384 bytes)
+    __shared__ alignas(16) char smem[24576];
+
+    // create accs (16 WMMAs and 4 output elements each) and zero
+    float4 acc_frag_0_0 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_1 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_2 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_3 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_4 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_5 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_6 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_0_7 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_0 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_1 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_2 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_3 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_4 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_5 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_6 = make_float4(0.0f,0.0f,0.0f,0.0f);
+    float4 acc_frag_1_7 = make_float4(0.0f,0.0f,0.0f,0.0f);
+
+    // create registers for block A elements (2)
+    half8 a_frag_0;
+    half8 a_frag_1;
+
+    // create register for block B elements (8)
+    half4 b_frag_0;
+    half4 b_frag_1;
+    half4 b_frag_2;
+    half4 b_frag_3;
+    half4 b_frag_4;
+    half4 b_frag_5;
+    half4 b_frag_6;
+    half4 b_frag_7;
+
+    half *smem_a = (half *)(smem);
+    half *smem_b = (half *)(smem + 8192);
+
+    // https://developer.nvidia.com/blog/controlling-data-movement-to-boost-performance-on-ampere-architecture/
+    // https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#asynchronous-data-copies
+
+    // start first pre-fetch load A
+    __pipeline_memcpy_async(&smem_a[store_smem_a_off +  (    0)], &data1[global_a_off + (   0)], 16);
+    __pipeline_memcpy_async(&smem_a[store_smem_a_off +  (16*64)], &data1[global_a_off + (16*K)], 16);
+    __pipeline_memcpy_async(&smem_a[store_smem_a_off +  (32*64)], &data1[global_a_off + (32*K)], 16);
+    __pipeline_memcpy_async(&smem_a[store_smem_a_off +  (48*64)], &data1[global_a_off + (48*K)], 16);
+
+    // start first pre-fetch load B
+    __pipeline_memcpy_async(&smem_b[store_smem_b_off + (     0)], &data2[global_b_off + (   0)], 16);
+    __pipeline_memcpy_async(&smem_b[store_smem_b_off + ( 8*128)], &data2[global_b_off + ( 8*N)], 16);
+    __pipeline_memcpy_async(&smem_b[store_smem_b_off + (16*128)], &data2[global_b_off + (16*N)], 16);
+    __pipeline_memcpy_async(&smem_b[store_smem_b_off + (24*128)], &data2[global_b_off + (24*N)], 16);
+    __pipeline_memcpy_async(&smem_b[store_smem_b_off + (32*128)], &data2[global_b_off + (32*N)], 16);
+    __pipeline_memcpy_async(&smem_b[store_smem_b_off + (40*128)], &data2[global_b_off + (40*N)], 16);
+    __pipeline_memcpy_async(&smem_b[store_smem_b_off + (48*128)], &data2[global_b_off + (48*N)], 16);
+    __pipeline_memcpy_async(&smem_b[store_smem_b_off + (56*128)], &data2[global_b_off + (56*N)], 16);
+    __pipeline_commit();
+
+    global_a_off += 64;
+    global_b_off += 64 * N;
+    __syncthreads();
+
+    for (int block_k = 0; block_k < num_k_blocks; block_k++) {
+        // wait on needed prefetch value
+        __pipeline_wait_prior(0);
+        __syncthreads();
+
+        // BLOCK_K==4: unroll 4 iterations of ldmatrix/wmma
+        half *smem_a_curr = smem_a;
+        half *smem_b_curr = smem_b;
+
+        // first load 16 K elements and 16 WMMAs: BLOCK_M==2 * BLOCK_N==8
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_0]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_0]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_0]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_0]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_0]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_0]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // next 16 K elements
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_1]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_1]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_1]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_1]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_1]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_1]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // next 16 K elements
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_2]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_2]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_2]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_2]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_2]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_2]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // last 16 K elements
+        __ldmatrix_a_elems(&a_frag_0,            &smem_a_curr[load_smem_a_0_k_3]);
+        __ldmatrix_a_elems(&a_frag_1,            &smem_a_curr[load_smem_a_1_k_3]);
+        __ldmatrix_b_elems(&b_frag_0, &b_frag_1, &smem_b_curr[load_smem_b_0_k_3]);
+        __ldmatrix_b_elems(&b_frag_2, &b_frag_3, &smem_b_curr[load_smem_b_1_k_3]);
+        __ldmatrix_b_elems(&b_frag_4, &b_frag_5, &smem_b_curr[load_smem_b_2_k_3]);
+        __ldmatrix_b_elems(&b_frag_6, &b_frag_7, &smem_b_curr[load_smem_b_3_k_3]);
+        acc_frag_0_0 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_0, acc_frag_0_0);
+        acc_frag_0_1 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_1, acc_frag_0_1);
+        acc_frag_0_2 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_2, acc_frag_0_2);
+        acc_frag_0_3 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_3, acc_frag_0_3);
+        acc_frag_0_4 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_4, acc_frag_0_4);
+        acc_frag_0_5 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_5, acc_frag_0_5);
+        acc_frag_0_6 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_6, acc_frag_0_6);
+        acc_frag_0_7 = __WMMA_8_16_16_half_float(a_frag_0, b_frag_7, acc_frag_0_7);
+        acc_frag_1_0 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_0, acc_frag_1_0);
+        acc_frag_1_1 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_1, acc_frag_1_1);
+        acc_frag_1_2 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_2, acc_frag_1_2);
+        acc_frag_1_3 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_3, acc_frag_1_3);
+        acc_frag_1_4 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_4, acc_frag_1_4);
+        acc_frag_1_5 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_5, acc_frag_1_5);
+        acc_frag_1_6 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_6, acc_frag_1_6);
+        acc_frag_1_7 = __WMMA_8_16_16_half_float(a_frag_1, b_frag_7, acc_frag_1_7);
+
+        // prefetch next iteration if needed
+        __syncthreads();
+        if (block_k < (num_k_blocks-1)) {
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (    0)], &data1[global_a_off + (   0)], 16);
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (16*64)], &data1[global_a_off + (16*K)], 16);
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (32*64)], &data1[global_a_off + (32*K)], 16);
+            __pipeline_memcpy_async(&smem_a_curr[store_smem_a_off +  (48*64)], &data1[global_a_off + (48*K)], 16);
+
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (     0)], &data2[global_b_off + (   0)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + ( 8*128)], &data2[global_b_off + ( 8*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (16*128)], &data2[global_b_off + (16*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (24*128)], &data2[global_b_off + (24*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (32*128)], &data2[global_b_off + (32*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (40*128)], &data2[global_b_off + (40*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (48*128)], &data2[global_b_off + (48*N)], 16);
+            __pipeline_memcpy_async(&smem_b_curr[store_smem_b_off + (56*128)], &data2[global_b_off + (56*N)], 16);
+
+            global_a_off += 64;
+            global_b_off += 64 * N;
+        }
+        __pipeline_commit();
+    }
+
+    // write accumulators to output
+    __pipeline_wait_prior(0);
+    __syncthreads();
+
+    // slower way: write floats one by one to data0
+    size_t wg_c_off     = ((grid_m * 64) * N) + (grid_n * 128) + (wg_m * 16 * N) + (wg_n * 16);
+    size_t thread_c_off = ((wg_threads % 4) * 2) + (((wg_threads / 4) % 8) * N);
+    data0[wg_c_off + thread_c_off           + 0 + ( 0*8)] = acc_frag_0_0.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 0*8)] = acc_frag_0_0.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 0*8)] = acc_frag_0_0.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 0*8)] = acc_frag_0_0.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 1*8)] = acc_frag_0_1.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 1*8)] = acc_frag_0_1.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 1*8)] = acc_frag_0_1.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 1*8)] = acc_frag_0_1.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 4*8)] = acc_frag_0_2.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 4*8)] = acc_frag_0_2.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 4*8)] = acc_frag_0_2.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 4*8)] = acc_frag_0_2.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 5*8)] = acc_frag_0_3.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 5*8)] = acc_frag_0_3.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 5*8)] = acc_frag_0_3.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 5*8)] = acc_frag_0_3.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 8*8)] = acc_frag_0_4.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 8*8)] = acc_frag_0_4.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 8*8)] = acc_frag_0_4.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 8*8)] = acc_frag_0_4.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 9*8)] = acc_frag_0_5.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 9*8)] = acc_frag_0_5.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 9*8)] = acc_frag_0_5.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 9*8)] = acc_frag_0_5.w;
+    data0[wg_c_off + thread_c_off           + 0 + (12*8)] = acc_frag_0_6.x;
+    data0[wg_c_off + thread_c_off           + 1 + (12*8)] = acc_frag_0_6.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (12*8)] = acc_frag_0_6.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (12*8)] = acc_frag_0_6.w;
+    data0[wg_c_off + thread_c_off           + 0 + (13*8)] = acc_frag_0_7.x;
+    data0[wg_c_off + thread_c_off           + 1 + (13*8)] = acc_frag_0_7.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (13*8)] = acc_frag_0_7.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (13*8)] = acc_frag_0_7.w;
+    wg_c_off += 32*N;
+    data0[wg_c_off + thread_c_off           + 0 + ( 0*8)] = acc_frag_1_0.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 0*8)] = acc_frag_1_0.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 0*8)] = acc_frag_1_0.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 0*8)] = acc_frag_1_0.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 1*8)] = acc_frag_1_1.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 1*8)] = acc_frag_1_1.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 1*8)] = acc_frag_1_1.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 1*8)] = acc_frag_1_1.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 4*8)] = acc_frag_1_2.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 4*8)] = acc_frag_1_2.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 4*8)] = acc_frag_1_2.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 4*8)] = acc_frag_1_2.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 5*8)] = acc_frag_1_3.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 5*8)] = acc_frag_1_3.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 5*8)] = acc_frag_1_3.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 5*8)] = acc_frag_1_3.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 8*8)] = acc_frag_1_4.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 8*8)] = acc_frag_1_4.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 8*8)] = acc_frag_1_4.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 8*8)] = acc_frag_1_4.w;
+    data0[wg_c_off + thread_c_off           + 0 + ( 9*8)] = acc_frag_1_5.x;
+    data0[wg_c_off + thread_c_off           + 1 + ( 9*8)] = acc_frag_1_5.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + ( 9*8)] = acc_frag_1_5.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + ( 9*8)] = acc_frag_1_5.w;
+    data0[wg_c_off + thread_c_off           + 0 + (12*8)] = acc_frag_1_6.x;
+    data0[wg_c_off + thread_c_off           + 1 + (12*8)] = acc_frag_1_6.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (12*8)] = acc_frag_1_6.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (12*8)] = acc_frag_1_6.w;
+    data0[wg_c_off + thread_c_off           + 0 + (13*8)] = acc_frag_1_7.x;
+    data0[wg_c_off + thread_c_off           + 1 + (13*8)] = acc_frag_1_7.y;
+    data0[wg_c_off + thread_c_off + (8 * N) + 0 + (13*8)] = acc_frag_1_7.z;
+    data0[wg_c_off + thread_c_off + (8 * N) + 1 + (13*8)] = acc_frag_1_7.w;
+}

tinygrad_repo/extra/gemm/max_matmul.py

@@ -0,0 +1,218 @@
+import numpy as np, os
+from tinygrad.helpers import getenv, flat_mv
+from tinygrad import dtypes
+
+# for copied uops
+from tinygrad import dtypes
+from tinygrad.dtype import DTYPES_DICT
+
+script_dir = os.path.dirname(os.path.abspath(__file__))
+
+# problem variations
+DTYPE_IN = DTYPES_DICT[getenv("DTYPE_IN", "half")]
+DTYPE_OUT = DTYPES_DICT[getenv("DTYPE_OUT", "half")]
+DTYPE_ACC = DTYPES_DICT[getenv("DTYPE_ACC", "float")]
+N = getenv("N", 4096)
+M = getenv("M", N)
+K = getenv("K", N)
+CNT = getenv("CNT", 10)
+ATOL = getenv("ATOL", 5e-3 if DTYPE_IN == dtypes.float else 1e-2)
+RTOL = getenv("RTOL", 1e-4 if DTYPE_IN == dtypes.float else 1e-3)
+FLOPS = M * N * K * 2
+BW = 2 * ((M*K) + (K*N) + (M*N))
+
+# algorithm variations
+INPUT = getenv("INPUT", "RAND")
+GEMM_VARIATION = getenv("GEMM_VARIATION", "nv_hcopt")
+
+def randoms():
+  if INPUT == "RAND":
+    na = np.random.default_rng().normal(scale=1.0, size=(M,K)).astype(dtype=np.float32)
+    nb = np.random.default_rng().normal(scale=1.0, size=(K,N)).astype(dtype=np.float32)
+  elif INPUT == "IDENTITY" and M==N==K:
+    na = np.identity(K, dtype=np.float32)
+    nb = np.identity(K, dtype=np.float32)
+  elif INPUT == "OUTPUTONES" and M==K:
+    na = np.identity(K, dtype=np.float32)
+    nb = np.ones((K,N), dtype=np.float32)
+  else:
+    na = np.ones((M,K), dtype=np.float32)
+    nb = np.ones((K,N), dtype=np.float32)
+  nc = np.zeros(M*N, np.float32)
+  if DTYPE_IN != dtypes.float:
+    na = na.astype(np.bfloat16 if DTYPE_IN == dtypes.bfloat16 else np.float16)
+    nb = nb.astype(np.bfloat16 if DTYPE_IN == dtypes.bfloat16 else np.float16)
+  if DTYPE_OUT != dtypes.float:
+    nc = nc.astype(np.bfloat16 if DTYPE_IN == dtypes.bfloat16 else np.float16)
+  return na, nb, nc
+
+if __name__ == "__main__":
+  print(f"gemm variation: {GEMM_VARIATION=} {M=} {N=} {K=} {DTYPE_IN=} {DTYPE_OUT=} {DTYPE_ACC=}")
+  prog, global_size, local_size = None, None, None
+
+  if getenv("CUDA") == 1:
+    from tinygrad.runtime.ops_cuda import CUDAAllocator, CUDADevice, CUDAProgram, CUDACompiler
+    device = CUDADevice("cuda:0")
+    compiler = CUDACompiler(device.arch)
+    cudaalloc = CUDAAllocator(device)
+
+    a = cudaalloc.alloc(M*K*DTYPE_IN.itemsize)
+    b = cudaalloc.alloc(K*N*DTYPE_IN.itemsize)
+    c = cudaalloc.alloc(M*N*DTYPE_OUT.itemsize)
+
+    if GEMM_VARIATION == "max" and (M%64)==0 and (N%128)==0 and (K%64)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.float and DTYPE_ACC == dtypes.float:
+      print("Using CUDA and triton-generated kernel")
+      # See nv_triton_gemm.annotated.ptx for PTX code which was generated from `PYTHONPATH=. DEBUG=6 CUDA=1 CUDA_PTX=1 python3 extra/gemm/triton_nv_matmul.py`
+      # this kernel with M=N=K=4096 does 162TFLOPS, vs torch at 144TFLOPS and BEAM=8 tinygrad at 138TFLOPS.  theo max is 165TFLOPS.
+
+      # WMMA element size is (M, N, K) = (16, 8, 16)
+      # warpgroup size in WMMA tiles is (B_M, B_N, B_K) = (2, 8, 4) so 64 HMMA calls per threadgroup reduce iteration
+      # thread block size is (T_M, T_N, T_K) = (2, 2, 1), i.e. macro blocks in M and N, so 256 HMMA calls per kernel reduce iteration
+      # kernel reduce iteration size in elements = (64, 128, 64)
+      # single iteration SMEM_A = (64 * 64) * (2 bytes / half) =  8192 bytes, SMEM_B = (128 * 64) * (2 bytes / half) = 16384 bytes
+      # double-buffer smem = (8192 + 16384) * 2 = 49152 bytes
+      # reduce for_loop size = [1, 1, (4096 // 16 // 4)==64]
+       # NOTE: T_K > 0 would be group_for_reduce
+      prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp32_fp32.max.cu')).read()))
+      args = (c, a, b)
+      kwargs = {
+        'global_size': [M//64, N//128, 1],
+        'local_size': [128, 1, 1], # 4 warpgroups == (T_M:=2) * (T_N:=2)
+        'wait': True,
+        'vals': (N, K),
+      }
+    elif GEMM_VARIATION == "2_stage_swizzled_smem_input" and (M%64)==0 and (N%128)==0 and (K%64)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.float and DTYPE_ACC == dtypes.float:
+      print("Using CUDA, 2-stage reduce pipeline, swizzled SMEM inputs")
+      prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp32_fp32.2_stage_swizzled_smem_input.cu')).read()))
+      args = (c, a, b)
+      kwargs = {
+        'global_size': [M//64, N//128, 1],
+        'local_size': [128, 1, 1], # 4 warpgroups == (T_M:=2) * (T_N:=2)
+        'wait': True,
+        'vals': (N, K),
+      }
+    elif GEMM_VARIATION == "swizzled_smem_input" and (M%64)==0 and (N%128)==0 and (K%64)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.float and DTYPE_ACC == dtypes.float:
+      print("Using CUDA, swizzled SMEM inputs")
+      prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp32_fp32.swizzled_smem_input.cu')).read()))
+      args = (c, a, b)
+      kwargs = {
+        'global_size': [M//64, N//128, 1],
+        'local_size': [128, 1, 1], # 4 warpgroups == (T_M:=2) * (T_N:=2)
+        'wait': True,
+        'vals': (N, K),
+      }
+    elif GEMM_VARIATION == "flat_smem_input" and (M%64)==0 and (N%128)==0 and (K%64)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.float and DTYPE_ACC == dtypes.float:
+      print("Using CUDA, flat SMEM inputs")
+      prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp32_fp32.flat_smem_input.cu')).read()))
+      args = (c, a, b)
+      kwargs = {
+        'global_size': [M//64, N//128, 1],
+        'local_size': [128, 1, 1], # 4 warpgroups == (T_M:=2) * (T_N:=2)
+        'wait': True,
+        'vals': (N, K),
+      }
+    elif GEMM_VARIATION == "hcopt" and M == N == K == 4096 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.half and DTYPE_ACC == dtypes.float:
+      print("Using CUDA and generated hcopt")
+      # [Opt(op=OptOps.TC, axis=0, amt=0), Opt(op=OptOps.UPCAST, axis=0, amt=4), Opt(op=OptOps.UPCAST, axis=1, amt=4), Opt(op=OptOps.LOCAL, axis=1, amt=4)]
+      prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp32_fp16.hcopt.cu')).read()))
+      args = (c, a, b)
+      kwargs = {
+        'global_size': [32, 64, 1],
+        'local_size': [16, 2, 4], # 16,2 are warp, 4 workgroups upcasted to axis=1
+        'wait': True,
+      }
+    elif GEMM_VARIATION == "2_stage" and (M%64)== 0 and (N%128)==0 and (K%64)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.half and DTYPE_ACC == dtypes.half:
+      print("Using CUDA and un-optimized 2-stage, swizzled SMEM inputs and direct acc to output kernel")
+      prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp16_fp16.2_stage.cu')).read()))
+      args = (c, a, b)
+      kwargs = {
+        'global_size': [M//64, N//128, 1],
+        'local_size': [128, 1, 1], # 4 warpgroups == (T_M:=2) * (T_N:=2)
+        'wait': True,
+        'vals': (N, K),
+      }
+    elif GEMM_VARIATION == "3_stage" and (M%256)== 0 and (N%128)==0 and (K%32)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.half and DTYPE_ACC == dtypes.half:
+      print("Using CUDA and 3-stage (interleave global copies and ldmatrix)")
+      prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp16_fp16.3_stage.cu')).read()), 73728)
+      args = (c, a, b)
+      kwargs = {
+        'global_size': [M//256, N//128, 1],
+        'local_size': [32, 4, 2], # 8 warpgroups, WG_M=4 and WG_N=2
+        'wait': True,
+        'vals': (N, K),
+      }
+    elif GEMM_VARIATION == "3_stage_swizzled" and (M%256)== 0 and (N%128)==0 and (K%32)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.half and DTYPE_ACC == dtypes.half:
+      print("Using CUDA and 3-stage (interleave global copies and ldmatrix) and swizzled SMEM inputs")
+      prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp16_fp16.3_stage_swizzled.cu')).read()), 73728)
+      args = (c, a, b)
+      kwargs = {
+        'global_size': [M//256, N//128, 1],
+        'local_size': [32, 4, 2], # 8 warpgroups, WG_M=4 and WG_N=2
+        'wait': True,
+        'vals': (N, K),
+      }
+    elif GEMM_VARIATION == "max" and (M%256)== 0 and (N%128)==0 and (K%32)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.half and DTYPE_ACC == dtypes.half:
+      print("Using CUDA and 3-stage (interleave global copies and ldmatrix), swizzled SMEM inputs and epilogue")
+      prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp16_fp16.max.cu')).read()), 73728)
+      args = (c, a, b)
+      kwargs = {
+        'global_size': [M//256, N//128, 1],
+        'local_size': [32, 4, 2], # 8 warpgroups, WG_M=4 and WG_N=2
+        'wait': True,
+        'vals': (N, K),
+      }
+    elif GEMM_VARIATION == "no_xor" and (M%256)== 0 and (N%128)==0 and (K%32)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.half and DTYPE_ACC == dtypes.half:
+      print("Using CUDA and 3-stage (interleave global copies and ldmatrix), swizzled SMEM inputs and epilogue")
+      prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp16_fp16.no_xor.cu')).read()), 73728)
+      args = (c, a, b)
+      kwargs = {
+        'global_size': [M//256, N//128, 1],
+        'local_size': [32, 4, 2], # 8 warpgroups, WG_M=4 and WG_N=2
+        'wait': True,
+        'vals': (N, K),
+      }
+    else:
+      raise RuntimeError(f"invalid gemm variation: {GEMM_VARIATION=} {M=} {N=} {K=} {DTYPE_IN=} {DTYPE_OUT=} {DTYPE_ACC=}")
+
+    tms = []
+    na, nb, nc = randoms()
+    cudaalloc._copyin(a, memoryview(bytearray(na)))
+    cudaalloc._copyin(b, memoryview(bytearray(nb)))
+    for i in range(CNT):
+      tms.append(prog(*args, **kwargs))
+    cudaalloc._copyout(flat_mv(nc.data), c)
+    comp = na.astype(np.float32) @ nb.astype(np.float32)
+    result = nc.reshape(M, N).astype(np.float32)
+
+    print(f"{N*N:10d} {min(tms)*1e6:9.2f} us, would be {FLOPS*1e-9/min(tms):9.2f} GFLOPS matmul, {BW*1e-9/min(tms):.2f} GB/s")
+    try:
+      np.testing.assert_allclose(result, comp, atol=ATOL, rtol=RTOL)
+    except AssertionError as e:
+      if getenv("DEBUG_VALUES") > 0:
+        indices = np.where(~np.isclose(result, comp, rtol=RTOL, atol=ATOL))
+        non_matching_elements_result = result[indices]
+        non_matching_elements_comp = comp[indices]
+        print("valid       :", np.where(np.isclose(result, comp, rtol=RTOL, atol=ATOL)))
+        print("invalid     :", indices)
+        print("result      :", non_matching_elements_result)
+        print("ground truth:", non_matching_elements_comp)
+        print("result sum  :", np.sum(result))
+        print("ground sum  :", np.sum(comp))
+      raise e
+
+    if getenv("DEBUG_VALUES") > 0:
+      print(comp)
+      print("ground sum  :", np.sum(comp))
+      print(result)
+      print("result sum  :", np.sum(result))
+
+  elif getenv("AMD") == 1:
+    # note: https://hipfft.readthedocs.io/en/rocm-6.1.2/how-to/fine-tuning-llms/optimizing-triton-kernel.html
+
+    # also this is different than the rocblas/tensile approach to GEMM
+    # see: https://github.com/ROCm/Tensile/blob/develop/Tensile/KernelWriterAssembly.py
+    raise RuntimeError("invalid max_matmul device")
+
+  else:
+    raise RuntimeError("invalid max_matmul device")
+

tinygrad_repo/extra/gemm/metal_conv.py

@@ -0,0 +1,49 @@
+import os
+#os.environ["METAL"] = "1"
+import numpy as np
+
+BS = 64
+CIN = 256
+COUT = 256
+HW = 32
+K = 3
+PADDING = 0
+# TODO: this is doing some trick, since with CIN=256 COUT=256 it's over 10.4 TFLOPS.
+# are winograd convs less flops? it appears so if they are batched
+# https://www.cse.ust.hk/~weiwa/papers/yan-ppopp20.pdf
+FLOPS = BS*K*K*CIN*HW*HW*COUT*2
+
+nb = np.random.default_rng().standard_normal(size=(BS,CIN,HW,HW), dtype=np.float32)
+nc = np.random.default_rng().standard_normal(size=(COUT,CIN,K,K), dtype=np.float32)
+
+try:
+  import time, torch, torch.mps
+  b = torch.from_numpy(nb).to('mps')
+  c = torch.from_numpy(nc).to('mps')
+
+  def torch_prog(b, c):
+    st = time.perf_counter()
+    a = torch.nn.functional.conv2d(b, c, padding=PADDING)
+    torch.mps.synchronize()
+    return time.perf_counter() - st
+  tm = min([torch_prog(b, c) for _ in range(20)])
+  print(f"{tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS conv in torch")
+except RuntimeError:
+  print("no torch metal conv")
+
+from tinygrad.tensor import Tensor
+from tinygrad.engine.jit import TinyJit
+from tinygrad import Device
+b = Tensor(nb)
+c = Tensor(nc)
+# TODO: slowness without the JIT I suspect comes from a lack of a caching allocator
+@TinyJit
+def tiny_jit(b, c):
+  return b.conv2d(c, padding=PADDING).realize()
+def tiny_prog(b, c):
+  st = time.perf_counter()
+  a = tiny_jit(b, c)
+  Device[a.device].synchronize()
+  return time.perf_counter() - st
+tm = min([tiny_prog(b, c) for _ in range(5)])
+print(f"{tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS conv in tinygrad")

tinygrad_repo/extra/gemm/metal_matmul.py

@@ -0,0 +1,132 @@
+import os
+os.environ["METAL"] = "1"
+import time
+import numpy as np
+from tinygrad import Device, dtypes
+from tinygrad.helpers import getenv, flat_mv
+from tinygrad.runtime.ops_metal import MetalAllocator, MetalDevice, MetalProgram, MetalCompiler
+
+N = getenv("N", 2048)
+LID = 2
+
+device = MetalDevice("METAL")
+metalalloc = MetalAllocator(device)
+
+a = metalalloc.alloc(N*N*4)
+b = metalalloc.alloc(N*N*4)
+c = metalalloc.alloc(N*N*4)
+
+na = np.zeros((N,N),dtype=np.float32)
+nb = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32) #.astype(np.int32).astype(np.float32)N
+nc = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32) #.astype(np.int32).astype(np.float32)
+
+metalalloc._copyin(b,nb.tobytes())
+metalalloc._copyin(c,nc.tobytes())
+
+FLOPS = N*N*N*2
+BW = N*N*3*4
+
+prog = MetalProgram(device, "test", MetalCompiler().compile(f"""
+#include <metal_stdlib>
+#include <metal_simdgroup_matrix>  // Available from Metal version 2.3 released with OS X 11.0+
+using namespace metal;
+kernel void test(device float *a, device const float *data1, device const float *data2, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {{
+  a += gid.x * 32 * {N} + (gid.y * {LID} + lid.y) * 32;
+  data1 += gid.x * 32 * {N};
+  data2 += (gid.y * {LID} + lid.y) * 32;
+
+  simdgroup_float8x8 acc[4][4];
+  for (uint i = 0; i < 4; i++) {{
+    for (uint j = 0; j < 4; j++) {{
+      acc[i][j] = simdgroup_float8x8(0);
+    }}
+  }}
+
+  simdgroup_float8x8 A[4];
+  simdgroup_float8x8 B[4];
+  for (uint k = 0; k < {N}; k+=8) {{
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    simdgroup_load(A[0], data1+k+{0*N}, {N}, ulong2(0, 0));
+    simdgroup_load(A[1], data1+k+{8*N}, {N}, ulong2(0, 0));
+    simdgroup_load(A[2], data1+k+{16*N}, {N}, ulong2(0, 0));
+    simdgroup_load(A[3], data1+k+{24*N}, {N}, ulong2(0, 0));
+    simdgroup_load(B[0], data2+0+k*{N}, {N}, ulong2(0, 0));
+    simdgroup_load(B[1], data2+8+k*{N}, {N}, ulong2(0, 0));
+    simdgroup_load(B[2], data2+16+k*{N}, {N}, ulong2(0, 0));
+    simdgroup_load(B[3], data2+24+k*{N}, {N}, ulong2(0, 0));
+
+    simdgroup_multiply_accumulate(acc[0][0], A[0], B[0], acc[0][0]);
+    simdgroup_multiply_accumulate(acc[0][1], A[1], B[0], acc[0][1]);
+    simdgroup_multiply_accumulate(acc[0][2], A[2], B[0], acc[0][2]);
+    simdgroup_multiply_accumulate(acc[0][3], A[3], B[0], acc[0][3]);
+    simdgroup_multiply_accumulate(acc[1][0], A[0], B[1], acc[1][0]);
+    simdgroup_multiply_accumulate(acc[1][1], A[1], B[1], acc[1][1]);
+    simdgroup_multiply_accumulate(acc[1][2], A[2], B[1], acc[1][2]);
+    simdgroup_multiply_accumulate(acc[1][3], A[3], B[1], acc[1][3]);
+    simdgroup_multiply_accumulate(acc[2][0], A[0], B[2], acc[2][0]);
+    simdgroup_multiply_accumulate(acc[2][1], A[1], B[2], acc[2][1]);
+    simdgroup_multiply_accumulate(acc[2][2], A[2], B[2], acc[2][2]);
+    simdgroup_multiply_accumulate(acc[2][3], A[3], B[2], acc[2][3]);
+    simdgroup_multiply_accumulate(acc[3][0], A[0], B[3], acc[3][0]);
+    simdgroup_multiply_accumulate(acc[3][1], A[1], B[3], acc[3][1]);
+    simdgroup_multiply_accumulate(acc[3][2], A[2], B[3], acc[3][2]);
+    simdgroup_multiply_accumulate(acc[3][3], A[3], B[3], acc[3][3]);
+  }}
+  simdgroup_store(acc[0][0], a+{0+0*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[1][0], a+{8+0*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[2][0], a+{16+0*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[3][0], a+{24+0*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[0][1], a+{0+8*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[1][1], a+{8+8*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[2][1], a+{16+8*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[3][1], a+{24+8*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[0][2], a+{0+16*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[1][2], a+{8+16*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[2][2], a+{16+16*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[3][2], a+{24+16*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[0][3], a+{0+24*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[1][3], a+{8+24*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[2][3], a+{16+24*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[3][3], a+{24+24*N}, {N}, ulong2(0, 0));
+}}"""))
+def timeit(fxn):
+  st = time.perf_counter()
+  et = fxn()
+  # NOTE: et doesn't contain the launch overhead
+  return time.perf_counter() - st
+tm = min([timeit(lambda: prog(a, b, c, global_size=[N//(8*4), N//(8*4*LID), 1], local_size=[32, LID, 1], wait=True)) for _ in range(20)])
+comp = nb@nc
+metalalloc._copyout(flat_mv(na.data), a)
+if N <= 32:
+  print(na)
+  print(comp)
+print(f"{N*N:10d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS matmul, {BW*1e-9/tm:.2f} GB/s")
+np.testing.assert_allclose(na, comp, atol=1e-3)
+
+import torch, torch.mps
+b = torch.from_numpy(nb).to('mps')
+c = torch.from_numpy(nc).to('mps')
+
+def torch_prog(b, c):
+  st = time.perf_counter()
+  a = b@c
+  torch.mps.synchronize()
+  return time.perf_counter() - st
+tm = min([torch_prog(b, c) for _ in range(20)])
+print(f"{N*N:10d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS matmul in torch")
+
+from tinygrad.tensor import Tensor
+from tinygrad.engine.jit import TinyJit
+b = Tensor(nb)
+c = Tensor(nc)
+# TODO: slowness without the JIT I suspect comes from a lack of a caching allocator
+@TinyJit
+def tiny_jit(b, c):
+  return (b@c).realize()
+def tiny_prog(b, c):
+  st = time.perf_counter()
+  a = tiny_jit(b, c)
+  Device["METAL"].synchronize()
+  return time.perf_counter() - st
+tm = min([tiny_prog(b, c) for _ in range(20)])
+print(f"{N*N:10d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS matmul in tinygrad")

tinygrad_repo/extra/gemm/metal_matvec.py

@@ -0,0 +1,113 @@
+import numpy as np
+import time, torch, torch.mps
+
+from tinygrad import Tensor, TinyJit, Device
+from tinygrad.helpers import flat_mv
+from tinygrad.runtime.ops_metal import MetalAllocator, MetalDevice, MetalProgram, MetalCompiler
+
+N = 16384
+M = 4096
+FLOPS = N*M*2
+
+nb = np.random.default_rng().standard_normal(size=(N), dtype=np.float32) #.astype(np.int32).astype(np.float32)
+nc = np.random.default_rng().standard_normal(size=(N,M), dtype=np.float32) #.astype(np.int32).astype(np.float32)
+
+b = torch.from_numpy(nb).to('mps')
+c = torch.from_numpy(nc).to('mps')
+
+def torch_prog(b, c):
+  st = time.perf_counter()
+  a = b@c
+  torch.mps.synchronize()
+  return time.perf_counter() - st
+tm = min([torch_prog(b, c) for _ in range(200)])
+print(f"{N:d}x{M:d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS matvec in torch")
+torch_a = (b@c).cpu()
+
+device = MetalDevice("METAL")
+metalalloc = MetalAllocator(device)
+WORKSIZE_ROW = 16
+WORKSIZE_COL = 1
+LOCAL_SIZE = [32, WORKSIZE_COL, WORKSIZE_ROW]
+GLOBAL_SIZE = [M//(LOCAL_SIZE[0]*LOCAL_SIZE[1]*4), 1, 1]
+prog = MetalProgram(device, "test", MetalCompiler().compile(f"""
+#include <metal_stdlib>
+using namespace metal;
+kernel void test(device float* data0, const device float* data1, const device float* data2, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {{
+  int gidx0 = gid.x;  /* {GLOBAL_SIZE[0]} */
+  int lidx1 = lid.x;  /* {LOCAL_SIZE[0]} */
+  int lidx2 = lid.y;  /* {LOCAL_SIZE[1]} */
+  int lidx3 = lid.z;  /* {LOCAL_SIZE[2]} */
+
+  // 4 rows per thread
+  threadgroup float4 acc0[{LOCAL_SIZE[0]*LOCAL_SIZE[1]*LOCAL_SIZE[2]}];
+  int acc0_index = ((lidx1*{LOCAL_SIZE[1]})+lidx2)+({LOCAL_SIZE[0]*LOCAL_SIZE[1]}*lidx3);
+  acc0[acc0_index] = float4(0.0f,0.0f,0.0f,0.0f);
+
+  threadgroup float4 val1[{LOCAL_SIZE[0]*LOCAL_SIZE[1]*LOCAL_SIZE[2]}];
+
+  // iterate over the columns
+  for (int ridx2 = 0; ridx2 < {N//(4*LOCAL_SIZE[0]*LOCAL_SIZE[1]*(LOCAL_SIZE[2]))}; ++ridx2) {{
+    // load 4*threadgroup_size columns into shared memory
+    int col_1 = (((lidx3*{N//(4*LOCAL_SIZE[0]*LOCAL_SIZE[1]*(LOCAL_SIZE[2]))})+ridx2)*{LOCAL_SIZE[0]*LOCAL_SIZE[1]})+(lidx1*{LOCAL_SIZE[1]})+lidx2;
+    val1[(lidx3*{LOCAL_SIZE[1]*LOCAL_SIZE[0]})+((lidx1*{LOCAL_SIZE[1]})+lidx2)] = *((device float4*)(data1+(col_1*4)));
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    for (int ridx3 = 0; ridx3 < {LOCAL_SIZE[0]*LOCAL_SIZE[1]}; ++ridx3) {{
+      int col = ((((lidx3*{N//(4*LOCAL_SIZE[0]*LOCAL_SIZE[1]*(LOCAL_SIZE[2]))})+ridx2)*{LOCAL_SIZE[0]*LOCAL_SIZE[1]})+ridx3);
+      float4 val1_0 = val1[(lidx3*{LOCAL_SIZE[1]*LOCAL_SIZE[0]})+ridx3];
+      float4 val2_0 = (float4)(*((device float4*)(data2+(gidx0*{M//GLOBAL_SIZE[0]})+(((lidx1*{LOCAL_SIZE[1]})+lidx2)*4)+(col*{M*4})+{M*0})));
+      float4 val2_1 = (float4)(*((device float4*)(data2+(gidx0*{M//GLOBAL_SIZE[0]})+(((lidx1*{LOCAL_SIZE[1]})+lidx2)*4)+(col*{M*4})+{M*1})));
+      float4 val2_2 = (float4)(*((device float4*)(data2+(gidx0*{M//GLOBAL_SIZE[0]})+(((lidx1*{LOCAL_SIZE[1]})+lidx2)*4)+(col*{M*4})+{M*2})));
+      float4 val2_3 = (float4)(*((device float4*)(data2+(gidx0*{M//GLOBAL_SIZE[0]})+(((lidx1*{LOCAL_SIZE[1]})+lidx2)*4)+(col*{M*4})+{M*3})));
+      acc0[acc0_index] = ((val1_0.x*val2_0)+acc0[acc0_index]);
+      acc0[acc0_index] = ((val1_0.y*val2_1)+acc0[acc0_index]);
+      acc0[acc0_index] = ((val1_0.z*val2_2)+acc0[acc0_index]);
+      acc0[acc0_index] = ((val1_0.w*val2_3)+acc0[acc0_index]);
+    }}
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+  }} /* reduce */
+
+  if (lidx3 == 0) {{
+    float4 out = float4(0.0f,0.0f,0.0f,0.0f);
+    for (int n = 0; n < {LOCAL_SIZE[2]}; n++) {{
+      out += acc0[((lidx1*{LOCAL_SIZE[1]})+lidx2)+({LOCAL_SIZE[0]*LOCAL_SIZE[1]}*n)];
+    }}
+    *( (device float4 *) (data0 + (gidx0*{M//GLOBAL_SIZE[0]}) + ( ( (lidx1*{LOCAL_SIZE[1]})+lidx2 ) * 4 ) ) ) = out;
+  }}
+}}
+"""))
+a = metalalloc.alloc(M*4)
+b = metalalloc.alloc(N*4)
+c = metalalloc.alloc(N*M*4)
+metalalloc._copyin(b,nb.tobytes())
+metalalloc._copyin(c,nc.tobytes())
+def metalrun():
+  prog(a, b, c, global_size=GLOBAL_SIZE, local_size=LOCAL_SIZE, wait=True)
+  return a
+def timeit(fxn):
+  st = time.perf_counter()
+  et = fxn()
+  # NOTE: et doesn't contain the launch overhead
+  return time.perf_counter() - st
+tm = min([timeit(metalrun) for _ in range(200)])
+print(f"{N:d}x{M:d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS matvec in metal")
+metal_a = np.zeros(M, dtype=np.float32)
+metalalloc._copyout(flat_mv(metal_a.data), a)
+np.testing.assert_allclose(metal_a, torch_a, atol=5e-3)
+
+b = Tensor(nb)
+c = Tensor(nc)
+# TODO: slowness without the JIT I suspect comes from a lack of a caching allocator
+@TinyJit
+def tiny_jit(b, c):
+  return (b@c).realize()
+def tiny_prog(b, c):
+  st = time.perf_counter()
+  a = tiny_jit(b, c)
+  Device["METAL"].synchronize()
+  return time.perf_counter() - st
+tm = min([tiny_prog(b, c) for _ in range(200)])
+print(f"{N:d}x{M:d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS matvec in tinygrad")
+tiny_a = tiny_jit(b, c).numpy()
+np.testing.assert_allclose(tiny_a, torch_a, atol=5e-3)

tinygrad_repo/extra/gemm/metal_uop_matmul.py

@@ -0,0 +1,42 @@
+from tinygrad import UOp, dtypes
+from tinygrad.uop.ops import AxisType, Ops, KernelInfo, AddrSpace
+from extra.gemm.amd_uop_matmul import test_matmul
+
+N = 2048
+
+# metal has an 8x8 tensor core. this is the indexing
+def mat_idx(buf, g0, g1, warp, u):
+  l = [(warp//2**i)%2 for i in range(5)]
+  return buf[g0, l[4]*4 + l[2]*2 + l[1], g1, l[3]*4 + l[0]*2 + u]
+
+def hand_spec_tc_cores():
+  gx = UOp.special(N // 8, "gidx0")
+  gy = UOp.special(N // 8, "gidx1")
+  warp = UOp.special(32, "lidx0")
+
+  c = UOp.placeholder((N, N), dtypes.float, slot=0).reshape((N//8, 8, N//8, 8))
+  a = UOp.placeholder((N, N), dtypes.float, slot=1).reshape((N//8, 8, N//8, 8))
+  b = UOp.placeholder((N, N), dtypes.float, slot=2).reshape((N//8, 8, N//8, 8))
+
+  gk = UOp.range(N // 8, 0, AxisType.REDUCE)
+
+  a_tc = UOp.vectorize(*[mat_idx(a, gx, gk, warp, i) for i in range(2)])
+  b_tc = UOp.vectorize(*[mat_idx(b, gk, gy, warp, i) for i in range(2)])
+
+  acc = UOp.placeholder((2,), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  acc = acc[0].set(0.0)
+  acc = acc[1].set(0.0)
+
+  # TODO: make this simple
+  wmma_arg = ('WMMA_8_8_8_float_float', (8, 8, 8), dtypes.float, dtypes.float, 'METAL', 32, (((3, 2),), ((3, 2),), ((3, 2),)), ())
+
+  acc_load = UOp.vectorize(acc.after(gk)[0], acc.after(gk)[1])
+  out = UOp(Ops.WMMA, dtypes.float.vec(2), (a_tc, b_tc, acc_load), arg=wmma_arg)
+
+  end_loop = UOp.group(*[acc[i].store(out.gep(i)) for i in range(2)]).end(gk)
+
+  sink = UOp.group(*[mat_idx(c.after(end_loop), gx, gy, warp, i).store(acc[i]) for i in range(2)])
+  return sink.sink(arg=KernelInfo(name="custom_metal_matmul", opts_to_apply=())).simplify()
+
+if __name__ == "__main__":
+  test_matmul(hand_spec_tc_cores(), N=N)

tinygrad_repo/extra/gemm/mi350x_uop_matmul.py

@@ -0,0 +1,229 @@
+import os
+import numpy as np
+np.set_printoptions(linewidth=1000000)
+os.environ["AMD_LLVM"] = "0"
+
+from tinygrad import Tensor, Context, dtypes, UOp, GlobalCounters
+from tinygrad.helpers import DEBUG, getenv
+from tinygrad.dtype import AddrSpace
+from tinygrad.uop.ops import AxisType, KernelInfo, Ops
+
+WARP_SIZE = 64
+
+# Reg tile sizes (tensor cores)
+TC_M = 16
+TC_N = 16
+TC_K = 32
+
+# 1024 matrix cores
+# 16 cycle mfma
+# 2.2 GHz
+# 16x16x32x2 FLOPS/mma = 16384
+# 2.2*1e9*16384*1024/16*1e-12 TFLOPS = 2306 TFLOPS
+
+#N,M,K = 256,256,64
+N,M,K = 4096,4096,4096
+
+# Threadblock tile sizes (block-level tile of C that a block computes)
+#BLOCK_M = 128   # rows of C (M-dim) per block
+#BLOCK_N = 128   # columns of C (N-dim) per block
+#BLOCK_K = 128   # K-slice per block iteration
+
+BLOCK_M = 64
+BLOCK_N = 64
+BLOCK_K = 128
+
+WARPGROUP_SIZE = 1
+BLOCK_M = BLOCK_M * WARPGROUP_SIZE
+
+# TODO: improve the syntax of this. better syntax, faster iteration
+#  -- DONE: add working slice a[gx, :, i] -> shape of the : (aka (16,16,32) becomes (16,))
+#  -- DONE(ish): add argfix to movement (traits shared with Tensor)
+#  -- fix WMMA to not require all the junk
+#  -- improve syntax for vectorized loads/stores (both with DEVECTORIZE and without)
+#  -- DONE: be able to use CONTRACT on a range
+#  -- fix upcasted RANGE on an already vectorized buffer
+#  -- improve "all ranges not ended error" / fix the bug with after on ended ranges (if you are after end of range, range is closed)
+
+CUS_PER_GPU = 256
+assert ((M//BLOCK_M) * (N//BLOCK_N)) >= CUS_PER_GPU, "not enough globals"
+
+def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
+  # A = (M x K)
+  # B = (K x N)
+  # C = (M x N)
+
+  # check it's proper matmul
+  assert C.shape[0] == A.shape[0]
+  assert C.shape[1] == B.shape[1]
+  assert A.shape[1] == B.shape[0]
+
+  gx, gy = UOp.special(M//BLOCK_M, "gidx0"), UOp.special(N//BLOCK_N, "gidx1")
+  warp = UOp.special(WARP_SIZE, "lidx0")
+  warpgroup = UOp.special(WARPGROUP_SIZE, "lidx1")
+
+  # generic copy logic (not good)
+  def generic_copy(glbl, gargs, lcl, rng):
+    # Fully coalesced 128-bit loads/stores.
+    INNER_SIZE = 8
+    cp_i = UOp.range(lcl.size//(WARPGROUP_SIZE*WARP_SIZE*INNER_SIZE), rng)
+    cp_inner = UOp.range(INNER_SIZE, rng+1, AxisType.UPCAST)
+    idx_i = cp_i*WARPGROUP_SIZE*WARP_SIZE*INNER_SIZE + warpgroup*WARP_SIZE*INNER_SIZE + warp*INNER_SIZE + cp_inner
+    return lcl[idx_i].store(glbl[*gargs, idx_i]).end(cp_i, cp_inner)
+
+  # split out the globals into blocks
+  C = C.reshape((M//BLOCK_M, BLOCK_M, N//BLOCK_N, BLOCK_N))
+  A = A.reshape((M//BLOCK_M, BLOCK_M, K//BLOCK_K, BLOCK_K))
+  B = B.reshape((K//BLOCK_K, BLOCK_K, N//BLOCK_N, BLOCK_N))
+
+  # this is the big accumulator
+  acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float.vec(4), 0, AddrSpace.REG)
+  assert acc.size*WARP_SIZE*WARPGROUP_SIZE*4 == BLOCK_M*BLOCK_N
+  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const(dtypes.float.vec(4), 0.0), end=init_l)
+
+  # create locals (note A is permuted, and the stride is changed to avoid bank conflicts)
+  def make_locals(slot) -> tuple[UOp, UOp]:
+    BM_As_stride = (BLOCK_M + 1)
+    BN_Bs_stride = (BLOCK_N + 0)
+    INNER_SLICE = 8
+    As = UOp.placeholder((BLOCK_K//INNER_SLICE, BM_As_stride, INNER_SLICE), dtypes.half, slot=slot, addrspace=AddrSpace.LOCAL)
+    INNER_SLICE = 1
+    Bs = UOp.placeholder((BLOCK_K//INNER_SLICE, BN_Bs_stride, INNER_SLICE), dtypes.half, slot=slot+1, addrspace=AddrSpace.LOCAL)
+    As = As.permute((0,2,1)).reshape((BLOCK_K, BM_As_stride)).shrink_to((BLOCK_K, BLOCK_M))
+    Bs = Bs.permute((0,2,1)).reshape((BLOCK_K, BN_Bs_stride)).shrink_to((BLOCK_K, BLOCK_N))
+    return As, Bs
+
+  # load from globals into locals (TODO: use the warpgroup)
+
+  def load_to_locals(l_K_outer_loop:UOp, Asl:UOp, Bsl:UOp, rng:int, barrier=True) -> tuple[UOp, UOp]:
+    if getenv("FAKE"):
+      return Asl[0].set(0), Bsl[0].set(0)
+    else:
+      pA = A.permute((0,2,1,3)).reshape((M//BLOCK_M, K//BLOCK_K, BLOCK_M*BLOCK_K))
+      pas = Asl.permute((1,0)).reshape((BLOCK_M*BLOCK_K,))
+      As_store = generic_copy(pA, (gx, l_K_outer_loop), pas, rng)
+
+      pB = B.permute((0,2,1,3)).reshape((K//BLOCK_K, N//BLOCK_N, BLOCK_K*BLOCK_N))
+      pbs = Bsl.reshape((BLOCK_K*BLOCK_N,))
+      Bs_store = generic_copy(pB, (l_K_outer_loop, gy), pbs, rng+2)
+
+      barrier = UOp.barrier(As_store, Bs_store) if barrier else UOp.group(As_store, Bs_store)
+      return Asl.after(barrier), Bsl.after(barrier)
+
+  def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...]=()) -> UOp:
+    K_inner_loop = UOp.range(BLOCK_K//TC_K, rng, AxisType.REDUCE)
+
+    # load from locals into registers
+    Ar = UOp.placeholder((BLOCK_M//TC_M//WARPGROUP_SIZE,), dtypes.half.vec(8), slot=1, addrspace=AddrSpace.REG)
+    Br = UOp.placeholder((BLOCK_N//TC_N,), dtypes.half.vec(8), slot=2, addrspace=AddrSpace.REG)
+
+    M_load_loop = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, rng+10)
+    Asl = Asl.reshape((BLOCK_K//TC_K, TC_K, BLOCK_M//TC_M//WARPGROUP_SIZE, WARPGROUP_SIZE, TC_M))
+    load_rng = UOp.range(8, rng+11, axis_type=AxisType.UPCAST)
+    A_in = Asl[K_inner_loop, (warp//16)*8+load_rng, M_load_loop, warpgroup, warp%16].contract(load_rng)
+    Ar = Ar[M_load_loop].set(A_in, end=M_load_loop)
+
+    N_load_loop = UOp.range(BLOCK_N//TC_N, rng+20)
+    Bsl = Bsl.reshape((BLOCK_K//TC_K, TC_K, BLOCK_N//TC_N, TC_N))
+    load_rng = UOp.range(8, rng+21, axis_type=AxisType.UPCAST)
+    B_in = Bsl[K_inner_loop, (warp//16)*8+load_rng, N_load_loop, warp%16].contract(load_rng)
+    Br = Br[N_load_loop].set(B_in, end=N_load_loop)
+
+    M_inner_loop = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, rng+30)
+    N_inner_loop = UOp.range(BLOCK_N//TC_N, rng+31)
+
+    # load values
+    acc_after = acc.after(*afters, M_inner_loop, N_inner_loop, K_inner_loop)
+    acc_load = acc_after[N_inner_loop, M_inner_loop]
+
+    # do WMMA
+    wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
+    out = UOp(Ops.WMMA, dtypes.float.vec(4), (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)
+
+    # store back the acc
+    acc_store = acc[N_inner_loop, M_inner_loop].store(out)
+    return acc_store.end(M_inner_loop, N_inner_loop, K_inner_loop)
+
+  # **** START INNER LOOP *****
+  # inner loop -- locals -> regs
+
+  # no pipeline
+  if not getenv("PIPELINE"):
+    As, Bs = make_locals(slot=0)
+
+    K_outer_loop = UOp.range(K//BLOCK_K, 0, AxisType.REDUCE)
+    As, Bs = load_to_locals(K_outer_loop, As, Bs, 1000, barrier=True)
+    acc_store = compute_on_locals(acc, As, Bs, 1500, afters=(K_outer_loop,))
+    acc = acc.after(acc_store.barrier().end(K_outer_loop))
+  else:
+    # this doesn't work
+    As0, Bs0 = make_locals(slot=0)
+    As1, Bs1 = make_locals(slot=2)
+    As0, Bs0 = load_to_locals(0, As0, Bs0, 1000)
+
+    K_outer_loop = UOp.range((K//BLOCK_K-2)//2, 0, AxisType.REDUCE)
+    As1, Bs1 = load_to_locals(K_outer_loop+1, As1, Bs1, 2000, barrier=False)
+    acc_store = compute_on_locals(acc, As0, Bs0, 1500, afters=(K_outer_loop,))
+    As0, Bs0 = load_to_locals(K_outer_loop+2, As0, Bs0, 3000, barrier=False)
+    acc_store = compute_on_locals(acc, As1, Bs1, 2500, afters=(acc_store, As0, Bs0))
+    acc = acc.after(acc_store.barrier().end(K_outer_loop))
+
+    #acc_store = compute_on_locals(acc, As0, Bs0, 3500, afters=(acc_store.barrier().end(K_outer_loop)))
+    """
+    As1, Bs1 = load_to_locals(K//BLOCK_K-1, As1, Bs1, 4000)
+    acc_store = compute_on_locals(acc, As1, Bs1, 4500, afters=(acc_store))
+    """
+    #acc = acc.after(acc_store)
+
+  # **** END LOOPS *****
+
+  # store the acc into gmem
+  cp_i, cp_j = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, 10004), UOp.range(BLOCK_N//TC_N, 10005)
+  c_load = lambda i: C[gx, cp_i*TC_M*WARPGROUP_SIZE + warpgroup*TC_M + (warp//16)*4+i, gy, cp_j*TC_N + warp%16]
+  store = UOp.group(*[c_load(i).store(acc[cp_j, cp_i].gep(i)) for i in range(4)])
+  store = store.end(cp_i, cp_j)
+
+  return store.sink(arg=KernelInfo(name="custom_gemm", opts_to_apply=())).simplify()
+
+# simplest WMMA
+"""
+# init the acc
+acc = UOp.placeholder((4,), dtypes.float, 0, AddrSpace.REG)
+acc = acc[init_l:=UOp.range(4, 1)].set(0.0, end=init_l)
+
+# do the wmma
+acc_load = UOp.vectorize(*[acc.after(K_loop)[i] for i in range(4)])
+wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
+out = UOp(Ops.WMMA, dtypes.float.vec(4), (A_in, B_in, acc_load), arg=wmma_arg)
+
+# store back the acc
+acc = acc.after(UOp.group(*[acc[i].store(out.gep(i)) for i in range(4)]).end(K_loop))
+
+# store the acc into gmem
+store = UOp.group(*[C[gx, (warp//16)*4+i, gy, warp%16].store(acc[i]) for i in range(4)])
+"""
+
+if __name__ == "__main__":
+  a = Tensor.randn(M, K, dtype=dtypes.half)
+  b = Tensor.randn(K, N, dtype=dtypes.half)
+
+  #a = Tensor.zeros(M, K, dtype=dtypes.half).contiguous()
+  #a[0,16] = 1
+  #b = Tensor.ones(K, N, dtype=dtypes.half).contiguous()
+
+  c = Tensor.empty(M, N, dtype=dtypes.float)
+  with Context(DEBUG=0): Tensor.realize(a,b)
+
+  ref = a.dot(b, dtype=dtypes.float)
+  ref.realize()
+
+  GlobalCounters.reset()
+  with Context(DEBUG=max(2, DEBUG.value)):
+    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
+    tst.realize()
+  print(f"{(N*M*K*2 / GlobalCounters.time_sum_s)*1e-12:.2f} REAL TFLOPS")
+
+  with Context(DEBUG=0):
+    #print(ref.numpy())
+    #print(tst.numpy())
+    assert Tensor.isclose(ref, tst, atol=1e-2).all().item(), "matrix not close"

tinygrad_repo/extra/gemm/mi350x_uop_matmul_2.py

@@ -0,0 +1,141 @@
+import os
+import numpy as np
+np.set_printoptions(linewidth=1000000)
+os.environ["AMD_LLVM"] = "0"
+
+from tinygrad import Tensor, Context, dtypes, UOp, GlobalCounters
+from tinygrad.helpers import DEBUG, getenv
+from tinygrad.dtype import AddrSpace
+from tinygrad.uop.ops import sint, AxisType, KernelInfo, Ops
+
+WARP_SIZE = 64
+
+# Reg tile sizes (tensor cores)
+TC_M = 16
+TC_N = 16
+TC_K = 32
+
+N,M,K = 4096,4096,4096
+
+# Threadblock tile sizes (block-level tile of C that a block computes)
+BLOCK_M = 64
+BLOCK_N = 64
+BLOCK_K = 64
+
+WARPGROUP_SIZE = 1
+BLOCK_M = BLOCK_M * WARPGROUP_SIZE
+
+TID_SIZE = WARPGROUP_SIZE*WARP_SIZE
+
+def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=()):
+  assert dest.shape == src.shape
+  rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.LOOP) for i,s in enumerate(src.shape)]
+  copy = dest[*rngs].store(src[*rngs]).end(*rngs)
+  return dest.after(copy) if set else copy
+
+def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...], warpgroup, warp) -> UOp:
+  K_inner_loop = UOp.range(BLOCK_K//TC_K, rng, AxisType.REDUCE)
+
+  # load from locals into registers
+  Ar = UOp.placeholder((BLOCK_M//TC_M//WARPGROUP_SIZE,), dtypes.half.vec(8), slot=1, addrspace=AddrSpace.REG)
+  Br = UOp.placeholder((BLOCK_N//TC_N,), dtypes.half.vec(8), slot=2, addrspace=AddrSpace.REG)
+
+  M_load_loop = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, rng+10)
+  Asl = Asl.reshape(BLOCK_K//TC_K, TC_K, BLOCK_M//TC_M//WARPGROUP_SIZE, WARPGROUP_SIZE, TC_M)
+  load_rng = UOp.range(8, rng+11, axis_type=AxisType.UPCAST)
+  A_in = Asl[K_inner_loop, (warp//16)*8+load_rng, M_load_loop, warpgroup, warp%16].contract(load_rng)
+  Ar = Ar[M_load_loop].set(A_in, end=M_load_loop)
+
+  N_load_loop = UOp.range(BLOCK_N//TC_N, rng+20)
+  Bsl = Bsl.reshape(BLOCK_K//TC_K, TC_K, BLOCK_N//TC_N, TC_N)
+  load_rng = UOp.range(8, rng+21, axis_type=AxisType.UPCAST)
+  B_in = Bsl[K_inner_loop, (warp//16)*8+load_rng, N_load_loop, warp%16].contract(load_rng)
+  Br = Br[N_load_loop].set(B_in, end=N_load_loop)
+
+  M_inner_loop = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, rng+30)
+  N_inner_loop = UOp.range(BLOCK_N//TC_N, rng+31)
+
+  # load values
+  acc_after = acc.after(*afters, M_inner_loop, N_inner_loop, K_inner_loop)
+  acc_load = acc_after[N_inner_loop, M_inner_loop]
+
+  # do WMMA
+  wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
+  out = UOp(Ops.WMMA, dtypes.float.vec(4), (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)
+
+  # store back the acc
+  acc_store = acc[N_inner_loop, M_inner_loop].store(out)
+  return acc_store.end(M_inner_loop, N_inner_loop, K_inner_loop)
+
+def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
+  gx, gy = UOp.special(M//BLOCK_M, "gidx0"), UOp.special(N//BLOCK_N, "gidx1")
+  K_outer_loop = UOp.range(K//BLOCK_K, 0, AxisType.REDUCE)
+
+  # split out the globals into blocks
+  C = C.src[0].cast(dtypes.float.vec(4).ptr(C.ptrdtype.size)).reshape((M//BLOCK_M, BLOCK_M, N//BLOCK_N, BLOCK_N))
+  A = A.reshape((M//BLOCK_M, BLOCK_M, K//BLOCK_K, BLOCK_K))[gx, :, K_outer_loop, :]
+  B = B.reshape((K//BLOCK_K, BLOCK_K, N//BLOCK_N, BLOCK_N))[K_outer_loop, :, gy, :]
+
+  # ---------------------------
+  # GLOBAL -> LOCAL (As, Bs)
+  # ---------------------------
+  tid = UOp.special(TID_SIZE, "lidx0")
+  warpgroup, warp = tid//WARP_SIZE, tid%WARP_SIZE
+
+  A_view = A.reshape(-1, TID_SIZE, 8)
+  B_view = B.reshape(-1, TID_SIZE, 8)
+
+  # A: read BM x BK tiles (permute on store into locals)
+  As = UOp.placeholder((BLOCK_K, BLOCK_M), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL).shrink_to(BLOCK_K, BLOCK_M)
+  As_view = As.reshape(-1, TID_SIZE, 8)
+
+  Bs = UOp.placeholder((BLOCK_K, BLOCK_N+4), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL).shrink_to(BLOCK_K, BLOCK_N)
+  Bs_view = Bs.reshape(-1, TID_SIZE, 8)
+
+  outer_copy = UOp.range(A_view.shape[0], 100, AxisType.UPCAST)
+  inner_copy = UOp.range(A_view.shape[2], 101, AxisType.UPCAST)
+  As_store = As_view[outer_copy, tid, inner_copy].store(A_view[outer_copy, tid, inner_copy])
+  Bs_store = Bs_view[outer_copy, tid, inner_copy].store(B_view[outer_copy, tid, inner_copy])
+
+  if getenv("NOLOAD"):
+    As_store = As[0,0].store(0)
+    Bs_store = Bs[0,0].store(0)
+
+  # TODO: can we automate barrier?
+  barrier = UOp.barrier(UOp.group(As_store, Bs_store).end(outer_copy, inner_copy))
+
+  if getenv("COMPUTE"):
+    As, Bs = As.after(barrier), Bs.after(barrier)
+
+    acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float.vec(4), 0, AddrSpace.REG)
+
+    sink = compute_on_locals(acc, As, Bs, 200, afters=(barrier,), warpgroup=warpgroup, warp=warp)
+    sink = sink.end(K_outer_loop)
+
+    C_view = C[gx, :, gy, :].reshape(BLOCK_M//TC_M//WARPGROUP_SIZE, WARPGROUP_SIZE, TC_M, BLOCK_N//TC_N, TC_N)[:, warpgroup, warp%16, :, (warp//16)*4]
+    sink = copy(C_view, acc.after(sink), rng=300)
+  else:
+    sink = C.after(barrier.end(K_outer_loop))[0,0,0,0].store(As[0,0]+Bs[0,0])
+
+  return sink.sink(arg=KernelInfo(name="custom_gemm", opts_to_apply=())).simplify()
+
+if __name__ == "__main__":
+  a = Tensor.randn(M, K, dtype=dtypes.half)
+  b = Tensor.randn(K, N, dtype=dtypes.half)
+  c = Tensor.empty(M, N, dtype=dtypes.float)
+  with Context(DEBUG=0): Tensor.realize(a,b)
+
+
+  GlobalCounters.reset()
+  with Context(DEBUG=max(2, DEBUG.value)):
+    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
+    tst.realize()
+  print(f"{(N*M*K*2 / GlobalCounters.time_sum_s)*1e-12:.2f} REAL TFLOPS")
+
+
+  with Context(DEBUG=0):
+    ref = a.dot(b, dtype=dtypes.float)
+    ref.realize()
+    #print(ref.numpy())
+    #print(tst.numpy())
+    assert Tensor.isclose(ref, tst, atol=1e-2).all().item(), "matrix not close"

tinygrad_repo/extra/gemm/rdna4_asm_matmul.py

@@ -0,0 +1,248 @@
+# RDNA4 128x128 GEMM using WMMA — optimized DS scheduling
+import numpy as np
+from tinygrad import Tensor, Device, Context, GlobalCounters
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.helpers import getenv, colored
+from tinygrad.dtype import dtypes, AddrSpace
+from tinygrad.engine.realize import Estimates, run_linear
+from tinygrad.renderer.amd.dsl import s, v, VCC_LO, NULL, src, ttmp
+from tinygrad.runtime.autogen.amd.rdna4.ins import *
+
+BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 16
+TILES_M, TILES_N = 4, 4
+THREADS, ELEM = 128, 2
+LDS_A_ROW = BLOCK_K*ELEM   # 32
+LDS_B_ROW = BLOCK_N*ELEM  # 256
+LDS_A_SIZE = BLOCK_M * LDS_A_ROW  # 4096
+LDS_B_SIZE = BLOCK_K * LDS_B_ROW  # 4096
+LDS_SIZE = LDS_A_SIZE + LDS_B_SIZE  # 8192
+LDS_B_OFF = LDS_A_SIZE
+ACC, DA, DB, FA, FB, ET = 60, 188, 196, 204, 44, 10
+
+def build_kernel(N, arch='gfx1200'):
+  assert N % BLOCK_M == 0 and N >= 256
+  NO_ALU, NO_DS, NO_GLOBAL = getenv("NO_ALU", 0), getenv("NO_DS", 0), getenv("NO_GLOBAL", 0)
+  I, L, B = [], {}, []
+  def e(i): I.append(i); return i
+  def label(n): L[n] = sum(i.size() for i in I)
+  def br(i, t): B.append((len(I)-1, t))
+
+  e(s_load_b128(sdata=s[4:7], sbase=s[0:1], ioffset=0, soffset=NULL))
+  e(s_load_b64(sdata=s[8:9], sbase=s[0:1], ioffset=0x10, soffset=NULL))
+  e(s_wait_kmcnt(simm16=0))
+  e(s_mov_b32(s[10], ttmp[9])); e(s_and_b32(s[11], ttmp[7], 0xFFFF))
+  e(s_lshl_b32(s[10], s[10], 7)); e(s_lshl_b32(s[11], s[11], 7))
+  e(s_mov_b32(s[12], N)); e(s_lshl_b32(s[13], s[12], 1))
+  e(s_mul_i32(s[14], s[12], BLOCK_K*ELEM))
+  e(s_add_co_i32(s[17], s[12], -2*BLOCK_K))  # loop bound
+
+  e(v_and_b32_e32(v[1], 31, v[0])); e(v_lshrrev_b32_e32(v[2], 5, v[0]))
+  e(v_and_b32_e32(v[3], 1, v[2])); e(v_lshrrev_b32_e32(v[2], 1, v[2]))
+
+  e(v_lshlrev_b32_e32(v[4], 5, v[0]))
+  # B store: transposed layout for stride-32 reads. addr = LDS_B_OFF + (tid%8)*512 + (tid/8)*32
+  e(v_and_b32_e32(v[48], 7, v[0])); e(v_lshlrev_b32_e32(v[5], 9, v[48]))   # (tid%8)*512
+  e(v_lshrrev_b32_e32(v[48], 3, v[0])); e(v_lshlrev_b32_e32(v[48], 5, v[48]))  # (tid/8)*32
+  e(v_add_nc_u32_e32(v[5], v[5], v[48])); e(v_add_nc_u32_e32(v[5], LDS_B_OFF, v[5]))
+
+  e(v_add_nc_u32_e32(v[48], s[11], v[0]))
+  e(v_mul_lo_u32(v[6], v[48], N*ELEM)); e(v_mov_b32_e32(v[7], 0))
+  e(v_lshrrev_b32_e32(v[48], 3, v[0])); e(v_mul_lo_u32(v[8], v[48], N*ELEM))
+  e(v_and_b32_e32(v[48], 7, v[0])); e(v_lshlrev_b32_e32(v[48], 5, v[48]))
+  e(v_add_nc_u32_e32(v[8], v[8], v[48]))
+  e(s_mul_i32(s[15], s[10], ELEM)); e(v_add_nc_u32_e32(v[8], s[15], v[8]))
+  e(v_mov_b32_e32(v[9], 0))
+
+  # LDS read addrs with padded strides (eliminates bank conflicts)
+  # A: (lane%16)*LDS_A_ROW + (lane/16)*16 + wave_m*64*LDS_A_ROW
+  # B: (lane%16)*LDS_B_ROW + (lane/16)*16 + wave_n*64*ELEM + LDS_B_OFF
+  LLA, LLB = 40, 43
+  e(v_and_b32_e32(v[50], 15, v[1])); e(v_lshrrev_b32_e32(v[51], 4, v[1]))
+  e(v_lshlrev_b32_e32(v[LLA], 5, v[50]))        # (lane%16) * 32
+  e(v_lshlrev_b32_e32(v[51], 4, v[51]))        # (lane/16) * 16
+  e(v_add_nc_u32_e32(v[LLA], v[LLA], v[51]))
+  e(v_lshlrev_b32_e32(v[52], 11, v[2]))         # wave_m * 2048
+  e(v_add_nc_u32_e32(v[LLA], v[LLA], v[52]))
+  # B read: transposed layout. addr = LDS_B_OFF + (lane%16)*32 + (lane/16)*16 + wave_n*2*512
+  # wave_n selects column panels: wave_n*2 panels (each panel=16 cols, wave_n covers 64 cols = 4 panels)
+  # But wave_n*2*512 = wave_n*1024. Hmm, wave_n covers cols [wave_n*64 : (wave_n+1)*64].
+  # Each panel = 16 cols = 512 bytes. wave_n*64/16 = wave_n*4 panels. Offset = wave_n*4*512 = wave_n*2048.
+  e(v_lshlrev_b32_e32(v[LLB], 5, v[50]))         # (lane%16) * 32 (stride 32!)
+  e(v_add_nc_u32_e32(v[LLB], v[LLB], v[51]))     # + (lane/16)*16
+  e(v_lshlrev_b32_e32(v[52], 11, v[3]))           # wave_n * 2048
+  e(v_add_nc_u32_e32(v[LLB], v[LLB], v[52]))
+  e(v_add_nc_u32_e32(v[LLB], LDS_B_OFF, v[LLB]))
+
+  for i in range(0, 128, 2):
+    e(VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[ACC+i], vdsty=v[ACC+i+1], srcx0=0, srcy0=0))
+  e(s_mov_b32(s[16], 0))
+
+  if not NO_GLOBAL:
+    for i in range(2): e(global_load_b128(vdst=v[DA+i*4:DA+i*4+3], vaddr=v[6:7], saddr=s[4:5], ioffset=i*16))
+    for i in range(2): e(global_load_b128(vdst=v[DB+i*4:DB+i*4+3], vaddr=v[8:9], saddr=s[6:7], ioffset=i*16))
+    e(s_wait_loadcnt(simm16=0))
+  if not NO_DS:
+    for i in range(2): e(ds_store_b128(addr=v[4], data0=v[DA+i*4:DA+i*4+3], offset0=(i*16)&0xFF, offset1=(i*16)>>8))
+    for i in range(2): e(ds_store_b128(addr=v[5], data0=v[DB+i*4:DB+i*4+3], offset0=(i*16)&0xFF, offset1=(i*16)>>8))
+  if not NO_GLOBAL:
+    e(v_add_nc_u32_e32(v[6], BLOCK_K*ELEM, v[6]))
+    e(v_add_nc_u32_e32(v[8], s[14], v[8]))
+
+  # =============================================================================
+  def emit_iter_body(load_set='AB'):
+    if not NO_DS:
+      e(s_wait_dscnt(simm16=0))
+      e(s_barrier_signal(ssrc0=src[193])); e(s_barrier_wait(simm16=0xFFFF))
+    if not NO_GLOBAL:
+      if 'A' in load_set:
+        for i in range(2): e(global_load_b128(vdst=v[DA+i*4:DA+i*4+3], vaddr=v[6:7], saddr=s[4:5], ioffset=i*16))
+        e(v_add_nc_u32_e32(v[6], BLOCK_K*ELEM, v[6]))
+      if 'B' in load_set:
+        for i in range(2): e(global_load_b128(vdst=v[DB+i*4:DB+i*4+3], vaddr=v[8:9], saddr=s[6:7], ioffset=i*16))
+        e(v_add_nc_u32_e32(v[8], s[14], v[8]))
+    if not NO_DS:
+      # Issue 6 loads: A[0:3] + B[0] + B[1]. B[2:3] interleaved with WMMAs.
+      for tm in range(TILES_M):
+        aoff = tm * 16 * LDS_A_ROW
+        e(ds_load_b128(vdst=v[FA+tm*4:FA+tm*4+3], addr=v[LLA], offset0=aoff&0xFF, offset1=aoff>>8))
+      e(ds_load_b128(vdst=v[FB:FB+3], addr=v[LLB], offset0=0, offset1=0))
+      e(ds_load_b128(vdst=v[FB+4:FB+7], addr=v[LLB], offset0=0, offset1=2))
+      e(s_wait_dscnt(simm16=0))  # wait for 6 loads (no stall!)
+    if not NO_ALU:
+      # B[0] WMMAs — issue B[2] during compute
+      if not NO_DS: e(ds_load_b128(vdst=v[FB+8:FB+11], addr=v[LLB], offset0=0, offset1=4))
+      for tm in range(TILES_M):
+        ac = ACC + (tm*TILES_N+0)*8
+        e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB:FB+3], src2=v[ac:ac+7]))
+      # B[1] WMMAs — issue B[3] during compute
+      if not NO_DS:
+        e(ds_load_b128(vdst=v[FB+12:FB+15], addr=v[LLB], offset0=0, offset1=6))
+      for tm in range(TILES_M):
+        ac = ACC + (tm*TILES_N+1)*8
+        e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB+4:FB+7], src2=v[ac:ac+7]))
+      # B[2] WMMAs — B[2] loaded during B[0] WMMAs (~100 cycles ago)
+      if not NO_DS: e(s_wait_dscnt(simm16=1))  # B[2] done, B[3] may still be loading
+      for tm in range(TILES_M):
+        ac = ACC + (tm*TILES_N+2)*8
+        e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB+8:FB+11], src2=v[ac:ac+7]))
+      # B[3] WMMAs
+      if not NO_DS: e(s_wait_dscnt(simm16=0))
+      for tm in range(TILES_M):
+        ac = ACC + (tm*TILES_N+3)*8
+        e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB+12:FB+15], src2=v[ac:ac+7]))
+    if not NO_GLOBAL and not NO_DS: e(s_wait_loadcnt(simm16=0))
+    if not NO_DS:
+      for i in range(2): e(ds_store_b128(addr=v[4], data0=v[DA+i*4:DA+i*4+3], offset0=(i*16)&0xFF, offset1=(i*16)>>8))
+      for i in range(2): e(ds_store_b128(addr=v[5], data0=v[DB+i*4:DB+i*4+3], offset0=(i*16)&0xFF, offset1=(i*16)>>8))
+    e(s_add_co_i32(s[16], s[16], BLOCK_K))
+
+  label('LOOP')
+  emit_iter_body(load_set='A')
+  emit_iter_body(load_set='B')
+  e(s_cmp_lt_i32(s[16], s[17])); e(s_cbranch_scc1(simm16=0)); br(I[-1], 'LOOP')
+
+  emit_iter_body(load_set='AB')  # tail with prefetch
+
+  # Final iteration: no prefetch, no ds_store needed
+  if not NO_DS:
+    e(s_wait_dscnt(simm16=0))
+    e(s_barrier_signal(ssrc0=src[193])); e(s_barrier_wait(simm16=0xFFFF))
+  if not NO_DS:
+    for tm in range(TILES_M):
+      aoff = tm * 16 * LDS_A_ROW
+      e(ds_load_b128(vdst=v[FA+tm*4:FA+tm*4+3], addr=v[LLA], offset0=aoff&0xFF, offset1=aoff>>8))
+    e(ds_load_b128(vdst=v[FB:FB+3], addr=v[LLB], offset0=0, offset1=0))
+    e(ds_load_b128(vdst=v[FB+4:FB+7], addr=v[LLB], offset0=0, offset1=2))
+    e(s_wait_dscnt(simm16=0))
+  if not NO_ALU:
+    if not NO_DS: e(ds_load_b128(vdst=v[FB+8:FB+11], addr=v[LLB], offset0=0, offset1=4))
+    for tm in range(TILES_M):
+      ac = ACC + (tm*TILES_N+0)*8
+      e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB:FB+3], src2=v[ac:ac+7]))
+    if not NO_DS: e(ds_load_b128(vdst=v[FB+12:FB+15], addr=v[LLB], offset0=0, offset1=6))
+    for tm in range(TILES_M):
+      ac = ACC + (tm*TILES_N+1)*8
+      e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB+4:FB+7], src2=v[ac:ac+7]))
+    if not NO_DS: e(s_wait_dscnt(simm16=1))
+    for tm in range(TILES_M):
+      ac = ACC + (tm*TILES_N+2)*8
+      e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB+8:FB+11], src2=v[ac:ac+7]))
+    if not NO_DS: e(s_wait_dscnt(simm16=0))
+    for tm in range(TILES_M):
+      ac = ACC + (tm*TILES_N+3)*8
+      e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB+12:FB+15], src2=v[ac:ac+7]))
+
+  label('EPILOGUE')
+  e(v_and_b32_e32(v[ET], 15, v[1]))
+  e(v_lshrrev_b32_e32(v[ET+1], 4, v[1])); e(v_lshlrev_b32_e32(v[ET+1], 3, v[ET+1]))
+  e(v_lshlrev_b32_e32(v[ET+2], 6, v[2])); e(v_add_nc_u32_e32(v[ET+2], s[11], v[ET+2]))
+  e(v_lshlrev_b32_e32(v[ET+3], 6, v[3])); e(v_add_nc_u32_e32(v[ET+3], s[10], v[ET+3]))
+  e(v_add_nc_u32_e32(v[ET+3], v[ET+3], v[ET])); e(v_mov_b32_e32(v[ET+5], 0))
+
+  for tm in range(TILES_M):
+    for tn in range(TILES_N):
+      ac = ACC + (tm*TILES_N+tn)*8; r_off, c_off = tm*16, tn*16
+      e(v_add_nc_u32_e32(v[ET+6], r_off, v[ET+2])); e(v_add_nc_u32_e32(v[ET+6], v[ET+1], v[ET+6]))
+      e(v_mul_lo_u32(v[ET+4], v[ET+6], s[12])); e(v_add_nc_u32_e32(v[ET+4], v[ET+4], v[ET+3]))
+      if c_off: e(v_add_nc_u32_e32(v[ET+4], c_off, v[ET+4]))
+      e(v_lshlrev_b32_e32(v[ET+4], 1, v[ET+4]))
+      for elem in range(8):
+        e(v_cvt_f16_f32_e32(v[ET+7], v[ac+elem]))
+        e(global_store_b16(vaddr=v[ET+4:ET+5], vsrc=v[ET+7], saddr=s[8:9]))
+        if elem < 7: e(v_add_nc_u32_e32(v[ET+4], s[13], v[ET+4]))
+
+  e(s_wait_storecnt(simm16=0)); e(s_sendmsg(simm16=3)); e(s_endpgm())
+
+  for idx, target in B:
+    off = (L[target] - sum(i.size() for i in I[:idx+1])) // 4
+    assert -32768 <= off <= 32767; I[idx].simm16 = off
+  return I
+
+N = getenv("N", 4096)
+
+def test_matmul():
+  dev = Device[Device.DEFAULT]
+  arch = getattr(dev.renderer, 'arch', 'gfx1200')
+  print(f"Device arch: {arch}")
+  insts = build_kernel(N, arch)
+
+  rng = np.random.default_rng(42)
+  a = Tensor(rng.random((N, N), dtype=np.float32).astype(np.float16))
+  b = Tensor(rng.random((N, N), dtype=np.float32).astype(np.float16))
+  c = Tensor.empty(N, N, dtype=dtypes.half)
+  Tensor.realize(a, b, c)
+
+  grid, local = (N//BLOCK_N, N//BLOCK_M, 1), (THREADS, 1, 1)
+  print(f"Grid: {grid}, Local: {local}")
+
+  dname = Device.DEFAULT
+  def asm_kernel(A, B, C):
+    gidxs = [UOp.special(n, f"gidx{i}") for i,n in enumerate(grid)]
+    lidxs = [UOp.special(THREADS, "lidx0")]
+    lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=max(LDS_SIZE, 65536//getenv("LIMIT_OCC",2)), addrspace=AddrSpace.LOCAL), (), 'lds')
+    sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs,
+                    arg=KernelInfo(name=colored("kernel","cyan"), estimates=Estimates(ops=N*N*N*2, mem=N*N*2*3)))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+
+  c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
+  linear = c.schedule_linear()
+
+  ets = []
+  with Context(DEBUG=2):
+    for _ in range(getenv("CNT", 5)):
+      start = GlobalCounters.time_sum_s
+      run_linear(linear)
+      ets.append(GlobalCounters.time_sum_s - start)
+  print(f"REAL TFLOPS {N*N*N*2 / min(ets) * 1e-12:.2f}")
+
+  if getenv("VERIFY", 1):
+    GlobalCounters.reset()
+    c_np = c.float().numpy()
+    a_np, b_np = a.float().numpy(), b.float().numpy()
+    ref = a_np @ b_np
+    err = np.sqrt(np.mean((c_np - ref)**2)) / np.sqrt(np.mean(ref**2))
+    print(f"relative RMSE {err:.6f}")
+    if err != err or err > 0.05: raise RuntimeError(f"matmul is wrong! RMSE={err}")
+
+if __name__ == "__main__":
+  test_matmul()

tinygrad_repo/extra/gemm/real_pmatmul.py

@@ -0,0 +1,20 @@
+import time
+from tinygrad import Tensor, Device, TinyJit
+from tinygrad.helpers import getenv
+
+if __name__ == "__main__":
+  DEVS = [f"NV:{i}" for i in range(getenv("GPUS", 2))]
+  N = getenv("N", 8192)
+  A = Tensor.rand(N, N).shard(DEVS, 0).realize()
+  B = Tensor.rand(N, N).shard(DEVS, 1).realize()
+  print("***** MUL *****")
+  jmatmul = TinyJit(Tensor.dot)
+  for i in range(10):
+    Device["NV:0"].synchronize()
+    Device["NV:1"].synchronize()
+    st = time.perf_counter()
+    jmatmul(A, B)
+    Device["NV:0"].synchronize()
+    Device["NV:1"].synchronize()
+    et = time.perf_counter()
+    print(f"{(N*N*N*2*1e-12)/(et-st):.2f} TFLOPS")

tinygrad_repo/extra/gemm/simple_conv.py

@@ -0,0 +1,33 @@
+from tinygrad.helpers import getenv
+from tinygrad import dtypes, Tensor
+
+dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float
+acc_dtype = dtypes.half if getenv("ACC_HALF") else dtypes.bfloat16 if getenv("ACC_BFLOAT16") else None
+
+CNT = getenv("CNT", 8)
+BS = getenv("BS", 16)
+CIN = getenv("CIN", 128)
+COUT = getenv("COUT", 128)
+HW = getenv("HW", 128)
+K = getenv("K", 3)
+PADDING = getenv("PADDING", 1)
+COMP = getenv("COMP", 0)
+ATOL = getenv("ATOL", 1e-4)
+RTOL = getenv("RTOL", 3e-2)
+
+FLOPS = BS*K*K*CIN*HW*HW*COUT*2
+def rand_input(): return Tensor.rand(BS, CIN, HW, HW, dtype=dtype_in).realize(), Tensor.rand(COUT, CIN, K, K, dtype=dtype_in).realize()
+
+if __name__ == "__main__":
+  a, b = rand_input()
+  for i in range(CNT):
+    if i > 0 and getenv("RAND", 0) != 0:
+      a, b = rand_input()
+    c = a.conv2d(b, padding=PADDING, dtype=acc_dtype).realize()
+
+  if COMP:
+    import numpy as np, time, torch
+    torch_device = "cuda:0" if torch.cuda.is_available() else ("mps" if getenv("MPS", 0) else "cpu")
+    ta, tb = torch.from_numpy(a.numpy()).to(torch_device), torch.from_numpy(b.numpy()).to(torch_device)
+    tc = torch.nn.functional.conv2d(ta, tb, padding=PADDING)
+    np.testing.assert_allclose(c.numpy(), tc.cpu(), atol=ATOL, rtol=RTOL)

tinygrad_repo/extra/gemm/simple_matmul.py

@@ -0,0 +1,57 @@
+import numpy as np
+from tinygrad import dtypes, Tensor
+from tinygrad.helpers import getenv, get_single_element
+from tinygrad.dtype import _to_np_dtype
+from tinygrad.engine.realize import compile_linear
+from tinygrad.codegen.opt import OptOps
+
+dtype_in = (dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else
+            dtypes.fp8e4m3 if getenv("FP8E4M3") else dtypes.fp8e5m2 if getenv("FP8E5M2") else dtypes.float)
+acc_dtype = (dtypes.half if getenv("ACC_HALF") else dtypes.bfloat16 if getenv("ACC_BFLOAT16") else
+            dtypes.fp8e4m3 if getenv("ACC_FP8E4M3") else dtypes.fp8e5m2 if getenv("ACC_FP8E5M2") else None)
+if getenv("INT"):  dtype_in, acc_dtype = dtypes.int8, dtypes.int32
+if getenv("UINT"): dtype_in, acc_dtype = dtypes.uint8, dtypes.int32
+
+N = getenv("N", 4096)
+M = getenv("M", N)
+K = getenv("K", N)
+CNT = getenv("CNT", 10)
+
+atol, rtol = {dtypes.half:{1e-3, 1e-2}, dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1)}.get(dtype_in, (1e-4, 3e-2))
+ATOL, RTOL = getenv("ATOL", atol), getenv("RTOL", rtol)
+
+INT_LOW = getenv("INT_LOW", 0)
+INT_HIGH = getenv("INT_HIGH", 10)
+
+if __name__ == "__main__":
+  def init_matrix(rows, cols):
+    rng = np.random.default_rng()
+    # NOTE: numpy does not support bfloat16
+    if (np_dtype := _to_np_dtype(dtype_in)) is None: np_dtype = np.float32
+    if dtype_in in dtypes.ints:
+      return Tensor(rng.integers(INT_LOW, INT_HIGH, (rows, cols), dtype=np_dtype)).realize()
+    return Tensor(rng.random((rows, cols), dtype=np.float32).astype(np_dtype)-0.5).cast(dtype_in).realize()
+
+  a, b = init_matrix(M, K), init_matrix(K, N)
+  for i in range(CNT):
+    if i > 0 and getenv("RAND", 0) != 0:
+      a, b = init_matrix(M, K), init_matrix(K, N)
+    c = a.matmul(b, dtype=acc_dtype).realize()
+
+  if getenv("SHOULD_USE_TC"):
+    linear = compile_linear(a.matmul(b, dtype=acc_dtype).schedule_linear())
+    call = get_single_element(list(linear.src))
+    applied_opts = call.src[0].src[0].arg.applied_opts
+    assert any(opt.op is OptOps.TC for opt in applied_opts), f"TC not triggered, {applied_opts}"
+
+  ref = a.numpy().astype(np.float32) @ b.numpy().astype(np.float32)
+  res = c.numpy()
+  try:
+    np.testing.assert_allclose(res, ref, rtol=RTOL, atol=ATOL)
+  except AssertionError as e:
+    if getenv("DEBUG_VALUES", 0) > 0:
+      mismatch = np.where(~np.isclose(res, ref, rtol=RTOL, atol=ATOL))
+      print("Mismatch indices:", mismatch)
+      print("Result          :", res[mismatch])
+      print("Ground truth    :", ref[mismatch])
+    raise e

tinygrad_repo/extra/gemm/simple_matvec.py

@@ -0,0 +1,30 @@
+import numpy as np
+from tinygrad.helpers import getenv
+from tinygrad import dtypes, Tensor, Device
+dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float
+acc_dtype = dtypes.half if getenv("ACC_HALF") else dtypes.bfloat16 if getenv("ACC_BFLOAT16") else None
+
+GPUS = getenv("GPUS", 0)
+M = getenv("M", 16384)
+N = getenv("N", 4096)
+CNT = getenv("CNT", 10)
+ATOL = getenv("ATOL", 1e-4)
+RTOL = getenv("RTOL", 3e-2)
+
+def _rand(device):
+  a, b = Tensor.rand(M, N, dtype=dtype_in).realize(), Tensor.rand(N, dtype=dtype_in).realize()
+  if isinstance(device, tuple):
+    a.shard_(device, axis=1)
+    b.shard_(device, axis=0)
+  return a, b
+
+if __name__ == "__main__":
+  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(GPUS)) if GPUS > 1 else Device.DEFAULT
+  a, b = _rand(device)
+  for i in range(CNT):
+    if i > 0 and getenv("RAND", 0) != 0:
+      a, b = _rand(device)
+    c = a.matmul(b, dtype=acc_dtype).realize()
+  nc = c.numpy()
+  comp = a.numpy().astype(np.float32) @ b.numpy().astype(np.float32)
+  np.testing.assert_allclose(nc, comp, atol=ATOL, rtol=RTOL)

tinygrad_repo/extra/gemm/tinygrad_nv_matmul.py

@@ -0,0 +1,34 @@
+from tinygrad import Tensor, dtypes, Context
+from tinygrad.helpers import getenv
+from tinygrad.codegen.opt import Opt, OptOps
+from tinygrad.engine.realize import run_linear
+from dataclasses import replace
+
+N = 4096
+if __name__ == "__main__":
+  if getenv("GEMV"):
+    A, B = Tensor.empty(1, N, dtype=dtypes.float), Tensor.empty(14336, N, dtype=dtypes.float16).T
+  else:
+    A, B = Tensor.empty(N, N, dtype=dtypes.float16), Tensor.empty(N, N, dtype=dtypes.float16)
+  C = A.matmul(B)
+  if getenv("GEMV"):
+    opts = [
+      Opt(op=OptOps.UNROLL, axis=0, amt=8),
+      Opt(op=OptOps.GROUP, axis=0, amt=32),
+    ]
+  else:
+    opts = [
+      Opt(op=OptOps.TC, axis=0, amt=0),
+      Opt(op=OptOps.UPCAST, axis=0, amt=4),
+      Opt(op=OptOps.UPCAST, axis=1, amt=8),
+      Opt(op=OptOps.LOCAL, axis=0, amt=2),
+      Opt(op=OptOps.LOCAL, axis=1, amt=2),
+      Opt(op=OptOps.LOCAL, axis=0, amt=2),
+    ]
+  linear = C.schedule_linear()
+  call = linear.src[-1]
+  new_ast = call.src[0].replace(arg=replace(call.src[0].arg, opts_to_apply=tuple(opts)))
+  new_call = call.replace(src=(new_ast, *call.src[1:]))
+  linear = linear.replace(src=tuple(new_call if c is call else c for c in linear.src))
+  with Context(DEBUG=2):
+    for i in range(5): run_linear(linear)

tinygrad_repo/extra/gemm/torch_gemm.py

@@ -0,0 +1,30 @@
+import os
+os.environ["NVIDIA_TF32_OVERRIDE"] = "0"
+os.environ["MKL_NUM_THREADS"] = "1"
+os.environ["NUMEXPR_NUM_THREADS"] = "1"
+os.environ["OMP_NUM_THREADS"] = "1"
+import time
+import torch
+torch.set_num_threads(1)
+from tinygrad.helpers import getenv
+CUDA = getenv("CUDA", 1)
+MPS = getenv("MPS", 0)
+if getenv("FP16_ACC"): torch.backends.cuda.matmul.allow_fp16_accumulation = True
+
+for dtype in [torch.float32, torch.float16, torch.bfloat16]:
+  for N in [256, 512, 1024, 2048, 4096] + ([6144, 8192] if getenv("BIG") else []):
+    FLOPS = N*N*N*2
+
+    b = torch.rand((N,N), dtype=dtype)
+    c = torch.rand((N,N), dtype=dtype)
+    if CUDA: b,c = b.cuda(),c.cuda()
+    if MPS: b,c = b.to('mps'),c.to('mps')
+
+    def torch_prog(b, c):
+      st = time.perf_counter()
+      a = b@c
+      if CUDA: torch.cuda.synchronize()
+      if MPS: torch.mps.synchronize()
+      return time.perf_counter() - st
+    tm = min([torch_prog(b, c) for _ in range(20)])
+    print(f"{N*N:10d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS {N:4d}x{N:4d}x{N:4d} matmul in {dtype}")

tinygrad_repo/extra/gemm/triton_nv_matmul.py

@@ -0,0 +1,117 @@
+import time
+import triton
+import triton.language as tl
+from triton.compiler import AttrsDescriptor, ASTSource, compile as triton_compile
+import numpy as np
+from tinygrad import Tensor, dtypes, Device
+from tinygrad.engine.realize import get_runtime
+from tinygrad.codegen import to_program
+from tinygrad.uop.ops import Ops, UOp, KernelInfo, ProgramInfo
+from tinygrad.helpers import getenv
+np.set_printoptions(suppress=True)
+
+@triton.jit
+def matmul_kernel(c_ptr, a_ptr, b_ptr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr):
+  pid_m = tl.program_id(axis=0)
+  pid_n = tl.program_id(axis=1)
+
+  M, N, K = 4096, 4096, 4096
+  stride_am = 4096
+  stride_ak = 1
+  stride_bk = 4096
+  stride_bn = 1
+  stride_cm = 4096
+  stride_cn = 1
+
+  offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+  offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
+  offs_k = tl.arange(0, BLOCK_SIZE_K)
+
+  a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
+  b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
+
+  accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+  for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+    a = tl.load(a_ptrs)
+    b = tl.load(b_ptrs)
+
+    accumulator = tl.dot(a, b, accumulator)
+    a_ptrs += BLOCK_SIZE_K * stride_ak
+    b_ptrs += BLOCK_SIZE_K * stride_bk
+
+  c = tl.cast(accumulator, tl.float16)
+  offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+  offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+  c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
+  tl.store(c_ptrs, c)
+
+# CUDA=1 CUDA_PTX=1 python3 extra/gemm/triton_nv_matmul.py
+if __name__ == "__main__":
+  BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 64, 128, 64
+  M, N, K = 4096, 4096, 4096
+
+  # **** torch test ****
+
+  if getenv("TORCH"):
+    import torch
+    c = torch.empty((M, N), device='cuda:0', dtype=torch.float16)
+    a = torch.empty((M, K), device='cuda:0', dtype=torch.float16)
+    b = torch.empty((K, N), device='cuda:0', dtype=torch.float16)
+
+    for i in range(5):
+      st = time.perf_counter()
+      matmul_kernel[triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N)](
+        c, a, b, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K)
+      torch.cuda.synchronize()
+      et = time.perf_counter() - st
+      print(f"TFLOPS {2*M*N*K*1e-12/et:.2f}")
+
+  # **** tinygrad test ****
+
+  compiled = triton_compile(ASTSource(matmul_kernel, "*fp16,*fp16,*fp16",
+                            attrs=AttrsDescriptor(divisible_by_16=(0, 1, 2, 3, 4, 5), equal_to_1=()),
+                            constants={"BLOCK_SIZE_M": BLOCK_SIZE_M, "BLOCK_SIZE_N": BLOCK_SIZE_N, "BLOCK_SIZE_K": BLOCK_SIZE_K}))
+  print(compiled.metadata)
+
+  A, B = Tensor.normal(M, K, std=1e-1, dtype=dtypes.float16).realize(), Tensor.normal(K, N, std=1e-1, dtype=dtypes.float16).realize()
+  C = A.matmul(B)
+  from tinygrad.uop.ops import Ops
+  linear, var_vals = C.linear_with_vars()
+  last_call = linear.src[-1]
+  ast = last_call.src[0]
+  bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]
+
+  src = compiled.asm["ptx"]
+  # specify the shared memory here so we don't need to do it dynamically
+  src = src.replace(".extern .shared .align 16 .b8 global_smem[];", f".shared .align 16 .b8 global_smem[{compiled.metadata.shared}];")
+  # useless comment spam
+  src = src.replace("\t// begin inline asm\n", "")
+  src = src.replace("\t// end inline asm\n", "")
+  # remove debug sections
+  src = src.split("\t.file")[0]
+  assert '.extern .shared' not in src
+  info = ProgramInfo(name="matmul_kernel",
+                     global_size=(M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1), local_size=(32*compiled.metadata.num_warps, 1, 1))
+  sink = UOp.sink(arg=KernelInfo(name="matmul_kernel"))
+  prg_uop = to_program(UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=src)), arg=info),
+                       Device.default.renderer)
+  rt = get_runtime(Device.DEFAULT, prg_uop)
+  all_bufs = [x.ensure_allocated() for x in bufs]
+  prg_bufs = [all_bufs[i] for i in info.globals]
+  gsize, lsize = info.launch_dims({})
+  tflops = []
+  for i in range(5):
+    tm = rt(*[b._buf for b in prg_bufs], global_size=gsize, local_size=lsize, vals=info.vals({}), wait=True)
+    tflops.append((2*M*K*N/tm)*1e-12)
+  print(f"TFLOPS: {max(tflops):.2f}")
+
+  # check correctness
+  if getenv("VERIFY"):
+    from tinygrad.engine.realize import run_linear
+    triton_buf = np.frombuffer(si.bufs[0].as_memoryview(), np.float16).reshape(M,N)
+    print(triton_buf)
+    run_linear(linear, var_vals)
+    tinygrad_buf = np.frombuffer(si.bufs[0].as_memoryview(), np.float16).reshape(M,N)
+    print(tinygrad_buf)
+    np.testing.assert_allclose(triton_buf, tinygrad_buf)
+    print("correct!")

tinygrad_repo/extra/gemm/tvm_gemm.py

@@ -0,0 +1,46 @@
+# https://tvm.apache.org/docs/tutorial/tensor_expr_get_started.html#example-2-manually-optimizing-matrix-multiplication-with-te
+
+M, N, K = 1024, 1024, 1024
+
+try:
+  import tvm
+  from tvm import te
+  #print(tvm.target.Target.list_kinds())
+
+  # c, opencl
+  target = tvm.target.Target(target="c")
+
+  # TVM Matrix Multiplication using TE
+  k = te.reduce_axis((0, K), "k")
+  A = te.placeholder((M, K), name="A")
+  B = te.placeholder((K, N), name="B")
+  C = te.compute((M, N), lambda x, y: te.sum(A[x, k] * B[k, y], axis=k), name="C")
+
+  # Default schedule
+  s = te.create_schedule(C.op)
+  #print(tvm.lower(s, [A, B, C], simple_mode=True))
+
+  # Output C code
+  func = tvm.build(s, [A, B, C], target=target, name="mmult")
+  print(func.get_source())
+except ImportError:
+  print("** please install TVM for TVM output")
+
+# tinygrad version
+
+import os
+from tinygrad.tensor import Tensor
+
+# define the compute
+A = Tensor.rand(M, K, device="CPU")
+B = Tensor.rand(K, N, device="CPU")
+C = (A.reshape(M, 1, K) * B.permute(1,0).reshape(1, N, K)).sum(axis=2)
+
+linear = C.schedule_linear()
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.device import CompilerOptions
+lin = Kernel(linear.src[-1].src[0], CompilerOptions(has_local=False, supports_float4=False))
+lin.to_program()
+from tinygrad.runtime.ops_cpu import renderer
+src = renderer("mmult", lin.uops)
+print(src)

tinygrad_repo/extra/gradcheck.py

@@ -0,0 +1,51 @@
+import numpy as np
+from tinygrad.tensor import Tensor, _to_np_dtype
+
+def mask_like(like, mask_inx, mask_value = 1.0):
+  mask = np.zeros(like.shape, dtype=_to_np_dtype(like.dtype)).reshape(-1)
+  mask[mask_inx] = mask_value
+  return mask.reshape(like.shape)
+
+def jacobian(func, input):
+  output = func(input)
+
+  ji = input.numpy().reshape(-1).shape[-1]
+  jo = output.numpy().reshape(-1).shape[-1]
+  J = np.zeros((jo,ji), dtype=np.float32)
+
+  for o in range(jo):
+    input.grad = None
+    output = func(input)
+
+    # tinygrad doesn't support slicing, tiny-hack to select
+    # the needed scalar an backpropagate only through it
+    o_scalar = Tensor(mask_like(output, o, 1.)).mul(output).sum()
+    o_scalar = Tensor(mask_like(output, o, 1.)).mul(output).sum()
+    o_scalar.backward()
+
+    for i, grad in enumerate(input.grad.numpy().reshape(-1)):
+      J[o,i] = grad
+  return J
+
+def numerical_jacobian(func, input, eps = 1e-3):
+  output = func(input)
+
+  ji = input.numpy().reshape(-1).shape[-1]
+  jo = output.numpy().reshape(-1).shape[-1]
+  NJ = np.zeros((jo, ji), dtype=np.float32)
+
+  for i in range(ji):
+    eps_perturb = mask_like(input, i, mask_value = eps)
+
+    output_perturb_add = func(Tensor(input.numpy() + eps_perturb)).numpy().reshape(-1)
+    output_perturb_sub = func(Tensor(input.numpy() - eps_perturb)).numpy().reshape(-1)
+
+    grad_approx = ((output_perturb_add) - (output_perturb_sub)) / (2*eps)
+
+    NJ[:,i] = grad_approx
+  return NJ
+
+def gradcheck(func, input, eps = 1e-3, atol = 1e-3, rtol = 1e-3):
+  NJ = numerical_jacobian(func, input, eps)
+  J = jacobian(func, input)
+  return np.allclose(J, NJ, atol = atol, rtol = rtol)

tinygrad_repo/extra/hcq/hcq_smi.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python3
+
+import argparse, glob, os, time, subprocess, sys
+from tinygrad.helpers import temp
+
+def scan_devs_based_on_lock(prefix:str, args) -> list[str]:
+  target_dev = args.pci_bus if 'pci_bus' in args.__dir__() else ""
+
+  devs = []
+  for dev in glob.glob(temp(f'{prefix}_*.lock')):
+    dev_id = dev.split('/')[-1][len(prefix)+1:-5]
+    if dev_id.startswith(target_dev): devs.append(dev_id)
+  return devs
+
+def _do_reset_device(pci_bus): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{pci_bus}/reset'")
+def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")
+
+def cmd_remove_module(args):
+  modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia", "ast"] if args.backend == "nv" else ["amdgpu"]
+  to_unload = [m for m in modules if _is_module_loaded(m)]
+  if not to_unload: print("Kernel modules are not loaded")
+  else:
+    print("Removing kernel modules:", ", ".join(to_unload))
+    try: subprocess.run(["sudo", "modprobe", "-r", *to_unload], check=True)
+    except subprocess.CalledProcessError as e:
+      print("Failed to unload all modules — they may be in use.", file=sys.stderr)
+      sys.exit(e.returncode)
+
+def cmd_insert_module(args):
+  cmd_remove_module(args)
+  cmd_reset_devices(args)
+
+  module = "nvidia" if args.backend == "nv" else "amdgpu"
+  if _is_module_loaded(module):
+    print(f"{module} kernel module already loaded")
+    return
+
+  print(f"Inserting kernel module: {module}")
+  if args.backend == "nv":
+    subprocess.run(["nvidia-smi"], check=True)
+  elif args.backend == "amd":
+    subprocess.run(["sudo", "modprobe", "amdgpu"], check=True)
+
+def cmd_reset_devices(args):
+  devs = scan_devs_based_on_lock({"amd":"am", "nv":"nv"}[args.backend], args)
+
+  for dev in devs:
+    print(f"Resetting device {dev}")
+    if args.backend != "amd": _do_reset_device(dev)
+    time.sleep(0.2)
+
+def cmd_show_pids(args):
+  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
+
+  for dev in devs:
+    try:
+      pid = subprocess.check_output(['sudo', 'lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
+      print(f"{dev}: {pid}")
+    except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
+
+def cmd_kill_pids(args):
+  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
+
+  for dev in devs:
+    for i in range(128):
+      if i > 0: time.sleep(0.2)
+
+      try:
+        try: pid = subprocess.check_output(['sudo', 'lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
+        except subprocess.CalledProcessError: break
+
+        print(f"Killing process {pid} (which uses {dev})")
+        subprocess.run(['sudo', 'kill', '-9', pid], check=True)
+      except subprocess.CalledProcessError as e:
+        print(f"Failed to kill process for device {dev}: {e}", file=sys.stderr)
+
+def add_common_commands(parent_subparsers):
+  p_insmod = parent_subparsers.add_parser("insmod", help="Insert a kernel module")
+  p_insmod.set_defaults(func=cmd_insert_module)
+
+  p_rmmod = parent_subparsers.add_parser("rmmod", help="Remove a kernel module")
+  p_rmmod.set_defaults(func=cmd_remove_module)
+
+  p_reset = parent_subparsers.add_parser("reset", help="Reset a device")
+  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device to reset")
+  p_reset.set_defaults(func=cmd_reset_devices)
+
+  p_reset = parent_subparsers.add_parser("pids", help="Show pids of processes using the device")
+  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
+  p_reset.set_defaults(func=cmd_show_pids)
+
+  p_reset = parent_subparsers.add_parser("kill_pids", help="Kill pids of processes using the device")
+  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
+  p_reset.set_defaults(func=cmd_kill_pids)
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  backend_subparsers = parser.add_subparsers(dest="backend", required=True, metavar="{nv,amd}", help="Hardware backend to target")
+
+  nv_parser = backend_subparsers.add_parser("nv", help="NVIDIA GPUs")
+  nv_commands = nv_parser.add_subparsers(dest="command", required=True)
+  add_common_commands(nv_commands)
+
+  amd_parser = backend_subparsers.add_parser("amd", help="AMD GPUs")
+  amd_commands = amd_parser.add_subparsers(dest="command", required=True)
+  add_common_commands(amd_commands)
+
+  args = parser.parse_args()
+  if args.command is None:
+    parser.print_help(sys.stderr)
+    sys.exit(1)
+
+  args.func(args)

tinygrad_repo/extra/hcq2/graph/hcq.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+import time
+from typing import cast
+from tinygrad.device import Buffer, BufferSpec, Compiled, Device, MultiBuffer
+from tinygrad.dtype import dtypes
+from tinygrad.engine.jit import GraphRunner
+from tinygrad.engine.realize import get_call_outs_ins, get_runtime
+from tinygrad.helpers import round_up, ceildiv
+from tinygrad.runtime.support.memory import BumpAllocator
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, graph_rewrite
+from extra.hcq2.hcq2 import HCQ2Compiled, HCQ2DeviceCtx, HCQ2LowerCtx, pm_prep_runtime, pm_lower_ops
+from extra.hcq2.hcq2 import pm_split_into_queues, pm_add_barriers, pm_add_signals
+from extra.hcq2.hcq2 import pm_bufferize, pm_lift_patches_to_cmdbuf, pm_resolve_patches, pm_parametrize_host_buffers
+from extra.hcq2.hcq2 import pm_add_timeline_inc, pm_callify, pm_calc_kernargs_sizes
+
+# **************** insert deps ****************
+
+def insert_deps(ctx:HCQ2Graph, linear:UOp) -> UOp:
+  src = []
+  for j, call in enumerate(linear.src):
+    call = call.replace(tag=j)
+    _, _, bufs, _ = ctx.calls[j]
+    outs, ins = get_call_outs_ins(call)
+    deps = ctx._access_resources([bufs[i] for i in outs + ins], list(range(len(outs))), call)
+    src.append(UOp(Ops.AFTER, call.dtype, (call, *deps), tag=call.tag))
+  return linear.replace(src=tuple(src))
+pm_insert_deps = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), insert_deps)])
+
+pm_replace_params = PatternMatcher([
+  (UPat(Ops.PARAM, name="p"), lambda ctx, p: ctx.input_addrs_uop.index(UOp.const(dtypes.int, p.arg))),
+  (UPat(Ops.SLICE, src=(UPat(Ops.INDEX, name="addr"), UPat(Ops.CONST, dtype=dtypes.weakint, name="off")), name="bv"),
+    lambda ctx, bv, addr, off: addr.cast(dtypes.uint64) + UOp.const(dtypes.uint64, off.arg * ctx.input_uops[addr.src[1].arg].dtype.itemsize)),
+])
+
+# **************** graph-only passes ****************
+
+def alloc_queue_sig(ctx:HCQ2Graph, q:UOp) -> None:
+  if q.arg in ctx.queue_sigs: return None
+  dev = q.arg[0][0]  # TODO: multi device
+  buf = Buffer(dev, 0x100, dtypes.uint8, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
+  ctx.queue_sig_bufs.append(buf)
+  ctx.queue_sigs[q.arg] = UOp.from_buffer(buf, dev)
+  return None
+pm_alloc_queue_sigs = PatternMatcher([(UPat(Ops.LINEAR, src=UPat({Ops.PROGRAM, Ops.COPY}), name="q"), alloc_queue_sig)])
+
+def lower_queue_deps(ctx:HCQ2Graph, after:UOp) -> UOp:
+  wrapper, deps, call_idx = after.src[0], after.src[1:], after.tag
+  def store(q_arg, v): return ctx.queue_sigs[q_arg].store(UOp.const(dtypes.uint32, v))
+  waits = tuple(UOp(Ops.WAIT, dtypes.void, (ctx.queue_sigs[dep.src[0].arg], UOp.const(dtypes.uint32, dep.tag),
+                                            store(dep.src[0].arg, dep.tag))) for dep in deps)
+  return wrapper.replace(src=tuple(q.replace(src=(*waits, *q.src, store(q.arg, call_idx))) for q in wrapper.src))
+pm_lower_queue_deps = PatternMatcher([(UPat(Ops.AFTER, src=UPat(Ops.LINEAR), name="after"), lower_queue_deps)])
+
+def optimize_queue_deps(ctx:HCQ2Graph, queue:UOp) -> UOp|None:
+  src, seen, pending, queue_sig = [], {}, {}, ctx.queue_sigs[queue.arg]
+  for x in queue.src:
+    if x.op is Ops.WAIT:
+      sig, val = x.src[0], x.src[1]
+      if sig is queue_sig or seen.get(sig, -1) >= val.arg: continue
+      if (old:=pending.get(sig)) is None or old.src[1].arg < val.arg: pending[sig] = x
+      continue
+    for wait in pending.values():
+      src.append(wait)
+      seen[wait.src[0]] = wait.src[1].arg
+    pending.clear()
+    src.append(x)
+  src += pending.values()
+  return queue.replace(src=tuple(src)) if tuple(src) != queue.src else None
+pm_optimize_queue_deps = PatternMatcher([
+  (UPat(Ops.LINEAR, src=UPat({Ops.BARRIER, Ops.WAIT, Ops.STORE, Ops.PROGRAM, Ops.COPY}), name="queue"), optimize_queue_deps),
+])
+
+def drop_dead_stores(ctx:HCQ2Graph, outer:UOp) -> UOp:
+  live = {u.src[2] for u in outer.toposort() if u.op is Ops.WAIT}
+  return outer.replace(src=tuple(q.replace(src=tuple(x for x in q.src if x.op is not Ops.STORE or x in live)) for q in outer.src))
+pm_drop_dead_stores = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR), name="outer"), drop_dead_stores)])
+
+def add_queue_sig_resets(ctx:HCQ2Graph, x:UOp, cmdbuf:UOp) -> UOp|None:
+  if not ctx.queue_sig_bufs or cmdbuf.tag not in ("compute", "copy"): return None
+  resets = tuple((b:=UOp.from_buffer(sig)).index(UOp.const(dtypes.int, 0), dtype=b.dtype.ptr())
+                 .cast(dtypes.uint64.ptr()).store(UOp.const(dtypes.uint64, 0)) for sig in ctx.queue_sig_bufs)
+  return x.replace(src=x.src + resets)
+pm_add_queue_sig_resets = PatternMatcher([(UPat(Ops.AFTER, src=(UPat(Ops.BUFFER, name="cmdbuf"),), allow_any_len=True, name="x"),
+                                           add_queue_sig_resets)])
+
+# **************** Graph ****************
+
+class HCQ2Graph(GraphRunner):
+  def __init__(self, linear:UOp, input_uops:tuple[UOp, ...]=()):
+    super().__init__(linear, input_uops)
+    self.dev = cast(HCQ2Compiled, Device[self.device])
+    self.hcq_ctx = HCQ2LowerCtx(name="hcq_graph")
+
+    self.input_addrs = Buffer("CPU", max(len(input_uops), 1), dtypes.uint64, preallocate=True)
+    self.input_addrs_uop = UOp.from_buffer(self.input_addrs, "CPU")
+
+    self.linear = graph_rewrite(self.linear, pm_insert_deps, ctx=self, name="hcq: insert deps", walk=True)
+    self.linear = graph_rewrite(self.linear, pm_replace_params, ctx=self, name="hcq: replace params", walk=True)
+    self.linear = graph_rewrite(self.linear, pm_prep_runtime, ctx=self.hcq_ctx, name="hcq: prepare runtime")
+    self.linear = graph_rewrite(self.linear, pm_lower_ops, ctx=self.hcq_ctx, name="hcq: lower ops")
+
+    # per-queue signal state — populated as a side-effect by pm_alloc_queue_sigs walking the lowered linear.
+    self.queue_sig_bufs:list[Buffer] = []
+    self.queue_sigs:dict[tuple[str, str], UOp] = {}
+    graph_rewrite(self.linear, pm_alloc_queue_sigs, ctx=self, name="hcq: alloc queue sigs", walk=True)
+
+    self.linear = graph_rewrite(self.linear, pm_lower_queue_deps, ctx=self, name="hcq: lower queue deps")
+    self.linear = graph_rewrite(self.linear, pm_split_into_queues, ctx=self.hcq_ctx, name="hcq: split into queues")
+    self.linear = graph_rewrite(self.linear, pm_add_barriers, ctx=self.hcq_ctx, name="hcq: add barriers", walk=True)
+    self.linear = graph_rewrite(self.linear, pm_optimize_queue_deps, ctx=self, name="hcq: optimize queue deps", walk=True)
+    self.linear = graph_rewrite(self.linear, pm_drop_dead_stores, ctx=self, name="hcq: drop dead stores")
+    self.linear = graph_rewrite(self.linear, pm_add_signals, ctx=self.hcq_ctx, name="hcq: add signals", walk=True)
+    self.linear = graph_rewrite(self.linear, pm_add_timeline_inc, ctx=self.hcq_ctx, name="hcq: add submit", walk=True)
+    self.linear = graph_rewrite(self.linear, self.dev.pm_lower, ctx=self.hcq_ctx, name=f"hcq: encode cmdbuf {self.dev.device}", walk=True)
+
+    graph_rewrite(self.linear, pm_calc_kernargs_sizes, ctx=(sizes:={}), name=None)
+    for dev_name, sz in sizes.items():
+      buf = Buffer(dev_name, sz, dtypes.uint8, options=BufferSpec(cpu_access=True), preallocate=True)
+      self.hcq_ctx.dev_ctx[dev_name] = HCQ2DeviceCtx(dev_name, UOp.from_buffer(buf, dev_name), UOp.const(dtypes.uint64, buf._buf.va_addr))
+
+    self.linear = graph_rewrite(self.linear, pm_bufferize, ctx=self.hcq_ctx, bottom_up=True, name="realize binaries")
+    self.linear = graph_rewrite(self.linear, pm_lift_patches_to_cmdbuf, ctx=self.hcq_ctx, bottom_up=False, name="lift patches to cmdbuf")
+    self.linear = graph_rewrite(self.linear, pm_resolve_patches, ctx=self.hcq_ctx, bottom_up=False, name="simplify patches")
+    self.linear = graph_rewrite(self.linear, pm_add_queue_sig_resets, ctx=self, name="hcq: add queue sig resets", walk=True)
+    self.linear = graph_rewrite(self.linear, pm_parametrize_host_buffers, ctx=self.hcq_ctx, bottom_up=True, name="parametrize host buffers")
+    self.host_call = graph_rewrite(self.linear, pm_callify, ctx=self.hcq_ctx, name="hcq: callify")
+
+    self.host_rt, self.host_globals = get_runtime("CPU", self.host_call.src[0]), self.host_call.src[0].arg.globals
+
+  def __call__(self, input_uops:tuple[UOp, ...], var_vals:dict[str, int], wait=False) -> float|None:
+    addrs = self.input_addrs.as_memoryview(force_zero_copy=True).cast('Q')
+    for i, u in enumerate(input_uops):
+      buf = next(b for b in u.buffer.bufs if b.device == self.dev.device) if isinstance(u.buffer, MultiBuffer) else u.buffer
+      addrs[i] = buf._buf.va_addr
+    self.host_rt(*[self.hcq_ctx.inputs[i].get_buf("CPU") for i in self.host_globals], vals=self.host_call.src[0].arg.vals(var_vals), wait=True)
+    if wait:
+      st = time.perf_counter()
+      self.dev.synchronize()
+      return time.perf_counter() - st
+    return None
+
+  @staticmethod
+  def supports_uop(batch_devs:list[Compiled], new_call:UOp) -> bool:
+    all_devs = GraphRunner._all_devs(batch_devs, new_call)
+    return new_call.src[0].op in (Ops.PROGRAM, Ops.COPY) and len(all_devs) == 1 and isinstance(all_devs[0], HCQ2Compiled)

tinygrad_repo/extra/hcq2/hcq2.py

@@ -0,0 +1,395 @@
+from __future__ import annotations
+from typing import cast, Callable, TypeVar, Generic, Any, TYPE_CHECKING
+import struct, functools, time, collections, importlib, itertools
+from dataclasses import replace
+if TYPE_CHECKING: from tinygrad.engine.realize import ExecContext
+from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, mv_address, round_up, DEBUG, dedup, pluralize
+from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer
+from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites, GroupOp
+from tinygrad.uop.symbolic import symbolic_simple, symbolic
+from tinygrad.dtype import dtypes, DType
+from dataclasses import dataclass, field
+from tinygrad.runtime.support.memory import BumpAllocator
+from tinygrad.runtime.support.hcq import MMIOInterface
+from tinygrad.renderer import Renderer, Estimates
+from tinygrad.engine.realize import to_program, track_stats, get_call_arg_uops, resolve_params, pm_flatten_linear
+
+HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled')
+
+class HCQ2Compiled(Compiled):
+  timestamp_divider: float = 1000.0 # GPU timestamp counter ticks per microsecond; override per device
+
+  def __init__(self, device:str, allocator:'HCQAllocator', compilers:list[type[Renderer]], runtime, can_recover:bool=False, arch=None):
+    self.device_id:int = int(device.split(":")[1]) if ":" in device else 0
+
+    # default pm bufferize
+    self.pm_bufferize = PatternMatcher([
+      (UPat(Ops.BUFFER, tag="timeline_signal"), lambda ctx: ctx.timeline_signal),
+      (UPat(Ops.BUFFER, tag="timeline_value"), lambda ctx: ctx.timeline_value),
+      (UPat(Ops.BUFFER, name="b"), lambda ctx, b: Buffer(ctx.device, b.arg, b.dtype, options=BufferSpec(host=True, uncached=True, cpu_access=True))),
+    ])
+
+    super().__init__(device, allocator, compilers, lambda *a, **kw: None, None, arch=arch)
+
+  @functools.cached_property
+  def timeline_signal(self) -> Buffer:
+    return Buffer(self.device, 0x100, dtypes.uint8, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
+
+  @functools.cached_property
+  def timestamps_buf(self) -> Buffer:
+    return Buffer(self.device, 0x100, dtypes.uint8, options=BufferSpec(cpu_access=True), preallocate=True)
+
+  @functools.cached_property
+  def timeline_value(self) -> Buffer:
+    buf = Buffer("CPU", 1, dtypes.uint64, preallocate=True)
+    buf.as_memoryview(force_zero_copy=True).cast('Q')[0] = 1
+    return buf
+
+  def synchronize(self, timeout:int|None=None):
+    if not hasattr(self, 'iface'): return
+    sig = self.timeline_signal._buf.cpu_view().mv.cast('Q')
+    tl = self.timeline_value.as_memoryview(force_zero_copy=True).cast('Q')
+    st = time.perf_counter()
+    while sig[0] < tl[0] - 1:
+      if time.perf_counter() - st > (timeout or 3000) / 1000: self.on_device_hang()
+
+  def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.
+
+  def count(self) -> int: return self.iface.count if hasattr(self, 'iface') else 1
+
+  def _select_iface(self):
+    assert (v:=getenv(k:=f'{type(self).__name__[:-6].upper()}_IFACE', "")) == "",  \
+      f"{k}={v} is deprecated, use DEV={replace(DEV.target(type(self).__name__[:-6]), interface=v)} instead"
+    assert hasattr(self, "ifaces"), "must have ifaces to select an iface"
+    t = DEV.target(dev:=type(self).__name__[:-6])
+    filtered = select_by_name(self.ifaces, lambda i: i.__name__[:-5], t.interface, f"{dev} has no interface {t.interface!r}")
+    filtered = [i for i in filtered if t.interface.startswith("MOCK") or not i.__name__[:-5].startswith("MOCK")] # never fall back to mock ifaces
+    return select_first_inited([functools.partial(cast(Callable, iface), self, self.device_id) for iface in filtered],
+                               f"No interface for {dev}:{self.device_id} is available")
+
+  def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"
+
+  def finalize(self):
+    try: self.synchronize() # try to finalize the device in any case
+    except RuntimeError as e: print(f"{self.device} synchronization failed before finalizing: {e}")
+
+    # if the device has an interface, call device_fini to clean up resources
+    if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini()
+
+class HCQ2Buffer:
+  def __init__(self, va_addr:sint, size:int, meta:Any=None, _base:HCQ2Buffer|None=None, view:MMIOInterface|None=None, owner:HCQ2Compiled|None=None):
+    self.va_addr, self.size, self.meta, self._base, self.view, self.owner = va_addr, size, meta, _base, view, owner
+
+  def offset(self, offset:int=0, size:int|None=None) -> HCQ2Buffer:
+    return HCQ2Buffer(self.va_addr+offset, size or (self.size - offset), owner=self.owner, meta=self.meta,
+      _base=self._base or self, view=(self.view.view(offset=offset, size=size) if self.view is not None else None))
+
+  def cpu_view(self) -> MMIOInterface:
+    assert self.view is not None, "buffer has no cpu_view"
+    return self.view
+
+  @property
+  def base(self) -> HCQ2Buffer: return self._base or self
+
+class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
+  def _map(self, buf:HCQ2Buffer) -> HCQ2Buffer:
+    if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented")
+    return self._do_map(buf)
+
+  @suppress_finalizing
+  def _free(self, buf:HCQ2Buffer, options:BufferSpec|None=None):
+    if options is not None and options.external_ptr is not None: return
+    if hasattr(self, '_do_free'): self._do_free(buf, options)
+
+  def _unmap(self, mb):
+    self.dev.synchronize()
+    self.dev.iface.dev_impl.mm.unmap_range(int(mb.va_addr), round_up(mb.size, 0x1000))
+
+  def _offset(self, buf, size:int, offset:int) -> HCQ2Buffer: return buf.offset(offset=offset, size=size)
+
+  def _wrap(self, dev:str, sz:int, opaque:HCQ2Buffer) -> Buffer:
+    return Buffer(dev, sz, dtypes.uint8, opaque=opaque, options=BufferSpec(external_ptr=1))
+
+  def _copy(self, dst:Buffer, src:Buffer):
+    from tinygrad.engine.realize import run_linear
+    su = UOp.from_buffer(src)
+    run_linear(UOp(Ops.LINEAR, dtypes.void, (su.copy_to_device(dst.device).call(UOp.from_buffer(dst), su),)), update_stats=False)
+
+  def _copyin(self, dest:HCQ2Buffer, src:memoryview):
+    s = Buffer(self.dev.device, len(src), dtypes.uint8, options=BufferSpec(host=True), preallocate=True)
+    s._buf.cpu_view()[:len(src)] = src
+    self._copy(self._wrap(self.dev.device, len(src), dest), s)
+
+  def _copyout(self, dest:memoryview, src:HCQ2Buffer):
+    d = Buffer(self.dev.device, len(dest), dtypes.uint8, options=BufferSpec(host=True), preallocate=True)
+    self._copy(d, self._wrap(self.dev.device, len(dest), src))
+    self.dev.synchronize()
+    dest[:] = d._buf.cpu_view()[:len(dest)]
+
+  # def _as_buffer(self, buf): return buf.cpu_view().mv
+
+def unwrap_after(uop):
+  while uop.op is Ops.AFTER: uop = uop.src[0]
+  return uop
+
+class HCQEncoder:
+  def __init__(self): self.blob, self.patches = b'', []
+
+  def get_dev_addr(self, uop:UOp) -> UOp:
+    if unwrap_after(uop).op not in (Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT): return uop
+    return UOp(Ops.GETADDR, dtypes.uint64, src=(uop, UOp(Ops.DEVICE, arg=self.dev.device)))
+
+  def append(self, *data, dtype=dtypes.uint32):
+    for d in data:
+      if isinstance(d, int): self.blob += struct.pack(f'<{dtype.fmt}', d)
+      else:
+        self.patches.append((len(self.blob), self.get_dev_addr(d), dtype))
+        self.blob += struct.pack(f'<{dtype.fmt}', 0)
+
+  def q(self, *values): self.append(*values)
+
+  def uop(self, dev:str|tuple[str, ...], tag:str|None=None) -> UOp:
+    buf = UOp.new_buffer(dev, len(self.blob), dtypes.uint8)
+    if tag: buf = buf.rtag(tag)
+    blob_uop = UOp(Ops.BINARY, dtypes.void, src=(), arg=self.blob)
+    stores = [buf.index(UOp.const(dtypes.int, off), dtype=buf.dtype.ptr()).cast(dt.ptr()).store(val.cast(dt)) for off, val, dt in self.patches]
+    return buf.after(buf.store(blob_uop), *stores)
+
+# *****************
+# 0. helpers
+
+HCQ_DEVS = frozenset(("AMD",))
+HCQ_P2P_DEVS = HCQ_DEVS | frozenset(("CPU",))
+
+def to_tuple(d): return d if isinstance(d, tuple) else (d,)
+
+def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
+
+# *****************
+# 1.1. prep runtimes: staging copies
+
+def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS) and not all_devices_in(b.device, HCQ_P2P_DEVS)
+
+def stage_copy(dst:UOp, src:UOp) -> UOp|None:
+  if not (_need_staging(src, dst) or _need_staging(dst, src)): return None
+
+  stage = UOp.new_buffer("CPU", src.buffer.nbytes, dtypes.uint8)
+  return UOp(Ops.LINEAR, dtypes.void, (src.copy_to_device("CPU").call(stage, src), stage.copy_to_device(dst.device).call(dst, stage)))
+pm_insert_copy_staging = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy)])
+
+# *****************
+# 1.2. prep runtimes: programs/kernargs
+
+@functools.cache
+def get_pm_prep_program(name:str) -> PatternMatcher|None:
+  try:
+    importlib.import_module(f'tinygrad.runtime.ops_{name.lower()}') # TODO: remove that
+    return importlib.import_module(f'extra.hcq2.ops_{name.lower()}2').pm_prep_program
+  except ImportError: return None
+
+def prep_program(call:UOp, prg:UOp) -> UOp|None:
+  dev = call.src[1].device
+  if (pm:=get_pm_prep_program(to_tuple(dev)[0].split(":")[0])) is None or (lowered:=pm.rewrite(prg)) is None: return None
+
+  data, image_bytes = lowered
+  buf = UOp.new_buffer(dev, len(image_bytes), dtypes.uint8).rtag("program")
+  blob = UOp(Ops.BINARY, dtypes.void, src=(), arg=image_bytes)
+  return call.replace(src=(prg.replace(src=(buf.after(buf.store(blob)),), arg=(data, prg.arg)),) + call.src[1:])
+
+def prep_kernargs(call:UOp, prg:UOp) -> UOp:
+  data, info = prg.arg
+  patches = [(i*dtypes.uint64.itemsize, UOp(Ops.GETADDR, dtypes.uint64, src=(call.src[1+gi], UOp(Ops.DEVICE, arg=call.src[1+gi].device))),
+              dtypes.uint64) for i,gi in enumerate(info.globals)] \
+          + [(len(info.globals)*dtypes.uint64.itemsize + i*dtypes.uint32.itemsize, v, dtypes.uint32) for i,v in enumerate(info.vars)]
+
+  buf = UOp.new_buffer(call.src[1].device, data.kernargs_alloc_size, dtypes.uint8).rtag("kernargs")
+  kernargs = buf.after(*tuple(buf.index(UOp.const(dtypes.int, o), dtype=buf.dtype.ptr()).cast(dt.ptr()).store(val.cast(dt)) for o, val, dt in patches))
+
+  return call.replace(src=(prg.replace(src=prg.src + (kernargs,), arg=(data, info)),) + call.src[1:])
+
+pm_prep_runtime = PatternMatcher([
+  # bind generic PROGRAM device to the call's actual dev(s), then run device-specific lowering
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"),),
+    name="call", allow_any_len=True), prep_program),
+
+  # lower kernargs (PROGRAM.src[0] is now AFTER(BUFFER, COPY) — the lowered program image)
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.AFTER),), name="prg"),), name="call", allow_any_len=True), prep_kernargs),
+])
+
+# *****************
+# 2.1. lowering to hcq ir
+
+def lower_program(call:UOp, prg:UOp) -> UOp:
+  q = UOp(Ops.LINEAR, dtypes.void, (prg,), arg=(call.src[1].device, "COMPUTE"))
+  return call.replace(src=(q,) + call.src[1:]).rtag('hcq')
+
+def lower_copy(call:UOp, copy:UOp) -> UOp|None:
+  dst, src = call.src[1], call.src[2]
+  if (hcq_dev:=next((b.device for b in (dst, src) if b.device.split(":")[0] in HCQ_DEVS), None)) is None: return None
+
+  q = UOp(Ops.LINEAR, dtypes.void, (UOp(Ops.COPY, dtypes.void, src=(dst, src), arg=src.buffer.nbytes),), arg=(hcq_dev, "COPY"))
+  return call.replace(src=(q,) + call.src[1:]).rtag('hcq')
+
+pm_lower_ops = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.AFTER), UPat(Ops.AFTER)), name="prg"),), name="call", allow_any_len=True), lower_program),
+  (UPat(Ops.CALL, src=(UPat(Ops.COPY, name="copy"),), name="call", allow_any_len=True), lower_copy),
+])
+
+# *****************
+# 2.2. queue split
+
+# def split_into_queues(linear:UOp) -> UOp:
+#   out = []
+#   for k, grp in itertools.groupby(linear.src, lambda c: c.src[0].arg if c.op is Ops.CALL and c.src[0].op is Ops.LINEAR else None):
+#     if k is None: out.extend(grp)
+#     else:
+#       calls = list(grp)
+#       items = tuple(x for c in calls for x in c.src[0].src)
+#       args = tuple(a for c in calls for a in c.src[1:])
+#       out.append(calls[0].replace(src=(UOp(Ops.LINEAR, dtypes.void, items, arg=k),) + args))
+#   return linear.replace(src=tuple(out))
+# pm_split_into_queues = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), split_into_queues)])
+
+# *****************
+# 2.3. barriers / signals / timeline inc
+
+def add_barriers(call:UOp, q:UOp) -> UOp:
+  return call.replace(src=(q.replace(src=(UOp(Ops.BARRIER, dtypes.void), *q.src)),) + call.src[1:])
+pm_add_barriers = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.LINEAR, name="q"),), name="call", allow_any_len=True), add_barriers)])
+
+def add_signals(call:UOp, q:UOp) -> UOp:
+  sig = UOp.new_buffer(q.arg[0], 0x100, dtypes.uint8).rtag("timeline_signal")
+  tl = UOp.new_buffer(q.arg[0], 1, dtypes.uint64).rtag("timeline_value").index(UOp.const(dtypes.int, 0))
+  return call.replace(src=(q.replace(src=(sig.wait(tl-1), *q.src, sig.store(tl)), arg=q.arg),) + call.src[1:])
+pm_add_signals = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.LINEAR, name="q"),), name="call", allow_any_len=True), add_signals)])
+
+# *****************
+# 3.1. encode cmdbufs
+
+@functools.cache
+def get_pm_lower(name:str) -> PatternMatcher|None:
+  try:
+    importlib.import_module(f'tinygrad.runtime.ops_{name.lower()}') # TODO: remove that
+    return importlib.import_module(f'extra.hcq2.ops_{name.lower()}2').pm_lower
+  except ImportError: return None
+
+def encode_cmdbuf(call:UOp, q:UOp) -> UOp|None:
+  if (pm:=get_pm_lower(to_tuple(q.arg[0])[0].split(":")[0])) is None or (encoded:=pm.rewrite(q)) is None: return None
+  return call.replace(src=(encoded,) + call.src[1:])
+pm_encode_cmdbufs = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.LINEAR, name="q"),), name="call", allow_any_len=True), encode_cmdbuf)])
+
+# *****************
+# 3.2. add timeline inc
+
+def add_timeline_inc(call:UOp, s:UOp) -> UOp:
+  tl = UOp.new_buffer(s.device, 1, dtypes.uint64).rtag("timeline_value")
+  return call.replace(src=(tl.after(s).index(UOp.const(dtypes.int, 0), dtype=tl.dtype.ptr()).store(tl.index(UOp.const(dtypes.int, 0)) + 1),) + call.src[1:])
+pm_add_timeline_inc = PatternMatcher([(UPat(Ops.CALL, tag="hcq", src=(UPat(name="s"),), name="call", allow_any_len=True), add_timeline_inc)])
+
+# *****************
+# 3.3. lift patches to the command buffer (root)
+
+def lift_patches_to_cmdbuf(cmdbuf:UOp) -> UOp|None:
+  if not (patches:=dedup(u for store in cmdbuf.src[1:] for u in store.toposort() if u.op is Ops.AFTER)): return None
+  deps = tuple(d for p in patches for d in p.src[1:])
+  return cmdbuf.replace(src=cmdbuf.src + deps).substitute({p: p.src[0] for p in patches})
+pm_lift_patches_to_cmdbuf = PatternMatcher([
+  (UPat(Ops.AFTER, src=(UPat(Ops.BUFFER, tag={"compute", "copy"}),), allow_any_len=True, name="cmdbuf"), lift_patches_to_cmdbuf),
+])
+
+# *****************
+# 4. bufferize placeholders: replace placeholders with real buffers.
+
+def bufferize_buf(buf:UOp) -> UOp|None:
+  if buf.tag is None: return None
+  uops = tuple(UOp.from_buffer((dv:=Device[dev]).pm_bufferize.rewrite(buf, ctx=dv), dev) for dev in to_tuple(buf.src[1].arg))
+  return uops[0] if len(uops) == 1 else UOp(Ops.MSTACK, uops[0].dtype, uops)
+pm_bufferize = PatternMatcher([(UPat(Ops.BUFFER, name="buf"), bufferize_buf)])
+
+# *****************
+# 5.1. capture buffers reachable from each hcq call as BIND, so we don't drop their refs
+
+def hold_call_buffers(call:UOp) -> UOp|None:
+  if not (bufs:=tuple(dedup(u for u in call.src[0].toposort() if u.op is Ops.BUFFER and u not in call.src))): return None
+  return call.replace(src=call.src + (UOp(Ops.BIND, dtypes.void, src=bufs),))
+pm_hold_call_buffers = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), hold_call_buffers)])
+
+# *****************
+# 5.2. resolve patches
+
+def push_stack(op, s): return UOp(Ops.STACK, op.dtype.scalar().vec(len(s.src)),
+  tuple(op.replace(dtype=op.dtype.scalar(), src=tuple(x if y is s else y for y in op.src)) for x in s.src))
+
+def fold_blob_store(buf:UOp, blob:UOp) -> UOp:
+  for b in (buf.src if buf.op is Ops.MSTACK else (buf,)): b.buffer.ensure_allocated()._buf.cpu_view().mv.cast('B')[:len(blob.arg)] = blob.arg
+  return UOp(Ops.NOOP)
+
+def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp:
+  for b, v in zip((buf.src if buf.op is Ops.MSTACK else (buf,)), (val.src if val.op is Ops.STACK else (val,))):
+    struct.pack_into(f'<{v.dtype.fmt}', b.buffer.ensure_allocated()._buf.cpu_view().mv.cast('B'), off.arg * b.dtype.base.itemsize, v.arg)
+  return UOp(Ops.NOOP)
+
+def resolve_getaddr(buf:UOp, g:UOp) -> UOp:
+  if isinstance(b:=buf.buffer, Buffer): return UOp.const(dtypes.uint64, b.get_buf(g.src[1].arg).va_addr)
+  return UOp(Ops.STACK, dtypes.uint64.vec(len(b.bufs)), tuple(UOp.const(dtypes.uint64, x.ensure_allocated()._buf.va_addr) for x in b.bufs))
+
+pm_resolve_patches = PatternMatcher([
+  # multi
+  (UPat(GroupOp.ALU, src=[UPat(Ops.STACK, name="s"), UPat(Ops.CONST)], name="op"), push_stack),
+  (UPat(Ops.CAST, src=(UPat(Ops.STACK, name="s"),), name="op"), push_stack),
+
+  # getaddr
+  (UPat(Ops.GETADDR, src=(UPat(Ops.SLICE, name="bv"), UPat(Ops.DEVICE, name="dev"))), # getaddr(slice(x)) -> offset+getaddr(x)
+    lambda bv, dev: UOp(Ops.GETADDR, dtypes.uint64, src=(bv.src[0], dev)) + UOp.const(dtypes.uint64, bv.src[1].arg * bv.src[0].dtype.itemsize)),
+  (UPat(Ops.GETADDR, src=(UPat({Ops.BUFFER, Ops.MSTACK, Ops.MSELECT}, name="buf"), UPat(Ops.DEVICE)), name="g"), resolve_getaddr),
+
+  # folders
+  (UPat({Ops.BUFFER, Ops.MSTACK}, name="buf").store(UPat(Ops.BINARY, name="blob")), fold_blob_store),
+  (UPat({Ops.BUFFER, Ops.MSTACK}, name="buf").index(UPat.cvar("off")).or_casted().store(UPat.any(UPat.cvar("val"), UPat(Ops.STACK, name="val"))),
+    fold_const_store),
+]) + symbolic_simple
+
+# *****************
+# 6. callify hcq programs
+
+pm_fixup = PatternMatcher([ # TODO: this should gone?
+  (UPat(Ops.CONST, name="c"), lambda c: c.replace(src=()) if len(c.src) else None),
+])
+
+def to_param(bufs:list[UOp], ref:UOp) -> UOp:
+  bufs.append(ref)
+  return UOp.placeholder((ref.buffer.size,), ref.dtype, len(bufs)-1)
+pm_to_param = PatternMatcher([(UPat({Ops.MSELECT, Ops.MSTACK, Ops.BUFFER}, name="r"), lambda ctx, r: to_param(ctx, r))])
+
+def parametrize_host_buffers(call:UOp) -> UOp:
+  body = graph_rewrite(call.src[0], pm_to_param, ctx=(bufs:=[]), bottom_up=True, name="parametrize host buffers")
+  return call.replace(src=(body, *bufs) + call.src[1:], tag="hcq_param")
+pm_parametrize_host_buffers = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), parametrize_host_buffers)])
+
+def callify_hcq(call:UOp) -> UOp:
+  sink = UOp.sink(call.src[0], arg=KernelInfo(name="hcq_submit", estimates=Estimates()), tag=1)
+  return to_program(sink, Device["CPU"].renderer).call(*call.src[1:])
+pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, tag="hcq_param", name="call"), callify_hcq)])
+
+@track_rewrites(lambda _,ret: f"HCQ Schedule {pluralize('Kernel', len(ret.src))}")
+def hcq_schedule(linear:UOp) -> UOp:
+  linear = graph_rewrite(linear, pm_insert_copy_staging + pm_flatten_linear, name="insert copy staging")
+  linear = graph_rewrite(linear, pm_prep_runtime, name="prepare runtime")
+
+  linear = graph_rewrite(linear, pm_lower_ops, name="lower ops into hcq ir")
+  # linear = graph_rewrite(linear, pm_split_into_queues, name="split into queues")
+  linear = graph_rewrite(linear, pm_add_barriers, walk=True, name="add barriers")
+  linear = graph_rewrite(linear, pm_add_signals, walk=True, name="add signals")
+  linear = graph_rewrite(linear, pm_encode_cmdbufs, walk=True, name="encode cmdbufs")
+  linear = graph_rewrite(linear, pm_add_timeline_inc, walk=True, name="add timeline inc")
+  linear = graph_rewrite(linear, pm_lift_patches_to_cmdbuf, name="lift patches to cmdbuf", enter_calls=True)
+
+  # realize starts from here
+  linear = graph_rewrite(linear, pm_bufferize, bottom_up=True, name="bufferize placeholders", enter_calls=True)
+  linear = graph_rewrite(linear, pm_hold_call_buffers, walk=True, name="hold call buffers")
+  linear = graph_rewrite(linear, pm_resolve_patches, bottom_up=False, name="simplify patches", enter_calls=True)
+  linear = graph_rewrite(linear, pm_fixup, bottom_up=False, name="fixup", enter_calls=True)
+  linear = graph_rewrite(linear, pm_parametrize_host_buffers, name="parametrize host buffers")
+  linear = graph_rewrite(linear, pm_callify_hcq, name="callify hcq")
+
+  return linear

tinygrad_repo/extra/hcq2/ops_amd2.py

@@ -0,0 +1,551 @@
+from __future__ import annotations
+from typing import cast
+import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit
+assert sys.platform != 'win32'
+from dataclasses import dataclass
+from extra.hcq2.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, HCQEncoder
+from tinygrad.uop.ops import sint, UOp
+from tinygrad.device import Compiled, BufferSpec, Buffer, Device
+from tinygrad.dtype import dtypes
+from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, ContextVar, TracingKey
+from tinygrad.helpers import VIZ, ceildiv, unwrap, pluralize
+from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer
+from tinygrad.renderer.llvmir import AMDLLVMRenderer
+from tinygrad.runtime.autogen import kfd, hsa, sqtt, amdgpu_kd, amdgpu_drm
+from tinygrad.runtime.autogen.am import am
+from tinygrad.runtime.support.elf import elf_loader
+from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
+from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc
+from tinygrad.runtime.support.system import PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
+from tinygrad.runtime.support.usb import USB3
+from tinygrad.runtime.support.memory import AddrSpace, BumpAllocator
+from tinygrad.runtime.ops_amd import SQTT, SQTT_ITRACE_SE_MASK, SQTT_LIMIT_SE, SQTT_SIMD_SEL, SQTT_TOKEN_EXCLUDE, PMC
+from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_EQ, WAIT_REG_MEM_FUNCTION_NEQ, WAIT_REG_MEM_FUNCTION_GEQ
+if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
+
+from tinygrad.engine.realize import get_runtime
+from tinygrad.uop.ops import Ops, UPat, PatternMatcher, graph_rewrite
+
+class AMDComputeQueue(HCQEncoder):
+  def __init__(self, dev:AMDDevice, devs:tuple[str, ...]|None=None):
+    super().__init__()
+    self.dev, self.devs = dev, devs or (dev.device,)
+    self.pm4, self.gc, self.nbio, self.soc = dev.pm4, dev.gc, dev.nbio, dev.soc
+
+  def pkt3(self, cmd, *vals): self.q(self.pm4.PACKET3(cmd, len(vals) - 1), *vals)
+
+  def wreg(self, reg:AMDReg, *args:sint, **kwargs:int):
+    if bool(args) == bool(kwargs): raise RuntimeError('One (and only one) of *args or **kwargs must be specified')
+    if self.pm4.PACKET3_SET_SH_REG_START <= reg.addr[0] < self.pm4.PACKET3_SET_SH_REG_END:
+      set_packet, set_packet_start = self.pm4.PACKET3_SET_SH_REG, self.pm4.PACKET3_SET_SH_REG_START
+    elif self.pm4.PACKET3_SET_UCONFIG_REG_START <= reg.addr[0] < self.pm4.PACKET3_SET_UCONFIG_REG_START + 2**16-1:
+      set_packet, set_packet_start = self.pm4.PACKET3_SET_UCONFIG_REG, self.pm4.PACKET3_SET_UCONFIG_REG_START
+    else: raise RuntimeError(f'Cannot set {reg.name} ({reg.addr[0]}) via pm4 packet')
+    self.pkt3(set_packet, reg.addr[0] - set_packet_start, *(args or (reg.encode(**kwargs),)))
+
+  def wait_reg_mem(self, value, mask=0xffffffff, mem=None, reg=None, reg_done=0, op=WAIT_REG_MEM_FUNCTION_GEQ):
+    wrm_info_dw = self.pm4.WAIT_REG_MEM_MEM_SPACE(int(mem is not None)) | self.pm4.WAIT_REG_MEM_OPERATION(int(mem is None and reg_done > 0)) \
+                | self.pm4.WAIT_REG_MEM_FUNCTION(op) | self.pm4.WAIT_REG_MEM_ENGINE(0)
+    self.pkt3(self.pm4.PACKET3_WAIT_REG_MEM, wrm_info_dw, *(data64_le(mem) if mem is not None else (reg, reg_done)), value, mask, 4)
+
+  def acquire_mem(self, addr=0x0, sz=(1 << 64)-1, gli=1, glm=1, glk=1, glv=1, gl1=1, gl2=1):
+    if self.dev.target[0] != 9:
+      cache_flags_dw = self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLI_INV(gli) \
+                     | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLM_INV(glm) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLM_WB(glm) \
+                     | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLK_INV(glk) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLK_WB(glk) \
+                     | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLV_INV(glv) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL1_INV(gl1) \
+                     | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL2_INV(gl2) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL2_WB(gl2)
+      self.pkt3(self.pm4.PACKET3_ACQUIRE_MEM, 0, *data64_le(sz), *data64_le(addr), 0, cache_flags_dw)
+    else:
+      cp_coher_cntl = self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_SH_ICACHE_ACTION_ENA(gli) | \
+                      self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_SH_KCACHE_ACTION_ENA(glk) | \
+                      self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TC_ACTION_ENA(gl2) | \
+                      self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TCL1_ACTION_ENA(gl1) | \
+                      self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TC_WB_ACTION_ENA(gl2)
+      self.pkt3(self.pm4.PACKET3_ACQUIRE_MEM, cp_coher_cntl, *data64_le(sz), *data64_le(addr), 0x0000000A)
+
+  def release_mem(self, address=0x0, value=0, data_sel=0, int_sel=2, ctxid=0, cache_flush=False):
+    if self.dev.target[0] != 9:
+      cache_flags_dw = 0 if not cache_flush else (self.pm4.PACKET3_RELEASE_MEM_GCR_GLV_INV | self.pm4.PACKET3_RELEASE_MEM_GCR_GL1_INV \
+                     | self.pm4.PACKET3_RELEASE_MEM_GCR_GL2_INV | self.pm4.PACKET3_RELEASE_MEM_GCR_GLM_WB \
+                     | self.pm4.PACKET3_RELEASE_MEM_GCR_GLM_INV | self.pm4.PACKET3_RELEASE_MEM_GCR_GL2_WB | self.pm4.PACKET3_RELEASE_MEM_GCR_SEQ)
+      event_dw = self.pm4.PACKET3_RELEASE_MEM_EVENT_TYPE(self.pm4.CACHE_FLUSH_AND_INV_TS_EVENT) \
+               | self.pm4.PACKET3_RELEASE_MEM_EVENT_INDEX(self.pm4.event_index__mec_release_mem__end_of_pipe)
+      memsel_dw = self.pm4.PACKET3_RELEASE_MEM_DATA_SEL(data_sel) | self.pm4.PACKET3_RELEASE_MEM_INT_SEL(int_sel) \
+                | self.pm4.PACKET3_RELEASE_MEM_DST_SEL(0)
+    else:
+      cache_flags_dw = 0 if not cache_flush else (self.pm4.EOP_TC_WB_ACTION_EN | self.pm4.EOP_TC_NC_ACTION_EN)
+      event_dw = self.pm4.EVENT_TYPE(self.pm4.CACHE_FLUSH_AND_INV_TS_EVENT) | self.pm4.EVENT_INDEX(self.pm4.event_index__mec_release_mem__end_of_pipe)
+      memsel_dw = self.pm4.DATA_SEL(data_sel) | self.pm4.INT_SEL(int_sel)
+      ctxid = 0
+    self.pkt3(self.pm4.PACKET3_RELEASE_MEM, event_dw | cache_flags_dw, memsel_dw, *data64_le(address), *data64_le(value), ctxid)
+
+  def memory_barrier(self):
+    pf = '' if self.nbio.version[0] == 2 else '0' if self.nbio.version[:2] != (7, 11) else '1'
+    self.wait_reg_mem(reg=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_REQ').addr[0],
+                      reg_done=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff)
+    self.acquire_mem()
+
+  def wait(self, x): self.wait_reg_mem(x.src[1], mem=self.get_dev_addr(x.src[0]))
+
+  def barrier(self, x): self.memory_barrier()
+
+  def store(self, x):
+    self.release_mem(self.get_dev_addr(x.src[0]), x.src[1], self.pm4.data_sel__mec_release_mem__send_32_bit_low,
+                     self.pm4.int_sel__mec_release_mem__send_interrupt_after_write_confirm, cache_flush=True)
+
+  def timestamp(self, x):
+    self.release_mem(self.get_dev_addr(x.src[0]), 0, self.pm4.data_sel__mec_release_mem__send_gpu_clock_counter,
+                     self.pm4.int_sel__mec_release_mem__none)
+
+  def program(self, x):
+    data, info = x.arg
+    lib_gpu, args = x.src
+    prog_addr = self.get_dev_addr(lib_gpu) + data.entry_point_offset
+
+    self.acquire_mem(gli=0, gl2=0)
+
+    scratch_addr = self.get_dev_addr(UOp.new_buffer(self.devs, data.private_segment_size, dtypes.uint8).rtag("scratch"))
+    args_addr = self.get_dev_addr(args)
+
+    user_regs = []
+    if data.enable_private_segment_sgpr:
+      scratch_hilo = data64_le(scratch_addr)
+      user_regs = [scratch_hilo[0], scratch_hilo[1] | 1 << 31, 0xffffffff, 0x20c14000]
+    if data.enable_dispatch_ptr: user_regs += [*data64_le(args_addr + data.kernargs_segment_size)]
+    user_regs += [*data64_le(args_addr)]
+
+    self.wreg(self.gc.regCOMPUTE_PGM_LO, *data64_le(prog_addr >> 8))
+    self.wreg(self.gc.regCOMPUTE_PGM_RSRC1, data.rsrc1, data.rsrc2)
+    self.wreg(self.gc.regCOMPUTE_PGM_RSRC3, data.rsrc3)
+    self.wreg(self.gc.regCOMPUTE_TMPRING_SIZE, self.dev.tmpring_size(data.private_segment_size))
+
+    for xcc_id in range(self.dev.xccs):
+      scratch_base = scratch_addr + (data.private_segment_size // self.dev.xccs * xcc_id)
+      self.wreg(self.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, *data64_le(scratch_base >> 8))
+
+    self.wreg(self.gc.regCOMPUTE_RESTART_X, 0, 0, 0)
+    self.wreg(self.gc.regCOMPUTE_USER_DATA_0, *user_regs)
+    self.wreg(self.gc.regCOMPUTE_RESOURCE_LIMITS, self.gc.regCOMPUTE_RESOURCE_LIMITS.encode(waves_per_sh=getenv("WAVES_PER_SH")))
+    self.wreg(self.gc.regCOMPUTE_START_X, 0, 0, 0, *(info.local_size or (1, 1, 1)), 0, 0)
+
+    dispatch_init = self.gc.regCOMPUTE_DISPATCH_INITIATOR.encode(
+      **({'cs_w32_en': int(data.wave32)} if self.dev.target[0] != 9 else {}), force_start_at_000=1, compute_shader_en=1)
+    self.pkt3(self.pm4.PACKET3_DISPATCH_DIRECT, *info.global_size, dispatch_init)
+    self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.CS_PARTIAL_FLUSH) | self.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))
+
+amd_inner_pm = PatternMatcher([
+  (UPat(Ops.LINEAR, src=(UPat(Ops.WAIT, name="x"),)),    lambda ctx, x: ctx.wait(x)),
+  (UPat(Ops.LINEAR, src=(UPat(Ops.BARRIER, name="x"),)), lambda ctx, x: ctx.barrier(x)),
+  (UPat(Ops.LINEAR, src=(UPat(Ops.PROGRAM, name="x"),)), lambda ctx, x: ctx.program(x)),
+  (UPat(Ops.LINEAR, src=(UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", name="x"),)), lambda ctx, x: ctx.timestamp(x)),
+  (UPat(Ops.LINEAR, src=(UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM)), UPat()), name="x"),)), lambda ctx, x: ctx.store(x)),
+])
+
+def amd_lower_pm4(linear, devs):
+  enc = AMDComputeQueue(Device[devs[0]], devs)
+  graph_rewrite(linear.replace(src=tuple(UOp(Ops.LINEAR, dtypes.void, (cmd,)) for cmd in linear.src)), amd_inner_pm, ctx=enc, name="amd: encode")
+  return enc.uop(dev=devs if len(devs) > 1 else devs[0], tag="compute")
+
+def amd_submit_pm4(cmdbuf, devs):
+  size, zero = UOp.const(dtypes.uint32, cmdbuf.src[0].arg // dtypes.uint32.itemsize), UOp.const(dtypes.int, 0)
+
+  # the compute queue's ring and its host-side ring/write/put pointers (placeholders, resolved in pm_bufferize)
+  q = Device['AMD'].compute_queue
+  ring, wptr, doorbell, put_ptr = (UOp.new_buffer(devs, b.size, b.dtype).rtag(("COMPUTE:0", name))
+    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
+
+  # place the cmdbuf at the ring's write offset, wrapping the ring
+  put = put_ptr.index(zero)
+  next_put = put + size.cast(put.dtype)
+  i = UOp.range(size, 0, dtype=dtypes.int, src=(cmdbuf,))
+  ring_idx = ((put + i.cast(put.dtype)) % q.ring.size).cast(dtypes.int)
+
+  # copy the cmdbuf into the ring and advance the put/write pointers
+  copy_to_ring = ring.index(ring_idx, dtype=ring.dtype.ptr()).store(
+    cmdbuf.index(i*4, dtype=cmdbuf.dtype.ptr()).cast(dtypes.uint32.ptr()).load()).end(i)
+  bump_put_ptr = put_ptr.index(zero, dtype=put_ptr.dtype.ptr()).store(next_put)
+  bump_wptr = wptr.index(zero, dtype=wptr.dtype.ptr()).store(next_put)
+
+  # ring the doorbell once the copy and pointer bumps have landed
+  flush = UOp.barrier(copy_to_ring, bump_put_ptr, bump_wptr)
+  return doorbell.after(flush).index(zero, dtype=doorbell.dtype.ptr()).store(next_put)
+
+class AMDCopyQueue(HCQEncoder):
+  def __init__(self, dev:AMDDevice, queue_idx=0):
+    super().__init__()
+    self.dev = dev
+    self.sdma, self.queue_idx, self.max_copy_size = dev.sdma, queue_idx, dev.max_copy_size
+
+  def copy(self, x):
+    dest, src, copy_size = self.get_dev_addr(x.src[0]), self.get_dev_addr(x.src[1]), x.arg
+    copied = 0
+    while copied < copy_size:
+      step = min(copy_size - copied, self.max_copy_size)
+      self.q(self.sdma.SDMA_OP_COPY | self.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(self.sdma.SDMA_SUBOP_COPY_LINEAR),
+             self.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(step - 1), 0, *data64_le(src + copied), *data64_le(dest + copied))
+      copied += step
+
+  def wait(self, x):
+    self.q(self.sdma.SDMA_OP_POLL_REGMEM | self.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) | \
+           self.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1), *data64_le(self.get_dev_addr(x.src[0])), x.src[1], 0xffffffff,
+           self.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | self.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))
+
+  def store(self, x):
+    fence_flags = self.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if self.dev.target[0] != 9 else 0
+    self.q(self.sdma.SDMA_OP_FENCE | fence_flags, *data64_le(self.get_dev_addr(x.src[0])), x.src[1])
+    self.q(self.sdma.SDMA_OP_TRAP, 0)
+
+  def timestamp(self, x):
+    self.q(self.sdma.SDMA_OP_TIMESTAMP | self.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(self.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL),
+           *data64_le(self.get_dev_addr(x.src[0])))
+
+def amd_lower_sdma(linear, devs):
+  enc = AMDCopyQueue(Device[devs[0]])
+  graph_rewrite(linear.replace(src=tuple(UOp(Ops.LINEAR, dtypes.void, (cmd,)) for cmd in linear.src)), amd_inner_sdma_pm, ctx=enc, name="amd: encode sdma")
+  return enc.uop(dev=devs if len(devs) > 1 else devs[0], tag="copy")
+
+amd_inner_sdma_pm = PatternMatcher([
+  (UPat(Ops.LINEAR, src=(UPat(Ops.WAIT, name="x"),)), lambda ctx, x: ctx.wait(x)),
+  (UPat(Ops.LINEAR, src=(UPat(Ops.BARRIER, name="x"),)), lambda ctx, x: None),
+  (UPat(Ops.LINEAR, src=(UPat(Ops.COPY, name="x"),)), lambda ctx, x: ctx.copy(x)),
+  (UPat(Ops.LINEAR, src=(UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", name="x"),)), lambda ctx, x: ctx.timestamp(x)),
+  (UPat(Ops.LINEAR, src=(UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM)), UPat()), name="x"),)), lambda ctx, x: ctx.store(x)),
+])
+
+def amd_submit_sdma(cmdbuf, devs):
+  # the cmdbuf to submit + the patch writes that fill it
+  size_dw, zero = cmdbuf.src[0].arg // dtypes.uint32.itemsize, UOp.const(dtypes.int, 0)
+
+  # the sdma queue's ring and its host-side ring/write/put pointers
+  q = Device['AMD'].sdma_queue(0)
+  ring, wptr, doorbell, put_ptr = (UOp.new_buffer(devs, b.size, b.dtype).rtag(("SDMA:0", name))
+    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
+
+  # sdma needs the cmdbuf contiguous: if it won't fit before the ring end, restart at 0 and zero the tail
+  put_b = put_ptr.index(zero)
+  tail_off_dw = ((put_b % (q.ring.size * 4)) // 4).cast(dtypes.int)
+  fits = (size_dw <= q.ring.size - tail_off_dw).cast(dtypes.int)
+  start_dw = fits * tail_off_dw
+  zero_amt_dw = (1 - fits) * (q.ring.size - tail_off_dw)
+
+  # zero the wrapped tail, then copy the cmdbuf into the ring
+  zi = UOp.range(zero_amt_dw, 0, dtype=dtypes.int, src=(cmdbuf,))
+  zero_tail = ring.index(tail_off_dw + zi, dtype=ring.dtype.ptr()).store(UOp.const(dtypes.uint32, 0)).end(zi)
+  i = UOp.range(UOp.const(dtypes.int, size_dw), 0, dtype=dtypes.int, src=(cmdbuf,))
+  copy_to_ring = ring.index(start_dw + i, dtype=ring.dtype.ptr()).store(
+    cmdbuf.index(i*4, dtype=cmdbuf.dtype.ptr()).cast(dtypes.uint32.ptr()).load()).end(i)
+
+  # advance the put/write pointers past the zeroed tail and the cmdbuf
+  next_put_b = put_b + ((zero_amt_dw + size_dw) * 4).cast(put_b.dtype)
+  bump_put_ptr = put_ptr.index(zero, dtype=put_ptr.dtype.ptr()).store(next_put_b)
+  bump_wptr = wptr.index(zero, dtype=wptr.dtype.ptr()).store(next_put_b)
+
+  # ring the doorbell once the writes have landed
+  flush = UOp.barrier(zero_tail, copy_to_ring, bump_put_ptr, bump_wptr)
+  return doorbell.after(flush).index(zero, dtype=doorbell.dtype.ptr()).store(next_put_b)
+
+@dataclass(frozen=True)
+class AMDProgramData:
+  entry_point_offset:int; rsrc1:int; rsrc2:int; rsrc3:int; wave32:bool
+  private_segment_size:int; kernargs_segment_size:int; kernargs_alloc_size:int
+  enable_dispatch_ptr:int; enable_private_segment_sgpr:int
+
+_amd_program_cache:dict[tuple[bytes,str], tuple[AMDProgramData,bytes]] = {}
+
+def amd_build_program(prg:UOp) -> UOp:
+  dev = Device[prg.src[1].arg] # TODO: rm this
+  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[4].arg, dev.device))) is None:
+    image, sections, relocs = elf_loader(lib)
+    rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata")
+    for off, sym, typ, addent in relocs:
+      assert typ == 5, f"unknown AMD reloc {typ}"  # R_AMDGPU_REL64
+      image[off:off+8] = struct.pack('<q', sym - off + addent)
+    desc = amdgpu_kd.llvm_amdhsa_kernel_descriptor_t.from_buffer_copy(bytes(image[rodata:rodata+ctypes.sizeof(amdgpu_kd.llvm_amdhsa_kernel_descriptor_t)]))
+    if (lds:=((desc.group_segment_fixed_size+511)//512)&0x1FF) > (dev.iface.props['lds_size_in_kb']*1024)//512:
+      raise RuntimeError("Too many resources requested: group_segment_size")
+    edp = desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_DISPATCH_PTR
+    cached = _amd_program_cache[key] = (AMDProgramData(
+      entry_point_offset=rodata + desc.kernel_code_entry_byte_offset,
+      rsrc1=desc.compute_pgm_rsrc1 | ((1<<20) if dev.target[0]==11 else 0),  # priv=1 on gfx11 for cwsr
+      rsrc2=desc.compute_pgm_rsrc2 | (lds<<15), rsrc3=desc.compute_pgm_rsrc3,
+      wave32=bool(desc.kernel_code_properties & 0x400),
+      private_segment_size=desc.private_segment_fixed_size,
+      kernargs_segment_size=desc.kernarg_size,
+      kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0),
+      enable_dispatch_ptr=edp,
+      enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER), bytes(image))
+  return cached
+
+pm_prep_program = PatternMatcher([
+  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE, arg="AMD"), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program),
+])
+
+class AMDAllocator(HCQAllocator['AMDDevice']):
+  def __init__(self, dev:AMDDevice):
+    super().__init__(dev, supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_queue and not dev.is_usb())
+
+  def _alloc(self, size:int, options:BufferSpec) -> HCQ2Buffer:
+    return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_sdma_queue)
+
+  def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque)
+
+  def _do_map(self, buf:HCQ2Buffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+
+@dataclass
+class AMDQueueDesc:
+  ring: Buffer        # uint32[ring_size//4]
+  read_ptr: Buffer    # uint64[1]
+  write_ptr: Buffer   # uint64[1]
+  doorbell: Buffer    # uint64[1]
+  put_value: Buffer   # uint64[1]
+  params: tuple|None = None  # setup_ring params for recovery
+
+class PCIIface(PCIIfaceBase):
+  def __init__(self, dev, dev_id):
+    super().__init__(dev, dev_id, vendor=0x1002, devices=((0xffff, (0x74a1,0x744c,0x7480,0x7550,0x7551,0x7590,0x75a0)),), vram_bar=0,
+      va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size, dev_impl_t=AMDev)
+    self._compute_props()
+
+  def p2p_paddrs(self, paddrs:list[tuple[int,int]]) -> tuple[list[tuple[int,int]], AddrSpace]:
+    return ([(self.dev_impl.paddr2xgmi(p), sz) for p, sz in paddrs], AddrSpace.PEER) if self.dev_impl.is_hive() else super().p2p_paddrs(paddrs)
+
+  def require_profile_mode(self): return True
+  def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return True # TODO: account for WGP disablement on some asics.
+
+  def _compute_props(self):
+    self.ip_versions = self.dev_impl.ip_ver
+
+    gfxver = int(f"{self.dev_impl.ip_ver[am.GC_HWIP][0]:02d}{self.dev_impl.ip_ver[am.GC_HWIP][1]:02d}{self.dev_impl.ip_ver[am.GC_HWIP][2]:02d}")
+    if self.dev_impl.gc_info.header.version_major == 2:
+      cu_per_sa = self.dev_impl.gc_info.gc_num_cu_per_sh
+      max_sh_per_se = self.dev_impl.gc_info.gc_num_sh_per_se
+    else:
+      cu_per_sa = 2 * (self.dev_impl.gc_info.gc_num_wgp0_per_sa + self.dev_impl.gc_info.gc_num_wgp1_per_sa)
+      max_sh_per_se = self.dev_impl.gc_info.gc_num_sa_per_se
+
+    array_count = max_sh_per_se * self.dev_impl.gc_info.gc_num_se * self.dev_impl.gfx.xccs
+    self.props = {'cu_per_simd_array': cu_per_sa, 'simd_count': 2 * cu_per_sa * array_count, 'simd_per_cu': 2, 'array_count': array_count,
+      'max_slots_scratch_cu': self.dev_impl.gc_info.gc_max_scratch_slots_per_cu, 'max_waves_per_simd': self.dev_impl.gc_info.gc_max_waves_per_simd,
+      'simd_arrays_per_engine': max_sh_per_se, 'lds_size_in_kb': self.dev_impl.gc_info.gc_lds_size, 'num_xcc': self.dev_impl.gfx.xccs,
+      'gfx_target_version': {90403: 90402}.get(gfxver, gfxver)}
+
+  def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
+                   xcc_id=0, idx=0):
+    assert cwsr_buffer is None, "no cwsr buffer for am"
+
+    rcvr_params: tuple
+    if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
+      doorbell_index = self.dev_impl.sdma.setup_ring(*(rcvr_params:=(ring.va_addr, ring.size, gart.va_addr+rptr, gart.va_addr+wptr, idx)))
+    else:
+      doorbell_index = self.dev_impl.gfx.setup_ring(*(rcvr_params:=(ring.va_addr, ring.size, gart.va_addr+rptr, gart.va_addr+wptr,
+        eop_buffer.va_addr, eop_buffer.size, is_aql:=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL), is_aql)))
+
+    ext = lambda addr,n,dt: Buffer("CPU", n, dt, options=BufferSpec(external_ptr=addr), preallocate=True)
+    (put_value := Buffer("CPU", 1, dtypes.uint64, preallocate=True))._buf.view.view(fmt='Q')[0] = 0
+    return AMDQueueDesc(ring=ext(ring.va_addr, ring.size//4, dtypes.uint32),
+      doorbell=ext(self.dev_impl.doorbell64.addr + doorbell_index*8, 1, dtypes.uint64),
+      read_ptr=ext(gart.va_addr+rptr, 1, dtypes.uint64), write_ptr=ext(gart.va_addr+wptr, 1, dtypes.uint64),
+      put_value=put_value, params=rcvr_params)
+
+  def _collect_interrupts(self, reset=False, drain_only=False):
+    d = self.dev
+    if drain_only: d.iface.dev_impl.ih.drain()
+    else: d.iface.dev_impl.ih.interrupt_handler()
+
+    if reset and d.iface.dev_impl.recover():
+      cq = d.compute_queue
+      for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0
+      d.iface.dev_impl.gfx.setup_ring(*cq.params)
+      d.timeline_signal._buf.cpu_view().mv.cast('Q')[0] = d.timeline_value.as_memoryview(force_zero_copy=True).cast('Q')[0] - 1
+
+  def sleep(self, timeout):
+    if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
+      self.pci_dev.irq_fd.read(8 * events_cnt)
+    self._collect_interrupts()
+    if self.dev_impl.is_err_state: raise RuntimeError("Device is in error state")
+
+  def on_device_hang(self):
+    self._collect_interrupts(reset=True)
+    raise RuntimeError("Device hang detected")
+
+  def device_fini(self): self.dev_impl.fini()
+
+def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {})
+
+def encode_queue(q:UOp) -> UOp|None:
+  if not (isinstance(q.arg, tuple) and len(q.arg) == 2 and q.arg[1] in ("COMPUTE", "COPY")): return None
+  devs = (q.arg[0],) if isinstance(q.arg[0], str) else q.arg[0]  # TODO: make this prettier
+  return amd_submit_pm4(amd_lower_pm4(q, devs), devs) if q.arg[1] == "COMPUTE" else amd_submit_sdma(amd_lower_sdma(q, devs), devs)
+
+pm_lower = PatternMatcher([
+  (UPat(Ops.LINEAR, name="q"), encode_queue),
+])
+
+class AMDDevice(HCQ2Compiled):
+  timestamp_divider = 100.0  # AMD GPU clock: ticks/us
+
+  ifaces = [PCIIface]
+
+  def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,))
+  def is_usb(self) -> bool: return False
+
+  def __init__(self, device:str=""):
+    self.device_id = int(device.split(":")[1]) if ":" in device else 0
+
+    self.iface = self._select_iface()
+
+    self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
+    self.arch = "gfx%d%x%x" % self.target
+    assert (self.target in ((9,4,2),(9,5,0))) or self.target[0] in (11, 12), f"Unsupported arch: {self.arch}"
+    if DEBUG >= 1: print(f"AMDDevice: opening {self.device_id} with target {self.target} arch {self.arch}")
+
+    self.xccs = self.iface.props.get('num_xcc', 1)
+    self.se_cnt = self.iface.props['array_count'] // self.iface.props['simd_arrays_per_engine'] // self.xccs
+    self.cu_cnt = self.iface.props['simd_count'] // self.iface.props['simd_per_cu'] // self.xccs
+    self.waves_per_cu = self.iface.props['max_waves_per_simd'] * self.iface.props['simd_per_cu']
+    self.wave_cnt = (self.cu_cnt * self.waves_per_cu) if self.target[0] != 9 else min(self.cu_cnt * 40, self.se_cnt * self.xccs * 512)
+
+    self.ip_off = importlib.import_module(f"tinygrad.runtime.autogen.am.{'vega' if self.target[0] == 9 else 'navi'}_offsets")
+    self.soc = import_soc(self.target)
+    self.pm4 = importlib.import_module(f"tinygrad.runtime.autogen.am.pm4_{'soc15' if self.target[0] == 9 else 'nv'}")
+    self.sdma = import_module('sdma', min(self.iface.ip_versions[am.SDMA0_HWIP], (6, 0, 0)))
+    self.gc = AMDIP('gc', self.iface.ip_versions[am.GC_HWIP],
+                    bases={i: tuple(getattr(self.ip_off, f'GC_BASE__INST{i}_SEG{s}', 0) for s in range(6)) for i in range(6)})
+
+    self.nbio = AMDIP('nbio' if self.target[0] < 12 else 'nbif', self.iface.ip_versions[am.NBIF_HWIP],
+                      bases={i: tuple(getattr(self.ip_off, f'NBIO_BASE__INST{i}_SEG{s}', 0) for s in range(9)) for i in range(6)})
+
+    self.is_aql = getenv("AMD_AQL", int(self.xccs > 1))
+    if self.is_aql:
+      self.pm4_ibs = self.iface.alloc(0x2000 if self.is_usb() else (16 << 20), uncached=True, cpu_access=True)
+      self.pm4_ib_alloc = BumpAllocator(self.pm4_ibs.size, wrap=True)
+
+    self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
+    self.sdma_queues:dict = {}
+    self.has_sdma_queue = True # self.sdma_queue(0) is not None, TODO: think of this
+
+    super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch)
+
+    # Scratch setup
+    self.max_private_segment_size = 0
+    self.pm_bufferize = PatternMatcher([(UPat(Ops.BUFFER, tag="scratch", name="b"), lambda ctx, b: ctx.scratch_buffer(b.arg))]) + self.pm_bufferize
+
+    self.pmc_enabled:bool = PROFILE > 0 and PMC > 0
+    if self.pmc_enabled:
+      self.iface.require_profile_mode()
+
+      self.pmc_sched:list[PMCSample] = []
+      self.pmc_counters = import_pmc(self.target)
+
+      # validate counters: SQ for SIMD busy/instruction counts, LDS stats, GRBM for GPU cycles, L2 cache hits/misses
+      l2, lds = ("TCC", "SQ") if self.target[0] == 9 else ("GL2C", "SQC")
+      pmc_default = f"SQ_BUSY_CYCLES,SQ_INSTS_VALU,SQ_INSTS_SALU,{lds}_LDS_IDX_ACTIVE,{lds}_LDS_BANK_CONFLICT,GRBM_GUI_ACTIVE,{l2}_HIT,{l2}_MISS"
+      for k in (PMC_COUNTERS:=getenv("PMC_COUNTERS", pmc_default).split(",")):
+        if k not in self.pmc_counters: raise RuntimeError(f"PMC counter {k} is not supported. Available: {','.join(self.pmc_counters.keys())}")
+
+      raise NotImplementedError("PMC start not migrated to hcq2 yet")
+
+    # SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them)
+    self.sqtt_enabled:bool = PROFILE > 0 and SQTT > 0
+    if self.sqtt_enabled:
+      self.iface.require_profile_mode()
+
+      SQTT_BUFFER_SIZE = getenv("SQTT_BUFFER_SIZE", 256) # in mb, per shader engine
+      self.sqtt_buffers = [self.allocator.alloc(SQTT_BUFFER_SIZE<<20, BufferSpec(nolru=True, uncached=True)) for _ in range(self.se_cnt * self.xccs)]
+      self.sqtt_wptrs = self.allocator.alloc(round_up(self.se_cnt * self.xccs * 4, 0x1000), BufferSpec(cpu_access=True, nolru=True))
+      self.sqtt_next_cmd_id = itertools.count(0)
+
+  def create_queue(self, queue_type, ring_size, ctx_save_restore_size=0, eop_buffer_size=0, ctl_stack_size=0, debug_memory_size=0, idx=0):
+    ring = self.iface.alloc(ring_size, uncached=True, cpu_access=True)
+    gart = self.iface.alloc(0x100, uncached=True, cpu_access=True)
+
+    if queue_type == kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL:
+      self.aql_gart = gart
+      self.aql_desc = hsa.amd_queue_t(queue_properties=hsa.AMD_QUEUE_PROPERTIES_IS_PTR64 | hsa.AMD_QUEUE_PROPERTIES_ENABLE_PROFILING,
+        read_dispatch_id_field_base_byte_offset=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset,
+        max_cu_id=(self.cu_cnt * self.xccs) - 1, max_wave_id=self.waves_per_cu - 1)
+      self.aql_gart.cpu_view().view(fmt='B')[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc)
+
+    cwsr_buffer_size = round_up((ctx_save_restore_size + debug_memory_size) * self.xccs, mmap.PAGESIZE)
+    cwsr_buffer = self.iface.alloc(cwsr_buffer_size) if ctx_save_restore_size else None
+    eop_buffer = self.iface.alloc(eop_buffer_size) if eop_buffer_size else None
+
+    queue = (self.iface.create_queue(queue_type, ring, gart, rptr=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset,
+             wptr=getattr(hsa.amd_queue_t, 'write_dispatch_id').offset, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer,
+             ctx_save_restore_size=ctx_save_restore_size, ctl_stack_size=ctl_stack_size, idx=idx))
+
+    qname = f"{'SDMA' if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA else 'COMPUTE'}:{idx}"
+    self.pm_bufferize = PatternMatcher([
+      (UPat(Ops.BUFFER, tag={(qname, name)}), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
+    ]) + self.pm_bufferize
+
+    return queue
+
+  @functools.cached_property
+  def compute_queue(self) -> AMDQueueDesc:
+    # https://gitlab.freedesktop.org/agd5f/linux/-/blob/a1fc9f584c4aaf8bc1ebfa459fc57a3f26a290d8/drivers/gpu/drm/amd/amdkfd/kfd_queue.c#L391
+    sgrp_size_per_cu, hwreg_size_per_cu = 0x4000, 0x1000
+    lds_size_per_cu = self.iface.props["lds_size_in_kb"] << 10 if self.target[:2] == (9,5) else 0x10000
+    vgpr_size_per_cu = 0x60000 if self.target in {(11,0,0), (11,0,1), (11,5,1), (12,0,0), (12,0,1)} else 0x80000 if self.target[0] == 9 else 0x40000
+    wg_data_size = round_up((vgpr_size_per_cu + sgrp_size_per_cu + lds_size_per_cu + hwreg_size_per_cu) * self.cu_cnt, mmap.PAGESIZE)
+    ctl_stack_size = round_up((12 if self.target[0] != 9 else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE)
+    return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL if self.is_aql else kfd.KFD_IOC_QUEUE_TYPE_COMPUTE,
+      0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000,
+      ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size,
+      debug_memory_size=round_up(self.wave_cnt * 32, 64))
+
+  def sdma_queue(self, idx:int):
+    if getenv("AMD_DISABLE_SDMA"): return None
+    if idx in self.sdma_queues: return self.sdma_queues[idx]
+    with contextlib.suppress(OSError):
+      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20), idx=idx)
+    return self.sdma_queues.get(idx, None)
+
+  def tmpring_size(self, private_segment_size):
+    private_segment_size = max(private_segment_size, 128)
+
+    lanes_per_wave = 64 # wave64
+    mem_alignment_size = 256 if self.target[0] != 9 else 1024
+    size_per_thread = round_up(private_segment_size, mem_alignment_size // lanes_per_wave)
+    size_per_xcc = size_per_thread * lanes_per_wave * self.iface.props['max_slots_scratch_cu'] * self.cu_cnt
+
+    # NOTE: xcc logic is correct only for GFX9.
+    max_scratch_waves = self.cu_cnt * self.iface.props['max_slots_scratch_cu'] * self.xccs
+    wave_scratch = ceildiv(lanes_per_wave * size_per_thread, mem_alignment_size)
+    num_waves = (size_per_xcc // (wave_scratch * mem_alignment_size)) // (self.se_cnt if self.target[0] != 9 else 1)
+
+    tmpring_t = getattr(hsa, f'union_COMPUTE_TMPRING_SIZE{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
+    tmpring = int.from_bytes(tmpring_t(WAVES=min(num_waves, max_scratch_waves), WAVESIZE=wave_scratch), 'little')
+
+    if hasattr(self, 'aql_desc'):
+      gfx9_rsrc = {'NUM_FORMAT':hsa.BUF_NUM_FORMAT_UINT, 'DATA_FORMAT':hsa.BUF_DATA_FORMAT_32, 'ELEMENT_SIZE':1, 'INDEX_STRIDE':3}
+      rsrc = {'DST_SEL_X':hsa.SQ_SEL_X, 'DST_SEL_Y':hsa.SQ_SEL_Y, 'DST_SEL_Z':hsa.SQ_SEL_Z, 'DST_SEL_W':hsa.SQ_SEL_W, 'ADD_TID_ENABLE':1,
+              'TYPE':hsa.SQ_RSRC_BUF, **(gfx9_rsrc if self.target[0] == 9 else {'FORMAT':hsa.BUF_FORMAT_32_UINT, 'OOB_SELECT':2})}
+      rsrc1_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD1{"_GFX11" if self.target[0] != 9 else ""}_bitfields')
+      rsrc3_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD3{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
+
+      self.aql_desc.scratch_backing_memory_location = int(self.scratch.get_buf().va_addr)
+      self.aql_desc.scratch_wave64_lane_byte_size = self.max_private_segment_size * lanes_per_wave // 64
+      self.aql_desc.scratch_resource_descriptor[:] = [lo32(self.scratch.get_buf().va_addr),
+        int.from_bytes(rsrc1_t(BASE_ADDRESS_HI=hi32(self.scratch.get_buf().va_addr), SWIZZLE_ENABLE=1), 'little'),
+        lo32(size_per_xcc), int.from_bytes(bytes(rsrc3_t(**rsrc)), 'little')]
+      self.aql_desc.compute_tmpring_size = tmpring
+      self.aql_gart.cpu_view()[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc)
+
+    return tmpring
+
+  def scratch_buffer(self, private_segment_size):
+    private_segment_size = max(private_segment_size, 128)
+    if self.max_private_segment_size < private_segment_size:
+      lanes_per_wave = 64 # wave64
+      mem_alignment_size = 256 if self.target[0] != 9 else 1024
+      size_per_thread = round_up(private_segment_size, mem_alignment_size // lanes_per_wave)
+      size_per_xcc = size_per_thread * lanes_per_wave * self.iface.props['max_slots_scratch_cu'] * self.cu_cnt
+      self.scratch = Buffer(self.device, size_per_xcc * self.xccs, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
+      self.max_private_segment_size = private_segment_size
+    return self.scratch
+
+  def on_device_hang(self): self.iface.on_device_hang()
+
+  def device_props(self): return self.iface.props

tinygrad_repo/extra/hcqfuzz/.gitignore

@@ -0,0 +1 @@
+reports

tinygrad_repo/extra/hcqfuzz/fuzzer.py

@@ -0,0 +1,72 @@
+import os, random, subprocess, shlex, datetime, time, signal
+from extra.hcqfuzz.tools import create_report, on_start_run, collect_tests, init_log, log
+from extra.hcqfuzz.spec import AMSpec
+
+def run_test(dev, test):
+  on_start_run(dev, test)
+
+  dev_env = dev.get_exec_state()
+  test_env, cmd, timeout = test.get_exec_state()
+  env = {**dev_env, **test_env}
+
+  if isinstance(cmd, str): cmd = shlex.split(cmd)
+  assert isinstance(cmd, list), "cmd must be list or str"
+
+  if env is None: env = os.environ.copy()
+  else:
+    env = {k: str(v) for k, v in env.items()}
+    env = {**os.environ, **env}
+
+  start_ts = datetime.datetime.now()
+  t0 = time.perf_counter()
+  log(f"[{start_ts:%Y-%m-%d %H:%M:%S}] running: {test.name()}: {' '.join(cmd)}", end="", flush=True)
+
+  proc = subprocess.Popen(cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+  try:
+    stdout, stderr = proc.communicate(timeout=timeout)
+    ret = proc.returncode
+  except KeyboardInterrupt:
+    print("\nExiting...", flush=True)
+    proc.send_signal(signal.SIGINT)
+    try: stdout, stderr = proc.communicate(timeout=5)
+    except subprocess.TimeoutExpired:
+      proc.kill()
+      stdout, stderr = proc.communicate()
+    raise
+  except subprocess.TimeoutExpired:
+    cur_time = datetime.datetime.now()
+    log(f"\r[{cur_time:%Y-%m-%d %H:%M:%S}] {test.name()} send SIGKILL", end="", flush=True)
+
+    proc.kill()
+    stdout, stderr = proc.communicate()
+    ret = -9
+
+  finish_time = datetime.datetime.now()
+  elapsed = time.perf_counter() - t0
+  if ret != 0:
+    log(f"\r[{finish_time:%Y-%m-%d %H:%M:%S}] {test.name()} failed with {ret} after {elapsed:.1f}s", flush=True)
+    create_report(dev, test, ret, stdout, stderr)
+  else:
+    log(f"\r[{finish_time:%Y-%m-%d %H:%M:%S}] {test.name()} exited {ret} after {elapsed:.1f}s", flush=True)
+
+if __name__ == "__main__":
+  init_log()
+  device_name = "AM"
+  dev = AMSpec()
+
+  start_seed = os.environ.get("SEED", 3332)
+  random.seed(start_seed)
+
+  log(f"Starting with seed {start_seed}")
+
+  test_set = collect_tests()
+  log(f"Found {len(test_set)} tests:")
+  for test in test_set: log(f" - {test.name()}")
+
+  while True:
+    seed = random.randint(0, 2**31)
+    test = random.choice(test_set)
+
+    dev.prepare(seed)
+    test.prepare(dev, seed)
+    run_test(dev, test)

tinygrad_repo/extra/hcqfuzz/readme

@@ -0,0 +1,12 @@
+# Fuzzing Infra
+
+To add a new test, define a `TestSpec`-based class in a file in the `tests/` folder.
+
+You can choose which tests to load from which file:
+```bash
+PYTHONPATH=. RUN_FILES="hcq,allocator" python3 extra/hcqfuzz/fuzzer.py
+```
+Or skip tests from any file:
+```bash
+PYTHONPATH=. SKIP_FILES="allocator" python3 extra/hcqfuzz/fuzzer.py
+```

tinygrad_repo/extra/hcqfuzz/spec.py

@@ -0,0 +1,42 @@
+import os, random
+
+class TestSpec:
+  def prepare(self, device, seed):
+    raise NotImplementedError("prepare must be implemented in the derived class")
+  def get_exec_state(self):
+    raise NotImplementedError("get_exec_state must be implemented in the derived class")
+  def name(self): return self.__class__.__name__
+
+class DeviceSpec:
+  def prepare(self, seed):
+    raise NotImplementedError("prepare must be implemented in the derived class")
+  def get_exec_state(self):
+    raise NotImplementedError("get_exec_state must be implemented in the derived class")
+  def name(self): return self.__class__.__name__
+
+class HCQSpec(DeviceSpec): pass
+class AMDSpec(HCQSpec):
+  def __init__(self):
+    assert os.path.exists('/sys/module/amdgpu'), "amdgpu module should be loaded"
+
+  def prepare(self, seed):
+    self.env = {
+      "AMD": 1,
+      "AMD_LLVM": 0
+    }
+
+  def get_exec_state(self): return self.env
+
+class AMSpec(AMDSpec):
+  def __init__(self):
+    assert not os.path.exists('/sys/module/amdgpu'), "amdgpu module should not be loaded"
+
+  def prepare(self, seed):
+    super().prepare(seed)
+
+    self.env = {
+      **self.env, # from AMDSpec
+      "AMD_SDMA_BIND": random.randint(0, 1),
+      "AMD_ALLOC_QUEUE_DEV_MEM": 0, # random.randint(0, 1) need to validate
+      "AMD_QUEUE_SIZE": 1 << random.randint(10, 26),
+    }

tinygrad_repo/extra/hcqfuzz/tests/allocator.py

@@ -0,0 +1,16 @@
+from extra.hcqfuzz.spec import TestSpec
+import random
+
+class TLSFAllocator(TestSpec):
+  def prepare(self, dev, seed):
+    random.seed(seed)
+
+    self.env = {
+        "SEED": seed,
+        "ITERS": random.randint(10000, 1000000),
+    }
+
+    self.cmd = "python3 test/external/external_fuzz_tlsf.py"
+    self.timeout = 60 * 60 # 60 minutes
+
+  def get_exec_state(self): return self.env, self.cmd, self.timeout

tinygrad_repo/extra/hcqfuzz/tests/allreduce.py

@@ -0,0 +1,17 @@
+from extra.hcqfuzz.spec import TestSpec
+import random
+
+class RingAllreduce(TestSpec):
+  def prepare(self, dev, seed):
+    random.seed(seed)
+
+    self.env = {
+      "GPUS": random.choice([2, 3, 4, 5, 6]),
+      "ITERS": random.randint(10, 1000),
+      "DEBUG": 2,
+    }
+
+    self.cmd = "python3 test/external/external_benchmark_multitensor_allreduce.py"
+    self.timeout = 10 * 60 # 10 minutes
+
+  def get_exec_state(self): return self.env, self.cmd, self.timeout

tinygrad_repo/extra/hcqfuzz/tests/bert.py

@@ -0,0 +1,81 @@
+from extra.hcqfuzz.spec import TestSpec
+import random
+
+bert_train_params = {
+  "DEFAULT_FLOAT": "HALF",
+  "SUM_DTYPE": "HALF",
+  "GPUS": 6,
+  "BS": 96,
+  "EVAL_BS": 96,
+  "BASEDIR": "/raid/datasets/wiki",
+}
+
+class TrainBert(TestSpec):
+  def prepare(self, dev, seed):
+    random.seed(seed)
+
+    self.env = {
+      **bert_train_params,
+      "IGNORE_BEAM_CACHE": 1,
+      "BEAM": 5,
+      "BEAM_UOPS_MAX": 10000,
+      "BEAM_UPCAST_MAX": 256,
+      "BEAM_LOCAL_MAX": 1024,
+      "BEAM_MIN_PROGRESS": 5,
+      "IGNORE_JIT_FIRST_BEAM": 1,
+      "LOGMLPERF": 0,
+      "SEED": seed,
+    }
+
+    self.cmd = "python3 examples/mlperf/model_train.py"
+    self.timeout = 7 * 60 * 60 # 7 hours
+
+  def get_exec_state(self): return self.env, self.cmd, self.timeout
+
+class TrainBertShort(TestSpec):
+  def prepare(self, dev, seed):
+    random.seed(seed)
+
+    self.env = {
+      **bert_train_params,
+      "IGNORE_BEAM_CACHE": 1,
+      "BEAM": 5,
+      "BEAM_UOPS_MAX": 10000,
+      "BEAM_UPCAST_MAX": 256,
+      "BEAM_LOCAL_MAX": 1024,
+      "BEAM_MIN_PROGRESS": 5,
+      "IGNORE_JIT_FIRST_BEAM": 1,
+      "SEED": seed,
+      "BENCHMARK": 4096,
+      "JIT": 2
+    }
+
+    self.cmd = "python3 examples/mlperf/model_train.py"
+    self.timeout = 2 * 60 * 60 # 2 hours
+
+  def get_exec_state(self): return self.env, self.cmd, self.timeout
+
+class BertBeam(TestSpec):
+  def prepare(self, dev, seed):
+    random.seed(seed)
+
+    self.env = {
+      **bert_train_params,
+      "IGNORE_BEAM_CACHE": 1,
+      "BEAM": random.choice([1, 2, 3, 4, 5]),
+      "BEAM_UOPS_MAX": 10000,
+      "BEAM_UPCAST_MAX": 256,
+      "BEAM_LOCAL_MAX": 1024,
+      "BEAM_MIN_PROGRESS": 5,
+      "IGNORE_JIT_FIRST_BEAM": 1,
+      "SEED": seed,
+      "RESET_STEP": 1,
+      "BENCHMARK": 10,
+      "BERT_LAYERS": 2,
+      "SEED": seed,
+    }
+
+    self.cmd = "python3 examples/mlperf/model_train.py"
+    self.timeout = 1 * 60 * 60 # 1 hour
+
+  def get_exec_state(self): return self.env, self.cmd, self.timeout

tinygrad_repo/extra/hcqfuzz/tests/hcq.py

@@ -0,0 +1,35 @@
+from extra.hcqfuzz.spec import TestSpec
+import random
+
+class HCQSignalFuzzer(TestSpec):
+  def prepare(self, dev, seed):
+    random.seed(seed)
+
+    self.env = {
+      "GPUS": random.choice([2, 3, 4, 5, 6]),
+      "ITERS": random.randint(1000000, 10000000),
+      "SEED": seed,
+    }
+
+    self.cmd = "python3 test/external/external_fuzz_hcq_signals.py"
+    self.timeout = 30 * 60 # 30 minutes
+
+  def get_exec_state(self): return self.env, self.cmd, self.timeout
+
+class HCQGraphFuzzer(TestSpec):
+  def prepare(self, dev, seed):
+    random.seed(seed)
+
+    self.env = {
+      "FUZZ_GRAPH_SPLIT_RUNS": random.randint(48, 64),
+      "FUZZ_GRAPH_MAX_SPLITS": random.randint(4, 16),
+      "FUZZ_GRAPH_SPLIT_RETRY_RUNS": random.randint(4, 8),
+      "MAX_KERNELS": random.randint(32, 512),
+      "MAX_DEVICES": random.choice([2, 3, 4, 5, 6]),
+      "ITERS": random.randint(100, 1000),
+    }
+
+    self.cmd = "python3 test/external/fuzz_graph.py"
+    self.timeout = 60 * 60 # 60 minutes
+
+  def get_exec_state(self): return self.env, self.cmd, self.timeout

tinygrad_repo/extra/hcqfuzz/tests/resnet.py

@@ -0,0 +1,66 @@
+from extra.hcqfuzz.spec import TestSpec
+import random
+
+resnet_train_params = {
+  "DEFAULT_FLOAT": "HALF",
+  "SUM_DTYPE": "HALF",
+  "GPUS": 6,
+  "BS": 1536,
+  "EVAL_BS": 192,
+  "TRAIN_BEAM": 4,
+  "IGNORE_JIT_FIRST_BEAM": 1,
+  "BEAM_UOPS_MAX": 2000,
+  "BEAM_UPCAST_MAX": 96,
+  "BEAM_LOCAL_MAX": 1024,
+  "BEAM_MIN_PROGRESS": 5,
+  "BEAM_PADTO": 0,
+  "EVAL_START_EPOCH": 3,
+  "EVAL_FREQ": 4
+}
+
+class TrainResnet(TestSpec):
+  def prepare(self, dev, seed):
+    random.seed(seed)
+
+    self.env = {
+      **resnet_train_params,
+      "IGNORE_BEAM_CACHE": 1,
+      "SEED": seed,
+    }
+
+    self.cmd = "python3 examples/mlperf/model_train.py"
+    self.timeout = 4 * 60 * 60 # 7 hours
+
+  def get_exec_state(self): return self.env, self.cmd, self.timeout
+
+class TrainResnetShort(TestSpec):
+  def prepare(self, dev, seed):
+    random.seed(seed)
+
+    self.env = {
+      **resnet_train_params,
+      "SEED": seed,
+      "BENCHMARK": 4096,
+      "JIT": 2,
+    }
+
+    self.cmd = "python3 examples/mlperf/model_train.py"
+    self.timeout = 2 * 60 * 60 # 2 hours
+
+  def get_exec_state(self): return self.env, self.cmd, self.timeout
+
+class ResnetBeam(TestSpec):
+  def prepare(self, dev, seed):
+    random.seed(seed)
+
+    self.env = {
+      **resnet_train_params,
+      "IGNORE_BEAM_CACHE": 1,
+      "BENCHMARK": 10,
+      "SEED": seed,
+    }
+
+    self.cmd = "python3 examples/mlperf/model_train.py"
+    self.timeout = 1 * 60 * 60 # 1 hour
+
+  def get_exec_state(self): return self.env, self.cmd, self.timeout

tinygrad_repo/extra/hcqfuzz/tools.py

@@ -0,0 +1,80 @@
+import pickle, datetime, os, tempfile, subprocess, zipfile, importlib.util
+from extra.hcqfuzz.spec import TestSpec
+from tinygrad.helpers import getenv
+
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+TEST_DIR = os.path.join(BASE_DIR, "tests")
+REPORTS_DIR = os.path.join(BASE_DIR, "reports")
+
+def collect_tests():
+  run_files = getenv("RUN_FILES", "").split(",")
+  skip_tests = getenv("SKIP_FILES", "").split(",")
+
+  tests = []
+  for filename in os.listdir(TEST_DIR):
+    if filename.endswith(".py") and not filename.startswith("__"):
+      if run_files and filename[:-3] not in run_files: continue
+      if skip_tests and filename[:-3] in skip_tests: continue
+
+      filepath = os.path.join(TEST_DIR, filename)
+      module_name = f"tests.{filename[:-3]}"
+      module = importlib.import_module(module_name)
+      for attr_name in dir(module):
+        attr = getattr(module, attr_name)
+        if isinstance(attr, type) and issubclass(attr, TestSpec) and attr is not TestSpec:
+          tests.append(attr())
+  return tests
+
+def on_start_run(dev, test):
+  os.makedirs(REPORTS_DIR, exist_ok=True)
+  pickle.dump((dev, test), open(f"{REPORTS_DIR}/last_launch.pkl", "wb"))
+
+def create_report(dev, test, result, stdout, stderr):
+  os.makedirs(REPORTS_DIR, exist_ok=True)
+
+  timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+  report_name = f"{timestamp}_{test.name()}_report"
+  report_path = os.path.join(REPORTS_DIR, report_name)
+
+  os.makedirs(report_path, exist_ok=False)
+
+  pickle_path = os.path.join(report_path, "repro.pkl")
+  with open(pickle_path, "wb") as f: pickle.dump((dev, test), f)
+
+  stdout_path = os.path.join(report_path, "stdout.txt")
+  with open(stdout_path, "w") as f: f.write(stdout)
+
+  stderr_path = os.path.join(report_path, "stderr.txt")
+  with open(stderr_path, "w") as f: f.write(stderr)
+
+  dmesg_path = os.path.join(report_path, "dmesg.txt")
+  dmesg_output = subprocess.check_output(["sudo", "dmesg", "--ctime", "--color=never"], text=True)
+  with open(dmesg_path, "w") as f: f.write(dmesg_output)
+
+  env_vars = " ".join(f"{k}={v}" for k, v in test.env.items())
+  reproduce_cmd = f"{env_vars} {test.cmd}"
+
+  summary_path = os.path.join(report_path, "summary.txt")
+  with open(summary_path, "w") as f:
+    f.write(f"Test: {test.name()}\n")
+    f.write(f"Dev params: {vars(dev)}\n")
+    f.write(f"Test params: {vars(test)}\n")
+    f.write(f"Reproduce cmd: {reproduce_cmd}\n")
+    f.write(f"Exit Code: {result}\n")
+
+  print(f"Crash report saved to {report_path}")
+
+_log_file = None
+def init_log():
+  global _log_file
+  os.makedirs(REPORTS_DIR, exist_ok=True)
+
+  ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+  name = f"log_{ts}.log"
+  _log_file = open(f"{REPORTS_DIR}/{name}", "a", buffering=1)
+
+def log(msg="", end="\n", flush=False):
+  global _log_file
+  _log_file.write(msg.replace("\r", "\n") + end)
+  if flush: _log_file.flush()
+  print(msg + " " * 60, end=end, flush=flush)

tinygrad_repo/extra/hevc/.gitignore

@@ -0,0 +1 @@
+out/

tinygrad_repo/extra/hevc/decode.py

@@ -0,0 +1,104 @@
+import argparse, os, hashlib, functools
+from typing import Iterator, Callable
+from tinygrad.helpers import getenv, DEBUG, round_up, Timing, tqdm, fetch, ceildiv
+from extra.hevc.hevc import parse_hevc_file_headers, untile_nv12, to_bgr, nv_gpu
+from tinygrad import Tensor, dtypes, Device, Variable, TinyJit
+
+# rounds up hevc input data to 32 bytes, so more optimal kernels can be generated
+HEVC_ROUNDUP = getenv("DATA_ROUNDUP", 32)
+
+@functools.cache
+def _hevc_jitted_decoder(out_image_size:tuple[int, int], max_hist:int, inplace:bool):
+  def hevc_decode_frame(pos:Variable, hevc_tensor:Tensor, offset:Variable, sz:Variable, opaque:Tensor, i:Variable, *hist:Tensor, outbuf:Tensor|None=None):
+    x = hevc_tensor[offset:offset+sz*HEVC_ROUNDUP].decode_hevc_frame(pos, out_image_size, opaque[i], hist).realize()
+    if outbuf is not None: outbuf.assign(x).realize()
+    return x
+  return TinyJit(hevc_decode_frame)
+
+def hevc_decode(hevc_tensor:Tensor, opaque:Tensor, frame_info:list, luma_h:int, luma_w:int,
+                history:list[Tensor]|None=None, preallocated_outputs:list[Tensor]|None=None, warmup=False) -> Iterator[Tensor]:
+  out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64)
+  max_hist = max((hs for _, _, _, hs, _ in frame_info), default=0)
+
+  v_pos = Variable("pos", 0, max_hist + 1)
+  v_offset = Variable("offset", 0, hevc_tensor.numel()-1)
+  v_sz = Variable("sz", 1, ceildiv(hevc_tensor.numel(), HEVC_ROUNDUP))
+  v_i = Variable("i", 0, len(frame_info)-1)
+
+  decode_jit = _hevc_jitted_decoder(out_image_size, max_hist, preallocated_outputs is not None)
+  history = history or [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV").contiguous().realize() for _ in range(max_hist)]
+  assert len(history) == max_hist, f"history length {len(history)} does not match max_hist {max_hist}"
+
+  for i, (offset, sz, frame_pos, _, is_hist) in enumerate(frame_info):
+    history = history[-max_hist:] if max_hist > 0 else []
+    img = decode_jit(v_pos.bind(frame_pos), hevc_tensor, v_offset.bind(offset), v_sz.bind(ceildiv(sz, HEVC_ROUNDUP)),
+                     opaque, v_i.bind(i), *history, outbuf=preallocated_outputs[i] if preallocated_outputs else None)
+    res = preallocated_outputs[i] if preallocated_outputs else img.clone().realize()
+    if is_hist: history.append(res)
+    yield res
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--input_file", type=str, default="")
+  parser.add_argument("--output_dir", type=str, default="extra/hevc/out")
+  args = parser.parse_args()
+
+  if args.input_file == "":
+    url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc"
+    hevc_tensor = Tensor.from_url(url, device="CPU")
+  else:
+    hevc_tensor = Tensor.empty(os.stat(args.input_file).st_size, dtype=dtypes.uint8, device=f"disk:{args.input_file}").to("CPU")
+
+  dat = bytes(hevc_tensor.data())
+  dat_hash = hashlib.md5(dat).hexdigest()
+
+  with Timing("prep infos: "):
+    opaque, frame_info, w, h, luma_w, luma_h, chroma_off = parse_hevc_file_headers(dat)
+
+  frame_info = frame_info[:getenv("MAX_FRAMES", len(frame_info))]
+
+  # move all needed data to gpu
+  with Timing("copy to gpu: "):
+    opaque_nv = opaque.to("NV").contiguous().realize()
+    hevc_tensor = hevc_tensor.to("NV")
+
+  out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64)
+
+  # preallocate output/hist buffers
+  max_hist = max((hs for _, _, _, hs, _ in frame_info), default=0)
+  hist = [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV").contiguous().realize() for _ in range(max_hist)]
+  out_images = [Tensor.zeros(*out_image_size, dtype=dtypes.uint8, device="NV").contiguous().realize() for _ in range(len(frame_info))]
+
+  # warmup decode
+  _ = list(hevc_decode(hevc_tensor, opaque_nv, frame_info[:3], luma_h, luma_w, history=hist, preallocated_outputs=out_images))
+  Device.default.synchronize()
+
+  # decode all frames using the iterator
+  tm = Timing("decoding whole file: ", on_exit=(lambda et: f", {len(frame_info)} frames, {len(frame_info)/(et/1e9):.2f} fps"))
+  with tm:
+    images = list(hevc_decode(hevc_tensor, opaque_nv, frame_info, luma_h, luma_w, history=hist, preallocated_outputs=out_images))
+    Device.default.synchronize()
+
+  fps = len(frame_info)/(tm.et/1e9)
+  assert fps >= getenv("ASSERT_FPS", 0), f"HEVC decode too slow: {fps:.2f} fps"
+
+  # validation
+  if getenv("VALIDATE", 0):
+    import pickle
+    if dat_hash == "b813bfdbec194fd17fdf0e3ceb8cea1c":
+      url = "https://github.com/nimlgen/hevc_validate_set/raw/refs/heads/main/decoded_frames_b813bfdbec194fd17fdf0e3ceb8cea1c.pkl"
+      decoded_frames = pickle.load(fetch(url).open("rb"))
+    else: decoded_frames = pickle.load(open(f"extra/hevc/decoded_frames_{dat_hash}.pkl", "rb"))
+  else: import cv2
+
+  for i, img in tqdm(enumerate(images)):
+    if getenv("VALIDATE", 0):
+      if i < len(decoded_frames) and len(decoded_frames[i]) > 0:
+        img = untile_nv12(img, h, w, luma_w, chroma_off).realize()
+        assert img.data() == decoded_frames[i], f"Frame {i} does not match reference decoder!"
+        print(f"Frame {i} matches reference decoder!")
+    else:
+      if len(args.output_dir):
+        os.makedirs(args.output_dir, exist_ok=True)
+        img = to_bgr(img, h, w, luma_w, chroma_off).realize()
+        cv2.imwrite(f"{args.output_dir}/out_frame_{i:04d}.png", img.numpy())

tinygrad_repo/extra/hevc/hevc.py

@@ -0,0 +1,450 @@
+import dataclasses, enum, argparse, os, itertools, time, ctypes
+from typing import Any
+from tinygrad import Tensor, dtypes, Device, TinyJit
+from tinygrad.helpers import DEBUG, round_up, ceildiv, Timing, prod
+from tinygrad.runtime.autogen import avcodec, nv_570 as nv_gpu
+
+class BitReader:
+  def __init__(self, data:bytes): self.reader, self.current_bits, self.bits, self.read_bits, self.total = iter(data), 0, 0, 0, len(data) * 8
+  def empty(self): return self.read_bits == self.total and self.current_bits == 0
+  def peak_bits(self, n):
+    while self.current_bits < n:
+      self.bits = (self.bits << 8) | next(self.reader)
+      self.current_bits += 8
+      self.read_bits += 8
+    return (self.bits >> (self.current_bits - n)) & ((1 << n) - 1)
+  def _next_bits(self, n):
+    val = self.peak_bits(n)
+    self.bits &= (1 << (self.current_bits - n)) - 1
+    self.current_bits -= n
+    return val
+
+  def u(self, n): return self._next_bits(n)
+
+  # 9.2 Parsing process for 0-th order Exp-Golomb codes
+  def ue_v(self):
+    leading_zero_bits = -1
+    while True:
+      bit = self.u(1)
+      leading_zero_bits += 1
+      if bit == 1: break
+
+    part = self.u(leading_zero_bits)
+
+    if leading_zero_bits == 0: return 0
+    return (1 << leading_zero_bits) - 1 + part
+
+  # 9.2.2 Mapping process for signed Exp-Golomb codes
+  def se_v(self):
+    k = self.ue_v()
+    return (-1 ** (k + 1)) * (k // 2)
+
+# 7.3.1.1 General NAL unit syntax
+def _hevc_get_rbsp(dat:bytes, off=0) -> bytes:
+  rbsp = bytes()
+  while off < len(dat):
+    if off + 2 < len(dat) and dat[off:off+3] == b'\x00\x00\x03':
+      rbsp += bytes([0, 0])
+      off += 3
+    else:
+      rbsp += bytes([dat[off]])
+      off += 1
+  return rbsp
+
+class HevcSlice:
+  # 7.3.3 Profile, tier and level syntax
+  def profile_tier_level(self, r:BitReader, enable:bool, max_sub_layers:int):
+    assert enable and max_sub_layers == 0, "no sublayers supported"
+    self._notimpl_profile_tier_level = r.u(88)
+    self.general_level_idc = r.u(8)
+
+  # 7.3.7 Short-term reference picture set syntax
+  def st_ref_pic_set(self, r:BitReader, stRpsIdx:int, num_short_term_ref_pic_sets:int=0, sps=None):
+    inter_ref_pic_set_prediction_flag = r.u(1) if stRpsIdx != 0 else 0
+
+    if inter_ref_pic_set_prediction_flag:
+      if stRpsIdx == num_short_term_ref_pic_sets:
+        delta_idx_minus1 = r.ue_v()
+      delta_rps_sign = r.u(1)
+      abs_delta_rps_minus1 = r.ue_v()
+
+      NumDeltaPocs = sps.num_negative_pics + sps.num_positive_pics
+      for i in range(NumDeltaPocs + 1):
+        used_by_curr_pic_flag = r.u(1)
+        if not used_by_curr_pic_flag:
+          use_delta_flag = r.u(1)
+    else:
+      self.num_negative_pics = r.ue_v()
+      self.num_positive_pics = r.ue_v()
+      for i in range(self.num_negative_pics):
+        delta_poc_s0_minus1 = r.ue_v()
+        used_by_curr_pic_s0_flag = r.u(1)
+      for i in range(self.num_positive_pics):
+        delta_poc_s1_minus1 = r.ue_v()
+        used_by_curr_pic_s1_flag = r.u(1)
+
+# 7.3.2.2 Sequence parameter set RBSP syntax
+class SPS(HevcSlice):
+  def __init__(self, r:BitReader):
+    self.sps_video_parameter_set_id = r.u(4)
+    self.sps_max_sub_layers_minus1 = r.u(3)
+    self.sps_temporal_id_nesting_flag = r.u(1)
+
+    self.profile_tier_level(r, True, self.sps_max_sub_layers_minus1)
+
+    self.sps_seq_parameter_set_id = r.ue_v()
+    self.chroma_format_idc = r.ue_v()
+    self.separate_colour_plane_flag = r.u(1) if self.chroma_format_idc == 3 else 0
+    self.pic_width_in_luma_samples = r.ue_v()
+    self.pic_height_in_luma_samples = r.ue_v()
+    self.conformance_window_flag = r.u(1)
+
+    if self.conformance_window_flag:
+      self.conf_win_left_offset = r.ue_v()
+      self.conf_win_right_offset = r.ue_v()
+      self.conf_win_top_offset = r.ue_v()
+      self.conf_win_bottom_offset = r.ue_v()
+    else: self.conf_win_left_offset = self.conf_win_right_offset = self.conf_win_top_offset = self.conf_win_bottom_offset = 0
+
+    self.bit_depth_luma = r.ue_v() + 8
+    self.bit_depth_chroma = r.ue_v() + 8
+    self.log2_max_pic_order_cnt_lsb_minus4 = r.ue_v()
+    self.sps_sub_layer_ordering_info_present_flag = r.u(1)
+    self.sps_max_dec_pic_buffering, self.sps_max_num_reorder_pics, self.sps_max_latency_increase_plus1 = [], [], []
+    for i in range((0 if self.sps_sub_layer_ordering_info_present_flag else self.sps_max_sub_layers_minus1), self.sps_max_sub_layers_minus1 + 1):
+      self.sps_max_dec_pic_buffering.append(r.ue_v() + 1)
+      self.sps_max_num_reorder_pics.append(r.ue_v())
+      self.sps_max_latency_increase_plus1.append(r.ue_v())
+    self.log2_min_luma_coding_block_size = r.ue_v() + 3
+    self.log2_max_luma_coding_block_size = self.log2_min_luma_coding_block_size + r.ue_v()
+    self.log2_min_transform_block_size = r.ue_v() + 2
+    self.log2_max_transform_block_size = self.log2_min_transform_block_size + r.ue_v()
+    self.max_transform_hierarchy_depth_inter = r.ue_v()
+    self.max_transform_hierarchy_depth_intra = r.ue_v()
+    if scaling_list_enabled_flag := r.u(1):
+      if sps_scaling_list_data_present_flag := r.u(1): assert False, "scaling_list_data parsing not implemented"
+    self.amp_enabled_flag = r.u(1)
+    self.sample_adaptive_offset_enabled_flag = r.u(1)
+    self.pcm_enabled_flag = r.u(1)
+    assert self.pcm_enabled_flag == 0, "pcm not implemented"
+    self.num_short_term_ref_pic_sets = r.ue_v()
+    for i in range(self.num_short_term_ref_pic_sets):
+      self.st_ref_pic_set(r, i, self.num_short_term_ref_pic_sets)
+    self.long_term_ref_pics_present_flag = r.u(1)
+    if self.long_term_ref_pics_present_flag: assert False, "long_term_ref_pics parsing not implemented"
+    self.sps_temporal_mvp_enabled_flag = r.u(1)
+    self.strong_intra_smoothing_enabled_flag = r.u(1)
+
+# 7.3.2.3 Picture parameter set RBSP syntax
+class PPS(HevcSlice):
+  def __init__(self, r:BitReader):
+    self.pps_pic_parameter_set_id = r.ue_v()
+    self.pps_seq_parameter_set_id = r.ue_v()
+    self.dependent_slice_segments_enabled_flag = r.u(1)
+    self.output_flag_present_flag = r.u(1)
+    self.num_extra_slice_header_bits = r.u(3)
+    self.sign_data_hiding_enabled_flag = r.u(1)
+    self.cabac_init_present_flag = r.u(1)
+    self.num_ref_idx_l0_default_active = r.ue_v() + 1
+    self.num_ref_idx_l1_default_active = r.ue_v() + 1
+    self.init_qp = r.se_v() + 26
+    self.constrained_intra_pred_flag = r.u(1)
+    self.transform_skip_enabled_flag = r.u(1)
+    self.cu_qp_delta_enabled_flag = r.u(1)
+    if self.cu_qp_delta_enabled_flag: self.diff_cu_qp_delta_depth = r.ue_v()
+
+    self.pps_cb_qp_offset = r.se_v()
+    self.pps_cr_qp_offset = r.se_v()
+    self.pps_slice_chroma_qp_offsets_present_flag = r.u(1)
+    self.weighted_pred_flag = r.u(1)
+    self.weighted_bipred_flag = r.u(1)
+    self.transquant_bypass_enabled_flag = r.u(1)
+    self.tiles_enabled_flag = r.u(1)
+    self.entropy_coding_sync_enabled_flag = r.u(1)
+    if self.tiles_enabled_flag:
+      self.num_tile_columns_minus1 = r.ue_v()
+      self.num_tile_rows_minus1 = r.ue_v()
+      self.uniform_spacing_flag = r.u(1)
+      self.column_width_minus1, self.row_height_minus1 = [], []
+      if not self.uniform_spacing_flag:
+        for i in range(self.num_tile_columns_minus1): self.column_width_minus1.append(r.ue_v())
+        for i in range(self.num_tile_rows_minus1): self.row_height_minus1.append(r.ue_v())
+      self.loop_filter_across_tiles_enabled_flag = r.u(1)
+    self.loop_filter_across_slices_enabled_flag = r.u(1)
+    self.deblocking_filter_control_present_flag = r.u(1)
+    if self.deblocking_filter_control_present_flag: assert False, "deblocking_filter parsing not implemented"
+    self.scaling_list_data_present_flag = r.u(1)
+    if self.scaling_list_data_present_flag: assert False, "scaling_list_data parsing not implemented"
+    self.lists_modification_present_flag = r.u(1)
+    self.log2_parallel_merge_level = r.ue_v() + 2
+
+# 7.3.6 Slice segment header syntax
+class SliceSegment(HevcSlice):
+  def __init__(self, r:BitReader, nal_unit_type:int, sps:SPS, pps:PPS):
+    self.first_slice_segment_in_pic_flag = r.u(1)
+    if nal_unit_type >= avcodec.HEVC_NAL_BLA_W_LP and nal_unit_type <= avcodec.HEVC_NAL_RSV_IRAP_VCL23:
+      self.no_output_of_prior_pics_flag = r.u(1)
+    self.slice_pic_parameter_set_id = r.ue_v()
+    if not self.first_slice_segment_in_pic_flag:
+      if pps.dependent_slice_segments_enabled_flag:
+        self.dependent_slice_segment_flag = r.u(1)
+      self.slice_segment_address = r.ue_v()
+    self.dependent_slice_segment_flag = 0
+    if not self.dependent_slice_segment_flag:
+      r.u(pps.num_extra_slice_header_bits) # extra bits ignored
+      self.slice_type = r.ue_v()
+
+      self.sw_skip_start = r.read_bits - r.current_bits
+      self.pic_output_flag = r.u(1) if pps.output_flag_present_flag else 0
+      self.colour_plane_id = r.u(2) if sps.separate_colour_plane_flag else 0
+
+      if nal_unit_type != avcodec.HEVC_NAL_IDR_W_RADL and nal_unit_type != avcodec.HEVC_NAL_IDR_N_LP:
+        self.slice_pic_order_cnt_lsb = r.u(sps.log2_max_pic_order_cnt_lsb_minus4 + 4)
+
+        self.short_term_ref_pic_set_sps_flag = r.u(1)
+        if not self.short_term_ref_pic_set_sps_flag:
+          self.short_term_ref_pics_in_slice_start = r.read_bits - r.current_bits
+          self.st_ref_pic_set(r, sps.num_short_term_ref_pic_sets, sps=sps)
+          self.short_term_ref_pics_in_slice_end = r.read_bits - r.current_bits
+        elif sps.num_short_term_ref_pic_sets > 1: assert False, "short_term_ref_pic_set parsing not implemented"
+
+        if sps.long_term_ref_pics_present_flag: assert False, "long_term_ref_pics parsing not implemented"
+
+        self.sw_skip_end = r.read_bits - r.current_bits
+        self.slice_temporal_mvp_enabled_flag = r.u(1) if sps.sps_temporal_mvp_enabled_flag else 0
+      else: self.slice_pic_order_cnt_lsb, self.sw_skip_end = 0, self.sw_skip_start
+
+      if sps.sample_adaptive_offset_enabled_flag:
+        slice_sao_luma_flag = r.u(1)
+        ChromaArrayType = sps.chroma_format_idc if sps.separate_colour_plane_flag == 0 else 0
+        slice_sao_chroma_flag = r.u(1) if ChromaArrayType != 0 else 0
+
+      if self.slice_type in {avcodec.HEVC_SLICE_B, avcodec.HEVC_SLICE_B}:
+        if num_ref_idx_active_override_flag := r.u(1):
+          num_ref_idx_l0_active_minus1 = r.ue_v()
+          num_ref_idx_l1_active_minus1 = r.ue_v() if self.slice_type == avcodec.HEVC_SLICE_B else 0
+
+def fill_sps_into_dev_context(device_ctx, sps:SPS):
+  device_ctx.chroma_format_idc = sps.chroma_format_idc
+  device_ctx.pic_width_in_luma_samples = sps.pic_width_in_luma_samples
+  device_ctx.pic_height_in_luma_samples = sps.pic_height_in_luma_samples
+  device_ctx.bit_depth_luma = sps.bit_depth_luma
+  device_ctx.bit_depth_chroma = sps.bit_depth_chroma
+  device_ctx.log2_max_pic_order_cnt_lsb_minus4 = sps.log2_max_pic_order_cnt_lsb_minus4
+  device_ctx.log2_min_luma_coding_block_size = sps.log2_min_luma_coding_block_size
+  device_ctx.log2_max_luma_coding_block_size = sps.log2_max_luma_coding_block_size
+  device_ctx.log2_min_transform_block_size = sps.log2_min_transform_block_size
+  device_ctx.log2_max_transform_block_size = sps.log2_max_transform_block_size
+  device_ctx.amp_enabled_flag = sps.amp_enabled_flag
+  device_ctx.pcm_enabled_flag = sps.pcm_enabled_flag
+  device_ctx.sample_adaptive_offset_enabled_flag = sps.sample_adaptive_offset_enabled_flag
+  device_ctx.sps_temporal_mvp_enabled_flag = sps.sps_temporal_mvp_enabled_flag
+  device_ctx.strong_intra_smoothing_enabled_flag = sps.strong_intra_smoothing_enabled_flag
+
+def fill_pps_into_dev_context(device_ctx, pps:PPS):
+  device_ctx.sign_data_hiding_enabled_flag = pps.sign_data_hiding_enabled_flag
+  device_ctx.cabac_init_present_flag = pps.cabac_init_present_flag
+  device_ctx.num_ref_idx_l0_default_active = pps.num_ref_idx_l0_default_active
+  device_ctx.num_ref_idx_l1_default_active = pps.num_ref_idx_l1_default_active
+  device_ctx.init_qp = pps.init_qp
+  device_ctx.cu_qp_delta_enabled_flag = pps.cu_qp_delta_enabled_flag
+  device_ctx.diff_cu_qp_delta_depth = getattr(pps, 'diff_cu_qp_delta_depth', 0)
+  device_ctx.pps_cb_qp_offset = pps.pps_cb_qp_offset
+  device_ctx.pps_cr_qp_offset = pps.pps_cr_qp_offset
+  device_ctx.pps_slice_chroma_qp_offsets_present_flag = pps.pps_slice_chroma_qp_offsets_present_flag
+  device_ctx.weighted_pred_flag = pps.weighted_pred_flag
+  device_ctx.weighted_bipred_flag = pps.weighted_bipred_flag
+  device_ctx.transquant_bypass_enabled_flag = pps.transquant_bypass_enabled_flag
+  device_ctx.tiles_enabled_flag = pps.tiles_enabled_flag
+  device_ctx.entropy_coding_sync_enabled_flag = pps.entropy_coding_sync_enabled_flag
+  device_ctx.loop_filter_across_slices_enabled_flag = pps.loop_filter_across_slices_enabled_flag
+  device_ctx.deblocking_filter_control_present_flag = pps.deblocking_filter_control_present_flag
+  device_ctx.scaling_list_data_present_flag = pps.scaling_list_data_present_flag
+  device_ctx.lists_modification_present_flag = pps.lists_modification_present_flag
+  device_ctx.log2_parallel_merge_level = pps.log2_parallel_merge_level
+  device_ctx.loop_filter_across_tiles_enabled_flag = getattr(pps, 'loop_filter_across_tiles_enabled_flag', 0)
+
+def parse_hevc_file_headers(dat:bytes, device="NV"):
+  res = []
+  nal_unit_start = 1
+  history:list[tuple[int, int, int]] = []
+  device_ctx = nv_gpu.nvdec_hevc_pic_s(gptimer_timeout_value=92720000, tileformat=1, sw_start_code_e=1, pattern_id=2)
+  nal_infos = []
+  ctx_bytes = bytes()
+  align_ctx_bytes_size = 0x300
+
+  def _flush_picture():
+    nonlocal res, history, device_ctx, nal_infos, ctx_bytes, align_ctx_bytes_size
+
+    if not len(nal_infos): return
+
+    hdr, nal_unit_type = nal_infos[0][0]
+    assert all(nal_unit_type == x[0][1] for x in nal_infos), "all NAL units in a picture must be of the same type"
+
+    device_ctx.curr_pic_idx = next(i for i in range(16) if all(d[0] != i for d in history))
+
+    if nal_unit_type in {avcodec.HEVC_NAL_IDR_W_RADL, avcodec.HEVC_NAL_IDR_N_LP}:
+      history = []
+
+    device_ctx.num_ref_frames = len(history)
+    device_ctx.IDR_picture_flag = int(nal_unit_type in {avcodec.HEVC_NAL_IDR_W_RADL, avcodec.HEVC_NAL_IDR_N_LP})
+    device_ctx.RAP_picture_flag = int(nal_unit_type >= avcodec.HEVC_NAL_BLA_W_LP and nal_unit_type <= avcodec.HEVC_NAL_RSV_IRAP_VCL23)
+    device_ctx.RefDiffPicOrderCnts=(ctypes.c_int16 * 16)()
+    device_ctx.colMvBuffersize = (round_up(sps.pic_width_in_luma_samples, 64) * round_up(sps.pic_height_in_luma_samples, 64) // 16) // 256
+    device_ctx.framestride=(ctypes.c_uint32 * 2)(round_up(sps.pic_width_in_luma_samples, 64), round_up(sps.pic_width_in_luma_samples, 64))
+    device_ctx.sw_hdr_skip_length = hdr.sw_skip_end - hdr.sw_skip_start
+    device_ctx.num_bits_short_term_ref_pics_in_slice = max(0, device_ctx.sw_hdr_skip_length - 9)
+    device_ctx.stream_len = sum(x[2] for x in nal_infos)
+
+    if pps.tiles_enabled_flag:
+      device_ctx.num_tile_columns = pps.num_tile_columns_minus1 + 1
+      device_ctx.num_tile_rows = pps.num_tile_rows_minus1 + 1
+
+    device_ctx.num_short_term_ref_pic_sets = sps.num_short_term_ref_pic_sets
+
+    luma_h_rounded = round_up(sps.pic_height_in_luma_samples, 64)
+    device_ctx.HevcSaoBufferOffset = (608 * luma_h_rounded) >> 8
+    device_ctx.HevcBsdCtrlOffset = ((device_ctx.HevcSaoBufferOffset<<8) + 4864 * luma_h_rounded) >> 8
+
+    device_ctx.v1.hevc_main10_444_ext.HevcFltAboveOffset = ((device_ctx.HevcBsdCtrlOffset<<8) + 152 * luma_h_rounded) >> 8
+    device_ctx.v1.hevc_main10_444_ext.HevcSaoAboveOffset = ((device_ctx.v1.hevc_main10_444_ext.HevcFltAboveOffset<<8) + 2000 * luma_h_rounded) >> 8
+    device_ctx.v3.HevcSliceEdgeOffset = device_ctx.v1.hevc_main10_444_ext.HevcSaoAboveOffset
+
+    before_list, after_list = [], []
+    for pic_idx, poc, _ in history:
+      device_ctx.RefDiffPicOrderCnts[pic_idx] = hdr.slice_pic_order_cnt_lsb - poc
+      if hdr.slice_pic_order_cnt_lsb < poc: after_list.append((poc - hdr.slice_pic_order_cnt_lsb, pic_idx))
+      else: before_list.append((hdr.slice_pic_order_cnt_lsb - poc, pic_idx))
+    before_list.sort()
+    after_list.sort()
+
+    device_ctx.initreflistidxl0 = (ctypes.c_uint8 * 16)(*[idx for _,idx in before_list + after_list])
+    if hdr.slice_type == avcodec.HEVC_SLICE_B: device_ctx.initreflistidxl1 = (ctypes.c_uint8 * 16)(*[idx for _,idx in after_list + before_list])
+
+    locl_ctx_bytes = bytes(device_ctx)
+    locl_ctx_bytes += b'\x00\x00\x00\x00\x00\x00\x00\x00\x10\x00\x00\x00' # blackwell extension
+    locl_ctx_bytes += bytes(0x200 - len(locl_ctx_bytes)) # pad to 512 bytes
+
+    pic_width_in_ctbs = ceildiv(sps.pic_width_in_luma_samples, (1 << sps.log2_max_luma_coding_block_size))
+    pic_height_in_ctbs = ceildiv(sps.pic_height_in_luma_samples, (1 << sps.log2_max_luma_coding_block_size))
+    # append tile sizes 0x200
+    if pps.tiles_enabled_flag and pps.uniform_spacing_flag:
+      assert device_ctx.num_tile_columns == 1 and device_ctx.num_tile_rows == 1, "not implemented: uniform spacing with multiple tiles"
+      locl_ctx_bytes += pic_width_in_ctbs.to_bytes(2, "little") + pic_height_in_ctbs.to_bytes(2, "little")
+    else:
+      if pps.tiles_enabled_flag and not getattr(pps, 'uniform_spacing_flag', 0):
+        column_width = [cw_minus1 + 1 for cw_minus1 in pps.column_width_minus1[0:pps.num_tile_columns_minus1]]
+        row_height = [rh_minus1 + 1 for rh_minus1 in pps.row_height_minus1[0:pps.num_tile_rows_minus1]]
+      else:
+        column_width = []
+        row_height = []
+
+      column_width.append(pic_width_in_ctbs - sum(column_width))
+      row_height.append(pic_height_in_ctbs - sum(row_height))
+
+      for c in column_width:
+        for r in row_height: locl_ctx_bytes += c.to_bytes(2, "little") + r.to_bytes(2, "little")
+
+    luma_size = round_up(sps.pic_width_in_luma_samples, 64) * round_up(sps.pic_height_in_luma_samples, 64)
+    chroma_size = round_up(sps.pic_width_in_luma_samples, 64) * round_up((sps.pic_height_in_luma_samples + 1) // 2, 64)
+    is_hist = nal_unit_type in {avcodec.HEVC_NAL_TRAIL_R, avcodec.HEVC_NAL_IDR_N_LP, avcodec.HEVC_NAL_IDR_W_RADL}
+
+    res.append((nal_infos[0][1], device_ctx.stream_len, device_ctx.curr_pic_idx, len(history), is_hist))
+
+    locl_ctx_bytes += (align_ctx_bytes_size - len(locl_ctx_bytes)) * b'\x00'
+    ctx_bytes += locl_ctx_bytes
+
+    if nal_unit_type in {avcodec.HEVC_NAL_TRAIL_R, avcodec.HEVC_NAL_IDR_N_LP, avcodec.HEVC_NAL_IDR_W_RADL}:
+      history.append((device_ctx.curr_pic_idx, hdr.slice_pic_order_cnt_lsb, None))
+
+    if len(history) >= sps.sps_max_dec_pic_buffering[0]:
+      # remove the oldest poc
+      history.pop(0)
+
+    nal_infos = []
+
+  cnt = 0
+  while nal_unit_start < len(dat):
+    assert dat[nal_unit_start:nal_unit_start+3] == b"\x00\x00\x01", "NAL unit start code not found"
+
+    pos = dat.find(b"\x00\x00\x01", nal_unit_start + 3)
+    nal_unit_len = (pos if pos != -1 else len(dat)) - nal_unit_start
+
+    # 7.3.1.1 General NAL unit syntax
+    nal_unit_type = (dat[nal_unit_start+3] >> 1) & 0x3F
+    slice_dat = dat[nal_unit_start+5:nal_unit_start+nal_unit_len]
+
+    if nal_unit_type == avcodec.HEVC_NAL_SPS:
+      sps = SPS(BitReader(_hevc_get_rbsp(slice_dat)))
+      fill_sps_into_dev_context(device_ctx, sps)
+    elif nal_unit_type == avcodec.HEVC_NAL_PPS:
+      pps = PPS(BitReader(_hevc_get_rbsp(slice_dat)))
+      fill_pps_into_dev_context(device_ctx, pps)
+    elif nal_unit_type in {avcodec.HEVC_NAL_IDR_N_LP, avcodec.HEVC_NAL_IDR_W_RADL, avcodec.HEVC_NAL_TRAIL_R, avcodec.HEVC_NAL_TRAIL_N}:
+      hdr = SliceSegment(BitReader(slice_dat), nal_unit_type, sps, pps)
+
+      if hdr.first_slice_segment_in_pic_flag == 1: _flush_picture()
+      nal_infos.append(((hdr, nal_unit_type), nal_unit_start, nal_unit_len))
+
+    nal_unit_start += nal_unit_len
+  _flush_picture()
+
+  w = sps.pic_width_in_luma_samples - 2 * (sps.conf_win_left_offset + sps.conf_win_right_offset)
+  h = sps.pic_height_in_luma_samples - 2 * (sps.conf_win_top_offset  + sps.conf_win_bottom_offset)
+  chroma_off = round_up(sps.pic_width_in_luma_samples, 64) * round_up(sps.pic_height_in_luma_samples, 64)
+  opaque = Tensor(ctx_bytes, device=device).reshape(len(res), align_ctx_bytes_size)
+  return opaque, res, w, h, sps.pic_width_in_luma_samples, sps.pic_height_in_luma_samples, chroma_off
+
+def _addr_table(h, w, w_aligned):
+  GOB_W, GOB_H = 64, 8
+  GOB_SIZE = GOB_W * GOB_H
+  BLOCK_H_GOBS = 2
+
+  xs = Tensor.arange(w, dtype=dtypes.uint32).reshape(1, w)
+  ys = Tensor.arange(h, dtype=dtypes.uint32).reshape(h, 1)
+
+  gob_x = xs // GOB_W
+  gob_y = ys // GOB_H
+  super_block_y = gob_y // BLOCK_H_GOBS
+  gob_y_in_block = gob_y  % BLOCK_H_GOBS
+  stride_gobs = w_aligned // GOB_W
+
+  base = ((super_block_y * stride_gobs + gob_x) * BLOCK_H_GOBS + gob_y_in_block) * GOB_SIZE
+
+  lx, ly = xs % GOB_W, ys % GOB_H
+  swiz = (lx & 0x0F) | ((ly & 0x03) << 4) | ((lx & 0x10) << 2) | ((ly & 0x04) << 5) | ((lx & 0x20) << 3)
+  return (base + swiz).reshape(-1)
+
+def nv12_to_bgr_from_planes(luma: Tensor, chroma: Tensor, h: int, w: int) -> Tensor:
+  Y = luma.reshape(h, w).cast(dtypes.float32)
+
+  uv = chroma.reshape(h // 2, w // 2, 2).cast(dtypes.float32)
+  U_small = uv[..., 0]
+  V_small = uv[..., 1]
+
+  U = U_small.reshape(h // 2, 1, w // 2, 1).expand(h // 2, 2, w // 2, 2).reshape(h, w)
+  V = V_small.reshape(h // 2, 1, w // 2, 1).expand(h // 2, 2, w // 2, 2).reshape(h, w)
+
+  C = Y - 16.0
+  D = U - 128.0
+  E = V - 128.0
+
+  R = 1.1643835616438356 * C + 1.5960267857142858 * E
+  G = 1.1643835616438356 * C - 0.39176229009491365 * D - 0.8129676472377708 * E
+  B = 1.1643835616438356 * C + 2.017232142857143  * D
+
+  R = R.maximum(0.0).minimum(255.0)
+  G = G.maximum(0.0).minimum(255.0)
+  B = B.maximum(0.0).minimum(255.0)
+
+  return Tensor.stack([B, G, R], dim=2).cast(dtypes.uint8)
+
+def untile_nv12(src:Tensor, h:int, w:int, luma_w:int, chroma_off:int) -> Tensor:
+  luma = src.reshape(-1)[_addr_table(h, w, round_up(luma_w, 64))]
+  chroma = src.reshape(-1)[chroma_off:][_addr_table((h + 1) // 2, w, round_up(luma_w, 64))]
+  return luma.cat(chroma).realize()
+
+def to_bgr(tensor:Tensor, h:int, w:int, luma_w:int, chroma_off:int) -> Tensor:
+  luma = tensor.reshape(-1)[_addr_table(h, w, round_up(luma_w, 64))]
+  chroma = tensor.reshape(-1)[chroma_off:][_addr_table((h + 1) // 2, w, round_up(luma_w, 64))]
+  return nv12_to_bgr_from_planes(luma, chroma, h, w).realize()

tinygrad_repo/extra/hip_gpu_driver/hip_ioctl.py

@@ -0,0 +1,117 @@
+# type: ignore
+import ctypes, ctypes.util, struct, platform, pathlib, re, time, os
+start = time.perf_counter()
+
+# *** ioctl lib ***
+libc = ctypes.CDLL(ctypes.util.find_library("c"))
+# platform.processor calls `uname -p` which can return `unknown` on some systems
+processor = os.getenv("IOCTL_PROCESSOR") or platform.processor() or platform.machine()
+IOCTL_SYSCALL = {"aarch64": 0x1d, "x86_64":16}[processor]
+
+def get_struct(argp, stype):
+  return ctypes.cast(ctypes.c_void_p(argp), ctypes.POINTER(stype)).contents
+
+def format_struct(s):
+  sdats = []
+  for field_name, *_ in s._real_fields_:
+    dat = getattr(s, field_name)
+    if isinstance(dat, int): sdats.append(f"{field_name}:0x{dat:X}")
+    else: sdats.append(f"{field_name}:{dat}")
+  return sdats
+
+def install_hook(c_function, python_function):
+  python_function_addr = ctypes.cast(ctypes.byref(python_function), ctypes.POINTER(ctypes.c_ulong)).contents.value
+  # AARCH64 trampoline to ioctl
+  if processor == "aarch64":
+    # 0x0000000000000000:  70 00 00 10    adr x16, #0xc
+    # 0x0000000000000004:  10 02 40 F9    ldr x16, [x16]
+    # 0x0000000000000008:  00 02 1F D6    br  x16
+    tramp = b"\x70\x00\x00\x10\x10\x02\x40\xf9\x00\x02\x1f\xd6"
+    tramp += struct.pack("Q", python_function_addr)
+  elif processor == "x86_64":
+    # 0x0000000000000000:  49 B8 aa aa aa aa aa aa aa aa    movabs r8, <address>
+    # 0x000000000000000a:  41 FF E0                         jmp    r8
+    tramp = b"\x49\xB8" + struct.pack("Q", python_function_addr) + b"\x41\xFF\xE0"
+  else:
+    raise Exception(f"processor {processor} not supported")
+
+  # get real ioctl address
+  ioctl_address = ctypes.cast(ctypes.byref(c_function), ctypes.POINTER(ctypes.c_ulong))
+
+  # hook ioctl
+  ret = libc.mprotect(ctypes.c_ulong((ioctl_address.contents.value//0x1000)*0x1000), 0x2000, 7)
+  assert ret == 0
+  libc.memcpy(ioctl_address.contents, ctypes.create_string_buffer(tramp), len(tramp))
+
+# *** ioctl lib end ***
+
+import tinygrad.runtime.autogen.kfd as kfd_ioctl
+import tinygrad.runtime.autogen.hsa as hsa
+
+def print_aql_queue(read_pointer_address):
+  rptr_offset = getattr(hsa.amd_queue_v2_t, 'read_dispatch_id').offset
+  queue_base = read_pointer_address - rptr_offset
+  queue = hsa.amd_queue_v2_t.from_address(queue_base)
+  print(f"  AQL Queue @ 0x{queue_base:X}:")
+  for field_name, *_ in hsa.amd_queue_v2_t._real_fields_:
+    val = getattr(queue, field_name)
+    if isinstance(val, int): print(f"    {field_name}: 0x{val:X}")
+    elif hasattr(val, '_length_'):
+      arr_vals = [f"{format_struct(v)}" if hasattr(v, '_real_fields_') else f"{v:#X}" for v in val]
+      print(f"    {field_name}: [{', '.join(arr_vals)}]")
+    elif hasattr(val, '_real_fields_'): print(f"    {field_name}: {format_struct(val)}")
+    else: print(f"    {field_name}: {val}")
+
+def ioctls_from_header():
+  hdr = (pathlib.Path(__file__).parent / "kfd_ioctl.h").read_text().replace("\\\n", "")
+  pattern = r'#define\s+(AMDKFD_IOC_[A-Z0-9_]+)\s+AMDKFD_IOW?R?\((0x[0-9a-fA-F]+),\s+struct\s([A-Za-z0-9_]+)\)'
+  matches = re.findall(pattern, hdr, re.MULTILINE)
+  return {int(nr, 0x10):(name, getattr(kfd_ioctl, "struct_"+sname, None)) for name, nr, sname in matches}
+nrs = ioctls_from_header()
+
+@ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int, ctypes.c_ulong, ctypes.c_void_p)
+def ioctl(fd, request, argp):
+  st = time.perf_counter()
+  ret = libc.syscall(IOCTL_SYSCALL, ctypes.c_int(fd), ctypes.c_ulong(request), ctypes.c_void_p(argp))
+  et = time.perf_counter()-st
+  idir, size, itype, nr = (request>>30), (request>>16)&0x3FFF, (request>>8)&0xFF, request&0xFF
+  if nr in nrs and itype == 75:
+    # /dev/kfd
+    name, stype = nrs[nr]
+    s = get_struct(argp, stype)
+    print(f"{(st-start)*1000:7.2f} ms +{et*1000.:7.2f} ms : {ret:2d} = {name:40s}", ' '.join(format_struct(s)))
+    if name == "AMDKFD_IOC_SVM":
+      out = ctypes.cast(s.attrs, ctypes.POINTER(kfd_ioctl.struct_kfd_ioctl_svm_attribute))
+      for i in range(s.nattr): print(f"{i}: {kfd_ioctl.enum_kfd_ioctl_svm_attr_type.get(out[i].type):40s}: {out[i].value:#x}")
+    if name == "AMDKFD_IOC_CREATE_QUEUE" and s.queue_type == kfd_ioctl.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL: print_aql_queue(s.read_pointer_address)
+  else:
+    print(f"{(st-start)*1000:7.2f} ms +{et*1000.:7.2f} ms : ioctl",
+          f"{idir=} {size=} {itype=} {nr=} {fd=} {ret=}", os.readlink(f"/proc/self/fd/{fd}") if fd >= 0 else "")
+  return ret
+
+install_hook(libc.ioctl, ioctl)
+
+# AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7
+if __name__ == "__main__":
+  print("***** import tinygrad")
+  from tinygrad import Tensor, Device, TinyJit
+  print("***** access HIP")
+  dev = Device["HIP"]
+  print("***** create tensor a")
+  a = Tensor([1.,2.]*1024*1024, device="HIP").realize()
+  print("***** create tensor b")
+  b = Tensor([3.,4.]*1024*1024, device="HIP").realize()
+  @TinyJit
+  def add(a, b): return (a+b).realize()
+  for i in range(4):
+    print(f"***** add tensors {i}")
+    c = add(a, b)
+    #dev.synchronize()
+    c = add(b, a)
+    dev.synchronize()
+  print(f"***** copyout")
+  nc = c.numpy()
+  print(f"***** delete")
+  del add, a, b, c, dev
+  print(f"***** done")
+  os._exit(0)

tinygrad_repo/extra/hip_gpu_driver/test_kfd_2.py

@@ -0,0 +1,207 @@
+import os, ctypes, pathlib, re, fcntl, functools, mmap, time
+import tinygrad.runtime.autogen.kfd as kfd
+from tinygrad.helpers import to_mv, getenv
+from extra.hip_gpu_driver import hip_ioctl
+import tinygrad.runtime.autogen.hsa as hsa
+from hexdump import hexdump
+
+libc = ctypes.CDLL("libc.so.6")
+libc.memset.argtypes = [ctypes.c_void_p, ctypes.c_char, ctypes.c_int]
+libc.mmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_long]
+libc.mmap.restype = ctypes.c_void_p
+MAP_NORESERVE = 0x4000
+MAP_FIXED = 0x10
+
+def kfd_ioctl(idir, nr, user_struct, fd, **kwargs):
+  made = user_struct(**kwargs)
+  ret = fcntl.ioctl(fd, (idir<<30) | (ctypes.sizeof(user_struct)<<16) | (ord('K')<<8) | nr, made)
+  if ret != 0: raise RuntimeError(f"ioctl returned {ret}")
+  return made
+
+def format_struct(s):
+  sdats = []
+  for field_name, field_type in s._fields_:
+    dat = getattr(s, field_name)
+    if isinstance(dat, int): sdats.append(f"{field_name}:0x{dat:X}")
+    else: sdats.append(f"{field_name}:{dat}")
+  return sdats
+
+idirs = {"IOW": 1, "IOR": 2, "IOWR": 3}
+def ioctls_from_header():
+  hdr = pathlib.Path("/usr/include/linux/kfd_ioctl.h").read_text().replace("\\\n", "")
+  pattern = r'#define\s+(AMDKFD_IOC_[A-Z0-9_]+)\s+AMDKFD_(IOW?R?)\((0x[0-9a-fA-F]+),\s+struct\s([A-Za-z0-9_]+)\)'
+  matches = re.findall(pattern, hdr, re.MULTILINE)
+
+  fxns = {}
+  for name, idir, nr, sname in matches:
+    fxns[name.replace("AMDKFD_IOC_", "").lower()] = functools.partial(kfd_ioctl, idirs[idir], int(nr, 0x10), getattr(kfd, "struct_"+sname))
+  return type("KIO", (object, ), fxns)
+kio = ioctls_from_header()
+
+# sudo su -c "echo 'file drivers/gpu/drm/amd/* +p' > /sys/kernel/debug/dynamic_debug/control"
+
+def gpu_alloc_userptr(fd, size, flags):
+  addr = libc.mmap(0, size, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED|mmap.MAP_ANONYMOUS, -1, 0)
+  assert addr != 0xffffffffffffffff
+  mem = kio.alloc_memory_of_gpu(fd, va_addr=addr, size=size, gpu_id=GPU_ID, flags=flags, mmap_offset=addr)
+  return mem
+
+def gpu_alloc(fd, size, flags):
+  addr = libc.mmap(0, size, 0, mmap.MAP_PRIVATE|mmap.MAP_ANONYMOUS|MAP_NORESERVE, -1, 0)
+  assert addr != 0xffffffffffffffff
+  mem = kio.alloc_memory_of_gpu(fd, va_addr=addr, size=size, gpu_id=GPU_ID, flags=flags)
+  buf = libc.mmap(mem.va_addr, mem.size, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED|MAP_FIXED, drm_fd, mem.mmap_offset)
+  assert buf != 0xffffffffffffffff
+  assert addr == buf == mem.va_addr
+  return mem
+
+if __name__ == "__main__":
+  fd = os.open("/dev/kfd", os.O_RDWR)
+  gpu_num = getenv("GPU", 0)
+  drm_fd = os.open(f"/dev/dri/renderD{128+gpu_num}", os.O_RDWR)
+  with open(f"/sys/devices/virtual/kfd/kfd/topology/nodes/{1+gpu_num}/gpu_id", "r") as f: GPU_ID = int(f.read())
+
+  #ver = kio.get_version(fd)
+  st = kio.acquire_vm(fd, drm_fd=drm_fd, gpu_id=GPU_ID)
+  #exit(0)
+
+  # 0xF0000001 = KFD_IOC_ALLOC_MEM_FLAGS_VRAM | KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE | KFD_IOC_ALLOC_MEM_FLAGS_EXECUTABLE | KFD_IOC_ALLOC_MEM_FLAGS_PUBLIC | KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE
+  # 0xD6000002 = KFD_IOC_ALLOC_MEM_FLAGS_GTT | KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE | KFD_IOC_ALLOC_MEM_FLAGS_EXECUTABLE | KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE
+  # 0xD6000004 = KFD_IOC_ALLOC_MEM_FLAGS_USERPTR | KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE | KFD_IOC_ALLOC_MEM_FLAGS_EXECUTABLE | KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE
+  # 0x94000010 = KFD_IOC_ALLOC_MEM_FLAGS_MMIO_REMAP | KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE | KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE
+  #addr = libc.mmap(0, 0x1000, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_PRIVATE|mmap.MAP_ANONYMOUS, -1, 0)
+  #addr = libc.mmap(0, 0x1000, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED|mmap.MAP_ANONYMOUS, -1, 0)
+  #mem = kio.AMDKFD_IOC_ALLOC_MEMORY_OF_GPU(fd, va_addr=addr, size=0x1000, gpu_id=GPU_ID, flags=0xD6000004)
+
+  #mem = gpu_alloc(fd, 0x1000, kfd.KFD_IOC_ALLOC_MEM_FLAGS_VRAM |
+  #                            kfd.KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_EXECUTABLE |
+  #                            kfd.KFD_IOC_ALLOC_MEM_FLAGS_PUBLIC | kfd.KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE)
+  #arr = (ctypes.c_int32 * 1)(GPU_ID)
+  #stm = kio.map_memory_to_gpu(fd, handle=mem.handle, device_ids_array_ptr=ctypes.addressof(arr), n_devices=1)
+
+  arr = (ctypes.c_int32 * 1)(GPU_ID)
+  rw_ptr = gpu_alloc(fd, 0x1000, kfd.KFD_IOC_ALLOC_MEM_FLAGS_GTT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE |
+                                 kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED |
+                                 kfd.KFD_IOC_ALLOC_MEM_FLAGS_EXECUTABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE)
+  stm = kio.map_memory_to_gpu(fd, handle=rw_ptr.handle, device_ids_array_ptr=ctypes.addressof(arr), n_devices=1)
+  assert stm.n_success == 1
+  event_page = gpu_alloc(fd, 0x8000, kfd.KFD_IOC_ALLOC_MEM_FLAGS_GTT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE |
+                                 kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED |
+                                 kfd.KFD_IOC_ALLOC_MEM_FLAGS_EXECUTABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE)
+  stm = kio.map_memory_to_gpu(fd, handle=event_page.handle, device_ids_array_ptr=ctypes.addressof(arr), n_devices=1)
+  assert stm.n_success == 1
+  ring_base = gpu_alloc_userptr(fd, 0x1000, kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR | kfd.KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE |
+                                    kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED |
+                                    kfd.KFD_IOC_ALLOC_MEM_FLAGS_EXECUTABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE)
+  stm = kio.map_memory_to_gpu(fd, handle=ring_base.handle, device_ids_array_ptr=ctypes.addressof(arr), n_devices=1)
+  assert stm.n_success == 1
+  signals = gpu_alloc_userptr(fd, 0x1000, kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR | kfd.KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE |
+                                    kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED |
+                                    kfd.KFD_IOC_ALLOC_MEM_FLAGS_EXECUTABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE)
+  stm = kio.map_memory_to_gpu(fd, handle=signals.handle, device_ids_array_ptr=ctypes.addressof(arr), n_devices=1)
+  assert stm.n_success == 1
+  eop_buffer = gpu_alloc(fd, 0x1000, kfd.KFD_IOC_ALLOC_MEM_FLAGS_VRAM |
+                                     kfd.KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_EXECUTABLE |
+                                     kfd.KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE)
+  stm = kio.map_memory_to_gpu(fd, handle=eop_buffer.handle, device_ids_array_ptr=ctypes.addressof(arr), n_devices=1)
+  assert stm.n_success == 1
+  ctx_save_restore_address = gpu_alloc(fd, 0x2C02000, kfd.KFD_IOC_ALLOC_MEM_FLAGS_VRAM |
+                                     kfd.KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_EXECUTABLE |
+                                     kfd.KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE)
+  stm = kio.map_memory_to_gpu(fd, handle=ctx_save_restore_address.handle, device_ids_array_ptr=ctypes.addressof(arr), n_devices=1)
+  assert stm.n_success == 1
+
+  #113.00 ms +   0.00 ms :  0 = AMDKFD_IOC_CREATE_QUEUE                  ring_base_address:0x797465200000 write_pointer_address:0x79751C068038 read_pointer_address:0x79751C068080 doorbell_offset:0x0 ring_size:0x800000 gpu_id:0x433D queue_type:0x2 queue_per
+  #centage:0x64 queue_priority:0x7 queue_id:0x0 eop_buffer_address:0x79751C064000 eop_buffer_size:0x1000 ctx_save_restore_address:0x796E52400000 ctx_save_restore_size:0x2BEA000 ctl_stack_size:0xA000
+
+  #113.84 ms +   0.59 ms :  0 = AMDKFD_IOC_CREATE_QUEUE                  ring_base_address:0x71AC3F600000 write_pointer_address:0x71B302AB0038 read_pointer_address:0x71B302AB0080 doorbell_offset:0xD0CF400000000008 ring_size:0x800000 gpu_id:0x433D queue_typ
+  #e:0x2 queue_percentage:0x64 queue_priority:0x7 queue_id:0x1 eop_buffer_address:0x71B302AAC000 eop_buffer_size:0x1000 ctx_save_restore_address:0x71AC3C800000 ctx_save_restore_size:0x2BEA000 ctl_stack_size:0xA000
+
+  #define KFD_MMAP_TYPE_SHIFT	62
+  #define KFD_MMAP_TYPE_DOORBELL	(0x3ULL << KFD_MMAP_TYPE_SHIFT)
+  evt = kio.create_event(fd, event_page_offset=event_page.handle, auto_reset=1)
+
+  nq = kio.create_queue(fd, ring_base_address=ring_base.va_addr, ring_size=0x1000, gpu_id=GPU_ID,
+                        queue_type=kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL, queue_percentage=kfd.KFD_MAX_QUEUE_PERCENTAGE,
+                        queue_priority=kfd.KFD_MAX_QUEUE_PRIORITY,
+                        eop_buffer_address=eop_buffer.va_addr, eop_buffer_size=0x1000,
+                        ctx_save_restore_address=ctx_save_restore_address.va_addr, ctx_save_restore_size=0x2C02000,
+                        ctl_stack_size = 0xa000,
+                        # write_pointer_address and read_pointer_address are on GART
+                        #write_pointer_address=0xaaaabbbb, read_pointer_address=0xaaaacccc)
+                        write_pointer_address=rw_ptr.va_addr+0, read_pointer_address=rw_ptr.va_addr+0x8)
+  doorbell = libc.mmap(0, 8192, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED, fd, nq.doorbell_offset)
+  print("doorbell", hex(doorbell))
+
+  to_mv(signals.va_addr, 0x40)
+
+  """
+  hexdump(to_mv(event_page.va_addr, 0x40))
+  kio.set_event(fd, event_id=evt.event_id)
+  hexdump(to_mv(event_page.va_addr, 0x40))
+  kio.reset_event(fd, event_id=evt.event_id)
+  hexdump(to_mv(event_page.va_addr, 0x40))
+  """
+
+  # KFD_EVENT_TYPE_SIGNAL
+
+  BARRIER_HEADER  = 1 << hsa.HSA_PACKET_HEADER_BARRIER
+  BARRIER_HEADER |= hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCACQUIRE_FENCE_SCOPE
+  BARRIER_HEADER |= hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCRELEASE_FENCE_SCOPE
+  BARRIER_HEADER |= hsa.HSA_PACKET_TYPE_BARRIER_AND << hsa.HSA_PACKET_HEADER_TYPE
+
+  AQL_PACKET_SIZE = ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t)
+  EMPTY_SIGNAL = hsa.hsa_signal_t()
+
+  ds = to_mv(rw_ptr.va_addr, 0x100).cast("Q")
+  ds[0] = 1 #ring_base.va_addr + AQL_PACKET_SIZE
+  ds[1] = 0 #ring_base.va_addr
+  #libc.memset(rw_ptr.va_addr, 0xaa, 0x100)
+  #hexdump(to_mv(rw_ptr.va_addr, 0x100))
+
+  #packet = hsa.hsa_barrier_and_packet_t.from_address(rw_ptr.va_addr+0x38)
+  packet = hsa.hsa_barrier_and_packet_t.from_address(ring_base.va_addr)
+  packet.reserved0 = 0
+  packet.reserved1 = 0
+  for i in range(5): packet.dep_signal[i] = EMPTY_SIGNAL
+  #packet.dep_signal[0] = hsa.hsa_signal_t(evt.event_id)
+  packet.reserved2 = 0
+  #packet.completion_signal = EMPTY_SIGNAL
+  packet.completion_signal = hsa.hsa_signal_t(signals.va_addr)
+  packet.header = BARRIER_HEADER
+  hexdump(to_mv(ring_base.va_addr, AQL_PACKET_SIZE))
+
+  # _HsaEventData
+  to_mv(signals.va_addr, 0x40).cast("Q")[0] = 1
+  to_mv(signals.va_addr, 0x40).cast("Q")[1] = 1
+  #to_mv(signals.va_addr, 0x40).cast("Q")[2] = event_page
+  to_mv(signals.va_addr, 0x40).cast("Q")[2] = event_page.va_addr + evt.event_slot_index*8  # HWData2=HWAddress
+  to_mv(signals.va_addr, 0x40).cast("Q")[3] = evt.event_trigger_data # HWData3=HWData
+  print(hex(ds[0]), hex(ds[1]), hex(ds[2]))
+  hexdump(to_mv(signals.va_addr, 0x40))
+
+  # 10 08 49 3E 46 77 00 00
+
+
+  # ring doorbell
+  print(hex(to_mv(doorbell, 0x10).cast("I")[0]))
+  #to_mv(doorbell, 0x10).cast("I")[0] = 0xffffffff
+  to_mv(doorbell, 0x10).cast("I")[0] = 0
+
+  evt_arr = (kfd.struct_kfd_event_data * 1)()
+  evt_arr[0].event_id = evt.event_id
+  kio.wait_events(fd, events_ptr=ctypes.addressof(evt_arr), num_events=1, wait_for_all=0, timeout=1000)
+
+  print(hex(ds[0]), hex(ds[1]), hex(ds[2]))
+  hexdump(to_mv(signals.va_addr, 0x40))
+
+  #nq = kio.create_queue(fd, ring_base_address=buf, ring_size=0x1000, gpu_id=GPU_ID,
+  #                      queue_type=kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL, queue_percentage=kfd.KFD_MAX_QUEUE_PERCENTAGE,
+  #                      queue_priority=kfd.KFD_MAX_QUEUE_PRIORITY, write_pointer_address=buf+8, read_pointer_address=buf+0x10)
+  #print(nq)
+
+  #mv = to_mv(buf, 0x1000)
+  #addr = libc.mmap(0, 0x1000, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_PRIVATE|mmap.MAP_ANONYMOUS, -1, 0)
+
+  #print('\n'.join(format_struct(ver)))
+  #print('\n'.join(format_struct(st)))

tinygrad_repo/extra/hip_gpu_driver/test_pm4.py

@@ -0,0 +1,156 @@
+import time
+from hexdump import hexdump
+from tinygrad import Tensor, Device
+import tinygrad.runtime.autogen.amd_gpu as amd_gpu
+import tinygrad.runtime.autogen.kfd as kfd
+import tinygrad.runtime.autogen.hsa as hsa
+from tinygrad.runtime.ops_amd import kio, AMDProgram
+from tinygrad.helpers import to_mv
+
+DISPATCH_INIT_VALUE = 0x21 | 0x8000
+
+#mmCOMPUTE_START_X = 0x2e04
+#mmCOMPUTE_PGM_LO = 0x2e0c
+
+BASE_ADDR = 0x00001260
+PACKET3_SET_SH_REG_START = 0x2c00
+SUB = PACKET3_SET_SH_REG_START - BASE_ADDR
+
+regCOMPUTE_PGM_LO = 0x1bac - SUB
+regCOMPUTE_START_X = 0x1ba4 - SUB
+regCOMPUTE_NUM_THREAD_X = 0x1ba7 - SUB
+regCOMPUTE_USER_DATA_0 = 0x1be0 - SUB
+regCOMPUTE_USER_DATA_8 = 0x1be8 - SUB
+
+regCOMPUTE_PGM_RSRC1 = 0x1bb2 - SUB
+regCOMPUTE_PGM_RSRC2 = 0x1bb3 - SUB
+
+# DEBUG=6 python3 extra/hip_gpu_driver/test_pm4.py
+# sudo umr -i 1 -s amd744c.gfx1100 --sbank 1 1 2 | grep regCOMPUTE
+
+# 0x00009025
+
+COMPUTE_SHADER_EN = 1
+USE_THREAD_DIMENSIONS = 1 << 5
+CS_W32_EN = 1 << 15
+
+def format_struct(s):
+  sdats = []
+  for field_name, field_type in s._fields_:
+    dat = getattr(s, field_name)
+    if isinstance(dat, int): sdats.append(f"{field_name}:0x{dat:X}")
+    else: sdats.append(f"{field_name}:{dat}")
+  return sdats
+
+if __name__ == "__main__":
+  dev = Device["KFD"]
+
+  a = Tensor([0.,1.,2.], device="KFD").realize()
+  b = a + 7
+  b.uop.buffer.allocate()
+  si = b.schedule()[-1]
+  runner = dev.get_runner(*si.ast)
+  prg: AMDProgram = runner.clprg
+  print("device initted")
+
+  # Compute Queue
+
+  gart_compute = dev._gpu_alloc(0x1000, kfd.KFD_IOC_ALLOC_MEM_FLAGS_GTT, uncached=True)
+  eop_buffer = dev._gpu_alloc(0x1000, kfd.KFD_IOC_ALLOC_MEM_FLAGS_VRAM)
+  compute_ring = dev._gpu_alloc(0x800000, kfd.KFD_IOC_ALLOC_MEM_FLAGS_GTT, uncached=True)
+  ctx_save_restore_address = dev._gpu_alloc(0x2C02000, kfd.KFD_IOC_ALLOC_MEM_FLAGS_VRAM)
+  compute_queue = kio.create_queue(dev.kfd, ring_base_address=compute_ring.va_addr, ring_size=compute_ring.size, gpu_id=dev.gpu_id,
+    queue_type=kfd.KFD_IOC_QUEUE_TYPE_COMPUTE, queue_percentage=kfd.KFD_MAX_QUEUE_PERCENTAGE, queue_priority=kfd.KFD_MAX_QUEUE_PRIORITY,
+    #eop_buffer_address=eop_buffer.va_addr, eop_buffer_size=eop_buffer.size,
+    #ctx_save_restore_address=ctx_save_restore_address.va_addr, ctx_save_restore_size=ctx_save_restore_address.size,
+    #ctl_stack_size = 0xa000,
+    write_pointer_address=gart_compute.va_addr, read_pointer_address=gart_compute.va_addr+8)
+  compute_doorbell = to_mv(dev.doorbells + compute_queue.doorbell_offset - dev.doorbells_base, 4).cast("I")
+
+  #scratch = dev._gpu_alloc(0x10000, kfd.KFD_IOC_ALLOC_MEM_FLAGS_VRAM)
+  ka = to_mv(dev.kernargs_ptr, 0x10).cast("Q")
+  ka[0] = b.uop.buffer._buf.va_addr
+  ka[1] = a.uop.buffer._buf.va_addr
+
+  compute_read_pointer = to_mv(compute_queue.read_pointer_address, 8).cast("Q")
+  compute_write_pointer = to_mv(compute_queue.write_pointer_address, 8).cast("Q")
+
+  hexdump(to_mv(prg.handle, 0x40))
+  code = hsa.amd_kernel_code_t.from_address(prg.handle)
+
+  #print(format_struct(code))
+  #print("code")
+  #hexdump(to_mv(code_ptr, 0x100))
+  #runner.local_size = [2,1,1]
+
+  print(runner.local_size, runner.global_size)
+
+  #pm4_cmd += [amd_gpu.PACKET3(amd_gpu.PACKET3_SET_SH_REG, 6), mmCOMPUTE_PGM_LO,
+  #  prg.handle&0xFFFFFFFF, prg.handle>>32, 0, 0, (scratch.va_addr>>8)&0xFFFFFFFF, scratch.va_addr>>40]
+  code_ptr = (prg.handle + code.kernel_code_entry_byte_offset) >> 8
+  pm4_cmd  = [amd_gpu.PACKET3(amd_gpu.PACKET3_SET_SH_REG, 6), regCOMPUTE_PGM_LO, code_ptr&0xFFFFFFFF, code_ptr>>32, 0, 0, 0, 0]
+  pm4_cmd += [amd_gpu.PACKET3(amd_gpu.PACKET3_SET_SH_REG, 2), regCOMPUTE_PGM_RSRC1, code.compute_pgm_rsrc1, code.compute_pgm_rsrc2]
+  pm4_cmd += [amd_gpu.PACKET3(amd_gpu.PACKET3_SET_SH_REG, 2), regCOMPUTE_USER_DATA_0, dev.kernargs_ptr&0xFFFFFFFF, dev.kernargs_ptr>>32]
+  #pm4_cmd += [amd_gpu.PACKET3(amd_gpu.PACKET3_SET_SH_REG, 2), regCOMPUTE_USER_DATA_0, 0, 0]
+  pm4_cmd += [amd_gpu.PACKET3(amd_gpu.PACKET3_SET_SH_REG, 8), regCOMPUTE_START_X, 0,0,0,
+              runner.local_size[0],runner.local_size[1],runner.local_size[2],0,0]
+  # disabled USE_THREAD_DIMENSIONS
+  pm4_cmd += [amd_gpu.PACKET3(amd_gpu.PACKET3_DISPATCH_DIRECT, 3),
+              runner.global_size[0],runner.global_size[1],runner.global_size[2], CS_W32_EN | COMPUTE_SHADER_EN]
+
+  #pm4_cmd = [amd_gpu.PACKET3(amd_gpu.PACKET3_NOP, 0x3fff)]*0x200
+
+  """
+  addr=0x0
+  sz=(1 << 64)-1
+  gli=0
+  glv=0
+  glk=0
+  gl1=0
+  gl2=0
+  pm4_cmd = [amd_gpu.PACKET3(amd_gpu.PACKET3_ACQUIRE_MEM, 6), 0,
+              sz & 0xffffffff, (sz >> 32) & 0xff, addr & 0xffffffff, (addr >> 32) & 0xffffff, 0,
+              amd_gpu.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLI_INV(gli) | amd_gpu.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLK_INV(glk) | \
+              amd_gpu.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLV_INV(glv) | amd_gpu.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL1_INV(gl1) | \
+              amd_gpu.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL2_INV(gl2)]
+  print(pm4_cmd)
+  """
+
+  wptr = 0
+  pm4_buffer_view = to_mv(compute_ring.va_addr, compute_ring.size).cast("I")
+
+  for j in range(0x80000):
+    for i, value in enumerate(pm4_cmd): pm4_buffer_view[wptr+i] = value
+    wptr += len(pm4_cmd)
+
+    compute_write_pointer[0] = wptr
+    compute_doorbell[0] = wptr
+    for k in range(10):
+      done = compute_read_pointer[0] == compute_write_pointer[0]
+      print(compute_read_pointer[0], compute_write_pointer[0], done)
+      if done: break
+      time.sleep(0.01)
+    break
+    #break
+
+    #print(compute_read_pointer[0])
+    #time.sleep(0.05)
+    #print(compute_read_pointer[0])
+
+  #time.sleep(100)
+
+  print(a.numpy())
+  print(b.numpy())
+  exit(0)
+
+  #pm4_cmd = [amd_gpu.PACKET3(amd_gpu.PACKET3_SET_SH_REG, 8), mmCOMPUTE_PGM_LO, 0,0,0,1,1,1,0,0]
+
+
+  #pm4_cmd += [amd_gpu.PACKET3(amd_gpu.PACKET3_DISPATCH_DIRECT, )]
+
+
+  #pm4_cmd = [amd_gpu.PACKET3(amd_gpu.PACKET3_ACQUIRE_MEM, 6), 0,
+  #          sz & 0xffffffff, (sz >> 32) & 0xff, addr & 0xffffffff, (addr >> 32) & 0xffffff, 0,
+  #          amd_gpu.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLI_INV(gli) | amd_gpu.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLK_INV(glk) | \
+  #          amd_gpu.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLV_INV(glv) | amd_gpu.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL1_INV(gl1) | \
+  #          amd_gpu.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL2_INV(gl2)]